diff --git a/.bazelrc b/.bazelrc
deleted file mode 100644
index fc2995dc838c5..0000000000000
--- a/.bazelrc
+++ /dev/null
@@ -1,114 +0,0 @@
-build --cxxopt=--std=c++17
-build --copt=-I.
-# Bazel does not support including its cc_library targets as system
-# headers. We work around this for generated code
-# (e.g. torch/headeronly/macros/cmake_macros.h) by making the generated directory a
-# system include path.
-build --copt=-isystem --copt bazel-out/k8-fastbuild/bin
-build --copt=-isystem --copt bazel-out/darwin-fastbuild/bin
-build --experimental_ui_max_stdouterr_bytes=2048576
-
-# Configuration to disable tty features for environments like CI
-build:no-tty --curses no
-build:no-tty --progress_report_interval 10
-build:no-tty --show_progress_rate_limit 10
-
-# Build with GPU support by default.
-build --define=cuda=true
-# rules_cuda configuration
-build --@rules_cuda//cuda:enable_cuda
-build --@rules_cuda//cuda:cuda_targets=sm_52
-build --@rules_cuda//cuda:compiler=nvcc
-build --repo_env=CUDA_PATH=/usr/local/cuda
-
-# Configuration to build without GPU support
-build:cpu-only --define=cuda=false
-# define a separate build folder for faster switching between configs
-build:cpu-only --platform_suffix=-cpu-only
-# See the note on the config-less build for details about why we are
-# doing this. We must also do it for the "-cpu-only" platform suffix.
-build --copt=-isystem --copt=bazel-out/k8-fastbuild-cpu-only/bin
-# rules_cuda configuration
-build:cpu-only --@rules_cuda//cuda:enable_cuda=False
-
-# Definition of --config=shell
-# interactive shell immediately before execution
-build:shell --run_under="//tools/bazel_tools:shellwrap"
-
-# Disable all warnings for external repositories. We don't care about
-# their warnings.
-build --per_file_copt=^external/@-w
-
-# Set additional warnings to error level.
-#
-# Implementation notes:
-#  * we use file extensions to determine if we are using the C++
-#    compiler or the cuda compiler
-#  * we use ^// at the start of the regex to only permit matching
-#    PyTorch files. This excludes external repos.
-#
-# Note that because this is logically a command-line flag, it is
-# considered the word on what warnings are enabled. This has the
-# unfortunate consequence of preventing us from disabling an error at
-# the target level because those flags will come before these flags in
-# the action invocation. Instead we provide per-file exceptions after
-# this.
-#
-# On the bright side, this means we don't have to more broadly apply
-# the exceptions to an entire target.
-#
-# Looking for CUDA flags? We have a cu_library macro that we can edit
-# directly. Look in //tools/rules:cu.bzl for details. Editing the
-# macro over this has the following advantages:
-#  * making changes does not require discarding the Bazel analysis
-#    cache
-#  * it allows for selective overrides on individual targets since the
-#    macro-level opts will come earlier than target level overrides
-
-build --per_file_copt='^//.*\.(cpp|cc)$'@-Werror=all
-# The following warnings come from -Wall. We downgrade them from error
-# to warnings here.
-#
-# We intentionally use #pragma unroll, which is compiler specific.
-build --per_file_copt='^//.*\.(cpp|cc)$'@-Wno-error=unknown-pragmas
-
-build --per_file_copt='^//.*\.(cpp|cc)$'@-Werror=extra
-# The following warnings come from -Wextra. We downgrade them from error
-# to warnings here.
-#
-# unused-parameter-compare has a tremendous amount of violations in the
-# codebase. It will be a lot of work to fix them, just disable it for
-# now.
-build --per_file_copt='^//.*\.(cpp|cc)$'@-Wno-unused-parameter
-# missing-field-parameters has both a large number of violations in
-# the codebase, but it also is used pervasively in the Python C
-# API. There are a couple of catches though:
-# * we use multiple versions of the Python API and hence have
-#   potentially multiple different versions of each relevant
-#   struct. They may have different numbers of fields. It will be
-#   unwieldy to support multiple versions in the same source file.
-# * Python itself for many of these structs recommends only
-#   initializing a subset of the fields. We should respect the API
-#   usage conventions of our dependencies.
-#
-# Hence, we just disable this warning altogether. We may want to clean
-# up some of the clear-cut cases that could be risky, but we still
-# likely want to have this disabled for the most part.
-build --per_file_copt='^//.*\.(cpp|cc)$'@-Wno-missing-field-initializers
-
-build --per_file_copt='^//.*\.(cpp|cc)$'@-Wno-unused-function
-build --per_file_copt='^//.*\.(cpp|cc)$'@-Wno-unused-variable
-
-build --per_file_copt='//:aten/src/ATen/RegisterCompositeExplicitAutograd\.cpp$'@-Wno-error=unused-function
-build --per_file_copt='//:aten/src/ATen/RegisterCompositeImplicitAutograd\.cpp$'@-Wno-error=unused-function
-build --per_file_copt='//:aten/src/ATen/RegisterMkldnnCPU\.cpp$'@-Wno-error=unused-function
-build --per_file_copt='//:aten/src/ATen/RegisterNestedTensorCPU\.cpp$'@-Wno-error=unused-function
-build --per_file_copt='//:aten/src/ATen/RegisterQuantizedCPU\.cpp$'@-Wno-error=unused-function
-build --per_file_copt='//:aten/src/ATen/RegisterSparseCPU\.cpp$'@-Wno-error=unused-function
-build --per_file_copt='//:aten/src/ATen/RegisterSparseCsrCPU\.cpp$'@-Wno-error=unused-function
-build --per_file_copt='//:aten/src/ATen/RegisterNestedTensorMeta\.cpp$'@-Wno-error=unused-function
-build --per_file_copt='//:aten/src/ATen/RegisterSparseMeta\.cpp$'@-Wno-error=unused-function
-build --per_file_copt='//:aten/src/ATen/RegisterQuantizedMeta\.cpp$'@-Wno-error=unused-function
-build --per_file_copt='//:aten/src/ATen/RegisterZeroTensor\.cpp$'@-Wno-error=unused-function
-build --per_file_copt='//:torch/csrc/lazy/generated/RegisterAutogradLazy\.cpp$'@-Wno-error=unused-function
-build --per_file_copt='//:torch/csrc/lazy/generated/RegisterLazy\.cpp$'@-Wno-error=unused-function
diff --git a/.bazelversion b/.bazelversion
deleted file mode 100644
index f22d756da39d4..0000000000000
--- a/.bazelversion
+++ /dev/null
@@ -1 +0,0 @@
-6.5.0
diff --git a/.ci/docker/README.md b/.ci/docker/README.md
index 32d3a21a1897b..56d6ff1f2fe27 100644
--- a/.ci/docker/README.md
+++ b/.ci/docker/README.md
@@ -25,7 +25,6 @@ See `build.sh` for valid build environments (it's the giant switch).
 
 * `conda` - Dockerfile and build.sh to build Docker images used in nightly conda builds
 * `manywheel` - Dockerfile and build.sh to build Docker images used in nightly manywheel builds
-* `libtorch` - Dockerfile and build.sh to build Docker images used in nightly libtorch builds
 
 ## Usage
 
@@ -108,8 +107,6 @@ If your new Docker image needs a library installed from a specific pinned commit
      GCC_VERSION=11
      VISION=yes
      KATEX=yes
-     UCX_COMMIT=${_UCX_COMMIT}
-     UCC_COMMIT=${_UCC_COMMIT}
      TRITON=yes
      NEW_ARG_1=yes
      ;;
diff --git a/.ci/docker/almalinux/Dockerfile b/.ci/docker/almalinux/Dockerfile
index 08a43513a0904..d061353187c55 100644
--- a/.ci/docker/almalinux/Dockerfile
+++ b/.ci/docker/almalinux/Dockerfile
@@ -19,15 +19,12 @@ RUN git config --global --add safe.directory '*'
 ENV PATH=/opt/rh/gcc-toolset-${DEVTOOLSET_VERSION}/root/usr/bin:$PATH
 
 # cmake-3.18.4 from pip
+# NS: Apr 1 2026 3.18.4 is gone, reported here https://github.com/scikit-build/cmake-python-distributions/issues/693
 RUN yum install -y python3-pip && \
-    python3 -mpip install cmake==3.18.4 && \
+    python3 -mpip install cmake==3.18.4.post1 && \
     ln -s /usr/local/bin/cmake /usr/bin/cmake3
 RUN rm -rf /usr/local/cuda-*
 
-FROM base as openssl
-ADD ./common/install_openssl.sh install_openssl.sh
-RUN bash ./install_openssl.sh && rm install_openssl.sh
-
 FROM base as patchelf
 # Install patchelf
 ADD ./common/install_patchelf.sh install_patchelf.sh
@@ -53,7 +50,6 @@ 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
 
-
 FROM cuda as cuda12.6
 RUN bash ./install_cuda.sh 12.6
 ENV DESIRED_CUDA=12.6
@@ -70,6 +66,10 @@ FROM cuda as cuda13.0
 RUN bash ./install_cuda.sh 13.0
 ENV DESIRED_CUDA=13.0
 
+FROM cuda as cuda13.2
+RUN bash ./install_cuda.sh 13.2
+ENV DESIRED_CUDA=13.2
+
 FROM ${ROCM_IMAGE} as rocm_base
 ARG DEVTOOLSET_VERSION=13
 ENV LC_ALL en_US.UTF-8
@@ -81,6 +81,8 @@ RUN yum -y update && \
     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 '*'
+# All rocm clang cfg files load the same rocm.cfg, make sure it points to the right toolchain.
+RUN echo "--gcc-toolchain=/opt/rh/gcc-toolset-${DEVTOOLSET_VERSION}/root/usr" >> /opt/rocm/llvm/bin/rocm.cfg
 ENV PATH=/opt/rh/gcc-toolset-${DEVTOOLSET_VERSION}/root/usr/bin:$PATH
 
 FROM rocm_base as rocm
@@ -101,11 +103,11 @@ COPY --from=cuda12.6  /usr/local/cuda-12.6 /usr/local/cuda-12.6
 COPY --from=cuda12.8  /usr/local/cuda-12.8 /usr/local/cuda-12.8
 COPY --from=cuda12.9  /usr/local/cuda-12.9 /usr/local/cuda-12.9
 COPY --from=cuda13.0  /usr/local/cuda-13.0 /usr/local/cuda-13.0
+COPY --from=cuda13.2  /usr/local/cuda-13.2 /usr/local/cuda-13.2
 
 # 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
 
diff --git a/.ci/docker/build.sh b/.ci/docker/build.sh
index 37c082e7d378e..1eca06471a110 100755
--- a/.ci/docker/build.sh
+++ b/.ci/docker/build.sh
@@ -81,13 +81,6 @@ elif [[ "$image" == *riscv* ]]; then
   DOCKERFILE="ubuntu-cross-riscv/Dockerfile"
 fi
 
-_UCX_COMMIT=7836b165abdbe468a2f607e7254011c07d788152
-_UCC_COMMIT=430e241bf5d38cbc73fc7a6b89155397232e3f96
-if [[ "$image" == *rocm* ]]; then
-  _UCX_COMMIT=29831d319e6be55cb8c768ca61de335c934ca39e
-  _UCC_COMMIT=9f4b242cbbd8b1462cbc732eb29316cdfa124b77
-fi
-
 tag=$(echo $image | awk -F':' '{print $2}')
 # If no tag (no colon in image name), use the image name itself
 if [[ -z "$tag" ]]; then
@@ -102,20 +95,14 @@ case "$tag" in
     CUDA_VERSION=12.4
     ANACONDA_PYTHON_VERSION=3.10
     GCC_VERSION=11
-    VISION=yes
     KATEX=yes
-    UCX_COMMIT=${_UCX_COMMIT}
-    UCC_COMMIT=${_UCC_COMMIT}
     TRITON=yes
     ;;
   pytorch-linux-jammy-cuda12.8-cudnn9-py3-gcc11)
     CUDA_VERSION=12.8.1
     ANACONDA_PYTHON_VERSION=3.10
     GCC_VERSION=11
-    VISION=yes
     KATEX=yes
-    UCX_COMMIT=${_UCX_COMMIT}
-    UCC_COMMIT=${_UCC_COMMIT}
     TRITON=yes
     INSTALL_MINGW=yes
     ;;
@@ -123,69 +110,48 @@ case "$tag" in
     CUDA_VERSION=13.0.2
     ANACONDA_PYTHON_VERSION=3.10
     GCC_VERSION=11
-    VISION=yes
     KATEX=yes
-    UCX_COMMIT=${_UCX_COMMIT}
-    UCC_COMMIT=${_UCC_COMMIT}
     TRITON=yes
-    ;;
-  pytorch-linux-jammy-cuda12.8-cudnn9-py3-gcc11-inductor-benchmarks)
-    CUDA_VERSION=12.8.1
-    ANACONDA_PYTHON_VERSION=3.10
-    GCC_VERSION=11
-    VISION=yes
-    KATEX=yes
-    UCX_COMMIT=${_UCX_COMMIT}
-    UCC_COMMIT=${_UCC_COMMIT}
-    TRITON=yes
-    INDUCTOR_BENCHMARKS=yes
+    INSTALL_MINGW=yes
     ;;
   pytorch-linux-jammy-cuda13.0-cudnn9-py3-gcc11-inductor-benchmarks)
     CUDA_VERSION=13.0.2
     ANACONDA_PYTHON_VERSION=3.10
     GCC_VERSION=11
-    VISION=yes
     KATEX=yes
-    UCX_COMMIT=${_UCX_COMMIT}
-    UCC_COMMIT=${_UCC_COMMIT}
     TRITON=yes
     INDUCTOR_BENCHMARKS=yes
     ;;
-  pytorch-linux-jammy-cuda12.9-cudnn9-py3.12-gcc11-vllm)
-    CUDA_VERSION=12.9.1
+  pytorch-linux-jammy-cuda13.0-cudnn9-py3.12-gcc11-vllm)
+    CUDA_VERSION=13.0.2
     ANACONDA_PYTHON_VERSION=3.12
     GCC_VERSION=11
-    VISION=yes
     KATEX=yes
-    UCX_COMMIT=${_UCX_COMMIT}
-    UCC_COMMIT=${_UCC_COMMIT}
     TRITON=yes
     ;;
-  pytorch-linux-jammy-py3-clang15-onnx)
+  pytorch-linux-jammy-py3.10-clang18)
     ANACONDA_PYTHON_VERSION=3.10
-    CLANG_VERSION=15
-    VISION=yes
+    CLANG_VERSION=18
+    GCC_VERSION=11
+    KATEX=yes
+    DOCS=yes
     ONNX=yes
     ;;
-  pytorch-linux-jammy-py3.10-clang15)
-    ANACONDA_PYTHON_VERSION=3.10
-    CLANG_VERSION=15
-    ;;
-  pytorch-linux-jammy-py3.11-clang15)
+  pytorch-linux-jammy-py3.11-clang18)
     ANACONDA_PYTHON_VERSION=3.11
-    CLANG_VERSION=15
+    CLANG_VERSION=18
     ;;
-  pytorch-linux-jammy-py3.12-clang15)
+  pytorch-linux-jammy-py3.12-clang18)
     ANACONDA_PYTHON_VERSION=3.12
-    CLANG_VERSION=15
+    CLANG_VERSION=18
     ;;
-  pytorch-linux-jammy-py3.13-clang15)
+  pytorch-linux-jammy-py3.13-clang18)
     ANACONDA_PYTHON_VERSION=3.13
-    CLANG_VERSION=15
+    CLANG_VERSION=18
     ;;
-  pytorch-linux-jammy-py3.14-clang15)
+  pytorch-linux-jammy-py3.14-clang18)
     ANACONDA_PYTHON_VERSION=3.14
-    CLANG_VERSION=15
+    CLANG_VERSION=18
     ;;
   pytorch-linux-jammy-rocm-n-py3 | pytorch-linux-jammy-rocm-n-py3-benchmarks | pytorch-linux-noble-rocm-n-py3)
     if [[ $tag =~ "jammy" ]]; then
@@ -193,14 +159,10 @@ case "$tag" in
     else
       ANACONDA_PYTHON_VERSION=3.12
     fi
-    GCC_VERSION=11
-    VISION=yes
-    ROCM_VERSION=7.1
-    NINJA_VERSION=1.9.0
+    GCC_VERSION=13
+    ROCM_VERSION=7.2
     TRITON=yes
     KATEX=yes
-    UCX_COMMIT=${_UCX_COMMIT}
-    UCC_COMMIT=${_UCC_COMMIT}
     PYTORCH_ROCM_ARCH="gfx90a;gfx942;gfx950;gfx1100"
     if [[ $tag =~ "benchmarks" ]]; then
       INDUCTOR_BENCHMARKS=yes
@@ -208,32 +170,28 @@ case "$tag" in
     ;;
   pytorch-linux-noble-rocm-nightly-py3)
     ANACONDA_PYTHON_VERSION=3.12
-    GCC_VERSION=11
-    VISION=yes
+    GCC_VERSION=13
     ROCM_VERSION=nightly
-    NINJA_VERSION=1.9.0
     TRITON=yes
     KATEX=yes
-    UCX_COMMIT=${_UCX_COMMIT}
-    UCC_COMMIT=${_UCC_COMMIT}
     PYTORCH_ROCM_ARCH="gfx942"
     ;;
   pytorch-linux-jammy-xpu-n-1-py3)
     ANACONDA_PYTHON_VERSION=3.10
     GCC_VERSION=11
-    VISION=yes
     XPU_VERSION=2025.2
     XPU_DRIVER_TYPE=LTS
-    NINJA_VERSION=1.9.0
     TRITON=yes
     ;;
-  pytorch-linux-noble-xpu-n-py3 | pytorch-linux-noble-xpu-n-py3-inductor-benchmarks)
+  pytorch-linux-noble-xpu-n-py3 | pytorch-linux-noble-xpu-n-py3-client | pytorch-linux-noble-xpu-n-py3-inductor-benchmarks)
     ANACONDA_PYTHON_VERSION=3.10
     GCC_VERSION=13
-    VISION=yes
     XPU_VERSION=2025.3
-    XPU_DRIVER_TYPE=LTS
-    NINJA_VERSION=1.9.0
+    if [[ $tag =~ "client" ]]; then
+      XPU_DRIVER_TYPE=CLIENT
+    else
+      XPU_DRIVER_TYPE=LTS
+    fi
     TRITON=yes
     if [[ $tag =~ "benchmarks" ]]; then
       INDUCTOR_BENCHMARKS=yes
@@ -242,36 +200,19 @@ case "$tag" in
   pytorch-linux-jammy-py3-gcc11-inductor-benchmarks)
     ANACONDA_PYTHON_VERSION=3.10
     GCC_VERSION=11
-    VISION=yes
     KATEX=yes
-    TRITON=yes
     DOCS=yes
     INDUCTOR_BENCHMARKS=yes
     ;;
-  pytorch-linux-jammy-cuda12.8-cudnn9-py3.10-clang15)
+  pytorch-linux-jammy-cuda12.8-cudnn9-py3.10-clang18)
     ANACONDA_PYTHON_VERSION=3.10
     CUDA_VERSION=12.8.1
-    CLANG_VERSION=15
-    VISION=yes
-    TRITON=yes
-    ;;
-  pytorch-linux-jammy-py3-clang18-asan)
-    ANACONDA_PYTHON_VERSION=3.10
     CLANG_VERSION=18
-    VISION=yes
-    ;;
-  pytorch-linux-jammy-py3.10-gcc11)
-    ANACONDA_PYTHON_VERSION=3.10
-    GCC_VERSION=11
-    VISION=yes
-    KATEX=yes
     TRITON=yes
-    DOCS=yes
-    UNINSTALL_DILL=yes
     ;;
-  pytorch-linux-jammy-py3-clang15-executorch)
+  pytorch-linux-jammy-py3-clang18-executorch)
     ANACONDA_PYTHON_VERSION=3.10
-    CLANG_VERSION=15
+    CLANG_VERSION=18
     EXECUTORCH=yes
     ;;
   pytorch-linux-jammy-py3.12-halide)
@@ -316,31 +257,13 @@ case "$tag" in
     ANACONDA_PYTHON_VERSION=3.10
     GCC_VERSION=13
     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-clang21)
-    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
     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
     INDUCTOR_BENCHMARKS=yes
     ;;
   pytorch-linux-noble-riscv64-py3.12-gcc14)
@@ -348,7 +271,6 @@ case "$tag" in
     ;;
   *)
     # Catch-all for builds that are not hardcoded.
-    VISION=yes
     echo "image '$image' did not match an existing build configuration"
     if [[ "$image" == *py* ]]; then
       extract_version_from_image_name py ANACONDA_PYTHON_VERSION
@@ -365,14 +287,10 @@ case "$tag" in
       if [[ -z "$ROCM_VERSION" ]]; then
         extract_version_from_image_name rocm ROCM_VERSION
       fi
-      NINJA_VERSION=1.9.0
       TRITON=yes
       # To ensure that any ROCm config will build using conda cmake
       # and thus have LAPACK/MKL enabled
       fi
-    if [[ "$image" == *centos7* ]]; then
-      NINJA_VERSION=1.10.2
-    fi
     if [[ "$image" == *gcc* ]]; then
       extract_version_from_image_name gcc GCC_VERSION
     fi
@@ -404,7 +322,6 @@ docker buildx build \
        ${progress_flag} \
        --build-arg "BUILD_ENVIRONMENT=${image}" \
        --build-arg "LLVMDEV=${LLVMDEV:-}" \
-       --build-arg "VISION=${VISION:-}" \
        --build-arg "UBUNTU_VERSION=${UBUNTU_VERSION}" \
        --build-arg "DEVTOOLSET_VERSION=${DEVTOOLSET_VERSION}" \
        --build-arg "GLIBC_VERSION=${GLIBC_VERSION}" \
@@ -414,13 +331,10 @@ docker buildx build \
        --build-arg "PYTHON_VERSION=${PYTHON_VERSION}" \
        --build-arg "GCC_VERSION=${GCC_VERSION}" \
        --build-arg "CUDA_VERSION=${CUDA_VERSION}" \
-       --build-arg "NINJA_VERSION=${NINJA_VERSION:-}" \
        --build-arg "KATEX=${KATEX:-}" \
        --build-arg "ROCM_VERSION=${ROCM_VERSION:-}" \
        --build-arg "PYTORCH_ROCM_ARCH=${PYTORCH_ROCM_ARCH}" \
        --build-arg "IMAGE_NAME=${IMAGE_NAME}" \
-       --build-arg "UCX_COMMIT=${UCX_COMMIT}" \
-       --build-arg "UCC_COMMIT=${UCC_COMMIT}" \
        --build-arg "TRITON=${TRITON}" \
        --build-arg "TRITON_CPU=${TRITON_CPU}" \
        --build-arg "ONNX=${ONNX}" \
@@ -432,11 +346,9 @@ docker buildx build \
        --build-arg "TPU=${TPU}" \
        --build-arg "XPU_VERSION=${XPU_VERSION}" \
        --build-arg "XPU_DRIVER_TYPE=${XPU_DRIVER_TYPE}" \
-       --build-arg "UNINSTALL_DILL=${UNINSTALL_DILL}" \
        --build-arg "ACL=${ACL:-}" \
        --build-arg "OPENBLAS=${OPENBLAS:-}" \
        --build-arg "SKIP_SCCACHE_INSTALL=${SKIP_SCCACHE_INSTALL:-}" \
-       --build-arg "SKIP_LLVM_SRC_BUILD_INSTALL=${SKIP_LLVM_SRC_BUILD_INSTALL:-}" \
        --build-arg "INSTALL_MINGW=${INSTALL_MINGW:-}" \
        -f $(dirname ${DOCKERFILE})/Dockerfile \
        --load \
diff --git a/.ci/docker/ci_commit_pins/huggingface-requirements.txt b/.ci/docker/ci_commit_pins/huggingface-requirements.txt
index 408343c9099c8..51a16f10e0632 100644
--- a/.ci/docker/ci_commit_pins/huggingface-requirements.txt
+++ b/.ci/docker/ci_commit_pins/huggingface-requirements.txt
@@ -1,2 +1,2 @@
-transformers==4.57.5
+transformers==5.5.3
 soxr==0.5.0
diff --git a/.ci/docker/ci_commit_pins/nccl-cu126.txt b/.ci/docker/ci_commit_pins/nccl-cu126.txt
new file mode 100644
index 0000000000000..1706c910183ce
--- /dev/null
+++ b/.ci/docker/ci_commit_pins/nccl-cu126.txt
@@ -0,0 +1 @@
+v2.29.3-1
diff --git a/.ci/docker/ci_commit_pins/nccl.txt b/.ci/docker/ci_commit_pins/nccl.txt
index 7c451d9fad29a..9ad2e5cfc6595 100644
--- a/.ci/docker/ci_commit_pins/nccl.txt
+++ b/.ci/docker/ci_commit_pins/nccl.txt
@@ -1 +1 @@
-v2.28.9-1
+v2.29.7-1
diff --git a/.ci/docker/ci_commit_pins/torchbench.txt b/.ci/docker/ci_commit_pins/torchbench.txt
index fdcdfa96ac643..36a3b62eaec96 100644
--- a/.ci/docker/ci_commit_pins/torchbench.txt
+++ b/.ci/docker/ci_commit_pins/torchbench.txt
@@ -1 +1 @@
-7be477a510cdbe9543b2f11357598db51c22e94e
+0fff8bf35400ec7b733d5080628fc79551205e36
diff --git a/.ci/docker/ci_commit_pins/triton-xpu.txt b/.ci/docker/ci_commit_pins/triton-xpu.txt
index 7453190b8a7f1..3bec45a60a706 100644
--- a/.ci/docker/ci_commit_pins/triton-xpu.txt
+++ b/.ci/docker/ci_commit_pins/triton-xpu.txt
@@ -1 +1 @@
-64bb0de394fe1b9e163758380ce8eed25511bb5e
+21033c4e2be9b42c9e6ce7a39a70ead2aba279b4
diff --git a/.ci/docker/ci_commit_pins/triton.txt b/.ci/docker/ci_commit_pins/triton.txt
index 23407b4d540c4..cdcf66491de62 100644
--- a/.ci/docker/ci_commit_pins/triton.txt
+++ b/.ci/docker/ci_commit_pins/triton.txt
@@ -1 +1 @@
-9844da955a9db14ec69c9aac828ee9803085e288
+4cff872ced001ea92d9fcf05b3f6517e2b486d19
diff --git a/.ci/docker/common/install_amdsmi.sh b/.ci/docker/common/install_amdsmi.sh
index 8e0ee620da679..759e2bababe25 100644
--- a/.ci/docker/common/install_amdsmi.sh
+++ b/.ci/docker/common/install_amdsmi.sh
@@ -7,7 +7,7 @@ source /etc/rocm_env.sh
 # For theRock nightly, amd_smi may already be installed or in a different location
 if [ -d "${ROCM_PATH}/share/amd_smi" ]; then
   echo "Installing amdsmi from: ${ROCM_PATH}/share/amd_smi"
-  cd ${ROCM_PATH}/share/amd_smi && pip install .
+  cd ${ROCM_PATH}/share/amd_smi && python3 -m pip install .
 else
   echo "AMD SMI not found at ${ROCM_PATH}/share/amd_smi - skipping (may already be installed via pip)"
 fi
diff --git a/.ci/docker/common/install_base.sh b/.ci/docker/common/install_base.sh
index 7d8ae247d7a0b..27d9ca0601779 100755
--- a/.ci/docker/common/install_base.sh
+++ b/.ci/docker/common/install_base.sh
@@ -11,36 +11,26 @@ install_ubuntu() {
   #   "$UBUNTU_VERSION" == "18.04"
   if [[ "$UBUNTU_VERSION" == "20.04"* ]]; then
     cmake3="cmake=3.16*"
-    maybe_libiomp_dev=""
   elif [[ "$UBUNTU_VERSION" == "22.04"* ]]; then
     cmake3="cmake=3.22*"
-    maybe_libiomp_dev=""
   elif [[ "$UBUNTU_VERSION" == "24.04"* ]]; then
     cmake3="cmake=3.28*"
-    maybe_libiomp_dev=""
   else
-    cmake3="cmake=3.5*"
-    maybe_libiomp_dev="libiomp-dev"
-  fi
-
-  if [[ "$CLANG_VERSION" == 15 ]]; then
-    maybe_libomp_dev="libomp-15-dev"
-  elif [[ "$CLANG_VERSION" == 12 ]]; then
-    maybe_libomp_dev="libomp-12-dev"
-  elif [[ "$CLANG_VERSION" == 10 ]]; then
-    maybe_libomp_dev="libomp-10-dev"
-  else
-    maybe_libomp_dev=""
+    echo "Unknown Ubuntu version $UBUNTU_VERSION"
+    exit 1
   fi
 
   # Install common dependencies
   apt-get update
+  # Install prerequisites for add-apt-repository (needs gpg-agent for PPA key import)
+  apt-get install -y --no-install-recommends software-properties-common gpg-agent
+  # Add git-core PPA for a newer version of git
+  add-apt-repository ppa:git-core/ppa -y
+  apt-get update
   # TODO: Some of these may not be necessary
-  ccache_deps="asciidoc docbook-xml docbook-xsl xsltproc"
   deploy_deps="libffi-dev libbz2-dev libreadline-dev libncurses5-dev libncursesw5-dev libgdbm-dev libsqlite3-dev uuid-dev tk-dev"
   numpy_deps="gfortran"
   apt-get install -y --no-install-recommends \
-    $ccache_deps \
     $numpy_deps \
     ${deploy_deps} \
     ${cmake3} \
@@ -53,14 +43,13 @@ install_ubuntu() {
     git \
     libatlas-base-dev \
     libc6-dbg \
-    ${maybe_libiomp_dev} \
     libyaml-dev \
     libz-dev \
     libjemalloc2 \
+    libgl1 \
     libjpeg-dev \
     libasound2-dev \
     libsndfile-dev \
-    ${maybe_libomp_dev} \
     software-properties-common \
     wget \
     sudo \
@@ -71,7 +60,9 @@ install_ubuntu() {
     unzip \
     gpg-agent \
     gdb \
-    bc
+    bc \
+    zip \
+    valgrind
 
   # Should resolve issues related to various apt package repository cert issues
   # see: https://github.com/pytorch/pytorch/issues/65931
@@ -82,70 +73,14 @@ install_ubuntu() {
   rm -rf /var/lib/apt/lists/* /tmp/* /var/tmp/*
 }
 
-install_centos() {
-  # Need EPEL for many packages we depend on.
-  # See http://fedoraproject.org/wiki/EPEL
-  yum --enablerepo=extras install -y epel-release
-
-  ccache_deps="asciidoc docbook-dtds docbook-style-xsl libxslt"
-  numpy_deps="gcc-gfortran"
-  yum install -y \
-    $ccache_deps \
-    $numpy_deps \
-    autoconf \
-    automake \
-    bzip2 \
-    cmake \
-    cmake3 \
-    curl \
-    gcc \
-    gcc-c++ \
-    gflags-devel \
-    git \
-    glibc-devel \
-    glibc-headers \
-    glog-devel \
-    libstdc++-devel \
-    libsndfile-devel \
-    make \
-    opencv-devel \
-    sudo \
-    wget \
-    vim \
-    unzip \
-    gdb
-
-  # Cleanup
-  yum clean all
-  rm -rf /var/cache/yum
-  rm -rf /var/lib/yum/yumdb
-  rm -rf /var/lib/yum/history
-}
-
 # Install base packages depending on the base OS
 ID=$(grep -oP '(?<=^ID=).+' /etc/os-release | tr -d '"')
 case "$ID" in
   ubuntu)
     install_ubuntu
     ;;
-  centos)
-    install_centos
-    ;;
   *)
     echo "Unable to determine OS..."
     exit 1
     ;;
 esac
-
-# Install Valgrind separately since the apt-get version is too old.
-mkdir valgrind_build && cd valgrind_build
-VALGRIND_VERSION=3.20.0
-wget https://ossci-linux.s3.amazonaws.com/valgrind-${VALGRIND_VERSION}.tar.bz2
-tar -xjf valgrind-${VALGRIND_VERSION}.tar.bz2
-cd valgrind-${VALGRIND_VERSION}
-./configure --prefix=/usr/local
-make -j$[$(nproc) - 2]
-sudo make install
-cd ../../
-rm -rf valgrind_build
-alias valgrind="/usr/local/bin/valgrind"
diff --git a/.ci/docker/common/install_cache.sh b/.ci/docker/common/install_cache.sh
index 3eab145967854..837ba50098029 100644
--- a/.ci/docker/common/install_cache.sh
+++ b/.ci/docker/common/install_cache.sh
@@ -18,7 +18,8 @@ install_ubuntu() {
   cp target/release/sccache-dist /opt/cache/bin
   echo "Cleaning up"
   cd ..
-  rm -rf sccache .cargo
+  rm -rf sccache
+  rustup self uninstall -y
   apt-get remove -y pkg-config libssl-dev
   apt-get autoclean && apt-get clean
 
@@ -77,10 +78,15 @@ EOF
   chmod a+x "/opt/cache/bin/$1"
 }
 
-write_sccache_stub cc
-write_sccache_stub c++
-write_sccache_stub gcc
-write_sccache_stub g++
+# Skip all sccache wrapping for theRock nightly: sccache PATH wrappers
+# intercept assembly (.s) compilation and fail because the assembler does not
+# produce the .d dependency file that sccache expects.
+if [ "$ROCM_VERSION" != "nightly" ]; then
+  write_sccache_stub cc
+  write_sccache_stub c++
+  write_sccache_stub gcc
+  write_sccache_stub g++
+fi
 
 # NOTE: See specific ROCM_VERSION case below.
 if [ "x$ROCM_VERSION" = x ]; then
diff --git a/.ci/docker/common/install_clang.sh b/.ci/docker/common/install_clang.sh
index 93daeee919b3d..20eea1ee157fe 100755
--- a/.ci/docker/common/install_clang.sh
+++ b/.ci/docker/common/install_clang.sh
@@ -15,7 +15,7 @@ if [ -n "$CLANG_VERSION" ]; then
 
   sudo apt-get update
   if [[ $CLANG_VERSION -ge 18 ]]; then
-    apt-get install -y libomp-${CLANG_VERSION}-dev libclang-rt-${CLANG_VERSION}-dev clang-"$CLANG_VERSION" llvm-"$CLANG_VERSION"
+    apt-get install -y --no-install-recommends libomp-${CLANG_VERSION}-dev libclang-rt-${CLANG_VERSION}-dev clang-"$CLANG_VERSION" llvm-"$CLANG_VERSION"
   else
     apt-get install -y --no-install-recommends clang-"$CLANG_VERSION" llvm-"$CLANG_VERSION"
   fi
diff --git a/.ci/docker/common/install_conda.sh b/.ci/docker/common/install_conda.sh
index 57c9845d76a0b..547eec0d401f5 100755
--- a/.ci/docker/common/install_conda.sh
+++ b/.ci/docker/common/install_conda.sh
@@ -67,13 +67,6 @@ if [ -n "$ANACONDA_PYTHON_VERSION" ]; then
     conda_install sqlite
   fi
 
-  # Install PyTorch conda deps, as per https://github.com/pytorch/pytorch README
-  if [[ $(uname -m) != "aarch64" ]]; then
-    pip_install mkl==2024.2.0
-    pip_install mkl-static==2024.2.0
-    pip_install mkl-include==2024.2.0
-  fi
-
   # Install llvm-8 as it is required to compile llvmlite-0.30.0 from source
   # and libpython-static for torch deploy
   conda_install llvmdev=8.0.0 "libpython-static=${ANACONDA_PYTHON_VERSION}"
@@ -103,5 +96,8 @@ if [ -n "$ANACONDA_PYTHON_VERSION" ]; then
     pip_install -r /opt/conda/requirements-docs.txt
   fi
 
+  # Clean conda package cache
+  as_jenkins conda clean -ya
+
   popd
 fi
diff --git a/.ci/docker/common/install_conda_docker.sh b/.ci/docker/common/install_conda_docker.sh
index dc377075750ac..9665a799b6ff5 100755
--- a/.ci/docker/common/install_conda_docker.sh
+++ b/.ci/docker/common/install_conda_docker.sh
@@ -13,8 +13,8 @@ rm $(basename "$MINICONDA_URL")
 export PATH=/opt/conda/bin:$PATH
 # See https://github.com/pytorch/builder/issues/1473
 # Pin conda to 23.5.2 as it's the last one compatible with openssl-1.1.1
-conda install -y conda=23.5.2 conda-build anaconda-client git ninja
+conda install -y conda=23.5.2 conda-build anaconda-client git
 # The cmake version here needs to match with the minimum version of cmake
-# supported by PyTorch (3.18). There is only 3.18.2 on anaconda
-/opt/conda/bin/pip3 install cmake==3.18.2
+# supported by PyTorch (3.18).
+/opt/conda/bin/pip3 install cmake==3.18.4.post1 ninja
 conda remove -y --force patchelf
diff --git a/.ci/docker/common/install_cuda.sh b/.ci/docker/common/install_cuda.sh
index 2d1db795d9cb4..9d7cd7ad78c05 100644
--- a/.ci/docker/common/install_cuda.sh
+++ b/.ci/docker/common/install_cuda.sh
@@ -82,21 +82,23 @@ function install_nvshmem {
 
 function install_124 {
   CUDNN_VERSION=9.1.0.70
-  echo "Installing CUDA 12.4.1 and cuDNN ${CUDNN_VERSION} and NCCL and cuSparseLt-0.6.2"
+  CUSPARSELT_VERSION=0.6.2.3
+  echo "Installing CUDA 12.4.1 and cuDNN ${CUDNN_VERSION} and NCCL and cuSparseLt-${CUSPARSELT_VERSION}"
   install_cuda 12.4.1 cuda_12.4.1_550.54.15_linux
 
   install_cudnn 12 $CUDNN_VERSION
 
   CUDA_VERSION=12.4 bash install_nccl.sh
 
-  CUDA_VERSION=12.4 bash install_cusparselt.sh
+  CUDA_VERSION=12.4 bash install_cusparselt.sh $CUSPARSELT_VERSION
 
   ldconfig
 }
 
 function install_126 {
   CUDNN_VERSION=9.10.2.21
-  echo "Installing CUDA 12.6.3 and cuDNN ${CUDNN_VERSION} and NVSHMEM and NCCL and cuSparseLt-0.7.1"
+  CUSPARSELT_VERSION=0.7.1.0
+  echo "Installing CUDA 12.6.3 and cuDNN ${CUDNN_VERSION} and NVSHMEM and NCCL and cuSparseLt-${CUSPARSELT_VERSION}"
   install_cuda 12.6.3 cuda_12.6.3_560.35.05_linux
 
   install_cudnn 12 $CUDNN_VERSION
@@ -105,14 +107,15 @@ function install_126 {
 
   CUDA_VERSION=12.6 bash install_nccl.sh
 
-  CUDA_VERSION=12.6 bash install_cusparselt.sh
+  CUDA_VERSION=12.6 bash install_cusparselt.sh $CUSPARSELT_VERSION
 
   ldconfig
 }
 
 function install_129 {
-  CUDNN_VERSION=9.17.1.4
-  echo "Installing CUDA 12.9.1 and cuDNN ${CUDNN_VERSION} and NVSHMEM and NCCL and cuSparseLt-0.7.1"
+  CUDNN_VERSION=9.20.0.48
+  CUSPARSELT_VERSION=0.8.1.1
+  echo "Installing CUDA 12.9.1 and cuDNN ${CUDNN_VERSION} and NVSHMEM and NCCL and cuSparseLt-${CUSPARSELT_VERSION}"
   # install CUDA 12.9.1 in the same container
   install_cuda 12.9.1 cuda_12.9.1_575.57.08_linux
 
@@ -123,14 +126,15 @@ function install_129 {
 
   CUDA_VERSION=12.9 bash install_nccl.sh
 
-  CUDA_VERSION=12.9 bash install_cusparselt.sh
+  CUDA_VERSION=12.9 bash install_cusparselt.sh $CUSPARSELT_VERSION
 
   ldconfig
 }
 
 function install_128 {
-  CUDNN_VERSION=9.17.1.4
-  echo "Installing CUDA 12.8.1 and cuDNN ${CUDNN_VERSION} and NVSHMEM and NCCL and cuSparseLt-0.7.1"
+  CUDNN_VERSION=9.20.0.48
+  CUSPARSELT_VERSION=0.7.1.0
+  echo "Installing CUDA 12.8.1 and cuDNN ${CUDNN_VERSION} and NVSHMEM and NCCL and cuSparseLt-${CUSPARSELT_VERSION}"
   # install CUDA 12.8.1 in the same container
   install_cuda 12.8.1 cuda_12.8.1_570.124.06_linux
 
@@ -141,14 +145,15 @@ function install_128 {
 
   CUDA_VERSION=12.8 bash install_nccl.sh
 
-  CUDA_VERSION=12.8 bash install_cusparselt.sh
+  CUDA_VERSION=12.8 bash install_cusparselt.sh $CUSPARSELT_VERSION
 
   ldconfig
 }
 
 function install_130 {
-  CUDNN_VERSION=9.17.1.4
-  echo "Installing CUDA 13.0 and cuDNN ${CUDNN_VERSION} and NVSHMEM and NCCL and cuSparseLt-0.7.1"
+  CUDNN_VERSION=9.20.0.48
+  CUSPARSELT_VERSION=0.8.1.1
+  echo "Installing CUDA 13.0 and cuDNN ${CUDNN_VERSION} and NVSHMEM and NCCL and cuSparseLt-${CUSPARSELT_VERSION}"
   # install CUDA 13.0 in the same container
   install_cuda 13.0.2 cuda_13.0.2_580.95.05_linux
 
@@ -159,7 +164,26 @@ function install_130 {
 
   CUDA_VERSION=13.0 bash install_nccl.sh
 
-  CUDA_VERSION=13.0 bash install_cusparselt.sh
+  CUDA_VERSION=13.0 bash install_cusparselt.sh $CUSPARSELT_VERSION
+
+  ldconfig
+}
+
+function install_132 {
+  CUDNN_VERSION=9.20.0.48
+  CUSPARSELT_VERSION=0.8.1.1
+  echo "Installing CUDA 13.2 and cuDNN ${CUDNN_VERSION} and NVSHMEM and NCCL and cuSparseLt-${CUSPARSELT_VERSION}"
+  # install CUDA 13.2 in the same container
+  install_cuda 13.2.1 cuda_13.2.1_595.58.03_linux
+
+  # cuDNN license: https://developer.nvidia.com/cudnn/license_agreement
+  install_cudnn 13 $CUDNN_VERSION
+
+  install_nvshmem 13 $NVSHMEM_VERSION
+
+  CUDA_VERSION=13.2 bash install_nccl.sh
+
+  CUDA_VERSION=13.2 bash install_cusparselt.sh $CUSPARSELT_VERSION
 
   ldconfig
 }
@@ -178,6 +202,8 @@ do
         ;;
     13.0|13.0.*) install_130;
         ;;
+    13.2|13.2.*) install_132;
+        ;;
     *) echo "bad argument $1"; exit 1
         ;;
     esac
diff --git a/.ci/docker/common/install_cusparselt.sh b/.ci/docker/common/install_cusparselt.sh
index b532c086371f1..0568dd1a18f55 100644
--- a/.ci/docker/common/install_cusparselt.sh
+++ b/.ci/docker/common/install_cusparselt.sh
@@ -5,34 +5,29 @@ set -ex
 # cuSPARSELt license: https://docs.nvidia.com/cuda/cusparselt/license.html
 mkdir tmp_cusparselt && cd tmp_cusparselt
 
-if [[ ${CUDA_VERSION:0:4} =~ "13" ]]; then
-    arch_path='sbsa'
-    export TARGETARCH=${TARGETARCH:-$(uname -m)}
-    if [ ${TARGETARCH} = 'amd64' ] || [ "${TARGETARCH}" = 'x86_64' ]; then
-        arch_path='x86_64'
-    fi
-    CUSPARSELT_NAME="libcusparse_lt-linux-${arch_path}-0.8.0.4_cuda13-archive"
-    curl --retry 3 -OLs https://developer.download.nvidia.com/compute/cusparselt/redist/libcusparse_lt/linux-${arch_path}/${CUSPARSELT_NAME}.tar.xz
-elif [[ ${CUDA_VERSION:0:4} =~ ^12\.[5-9]$ ]]; then
-    arch_path='sbsa'
-    export TARGETARCH=${TARGETARCH:-$(uname -m)}
-    if [ ${TARGETARCH} = 'amd64' ] || [ "${TARGETARCH}" = 'x86_64' ]; then
-        arch_path='x86_64'
-    fi
-    CUSPARSELT_NAME="libcusparse_lt-linux-${arch_path}-0.7.1.0-archive"
-    curl --retry 3 -OLs https://developer.download.nvidia.com/compute/cusparselt/redist/libcusparse_lt/linux-${arch_path}/${CUSPARSELT_NAME}.tar.xz
-elif [[ ${CUDA_VERSION:0:4} == "12.4" ]]; then
-    arch_path='sbsa'
-    export TARGETARCH=${TARGETARCH:-$(uname -m)}
-    if [ ${TARGETARCH} = 'amd64' ] || [ "${TARGETARCH}" = 'x86_64' ]; then
-        arch_path='x86_64'
-    fi
-    CUSPARSELT_NAME="libcusparse_lt-linux-${arch_path}-0.6.2.3-archive"
-    curl --retry 3 -OLs https://developer.download.nvidia.com/compute/cusparselt/redist/libcusparse_lt/linux-${arch_path}/${CUSPARSELT_NAME}.tar.xz
+cusparselt_version=$1
+
+arch_path='sbsa'
+export TARGETARCH=${TARGETARCH:-$(uname -m)}
+if [ ${TARGETARCH} = 'amd64' ] || [ "${TARGETARCH}" = 'x86_64' ]; then
+    arch_path='x86_64'
+fi
+
+if [[ -z "${cusparselt_version}" ]]; then
+    echo "Usage: install_cusparselt.sh <cusparselt_version>"
+    exit 1
+fi
+
+cuda_major_version=${CUDA_VERSION%%.*}
+cusparselt_minor=$(echo "${cusparselt_version}" | cut -d. -f2)
+# Starting from 0.8.0, NVIDIA ships separate archives per CUDA major version
+if [[ "${cusparselt_minor}" -ge 8 ]]; then
+    CUSPARSELT_NAME="libcusparse_lt-linux-${arch_path}-${cusparselt_version}_cuda${cuda_major_version}-archive"
 else
-    echo "Not sure which libcusparselt version to install for this ${CUDA_VERSION}"
+    CUSPARSELT_NAME="libcusparse_lt-linux-${arch_path}-${cusparselt_version}-archive"
 fi
 
+curl --retry 3 -OLs https://developer.download.nvidia.com/compute/cusparselt/redist/libcusparse_lt/linux-${arch_path}/${CUSPARSELT_NAME}.tar.xz
 tar xf ${CUSPARSELT_NAME}.tar.xz
 cp -a ${CUSPARSELT_NAME}/include/* /usr/local/cuda/include/
 cp -a ${CUSPARSELT_NAME}/lib/* /usr/local/cuda/lib64/
diff --git a/.ci/docker/common/install_devtoolset.sh b/.ci/docker/common/install_devtoolset.sh
deleted file mode 100755
index bdae637598138..0000000000000
--- a/.ci/docker/common/install_devtoolset.sh
+++ /dev/null
@@ -1,10 +0,0 @@
-#!/bin/bash
-
-set -ex
-
-[ -n "$DEVTOOLSET_VERSION" ]
-
-yum install -y centos-release-scl
-yum install -y devtoolset-$DEVTOOLSET_VERSION
-
-echo "source scl_source enable devtoolset-$DEVTOOLSET_VERSION" > "/etc/profile.d/devtoolset-$DEVTOOLSET_VERSION.sh"
diff --git a/.ci/docker/common/install_docs_reqs.sh b/.ci/docker/common/install_docs_reqs.sh
index c907145f2ec62..c06160373a05e 100644
--- a/.ci/docker/common/install_docs_reqs.sh
+++ b/.ci/docker/common/install_docs_reqs.sh
@@ -17,7 +17,7 @@ if [ -n "$KATEX" ]; then
   apt-get install -y --no-install-recommends yarn
   yarn global add katex --prefix /usr/local
 
-  sudo apt-get -y install doxygen
+  sudo apt-get -y install doxygen lcov
 
   apt-get autoclean && apt-get clean
   rm -rf /var/lib/apt/lists/* /tmp/* /var/tmp/*
diff --git a/.ci/docker/common/install_glibc.sh b/.ci/docker/common/install_glibc.sh
deleted file mode 100755
index c98791e2bf85b..0000000000000
--- a/.ci/docker/common/install_glibc.sh
+++ /dev/null
@@ -1,34 +0,0 @@
-#!/bin/bash
-
-set -ex
-
-[ -n "$GLIBC_VERSION" ]
-if [[ -n "$CENTOS_VERSION" ]]; then
-  [ -n "$DEVTOOLSET_VERSION" ]
-fi
-
-yum install -y wget sed
-
-mkdir -p /packages && cd /packages
-wget -q http://ftp.gnu.org/gnu/glibc/glibc-$GLIBC_VERSION.tar.gz
-tar xzf glibc-$GLIBC_VERSION.tar.gz
-if [[ "$GLIBC_VERSION" == "2.26" ]]; then
-  cd glibc-$GLIBC_VERSION
-  sed -i 's/$name ne "nss_test1"/$name ne "nss_test1" \&\& $name ne "nss_test2"/' scripts/test-installation.pl
-  cd ..
-fi
-mkdir -p glibc-$GLIBC_VERSION-build && cd glibc-$GLIBC_VERSION-build
-
-if [[ -n "$CENTOS_VERSION" ]]; then
-  export PATH=/opt/rh/devtoolset-$DEVTOOLSET_VERSION/root/usr/bin:$PATH
-fi
-
-../glibc-$GLIBC_VERSION/configure --prefix=/usr CFLAGS='-Wno-stringop-truncation -Wno-format-overflow -Wno-restrict -Wno-format-truncation -g -O2'
-make -j$(nproc)
-make install
-
-# Cleanup
-rm -rf /packages
-rm -rf /var/cache/yum/*
-rm -rf /var/lib/rpm/__db.*
-yum clean all
diff --git a/.ci/docker/common/install_inductor_benchmark_deps.sh b/.ci/docker/common/install_inductor_benchmark_deps.sh
index 674b141efcfb2..c54b8a44f0632 100644
--- a/.ci/docker/common/install_inductor_benchmark_deps.sh
+++ b/.ci/docker/common/install_inductor_benchmark_deps.sh
@@ -18,18 +18,16 @@ function install_timm() {
 function install_torchbench() {
   local commit
   commit=$(get_pinned_commit torchbench)
-  git clone https://github.com/pytorch/benchmark torchbench
+  mkdir torchbench && chown jenkins torchbench
+  as_jenkins git clone https://github.com/pytorch/benchmark torchbench
   pushd torchbench
-  git checkout "$commit"
+  as_jenkins git checkout "$commit"
 
-  python install.py --continue_on_fail
+  conda_run python install.py --continue_on_fail
 
   echo "Print all dependencies after TorchBench is installed"
-  python -mpip freeze
+  conda_run python -mpip freeze
   popd
-
-  chown -R jenkins torchbench
-  chown -R jenkins /opt/conda
 }
 
 # Pango is needed for weasyprint which is needed for doctr
diff --git a/.ci/docker/common/install_mingw.sh b/.ci/docker/common/install_mingw.sh
index 6232a0d0245c7..e82a666ff4352 100644
--- a/.ci/docker/common/install_mingw.sh
+++ b/.ci/docker/common/install_mingw.sh
@@ -4,7 +4,7 @@ set -ex
 
 # Install MinGW-w64 for Windows cross-compilation
 apt-get update
-apt-get install -y g++-mingw-w64-x86-64-posix
+apt-get install -y g++-mingw-w64-x86-64-posix mingw-w64-tools
 
 echo "MinGW-w64 installed successfully"
 x86_64-w64-mingw32-g++ --version
diff --git a/.ci/docker/common/install_miopen.sh b/.ci/docker/common/install_miopen.sh
index 3dbc67b90abaf..039458add8406 100644
--- a/.ci/docker/common/install_miopen.sh
+++ b/.ci/docker/common/install_miopen.sh
@@ -16,7 +16,7 @@ case "$ID" in
   ubuntu)
     IS_UBUNTU=1
     ;;
-  centos|almalinux)
+  almalinux)
     IS_UBUNTU=0
     ;;
   *)
diff --git a/.ci/docker/common/install_nccl.sh b/.ci/docker/common/install_nccl.sh
index 486604140a983..f505e4c3f249a 100644
--- a/.ci/docker/common/install_nccl.sh
+++ b/.ci/docker/common/install_nccl.sh
@@ -18,6 +18,11 @@ NCCL_VERSION=$(cat ci_commit_pins/nccl.txt)
 #   exit 1
 # fi
 
+# Use the NCCL version for CUDA 12.6 due to sm50 support
+if [[ ${CUDA_VERSION:0:4} == "12.6" ]]; then
+  NCCL_VERSION=$(cat ci_commit_pins/nccl-cu126.txt)
+fi
+
 if [[ -n "${NCCL_VERSION}" ]]; then
   # NCCL license: https://docs.nvidia.com/deeplearning/nccl/#licenses
   # Follow build: https://github.com/NVIDIA/nccl/tree/master?tab=readme-ov-file#build
diff --git a/.ci/docker/common/install_ninja.sh b/.ci/docker/common/install_ninja.sh
deleted file mode 100644
index fa380722bdc2f..0000000000000
--- a/.ci/docker/common/install_ninja.sh
+++ /dev/null
@@ -1,18 +0,0 @@
-#!/bin/bash
-
-set -ex
-
-[ -n "$NINJA_VERSION" ]
-
-arch=$(uname -m)
-if [ "$arch" == "aarch64" ]; then
-    url="https://github.com/ninja-build/ninja/releases/download/v${NINJA_VERSION}/ninja-linux-aarch64.zip"
-else
-    url="https://github.com/ninja-build/ninja/releases/download/v${NINJA_VERSION}/ninja-linux.zip"
-fi
-
-pushd /tmp
-wget --no-verbose --output-document=ninja-linux.zip "$url"
-unzip ninja-linux.zip -d /usr/local/bin
-rm -f ninja-linux.zip
-popd
\ No newline at end of file
diff --git a/.ci/docker/common/install_onnx.sh b/.ci/docker/common/install_onnx.sh
index 36ce5b11d9135..eacc7bbcad157 100755
--- a/.ci/docker/common/install_onnx.sh
+++ b/.ci/docker/common/install_onnx.sh
@@ -11,15 +11,11 @@ retry () {
 # ONNXRuntime should be installed before installing
 # onnx-weekly. Otherwise, onnx-weekly could be
 # overwritten by onnx.
+# Note: parameterized, pytest-subtests, tabulate, packaging are already
+# installed via requirements-ci.txt
 pip_install \
-  parameterized==0.8.1 \
-  pytest-cov==4.0.0 \
-  pytest-subtests==0.10.0 \
-  tabulate==0.9.0 \
-  transformers==4.36.2
-
-pip_install coloredlogs packaging
-pip_install onnxruntime==1.23.1
+  transformers==4.36.2 \
+  onnxruntime==1.23.1
 
 # Cache the transformers model to be used later by ONNX tests. We need to run the transformers
 # package to download the model. By default, the model is cached at ~/.cache/huggingface/hub/
@@ -34,4 +30,5 @@ conda_run python "${IMPORT_SCRIPT_FILENAME}"
 
 # Cleaning up
 conda_run pip uninstall -y torch
+conda_run pip cache purge
 rm "${IMPORT_SCRIPT_FILENAME}" || true
diff --git a/.ci/docker/common/install_openssl.sh b/.ci/docker/common/install_openssl.sh
deleted file mode 100644
index c73c9c333c002..0000000000000
--- a/.ci/docker/common/install_openssl.sh
+++ /dev/null
@@ -1,17 +0,0 @@
-#!/bin/bash
-
-set -ex
-
-OPENSSL=openssl-1.1.1k
-
-wget -q -O "${OPENSSL}.tar.gz" "https://ossci-linux.s3.amazonaws.com/${OPENSSL}.tar.gz"
-tar xf "${OPENSSL}.tar.gz"
-cd "${OPENSSL}"
-./config --prefix=/opt/openssl -d '-Wl,--enable-new-dtags,-rpath,$(LIBRPATH)'
-# NOTE: openssl install errors out when built with the -j option
-NPROC=$[$(nproc) - 2]
-make -j${NPROC}; make install_sw
-# Link the ssl libraries to the /usr/lib folder.
-sudo ln -s /opt/openssl/lib/lib* /usr/lib
-cd ..
-rm -rf "${OPENSSL}"
diff --git a/.ci/docker/common/install_rocSHMEM.sh b/.ci/docker/common/install_rocSHMEM.sh
new file mode 100644
index 0000000000000..ff59951f94875
--- /dev/null
+++ b/.ci/docker/common/install_rocSHMEM.sh
@@ -0,0 +1,38 @@
+#!/usr/bin/env bash
+# Script used only in CD pipeline to build and install rocSHMEM
+
+set -eou pipefail
+
+function do_install() {
+    ROCSHMEM_VERSION=ea5c137103f18a9aadd570d09d72e78ec52f0a3a
+    rocm_dir="${ROCM_HOME:-}"
+    if [[ -z "${rocm_dir}" && -f /etc/rocm_env.sh ]]; then
+        source /etc/rocm_env.sh
+        rocm_dir="${ROCM_HOME:-}"
+    fi
+    rocm_dir="${rocm_dir:-/opt/rocm}"
+    echo "install_rocSHMEM.sh: using ROCM install prefix ${rocm_dir}"
+    if [[ -f "${rocm_dir}/lib/librocshmem.a" ]]; then
+        echo "install_rocSHMEM.sh: librocshmem.a already present in ${rocm_dir}/lib, skipping build"
+        return
+    fi
+    (
+        set -x
+        curr_dir=$(pwd)
+        tmp_dir=$(mktemp -d)
+
+        git clone --no-checkout --filter=blob:none https://github.com/ROCm/rocm-systems.git ${tmp_dir}/rocm-systems
+        cd ${tmp_dir}/rocm-systems
+        git sparse-checkout set --cone projects/rocshmem
+        git checkout ${ROCSHMEM_VERSION}
+
+        cd ${tmp_dir}/rocm-systems/projects/rocshmem
+        mkdir build
+        cd build
+        INSTALL_PREFIX="${rocm_dir}" ../scripts/build_configs/all_backends
+        cd ${curr_dir}
+
+    )
+}
+
+do_install
diff --git a/.ci/docker/common/install_rocm.sh b/.ci/docker/common/install_rocm.sh
index 21e5968016bd6..b4f6d70af3676 100644
--- a/.ci/docker/common/install_rocm.sh
+++ b/.ci/docker/common/install_rocm.sh
@@ -26,87 +26,131 @@ install_ubuntu() {
     apt-get install -y libc++1
     apt-get install -y libc++abi1
 
-    # When ROCM_VERSION=nightly, install ROCm from TheRock nightly wheels
+    # When ROCM_VERSION=nightly, install ROCm from TheRock nightly tarballs
+    # Mirrors: https://github.com/ROCm/TheRock/blob/main/dockerfiles/install_rocm_tarball.sh
     if [[ "${ROCM_VERSION}" == "nightly" ]]; then
-      echo "install_rocm.sh: installing ROCm from TheRock nightly wheels"
+      apt-get install -y --no-install-recommends pkg-config
 
-      # Clean any previous ROCm installation in the base CI image.
       if [[ -d /opt/rocm ]]; then
-        echo "Removing existing /opt/rocm from base image"
         rm -rf /opt/rocm
       fi
 
-      # Determine theRock nightly URL based on GPU architecture
-      # Check BUILD_ENVIRONMENT or PYTORCH_ROCM_ARCH for the target GPU
-      if [[ -z "${THEROCK_NIGHTLY_INDEX_URL:-}" ]]; then
+      # Determine GPU family based on target architecture
+      AMDGPU_FAMILY="${THEROCK_AMDGPU_FAMILY:-}"
+      if [[ -z "${AMDGPU_FAMILY}" ]]; then
         if [[ "${BUILD_ENVIRONMENT}" == *"gfx950"* ]] || [[ "${PYTORCH_ROCM_ARCH}" == *"gfx950"* ]]; then
-          # MI350 (gfx950)
-          THEROCK_NIGHTLY_INDEX_URL="https://rocm.nightlies.amd.com/v2/gfx950-dcgpu/"
-          echo "Detected gfx950 architecture - using MI350 theRock nightly repository"
+          AMDGPU_FAMILY="gfx950-dcgpu"
         else
-          # Default to MI300 (gfx942/gfx94X)
-          THEROCK_NIGHTLY_INDEX_URL="https://rocm.nightlies.amd.com/v2/gfx94X-dcgpu/"
-          echo "Using gfx94X (MI300) theRock nightly repository"
+          AMDGPU_FAMILY="gfx94X-dcgpu"
         fi
       fi
 
-      export THEROCK_NIGHTLY_INDEX_URL
-      echo "TheRock Index URL: ${THEROCK_NIGHTLY_INDEX_URL}"
+      # Auto-detect latest nightly version if not pinned
+      VERSION="${THEROCK_VERSION:-}"
+      if [[ -z "${VERSION}" ]]; then
+        VERSION=$(curl -fsSL "https://rocm.nightlies.amd.com/tarball/" \
+          | grep -oP "therock-dist-linux-${AMDGPU_FAMILY}-\K[^\"]+(?=\.tar\.gz)" \
+          | grep -v ADHOCBUILD \
+          | sort -V \
+          | tail -1)
+        if [[ -z "${VERSION}" ]]; then
+          echo "Error: Could not find a nightly tarball for ${AMDGPU_FAMILY}"
+          exit 1
+        fi
+      fi
+
+      # URL-encode '+' as '%2B' in VERSION (required for devreleases)
+      VERSION_ENCODED="${VERSION//+/%2B}"
 
-      python3 -m pip install \
-        --index-url "${THEROCK_NIGHTLY_INDEX_URL}" \
-        "rocm[libraries,devel]"
+      TARBALL_URL="https://rocm.nightlies.amd.com/tarball/therock-dist-linux-${AMDGPU_FAMILY}-${VERSION_ENCODED}.tar.gz"
 
-      # Use the rocm-sdk CLI helper to populate environment defaults
-      ROCM_HOME="$(rocm-sdk path --root)"
-      ROCM_BIN="$(rocm-sdk path --bin)"
-      ROCM_CMAKE_PREFIX="$(rocm-sdk path --cmake)"
+      echo "=============================================="
+      echo "ROCm Tarball Installation"
+      echo "=============================================="
+      echo "Version:         ${VERSION}"
+      echo "AMDGPU Family:   ${AMDGPU_FAMILY}"
+      echo "Tarball URL:     ${TARBALL_URL}"
+      echo "=============================================="
 
-      echo "ROCM_HOME=${ROCM_HOME}"
-      echo "ROCM_BIN=${ROCM_BIN}"
-      echo "ROCM_CMAKE_PREFIX=${ROCM_CMAKE_PREFIX}"
+      # Download tarball
+      TARBALL_FILE="/tmp/rocm-tarball.tar.gz"
 
-      export ROCM_HOME
-      export ROCM_PATH="${ROCM_HOME}"
-      export PATH="${ROCM_BIN}:${PATH}"
-      export CMAKE_PREFIX_PATH="${ROCM_CMAKE_PREFIX}:${CMAKE_PREFIX_PATH:-}"
+      echo "Downloading tarball..."
+      curl -fsSL -o "$TARBALL_FILE" "$TARBALL_URL" || {
+        echo "Error: Failed to download tarball from $TARBALL_URL"
+        exit 1
+      }
+
+      # Verify download
+      if [ ! -f "$TARBALL_FILE" ] || [ ! -s "$TARBALL_FILE" ]; then
+        echo "Error: Downloaded file is empty or does not exist"
+        exit 1
+      fi
 
-      # theRock bundles system dependencies like libdrm, liblzma in rocm_sysdeps
-      ROCM_SYSDEPS="${ROCM_HOME}/lib/rocm_sysdeps"
-      ROCM_SYSDEPS_INCLUDE="${ROCM_SYSDEPS}/include"
-      ROCM_SYSDEPS_PKGCONFIG="${ROCM_SYSDEPS}/lib/pkgconfig"
+      # Install directory is fixed to /opt/rocm-{VERSION}
+      ROCM_INSTALL_DIR="/opt/rocm-${VERSION}"
 
-      # Write environment to file that can be sourced by CI scripts and users
+      # Extract tarball to versioned directory
+      echo "Extracting tarball to ${ROCM_INSTALL_DIR}..."
+      mkdir -p "$ROCM_INSTALL_DIR"
+      tar -xzf "$TARBALL_FILE" -C "$ROCM_INSTALL_DIR"
+
+      # Clean up downloaded file
+      rm -f "$TARBALL_FILE"
+      echo "Tarball extracted and cleaned up"
+
+      # Create symlink /opt/rocm -> /opt/rocm-{VERSION} for compatibility
+      ln -sfn "$ROCM_INSTALL_DIR" /opt/rocm
+      echo "Created symlink: /opt/rocm -> $ROCM_INSTALL_DIR"
+
+      # Verify bin and lib folder exists after extraction
+      echo "Verifying installation..."
+      for dir in bin clients include lib libexec share; do
+        if [ ! -d "$ROCM_INSTALL_DIR/$dir" ]; then
+          echo "Error: ROCm $dir directory not found"
+          exit 1
+        fi
+        echo "ROCm $dir found in $ROCM_INSTALL_DIR/$dir"
+      done
+
+      echo "=============================================="
+      echo "ROCm installed successfully to $ROCM_INSTALL_DIR"
+      echo "ROCM_PATH=$ROCM_INSTALL_DIR"
+      echo "PATH should include: $ROCM_INSTALL_DIR/bin"
+      echo "=============================================="
+
+      # Write environment file (sourced by CI scripts and interactive shells)
       cat > /etc/rocm_env.sh << ROCM_ENV
 # ROCm paths
-export ROCM_PATH="${ROCM_HOME}"
-export ROCM_HOME="${ROCM_HOME}"
-export ROCM_SOURCE_DIR="${ROCM_HOME}"
-export ROCM_BIN="${ROCM_BIN}"
-export ROCM_CMAKE="${ROCM_CMAKE_PREFIX}"
-export PATH="${ROCM_BIN}:\${PATH}"
-export CMAKE_PREFIX_PATH="${ROCM_CMAKE_PREFIX}:\${CMAKE_PREFIX_PATH:-}"
-# Device library paths
-export HIP_DEVICE_LIB_PATH="${ROCM_HOME}/lib/llvm/amdgcn/bitcode"
-export ROCM_DEVICE_LIB_PATH="${ROCM_HOME}/lib/llvm/amdgcn/bitcode"
-# theRock system dependencies
-export ROCM_SYSDEPS_INCLUDE="${ROCM_SYSDEPS_INCLUDE}"
-export CPLUS_INCLUDE_PATH="${ROCM_SYSDEPS_INCLUDE}:\${CPLUS_INCLUDE_PATH:-}"
-export C_INCLUDE_PATH="${ROCM_SYSDEPS_INCLUDE}:\${C_INCLUDE_PATH:-}"
-export PKG_CONFIG_PATH="${ROCM_SYSDEPS_PKGCONFIG}:\${PKG_CONFIG_PATH:-}"
-export LD_LIBRARY_PATH="${ROCM_SYSDEPS}/lib:\${LD_LIBRARY_PATH:-}"
-export LIBRARY_PATH="${ROCM_SYSDEPS}/lib:\${LIBRARY_PATH:-}"
-export MAGMA_HOME="${ROCM_HOME}/magma"
+export ROCM_PATH=/opt/rocm
+export ROCM_HOME=/opt/rocm
+export ROCM_SOURCE_DIR=/opt/rocm
+export ROCM_BIN=/opt/rocm/bin
+export ROCM_CMAKE=/opt/rocm
+export PATH=/opt/rocm/bin:/opt/rocm/llvm/bin:\${PATH}
+export LD_LIBRARY_PATH=/opt/rocm/lib:\${LD_LIBRARY_PATH:-}
+# Sysdeps include paths (libdrm headers, etc.)
+export CPLUS_INCLUDE_PATH=/opt/rocm/lib/rocm_sysdeps/include:\${CPLUS_INCLUDE_PATH:-}
+export C_INCLUDE_PATH=/opt/rocm/lib/rocm_sysdeps/include:\${C_INCLUDE_PATH:-}
+# Device library path
+export HIP_DEVICE_LIB_PATH=/opt/rocm/amdgcn/bitcode
+export MAGMA_HOME=/opt/rocm/magma
+# Tarball bundles sysdeps (libdrm, liblzma, etc.); expose their libs and .pc files
+if [ -d /opt/rocm/lib/rocm_sysdeps/lib ]; then
+  export LD_LIBRARY_PATH=/opt/rocm/lib/rocm_sysdeps/lib:\${LD_LIBRARY_PATH}
+  export PKG_CONFIG_PATH=/opt/rocm/lib/rocm_sysdeps/lib/pkgconfig:\${PKG_CONFIG_PATH:-}
+fi
 # Disable MSLK for theRock nightly (not yet supported)
 export USE_MSLK=0
 ROCM_ENV
 
-      # Append to bash.bashrc so interactive shells get the env vars
       echo "source /etc/rocm_env.sh" >> /etc/bash.bashrc
 
-      echo "install_rocm.sh: TheRock nightly ROCm install complete"
-      exit 0
-    fi
+      # --- End of theRock nightly tarball installation ---
+    else
+      # =========================================================================
+      # Non-nightly: install ROCm from repo.radeon.com apt packages
+      # =========================================================================
 
     # Make sure rocm packages from repo.radeon.com have highest priority
     cat << EOF > /etc/apt/preferences.d/rocm-pin-600
@@ -120,6 +164,11 @@ EOF
         ROCM_VERSION="${ROCM_VERSION}.2"
     fi
 
+    # we want the patch version of 7.2 instead
+    if [[ $(ver $ROCM_VERSION) -eq $(ver 7.2) ]]; then
+        ROCM_VERSION="${ROCM_VERSION}.1"
+    fi
+
     # Default url values
     rocm_baseurl="http://repo.radeon.com/rocm/apt/${ROCM_VERSION}"
     UBUNTU_VERSION_NAME=`cat /etc/os-release | grep UBUNTU_CODENAME | awk -F= '{print $2}'`
@@ -194,7 +243,7 @@ EOF
 
     pip_install "git+https://github.com/rocm/composable_kernel@$ROCM_COMPOSABLE_KERNEL_VERSION"
 
-    # Write environment to file that can be sourced by CI scripts and users
+    # Write environment file (sourced by CI scripts and interactive shells)
     cat > /etc/rocm_env.sh << ROCM_ENV
 # ROCm paths
 export ROCM_PATH=/opt/rocm
@@ -203,85 +252,18 @@ export ROCM_SOURCE_DIR=/opt/rocm
 export ROCM_BIN=/opt/rocm/bin
 export ROCM_CMAKE=/opt/rocm
 export PATH=/opt/rocm/bin:/opt/rocm/llvm/bin:\${PATH}
-# Device library paths
-export ROCM_DEVICE_LIB_PATH=/opt/rocm/amdgcn/bitcode
+export LD_LIBRARY_PATH=/opt/rocm/lib:\${LD_LIBRARY_PATH:-}
+# Device library path
 export HIP_DEVICE_LIB_PATH=/opt/rocm/amdgcn/bitcode
 export MAGMA_HOME=/opt/rocm/magma
 ROCM_ENV
 
-    # Append to bash.bashrc so interactive shells get the env vars
     echo "source /etc/rocm_env.sh" >> /etc/bash.bashrc
 
     # Cleanup
     apt-get autoclean && apt-get clean
     rm -rf /var/lib/apt/lists/* /tmp/* /var/tmp/*
-}
-
-install_centos() {
-
-  yum update -y
-  yum install -y kmod
-  yum install -y wget
-  yum install -y openblas-devel
-
-  yum install -y epel-release
-  yum install -y dkms kernel-headers-`uname -r` kernel-devel-`uname -r`
-
-  # Add amdgpu repository
-  local amdgpu_baseurl
-  if [[ $OS_VERSION == 9 ]]; then
-      amdgpu_baseurl="https://repo.radeon.com/amdgpu/${ROCM_VERSION}/rhel/9.0/main/x86_64"
-  else
-      amdgpu_baseurl="https://repo.radeon.com/amdgpu/${ROCM_VERSION}/rhel/7.9/main/x86_64"
-  fi
-  echo "[AMDGPU]" > /etc/yum.repos.d/amdgpu.repo
-  echo "name=AMDGPU" >> /etc/yum.repos.d/amdgpu.repo
-  echo "baseurl=${amdgpu_baseurl}" >> /etc/yum.repos.d/amdgpu.repo
-  echo "enabled=1" >> /etc/yum.repos.d/amdgpu.repo
-  echo "gpgcheck=1" >> /etc/yum.repos.d/amdgpu.repo
-  echo "gpgkey=http://repo.radeon.com/rocm/rocm.gpg.key" >> /etc/yum.repos.d/amdgpu.repo
-
-  local rocm_baseurl="http://repo.radeon.com/rocm/yum/${ROCM_VERSION}"
-  echo "[ROCm]" > /etc/yum.repos.d/rocm.repo
-  echo "name=ROCm" >> /etc/yum.repos.d/rocm.repo
-  echo "baseurl=${rocm_baseurl}" >> /etc/yum.repos.d/rocm.repo
-  echo "enabled=1" >> /etc/yum.repos.d/rocm.repo
-  echo "gpgcheck=1" >> /etc/yum.repos.d/rocm.repo
-  echo "gpgkey=http://repo.radeon.com/rocm/rocm.gpg.key" >> /etc/yum.repos.d/rocm.repo
-
-  yum update -y
-
-  yum install -y \
-                   rocm-dev \
-                   rocm-utils \
-                   rocm-libs \
-                   rccl \
-                   rocprofiler-dev \
-                   roctracer-dev \
-                   amd-smi-lib
-
-  # precompiled miopen kernels; search for all unversioned packages
-  # if search fails it will abort this script; use true to avoid case where search fails
-  MIOPENHIPGFX=$(yum -q search miopen-hip-gfx | grep miopen-hip-gfx | awk '{print $1}'| grep -F kdb. || true)
-  if [[ "x${MIOPENHIPGFX}" = x ]]; then
-    echo "miopen-hip-gfx package not available" && exit 1
-  else
-    yum install -y ${MIOPENHIPGFX}
-  fi
-
-  # ROCm 6.0 had a regression where journal_mode was enabled on the kdb files resulting in permission errors at runtime
-  for kdb in /opt/rocm/share/miopen/db/*.kdb
-  do
-      sqlite3 $kdb "PRAGMA journal_mode=off; PRAGMA VACUUM;"
-  done
-
-  pip_install "git+https://github.com/rocm/composable_kernel@$ROCM_COMPOSABLE_KERNEL_VERSION"
-
-  # Cleanup
-  yum clean all
-  rm -rf /var/cache/yum
-  rm -rf /var/lib/yum/yumdb
-  rm -rf /var/lib/yum/history
+    fi
 }
 
 # Install Python packages depending on the base OS
@@ -290,9 +272,6 @@ case "$ID" in
   ubuntu)
     install_ubuntu
     ;;
-  centos)
-    install_centos
-    ;;
   *)
     echo "Unable to determine OS..."
     exit 1
diff --git a/.ci/docker/common/install_rocm_drm.sh b/.ci/docker/common/install_rocm_drm.sh
index c70f5880f2c5c..a6b0fe2c03924 100644
--- a/.ci/docker/common/install_rocm_drm.sh
+++ b/.ci/docker/common/install_rocm_drm.sh
@@ -14,7 +14,7 @@ case "$ID" in
     apt-get install -y libpciaccess-dev pkg-config
     apt-get clean
     ;;
-  centos|almalinux)
+  almalinux)
     yum install -y libpciaccess-devel pkgconfig
     ;;
   *)
diff --git a/.ci/docker/common/install_torch_tpu.sh b/.ci/docker/common/install_torch_tpu.sh
index cfddd4badc741..c4e4104edfe42 100644
--- a/.ci/docker/common/install_torch_tpu.sh
+++ b/.ci/docker/common/install_torch_tpu.sh
@@ -60,7 +60,7 @@ fetch_secret() {
         set +x
     fi
 
-    if ! gcloud secrets versions access latest --secret="torchtpu-readonly-key" --project="ml-velocity-actions-testing" > "temp_ssh_key"; then
+    if ! gcloud secrets versions access latest --secret="torchtpu-read-key" --project="ml-velocity-actions-testing" > "temp_ssh_key"; then
         echo "Error: Failed to fetch secret. Ensure you are authenticated with gcloud."
 
         # Restore xtrace if it was enabled, before exiting
@@ -82,7 +82,7 @@ clone_repo() {
 
     # Use GIT_SSH_COMMAND to specify the key and disable strict host key checking for automation
     export GIT_SSH_COMMAND="ssh -i temp_ssh_key -o IdentitiesOnly=yes -o StrictHostKeyChecking=no"
-    if git clone --recursive "git@github.com:google-ml-infra/torch_tpu.git"; then
+    if git clone --recursive "git@github.com:google-pytorch/torch_tpu.git"; then
         echo "Repository cloned successfully."
     else
         echo "Error: Failed to clone repository."
@@ -110,7 +110,7 @@ pull_torch_tpu() {
 # sleep 28800 # Debug sleep to connect to runner to streamline debugging, do not submit
 
 # 3. Configuration
-TORCH_TPU_REPO="${TORCH_TPU_REPO:-https://github.com/google-ml-infra/torch_tpu.git}"
+TORCH_TPU_REPO="${TORCH_TPU_REPO:-https://github.com/google-pytorch/torch_tpu.git}"
 TORCH_TPU_BRANCH="${TORCH_TPU_BRANCH:-main}"
 
 # Pin File Configuration
@@ -124,7 +124,11 @@ fi
 if ! command -v bazel &> /dev/null; then
     echo "Bazel not found. Installing Bazelisk..."
     temp_dir=$(mktemp -d)
+    # Download Bazelisk v1.27.0
     curl -L https://github.com/bazelbuild/bazelisk/releases/download/v1.27.0/bazelisk-linux-amd64 -o "${temp_dir}/bazel"
+    # Verify Checksum (SHA256 for v1.27.0 linux-amd64)
+    # Source: https://github.com/bazelbuild/bazelisk/releases/tag/v1.27.0
+    echo "e1508323f347ad1465a887bc5d2bfb91cffc232d11e8e997b623227c6b32fb76  ${temp_dir}/bazel" | sha256sum --check
     sudo mv "${temp_dir}/bazel" /usr/local/bin/bazel
     sudo chmod +x /usr/local/bin/bazel
     rm -rf "${temp_dir}"
diff --git a/.ci/docker/common/install_triton.sh b/.ci/docker/common/install_triton.sh
index 1b68e3c247839..b2fdebdcc4747 100755
--- a/.ci/docker/common/install_triton.sh
+++ b/.ci/docker/common/install_triton.sh
@@ -21,7 +21,7 @@ elif [ -n "${TRITON_CPU}" ]; then
   TRITON_REPO="https://github.com/triton-lang/triton-cpu"
   TRITON_TEXT_FILE="triton-cpu"
 else
-  TRITON_REPO="https://github.com/triton-lang/triton"
+  TRITON_REPO="https://github.com/ROCm/triton"
   TRITON_TEXT_FILE="triton"
 fi
 
diff --git a/.ci/docker/common/install_ucc.sh b/.ci/docker/common/install_ucc.sh
deleted file mode 100755
index 6c97edb0e0287..0000000000000
--- a/.ci/docker/common/install_ucc.sh
+++ /dev/null
@@ -1,94 +0,0 @@
-#!/bin/bash
-
-set -ex
-
-if [[ -d "/usr/local/cuda/" ]];  then
-  with_cuda=/usr/local/cuda/
-else
-  with_cuda=no
-fi
-
-if [[ -f /etc/rocm_env.sh ]]; then
-  source /etc/rocm_env.sh
-fi
-
-if [[ -d "${ROCM_PATH}" ]]; then
-  with_rocm="${ROCM_PATH}"
-else
-  with_rocm=no
-fi
-
-function install_ucx() {
-  set -ex
-  git clone --recursive https://github.com/openucx/ucx.git
-  pushd ucx
-  git checkout ${UCX_COMMIT}
-  git submodule update --init --recursive
-
-  ./autogen.sh
-  ./configure --prefix=$UCX_HOME      \
-      --enable-mt                     \
-      --with-cuda=$with_cuda          \
-      --with-rocm=$with_rocm          \
-      --enable-profiling              \
-      --enable-stats
-  time make -j
-  sudo make install
-
-  popd
-  rm -rf ucx
-}
-
-function install_ucc() {
-  set -ex
-  git clone --recursive https://github.com/openucx/ucc.git
-  pushd ucc
-  git checkout ${UCC_COMMIT}
-  git submodule update --init --recursive
-
-  ./autogen.sh
-
-  if [[ -n "$CUDA_VERSION"  && $CUDA_VERSION == 13* ]]; then
-    NVCC_GENCODE="-gencode=arch=compute_86,code=compute_86"
-  else
-    # We only run distributed tests on Tesla M60 and A10G
-    NVCC_GENCODE="-gencode=arch=compute_52,code=sm_52 -gencode=arch=compute_86,code=compute_86"
-  fi
-
-  if [[ -n "$ROCM_VERSION" ]]; then
-    if [[ -n "$PYTORCH_ROCM_ARCH" ]]; then
-      amdgpu_targets=`echo $PYTORCH_ROCM_ARCH | sed 's/;/ /g'`
-    else
-      amdgpu_targets=`rocm_agent_enumerator | grep -v gfx000 | sort -u | xargs`
-    fi
-    for arch in $amdgpu_targets; do
-      HIP_OFFLOAD="$HIP_OFFLOAD --offload-arch=$arch"
-    done
-    HIP_OFFLOAD="$HIP_OFFLOAD --rocm-path=${ROCM_PATH}"
-
-    # Set device library path if detected (handles TheRock vs traditional ROCm)
-    if [ -n "${ROCM_DEVICE_LIB_PATH}" ] && [ -d "${ROCM_DEVICE_LIB_PATH}" ]; then
-      HIP_OFFLOAD="$HIP_OFFLOAD --rocm-device-lib-path=${ROCM_DEVICE_LIB_PATH}"
-    fi
-  else
-    HIP_OFFLOAD="all-arch-no-native"
-  fi
-
-  ./configure --prefix=$UCC_HOME          \
-    --with-ucx=$UCX_HOME                  \
-    --with-cuda=$with_cuda                \
-    --with-nvcc-gencode="${NVCC_GENCODE}" \
-    --with-rocm=$with_rocm                \
-    --with-rocm-arch="${HIP_OFFLOAD}"
-  # First observed by ROCm nightly builds, ucc rccl sources fail compile with
-  # error: #warning "NCCL C++ API is disabled because C compiler is being used. [-Werror=cpp]
-  # Work-around by adding make CFLAGS=-Wno-error=cpp
-  time make -j CFLAGS=-Wno-error=cpp
-  sudo make install
-
-  popd
-  rm -rf ucc
-}
-
-install_ucx
-install_ucc
diff --git a/.ci/docker/common/install_vision.sh b/.ci/docker/common/install_vision.sh
deleted file mode 100755
index 78c445568ddcd..0000000000000
--- a/.ci/docker/common/install_vision.sh
+++ /dev/null
@@ -1,46 +0,0 @@
-#!/bin/bash
-
-set -ex
-
-install_ubuntu() {
-  apt-get update
-  apt-get install -y --no-install-recommends \
-          libopencv-dev
-
-  # Cleanup
-  apt-get autoclean && apt-get clean
-  rm -rf /var/lib/apt/lists/* /tmp/* /var/tmp/*
-}
-
-install_centos() {
-  # Need EPEL for many packages we depend on.
-  # See http://fedoraproject.org/wiki/EPEL
-  yum --enablerepo=extras install -y epel-release
-
-  yum install -y \
-      opencv-devel
-
-  # Cleanup
-  yum clean all
-  rm -rf /var/cache/yum
-  rm -rf /var/lib/yum/yumdb
-  rm -rf /var/lib/yum/history
-}
-
-# Install base packages depending on the base OS
-ID=$(grep -oP '(?<=^ID=).+' /etc/os-release | tr -d '"')
-case "$ID" in
-  ubuntu)
-    install_ubuntu
-    ;;
-  centos)
-    install_centos
-    ;;
-  *)
-    echo "Unable to determine OS..."
-    exit 1
-    ;;
-esac
-
-# Cache vision models used by the test
-source "$(dirname "${BASH_SOURCE[0]}")/cache_vision_models.sh"
diff --git a/.ci/docker/libtorch/Dockerfile b/.ci/docker/libtorch/Dockerfile
deleted file mode 100644
index 9c4be5abe459d..0000000000000
--- a/.ci/docker/libtorch/Dockerfile
+++ /dev/null
@@ -1,117 +0,0 @@
-ARG BASE_TARGET=base
-ARG GPU_IMAGE=ubuntu:20.04
-FROM ${GPU_IMAGE} as base
-
-ENV DEBIAN_FRONTEND=noninteractive
-
-RUN apt-get clean && apt-get update
-RUN apt-get install -y curl locales g++ git-all autoconf automake make cmake wget unzip sudo
-# 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 '*'
-
-RUN locale-gen en_US.UTF-8
-
-ENV LC_ALL en_US.UTF-8
-ENV LANG en_US.UTF-8
-ENV LANGUAGE en_US.UTF-8
-
-# Install openssl
-FROM base as openssl
-ADD ./common/install_openssl.sh install_openssl.sh
-RUN bash ./install_openssl.sh && rm install_openssl.sh
-
-# Install python
-FROM base as python
-ADD common/install_cpython.sh install_cpython.sh
-RUN apt-get update -y && \
-    apt-get install build-essential gdb lcov libbz2-dev libffi-dev \
-        libgdbm-dev liblzma-dev libncurses5-dev libreadline6-dev \
-        libsqlite3-dev libssl-dev lzma lzma-dev tk-dev uuid-dev zlib1g-dev -y && \
-    bash ./install_cpython.sh && \
-    rm install_cpython.sh && \
-    apt-get clean
-
-FROM base as conda
-ADD ./common/install_conda_docker.sh install_conda.sh
-RUN bash ./install_conda.sh && rm install_conda.sh
-
-FROM base as cpu
-# Install Anaconda
-COPY --from=conda /opt/conda /opt/conda
-# Install python
-COPY --from=python /opt/python    /opt/python
-COPY --from=python /opt/_internal /opt/_internal
-ENV PATH=/opt/conda/bin:/usr/local/cuda/bin:$PATH
-# Install MKL
-ADD ./common/install_mkl.sh install_mkl.sh
-RUN bash ./install_mkl.sh && rm install_mkl.sh
-
-FROM cpu as cuda
-ADD ./common/install_cuda.sh install_cuda.sh
-ADD ./common/install_magma.sh install_magma.sh
-COPY ./common/install_nccl.sh install_nccl.sh
-COPY ./ci_commit_pins/nccl* /ci_commit_pins/
-COPY ./common/install_cusparselt.sh install_cusparselt.sh
-ENV CUDA_HOME /usr/local/cuda
-
-FROM cuda as cuda12.6
-RUN bash ./install_cuda.sh 12.6
-RUN bash ./install_magma.sh 12.6
-RUN ln -sf /usr/local/cuda-12.6 /usr/local/cuda
-
-FROM cuda as cuda12.8
-RUN bash ./install_cuda.sh 12.8
-RUN bash ./install_magma.sh 12.8
-RUN ln -sf /usr/local/cuda-12.8 /usr/local/cuda
-
-FROM cuda as cuda12.9
-RUN bash ./install_cuda.sh 12.9
-RUN bash ./install_magma.sh 12.9
-RUN ln -sf /usr/local/cuda-12.9 /usr/local/cuda
-
-FROM cuda as cuda13.0
-RUN bash ./install_cuda.sh 13.0
-RUN bash ./install_magma.sh 13.0
-RUN ln -sf /usr/local/cuda-13.0 /usr/local/cuda
-
-# Install libibverbs for libtorch and copy to CUDA directory
-RUN apt-get update -y && \
-    apt-get install -y libibverbs-dev librdmacm-dev && \
-    cp /usr/lib/x86_64-linux-gnu/libmlx5.so* /usr/local/cuda/lib64/ && \
-    cp /usr/lib/x86_64-linux-gnu/librdmacm.so* /usr/local/cuda/lib64/ && \
-    cp /usr/lib/x86_64-linux-gnu/libibverbs.so* /usr/local/cuda/lib64/ && \
-    cp /usr/lib/x86_64-linux-gnu/libnl* /usr/local/cuda/lib64/
-
-FROM cpu as rocm
-ARG ROCM_VERSION
-ARG PYTORCH_ROCM_ARCH
-ENV PYTORCH_ROCM_ARCH ${PYTORCH_ROCM_ARCH}
-ENV MKLROOT /opt/intel
-# Adding ROCM_PATH env var so that LoadHip.cmake (even with logic updated for ROCm6.0)
-# find HIP works for ROCm5.7. Not needed for ROCm6.0 and above.
-# Remove below when ROCm5.7 is not in support matrix anymore.
-ENV ROCM_PATH /opt/rocm
-# No need to install ROCm as base docker image should have full ROCm install
-#ADD ./common/install_rocm.sh install_rocm.sh
-ADD ./common/install_rocm_drm.sh install_rocm_drm.sh
-ADD ./common/install_rocm_magma.sh install_rocm_magma.sh
-# gfortran and python needed for building magma from source for ROCm
-RUN apt-get update -y && \
-    apt-get install gfortran -y && \
-    apt-get install python3 python-is-python3 -y && \
-    apt-get clean
-
-RUN bash ./install_rocm_drm.sh /opt/amdgpu && rm install_rocm_drm.sh
-RUN bash ./install_rocm_magma.sh ${ROCM_VERSION} && rm install_rocm_magma.sh
-
-FROM ${BASE_TARGET} as final
-COPY --from=openssl            /opt/openssl           /opt/openssl
-# Install patchelf
-ADD ./common/install_patchelf.sh install_patchelf.sh
-RUN bash ./install_patchelf.sh && rm install_patchelf.sh
-# Install Anaconda
-COPY --from=conda /opt/conda /opt/conda
-# Install python
-COPY --from=python /opt/python    /opt/python
-COPY --from=python /opt/_internal /opt/_internal
-ENV PATH=/opt/conda/bin:/usr/local/cuda/bin:$PATH
diff --git a/.ci/docker/libtorch/build.sh b/.ci/docker/libtorch/build.sh
deleted file mode 100755
index 5bfe70f34347e..0000000000000
--- a/.ci/docker/libtorch/build.sh
+++ /dev/null
@@ -1,75 +0,0 @@
-#!/usr/bin/env bash
-# Script used only in CD pipeline
-
-set -eoux pipefail
-
-image="$1"
-shift
-
-if [ -z "${image}" ]; then
-  echo "Usage: $0 IMAGENAME:ARCHTAG"
-  exit 1
-fi
-
-TOPDIR=$(git rev-parse --show-toplevel)
-
-DOCKER=${DOCKER:-docker}
-
-# Go from imagename:tag to tag
-DOCKER_TAG_PREFIX=$(echo "${image}" | awk -F':' '{print $2}')
-
-GPU_ARCH_VERSION=""
-if [[ "${DOCKER_TAG_PREFIX}" == cuda* ]]; then
-    # extract cuda version from image name.  e.g. manylinux2_28-builder:cuda12.8 returns 12.8
-    GPU_ARCH_VERSION=$(echo "${DOCKER_TAG_PREFIX}" | awk -F'cuda' '{print $2}')
-elif [[ "${DOCKER_TAG_PREFIX}" == rocm* ]]; then
-    # extract rocm version from image name.  e.g. manylinux2_28-builder:rocm6.2.4 returns 6.2.4
-    GPU_ARCH_VERSION=$(echo "${DOCKER_TAG_PREFIX}" | awk -F'rocm' '{print $2}')
-fi
-
-case ${DOCKER_TAG_PREFIX} in
-    cpu)
-        BASE_TARGET=cpu
-        GPU_IMAGE=ubuntu:20.04
-        DOCKER_GPU_BUILD_ARG=""
-        ;;
-    cuda*)
-        BASE_TARGET=cuda${GPU_ARCH_VERSION}
-        GPU_IMAGE=ubuntu:20.04
-        DOCKER_GPU_BUILD_ARG=""
-        ;;
-    rocm*)
-        # we want the patch version of 7.1 instead
-        if [[ "$GPU_ARCH_VERSION" == *"7.1"* ]]; then
-            GPU_ARCH_VERSION="${GPU_ARCH_VERSION}.1"
-        fi
-        # we want the patch version of 7.0 instead
-        if [[ "$GPU_ARCH_VERSION" == *"7.0"* ]]; then
-            GPU_ARCH_VERSION="${GPU_ARCH_VERSION}.2"
-        fi
-        # we want the patch version of 6.4 instead
-        if [[ "$GPU_ARCH_VERSION" == *"6.4"* ]]; then
-            GPU_ARCH_VERSION="${GPU_ARCH_VERSION}.4"
-        fi
-        BASE_TARGET=rocm
-        GPU_IMAGE=rocm/dev-ubuntu-22.04:${GPU_ARCH_VERSION}-complete
-        PYTORCH_ROCM_ARCH="gfx900;gfx906;gfx908;gfx90a;gfx942;gfx1030;gfx1100;gfx1101;gfx1102;gfx1200;gfx1201;gfx950;gfx1150;gfx1151"
-        DOCKER_GPU_BUILD_ARG="--build-arg PYTORCH_ROCM_ARCH=${PYTORCH_ROCM_ARCH} --build-arg ROCM_VERSION=${GPU_ARCH_VERSION}"
-        ;;
-    *)
-        echo "ERROR: Unrecognized DOCKER_TAG_PREFIX: ${DOCKER_TAG_PREFIX}"
-        exit 1
-        ;;
-esac
-
-tmp_tag=$(basename "$(mktemp -u)" | tr '[:upper:]' '[:lower:]')
-
-DOCKER_BUILDKIT=1 ${DOCKER} build \
-    --target final \
-    ${DOCKER_GPU_BUILD_ARG} \
-    --build-arg "GPU_IMAGE=${GPU_IMAGE}" \
-    --build-arg "BASE_TARGET=${BASE_TARGET}" \
-    -t "${tmp_tag}" \
-    $@ \
-    -f "${TOPDIR}/.ci/docker/libtorch/Dockerfile" \
-    "${TOPDIR}/.ci/docker/"
diff --git a/.ci/docker/manywheel/Dockerfile_2_28 b/.ci/docker/manywheel/Dockerfile_2_28
index 4055e6b872539..2ad5cd6498249 100644
--- a/.ci/docker/manywheel/Dockerfile_2_28
+++ b/.ci/docker/manywheel/Dockerfile_2_28
@@ -12,20 +12,12 @@ RUN yum install -y sudo wget curl perl util-linux xz bzip2 git patch which perl
 ENV PATH=/opt/rh/gcc-toolset-${DEVTOOLSET_VERSION}/root/usr/bin:$PATH
 ENV LD_LIBRARY_PATH=/opt/rh/gcc-toolset-${DEVTOOLSET_VERSION}/root/usr/lib64:/opt/rh/gcc-toolset-${DEVTOOLSET_VERSION}/root/usr/lib:$LD_LIBRARY_PATH
 
-# cmake-3.18.4 from pip
+# cmake-3.18.4.post1 from pip
+# NS: Apr 1 2026 3.18.4 is gone, reported here https://github.com/scikit-build/cmake-python-distributions/issues/693
 RUN yum install -y python3-pip && \
-    python3 -mpip install cmake==3.18.4 && \
+    python3 -mpip install cmake==3.18.4.post1 && \
     ln -s /usr/local/bin/cmake /usr/bin/cmake3
 
-FROM base as openssl
-# Install openssl (this must precede `build python` step)
-# (In order to have a proper SSL module, Python is compiled
-# against a recent openssl [see env vars above], which is linked
-# statically. We delete openssl afterwards.)
-ADD ./common/install_openssl.sh install_openssl.sh
-RUN bash ./install_openssl.sh && rm install_openssl.sh
-
-
 FROM base as cuda
 ARG BASE_CUDA_VERSION=12.6
 # Install CUDA
@@ -95,7 +87,6 @@ RUN git config --global --add safe.directory "*"
 
 ENV SSL_CERT_FILE=/opt/_internal/certs.pem
 # Install LLVM version
-COPY --from=openssl            /opt/openssl                          /opt/openssl
 COPY --from=base               /opt/python                           /opt/python
 COPY --from=base               /usr/local/lib/                       /usr/local/lib/
 COPY --from=base               /opt/_internal                        /opt/_internal
@@ -127,9 +118,9 @@ RUN for cpython_version in "cp312-cp312" "cp313-cp313" "cp313-cp313t"; do \
 ADD ./common/patch_libstdc.sh patch_libstdc.sh
 RUN bash ./patch_libstdc.sh && rm patch_libstdc.sh
 
-# cmake-3.18.4 from pip; force in case cmake3 already exists
+# cmake-3.18.4.post1 from pip; force in case cmake3 already exists
 RUN yum install -y python3-pip && \
-    python3 -mpip install cmake==3.18.4 && \
+    python3 -mpip install cmake==3.18.4.post1 && \
     ln -sf /usr/local/bin/cmake /usr/bin/cmake3
 
 FROM cpu_final as cuda_final
@@ -144,7 +135,8 @@ ARG ROCM_VERSION=6.0
 ARG PYTORCH_ROCM_ARCH
 ENV PYTORCH_ROCM_ARCH ${PYTORCH_ROCM_ARCH}
 ARG DEVTOOLSET_VERSION=13
-ENV LDFLAGS="-Wl,-rpath=/opt/rh/gcc-toolset-${DEVTOOLSET_VERSION}/root/usr/lib64 -Wl,-rpath=/opt/rh/gcc-toolset-${DEVTOOLSET_VERSION}/root/usr/lib"
+# All rocm clang cfg files load the same rocm.cfg, make sure it points to the right toolchain.
+RUN echo "--gcc-toolchain=/opt/rh/gcc-toolset-${DEVTOOLSET_VERSION}/root/usr" >> /opt/rocm/llvm/bin/rocm.cfg
 # Somewhere in ROCm stack, we still use non-existing /opt/rocm/hip path,
 # below workaround helps avoid error
 ENV ROCM_PATH /opt/rocm
@@ -160,6 +152,10 @@ RUN yum install -y libdrm-devel
 ENV MKLROOT /opt/intel
 ADD ./common/install_rocm_magma.sh install_rocm_magma.sh
 RUN bash ./install_rocm_magma.sh ${ROCM_VERSION} && rm install_rocm_magma.sh
+
+ADD ./common/install_rocSHMEM.sh install_rocSHMEM.sh
+RUN bash ./install_rocSHMEM.sh ${ROCM_VERSION} && rm install_rocSHMEM.sh
+
 ADD ./common/install_miopen.sh install_miopen.sh
 RUN bash ./install_miopen.sh ${ROCM_VERSION} && rm install_miopen.sh
 
diff --git a/.ci/docker/manywheel/Dockerfile_2_28_aarch64 b/.ci/docker/manywheel/Dockerfile_2_28_aarch64
index b5bf2ffc1c081..477e4221cb49f 100644
--- a/.ci/docker/manywheel/Dockerfile_2_28_aarch64
+++ b/.ci/docker/manywheel/Dockerfile_2_28_aarch64
@@ -39,12 +39,7 @@ RUN yum install -y \
   gcc-toolset-${GCCTOOLSET_VERSION}-gcc-c++ \
   gcc-toolset-${GCCTOOLSET_VERSION}-gcc-gfortran \
   gcc-toolset-${GCCTOOLSET_VERSION}-gdb
-
-# (optional) Install non-default Ninja version
-ARG NINJA_VERSION
-COPY ./common/install_ninja.sh install_ninja.sh
-RUN if [ -n "${NINJA_VERSION}" ]; then bash ./install_ninja.sh; fi
-RUN rm install_ninja.sh
+RUN yum install -y --enablerepo=powertools ninja-build
 
 # Ensure the expected devtoolset is used
 ENV PATH=/opt/rh/gcc-toolset-${GCCTOOLSET_VERSION}/root/usr/bin:$PATH
diff --git a/.ci/docker/manywheel/Dockerfile_cuda_aarch64 b/.ci/docker/manywheel/Dockerfile_cuda_aarch64
index 794a791b2721a..2a3f266fc413f 100644
--- a/.ci/docker/manywheel/Dockerfile_cuda_aarch64
+++ b/.ci/docker/manywheel/Dockerfile_cuda_aarch64
@@ -50,16 +50,7 @@ ENV LD_LIBRARY_PATH=/opt/rh/gcc-toolset-${DEVTOOLSET_VERSION}/root/usr/lib64:/op
 RUN git config --global --add safe.directory "*"
 
 
-FROM base as openssl
-# Install openssl (this must precede `build python` step)
-# (In order to have a proper SSL module, Python is compiled
-# against a recent openssl [see env vars above], which is linked
-# statically. We delete openssl afterwards.)
-ADD ./common/install_openssl.sh install_openssl.sh
-RUN bash ./install_openssl.sh && rm install_openssl.sh
-ENV SSL_CERT_FILE=/opt/_internal/certs.pem
-
-FROM openssl as final
+FROM base as final
 
 FROM base as cuda
 ARG BASE_CUDA_VERSION
diff --git a/.ci/docker/manywheel/Dockerfile_s390x b/.ci/docker/manywheel/Dockerfile_s390x
index 1cf83acb1c736..1367b004ee8a3 100644
--- a/.ci/docker/manywheel/Dockerfile_s390x
+++ b/.ci/docker/manywheel/Dockerfile_s390x
@@ -84,7 +84,7 @@ RUN cp $(which patchelf) /patchelf
 
 FROM patchelf as python
 # build python
-COPY manywheel/build_scripts /build_scripts
+COPY manywheel/s390_scripts /build_scripts
 ADD ./common/install_cpython.sh /build_scripts/install_cpython.sh
 ENV SSL_CERT_FILE=
 RUN bash build_scripts/build.sh && rm -r build_scripts
diff --git a/.ci/docker/manywheel/build.sh b/.ci/docker/manywheel/build.sh
index b0047f98290b3..05e468d72dda9 100755
--- a/.ci/docker/manywheel/build.sh
+++ b/.ci/docker/manywheel/build.sh
@@ -40,7 +40,7 @@ case ${image} in
     manylinux2_28_aarch64-builder:cpu-aarch64)
         TARGET=final
         GPU_IMAGE=arm64v8/almalinux:8
-        DOCKER_GPU_BUILD_ARG=" --build-arg DEVTOOLSET_VERSION=13 --build-arg NINJA_VERSION=1.12.1"
+        DOCKER_GPU_BUILD_ARG=" --build-arg DEVTOOLSET_VERSION=13"
         MANY_LINUX_VERSION="2_28_aarch64"
         ;;
     manylinuxs390x-builder:cpu-s390x)
@@ -75,6 +75,10 @@ case ${image} in
         DOCKERFILE_SUFFIX="_cuda_aarch64"
         ;;
     manylinux2_28-builder:rocm*)
+        # we want the patch version of 7.2 instead
+        if [[ "$GPU_ARCH_VERSION" == *"7.2"* ]]; then
+            GPU_ARCH_VERSION="${GPU_ARCH_VERSION}.1"
+        fi
         # we want the patch version of 7.1 instead
         if [[ "$GPU_ARCH_VERSION" == *"7.1"* ]]; then
             GPU_ARCH_VERSION="${GPU_ARCH_VERSION}.1"
@@ -89,7 +93,7 @@ case ${image} in
         fi
         TARGET=rocm_final
         MANY_LINUX_VERSION="2_28"
-        DEVTOOLSET_VERSION="11"
+        DEVTOOLSET_VERSION="13"
         GPU_IMAGE=rocm/dev-almalinux-8:${GPU_ARCH_VERSION}-complete
         PYTORCH_ROCM_ARCH="gfx900;gfx906;gfx908;gfx90a;gfx942;gfx1030;gfx1100;gfx1101;gfx1102;gfx1200;gfx1201;gfx950;gfx1150;gfx1151"
         DOCKER_GPU_BUILD_ARG="--build-arg ROCM_VERSION=${GPU_ARCH_VERSION} --build-arg PYTORCH_ROCM_ARCH=${PYTORCH_ROCM_ARCH} --build-arg DEVTOOLSET_VERSION=${DEVTOOLSET_VERSION}"
diff --git a/.ci/docker/manywheel/build_scripts/manylinux1-check.py b/.ci/docker/manywheel/build_scripts/manylinux1-check.py
deleted file mode 100644
index f6b9b9fc2393e..0000000000000
--- a/.ci/docker/manywheel/build_scripts/manylinux1-check.py
+++ /dev/null
@@ -1,63 +0,0 @@
-# Logic copied from PEP 513
-
-
-def is_manylinux1_compatible():
-    # Only Linux, and only x86-64 / i686
-    from distutils.util import get_platform
-
-    if get_platform() not in ["linux-x86_64", "linux-i686", "linux-s390x"]:
-        return False
-
-    # Check for presence of _manylinux module
-    try:
-        import _manylinux
-
-        return bool(_manylinux.manylinux1_compatible)
-    except (ImportError, AttributeError):
-        # Fall through to heuristic check below
-        pass
-
-    # Check glibc version. CentOS 5 uses glibc 2.5.
-    return have_compatible_glibc(2, 5)
-
-
-def have_compatible_glibc(major, minimum_minor):
-    import ctypes
-
-    process_namespace = ctypes.CDLL(None)
-    try:
-        gnu_get_libc_version = process_namespace.gnu_get_libc_version
-    except AttributeError:
-        # Symbol doesn't exist -> therefore, we are not linked to
-        # glibc.
-        return False
-
-    # Call gnu_get_libc_version, which returns a string like "2.5".
-    gnu_get_libc_version.restype = ctypes.c_char_p
-    version_str = gnu_get_libc_version()
-    # py2 / py3 compatibility:
-    if not isinstance(version_str, str):
-        version_str = version_str.decode("ascii")
-
-    # Parse string and check against requested version.
-    version = [int(piece) for piece in version_str.split(".")]
-    if len(version) != 2:
-        raise AssertionError(
-            f"Expected version to have 2 components (major.minor), got {len(version)}: {version_str}"
-        )
-    if major != version[0]:
-        return False
-    if minimum_minor > version[1]:
-        return False
-    return True
-
-
-import sys
-
-
-if is_manylinux1_compatible():
-    print(f"{sys.executable} is manylinux1 compatible")
-    sys.exit(0)
-else:
-    print(f"{sys.executable} is NOT manylinux1 compatible")
-    sys.exit(1)
diff --git a/.ci/docker/manywheel/build_scripts/ssl-check.py b/.ci/docker/manywheel/build_scripts/ssl-check.py
deleted file mode 100644
index c4df0eacbb7fd..0000000000000
--- a/.ci/docker/manywheel/build_scripts/ssl-check.py
+++ /dev/null
@@ -1,26 +0,0 @@
-# cf. https://github.com/pypa/manylinux/issues/53
-
-import sys
-from urllib.request import urlopen
-
-
-GOOD_SSL = "https://google.com"
-BAD_SSL = "https://self-signed.badssl.com"
-
-
-print("Testing SSL certificate checking for Python:", sys.version)
-
-EXC = OSError
-
-print(f"Connecting to {GOOD_SSL} should work")
-urlopen(GOOD_SSL)
-print("...it did, yay.")
-
-print(f"Connecting to {BAD_SSL} should fail")
-try:
-    urlopen(BAD_SSL)
-    # If we get here then we failed:
-    print("...it DIDN'T!!!!!11!!1one!")
-    sys.exit(1)
-except EXC:
-    print("...it did, yay.")
diff --git a/.ci/docker/manywheel/build_scripts/build.sh b/.ci/docker/manywheel/s390_scripts/build.sh
similarity index 90%
rename from .ci/docker/manywheel/build_scripts/build.sh
rename to .ci/docker/manywheel/s390_scripts/build.sh
index b6a70f0a72787..13141dfd4ae33 100644
--- a/.ci/docker/manywheel/build_scripts/build.sh
+++ b/.ci/docker/manywheel/s390_scripts/build.sh
@@ -18,13 +18,7 @@ AUTOCONF_HASH=954bd69b391edc12d6a4a51a2dd1476543da5c6bbf05a95b59dc0dd6fd4c2969
 
 # Dependencies for compiling Python that we want to remove from
 # the final image after compiling Python
-PYTHON_COMPILE_DEPS="zlib-devel bzip2-devel ncurses-devel sqlite-devel readline-devel tk-devel gdbm-devel libpcap-devel xz-devel libffi-devel"
-
-if [ "$(uname -m)" != "s390x" ] ; then
-    PYTHON_COMPILE_DEPS="${PYTHON_COMPILE_DEPS} db4-devel"
-else
-    PYTHON_COMPILE_DEPS="${PYTHON_COMPILE_DEPS} libdb-devel"
-fi
+PYTHON_COMPILE_DEPS="zlib-devel bzip2-devel ncurses-devel sqlite-devel readline-devel tk-devel gdbm-devel libpcap-devel xz-devel libffi-devel libdb-devel"
 
 # Libraries that are allowed as part of the manylinux1 profile
 MANYLINUX1_DEPS="glibc-devel libstdc++-devel glib2-devel libX11-devel libXext-devel libXrender-devel  mesa-libGL-devel libICE-devel libSM-devel ncurses-devel"
@@ -103,13 +97,6 @@ find /opt/_internal \
   -o \( -type f -a -name '*.pyc' -o -name '*.pyo' \) \
   -print0 | xargs -0 rm -f
 
-for PYTHON in /opt/python/*/bin/python; do
-    # Smoke test to make sure that our Pythons work, and do indeed detect as
-    # being manylinux compatible:
-    $PYTHON $MY_DIR/manylinux1-check.py
-    # Make sure that SSL cert checking works
-    $PYTHON $MY_DIR/ssl-check.py
-done
 
 # Fix libc headers to remain compatible with C99 compilers.
 find /usr/include/ -type f -exec sed -i 's/\bextern _*inline_*\b/extern __inline __attribute__ ((__gnu_inline__))/g' {} +
diff --git a/.ci/docker/manywheel/build_scripts/build_utils.sh b/.ci/docker/manywheel/s390_scripts/build_utils.sh
similarity index 100%
rename from .ci/docker/manywheel/build_scripts/build_utils.sh
rename to .ci/docker/manywheel/s390_scripts/build_utils.sh
diff --git a/.ci/docker/requirements-ci.txt b/.ci/docker/requirements-ci.txt
index 9a033b90fcb46..2ed9f431e7ae6 100644
--- a/.ci/docker/requirements-ci.txt
+++ b/.ci/docker/requirements-ci.txt
@@ -15,13 +15,13 @@ build==1.3.0
 #Pinned versions: 1.3.0
 #test that import:
 
-click
+click==8.3.1
 #Description: Command Line Interface Creation Kit
-#Pinned versions:
+#Pinned versions: 8.3.1
 #test that import:
 
 coremltools==5.0b5 ; python_version < "3.12"
-coremltools==8.3 ; python_version == "3.12"
+coremltools==8.3.0 ; python_version == "3.12"
 #Description: Apple framework for ML integration
 #Pinned versions: 5.0b5
 #test that import:
@@ -42,7 +42,7 @@ expecttest==0.3.0
 #Pinned versions: 0.3.0
 #test that import:
 
-fbscribelogger==0.1.7
+fbscribelogger==0.1.7 ; python_version < "3.14"
 #Description: write to scribe from authenticated jobs on CI
 #Pinned versions: 0.1.6
 #test that import:
@@ -68,23 +68,23 @@ lark==0.12.0
 #Pinned versions: 0.12.0
 #test that import:
 
-librosa>=0.6.2 ; python_version < "3.11" and platform_machine != "s390x"
+librosa==0.10.2 ; python_version < "3.11" and platform_machine != "s390x"
 librosa==0.10.2 ; python_version == "3.12" and platform_machine != "s390x"
 #Description: A python package for music and audio analysis
 #Pinned versions: >=0.6.2
 #test that import: test_spectral_ops.py
 #librosa depends on numba; disable it for s390x while numba is disabled too
 
-#mkl #this breaks linux-bionic-rocm4.5-py3.7
+# Only mkl-static and mkl-include are needed; the mkl package contains
+# dynamic libraries that are not discoverable by our build scripts.
+mkl-static==2024.2.0 ; platform_machine != "aarch64" and sys_platform != "darwin"
+mkl-include==2024.2.0 ; platform_machine != "aarch64" and sys_platform != "darwin"
 #Description: Intel oneAPI Math Kernel Library
-#Pinned versions:
+#Pinned versions: 2024.2.0
 #test that import: test_profiler.py, test_public_bindings.py, test_testing.py,
 #test_nn.py, test_mkldnn.py, test_jit.py, test_fx_experimental.py,
 #test_autograd.py
 
-#mkl-devel
-# see mkl
-
 #mock
 #Description: A testing library that allows you to replace parts of your
 #system under test with mock objects
@@ -105,7 +105,8 @@ mypy==1.16.0 ; platform_system == "Linux"
 #Pinned versions: 1.16.0
 #test that import: test_typing.py, test_type_hints.py
 
-networkx==2.8.8
+networkx==2.8.8 ; python_version < "3.13"
+networkx==3.0.0 ; python_version >= "3.13"
 #Description: creation, manipulation, and study of
 #the structure, dynamics, and functions of complex networks
 #Pinned versions: 2.8.8
@@ -117,8 +118,9 @@ ninja==1.11.1.4
 #Pinned versions: 1.11.1.4
 #test that import: run_test.py, test_cpp_extensions_aot.py,test_determination.py
 
-numba==0.57.1 ; python_version == "3.10" and platform_machine != "s390x"
-numba==0.60.0 ; python_version == "3.12" and platform_machine != "s390x"
+numba==0.61.2 ; python_version < "3.14" and platform_machine != "s390x"
+numba==0.64.0 ; python_version >= "3.14" and platform_machine != "s390x"
+
 #Description: Just-In-Time Compiler for Numerical Functions
 #Pinned versions: 0.55.2, 0.60.0
 #test that import: test_numba_integration.py
@@ -136,13 +138,10 @@ numba==0.60.0 ; python_version == "3.12" and platform_machine != "s390x"
 #test_nn.py, test_namedtensor.py, test_linalg.py, test_jit_cuda_fuser.py,
 #test_jit.py, test_indexing.py, test_datapipe.py, test_dataloader.py,
 #test_binary_ufuncs.py
-numpy==1.23.2; python_version == "3.10"
-numpy==1.26.2; python_version == "3.11" or python_version == "3.12"
-numpy==2.1.2; python_version >= "3.13" and python_version < "3.14"
+numpy==2.1.2 ; python_version < "3.14"
 numpy==2.3.4; python_version >= "3.14"
 
-pandas==2.0.3; python_version < "3.12"
-pandas==2.2.3; python_version >= "3.12" and python_version < "3.14"
+pandas==2.2.3; python_version < "3.14"
 pandas==2.3.3; python_version >= "3.14"
 
 #onnxruntime
@@ -150,9 +149,9 @@ pandas==2.3.3; python_version >= "3.14"
 #Pinned versions: 1.9.0
 #test that import:
 
-opt-einsum==3.3
+opt-einsum==3.3.0
 #Description: Python library to optimize tensor contraction order, used in einsum
-#Pinned versions: 3.3
+#Pinned versions: 3.3.0
 #test that import: test_linalg.py
 
 optree==0.13.0 ; python_version < "3.14"
@@ -169,7 +168,7 @@ optree==0.17.0 ; python_version >= "3.14"
 #test_pointwise_ops.py, test_dtensor_ops.py, test_torchinductor.py, test_fx.py,
 #test_fake_tensor.py, test_mps.py
 
-pillow==11.0.0
+pillow==12.2.0
 #Description:  Python Imaging Library fork
 #Pinned versions: 11.0.0
 #test that import:
@@ -179,9 +178,9 @@ protobuf==6.33.5
 #Pinned versions: 6.33.2
 #test that import: test_tensorboard.py, test/onnx/*
 
-psutil
+psutil==7.2.2
 #Description: information on running processes and system utilization
-#Pinned versions:
+#Pinned versions: 7.2.2
 #test that import: test_profiler.py, test_openmp.py, test_dataloader.py
 
 pytest==7.3.2
@@ -199,9 +198,9 @@ pytest-flakefinder==1.1.0
 #Pinned versions: 1.1.0
 #test that import:
 
-pytest-rerunfailures>=10.3
+pytest-rerunfailures==14.0
 #Description: plugin for rerunning failure tests in pytest
-#Pinned versions:
+#Pinned versions: 14.0
 #test that import:
 
 pytest-subtests==0.13.1
@@ -224,7 +223,7 @@ xdoctest==1.3.0
 #Pinned versions: 1.1.0
 #test that import:
 
-pygments==2.15.0
+pygments==2.20.0
 #Description: support doctest highlighting
 #Pinned versions: 2.12.0
 #test that import: the doctests
@@ -244,7 +243,8 @@ pygments==2.15.0
 #Pinned versions: 14.1.0
 #test that import:
 
-scikit-image==0.22.0
+scikit-image==0.22.0 ; python_version < "3.13"
+scikit-image==0.26.0 ; python_version >= "3.13"
 #Description: image processing routines
 #Pinned versions: 0.22.0
 #test that import: test_nn.py
@@ -254,9 +254,9 @@ scikit-image==0.22.0
 #Pinned versions: 0.20.3
 #test that import:
 
-scipy==1.10.1 ; python_version <= "3.11"
-scipy==1.14.1 ; python_version > "3.11" and python_version < "3.14"
+scipy==1.14.1 ; python_version < "3.14"
 scipy==1.16.2 ; python_version >= "3.14"
+
 # Pin SciPy because of failing distribution tests (see #60347)
 #Description: scientific python
 #Pinned versions: 1.10.1
@@ -270,8 +270,7 @@ scipy==1.16.2 ; python_version >= "3.14"
 #test that import:
 
 # needed by torchgen utils
-typing-extensions==4.12.2 ; python_version < "3.14"
-typing-extensions==4.15.0 ; python_version >= "3.14"
+typing-extensions==4.15.0
 #Description: type hints for python
 #Pinned versions:
 #test that import:
@@ -281,18 +280,22 @@ typing-extensions==4.15.0 ; python_version >= "3.14"
 #Pinned versions:
 #test that import:
 
-unittest-xml-reporting<=3.2.0,>=2.0.0
+unittest-xml-reporting==3.2.0
 #Description: saves unit test results to xml
 #Pinned versions:
 #test that import:
 
-#lintrunner is supported on aarch64-linux only from 0.12.4 version
-lintrunner==0.12.11
+lintrunner==0.13.0
 #Description: all about linters!
-#Pinned versions: 0.12.11
+#Pinned versions: 0.13.0
+#test that import:
+
+spin==0.17
+#Description: developer CLI for common build/lint tasks
+#Pinned versions: 0.17
 #test that import:
 
-redis>=4.0.0
+redis==7.4.0
 #Description: redis database
 #test that import: anything that tests OSS caching/mocking (inductor/test_codecache.py, inductor/test_max_autotune.py)
 
@@ -316,14 +319,14 @@ z3-solver==4.15.1.0 ; platform_machine != "s390x"
 #Pinned versions:
 #test that import:
 
-tensorboard==2.13.0 ; python_version < "3.13"
-tensorboard==2.18.0 ; python_version >= "3.13"
+tensorboard==2.18.0
 #Description: Also included in .ci/docker/requirements-docs.txt
 #Pinned versions:
 #test that import: test_tensorboard
 
 pywavelets==1.4.1 ; python_version < "3.12"
-pywavelets==1.7.0 ; python_version >= "3.12"
+pywavelets==1.7.0 ; python_version >= "3.12" and python_version < "3.14"
+pywavelets==1.9.0 ; python_version >= "3.14"
 #Description: This is a requirement of scikit-image, we need to pin
 # it here because 1.5.0 conflicts with numpy 1.21.2 used in CI
 #Pinned versions: 1.4.1
@@ -340,7 +343,7 @@ sympy==1.13.3
 #Pinned versions:
 #test that import:
 
-onnx==1.20.0
+onnx==1.21.0
 #Description: Required by the torch.onnx exporter
 #Pinned versions:
 #test that import:
@@ -370,11 +373,13 @@ pwlf==2.2.1
 #test that import: test_sac_estimator.py
 
 # To build PyTorch itself
+pip==26.0.1
 pyyaml==6.0.3
-pyzstd
-setuptools==78.1.1
-packaging==24.0
-six
+pyzstd==0.16.2 ; python_version < "3.14"
+pyzstd==0.18.0 ; python_version >= "3.14"
+setuptools==79.0.1
+packaging==25.0
+six==1.17.0
 
 scons==4.5.2 ; platform_machine == "aarch64"
 
@@ -397,7 +402,7 @@ tlparse==0.4.0
 filelock==3.20.3
 #Description: required for inductor testing
 
-cuda-bindings>=12.0,<13.0 ; platform_machine != "s390x" and platform_system != "Darwin"
+cuda-bindings==12.9.6 ; platform_machine != "s390x" and platform_system != "Darwin"
 #Description: required for testing CUDAGraph::raw_cuda_graph(). See https://nvidia.github.io/cuda-python/cuda-bindings/latest/support.html for how this version was chosen. Note "Any fix in the latest bindings would be backported to the prior major version" means that only the newest version of cuda-bindings will get fixes. Depending on the latest version of 12.x is okay because all 12.y versions will be supported via "CUDA minor version compatibility". Pytorch builds against 13.z versions of cuda toolkit work with 12.x versions of cuda-bindings as well because newer drivers work with old toolkits.
 #test that import: test_cuda.py
 
@@ -407,9 +412,9 @@ pyre-extensions==0.0.32
 tabulate==0.9.0
 #Description: These package are needed to build FBGEMM and torchrec on PyTorch CI
 
-tqdm>=4.66.0
+tqdm==4.67.3
 #Description: progress bar library required for dynamo benchmarks
 #test that import: benchmarks/dynamo/*
 
-aiohttp==3.13.3
+aiohttp==3.13.4
 #Description: required for torch.distributed.debug
diff --git a/.ci/docker/requirements-docs.txt b/.ci/docker/requirements-docs.txt
index 7f3e0b5cc9215..484d99ec1152e 100644
--- a/.ci/docker/requirements-docs.txt
+++ b/.ci/docker/requirements-docs.txt
@@ -2,17 +2,14 @@ sphinx==7.2.6
 #Description: This is used to generate PyTorch docs
 #Pinned versions: 7.2.6
 
-pytorch_sphinx_theme2==0.4.3
+pytorch_sphinx_theme2==0.4.9
 #Description: This is needed to generate PyTorch docs
-#Pinned versions: 0.4.3
+#Pinned versions: 0.4.9
 
-# 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.
-
-sphinxcontrib.katex==0.8.6
+sphinxcontrib.katex==0.9.11
 #Description: This is used to generate PyTorch docs
-#Pinned versions: 0.8.6
+#Pinned versions: 0.9.11 (0.9.0+ uses a persistent KaTeX server instead of
+#  spawning a subprocess per math expression, ~20% faster writes)
 
 sphinxext-opengraph==0.9.1
 #Description: This is used to generate PyTorch docs
@@ -48,6 +45,10 @@ docutils==0.20
 #Description: This is used to generate PyTorch C++ docs
 #Pinned versions: 0.20
 
+coverxygen==1.8.1
+#Description: This is used to measure C++ API doc coverage from Doxygen XML
+#Pinned versions: 1.8.1
+
 bs4==0.0.1
 #Description: This is used to generate PyTorch C++ docs
 #Pinned versions: 0.0.1
diff --git a/.ci/docker/triton_version.txt b/.ci/docker/triton_version.txt
index 40c341bdcdbe8..19811903a7f75 100644
--- a/.ci/docker/triton_version.txt
+++ b/.ci/docker/triton_version.txt
@@ -1 +1 @@
-3.6.0
+3.8.0
diff --git a/.ci/docker/triton_xpu_version.txt b/.ci/docker/triton_xpu_version.txt
index 7c69a55dbb185..a76ccff2a6e0d 100644
--- a/.ci/docker/triton_xpu_version.txt
+++ b/.ci/docker/triton_xpu_version.txt
@@ -1 +1 @@
-3.7.0
+3.7.1
diff --git a/.ci/docker/ubuntu-cross-riscv/Dockerfile b/.ci/docker/ubuntu-cross-riscv/Dockerfile
index 08201dc83216c..8cf540f8414af 100644
--- a/.ci/docker/ubuntu-cross-riscv/Dockerfile
+++ b/.ci/docker/ubuntu-cross-riscv/Dockerfile
@@ -38,7 +38,7 @@ COPY ./common/install_user.sh install_user.sh
 RUN bash ./install_user.sh && rm install_user.sh
 
 FROM base as python
-ARG ZLIB_VERSION=1.3.1
+ARG ZLIB_VERSION=1.3.2
 ARG FFI_VERSION=3.4.6
 ARG BZ2_VERSION=1.0.8
 ARG XZ_VERSION=5.4.6
diff --git a/.ci/docker/ubuntu-rocm/Dockerfile b/.ci/docker/ubuntu-rocm/Dockerfile
index 3f487eb12809e..a3d697f1b27f2 100644
--- a/.ci/docker/ubuntu-rocm/Dockerfile
+++ b/.ci/docker/ubuntu-rocm/Dockerfile
@@ -13,6 +13,7 @@ ENV PYTORCH_ROCM_ARCH ${PYTORCH_ROCM_ARCH}
 # Install common dependencies (so that this step can be cached separately)
 COPY ./common/install_base.sh install_base.sh
 RUN bash ./install_base.sh && rm install_base.sh
+RUN apt-get update && apt-get install -y --no-install-recommends libtbb-dev && rm -rf /var/lib/apt/lists/*
 
 # Install user
 COPY ./common/install_user.sh install_user.sh
@@ -43,13 +44,6 @@ ARG CLANG_VERSION
 COPY ./common/install_clang.sh install_clang.sh
 RUN bash ./install_clang.sh && rm install_clang.sh
 
-# (optional) Install vision packages like OpenCV
-ARG VISION
-COPY ./common/install_vision.sh ./common/cache_vision_models.sh ./common/common_utils.sh ./
-RUN if [ -n "${VISION}" ]; then bash ./install_vision.sh; fi
-RUN rm install_vision.sh cache_vision_models.sh common_utils.sh
-ENV INSTALLED_VISION ${VISION}
-
 # Install rocm
 ARG ROCM_VERSION
 ENV ROCM_VERSION=${ROCM_VERSION}
@@ -63,14 +57,27 @@ RUN rm -r ci_commit_pins
 COPY ./common/install_rocm_magma.sh install_rocm_magma.sh
 RUN if [ "${ROCM_VERSION}" != "nightly" ]; then bash ./install_rocm_magma.sh ${ROCM_VERSION}; fi
 RUN rm install_rocm_magma.sh
+COPY ./common/install_rocSHMEM.sh install_rocSHMEM.sh
+RUN bash ./install_rocSHMEM.sh ${ROCM_VERSION}
+RUN rm install_rocSHMEM.sh
 ADD ./common/install_miopen.sh install_miopen.sh
 RUN if [ "${ROCM_VERSION}" != "nightly" ]; then bash ./install_miopen.sh ${ROCM_VERSION}; fi && rm install_miopen.sh
 ADD ./common/install_rocm_drm.sh install_rocm_drm.sh
 RUN if [ "${ROCM_VERSION}" != "nightly" ]; then bash ./install_rocm_drm.sh /usr ; fi && rm install_rocm_drm.sh
 
-# ROCm environment variables are set in /etc/rocm_env.sh by install_rocm.sh
-# and sourced via /etc/bash.bashrc for interactive shells.
-# CI scripts should source /etc/rocm_env.sh directly.
+# Default ROCm environment; /etc/rocm_env.sh (created by install_rocm.sh) may
+# override these at runtime for different install methods (tarballs vs wheels).
+ENV ROCM_PATH=/opt/rocm \
+    ROCM_HOME=/opt/rocm \
+    ROCM_SOURCE_DIR=/opt/rocm \
+    ROCM_BIN=/opt/rocm/bin \
+    ROCM_CMAKE=/opt/rocm \
+    ROCM_DEVICE_LIB_PATH=/opt/rocm/amdgcn/bitcode \
+    HIP_DEVICE_LIB_PATH=/opt/rocm/amdgcn/bitcode \
+    MAGMA_HOME=/opt/rocm/magma
+ENV PATH=/opt/rocm/bin:/opt/rocm/llvm/bin:$PATH
+ENV LD_LIBRARY_PATH=/opt/rocm/lib:${LD_LIBRARY_PATH:-}
+
 ENV LANG C.UTF-8
 ENV LC_ALL C.UTF-8
 
@@ -79,22 +86,6 @@ COPY ./common/install_amdsmi.sh install_amdsmi.sh
 RUN bash ./install_amdsmi.sh
 RUN rm install_amdsmi.sh
 
-# (optional) Install UCC
-ARG UCX_COMMIT
-ARG UCC_COMMIT
-ENV UCX_COMMIT $UCX_COMMIT
-ENV UCC_COMMIT $UCC_COMMIT
-ENV UCX_HOME /usr
-ENV UCC_HOME /usr
-ADD ./common/install_ucc.sh install_ucc.sh
-RUN if [ -n "${UCX_COMMIT}" ] && [ -n "${UCC_COMMIT}" ]; then bash ./install_ucc.sh; fi
-RUN rm install_ucc.sh
-
-COPY ./common/install_openssl.sh install_openssl.sh
-ENV OPENSSL_ROOT_DIR /opt/openssl
-RUN bash ./install_openssl.sh
-ENV OPENSSL_DIR /opt/openssl
-
 ARG INDUCTOR_BENCHMARKS
 ARG ANACONDA_PYTHON_VERSION
 ENV ANACONDA_PYTHON_VERSION=$ANACONDA_PYTHON_VERSION
@@ -106,12 +97,6 @@ COPY ci_commit_pins/torchbench.txt torchbench.txt
 RUN if [ -n "${INDUCTOR_BENCHMARKS}" ]; then bash ./install_inductor_benchmark_deps.sh; fi
 RUN rm install_inductor_benchmark_deps.sh common_utils.sh timm.txt huggingface-requirements.txt torchbench.txt
 
-# (optional) Install non-default Ninja version
-ARG NINJA_VERSION
-COPY ./common/install_ninja.sh install_ninja.sh
-RUN if [ -n "${NINJA_VERSION}" ]; then bash ./install_ninja.sh; fi
-RUN rm install_ninja.sh
-
 ARG TRITON
 # Install triton, this needs to be done before sccache because the latter will
 # try to reach out to S3, which docker build runners don't have access
@@ -137,8 +122,5 @@ RUN rm install_openmpi.sh
 ARG BUILD_ENVIRONMENT
 ENV BUILD_ENVIRONMENT ${BUILD_ENVIRONMENT}
 
-# Install LLVM dev version (Defined in the pytorch/builder github repository)
-COPY --from=pytorch/llvm:9.0.1 /opt/llvm /opt/llvm
-
 USER jenkins
 CMD ["bash"]
diff --git a/.ci/docker/ubuntu-xpu/Dockerfile b/.ci/docker/ubuntu-xpu/Dockerfile
index c61612882032d..3aba712da78d7 100644
--- a/.ci/docker/ubuntu-xpu/Dockerfile
+++ b/.ci/docker/ubuntu-xpu/Dockerfile
@@ -47,12 +47,6 @@ RUN bash ./install_gcc.sh && rm install_gcc.sh
 COPY ./common/install_lcov.sh install_lcov.sh
 RUN  bash ./install_lcov.sh && rm install_lcov.sh
 
-COPY ./common/install_openssl.sh install_openssl.sh
-RUN bash ./install_openssl.sh
-ENV OPENSSL_ROOT_DIR /opt/openssl
-ENV OPENSSL_DIR /opt/openssl
-RUN rm install_openssl.sh
-
 ARG INDUCTOR_BENCHMARKS
 ARG ANACONDA_PYTHON_VERSION
 ENV ANACONDA_PYTHON_VERSION=$ANACONDA_PYTHON_VERSION
@@ -80,19 +74,6 @@ COPY triton_xpu_version.txt triton_version.txt
 RUN if [ -n "${TRITON}" ]; then bash ./install_triton.sh; fi
 RUN rm install_triton.sh common_utils.sh triton-xpu.txt triton_version.txt
 
-# (optional) Install vision packages like OpenCV
-ARG VISION
-COPY ./common/install_vision.sh ./common/cache_vision_models.sh ./common/common_utils.sh ./
-RUN if [ -n "${VISION}" ]; then bash ./install_vision.sh; fi
-RUN rm install_vision.sh cache_vision_models.sh common_utils.sh
-ENV INSTALLED_VISION ${VISION}
-
-# (optional) Install non-default Ninja version
-ARG NINJA_VERSION
-COPY ./common/install_ninja.sh install_ninja.sh
-RUN if [ -n "${NINJA_VERSION}" ]; then bash ./install_ninja.sh; fi
-RUN rm install_ninja.sh
-
 # Install ccache/sccache (do this last, so we get priority in PATH)
 COPY ./common/install_cache.sh install_cache.sh
 ENV PATH /opt/cache/bin:$PATH
@@ -102,8 +83,5 @@ RUN bash ./install_cache.sh && rm install_cache.sh
 ARG BUILD_ENVIRONMENT
 ENV BUILD_ENVIRONMENT ${BUILD_ENVIRONMENT}
 
-# Install LLVM dev version (Defined in the pytorch/builder github repository)
-COPY --from=pytorch/llvm:9.0.1 /opt/llvm /opt/llvm
-
 USER jenkins
 CMD ["bash"]
diff --git a/.ci/docker/ubuntu/Dockerfile b/.ci/docker/ubuntu/Dockerfile
index e719df220dbaf..6c98d6639fb16 100644
--- a/.ci/docker/ubuntu/Dockerfile
+++ b/.ci/docker/ubuntu/Dockerfile
@@ -42,7 +42,6 @@ COPY ./common/install_conda.sh install_conda.sh
 COPY ./common/common_utils.sh common_utils.sh
 COPY ./common/install_magma_conda.sh install_magma_conda.sh
 RUN bash ./install_conda.sh && rm install_conda.sh install_magma_conda.sh common_utils.sh /opt/conda/requirements-ci.txt /opt/conda/requirements-docs.txt
-RUN if [ -n "${UNINSTALL_DILL}" ]; then pip uninstall -y dill; fi
 
 # Install gcc
 ARG GCC_VERSION
@@ -68,36 +67,7 @@ ENV NCCL_INCLUDE_DIR="/usr/local/cuda/include/"
 ENV NCCL_LIB_DIR="/usr/local/cuda/lib64/"
 
 
-# (optional) Install UCC
-ARG UCX_COMMIT
-ARG UCC_COMMIT
 ARG CUDA_VERSION
-ENV UCX_COMMIT $UCX_COMMIT
-ENV UCC_COMMIT $UCC_COMMIT
-ENV UCX_HOME /usr
-ENV UCC_HOME /usr
-ADD ./common/install_ucc.sh install_ucc.sh
-RUN if [ -n "${UCX_COMMIT}" ] && [ -n "${UCC_COMMIT}" ]; then bash ./install_ucc.sh; fi
-RUN rm install_ucc.sh
-
-# (optional) Install vision packages like OpenCV
-ARG VISION
-COPY ./common/install_vision.sh ./common/cache_vision_models.sh ./common/common_utils.sh ./
-RUN if [ -n "${VISION}" ]; then bash ./install_vision.sh; fi
-RUN rm install_vision.sh cache_vision_models.sh common_utils.sh
-ENV INSTALLED_VISION ${VISION}
-
-# (optional) Install non-default Ninja version
-ARG NINJA_VERSION
-COPY ./common/install_ninja.sh install_ninja.sh
-RUN if [ -n "${NINJA_VERSION}" ]; then bash ./install_ninja.sh; fi
-RUN rm install_ninja.sh
-
-COPY ./common/install_openssl.sh install_openssl.sh
-RUN bash ./install_openssl.sh
-ENV OPENSSL_ROOT_DIR /opt/openssl
-ENV OPENSSL_DIR /opt/openssl
-RUN rm install_openssl.sh
 
 ARG INDUCTOR_BENCHMARKS
 COPY ./common/install_inductor_benchmark_deps.sh install_inductor_benchmark_deps.sh
@@ -207,11 +177,6 @@ RUN rm install_openmpi.sh
 ARG BUILD_ENVIRONMENT
 ENV BUILD_ENVIRONMENT ${BUILD_ENVIRONMENT}
 
-# Install LLVM dev version (Defined in the pytorch/builder github repository)
-ARG SKIP_LLVM_SRC_BUILD_INSTALL
-COPY --from=pytorch/llvm:9.0.1 /opt/llvm /opt/llvm
-RUN if [ -n "${SKIP_LLVM_SRC_BUILD_INSTALL}" ]; then set -eu; rm -rf /opt/llvm; fi
-
 # AWS specific CUDA build guidance
 ENV TORCH_NVCC_FLAGS "-Xfatbin -compress-all"
 ENV CUDA_PATH /usr/local/cuda
diff --git a/.ci/libtorch/extract_libtorch_from_wheel.py b/.ci/libtorch/extract_libtorch_from_wheel.py
new file mode 100644
index 0000000000000..423a1b82cb0b5
--- /dev/null
+++ b/.ci/libtorch/extract_libtorch_from_wheel.py
@@ -0,0 +1,334 @@
+#!/usr/bin/env python3
+"""Extract libtorch package from a PyTorch wheel.
+
+Creates a libtorch zip from a pre-built wheel by copying the C++ libraries,
+headers, and CMake files. On Linux, optionally splits debug symbols from
+libtorch_cpu.so into a separate debug zip.
+
+Usage:
+    python extract_libtorch_from_wheel.py \
+        --wheel-dir DIR --output-dir DIR --platform linux|macos|windows
+"""
+
+import argparse
+import glob
+import os
+import re
+import shutil
+import subprocess
+import sys
+import zipfile
+from pathlib import Path
+
+
+def find_wheel(wheel_dir: str) -> Path:
+    wheels = glob.glob(os.path.join(wheel_dir, "*.whl"))
+    if not wheels:
+        raise FileNotFoundError(f"No .whl files found in {wheel_dir}")
+    if len(wheels) > 1:
+        raise RuntimeError(f"Multiple .whl files found in {wheel_dir}: {wheels}")
+    return Path(wheels[0])
+
+
+def parse_version_from_wheel(wheel_path: Path) -> str:
+    # Wheel filename format: {name}-{version}(-{build})?-{python}-{abi}-{platform}.whl
+    name = wheel_path.stem
+    parts = name.split("-")
+    if len(parts) < 3:
+        raise ValueError(f"Cannot parse version from wheel filename: {wheel_path.name}")
+    return parts[1]
+
+
+def extract_wheel(wheel_path: Path, extract_dir: Path) -> Path:
+    with zipfile.ZipFile(wheel_path, "r") as zf:
+        zf.extractall(extract_dir)
+    # Find the torch directory
+    torch_dir = extract_dir / "torch"
+    if not torch_dir.is_dir():
+        raise FileNotFoundError(
+            f"No 'torch' directory found in extracted wheel at {extract_dir}"
+        )
+    return torch_dir
+
+
+def should_exclude_lib(filename: str) -> bool:
+    """Return True for files that should not go into the libtorch package."""
+    if filename.startswith("libtorch_python"):
+        return True
+    if re.match(r"_C\.cpython.*", filename):
+        return True
+    if filename.endswith((".py", ".pyc")):
+        return True
+    if filename == "__init__.py":
+        return True
+    return False
+
+
+def _is_lib_file(name: str, platform: str) -> bool:
+    """Return True if the file looks like a library or header to include."""
+    if platform == "linux":
+        return ".so" in name or name.endswith(".a")
+    elif platform == "macos":
+        return name.endswith((".dylib", ".a"))
+    elif platform == "windows":
+        return name.endswith((".dll", ".lib", ".pdb"))
+    return False
+
+
+def copy_libraries(torch_dir: Path, libtorch_lib: Path, platform: str) -> None:
+    """Copy libraries from torch/lib/ to libtorch/lib/."""
+    torch_lib = torch_dir / "lib"
+    if not torch_lib.is_dir():
+        raise FileNotFoundError(f"torch/lib/ not found at {torch_lib}")
+
+    for item in torch_lib.iterdir():
+        if item.is_dir():
+            # Copy subdirectories (e.g. libshm/) as-is
+            shutil.copytree(item, libtorch_lib / item.name, dirs_exist_ok=True)
+            continue
+        if should_exclude_lib(item.name):
+            continue
+        if _is_lib_file(item.name, platform):
+            shutil.copy2(item, libtorch_lib / item.name)
+
+    # On macOS, also copy delocated dylibs from torch/.dylibs/ if present
+    if platform == "macos":
+        dylibs_dir = torch_dir / ".dylibs"
+        if dylibs_dir.is_dir():
+            for item in dylibs_dir.iterdir():
+                if item.suffix == ".dylib" and not should_exclude_lib(item.name):
+                    shutil.copy2(item, libtorch_lib / item.name)
+
+
+def copy_includes(torch_dir: Path, libtorch_include: Path) -> None:
+    torch_include = torch_dir / "include"
+    if not torch_include.is_dir():
+        # Some older wheels might have include under torch/lib/include
+        torch_include = torch_dir / "lib" / "include"
+    if not torch_include.is_dir():
+        raise FileNotFoundError("include/ not found in torch directory")
+    shutil.copytree(torch_include, libtorch_include, dirs_exist_ok=True)
+
+
+def copy_cmake(torch_dir: Path, libtorch_share: Path) -> None:
+    torch_cmake = torch_dir / "share" / "cmake"
+    if not torch_cmake.is_dir():
+        print(f"Warning: share/cmake/ not found at {torch_cmake}", file=sys.stderr)
+        return
+    cmake_dest = libtorch_share / "cmake"
+    shutil.copytree(torch_cmake, cmake_dest, dirs_exist_ok=True)
+
+
+def copy_bin(torch_dir: Path, libtorch_bin: Path, platform: str) -> None:
+    """Copy binary executables (mainly relevant for Windows)."""
+    if platform == "windows":
+        torch_lib = torch_dir / "lib"
+        if torch_lib.is_dir():
+            for item in torch_lib.iterdir():
+                if item.suffix == ".dll" and not should_exclude_lib(item.name):
+                    shutil.copy2(item, libtorch_bin / item.name)
+
+
+def write_metadata(libtorch_dir: Path, version: str, git_hash: str) -> None:
+    (libtorch_dir / "build-version").write_text(version + "\n")
+    (libtorch_dir / "build-hash").write_text(git_hash + "\n")
+
+
+def get_git_hash(torch_dir: Path) -> str:
+    """Read git_version from the wheel's torch/version.py."""
+    version_file = torch_dir / "version.py"
+    if not version_file.exists():
+        return "unknown"
+    from ast import literal_eval
+
+    for line in version_file.read_text().splitlines():
+        if line.strip().startswith("git_version"):
+            try:
+                return literal_eval(line.partition("=")[2].strip())
+            except Exception:
+                pass
+    return "unknown"
+
+
+def split_debug_symbols(
+    libtorch_dir: Path, output_dir: Path, zip_prefix: str, version: str
+) -> None:
+    """Split debug symbols from libtorch_cpu.so (Linux only)."""
+    libtorch_cpu = libtorch_dir / "lib" / "libtorch_cpu.so"
+    if not libtorch_cpu.exists():
+        print(
+            "Warning: libtorch_cpu.so not found, skipping debug symbol split",
+            file=sys.stderr,
+        )
+        return
+
+    debug_dir = libtorch_dir.parent / "debug"
+    debug_dir.mkdir(exist_ok=True)
+    dbg_file = debug_dir / "libtorch_cpu.so.dbg"
+
+    # Copy to create debug file
+    shutil.copy2(libtorch_cpu, dbg_file)
+
+    # Keep only debug symbols
+    subprocess.run(
+        ["strip", "--only-keep-debug", str(dbg_file)],
+        check=True,
+    )
+
+    # Strip debug info from release lib
+    subprocess.run(
+        ["strip", "--strip-debug", str(libtorch_cpu)],
+        check=True,
+    )
+
+    # Add debug link
+    subprocess.run(
+        ["objcopy", str(libtorch_cpu), f"--add-gnu-debuglink={dbg_file}"],
+        check=True,
+        cwd=str(libtorch_dir / "lib"),
+    )
+
+    # Extract CRC32 from the debug link section
+    try:
+        result = subprocess.run(
+            [
+                "bash",
+                "-c",
+                f"objcopy --dump-section .gnu_debuglink=>(tail -c4 | od -t x4 -An | xargs echo) {libtorch_cpu}",
+            ],
+            capture_output=True,
+            text=True,
+            check=True,
+        )
+        crc32 = result.stdout.strip()
+    except subprocess.CalledProcessError:
+        crc32 = "unknown"
+
+    # Create debug zip
+    debug_zip = output_dir / f"debug-{zip_prefix}-{version}-{crc32}.zip"
+    with zipfile.ZipFile(debug_zip, "w", zipfile.ZIP_DEFLATED) as zf:
+        zf.write(dbg_file, "debug/libtorch_cpu.so.dbg")
+
+    print(f"Debug symbols zip: {debug_zip}")
+
+
+def create_libtorch_zip(
+    libtorch_dir: Path,
+    output_dir: Path,
+    zip_prefix: str,
+    version: str,
+) -> Path:
+    zip_path = output_dir / f"{zip_prefix}-{version}.zip"
+    with zipfile.ZipFile(zip_path, "w", zipfile.ZIP_DEFLATED) as zf:
+        for root, dirs, files in os.walk(libtorch_dir):
+            for f in files:
+                filepath = Path(root) / f
+                arcname = filepath.relative_to(libtorch_dir.parent)
+                zf.write(filepath, arcname)
+    # Create latest symlink
+    latest_zip = output_dir / f"{zip_prefix}-latest.zip"
+    latest_zip.symlink_to(zip_path.name)
+
+    print(f"Libtorch zip: {zip_path}")
+    print(f"Libtorch latest zip: {latest_zip}")
+    return zip_path
+
+
+def compute_zip_prefix(platform: str, desired_cuda: str, libtorch_variant: str) -> str:
+    """Compute the zip filename prefix matching existing naming conventions.
+
+    Linux:  libtorch-shared-with-deps
+    macOS:  libtorch-macos-arm64
+    Windows: libtorch-win-shared-with-deps (or libtorch-win-arm64-shared-with-deps)
+    """
+    if platform == "macos":
+        return "libtorch-macos-arm64"
+    elif platform == "windows":
+        return f"libtorch-win-{libtorch_variant}"
+    else:
+        return f"libtorch-{libtorch_variant}"
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Extract libtorch from PyTorch wheel")
+    parser.add_argument(
+        "--wheel-dir", required=True, help="Directory containing the .whl file"
+    )
+    parser.add_argument(
+        "--output-dir", required=True, help="Directory for output zip files"
+    )
+    parser.add_argument(
+        "--platform",
+        required=True,
+        choices=["linux", "macos", "windows"],
+        help="Target platform",
+    )
+    parser.add_argument(
+        "--desired-cuda",
+        default="cpu",
+        help="CUDA variant (cpu, cu126, cu128, rocm7.1, etc.)",
+    )
+    parser.add_argument(
+        "--libtorch-variant",
+        default="shared-with-deps",
+        help="Libtorch variant (shared-with-deps, etc.)",
+    )
+    parser.add_argument(
+        "--git-hash",
+        default="",
+        help="Git hash to use for build-hash (auto-detected if not set)",
+    )
+    args = parser.parse_args()
+
+    wheel_dir = Path(args.wheel_dir)
+    output_dir = Path(args.output_dir)
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    # Find and extract wheel
+    wheel_path = find_wheel(str(wheel_dir))
+    version = parse_version_from_wheel(wheel_path)
+    print(f"Found wheel: {wheel_path}")
+    print(f"Version: {version}")
+
+    extract_dir = wheel_dir / "_extract_tmp"
+    if extract_dir.exists():
+        shutil.rmtree(extract_dir)
+    extract_dir.mkdir()
+
+    try:
+        torch_dir = extract_wheel(wheel_path, extract_dir)
+
+        # Create libtorch directory structure
+        libtorch_dir = extract_dir / "libtorch"
+        libtorch_dir.mkdir()
+        for subdir in ["lib", "bin", "include", "share"]:
+            (libtorch_dir / subdir).mkdir()
+
+        # Copy components
+        copy_libraries(torch_dir, libtorch_dir / "lib", args.platform)
+        copy_includes(torch_dir, libtorch_dir / "include")
+        copy_cmake(torch_dir, libtorch_dir / "share")
+        copy_bin(torch_dir, libtorch_dir / "bin", args.platform)
+
+        # Write metadata
+        git_hash = args.git_hash or get_git_hash(torch_dir)
+        write_metadata(libtorch_dir, version, git_hash)
+
+        # Compute zip prefix
+        zip_prefix = compute_zip_prefix(
+            args.platform, args.desired_cuda, args.libtorch_variant
+        )
+
+        # Split debug symbols on Linux
+        if args.platform == "linux":
+            split_debug_symbols(libtorch_dir, output_dir, zip_prefix, version)
+
+        # Create the zip
+        create_libtorch_zip(libtorch_dir, output_dir, zip_prefix, version)
+
+    finally:
+        shutil.rmtree(extract_dir)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/.ci/lumen_cli/cli/build_cli/register_build.py b/.ci/lumen_cli/cli/build_cli/register_build.py
index 9f35a9c8165dc..277abce662bcd 100644
--- a/.ci/lumen_cli/cli/build_cli/register_build.py
+++ b/.ci/lumen_cli/cli/build_cli/register_build.py
@@ -1,4 +1,6 @@
-import argparse
+from __future__ import annotations
+
+import argparse  # noqa: TC003
 import logging
 
 from cli.lib.common.cli_helper import register_targets, RichHelp, TargetSpec
diff --git a/.ci/lumen_cli/cli/lib/common/cli_helper.py b/.ci/lumen_cli/cli/lib/common/cli_helper.py
index 4086eb7d46e81..b19ad297ca1f0 100644
--- a/.ci/lumen_cli/cli/lib/common/cli_helper.py
+++ b/.ci/lumen_cli/cli/lib/common/cli_helper.py
@@ -3,6 +3,8 @@
 
 """
 
+from __future__ import annotations
+
 import argparse
 from abc import ABC, abstractmethod
 
@@ -11,7 +13,7 @@
     from collections.abc import Callable  # Python 3.11+
     from typing import Any, Required, TypedDict
 except ImportError:
-    from collections.abc import Callable
+    from collections.abc import Callable  # noqa: TC003
     from typing import Any, TypedDict
 
     from typing_extensions import Required  # Fallback for Python <3.11
diff --git a/.ci/lumen_cli/cli/lib/common/docker_helper.py b/.ci/lumen_cli/cli/lib/common/docker_helper.py
index b5f0a90e2d47a..1c96ab37e9219 100644
--- a/.ci/lumen_cli/cli/lib/common/docker_helper.py
+++ b/.ci/lumen_cli/cli/lib/common/docker_helper.py
@@ -2,8 +2,9 @@
 Docker Utility helpers for CLI tasks.
 """
 
+from __future__ import annotations
+
 import logging
-from typing import Optional
 
 import docker
 from docker.errors import APIError, NotFound
@@ -12,7 +13,7 @@
 logger = logging.getLogger(__name__)
 
 # lazy singleton so we don't reconnect every call
-_docker_client: Optional[docker.DockerClient] = None
+_docker_client: docker.DockerClient | None = None
 
 
 def _get_client() -> docker.DockerClient:
@@ -23,7 +24,7 @@ def _get_client() -> docker.DockerClient:
 
 
 def local_image_exists(
-    image_name: str, client: Optional[docker.DockerClient] = None
+    image_name: str, client: docker.DockerClient | None = None
 ) -> bool:
     """Return True if a local Docker image exists."""
     if not image_name:
diff --git a/.ci/lumen_cli/cli/lib/common/envs_helper.py b/.ci/lumen_cli/cli/lib/common/envs_helper.py
index a654e7f18ed9f..06c33163711a1 100644
--- a/.ci/lumen_cli/cli/lib/common/envs_helper.py
+++ b/.ci/lumen_cli/cli/lib/common/envs_helper.py
@@ -2,11 +2,12 @@
 Environment Variables and Dataclasses Utility helpers for CLI tasks.
 """
 
+from __future__ import annotations
+
 import os
 from dataclasses import field, fields, is_dataclass, MISSING
 from pathlib import Path
 from textwrap import indent
-from typing import Optional, Union
 
 from cli.lib.common.utils import str2bool
 
@@ -18,9 +19,9 @@ def get_env(name: str, default: str = "") -> str:
 
 def env_path_optional(
     name: str,
-    default: Optional[Union[str, Path]] = None,
+    default: str | Path | None = None,
     resolve: bool = True,
-) -> Optional[Path]:
+) -> Path | None:
     """Get environment variable as optional Path."""
     val = get_env(name) or default
     if not val:
@@ -32,7 +33,7 @@ def env_path_optional(
 
 def env_path(
     name: str,
-    default: Optional[Union[str, Path]] = None,
+    default: str | Path | None = None,
     resolve: bool = True,
 ) -> Path:
     """Get environment variable as Path, raise if missing."""
@@ -61,7 +62,7 @@ def env_bool_field(
 
 def env_path_field(
     name: str,
-    default: Union[str, Path] = "",
+    default: str | Path = "",
     *,
     resolve: bool = True,
 ) -> Path:
diff --git a/.ci/lumen_cli/cli/lib/common/path_helper.py b/.ci/lumen_cli/cli/lib/common/path_helper.py
index 4f74aa6e509de..3102d66ff7131 100644
--- a/.ci/lumen_cli/cli/lib/common/path_helper.py
+++ b/.ci/lumen_cli/cli/lib/common/path_helper.py
@@ -1,15 +1,16 @@
 """Path utility helpers for CLI tasks."""
 
+from __future__ import annotations
+
 import logging
 import shutil
 from pathlib import Path
-from typing import Union
 
 
 logger = logging.getLogger(__name__)
 
 
-def get_path(path: Union[str, Path], resolve: bool = False) -> Path:
+def get_path(path: str | Path, resolve: bool = False) -> Path:
     """Convert to Path object, optionally resolving to absolute path."""
     if not path:
         raise ValueError("Path cannot be None or empty")
@@ -17,14 +18,14 @@ def get_path(path: Union[str, Path], resolve: bool = False) -> Path:
     return result.resolve() if resolve else result
 
 
-def ensure_dir_exists(path: Union[str, Path]) -> Path:
+def ensure_dir_exists(path: str | Path) -> Path:
     """Create directory if it doesn't exist."""
     path_obj = get_path(path)
     path_obj.mkdir(parents=True, exist_ok=True)
     return path_obj
 
 
-def remove_dir(path: Union[str, Path, None]) -> None:
+def remove_dir(path: str | Path | None) -> None:
     """Remove directory if it exists."""
     if not path:
         return
@@ -33,13 +34,13 @@ def remove_dir(path: Union[str, Path, None]) -> None:
         shutil.rmtree(path_obj)
 
 
-def force_create_dir(path: Union[str, Path]) -> Path:
+def force_create_dir(path: str | Path) -> Path:
     """Remove directory if exists, then create fresh empty directory."""
     remove_dir(path)
     return ensure_dir_exists(path)
 
 
-def copy(src: Union[str, Path], dst: Union[str, Path]) -> None:
+def copy(src: str | Path, dst: str | Path) -> None:
     """Copy file or directory from src to dst."""
     src_path = get_path(src, resolve=True)
     dst_path = get_path(dst, resolve=True)
@@ -57,6 +58,6 @@ def copy(src: Union[str, Path], dst: Union[str, Path]) -> None:
         raise ValueError(f"Unsupported path type: {src_path}")
 
 
-def is_path_exist(path: Union[str, Path, None]) -> bool:
+def is_path_exist(path: str | Path | None) -> bool:
     """Check if path exists."""
     return bool(path and get_path(path).exists())
diff --git a/.ci/lumen_cli/cli/lib/common/pip_helper.py b/.ci/lumen_cli/cli/lib/common/pip_helper.py
index a53747e24d256..a0cb1e17840e1 100644
--- a/.ci/lumen_cli/cli/lib/common/pip_helper.py
+++ b/.ci/lumen_cli/cli/lib/common/pip_helper.py
@@ -1,11 +1,12 @@
+from __future__ import annotations
+
 import glob
 import logging
 import shlex
 import shutil
 import sys
-from collections.abc import Iterable
+from collections.abc import Iterable  # noqa: TC003
 from importlib.metadata import PackageNotFoundError, version  # noqa: UP035
-from typing import Optional, Union
 
 from cli.lib.common.utils import run_command
 
@@ -17,8 +18,8 @@ def pip_install_packages(
     packages: Iterable[str] = (),
     env=None,
     *,
-    requirements: Optional[str] = None,
-    constraints: Optional[str] = None,
+    requirements: str | None = None,
+    constraints: str | None = None,
     prefer_uv: bool = False,
 ) -> None:
     use_uv = prefer_uv and shutil.which("uv") is not None
@@ -37,14 +38,14 @@ def pip_install_packages(
     run_command(" ".join(map(shlex.quote, cmd)), env=env)
 
 
-def pip_install_first_match(pattern: str, extras: Optional[str] = None, pref_uv=False):
+def pip_install_first_match(pattern: str, extras: str | None = None, pref_uv=False):
     wheel = first_matching_pkg(pattern)
     target = f"{wheel}[{extras}]" if extras else wheel
     logger.info("Installing %s...", target)
     pip_install_packages([target], prefer_uv=pref_uv)
 
 
-def run_python(args: Union[str, list[str]], env=None):
+def run_python(args: str | list[str], env=None):
     """
     Run the python in the current environment.
     """
diff --git a/.ci/lumen_cli/cli/lib/common/utils.py b/.ci/lumen_cli/cli/lib/common/utils.py
index b03309810d986..0798d7d1d369b 100644
--- a/.ci/lumen_cli/cli/lib/common/utils.py
+++ b/.ci/lumen_cli/cli/lib/common/utils.py
@@ -2,6 +2,8 @@
 General Utility helpers for CLI tasks.
 """
 
+from __future__ import annotations
+
 import logging
 import os
 import shlex
@@ -9,7 +11,6 @@
 import sys
 from contextlib import contextmanager
 from pathlib import Path
-from typing import Optional
 
 
 logger = logging.getLogger(__name__)
@@ -19,8 +20,8 @@ def run_command(
     cmd: str,
     use_shell: bool = False,
     log_cmd: bool = True,
-    cwd: Optional[str] = None,
-    env: Optional[dict] = None,
+    cwd: str | None = None,
+    env: dict | None = None,
     check: bool = True,
 ) -> int:
     """Run a command with optional shell execution."""
@@ -61,7 +62,7 @@ def run_command(
     return proc.returncode
 
 
-def str2bool(value: Optional[str]) -> bool:
+def str2bool(value: str | None) -> bool:
     """Convert environment variables to boolean values."""
     if not value:
         return False
@@ -120,7 +121,7 @@ def working_directory(path: str):
 
 def get_wheels(
     output_dir: Path,
-    max_depth: Optional[int] = None,
+    max_depth: int | None = None,
 ) -> list[str]:
     """Return a list of wheels found in the given output directory."""
     root = Path(output_dir)
diff --git a/.ci/lumen_cli/cli/lib/core/torchtitan/lib.py b/.ci/lumen_cli/cli/lib/core/torchtitan/lib.py
new file mode 100644
index 0000000000000..3b4a818ff04ac
--- /dev/null
+++ b/.ci/lumen_cli/cli/lib/core/torchtitan/lib.py
@@ -0,0 +1,65 @@
+import logging
+from pathlib import Path
+from typing import Any
+
+import yaml
+from cli.lib.common.git_helper import clone_external_repo
+from cli.lib.common.utils import run_command, temp_environ, working_directory
+
+
+logger = logging.getLogger(__name__)
+
+_TORCHTITAN_TEST_LIBRARY_PATH = Path(__file__).parent / "torchtitan_test_library.yaml"
+
+
+def _load_torchtitan_test_library_yaml() -> dict[str, Any]:
+    if not _TORCHTITAN_TEST_LIBRARY_PATH.exists():
+        raise FileNotFoundError(
+            f"torchtitan test library YAML not found: {_TORCHTITAN_TEST_LIBRARY_PATH}"
+        )
+    with open(_TORCHTITAN_TEST_LIBRARY_PATH, encoding="utf-8") as f:
+        return yaml.safe_load(f)
+
+
+def load_torchtitan_test_library() -> dict[str, Any]:
+    return _load_torchtitan_test_library_yaml()
+
+
+def clone_torchtitan(dst: str = "torchtitan"):
+    _, commit = clone_external_repo(
+        target="torchtitan",
+        repo="https://github.com/pytorch/torchtitan.git",
+        dst=dst,
+    )
+    return commit
+
+
+def run_test_plan(
+    test_plan: str,
+    tests_map: dict[str, Any],
+):
+    logger.info("Running torchtitan test plan: %s", test_plan)
+    if test_plan not in tests_map:
+        raise RuntimeError(
+            f"test plan '{test_plan}' not found in torchtitan test library"
+        )
+
+    tests = tests_map[test_plan]
+    title = tests.get("title", "unknown test")
+    logger.info("Running tests: %s", title)
+
+    with (
+        working_directory(tests.get("working_directory", "")),
+        temp_environ(tests.get("env_vars", {})),
+    ):
+        failures = []
+        for step in tests["steps"]:
+            logger.info("Running step: %s", step)
+            code = run_command(cmd=step, check=False, use_shell=True)
+            if code != 0:
+                failures.append(step)
+            logger.info("Finished step: %s", step)
+        if failures:
+            logger.error("Failed steps: %s", failures)
+            raise RuntimeError(f"{len(failures)} test steps failed: {failures}")
+        logger.info("All tests passed for plan: %s", test_plan)
diff --git a/.ci/lumen_cli/cli/lib/core/torchtitan/torchtitan_test.py b/.ci/lumen_cli/cli/lib/core/torchtitan/torchtitan_test.py
new file mode 100644
index 0000000000000..d42fc93d6f649
--- /dev/null
+++ b/.ci/lumen_cli/cli/lib/core/torchtitan/torchtitan_test.py
@@ -0,0 +1,41 @@
+import logging
+from typing import Any
+
+from cli.lib.common.cli_helper import BaseRunner
+from cli.lib.common.pip_helper import pip_install_packages
+from cli.lib.common.utils import working_directory
+from cli.lib.core.torchtitan.lib import (
+    clone_torchtitan,
+    load_torchtitan_test_library,
+    run_test_plan,
+)
+
+
+logger = logging.getLogger(__name__)
+
+
+class TorchtitanTestRunner(BaseRunner):
+    def __init__(self, args: Any):
+        self.work_directory = "torchtitan"
+        self.test_plan = args.test_plan
+
+    def prepare(self):
+        clone_torchtitan(dst=self.work_directory)
+        # torchao and torchcomms nightlies are required by torchtitan
+        pip_install_packages(
+            packages=[
+                "--pre",
+                "torchao",
+                "torchcomms",
+                "--index-url",
+                "https://download.pytorch.org/whl/nightly/cu129",
+            ],
+        )
+        with working_directory(self.work_directory):
+            pip_install_packages(packages=["-e", "."])
+            pip_install_packages(packages=["pytest", "pytest-cov"])
+
+    def run(self):
+        self.prepare()
+        with working_directory(self.work_directory):
+            run_test_plan(self.test_plan, load_torchtitan_test_library())
diff --git a/.ci/lumen_cli/cli/lib/core/torchtitan/torchtitan_test_library.yaml b/.ci/lumen_cli/cli/lib/core/torchtitan/torchtitan_test_library.yaml
new file mode 100644
index 0000000000000..eba45310855c3
--- /dev/null
+++ b/.ci/lumen_cli/cli/lib/core/torchtitan/torchtitan_test_library.yaml
@@ -0,0 +1,21 @@
+# torchtitan Test Library Configuration
+# Each test plan maps to torchtitan's own test runners.
+# When tests are added/removed in torchtitan, the daily pin bump picks them up.
+# Test filtering and exclusions are managed in the torchtitan repo via
+# scripts/ci/pytorch_ci_test_runner.sh.
+
+torchtitan_features_integration:
+  title: torchtitan Feature Integration Tests (8 GPU)
+  id: torchtitan_features_integration
+  env_vars:
+    NGPU: "8"
+  steps:
+    - scripts/ci/pytorch_ci_test_runner.sh feature_tests
+
+torchtitan_models_integration:
+  title: torchtitan Model Integration Tests (8 GPU)
+  id: torchtitan_models_integration
+  env_vars:
+    NGPU: "8"
+  steps:
+    - scripts/ci/pytorch_ci_test_runner.sh model_tests
diff --git a/.ci/lumen_cli/cli/lib/core/vllm/disabled_vllm_tests.yaml b/.ci/lumen_cli/cli/lib/core/vllm/disabled_vllm_tests.yaml
new file mode 100644
index 0000000000000..a350ced6be862
--- /dev/null
+++ b/.ci/lumen_cli/cli/lib/core/vllm/disabled_vllm_tests.yaml
@@ -0,0 +1,21 @@
+# Disabled vLLM tests for PyTorch CI.
+# Node IDs are relative to vLLM's tests/ directory.
+# Each entry requires 'test' (node ID) and 'issue' (tracking URL).
+#
+# To disable a test, copy one of the entries below into the disabled_tests list.
+#
+# Disable an entire test file (produces --ignore):
+#   - test: basic_correctness/test_basic_correctness.py
+#     issue: https://github.com/pytorch/pytorch/issues/12345
+#
+# Disable a single test (produces --deselect):
+#   - test: basic_correctness/test_basic_correctness.py::test_something
+#     issue: https://github.com/pytorch/pytorch/issues/12345
+#
+# Disable only in specific test plans via 'configs':
+#   - test: compile/test_fusion.py
+#     issue: https://github.com/pytorch/pytorch/issues/12345
+#     configs:
+#       - vllm_pytorch_compilation_unit_tests
+
+disabled_tests: []
diff --git a/.ci/lumen_cli/cli/lib/core/vllm/lib.py b/.ci/lumen_cli/cli/lib/core/vllm/lib.py
index 2532a97ebf345..f197080857f72 100644
--- a/.ci/lumen_cli/cli/lib/core/vllm/lib.py
+++ b/.ci/lumen_cli/cli/lib/core/vllm/lib.py
@@ -1,6 +1,11 @@
+from __future__ import annotations
+
+import json
 import logging
 import os
+import re
 import textwrap
+import urllib.request
 from pathlib import Path
 from typing import Any
 
@@ -19,6 +24,8 @@
 logger = logging.getLogger(__name__)
 
 _VLLM_TEST_LIBRARY_PATH = Path(__file__).parent / "vllm_test_library.yaml"
+_DISABLED_VLLM_TESTS_PATH = Path(__file__).parent / "disabled_vllm_tests.yaml"
+_DISABLED_VLLM_TESTS_ISSUE = 175899
 
 
 def _load_vllm_test_library_yaml() -> dict[str, Any]:
@@ -87,6 +94,92 @@ def check_parallelism(tests: Any, title: str, shard_id: int = 0, num_shards: int
     return True
 
 
+def _load_disabled_vllm_tests_from_yaml() -> list[dict[str, Any]]:
+    if not _DISABLED_VLLM_TESTS_PATH.exists():
+        return []
+    with open(_DISABLED_VLLM_TESTS_PATH, encoding="utf-8") as f:
+        data = yaml.safe_load(f)
+    if not data or "disabled_tests" not in data:
+        return []
+    entries = data["disabled_tests"]
+    if not entries:
+        return []
+    for entry in entries:
+        if "test" not in entry or "issue" not in entry:
+            raise ValueError(
+                f"disabled_vllm_tests.yaml: each entry must have 'test' and 'issue' keys, got {entry}"
+            )
+    return entries
+
+
+def _parse_disabled_tests_from_issue_body(body: str) -> list[dict[str, Any]]:
+    match = re.search(r"```yaml\s*\n(.*?)```", body, re.DOTALL)
+    if not match:
+        return []
+    block = match.group(1)
+    data = yaml.safe_load(block)
+    if not data or "disabled_tests" not in data:
+        return []
+    return data["disabled_tests"] or []
+
+
+def _load_disabled_vllm_tests_from_github() -> list[dict[str, Any]]:
+    if not _DISABLED_VLLM_TESTS_ISSUE:
+        return []
+    url = f"https://api.github.com/repos/pytorch/pytorch/issues/{_DISABLED_VLLM_TESTS_ISSUE}"
+    headers = {"Accept": "application/vnd.github.v3+json"}
+    token = os.environ.get("GITHUB_TOKEN")
+    if token:
+        headers["Authorization"] = f"token {token}"
+    try:
+        req = urllib.request.Request(url, headers=headers)
+        with urllib.request.urlopen(req, timeout=30) as resp:
+            issue = json.loads(resp.read())
+        body = issue.get("body", "") or ""
+        entries = _parse_disabled_tests_from_issue_body(body)
+        # Filter out malformed entries — the issue body is user-editable
+        entries = [e for e in entries if "test" in e]
+        issue_url = issue.get("html_url", url)
+        for entry in entries:
+            entry.setdefault("issue", issue_url)
+        return entries
+    except Exception:
+        logger.warning(
+            "Failed to fetch disabled vLLM tests from GitHub issue #%d",
+            _DISABLED_VLLM_TESTS_ISSUE,
+            exc_info=True,
+        )
+        return []
+
+
+def _load_disabled_vllm_tests() -> list[dict[str, Any]]:
+    yaml_entries = _load_disabled_vllm_tests_from_yaml()
+    github_entries = _load_disabled_vllm_tests_from_github()
+    seen = {e["test"] for e in yaml_entries}
+    merged = list(yaml_entries)
+    for entry in github_entries:
+        if entry["test"] not in seen:
+            seen.add(entry["test"])
+            merged.append(entry)
+    return merged
+
+
+def _build_disabled_test_flags(
+    disabled_tests: list[dict[str, Any]], test_plan: str
+) -> str:
+    flags = []
+    for entry in disabled_tests:
+        configs = entry.get("configs")
+        if configs and test_plan not in configs:
+            continue
+        node_id = entry["test"]
+        if "::" in node_id:
+            flags.append(f"--deselect={node_id}")
+        else:
+            flags.append(f"--ignore={node_id}")
+    return " ".join(flags)
+
+
 def run_test_plan(
     test_plan: str,
     test_target: str,
@@ -110,6 +203,11 @@ def run_test_plan(
     if is_parallel:
         title = title.replace("%N", f"{shard_id}/{num_shards}")
 
+    disabled_tests = _load_disabled_vllm_tests()
+    disabled_flags = _build_disabled_test_flags(disabled_tests, test_plan)
+    if disabled_flags:
+        logger.info("Disabled test flags for %s: %s", test_plan, disabled_flags)
+
     logger.info("Running tests: %s", title)
     if pkgs:
         logger.info("Installing packages: %s", pkgs)
@@ -124,11 +222,14 @@ def run_test_plan(
             if is_parallel:
                 step = replace_buildkite_placeholders(step, shard_id, num_shards)
                 logger.info("Running parallel step: %s", step)
-            # Support retry with delay for all pytest commands, pytest-rerunfailures
-            # is already a dependency of vLLM. This is needed as a stop gap to reduce
-            # the number of requests to HF until #172300 can be landed to enable
-            # HF offline mode
             if "pytest" in step:
+                # Inject disabled test flags before rerun flags
+                if disabled_flags:
+                    step = step.replace("pytest", f"pytest {disabled_flags}", 1)
+                # Support retry with delay for all pytest commands, pytest-rerunfailures
+                # is already a dependency of vLLM. This is needed as a stop gap to reduce
+                # the number of requests to HF until #172300 can be landed to enable
+                # HF offline mode.
                 # Use a low retry count and a high delay value to lower the risk of
                 # having a retry storm and make thing worse
                 rerun_count = os.getenv(
diff --git a/.ci/lumen_cli/cli/lib/core/vllm/vllm_build.py b/.ci/lumen_cli/cli/lib/core/vllm/vllm_build.py
index 63e5f7a28de54..f0780f98541f3 100644
--- a/.ci/lumen_cli/cli/lib/core/vllm/vllm_build.py
+++ b/.ci/lumen_cli/cli/lib/core/vllm/vllm_build.py
@@ -1,9 +1,10 @@
+from __future__ import annotations
+
 import logging
 import os
 import textwrap
 from dataclasses import dataclass
 from pathlib import Path
-from typing import Optional
 
 from cli.lib.common.cli_helper import BaseRunner
 from cli.lib.common.docker_helper import local_image_exists
@@ -235,7 +236,7 @@ def get_result_path(self, path):
         abs_path = get_path(path, resolve=True)
         return abs_path
 
-    def _get_torch_wheel_path_arg(self, torch_whl_dir: Optional[Path]) -> str:
+    def _get_torch_wheel_path_arg(self, torch_whl_dir: Path | None) -> str:
         if not torch_whl_dir:
             return ""
         return f"--build-arg TORCH_WHEELS_PATH={_VLLM_TEMP_FOLDER}"
diff --git a/.ci/lumen_cli/cli/lib/core/vllm/vllm_test_library.yaml b/.ci/lumen_cli/cli/lib/core/vllm/vllm_test_library.yaml
index f2f450b6f9004..402f2d8bf0e69 100644
--- a/.ci/lumen_cli/cli/lib/core/vllm/vllm_test_library.yaml
+++ b/.ci/lumen_cli/cli/lib/core/vllm/vllm_test_library.yaml
@@ -20,7 +20,7 @@ vllm_basic_models_test:
     - pytest -v -s models/test_registry.py
     - pytest -v -s models/test_utils.py
     - pytest -v -s models/test_vision.py
-    - pytest -v -s models/test_initialization.py
+    - HF_DATASETS_OFFLINE=0 TRANSFORMERS_OFFLINE=0 pytest -v -s models/test_initialization.py
 
 vllm_entrypoints_test:
   title: Entrypoints Test
@@ -60,7 +60,7 @@ vllm_distributed_test_28_failure_test:
     VLLM_WORKER_MULTIPROC_METHOD: spawn
   num_gpus: 4
   steps:
-    - pytest -v -s distributed/test_sequence_parallel.py
+    - pytest -v -s compile/correctness_e2e/test_sequence_parallel.py
 
 vllm_lora_28_failure_test:
   title: LoRA pytorch 2.8 failure test
@@ -85,29 +85,28 @@ vllm_multi_model_test_28_failure_test:
   package_install:
     - git+https://github.com/TIGER-AI-Lab/Mantis.git
   steps:
-    - pytest -v -s models/multimodal/generation/test_voxtral.py -k 'not 5-128-half'
-    - HF_DATASETS_OFFLINE=0 TRANSFORMERS_OFFLINE=0 pytest -v -s models/multimodal/generation/test_voxtral.py -k 5-128-half
+    - HF_DATASETS_OFFLINE=0 TRANSFORMERS_OFFLINE=0 pytest -v -s models/multimodal/generation/test_voxtral.py
     - pytest -v -s models/multimodal/pooling
 
 vllm_pytorch_compilation_unit_tests:
   title: PyTorch Compilation Unit Tests
   id: vllm_pytorch_compilation_unit_tests
   steps:
-    - pytest -v -s compile/test_pass_manager.py
-    - pytest -v -s compile/test_fusion.py
-    - pytest -v -s compile/test_fusion_attn.py
-    - pytest -v -s compile/test_silu_mul_quant_fusion.py
-    - pytest -v -s compile/distributed/test_sequence_parallelism.py
-    - pytest -v -s compile/distributed/test_async_tp.py
-    - pytest -v -s compile/distributed/test_fusion_all_reduce.py
+    - pytest -v -s compile/passes/test_pass_manager.py
+    - pytest -v -s compile/passes/test_fusion.py
+    - pytest -v -s compile/passes/test_fusion_attn.py
+    - pytest -v -s compile/passes/test_silu_mul_quant_fusion.py
+    - pytest -v -s compile/passes/distributed/test_sequence_parallelism.py
+    - pytest -v -s compile/passes/distributed/test_async_tp.py
+    - pytest -v -s compile/passes/distributed/test_fusion_all_reduce.py
     - pytest -v -s compile/test_decorator.py
 
 vllm_language_model_test_extended_generation_28_failure_test:
   title: Language Models Test (Extended Generation) 2.8 release failure
-  id: vllm_languagde_model_test_extended_generation_28_failure_test
+  id: vllm_language_model_test_extended_generation_28_failure_test
   package_install:
     - --no-build-isolation
-    - git+https://github.com/Dao-AILab/causal-conv1d@v1.5.0.post8
+    - git+https://github.com/Dao-AILab/causal-conv1d@v1.6.0
   steps:
     - pytest -v -s models/language/generation/test_mistral.py
 
@@ -118,7 +117,7 @@ vllm_distributed_test_2_gpu_28_failure_test:
     VLLM_WORKER_MULTIPROC_METHOD: spawn
   num_gpus: 4
   steps:
-    - pytest -v -s distributed/test_sequence_parallel.py
+    - pytest -v -s compile/correctness_e2e/test_sequence_parallel.py
 
 vllm_lora_test:
   title: LoRA Test %N
diff --git a/.ci/lumen_cli/cli/test_cli/register_test.py b/.ci/lumen_cli/cli/test_cli/register_test.py
index 2973341b83ed2..12088da4f2ef9 100644
--- a/.ci/lumen_cli/cli/test_cli/register_test.py
+++ b/.ci/lumen_cli/cli/test_cli/register_test.py
@@ -1,7 +1,10 @@
-import argparse
+from __future__ import annotations
+
+import argparse  # noqa: TC003
 import logging
 
 from cli.lib.common.cli_helper import register_targets, RichHelp, TargetSpec
+from cli.lib.core.torchtitan.torchtitan_test import TorchtitanTestRunner
 from cli.lib.core.vllm.vllm_test import VllmTestRunner
 
 
@@ -13,8 +16,11 @@
     "vllm": {
         "runner": VllmTestRunner,
         "help": "test vLLM with pytorch main",
-    }
-    # add yours ...
+    },
+    "torchtitan": {
+        "runner": TorchtitanTestRunner,
+        "help": "test torchtitan with pytorch main",
+    },
 }
 
 
diff --git a/.ci/lumen_cli/pyproject.toml b/.ci/lumen_cli/pyproject.toml
index b2ac379e34ab0..ce8cf59d99fda 100644
--- a/.ci/lumen_cli/pyproject.toml
+++ b/.ci/lumen_cli/pyproject.toml
@@ -6,7 +6,7 @@ dependencies = [
     "GitPython==3.1.45",
     "docker==7.1.0",
     "pytest==7.3.2",
-    "uv==0.9.6"
+    "uv==0.11.6"
 ]
 
 [tool.setuptools]
diff --git a/.ci/lumen_cli/tests/test_cli_helper.py b/.ci/lumen_cli/tests/test_cli_helper.py
index 848f22d6be200..b839c8ddf62a6 100644
--- a/.ci/lumen_cli/tests/test_cli_helper.py
+++ b/.ci/lumen_cli/tests/test_cli_helper.py
@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import argparse
 import io
 import unittest
diff --git a/.ci/lumen_cli/tests/test_run_plan.py b/.ci/lumen_cli/tests/test_run_plan.py
index 1ad3433cd6661..13d2695d21c0c 100644
--- a/.ci/lumen_cli/tests/test_run_plan.py
+++ b/.ci/lumen_cli/tests/test_run_plan.py
@@ -1,8 +1,11 @@
 # tests/test_run_test_plan.py
+from __future__ import annotations
+
 import importlib
+import json
 from contextlib import nullcontext
 from types import SimpleNamespace
-from unittest.mock import MagicMock
+from unittest.mock import MagicMock, patch as mock_patch
 
 import pytest
 
@@ -29,13 +32,15 @@ def _get_check(c):
 def patch_module(monkeypatch):
     """
     Patch helpers ('pip_install_packages', 'temp_environ', 'working_directory',
-    'run_command', 'logger') inside the target module and expose them.
+    'run_command', 'logger', '_load_disabled_vllm_tests') inside the target
+    module and expose them.
     """
     module = importlib.import_module(MOD)
 
     # Create fakes/mocks
     pip_install_packages = MagicMock(name="pip_install_packages")
     run_command = MagicMock(name="run_command", return_value=0)
+    disabled_mock = MagicMock(name="_load_disabled_vllm_tests", return_value=[])
 
     # temp_environ / working_directory: record calls but act as context managers
     temp_calls: list[dict] = []
@@ -51,6 +56,7 @@ def fake_temp_env(map: dict[str, str]):
 
     logger = SimpleNamespace(
         info=MagicMock(name="logger.info"),
+        warning=MagicMock(name="logger.warning"),
         error=MagicMock(name="logger.error"),
     )
 
@@ -64,6 +70,9 @@ def fake_temp_env(map: dict[str, str]):
     )
     monkeypatch.setattr(module, "temp_environ", fake_temp_env, raising=True)
     monkeypatch.setattr(module, "logger", logger, raising=True)
+    monkeypatch.setattr(
+        module, "_load_disabled_vllm_tests", disabled_mock, raising=True
+    )
 
     return SimpleNamespace(
         module=module,
@@ -73,6 +82,7 @@ def fake_temp_env(map: dict[str, str]):
         temp_calls=temp_calls,
         workdir_calls=workdir_calls,
         logger=logger,
+        disabled_mock=disabled_mock,
     )
 
 
@@ -101,12 +111,10 @@ def test_success_runs_all_steps_and_uses_env_and_workdir(monkeypatch, patch_modu
     cmds = [_get_cmd(c) for c in calls]
     checks = [_get_check(c) for c in calls]
 
-    expected_cmds = [
-        "export A=x && pytest -q",
-        "export B=y && pytest -q tests/unit",
-    ]
-    if cmds != expected_cmds:
-        raise AssertionError(f"Expected cmds={expected_cmds}, got cmds={cmds}")
+    if len(cmds) != 2:
+        raise AssertionError(f"Expected 2 commands, got {len(cmds)}: {cmds}")
+    if "pytest" not in cmds[0] or "pytest" not in cmds[1]:
+        raise AssertionError(f"Expected pytest in both commands, got {cmds}")
     if not all(chk is False for chk in checks):
         raise AssertionError(f"Expected all checks to be False, got checks={checks}")
 
@@ -202,3 +210,313 @@ def test_custom_working_directory_used(patch_module):
         raise AssertionError(
             f"Expected workdir_calls=['examples/ci'], got {patch_module.workdir_calls}"
         )
+
+
+# -- Disabled vLLM test injection (integration tests) -------------------------
+
+_SIMPLE_TESTS_MAP = {
+    "plan_a": {
+        "title": "Plan A",
+        "steps": ["pytest -v -s test_foo.py", "pytest -v -s test_bar.py"],
+    }
+}
+
+
+def _cmds(patch_module):
+    return [_get_cmd(c) for c in patch_module.run_command.call_args_list]
+
+
+def test_deselect_injected_for_test_level(patch_module):
+    """Test-level node IDs (with ::) produce --deselect."""
+    patch_module.disabled_mock.return_value = [
+        {
+            "test": "test_foo.py::test_x",
+            "issue": "https://github.com/pytorch/pytorch/issues/175899",
+        }
+    ]
+    patch_module.run_test_plan("plan_a", "cpu", _SIMPLE_TESTS_MAP)
+    cmds = _cmds(patch_module)
+    if not any("--deselect=test_foo.py::test_x" in c for c in cmds):
+        raise AssertionError(
+            f"Expected --deselect=test_foo.py::test_x in commands, got {cmds}"
+        )
+
+
+def test_ignore_injected_for_file_level(patch_module):
+    """File-level node IDs (no ::) produce --ignore."""
+    patch_module.disabled_mock.return_value = [
+        {
+            "test": "test_foo.py",
+            "issue": "https://github.com/pytorch/pytorch/issues/175899",
+        }
+    ]
+    patch_module.run_test_plan("plan_a", "cpu", _SIMPLE_TESTS_MAP)
+    cmds = _cmds(patch_module)
+    if not any("--ignore=test_foo.py" in c for c in cmds):
+        raise AssertionError(f"Expected --ignore=test_foo.py in commands, got {cmds}")
+
+
+def test_empty_disabled_list_no_modification(patch_module):
+    """No flags injected when disabled list is empty."""
+    patch_module.run_test_plan("plan_a", "cpu", _SIMPLE_TESTS_MAP)
+    for cmd in _cmds(patch_module):
+        if "--ignore" in cmd or "--deselect" in cmd:
+            raise AssertionError(f"Expected no --ignore/--deselect flags, got {cmd}")
+
+
+def test_non_pytest_steps_not_modified(patch_module):
+    """Only pytest steps get disabled flags; other commands are left alone."""
+    tests_map = {
+        "mixed_steps": {
+            "title": "Mixed",
+            "steps": ["echo hello", "pytest -v test.py"],
+        }
+    }
+    patch_module.disabled_mock.return_value = [
+        {
+            "test": "test.py::test_z",
+            "issue": "https://github.com/pytorch/pytorch/issues/175899",
+        }
+    ]
+    patch_module.run_test_plan("mixed_steps", "cpu", tests_map)
+    cmds = _cmds(patch_module)
+    if "--deselect" in cmds[0]:
+        raise AssertionError(
+            f"Non-pytest step should not have --deselect, got {cmds[0]}"
+        )
+    if "--deselect=test.py::test_z" not in cmds[1]:
+        raise AssertionError(
+            f"Expected --deselect=test.py::test_z in pytest step, got {cmds[1]}"
+        )
+
+
+def test_disabled_and_rerun_flags_both_present(patch_module):
+    """Disabled flags and rerun flags compose into a single pytest invocation."""
+    patch_module.disabled_mock.return_value = [
+        {
+            "test": "skip.py::test_x",
+            "issue": "https://github.com/pytorch/pytorch/issues/175899",
+        }
+    ]
+    patch_module.run_test_plan("plan_a", "cpu", _SIMPLE_TESTS_MAP)
+    for cmd in _cmds(patch_module):
+        if "--deselect=skip.py::test_x" not in cmd:
+            raise AssertionError(f"Expected --deselect=skip.py::test_x, got {cmd}")
+        if "--reruns" not in cmd:
+            raise AssertionError(f"Expected --reruns in command, got {cmd}")
+        if cmd.count("pytest") != 1:
+            raise AssertionError(f"Expected exactly one 'pytest' token, got {cmd}")
+
+
+# -- Disabled vLLM test helpers (unit tests) -----------------------------------
+
+
+class FakeResponse:
+    """Minimal fake for urllib.request.urlopen return value."""
+
+    def __init__(self, body):
+        self._data = json.dumps(body).encode()
+
+    def read(self):
+        return self._data
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, *a):
+        pass
+
+
+@pytest.fixture
+def vllm_module():
+    return importlib.import_module(MOD)
+
+
+@pytest.mark.parametrize(
+    ("node_id", "expected_flag"),
+    (
+        ("a.py::test_1", "--deselect=a.py::test_1"),
+        ("a.py", "--ignore=a.py"),
+    ),
+)
+def test_build_flags_single_entry(vllm_module, node_id, expected_flag):
+    """:: in node_id produces --deselect, otherwise --ignore."""
+    flags = vllm_module._build_disabled_test_flags(
+        [{"test": node_id, "issue": "url"}], "plan"
+    )
+    if flags != expected_flag:
+        raise AssertionError(f"Expected '{expected_flag}', got '{flags}'")
+
+
+def test_build_flags_config_filter(vllm_module):
+    """Entries with matching configs are included; entries without configs apply to all."""
+    entries = [
+        {"test": "a.py", "issue": "url", "configs": ["plan_x"]},
+        {"test": "b.py::test_1", "issue": "url"},
+    ]
+    flags = vllm_module._build_disabled_test_flags(entries, "plan_x")
+    if "--ignore=a.py" not in flags:
+        raise AssertionError(f"Expected '--ignore=a.py' in '{flags}'")
+    if "--deselect=b.py::test_1" not in flags:
+        raise AssertionError(f"Expected '--deselect=b.py::test_1' in '{flags}'")
+
+
+def test_build_flags_config_excludes(vllm_module):
+    """Entries with non-matching configs are excluded."""
+    entries = [{"test": "a.py", "issue": "url", "configs": ["other"]}]
+    flags = vllm_module._build_disabled_test_flags(entries, "plan_x")
+    if flags != "":
+        raise AssertionError(f"Expected empty flags, got '{flags}'")
+
+
+def test_parse_issue_body(vllm_module):
+    """Extracts disabled_tests from a ```yaml code block in issue body."""
+    body = (
+        "## Disabled vLLM Tests\n"
+        "```yaml\n"
+        "disabled_tests:\n"
+        "  - test: foo.py::test_bar\n"
+        "  - test: baz.py\n"
+        "    configs:\n"
+        "      - plan_a\n"
+        "```\n"
+    )
+    entries = vllm_module._parse_disabled_tests_from_issue_body(body)
+    if len(entries) != 2:
+        raise AssertionError(f"Expected 2 entries, got {len(entries)}")
+    if entries[0]["test"] != "foo.py::test_bar":
+        raise AssertionError(f"Expected 'foo.py::test_bar', got '{entries[0]['test']}'")
+    if entries[1]["test"] != "baz.py":
+        raise AssertionError(f"Expected 'baz.py', got '{entries[1]['test']}'")
+    if entries[1]["configs"] != ["plan_a"]:
+        raise AssertionError(f"Expected ['plan_a'], got {entries[1]['configs']}")
+
+
+def test_parse_issue_body_no_yaml(vllm_module):
+    """Returns [] when issue body has no yaml code block."""
+    entries = vllm_module._parse_disabled_tests_from_issue_body("no yaml here")
+    if entries != []:
+        raise AssertionError(f"Expected [], got {entries}")
+
+
+def test_load_yaml_valid_entries(tmp_path, vllm_module):
+    """Loads and validates entries from a well-formed YAML file."""
+    yaml_file = tmp_path / "disabled.yaml"
+    yaml_file.write_text(
+        "disabled_tests:\n"
+        "  - test: a.py\n"
+        "    issue: https://github.com/pytorch/pytorch/issues/1\n"
+        "  - test: b.py::test_x\n"
+        "    issue: https://github.com/pytorch/pytorch/issues/2\n"
+        "    configs:\n"
+        "      - plan_a\n"
+    )
+    with mock_patch.object(vllm_module, "_DISABLED_VLLM_TESTS_PATH", yaml_file):
+        entries = vllm_module._load_disabled_vllm_tests_from_yaml()
+
+    if len(entries) != 2:
+        raise AssertionError(f"Expected 2 entries, got {len(entries)}")
+    if entries[0] != {
+        "test": "a.py",
+        "issue": "https://github.com/pytorch/pytorch/issues/175899",
+    }:
+        raise AssertionError(f"Unexpected first entry: {entries[0]}")
+    if entries[1]["test"] != "b.py::test_x":
+        raise AssertionError(f"Expected 'b.py::test_x', got '{entries[1]['test']}'")
+    if entries[1]["configs"] != ["plan_a"]:
+        raise AssertionError(f"Expected ['plan_a'], got {entries[1]['configs']}")
+
+
+def test_load_yaml_missing_keys(tmp_path, vllm_module):
+    """Raises ValueError when YAML entry is missing required 'issue' key."""
+    yaml_file = tmp_path / "disabled.yaml"
+    yaml_file.write_text("disabled_tests:\n  - test: foo.py\n")
+    with mock_patch.object(vllm_module, "_DISABLED_VLLM_TESTS_PATH", yaml_file):
+        with pytest.raises(ValueError, match="must have 'test' and 'issue' keys"):
+            vllm_module._load_disabled_vllm_tests_from_yaml()
+
+
+def test_load_yaml_missing_file(tmp_path, vllm_module):
+    """Returns [] when YAML file doesn't exist."""
+    with mock_patch.object(
+        vllm_module, "_DISABLED_VLLM_TESTS_PATH", tmp_path / "nope.yaml"
+    ):
+        result = vllm_module._load_disabled_vllm_tests_from_yaml()
+    if result != []:
+        raise AssertionError(f"Expected [], got {result}")
+
+
+def test_load_github_skipped_when_issue_unset(vllm_module):
+    """Returns [] immediately when _DISABLED_VLLM_TESTS_ISSUE is 0 (not configured)."""
+    with mock_patch.object(vllm_module, "_DISABLED_VLLM_TESTS_ISSUE", 0):
+        result = vllm_module._load_disabled_vllm_tests_from_github()
+    if result != []:
+        raise AssertionError(f"Expected [], got {result}")
+
+
+def test_load_github_failure_returns_empty(vllm_module):
+    """Network errors are swallowed and return []."""
+
+    def _raise(*args, **kwargs):
+        raise OSError("network error")
+
+    with (
+        mock_patch.object(vllm_module, "_DISABLED_VLLM_TESTS_ISSUE", 175899),
+        mock_patch.object(vllm_module.urllib.request, "urlopen", _raise),
+    ):
+        result = vllm_module._load_disabled_vllm_tests_from_github()
+    if result != []:
+        raise AssertionError(f"Expected [], got {result}")
+
+
+def test_load_github_filters_malformed_entries(vllm_module):
+    """Entries without 'test' key are filtered out; issue URL is auto-filled."""
+    fake_resp = FakeResponse(
+        {
+            "body": "```yaml\ndisabled_tests:\n  - bad_key: oops\n  - test: good.py\n```",
+            "html_url": "https://github.com/pytorch/pytorch/issues/175899",
+        }
+    )
+    with (
+        mock_patch.object(vllm_module, "_DISABLED_VLLM_TESTS_ISSUE", 175899),
+        mock_patch.object(
+            vllm_module.urllib.request, "urlopen", lambda *a, **kw: fake_resp
+        ),
+    ):
+        entries = vllm_module._load_disabled_vllm_tests_from_github()
+
+    if len(entries) != 1:
+        raise AssertionError(f"Expected 1 entry, got {len(entries)}")
+    if entries[0]["test"] != "good.py":
+        raise AssertionError(f"Expected 'good.py', got '{entries[0]['test']}'")
+    if entries[0]["issue"] != "https://github.com/pytorch/pytorch/issues/175899":
+        raise AssertionError(f"Expected issue URL, got '{entries[0]['issue']}'")
+
+
+def test_deduplication_yaml_wins(vllm_module):
+    """YAML entries take precedence over GitHub entries with the same test key."""
+    yaml_entries = [{"test": "a.py", "issue": "yaml-url"}]
+    github_entries = [
+        {"test": "a.py", "issue": "github-url"},
+        {"test": "b.py::test_1", "issue": "github-url"},
+    ]
+    with (
+        mock_patch.object(
+            vllm_module,
+            "_load_disabled_vllm_tests_from_yaml",
+            return_value=yaml_entries,
+        ),
+        mock_patch.object(
+            vllm_module,
+            "_load_disabled_vllm_tests_from_github",
+            return_value=github_entries,
+        ),
+    ):
+        merged = vllm_module._load_disabled_vllm_tests()
+
+    if len(merged) != 2:
+        raise AssertionError(f"Expected 2 merged entries, got {len(merged)}")
+    if merged[0] != {"test": "a.py", "issue": "yaml-url"}:
+        raise AssertionError(f"Expected yaml entry to win, got {merged[0]}")
+    if merged[1] != {"test": "b.py::test_1", "issue": "github-url"}:
+        raise AssertionError(f"Unexpected second entry: {merged[1]}")
diff --git a/.ci/magma/Makefile b/.ci/magma/Makefile
index 4169aedd03fa5..848b567cb8992 100644
--- a/.ci/magma/Makefile
+++ b/.ci/magma/Makefile
@@ -16,6 +16,7 @@ DOCKER_RUN = set -eou pipefail; ${DOCKER_CMD} run --rm -i \
 	magma/build_magma.sh
 
 .PHONY: all
+all: magma-cuda132
 all: magma-cuda130
 all: magma-cuda129
 all: magma-cuda128
@@ -26,6 +27,12 @@ clean:
 	$(RM) -r magma-*
 	$(RM) -r output
 
+.PHONY: magma-cuda132
+magma-cuda132: DESIRED_CUDA := 13.2
+magma-cuda132: CUDA_ARCH_LIST := -gencode arch=compute_80,code=sm_80 -gencode arch=compute_86,code=sm_86 -gencode arch=compute_90,code=sm_90 -gencode arch=compute_100,code=sm_100 -gencode arch=compute_120,code=sm_120
+magma-cuda132:
+	$(DOCKER_RUN)
+
 .PHONY: magma-cuda130
 magma-cuda130: DESIRED_CUDA := 13.0
 magma-cuda130: CUDA_ARCH_LIST := -gencode arch=compute_80,code=sm_80 -gencode arch=compute_86,code=sm_86 -gencode arch=compute_90,code=sm_90 -gencode arch=compute_100,code=sm_100 -gencode arch=compute_120,code=sm_120
diff --git a/.ci/manywheel/build_common.sh b/.ci/manywheel/build_common.sh
index d50bd623dace0..d76b75dea6850 100644
--- a/.ci/manywheel/build_common.sh
+++ b/.ci/manywheel/build_common.sh
@@ -91,11 +91,6 @@ export PYTORCH_BUILD_NUMBER=$build_number
 export CMAKE_LIBRARY_PATH="/opt/intel/lib:/lib:$CMAKE_LIBRARY_PATH"
 export CMAKE_INCLUDE_PATH="/opt/intel/include:$CMAKE_INCLUDE_PATH"
 
-if [[ -e /opt/openssl ]]; then
-    export OPENSSL_ROOT_DIR=/opt/openssl
-    export CMAKE_INCLUDE_PATH="/opt/openssl/include":$CMAKE_INCLUDE_PATH
-fi
-
 mkdir -p /tmp/$WHEELHOUSE_DIR
 
 export PATCHELF_BIN=/usr/local/bin/patchelf
@@ -118,6 +113,9 @@ retry pip install -qUr requirements-build.txt
 python setup.py clean
 retry pip install -qr requirements.txt
 case ${DESIRED_PYTHON} in
+  cp314*)
+    retry pip install -q --pre numpy==2.3.4
+    ;;
   cp31*)
     retry pip install -q --pre numpy==2.1.0
     ;;
diff --git a/.ci/manywheel/build_cuda.sh b/.ci/manywheel/build_cuda.sh
index 94bf6a6b4b26c..613301059f2a6 100644
--- a/.ci/manywheel/build_cuda.sh
+++ b/.ci/manywheel/build_cuda.sh
@@ -13,7 +13,7 @@ export ATEN_STATIC_CUDA=1
 export USE_CUDA_STATIC_LINK=1
 export INSTALL_TEST=0 # dont install test binaries into site-packages
 export USE_CUPTI_SO=0
-export USE_CUSPARSELT=${USE_CUSPARSELT:-1} # Enable if not disabled by libtorch build
+export USE_CUSPARSELT=${USE_CUSPARSELT:-1} # Enable if not disabled by libtorch build.
 export USE_CUFILE=${USE_CUFILE:-1}
 export USE_SYSTEM_NCCL=1
 export NCCL_INCLUDE_DIR="/usr/local/cuda/include/"
@@ -109,13 +109,13 @@ TORCH_CUDA_ARCH_LIST="7.5;8.0;8.6;9.0;10.0"
 case ${CUDA_VERSION} in
     12.6) TORCH_CUDA_ARCH_LIST="5.0;6.0;7.0;${TORCH_CUDA_ARCH_LIST//10.0/}" ;;  # Only 12.6 includes legacy Maxwell/Pascal/Volta, -Hopper support
     12.8) TORCH_CUDA_ARCH_LIST="${TORCH_CUDA_ARCH_LIST};12.0" ;;  # +Blackwell support
-    12.9) TORCH_CUDA_ARCH_LIST="${TORCH_CUDA_ARCH_LIST};12.0+PTX" # +Blackwell support + PTX for forward compatibility
+    12.9) TORCH_CUDA_ARCH_LIST="${TORCH_CUDA_ARCH_LIST};12.0" # +Blackwell support + PTX for forward compatibility
         if [[ "$PACKAGE_TYPE" == "libtorch" ]]; then
             TORCH_CUDA_ARCH_LIST="${TORCH_CUDA_ARCH_LIST//8.6;/}"  # Remove 8.6 for libtorch
         fi
         ;;
-    13.0)
-        TORCH_CUDA_ARCH_LIST="${TORCH_CUDA_ARCH_LIST};$([[ "$ARCH" == "aarch64" ]] && echo "11.0;" || echo "")12.0+PTX"
+    13.0|13.2)
+        TORCH_CUDA_ARCH_LIST="${TORCH_CUDA_ARCH_LIST};$([[ "$ARCH" == "aarch64" ]] && echo "11.0;" || echo "")12.0"
         export TORCH_NVCC_FLAGS="-compress-mode=size"
         export BUILD_BUNDLE_PTXAS=1
         ;;
diff --git a/.ci/manywheel/build_libtorch.sh b/.ci/manywheel/build_libtorch.sh
index d78fbd5c3ed36..852ecf7500604 100644
--- a/.ci/manywheel/build_libtorch.sh
+++ b/.ci/manywheel/build_libtorch.sh
@@ -59,12 +59,6 @@ export PYTORCH_BUILD_NUMBER=$build_number
 export CMAKE_LIBRARY_PATH="/opt/intel/lib:/lib:$CMAKE_LIBRARY_PATH"
 export CMAKE_INCLUDE_PATH="/opt/intel/include:$CMAKE_INCLUDE_PATH"
 
-# set OPENSSL_ROOT_DIR=/opt/openssl if it exists
-if [[ -e /opt/openssl ]]; then
-    export OPENSSL_ROOT_DIR=/opt/openssl
-    export CMAKE_INCLUDE_PATH="/opt/openssl/include":$CMAKE_INCLUDE_PATH
-fi
-
 # If given a python version like 3.6m or 2.7mu, convert this to the format we
 # expect. The binary CI jobs pass in python versions like this; they also only
 # ever pass one python version, so we assume that DESIRED_PYTHON is not a list
diff --git a/.ci/manywheel/build_rocm.sh b/.ci/manywheel/build_rocm.sh
index bac56746f4501..fa5724dca25a7 100755
--- a/.ci/manywheel/build_rocm.sh
+++ b/.ci/manywheel/build_rocm.sh
@@ -97,20 +97,13 @@ ROCM_SO_FILES=(
     "libhipblaslt.so"
     "libhipsparselt.so"
     "libhiprtc.so"
+    "librocprofiler-sdk.so"
+    "librocprofiler-register.so"
+    "libhsa-amd-aqlprofile64.so"
+    "librocm-core.so"
+    "librocroller.so"
 )
 
-if [[ $ROCM_INT -ge 60100 ]]; then
-    ROCM_SO_FILES+=("librocprofiler-register.so")
-fi
-
-if [[ $ROCM_INT -ge 60200 ]]; then
-    ROCM_SO_FILES+=("librocm-core.so")
-fi
-
-if [[ $ROCM_INT -ge 70000 ]]; then
-    ROCM_SO_FILES+=("librocroller.so")
-fi
-
 OS_NAME=`awk -F= '/^NAME/{print $2}' /etc/os-release`
 if [[ "$OS_NAME" == *"CentOS Linux"* || "$OS_NAME" == *"AlmaLinux"* ]]; then
     LIBGOMP_PATH="/usr/lib64/libgomp.so.1"
@@ -121,61 +114,23 @@ if [[ "$OS_NAME" == *"CentOS Linux"* || "$OS_NAME" == *"AlmaLinux"* ]]; then
     else
         LIBTINFO_PATH="/usr/lib64/libtinfo.so.6"
     fi
+    LIBDW_PATH="/usr/lib64/libdw.so.1"
     LIBDRM_PATH="/opt/amdgpu/lib64/libdrm.so.2"
     LIBDRM_AMDGPU_PATH="/opt/amdgpu/lib64/libdrm_amdgpu.so.1"
-    if [[ $ROCM_INT -ge 60100 && $ROCM_INT -lt 60300 ]]; then
-        # Below libs are direct dependencies of libhipsolver
-        LIBSUITESPARSE_CONFIG_PATH="/lib64/libsuitesparseconfig.so.4"
-        if [[ "$OS_NAME" == *"CentOS Linux"* ]]; then
-            LIBCHOLMOD_PATH="/lib64/libcholmod.so.2"
-            # Below libs are direct dependencies of libsatlas
-            LIBGFORTRAN_PATH="/lib64/libgfortran.so.3"
-        else
-            LIBCHOLMOD_PATH="/lib64/libcholmod.so.3"
-            # Below libs are direct dependencies of libsatlas
-            LIBGFORTRAN_PATH="/lib64/libgfortran.so.5"
-        fi
-        # Below libs are direct dependencies of libcholmod
-        LIBAMD_PATH="/lib64/libamd.so.2"
-        LIBCAMD_PATH="/lib64/libcamd.so.2"
-        LIBCCOLAMD_PATH="/lib64/libccolamd.so.2"
-        LIBCOLAMD_PATH="/lib64/libcolamd.so.2"
-        LIBSATLAS_PATH="/lib64/atlas/libsatlas.so.3"
-        # Below libs are direct dependencies of libsatlas
-        LIBQUADMATH_PATH="/lib64/libquadmath.so.0"
-    fi
     MAYBE_LIB64=lib64
 elif [[ "$OS_NAME" == *"Ubuntu"* ]]; then
     LIBGOMP_PATH="/usr/lib/x86_64-linux-gnu/libgomp.so.1"
     LIBNUMA_PATH="/usr/lib/x86_64-linux-gnu/libnuma.so.1"
     LIBELF_PATH="/usr/lib/x86_64-linux-gnu/libelf.so.1"
-    if [[ $ROCM_INT -ge 50300 ]]; then
-        LIBTINFO_PATH="/lib/x86_64-linux-gnu/libtinfo.so.6"
-    else
-        LIBTINFO_PATH="/lib/x86_64-linux-gnu/libtinfo.so.5"
-    fi
+    LIBTINFO_PATH="/lib/x86_64-linux-gnu/libtinfo.so.6"
+    LIBDW_PATH="/usr/lib/x86_64-linux-gnu/libdw.so.1"
     LIBDRM_PATH="/usr/lib/x86_64-linux-gnu/libdrm.so.2"
     LIBDRM_AMDGPU_PATH="/usr/lib/x86_64-linux-gnu/libdrm_amdgpu.so.1"
-    if [[ $ROCM_INT -ge 60100 && $ROCM_INT -lt 60300 ]]; then
-        # Below libs are direct dependencies of libhipsolver
-        LIBCHOLMOD_PATH="/lib/x86_64-linux-gnu/libcholmod.so.3"
-        # Below libs are direct dependencies of libcholmod
-        LIBSUITESPARSE_CONFIG_PATH="/lib/x86_64-linux-gnu/libsuitesparseconfig.so.5"
-        LIBAMD_PATH="/lib/x86_64-linux-gnu/libamd.so.2"
-        LIBCAMD_PATH="/lib/x86_64-linux-gnu/libcamd.so.2"
-        LIBCCOLAMD_PATH="/lib/x86_64-linux-gnu/libccolamd.so.2"
-        LIBCOLAMD_PATH="/lib/x86_64-linux-gnu/libcolamd.so.2"
-        LIBMETIS_PATH="/lib/x86_64-linux-gnu/libmetis.so.5"
-        LIBLAPACK_PATH="/lib/x86_64-linux-gnu/liblapack.so.3"
-        LIBBLAS_PATH="/lib/x86_64-linux-gnu/libblas.so.3"
-        # Below libs are direct dependencies of libblas
-        LIBGFORTRAN_PATH="/lib/x86_64-linux-gnu/libgfortran.so.5"
-        LIBQUADMATH_PATH="/lib/x86_64-linux-gnu/libquadmath.so.0"
-    fi
     MAYBE_LIB64=lib
 fi
 OS_SO_PATHS=($LIBGOMP_PATH $LIBNUMA_PATH\
              $LIBELF_PATH $LIBTINFO_PATH\
+             $LIBDW_PATH\
              $LIBDRM_PATH $LIBDRM_AMDGPU_PATH\
              $LIBSUITESPARSE_CONFIG_PATH\
              $LIBCHOLMOD_PATH $LIBAMD_PATH\
diff --git a/.ci/onnx/README.md b/.ci/onnx/README.md
deleted file mode 100644
index 47739136aabdf..0000000000000
--- a/.ci/onnx/README.md
+++ /dev/null
@@ -1,12 +0,0 @@
-# Jenkins
-
-The scripts in this directory are the entrypoint for testing ONNX exporter.
-
-The environment variable `BUILD_ENVIRONMENT` is expected to be set to
-the build environment you intend to test. It is a hint for the build
-and test scripts to configure Caffe2 a certain way and include/exclude
-tests. Docker images, they equal the name of the image itself. For
-example: `py2-cuda9.0-cudnn7-ubuntu16.04`. The Docker images that are
-built on Jenkins and are used in triggered builds already have this
-environment variable set in their manifest. Also see
-`./docker/jenkins/*/Dockerfile` and search for `BUILD_ENVIRONMENT`.
diff --git a/.ci/onnx/common.sh b/.ci/onnx/common.sh
deleted file mode 100644
index b8f912fbbb4e6..0000000000000
--- a/.ci/onnx/common.sh
+++ /dev/null
@@ -1,107 +0,0 @@
-#!/bin/bash
-
-set -ex
-
-source "$(dirname "${BASH_SOURCE[0]}")/../pytorch/common_utils.sh"
-
-LOCAL_DIR=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)
-ROOT_DIR=$(cd "$LOCAL_DIR"/../.. && pwd)
-TEST_DIR="$ROOT_DIR/test"
-pytest_reports_dir="${TEST_DIR}/test-reports/python"
-
-# Figure out which Python to use
-PYTHON="$(which python)"
-if [[ "${BUILD_ENVIRONMENT}" =~ py((2|3)\.?[0-9]?\.?[0-9]?) ]]; then
-  PYTHON=$(which "python${BASH_REMATCH[1]}")
-fi
-
-if [[ "${BUILD_ENVIRONMENT}" == *rocm* ]]; then
-    # HIP_PLATFORM is auto-detected by hipcc; unset to avoid build errors
-    unset HIP_PLATFORM
-fi
-
-mkdir -p "$pytest_reports_dir" || true
-
-##########################################
-# copied from .ci/pytorch/common_utils.sh
-##########################################
-
-function get_pinned_commit() {
-  cat .github/ci_commit_pins/"${1}".txt
-}
-
-function pip_install_whl() {
-  # This is used to install PyTorch and other build artifacts wheel locally
-  # without using any network connection
-
-  # Convert the input arguments into an array
-  local args=("$@")
-
-  # Check if the first argument contains multiple paths separated by spaces
-  if [[ "${args[0]}" == *" "* ]]; then
-    # Split the string by spaces into an array
-    IFS=' ' read -r -a paths <<< "${args[0]}"
-    # Loop through each path and install individually
-    for path in "${paths[@]}"; do
-      echo "Installing $path"
-      python3 -mpip install --no-index --no-deps "$path"
-    done
-  else
-    # Loop through each argument and install individually
-    for path in "${args[@]}"; do
-      echo "Installing $path"
-      python3 -mpip install --no-index --no-deps "$path"
-    done
-  fi
-}
-
-function pip_build_and_install() {
-  local build_target=$1
-  local wheel_dir=$2
-
-  local found_whl=0
-  for file in "${wheel_dir}"/*.whl
-  do
-    if [[ -f "${file}" ]]; then
-      found_whl=1
-      break
-    fi
-  done
-
-  # Build the wheel if it doesn't exist
-  if [ "${found_whl}" == "0" ]; then
-    python3 -m pip wheel \
-      --no-build-isolation \
-      --no-deps \
-      -w "${wheel_dir}" \
-      "${build_target}"
-  fi
-
-  for file in "${wheel_dir}"/*.whl
-  do
-    pip_install_whl "${file}"
-  done
-}
-
-function install_torchvision() {
-  local orig_preload
-  local commit
-  commit=$(get_pinned_commit vision)
-  orig_preload=${LD_PRELOAD}
-  if [ -n "${LD_PRELOAD}" ]; then
-    # Silence dlerror to work-around glibc ASAN bug, see https://sourceware.org/bugzilla/show_bug.cgi?id=27653#c9
-    echo 'char* dlerror(void) { return "";}'|gcc -fpic -shared -o "${HOME}/dlerror.so" -x c -
-    LD_PRELOAD=${orig_preload}:${HOME}/dlerror.so
-  fi
-
-  if [[ "${BUILD_ENVIRONMENT}" == *cuda* ]]; then
-    # Not sure if both are needed, but why not
-    export FORCE_CUDA=1
-    export WITH_CUDA=1
-  fi
-  pip_build_and_install "git+https://github.com/pytorch/vision.git@${commit}" dist/vision
-
-  if [ -n "${LD_PRELOAD}" ]; then
-    LD_PRELOAD=${orig_preload}
-  fi
-}
diff --git a/.ci/onnx/test.sh b/.ci/onnx/test.sh
deleted file mode 100755
index 1f2a23b49dc45..0000000000000
--- a/.ci/onnx/test.sh
+++ /dev/null
@@ -1,29 +0,0 @@
-#!/bin/bash
-
-# shellcheck source=./common.sh
-source "$(dirname "${BASH_SOURCE[0]}")/common.sh"
-
-# Workaround for dind-rootless userid mapping (https://github.com/pytorch/ci-infra/issues/96)
-WORKSPACE_ORIGINAL_OWNER_ID=$(stat -c '%u' "/var/lib/jenkins/workspace")
-cleanup_workspace() {
-  echo "sudo may print the following warning message that can be ignored. The chown command will still run."
-  echo "    sudo: setrlimit(RLIMIT_STACK): Operation not permitted"
-  echo "For more details refer to https://github.com/sudo-project/sudo/issues/42"
-  sudo chown -R "$WORKSPACE_ORIGINAL_OWNER_ID" /var/lib/jenkins/workspace
-}
-# Disable shellcheck SC2064 as we want to parse the original owner immediately.
-# shellcheck disable=SC2064
-trap_add cleanup_workspace EXIT
-sudo chown -R jenkins /var/lib/jenkins/workspace
-git config --global --add safe.directory /var/lib/jenkins/workspace
-
-if [[ "$BUILD_ENVIRONMENT" == *onnx* ]]; then
-  # TODO: This can be removed later once vision is also part of the Docker image
-  install_torchvision
-  # JIT C++ extensions require ninja, so put it into PATH.
-  export PATH="/var/lib/jenkins/.local/bin:$PATH"
-  # NB: ONNX test is fast (~15m) so it's ok to retry it few more times to avoid any flaky issue, we
-  # need to bring this to the standard PyTorch run_test eventually. The issue will be tracked in
-  # https://github.com/pytorch/pytorch/issues/98626
-  "$ROOT_DIR/scripts/onnx/test.sh"
-fi
diff --git a/.circleci/scripts/binary_linux_test.sh b/.ci/pytorch/binary_linux_test.sh
similarity index 98%
rename from .circleci/scripts/binary_linux_test.sh
rename to .ci/pytorch/binary_linux_test.sh
index 3771ecc108f87..1180012458dcc 100755
--- a/.circleci/scripts/binary_linux_test.sh
+++ b/.ci/pytorch/binary_linux_test.sh
@@ -1,4 +1,5 @@
 #!/bin/bash
+# shellcheck disable=SC1091,SC2012,SC2154
 
 OUTPUT_SCRIPT=${OUTPUT_SCRIPT:-/home/circleci/project/ci_test_script.sh}
 
diff --git a/.circleci/scripts/binary_populate_env.sh b/.ci/pytorch/binary_populate_env.sh
similarity index 98%
rename from .circleci/scripts/binary_populate_env.sh
rename to .ci/pytorch/binary_populate_env.sh
index 74ad225db933b..53914914c8c93 100755
--- a/.circleci/scripts/binary_populate_env.sh
+++ b/.ci/pytorch/binary_populate_env.sh
@@ -1,4 +1,5 @@
 #!/bin/bash
+# shellcheck disable=SC2002,SC2004,SC2086,SC2129,SC2155
 set -eux -o pipefail
 export TZ=UTC
 
@@ -111,9 +112,8 @@ if [[ "$PACKAGE_TYPE" =~ .*wheel.* && -n "$PYTORCH_BUILD_VERSION" && "$PYTORCH_B
     fi
 fi
 
-USE_GLOO_WITH_OPENSSL="ON"
+USE_GLOO_WITH_OPENSSL="OFF"
 if [[ "$GPU_ARCH_TYPE" =~ .*aarch64.* ]]; then
-  USE_GLOO_WITH_OPENSSL="OFF"
   USE_GOLD_LINKER="OFF"
 fi
 
diff --git a/.circleci/scripts/binary_upload.sh b/.ci/pytorch/binary_upload.sh
similarity index 62%
rename from .circleci/scripts/binary_upload.sh
rename to .ci/pytorch/binary_upload.sh
index d48077e112455..a0ca32be72852 100755
--- a/.circleci/scripts/binary_upload.sh
+++ b/.ci/pytorch/binary_upload.sh
@@ -1,5 +1,5 @@
 #!/usr/bin/env bash
-
+# shellcheck disable=SC2231
 set -euo pipefail
 
 PACKAGE_TYPE=${PACKAGE_TYPE:-wheel}
@@ -62,17 +62,61 @@ s3_upload() {
   )
 }
 
+R2_UPLOAD=${R2_UPLOAD:-}
+R2_BUCKET="s3://pytorch-downloads"
+R2_ACCOUNT_ID=${R2_ACCOUNT_ID:-}
+R2_ACCESS_KEY_ID=${R2_ACCESS_KEY_ID:-}
+R2_SECRET_ACCESS_KEY=${R2_SECRET_ACCESS_KEY:-}
+
+r2_upload() {
+  if [[ -z "${R2_ACCOUNT_ID}" || -z "${R2_ACCESS_KEY_ID}" || -z "${R2_SECRET_ACCESS_KEY}" ]]; then
+    echo "WARNING: R2 credentials not configured, skipping R2 upload"
+    return
+  fi
+  local extension
+  local pkg_type
+  extension="$1"
+  pkg_type="$2"
+  r2_root_dir="${R2_BUCKET}/${pkg_type}/${UPLOAD_CHANNEL}"
+  if [[ -z ${UPLOAD_SUBFOLDER:-} ]]; then
+    r2_upload_dir="${r2_root_dir}/"
+  else
+    r2_upload_dir="${r2_root_dir}/${UPLOAD_SUBFOLDER}/"
+  fi
+  (
+    for pkg in ${PKG_DIR}/*.${extension}; do
+      (
+        set -x
+        shm_id=$(sha256sum "${pkg}" | awk '{print $1}')
+        AWS_ACCESS_KEY_ID="${R2_ACCESS_KEY_ID}" \
+        AWS_SECRET_ACCESS_KEY="${R2_SECRET_ACCESS_KEY}" \
+        AWS_SESSION_TOKEN="" \
+        AWS_DEFAULT_REGION="auto" \
+        ${AWS_S3_CP} --no-progress "${pkg}" "${r2_upload_dir}" \
+          --metadata "checksum-sha256=${shm_id}" \
+          --endpoint-url "https://${R2_ACCOUNT_ID}.r2.cloudflarestorage.com"
+      )
+    done
+  )
+}
+
 # Install dependencies (should be a no-op if previously installed)
 pip install -q awscli uv
 
 case "${PACKAGE_TYPE}" in
   libtorch)
     s3_upload "zip" "libtorch"
+    if [[ "${R2_UPLOAD}" == "true" ]]; then
+      r2_upload "zip" "libtorch"
+    fi
     BACKUP_DIR="libtorch/${UPLOAD_CHANNEL}/${UPLOAD_SUBFOLDER}"
     ;;
   # wheel can either refer to wheel/manywheel
   *wheel)
     s3_upload "whl" "whl"
+    if [[ "${R2_UPLOAD}" == "true" ]]; then
+      r2_upload "whl" "whl"
+    fi
     BACKUP_DIR="whl/${UPLOAD_CHANNEL}/${UPLOAD_SUBFOLDER}"
     ;;
   *)
diff --git a/.circleci/scripts/binary_windows_build.sh b/.ci/pytorch/binary_windows_build.sh
similarity index 97%
rename from .circleci/scripts/binary_windows_build.sh
rename to .ci/pytorch/binary_windows_build.sh
index 59dbbb3d9b6a8..2f560233b1201 100644
--- a/.circleci/scripts/binary_windows_build.sh
+++ b/.ci/pytorch/binary_windows_build.sh
@@ -1,4 +1,5 @@
 #!/bin/bash
+# shellcheck disable=SC1090
 set -eux -o pipefail
 
 source "${BINARY_ENV_FILE:-/c/w/env}"
diff --git a/.circleci/scripts/binary_windows_test.sh b/.ci/pytorch/binary_windows_test.sh
similarity index 93%
rename from .circleci/scripts/binary_windows_test.sh
rename to .ci/pytorch/binary_windows_test.sh
index b8b82979caf48..f64a9447193bb 100644
--- a/.circleci/scripts/binary_windows_test.sh
+++ b/.ci/pytorch/binary_windows_test.sh
@@ -1,4 +1,5 @@
 #!/bin/bash
+# shellcheck disable=SC1090
 set -eux -o pipefail
 
 source "${BINARY_ENV_FILE:-/c/w/env}"
diff --git a/.ci/pytorch/build.sh b/.ci/pytorch/build.sh
index eb3529a0c43f3..265a6f579f427 100755
--- a/.ci/pytorch/build.sh
+++ b/.ci/pytorch/build.sh
@@ -23,12 +23,6 @@ cmake --version
 echo "Environment variables:"
 env
 
-# The sccache wrapped version of nvcc gets put in /opt/cache/lib in docker since
-# there are some issues if it is always wrapped, so we need to add it to PATH
-# during CI builds.
-# https://github.com/pytorch/pytorch/blob/0b6c0898e6c352c8ea93daec854e704b41485375/.ci/docker/common/install_cache.sh#L97
-export PATH="/opt/cache/lib:$PATH"
-
 if [[ "$BUILD_ENVIRONMENT" == *cuda* ]]; then
   # Use jemalloc during compilation to mitigate https://github.com/pytorch/pytorch/issues/116289
   export LD_PRELOAD=/usr/lib/x86_64-linux-gnu/libjemalloc.so.2
@@ -41,16 +35,6 @@ if [[ "$BUILD_ENVIRONMENT" == *cuda13* ]]; then
   export USE_FBGEMM=0
 fi
 
-if [[ "$BUILD_ENVIRONMENT" == *cuda11* ]]; then
-  if [[ "$BUILD_ENVIRONMENT" != *clang* ]]; then
-    # TODO: there is a linking issue when building with UCC using clang,
-    # disable it for now and to be fix later.
-    # TODO: disable UCC temporarily to enable CUDA 12.1 in CI
-    export USE_UCC=1
-    export USE_SYSTEM_UCC=1
-  fi
-fi
-
 if [[ ${BUILD_ENVIRONMENT} == *"parallelnative"* ]]; then
   export ATEN_THREADING=NATIVE
 fi
@@ -118,30 +102,6 @@ if [[ "$BUILD_ENVIRONMENT" == *riscv64* ]]; then
 
 fi
 
-if [[ "$BUILD_ENVIRONMENT" == *libtorch* ]]; then
-  POSSIBLE_JAVA_HOMES=()
-  POSSIBLE_JAVA_HOMES+=(/usr/local)
-  POSSIBLE_JAVA_HOMES+=(/usr/lib/jvm/java-8-openjdk-amd64)
-  POSSIBLE_JAVA_HOMES+=(/Library/Java/JavaVirtualMachines/*.jdk/Contents/Home)
-  # Add the Windows-specific JNI
-  POSSIBLE_JAVA_HOMES+=("$PWD/.circleci/windows-jni/")
-  for JH in "${POSSIBLE_JAVA_HOMES[@]}" ; do
-    if [[ -e "$JH/include/jni.h" ]] ; then
-      # Skip if we're not on Windows but haven't found a JAVA_HOME
-      if [[ "$JH" == "$PWD/.circleci/windows-jni/" && "$OSTYPE" != "msys" ]] ; then
-        break
-      fi
-      echo "Found jni.h under $JH"
-      export JAVA_HOME="$JH"
-      export BUILD_JNI=ON
-      break
-    fi
-  done
-  if [ -z "$JAVA_HOME" ]; then
-    echo "Did not find jni.h"
-  fi
-fi
-
 # Use special scripts for Android builds
 
 if [[ "$BUILD_ENVIRONMENT" == *vulkan* ]]; then
@@ -184,6 +144,7 @@ if [[ "$BUILD_ENVIRONMENT" == *xpu* ]]; then
   export USE_XCCL=1
   export USE_MPI=0
   export TORCH_XPU_ARCH_LIST=pvc
+  export USE_STATIC_MKL=1
 fi
 
 # sccache will fail for CUDA builds if all cores are used for compiling
@@ -214,9 +175,9 @@ if [[ "$BUILD_ENVIRONMENT" == *cuda* ]] && echo "${TORCH_CUDA_ARCH_LIST}" | tr '
   export BUILD_CUSTOM_STEP="ninja -C build flash_attention -j ${J}"
 fi
 
-if [[ "${BUILD_ENVIRONMENT}" == *clang* ]]; then
-  export CC=clang
-  export CXX=clang++
+# TODO: Removeme once all the wrappers are gone
+if [[ "$BUILD_ENVIRONMENT" == *clang* ]] && [[ "$BUILD_ENVIRONMENT" == *cuda* ]]; then
+  sudo rm -f /opt/cache/bin/clang++
 fi
 
 if [[ "$BUILD_ENVIRONMENT" == *-clang*-asan* ]]; then
@@ -232,10 +193,6 @@ if [[ "${BUILD_ENVIRONMENT}" == *no-ops* ]]; then
   export USE_PER_OPERATOR_HEADERS=0
 fi
 
-if [[ "${BUILD_ENVIRONMENT}" == *-pch* ]]; then
-    export USE_PRECOMPILED_HEADERS=1
-fi
-
 if [[ "${BUILD_ENVIRONMENT}" != *cuda* ]]; then
   export BUILD_STATIC_RUNTIME_BENCHMARK=ON
 fi
@@ -264,187 +221,168 @@ if [[ "$BUILD_ENVIRONMENT" != *rocm* && "$BUILD_ENVIRONMENT" != *s390x* && "$BUI
   git config --global --add safe.directory /var/lib/jenkins/workspace
 fi
 
-if [[ "$BUILD_ENVIRONMENT" == *-bazel-* ]]; then
-  set -e -o pipefail
-
-  get_bazel
-  python3 tools/optional_submodules.py checkout_eigen
+# check that setup.py would fail with bad arguments
+echo "The next three invocations are expected to fail with invalid command error messages."
+( ! get_exit_code python setup.py bad_argument )
+( ! get_exit_code python setup.py clean] )
+( ! get_exit_code python setup.py clean bad_argument )
+
+if [[ "$BUILD_ENVIRONMENT" != *libtorch* ]]; then
+  # rocm builds fail when WERROR=1
+  # XLA test build fails when WERROR=1
+  # set only when building other architectures
+  # or building non-XLA tests.
+  if [[ "$BUILD_ENVIRONMENT" != *rocm*  && "$BUILD_ENVIRONMENT" != *xla* && "$BUILD_ENVIRONMENT" != *riscv64* ]]; then
+    # TODO: Remove me and may be just focus on numpy-2.x testing
+    if [[ "$ANACONDA_PYTHON_VERSION" =~ ^3\.1[0-2]$ ]]; then
+      # Install numpy-2.0.2 for builds which are backward compatible with 1.X
+      # In relality it's only needed for numpy_2_x and vllm shards (where vllm depends on numpy-2)
+      python -mpip install numpy==2.0.2
+    fi
 
-  # Leave 1 CPU free and use only up to 80% of memory to reduce the change of crashing
-  # the runner
-  BAZEL_MEM_LIMIT="--local_ram_resources=HOST_RAM*.8"
-  BAZEL_CPU_LIMIT="--local_cpu_resources=HOST_CPUS-1"
+    WERROR=1 python setup.py clean
 
-  if [[ "$CUDA_VERSION" == "cpu" ]]; then
-    # Build torch, the Python module, and tests for CPU-only
-    tools/bazel build --config=no-tty "${BAZEL_MEM_LIMIT}" "${BAZEL_CPU_LIMIT}" --config=cpu-only :torch :torch/_C.so :all_tests
+    WERROR=1 python -m build --wheel --no-isolation
   else
-    tools/bazel build --config=no-tty "${BAZEL_MEM_LIMIT}" "${BAZEL_CPU_LIMIT}" //...
-  fi
-else
-  # check that setup.py would fail with bad arguments
-  echo "The next three invocations are expected to fail with invalid command error messages."
-  ( ! get_exit_code python setup.py bad_argument )
-  ( ! get_exit_code python setup.py clean] )
-  ( ! get_exit_code python setup.py clean bad_argument )
-
-  if [[ "$BUILD_ENVIRONMENT" != *libtorch* ]]; then
-    # rocm builds fail when WERROR=1
-    # XLA test build fails when WERROR=1
-    # set only when building other architectures
-    # or building non-XLA tests.
-    if [[ "$BUILD_ENVIRONMENT" != *rocm*  && "$BUILD_ENVIRONMENT" != *xla* && "$BUILD_ENVIRONMENT" != *riscv64* ]]; then
-      # TODO: Remove me and may be just focus on numpy-2.x testing
-      if [[ "$ANACONDA_PYTHON_VERSION" =~ ^3\.1[0-2]$ ]]; then
-        # Install numpy-2.0.2 for builds which are backward compatible with 1.X
-        # In relality it's only needed for numpy_2_x and vllm shards (where vllm depends on numpy-2)
-        python -mpip install numpy==2.0.2
-      fi
-
-      WERROR=1 python setup.py clean
-
-      WERROR=1 python -m build --wheel --no-isolation
-    else
-      python setup.py clean
-      if [[ "$BUILD_ENVIRONMENT" == *xla* ]]; then
-        source .ci/pytorch/install_cache_xla.sh
-      fi
-      python -m build --wheel --no-isolation
-    fi
-    pip_install_whl "$(echo dist/*.whl)"
-    if [[ "$BUILD_ENVIRONMENT" == *full-debug* ]]; then
-      # Regression test for https://github.com/pytorch/pytorch/issues/164297
-      # Torch should be importable and that's about it
-      pushd /; python -c "import torch;print(torch.__config__.show(), torch.randn(5) + 1.7)"; popd
+    python setup.py clean
+    if [[ "$BUILD_ENVIRONMENT" == *xla* ]]; then
+      source .ci/pytorch/install_cache_xla.sh
     fi
+    python -m build --wheel --no-isolation
+  fi
+  pip_install_whl "$(echo dist/*.whl)"
+  if [[ "$BUILD_ENVIRONMENT" == *full-debug* ]]; then
+    # Regression test for https://github.com/pytorch/pytorch/issues/164297
+    # Torch should be importable and that's about it
+    pushd /; python -c "import torch;print(torch.__config__.show(), torch.randn(5) + 1.7)"; popd
+  fi
 
-    if [[ "${BUILD_ADDITIONAL_PACKAGES:-}" == *vision* ]]; then
-      install_torchvision
-    fi
+  if [[ "${BUILD_ADDITIONAL_PACKAGES:-}" == *vision* ]]; then
+    install_torchvision
+  fi
 
-    if [[ "${BUILD_ADDITIONAL_PACKAGES:-}" == *audio* ]]; then
-      install_torchaudio
-    fi
+  if [[ "${BUILD_ADDITIONAL_PACKAGES:-}" == *audio* ]]; then
+    install_torchaudio
+  fi
 
-    if [[ "${BUILD_ADDITIONAL_PACKAGES:-}" == *torchrec* || "${BUILD_ADDITIONAL_PACKAGES:-}" == *fbgemm* ]]; then
-      install_torchrec_and_fbgemm
-    fi
+  if [[ "${BUILD_ADDITIONAL_PACKAGES:-}" == *torchrec* || "${BUILD_ADDITIONAL_PACKAGES:-}" == *fbgemm* ]]; then
+    install_torchrec_and_fbgemm
+  fi
 
-    if [[ "${BUILD_ADDITIONAL_PACKAGES:-}" == *torchao* ]]; then
-      install_torchao
-    fi
+  if [[ "${BUILD_ADDITIONAL_PACKAGES:-}" == *torchao* ]]; then
+    install_torchao
+  fi
 
-    if [[ "$BUILD_ENVIRONMENT" == *xpu* ]]; then
-      echo "Checking that xpu is compiled"
-      pushd dist/
-      if python -c 'import torch; exit(0 if torch.xpu._is_compiled() else 1)'; then
-        echo "XPU support is compiled in."
-      else
-        echo "XPU support is NOT compiled in."
-        exit 1
-      fi
-      popd
+  if [[ "$BUILD_ENVIRONMENT" == *xpu* ]]; then
+    echo "Checking that xpu is compiled"
+    pushd dist/
+    if python -c 'import torch; exit(0 if torch.xpu._is_compiled() else 1)'; then
+      echo "XPU support is compiled in."
+    else
+      echo "XPU support is NOT compiled in."
+      exit 1
     fi
+    popd
+  fi
 
-    # TODO: I'm not sure why, but somehow we lose verbose commands
-    set -x
+  # TODO: I'm not sure why, but somehow we lose verbose commands
+  set -x
 
-    assert_git_not_dirty
-    # Copy ninja build logs to dist folder
-    mkdir -p dist
-    if [ -f build/.ninja_log ]; then
-      cp build/.ninja_log dist
-    fi
+  assert_git_not_dirty
+  # Copy ninja build logs to dist folder
+  mkdir -p dist
+  if [ -f build/.ninja_log ]; then
+    cp build/.ninja_log dist
+  fi
 
-    if [[ "$BUILD_ENVIRONMENT" == *rocm* ]]; then
-      # remove sccache wrappers post-build; runtime compilation of MIOpen kernels does not yet fully support them
-      sudo rm -f /opt/cache/bin/cc
-      sudo rm -f /opt/cache/bin/c++
-      sudo rm -f /opt/cache/bin/gcc
-      sudo rm -f /opt/cache/bin/g++
-      # Restore original clang compilers that were backed up during sccache wrapping.
-      # Skip for theRock nightly: sccache wrapping is disabled, so no backup exists.
-      # theRock also uses ${ROCM_PATH}/lib/llvm/bin instead of /opt/rocm/llvm/bin.
-      if [[ -d /opt/rocm/llvm/bin ]]; then
-        pushd /opt/rocm/llvm/bin
-        if [[ -d original ]]; then
-          sudo mv original/clang .
-          sudo mv original/clang++ .
-        fi
-        sudo rm -rf original
-        popd
+  if [[ "$BUILD_ENVIRONMENT" == *rocm* ]]; then
+    # remove sccache wrappers post-build; runtime compilation of MIOpen kernels does not yet fully support them
+    sudo rm -f /opt/cache/bin/cc
+    sudo rm -f /opt/cache/bin/c++
+    sudo rm -f /opt/cache/bin/gcc
+    sudo rm -f /opt/cache/bin/g++
+    # Restore original clang compilers that were backed up during sccache wrapping.
+    # Skip for theRock nightly: sccache wrapping is disabled, so no backup exists.
+    # theRock also uses ${ROCM_PATH}/lib/llvm/bin instead of /opt/rocm/llvm/bin.
+    if [[ -d /opt/rocm/llvm/bin ]]; then
+      pushd /opt/rocm/llvm/bin
+      if [[ -d original ]]; then
+        sudo mv original/clang .
+        sudo mv original/clang++ .
       fi
+      sudo rm -rf original
+      popd
     fi
-
-    CUSTOM_TEST_ARTIFACT_BUILD_DIR=${CUSTOM_TEST_ARTIFACT_BUILD_DIR:-"build/custom_test_artifacts"}
-    CUSTOM_TEST_USE_ROCM=$([[ "$BUILD_ENVIRONMENT" == *rocm* ]] && echo "ON" || echo "OFF")
-    CUSTOM_TEST_MODULE_PATH="${PWD}/cmake/public"
-    mkdir -pv "${CUSTOM_TEST_ARTIFACT_BUILD_DIR}"
-
-    # Build custom operator tests.
-    CUSTOM_OP_BUILD="${CUSTOM_TEST_ARTIFACT_BUILD_DIR}/custom-op-build"
-    CUSTOM_OP_TEST="$PWD/test/custom_operator"
-    python --version
-    SITE_PACKAGES="$(python -c 'import site; print(";".join([x for x in site.getsitepackages()] + [x + "/torch" for x in site.getsitepackages()]))')"
-
-    mkdir -p "$CUSTOM_OP_BUILD"
-    pushd "$CUSTOM_OP_BUILD"
-    cmake "$CUSTOM_OP_TEST" -DCMAKE_PREFIX_PATH="$SITE_PACKAGES" -DPython_EXECUTABLE="$(which python)" \
-          -DCMAKE_MODULE_PATH="$CUSTOM_TEST_MODULE_PATH" -DUSE_ROCM="$CUSTOM_TEST_USE_ROCM"
-    make VERBOSE=1
-    popd
-    assert_git_not_dirty
-
-    # Build jit hook tests
-    JIT_HOOK_BUILD="${CUSTOM_TEST_ARTIFACT_BUILD_DIR}/jit-hook-build"
-    JIT_HOOK_TEST="$PWD/test/jit_hooks"
-    python --version
-    SITE_PACKAGES="$(python -c 'import site; print(";".join([x for x in site.getsitepackages()] + [x + "/torch" for x in site.getsitepackages()]))')"
-    mkdir -p "$JIT_HOOK_BUILD"
-    pushd "$JIT_HOOK_BUILD"
-    cmake "$JIT_HOOK_TEST" -DCMAKE_PREFIX_PATH="$SITE_PACKAGES" -DPython_EXECUTABLE="$(which python)" \
-          -DCMAKE_MODULE_PATH="$CUSTOM_TEST_MODULE_PATH" -DUSE_ROCM="$CUSTOM_TEST_USE_ROCM"
-    make VERBOSE=1
-    popd
-    assert_git_not_dirty
-
-    # Build custom backend tests.
-    CUSTOM_BACKEND_BUILD="${CUSTOM_TEST_ARTIFACT_BUILD_DIR}/custom-backend-build"
-    CUSTOM_BACKEND_TEST="$PWD/test/custom_backend"
-    python --version
-    mkdir -p "$CUSTOM_BACKEND_BUILD"
-    pushd "$CUSTOM_BACKEND_BUILD"
-    cmake "$CUSTOM_BACKEND_TEST" -DCMAKE_PREFIX_PATH="$SITE_PACKAGES" -DPython_EXECUTABLE="$(which python)" \
-          -DCMAKE_MODULE_PATH="$CUSTOM_TEST_MODULE_PATH" -DUSE_ROCM="$CUSTOM_TEST_USE_ROCM"
-    make VERBOSE=1
-    popd
-    assert_git_not_dirty
-  else
-    # Test no-Python build
-    echo "Building libtorch"
-
-    # This is an attempt to mitigate flaky libtorch build OOM error. By default, the build parallelization
-    # is set to be the number of CPU minus 2. So, let's try a more conservative value here. A 4xlarge has
-    # 16 CPUs
-    MAX_JOBS=$(nproc --ignore=4)
-    export MAX_JOBS
-
-    # NB: Install outside of source directory (at the same level as the root
-    # pytorch folder) so that it doesn't get cleaned away prior to docker push.
-    BUILD_LIBTORCH_PY=$PWD/tools/build_libtorch.py
-    mkdir -p ../cpp-build/caffe2
-    pushd ../cpp-build/caffe2
-    WERROR=1 VERBOSE=1 DEBUG=1 python "$BUILD_LIBTORCH_PY"
-    popd
   fi
+
+  CUSTOM_TEST_ARTIFACT_BUILD_DIR=${CUSTOM_TEST_ARTIFACT_BUILD_DIR:-"build/custom_test_artifacts"}
+  CUSTOM_TEST_USE_ROCM=$([[ "$BUILD_ENVIRONMENT" == *rocm* ]] && echo "ON" || echo "OFF")
+  CUSTOM_TEST_MODULE_PATH="${PWD}/cmake/public"
+  mkdir -pv "${CUSTOM_TEST_ARTIFACT_BUILD_DIR}"
+
+  # Build custom operator tests.
+  CUSTOM_OP_BUILD="${CUSTOM_TEST_ARTIFACT_BUILD_DIR}/custom-op-build"
+  CUSTOM_OP_TEST="$PWD/test/custom_operator"
+  python --version
+  SITE_PACKAGES="$(python -c 'import site; print(";".join([x for x in site.getsitepackages()] + [x + "/torch" for x in site.getsitepackages()]))')"
+
+  mkdir -p "$CUSTOM_OP_BUILD"
+  pushd "$CUSTOM_OP_BUILD"
+  cmake "$CUSTOM_OP_TEST" -DCMAKE_PREFIX_PATH="$SITE_PACKAGES" -DPython_EXECUTABLE="$(which python)" \
+        -DCMAKE_MODULE_PATH="$CUSTOM_TEST_MODULE_PATH" -DUSE_ROCM="$CUSTOM_TEST_USE_ROCM"
+  make VERBOSE=1
+  popd
+  assert_git_not_dirty
+
+  # Build jit hook tests
+  JIT_HOOK_BUILD="${CUSTOM_TEST_ARTIFACT_BUILD_DIR}/jit-hook-build"
+  JIT_HOOK_TEST="$PWD/test/jit_hooks"
+  python --version
+  SITE_PACKAGES="$(python -c 'import site; print(";".join([x for x in site.getsitepackages()] + [x + "/torch" for x in site.getsitepackages()]))')"
+  mkdir -p "$JIT_HOOK_BUILD"
+  pushd "$JIT_HOOK_BUILD"
+  cmake "$JIT_HOOK_TEST" -DCMAKE_PREFIX_PATH="$SITE_PACKAGES" -DPython_EXECUTABLE="$(which python)" \
+        -DCMAKE_MODULE_PATH="$CUSTOM_TEST_MODULE_PATH" -DUSE_ROCM="$CUSTOM_TEST_USE_ROCM"
+  make VERBOSE=1
+  popd
+  assert_git_not_dirty
+
+  # Build custom backend tests.
+  CUSTOM_BACKEND_BUILD="${CUSTOM_TEST_ARTIFACT_BUILD_DIR}/custom-backend-build"
+  CUSTOM_BACKEND_TEST="$PWD/test/custom_backend"
+  python --version
+  mkdir -p "$CUSTOM_BACKEND_BUILD"
+  pushd "$CUSTOM_BACKEND_BUILD"
+  cmake "$CUSTOM_BACKEND_TEST" -DCMAKE_PREFIX_PATH="$SITE_PACKAGES" -DPython_EXECUTABLE="$(which python)" \
+        -DCMAKE_MODULE_PATH="$CUSTOM_TEST_MODULE_PATH" -DUSE_ROCM="$CUSTOM_TEST_USE_ROCM"
+  make VERBOSE=1
+  popd
+  assert_git_not_dirty
+else
+  # Test no-Python build
+  echo "Building libtorch"
+
+  # This is an attempt to mitigate flaky libtorch build OOM error. By default, the build parallelization
+  # is set to be the number of CPU minus 2. So, let's try a more conservative value here. A 4xlarge has
+  # 16 CPUs
+  MAX_JOBS=$(nproc --ignore=4)
+  export MAX_JOBS
+
+  BUILD_LIBTORCH_PY=$PWD/tools/build_libtorch.py
+  # Build outside the source tree so the artifacts don't interfere with
+  # the workspace. /tmp is writable on both EC2 and OSDC runners.
+  mkdir -p /tmp/cpp-build/caffe2
+  pushd /tmp/cpp-build/caffe2
+  WERROR=1 VERBOSE=1 DEBUG=1 python "$BUILD_LIBTORCH_PY"
+  popd
 fi
 
-if [[ "$BUILD_ENVIRONMENT" != *libtorch* && "$BUILD_ENVIRONMENT" != *bazel* ]]; then
+if [[ "$BUILD_ENVIRONMENT" != *libtorch* ]]; then
   # export test times so that potential sharded tests that'll branch off this build will use consistent data
   # don't do this for libtorch as libtorch is C++ only and thus won't have python tests run on its build
   PYTHONPATH=. python tools/stats/export_test_times.py
 fi
-# don't do this for bazel or s390x or riscv64 as they don't use sccache
-if [[ "$BUILD_ENVIRONMENT" != *s390x* && "$BUILD_ENVIRONMENT" != *riscv64* && "$BUILD_ENVIRONMENT" != *-bazel-* ]]; then
+# don't do this for s390x or riscv64 as they don't use sccache
+if [[ "$BUILD_ENVIRONMENT" != *s390x* && "$BUILD_ENVIRONMENT" != *riscv64* ]]; then
   print_sccache_stats
 fi
diff --git a/.ci/pytorch/check_binary.sh b/.ci/pytorch/check_binary.sh
index c8c89fe871fe3..9356970394e81 100755
--- a/.ci/pytorch/check_binary.sh
+++ b/.ci/pytorch/check_binary.sh
@@ -154,7 +154,7 @@ setup_link_flags () {
 TEST_CODE_DIR="$(dirname $(realpath ${BASH_SOURCE[0]}))/test_example_code"
 build_and_run_example_cpp () {
   setup_link_flags
-  g++ ${TEST_CODE_DIR}/$1.cpp -I${install_root}/include -I${install_root}/include/torch/csrc/api/include -std=gnu++17 -L${install_root}/lib ${REF_LIB} ${ADDITIONAL_LINKER_FLAGS} -ltorch $TORCH_CPU_LINK_FLAGS $TORCH_CUDA_LINK_FLAGS $C10_LINK_FLAGS -o $1
+  g++ ${TEST_CODE_DIR}/$1.cpp -I${install_root}/include -I${install_root}/include/torch/csrc/api/include -std=gnu++20 -L${install_root}/lib ${REF_LIB} ${ADDITIONAL_LINKER_FLAGS} -ltorch $TORCH_CPU_LINK_FLAGS $TORCH_CUDA_LINK_FLAGS $C10_LINK_FLAGS -o $1
   ./$1
 }
 
@@ -292,25 +292,6 @@ if [[ "$PACKAGE_TYPE" != 'libtorch' ]]; then
   popd
 fi
 
-###############################################################################
-# Check PyTorch supports TCP_TLS gloo transport
-###############################################################################
-
-if [[ "$(uname)" == 'Linux' && "$PACKAGE_TYPE" != 'libtorch' ]]; then
-  GLOO_CHECK="import torch.distributed as dist
-try:
-    dist.init_process_group('gloo', rank=0, world_size=1)
-except RuntimeError as e:
-    print(e)
-"
-  RESULT=`GLOO_DEVICE_TRANSPORT=TCP_TLS MASTER_ADDR=localhost MASTER_PORT=63945 python -c "$GLOO_CHECK"`
-  GLOO_TRANSPORT_IS_NOT_SUPPORTED='gloo transport is not supported'
-  if [[ "$RESULT" =~ "$GLOO_TRANSPORT_IS_NOT_SUPPORTED" ]]; then
-    echo "PyTorch doesn't support TLS_TCP transport, please build with USE_GLOO_WITH_OPENSSL=1"
-    exit 1
-  fi
-fi
-
 ###############################################################################
 # Restore LD_LIBRARY_PATH to its original value
 ###############################################################################
diff --git a/.ci/pytorch/common-build.sh b/.ci/pytorch/common-build.sh
index 8ca9fdb34c77a..5ccf317711553 100644
--- a/.ci/pytorch/common-build.sh
+++ b/.ci/pytorch/common-build.sh
@@ -6,12 +6,6 @@ if [[ "$BUILD_ENVIRONMENT" != *win-* ]]; then
     # Save the absolute path in case later we chdir (as occurs in the gpu perf test)
     script_dir="$( cd "$(dirname "${BASH_SOURCE[0]}")" || exit ; pwd -P )"
 
-    if [[ "${BUILD_ENVIRONMENT}" == *-pch* ]]; then
-        # This is really weird, but newer sccache somehow produces broken binary
-        # see https://github.com/pytorch/pytorch/issues/139188
-        sudo mv /opt/cache/bin/sccache-0.2.14a /opt/cache/bin/sccache
-    fi
-
     if which sccache > /dev/null; then
         # Clear SCCACHE_BUCKET and SCCACHE_REGION if they are empty, otherwise
         # sccache will complain about invalid bucket configuration
diff --git a/.ci/pytorch/common.sh b/.ci/pytorch/common.sh
index 072b8da9b10c6..94d9629eac519 100644
--- a/.ci/pytorch/common.sh
+++ b/.ci/pytorch/common.sh
@@ -5,8 +5,8 @@
 source "$(dirname "${BASH_SOURCE[0]}")/common_utils.sh"
 set -ex -o pipefail
 
-# for ROCm environment variables
-if [[ "${BUILD_ENVIRONMENT}" == *rocm* ]]; then
+# Source ROCm environment variables (paths may vary between tarball/wheel installs)
+if [[ "${BUILD_ENVIRONMENT}" == *rocm* ]] && [[ -f /etc/rocm_env.sh ]]; then
   # shellcheck disable=SC1091
   source /etc/rocm_env.sh
 fi
@@ -14,6 +14,19 @@ fi
 # Required environment variables:
 #   $BUILD_ENVIRONMENT (should be set by your Docker image)
 
+# Select compiler based on build environment name. Images that have both
+# GCC and Clang installed default cc/c++ to Clang (via install_clang.sh),
+# so we need to override when a gcc build is requested.
+if [[ "${BUILD_ENVIRONMENT}" == *clang* ]]; then
+  export CC=clang
+  export CXX=clang++
+elif [[ "${BUILD_ENVIRONMENT}" == *gcc* ]]; then
+  export CC=gcc
+  export CXX=g++
+  sudo update-alternatives --install /usr/bin/cc cc /usr/bin/gcc 100
+  sudo update-alternatives --install /usr/bin/c++ c++ /usr/bin/g++ 100
+fi
+
 # Figure out which Python to use for ROCm
 if [[ "${BUILD_ENVIRONMENT}" == *rocm* ]]; then
   # HIP_PLATFORM is auto-detected by hipcc; unset to avoid build errors
diff --git a/.ci/pytorch/common_utils.sh b/.ci/pytorch/common_utils.sh
index c4a92c997561e..354841db899f8 100644
--- a/.ci/pytorch/common_utils.sh
+++ b/.ci/pytorch/common_utils.sh
@@ -127,17 +127,6 @@ function get_exit_code() {
   return $retcode
 }
 
-function get_bazel() {
-  # Download and use the cross-platform, dependency-free Python
-  # version of Bazelisk to fetch the platform specific version of
-  # Bazel to use from .bazelversion.
-  retry curl --location --output tools/bazel \
-    https://raw.githubusercontent.com/bazelbuild/bazelisk/v1.23.0/bazelisk.py
-  shasum --algorithm=1 --check \
-    <(echo '01df9cf7f08dd80d83979ed0d0666a99349ae93c  tools/bazel')
-  chmod u+x tools/bazel
-}
-
 function install_monkeytype {
   # Install MonkeyType
   pip_install MonkeyType
@@ -290,7 +279,7 @@ function install_torchrec_and_fbgemm() {
 
 function clone_pytorch_xla() {
   if [[ ! -d ./xla ]]; then
-    git clone --recursive --quiet https://github.com/pytorch/xla.git
+    git clone --recursive -b r2.12 https://github.com/pytorch/xla.git
     pushd xla
     # pin the xla hash so that we don't get broken by changes to xla
     git checkout "$(cat ../.github/ci_commit_pins/xla.txt)"
@@ -307,25 +296,9 @@ function install_torchao() {
 }
 
 function install_flash_attn_cute() {
-  echo "Installing FlashAttention CuTe from GitHub..."
-  # Grab latest main til we have a pinned commit
-  local flash_attn_commit
-  flash_attn_commit=$(git ls-remote https://github.com/Dao-AILab/flash-attention.git HEAD | cut -f1)
-
-  # Clone the repo to a temporary directory
-  rm -rf flash-attention-build
-  git clone --depth 1 --recursive https://github.com/Dao-AILab/flash-attention.git flash-attention-build
-
-  pushd flash-attention-build
-  git checkout "${flash_attn_commit}"
-
-  # Install only the 'cute' sub-directory
-  pip_install -e flash_attn/cute/
-  popd
-
-  # remove the local repo
-  rm -rf flash-attention-build
-  echo "FlashAttention CuTe installation complete."
+  echo "Installing FlashAttention 4 from PyPI..."
+  pip_install flash-attn-4==4.0.0b5
+  echo "FlashAttention 4 installation complete."
 }
 
 function install_cutlass_dsl() {
@@ -367,7 +340,7 @@ function install_cutlass_api() {
   git checkout "${cutlass_commit}"
 
   # Install cutlass_api with torch extras
-  pip_install -e "python/cutlass_api[torch]"
+  pip_install "python/cutlass_api[torch]"
   popd
 
   rm -rf cutlass-build
diff --git a/.ci/pytorch/cpp_doc_push_script.sh b/.ci/pytorch/cpp_doc_push_script.sh
index f085fa78bebe9..d0b4fd38826fa 100755
--- a/.ci/pytorch/cpp_doc_push_script.sh
+++ b/.ci/pytorch/cpp_doc_push_script.sh
@@ -1,7 +1,7 @@
 #!/bin/bash
 
 # This is where the local pytorch install in the docker image is located
-pt_checkout="/var/lib/jenkins/workspace"
+pt_checkout="${GITHUB_WORKSPACE:-/var/lib/jenkins/workspace}"
 
 # Since we're cat-ing this file, we need to escape all $'s
 echo "cpp_doc_push_script.sh: Invoked with $*"
@@ -60,6 +60,34 @@ time python tools/setup_helpers/generate_code.py \
 pushd docs/cpp
 time make VERBOSE=1 html
 
+# Run C++ API coverage check (allowlist-based + HTML formatting)
+echo "Running C++ docs coverage check..."
+python check_coverage.py --coverxygen || coverage_exit=$?
+
+# Generate coverxygen HTML report if coverxygen produced output
+if [ -f coverxygen.info ] && command -v genhtml &> /dev/null; then
+  genhtml --no-function-coverage coverxygen.info -o build/html/_coverage \
+    --title "PyTorch C++ API Doc Coverage" \
+    --legend --highlight 2>/dev/null || true
+fi
+
+# Copy coverage reports into the build output so they get uploaded
+mkdir -p build/html/_coverage
+cp -f cpp_coverage.txt cpp_html_issues.txt build/html/_coverage/ 2>/dev/null || true
+cp -f coverxygen.info build/html/_coverage/ 2>/dev/null || true
+
+if [ "${coverage_exit:-0}" -ne 0 ]; then
+  echo ""
+  echo "========================================"
+  echo "C++ DOCS COVERAGE: HIGH-PRIORITY GAPS"
+  echo "========================================"
+  echo ""
+  cat cpp_coverage.txt
+  echo ""
+  echo "See the full coverage report at: _coverage/cpp_coverage.txt"
+  echo "See the HTML issues report at: _coverage/cpp_html_issues.txt"
+fi
+
 popd
 popd
 
@@ -76,6 +104,7 @@ cp -r "${pt_checkout}"/docs/cpp/build/html/* .
 # Copy back _config.yml
 rm -rf _config.yml
 mv /tmp/cppdocs-sync/* .
+touch .nojekyll
 
 # Make a new commit
 git add . || true
diff --git a/.ci/pytorch/macos-build.sh b/.ci/pytorch/macos-build.sh
index f2bcb486cf95b..3259c62149b61 100755
--- a/.ci/pytorch/macos-build.sh
+++ b/.ci/pytorch/macos-build.sh
@@ -48,7 +48,7 @@ if [[ ${BUILD_ENVIRONMENT} == *"distributed"* ]]; then
 else
   # Explicitly set USE_DISTRIBUTED=0 to align with the default build config on mac. This also serves as the sole CI config that tests
   # that building with USE_DISTRIBUTED=0 works at all. See https://github.com/pytorch/pytorch/issues/86448
-  USE_DISTRIBUTED=0 USE_OPENMP=1 MACOSX_DEPLOYMENT_TARGET=11.0 WERROR=1 BUILD_TEST=OFF USE_PYTORCH_METAL=1 python -m build --wheel --no-isolation -C--build-option=--plat-name=macosx_11_0_arm64
+  USE_DISTRIBUTED=0 USE_OPENMP=1 WERROR=1 BUILD_TEST=OFF USE_PYTORCH_METAL=1 python -m build --wheel --no-isolation
 fi
 if which sccache > /dev/null; then
   print_sccache_stats
diff --git a/.ci/pytorch/macos-common.sh b/.ci/pytorch/macos-common.sh
index 6826a52577a29..d21ee458f6c99 100755
--- a/.ci/pytorch/macos-common.sh
+++ b/.ci/pytorch/macos-common.sh
@@ -9,6 +9,6 @@ sysctl -a | grep machdep.cpu
 
 # These are required for both the build job and the test job.
 # In the latter to test cpp extensions.
-export MACOSX_DEPLOYMENT_TARGET=11.1
+export MACOSX_DEPLOYMENT_TARGET=14.0
 export CXX=clang++
 export CC=clang
diff --git a/.ci/pytorch/macos-test.sh b/.ci/pytorch/macos-test.sh
index f6173f64224e2..332a52b11b10a 100755
--- a/.ci/pytorch/macos-test.sh
+++ b/.ci/pytorch/macos-test.sh
@@ -49,6 +49,8 @@ test_python_mps() {
 test_python_openreg() {
   setup_test_python
 
+  git submodule update --init --depth 1 third_party/googletest
+
   time python test/run_test.py --openreg --verbose
 
   assert_git_not_dirty
diff --git a/.ci/pytorch/python_doc_push_script.sh b/.ci/pytorch/python_doc_push_script.sh
index 6bcd46c4815a6..dde476a9392ee 100755
--- a/.ci/pytorch/python_doc_push_script.sh
+++ b/.ci/pytorch/python_doc_push_script.sh
@@ -1,7 +1,7 @@
 #!/bin/bash
 
 # This is where the local pytorch install in the docker image is located
-pt_checkout="/var/lib/jenkins/workspace"
+pt_checkout="${GITHUB_WORKSPACE:-/var/lib/jenkins/workspace}"
 
 source "$pt_checkout/.ci/pytorch/common_utils.sh"
 
@@ -50,8 +50,14 @@ echo "install_path: $install_path  version: $version"
 
 build_docs () {
   set +e
-  set -o pipefail
-  make "$1" 2>&1 | tee /tmp/docs_build.txt
+  # Don't pipe through tee: sphinx -j auto forks workers that inherit
+  # the pipe fd and hold it open after sphinx exits, causing tee to
+  # block forever. Write to a file and tail with --pid so it exits
+  # (after draining) when make finishes.
+  make "$1" > /tmp/docs_build.txt 2>&1 &
+  local make_pid=$!
+  tail -f --pid=$make_pid /tmp/docs_build.txt
+  wait $make_pid
   code=$?
   if [ $code -ne 0 ]; then
     set +x
@@ -87,44 +93,36 @@ pushd docs
 if [ "$is_main_doc" = true ]; then
   build_docs html || exit $?
 
-  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 |
-  #
-  # 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
-    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
+  # Coverage check is only needed on PR builds (push=false) to catch
+  # undocumented APIs early. Nightly/release builds (push=true) skip it
+  # since PRs already enforce coverage.
+  if [[ "${WITH_PUSH:-}" != true ]]; then
+    SPHINXOPTS="-WT --keep-going" make coverage
+
+    undocumented=$(grep "| TOTAL" build/coverage/python.txt | awk -F'|' '{print $4}' | tr -d ' ')
+
+    if [ -z "$undocumented" ] || ! [[ "$undocumented" =~ ^[0-9]+$ ]]; then
+      echo coverage output not found
+      exit 1
+    elif [ "$undocumented" -gt 0 ]; then
+      set +x
+      echo ""
+      echo "====================="
+      echo "UNDOCUMENTED OBJECTS:"
+      echo "====================="
+      echo ""
+      total_line=$(grep -n "| TOTAL" build/coverage/python.txt | cut -d: -f1)
+      if [ -n "$total_line" ]; then
+        tail -n +$((total_line + 2)) build/coverage/python.txt
+      else
+        cat build/coverage/python.txt
+      fi
+      echo ""
+      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
+      exit 1
     fi
-    echo ""
-    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
-    exit 1
   fi
 else
   # skip coverage, format for stable or tags
diff --git a/.ci/pytorch/smoke_test/check_binary_symbols.py b/.ci/pytorch/smoke_test/check_binary_symbols.py
index 95f1e51a88a13..4a49be6db5730 100755
--- a/.ci/pytorch/smoke_test/check_binary_symbols.py
+++ b/.ci/pytorch/smoke_test/check_binary_symbols.py
@@ -1,4 +1,6 @@
 #!/usr/bin/env python3
+from __future__ import annotations
+
 import concurrent.futures
 import distutils.sysconfig
 import functools
diff --git a/.ci/pytorch/smoke_test/check_wheel_tags.py b/.ci/pytorch/smoke_test/check_wheel_tags.py
new file mode 100644
index 0000000000000..901304657b531
--- /dev/null
+++ b/.ci/pytorch/smoke_test/check_wheel_tags.py
@@ -0,0 +1,269 @@
+"""Validate wheel platform tags and macOS dylib minos.
+Supports two modes:
+1. Pre-install: reads .whl files from PYTORCH_FINAL_PACKAGE_DIR
+2. Post-install: reads metadata from installed torch package (soft warnings)
+- (macOS only) dylib minos matches the wheel platform tag
+"""
+
+import os
+import platform
+import re
+import subprocess
+import sys
+import tempfile
+import zipfile
+from pathlib import Path
+
+
+EXPECTED_PLATFORM_TAGS: dict[str, str] = {
+    "linux": r"_x86_64$",
+    "linux-aarch64": r"_aarch64$",
+    "windows": r"^win_amd64$",
+    "win32": r"^win_amd64$",
+    "macos-arm64": r"^macosx_\d+_\d+_arm64$",
+    "darwin": r"^macosx_\d+_\d+_(arm64|x86_64)$",
+}
+
+
+def _extract_wheel_tags(whl_path: Path) -> list[str]:
+    """Extract Tag values from the WHEEL metadata file inside a .whl archive."""
+    tags = []
+    with zipfile.ZipFile(whl_path, "r") as zf:
+        wheel_files = [n for n in zf.namelist() if n.endswith("/WHEEL")]
+        if not wheel_files:
+            return tags
+        content = zf.read(wheel_files[0]).decode("utf-8")
+        for line in content.splitlines():
+            if line.startswith("Tag:"):
+                tags.append(line.split(":", 1)[1].strip())
+    return tags
+
+
+def _extract_installed_wheel_tags(package: str = "torch") -> list[str]:
+    """Extract Tag values from an installed package's WHEEL metadata."""
+    from importlib.metadata import distribution
+
+    dist = distribution(package)
+    wheel_text = dist.read_text("WHEEL")
+    if not wheel_text:
+        return []
+    tags = []
+    for line in wheel_text.splitlines():
+        if line.startswith("Tag:"):
+            tags.append(line.split(":", 1)[1].strip())
+    return tags
+
+
+def check_wheel_platform_tag() -> None:
+    """Validate that wheel Tags in WHEEL metadata match the expected platform.
+
+    Mode 1: PYTORCH_FINAL_PACKAGE_DIR set → read .whl file (strict, raises on mismatch)
+    Mode 2: No wheel dir → read from installed torch package (soft, prints warnings)
+    """
+    wheel_dir = os.getenv("PYTORCH_FINAL_PACKAGE_DIR", "")
+
+    target_os = os.getenv("TARGET_OS", sys.platform)
+    if target_os == "linux" and platform.machine() == "aarch64":
+        target_os = "linux-aarch64"
+    expected_python = f"cp{sys.version_info.major}{sys.version_info.minor}"
+    import sysconfig
+
+    abiflags = getattr(sys, "abiflags", "")
+    if not abiflags and (
+        os.getenv("MATRIX_PYTHON_VERSION", "").endswith("t")
+        or bool(sysconfig.get_config_var("Py_GIL_DISABLED"))
+        or not getattr(sys, "_is_gil_enabled", lambda: True)()
+    ):
+        abiflags = "t"
+    expected_abi = f"cp{sys.version_info.major}{sys.version_info.minor}{abiflags}"
+    print(f"Expected ABI tag: {expected_abi}")
+
+    platform_pattern = EXPECTED_PLATFORM_TAGS.get(target_os)
+    if not platform_pattern:
+        print(
+            f"No expected platform pattern for TARGET_OS={target_os}, "
+            "skipping wheel tag check"
+        )
+        return
+
+    # Mode 1: Read from .whl file
+    if wheel_dir and os.path.isdir(wheel_dir):
+        whls = list(Path(wheel_dir).glob("torch-*.whl"))
+        if not whls:
+            print(f"No torch wheel found in {wheel_dir}, skipping wheel tag check")
+            return
+        if len(whls) > 1:
+            raise RuntimeError(
+                f"Expected exactly one torch wheel in {wheel_dir}, "
+                f"found {len(whls)}: {[w.name for w in whls]}"
+            )
+        whl = whls[0]
+        print(f"Checking wheel platform tag for: {whl.name}")
+        tags = _extract_wheel_tags(whl)
+        source = whl.name
+    else:
+        # Mode 2: Read from installed package (soft)
+        print("PYTORCH_FINAL_PACKAGE_DIR not set, reading from installed torch package")
+        try:
+            tags = _extract_installed_wheel_tags("torch")
+            source = "installed torch"
+        except Exception as e:
+            print(f"Could not read installed torch metadata: {e}, skipping")
+            return
+
+    if not tags:
+        raise RuntimeError(f"No Tag found in WHEEL metadata of {source}")
+
+    for tag_str in tags:
+        parts = tag_str.split("-")
+        if len(parts) != 3:
+            msg = (
+                f"Malformed wheel tag '{tag_str}' in {source}, "
+                f"expected format: <python>-<abi>-<platform>"
+            )
+            raise RuntimeError(msg)
+
+        python_tag, abi_tag, platform_tag = parts
+
+        print(f"Checking tag: {tag_str} (from {source})")
+        if python_tag != expected_python:
+            msg: str = (
+                f"Python tag mismatch in {source}: "
+                f"got '{python_tag}', expected '{expected_python}'"
+            )
+            raise RuntimeError(msg)
+
+        if abi_tag != expected_abi:
+            msg = (
+                f"ABI tag mismatch in {source}: "
+                f"got '{abi_tag}', expected '{expected_abi}'"
+            )
+            raise RuntimeError(msg)
+
+        if not re.search(platform_pattern, platform_tag):
+            msg = (
+                f"Platform tag mismatch in {source}: "
+                f"got '{platform_tag}', expected pattern matching "
+                f"'{platform_pattern}' for TARGET_OS={target_os}"
+            )
+            raise RuntimeError(msg)
+
+    print(f"OK: Wheel tag(s) valid for {source}: {', '.join(tags)}")
+
+
+def _check_dylibs_minos(dylibs: list, expected_minos: str, source: str) -> None:
+    mismatches = []
+    for dylib in dylibs:
+        try:
+            result = subprocess.run(
+                ["otool", "-l", str(dylib)],
+                capture_output=True,
+                text=True,
+                timeout=30,
+            )
+        except Exception:
+            continue
+
+        minos = None
+        lines = result.stdout.splitlines()
+        for i, line in enumerate(lines):
+            s = line.strip()
+            if "LC_BUILD_VERSION" in s:
+                for j in range(i + 1, min(i + 6, len(lines))):
+                    if lines[j].strip().startswith("minos"):
+                        minos = lines[j].strip().split()[1]
+                        break
+                break
+            if "LC_VERSION_MIN_MACOSX" in s:
+                for j in range(i + 1, min(i + 4, len(lines))):
+                    if lines[j].strip().startswith("version"):
+                        minos = lines[j].strip().split()[1]
+                        break
+                break
+
+        # A dylib with a lower minos than the wheel tag is safe (forward compatible).
+        # Only flag dylibs that require a *higher* macOS than the wheel claims to support.
+        if minos and tuple(int(x) for x in minos.split(".")) > tuple(
+            int(x) for x in expected_minos.split(".")
+        ):
+            mismatches.append(
+                f"{dylib.name}: minos={minos}, expected<={expected_minos}"
+            )
+
+    if mismatches:
+        raise RuntimeError(
+            f"minos/platform tag mismatch in {len(mismatches)} dylib(s):\n"
+            + "\n".join(f"  {m}" for m in mismatches)
+        )
+    print(
+        f"OK: All {len(dylibs)} dylib(s) have minos matching "
+        f"platform tag ({expected_minos}) for {source}"
+    )
+
+
+def check_mac_wheel_minos() -> None:
+    if sys.platform != "darwin":
+        return
+
+    wheel_dir = os.getenv("PYTORCH_FINAL_PACKAGE_DIR", "")
+
+    if wheel_dir and os.path.isdir(wheel_dir):
+        # Mode 1: extract dylibs from .whl file
+        whls = list(Path(wheel_dir).glob("*.whl"))
+        if not whls:
+            print(f"No .whl files in {wheel_dir}, skipping wheel minos check")
+            return
+
+        macos_whl_re = re.compile(r"macosx_(\d+)_(\d+)_(\w+)\.whl$")
+        for whl in whls:
+            print(f"Checking wheel tag minos for: {whl.name}")
+            m = macos_whl_re.search(whl.name)
+            if not m:
+                print(f"No macOS platform tag in {whl.name}, skipping")
+                continue
+            expected_minos = f"{m.group(1)}.{m.group(2)}"
+
+            with tempfile.TemporaryDirectory() as tmpdir:
+                with zipfile.ZipFile(whl, "r") as zf:
+                    dylib_names = [n for n in zf.namelist() if n.endswith(".dylib")]
+                    if not dylib_names:
+                        print("No .dylib files in wheel, skipping minos check")
+                        continue
+                    for name in dylib_names:
+                        zf.extract(name, tmpdir)
+                dylibs = list(Path(tmpdir).rglob("*.dylib"))
+                _check_dylibs_minos(dylibs, expected_minos, whl.name)
+    else:
+        # Mode 2: read from installed torch package
+        print("PYTORCH_FINAL_PACKAGE_DIR not set, checking installed torch dylibs")
+        try:
+            tags = _extract_installed_wheel_tags("torch")
+        except Exception as e:
+            print(f"Could not read installed torch metadata: {e}, skipping")
+            return
+
+        expected_minos = None
+        for tag_str in tags:
+            m = re.search(r"macosx_(\d+)_(\d+)_\w+", tag_str)
+            if m:
+                expected_minos = f"{m.group(1)}.{m.group(2)}"
+                break
+
+        if not expected_minos:
+            print("No macOS platform tag found in installed torch metadata, skipping")
+            return
+
+        print(f"Expected minos from installed wheel tag: {expected_minos}")
+
+        import torch
+
+        torch_dir = Path(torch.__file__).parent
+        dylibs = list(torch_dir.rglob("*.dylib"))
+        if not dylibs:
+            raise RuntimeError("No .dylib files found in installed torch")
+        _check_dylibs_minos(dylibs, expected_minos, "installed torch")
+
+
+if __name__ == "__main__":
+    check_wheel_platform_tag()
+    check_mac_wheel_minos()
diff --git a/.ci/pytorch/smoke_test/smoke_test.py b/.ci/pytorch/smoke_test/smoke_test.py
index d916f6b49a2f3..877218e71e307 100644
--- a/.ci/pytorch/smoke_test/smoke_test.py
+++ b/.ci/pytorch/smoke_test/smoke_test.py
@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import argparse
 import importlib
 import json
@@ -7,7 +9,8 @@
 import sys
 from pathlib import Path
 from tempfile import NamedTemporaryFile
-from typing import Optional
+
+from check_wheel_tags import check_mac_wheel_minos, check_wheel_platform_tag
 
 import torch
 import torch._dynamo
@@ -24,6 +27,7 @@
 package_type = os.getenv("MATRIX_PACKAGE_TYPE")
 target_os = os.getenv("TARGET_OS", sys.platform)
 BASE_DIR = Path(__file__).parent.parent.parent
+PYTORCH_ROOT = BASE_DIR.parent
 
 is_cuda_system = gpu_arch_type == "cuda"
 NIGHTLY_ALLOWED_DELTA = 3
@@ -205,7 +209,7 @@ def test_cuda_gds_errors_captured() -> None:
         )
 
 
-def find_pypi_package_version(package: str) -> Optional[str]:
+def find_pypi_package_version(package: str) -> str | None:
     from importlib import metadata
 
     dists = metadata.distributions()
@@ -215,6 +219,89 @@ def find_pypi_package_version(package: str) -> Optional[str]:
     return None
 
 
+def get_expected_cudnn_version_linux(cuda_version: str) -> str | None:
+    """Parse expected cuDNN version from generate_binary_build_matrix.py for Linux.
+
+    Reads PYTORCH_EXTRA_INSTALL_REQUIREMENTS and extracts the cudnn version
+    for the given CUDA version (e.g. "12.6").
+    """
+    matrix_script = (
+        PYTORCH_ROOT / ".github" / "scripts" / "generate_binary_build_matrix.py"
+    )
+    if not matrix_script.exists():
+        print(f"Warning: {matrix_script} not found, skipping cuDNN version check")
+        return None
+
+    content = matrix_script.read_text()
+    # Match the full cudnn package version like nvidia-cudnn-cu12==9.10.2.21
+    # and extract major.minor.patch (dropping the build number)
+    pattern = (
+        rf'"{re.escape(cuda_version)}":\s*\(\s*'
+        r"[\s\S]*?nvidia-cudnn-cu\d+==(\d+\.\d+\.\d+)\.\d+"
+    )
+    match = re.search(pattern, content)
+    if match:
+        return match.group(1)
+    return None
+
+
+def get_expected_cudnn_version_windows(cuda_version: str) -> str | None:
+    """Parse expected cuDNN version from cuda_install.bat for Windows.
+
+    Reads the batch file and extracts EXPECTED_CUDNN_VERSION for the given
+    CUDA version (e.g. "12.6" maps to CUDA_VER 126).
+    """
+    bat_file = (
+        PYTORCH_ROOT / ".ci" / "pytorch" / "windows" / "internal" / "cuda_install.bat"
+    )
+    if not bat_file.exists():
+        print(f"Warning: {bat_file} not found, skipping cuDNN version check")
+        return None
+
+    content = bat_file.read_text()
+    # Convert "12.6" to "126" to match batch file's CUDA_VER format
+    cuda_ver_nodot = cuda_version.replace(".", "")
+    # Match: if %CUDA_VER% EQU 126 ( ... set EXPECTED_CUDNN_VERSION=9.10.2 )
+    pattern = (
+        rf"if %CUDA_VER% EQU {re.escape(cuda_ver_nodot)}\s*\("
+        r"[\s\S]*?set EXPECTED_CUDNN_VERSION=(\d+\.\d+\.\d+)"
+    )
+    match = re.search(pattern, content)
+    if match:
+        return match.group(1)
+    return None
+
+
+def check_cudnn_version(cuda_version: str, actual_cudnn_version: str) -> None:
+    """Validate cuDNN version matches expected version from build config files."""
+    if sys.platform in ["linux", "linux2"]:
+        expected = get_expected_cudnn_version_linux(cuda_version)
+        source = "generate_binary_build_matrix.py"
+    elif sys.platform == "win32":
+        expected = get_expected_cudnn_version_windows(cuda_version)
+        source = "cuda_install.bat"
+    else:
+        print(f"cuDNN version check not supported on platform {sys.platform}")
+        return
+
+    if expected is None:
+        print(
+            f"Warning: Could not determine expected cuDNN version for CUDA {cuda_version} "
+            f"from {source}, skipping validation"
+        )
+        return
+
+    if not actual_cudnn_version.startswith(expected):
+        raise RuntimeError(
+            f"cuDNN version mismatch for CUDA {cuda_version}. "
+            f"Loaded: {actual_cudnn_version} Expected: {expected} (from {source})"
+        )
+    print(
+        f"cuDNN version check passed: {actual_cudnn_version} matches "
+        f"expected {expected} from {source}"
+    )
+
+
 def cudnn_to_version_str(cudnn_version: int) -> str:
     patch = int(cudnn_version % 10)
     minor = int((cudnn_version / 100) % 100)
@@ -281,6 +368,20 @@ 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}"
+            )
+
+        check_cudnn_version(gpu_arch_ver, torch_cudnn_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}")
@@ -538,6 +639,9 @@ def main() -> None:
 
     smoke_test_nvshmem()
 
+    check_wheel_platform_tag()
+    check_mac_wheel_minos()
+
 
 if __name__ == "__main__":
     main()
diff --git a/.ci/pytorch/test.sh b/.ci/pytorch/test.sh
index 7bc94541a7558..0b10d1dda82b3 100755
--- a/.ci/pytorch/test.sh
+++ b/.ci/pytorch/test.sh
@@ -14,9 +14,10 @@ source "$(dirname "${BASH_SOURCE[0]}")/common.sh"
 # shellcheck source=./common-build.sh
 source "$(dirname "${BASH_SOURCE[0]}")/common-build.sh"
 
-# Do not change workspace permissions for ROCm and s390x CI jobs
-# as it can leave workspace with bad permissions for cancelled jobs
-if [[ "$BUILD_ENVIRONMENT" != *rocm* && "$BUILD_ENVIRONMENT" != *s390x* && -d /var/lib/jenkins/workspace ]]; then
+# Only change workspace permissions if passwordless sudo is available
+# (e.g. ROCm and s390x CI jobs lack it, and changing permissions
+# can leave the workspace in a bad state for cancelled jobs)
+if sudo -n true 2>/dev/null && [[ -d /var/lib/jenkins/workspace ]]; then
   # Workaround for dind-rootless userid mapping (https://github.com/pytorch/ci-infra/issues/96)
   WORKSPACE_ORIGINAL_OWNER_ID=$(stat -c '%u' "/var/lib/jenkins/workspace")
   cleanup_workspace() {
@@ -44,6 +45,16 @@ if [[ "$BUILD_ENVIRONMENT" == *cuda* ]]; then
   fi
 fi
 
+# Remove onnxruntime if present to avoid interference with non-ONNX tests
+if [[ "$TEST_CONFIG" != "onnx" ]]; then
+  pip uninstall -y onnxruntime 2>/dev/null || true
+fi
+
+# Remove dill to test that serialization works without it
+if [[ "$BUILD_ENVIRONMENT" == *py3.10-gcc11 ]]; then
+  pip uninstall -y dill 2>/dev/null || true
+fi
+
 echo "Environment variables:"
 env
 
@@ -129,9 +140,7 @@ if [[ "${PYTORCH_TEST_RERUN_DISABLED_TESTS}" == "1" ]] || [[ "${CONTINUE_THROUGH
 fi
 
 # Get fully qualified path using realpath
-if [[ "$BUILD_ENVIRONMENT" != *bazel* ]]; then
-  CUSTOM_TEST_ARTIFACT_BUILD_DIR=$(realpath "${CUSTOM_TEST_ARTIFACT_BUILD_DIR:-"build/custom_test_artifacts"}")
-fi
+CUSTOM_TEST_ARTIFACT_BUILD_DIR=$(realpath "${CUSTOM_TEST_ARTIFACT_BUILD_DIR:-"build/custom_test_artifacts"}")
 
 # Reduce set of tests to include when running run_test.py
 if [[ -n $TESTS_TO_INCLUDE ]]; then
@@ -144,11 +153,34 @@ env
 
 echo "Testing pytorch"
 
+# Set OMP_NUM_THREADS to nproc/4 on k8s ARC runners if not already set.
+#
+# We use nproc (cgroup-aware) rather than os.cpu_count() because on k8s (ARC)
+# pods, os.cpu_count() returns the host's CPU count (e.g., 192) rather than
+# the pod's cpuset allocation (e.g., 16).
+#
+# We use nproc/4 rather than nproc because OpenMP spin-waits at thread barriers.
+# When thread count equals cpuset size (e.g., 16 threads on 16 CPUs), spinning
+# barrier threads monopolize all CPUs and the OS must context-switch to let
+# actual work complete. This causes ~5000x slowdowns on small tensor ops
+# (e.g., aten::copy_ on 147KB: ~34ms instead of ~7us). Using nproc/4 leaves
+# headroom for the main thread and for NUM_PROCS=3 parallel test processes.
+if [[ -z "${OMP_NUM_THREADS:-}" ]] && [[ -n "${USE_ARC:-}" ]]; then
+  OMP_NUM_THREADS=$(( $(nproc) / 4 ))
+  # Floor of 4: low OMP_NUM_THREADS (1-2) changes floating-point reduction
+  # order, causing numerical mismatches in tests with tight tolerances
+  # (e.g., test_batchnorm_nhwc_cpu).
+  if [[ "$OMP_NUM_THREADS" -lt 4 ]]; then
+    OMP_NUM_THREADS=4
+  fi
+  export OMP_NUM_THREADS
+fi
+
 export LANG=C.UTF-8
 
 PR_NUMBER=${PR_NUMBER:-${CIRCLE_PR_NUMBER:-}}
 
-if [[ -d "${HF_CACHE}" ]]; then
+if [[ -d "${HF_CACHE}" && "$TEST_CONFIG" != "onnx" ]]; then
   export HF_HOME="${HF_CACHE}"
 fi
 
@@ -222,12 +254,10 @@ if [[ "$BUILD_ENVIRONMENT" == *xpu* ]]; then
   timeout 30 xpu-smi discovery || true
 fi
 
-if [[ "$BUILD_ENVIRONMENT" != *-bazel-* ]] ; then
-  # JIT C++ extensions require ninja (installed from requirements-ci.txt).
-  # ninja is installed in $HOME/.local/bin, e.g., /var/lib/jenkins/.local/bin for CI user jenkins
-  # but this script should be runnable by any user, including root
-  export PATH="$HOME/.local/bin:$PATH"
-fi
+# JIT C++ extensions require ninja (installed from requirements-ci.txt).
+# ninja is installed in $HOME/.local/bin, e.g., /var/lib/jenkins/.local/bin for CI user jenkins
+# but this script should be runnable by any user, including root
+export PATH="$HOME/.local/bin:$PATH"
 
 if [[ "$BUILD_ENVIRONMENT" == *aarch64* ]]; then
   # TODO: revisit this once the CI is stabilized on aarch64 linux
@@ -299,8 +329,7 @@ fi
 if [[ "$BUILD_ENVIRONMENT" == *-debug* ]]; then
     echo "We are in debug mode: $BUILD_ENVIRONMENT. Expect the python assertion to fail"
     (cd test && ! get_exit_code python -c "import torch; torch._C._crash_if_debug_asserts_fail(424242)")
-elif [[ "$BUILD_ENVIRONMENT" != *-bazel-* ]]; then
-    # Noop when debug is disabled. Skip bazel jobs because torch isn't available there yet.
+else
     echo "We are not in debug mode: $BUILD_ENVIRONMENT. Expect the assertion to pass"
     (cd test && python -c "import torch; torch._C._crash_if_debug_asserts_fail(424242)")
 fi
@@ -311,12 +340,6 @@ elif [[ $TEST_CONFIG == 'nogpu_AVX512' ]]; then
   export ATEN_CPU_CAPABILITY=avx2
 fi
 
-if [[ "${TEST_CONFIG}" == "legacy_nvidia_driver" ]]; then
-  # Make sure that CUDA can be initialized
-  (cd test && python -c "import torch; torch.rand(2, 2, device='cuda')")
-  export USE_LEGACY_DRIVER=1
-fi
-
 test_python_legacy_jit() {
   time python test/run_test.py --include test_jit_legacy test_jit_fuser_legacy --verbose
   assert_git_not_dirty
@@ -365,8 +388,28 @@ test_python_smoke_b200() {
       inductor/test_flex_flash \
       inductor/test_torchinductor \
       inductor/test_nv_universal_gemm \
-    $PYTHON_TEST_EXTRA_OPTION \
-    --upload-artifacts-while-running
+      inductor/test_fused_attention \
+      test_varlen_attention \
+      $PYTHON_TEST_EXTRA_OPTION \
+      --upload-artifacts-while-running
+  assert_git_not_dirty
+}
+
+
+test_python_smoke_xpu() {
+  # Smoke tests for XPU client
+  time python test/run_test.py --include test_transformers $PYTHON_TEST_EXTRA_OPTION --upload-artifacts-while-running
+  time test_xpu_sycl_tla_backend
+  assert_git_not_dirty
+}
+
+test_dtensor() {
+  # Dynamically discover all test files under test/distributed/tensor/
+  # so new tests are automatically picked up.
+  # shellcheck disable=SC2046
+  time python test/run_test.py \
+    --include $(find test/distributed/tensor -name 'test_*.py' -printf '%P\n' | sed 's|\.py$||; s|^|distributed/tensor/|' | sort | tr '\n' ' ') \
+    --verbose $PYTHON_TEST_EXTRA_OPTION --upload-artifacts-while-running
   assert_git_not_dirty
 }
 
@@ -393,6 +436,7 @@ test_h100_symm_mem() {
   export NVSHMEM_SYMMETRIC_SIZE=4G
   # Disable NVLink Switch features (not available on AWS H100 instances)
   export NVSHMEM_DISABLE_NVLS=1
+  export NCCL_NVLS_ENABLE=0
   _run_symm_mem_tests
 }
 
@@ -402,10 +446,20 @@ test_b200_symm_mem() {
 
 test_h100_cutlass_backend() {
   # cutlass backend tests for H100
+  git submodule update --init --depth 1 third_party/cutlass
   TORCHINDUCTOR_CUTLASS_DIR=$(realpath "./third_party/cutlass") python test/run_test.py --include inductor/test_cutlass_backend -k "not addmm" $PYTHON_TEST_EXTRA_OPTION --upload-artifacts-while-running
   TORCHINDUCTOR_CUTLASS_DIR=$(realpath "./third_party/cutlass") python test/run_test.py --include inductor/test_cutlass_evt $PYTHON_TEST_EXTRA_OPTION --upload-artifacts-while-running
 }
 
+test_xpu_sycl_tla_backend() {
+  # Inductor sycl-tla backend tests for XPU
+  # shellcheck disable=SC1091
+  source  /opt/intel/oneapi/mkl/latest/env/vars.sh
+  sycl_tla_dir=$(realpath "./third_party/sycl-tla")
+  rm -rf "${sycl_tla_dir}" && git clone --depth 1 --single-branch -b v0.8 --quiet https://github.com/intel/sycl-tla.git "${sycl_tla_dir}"
+  TORCHINDUCTOR_CUTLASS_DIR=$(realpath "./third_party/sycl-tla") python test/run_test.py --include inductor/test_cutlass_backend -k "not addmm" $PYTHON_TEST_EXTRA_OPTION --upload-artifacts-while-running
+}
+
 test_lazy_tensor_meta_reference_disabled() {
   export TORCH_DISABLE_FUNCTIONALIZATION_META_REFERENCE=1
   echo "Testing lazy tensor operations without meta reference"
@@ -422,12 +476,16 @@ test_dynamo_core() {
 }
 
 test_dynamo_cpython() {
+  # Disable TD for cpython since it's pretty cheap to run the cpython tests (< 10 min)
+  # and if TD is enabled, only 25% of the tests will be executed
+  export NO_TD=1
   time python test/run_test.py \
     --include-cpython-tests \
     --dynamo \
     --verbose \
     --upload-artifacts-while-running
   assert_git_not_dirty
+  unset NO_TD
 }
 
 test_dynamo_wrapped_shard() {
@@ -452,6 +510,8 @@ test_dynamo_wrapped_shard() {
 }
 
 test_einops() {
+  pip install einops==0.5.0
+  time python test/run_test.py --einops --verbose --upload-artifacts-while-running
   pip install einops==0.6.1
   time python test/run_test.py --einops --verbose --upload-artifacts-while-running
   pip install einops==0.7.0
@@ -532,7 +592,7 @@ test_inductor_shard() {
 
   # Do not add --inductor for the following inductor unit tests, otherwise we will fail because of nested dynamo state
   python test/run_test.py \
-    --include inductor/test_torchinductor inductor/test_torchinductor_opinfo inductor/test_aot_inductor \
+    --include inductor/test_torchinductor inductor/test_torchinductor_opinfo inductor/test_aot_inductor inductor/test_cpu_select_algorithm \
     --shard "$1" "$NUM_TEST_SHARDS" \
     --verbose
 }
@@ -590,7 +650,7 @@ test_inductor_cpp_wrapper_shard() {
     --shard "$1" "$NUM_TEST_SHARDS" \
     --verbose
   python test/run_test.py \
-    --include inductor/test_torchinductor inductor/test_max_autotune inductor/test_cpu_repro \
+    --include inductor/test_torchinductor inductor/test_max_autotune inductor/test_cpu_repro inductor/test_triton_kernels \
     --shard "$1" "$NUM_TEST_SHARDS" \
     --verbose
   python test/run_test.py --inductor \
@@ -598,7 +658,12 @@ test_inductor_cpp_wrapper_shard() {
     -k 'take' \
     --shard "$1" "$NUM_TEST_SHARDS" \
     --verbose
-
+  # Keep testing TORCHINDUCTOR_AUTOTUNE_AT_COMPILE_TIME=1 for the near future.
+  # Will drop this after AOTInductor also switches to lazy Triton compilation.
+  TORCHINDUCTOR_AUTOTUNE_AT_COMPILE_TIME=1 python test/run_test.py \
+    --include inductor/test_torchinductor inductor/test_triton_kernels inductor/test_max_autotune \
+    --shard "$1" "$NUM_TEST_SHARDS" \
+    --verbose
   if [[ "${BUILD_ENVIRONMENT}" == *xpu* ]]; then
     python test/run_test.py \
       --include inductor/test_mkldnn_pattern_matcher \
@@ -712,6 +777,10 @@ test_perf_for_dashboard() {
   TEST_REPORTS_DIR=$(pwd)/test/test-reports
   mkdir -p "$TEST_REPORTS_DIR"
 
+  if [[ "${EXPORT_PROFILER_TRACE:-0}" == "1" ]]; then
+    mkdir -p "$TEST_REPORTS_DIR/profiler_traces"
+  fi
+
   local suite="$1"
   shift
 
@@ -771,52 +840,97 @@ test_perf_for_dashboard() {
       fi
 
       if [[ "$DASHBOARD_TAG" == *default-true* ]]; then
+        local profiler_trace_flags=()
+        if [[ "${EXPORT_PROFILER_TRACE:-0}" == "1" && "$target" == "performance" ]]; then
+          profiler_trace_flags=(--export-profiler-trace --profiler-trace-name "$TEST_REPORTS_DIR/profiler_traces/${backend}_no_cudagraphs_${suite}_${dtype}_${mode}_${device}")
+        fi
         $TASKSET python "benchmarks/dynamo/$suite.py" \
             "${target_flag[@]}" --"$mode" --"$dtype" --backend "$backend" --disable-cudagraphs "$@" \
+            "${profiler_trace_flags[@]}" \
             --output "$TEST_REPORTS_DIR/${backend}_no_cudagraphs_${suite}_${dtype}_${mode}_${device}_${target}.csv"
       fi
       if [[ "$DASHBOARD_TAG" == *cudagraphs-true* ]]; then
+        local profiler_trace_flags=()
+        if [[ "${EXPORT_PROFILER_TRACE:-0}" == "1" && "$target" == "performance" ]]; then
+          profiler_trace_flags=(--export-profiler-trace --profiler-trace-name "$TEST_REPORTS_DIR/profiler_traces/${backend}_with_cudagraphs_${suite}_${dtype}_${mode}_${device}")
+        fi
         $TASKSET python "benchmarks/dynamo/$suite.py" \
             "${target_flag[@]}" --"$mode" --"$dtype" --backend "$backend" "$@" \
+            "${profiler_trace_flags[@]}" \
             --output "$TEST_REPORTS_DIR/${backend}_with_cudagraphs_${suite}_${dtype}_${mode}_${device}_${target}.csv"
       fi
       if [[ "$DASHBOARD_TAG" == *dynamic-true* ]]; then
+        local profiler_trace_flags=()
+        if [[ "${EXPORT_PROFILER_TRACE:-0}" == "1" && "$target" == "performance" ]]; then
+          profiler_trace_flags=(--export-profiler-trace --profiler-trace-name "$TEST_REPORTS_DIR/profiler_traces/${backend}_dynamic_${suite}_${dtype}_${mode}_${device}")
+        fi
         $TASKSET python "benchmarks/dynamo/$suite.py" \
             "${target_flag[@]}" --"$mode" --"$dtype" --backend "$backend" --dynamic-shapes \
             --dynamic-batch-only "$@" \
+            "${profiler_trace_flags[@]}" \
             --output "$TEST_REPORTS_DIR/${backend}_dynamic_${suite}_${dtype}_${mode}_${device}_${target}.csv"
       fi
       if [[ "$DASHBOARD_TAG" == *cppwrapper-true* ]]; then
+        local profiler_trace_flags=()
+        if [[ "${EXPORT_PROFILER_TRACE:-0}" == "1" && "$target" == "performance" ]]; then
+          profiler_trace_flags=(--export-profiler-trace --profiler-trace-name "$TEST_REPORTS_DIR/profiler_traces/${backend}_cpp_wrapper_${suite}_${dtype}_${mode}_${device}")
+        fi
         TORCHINDUCTOR_CPP_WRAPPER=1 $TASKSET python "benchmarks/dynamo/$suite.py" \
             "${target_flag[@]}" --"$mode" --"$dtype" --backend "$backend" --disable-cudagraphs "$@" \
+            "${profiler_trace_flags[@]}" \
             --output "$TEST_REPORTS_DIR/${backend}_cpp_wrapper_${suite}_${dtype}_${mode}_${device}_${target}.csv"
       fi
       if [[ "$DASHBOARD_TAG" == *freezing_cudagraphs-true* ]] && [[ "$mode" == "inference" ]]; then
+        local profiler_trace_flags=()
+        if [[ "${EXPORT_PROFILER_TRACE:-0}" == "1" && "$target" == "performance" ]]; then
+          profiler_trace_flags=(--export-profiler-trace --profiler-trace-name "$TEST_REPORTS_DIR/profiler_traces/${backend}_with_cudagraphs_freezing_${suite}_${dtype}_${mode}_${device}")
+        fi
         $TASKSET python "benchmarks/dynamo/$suite.py" \
             "${target_flag[@]}" --"$mode" --"$dtype" --backend "$backend" "$@" --freezing \
+            "${profiler_trace_flags[@]}" \
             --output "$TEST_REPORTS_DIR/${backend}_with_cudagraphs_freezing_${suite}_${dtype}_${mode}_${device}_${target}.csv"
       fi
       if [[ "$DASHBOARD_TAG" == *freeze_autotune_cudagraphs-true* ]] && [[ "$mode" == "inference" ]]; then
+        local profiler_trace_flags=()
+        if [[ "${EXPORT_PROFILER_TRACE:-0}" == "1" && "$target" == "performance" ]]; then
+          profiler_trace_flags=(--export-profiler-trace --profiler-trace-name "$TEST_REPORTS_DIR/profiler_traces/${backend}_with_cudagraphs_freezing_autotune_${suite}_${dtype}_${mode}_${device}")
+        fi
         TORCHINDUCTOR_MAX_AUTOTUNE=1 $TASKSET python "benchmarks/dynamo/$suite.py" \
             "${target_flag[@]}" --"$mode" --"$dtype" --backend "$backend" "$@" --freezing \
+            "${profiler_trace_flags[@]}" \
             --output "$TEST_REPORTS_DIR/${backend}_with_cudagraphs_freezing_autotune_${suite}_${dtype}_${mode}_${device}_${target}.csv"
       fi
       if [[ "$DASHBOARD_TAG" == *aotinductor-true* ]] && [[ "$mode" == "inference" ]]; then
-        if [[ "$target" == "accuracy" ]]; then
           # Also collect Export pass rate and display as a separate row
+          if [[ "$target" == "accuracy" ]]; then
           $TASKSET python "benchmarks/dynamo/$suite.py" \
               "${target_flag[@]}" --"$mode" --"$dtype" --export --disable-cudagraphs "$@" \
               --output "$TEST_REPORTS_DIR/${backend}_export_${suite}_${dtype}_${mode}_${device}_${target}.csv"
         fi
+        local profiler_trace_flags=()
+        if [[ "${EXPORT_PROFILER_TRACE:-0}" == "1" && "$target" == "performance" ]]; then
+          profiler_trace_flags=(--export-profiler-trace --profiler-trace-name "$TEST_REPORTS_DIR/profiler_traces/${backend}_aot_inductor_${suite}_${dtype}_${mode}_${device}")
+        fi
         $TASKSET python "benchmarks/dynamo/$suite.py" \
             "${target_flag[@]}" --"$mode" --"$dtype" --export-aot-inductor --disable-cudagraphs "$@" \
+            "${profiler_trace_flags[@]}" \
             --output "$TEST_REPORTS_DIR/${backend}_aot_inductor_${suite}_${dtype}_${mode}_${device}_${target}.csv"
       fi
       if [[ "$DASHBOARD_TAG" == *maxautotune-true* ]]; then
+        local profiler_trace_flags=()
+        if [[ "${EXPORT_PROFILER_TRACE:-0}" == "1" && "$target" == "performance" ]]; then
+          profiler_trace_flags=(--export-profiler-trace --profiler-trace-name "$TEST_REPORTS_DIR/profiler_traces/${backend}_max_autotune_${suite}_${dtype}_${mode}_${device}")
+        fi
         TORCHINDUCTOR_MAX_AUTOTUNE=1 $TASKSET python "benchmarks/dynamo/$suite.py" \
             "${target_flag[@]}" --"$mode" --"$dtype" --backend "$backend" "$@" \
+            "${profiler_trace_flags[@]}" \
             --output "$TEST_REPORTS_DIR/${backend}_max_autotune_${suite}_${dtype}_${mode}_${device}_${target}.csv"
       fi
+      if [[ "$DASHBOARD_TAG" == *deterministic_perf-true* ]]; then
+        $TASKSET python "benchmarks/dynamo/$suite.py" \
+            "${target_flag[@]}" --"$mode" --"$dtype" --backend "$backend" --disable-cudagraphs --deterministic "$@" \
+            --output "$TEST_REPORTS_DIR/${backend}_deterministic_perf_${suite}_${dtype}_${mode}_${device}_${target}.csv"
+      fi
     done
   done
 }
@@ -844,9 +958,15 @@ test_single_dynamo_benchmark() {
   fi
 
   if [[ "${TEST_CONFIG}" == *perf_compare* ]]; then
+    local profiler_trace_flags=()
+    if [[ "${EXPORT_PROFILER_TRACE:-0}" == "1" ]]; then
+      mkdir -p "$TEST_REPORTS_DIR/profiler_traces"
+      profiler_trace_flags=(--export-profiler-trace --profiler-trace-name "$TEST_REPORTS_DIR/profiler_traces/${name}_${suite}")
+    fi
     python "benchmarks/dynamo/$suite.py" \
       --ci --performance --disable-cudagraphs --inductor \
       "${DYNAMO_BENCHMARK_FLAGS[@]}" "$@" "${partition_flags[@]}" \
+      "${profiler_trace_flags[@]}" \
       --output "$TEST_REPORTS_DIR/${name}_${suite}.csv"
   elif [[ "${TEST_CONFIG}" == *perf* ]]; then
     test_perf_for_dashboard "$suite" \
@@ -900,6 +1020,49 @@ test_inductor_triton_cpu() {
   assert_git_not_dirty
 }
 
+setup_torch_trace() {
+  if [[ "${ENABLE_TORCH_TRACE:-0}" != "1" ]]; then
+    return
+  fi
+  local trace_dir="${RUNNER_TEMP:-/tmp}/torch_traces"
+  mkdir -p "$trace_dir"
+  export TORCH_TRACE="$trace_dir"
+  echo "TORCH_TRACE enabled: writing structured trace logs to $trace_dir"
+}
+
+collect_tlparse_output() {
+  if [[ "${ENABLE_TORCH_TRACE:-0}" != "1" ]]; then
+    return
+  fi
+  local trace_dir="${RUNNER_TEMP:-/tmp}/torch_traces"
+  local test_reports_dir
+  test_reports_dir=$(pwd)/test/test-reports
+
+  if [[ ! -d "$trace_dir" ]] || [[ -z "$(ls -A "$trace_dir" 2>/dev/null)" ]]; then
+    echo "No torch trace files found in $trace_dir, skipping tlparse"
+    return
+  fi
+
+  echo "Collecting tlparse output from $trace_dir"
+
+  # Install tlparse if not already available
+  if ! command -v tlparse &>/dev/null; then
+    pip install tlparse 2>/dev/null || {
+      echo "Warning: failed to install tlparse, skipping HTML generation"
+      return
+    }
+  fi
+
+  # Run tlparse to generate HTML report
+  mkdir -p "$test_reports_dir/tlparse_output"
+  tlparse -o "$test_reports_dir/tlparse_output/" --no-browser --overwrite "$trace_dir" 2>&1 || {
+    echo "Warning: tlparse failed to generate HTML output"
+    return
+  }
+
+  echo "TLParse output generated in $test_reports_dir/tlparse_output/"
+}
+
 test_dynamo_benchmark() {
   # Usage: test_dynamo_benchmark huggingface 0
   TEST_REPORTS_DIR=$(pwd)/test/test-reports
@@ -982,6 +1145,147 @@ test_inductor_torchbench_smoketest_perf() {
   done
 }
 
+test_unbacked_parity_smoketest() {
+  # Check that unbacked batch-only has performance parity with backed batch-only
+  # Fails if any model regresses >THRESHOLD% consistently across 3 retries
+  TEST_REPORTS_DIR=$(pwd)/test/test-reports
+  mkdir -p "$TEST_REPORTS_DIR"
+
+  local THRESHOLD=1.0
+  local MAX_RETRIES=3
+  local MODELS="MobileBertForMaskedLM|DistilBertForMaskedLM|DistillGPT2|T5Small"
+
+  # Issue 6: Write per-run output files for post-failure debugging
+  run_comparison() {
+    local run_num=$1
+    local output_file="$TEST_REPORTS_DIR/unbacked_parity_results_run${run_num}.txt"
+    python benchmarks/dynamo/huggingface.py \
+      --compare-backed-unbacked \
+      --performance --inference --inductor --device cuda \
+      --filter "$MODELS" 2>&1 | tee "$output_file"
+  }
+
+  check_regressions() {
+    local run_num=$1
+    local output_file="$TEST_REPORTS_DIR/unbacked_parity_results_run${run_num}.txt"
+    # Parse the comparison table and check for regressions > threshold
+    # Returns 0 if regressions found, 1 if no regressions
+    local regressions=()
+    while IFS= read -r line; do
+      # Issue 3: Broadened regex to match model names with hyphens, slashes, dots
+      # Match lines like: "  ModelName                      10.000      10.500    +5.0%"
+      if [[ "$line" =~ ^[[:space:]]+([A-Za-z0-9_./-]+)[[:space:]]+([0-9.]+)[[:space:]]+([0-9.]+)[[:space:]]+\+([0-9.]+)% ]]; then
+        local model="${BASH_REMATCH[1]}"
+        local diff="${BASH_REMATCH[4]}"
+        # Nit: Use awk instead of bc -l to avoid dependency on bc
+        if awk "BEGIN{exit !($diff > $THRESHOLD)}"; then
+          regressions+=("$model:+${diff}%")
+        fi
+      fi
+    done < "$output_file"
+
+    if [[ ${#regressions[@]} -gt 0 ]]; then
+      echo "Regressions found: ${regressions[*]}"
+      return 0
+    fi
+    return 1
+  }
+
+  check_failures() {
+    local run_num=$1
+    local output_file="$TEST_REPORTS_DIR/unbacked_parity_results_run${run_num}.txt"
+    # Issue 2: Check for any model failure — not just paired failures.
+    # Specifically flags when unbacked fails but backed succeeds (regression signal).
+    # Returns 0 if failures found, 1 if no failures
+    local current_model=""
+    local backed_failed=false
+    local unbacked_failed=false
+    local both_failures=()
+    local unbacked_only_failures=()
+
+    # Append a sentinel header so the loop naturally evaluates the last real model
+    while IFS= read -r line; do
+      if [[ "$line" =~ ^---[[:space:]]+([A-Za-z0-9_./-]+)[[:space:]]+--- ]]; then
+        if [[ -n "$current_model" ]]; then
+          if $backed_failed && $unbacked_failed; then
+            both_failures+=("$current_model")
+          elif $unbacked_failed && ! $backed_failed; then
+            unbacked_only_failures+=("$current_model")
+          fi
+        fi
+        current_model="${BASH_REMATCH[1]}"
+        backed_failed=false
+        unbacked_failed=false
+      elif [[ "$line" =~ backed.*FAILED|backed.*TIMEOUT|backed.*ERROR ]]; then
+        backed_failed=true
+      elif [[ "$line" =~ unbacked.*FAILED|unbacked.*TIMEOUT|unbacked.*ERROR ]]; then
+        unbacked_failed=true
+      fi
+    done < <(cat "$output_file"; echo "--- END ---")
+
+    local has_failures=false
+    if [[ ${#both_failures[@]} -gt 0 ]]; then
+      echo "❌ FAILURES DETECTED: Both backed and unbacked failed for: ${both_failures[*]}"
+      has_failures=true
+    fi
+    if [[ ${#unbacked_only_failures[@]} -gt 0 ]]; then
+      echo "❌ FAILURES DETECTED: Unbacked failed (but backed succeeded) for: ${unbacked_only_failures[*]}"
+      has_failures=true
+    fi
+
+    if $has_failures; then
+      return 0
+    fi
+    return 1
+  }
+
+  # Run initial comparison
+  echo "=== Run 1/$MAX_RETRIES ==="
+  run_comparison 1
+
+  # Check for failures first
+  if check_failures 1; then
+    echo "❌ Test failed: Models failed to run (see above for details)"
+    exit 1
+  fi
+
+  # Check for regressions
+  if ! check_regressions 1; then
+    echo "✅ PASSED: No regressions above ${THRESHOLD}% threshold"
+    exit 0
+  fi
+
+  # Regression detected - retry to confirm
+  local regression_count=1
+  for ((retry=2; retry<=MAX_RETRIES; retry++)); do
+    echo ""
+    echo "=== Retry $retry/$MAX_RETRIES (potential regression detected) ==="
+    run_comparison "$retry"
+
+    # Issue 4: Also check for failures on retries (e.g., intermittent OOM)
+    if check_failures "$retry"; then
+      echo "❌ Test failed: Models failed on retry $retry (see above for details)"
+      exit 1
+    fi
+
+    if check_regressions "$retry"; then
+      ((regression_count++))
+    fi
+  done
+
+  # Check if regression was consistent (majority of runs)
+  local required=$((MAX_RETRIES / 2 + 1))
+  if [[ $regression_count -ge $required ]]; then
+    echo ""
+    echo "❌ REGRESSION CONFIRMED: Detected in $regression_count/$MAX_RETRIES runs (threshold: ${THRESHOLD}%)"
+    exit 1
+  else
+    echo ""
+    echo "✅ PASSED: Regressions were not consistent ($regression_count/$MAX_RETRIES runs, needed $required)"
+    exit 0
+  fi
+}
+
 test_inductor_set_cpu_affinity(){
   JEMALLOC_LIB="$(find /usr/lib -name libjemalloc.so.2)"
   export LD_PRELOAD="$JEMALLOC_LIB":"$LD_PRELOAD"
@@ -1156,6 +1460,18 @@ test_libtorch_jit() {
   popd
 }
 
+test_libtorch_profiler() {
+  echo "Testing profiler C++ tests"
+  export CPP_TESTS_DIR="${TORCH_BIN_DIR}"
+  export LD_LIBRARY_PATH="${TORCH_LIB_DIR}:${LD_LIBRARY_PATH}"
+
+  # Run E2E test first (needs clean Kineto state)
+  python test/run_test.py --cpp --verbose -i cpp/test_privateuse1_profiler -k "EndToEndProfiling"
+
+  # Run all other tests
+  python test/run_test.py --cpp --verbose -i cpp/test_privateuse1_profiler -k "not EndToEndProfiling"
+}
+
 test_libtorch_api() {
   # Start background download
   MNIST_DIR="${PWD}/test/cpp/api/mnist"
@@ -1625,74 +1941,6 @@ EOF
   assert_git_not_dirty
 }
 
-test_bazel() {
-  set -e -o pipefail
-
-  # bazel test needs sccache setup.
-  # shellcheck source=./common-build.sh
-  source "$(dirname "${BASH_SOURCE[0]}")/common-build.sh"
-
-  get_bazel
-
-  if [[ "$CUDA_VERSION" == "cpu" ]]; then
-    # Test //c10/... without Google flags and logging libraries. The
-    # :all_tests target in the subsequent Bazel invocation tests
-    # //c10/... with the Google libraries.
-    tools/bazel test --config=cpu-only --test_timeout=480 --test_output=all --test_tag_filters=-gpu-required --test_filter=-*CUDA \
-      --no//c10:use_gflags --no//c10:use_glog //c10/...
-
-    tools/bazel test --config=cpu-only --test_timeout=480 --test_output=all --test_tag_filters=-gpu-required --test_filter=-*CUDA :all_tests
-  else
-    # Increase the test timeout to 480 like CPU tests because modules_test frequently timeout
-    tools/bazel test --test_timeout=480 --test_output=errors \
-      //:any_test \
-      //:autograd_test \
-      //:dataloader_test \
-      //:dispatch_test \
-      //:enum_test \
-      //:expanding_array_test \
-      //:fft_test \
-      //:functional_test \
-      //:grad_mode_test \
-      //:inference_mode_test \
-      //:init_test \
-      //:jit_test \
-      //:memory_test \
-      //:meta_tensor_test \
-      //:misc_test \
-      //:moduledict_test \
-      //:modulelist_test \
-      //:modules_test \
-      //:namespace_test \
-      //:nested_test \
-      //:nn_utils_test \
-      //:operations_test \
-      //:ordered_dict_test \
-      //:parallel_benchmark_test \
-      //:parameterdict_test \
-      //:parameterlist_test \
-      //:sequential_test \
-      //:serialize_test \
-      //:special_test \
-      //:static_test \
-      //:support_test \
-      //:tensor_flatten_test \
-      //:tensor_indexing_test \
-      //:tensor_options_cuda_test \
-      //:tensor_options_test \
-      //:tensor_test \
-      //:torch_dist_autograd_test \
-      //:torch_include_test \
-      //:transformer_test \
-      //:test_bazel \
-      //c10/cuda/test:test \
-      //c10/test:core_tests \
-      //c10/test:typeid_test \
-      //c10/test:util/ssize_test \
-      //c10/test:util_base_tests
-  fi
-}
-
 test_benchmarks() {
   if [[ "$BUILD_ENVIRONMENT" == *cuda* && $TEST_CONFIG != *nogpu* ]]; then
     pip_install "pytest-benchmark==3.2.3"
@@ -1763,34 +2011,6 @@ test_executorch() {
   assert_git_not_dirty
 }
 
-test_linux_aarch64() {
-  python test/run_test.py --include test_modules test_utils test_mkldnn test_mkldnn_fusion test_openmp test_torch test_dynamic_shapes \
-        test_transformers test_multiprocessing test_numpy_interop test_autograd test_binary_ufuncs test_complex test_spectral_ops \
-        test_foreach test_reductions test_unary_ufuncs test_tensor_creation_ops test_ops profiler/test_memory_profiler \
-        distributed/elastic/timer/api_test distributed/elastic/timer/local_timer_example distributed/elastic/timer/local_timer_test \
-        test_linalg \
-        --shard "$SHARD_NUMBER" "$NUM_TEST_SHARDS" --verbose
-
-  # Dynamo tests
-  python test/run_test.py --include dynamo/test_compile dynamo/test_backends dynamo/test_comptime dynamo/test_config \
-       dynamo/test_functions dynamo/test_fx_passes_pre_grad dynamo/test_interop dynamo/test_model_output dynamo/test_modules \
-       dynamo/test_optimizers dynamo/test_recompile_ux dynamo/test_recompiles \
-       --shard "$SHARD_NUMBER" "$NUM_TEST_SHARDS" --verbose
-
-  # Inductor tests
-  python test/run_test.py --include inductor/test_torchinductor inductor/test_benchmark_fusion inductor/test_codecache \
-       inductor/test_config inductor/test_control_flow inductor/test_coordinate_descent_tuner inductor/test_fx_fusion \
-       inductor/test_group_batch_fusion inductor/test_inductor_freezing inductor/test_inductor_utils \
-       inductor/test_inplacing_pass inductor/test_kernel_benchmark inductor/test_layout_optim \
-       inductor/test_max_autotune inductor/test_memory_planning inductor/test_metrics inductor/test_multi_kernel inductor/test_pad_mm \
-       inductor/test_pattern_matcher inductor/test_perf inductor/test_profiler inductor/test_select_algorithm inductor/test_smoke \
-       inductor/test_split_cat_fx_passes inductor/test_compile inductor/test_torchinductor \
-       inductor/test_torchinductor_codegen_dynamic_shapes inductor/test_torchinductor_dynamic_shapes inductor/test_memory \
-       inductor/test_triton_cpu_backend inductor/test_triton_extension_backend inductor/test_mkldnn_pattern_matcher inductor/test_cpu_cpp_wrapper \
-       inductor/test_cpu_select_algorithm \
-       --shard "$SHARD_NUMBER" "$NUM_TEST_SHARDS" --verbose
-}
-
 test_operator_benchmark() {
   TEST_REPORTS_DIR=$(pwd)/test/test-reports
   mkdir -p "$TEST_REPORTS_DIR"
@@ -1853,15 +2073,19 @@ test_attention_microbenchmark() {
 }
 
 test_openreg() {
+  git submodule update --init --depth 1 third_party/googletest
   python test/run_test.py --openreg --verbose
   assert_git_not_dirty
 }
 
-if ! [[ "${BUILD_ENVIRONMENT}" == *libtorch* || "${BUILD_ENVIRONMENT}" == *-bazel-* ]]; then
+if ! [[ "${BUILD_ENVIRONMENT}" == *libtorch* ]]; then
   (cd test && python -c "import torch; print(torch.__config__.show())")
   (cd test && python -c "import torch; print(torch.__config__.parallel_info())")
 fi
-if [[ "${TEST_CONFIG}" == *numpy_2* ]]; then
+if [[ "${TEST_CONFIG}" == "onnx" ]]; then
+  install_torchvision
+  "$(dirname "${BASH_SOURCE[0]}")/../../scripts/onnx/test.sh"
+elif [[ "${TEST_CONFIG}" == *numpy_2* ]]; then
   # Install numpy-2.0.2 and compatible scipy & numba versions
   # Force re-install of pandas to avoid error where pandas checks numpy version from initial install and fails upon import
   TMP_PANDAS_VERSION=$(python -c "import pandas; print(pandas.__version__)" 2>/dev/null)
@@ -1871,8 +2095,6 @@ if [[ "${TEST_CONFIG}" == *numpy_2* ]]; then
     python -m pip install --pre numpy==2.0.2 scipy==1.13.1 numba==0.60.0
   fi
   python test/run_test.py --include dynamo/test_functions.py dynamo/test_unspec.py test_binary_ufuncs.py test_fake_tensor.py test_linalg.py test_numpy_interop.py test_tensor_creation_ops.py test_torch.py torch_np/test_basic.py
-elif [[ "${BUILD_ENVIRONMENT}" == *aarch64* && "${TEST_CONFIG}" == 'default' ]]; then
-  test_linux_aarch64
 elif [[ "${TEST_CONFIG}" == *backward* ]]; then
   test_forward_backward_compatibility
   # Do NOT add tests after bc check tests, see its comment.
@@ -1887,6 +2109,9 @@ elif [[ "$TEST_CONFIG" == *vllm* ]]; then
     (cd .ci/lumen_cli && python -m pip install -e .)
 
     python -m cli.run test external vllm --test-plan "$TEST_CONFIG" --shard-id "$SHARD_NUMBER" --num-shards "$NUM_TEST_SHARDS"
+elif [[ "$TEST_CONFIG" == *torchtitan* ]]; then
+    (cd .ci/lumen_cli && python -m pip install -e .)
+    python -m cli.run test external torchtitan --test-plan "$TEST_CONFIG" --shard-id "$SHARD_NUMBER" --num-shards "$NUM_TEST_SHARDS"
 elif [[ "${TEST_CONFIG}" == *executorch* ]]; then
   test_executorch
 elif [[ "$TEST_CONFIG" == 'jit_legacy' ]]; then
@@ -1920,7 +2145,9 @@ elif [[ "${TEST_CONFIG}" == *operator_microbenchmark* ]]; then
 elif [[ "${TEST_CONFIG}" == *attention_microbenchmark* ]]; then
   test_attention_microbenchmark
 elif [[ "${TEST_CONFIG}" == *inductor_distributed* ]]; then
+  setup_torch_trace
   test_inductor_distributed
+  collect_tlparse_output
 elif [[ "${TEST_CONFIG}" == *inductor-halide* ]]; then
   test_inductor_halide
 elif [[ "${TEST_CONFIG}" == *inductor-pallas* ]]; then
@@ -1934,17 +2161,19 @@ elif [[ "${TEST_CONFIG}" == *aoti_cross_compile_for_windows* ]]; then
 elif [[ "${TEST_CONFIG}" == *huggingface* ]]; then
   install_torchvision
   id=$((SHARD_NUMBER-1))
-  test_dynamo_benchmark huggingface "$id"
+  setup_torch_trace
+  if [[ "${TEST_CONFIG}" == *unbacked_parity* ]]; then
+    test_unbacked_parity_smoketest
+  else
+    test_dynamo_benchmark huggingface "$id"
+  fi
+  collect_tlparse_output
 elif [[ "${TEST_CONFIG}" == *timm* ]]; then
   install_torchvision
-  TIMM_PIN="$(< .ci/docker/ci_commit_pins/timm.txt)"
-  export HF_HOME="${HF_HOME}/timm_${TIMM_PIN}"
-  if [[ "${TRANSFORMERS_OFFLINE:-1}" == "0" ]]; then
-    python benchmarks/dynamo/timm_models.py --download-only \
-      && touch "${HF_HOME}/.timm_cache_complete"
-  fi
   id=$((SHARD_NUMBER-1))
+  setup_torch_trace
   test_dynamo_benchmark timm_models "$id"
+  collect_tlparse_output
 elif [[ "${TEST_CONFIG}" == cachebench ]]; then
   install_torchaudio
   install_torchvision
@@ -1975,19 +2204,27 @@ elif [[ "${TEST_CONFIG}" == *torchbench* ]]; then
       LIBTBB_PATH="$(find "$(dirname "$(which python)")/../lib/" -name libtbb.so.12)"
       export LD_PRELOAD="$LIBTBB_PATH":"$LD_PRELOAD"
     fi
+    setup_torch_trace
     PYTHONPATH=/torchbench test_dynamo_benchmark torchbench "$id"
+    collect_tlparse_output
   fi
 elif [[ "${TEST_CONFIG}" == *inductor_cpp_wrapper* ]]; then
   install_torchvision
+  setup_torch_trace
   PYTHONPATH=/torchbench test_inductor_cpp_wrapper_shard "$SHARD_NUMBER"
   if [[ "$SHARD_NUMBER" -eq "1" ]]; then
     test_inductor_aoti_cpp
   fi
+  collect_tlparse_output
 elif [[ "${TEST_CONFIG}" == *inductor_core* ]]; then
+  setup_torch_trace
   test_inductor_core
+  collect_tlparse_output
 elif [[ "${TEST_CONFIG}" == *inductor* ]]; then
   install_torchvision
+  setup_torch_trace
   test_inductor_shard "${SHARD_NUMBER}"
+  collect_tlparse_output
 elif [[ "${TEST_CONFIG}" == *einops* ]]; then
   test_einops
 elif [[ "${TEST_CONFIG}" == *dynamo_core* ]]; then
@@ -2022,27 +2259,25 @@ elif [[ "${SHARD_NUMBER}" == 2 && $NUM_TEST_SHARDS -gt 1 ]]; then
   test_custom_script_ops
   test_custom_backend
   test_torch_function_benchmark
+  test_libtorch_profiler
 elif [[ "${SHARD_NUMBER}" -gt 2 ]]; then
   # Handle arbitrary number of shards
   install_torchvision
   test_python_shard "$SHARD_NUMBER"
 elif [[ "${BUILD_ENVIRONMENT}" == *vulkan* ]]; then
   test_vulkan
-elif [[ "${BUILD_ENVIRONMENT}" == *-bazel-* ]]; then
-  test_bazel
 elif [[ "${BUILD_ENVIRONMENT}" == *-mobile-lightweight-dispatch* ]]; then
   test_libtorch
 elif [[ "${TEST_CONFIG}" = docs_test ]]; then
   test_docs_test
-elif [[ "${BUILD_ENVIRONMENT}" == *xpu* ]]; then
-  install_torchvision
-  test_python
-  test_aten
-  test_xpu_bin
 elif [[ "${TEST_CONFIG}" == smoke ]]; then
   test_python_smoke
 elif [[ "${TEST_CONFIG}" == smoke_b200 ]]; then
   test_python_smoke_b200
+elif [[ "${TEST_CONFIG}" == smoke_xpu ]]; then
+  test_python_smoke_xpu
+elif [[ "${TEST_CONFIG}" == dtensor ]]; then
+  test_dtensor
 elif [[ "${TEST_CONFIG}" == h100_distributed ]]; then
   test_h100_distributed
 elif [[ "${TEST_CONFIG}" == "h100-symm-mem" ]]; then
diff --git a/.ci/pytorch/test_example_code/CMakeLists.txt b/.ci/pytorch/test_example_code/CMakeLists.txt
index e87f37ae61fb4..c89a728f27649 100644
--- a/.ci/pytorch/test_example_code/CMakeLists.txt
+++ b/.ci/pytorch/test_example_code/CMakeLists.txt
@@ -1,4 +1,4 @@
-cmake_minimum_required(VERSION 3.0 FATAL_ERROR)
+cmake_minimum_required(VERSION 3.10 FATAL_ERROR)
 project(simple-torch-test)
 
 find_package(Torch REQUIRED)
@@ -8,7 +8,7 @@ set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} ${TORCH_CXX_FLAGS}")
 add_executable(simple-torch-test simple-torch-test.cpp)
 target_include_directories(simple-torch-test PRIVATE  ${TORCH_INCLUDE_DIRS})
 target_link_libraries(simple-torch-test "${TORCH_LIBRARIES}")
-set_property(TARGET simple-torch-test PROPERTY CXX_STANDARD 17)
+set_property(TARGET simple-torch-test PROPERTY CXX_STANDARD 20)
 
 find_package(CUDAToolkit 11.8)
 
diff --git a/.ci/pytorch/win-test-helpers/test_openreg.bat b/.ci/pytorch/win-test-helpers/test_openreg.bat
index 0470057daf641..db3c0758f44e6 100644
--- a/.ci/pytorch/win-test-helpers/test_openreg.bat
+++ b/.ci/pytorch/win-test-helpers/test_openreg.bat
@@ -6,6 +6,8 @@ if not errorlevel 0 (
   exit /b
 )
 
+git submodule update --init --depth 1 third_party/googletest
+
 pushd test
 
 echo Run openreg tests
diff --git a/.ci/pytorch/windows/arm64/build_libtorch.bat b/.ci/pytorch/windows/arm64/build_libtorch.bat
index 1ac14ff697730..33080af122c73 100644
--- a/.ci/pytorch/windows/arm64/build_libtorch.bat
+++ b/.ci/pytorch/windows/arm64/build_libtorch.bat
@@ -27,8 +27,8 @@ where cl.exe
 :: change to source directory
 cd %PYTORCH_ROOT%
 
-:: copy libuv.dll
-copy %libuv_ROOT%\lib\Release\uv.dll torch\lib\uv.dll
+:: copy libuv.dll (cmake installs the dll to bin/, not lib/Release/)
+copy %libuv_ROOT%\bin\uv.dll torch\lib\uv.dll
 
 :: create virtual environment
 python -m venv .venv
diff --git a/.ci/pytorch/windows/arm64/build_pytorch.bat b/.ci/pytorch/windows/arm64/build_pytorch.bat
index b5c2ef65b84ad..7d10b26339d25 100644
--- a/.ci/pytorch/windows/arm64/build_pytorch.bat
+++ b/.ci/pytorch/windows/arm64/build_pytorch.bat
@@ -5,6 +5,7 @@ set CMAKE_BUILD_TYPE=%BUILD_TYPE%
 set CMAKE_C_COMPILER_LAUNCHER=sccache
 set CMAKE_CXX_COMPILER_LAUNCHER=sccache
 set libuv_ROOT=%DEPENDENCIES_DIR%\libuv\install
+set INSTALL_TEST=0
 set MSSdk=1
 if defined PYTORCH_BUILD_VERSION (
   set PYTORCH_BUILD_VERSION=%PYTORCH_BUILD_VERSION%
@@ -27,8 +28,8 @@ where cl.exe
 :: change to source directory
 cd %PYTORCH_ROOT%
 
-:: copy libuv.dll
-copy %libuv_ROOT%\lib\Release\uv.dll torch\lib\uv.dll
+:: copy libuv.dll (cmake installs the dll to bin/, not lib/Release/)
+copy %libuv_ROOT%\bin\uv.dll torch\lib\uv.dll
 
 :: create virtual environment
 python -m venv .venv
@@ -57,4 +58,4 @@ sccache --show-stats
 if %errorlevel% neq 0 (
     echo "Failed on build_pytorch. (exitcode = %errorlevel%)"
     exit /b 1
-)
\ No newline at end of file
+)
diff --git a/.ci/pytorch/windows/arm64/smoke_test.bat b/.ci/pytorch/windows/arm64/smoke_test.bat
index 2c90f6158062f..202f6eb8cc9c5 100644
--- a/.ci/pytorch/windows/arm64/smoke_test.bat
+++ b/.ci/pytorch/windows/arm64/smoke_test.bat
@@ -40,10 +40,10 @@ set INCLUDE=%INCLUDE%;%install_root%\include;%install_root%\include\torch\csrc\a
 set LIB=%LIB%;%install_root%\lib
 set PATH=%PATH%;%install_root%\lib
 
-cl %PYTORCH_ROOT%\.ci\pytorch\test_example_code\simple-torch-test.cpp c10.lib torch_cpu.lib /EHsc /std:c++17
+cl %PYTORCH_ROOT%\.ci\pytorch\test_example_code\simple-torch-test.cpp c10.lib torch_cpu.lib /EHsc /std:c++20
 if ERRORLEVEL 1 exit /b 1
 
 .\simple-torch-test.exe
 if ERRORLEVEL 1 exit /b 1
 
-:end
\ No newline at end of file
+:end
diff --git a/.ci/pytorch/windows/internal/cuda_config.bat b/.ci/pytorch/windows/internal/cuda_config.bat
index 352eb8a3391bf..69059e2e46ee1 100644
--- a/.ci/pytorch/windows/internal/cuda_config.bat
+++ b/.ci/pytorch/windows/internal/cuda_config.bat
@@ -22,6 +22,10 @@ if "%CUDA_VER%"=="126" (
     set "CUDA_DOTTED_VERSION=13.0"
     set "CUDA_ARCH_LIST=7.5;8.0;8.6;9.0;10.0;12.0"
     set "VISION_GENCODE=-gencode=arch=compute_75,code=sm_75 -gencode=arch=compute_80,code=compute_80 -gencode=arch=compute_86,code=compute_86 -gencode=arch=compute_90,code=compute_90 -gencode=arch=compute_100,code=compute_100 -gencode=arch=compute_120,code=compute_120"
+) else if "%CUDA_VER%"=="132" (
+    set "CUDA_DOTTED_VERSION=13.2"
+    set "CUDA_ARCH_LIST=7.5;8.0;8.6;9.0;10.0;12.0"
+    set "VISION_GENCODE=-gencode=arch=compute_75,code=sm_75 -gencode=arch=compute_80,code=compute_80 -gencode=arch=compute_86,code=compute_86 -gencode=arch=compute_90,code=compute_90 -gencode=arch=compute_100,code=compute_100 -gencode=arch=compute_120,code=compute_120"
 ) else (
     echo Unknown CUDA version: %CUDA_VER%
     exit /b 1
diff --git a/.ci/pytorch/windows/internal/cuda_install.bat b/.ci/pytorch/windows/internal/cuda_install.bat
index 1349d3e661f55..c1050edecc0b9 100644
--- a/.ci/pytorch/windows/internal/cuda_install.bat
+++ b/.ci/pytorch/windows/internal/cuda_install.bat
@@ -17,16 +17,17 @@ set /a CUDA_VER=%CUDA_VERSION%
 set CUDA_VER_MAJOR=%CUDA_VERSION:~0,-1%
 set CUDA_VER_MINOR=%CUDA_VERSION:~-1,1%
 set CUDA_VERSION_STR=%CUDA_VER_MAJOR%.%CUDA_VER_MINOR%
-set CUDNN_FOLDER="cuda"
-set CUDNN_LIB_FOLDER="lib\x64"
+set CUDNN_FOLDER=cuda
+set CUDNN_LIB_FOLDER=lib\x64
 
-:: Skip all of this if we already have cuda installed
-if exist "C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v%CUDA_VERSION_STR%\bin\nvcc.exe" goto set_cuda_env_vars
+:: If CUDA is already installed, skip CUDA installation but still verify cuDNN
+if exist "C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v%CUDA_VERSION_STR%\bin\nvcc.exe" goto check_cudnn
 
 if %CUDA_VER% EQU 126 goto cuda126
 if %CUDA_VER% EQU 128 goto cuda128
 if %CUDA_VER% EQU 129 goto cuda129
 if %CUDA_VER% EQU 130 goto cuda130
+if %CUDA_VER% EQU 132 goto cuda132
 
 echo CUDA %CUDA_VERSION_STR% is not supported
 exit /b 1
@@ -34,110 +35,53 @@ exit /b 1
 goto cuda_common
 
 :cuda126
-
 set CUDA_INSTALL_EXE=cuda_12.6.2_560.94_windows.exe
-if not exist "%SRC_DIR%\temp_build\%CUDA_INSTALL_EXE%" (
-    curl -k -L "https://ossci-windows.s3.amazonaws.com/%CUDA_INSTALL_EXE%" --output "%SRC_DIR%\temp_build\%CUDA_INSTALL_EXE%" & REM @lint-ignore
-    if errorlevel 1 exit /b 1
-    set "CUDA_SETUP_FILE=%SRC_DIR%\temp_build\%CUDA_INSTALL_EXE%"
-    set "ARGS=cuda_profiler_api_12.6 thrust_12.6 nvcc_12.6 cuobjdump_12.6 nvprune_12.6 nvprof_12.6 cupti_12.6 cublas_12.6 cublas_dev_12.6 cudart_12.6 cufft_12.6 cufft_dev_12.6 curand_12.6 curand_dev_12.6 cusolver_12.6 cusolver_dev_12.6 cusparse_12.6 cusparse_dev_12.6 npp_12.6 npp_dev_12.6 nvrtc_12.6 nvrtc_dev_12.6 nvml_dev_12.6 nvjitlink_12.6 nvtx_12.6"
-)
-
-set CUDNN_FOLDER=cudnn-windows-x86_64-9.5.0.50_cuda12-archive
-set CUDNN_LIB_FOLDER="lib"
-set "CUDNN_INSTALL_ZIP=%CUDNN_FOLDER%.zip"
-if not exist "%SRC_DIR%\temp_build\%CUDNN_INSTALL_ZIP%" (
-    curl -k -L "http://s3.amazonaws.com/ossci-windows/%CUDNN_INSTALL_ZIP%" --output "%SRC_DIR%\temp_build\%CUDNN_INSTALL_ZIP%" & REM @lint-ignore
-    if errorlevel 1 exit /b 1
-    set "CUDNN_SETUP_FILE=%SRC_DIR%\temp_build\%CUDNN_INSTALL_ZIP%"
-)
-
-@REM cuDNN 8.3+ required zlib to be installed on the path
-echo Installing ZLIB dlls
-curl -k -L "http://s3.amazonaws.com/ossci-windows/zlib123dllx64.zip" --output "%SRC_DIR%\temp_build\zlib123dllx64.zip"
-7z x "%SRC_DIR%\temp_build\zlib123dllx64.zip" -o"%SRC_DIR%\temp_build\zlib"
-xcopy /Y "%SRC_DIR%\temp_build\zlib\dll_x64\*.dll" "C:\Windows\System32"
-
-goto cuda_common
+set "ARGS=cuda_profiler_api_12.6 thrust_12.6 nvcc_12.6 cuobjdump_12.6 nvprune_12.6 nvprof_12.6 cupti_12.6 cublas_12.6 cublas_dev_12.6 cudart_12.6 cufft_12.6 cufft_dev_12.6 curand_12.6 curand_dev_12.6 cusolver_12.6 cusolver_dev_12.6 cusparse_12.6 cusparse_dev_12.6 npp_12.6 npp_dev_12.6 nvrtc_12.6 nvrtc_dev_12.6 nvml_dev_12.6 nvjitlink_12.6 nvtx_12.6"
+set CUDNN_FOLDER=cudnn-windows-x86_64-9.10.2.21_cuda12-archive
+goto cuda_download
 
 :cuda128
-
 set CUDA_INSTALL_EXE=cuda_12.8.0_571.96_windows.exe
-if not exist "%SRC_DIR%\temp_build\%CUDA_INSTALL_EXE%" (
-    curl -k -L "https://ossci-windows.s3.amazonaws.com/%CUDA_INSTALL_EXE%" --output "%SRC_DIR%\temp_build\%CUDA_INSTALL_EXE%" & REM @lint-ignore
-    if errorlevel 1 exit /b 1
-    set "CUDA_SETUP_FILE=%SRC_DIR%\temp_build\%CUDA_INSTALL_EXE%"
-    set "ARGS=cuda_profiler_api_12.8 thrust_12.8 nvcc_12.8 cuobjdump_12.8 nvprune_12.8 nvprof_12.8 cupti_12.8 cublas_12.8 cublas_dev_12.8 cudart_12.8 cufft_12.8 cufft_dev_12.8 curand_12.8 curand_dev_12.8 cusolver_12.8 cusolver_dev_12.8 cusparse_12.8 cusparse_dev_12.8 npp_12.8 npp_dev_12.8 nvrtc_12.8 nvrtc_dev_12.8 nvml_dev_12.8 nvjitlink_12.8 nvtx_12.8"
-)
-
-set CUDNN_FOLDER=cudnn-windows-x86_64-9.7.0.66_cuda12-archive
-set CUDNN_LIB_FOLDER="lib"
-set "CUDNN_INSTALL_ZIP=%CUDNN_FOLDER%.zip"
-if not exist "%SRC_DIR%\temp_build\%CUDNN_INSTALL_ZIP%" (
-    curl -k -L "http://s3.amazonaws.com/ossci-windows/%CUDNN_INSTALL_ZIP%" --output "%SRC_DIR%\temp_build\%CUDNN_INSTALL_ZIP%" & REM @lint-ignore
-    if errorlevel 1 exit /b 1
-    set "CUDNN_SETUP_FILE=%SRC_DIR%\temp_build\%CUDNN_INSTALL_ZIP%"
-)
-
-@REM cuDNN 8.3+ required zlib to be installed on the path
-echo Installing ZLIB dlls
-curl -k -L "http://s3.amazonaws.com/ossci-windows/zlib123dllx64.zip" --output "%SRC_DIR%\temp_build\zlib123dllx64.zip"
-7z x "%SRC_DIR%\temp_build\zlib123dllx64.zip" -o"%SRC_DIR%\temp_build\zlib"
-xcopy /Y "%SRC_DIR%\temp_build\zlib\dll_x64\*.dll" "C:\Windows\System32"
-
-goto cuda_common
+set "ARGS=cuda_profiler_api_12.8 thrust_12.8 nvcc_12.8 cuobjdump_12.8 nvprune_12.8 nvprof_12.8 cupti_12.8 cublas_12.8 cublas_dev_12.8 cudart_12.8 cufft_12.8 cufft_dev_12.8 curand_12.8 curand_dev_12.8 cusolver_12.8 cusolver_dev_12.8 cusparse_12.8 cusparse_dev_12.8 npp_12.8 npp_dev_12.8 nvrtc_12.8 nvrtc_dev_12.8 nvml_dev_12.8 nvjitlink_12.8 nvtx_12.8"
+set CUDNN_FOLDER=cudnn-windows-x86_64-9.20.0.48_cuda12-archive
+goto cuda_download
 
 :cuda129
-
 set CUDA_INSTALL_EXE=cuda_12.9.1_576.57_windows.exe
-if not exist "%SRC_DIR%\temp_build\%CUDA_INSTALL_EXE%" (
-    curl -k -L "https://ossci-windows.s3.amazonaws.com/%CUDA_INSTALL_EXE%" --output "%SRC_DIR%\temp_build\%CUDA_INSTALL_EXE%" & REM @lint-ignore
-    if errorlevel 1 exit /b 1
-    set "CUDA_SETUP_FILE=%SRC_DIR%\temp_build\%CUDA_INSTALL_EXE%"
-    set "ARGS=cuda_profiler_api_12.9 thrust_12.9 nvcc_12.9 cuobjdump_12.9 nvprune_12.9 nvprof_12.9 cupti_12.9 cublas_12.9 cublas_dev_12.9 cudart_12.9 cufft_12.9 cufft_dev_12.9 curand_12.9 curand_dev_12.9 cusolver_12.9 cusolver_dev_12.9 cusparse_12.9 cusparse_dev_12.9 npp_12.9 npp_dev_12.9 nvrtc_12.9 nvrtc_dev_12.9 nvml_dev_12.9 nvjitlink_12.9 nvtx_12.9"
-)
-
-set CUDNN_FOLDER=cudnn-windows-x86_64-9.10.2.21_cuda12-archive
-set CUDNN_LIB_FOLDER="lib"
-set "CUDNN_INSTALL_ZIP=%CUDNN_FOLDER%.zip"
-if not exist "%SRC_DIR%\temp_build\%CUDNN_INSTALL_ZIP%" (
-    curl -k -L "http://s3.amazonaws.com/ossci-windows/%CUDNN_INSTALL_ZIP%" --output "%SRC_DIR%\temp_build\%CUDNN_INSTALL_ZIP%" & REM @lint-ignore
-    if errorlevel 1 exit /b 1
-    set "CUDNN_SETUP_FILE=%SRC_DIR%\temp_build\%CUDNN_INSTALL_ZIP%"
-)
-
-@REM cuDNN 8.3+ required zlib to be installed on the path
-echo Installing ZLIB dlls
-curl -k -L "http://s3.amazonaws.com/ossci-windows/zlib123dllx64.zip" --output "%SRC_DIR%\temp_build\zlib123dllx64.zip"
-7z x "%SRC_DIR%\temp_build\zlib123dllx64.zip" -o"%SRC_DIR%\temp_build\zlib"
-xcopy /Y "%SRC_DIR%\temp_build\zlib\dll_x64\*.dll" "C:\Windows\System32"
-
-goto cuda_common
+set "ARGS=cuda_profiler_api_12.9 thrust_12.9 nvcc_12.9 cuobjdump_12.9 nvprune_12.9 nvprof_12.9 cupti_12.9 cublas_12.9 cublas_dev_12.9 cudart_12.9 cufft_12.9 cufft_dev_12.9 curand_12.9 curand_dev_12.9 cusolver_12.9 cusolver_dev_12.9 cusparse_12.9 cusparse_dev_12.9 npp_12.9 npp_dev_12.9 nvrtc_12.9 nvrtc_dev_12.9 nvml_dev_12.9 nvjitlink_12.9 nvtx_12.9"
+set CUDNN_FOLDER=cudnn-windows-x86_64-9.20.0.48_cuda12-archive
+goto cuda_download
 
 :cuda130
-
 set CUDA_INSTALL_EXE=cuda_13.0.0_windows.exe
+set "ARGS="
+set CUDNN_FOLDER=cudnn-windows-x86_64-9.20.0.48_cuda13-archive
+goto cuda_download
+
+:cuda132
+set CUDA_INSTALL_EXE=cuda_13.2.1_windows.exe
+set "ARGS="
+set CUDNN_FOLDER=cudnn-windows-x86_64-9.20.0.48_cuda13-archive
+goto cuda_download
+
+:: Common download logic for CUDA toolkit, cuDNN, and ZLIB
+:cuda_download
+set CUDNN_LIB_FOLDER=lib
+set "CUDNN_INSTALL_ZIP=%CUDNN_FOLDER%.zip"
+
 if not exist "%SRC_DIR%\temp_build\%CUDA_INSTALL_EXE%" (
     curl -k -L "https://ossci-windows.s3.amazonaws.com/%CUDA_INSTALL_EXE%" --output "%SRC_DIR%\temp_build\%CUDA_INSTALL_EXE%" & REM @lint-ignore
     if errorlevel 1 exit /b 1
     set "CUDA_SETUP_FILE=%SRC_DIR%\temp_build\%CUDA_INSTALL_EXE%"
-    set "ARGS="
 )
 
-set CUDNN_FOLDER=cudnn-windows-x86_64-9.12.0.46_cuda13-archive
-set CUDNN_LIB_FOLDER="lib"
-set "CUDNN_INSTALL_ZIP=%CUDNN_FOLDER%.zip"
 if not exist "%SRC_DIR%\temp_build\%CUDNN_INSTALL_ZIP%" (
     curl -k -L "http://s3.amazonaws.com/ossci-windows/%CUDNN_INSTALL_ZIP%" --output "%SRC_DIR%\temp_build\%CUDNN_INSTALL_ZIP%" & REM @lint-ignore
     if errorlevel 1 exit /b 1
     set "CUDNN_SETUP_FILE=%SRC_DIR%\temp_build\%CUDNN_INSTALL_ZIP%"
 )
 
-@REM cuDNN 8.3+ required zlib to be installed on the path
-echo Installing ZLIB dlls
-curl -k -L "http://s3.amazonaws.com/ossci-windows/zlib123dllx64.zip" --output "%SRC_DIR%\temp_build\zlib123dllx64.zip"
-7z x "%SRC_DIR%\temp_build\zlib123dllx64.zip" -o"%SRC_DIR%\temp_build\zlib"
-xcopy /Y "%SRC_DIR%\temp_build\zlib\dll_x64\*.dll" "C:\Windows\System32"
+call :install_zlib
 
 goto cuda_common
 
@@ -189,9 +133,12 @@ if not exist "C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v%CUDA_VERSION_
 
     echo Installing cuDNN...
     7z x %CUDNN_SETUP_FILE% -o"%SRC_DIR%\temp_build\cudnn"
-    xcopy /Y "%SRC_DIR%\temp_build\cudnn\%CUDNN_FOLDER%\bin\*.*" "%ProgramFiles%\NVIDIA GPU Computing Toolkit\CUDA\v%CUDA_VERSION_STR%\bin"
-    xcopy /Y "%SRC_DIR%\temp_build\cudnn\%CUDNN_FOLDER%\%CUDNN_LIB_FOLDER%\*.*" "%ProgramFiles%\NVIDIA GPU Computing Toolkit\CUDA\v%CUDA_VERSION_STR%\lib\x64"
-    xcopy /Y "%SRC_DIR%\temp_build\cudnn\%CUDNN_FOLDER%\include\*.*" "%ProgramFiles%\NVIDIA GPU Computing Toolkit\CUDA\v%CUDA_VERSION_STR%\include"
+    xcopy /Y /S "%SRC_DIR%\temp_build\cudnn\%CUDNN_FOLDER%\bin\*.*" "%ProgramFiles%\NVIDIA GPU Computing Toolkit\CUDA\v%CUDA_VERSION_STR%\bin\"
+    if exist "%SRC_DIR%\temp_build\cudnn\%CUDNN_FOLDER%\bin\x64\*.*" (
+        xcopy /Y "%SRC_DIR%\temp_build\cudnn\%CUDNN_FOLDER%\bin\x64\*.*" "%ProgramFiles%\NVIDIA GPU Computing Toolkit\CUDA\v%CUDA_VERSION_STR%\bin\"
+    )
+    xcopy /Y /S "%SRC_DIR%\temp_build\cudnn\%CUDNN_FOLDER%\%CUDNN_LIB_FOLDER%\*.*" "%ProgramFiles%\NVIDIA GPU Computing Toolkit\CUDA\v%CUDA_VERSION_STR%\lib\x64\"
+    xcopy /Y /S "%SRC_DIR%\temp_build\cudnn\%CUDNN_FOLDER%\include\*.*" "%ProgramFiles%\NVIDIA GPU Computing Toolkit\CUDA\v%CUDA_VERSION_STR%\include\"
 
     echo Installing GPU driver DLLs
     7z x %SRC_DIR%\temp_build\gpu_driver_dlls.zip -o"C:\Windows\System32"
@@ -211,6 +158,85 @@ if not exist "C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v%CUDA_VERSION_
 
 goto set_cuda_env_vars
 
+:check_cudnn
+:: When CUDA is pre-installed on the AMI, cuDNN may still be missing.
+:: Set the correct cuDNN variables for the CUDA version, then install if needed.
+
+set CUDNN_LIB_FOLDER=lib
+if %CUDA_VER% EQU 126 (
+    set CUDNN_FOLDER=cudnn-windows-x86_64-9.10.2.21_cuda12-archive
+    set EXPECTED_CUDNN_VERSION=9.10.2
+)
+if %CUDA_VER% EQU 128 (
+    set CUDNN_FOLDER=cudnn-windows-x86_64-9.20.0.48_cuda12-archive
+    set EXPECTED_CUDNN_VERSION=9.20.0
+)
+if %CUDA_VER% EQU 129 (
+    set CUDNN_FOLDER=cudnn-windows-x86_64-9.20.0.48_cuda12-archive
+    set EXPECTED_CUDNN_VERSION=9.20.0
+)
+if %CUDA_VER% EQU 130 (
+    set CUDNN_FOLDER=cudnn-windows-x86_64-9.20.0.48_cuda13-archive
+    set EXPECTED_CUDNN_VERSION=9.20.0
+)
+if %CUDA_VER% EQU 132 (
+    set CUDNN_FOLDER=cudnn-windows-x86_64-9.20.0.48_cuda13-archive
+    set EXPECTED_CUDNN_VERSION=9.20.0
+)
+set "CUDNN_INSTALL_ZIP=%CUDNN_FOLDER%.zip"
+
+set "CUDNN_VERSION_FILE=%ProgramFiles%\NVIDIA GPU Computing Toolkit\CUDA\v%CUDA_VERSION_STR%\include\cudnn_version.h"
+
+if not exist "%CUDNN_VERSION_FILE%" (
+    echo cuDNN not found, installing %CUDNN_FOLDER%...
+    goto install_cudnn
+)
+
+for /f "tokens=3" %%a in ('findstr /C:"#define CUDNN_MAJOR " "%CUDNN_VERSION_FILE%"') do set INSTALLED_MAJOR=%%a
+for /f "tokens=3" %%a in ('findstr /C:"#define CUDNN_MINOR " "%CUDNN_VERSION_FILE%"') do set INSTALLED_MINOR=%%a
+for /f "tokens=3" %%a in ('findstr /C:"#define CUDNN_PATCHLEVEL " "%CUDNN_VERSION_FILE%"') do set INSTALLED_PATCHLEVEL=%%a
+set "INSTALLED_CUDNN_VERSION=%INSTALLED_MAJOR%.%INSTALLED_MINOR%.%INSTALLED_PATCHLEVEL%"
+
+if "%INSTALLED_CUDNN_VERSION%" == "%EXPECTED_CUDNN_VERSION%" (
+    echo cuDNN %INSTALLED_CUDNN_VERSION% already installed at %ProgramFiles%\NVIDIA GPU Computing Toolkit\CUDA\v%CUDA_VERSION_STR%
+    goto set_cuda_env_vars
+)
+
+echo cuDNN version mismatch: installed %INSTALLED_CUDNN_VERSION%, expected %EXPECTED_CUDNN_VERSION%. Reinstalling...
+
+:: Remove old cuDNN DLLs so they don't shadow the new version at runtime.
+:: AMI-installed cuDNN places DLLs directly in bin\, while newer archives
+:: use bin\x64\. Without cleanup the old DLLs in bin\ are found first.
+del /Q "%ProgramFiles%\NVIDIA GPU Computing Toolkit\CUDA\v%CUDA_VERSION_STR%\bin\cudnn*.dll" 2>nul
+
+:install_cudnn
+
+if not exist "%SRC_DIR%\temp_build" mkdir "%SRC_DIR%\temp_build"
+
+curl -k -L "http://s3.amazonaws.com/ossci-windows/%CUDNN_INSTALL_ZIP%" --output "%SRC_DIR%\temp_build\%CUDNN_INSTALL_ZIP%" & REM @lint-ignore
+if errorlevel 1 exit /b 1
+
+7z x "%SRC_DIR%\temp_build\%CUDNN_INSTALL_ZIP%" -o"%SRC_DIR%\temp_build\cudnn"
+if errorlevel 1 (
+    echo Failed to extract cuDNN archive %CUDNN_INSTALL_ZIP%
+    exit /b 1
+)
+echo Listing extracted cuDNN archive contents:
+dir /S /B "%SRC_DIR%\temp_build\cudnn\%CUDNN_FOLDER%"
+xcopy /Y /S "%SRC_DIR%\temp_build\cudnn\%CUDNN_FOLDER%\bin\*.*" "%ProgramFiles%\NVIDIA GPU Computing Toolkit\CUDA\v%CUDA_VERSION_STR%\bin\"
+:: Newer cuDNN archives place DLLs under bin\x64\. Flatten them into bin\
+:: so they are found via PATH (which only includes bin\, not bin\x64\).
+if exist "%SRC_DIR%\temp_build\cudnn\%CUDNN_FOLDER%\bin\x64\*.*" (
+    xcopy /Y "%SRC_DIR%\temp_build\cudnn\%CUDNN_FOLDER%\bin\x64\*.*" "%ProgramFiles%\NVIDIA GPU Computing Toolkit\CUDA\v%CUDA_VERSION_STR%\bin\"
+)
+xcopy /Y /S "%SRC_DIR%\temp_build\cudnn\%CUDNN_FOLDER%\%CUDNN_LIB_FOLDER%\*.*" "%ProgramFiles%\NVIDIA GPU Computing Toolkit\CUDA\v%CUDA_VERSION_STR%\lib\x64\"
+xcopy /Y /S "%SRC_DIR%\temp_build\cudnn\%CUDNN_FOLDER%\include\*.*" "%ProgramFiles%\NVIDIA GPU Computing Toolkit\CUDA\v%CUDA_VERSION_STR%\include\"
+
+call :install_zlib
+
+echo Cleaning temp files
+rd /s /q "%SRC_DIR%\temp_build" || ver > nul
+
 :set_cuda_env_vars
 
 echo Setting up environment...
@@ -218,3 +244,13 @@ set "PATH=%ProgramFiles%\NVIDIA GPU Computing Toolkit\CUDA\v%CUDA_VERSION_STR%\b
 set "CUDA_PATH=%ProgramFiles%\NVIDIA GPU Computing Toolkit\CUDA\v%CUDA_VERSION_STR%"
 set "CUDA_PATH_V%CUDA_VER_MAJOR%_%CUDA_VER_MINOR%=%ProgramFiles%\NVIDIA GPU Computing Toolkit\CUDA\v%CUDA_VERSION_STR%"
 set "NVTOOLSEXT_PATH=%ProgramFiles%\NVIDIA Corporation\NvToolsExt"
+
+goto :eof
+
+@REM cuDNN 8.3+ requires zlib to be installed on the path
+:install_zlib
+echo Installing ZLIB dlls
+curl -k -L "http://s3.amazonaws.com/ossci-windows/zlib123dllx64.zip" --output "%SRC_DIR%\temp_build\zlib123dllx64.zip"
+7z x "%SRC_DIR%\temp_build\zlib123dllx64.zip" -o"%SRC_DIR%\temp_build\zlib"
+xcopy /Y "%SRC_DIR%\temp_build\zlib\dll_x64\*.dll" "C:\Windows\System32"
+goto :eof
diff --git a/.ci/pytorch/windows/internal/smoke_test.bat b/.ci/pytorch/windows/internal/smoke_test.bat
index f671a9d0e0abb..8cce63a81693d 100644
--- a/.ci/pytorch/windows/internal/smoke_test.bat
+++ b/.ci/pytorch/windows/internal/smoke_test.bat
@@ -1,10 +1,15 @@
 set SRC_DIR=%~dp0
+set TARGET_OS=windows
 
 pushd %SRC_DIR%\..
 
 if not "%CUDA_VERSION%" == "cpu" if not "%CUDA_VERSION%" == "xpu" call internal\driver_update.bat
 if errorlevel 1 exit /b 1
 
+echo "Check if CUDA and CUDNN versions need to be updated"
+call internal\cuda_install.bat
+if errorlevel 1 exit /b 1
+
 if "%CUDA_VERSION%" == "xpu" (
     call internal\xpu_install.bat
     if errorlevel 1 exit /b 1
@@ -94,6 +99,10 @@ echo Checking that basic CNN works
 %PYTHON_EXEC% %PYTORCH_ROOT%\.ci\pytorch\test_example_code\cnn_smoke.py
 if ERRORLEVEL 1 exit /b 1
 
+echo Running smoke_test.py
+%PYTHON_EXEC% %PYTORCH_ROOT%\.ci\pytorch\smoke_test\smoke_test.py --package=torchonly --torch-compile-check disabled --runtime-error-check disabled
+if ERRORLEVEL 1 exit /b 1
+
 goto end
 
 :libtorch
@@ -131,13 +140,13 @@ set INCLUDE=%INCLUDE%;%install_root%\include;%install_root%\include\torch\csrc\a
 set LIB=%LIB%;%install_root%\lib
 set PATH=%PATH%;%install_root%\lib
 
-cl %PYTORCH_ROOT%\.ci\pytorch\test_example_code\simple-torch-test.cpp c10.lib torch_cpu.lib /EHsc /std:c++17
+cl %PYTORCH_ROOT%\.ci\pytorch\test_example_code\simple-torch-test.cpp c10.lib torch_cpu.lib /EHsc /std:c++20
 if ERRORLEVEL 1 exit /b 1
 
 .\simple-torch-test.exe
 if ERRORLEVEL 1 exit /b 1
 
-cl %PYTORCH_ROOT%\.ci\pytorch\test_example_code\check-torch-mkl.cpp c10.lib torch_cpu.lib /EHsc /std:c++17
+cl %PYTORCH_ROOT%\.ci\pytorch\test_example_code\check-torch-mkl.cpp c10.lib torch_cpu.lib /EHsc /std:c++20
 if ERRORLEVEL 1 exit /b 1
 
 .\check-torch-mkl.exe
@@ -148,7 +157,7 @@ if "%NVIDIA_GPU_EXISTS%" == "0" (
     goto end
 )
 
-cl %PYTORCH_ROOT%\.ci\pytorch\test_example_code\check-torch-cuda.cpp torch_cpu.lib c10.lib torch_cuda.lib /EHsc /std:c++17 /link /INCLUDE:?warp_size@cuda@at@@YAHXZ
+cl %PYTORCH_ROOT%\.ci\pytorch\test_example_code\check-torch-cuda.cpp torch_cpu.lib c10.lib torch_cuda.lib /EHsc /std:c++20 /link /INCLUDE:?warp_size@cuda@at@@YAHXZ
 .\check-torch-cuda.exe
 if ERRORLEVEL 1 exit /b 1
 
diff --git a/.ci/wheel/build_wheel.sh b/.ci/wheel/build_wheel.sh
index afd1faf2a5f7c..2563d5ba31765 100755
--- a/.ci/wheel/build_wheel.sh
+++ b/.ci/wheel/build_wheel.sh
@@ -97,7 +97,7 @@ fi
 whl_tmp_dir="${MAC_PACKAGE_WORK_DIR}/dist"
 mkdir -p "$whl_tmp_dir"
 
-mac_version='macosx-11.0-arm64'
+mac_version='macosx-14.0-arm64'
 libtorch_arch='arm64'
 
 # Create a consistent wheel package name to rename the wheel to
@@ -125,27 +125,19 @@ popd
 
 export TH_BINARY_BUILD=1
 export INSTALL_TEST=0 # dont install test binaries into site-packages
-export MACOSX_DEPLOYMENT_TARGET=11.0
+export MACOSX_DEPLOYMENT_TARGET=14.0
 
 EXTRA_CONDA_INSTALL_FLAGS=""
 CONDA_ENV_CREATE_FLAGS=""
 RENAME_WHEEL=false
 VERIFY_WHEELNAME=true
 case $desired_python in
-    3.14t)
-        echo "Using 3.14 deps"
-        NUMPY_PINNED_VERSION="==2.1.0"
-        ;;
-    3.14)
-        echo "Using 3.14t deps"
-        NUMPY_PINNED_VERSION="==2.1.0"
-        ;;
-    3.13t)
-        echo "Using 3.13t deps"
-        NUMPY_PINNED_VERSION="==2.1.0"
+    3.14*)
+        echo "Using ${desired_python} deps"
+        NUMPY_PINNED_VERSION="==2.3.4"
         ;;
-    3.13)
-        echo "Using 3.13 deps"
+    3.13*)
+        echo "Using ${desired_python} deps"
         NUMPY_PINNED_VERSION="==2.1.0"
         ;;
     3.12)
@@ -179,6 +171,7 @@ retry pip install "${PINNED_PACKAGES[@]}" -r "${pytorch_rootdir}/requirements.tx
 if [[ -d "/opt/llvm-openmp" ]]; then
   export OMP_PREFIX=/opt/llvm-openmp
 else
+  echo "libomp not found, installing via brew"
   retry brew install libomp
 fi
 
@@ -217,7 +210,7 @@ if [[ -z "$BUILD_PYTHONLESS" && $RENAME_WHEEL == true  ]]; then
     # Copy the whl to a final destination before tests are run
     echo "Renaming Wheel file: $wheel_filename_gen to $wheel_filename_new"
     cp "$whl_tmp_dir/$wheel_filename_gen" "$PYTORCH_FINAL_PACKAGE_DIR/$wheel_filename_new"
-elif [[ $RENAME_WHEEL == false ]]; then
+elif [[ -z "$BUILD_PYTHONLESS" && $RENAME_WHEEL == false ]]; then
     echo "Copying Wheel file: $wheel_filename_gen to $PYTORCH_FINAL_PACKAGE_DIR"
     cp "$whl_tmp_dir/$wheel_filename_gen" "$PYTORCH_FINAL_PACKAGE_DIR/$wheel_filename_gen"
     if [[ "$VERIFY_WHEELNAME" == "true" && "$wheel_filename_gen" != "$wheel_filename_new" ]]; then
diff --git a/.circleci/.gitignore b/.circleci/.gitignore
deleted file mode 100644
index c2153925851f5..0000000000000
--- a/.circleci/.gitignore
+++ /dev/null
@@ -1,2 +0,0 @@
-*.svg
-*.png
diff --git a/.circleci/README.md b/.circleci/README.md
deleted file mode 100644
index 24dde8b47666f..0000000000000
--- a/.circleci/README.md
+++ /dev/null
@@ -1,4 +0,0 @@
-Warning
-=======
-
-PyTorch migration from CircleCI to github actions has been completed. All continuous integration & deployment workflows are defined in  `.github/workflows` folder
diff --git a/.circleci/codegen_validation/compare_normalized_yaml.sh b/.circleci/codegen_validation/compare_normalized_yaml.sh
deleted file mode 100755
index af2a8268d3c55..0000000000000
--- a/.circleci/codegen_validation/compare_normalized_yaml.sh
+++ /dev/null
@@ -1,17 +0,0 @@
-#!/bin/bash -xe
-
-
-YAML_FILENAME=verbatim-sources/workflows-pytorch-ge-config-tests.yml
-DIFF_TOOL=meld
-
-
-# Allows this script to be invoked from any directory:
-cd $(dirname "$0")
-
-pushd ..
-
-
-$DIFF_TOOL $YAML_FILENAME <(./codegen_validation/normalize_yaml_fragment.py < $YAML_FILENAME)
-
-
-popd
diff --git a/.circleci/codegen_validation/normalize_yaml_fragment.py b/.circleci/codegen_validation/normalize_yaml_fragment.py
deleted file mode 100755
index 232eaa833b932..0000000000000
--- a/.circleci/codegen_validation/normalize_yaml_fragment.py
+++ /dev/null
@@ -1,26 +0,0 @@
-#!/usr/bin/env python3
-
-import os
-import sys
-
-import yaml
-
-
-# Need to import modules that lie on an upward-relative path
-sys.path.append(os.path.dirname(sys.path[0]))
-
-import cimodel.lib.miniyaml as miniyaml
-
-
-def regurgitate(depth, use_pyyaml_formatter=False):
-    data = yaml.safe_load(sys.stdin)
-
-    if use_pyyaml_formatter:
-        output = yaml.dump(data, sort_keys=True)
-        sys.stdout.write(output)
-    else:
-        miniyaml.render(sys.stdout, data, depth)
-
-
-if __name__ == "__main__":
-    regurgitate(3)
diff --git a/.circleci/codegen_validation/overwrite_with_normalized.sh b/.circleci/codegen_validation/overwrite_with_normalized.sh
deleted file mode 100755
index 7665984cedb69..0000000000000
--- a/.circleci/codegen_validation/overwrite_with_normalized.sh
+++ /dev/null
@@ -1,15 +0,0 @@
-#!/bin/bash -xe
-
-YAML_FILENAME=$1
-
-# Allows this script to be invoked from any directory:
-cd $(dirname "$0")
-
-pushd ..
-
-TEMP_FILENAME=$(mktemp)
-
-cat $YAML_FILENAME | ./codegen_validation/normalize_yaml_fragment.py > $TEMP_FILENAME
-mv $TEMP_FILENAME $YAML_FILENAME
-
-popd
diff --git a/.circleci/scripts/README.md b/.circleci/scripts/README.md
deleted file mode 100644
index c06504fd43042..0000000000000
--- a/.circleci/scripts/README.md
+++ /dev/null
@@ -1,4 +0,0 @@
-All the scripts in this directory are callable from `~/workspace/.circleci/scripts/foo.sh`.
-Don't try to call them as `.circleci/scripts/foo.sh`, that won't
-(necessarily) work.  See Note [Workspace for CircleCI scripts] in
-job-specs-setup.yml for more details.
diff --git a/.circleci/scripts/driver_update.bat b/.circleci/scripts/driver_update.bat
deleted file mode 100644
index fb87743666213..0000000000000
--- a/.circleci/scripts/driver_update.bat
+++ /dev/null
@@ -1,8 +0,0 @@
-set "DRIVER_DOWNLOAD_LINK=https://s3.amazonaws.com/ossci-windows/452.39-data-center-tesla-desktop-win10-64bit-international.exe"
-curl --retry 3 --retry-all-errors -kL %DRIVER_DOWNLOAD_LINK% --output 452.39-data-center-tesla-desktop-win10-64bit-international.exe
-if errorlevel 1 exit /b 1
-
-start /wait 452.39-data-center-tesla-desktop-win10-64bit-international.exe -s -noreboot
-if errorlevel 1 exit /b 1
-
-del 452.39-data-center-tesla-desktop-win10-64bit-international.exe || ver > NUL
diff --git a/.circleci/scripts/publish_android_snapshot.sh b/.circleci/scripts/publish_android_snapshot.sh
deleted file mode 100755
index 352b9d54f871a..0000000000000
--- a/.circleci/scripts/publish_android_snapshot.sh
+++ /dev/null
@@ -1,46 +0,0 @@
-#!/usr/bin/env bash
-# DO NOT ADD 'set -x' not to reveal CircleCI secret context environment variables
-set -eu -o pipefail
-
-export ANDROID_NDK_HOME=/opt/ndk
-export ANDROID_HOME=/opt/android/sdk
-
-export GRADLE_VERSION=6.8.3
-export GRADLE_HOME=/opt/gradle/gradle-$GRADLE_VERSION
-export GRADLE_PATH=$GRADLE_HOME/bin/gradle
-
-echo "BUILD_ENVIRONMENT:$BUILD_ENVIRONMENT"
-ls -la ~/workspace
-
-GRADLE_PROPERTIES=~/workspace/android/gradle.properties
-
-IS_SNAPSHOT="$(grep 'VERSION_NAME=[0-9\.]\+-SNAPSHOT' "$GRADLE_PROPERTIES")"
-echo "IS_SNAPSHOT:$IS_SNAPSHOT"
-
-if [ -z "$IS_SNAPSHOT" ]; then
-  echo "Error: version is not snapshot."
-elif [ -z "$SONATYPE_NEXUS_USERNAME" ]; then
-  echo "Error: missing env variable SONATYPE_NEXUS_USERNAME."
-elif [ -z "$SONATYPE_NEXUS_PASSWORD" ]; then
-  echo "Error: missing env variable SONATYPE_NEXUS_PASSWORD."
-elif [ -z "$ANDROID_SIGN_KEY" ]; then
-  echo "Error: missing env variable ANDROID_SIGN_KEY."
-elif [ -z "$ANDROID_SIGN_PASS" ]; then
-  echo "Error: missing env variable ANDROID_SIGN_PASS."
-else
-  GRADLE_LOCAL_PROPERTIES=~/workspace/android/local.properties
-  rm -f $GRADLE_LOCAL_PROPERTIES
-
-  echo "sdk.dir=/opt/android/sdk" >> $GRADLE_LOCAL_PROPERTIES
-  echo "ndk.dir=/opt/ndk" >> $GRADLE_LOCAL_PROPERTIES
-
-  echo "SONATYPE_NEXUS_USERNAME=${SONATYPE_NEXUS_USERNAME}" >> $GRADLE_PROPERTIES
-  echo "mavenCentralRepositoryUsername=${SONATYPE_NEXUS_USERNAME}" >> $GRADLE_PROPERTIES
-  echo "SONATYPE_NEXUS_PASSWORD=${SONATYPE_NEXUS_PASSWORD}" >> $GRADLE_PROPERTIES
-  echo "mavenCentralRepositoryPassword=${SONATYPE_NEXUS_PASSWORD}" >> $GRADLE_PROPERTIES
-
-  echo "signing.keyId=${ANDROID_SIGN_KEY}" >> $GRADLE_PROPERTIES
-  echo "signing.password=${ANDROID_SIGN_PASS}" >> $GRADLE_PROPERTIES
-
-  $GRADLE_PATH -p ~/workspace/android/ uploadArchives
-fi
diff --git a/.circleci/windows-jni/include/jni.h b/.circleci/windows-jni/include/jni.h
deleted file mode 100644
index 1fdc6f3ad69d3..0000000000000
--- a/.circleci/windows-jni/include/jni.h
+++ /dev/null
@@ -1,1131 +0,0 @@
-/*
- * Copyright (C) 2006 The Android Open Source Project
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *      http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-/*
- * JNI specification, as defined by Sun:
- * http://java.sun.com/javase/6/docs/technotes/guides/jni/spec/jniTOC.html
- *
- * Everything here is expected to be VM-neutral.
- */
-
-#ifndef JNI_H_
-#define JNI_H_
-
-#include <stdarg.h>
-#include <stdint.h>
-
-/* Primitive types that match up with Java equivalents. */
-typedef uint8_t  jboolean; /* unsigned 8 bits */
-typedef int8_t   jbyte;    /* signed 8 bits */
-typedef uint16_t jchar;    /* unsigned 16 bits */
-typedef int16_t  jshort;   /* signed 16 bits */
-typedef int32_t  jint;     /* signed 32 bits */
-typedef int64_t  jlong;    /* signed 64 bits */
-typedef float    jfloat;   /* 32-bit IEEE 754 */
-typedef double   jdouble;  /* 64-bit IEEE 754 */
-
-/* "cardinal indices and sizes" */
-typedef jint     jsize;
-
-#ifdef __cplusplus
-/*
- * Reference types, in C++
- */
-class _jobject {};
-class _jclass : public _jobject {};
-class _jstring : public _jobject {};
-class _jarray : public _jobject {};
-class _jobjectArray : public _jarray {};
-class _jbooleanArray : public _jarray {};
-class _jbyteArray : public _jarray {};
-class _jcharArray : public _jarray {};
-class _jshortArray : public _jarray {};
-class _jintArray : public _jarray {};
-class _jlongArray : public _jarray {};
-class _jfloatArray : public _jarray {};
-class _jdoubleArray : public _jarray {};
-class _jthrowable : public _jobject {};
-
-typedef _jobject*       jobject;
-typedef _jclass*        jclass;
-typedef _jstring*       jstring;
-typedef _jarray*        jarray;
-typedef _jobjectArray*  jobjectArray;
-typedef _jbooleanArray* jbooleanArray;
-typedef _jbyteArray*    jbyteArray;
-typedef _jcharArray*    jcharArray;
-typedef _jshortArray*   jshortArray;
-typedef _jintArray*     jintArray;
-typedef _jlongArray*    jlongArray;
-typedef _jfloatArray*   jfloatArray;
-typedef _jdoubleArray*  jdoubleArray;
-typedef _jthrowable*    jthrowable;
-typedef _jobject*       jweak;
-
-
-#else /* not __cplusplus */
-
-/*
- * Reference types, in C.
- */
-typedef void*           jobject;
-typedef jobject         jclass;
-typedef jobject         jstring;
-typedef jobject         jarray;
-typedef jarray          jobjectArray;
-typedef jarray          jbooleanArray;
-typedef jarray          jbyteArray;
-typedef jarray          jcharArray;
-typedef jarray          jshortArray;
-typedef jarray          jintArray;
-typedef jarray          jlongArray;
-typedef jarray          jfloatArray;
-typedef jarray          jdoubleArray;
-typedef jobject         jthrowable;
-typedef jobject         jweak;
-
-#endif /* not __cplusplus */
-
-struct _jfieldID;                       /* opaque structure */
-typedef struct _jfieldID* jfieldID;     /* field IDs */
-
-struct _jmethodID;                      /* opaque structure */
-typedef struct _jmethodID* jmethodID;   /* method IDs */
-
-struct JNIInvokeInterface;
-
-typedef union jvalue {
-    jboolean    z;
-    jbyte       b;
-    jchar       c;
-    jshort      s;
-    jint        i;
-    jlong       j;
-    jfloat      f;
-    jdouble     d;
-    jobject     l;
-} jvalue;
-
-typedef enum jobjectRefType {
-    JNIInvalidRefType = 0,
-    JNILocalRefType = 1,
-    JNIGlobalRefType = 2,
-    JNIWeakGlobalRefType = 3
-} jobjectRefType;
-
-typedef struct {
-    const char* name;
-    const char* signature;
-    void*       fnPtr;
-} JNINativeMethod;
-
-struct _JNIEnv;
-struct _JavaVM;
-typedef const struct JNINativeInterface* C_JNIEnv;
-
-#if defined(__cplusplus)
-typedef _JNIEnv JNIEnv;
-typedef _JavaVM JavaVM;
-#else
-typedef const struct JNINativeInterface* JNIEnv;
-typedef const struct JNIInvokeInterface* JavaVM;
-#endif
-
-/*
- * Table of interface function pointers.
- */
-struct JNINativeInterface {
-    void*       reserved0;
-    void*       reserved1;
-    void*       reserved2;
-    void*       reserved3;
-
-    jint        (*GetVersion)(JNIEnv *);
-
-    jclass      (*DefineClass)(JNIEnv*, const char*, jobject, const jbyte*,
-                        jsize);
-    jclass      (*FindClass)(JNIEnv*, const char*);
-
-    jmethodID   (*FromReflectedMethod)(JNIEnv*, jobject);
-    jfieldID    (*FromReflectedField)(JNIEnv*, jobject);
-    /* spec doesn't show jboolean parameter */
-    jobject     (*ToReflectedMethod)(JNIEnv*, jclass, jmethodID, jboolean);
-
-    jclass      (*GetSuperclass)(JNIEnv*, jclass);
-    jboolean    (*IsAssignableFrom)(JNIEnv*, jclass, jclass);
-
-    /* spec doesn't show jboolean parameter */
-    jobject     (*ToReflectedField)(JNIEnv*, jclass, jfieldID, jboolean);
-
-    jint        (*Throw)(JNIEnv*, jthrowable);
-    jint        (*ThrowNew)(JNIEnv *, jclass, const char *);
-    jthrowable  (*ExceptionOccurred)(JNIEnv*);
-    void        (*ExceptionDescribe)(JNIEnv*);
-    void        (*ExceptionClear)(JNIEnv*);
-    void        (*FatalError)(JNIEnv*, const char*);
-
-    jint        (*PushLocalFrame)(JNIEnv*, jint);
-    jobject     (*PopLocalFrame)(JNIEnv*, jobject);
-
-    jobject     (*NewGlobalRef)(JNIEnv*, jobject);
-    void        (*DeleteGlobalRef)(JNIEnv*, jobject);
-    void        (*DeleteLocalRef)(JNIEnv*, jobject);
-    jboolean    (*IsSameObject)(JNIEnv*, jobject, jobject);
-
-    jobject     (*NewLocalRef)(JNIEnv*, jobject);
-    jint        (*EnsureLocalCapacity)(JNIEnv*, jint);
-
-    jobject     (*AllocObject)(JNIEnv*, jclass);
-    jobject     (*NewObject)(JNIEnv*, jclass, jmethodID, ...);
-    jobject     (*NewObjectV)(JNIEnv*, jclass, jmethodID, va_list);
-    jobject     (*NewObjectA)(JNIEnv*, jclass, jmethodID, jvalue*);
-
-    jclass      (*GetObjectClass)(JNIEnv*, jobject);
-    jboolean    (*IsInstanceOf)(JNIEnv*, jobject, jclass);
-    jmethodID   (*GetMethodID)(JNIEnv*, jclass, const char*, const char*);
-
-    jobject     (*CallObjectMethod)(JNIEnv*, jobject, jmethodID, ...);
-    jobject     (*CallObjectMethodV)(JNIEnv*, jobject, jmethodID, va_list);
-    jobject     (*CallObjectMethodA)(JNIEnv*, jobject, jmethodID, jvalue*);
-    jboolean    (*CallBooleanMethod)(JNIEnv*, jobject, jmethodID, ...);
-    jboolean    (*CallBooleanMethodV)(JNIEnv*, jobject, jmethodID, va_list);
-    jboolean    (*CallBooleanMethodA)(JNIEnv*, jobject, jmethodID, jvalue*);
-    jbyte       (*CallByteMethod)(JNIEnv*, jobject, jmethodID, ...);
-    jbyte       (*CallByteMethodV)(JNIEnv*, jobject, jmethodID, va_list);
-    jbyte       (*CallByteMethodA)(JNIEnv*, jobject, jmethodID, jvalue*);
-    jchar       (*CallCharMethod)(JNIEnv*, jobject, jmethodID, ...);
-    jchar       (*CallCharMethodV)(JNIEnv*, jobject, jmethodID, va_list);
-    jchar       (*CallCharMethodA)(JNIEnv*, jobject, jmethodID, jvalue*);
-    jshort      (*CallShortMethod)(JNIEnv*, jobject, jmethodID, ...);
-    jshort      (*CallShortMethodV)(JNIEnv*, jobject, jmethodID, va_list);
-    jshort      (*CallShortMethodA)(JNIEnv*, jobject, jmethodID, jvalue*);
-    jint        (*CallIntMethod)(JNIEnv*, jobject, jmethodID, ...);
-    jint        (*CallIntMethodV)(JNIEnv*, jobject, jmethodID, va_list);
-    jint        (*CallIntMethodA)(JNIEnv*, jobject, jmethodID, jvalue*);
-    jlong       (*CallLongMethod)(JNIEnv*, jobject, jmethodID, ...);
-    jlong       (*CallLongMethodV)(JNIEnv*, jobject, jmethodID, va_list);
-    jlong       (*CallLongMethodA)(JNIEnv*, jobject, jmethodID, jvalue*);
-    jfloat      (*CallFloatMethod)(JNIEnv*, jobject, jmethodID, ...);
-    jfloat      (*CallFloatMethodV)(JNIEnv*, jobject, jmethodID, va_list);
-    jfloat      (*CallFloatMethodA)(JNIEnv*, jobject, jmethodID, jvalue*);
-    jdouble     (*CallDoubleMethod)(JNIEnv*, jobject, jmethodID, ...);
-    jdouble     (*CallDoubleMethodV)(JNIEnv*, jobject, jmethodID, va_list);
-    jdouble     (*CallDoubleMethodA)(JNIEnv*, jobject, jmethodID, jvalue*);
-    void        (*CallVoidMethod)(JNIEnv*, jobject, jmethodID, ...);
-    void        (*CallVoidMethodV)(JNIEnv*, jobject, jmethodID, va_list);
-    void        (*CallVoidMethodA)(JNIEnv*, jobject, jmethodID, jvalue*);
-
- jobject     (*CallNonvirtualObjectMethod)(JNIEnv*, jobject, jclass,
-                        jmethodID, ...);
-    jobject     (*CallNonvirtualObjectMethodV)(JNIEnv*, jobject, jclass,
-                        jmethodID, va_list);
-    jobject     (*CallNonvirtualObjectMethodA)(JNIEnv*, jobject, jclass,
-                        jmethodID, jvalue*);
-    jboolean    (*CallNonvirtualBooleanMethod)(JNIEnv*, jobject, jclass,
-                        jmethodID, ...);
-    jboolean    (*CallNonvirtualBooleanMethodV)(JNIEnv*, jobject, jclass,
-                         jmethodID, va_list);
-    jboolean    (*CallNonvirtualBooleanMethodA)(JNIEnv*, jobject, jclass,
-                         jmethodID, jvalue*);
-    jbyte       (*CallNonvirtualByteMethod)(JNIEnv*, jobject, jclass,
-                        jmethodID, ...);
-    jbyte       (*CallNonvirtualByteMethodV)(JNIEnv*, jobject, jclass,
-                        jmethodID, va_list);
-    jbyte       (*CallNonvirtualByteMethodA)(JNIEnv*, jobject, jclass,
-                        jmethodID, jvalue*);
-    jchar       (*CallNonvirtualCharMethod)(JNIEnv*, jobject, jclass,
-                        jmethodID, ...);
-    jchar       (*CallNonvirtualCharMethodV)(JNIEnv*, jobject, jclass,
-                        jmethodID, va_list);
-    jchar       (*CallNonvirtualCharMethodA)(JNIEnv*, jobject, jclass,
-                        jmethodID, jvalue*);
-    jshort      (*CallNonvirtualShortMethod)(JNIEnv*, jobject, jclass,
-                        jmethodID, ...);
-    jshort      (*CallNonvirtualShortMethodV)(JNIEnv*, jobject, jclass,
-                        jmethodID, va_list);
-    jshort      (*CallNonvirtualShortMethodA)(JNIEnv*, jobject, jclass,
-                        jmethodID, jvalue*);
-    jint        (*CallNonvirtualIntMethod)(JNIEnv*, jobject, jclass,
-                        jmethodID, ...);
-    jint        (*CallNonvirtualIntMethodV)(JNIEnv*, jobject, jclass,
-                        jmethodID, va_list);
-    jint        (*CallNonvirtualIntMethodA)(JNIEnv*, jobject, jclass,
-                        jmethodID, jvalue*);
-    jlong       (*CallNonvirtualLongMethod)(JNIEnv*, jobject, jclass,
-                        jmethodID, ...);
-    jlong       (*CallNonvirtualLongMethodV)(JNIEnv*, jobject, jclass,
-                        jmethodID, va_list);
-    jlong       (*CallNonvirtualLongMethodA)(JNIEnv*, jobject, jclass,
-                        jmethodID, jvalue*);
-    jfloat      (*CallNonvirtualFloatMethod)(JNIEnv*, jobject, jclass,
-                        jmethodID, ...);
-    jfloat      (*CallNonvirtualFloatMethodV)(JNIEnv*, jobject, jclass,
-                        jmethodID, va_list);
-    jfloat      (*CallNonvirtualFloatMethodA)(JNIEnv*, jobject, jclass,
-                        jmethodID, jvalue*);
-    jdouble     (*CallNonvirtualDoubleMethod)(JNIEnv*, jobject, jclass,
-                        jmethodID, ...);
-    jdouble     (*CallNonvirtualDoubleMethodV)(JNIEnv*, jobject, jclass,
-                        jmethodID, va_list);
-    jdouble     (*CallNonvirtualDoubleMethodA)(JNIEnv*, jobject, jclass,
-                        jmethodID, jvalue*);
-    void        (*CallNonvirtualVoidMethod)(JNIEnv*, jobject, jclass,
-                        jmethodID, ...);
-    void        (*CallNonvirtualVoidMethodV)(JNIEnv*, jobject, jclass,
-                        jmethodID, va_list);
-    void        (*CallNonvirtualVoidMethodA)(JNIEnv*, jobject, jclass,
-                        jmethodID, jvalue*);
-
-    jfieldID    (*GetFieldID)(JNIEnv*, jclass, const char*, const char*);
-
-    jobject     (*GetObjectField)(JNIEnv*, jobject, jfieldID);
-    jboolean    (*GetBooleanField)(JNIEnv*, jobject, jfieldID);
-    jbyte       (*GetByteField)(JNIEnv*, jobject, jfieldID);
-    jchar       (*GetCharField)(JNIEnv*, jobject, jfieldID);
-    jshort      (*GetShortField)(JNIEnv*, jobject, jfieldID);
-    jint        (*GetIntField)(JNIEnv*, jobject, jfieldID);
-    jlong       (*GetLongField)(JNIEnv*, jobject, jfieldID);
-    jfloat      (*GetFloatField)(JNIEnv*, jobject, jfieldID);
-    jdouble     (*GetDoubleField)(JNIEnv*, jobject, jfieldID);
-
-    void        (*SetObjectField)(JNIEnv*, jobject, jfieldID, jobject);
-    void        (*SetBooleanField)(JNIEnv*, jobject, jfieldID, jboolean);
-    void        (*SetByteField)(JNIEnv*, jobject, jfieldID, jbyte);
-    void        (*SetCharField)(JNIEnv*, jobject, jfieldID, jchar);
-    void        (*SetShortField)(JNIEnv*, jobject, jfieldID, jshort);
-    void        (*SetIntField)(JNIEnv*, jobject, jfieldID, jint);
-    void        (*SetLongField)(JNIEnv*, jobject, jfieldID, jlong);
-    void        (*SetFloatField)(JNIEnv*, jobject, jfieldID, jfloat);
-    void        (*SetDoubleField)(JNIEnv*, jobject, jfieldID, jdouble);
-
-    jmethodID   (*GetStaticMethodID)(JNIEnv*, jclass, const char*, const char*);
-
-    jobject     (*CallStaticObjectMethod)(JNIEnv*, jclass, jmethodID, ...);
-    jobject     (*CallStaticObjectMethodV)(JNIEnv*, jclass, jmethodID, va_list);
-    jobject     (*CallStaticObjectMethodA)(JNIEnv*, jclass, jmethodID, jvalue*);
-    jboolean    (*CallStaticBooleanMethod)(JNIEnv*, jclass, jmethodID, ...);
-    jboolean    (*CallStaticBooleanMethodV)(JNIEnv*, jclass, jmethodID,
-                        va_list);
-    jboolean    (*CallStaticBooleanMethodA)(JNIEnv*, jclass, jmethodID,
-                        jvalue*);
-    jbyte       (*CallStaticByteMethod)(JNIEnv*, jclass, jmethodID, ...);
-    jbyte       (*CallStaticByteMethodV)(JNIEnv*, jclass, jmethodID, va_list);
-    jbyte       (*CallStaticByteMethodA)(JNIEnv*, jclass, jmethodID, jvalue*);
-    jchar       (*CallStaticCharMethod)(JNIEnv*, jclass, jmethodID, ...);
-    jchar       (*CallStaticCharMethodV)(JNIEnv*, jclass, jmethodID, va_list);
-    jchar       (*CallStaticCharMethodA)(JNIEnv*, jclass, jmethodID, jvalue*);
-    jshort      (*CallStaticShortMethod)(JNIEnv*, jclass, jmethodID, ...);
-    jshort      (*CallStaticShortMethodV)(JNIEnv*, jclass, jmethodID, va_list);
-    jshort      (*CallStaticShortMethodA)(JNIEnv*, jclass, jmethodID, jvalue*);
-    jint        (*CallStaticIntMethod)(JNIEnv*, jclass, jmethodID, ...);
-    jint        (*CallStaticIntMethodV)(JNIEnv*, jclass, jmethodID, va_list);
-    jint        (*CallStaticIntMethodA)(JNIEnv*, jclass, jmethodID, jvalue*);
-    jlong       (*CallStaticLongMethod)(JNIEnv*, jclass, jmethodID, ...);
-    jlong       (*CallStaticLongMethodV)(JNIEnv*, jclass, jmethodID, va_list);
-    jlong       (*CallStaticLongMethodA)(JNIEnv*, jclass, jmethodID, jvalue*);
-    jfloat      (*CallStaticFloatMethod)(JNIEnv*, jclass, jmethodID, ...);
-    jfloat      (*CallStaticFloatMethodV)(JNIEnv*, jclass, jmethodID, va_list);
-    jfloat      (*CallStaticFloatMethodA)(JNIEnv*, jclass, jmethodID, jvalue*);
-    jdouble     (*CallStaticDoubleMethod)(JNIEnv*, jclass, jmethodID, ...);
-    jdouble     (*CallStaticDoubleMethodV)(JNIEnv*, jclass, jmethodID, va_list);
-    jdouble     (*CallStaticDoubleMethodA)(JNIEnv*, jclass, jmethodID, jvalue*);
-    void        (*CallStaticVoidMethod)(JNIEnv*, jclass, jmethodID, ...);
-    void        (*CallStaticVoidMethodV)(JNIEnv*, jclass, jmethodID, va_list);
-    void        (*CallStaticVoidMethodA)(JNIEnv*, jclass, jmethodID, jvalue*);
-
-    jfieldID    (*GetStaticFieldID)(JNIEnv*, jclass, const char*,
-                        const char*);
-
-    jobject     (*GetStaticObjectField)(JNIEnv*, jclass, jfieldID);
-    jboolean    (*GetStaticBooleanField)(JNIEnv*, jclass, jfieldID);
-    jbyte       (*GetStaticByteField)(JNIEnv*, jclass, jfieldID);
-    jchar       (*GetStaticCharField)(JNIEnv*, jclass, jfieldID);
-    jshort      (*GetStaticShortField)(JNIEnv*, jclass, jfieldID);
-    jint        (*GetStaticIntField)(JNIEnv*, jclass, jfieldID);
-    jlong       (*GetStaticLongField)(JNIEnv*, jclass, jfieldID);
-    jfloat      (*GetStaticFloatField)(JNIEnv*, jclass, jfieldID);
-    jdouble     (*GetStaticDoubleField)(JNIEnv*, jclass, jfieldID);
-
-    void        (*SetStaticObjectField)(JNIEnv*, jclass, jfieldID, jobject);
-    void        (*SetStaticBooleanField)(JNIEnv*, jclass, jfieldID, jboolean);
-    void        (*SetStaticByteField)(JNIEnv*, jclass, jfieldID, jbyte);
-    void        (*SetStaticCharField)(JNIEnv*, jclass, jfieldID, jchar);
-    void        (*SetStaticShortField)(JNIEnv*, jclass, jfieldID, jshort);
-    void        (*SetStaticIntField)(JNIEnv*, jclass, jfieldID, jint);
-    void        (*SetStaticLongField)(JNIEnv*, jclass, jfieldID, jlong);
-    void        (*SetStaticFloatField)(JNIEnv*, jclass, jfieldID, jfloat);
-    void        (*SetStaticDoubleField)(JNIEnv*, jclass, jfieldID, jdouble);
-
-    jstring     (*NewString)(JNIEnv*, const jchar*, jsize);
-    jsize       (*GetStringLength)(JNIEnv*, jstring);
-    const jchar* (*GetStringChars)(JNIEnv*, jstring, jboolean*);
-    void        (*ReleaseStringChars)(JNIEnv*, jstring, const jchar*);
-    jstring     (*NewStringUTF)(JNIEnv*, const char*);
-    jsize       (*GetStringUTFLength)(JNIEnv*, jstring);
-    /* JNI spec says this returns const jbyte*, but that's inconsistent */
-    const char* (*GetStringUTFChars)(JNIEnv*, jstring, jboolean*);
-    void        (*ReleaseStringUTFChars)(JNIEnv*, jstring, const char*);
-    jsize       (*GetArrayLength)(JNIEnv*, jarray);
-    jobjectArray (*NewObjectArray)(JNIEnv*, jsize, jclass, jobject);
-    jobject     (*GetObjectArrayElement)(JNIEnv*, jobjectArray, jsize);
-    void        (*SetObjectArrayElement)(JNIEnv*, jobjectArray, jsize, jobject);
-
-    jbooleanArray (*NewBooleanArray)(JNIEnv*, jsize);
-    jbyteArray    (*NewByteArray)(JNIEnv*, jsize);
-    jcharArray    (*NewCharArray)(JNIEnv*, jsize);
-    jshortArray   (*NewShortArray)(JNIEnv*, jsize);
-    jintArray     (*NewIntArray)(JNIEnv*, jsize);
-    jlongArray    (*NewLongArray)(JNIEnv*, jsize);
-    jfloatArray   (*NewFloatArray)(JNIEnv*, jsize);
-    jdoubleArray  (*NewDoubleArray)(JNIEnv*, jsize);
-
-    jboolean*   (*GetBooleanArrayElements)(JNIEnv*, jbooleanArray, jboolean*);
-    jbyte*      (*GetByteArrayElements)(JNIEnv*, jbyteArray, jboolean*);
-    jchar*      (*GetCharArrayElements)(JNIEnv*, jcharArray, jboolean*);
-    jshort*     (*GetShortArrayElements)(JNIEnv*, jshortArray, jboolean*);
-    jint*       (*GetIntArrayElements)(JNIEnv*, jintArray, jboolean*);
-    jlong*      (*GetLongArrayElements)(JNIEnv*, jlongArray, jboolean*);
-    jfloat*     (*GetFloatArrayElements)(JNIEnv*, jfloatArray, jboolean*);
-    jdouble*    (*GetDoubleArrayElements)(JNIEnv*, jdoubleArray, jboolean*);
-
-    void        (*ReleaseBooleanArrayElements)(JNIEnv*, jbooleanArray,
-                        jboolean*, jint);
-    void        (*ReleaseByteArrayElements)(JNIEnv*, jbyteArray,
-                        jbyte*, jint);
-    void        (*ReleaseCharArrayElements)(JNIEnv*, jcharArray,
-                        jchar*, jint);
-    void        (*ReleaseShortArrayElements)(JNIEnv*, jshortArray,
-                        jshort*, jint);
-    void        (*ReleaseIntArrayElements)(JNIEnv*, jintArray,
-                        jint*, jint);
-    void        (*ReleaseLongArrayElements)(JNIEnv*, jlongArray,
-                        jlong*, jint);
-    void        (*ReleaseFloatArrayElements)(JNIEnv*, jfloatArray,
-                        jfloat*, jint);
-    void        (*ReleaseDoubleArrayElements)(JNIEnv*, jdoubleArray,
-                        jdouble*, jint);
-
-    void        (*GetBooleanArrayRegion)(JNIEnv*, jbooleanArray,
-                        jsize, jsize, jboolean*);
-    void        (*GetByteArrayRegion)(JNIEnv*, jbyteArray,
-                        jsize, jsize, jbyte*);
-    void        (*GetCharArrayRegion)(JNIEnv*, jcharArray,
-                        jsize, jsize, jchar*);
-    void        (*GetShortArrayRegion)(JNIEnv*, jshortArray,
-                        jsize, jsize, jshort*);
-    void        (*GetIntArrayRegion)(JNIEnv*, jintArray,
-                        jsize, jsize, jint*);
-    void        (*GetLongArrayRegion)(JNIEnv*, jlongArray,
-                        jsize, jsize, jlong*);
-    void        (*GetFloatArrayRegion)(JNIEnv*, jfloatArray,
-                        jsize, jsize, jfloat*);
-    void        (*GetDoubleArrayRegion)(JNIEnv*, jdoubleArray,
-                        jsize, jsize, jdouble*);
-
-    /* spec shows these without const; some jni.h do, some don't */
-    void        (*SetBooleanArrayRegion)(JNIEnv*, jbooleanArray,
-                        jsize, jsize, const jboolean*);
-    void        (*SetByteArrayRegion)(JNIEnv*, jbyteArray,
-                        jsize, jsize, const jbyte*);
-    void        (*SetCharArrayRegion)(JNIEnv*, jcharArray,
-                        jsize, jsize, const jchar*);
-    void        (*SetShortArrayRegion)(JNIEnv*, jshortArray,
-                        jsize, jsize, const jshort*);
-    void        (*SetIntArrayRegion)(JNIEnv*, jintArray,
-                        jsize, jsize, const jint*);
-    void        (*SetLongArrayRegion)(JNIEnv*, jlongArray,
-                        jsize, jsize, const jlong*);
-    void        (*SetFloatArrayRegion)(JNIEnv*, jfloatArray,
-                        jsize, jsize, const jfloat*);
-    void        (*SetDoubleArrayRegion)(JNIEnv*, jdoubleArray,
-                        jsize, jsize, const jdouble*);
-
-    jint        (*RegisterNatives)(JNIEnv*, jclass, const JNINativeMethod*,
-                        jint);
-    jint        (*UnregisterNatives)(JNIEnv*, jclass);
-    jint        (*MonitorEnter)(JNIEnv*, jobject);
-    jint        (*MonitorExit)(JNIEnv*, jobject);
-    jint        (*GetJavaVM)(JNIEnv*, JavaVM**);
-
-    void        (*GetStringRegion)(JNIEnv*, jstring, jsize, jsize, jchar*);
-    void        (*GetStringUTFRegion)(JNIEnv*, jstring, jsize, jsize, char*);
-
-    void*       (*GetPrimitiveArrayCritical)(JNIEnv*, jarray, jboolean*);
-    void        (*ReleasePrimitiveArrayCritical)(JNIEnv*, jarray, void*, jint);
-
-    const jchar* (*GetStringCritical)(JNIEnv*, jstring, jboolean*);
-    void        (*ReleaseStringCritical)(JNIEnv*, jstring, const jchar*);
-
-    jweak       (*NewWeakGlobalRef)(JNIEnv*, jobject);
-    void        (*DeleteWeakGlobalRef)(JNIEnv*, jweak);
-
-    jboolean    (*ExceptionCheck)(JNIEnv*);
-    jobject     (*NewDirectByteBuffer)(JNIEnv*, void*, jlong);
-
-    void*       (*GetDirectBufferAddress)(JNIEnv*, jobject);
-    jlong       (*GetDirectBufferCapacity)(JNIEnv*, jobject);
-
-    /* added in JNI 1.6 */
-    jobjectRefType (*GetObjectRefType)(JNIEnv*, jobject);
-};
-
-/*
- * C++ object wrapper.
- *
- * This is usually overlaid on a C struct whose first element is a
- * JNINativeInterface*.  We rely somewhat on compiler behavior.
- */
-struct _JNIEnv {
-    /* do not rename this; it does not seem to be entirely opaque */
-    const struct JNINativeInterface* functions;
-
-#if defined(__cplusplus)
-    jint GetVersion()
-    { return functions->GetVersion(this); }
-
-    jclass DefineClass(const char *name, jobject loader, const jbyte* buf,
-        jsize bufLen)
-    { return functions->DefineClass(this, name, loader, buf, bufLen); }
-
-    jclass FindClass(const char* name)
-    { return functions->FindClass(this, name); }
-
-    jmethodID FromReflectedMethod(jobject method)
-    { return functions->FromReflectedMethod(this, method); }
-
-    jfieldID FromReflectedField(jobject field)
-    { return functions->FromReflectedField(this, field); }
-
-    jobject ToReflectedMethod(jclass cls, jmethodID methodID, jboolean isStatic)
-    { return functions->ToReflectedMethod(this, cls, methodID, isStatic); }
-
-    jclass GetSuperclass(jclass clazz)
-    { return functions->GetSuperclass(this, clazz); }
-
-    jboolean IsAssignableFrom(jclass clazz1, jclass clazz2)
-    { return functions->IsAssignableFrom(this, clazz1, clazz2); }
-
-    jobject ToReflectedField(jclass cls, jfieldID fieldID, jboolean isStatic)
-    { return functions->ToReflectedField(this, cls, fieldID, isStatic); }
-
-    jint Throw(jthrowable obj)
-    { return functions->Throw(this, obj); }
-
-    jint ThrowNew(jclass clazz, const char* message)
-    { return functions->ThrowNew(this, clazz, message); }
-
-    jthrowable ExceptionOccurred()
-    { return functions->ExceptionOccurred(this); }
-
-    void ExceptionDescribe()
-    { functions->ExceptionDescribe(this); }
-
-    void ExceptionClear()
-    { functions->ExceptionClear(this); }
-
-    void FatalError(const char* msg)
-    { functions->FatalError(this, msg); }
-
-    jint PushLocalFrame(jint capacity)
-    { return functions->PushLocalFrame(this, capacity); }
-
-    jobject PopLocalFrame(jobject result)
-    { return functions->PopLocalFrame(this, result); }
-
-    jobject NewGlobalRef(jobject obj)
-    { return functions->NewGlobalRef(this, obj); }
-
-    void DeleteGlobalRef(jobject globalRef)
-    { functions->DeleteGlobalRef(this, globalRef); }
-
-    void DeleteLocalRef(jobject localRef)
-    { functions->DeleteLocalRef(this, localRef); }
-
-    jboolean IsSameObject(jobject ref1, jobject ref2)
-    { return functions->IsSameObject(this, ref1, ref2); }
-
-    jobject NewLocalRef(jobject ref)
-    { return functions->NewLocalRef(this, ref); }
-
-    jint EnsureLocalCapacity(jint capacity)
-    { return functions->EnsureLocalCapacity(this, capacity); }
-
-    jobject AllocObject(jclass clazz)
-    { return functions->AllocObject(this, clazz); }
-
-    jobject NewObject(jclass clazz, jmethodID methodID, ...)
-    {
-        va_list args;
-        va_start(args, methodID);
-        jobject result = functions->NewObjectV(this, clazz, methodID, args);
-        va_end(args);
-        return result;
-    }
-
-    jobject NewObjectV(jclass clazz, jmethodID methodID, va_list args)
-    { return functions->NewObjectV(this, clazz, methodID, args); }
-
-    jobject NewObjectA(jclass clazz, jmethodID methodID, jvalue* args)
-    { return functions->NewObjectA(this, clazz, methodID, args); }
-
-    jclass GetObjectClass(jobject obj)
-    { return functions->GetObjectClass(this, obj); }
-
-    jboolean IsInstanceOf(jobject obj, jclass clazz)
-    { return functions->IsInstanceOf(this, obj, clazz); }
-
-    jmethodID GetMethodID(jclass clazz, const char* name, const char* sig)
-    { return functions->GetMethodID(this, clazz, name, sig); }
-
-#define CALL_TYPE_METHOD(_jtype, _jname)                                    \
-    _jtype Call##_jname##Method(jobject obj, jmethodID methodID, ...)       \
-    {                                                                       \
-        _jtype result;                                                      \
-        va_list args;                                                       \
-        va_start(args, methodID);                                           \
-        result = functions->Call##_jname##MethodV(this, obj, methodID,      \
-                    args);                                                  \
-        va_end(args);                                                       \
-        return result;                                                      \
-    }
-
-#define CALL_TYPE_METHODV(_jtype, _jname)                                   \
-    _jtype Call##_jname##MethodV(jobject obj, jmethodID methodID,           \
-        va_list args)                                                       \
-    { return functions->Call##_jname##MethodV(this, obj, methodID, args); }
-
-#define CALL_TYPE_METHODA(_jtype, _jname)                                   \
-    _jtype Call##_jname##MethodA(jobject obj, jmethodID methodID,           \
-        jvalue* args)                                                       \
-    { return functions->Call##_jname##MethodA(this, obj, methodID, args); }
-
-#define CALL_TYPE(_jtype, _jname)                                           \
-    CALL_TYPE_METHOD(_jtype, _jname)                                        \
-    CALL_TYPE_METHODV(_jtype, _jname)                                       \
-    CALL_TYPE_METHODA(_jtype, _jname)
-    CALL_TYPE(jobject, Object)
-    CALL_TYPE(jboolean, Boolean)
-    CALL_TYPE(jbyte, Byte)
-    CALL_TYPE(jchar, Char)
-    CALL_TYPE(jshort, Short)
-    CALL_TYPE(jint, Int)
-    CALL_TYPE(jlong, Long)
-    CALL_TYPE(jfloat, Float)
-    CALL_TYPE(jdouble, Double)
-
-    void CallVoidMethod(jobject obj, jmethodID methodID, ...)
-    {
-        va_list args;
-        va_start(args, methodID);
-        functions->CallVoidMethodV(this, obj, methodID, args);
-        va_end(args);
-    }
-
-    void CallVoidMethodV(jobject obj, jmethodID methodID, va_list args)
-    { functions->CallVoidMethodV(this, obj, methodID, args); }
-
-    void CallVoidMethodA(jobject obj, jmethodID methodID, jvalue* args)
-    { functions->CallVoidMethodA(this, obj, methodID, args); }
-#define CALL_NONVIRT_TYPE_METHOD(_jtype, _jname)                            \
-    _jtype CallNonvirtual##_jname##Method(jobject obj, jclass clazz,        \
-        jmethodID methodID, ...)                                            \
-    {                                                                       \
-        _jtype result;                                                      \
-        va_list args;                                                       \
-        va_start(args, methodID);                                           \
-        result = functions->CallNonvirtual##_jname##MethodV(this, obj,      \
-                    clazz, methodID, args);                                 \
-        va_end(args);                                                       \
-        return result;                                                      \
-    }
-#define CALL_NONVIRT_TYPE_METHODV(_jtype, _jname)                           \
-    _jtype CallNonvirtual##_jname##MethodV(jobject obj, jclass clazz,       \
-        jmethodID methodID, va_list args)                                   \
-    { return functions->CallNonvirtual##_jname##MethodV(this, obj, clazz,   \
-        methodID, args); }
-#define CALL_NONVIRT_TYPE_METHODA(_jtype, _jname)                           \
-    _jtype CallNonvirtual##_jname##MethodA(jobject obj, jclass clazz,       \
-        jmethodID methodID, jvalue* args)                                   \
-    { return functions->CallNonvirtual##_jname##MethodA(this, obj, clazz,   \
-        methodID, args); }
-#define CALL_NONVIRT_TYPE(_jtype, _jname)                                   \
-    CALL_NONVIRT_TYPE_METHOD(_jtype, _jname)                                \
-    CALL_NONVIRT_TYPE_METHODV(_jtype, _jname)                               \
-    CALL_NONVIRT_TYPE_METHODA(_jtype, _jname)
-    CALL_NONVIRT_TYPE(jobject, Object)
-    CALL_NONVIRT_TYPE(jboolean, Boolean)
-    CALL_NONVIRT_TYPE(jbyte, Byte)
-    CALL_NONVIRT_TYPE(jchar, Char)
-    CALL_NONVIRT_TYPE(jshort, Short)
-    CALL_NONVIRT_TYPE(jint, Int)
-    CALL_NONVIRT_TYPE(jlong, Long)
-    CALL_NONVIRT_TYPE(jfloat, Float)
-    CALL_NONVIRT_TYPE(jdouble, Double)
-    void CallNonvirtualVoidMethod(jobject obj, jclass clazz,
-        jmethodID methodID, ...)
-    {
-        va_list args;
-        va_start(args, methodID);
-        functions->CallNonvirtualVoidMethodV(this, obj, clazz, methodID, args);
-        va_end(args);
-    }
-    void CallNonvirtualVoidMethodV(jobject obj, jclass clazz,
-        jmethodID methodID, va_list args)
-    { functions->CallNonvirtualVoidMethodV(this, obj, clazz, methodID, args); }
-    void CallNonvirtualVoidMethodA(jobject obj, jclass clazz,
-        jmethodID methodID, jvalue* args)
-    { functions->CallNonvirtualVoidMethodA(this, obj, clazz, methodID, args); }
-    jfieldID GetFieldID(jclass clazz, const char* name, const char* sig)
-    { return functions->GetFieldID(this, clazz, name, sig); }
-    jobject GetObjectField(jobject obj, jfieldID fieldID)
-    { return functions->GetObjectField(this, obj, fieldID); }
-    jboolean GetBooleanField(jobject obj, jfieldID fieldID)
-    { return functions->GetBooleanField(this, obj, fieldID); }
-    jbyte GetByteField(jobject obj, jfieldID fieldID)
-    { return functions->GetByteField(this, obj, fieldID); }
-    jchar GetCharField(jobject obj, jfieldID fieldID)
-    { return functions->GetCharField(this, obj, fieldID); }
-    jshort GetShortField(jobject obj, jfieldID fieldID)
-    { return functions->GetShortField(this, obj, fieldID); }
-    jint GetIntField(jobject obj, jfieldID fieldID)
-    { return functions->GetIntField(this, obj, fieldID); }
-    jlong GetLongField(jobject obj, jfieldID fieldID)
-    { return functions->GetLongField(this, obj, fieldID); }
-    jfloat GetFloatField(jobject obj, jfieldID fieldID)
-    { return functions->GetFloatField(this, obj, fieldID); }
-    jdouble GetDoubleField(jobject obj, jfieldID fieldID)
-    { return functions->GetDoubleField(this, obj, fieldID); }
-    void SetObjectField(jobject obj, jfieldID fieldID, jobject value)
-    { functions->SetObjectField(this, obj, fieldID, value); }
-    void SetBooleanField(jobject obj, jfieldID fieldID, jboolean value)
-    { functions->SetBooleanField(this, obj, fieldID, value); }
-    void SetByteField(jobject obj, jfieldID fieldID, jbyte value)
-    { functions->SetByteField(this, obj, fieldID, value); }
-    void SetCharField(jobject obj, jfieldID fieldID, jchar value)
-    { functions->SetCharField(this, obj, fieldID, value); }
-    void SetShortField(jobject obj, jfieldID fieldID, jshort value)
-    { functions->SetShortField(this, obj, fieldID, value); }
-    void SetIntField(jobject obj, jfieldID fieldID, jint value)
-    { functions->SetIntField(this, obj, fieldID, value); }
-    void SetLongField(jobject obj, jfieldID fieldID, jlong value)
-    { functions->SetLongField(this, obj, fieldID, value); }
-    void SetFloatField(jobject obj, jfieldID fieldID, jfloat value)
-    { functions->SetFloatField(this, obj, fieldID, value); }
-    void SetDoubleField(jobject obj, jfieldID fieldID, jdouble value)
-    { functions->SetDoubleField(this, obj, fieldID, value); }
-    jmethodID GetStaticMethodID(jclass clazz, const char* name, const char* sig)
-    { return functions->GetStaticMethodID(this, clazz, name, sig); }
-
-#define CALL_STATIC_TYPE_METHOD(_jtype, _jname)                             \
-    _jtype CallStatic##_jname##Method(jclass clazz, jmethodID methodID,     \
-        ...)                                                                \
-    {                                                                       \
-        _jtype result;                                                      \
-        va_list args;                                                       \
-        va_start(args, methodID);                                           \
-        result = functions->CallStatic##_jname##MethodV(this, clazz,        \
-                    methodID, args);                                        \
-        va_end(args);                                                       \
-        return result;                                                      \
-    }
-#define CALL_STATIC_TYPE_METHODV(_jtype, _jname)                            \
-    _jtype CallStatic##_jname##MethodV(jclass clazz, jmethodID methodID,    \
-        va_list args)                                                       \
-    { return functions->CallStatic##_jname##MethodV(this, clazz, methodID,  \
-        args); }
-#define CALL_STATIC_TYPE_METHODA(_jtype, _jname)                            \
-    _jtype CallStatic##_jname##MethodA(jclass clazz, jmethodID methodID,    \
-        jvalue* args)                                                       \
-    { return functions->CallStatic##_jname##MethodA(this, clazz, methodID,  \
-        args); }
-
-#define CALL_STATIC_TYPE(_jtype, _jname)                                    \
-    CALL_STATIC_TYPE_METHOD(_jtype, _jname)                                 \
-    CALL_STATIC_TYPE_METHODV(_jtype, _jname)                                \
-    CALL_STATIC_TYPE_METHODA(_jtype, _jname)
-    CALL_STATIC_TYPE(jobject, Object)
-    CALL_STATIC_TYPE(jboolean, Boolean)
-    CALL_STATIC_TYPE(jbyte, Byte)
-    CALL_STATIC_TYPE(jchar, Char)
-    CALL_STATIC_TYPE(jshort, Short)
-    CALL_STATIC_TYPE(jint, Int)
-    CALL_STATIC_TYPE(jlong, Long)
-    CALL_STATIC_TYPE(jfloat, Float)
-    CALL_STATIC_TYPE(jdouble, Double)
-    void CallStaticVoidMethod(jclass clazz, jmethodID methodID, ...)
-    {
-        va_list args;
-        va_start(args, methodID);
-        functions->CallStaticVoidMethodV(this, clazz, methodID, args);
-        va_end(args);
-    }
-    void CallStaticVoidMethodV(jclass clazz, jmethodID methodID, va_list args)
-    { functions->CallStaticVoidMethodV(this, clazz, methodID, args); }
-    void CallStaticVoidMethodA(jclass clazz, jmethodID methodID, jvalue* args)
-    { functions->CallStaticVoidMethodA(this, clazz, methodID, args); }
-
-    jfieldID GetStaticFieldID(jclass clazz, const char* name, const char* sig)
-    { return functions->GetStaticFieldID(this, clazz, name, sig); }
-
-    jobject GetStaticObjectField(jclass clazz, jfieldID fieldID)
-    { return functions->GetStaticObjectField(this, clazz, fieldID); }
-    jboolean GetStaticBooleanField(jclass clazz, jfieldID fieldID)
-    { return functions->GetStaticBooleanField(this, clazz, fieldID); }
-    jbyte GetStaticByteField(jclass clazz, jfieldID fieldID)
-    { return functions->GetStaticByteField(this, clazz, fieldID); }
-    jchar GetStaticCharField(jclass clazz, jfieldID fieldID)
-    { return functions->GetStaticCharField(this, clazz, fieldID); }
-    jshort GetStaticShortField(jclass clazz, jfieldID fieldID)
-    { return functions->GetStaticShortField(this, clazz, fieldID); }
-    jint GetStaticIntField(jclass clazz, jfieldID fieldID)
-    { return functions->GetStaticIntField(this, clazz, fieldID); }
-    jlong GetStaticLongField(jclass clazz, jfieldID fieldID)
-    { return functions->GetStaticLongField(this, clazz, fieldID); }
-    jfloat GetStaticFloatField(jclass clazz, jfieldID fieldID)
-    { return functions->GetStaticFloatField(this, clazz, fieldID); }
-    jdouble GetStaticDoubleField(jclass clazz, jfieldID fieldID)
-    { return functions->GetStaticDoubleField(this, clazz, fieldID); }
-
-    void SetStaticObjectField(jclass clazz, jfieldID fieldID, jobject value)
-    { functions->SetStaticObjectField(this, clazz, fieldID, value); }
-    void SetStaticBooleanField(jclass clazz, jfieldID fieldID, jboolean value)
-    { functions->SetStaticBooleanField(this, clazz, fieldID, value); }
-    void SetStaticByteField(jclass clazz, jfieldID fieldID, jbyte value)
-    { functions->SetStaticByteField(this, clazz, fieldID, value); }
-    void SetStaticCharField(jclass clazz, jfieldID fieldID, jchar value)
-    { functions->SetStaticCharField(this, clazz, fieldID, value); }
-    void SetStaticShortField(jclass clazz, jfieldID fieldID, jshort value)
-    { functions->SetStaticShortField(this, clazz, fieldID, value); }
-    void SetStaticIntField(jclass clazz, jfieldID fieldID, jint value)
-    { functions->SetStaticIntField(this, clazz, fieldID, value); }
-    void SetStaticLongField(jclass clazz, jfieldID fieldID, jlong value)
-    { functions->SetStaticLongField(this, clazz, fieldID, value); }
-    void SetStaticFloatField(jclass clazz, jfieldID fieldID, jfloat value)
-    { functions->SetStaticFloatField(this, clazz, fieldID, value); }
-    void SetStaticDoubleField(jclass clazz, jfieldID fieldID, jdouble value)
-    { functions->SetStaticDoubleField(this, clazz, fieldID, value); }
-
-    jstring NewString(const jchar* unicodeChars, jsize len)
-    { return functions->NewString(this, unicodeChars, len); }
-
-    jsize GetStringLength(jstring string)
-    { return functions->GetStringLength(this, string); }
-
-    const jchar* GetStringChars(jstring string, jboolean* isCopy)
-    { return functions->GetStringChars(this, string, isCopy); }
-
-    void ReleaseStringChars(jstring string, const jchar* chars)
-    { functions->ReleaseStringChars(this, string, chars); }
-
-    jstring NewStringUTF(const char* bytes)
-    { return functions->NewStringUTF(this, bytes); }
-
-    jsize GetStringUTFLength(jstring string)
-    { return functions->GetStringUTFLength(this, string); }
-
-    const char* GetStringUTFChars(jstring string, jboolean* isCopy)
-    { return functions->GetStringUTFChars(this, string, isCopy); }
-
-    void ReleaseStringUTFChars(jstring string, const char* utf)
-    { functions->ReleaseStringUTFChars(this, string, utf); }
-
-    jsize GetArrayLength(jarray array)
-    { return functions->GetArrayLength(this, array); }
-
-    jobjectArray NewObjectArray(jsize length, jclass elementClass,
-        jobject initialElement)
-    { return functions->NewObjectArray(this, length, elementClass,
-        initialElement); }
-
-    jobject GetObjectArrayElement(jobjectArray array, jsize index)
-    { return functions->GetObjectArrayElement(this, array, index); }
-
-    void SetObjectArrayElement(jobjectArray array, jsize index, jobject value)
-    { functions->SetObjectArrayElement(this, array, index, value); }
-
-    jbooleanArray NewBooleanArray(jsize length)
-    { return functions->NewBooleanArray(this, length); }
-    jbyteArray NewByteArray(jsize length)
-    { return functions->NewByteArray(this, length); }
-    jcharArray NewCharArray(jsize length)
-    { return functions->NewCharArray(this, length); }
-    jshortArray NewShortArray(jsize length)
-    { return functions->NewShortArray(this, length); }
-    jintArray NewIntArray(jsize length)
-    { return functions->NewIntArray(this, length); }
-    jlongArray NewLongArray(jsize length)
-    { return functions->NewLongArray(this, length); }
-    jfloatArray NewFloatArray(jsize length)
-    { return functions->NewFloatArray(this, length); }
-    jdoubleArray NewDoubleArray(jsize length)
-    { return functions->NewDoubleArray(this, length); }
-
-    jboolean* GetBooleanArrayElements(jbooleanArray array, jboolean* isCopy)
-    { return functions->GetBooleanArrayElements(this, array, isCopy); }
-    jbyte* GetByteArrayElements(jbyteArray array, jboolean* isCopy)
-    { return functions->GetByteArrayElements(this, array, isCopy); }
-    jchar* GetCharArrayElements(jcharArray array, jboolean* isCopy)
-    { return functions->GetCharArrayElements(this, array, isCopy); }
-    jshort* GetShortArrayElements(jshortArray array, jboolean* isCopy)
-    { return functions->GetShortArrayElements(this, array, isCopy); }
-    jint* GetIntArrayElements(jintArray array, jboolean* isCopy)
-    { return functions->GetIntArrayElements(this, array, isCopy); }
-    jlong* GetLongArrayElements(jlongArray array, jboolean* isCopy)
-    { return functions->GetLongArrayElements(this, array, isCopy); }
-    jfloat* GetFloatArrayElements(jfloatArray array, jboolean* isCopy)
-    { return functions->GetFloatArrayElements(this, array, isCopy); }
-    jdouble* GetDoubleArrayElements(jdoubleArray array, jboolean* isCopy)
-    { return functions->GetDoubleArrayElements(this, array, isCopy); }
-
-    void ReleaseBooleanArrayElements(jbooleanArray array, jboolean* elems,
-        jint mode)
-    { functions->ReleaseBooleanArrayElements(this, array, elems, mode); }
-    void ReleaseByteArrayElements(jbyteArray array, jbyte* elems,
-        jint mode)
-    { functions->ReleaseByteArrayElements(this, array, elems, mode); }
-    void ReleaseCharArrayElements(jcharArray array, jchar* elems,
-        jint mode)
-    { functions->ReleaseCharArrayElements(this, array, elems, mode); }
-    void ReleaseShortArrayElements(jshortArray array, jshort* elems,
-        jint mode)
-    { functions->ReleaseShortArrayElements(this, array, elems, mode); }
-    void ReleaseIntArrayElements(jintArray array, jint* elems,
-        jint mode)
-    { functions->ReleaseIntArrayElements(this, array, elems, mode); }
-    void ReleaseLongArrayElements(jlongArray array, jlong* elems,
-        jint mode)
-    { functions->ReleaseLongArrayElements(this, array, elems, mode); }
-    void ReleaseFloatArrayElements(jfloatArray array, jfloat* elems,
-        jint mode)
-    { functions->ReleaseFloatArrayElements(this, array, elems, mode); }
-    void ReleaseDoubleArrayElements(jdoubleArray array, jdouble* elems,
-        jint mode)
-    { functions->ReleaseDoubleArrayElements(this, array, elems, mode); }
-
-    void GetBooleanArrayRegion(jbooleanArray array, jsize start, jsize len,
-        jboolean* buf)
-    { functions->GetBooleanArrayRegion(this, array, start, len, buf); }
-    void GetByteArrayRegion(jbyteArray array, jsize start, jsize len,
-        jbyte* buf)
-    { functions->GetByteArrayRegion(this, array, start, len, buf); }
-    void GetCharArrayRegion(jcharArray array, jsize start, jsize len,
-        jchar* buf)
-    { functions->GetCharArrayRegion(this, array, start, len, buf); }
-    void GetShortArrayRegion(jshortArray array, jsize start, jsize len,
-        jshort* buf)
-    { functions->GetShortArrayRegion(this, array, start, len, buf); }
-    void GetIntArrayRegion(jintArray array, jsize start, jsize len,
-        jint* buf)
-    { functions->GetIntArrayRegion(this, array, start, len, buf); }
-    void GetLongArrayRegion(jlongArray array, jsize start, jsize len,
-        jlong* buf)
-    { functions->GetLongArrayRegion(this, array, start, len, buf); }
-    void GetFloatArrayRegion(jfloatArray array, jsize start, jsize len,
-        jfloat* buf)
-    { functions->GetFloatArrayRegion(this, array, start, len, buf); }
-    void GetDoubleArrayRegion(jdoubleArray array, jsize start, jsize len,
-        jdouble* buf)
-    { functions->GetDoubleArrayRegion(this, array, start, len, buf); }
-
-    void SetBooleanArrayRegion(jbooleanArray array, jsize start, jsize len,
-        const jboolean* buf)
-    { functions->SetBooleanArrayRegion(this, array, start, len, buf); }
-    void SetByteArrayRegion(jbyteArray array, jsize start, jsize len,
-        const jbyte* buf)
-    { functions->SetByteArrayRegion(this, array, start, len, buf); }
-    void SetCharArrayRegion(jcharArray array, jsize start, jsize len,
-        const jchar* buf)
-    { functions->SetCharArrayRegion(this, array, start, len, buf); }
-    void SetShortArrayRegion(jshortArray array, jsize start, jsize len,
-        const jshort* buf)
-    { functions->SetShortArrayRegion(this, array, start, len, buf); }
-    void SetIntArrayRegion(jintArray array, jsize start, jsize len,
-        const jint* buf)
-    { functions->SetIntArrayRegion(this, array, start, len, buf); }
-    void SetLongArrayRegion(jlongArray array, jsize start, jsize len,
-        const jlong* buf)
-    { functions->SetLongArrayRegion(this, array, start, len, buf); }
-    void SetFloatArrayRegion(jfloatArray array, jsize start, jsize len,
-        const jfloat* buf)
-    { functions->SetFloatArrayRegion(this, array, start, len, buf); }
-    void SetDoubleArrayRegion(jdoubleArray array, jsize start, jsize len,
-        const jdouble* buf)
-    { functions->SetDoubleArrayRegion(this, array, start, len, buf); }
-
-    jint RegisterNatives(jclass clazz, const JNINativeMethod* methods,
-        jint nMethods)
-    { return functions->RegisterNatives(this, clazz, methods, nMethods); }
-
-    jint UnregisterNatives(jclass clazz)
-    { return functions->UnregisterNatives(this, clazz); }
-
-    jint MonitorEnter(jobject obj)
-    { return functions->MonitorEnter(this, obj); }
-
-    jint MonitorExit(jobject obj)
-    { return functions->MonitorExit(this, obj); }
-
-    jint GetJavaVM(JavaVM** vm)
-    { return functions->GetJavaVM(this, vm); }
-
-    void GetStringRegion(jstring str, jsize start, jsize len, jchar* buf)
-    { functions->GetStringRegion(this, str, start, len, buf); }
-
-    void GetStringUTFRegion(jstring str, jsize start, jsize len, char* buf)
-    { return functions->GetStringUTFRegion(this, str, start, len, buf); }
-
-    void* GetPrimitiveArrayCritical(jarray array, jboolean* isCopy)
-    { return functions->GetPrimitiveArrayCritical(this, array, isCopy); }
-
-    void ReleasePrimitiveArrayCritical(jarray array, void* carray, jint mode)
-    { functions->ReleasePrimitiveArrayCritical(this, array, carray, mode); }
-
-    const jchar* GetStringCritical(jstring string, jboolean* isCopy)
-    { return functions->GetStringCritical(this, string, isCopy); }
-
-    void ReleaseStringCritical(jstring string, const jchar* carray)
-    { functions->ReleaseStringCritical(this, string, carray); }
-
-    jweak NewWeakGlobalRef(jobject obj)
-    { return functions->NewWeakGlobalRef(this, obj); }
-
-    void DeleteWeakGlobalRef(jweak obj)
-    { functions->DeleteWeakGlobalRef(this, obj); }
-
-    jboolean ExceptionCheck()
-    { return functions->ExceptionCheck(this); }
-
-    jobject NewDirectByteBuffer(void* address, jlong capacity)
-    { return functions->NewDirectByteBuffer(this, address, capacity); }
-
-    void* GetDirectBufferAddress(jobject buf)
-    { return functions->GetDirectBufferAddress(this, buf); }
-
-    jlong GetDirectBufferCapacity(jobject buf)
-    { return functions->GetDirectBufferCapacity(this, buf); }
-
-    /* added in JNI 1.6 */
-    jobjectRefType GetObjectRefType(jobject obj)
-    { return functions->GetObjectRefType(this, obj); }
-#endif /*__cplusplus*/
-};
-
-
-/*
- * JNI invocation interface.
- */
-struct JNIInvokeInterface {
-    void*       reserved0;
-    void*       reserved1;
-    void*       reserved2;
-    jint        (*DestroyJavaVM)(JavaVM*);
-    jint        (*AttachCurrentThread)(JavaVM*, JNIEnv**, void*);
-    jint        (*DetachCurrentThread)(JavaVM*);
-    jint        (*GetEnv)(JavaVM*, void**, jint);
-    jint        (*AttachCurrentThreadAsDaemon)(JavaVM*, JNIEnv**, void*);
-};
-
-/*
- * C++ version.
- */
-struct _JavaVM {
-    const struct JNIInvokeInterface* functions;
-
-#if defined(__cplusplus)
-    jint DestroyJavaVM()
-    { return functions->DestroyJavaVM(this); }
-    jint AttachCurrentThread(JNIEnv** p_env, void* thr_args)
-    { return functions->AttachCurrentThread(this, p_env, thr_args); }
-    jint DetachCurrentThread()
-    { return functions->DetachCurrentThread(this); }
-    jint GetEnv(void** env, jint version)
-    { return functions->GetEnv(this, env, version); }
-    jint AttachCurrentThreadAsDaemon(JNIEnv** p_env, void* thr_args)
-    { return functions->AttachCurrentThreadAsDaemon(this, p_env, thr_args); }
-#endif /*__cplusplus*/
-};
-
-struct JavaVMAttachArgs {
-    jint        version;    /* must be >= JNI_VERSION_1_2 */
-    const char* name;       /* NULL or name of thread as modified UTF-8 str */
-    jobject     group;      /* global ref of a ThreadGroup object, or NULL */
-};
-typedef struct JavaVMAttachArgs JavaVMAttachArgs;
-
-/*
- * JNI 1.2+ initialization.  (As of 1.6, the pre-1.2 structures are no
- * longer supported.)
- */
-typedef struct JavaVMOption {
-    const char* optionString;
-    void*       extraInfo;
-} JavaVMOption;
-
-typedef struct JavaVMInitArgs {
-    jint        version;    /* use JNI_VERSION_1_2 or later */
-    jint        nOptions;
-    JavaVMOption* options;
-    jboolean    ignoreUnrecognized;
-} JavaVMInitArgs;
-
-#ifdef __cplusplus
-extern "C" {
-#endif
-/*
- * VM initialization functions.
- *
- * Note these are the only symbols exported for JNI by the VM.
- */
-jint JNI_GetDefaultJavaVMInitArgs(void*);
-jint JNI_CreateJavaVM(JavaVM**, JNIEnv**, void*);
-jint JNI_GetCreatedJavaVMs(JavaVM**, jsize, jsize*);
-
-#define JNIIMPORT
-// To match the JNIEXPORT on Windows
-#define JNIEXPORT __declspec(dllexport)
-#define JNICALL
-
-/*
- * Prototypes for functions exported by loadable shared libs.  These are
- * called by JNI, not provided by JNI.
- */
-JNIEXPORT jint JNI_OnLoad(JavaVM* vm, void* reserved);
-JNIEXPORT void JNI_OnUnload(JavaVM* vm, void* reserved);
-
-#ifdef __cplusplus
-}
-#endif
-
-/*
- * Manifest constants.
- */
-#define JNI_FALSE   0
-#define JNI_TRUE    1
-
-#define JNI_VERSION_1_1 0x00010001
-#define JNI_VERSION_1_2 0x00010002
-#define JNI_VERSION_1_4 0x00010004
-#define JNI_VERSION_1_6 0x00010006
-
-#define JNI_OK          (0)         /* no error */
-#define JNI_ERR         (-1)        /* generic error */
-#define JNI_EDETACHED   (-2)        /* thread detached from the VM */
-#define JNI_EVERSION    (-3)        /* JNI version error */
-
-#define JNI_COMMIT      1           /* copy content, do not free buffer */
-#define JNI_ABORT       2           /* free buffer w/o copying back */
-
-#endif  /* JNI_H_ */
diff --git a/.claude/skills/aoti-debug/SKILL.md b/.claude/skills/aoti-debug/SKILL.md
index 56a17e08393e9..a790d1b15c26e 100644
--- a/.claude/skills/aoti-debug/SKILL.md
+++ b/.claude/skills/aoti-debug/SKILL.md
@@ -7,6 +7,22 @@ description: Debug AOTInductor (AOTI) errors and crashes. Use when encountering
 
 This skill helps diagnose and fix common AOTInductor issues.
 
+## Error Pattern Routing
+
+**Check the error message and route to the appropriate sub-guide:**
+
+### Triton Index Out of Bounds
+If the error matches this pattern:
+```
+Assertion `index out of bounds: 0 <= tmpN < ksM` failed
+```
+**→ Follow the guide in `triton-index-out-of-bounds.md`**
+
+### All Other Errors
+Continue with the sections below.
+
+---
+
 ## First Step: Always Check Device and Shape Matching
 
 **For ANY AOTI error (segfault, exception, crash, wrong output), ALWAYS check these first:**
diff --git a/.claude/skills/aoti-debug/triton-index-out-of-bounds.md b/.claude/skills/aoti-debug/triton-index-out-of-bounds.md
new file mode 100644
index 0000000000000..fee827e32feb2
--- /dev/null
+++ b/.claude/skills/aoti-debug/triton-index-out-of-bounds.md
@@ -0,0 +1,265 @@
+# AOTI Triton Index Out of Bounds Debug Guide
+
+This guide helps debug AOTI Triton kernel assertion errors with the `index out of bounds` pattern.
+
+## Error Pattern
+
+This guide applies when you see errors like:
+
+```
+/var/tmp/torchinductor_*/.../*.py:NN: unknown: block: [X,Y,Z], thread: [X,Y,Z]
+Assertion `index out of bounds: 0 <= tmpN < ksM` failed.
+```
+
+### Key Information from Error
+
+| Field | Value | Meaning |
+|-------|-------|---------|
+| File Path | `/var/tmp/torchinductor_*/*.py` | Generated Triton kernel file (runtime) |
+| Line Number | `:NN` | Line in the generated kernel where assertion failed |
+| Block/Thread | `[X,Y,Z]` | CUDA block and thread indices |
+| Assertion | `0 <= tmpN < ksM` | Index `tmpN` must be within bounds `[0, ksM)` |
+
+### Understanding the Assertion
+
+- `tmpN`: A computed index value in the Triton kernel
+- `ksM`: A dynamic kernel size parameter (runtime value)
+- The assertion fails when `tmpN < 0` or `tmpN >= ksM`
+
+---
+
+## Step 1: Collect AOTI Package
+
+You need access to the AOTI package that was compiled. This is typically a `.pt2` package or extracted archive containing a `wrapper.cpp` file.
+
+**Key File**: `*.wrapper.cpp` contains:
+- All Triton kernel source code (embedded as comments)
+- Kernel launch configurations
+- Input/output tensor mappings
+- Dynamic shape variable definitions
+
+---
+
+## Step 2: Locate the Failing Kernel in C++ Wrapper
+
+### Search for the Assertion Pattern
+
+Extract the assertion pattern from the error (e.g., `tmp18 < ks0`) and search:
+
+```bash
+# Search for the specific assertion
+grep -n "tmpN < ksM" /path/to/*.wrapper.cpp
+
+# Get context around the assertion (80 lines before, 20 after)
+grep -n -B80 -A20 "tmpN < ksM" /path/to/*.wrapper.cpp
+```
+
+### Find the Full Kernel Definition
+
+The kernel is embedded as a Python docstring comment in the C++ wrapper:
+
+```cpp
+    /*
+    async_compile.triton('triton_red_fused_...', '''
+    import triton
+    import triton.language as tl
+    ...
+    def triton_red_fused_...(in_ptr0, out_ptr1, ks0, xnumel, r0_numel, ...):
+```
+
+---
+
+## Step 3: Understand the Kernel Logic
+
+Analyze the code path leading to the assertion. Common patterns that cause index out of bounds:
+
+### Pattern: Empty Tensor with ks0 = 0
+
+When a dynamic shape `ks0 = 0`:
+1. `tmp13 = (-1) + 0 = -1`
+2. Index wrapping logic produces `-1`
+3. Assertion `0 <= -1 < 0` fails
+
+### Example Kernel Pattern
+
+```python
+tmp13 = (-1) + ks0           # ks0 - 1
+tmp14 = tl.where(tmp12, tmp10, tmp13)  # if condition: use tmp10, else: ks0-1
+tmp15 = ks0
+tmp16 = tmp14 + tmp15        # wrap-around for negative indices
+tmp17 = tmp14 < 0
+tmp18 = tl.where(tmp17, tmp16, tmp14)  # if negative: add ks0
+
+# ASSERTION: 0 <= tmp18 < ks0
+tl.device_assert(((0 <= tmp18) & (tmp18 < ks0)), "index out of bounds")
+```
+
+---
+
+## Step 4: Identify the Dynamic Shape Variable
+
+### Find Where the Kernel is Called
+
+```bash
+grep -n "call_triton_KERNEL_NAME" /path/to/*.wrapper.cpp
+```
+
+### Example Output
+
+```cpp
+call_triton_red_fused_...(arg1415_1, buf696, s607, 1L, s13, ...);
+```
+
+### Parameter Mapping
+
+| Parameter | Value | Meaning |
+|-----------|-------|---------|
+| `in_ptr0` | `arg1415_1` | Input tensor |
+| `out_ptr1` | `buf696` | Output buffer |
+| `ks0` | `s607` | **Dynamic shape - this is the failing bound** |
+
+### Find the Definition of the Shape Variable
+
+```bash
+grep -n "int64_t s607 = " /path/to/*.wrapper.cpp
+```
+
+This shows which input tensor dimension defines the shape:
+
+```cpp
+int64_t s607 = arg1416_1_size[0];
+```
+
+---
+
+## Step 5: Trace Back to Model Input
+
+### Find Input Index
+
+Inputs are numbered sequentially. Find which input the argument corresponds to:
+
+```bash
+grep -n 'inputs_info_\[INDEX\].name = "argNNN_1"' /path/to/*.wrapper.cpp
+```
+
+### Check Input Constraints
+
+```bash
+grep -n "argNNN_1_size\[0\]" /path/to/*.wrapper.cpp
+```
+
+Look for guards like:
+```cpp
+if (arg_size[0] > 230400) {  // Upper bound check only - no lower bound!
+```
+
+**Common Issue**: Upper bound checks exist but no lower bound checks for `>= 1`.
+
+---
+
+## Step 6: Map to Model Code
+
+### Use Source Node Comments
+
+The C++ wrapper includes comments showing which PyTorch operations generated each kernel:
+
+```bash
+grep -n -B5 "call_triton_KERNEL_NAME" /path/to/*.wrapper.cpp | grep "Source Nodes"
+```
+
+### Example Output
+
+```cpp
+// Topologically Sorted Source Nodes: [slice_1, sub_89, cumsum, ge_231, where_2, index_copy]
+```
+
+### Map Operations to Python Code
+
+| ATen Operation | Python Code Pattern |
+|----------------|---------------------|
+| `cumsum` | `torch.cumsum(tensor, dim=0)` |
+| `sub` | `idx - 1` |
+| `ge` | `idx >= 0` |
+| `where` | `torch.where(condition, ...)` |
+| `index_copy` | `tensor.index_copy(0, indices, source)` |
+
+---
+
+## Root Cause Analysis
+
+### Common Root Causes
+
+1. **Empty tensor at runtime**: A jagged/variable-length tensor has size 0 at runtime but wasn't tested during compilation
+2. **Missing lower bound guards**: AOTI only generates upper bound checks, not lower bound checks
+3. **Edge case not in sample inputs**: Sample inputs during AOTI export never included the edge case
+
+---
+
+## Fix Recommendations
+
+### Option 1: Add Guard in Forward Method
+
+```python
+def forward(self, lengths: torch.Tensor, ...) -> torch.Tensor:
+    if lengths.numel() == 0:
+        device = lengths.device
+        return torch.empty(0, self.output_dim, device=device)
+    # ... rest of method
+```
+
+### Option 2: Fix the Specific Operation
+
+Add handling for empty tensors in the problematic operation:
+
+```python
+def process_events(self, lengths: torch.Tensor, ...):
+    if lengths.numel() == 0:
+        return torch.empty(0, self.emb_dim, device=lengths.device)
+    # ... rest of method
+```
+
+### Option 3: Include Edge Cases in AOTI Export
+
+Ensure sample inputs during AOTI export include:
+- Empty tensors (size 0)
+- Minimum size tensors (size 1)
+- Maximum expected sizes
+
+---
+
+## Useful Commands Summary
+
+### Searching in AOTI Wrapper
+
+```bash
+# Find kernel by assertion pattern
+grep -n "tmpN < ksM" *.wrapper.cpp
+
+# Get full kernel context
+grep -n -B80 -A20 "ASSERTION_PATTERN" *.wrapper.cpp
+
+# Find kernel call site
+grep -n "call_KERNEL_NAME" *.wrapper.cpp
+
+# Find dynamic shape definition
+grep -n "int64_t SHAPE_VAR = " *.wrapper.cpp
+
+# Find input mapping
+grep -n 'inputs_info_\[INDEX\].name' *.wrapper.cpp
+
+# Find size constraints
+grep -n "SHAPE_VAR_size\[0\]" *.wrapper.cpp
+```
+
+### Environment Variables for Debugging
+
+```bash
+# Enable debug output during torch.compile
+export TORCH_COMPILE_DEBUG=1
+
+# Save generated kernels to persistent location
+export TORCHINDUCTOR_CACHE_DIR=/path/to/save/kernels
+
+# Enable CUDA launch blocking for accurate stack traces
+export CUDA_LAUNCH_BLOCKING=1
+```
diff --git a/.claude/skills/document-public-apis/SKILL.md b/.claude/skills/document-public-apis/SKILL.md
new file mode 100644
index 0000000000000..211b444fc5461
--- /dev/null
+++ b/.claude/skills/document-public-apis/SKILL.md
@@ -0,0 +1,342 @@
+---
+name: document-public-apis
+description: Document undocumented public APIs in PyTorch by removing functions from coverage_ignore_functions and coverage_ignore_classes in docs/source/conf.py, running Sphinx coverage, and adding the appropriate autodoc directives to the correct .md or .rst doc files. Use when a user asks to remove functions from conf.py ignore lists.
+---
+
+# Document Public APIs
+
+This skill documents undocumented public APIs in PyTorch by removing entries from the coverage ignore lists in `docs/source/conf.py` and adding Sphinx autodoc directives (e.g., `autosummary`, `currentmodule`, `autoclass`, `automodule`) to the corresponding `.md` or `.rst` doc source files in `docs/source/`.
+
+**"Documenting" means adding autodoc directives to doc source files — NEVER modifying Python source code.** Do not add or edit docstrings in `.py` files. Your only job is to add the correct directive to the correct doc file.
+
+**IMPORTANT: Before adding any function to the sphinx doctree, verify it has a real docstring.** Use a quick Python check (e.g., `python -c "from torch.module import func; print(bool(func.__doc__))"`) to confirm the function has actual documentation content — not just an empty docstring or a bare `.. warning:: This API is experimental` stub. Functions without meaningful docstrings should be left in the `coverage_ignore_functions`/`coverage_ignore_classes` lists. Adding undocumented functions to the doctree creates empty or near-empty pages that degrade documentation quality.
+
+## Overview
+
+`docs/source/conf.py` contains two lists that suppress Sphinx coverage warnings for undocumented APIs:
+
+- `coverage_ignore_functions`: undocumented functions
+- `coverage_ignore_classes`: undocumented classes
+
+Entries are organized by **module comment groups**. Each group has a module label comment followed by the function/class names that belong to that module:
+
+```python
+coverage_ignore_functions = [
+    # torch.ao.quantization.fx.convert              <-- module label comment
+    "convert",                                       # <-- entries belonging to this module
+    "convert_custom_module",
+    "convert_standalone_module",
+    "convert_weighted_module",
+    # torch.ao.quantization.fx.fuse                 <-- next module group
+    "fuse",
+    # torch.nn.functional
+    "assert_int_or_pair",  # looks unintentionally public   <-- entry with inline comment
+    "constant",  # deprecated                                <-- entry with inline comment
+]
+```
+
+There are two kinds of comments:
+- **Module label comments** (`# torch.ao.quantization.fx.convert`): these label which module the entries below belong to. They appear on their own line before a group of entries.
+- **Inline comments** (`# deprecated`, `# documented as adaptive_max_pool1d`): these appear after a string entry on the same line and explain *why* the entry is in the ignore list.
+
+The module label comment directly tells you:
+1. Which module the functions belong to
+2. Where to add them in the docs (e.g., `# torch.ao.quantization.fx.convert` → the functions go under `torch.ao.quantization.fx.convert` in the doc file)
+
+## Instructions
+
+Each invocation of this skill processes **one batch** of module groups. Pick one or more complete module groups from the ignore lists, document their functions, and verify.
+
+### Step 1: Select module groups to document
+
+Read `docs/source/conf.py` and select one or more **complete module groups** to document. A module group is a module label comment and all entries beneath it up to the next module label comment. Process entire groups — never split a group across batches.
+
+For example, selecting the `torch.ao.quantization.fx.convert` group means taking all of:
+
+```python
+# torch.ao.quantization.fx.convert
+"convert",
+"convert_custom_module",
+"convert_standalone_module",
+"convert_weighted_module",
+```
+
+Work through the lists top-to-bottom. Choose enough groups to make meaningful progress (aim for 5–15 functions total, but always include complete groups even if that means going slightly over).
+
+**Check inline comments before including an entry.** Some entries have inline comments that indicate they should not be documented:
+
+- `# deprecated` — The function is deprecated. Leave it in the ignore list.
+- `# documented as <other_name>` — Already documented under a different name. Leave it.
+- `# looks unintentionally public` — Probably not meant to be public API. Leave it.
+- `# legacy helper for ...` — Same as deprecated. Leave it.
+- `# utility function` - Leave it.
+
+If a module group has a **mix** of regular entries and entries with inline comments, still process the group — but only comment out the regular entries. Leave entries with inline comments untouched in the ignore list.
+
+### Step 1b: Verify functions have actual docstrings
+
+For each function selected in Step 1, check that it has a meaningful docstring by running:
+
+```bash
+python -c "from torch.module.path import func_name; doc = func_name.__doc__; print('HAS DOC' if doc and len(doc.strip()) > 80 else 'NO DOC'); print(repr(doc[:120]) if doc else 'None')"
+```
+
+A function has a **meaningful docstring** if it has real descriptive content — not just:
+- `None` or empty string
+- Only a `.. warning:: This API is experimental` stub with no description
+- Only a one-line auto-generated signature
+
+**Functions without meaningful docstrings must stay in the ignore list.** Remove them from your batch. If an entire module group has no functions with docstrings, skip the whole group.
+
+### Step 2: Present the batch to the user
+
+**Before making any edits**, present the selected module groups and their functions to the user. Indicate which functions passed the docstring check and which were excluded. Show them organized by module:
+
+```
+Module: torch.ao.quantization.fx.convert
+  - convert
+  - convert_custom_module
+  - convert_standalone_module
+  - convert_weighted_module
+
+Module: torch.ao.quantization.fx.fuse
+  - fuse
+```
+
+Then use the `AskUserQuestion` tool to let the user confirm, with options like:
+- "Proceed with this batch"
+- "Skip some entries" (user can specify which to remove)
+- "Pick a different batch"
+
+### Step 3: Comment out entries in conf.py
+
+After the user confirms, edit `docs/source/conf.py` and **comment out** (do not delete) the selected entries. Use a `#` prefix on each string entry line:
+
+```python
+# torch.ao.quantization.fx.convert
+# "convert",
+# "convert_custom_module",
+# "convert_standalone_module",
+# "convert_weighted_module",
+```
+
+This preserves the original entries so they can be restored if verification fails.
+
+### Step 4: Run Sphinx coverage
+
+```bash
+cd docs && make coverage
+```
+
+**Ignore the terminal output of `make coverage`.** It often contains unrelated tracebacks and errors from Sphinx extensions (e.g., `onnx_ir`, `katex`, `sphinxcontrib`) that have nothing to do with coverage. The only thing that matters is whether `docs/build/coverage/python.txt` was generated. Read that file to see the specific undocumented APIs.
+
+The format of `python.txt` lists each undocumented API as:
+
+```
+torch.ao.quantization.fx.convert
+   * convert
+   * convert_custom_module
+   * convert_standalone_module
+   * convert_weighted_module
+```
+
+**Not all commented-out functions will appear in `python.txt`.** Some may already be documented elsewhere. This is fine — only add directives for functions that actually appear in `python.txt`.
+
+If `make coverage` fails due to missing dependencies, first run:
+
+```bash
+cd docs && pip install -r requirements.txt
+```
+
+### Step 5: Add documentation directives
+
+For each function listed in `python.txt`, use the **module label comment** from `conf.py` to determine where it should be added. The module comment gives you the full module path, which maps to a doc source file and a section within that file.
+
+#### Finding the correct doc file
+
+The module comment maps to a doc source file in `docs/source/`. When unsure, search for other functions from the same module:
+
+```bash
+grep -rn "torch.module_name" docs/source/*.md docs/source/*.rst
+```
+
+Or list candidate files:
+
+```bash
+ls docs/source/*module_name*
+```
+
+If no doc file exists for a submodule, check whether a parent module's doc file has a section for it (e.g., `backends.md` has sections for `torch.backends.cuda`, `torch.backends.cudnn`, etc.). If not, add a new section to the parent file following existing patterns.
+
+#### Adding the directives
+
+**Read the target doc file first** and match the exact patterns already used there. Do not invent new patterns or use bare `autofunction` with fully qualified names — always use the proper hierarchical structure with `automodule`, `currentmodule`, and short names. Do not use `. py:module::` since that just suppresses errors and doesn't actually document the function. Look at other files that match the target file's format (e.g., `.md` vs. `.rst`) under `docs/source/` to see examples.
+
+There are two file formats. Match the one used in the target file.
+
+**Pattern A — MyST Markdown files (`.md`):** Used in files like `accelerator.md`, `backends.md`, `cuda.md`.
+
+The hierarchical structure uses `automodule` to register the module, `currentmodule` to set context, then short names:
+
+```markdown
+## torch.ao.quantization.fx.convert
+
+```{eval-rst}
+.. automodule:: torch.ao.quantization.fx.convert
+```
+
+```{eval-rst}
+.. currentmodule:: torch.ao.quantization.fx.convert
+```
+
+```{eval-rst}
+.. autofunction:: convert
+```
+
+```{eval-rst}
+.. autofunction:: convert_custom_module
+```
+```
+
+For `autosummary` blocks (used in some files instead of individual directives):
+
+```markdown
+```{eval-rst}
+.. autosummary::
+    :toctree: generated
+    :nosignatures:
+
+    existing_function
+    your_new_function
+`` `
+```
+
+For classes:
+
+```markdown
+```{eval-rst}
+.. autoclass:: YourClass
+    :members:
+`` `
+```
+
+**Pattern B — reStructuredText files (`.rst`):** Used in files like `torch.rst`, `nn.rst`.
+
+Same hierarchical structure without the markdown fences:
+
+```rst
+torch.ao.quantization.fx.convert
+---------------------------------
+
+.. automodule:: torch.ao.quantization.fx.convert
+
+.. currentmodule:: torch.ao.quantization.fx.convert
+
+.. autosummary::
+    :toctree: generated
+    :nosignatures:
+
+    convert
+    convert_custom_module
+    convert_standalone_module
+    convert_weighted_module
+```
+
+For individual directives:
+
+```rst
+.. automodule:: torch.submodule
+
+.. currentmodule:: torch.submodule
+
+.. autofunction:: function_name
+
+.. autoclass:: ClassName
+    :members:
+```
+
+**Key rules:**
+- The module label comment from `conf.py` (e.g., `# torch.ao.quantization.fx.convert`) tells you exactly which `automodule` and `currentmodule` to use.
+- Always set `.. automodule::` and `.. currentmodule::` before documenting functions from a module.
+- Use **short names** (e.g., `convert`, not `torch.ao.quantization.fx.convert.convert`) after `currentmodule` is set.
+- If the module already has an `automodule`/`currentmodule` in the file, don't add another — just add your function under the existing one.
+- Match whichever style the file already uses (`autosummary` blocks vs. individual `autofunction` directives).
+
+#### Placing in the right section
+
+Read the target doc file and find the appropriate section. If the module already has a section (e.g., `## torch.backends.cuda` in `backends.md`), add the functions there. If no section exists yet, create one following the existing section patterns in the file. Group all functions from the same module group together.
+
+### Step 6: Verify with coverage
+
+Run coverage again:
+
+```bash
+cd docs && make coverage
+```
+
+Ignore the terminal output — only read `docs/build/coverage/python.txt`. Verification passes when `python.txt` contains **zero undocumented functions across ALL modules**. It should only have the statistics table with 100% coverage and 0 undocumented for every module. For example:
+
+```
+Undocumented Python objects
+===========================
+
+Statistics
+----------
+
++---------------------------+----------+--------------+
+| Module                    | Coverage | Undocumented |
++===========================+==========+==============+
+| torch                     | 100.00%  | 0            |
++---------------------------+----------+--------------+
+| torch.accelerator         | 100.00%  | 0            |
++---------------------------+----------+--------------+
+```
+
+If any module shows undocumented functions (coverage below 100% or undocumented count > 0), verification has failed.
+
+**If verification succeeds (zero undocumented across all modules):** Go to Step 7.
+
+**If verification fails (any undocumented functions remain):** Read `docs/build/coverage/python.txt` to see which functions are still listed as undocumented. Common issues include:
+
+- Wrong doc file: the function was added to the wrong `.md`/`.rst` file. Move the directive to the correct file.
+- Wrong directive type: e.g., used `autofunction` for a class, or `autoclass` for a function. Fix the directive.
+- Wrong module path in the directive: e.g., `torch.foo.bar` should be `torch.foo.baz.bar`. Correct the qualified name.
+- Function added to an `autosummary` block with the wrong `currentmodule`: make sure the `.. currentmodule::` directive above the block matches.
+- Missing `automodule` for a submodule that hasn't been registered yet. Add a `.. automodule:: torch.submodule` directive before documenting functions from that submodule.
+
+Fix the doc directive based on the error, then re-run `make coverage`. Repeat until verification passes.
+
+If a function still fails after multiple attempts, **stop and show the error to the user.** Present the function name and the error, then use the `AskUserQuestion` tool with options like:
+- "Uncomment it to restore to ignore list (skip for now)"
+- "Try a different approach"
+- "Investigate further"
+
+### Step 7: Report progress
+
+**Present a progress summary to the user** showing:
+
+- Which module groups were processed and how many functions were documented
+- Which functions were skipped or restored to the ignore list (and why)
+- How many entries remain in `coverage_ignore_functions` and `coverage_ignore_classes`
+
+### Step 8: Clean up commented-out entries in conf.py
+
+Now that verification has passed, delete the commented-out string entries from Step 3. These are lines that start with `# "` inside `coverage_ignore_functions` and `coverage_ignore_classes`. Commented-out string entries always contain **quotes** — that's how you distinguish them from module label comments:
+
+```python
+# "disable_global_flags",       <-- commented-out string entry (has quotes) → DELETE
+# torch.backends                <-- module label comment (no quotes) → KEEP if it has active entries
+```
+
+Also delete any module label comments that no longer have active entries beneath them (i.e., all their entries were either commented out and now deleted, or had inline comments and were left in place but the module label is otherwise empty).
+
+## Important notes
+
+- **Follow the steps exactly as written.** The `make coverage` step is the primary verification for correct Sphinx directives, and Step 1b's docstring check ensures you only document functions that have real content.
+- **Never modify Python source files (`.py`).** This skill only edits `docs/source/conf.py` and doc source files (`.md`/`.rst`) in `docs/source/`. Do not add or edit docstrings. The only reason to inspect Python modules is in Step 1b to check whether a docstring exists — never to modify source code.
+- Entries are commented out in Step 3, verified in Step 6, and cleaned up in Step 8 after verification passes. Never delete uncommented entries directly.
+- **Read inline comments** on entries before deciding to document them. Entries marked `# deprecated`, `# documented as ...`, `# looks unintentionally public`, or `# legacy helper` should stay in the ignore list.
+- The `coverage_ignore_functions` list uses bare function names (not fully qualified), so the same name can appear multiple times for different modules. Use the module label comment above each entry to identify which module it belongs to. Be careful during Step 8 cleanup to only delete the correct commented-out lines — commented-out string entries have **quotes** (`# "func_name",`), module label comments do not.
+- Always match the existing style of the target doc file — don't mix `.md` style directives into `.rst` files or vice versa.
+- **Use the module label comment** (e.g., `# torch.ao.quantization.fx.convert`) as the primary guide for both the `automodule`/`currentmodule` directives and for finding the right section in the doc file.
+- Always process complete module groups — never split a group across invocations.
diff --git a/.claude/skills/metal-kernel/SKILL.md b/.claude/skills/metal-kernel/SKILL.md
index 75e6684b73a40..18ba4f64570ea 100644
--- a/.claude/skills/metal-kernel/SKILL.md
+++ b/.claude/skills/metal-kernel/SKILL.md
@@ -325,6 +325,83 @@ python test/test_mps.py -k test_output_match_my_op
 python test/test_mps.py
 ```
 
+## Debugging Metal Kernels with `torch.mps.compile_shader`
+
+Use `torch.mps.compile_shader` to JIT-compile and test individual Metal kernels in isolation. This is invaluable for debugging multi-kernel pipelines where you need to verify each stage independently.
+
+### Basic Usage
+
+```python
+import torch
+
+source = '''
+#include <metal_stdlib>
+using namespace metal;
+
+kernel void my_kernel(
+    const device float* input [[buffer(0)]],
+    device float* output [[buffer(1)]],
+    uint tid [[thread_position_in_grid]]) {
+  output[tid] = input[tid] * 2.0;
+}
+'''
+
+lib = torch.mps.compile_shader(source)
+
+inp = torch.tensor([1.0, 2.0, 3.0], device='mps')
+out = torch.zeros(3, device='mps')
+lib.my_kernel(inp, out, threads=[3, 1, 1], group_size=[3, 1, 1])
+torch.mps.synchronize()
+print(out)  # tensor([2., 4., 6.], device='mps:0')
+```
+
+### Dispatch Semantics
+
+`compile_shader` uses **`dispatchThreads`** semantics (same as `mtl_dispatch1DJob` in PyTorch):
+- `threads=[N, 1, 1]` — total number of threads (NOT threadgroups)
+- `group_size=[G, 1, 1]` — threads per threadgroup
+
+This differs from the `dispatchThreadgroups` API used by some host-side code. To match `dispatchThreadgroups:MTLSizeMake(num_tgs, num_slices, 1) threadsPerThreadgroup:MTLSizeMake(TG_SIZE, 1, 1)`:
+
+```python
+# Equivalent compile_shader call:
+lib.kernel(args...,
+    threads=[num_tgs * TG_SIZE, num_slices, 1],
+    group_size=[TG_SIZE, 1, 1])
+```
+
+### Constant Buffer Parameters
+
+Pass scalar constants as single-element tensors:
+
+```python
+slice_size = torch.tensor([1024], dtype=torch.int32, device='mps')
+lib.my_kernel(data, output, slice_size, threads=[1024, 1, 1], group_size=[256, 1, 1])
+```
+
+### Debugging Strategy for Multi-Kernel Pipelines
+
+When a pipeline of kernels (e.g., histogram → prefix_sum → scatter) produces wrong results, test each kernel individually and verify its output against a Python/NumPy reference:
+
+```python
+# 1. Run GPU kernel
+lib.histogram(keys, hist, ..., threads=[N, 1, 1], group_size=[256, 1, 1])
+torch.mps.synchronize()
+
+# 2. Compute reference in Python
+ref_hist = compute_histogram_cpu(keys.cpu().numpy(), ...)
+
+# 3. Compare
+assert np.array_equal(hist.cpu().numpy(), ref_hist), "Histogram mismatch!"
+```
+
+This isolates which kernel in the pipeline is broken, rather than debugging the entire pipeline at once.
+
+### Common Pitfalls
+
+- **Wrong `threads` count** — `threads` is total threads, not threadgroups. For 5 threadgroups of 256, use `threads=[1280, 1, 1]`.
+- **Threadgroup memory** — `compile_shader` doesn't support `[[threadgroup(N)]]` parameters directly. If your kernel needs threadgroup memory, restructure to use `threadgroup` arrays declared inside the kernel body instead.
+
 ## Checklist
 
 - [ ] Added MPS dispatch to `native_functions.yaml`
diff --git a/.claude/skills/pr-review/SKILL.md b/.claude/skills/pr-review/SKILL.md
index 020ad83671bbd..01d68fb6a7610 100644
--- a/.claude/skills/pr-review/SKILL.md
+++ b/.claude/skills/pr-review/SKILL.md
@@ -5,7 +5,7 @@ description: Review PyTorch pull requests for code quality, test coverage, secur
 
 # PyTorch PR Review Skill
 
-Review PyTorch pull requests focusing on what CI cannot check: code quality, test coverage adequacy, security vulnerabilities, and backward compatibility. Linting, formatting, type checking, and import ordering are handled by CI.
+Review PyTorch pull requests focusing on what CI cannot check: code quality, test coverage adequacy, security vulnerabilities, and backward compatibility.
 
 ## Usage Modes
 
@@ -80,60 +80,79 @@ For local branch reviews:
 
 ### GitHub Actions Mode
 
-When invoked via workflow, PR data is passed as context. The PR number or diff will be available in the prompt.
+When invoked via `@claude /pr-review` on a GitHub PR, the action pre-fetches PR
+metadata and injects it into the prompt. Detect this mode by the presence of
+`<formatted_context>`, `<pr_or_issue_body>`, and `<comments>` tags in the prompt.
 
-## Review Workflow
+The prompt already contains:
+- PR metadata (title, author, branch names, additions/deletions, file count)
+- PR body/description
+- All comments and review comments (with file/line references)
+- List of changed files with paths and change types
+
+Use git commands to get the diff and commit history. The base branch name is in the
+prompt context (look for `PR Branch: <head> -> <base>` or the `baseBranch` field).
+
+```bash
+# Get the full diff against the base branch
+git diff origin/<baseBranch>...HEAD
+
+# Get diff stats
+git diff --stat origin/<baseBranch>...HEAD
+
+# Get commit history for this PR
+git log origin/<baseBranch>..HEAD --oneline
+
+# If the base branch ref is not available, fetch it first
+git fetch origin <baseBranch> --depth=1
+```
+
+Do NOT use `gh` CLI commands in this mode -- only git commands are available.
+All PR metadata, comments, and reviews are already in the prompt context;
+only the diff and commit log need to be fetched via git.
+
+## Review Philosophy
+
+A single line of code can have deep cross-cutting implications: a missing device guard causes silent data corruption on multi-GPU, a missing `Composite` dispatch key breaks every out-of-tree backend, a manual dtype check instead of `TensorIterator` silently skips type promotion. **Treat every line as potentially load-bearing.**
 
-### Step 1: Fetch PR Information
+1. **Investigate, don't guess** — When uncertain whether a checklist item applies, spawn a sub-agent to read the relevant code. A reviewer who guesses wrong provides negative value.
+2. **Review the design, not just the implementation** — A PR can have perfectly correct implementation of a bad design. Question side-channel communication, on/off private flags, and demand concrete interface documentation for new contracts between components.
+3. **Focus on what CI cannot check** — Don't comment on formatting, linting, type errors, or CI failures. Focus on design quality, interface correctness, thread safety, BC implications, test adequacy, and pattern adherence.
+4. **Everything is a must-fix** — There are no "nits." If it's worth mentioning, it's worth fixing. Every inconsistency degrades the codebase over time.
+5. **Be specific and actionable** — Reference file paths and line numbers. Name the function/class/file the author should use.
+6. **Match the immediate context** — Read how similar features are already implemented in the same file. Pattern mismatches within a file are always wrong.
+7. **Assume competence** — The author knows PyTorch; explain only non-obvious context.
+8. **No repetition** — Each observation appears in exactly one section of the review output.
 
-For local mode, use `gh` commands to get:
-1. PR metadata (title, description, author)
-2. List of changed files
-3. Full diff of changes
-4. Existing comments/reviews
-5. Fetch associated issue information when applicable
+### Using sub-agents
 
-### Step 2: Analyze Changes
+The review checklist is large. You cannot hold the full context of every infrastructure system in your head. **Spawn sub-agents** to investigate whether checklist items apply: read surrounding code, infrastructure the PR should be using, or tests that should exist. Spawn them in parallel for independent areas. A typical medium PR should spawn 3-8 sub-agents.
 
-Read through the diff systematically:
+## Review Workflow
+
+### Step 1: Understand Context
+
+Before reviewing, build understanding of what the PR touches and why:
 1. Identify the purpose of the change from title/description/issue
 2. Group changes by type (new code, tests, config, docs)
 3. Note the scope of changes (files affected, lines changed)
+4. Spawn sub-agents to read the unchanged code surrounding each significantly changed file to understand existing patterns and infrastructure
 
-### Step 3: Deep Review
+### Step 2: Deep Review
 
-Perform thorough line-by-line analysis using the review checklist. See [review-checklist.md](review-checklist.md) for detailed criteria covering:
-- Code quality and design
-- Testing adequacy
-- Security considerations
-- Performance implications
-- Any behavior change not expected by author
+Go through **every changed line** in the diff and evaluate it against the review checklist in [review-checklist.md](review-checklist.md).
 
-### Step 4: Check Backward Compatibility
+### Step 3: Check Backward Compatibility
 
-Evaluate BC implications. See [bc-guidelines.md](bc-guidelines.md) for:
-- What constitutes a BC-breaking change
-- Required deprecation patterns
-- Common BC pitfalls
+Evaluate BC implications per [bc-guidelines.md](bc-guidelines.md). For non-trivial BC questions, spawn a sub-agent to search for existing callers of the modified API.
 
-### Step 5: Formulate Review
+### Step 4: Formulate Review
 
-Structure your review with actionable feedback organized by category.
-
-## Review Areas
-
-| Area | Focus | Reference |
-|------|-------|-----------|
-| Code Quality | Abstractions, patterns, complexity | [review-checklist.md](review-checklist.md) |
-| API Design | New patterns, flag-based access, broader implications | [review-checklist.md](review-checklist.md) |
-| Testing | Coverage, patterns, edge cases | [review-checklist.md](review-checklist.md) |
-| Security | Injection, credentials, input handling | [review-checklist.md](review-checklist.md) |
-| Performance | Regressions, device handling, memory | [review-checklist.md](review-checklist.md) |
-| BC | Breaking changes, deprecation | [bc-guidelines.md](bc-guidelines.md) |
+Structure your review with actionable feedback organized by category. Every finding should be traceable to a specific line in the diff.
 
 ## Output Format
 
-Structure your review as follows:
+Structure your review as follows. **Omit sections where you have no findings** — don't write "No concerns" for every empty section. Only include sections with actual observations.
 
 ```markdown
 ## PR Review: #<number>
@@ -144,25 +163,29 @@ Structure your review as follows:
 Brief overall assessment of the changes (1-2 sentences).
 
 ### Code Quality
-[Issues and suggestions, or "No concerns" if none]
+[Issues and suggestions]
 
-### API Design
-[Flag new patterns, internal-access flags, or broader implications if any. Otherwise omit this section.]
+### Infrastructure
+[Flag any checklist items from the PyTorch Infrastructure section that apply.
+Reference the specific infrastructure the PR should be using.]
 
 ### Testing
-- [ ] Tests exist for new functionality
-- [ ] Edge cases covered
-- [ ] Tests follow PyTorch patterns (TestCase, assertEqual)
-[Additional testing feedback]
+[Testing adequacy findings — missing OpInfo usage, non-device-generic tests, etc.]
+
+### API Design
+[Flag new patterns, internal-access flags, or broader implications if any.]
 
 ### Security
-[Issues if any, or "No security concerns identified"]
+[Issues if any]
+
+### Thread Safety
+[Threading concerns if any]
 
 ### Backward Compatibility
-[BC concerns if any, or "No BC-breaking changes"]
+[BC concerns if any]
 
 ### Performance
-[Performance concerns if any, or "No performance concerns"]
+[Performance concerns if any]
 
 ### Recommendation
 **Approve** / **Request Changes** / **Needs Discussion**
@@ -185,19 +208,13 @@ When requested, add file-specific feedback with line references:
 - `torch/nn/modules/linear.py:78` - This allocation could be moved outside the loop
 ```
 
-## Key Principles
-
-1. **No repetition** - Each observation appears in exactly one section. Never repeat the same issue, concern, or suggestion across multiple sections. If an issue spans categories (e.g., a security issue that also affects performance), place it in the most relevant section only.
-2. **Focus on what CI cannot check** - Don't comment on formatting, linting, or type errors
-3. **Be specific** - Reference file paths and line numbers
-4. **Be actionable** - Provide concrete suggestions, not vague concerns
-5. **Be proportionate** - Minor issues shouldn't block, but note them
-6. **Assume competence** - The author knows PyTorch; explain only non-obvious context
-
 ## Files to Reference
 
-When reviewing, consult these project files for context:
+When reviewing, consult these project files for context — read them rather than relying on memory, as they change frequently:
 - `CLAUDE.md` - Coding style philosophy and testing patterns
 - `CONTRIBUTING.md` - PR requirements and review process
 - `torch/testing/_internal/common_utils.py` - Test patterns and utilities
 - `torch/testing/_internal/opinfo/core.py` - OpInfo test framework
+- `aten/src/ATen/native/native_functions.yaml` - Operator declarations (for checking tags, dispatch keys, structured kernels)
+- `tools/autograd/derivatives.yaml` - Backward formulas (for checking if an op should register here)
+- `aten/src/ATen/native/tags.yaml` - Operator semantic tags
diff --git a/.claude/skills/pr-review/review-checklist.md b/.claude/skills/pr-review/review-checklist.md
index 3f4f8929d7a74..7c588027cbdca 100644
--- a/.claude/skills/pr-review/review-checklist.md
+++ b/.claude/skills/pr-review/review-checklist.md
@@ -7,7 +7,10 @@ This checklist covers areas that CI cannot check. Skip items related to linting,
 ### Abstractions and Design
 
 - [ ] **Clear abstractions** - State management is explicit; no dynamic attribute setting/getting
-- [ ] **Match existing patterns** - Code follows architectural patterns already in the codebase
+- [ ] **No side-channel communication** - If behavior changes based on a hidden flag or dynamically-set attribute, the interface itself should change instead (different function signature, different class, different code path). Side-channel patterns (set a private flag in one place, check it in another via `getattr`) create undocumented behavioral modes
+- [ ] **Proper interface, not on/off flags** - A private boolean that switches between two fundamentally different behaviors should be two separate code paths or a proper interface change, not a flag
+- [ ] **Interface documentation** - New internal calling conventions, protocols, or contracts between components must have concrete documentation: what the caller provides, what the callee receives, what invariants hold, and cleanup responsibilities. Motivational comments ("this allows X") are not interface documentation
+- [ ] **Match existing patterns in the same file** - Before accepting new code in a file, read how similar features are already implemented in that same file. If the file uses class attributes for boolean flags, new boolean flags must use class attributes. If the file uses a specific setter pattern, new setters must use the same pattern
 - [ ] **No over-engineering** - Only requested changes are made; no speculative features
 - [ ] **No premature abstraction** - Helpers and utilities are only created when reused; three similar lines is better than a one-use helper
 - [ ] **No trivial helpers** - Avoid 1-2 LOC helper functions used only once (unless significantly improves readability)
@@ -29,14 +32,128 @@ When a PR introduces new API patterns, carefully evaluate the broader implicatio
 - [ ] **Useful comments only** - Comments explain non-obvious context that cannot be inferred locally. For large comment use the `# Note [Good note title]` and `See Note [Good note title]` to write larger comments that can be referenced from multiple places in the codebase.
 - [ ] **No backward-compatibility hacks** - Unused code is deleted completely, not renamed with underscores or marked with "removed" comments
 - [ ] **Appropriate complexity** - Solutions are as simple as possible for the current requirements
+- [ ] **Documentation shows correct patterns only** - Docs and markdown files should show the right way to do things directly, not anti-patterns followed by corrections. Code examples must have correct indentation, names, and syntax
+
+### Initialization and Module Design
+
+- [ ] **No fragile init ordering** - If multiple imports/calls must happen in a specific undocumented order, flag the design. Dependencies should be explicit or combined into a single entry point
+- [ ] **Idempotent global state** - Registries and global lists that accumulate entries must handle multiple calls safely (no duplicate registration, clear cleanup story)
+
+## PyTorch Infrastructure
+
+When a PR touches code in the scope of any item below, **stop and investigate** whether the established infrastructure should be used.
 
-### Common Issues to Flag
+### C++ Kernel Infrastructure
 
-- Dynamic `setattr`/`getattr` for state management (prefer explicit class members)
-- Unused imports, variables, or dead code paths
-- Copy-pasted code that could be a shared helper
-- Magic numbers without explanation
-- Overly defensive error handling for impossible cases
+- [ ] **TensorIterator** — PR adds or modifies C++ kernel code that iterates over tensor data (raw pointers, `at::parallel_for`, manual contiguity checks, manual output reshape/resize)
+- [ ] **DispatchStub** — PR adds C++ kernel code with manual `if (device_type == kCPU) ... else if (device_type == kCUDA)` dispatch instead of using `DECLARE_DISPATCH` / `DEFINE_DISPATCH` / `REGISTER_DISPATCH` from `aten/src/ATen/native/DispatchStub.h`
+- [ ] **Structured Kernels** — PR adds a new ATen operator with separate hand-written functional, inplace, and out= variants instead of using `structured: True` + `structured_delegate` in `native_functions.yaml` to generate boilerplate
+- [ ] **TORCH_CHECK variants** — PR uses generic `TORCH_CHECK` for conditions that have a more specific variant: `ValueError` → `TORCH_CHECK_VALUE`, `IndexError` → `TORCH_CHECK_INDEX`, `TypeError` → `TORCH_CHECK_TYPE`, `NotImplementedError` → `TORCH_CHECK_NOT_IMPLEMENTED`
+- [ ] **AT_DISPATCH macros** — PR manually switches on `dtype` with `if (dtype == kFloat) ... else if (dtype == kDouble)` instead of using `AT_DISPATCH_FLOATING_TYPES`, `AT_DISPATCH_ALL_TYPES_AND`, or the `AT_DISPATCH_SWITCH` / `AT_DISPATCH_CASE` pattern from `aten/src/ATen/Dispatch.h`
+- [ ] **Device guards (RAII)** — PR manually saves/restores device context (`cudaSetDevice` + try/catch) instead of using `DeviceGuard` or `OptionalDeviceGuard` from `c10/core/DeviceGuard.h`. **Note:** Operators registered in `native_functions.yaml` get automatic `DeviceGuard` insertion from codegen (controlled by `device_guard: True`, the default) — do NOT flag missing device guards for these ops unless they explicitly set `device_guard: False`
+- [ ] **Memory format propagation** — PR allocates output tensors with `at::empty(shape, options)` (defaulting to contiguous) without calling `input.suggest_memory_format()` to preserve ChannelsLast or other input formats
+- [ ] **Subclass-safe tensor allocation** — PR uses `at::empty(shape, input.options())` instead of `input.new_empty(shape)` or `at::empty_like(input)`, which don't propagate tensor subclass metadata
+- [ ] **TORCH_LIBRARY operator registration** — PR registers operators using manual dispatcher calls instead of `TORCH_LIBRARY` / `TORCH_LIBRARY_IMPL` macros from `torch/library.h`
+- [ ] **TORCH_WARN_DEPRECATION** — PR uses `TORCH_WARN` for deprecation notices instead of `TORCH_WARN_DEPRECATION` which issues a proper `DeprecationWarning`
+
+### CUDA & Device Management
+
+- [ ] **C10_CUDA_CHECK** — PR calls raw CUDA APIs (`cudaMalloc`, `cudaMemcpy`, etc.) without wrapping in `C10_CUDA_CHECK()` from `c10/cuda/CUDAException.h`
+- [ ] **C10_CUDA_KERNEL_LAUNCH_CHECK** — PR launches CUDA kernels with `<<<>>>` syntax but doesn't follow with `C10_CUDA_KERNEL_LAUNCH_CHECK()` immediately after to detect launch errors early
+- [ ] **CUDAStreamGuard** — PR manually manages CUDA streams (`cudaStreamCreate`/`cudaStreamDestroy`) instead of using `CUDAStreamGuard` or getting streams from `at::cuda::getCurrentCUDAStream()` / `getStreamFromPool()`
+- [ ] **CUDAEvent synchronization** — PR uses `cudaDeviceSynchronize()` or `cudaStreamSynchronize()` for cross-stream ordering instead of `CUDAEvent::record()` + `CUDAEvent::block()` which avoids unnecessary full synchronization
+- [ ] **recordStream for allocator** — PR uses tensors on non-default CUDA streams without calling `c10::cuda::CUDACachingAllocator::recordStream()` to prevent premature memory reuse
+- [ ] **CUDA graph compatibility** — PR adds host-GPU synchronization, unpinned memory transfers, or other graph-unsafe operations without checking `currentStreamCaptureStatusMayInitCtx()` to detect CUDA graph capture mode
+- [ ] **AcceleratorHooksInterface** — PR adds device-specific `#ifdef USE_CUDA` blocks in generic code instead of using `AcceleratorHooksInterface` from `aten/src/ATen/detail/AcceleratorHooksInterface.h` for device-agnostic behavior
+- [ ] **DeviceGuardImplInterface** — PR implements custom device management without going through `DeviceGuardImplInterface` from `c10/core/impl/DeviceGuardImplInterface.h`, bypassing the standard device abstraction layer
+
+### Operator Registration & Codegen
+
+- [ ] **native_functions.yaml** — PR adds a new ATen operator by writing manual C++ bindings and Python wrappers instead of declaring it in `aten/src/ATen/native/native_functions.yaml` and letting codegen produce the boilerplate
+- [ ] **Operator tags** — PR adds an operator to `native_functions.yaml` without appropriate tags from `tags.yaml` (e.g., `pointwise`, `reduction`, `view_copy`, `core`, `pt2_compliant_tag`)
+- [ ] **Missing Composite fallback** — PR adds a new operator to `native_functions.yaml` with only backend-specific dispatch keys (e.g., `CPU`, `CUDA`) but no `CompositeImplicitAutograd` or `CompositeExplicitAutograd` fallback. Without a Composite entry, the op will fail on all backends that don't have an explicit registration (XLA, MPS, HPU, PrivateUse1, etc.). Every new op should either have a Composite implementation or a clear justification for why it can only work on specific backends
+- [ ] **Meta function registration** — PR adds a new operator without a meta (shape-only) implementation, blocking `torch.compile` and `torch.export`. Meta implementations can be registered in Python via `@register_meta` from `torch/_meta_registrations.py`, via `torch.library.impl(..., "Meta")`, or in C++ as a structured kernel with a `meta` dispatch key or via any `Composite` dispatch key in `native_functions.yaml` (since Composite kernels automatically work on Meta tensors)
+- [ ] **Fake tensor implementation** — PR adds a custom op (registered via `torch.library`) without a fake implementation. Custom ops need `@register_fake` / `my_op.register_fake()` for `torch.compile` to trace through the op. For C++ ops registered via `native_functions.yaml`, the meta kernel serves this purpose. For Python `torch.library` custom ops, use `@torch.library.register_fake("mylib::my_op")` or `@my_op.register_fake` to provide a shape/dtype-only implementation. The fake impl receives `FakeImplCtx` with `ctx.new_dynamic_size()` for data-dependent output shapes
+- [ ] **Schema annotations** — PR defines operator schemas without proper alias annotations (`Tensor(a)`, `Tensor(a!)`) for view and in-place ops, which breaks functionalization and autograd's alias tracking
+
+### Autograd
+
+- [ ] **derivatives.yaml** — PR writes a custom `autograd.Function` subclass for an operation that should have its backward formula registered in `tools/autograd/derivatives.yaml` (the centralized backward formula registry for ATen ops)
+- [ ] **setup_context pattern** — PR writes `autograd.Function` with `forward(ctx, ...)` (legacy pattern) instead of separated `forward(...)` + `setup_context(ctx, inputs, output)` which is required for functorch compatibility (vmap, grad)
+- [ ] **ctx.save_for_backward** — PR saves tensors in `autograd.Function` via `ctx.my_tensor = tensor` instead of `ctx.save_for_backward(tensor)`, causing memory leaks by keeping tensors alive longer than needed
+- [ ] **gradcheck testing** — PR adds custom backward logic but doesn't test it with `torch.autograd.gradcheck()` / `gradgradcheck()` which verify numerical correctness of gradients via finite differences
+- [ ] **Forward-mode AD** — PR adds a new differentiable op with backward formula in `derivatives.yaml` but doesn't add a `result:` entry for forward-mode AD (JVP). Can often use `auto_element_wise` or `auto_linear` for automatic generation
+- [ ] **register_autograd for custom ops** — PR writes a full `autograd.Function` subclass for a custom op registered via `torch.library` instead of using the simpler `@my_op.register_autograd(backward, setup_context=...)` API
+- [ ] **Vmap rule for custom ops** — PR adds a custom op or `autograd.Function` without a vmap rule (`generate_vmap_rule = True` or manual `vmap()` static method), breaking `torch.vmap` support
+
+### Python Utilities
+
+- [ ] **__torch_function__ support** — PR adds a new Python-level function that takes tensors but doesn't check `has_torch_function()` / call `handle_torch_function()`, breaking tensor subclass dispatch
+- [ ] **Pytree registration** — PR manually flattens/unflattens custom container types (dataclasses, named tuples) instead of registering them with `torch.utils._pytree.register_pytree_node()` or `register_dataclass()`
+- [ ] **tree_map** — PR manually walks nested structures of tensors with recursive functions instead of using `torch.utils._pytree.tree_map()`
+- [ ] **_DecoratorContextManager** — PR implements a context manager that should also work as a decorator but doesn't inherit from `torch.utils._contextlib._DecoratorContextManager`
+- [ ] **Deprecation utilities** — PR deprecates a function using ad-hoc `warnings.warn()` calls instead of PyTorch's deprecation infrastructure (`lazy_deprecated_import` for module-level, `TORCH_WARN_DEPRECATION` for C++)
+- [ ] **No print statements** — PR adds `print()` calls for debugging or diagnostics. Use `torch._logging` utilities instead (`getArtifactLogger`, `LazyString`, `warning_once`). For the `torch.compile` stack specifically, use `trace_structured()` for structured artifacts that integrate with `tlparse` for production debugging. No bare `print()` should ever land in production code
+- [ ] **torch.backends context** — PR manually saves/restores backend flags (`cudnn.deterministic`, etc.) instead of using the `torch.backends.cudnn.flags()` context manager
+
+### nn Module Patterns
+
+- [ ] **ModuleList / ModuleDict** — PR stores submodules in plain Python `list` or `dict` instead of `nn.ModuleList` or `nn.ModuleDict`, causing them to be invisible to `parameters()`, `to()`, `state_dict()`, etc.
+- [ ] **nn.init methods** — PR manually initializes weights with `self.weight.data.normal_(0, 0.01)` instead of using `torch.nn.init.kaiming_uniform_()`, `xavier_uniform_()`, etc., which handle fan-in/fan-out calculations correctly
+- [ ] **Parametrization framework** — PR implements custom weight reparameterization via forward pre-hooks (the deprecated pattern) instead of using `torch.nn.utils.parametrize.register_parametrization()`
+- [ ] **_load_from_state_dict versioning** — PR changes a module's parameter layout without implementing `_load_from_state_dict()` for backward-compatible loading of old checkpoints (see BatchNorm's `_version = 2` pattern)
+- [ ] **clip_grad_norm_** — PR manually computes gradient norms and clips in training loops instead of using `torch.nn.utils.clip_grad_norm_()` or `clip_grad_value_()`
+- [ ] **LazyModule pattern** — PR implements deferred parameter initialization with manual shape inference in `forward()` instead of using `LazyModuleMixin` with `UninitializedParameter`
+
+### Dynamo / Inductor / Compile
+
+- [ ] **@register_lowering** — PR adds Inductor support for an op by modifying core lowering code instead of using `@register_lowering(aten.my_op)` from `torch/_inductor/lowering.py` with automatic type promotion and broadcasting
+- [ ] **Inductor decompositions** — PR writes a full Inductor lowering for a complex op that can be decomposed into simpler already-lowered ops via `@register_decomposition` in `torch/_inductor/decomposition.py`
+- [ ] **CustomGraphPass** — PR writes ad-hoc FX graph iteration for Inductor optimization instead of implementing `CustomGraphPass` (with `__call__` and `uuid()`) from `torch/_inductor/custom_graph_pass.py`
+- [ ] **config.patch** — PR manually saves/restores Dynamo or Inductor config values in tests instead of using `torch._dynamo.config.patch()` as a decorator or context manager
+- [ ] **Graph break hints** — PR calls `unimplemented()` in Dynamo without providing `gb_type`, `explanation`, or `hints` (like `SUPPORTABLE`, `FUNDAMENTAL`), making it hard for users to understand and fix graph breaks
+- [ ] **Dynamo trace rules** — PR adds manual skip/inline logic in Dynamo variable tracking instead of updating `manual_torch_name_rule_map`, `MOD_INLINELIST`, or `MOD_SKIPLIST` in `torch/_dynamo/trace_rules.py`
+- [ ] **torch.compile compatibility** — PR adds a new op or modifies an existing one without verifying it works under `torch.compile` (should test with `pt2_compliant_tag` and run opcheck)
+
+### FX / Export
+
+- [ ] **FX PassBase** — PR writes a custom FX graph transformation with manual graph walking instead of inheriting from `PassBase` (with `requires()`, `call()`, `ensures()`) from `torch/fx/passes/infra/pass_base.py`
+- [ ] **FX PassManager** — PR manually orders and applies multiple FX passes instead of using `PassManager` with `this_before_that_pass_constraint` from `torch/fx/passes/infra/pass_manager.py`
+- [ ] **FX Interpreter** — PR manually iterates FX graph nodes and tracks values in a dict instead of subclassing `torch.fx.Interpreter` which provides structured `run_node()` / `call_function()` / `call_module()` overrides
+- [ ] **Subgraph rewriter** — PR manually matches and replaces graph patterns instead of using `replace_pattern()` from `torch/fx/subgraph_rewriter.py`
+- [ ] **ShapeProp** — PR manually executes FX graphs to annotate shapes on nodes instead of using `ShapeProp(gm).propagate(*args)` from `torch/fx/passes/shape_prop.py`
+- [ ] **torch.export dynamic shapes** — PR hard-codes tensor shapes in export constraints instead of using `Dim(name, min, max)` and `dims()` from `torch/export/dynamic_shapes.py`
+- [ ] **make_fx** — PR manually initializes `torch.fx.Tracer` for proxy-based tracing instead of using `make_fx(f, tracing_mode="symbolic")` from `torch/fx/experimental/proxy_tensor.py`
+
+### Type Promotion & Dtypes
+
+- [ ] **elementwise_dtypes / TensorIterator** — PR manually implements type promotion logic for elementwise ops. In Python, use `elementwise_dtypes()` from `torch/_prims_common/` with the appropriate `ELEMENTWISE_TYPE_PROMOTION_KIND`. In C++, use `TensorIteratorConfig` which handles type promotion automatically: call `.promote_inputs_to_common_dtype(true)` and `.cast_common_dtype_to_outputs(true)` on the config builder, then `TensorIterator` computes `common_dtype()` for the kernel and handles all input/output casting. Kernels should operate on `iter.common_dtype()` via `AT_DISPATCH` rather than manually checking and promoting dtypes
+- [ ] **result_type** — PR manually resolves output dtype from mixed-dtype inputs instead of using `torch.result_type()` (Python) or `at::result_type()` / `update_result_type_state()` (C++)
+- [ ] **Complex dtype handling** — PR manually maps between complex and real dtypes (e.g., `complex64` to `float32`) instead of using `corresponding_real_dtype()` / `corresponding_complex_dtype()` from `torch/_prims_common/`
+- [ ] **promoteTypes** — PR writes manual dtype promotion tables instead of using `c10::promoteTypes(a, b)` from `c10/core/ScalarType.h`
+
+### Serialization
+
+- [ ] **weights_only=False** — PR adds `torch.load(..., weights_only=False)`, explicitly opting out of safe deserialization. `weights_only=True` is already the default; setting it to `False` enables arbitrary code execution via pickle and is almost never the right thing to do. Flag this and ask the author to register safe globals via `torch.serialization.add_safe_globals()` instead
+- [ ] **safe_globals** — PR adds new types to serialization that should be loadable with `weights_only=True` but doesn't register them via `torch.serialization.add_safe_globals()`
+- [ ] **skip_data context** — PR implements metadata-only checkpoint inspection by reading full tensors instead of using `torch.serialization.skip_data()` context manager
+
+### Distributed
+
+- [ ] **DeviceMesh** — PR manually creates multiple `ProcessGroup`s for multi-dimensional parallelism (TP + DP) instead of using `DeviceMesh` from `torch/distributed/device_mesh.py` which manages this automatically
+- [ ] **Distributed testing** — PR spawns multiple real processes for distributed unit tests instead of using `MultiThreadedPG` from `torch/testing/_internal/distributed/multi_threaded_pg.py` for single-process testing
+
+### Tensor Subclasses
+
+- [ ] **_make_wrapper_subclass** — PR creates tensor subclasses by calling `torch.Tensor.__new__()` directly instead of using `torch.Tensor._make_wrapper_subclass()` which properly sets up the subclass wrapper
+- [ ] **__tensor_flatten__ / __tensor_unflatten__** — PR adds a tensor subclass without implementing `__tensor_flatten__()` and `__tensor_unflatten__()`, breaking serialization and `torch.compile` support
+
+### Miscellaneous
+
+- [ ] **torch._check** — PR uses `assert` or `if not cond: raise` in Python op implementations instead of `torch._check()` / `torch._check_is_size()` which work correctly with meta tensors and symbolic shapes
+- [ ] **C++ extension building** — PR uses raw `setuptools` or `distutils` for building C++ extensions instead of `torch.utils.cpp_extension.CppExtension` / `CUDAExtension` / `load_inline()` which handle compiler flags, ABI, and includes
+- [ ] **register_package for custom devices** — PR adds custom device serialization handling by monkey-patching `torch.save`/`torch.load` instead of using `torch.serialization.register_package()` to register location tag and restore functions
+- [ ] **@register_backend** — PR adds a `torch.compile` backend by manually modifying internal dispatch tables instead of using `@torch._dynamo.backends.registry.register_backend(name=...)`
 
 ## Testing
 
@@ -48,41 +165,29 @@ When a PR introduces new API patterns, carefully evaluate the broader implicatio
 
 ### Test Patterns
 
-- [ ] **Use OpInfo** - Any testing for an operator or a cross cutting feature must be done via OpInfo
+- [ ] **Proper module ownership** - Test files must have a real `# Owner(s): ["module: ..."]` label, not `"module: unknown"`. The author should create a new module label if needed and add themselves as owner
+- [ ] **Use OpInfo** - Any testing for an operator or a cross-cutting feature must be done via OpInfo. Flag manual tests (e.g., `assertEqual(a + b, expected)`) for operators that already have OpInfo entries — these are redundant and will rot. When a PR adds dtype/device support to an operator, the testing should come from existing OpInfo infrastructure automatically (e.g., by adding the dtype to the operator's OpInfo `dtypes`), not from new manual tests. Likewise, a test checking a specific behavior for a single operator should not be a standalone test — the OpInfo infrastructure for that test category should be updated to cover the behavior across all applicable operators
+- [ ] **Use ModuleInfo** - Manual forward/backward tests for `nn.Module` subclasses should use `ModuleInfo` from `torch/testing/_internal/common_modules.py` and the `@modules` decorator instead of hand-written per-module tests
 - [ ] **Use TestCase** - Tests inherit from `torch.testing._internal.common_utils.TestCase`
 - [ ] **Use run_tests** - Test file ends with `if __name__ == "__main__": run_tests()`
-- [ ] **Use assertEqual for tensors** - Tensor comparisons use `assertEqual`, not raw assertions
+- [ ] **Use assertEqual for tensors** - Tensor comparisons use `assertEqual`, not raw assertions or `torch.allclose`
+- [ ] **Device generic** - Any test checking compute result should happen in a device-generic test class (taking device as an argument) via `instantiate_device_type_tests`. Device-specific tests should be very rare and in device-specific test files
+- [ ] **Use @dtypes** - PR writes separate test methods per dtype or manual `for dtype in [...]` loops instead of using the `@dtypes(...)` decorator from `common_device_type.py`
+- [ ] **Use @parametrize** - PR duplicates test methods that differ only in a parameter instead of using `@parametrize` from `common_utils.py`
+- [ ] **Use @ops for operator tests** - PR writes manual per-operator test iterations instead of using the `@ops(op_db)` decorator which automatically parametrizes tests over OpInfo entries
+- [ ] **Use make_tensor** - PR creates test tensors with `torch.rand(shape)` (implicit CPU, implicit dtype) instead of `make_tensor(shape, device=device, dtype=dtype)` from `torch.testing` which enforces explicit device/dtype
+- [ ] **Use common dtype groups** - PR manually lists dtypes like `[torch.float32, torch.float64]` instead of using helpers like `floating_types()`, `all_types_and_complex()`, etc. from `common_dtype.py`
+- [ ] **Use toleranceOverride** - PR hard-codes tolerance values in individual assertions instead of using `@toleranceOverride` / `@precisionOverride` decorators which set per-dtype tolerances
+- [ ] **Use DecorateInfo for OpInfo skips** - PR adds `@skipIf` conditionals inside OpInfo test methods instead of using `DecorateInfo` in the OpInfo's `skips` or `decorators` tuple
+- [ ] **Use largeTensorTest** - PR manually checks free memory before large-tensor tests instead of using `@largeTensorTest("4 GB")` decorator from `common_device_type.py`
 - [ ] **Descriptive test names** - Test method names describe what is being tested
-- [ ] **Device generic** - Any test checking compute result should happen in a Device-generic test class (taking device as an argument). Device-specific test should be very rare and in device-specific test files.
 
 ### Test Quality
 
 - [ ] **Edge cases covered** - Tests include boundary conditions, empty inputs, error cases
-- [ ] **Error conditions tested** - Expected exceptions are tested with `assertRaises` or `assertRaisesRegex`
-- [ ] **No duplicated test logic** - Similar tests share a private helper method (e.g., `_test_foo(config)`) called from individual tests with different configs
-
-**Example of good test structure:**
-```python
-def _test_feature_with_config(self, flag, expected_shape):
-    """Shared test logic called by device-specific tests."""
-    x = torch.randn(10)
-    result = my_feature(x, flag)
-    self.assertEqual(result.shape, expected_shape)
-
-def test_feature_enabled(self):
-    self._test_feature_with_config(True, (10, 10))
-
-def test_feature_disabled(self):
-    self._test_feature_with_config(False, (10, 5))
-```
-
-### Common Testing Issues
-
-- Tests that only check the happy path without error cases
-- Duplicated test code that should be a parameterized helper
-- Tests that don't clean up resources (files, CUDA memory)
-- Flaky tests (timing-dependent, order-dependent, golden value)
-- Tests that skip without clear justification
+- [ ] **Error conditions tested** - Expected exceptions are tested with `assertRaisesRegex`, not bare `assertRaises`. `assertRaisesRegex` verifies both the exception type and message, catching cases where the right exception is raised for the wrong reason. Bare `assertRaises` should be flagged — always require a message pattern match
+- [ ] **No duplicated test logic** - Similar tests share a private helper method called from individual tests with different configs
+- [ ] **Use weakref for lifetime testing** - PR uses `sys.getrefcount()` to test whether objects are kept alive. Use `weakref.ref()` instead — create a weak reference, delete the strong references, then check if the weakref is dead (`wr() is None`). `sys.getrefcount` is a CPython implementation detail that varies across versions and is fragile
 
 ## Security
 
@@ -105,6 +210,43 @@ When reviewing changes to PyTorch APIs and user-facing code:
 - [ ] **Distributed primitives** - `torch.distributed`, RPC, and TCPStore have no auth/encryption and accept connections from anywhere; they are for internal networks only, not untrusted environments
 - [ ] **No new pickle usage** - Avoid adding `pickle.load` or `torch.load` without `weights_only=True` on paths that could receive untrusted data
 
+## Thread Safety & Concurrency
+
+### Python Threading
+
+- [ ] **No unprotected shared mutable state** - Shared data structures accessed from multiple threads are protected by locks or are inherently thread-safe
+- [ ] **Lock ordering** - When multiple locks are acquired, ordering is consistent to avoid deadlocks
+- [ ] **No GIL-reliant correctness** - Code that mutates shared state should not rely on the GIL for thread safety, since the GIL may not be present in free-threaded builds
+
+### C++ Threading
+
+- [ ] **No data races** - Shared mutable state is protected by mutexes or uses atomics with appropriate memory ordering
+- [ ] **RAII lock guards** - Prefer `std::lock_guard` or `std::unique_lock` over manual `lock()`/`unlock()` to ensure exception-safe unlocking
+- [ ] **No lock-order inversions** - When acquiring multiple locks, a consistent global ordering is followed
+- [ ] **Correct atomic memory ordering** - `std::memory_order_relaxed` is only used when ordering with other operations is genuinely unnecessary; default to `seq_cst` or use `acquire`/`release` pairs
+
+### CPython C API Thread Safety
+
+This is particularly important for PyTorch's autograd, which has multi-threaded C++ code calling into the CPython C API.
+
+- [ ] **GIL held for Python object access** - Any code that touches `PyObject*` (incref, decref, attribute access, container mutation) must hold the GIL. When releasing the GIL for long-running C++ work (`Py_BEGIN_ALLOW_THREADS`), verify no Python objects are accessed in that region
+- [ ] **Borrowed references across GIL release** - Borrowed references (`PyTuple_GET_ITEM`, `PyList_GET_ITEM`) become unsafe if the GIL is released and reacquired, since another thread may have mutated the container
+- [ ] **Decref-before-update hazard** - When replacing an item in a container (tuple, list, dict), update the container slot first, then `Py_DECREF` the old value. Decref can trigger `__del__` finalizers that re-enter and observe the container in an inconsistent state. Without the GIL (free-threaded builds), this is also a data race. This is **always** a must-fix — even if "safe in practice" because of refcount guarantees, the pattern is wrong and breaks under NoGIL. The correct pattern costs nothing extra
+
+### Free-Threaded Python (NoGIL, PEP 703)
+
+CPython 3.13t+ can run without the GIL. Code that was previously safe under the GIL may have races in free-threaded builds:
+
+- [ ] **No implicit GIL serialization assumptions** - Code paths that assume only one thread can execute Python at a time are broken under NoGIL. Look for shared mutable state accessed from C extensions without explicit locking
+- [ ] **Raw `PyTuple_SET_ITEM` / `PyList_SET_ITEM`** - These are raw slot writes with no memory ordering guarantees. In free-threaded builds, concurrent reads from other threads may see stale or torn values. Consider whether the data structure could be accessed concurrently and whether atomic operations or the thread-safe API alternatives are needed
+- [ ] **Module-level mutable state in C extensions** - Global/static `PyObject*` variables or C-level caches accessed from multiple threads need synchronization in NoGIL builds
+
+### PyTorch-Specific Concurrency
+
+- [ ] **Autograd engine multi-threading** - The autograd engine runs node `apply()` methods from worker threads. Code in custom autograd node implementations must be safe for concurrent execution across different nodes, and must hold the GIL when accessing Python objects
+- [ ] **CUDA stream synchronization** - Operations across different CUDA streams require explicit synchronization (`cudaStreamSynchronize`, `cudaEventRecord`/`cudaStreamWaitEvent`). Missing synchronization can cause silent data corruption
+- [ ] **DataLoader worker safety** - Objects shared between the main process and DataLoader worker processes (or threads) must be fork-safe or use appropriate IPC mechanisms
+
 ## Performance
 
 ### Obvious Regressions
@@ -125,9 +267,7 @@ When reviewing changes to PyTorch APIs and user-facing code:
 - [ ] **Efficient data structures** - Appropriate containers for access patterns
 - [ ] **Gradient memory** - Proper use of `no_grad()`, `detach()` to avoid unnecessary graph retention
 
-### Common Performance Issues
+### Profiling & Benchmarking
 
-- Creating new tensors inside training loops instead of pre-allocating
-- Synchronous CUDA operations where async would work
-- Keeping computation graph alive longer than needed
-- Redundant clones or copies
+- [ ] **Use torch.profiler** - PR adds manual `time.time()` instrumentation instead of using `torch.profiler.profile()` context manager with `schedule()` and `tensorboard_trace_handler()`
+- [ ] **Use torch.utils.benchmark.Timer** - PR benchmarks with `time.time()` loops instead of `torch.utils.benchmark.Timer` which handles warmup, statistics, and proper CUDA synchronization
diff --git a/.claude/skills/pt2-bug-basher/SKILL.md b/.claude/skills/pt2-bug-basher/SKILL.md
new file mode 100644
index 0000000000000..5c30f9002f79c
--- /dev/null
+++ b/.claude/skills/pt2-bug-basher/SKILL.md
@@ -0,0 +1,273 @@
+---
+name: pt2-bug-basher
+disable-model-invocation: true
+description: Debug PyTorch 2 compiler stack failures including Dynamo graph breaks, Inductor codegen errors, AOTAutograd crashes, and accuracy mismatches. Use when encountering torch.compile errors, BackendCompilerFailed exceptions, recompilation issues, Triton kernel failures, FX graph problems, or when the user mentions debugging PT2, Dynamo, Inductor, or compiled model issues.
+---
+
+# PT2 Bug Basher
+
+Debug test failures and runtime errors in the PyTorch 2 compiler stack (Dynamo, Inductor, AOTAutograd, FX graphs).
+
+## Workflow Summary
+
+1. **Environment check** -- Ask the user which conda environment to use. Verify it is active by checking `$CONDA_DEFAULT_ENV`. Then run `python -c "import torch; print(torch.__version__)"` to confirm torch is importable and report the version. If the environment is not active or torch cannot be imported, stop and ask the user to activate the correct environment before proceeding.
+2. **Reproduce** -- Get a consistent reproduction of the failure
+3. **Minimize** -- Reduce the repro to the smallest possible standalone case. Strip away unrelated model logic, use minimal tensor shapes, and isolate the specific op or pattern that triggers the bug.
+4. **Add a unit test** -- **Do this BEFORE diving into code search or root cause investigation.** Add a failing test to the codebase that captures the bug. Place it in a specific, topic-appropriate test file (e.g., `test/dynamo/test_repros.py`, `test/inductor/test_torchinductor.py`, `test/export/test_export.py`). **Avoid `test/dynamo/test_misc.py`** — it is already oversized; find a more specific test file that matches the area of the bug. Use `torch.testing._internal.common_utils.TestCase` and `run_tests`. The test must fail before the fix and pass after. Having the test first keeps you grounded — you know exactly what "fixed" looks like before you start exploring the codebase.
+5. **Validate on main** -- Use `EnterWorktree` to create a worktree checked out at `main`. Copy the new test file into the worktree and run the test there to confirm it **fails** on main. If the test passes on main, stop — the test may not be capturing the right bug, or the bug may already be fixed. Exit the worktree with `ExitWorktree` (action: remove) and return to the working branch before continuing.
+6. **Gather logs** -- Run with appropriate `TORCH_LOGS` settings
+7. **Classify** -- Use the [Error Triage](#error-triage) table to identify the category
+8. **Inspect artifacts** -- Check FX graphs, IR, and generated code via `TORCH_COMPILE_DEBUG=1`
+9. **Identify root cause** -- Trace from the error back through the compilation pipeline
+10. **Fix** -- Apply the fix
+11. **Verify** -- Run the new unit test AND nearby related existing tests (e.g., if you changed how `is_exporting` works, also run the existing `test_is_exporting` export test). Use `pytest -k` to quickly run related tests by name. The task is not complete until all pass.
+12. **Self-review** -- Use the `/pr-review` skill to review your own changes before presenting them. Fix any issues it flags.
+13. **Celebrate** -- Summarize the changes: explain the root cause, what was changed and why, and which tests were added/verified. Then tell the user the bug is squashed. Include a fun, varied motivational message or easter egg to keep spirits high (e.g., a pun, a quote, an ASCII art bug getting squashed). Keep it short and different each time.
+
+## Investigation Strategy
+
+### Prefer direct tools over meta_codesearch
+
+Use `Grep`, `Glob`, and `Read` directly for code exploration. **Do not spawn `meta_codesearch` agents** — they are slow and expensive. The [Architectural Knowledge](#architectural-knowledge) and [Key Source Files](#key-source-files) sections below should give you enough context to know where to look. A targeted `Grep` for a function name is always faster.
+
+### Know which compilation mode you're in
+
+Before reading implementation code, determine the compilation mode. These share code but diverge in important ways:
+- **`torch.compile`** -- Dynamo + Inductor. `tx.export=False`, no `_compiling_state_context()`.
+- **`torch.export` (strict)** -- `tx.export=True`, `_compiling_state_context()` active.
+- **`torch.export` (non-strict, **the default**)** -- Uses Dynamo via `fullgraph_capture` but `tx.export` may differ from strict. `_compiling_state_context()` active. Check `torch._export.config.use_new_tracer_experimental` — it changes which code path is used.
+
+### Distinguish trace-time vs runtime
+
+Many PT2 bugs come from confusing these two:
+- **Trace-time**: Inside Dynamo's symbolic interpreter. Dynamo intercepts function calls and may constant-fold them (e.g., `is_exporting()` → `ConstantVariable(True)`).
+- **Runtime**: Real tensors, real Python calls, module-level flags like `torch.compiler._is_exporting_flag`.
+
+When debugging, add temporary `print()` statements directly in the source file rather than monkey-patching from outside — dispatch chains make monkey-patching unreliable.
+
+## Gathering Information
+
+Pick the right diagnostic tool based on the error category:
+
+- **Quick overview**: `TORCH_LOGS="+dynamo,graph_breaks,recompiles" python your_script.py`
+- **Full debug artifacts**: `TORCH_COMPILE_DEBUG=1 python your_script.py` — creates `torch_compile_debug/` with FX graphs, Inductor IR, and generated code
+- **Generated code only**: `TORCH_LOGS="output_code" python your_script.py`
+- **Structured tracing**: `TORCH_TRACE=/path/to/trace python your_script.py` then `tlparse /path/to/trace`
+- **Single-threaded (for pdb)**: `TORCHINDUCTOR_COMPILE_THREADS=1 python your_script.py`
+
+## Error Triage
+
+Classify the failure using the error message and traceback:
+
+| Error Pattern | Category | Jump To |
+|---|---|---|
+| `Unsupported: ...` or `graph break` in logs | Graph break | [Graph Breaks](#graph-breaks) |
+| `BackendCompilerFailed` | Inductor/backend crash | [Backend Failures](#backend-compiler-failures) |
+| `RecompileError` or `cache_size_limit` | Recompilation | [Recompilation](#recompilation-issues) |
+| Accuracy mismatch / wrong numerical output | Accuracy | [Accuracy](#accuracy-issues) |
+| `InternalTorchDynamoError` | Dynamo bug | [Internal Errors](#internal-dynamo-errors) |
+| Segfault or CUDA IMA | Runtime crash | [Runtime Crashes](#runtime-crashes) |
+| Triton assertion / index out of bounds | Triton kernel bug | [Triton Failures](#triton-kernel-failures) |
+
+## Debugging by Category
+
+### Graph Breaks
+
+Graph breaks split the compiled graph into smaller subgraphs, often causing performance regressions or unexpected behavior.
+
+**Diagnosis:**
+```bash
+TORCH_LOGS="graph_breaks" python your_script.py
+```
+
+**Key files:**
+- `torch/_dynamo/exc.py` -- `Unsupported` exception class
+- `torch/_dynamo/variables/` -- where most graph break decisions happen
+
+**Common causes:**
+- Unsupported Python constructs (data-dependent control flow, unsupported builtins)
+- Tensor operations that can't be traced (in-place ops on inputs, unsupported dtypes)
+- Calls to non-traceable functions
+
+**Fix approach:**
+1. Read the graph break message to identify the unsupported operation
+2. Check if there's a decomposition or supported alternative
+3. If the operation genuinely can't be traced, consider `torch._dynamo.allow_in_graph` or restructuring user code
+
+### Backend Compiler Failures
+
+`BackendCompilerFailed` means Inductor (or another backend) crashed during compilation.
+
+**Diagnosis:**
+```bash
+TORCHDYNAMO_REPRO_AFTER=aot TORCHDYNAMO_REPRO_LEVEL=2 python your_script.py
+```
+
+This generates `minifier_launcher.py` that isolates the minimal failing graph.
+
+**Key files:**
+- `torch/_dynamo/repro/after_aot.py` -- repro/minifier for post-AOT failures
+- `torch/_inductor/` -- the backend itself
+
+**Fix approach:**
+1. Run the minifier to get a minimal reproduction
+2. Inspect the FX graph (`TORCH_COMPILE_DEBUG=1`) to understand what ops are involved
+3. Check if it's a lowering issue (`torch/_inductor/lowering.py`), scheduling issue, or codegen issue
+4. Look at the generated output code if the error is in codegen
+
+### Recompilation Issues
+
+Excessive recompilation happens when guards are too specific, causing cache misses.
+
+**Diagnosis:**
+```bash
+TORCH_LOGS="recompiles,recompiles_verbose,guards" python your_script.py
+```
+
+**Key config:**
+- `torch._dynamo.config.recompile_limit` (default: 8)
+- `torch._dynamo.config.fail_on_recompile_limit_hit` -- set to `True` to get a hard error
+
+**Common causes:**
+- Changing tensor shapes without marking them dynamic
+- Python scalar values that change between calls
+- Global state mutations between calls
+
+**Fix approach:**
+1. Read the recompilation reason from logs
+2. Identify the failing guard
+3. Either mark the relevant dimension as dynamic with `torch._dynamo.mark_dynamic()` or fix the source of guard instability
+
+### Accuracy Issues
+
+The compiled model produces different numerical results than eager mode.
+
+**Diagnosis:**
+```bash
+TORCHDYNAMO_REPRO_AFTER=aot TORCHDYNAMO_REPRO_LEVEL=4 python your_script.py
+```
+
+This compares compiled vs. eager with an fp64 reference and dumps a repro if accuracy fails.
+
+**Key utilities:**
+- `torch/_dynamo/debug_utils.py` -- `same_two_models()`, `backend_accuracy_fails()`, `cast_to_fp64()`
+- `torch._dynamo.config.repro_tolerance` (default: 1e-3)
+
+**Fix approach:**
+1. Get the minimal failing graph from the minifier
+2. Compare eager vs. compiled output at fp64 precision
+3. Binary search through ops to find the diverging operation
+4. Check for known numerical issues (reduction order, fused kernels, dtype promotions)
+
+### Internal Dynamo Errors
+
+`InternalTorchDynamoError` indicates a bug in Dynamo itself.
+
+**Diagnosis:**
+```bash
+TORCHDYNAMO_VERBOSE=1 python your_script.py
+# or equivalently:
+TORCH_LOGS="+dynamo" python your_script.py
+```
+
+**Key files:**
+- `torch/_dynamo/symbolic_convert.py` -- bytecode interpreter
+- `torch/_dynamo/variables/` -- variable tracking system
+- `torch/_dynamo/guards.py` -- guard generation
+
+**Fix approach:**
+1. Get the full stack trace with `TORCHDYNAMO_VERBOSE=1`
+2. Identify which bytecode instruction or variable type caused the crash
+3. Create a minimal repro (the error message often includes a minifier path)
+4. Debug with `TORCHINDUCTOR_COMPILE_THREADS=1` and pdb if needed
+
+### Runtime Crashes
+
+Segfaults and CUDA illegal memory access errors during execution of compiled code.
+
+**Diagnosis (make crash deterministic):**
+```bash
+PYTORCH_NO_CUDA_MEMORY_CACHING=1 CUDA_LAUNCH_BLOCKING=1 python your_script.py
+```
+
+**For CUDA IMA, add NaN checks:**
+```bash
+TORCHINDUCTOR_NAN_ASSERTS=1 python your_script.py
+```
+
+**For Inductor-level sync debugging:**
+```python
+torch._inductor.config.triton.debug_sync_kernel = True  # sync after every kernel
+torch._inductor.config.triton.debug_sync_graph = True   # sync before/after graph
+```
+
+**Fix approach:**
+1. Make the crash deterministic with `PYTORCH_NO_CUDA_MEMORY_CACHING=1 CUDA_LAUNCH_BLOCKING=1`
+2. Check if it's an input mismatch (shapes, devices, dtypes)
+3. Inspect the generated kernel code with `TORCH_LOGS="output_code"`
+4. Use `TORCHINDUCTOR_NAN_ASSERTS=1` to find the first kernel producing bad values
+5. Check for dynamic shapes issues (historically a common source of IMA)
+
+### Triton Kernel Failures
+
+Triton assertion failures or index-out-of-bounds in generated kernels.
+
+**Diagnosis:**
+```bash
+TORCH_LOGS="output_code,schedule" python your_script.py
+```
+
+**Key files:**
+- `torch/_inductor/codegen/triton.py` -- Triton codegen
+- `torch/_inductor/scheduler.py` -- kernel fusion decisions
+
+**Fix approach:**
+1. Get the generated Triton kernel from `output_code` logs
+2. Check index computations for off-by-one or wrong stride calculations
+3. Look at the IR (`TORCH_COMPILE_DEBUG=1`) to trace back to the FX op
+4. Check if fusion decisions created invalid index combinations
+
+## Key Source Files
+
+| File | Purpose |
+|---|---|
+| `torch/_dynamo/exc.py` | Exception hierarchy and error formatting |
+| `torch/_dynamo/debug_utils.py` | Minifier support, accuracy checking, input serialization |
+| `torch/_dynamo/repro/after_dynamo.py` | Repro/minifier for Dynamo-stage failures |
+| `torch/_dynamo/repro/after_aot.py` | Repro/minifier for post-AOTAutograd failures |
+| `torch/_dynamo/repro/aoti.py` | Repro/minifier for AOTI failures |
+| `torch/_dynamo/config.py` | Dynamo config (repro levels, recompile limits) |
+| `torch/_dynamo/variables/torch.py` | Torch function handling, tracing state functions |
+| `torch/_dynamo/variables/higher_order_ops.py` | HOP tracing (cond, map, etc.) |
+| `torch/_dynamo/symbolic_convert.py` | Bytecode interpreter, InstructionTranslator |
+| `torch/_dynamo/convert_frame.py` | Frame compilation, `fullgraph_capture` entry point |
+| `torch/_dynamo/functional_export.py` | New export tracer (`_dynamo_graph_capture_for_export`) |
+| `torch/_dynamo/eval_frame.py` | `torch._dynamo.export`, `optimize_assert` |
+| `torch/_export/_trace.py` | Export pipeline (`_export`, `_strict_export`, `_non_strict_export`, `_export_to_aten_ir`) |
+| `torch/_export/utils.py` | `_compiling_state_context()` |
+| `torch/compiler/__init__.py` | `is_compiling()`, `is_exporting()`, runtime flags |
+| `torch/_higher_order_ops/cond.py` | `torch.cond` implementation and proxy tracing |
+| `torch/_higher_order_ops/utils.py` | `reenter_make_fx` for HOP branch tracing |
+| `torch/_inductor/config.py` | Inductor config (debug flags, trace settings) |
+| `torch/_inductor/debug.py` | DebugContext, graph visualization, IR logging |
+| `torch/_logging/_registrations.py` | All registered log aliases and artifacts |
+
+## Using the Minifier
+
+The minifier reduces a failing graph to the smallest reproduction:
+
+```bash
+# Step 1: Generate the minifier launcher
+TORCHDYNAMO_REPRO_AFTER=aot TORCHDYNAMO_REPRO_LEVEL=2 python your_script.py
+
+# Step 2: Run the minifier
+python minifier_launcher.py minify
+
+# Step 3: Run the minimized repro
+python minifier_launcher.py run
+```
+
+For accuracy issues, use level 4:
+```bash
+TORCHDYNAMO_REPRO_AFTER=aot TORCHDYNAMO_REPRO_LEVEL=4 python your_script.py
+```
diff --git a/.claude/skills/scrub-issue/SKILL.md b/.claude/skills/scrub-issue/SKILL.md
new file mode 100644
index 0000000000000..aeb5e0d6839d4
--- /dev/null
+++ b/.claude/skills/scrub-issue/SKILL.md
@@ -0,0 +1,362 @@
+---
+name: scrub-issue
+disable-model-invocation: true
+description: Fetch, analyze, reproduce, and minimize GitHub issue reproductions. Use when asked to check if an issue reproduces, minimize a repro, analyze a bug report, or scrub/triage an issue for reproducibility.
+---
+
+# Minimize Issue Reproduction
+
+Fetch a GitHub issue, evaluate whether it has a reasonable repro, check if it
+still reproduces, and systematically minimize the repro to the smallest possible
+self-contained script.
+
+## Tools
+
+Assume the current environment is correct and run `python` directly. Only use
+`conda run -n <env>` for version bisection (step 5a) where you need to
+temporarily use a different environment. Use the Bash tool's `timeout`
+parameter to enforce timeouts when running repro scripts.
+
+- `gh issue view <NUMBER> --repo pytorch/pytorch` to fetch the issue body
+- `gh issue view --comments <NUMBER> --repo pytorch/pytorch` to fetch comments
+- `python <script>` to run repro scripts
+- `gh issue comment <NUMBER> --repo pytorch/pytorch --body <BODY>` to comment
+- `gh issue edit <NUMBER> --repo pytorch/pytorch --add-label <LABEL>` to add labels
+
+Multiple `gh issue edit` flags can be combined in a single command (e.g.
+`--add-label "bug,help wanted" --add-assignee "@me"`). Prefer batching
+edits into one command to minimize API calls and reduce the chance of
+auto-subscribing to notifications.
+
+### Preserving notification subscription state
+
+Modifying an issue (commenting, adding labels) auto-subscribes you to
+notifications. Use `tools/stale_issues.py` to save and restore subscription
+state:
+
+1. **Before the first modification**, save the current state. This can be
+   run in parallel with fetching the issue body/comments, but it must
+   **complete** before any `gh issue edit`, `gh issue comment`, or
+   `gh issue close` command is executed:
+   ```
+   python tools/stale_issues.py subscription save <NUMBER>
+   ```
+2. **After the last modification**, restore the saved state — **unless the
+   issue was closed** (via `gh issue close`), in which case skip the restore
+   so the user stays subscribed to follow any responses to the closure. The
+   restore must be the very last GitHub API call — do not run it in the
+   background or in parallel with any `gh issue edit` or `gh issue comment`
+   commands:
+   ```
+   python tools/stale_issues.py subscription restore <NUMBER>
+   ```
+
+If either the save or restore command fails, warn the user and continue
+without the save/restore mechanism.
+
+**Important**: Never run `gh issue edit`, `gh issue comment`, `gh issue close`,
+or subscription save/restore commands in the background. These must all run in
+the foreground so their completion can be verified before proceeding.
+If commenting or editing fails because the issue is locked, report this to the
+user and skip the modification.
+
+### Security review checklist
+
+Before running any repro code, check for the following concerns:
+
+- Network requests to untrusted URLs (requests, urllib, curl, wget)
+- File operations outside `/tmp/`
+- Shell command execution (os.system, subprocess, eval, exec) — but
+  `subprocess` used to launch `torchrun` or `mp.spawn` for distributed repros
+  is expected and not a concern
+- Downloading or loading external files (model weights, pickled objects, data
+  files) — especially `torch.load` on untrusted `.pt`/`.pth` files
+- Obfuscated code (base64-encoded strings, encoded bytes, unusual escapes)
+- Package installation (pip install, conda install)
+- Environment variable manipulation that could affect the host system — but
+  setting `MASTER_ADDR`, `MASTER_PORT`, `RANK`, `WORLD_SIZE`,
+  `CUDA_VISIBLE_DEVICES`, or other standard PyTorch/CUDA env vars is expected
+  and not a concern
+
+If any of these are present, explain the concern to the user and ask whether
+to proceed, skip, or modify the repro to remove the risky parts. If the user
+chooses to skip, still refresh the `triaged` label timestamp (remove and
+re-add, or just add if not present) before reporting that the analysis is
+finished.
+
+Even if the repro passes the checklist above, check whether the author of the
+repro code is a PyTorch collaborator by running
+`python tools/stale_issues.py collaborator-check <username>`. If the command
+exits with a non-zero status (user is not a collaborator), show the repro code
+to the user and ask them to verify it is safe to run before executing it.
+
+## Steps
+
+### 1) Fetch the issue
+
+Fetch the issue body and comments in parallel. Identify the reported repro
+script and error. If multiple repros are present, prefer the most recent one
+from the original poster. If a commenter has posted a strictly shorter and
+more self-contained repro that doesn't require additional context from the
+issue description, prefer that one. Note which repro you selected. If the
+repro code is only present in screenshots or images rather than copyable text,
+stop and report this to the user.
+
+### 2) Check if the issue is actionable
+
+Before investing effort in reproduction, check whether the issue is actionable.
+**Do not proceed to later steps** if any of the following apply:
+
+- **Already closed or resolved** in comments. Report to the user and stop.
+  Do not modify labels on closed issues.
+- **Duplicate** of another issue (linked or obviously the same bug)
+- **Not a bug report** (feature request, question, discussion, refactoring /
+  code cleanup task). If the issue is a feature request and doesn't already
+  have the `feature` label, add it via `gh issue edit`. If the issue is a
+  better-engineering / refactoring task and doesn't already have the
+  `better-engineering` label, add it via `gh issue edit`. If the issue
+  includes a repro script that demonstrates the current behavior, apply the
+  security review checklist first, then run it to verify the behavior
+  persists. If the repro needed to be modernized (e.g. updated imports for
+  renamed APIs), or if you verified that the behavior still persists, comment
+  on the issue with the findings and updated repro. After adding any
+  applicable labels (and optionally running/commenting on a repro), report
+  the analysis to the user and do not proceed to later steps.
+- **Tracking/meta issue** (umbrella issue tracking multiple bugs, burn-down
+  lists, improvement proposals without a specific repro). If the issue doesn't
+  already have the `tracker` label, add it via `gh issue edit`. Then stop and
+  report to the user.
+- **Requires unavailable hardware** (specific GPU models, TPUs, multi-node) with
+  no path to simplify. Note: CUDA is available on the current machine, so
+  single-GPU CUDA repros can be run directly.
+
+For non-actionable issues that are old (more than ~1 year), have no assignees,
+no recent progress, and are underspecified or lack concrete motivation, suggest
+closing them as "not planned" (`gh issue close --reason "not planned"`) with a
+comment explaining the rationale. Ask the user before closing.
+
+If the issue is not actionable and no GitHub-visible modification was made (no
+label added, no comment posted — saving subscription state does not count),
+refresh the `triaged` label to update the issue's
+"last updated" timestamp. A single `gh issue edit` with both `--remove-label`
+and `--add-label` doesn't work because the remove and add cancel each other
+out. Instead, chain both edits in a single Bash tool call:
+`gh issue edit ... --remove-label triaged && gh issue edit ... --add-label triaged`.
+If the issue doesn't have the `triaged` label, just add it.
+
+Then summarize why the issue is not actionable. If a label or other update was
+made, just report that the analysis is finished. Only ask the user how to
+proceed if no update was made and the situation is ambiguous. Always ask the
+user before closing an issue.
+
+### 3) Analyze the repro
+
+Evaluate whether the issue has a reasonable repro:
+
+- Is there a code snippet that can be run?
+- Are the dependencies available (CUDA, distributed, specific hardware)?
+  Note: `torch.distributed` repros often don't require special hardware — they
+  can be launched with `torchrun --nproc_per_node=1` or `mp.spawn` on a single
+  machine.
+- Is the expected error described?
+- Is the repro self-contained or does it need external data/models?
+
+If there is no repro code at all, or the issue is missing critical info (no
+error message, no description of expected vs actual behavior), add the
+`needs reproduction` label (if not already present) via `gh issue edit`, then
+stop and report to the user. Do not attempt to write a repro from the
+description without being asked.
+
+Conversely, if the issue already has the `needs reproduction` label but does
+have a valid repro, remove the label via `gh issue edit --remove-label
+"needs reproduction"`.
+
+Apply the security review checklist (see above) to the repro before running it.
+
+If the repro requires third-party packages that are not installed (e.g.
+`transformers`, `torchvision`, `numpy`), stop and ask the user how to proceed
+rather than installing them yourself.
+
+Summarize your assessment before proceeding.
+
+### 4) Check for recent verification
+
+If a comment from the last six months already confirms the issue still
+reproduces (with a matching error and a reasonable repro), stop and ask the
+user whether they want to re-verify or skip ahead to minimization.
+
+### 5) Check if it still reproduces
+
+Extract the repro code into a temporary file under `/tmp/` and run it with a
+timeout of 120 seconds. For repros involving `torch.compile` that call compile
+multiple times in the same process, add `torch._dynamo.reset()` between
+invocations to reset in-memory Dynamo state. This is unnecessary for scripts
+that compile once and exit.
+
+If the script times out, consider whether a hang is the reported bug or an
+unrelated issue, and report to the user.
+
+Record the PyTorch version before running
+(`python -c "import torch; print(torch.__version__)"`) for inclusion in the
+report.
+
+Check both the exit code and output to determine the result. An exit code > 128
+indicates the process was killed by a signal (e.g. segfault = 139, OOM kill =
+137) — this is a valid crash reproduction even without a Python traceback. If
+the repro crashes with a CUDA out-of-memory error and OOM is not the reported
+bug, try reducing tensor sizes before concluding it doesn't reproduce. For
+correctness bugs (wrong numerical results rather than crashes), the repro
+should include an assertion that fails when the bug is present. If the
+original repro only prints output without asserting, add a simple assertion
+based on the expected behavior described in the issue (e.g.
+`assert torch.allclose(actual, expected)`) so the repro has a clear
+pass/fail signal.
+
+If the result is inconsistent across runs, run 3-5 times to assess flakiness.
+For non-deterministic bugs, try setting `PYTHONHASHSEED=0` and a fixed
+`torch.manual_seed` to stabilize reproduction. Report the success/failure ratio
+to the user. Flaky repros are still valid bugs — note the flakiness in the
+issue comment (step 8) and include the success/failure ratio.
+
+Three possible outcomes (to determine which outcome applies, match on the
+exception class and a distinctive substring of the error message — the
+substring should be specific enough to identify the bug, e.g.
+`RuntimeError: expected scalar type Float` rather than just `RuntimeError`):
+
+- **Same error as reported**: the bug still reproduces. Continue to step 6.
+- **No error (passes)**: the bug may have been fixed. Try to identify the
+  fixing PR (see step 5a), then report to the user. The issue will be closed
+  in step 8.
+- **Different error**: distinguish between setup issues (missing import, renamed
+  API) that can be fixed and genuinely different bugs. If the error is due to
+  API changes between the reported version and the current version (renamed
+  functions, moved modules, changed signatures), adapt the repro to use the
+  current API while preserving the original intent. If the error is unclear,
+  consider re-running with `TORCH_LOGS=+dynamo` or other relevant logging flags
+  for more diagnostic output. Report genuinely different errors to the user.
+
+### 5a) Identify when and how it was fixed
+
+When the bug no longer reproduces, try to determine which version fixed it and
+which PR introduced the fix.
+
+**Version bisection**: Check if conda environments named `pytorch-<version>`
+(e.g. `pytorch-2.6`, `pytorch-2.8`) are available (`conda env list | grep
+pytorch-`). If they exist, binary-search across them to find the first version
+where the bug is fixed. Run from `/tmp` in a subshell and clear `PYTHONPATH`
+to avoid picking up the local source tree (which would cause `torch._C` import
+errors): `(cd /tmp && PYTHONPATH= conda run -n pytorch-<version> python
+/tmp/repro_....py)`. To speed up bisection, pick 2-3 evenly-spaced probe
+points from the candidate range and test them in parallel each round (e.g.
+if candidates are 2.2 through 2.8, test 2.4 and 2.6 simultaneously to split
+the range into thirds). If no versioned conda environments are
+available, skip bisection and just report that the bug no longer reproduces on
+the current version.
+
+**PR identification**: Try to find the specific PR that fixed the bug. Use
+version control blame on the relevant fix code to find the changeset, then look
+up the commit message for the PR number (format `(#NNNNN)`). Alternatively,
+search version control history for commits touching the relevant file with a
+related keyword. If this doesn't yield a clear answer quickly, just report the
+version — don't spend extra time on PR identification.
+
+### 6) Minimize the repro
+
+Save the original working repro to `/tmp/repro_<issue_number>_original.py`
+before making any changes.
+
+First, assess whether the repro is already reasonably minimal. Only minimize if
+the repro has significant unnecessary complexity (large models, unused code
+paths, unnecessary dependencies, etc.). When counting complexity, don't count
+irreducible boilerplate — e.g. a tensor subclass definition that only contains
+the required dunder methods (`__new__`, `__init__`, `__torch_dispatch__`,
+`__tensor_flatten__`, `__tensor_unflatten__`) is not reducible even if it's
+20+ lines. Focus on whether the trigger code and model/setup complexity can
+be meaningfully reduced. If the only possible "simplifications" are cosmetic
+(inlining variables, removing `__repr__`), skip minimization.
+
+Systematically reduce the repro by testing whether each simplification still
+triggers the same error. Use a shorter timeout (30-60 seconds) during
+minimization since simplified repros should run faster. Run multiple candidate
+simplifications in parallel when they are independent (i.e. they modify
+non-overlapping parts of the code and neither depends on what the other
+removes).
+
+**Reduction strategies (apply in roughly this order):**
+
+1. **Remove unnecessary imports and setup** (distributed init, env vars, logging)
+2. **Shrink the model** (replace large modules with minimal equivalents)
+3. **Remove the class/module wrapper** if a bare function suffices
+4. **Reduce tensor sizes** (large dims → small dims like 4 or 8)
+5. **Remove device/dtype requirements** (try CPU and float32 first)
+6. **Simplify the computation** (replace complex ops with minimal ones that still trigger the bug)
+7. **Remove unnecessary control flow** (branches, loops, conditions)
+8. **Try simpler backends** (e.g. `aot_eager` instead of inductor) if the bug is not backend-specific
+
+After each round, verify the error still reproduces — same exception class and
+a distinctive error message substring as described in step 5 (minor traceback
+differences are fine). For correctness bugs, preserve the assertion that demonstrates the
+wrong result and verify it still shows the same incorrect behavior. When merging multiple successful parallel simplifications, verify
+the combined result still reproduces since independent simplifications can
+interact. Stop minimizing when the repro is under ~20 lines of non-blank
+non-import code, or when two consecutive rounds (where a round is one full
+pass through the applicable reduction strategies) fail to simplify further.
+
+### 7) Apply trivial fixes
+
+If the analysis reveals a trivial fix (e.g. removing a stale `xfailIfTorchDynamo`
+or `expectedFailure` annotation from a test because the underlying issue is
+fixed), report the fix to the user and ask whether to apply it. Do not modify
+source files without the user's approval. In step 8, mention the fix
+regardless of whether it was applied or declined.
+
+### 8) Report findings
+
+If the repro was minimized, save it to `/tmp/repro_<issue_number>.py` so the
+user can run it directly.
+
+Present findings to the user including:
+- Whether the bug still reproduces (and on what PyTorch version)
+- The fixing PR, if identified (only if the bug no longer reproduces)
+- The minimized repro script (only if we minimized it)
+- The necessary conditions to trigger the bug
+- Any trivial fix identified in step 7 (whether applied or not)
+- Recommended next action (e.g. "still a valid bug", "appears fixed", "needs
+  more info from reporter")
+
+After presenting findings, always comment on the issue with the results.
+Keep the comment concise — don't repeat information already on the issue.
+Only include a repro if it was materially changed from the original (e.g.
+minimized, modernized imports, fixed to run on current API). Only include
+trigger conditions if they are new findings not already discussed in prior
+comments. If the only finding is "still reproduces" or "no longer reproduces",
+a short comment is sufficient.
+
+```
+This issue [still reproduces / no longer reproduces] on PyTorch <version>.
+
+[If fixed and PR identified:]
+Fixed by #NNNNN.
+
+[If minimized or modernized — only include repro if changed from original:]
+Minimized repro:
+
+\`\`\`python
+<repro script>
+\`\`\`
+
+[Only if new findings about trigger conditions:]
+All of the following are necessary to trigger the bug:
+- <condition 1>
+- <condition 2>
+
+(Analysis done by Claude.)
+```
+
+If the bug no longer reproduces, after commenting close the issue:
+`gh issue close <NUMBER> --repo pytorch/pytorch --reason completed --comment "Closing as this was fixed in PyTorch <version>."`.
+Do not ask the user before closing — fixed bugs should always be closed.
+
+After the last GitHub modification, restore the notification subscription
+state (see "Preserving notification subscription state" above) — unless the
+issue was closed, in which case skip the restore.
diff --git a/.claude/skills/triaging-issues/SKILL.md b/.claude/skills/triaging-issues/SKILL.md
index c1404f81a532f..e774d6dd40cb6 100644
--- a/.claude/skills/triaging-issues/SKILL.md
+++ b/.claude/skills/triaging-issues/SKILL.md
@@ -67,11 +67,16 @@ Use these GitHub MCP tools for triage:
 | `sev*` | Severity labels require human decision |
 | `merge blocking` | Requires human decision |
 | Any label containing "deprecated" | Obsolete |
+| `oncall: releng` | Not a triage redirect target. Use `module: ci` instead |
 
 **If blocked:** When a label is blocked by the hook, add ONLY `triage review` and stop. A human will handle it.
 
 These rules are enforced by a PreToolUse hook that validates all labels against `labels.json`.
 
+### Never Override Human Labels
+
+If a human has already applied labels (especially `ci: sev`, severity labels, or priority labels), do NOT remove or replace them. Your job is to supplement, not override.
+
 ---
 
 ## Issue Triage (for each issue)
@@ -107,6 +112,14 @@ Check if the issue body contains links to external files that users would need t
 3. Use the `needs_reproduction` template from `templates.json` to request a self-contained reproduction
 4. Do NOT add `triaged` — wait for the user to provide a reproducible example
 
+### 1.55) Needs Reproduction — Other Cases
+
+Also add `needs reproduction` when:
+- The user reports a hardware-specific issue (e.g., specific GPU model) without a self-contained repro script
+- The user references a specific model/checkpoint/dataset that is not publicly runnable in a few lines
+- The issue describes version-upgrade breakage but only provides a high-level description without a minimal script
+- The repro depends on a specific training setup, distributed environment, or non-trivial infrastructure
+
 ### 1.6) Edge Cases & Numerical Accuracy
 
 If the issue involves extremal values or numerical precision differences:
@@ -118,6 +131,17 @@ If the issue involves extremal values or numerical precision differences:
 - Precision differences between dtypes (e.g., fp32 vs fp16)
 - Fuzzer-generated edge cases
 
+**IMPORTANT — avoid keyword-triggered mislabeling:**
+
+Label based on the **root cause**, not keywords that appear in the error or title. A keyword tells you what failed, not why.
+
+- An `undefined symbol: ncclAlltoAll` error at `import torch` is a **packaging** issue (`module: binaries`), not a distributed training bug — the user never ran distributed code.
+- A `nan` in a parameter name or tolerance check is not `module: NaNs and Infs` unless the bug is actually about NaN propagation.
+- A stack trace mentioning `autograd` does not mean `module: autograd` — check whether the bug is in autograd itself or just on the call path.
+- A test failure with tolerance thresholds is `module: tests`, not `module: numerical-stability`.
+
+Ask: "Where would the fix need to be made?" That determines the label.
+
 **Action:**
 1. Add `module: edge cases` label
 2. If from a fuzzer, also add `topic: fuzzer`
@@ -130,16 +154,13 @@ If the issue belongs in another repo (vision/text/audio/RL/ExecuTorch/etc.), tra
 
 ### 2.5) PT2 Issues — Special Handling
 
-When triaging PT2 issues (torch.compile, dynamo, inductor), see [pt2-triage-rubric.md](pt2-triage-rubric.md) for detailed labeling decisions.
+**PT2 is NOT a redirect.** `oncall: pt2` is not like the other oncall labels in Step 3. PT2 issues continue through Steps 4–7 for full triage — add `oncall: pt2`, then proceed to label with `module:` labels, mark `triaged`, etc.
 
-**Key differences from general triage:**
-- For PT2 issues, you MAY apply `module:` labels (e.g., `module: dynamo`, `module: inductor`, `module: dynamic shapes`)
-- Use the rubric to determine the correct component labels
-- Only redirect to `oncall: cpu inductor` for MKLDNN-specific issues; otherwise keep in PT2 queue
+**Every `oncall: pt2` issue MUST have at least one `module:` label.** The PT2 oncall queue is too broad without a module label — the team needs to know which component is affected (e.g., `module: dynamo`, `module: inductor`, `module: helion`, `module: dynamic shapes`). If you cannot determine the specific module, use `module: compile ux` as a fallback, but always try to be specific first. See [pt2-triage-rubric.md](pt2-triage-rubric.md) for detailed guidance.
 
 ### 3) Redirect to Secondary Oncall
 
-**CRITICAL:** When redirecting issues to an oncall queue (**critical** with the exception of PT2), apply exactly one `oncall: ...` label and **STOP**. Do NOT:
+**CRITICAL:** When redirecting issues to a **non-PT2** oncall queue, apply exactly one `oncall: ...` label and **STOP**. Do NOT:
 - Add any `module:` labels
 - Mark it `triaged`
 - Do any further triage work
@@ -158,14 +179,40 @@ The sub-oncall team will handle their own triage. Your job is only to route it t
 | `oncall: profiler` | Profiler issues (CPU, GPU, Kineto) |
 | `oncall: visualization` | TensorBoard integration |
 
+**Common routing mistakes to avoid:**
+- **MPS ≠ Mobile.** MPS (Metal Performance Shaders) is the macOS/Apple Silicon GPU backend. Do NOT route MPS issues to `oncall: mobile`. MPS issues stay in the general queue with `module: mps`.
+- **DTensor → `oncall: distributed`.** DTensor issues should always be routed to `oncall: distributed`, even if they don't mention DDP/FSDP.
+- **ONNX → `module: onnx`.** There is no `oncall: onnx`. Use `module: onnx` and keep in the general queue.
+- **CI/releng → `module: ci`.** Do not use `oncall: releng`. Use `module: ci` for CI infrastructure issues.
+- **torch.compile + distributed.** When `torch.compile` mishandles a distributed op (e.g., `dist.all_reduce`), the issue typically needs BOTH `oncall: pt2` and `oncall: distributed` since the fix may span both codebases.
+
 **Note:** `oncall: cpu inductor` is a sub-queue of PT2. For general triage, just use `oncall: pt2`.
 
 ### 4) Label the issue (if NOT transferred/redirected)
 
 Only if the issue stays in the general queue:
 - Add 1+ `module: ...` labels based on the affected area
-- If feature request: add `feature` (or `function request` for a new function or new arguments/modes)
-- If small improvement: add `enhancement`
+- Prefer specific labels over general ones when both exist. Check `labels.json` descriptions for guidance on when a specific label supersedes a general one (e.g., `module: sdpa` instead of `module: nn` for SDPA issues, `module: flex attention` instead of `module: nn` for flex attention).
+- `feature` — wholly new functionality that does not exist today in any form
+- `enhancement` — improvement to something that already works (e.g., adding a native backend kernel for an op that already runs via fallback/composite, performance optimization, better error messages). If the enhancement is about performance, also add `module: performance`.
+- `function request` — a new function or new arguments/modes for an existing function
+- If the issue says the operation "currently works" or "falls back to" a slower path, that is `enhancement`, not `feature`
+
+**Commonly missed labels — always check for these:**
+
+| Condition | Label |
+|-----------|-------|
+| Segfault, illegal memory access, SIGSEGV | `module: crash` |
+| Performance issue: regression, slowdown, or optimization request | `module: performance` |
+| Issue on Windows | `module: windows` |
+| Previously working feature now broken | `module: regression` |
+| Broken docs/links that previously worked | `module: docs` + `module: regression` (NOT `enhancement`) |
+| Issue about a test failing (not the underlying functionality) | `module: tests` |
+| Backward pass / gradient computation bug | `module: autograd` (in addition to the op's module label) |
+| `torch.linalg` ops or linear algebra ops (solve, svd, eig, inv, etc.) | `module: linear algebra` |
+| `has workaround` | Only add when the workaround is **non-trivial and non-obvious**. If the issue is "X doesn't work for non-contiguous tensors," calling `.contiguous()` is the tautological inverse of the bug, not a workaround. A real workaround is something like installing a specific package version, adding a synchronization point, inserting `gc.collect()`, or using a different API that isn't obviously implied by the bug description. |
+
+**Label based on the actual bug, not keywords.** Read the issue to understand what is actually broken. A bug about broadcasting that happens to mention "nan" in a parameter name is a frontend bug, not a NaN/Inf bug.
 
 ### 5) High Priority — REQUIRES HUMAN REVIEW
 
diff --git a/.claude/skills/triaging-issues/labels.json b/.claude/skills/triaging-issues/labels.json
index e99651488212e..fa3703186472b 100644
--- a/.claude/skills/triaging-issues/labels.json
+++ b/.claude/skills/triaging-issues/labels.json
@@ -21,11 +21,11 @@
     },
     {
       "name": "enhancement",
-      "description": "Not as big of a feature, but technically not a bug. Should be easy to fix"
+      "description": "Improvement to something that already works: performance optimization, better error messages, adding a backend-native implementation for an op that already runs via fallback/composite"
     },
     {
       "name": "feature",
-      "description": "A request for a proper, new feature."
+      "description": "A request for wholly new functionality that does not exist today in any form (no fallback, no composite, no workaround)"
     },
     {
       "name": "function request",
@@ -105,7 +105,7 @@
     },
     {
       "name": "module: aotdispatch",
-      "description": "umbrella label for AOTAutograd issues"
+      "description": "Issues in the AOTAutograd subsystem (torch/_functorch/aot_autograd.py, torch/_functorch/_aot_autograd/). Covers tracing of the forward+backward graph, partitioning, runtime wrappers, and autograd cache serialization. If a traceback includes _aot_autograd/runtime_wrappers.py or _aot_autograd/functional_utils.py, this label applies. Prefer this over 'module: inductor' when the crash is in the autograd dispatch layer, not in code generation."
     },
     {
       "name": "module: aotinductor",
@@ -117,7 +117,7 @@
     },
     {
       "name": "module: autograd",
-      "description": "Related to torch.autograd, and the autograd engine in general"
+      "description": "Issues in the eager-mode autograd engine (torch/autograd/, torch/csrc/autograd/). Covers gradient computation, autograd functions, hooks, and the backward pass in eager mode. Do NOT apply when the issue is about torch.compile silently dropping or mishandling an autograd operation — if the bug only reproduces under torch.compile (not in eager), the root cause is likely in Dynamo tracing ('module: dynamo') or AOTAutograd ('module: aotdispatch'), not in the eager autograd engine."
     },
     {
       "name": "module: backend",
@@ -141,7 +141,7 @@
     },
     {
       "name": "module: binaries",
-      "description": "Anything related to official binaries that we release to users"
+      "description": "Packaging, wheel, and install issues: undefined symbols at import time, version mismatches between pip-installed packages, missing shared libraries (libtorch_cuda.so, libnccl, libcudnn), and broken official binaries"
     },
     {
       "name": "module: bottleneck",
@@ -177,7 +177,7 @@
     },
     {
       "name": "module: compiled autograd",
-      "description": "compiled_autograd"
+      "description": "Issues in the compiled autograd subsystem (torch/csrc/autograd/compiled_autograd.*, torch/_dynamo/compiled_autograd.py). This is the system that compiles the autograd backward graph with Dynamo+Inductor. Do NOT apply just because 'autograd' appears in the issue — use 'module: aotdispatch' for AOTAutograd tracing/partitioning issues, and 'module: autograd' for eager autograd issues."
     },
     {
       "name": "module: complex",
@@ -269,7 +269,7 @@
     },
     {
       "name": "module: decompositions",
-      "description": "Topics related to decomposition (excluding PrimTorch)"
+      "description": "Issues with operator decompositions (torch/_decomp/). Decompositions rewrite higher-level ops into lower-level ones for tracing/compilation. Apply when traced/compiled results diverge from eager due to a decomposition producing wrong gradients, numerically incorrect results, or missing higher-order derivative support. Key signal: 'eager vs traced results diverge' + a specific op (e.g., sqrt, layer_norm) suggests a bad decomposition."
     },
     {
       "name": "module: dependency bug",
@@ -289,7 +289,7 @@
     },
     {
       "name": "module: determinism",
-      "description": ""
+      "description": "Issues related to torch.use_deterministic_algorithms(), deterministic mode, and reproducibility of results"
     },
     {
       "name": "module: DeviceMesh",
@@ -345,7 +345,7 @@
     },
     {
       "name": "module: dynamo",
-      "description": "Related to TorchDynamo (torch.compile frontend/tracing)"
+      "description": "Issues in TorchDynamo, the torch.compile frontend/tracing system (torch/_dynamo/). Covers graph capture, graph breaks, symbolic tracing of Python bytecode, and variable tracking. Key signal: if torch.compile silently drops, ignores, or mishandles an operation that works correctly in eager mode (e.g., in-place mutations like detach_() being silently skipped, side-effects not captured), the root cause is almost always Dynamo's tracing."
     },
     {
       "name": "module: edge cases",
@@ -365,7 +365,7 @@
     },
     {
       "name": "module: error checking",
-      "description": "Bugs related to incorrect/lacking error checking"
+      "description": "Bugs where PyTorch fails to validate inputs or raises a wrong/missing error message at a validation boundary (e.g., shape checks, dtype checks, device checks). Do NOT apply just because the user complains about a bad error message — if the root cause is a bug in a specific feature, use the feature-specific label instead."
     },
     {
       "name": "module: expecttest",
@@ -385,7 +385,7 @@
     },
     {
       "name": "module: flex attention",
-      "description": ""
+      "description": "torch.nn.attention.flex_attention API. Use instead of module: nn for flex attention issues"
     },
     {
       "name": "module: floatx (formerly float8)",
@@ -417,7 +417,7 @@
     },
     {
       "name": "module: fx",
-      "description": ""
+      "description": "Issues in torch.fx (torch/fx/) — the graph IR, symbolic tracing, graph transforms, and make_fx. Note: make_fx with symbolic tracing is PT2 infrastructure, so issues involving make_fx should also get 'oncall: pt2'."
     },
     {
       "name": "module: fx.passes",
@@ -437,7 +437,7 @@
     },
     {
       "name": "module: helion",
-      "description": ""
+      "description": "Helion (pytorch/helion) is a high-level Python DSL for writing GPU kernels that compiles to Triton. Users write kernels with @helion.kernel decorator and helion.language (hl.tile(), etc.) using familiar PyTorch operations. Helion kernels integrate with torch.compile/inductor via template fusion hooks. Apply when the issue involves helion kernel authoring, helion compilation, helion-to-Triton codegen, or helion+inductor template fusion. Do NOT apply for plain Triton kernel issues (use module: inductor) or general torch.compile issues unless helion is specifically involved."
     },
     {
       "name": "module: higher order operators",
@@ -453,7 +453,7 @@
     },
     {
       "name": "module: inductor",
-      "description": "Related to TorchInductor (torch.compile codegen backend)"
+      "description": "Issues in the TorchInductor code generation backend (torch/_inductor/). Covers kernel codegen, scheduling, fusion, and inductor-specific caching (TORCHINDUCTOR_* env vars). Do NOT apply just because the user mentions 'inductor cache' or 'torch.compile' — if the traceback points to _aot_autograd/ or _functorch/, use 'module: aotdispatch' instead."
     },
     {
       "name": "module: infallible views",
@@ -769,7 +769,7 @@
     },
     {
       "name": "module: pt2-dispatcher",
-      "description": "PT2  dispatcher-related issues (e.g., aotdispatch, functionalization, faketensor, custom-op,"
+      "description": "Umbrella label for PT2 dispatcher-layer issues: aotdispatch, functionalization, FakeTensor, custom-op registration, proxy tensor tracing. Apply alongside the more specific label (e.g., 'module: aotdispatch') when the issue touches the dispatcher layer. If unsure whether a PT2 issue is inductor vs dispatcher, check the traceback: _functorch/ and _aot_autograd/ paths mean dispatcher, _inductor/ paths mean inductor."
     },
     {
       "name": "module: pybind",
@@ -789,7 +789,7 @@
     },
     {
       "name": "module: python version",
-      "description": "Issues related to specific Python versions"
+      "description": "Issues where behavior differs across Python versions or where a specific Python version causes breakage (e.g., 'works on 3.11 but fails on 3.12'). Do NOT apply just because the reporter mentions a Python version in their environment info, test path, or conformance test name — the bug must be specifically about Python version compatibility."
     },
     {
       "name": "module: pytree",
@@ -849,7 +849,7 @@
     },
     {
       "name": "module: sdpa",
-      "description": "All things related to torch.nn.functional.scaled_dot_product_attentiion"
+      "description": "scaled_dot_product_attention and its backends (math, flash, mem_efficient, cudnn). Use instead of module: nn for SDPA issues"
     },
     {
       "name": "module: selective build",
diff --git a/.claude/skills/triaging-issues/pt2-triage-rubric.md b/.claude/skills/triaging-issues/pt2-triage-rubric.md
index 883f40fdccc3c..967725380f7be 100644
--- a/.claude/skills/triaging-issues/pt2-triage-rubric.md
+++ b/.claude/skills/triaging-issues/pt2-triage-rubric.md
@@ -2,6 +2,8 @@
 
 This rubric guides labeling decisions for PT2 oncall triage.
 
+**Every `oncall: pt2` issue MUST have at least one `module:` label.** The PT2 queue is too broad without one — the team needs to know which component is affected. Use the sections below to determine the right module label(s).
+
 ## 1. Component Isolation - Be Precise, Don't Over-Tag
 
 ### Dynamo vs Dynamic Shapes
@@ -31,6 +33,31 @@ When component isn't clear from the issue body:
 
 ---
 
+**This is critical** when you have identified an issue as inductor, and the failing device is "cpu" ONLY, then this is a CPU inductor issue, and should be redirected to `oncall: cpu inductor`
+
+### Silently Dropped Operations = Dynamo
+
+If torch.compile silently drops or ignores an operation that works in eager, the bug is in Dynamo's tracing.
+
+| Signal | Label |
+|--------|-------|
+| In-place mutation skipped under compile (`detach_()`, `requires_grad_()`) | `module: dynamo` |
+| Side-effect not captured (global state, tensor metadata) | `module: dynamo` |
+
+Don't apply `module: autograd` just because the dropped operation involves autograd. If eager works fine, the autograd engine is fine.
+
+### Decomposition Bugs
+
+If eager and traced results diverge numerically for a specific op, suspect a bad decomposition.
+
+| Signal | Label |
+|--------|-------|
+| Eager vs traced diverge for a specific op | `module: decompositions` |
+| Higher-order gradients wrong under tracing | `module: decompositions` |
+| `make_fx` symbolic tracing diverges from eager | `module: fx` + `module: decompositions` + `oncall: pt2` |
+
+---
+
 ## 3. Don't Over-Tag pt2-dispatcher
 
 `module: pt2-dispatcher` is for bugs **IN** the dispatcher code, not just when it appears in a stack trace.
@@ -65,6 +92,8 @@ When component isn't clear from the issue body:
 | DTensor has a bad error message under compile | Bug is in PT2's error handling → PT2 owns UX |
 | Distributed training fails, but stack trace shows inductor issue | Bug is in inductor → PT2 owns it |
 
+For PT2-D issues, you may also add `oncall: distributed` but DO NOT hand this off fully - keep the `oncall: pt2` label.
+
 **Examples - DO redirect:**
 
 | Situation | Why redirect |
@@ -77,6 +106,8 @@ When component isn't clear from the issue body:
 
 **Adding labels for visibility:** You CAN add domain labels (e.g., `module: dtensor`) so domain experts see the issue, but don't ADD the oncall redirect label unless you're actually handing it off.
 
+**This is critical** when you have identified an issue as inductor, and the failing device is "cpu" ONLY, then this is a CPU inductor issue, and should be redirected to `oncall: cpu inductor`
+
 ---
 
 ## 5. Add Domain-Specific Labels for Visibility
@@ -106,7 +137,34 @@ Check for existing labels before inventing categories:
 
 ---
 
-## 7. functorch + compile
+## 7. Helion Kernel Issues
+
+[Helion](https://github.com/pytorch/helion) is a high-level Python DSL for writing GPU kernels. Users write kernels with `@helion.kernel` and `helion.language` (`hl.tile()`, etc.) using standard PyTorch ops, and Helion compiles them to Triton. Helion kernels can also be fused into `torch.compile` graphs via inductor's template fusion hooks.
+
+### Identifying Helion Issues
+
+| Signal | Label |
+|--------|-------|
+| Issue mentions `helion`, `@helion.kernel`, `helion.language`, `hl.tile` | `module: helion` |
+| Error traceback includes `helion/` or `helion.` frames | `module: helion` |
+| Helion kernel produces wrong results (standalone, no torch.compile) | `module: helion` only |
+| Helion kernel fails or miscompiles under `torch.compile` | `module: helion` + `module: inductor` |
+| Inductor template fusion bug triggered by a Helion template | `module: helion` + `module: inductor` |
+
+### Routing
+
+- Helion issues get `oncall: pt2` — Helion is a PT2 component.
+- If the bug is purely in Helion's own compilation (standalone kernel, not under torch.compile), apply `module: helion` without `module: inductor`.
+- If the bug is in how inductor fuses or emits Helion templates, apply both `module: helion` and `module: inductor`.
+
+### Common Mistakes
+
+- **Don't confuse Helion with raw Triton**: If the user is writing Triton kernels directly (using `@triton.jit`, `tl.load`, etc.) without Helion, that's `module: inductor`, not `module: helion`.
+- **Don't apply `module: helion` for general inductor codegen bugs**: Just because inductor generates Triton code doesn't make it a Helion issue. Helion is a specific DSL — look for explicit `helion` imports or mentions.
+
+---
+
+## 8. Functorch + Compile
 
 | Situation | Labels |
 |-----------|--------|
@@ -115,9 +173,12 @@ Check for existing labels before inventing categories:
 
 ---
 
-## 8. High Priority Criteria
+## 9. High Priority Criteria
+
+**This is critical** You should not explicitly add `high priority` - add `triage review` instead
+so that it is reviewed at the next triage meeting by the oncall.
 
-Mark `high priority` if ANY of these apply:
+Mark `triage review` if ANY of these apply:
 
 | Criteria | Example |
 |----------|---------|
@@ -128,11 +189,9 @@ Mark `high priority` if ANY of these apply:
 | **Important model regressed** (>10% perf) | Common model, significant slowdown |
 | **Important customer** | Huggingface, common usage patterns |
 
-Note: Adding `high priority` auto-adds `triage review` for discussion.
-
 ---
 
-## 9. Fuzzer Issues
+## 10. Fuzzer Issues
 
 For `topic: fuzzer` issues:
 
@@ -144,7 +203,7 @@ For `topic: fuzzer` issues:
 
 ---
 
-## 10. Quick Label Reference
+## 11. Quick Label Reference
 
 ### Core Components
 - `module: dynamo` - Tracing, bytecode, graph breaks
@@ -153,6 +212,7 @@ For `topic: fuzzer` issues:
 - `module: pt2-dispatcher` - AOT autograd, functionalization, FakeTensor
 - `module: cuda graphs` - CUDA graph capture/replay
 - `module: flex attention` - Flex attention API
+- `module: helion` - Helion DSL kernel authoring, compilation, and inductor fusion
 
 ### Holistic Areas
 - `module: compile ux` - Error messages, APIs, programming model
@@ -171,3 +231,13 @@ For `topic: fuzzer` issues:
 - `oncall: cpu inductor` - CPU/MKLDNN issues
 - `oncall: export` - Export-specific issues
 - `oncall: distributed` - Distributed training issues
+
+### CPU Inductor Routing
+
+Route to `oncall: cpu inductor` (not generic `oncall: pt2`) when the issue is specific to CPU backend in inductor:
+- Title or body mentions `[CPU]`, `cpu`, or `MKLDNN`
+- CPU-specific codegen bugs (e.g., float16 handling on CPU)
+- Issues that only reproduce on CPU, not CUDA
+- MKLDNN-specific kernel issues
+
+Example: "[Inductor][CPU][float16] LayerNorm outputs NaN" → `oncall: cpu inductor`, NOT `oncall: pt2`
diff --git a/.claude/skills/triaging-issues/scripts/validate_labels.py b/.claude/skills/triaging-issues/scripts/validate_labels.py
index b8bb93dff0567..df489189a20b5 100755
--- a/.claude/skills/triaging-issues/scripts/validate_labels.py
+++ b/.claude/skills/triaging-issues/scripts/validate_labels.py
@@ -39,6 +39,7 @@
 
 FORBIDDEN_EXACT = [
     "merge blocking",
+    "oncall: releng",  # Not a triage redirect target; use module: ci instead
 ]
 
 REDUNDANT_PAIRS = [
@@ -207,12 +208,18 @@ def main():
     except json.JSONDecodeError as e:
         debug_log(f"JSON decode error: {e}")
         print(f"Hook error: Invalid JSON input: {e}", file=sys.stderr)
-        print("Hook was unable to validate labels; stopping triage.", file=sys.stderr)
+        print(
+            "Hook was unable to validate labels; stopping triage.",
+            file=sys.stderr,
+        )
         sys.exit(2)
     except Exception as e:
         debug_log(f"Unexpected error: {type(e).__name__}: {e}")
         print(f"Hook error: {e}", file=sys.stderr)
-        print("Hook was unable to validate labels; stopping triage.", file=sys.stderr)
+        print(
+            "Hook was unable to validate labels; stopping triage.",
+            file=sys.stderr,
+        )
         sys.exit(2)
 
 
diff --git a/.flake8 b/.flake8
index c340ef9903870..f5507047d7f39 100644
--- a/.flake8
+++ b/.flake8
@@ -12,13 +12,21 @@ ignore =
     # to line this up with executable bit
     EXE001,
     # these ignores are from flake8-bugbear; please fix!
-    B007,B008,B017,B019,B023,B028,B903,B905,B906,B907,B908,B910
+    B007,B008,B017,B019,B023,B028,B903,B905,B906,B907,B908,B910,
     # these ignores are from flake8-simplify. please fix or ignore with commented reason
     SIM105,SIM108,SIM110,SIM111,SIM113,SIM114,SIM115,SIM116,SIM117,SIM118,SIM119,SIM12,
     # SIM104 is already covered by pyupgrade ruff
     SIM104,
     # flake8-simplify code styles
-    SIM102,SIM103,SIM106,SIM112
+    SIM102,SIM103,SIM106,SIM112,
+    # Codes where ruff is the source of truth. Ruff handles these (sometimes
+    # differently than flake8). As ruff promotes preview codes, move them here
+    # from ruff's external list.
+    B001,B036,B902,B904,B950,
+    C406,C417,C419,C901,
+    E121,E122,E124,E128,E131,E712,E722,E723,E731,
+    F401,F722,F723,F811,F812,
+    G001,G003,G004,G010,G200,G201,G202
 per-file-ignores =
     __init__.py: F401
     test/**: F821
@@ -43,6 +51,7 @@ exclude =
     ./functorch/docs,
     ./functorch/examples,
     ./functorch/docs/source/tutorials,
+    ./torch/testing/_internal/py312_intrinsics.py
     ./scripts,
     ./test/generated_type_hints_smoketest.py,
     ./third_party,
diff --git a/.gitattributes b/.gitattributes
index f20bd8f1b31b4..b58735e17fe55 100644
--- a/.gitattributes
+++ b/.gitattributes
@@ -1,5 +1,4 @@
 *.bat text eol=crlf
-.circleci/config.yml linguist-generated=true
 .github/workflows/generated-*.yml linguist-generated=true
 .github/generated-* linguist-generated=true
 .github/scripts/gql_mocks.json linguist-generated=true
diff --git a/.github/ISSUE_TEMPLATE/bug-report.yml b/.github/ISSUE_TEMPLATE/bug-report.yml
index e94605e1ef065..bd0012e73ab7f 100644
--- a/.github/ISSUE_TEMPLATE/bug-report.yml
+++ b/.github/ISSUE_TEMPLATE/bug-report.yml
@@ -51,7 +51,11 @@ body:
     description: |
       Please run the following and paste the output below.
       ```sh
-      wget https://raw.githubusercontent.com/pytorch/pytorch/main/torch/utils/collect_env.py
+      curl -sL https://raw.githubusercontent.com/pytorch/pytorch/main/torch/utils/collect_env.py | python
+      ```
+      Or, if you prefer to inspect the script first:
+      ```sh
+      curl -OL https://raw.githubusercontent.com/pytorch/pytorch/main/torch/utils/collect_env.py
       # For security purposes, please check the contents of collect_env.py before running it.
       python collect_env.py
       ```
diff --git a/.github/ISSUE_TEMPLATE/pt2-bug-report.yml b/.github/ISSUE_TEMPLATE/pt2-bug-report.yml
index 51255fde13cca..c8080bb70e3d4 100644
--- a/.github/ISSUE_TEMPLATE/pt2-bug-report.yml
+++ b/.github/ISSUE_TEMPLATE/pt2-bug-report.yml
@@ -64,6 +64,10 @@ body:
       description: |
         Please run the following and paste the output below.
         ```sh
+        curl -sL https://raw.githubusercontent.com/pytorch/pytorch/main/torch/utils/collect_env.py | python3
+        ```
+        Or, if you prefer to inspect the script first:
+        ```sh
         curl -OL https://raw.githubusercontent.com/pytorch/pytorch/main/torch/utils/collect_env.py
         # For security purposes, please check the contents of collect_env.py before running it.
         python3 collect_env.py
diff --git a/.github/actionlint.yaml b/.github/actionlint.yaml
index 95637501e1069..b894d7214fced 100644
--- a/.github/actionlint.yaml
+++ b/.github/actionlint.yaml
@@ -25,9 +25,6 @@ self-hosted-runner:
     - linux.aws.h100
     - linux.aws.h100.4
     - linux.aws.h100.8
-    - linux.4xlarge.nvidia.gpu
-    - linux.8xlarge.nvidia.gpu
-    - linux.16xlarge.nvidia.gpu
     - linux.g5.4xlarge.nvidia.gpu
     - linux.c7i.2xlarge
     # Pytorch/pytorch AWS Linux Runners on Linux Foundation account
@@ -37,9 +34,6 @@ self-hosted-runner:
     - lf.linux.12xlarge
     - lf.linux.24xlarge
     - lf.linux.arm64.2xlarge
-    - lf.linux.4xlarge.nvidia.gpu
-    - lf.linux.8xlarge.nvidia.gpu
-    - lf.linux.16xlarge.nvidia.gpu
     - lf.linux.g5.4xlarge.nvidia.gpu
     - lf.linux.c7i.2xlarge
     # Repo-specific IBM hosted S390x runner
@@ -63,6 +57,7 @@ self-hosted-runner:
     - linux.rocm.gpu.mi250.4
     - linux.rocm.gpu.2
     - linux.rocm.gpu.4
+    - linux.rocm.mi210.docker-cache
     - linux.rocm.mi250.docker-cache
     # gfx942 runners
     - linux.rocm.gpu.gfx942.1
@@ -73,10 +68,13 @@ self-hosted-runner:
     # Org wise AWS `mac2.metal` runners (2020 Mac mini hardware powered by Apple silicon M1 processors)
     - macos-m1-stable
     - macos-m1-14
+    - macos-m2-15
+    - macos-m2-26
     # GitHub-hosted MacOS runners
     - macos-latest-xlarge
     - macos-13-xlarge
     - macos-14-xlarge
+    - macos-26-xlarge
     # Organization-wide Intel hosted XPU runners
     - linux.idc.xpu
     # Organization-wide Google Cloud TPU runners
diff --git a/.github/actions/checkout-pytorch/action.yml b/.github/actions/checkout-pytorch/action.yml
index 3ea7d295fb006..97303742c499c 100644
--- a/.github/actions/checkout-pytorch/action.yml
+++ b/.github/actions/checkout-pytorch/action.yml
@@ -60,6 +60,7 @@ runs:
         # --depth=1 for speed, manually fetch history and other refs as necessary
         fetch-depth: ${{ inputs.fetch-depth }}
         single-branch: true
+        fetch-tags: ${{ github.ref_type == 'tag' }}
         submodules: ${{ inputs.submodules }}
         show-progress: false
         filter: ${{ env.filter }}
@@ -107,6 +108,7 @@ runs:
         ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
         fetch-depth: ${{ inputs.fetch-depth }}
         single-branch: true
+        fetch-tags: ${{ github.ref_type == 'tag' }}
         submodules: ${{ inputs.submodules }}
         show-progress: false
         filter: ${{ env.filter }}
diff --git a/.github/actions/download-build-artifacts/action.yml b/.github/actions/download-build-artifacts/action.yml
index c44b6a4083448..ea78440e87686 100644
--- a/.github/actions/download-build-artifacts/action.yml
+++ b/.github/actions/download-build-artifacts/action.yml
@@ -18,14 +18,14 @@ runs:
   using: composite
   steps:
     - name: Download PyTorch Build Artifacts from S3
-      if: ${{ !inputs.use-gha }}
+      if: inputs.use-gha == '' || inputs.use-gha == 'false'
       uses: seemethere/download-artifact-s3@v4
       with:
         name: ${{ inputs.name }}
         s3-bucket: ${{ inputs.s3-bucket }}
 
     - name: Download PyTorch Build Artifacts from GHA
-      if: ${{ inputs.use-gha }}
+      if: inputs.use-gha != '' && inputs.use-gha != 'false'
       uses: actions/download-artifact@v4
       with:
         name: ${{ inputs.name }}
diff --git a/.github/actions/download-td-artifacts/action.yml b/.github/actions/download-td-artifacts/action.yml
index 18766bf670f63..3b25f14ef6520 100644
--- a/.github/actions/download-td-artifacts/action.yml
+++ b/.github/actions/download-td-artifacts/action.yml
@@ -11,13 +11,13 @@ runs:
   using: composite
   steps:
     - name: Download TD Artifacts from S3
-      if: ${{ !inputs.use-gha }}
+      if: inputs.use-gha == '' || inputs.use-gha == 'false'
       uses: seemethere/download-artifact-s3@v4
       with:
         name: td_results
 
     - name: Download TD Artifacts from GHA
-      if: inputs.use-gha
+      if: inputs.use-gha != '' && inputs.use-gha != 'false'
       uses: actions/download-artifact@v4
       with:
         name: td_results.json
diff --git a/.github/actions/filter-test-configs/action.yml b/.github/actions/filter-test-configs/action.yml
index 8afbc40cebeaa..36ccd839f26be 100644
--- a/.github/actions/filter-test-configs/action.yml
+++ b/.github/actions/filter-test-configs/action.yml
@@ -61,6 +61,7 @@ runs:
   using: composite
   steps:
     - name: Setup uv
+      if: ${{ !env.UV_PYTHON }}
       uses: pytorch/test-infra/.github/actions/setup-uv@main
       with:
         python-version: "3.12"
diff --git a/.github/actions/linux-test/action.yml b/.github/actions/linux-test/action.yml
deleted file mode 100644
index 8ad74e308168a..0000000000000
--- a/.github/actions/linux-test/action.yml
+++ /dev/null
@@ -1,406 +0,0 @@
-name: linux-test
-
-inputs:
-  build-environment:
-    required: true
-    type: string
-    description: Top-level label for what's being built/tested.
-  test-matrix:
-    required: true
-    type: string
-    description: JSON description of what test configs to run.
-  docker-image:
-    required: true
-    type: string
-    description: Docker image to run in.
-  sync-tag:
-    required: false
-    type: string
-    default: ""
-    description: |
-      If this is set, our linter will use this to make sure that every other
-      job with the same `sync-tag` is identical.
-  use-gha:
-    required: false
-    type: string
-    default: ""
-    description: If set to any value, upload to GHA. Otherwise upload to S3.
-  dashboard-tag:
-    required: false
-    type: string
-    default: ""
-  s3-bucket:
-    description: S3 bucket to download artifact
-    required: false
-    type: string
-    default: "gha-artifacts"
-  aws-role-to-assume:
-    description: role to assume for downloading artifacts
-    required: false
-    type: string
-    default: ""
-  HUGGING_FACE_HUB_TOKEN:
-    description: |
-      HF Auth token to avoid rate limits when downloading models or datasets from hub
-    required: false
-    default: ""
-  GITHUB_TOKEN:
-    description: GitHub token
-    required: true
-  disable-monitor:
-    description: |
-      [Experimental] Disable utilization monitoring for tests.
-      Currently, by default we disable the monitor job and only look for specific tests,
-      since we are investigating the behaviour of the monitor script with different tests.
-    required: false
-    type: boolean
-    default: true
-#env:
-#  GIT_DEFAULT_BRANCH: ${{ inputs.default_branch }}
-
-runs:
-  using: composite
-  steps:
-    - name: Setup Linux
-      uses: ./.github/actions/setup-linux
-
-    - name: Login to ECR
-      uses: ./.github/actions/ecr-login
-      with:
-        aws-role-to-assume: ${{ inputs.aws-role-to-assume }}
-
-    - name: Calculate docker image
-      id: calculate-docker-image
-      uses: pytorch/test-infra/.github/actions/calculate-docker-image@main
-      with:
-        docker-image-name: ${{ inputs.docker-image }}
-
-    - name: Use following to pull public copy of the image
-      id: print-ghcr-mirror
-      env:
-        ECR_DOCKER_IMAGE: ${{ steps.calculate-docker-image.outputs.docker-image }}
-      shell: bash
-      run: |
-        tag=${ECR_DOCKER_IMAGE##*/}
-        echo "docker pull ghcr.io/pytorch/ci-image:${tag/:/-}"
-
-    - name: Pull docker image
-      uses: pytorch/test-infra/.github/actions/pull-docker-image@main
-      with:
-        docker-image: ${{ steps.calculate-docker-image.outputs.docker-image }}
-
-    - name: Check if in a container runner
-      shell: bash
-      id: check_container_runner
-      run: echo "IN_CONTAINER_RUNNER=$(if [ -f /.inarc ] || [ -f /.incontainer ]; then echo true ; else echo false; fi)" >> "$GITHUB_OUTPUT"
-
-    - name: Install nvidia driver, nvidia-docker runtime, set GPU_FLAG
-      id: install-nvidia-driver
-      uses: pytorch/test-infra/.github/actions/setup-nvidia@main
-      if: ${{ contains(inputs.build-environment, 'cuda') && !contains(matrix.config, 'nogpu') && steps.check_container_runner.outputs.IN_CONTAINER_RUNNER == 'false' }}
-
-    - name: Setup GPU_FLAG for docker run
-      id: setup-gpu-flag
-      run: echo "GPU_FLAG=--gpus all -e NVIDIA_DRIVER_CAPABILITIES=all" >> "${GITHUB_ENV}"
-      if: ${{ contains(inputs.build-environment, 'cuda') && !contains(matrix.config, 'nogpu') && steps.check_container_runner.outputs.IN_CONTAINER_RUNNER == 'true' }}
-
-    - name: Setup SCCACHE_SERVER_PORT environment for docker run when on container
-      id: setup-sscache-port-flag
-      run: echo "SCCACHE_SERVER_PORT_DOCKER_FLAG=-e SCCACHE_SERVER_PORT=$((RUNNER_UID + 4226))" >> "${GITHUB_ENV}"
-      if: ${{ steps.check_container_runner.outputs.IN_CONTAINER_RUNNER == 'true' }}
-
-    - name: Lock NVIDIA A100 40GB Frequency
-      shell: bash
-      run: |
-        sudo nvidia-smi -pm 1
-        sudo nvidia-smi -ac 1215,1410
-        nvidia-smi
-      if: ${{ contains(matrix.runner, 'a100') && steps.check_container_runner.outputs.IN_CONTAINER_RUNNER == 'false' }}
-
-    - name: Start monitoring script
-      id: monitor-script
-      if: ${{ !inputs.disable-monitor }}
-      shell: bash
-      continue-on-error: true
-      run: |
-        python3 -m pip install psutil==5.9.8 nvidia-ml-py==11.525.84
-        python3 -m tools.stats.monitor > usage_log.txt 2>&1 &
-        echo "monitor-script-pid=${!}" >> "${GITHUB_OUTPUT}"
-
-    - name: Download build artifacts
-      uses: ./.github/actions/download-build-artifacts
-      with:
-        name: ${{ inputs.build-environment }}
-        s3-bucket: ${{ inputs.s3-bucket }}
-
-    - name: Download TD artifacts
-      continue-on-error: true
-      uses: ./.github/actions/download-td-artifacts
-
-    - name: Parse ref
-      id: parse-ref
-      shell: bash
-      run: .github/scripts/parse_ref.py
-
-    - name: Get workflow job id
-      id: get-job-id
-      uses: ./.github/actions/get-workflow-job-id
-      if: always()
-      with:
-        github-token: ${{ inputs.GITHUB_TOKEN }}
-
-    - name: Check for keep-going label and re-enabled test issues
-      # This uses the filter-test-configs action because it conveniently
-      # checks for labels and re-enabled test issues.  It does not actually do
-      # any filtering.  All filtering is done in the build step.
-      id: keep-going
-      uses: ./.github/actions/filter-test-configs
-      with:
-        github-token: ${{ inputs.GITHUB_TOKEN }}
-        test-matrix: ${{ inputs.test-matrix }}
-        job-name: ${{ steps.get-job-id.outputs.job-name }}
-
-    - name: Test
-      id: test
-      env:
-        BUILD_ENVIRONMENT: ${{ inputs.build-environment }}
-        PR_NUMBER: ${{ github.event.pull_request.number }}
-        GITHUB_REPOSITORY: ${{ github.repository }}
-        GITHUB_WORKFLOW: ${{ github.workflow }}
-        GITHUB_JOB: ${{ github.job }}
-        GITHUB_RUN_ID: ${{ github.run_id }}
-        GITHUB_RUN_NUMBER: ${{ github.run_number }}
-        GITHUB_RUN_ATTEMPT: ${{ github.run_attempt }}
-        JOB_ID: ${{ steps.get-job-id.outputs.job-id }}
-        JOB_NAME: ${{ steps.get-job-id.outputs.job-name }}
-        BRANCH: ${{ steps.parse-ref.outputs.branch }}
-        SHA1: ${{ github.event.pull_request.head.sha || github.sha }}
-        BASE_SHA: ${{ github.event.pull_request.base.sha || github.sha }}
-        TEST_CONFIG: ${{ matrix.config }}
-        SHARD_NUMBER: ${{ matrix.shard }}
-        NUM_TEST_SHARDS: ${{ matrix.num_shards }}
-        REENABLED_ISSUES: ${{ steps.keep-going.outputs.reenabled-issues }}
-        CONTINUE_THROUGH_ERROR: ${{ steps.keep-going.outputs.keep-going }}
-        VERBOSE_TEST_LOGS: ${{ steps.keep-going.outputs.ci-verbose-test-logs }}
-        TEST_SHOWLOCALS: ${{ steps.keep-going.outputs.ci-test-showlocals }}
-        NO_TEST_TIMEOUT: ${{ steps.keep-going.outputs.ci-no-test-timeout }}
-        NO_TD: ${{ steps.keep-going.outputs.ci-no-td }}
-        TD_DISTRIBUTED: ${{ steps.keep-going.outputs.ci-td-distributed }}
-        SCCACHE_BUCKET: ossci-compiler-cache-circleci-v2
-        SCCACHE_REGION: us-east-1
-        SCCACHE_S3_KEY_PREFIX: ${{ github.workflow }}
-        SHM_SIZE: ${{ contains(inputs.build-environment, 'cuda') && '2g' || '1g' }}
-        DOCKER_IMAGE: ${{ inputs.docker-image }}
-        XLA_CUDA: ${{ contains(inputs.build-environment, 'xla') && '0' || '' }}
-        XLA_CLANG_CACHE_S3_BUCKET_NAME: ossci-compiler-clang-cache-circleci-xla
-        PYTORCH_TEST_CUDA_MEM_LEAK_CHECK: ${{ matrix.mem_leak_check && '1' || '0' }}
-        PYTORCH_TEST_RERUN_DISABLED_TESTS: ${{ matrix.rerun_disabled_tests && '1' || '0' }}
-        DASHBOARD_TAG: ${{ inputs.dashboard-tag }}
-        HUGGING_FACE_HUB_TOKEN: ${{ inputs.HUGGING_FACE_HUB_TOKEN }}
-        SCRIBE_GRAPHQL_ACCESS_TOKEN: ${{ secrets.SCRIBE_GRAPHQL_ACCESS_TOKEN }}
-        IS_A100_RUNNER: ${{ contains(matrix.runner, 'a100') && '1' || '0' }}
-
-      shell: bash
-      run: |
-        set -x
-
-        if [[ $TEST_CONFIG == 'multigpu' ]]; then
-          TEST_COMMAND=.ci/pytorch/multigpu-test.sh
-        elif [[ $BUILD_ENVIRONMENT == *onnx* ]]; then
-          TEST_COMMAND=.ci/onnx/test.sh
-        else
-          TEST_COMMAND=.ci/pytorch/test.sh
-        fi
-
-        # detached container should get cleaned up by teardown_ec2_linux
-        # TODO: Stop building test binaries as part of the build phase
-        # Used for GPU_FLAG since that doesn't play nice
-        # shellcheck disable=SC2086,SC2090
-        container_name=$(docker run \
-          ${GPU_FLAG:-} \
-          ${SCCACHE_SERVER_PORT_DOCKER_FLAG:-} \
-          -e BUILD_ENVIRONMENT \
-          -e PR_NUMBER \
-          -e GITHUB_ACTIONS \
-          -e GITHUB_REPOSITORY \
-          -e GITHUB_WORKFLOW \
-          -e GITHUB_JOB \
-          -e GITHUB_RUN_ID \
-          -e GITHUB_RUN_NUMBER \
-          -e GITHUB_RUN_ATTEMPT \
-          -e JOB_ID \
-          -e JOB_NAME \
-          -e BASE_SHA \
-          -e BRANCH \
-          -e SHA1 \
-          -e AWS_DEFAULT_REGION \
-          -e IN_WHEEL_TEST \
-          -e SHARD_NUMBER \
-          -e TEST_CONFIG \
-          -e NUM_TEST_SHARDS \
-          -e REENABLED_ISSUES \
-          -e CONTINUE_THROUGH_ERROR \
-          -e VERBOSE_TEST_LOGS \
-          -e NO_TEST_TIMEOUT \
-          -e NO_TD \
-          -e TD_DISTRIBUTED \
-          -e PR_LABELS \
-          -e MAX_JOBS="$(nproc --ignore=2)" \
-          -e SCCACHE_BUCKET \
-          -e SCCACHE_REGION \
-          -e SCCACHE_S3_KEY_PREFIX \
-          -e XLA_CUDA \
-          -e XLA_CLANG_CACHE_S3_BUCKET_NAME \
-          -e PYTORCH_TEST_CUDA_MEM_LEAK_CHECK \
-          -e PYTORCH_TEST_RERUN_DISABLED_TESTS \
-          -e SKIP_SCCACHE_INITIALIZATION=1 \
-          -e HUGGING_FACE_HUB_TOKEN \
-          -e SCRIBE_GRAPHQL_ACCESS_TOKEN \
-          -e DASHBOARD_TAG \
-          -e IS_A100_RUNNER \
-          --env-file="/tmp/github_env_${GITHUB_RUN_ID}" \
-          --security-opt seccomp=unconfined \
-          --cap-add=SYS_PTRACE \
-          --ipc=host \
-          --shm-size="${SHM_SIZE}" \
-          --tty \
-          --detach \
-          --name="${container_name}" \
-          --user jenkins \
-          -v "${GITHUB_WORKSPACE}:/var/lib/jenkins/workspace" \
-          -w /var/lib/jenkins/workspace \
-          "${DOCKER_IMAGE}"
-        )
-        echo "DOCKER_CONTAINER_ID=${container_name}" >> "${GITHUB_ENV}"
-        docker exec -t "${container_name}" sh -c "pip install $(echo dist/*.whl)[opt-einsum] && ${TEST_COMMAND}"
-
-    - name: Upload pytest cache if tests failed
-      uses: ./.github/actions/pytest-cache-upload
-      continue-on-error: true
-      if: failure() && steps.test.conclusion && steps.test.conclusion == 'failure'
-      with:
-        cache_dir: .pytest_cache
-        shard: ${{ matrix.shard }}
-        sha: ${{ github.event.pull_request.head.sha || github.sha }}
-        test_config: ${{ matrix.config }}
-        job_identifier: ${{ github.workflow }}_${{ inputs.build-environment }}
-
-    - name: Print remaining test logs
-      shell: bash
-      if: always() && steps.test.conclusion
-      run: |
-        cat test/**/*_toprint.log || true
-
-    - name: Stop monitoring script
-      if: ${{ always() && steps.monitor-script.outputs.monitor-script-pid }}
-      shell: bash
-      continue-on-error: true
-      env:
-        MONITOR_SCRIPT_PID: ${{ steps.monitor-script.outputs.monitor-script-pid }}
-      run: |
-        kill "$MONITOR_SCRIPT_PID"
-
-    - name: Upload test artifacts
-      uses: ./.github/actions/upload-test-artifacts
-      if: always() && steps.test.conclusion && steps.test.conclusion != 'skipped'
-      with:
-        file-suffix: ${{ github.job }}-${{ matrix.config }}-${{ matrix.shard }}-${{ matrix.num_shards }}-${{ matrix.runner }}_${{ steps.get-job-id.outputs.job-id }}
-        use-gha: ${{ inputs.use-gha }}
-        s3-bucket: ${{ inputs.s3-bucket }}
-
-    - name: Collect backtraces from coredumps (if any)
-      if: always()
-      shell: bash
-      run: |
-        # shellcheck disable=SC2156
-        find . -iname "core.[1-9]*" -exec docker exec "${DOCKER_CONTAINER_ID}" sh -c "gdb python {} -ex 'bt' -ex 'q'" \;
-
-    - name: Store Core dumps on S3
-      uses: seemethere/upload-artifact-s3@v5
-      if: failure()
-      with:
-        name: coredumps-${{ matrix.config }}-${{ matrix.shard }}-${{ matrix.num_shards }}-${{ matrix.runner }}
-        retention-days: 14
-        if-no-files-found: ignore
-        path: ./**/core.[1-9]*
-
-    - name: Teardown Linux
-      uses: pytorch/test-infra/.github/actions/teardown-linux@main
-      if: always() && steps.check_container_runner.outputs.IN_CONTAINER_RUNNER == 'false'
-
-    # NB: We are currently having an intermittent GPU-related issue on G5 runners with
-    # A10G GPU. Once this happens, trying to reset the GPU as done in setup-nvidia does
-    # not seem to help. Here are some symptoms:
-    #   * Calling nvidia-smi timeouts after 60 second
-    #   * Fail to run nvidia-smi with an unable to determine the device handle for GPU
-    #     unknown error
-    #   * Test fails with a missing CUDA GPU error when initializing CUDA in PyTorch
-    #   * Run docker --gpus all fails with error response from daemon
-    #
-    # As both the root cause and recovery path are unclear, let's take the runner out of
-    # service so that it doesn't get any more jobs
-    - name: Check NVIDIA driver installation step
-      if: failure() && steps.install-nvidia-driver.outcome && steps.install-nvidia-driver.outcome != 'skipped'
-      shell: bash
-      env:
-        RUNNER_WORKSPACE: ${{ runner.workspace }}
-      run: |
-        set +e
-        set -x
-
-        nvidia-smi
-        # NB: Surprisingly, nvidia-smi command returns successfully with return code 0 even in
-        # the case where the driver has already crashed as it still can get the driver version
-        # and some basic information like the bus ID.  However, the rest of the information
-        # would be missing (ERR!), for example:
-        #
-        # +-----------------------------------------------------------------------------+
-        # | NVIDIA-SMI 525.89.02    Driver Version: 525.89.02    CUDA Version: 12.0     |
-        # |-------------------------------+----------------------+----------------------+
-        # | GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
-        # | Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
-        # |                               |                      |               MIG M. |
-        # |===============================+======================+======================|
-        # |   0  ERR!                Off  | 00000000:00:1E.0 Off |                 ERR! |
-        # |ERR!  ERR! ERR!    ERR! / ERR! |   4184MiB / 23028MiB |    ERR!      Default |
-        # |                               |                      |                 ERR! |
-        # +-------------------------------+----------------------+----------------------+
-        #
-        # +-----------------------------------------------------------------------------+
-        # | Processes:                                                                  |
-        # |  GPU   GI   CI        PID   Type   Process name                  GPU Memory |
-        # |        ID   ID                                                   Usage      |
-        # |=============================================================================|
-        # +-----------------------------------------------------------------------------+
-        #
-        # This should be reported as a failure instead as it will guarantee to fail when
-        # Docker tries to run with --gpus all
-        #
-        # So, the correct check here is to query one of the missing piece of info like
-        # GPU name, so that the command can fail accordingly
-        nvidia-smi --query-gpu=gpu_name --format=csv,noheader --id=0
-        NVIDIA_SMI_STATUS=$?
-
-        # These are acceptable return code from nvidia-smi as copied from setup-nvidia GitHub action
-        if [ "$NVIDIA_SMI_STATUS" -ne 0 ] && [ "$NVIDIA_SMI_STATUS" -ne 14 ]; then
-          echo "NVIDIA driver installation has failed, shutting down the runner..."
-          .github/scripts/stop_runner_service.sh
-        fi
-
-        # For runner with multiple GPUs, we also want to confirm that the number of GPUs are the
-        # power of 2, i.e. 1, 2, 4, or 8. This is to avoid flaky test issue when one GPU fails
-        # https://github.com/pytorch/test-infra/issues/4000
-        GPU_COUNT=$(nvidia-smi --list-gpus | wc -l)
-        NVIDIA_SMI_STATUS=$?
-
-        # These are acceptable return code from nvidia-smi as copied from setup-nvidia GitHub action
-        if [ "$NVIDIA_SMI_STATUS" -ne 0 ] && [ "$NVIDIA_SMI_STATUS" -ne 14 ]; then
-          echo "NVIDIA driver installation has failed, shutting down the runner..."
-          .github/scripts/stop_runner_service.sh
-        fi
-
-        # Check the GPU count to be a power of 2
-        if [ "$GPU_COUNT" -le 8 ] && [ "$GPU_COUNT" -ne 1 ] && [ "$GPU_COUNT" -ne 2 ] && [ "$GPU_COUNT" -ne 4 ] && [ "$GPU_COUNT" -ne 8 ]; then
-          echo "NVIDIA driver detects $GPU_COUNT GPUs. The runner has a broken GPU, shutting it down..."
-          .github/scripts/stop_runner_service.sh
-        fi
diff --git a/.github/actions/pytest-cache-download/action.yml b/.github/actions/pytest-cache-download/action.yml
index e23e0a5eaba70..bb9f3f33aef62 100644
--- a/.github/actions/pytest-cache-download/action.yml
+++ b/.github/actions/pytest-cache-download/action.yml
@@ -18,6 +18,7 @@ runs:
   using: composite
   steps:
     - name: Setup uv
+      if: ${{ !env.UV_PYTHON }}
       uses: pytorch/test-infra/.github/actions/setup-uv@main
       with:
         python-version: "3.12"
diff --git a/.github/actions/pytest-cache-upload/action.yml b/.github/actions/pytest-cache-upload/action.yml
index 46e7d2db7e935..ac6c1a709c15e 100644
--- a/.github/actions/pytest-cache-upload/action.yml
+++ b/.github/actions/pytest-cache-upload/action.yml
@@ -25,6 +25,7 @@ runs:
   using: composite
   steps:
     - name: Setup uv
+      if: ${{ !env.UV_PYTHON }}
       uses: pytorch/test-infra/.github/actions/setup-uv@main
       with:
         python-version: "3.12"
diff --git a/.github/actions/reuse-old-whl/action.yml b/.github/actions/reuse-old-whl/action.yml
index 46e7d0e278ba4..584c3a2ab64af 100644
--- a/.github/actions/reuse-old-whl/action.yml
+++ b/.github/actions/reuse-old-whl/action.yml
@@ -30,6 +30,7 @@ runs:
 
   steps:
     - name: Setup uv
+      if: ${{ !env.UV_PYTHON }}
       uses: pytorch/test-infra/.github/actions/setup-uv@main
       with:
         python-version: "3.12"
diff --git a/.github/actions/reuse-old-whl/reuse_old_whl.py b/.github/actions/reuse-old-whl/reuse_old_whl.py
index 48a8490985946..725eec360bb26 100644
--- a/.github/actions/reuse-old-whl/reuse_old_whl.py
+++ b/.github/actions/reuse-old-whl/reuse_old_whl.py
@@ -1,10 +1,12 @@
+from __future__ import annotations
+
 import argparse
 import os
 import subprocess
 import sys
 from functools import lru_cache
 from pathlib import Path
-from typing import Any, cast, Optional, Union
+from typing import Any, cast
 
 import requests
 
@@ -100,7 +102,7 @@ def check_issue_open() -> bool:
         return False
 
 
-def get_workflow_id(run_id: str) -> Optional[str]:
+def get_workflow_id(run_id: str) -> str | None:
     # Get the workflow ID that corresponds to the file for the run ID
     url = f"https://api.github.com/repos/pytorch/pytorch/actions/runs/{run_id}"
     response = query_github_api(url)
@@ -228,13 +230,13 @@ def unzip_artifact_and_replace_files() -> None:
         old_version = f"+git{path.stem.split('+')[1].split('-')[0][3:]}"
         new_version = f"+git{head_sha[:7]}"
 
-        def rename_to_new_version(file: Union[str, Path]) -> None:
+        def rename_to_new_version(file: str | Path) -> None:
             # Rename file with old_version to new_version
             subprocess.check_output(
                 ["mv", file, str(file).replace(old_version, new_version)]
             )
 
-        def change_content_to_new_version(file: Union[str, Path]) -> None:
+        def change_content_to_new_version(file: str | Path) -> None:
             # Check if is a file
             if os.path.isdir(file):
                 return
diff --git a/.github/actions/setup-linux/action.yml b/.github/actions/setup-linux/action.yml
index 6314e0e1f7db4..1fb7a8e7e8b66 100644
--- a/.github/actions/setup-linux/action.yml
+++ b/.github/actions/setup-linux/action.yml
@@ -1,11 +1,104 @@
 name: Setup Linux
 
-description: Set up Docker workspace on EC2
+description: Set up Linux workspace on EC2 or OSDC ARC runners
+
+inputs:
+  use-arc:
+    description: Whether the runner is an OSDC ARC runner (returned by runner_determinator)
+    required: false
+    default: ''
+  python-version:
+    description: Python version to install (e.g. "3.12"). Empty string keeps the default.
+    required: false
+    default: ''
+  compiler:
+    description: Compiler name and version (e.g. "gcc11", "clang15"). Empty string keeps the base image default.
+    required: false
+    default: ''
+  cuda-version:
+    description: CUDA version to activate (e.g. "13.0"). Empty string keeps the default.
+    required: false
+    default: ''
+  submodules:
+    description: Submodule checkout mode passed to checkout-pytorch (default "recursive", use "false" for test jobs).
+    required: false
+    default: 'recursive'
+  checkout-mode:
+    description: Checkout mode passed to checkout-pytorch (one of "normal", "blobless", "treeless").
+    required: false
+    default: 'treeless'
+  github-token:
+    description: GITHUB_TOKEN, needed to retrieve the workflow job id.
+    required: false
+    default: ''
+
+outputs:
+  branch:
+    description: Parsed branch name from GITHUB_REF
+    value: ${{ steps.parse-ref.outputs.branch }}
+  tag:
+    description: Parsed tag name from GITHUB_REF (if applicable)
+    value: ${{ steps.parse-ref.outputs.tag }}
+  job-id:
+    description: The workflow job id
+    value: ${{ steps.get-job-id.outputs.job-id }}
+  job-name:
+    description: The workflow job name
+    value: ${{ steps.get-job-id.outputs.job-name }}
 
 runs:
   using: composite
   steps:
+    # ── ARC-only steps ──────────────────────────────────────────────────
+    - name: Fix workspace permissions
+      if: ${{ inputs.use-arc == 'true' }}
+      shell: bash
+      run: |
+        # GH runner image has switched to uid 1001 https://github.com/actions/runner-images/issues/10936
+        # while current PyTorch CI image are still using uid 1000 (ec2-user). We
+        # can update the uid to 1001 eventually when everything migrates to ARC.
+        # In the meantime, this is a quick fix to ensure that CI has the permission
+        # to use the GITHUB_WORKSPACE while still allowing the GH hook (uid 1001,
+        # gid 1001) to clean up the directory after the job
+        sudo chmod -R 777 "$GITHUB_WORKSPACE"
+
+    - name: Ack Git cache ownership
+      if: ${{ inputs.use-arc == 'true' }}
+      shell: bash
+      run: |
+        git config --global --add safe.directory "$GITHUB_WORKSPACE"
+
+    # ── Common steps (shared) ────────────────────────────────────────
+    - name: Checkout PyTorch
+      uses: pytorch/pytorch/.github/actions/checkout-pytorch@main
+      with:
+        no-sudo: true
+        checkout-mode: ${{ inputs.checkout-mode }}
+        submodules: ${{ inputs.submodules }}
+
+    - name: Parse ref
+      id: parse-ref
+      shell: bash
+      run: |
+        if [ -f "${{ github.action_path }}/../../scripts/parse_ref.py" ]; then
+          python3 "${{ github.action_path }}/../../scripts/parse_ref.py"
+        elif [ -f .github/scripts/parse_ref.py ]; then
+          python3 .github/scripts/parse_ref.py
+        else
+          echo "ERROR: parse_ref.py not found" >&2
+          exit 1
+        fi
+
+    - name: Get workflow job id
+      id: get-job-id
+      if: ${{ always() && inputs.github-token != '' }}
+      uses: pytorch/pytorch/.github/actions/get-workflow-job-id@main
+      with:
+        github-token: ${{ inputs.github-token }}
+
+    # ── EC2-only steps ──────────────────────────────────────────────────
     - name: Display EC2 information
+      if: ${{ inputs.use-arc != 'true' }}
       shell: bash
       run: |
         set -euo pipefail
@@ -28,11 +121,8 @@ runs:
         echo "instance-type: $(get_ec2_metadata instance-type)"
         echo "system info $(uname -a)"
 
-    - name: Print GPU info (if present)
-      shell: bash
-      run: if [ -f /usr/bin/nvidia-smi ]; then nvidia-smi; fi
-
     - name: Check if in a container runner
+      if: ${{ inputs.use-arc != 'true' }}
       shell: bash
       id: check_container_runner
       run: echo "IN_CONTAINER_RUNNER=$(if [ -f /.inarc ] || [ -f /.incontainer ]; then echo true ; else echo false; fi)" >> "$GITHUB_OUTPUT"
@@ -49,7 +139,7 @@ runs:
           fi
         fi
 
-    - name: Install uv
+    - name: Install uv (EC2)
       uses: pytorch/test-infra/.github/actions/setup-uv@main
       if: ${{ steps.check_container_runner.outputs.IN_CONTAINER_RUNNER == 'false' }}
       with:
@@ -67,12 +157,6 @@ runs:
           uv tool install pip
         fi
 
-    - name: Preserve github env variables for use in docker
-      shell: bash
-      run: |
-        env | grep '^GITHUB' >> "/tmp/github_env_${GITHUB_RUN_ID}"
-        env | grep '^CI' >> "/tmp/github_env_${GITHUB_RUN_ID}"
-
     - name: Kill any existing containers, clean up images
       if: ${{ steps.check_container_runner.outputs.IN_CONTAINER_RUNNER == 'false' }}
       shell: bash
@@ -107,3 +191,14 @@ runs:
         done
         echo "Reached maximum attempts to connect to Docker. Exiting."
         exit 1
+
+    # ── Shared steps ────────────────────────────────────────────────────
+    - name: Print GPU info (if present)
+      shell: bash
+      run: if [ -f /usr/bin/nvidia-smi ]; then nvidia-smi; fi
+
+    - name: Preserve github env variables for use in docker
+      shell: bash
+      run: |
+        env | grep '^GITHUB' >> "/tmp/github_env_${GITHUB_RUN_ID}"
+        env | grep '^CI' >> "/tmp/github_env_${GITHUB_RUN_ID}"
diff --git a/.github/actions/setup-rocm/action.yml b/.github/actions/setup-rocm/action.yml
index 5bb982a4085b1..3c1445b85fa84 100644
--- a/.github/actions/setup-rocm/action.yml
+++ b/.github/actions/setup-rocm/action.yml
@@ -110,7 +110,7 @@ runs:
         # This is due to the device files (/dev/kfd & /dev/dri) being owned by video group on bare metal.
         # This video group ID maps to subgid 1 inside the docker image due to the /etc/subgid entries.
         # The group name corresponding to group ID 1 can change depending on the OS, so both are necessary.
-        echo "GPU_FLAG=--device=/dev/mem --device=/dev/kfd $DEVICE_FLAG --group-add video --group-add $render_gid --group-add daemon --group-add bin --cap-add=SYS_PTRACE --security-opt seccomp=unconfined --network=host" >> "${GITHUB_ENV}"
+        echo "GPU_FLAG=--device=/dev/mem --device=/dev/kfd $DEVICE_FLAG --group-add video --group-add $render_gid --group-add daemon --group-add bin --cap-add=SYS_PTRACE --security-opt seccomp=unconfined" >> "${GITHUB_ENV}"
 
     - name: Login to ECR
       uses: pytorch/pytorch/.github/actions/ecr-login@main
diff --git a/.github/actions/setup-xpu/action.yml b/.github/actions/setup-xpu/action.yml
index 740492475d6e2..23b750205ab08 100644
--- a/.github/actions/setup-xpu/action.yml
+++ b/.github/actions/setup-xpu/action.yml
@@ -44,7 +44,7 @@ runs:
         fi
 
     - name: Runner diskspace health check
-      uses: ./.github/actions/diskspace-cleanup
+      uses: pytorch/pytorch/.github/actions/diskspace-cleanup@main
       if: always()
 
     - name: Runner health check disconnect on failure
@@ -53,11 +53,30 @@ runs:
       run: |
         killall runsvc.sh
 
+    - name: Setup useful environment variables
+      shell: bash
+      run: |
+        RUNNER_ARTIFACT_DIR="${RUNNER_TEMP}/artifacts"
+        rm -rf "${RUNNER_ARTIFACT_DIR}"
+        mkdir -p "${RUNNER_ARTIFACT_DIR}"
+        echo "RUNNER_ARTIFACT_DIR=${RUNNER_ARTIFACT_DIR}" >> "${GITHUB_ENV}"
+
+        RUNNER_TEST_RESULTS_DIR="${RUNNER_TEMP}/test-results"
+        rm -rf "${RUNNER_TEST_RESULTS_DIR}"
+        mkdir -p "${RUNNER_TEST_RESULTS_DIR}"
+        echo "RUNNER_TEST_RESULTS_DIR=${RUNNER_TEST_RESULTS_DIR}" >> "${GITHUB_ENV}"
+
+        RUNNER_DOCS_DIR="${RUNNER_TEMP}/docs"
+        rm -rf "${RUNNER_DOCS_DIR}"
+        mkdir -p "${RUNNER_DOCS_DIR}"
+        echo "RUNNER_DOCS_DIR=${RUNNER_DOCS_DIR}" >> "${GITHUB_ENV}"
+
     - name: Preserve github env variables for use in docker
       shell: bash
       run: |
-        env | grep '^GITHUB' >> "/tmp/github_env_${GITHUB_RUN_ID}"
-        env | grep '^CI' >> "/tmp/github_env_${GITHUB_RUN_ID}"
+        env | grep '^GITHUB' >> "${RUNNER_TEMP}/github_env_${GITHUB_RUN_ID}"
+        env | grep '^CI' >> "${RUNNER_TEMP}/github_env_${GITHUB_RUN_ID}"
+        env | grep '^RUNNER' >> "${RUNNER_TEMP}/github_env_${GITHUB_RUN_ID}"
 
     - name: XPU set GPU_FLAG
       shell: bash
@@ -65,3 +84,6 @@ runs:
         # Add render group for container creation.
         render_gid=`cat /etc/group | grep render | cut -d: -f3`
         echo "GPU_FLAG=--device=/dev/mem --device=/dev/dri --group-add video --group-add $render_gid" >> "${GITHUB_ENV}"
+
+    - name: Login to ECR
+      uses: pytorch/pytorch/.github/actions/ecr-login@main
diff --git a/.github/actions/test-pytorch-binary/action.yml b/.github/actions/test-pytorch-binary/action.yml
index 991cf9fb87eff..4c25c8450bce8 100644
--- a/.github/actions/test-pytorch-binary/action.yml
+++ b/.github/actions/test-pytorch-binary/action.yml
@@ -26,7 +26,7 @@ runs:
           -e SKIP_ALL_TESTS \
           --tty \
           --detach \
-          -v "${GITHUB_WORKSPACE}/pytorch:/pytorch" \
+          -v "${GITHUB_WORKSPACE}:/pytorch" \
           -v "${RUNNER_TEMP}/artifacts:/final_pkgs" \
           -w / \
           "${DOCKER_IMAGE}"
@@ -34,9 +34,9 @@ runs:
 
         echo "CONTAINER_NAME=${container_name}" >> "$GITHUB_ENV"
 
-        docker exec -t -w "${PYTORCH_ROOT}" "${container_name}" bash -c "bash .circleci/scripts/binary_populate_env.sh"
+        docker exec -t -w "${PYTORCH_ROOT}" "${container_name}" bash -c "bash .ci/pytorch/binary_populate_env.sh"
         # Generate test script
-        docker exec -t -w "${PYTORCH_ROOT}" -e OUTPUT_SCRIPT="/run.sh" "${container_name}" bash -c "bash .circleci/scripts/binary_linux_test.sh"
+        docker exec -t -w "${PYTORCH_ROOT}" -e OUTPUT_SCRIPT="/run.sh" "${container_name}" bash -c "bash .ci/pytorch/binary_linux_test.sh"
         docker exec -t "${container_name}" bash -c "source ${BINARY_ENV_FILE} && bash -x /run.sh"
 
     - name: Cleanup docker
diff --git a/.github/actions/upload-build-artifacts/action.yml b/.github/actions/upload-build-artifacts/action.yml
new file mode 100644
index 0000000000000..62413131581ea
--- /dev/null
+++ b/.github/actions/upload-build-artifacts/action.yml
@@ -0,0 +1,37 @@
+name: Upload PyTorch Build Artifacts
+
+description: Upload build artifacts to S3 or GHA.
+
+inputs:
+  name:
+    description: Name of the artifact
+    required: true
+  use-gha:
+    description: If set to any value, use GHA to upload the artifact. Otherwise use S3.
+    required: false
+  s3-bucket:
+    description: S3 bucket to upload builds
+    required: false
+    default: "gha-artifacts"
+
+runs:
+  using: composite
+  steps:
+    - name: Upload PyTorch Build Artifacts to S3
+      if: inputs.use-gha == '' || inputs.use-gha == 'false'
+      uses: seemethere/upload-artifact-s3@baba72d0712b404f646cebe0730933554ebce96a # v5.1.0
+      with:
+        name: ${{ inputs.name }}
+        retention-days: 14
+        if-no-files-found: error
+        path: artifacts.zip
+        s3-bucket: ${{ inputs.s3-bucket }}
+
+    - name: Upload PyTorch Build Artifacts to GHA
+      if: inputs.use-gha != '' && inputs.use-gha != 'false'
+      uses: actions/upload-artifact@ea165f8d65b6e75b540449e92b4886f43607fa02 # v4.6.2
+      with:
+        name: ${{ inputs.name }}
+        retention-days: 14
+        if-no-files-found: error
+        path: artifacts.zip
diff --git a/.github/actions/upload-test-artifacts/action.yml b/.github/actions/upload-test-artifacts/action.yml
index 6f4c8a89c0f45..9fcb125a42313 100644
--- a/.github/actions/upload-test-artifacts/action.yml
+++ b/.github/actions/upload-test-artifacts/action.yml
@@ -25,7 +25,7 @@ runs:
   steps:
     # Always set up uv and zip files first (needed for S3, reusable for GHA fallback)
     - name: Setup uv
-      if: runner.os != 'Windows'
+      if: ${{ runner.os != 'Windows' && !env.UV_PYTHON }}
       uses: pytorch/test-infra/.github/actions/setup-uv@main
       with:
         python-version: "3.12"
@@ -37,7 +37,7 @@ runs:
         FILE_SUFFIX: ${{ inputs.file-suffix }}
       run: |
         set -euo pipefail
-        rm -f test-jsons-*.zip test-reports-*.zip logs-*.zip debug-*.zip
+        rm -f test-jsons-*.zip test-reports-*.zip logs-*.zip debug-*.zip profiler-traces-*.zip tlparse-*.zip
         ZIP_CMD='
         import sys, zipfile, os
         from pathlib import Path
@@ -60,6 +60,14 @@ runs:
           uv run --no-project python -c "$ZIP_CMD" \
             test/debug "debug-${FILE_SUFFIX}.zip" '**/*'
         fi
+        if [ -d 'test/test-reports/profiler_traces' ]; then
+          uv run --no-project python -c "$ZIP_CMD" \
+            test/test-reports/profiler_traces "profiler-traces-${FILE_SUFFIX}.zip" '**/*'
+        fi
+        if [ -d 'test/test-reports/tlparse_output' ]; then
+          uv run --no-project python -c "$ZIP_CMD" \
+            test/test-reports/tlparse_output "tlparse-${FILE_SUFFIX}.zip" '**/*'
+        fi
 
     # Windows zip
     - name: Zip JSONs for upload
@@ -141,6 +149,30 @@ runs:
         if-no-files-found: ignore
         path: debug-*.zip
 
+    - name: Store Profiler Traces on S3
+      id: s3-upload-profiler-traces
+      uses: seemethere/upload-artifact-s3@v5
+      continue-on-error: true
+      with:
+        s3-bucket: ${{ inputs.s3-bucket }}
+        s3-prefix: |
+          ${{ github.repository }}/${{ github.run_id }}/${{ github.run_attempt }}/artifact
+        retention-days: 14
+        if-no-files-found: ignore
+        path: profiler-traces-*.zip
+
+    - name: Store TLParse Output on S3
+      id: s3-upload-tlparse
+      uses: seemethere/upload-artifact-s3@v5
+      continue-on-error: true
+      with:
+        s3-bucket: ${{ inputs.s3-bucket }}
+        s3-prefix: |
+          ${{ github.repository }}/${{ github.run_id }}/${{ github.run_attempt }}/artifact
+        retention-days: 14
+        if-no-files-found: ignore
+        path: tlparse-*.zip
+
     # Check if S3 upload failed (test-reports is the critical one)
     - name: Check S3 upload status
       id: check-s3
@@ -193,3 +225,21 @@ runs:
         path: |
           usage_log.txt
           test/**/*.log
+
+    - name: Store Profiler Traces on Github
+      uses: actions/upload-artifact@v4
+      continue-on-error: true
+      with:
+        name: profiler-traces-runattempt${{ github.run_attempt }}-${{ inputs.file-suffix }}
+        retention-days: 14
+        if-no-files-found: ignore
+        path: test/test-reports/profiler_traces/**/*
+
+    - name: Store TLParse Output on Github
+      uses: actions/upload-artifact@v4
+      continue-on-error: true
+      with:
+        name: tlparse-runattempt${{ github.run_attempt }}-${{ inputs.file-suffix }}
+        retention-days: 14
+        if-no-files-found: ignore
+        path: test/test-reports/tlparse_output/**/*
diff --git a/.github/actions/upload-utilization-stats/action.yml b/.github/actions/upload-utilization-stats/action.yml
index 61332feaec688..007abb7c4bdf4 100644
--- a/.github/actions/upload-utilization-stats/action.yml
+++ b/.github/actions/upload-utilization-stats/action.yml
@@ -39,6 +39,7 @@ runs:
   using: composite
   steps:
     - name: Setup uv
+      if: ${{ !env.UV_PYTHON }}
       uses: pytorch/test-infra/.github/actions/setup-uv@main
       with:
         python-version: "3.12"
diff --git a/.github/allowlist.yml b/.github/allowlist.yml
new file mode 100644
index 0000000000000..c6a0bd196837c
--- /dev/null
+++ b/.github/allowlist.yml
@@ -0,0 +1,36 @@
+# Cross Repo CI Relay (CRCR) Allowlist
+#
+# This document defines which downstream repositories can receive PyTorch
+# PR events or feedback through CRCR. Each repository is assigned a level
+# (L1-L4), which determines the depth to which downstream repositories can
+# participate.
+#
+# Trust Levels:
+#
+#   L1 – Onboarding
+#     Events are forwarded to downstream, but upstream receives no feedback.
+#
+#   L2 – Observation
+#     Downstream CI results are displayed on the HUD page, but not on PRs.
+#
+#   L3 – Stable
+#     Adds a non-blocking check run on PRs when ciflow/oot/<name> label is applied.
+#
+#   L4 – Mature
+#     Adds a blocking check run on every PR; reserved for critical accelerators.
+#
+# For more information, see the RFC: https://github.com/pytorch/rfcs/pull/90
+
+# Example:
+#
+# L1:
+#     - org1/downstream-repo1
+#
+# L2:
+#     - org2/downstream-repo2
+#
+# L3:
+#     - org3/downstream-repo3: @oncall1,oncall2
+#
+# L4:
+#     - org4/downstream-repo4: @oncall1,oncall2
diff --git a/.github/arc.yaml b/.github/arc.yaml
new file mode 100644
index 0000000000000..497f4f5f01e9c
--- /dev/null
+++ b/.github/arc.yaml
@@ -0,0 +1,89 @@
+# ARC (Actions Runner Controller) Runner Label Mapping
+#
+# Maps current GitHub Actions runner labels to new ARC runner labels.
+# Reference: https://github.com/pytorch/ci-infra/issues/396
+#
+# New label format:
+#   {os}-[b]{arch}{vendor}{features}-{vcpu}-{memory}[-{gpu_type}[-{gpu_count}]]
+#
+# Fields:
+#   os       - l=Linux, w=Windows, m=MacOS
+#   b        - (optional) bare-metal instance
+#   arch     - x86=x86_64, arm64=AArch64
+#   vendor   - i=Intel, a=AMD, g2/g3/g4=Graviton gen
+#   features - (x86 only) avx2, avx512, amx
+#   vcpu     - vCPU count
+#   memory   - RAM in GiB
+#   gpu_type - (optional) t4, a10g, l4
+#   gpu_count- (optional, omitted when 1)
+#
+# Entries marked "# upgraded" had no exact ARC equivalent and were mapped to
+# the next larger available runner.
+
+runner_mapping:
+
+  # ---- x86 CPU — Intel AVX-512 (c5, c7i families) ----
+
+  linux.large: l-x86iavx512-2-4                        # c5.large
+  linux.2xlarge: l-x86iavx512-8-64                     # c5.2xlarge
+  linux.c7i.2xlarge: l-x86iavx512-8-64                 # c7i.2xlarge
+  linux.4xlarge: l-x86iavx512-16-128                   # c5.4xlarge
+  linux.c7i.4xlarge: l-x86iavx512-16-128               # c7i.4xlarge
+  linux.12xlarge: l-x86iavx512-48-384                  # c5.12xlarge
+  linux.c7i.12xlarge: l-x86iavx512-48-384              # c7i.12xlarge
+  linux.24xl.spr-metal: l-bx86iamx-92-167              # c7i.metal-24xl
+
+  # ---- x86 CPU — Intel AMX (m7i-flex family) ----
+
+  linux.2xlarge.amx: l-x86iamx-8-64                    # m7i-flex.2xlarge
+  linux.8xlarge.amx: l-x86iamx-32-128                  # m7i-flex.8xlarge
+
+  # ---- x86 CPU — Intel AVX2 (m4 family) ----
+
+  linux.2xlarge.avx2: l-x86iavx2-8-32                  # m4.2xlarge
+  linux.10xlarge.avx2: l-x86iavx2-40-160               # m4.10xlarge
+
+  # ---- x86 CPU — Memory-optimized (r5, r7i families) ----
+
+  linux.r7i.2xlarge: l-x86iavx512-8-64                 # r7i.2xlarge
+  linux.r7i.4xlarge: l-x86iavx512-16-128               # r7i.4xlarge
+  linux.4xlarge.memory: l-x86iavx512-16-128            # r5.4xlarge
+  linux.8xlarge.memory: l-x86iavx512-32-256            # r5.8xlarge
+  linux.12xlarge.memory: l-x86iavx512-48-384           # r5.12xlarge
+  linux.24xlarge.memory: l-x86iavx512-94-768           # r5.24xlarge
+
+  # ---- x86 CPU — AMD (m6a, m7a families) ----
+
+  linux.24xlarge.amd: l-x86aavx512-125-463             # m7a.24xlarge
+
+  # ---- x86 GPU — T4 (g4dn family) ----
+
+  linux.g4dn.4xlarge.nvidia.gpu: l-x86iavx512-29-115-t4          # g4dn.4xlarge
+  linux.g4dn.12xlarge.nvidia.gpu: l-x86iavx512-45-172-t4-4       # g4dn.12xlarge
+  linux.g4dn.metal.nvidia.gpu: l-bx86iavx512-94-344-t4-8         # g4dn.metal
+
+  # ---- x86 GPU — A10G (g5 family) ----
+
+  linux.g5.4xlarge.nvidia.gpu: l-x86aavx2-29-113-a10g            # g5.4xlarge
+  linux.g5.12xlarge.nvidia.gpu: l-x86aavx2-45-167-a10g-4         # g5.12xlarge
+  linux.g5.48xlarge.nvidia.gpu: l-x86aavx2-189-704-a10g-8        # g5.48xlarge
+
+  # ---- x86 GPU — L4 (g6 family) ----
+
+  linux.g6.4xlarge.experimental.nvidia.gpu: l-x86aavx2-29-113-l4 # g6.4xlarge
+  linux.g6.12xlarge.nvidia.gpu: l-x86aavx2-45-172-l4-4           # g6.12xlarge
+
+  # ---- ARM64 — Graviton ----
+
+  linux.arm64.2xlarge: l-arm64g2-6-32                            # t4g.2xlarge
+  linux.arm64.2xlarge.ephemeral: l-arm64g2-6-32                  # t4g.2xlarge
+  linux.arm64.m7g.4xlarge: l-arm64g3-16-62                       # m7g.4xlarge
+  linux.arm64.m8g.4xlarge: l-arm64g4-16-62                       # m8g.4xlarge
+  linux.arm64.r7g.12xlarge.memory: l-arm64g3-61-463              # r7g.12xlarge
+  linux.arm64.m7g.metal: l-barm64g4-62-226                       # m7g.metal
+
+
+  # ---- Partner hardwares ----
+
+  linux.idc.xpu: linux.idc.xpu
+  linux.rocm.gpu.2: linux.rocm.gpu.2
diff --git a/.github/ci_commit_pins/audio.txt b/.github/ci_commit_pins/audio.txt
index c44ee060b2c0c..e5487b7ac5eff 100644
--- a/.github/ci_commit_pins/audio.txt
+++ b/.github/ci_commit_pins/audio.txt
@@ -1 +1 @@
-96200d42e795098ee5e3f0106bcf9c370cf126a1
+c0cbdb95674556cdff7266f2d44bb855f634cfde
diff --git a/.github/ci_commit_pins/torch_tpu.txt b/.github/ci_commit_pins/torch_tpu.txt
index 6835895a298dd..0f959421d9453 100644
--- a/.github/ci_commit_pins/torch_tpu.txt
+++ b/.github/ci_commit_pins/torch_tpu.txt
@@ -1 +1 @@
-efec24ca8adca3d742e7fa3e0483922de20b4ae9
+c1f09e2a7d5a595abe600089306cb73bce91db72
diff --git a/.github/ci_commit_pins/torchao.txt b/.github/ci_commit_pins/torchao.txt
index f33eac1d8778d..261d79443851c 100644
--- a/.github/ci_commit_pins/torchao.txt
+++ b/.github/ci_commit_pins/torchao.txt
@@ -1 +1 @@
-985d970b5e16b58c1e5b8bab440169d3da78cf16
+fbce9178697ca22d604314b316feb765360e0dec
diff --git a/.github/ci_commit_pins/torchtitan.txt b/.github/ci_commit_pins/torchtitan.txt
new file mode 100644
index 0000000000000..0c66dcaabc888
--- /dev/null
+++ b/.github/ci_commit_pins/torchtitan.txt
@@ -0,0 +1 @@
+7c8fc3f540edaec84d4610066ccff423141f7aa7
diff --git a/.github/ci_commit_pins/vision.txt b/.github/ci_commit_pins/vision.txt
index e7ac8b6c6bbd7..516c2c1931b15 100644
--- a/.github/ci_commit_pins/vision.txt
+++ b/.github/ci_commit_pins/vision.txt
@@ -1 +1 @@
-4b0a90c3616a348ee0cc9ce7fb36eb1b0e40ce54
+601776dc14ab12179412fa7fb08762e20862720c
diff --git a/.github/ci_commit_pins/vllm.txt b/.github/ci_commit_pins/vllm.txt
index 235f99edfa759..21f23c297a01d 100644
--- a/.github/ci_commit_pins/vllm.txt
+++ b/.github/ci_commit_pins/vllm.txt
@@ -1 +1 @@
-52ee21021a87735d46c4245c60bc0be42dd58c73
+ba4a78eb5d2ea30477b58a0bb8109b129f35c8b1
diff --git a/.github/ci_commit_pins/xla.txt b/.github/ci_commit_pins/xla.txt
index 7d3a8548d145b..42fb03f58599d 100644
--- a/.github/ci_commit_pins/xla.txt
+++ b/.github/ci_commit_pins/xla.txt
@@ -1 +1 @@
-c04e61c3424142c0eebcc9e59984b9d8fced18c0
+r2.12
diff --git a/.github/ci_configs/vllm/Dockerfile b/.github/ci_configs/vllm/Dockerfile
index 1d6ac16926834..549c336a444cb 100644
--- a/.github/ci_configs/vllm/Dockerfile
+++ b/.github/ci_configs/vllm/Dockerfile
@@ -146,6 +146,9 @@ ARG SCCACHE_BUCKET_NAME=vllm-build-sccache
 ARG SCCACHE_REGION_NAME=us-west-2
 ARG SCCACHE_S3_NO_CREDENTIALS=0
 
+ARG torch_cuda_arch_list='8.0;8.9;9.0;10.0;12.0'
+ENV TORCH_CUDA_ARCH_LIST=${torch_cuda_arch_list}
+
 # Use sccache to speed up compilation
 RUN --mount=type=cache,target=/root/.cache/uv \
     --mount=type=bind,source=.git,target=.git \
@@ -171,9 +174,6 @@ RUN --mount=type=cache,target=/root/.cache/uv \
         && sccache --show-stats; \
     fi
 
-ARG torch_cuda_arch_list='8.0 8.6 8.9 9.0'
-ENV TORCH_CUDA_ARCH_LIST=${torch_cuda_arch_list}
-
 ARG vllm_target_device="cuda"
 ENV VLLM_TARGET_DEVICE=${vllm_target_device}
 ENV CCACHE_DIR=/root/.cache/ccache
diff --git a/.github/labeler.yml b/.github/labeler.yml
index 39ebe16318560..c8f77e91be65d 100644
--- a/.github/labeler.yml
+++ b/.github/labeler.yml
@@ -62,16 +62,6 @@
 - test/test_jit_llga_fuser.py
 - test/test_mkldnn.py
 
-"ciflow/linux-aarch64":
-- third_party/ideep
-- caffe2/ideep/**
-- caffe2/python/ideep/**
-- cmake/Modules/FindMKLDNN.cmake
-- third_party/mkl-dnn.BUILD
-- torch/csrc/jit/codegen/onednn/**
-- test/test_jit_llga_fuser.py
-- test/test_mkldnn.py
-
 "module: amp (automated mixed precision)":
 - torch/amp/**
 - aten/src/ATen/autocast_mode.*
@@ -128,17 +118,31 @@
 - torch/distributed/_tensor/**
 - torch/distributed/tensor/**
 
+"ciflow/dtensor":
+- test/distributed/tensor/**
+- torch/distributed/tensor/**
+
 "ciflow/vllm":
 - .github/ci_commit_pins/vllm.txt
 
 "ci-refresh-hf-cache":
 - .github/ci_commit_pins/vllm.txt
+- .ci/docker/ci_commit_pins/timm.txt
+- .ci/docker/ci_commit_pins/huggingface-requirements.txt
+- .ci/docker/ci_commit_pins/torchbench.txt
 
 "ciflow/inductor-pallas":
 - torch/_inductor/codegen/pallas.py
+- torch/utils/_pallas.py
 - test/inductor/test_pallas.py
 - test/inductor/pallas_expected_failures/**
+- test/inductor/pallas_skip_tests/**
 - .github/workflows/inductor-pallas.yml
+- .ci/docker/ci_commit_pins/jax.txt
+- .ci/docker/common/install_jax.sh
+- .ci/docker/common/install_torch_tpu.sh
+- .ci/docker/common/requirements_tpu.txt
+- .github/ci_commit_pins/torch_tpu.txt
 
 "ciflow/b200":
 - test/test_matmul_cuda.py
@@ -184,3 +188,21 @@
 - torch/distributed/_symmetric_memory/**
 - test/distributed/**/*mem*
 - test/distributed/**/*mem*/**
+
+"ciflow/torchtitan":
+# torchtitan commit pin updates
+- .github/ci_commit_pins/torchtitan.txt
+# torch.distributed (FSDP, DTensor, etc.)
+- torch/distributed/**
+# torch.compile
+- torch/_dynamo/**
+- torch/_inductor/**
+# activation checkpointing / selective activation checkpointing
+- torch/utils/checkpoint.py
+- torch/_functorch/_activation_checkpointing/**
+- torch/_functorch/partitioners.py
+- torch/distributed/algorithms/_checkpoint/checkpoint_wrapper.py
+# flexattention
+- torch/nn/attention/flex_attention.py
+- torch/_higher_order_ops/flex_attention.py
+- torch/_inductor/kernel/flex/**
diff --git a/.github/merge_rules.yaml b/.github/merge_rules.yaml
index 59122a044cbe1..a200ba08206f5 100644
--- a/.github/merge_rules.yaml
+++ b/.github/merge_rules.yaml
@@ -38,33 +38,14 @@
   - Lint
   - pull
 
-- name: NVFuser
-  patterns:
-  - test/test_jit_cuda_fuser.py
-  - torch/csrc/jit/codegen/fuser/cuda/**
-  - torch/csrc/jit/codegen/cuda/**
-  - benchmarks/cpp/nvfuser/**
-  approved_by:
-  - csarofeen
-  - ngimel
-  - jjsjann123
-  - kevinstephano
-  - ptrblck
-  mandatory_checks_name:
-  - EasyCLA
-  - Lint
-  - pull
-
 - name: OSS CI
   patterns:
   - .github/**
-  - .circleci/**
   - .ci/**
   - scripts/**
   - tools/**
   approved_by:
   - alband
-  - dagitses
   - pytorch/pytorch-dev-infra
   mandatory_checks_name:
   - EasyCLA
@@ -216,15 +197,17 @@
 - name: MPS
   patterns:
   - test/test_mps.py
+  - test/inductor/test_mps_basic.py
+  - torch/testing/_internal/common_mps.py
   - aten/src/ATen/native/native_functions.yaml
   - aten/src/ATen/mps/**
   - aten/src/ATen/native/mps/**
+  - c10/metal/**
   approved_by:
-  - kulinseth
-  - alband
   - malfet
-  - razarmehr
   - DenisVieriu97
+  - kurtamohler
+
   mandatory_checks_name:
   - EasyCLA
   - Lint
@@ -602,10 +585,12 @@
   - soumith
   - gchanan
   - ezyang
-  - dzhulgakov
   - malfet
   - albanD
   - ptrblck
+  - jspisak
+  - ngimel
+  - Chillee
   mandatory_checks_name:
   - EasyCLA
   - Lint
diff --git a/.github/pytorch-probot.yml b/.github/pytorch-probot.yml
index 1b7771040a588..6f875dcf0716d 100644
--- a/.github/pytorch-probot.yml
+++ b/.github/pytorch-probot.yml
@@ -7,6 +7,7 @@ ciflow_push_tags:
 - ciflow/binaries
 - ciflow/binaries_libtorch
 - ciflow/binaries_wheel
+- ciflow/dtensor
 - ciflow/dynamo
 - ciflow/h100
 - ciflow/h100-cutlass-backend
@@ -26,7 +27,6 @@ ciflow_push_tags:
 - ciflow/inductor-rocm-mi200
 - ciflow/inductor-rocm-mi300
 - ciflow/inductor-rocm-mi355
-- ciflow/linux-aarch64
 - ciflow/mps
 - ciflow/nightly
 - ciflow/op-benchmark
@@ -41,10 +41,12 @@ ciflow_push_tags:
 - ciflow/rocm-mi300
 - ciflow/rocm-mi355
 - ciflow/rocm-navi31
+- ciflow/rocm-nightly
 - ciflow/s390
 - ciflow/slow
 - ciflow/slow-rocm-mi200
 - ciflow/torchbench
+- ciflow/torchtitan
 - ciflow/triton_binaries
 - ciflow/trunk
 - ciflow/unstable
@@ -57,6 +59,8 @@ retryable_workflows:
 - linux-binary
 - windows-binary
 - inductor-A100-perf-nightly
+retryable_step_names:
+- Initialize containers
 labeler_config: labeler.yml
 label_to_label_config: label_to_label.yml
 mergebot: true
diff --git a/.github/requirements-gha-cache.txt b/.github/requirements-gha-cache.txt
index c274ca1e5914d..9bd898687ba67 100644
--- a/.github/requirements-gha-cache.txt
+++ b/.github/requirements-gha-cache.txt
@@ -11,5 +11,5 @@ lintrunner==0.12.7
 ninja==1.10.0.post1
 nvidia-ml-py==11.525.84
 pyyaml==6.0.2
-requests==2.32.4
+requests==2.33.0
 rich==14.1.0
diff --git a/.github/scripts/amd/package_triton_wheel.sh b/.github/scripts/amd/package_triton_wheel.sh
index 501e50e2fe2f1..549bacc5a7573 100755
--- a/.github/scripts/amd/package_triton_wheel.sh
+++ b/.github/scripts/amd/package_triton_wheel.sh
@@ -28,9 +28,7 @@ if [[ -z "${TRITON_ROCM_DIR}" ]]; then
     export TRITON_ROCM_DIR=third_party/amd/backend
 fi
 
-# Remove packaged libs and headers
-rm -rf $TRITON_ROCM_DIR/include/*
-
+# Remove packaged libs
 LIBNUMA_PATH="/usr/lib64/libnuma.so.1"
 LIBELF_PATH="/usr/lib64/libelf.so.1"
 OS_NAME=`awk -F= '/^NAME/{print $2}' /etc/os-release`
diff --git a/.github/scripts/build_triton_wheel.py b/.github/scripts/build_triton_wheel.py
index 34c6c3549f9c7..de0fa7739edc5 100644
--- a/.github/scripts/build_triton_wheel.py
+++ b/.github/scripts/build_triton_wheel.py
@@ -1,12 +1,13 @@
 #!/usr/bin/env python3
 
+from __future__ import annotations
+
 import os
 import shutil
 import sys
 from pathlib import Path
 from subprocess import check_call
 from tempfile import TemporaryDirectory
-from typing import Optional
 
 
 SCRIPT_DIR = Path(__file__).parent
@@ -37,7 +38,7 @@ def check_and_replace(inp: str, src: str, dst: str) -> str:
 
 
 def patch_init_py(
-    path: Path, *, version: str, expected_version: Optional[str] = None
+    path: Path, *, version: str, expected_version: str | None = None
 ) -> None:
     if not expected_version:
         expected_version = read_triton_version()
@@ -56,7 +57,7 @@ def build_triton(
     version: str,
     commit_hash: str,
     device: str = "cuda",
-    py_version: Optional[str] = None,
+    py_version: str | None = None,
     release: bool = False,
     with_clang_ldd: bool = False,
 ) -> Path:
@@ -80,15 +81,23 @@ def build_triton(
         check_call(["git", "clone", triton_repo, "triton"], cwd=tmpdir)
         if release:
             ver, rev, patch = version.split(".")
-            check_call(
-                ["git", "checkout", f"release/{ver}.{rev}.x"], cwd=triton_basedir
-            )
+            if device == "xpu":
+                # XPU uses the patch version in the release branch name
+                check_call(
+                    ["git", "checkout", f"release/{ver}.{rev}.{patch}"],
+                    cwd=triton_basedir,
+                )
+            else:
+                check_call(
+                    ["git", "checkout", f"release/{ver}.{rev}.x"], cwd=triton_basedir
+                )
         else:
             check_call(["git", "fetch", "origin", commit_hash], cwd=triton_basedir)
             check_call(["git", "checkout", commit_hash], cwd=triton_basedir)
 
         # change built wheel name and version
         env["TRITON_WHEEL_NAME"] = triton_pkg_name
+        env["TRITON_EXT_ENABLED"] = "ON"
         if with_clang_ldd:
             env["TRITON_BUILD_WITH_CLANG_LLD"] = "1"
 
diff --git a/.github/scripts/check_doc_redirects.py b/.github/scripts/check_doc_redirects.py
new file mode 100644
index 0000000000000..b6825f42bc53f
--- /dev/null
+++ b/.github/scripts/check_doc_redirects.py
@@ -0,0 +1,252 @@
+#!/usr/bin/env python3
+"""
+Check for moved/deleted doc files and optionally auto-update redirects.py.
+
+This script detects when documentation files in docs/source/ are moved or deleted
+and verifies that corresponding redirects exist in docs/source/redirects.py.
+
+Usage:
+    # Check only (CI mode) - reports missing redirects
+    python check_doc_redirects.py --base-ref origin/main
+
+    # Auto-update redirects.py with missing entries
+    python check_doc_redirects.py --base-ref origin/main --auto-fix
+"""
+
+from __future__ import annotations
+
+import argparse
+import ast
+import re
+import subprocess
+import sys
+from pathlib import Path
+
+
+def run_git(args: list[str]) -> str:
+    """Run a git command and return stdout."""
+    result = subprocess.run(["git"] + args, capture_output=True, text=True)
+    return result.stdout.strip()
+
+
+def get_doc_changes(base_ref: str) -> list[tuple[str, str, str | None]]:
+    """
+    Get moved/deleted doc files between base_ref and HEAD.
+
+    Returns:
+        List of (status, old_path, new_path) tuples.
+        For deletions, new_path is None.
+    """
+    diff = run_git(
+        [
+            "diff",
+            "--name-status",
+            "-M",  # Enable rename detection
+            f"{base_ref}...HEAD",
+            "--",
+            "docs/source/*.rst",
+            "docs/source/*.md",
+        ]
+    )
+
+    changes = []
+    for line in diff.split("\n"):
+        if not line:
+            continue
+        parts = line.split("\t")
+        status = parts[0]
+
+        # Renames: R100 (100% similar), R095 (95% similar), etc.
+        if status.startswith("R") and len(parts) >= 3:
+            changes.append((status, parts[1], parts[2]))
+        # Deletions
+        elif status == "D" and len(parts) >= 2:
+            changes.append((status, parts[1], None))
+
+    return changes
+
+
+def path_to_key(path: str) -> str:
+    """
+    Convert a file path to a redirect key.
+
+    Example: docs/source/torch.compiler.rst -> torch.compiler
+    Example: docs/source/user_guide/foo.rst -> user_guide/foo
+    """
+    return re.sub(r"\.(rst|md)$", "", path.replace("docs/source/", ""))
+
+
+def path_to_url(path: str) -> str:
+    """
+    Convert a file path to an HTML URL for redirects.
+
+    Example: docs/source/user_guide/foo.rst -> user_guide/foo.html
+    """
+    return re.sub(r"\.(rst|md)$", ".html", path.replace("docs/source/", ""))
+
+
+def parse_existing_redirects(redirects_file: Path) -> dict[str, str]:
+    """
+    Parse redirects.py and return the existing redirects dictionary.
+
+    Uses AST parsing for robustness.
+    """
+    content = redirects_file.read_text()
+    tree = ast.parse(content)
+
+    for node in ast.walk(tree):
+        if isinstance(node, ast.Assign):
+            for target in node.targets:
+                if isinstance(target, ast.Name) and target.id == "redirects":
+                    if isinstance(node.value, ast.Dict):
+                        return {
+                            k.value: v.value
+                            for k, v in zip(node.value.keys, node.value.values)
+                            if isinstance(k, ast.Constant)
+                            and isinstance(v, ast.Constant)
+                        }
+    return {}
+
+
+def find_missing_redirects(
+    changes: list[tuple[str, str, str | None]],
+    existing: dict[str, str],
+) -> list[tuple[str, str | None]]:
+    """
+    Find file changes that don't have corresponding redirects.
+
+    Returns:
+        List of (old_key, new_url) tuples. new_url is None for deletions.
+    """
+    missing = []
+    for status, old_path, new_path in changes:
+        old_key = path_to_key(old_path)
+        if old_key not in existing:
+            new_url = path_to_url(new_path) if new_path else None
+            missing.append((old_key, new_url))
+    return missing
+
+
+def update_redirects_file(
+    redirects_file: Path, new_entries: list[tuple[str, str]]
+) -> None:
+    """
+    Add new redirect entries to redirects.py.
+
+    Inserts entries just before the closing brace of the redirects dict.
+    """
+    content = redirects_file.read_text()
+    lines = content.split("\n")
+
+    # Find the line with the closing brace
+    insert_idx = None
+    for i in range(len(lines) - 1, -1, -1):
+        if lines[i].strip() == "}":
+            insert_idx = i
+            break
+
+    if insert_idx is None:
+        print("Error: Could not find closing brace in redirects.py", file=sys.stderr)
+        sys.exit(1)
+
+    # Generate new entries
+    new_lines = ["    # Auto-generated redirects for moved files"]
+    for old_key, new_url in new_entries:
+        new_lines.append(f'    "{old_key}": "{new_url}",')
+
+    # Insert the new entries before the closing brace
+    lines = lines[:insert_idx] + new_lines + lines[insert_idx:]
+
+    redirects_file.write_text("\n".join(lines))
+    print(f"✅ Added {len(new_entries)} redirect(s) to {redirects_file}")
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(
+        description="Check for missing doc redirects and optionally auto-fix"
+    )
+    parser.add_argument(
+        "--base-ref",
+        default="origin/main",
+        help="Base git ref to compare against (default: origin/main)",
+    )
+    parser.add_argument(
+        "--redirects-file",
+        default="docs/source/redirects.py",
+        help="Path to redirects.py (default: docs/source/redirects.py)",
+    )
+    parser.add_argument(
+        "--auto-fix",
+        action="store_true",
+        help="Automatically add missing redirects for moved files",
+    )
+    args = parser.parse_args()
+
+    redirects_file = Path(args.redirects_file)
+    if not redirects_file.exists():
+        print(f"Error: {redirects_file} not found", file=sys.stderr)
+        sys.exit(1)
+
+    # Get file changes
+    changes = get_doc_changes(args.base_ref)
+    if not changes:
+        print("✅ No doc files were moved or deleted")
+        return
+
+    print(f"Found {len(changes)} moved/deleted doc file(s)")
+
+    # Parse existing redirects
+    existing = parse_existing_redirects(redirects_file)
+    print(f"Found {len(existing)} existing redirect(s)")
+
+    # Find missing redirects
+    missing = find_missing_redirects(changes, existing)
+
+    if not missing:
+        print("✅ All moved/deleted doc files have corresponding redirects")
+        return
+
+    # Separate auto-fixable (moves with known destination) from manual (deletes)
+    auto_fixable = [(k, v) for k, v in missing if v is not None]
+    manual_needed = [(k, v) for k, v in missing if v is None]
+
+    # Auto-fix mode
+    if args.auto_fix and auto_fixable:
+        update_redirects_file(redirects_file, auto_fixable)
+
+        if manual_needed:
+            print(f"\n⚠️  {len(manual_needed)} deleted file(s) need manual redirects:")
+            for old_key, _ in manual_needed:
+                print(f"  • {old_key}")
+            print("\nPlease add redirects for deleted files manually.")
+            sys.exit(1)
+        return
+
+    # Report mode - show what's missing
+    print("\n❌ Missing redirects detected!\n")
+
+    for old_key, new_url in missing:
+        if new_url:
+            print(f"  • MOVED: {old_key} → {new_url}")
+        else:
+            print(f"  • DELETED: {old_key} (needs manual redirect target)")
+
+    if auto_fixable:
+        print("\n📝 Suggested additions to docs/source/redirects.py:\n")
+        for old_key, new_url in auto_fixable:
+            print(f'    "{old_key}": "{new_url}",')
+        print("\n💡 To auto-fix, run:")
+        print(
+            f"    python3 .github/scripts/check_doc_redirects.py "
+            f"--base-ref {args.base_ref} --auto-fix"
+        )
+
+    if manual_needed:
+        print(f"\n⚠️  {len(manual_needed)} deleted file(s) need manual redirects.")
+        print("Please determine appropriate redirect targets for deleted files.")
+
+    sys.exit(1)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/.github/scripts/cherry_pick.py b/.github/scripts/cherry_pick.py
index c2776040d81f1..e2b77a25f7fc2 100755
--- a/.github/scripts/cherry_pick.py
+++ b/.github/scripts/cherry_pick.py
@@ -1,9 +1,11 @@
 #!/usr/bin/env python3
 
+from __future__ import annotations
+
 import json
 import os
 import re
-from typing import Any, cast, Optional
+from typing import Any, cast
 from urllib.error import HTTPError
 
 from github_utils import gh_fetch_url, gh_post_pr_comment, gh_query_issues_by_labels
@@ -48,7 +50,7 @@ def parse_args() -> Any:
     return parser.parse_args()
 
 
-def get_merge_commit_sha(repo: GitRepo, pr: GitHubPR) -> Optional[str]:
+def get_merge_commit_sha(repo: GitRepo, pr: GitHubPR) -> str | None:
     """
     Return the merge commit SHA iff the PR has been merged. For simplicity, we
     will only cherry pick PRs that have been merged into main
@@ -57,7 +59,7 @@ def get_merge_commit_sha(repo: GitRepo, pr: GitHubPR) -> Optional[str]:
     return commit_sha if pr.is_closed() else None
 
 
-def get_release_version(onto_branch: str) -> Optional[str]:
+def get_release_version(onto_branch: str) -> str | None:
     """
     Return the release version if the target branch is a release branch
     """
diff --git a/.github/scripts/convert_lintrunner_annotations_to_github.py b/.github/scripts/convert_lintrunner_annotations_to_github.py
index 884a8205f53c4..57f54561ecc0f 100644
--- a/.github/scripts/convert_lintrunner_annotations_to_github.py
+++ b/.github/scripts/convert_lintrunner_annotations_to_github.py
@@ -1,9 +1,11 @@
+from __future__ import annotations
+
 import json
 import subprocess
 import sys
 from enum import Enum
 from pathlib import Path
-from typing import NamedTuple, Optional
+from typing import NamedTuple
 
 
 # From: https://docs.github.com/en/rest/reference/checks
@@ -17,12 +19,12 @@ class GitHubAnnotation(NamedTuple):
     path: str
     start_line: int
     end_line: int
-    start_column: Optional[int]
-    end_column: Optional[int]
+    start_column: int | None
+    end_column: int | None
     annotation_level: GitHubAnnotationLevel
     message: str
-    title: Optional[str]
-    raw_details: Optional[str]
+    title: str | None
+    raw_details: str | None
 
 
 PYTORCH_ROOT = Path(
diff --git a/.github/scripts/filter_test_configs.py b/.github/scripts/filter_test_configs.py
index 182a75f13cad6..0dbbb4aa3761b 100755
--- a/.github/scripts/filter_test_configs.py
+++ b/.github/scripts/filter_test_configs.py
@@ -1,5 +1,6 @@
 #!/usr/bin/env python3
 # ruff: noqa: LOG015
+from __future__ import annotations
 
 import json
 import logging
@@ -8,16 +9,19 @@
 import subprocess
 import sys
 import warnings
-from collections.abc import Callable
 from enum import Enum
 from functools import cache
 from logging import info
-from typing import Any, Optional
+from typing import Any, TYPE_CHECKING
 from urllib.request import Request, urlopen
 
 import yaml
 
 
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+
 REENABLE_TEST_REGEX = "(?i)(Close(d|s)?|Resolve(d|s)?|Fix(ed|es)?) (#|https://github.com/pytorch/pytorch/issues/)([0-9]+)"
 MAIN_BRANCH = "main"
 
@@ -26,25 +30,37 @@
 logging.basicConfig(level=logging.INFO)
 
 
-def is_cuda_or_rocm_job(job_name: Optional[str]) -> bool:
-    if not job_name:
-        return False
+def is_cuda_or_rocm_job(
+    job_name: str | None, config: dict[str, Any] | None = None
+) -> bool:
+    if job_name and ("cuda" in job_name or "rocm" in job_name):
+        return True
+
+    # Also check the runner name in the config, since some workflows (e.g.
+    # inductor-unittest) use job names that don't include "cuda" even though
+    # they target CUDA runners.
+    if config:
+        runner = config.get("runner", "")
+        if "nvidia.gpu" in runner or "rocm.gpu" in runner:
+            return True
 
-    return "cuda" in job_name or "rocm" in job_name
+    return False
 
 
 # Supported modes when running periodically. Only applying the mode when
-# its lambda condition returns true
-SUPPORTED_PERIODICAL_MODES: dict[str, Callable[[Optional[str]], bool]] = {
+# its lambda condition returns true. Each callable receives (job_name, config).
+SUPPORTED_PERIODICAL_MODES: dict[
+    str, Callable[[str | None, dict[str, Any] | None], bool]
+] = {
     # Memory leak check is only needed for CUDA and ROCm jobs which utilize GPU memory
     "mem_leak_check": is_cuda_or_rocm_job,
-    "rerun_disabled_tests": lambda job_name: True,
+    "rerun_disabled_tests": lambda job_name, config=None: True,
 }
 
 # The link to the published list of disabled jobs
-DISABLED_JOBS_URL = "https://ossci-metrics.s3.amazonaws.com/disabled-jobs.json"
+DISABLED_JOBS_URL = "https://ossci-metrics.s3.amazonaws.com/disabled-jobs.json?versionId=wTGHVmBBfqYt9mKCiXeWHzek1ZaYHJeH"
 # and unstable jobs
-UNSTABLE_JOBS_URL = "https://ossci-metrics.s3.amazonaws.com/unstable-jobs.json"
+UNSTABLE_JOBS_URL = "https://ossci-metrics.s3.amazonaws.com/unstable-jobs.json?versionId=6zhdrg_i6w2bx_KfIfy8SKn5WdSO0ZAf"
 
 # Some constants used to handle disabled and unstable jobs
 JOB_NAME_SEP = "/"
@@ -210,7 +226,7 @@ def filter_selected_test_configs(
 
 
 def set_periodic_modes(
-    test_matrix: dict[str, list[Any]], job_name: Optional[str]
+    test_matrix: dict[str, list[Any]], job_name: str | None
 ) -> dict[str, list[Any]]:
     """
     Apply all periodic modes when running under a schedule
@@ -221,7 +237,7 @@ def set_periodic_modes(
 
     for config in test_matrix.get("include", []):
         for mode, cond in SUPPORTED_PERIODICAL_MODES.items():
-            if not cond(job_name):
+            if not cond(job_name, config):
                 continue
 
             cfg = config.copy()
@@ -266,7 +282,7 @@ def remove_disabled_jobs(
 def _filter_jobs(
     test_matrix: dict[str, list[Any]],
     issue_type: IssueType,
-    target_cfg: Optional[str] = None,
+    target_cfg: str | None = None,
 ) -> dict[str, list[Any]]:
     """
     An utility function used to actually apply the job filter
@@ -466,7 +482,7 @@ def set_output(name: str, val: Any) -> None:
         print(f"::set-output name={name}::{val}")
 
 
-def parse_reenabled_issues(s: Optional[str]) -> list[str]:
+def parse_reenabled_issues(s: str | None) -> list[str]:
     # NB: When the PR body is empty, GitHub API returns a None value, which is
     # passed into this function
     if not s:
@@ -502,8 +518,8 @@ def perform_misc_tasks(
     test_matrix: dict[str, list[Any]],
     job_name: str,
     pr_body: str,
-    branch: Optional[str] = None,
-    tag: Optional[str] = None,
+    branch: str | None = None,
+    tag: str | None = None,
 ) -> None:
     """
     In addition to apply the filter logic, the script also does the following
diff --git a/.github/scripts/generate_binary_build_matrix.py b/.github/scripts/generate_binary_build_matrix.py
index 267afbe216fa4..aaf19e57b8d8d 100644
--- a/.github/scripts/generate_binary_build_matrix.py
+++ b/.github/scripts/generate_binary_build_matrix.py
@@ -11,30 +11,29 @@
     * Latest XPU
 """
 
+from __future__ import annotations
+
 import json
 import os
 import re
 from pathlib import Path
-from typing import Optional
 
 
 SCRIPT_DIR = Path(__file__).absolute().parent
 REPO_ROOT = SCRIPT_DIR.parent.parent
 
 
-CUDA_ARCHES = ["12.6", "12.8", "12.9", "13.0"]
-CUDA_STABLE = "12.8"
+CUDA_ARCHES = ["12.6", "13.0", "13.2"]
+CUDA_STABLE = "13.0"
 CUDA_ARCHES_FULL_VERSION = {
     "12.6": "12.6.3",
-    "12.8": "12.8.1",
-    "12.9": "12.9.1",
     "13.0": "13.0.2",
+    "13.2": "13.2.1",
 }
 CUDA_ARCHES_CUDNN_VERSION = {
     "12.6": "9",
-    "12.8": "9",
-    "12.9": "9",
     "13.0": "9",
+    "13.2": "9",
 }
 
 ROCM_ARCHES = ["7.1", "7.2"]
@@ -45,40 +44,40 @@
 
 CPU_S390X_ARCH = ["cpu-s390x"]
 
-CUDA_AARCH64_ARCHES = ["12.6-aarch64", "12.8-aarch64", "12.9-aarch64", "13.0-aarch64"]
+CUDA_AARCH64_ARCHES = [
+    "12.6-aarch64",
+    "13.0-aarch64",
+    "13.2-aarch64",
+]
 
 
+# WARNING: For CUDA 13.0, cublas is pinned to a version range rather
+# than an exact version. A broken cublas release within that range will be
+# silently pulled in.
 PYTORCH_EXTRA_INSTALL_REQUIREMENTS = {
     "12.6": (
         "cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cublas,cufile,nvjitlink,nvtx]==12.6.3; platform_system == 'Linux' | "  # noqa: B950
-        "cuda-bindings==12.9.4; platform_system == 'Linux' | "
+        "cuda-bindings>=12.9.4,<13; platform_system == 'Linux' | "
         "nvidia-cudnn-cu12==9.10.2.21; platform_system == 'Linux' | "
         "nvidia-cusparselt-cu12==0.7.1; platform_system == 'Linux' | "
-        "nvidia-nccl-cu12==2.28.9; platform_system == 'Linux' | "
-        "nvidia-nvshmem-cu12==3.4.5; platform_system == 'Linux'"
-    ),
-    "12.8": (
-        "cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cublas,cufile,nvjitlink,nvtx]==12.8.1; platform_system == 'Linux' | "  # noqa: B950
-        "cuda-bindings==12.9.4; platform_system == 'Linux' | "
-        "nvidia-cudnn-cu12==9.17.1.4; platform_system == 'Linux' | "
-        "nvidia-cusparselt-cu12==0.7.1; platform_system == 'Linux' | "
-        "nvidia-nccl-cu12==2.28.9; platform_system == 'Linux' | "
-        "nvidia-nvshmem-cu12==3.4.5; platform_system == 'Linux'"
-    ),
-    "12.9": (
-        "cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cublas,cufile,nvjitlink,nvtx]==12.9.1; platform_system == 'Linux' | "  # noqa: B950
-        "cuda-bindings==12.9.4; platform_system == 'Linux' | "
-        "nvidia-cudnn-cu12==9.17.1.4; platform_system == 'Linux' | "
-        "nvidia-cusparselt-cu12==0.7.1; platform_system == 'Linux' | "
-        "nvidia-nccl-cu12==2.28.9; platform_system == 'Linux' | "
+        "nvidia-nccl-cu12==2.29.3; platform_system == 'Linux' | "
         "nvidia-nvshmem-cu12==3.4.5; platform_system == 'Linux'"
     ),
     "13.0": (
-        "cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cublas,cufile,nvjitlink,nvtx]==13.0.2; platform_system == 'Linux' | "  # noqa: B950
-        "cuda-bindings==13.0.3; platform_system == 'Linux' | "
-        "nvidia-cudnn-cu13==9.17.1.4; platform_system == 'Linux' | "
-        "nvidia-cusparselt-cu13==0.8.0; platform_system == 'Linux' | "
-        "nvidia-nccl-cu13==2.28.9; platform_system == 'Linux' | "
+        "cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cufile,nvjitlink,nvtx]==13.0.2; platform_system == 'Linux' | "  # noqa: B950
+        "nvidia-cublas>=13.1.0.3,<=13.1.1.3; platform_system == 'Linux' | "
+        "cuda-bindings>=13.0.3,<14; platform_system == 'Linux' | "
+        "nvidia-cudnn-cu13==9.20.0.48; platform_system == 'Linux' | "
+        "nvidia-cusparselt-cu13==0.8.1; platform_system == 'Linux' | "
+        "nvidia-nccl-cu13==2.29.7; platform_system == 'Linux' | "
+        "nvidia-nvshmem-cu13==3.4.5; platform_system == 'Linux'"
+    ),
+    "13.2": (
+        "cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cublas,cufile,nvjitlink,nvtx]==13.2.1; platform_system == 'Linux' | "  # noqa: B950
+        "cuda-bindings>=13.0.3,<14; platform_system == 'Linux' | "
+        "nvidia-cudnn-cu13==9.20.0.48; platform_system == 'Linux' | "
+        "nvidia-cusparselt-cu13==0.8.1; platform_system == 'Linux' | "
+        "nvidia-nccl-cu13==2.29.7; platform_system == 'Linux' | "
         "nvidia-nvshmem-cu13==3.4.5; platform_system == 'Linux'"
     ),
     "xpu": (
@@ -162,7 +161,7 @@ def read_nccl_pin(arch_version: str) -> str:
     # Single source of truth for NCCL version
     from optional_submodules import read_nccl_pin
 
-    return read_nccl_pin()
+    return read_nccl_pin(arch_version)
 
 
 def validate_nccl_dep_consistency(arch_version: str) -> None:
@@ -175,6 +174,67 @@ def validate_nccl_dep_consistency(arch_version: str) -> None:
         )
 
 
+def _parse_linux_cudnn_versions() -> dict[str, str]:
+    """Return {cuda_short_version: cudnn_version} from install_cuda.sh."""
+    text = (REPO_ROOT / ".ci" / "docker" / "common" / "install_cuda.sh").read_text()
+    results: dict[str, str] = {}
+    func_re = re.compile(r"^function install_(\d+)\s*\{")
+    cudnn_re = re.compile(r"^\s*CUDNN_VERSION=(\S+)")
+    current_func: str | None = None
+    for line in text.splitlines():
+        m = func_re.match(line)
+        if m:
+            digits = m.group(1)
+            current_func = digits[:-1] + "." + digits[-1]
+            continue
+        if current_func is not None:
+            m = cudnn_re.match(line)
+            if m:
+                results[current_func] = m.group(1)
+                current_func = None
+    return results
+
+
+def _parse_windows_cudnn_versions() -> dict[str, str]:
+    """Return {cuda_short_version: cudnn_version} from cuda_install.bat."""
+    text = (
+        REPO_ROOT / ".ci" / "pytorch" / "windows" / "internal" / "cuda_install.bat"
+    ).read_text()
+    results: dict[str, str] = {}
+    label_re = re.compile(r"^:cuda(\d+)\s*$")
+    cudnn_re = re.compile(
+        r"^set CUDNN_FOLDER=cudnn-windows-x86_64-([0-9.]+)_cuda\d+-archive"
+    )
+    current_label: str | None = None
+    for line in text.splitlines():
+        m = label_re.match(line)
+        if m:
+            digits = m.group(1)
+            current_label = digits[:-1] + "." + digits[-1]
+            continue
+        if current_label is not None:
+            m = cudnn_re.match(line)
+            if m:
+                results[current_label] = m.group(1)
+                current_label = None
+    return results
+
+
+def validate_cudnn_version_consistency(arch_version: str) -> None:
+    linux_versions = _parse_linux_cudnn_versions()
+    windows_versions = _parse_windows_cudnn_versions()
+    linux_ver = linux_versions.get(arch_version)
+    windows_ver = windows_versions.get(arch_version)
+    if linux_ver is None or windows_ver is None:
+        return
+    if linux_ver != windows_ver:
+        raise RuntimeError(
+            f"cuDNN version mismatch for CUDA {arch_version}: "
+            f"Linux has {linux_ver} (.ci/docker/common/install_cuda.sh) "
+            f"but Windows has {windows_ver} (.ci/pytorch/windows/internal/cuda_install.bat)"
+        )
+
+
 def arch_type(arch_version: str) -> str:
     if arch_version in CUDA_ARCHES:
         return "cuda"
@@ -210,12 +270,6 @@ def arch_type(arch_version: str) -> str:
 RELEASE = "release"
 DEBUG = "debug"
 
-LIBTORCH_CONTAINER_IMAGES: dict[str, str] = {
-    **{gpu_arch: f"libtorch-cxx11-builder:cuda{gpu_arch}" for gpu_arch in CUDA_ARCHES},
-    **{gpu_arch: f"libtorch-cxx11-builder:rocm{gpu_arch}" for gpu_arch in ROCM_ARCHES},
-    "cpu": "libtorch-cxx11-builder:cpu",
-}
-
 FULL_PYTHON_VERSIONS = ["3.10", "3.11", "3.12", "3.13", "3.13t", "3.14", "3.14t"]
 
 
@@ -238,20 +292,13 @@ def list_without(in_list: list[str], without: list[str]) -> list[str]:
 def generate_libtorch_matrix(
     os: str,
     release_type: str,
-    arches: Optional[list[str]] = None,
-    libtorch_variants: Optional[list[str]] = None,
+    arches: list[str] | None = None,
+    libtorch_variants: list[str] | None = None,
 ) -> list[dict[str, str]]:
     if arches is None:
         arches = ["cpu"]
-        if os == "linux":
+        if os == "windows":
             arches += CUDA_ARCHES
-            arches += ROCM_ARCHES
-        elif os == "windows":
-            # TODO (huydhn): Only build CUDA 12.9 for Linux. This logic is to be cleaned up
-            # in 2.10
-            windows_cuda_arches = CUDA_ARCHES.copy()
-            windows_cuda_arches.remove("12.9")
-            arches += windows_cuda_arches
     if libtorch_variants is None:
         libtorch_variants = [
             "shared-with-deps",
@@ -265,9 +312,6 @@ def generate_libtorch_matrix(
         for libtorch_variant in libtorch_variants:
             gpu_arch_type = arch_type(arch_version)
             gpu_arch_version = "" if arch_version == "cpu" else arch_version
-            # ROCm builds without-deps failed even in ROCm runners; skip for now
-            if gpu_arch_type == "rocm" and ("without-deps" in libtorch_variant):
-                continue
             ret.append(
                 {
                     "gpu_arch_type": gpu_arch_type,
@@ -277,16 +321,8 @@ def generate_libtorch_matrix(
                     ),
                     "libtorch_config": release_type,
                     "libtorch_variant": libtorch_variant,
-                    "container_image": (
-                        LIBTORCH_CONTAINER_IMAGES[arch_version].split(":")[0]
-                        if os not in ("windows", "windows-arm64")
-                        else ""
-                    ),
-                    "container_image_tag_prefix": (
-                        LIBTORCH_CONTAINER_IMAGES[arch_version].split(":")[1]
-                        if os not in ("windows", "windows-arm64")
-                        else ""
-                    ),
+                    "container_image": "",
+                    "container_image_tag_prefix": "",
                     "package_type": "libtorch",
                     "build_name": f"libtorch-{gpu_arch_type}{gpu_arch_version}-{libtorch_variant}-{release_type}".replace(
                         ".", "_"
@@ -298,8 +334,8 @@ def generate_libtorch_matrix(
 
 def generate_wheels_matrix(
     os: str,
-    arches: Optional[list[str]] = None,
-    python_versions: Optional[list[str]] = None,
+    arches: list[str] | None = None,
+    python_versions: list[str] | None = None,
 ) -> list[dict[str, str]]:
     package_type = "wheel"
     if os == "linux" or os == "linux-aarch64" or os == "linux-s390x":
@@ -315,11 +351,7 @@ def generate_wheels_matrix(
         if os == "linux":
             arches += CUDA_ARCHES + ROCM_ARCHES + XPU_ARCHES
         elif os == "windows":
-            # TODO (huydhn): Only build CUDA 12.9 for Linux. This logic is to be cleaned up
-            # in 2.10
-            windows_cuda_arches = CUDA_ARCHES.copy()
-            windows_cuda_arches.remove("12.9")
-            arches += windows_cuda_arches + XPU_ARCHES
+            arches += CUDA_ARCHES + XPU_ARCHES
         elif os == "linux-aarch64":
             # Separate new if as the CPU type is different and
             # uses different build/test scripts
@@ -355,7 +387,7 @@ def generate_wheels_matrix(
             # cuda linux wheels require PYTORCH_EXTRA_INSTALL_REQUIREMENTS to install
 
             if (
-                arch_version in ["13.0", "12.9", "12.8", "12.6"]
+                arch_version in ["13.2", "13.0", "12.6"]
                 and os == "linux"
                 or arch_version in CUDA_AARCH64_ARCHES
             ):
@@ -410,6 +442,9 @@ def generate_wheels_matrix(
                         "pytorch_extra_install_requirements": (
                             PYTORCH_EXTRA_INSTALL_REQUIREMENTS["xpu"]
                             if gpu_arch_type == "xpu"
+                            else PYTORCH_EXTRA_INSTALL_REQUIREMENTS[CUDA_STABLE]
+                            if gpu_arch_type == "cpu"
+                            and os in ("windows", "macos-arm64")
                             else ""
                         ),
                     }
@@ -418,9 +453,57 @@ def generate_wheels_matrix(
     return ret
 
 
+def generate_libtorch_extraction_configs(
+    os: str,
+    wheel_configs: list[dict[str, str]],
+) -> list[dict[str, str]]:
+    """Generate libtorch extraction configs from existing wheel build configs.
+
+    For each unique arch variant in wheel_configs, find the py3.10 config
+    (py3.11 for windows-arm64) and produce a config that the CI template
+    uses to add an extraction job that depends on that wheel's build job.
+    """
+    preferred_python = "3.11" if os == "windows-arm64" else "3.10"
+
+    # Group wheel configs by (gpu_arch_type, gpu_arch_version)
+    arch_to_config: dict[tuple[str, str], dict[str, str]] = {}
+    for config in wheel_configs:
+        key = (config["gpu_arch_type"], config.get("gpu_arch_version", ""))
+        if config.get("python_version") == preferred_python:
+            arch_to_config[key] = config
+
+    ret: list[dict[str, str]] = []
+    for (gpu_arch_type, gpu_arch_version), source_config in arch_to_config.items():
+        # No libtorch for XPU
+        if gpu_arch_type == "xpu":
+            continue
+
+        desired_cuda = source_config["desired_cuda"]
+        libtorch_variant = "shared-with-deps"
+        build_name = f"libtorch-{gpu_arch_type}{gpu_arch_version}-{libtorch_variant}-release".replace(
+            ".", "_"
+        )
+
+        ret.append(
+            {
+                "source_wheel_build_name": source_config["build_name"],
+                "build_name": build_name,
+                "package_type": "libtorch",
+                "libtorch_variant": libtorch_variant,
+                "libtorch_config": RELEASE,
+                "desired_cuda": desired_cuda,
+                "gpu_arch_type": gpu_arch_type,
+                "gpu_arch_version": gpu_arch_version,
+            }
+        )
+
+    return ret
+
+
 arch_version = ""
 for arch_version in CUDA_ARCHES:
     validate_nccl_dep_consistency(arch_version)
+    validate_cudnn_version_consistency(arch_version)
 del arch_version
 
 
diff --git a/.github/scripts/generate_ci_workflows.py b/.github/scripts/generate_ci_workflows.py
index 1fc2eb59232c5..b1e33b6038ec7 100755
--- a/.github/scripts/generate_ci_workflows.py
+++ b/.github/scripts/generate_ci_workflows.py
@@ -62,6 +62,9 @@ class BinaryBuildWorkflow:
     macos_runner: str = "macos-14-xlarge"
     # Mainly used for libtorch builds
     build_variant: str = ""
+    # Libtorch extraction configs: lightweight jobs that extract libtorch
+    # from a wheel build instead of building libtorch from scratch
+    libtorch_extraction_configs: list[dict[str, str]] = field(default_factory=list)
 
     def __post_init__(self) -> None:
         if self.build_environment == "":
@@ -100,57 +103,55 @@ class OperatingSystem:
     LINUX_S390X = "linux-s390x"
 
 
+_LINUX_WHEEL_CONFIGS = generate_binary_build_matrix.generate_wheels_matrix(
+    OperatingSystem.LINUX
+)
+
 LINUX_BINARY_BUILD_WORFKLOWS = [
     BinaryBuildWorkflow(
         os=OperatingSystem.LINUX,
         package_type="manywheel",
-        build_configs=generate_binary_build_matrix.generate_wheels_matrix(
-            OperatingSystem.LINUX
-        ),
+        build_configs=_LINUX_WHEEL_CONFIGS,
         ciflow_config=CIFlowConfig(
-            labels={LABEL_CIFLOW_BINARIES, LABEL_CIFLOW_BINARIES_WHEEL},
+            labels={
+                LABEL_CIFLOW_BINARIES,
+                LABEL_CIFLOW_BINARIES_WHEEL,
+                LABEL_CIFLOW_BINARIES_LIBTORCH,
+            },
             isolated_workflow=True,
         ),
-    ),
-    BinaryBuildWorkflow(
-        os=OperatingSystem.LINUX,
-        package_type="libtorch",
-        build_configs=generate_binary_build_matrix.generate_libtorch_matrix(
+        libtorch_extraction_configs=generate_binary_build_matrix.generate_libtorch_extraction_configs(
             OperatingSystem.LINUX,
-            generate_binary_build_matrix.RELEASE,
-            libtorch_variants=["shared-with-deps"],
-        ),
-        ciflow_config=CIFlowConfig(
-            labels={LABEL_CIFLOW_BINARIES, LABEL_CIFLOW_BINARIES_LIBTORCH},
-            isolated_workflow=True,
+            _LINUX_WHEEL_CONFIGS,
         ),
     ),
 ]
 
+_WINDOWS_WHEEL_CONFIGS = generate_binary_build_matrix.generate_wheels_matrix(
+    OperatingSystem.WINDOWS
+)
+_WINDOWS_ARM64_WHEEL_CONFIGS = generate_binary_build_matrix.generate_wheels_matrix(
+    OperatingSystem.WINDOWS_ARM64,
+    arches=["cpu"],
+    python_versions=["3.11", "3.12", "3.13"],
+)
+
 WINDOWS_BINARY_BUILD_WORKFLOWS = [
     BinaryBuildWorkflow(
         os=OperatingSystem.WINDOWS,
         package_type="wheel",
-        build_configs=generate_binary_build_matrix.generate_wheels_matrix(
-            OperatingSystem.WINDOWS
-        ),
+        build_configs=_WINDOWS_WHEEL_CONFIGS,
         ciflow_config=CIFlowConfig(
-            labels={LABEL_CIFLOW_BINARIES, LABEL_CIFLOW_BINARIES_WHEEL},
+            labels={
+                LABEL_CIFLOW_BINARIES,
+                LABEL_CIFLOW_BINARIES_WHEEL,
+                LABEL_CIFLOW_BINARIES_LIBTORCH,
+            },
             isolated_workflow=True,
         ),
-    ),
-    BinaryBuildWorkflow(
-        os=OperatingSystem.WINDOWS,
-        package_type="libtorch",
-        build_variant=generate_binary_build_matrix.RELEASE,
-        build_configs=generate_binary_build_matrix.generate_libtorch_matrix(
+        libtorch_extraction_configs=generate_binary_build_matrix.generate_libtorch_extraction_configs(
             OperatingSystem.WINDOWS,
-            generate_binary_build_matrix.RELEASE,
-            libtorch_variants=["shared-with-deps"],
-        ),
-        ciflow_config=CIFlowConfig(
-            labels={LABEL_CIFLOW_BINARIES, LABEL_CIFLOW_BINARIES_LIBTORCH},
-            isolated_workflow=True,
+            _WINDOWS_WHEEL_CONFIGS,
         ),
     ),
     BinaryBuildWorkflow(
@@ -170,29 +171,18 @@ class OperatingSystem:
     BinaryBuildWorkflow(
         os=OperatingSystem.WINDOWS_ARM64,
         package_type="wheel",
-        build_configs=generate_binary_build_matrix.generate_wheels_matrix(
-            OperatingSystem.WINDOWS_ARM64,
-            arches=["cpu"],
-            python_versions=["3.11", "3.12", "3.13"],
-        ),
+        build_configs=_WINDOWS_ARM64_WHEEL_CONFIGS,
         ciflow_config=CIFlowConfig(
-            labels={LABEL_CIFLOW_BINARIES, LABEL_CIFLOW_BINARIES_WHEEL},
+            labels={
+                LABEL_CIFLOW_BINARIES,
+                LABEL_CIFLOW_BINARIES_WHEEL,
+                LABEL_CIFLOW_BINARIES_LIBTORCH,
+            },
             isolated_workflow=True,
         ),
-    ),
-    BinaryBuildWorkflow(
-        os=OperatingSystem.WINDOWS_ARM64,
-        package_type="libtorch",
-        build_variant=generate_binary_build_matrix.RELEASE,
-        build_configs=generate_binary_build_matrix.generate_libtorch_matrix(
+        libtorch_extraction_configs=generate_binary_build_matrix.generate_libtorch_extraction_configs(
             OperatingSystem.WINDOWS_ARM64,
-            generate_binary_build_matrix.RELEASE,
-            arches=["cpu"],
-            libtorch_variants=["shared-with-deps"],
-        ),
-        ciflow_config=CIFlowConfig(
-            labels={LABEL_CIFLOW_BINARIES, LABEL_CIFLOW_BINARIES_LIBTORCH},
-            isolated_workflow=True,
+            _WINDOWS_ARM64_WHEEL_CONFIGS,
         ),
     ),
     BinaryBuildWorkflow(
@@ -212,33 +202,28 @@ class OperatingSystem:
     ),
 ]
 
+_MACOS_ARM64_WHEEL_CONFIGS = generate_binary_build_matrix.generate_wheels_matrix(
+    OperatingSystem.MACOS_ARM64
+)
+
 MACOS_BINARY_BUILD_WORKFLOWS = [
-    BinaryBuildWorkflow(
-        os=OperatingSystem.MACOS_ARM64,
-        package_type="libtorch",
-        build_variant=generate_binary_build_matrix.RELEASE,
-        build_configs=generate_binary_build_matrix.generate_libtorch_matrix(
-            OperatingSystem.MACOS,
-            generate_binary_build_matrix.RELEASE,
-            libtorch_variants=["shared-with-deps"],
-        ),
-        macos_runner="macos-14-xlarge",
-        ciflow_config=CIFlowConfig(
-            labels={LABEL_CIFLOW_BINARIES, LABEL_CIFLOW_BINARIES_LIBTORCH},
-            isolated_workflow=True,
-        ),
-    ),
     BinaryBuildWorkflow(
         os=OperatingSystem.MACOS_ARM64,
         package_type="wheel",
-        build_configs=generate_binary_build_matrix.generate_wheels_matrix(
-            OperatingSystem.MACOS_ARM64
-        ),
-        macos_runner="macos-14-xlarge",
+        build_configs=_MACOS_ARM64_WHEEL_CONFIGS,
+        macos_runner="macos-26-xlarge",
         ciflow_config=CIFlowConfig(
-            labels={LABEL_CIFLOW_BINARIES, LABEL_CIFLOW_BINARIES_WHEEL},
+            labels={
+                LABEL_CIFLOW_BINARIES,
+                LABEL_CIFLOW_BINARIES_WHEEL,
+                LABEL_CIFLOW_BINARIES_LIBTORCH,
+            },
             isolated_workflow=True,
         ),
+        libtorch_extraction_configs=generate_binary_build_matrix.generate_libtorch_extraction_configs(
+            OperatingSystem.MACOS_ARM64,
+            _MACOS_ARM64_WHEEL_CONFIGS,
+        ),
     ),
 ]
 
diff --git a/.github/scripts/get_workflow_job_id.py b/.github/scripts/get_workflow_job_id.py
index 3a1e974265998..51258b9868e3e 100644
--- a/.github/scripts/get_workflow_job_id.py
+++ b/.github/scripts/get_workflow_job_id.py
@@ -1,6 +1,7 @@
 # Helper to get the id of the currently running job in a GitHub Actions
 # workflow. GitHub does not provide this information to workflow runs, so we
 # need to figure it out based on what they *do* provide.
+from __future__ import annotations
 
 import argparse
 import json
@@ -11,11 +12,14 @@
 import time
 import urllib
 import urllib.parse
-from collections.abc import Callable
-from typing import Any, Optional
+from typing import Any, TYPE_CHECKING
 from urllib.request import Request, urlopen
 
 
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+
 def parse_json_and_links(conn: Any) -> tuple[Any, dict[str, dict[str, str]]]:
     links = {}
     # Extract links which GH uses for pagination
@@ -43,9 +47,9 @@ def parse_json_and_links(conn: Any) -> tuple[Any, dict[str, dict[str, str]]]:
 def fetch_url(
     url: str,
     *,
-    headers: Optional[dict[str, str]] = None,
+    headers: dict[str, str] | None = None,
     reader: Callable[[Any], Any] = lambda x: x.read(),
-    retries: Optional[int] = 3,
+    retries: int | None = 3,
     backoff_timeout: float = 0.5,
 ) -> Any:
     if headers is None:
diff --git a/.github/scripts/github_utils.py b/.github/scripts/github_utils.py
index 6479fb64ddbaf..85a3243c0db64 100644
--- a/.github/scripts/github_utils.py
+++ b/.github/scripts/github_utils.py
@@ -1,16 +1,21 @@
 """GitHub Utilities"""
 
+from __future__ import annotations
+
 import json
 import os
 import warnings
-from collections.abc import Callable
 from dataclasses import dataclass
-from typing import Any, cast, Optional, Union
+from typing import Any, cast, TYPE_CHECKING
 from urllib.error import HTTPError
 from urllib.parse import quote
 from urllib.request import Request, urlopen
 
 
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+
 GITHUB_API_URL = "https://api.github.com"
 
 
@@ -19,9 +24,9 @@ class GitHubComment:
     body_text: str
     created_at: str
     author_login: str
-    author_url: Optional[str]
+    author_url: str | None
     author_association: str
-    editor_login: Optional[str]
+    editor_login: str | None
     database_id: int
     url: str
 
@@ -29,9 +34,9 @@ class GitHubComment:
 def gh_fetch_url_and_headers(
     url: str,
     *,
-    headers: Optional[dict[str, str]] = None,
-    data: Union[Optional[dict[str, Any]], str] = None,
-    method: Optional[str] = None,
+    headers: dict[str, str] | None = None,
+    data: dict[str, Any] | None | str = None,
+    method: str | None = None,
     reader: Callable[[Any], Any] = lambda x: x.read(),
 ) -> tuple[Any, Any]:
     if headers is None:
@@ -72,9 +77,9 @@ def gh_fetch_url_and_headers(
 def gh_fetch_url(
     url: str,
     *,
-    headers: Optional[dict[str, str]] = None,
-    data: Union[Optional[dict[str, Any]], str] = None,
-    method: Optional[str] = None,
+    headers: dict[str, str] | None = None,
+    data: dict[str, Any] | None | str = None,
+    method: str | None = None,
     reader: Callable[[Any], Any] = json.load,
 ) -> Any:
     return gh_fetch_url_and_headers(
@@ -84,9 +89,9 @@ def gh_fetch_url(
 
 def gh_fetch_json(
     url: str,
-    params: Optional[dict[str, Any]] = None,
-    data: Optional[dict[str, Any]] = None,
-    method: Optional[str] = None,
+    params: dict[str, Any] | None = None,
+    data: dict[str, Any] | None = None,
+    method: str | None = None,
 ) -> list[dict[str, Any]]:
     headers = {"Accept": "application/vnd.github.v3+json"}
     if params is not None and len(params) > 0:
@@ -101,8 +106,8 @@ def gh_fetch_json(
 
 def _gh_fetch_json_any(
     url: str,
-    params: Optional[dict[str, Any]] = None,
-    data: Optional[dict[str, Any]] = None,
+    params: dict[str, Any] | None = None,
+    data: dict[str, Any] | None = None,
 ) -> Any:
     headers = {"Accept": "application/vnd.github.v3+json"}
     if params is not None and len(params) > 0:
@@ -114,16 +119,16 @@ def _gh_fetch_json_any(
 
 def gh_fetch_json_list(
     url: str,
-    params: Optional[dict[str, Any]] = None,
-    data: Optional[dict[str, Any]] = None,
+    params: dict[str, Any] | None = None,
+    data: dict[str, Any] | None = None,
 ) -> list[dict[str, Any]]:
     return cast(list[dict[str, Any]], _gh_fetch_json_any(url, params, data))
 
 
 def gh_fetch_json_dict(
     url: str,
-    params: Optional[dict[str, Any]] = None,
-    data: Optional[dict[str, Any]] = None,
+    params: dict[str, Any] | None = None,
+    data: dict[str, Any] | None = None,
 ) -> dict[str, Any]:
     return cast(dict[str, Any], _gh_fetch_json_any(url, params, data))
 
diff --git a/.github/scripts/gitutils.py b/.github/scripts/gitutils.py
index 9cc8f594e531a..ba406625d7f7c 100644
--- a/.github/scripts/gitutils.py
+++ b/.github/scripts/gitutils.py
@@ -1,5 +1,7 @@
 #!/usr/bin/env python3
 
+from __future__ import annotations
+
 import os
 import re
 import tempfile
@@ -7,7 +9,7 @@
 from collections.abc import Callable, Iterator
 from datetime import datetime
 from functools import wraps
-from typing import Any, cast, Optional, TypeVar, Union
+from typing import Any, cast, TypeVar
 
 
 T = TypeVar("T")
@@ -60,7 +62,7 @@ class GitCommit:
     body: str
     author: str
     author_date: datetime
-    commit_date: Optional[datetime]
+    commit_date: datetime | None
 
     def __init__(
         self,
@@ -69,7 +71,7 @@ def __init__(
         author_date: datetime,
         title: str,
         body: str,
-        commit_date: Optional[datetime] = None,
+        commit_date: datetime | None = None,
     ) -> None:
         self.commit_hash = commit_hash
         self.author = author
@@ -85,7 +87,7 @@ def __contains__(self, item: Any) -> bool:
         return item in self.body or item in self.title
 
 
-def parse_fuller_format(lines: Union[str, list[str]]) -> GitCommit:
+def parse_fuller_format(lines: str | list[str]) -> GitCommit:
     """
     Expect commit message generated using `--format=fuller --date=unix` format, i.e.:
         commit <sha1>
@@ -163,7 +165,7 @@ def branches_containing_ref(
         )
         return [x.strip() for x in rc.split("\n") if x.strip()] if len(rc) > 0 else []
 
-    def current_branch(self) -> Optional[str]:
+    def current_branch(self) -> str | None:
         try:
             return self._run_git("symbolic-ref", "--short", "HEAD").strip()
         except RuntimeError:
@@ -176,7 +178,7 @@ def checkout(self, branch: str) -> None:
     def create_branch_and_checkout(self, branch: str) -> None:
         self._run_git("checkout", "-b", branch)
 
-    def fetch(self, ref: Optional[str] = None, branch: Optional[str] = None) -> None:
+    def fetch(self, ref: str | None = None, branch: str | None = None) -> None:
         if branch is None and ref is None:
             self._run_git("fetch", self.remote)
         elif branch is None:
@@ -196,7 +198,7 @@ def rev_parse(self, name: str) -> str:
     def get_merge_base(self, from_ref: str, to_ref: str) -> str:
         return self._run_git("merge-base", from_ref, to_ref).strip()
 
-    def patch_id(self, ref: Union[str, list[str]]) -> list[tuple[str, str]]:
+    def patch_id(self, ref: str | list[str]) -> list[tuple[str, str]]:
         is_list = isinstance(ref, list)
         if is_list:
             if len(ref) == 0:
@@ -334,7 +336,7 @@ def commit_message(self, ref: str) -> str:
     def amend_commit_message(self, msg: str) -> None:
         self._run_git("commit", "--amend", "-m", msg)
 
-    def diff(self, from_ref: str, to_ref: Optional[str] = None) -> str:
+    def diff(self, from_ref: str, to_ref: str | None = None) -> str:
         if to_ref is None:
             return self._run_git("diff", f"{from_ref}^!")
         return self._run_git("diff", f"{from_ref}..{to_ref}")
@@ -358,12 +360,12 @@ def __init__(self, val: str) -> None:
         self._val = val
         self._idx = -1
 
-    def peek(self) -> Optional[str]:
+    def peek(self) -> str | None:
         if self._idx + 1 >= len(self._val):
             return None
         return self._val[self._idx + 1]
 
-    def __iter__(self) -> "PeekableIterator":
+    def __iter__(self) -> PeekableIterator:
         return self
 
     def __next__(self) -> str:
@@ -427,7 +429,7 @@ def is_commit_hash(ref: str) -> bool:
 
 
 def are_ghstack_branches_in_sync(
-    repo: GitRepo, head_ref: str, base_ref: Optional[str] = None
+    repo: GitRepo, head_ref: str, base_ref: str | None = None
 ) -> bool:
     """Checks that diff between base and head is the same as diff between orig and its parent"""
     orig_ref = re.sub(r"/head$", "/orig", head_ref)
diff --git a/.github/scripts/label_utils.py b/.github/scripts/label_utils.py
index a3ea274f8304f..bb97c65fdefe3 100644
--- a/.github/scripts/label_utils.py
+++ b/.github/scripts/label_utils.py
@@ -1,8 +1,10 @@
 """GitHub Label Utilities."""
 
+from __future__ import annotations
+
 import json
 from functools import lru_cache
-from typing import Any, TYPE_CHECKING, Union
+from typing import Any, TYPE_CHECKING
 
 from github_utils import gh_fetch_url_and_headers, GitHubComment
 
@@ -75,7 +77,7 @@ def gh_get_labels(org: str, repo: str) -> list[str]:
 
 
 def gh_add_labels(
-    org: str, repo: str, pr_num: int, labels: Union[str, list[str]], dry_run: bool
+    org: str, repo: str, pr_num: int, labels: str | list[str], dry_run: bool
 ) -> None:
     if dry_run:
         print(f"Dryrun: Adding labels {labels} to PR {pr_num}")
@@ -106,7 +108,7 @@ def get_release_notes_labels(org: str, repo: str) -> list[str]:
     ]
 
 
-def has_required_labels(pr: "GitHubPR") -> bool:
+def has_required_labels(pr: GitHubPR) -> bool:
     pr_labels = pr.get_labels()
     # Check if PR is not user facing
     is_not_user_facing_pr = any(
diff --git a/.github/scripts/map_ec2_to_arc.py b/.github/scripts/map_ec2_to_arc.py
new file mode 100644
index 0000000000000..b32cd1c784775
--- /dev/null
+++ b/.github/scripts/map_ec2_to_arc.py
@@ -0,0 +1,100 @@
+#!/usr/bin/env python3
+"""Map EC2 runner labels to ARC equivalents using .github/arc.yaml.
+
+Takes a GitHub Actions test matrix, replaces each runner with its ARC
+equivalent, and prints the updated matrix as JSON.
+
+Usage:
+    python map_ec2_to_arc.py --prefix mt- '{ include: [
+      { config: "default", shard: 1, num_shards: 5, runner: "mt-linux.4xlarge" },
+    ]}'
+"""
+
+import argparse
+import json
+import os
+import sys
+from pathlib import Path
+
+import yaml
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(
+        description="Map EC2 runner labels to ARC runner labels in a test matrix"
+    )
+    parser.add_argument(
+        "matrix",
+        help="GitHub Actions test matrix string to transform",
+    )
+    parser.add_argument(
+        "--prefix",
+        default="",
+        help="Runner prefix to strip from labels (e.g. 'mt-')",
+    )
+    return parser.parse_args()
+
+
+def strip_prefix(label: str, prefix: str) -> str:
+    if prefix and label.startswith(prefix):
+        return label[len(prefix) :]
+    return label
+
+
+def load_mapping(arc_yaml: Path) -> dict[str, str]:
+    with open(arc_yaml) as f:
+        data = yaml.safe_load(f)
+    return data["runner_mapping"]
+
+
+def set_output(name: str, val: str) -> None:
+    print(f"Setting {name}={val}")
+    github_output = os.getenv("GITHUB_OUTPUT")
+    if github_output:
+        with open(github_output, "a") as f:
+            print(f"{name}={val}", file=f)
+
+
+def main() -> None:
+    args = parse_args()
+    arc_yaml = Path(__file__).resolve().parent.parent / "arc.yaml"
+    mapping = load_mapping(arc_yaml)
+
+    matrix = yaml.safe_load(args.matrix)
+    if not matrix:
+        set_output("test-matrix", args.matrix)
+        return
+
+    entries = matrix.get("include", [])
+    if not entries:
+        set_output("test-matrix", json.dumps(matrix))
+        return
+
+    # TODO(huydo): onnxruntime uses hardware_concurrency() to size its thread
+    # pool, which sees all host CPUs (e.g., 192) on ARC k8s instead of the
+    # container's cpuset (e.g., 16). This causes pthread_setaffinity_np errors.
+    # Skip onnx tests on ARC until the onnxruntime session options are fixed to
+    # use cgroup-aware CPU counts.
+    excluded_configs = {"onnx"}
+    filtered = []
+    for entry in entries:
+        if entry.get("config") in excluded_configs:
+            print(f"Excluding config '{entry['config']}' from ARC test matrix")
+            continue
+        filtered.append(entry)
+    matrix["include"] = filtered
+
+    for entry in filtered:
+        if "runner" not in entry:
+            continue
+        clean = strip_prefix(entry["runner"].strip(), args.prefix)
+        if clean not in mapping:
+            print(f"error: no ARC runner found for '{clean}'", file=sys.stderr)
+            sys.exit(1)
+        entry["runner"] = args.prefix + mapping[clean]
+
+    set_output("test-matrix", json.dumps(matrix))
+
+
+if __name__ == "__main__":
+    main()
diff --git a/.github/scripts/runner_determinator.py b/.github/scripts/runner_determinator.py
index baf560234549b..40895b84c488f 100644
--- a/.github/scripts/runner_determinator.py
+++ b/.github/scripts/runner_determinator.py
@@ -1,9 +1,5 @@
 # flake8: noqa: G004
 
-# Note: Copies of this script in runner_determinator.py and _runner-determinator.yml
-#       must be kept in sync. You can do it easily by running the following command:
-#           python .github/scripts/update_runner_determinator.py
-
 """
 This runner determinator is used to determine which set of runners to run a
 GitHub job on. It uses the first comment of a GitHub issue (by default
@@ -79,13 +75,18 @@
 GITHUB_OUTPUT = os.getenv("GITHUB_OUTPUT", "")
 GH_OUTPUT_KEY_AMI = "runner-ami"
 GH_OUTPUT_KEY_LABEL_TYPE = "label-type"
+GH_OUTPUT_KEY_USE_ARC = "use-arc"
 OPT_OUT_LABEL = "no-runner-experiments"
 
 SETTING_EXPERIMENTS = "experiments"
 
 LF_FLEET_EXPERIMENT = "lf"
+ARC_FLEET_EXPERIMENT = "arc"
 CANARY_FLEET_SUFFIX = ".c"
 
+ARC_LABEL_PREFIX = "mt-"
+ARC_CANARY_LABEL_PREFIX = "c-"
+
 
 class Experiment(NamedTuple):
     rollout_perc: float = (
@@ -101,6 +102,11 @@ class Experiment(NamedTuple):
     # Add more fields as needed
 
 
+class RunnerPrefixResult(NamedTuple):
+    prefix: str
+    use_arc: bool = False
+
+
 class Settings(NamedTuple):
     """
     Settings for the experiments that can be opted into.
@@ -352,11 +358,8 @@ def parse_settings_from_text(settings_text: str) -> Settings:
     """
     try:
         if settings_text:
-            # Escape the backtick as well so that we can have the settings in a code block on the GH issue
-            # for easy reading
-            # Note: Using ascii for the backtick so that the cat step in _runner-determinator.yml doesn't choke on
-            #       the backtick character in shell commands.
-            backtick = chr(96)  # backtick character
+            # Strip backticks so settings can be in a code block on the GH issue
+            backtick = chr(96)
             settings_text = settings_text.strip(f"\r\n\t{backtick} ")
             settings = load_yaml(settings_text)
 
@@ -439,12 +442,13 @@ def get_runner_prefix(
     eligible_experiments: frozenset[str] = frozenset(),
     opt_out_experiments: frozenset[str] = frozenset(),
     is_canary: bool = False,
-) -> str:
+) -> RunnerPrefixResult:
     settings = parse_settings(rollout_state)
     user_optins = parse_users(rollout_state)
 
     fleet_prefix = ""
     prefixes = []
+    use_arc = False
     for experiment_name, experiment_settings in settings.experiments.items():
         if not experiment_settings.all_branches and is_exception_branch(branch):
             log.info(
@@ -510,7 +514,12 @@ def get_runner_prefix(
 
         if enabled:
             label = experiment_name
-            if experiment_name == LF_FLEET_EXPERIMENT:
+            if experiment_name == ARC_FLEET_EXPERIMENT:
+                use_arc = True
+                log.info(
+                    f"ARC experiment enabled. Using ARC runner prefix ({'canary' if is_canary else 'production'})."
+                )
+            elif experiment_name == LF_FLEET_EXPERIMENT:
                 # We give some special treatment to the "lf" experiment since determines the fleet we use
                 #  - If it's enabled, then we always list it's prefix first
                 #  - If we're in the canary branch, then we append ".c" to the lf prefix
@@ -520,6 +529,15 @@ def get_runner_prefix(
             else:
                 prefixes.append(label)
 
+    # ARC experiment takes precedence: return a fixed label prefix
+    if use_arc:
+        arc_prefix = (
+            ARC_CANARY_LABEL_PREFIX + ARC_LABEL_PREFIX
+            if is_canary
+            else ARC_LABEL_PREFIX
+        )
+        return RunnerPrefixResult(prefix=arc_prefix, use_arc=True)
+
     if len(prefixes) > 1:
         log.error(
             f"Only a fleet and one other experiment can be enabled for a job at any time. Enabling {prefixes[0]} and ignoring the rest, which are {', '.join(prefixes[1:])}"
@@ -530,7 +548,8 @@ def get_runner_prefix(
     if fleet_prefix:
         prefixes.insert(0, fleet_prefix)
 
-    return ".".join(prefixes) + "." if prefixes else ""
+    prefix = ".".join(prefixes) + "." if prefixes else ""
+    return RunnerPrefixResult(prefix=prefix)
 
 
 def get_rollout_state_from_issue(github_token: str, repo: str, issue_num: int) -> str:
@@ -619,7 +638,7 @@ def main() -> None:
 
         is_canary = args.github_repo == "pytorch/pytorch-canary"
 
-        runner_label_prefix = get_runner_prefix(
+        result = get_runner_prefix(
             rollout_state,
             (args.github_issue_owner, username),
             args.github_branch,
@@ -627,6 +646,8 @@ def main() -> None:
             args.opt_out_experiments,
             is_canary,
         )
+        runner_label_prefix = result.prefix
+        set_github_output(GH_OUTPUT_KEY_USE_ARC, str(result.use_arc).lower())
 
     except Exception as e:
         log.error(
diff --git a/.github/scripts/test_filter_test_configs.py b/.github/scripts/test_filter_test_configs.py
index 26e38828b7865..1463f82f61eda 100755
--- a/.github/scripts/test_filter_test_configs.py
+++ b/.github/scripts/test_filter_test_configs.py
@@ -376,13 +376,31 @@ def test_set_periodic_modes(self) -> None:
             scheduled_test_matrix = set_periodic_modes(test_matrix, job_name)
 
             expected_modes = [
-                m for m, c in SUPPORTED_PERIODICAL_MODES.items() if c(job_name)
+                m for m, c in SUPPORTED_PERIODICAL_MODES.items() if c(job_name, None)
             ]
             self.assertEqual(
                 len(test_matrix["include"]) * len(expected_modes),
                 len(scheduled_test_matrix["include"]),
             )
 
+    def test_set_periodic_modes_gpu_runner(self) -> None:
+        """Job name without 'cuda' but runner indicates a CUDA job (e.g. inductor-unittest)."""
+        test_matrix = yaml.safe_load(
+            "{include: ["
+            '{config: "inductor", shard: 1, num_shards: 2, runner: "linux.g5.4xlarge.nvidia.gpu"}, '
+            '{config: "inductor", shard: 2, num_shards: 2, runner: "linux.g5.4xlarge.nvidia.gpu"}'
+            "]}"
+        )
+        scheduled = set_periodic_modes(test_matrix, "inductor-build / build")
+
+        modes_per_config = [
+            entry.get("mem_leak_check") or entry.get("rerun_disabled_tests")
+            for entry in scheduled["include"]
+        ]
+        self.assertIn("mem_leak_check", modes_per_config)
+        self.assertIn("rerun_disabled_tests", modes_per_config)
+        self.assertEqual(len(scheduled["include"]), 4)
+
     @mock.patch("filter_test_configs.download_json")
     def test_remove_disabled_jobs(self, mock_download_json: Any) -> None:
         mock_download_json.return_value = MOCKED_DISABLED_UNSTABLE_JOBS
diff --git a/.github/scripts/test_map_ec2_to_arc.py b/.github/scripts/test_map_ec2_to_arc.py
new file mode 100644
index 0000000000000..4414532be546e
--- /dev/null
+++ b/.github/scripts/test_map_ec2_to_arc.py
@@ -0,0 +1,188 @@
+#!/usr/bin/env python3
+
+import json
+import os
+import subprocess
+import sys
+import tempfile
+from pathlib import Path
+
+
+SCRIPT = Path(__file__).resolve().parent / "map_ec2_to_arc.py"
+
+
+def run(
+    matrix: str, prefix: str = "", github_output: str | None = None
+) -> subprocess.CompletedProcess:
+    cmd = [sys.executable, str(SCRIPT)]
+    if prefix:
+        cmd += ["--prefix", prefix]
+    cmd.append(matrix)
+
+    env = os.environ.copy()
+    if github_output is not None:
+        env["GITHUB_OUTPUT"] = github_output
+    else:
+        env.pop("GITHUB_OUTPUT", None)
+
+    return subprocess.run(cmd, capture_output=True, text=True, env=env)
+
+
+def parse_output(stdout: str) -> dict:
+    """Extract the JSON matrix from the 'Setting test-matrix=...' line."""
+    prefix = "Setting test-matrix="
+    for line in stdout.splitlines():
+        if line.startswith(prefix):
+            return json.loads(line[len(prefix) :])
+    raise ValueError(f"no test-matrix output found in: {stdout}")
+
+
+def check(condition: bool, msg: str = "") -> None:
+    if not condition:
+        raise AssertionError(msg)
+
+
+def test_basic_matrix():
+    matrix = """{ include: [
+      { config: "default", shard: 1, num_shards: 1, runner: "linux.4xlarge" },
+      { config: "openreg", shard: 1, num_shards: 1, runner: "linux.2xlarge" },
+    ]}"""
+    result = run(matrix)
+    check(result.returncode == 0, result.stderr)
+    output = parse_output(result.stdout)
+    runners = [e["runner"] for e in output["include"]]
+    check(runners == ["l-x86iavx512-16-128", "l-x86iavx512-8-64"])
+
+
+def test_matrix_with_prefix():
+    matrix = """{ include: [
+      { config: "default", shard: 1, num_shards: 7, runner: "mt-linux.4xlarge" },
+      { config: "default", shard: 2, num_shards: 7, runner: "mt-linux.4xlarge" },
+      { config: "openreg", shard: 1, num_shards: 1, runner: "mt-linux.2xlarge" },
+    ]}"""
+    result = run(matrix, prefix="mt-")
+    check(result.returncode == 0, result.stderr)
+    output = parse_output(result.stdout)
+    runners = [e["runner"] for e in output["include"]]
+    check(
+        runners
+        == [
+            "mt-l-x86iavx512-16-128",
+            "mt-l-x86iavx512-16-128",
+            "mt-l-x86iavx512-8-64",
+        ]
+    )
+
+
+def test_matrix_without_prefix_when_none_present():
+    matrix = """{ include: [
+      { config: "default", shard: 1, num_shards: 1, runner: "linux.g5.4xlarge.nvidia.gpu" },
+    ]}"""
+    result = run(matrix)
+    check(result.returncode == 0, result.stderr)
+    output = parse_output(result.stdout)
+    check(output["include"][0]["runner"] == "l-x86aavx2-29-113-a10g")
+
+
+def test_unknown_runner_fails():
+    matrix = """{ include: [
+      { config: "default", shard: 1, num_shards: 1, runner: "bogus.runner" },
+    ]}"""
+    result = run(matrix)
+    check(result.returncode == 1)
+    check("no ARC runner found for 'bogus.runner'" in result.stderr)
+
+
+def test_prefix_not_present_on_runner():
+    """When --prefix is given but a runner doesn't have it, the raw label is used."""
+    matrix = """{ include: [
+      { config: "default", shard: 1, num_shards: 1, runner: "linux.4xlarge" },
+    ]}"""
+    result = run(matrix, prefix="mt-")
+    check(result.returncode == 0, result.stderr)
+    output = parse_output(result.stdout)
+    check(output["include"][0]["runner"] == "mt-l-x86iavx512-16-128")
+
+
+def test_preserves_non_runner_fields():
+    matrix = """{ include: [
+      { config: "default", shard: 3, num_shards: 7, runner: "linux.large" },
+    ]}"""
+    result = run(matrix)
+    check(result.returncode == 0, result.stderr)
+    entry = parse_output(result.stdout)["include"][0]
+    check(entry["config"] == "default")
+    check(entry["shard"] == 3)
+    check(entry["num_shards"] == 7)
+    check(entry["runner"] == "l-x86iavx512-2-4")
+
+
+def test_empty_include_passes_through():
+    matrix = """{ include: [] }"""
+    result = run(matrix)
+    check(result.returncode == 0, result.stderr)
+    output = parse_output(result.stdout)
+    check(output == {"include": []}, f"expected empty include, got {output}")
+
+
+def test_empty_string_passes_through():
+    result = run("")
+    check(result.returncode == 0, result.stderr)
+
+
+def test_mixed_runners():
+    matrix = """{ include: [
+      { config: "default", shard: 1, num_shards: 1, runner: "linux.4xlarge" },
+      { config: "gpu", shard: 1, num_shards: 1, runner: "linux.g5.4xlarge.nvidia.gpu" },
+      { config: "arm", shard: 1, num_shards: 1, runner: "linux.arm64.2xlarge" },
+    ]}"""
+    result = run(matrix)
+    check(result.returncode == 0, result.stderr)
+    output = parse_output(result.stdout)
+    runners = [e["runner"] for e in output["include"]]
+    check(
+        runners
+        == [
+            "l-x86iavx512-16-128",
+            "l-x86aavx2-29-113-a10g",
+            "l-arm64g2-6-32",
+        ]
+    )
+
+
+def test_github_output_file():
+    """When GITHUB_OUTPUT is set, the script writes test-matrix to that file."""
+    matrix = """{ include: [
+      { config: "default", shard: 1, num_shards: 1, runner: "linux.2xlarge" },
+    ]}"""
+    with tempfile.NamedTemporaryFile(mode="w", suffix=".txt", delete=False) as f:
+        tmp_path = f.name
+
+    try:
+        result = run(matrix, github_output=tmp_path)
+        check(result.returncode == 0, result.stderr)
+
+        contents = Path(tmp_path).read_text()
+        check(
+            contents.startswith("test-matrix="), f"unexpected file contents: {contents}"
+        )
+        written = json.loads(contents[len("test-matrix=") :].strip())
+        check(written["include"][0]["runner"] == "l-x86iavx512-8-64")
+    finally:
+        os.unlink(tmp_path)
+
+
+if __name__ == "__main__":
+    tests = [v for k, v in sorted(globals().items()) if k.startswith("test_")]
+    failed = 0
+    for t in tests:
+        try:
+            t()
+            print(f"  PASS  {t.__name__}")
+        except AssertionError as e:
+            print(f"  FAIL  {t.__name__}: {e}")
+            failed += 1
+    if failed:
+        print(f"\n{failed}/{len(tests)} tests failed")
+        sys.exit(1)
+    print(f"\nAll {len(tests)} tests passed")
diff --git a/.github/scripts/test_runner_determinator.py b/.github/scripts/test_runner_determinator.py
index e8f9f1b8b4aa6..47d54af377f45 100644
--- a/.github/scripts/test_runner_determinator.py
+++ b/.github/scripts/test_runner_determinator.py
@@ -183,8 +183,8 @@ def test_opted_in_user(self) -> None:
         @User2,lf,otherExp
 
         """
-        prefix = rd.get_runner_prefix(settings_text, ["User1"], USER_BRANCH)
-        self.assertEqual("lf.", prefix, "Runner prefix not correct for User1")
+        result = rd.get_runner_prefix(settings_text, ["User1"], USER_BRANCH)
+        self.assertEqual("lf.", result.prefix, "Runner prefix not correct for User1")
 
     def test_explicitly_opted_out_user(self) -> None:
         settings_text = """
@@ -200,8 +200,8 @@ def test_explicitly_opted_out_user(self) -> None:
         @User2,lf,otherExp
 
         """
-        prefix = rd.get_runner_prefix(settings_text, ["User1"], USER_BRANCH)
-        self.assertEqual("", prefix, "Runner prefix not correct for User1")
+        result = rd.get_runner_prefix(settings_text, ["User1"], USER_BRANCH)
+        self.assertEqual("", result.prefix, "Runner prefix not correct for User1")
 
     def test_explicitly_opted_in_and_out_user_should_opt_out(self) -> None:
         settings_text = """
@@ -217,8 +217,8 @@ def test_explicitly_opted_in_and_out_user_should_opt_out(self) -> None:
         @User2,lf,otherExp
 
         """
-        prefix = rd.get_runner_prefix(settings_text, ["User1"], USER_BRANCH)
-        self.assertEqual("", prefix, "Runner prefix not correct for User1")
+        result = rd.get_runner_prefix(settings_text, ["User1"], USER_BRANCH)
+        self.assertEqual("", result.prefix, "Runner prefix not correct for User1")
 
     def test_opted_in_user_two_experiments(self) -> None:
         settings_text = """
@@ -234,8 +234,10 @@ def test_opted_in_user_two_experiments(self) -> None:
         @User2,lf,otherExp
 
         """
-        prefix = rd.get_runner_prefix(settings_text, ["User2"], USER_BRANCH)
-        self.assertEqual("lf.otherExp.", prefix, "Runner prefix not correct for User2")
+        result = rd.get_runner_prefix(settings_text, ["User2"], USER_BRANCH)
+        self.assertEqual(
+            "lf.otherExp.", result.prefix, "Runner prefix not correct for User2"
+        )
 
     def test_opted_in_user_two_experiments_default(self) -> None:
         settings_text = """
@@ -252,8 +254,8 @@ def test_opted_in_user_two_experiments_default(self) -> None:
         @User2,lf,otherExp
 
         """
-        prefix = rd.get_runner_prefix(settings_text, ["User2"], USER_BRANCH)
-        self.assertEqual("lf.", prefix, "Runner prefix not correct for User2")
+        result = rd.get_runner_prefix(settings_text, ["User2"], USER_BRANCH)
+        self.assertEqual("lf.", result.prefix, "Runner prefix not correct for User2")
 
     def test_opted_in_user_two_experiments_default_exp(self) -> None:
         settings_text = """
@@ -270,10 +272,12 @@ def test_opted_in_user_two_experiments_default_exp(self) -> None:
         @User2,lf,otherExp
 
         """
-        prefix = rd.get_runner_prefix(
+        result = rd.get_runner_prefix(
             settings_text, ["User2"], USER_BRANCH, frozenset(["lf", "otherExp"])
         )
-        self.assertEqual("lf.otherExp.", prefix, "Runner prefix not correct for User2")
+        self.assertEqual(
+            "lf.otherExp.", result.prefix, "Runner prefix not correct for User2"
+        )
 
     def test_opted_in_user_two_experiments_default_exp_2(self) -> None:
         settings_text = """
@@ -290,10 +294,12 @@ def test_opted_in_user_two_experiments_default_exp_2(self) -> None:
         @User2,lf,otherExp
 
         """
-        prefix = rd.get_runner_prefix(
+        result = rd.get_runner_prefix(
             settings_text, ["User2"], USER_BRANCH, frozenset(["otherExp"])
         )
-        self.assertEqual("otherExp.", prefix, "Runner prefix not correct for User2")
+        self.assertEqual(
+            "otherExp.", result.prefix, "Runner prefix not correct for User2"
+        )
 
     @patch("random.uniform", return_value=50)
     def test_opted_out_user(self, mock_uniform: Mock) -> None:
@@ -310,8 +316,8 @@ def test_opted_out_user(self, mock_uniform: Mock) -> None:
         @User2,lf,otherExp
 
         """
-        prefix = rd.get_runner_prefix(settings_text, ["User3"], USER_BRANCH)
-        self.assertEqual("", prefix, "Runner prefix not correct for user")
+        result = rd.get_runner_prefix(settings_text, ["User3"], USER_BRANCH)
+        self.assertEqual("", result.prefix, "Runner prefix not correct for user")
 
     @patch("random.uniform", return_value=10)
     def test_opted_out_user_was_pulled_in_by_rollout(self, mock_uniform: Mock) -> None:
@@ -330,8 +336,10 @@ def test_opted_out_user_was_pulled_in_by_rollout(self, mock_uniform: Mock) -> No
         """
 
         # User3 is opted out, but is pulled into both experiments by the 10% rollout
-        prefix = rd.get_runner_prefix(settings_text, ["User3"], USER_BRANCH)
-        self.assertEqual("lf.otherExp.", prefix, "Runner prefix not correct for user")
+        result = rd.get_runner_prefix(settings_text, ["User3"], USER_BRANCH)
+        self.assertEqual(
+            "lf.otherExp.", result.prefix, "Runner prefix not correct for user"
+        )
 
     @patch("random.uniform", return_value=10)
     def test_opted_out_user_was_pulled_in_by_rollout_excl_nondefault(
@@ -353,8 +361,8 @@ def test_opted_out_user_was_pulled_in_by_rollout_excl_nondefault(
         """
 
         # User3 is opted out, but is pulled into default experiments by the 10% rollout
-        prefix = rd.get_runner_prefix(settings_text, ["User3"], USER_BRANCH)
-        self.assertEqual("lf.", prefix, "Runner prefix not correct for user")
+        result = rd.get_runner_prefix(settings_text, ["User3"], USER_BRANCH)
+        self.assertEqual("lf.", result.prefix, "Runner prefix not correct for user")
 
     @patch("random.uniform", return_value=10)
     def test_opted_out_user_was_pulled_in_by_rollout_filter_exp(
@@ -376,10 +384,12 @@ def test_opted_out_user_was_pulled_in_by_rollout_filter_exp(
         """
 
         # User3 is opted out, but is pulled into default experiments by the 10% rollout
-        prefix = rd.get_runner_prefix(
+        result = rd.get_runner_prefix(
             settings_text, ["User3"], USER_BRANCH, frozenset(["otherExp"])
         )
-        self.assertEqual("otherExp.", prefix, "Runner prefix not correct for user")
+        self.assertEqual(
+            "otherExp.", result.prefix, "Runner prefix not correct for user"
+        )
 
     @patch("random.uniform", return_value=25)
     def test_opted_out_user_was_pulled_out_by_rollout_filter_exp(
@@ -401,8 +411,8 @@ def test_opted_out_user_was_pulled_out_by_rollout_filter_exp(
         """
 
         # User3 is opted out, but is pulled into default experiments by the 10% rollout
-        prefix = rd.get_runner_prefix(settings_text, ["User3"], USER_BRANCH)
-        self.assertEqual("", prefix, "Runner prefix not correct for user")
+        result = rd.get_runner_prefix(settings_text, ["User3"], USER_BRANCH)
+        self.assertEqual("", result.prefix, "Runner prefix not correct for user")
 
     def test_lf_prefix_always_comes_first(self) -> None:
         settings_text = """
@@ -419,8 +429,10 @@ def test_lf_prefix_always_comes_first(self) -> None:
 
         """
 
-        prefix = rd.get_runner_prefix(settings_text, ["User2"], USER_BRANCH)
-        self.assertEqual("lf.otherExp.", prefix, "Runner prefix not correct for user")
+        result = rd.get_runner_prefix(settings_text, ["User2"], USER_BRANCH)
+        self.assertEqual(
+            "lf.otherExp.", result.prefix, "Runner prefix not correct for user"
+        )
 
     def test_ignores_commented_users(self) -> None:
         settings_text = """
@@ -437,8 +449,8 @@ def test_ignores_commented_users(self) -> None:
 
         """
 
-        prefix = rd.get_runner_prefix(settings_text, ["User1"], USER_BRANCH)
-        self.assertEqual("", prefix, "Runner prefix not correct for user")
+        result = rd.get_runner_prefix(settings_text, ["User1"], USER_BRANCH)
+        self.assertEqual("", result.prefix, "Runner prefix not correct for user")
 
     def test_ignores_extra_experiments(self) -> None:
         settings_text = """
@@ -456,8 +468,10 @@ def test_ignores_extra_experiments(self) -> None:
 
         """
 
-        prefix = rd.get_runner_prefix(settings_text, ["User1"], USER_BRANCH)
-        self.assertEqual("lf.otherExp.", prefix, "Runner prefix not correct for user")
+        result = rd.get_runner_prefix(settings_text, ["User1"], USER_BRANCH)
+        self.assertEqual(
+            "lf.otherExp.", result.prefix, "Runner prefix not correct for user"
+        )
 
     def test_disables_experiment_on_exception_branches_when_not_explicitly_opted_in(
         self,
@@ -473,8 +487,8 @@ def test_disables_experiment_on_exception_branches_when_not_explicitly_opted_in(
 
         """
 
-        prefix = rd.get_runner_prefix(settings_text, ["User1"], EXCEPTION_BRANCH)
-        self.assertEqual("", prefix, "Runner prefix not correct for user")
+        result = rd.get_runner_prefix(settings_text, ["User1"], EXCEPTION_BRANCH)
+        self.assertEqual("", result.prefix, "Runner prefix not correct for user")
 
     def test_allows_experiment_on_exception_branches_when_explicitly_opted_in(
         self,
@@ -491,8 +505,136 @@ def test_allows_experiment_on_exception_branches_when_explicitly_opted_in(
 
         """
 
-        prefix = rd.get_runner_prefix(settings_text, ["User1"], EXCEPTION_BRANCH)
-        self.assertEqual("lf.", prefix, "Runner prefix not correct for user")
+        result = rd.get_runner_prefix(settings_text, ["User1"], EXCEPTION_BRANCH)
+        self.assertEqual("lf.", result.prefix, "Runner prefix not correct for user")
+
+
+class TestRunnerDeterminatorArcExperiment(TestCase):
+    ARC_SETTINGS = """
+        experiments:
+            arc:
+                rollout_perc: 0
+        ---
+
+        Users:
+        @User1,arc
+        @User2,lf
+
+        """
+
+    def test_arc_opted_in_user_returns_mt_prefix(self) -> None:
+        result = rd.get_runner_prefix(self.ARC_SETTINGS, ["User1"], USER_BRANCH)
+        self.assertEqual("mt-", result.prefix)
+        self.assertTrue(result.use_arc)
+
+    def test_arc_opted_in_user_canary_returns_c_mt_prefix(self) -> None:
+        result = rd.get_runner_prefix(
+            self.ARC_SETTINGS, ["User1"], USER_BRANCH, is_canary=True
+        )
+        self.assertEqual("c-mt-", result.prefix)
+        self.assertTrue(result.use_arc)
+
+    def test_arc_not_enabled_returns_use_arc_false(self) -> None:
+        result = rd.get_runner_prefix(self.ARC_SETTINGS, ["User2"], USER_BRANCH)
+        self.assertFalse(result.use_arc)
+
+    def test_arc_not_enabled_no_experiments_returns_use_arc_false(self) -> None:
+        settings_text = """
+        experiments:
+            arc:
+                rollout_perc: 0
+        ---
+
+        Users:
+        @User1,arc
+
+        """
+        result = rd.get_runner_prefix(settings_text, ["User3"], USER_BRANCH)
+        self.assertEqual("", result.prefix)
+        self.assertFalse(result.use_arc)
+
+    @patch("random.uniform", return_value=10)
+    def test_arc_rollout_percentage(self, mock_uniform: Mock) -> None:
+        settings_text = """
+        experiments:
+            arc:
+                rollout_perc: 25
+        ---
+
+        Users:
+
+        """
+        result = rd.get_runner_prefix(settings_text, ["User3"], USER_BRANCH)
+        self.assertEqual("mt-", result.prefix)
+        self.assertTrue(result.use_arc)
+
+    @patch("random.uniform", return_value=50)
+    def test_arc_rollout_percentage_not_selected(self, mock_uniform: Mock) -> None:
+        settings_text = """
+        experiments:
+            arc:
+                rollout_perc: 25
+        ---
+
+        Users:
+
+        """
+        result = rd.get_runner_prefix(settings_text, ["User3"], USER_BRANCH)
+        self.assertEqual("", result.prefix)
+        self.assertFalse(result.use_arc)
+
+    def test_arc_opted_out_user(self) -> None:
+        settings_text = """
+        experiments:
+            arc:
+                rollout_perc: 100
+        ---
+
+        Users:
+        @User1,-arc
+
+        """
+        result = rd.get_runner_prefix(settings_text, ["User1"], USER_BRANCH)
+        self.assertEqual("", result.prefix)
+        self.assertFalse(result.use_arc)
+
+    def test_arc_exception_branch_not_enabled(self) -> None:
+        result = rd.get_runner_prefix(self.ARC_SETTINGS, ["User1"], EXCEPTION_BRANCH)
+        self.assertEqual("", result.prefix)
+        self.assertFalse(result.use_arc)
+
+    def test_arc_exception_branch_all_branches(self) -> None:
+        settings_text = """
+        experiments:
+            arc:
+                rollout_perc: 0
+                all_branches: true
+        ---
+
+        Users:
+        @User1,arc
+
+        """
+        result = rd.get_runner_prefix(settings_text, ["User1"], EXCEPTION_BRANCH)
+        self.assertEqual("mt-", result.prefix)
+        self.assertTrue(result.use_arc)
+
+    def test_arc_takes_precedence_over_lf(self) -> None:
+        settings_text = """
+        experiments:
+            lf:
+                rollout_perc: 0
+            arc:
+                rollout_perc: 0
+        ---
+
+        Users:
+        @User1,lf,arc
+
+        """
+        result = rd.get_runner_prefix(settings_text, ["User1"], USER_BRANCH)
+        self.assertEqual("mt-", result.prefix)
+        self.assertTrue(result.use_arc)
 
 
 if __name__ == "__main__":
diff --git a/.github/scripts/test_trymerge.py b/.github/scripts/test_trymerge.py
index 85e6235a23382..74803f450592c 100755
--- a/.github/scripts/test_trymerge.py
+++ b/.github/scripts/test_trymerge.py
@@ -1,4 +1,5 @@
 #!/usr/bin/env python3
+
 # Tests implemented in this file are relying on GitHub GraphQL APIs
 # In order to avoid test flakiness, results of the queries
 # are cached in gql_mocks.json
@@ -7,12 +8,14 @@
 # GraphQL queries in trymerge.py, please make sure to delete `gql_mocks.json`
 # And re-run the test locally with ones PAT
 
+from __future__ import annotations
+
 import gzip
 import json
 import os
 import warnings
 from hashlib import sha256
-from typing import Any, Optional
+from typing import Any
 from unittest import main, mock, skip, TestCase
 from urllib.error import HTTPError
 
@@ -147,8 +150,8 @@ def mock_revert(
     pr: GitHubPR,
     *,
     dry_run: bool = False,
-    comment_id: Optional[int] = None,
-    reason: Optional[str] = None,
+    comment_id: int | None = None,
+    reason: str | None = None,
 ) -> None:
     pass
 
diff --git a/.github/scripts/trymerge.py b/.github/scripts/trymerge.py
index 7def9d85e06c7..658728b096d47 100755
--- a/.github/scripts/trymerge.py
+++ b/.github/scripts/trymerge.py
@@ -1,14 +1,6 @@
 #!/usr/bin/env python3
 
-# NB: the following functions are used in Meta-internal workflows
-# (github_first_try_merge/my_handler.py) and thus have functionality limitations
-# (no `git` command access, no network access besides the strict allow list):
-#
-# find_matching_merge_rule
-# read_merge_rules
-#
-# Also any signature changes of these functions, as well as changes to the `GitHubPR`
-# class, will likely require corresponding changes for the internal workflows.
+from __future__ import annotations
 
 import base64
 import json
@@ -17,12 +9,11 @@
 import time
 import urllib.parse
 from collections import defaultdict
-from collections.abc import Callable, Iterable
 from dataclasses import dataclass
 from functools import cache
 from pathlib import Path
 from re import Pattern
-from typing import Any, cast, NamedTuple, Optional
+from typing import Any, cast, NamedTuple, TYPE_CHECKING
 from warnings import warn
 
 import yaml
@@ -54,6 +45,10 @@
 from trymerge_explainer import get_revert_message, TryMergeExplainer
 
 
+if TYPE_CHECKING:
+    from collections.abc import Callable, Iterable
+
+
 # labels
 MERGE_IN_PROGRESS_LABEL = "merging"
 MERGE_COMPLETE_LABEL = "merged"
@@ -62,22 +57,22 @@
 class JobCheckState(NamedTuple):
     name: str
     url: str
-    status: Optional[str]
-    classification: Optional[str]
-    job_id: Optional[int]
-    title: Optional[str]
-    summary: Optional[str]
+    status: str | None
+    classification: str | None
+    job_id: int | None
+    title: str | None
+    summary: str | None
 
 
 JobNameToStateDict = dict[str, JobCheckState]
 
 
 class WorkflowCheckState:
-    def __init__(self, name: str, url: str, run_id: int, status: Optional[str]):
+    def __init__(self, name: str, url: str, run_id: int, status: str | None):
         self.name: str = name
         self.url: str = url
         self.run_id: int = run_id
-        self.status: Optional[str] = status
+        self.status: str | None = status
         self.jobs: JobNameToStateDict = {}
 
 
@@ -489,12 +484,12 @@ def iter_issue_timeline_until_comment(
     )
 
 
-def sha_from_committed_event(ev: dict[str, Any]) -> Optional[str]:
+def sha_from_committed_event(ev: dict[str, Any]) -> str | None:
     """Extract SHA from committed event in timeline"""
     return ev.get("sha")
 
 
-def sha_from_force_push_after(ev: dict[str, Any]) -> Optional[str]:
+def sha_from_force_push_after(ev: dict[str, Any]) -> str | None:
     """Extract SHA from force push event in timeline"""
     # The current GitHub API format
     commit_id = ev.get("commit_id")
@@ -542,7 +537,7 @@ def get_check_run_name_prefix(workflow_run: Any) -> str:
         return f"{workflow_run['workflow']['name']} / "
 
 
-def is_passing_status(status: Optional[str]) -> bool:
+def is_passing_status(status: str | None) -> bool:
     return status is not None and status.upper() in ["SUCCESS", "SKIPPED", "NEUTRAL"]
 
 
@@ -664,7 +659,7 @@ def parse_args() -> Any:
     return parser.parse_args()
 
 
-def can_skip_internal_checks(pr: "GitHubPR", comment_id: Optional[int] = None) -> bool:
+def can_skip_internal_checks(pr: GitHubPR, comment_id: int | None = None) -> bool:
     if comment_id is None:
         return False
     comment = pr.get_comment_by_id(comment_id)
@@ -675,10 +670,10 @@ def can_skip_internal_checks(pr: "GitHubPR", comment_id: Optional[int] = None) -
 
 def _revlist_to_prs(
     repo: GitRepo,
-    pr: "GitHubPR",
+    pr: GitHubPR,
     rev_list: Iterable[str],
-    should_skip: Optional[Callable[[int, "GitHubPR"], bool]] = None,
-) -> list[tuple["GitHubPR", str]]:
+    should_skip: Callable[[int, GitHubPR], bool] | None = None,
+) -> list[tuple[GitHubPR, str]]:
     rc: list[tuple[GitHubPR, str]] = []
     for idx, rev in enumerate(rev_list):
         msg = repo.commit_message(rev)
@@ -706,8 +701,8 @@ def _revlist_to_prs(
 
 
 def get_ghstack_prs(
-    repo: GitRepo, pr: "GitHubPR", open_only: bool = True
-) -> list[tuple["GitHubPR", str]]:
+    repo: GitRepo, pr: GitHubPR, open_only: bool = True
+) -> list[tuple[GitHubPR, str]]:
     """
     Get the PRs in the stack that are below this PR (inclusive).  Throws error if any of the open PRs are out of sync.
     @:param open_only: Only return open PRs
@@ -716,7 +711,7 @@ def get_ghstack_prs(
     orig_ref = f"{repo.remote}/{pr.get_ghstack_orig_ref()}"
     rev_list = repo.revlist(f"{pr.default_branch()}..{orig_ref}")
 
-    def skip_func(idx: int, candidate: "GitHubPR") -> bool:
+    def skip_func(idx: int, candidate: GitHubPR) -> bool:
         if not open_only or not candidate.is_closed():
             return False
         print(
@@ -761,14 +756,14 @@ def __init__(self, org: str, project: str, pr_num: int) -> None:
         self.project = project
         self.pr_num = pr_num
         self.info = gh_get_pr_info(org, project, pr_num)
-        self.changed_files: Optional[list[str]] = None
-        self.labels: Optional[list[str]] = None
-        self.conclusions: Optional[JobNameToStateDict] = None
-        self.comments: Optional[list[GitHubComment]] = None
-        self._authors: Optional[list[tuple[str, str]]] = None
-        self._reviews: Optional[list[tuple[str, str]]] = None
-        self.merge_base: Optional[str] = None
-        self.submodules: Optional[list[str]] = None
+        self.changed_files: list[str] | None = None
+        self.labels: list[str] | None = None
+        self.conclusions: JobNameToStateDict | None = None
+        self.comments: list[GitHubComment] | None = None
+        self._authors: list[tuple[str, str]] | None = None
+        self._reviews: list[tuple[str, str]] | None = None
+        self.merge_base: str | None = None
+        self.submodules: list[str] | None = None
 
     def is_closed(self) -> bool:
         return bool(self.info["closed"])
@@ -804,7 +799,7 @@ def get_changed_files_count(self) -> int:
     def last_commit(self) -> Any:
         return self.info["commits"]["nodes"][-1]["commit"]
 
-    def last_commit_sha(self, default: Optional[str] = None) -> str:
+    def last_commit_sha(self, default: str | None = None) -> str:
         # for commits, the oid is the sha
 
         if default is None:
@@ -910,7 +905,7 @@ def get_approved_by(self) -> list[str]:
     def get_commit_count(self) -> int:
         return int(self.info["commits_with_authors"]["totalCount"])
 
-    def get_commit_sha_at_comment(self, comment_id: int) -> Optional[str]:
+    def get_commit_sha_at_comment(self, comment_id: int) -> str | None:
         """
         Get the PR head commit SHA that was present when a specific comment was posted.
         This ensures we only merge the state of the PR at the time the merge command was issued,
@@ -1057,11 +1052,10 @@ def get_pr_next_checksuites(checksuites: Any) -> Any:
                     summary=None,
                 )
 
-        # Making an exception for Apply lint auggestions/autoformat because the
-        # bot adds a merged label -> triggers workflow -> sometimes needs
-        # approval -> is read as failure, which results in a blocked merge, but
-        # this workflow doesn't provide mergability info
-        self.conclusions.pop("Apply lint suggestions", None)
+        # Same issue for Claude Code: triggered by comment events, uses a
+        # protected environment (bedrock), resulting in ACTION_REQUIRED
+        # check suite conclusions that block merges
+        self.conclusions.pop("Claude Code", None)
 
         return self.conclusions
 
@@ -1089,7 +1083,7 @@ def get_title(self) -> str:
     def get_body(self) -> str:
         return cast(str, self.info["body"])
 
-    def get_merge_commit(self) -> Optional[str]:
+    def get_merge_commit(self) -> str | None:
         mc = self.info["mergeCommit"]
         return mc["oid"] if mc is not None else None
 
@@ -1158,7 +1152,7 @@ def get_comment_by_id(self, database_id: int) -> GitHubComment:
 
         raise RuntimeError(f"Comment with id {database_id} not found")
 
-    def get_diff_revision(self) -> Optional[str]:
+    def get_diff_revision(self) -> str | None:
         rc = RE_DIFF_REV.search(self.get_body())
         return rc.group(1) if rc is not None else None
 
@@ -1182,9 +1176,9 @@ def merge_ghstack_into(
         self,
         repo: GitRepo,
         skip_mandatory_checks: bool,
-        comment_id: Optional[int] = None,
+        comment_id: int | None = None,
         skip_all_rule_checks: bool = False,
-    ) -> list["GitHubPR"]:
+    ) -> list[GitHubPR]:
         if not self.is_ghstack_pr():
             raise AssertionError(
                 f"merge_ghstack_into called on non-ghstack PR #{self.pr_num}"
@@ -1217,7 +1211,7 @@ def merge_ghstack_into(
     def gen_commit_message(
         self,
         filter_ghstack: bool = False,
-        ghstack_deps: Optional[list["GitHubPR"]] = None,
+        ghstack_deps: list[GitHubPR] | None = None,
     ) -> str:
         """Fetches title and body from PR description
         adds reviewed by, pull request resolved and optionally
@@ -1269,7 +1263,7 @@ def merge_into(
         skip_mandatory_checks: bool = False,
         dry_run: bool = False,
         comment_id: int,
-        ignore_current_checks: Optional[list[str]] = None,
+        ignore_current_checks: list[str] | None = None,
     ) -> None:
         # Raises exception if matching rule is not found
         (
@@ -1338,10 +1332,10 @@ def merge_changes_locally(
         self,
         repo: GitRepo,
         skip_mandatory_checks: bool = False,
-        comment_id: Optional[int] = None,
-        branch: Optional[str] = None,
+        comment_id: int | None = None,
+        branch: str | None = None,
         skip_all_rule_checks: bool = False,
-    ) -> list["GitHubPR"]:
+    ) -> list[GitHubPR]:
         """
         :param skip_all_rule_checks: If true, skips all rule checks on ghstack PRs, useful for dry-running merge locally
         """
@@ -1403,7 +1397,7 @@ def merge_changes_locally(
 
 
 class MergeRuleFailedError(RuntimeError):
-    def __init__(self, message: str, rule: Optional["MergeRule"] = None) -> None:
+    def __init__(self, message: str, rule: MergeRule | None = None) -> None:
         super().__init__(message)
         self.rule = rule
 
@@ -1421,7 +1415,7 @@ class MergeRule:
     name: str
     patterns: list[str]
     approved_by: list[str]
-    mandatory_checks_name: Optional[list[str]]
+    mandatory_checks_name: list[str] | None
     ignore_flaky_failures: bool = True
 
 
@@ -1436,9 +1430,7 @@ def gen_new_issue_link(
     )
 
 
-def read_merge_rules(
-    repo: Optional[GitRepo], org: str, project: str
-) -> list[MergeRule]:
+def read_merge_rules(repo: GitRepo | None, org: str, project: str) -> list[MergeRule]:
     """Returns the list of all merge rules for the repo or project.
 
     NB: this function is used in Meta-internal workflows, see the comment
@@ -1465,14 +1457,14 @@ def read_merge_rules(
 
 def find_matching_merge_rule(
     pr: GitHubPR,
-    repo: Optional[GitRepo] = None,
+    repo: GitRepo | None = None,
     skip_mandatory_checks: bool = False,
     skip_internal_checks: bool = False,
-    ignore_current_checks: Optional[list[str]] = None,
+    ignore_current_checks: list[str] | None = None,
 ) -> tuple[
     MergeRule,
-    list[tuple[str, Optional[str], Optional[int]]],
-    list[tuple[str, Optional[str], Optional[int]]],
+    list[tuple[str, str | None, int | None]],
+    list[tuple[str, str | None, int | None]],
     dict[str, list[Any]],
 ]:
     """
@@ -1660,12 +1652,12 @@ def find_matching_merge_rule(
     raise MergeRuleFailedError(reject_reason, rule)
 
 
-def checks_to_str(checks: list[tuple[str, Optional[str]]]) -> str:
+def checks_to_str(checks: list[tuple[str, str | None]]) -> str:
     return ", ".join(f"[{c[0]}]({c[1]})" if c[1] is not None else c[0] for c in checks)
 
 
 def checks_to_markdown_bullets(
-    checks: list[tuple[str, Optional[str], Optional[int]]],
+    checks: list[tuple[str, str | None, int | None]],
 ) -> list[str]:
     return [
         f"- [{c[0]}]({c[1]})" if c[1] is not None else f"- {c[0]}" for c in checks[:5]
@@ -1677,9 +1669,7 @@ def post_starting_merge_comment(
     pr: GitHubPR,
     explainer: TryMergeExplainer,
     dry_run: bool,
-    ignore_current_checks_info: Optional[
-        list[tuple[str, Optional[str], Optional[int]]]
-    ] = None,
+    ignore_current_checks_info: list[tuple[str, str | None, int | None]] | None = None,
 ) -> None:
     """Post the initial merge starting message on the PR. Also post a short
     message on all PRs in the stack."""
@@ -1737,12 +1727,12 @@ def save_merge_record(
     owner: str,
     project: str,
     author: str,
-    pending_checks: list[tuple[str, Optional[str], Optional[int]]],
-    failed_checks: list[tuple[str, Optional[str], Optional[int]]],
-    ignore_current_checks: list[tuple[str, Optional[str], Optional[int]]],
-    broken_trunk_checks: list[tuple[str, Optional[str], Optional[int]]],
-    flaky_checks: list[tuple[str, Optional[str], Optional[int]]],
-    unstable_checks: list[tuple[str, Optional[str], Optional[int]]],
+    pending_checks: list[tuple[str, str | None, int | None]],
+    failed_checks: list[tuple[str, str | None, int | None]],
+    ignore_current_checks: list[tuple[str, str | None, int | None]],
+    broken_trunk_checks: list[tuple[str, str | None, int | None]],
+    flaky_checks: list[tuple[str, str | None, int | None]],
+    unstable_checks: list[tuple[str, str | None, int | None]],
     last_commit_sha: str,
     merge_base_sha: str,
     merge_commit_sha: str = "",
@@ -1875,7 +1865,7 @@ def is_flaky(
 
 def is_invalid_cancel(
     name: str,
-    conclusion: Optional[str],
+    conclusion: str | None,
     drci_classifications: Any,
 ) -> bool:
     """
@@ -1902,7 +1892,7 @@ def get_classifications(
     pr_num: int,
     project: str,
     checks: dict[str, JobCheckState],
-    ignore_current_checks: Optional[list[str]],
+    ignore_current_checks: list[str] | None,
 ) -> dict[str, JobCheckState]:
     # Get the failure classification from Dr.CI, which is the source of truth
     # going forward. It's preferable to try calling Dr.CI API directly first
@@ -2011,7 +2001,7 @@ def get_readable_drci_results(drci_classifications: Any) -> str:
 
 
 def filter_checks_with_lambda(
-    checks: JobNameToStateDict, status_filter: Callable[[Optional[str]], bool]
+    checks: JobNameToStateDict, status_filter: Callable[[str | None], bool]
 ) -> list[JobCheckState]:
     return [check for check in checks.values() if status_filter(check.status)]
 
@@ -2027,7 +2017,7 @@ def get_pr_commit_sha(repo: GitRepo, pr: GitHubPR) -> str:
 
 
 def validate_revert(
-    repo: GitRepo, pr: GitHubPR, *, comment_id: Optional[int] = None
+    repo: GitRepo, pr: GitHubPR, *, comment_id: int | None = None
 ) -> tuple[str, str]:
     comment = (
         pr.get_last_comment()
@@ -2044,9 +2034,13 @@ def validate_revert(
     # For some reason, one can not be a member of private repo, only CONTRIBUTOR
     if pr.is_base_repo_private():
         allowed_reverters.append("CONTRIBUTOR")
-    # Special case the pytorch-auto-revert app, whose does not have association
-    # But should be able to issue revert command
-    if comment.author_url == "https://github.com/apps/pytorch-auto-revert":
+    # Special case GitHub Apps that don't have a repo association
+    # but should be able to issue revert commands
+    allowed_apps = {
+        "https://github.com/apps/pytorch-auto-revert",
+        "https://github.com/apps/facebook-github-tools",
+    }
+    if comment.author_url in allowed_apps:
         allowed_reverters.append("NONE")
 
     if author_association not in allowed_reverters:
@@ -2158,8 +2152,8 @@ def try_revert(
     pr: GitHubPR,
     *,
     dry_run: bool = False,
-    comment_id: Optional[int] = None,
-    reason: Optional[str] = None,
+    comment_id: int | None = None,
+    reason: str | None = None,
 ) -> None:
     try:
         author_login, commit_sha = validate_revert(repo, pr, comment_id=comment_id)
@@ -2184,6 +2178,13 @@ def try_revert(
                 f"Failed to fetch dependent PRs: {str(e)}, fall over to single revert"
             )
 
+    if not shas_and_prs:
+        raise RuntimeError(
+            f"No revertable PRs found in ghstack for #{pr.pr_num}. "
+            f"This typically means the PR is still open (not merged) or "
+            f"its GitHub state is inconsistent. Only closed/merged PRs can be reverted."
+        )
+
     do_revert_prs(
         repo,
         pr,
@@ -2228,10 +2229,10 @@ def has_label(labels: list[str], pattern: Pattern[str] = CIFLOW_LABEL) -> bool:
 def categorize_checks(
     check_runs: JobNameToStateDict,
     required_checks: list[str],
-    ok_failed_checks_threshold: Optional[int] = None,
+    ok_failed_checks_threshold: int | None = None,
 ) -> tuple[
-    list[tuple[str, Optional[str], Optional[int]]],
-    list[tuple[str, Optional[str], Optional[int]]],
+    list[tuple[str, str | None, int | None]],
+    list[tuple[str, str | None, int | None]],
     dict[str, list[Any]],
 ]:
     """
@@ -2239,8 +2240,8 @@ def categorize_checks(
     failures and broken trunk are ignored by defaults when ok_failed_checks_threshold
     is not set (unlimited)
     """
-    pending_checks: list[tuple[str, Optional[str], Optional[int]]] = []
-    failed_checks: list[tuple[str, Optional[str], Optional[int]]] = []
+    pending_checks: list[tuple[str, str | None, int | None]] = []
+    failed_checks: list[tuple[str, str | None, int | None]] = []
 
     # failed_checks_categorization is used to keep track of all ignorable failures when saving the merge record on s3
     failed_checks_categorization: dict[str, list[Any]] = defaultdict(list)
diff --git a/.github/scripts/trymerge_explainer.py b/.github/scripts/trymerge_explainer.py
index bbc85f020a066..67638616f3778 100644
--- a/.github/scripts/trymerge_explainer.py
+++ b/.github/scripts/trymerge_explainer.py
@@ -1,7 +1,8 @@
+from __future__ import annotations
+
 import os
 import re
 from re import Pattern
-from typing import Optional
 
 
 BOT_COMMANDS_WIKI = "https://github.com/pytorch/pytorch/wiki/Bot-commands"
@@ -47,9 +48,7 @@ def __init__(
 
     def _get_flag_msg(
         self,
-        ignore_current_checks: Optional[
-            list[tuple[str, Optional[str], Optional[int]]]
-        ] = None,
+        ignore_current_checks: list[tuple[str, str | None, int | None]] | None = None,
     ) -> str:
         if self.force:
             return (
@@ -68,9 +67,7 @@ def _get_flag_msg(
 
     def get_merge_message(
         self,
-        ignore_current_checks: Optional[
-            list[tuple[str, Optional[str], Optional[int]]]
-        ] = None,
+        ignore_current_checks: list[tuple[str, str | None, int | None]] | None = None,
     ) -> str:
         title = "### Merge started"
         main_message = self._get_flag_msg(ignore_current_checks)
diff --git a/.github/scripts/update_runner_determinator.py b/.github/scripts/update_runner_determinator.py
deleted file mode 100755
index 772df87c6405a..0000000000000
--- a/.github/scripts/update_runner_determinator.py
+++ /dev/null
@@ -1,31 +0,0 @@
-#!/usr/bin/env python3
-
-import re
-
-
-# Read the contents of runner_determinator.py
-with open(".github/scripts/runner_determinator.py") as script_file:
-    script_content = script_file.read()
-
-# Indent the script content by 10 spaces to match destination indentation
-indented_script_content = "\n".join(
-    [" " * 10 + line if line else line for line in script_content.splitlines()]
-)
-
-# Read the contents of _runner-determinator.yml
-with open(".github/workflows/_runner-determinator.yml") as yml_file:
-    yml_content = yml_file.read()
-
-# Replace the content between the markers
-new_yml_content = re.sub(
-    r"(cat <<EOF > runner_determinator.py\n)(.*?)(\n\s+EOF)",
-    lambda match: match.group(1) + indented_script_content + match.group(3),
-    yml_content,
-    flags=re.DOTALL,
-)
-
-# Save the modified content back to _runner-determinator.yml
-with open(".github/workflows/_runner-determinator.yml", "w") as yml_file:
-    yml_file.write(new_yml_content)
-
-print("Updated _runner-determinator.yml with the contents of runner_determinator.py")
diff --git a/.github/scripts/windows/build_magma.bat b/.github/scripts/windows/build_magma.bat
index 75c916ecdbef7..ddb38c9ed9e6d 100644
--- a/.github/scripts/windows/build_magma.bat
+++ b/.github/scripts/windows/build_magma.bat
@@ -35,6 +35,9 @@ cd magma
 mkdir build && cd build
 
 set GPU_TARGET=All
+if "%CUVER_NODOT%" == "132" (
+  set CUDA_ARCH_LIST=-gencode=arch=compute_75,code=sm_75 -gencode arch=compute_80,code=sm_80 -gencode arch=compute_86,code=sm_86 -gencode arch=compute_90,code=sm_90 -gencode arch=compute_100,code=sm_100 -gencode arch=compute_120,code=sm_120
+)
 if "%CUVER_NODOT%" == "130" (
   set CUDA_ARCH_LIST=-gencode=arch=compute_75,code=sm_75 -gencode arch=compute_80,code=sm_80 -gencode arch=compute_86,code=sm_86 -gencode arch=compute_90,code=sm_90 -gencode arch=compute_100,code=sm_100 -gencode arch=compute_120,code=sm_120
 )
diff --git a/.github/templates/common.yml.j2 b/.github/templates/common.yml.j2
index 064eea7592230..acbd4aee44736 100644
--- a/.github/templates/common.yml.j2
+++ b/.github/templates/common.yml.j2
@@ -32,7 +32,7 @@ concurrency:
 {%- macro setup_ec2_windows() -%}
       !{{ display_ec2_information() }}
       - name: "[FB EMPLOYEES] Enable SSH (Click me for login details)"
-        uses: pytorch/test-infra/.github/actions/setup-ssh@main
+        uses: pytorch/test-infra/.github/actions/setup-ssh@release/2.12
         continue-on-error: true
         with:
           github-secret: ${{ secrets.GITHUB_TOKEN }}
diff --git a/.github/templates/linux_binary_build_workflow.yml.j2 b/.github/templates/linux_binary_build_workflow.yml.j2
index b9d3a51354e06..349207af1d7de 100644
--- a/.github/templates/linux_binary_build_workflow.yml.j2
+++ b/.github/templates/linux_binary_build_workflow.yml.j2
@@ -56,7 +56,7 @@ jobs:
   get-label-type:
     if: github.repository_owner == 'pytorch'
     name: get-label-type
-    uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@main
+    uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@release/2.12
     with:
       triggering_actor: ${{ github.triggering_actor }}
       issue_owner: ${{ github.event.pull_request.user.login || github.event.issue.user.login }}
@@ -79,7 +79,7 @@ jobs:
       timeout-minutes: 420
       {%- elif config["gpu_arch_type"] == "rocm" %}
       runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
-      timeout-minutes: 300
+      timeout-minutes: 420
       {%- elif "conda" in build_environment and config["gpu_arch_type"] == "cuda" %}
       runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
       runs_on: linux.24xlarge.ephemeral
@@ -116,12 +116,9 @@ jobs:
       ALPINE_IMAGE: "docker.io/s390x/alpine"
       {%- elif config["gpu_arch_type"] == "rocm" %}
       runs_on: linux.rocm.gpu
-      {%- elif config["gpu_arch_type"] == "cuda" and config["gpu_arch_version"] in ["12.6"] %}
-      runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
-      runs_on: linux.4xlarge.nvidia.gpu  # 12.6 build can use maxwell (sm_50) runner
       {%- elif config["gpu_arch_type"] == "cuda" %}
       runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
-      runs_on: linux.g4dn.4xlarge.nvidia.gpu # 12.8+ builds need sm_70+ runner
+      runs_on: linux.g4dn.4xlarge.nvidia.gpu
       {%- else %}
       runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
       runs_on: linux.4xlarge
@@ -136,43 +133,43 @@ jobs:
       id-token: write
       contents: read
     steps:
+      !{{ common.checkout(deep_clone=False, checkout_pr_head=False) }}
       - name: Setup XPU
-        uses: pytorch/pytorch/.github/actions/setup-xpu@main
+        uses: pytorch/pytorch/.github/actions/setup-xpu@release/2.12
       - name: Login to ECR
-        uses: pytorch/pytorch/.github/actions/ecr-login@main
+        uses: pytorch/pytorch/.github/actions/ecr-login@release/2.12
       - uses: !{{ common.download_artifact_action }}
         name: Download Build Artifacts
         with:
           name: !{{ config["build_name"] }}
           path: "${{ runner.temp }}/artifacts/"
-      !{{ common.checkout(deep_clone=False, directory="pytorch") }}
       - name: Calculate docker image
         id: calculate-docker-image
-        uses: pytorch/test-infra/.github/actions/calculate-docker-image@main
+        uses: pytorch/test-infra/.github/actions/calculate-docker-image@release/2.12
         with:
           docker-registry: ${{ startsWith(github.event.ref, 'refs/tags/ciflow/') && '308535385114.dkr.ecr.us-east-1.amazonaws.com' || 'docker.io' }}
           docker-image-name: !{{ config["container_image"] }}
           custom-tag-prefix: !{{ config["container_image_tag_prefix"] }}
           docker-build-dir: .ci/docker
-          working-directory: pytorch
       - name: Pull Docker image
-        uses: pytorch/test-infra/.github/actions/pull-docker-image@main
+        uses: pytorch/test-infra/.github/actions/pull-docker-image@release/2.12
         with:
           docker-image: ${{ steps.calculate-docker-image.outputs.docker-image }}
       - name: Test Pytorch binary
-        uses: ./pytorch/.github/actions/test-pytorch-binary
+        uses: ./.github/actions/test-pytorch-binary
         env:
           DOCKER_IMAGE: ${{ steps.calculate-docker-image.outputs.docker-image }}
       - name: Teardown XPU
         uses: ./.github/actions/teardown-xpu
     {%- else %}
-    runs-on: linux.rocm.gpu.mi250.1
+    runs-on: linux.rocm.gpu.gfx942.1
     timeout-minutes: !{{ common.timeout_minutes }}
     !{{ upload.binary_env(config) }}
     permissions:
       id-token: write
       contents: read
     steps:
+      !{{ common.checkout(deep_clone=False, checkout_pr_head=False) }}
       - name: Setup ROCm
         uses: ./.github/actions/setup-rocm
       - uses: !{{ common.download_artifact_action }}
@@ -180,33 +177,31 @@ jobs:
         with:
           name: !{{ config["build_name"] }}
           path: "${{ runner.temp }}/artifacts/"
-      !{{ common.checkout(deep_clone=False, directory="pytorch") }}
       - name: ROCm set GPU_FLAG
         run: |
           echo "GPU_FLAG=--device=/dev/mem --device=/dev/kfd --device=/dev/dri --group-add video --group-add daemon" >> "${GITHUB_ENV}"
       - name: configure aws credentials
         id: aws_creds
         if: ${{ startsWith(github.event.ref, 'refs/tags/ciflow/') }}
-        uses: aws-actions/configure-aws-credentials@v4
+        uses: aws-actions/configure-aws-credentials@7474bc4690e29a8392af63c5b98e7449536d5c3a # v4
         with:
           role-to-assume: arn:aws:iam::308535385114:role/gha_workflow_s3_and_ecr_read_only
           aws-region: us-east-1
           role-duration-seconds: 18000
       - name: Calculate docker image
         id: calculate-docker-image
-        uses: pytorch/test-infra/.github/actions/calculate-docker-image@main
+        uses: pytorch/test-infra/.github/actions/calculate-docker-image@release/2.12
         with:
           docker-registry: ${{ startsWith(github.event.ref, 'refs/tags/ciflow/') && '308535385114.dkr.ecr.us-east-1.amazonaws.com' || 'docker.io' }}
           docker-image-name: !{{ config["container_image"] }}
           custom-tag-prefix: !{{ config["container_image_tag_prefix"] }}
           docker-build-dir: .ci/docker
-          working-directory: pytorch
       - name: Pull Docker image
-        uses: pytorch/test-infra/.github/actions/pull-docker-image@main
+        uses: pytorch/test-infra/.github/actions/pull-docker-image@release/2.12
         with:
           docker-image: ${{ steps.calculate-docker-image.outputs.docker-image }}
       - name: Test Pytorch binary
-        uses: ./pytorch/.github/actions/test-pytorch-binary
+        uses: ./.github/actions/test-pytorch-binary
         env:
           DOCKER_IMAGE: ${{ steps.calculate-docker-image.outputs.docker-image }}
       - name: Teardown ROCm
@@ -217,3 +212,48 @@ jobs:
   !{{ upload.upload_binaries(config) }}
 {%- endif %}
 {% endfor %}
+
+{%- for lt_config in libtorch_extraction_configs %}
+  !{{ lt_config["build_name"] }}-extract:
+    if: ${{ github.repository_owner == 'pytorch' }}
+    needs:
+      - !{{ lt_config["source_wheel_build_name"] }}-build
+      - get-label-type
+    runs-on: "${{ needs.get-label-type.outputs.label-type }}linux.4xlarge"
+    timeout-minutes: 60
+    env:
+      DESIRED_CUDA: !{{ lt_config["desired_cuda"] }}
+      LIBTORCH_VARIANT: !{{ lt_config["libtorch_variant"] }}
+    steps:
+      - name: Checkout PyTorch
+        uses: actions/checkout@v4
+        with:
+          ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
+          sparse-checkout: .ci/libtorch/
+          show-progress: false
+      - uses: !{{ common.download_artifact_action }}
+        name: Download Wheel Artifact
+        with:
+          name: !{{ lt_config["source_wheel_build_name"] }}
+          path: "${{ runner.temp }}/wheel_artifact/"
+      - name: Extract libtorch from wheel
+        run: |
+          set -eux
+          mkdir -p "${{ runner.temp }}/libtorch_output"
+          python3 .ci/libtorch/extract_libtorch_from_wheel.py \
+            --wheel-dir "${{ runner.temp }}/wheel_artifact" \
+            --output-dir "${{ runner.temp }}/libtorch_output" \
+            --platform linux \
+            --desired-cuda "$DESIRED_CUDA" \
+            --libtorch-variant "$LIBTORCH_VARIANT"
+      - uses: !{{ common.upload_artifact_action }}
+        if: always()
+        with:
+          name: !{{ lt_config["build_name"] }}
+          retention-days: 14
+          if-no-files-found: error
+          path: "${{ runner.temp }}/libtorch_output/"
+{%- if branches == "nightly" %}
+  !{{ upload.upload_extracted_libtorch(lt_config) }}
+{%- endif %}
+{% endfor %}
diff --git a/.github/templates/macos_binary_build_workflow.yml.j2 b/.github/templates/macos_binary_build_workflow.yml.j2
index 958c6b85902c2..2bb6f726003ba 100644
--- a/.github/templates/macos_binary_build_workflow.yml.j2
+++ b/.github/templates/macos_binary_build_workflow.yml.j2
@@ -71,10 +71,10 @@ jobs:
     steps:
       !{{ set_runner_specific_vars() }}
       !{{ setup_python(config.get("python_version", "3.10")) }}
-      !{{ common.checkout(deep_clone=False, directory="pytorch") }}
+      !{{ common.checkout(deep_clone=False, checkout_pr_head=False) }}
       - name: Populate binary env
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_populate_env.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_populate_env.sh"
       - name: Build PyTorch binary
         run: |
           set -eux -o pipefail
@@ -121,3 +121,45 @@ jobs:
           path: "${{ env.PYTORCH_FINAL_PACKAGE_DIR }}"
   !{{ upload.upload_binaries(config, has_test=False, use_s3=False) }}
 {%- endfor %}
+
+{%- for lt_config in libtorch_extraction_configs %}
+  !{{ lt_config["build_name"] }}-extract:
+    if: ${{ github.repository_owner == 'pytorch' }}
+    needs: !{{ lt_config["source_wheel_build_name"] }}-build
+    runs-on: !{{ macos_runner }}
+    timeout-minutes: 60
+    env:
+      DESIRED_CUDA: !{{ lt_config["desired_cuda"] }}
+      LIBTORCH_VARIANT: !{{ lt_config["libtorch_variant"] }}
+    steps:
+      - name: Checkout PyTorch
+        uses: actions/checkout@v4
+        with:
+          ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
+          sparse-checkout: .ci/libtorch/
+          show-progress: false
+      !{{ setup_python("3.10") }}
+      - uses: actions/download-artifact@v4.1.7
+        name: Download Wheel Artifact
+        with:
+          name: !{{ lt_config["source_wheel_build_name"] }}
+          path: "${{ runner.temp }}/wheel_artifact/"
+      - name: Extract libtorch from wheel
+        run: |
+          set -eux
+          mkdir -p "${{ runner.temp }}/libtorch_output"
+          python3 .ci/libtorch/extract_libtorch_from_wheel.py \
+            --wheel-dir "${{ runner.temp }}/wheel_artifact" \
+            --output-dir "${{ runner.temp }}/libtorch_output" \
+            --platform macos \
+            --desired-cuda "$DESIRED_CUDA" \
+            --libtorch-variant "$LIBTORCH_VARIANT"
+      - uses: actions/upload-artifact@v4.4.0
+        if: always()
+        with:
+          name: !{{ lt_config["build_name"] }}
+          retention-days: 14
+          if-no-files-found: error
+          path: "${{ runner.temp }}/libtorch_output/"
+  !{{ upload.upload_extracted_libtorch(lt_config, use_s3=False) }}
+{%- endfor %}
diff --git a/.github/templates/upload.yml.j2 b/.github/templates/upload.yml.j2
index 5e3798f8e2377..8c36f9553757c 100644
--- a/.github/templates/upload.yml.j2
+++ b/.github/templates/upload.yml.j2
@@ -6,7 +6,7 @@
 
 {%- macro binary_env_as_input(config, is_windows=False, include_skip_tests=False) -%}
 {%- if is_windows %}
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
+      PYTORCH_ROOT: ${{ github.workspace }}
 {%- else %}
       PYTORCH_ROOT: /pytorch
 {%- endif %}
@@ -64,5 +64,43 @@
       {%- endif %}
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
+    uses: ./.github/workflows/_binary-upload.yml
+{%- endmacro %}
+
+{%- macro upload_extracted_libtorch(lt_config, is_windows=False, use_s3=True) -%}
+!{{ lt_config["build_name"] }}-upload:  # Uploading
+    if: ${{ github.repository_owner == 'pytorch' }}
+    permissions:
+      id-token: write
+      contents: read
+    needs: !{{ lt_config["build_name"] }}-extract
+    with:
+{%- if is_windows %}
+      PYTORCH_ROOT: ${{ github.workspace }}
+{%- else %}
+      PYTORCH_ROOT: /pytorch
+{%- endif %}
+      PACKAGE_TYPE: libtorch
+      DESIRED_CUDA: !{{ lt_config["desired_cuda"] }}
+{%- if lt_config["gpu_arch_version"] %}
+      GPU_ARCH_VERSION: "!{{ lt_config["gpu_arch_version"] }}"
+{%- endif %}
+      GPU_ARCH_TYPE: !{{ lt_config["gpu_arch_type"] }}
+{%- if lt_config["libtorch_config"] %}
+      LIBTORCH_CONFIG: !{{ lt_config["libtorch_config"] }}
+{%- endif %}
+      LIBTORCH_VARIANT: !{{ lt_config["libtorch_variant"] }}
+      build_name: !{{ lt_config["build_name"] }}
+      {%- if not use_s3 %}
+      use_s3: False
+      {%- endif %}
+    secrets:
+      github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
     uses: ./.github/workflows/_binary-upload.yml
 {%- endmacro %}
diff --git a/.github/templates/windows_binary_build_workflow.yml.j2 b/.github/templates/windows_binary_build_workflow.yml.j2
index 34c148270c6bc..b24a763656428 100644
--- a/.github/templates/windows_binary_build_workflow.yml.j2
+++ b/.github/templates/windows_binary_build_workflow.yml.j2
@@ -64,7 +64,7 @@ jobs:
   get-label-type:
     if: github.repository_owner == 'pytorch'
     name: get-label-type
-    uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@main
+    uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@release/2.12
     with:
       triggering_actor: ${{ github.triggering_actor }}
       issue_owner: ${{ github.event.pull_request.user.login || github.event.issue.user.login }}
@@ -110,41 +110,40 @@ jobs:
       - name: Git checkout PyTorch
         uses: actions/checkout@v4
         with:
-          path: "pytorch"
           submodules: recursive
       - name: Bootstrap Python
         shell: cmd
         run: |
-          "pytorch/.ci/pytorch/windows/arm64/bootstrap_python.bat"
+          ".ci/pytorch/windows/arm64/bootstrap_python.bat"
       - name: Bootstrap APL
         shell: cmd
         run: |
-          "pytorch/.ci/pytorch/windows/arm64/bootstrap_apl.bat"
+          ".ci/pytorch/windows/arm64/bootstrap_apl.bat"
       - name: Bootstrap Rust
         shell: cmd
         run: |
-          "pytorch/.ci/pytorch/windows/arm64/bootstrap_rust.bat"
+          ".ci/pytorch/windows/arm64/bootstrap_rust.bat"
       - name: Bootstrap sccache
         shell: cmd
         run: |
-          "pytorch/.ci/pytorch/windows/arm64/bootstrap_sccache.bat"
+          ".ci/pytorch/windows/arm64/bootstrap_sccache.bat"
       - name: Bootstrap Libuv
         shell: cmd
         run: |
-          "pytorch/.ci/pytorch/windows/arm64/bootstrap_libuv.bat"
+          ".ci/pytorch/windows/arm64/bootstrap_libuv.bat"
 {%- else %}
       !{{ set_runner_specific_vars() }}
       !{{ common.setup_ec2_windows() }}
-      !{{ common.checkout(deep_clone=False, directory="pytorch") }}
+      !{{ common.checkout(deep_clone=False, checkout_pr_head=False) }}
 {%- endif %}
       - name: Populate binary env
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_populate_env.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_populate_env.sh"
       - name: Build PyTorch binary
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_windows_build.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_windows_build.sh"
       - uses: !{{ common.upload_artifact_action }}
         if: always()
         with:
@@ -153,7 +152,7 @@ jobs:
           if-no-files-found: error
           path: "${{ env.PYTORCH_FINAL_PACKAGE_DIR }}"
 {%- if os != "windows-arm64" %}
-      !{{ common.wait_and_kill_ssh_windows('pytorch') }}
+      !{{ common.wait_and_kill_ssh_windows() }}
 {% endif %}
   !{{ config["build_name"] }}-test:  # Testing
     if: ${{ github.repository_owner == 'pytorch' }}
@@ -195,23 +194,22 @@ jobs:
       - name: Git checkout PyTorch
         uses: actions/checkout@v4
         with:
-          path: "pytorch"
           submodules: recursive
       - name: Bootstrap APL
         shell: cmd
         run: |
-          "pytorch/.ci/pytorch/windows/arm64/bootstrap_apl.bat"
+          ".ci/pytorch/windows/arm64/bootstrap_apl.bat"
       - name: Bootstrap Python
         shell: cmd
         run: |
-          "pytorch/.ci/pytorch/windows/arm64/bootstrap_python.bat"
+          ".ci/pytorch/windows/arm64/bootstrap_python.bat"
       - name: Bootstrap Rust
         shell: cmd
         run: |
-          "pytorch/.ci/pytorch/windows/arm64/bootstrap_rust.bat"
+          ".ci/pytorch/windows/arm64/bootstrap_rust.bat"
 {%- else %}
       !{{ common.setup_ec2_windows() }}
-      !{{ common.checkout(deep_clone=False, directory="pytorch") }}
+      !{{ common.checkout(deep_clone=False, checkout_pr_head=False) }}
       !{{ set_runner_specific_vars() }}
 {%- endif %}
       - uses: !{{ common.download_artifact_action }}
@@ -222,15 +220,61 @@ jobs:
       - name: Populate binary env
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_populate_env.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_populate_env.sh"
       - name: Test PyTorch binary
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_windows_test.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_windows_test.sh"
 {%- if os != "windows-arm64" %}
-      !{{ common.wait_and_kill_ssh_windows('pytorch') }}
+      !{{ common.wait_and_kill_ssh_windows() }}
 {%- endif %}
   {%- if branches == "nightly" %}
   !{{ upload.upload_binaries(config, True) }}
   {%- endif %}
-{%- endfor %}
\ No newline at end of file
+{%- endfor %}
+
+{%- for lt_config in libtorch_extraction_configs %}
+  !{{ lt_config["build_name"] }}-extract:
+    if: ${{ github.repository_owner == 'pytorch' }}
+    needs:
+      - !{{ lt_config["source_wheel_build_name"] }}-build
+      - get-label-type
+    runs-on: "${{ needs.get-label-type.outputs.label-type }}linux.4xlarge"
+    timeout-minutes: 60
+    env:
+      DESIRED_CUDA: !{{ lt_config["desired_cuda"] }}
+      LIBTORCH_VARIANT: !{{ lt_config["libtorch_variant"] }}
+    steps:
+      - name: Checkout PyTorch
+        uses: actions/checkout@v4
+        with:
+          ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
+          sparse-checkout: .ci/libtorch/
+          show-progress: false
+      - uses: !{{ common.download_artifact_action }}
+        name: Download Wheel Artifact
+        with:
+          name: !{{ lt_config["source_wheel_build_name"] }}
+          path: "${{ runner.temp }}/wheel_artifact/"
+      - name: Extract libtorch from wheel
+        shell: bash
+        run: |
+          set -eux
+          mkdir -p "${{ runner.temp }}/libtorch_output"
+          python3 .ci/libtorch/extract_libtorch_from_wheel.py \
+            --wheel-dir "${{ runner.temp }}/wheel_artifact" \
+            --output-dir "${{ runner.temp }}/libtorch_output" \
+            --platform windows \
+            --desired-cuda "$DESIRED_CUDA" \
+            --libtorch-variant "$LIBTORCH_VARIANT"
+      - uses: !{{ common.upload_artifact_action }}
+        if: always()
+        with:
+          name: !{{ lt_config["build_name"] }}
+          retention-days: 14
+          if-no-files-found: error
+          path: "${{ runner.temp }}/libtorch_output/"
+  {%- if branches == "nightly" %}
+  !{{ upload.upload_extracted_libtorch(lt_config, True) }}
+  {%- endif %}
+{%- endfor %}
diff --git a/.github/workflows/_bazel-build-test.yml b/.github/workflows/_bazel-build-test.yml
deleted file mode 100644
index eaebce92ba898..0000000000000
--- a/.github/workflows/_bazel-build-test.yml
+++ /dev/null
@@ -1,215 +0,0 @@
-name: bazel
-
-on:
-  workflow_call:
-    inputs:
-      build-environment:
-        required: true
-        type: string
-        description: Top-level label for what's being built/tested.
-      docker-image-name:
-        required: true
-        type: string
-        description: Name of the base docker image to build with.
-      cuda-version:
-        required: true
-        type: string
-        description: What CUDA version to build with (i.e. "11.7"), "cpu" for none.
-      sync-tag:
-        required: false
-        type: string
-        default: ""
-        description: |
-          If this is set, our linter will use this to make sure that every other
-          job with the same `sync-tag` is identical.
-      test-matrix:
-        required: true
-        type: string
-        description: |
-          A JSON description of what configs to run later on.
-      runner:
-        required: false
-        type: string
-        default: "linux.large"
-        description: Runner type
-
-env:
-  GIT_DEFAULT_BRANCH: ${{ github.event.repository.default_branch }}
-
-jobs:
-  filter:
-    if: github.repository_owner == 'pytorch'
-    runs-on: ${{ inputs.runner }}
-    outputs:
-      test-matrix: ${{ steps.filter.outputs.test-matrix }}
-      is-test-matrix-empty: ${{ steps.filter.outputs.is-test-matrix-empty }}
-      keep-going: ${{ steps.filter.outputs.keep-going }}
-      reenabled-issues: ${{ steps.filter.outputs.reenabled-issues }}
-    steps:
-      - name: Checkout PyTorch
-        uses: pytorch/pytorch/.github/actions/checkout-pytorch@main
-        with:
-          fetch-depth: 1
-          submodules: false
-
-      - name: Select all requested test configurations
-        id: filter
-        uses: ./.github/actions/filter-test-configs
-        with:
-          github-token: ${{ secrets.GITHUB_TOKEN }}
-          test-matrix: ${{ inputs.test-matrix }}
-
-  build-and-test:
-    needs: filter
-    # Don't run on forked repos.
-    if: github.repository_owner == 'pytorch' && needs.filter.outputs.is-test-matrix-empty == 'False'
-    strategy:
-      matrix: ${{ fromJSON(needs.filter.outputs.test-matrix) }}
-      fail-fast: false
-    runs-on: ${{ matrix.runner }}
-    steps:
-      - name: Setup SSH (Click me for login details)
-        uses: pytorch/test-infra/.github/actions/setup-ssh@main
-        with:
-          github-secret: ${{ secrets.GITHUB_TOKEN }}
-
-      # [see note: pytorch repo ref]
-      - name: Checkout PyTorch
-        uses: pytorch/pytorch/.github/actions/checkout-pytorch@main
-
-      - name: Setup Linux
-        uses: ./.github/actions/setup-linux
-
-      - name: Login to ECR
-        uses: ./.github/actions/ecr-login
-
-      - name: Calculate docker image
-        id: calculate-docker-image
-        uses: pytorch/test-infra/.github/actions/calculate-docker-image@main
-        with:
-          docker-image-name: ${{ inputs.docker-image-name }}
-
-      - name: Pull docker image
-        uses: pytorch/test-infra/.github/actions/pull-docker-image@main
-        with:
-          docker-image: ${{ steps.calculate-docker-image.outputs.docker-image }}
-
-      - name: Check if in a container runner
-        shell: bash
-        id: check_container_runner
-        run: echo "IN_CONTAINER_RUNNER=$(if [ -f /.inarc ] || [ -f /.incontainer ]; then echo true ; else echo false; fi)" >> "$GITHUB_OUTPUT"
-
-      - name: Install nvidia driver, nvidia-docker runtime, set GPU_FLAG
-        uses: pytorch/test-infra/.github/actions/setup-nvidia@main
-
-      - name: Output disk space left
-        run: |
-          sudo df -H
-
-      - name: Preserve github env variables for use in docker
-        run: |
-          env | grep '^GITHUB' >> "/tmp/github_env_${GITHUB_RUN_ID}"
-          env | grep '^CI' >> "/tmp/github_env_${GITHUB_RUN_ID}"
-
-      - name: Parse ref
-        id: parse-ref
-        run: .github/scripts/parse_ref.py
-
-      - name: Get workflow job id
-        id: get-job-id
-        uses: ./.github/actions/get-workflow-job-id
-        if: always()
-        with:
-          github-token: ${{ secrets.GITHUB_TOKEN }}
-
-      - name: Build
-        env:
-          BUILD_ENVIRONMENT: ${{ inputs.build-environment }}
-          PR_NUMBER: ${{ github.event.pull_request.number }}
-          BRANCH: ${{ steps.parse-ref.outputs.branch }}
-          GITHUB_REPOSITORY: ${{ github.repository }}
-          GITHUB_WORKFLOW: ${{ github.workflow }}
-          GITHUB_JOB: ${{ github.job }}
-          GITHUB_RUN_ID: ${{ github.run_id }}
-          GITHUB_RUN_NUMBER: ${{ github.run_number }}
-          GITHUB_RUN_ATTEMPT: ${{ github.run_attempt }}
-          JOB_ID: ${{ steps.get-job-id.outputs.job-id }}
-          REENABLED_ISSUES: ${{ needs.filter.outputs.reenabled-issues }}
-          # TODO duplicated
-          AWS_DEFAULT_REGION: us-east-1
-          SHA1: ${{ github.event.pull_request.head.sha || github.sha }}
-          SCCACHE_BUCKET: ossci-compiler-cache-circleci-v2
-          SCCACHE_REGION: us-east-1
-          TORCH_CUDA_ARCH_LIST: 5.2
-          DOCKER_IMAGE: ${{ steps.calculate-docker-image.outputs.docker-image }}
-          OUR_GITHUB_JOB_ID: ${{ steps.get-job-id.outputs.job-id }}
-          CUDA_VERSION: ${{ inputs.cuda-version }}
-        run: |
-          export SHARD_NUMBER=0
-          # detached container should get cleaned up by teardown_ec2_linux
-          # TODO: Stop building test binaries as part of the build phase
-          # Make sure we copy test results from bazel-testlogs symlink to
-          # a regular directory ./test/test-reports
-          # shellcheck disable=SC2086
-          container_name=$(docker run \
-            ${GPU_FLAG:-} \
-            -e AWS_DEFAULT_REGION \
-            -e BUILD_ENVIRONMENT \
-            -e GITHUB_ACTIONS \
-            -e GITHUB_REPOSITORY \
-            -e GITHUB_WORKFLOW \
-            -e GITHUB_JOB \
-            -e GITHUB_RUN_NUMBER \
-            -e GITHUB_RUN_ATTEMPT \
-            -e JOB_ID \
-            -e GIT_DEFAULT_BRANCH="$GIT_DEFAULT_BRANCH" \
-            -e SHARD_NUMBER \
-            -e NUM_TEST_SHARDS \
-            -e MAX_JOBS="$(nproc --ignore=2)" \
-            -e SCCACHE_BUCKET \
-            -e SCCACHE_REGION \
-            -e SKIP_SCCACHE_INITIALIZATION=1 \
-            -e REENABLED_ISSUES \
-            -e TORCH_CUDA_ARCH_LIST \
-            -e OUR_GITHUB_JOB_ID \
-            -e CUDA_VERSION \
-            --env-file="/tmp/github_env_${GITHUB_RUN_ID}" \
-            --security-opt seccomp=unconfined \
-            --cap-add=SYS_PTRACE \
-            --shm-size="1g" \
-            --tty \
-            --detach \
-            --user jenkins \
-            -v "${GITHUB_WORKSPACE}:/var/lib/jenkins/workspace" \
-            -w /var/lib/jenkins/workspace \
-            "${DOCKER_IMAGE}"
-          )
-          docker exec -t "${container_name}" sh -c '.ci/pytorch/build.sh'
-          echo "container_id=${container_name}" >> "${GITHUB_ENV}"
-
-      - name: Test
-        id: test
-        # Time out the test phase after 3.5 hours
-        timeout-minutes: 120
-        run: |
-          docker exec -t "${container_id}" sh -c '.ci/pytorch/test.sh && cp -Lr ./bazel-testlogs ./test/test-reports'
-
-      - name: Print remaining test logs
-        shell: bash
-        if: always() && steps.test.conclusion
-        run: |
-          cat test/**/*_toprint.log || true
-
-      - name: Chown workspace
-        uses: ./.github/actions/chown-workspace
-        if: always()
-
-      - name: Upload test artifacts
-        uses: ./.github/actions/upload-test-artifacts
-        if: always() && steps.test.conclusion && steps.test.conclusion != 'skipped'
-        with:
-          file-suffix: bazel-${{ github.job }}_${{ steps.get-job-id.outputs.job-id }}
-
-      - name: Teardown Linux
-        uses: pytorch/test-infra/.github/actions/teardown-linux@main
-        if: always()
diff --git a/.github/workflows/_binary-build-flash-attention-wheel-linux.yml b/.github/workflows/_binary-build-flash-attention-wheel-linux.yml
index 3fdc1dc4175c9..6836dc19edf9e 100644
--- a/.github/workflows/_binary-build-flash-attention-wheel-linux.yml
+++ b/.github/workflows/_binary-build-flash-attention-wheel-linux.yml
@@ -23,7 +23,7 @@ jobs:
   get-label-type:
     if: github.repository_owner == 'pytorch'
     name: get-label-type
-    uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@main
+    uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@release/2.12
     with:
       triggering_actor: ${{ github.triggering_actor }}
       issue_owner: ${{ github.event.pull_request.user.login || github.event.issue.user.login }}
@@ -83,21 +83,18 @@ jobs:
       TORCH_VERSION: "2.10.0"
     steps:
       - name: Setup SSH (Click me for login details)
-        uses: pytorch/test-infra/.github/actions/setup-ssh@main
+        uses: pytorch/test-infra/.github/actions/setup-ssh@release/2.12
         with:
           github-secret: ${{ secrets.GITHUB_TOKEN }}
           fail-silently: false
 
-      - name: Checkout PyTorch
-        uses: pytorch/pytorch/.github/actions/checkout-pytorch@main
+      - name: Setup Linux
+        uses: pytorch/pytorch/.github/actions/setup-linux@release/2.12
         with:
           submodules: true
 
-      - name: Setup Linux
-        uses: ./.github/actions/setup-linux
-
       - name: Pull Docker image
-        uses: pytorch/test-infra/.github/actions/pull-docker-image@main
+        uses: pytorch/test-infra/.github/actions/pull-docker-image@release/2.12
         with:
           docker-image: ${{ env.DOCKER_IMAGE }}
 
@@ -138,5 +135,5 @@ jobs:
           path: ${{ runner.temp }}/artifacts/*.whl
 
       - name: Teardown Linux
-        uses: pytorch/test-infra/.github/actions/teardown-linux@main
+        uses: pytorch/test-infra/.github/actions/teardown-linux@release/2.12
         if: always()
diff --git a/.github/workflows/_binary-build-flash-attention-wheel-windows.yml b/.github/workflows/_binary-build-flash-attention-wheel-windows.yml
index 4fc1dc8a53367..d3d5e96a42eb1 100644
--- a/.github/workflows/_binary-build-flash-attention-wheel-windows.yml
+++ b/.github/workflows/_binary-build-flash-attention-wheel-windows.yml
@@ -22,7 +22,7 @@ jobs:
   get-label-type:
     if: github.repository_owner == 'pytorch'
     name: get-label-type
-    uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@main
+    uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@release/2.12
     with:
       triggering_actor: ${{ github.triggering_actor }}
       issue_owner: ${{ github.event.pull_request.user.login || github.event.issue.user.login }}
@@ -71,12 +71,12 @@ jobs:
           git config --global core.ignorecase false
           git config --global core.fsmonitor false
       - name: Setup SSH (Click me for login details)
-        uses: pytorch/test-infra/.github/actions/setup-ssh@main
+        uses: pytorch/test-infra/.github/actions/setup-ssh@release/2.12
         with:
           github-secret: ${{ secrets.GITHUB_TOKEN }}
 
       - name: Checkout PyTorch
-        uses: pytorch/pytorch/.github/actions/checkout-pytorch@main
+        uses: pytorch/pytorch/.github/actions/checkout-pytorch@release/2.12
         with:
           no-sudo: true
           submodules: true
@@ -90,7 +90,7 @@ jobs:
           echo "CUDA_HOME=${CUDA_PATH}" >> "${GITHUB_ENV}"
           echo "${CUDA_PATH}/bin" >> "${GITHUB_PATH}"
       - name: Setup MSVC
-        uses: ilammy/msvc-dev-cmd@v1
+        uses: ilammy/msvc-dev-cmd@0b201ec74fa43914dc39ae48a89fd1d8cb592756 # v1
 
       - name: Remove link.exe conflict
         run: rm -f /usr/bin/link
diff --git a/.github/workflows/_binary-build-linux.yml b/.github/workflows/_binary-build-linux.yml
index 115f6572441a1..88d5177ee093d 100644
--- a/.github/workflows/_binary-build-linux.yml
+++ b/.github/workflows/_binary-build-linux.yml
@@ -142,19 +142,19 @@ jobs:
 
       - name: "[FB EMPLOYEES] Enable SSH (Click me for login details)"
         if: inputs.build_environment != 'linux-s390x-binary-manywheel'
-        uses: pytorch/test-infra/.github/actions/setup-ssh@main
+        uses: pytorch/test-infra/.github/actions/setup-ssh@release/2.12
         continue-on-error: true
         with:
           github-secret: ${{ secrets.github-token }}
 
       - name: Checkout PyTorch
-        uses: pytorch/pytorch/.github/actions/checkout-pytorch@main
+        uses: pytorch/pytorch/.github/actions/checkout-pytorch@release/2.12
         with:
           no-sudo: ${{ inputs.build_environment == 'linux-aarch64-binary-manywheel' || inputs.build_environment == 'linux-s390x-binary-manywheel' }}
 
       - name: Setup Linux
         if: inputs.build_environment != 'linux-s390x-binary-manywheel'
-        uses: ./.github/actions/setup-linux
+        uses: pytorch/pytorch/.github/actions/setup-linux@release/2.12
 
       - name: Login to ECR
         if: inputs.build_environment != 'linux-s390x-binary-manywheel'
@@ -179,7 +179,7 @@ jobs:
       - name: Checkout PyTorch to pytorch dir
         uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
         with:
-          ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
+          fetch-depth: 2
           submodules: recursive
           path: pytorch
           show-progress: false
@@ -214,7 +214,7 @@ jobs:
       - name: Calculate docker image
         id: calculate-docker-image
         if: ${{ steps.filter.outputs.is-test-matrix-empty == 'False' && inputs.build_environment != 'linux-s390x-binary-manywheel' }}
-        uses: pytorch/test-infra/.github/actions/calculate-docker-image@main
+        uses: pytorch/test-infra/.github/actions/calculate-docker-image@release/2.12
         with:
           # If doing this in main or release branch, use docker.io. Otherwise
           # use ECR
@@ -224,11 +224,12 @@ jobs:
           # The build.sh script in this folder is not actually the correct one,
           # this is just needed for sha calculation
           docker-build-dir: .ci/docker
+          docker-build-script: ${{ contains(inputs.DOCKER_IMAGE, 'manylinux') && 'manywheel/build.sh' || 'build.sh' }}
           working-directory: pytorch
 
       - name: Pull Docker image
         if: ${{ steps.filter.outputs.is-test-matrix-empty == 'False' && inputs.build_environment != 'linux-s390x-binary-manywheel' }}
-        uses: pytorch/test-infra/.github/actions/pull-docker-image@main
+        uses: pytorch/test-infra/.github/actions/pull-docker-image@release/2.12
         with:
           docker-image: ${{ steps.calculate-docker-image.outputs.docker-image }}
 
@@ -255,12 +256,12 @@ jobs:
             -e PYTORCH_EXTRA_INSTALL_REQUIREMENTS \
             --tty \
             --detach \
-            -v "${GITHUB_WORKSPACE}/pytorch:/pytorch" \
+            -v "${GITHUB_WORKSPACE}:/pytorch" \
             -v "${RUNNER_TEMP}/artifacts:/artifacts" \
             -w / \
             "${DOCKER_IMAGE}"
           )
-          docker exec -t -w "${PYTORCH_ROOT}" "${container_name}" bash -c "bash .circleci/scripts/binary_populate_env.sh"
+          docker exec -t -w "${PYTORCH_ROOT}" "${container_name}" bash -c "bash .ci/pytorch/binary_populate_env.sh"
           # Unified build script for all architectures (x86_64, aarch64, s390x)
           docker exec -t "${container_name}" bash -c "source ${BINARY_ENV_FILE} && bash /pytorch/.ci/${{ inputs.PACKAGE_TYPE }}/build.sh"
 
@@ -281,7 +282,7 @@ jobs:
 
       - name: Teardown Linux
         if: always() && inputs.build_environment != 'linux-s390x-binary-manywheel'
-        uses: pytorch/test-infra/.github/actions/teardown-linux@main
+        uses: pytorch/test-infra/.github/actions/teardown-linux@release/2.12
 
       - name: Chown workspace
         if: always() && inputs.build_environment != 'linux-s390x-binary-manywheel'
diff --git a/.github/workflows/_binary-test-linux.yml b/.github/workflows/_binary-test-linux.yml
index 3dd8235a4c4f2..d9270ea39313d 100644
--- a/.github/workflows/_binary-test-linux.yml
+++ b/.github/workflows/_binary-test-linux.yml
@@ -63,7 +63,7 @@ on:
       runs_on:
         required: true
         type: string
-        description: Hardware to run this job on. Valid values are linux.4xlarge, linux.4xlarge.nvidia.gpu, linux.arm64.2xlarge, and linux.rocm.gpu
+        description: Hardware to run this job on. Valid values are linux.4xlarge, linux.g4dn.4xlarge.nvidia.gpu, linux.arm64.2xlarge, and linux.rocm.gpu
     secrets:
       github-token:
         required: true
@@ -125,20 +125,20 @@ jobs:
 
       - name: "[FB EMPLOYEES] Enable SSH (Click me for login details)"
         if: inputs.build_environment != 'linux-s390x-binary-manywheel'
-        uses: pytorch/test-infra/.github/actions/setup-ssh@main
+        uses: pytorch/test-infra/.github/actions/setup-ssh@release/2.12
         continue-on-error: true
         with:
           github-secret: ${{ secrets.github-token }}
 
         # Setup the environment
       - name: Checkout PyTorch
-        uses: pytorch/pytorch/.github/actions/checkout-pytorch@main
+        uses: pytorch/pytorch/.github/actions/checkout-pytorch@release/2.12
         with:
           no-sudo: ${{ inputs.build_environment == 'linux-aarch64-binary-manywheel' || inputs.build_environment == 'linux-s390x-binary-manywheel' }}
 
       - name: Setup Linux
         if: inputs.build_environment != 'linux-s390x-binary-manywheel'
-        uses: ./.github/actions/setup-linux
+        uses: pytorch/pytorch/.github/actions/setup-linux@release/2.12
 
       - name: Login to ECR
         if: inputs.build_environment != 'linux-s390x-binary-manywheel'
@@ -153,7 +153,6 @@ jobs:
       - name: Checkout PyTorch to pytorch dir
         uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
         with:
-          ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
           submodules: recursive
           show-progress: false
           path: pytorch
@@ -185,7 +184,7 @@ jobs:
 
       - name: Install nvidia driver, nvidia-docker runtime, set GPU_FLAG
         id: install-nvidia-driver
-        uses: pytorch/test-infra/.github/actions/setup-nvidia@main
+        uses: pytorch/test-infra/.github/actions/setup-nvidia@release/2.12
         if: ${{ steps.filter.outputs.is-test-matrix-empty == 'False' }}
 
       - name: configure aws credentials
@@ -200,7 +199,7 @@ jobs:
       - name: Calculate docker image
         id: calculate-docker-image
         if: ${{ steps.filter.outputs.is-test-matrix-empty == 'False' && inputs.build_environment != 'linux-s390x-binary-manywheel' }}
-        uses: pytorch/test-infra/.github/actions/calculate-docker-image@main
+        uses: pytorch/test-infra/.github/actions/calculate-docker-image@release/2.12
         with:
           docker-registry: ${{ startsWith(github.event.ref, 'refs/tags/ciflow/') && '308535385114.dkr.ecr.us-east-1.amazonaws.com' || 'docker.io' }}
           docker-image-name: ${{ inputs.DOCKER_IMAGE }}
@@ -210,7 +209,7 @@ jobs:
 
       - name: Pull Docker image
         if: ${{ steps.filter.outputs.is-test-matrix-empty == 'False' && inputs.build_environment != 'linux-s390x-binary-manywheel' }}
-        uses: pytorch/test-infra/.github/actions/pull-docker-image@main
+        uses: pytorch/test-infra/.github/actions/pull-docker-image@release/2.12
         with:
           docker-image: ${{ steps.calculate-docker-image.outputs.docker-image }}
 
@@ -222,7 +221,7 @@ jobs:
 
       - name: Teardown Linux
         if: always() && inputs.build_environment != 'linux-s390x-binary-manywheel'
-        uses: pytorch/test-infra/.github/actions/teardown-linux@main
+        uses: pytorch/test-infra/.github/actions/teardown-linux@release/2.12
 
       - name: Chown workspace
         if: always() && inputs.build_environment != 'linux-s390x-binary-manywheel'
diff --git a/.github/workflows/_binary-upload-flash-attention-wheel.yml b/.github/workflows/_binary-upload-flash-attention-wheel.yml
index 2d66082bcfc43..a477325dfe103 100644
--- a/.github/workflows/_binary-upload-flash-attention-wheel.yml
+++ b/.github/workflows/_binary-upload-flash-attention-wheel.yml
@@ -67,7 +67,7 @@ jobs:
         run: |
           set -ex
           for subfolder in ${{ matrix.upload_subfolders }}; do
-            UPLOAD_SUBFOLDER="${subfolder}" bash .circleci/scripts/binary_upload.sh
+            UPLOAD_SUBFOLDER="${subfolder}" bash .ci/pytorch/binary_upload.sh
           done
 
   upload-wheel-windows:
@@ -114,5 +114,5 @@ jobs:
         run: |
           set -ex
           for subfolder in ${{ matrix.upload_subfolders }}; do
-            UPLOAD_SUBFOLDER="${subfolder}" bash .circleci/scripts/binary_upload.sh
+            UPLOAD_SUBFOLDER="${subfolder}" bash .ci/pytorch/binary_upload.sh
           done
diff --git a/.github/workflows/_binary-upload.yml b/.github/workflows/_binary-upload.yml
index 636b76d42931a..c7eb31efdb050 100644
--- a/.github/workflows/_binary-upload.yml
+++ b/.github/workflows/_binary-upload.yml
@@ -55,10 +55,20 @@ on:
       github-token:
         required: true
         description: Github Token
+      R2_ACCOUNT_ID:
+        required: false
+        description: Cloudflare R2 account ID for nightly uploads
+      R2_ACCESS_KEY_ID:
+        required: false
+        description: Cloudflare R2 access key ID for nightly uploads
+      R2_SECRET_ACCESS_KEY:
+        required: false
+        description: Cloudflare R2 secret access key for nightly uploads
 
 jobs:
   upload:
     runs-on: ubuntu-22.04
+    environment: ${{ (github.event_name == 'push' && github.event.ref == 'refs/heads/nightly') && 'nightly-wheel-upload' || '' }}
     container:
       image: continuumio/miniconda3:4.12.0
     env:
@@ -81,7 +91,7 @@ jobs:
       SHA1: ${{ github.event.pull_request.head.sha || github.sha }}
     steps:
       - name: Checkout PyTorch
-        uses: pytorch/pytorch/.github/actions/checkout-pytorch@main
+        uses: pytorch/pytorch/.github/actions/checkout-pytorch@release/2.12
         with:
           no-sudo: true
 
@@ -130,6 +140,10 @@ jobs:
           PKG_DIR: "${{ runner.temp }}/artifacts"
           UPLOAD_SUBFOLDER: "${{ env.DESIRED_CUDA }}"
           BUILD_NAME: ${{ inputs.build_name }}
+          R2_UPLOAD: ${{ (github.event_name == 'push' && github.event.ref == 'refs/heads/nightly') && 'true' || '' }}
+          R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+          R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+          R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
         run: |
             set -ex
-            bash .circleci/scripts/binary_upload.sh
+            bash .ci/pytorch/binary_upload.sh
diff --git a/.github/workflows/_docs.yml b/.github/workflows/_docs.yml
index f5cb186c6f189..a4e51b01ee4a8 100644
--- a/.github/workflows/_docs.yml
+++ b/.github/workflows/_docs.yml
@@ -47,6 +47,26 @@ on:
         description: prefix for runner label
         type: string
         default: ""
+      use-arc:
+        required: false
+        type: boolean
+        default: false
+        description: If true, use ARC (OSDC) runner path instead of EC2.
+      python-version:
+        required: false
+        type: string
+        default: ""
+        description: Python version to use for the OSDC build.
+      compiler:
+        required: false
+        type: string
+        default: ""
+        description: Compiler to use for the OSDC build.
+      cuda-version:
+        required: false
+        type: string
+        default: ""
+        description: CUDA version to use for the OSDC build.
     secrets:
       GH_PYTORCHBOT_TOKEN:
         required: false
@@ -55,7 +75,7 @@ on:
 jobs:
   build-docs:
     # Don't run on forked repos.
-    if: github.repository_owner == 'pytorch'
+    if: github.repository_owner == 'pytorch' && !inputs.use-arc
     runs-on: ${{ matrix.runner }}
     environment: ${{ (github.ref == 'refs/heads/main' || startsWith(github.event.ref, 'refs/tags/v')) && 'pytorchbot-env' || '' }}
     strategy:
@@ -73,14 +93,13 @@ jobs:
             timeout-minutes: 360
           - docs_type: python
             runner: ${{ inputs.runner_prefix }}linux.c7i.2xlarge
-            # It takes less than 30m to finish python docs unless there are issues
-            timeout-minutes: 30
+            timeout-minutes: 60
     # Set a fixed name for this job instead of using the current matrix-generated name, i.e. build-docs (cpp, linux.12xlarge, 180)
     # The current name requires updating the database last docs push query from test-infra every time the matrix is updated
     name: build-docs-${{ matrix.docs_type }}-${{ inputs.push }}
     steps:
       - name: Setup SSH (Click me for login details)
-        uses: pytorch/test-infra/.github/actions/setup-ssh@main
+        uses: pytorch/test-infra/.github/actions/setup-ssh@release/2.12
         with:
           github-secret: ${{ secrets.GITHUB_TOKEN }}
           instructions: |
@@ -89,12 +108,13 @@ jobs:
             To start Python docs build type:
               cd docs && make html && make coverage
 
-      # [see note: pytorch repo ref]
-      - name: Checkout PyTorch
-        uses: pytorch/pytorch/.github/actions/checkout-pytorch@main
-
       - name: Setup Linux
-        uses: ./.github/actions/setup-linux
+        uses: pytorch/pytorch/.github/actions/setup-linux@release/2.12
+        with:
+          # Docs build runs `git log --follow` per page for "Created/Updated"
+          # dates; a treeless clone triggers lazy fetches and makes the build
+          # hang. Use a full checkout for this job.
+          checkout-mode: normal
 
       - name: Login to ECR
         uses: ./.github/actions/ecr-login
@@ -103,12 +123,12 @@ jobs:
 
       - name: Calculate docker image
         id: calculate-docker-image
-        uses: pytorch/test-infra/.github/actions/calculate-docker-image@main
+        uses: pytorch/test-infra/.github/actions/calculate-docker-image@release/2.12
         with:
           docker-image-name: ${{ inputs.docker-image }}
 
       - name: Pull docker image
-        uses: pytorch/test-infra/.github/actions/pull-docker-image@main
+        uses: pytorch/test-infra/.github/actions/pull-docker-image@release/2.12
         with:
           docker-image: ${{ steps.calculate-docker-image.outputs.docker-image }}
 
@@ -160,6 +180,7 @@ jobs:
             -e RUN_DOXYGEN \
             -e WITH_PUSH \
             --env-file="/tmp/github_env_${GITHUB_RUN_ID}" \
+            -e GITHUB_WORKSPACE=/var/lib/jenkins/workspace \
             --security-opt seccomp=unconfined \
             --cap-add=SYS_PTRACE \
             --tty \
@@ -204,6 +225,40 @@ jobs:
           path: cppdocs/
           s3-prefix: pytorch/pytorch/${{ github.event.pull_request.number }}/cppdocs
 
+      - name: Post C++ Docs Coverage Comment
+        if: ${{ github.event_name == 'pull_request' && matrix.docs_type == 'cpp' && steps.build-docs.outcome == 'success' }}
+        env:
+          GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+        run: |
+          PR_NUM="${{ github.event.pull_request.number }}"
+          MARKER="<!-- cpp-docs-coverage -->"
+
+          # Only post if the PR touches docs/cpp
+          if ! gh pr diff "$PR_NUM" --name-only | grep -q '^docs/cpp/'; then
+            echo "No changes to docs/cpp/, skipping coverage comment."
+            exit 0
+          fi
+
+          # Skip if we already posted
+          if gh pr view "$PR_NUM" --json comments --jq '.comments[].body' | grep -q "$MARKER"; then
+            echo "Coverage comment already posted, skipping."
+            exit 0
+          fi
+
+          body="${MARKER}
+          ## C++ Docs
+
+          - [Doc Preview](https://docs-preview.pytorch.org/pytorch/pytorch/${PR_NUM}/cppdocs/index.html)
+          - [API Coverage Report](https://docs-preview.pytorch.org/pytorch/pytorch/${PR_NUM}/cppdocs/_coverage/cpp_coverage.txt)
+          - [HTML Issues Report](https://docs-preview.pytorch.org/pytorch/pytorch/${PR_NUM}/cppdocs/_coverage/cpp_html_issues.txt)"
+
+          if [ -f cppdocs/_coverage/index.html ]; then
+            body="${body}
+          - [Coverxygen Report (interactive)](https://docs-preview.pytorch.org/pytorch/pytorch/${PR_NUM}/cppdocs/_coverage/index.html)"
+          fi
+
+          gh pr comment "$PR_NUM" --body "$body"
+
       - name: Upload C++ Docs Preview (nightly dry-run)
         if: ${{ !inputs.push && github.event_name != 'pull_request' && matrix.docs_type == 'cpp' && steps.build-docs.outcome == 'success' }}
         run: |
@@ -213,5 +268,161 @@ jobs:
           echo "https://docs-preview.pytorch.org/pytorch/pytorch/nightly-${{ github.sha }}/cppdocs/index.html"
 
       - name: Teardown Linux
-        uses: pytorch/test-infra/.github/actions/teardown-linux@main
+        uses: pytorch/test-infra/.github/actions/teardown-linux@release/2.12
         if: always()
+
+  build-docs-osdc:
+    if: github.repository_owner == 'pytorch' && inputs.use-arc
+    permissions:
+      id-token: write
+      contents: read
+      actions: read
+    runs-on: ${{ matrix.runner }}
+    container:
+      image: ${{ inputs.docker-image }}
+    strategy:
+      fail-fast: false
+      matrix:
+        include:
+          - docs_type: cpp
+            runner: ${{ inputs.runner_prefix }}l-x86iavx512-48-384
+            timeout-minutes: 360
+          - docs_type: python
+            runner: ${{ inputs.runner_prefix }}l-x86iavx512-16-128
+            timeout-minutes: 60
+    name: build-docs-${{ matrix.docs_type }}-${{ inputs.push }}
+    steps:
+      - name: Setup Linux
+        id: setup-linux
+        uses: pytorch/pytorch/.github/actions/setup-linux@release/2.12
+        with:
+          use-arc: true
+          python-version: ${{ inputs.python-version }}
+          compiler: ${{ inputs.compiler }}
+          cuda-version: ${{ inputs.cuda-version }}
+          submodules: false
+          # See comment in the EC2 build-docs job above.
+          checkout-mode: normal
+          github-token: ${{ secrets.GITHUB_TOKEN }}
+
+      - name: Configure AWS credentials
+        id: aws-creds
+        continue-on-error: true
+        uses: aws-actions/configure-aws-credentials@8df5847569e6427dd6c4fb1cf565c83acfa8afa7
+        with:
+          role-to-assume: arn:aws:iam::308535385114:role/arc
+          aws-region: us-east-1
+          role-duration-seconds: 18000
+
+      - name: Download build artifacts
+        uses: pytorch/pytorch/.github/actions/download-build-artifacts@release/2.12
+        with:
+          name: ${{ inputs.build-environment }}
+          s3-bucket: ${{ inputs.s3-bucket }}
+          use-gha: ${{ steps.aws-creds.outcome != 'success' }}
+
+      - name: Generate netrc (only for docs-push)
+        if: inputs.push
+        env:
+          GITHUB_PYTORCHBOT_TOKEN: ${{ secrets.GH_PYTORCHBOT_TOKEN }}
+        run: |
+          rm -rf "${HOME}/.netrc"
+          echo "machine github.com" > "${HOME}/.netrc"
+          echo "login pytorchbot" >> "${HOME}/.netrc"
+          echo "password ${GITHUB_PYTORCHBOT_TOKEN}" >> "${HOME}/.netrc"
+
+      - name: Build ${{ matrix.docs_type }} docs
+        timeout-minutes: ${{ matrix.timeout-minutes }}
+        id: build-docs
+        env:
+          WITH_PUSH: ${{ inputs.push }}
+          DOCS_TYPE: ${{ matrix.docs_type }}
+          RUN_DOXYGEN: ${{ inputs.run-doxygen }}
+          BUILD_ENVIRONMENT: ${{ inputs.build-environment }}
+          SHA1: ${{ github.sha }}
+        shell: bash
+        run: |
+          set -ex
+          if [[ "${GITHUB_REF}" =~ ^refs/tags/v([0-9]+\.[0-9]+)\.* ]]; then
+            target="${BASH_REMATCH[1]}"
+          else
+            target="main"
+          fi
+          export DOCS_VERSION="${target}"
+          pip install $(echo dist/*.whl)[opt-einsum]
+          ./.ci/pytorch/${DOCS_TYPE}_doc_push_script.sh
+
+      - name: Upload Python Docs Preview
+        uses: seemethere/upload-artifact-s3@baba72d0712b404f646cebe0730933554ebce96a # v5.1.0
+        if: ${{ steps.aws-creds.outcome == 'success' && github.event_name == 'pull_request' && matrix.docs_type == 'python' && steps.build-docs.outcome == 'success' }}
+        with:
+          retention-days: 14
+          s3-bucket: doc-previews
+          if-no-files-found: error
+          path: pytorch_docs/main/
+          s3-prefix: pytorch/pytorch/${{ github.event.pull_request.number }}
+
+      - name: Upload C++ Docs Preview
+        uses: seemethere/upload-artifact-s3@baba72d0712b404f646cebe0730933554ebce96a # v5.1.0
+        if: ${{ steps.aws-creds.outcome == 'success' && github.event_name == 'pull_request' && matrix.docs_type == 'cpp' && steps.build-docs.outcome == 'success' }}
+        with:
+          retention-days: 14
+          if-no-files-found: error
+          s3-bucket: doc-previews
+          path: cppdocs/
+          s3-prefix: pytorch/pytorch/${{ github.event.pull_request.number }}/cppdocs
+
+      - name: Post C++ Docs Coverage Comment
+        if: ${{ steps.aws-creds.outcome == 'success' && github.event_name == 'pull_request' && matrix.docs_type == 'cpp' && steps.build-docs.outcome == 'success' }}
+        env:
+          GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+        run: |
+          PR_NUM="${{ github.event.pull_request.number }}"
+          MARKER="<!-- cpp-docs-coverage -->"
+
+          # Only post if the PR touches docs/cpp
+          if ! gh pr diff "$PR_NUM" --name-only | grep -q '^docs/cpp/'; then
+            echo "No changes to docs/cpp/, skipping coverage comment."
+            exit 0
+          fi
+
+          # Skip if we already posted
+          if gh pr view "$PR_NUM" --json comments --jq '.comments[].body' | grep -q "$MARKER"; then
+            echo "Coverage comment already posted, skipping."
+            exit 0
+          fi
+
+          body="${MARKER}
+          ## C++ Docs
+
+          - [Doc Preview](https://docs-preview.pytorch.org/pytorch/pytorch/${PR_NUM}/cppdocs/index.html)
+          - [API Coverage Report](https://docs-preview.pytorch.org/pytorch/pytorch/${PR_NUM}/cppdocs/_coverage/cpp_coverage.txt)
+          - [HTML Issues Report](https://docs-preview.pytorch.org/pytorch/pytorch/${PR_NUM}/cppdocs/_coverage/cpp_html_issues.txt)"
+
+          if [ -f cppdocs/_coverage/index.html ]; then
+            body="${body}
+          - [Coverxygen Report (interactive)](https://docs-preview.pytorch.org/pytorch/pytorch/${PR_NUM}/cppdocs/_coverage/index.html)"
+          fi
+
+          gh pr comment "$PR_NUM" --body "$body"
+
+      - name: Upload C++ Docs Preview (nightly dry-run)
+        if: ${{ steps.aws-creds.outcome == 'success' && !inputs.push && github.event_name != 'pull_request' && matrix.docs_type == 'cpp' && steps.build-docs.outcome == 'success' }}
+        run: |
+          # Unlike EC2 runners, the OSDC container doesn't have aws CLI pre-installed
+          pip install awscli==1.29.40
+          aws s3 cp cppdocs/ s3://doc-previews/pytorch/pytorch/nightly-${{ github.sha }}/cppdocs --recursive --quiet
+          echo "C++ docs preview available at:"
+          echo "https://docs-preview.pytorch.org/pytorch/pytorch/nightly-${{ github.sha }}/cppdocs/index.html"
+
+  memory-viz-tests:
+    # Tests for torch memory visualizer
+    # tag shangdiy for any issue with this test
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - uses: actions/setup-node@v4
+        with:
+          node-version: '20'
+      - name: Run MemoryViz JS tests (nonretryable)
+        run: node test/profiler/test_memory_viz.js
diff --git a/.github/workflows/_get-changed-files.yml b/.github/workflows/_get-changed-files.yml
index 599d7a3277fe1..f3367916f8783 100644
--- a/.github/workflows/_get-changed-files.yml
+++ b/.github/workflows/_get-changed-files.yml
@@ -25,6 +25,7 @@ jobs:
         env:
           GH_TOKEN: ${{ github.token }}
         run: |
+          set -eo pipefail
           # Check if we're in a pull request context
           if [ "${{ github.event_name }}" = "pull_request" ] || [ "${{ github.event_name }}" = "pull_request_target" ]; then
             echo "Running in PR context"
@@ -37,8 +38,15 @@ jobs:
               echo "all_files input is true, returning all files"
               echo "changed-files=*" >> "$GITHUB_OUTPUT"
             else
-              # Use gh CLI to get changed files in the PR with explicit repo
-              CHANGED_FILES=$(gh api repos/${{ github.repository }}/pulls/$PR_NUMBER/files --paginate --jq '.[] | select(.status != "removed") | .filename' | tr '\n' ' ' | sed 's/ $//')
+              # Use gh CLI to get changed files in the PR with explicit repo.
+              # pipefail ensures that if gh fails (e.g. 403 rate limit), the
+              # error propagates through the pipe instead of being silently
+              # captured as the file list and injected into downstream lint jobs.
+              if ! CHANGED_FILES=$(gh api repos/${{ github.repository }}/pulls/$PR_NUMBER/files --paginate --jq '.[] | select(.status != "removed") | .filename' | tr '\n' ' ' | sed 's/ $//'); then
+                echo "Failed to get changed files from GitHub API, falling back to all files"
+                echo "changed-files=*" >> "$GITHUB_OUTPUT"
+                exit 0
+              fi
 
               # See https://github.com/pytorch/pytorch/pull/134215#issuecomment-2332128790
               PYI_FILES_TO_ADD=""
diff --git a/.github/workflows/_link_check.yml b/.github/workflows/_link_check.yml
index 014e6106b0730..efa5b433947ef 100644
--- a/.github/workflows/_link_check.yml
+++ b/.github/workflows/_link_check.yml
@@ -11,15 +11,10 @@ on:
 jobs:
   lint-urls:
     if: ${{ github.event_name != 'pull_request' || !contains(github.event.pull_request.labels.*.name, 'skip-url-lint') }}
-    uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main
+    uses: ./.github/workflows/_lint.yml
     with:
-      job-name: lint-urls
-      timeout: 120
-      runner: ${{ inputs.runner }}linux.2xlarge
-      docker-image: ci-image:pytorch-linux-jammy-linter
-      fetch-depth: 0
-      submodules: false
-      ref: ${{ inputs.ref }}
+      runner: mt-l-x86iamx-8-16
+      docker-image: ghcr.io/pytorch/test-infra:cpu-x86_64-810d48d
       script: |
         ./scripts/lint_urls.sh $(
           if [ "${{ github.event_name }}" = "pull_request" ]; then
@@ -37,15 +32,10 @@ jobs:
 
   lint-xrefs:
     if: ${{ github.event_name != 'pull_request' || !contains(github.event.pull_request.labels.*.name, 'skip-xref-lint') }}
-    uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main
+    uses: ./.github/workflows/_lint.yml
     with:
-      job-name: lint-xrefs
-      timeout: 60
-      runner: ${{ inputs.runner }}linux.2xlarge
-      docker-image: ci-image:pytorch-linux-jammy-linter
-      fetch-depth: 0
-      submodules: false
-      ref: ${{ inputs.ref }}
+      runner: mt-l-x86iamx-8-16
+      docker-image: ghcr.io/pytorch/test-infra:cpu-x86_64-810d48d
       script: |
         ./scripts/lint_xrefs.sh $(
           if [ "${{ github.event_name }}" = "pull_request" ]; then
diff --git a/.github/workflows/_lint.yml b/.github/workflows/_lint.yml
new file mode 100644
index 0000000000000..4ae37c0476ac8
--- /dev/null
+++ b/.github/workflows/_lint.yml
@@ -0,0 +1,82 @@
+name: Run linters
+
+on:
+  workflow_call:
+    inputs:
+      runner:
+        required: true
+        type: string
+        description: The runner to use
+      docker-image:
+        required: true
+        type: string
+        description: The Docker image to use
+      script:
+        required: true
+        type: string
+        description: The linter script to run
+
+jobs:
+  lint:
+    runs-on: ${{ inputs.runner }}
+    container:
+      image: ${{ inputs.docker-image }}
+    timeout-minutes: 120
+    steps:
+      - name: Fix Git ownership
+        shell: bash
+        run: |
+          git config --global --add safe.directory "$GITHUB_WORKSPACE"
+
+      - name: Checkout PyTorch
+        uses: pytorch/pytorch/.github/actions/checkout-pytorch@release/2.12
+        with:
+          fetch-depth: 0
+          submodules: true
+          no-sudo: true
+          checkout-mode: treeless
+
+      - name: Setup uv
+        uses: pytorch/test-infra/.github/actions/setup-uv@release/2.12
+        with:
+          python-version: "3.12"
+          activate-environment: true
+
+      - name: Install pip requirements
+        shell: bash
+        run: |
+          set -eux
+          uv pip install -r .ci/docker/requirements-ci.txt
+
+      - name: Install system requirements
+        shell: bash
+        run: |
+          set -eux
+          # Update repository
+          dnf install -y doxygen graphviz nodejs npm
+
+      - name: Install Node.js packages
+        shell: bash
+        run: |
+          set -eux
+          npm install -g markdown-toc
+
+      - name: Prepare lintrunner
+        shell: bash
+        run: |
+          set -eux
+          lintrunner init
+
+      - name: Run linter
+        shell: bash
+        env:
+          SCRIPT: ${{ inputs.script }}
+        run: |
+          {
+            echo "#!/usr/bin/env bash";
+            echo "set -eou pipefail";
+            echo "${SCRIPT}";
+          } > "${RUNNER_TEMP}/linter_script"
+
+          # Execute the linter script
+          bash "${RUNNER_TEMP}/linter_script"
diff --git a/.github/workflows/_linux-build.yml b/.github/workflows/_linux-build.yml
index 25501b59c7851..07872fb5038d5 100644
--- a/.github/workflows/_linux-build.yml
+++ b/.github/workflows/_linux-build.yml
@@ -26,7 +26,7 @@ on:
       cuda-arch-list:
         required: false
         type: string
-        default: "5.2"
+        default: "7.5"
         description: |
           List of CUDA architectures CI build should target.
       runner_prefix:
@@ -103,34 +103,49 @@ on:
         required: false
         type: string
         default: ""
+      use-arc:
+        required: false
+        type: boolean
+        default: false
+        description: If true, use ARC (OSDC) runner path instead of EC2.
+      python-version:
+        required: false
+        type: string
+        default: ""
+        description: Python version to use for the OSDC build.
+      compiler:
+        required: false
+        type: string
+        default: ""
+        description: Compiler to use for the OSDC build.
+      cuda-version:
+        required: false
+        type: string
+        default: ""
+        description: CUDA version to use for the OSDC build.
 
     secrets:
       HUGGING_FACE_HUB_TOKEN:
         required: false
         description: |
           HF Auth token to avoid rate limits when downloading models or datasets from hub
-      SCRIBE_GRAPHQL_ACCESS_TOKEN:
-        required: false
-        description: |
-          FB app token to write to scribe endpoint
 
     outputs:
       docker-image:
-        value: ${{ jobs.build.outputs.docker-image }}
+        value: ${{ jobs.build.outputs.docker-image || jobs.build-osdc.outputs.docker-image }}
         description: The docker image containing the built PyTorch.
       test-matrix:
-        value: ${{ jobs.build.outputs.test-matrix }}
+        value: ${{ jobs.build.outputs.test-matrix || jobs.build-osdc.outputs.test-matrix }}
         description: An optional JSON description of what test configs to run later on.
       build-environment:
-        value: ${{ jobs.build.outputs.build-environment }}
+        value: ${{ jobs.build.outputs.build-environment || jobs.build-osdc.outputs.build-environment }}
         description: Top-level label for what's being built/tested.
 
 jobs:
   build:
-    environment: ${{ github.ref == 'refs/heads/main' && 'scribe-protected' || startsWith(github.ref, 'refs/heads/release/') && 'scribe-protected' || contains(github.event.pull_request.labels.*.name, 'ci-scribe') && 'scribe-pr' || '' }}
     # Don't run on forked repos
-    if: github.repository_owner == 'pytorch'
-    runs-on: ${{ inputs.runner_prefix}}${{ inputs.runner }}
+    if: github.repository_owner == 'pytorch' && !inputs.use-arc
+    runs-on: ${{ inputs.runner_prefix }}${{ inputs.runner }}
     timeout-minutes: 480
     outputs:
       docker-image: ${{ steps.calculate-docker-image.outputs.docker-image }}
@@ -138,7 +153,7 @@ jobs:
       build-environment: ${{ inputs.build-environment }}
     steps:
       - name: Setup SSH (Click me for login details)
-        uses: pytorch/test-infra/.github/actions/setup-ssh@main
+        uses: pytorch/test-infra/.github/actions/setup-ssh@release/2.12
         if: inputs.build-environment != 'linux-s390x-binary-manywheel'
         with:
           github-secret: ${{ secrets.GITHUB_TOKEN }}
@@ -146,18 +161,14 @@ jobs:
             Build is done inside the container, to start an interactive session run:
               docker exec -it $(docker container ps --format '{{.ID}}') bash
 
-      # [pytorch repo ref]
-      # Use a pytorch/pytorch reference instead of a reference to the local
-      # checkout because when we run this action we don't *have* a local
-      # checkout. In other cases you should prefer a local checkout.
-      - name: Checkout PyTorch
-        uses: pytorch/pytorch/.github/actions/checkout-pytorch@main
-        with:
-          no-sudo: true
-
       - name: Setup Linux
-        uses: ./.github/actions/setup-linux
-        if: inputs.build-environment != 'linux-s390x-binary-manywheel'
+        id: setup-linux
+        uses: pytorch/pytorch/.github/actions/setup-linux@release/2.12
+        with:
+          python-version: ${{ inputs.python-version }}
+          compiler: ${{ inputs.compiler }}
+          cuda-version: ${{ inputs.cuda-version }}
+          github-token: ${{ secrets.GITHUB_TOKEN }}
 
       - name: Login to ECR
         if: inputs.build-environment != 'linux-s390x-binary-manywheel'
@@ -165,13 +176,6 @@ jobs:
         with:
           aws-role-to-assume: ${{ inputs.aws-role-to-assume }}
 
-      - name: Get workflow job id
-        id: get-job-id
-        uses: ./.github/actions/get-workflow-job-id
-        if: always()
-        with:
-          github-token: ${{ secrets.GITHUB_TOKEN }}
-
       - name: Check if can use old whl build
         id: use-old-whl
         uses: ./.github/actions/reuse-old-whl
@@ -180,12 +184,12 @@ jobs:
           build-environment: ${{ inputs.build-environment }}
           run-id: ${{ github.run_id }}
           github-token: ${{ secrets.GITHUB_TOKEN }}
-          job-id: ${{ steps.get-job-id.outputs.job-id }}
-          job-name: ${{ steps.get-job-id.outputs.job-name }}
+          job-id: ${{ steps.setup-linux.outputs.job-id }}
+          job-name: ${{ steps.setup-linux.outputs.job-name }}
 
       - name: Calculate docker image
         id: calculate-docker-image
-        uses: pytorch/test-infra/.github/actions/calculate-docker-image@main
+        uses: pytorch/test-infra/.github/actions/calculate-docker-image@release/2.12
         if: inputs.build-environment != 'linux-s390x-binary-manywheel'
         with:
           docker-image-name: ${{ inputs.docker-image-name }}
@@ -201,15 +205,11 @@ jobs:
           echo "docker pull ghcr.io/pytorch/ci-image:${tag/:/-}"
 
       - name: Pull docker image
-        uses: pytorch/test-infra/.github/actions/pull-docker-image@main
+        uses: pytorch/test-infra/.github/actions/pull-docker-image@release/2.12
         if: inputs.build-environment != 'linux-s390x-binary-manywheel' && steps.use-old-whl.outputs.reuse != 'true'
         with:
           docker-image: ${{ steps.calculate-docker-image.outputs.docker-image }}
 
-      - name: Parse ref
-        id: parse-ref
-        run: .github/scripts/parse_ref.py
-
       # Apply the filter logic to the build step too if the test-config label is already there
       - name: Select all requested test configurations (if the test matrix is available)
         id: filter
@@ -218,7 +218,7 @@ jobs:
           github-token: ${{ secrets.GITHUB_TOKEN }}
           test-matrix: ${{ inputs.test-matrix }}
           selected-test-configs: ${{ inputs.selected-test-configs }}
-          job-name: ${{ steps.get-job-id.outputs.job-name }}
+          job-name: ${{ steps.setup-linux.outputs.job-name }}
 
       - name: Start monitoring script
         id: monitor-script
@@ -226,8 +226,8 @@ jobs:
         shell: bash
         continue-on-error: true
         env:
-          JOB_ID: ${{ steps.get-job-id.outputs.job-id }}
-          JOB_NAME: ${{ steps.get-job-id.outputs.job-name }}
+          JOB_ID: ${{ steps.setup-linux.outputs.job-id }}
+          JOB_NAME: ${{ steps.setup-linux.outputs.job-name }}
           WORKFLOW_NAME: ${{ github.workflow }}
           WORKFLOW_RUN_ID: ${{github.run_id}}
           MONITOR_LOG_INTERVAL: ${{ inputs.monitor-log-interval }}
@@ -241,21 +241,12 @@ jobs:
           > "../../usage_logs/usage_log_build_${JOB_ID}.txt" 2>&1 &
           echo "monitor-script-pid=${!}" >> "${GITHUB_OUTPUT}"
 
-      - name: Download pytest cache
-        uses: ./.github/actions/pytest-cache-download
-        continue-on-error: true
-        if: inputs.build-environment != 'linux-s390x-binary-manywheel' && steps.use-old-whl.outputs.reuse != 'true'
-        with:
-          cache_dir: .pytest_cache
-          job_identifier: ${{ github.workflow }}_${{ inputs.build-environment }}
-          s3_bucket: ${{ inputs.s3-bucket }}
-
       - name: Build
         if: (steps.filter.outputs.is-test-matrix-empty == 'False' || inputs.test-matrix == '') && steps.use-old-whl.outputs.reuse != 'true'
         id: build
         env:
           BUILD_ENVIRONMENT: ${{ inputs.build-environment }}
-          BRANCH: ${{ steps.parse-ref.outputs.branch }}
+          BRANCH: ${{ steps.setup-linux.outputs.branch }}
           PR_NUMBER: ${{ github.event.pull_request.number }}
           SHA1: ${{ github.event.pull_request.head.sha || github.sha }}
           # Do not set SCCACHE_S3_KEY_PREFIX to share the cache between all build jobs
@@ -267,9 +258,8 @@ jobs:
           DOCKER_IMAGE: ${{ steps.calculate-docker-image.outputs.docker-image }}
           DOCKER_IMAGE_S390X: ${{ inputs.docker-image-name }}
           XLA_CUDA: ${{ contains(inputs.build-environment, 'xla') && '0' || '' }}
-          OUR_GITHUB_JOB_ID: ${{ steps.get-job-id.outputs.job-id }}
+          OUR_GITHUB_JOB_ID: ${{ steps.setup-linux.outputs.job-id }}
           HUGGING_FACE_HUB_TOKEN: ${{ secrets.HUGGING_FACE_HUB_TOKEN }}
-          SCRIBE_GRAPHQL_ACCESS_TOKEN: ${{ secrets.SCRIBE_GRAPHQL_ACCESS_TOKEN }}
           BUILD_ADDITIONAL_PACKAGES: ${{ inputs.build-additional-packages }}
           RUNNER: ${{ inputs.runner }}
         run: |
@@ -341,7 +331,6 @@ jobs:
             -e PR_LABELS \
             -e OUR_GITHUB_JOB_ID \
             -e HUGGING_FACE_HUB_TOKEN \
-            -e SCRIBE_GRAPHQL_ACCESS_TOKEN \
             -e BUILD_ADDITIONAL_PACKAGES \
             -e RUNNER \
             --memory="${TOTAL_AVAILABLE_MEMORY_IN_GB%.*}g" \
@@ -425,7 +414,7 @@ jobs:
           if-no-files-found: error
           path: artifacts.zip
 
-      - name: copy logs
+      - name: Copy logs
         shell: bash
         if: ${{ always() && steps.build.outcome != 'skipped' && !inputs.disable-monitor && inputs.build-environment != 'linux-s390x-binary-manywheel'}}
         continue-on-error: true
@@ -436,7 +425,7 @@ jobs:
 
       - name: Upload raw usage log to s3
         if: ${{ always() && steps.build.outcome != 'skipped' && !inputs.disable-monitor && inputs.build-environment != 'linux-s390x-binary-manywheel'}}
-        uses: seemethere/upload-artifact-s3@v5
+        uses: seemethere/upload-artifact-s3@baba72d0712b404f646cebe0730933554ebce96a # v5
         with:
           s3-prefix: |
             ${{ github.repository }}/${{ github.run_id }}/${{ github.run_attempt }}/artifact
@@ -456,15 +445,15 @@ jobs:
         continue-on-error: true
         uses: ./.github/actions/upload-utilization-stats
         with:
-          job_id: ${{ steps.get-job-id.outputs.job-id }}
-          job_name: ${{ steps.get-job-id.outputs.job-name }}
+          job_id: ${{ steps.setup-linux.outputs.job-id }}
+          job_name: ${{ steps.setup-linux.outputs.job-name }}
           workflow_name: ${{ github.workflow }}
           workflow_run_id: ${{github.run_id}}
           workflow_attempt: ${{github.run_attempt}}
-          artifact_prefix: usage_log_build_${{ steps.get-job-id.outputs.job-id }}
+          artifact_prefix: usage_log_build_${{ steps.setup-linux.outputs.job-id }}
 
       - name: Teardown Linux
-        uses: pytorch/test-infra/.github/actions/teardown-linux@main
+        uses: pytorch/test-infra/.github/actions/teardown-linux@release/2.12
         if: always() && inputs.build-environment != 'linux-s390x-binary-manywheel'
 
       - name: Cleanup docker
@@ -474,3 +463,136 @@ jobs:
           # on s390x stop the container for clean worker stop
           docker stop -a || true
           docker kill -a || true
+
+  build-osdc:
+    permissions:
+      id-token: write
+      contents: read
+      actions: read
+    # Don't run on forked repos
+    if: github.repository_owner == 'pytorch' && inputs.use-arc
+    runs-on: ${{ inputs.runner_prefix }}${{ startsWith(inputs.runner, 'l-') && inputs.runner || contains(inputs.runner, 'arm64') && 'l-arm64g4-16-62' || 'l-x86iavx512-8-64' }}
+    container:
+      image: ghcr.io/pytorch/${{ inputs.docker-image-name }}
+    timeout-minutes: 480
+    outputs:
+      docker-image: ghcr.io/pytorch/${{ inputs.docker-image-name }}
+      test-matrix: ${{ steps.map-runners.outputs.test-matrix }}
+      build-environment: ${{ inputs.build-environment }}
+    steps:
+      - name: Setup Linux
+        id: setup-linux
+        uses: pytorch/pytorch/.github/actions/setup-linux@release/2.12
+        with:
+          use-arc: true
+          python-version: ${{ inputs.python-version }}
+          compiler: ${{ inputs.compiler }}
+          cuda-version: ${{ inputs.cuda-version }}
+          github-token: ${{ secrets.GITHUB_TOKEN }}
+
+      - name: Check if can use old whl build
+        id: use-old-whl
+        uses: pytorch/pytorch/.github/actions/reuse-old-whl@release/2.12
+        if: ${{ inputs.allow-reuse-old-whl }}
+        with:
+          build-environment: ${{ inputs.build-environment }}
+          run-id: ${{ github.run_id }}
+          github-token: ${{ secrets.GITHUB_TOKEN }}
+          job-id: ${{ steps.setup-linux.outputs.job-id }}
+          job-name: ${{ steps.setup-linux.outputs.job-name }}
+
+      # Apply the filter logic to the build step too if the test-config label is already there
+      - name: Select all requested test configurations (if the test matrix is available)
+        id: filter
+        uses: pytorch/pytorch/.github/actions/filter-test-configs@release/2.12
+        with:
+          github-token: ${{ secrets.GITHUB_TOKEN }}
+          test-matrix: ${{ inputs.test-matrix }}
+          selected-test-configs: ${{ inputs.selected-test-configs }}
+          job-name: ${{ steps.setup-linux.outputs.job-name }}
+
+      - name: Map EC2 runners to ARC runners
+        id: map-runners
+        env:
+          FILTERED_TEST_MATRIX: ${{ steps.filter.outputs.test-matrix }}
+          RUNNER_PREFIX: ${{ inputs.runner_prefix }}
+        shell: bash
+        run: |
+          python3 .github/scripts/map_ec2_to_arc.py --prefix "${RUNNER_PREFIX}" "${FILTERED_TEST_MATRIX}"
+
+      - name: Configure AWS credentials
+        id: aws-creds
+        continue-on-error: true
+        uses: aws-actions/configure-aws-credentials@8df5847569e6427dd6c4fb1cf565c83acfa8afa7
+        with:
+          role-to-assume: arn:aws:iam::308535385114:role/arc
+          aws-region: us-east-1
+          # The max duration enforced by the server side
+          role-duration-seconds: 18000
+
+      - name: Build
+        if: (steps.filter.outputs.is-test-matrix-empty == 'False' || inputs.test-matrix == '') && steps.use-old-whl.outputs.reuse != 'true'
+        id: build
+        env:
+          BUILD_ENVIRONMENT: ${{ inputs.build-environment }}
+          BRANCH: ${{ steps.setup-linux.outputs.branch }}
+          PR_NUMBER: ${{ github.event.pull_request.number }}
+          SHA1: ${{ github.event.pull_request.head.sha || github.sha }}
+          SCCACHE_BUCKET: ossci-compiler-cache-circleci-v2
+          SCCACHE_REGION: us-east-1
+          SCCACHE_S3_NO_CREDENTIALS: ${{ steps.aws-creds.outcome != 'success' && 'true' || 'false' }}
+          XLA_CLANG_CACHE_S3_BUCKET_NAME: ossci-compiler-clang-cache-circleci-xla
+          PR_LABELS: ${{ toJson(github.event.pull_request.labels.*.name) }}
+          TORCH_CUDA_ARCH_LIST: ${{ inputs.cuda-arch-list }}
+          XLA_CUDA: ${{ contains(inputs.build-environment, 'xla') && '0' || '' }}
+          OUR_GITHUB_JOB_ID: ${{ steps.setup-linux.outputs.job-id }}
+          HUGGING_FACE_HUB_TOKEN: ${{ secrets.HUGGING_FACE_HUB_TOKEN }}
+          BUILD_ADDITIONAL_PACKAGES: ${{ inputs.build-additional-packages }}
+          RUNNER: ${{ inputs.runner }}
+          SKIP_SCCACHE_INITIALIZATION: 1
+        shell: bash
+        run: |
+          START_TIME=$(date +%s)
+          .ci/pytorch/build.sh
+          END_TIME=$(date +%s)
+          echo "build_time=$((END_TIME - START_TIME))" >> "$GITHUB_OUTPUT"
+
+      - name: Build external packages
+        id: build-external-packages
+        if: inputs.build-external-packages != '' &&  steps.build.outcome != 'skipped'
+        uses: pytorch/pytorch/.github/actions/build-external-packages@release/2.12
+        with:
+          build-targets: ${{ inputs.build-external-packages }}
+          docker-image: ghcr.io/pytorch/${{ inputs.docker-image-name }}
+          cuda-arch-list: ${{ inputs.cuda-arch-list }}
+          output-dir: external
+
+      - name: Move external packages to dist
+        if: steps.build-external-packages.outputs.output_dir != '' && steps.build-external-packages.outcome != 'skipped'
+        shell: bash
+        run: |
+          src="${{ steps.build-external-packages.outputs.output_dir }}"
+          if [ -d "$src" ]; then
+            mkdir -p "dist/$(dirname "$src")"
+            mv "$src" "dist/$(dirname "$src")/"
+          fi
+
+      - name: Archive artifacts into zip
+        if: inputs.build-generates-artifacts && steps.build.outcome != 'skipped' && steps.use-old-whl.outputs.reuse != 'true'
+        run: |
+          zip -1 -r artifacts.zip dist/ build/custom_test_artifacts build/lib build/bin .additional_ci_files
+
+      - name: Store build artifacts
+        if: inputs.build-generates-artifacts && (steps.build.outcome != 'skipped' || steps.use-old-whl.outputs.reuse == 'true')
+        uses: pytorch/pytorch/.github/actions/upload-build-artifacts@release/2.12
+        with:
+          name: ${{ inputs.build-environment }}
+          s3-bucket: ${{ inputs.s3-bucket }}
+          use-gha: ${{ steps.aws-creds.outcome != 'success' }}
+
+      - name: Upload sccache stats
+        if: steps.build.outcome != 'skipped' && steps.aws-creds.outcome == 'success'
+        uses: pytorch/pytorch/.github/actions/upload-sccache-stats@release/2.12
+        with:
+          github-token: ${{ secrets.GITHUB_TOKEN }}
+          build-time: ${{ steps.build.outputs.build_time }}
diff --git a/.github/workflows/_linux-test-stable-fa3.yml b/.github/workflows/_linux-test-stable-fa3.yml
index f2e16712ff447..d45a38c424dcd 100644
--- a/.github/workflows/_linux-test-stable-fa3.yml
+++ b/.github/workflows/_linux-test-stable-fa3.yml
@@ -54,15 +54,13 @@ jobs:
     # Don't run on forked repos
     if: github.repository_owner == 'pytorch'
     runs-on: linux.aws.h100
-    timeout-minutes: ${{ inputs.timeout-minutes || 30 }}
+    timeout-minutes: ${{ inputs.timeout-minutes || 60 }}
     permissions:
       id-token: write
       contents: read
     steps:
-      - name: Checkout PyTorch
-        uses: pytorch/pytorch/.github/actions/checkout-pytorch@main
-        with:
-          no-sudo: true
+      - name: Setup Linux
+        uses: pytorch/pytorch/.github/actions/setup-linux@release/2.12
 
       - name: Checkout flash-attention as a secondary repository
         uses: actions/checkout@v4
@@ -70,15 +68,12 @@ jobs:
           repository: Dao-AILab/flash-attention
           path: flash-attention
 
-      - name: Setup Linux
-        uses: ./.github/actions/setup-linux
-
       - name: Login to ECR
         uses: ./.github/actions/ecr-login
 
       - name: Calculate docker image
         id: calculate-docker-image
-        uses: pytorch/test-infra/.github/actions/calculate-docker-image@main
+        uses: pytorch/test-infra/.github/actions/calculate-docker-image@release/2.12
         with:
           docker-image-name: ${{ inputs.docker-image }}
 
@@ -92,7 +87,7 @@ jobs:
           echo "docker pull ghcr.io/pytorch/ci-image:${tag/:/-}"
 
       - name: Pull docker image
-        uses: pytorch/test-infra/.github/actions/pull-docker-image@main
+        uses: pytorch/test-infra/.github/actions/pull-docker-image@release/2.12
         with:
           docker-image: ${{ steps.calculate-docker-image.outputs.docker-image }}
 
@@ -254,5 +249,5 @@ jobs:
           workflow_attempt: ${{github.run_attempt}}
 
       - name: Teardown Linux
-        uses: pytorch/test-infra/.github/actions/teardown-linux@main
+        uses: pytorch/test-infra/.github/actions/teardown-linux@release/2.12
         if: always() && steps.check_container_runner.outputs.IN_CONTAINER_RUNNER == 'false'
diff --git a/.github/workflows/_linux-test.yml b/.github/workflows/_linux-test.yml
index 979b7a8e55a08..46a591ba0ba6e 100644
--- a/.github/workflows/_linux-test.yml
+++ b/.github/workflows/_linux-test.yml
@@ -72,6 +72,38 @@ on:
         required: false
         type: number
         default: 1
+      use-arc:
+        required: false
+        type: boolean
+        default: false
+        description: If true, use ARC (OSDC) runner path instead of EC2.
+      python-version:
+        required: false
+        type: string
+        default: ""
+        description: Python version to use for the OSDC test.
+      compiler:
+        required: false
+        type: string
+        default: ""
+        description: Compiler to use for the OSDC test.
+      cuda-version:
+        required: false
+        type: string
+        default: ""
+        description: CUDA version to use for the OSDC test.
+      export-profiler-trace:
+        description: |
+          If set to "1", export Chrome profiler traces from performance benchmarks.
+        required: false
+        type: string
+        default: ""
+      enable-torch-trace:
+        description: |
+          If set to "1", enable TORCH_TRACE structured logging and collect tlparse output.
+        required: false
+        type: string
+        default: ""
     secrets:
       HUGGING_FACE_HUB_TOKEN:
         required: false
@@ -92,7 +124,7 @@ env:
 jobs:
   test:
     # Don't run on forked repos or empty test matrix
-    if: github.repository_owner == 'pytorch' && toJSON(fromJSON(inputs.test-matrix).include) != '[]'
+    if: github.repository_owner == 'pytorch' && toJSON(fromJSON(inputs.test-matrix).include) != '[]' && !inputs.use-arc
     strategy:
       matrix: ${{ fromJSON(inputs.test-matrix) }}
       fail-fast: false
@@ -104,7 +136,7 @@ jobs:
       contents: read
     steps:
       - name: Setup SSH (Click me for login details)
-        uses: pytorch/test-infra/.github/actions/setup-ssh@main
+        uses: pytorch/test-infra/.github/actions/setup-ssh@release/2.12
         if: ${{ !contains(matrix.runner, 'b200') && inputs.build-environment != 'linux-s390x-binary-manywheel' }}
         with:
           github-secret: ${{ secrets.GITHUB_TOKEN }}
@@ -112,21 +144,15 @@ jobs:
             All testing is done inside the container, to start an interactive session run:
               docker exec -it $(docker container ps --format '{{.ID}}') bash
 
-      - name: Checkout PyTorch
-        uses: pytorch/pytorch/.github/actions/checkout-pytorch@main
-        with:
-          no-sudo: true
-
-      - name: Setup Python
-        if: contains(matrix.runner, 'b200')
-        uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
-        with:
-          python-version: '3.12'
-          cache: pip
-
       - name: Setup Linux
-        uses: ./.github/actions/setup-linux
-        if: inputs.build-environment != 'linux-s390x-binary-manywheel' && !contains(matrix.runner, 'b200')
+        id: setup-linux
+        uses: pytorch/pytorch/.github/actions/setup-linux@release/2.12
+        with:
+          python-version: ${{ inputs.python-version }}
+          compiler: ${{ inputs.compiler }}
+          cuda-version: ${{ inputs.cuda-version }}
+          submodules: 'false'
+          github-token: ${{ secrets.GITHUB_TOKEN }}
 
       - name: Check TPU Availability
         id: check-tpu
@@ -147,7 +173,7 @@ jobs:
 
       - name: Calculate docker image
         id: calculate-docker-image
-        uses: pytorch/test-infra/.github/actions/calculate-docker-image@main
+        uses: pytorch/test-infra/.github/actions/calculate-docker-image@release/2.12
         if: inputs.build-environment != 'linux-s390x-binary-manywheel'
         with:
           docker-image-name: ${{ inputs.docker-image }}
@@ -163,7 +189,7 @@ jobs:
           echo "docker pull ghcr.io/pytorch/ci-image:${tag/:/-}"
 
       - name: Pull docker image
-        uses: pytorch/test-infra/.github/actions/pull-docker-image@main
+        uses: pytorch/test-infra/.github/actions/pull-docker-image@release/2.12
         if: inputs.build-environment != 'linux-s390x-binary-manywheel'
         with:
           docker-image: ${{ steps.calculate-docker-image.outputs.docker-image }}
@@ -175,9 +201,9 @@ jobs:
 
       - name: Install nvidia driver, nvidia-docker runtime, set GPU_FLAG
         id: install-nvidia-driver
-        uses: pytorch/test-infra/.github/actions/setup-nvidia@main
+        uses: pytorch/test-infra/.github/actions/setup-nvidia@release/2.12
         with:
-          driver-version: ${{ matrix.config == 'legacy_nvidia_driver' && '525.105.17' || '580.82.07' }}
+          driver-version: '580.82.07'
         if: ${{ !contains(matrix.runner, 'b200') }}
 
       - name: Setup GPU_FLAG for docker run
@@ -190,21 +216,14 @@ jobs:
         run: echo "SCCACHE_SERVER_PORT_DOCKER_FLAG=-e SCCACHE_SERVER_PORT=$((RUNNER_UID + 4226))" >> "${GITHUB_ENV}"
         if: ${{ steps.check_container_runner.outputs.IN_CONTAINER_RUNNER == 'true' && !contains(matrix.runner, 'b200') }}
 
-      - name: Get workflow job id
-        id: get-job-id
-        uses: ./.github/actions/get-workflow-job-id
-        if: always()
-        with:
-          github-token: ${{ secrets.GITHUB_TOKEN }}
-
       - name: Start monitoring script
         id: monitor-script
         if: ${{ !inputs.disable-monitor }}
         shell: bash
         continue-on-error: true
         env:
-          JOB_ID: ${{ steps.get-job-id.outputs.job-id }}
-          JOB_NAME: ${{ steps.get-job-id.outputs.job-name }}
+          JOB_ID: ${{ steps.setup-linux.outputs.job-id }}
+          JOB_NAME: ${{ steps.setup-linux.outputs.job-name }}
           WORKFLOW_NAME: ${{ github.workflow }}
           WORKFLOW_RUN_ID: ${{github.run_id}}
           MONITOR_LOG_INTERVAL: ${{ inputs.monitor-log-interval }}
@@ -260,13 +279,20 @@ jobs:
             echo "Moved cudart.lib to win-torch-wheel-extracted/lib/x64/"
           fi
 
-          # Verify CUDA libraries are present
+          # Setup CUDA runtime DLL (needed for MinGW import lib generation on CUDA 13.0+)
+          mkdir -p win-torch-wheel-extracted/bin/x64
+          for dll in win-torch-wheel/cudart64_*.dll; do
+            if [ -f "$dll" ]; then
+              mv "$dll" win-torch-wheel-extracted/bin/x64/
+              echo "Moved $(basename $dll) to win-torch-wheel-extracted/bin/x64/"
+            fi
+          done
+
+          # Verify CUDA libraries and DLLs are present
           echo "CUDA libraries:"
           ls -la win-torch-wheel-extracted/lib/x64/ || echo "No CUDA libraries found"
-
-      - name: Parse ref
-        id: parse-ref
-        run: .github/scripts/parse_ref.py
+          echo "CUDA DLLs:"
+          ls -la win-torch-wheel-extracted/bin/x64/ || echo "No CUDA DLLs found"
 
       - name: Check for keep-going label and re-enabled test issues
         # This uses the filter-test-configs action because it conveniently
@@ -277,7 +303,7 @@ jobs:
         with:
           github-token: ${{ secrets.GITHUB_TOKEN }}
           test-matrix: ${{ inputs.test-matrix }}
-          job-name: ${{ steps.get-job-id.outputs.job-name }}
+          job-name: ${{ steps.setup-linux.outputs.job-name }}
 
       - name: Set Test step time
         id: test-timeout
@@ -293,6 +319,17 @@ jobs:
           env | grep '^GITHUB' >> "/tmp/github_env_${GITHUB_RUN_ID}"
           env | grep '^CI' >> "/tmp/github_env_${GITHUB_RUN_ID}"
 
+      # Allow the test results to be uploaded S3 right after they finish, which are
+      # then used in the test insight dashboard
+      - name: Authenticate with AWS
+        if: ${{ contains(matrix.runner, 'b200') || steps.check-tpu.outputs.has_tpu == 'true' }}
+        uses: aws-actions/configure-aws-credentials@ececac1a45f3b08a01d2dd070d28d111c5fe6722 # v4.1.0
+        with:
+          role-to-assume: arn:aws:iam::308535385114:role/gha_workflow_upload-benchmark-results
+          # The max duration enforced by the server side
+          role-duration-seconds: 18000
+          aws-region: us-east-1
+
       - name: Test
         id: test
         timeout-minutes: ${{ fromJson(steps.test-timeout.outputs.timeout) }}
@@ -305,9 +342,9 @@ jobs:
           GITHUB_RUN_ID: ${{ github.run_id }}
           GITHUB_RUN_NUMBER: ${{ github.run_number }}
           GITHUB_RUN_ATTEMPT: ${{ github.run_attempt }}
-          JOB_ID: ${{ steps.get-job-id.outputs.job-id }}
-          JOB_NAME: ${{ steps.get-job-id.outputs.job-name }}
-          BRANCH: ${{ steps.parse-ref.outputs.branch }}
+          JOB_ID: ${{ steps.setup-linux.outputs.job-id }}
+          JOB_NAME: ${{ steps.setup-linux.outputs.job-name }}
+          BRANCH: ${{ steps.setup-linux.outputs.branch }}
           SHA1: ${{ github.event.pull_request.head.sha || github.sha }}
           BASE_SHA: ${{ github.event.pull_request.base.sha || github.sha }}
           TEST_CONFIG: ${{ matrix.config }}
@@ -334,17 +371,20 @@ jobs:
           PYTORCH_TEST_RERUN_DISABLED_TESTS: ${{ matrix.rerun_disabled_tests && '1' || '0' }}
           TESTS_TO_INCLUDE: ${{ inputs.tests-to-include }}
           DASHBOARD_TAG: ${{ inputs.dashboard-tag }}
+          EXPORT_PROFILER_TRACE: ${{ inputs.export-profiler-trace }}
+          ENABLE_TORCH_TRACE: ${{ inputs.enable-torch-trace }}
           VLLM_TEST_HUGGING_FACE_TOKEN: ${{ secrets.VLLM_TEST_HUGGING_FACE_TOKEN }}
           HF_CACHE: /mnt/hf_cache
           # Use offline mode by default, only enable online mode to refresh the models
-          # from HF when the PR has a special ci-refresh-hf-cache label. This label is
+          # from HF when the PR has a special ci-refresh-hf-cache label or when the job
+          # is a nightly scheduled run on main. The ci-refresh-hf-cache label is
           # automatically added to the vLLM pinned commit hash update, which effectively
           # refreshes the cache daily
-          TRANSFORMERS_OFFLINE: ${{ contains(steps.keep-going.outputs.labels, 'ci-refresh-hf-cache') && '0' || '1' }}
-          HF_DATASETS_OFFLINE: ${{ contains(steps.keep-going.outputs.labels, 'ci-refresh-hf-cache') && '0' || '1' }}
+          TRANSFORMERS_OFFLINE: ${{ (github.event_name == 'schedule' || contains(steps.keep-going.outputs.labels, 'ci-refresh-hf-cache')) && '0' || '1' }}
+          HF_DATASETS_OFFLINE: ${{ (github.event_name == 'schedule' || contains(steps.keep-going.outputs.labels, 'ci-refresh-hf-cache')) && '0' || '1' }}
           HUGGING_FACE_HUB_TOKEN: ${{ secrets.HUGGING_FACE_HUB_TOKEN }}
           SCRIBE_GRAPHQL_ACCESS_TOKEN: ${{ secrets.SCRIBE_GRAPHQL_ACCESS_TOKEN }}
-          ARTIFACTS_FILE_SUFFIX: ${{ github.job }}-${{ matrix.config }}-${{ matrix.shard }}-${{ matrix.num_shards }}-${{ matrix.runner }}_${{ steps.get-job-id.outputs.job-id }}
+          ARTIFACTS_FILE_SUFFIX: ${{ github.job }}-${{ matrix.config }}-${{ matrix.shard }}-${{ matrix.num_shards }}-${{ matrix.runner }}_${{ steps.setup-linux.outputs.job-id }}
           TORCH_TPU: ${{ steps.check-tpu.outputs.has_tpu == 'true' && '1' || '' }}
           TORCH_TPU_TEXT_FILE: /var/lib/jenkins/workspace/.github/ci_commit_pins/torch_tpu.txt
         run: |
@@ -352,8 +392,6 @@ jobs:
 
           if [[ $TEST_CONFIG == 'multigpu' ]]; then
             TEST_COMMAND=.ci/pytorch/multigpu-test.sh
-          elif [[ $BUILD_ENVIRONMENT == *onnx* ]]; then
-            TEST_COMMAND=.ci/onnx/test.sh
           else
             TEST_COMMAND=.ci/pytorch/test.sh
           fi
@@ -391,22 +429,6 @@ jobs:
             export HF_DATASETS_OFFLINE=0
           fi
 
-          # For TIMM jobs, create a pin-specific cache directory and make it
-          # world-writable so the jenkins user inside docker can write to it.
-          # Only enable online mode if the cache hasn't been fully populated.
-          if [[ "${TEST_CONFIG}" == *timm* ]]; then
-            TIMM_PIN="$(< .ci/docker/ci_commit_pins/timm.txt)"
-            TIMM_CACHE_DIR="${HF_CACHE}/timm_${TIMM_PIN}"
-            if [[ ! -d "${TIMM_CACHE_DIR}" ]]; then
-              mkdir -p "${TIMM_CACHE_DIR}"
-              chmod -R a+rwX "${TIMM_CACHE_DIR}"
-            fi
-            if [[ ! -f "${TIMM_CACHE_DIR}/.timm_cache_complete" ]]; then
-              export TRANSFORMERS_OFFLINE=0
-              export HF_DATASETS_OFFLINE=0
-            fi
-          fi
-
           # detached container should get cleaned up by teardown_ec2_linux
           # TODO: Stop building test binaries as part of the build phase
           # Used for GPU_FLAG, SHM_OPTS, JENKINS_USER and DOCKER_SHELL_CMD since that doesn't play nice
@@ -429,6 +451,9 @@ jobs:
             -e BRANCH \
             -e SHA1 \
             -e AWS_DEFAULT_REGION \
+            -e AWS_ACCESS_KEY_ID \
+            -e AWS_SECRET_ACCESS_KEY \
+            -e AWS_SESSION_TOKEN \
             -e IN_WHEEL_TEST \
             -e SHARD_NUMBER \
             -e TEST_CONFIG \
@@ -457,6 +482,8 @@ jobs:
             -e HF_DATASETS_OFFLINE \
             -e SCRIBE_GRAPHQL_ACCESS_TOKEN \
             -e DASHBOARD_TAG \
+            -e EXPORT_PROFILER_TRACE \
+            -e ENABLE_TORCH_TRACE \
             -e ARTIFACTS_FILE_SUFFIX \
             -e TORCH_TPU \
             -e TORCH_TPU_TEXT_FILE \
@@ -512,7 +539,7 @@ jobs:
           aws-region: us-east-1
 
       - name: Upload the benchmark results
-        uses: pytorch/test-infra/.github/actions/upload-benchmark-results@main
+        uses: pytorch/test-infra/.github/actions/upload-benchmark-results@release/2.12
         if: inputs.build-environment != 'linux-s390x-binary-manywheel' && steps.check-tpu.outputs.has_tpu != 'true'
         with:
           benchmark-results-dir: test/test-reports
@@ -539,7 +566,7 @@ jobs:
         uses: ./.github/actions/upload-test-artifacts
         if: always() && steps.test.conclusion && steps.test.conclusion != 'skipped'
         with:
-          file-suffix: ${{ github.job }}-${{ matrix.config }}-${{ matrix.shard }}-${{ matrix.num_shards }}-${{ matrix.runner }}_${{ steps.get-job-id.outputs.job-id }}
+          file-suffix: ${{ github.job }}-${{ matrix.config }}-${{ matrix.shard }}-${{ matrix.num_shards }}-${{ matrix.runner }}_${{ steps.setup-linux.outputs.job-id }}
           use-gha: ${{ inputs.use-gha }}
           s3-bucket: ${{ inputs.s3-bucket }}
 
@@ -563,14 +590,14 @@ jobs:
         continue-on-error: true
         uses: ./.github/actions/upload-utilization-stats
         with:
-          job_id: ${{ steps.get-job-id.outputs.job-id }}
-          job_name: ${{ steps.get-job-id.outputs.job-name }}
+          job_id: ${{ steps.setup-linux.outputs.job-id }}
+          job_name: ${{ steps.setup-linux.outputs.job-name }}
           workflow_name: ${{ github.workflow }}
           workflow_run_id: ${{github.run_id}}
           workflow_attempt: ${{github.run_attempt}}
 
       - name: Teardown Linux
-        uses: pytorch/test-infra/.github/actions/teardown-linux@main
+        uses: pytorch/test-infra/.github/actions/teardown-linux@release/2.12
         if: always() && steps.check_container_runner.outputs.IN_CONTAINER_RUNNER == 'false'
 
       - name: Cleanup docker
@@ -580,3 +607,237 @@ jobs:
           # on s390x stop the container for clean worker stop
           docker stop -a || true
           docker kill -a || true
+
+  test-osdc:
+    # Don't run on forked repos or empty test matrix
+    if: github.repository_owner == 'pytorch' && toJSON(fromJSON(inputs.test-matrix).include) != '[]' && inputs.use-arc
+    strategy:
+      matrix: ${{ fromJSON(inputs.test-matrix) }}
+      fail-fast: false
+    environment: ${{ github.ref == 'refs/heads/main' && 'scribe-protected' || startsWith(github.ref, 'refs/heads/release/') && 'scribe-protected' || contains(github.event.pull_request.labels.*.name, 'ci-scribe') && 'scribe-pr' || '' }}
+    runs-on: ${{ matrix.runner }}
+    container:
+      image: ${{ inputs.docker-image }}
+      options: "--gpus all"
+    timeout-minutes: ${{ matrix.mem_leak_check == 'mem_leak_check' && 600 || inputs.timeout-minutes }}
+    permissions:
+      id-token: write
+      contents: read
+      actions: read
+    steps:
+      - name: Setup Linux
+        id: setup-linux
+        uses: pytorch/pytorch/.github/actions/setup-linux@release/2.12
+        with:
+          use-arc: true
+          python-version: ${{ inputs.python-version }}
+          compiler: ${{ inputs.compiler }}
+          cuda-version: ${{ inputs.cuda-version }}
+          submodules: 'false'
+          github-token: ${{ secrets.GITHUB_TOKEN }}
+
+      - name: Configure AWS credentials
+        id: aws-creds
+        continue-on-error: true
+        uses: aws-actions/configure-aws-credentials@8df5847569e6427dd6c4fb1cf565c83acfa8afa7
+        with:
+          role-to-assume: arn:aws:iam::308535385114:role/arc
+          aws-region: us-east-1
+          # The max duration enforced by the server side
+          role-duration-seconds: 18000
+
+      - name: Download build artifacts
+        uses: pytorch/pytorch/.github/actions/download-build-artifacts@release/2.12
+        with:
+          name: ${{ inputs.build-environment }}
+          s3-bucket: ${{ inputs.s3-bucket }}
+          use-gha: ${{ steps.aws-creds.outcome != 'success' || inputs.use-gha }}
+
+      - name: Download TD artifacts
+        continue-on-error: true
+        uses: pytorch/pytorch/.github/actions/download-td-artifacts@release/2.12
+
+      - name: Download Windows torch wheel for cross-compilation
+        if: matrix.win_torch_wheel_artifact != ''
+        uses: seemethere/download-artifact-s3@1da556a7aa0a088e3153970611f6c432d58e80e6 # v4.2.0
+        with:
+          name: ${{ matrix.win_torch_wheel_artifact }}
+          path: win-torch-wheel
+
+      - name: Extract Windows wheel and setup CUDA libraries
+        if: matrix.win_torch_wheel_artifact != ''
+        shell: bash
+        run: |
+          set -x
+
+          # Find the wheel file
+          WHEEL_FILE=$(find win-torch-wheel -name "*.whl" -type f | head -n 1)
+          if [ -z "$WHEEL_FILE" ]; then
+            echo "Error: No wheel file found in win-torch-wheel directory"
+            exit 1
+          fi
+          echo "Found wheel file: $WHEEL_FILE"
+
+          # Unzip the wheel file
+          unzip -q "$WHEEL_FILE" -d win-torch-wheel-extracted
+          echo "Extracted wheel contents"
+
+          # Setup CUDA libraries (cuda.lib and cudart.lib) directory
+          mkdir -p win-torch-wheel-extracted/lib/x64
+          if [ -f "win-torch-wheel/cuda.lib" ]; then
+            mv win-torch-wheel/cuda.lib win-torch-wheel-extracted/lib/x64/
+            echo "Moved cuda.lib to win-torch-wheel-extracted/lib/x64/"
+          fi
+          if [ -f "win-torch-wheel/cudart.lib" ]; then
+            mv win-torch-wheel/cudart.lib win-torch-wheel-extracted/lib/x64/
+            echo "Moved cudart.lib to win-torch-wheel-extracted/lib/x64/"
+          fi
+
+          # Setup CUDA runtime DLL (needed for MinGW import lib generation on CUDA 13.0+)
+          mkdir -p win-torch-wheel-extracted/bin/x64
+          for dll in win-torch-wheel/cudart64_*.dll; do
+            if [ -f "$dll" ]; then
+              mv "$dll" win-torch-wheel-extracted/bin/x64/
+              echo "Moved $(basename $dll) to win-torch-wheel-extracted/bin/x64/"
+            fi
+          done
+
+          # Verify CUDA libraries and DLLs are present
+          echo "CUDA libraries:"
+          ls -la win-torch-wheel-extracted/lib/x64/ || echo "No CUDA libraries found"
+          echo "CUDA DLLs:"
+          ls -la win-torch-wheel-extracted/bin/x64/ || echo "No CUDA DLLs found"
+
+      - name: Check for keep-going label and re-enabled test issues
+        # This uses the filter-test-configs action because it conveniently
+        # checks for labels and re-enabled test issues. It does not actually do
+        # any filtering.  All filtering is done in the build step.
+        id: keep-going
+        uses: pytorch/pytorch/.github/actions/filter-test-configs@release/2.12
+        with:
+          github-token: ${{ secrets.GITHUB_TOKEN }}
+          test-matrix: ${{ inputs.test-matrix }}
+          job-name: ${{ steps.setup-linux.outputs.job-name }}
+
+      - name: Set test step time
+        id: test-timeout
+        shell: bash
+        env:
+          JOB_TIMEOUT: ${{ matrix.mem_leak_check == 'mem_leak_check' && 600 || inputs.timeout-minutes }}
+        run: |
+          echo "timeout=$((JOB_TIMEOUT-30))" >> "${GITHUB_OUTPUT}"
+
+      - name: Test
+        id: test
+        timeout-minutes: ${{ fromJson(steps.test-timeout.outputs.timeout) }}
+        env:
+          BUILD_ENVIRONMENT: ${{ inputs.build-environment }}
+          PR_NUMBER: ${{ github.event.pull_request.number }}
+          GITHUB_REPOSITORY: ${{ github.repository }}
+          GITHUB_WORKFLOW: ${{ github.workflow }}
+          GITHUB_JOB: ${{ github.job }}
+          GITHUB_RUN_ID: ${{ github.run_id }}
+          GITHUB_RUN_NUMBER: ${{ github.run_number }}
+          GITHUB_RUN_ATTEMPT: ${{ github.run_attempt }}
+          JOB_ID: ${{ steps.setup-linux.outputs.job-id }}
+          JOB_NAME: ${{ steps.setup-linux.outputs.job-name }}
+          BRANCH: ${{ steps.setup-linux.outputs.branch }}
+          SHA1: ${{ github.event.pull_request.head.sha || github.sha }}
+          BASE_SHA: ${{ github.event.pull_request.base.sha || github.sha }}
+          TEST_CONFIG: ${{ matrix.config }}
+          SHARD_NUMBER: ${{ matrix.shard }}
+          NUM_TEST_SHARDS: ${{ matrix.num_shards }}
+          EXTRA_FLAGS: ${{ matrix.extra_flags || '' }}
+          OP_BENCHMARK_TESTS: ${{ matrix.op_benchmark_tests }}
+          REENABLED_ISSUES: ${{ steps.keep-going.outputs.reenabled-issues }}
+          CONTINUE_THROUGH_ERROR: ${{ steps.keep-going.outputs.keep-going }}
+          VERBOSE_TEST_LOGS: ${{ steps.keep-going.outputs.ci-verbose-test-logs }}
+          TEST_SHOWLOCALS: ${{ steps.keep-going.outputs.ci-test-showlocals }}
+          NO_TEST_TIMEOUT: ${{ steps.keep-going.outputs.ci-no-test-timeout }}
+          NO_TD: ${{ steps.keep-going.outputs.ci-no-td }}
+          TD_DISTRIBUTED: ${{ steps.keep-going.outputs.ci-td-distributed }}
+          SCCACHE_BUCKET: ossci-compiler-cache-circleci-v2
+          SCCACHE_REGION: us-east-1
+          SCCACHE_S3_NO_CREDENTIALS: ${{ steps.aws-creds.outcome != 'success' && 'true' || 'false' }}
+          XLA_CUDA: ${{ contains(inputs.build-environment, 'xla') && '0' || '' }}
+          XLA_CLANG_CACHE_S3_BUCKET_NAME: ossci-compiler-clang-cache-circleci-xla
+          PYTORCH_TEST_CUDA_MEM_LEAK_CHECK: ${{ matrix.mem_leak_check && '1' || '0' }}
+          PYTORCH_TEST_RERUN_DISABLED_TESTS: ${{ matrix.rerun_disabled_tests && '1' || '0' }}
+          TESTS_TO_INCLUDE: ${{ inputs.tests-to-include }}
+          DASHBOARD_TAG: ${{ inputs.dashboard-tag }}
+          VLLM_TEST_HUGGING_FACE_TOKEN: ${{ secrets.VLLM_TEST_HUGGING_FACE_TOKEN }}
+          HF_CACHE: /mnt/hf_cache
+          TRANSFORMERS_OFFLINE: ${{ (github.event_name == 'schedule' || contains(steps.keep-going.outputs.labels, 'ci-refresh-hf-cache')) && '0' || '1' }}
+          HF_DATASETS_OFFLINE: ${{ (github.event_name == 'schedule' || contains(steps.keep-going.outputs.labels, 'ci-refresh-hf-cache')) && '0' || '1' }}
+          HUGGING_FACE_HUB_TOKEN: ${{ secrets.HUGGING_FACE_HUB_TOKEN }}
+          SCRIBE_GRAPHQL_ACCESS_TOKEN: ${{ secrets.SCRIBE_GRAPHQL_ACCESS_TOKEN }}
+          ARTIFACTS_FILE_SUFFIX: ${{ github.job }}-${{ matrix.config }}-${{ matrix.shard }}-${{ matrix.num_shards }}-${{ matrix.runner }}_${{ steps.setup-linux.outputs.job-id }}
+          TORCH_TPU_TEXT_FILE: /var/lib/jenkins/workspace/.github/ci_commit_pins/torch_tpu.txt
+          USE_ARC: "1"
+        shell: bash
+        run: |
+          set -x
+
+          if [[ "${TEST_CONFIG}" == 'multigpu' ]]; then
+            TEST_COMMAND=.ci/pytorch/multigpu-test.sh
+          else
+            TEST_COMMAND=.ci/pytorch/test.sh
+          fi
+
+          # Just create an empty HF_CACHE dir if it doesn't exist. This dir is not
+          # used for anything besides vLLM jobs
+          if [[ ! -d "${HF_CACHE}" ]]; then
+            export HF_CACHE="${RUNNER_TEMP}/hf_cache"
+            mkdir -p "${HF_CACHE}"
+
+            # When there is no cache directory, e.g. benchmark, the job has no
+            # way but to reach out to HF if needed
+            export TRANSFORMERS_OFFLINE=0
+            export HF_DATASETS_OFFLINE=0
+          fi
+
+          # shellcheck disable=SC2046
+          python3 -m pip install $(echo dist/*.whl)[opt-einsum]
+          ${TEST_COMMAND}
+
+      - name: Configure AWS credentials
+        id: aws-creds-benchmark
+        continue-on-error: true
+        uses: aws-actions/configure-aws-credentials@8df5847569e6427dd6c4fb1cf565c83acfa8afa7
+        with:
+          role-to-assume: arn:aws:iam::308535385114:role/gha_workflow_upload-benchmark-results
+          aws-region: us-east-1
+
+      - name: Upload pytest cache if tests failed
+        uses: pytorch/pytorch/.github/actions/pytest-cache-upload@release/2.12
+        continue-on-error: true
+        if: failure() && steps.test.conclusion && steps.test.conclusion == 'failure' && steps.aws-creds-benchmark.outcome == 'success'
+        with:
+          cache_dir: .pytest_cache
+          shard: ${{ matrix.shard }}
+          sha: ${{ github.event.pull_request.head.sha || github.sha }}
+          test_config: ${{ matrix.config }}
+          job_identifier: ${{ github.workflow }}_${{ inputs.build-environment }}
+
+      - name: Upload the benchmark results
+        if: steps.aws-creds-benchmark.outcome == 'success'
+        uses: pytorch/test-infra/.github/actions/upload-benchmark-results@release/2.12
+        with:
+          benchmark-results-dir: test/test-reports
+          dry-run: false
+          schema-version: v3
+          github-token: ${{ secrets.GITHUB_TOKEN }}
+
+      - name: Print remaining test logs
+        shell: bash
+        if: always() && steps.test.conclusion
+        run: |
+          cat test/**/*_toprint.log || true
+
+      - name: Upload test artifacts
+        uses: pytorch/pytorch/.github/actions/upload-test-artifacts@release/2.12
+        if: always() && steps.test.conclusion && steps.test.conclusion != 'skipped'
+        with:
+          file-suffix: ${{ github.job }}-${{ matrix.config }}-${{ matrix.shard }}-${{ matrix.num_shards }}-${{ matrix.runner }}_${{ steps.setup-linux.outputs.job-id }}
+          use-gha: ${{ inputs.use-gha }}
+          s3-bucket: ${{ inputs.s3-bucket }}
diff --git a/.github/workflows/_mac-build.yml b/.github/workflows/_mac-build.yml
index 4fd7874ee0c4d..e4ffcef200656 100644
--- a/.github/workflows/_mac-build.yml
+++ b/.github/workflows/_mac-build.yml
@@ -71,11 +71,11 @@ jobs:
       build-environment: ${{ inputs.build-environment }}
     steps:
       - name: Clean up disk space before running MacOS workflow
-        uses: pytorch/test-infra/.github/actions/check-disk-space@main
+        uses: pytorch/test-infra/.github/actions/check-disk-space@release/2.12
 
       # [see note: pytorch repo ref]
       - name: Checkout PyTorch
-        uses: pytorch/pytorch/.github/actions/checkout-pytorch@main
+        uses: pytorch/pytorch/.github/actions/checkout-pytorch@release/2.12
 
       - name: Set xcode version
         env:
@@ -86,7 +86,7 @@ jobs:
           fi
 
       - name: Setup Python
-        uses: pytorch/test-infra/.github/actions/setup-python@main
+        uses: pytorch/test-infra/.github/actions/setup-python@release/2.12
         with:
           python-version: ${{ inputs.python-version }}
           pip-requirements-file: .ci/docker/requirements-ci.txt
@@ -192,4 +192,4 @@ jobs:
       - name: Clean up disk space
         if: always()
         continue-on-error: true
-        uses: pytorch/test-infra/.github/actions/check-disk-space@main
+        uses: pytorch/test-infra/.github/actions/check-disk-space@release/2.12
diff --git a/.github/workflows/_mac-test.yml b/.github/workflows/_mac-test.yml
index 82eb3c4bf2c75..8a1464e207703 100644
--- a/.github/workflows/_mac-test.yml
+++ b/.github/workflows/_mac-test.yml
@@ -105,11 +105,11 @@ jobs:
           done
 
       - name: Clean up disk space before running MacOS workflow
-        uses: pytorch/test-infra/.github/actions/check-disk-space@main
+        uses: pytorch/test-infra/.github/actions/check-disk-space@release/2.12
 
       # [see note: pytorch repo ref]
       - name: Checkout PyTorch
-        uses: pytorch/pytorch/.github/actions/checkout-pytorch@main
+        uses: pytorch/pytorch/.github/actions/checkout-pytorch@release/2.12
 
       - name: Get workflow job id
         id: get-job-id
@@ -119,7 +119,7 @@ jobs:
           github-token: ${{ secrets.GITHUB_TOKEN }}
 
       - name: Setup Python
-        uses: pytorch/test-infra/.github/actions/setup-python@main
+        uses: pytorch/test-infra/.github/actions/setup-python@release/2.12
         with:
           python-version: ${{ inputs.python-version }}
           pip-requirements-file: .ci/docker/requirements-ci.txt
@@ -257,7 +257,7 @@ jobs:
           file-suffix: ${{ github.job }}-${{ matrix.config }}-${{ matrix.shard }}-${{ matrix.num_shards }}-${{ matrix.runner }}_${{ steps.get-job-id.outputs.job-id }}
 
       - name: Upload the benchmark results
-        uses: pytorch/test-infra/.github/actions/upload-benchmark-results@main
+        uses: pytorch/test-infra/.github/actions/upload-benchmark-results@release/2.12
         with:
           benchmark-results-dir: test/test-reports
           dry-run: false
@@ -287,4 +287,4 @@ jobs:
       - name: Clean up disk space
         if: always()
         continue-on-error: true
-        uses: pytorch/test-infra/.github/actions/check-disk-space@main
+        uses: pytorch/test-infra/.github/actions/check-disk-space@release/2.12
diff --git a/.github/workflows/_rocm-test.yml b/.github/workflows/_rocm-test.yml
index 38f3bff66c14e..1cd9cc3381610 100644
--- a/.github/workflows/_rocm-test.yml
+++ b/.github/workflows/_rocm-test.yml
@@ -85,7 +85,7 @@ jobs:
     timeout-minutes: ${{ matrix.mem_leak_check == 'mem_leak_check' && 600 || inputs.timeout-minutes }}
     steps:
       - name: Checkout PyTorch
-        uses: pytorch/pytorch/.github/actions/checkout-pytorch@main
+        uses: pytorch/pytorch/.github/actions/checkout-pytorch@release/2.12
         with:
           no-sudo: true
 
@@ -104,12 +104,12 @@ jobs:
 
       - name: Calculate docker image
         id: calculate-docker-image
-        uses: pytorch/test-infra/.github/actions/calculate-docker-image@main
+        uses: pytorch/test-infra/.github/actions/calculate-docker-image@release/2.12
         with:
           docker-image-name: ${{ inputs.docker-image }}
 
       - name: Pull docker image
-        uses: pytorch/test-infra/.github/actions/pull-docker-image@main
+        uses: pytorch/test-infra/.github/actions/pull-docker-image@release/2.12
         with:
           docker-image: ${{ steps.calculate-docker-image.outputs.docker-image }}
 
@@ -282,7 +282,7 @@ jobs:
           docker exec -t "${{ env.CONTAINER_NAME }}" sh -c "cd ../pytorch && sudo cp -R test/test-reports ../workspace/test"
 
       - name: Change permissions (only needed for kubernetes runners for now)
-        if: ${{ always() && steps.test.conclusion && (contains(matrix.runner, 'gfx942') || contains(matrix.runner, 'gfx950')) }}
+        if: ${{ always() && steps.test.conclusion && (contains(matrix.runner, 'gfx942') || contains(matrix.runner, 'gfx950') || contains(matrix.runner, 'mi210')) }}
         run: |
           docker exec -t "${{ env.CONTAINER_NAME }}" sh -c "sudo chown -R 1001:1001 test"
 
@@ -332,7 +332,7 @@ jobs:
           aws-region: us-east-1
 
       - name: Upload the benchmark results
-        uses: pytorch/test-infra/.github/actions/upload-benchmark-results@main
+        uses: pytorch/test-infra/.github/actions/upload-benchmark-results@release/2.12
         with:
           benchmark-results-dir: test/test-reports
           dry-run: false
diff --git a/.github/workflows/_runner-determinator.yml b/.github/workflows/_runner-determinator.yml
index 0d674f044ec42..06aefad9dd634 100644
--- a/.github/workflows/_runner-determinator.yml
+++ b/.github/workflows/_runner-determinator.yml
@@ -7,7 +7,8 @@ on:
         required: false
         type: string
         description: |
-          List of experiments for this workflow. If not defined, all default experiments are included.
+          Comma-separated list of non-default experiments to opt into for this workflow.
+          These are added on top of all default experiments. If not defined, only default experiments are included.
       opt_out_experiments:
         required: false
         type: string
@@ -41,6 +42,9 @@ on:
       label-type:
         description: Type of runners to use
         value: ${{ jobs.runner-determinator.outputs.label-type }}
+      use-arc:
+        description: Whether to use ARC runners
+        value: ${{ jobs.runner-determinator.outputs.use-arc }}
 
 jobs:
   runner-determinator:
@@ -49,6 +53,7 @@ jobs:
     runs-on: ubuntu-latest
     outputs:
       label-type: ${{ steps.set-condition.outputs.label-type }}
+      use-arc: ${{ steps.set-condition.outputs.use-arc }}
     env:
       GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
       ISSUE_NUMBER: ${{ inputs.issue_number }}
@@ -58,658 +63,8 @@ jobs:
       OPT_OUT_EXPERIMENTS: ${{ inputs.opt_out_experiments }}
       PR_NUMBER: ${{ github.event.pull_request.number }}
     steps:
-      # - name: Checkout PyTorch
-      #   uses: pytorch/pytorch/.github/actions/checkout-pytorch@main
-      #   with:
-      #     fetch-depth: 1
-      #     submodules: true
-
-      # TODO: Remove the hardcoded step below
-      # Hardcoding below is temporary for testing ALI runners
-      # This file below should match the script found in .github/scripts/runner_determinator.py
-      - name: Hardcode runner-determinator script
-        id: hardcode-script
-        run: |
-          cat <<EOF > runner_determinator.py
-          # flake8: noqa: G004
-
-          # Note: Copies of this script in runner_determinator.py and _runner-determinator.yml
-          #       must be kept in sync. You can do it easily by running the following command:
-          #           python .github/scripts/update_runner_determinator.py
-
-          """
-          This runner determinator is used to determine which set of runners to run a
-          GitHub job on. It uses the first comment of a GitHub issue (by default
-          https://github.com/pytorch/test-infra/issues/5132) to define the configuration
-          of which runners should be used to run which job.
-
-          The configuration has two parts, the settings and a list of opted-in users,
-          separated by a line containing "---".  If the line is not present, the
-          settings are considered to be empty with only the second part, the user
-          list, defined.
-
-          The first part is a YAML block that defines the rollout settings. This can be
-          used to define any settings that are needed to determine which runners to use.
-          It's fields are defined by the RolloutSettings class below.
-
-          The second part is a list of users who are explicitly opted in to the LF fleet.
-          The user list is also a comma separated list of additional features or
-          experiments which the user could be opted in to.
-
-          The user list has the following rules:
-
-          - Users are GitHub usernames, which must start with the @ prefix
-          - Each user is also a comma-separated list of features/experiments to enable
-          - A "#" prefix opts the user out of all experiments
-
-          Example config:
-              # A list of experiments that can be opted into.
-              # This defines the behavior they'll induce when opted into.
-              # Expected syntax is:
-              #   [experiment_name]: # Name of the experiment. Also used for the label prefix.
-              #      rollout_perc: [int] # % of workflows to run with this experiment when users are not opted in.
-
-              experiments:
-                lf:
-                  rollout_percent: 25
-                  all_branches: false
-                  default: true
-              ---
-
-              # Opt-ins:
-              # Users can opt into the LF fleet by adding their GitHub username to this list
-              # and specifying experiments to enable in a comma-separated list.
-              # To always opt out of an experiment, prefix it with a "-".
-              # Experiments should be from the above list.
-
-              @User1,-lf,split_build
-              @User2,lf
-              @User3,split_build
-          """
-
-          import json
-          import logging
-          import os
-          import random
-          import re
-          import sys
-          from argparse import ArgumentParser
-          from collections.abc import Iterable
-          from functools import cache
-          from logging import LogRecord
-          from typing import Any, NamedTuple
-          from urllib.request import Request, urlopen
-
-          import yaml
-          from github import Auth, Github
-          from github.Issue import Issue
-
-
-          DEFAULT_LABEL_PREFIX = ""  # use meta runners
-          WORKFLOW_LABEL_LF = "lf."  # use runners from the linux foundation
-          WORKFLOW_LABEL_LF_CANARY = "lf.c."  # use canary runners from the linux foundation
-
-          GITHUB_OUTPUT = os.getenv("GITHUB_OUTPUT", "")
-          GH_OUTPUT_KEY_AMI = "runner-ami"
-          GH_OUTPUT_KEY_LABEL_TYPE = "label-type"
-          OPT_OUT_LABEL = "no-runner-experiments"
-
-          SETTING_EXPERIMENTS = "experiments"
-
-          LF_FLEET_EXPERIMENT = "lf"
-          CANARY_FLEET_SUFFIX = ".c"
-
-
-          class Experiment(NamedTuple):
-              rollout_perc: float = (
-                  0  # Percentage of workflows to experiment on when user is not opted-in.
-              )
-              all_branches: bool = (
-                  False  # If True, the experiment is also enabled on the exception branches
-              )
-              default: bool = (
-                  True  # If True, the experiment is enabled by default for all queries
-              )
-
-              # Add more fields as needed
-
-
-          class Settings(NamedTuple):
-              """
-              Settings for the experiments that can be opted into.
-              """
-
-              experiments: dict[str, Experiment] = {}
-
-
-          class ColorFormatter(logging.Formatter):
-              """Color codes the log messages based on the log level"""
-
-              COLORS = {
-                  "WARNING": "\033[33m",  # Yellow
-                  "ERROR": "\033[31m",  # Red
-                  "CRITICAL": "\033[31m",  # Red
-                  "INFO": "\033[0m",  # Reset
-                  "DEBUG": "\033[0m",  # Reset
-              }
-
-              def format(self, record: LogRecord) -> str:
-                  log_color = self.COLORS.get(record.levelname, "\033[0m")  # Default to reset
-                  record.msg = f"{log_color}{record.msg}\033[0m"
-                  return super().format(record)
-
-
-          handler = logging.StreamHandler()
-          handler.setFormatter(ColorFormatter(fmt="%(levelname)-8s: %(message)s"))
-
-          log = logging.getLogger(os.path.basename(__file__))
-          log.addHandler(handler)
-          log.setLevel(logging.INFO)
-
-
-          def set_github_output(key: str, value: str) -> None:
-              """
-              Defines outputs of the github action that invokes this script
-              """
-              if not GITHUB_OUTPUT:
-                  # See https://github.blog/changelog/2022-10-11-github-actions-deprecating-save-state-and-set-output-commands/ for deprecation notice
-                  log.warning(
-                      "No env var found for GITHUB_OUTPUT, you must be running this code locally. Falling back to the deprecated print method."
-                  )
-                  print(f"::set-output name={key}::{value}")
-                  return
-
-              with open(GITHUB_OUTPUT, "a") as f:
-                  log.info(f"Setting output: {key}='{value}'")
-                  f.write(f"{key}={value}\n")
-
-
-          def _str_comma_separated_to_set(value: str) -> frozenset[str]:
-              return frozenset(
-                  filter(lambda itm: itm != "", map(str.strip, value.strip(" \n\t").split(",")))
-              )
-
-
-          def parse_args() -> Any:
-              parser = ArgumentParser("Get dynamic rollout settings")
-              parser.add_argument("--github-token", type=str, required=True, help="GitHub token")
-              parser.add_argument(
-                  "--github-issue-repo",
-                  type=str,
-                  required=False,
-                  default="pytorch/test-infra",
-                  help="GitHub repo to get the issue",
-              )
-              parser.add_argument(
-                  "--github-repo",
-                  type=str,
-                  required=True,
-                  help="GitHub repo where CI is running",
-              )
-              parser.add_argument(
-                  "--github-issue", type=int, required=True, help="GitHub issue number"
-              )
-              parser.add_argument(
-                  "--github-actor", type=str, required=True, help="GitHub triggering_actor"
-              )
-              parser.add_argument(
-                  "--github-issue-owner", type=str, required=True, help="GitHub issue owner"
-              )
-              parser.add_argument(
-                  "--github-branch", type=str, required=True, help="Current GitHub branch or tag"
-              )
-              parser.add_argument(
-                  "--github-ref-type",
-                  type=str,
-                  required=True,
-                  help="Current GitHub ref type, branch or tag",
-              )
-              parser.add_argument(
-                  "--eligible-experiments",
-                  type=_str_comma_separated_to_set,
-                  required=False,
-                  default="",
-                  help="comma separated list of experiments to check, if omitted all experiments marked with default=True are checked",
-              )
-              parser.add_argument(
-                  "--opt-out-experiments",
-                  type=_str_comma_separated_to_set,
-                  required=False,
-                  default="",
-                  help=(
-                      "comma separated list of experiments to opt-out of. If unset, no opt-outs will occur. "
-                      "If the same experiment is listed both here and in '--eligible-experiments' opt-out will take priority."
-                  ),
-              )
-              parser.add_argument(
-                  "--pr-number",
-                  type=str,
-                  required=False,
-                  default="",
-                  help="the optional PR number where this is run",
-              )
-
-              return parser.parse_args()
-
-
-          def get_gh_client(github_token: str) -> Github:  # type: ignore[no-any-unimported]
-              auth = Auth.Token(github_token)
-              return Github(auth=auth)
-
-
-          def get_issue(gh: Github, repo: str, issue_num: int) -> Issue:  # type: ignore[no-any-unimported]
-              repo = gh.get_repo(repo)
-              return repo.get_issue(number=issue_num)
-
-
-          def get_potential_pr_author(
-              github_token: str, repo: str, username: str, ref_type: str, ref_name: str
-          ) -> str:
-              # If the trigger was a new tag added by a bot, this is a ciflow case
-              # Fetch the actual username from the original PR. The PR number is
-              # embedded in the tag name: ciflow/<name>/<pr-number>
-
-              gh = get_gh_client(github_token)
-
-              if username == "pytorch-bot[bot]" and ref_type == "tag":
-                  split_tag = ref_name.split("/")
-                  if (
-                      len(split_tag) == 3
-                      and split_tag[0] == "ciflow"
-                      and split_tag[2].isnumeric()
-                  ):
-                      pr_number = split_tag[2]
-                      try:
-                          repository = gh.get_repo(repo)
-                          pull = repository.get_pull(number=int(pr_number))
-                      except Exception as e:
-                          raise Exception(  # noqa: TRY002
-                              f"issue with pull request {pr_number} from repo {repository}"
-                          ) from e
-                      return pull.user.login  # type: ignore[no-any-return]
-              # In all other cases, return the original input username
-              return username
-
-
-          def is_exception_branch(branch: str) -> bool:
-              """
-              Branches that get opted out of experiments by default, until they're explicitly enabled.
-              """
-              return branch.split("/")[0] in {"main", "nightly", "release", "landchecks"}
-
-
-          def load_yaml(yaml_text: str) -> Any:
-              try:
-                  data = yaml.safe_load(yaml_text)
-                  return data
-              except yaml.YAMLError:
-                  log.exception("Error loading YAML")
-                  raise
-
-
-          def extract_settings_user_opt_in_from_text(rollout_state: str) -> tuple[str, str]:
-              """
-              Extracts the text with settings, if any, and the opted in users from the rollout state.
-
-              If the issue body contains "---" then the text above that is the settings
-              and the text below is the list of opted in users.
-
-              If it doesn't contain "---" then the settings are empty and the rest is the users.
-              """
-              rollout_state_parts = rollout_state.split("---")
-              if len(rollout_state_parts) >= 2:
-                  return rollout_state_parts[0], rollout_state_parts[1]
-              else:
-                  return "", rollout_state
-
-
-          class UserOptins(dict[str, list[str]]):
-              """
-              Dictionary of users with a list of features they have opted into
-              """
-
-
-          def parse_user_opt_in_from_text(user_optin_text: str) -> UserOptins:
-              """
-              Parse the user opt-in text into a key value pair of username and the list of features they have opted into
-
-              Users are GitHub usernames with the @ prefix. Each user is also a comma-separated list of features/experiments to enable.
-                  - Example line: "@User1,lf,split_build"
-                  - A "#" prefix indicates the user is opted out of all experiments
-
-
-              """
-              optins = UserOptins()
-              for user in user_optin_text.split("\n"):
-                  user = user.strip("\r\n\t -")
-                  if not user or not user.startswith("@"):
-                      # Not a valid user. Skip
-                      continue
-
-                  if user:
-                      usr_name = user.split(",")[0].strip("@")
-                      optins[usr_name] = [exp.strip(" ") for exp in user.split(",")[1:]]
-
-              return optins
-
-
-          def is_valid_experiment_name(experiment_name: str) -> bool:
-              """
-              Check if the experiment name is valid.
-              A valid name:
-                  - Contains only alphanumeric characters and the special characters "_" & "-"
-                  - The special characters "_" & "-" shouldn't be the first or last characters
-                  - Cannot contain spaces
-              """
-
-              valid_char_regex = r"^[a-zA-Z0-9]([\w-]*[a-zA-Z0-9])?$"
-              valid = bool(re.match(valid_char_regex, experiment_name))
-
-              if valid:
-                  return True
-
-              log.error(
-                  f"Invalid experiment name: {experiment_name}. Experiment names should only contain alphanumeric characters, '_', and '-'. They cannot contain spaces, and the special characters '_' and '-' cannot be the first or last characters."
-              )
-              return False
-
-
-          def parse_settings_from_text(settings_text: str) -> Settings:
-              """
-              Parse the experiments from the issue body into a list of ExperimentSettings
-              """
-              try:
-                  if settings_text:
-                      # Escape the backtick as well so that we can have the settings in a code block on the GH issue
-                      # for easy reading
-                      # Note: Using ascii for the backtick so that the cat step in _runner-determinator.yml doesn't choke on
-                      #       the backtick character in shell commands.
-                      backtick = chr(96)  # backtick character
-                      settings_text = settings_text.strip(f"\r\n\t{backtick} ")
-                      settings = load_yaml(settings_text)
-
-                      # For now we just load experiments. We can expand this if/when we add more settings
-                      experiments = {}
-
-                      for exp_name, exp_settings in settings.get(SETTING_EXPERIMENTS).items():
-                          if not is_valid_experiment_name(exp_name):
-                              # Exclude invalid experiments from the list. We log an error, but don't raise an exception so that other experiments can still be processed.
-                              continue
-
-                          valid_settings = {}
-                          for setting in exp_settings:
-                              if setting not in Experiment._fields:
-                                  log.warning(
-                                      f"Unexpected setting in experiment: {setting} = {exp_settings[setting]}"
-                                  )
-                              else:
-                                  valid_settings[setting] = exp_settings[setting]
-
-                          experiments[exp_name] = Experiment(**valid_settings)
-                      return Settings(experiments)
-
-              except Exception:
-                  log.exception("Failed to parse settings")
-
-              return Settings()
-
-
-          def parse_settings(rollout_state: str) -> Settings:
-              """
-              Parse settings, if any, from the rollout state.
-
-              If the issue body contains "---" then the text above that is the settings
-              and the text below is the list of opted in users.
-
-              If it doesn't contain "---" then the settings are empty and the default values are used.
-              """
-              settings_text, _ = extract_settings_user_opt_in_from_text(rollout_state)
-              return parse_settings_from_text(settings_text)
-
-
-          def parse_users(rollout_state: str) -> UserOptins:
-              """
-              Parse users from the rollout state.
-
-              """
-              _, users_text = extract_settings_user_opt_in_from_text(rollout_state)
-              return parse_user_opt_in_from_text(users_text)
-
-
-          def is_user_opted_in(user: str, user_optins: UserOptins, experiment_name: str) -> bool:
-              """
-              Check if a user is opted into an experiment
-              """
-              return experiment_name in user_optins.get(user, [])
-
-
-          def is_user_opted_out(user: str, user_optins: UserOptins, experiment_name: str) -> bool:
-              """
-              Check if a user explicitly opted out of an experiment
-              """
-              # if the experiment is prefixed with a "-", then it's an opt-out
-              experiment_optout = "-" + experiment_name
-              if experiment_optout not in user_optins.get(user, []):
-                  return False
-
-              if is_user_opted_in(user, user_optins, experiment_name):
-                  log.warning(
-                      f"User {user} is opted into experiment {experiment_name}, but also opted out of it. Defaulting to opting out"
-                  )
-
-              return True
-
-
-          def get_runner_prefix(
-              rollout_state: str,
-              workflow_requestors: Iterable[str],
-              branch: str,
-              eligible_experiments: frozenset[str] = frozenset(),
-              opt_out_experiments: frozenset[str] = frozenset(),
-              is_canary: bool = False,
-          ) -> str:
-              settings = parse_settings(rollout_state)
-              user_optins = parse_users(rollout_state)
-
-              fleet_prefix = ""
-              prefixes = []
-              for experiment_name, experiment_settings in settings.experiments.items():
-                  if not experiment_settings.all_branches and is_exception_branch(branch):
-                      log.info(
-                          f"Branch {branch} is an exception branch. Not enabling experiment {experiment_name}."
-                      )
-                      continue
-
-                  if opt_out_experiments:
-                      if experiment_name in opt_out_experiments:
-                          opt_out_exp_list = ", ".join(opt_out_experiments)
-                          log.info(
-                              f"Skipping experiment '{experiment_name}', as this workflow has opted-out (opted out experiments are: {opt_out_exp_list})"
-                          )
-                          continue
-
-                  if eligible_experiments:
-                      if experiment_name not in eligible_experiments:
-                          exp_list = ", ".join(eligible_experiments)
-                          log.info(
-                              f"Skipping experiment '{experiment_name}', as it is not in the eligible_experiments list: {exp_list}"
-                          )
-                          continue
-                  elif not experiment_settings.default:
-                      log.info(
-                          f"Skipping experiment '{experiment_name}', as it is not a default experiment"
-                      )
-                      continue
-
-                  # Is any workflow_requestor opted out to this experiment?
-                  opted_out_users = [
-                      requestor
-                      for requestor in workflow_requestors
-                      if is_user_opted_out(requestor, user_optins, experiment_name)
-                  ]
-
-                  if opted_out_users:
-                      log.info(
-                          f"{', '.join(opted_out_users)} have opted out of experiment {experiment_name}."
-                      )
-                      continue
-
-                  # Is any workflow_requestor opted in to this experiment?
-                  opted_in_users = [
-                      requestor
-                      for requestor in workflow_requestors
-                      if is_user_opted_in(requestor, user_optins, experiment_name)
-                  ]
-
-                  enabled = False
-                  if opted_in_users:
-                      log.info(
-                          f"{', '.join(opted_in_users)} have opted into experiment {experiment_name}."
-                      )
-                      enabled = True
-
-                  elif experiment_settings.rollout_perc:
-                      # If no user is opted in, then we randomly enable the experiment based on the rollout percentage
-                      if random.uniform(0, 100) <= experiment_settings.rollout_perc:
-                          log.info(
-                              f"Based on rollout percentage of {experiment_settings.rollout_perc}%, enabling experiment {experiment_name}."
-                          )
-                          enabled = True
-
-                  if enabled:
-                      label = experiment_name
-                      if experiment_name == LF_FLEET_EXPERIMENT:
-                          # We give some special treatment to the "lf" experiment since determines the fleet we use
-                          #  - If it's enabled, then we always list it's prefix first
-                          #  - If we're in the canary branch, then we append ".c" to the lf prefix
-                          if is_canary:
-                              label += CANARY_FLEET_SUFFIX
-                          fleet_prefix = label
-                      else:
-                          prefixes.append(label)
-
-              if len(prefixes) > 1:
-                  log.error(
-                      f"Only a fleet and one other experiment can be enabled for a job at any time. Enabling {prefixes[0]} and ignoring the rest, which are {', '.join(prefixes[1:])}"
-                  )
-                  prefixes = prefixes[:1]
-
-              # Fleet always comes first
-              if fleet_prefix:
-                  prefixes.insert(0, fleet_prefix)
-
-              return ".".join(prefixes) + "." if prefixes else ""
-
-
-          def get_rollout_state_from_issue(github_token: str, repo: str, issue_num: int) -> str:
-              """
-              Gets the first comment of the issue, which contains the desired rollout state.
-
-              The default issue we use - https://github.com/pytorch/test-infra/issues/5132
-              """
-              gh = get_gh_client(github_token)
-              issue = get_issue(gh, repo, issue_num)
-              return str(issue.get_comments()[0].body.strip("\n\t "))
-
-
-          def download_json(url: str, headers: dict[str, str], num_retries: int = 3) -> Any:
-              for _ in range(num_retries):
-                  try:
-                      req = Request(url=url, headers=headers)
-                      content = urlopen(req, timeout=5).read().decode("utf-8")
-                      return json.loads(content)
-                  except Exception as e:
-                      log.warning(f"Could not download {url}: {e}")
-
-              log.warning(f"All {num_retries} retries exhausted, downloading {url} failed")
-              return {}
-
-
-          @cache
-          def get_pr_info(github_repo: str, github_token: str, pr_number: int) -> dict[str, Any]:
-              """
-              Dynamically get PR information
-              """
-              github_api = f"https://api.github.com/repos/{github_repo}"
-              headers = {
-                  "Accept": "application/vnd.github.v3+json",
-                  "Authorization": f"token {github_token}",
-              }
-              json_response: dict[str, Any] = download_json(
-                  url=f"{github_api}/issues/{pr_number}",
-                  headers=headers,
-              )
-
-              if not json_response:
-                  log.warning(f"Failed to get the labels for #{pr_number}")
-                  return {}
-
-              return json_response
-
-
-          def get_labels(github_repo: str, github_token: str, pr_number: int) -> set[str]:
-              """
-              Dynamically get the latest list of labels from the pull request
-              """
-              pr_info = get_pr_info(github_repo, github_token, pr_number)
-              return {
-                  label.get("name") for label in pr_info.get("labels", []) if label.get("name")
-              }
-
-
-          def main() -> None:
-              args = parse_args()
-
-              runner_label_prefix = DEFAULT_LABEL_PREFIX
-
-              # Check if the PR is opt-out
-              if args.pr_number:
-                  labels = get_labels(args.github_repo, args.github_token, int(args.pr_number))
-                  if OPT_OUT_LABEL in labels:
-                      log.info(
-                          f"Opt-out runner determinator because #{args.pr_number} has {OPT_OUT_LABEL} label"
-                      )
-                      set_github_output(GH_OUTPUT_KEY_LABEL_TYPE, runner_label_prefix)
-                      sys.exit()
-
-              try:
-                  rollout_state = get_rollout_state_from_issue(
-                      args.github_token, args.github_issue_repo, args.github_issue
-                  )
-
-                  username = get_potential_pr_author(
-                      args.github_token,
-                      args.github_repo,
-                      args.github_actor,
-                      args.github_ref_type,
-                      args.github_branch,
-                  )
-
-                  is_canary = args.github_repo == "pytorch/pytorch-canary"
-
-                  runner_label_prefix = get_runner_prefix(
-                      rollout_state,
-                      (args.github_issue_owner, username),
-                      args.github_branch,
-                      args.eligible_experiments,
-                      args.opt_out_experiments,
-                      is_canary,
-                  )
-
-              except Exception as e:
-                  log.error(
-                      f"Failed to get issue. Defaulting to Meta runners and no experiments. Exception: {e}"
-                  )
-
-              set_github_output(GH_OUTPUT_KEY_LABEL_TYPE, runner_label_prefix)
-
-
-          if __name__ == "__main__":
-              main()
-
-          EOF
-
-          cat runner_determinator.py
+      - name: Checkout PyTorch
+        uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
 
       - name: Install dependencies
         run: python3 -m pip install urllib3==1.26.18 PyGithub==2.3.0
@@ -721,7 +76,7 @@ jobs:
           curr_ref_type="${{ inputs.curr_ref_type }}"
           echo "Current branch is '$curr_branch'"
 
-          python3 runner_determinator.py \
+          python3 .github/scripts/runner_determinator.py \
             --github-token "$GITHUB_TOKEN" \
             --github-issue "$ISSUE_NUMBER" \
             --github-branch "$curr_branch" \
diff --git a/.github/workflows/_vllm-benchmark.yml b/.github/workflows/_vllm-benchmark.yml
index d25c3c80767b1..7a69ccaa10d08 100644
--- a/.github/workflows/_vllm-benchmark.yml
+++ b/.github/workflows/_vllm-benchmark.yml
@@ -14,7 +14,7 @@ on:
       build_environment:
         required: true
         type: string
-        description: The build environment name, e.g. linux-jammy-cuda12.9-py3.12-gcc11
+        description: The build environment name, e.g. linux-jammy-cuda13.0-py3.12-gcc11
       pytorch_branch:
         required: false
         type: string
@@ -28,6 +28,16 @@ on:
         type: string
         description: The list of models to benchmark. Default to all models.
         default: ''
+      compilation_config:
+        required: false
+        type: string
+        description: A JSON string for vLLM --compilation-config, applied to all benchmark tests.
+        default: ''
+      is_nightly:
+        required: false
+        type: boolean
+        description: Whether this is a nightly scheduled run. When true, TRANSFORMERS_OFFLINE is disabled.
+        default: false
 
 jobs:
   benchmark:
@@ -84,7 +94,7 @@ jobs:
           name: ${{ inputs.build_environment }}
           s3-bucket: gha-artifacts
 
-      - uses: pytorch/test-infra/.github/actions/setup-uv@main
+      - uses: pytorch/test-infra/.github/actions/setup-uv@release/2.12
         with:
           python-version: "3.12"
           activate-environment: "true"
@@ -106,7 +116,7 @@ jobs:
             dist/ao/torchao-*.whl \
             dist/vllm/vllm-*.whl \
             dist/deepgemm/deep_gemm-*.whl \
-            --extra-index-url https://download.pytorch.org/whl/cu129 \
+            --extra-index-url https://download.pytorch.org/whl/cu130 \
             --index-strategy unsafe-best-match
 
       - name: Print some debug information
@@ -163,6 +173,7 @@ jobs:
       - name: Setup benchmark tests
         env:
           MODELS: ${{ inputs.models }}
+          COMPILATION_CONFIG: ${{ inputs.compilation_config }}
         shell: bash
         run: |
           set -eux
@@ -178,7 +189,9 @@ jobs:
             --to-benchmark-configs-dir ../../vllm-project/vllm/.buildkite/performance-benchmarks/tests \
             --models "${MODELS}" \
             --device "${DEVICE_NAME}" \
-            --include-eager-mode
+            --include-eager-mode \
+            --include-inductor-graph-partition \
+            --compilation-config "${COMPILATION_CONFIG}"
           popd
 
           pushd vllm-project/vllm
@@ -189,9 +202,11 @@ jobs:
       - name: Run vLLM benchmark
         working-directory: vllm-project/vllm
         env:
+          FLASHINFER_WORKSPACE_BASE: /mnt/hf_cache
           HF_TOKEN: ${{ secrets.VLLM_TEST_HUGGING_FACE_TOKEN }}
           HF_HOME: /mnt/hf_cache
-          TRANSFORMERS_OFFLINE: 1
+          # Nightly runs download from HF to periodically refresh the local cache
+          TRANSFORMERS_OFFLINE: ${{ inputs.is_nightly && '0' || '1' }}
           # vLLM-related environment variables
           ENGINE_VERSION: v1
           SAVE_TO_PYTORCH_BENCHMARK_FORMAT: 1
@@ -208,7 +223,8 @@ jobs:
 
           # Fail when there is no result or if the metrics are all zero
           python3 pytorch/pytorch-integration-testing/.github/scripts/check_benchmark_results.py \
-            --benchmark-results "${BENCHMARK_RESULTS}"
+            --benchmark-results "${BENCHMARK_RESULTS}" \
+            --strict
 
       - name: Authenticate with AWS
         uses: aws-actions/configure-aws-credentials@ececac1a45f3b08a01d2dd070d28d111c5fe6722 # v4.1.0
@@ -219,7 +235,7 @@ jobs:
           aws-region: us-east-1
 
       - name: Upload the benchmark results to OSS benchmark database for the dashboard
-        uses: pytorch/test-infra/.github/actions/upload-benchmark-results@main
+        uses: pytorch/test-infra/.github/actions/upload-benchmark-results@release/2.12
         with:
           benchmark-results-dir: vllm-project/vllm/benchmarks/results
           benchmark-name: 'PyTorch x vLLM benchmark'
diff --git a/.github/workflows/_vllm-build.yml b/.github/workflows/_vllm-build.yml
index 630b5e8b6075d..b3e8c546c66f0 100644
--- a/.github/workflows/_vllm-build.yml
+++ b/.github/workflows/_vllm-build.yml
@@ -14,7 +14,7 @@ on:
       build_environment:
         required: true
         type: string
-        description: The build environment name, e.g. linux-jammy-cuda12.9-py3.12-gcc11
+        description: The build environment name, e.g. linux-jammy-cuda13.0-py3.12-gcc11
       pytorch_branch:
         required: false
         type: string
diff --git a/.github/workflows/_win-build.yml b/.github/workflows/_win-build.yml
index 005d68ece857d..d763e895596d9 100644
--- a/.github/workflows/_win-build.yml
+++ b/.github/workflows/_win-build.yml
@@ -89,7 +89,7 @@ jobs:
           git config --global core.fsmonitor false
 
       - name: Setup SSH (Click me for login details)
-        uses: pytorch/test-infra/.github/actions/setup-ssh@main
+        uses: pytorch/test-infra/.github/actions/setup-ssh@release/2.12
         with:
           github-secret: ${{ secrets.GITHUB_TOKEN }}
           instructions: |
@@ -104,7 +104,7 @@ jobs:
 
       # [see note: pytorch repo ref]
       - name: Checkout PyTorch
-        uses: pytorch/pytorch/.github/actions/checkout-pytorch@main
+        uses: pytorch/pytorch/.github/actions/checkout-pytorch@release/2.12
         with:
           no-sudo: true
 
@@ -193,9 +193,19 @@ jobs:
             cp "${CUDA_PATH}/lib/x64/cudart.lib" /c/${{ github.run_id }}/build-results/
           fi
 
+          # Copy CUDA runtime DLL (needed for MinGW import lib generation on CUDA 13.0+)
+          echo "Searching for CUDA runtime DLLs in ${CUDA_PATH}/bin/:"
+          ls -la "${CUDA_PATH}"/bin/cudart*.dll 2>/dev/null || echo "No cudart*.dll found in ${CUDA_PATH}/bin/"
+          for dll in "${CUDA_PATH}"/bin/cudart64_*.dll "${CUDA_PATH}"/bin/cudart.dll; do
+            if [ -f "$dll" ]; then
+              cp "$dll" /c/${{ github.run_id }}/build-results/
+              echo "Copied $(basename $dll)"
+            fi
+          done
+
           # List collected files
-          echo "Collected CUDA libs:"
-          ls -lah /c/${{ github.run_id }}/build-results/*.lib
+          echo "Collected CUDA libs and DLLs:"
+          ls -lah /c/${{ github.run_id }}/build-results/*.lib /c/${{ github.run_id }}/build-results/*.dll 2>/dev/null || true
 
       # Upload to github so that people can click and download artifacts
       - name: Upload artifacts to s3
diff --git a/.github/workflows/_win-test.yml b/.github/workflows/_win-test.yml
index 3d2fe8a4b3fac..b0d4b472debe4 100644
--- a/.github/workflows/_win-test.yml
+++ b/.github/workflows/_win-test.yml
@@ -78,7 +78,7 @@ jobs:
           git config --global core.fsmonitor false
 
       - name: Setup SSH (Click me for login details)
-        uses: pytorch/test-infra/.github/actions/setup-ssh@main
+        uses: pytorch/test-infra/.github/actions/setup-ssh@release/2.12
         with:
           github-secret: ${{ secrets.GITHUB_TOKEN }}
           instructions: |
@@ -94,9 +94,10 @@ jobs:
 
       # [see note: pytorch repo ref]
       - name: Checkout PyTorch
-        uses: pytorch/pytorch/.github/actions/checkout-pytorch@main
+        uses: pytorch/pytorch/.github/actions/checkout-pytorch@release/2.12
         with:
           no-sudo: true
+          submodules: false
 
       - name: Setup Windows
         uses: ./.github/actions/setup-win
diff --git a/.github/workflows/_xpu-test.yml b/.github/workflows/_xpu-test.yml
index 5724403e6de44..4caf910760f27 100644
--- a/.github/workflows/_xpu-test.yml
+++ b/.github/workflows/_xpu-test.yml
@@ -86,17 +86,14 @@ jobs:
     steps:
       # [see note: pytorch repo ref]
       - name: Checkout PyTorch
-        uses: pytorch/pytorch/.github/actions/checkout-pytorch@main
+        uses: pytorch/pytorch/.github/actions/checkout-pytorch@release/2.12
 
       - name: Setup XPU
         uses: ./.github/actions/setup-xpu
 
-      - name: Login to ECR
-        uses: ./.github/actions/ecr-login
-
       - name: Calculate docker image
         id: calculate-docker-image
-        uses: pytorch/test-infra/.github/actions/calculate-docker-image@main
+        uses: pytorch/test-infra/.github/actions/calculate-docker-image@release/2.12
         with:
           docker-image-name: ${{ inputs.docker-image }}
 
@@ -110,7 +107,7 @@ jobs:
           echo "docker pull ghcr.io/pytorch/ci-image:${tag/:/-}"
 
       - name: Pull docker image
-        uses: pytorch/test-infra/.github/actions/pull-docker-image@main
+        uses: pytorch/test-infra/.github/actions/pull-docker-image@release/2.12
         with:
           docker-image: ${{ steps.calculate-docker-image.outputs.docker-image }}
 
@@ -251,7 +248,7 @@ jobs:
             -e ZE_AFFINITY_MASK \
             -e HUGGING_FACE_HUB_TOKEN \
             -e DASHBOARD_TAG \
-            --env-file="/tmp/github_env_${GITHUB_RUN_ID}" \
+            --env-file="${RUNNER_TEMP}/github_env_${GITHUB_RUN_ID}" \
             --ulimit stack=10485760:83886080 \
             --ulimit core=0 \
             --security-opt seccomp=unconfined \
@@ -345,7 +342,7 @@ jobs:
           aws-region: us-east-1
 
       - name: Upload the benchmark results
-        uses: pytorch/test-infra/.github/actions/upload-benchmark-results@main
+        uses: pytorch/test-infra/.github/actions/upload-benchmark-results@release/2.12
         with:
           benchmark-results-dir: test/test-reports
           dry-run: false
diff --git a/.github/workflows/apply-lint.yml b/.github/workflows/apply-lint.yml
new file mode 100644
index 0000000000000..d30c0a2b316e7
--- /dev/null
+++ b/.github/workflows/apply-lint.yml
@@ -0,0 +1,138 @@
+name: Apply lint patches
+
+on:
+  repository_dispatch:
+    types: [apply-lint]
+
+jobs:
+  apply-lint:
+    runs-on: ubuntu-24.04
+    environment: mergebot
+    env:
+      GH_RUN_URL: ${{ github.server_url }}/${{ github.repository }}/actions/runs/${{ github.run_id }}
+      LINTERS: CLANGFORMAT,NEWLINE,PYFMT,RUFF,ATEN_CPU_GPU_AGNOSTIC
+    steps:
+      - name: Checkout repo
+        id: checkout
+        uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
+        with:
+          fetch-depth: 0
+          token: ${{ secrets.MERGEBOT_TOKEN }}
+
+      - name: Checkout PR branch
+        env:
+          GITHUB_TOKEN: ${{ secrets.MERGEBOT_TOKEN }}
+          PR_NUM: ${{ github.event.client_payload.pr_num }}
+        run: |
+          HEAD_REF=$(gh pr view "$PR_NUM" --json headRefName -q .headRefName)
+          if [[ "$HEAD_REF" =~ ^gh/[^/]+/[0-9]+/head$ ]]; then
+            echo "IS_GHSTACK=true" >> "$GITHUB_ENV"
+            ORIG_REF="${HEAD_REF/%head/orig}"
+            echo "ORIG_REF=$ORIG_REF" >> "$GITHUB_ENV"
+            BASE_BRANCH=$(gh pr view "$PR_NUM" --json baseRefName -q .baseRefName)
+            git fetch origin "$ORIG_REF" "$BASE_BRANCH"
+            git checkout -b "$ORIG_REF" "origin/$ORIG_REF"
+          else
+            gh pr checkout "$PR_NUM"
+          fi
+
+      - name: Setup Python
+        uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
+        with:
+          python-version: '3.10'
+          architecture: x64
+          check-latest: false
+          cache: pip
+
+      - name: Install and initialize lintrunner
+        run: |
+          pip install lintrunner pyyaml==6.0.2 uv ghstack
+          lintrunner init --take "$LINTERS"
+
+      - name: Setup committer id
+        run: |
+          if [ "$IS_GHSTACK" = "true" ]; then
+            git config --global user.email "$(git log HEAD --pretty=format:%ae -1)"
+            git config --global user.name "$(git log HEAD --pretty=format:%an -1)"
+          else
+            git config --global user.email "pytorchmergebot@users.noreply.github.com"
+            git config --global user.name "PyTorch MergeBot"
+          fi
+
+      - name: Apply lint patches
+        env:
+          GITHUB_TOKEN: ${{ secrets.MERGEBOT_TOKEN }}
+          PR_NUM: ${{ github.event.client_payload.pr_num }}
+        run: |
+          set -x
+
+          # Determine the base branch for this PR
+          BASE_BRANCH=$(gh pr view "$PR_NUM" --json baseRefName -q .baseRefName)
+
+          # Only run linters that produce auto-fixable patches (is_formatter = true)
+          lintrunner -a \
+            --take "$LINTERS" \
+            -m "origin/${BASE_BRANCH}" 2>/dev/null || true
+
+          if git diff --quiet; then
+            echo "No lint changes to apply"
+            gh pr comment "$PR_NUM" --body \
+              "No auto-fixable lint issues found, your PR is clean."
+          else
+            git add -u
+            if [ "$IS_GHSTACK" = "true" ]; then
+              git commit --amend --no-edit
+            else
+              git commit -m "Apply lint patches
+
+          Generated by apply-lint bot, see ${GH_RUN_URL}"
+            fi
+            echo "HAS_LINT_CHANGES=true" >> "$GITHUB_ENV"
+          fi
+
+      - name: Push lint changes
+        if: env.HAS_LINT_CHANGES == 'true' && env.IS_GHSTACK != 'true'
+        env:
+          GITHUB_TOKEN: ${{ secrets.MERGEBOT_TOKEN }}
+          PR_NUM: ${{ github.event.client_payload.pr_num }}
+        run: |
+          git push
+          gh pr comment "$PR_NUM" --body \
+            "Successfully applied lint patches in ${GH_RUN_URL}. Please pull locally before pushing more changes."
+
+      - name: Push lint changes (ghstack)
+        if: env.HAS_LINT_CHANGES == 'true' && env.IS_GHSTACK == 'true'
+        env:
+          GITHUB_TOKEN: ${{ secrets.MERGEBOT_TOKEN }}
+          OAUTH_TOKEN: ${{ secrets.MERGEBOT_TOKEN }}
+          PR_NUM: ${{ github.event.client_payload.pr_num }}
+        run: |
+          cat > .ghstackrc <<EOF
+          [ghstack]
+          github_url=github.com
+          github_username=pytorchmergebot
+          remote_name=origin
+          EOF
+          ghstack
+          gh pr comment "$PR_NUM" --body \
+            "Successfully applied lint patches in ${GH_RUN_URL}. Please pull locally before pushing more changes."
+
+      - name: Comment on Canceled
+        if: ${{ cancelled() && steps.checkout.outcome == 'success' }}
+        continue-on-error: true
+        env:
+          GITHUB_TOKEN: ${{ secrets.MERGEBOT_TOKEN }}
+          PR_NUM: ${{ github.event.client_payload.pr_num }}
+        run: |
+          gh pr comment "$PR_NUM" --body \
+            "Apply lint patches was cancelled, see ${GH_RUN_URL}"
+
+      - name: Comment on failure
+        if: ${{ failure() && steps.checkout.outcome == 'success' }}
+        continue-on-error: true
+        env:
+          GITHUB_TOKEN: ${{ secrets.MERGEBOT_TOKEN }}
+          PR_NUM: ${{ github.event.client_payload.pr_num }}
+        run: |
+          gh pr comment "$PR_NUM" --body \
+            "Failed to apply lint patches, see ${GH_RUN_URL}"
diff --git a/.github/workflows/attention_op_microbenchmark.yml b/.github/workflows/attention_op_microbenchmark.yml
index cd04a48223ce1..085f4cda787e8 100644
--- a/.github/workflows/attention_op_microbenchmark.yml
+++ b/.github/workflows/attention_op_microbenchmark.yml
@@ -16,6 +16,7 @@ concurrency:
 permissions:
   id-token: write
   contents: read
+  actions: read
 
 jobs:
   attn-microbenchmark-build:
diff --git a/.github/workflows/b200-distributed.yml b/.github/workflows/b200-distributed.yml
index e52c7a4b5f5c5..301e4e966563a 100644
--- a/.github/workflows/b200-distributed.yml
+++ b/.github/workflows/b200-distributed.yml
@@ -18,13 +18,14 @@ concurrency:
 permissions:
   id-token: write
   contents: read
+  actions: read
 
 jobs:
 
   get-label-type:
     if: github.repository_owner == 'pytorch'
     name: get-label-type
-    uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@main
+    uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@release/2.12
     with:
       triggering_actor: ${{ github.triggering_actor }}
       issue_owner: ${{ github.event.pull_request.user.login || github.event.issue.user.login }}
diff --git a/.github/workflows/b200-symm-mem.yml b/.github/workflows/b200-symm-mem.yml
index 62367b61b07b9..284b8cef3af98 100644
--- a/.github/workflows/b200-symm-mem.yml
+++ b/.github/workflows/b200-symm-mem.yml
@@ -18,13 +18,14 @@ concurrency:
 permissions:
   id-token: write
   contents: read
+  actions: read
 
 jobs:
 
   get-label-type:
     if: github.repository_owner == 'pytorch'
     name: get-label-type
-    uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@main
+    uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@release/2.12
     with:
       triggering_actor: ${{ github.triggering_actor }}
       issue_owner: ${{ github.event.pull_request.user.login || github.event.issue.user.login }}
diff --git a/.github/workflows/build-almalinux-images.yml b/.github/workflows/build-almalinux-images.yml
index 9090fba0a9773..2d11a97ab692c 100644
--- a/.github/workflows/build-almalinux-images.yml
+++ b/.github/workflows/build-almalinux-images.yml
@@ -36,10 +36,10 @@ jobs:
     runs-on: linux.9xlarge.ephemeral
     strategy:
       matrix:
-        tag: ["cuda12.6", "cuda12.8", "cuda12.9", "cuda13.0", "rocm7.0", "rocm7.1", "rocm7.2", "cpu"]
+        tag: ["cuda12.6", "cuda12.8", "cuda13.0", "cuda13.2", "rocm7.0", "rocm7.1", "rocm7.2", "cpu"]
     steps:
       - name: Build docker image
-        uses: pytorch/pytorch/.github/actions/binary-docker-build@main
+        uses: pytorch/pytorch/.github/actions/binary-docker-build@release/2.12
         with:
           docker-image-name: almalinux-builder
           custom-tag-prefix: ${{matrix.tag}}
diff --git a/.github/workflows/build-libtorch-images.yml b/.github/workflows/build-libtorch-images.yml
deleted file mode 100644
index 47bf15e1db3ab..0000000000000
--- a/.github/workflows/build-libtorch-images.yml
+++ /dev/null
@@ -1,68 +0,0 @@
-name: Build libtorch docker images
-
-on:
-  push:
-    branches:
-      - main
-      - release/*
-    tags:
-      # NOTE: Binary build pipelines should only get triggered on release candidate or nightly builds
-      # Release candidate tags look like: v1.11.0-rc1
-      - v[0-9]+.[0-9]+.[0-9]+-rc[0-9]+
-    paths:
-      - .ci/docker/**
-      - .github/workflows/build-libtorch-images.yml
-      - .github/actions/binary-docker-build/**
-  pull_request:
-    paths:
-      - .ci/docker/**
-      - .github/workflows/build-libtorch-images.yml
-      - .github/actions/binary-docker-build/**
-
-env:
-  DOCKER_REGISTRY: "docker.io"
-  DOCKER_BUILDKIT: 1
-  WITH_PUSH: ${{ github.event_name == 'push' && (github.ref == 'refs/heads/main' || startsWith(github.ref, 'refs/heads/release') || startsWith(github.ref, 'refs/tags/v')) }}
-
-concurrency:
-  group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}-${{ github.event_name == 'workflow_dispatch' }}
-  cancel-in-progress: true
-
-jobs:
-  get-label-type:
-    if: github.repository_owner == 'pytorch'
-    name: get-label-type
-    uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@main
-    with:
-      triggering_actor: ${{ github.triggering_actor }}
-      issue_owner: ${{ github.event.pull_request.user.login || github.event.issue.user.login }}
-      curr_branch: ${{ github.head_ref || github.ref_name }}
-      curr_ref_type: ${{ github.ref_type }}
-
-  build:
-    environment: ${{ (github.event_name == 'push' && (github.ref == 'refs/heads/main' || startsWith(github.ref, 'refs/heads/release') || startsWith(github.ref, 'refs/tags/v')) && 'docker-build') || '' }}
-    needs: get-label-type
-    runs-on: ${{ needs.get-label-type.outputs.label-type }}linux.9xlarge.ephemeral
-    name: libtorch-cxx11-builder:${{ matrix.tag }}
-    strategy:
-      fail-fast: false
-      matrix:
-        include: [
-          { tag: "cuda13.0" },
-          { tag: "cuda12.9" },
-          { tag: "cuda12.8" },
-          { tag: "cuda12.6" },
-          { tag: "rocm7.0"  },
-          { tag: "rocm7.1"  },
-          { tag: "rocm7.2"  },
-          { tag: "cpu"      },
-        ]
-    steps:
-      - name: Build docker image
-        uses: pytorch/pytorch/.github/actions/binary-docker-build@main
-        with:
-          docker-image-name: libtorch-cxx11-builder
-          custom-tag-prefix: ${{ matrix.tag }}
-          docker-build-dir: libtorch
-          DOCKER_TOKEN: ${{ secrets.DOCKER_TOKEN }}
-          DOCKER_ID: ${{ secrets.DOCKER_ID }}
diff --git a/.github/workflows/build-magma-linux.yml b/.github/workflows/build-magma-linux.yml
index be8f613169e8c..49170ca684b5f 100644
--- a/.github/workflows/build-magma-linux.yml
+++ b/.github/workflows/build-magma-linux.yml
@@ -34,7 +34,7 @@ jobs:
       id-token: write
     strategy:
       matrix:
-        cuda_version: ["130", "129", "128", "126"]
+        cuda_version: ["132", "130", "129", "128", "126"]
     steps:
       - name: Checkout PyTorch
         uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
diff --git a/.github/workflows/build-magma-windows.yml b/.github/workflows/build-magma-windows.yml
index b7d293a5cec11..39a705338f189 100644
--- a/.github/workflows/build-magma-windows.yml
+++ b/.github/workflows/build-magma-windows.yml
@@ -22,7 +22,7 @@ jobs:
     runs-on: windows-2022
     strategy:
       matrix:
-        cuda_version: ["130", "129", "128", "126"]
+        cuda_version: ["132", "130", "129", "128", "126"]
         config: ["Release", "Debug"]
     env:
       CUDA_VERSION: ${{ matrix.cuda_version }}
diff --git a/.github/workflows/build-manywheel-images-s390x.yml b/.github/workflows/build-manywheel-images-s390x.yml
index c498e169f1aa5..b15d0b4f433f0 100644
--- a/.github/workflows/build-manywheel-images-s390x.yml
+++ b/.github/workflows/build-manywheel-images-s390x.yml
@@ -25,7 +25,7 @@ jobs:
     runs-on: linux.s390x
     steps:
       - name: Checkout PyTorch
-        uses: pytorch/pytorch/.github/actions/checkout-pytorch@main
+        uses: pytorch/pytorch/.github/actions/checkout-pytorch@release/2.12
         with:
           submodules: false
           no-sudo: true
diff --git a/.github/workflows/build-manywheel-images.yml b/.github/workflows/build-manywheel-images.yml
index f86cefd7c7a1a..28e5739fba62e 100644
--- a/.github/workflows/build-manywheel-images.yml
+++ b/.github/workflows/build-manywheel-images.yml
@@ -32,7 +32,7 @@ jobs:
   get-label-type:
     if: github.repository_owner == 'pytorch'
     name: get-label-type
-    uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@main
+    uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@release/2.12
     with:
       triggering_actor: ${{ github.triggering_actor }}
       issue_owner: ${{ github.event.pull_request.user.login || github.event.issue.user.login }}
@@ -46,14 +46,14 @@ jobs:
       fail-fast: false
       matrix:
         include: [
+          { name: "manylinux2_28-builder",          tag: "cuda13.2",          runner: "linux.9xlarge.ephemeral" },
           { name: "manylinux2_28-builder",          tag: "cuda13.0",          runner: "linux.9xlarge.ephemeral" },
           { name: "manylinux2_28-builder",          tag: "cuda12.8",          runner: "linux.9xlarge.ephemeral" },
-          { name: "manylinux2_28-builder",          tag: "cuda12.9",          runner: "linux.9xlarge.ephemeral" },
           { name: "manylinux2_28-builder",          tag: "cuda12.6",          runner: "linux.9xlarge.ephemeral" },
-          { name: "manylinuxaarch64-builder",       tag: "cuda13.0",          runner: "linux.arm64.2xlarge.ephemeral" },
-          { name: "manylinuxaarch64-builder",       tag: "cuda12.9",          runner: "linux.arm64.2xlarge.ephemeral" },
-          { name: "manylinuxaarch64-builder",       tag: "cuda12.8",          runner: "linux.arm64.2xlarge.ephemeral" },
-          { name: "manylinuxaarch64-builder",       tag: "cuda12.6",          runner: "linux.arm64.2xlarge.ephemeral" },
+          { name: "manylinuxaarch64-builder",       tag: "cuda13.2",          runner: "linux.arm64.m7g.4xlarge.ephemeral" },
+          { name: "manylinuxaarch64-builder",       tag: "cuda13.0",          runner: "linux.arm64.m7g.4xlarge.ephemeral" },
+          { name: "manylinuxaarch64-builder",       tag: "cuda12.8",          runner: "linux.arm64.m7g.4xlarge.ephemeral" },
+          { name: "manylinuxaarch64-builder",       tag: "cuda12.6",          runner: "linux.arm64.m7g.4xlarge.ephemeral" },
           { name: "manylinux2_28-builder",          tag: "rocm7.0",           runner: "linux.9xlarge.ephemeral" },
           { name: "manylinux2_28-builder",          tag: "rocm7.1",           runner: "linux.9xlarge.ephemeral" },
           { name: "manylinux2_28-builder",          tag: "rocm7.2",           runner: "linux.9xlarge.ephemeral" },
@@ -65,7 +65,7 @@ jobs:
     name: ${{ matrix.name }}:${{ matrix.tag }}
     steps:
       - name: Build docker image
-        uses: pytorch/pytorch/.github/actions/binary-docker-build@main
+        uses: pytorch/pytorch/.github/actions/binary-docker-build@release/2.12
         with:
           docker-image-name: ${{ matrix.name }}
           custom-tag-prefix: ${{ matrix.tag }}
diff --git a/.github/workflows/build-triton-wheel.yml b/.github/workflows/build-triton-wheel.yml
index 263745ff0fe23..6f8aaca9d1f86 100644
--- a/.github/workflows/build-triton-wheel.yml
+++ b/.github/workflows/build-triton-wheel.yml
@@ -3,7 +3,7 @@ name: Build Triton wheels
 on:
   push:
     branches:
-      - main
+      - release/2.12
     tags:
       # NOTE: Binary build pipelines should only get triggered on release candidate builds
       # Release candidate tags look like: v1.11.0-rc1
@@ -36,7 +36,7 @@ jobs:
   get-label-type:
     if: github.repository_owner == 'pytorch'
     name: get-label-type
-    uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@main
+    uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@release/2.12
     with:
       triggering_actor: ${{ github.triggering_actor }}
       issue_owner: ${{ github.event.pull_request.user.login || github.event.issue.user.login }}
@@ -74,23 +74,20 @@ jobs:
       PLATFORM: 'manylinux_2_28_x86_64'
     steps:
       - name: Setup SSH (Click me for login details)
-        uses: pytorch/test-infra/.github/actions/setup-ssh@main
+        uses: pytorch/test-infra/.github/actions/setup-ssh@release/2.12
         with:
           github-secret: ${{ secrets.GITHUB_TOKEN }}
 
-      - name: Checkout PyTorch
-        uses: pytorch/pytorch/.github/actions/checkout-pytorch@main
+      - name: Setup Linux
+        uses: pytorch/pytorch/.github/actions/setup-linux@release/2.12
         with:
           submodules: false
 
-      - name: Setup Linux
-        uses: ./.github/actions/setup-linux
-
       - name: Login to ECR
         uses: ./.github/actions/ecr-login
 
       - name: Pull Docker image
-        uses: pytorch/test-infra/.github/actions/pull-docker-image@main
+        uses: pytorch/test-infra/.github/actions/pull-docker-image@release/2.12
         with:
           docker-image: ${{ env.DOCKER_IMAGE }}
 
@@ -165,7 +162,7 @@ jobs:
           docker exec -t "${container_name}" bash -c "${PYTHON_EXECUTABLE} /pytorch/.github/scripts/build_triton_wheel.py --device=$BUILD_DEVICE $RELEASE $WITH_CLANG_LDD"
 
           if [[ ("${{ matrix.device }}" == "cuda" || "${{ matrix.device }}" == "xpu") ]]; then
-            docker exec -t "${container_name}"  bash -c "auditwheel repair --plat ${PLATFORM} //artifacts/*.whl"
+            docker exec -t "${container_name}"  bash -c "auditwheel repair --plat ${PLATFORM} --exclude libtriton.so //artifacts/*.whl"
           else
             docker exec -t "${container_name}"  bash -c "mkdir //artifacts/wheelhouse"
             docker exec -t "${container_name}"  bash -c "mv //artifacts/*.whl //artifacts/wheelhouse/"
@@ -179,7 +176,7 @@ jobs:
           path: ${{ runner.temp }}/artifacts/wheelhouse/*
 
       - name: Teardown Linux
-        uses: pytorch/test-infra/.github/actions/teardown-linux@main
+        uses: pytorch/test-infra/.github/actions/teardown-linux@release/2.12
         if: always()
 
   build-wheel-win:
@@ -212,7 +209,7 @@ jobs:
           echo "instance-type: $(get_ec2_metadata instance-type)"
           echo "system info $(uname -a)"
       - name: "[FB EMPLOYEES] Enable SSH (Click me for login details)"
-        uses: pytorch/test-infra/.github/actions/setup-ssh@main
+        uses: pytorch/test-infra/.github/actions/setup-ssh@release/2.12
         continue-on-error: true
         with:
           github-secret: ${{ secrets.GITHUB_TOKEN }}
@@ -332,7 +329,11 @@ jobs:
           # to nightly or test
           UPLOAD_SUBFOLDER: ""
           PKG_DIR: ${{ runner.temp }}/artifacts
+          R2_UPLOAD: ${{ (github.event_name == 'push' && github.event.ref == 'refs/heads/main') && 'true' || '' }}
+          R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+          R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+          R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
         shell: bash
         run: |
           set -ex
-          bash .circleci/scripts/binary_upload.sh
+          bash .ci/pytorch/binary_upload.sh
diff --git a/.github/workflows/build-vllm-wheel.yml b/.github/workflows/build-vllm-wheel.yml
index 1ef0684688218..458e4a101acfe 100644
--- a/.github/workflows/build-vllm-wheel.yml
+++ b/.github/workflows/build-vllm-wheel.yml
@@ -65,18 +65,15 @@ jobs:
       BUILD_DEVICE: ${{ matrix.device }}
     steps:
       - name: Setup SSH (Click me for login details)
-        uses: pytorch/test-infra/.github/actions/setup-ssh@main
+        uses: pytorch/test-infra/.github/actions/setup-ssh@release/2.12
         with:
           github-secret: ${{ secrets.GITHUB_TOKEN }}
 
-      - name: Checkout PyTorch
-        uses: pytorch/pytorch/.github/actions/checkout-pytorch@main
+      - name: Setup Linux
+        uses: pytorch/pytorch/.github/actions/setup-linux@release/2.12
         with:
           submodules: false
 
-      - name: Setup Linux
-        uses: ./.github/actions/setup-linux
-
       - name: Login to ECR
         uses: ./.github/actions/ecr-login
 
@@ -167,7 +164,7 @@ jobs:
           path: ${{ runner.temp }}/artifacts/externals/vllm/wheels/*.whl
 
       - name: Teardown Linux
-        uses: pytorch/test-infra/.github/actions/teardown-linux@main
+        uses: pytorch/test-infra/.github/actions/teardown-linux@release/2.12
         if: always()
 
   # Copied from build-triton-wheel workflow (mostly)
@@ -246,7 +243,11 @@ jobs:
           PACKAGE_TYPE: wheel
           UPLOAD_SUBFOLDER: ${{ env.BUILD_DEVICE }}
           PKG_DIR: ${{ runner.temp }}/artifacts
+          R2_UPLOAD: ${{ ((github.event_name == 'push' && github.event.ref == 'refs/heads/main') || github.event_name == 'schedule' || github.event_name == 'workflow_dispatch') && 'true' || '' }}
+          R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+          R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+          R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
         shell: bash
         run: |
           set -ex
-          bash .circleci/scripts/binary_upload.sh
+          bash .ci/pytorch/binary_upload.sh
diff --git a/.github/workflows/claude-autorevert-advisor.yml b/.github/workflows/claude-autorevert-advisor.yml
new file mode 100644
index 0000000000000..483e4189fbf48
--- /dev/null
+++ b/.github/workflows/claude-autorevert-advisor.yml
@@ -0,0 +1,173 @@
+name: Claude Autorevert Advisor
+
+on:
+  workflow_dispatch:
+    inputs:
+      suspect_commit:
+        description: "SHA of the commit autorevert suspects caused the failure"
+        required: true
+        type: string
+      pr_number:
+        description: "PR number of the suspect commit"
+        required: true
+        type: string
+      signal_pattern:
+        description: "JSON object with signal metadata, failed partition (with job links/log URLs), and successful partition (baseline)"
+        required: true
+        type: string
+
+jobs:
+  analyze:
+    if: github.repository == 'pytorch/pytorch'
+    runs-on: ubuntu-latest
+    timeout-minutes: 15
+    environment: bedrock
+    permissions:
+      contents: read
+      id-token: write
+
+    steps:
+      - name: Checkout pytorch/pytorch trunk
+        uses: actions/checkout@v4
+        with:
+          ref: main
+          fetch-depth: 128
+
+      - name: Fetch suspect commit
+        run: |
+          git fetch origin ${{ inputs.suspect_commit }} --depth=32 || true
+
+      - name: Configure AWS credentials via OIDC
+        uses: aws-actions/configure-aws-credentials@7474bc4690e29a8392af63c5b98e7449536d5c3a # v4
+        with:
+          role-to-assume: arn:aws:iam::308535385114:role/gha_workflow_claude_code
+          aws-region: us-east-1
+
+      - name: Run AI Advisor
+        id: claude
+        uses: anthropics/claude-code-action@6e2bd52842c65e914eba5c8badd17560bd26b5de # v1.0.89
+        with:
+          allowed_bots: "pytorch-auto-revert[bot]"
+          use_bedrock: "true"
+          claude_args: |
+            --model global.anthropic.claude-opus-4-6-v1
+            --allowedTools "Bash,Read,Glob,Grep,WebFetch"
+            --json-schema '{"type":"object","required":["verdict","confidence","summary","causal_reasoning"],"properties":{"verdict":{"type":"string","enum":["revert","unsure","not_related","garbage"]},"confidence":{"type":"number","minimum":0,"maximum":1},"summary":{"type":"string","description":"1-3 sentence explanation of the verdict"},"causal_reasoning":{"type":"string","description":"Detailed explanation of the causal chain (or lack thereof) between the code change and the observed failures"}}}'
+          settings: '{"alwaysThinkingEnabled": true}'
+          prompt: |
+            You are an AI advisor for PyTorch's autorevert system. Your job is to analyze a suspect commit and CI failures to determine whether the commit caused the failures.
+
+            ## Context
+
+            **Suspect commit:** ${{ inputs.suspect_commit }}
+            **PR:** https://github.com/pytorch/pytorch/pull/${{ inputs.pr_number }}
+
+            **Signal pattern (partition data from autorevert):**
+            ```json
+            ${{ inputs.signal_pattern }}
+            ```
+
+            The signal pattern contains:
+            - `signal_key`: the test or job name that defines this signal
+            - `signal_source`: "test" (individual test) or "job" (whole job)
+            - `workflow_name`: the CI workflow where this signal lives
+            - `failed_partition`: commits where this signal FAILS (newest first, suspect commit marked with `is_suspect: true`). Each has job links and log URLs.
+            - `successful_partition`: baseline commits where this signal was GREEN (most recent first). These prove the signal was passing before the suspect commit.
+
+            The partition is ground truth from CI data. The signal transitioned from SUCCESS to FAILURE at the suspect commit. Your job is to determine whether the CODE CHANGE caused the transition, or whether it's coincidence (upstream dependency change, flaky environment, infra issue).
+
+            ## Investigation Steps
+
+            1. **Read the suspect commit diff.** Run `git show ${{ inputs.suspect_commit }}` to understand what code changed. Read the modified files to understand the broader context around the changes.
+
+            2. **Analyze the PR.** Fetch the PR description and understand the intent of the change: `gh pr view ${{ inputs.pr_number }} --repo pytorch/pytorch --json title,body,files`
+
+            3. **Investigate the failed jobs.** For each job in the failed_partition:
+               - Fetch the job logs from the log_url to find error messages, stack traces, and assertion failures
+               - Identify the specific test(s) or build step(s) that failed
+               - Read the relevant test source code in this checkout to understand what the test exercises
+
+            4. **Trace the causal chain.** This is the critical step. Do NOT take shortcuts:
+               - Trace imports, function calls, class hierarchies, and runtime dispatch paths from the changed code to the failing tests
+               - Consider indirect dependencies: a change to a utility function may break tests through several layers of abstraction
+               - Consider runtime behavior: changes to tensor operations, kernel dispatch, autograd behavior, or compilation paths can have non-obvious downstream effects
+               - Check if the error message or stack trace references any of the changed code, even indirectly
+
+            5. **Check for alternative explanations:**
+               - Is this a known flaky test? (Look for skip/xfail decorators, recent similar failures on other commits)
+               - Is this an infrastructure issue? (OOM, timeout, network error, GPU driver issue)
+               - Did a different concurrent commit plausibly cause this? (Check the commit timestamp and other recent merges)
+               - Is this signal fundamentally unreliable? (Persistent timeouts, environment-dependent flakes, resource exhaustion not caused by any code change)
+
+            6. **Render your verdict:**
+               - `revert` — You found a plausible causal chain from the code change to the failure. Confidence should reflect how direct and clear the chain is.
+               - `not_related` — After thorough investigation, there is no plausible causal chain. The failure has a clear alternative explanation.
+               - `unsure` — The investigation is inconclusive. There may be a connection but you cannot confirm it with sufficient confidence.
+               - `garbage` — The signal itself is unreliable and should be temporarily suppressed. Use this when the failure is clearly not caused by ANY code change: persistent infra flakes (runner timeouts, network errors, GPU driver crashes), environment issues (disk full, dependency mirror outage), or tests so flaky that the signal has no diagnostic value. This tells autorevert to ignore this signal for ~2 hours.
+
+            When in doubt between `unsure` and `not_related`, prefer `unsure`. A false `not_related` verdict is more costly than an `unsure` (which just lets autorevert continue its normal restart-to-confirm flow).
+
+            Use `garbage` only when the failure clearly has nothing to do with any code change — it's a statement about the signal quality, not about the suspect commit.
+
+            IMPORTANT: Be thorough. Read actual code. Trace actual call paths. Do not guess based on file names alone.
+
+      - name: Save verdict artifact
+        if: always() && steps.claude.outputs.structured_output != ''
+        env:
+          VERDICT_JSON: ${{ steps.claude.outputs.structured_output }}
+        run: |
+          mkdir -p /tmp/verdict
+          printf '%s' "$VERDICT_JSON" > /tmp/verdict/verdict.json
+          cat /tmp/verdict/verdict.json
+
+      - name: Upload verdict artifact
+        if: always() && steps.claude.outputs.structured_output != ''
+        uses: actions/upload-artifact@v4
+        with:
+          name: autorevert-verdict
+          path: /tmp/verdict/verdict.json
+          retention-days: 30
+
+      - name: Upload verdict to S3 for ClickHouse ingestion
+        if: always() && steps.claude.outputs.structured_output != ''
+        env:
+          VERDICT_JSON: ${{ steps.claude.outputs.structured_output }}
+          SIGNAL_PATTERN: ${{ inputs.signal_pattern }}
+        run: |
+          # Build enriched verdict JSON with signal metadata for CH ingestion
+          jq -n \
+            --arg repo "${{ github.repository }}" \
+            --argjson run_id "${{ github.run_id }}" \
+            --argjson run_attempt "${{ github.run_attempt }}" \
+            --arg timestamp "$(date -u +%Y-%m-%dT%H:%M:%S.000)" \
+            --arg suspect_commit "${{ inputs.suspect_commit }}" \
+            --argjson pr_number "${{ inputs.pr_number }}" \
+            --arg signal_key "$(echo "$SIGNAL_PATTERN" | jq -r '.signal_key // ""')" \
+            --arg signal_source "$(echo "$SIGNAL_PATTERN" | jq -r '.signal_source // ""')" \
+            --arg workflow_name "$(echo "$SIGNAL_PATTERN" | jq -r '.workflow_name // ""')" \
+            --arg verdict "$(echo "$VERDICT_JSON" | jq -r '.verdict // ""')" \
+            --argjson confidence "$(echo "$VERDICT_JSON" | jq -r '.confidence // 0')" \
+            --arg summary "$(echo "$VERDICT_JSON" | jq -r '.summary // ""')" \
+            --arg causal_reasoning "$(echo "$VERDICT_JSON" | jq -r '.causal_reasoning // ""')" \
+            '{
+              repo: $repo,
+              run_id: $run_id,
+              run_attempt: $run_attempt,
+              timestamp: $timestamp,
+              suspect_commit: $suspect_commit,
+              pr_number: $pr_number,
+              signal_key: $signal_key,
+              signal_source: $signal_source,
+              workflow_name: $workflow_name,
+              verdict: $verdict,
+              confidence: $confidence,
+              summary: $summary,
+              causal_reasoning: $causal_reasoning
+            }' > /tmp/advisor_verdict.json
+
+          aws s3 cp /tmp/advisor_verdict.json \
+            "s3://ossci-raw-job-status/autorevert_advisor_verdicts/${{ github.repository }}/${{ github.run_id }}_${{ github.run_attempt }}.json"
+
+      - name: Upload usage metrics
+        if: always()
+        uses: pytorch/test-infra/.github/actions/upload-claude-usage@release/2.12
diff --git a/.github/workflows/claude-code.yml b/.github/workflows/claude-code.yml
index 6ba4f6c26165e..e0a565ddfd3f9 100644
--- a/.github/workflows/claude-code.yml
+++ b/.github/workflows/claude-code.yml
@@ -3,76 +3,20 @@ name: Claude Code
 on:
   issue_comment:
     types: [created]
-  pull_request_review_comment:
-    types: [created]
   issues:
     types: [opened]
 
 jobs:
   claude-code:
-    # Early exit conditions:
-    # 1. Must be pytorch org
-    # 2. Must be triggered by pilot user
-    # 3. Must mention @claude
-    if: |
-      github.repository_owner == 'pytorch' &&
-      contains(fromJSON('[
-        "huydhn",
-        "seemethere",
-        "malfet",
-        "ZainRizvi",
-        "jeanschmidt",
-        "atalman",
-        "wdvr",
-        "izaitsevfb",
-        "yangw-dev",
-        "ezyang",
-        "drisspg",
-        "albanD",
-        "eellison",
-        "pytorch-auto-revert[bot]",
-        "janeyx99",
-        "SherlockNoMad",
-        "svekars",
-        "sekyondaMeta",
-        "AlannaBurke",
-        "ngimel"
-      ]'), github.actor) &&
-      (
-        (github.event_name == 'issue_comment' && contains(github.event.comment.body, '@claude')) ||
-        (github.event_name == 'pull_request_review_comment' && contains(github.event.comment.body, '@claude')) ||
-        (github.event_name == 'issues' && contains(github.event.issue.body, '@claude'))
-      )
-    runs-on: ubuntu-latest
-    timeout-minutes: 60
-    environment: bedrock
+    uses: pytorch/test-infra/.github/workflows/_claude-code.yml@release/2.12
     permissions:
       contents: read
       pull-requests: write
       issues: write
       id-token: write
-    steps:
-      # Fork PR support enabled by using izaitsevfb/claude-code-action@forked-pr-fix
-
-      - uses: actions/checkout@v4
-        with:
-          fetch-depth: 1
-
-      - name: Configure AWS credentials via OIDC
-        uses: aws-actions/configure-aws-credentials@v4
-        with:
-          role-to-assume: arn:aws:iam::308535385114:role/gha_workflow_claude_code
-          aws-region: us-east-1
-
-      - name: Run Claude Code
-        uses: izaitsevfb/claude-code-action@forked-pr-fix
-        with:
-          # We filter by github.actor at workflow level, there is no point of filtering here as well
-          allowed_bots: "*"
-          claude_args: "--model global.anthropic.claude-opus-4-5-20251101-v1:0"
-          settings: '{"alwaysThinkingEnabled": true}'
-          use_bedrock: "true"
-
-      - name: Upload usage metrics
-        if: always()
-        uses: pytorch/test-infra/.github/actions/upload-claude-usage@main
+    secrets: inherit
+    with:
+      additional_claude_args: '--allowedTools Skill'
+      append_system_prompt: |
+        When asked to review a PR, always use the /pr-review skill first.
+        It contains PyTorch-specific review guidelines, output format, and critical checks.
diff --git a/.github/workflows/claude-issue-triage-run.yml b/.github/workflows/claude-issue-triage-run.yml
index 6d63695b97ecf..7751584b14322 100644
--- a/.github/workflows/claude-issue-triage-run.yml
+++ b/.github/workflows/claude-issue-triage-run.yml
@@ -60,22 +60,24 @@ jobs:
                   "ghcr.io/github/github-mcp-server:v0.30.1"
                 ],
                 "env": {
-                  "GITHUB_PERSONAL_ACCESS_TOKEN": "${{ secrets.GITHUB_TOKEN }}"
+                  "GITHUB_PERSONAL_ACCESS_TOKEN": "${GITHUB_TOKEN}"
                 }
               }
             }
           }
           EOF
 
+        env:
+          GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
       - name: Configure AWS credentials via OIDC
-        uses: aws-actions/configure-aws-credentials@v4
+        uses: aws-actions/configure-aws-credentials@7474bc4690e29a8392af63c5b98e7449536d5c3a # v4
         with:
           role-to-assume: arn:aws:iam::308535385114:role/gha_workflow_claude_code
           aws-region: us-east-1
 
       - name: Run Issue Triage
         timeout-minutes: 5
-        uses: anthropics/claude-code-action@v1
+        uses: anthropics/claude-code-action@6e2bd52842c65e914eba5c8badd17560bd26b5de # v1.0.89
         env:
           GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
           TRIAGE_HOOK_DEBUG_LOG: /tmp/triage_hooks.log
@@ -106,4 +108,27 @@ jobs:
           fi
 
       - name: Upload usage metrics
-        uses: pytorch/test-infra/.github/actions/upload-claude-usage@main
+        uses: pytorch/test-infra/.github/actions/upload-claude-usage@release/2.12
+
+      - name: Upload execution output to S3
+        if: always()
+        continue-on-error: true
+        run: |
+          S3_BUCKET="ossci-raw-job-status"
+          ISSUE_NUM="${{ steps.issue.outputs.number }}"
+          S3_KEY="review-logs/${ISSUE_NUM}.json"
+          OUTPUT_FILE="${RUNNER_TEMP}/claude-execution-output.json"
+          if [ -f "$OUTPUT_FILE" ]; then
+            COMBINED="${RUNNER_TEMP}/combined_output.json"
+            # Download existing log if present, then append new content
+            if aws s3 cp "s3://${S3_BUCKET}/${S3_KEY}" "$COMBINED" 2>/dev/null; then
+              echo "" >> "$COMBINED"
+              cat "$OUTPUT_FILE" >> "$COMBINED"
+            else
+              cp "$OUTPUT_FILE" "$COMBINED"
+            fi
+            aws s3 cp "$COMBINED" "s3://${S3_BUCKET}/${S3_KEY}"
+            echo "Uploaded execution output to s3://${S3_BUCKET}/${S3_KEY}"
+          else
+            echo "No execution output found at ${OUTPUT_FILE}"
+          fi
diff --git a/.github/workflows/close-nonexistent-disable-issues.yml b/.github/workflows/close-nonexistent-disable-issues.yml
index bef3d8797149c..fd3693178786c 100644
--- a/.github/workflows/close-nonexistent-disable-issues.yml
+++ b/.github/workflows/close-nonexistent-disable-issues.yml
@@ -13,7 +13,7 @@ jobs:
     runs-on: ubuntu-latest
     steps:
       - name: Checkout PyTorch
-        uses: pytorch/pytorch/.github/actions/checkout-pytorch@main
+        uses: pytorch/pytorch/.github/actions/checkout-pytorch@release/2.12
         with:
           submodules: false
           fetch-depth: 1
diff --git a/.github/workflows/create_release.yml b/.github/workflows/create_release.yml
index d5e0d96fe19f2..2058df9f15967 100644
--- a/.github/workflows/create_release.yml
+++ b/.github/workflows/create_release.yml
@@ -19,7 +19,7 @@ jobs:
   get-label-type:
     if: github.repository_owner == 'pytorch'
     name: get-label-type
-    uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@main
+    uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@release/2.12
     with:
       triggering_actor: ${{ github.triggering_actor }}
       issue_owner: ${{ github.event.pull_request.user.login || github.event.issue.user.login }}
@@ -54,11 +54,19 @@ jobs:
           tag_or_branch="${tag_or_branch#refs/heads/}"
           # replace directory separators with _ in branch name
           tag_or_branch="${tag_or_branch//\//_}"
+          # Set PYTORCH_BUILD_VERSION only for release/RC tags; convert to PEP 440
+          if [[ "$PT_GITHUB_REF" =~ ^refs/tags/v[0-9]+\.[0-9]+\.[0-9]+(-rc[0-9]+)?$ ]]; then
+            ver="${PT_GITHUB_REF#refs/tags/v}"
+            export PYTORCH_BUILD_VERSION="${ver/-rc/rc}"
+            export PYTORCH_BUILD_NUMBER=0
+          fi
           torch_version="$(python -c 'from tools.generate_torch_version import get_torch_version; print(get_torch_version())')"
           {
             echo "PT_RELEASE_NAME=pytorch-$tag_or_branch";
             echo "PT_RELEASE_FILE=pytorch-$tag_or_branch.tar.gz";
             echo "PT_PEP517_RELEASE_FILE=torch-${torch_version}.tar.gz";
+            echo "PYTORCH_BUILD_VERSION=${PYTORCH_BUILD_VERSION:-}";
+            echo "PYTORCH_BUILD_NUMBER=${PYTORCH_BUILD_NUMBER:-}";
           } >> "$GITHUB_ENV"
       - name: Checkout optional submodules
         run: python3 tools/optional_submodules.py
@@ -74,7 +82,7 @@ jobs:
           cp -r "$PWD" "/tmp/$PT_RELEASE_NAME"
           mv "/tmp/$PT_RELEASE_NAME" .
           # Cleanup
-          rm -rf "$PT_RELEASE_NAME"/{.circleci,.ci}
+          rm -rf "$PT_RELEASE_NAME"/.ci
           find "$PT_RELEASE_NAME" -name '.git*' -exec rm -rv {} \; || true
           # Create archive
           tar -czf "$PT_RELEASE_FILE" "$PT_RELEASE_NAME"
@@ -83,7 +91,7 @@ jobs:
         run: |
           pip install build==1.2.2.post1 || exit 1
           python -m build --sdist || exit 1
-          cd dist || exit 1
+          mv dist/$PT_PEP517_RELEASE_FILE . || exit 1
       - name: Upload source distribution for release
         if: ${{ github.event_name == 'release' }}
         uses: softprops/action-gh-release@da05d552573ad5aba039eaac05058a918a7bf631 # v2.2.2
@@ -102,7 +110,7 @@ jobs:
         uses: actions/upload-artifact@50769540e7f4bd5e21e526ee35c689e35e0d6874 # v4.4.0
         with:
           name: ${{ env.PT_PEP517_RELEASE_FILE }}
-          path: dist/${{ env.PT_PEP517_RELEASE_FILE }}
+          path: ${{ env.PT_PEP517_RELEASE_FILE }}
       - name: Set output
         id: release_name
         run: |
diff --git a/.github/workflows/docker-builds.yml b/.github/workflows/docker-builds.yml
index c4c989af980e6..bfe3489526445 100644
--- a/.github/workflows/docker-builds.yml
+++ b/.github/workflows/docker-builds.yml
@@ -33,7 +33,7 @@ jobs:
   get-label-type:
     if: github.repository_owner == 'pytorch'
     name: get-label-type
-    uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@main
+    uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@release/2.12
     with:
       triggering_actor: ${{ github.triggering_actor }}
       issue_owner: ${{ github.event.pull_request.user.login || github.event.issue.user.login }}
@@ -51,22 +51,19 @@ jobs:
         docker-image-name: [
           pytorch-linux-jammy-cuda12.8-cudnn9-py3-gcc11,
           pytorch-linux-jammy-cuda13.0-cudnn9-py3-gcc11,
-          pytorch-linux-jammy-cuda12.9-cudnn9-py3.12-gcc11-vllm,
-          pytorch-linux-jammy-cuda12.8-cudnn9-py3-gcc11-inductor-benchmarks,
+          pytorch-linux-jammy-cuda13.0-cudnn9-py3.12-gcc11-vllm,
           pytorch-linux-jammy-cuda13.0-cudnn9-py3-gcc11-inductor-benchmarks,
-          pytorch-linux-jammy-cuda12.4-cudnn9-py3-gcc11,
-          pytorch-linux-jammy-py3.10-clang15,
-          pytorch-linux-jammy-py3.11-clang15,
-          pytorch-linux-jammy-py3.12-clang15,
-          pytorch-linux-jammy-py3.13-clang15,
-          pytorch-linux-jammy-py3.14-clang15,
-          pytorch-linux-jammy-py3.14t-clang15,
+          pytorch-linux-jammy-py3.10-clang18,
+          pytorch-linux-jammy-py3.11-clang18,
+          pytorch-linux-jammy-py3.12-clang18,
+          pytorch-linux-jammy-py3.13-clang18,
+          pytorch-linux-jammy-py3.14-clang18,
+          pytorch-linux-jammy-py3.14t-clang18,
           pytorch-linux-jammy-rocm-n-py3,
           pytorch-linux-noble-rocm-n-py3,
           pytorch-linux-noble-rocm-nightly-py3,
           pytorch-linux-jammy-rocm-n-py3-benchmarks,
-          pytorch-linux-jammy-cuda12.8-cudnn9-py3.10-clang15,
-          pytorch-linux-jammy-py3.10-gcc11,
+          pytorch-linux-jammy-cuda12.8-cudnn9-py3.10-clang18,
           pytorch-linux-jammy-py3-gcc11-inductor-benchmarks,
           pytorch-linux-jammy-py3.12-halide,
           pytorch-linux-jammy-py3.12-pallas,
@@ -74,21 +71,18 @@ jobs:
           pytorch-linux-jammy-tpu-py3.12-pallas,
           pytorch-linux-jammy-xpu-n-1-py3,
           pytorch-linux-noble-xpu-n-py3,
+          pytorch-linux-noble-xpu-n-py3-client,
           pytorch-linux-noble-xpu-n-py3-inductor-benchmarks,
-          pytorch-linux-jammy-py3-clang18-asan,
-          pytorch-linux-jammy-py3-clang15-onnx,
           pytorch-linux-jammy-linter,
           pytorch-linux-jammy-cuda12.8-cudnn9-py3.10-linter,
           # TODO: Re-enable me when docker pin update happens
-          # pytorch-linux-jammy-py3-clang15-executorch,
+          # pytorch-linux-jammy-py3-clang18-executorch,
           pytorch-linux-jammy-py3.12-triton-cpu,
           pytorch-linux-noble-riscv64-py3.12-gcc14
         ]
         include:
           - docker-image-name: pytorch-linux-jammy-aarch64-py3.10-gcc13
             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
             runner: linux.arm64.m7g.4xlarge
             timeout-minutes: 600
@@ -103,27 +97,22 @@ jobs:
           sudo rm -rf "${GITHUB_WORKSPACE}"
           mkdir "${GITHUB_WORKSPACE}"
 
-      # [see note: pytorch repo ref]
-      # deep clone (fetch-depth 0) required for git merge-base
-      - name: Checkout PyTorch
-        uses: pytorch/pytorch/.github/actions/checkout-pytorch@main
-
       - name: Setup Linux
-        uses: ./.github/actions/setup-linux
+        uses: pytorch/pytorch/.github/actions/setup-linux@release/2.12
 
       - name: Login to ECR
         uses: ./.github/actions/ecr-login
 
       - name: Build docker image
         id: build-docker-image
-        uses: pytorch/test-infra/.github/actions/calculate-docker-image@main
+        uses: pytorch/test-infra/.github/actions/calculate-docker-image@release/2.12
         with:
           docker-image-name: ci-image:${{ matrix.docker-image-name }}
           always-rebuild: true
           push: true
 
       - name: Pull docker image
-        uses: pytorch/test-infra/.github/actions/pull-docker-image@main
+        uses: pytorch/test-infra/.github/actions/pull-docker-image@release/2.12
         with:
           docker-image: ${{ steps.build-docker-image.outputs.docker-image }}
 
@@ -162,6 +151,9 @@ jobs:
             echo $GHCR_PAT | docker login ghcr.io -u pytorch --password-stdin
             docker tag "${ECR_DOCKER_IMAGE}" "${ghcr_image}:${tag}"
             docker push "${ghcr_image}:${tag}"
+            # Also push a tag without the hash for easier reference
+            docker tag "${ECR_DOCKER_IMAGE}" "${ghcr_image}:${{ matrix.docker-image-name }}"
+            docker push "${ghcr_image}:${{ matrix.docker-image-name }}"
 
       - name: Chown workspace
         uses: ./.github/actions/chown-workspace
@@ -170,5 +162,5 @@ jobs:
         if: always()
 
       - name: Teardown Linux
-        uses: pytorch/test-infra/.github/actions/teardown-linux@main
+        uses: pytorch/test-infra/.github/actions/teardown-linux@release/2.12
         if: always()
diff --git a/.github/workflows/docker-cache-rocm.yml b/.github/workflows/docker-cache-rocm.yml
index 2c2a9ba16647f..067a953790923 100644
--- a/.github/workflows/docker-cache-rocm.yml
+++ b/.github/workflows/docker-cache-rocm.yml
@@ -57,7 +57,7 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        runner: [linux.rocm.mi250.docker-cache]
+        runner: [linux.rocm.mi250.docker-cache, linux.rocm.mi210.docker-cache]
         docker-image: [
           "${{ needs.download-docker-builds-artifacts.outputs.pytorch-linux-jammy-rocm-n-py3 }}",
           "${{ needs.download-docker-builds-artifacts.outputs.pytorch-linux-noble-rocm-n-py3 }}"
@@ -71,7 +71,7 @@ jobs:
           echo "Outputs of download-docker-builds-artifacts job: ${JSON_STRINGIFIED}"
 
       - name: Checkout PyTorch
-        uses: pytorch/pytorch/.github/actions/checkout-pytorch@main
+        uses: pytorch/pytorch/.github/actions/checkout-pytorch@release/2.12
         with:
           no-sudo: true
 
@@ -86,7 +86,7 @@ jobs:
             echo "ghcr_image=${ghcr_image}" >> "$GITHUB_OUTPUT"
 
       - name: Pull docker image
-        uses: pytorch/test-infra/.github/actions/pull-docker-image@main
+        uses: pytorch/test-infra/.github/actions/pull-docker-image@release/2.12
         with:
           docker-image: ${{ steps.ghcr-io-tag.outputs.ghcr_image }}
 
@@ -95,7 +95,7 @@ jobs:
           docker_image_tag=${{ matrix.docker-image }}
           docker_image_tag="${docker_image_tag#*:}" # Remove everything before and including first ":"
           docker_image_tag="${docker_image_tag%-*}" # Remove everything after and including last "-"
-          ref_name=${{ github.event.workflow_run.head_branch || github.event.inputs.branch }}
+          ref_name=${WORKFLOW_RUN_HEAD_BRANCH}
           if [[ $ref_name =~ "release/" ]]; then
             ref_suffix="release"
           elif [[ $ref_name == "main" ]]; then
@@ -107,3 +107,6 @@ jobs:
           # mv is atomic operation, so we use intermediate tar.tmp file to prevent read-write contention
           docker save -o ~/pytorch-data/docker/${docker_image_tag}.tar.tmp ${{ matrix.docker-image }}
           mv ~/pytorch-data/docker/${docker_image_tag}.tar.tmp ~/pytorch-data/docker/${docker_image_tag}_${ref_suffix}.tar
+
+        env:
+          WORKFLOW_RUN_HEAD_BRANCH: ${{ github.event.workflow_run.head_branch || github.event.inputs.branch }}
diff --git a/.github/workflows/docker-release.yml b/.github/workflows/docker-release.yml
index 577a8acb5203f..537eb7f4cf91d 100644
--- a/.github/workflows/docker-release.yml
+++ b/.github/workflows/docker-release.yml
@@ -38,7 +38,7 @@ jobs:
   get-label-type:
     if: github.repository_owner == 'pytorch'
     name: get-label-type
-    uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@main
+    uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@release/2.12
     with:
       triggering_actor: ${{ github.triggering_actor }}
       issue_owner: ${{ github.event.pull_request.user.login || github.event.issue.user.login }}
@@ -53,7 +53,7 @@ jobs:
       matrix: ${{ steps.generate-matrix.outputs.matrix }}
     steps:
       - name: Checkout PyTorch
-        uses: pytorch/pytorch/.github/actions/checkout-pytorch@main
+        uses: pytorch/pytorch/.github/actions/checkout-pytorch@release/2.12
         with:
           fetch-depth: 1
           submodules: true
@@ -83,18 +83,12 @@ jobs:
       CUDNN_VERSION: ${{ matrix.cudnn_version }}
     steps:
       - name: Setup SSH (Click me for login details)
-        uses: pytorch/test-infra/.github/actions/setup-ssh@main
+        uses: pytorch/test-infra/.github/actions/setup-ssh@release/2.12
         with:
           github-secret: ${{ secrets.GITHUB_TOKEN }}
-      # [see note: pytorch repo ref]
-      # deep clone (fetch-depth 0) required for git merge-base
-      - name: Checkout PyTorch
-        uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
-        with:
-          fetch-depth: 0
-          submodules: 'recursive'
+
       - name: Setup Linux
-        uses: ./.github/actions/setup-linux
+        uses: pytorch/pytorch/.github/actions/setup-linux@release/2.12
 
       - name: Login to ECR
         uses: ./.github/actions/ecr-login
@@ -137,7 +131,7 @@ jobs:
             } >> "${GITHUB_ENV}"
           fi
       - name: Setup nightly specific variables
-        if: ${{ github.event.ref == 'refs/heads/nightly' || startsWith(github.event.ref, 'refs/tags/ciflow/nightly/') }}
+        if: ${{ github.event_name == 'pull_request' || github.event.ref == 'refs/heads/nightly' || startsWith(github.event.ref, 'refs/tags/ciflow/nightly/') }}
         run: |
           {
             echo "DOCKER_IMAGE=pytorch-nightly";
@@ -169,13 +163,13 @@ jobs:
           fi
 
       - name: Teardown Linux
-        uses: pytorch/test-infra/.github/actions/teardown-linux@main
+        uses: pytorch/test-infra/.github/actions/teardown-linux@release/2.12
         if: always()
 
   validate:
     needs: build
     if: ${{ github.event_name == 'push' && (github.event.ref == 'refs/heads/nightly' || startsWith(github.event.ref, 'refs/tags/v')) }}
-    uses: pytorch/test-infra/.github/workflows/validate-docker-images.yml@main
+    uses: pytorch/test-infra/.github/workflows/validate-docker-images.yml@release/2.12
     with:
       channel: nightly
       ref: main
diff --git a/.github/workflows/dtensor.yml b/.github/workflows/dtensor.yml
new file mode 100644
index 0000000000000..18b1f6c565d72
--- /dev/null
+++ b/.github/workflows/dtensor.yml
@@ -0,0 +1,60 @@
+# Workflow: DTensor Unit Test
+# Runs unit tests for DTensor (test/distributed/tensor/).
+name: dtensor
+
+on:
+  push:
+    tags:
+      - ciflow/dtensor/*
+  workflow_call:
+  pull_request:
+    paths:
+      - .github/workflows/dtensor.yml
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}-${{ github.event_name == 'workflow_dispatch' }}
+  cancel-in-progress: true
+
+permissions:
+  id-token: write
+  contents: read
+  actions: read
+
+jobs:
+  get-label-type:
+    name: get-label-type
+    uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@release/2.12
+    if: github.repository_owner == 'pytorch'
+    with:
+      triggering_actor: ${{ github.triggering_actor }}
+      issue_owner: ${{ github.event.pull_request.user.login || github.event.issue.user.login }}
+      curr_branch: ${{ github.head_ref || github.ref_name }}
+      curr_ref_type: ${{ github.ref_type }}
+
+  dtensor-build:
+    name: dtensor-build
+    uses: ./.github/workflows/_linux-build.yml
+    needs: get-label-type
+    with:
+      runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
+      build-environment: linux-jammy-cuda13.0-py3.10-gcc11
+      docker-image-name: ci-image:pytorch-linux-jammy-cuda13.0-cudnn9-py3-gcc11
+      cuda-arch-list: '7.5 8.6 8.9'
+      test-matrix: |
+        { include: [
+          { config: "dtensor", shard: 1, num_shards: 1, runner: "${{ needs.get-label-type.outputs.label-type }}linux.g5.12xlarge.nvidia.gpu" },
+        ]}
+    secrets: inherit
+
+  dtensor-test:
+    name: dtensor-test
+    uses: ./.github/workflows/_linux-test.yml
+    needs: [get-label-type, dtensor-build]
+    with:
+      build-environment: linux-jammy-cuda13.0-py3.10-gcc11
+      docker-image: ci-image:pytorch-linux-jammy-cuda13.0-cudnn9-py3-gcc11
+      test-matrix: |
+        { include: [
+          { config: "dtensor", shard: 1, num_shards: 1, runner: "${{ needs.get-label-type.outputs.label-type }}linux.g5.12xlarge.nvidia.gpu" },
+        ]}
+    secrets: inherit
diff --git a/.github/workflows/dynamo-unittest.yml b/.github/workflows/dynamo-unittest.yml
index f7eea350b5644..d822926765d1c 100644
--- a/.github/workflows/dynamo-unittest.yml
+++ b/.github/workflows/dynamo-unittest.yml
@@ -17,11 +17,12 @@ concurrency:
 permissions:
   id-token: write
   contents: read
+  actions: read
 
 jobs:
   get-label-type:
     name: get-label-type
-    uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@main
+    uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@release/2.12
     if: ${{ (github.event_name != 'schedule' || github.repository == 'pytorch/pytorch') && github.repository_owner == 'pytorch' }}
     with:
       triggering_actor: ${{ github.triggering_actor }}
@@ -39,8 +40,8 @@ jobs:
         python-version: ['3.11', '3.12', '3.13']
     with:
       runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
-      build-environment: linux-jammy-py${{ matrix.python-version }}-clang15
-      docker-image-name: ci-image:pytorch-linux-jammy-py${{ matrix.python-version }}-clang15
+      build-environment: linux-jammy-py${{ matrix.python-version }}-clang18
+      docker-image-name: ci-image:pytorch-linux-jammy-py${{ matrix.python-version }}-clang18
       test-matrix: |
         { include: [
           { config: "dynamo_core", shard: 1, num_shards: 1, runner: "${{ needs.get-label-type.outputs.label-type }}linux.c7i.2xlarge" },
@@ -58,8 +59,8 @@ jobs:
       matrix:
         python-version: ['3.11', '3.12', '3.13']
     with:
-      build-environment: linux-jammy-py${{ matrix.python-version }}-clang15
-      docker-image: ci-image:pytorch-linux-jammy-py${{ matrix.python-version }}-clang15
+      build-environment: linux-jammy-py${{ matrix.python-version }}-clang18
+      docker-image: ci-image:pytorch-linux-jammy-py${{ matrix.python-version }}-clang18
       test-matrix: |
         { include: [
           { config: "dynamo_core", shard: 1, num_shards: 1, runner: "${{ needs.get-label-type.outputs.label-type }}linux.c7i.2xlarge" },
diff --git a/.github/workflows/generated-linux-aarch64-binary-manywheel-nightly.yml b/.github/workflows/generated-linux-aarch64-binary-manywheel-nightly.yml
index e4a8e5e96a88a..a1d1e3371562f 100644
--- a/.github/workflows/generated-linux-aarch64-binary-manywheel-nightly.yml
+++ b/.github/workflows/generated-linux-aarch64-binary-manywheel-nightly.yml
@@ -41,7 +41,7 @@ jobs:
   get-label-type:
     if: github.repository_owner == 'pytorch'
     name: get-label-type
-    uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@main
+    uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@release/2.12
     with:
       triggering_actor: ${{ github.triggering_actor }}
       issue_owner: ${{ github.event.pull_request.user.login || github.event.issue.user.login }}
@@ -111,6 +111,9 @@ jobs:
       build_name: manywheel-py3_10-cpu-aarch64
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
     uses: ./.github/workflows/_binary-upload.yml
 
   manywheel-py3_10-cuda-aarch64-12_6-build:
@@ -133,7 +136,7 @@ jobs:
       ALPINE_IMAGE: "arm64v8/alpine"
       build_name: manywheel-py3_10-cuda-aarch64-12_6
       build_environment: linux-aarch64-binary-manywheel
-      PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cublas,cufile,nvjitlink,nvtx]==12.6.3; platform_system == 'Linux' | cuda-bindings==12.9.4; platform_system == 'Linux' | nvidia-cudnn-cu12==9.10.2.21; platform_system == 'Linux' | nvidia-cusparselt-cu12==0.7.1; platform_system == 'Linux' | nvidia-nccl-cu12==2.28.9; platform_system == 'Linux' | nvidia-nvshmem-cu12==3.4.5; platform_system == 'Linux'
+      PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cublas,cufile,nvjitlink,nvtx]==12.6.3; platform_system == 'Linux' | cuda-bindings>=12.9.4,<13; platform_system == 'Linux' | nvidia-cudnn-cu12==9.10.2.21; platform_system == 'Linux' | nvidia-cusparselt-cu12==0.7.1; platform_system == 'Linux' | nvidia-nccl-cu12==2.29.3; platform_system == 'Linux' | nvidia-nvshmem-cu12==3.4.5; platform_system == 'Linux'
       timeout-minutes: 420
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
@@ -182,80 +185,12 @@ jobs:
       build_name: manywheel-py3_10-cuda-aarch64-12_6
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
     uses: ./.github/workflows/_binary-upload.yml
 
-  manywheel-py3_10-cuda-aarch64-12_8-build:
-    if: ${{ github.repository_owner == 'pytorch' }}
-    uses: ./.github/workflows/_binary-build-linux.yml
-    needs: get-label-type
-    with:
-      PYTORCH_ROOT: /pytorch
-      PACKAGE_TYPE: manywheel
-      # TODO: This is a legacy variable that we eventually want to get rid of in
-      #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu128
-      GPU_ARCH_VERSION: "12.8-aarch64"
-      GPU_ARCH_TYPE: cuda-aarch64
-      DOCKER_IMAGE: manylinuxaarch64-builder
-      DOCKER_IMAGE_TAG_PREFIX: cuda12.8
-      DESIRED_PYTHON: "3.10"
-      runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
-      runs_on: linux.arm64.r7g.12xlarge.memory
-      ALPINE_IMAGE: "arm64v8/alpine"
-      build_name: manywheel-py3_10-cuda-aarch64-12_8
-      build_environment: linux-aarch64-binary-manywheel
-      PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cublas,cufile,nvjitlink,nvtx]==12.8.1; platform_system == 'Linux' | cuda-bindings==12.9.4; platform_system == 'Linux' | nvidia-cudnn-cu12==9.17.1.4; platform_system == 'Linux' | nvidia-cusparselt-cu12==0.7.1; platform_system == 'Linux' | nvidia-nccl-cu12==2.28.9; platform_system == 'Linux' | nvidia-nvshmem-cu12==3.4.5; platform_system == 'Linux'
-      timeout-minutes: 420
-    secrets:
-      github-token: ${{ secrets.GITHUB_TOKEN }}
-
-  manywheel-py3_10-cuda-aarch64-12_8-test:  # Testing
-    if: ${{ github.repository_owner == 'pytorch' }}
-    needs:
-      - manywheel-py3_10-cuda-aarch64-12_8-build
-      - get-label-type
-    uses: ./.github/workflows/_binary-test-linux.yml
-    with:
-      PYTORCH_ROOT: /pytorch
-      PACKAGE_TYPE: manywheel
-      # TODO: This is a legacy variable that we eventually want to get rid of in
-      #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu128
-      GPU_ARCH_VERSION: "12.8-aarch64"
-      GPU_ARCH_TYPE: cuda-aarch64
-      DOCKER_IMAGE: manylinuxaarch64-builder
-      DOCKER_IMAGE_TAG_PREFIX: cuda12.8
-      DESIRED_PYTHON: "3.10"
-      build_name: manywheel-py3_10-cuda-aarch64-12_8
-      build_environment: linux-aarch64-binary-manywheel
-      runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
-      runs_on: linux.arm64.2xlarge
-      ALPINE_IMAGE: "arm64v8/alpine"
-    secrets:
-      github-token: ${{ secrets.GITHUB_TOKEN }}
-  manywheel-py3_10-cuda-aarch64-12_8-upload:  # Uploading
-    if: ${{ github.repository_owner == 'pytorch' }}
-    permissions:
-      id-token: write
-      contents: read
-    needs: manywheel-py3_10-cuda-aarch64-12_8-build
-    with:
-      PYTORCH_ROOT: /pytorch
-      PACKAGE_TYPE: manywheel
-      # TODO: This is a legacy variable that we eventually want to get rid of in
-      #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu128
-      GPU_ARCH_VERSION: "12.8-aarch64"
-      GPU_ARCH_TYPE: cuda-aarch64
-      DOCKER_IMAGE: manylinuxaarch64-builder
-      DOCKER_IMAGE_TAG_PREFIX: cuda12.8
-      DESIRED_PYTHON: "3.10"
-      build_name: manywheel-py3_10-cuda-aarch64-12_8
-    secrets:
-      github-token: ${{ secrets.GITHUB_TOKEN }}
-    uses: ./.github/workflows/_binary-upload.yml
-
-  manywheel-py3_10-cuda-aarch64-12_9-build:
+  manywheel-py3_10-cuda-aarch64-13_0-build:
     if: ${{ github.repository_owner == 'pytorch' }}
     uses: ./.github/workflows/_binary-build-linux.yml
     needs: get-label-type
@@ -264,26 +199,26 @@ jobs:
       PACKAGE_TYPE: manywheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu129
-      GPU_ARCH_VERSION: "12.9-aarch64"
+      DESIRED_CUDA: cu130
+      GPU_ARCH_VERSION: "13.0-aarch64"
       GPU_ARCH_TYPE: cuda-aarch64
       DOCKER_IMAGE: manylinuxaarch64-builder
-      DOCKER_IMAGE_TAG_PREFIX: cuda12.9
+      DOCKER_IMAGE_TAG_PREFIX: cuda13.0
       DESIRED_PYTHON: "3.10"
       runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
       runs_on: linux.arm64.r7g.12xlarge.memory
       ALPINE_IMAGE: "arm64v8/alpine"
-      build_name: manywheel-py3_10-cuda-aarch64-12_9
+      build_name: manywheel-py3_10-cuda-aarch64-13_0
       build_environment: linux-aarch64-binary-manywheel
-      PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cublas,cufile,nvjitlink,nvtx]==12.9.1; platform_system == 'Linux' | cuda-bindings==12.9.4; platform_system == 'Linux' | nvidia-cudnn-cu12==9.17.1.4; platform_system == 'Linux' | nvidia-cusparselt-cu12==0.7.1; platform_system == 'Linux' | nvidia-nccl-cu12==2.28.9; platform_system == 'Linux' | nvidia-nvshmem-cu12==3.4.5; platform_system == 'Linux'
+      PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cufile,nvjitlink,nvtx]==13.0.2; platform_system == 'Linux' | nvidia-cublas>=13.1.0.3,<=13.1.1.3; platform_system == 'Linux' | cuda-bindings>=13.0.3,<14; platform_system == 'Linux' | nvidia-cudnn-cu13==9.20.0.48; platform_system == 'Linux' | nvidia-cusparselt-cu13==0.8.1; platform_system == 'Linux' | nvidia-nccl-cu13==2.29.7; platform_system == 'Linux' | nvidia-nvshmem-cu13==3.4.5; platform_system == 'Linux'
       timeout-minutes: 420
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
 
-  manywheel-py3_10-cuda-aarch64-12_9-test:  # Testing
+  manywheel-py3_10-cuda-aarch64-13_0-test:  # Testing
     if: ${{ github.repository_owner == 'pytorch' }}
     needs:
-      - manywheel-py3_10-cuda-aarch64-12_9-build
+      - manywheel-py3_10-cuda-aarch64-13_0-build
       - get-label-type
     uses: ./.github/workflows/_binary-test-linux.yml
     with:
@@ -291,42 +226,45 @@ jobs:
       PACKAGE_TYPE: manywheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu129
-      GPU_ARCH_VERSION: "12.9-aarch64"
+      DESIRED_CUDA: cu130
+      GPU_ARCH_VERSION: "13.0-aarch64"
       GPU_ARCH_TYPE: cuda-aarch64
       DOCKER_IMAGE: manylinuxaarch64-builder
-      DOCKER_IMAGE_TAG_PREFIX: cuda12.9
+      DOCKER_IMAGE_TAG_PREFIX: cuda13.0
       DESIRED_PYTHON: "3.10"
-      build_name: manywheel-py3_10-cuda-aarch64-12_9
+      build_name: manywheel-py3_10-cuda-aarch64-13_0
       build_environment: linux-aarch64-binary-manywheel
       runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
       runs_on: linux.arm64.2xlarge
       ALPINE_IMAGE: "arm64v8/alpine"
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
-  manywheel-py3_10-cuda-aarch64-12_9-upload:  # Uploading
+  manywheel-py3_10-cuda-aarch64-13_0-upload:  # Uploading
     if: ${{ github.repository_owner == 'pytorch' }}
     permissions:
       id-token: write
       contents: read
-    needs: manywheel-py3_10-cuda-aarch64-12_9-build
+    needs: manywheel-py3_10-cuda-aarch64-13_0-build
     with:
       PYTORCH_ROOT: /pytorch
       PACKAGE_TYPE: manywheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu129
-      GPU_ARCH_VERSION: "12.9-aarch64"
+      DESIRED_CUDA: cu130
+      GPU_ARCH_VERSION: "13.0-aarch64"
       GPU_ARCH_TYPE: cuda-aarch64
       DOCKER_IMAGE: manylinuxaarch64-builder
-      DOCKER_IMAGE_TAG_PREFIX: cuda12.9
+      DOCKER_IMAGE_TAG_PREFIX: cuda13.0
       DESIRED_PYTHON: "3.10"
-      build_name: manywheel-py3_10-cuda-aarch64-12_9
+      build_name: manywheel-py3_10-cuda-aarch64-13_0
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
     uses: ./.github/workflows/_binary-upload.yml
 
-  manywheel-py3_10-cuda-aarch64-13_0-build:
+  manywheel-py3_10-cuda-aarch64-13_2-build:
     if: ${{ github.repository_owner == 'pytorch' }}
     uses: ./.github/workflows/_binary-build-linux.yml
     needs: get-label-type
@@ -335,26 +273,26 @@ jobs:
       PACKAGE_TYPE: manywheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu130
-      GPU_ARCH_VERSION: "13.0-aarch64"
+      DESIRED_CUDA: cu132
+      GPU_ARCH_VERSION: "13.2-aarch64"
       GPU_ARCH_TYPE: cuda-aarch64
       DOCKER_IMAGE: manylinuxaarch64-builder
-      DOCKER_IMAGE_TAG_PREFIX: cuda13.0
+      DOCKER_IMAGE_TAG_PREFIX: cuda13.2
       DESIRED_PYTHON: "3.10"
       runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
       runs_on: linux.arm64.r7g.12xlarge.memory
       ALPINE_IMAGE: "arm64v8/alpine"
-      build_name: manywheel-py3_10-cuda-aarch64-13_0
+      build_name: manywheel-py3_10-cuda-aarch64-13_2
       build_environment: linux-aarch64-binary-manywheel
-      PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cublas,cufile,nvjitlink,nvtx]==13.0.2; platform_system == 'Linux' | cuda-bindings==13.0.3; platform_system == 'Linux' | nvidia-cudnn-cu13==9.17.1.4; platform_system == 'Linux' | nvidia-cusparselt-cu13==0.8.0; platform_system == 'Linux' | nvidia-nccl-cu13==2.28.9; platform_system == 'Linux' | nvidia-nvshmem-cu13==3.4.5; platform_system == 'Linux'
+      PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cublas,cufile,nvjitlink,nvtx]==13.2.1; platform_system == 'Linux' | cuda-bindings>=13.0.3,<14; platform_system == 'Linux' | nvidia-cudnn-cu13==9.20.0.48; platform_system == 'Linux' | nvidia-cusparselt-cu13==0.8.1; platform_system == 'Linux' | nvidia-nccl-cu13==2.29.7; platform_system == 'Linux' | nvidia-nvshmem-cu13==3.4.5; platform_system == 'Linux'
       timeout-minutes: 420
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
 
-  manywheel-py3_10-cuda-aarch64-13_0-test:  # Testing
+  manywheel-py3_10-cuda-aarch64-13_2-test:  # Testing
     if: ${{ github.repository_owner == 'pytorch' }}
     needs:
-      - manywheel-py3_10-cuda-aarch64-13_0-build
+      - manywheel-py3_10-cuda-aarch64-13_2-build
       - get-label-type
     uses: ./.github/workflows/_binary-test-linux.yml
     with:
@@ -362,39 +300,42 @@ jobs:
       PACKAGE_TYPE: manywheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu130
-      GPU_ARCH_VERSION: "13.0-aarch64"
+      DESIRED_CUDA: cu132
+      GPU_ARCH_VERSION: "13.2-aarch64"
       GPU_ARCH_TYPE: cuda-aarch64
       DOCKER_IMAGE: manylinuxaarch64-builder
-      DOCKER_IMAGE_TAG_PREFIX: cuda13.0
+      DOCKER_IMAGE_TAG_PREFIX: cuda13.2
       DESIRED_PYTHON: "3.10"
-      build_name: manywheel-py3_10-cuda-aarch64-13_0
+      build_name: manywheel-py3_10-cuda-aarch64-13_2
       build_environment: linux-aarch64-binary-manywheel
       runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
       runs_on: linux.arm64.2xlarge
       ALPINE_IMAGE: "arm64v8/alpine"
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
-  manywheel-py3_10-cuda-aarch64-13_0-upload:  # Uploading
+  manywheel-py3_10-cuda-aarch64-13_2-upload:  # Uploading
     if: ${{ github.repository_owner == 'pytorch' }}
     permissions:
       id-token: write
       contents: read
-    needs: manywheel-py3_10-cuda-aarch64-13_0-build
+    needs: manywheel-py3_10-cuda-aarch64-13_2-build
     with:
       PYTORCH_ROOT: /pytorch
       PACKAGE_TYPE: manywheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu130
-      GPU_ARCH_VERSION: "13.0-aarch64"
+      DESIRED_CUDA: cu132
+      GPU_ARCH_VERSION: "13.2-aarch64"
       GPU_ARCH_TYPE: cuda-aarch64
       DOCKER_IMAGE: manylinuxaarch64-builder
-      DOCKER_IMAGE_TAG_PREFIX: cuda13.0
+      DOCKER_IMAGE_TAG_PREFIX: cuda13.2
       DESIRED_PYTHON: "3.10"
-      build_name: manywheel-py3_10-cuda-aarch64-13_0
+      build_name: manywheel-py3_10-cuda-aarch64-13_2
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
     uses: ./.github/workflows/_binary-upload.yml
 
   manywheel-py3_11-cpu-aarch64-build:
@@ -461,6 +402,9 @@ jobs:
       build_name: manywheel-py3_11-cpu-aarch64
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
     uses: ./.github/workflows/_binary-upload.yml
 
   manywheel-py3_11-cuda-aarch64-12_6-build:
@@ -483,7 +427,7 @@ jobs:
       ALPINE_IMAGE: "arm64v8/alpine"
       build_name: manywheel-py3_11-cuda-aarch64-12_6
       build_environment: linux-aarch64-binary-manywheel
-      PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cublas,cufile,nvjitlink,nvtx]==12.6.3; platform_system == 'Linux' | cuda-bindings==12.9.4; platform_system == 'Linux' | nvidia-cudnn-cu12==9.10.2.21; platform_system == 'Linux' | nvidia-cusparselt-cu12==0.7.1; platform_system == 'Linux' | nvidia-nccl-cu12==2.28.9; platform_system == 'Linux' | nvidia-nvshmem-cu12==3.4.5; platform_system == 'Linux'
+      PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cublas,cufile,nvjitlink,nvtx]==12.6.3; platform_system == 'Linux' | cuda-bindings>=12.9.4,<13; platform_system == 'Linux' | nvidia-cudnn-cu12==9.10.2.21; platform_system == 'Linux' | nvidia-cusparselt-cu12==0.7.1; platform_system == 'Linux' | nvidia-nccl-cu12==2.29.3; platform_system == 'Linux' | nvidia-nvshmem-cu12==3.4.5; platform_system == 'Linux'
       timeout-minutes: 420
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
@@ -532,80 +476,12 @@ jobs:
       build_name: manywheel-py3_11-cuda-aarch64-12_6
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
     uses: ./.github/workflows/_binary-upload.yml
 
-  manywheel-py3_11-cuda-aarch64-12_8-build:
-    if: ${{ github.repository_owner == 'pytorch' }}
-    uses: ./.github/workflows/_binary-build-linux.yml
-    needs: get-label-type
-    with:
-      PYTORCH_ROOT: /pytorch
-      PACKAGE_TYPE: manywheel
-      # TODO: This is a legacy variable that we eventually want to get rid of in
-      #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu128
-      GPU_ARCH_VERSION: "12.8-aarch64"
-      GPU_ARCH_TYPE: cuda-aarch64
-      DOCKER_IMAGE: manylinuxaarch64-builder
-      DOCKER_IMAGE_TAG_PREFIX: cuda12.8
-      DESIRED_PYTHON: "3.11"
-      runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
-      runs_on: linux.arm64.r7g.12xlarge.memory
-      ALPINE_IMAGE: "arm64v8/alpine"
-      build_name: manywheel-py3_11-cuda-aarch64-12_8
-      build_environment: linux-aarch64-binary-manywheel
-      PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cublas,cufile,nvjitlink,nvtx]==12.8.1; platform_system == 'Linux' | cuda-bindings==12.9.4; platform_system == 'Linux' | nvidia-cudnn-cu12==9.17.1.4; platform_system == 'Linux' | nvidia-cusparselt-cu12==0.7.1; platform_system == 'Linux' | nvidia-nccl-cu12==2.28.9; platform_system == 'Linux' | nvidia-nvshmem-cu12==3.4.5; platform_system == 'Linux'
-      timeout-minutes: 420
-    secrets:
-      github-token: ${{ secrets.GITHUB_TOKEN }}
-
-  manywheel-py3_11-cuda-aarch64-12_8-test:  # Testing
-    if: ${{ github.repository_owner == 'pytorch' }}
-    needs:
-      - manywheel-py3_11-cuda-aarch64-12_8-build
-      - get-label-type
-    uses: ./.github/workflows/_binary-test-linux.yml
-    with:
-      PYTORCH_ROOT: /pytorch
-      PACKAGE_TYPE: manywheel
-      # TODO: This is a legacy variable that we eventually want to get rid of in
-      #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu128
-      GPU_ARCH_VERSION: "12.8-aarch64"
-      GPU_ARCH_TYPE: cuda-aarch64
-      DOCKER_IMAGE: manylinuxaarch64-builder
-      DOCKER_IMAGE_TAG_PREFIX: cuda12.8
-      DESIRED_PYTHON: "3.11"
-      build_name: manywheel-py3_11-cuda-aarch64-12_8
-      build_environment: linux-aarch64-binary-manywheel
-      runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
-      runs_on: linux.arm64.2xlarge
-      ALPINE_IMAGE: "arm64v8/alpine"
-    secrets:
-      github-token: ${{ secrets.GITHUB_TOKEN }}
-  manywheel-py3_11-cuda-aarch64-12_8-upload:  # Uploading
-    if: ${{ github.repository_owner == 'pytorch' }}
-    permissions:
-      id-token: write
-      contents: read
-    needs: manywheel-py3_11-cuda-aarch64-12_8-build
-    with:
-      PYTORCH_ROOT: /pytorch
-      PACKAGE_TYPE: manywheel
-      # TODO: This is a legacy variable that we eventually want to get rid of in
-      #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu128
-      GPU_ARCH_VERSION: "12.8-aarch64"
-      GPU_ARCH_TYPE: cuda-aarch64
-      DOCKER_IMAGE: manylinuxaarch64-builder
-      DOCKER_IMAGE_TAG_PREFIX: cuda12.8
-      DESIRED_PYTHON: "3.11"
-      build_name: manywheel-py3_11-cuda-aarch64-12_8
-    secrets:
-      github-token: ${{ secrets.GITHUB_TOKEN }}
-    uses: ./.github/workflows/_binary-upload.yml
-
-  manywheel-py3_11-cuda-aarch64-12_9-build:
+  manywheel-py3_11-cuda-aarch64-13_0-build:
     if: ${{ github.repository_owner == 'pytorch' }}
     uses: ./.github/workflows/_binary-build-linux.yml
     needs: get-label-type
@@ -614,26 +490,26 @@ jobs:
       PACKAGE_TYPE: manywheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu129
-      GPU_ARCH_VERSION: "12.9-aarch64"
+      DESIRED_CUDA: cu130
+      GPU_ARCH_VERSION: "13.0-aarch64"
       GPU_ARCH_TYPE: cuda-aarch64
       DOCKER_IMAGE: manylinuxaarch64-builder
-      DOCKER_IMAGE_TAG_PREFIX: cuda12.9
+      DOCKER_IMAGE_TAG_PREFIX: cuda13.0
       DESIRED_PYTHON: "3.11"
       runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
       runs_on: linux.arm64.r7g.12xlarge.memory
       ALPINE_IMAGE: "arm64v8/alpine"
-      build_name: manywheel-py3_11-cuda-aarch64-12_9
+      build_name: manywheel-py3_11-cuda-aarch64-13_0
       build_environment: linux-aarch64-binary-manywheel
-      PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cublas,cufile,nvjitlink,nvtx]==12.9.1; platform_system == 'Linux' | cuda-bindings==12.9.4; platform_system == 'Linux' | nvidia-cudnn-cu12==9.17.1.4; platform_system == 'Linux' | nvidia-cusparselt-cu12==0.7.1; platform_system == 'Linux' | nvidia-nccl-cu12==2.28.9; platform_system == 'Linux' | nvidia-nvshmem-cu12==3.4.5; platform_system == 'Linux'
+      PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cufile,nvjitlink,nvtx]==13.0.2; platform_system == 'Linux' | nvidia-cublas>=13.1.0.3,<=13.1.1.3; platform_system == 'Linux' | cuda-bindings>=13.0.3,<14; platform_system == 'Linux' | nvidia-cudnn-cu13==9.20.0.48; platform_system == 'Linux' | nvidia-cusparselt-cu13==0.8.1; platform_system == 'Linux' | nvidia-nccl-cu13==2.29.7; platform_system == 'Linux' | nvidia-nvshmem-cu13==3.4.5; platform_system == 'Linux'
       timeout-minutes: 420
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
 
-  manywheel-py3_11-cuda-aarch64-12_9-test:  # Testing
+  manywheel-py3_11-cuda-aarch64-13_0-test:  # Testing
     if: ${{ github.repository_owner == 'pytorch' }}
     needs:
-      - manywheel-py3_11-cuda-aarch64-12_9-build
+      - manywheel-py3_11-cuda-aarch64-13_0-build
       - get-label-type
     uses: ./.github/workflows/_binary-test-linux.yml
     with:
@@ -641,42 +517,45 @@ jobs:
       PACKAGE_TYPE: manywheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu129
-      GPU_ARCH_VERSION: "12.9-aarch64"
+      DESIRED_CUDA: cu130
+      GPU_ARCH_VERSION: "13.0-aarch64"
       GPU_ARCH_TYPE: cuda-aarch64
       DOCKER_IMAGE: manylinuxaarch64-builder
-      DOCKER_IMAGE_TAG_PREFIX: cuda12.9
+      DOCKER_IMAGE_TAG_PREFIX: cuda13.0
       DESIRED_PYTHON: "3.11"
-      build_name: manywheel-py3_11-cuda-aarch64-12_9
+      build_name: manywheel-py3_11-cuda-aarch64-13_0
       build_environment: linux-aarch64-binary-manywheel
       runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
       runs_on: linux.arm64.2xlarge
       ALPINE_IMAGE: "arm64v8/alpine"
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
-  manywheel-py3_11-cuda-aarch64-12_9-upload:  # Uploading
+  manywheel-py3_11-cuda-aarch64-13_0-upload:  # Uploading
     if: ${{ github.repository_owner == 'pytorch' }}
     permissions:
       id-token: write
       contents: read
-    needs: manywheel-py3_11-cuda-aarch64-12_9-build
+    needs: manywheel-py3_11-cuda-aarch64-13_0-build
     with:
       PYTORCH_ROOT: /pytorch
       PACKAGE_TYPE: manywheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu129
-      GPU_ARCH_VERSION: "12.9-aarch64"
+      DESIRED_CUDA: cu130
+      GPU_ARCH_VERSION: "13.0-aarch64"
       GPU_ARCH_TYPE: cuda-aarch64
       DOCKER_IMAGE: manylinuxaarch64-builder
-      DOCKER_IMAGE_TAG_PREFIX: cuda12.9
+      DOCKER_IMAGE_TAG_PREFIX: cuda13.0
       DESIRED_PYTHON: "3.11"
-      build_name: manywheel-py3_11-cuda-aarch64-12_9
+      build_name: manywheel-py3_11-cuda-aarch64-13_0
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
     uses: ./.github/workflows/_binary-upload.yml
 
-  manywheel-py3_11-cuda-aarch64-13_0-build:
+  manywheel-py3_11-cuda-aarch64-13_2-build:
     if: ${{ github.repository_owner == 'pytorch' }}
     uses: ./.github/workflows/_binary-build-linux.yml
     needs: get-label-type
@@ -685,26 +564,26 @@ jobs:
       PACKAGE_TYPE: manywheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu130
-      GPU_ARCH_VERSION: "13.0-aarch64"
+      DESIRED_CUDA: cu132
+      GPU_ARCH_VERSION: "13.2-aarch64"
       GPU_ARCH_TYPE: cuda-aarch64
       DOCKER_IMAGE: manylinuxaarch64-builder
-      DOCKER_IMAGE_TAG_PREFIX: cuda13.0
+      DOCKER_IMAGE_TAG_PREFIX: cuda13.2
       DESIRED_PYTHON: "3.11"
       runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
       runs_on: linux.arm64.r7g.12xlarge.memory
       ALPINE_IMAGE: "arm64v8/alpine"
-      build_name: manywheel-py3_11-cuda-aarch64-13_0
+      build_name: manywheel-py3_11-cuda-aarch64-13_2
       build_environment: linux-aarch64-binary-manywheel
-      PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cublas,cufile,nvjitlink,nvtx]==13.0.2; platform_system == 'Linux' | cuda-bindings==13.0.3; platform_system == 'Linux' | nvidia-cudnn-cu13==9.17.1.4; platform_system == 'Linux' | nvidia-cusparselt-cu13==0.8.0; platform_system == 'Linux' | nvidia-nccl-cu13==2.28.9; platform_system == 'Linux' | nvidia-nvshmem-cu13==3.4.5; platform_system == 'Linux'
+      PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cublas,cufile,nvjitlink,nvtx]==13.2.1; platform_system == 'Linux' | cuda-bindings>=13.0.3,<14; platform_system == 'Linux' | nvidia-cudnn-cu13==9.20.0.48; platform_system == 'Linux' | nvidia-cusparselt-cu13==0.8.1; platform_system == 'Linux' | nvidia-nccl-cu13==2.29.7; platform_system == 'Linux' | nvidia-nvshmem-cu13==3.4.5; platform_system == 'Linux'
       timeout-minutes: 420
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
 
-  manywheel-py3_11-cuda-aarch64-13_0-test:  # Testing
+  manywheel-py3_11-cuda-aarch64-13_2-test:  # Testing
     if: ${{ github.repository_owner == 'pytorch' }}
     needs:
-      - manywheel-py3_11-cuda-aarch64-13_0-build
+      - manywheel-py3_11-cuda-aarch64-13_2-build
       - get-label-type
     uses: ./.github/workflows/_binary-test-linux.yml
     with:
@@ -712,39 +591,42 @@ jobs:
       PACKAGE_TYPE: manywheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu130
-      GPU_ARCH_VERSION: "13.0-aarch64"
+      DESIRED_CUDA: cu132
+      GPU_ARCH_VERSION: "13.2-aarch64"
       GPU_ARCH_TYPE: cuda-aarch64
       DOCKER_IMAGE: manylinuxaarch64-builder
-      DOCKER_IMAGE_TAG_PREFIX: cuda13.0
+      DOCKER_IMAGE_TAG_PREFIX: cuda13.2
       DESIRED_PYTHON: "3.11"
-      build_name: manywheel-py3_11-cuda-aarch64-13_0
+      build_name: manywheel-py3_11-cuda-aarch64-13_2
       build_environment: linux-aarch64-binary-manywheel
       runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
       runs_on: linux.arm64.2xlarge
       ALPINE_IMAGE: "arm64v8/alpine"
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
-  manywheel-py3_11-cuda-aarch64-13_0-upload:  # Uploading
+  manywheel-py3_11-cuda-aarch64-13_2-upload:  # Uploading
     if: ${{ github.repository_owner == 'pytorch' }}
     permissions:
       id-token: write
       contents: read
-    needs: manywheel-py3_11-cuda-aarch64-13_0-build
+    needs: manywheel-py3_11-cuda-aarch64-13_2-build
     with:
       PYTORCH_ROOT: /pytorch
       PACKAGE_TYPE: manywheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu130
-      GPU_ARCH_VERSION: "13.0-aarch64"
+      DESIRED_CUDA: cu132
+      GPU_ARCH_VERSION: "13.2-aarch64"
       GPU_ARCH_TYPE: cuda-aarch64
       DOCKER_IMAGE: manylinuxaarch64-builder
-      DOCKER_IMAGE_TAG_PREFIX: cuda13.0
+      DOCKER_IMAGE_TAG_PREFIX: cuda13.2
       DESIRED_PYTHON: "3.11"
-      build_name: manywheel-py3_11-cuda-aarch64-13_0
+      build_name: manywheel-py3_11-cuda-aarch64-13_2
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
     uses: ./.github/workflows/_binary-upload.yml
 
   manywheel-py3_12-cpu-aarch64-build:
@@ -811,6 +693,9 @@ jobs:
       build_name: manywheel-py3_12-cpu-aarch64
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
     uses: ./.github/workflows/_binary-upload.yml
 
   manywheel-py3_12-cuda-aarch64-12_6-build:
@@ -833,7 +718,7 @@ jobs:
       ALPINE_IMAGE: "arm64v8/alpine"
       build_name: manywheel-py3_12-cuda-aarch64-12_6
       build_environment: linux-aarch64-binary-manywheel
-      PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cublas,cufile,nvjitlink,nvtx]==12.6.3; platform_system == 'Linux' | cuda-bindings==12.9.4; platform_system == 'Linux' | nvidia-cudnn-cu12==9.10.2.21; platform_system == 'Linux' | nvidia-cusparselt-cu12==0.7.1; platform_system == 'Linux' | nvidia-nccl-cu12==2.28.9; platform_system == 'Linux' | nvidia-nvshmem-cu12==3.4.5; platform_system == 'Linux'
+      PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cublas,cufile,nvjitlink,nvtx]==12.6.3; platform_system == 'Linux' | cuda-bindings>=12.9.4,<13; platform_system == 'Linux' | nvidia-cudnn-cu12==9.10.2.21; platform_system == 'Linux' | nvidia-cusparselt-cu12==0.7.1; platform_system == 'Linux' | nvidia-nccl-cu12==2.29.3; platform_system == 'Linux' | nvidia-nvshmem-cu12==3.4.5; platform_system == 'Linux'
       timeout-minutes: 420
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
@@ -882,80 +767,12 @@ jobs:
       build_name: manywheel-py3_12-cuda-aarch64-12_6
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
     uses: ./.github/workflows/_binary-upload.yml
 
-  manywheel-py3_12-cuda-aarch64-12_8-build:
-    if: ${{ github.repository_owner == 'pytorch' }}
-    uses: ./.github/workflows/_binary-build-linux.yml
-    needs: get-label-type
-    with:
-      PYTORCH_ROOT: /pytorch
-      PACKAGE_TYPE: manywheel
-      # TODO: This is a legacy variable that we eventually want to get rid of in
-      #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu128
-      GPU_ARCH_VERSION: "12.8-aarch64"
-      GPU_ARCH_TYPE: cuda-aarch64
-      DOCKER_IMAGE: manylinuxaarch64-builder
-      DOCKER_IMAGE_TAG_PREFIX: cuda12.8
-      DESIRED_PYTHON: "3.12"
-      runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
-      runs_on: linux.arm64.r7g.12xlarge.memory
-      ALPINE_IMAGE: "arm64v8/alpine"
-      build_name: manywheel-py3_12-cuda-aarch64-12_8
-      build_environment: linux-aarch64-binary-manywheel
-      PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cublas,cufile,nvjitlink,nvtx]==12.8.1; platform_system == 'Linux' | cuda-bindings==12.9.4; platform_system == 'Linux' | nvidia-cudnn-cu12==9.17.1.4; platform_system == 'Linux' | nvidia-cusparselt-cu12==0.7.1; platform_system == 'Linux' | nvidia-nccl-cu12==2.28.9; platform_system == 'Linux' | nvidia-nvshmem-cu12==3.4.5; platform_system == 'Linux'
-      timeout-minutes: 420
-    secrets:
-      github-token: ${{ secrets.GITHUB_TOKEN }}
-
-  manywheel-py3_12-cuda-aarch64-12_8-test:  # Testing
-    if: ${{ github.repository_owner == 'pytorch' }}
-    needs:
-      - manywheel-py3_12-cuda-aarch64-12_8-build
-      - get-label-type
-    uses: ./.github/workflows/_binary-test-linux.yml
-    with:
-      PYTORCH_ROOT: /pytorch
-      PACKAGE_TYPE: manywheel
-      # TODO: This is a legacy variable that we eventually want to get rid of in
-      #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu128
-      GPU_ARCH_VERSION: "12.8-aarch64"
-      GPU_ARCH_TYPE: cuda-aarch64
-      DOCKER_IMAGE: manylinuxaarch64-builder
-      DOCKER_IMAGE_TAG_PREFIX: cuda12.8
-      DESIRED_PYTHON: "3.12"
-      build_name: manywheel-py3_12-cuda-aarch64-12_8
-      build_environment: linux-aarch64-binary-manywheel
-      runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
-      runs_on: linux.arm64.2xlarge
-      ALPINE_IMAGE: "arm64v8/alpine"
-    secrets:
-      github-token: ${{ secrets.GITHUB_TOKEN }}
-  manywheel-py3_12-cuda-aarch64-12_8-upload:  # Uploading
-    if: ${{ github.repository_owner == 'pytorch' }}
-    permissions:
-      id-token: write
-      contents: read
-    needs: manywheel-py3_12-cuda-aarch64-12_8-build
-    with:
-      PYTORCH_ROOT: /pytorch
-      PACKAGE_TYPE: manywheel
-      # TODO: This is a legacy variable that we eventually want to get rid of in
-      #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu128
-      GPU_ARCH_VERSION: "12.8-aarch64"
-      GPU_ARCH_TYPE: cuda-aarch64
-      DOCKER_IMAGE: manylinuxaarch64-builder
-      DOCKER_IMAGE_TAG_PREFIX: cuda12.8
-      DESIRED_PYTHON: "3.12"
-      build_name: manywheel-py3_12-cuda-aarch64-12_8
-    secrets:
-      github-token: ${{ secrets.GITHUB_TOKEN }}
-    uses: ./.github/workflows/_binary-upload.yml
-
-  manywheel-py3_12-cuda-aarch64-12_9-build:
+  manywheel-py3_12-cuda-aarch64-13_0-build:
     if: ${{ github.repository_owner == 'pytorch' }}
     uses: ./.github/workflows/_binary-build-linux.yml
     needs: get-label-type
@@ -964,26 +781,26 @@ jobs:
       PACKAGE_TYPE: manywheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu129
-      GPU_ARCH_VERSION: "12.9-aarch64"
+      DESIRED_CUDA: cu130
+      GPU_ARCH_VERSION: "13.0-aarch64"
       GPU_ARCH_TYPE: cuda-aarch64
       DOCKER_IMAGE: manylinuxaarch64-builder
-      DOCKER_IMAGE_TAG_PREFIX: cuda12.9
+      DOCKER_IMAGE_TAG_PREFIX: cuda13.0
       DESIRED_PYTHON: "3.12"
       runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
       runs_on: linux.arm64.r7g.12xlarge.memory
       ALPINE_IMAGE: "arm64v8/alpine"
-      build_name: manywheel-py3_12-cuda-aarch64-12_9
+      build_name: manywheel-py3_12-cuda-aarch64-13_0
       build_environment: linux-aarch64-binary-manywheel
-      PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cublas,cufile,nvjitlink,nvtx]==12.9.1; platform_system == 'Linux' | cuda-bindings==12.9.4; platform_system == 'Linux' | nvidia-cudnn-cu12==9.17.1.4; platform_system == 'Linux' | nvidia-cusparselt-cu12==0.7.1; platform_system == 'Linux' | nvidia-nccl-cu12==2.28.9; platform_system == 'Linux' | nvidia-nvshmem-cu12==3.4.5; platform_system == 'Linux'
+      PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cufile,nvjitlink,nvtx]==13.0.2; platform_system == 'Linux' | nvidia-cublas>=13.1.0.3,<=13.1.1.3; platform_system == 'Linux' | cuda-bindings>=13.0.3,<14; platform_system == 'Linux' | nvidia-cudnn-cu13==9.20.0.48; platform_system == 'Linux' | nvidia-cusparselt-cu13==0.8.1; platform_system == 'Linux' | nvidia-nccl-cu13==2.29.7; platform_system == 'Linux' | nvidia-nvshmem-cu13==3.4.5; platform_system == 'Linux'
       timeout-minutes: 420
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
 
-  manywheel-py3_12-cuda-aarch64-12_9-test:  # Testing
+  manywheel-py3_12-cuda-aarch64-13_0-test:  # Testing
     if: ${{ github.repository_owner == 'pytorch' }}
     needs:
-      - manywheel-py3_12-cuda-aarch64-12_9-build
+      - manywheel-py3_12-cuda-aarch64-13_0-build
       - get-label-type
     uses: ./.github/workflows/_binary-test-linux.yml
     with:
@@ -991,42 +808,45 @@ jobs:
       PACKAGE_TYPE: manywheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu129
-      GPU_ARCH_VERSION: "12.9-aarch64"
+      DESIRED_CUDA: cu130
+      GPU_ARCH_VERSION: "13.0-aarch64"
       GPU_ARCH_TYPE: cuda-aarch64
       DOCKER_IMAGE: manylinuxaarch64-builder
-      DOCKER_IMAGE_TAG_PREFIX: cuda12.9
+      DOCKER_IMAGE_TAG_PREFIX: cuda13.0
       DESIRED_PYTHON: "3.12"
-      build_name: manywheel-py3_12-cuda-aarch64-12_9
+      build_name: manywheel-py3_12-cuda-aarch64-13_0
       build_environment: linux-aarch64-binary-manywheel
       runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
       runs_on: linux.arm64.2xlarge
       ALPINE_IMAGE: "arm64v8/alpine"
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
-  manywheel-py3_12-cuda-aarch64-12_9-upload:  # Uploading
+  manywheel-py3_12-cuda-aarch64-13_0-upload:  # Uploading
     if: ${{ github.repository_owner == 'pytorch' }}
     permissions:
       id-token: write
       contents: read
-    needs: manywheel-py3_12-cuda-aarch64-12_9-build
+    needs: manywheel-py3_12-cuda-aarch64-13_0-build
     with:
       PYTORCH_ROOT: /pytorch
       PACKAGE_TYPE: manywheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu129
-      GPU_ARCH_VERSION: "12.9-aarch64"
+      DESIRED_CUDA: cu130
+      GPU_ARCH_VERSION: "13.0-aarch64"
       GPU_ARCH_TYPE: cuda-aarch64
       DOCKER_IMAGE: manylinuxaarch64-builder
-      DOCKER_IMAGE_TAG_PREFIX: cuda12.9
+      DOCKER_IMAGE_TAG_PREFIX: cuda13.0
       DESIRED_PYTHON: "3.12"
-      build_name: manywheel-py3_12-cuda-aarch64-12_9
+      build_name: manywheel-py3_12-cuda-aarch64-13_0
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
     uses: ./.github/workflows/_binary-upload.yml
 
-  manywheel-py3_12-cuda-aarch64-13_0-build:
+  manywheel-py3_12-cuda-aarch64-13_2-build:
     if: ${{ github.repository_owner == 'pytorch' }}
     uses: ./.github/workflows/_binary-build-linux.yml
     needs: get-label-type
@@ -1035,26 +855,26 @@ jobs:
       PACKAGE_TYPE: manywheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu130
-      GPU_ARCH_VERSION: "13.0-aarch64"
+      DESIRED_CUDA: cu132
+      GPU_ARCH_VERSION: "13.2-aarch64"
       GPU_ARCH_TYPE: cuda-aarch64
       DOCKER_IMAGE: manylinuxaarch64-builder
-      DOCKER_IMAGE_TAG_PREFIX: cuda13.0
+      DOCKER_IMAGE_TAG_PREFIX: cuda13.2
       DESIRED_PYTHON: "3.12"
       runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
       runs_on: linux.arm64.r7g.12xlarge.memory
       ALPINE_IMAGE: "arm64v8/alpine"
-      build_name: manywheel-py3_12-cuda-aarch64-13_0
+      build_name: manywheel-py3_12-cuda-aarch64-13_2
       build_environment: linux-aarch64-binary-manywheel
-      PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cublas,cufile,nvjitlink,nvtx]==13.0.2; platform_system == 'Linux' | cuda-bindings==13.0.3; platform_system == 'Linux' | nvidia-cudnn-cu13==9.17.1.4; platform_system == 'Linux' | nvidia-cusparselt-cu13==0.8.0; platform_system == 'Linux' | nvidia-nccl-cu13==2.28.9; platform_system == 'Linux' | nvidia-nvshmem-cu13==3.4.5; platform_system == 'Linux'
+      PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cublas,cufile,nvjitlink,nvtx]==13.2.1; platform_system == 'Linux' | cuda-bindings>=13.0.3,<14; platform_system == 'Linux' | nvidia-cudnn-cu13==9.20.0.48; platform_system == 'Linux' | nvidia-cusparselt-cu13==0.8.1; platform_system == 'Linux' | nvidia-nccl-cu13==2.29.7; platform_system == 'Linux' | nvidia-nvshmem-cu13==3.4.5; platform_system == 'Linux'
       timeout-minutes: 420
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
 
-  manywheel-py3_12-cuda-aarch64-13_0-test:  # Testing
+  manywheel-py3_12-cuda-aarch64-13_2-test:  # Testing
     if: ${{ github.repository_owner == 'pytorch' }}
     needs:
-      - manywheel-py3_12-cuda-aarch64-13_0-build
+      - manywheel-py3_12-cuda-aarch64-13_2-build
       - get-label-type
     uses: ./.github/workflows/_binary-test-linux.yml
     with:
@@ -1062,39 +882,42 @@ jobs:
       PACKAGE_TYPE: manywheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu130
-      GPU_ARCH_VERSION: "13.0-aarch64"
+      DESIRED_CUDA: cu132
+      GPU_ARCH_VERSION: "13.2-aarch64"
       GPU_ARCH_TYPE: cuda-aarch64
       DOCKER_IMAGE: manylinuxaarch64-builder
-      DOCKER_IMAGE_TAG_PREFIX: cuda13.0
+      DOCKER_IMAGE_TAG_PREFIX: cuda13.2
       DESIRED_PYTHON: "3.12"
-      build_name: manywheel-py3_12-cuda-aarch64-13_0
+      build_name: manywheel-py3_12-cuda-aarch64-13_2
       build_environment: linux-aarch64-binary-manywheel
       runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
       runs_on: linux.arm64.2xlarge
       ALPINE_IMAGE: "arm64v8/alpine"
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
-  manywheel-py3_12-cuda-aarch64-13_0-upload:  # Uploading
+  manywheel-py3_12-cuda-aarch64-13_2-upload:  # Uploading
     if: ${{ github.repository_owner == 'pytorch' }}
     permissions:
       id-token: write
       contents: read
-    needs: manywheel-py3_12-cuda-aarch64-13_0-build
+    needs: manywheel-py3_12-cuda-aarch64-13_2-build
     with:
       PYTORCH_ROOT: /pytorch
       PACKAGE_TYPE: manywheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu130
-      GPU_ARCH_VERSION: "13.0-aarch64"
+      DESIRED_CUDA: cu132
+      GPU_ARCH_VERSION: "13.2-aarch64"
       GPU_ARCH_TYPE: cuda-aarch64
       DOCKER_IMAGE: manylinuxaarch64-builder
-      DOCKER_IMAGE_TAG_PREFIX: cuda13.0
+      DOCKER_IMAGE_TAG_PREFIX: cuda13.2
       DESIRED_PYTHON: "3.12"
-      build_name: manywheel-py3_12-cuda-aarch64-13_0
+      build_name: manywheel-py3_12-cuda-aarch64-13_2
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
     uses: ./.github/workflows/_binary-upload.yml
 
   manywheel-py3_13-cpu-aarch64-build:
@@ -1161,6 +984,9 @@ jobs:
       build_name: manywheel-py3_13-cpu-aarch64
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
     uses: ./.github/workflows/_binary-upload.yml
 
   manywheel-py3_13-cuda-aarch64-12_6-build:
@@ -1183,7 +1009,7 @@ jobs:
       ALPINE_IMAGE: "arm64v8/alpine"
       build_name: manywheel-py3_13-cuda-aarch64-12_6
       build_environment: linux-aarch64-binary-manywheel
-      PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cublas,cufile,nvjitlink,nvtx]==12.6.3; platform_system == 'Linux' | cuda-bindings==12.9.4; platform_system == 'Linux' | nvidia-cudnn-cu12==9.10.2.21; platform_system == 'Linux' | nvidia-cusparselt-cu12==0.7.1; platform_system == 'Linux' | nvidia-nccl-cu12==2.28.9; platform_system == 'Linux' | nvidia-nvshmem-cu12==3.4.5; platform_system == 'Linux'
+      PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cublas,cufile,nvjitlink,nvtx]==12.6.3; platform_system == 'Linux' | cuda-bindings>=12.9.4,<13; platform_system == 'Linux' | nvidia-cudnn-cu12==9.10.2.21; platform_system == 'Linux' | nvidia-cusparselt-cu12==0.7.1; platform_system == 'Linux' | nvidia-nccl-cu12==2.29.3; platform_system == 'Linux' | nvidia-nvshmem-cu12==3.4.5; platform_system == 'Linux'
       timeout-minutes: 420
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
@@ -1232,148 +1058,9 @@ jobs:
       build_name: manywheel-py3_13-cuda-aarch64-12_6
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
-    uses: ./.github/workflows/_binary-upload.yml
-
-  manywheel-py3_13-cuda-aarch64-12_8-build:
-    if: ${{ github.repository_owner == 'pytorch' }}
-    uses: ./.github/workflows/_binary-build-linux.yml
-    needs: get-label-type
-    with:
-      PYTORCH_ROOT: /pytorch
-      PACKAGE_TYPE: manywheel
-      # TODO: This is a legacy variable that we eventually want to get rid of in
-      #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu128
-      GPU_ARCH_VERSION: "12.8-aarch64"
-      GPU_ARCH_TYPE: cuda-aarch64
-      DOCKER_IMAGE: manylinuxaarch64-builder
-      DOCKER_IMAGE_TAG_PREFIX: cuda12.8
-      DESIRED_PYTHON: "3.13"
-      runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
-      runs_on: linux.arm64.r7g.12xlarge.memory
-      ALPINE_IMAGE: "arm64v8/alpine"
-      build_name: manywheel-py3_13-cuda-aarch64-12_8
-      build_environment: linux-aarch64-binary-manywheel
-      PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cublas,cufile,nvjitlink,nvtx]==12.8.1; platform_system == 'Linux' | cuda-bindings==12.9.4; platform_system == 'Linux' | nvidia-cudnn-cu12==9.17.1.4; platform_system == 'Linux' | nvidia-cusparselt-cu12==0.7.1; platform_system == 'Linux' | nvidia-nccl-cu12==2.28.9; platform_system == 'Linux' | nvidia-nvshmem-cu12==3.4.5; platform_system == 'Linux'
-      timeout-minutes: 420
-    secrets:
-      github-token: ${{ secrets.GITHUB_TOKEN }}
-
-  manywheel-py3_13-cuda-aarch64-12_8-test:  # Testing
-    if: ${{ github.repository_owner == 'pytorch' }}
-    needs:
-      - manywheel-py3_13-cuda-aarch64-12_8-build
-      - get-label-type
-    uses: ./.github/workflows/_binary-test-linux.yml
-    with:
-      PYTORCH_ROOT: /pytorch
-      PACKAGE_TYPE: manywheel
-      # TODO: This is a legacy variable that we eventually want to get rid of in
-      #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu128
-      GPU_ARCH_VERSION: "12.8-aarch64"
-      GPU_ARCH_TYPE: cuda-aarch64
-      DOCKER_IMAGE: manylinuxaarch64-builder
-      DOCKER_IMAGE_TAG_PREFIX: cuda12.8
-      DESIRED_PYTHON: "3.13"
-      build_name: manywheel-py3_13-cuda-aarch64-12_8
-      build_environment: linux-aarch64-binary-manywheel
-      runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
-      runs_on: linux.arm64.2xlarge
-      ALPINE_IMAGE: "arm64v8/alpine"
-    secrets:
-      github-token: ${{ secrets.GITHUB_TOKEN }}
-  manywheel-py3_13-cuda-aarch64-12_8-upload:  # Uploading
-    if: ${{ github.repository_owner == 'pytorch' }}
-    permissions:
-      id-token: write
-      contents: read
-    needs: manywheel-py3_13-cuda-aarch64-12_8-build
-    with:
-      PYTORCH_ROOT: /pytorch
-      PACKAGE_TYPE: manywheel
-      # TODO: This is a legacy variable that we eventually want to get rid of in
-      #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu128
-      GPU_ARCH_VERSION: "12.8-aarch64"
-      GPU_ARCH_TYPE: cuda-aarch64
-      DOCKER_IMAGE: manylinuxaarch64-builder
-      DOCKER_IMAGE_TAG_PREFIX: cuda12.8
-      DESIRED_PYTHON: "3.13"
-      build_name: manywheel-py3_13-cuda-aarch64-12_8
-    secrets:
-      github-token: ${{ secrets.GITHUB_TOKEN }}
-    uses: ./.github/workflows/_binary-upload.yml
-
-  manywheel-py3_13-cuda-aarch64-12_9-build:
-    if: ${{ github.repository_owner == 'pytorch' }}
-    uses: ./.github/workflows/_binary-build-linux.yml
-    needs: get-label-type
-    with:
-      PYTORCH_ROOT: /pytorch
-      PACKAGE_TYPE: manywheel
-      # TODO: This is a legacy variable that we eventually want to get rid of in
-      #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu129
-      GPU_ARCH_VERSION: "12.9-aarch64"
-      GPU_ARCH_TYPE: cuda-aarch64
-      DOCKER_IMAGE: manylinuxaarch64-builder
-      DOCKER_IMAGE_TAG_PREFIX: cuda12.9
-      DESIRED_PYTHON: "3.13"
-      runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
-      runs_on: linux.arm64.r7g.12xlarge.memory
-      ALPINE_IMAGE: "arm64v8/alpine"
-      build_name: manywheel-py3_13-cuda-aarch64-12_9
-      build_environment: linux-aarch64-binary-manywheel
-      PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cublas,cufile,nvjitlink,nvtx]==12.9.1; platform_system == 'Linux' | cuda-bindings==12.9.4; platform_system == 'Linux' | nvidia-cudnn-cu12==9.17.1.4; platform_system == 'Linux' | nvidia-cusparselt-cu12==0.7.1; platform_system == 'Linux' | nvidia-nccl-cu12==2.28.9; platform_system == 'Linux' | nvidia-nvshmem-cu12==3.4.5; platform_system == 'Linux'
-      timeout-minutes: 420
-    secrets:
-      github-token: ${{ secrets.GITHUB_TOKEN }}
-
-  manywheel-py3_13-cuda-aarch64-12_9-test:  # Testing
-    if: ${{ github.repository_owner == 'pytorch' }}
-    needs:
-      - manywheel-py3_13-cuda-aarch64-12_9-build
-      - get-label-type
-    uses: ./.github/workflows/_binary-test-linux.yml
-    with:
-      PYTORCH_ROOT: /pytorch
-      PACKAGE_TYPE: manywheel
-      # TODO: This is a legacy variable that we eventually want to get rid of in
-      #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu129
-      GPU_ARCH_VERSION: "12.9-aarch64"
-      GPU_ARCH_TYPE: cuda-aarch64
-      DOCKER_IMAGE: manylinuxaarch64-builder
-      DOCKER_IMAGE_TAG_PREFIX: cuda12.9
-      DESIRED_PYTHON: "3.13"
-      build_name: manywheel-py3_13-cuda-aarch64-12_9
-      build_environment: linux-aarch64-binary-manywheel
-      runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
-      runs_on: linux.arm64.2xlarge
-      ALPINE_IMAGE: "arm64v8/alpine"
-    secrets:
-      github-token: ${{ secrets.GITHUB_TOKEN }}
-  manywheel-py3_13-cuda-aarch64-12_9-upload:  # Uploading
-    if: ${{ github.repository_owner == 'pytorch' }}
-    permissions:
-      id-token: write
-      contents: read
-    needs: manywheel-py3_13-cuda-aarch64-12_9-build
-    with:
-      PYTORCH_ROOT: /pytorch
-      PACKAGE_TYPE: manywheel
-      # TODO: This is a legacy variable that we eventually want to get rid of in
-      #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu129
-      GPU_ARCH_VERSION: "12.9-aarch64"
-      GPU_ARCH_TYPE: cuda-aarch64
-      DOCKER_IMAGE: manylinuxaarch64-builder
-      DOCKER_IMAGE_TAG_PREFIX: cuda12.9
-      DESIRED_PYTHON: "3.13"
-      build_name: manywheel-py3_13-cuda-aarch64-12_9
-    secrets:
-      github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
     uses: ./.github/workflows/_binary-upload.yml
 
   manywheel-py3_13-cuda-aarch64-13_0-build:
@@ -1396,7 +1083,7 @@ jobs:
       ALPINE_IMAGE: "arm64v8/alpine"
       build_name: manywheel-py3_13-cuda-aarch64-13_0
       build_environment: linux-aarch64-binary-manywheel
-      PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cublas,cufile,nvjitlink,nvtx]==13.0.2; platform_system == 'Linux' | cuda-bindings==13.0.3; platform_system == 'Linux' | nvidia-cudnn-cu13==9.17.1.4; platform_system == 'Linux' | nvidia-cusparselt-cu13==0.8.0; platform_system == 'Linux' | nvidia-nccl-cu13==2.28.9; platform_system == 'Linux' | nvidia-nvshmem-cu13==3.4.5; platform_system == 'Linux'
+      PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cufile,nvjitlink,nvtx]==13.0.2; platform_system == 'Linux' | nvidia-cublas>=13.1.0.3,<=13.1.1.3; platform_system == 'Linux' | cuda-bindings>=13.0.3,<14; platform_system == 'Linux' | nvidia-cudnn-cu13==9.20.0.48; platform_system == 'Linux' | nvidia-cusparselt-cu13==0.8.1; platform_system == 'Linux' | nvidia-nccl-cu13==2.29.7; platform_system == 'Linux' | nvidia-nvshmem-cu13==3.4.5; platform_system == 'Linux'
       timeout-minutes: 420
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
@@ -1445,9 +1132,12 @@ jobs:
       build_name: manywheel-py3_13-cuda-aarch64-13_0
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
     uses: ./.github/workflows/_binary-upload.yml
 
-  manywheel-py3_13t-cpu-aarch64-build:
+  manywheel-py3_13-cuda-aarch64-13_2-build:
     if: ${{ github.repository_owner == 'pytorch' }}
     uses: ./.github/workflows/_binary-build-linux.yml
     needs: get-label-type
@@ -1456,23 +1146,26 @@ jobs:
       PACKAGE_TYPE: manywheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cpu
-      GPU_ARCH_TYPE: cpu-aarch64
-      DOCKER_IMAGE: manylinux2_28_aarch64-builder
-      DOCKER_IMAGE_TAG_PREFIX: cpu-aarch64
-      DESIRED_PYTHON: "3.13t"
+      DESIRED_CUDA: cu132
+      GPU_ARCH_VERSION: "13.2-aarch64"
+      GPU_ARCH_TYPE: cuda-aarch64
+      DOCKER_IMAGE: manylinuxaarch64-builder
+      DOCKER_IMAGE_TAG_PREFIX: cuda13.2
+      DESIRED_PYTHON: "3.13"
       runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
       runs_on: linux.arm64.r7g.12xlarge.memory
       ALPINE_IMAGE: "arm64v8/alpine"
-      build_name: manywheel-py3_13t-cpu-aarch64
+      build_name: manywheel-py3_13-cuda-aarch64-13_2
       build_environment: linux-aarch64-binary-manywheel
+      PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cublas,cufile,nvjitlink,nvtx]==13.2.1; platform_system == 'Linux' | cuda-bindings>=13.0.3,<14; platform_system == 'Linux' | nvidia-cudnn-cu13==9.20.0.48; platform_system == 'Linux' | nvidia-cusparselt-cu13==0.8.1; platform_system == 'Linux' | nvidia-nccl-cu13==2.29.7; platform_system == 'Linux' | nvidia-nvshmem-cu13==3.4.5; platform_system == 'Linux'
+      timeout-minutes: 420
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
 
-  manywheel-py3_13t-cpu-aarch64-test:  # Testing
+  manywheel-py3_13-cuda-aarch64-13_2-test:  # Testing
     if: ${{ github.repository_owner == 'pytorch' }}
     needs:
-      - manywheel-py3_13t-cpu-aarch64-build
+      - manywheel-py3_13-cuda-aarch64-13_2-build
       - get-label-type
     uses: ./.github/workflows/_binary-test-linux.yml
     with:
@@ -1480,40 +1173,45 @@ jobs:
       PACKAGE_TYPE: manywheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cpu
-      GPU_ARCH_TYPE: cpu-aarch64
-      DOCKER_IMAGE: manylinux2_28_aarch64-builder
-      DOCKER_IMAGE_TAG_PREFIX: cpu-aarch64
-      DESIRED_PYTHON: "3.13t"
-      build_name: manywheel-py3_13t-cpu-aarch64
+      DESIRED_CUDA: cu132
+      GPU_ARCH_VERSION: "13.2-aarch64"
+      GPU_ARCH_TYPE: cuda-aarch64
+      DOCKER_IMAGE: manylinuxaarch64-builder
+      DOCKER_IMAGE_TAG_PREFIX: cuda13.2
+      DESIRED_PYTHON: "3.13"
+      build_name: manywheel-py3_13-cuda-aarch64-13_2
       build_environment: linux-aarch64-binary-manywheel
       runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
       runs_on: linux.arm64.2xlarge
       ALPINE_IMAGE: "arm64v8/alpine"
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
-  manywheel-py3_13t-cpu-aarch64-upload:  # Uploading
+  manywheel-py3_13-cuda-aarch64-13_2-upload:  # Uploading
     if: ${{ github.repository_owner == 'pytorch' }}
     permissions:
       id-token: write
       contents: read
-    needs: manywheel-py3_13t-cpu-aarch64-test
+    needs: manywheel-py3_13-cuda-aarch64-13_2-build
     with:
       PYTORCH_ROOT: /pytorch
       PACKAGE_TYPE: manywheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cpu
-      GPU_ARCH_TYPE: cpu-aarch64
-      DOCKER_IMAGE: manylinux2_28_aarch64-builder
-      DOCKER_IMAGE_TAG_PREFIX: cpu-aarch64
-      DESIRED_PYTHON: "3.13t"
-      build_name: manywheel-py3_13t-cpu-aarch64
+      DESIRED_CUDA: cu132
+      GPU_ARCH_VERSION: "13.2-aarch64"
+      GPU_ARCH_TYPE: cuda-aarch64
+      DOCKER_IMAGE: manylinuxaarch64-builder
+      DOCKER_IMAGE_TAG_PREFIX: cuda13.2
+      DESIRED_PYTHON: "3.13"
+      build_name: manywheel-py3_13-cuda-aarch64-13_2
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
     uses: ./.github/workflows/_binary-upload.yml
 
-  manywheel-py3_13t-cuda-aarch64-12_6-build:
+  manywheel-py3_13t-cpu-aarch64-build:
     if: ${{ github.repository_owner == 'pytorch' }}
     uses: ./.github/workflows/_binary-build-linux.yml
     needs: get-label-type
@@ -1522,26 +1220,23 @@ jobs:
       PACKAGE_TYPE: manywheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu126
-      GPU_ARCH_VERSION: "12.6-aarch64"
-      GPU_ARCH_TYPE: cuda-aarch64
-      DOCKER_IMAGE: manylinuxaarch64-builder
-      DOCKER_IMAGE_TAG_PREFIX: cuda12.6
+      DESIRED_CUDA: cpu
+      GPU_ARCH_TYPE: cpu-aarch64
+      DOCKER_IMAGE: manylinux2_28_aarch64-builder
+      DOCKER_IMAGE_TAG_PREFIX: cpu-aarch64
       DESIRED_PYTHON: "3.13t"
       runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
       runs_on: linux.arm64.r7g.12xlarge.memory
       ALPINE_IMAGE: "arm64v8/alpine"
-      build_name: manywheel-py3_13t-cuda-aarch64-12_6
+      build_name: manywheel-py3_13t-cpu-aarch64
       build_environment: linux-aarch64-binary-manywheel
-      PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cublas,cufile,nvjitlink,nvtx]==12.6.3; platform_system == 'Linux' | cuda-bindings==12.9.4; platform_system == 'Linux' | nvidia-cudnn-cu12==9.10.2.21; platform_system == 'Linux' | nvidia-cusparselt-cu12==0.7.1; platform_system == 'Linux' | nvidia-nccl-cu12==2.28.9; platform_system == 'Linux' | nvidia-nvshmem-cu12==3.4.5; platform_system == 'Linux'
-      timeout-minutes: 420
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
 
-  manywheel-py3_13t-cuda-aarch64-12_6-test:  # Testing
+  manywheel-py3_13t-cpu-aarch64-test:  # Testing
     if: ${{ github.repository_owner == 'pytorch' }}
     needs:
-      - manywheel-py3_13t-cuda-aarch64-12_6-build
+      - manywheel-py3_13t-cpu-aarch64-build
       - get-label-type
     uses: ./.github/workflows/_binary-test-linux.yml
     with:
@@ -1549,42 +1244,43 @@ jobs:
       PACKAGE_TYPE: manywheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu126
-      GPU_ARCH_VERSION: "12.6-aarch64"
-      GPU_ARCH_TYPE: cuda-aarch64
-      DOCKER_IMAGE: manylinuxaarch64-builder
-      DOCKER_IMAGE_TAG_PREFIX: cuda12.6
+      DESIRED_CUDA: cpu
+      GPU_ARCH_TYPE: cpu-aarch64
+      DOCKER_IMAGE: manylinux2_28_aarch64-builder
+      DOCKER_IMAGE_TAG_PREFIX: cpu-aarch64
       DESIRED_PYTHON: "3.13t"
-      build_name: manywheel-py3_13t-cuda-aarch64-12_6
+      build_name: manywheel-py3_13t-cpu-aarch64
       build_environment: linux-aarch64-binary-manywheel
       runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
       runs_on: linux.arm64.2xlarge
       ALPINE_IMAGE: "arm64v8/alpine"
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
-  manywheel-py3_13t-cuda-aarch64-12_6-upload:  # Uploading
+  manywheel-py3_13t-cpu-aarch64-upload:  # Uploading
     if: ${{ github.repository_owner == 'pytorch' }}
     permissions:
       id-token: write
       contents: read
-    needs: manywheel-py3_13t-cuda-aarch64-12_6-build
+    needs: manywheel-py3_13t-cpu-aarch64-test
     with:
       PYTORCH_ROOT: /pytorch
       PACKAGE_TYPE: manywheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu126
-      GPU_ARCH_VERSION: "12.6-aarch64"
-      GPU_ARCH_TYPE: cuda-aarch64
-      DOCKER_IMAGE: manylinuxaarch64-builder
-      DOCKER_IMAGE_TAG_PREFIX: cuda12.6
+      DESIRED_CUDA: cpu
+      GPU_ARCH_TYPE: cpu-aarch64
+      DOCKER_IMAGE: manylinux2_28_aarch64-builder
+      DOCKER_IMAGE_TAG_PREFIX: cpu-aarch64
       DESIRED_PYTHON: "3.13t"
-      build_name: manywheel-py3_13t-cuda-aarch64-12_6
+      build_name: manywheel-py3_13t-cpu-aarch64
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
     uses: ./.github/workflows/_binary-upload.yml
 
-  manywheel-py3_13t-cuda-aarch64-12_8-build:
+  manywheel-py3_13t-cuda-aarch64-12_6-build:
     if: ${{ github.repository_owner == 'pytorch' }}
     uses: ./.github/workflows/_binary-build-linux.yml
     needs: get-label-type
@@ -1593,26 +1289,26 @@ jobs:
       PACKAGE_TYPE: manywheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu128
-      GPU_ARCH_VERSION: "12.8-aarch64"
+      DESIRED_CUDA: cu126
+      GPU_ARCH_VERSION: "12.6-aarch64"
       GPU_ARCH_TYPE: cuda-aarch64
       DOCKER_IMAGE: manylinuxaarch64-builder
-      DOCKER_IMAGE_TAG_PREFIX: cuda12.8
+      DOCKER_IMAGE_TAG_PREFIX: cuda12.6
       DESIRED_PYTHON: "3.13t"
       runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
       runs_on: linux.arm64.r7g.12xlarge.memory
       ALPINE_IMAGE: "arm64v8/alpine"
-      build_name: manywheel-py3_13t-cuda-aarch64-12_8
+      build_name: manywheel-py3_13t-cuda-aarch64-12_6
       build_environment: linux-aarch64-binary-manywheel
-      PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cublas,cufile,nvjitlink,nvtx]==12.8.1; platform_system == 'Linux' | cuda-bindings==12.9.4; platform_system == 'Linux' | nvidia-cudnn-cu12==9.17.1.4; platform_system == 'Linux' | nvidia-cusparselt-cu12==0.7.1; platform_system == 'Linux' | nvidia-nccl-cu12==2.28.9; platform_system == 'Linux' | nvidia-nvshmem-cu12==3.4.5; platform_system == 'Linux'
+      PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cublas,cufile,nvjitlink,nvtx]==12.6.3; platform_system == 'Linux' | cuda-bindings>=12.9.4,<13; platform_system == 'Linux' | nvidia-cudnn-cu12==9.10.2.21; platform_system == 'Linux' | nvidia-cusparselt-cu12==0.7.1; platform_system == 'Linux' | nvidia-nccl-cu12==2.29.3; platform_system == 'Linux' | nvidia-nvshmem-cu12==3.4.5; platform_system == 'Linux'
       timeout-minutes: 420
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
 
-  manywheel-py3_13t-cuda-aarch64-12_8-test:  # Testing
+  manywheel-py3_13t-cuda-aarch64-12_6-test:  # Testing
     if: ${{ github.repository_owner == 'pytorch' }}
     needs:
-      - manywheel-py3_13t-cuda-aarch64-12_8-build
+      - manywheel-py3_13t-cuda-aarch64-12_6-build
       - get-label-type
     uses: ./.github/workflows/_binary-test-linux.yml
     with:
@@ -1620,42 +1316,45 @@ jobs:
       PACKAGE_TYPE: manywheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu128
-      GPU_ARCH_VERSION: "12.8-aarch64"
+      DESIRED_CUDA: cu126
+      GPU_ARCH_VERSION: "12.6-aarch64"
       GPU_ARCH_TYPE: cuda-aarch64
       DOCKER_IMAGE: manylinuxaarch64-builder
-      DOCKER_IMAGE_TAG_PREFIX: cuda12.8
+      DOCKER_IMAGE_TAG_PREFIX: cuda12.6
       DESIRED_PYTHON: "3.13t"
-      build_name: manywheel-py3_13t-cuda-aarch64-12_8
+      build_name: manywheel-py3_13t-cuda-aarch64-12_6
       build_environment: linux-aarch64-binary-manywheel
       runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
       runs_on: linux.arm64.2xlarge
       ALPINE_IMAGE: "arm64v8/alpine"
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
-  manywheel-py3_13t-cuda-aarch64-12_8-upload:  # Uploading
+  manywheel-py3_13t-cuda-aarch64-12_6-upload:  # Uploading
     if: ${{ github.repository_owner == 'pytorch' }}
     permissions:
       id-token: write
       contents: read
-    needs: manywheel-py3_13t-cuda-aarch64-12_8-build
+    needs: manywheel-py3_13t-cuda-aarch64-12_6-build
     with:
       PYTORCH_ROOT: /pytorch
       PACKAGE_TYPE: manywheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu128
-      GPU_ARCH_VERSION: "12.8-aarch64"
+      DESIRED_CUDA: cu126
+      GPU_ARCH_VERSION: "12.6-aarch64"
       GPU_ARCH_TYPE: cuda-aarch64
       DOCKER_IMAGE: manylinuxaarch64-builder
-      DOCKER_IMAGE_TAG_PREFIX: cuda12.8
+      DOCKER_IMAGE_TAG_PREFIX: cuda12.6
       DESIRED_PYTHON: "3.13t"
-      build_name: manywheel-py3_13t-cuda-aarch64-12_8
+      build_name: manywheel-py3_13t-cuda-aarch64-12_6
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
     uses: ./.github/workflows/_binary-upload.yml
 
-  manywheel-py3_13t-cuda-aarch64-12_9-build:
+  manywheel-py3_13t-cuda-aarch64-13_0-build:
     if: ${{ github.repository_owner == 'pytorch' }}
     uses: ./.github/workflows/_binary-build-linux.yml
     needs: get-label-type
@@ -1664,26 +1363,26 @@ jobs:
       PACKAGE_TYPE: manywheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu129
-      GPU_ARCH_VERSION: "12.9-aarch64"
+      DESIRED_CUDA: cu130
+      GPU_ARCH_VERSION: "13.0-aarch64"
       GPU_ARCH_TYPE: cuda-aarch64
       DOCKER_IMAGE: manylinuxaarch64-builder
-      DOCKER_IMAGE_TAG_PREFIX: cuda12.9
+      DOCKER_IMAGE_TAG_PREFIX: cuda13.0
       DESIRED_PYTHON: "3.13t"
       runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
       runs_on: linux.arm64.r7g.12xlarge.memory
       ALPINE_IMAGE: "arm64v8/alpine"
-      build_name: manywheel-py3_13t-cuda-aarch64-12_9
+      build_name: manywheel-py3_13t-cuda-aarch64-13_0
       build_environment: linux-aarch64-binary-manywheel
-      PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cublas,cufile,nvjitlink,nvtx]==12.9.1; platform_system == 'Linux' | cuda-bindings==12.9.4; platform_system == 'Linux' | nvidia-cudnn-cu12==9.17.1.4; platform_system == 'Linux' | nvidia-cusparselt-cu12==0.7.1; platform_system == 'Linux' | nvidia-nccl-cu12==2.28.9; platform_system == 'Linux' | nvidia-nvshmem-cu12==3.4.5; platform_system == 'Linux'
+      PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cufile,nvjitlink,nvtx]==13.0.2; platform_system == 'Linux' | nvidia-cublas>=13.1.0.3,<=13.1.1.3; platform_system == 'Linux' | cuda-bindings>=13.0.3,<14; platform_system == 'Linux' | nvidia-cudnn-cu13==9.20.0.48; platform_system == 'Linux' | nvidia-cusparselt-cu13==0.8.1; platform_system == 'Linux' | nvidia-nccl-cu13==2.29.7; platform_system == 'Linux' | nvidia-nvshmem-cu13==3.4.5; platform_system == 'Linux'
       timeout-minutes: 420
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
 
-  manywheel-py3_13t-cuda-aarch64-12_9-test:  # Testing
+  manywheel-py3_13t-cuda-aarch64-13_0-test:  # Testing
     if: ${{ github.repository_owner == 'pytorch' }}
     needs:
-      - manywheel-py3_13t-cuda-aarch64-12_9-build
+      - manywheel-py3_13t-cuda-aarch64-13_0-build
       - get-label-type
     uses: ./.github/workflows/_binary-test-linux.yml
     with:
@@ -1691,42 +1390,45 @@ jobs:
       PACKAGE_TYPE: manywheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu129
-      GPU_ARCH_VERSION: "12.9-aarch64"
+      DESIRED_CUDA: cu130
+      GPU_ARCH_VERSION: "13.0-aarch64"
       GPU_ARCH_TYPE: cuda-aarch64
       DOCKER_IMAGE: manylinuxaarch64-builder
-      DOCKER_IMAGE_TAG_PREFIX: cuda12.9
+      DOCKER_IMAGE_TAG_PREFIX: cuda13.0
       DESIRED_PYTHON: "3.13t"
-      build_name: manywheel-py3_13t-cuda-aarch64-12_9
+      build_name: manywheel-py3_13t-cuda-aarch64-13_0
       build_environment: linux-aarch64-binary-manywheel
       runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
       runs_on: linux.arm64.2xlarge
       ALPINE_IMAGE: "arm64v8/alpine"
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
-  manywheel-py3_13t-cuda-aarch64-12_9-upload:  # Uploading
+  manywheel-py3_13t-cuda-aarch64-13_0-upload:  # Uploading
     if: ${{ github.repository_owner == 'pytorch' }}
     permissions:
       id-token: write
       contents: read
-    needs: manywheel-py3_13t-cuda-aarch64-12_9-build
+    needs: manywheel-py3_13t-cuda-aarch64-13_0-build
     with:
       PYTORCH_ROOT: /pytorch
       PACKAGE_TYPE: manywheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu129
-      GPU_ARCH_VERSION: "12.9-aarch64"
+      DESIRED_CUDA: cu130
+      GPU_ARCH_VERSION: "13.0-aarch64"
       GPU_ARCH_TYPE: cuda-aarch64
       DOCKER_IMAGE: manylinuxaarch64-builder
-      DOCKER_IMAGE_TAG_PREFIX: cuda12.9
+      DOCKER_IMAGE_TAG_PREFIX: cuda13.0
       DESIRED_PYTHON: "3.13t"
-      build_name: manywheel-py3_13t-cuda-aarch64-12_9
+      build_name: manywheel-py3_13t-cuda-aarch64-13_0
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
     uses: ./.github/workflows/_binary-upload.yml
 
-  manywheel-py3_13t-cuda-aarch64-13_0-build:
+  manywheel-py3_13t-cuda-aarch64-13_2-build:
     if: ${{ github.repository_owner == 'pytorch' }}
     uses: ./.github/workflows/_binary-build-linux.yml
     needs: get-label-type
@@ -1735,26 +1437,26 @@ jobs:
       PACKAGE_TYPE: manywheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu130
-      GPU_ARCH_VERSION: "13.0-aarch64"
+      DESIRED_CUDA: cu132
+      GPU_ARCH_VERSION: "13.2-aarch64"
       GPU_ARCH_TYPE: cuda-aarch64
       DOCKER_IMAGE: manylinuxaarch64-builder
-      DOCKER_IMAGE_TAG_PREFIX: cuda13.0
+      DOCKER_IMAGE_TAG_PREFIX: cuda13.2
       DESIRED_PYTHON: "3.13t"
       runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
       runs_on: linux.arm64.r7g.12xlarge.memory
       ALPINE_IMAGE: "arm64v8/alpine"
-      build_name: manywheel-py3_13t-cuda-aarch64-13_0
+      build_name: manywheel-py3_13t-cuda-aarch64-13_2
       build_environment: linux-aarch64-binary-manywheel
-      PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cublas,cufile,nvjitlink,nvtx]==13.0.2; platform_system == 'Linux' | cuda-bindings==13.0.3; platform_system == 'Linux' | nvidia-cudnn-cu13==9.17.1.4; platform_system == 'Linux' | nvidia-cusparselt-cu13==0.8.0; platform_system == 'Linux' | nvidia-nccl-cu13==2.28.9; platform_system == 'Linux' | nvidia-nvshmem-cu13==3.4.5; platform_system == 'Linux'
+      PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cublas,cufile,nvjitlink,nvtx]==13.2.1; platform_system == 'Linux' | cuda-bindings>=13.0.3,<14; platform_system == 'Linux' | nvidia-cudnn-cu13==9.20.0.48; platform_system == 'Linux' | nvidia-cusparselt-cu13==0.8.1; platform_system == 'Linux' | nvidia-nccl-cu13==2.29.7; platform_system == 'Linux' | nvidia-nvshmem-cu13==3.4.5; platform_system == 'Linux'
       timeout-minutes: 420
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
 
-  manywheel-py3_13t-cuda-aarch64-13_0-test:  # Testing
+  manywheel-py3_13t-cuda-aarch64-13_2-test:  # Testing
     if: ${{ github.repository_owner == 'pytorch' }}
     needs:
-      - manywheel-py3_13t-cuda-aarch64-13_0-build
+      - manywheel-py3_13t-cuda-aarch64-13_2-build
       - get-label-type
     uses: ./.github/workflows/_binary-test-linux.yml
     with:
@@ -1762,39 +1464,42 @@ jobs:
       PACKAGE_TYPE: manywheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu130
-      GPU_ARCH_VERSION: "13.0-aarch64"
+      DESIRED_CUDA: cu132
+      GPU_ARCH_VERSION: "13.2-aarch64"
       GPU_ARCH_TYPE: cuda-aarch64
       DOCKER_IMAGE: manylinuxaarch64-builder
-      DOCKER_IMAGE_TAG_PREFIX: cuda13.0
+      DOCKER_IMAGE_TAG_PREFIX: cuda13.2
       DESIRED_PYTHON: "3.13t"
-      build_name: manywheel-py3_13t-cuda-aarch64-13_0
+      build_name: manywheel-py3_13t-cuda-aarch64-13_2
       build_environment: linux-aarch64-binary-manywheel
       runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
       runs_on: linux.arm64.2xlarge
       ALPINE_IMAGE: "arm64v8/alpine"
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
-  manywheel-py3_13t-cuda-aarch64-13_0-upload:  # Uploading
+  manywheel-py3_13t-cuda-aarch64-13_2-upload:  # Uploading
     if: ${{ github.repository_owner == 'pytorch' }}
     permissions:
       id-token: write
       contents: read
-    needs: manywheel-py3_13t-cuda-aarch64-13_0-build
+    needs: manywheel-py3_13t-cuda-aarch64-13_2-build
     with:
       PYTORCH_ROOT: /pytorch
       PACKAGE_TYPE: manywheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu130
-      GPU_ARCH_VERSION: "13.0-aarch64"
+      DESIRED_CUDA: cu132
+      GPU_ARCH_VERSION: "13.2-aarch64"
       GPU_ARCH_TYPE: cuda-aarch64
       DOCKER_IMAGE: manylinuxaarch64-builder
-      DOCKER_IMAGE_TAG_PREFIX: cuda13.0
+      DOCKER_IMAGE_TAG_PREFIX: cuda13.2
       DESIRED_PYTHON: "3.13t"
-      build_name: manywheel-py3_13t-cuda-aarch64-13_0
+      build_name: manywheel-py3_13t-cuda-aarch64-13_2
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
     uses: ./.github/workflows/_binary-upload.yml
 
   manywheel-py3_14-cpu-aarch64-build:
@@ -1861,6 +1566,9 @@ jobs:
       build_name: manywheel-py3_14-cpu-aarch64
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
     uses: ./.github/workflows/_binary-upload.yml
 
   manywheel-py3_14-cuda-aarch64-12_6-build:
@@ -1883,7 +1591,7 @@ jobs:
       ALPINE_IMAGE: "arm64v8/alpine"
       build_name: manywheel-py3_14-cuda-aarch64-12_6
       build_environment: linux-aarch64-binary-manywheel
-      PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cublas,cufile,nvjitlink,nvtx]==12.6.3; platform_system == 'Linux' | cuda-bindings==12.9.4; platform_system == 'Linux' | nvidia-cudnn-cu12==9.10.2.21; platform_system == 'Linux' | nvidia-cusparselt-cu12==0.7.1; platform_system == 'Linux' | nvidia-nccl-cu12==2.28.9; platform_system == 'Linux' | nvidia-nvshmem-cu12==3.4.5; platform_system == 'Linux'
+      PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cublas,cufile,nvjitlink,nvtx]==12.6.3; platform_system == 'Linux' | cuda-bindings>=12.9.4,<13; platform_system == 'Linux' | nvidia-cudnn-cu12==9.10.2.21; platform_system == 'Linux' | nvidia-cusparselt-cu12==0.7.1; platform_system == 'Linux' | nvidia-nccl-cu12==2.29.3; platform_system == 'Linux' | nvidia-nvshmem-cu12==3.4.5; platform_system == 'Linux'
       timeout-minutes: 420
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
@@ -1932,80 +1640,12 @@ jobs:
       build_name: manywheel-py3_14-cuda-aarch64-12_6
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
     uses: ./.github/workflows/_binary-upload.yml
 
-  manywheel-py3_14-cuda-aarch64-12_8-build:
-    if: ${{ github.repository_owner == 'pytorch' }}
-    uses: ./.github/workflows/_binary-build-linux.yml
-    needs: get-label-type
-    with:
-      PYTORCH_ROOT: /pytorch
-      PACKAGE_TYPE: manywheel
-      # TODO: This is a legacy variable that we eventually want to get rid of in
-      #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu128
-      GPU_ARCH_VERSION: "12.8-aarch64"
-      GPU_ARCH_TYPE: cuda-aarch64
-      DOCKER_IMAGE: manylinuxaarch64-builder
-      DOCKER_IMAGE_TAG_PREFIX: cuda12.8
-      DESIRED_PYTHON: "3.14"
-      runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
-      runs_on: linux.arm64.r7g.12xlarge.memory
-      ALPINE_IMAGE: "arm64v8/alpine"
-      build_name: manywheel-py3_14-cuda-aarch64-12_8
-      build_environment: linux-aarch64-binary-manywheel
-      PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cublas,cufile,nvjitlink,nvtx]==12.8.1; platform_system == 'Linux' | cuda-bindings==12.9.4; platform_system == 'Linux' | nvidia-cudnn-cu12==9.17.1.4; platform_system == 'Linux' | nvidia-cusparselt-cu12==0.7.1; platform_system == 'Linux' | nvidia-nccl-cu12==2.28.9; platform_system == 'Linux' | nvidia-nvshmem-cu12==3.4.5; platform_system == 'Linux'
-      timeout-minutes: 420
-    secrets:
-      github-token: ${{ secrets.GITHUB_TOKEN }}
-
-  manywheel-py3_14-cuda-aarch64-12_8-test:  # Testing
-    if: ${{ github.repository_owner == 'pytorch' }}
-    needs:
-      - manywheel-py3_14-cuda-aarch64-12_8-build
-      - get-label-type
-    uses: ./.github/workflows/_binary-test-linux.yml
-    with:
-      PYTORCH_ROOT: /pytorch
-      PACKAGE_TYPE: manywheel
-      # TODO: This is a legacy variable that we eventually want to get rid of in
-      #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu128
-      GPU_ARCH_VERSION: "12.8-aarch64"
-      GPU_ARCH_TYPE: cuda-aarch64
-      DOCKER_IMAGE: manylinuxaarch64-builder
-      DOCKER_IMAGE_TAG_PREFIX: cuda12.8
-      DESIRED_PYTHON: "3.14"
-      build_name: manywheel-py3_14-cuda-aarch64-12_8
-      build_environment: linux-aarch64-binary-manywheel
-      runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
-      runs_on: linux.arm64.2xlarge
-      ALPINE_IMAGE: "arm64v8/alpine"
-    secrets:
-      github-token: ${{ secrets.GITHUB_TOKEN }}
-  manywheel-py3_14-cuda-aarch64-12_8-upload:  # Uploading
-    if: ${{ github.repository_owner == 'pytorch' }}
-    permissions:
-      id-token: write
-      contents: read
-    needs: manywheel-py3_14-cuda-aarch64-12_8-build
-    with:
-      PYTORCH_ROOT: /pytorch
-      PACKAGE_TYPE: manywheel
-      # TODO: This is a legacy variable that we eventually want to get rid of in
-      #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu128
-      GPU_ARCH_VERSION: "12.8-aarch64"
-      GPU_ARCH_TYPE: cuda-aarch64
-      DOCKER_IMAGE: manylinuxaarch64-builder
-      DOCKER_IMAGE_TAG_PREFIX: cuda12.8
-      DESIRED_PYTHON: "3.14"
-      build_name: manywheel-py3_14-cuda-aarch64-12_8
-    secrets:
-      github-token: ${{ secrets.GITHUB_TOKEN }}
-    uses: ./.github/workflows/_binary-upload.yml
-
-  manywheel-py3_14-cuda-aarch64-12_9-build:
+  manywheel-py3_14-cuda-aarch64-13_0-build:
     if: ${{ github.repository_owner == 'pytorch' }}
     uses: ./.github/workflows/_binary-build-linux.yml
     needs: get-label-type
@@ -2014,26 +1654,26 @@ jobs:
       PACKAGE_TYPE: manywheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu129
-      GPU_ARCH_VERSION: "12.9-aarch64"
+      DESIRED_CUDA: cu130
+      GPU_ARCH_VERSION: "13.0-aarch64"
       GPU_ARCH_TYPE: cuda-aarch64
       DOCKER_IMAGE: manylinuxaarch64-builder
-      DOCKER_IMAGE_TAG_PREFIX: cuda12.9
+      DOCKER_IMAGE_TAG_PREFIX: cuda13.0
       DESIRED_PYTHON: "3.14"
       runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
       runs_on: linux.arm64.r7g.12xlarge.memory
       ALPINE_IMAGE: "arm64v8/alpine"
-      build_name: manywheel-py3_14-cuda-aarch64-12_9
+      build_name: manywheel-py3_14-cuda-aarch64-13_0
       build_environment: linux-aarch64-binary-manywheel
-      PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cublas,cufile,nvjitlink,nvtx]==12.9.1; platform_system == 'Linux' | cuda-bindings==12.9.4; platform_system == 'Linux' | nvidia-cudnn-cu12==9.17.1.4; platform_system == 'Linux' | nvidia-cusparselt-cu12==0.7.1; platform_system == 'Linux' | nvidia-nccl-cu12==2.28.9; platform_system == 'Linux' | nvidia-nvshmem-cu12==3.4.5; platform_system == 'Linux'
+      PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cufile,nvjitlink,nvtx]==13.0.2; platform_system == 'Linux' | nvidia-cublas>=13.1.0.3,<=13.1.1.3; platform_system == 'Linux' | cuda-bindings>=13.0.3,<14; platform_system == 'Linux' | nvidia-cudnn-cu13==9.20.0.48; platform_system == 'Linux' | nvidia-cusparselt-cu13==0.8.1; platform_system == 'Linux' | nvidia-nccl-cu13==2.29.7; platform_system == 'Linux' | nvidia-nvshmem-cu13==3.4.5; platform_system == 'Linux'
       timeout-minutes: 420
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
 
-  manywheel-py3_14-cuda-aarch64-12_9-test:  # Testing
+  manywheel-py3_14-cuda-aarch64-13_0-test:  # Testing
     if: ${{ github.repository_owner == 'pytorch' }}
     needs:
-      - manywheel-py3_14-cuda-aarch64-12_9-build
+      - manywheel-py3_14-cuda-aarch64-13_0-build
       - get-label-type
     uses: ./.github/workflows/_binary-test-linux.yml
     with:
@@ -2041,42 +1681,45 @@ jobs:
       PACKAGE_TYPE: manywheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu129
-      GPU_ARCH_VERSION: "12.9-aarch64"
+      DESIRED_CUDA: cu130
+      GPU_ARCH_VERSION: "13.0-aarch64"
       GPU_ARCH_TYPE: cuda-aarch64
       DOCKER_IMAGE: manylinuxaarch64-builder
-      DOCKER_IMAGE_TAG_PREFIX: cuda12.9
+      DOCKER_IMAGE_TAG_PREFIX: cuda13.0
       DESIRED_PYTHON: "3.14"
-      build_name: manywheel-py3_14-cuda-aarch64-12_9
+      build_name: manywheel-py3_14-cuda-aarch64-13_0
       build_environment: linux-aarch64-binary-manywheel
       runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
       runs_on: linux.arm64.2xlarge
       ALPINE_IMAGE: "arm64v8/alpine"
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
-  manywheel-py3_14-cuda-aarch64-12_9-upload:  # Uploading
+  manywheel-py3_14-cuda-aarch64-13_0-upload:  # Uploading
     if: ${{ github.repository_owner == 'pytorch' }}
     permissions:
       id-token: write
       contents: read
-    needs: manywheel-py3_14-cuda-aarch64-12_9-build
+    needs: manywheel-py3_14-cuda-aarch64-13_0-build
     with:
       PYTORCH_ROOT: /pytorch
       PACKAGE_TYPE: manywheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu129
-      GPU_ARCH_VERSION: "12.9-aarch64"
+      DESIRED_CUDA: cu130
+      GPU_ARCH_VERSION: "13.0-aarch64"
       GPU_ARCH_TYPE: cuda-aarch64
       DOCKER_IMAGE: manylinuxaarch64-builder
-      DOCKER_IMAGE_TAG_PREFIX: cuda12.9
+      DOCKER_IMAGE_TAG_PREFIX: cuda13.0
       DESIRED_PYTHON: "3.14"
-      build_name: manywheel-py3_14-cuda-aarch64-12_9
+      build_name: manywheel-py3_14-cuda-aarch64-13_0
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
     uses: ./.github/workflows/_binary-upload.yml
 
-  manywheel-py3_14-cuda-aarch64-13_0-build:
+  manywheel-py3_14-cuda-aarch64-13_2-build:
     if: ${{ github.repository_owner == 'pytorch' }}
     uses: ./.github/workflows/_binary-build-linux.yml
     needs: get-label-type
@@ -2085,26 +1728,26 @@ jobs:
       PACKAGE_TYPE: manywheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu130
-      GPU_ARCH_VERSION: "13.0-aarch64"
+      DESIRED_CUDA: cu132
+      GPU_ARCH_VERSION: "13.2-aarch64"
       GPU_ARCH_TYPE: cuda-aarch64
       DOCKER_IMAGE: manylinuxaarch64-builder
-      DOCKER_IMAGE_TAG_PREFIX: cuda13.0
+      DOCKER_IMAGE_TAG_PREFIX: cuda13.2
       DESIRED_PYTHON: "3.14"
       runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
       runs_on: linux.arm64.r7g.12xlarge.memory
       ALPINE_IMAGE: "arm64v8/alpine"
-      build_name: manywheel-py3_14-cuda-aarch64-13_0
+      build_name: manywheel-py3_14-cuda-aarch64-13_2
       build_environment: linux-aarch64-binary-manywheel
-      PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cublas,cufile,nvjitlink,nvtx]==13.0.2; platform_system == 'Linux' | cuda-bindings==13.0.3; platform_system == 'Linux' | nvidia-cudnn-cu13==9.17.1.4; platform_system == 'Linux' | nvidia-cusparselt-cu13==0.8.0; platform_system == 'Linux' | nvidia-nccl-cu13==2.28.9; platform_system == 'Linux' | nvidia-nvshmem-cu13==3.4.5; platform_system == 'Linux'
+      PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cublas,cufile,nvjitlink,nvtx]==13.2.1; platform_system == 'Linux' | cuda-bindings>=13.0.3,<14; platform_system == 'Linux' | nvidia-cudnn-cu13==9.20.0.48; platform_system == 'Linux' | nvidia-cusparselt-cu13==0.8.1; platform_system == 'Linux' | nvidia-nccl-cu13==2.29.7; platform_system == 'Linux' | nvidia-nvshmem-cu13==3.4.5; platform_system == 'Linux'
       timeout-minutes: 420
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
 
-  manywheel-py3_14-cuda-aarch64-13_0-test:  # Testing
+  manywheel-py3_14-cuda-aarch64-13_2-test:  # Testing
     if: ${{ github.repository_owner == 'pytorch' }}
     needs:
-      - manywheel-py3_14-cuda-aarch64-13_0-build
+      - manywheel-py3_14-cuda-aarch64-13_2-build
       - get-label-type
     uses: ./.github/workflows/_binary-test-linux.yml
     with:
@@ -2112,39 +1755,42 @@ jobs:
       PACKAGE_TYPE: manywheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu130
-      GPU_ARCH_VERSION: "13.0-aarch64"
+      DESIRED_CUDA: cu132
+      GPU_ARCH_VERSION: "13.2-aarch64"
       GPU_ARCH_TYPE: cuda-aarch64
       DOCKER_IMAGE: manylinuxaarch64-builder
-      DOCKER_IMAGE_TAG_PREFIX: cuda13.0
+      DOCKER_IMAGE_TAG_PREFIX: cuda13.2
       DESIRED_PYTHON: "3.14"
-      build_name: manywheel-py3_14-cuda-aarch64-13_0
+      build_name: manywheel-py3_14-cuda-aarch64-13_2
       build_environment: linux-aarch64-binary-manywheel
       runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
       runs_on: linux.arm64.2xlarge
       ALPINE_IMAGE: "arm64v8/alpine"
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
-  manywheel-py3_14-cuda-aarch64-13_0-upload:  # Uploading
+  manywheel-py3_14-cuda-aarch64-13_2-upload:  # Uploading
     if: ${{ github.repository_owner == 'pytorch' }}
     permissions:
       id-token: write
       contents: read
-    needs: manywheel-py3_14-cuda-aarch64-13_0-build
+    needs: manywheel-py3_14-cuda-aarch64-13_2-build
     with:
       PYTORCH_ROOT: /pytorch
       PACKAGE_TYPE: manywheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu130
-      GPU_ARCH_VERSION: "13.0-aarch64"
+      DESIRED_CUDA: cu132
+      GPU_ARCH_VERSION: "13.2-aarch64"
       GPU_ARCH_TYPE: cuda-aarch64
       DOCKER_IMAGE: manylinuxaarch64-builder
-      DOCKER_IMAGE_TAG_PREFIX: cuda13.0
+      DOCKER_IMAGE_TAG_PREFIX: cuda13.2
       DESIRED_PYTHON: "3.14"
-      build_name: manywheel-py3_14-cuda-aarch64-13_0
+      build_name: manywheel-py3_14-cuda-aarch64-13_2
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
     uses: ./.github/workflows/_binary-upload.yml
 
   manywheel-py3_14t-cpu-aarch64-build:
@@ -2211,6 +1857,9 @@ jobs:
       build_name: manywheel-py3_14t-cpu-aarch64
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
     uses: ./.github/workflows/_binary-upload.yml
 
   manywheel-py3_14t-cuda-aarch64-12_6-build:
@@ -2233,7 +1882,7 @@ jobs:
       ALPINE_IMAGE: "arm64v8/alpine"
       build_name: manywheel-py3_14t-cuda-aarch64-12_6
       build_environment: linux-aarch64-binary-manywheel
-      PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cublas,cufile,nvjitlink,nvtx]==12.6.3; platform_system == 'Linux' | cuda-bindings==12.9.4; platform_system == 'Linux' | nvidia-cudnn-cu12==9.10.2.21; platform_system == 'Linux' | nvidia-cusparselt-cu12==0.7.1; platform_system == 'Linux' | nvidia-nccl-cu12==2.28.9; platform_system == 'Linux' | nvidia-nvshmem-cu12==3.4.5; platform_system == 'Linux'
+      PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cublas,cufile,nvjitlink,nvtx]==12.6.3; platform_system == 'Linux' | cuda-bindings>=12.9.4,<13; platform_system == 'Linux' | nvidia-cudnn-cu12==9.10.2.21; platform_system == 'Linux' | nvidia-cusparselt-cu12==0.7.1; platform_system == 'Linux' | nvidia-nccl-cu12==2.29.3; platform_system == 'Linux' | nvidia-nvshmem-cu12==3.4.5; platform_system == 'Linux'
       timeout-minutes: 420
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
@@ -2282,80 +1931,12 @@ jobs:
       build_name: manywheel-py3_14t-cuda-aarch64-12_6
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
     uses: ./.github/workflows/_binary-upload.yml
 
-  manywheel-py3_14t-cuda-aarch64-12_8-build:
-    if: ${{ github.repository_owner == 'pytorch' }}
-    uses: ./.github/workflows/_binary-build-linux.yml
-    needs: get-label-type
-    with:
-      PYTORCH_ROOT: /pytorch
-      PACKAGE_TYPE: manywheel
-      # TODO: This is a legacy variable that we eventually want to get rid of in
-      #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu128
-      GPU_ARCH_VERSION: "12.8-aarch64"
-      GPU_ARCH_TYPE: cuda-aarch64
-      DOCKER_IMAGE: manylinuxaarch64-builder
-      DOCKER_IMAGE_TAG_PREFIX: cuda12.8
-      DESIRED_PYTHON: "3.14t"
-      runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
-      runs_on: linux.arm64.r7g.12xlarge.memory
-      ALPINE_IMAGE: "arm64v8/alpine"
-      build_name: manywheel-py3_14t-cuda-aarch64-12_8
-      build_environment: linux-aarch64-binary-manywheel
-      PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cublas,cufile,nvjitlink,nvtx]==12.8.1; platform_system == 'Linux' | cuda-bindings==12.9.4; platform_system == 'Linux' | nvidia-cudnn-cu12==9.17.1.4; platform_system == 'Linux' | nvidia-cusparselt-cu12==0.7.1; platform_system == 'Linux' | nvidia-nccl-cu12==2.28.9; platform_system == 'Linux' | nvidia-nvshmem-cu12==3.4.5; platform_system == 'Linux'
-      timeout-minutes: 420
-    secrets:
-      github-token: ${{ secrets.GITHUB_TOKEN }}
-
-  manywheel-py3_14t-cuda-aarch64-12_8-test:  # Testing
-    if: ${{ github.repository_owner == 'pytorch' }}
-    needs:
-      - manywheel-py3_14t-cuda-aarch64-12_8-build
-      - get-label-type
-    uses: ./.github/workflows/_binary-test-linux.yml
-    with:
-      PYTORCH_ROOT: /pytorch
-      PACKAGE_TYPE: manywheel
-      # TODO: This is a legacy variable that we eventually want to get rid of in
-      #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu128
-      GPU_ARCH_VERSION: "12.8-aarch64"
-      GPU_ARCH_TYPE: cuda-aarch64
-      DOCKER_IMAGE: manylinuxaarch64-builder
-      DOCKER_IMAGE_TAG_PREFIX: cuda12.8
-      DESIRED_PYTHON: "3.14t"
-      build_name: manywheel-py3_14t-cuda-aarch64-12_8
-      build_environment: linux-aarch64-binary-manywheel
-      runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
-      runs_on: linux.arm64.2xlarge
-      ALPINE_IMAGE: "arm64v8/alpine"
-    secrets:
-      github-token: ${{ secrets.GITHUB_TOKEN }}
-  manywheel-py3_14t-cuda-aarch64-12_8-upload:  # Uploading
-    if: ${{ github.repository_owner == 'pytorch' }}
-    permissions:
-      id-token: write
-      contents: read
-    needs: manywheel-py3_14t-cuda-aarch64-12_8-build
-    with:
-      PYTORCH_ROOT: /pytorch
-      PACKAGE_TYPE: manywheel
-      # TODO: This is a legacy variable that we eventually want to get rid of in
-      #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu128
-      GPU_ARCH_VERSION: "12.8-aarch64"
-      GPU_ARCH_TYPE: cuda-aarch64
-      DOCKER_IMAGE: manylinuxaarch64-builder
-      DOCKER_IMAGE_TAG_PREFIX: cuda12.8
-      DESIRED_PYTHON: "3.14t"
-      build_name: manywheel-py3_14t-cuda-aarch64-12_8
-    secrets:
-      github-token: ${{ secrets.GITHUB_TOKEN }}
-    uses: ./.github/workflows/_binary-upload.yml
-
-  manywheel-py3_14t-cuda-aarch64-12_9-build:
+  manywheel-py3_14t-cuda-aarch64-13_0-build:
     if: ${{ github.repository_owner == 'pytorch' }}
     uses: ./.github/workflows/_binary-build-linux.yml
     needs: get-label-type
@@ -2364,26 +1945,26 @@ jobs:
       PACKAGE_TYPE: manywheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu129
-      GPU_ARCH_VERSION: "12.9-aarch64"
+      DESIRED_CUDA: cu130
+      GPU_ARCH_VERSION: "13.0-aarch64"
       GPU_ARCH_TYPE: cuda-aarch64
       DOCKER_IMAGE: manylinuxaarch64-builder
-      DOCKER_IMAGE_TAG_PREFIX: cuda12.9
+      DOCKER_IMAGE_TAG_PREFIX: cuda13.0
       DESIRED_PYTHON: "3.14t"
       runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
       runs_on: linux.arm64.r7g.12xlarge.memory
       ALPINE_IMAGE: "arm64v8/alpine"
-      build_name: manywheel-py3_14t-cuda-aarch64-12_9
+      build_name: manywheel-py3_14t-cuda-aarch64-13_0
       build_environment: linux-aarch64-binary-manywheel
-      PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cublas,cufile,nvjitlink,nvtx]==12.9.1; platform_system == 'Linux' | cuda-bindings==12.9.4; platform_system == 'Linux' | nvidia-cudnn-cu12==9.17.1.4; platform_system == 'Linux' | nvidia-cusparselt-cu12==0.7.1; platform_system == 'Linux' | nvidia-nccl-cu12==2.28.9; platform_system == 'Linux' | nvidia-nvshmem-cu12==3.4.5; platform_system == 'Linux'
+      PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cufile,nvjitlink,nvtx]==13.0.2; platform_system == 'Linux' | nvidia-cublas>=13.1.0.3,<=13.1.1.3; platform_system == 'Linux' | cuda-bindings>=13.0.3,<14; platform_system == 'Linux' | nvidia-cudnn-cu13==9.20.0.48; platform_system == 'Linux' | nvidia-cusparselt-cu13==0.8.1; platform_system == 'Linux' | nvidia-nccl-cu13==2.29.7; platform_system == 'Linux' | nvidia-nvshmem-cu13==3.4.5; platform_system == 'Linux'
       timeout-minutes: 420
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
 
-  manywheel-py3_14t-cuda-aarch64-12_9-test:  # Testing
+  manywheel-py3_14t-cuda-aarch64-13_0-test:  # Testing
     if: ${{ github.repository_owner == 'pytorch' }}
     needs:
-      - manywheel-py3_14t-cuda-aarch64-12_9-build
+      - manywheel-py3_14t-cuda-aarch64-13_0-build
       - get-label-type
     uses: ./.github/workflows/_binary-test-linux.yml
     with:
@@ -2391,42 +1972,45 @@ jobs:
       PACKAGE_TYPE: manywheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu129
-      GPU_ARCH_VERSION: "12.9-aarch64"
+      DESIRED_CUDA: cu130
+      GPU_ARCH_VERSION: "13.0-aarch64"
       GPU_ARCH_TYPE: cuda-aarch64
       DOCKER_IMAGE: manylinuxaarch64-builder
-      DOCKER_IMAGE_TAG_PREFIX: cuda12.9
+      DOCKER_IMAGE_TAG_PREFIX: cuda13.0
       DESIRED_PYTHON: "3.14t"
-      build_name: manywheel-py3_14t-cuda-aarch64-12_9
+      build_name: manywheel-py3_14t-cuda-aarch64-13_0
       build_environment: linux-aarch64-binary-manywheel
       runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
       runs_on: linux.arm64.2xlarge
       ALPINE_IMAGE: "arm64v8/alpine"
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
-  manywheel-py3_14t-cuda-aarch64-12_9-upload:  # Uploading
+  manywheel-py3_14t-cuda-aarch64-13_0-upload:  # Uploading
     if: ${{ github.repository_owner == 'pytorch' }}
     permissions:
       id-token: write
       contents: read
-    needs: manywheel-py3_14t-cuda-aarch64-12_9-build
+    needs: manywheel-py3_14t-cuda-aarch64-13_0-build
     with:
       PYTORCH_ROOT: /pytorch
       PACKAGE_TYPE: manywheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu129
-      GPU_ARCH_VERSION: "12.9-aarch64"
+      DESIRED_CUDA: cu130
+      GPU_ARCH_VERSION: "13.0-aarch64"
       GPU_ARCH_TYPE: cuda-aarch64
       DOCKER_IMAGE: manylinuxaarch64-builder
-      DOCKER_IMAGE_TAG_PREFIX: cuda12.9
+      DOCKER_IMAGE_TAG_PREFIX: cuda13.0
       DESIRED_PYTHON: "3.14t"
-      build_name: manywheel-py3_14t-cuda-aarch64-12_9
+      build_name: manywheel-py3_14t-cuda-aarch64-13_0
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
     uses: ./.github/workflows/_binary-upload.yml
 
-  manywheel-py3_14t-cuda-aarch64-13_0-build:
+  manywheel-py3_14t-cuda-aarch64-13_2-build:
     if: ${{ github.repository_owner == 'pytorch' }}
     uses: ./.github/workflows/_binary-build-linux.yml
     needs: get-label-type
@@ -2435,26 +2019,26 @@ jobs:
       PACKAGE_TYPE: manywheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu130
-      GPU_ARCH_VERSION: "13.0-aarch64"
+      DESIRED_CUDA: cu132
+      GPU_ARCH_VERSION: "13.2-aarch64"
       GPU_ARCH_TYPE: cuda-aarch64
       DOCKER_IMAGE: manylinuxaarch64-builder
-      DOCKER_IMAGE_TAG_PREFIX: cuda13.0
+      DOCKER_IMAGE_TAG_PREFIX: cuda13.2
       DESIRED_PYTHON: "3.14t"
       runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
       runs_on: linux.arm64.r7g.12xlarge.memory
       ALPINE_IMAGE: "arm64v8/alpine"
-      build_name: manywheel-py3_14t-cuda-aarch64-13_0
+      build_name: manywheel-py3_14t-cuda-aarch64-13_2
       build_environment: linux-aarch64-binary-manywheel
-      PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cublas,cufile,nvjitlink,nvtx]==13.0.2; platform_system == 'Linux' | cuda-bindings==13.0.3; platform_system == 'Linux' | nvidia-cudnn-cu13==9.17.1.4; platform_system == 'Linux' | nvidia-cusparselt-cu13==0.8.0; platform_system == 'Linux' | nvidia-nccl-cu13==2.28.9; platform_system == 'Linux' | nvidia-nvshmem-cu13==3.4.5; platform_system == 'Linux'
+      PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cublas,cufile,nvjitlink,nvtx]==13.2.1; platform_system == 'Linux' | cuda-bindings>=13.0.3,<14; platform_system == 'Linux' | nvidia-cudnn-cu13==9.20.0.48; platform_system == 'Linux' | nvidia-cusparselt-cu13==0.8.1; platform_system == 'Linux' | nvidia-nccl-cu13==2.29.7; platform_system == 'Linux' | nvidia-nvshmem-cu13==3.4.5; platform_system == 'Linux'
       timeout-minutes: 420
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
 
-  manywheel-py3_14t-cuda-aarch64-13_0-test:  # Testing
+  manywheel-py3_14t-cuda-aarch64-13_2-test:  # Testing
     if: ${{ github.repository_owner == 'pytorch' }}
     needs:
-      - manywheel-py3_14t-cuda-aarch64-13_0-build
+      - manywheel-py3_14t-cuda-aarch64-13_2-build
       - get-label-type
     uses: ./.github/workflows/_binary-test-linux.yml
     with:
@@ -2462,37 +2046,40 @@ jobs:
       PACKAGE_TYPE: manywheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu130
-      GPU_ARCH_VERSION: "13.0-aarch64"
+      DESIRED_CUDA: cu132
+      GPU_ARCH_VERSION: "13.2-aarch64"
       GPU_ARCH_TYPE: cuda-aarch64
       DOCKER_IMAGE: manylinuxaarch64-builder
-      DOCKER_IMAGE_TAG_PREFIX: cuda13.0
+      DOCKER_IMAGE_TAG_PREFIX: cuda13.2
       DESIRED_PYTHON: "3.14t"
-      build_name: manywheel-py3_14t-cuda-aarch64-13_0
+      build_name: manywheel-py3_14t-cuda-aarch64-13_2
       build_environment: linux-aarch64-binary-manywheel
       runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
       runs_on: linux.arm64.2xlarge
       ALPINE_IMAGE: "arm64v8/alpine"
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
-  manywheel-py3_14t-cuda-aarch64-13_0-upload:  # Uploading
+  manywheel-py3_14t-cuda-aarch64-13_2-upload:  # Uploading
     if: ${{ github.repository_owner == 'pytorch' }}
     permissions:
       id-token: write
       contents: read
-    needs: manywheel-py3_14t-cuda-aarch64-13_0-build
+    needs: manywheel-py3_14t-cuda-aarch64-13_2-build
     with:
       PYTORCH_ROOT: /pytorch
       PACKAGE_TYPE: manywheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu130
-      GPU_ARCH_VERSION: "13.0-aarch64"
+      DESIRED_CUDA: cu132
+      GPU_ARCH_VERSION: "13.2-aarch64"
       GPU_ARCH_TYPE: cuda-aarch64
       DOCKER_IMAGE: manylinuxaarch64-builder
-      DOCKER_IMAGE_TAG_PREFIX: cuda13.0
+      DOCKER_IMAGE_TAG_PREFIX: cuda13.2
       DESIRED_PYTHON: "3.14t"
-      build_name: manywheel-py3_14t-cuda-aarch64-13_0
+      build_name: manywheel-py3_14t-cuda-aarch64-13_2
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
     uses: ./.github/workflows/_binary-upload.yml
diff --git a/.github/workflows/generated-linux-binary-libtorch-nightly.yml b/.github/workflows/generated-linux-binary-libtorch-nightly.yml
deleted file mode 100644
index 771681dc123d0..0000000000000
--- a/.github/workflows/generated-linux-binary-libtorch-nightly.yml
+++ /dev/null
@@ -1,628 +0,0 @@
-# @generated DO NOT EDIT MANUALLY
-
-# Template is at:    .github/templates/linux_binary_build_workflow.yml.j2
-# Generation script: .github/scripts/generate_ci_workflows.py
-name: linux-binary-libtorch
-
-
-on:
-  push:
-    # NOTE: Meta Employees can trigger new nightlies using: https://fburl.com/trigger_pytorch_nightly_build
-    branches:
-      - nightly
-    tags:
-      # NOTE: Binary build pipelines should only get triggered on release candidate builds
-      # Release candidate tags look like: v1.11.0-rc1
-      - v[0-9]+.[0-9]+.[0-9]+-rc[0-9]+
-      - 'ciflow/binaries/*'
-      - 'ciflow/binaries_libtorch/*'
-  workflow_dispatch:
-
-permissions:
-  id-token: write
-
-env:
-  # Needed for conda builds
-  ALPINE_IMAGE: "308535385114.dkr.ecr.us-east-1.amazonaws.com/tool/alpine"
-  AWS_DEFAULT_REGION: us-east-1
-  BINARY_ENV_FILE: /tmp/env
-  BUILD_ENVIRONMENT: linux-binary-libtorch
-  GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
-  PR_NUMBER: ${{ github.event.pull_request.number }}
-  PYTORCH_FINAL_PACKAGE_DIR: /artifacts
-  PYTORCH_ROOT: /pytorch
-  SHA1: ${{ github.event.pull_request.head.sha || github.sha }}
-  SKIP_ALL_TESTS: 0
-concurrency:
-  group: linux-binary-libtorch-${{ github.event.pull_request.number || github.ref_name }}-${{ github.ref_type == 'branch' && github.sha }}-${{ github.event_name == 'workflow_dispatch' }}
-  cancel-in-progress: true
-
-jobs:
-  get-label-type:
-    if: github.repository_owner == 'pytorch'
-    name: get-label-type
-    uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@main
-    with:
-      triggering_actor: ${{ github.triggering_actor }}
-      issue_owner: ${{ github.event.pull_request.user.login || github.event.issue.user.login }}
-      curr_branch: ${{ github.head_ref || github.ref_name }}
-      curr_ref_type: ${{ github.ref_type }}
-  libtorch-cpu-shared-with-deps-release-build:
-    if: ${{ github.repository_owner == 'pytorch' }}
-    uses: ./.github/workflows/_binary-build-linux.yml
-    needs: get-label-type
-    with:
-      PYTORCH_ROOT: /pytorch
-      PACKAGE_TYPE: libtorch
-      # TODO: This is a legacy variable that we eventually want to get rid of in
-      #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cpu
-      GPU_ARCH_TYPE: cpu
-      DOCKER_IMAGE: libtorch-cxx11-builder
-      DOCKER_IMAGE_TAG_PREFIX: cpu
-      LIBTORCH_CONFIG: release
-      LIBTORCH_VARIANT: shared-with-deps
-      runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
-      build_name: libtorch-cpu-shared-with-deps-release
-      build_environment: linux-binary-libtorch
-    secrets:
-      github-token: ${{ secrets.GITHUB_TOKEN }}
-
-  libtorch-cpu-shared-with-deps-release-test:  # Testing
-    if: ${{ github.repository_owner == 'pytorch' }}
-    needs:
-      - libtorch-cpu-shared-with-deps-release-build
-      - get-label-type
-    uses: ./.github/workflows/_binary-test-linux.yml
-    with:
-      PYTORCH_ROOT: /pytorch
-      PACKAGE_TYPE: libtorch
-      # TODO: This is a legacy variable that we eventually want to get rid of in
-      #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cpu
-      GPU_ARCH_TYPE: cpu
-      DOCKER_IMAGE: libtorch-cxx11-builder
-      DOCKER_IMAGE_TAG_PREFIX: cpu
-      LIBTORCH_CONFIG: release
-      LIBTORCH_VARIANT: shared-with-deps
-      build_name: libtorch-cpu-shared-with-deps-release
-      build_environment: linux-binary-libtorch
-      runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
-      runs_on: linux.4xlarge
-    secrets:
-      github-token: ${{ secrets.GITHUB_TOKEN }}
-  libtorch-cpu-shared-with-deps-release-upload:  # Uploading
-    if: ${{ github.repository_owner == 'pytorch' }}
-    permissions:
-      id-token: write
-      contents: read
-    needs: libtorch-cpu-shared-with-deps-release-test
-    with:
-      PYTORCH_ROOT: /pytorch
-      PACKAGE_TYPE: libtorch
-      # TODO: This is a legacy variable that we eventually want to get rid of in
-      #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cpu
-      GPU_ARCH_TYPE: cpu
-      DOCKER_IMAGE: libtorch-cxx11-builder
-      DOCKER_IMAGE_TAG_PREFIX: cpu
-      LIBTORCH_CONFIG: release
-      LIBTORCH_VARIANT: shared-with-deps
-      build_name: libtorch-cpu-shared-with-deps-release
-    secrets:
-      github-token: ${{ secrets.GITHUB_TOKEN }}
-    uses: ./.github/workflows/_binary-upload.yml
-
-  libtorch-cuda12_6-shared-with-deps-release-build:
-    if: ${{ github.repository_owner == 'pytorch' }}
-    uses: ./.github/workflows/_binary-build-linux.yml
-    needs: get-label-type
-    with:
-      PYTORCH_ROOT: /pytorch
-      PACKAGE_TYPE: libtorch
-      # TODO: This is a legacy variable that we eventually want to get rid of in
-      #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu126
-      GPU_ARCH_VERSION: "12.6"
-      GPU_ARCH_TYPE: cuda
-      DOCKER_IMAGE: libtorch-cxx11-builder
-      DOCKER_IMAGE_TAG_PREFIX: cuda12.6
-      LIBTORCH_CONFIG: release
-      LIBTORCH_VARIANT: shared-with-deps
-      runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
-      build_name: libtorch-cuda12_6-shared-with-deps-release
-      build_environment: linux-binary-libtorch
-    secrets:
-      github-token: ${{ secrets.GITHUB_TOKEN }}
-
-  libtorch-cuda12_6-shared-with-deps-release-test:  # Testing
-    if: ${{ github.repository_owner == 'pytorch' }}
-    needs:
-      - libtorch-cuda12_6-shared-with-deps-release-build
-      - get-label-type
-    uses: ./.github/workflows/_binary-test-linux.yml
-    with:
-      PYTORCH_ROOT: /pytorch
-      PACKAGE_TYPE: libtorch
-      # TODO: This is a legacy variable that we eventually want to get rid of in
-      #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu126
-      GPU_ARCH_VERSION: "12.6"
-      GPU_ARCH_TYPE: cuda
-      DOCKER_IMAGE: libtorch-cxx11-builder
-      DOCKER_IMAGE_TAG_PREFIX: cuda12.6
-      LIBTORCH_CONFIG: release
-      LIBTORCH_VARIANT: shared-with-deps
-      build_name: libtorch-cuda12_6-shared-with-deps-release
-      build_environment: linux-binary-libtorch
-      runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
-      runs_on: linux.4xlarge.nvidia.gpu  # 12.6 build can use maxwell (sm_50) runner
-    secrets:
-      github-token: ${{ secrets.GITHUB_TOKEN }}
-  libtorch-cuda12_6-shared-with-deps-release-upload:  # Uploading
-    if: ${{ github.repository_owner == 'pytorch' }}
-    permissions:
-      id-token: write
-      contents: read
-    needs: libtorch-cuda12_6-shared-with-deps-release-test
-    with:
-      PYTORCH_ROOT: /pytorch
-      PACKAGE_TYPE: libtorch
-      # TODO: This is a legacy variable that we eventually want to get rid of in
-      #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu126
-      GPU_ARCH_VERSION: "12.6"
-      GPU_ARCH_TYPE: cuda
-      DOCKER_IMAGE: libtorch-cxx11-builder
-      DOCKER_IMAGE_TAG_PREFIX: cuda12.6
-      LIBTORCH_CONFIG: release
-      LIBTORCH_VARIANT: shared-with-deps
-      build_name: libtorch-cuda12_6-shared-with-deps-release
-    secrets:
-      github-token: ${{ secrets.GITHUB_TOKEN }}
-    uses: ./.github/workflows/_binary-upload.yml
-
-  libtorch-cuda12_8-shared-with-deps-release-build:
-    if: ${{ github.repository_owner == 'pytorch' }}
-    uses: ./.github/workflows/_binary-build-linux.yml
-    needs: get-label-type
-    with:
-      PYTORCH_ROOT: /pytorch
-      PACKAGE_TYPE: libtorch
-      # TODO: This is a legacy variable that we eventually want to get rid of in
-      #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu128
-      GPU_ARCH_VERSION: "12.8"
-      GPU_ARCH_TYPE: cuda
-      DOCKER_IMAGE: libtorch-cxx11-builder
-      DOCKER_IMAGE_TAG_PREFIX: cuda12.8
-      LIBTORCH_CONFIG: release
-      LIBTORCH_VARIANT: shared-with-deps
-      runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
-      build_name: libtorch-cuda12_8-shared-with-deps-release
-      build_environment: linux-binary-libtorch
-    secrets:
-      github-token: ${{ secrets.GITHUB_TOKEN }}
-
-  libtorch-cuda12_8-shared-with-deps-release-test:  # Testing
-    if: ${{ github.repository_owner == 'pytorch' }}
-    needs:
-      - libtorch-cuda12_8-shared-with-deps-release-build
-      - get-label-type
-    uses: ./.github/workflows/_binary-test-linux.yml
-    with:
-      PYTORCH_ROOT: /pytorch
-      PACKAGE_TYPE: libtorch
-      # TODO: This is a legacy variable that we eventually want to get rid of in
-      #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu128
-      GPU_ARCH_VERSION: "12.8"
-      GPU_ARCH_TYPE: cuda
-      DOCKER_IMAGE: libtorch-cxx11-builder
-      DOCKER_IMAGE_TAG_PREFIX: cuda12.8
-      LIBTORCH_CONFIG: release
-      LIBTORCH_VARIANT: shared-with-deps
-      build_name: libtorch-cuda12_8-shared-with-deps-release
-      build_environment: linux-binary-libtorch
-      runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
-      runs_on: linux.g4dn.4xlarge.nvidia.gpu # 12.8+ builds need sm_70+ runner
-    secrets:
-      github-token: ${{ secrets.GITHUB_TOKEN }}
-  libtorch-cuda12_8-shared-with-deps-release-upload:  # Uploading
-    if: ${{ github.repository_owner == 'pytorch' }}
-    permissions:
-      id-token: write
-      contents: read
-    needs: libtorch-cuda12_8-shared-with-deps-release-test
-    with:
-      PYTORCH_ROOT: /pytorch
-      PACKAGE_TYPE: libtorch
-      # TODO: This is a legacy variable that we eventually want to get rid of in
-      #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu128
-      GPU_ARCH_VERSION: "12.8"
-      GPU_ARCH_TYPE: cuda
-      DOCKER_IMAGE: libtorch-cxx11-builder
-      DOCKER_IMAGE_TAG_PREFIX: cuda12.8
-      LIBTORCH_CONFIG: release
-      LIBTORCH_VARIANT: shared-with-deps
-      build_name: libtorch-cuda12_8-shared-with-deps-release
-    secrets:
-      github-token: ${{ secrets.GITHUB_TOKEN }}
-    uses: ./.github/workflows/_binary-upload.yml
-
-  libtorch-cuda12_9-shared-with-deps-release-build:
-    if: ${{ github.repository_owner == 'pytorch' }}
-    uses: ./.github/workflows/_binary-build-linux.yml
-    needs: get-label-type
-    with:
-      PYTORCH_ROOT: /pytorch
-      PACKAGE_TYPE: libtorch
-      # TODO: This is a legacy variable that we eventually want to get rid of in
-      #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu129
-      GPU_ARCH_VERSION: "12.9"
-      GPU_ARCH_TYPE: cuda
-      DOCKER_IMAGE: libtorch-cxx11-builder
-      DOCKER_IMAGE_TAG_PREFIX: cuda12.9
-      LIBTORCH_CONFIG: release
-      LIBTORCH_VARIANT: shared-with-deps
-      runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
-      build_name: libtorch-cuda12_9-shared-with-deps-release
-      build_environment: linux-binary-libtorch
-    secrets:
-      github-token: ${{ secrets.GITHUB_TOKEN }}
-
-  libtorch-cuda12_9-shared-with-deps-release-test:  # Testing
-    if: ${{ github.repository_owner == 'pytorch' }}
-    needs:
-      - libtorch-cuda12_9-shared-with-deps-release-build
-      - get-label-type
-    uses: ./.github/workflows/_binary-test-linux.yml
-    with:
-      PYTORCH_ROOT: /pytorch
-      PACKAGE_TYPE: libtorch
-      # TODO: This is a legacy variable that we eventually want to get rid of in
-      #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu129
-      GPU_ARCH_VERSION: "12.9"
-      GPU_ARCH_TYPE: cuda
-      DOCKER_IMAGE: libtorch-cxx11-builder
-      DOCKER_IMAGE_TAG_PREFIX: cuda12.9
-      LIBTORCH_CONFIG: release
-      LIBTORCH_VARIANT: shared-with-deps
-      build_name: libtorch-cuda12_9-shared-with-deps-release
-      build_environment: linux-binary-libtorch
-      runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
-      runs_on: linux.g4dn.4xlarge.nvidia.gpu # 12.8+ builds need sm_70+ runner
-    secrets:
-      github-token: ${{ secrets.GITHUB_TOKEN }}
-  libtorch-cuda12_9-shared-with-deps-release-upload:  # Uploading
-    if: ${{ github.repository_owner == 'pytorch' }}
-    permissions:
-      id-token: write
-      contents: read
-    needs: libtorch-cuda12_9-shared-with-deps-release-test
-    with:
-      PYTORCH_ROOT: /pytorch
-      PACKAGE_TYPE: libtorch
-      # TODO: This is a legacy variable that we eventually want to get rid of in
-      #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu129
-      GPU_ARCH_VERSION: "12.9"
-      GPU_ARCH_TYPE: cuda
-      DOCKER_IMAGE: libtorch-cxx11-builder
-      DOCKER_IMAGE_TAG_PREFIX: cuda12.9
-      LIBTORCH_CONFIG: release
-      LIBTORCH_VARIANT: shared-with-deps
-      build_name: libtorch-cuda12_9-shared-with-deps-release
-    secrets:
-      github-token: ${{ secrets.GITHUB_TOKEN }}
-    uses: ./.github/workflows/_binary-upload.yml
-
-  libtorch-cuda13_0-shared-with-deps-release-build:
-    if: ${{ github.repository_owner == 'pytorch' }}
-    uses: ./.github/workflows/_binary-build-linux.yml
-    needs: get-label-type
-    with:
-      PYTORCH_ROOT: /pytorch
-      PACKAGE_TYPE: libtorch
-      # TODO: This is a legacy variable that we eventually want to get rid of in
-      #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu130
-      GPU_ARCH_VERSION: "13.0"
-      GPU_ARCH_TYPE: cuda
-      DOCKER_IMAGE: libtorch-cxx11-builder
-      DOCKER_IMAGE_TAG_PREFIX: cuda13.0
-      LIBTORCH_CONFIG: release
-      LIBTORCH_VARIANT: shared-with-deps
-      runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
-      build_name: libtorch-cuda13_0-shared-with-deps-release
-      build_environment: linux-binary-libtorch
-    secrets:
-      github-token: ${{ secrets.GITHUB_TOKEN }}
-
-  libtorch-cuda13_0-shared-with-deps-release-test:  # Testing
-    if: ${{ github.repository_owner == 'pytorch' }}
-    needs:
-      - libtorch-cuda13_0-shared-with-deps-release-build
-      - get-label-type
-    uses: ./.github/workflows/_binary-test-linux.yml
-    with:
-      PYTORCH_ROOT: /pytorch
-      PACKAGE_TYPE: libtorch
-      # TODO: This is a legacy variable that we eventually want to get rid of in
-      #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu130
-      GPU_ARCH_VERSION: "13.0"
-      GPU_ARCH_TYPE: cuda
-      DOCKER_IMAGE: libtorch-cxx11-builder
-      DOCKER_IMAGE_TAG_PREFIX: cuda13.0
-      LIBTORCH_CONFIG: release
-      LIBTORCH_VARIANT: shared-with-deps
-      build_name: libtorch-cuda13_0-shared-with-deps-release
-      build_environment: linux-binary-libtorch
-      runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
-      runs_on: linux.g4dn.4xlarge.nvidia.gpu # 12.8+ builds need sm_70+ runner
-    secrets:
-      github-token: ${{ secrets.GITHUB_TOKEN }}
-  libtorch-cuda13_0-shared-with-deps-release-upload:  # Uploading
-    if: ${{ github.repository_owner == 'pytorch' }}
-    permissions:
-      id-token: write
-      contents: read
-    needs: libtorch-cuda13_0-shared-with-deps-release-test
-    with:
-      PYTORCH_ROOT: /pytorch
-      PACKAGE_TYPE: libtorch
-      # TODO: This is a legacy variable that we eventually want to get rid of in
-      #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu130
-      GPU_ARCH_VERSION: "13.0"
-      GPU_ARCH_TYPE: cuda
-      DOCKER_IMAGE: libtorch-cxx11-builder
-      DOCKER_IMAGE_TAG_PREFIX: cuda13.0
-      LIBTORCH_CONFIG: release
-      LIBTORCH_VARIANT: shared-with-deps
-      build_name: libtorch-cuda13_0-shared-with-deps-release
-    secrets:
-      github-token: ${{ secrets.GITHUB_TOKEN }}
-    uses: ./.github/workflows/_binary-upload.yml
-
-  libtorch-rocm7_1-shared-with-deps-release-build:
-    if: ${{ github.repository_owner == 'pytorch' }}
-    uses: ./.github/workflows/_binary-build-linux.yml
-    needs: get-label-type
-    with:
-      PYTORCH_ROOT: /pytorch
-      PACKAGE_TYPE: libtorch
-      # TODO: This is a legacy variable that we eventually want to get rid of in
-      #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: rocm7.1
-      GPU_ARCH_VERSION: "7.1"
-      GPU_ARCH_TYPE: rocm
-      DOCKER_IMAGE: libtorch-cxx11-builder
-      DOCKER_IMAGE_TAG_PREFIX: rocm7.1
-      LIBTORCH_CONFIG: release
-      LIBTORCH_VARIANT: shared-with-deps
-      runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
-      timeout-minutes: 300
-      build_name: libtorch-rocm7_1-shared-with-deps-release
-      build_environment: linux-binary-libtorch
-    secrets:
-      github-token: ${{ secrets.GITHUB_TOKEN }}
-
-  libtorch-rocm7_1-shared-with-deps-release-test:  # Testing
-    if: ${{ github.repository_owner == 'pytorch' }}
-    needs:
-      - libtorch-rocm7_1-shared-with-deps-release-build
-      - get-label-type
-    runs-on: linux.rocm.gpu.mi250.1
-    timeout-minutes: 240
-    env:
-      PYTORCH_ROOT: /pytorch
-      PACKAGE_TYPE: libtorch
-      # TODO: This is a legacy variable that we eventually want to get rid of in
-      #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: rocm7.1
-      GPU_ARCH_VERSION: "7.1"
-      GPU_ARCH_TYPE: rocm
-      SKIP_ALL_TESTS: 1
-      DOCKER_IMAGE: libtorch-cxx11-builder
-      DOCKER_IMAGE_TAG_PREFIX: rocm7.1
-      LIBTORCH_CONFIG: release
-      LIBTORCH_VARIANT: shared-with-deps
-    permissions:
-      id-token: write
-      contents: read
-    steps:
-      - name: Setup ROCm
-        uses: ./.github/actions/setup-rocm
-      - uses: actions/download-artifact@v4.1.7
-        name: Download Build Artifacts
-        with:
-          name: libtorch-rocm7_1-shared-with-deps-release
-          path: "${{ runner.temp }}/artifacts/"
-      - name: Checkout PyTorch
-        uses: actions/checkout@v4
-        with:
-          ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
-          submodules: recursive
-          path: pytorch
-          show-progress: false
-      - name: Clean PyTorch checkout
-        run: |
-          # Remove any artifacts from the previous checkouts
-          git clean -fxd
-        working-directory: pytorch
-      - name: ROCm set GPU_FLAG
-        run: |
-          echo "GPU_FLAG=--device=/dev/mem --device=/dev/kfd --device=/dev/dri --group-add video --group-add daemon" >> "${GITHUB_ENV}"
-      - name: configure aws credentials
-        id: aws_creds
-        if: ${{ startsWith(github.event.ref, 'refs/tags/ciflow/') }}
-        uses: aws-actions/configure-aws-credentials@v4
-        with:
-          role-to-assume: arn:aws:iam::308535385114:role/gha_workflow_s3_and_ecr_read_only
-          aws-region: us-east-1
-          role-duration-seconds: 18000
-      - name: Calculate docker image
-        id: calculate-docker-image
-        uses: pytorch/test-infra/.github/actions/calculate-docker-image@main
-        with:
-          docker-registry: ${{ startsWith(github.event.ref, 'refs/tags/ciflow/') && '308535385114.dkr.ecr.us-east-1.amazonaws.com' || 'docker.io' }}
-          docker-image-name: libtorch-cxx11-builder
-          custom-tag-prefix: rocm7.1
-          docker-build-dir: .ci/docker
-          working-directory: pytorch
-      - name: Pull Docker image
-        uses: pytorch/test-infra/.github/actions/pull-docker-image@main
-        with:
-          docker-image: ${{ steps.calculate-docker-image.outputs.docker-image }}
-      - name: Test Pytorch binary
-        uses: ./pytorch/.github/actions/test-pytorch-binary
-        env:
-          DOCKER_IMAGE: ${{ steps.calculate-docker-image.outputs.docker-image }}
-      - name: Teardown ROCm
-        uses: ./.github/actions/teardown-rocm
-  libtorch-rocm7_1-shared-with-deps-release-upload:  # Uploading
-    if: ${{ github.repository_owner == 'pytorch' }}
-    permissions:
-      id-token: write
-      contents: read
-    needs: libtorch-rocm7_1-shared-with-deps-release-test
-    with:
-      PYTORCH_ROOT: /pytorch
-      PACKAGE_TYPE: libtorch
-      # TODO: This is a legacy variable that we eventually want to get rid of in
-      #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: rocm7.1
-      GPU_ARCH_VERSION: "7.1"
-      GPU_ARCH_TYPE: rocm
-      DOCKER_IMAGE: libtorch-cxx11-builder
-      DOCKER_IMAGE_TAG_PREFIX: rocm7.1
-      LIBTORCH_CONFIG: release
-      LIBTORCH_VARIANT: shared-with-deps
-      build_name: libtorch-rocm7_1-shared-with-deps-release
-    secrets:
-      github-token: ${{ secrets.GITHUB_TOKEN }}
-    uses: ./.github/workflows/_binary-upload.yml
-
-  libtorch-rocm7_2-shared-with-deps-release-build:
-    if: ${{ github.repository_owner == 'pytorch' }}
-    uses: ./.github/workflows/_binary-build-linux.yml
-    needs: get-label-type
-    with:
-      PYTORCH_ROOT: /pytorch
-      PACKAGE_TYPE: libtorch
-      # TODO: This is a legacy variable that we eventually want to get rid of in
-      #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: rocm7.2
-      GPU_ARCH_VERSION: "7.2"
-      GPU_ARCH_TYPE: rocm
-      DOCKER_IMAGE: libtorch-cxx11-builder
-      DOCKER_IMAGE_TAG_PREFIX: rocm7.2
-      LIBTORCH_CONFIG: release
-      LIBTORCH_VARIANT: shared-with-deps
-      runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
-      timeout-minutes: 300
-      build_name: libtorch-rocm7_2-shared-with-deps-release
-      build_environment: linux-binary-libtorch
-    secrets:
-      github-token: ${{ secrets.GITHUB_TOKEN }}
-
-  libtorch-rocm7_2-shared-with-deps-release-test:  # Testing
-    if: ${{ github.repository_owner == 'pytorch' }}
-    needs:
-      - libtorch-rocm7_2-shared-with-deps-release-build
-      - get-label-type
-    runs-on: linux.rocm.gpu.mi250.1
-    timeout-minutes: 240
-    env:
-      PYTORCH_ROOT: /pytorch
-      PACKAGE_TYPE: libtorch
-      # TODO: This is a legacy variable that we eventually want to get rid of in
-      #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: rocm7.2
-      GPU_ARCH_VERSION: "7.2"
-      GPU_ARCH_TYPE: rocm
-      SKIP_ALL_TESTS: 1
-      DOCKER_IMAGE: libtorch-cxx11-builder
-      DOCKER_IMAGE_TAG_PREFIX: rocm7.2
-      LIBTORCH_CONFIG: release
-      LIBTORCH_VARIANT: shared-with-deps
-    permissions:
-      id-token: write
-      contents: read
-    steps:
-      - name: Setup ROCm
-        uses: ./.github/actions/setup-rocm
-      - uses: actions/download-artifact@v4.1.7
-        name: Download Build Artifacts
-        with:
-          name: libtorch-rocm7_2-shared-with-deps-release
-          path: "${{ runner.temp }}/artifacts/"
-      - name: Checkout PyTorch
-        uses: actions/checkout@v4
-        with:
-          ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
-          submodules: recursive
-          path: pytorch
-          show-progress: false
-      - name: Clean PyTorch checkout
-        run: |
-          # Remove any artifacts from the previous checkouts
-          git clean -fxd
-        working-directory: pytorch
-      - name: ROCm set GPU_FLAG
-        run: |
-          echo "GPU_FLAG=--device=/dev/mem --device=/dev/kfd --device=/dev/dri --group-add video --group-add daemon" >> "${GITHUB_ENV}"
-      - name: configure aws credentials
-        id: aws_creds
-        if: ${{ startsWith(github.event.ref, 'refs/tags/ciflow/') }}
-        uses: aws-actions/configure-aws-credentials@v4
-        with:
-          role-to-assume: arn:aws:iam::308535385114:role/gha_workflow_s3_and_ecr_read_only
-          aws-region: us-east-1
-          role-duration-seconds: 18000
-      - name: Calculate docker image
-        id: calculate-docker-image
-        uses: pytorch/test-infra/.github/actions/calculate-docker-image@main
-        with:
-          docker-registry: ${{ startsWith(github.event.ref, 'refs/tags/ciflow/') && '308535385114.dkr.ecr.us-east-1.amazonaws.com' || 'docker.io' }}
-          docker-image-name: libtorch-cxx11-builder
-          custom-tag-prefix: rocm7.2
-          docker-build-dir: .ci/docker
-          working-directory: pytorch
-      - name: Pull Docker image
-        uses: pytorch/test-infra/.github/actions/pull-docker-image@main
-        with:
-          docker-image: ${{ steps.calculate-docker-image.outputs.docker-image }}
-      - name: Test Pytorch binary
-        uses: ./pytorch/.github/actions/test-pytorch-binary
-        env:
-          DOCKER_IMAGE: ${{ steps.calculate-docker-image.outputs.docker-image }}
-      - name: Teardown ROCm
-        uses: ./.github/actions/teardown-rocm
-  libtorch-rocm7_2-shared-with-deps-release-upload:  # Uploading
-    if: ${{ github.repository_owner == 'pytorch' }}
-    permissions:
-      id-token: write
-      contents: read
-    needs: libtorch-rocm7_2-shared-with-deps-release-test
-    with:
-      PYTORCH_ROOT: /pytorch
-      PACKAGE_TYPE: libtorch
-      # TODO: This is a legacy variable that we eventually want to get rid of in
-      #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: rocm7.2
-      GPU_ARCH_VERSION: "7.2"
-      GPU_ARCH_TYPE: rocm
-      DOCKER_IMAGE: libtorch-cxx11-builder
-      DOCKER_IMAGE_TAG_PREFIX: rocm7.2
-      LIBTORCH_CONFIG: release
-      LIBTORCH_VARIANT: shared-with-deps
-      build_name: libtorch-rocm7_2-shared-with-deps-release
-    secrets:
-      github-token: ${{ secrets.GITHUB_TOKEN }}
-    uses: ./.github/workflows/_binary-upload.yml
diff --git a/.github/workflows/generated-linux-binary-manywheel-nightly.yml b/.github/workflows/generated-linux-binary-manywheel-nightly.yml
index 4e49ac2223ada..73fcb91f9f74e 100644
--- a/.github/workflows/generated-linux-binary-manywheel-nightly.yml
+++ b/.github/workflows/generated-linux-binary-manywheel-nightly.yml
@@ -15,6 +15,7 @@ on:
       # Release candidate tags look like: v1.11.0-rc1
       - v[0-9]+.[0-9]+.[0-9]+-rc[0-9]+
       - 'ciflow/binaries/*'
+      - 'ciflow/binaries_libtorch/*'
       - 'ciflow/binaries_wheel/*'
   workflow_dispatch:
 
@@ -41,7 +42,7 @@ jobs:
   get-label-type:
     if: github.repository_owner == 'pytorch'
     name: get-label-type
-    uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@main
+    uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@release/2.12
     with:
       triggering_actor: ${{ github.triggering_actor }}
       issue_owner: ${{ github.event.pull_request.user.login || github.event.issue.user.login }}
@@ -108,6 +109,9 @@ jobs:
       build_name: manywheel-py3_10-cpu
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
     uses: ./.github/workflows/_binary-upload.yml
 
   manywheel-py3_10-cuda12_6-build:
@@ -128,7 +132,7 @@ jobs:
       runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
       build_name: manywheel-py3_10-cuda12_6
       build_environment: linux-binary-manywheel
-      PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cublas,cufile,nvjitlink,nvtx]==12.6.3; platform_system == 'Linux' | cuda-bindings==12.9.4; platform_system == 'Linux' | nvidia-cudnn-cu12==9.10.2.21; platform_system == 'Linux' | nvidia-cusparselt-cu12==0.7.1; platform_system == 'Linux' | nvidia-nccl-cu12==2.28.9; platform_system == 'Linux' | nvidia-nvshmem-cu12==3.4.5; platform_system == 'Linux'
+      PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cublas,cufile,nvjitlink,nvtx]==12.6.3; platform_system == 'Linux' | cuda-bindings>=12.9.4,<13; platform_system == 'Linux' | nvidia-cudnn-cu12==9.10.2.21; platform_system == 'Linux' | nvidia-cusparselt-cu12==0.7.1; platform_system == 'Linux' | nvidia-nccl-cu12==2.29.3; platform_system == 'Linux' | nvidia-nvshmem-cu12==3.4.5; platform_system == 'Linux'
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
 
@@ -152,7 +156,7 @@ jobs:
       build_name: manywheel-py3_10-cuda12_6
       build_environment: linux-binary-manywheel
       runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
-      runs_on: linux.4xlarge.nvidia.gpu  # 12.6 build can use maxwell (sm_50) runner
+      runs_on: linux.g4dn.4xlarge.nvidia.gpu
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
   manywheel-py3_10-cuda12_6-upload:  # Uploading
@@ -175,76 +179,12 @@ jobs:
       build_name: manywheel-py3_10-cuda12_6
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
     uses: ./.github/workflows/_binary-upload.yml
 
-  manywheel-py3_10-cuda12_8-build:
-    if: ${{ github.repository_owner == 'pytorch' }}
-    uses: ./.github/workflows/_binary-build-linux.yml
-    needs: get-label-type
-    with:
-      PYTORCH_ROOT: /pytorch
-      PACKAGE_TYPE: manywheel
-      # TODO: This is a legacy variable that we eventually want to get rid of in
-      #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu128
-      GPU_ARCH_VERSION: "12.8"
-      GPU_ARCH_TYPE: cuda
-      DOCKER_IMAGE: manylinux2_28-builder
-      DOCKER_IMAGE_TAG_PREFIX: cuda12.8
-      DESIRED_PYTHON: "3.10"
-      runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
-      build_name: manywheel-py3_10-cuda12_8
-      build_environment: linux-binary-manywheel
-      PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cublas,cufile,nvjitlink,nvtx]==12.8.1; platform_system == 'Linux' | cuda-bindings==12.9.4; platform_system == 'Linux' | nvidia-cudnn-cu12==9.17.1.4; platform_system == 'Linux' | nvidia-cusparselt-cu12==0.7.1; platform_system == 'Linux' | nvidia-nccl-cu12==2.28.9; platform_system == 'Linux' | nvidia-nvshmem-cu12==3.4.5; platform_system == 'Linux'
-    secrets:
-      github-token: ${{ secrets.GITHUB_TOKEN }}
-
-  manywheel-py3_10-cuda12_8-test:  # Testing
-    if: ${{ github.repository_owner == 'pytorch' }}
-    needs:
-      - manywheel-py3_10-cuda12_8-build
-      - get-label-type
-    uses: ./.github/workflows/_binary-test-linux.yml
-    with:
-      PYTORCH_ROOT: /pytorch
-      PACKAGE_TYPE: manywheel
-      # TODO: This is a legacy variable that we eventually want to get rid of in
-      #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu128
-      GPU_ARCH_VERSION: "12.8"
-      GPU_ARCH_TYPE: cuda
-      DOCKER_IMAGE: manylinux2_28-builder
-      DOCKER_IMAGE_TAG_PREFIX: cuda12.8
-      DESIRED_PYTHON: "3.10"
-      build_name: manywheel-py3_10-cuda12_8
-      build_environment: linux-binary-manywheel
-      runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
-      runs_on: linux.g4dn.4xlarge.nvidia.gpu # 12.8+ builds need sm_70+ runner
-    secrets:
-      github-token: ${{ secrets.GITHUB_TOKEN }}
-  manywheel-py3_10-cuda12_8-upload:  # Uploading
-    if: ${{ github.repository_owner == 'pytorch' }}
-    permissions:
-      id-token: write
-      contents: read
-    needs: manywheel-py3_10-cuda12_8-test
-    with:
-      PYTORCH_ROOT: /pytorch
-      PACKAGE_TYPE: manywheel
-      # TODO: This is a legacy variable that we eventually want to get rid of in
-      #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu128
-      GPU_ARCH_VERSION: "12.8"
-      GPU_ARCH_TYPE: cuda
-      DOCKER_IMAGE: manylinux2_28-builder
-      DOCKER_IMAGE_TAG_PREFIX: cuda12.8
-      DESIRED_PYTHON: "3.10"
-      build_name: manywheel-py3_10-cuda12_8
-    secrets:
-      github-token: ${{ secrets.GITHUB_TOKEN }}
-    uses: ./.github/workflows/_binary-upload.yml
-
-  manywheel-py3_10-cuda12_9-build:
+  manywheel-py3_10-cuda13_0-build:
     if: ${{ github.repository_owner == 'pytorch' }}
     uses: ./.github/workflows/_binary-build-linux.yml
     needs: get-label-type
@@ -253,23 +193,23 @@ jobs:
       PACKAGE_TYPE: manywheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu129
-      GPU_ARCH_VERSION: "12.9"
+      DESIRED_CUDA: cu130
+      GPU_ARCH_VERSION: "13.0"
       GPU_ARCH_TYPE: cuda
       DOCKER_IMAGE: manylinux2_28-builder
-      DOCKER_IMAGE_TAG_PREFIX: cuda12.9
+      DOCKER_IMAGE_TAG_PREFIX: cuda13.0
       DESIRED_PYTHON: "3.10"
       runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
-      build_name: manywheel-py3_10-cuda12_9
+      build_name: manywheel-py3_10-cuda13_0
       build_environment: linux-binary-manywheel
-      PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cublas,cufile,nvjitlink,nvtx]==12.9.1; platform_system == 'Linux' | cuda-bindings==12.9.4; platform_system == 'Linux' | nvidia-cudnn-cu12==9.17.1.4; platform_system == 'Linux' | nvidia-cusparselt-cu12==0.7.1; platform_system == 'Linux' | nvidia-nccl-cu12==2.28.9; platform_system == 'Linux' | nvidia-nvshmem-cu12==3.4.5; platform_system == 'Linux'
+      PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cufile,nvjitlink,nvtx]==13.0.2; platform_system == 'Linux' | nvidia-cublas>=13.1.0.3,<=13.1.1.3; platform_system == 'Linux' | cuda-bindings>=13.0.3,<14; platform_system == 'Linux' | nvidia-cudnn-cu13==9.20.0.48; platform_system == 'Linux' | nvidia-cusparselt-cu13==0.8.1; platform_system == 'Linux' | nvidia-nccl-cu13==2.29.7; platform_system == 'Linux' | nvidia-nvshmem-cu13==3.4.5; platform_system == 'Linux'
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
 
-  manywheel-py3_10-cuda12_9-test:  # Testing
+  manywheel-py3_10-cuda13_0-test:  # Testing
     if: ${{ github.repository_owner == 'pytorch' }}
     needs:
-      - manywheel-py3_10-cuda12_9-build
+      - manywheel-py3_10-cuda13_0-build
       - get-label-type
     uses: ./.github/workflows/_binary-test-linux.yml
     with:
@@ -277,41 +217,44 @@ jobs:
       PACKAGE_TYPE: manywheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu129
-      GPU_ARCH_VERSION: "12.9"
+      DESIRED_CUDA: cu130
+      GPU_ARCH_VERSION: "13.0"
       GPU_ARCH_TYPE: cuda
       DOCKER_IMAGE: manylinux2_28-builder
-      DOCKER_IMAGE_TAG_PREFIX: cuda12.9
+      DOCKER_IMAGE_TAG_PREFIX: cuda13.0
       DESIRED_PYTHON: "3.10"
-      build_name: manywheel-py3_10-cuda12_9
+      build_name: manywheel-py3_10-cuda13_0
       build_environment: linux-binary-manywheel
       runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
-      runs_on: linux.g4dn.4xlarge.nvidia.gpu # 12.8+ builds need sm_70+ runner
+      runs_on: linux.g4dn.4xlarge.nvidia.gpu
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
-  manywheel-py3_10-cuda12_9-upload:  # Uploading
+  manywheel-py3_10-cuda13_0-upload:  # Uploading
     if: ${{ github.repository_owner == 'pytorch' }}
     permissions:
       id-token: write
       contents: read
-    needs: manywheel-py3_10-cuda12_9-test
+    needs: manywheel-py3_10-cuda13_0-test
     with:
       PYTORCH_ROOT: /pytorch
       PACKAGE_TYPE: manywheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu129
-      GPU_ARCH_VERSION: "12.9"
+      DESIRED_CUDA: cu130
+      GPU_ARCH_VERSION: "13.0"
       GPU_ARCH_TYPE: cuda
       DOCKER_IMAGE: manylinux2_28-builder
-      DOCKER_IMAGE_TAG_PREFIX: cuda12.9
+      DOCKER_IMAGE_TAG_PREFIX: cuda13.0
       DESIRED_PYTHON: "3.10"
-      build_name: manywheel-py3_10-cuda12_9
+      build_name: manywheel-py3_10-cuda13_0
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
     uses: ./.github/workflows/_binary-upload.yml
 
-  manywheel-py3_10-cuda13_0-build:
+  manywheel-py3_10-cuda13_2-build:
     if: ${{ github.repository_owner == 'pytorch' }}
     uses: ./.github/workflows/_binary-build-linux.yml
     needs: get-label-type
@@ -320,23 +263,23 @@ jobs:
       PACKAGE_TYPE: manywheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu130
-      GPU_ARCH_VERSION: "13.0"
+      DESIRED_CUDA: cu132
+      GPU_ARCH_VERSION: "13.2"
       GPU_ARCH_TYPE: cuda
       DOCKER_IMAGE: manylinux2_28-builder
-      DOCKER_IMAGE_TAG_PREFIX: cuda13.0
+      DOCKER_IMAGE_TAG_PREFIX: cuda13.2
       DESIRED_PYTHON: "3.10"
       runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
-      build_name: manywheel-py3_10-cuda13_0
+      build_name: manywheel-py3_10-cuda13_2
       build_environment: linux-binary-manywheel
-      PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cublas,cufile,nvjitlink,nvtx]==13.0.2; platform_system == 'Linux' | cuda-bindings==13.0.3; platform_system == 'Linux' | nvidia-cudnn-cu13==9.17.1.4; platform_system == 'Linux' | nvidia-cusparselt-cu13==0.8.0; platform_system == 'Linux' | nvidia-nccl-cu13==2.28.9; platform_system == 'Linux' | nvidia-nvshmem-cu13==3.4.5; platform_system == 'Linux'
+      PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cublas,cufile,nvjitlink,nvtx]==13.2.1; platform_system == 'Linux' | cuda-bindings>=13.0.3,<14; platform_system == 'Linux' | nvidia-cudnn-cu13==9.20.0.48; platform_system == 'Linux' | nvidia-cusparselt-cu13==0.8.1; platform_system == 'Linux' | nvidia-nccl-cu13==2.29.7; platform_system == 'Linux' | nvidia-nvshmem-cu13==3.4.5; platform_system == 'Linux'
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
 
-  manywheel-py3_10-cuda13_0-test:  # Testing
+  manywheel-py3_10-cuda13_2-test:  # Testing
     if: ${{ github.repository_owner == 'pytorch' }}
     needs:
-      - manywheel-py3_10-cuda13_0-build
+      - manywheel-py3_10-cuda13_2-build
       - get-label-type
     uses: ./.github/workflows/_binary-test-linux.yml
     with:
@@ -344,38 +287,41 @@ jobs:
       PACKAGE_TYPE: manywheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu130
-      GPU_ARCH_VERSION: "13.0"
+      DESIRED_CUDA: cu132
+      GPU_ARCH_VERSION: "13.2"
       GPU_ARCH_TYPE: cuda
       DOCKER_IMAGE: manylinux2_28-builder
-      DOCKER_IMAGE_TAG_PREFIX: cuda13.0
+      DOCKER_IMAGE_TAG_PREFIX: cuda13.2
       DESIRED_PYTHON: "3.10"
-      build_name: manywheel-py3_10-cuda13_0
+      build_name: manywheel-py3_10-cuda13_2
       build_environment: linux-binary-manywheel
       runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
-      runs_on: linux.g4dn.4xlarge.nvidia.gpu # 12.8+ builds need sm_70+ runner
+      runs_on: linux.g4dn.4xlarge.nvidia.gpu
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
-  manywheel-py3_10-cuda13_0-upload:  # Uploading
+  manywheel-py3_10-cuda13_2-upload:  # Uploading
     if: ${{ github.repository_owner == 'pytorch' }}
     permissions:
       id-token: write
       contents: read
-    needs: manywheel-py3_10-cuda13_0-test
+    needs: manywheel-py3_10-cuda13_2-test
     with:
       PYTORCH_ROOT: /pytorch
       PACKAGE_TYPE: manywheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu130
-      GPU_ARCH_VERSION: "13.0"
+      DESIRED_CUDA: cu132
+      GPU_ARCH_VERSION: "13.2"
       GPU_ARCH_TYPE: cuda
       DOCKER_IMAGE: manylinux2_28-builder
-      DOCKER_IMAGE_TAG_PREFIX: cuda13.0
+      DOCKER_IMAGE_TAG_PREFIX: cuda13.2
       DESIRED_PYTHON: "3.10"
-      build_name: manywheel-py3_10-cuda13_0
+      build_name: manywheel-py3_10-cuda13_2
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
     uses: ./.github/workflows/_binary-upload.yml
 
   manywheel-py3_10-rocm7_1-build:
@@ -394,7 +340,7 @@ jobs:
       DOCKER_IMAGE_TAG_PREFIX: rocm7.1
       DESIRED_PYTHON: "3.10"
       runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
-      timeout-minutes: 300
+      timeout-minutes: 420
       build_name: manywheel-py3_10-rocm7_1
       build_environment: linux-binary-manywheel
     secrets:
@@ -405,7 +351,7 @@ jobs:
     needs:
       - manywheel-py3_10-rocm7_1-build
       - get-label-type
-    runs-on: linux.rocm.gpu.mi250.1
+    runs-on: linux.rocm.gpu.gfx942.1
     timeout-minutes: 240
     env:
       PYTORCH_ROOT: /pytorch
@@ -423,51 +369,47 @@ jobs:
       id-token: write
       contents: read
     steps:
-      - name: Setup ROCm
-        uses: ./.github/actions/setup-rocm
-      - uses: actions/download-artifact@v4.1.7
-        name: Download Build Artifacts
-        with:
-          name: manywheel-py3_10-rocm7_1
-          path: "${{ runner.temp }}/artifacts/"
       - name: Checkout PyTorch
         uses: actions/checkout@v4
         with:
-          ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
           submodules: recursive
-          path: pytorch
           show-progress: false
       - name: Clean PyTorch checkout
         run: |
           # Remove any artifacts from the previous checkouts
           git clean -fxd
-        working-directory: pytorch
+      - name: Setup ROCm
+        uses: ./.github/actions/setup-rocm
+      - uses: actions/download-artifact@v4.1.7
+        name: Download Build Artifacts
+        with:
+          name: manywheel-py3_10-rocm7_1
+          path: "${{ runner.temp }}/artifacts/"
       - name: ROCm set GPU_FLAG
         run: |
           echo "GPU_FLAG=--device=/dev/mem --device=/dev/kfd --device=/dev/dri --group-add video --group-add daemon" >> "${GITHUB_ENV}"
       - name: configure aws credentials
         id: aws_creds
         if: ${{ startsWith(github.event.ref, 'refs/tags/ciflow/') }}
-        uses: aws-actions/configure-aws-credentials@v4
+        uses: aws-actions/configure-aws-credentials@7474bc4690e29a8392af63c5b98e7449536d5c3a # v4
         with:
           role-to-assume: arn:aws:iam::308535385114:role/gha_workflow_s3_and_ecr_read_only
           aws-region: us-east-1
           role-duration-seconds: 18000
       - name: Calculate docker image
         id: calculate-docker-image
-        uses: pytorch/test-infra/.github/actions/calculate-docker-image@main
+        uses: pytorch/test-infra/.github/actions/calculate-docker-image@release/2.12
         with:
           docker-registry: ${{ startsWith(github.event.ref, 'refs/tags/ciflow/') && '308535385114.dkr.ecr.us-east-1.amazonaws.com' || 'docker.io' }}
           docker-image-name: manylinux2_28-builder
           custom-tag-prefix: rocm7.1
           docker-build-dir: .ci/docker
-          working-directory: pytorch
       - name: Pull Docker image
-        uses: pytorch/test-infra/.github/actions/pull-docker-image@main
+        uses: pytorch/test-infra/.github/actions/pull-docker-image@release/2.12
         with:
           docker-image: ${{ steps.calculate-docker-image.outputs.docker-image }}
       - name: Test Pytorch binary
-        uses: ./pytorch/.github/actions/test-pytorch-binary
+        uses: ./.github/actions/test-pytorch-binary
         env:
           DOCKER_IMAGE: ${{ steps.calculate-docker-image.outputs.docker-image }}
       - name: Teardown ROCm
@@ -492,6 +434,9 @@ jobs:
       build_name: manywheel-py3_10-rocm7_1
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
     uses: ./.github/workflows/_binary-upload.yml
 
   manywheel-py3_10-rocm7_2-build:
@@ -510,7 +455,7 @@ jobs:
       DOCKER_IMAGE_TAG_PREFIX: rocm7.2
       DESIRED_PYTHON: "3.10"
       runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
-      timeout-minutes: 300
+      timeout-minutes: 420
       build_name: manywheel-py3_10-rocm7_2
       build_environment: linux-binary-manywheel
     secrets:
@@ -521,7 +466,7 @@ jobs:
     needs:
       - manywheel-py3_10-rocm7_2-build
       - get-label-type
-    runs-on: linux.rocm.gpu.mi250.1
+    runs-on: linux.rocm.gpu.gfx942.1
     timeout-minutes: 240
     env:
       PYTORCH_ROOT: /pytorch
@@ -539,51 +484,47 @@ jobs:
       id-token: write
       contents: read
     steps:
-      - name: Setup ROCm
-        uses: ./.github/actions/setup-rocm
-      - uses: actions/download-artifact@v4.1.7
-        name: Download Build Artifacts
-        with:
-          name: manywheel-py3_10-rocm7_2
-          path: "${{ runner.temp }}/artifacts/"
       - name: Checkout PyTorch
         uses: actions/checkout@v4
         with:
-          ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
           submodules: recursive
-          path: pytorch
           show-progress: false
       - name: Clean PyTorch checkout
         run: |
           # Remove any artifacts from the previous checkouts
           git clean -fxd
-        working-directory: pytorch
+      - name: Setup ROCm
+        uses: ./.github/actions/setup-rocm
+      - uses: actions/download-artifact@v4.1.7
+        name: Download Build Artifacts
+        with:
+          name: manywheel-py3_10-rocm7_2
+          path: "${{ runner.temp }}/artifacts/"
       - name: ROCm set GPU_FLAG
         run: |
           echo "GPU_FLAG=--device=/dev/mem --device=/dev/kfd --device=/dev/dri --group-add video --group-add daemon" >> "${GITHUB_ENV}"
       - name: configure aws credentials
         id: aws_creds
         if: ${{ startsWith(github.event.ref, 'refs/tags/ciflow/') }}
-        uses: aws-actions/configure-aws-credentials@v4
+        uses: aws-actions/configure-aws-credentials@7474bc4690e29a8392af63c5b98e7449536d5c3a # v4
         with:
           role-to-assume: arn:aws:iam::308535385114:role/gha_workflow_s3_and_ecr_read_only
           aws-region: us-east-1
           role-duration-seconds: 18000
       - name: Calculate docker image
         id: calculate-docker-image
-        uses: pytorch/test-infra/.github/actions/calculate-docker-image@main
+        uses: pytorch/test-infra/.github/actions/calculate-docker-image@release/2.12
         with:
           docker-registry: ${{ startsWith(github.event.ref, 'refs/tags/ciflow/') && '308535385114.dkr.ecr.us-east-1.amazonaws.com' || 'docker.io' }}
           docker-image-name: manylinux2_28-builder
           custom-tag-prefix: rocm7.2
           docker-build-dir: .ci/docker
-          working-directory: pytorch
       - name: Pull Docker image
-        uses: pytorch/test-infra/.github/actions/pull-docker-image@main
+        uses: pytorch/test-infra/.github/actions/pull-docker-image@release/2.12
         with:
           docker-image: ${{ steps.calculate-docker-image.outputs.docker-image }}
       - name: Test Pytorch binary
-        uses: ./pytorch/.github/actions/test-pytorch-binary
+        uses: ./.github/actions/test-pytorch-binary
         env:
           DOCKER_IMAGE: ${{ steps.calculate-docker-image.outputs.docker-image }}
       - name: Teardown ROCm
@@ -608,6 +549,9 @@ jobs:
       build_name: manywheel-py3_10-rocm7_2
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
     uses: ./.github/workflows/_binary-upload.yml
 
   manywheel-py3_10-xpu-build:
@@ -653,42 +597,38 @@ jobs:
       id-token: write
       contents: read
     steps:
-      - name: Setup XPU
-        uses: pytorch/pytorch/.github/actions/setup-xpu@main
-      - name: Login to ECR
-        uses: pytorch/pytorch/.github/actions/ecr-login@main
-      - uses: actions/download-artifact@v4.1.7
-        name: Download Build Artifacts
-        with:
-          name: manywheel-py3_10-xpu
-          path: "${{ runner.temp }}/artifacts/"
       - name: Checkout PyTorch
         uses: actions/checkout@v4
         with:
-          ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
           submodules: recursive
-          path: pytorch
           show-progress: false
       - name: Clean PyTorch checkout
         run: |
           # Remove any artifacts from the previous checkouts
           git clean -fxd
-        working-directory: pytorch
+      - name: Setup XPU
+        uses: pytorch/pytorch/.github/actions/setup-xpu@release/2.12
+      - name: Login to ECR
+        uses: pytorch/pytorch/.github/actions/ecr-login@release/2.12
+      - uses: actions/download-artifact@v4.1.7
+        name: Download Build Artifacts
+        with:
+          name: manywheel-py3_10-xpu
+          path: "${{ runner.temp }}/artifacts/"
       - name: Calculate docker image
         id: calculate-docker-image
-        uses: pytorch/test-infra/.github/actions/calculate-docker-image@main
+        uses: pytorch/test-infra/.github/actions/calculate-docker-image@release/2.12
         with:
           docker-registry: ${{ startsWith(github.event.ref, 'refs/tags/ciflow/') && '308535385114.dkr.ecr.us-east-1.amazonaws.com' || 'docker.io' }}
           docker-image-name: manylinux2_28-builder
           custom-tag-prefix: xpu
           docker-build-dir: .ci/docker
-          working-directory: pytorch
       - name: Pull Docker image
-        uses: pytorch/test-infra/.github/actions/pull-docker-image@main
+        uses: pytorch/test-infra/.github/actions/pull-docker-image@release/2.12
         with:
           docker-image: ${{ steps.calculate-docker-image.outputs.docker-image }}
       - name: Test Pytorch binary
-        uses: ./pytorch/.github/actions/test-pytorch-binary
+        uses: ./.github/actions/test-pytorch-binary
         env:
           DOCKER_IMAGE: ${{ steps.calculate-docker-image.outputs.docker-image }}
       - name: Teardown XPU
@@ -712,6 +652,9 @@ jobs:
       build_name: manywheel-py3_10-xpu
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
     uses: ./.github/workflows/_binary-upload.yml
 
   manywheel-py3_11-cpu-build:
@@ -775,6 +718,9 @@ jobs:
       build_name: manywheel-py3_11-cpu
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
     uses: ./.github/workflows/_binary-upload.yml
 
   manywheel-py3_11-cuda12_6-build:
@@ -795,7 +741,7 @@ jobs:
       runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
       build_name: manywheel-py3_11-cuda12_6
       build_environment: linux-binary-manywheel
-      PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cublas,cufile,nvjitlink,nvtx]==12.6.3; platform_system == 'Linux' | cuda-bindings==12.9.4; platform_system == 'Linux' | nvidia-cudnn-cu12==9.10.2.21; platform_system == 'Linux' | nvidia-cusparselt-cu12==0.7.1; platform_system == 'Linux' | nvidia-nccl-cu12==2.28.9; platform_system == 'Linux' | nvidia-nvshmem-cu12==3.4.5; platform_system == 'Linux'
+      PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cublas,cufile,nvjitlink,nvtx]==12.6.3; platform_system == 'Linux' | cuda-bindings>=12.9.4,<13; platform_system == 'Linux' | nvidia-cudnn-cu12==9.10.2.21; platform_system == 'Linux' | nvidia-cusparselt-cu12==0.7.1; platform_system == 'Linux' | nvidia-nccl-cu12==2.29.3; platform_system == 'Linux' | nvidia-nvshmem-cu12==3.4.5; platform_system == 'Linux'
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
 
@@ -819,7 +765,7 @@ jobs:
       build_name: manywheel-py3_11-cuda12_6
       build_environment: linux-binary-manywheel
       runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
-      runs_on: linux.4xlarge.nvidia.gpu  # 12.6 build can use maxwell (sm_50) runner
+      runs_on: linux.g4dn.4xlarge.nvidia.gpu
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
   manywheel-py3_11-cuda12_6-upload:  # Uploading
@@ -842,76 +788,12 @@ jobs:
       build_name: manywheel-py3_11-cuda12_6
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
     uses: ./.github/workflows/_binary-upload.yml
 
-  manywheel-py3_11-cuda12_8-build:
-    if: ${{ github.repository_owner == 'pytorch' }}
-    uses: ./.github/workflows/_binary-build-linux.yml
-    needs: get-label-type
-    with:
-      PYTORCH_ROOT: /pytorch
-      PACKAGE_TYPE: manywheel
-      # TODO: This is a legacy variable that we eventually want to get rid of in
-      #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu128
-      GPU_ARCH_VERSION: "12.8"
-      GPU_ARCH_TYPE: cuda
-      DOCKER_IMAGE: manylinux2_28-builder
-      DOCKER_IMAGE_TAG_PREFIX: cuda12.8
-      DESIRED_PYTHON: "3.11"
-      runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
-      build_name: manywheel-py3_11-cuda12_8
-      build_environment: linux-binary-manywheel
-      PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cublas,cufile,nvjitlink,nvtx]==12.8.1; platform_system == 'Linux' | cuda-bindings==12.9.4; platform_system == 'Linux' | nvidia-cudnn-cu12==9.17.1.4; platform_system == 'Linux' | nvidia-cusparselt-cu12==0.7.1; platform_system == 'Linux' | nvidia-nccl-cu12==2.28.9; platform_system == 'Linux' | nvidia-nvshmem-cu12==3.4.5; platform_system == 'Linux'
-    secrets:
-      github-token: ${{ secrets.GITHUB_TOKEN }}
-
-  manywheel-py3_11-cuda12_8-test:  # Testing
-    if: ${{ github.repository_owner == 'pytorch' }}
-    needs:
-      - manywheel-py3_11-cuda12_8-build
-      - get-label-type
-    uses: ./.github/workflows/_binary-test-linux.yml
-    with:
-      PYTORCH_ROOT: /pytorch
-      PACKAGE_TYPE: manywheel
-      # TODO: This is a legacy variable that we eventually want to get rid of in
-      #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu128
-      GPU_ARCH_VERSION: "12.8"
-      GPU_ARCH_TYPE: cuda
-      DOCKER_IMAGE: manylinux2_28-builder
-      DOCKER_IMAGE_TAG_PREFIX: cuda12.8
-      DESIRED_PYTHON: "3.11"
-      build_name: manywheel-py3_11-cuda12_8
-      build_environment: linux-binary-manywheel
-      runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
-      runs_on: linux.g4dn.4xlarge.nvidia.gpu # 12.8+ builds need sm_70+ runner
-    secrets:
-      github-token: ${{ secrets.GITHUB_TOKEN }}
-  manywheel-py3_11-cuda12_8-upload:  # Uploading
-    if: ${{ github.repository_owner == 'pytorch' }}
-    permissions:
-      id-token: write
-      contents: read
-    needs: manywheel-py3_11-cuda12_8-test
-    with:
-      PYTORCH_ROOT: /pytorch
-      PACKAGE_TYPE: manywheel
-      # TODO: This is a legacy variable that we eventually want to get rid of in
-      #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu128
-      GPU_ARCH_VERSION: "12.8"
-      GPU_ARCH_TYPE: cuda
-      DOCKER_IMAGE: manylinux2_28-builder
-      DOCKER_IMAGE_TAG_PREFIX: cuda12.8
-      DESIRED_PYTHON: "3.11"
-      build_name: manywheel-py3_11-cuda12_8
-    secrets:
-      github-token: ${{ secrets.GITHUB_TOKEN }}
-    uses: ./.github/workflows/_binary-upload.yml
-
-  manywheel-py3_11-cuda12_9-build:
+  manywheel-py3_11-cuda13_0-build:
     if: ${{ github.repository_owner == 'pytorch' }}
     uses: ./.github/workflows/_binary-build-linux.yml
     needs: get-label-type
@@ -920,23 +802,23 @@ jobs:
       PACKAGE_TYPE: manywheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu129
-      GPU_ARCH_VERSION: "12.9"
+      DESIRED_CUDA: cu130
+      GPU_ARCH_VERSION: "13.0"
       GPU_ARCH_TYPE: cuda
       DOCKER_IMAGE: manylinux2_28-builder
-      DOCKER_IMAGE_TAG_PREFIX: cuda12.9
+      DOCKER_IMAGE_TAG_PREFIX: cuda13.0
       DESIRED_PYTHON: "3.11"
       runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
-      build_name: manywheel-py3_11-cuda12_9
+      build_name: manywheel-py3_11-cuda13_0
       build_environment: linux-binary-manywheel
-      PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cublas,cufile,nvjitlink,nvtx]==12.9.1; platform_system == 'Linux' | cuda-bindings==12.9.4; platform_system == 'Linux' | nvidia-cudnn-cu12==9.17.1.4; platform_system == 'Linux' | nvidia-cusparselt-cu12==0.7.1; platform_system == 'Linux' | nvidia-nccl-cu12==2.28.9; platform_system == 'Linux' | nvidia-nvshmem-cu12==3.4.5; platform_system == 'Linux'
+      PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cufile,nvjitlink,nvtx]==13.0.2; platform_system == 'Linux' | nvidia-cublas>=13.1.0.3,<=13.1.1.3; platform_system == 'Linux' | cuda-bindings>=13.0.3,<14; platform_system == 'Linux' | nvidia-cudnn-cu13==9.20.0.48; platform_system == 'Linux' | nvidia-cusparselt-cu13==0.8.1; platform_system == 'Linux' | nvidia-nccl-cu13==2.29.7; platform_system == 'Linux' | nvidia-nvshmem-cu13==3.4.5; platform_system == 'Linux'
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
 
-  manywheel-py3_11-cuda12_9-test:  # Testing
+  manywheel-py3_11-cuda13_0-test:  # Testing
     if: ${{ github.repository_owner == 'pytorch' }}
     needs:
-      - manywheel-py3_11-cuda12_9-build
+      - manywheel-py3_11-cuda13_0-build
       - get-label-type
     uses: ./.github/workflows/_binary-test-linux.yml
     with:
@@ -944,41 +826,44 @@ jobs:
       PACKAGE_TYPE: manywheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu129
-      GPU_ARCH_VERSION: "12.9"
+      DESIRED_CUDA: cu130
+      GPU_ARCH_VERSION: "13.0"
       GPU_ARCH_TYPE: cuda
       DOCKER_IMAGE: manylinux2_28-builder
-      DOCKER_IMAGE_TAG_PREFIX: cuda12.9
+      DOCKER_IMAGE_TAG_PREFIX: cuda13.0
       DESIRED_PYTHON: "3.11"
-      build_name: manywheel-py3_11-cuda12_9
+      build_name: manywheel-py3_11-cuda13_0
       build_environment: linux-binary-manywheel
       runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
-      runs_on: linux.g4dn.4xlarge.nvidia.gpu # 12.8+ builds need sm_70+ runner
+      runs_on: linux.g4dn.4xlarge.nvidia.gpu
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
-  manywheel-py3_11-cuda12_9-upload:  # Uploading
+  manywheel-py3_11-cuda13_0-upload:  # Uploading
     if: ${{ github.repository_owner == 'pytorch' }}
     permissions:
       id-token: write
       contents: read
-    needs: manywheel-py3_11-cuda12_9-test
+    needs: manywheel-py3_11-cuda13_0-test
     with:
       PYTORCH_ROOT: /pytorch
       PACKAGE_TYPE: manywheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu129
-      GPU_ARCH_VERSION: "12.9"
+      DESIRED_CUDA: cu130
+      GPU_ARCH_VERSION: "13.0"
       GPU_ARCH_TYPE: cuda
       DOCKER_IMAGE: manylinux2_28-builder
-      DOCKER_IMAGE_TAG_PREFIX: cuda12.9
+      DOCKER_IMAGE_TAG_PREFIX: cuda13.0
       DESIRED_PYTHON: "3.11"
-      build_name: manywheel-py3_11-cuda12_9
+      build_name: manywheel-py3_11-cuda13_0
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
     uses: ./.github/workflows/_binary-upload.yml
 
-  manywheel-py3_11-cuda13_0-build:
+  manywheel-py3_11-cuda13_2-build:
     if: ${{ github.repository_owner == 'pytorch' }}
     uses: ./.github/workflows/_binary-build-linux.yml
     needs: get-label-type
@@ -987,23 +872,23 @@ jobs:
       PACKAGE_TYPE: manywheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu130
-      GPU_ARCH_VERSION: "13.0"
+      DESIRED_CUDA: cu132
+      GPU_ARCH_VERSION: "13.2"
       GPU_ARCH_TYPE: cuda
       DOCKER_IMAGE: manylinux2_28-builder
-      DOCKER_IMAGE_TAG_PREFIX: cuda13.0
+      DOCKER_IMAGE_TAG_PREFIX: cuda13.2
       DESIRED_PYTHON: "3.11"
       runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
-      build_name: manywheel-py3_11-cuda13_0
+      build_name: manywheel-py3_11-cuda13_2
       build_environment: linux-binary-manywheel
-      PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cublas,cufile,nvjitlink,nvtx]==13.0.2; platform_system == 'Linux' | cuda-bindings==13.0.3; platform_system == 'Linux' | nvidia-cudnn-cu13==9.17.1.4; platform_system == 'Linux' | nvidia-cusparselt-cu13==0.8.0; platform_system == 'Linux' | nvidia-nccl-cu13==2.28.9; platform_system == 'Linux' | nvidia-nvshmem-cu13==3.4.5; platform_system == 'Linux'
+      PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cublas,cufile,nvjitlink,nvtx]==13.2.1; platform_system == 'Linux' | cuda-bindings>=13.0.3,<14; platform_system == 'Linux' | nvidia-cudnn-cu13==9.20.0.48; platform_system == 'Linux' | nvidia-cusparselt-cu13==0.8.1; platform_system == 'Linux' | nvidia-nccl-cu13==2.29.7; platform_system == 'Linux' | nvidia-nvshmem-cu13==3.4.5; platform_system == 'Linux'
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
 
-  manywheel-py3_11-cuda13_0-test:  # Testing
+  manywheel-py3_11-cuda13_2-test:  # Testing
     if: ${{ github.repository_owner == 'pytorch' }}
     needs:
-      - manywheel-py3_11-cuda13_0-build
+      - manywheel-py3_11-cuda13_2-build
       - get-label-type
     uses: ./.github/workflows/_binary-test-linux.yml
     with:
@@ -1011,38 +896,41 @@ jobs:
       PACKAGE_TYPE: manywheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu130
-      GPU_ARCH_VERSION: "13.0"
+      DESIRED_CUDA: cu132
+      GPU_ARCH_VERSION: "13.2"
       GPU_ARCH_TYPE: cuda
       DOCKER_IMAGE: manylinux2_28-builder
-      DOCKER_IMAGE_TAG_PREFIX: cuda13.0
+      DOCKER_IMAGE_TAG_PREFIX: cuda13.2
       DESIRED_PYTHON: "3.11"
-      build_name: manywheel-py3_11-cuda13_0
+      build_name: manywheel-py3_11-cuda13_2
       build_environment: linux-binary-manywheel
       runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
-      runs_on: linux.g4dn.4xlarge.nvidia.gpu # 12.8+ builds need sm_70+ runner
+      runs_on: linux.g4dn.4xlarge.nvidia.gpu
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
-  manywheel-py3_11-cuda13_0-upload:  # Uploading
+  manywheel-py3_11-cuda13_2-upload:  # Uploading
     if: ${{ github.repository_owner == 'pytorch' }}
     permissions:
       id-token: write
       contents: read
-    needs: manywheel-py3_11-cuda13_0-test
+    needs: manywheel-py3_11-cuda13_2-test
     with:
       PYTORCH_ROOT: /pytorch
       PACKAGE_TYPE: manywheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu130
-      GPU_ARCH_VERSION: "13.0"
+      DESIRED_CUDA: cu132
+      GPU_ARCH_VERSION: "13.2"
       GPU_ARCH_TYPE: cuda
       DOCKER_IMAGE: manylinux2_28-builder
-      DOCKER_IMAGE_TAG_PREFIX: cuda13.0
+      DOCKER_IMAGE_TAG_PREFIX: cuda13.2
       DESIRED_PYTHON: "3.11"
-      build_name: manywheel-py3_11-cuda13_0
+      build_name: manywheel-py3_11-cuda13_2
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
     uses: ./.github/workflows/_binary-upload.yml
 
   manywheel-py3_11-rocm7_1-build:
@@ -1061,7 +949,7 @@ jobs:
       DOCKER_IMAGE_TAG_PREFIX: rocm7.1
       DESIRED_PYTHON: "3.11"
       runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
-      timeout-minutes: 300
+      timeout-minutes: 420
       build_name: manywheel-py3_11-rocm7_1
       build_environment: linux-binary-manywheel
     secrets:
@@ -1072,7 +960,7 @@ jobs:
     needs:
       - manywheel-py3_11-rocm7_1-build
       - get-label-type
-    runs-on: linux.rocm.gpu.mi250.1
+    runs-on: linux.rocm.gpu.gfx942.1
     timeout-minutes: 240
     env:
       PYTORCH_ROOT: /pytorch
@@ -1090,51 +978,47 @@ jobs:
       id-token: write
       contents: read
     steps:
-      - name: Setup ROCm
-        uses: ./.github/actions/setup-rocm
-      - uses: actions/download-artifact@v4.1.7
-        name: Download Build Artifacts
-        with:
-          name: manywheel-py3_11-rocm7_1
-          path: "${{ runner.temp }}/artifacts/"
       - name: Checkout PyTorch
         uses: actions/checkout@v4
         with:
-          ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
           submodules: recursive
-          path: pytorch
           show-progress: false
       - name: Clean PyTorch checkout
         run: |
           # Remove any artifacts from the previous checkouts
           git clean -fxd
-        working-directory: pytorch
+      - name: Setup ROCm
+        uses: ./.github/actions/setup-rocm
+      - uses: actions/download-artifact@v4.1.7
+        name: Download Build Artifacts
+        with:
+          name: manywheel-py3_11-rocm7_1
+          path: "${{ runner.temp }}/artifacts/"
       - name: ROCm set GPU_FLAG
         run: |
           echo "GPU_FLAG=--device=/dev/mem --device=/dev/kfd --device=/dev/dri --group-add video --group-add daemon" >> "${GITHUB_ENV}"
       - name: configure aws credentials
         id: aws_creds
         if: ${{ startsWith(github.event.ref, 'refs/tags/ciflow/') }}
-        uses: aws-actions/configure-aws-credentials@v4
+        uses: aws-actions/configure-aws-credentials@7474bc4690e29a8392af63c5b98e7449536d5c3a # v4
         with:
           role-to-assume: arn:aws:iam::308535385114:role/gha_workflow_s3_and_ecr_read_only
           aws-region: us-east-1
           role-duration-seconds: 18000
       - name: Calculate docker image
         id: calculate-docker-image
-        uses: pytorch/test-infra/.github/actions/calculate-docker-image@main
+        uses: pytorch/test-infra/.github/actions/calculate-docker-image@release/2.12
         with:
           docker-registry: ${{ startsWith(github.event.ref, 'refs/tags/ciflow/') && '308535385114.dkr.ecr.us-east-1.amazonaws.com' || 'docker.io' }}
           docker-image-name: manylinux2_28-builder
           custom-tag-prefix: rocm7.1
           docker-build-dir: .ci/docker
-          working-directory: pytorch
       - name: Pull Docker image
-        uses: pytorch/test-infra/.github/actions/pull-docker-image@main
+        uses: pytorch/test-infra/.github/actions/pull-docker-image@release/2.12
         with:
           docker-image: ${{ steps.calculate-docker-image.outputs.docker-image }}
       - name: Test Pytorch binary
-        uses: ./pytorch/.github/actions/test-pytorch-binary
+        uses: ./.github/actions/test-pytorch-binary
         env:
           DOCKER_IMAGE: ${{ steps.calculate-docker-image.outputs.docker-image }}
       - name: Teardown ROCm
@@ -1159,6 +1043,9 @@ jobs:
       build_name: manywheel-py3_11-rocm7_1
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
     uses: ./.github/workflows/_binary-upload.yml
 
   manywheel-py3_11-rocm7_2-build:
@@ -1177,7 +1064,7 @@ jobs:
       DOCKER_IMAGE_TAG_PREFIX: rocm7.2
       DESIRED_PYTHON: "3.11"
       runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
-      timeout-minutes: 300
+      timeout-minutes: 420
       build_name: manywheel-py3_11-rocm7_2
       build_environment: linux-binary-manywheel
     secrets:
@@ -1188,7 +1075,7 @@ jobs:
     needs:
       - manywheel-py3_11-rocm7_2-build
       - get-label-type
-    runs-on: linux.rocm.gpu.mi250.1
+    runs-on: linux.rocm.gpu.gfx942.1
     timeout-minutes: 240
     env:
       PYTORCH_ROOT: /pytorch
@@ -1206,51 +1093,47 @@ jobs:
       id-token: write
       contents: read
     steps:
-      - name: Setup ROCm
-        uses: ./.github/actions/setup-rocm
-      - uses: actions/download-artifact@v4.1.7
-        name: Download Build Artifacts
-        with:
-          name: manywheel-py3_11-rocm7_2
-          path: "${{ runner.temp }}/artifacts/"
       - name: Checkout PyTorch
         uses: actions/checkout@v4
         with:
-          ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
           submodules: recursive
-          path: pytorch
           show-progress: false
       - name: Clean PyTorch checkout
         run: |
           # Remove any artifacts from the previous checkouts
           git clean -fxd
-        working-directory: pytorch
+      - name: Setup ROCm
+        uses: ./.github/actions/setup-rocm
+      - uses: actions/download-artifact@v4.1.7
+        name: Download Build Artifacts
+        with:
+          name: manywheel-py3_11-rocm7_2
+          path: "${{ runner.temp }}/artifacts/"
       - name: ROCm set GPU_FLAG
         run: |
           echo "GPU_FLAG=--device=/dev/mem --device=/dev/kfd --device=/dev/dri --group-add video --group-add daemon" >> "${GITHUB_ENV}"
       - name: configure aws credentials
         id: aws_creds
         if: ${{ startsWith(github.event.ref, 'refs/tags/ciflow/') }}
-        uses: aws-actions/configure-aws-credentials@v4
+        uses: aws-actions/configure-aws-credentials@7474bc4690e29a8392af63c5b98e7449536d5c3a # v4
         with:
           role-to-assume: arn:aws:iam::308535385114:role/gha_workflow_s3_and_ecr_read_only
           aws-region: us-east-1
           role-duration-seconds: 18000
       - name: Calculate docker image
         id: calculate-docker-image
-        uses: pytorch/test-infra/.github/actions/calculate-docker-image@main
+        uses: pytorch/test-infra/.github/actions/calculate-docker-image@release/2.12
         with:
           docker-registry: ${{ startsWith(github.event.ref, 'refs/tags/ciflow/') && '308535385114.dkr.ecr.us-east-1.amazonaws.com' || 'docker.io' }}
           docker-image-name: manylinux2_28-builder
           custom-tag-prefix: rocm7.2
           docker-build-dir: .ci/docker
-          working-directory: pytorch
       - name: Pull Docker image
-        uses: pytorch/test-infra/.github/actions/pull-docker-image@main
+        uses: pytorch/test-infra/.github/actions/pull-docker-image@release/2.12
         with:
           docker-image: ${{ steps.calculate-docker-image.outputs.docker-image }}
       - name: Test Pytorch binary
-        uses: ./pytorch/.github/actions/test-pytorch-binary
+        uses: ./.github/actions/test-pytorch-binary
         env:
           DOCKER_IMAGE: ${{ steps.calculate-docker-image.outputs.docker-image }}
       - name: Teardown ROCm
@@ -1275,6 +1158,9 @@ jobs:
       build_name: manywheel-py3_11-rocm7_2
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
     uses: ./.github/workflows/_binary-upload.yml
 
   manywheel-py3_11-xpu-build:
@@ -1320,42 +1206,38 @@ jobs:
       id-token: write
       contents: read
     steps:
-      - name: Setup XPU
-        uses: pytorch/pytorch/.github/actions/setup-xpu@main
-      - name: Login to ECR
-        uses: pytorch/pytorch/.github/actions/ecr-login@main
-      - uses: actions/download-artifact@v4.1.7
-        name: Download Build Artifacts
-        with:
-          name: manywheel-py3_11-xpu
-          path: "${{ runner.temp }}/artifacts/"
       - name: Checkout PyTorch
         uses: actions/checkout@v4
         with:
-          ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
           submodules: recursive
-          path: pytorch
           show-progress: false
       - name: Clean PyTorch checkout
         run: |
           # Remove any artifacts from the previous checkouts
           git clean -fxd
-        working-directory: pytorch
+      - name: Setup XPU
+        uses: pytorch/pytorch/.github/actions/setup-xpu@release/2.12
+      - name: Login to ECR
+        uses: pytorch/pytorch/.github/actions/ecr-login@release/2.12
+      - uses: actions/download-artifact@v4.1.7
+        name: Download Build Artifacts
+        with:
+          name: manywheel-py3_11-xpu
+          path: "${{ runner.temp }}/artifacts/"
       - name: Calculate docker image
         id: calculate-docker-image
-        uses: pytorch/test-infra/.github/actions/calculate-docker-image@main
+        uses: pytorch/test-infra/.github/actions/calculate-docker-image@release/2.12
         with:
           docker-registry: ${{ startsWith(github.event.ref, 'refs/tags/ciflow/') && '308535385114.dkr.ecr.us-east-1.amazonaws.com' || 'docker.io' }}
           docker-image-name: manylinux2_28-builder
           custom-tag-prefix: xpu
           docker-build-dir: .ci/docker
-          working-directory: pytorch
       - name: Pull Docker image
-        uses: pytorch/test-infra/.github/actions/pull-docker-image@main
+        uses: pytorch/test-infra/.github/actions/pull-docker-image@release/2.12
         with:
           docker-image: ${{ steps.calculate-docker-image.outputs.docker-image }}
       - name: Test Pytorch binary
-        uses: ./pytorch/.github/actions/test-pytorch-binary
+        uses: ./.github/actions/test-pytorch-binary
         env:
           DOCKER_IMAGE: ${{ steps.calculate-docker-image.outputs.docker-image }}
       - name: Teardown XPU
@@ -1379,6 +1261,9 @@ jobs:
       build_name: manywheel-py3_11-xpu
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
     uses: ./.github/workflows/_binary-upload.yml
 
   manywheel-py3_12-cpu-build:
@@ -1442,6 +1327,9 @@ jobs:
       build_name: manywheel-py3_12-cpu
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
     uses: ./.github/workflows/_binary-upload.yml
 
   manywheel-py3_12-cuda12_6-build:
@@ -1462,7 +1350,7 @@ jobs:
       runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
       build_name: manywheel-py3_12-cuda12_6
       build_environment: linux-binary-manywheel
-      PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cublas,cufile,nvjitlink,nvtx]==12.6.3; platform_system == 'Linux' | cuda-bindings==12.9.4; platform_system == 'Linux' | nvidia-cudnn-cu12==9.10.2.21; platform_system == 'Linux' | nvidia-cusparselt-cu12==0.7.1; platform_system == 'Linux' | nvidia-nccl-cu12==2.28.9; platform_system == 'Linux' | nvidia-nvshmem-cu12==3.4.5; platform_system == 'Linux'
+      PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cublas,cufile,nvjitlink,nvtx]==12.6.3; platform_system == 'Linux' | cuda-bindings>=12.9.4,<13; platform_system == 'Linux' | nvidia-cudnn-cu12==9.10.2.21; platform_system == 'Linux' | nvidia-cusparselt-cu12==0.7.1; platform_system == 'Linux' | nvidia-nccl-cu12==2.29.3; platform_system == 'Linux' | nvidia-nvshmem-cu12==3.4.5; platform_system == 'Linux'
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
 
@@ -1486,7 +1374,7 @@ jobs:
       build_name: manywheel-py3_12-cuda12_6
       build_environment: linux-binary-manywheel
       runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
-      runs_on: linux.4xlarge.nvidia.gpu  # 12.6 build can use maxwell (sm_50) runner
+      runs_on: linux.g4dn.4xlarge.nvidia.gpu
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
   manywheel-py3_12-cuda12_6-upload:  # Uploading
@@ -1509,9 +1397,12 @@ jobs:
       build_name: manywheel-py3_12-cuda12_6
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
     uses: ./.github/workflows/_binary-upload.yml
 
-  manywheel-py3_12-cuda12_8-build:
+  manywheel-py3_12-cuda13_0-build:
     if: ${{ github.repository_owner == 'pytorch' }}
     uses: ./.github/workflows/_binary-build-linux.yml
     needs: get-label-type
@@ -1520,90 +1411,23 @@ jobs:
       PACKAGE_TYPE: manywheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu128
-      GPU_ARCH_VERSION: "12.8"
-      GPU_ARCH_TYPE: cuda
-      DOCKER_IMAGE: manylinux2_28-builder
-      DOCKER_IMAGE_TAG_PREFIX: cuda12.8
-      DESIRED_PYTHON: "3.12"
-      runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
-      build_name: manywheel-py3_12-cuda12_8
-      build_environment: linux-binary-manywheel
-      PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cublas,cufile,nvjitlink,nvtx]==12.8.1; platform_system == 'Linux' | cuda-bindings==12.9.4; platform_system == 'Linux' | nvidia-cudnn-cu12==9.17.1.4; platform_system == 'Linux' | nvidia-cusparselt-cu12==0.7.1; platform_system == 'Linux' | nvidia-nccl-cu12==2.28.9; platform_system == 'Linux' | nvidia-nvshmem-cu12==3.4.5; platform_system == 'Linux'
-    secrets:
-      github-token: ${{ secrets.GITHUB_TOKEN }}
-
-  manywheel-py3_12-cuda12_8-test:  # Testing
-    if: ${{ github.repository_owner == 'pytorch' }}
-    needs:
-      - manywheel-py3_12-cuda12_8-build
-      - get-label-type
-    uses: ./.github/workflows/_binary-test-linux.yml
-    with:
-      PYTORCH_ROOT: /pytorch
-      PACKAGE_TYPE: manywheel
-      # TODO: This is a legacy variable that we eventually want to get rid of in
-      #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu128
-      GPU_ARCH_VERSION: "12.8"
-      GPU_ARCH_TYPE: cuda
-      DOCKER_IMAGE: manylinux2_28-builder
-      DOCKER_IMAGE_TAG_PREFIX: cuda12.8
-      DESIRED_PYTHON: "3.12"
-      build_name: manywheel-py3_12-cuda12_8
-      build_environment: linux-binary-manywheel
-      runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
-      runs_on: linux.g4dn.4xlarge.nvidia.gpu # 12.8+ builds need sm_70+ runner
-    secrets:
-      github-token: ${{ secrets.GITHUB_TOKEN }}
-  manywheel-py3_12-cuda12_8-upload:  # Uploading
-    if: ${{ github.repository_owner == 'pytorch' }}
-    permissions:
-      id-token: write
-      contents: read
-    needs: manywheel-py3_12-cuda12_8-test
-    with:
-      PYTORCH_ROOT: /pytorch
-      PACKAGE_TYPE: manywheel
-      # TODO: This is a legacy variable that we eventually want to get rid of in
-      #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu128
-      GPU_ARCH_VERSION: "12.8"
-      GPU_ARCH_TYPE: cuda
-      DOCKER_IMAGE: manylinux2_28-builder
-      DOCKER_IMAGE_TAG_PREFIX: cuda12.8
-      DESIRED_PYTHON: "3.12"
-      build_name: manywheel-py3_12-cuda12_8
-    secrets:
-      github-token: ${{ secrets.GITHUB_TOKEN }}
-    uses: ./.github/workflows/_binary-upload.yml
-
-  manywheel-py3_12-cuda12_9-build:
-    if: ${{ github.repository_owner == 'pytorch' }}
-    uses: ./.github/workflows/_binary-build-linux.yml
-    needs: get-label-type
-    with:
-      PYTORCH_ROOT: /pytorch
-      PACKAGE_TYPE: manywheel
-      # TODO: This is a legacy variable that we eventually want to get rid of in
-      #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu129
-      GPU_ARCH_VERSION: "12.9"
+      DESIRED_CUDA: cu130
+      GPU_ARCH_VERSION: "13.0"
       GPU_ARCH_TYPE: cuda
       DOCKER_IMAGE: manylinux2_28-builder
-      DOCKER_IMAGE_TAG_PREFIX: cuda12.9
+      DOCKER_IMAGE_TAG_PREFIX: cuda13.0
       DESIRED_PYTHON: "3.12"
       runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
-      build_name: manywheel-py3_12-cuda12_9
+      build_name: manywheel-py3_12-cuda13_0
       build_environment: linux-binary-manywheel
-      PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cublas,cufile,nvjitlink,nvtx]==12.9.1; platform_system == 'Linux' | cuda-bindings==12.9.4; platform_system == 'Linux' | nvidia-cudnn-cu12==9.17.1.4; platform_system == 'Linux' | nvidia-cusparselt-cu12==0.7.1; platform_system == 'Linux' | nvidia-nccl-cu12==2.28.9; platform_system == 'Linux' | nvidia-nvshmem-cu12==3.4.5; platform_system == 'Linux'
+      PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cufile,nvjitlink,nvtx]==13.0.2; platform_system == 'Linux' | nvidia-cublas>=13.1.0.3,<=13.1.1.3; platform_system == 'Linux' | cuda-bindings>=13.0.3,<14; platform_system == 'Linux' | nvidia-cudnn-cu13==9.20.0.48; platform_system == 'Linux' | nvidia-cusparselt-cu13==0.8.1; platform_system == 'Linux' | nvidia-nccl-cu13==2.29.7; platform_system == 'Linux' | nvidia-nvshmem-cu13==3.4.5; platform_system == 'Linux'
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
 
-  manywheel-py3_12-cuda12_9-test:  # Testing
+  manywheel-py3_12-cuda13_0-test:  # Testing
     if: ${{ github.repository_owner == 'pytorch' }}
     needs:
-      - manywheel-py3_12-cuda12_9-build
+      - manywheel-py3_12-cuda13_0-build
       - get-label-type
     uses: ./.github/workflows/_binary-test-linux.yml
     with:
@@ -1611,41 +1435,44 @@ jobs:
       PACKAGE_TYPE: manywheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu129
-      GPU_ARCH_VERSION: "12.9"
+      DESIRED_CUDA: cu130
+      GPU_ARCH_VERSION: "13.0"
       GPU_ARCH_TYPE: cuda
       DOCKER_IMAGE: manylinux2_28-builder
-      DOCKER_IMAGE_TAG_PREFIX: cuda12.9
+      DOCKER_IMAGE_TAG_PREFIX: cuda13.0
       DESIRED_PYTHON: "3.12"
-      build_name: manywheel-py3_12-cuda12_9
+      build_name: manywheel-py3_12-cuda13_0
       build_environment: linux-binary-manywheel
       runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
-      runs_on: linux.g4dn.4xlarge.nvidia.gpu # 12.8+ builds need sm_70+ runner
+      runs_on: linux.g4dn.4xlarge.nvidia.gpu
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
-  manywheel-py3_12-cuda12_9-upload:  # Uploading
+  manywheel-py3_12-cuda13_0-upload:  # Uploading
     if: ${{ github.repository_owner == 'pytorch' }}
     permissions:
       id-token: write
       contents: read
-    needs: manywheel-py3_12-cuda12_9-test
+    needs: manywheel-py3_12-cuda13_0-test
     with:
       PYTORCH_ROOT: /pytorch
       PACKAGE_TYPE: manywheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu129
-      GPU_ARCH_VERSION: "12.9"
+      DESIRED_CUDA: cu130
+      GPU_ARCH_VERSION: "13.0"
       GPU_ARCH_TYPE: cuda
       DOCKER_IMAGE: manylinux2_28-builder
-      DOCKER_IMAGE_TAG_PREFIX: cuda12.9
+      DOCKER_IMAGE_TAG_PREFIX: cuda13.0
       DESIRED_PYTHON: "3.12"
-      build_name: manywheel-py3_12-cuda12_9
+      build_name: manywheel-py3_12-cuda13_0
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
     uses: ./.github/workflows/_binary-upload.yml
 
-  manywheel-py3_12-cuda13_0-build:
+  manywheel-py3_12-cuda13_2-build:
     if: ${{ github.repository_owner == 'pytorch' }}
     uses: ./.github/workflows/_binary-build-linux.yml
     needs: get-label-type
@@ -1654,23 +1481,23 @@ jobs:
       PACKAGE_TYPE: manywheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu130
-      GPU_ARCH_VERSION: "13.0"
+      DESIRED_CUDA: cu132
+      GPU_ARCH_VERSION: "13.2"
       GPU_ARCH_TYPE: cuda
       DOCKER_IMAGE: manylinux2_28-builder
-      DOCKER_IMAGE_TAG_PREFIX: cuda13.0
+      DOCKER_IMAGE_TAG_PREFIX: cuda13.2
       DESIRED_PYTHON: "3.12"
       runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
-      build_name: manywheel-py3_12-cuda13_0
+      build_name: manywheel-py3_12-cuda13_2
       build_environment: linux-binary-manywheel
-      PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cublas,cufile,nvjitlink,nvtx]==13.0.2; platform_system == 'Linux' | cuda-bindings==13.0.3; platform_system == 'Linux' | nvidia-cudnn-cu13==9.17.1.4; platform_system == 'Linux' | nvidia-cusparselt-cu13==0.8.0; platform_system == 'Linux' | nvidia-nccl-cu13==2.28.9; platform_system == 'Linux' | nvidia-nvshmem-cu13==3.4.5; platform_system == 'Linux'
+      PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cublas,cufile,nvjitlink,nvtx]==13.2.1; platform_system == 'Linux' | cuda-bindings>=13.0.3,<14; platform_system == 'Linux' | nvidia-cudnn-cu13==9.20.0.48; platform_system == 'Linux' | nvidia-cusparselt-cu13==0.8.1; platform_system == 'Linux' | nvidia-nccl-cu13==2.29.7; platform_system == 'Linux' | nvidia-nvshmem-cu13==3.4.5; platform_system == 'Linux'
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
 
-  manywheel-py3_12-cuda13_0-test:  # Testing
+  manywheel-py3_12-cuda13_2-test:  # Testing
     if: ${{ github.repository_owner == 'pytorch' }}
     needs:
-      - manywheel-py3_12-cuda13_0-build
+      - manywheel-py3_12-cuda13_2-build
       - get-label-type
     uses: ./.github/workflows/_binary-test-linux.yml
     with:
@@ -1678,38 +1505,41 @@ jobs:
       PACKAGE_TYPE: manywheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu130
-      GPU_ARCH_VERSION: "13.0"
+      DESIRED_CUDA: cu132
+      GPU_ARCH_VERSION: "13.2"
       GPU_ARCH_TYPE: cuda
       DOCKER_IMAGE: manylinux2_28-builder
-      DOCKER_IMAGE_TAG_PREFIX: cuda13.0
+      DOCKER_IMAGE_TAG_PREFIX: cuda13.2
       DESIRED_PYTHON: "3.12"
-      build_name: manywheel-py3_12-cuda13_0
+      build_name: manywheel-py3_12-cuda13_2
       build_environment: linux-binary-manywheel
       runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
-      runs_on: linux.g4dn.4xlarge.nvidia.gpu # 12.8+ builds need sm_70+ runner
+      runs_on: linux.g4dn.4xlarge.nvidia.gpu
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
-  manywheel-py3_12-cuda13_0-upload:  # Uploading
+  manywheel-py3_12-cuda13_2-upload:  # Uploading
     if: ${{ github.repository_owner == 'pytorch' }}
     permissions:
       id-token: write
       contents: read
-    needs: manywheel-py3_12-cuda13_0-test
+    needs: manywheel-py3_12-cuda13_2-test
     with:
       PYTORCH_ROOT: /pytorch
       PACKAGE_TYPE: manywheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu130
-      GPU_ARCH_VERSION: "13.0"
+      DESIRED_CUDA: cu132
+      GPU_ARCH_VERSION: "13.2"
       GPU_ARCH_TYPE: cuda
       DOCKER_IMAGE: manylinux2_28-builder
-      DOCKER_IMAGE_TAG_PREFIX: cuda13.0
+      DOCKER_IMAGE_TAG_PREFIX: cuda13.2
       DESIRED_PYTHON: "3.12"
-      build_name: manywheel-py3_12-cuda13_0
+      build_name: manywheel-py3_12-cuda13_2
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
     uses: ./.github/workflows/_binary-upload.yml
 
   manywheel-py3_12-rocm7_1-build:
@@ -1728,7 +1558,7 @@ jobs:
       DOCKER_IMAGE_TAG_PREFIX: rocm7.1
       DESIRED_PYTHON: "3.12"
       runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
-      timeout-minutes: 300
+      timeout-minutes: 420
       build_name: manywheel-py3_12-rocm7_1
       build_environment: linux-binary-manywheel
     secrets:
@@ -1739,7 +1569,7 @@ jobs:
     needs:
       - manywheel-py3_12-rocm7_1-build
       - get-label-type
-    runs-on: linux.rocm.gpu.mi250.1
+    runs-on: linux.rocm.gpu.gfx942.1
     timeout-minutes: 240
     env:
       PYTORCH_ROOT: /pytorch
@@ -1757,51 +1587,47 @@ jobs:
       id-token: write
       contents: read
     steps:
-      - name: Setup ROCm
-        uses: ./.github/actions/setup-rocm
-      - uses: actions/download-artifact@v4.1.7
-        name: Download Build Artifacts
-        with:
-          name: manywheel-py3_12-rocm7_1
-          path: "${{ runner.temp }}/artifacts/"
       - name: Checkout PyTorch
         uses: actions/checkout@v4
         with:
-          ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
           submodules: recursive
-          path: pytorch
           show-progress: false
       - name: Clean PyTorch checkout
         run: |
           # Remove any artifacts from the previous checkouts
           git clean -fxd
-        working-directory: pytorch
+      - name: Setup ROCm
+        uses: ./.github/actions/setup-rocm
+      - uses: actions/download-artifact@v4.1.7
+        name: Download Build Artifacts
+        with:
+          name: manywheel-py3_12-rocm7_1
+          path: "${{ runner.temp }}/artifacts/"
       - name: ROCm set GPU_FLAG
         run: |
           echo "GPU_FLAG=--device=/dev/mem --device=/dev/kfd --device=/dev/dri --group-add video --group-add daemon" >> "${GITHUB_ENV}"
       - name: configure aws credentials
         id: aws_creds
         if: ${{ startsWith(github.event.ref, 'refs/tags/ciflow/') }}
-        uses: aws-actions/configure-aws-credentials@v4
+        uses: aws-actions/configure-aws-credentials@7474bc4690e29a8392af63c5b98e7449536d5c3a # v4
         with:
           role-to-assume: arn:aws:iam::308535385114:role/gha_workflow_s3_and_ecr_read_only
           aws-region: us-east-1
           role-duration-seconds: 18000
       - name: Calculate docker image
         id: calculate-docker-image
-        uses: pytorch/test-infra/.github/actions/calculate-docker-image@main
+        uses: pytorch/test-infra/.github/actions/calculate-docker-image@release/2.12
         with:
           docker-registry: ${{ startsWith(github.event.ref, 'refs/tags/ciflow/') && '308535385114.dkr.ecr.us-east-1.amazonaws.com' || 'docker.io' }}
           docker-image-name: manylinux2_28-builder
           custom-tag-prefix: rocm7.1
           docker-build-dir: .ci/docker
-          working-directory: pytorch
       - name: Pull Docker image
-        uses: pytorch/test-infra/.github/actions/pull-docker-image@main
+        uses: pytorch/test-infra/.github/actions/pull-docker-image@release/2.12
         with:
           docker-image: ${{ steps.calculate-docker-image.outputs.docker-image }}
       - name: Test Pytorch binary
-        uses: ./pytorch/.github/actions/test-pytorch-binary
+        uses: ./.github/actions/test-pytorch-binary
         env:
           DOCKER_IMAGE: ${{ steps.calculate-docker-image.outputs.docker-image }}
       - name: Teardown ROCm
@@ -1826,6 +1652,9 @@ jobs:
       build_name: manywheel-py3_12-rocm7_1
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
     uses: ./.github/workflows/_binary-upload.yml
 
   manywheel-py3_12-rocm7_2-build:
@@ -1844,7 +1673,7 @@ jobs:
       DOCKER_IMAGE_TAG_PREFIX: rocm7.2
       DESIRED_PYTHON: "3.12"
       runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
-      timeout-minutes: 300
+      timeout-minutes: 420
       build_name: manywheel-py3_12-rocm7_2
       build_environment: linux-binary-manywheel
     secrets:
@@ -1855,7 +1684,7 @@ jobs:
     needs:
       - manywheel-py3_12-rocm7_2-build
       - get-label-type
-    runs-on: linux.rocm.gpu.mi250.1
+    runs-on: linux.rocm.gpu.gfx942.1
     timeout-minutes: 240
     env:
       PYTORCH_ROOT: /pytorch
@@ -1873,51 +1702,47 @@ jobs:
       id-token: write
       contents: read
     steps:
-      - name: Setup ROCm
-        uses: ./.github/actions/setup-rocm
-      - uses: actions/download-artifact@v4.1.7
-        name: Download Build Artifacts
-        with:
-          name: manywheel-py3_12-rocm7_2
-          path: "${{ runner.temp }}/artifacts/"
       - name: Checkout PyTorch
         uses: actions/checkout@v4
         with:
-          ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
           submodules: recursive
-          path: pytorch
           show-progress: false
       - name: Clean PyTorch checkout
         run: |
           # Remove any artifacts from the previous checkouts
           git clean -fxd
-        working-directory: pytorch
+      - name: Setup ROCm
+        uses: ./.github/actions/setup-rocm
+      - uses: actions/download-artifact@v4.1.7
+        name: Download Build Artifacts
+        with:
+          name: manywheel-py3_12-rocm7_2
+          path: "${{ runner.temp }}/artifacts/"
       - name: ROCm set GPU_FLAG
         run: |
           echo "GPU_FLAG=--device=/dev/mem --device=/dev/kfd --device=/dev/dri --group-add video --group-add daemon" >> "${GITHUB_ENV}"
       - name: configure aws credentials
         id: aws_creds
         if: ${{ startsWith(github.event.ref, 'refs/tags/ciflow/') }}
-        uses: aws-actions/configure-aws-credentials@v4
+        uses: aws-actions/configure-aws-credentials@7474bc4690e29a8392af63c5b98e7449536d5c3a # v4
         with:
           role-to-assume: arn:aws:iam::308535385114:role/gha_workflow_s3_and_ecr_read_only
           aws-region: us-east-1
           role-duration-seconds: 18000
       - name: Calculate docker image
         id: calculate-docker-image
-        uses: pytorch/test-infra/.github/actions/calculate-docker-image@main
+        uses: pytorch/test-infra/.github/actions/calculate-docker-image@release/2.12
         with:
           docker-registry: ${{ startsWith(github.event.ref, 'refs/tags/ciflow/') && '308535385114.dkr.ecr.us-east-1.amazonaws.com' || 'docker.io' }}
           docker-image-name: manylinux2_28-builder
           custom-tag-prefix: rocm7.2
           docker-build-dir: .ci/docker
-          working-directory: pytorch
       - name: Pull Docker image
-        uses: pytorch/test-infra/.github/actions/pull-docker-image@main
+        uses: pytorch/test-infra/.github/actions/pull-docker-image@release/2.12
         with:
           docker-image: ${{ steps.calculate-docker-image.outputs.docker-image }}
       - name: Test Pytorch binary
-        uses: ./pytorch/.github/actions/test-pytorch-binary
+        uses: ./.github/actions/test-pytorch-binary
         env:
           DOCKER_IMAGE: ${{ steps.calculate-docker-image.outputs.docker-image }}
       - name: Teardown ROCm
@@ -1942,6 +1767,9 @@ jobs:
       build_name: manywheel-py3_12-rocm7_2
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
     uses: ./.github/workflows/_binary-upload.yml
 
   manywheel-py3_12-xpu-build:
@@ -1987,42 +1815,38 @@ jobs:
       id-token: write
       contents: read
     steps:
-      - name: Setup XPU
-        uses: pytorch/pytorch/.github/actions/setup-xpu@main
-      - name: Login to ECR
-        uses: pytorch/pytorch/.github/actions/ecr-login@main
-      - uses: actions/download-artifact@v4.1.7
-        name: Download Build Artifacts
-        with:
-          name: manywheel-py3_12-xpu
-          path: "${{ runner.temp }}/artifacts/"
       - name: Checkout PyTorch
         uses: actions/checkout@v4
         with:
-          ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
           submodules: recursive
-          path: pytorch
           show-progress: false
       - name: Clean PyTorch checkout
         run: |
           # Remove any artifacts from the previous checkouts
           git clean -fxd
-        working-directory: pytorch
+      - name: Setup XPU
+        uses: pytorch/pytorch/.github/actions/setup-xpu@release/2.12
+      - name: Login to ECR
+        uses: pytorch/pytorch/.github/actions/ecr-login@release/2.12
+      - uses: actions/download-artifact@v4.1.7
+        name: Download Build Artifacts
+        with:
+          name: manywheel-py3_12-xpu
+          path: "${{ runner.temp }}/artifacts/"
       - name: Calculate docker image
         id: calculate-docker-image
-        uses: pytorch/test-infra/.github/actions/calculate-docker-image@main
+        uses: pytorch/test-infra/.github/actions/calculate-docker-image@release/2.12
         with:
           docker-registry: ${{ startsWith(github.event.ref, 'refs/tags/ciflow/') && '308535385114.dkr.ecr.us-east-1.amazonaws.com' || 'docker.io' }}
           docker-image-name: manylinux2_28-builder
           custom-tag-prefix: xpu
           docker-build-dir: .ci/docker
-          working-directory: pytorch
       - name: Pull Docker image
-        uses: pytorch/test-infra/.github/actions/pull-docker-image@main
+        uses: pytorch/test-infra/.github/actions/pull-docker-image@release/2.12
         with:
           docker-image: ${{ steps.calculate-docker-image.outputs.docker-image }}
       - name: Test Pytorch binary
-        uses: ./pytorch/.github/actions/test-pytorch-binary
+        uses: ./.github/actions/test-pytorch-binary
         env:
           DOCKER_IMAGE: ${{ steps.calculate-docker-image.outputs.docker-image }}
       - name: Teardown XPU
@@ -2046,6 +1870,9 @@ jobs:
       build_name: manywheel-py3_12-xpu
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
     uses: ./.github/workflows/_binary-upload.yml
 
   manywheel-py3_13-cpu-build:
@@ -2109,6 +1936,9 @@ jobs:
       build_name: manywheel-py3_13-cpu
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
     uses: ./.github/workflows/_binary-upload.yml
 
   manywheel-py3_13-cuda12_6-build:
@@ -2129,7 +1959,7 @@ jobs:
       runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
       build_name: manywheel-py3_13-cuda12_6
       build_environment: linux-binary-manywheel
-      PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cublas,cufile,nvjitlink,nvtx]==12.6.3; platform_system == 'Linux' | cuda-bindings==12.9.4; platform_system == 'Linux' | nvidia-cudnn-cu12==9.10.2.21; platform_system == 'Linux' | nvidia-cusparselt-cu12==0.7.1; platform_system == 'Linux' | nvidia-nccl-cu12==2.28.9; platform_system == 'Linux' | nvidia-nvshmem-cu12==3.4.5; platform_system == 'Linux'
+      PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cublas,cufile,nvjitlink,nvtx]==12.6.3; platform_system == 'Linux' | cuda-bindings>=12.9.4,<13; platform_system == 'Linux' | nvidia-cudnn-cu12==9.10.2.21; platform_system == 'Linux' | nvidia-cusparselt-cu12==0.7.1; platform_system == 'Linux' | nvidia-nccl-cu12==2.29.3; platform_system == 'Linux' | nvidia-nvshmem-cu12==3.4.5; platform_system == 'Linux'
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
 
@@ -2153,7 +1983,7 @@ jobs:
       build_name: manywheel-py3_13-cuda12_6
       build_environment: linux-binary-manywheel
       runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
-      runs_on: linux.4xlarge.nvidia.gpu  # 12.6 build can use maxwell (sm_50) runner
+      runs_on: linux.g4dn.4xlarge.nvidia.gpu
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
   manywheel-py3_13-cuda12_6-upload:  # Uploading
@@ -2176,76 +2006,12 @@ jobs:
       build_name: manywheel-py3_13-cuda12_6
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
     uses: ./.github/workflows/_binary-upload.yml
 
-  manywheel-py3_13-cuda12_8-build:
-    if: ${{ github.repository_owner == 'pytorch' }}
-    uses: ./.github/workflows/_binary-build-linux.yml
-    needs: get-label-type
-    with:
-      PYTORCH_ROOT: /pytorch
-      PACKAGE_TYPE: manywheel
-      # TODO: This is a legacy variable that we eventually want to get rid of in
-      #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu128
-      GPU_ARCH_VERSION: "12.8"
-      GPU_ARCH_TYPE: cuda
-      DOCKER_IMAGE: manylinux2_28-builder
-      DOCKER_IMAGE_TAG_PREFIX: cuda12.8
-      DESIRED_PYTHON: "3.13"
-      runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
-      build_name: manywheel-py3_13-cuda12_8
-      build_environment: linux-binary-manywheel
-      PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cublas,cufile,nvjitlink,nvtx]==12.8.1; platform_system == 'Linux' | cuda-bindings==12.9.4; platform_system == 'Linux' | nvidia-cudnn-cu12==9.17.1.4; platform_system == 'Linux' | nvidia-cusparselt-cu12==0.7.1; platform_system == 'Linux' | nvidia-nccl-cu12==2.28.9; platform_system == 'Linux' | nvidia-nvshmem-cu12==3.4.5; platform_system == 'Linux'
-    secrets:
-      github-token: ${{ secrets.GITHUB_TOKEN }}
-
-  manywheel-py3_13-cuda12_8-test:  # Testing
-    if: ${{ github.repository_owner == 'pytorch' }}
-    needs:
-      - manywheel-py3_13-cuda12_8-build
-      - get-label-type
-    uses: ./.github/workflows/_binary-test-linux.yml
-    with:
-      PYTORCH_ROOT: /pytorch
-      PACKAGE_TYPE: manywheel
-      # TODO: This is a legacy variable that we eventually want to get rid of in
-      #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu128
-      GPU_ARCH_VERSION: "12.8"
-      GPU_ARCH_TYPE: cuda
-      DOCKER_IMAGE: manylinux2_28-builder
-      DOCKER_IMAGE_TAG_PREFIX: cuda12.8
-      DESIRED_PYTHON: "3.13"
-      build_name: manywheel-py3_13-cuda12_8
-      build_environment: linux-binary-manywheel
-      runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
-      runs_on: linux.g4dn.4xlarge.nvidia.gpu # 12.8+ builds need sm_70+ runner
-    secrets:
-      github-token: ${{ secrets.GITHUB_TOKEN }}
-  manywheel-py3_13-cuda12_8-upload:  # Uploading
-    if: ${{ github.repository_owner == 'pytorch' }}
-    permissions:
-      id-token: write
-      contents: read
-    needs: manywheel-py3_13-cuda12_8-test
-    with:
-      PYTORCH_ROOT: /pytorch
-      PACKAGE_TYPE: manywheel
-      # TODO: This is a legacy variable that we eventually want to get rid of in
-      #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu128
-      GPU_ARCH_VERSION: "12.8"
-      GPU_ARCH_TYPE: cuda
-      DOCKER_IMAGE: manylinux2_28-builder
-      DOCKER_IMAGE_TAG_PREFIX: cuda12.8
-      DESIRED_PYTHON: "3.13"
-      build_name: manywheel-py3_13-cuda12_8
-    secrets:
-      github-token: ${{ secrets.GITHUB_TOKEN }}
-    uses: ./.github/workflows/_binary-upload.yml
-
-  manywheel-py3_13-cuda12_9-build:
+  manywheel-py3_13-cuda13_0-build:
     if: ${{ github.repository_owner == 'pytorch' }}
     uses: ./.github/workflows/_binary-build-linux.yml
     needs: get-label-type
@@ -2254,23 +2020,23 @@ jobs:
       PACKAGE_TYPE: manywheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu129
-      GPU_ARCH_VERSION: "12.9"
+      DESIRED_CUDA: cu130
+      GPU_ARCH_VERSION: "13.0"
       GPU_ARCH_TYPE: cuda
       DOCKER_IMAGE: manylinux2_28-builder
-      DOCKER_IMAGE_TAG_PREFIX: cuda12.9
+      DOCKER_IMAGE_TAG_PREFIX: cuda13.0
       DESIRED_PYTHON: "3.13"
       runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
-      build_name: manywheel-py3_13-cuda12_9
+      build_name: manywheel-py3_13-cuda13_0
       build_environment: linux-binary-manywheel
-      PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cublas,cufile,nvjitlink,nvtx]==12.9.1; platform_system == 'Linux' | cuda-bindings==12.9.4; platform_system == 'Linux' | nvidia-cudnn-cu12==9.17.1.4; platform_system == 'Linux' | nvidia-cusparselt-cu12==0.7.1; platform_system == 'Linux' | nvidia-nccl-cu12==2.28.9; platform_system == 'Linux' | nvidia-nvshmem-cu12==3.4.5; platform_system == 'Linux'
+      PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cufile,nvjitlink,nvtx]==13.0.2; platform_system == 'Linux' | nvidia-cublas>=13.1.0.3,<=13.1.1.3; platform_system == 'Linux' | cuda-bindings>=13.0.3,<14; platform_system == 'Linux' | nvidia-cudnn-cu13==9.20.0.48; platform_system == 'Linux' | nvidia-cusparselt-cu13==0.8.1; platform_system == 'Linux' | nvidia-nccl-cu13==2.29.7; platform_system == 'Linux' | nvidia-nvshmem-cu13==3.4.5; platform_system == 'Linux'
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
 
-  manywheel-py3_13-cuda12_9-test:  # Testing
+  manywheel-py3_13-cuda13_0-test:  # Testing
     if: ${{ github.repository_owner == 'pytorch' }}
     needs:
-      - manywheel-py3_13-cuda12_9-build
+      - manywheel-py3_13-cuda13_0-build
       - get-label-type
     uses: ./.github/workflows/_binary-test-linux.yml
     with:
@@ -2278,41 +2044,44 @@ jobs:
       PACKAGE_TYPE: manywheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu129
-      GPU_ARCH_VERSION: "12.9"
+      DESIRED_CUDA: cu130
+      GPU_ARCH_VERSION: "13.0"
       GPU_ARCH_TYPE: cuda
       DOCKER_IMAGE: manylinux2_28-builder
-      DOCKER_IMAGE_TAG_PREFIX: cuda12.9
+      DOCKER_IMAGE_TAG_PREFIX: cuda13.0
       DESIRED_PYTHON: "3.13"
-      build_name: manywheel-py3_13-cuda12_9
+      build_name: manywheel-py3_13-cuda13_0
       build_environment: linux-binary-manywheel
       runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
-      runs_on: linux.g4dn.4xlarge.nvidia.gpu # 12.8+ builds need sm_70+ runner
+      runs_on: linux.g4dn.4xlarge.nvidia.gpu
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
-  manywheel-py3_13-cuda12_9-upload:  # Uploading
+  manywheel-py3_13-cuda13_0-upload:  # Uploading
     if: ${{ github.repository_owner == 'pytorch' }}
     permissions:
       id-token: write
       contents: read
-    needs: manywheel-py3_13-cuda12_9-test
+    needs: manywheel-py3_13-cuda13_0-test
     with:
       PYTORCH_ROOT: /pytorch
       PACKAGE_TYPE: manywheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu129
-      GPU_ARCH_VERSION: "12.9"
+      DESIRED_CUDA: cu130
+      GPU_ARCH_VERSION: "13.0"
       GPU_ARCH_TYPE: cuda
       DOCKER_IMAGE: manylinux2_28-builder
-      DOCKER_IMAGE_TAG_PREFIX: cuda12.9
+      DOCKER_IMAGE_TAG_PREFIX: cuda13.0
       DESIRED_PYTHON: "3.13"
-      build_name: manywheel-py3_13-cuda12_9
+      build_name: manywheel-py3_13-cuda13_0
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
     uses: ./.github/workflows/_binary-upload.yml
 
-  manywheel-py3_13-cuda13_0-build:
+  manywheel-py3_13-cuda13_2-build:
     if: ${{ github.repository_owner == 'pytorch' }}
     uses: ./.github/workflows/_binary-build-linux.yml
     needs: get-label-type
@@ -2321,23 +2090,23 @@ jobs:
       PACKAGE_TYPE: manywheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu130
-      GPU_ARCH_VERSION: "13.0"
+      DESIRED_CUDA: cu132
+      GPU_ARCH_VERSION: "13.2"
       GPU_ARCH_TYPE: cuda
       DOCKER_IMAGE: manylinux2_28-builder
-      DOCKER_IMAGE_TAG_PREFIX: cuda13.0
+      DOCKER_IMAGE_TAG_PREFIX: cuda13.2
       DESIRED_PYTHON: "3.13"
       runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
-      build_name: manywheel-py3_13-cuda13_0
+      build_name: manywheel-py3_13-cuda13_2
       build_environment: linux-binary-manywheel
-      PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cublas,cufile,nvjitlink,nvtx]==13.0.2; platform_system == 'Linux' | cuda-bindings==13.0.3; platform_system == 'Linux' | nvidia-cudnn-cu13==9.17.1.4; platform_system == 'Linux' | nvidia-cusparselt-cu13==0.8.0; platform_system == 'Linux' | nvidia-nccl-cu13==2.28.9; platform_system == 'Linux' | nvidia-nvshmem-cu13==3.4.5; platform_system == 'Linux'
+      PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cublas,cufile,nvjitlink,nvtx]==13.2.1; platform_system == 'Linux' | cuda-bindings>=13.0.3,<14; platform_system == 'Linux' | nvidia-cudnn-cu13==9.20.0.48; platform_system == 'Linux' | nvidia-cusparselt-cu13==0.8.1; platform_system == 'Linux' | nvidia-nccl-cu13==2.29.7; platform_system == 'Linux' | nvidia-nvshmem-cu13==3.4.5; platform_system == 'Linux'
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
 
-  manywheel-py3_13-cuda13_0-test:  # Testing
+  manywheel-py3_13-cuda13_2-test:  # Testing
     if: ${{ github.repository_owner == 'pytorch' }}
     needs:
-      - manywheel-py3_13-cuda13_0-build
+      - manywheel-py3_13-cuda13_2-build
       - get-label-type
     uses: ./.github/workflows/_binary-test-linux.yml
     with:
@@ -2345,38 +2114,41 @@ jobs:
       PACKAGE_TYPE: manywheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu130
-      GPU_ARCH_VERSION: "13.0"
+      DESIRED_CUDA: cu132
+      GPU_ARCH_VERSION: "13.2"
       GPU_ARCH_TYPE: cuda
       DOCKER_IMAGE: manylinux2_28-builder
-      DOCKER_IMAGE_TAG_PREFIX: cuda13.0
+      DOCKER_IMAGE_TAG_PREFIX: cuda13.2
       DESIRED_PYTHON: "3.13"
-      build_name: manywheel-py3_13-cuda13_0
+      build_name: manywheel-py3_13-cuda13_2
       build_environment: linux-binary-manywheel
       runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
-      runs_on: linux.g4dn.4xlarge.nvidia.gpu # 12.8+ builds need sm_70+ runner
+      runs_on: linux.g4dn.4xlarge.nvidia.gpu
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
-  manywheel-py3_13-cuda13_0-upload:  # Uploading
+  manywheel-py3_13-cuda13_2-upload:  # Uploading
     if: ${{ github.repository_owner == 'pytorch' }}
     permissions:
       id-token: write
       contents: read
-    needs: manywheel-py3_13-cuda13_0-test
+    needs: manywheel-py3_13-cuda13_2-test
     with:
       PYTORCH_ROOT: /pytorch
       PACKAGE_TYPE: manywheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu130
-      GPU_ARCH_VERSION: "13.0"
+      DESIRED_CUDA: cu132
+      GPU_ARCH_VERSION: "13.2"
       GPU_ARCH_TYPE: cuda
       DOCKER_IMAGE: manylinux2_28-builder
-      DOCKER_IMAGE_TAG_PREFIX: cuda13.0
+      DOCKER_IMAGE_TAG_PREFIX: cuda13.2
       DESIRED_PYTHON: "3.13"
-      build_name: manywheel-py3_13-cuda13_0
+      build_name: manywheel-py3_13-cuda13_2
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
     uses: ./.github/workflows/_binary-upload.yml
 
   manywheel-py3_13-rocm7_1-build:
@@ -2395,7 +2167,7 @@ jobs:
       DOCKER_IMAGE_TAG_PREFIX: rocm7.1
       DESIRED_PYTHON: "3.13"
       runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
-      timeout-minutes: 300
+      timeout-minutes: 420
       build_name: manywheel-py3_13-rocm7_1
       build_environment: linux-binary-manywheel
     secrets:
@@ -2406,7 +2178,7 @@ jobs:
     needs:
       - manywheel-py3_13-rocm7_1-build
       - get-label-type
-    runs-on: linux.rocm.gpu.mi250.1
+    runs-on: linux.rocm.gpu.gfx942.1
     timeout-minutes: 240
     env:
       PYTORCH_ROOT: /pytorch
@@ -2424,51 +2196,47 @@ jobs:
       id-token: write
       contents: read
     steps:
-      - name: Setup ROCm
-        uses: ./.github/actions/setup-rocm
-      - uses: actions/download-artifact@v4.1.7
-        name: Download Build Artifacts
-        with:
-          name: manywheel-py3_13-rocm7_1
-          path: "${{ runner.temp }}/artifacts/"
       - name: Checkout PyTorch
         uses: actions/checkout@v4
         with:
-          ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
           submodules: recursive
-          path: pytorch
           show-progress: false
       - name: Clean PyTorch checkout
         run: |
           # Remove any artifacts from the previous checkouts
           git clean -fxd
-        working-directory: pytorch
+      - name: Setup ROCm
+        uses: ./.github/actions/setup-rocm
+      - uses: actions/download-artifact@v4.1.7
+        name: Download Build Artifacts
+        with:
+          name: manywheel-py3_13-rocm7_1
+          path: "${{ runner.temp }}/artifacts/"
       - name: ROCm set GPU_FLAG
         run: |
           echo "GPU_FLAG=--device=/dev/mem --device=/dev/kfd --device=/dev/dri --group-add video --group-add daemon" >> "${GITHUB_ENV}"
       - name: configure aws credentials
         id: aws_creds
         if: ${{ startsWith(github.event.ref, 'refs/tags/ciflow/') }}
-        uses: aws-actions/configure-aws-credentials@v4
+        uses: aws-actions/configure-aws-credentials@7474bc4690e29a8392af63c5b98e7449536d5c3a # v4
         with:
           role-to-assume: arn:aws:iam::308535385114:role/gha_workflow_s3_and_ecr_read_only
           aws-region: us-east-1
           role-duration-seconds: 18000
       - name: Calculate docker image
         id: calculate-docker-image
-        uses: pytorch/test-infra/.github/actions/calculate-docker-image@main
+        uses: pytorch/test-infra/.github/actions/calculate-docker-image@release/2.12
         with:
           docker-registry: ${{ startsWith(github.event.ref, 'refs/tags/ciflow/') && '308535385114.dkr.ecr.us-east-1.amazonaws.com' || 'docker.io' }}
           docker-image-name: manylinux2_28-builder
           custom-tag-prefix: rocm7.1
           docker-build-dir: .ci/docker
-          working-directory: pytorch
       - name: Pull Docker image
-        uses: pytorch/test-infra/.github/actions/pull-docker-image@main
+        uses: pytorch/test-infra/.github/actions/pull-docker-image@release/2.12
         with:
           docker-image: ${{ steps.calculate-docker-image.outputs.docker-image }}
       - name: Test Pytorch binary
-        uses: ./pytorch/.github/actions/test-pytorch-binary
+        uses: ./.github/actions/test-pytorch-binary
         env:
           DOCKER_IMAGE: ${{ steps.calculate-docker-image.outputs.docker-image }}
       - name: Teardown ROCm
@@ -2493,6 +2261,9 @@ jobs:
       build_name: manywheel-py3_13-rocm7_1
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
     uses: ./.github/workflows/_binary-upload.yml
 
   manywheel-py3_13-rocm7_2-build:
@@ -2511,7 +2282,7 @@ jobs:
       DOCKER_IMAGE_TAG_PREFIX: rocm7.2
       DESIRED_PYTHON: "3.13"
       runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
-      timeout-minutes: 300
+      timeout-minutes: 420
       build_name: manywheel-py3_13-rocm7_2
       build_environment: linux-binary-manywheel
     secrets:
@@ -2522,7 +2293,7 @@ jobs:
     needs:
       - manywheel-py3_13-rocm7_2-build
       - get-label-type
-    runs-on: linux.rocm.gpu.mi250.1
+    runs-on: linux.rocm.gpu.gfx942.1
     timeout-minutes: 240
     env:
       PYTORCH_ROOT: /pytorch
@@ -2540,51 +2311,47 @@ jobs:
       id-token: write
       contents: read
     steps:
-      - name: Setup ROCm
-        uses: ./.github/actions/setup-rocm
-      - uses: actions/download-artifact@v4.1.7
-        name: Download Build Artifacts
-        with:
-          name: manywheel-py3_13-rocm7_2
-          path: "${{ runner.temp }}/artifacts/"
       - name: Checkout PyTorch
         uses: actions/checkout@v4
         with:
-          ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
           submodules: recursive
-          path: pytorch
           show-progress: false
       - name: Clean PyTorch checkout
         run: |
           # Remove any artifacts from the previous checkouts
           git clean -fxd
-        working-directory: pytorch
+      - name: Setup ROCm
+        uses: ./.github/actions/setup-rocm
+      - uses: actions/download-artifact@v4.1.7
+        name: Download Build Artifacts
+        with:
+          name: manywheel-py3_13-rocm7_2
+          path: "${{ runner.temp }}/artifacts/"
       - name: ROCm set GPU_FLAG
         run: |
           echo "GPU_FLAG=--device=/dev/mem --device=/dev/kfd --device=/dev/dri --group-add video --group-add daemon" >> "${GITHUB_ENV}"
       - name: configure aws credentials
         id: aws_creds
         if: ${{ startsWith(github.event.ref, 'refs/tags/ciflow/') }}
-        uses: aws-actions/configure-aws-credentials@v4
+        uses: aws-actions/configure-aws-credentials@7474bc4690e29a8392af63c5b98e7449536d5c3a # v4
         with:
           role-to-assume: arn:aws:iam::308535385114:role/gha_workflow_s3_and_ecr_read_only
           aws-region: us-east-1
           role-duration-seconds: 18000
       - name: Calculate docker image
         id: calculate-docker-image
-        uses: pytorch/test-infra/.github/actions/calculate-docker-image@main
+        uses: pytorch/test-infra/.github/actions/calculate-docker-image@release/2.12
         with:
           docker-registry: ${{ startsWith(github.event.ref, 'refs/tags/ciflow/') && '308535385114.dkr.ecr.us-east-1.amazonaws.com' || 'docker.io' }}
           docker-image-name: manylinux2_28-builder
           custom-tag-prefix: rocm7.2
           docker-build-dir: .ci/docker
-          working-directory: pytorch
       - name: Pull Docker image
-        uses: pytorch/test-infra/.github/actions/pull-docker-image@main
+        uses: pytorch/test-infra/.github/actions/pull-docker-image@release/2.12
         with:
           docker-image: ${{ steps.calculate-docker-image.outputs.docker-image }}
       - name: Test Pytorch binary
-        uses: ./pytorch/.github/actions/test-pytorch-binary
+        uses: ./.github/actions/test-pytorch-binary
         env:
           DOCKER_IMAGE: ${{ steps.calculate-docker-image.outputs.docker-image }}
       - name: Teardown ROCm
@@ -2609,6 +2376,9 @@ jobs:
       build_name: manywheel-py3_13-rocm7_2
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
     uses: ./.github/workflows/_binary-upload.yml
 
   manywheel-py3_13-xpu-build:
@@ -2654,42 +2424,38 @@ jobs:
       id-token: write
       contents: read
     steps:
-      - name: Setup XPU
-        uses: pytorch/pytorch/.github/actions/setup-xpu@main
-      - name: Login to ECR
-        uses: pytorch/pytorch/.github/actions/ecr-login@main
-      - uses: actions/download-artifact@v4.1.7
-        name: Download Build Artifacts
-        with:
-          name: manywheel-py3_13-xpu
-          path: "${{ runner.temp }}/artifacts/"
       - name: Checkout PyTorch
         uses: actions/checkout@v4
         with:
-          ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
           submodules: recursive
-          path: pytorch
           show-progress: false
       - name: Clean PyTorch checkout
         run: |
           # Remove any artifacts from the previous checkouts
           git clean -fxd
-        working-directory: pytorch
+      - name: Setup XPU
+        uses: pytorch/pytorch/.github/actions/setup-xpu@release/2.12
+      - name: Login to ECR
+        uses: pytorch/pytorch/.github/actions/ecr-login@release/2.12
+      - uses: actions/download-artifact@v4.1.7
+        name: Download Build Artifacts
+        with:
+          name: manywheel-py3_13-xpu
+          path: "${{ runner.temp }}/artifacts/"
       - name: Calculate docker image
         id: calculate-docker-image
-        uses: pytorch/test-infra/.github/actions/calculate-docker-image@main
+        uses: pytorch/test-infra/.github/actions/calculate-docker-image@release/2.12
         with:
           docker-registry: ${{ startsWith(github.event.ref, 'refs/tags/ciflow/') && '308535385114.dkr.ecr.us-east-1.amazonaws.com' || 'docker.io' }}
           docker-image-name: manylinux2_28-builder
           custom-tag-prefix: xpu
           docker-build-dir: .ci/docker
-          working-directory: pytorch
       - name: Pull Docker image
-        uses: pytorch/test-infra/.github/actions/pull-docker-image@main
+        uses: pytorch/test-infra/.github/actions/pull-docker-image@release/2.12
         with:
           docker-image: ${{ steps.calculate-docker-image.outputs.docker-image }}
       - name: Test Pytorch binary
-        uses: ./pytorch/.github/actions/test-pytorch-binary
+        uses: ./.github/actions/test-pytorch-binary
         env:
           DOCKER_IMAGE: ${{ steps.calculate-docker-image.outputs.docker-image }}
       - name: Teardown XPU
@@ -2713,6 +2479,9 @@ jobs:
       build_name: manywheel-py3_13-xpu
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
     uses: ./.github/workflows/_binary-upload.yml
 
   manywheel-py3_13t-cpu-build:
@@ -2776,6 +2545,9 @@ jobs:
       build_name: manywheel-py3_13t-cpu
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
     uses: ./.github/workflows/_binary-upload.yml
 
   manywheel-py3_13t-cuda12_6-build:
@@ -2796,7 +2568,7 @@ jobs:
       runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
       build_name: manywheel-py3_13t-cuda12_6
       build_environment: linux-binary-manywheel
-      PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cublas,cufile,nvjitlink,nvtx]==12.6.3; platform_system == 'Linux' | cuda-bindings==12.9.4; platform_system == 'Linux' | nvidia-cudnn-cu12==9.10.2.21; platform_system == 'Linux' | nvidia-cusparselt-cu12==0.7.1; platform_system == 'Linux' | nvidia-nccl-cu12==2.28.9; platform_system == 'Linux' | nvidia-nvshmem-cu12==3.4.5; platform_system == 'Linux'
+      PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cublas,cufile,nvjitlink,nvtx]==12.6.3; platform_system == 'Linux' | cuda-bindings>=12.9.4,<13; platform_system == 'Linux' | nvidia-cudnn-cu12==9.10.2.21; platform_system == 'Linux' | nvidia-cusparselt-cu12==0.7.1; platform_system == 'Linux' | nvidia-nccl-cu12==2.29.3; platform_system == 'Linux' | nvidia-nvshmem-cu12==3.4.5; platform_system == 'Linux'
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
 
@@ -2820,7 +2592,7 @@ jobs:
       build_name: manywheel-py3_13t-cuda12_6
       build_environment: linux-binary-manywheel
       runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
-      runs_on: linux.4xlarge.nvidia.gpu  # 12.6 build can use maxwell (sm_50) runner
+      runs_on: linux.g4dn.4xlarge.nvidia.gpu
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
   manywheel-py3_13t-cuda12_6-upload:  # Uploading
@@ -2843,9 +2615,12 @@ jobs:
       build_name: manywheel-py3_13t-cuda12_6
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
     uses: ./.github/workflows/_binary-upload.yml
 
-  manywheel-py3_13t-cuda12_8-build:
+  manywheel-py3_13t-cuda13_0-build:
     if: ${{ github.repository_owner == 'pytorch' }}
     uses: ./.github/workflows/_binary-build-linux.yml
     needs: get-label-type
@@ -2854,90 +2629,23 @@ jobs:
       PACKAGE_TYPE: manywheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu128
-      GPU_ARCH_VERSION: "12.8"
+      DESIRED_CUDA: cu130
+      GPU_ARCH_VERSION: "13.0"
       GPU_ARCH_TYPE: cuda
       DOCKER_IMAGE: manylinux2_28-builder
-      DOCKER_IMAGE_TAG_PREFIX: cuda12.8
+      DOCKER_IMAGE_TAG_PREFIX: cuda13.0
       DESIRED_PYTHON: "3.13t"
       runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
-      build_name: manywheel-py3_13t-cuda12_8
-      build_environment: linux-binary-manywheel
-      PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cublas,cufile,nvjitlink,nvtx]==12.8.1; platform_system == 'Linux' | cuda-bindings==12.9.4; platform_system == 'Linux' | nvidia-cudnn-cu12==9.17.1.4; platform_system == 'Linux' | nvidia-cusparselt-cu12==0.7.1; platform_system == 'Linux' | nvidia-nccl-cu12==2.28.9; platform_system == 'Linux' | nvidia-nvshmem-cu12==3.4.5; platform_system == 'Linux'
-    secrets:
-      github-token: ${{ secrets.GITHUB_TOKEN }}
-
-  manywheel-py3_13t-cuda12_8-test:  # Testing
-    if: ${{ github.repository_owner == 'pytorch' }}
-    needs:
-      - manywheel-py3_13t-cuda12_8-build
-      - get-label-type
-    uses: ./.github/workflows/_binary-test-linux.yml
-    with:
-      PYTORCH_ROOT: /pytorch
-      PACKAGE_TYPE: manywheel
-      # TODO: This is a legacy variable that we eventually want to get rid of in
-      #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu128
-      GPU_ARCH_VERSION: "12.8"
-      GPU_ARCH_TYPE: cuda
-      DOCKER_IMAGE: manylinux2_28-builder
-      DOCKER_IMAGE_TAG_PREFIX: cuda12.8
-      DESIRED_PYTHON: "3.13t"
-      build_name: manywheel-py3_13t-cuda12_8
-      build_environment: linux-binary-manywheel
-      runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
-      runs_on: linux.g4dn.4xlarge.nvidia.gpu # 12.8+ builds need sm_70+ runner
-    secrets:
-      github-token: ${{ secrets.GITHUB_TOKEN }}
-  manywheel-py3_13t-cuda12_8-upload:  # Uploading
-    if: ${{ github.repository_owner == 'pytorch' }}
-    permissions:
-      id-token: write
-      contents: read
-    needs: manywheel-py3_13t-cuda12_8-test
-    with:
-      PYTORCH_ROOT: /pytorch
-      PACKAGE_TYPE: manywheel
-      # TODO: This is a legacy variable that we eventually want to get rid of in
-      #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu128
-      GPU_ARCH_VERSION: "12.8"
-      GPU_ARCH_TYPE: cuda
-      DOCKER_IMAGE: manylinux2_28-builder
-      DOCKER_IMAGE_TAG_PREFIX: cuda12.8
-      DESIRED_PYTHON: "3.13t"
-      build_name: manywheel-py3_13t-cuda12_8
-    secrets:
-      github-token: ${{ secrets.GITHUB_TOKEN }}
-    uses: ./.github/workflows/_binary-upload.yml
-
-  manywheel-py3_13t-cuda12_9-build:
-    if: ${{ github.repository_owner == 'pytorch' }}
-    uses: ./.github/workflows/_binary-build-linux.yml
-    needs: get-label-type
-    with:
-      PYTORCH_ROOT: /pytorch
-      PACKAGE_TYPE: manywheel
-      # TODO: This is a legacy variable that we eventually want to get rid of in
-      #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu129
-      GPU_ARCH_VERSION: "12.9"
-      GPU_ARCH_TYPE: cuda
-      DOCKER_IMAGE: manylinux2_28-builder
-      DOCKER_IMAGE_TAG_PREFIX: cuda12.9
-      DESIRED_PYTHON: "3.13t"
-      runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
-      build_name: manywheel-py3_13t-cuda12_9
+      build_name: manywheel-py3_13t-cuda13_0
       build_environment: linux-binary-manywheel
-      PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cublas,cufile,nvjitlink,nvtx]==12.9.1; platform_system == 'Linux' | cuda-bindings==12.9.4; platform_system == 'Linux' | nvidia-cudnn-cu12==9.17.1.4; platform_system == 'Linux' | nvidia-cusparselt-cu12==0.7.1; platform_system == 'Linux' | nvidia-nccl-cu12==2.28.9; platform_system == 'Linux' | nvidia-nvshmem-cu12==3.4.5; platform_system == 'Linux'
+      PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cufile,nvjitlink,nvtx]==13.0.2; platform_system == 'Linux' | nvidia-cublas>=13.1.0.3,<=13.1.1.3; platform_system == 'Linux' | cuda-bindings>=13.0.3,<14; platform_system == 'Linux' | nvidia-cudnn-cu13==9.20.0.48; platform_system == 'Linux' | nvidia-cusparselt-cu13==0.8.1; platform_system == 'Linux' | nvidia-nccl-cu13==2.29.7; platform_system == 'Linux' | nvidia-nvshmem-cu13==3.4.5; platform_system == 'Linux'
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
 
-  manywheel-py3_13t-cuda12_9-test:  # Testing
+  manywheel-py3_13t-cuda13_0-test:  # Testing
     if: ${{ github.repository_owner == 'pytorch' }}
     needs:
-      - manywheel-py3_13t-cuda12_9-build
+      - manywheel-py3_13t-cuda13_0-build
       - get-label-type
     uses: ./.github/workflows/_binary-test-linux.yml
     with:
@@ -2945,41 +2653,44 @@ jobs:
       PACKAGE_TYPE: manywheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu129
-      GPU_ARCH_VERSION: "12.9"
+      DESIRED_CUDA: cu130
+      GPU_ARCH_VERSION: "13.0"
       GPU_ARCH_TYPE: cuda
       DOCKER_IMAGE: manylinux2_28-builder
-      DOCKER_IMAGE_TAG_PREFIX: cuda12.9
+      DOCKER_IMAGE_TAG_PREFIX: cuda13.0
       DESIRED_PYTHON: "3.13t"
-      build_name: manywheel-py3_13t-cuda12_9
+      build_name: manywheel-py3_13t-cuda13_0
       build_environment: linux-binary-manywheel
       runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
-      runs_on: linux.g4dn.4xlarge.nvidia.gpu # 12.8+ builds need sm_70+ runner
+      runs_on: linux.g4dn.4xlarge.nvidia.gpu
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
-  manywheel-py3_13t-cuda12_9-upload:  # Uploading
+  manywheel-py3_13t-cuda13_0-upload:  # Uploading
     if: ${{ github.repository_owner == 'pytorch' }}
     permissions:
       id-token: write
       contents: read
-    needs: manywheel-py3_13t-cuda12_9-test
+    needs: manywheel-py3_13t-cuda13_0-test
     with:
       PYTORCH_ROOT: /pytorch
       PACKAGE_TYPE: manywheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu129
-      GPU_ARCH_VERSION: "12.9"
+      DESIRED_CUDA: cu130
+      GPU_ARCH_VERSION: "13.0"
       GPU_ARCH_TYPE: cuda
       DOCKER_IMAGE: manylinux2_28-builder
-      DOCKER_IMAGE_TAG_PREFIX: cuda12.9
+      DOCKER_IMAGE_TAG_PREFIX: cuda13.0
       DESIRED_PYTHON: "3.13t"
-      build_name: manywheel-py3_13t-cuda12_9
+      build_name: manywheel-py3_13t-cuda13_0
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
     uses: ./.github/workflows/_binary-upload.yml
 
-  manywheel-py3_13t-cuda13_0-build:
+  manywheel-py3_13t-cuda13_2-build:
     if: ${{ github.repository_owner == 'pytorch' }}
     uses: ./.github/workflows/_binary-build-linux.yml
     needs: get-label-type
@@ -2988,23 +2699,23 @@ jobs:
       PACKAGE_TYPE: manywheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu130
-      GPU_ARCH_VERSION: "13.0"
+      DESIRED_CUDA: cu132
+      GPU_ARCH_VERSION: "13.2"
       GPU_ARCH_TYPE: cuda
       DOCKER_IMAGE: manylinux2_28-builder
-      DOCKER_IMAGE_TAG_PREFIX: cuda13.0
+      DOCKER_IMAGE_TAG_PREFIX: cuda13.2
       DESIRED_PYTHON: "3.13t"
       runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
-      build_name: manywheel-py3_13t-cuda13_0
+      build_name: manywheel-py3_13t-cuda13_2
       build_environment: linux-binary-manywheel
-      PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cublas,cufile,nvjitlink,nvtx]==13.0.2; platform_system == 'Linux' | cuda-bindings==13.0.3; platform_system == 'Linux' | nvidia-cudnn-cu13==9.17.1.4; platform_system == 'Linux' | nvidia-cusparselt-cu13==0.8.0; platform_system == 'Linux' | nvidia-nccl-cu13==2.28.9; platform_system == 'Linux' | nvidia-nvshmem-cu13==3.4.5; platform_system == 'Linux'
+      PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cublas,cufile,nvjitlink,nvtx]==13.2.1; platform_system == 'Linux' | cuda-bindings>=13.0.3,<14; platform_system == 'Linux' | nvidia-cudnn-cu13==9.20.0.48; platform_system == 'Linux' | nvidia-cusparselt-cu13==0.8.1; platform_system == 'Linux' | nvidia-nccl-cu13==2.29.7; platform_system == 'Linux' | nvidia-nvshmem-cu13==3.4.5; platform_system == 'Linux'
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
 
-  manywheel-py3_13t-cuda13_0-test:  # Testing
+  manywheel-py3_13t-cuda13_2-test:  # Testing
     if: ${{ github.repository_owner == 'pytorch' }}
     needs:
-      - manywheel-py3_13t-cuda13_0-build
+      - manywheel-py3_13t-cuda13_2-build
       - get-label-type
     uses: ./.github/workflows/_binary-test-linux.yml
     with:
@@ -3012,38 +2723,41 @@ jobs:
       PACKAGE_TYPE: manywheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu130
-      GPU_ARCH_VERSION: "13.0"
+      DESIRED_CUDA: cu132
+      GPU_ARCH_VERSION: "13.2"
       GPU_ARCH_TYPE: cuda
       DOCKER_IMAGE: manylinux2_28-builder
-      DOCKER_IMAGE_TAG_PREFIX: cuda13.0
+      DOCKER_IMAGE_TAG_PREFIX: cuda13.2
       DESIRED_PYTHON: "3.13t"
-      build_name: manywheel-py3_13t-cuda13_0
+      build_name: manywheel-py3_13t-cuda13_2
       build_environment: linux-binary-manywheel
       runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
-      runs_on: linux.g4dn.4xlarge.nvidia.gpu # 12.8+ builds need sm_70+ runner
+      runs_on: linux.g4dn.4xlarge.nvidia.gpu
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
-  manywheel-py3_13t-cuda13_0-upload:  # Uploading
+  manywheel-py3_13t-cuda13_2-upload:  # Uploading
     if: ${{ github.repository_owner == 'pytorch' }}
     permissions:
       id-token: write
       contents: read
-    needs: manywheel-py3_13t-cuda13_0-test
+    needs: manywheel-py3_13t-cuda13_2-test
     with:
       PYTORCH_ROOT: /pytorch
       PACKAGE_TYPE: manywheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu130
-      GPU_ARCH_VERSION: "13.0"
+      DESIRED_CUDA: cu132
+      GPU_ARCH_VERSION: "13.2"
       GPU_ARCH_TYPE: cuda
       DOCKER_IMAGE: manylinux2_28-builder
-      DOCKER_IMAGE_TAG_PREFIX: cuda13.0
+      DOCKER_IMAGE_TAG_PREFIX: cuda13.2
       DESIRED_PYTHON: "3.13t"
-      build_name: manywheel-py3_13t-cuda13_0
+      build_name: manywheel-py3_13t-cuda13_2
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
     uses: ./.github/workflows/_binary-upload.yml
 
   manywheel-py3_13t-rocm7_1-build:
@@ -3062,7 +2776,7 @@ jobs:
       DOCKER_IMAGE_TAG_PREFIX: rocm7.1
       DESIRED_PYTHON: "3.13t"
       runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
-      timeout-minutes: 300
+      timeout-minutes: 420
       build_name: manywheel-py3_13t-rocm7_1
       build_environment: linux-binary-manywheel
     secrets:
@@ -3073,7 +2787,7 @@ jobs:
     needs:
       - manywheel-py3_13t-rocm7_1-build
       - get-label-type
-    runs-on: linux.rocm.gpu.mi250.1
+    runs-on: linux.rocm.gpu.gfx942.1
     timeout-minutes: 240
     env:
       PYTORCH_ROOT: /pytorch
@@ -3091,51 +2805,47 @@ jobs:
       id-token: write
       contents: read
     steps:
-      - name: Setup ROCm
-        uses: ./.github/actions/setup-rocm
-      - uses: actions/download-artifact@v4.1.7
-        name: Download Build Artifacts
-        with:
-          name: manywheel-py3_13t-rocm7_1
-          path: "${{ runner.temp }}/artifacts/"
       - name: Checkout PyTorch
         uses: actions/checkout@v4
         with:
-          ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
           submodules: recursive
-          path: pytorch
           show-progress: false
       - name: Clean PyTorch checkout
         run: |
           # Remove any artifacts from the previous checkouts
           git clean -fxd
-        working-directory: pytorch
+      - name: Setup ROCm
+        uses: ./.github/actions/setup-rocm
+      - uses: actions/download-artifact@v4.1.7
+        name: Download Build Artifacts
+        with:
+          name: manywheel-py3_13t-rocm7_1
+          path: "${{ runner.temp }}/artifacts/"
       - name: ROCm set GPU_FLAG
         run: |
           echo "GPU_FLAG=--device=/dev/mem --device=/dev/kfd --device=/dev/dri --group-add video --group-add daemon" >> "${GITHUB_ENV}"
       - name: configure aws credentials
         id: aws_creds
         if: ${{ startsWith(github.event.ref, 'refs/tags/ciflow/') }}
-        uses: aws-actions/configure-aws-credentials@v4
+        uses: aws-actions/configure-aws-credentials@7474bc4690e29a8392af63c5b98e7449536d5c3a # v4
         with:
           role-to-assume: arn:aws:iam::308535385114:role/gha_workflow_s3_and_ecr_read_only
           aws-region: us-east-1
           role-duration-seconds: 18000
       - name: Calculate docker image
         id: calculate-docker-image
-        uses: pytorch/test-infra/.github/actions/calculate-docker-image@main
+        uses: pytorch/test-infra/.github/actions/calculate-docker-image@release/2.12
         with:
           docker-registry: ${{ startsWith(github.event.ref, 'refs/tags/ciflow/') && '308535385114.dkr.ecr.us-east-1.amazonaws.com' || 'docker.io' }}
           docker-image-name: manylinux2_28-builder
           custom-tag-prefix: rocm7.1
           docker-build-dir: .ci/docker
-          working-directory: pytorch
       - name: Pull Docker image
-        uses: pytorch/test-infra/.github/actions/pull-docker-image@main
+        uses: pytorch/test-infra/.github/actions/pull-docker-image@release/2.12
         with:
           docker-image: ${{ steps.calculate-docker-image.outputs.docker-image }}
       - name: Test Pytorch binary
-        uses: ./pytorch/.github/actions/test-pytorch-binary
+        uses: ./.github/actions/test-pytorch-binary
         env:
           DOCKER_IMAGE: ${{ steps.calculate-docker-image.outputs.docker-image }}
       - name: Teardown ROCm
@@ -3160,6 +2870,9 @@ jobs:
       build_name: manywheel-py3_13t-rocm7_1
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
     uses: ./.github/workflows/_binary-upload.yml
 
   manywheel-py3_13t-rocm7_2-build:
@@ -3178,7 +2891,7 @@ jobs:
       DOCKER_IMAGE_TAG_PREFIX: rocm7.2
       DESIRED_PYTHON: "3.13t"
       runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
-      timeout-minutes: 300
+      timeout-minutes: 420
       build_name: manywheel-py3_13t-rocm7_2
       build_environment: linux-binary-manywheel
     secrets:
@@ -3189,7 +2902,7 @@ jobs:
     needs:
       - manywheel-py3_13t-rocm7_2-build
       - get-label-type
-    runs-on: linux.rocm.gpu.mi250.1
+    runs-on: linux.rocm.gpu.gfx942.1
     timeout-minutes: 240
     env:
       PYTORCH_ROOT: /pytorch
@@ -3207,51 +2920,47 @@ jobs:
       id-token: write
       contents: read
     steps:
-      - name: Setup ROCm
-        uses: ./.github/actions/setup-rocm
-      - uses: actions/download-artifact@v4.1.7
-        name: Download Build Artifacts
-        with:
-          name: manywheel-py3_13t-rocm7_2
-          path: "${{ runner.temp }}/artifacts/"
       - name: Checkout PyTorch
         uses: actions/checkout@v4
         with:
-          ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
           submodules: recursive
-          path: pytorch
           show-progress: false
       - name: Clean PyTorch checkout
         run: |
           # Remove any artifacts from the previous checkouts
           git clean -fxd
-        working-directory: pytorch
+      - name: Setup ROCm
+        uses: ./.github/actions/setup-rocm
+      - uses: actions/download-artifact@v4.1.7
+        name: Download Build Artifacts
+        with:
+          name: manywheel-py3_13t-rocm7_2
+          path: "${{ runner.temp }}/artifacts/"
       - name: ROCm set GPU_FLAG
         run: |
           echo "GPU_FLAG=--device=/dev/mem --device=/dev/kfd --device=/dev/dri --group-add video --group-add daemon" >> "${GITHUB_ENV}"
       - name: configure aws credentials
         id: aws_creds
         if: ${{ startsWith(github.event.ref, 'refs/tags/ciflow/') }}
-        uses: aws-actions/configure-aws-credentials@v4
+        uses: aws-actions/configure-aws-credentials@7474bc4690e29a8392af63c5b98e7449536d5c3a # v4
         with:
           role-to-assume: arn:aws:iam::308535385114:role/gha_workflow_s3_and_ecr_read_only
           aws-region: us-east-1
           role-duration-seconds: 18000
       - name: Calculate docker image
         id: calculate-docker-image
-        uses: pytorch/test-infra/.github/actions/calculate-docker-image@main
+        uses: pytorch/test-infra/.github/actions/calculate-docker-image@release/2.12
         with:
           docker-registry: ${{ startsWith(github.event.ref, 'refs/tags/ciflow/') && '308535385114.dkr.ecr.us-east-1.amazonaws.com' || 'docker.io' }}
           docker-image-name: manylinux2_28-builder
           custom-tag-prefix: rocm7.2
           docker-build-dir: .ci/docker
-          working-directory: pytorch
       - name: Pull Docker image
-        uses: pytorch/test-infra/.github/actions/pull-docker-image@main
+        uses: pytorch/test-infra/.github/actions/pull-docker-image@release/2.12
         with:
           docker-image: ${{ steps.calculate-docker-image.outputs.docker-image }}
       - name: Test Pytorch binary
-        uses: ./pytorch/.github/actions/test-pytorch-binary
+        uses: ./.github/actions/test-pytorch-binary
         env:
           DOCKER_IMAGE: ${{ steps.calculate-docker-image.outputs.docker-image }}
       - name: Teardown ROCm
@@ -3276,6 +2985,9 @@ jobs:
       build_name: manywheel-py3_13t-rocm7_2
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
     uses: ./.github/workflows/_binary-upload.yml
 
   manywheel-py3_13t-xpu-build:
@@ -3321,42 +3033,38 @@ jobs:
       id-token: write
       contents: read
     steps:
-      - name: Setup XPU
-        uses: pytorch/pytorch/.github/actions/setup-xpu@main
-      - name: Login to ECR
-        uses: pytorch/pytorch/.github/actions/ecr-login@main
-      - uses: actions/download-artifact@v4.1.7
-        name: Download Build Artifacts
-        with:
-          name: manywheel-py3_13t-xpu
-          path: "${{ runner.temp }}/artifacts/"
       - name: Checkout PyTorch
         uses: actions/checkout@v4
         with:
-          ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
           submodules: recursive
-          path: pytorch
           show-progress: false
       - name: Clean PyTorch checkout
         run: |
           # Remove any artifacts from the previous checkouts
           git clean -fxd
-        working-directory: pytorch
+      - name: Setup XPU
+        uses: pytorch/pytorch/.github/actions/setup-xpu@release/2.12
+      - name: Login to ECR
+        uses: pytorch/pytorch/.github/actions/ecr-login@release/2.12
+      - uses: actions/download-artifact@v4.1.7
+        name: Download Build Artifacts
+        with:
+          name: manywheel-py3_13t-xpu
+          path: "${{ runner.temp }}/artifacts/"
       - name: Calculate docker image
         id: calculate-docker-image
-        uses: pytorch/test-infra/.github/actions/calculate-docker-image@main
+        uses: pytorch/test-infra/.github/actions/calculate-docker-image@release/2.12
         with:
           docker-registry: ${{ startsWith(github.event.ref, 'refs/tags/ciflow/') && '308535385114.dkr.ecr.us-east-1.amazonaws.com' || 'docker.io' }}
           docker-image-name: manylinux2_28-builder
           custom-tag-prefix: xpu
           docker-build-dir: .ci/docker
-          working-directory: pytorch
       - name: Pull Docker image
-        uses: pytorch/test-infra/.github/actions/pull-docker-image@main
+        uses: pytorch/test-infra/.github/actions/pull-docker-image@release/2.12
         with:
           docker-image: ${{ steps.calculate-docker-image.outputs.docker-image }}
       - name: Test Pytorch binary
-        uses: ./pytorch/.github/actions/test-pytorch-binary
+        uses: ./.github/actions/test-pytorch-binary
         env:
           DOCKER_IMAGE: ${{ steps.calculate-docker-image.outputs.docker-image }}
       - name: Teardown XPU
@@ -3380,6 +3088,9 @@ jobs:
       build_name: manywheel-py3_13t-xpu
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
     uses: ./.github/workflows/_binary-upload.yml
 
   manywheel-py3_14-cpu-build:
@@ -3443,6 +3154,9 @@ jobs:
       build_name: manywheel-py3_14-cpu
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
     uses: ./.github/workflows/_binary-upload.yml
 
   manywheel-py3_14-cuda12_6-build:
@@ -3463,7 +3177,7 @@ jobs:
       runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
       build_name: manywheel-py3_14-cuda12_6
       build_environment: linux-binary-manywheel
-      PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cublas,cufile,nvjitlink,nvtx]==12.6.3; platform_system == 'Linux' | cuda-bindings==12.9.4; platform_system == 'Linux' | nvidia-cudnn-cu12==9.10.2.21; platform_system == 'Linux' | nvidia-cusparselt-cu12==0.7.1; platform_system == 'Linux' | nvidia-nccl-cu12==2.28.9; platform_system == 'Linux' | nvidia-nvshmem-cu12==3.4.5; platform_system == 'Linux'
+      PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cublas,cufile,nvjitlink,nvtx]==12.6.3; platform_system == 'Linux' | cuda-bindings>=12.9.4,<13; platform_system == 'Linux' | nvidia-cudnn-cu12==9.10.2.21; platform_system == 'Linux' | nvidia-cusparselt-cu12==0.7.1; platform_system == 'Linux' | nvidia-nccl-cu12==2.29.3; platform_system == 'Linux' | nvidia-nvshmem-cu12==3.4.5; platform_system == 'Linux'
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
 
@@ -3487,7 +3201,7 @@ jobs:
       build_name: manywheel-py3_14-cuda12_6
       build_environment: linux-binary-manywheel
       runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
-      runs_on: linux.4xlarge.nvidia.gpu  # 12.6 build can use maxwell (sm_50) runner
+      runs_on: linux.g4dn.4xlarge.nvidia.gpu
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
   manywheel-py3_14-cuda12_6-upload:  # Uploading
@@ -3510,76 +3224,12 @@ jobs:
       build_name: manywheel-py3_14-cuda12_6
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
     uses: ./.github/workflows/_binary-upload.yml
 
-  manywheel-py3_14-cuda12_8-build:
-    if: ${{ github.repository_owner == 'pytorch' }}
-    uses: ./.github/workflows/_binary-build-linux.yml
-    needs: get-label-type
-    with:
-      PYTORCH_ROOT: /pytorch
-      PACKAGE_TYPE: manywheel
-      # TODO: This is a legacy variable that we eventually want to get rid of in
-      #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu128
-      GPU_ARCH_VERSION: "12.8"
-      GPU_ARCH_TYPE: cuda
-      DOCKER_IMAGE: manylinux2_28-builder
-      DOCKER_IMAGE_TAG_PREFIX: cuda12.8
-      DESIRED_PYTHON: "3.14"
-      runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
-      build_name: manywheel-py3_14-cuda12_8
-      build_environment: linux-binary-manywheel
-      PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cublas,cufile,nvjitlink,nvtx]==12.8.1; platform_system == 'Linux' | cuda-bindings==12.9.4; platform_system == 'Linux' | nvidia-cudnn-cu12==9.17.1.4; platform_system == 'Linux' | nvidia-cusparselt-cu12==0.7.1; platform_system == 'Linux' | nvidia-nccl-cu12==2.28.9; platform_system == 'Linux' | nvidia-nvshmem-cu12==3.4.5; platform_system == 'Linux'
-    secrets:
-      github-token: ${{ secrets.GITHUB_TOKEN }}
-
-  manywheel-py3_14-cuda12_8-test:  # Testing
-    if: ${{ github.repository_owner == 'pytorch' }}
-    needs:
-      - manywheel-py3_14-cuda12_8-build
-      - get-label-type
-    uses: ./.github/workflows/_binary-test-linux.yml
-    with:
-      PYTORCH_ROOT: /pytorch
-      PACKAGE_TYPE: manywheel
-      # TODO: This is a legacy variable that we eventually want to get rid of in
-      #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu128
-      GPU_ARCH_VERSION: "12.8"
-      GPU_ARCH_TYPE: cuda
-      DOCKER_IMAGE: manylinux2_28-builder
-      DOCKER_IMAGE_TAG_PREFIX: cuda12.8
-      DESIRED_PYTHON: "3.14"
-      build_name: manywheel-py3_14-cuda12_8
-      build_environment: linux-binary-manywheel
-      runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
-      runs_on: linux.g4dn.4xlarge.nvidia.gpu # 12.8+ builds need sm_70+ runner
-    secrets:
-      github-token: ${{ secrets.GITHUB_TOKEN }}
-  manywheel-py3_14-cuda12_8-upload:  # Uploading
-    if: ${{ github.repository_owner == 'pytorch' }}
-    permissions:
-      id-token: write
-      contents: read
-    needs: manywheel-py3_14-cuda12_8-test
-    with:
-      PYTORCH_ROOT: /pytorch
-      PACKAGE_TYPE: manywheel
-      # TODO: This is a legacy variable that we eventually want to get rid of in
-      #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu128
-      GPU_ARCH_VERSION: "12.8"
-      GPU_ARCH_TYPE: cuda
-      DOCKER_IMAGE: manylinux2_28-builder
-      DOCKER_IMAGE_TAG_PREFIX: cuda12.8
-      DESIRED_PYTHON: "3.14"
-      build_name: manywheel-py3_14-cuda12_8
-    secrets:
-      github-token: ${{ secrets.GITHUB_TOKEN }}
-    uses: ./.github/workflows/_binary-upload.yml
-
-  manywheel-py3_14-cuda12_9-build:
+  manywheel-py3_14-cuda13_0-build:
     if: ${{ github.repository_owner == 'pytorch' }}
     uses: ./.github/workflows/_binary-build-linux.yml
     needs: get-label-type
@@ -3588,23 +3238,23 @@ jobs:
       PACKAGE_TYPE: manywheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu129
-      GPU_ARCH_VERSION: "12.9"
+      DESIRED_CUDA: cu130
+      GPU_ARCH_VERSION: "13.0"
       GPU_ARCH_TYPE: cuda
       DOCKER_IMAGE: manylinux2_28-builder
-      DOCKER_IMAGE_TAG_PREFIX: cuda12.9
+      DOCKER_IMAGE_TAG_PREFIX: cuda13.0
       DESIRED_PYTHON: "3.14"
       runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
-      build_name: manywheel-py3_14-cuda12_9
+      build_name: manywheel-py3_14-cuda13_0
       build_environment: linux-binary-manywheel
-      PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cublas,cufile,nvjitlink,nvtx]==12.9.1; platform_system == 'Linux' | cuda-bindings==12.9.4; platform_system == 'Linux' | nvidia-cudnn-cu12==9.17.1.4; platform_system == 'Linux' | nvidia-cusparselt-cu12==0.7.1; platform_system == 'Linux' | nvidia-nccl-cu12==2.28.9; platform_system == 'Linux' | nvidia-nvshmem-cu12==3.4.5; platform_system == 'Linux'
+      PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cufile,nvjitlink,nvtx]==13.0.2; platform_system == 'Linux' | nvidia-cublas>=13.1.0.3,<=13.1.1.3; platform_system == 'Linux' | cuda-bindings>=13.0.3,<14; platform_system == 'Linux' | nvidia-cudnn-cu13==9.20.0.48; platform_system == 'Linux' | nvidia-cusparselt-cu13==0.8.1; platform_system == 'Linux' | nvidia-nccl-cu13==2.29.7; platform_system == 'Linux' | nvidia-nvshmem-cu13==3.4.5; platform_system == 'Linux'
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
 
-  manywheel-py3_14-cuda12_9-test:  # Testing
+  manywheel-py3_14-cuda13_0-test:  # Testing
     if: ${{ github.repository_owner == 'pytorch' }}
     needs:
-      - manywheel-py3_14-cuda12_9-build
+      - manywheel-py3_14-cuda13_0-build
       - get-label-type
     uses: ./.github/workflows/_binary-test-linux.yml
     with:
@@ -3612,41 +3262,44 @@ jobs:
       PACKAGE_TYPE: manywheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu129
-      GPU_ARCH_VERSION: "12.9"
+      DESIRED_CUDA: cu130
+      GPU_ARCH_VERSION: "13.0"
       GPU_ARCH_TYPE: cuda
       DOCKER_IMAGE: manylinux2_28-builder
-      DOCKER_IMAGE_TAG_PREFIX: cuda12.9
+      DOCKER_IMAGE_TAG_PREFIX: cuda13.0
       DESIRED_PYTHON: "3.14"
-      build_name: manywheel-py3_14-cuda12_9
+      build_name: manywheel-py3_14-cuda13_0
       build_environment: linux-binary-manywheel
       runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
-      runs_on: linux.g4dn.4xlarge.nvidia.gpu # 12.8+ builds need sm_70+ runner
+      runs_on: linux.g4dn.4xlarge.nvidia.gpu
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
-  manywheel-py3_14-cuda12_9-upload:  # Uploading
+  manywheel-py3_14-cuda13_0-upload:  # Uploading
     if: ${{ github.repository_owner == 'pytorch' }}
     permissions:
       id-token: write
       contents: read
-    needs: manywheel-py3_14-cuda12_9-test
+    needs: manywheel-py3_14-cuda13_0-test
     with:
       PYTORCH_ROOT: /pytorch
       PACKAGE_TYPE: manywheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu129
-      GPU_ARCH_VERSION: "12.9"
+      DESIRED_CUDA: cu130
+      GPU_ARCH_VERSION: "13.0"
       GPU_ARCH_TYPE: cuda
       DOCKER_IMAGE: manylinux2_28-builder
-      DOCKER_IMAGE_TAG_PREFIX: cuda12.9
+      DOCKER_IMAGE_TAG_PREFIX: cuda13.0
       DESIRED_PYTHON: "3.14"
-      build_name: manywheel-py3_14-cuda12_9
+      build_name: manywheel-py3_14-cuda13_0
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
     uses: ./.github/workflows/_binary-upload.yml
 
-  manywheel-py3_14-cuda13_0-build:
+  manywheel-py3_14-cuda13_2-build:
     if: ${{ github.repository_owner == 'pytorch' }}
     uses: ./.github/workflows/_binary-build-linux.yml
     needs: get-label-type
@@ -3655,23 +3308,23 @@ jobs:
       PACKAGE_TYPE: manywheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu130
-      GPU_ARCH_VERSION: "13.0"
+      DESIRED_CUDA: cu132
+      GPU_ARCH_VERSION: "13.2"
       GPU_ARCH_TYPE: cuda
       DOCKER_IMAGE: manylinux2_28-builder
-      DOCKER_IMAGE_TAG_PREFIX: cuda13.0
+      DOCKER_IMAGE_TAG_PREFIX: cuda13.2
       DESIRED_PYTHON: "3.14"
       runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
-      build_name: manywheel-py3_14-cuda13_0
+      build_name: manywheel-py3_14-cuda13_2
       build_environment: linux-binary-manywheel
-      PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cublas,cufile,nvjitlink,nvtx]==13.0.2; platform_system == 'Linux' | cuda-bindings==13.0.3; platform_system == 'Linux' | nvidia-cudnn-cu13==9.17.1.4; platform_system == 'Linux' | nvidia-cusparselt-cu13==0.8.0; platform_system == 'Linux' | nvidia-nccl-cu13==2.28.9; platform_system == 'Linux' | nvidia-nvshmem-cu13==3.4.5; platform_system == 'Linux'
+      PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cublas,cufile,nvjitlink,nvtx]==13.2.1; platform_system == 'Linux' | cuda-bindings>=13.0.3,<14; platform_system == 'Linux' | nvidia-cudnn-cu13==9.20.0.48; platform_system == 'Linux' | nvidia-cusparselt-cu13==0.8.1; platform_system == 'Linux' | nvidia-nccl-cu13==2.29.7; platform_system == 'Linux' | nvidia-nvshmem-cu13==3.4.5; platform_system == 'Linux'
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
 
-  manywheel-py3_14-cuda13_0-test:  # Testing
+  manywheel-py3_14-cuda13_2-test:  # Testing
     if: ${{ github.repository_owner == 'pytorch' }}
     needs:
-      - manywheel-py3_14-cuda13_0-build
+      - manywheel-py3_14-cuda13_2-build
       - get-label-type
     uses: ./.github/workflows/_binary-test-linux.yml
     with:
@@ -3679,38 +3332,41 @@ jobs:
       PACKAGE_TYPE: manywheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu130
-      GPU_ARCH_VERSION: "13.0"
+      DESIRED_CUDA: cu132
+      GPU_ARCH_VERSION: "13.2"
       GPU_ARCH_TYPE: cuda
       DOCKER_IMAGE: manylinux2_28-builder
-      DOCKER_IMAGE_TAG_PREFIX: cuda13.0
+      DOCKER_IMAGE_TAG_PREFIX: cuda13.2
       DESIRED_PYTHON: "3.14"
-      build_name: manywheel-py3_14-cuda13_0
+      build_name: manywheel-py3_14-cuda13_2
       build_environment: linux-binary-manywheel
       runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
-      runs_on: linux.g4dn.4xlarge.nvidia.gpu # 12.8+ builds need sm_70+ runner
+      runs_on: linux.g4dn.4xlarge.nvidia.gpu
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
-  manywheel-py3_14-cuda13_0-upload:  # Uploading
+  manywheel-py3_14-cuda13_2-upload:  # Uploading
     if: ${{ github.repository_owner == 'pytorch' }}
     permissions:
       id-token: write
       contents: read
-    needs: manywheel-py3_14-cuda13_0-test
+    needs: manywheel-py3_14-cuda13_2-test
     with:
       PYTORCH_ROOT: /pytorch
       PACKAGE_TYPE: manywheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu130
-      GPU_ARCH_VERSION: "13.0"
+      DESIRED_CUDA: cu132
+      GPU_ARCH_VERSION: "13.2"
       GPU_ARCH_TYPE: cuda
       DOCKER_IMAGE: manylinux2_28-builder
-      DOCKER_IMAGE_TAG_PREFIX: cuda13.0
+      DOCKER_IMAGE_TAG_PREFIX: cuda13.2
       DESIRED_PYTHON: "3.14"
-      build_name: manywheel-py3_14-cuda13_0
+      build_name: manywheel-py3_14-cuda13_2
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
     uses: ./.github/workflows/_binary-upload.yml
 
   manywheel-py3_14-rocm7_1-build:
@@ -3729,7 +3385,7 @@ jobs:
       DOCKER_IMAGE_TAG_PREFIX: rocm7.1
       DESIRED_PYTHON: "3.14"
       runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
-      timeout-minutes: 300
+      timeout-minutes: 420
       build_name: manywheel-py3_14-rocm7_1
       build_environment: linux-binary-manywheel
     secrets:
@@ -3740,7 +3396,7 @@ jobs:
     needs:
       - manywheel-py3_14-rocm7_1-build
       - get-label-type
-    runs-on: linux.rocm.gpu.mi250.1
+    runs-on: linux.rocm.gpu.gfx942.1
     timeout-minutes: 240
     env:
       PYTORCH_ROOT: /pytorch
@@ -3758,51 +3414,47 @@ jobs:
       id-token: write
       contents: read
     steps:
-      - name: Setup ROCm
-        uses: ./.github/actions/setup-rocm
-      - uses: actions/download-artifact@v4.1.7
-        name: Download Build Artifacts
-        with:
-          name: manywheel-py3_14-rocm7_1
-          path: "${{ runner.temp }}/artifacts/"
       - name: Checkout PyTorch
         uses: actions/checkout@v4
         with:
-          ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
           submodules: recursive
-          path: pytorch
           show-progress: false
       - name: Clean PyTorch checkout
         run: |
           # Remove any artifacts from the previous checkouts
           git clean -fxd
-        working-directory: pytorch
+      - name: Setup ROCm
+        uses: ./.github/actions/setup-rocm
+      - uses: actions/download-artifact@v4.1.7
+        name: Download Build Artifacts
+        with:
+          name: manywheel-py3_14-rocm7_1
+          path: "${{ runner.temp }}/artifacts/"
       - name: ROCm set GPU_FLAG
         run: |
           echo "GPU_FLAG=--device=/dev/mem --device=/dev/kfd --device=/dev/dri --group-add video --group-add daemon" >> "${GITHUB_ENV}"
       - name: configure aws credentials
         id: aws_creds
         if: ${{ startsWith(github.event.ref, 'refs/tags/ciflow/') }}
-        uses: aws-actions/configure-aws-credentials@v4
+        uses: aws-actions/configure-aws-credentials@7474bc4690e29a8392af63c5b98e7449536d5c3a # v4
         with:
           role-to-assume: arn:aws:iam::308535385114:role/gha_workflow_s3_and_ecr_read_only
           aws-region: us-east-1
           role-duration-seconds: 18000
       - name: Calculate docker image
         id: calculate-docker-image
-        uses: pytorch/test-infra/.github/actions/calculate-docker-image@main
+        uses: pytorch/test-infra/.github/actions/calculate-docker-image@release/2.12
         with:
           docker-registry: ${{ startsWith(github.event.ref, 'refs/tags/ciflow/') && '308535385114.dkr.ecr.us-east-1.amazonaws.com' || 'docker.io' }}
           docker-image-name: manylinux2_28-builder
           custom-tag-prefix: rocm7.1
           docker-build-dir: .ci/docker
-          working-directory: pytorch
       - name: Pull Docker image
-        uses: pytorch/test-infra/.github/actions/pull-docker-image@main
+        uses: pytorch/test-infra/.github/actions/pull-docker-image@release/2.12
         with:
           docker-image: ${{ steps.calculate-docker-image.outputs.docker-image }}
       - name: Test Pytorch binary
-        uses: ./pytorch/.github/actions/test-pytorch-binary
+        uses: ./.github/actions/test-pytorch-binary
         env:
           DOCKER_IMAGE: ${{ steps.calculate-docker-image.outputs.docker-image }}
       - name: Teardown ROCm
@@ -3827,6 +3479,9 @@ jobs:
       build_name: manywheel-py3_14-rocm7_1
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
     uses: ./.github/workflows/_binary-upload.yml
 
   manywheel-py3_14-rocm7_2-build:
@@ -3845,7 +3500,7 @@ jobs:
       DOCKER_IMAGE_TAG_PREFIX: rocm7.2
       DESIRED_PYTHON: "3.14"
       runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
-      timeout-minutes: 300
+      timeout-minutes: 420
       build_name: manywheel-py3_14-rocm7_2
       build_environment: linux-binary-manywheel
     secrets:
@@ -3856,7 +3511,7 @@ jobs:
     needs:
       - manywheel-py3_14-rocm7_2-build
       - get-label-type
-    runs-on: linux.rocm.gpu.mi250.1
+    runs-on: linux.rocm.gpu.gfx942.1
     timeout-minutes: 240
     env:
       PYTORCH_ROOT: /pytorch
@@ -3874,51 +3529,47 @@ jobs:
       id-token: write
       contents: read
     steps:
-      - name: Setup ROCm
-        uses: ./.github/actions/setup-rocm
-      - uses: actions/download-artifact@v4.1.7
-        name: Download Build Artifacts
-        with:
-          name: manywheel-py3_14-rocm7_2
-          path: "${{ runner.temp }}/artifacts/"
       - name: Checkout PyTorch
         uses: actions/checkout@v4
         with:
-          ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
           submodules: recursive
-          path: pytorch
           show-progress: false
       - name: Clean PyTorch checkout
         run: |
           # Remove any artifacts from the previous checkouts
           git clean -fxd
-        working-directory: pytorch
+      - name: Setup ROCm
+        uses: ./.github/actions/setup-rocm
+      - uses: actions/download-artifact@v4.1.7
+        name: Download Build Artifacts
+        with:
+          name: manywheel-py3_14-rocm7_2
+          path: "${{ runner.temp }}/artifacts/"
       - name: ROCm set GPU_FLAG
         run: |
           echo "GPU_FLAG=--device=/dev/mem --device=/dev/kfd --device=/dev/dri --group-add video --group-add daemon" >> "${GITHUB_ENV}"
       - name: configure aws credentials
         id: aws_creds
         if: ${{ startsWith(github.event.ref, 'refs/tags/ciflow/') }}
-        uses: aws-actions/configure-aws-credentials@v4
+        uses: aws-actions/configure-aws-credentials@7474bc4690e29a8392af63c5b98e7449536d5c3a # v4
         with:
           role-to-assume: arn:aws:iam::308535385114:role/gha_workflow_s3_and_ecr_read_only
           aws-region: us-east-1
           role-duration-seconds: 18000
       - name: Calculate docker image
         id: calculate-docker-image
-        uses: pytorch/test-infra/.github/actions/calculate-docker-image@main
+        uses: pytorch/test-infra/.github/actions/calculate-docker-image@release/2.12
         with:
           docker-registry: ${{ startsWith(github.event.ref, 'refs/tags/ciflow/') && '308535385114.dkr.ecr.us-east-1.amazonaws.com' || 'docker.io' }}
           docker-image-name: manylinux2_28-builder
           custom-tag-prefix: rocm7.2
           docker-build-dir: .ci/docker
-          working-directory: pytorch
       - name: Pull Docker image
-        uses: pytorch/test-infra/.github/actions/pull-docker-image@main
+        uses: pytorch/test-infra/.github/actions/pull-docker-image@release/2.12
         with:
           docker-image: ${{ steps.calculate-docker-image.outputs.docker-image }}
       - name: Test Pytorch binary
-        uses: ./pytorch/.github/actions/test-pytorch-binary
+        uses: ./.github/actions/test-pytorch-binary
         env:
           DOCKER_IMAGE: ${{ steps.calculate-docker-image.outputs.docker-image }}
       - name: Teardown ROCm
@@ -3943,6 +3594,9 @@ jobs:
       build_name: manywheel-py3_14-rocm7_2
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
     uses: ./.github/workflows/_binary-upload.yml
 
   manywheel-py3_14-xpu-build:
@@ -3988,42 +3642,38 @@ jobs:
       id-token: write
       contents: read
     steps:
-      - name: Setup XPU
-        uses: pytorch/pytorch/.github/actions/setup-xpu@main
-      - name: Login to ECR
-        uses: pytorch/pytorch/.github/actions/ecr-login@main
-      - uses: actions/download-artifact@v4.1.7
-        name: Download Build Artifacts
-        with:
-          name: manywheel-py3_14-xpu
-          path: "${{ runner.temp }}/artifacts/"
       - name: Checkout PyTorch
         uses: actions/checkout@v4
         with:
-          ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
           submodules: recursive
-          path: pytorch
           show-progress: false
       - name: Clean PyTorch checkout
         run: |
           # Remove any artifacts from the previous checkouts
           git clean -fxd
-        working-directory: pytorch
+      - name: Setup XPU
+        uses: pytorch/pytorch/.github/actions/setup-xpu@release/2.12
+      - name: Login to ECR
+        uses: pytorch/pytorch/.github/actions/ecr-login@release/2.12
+      - uses: actions/download-artifact@v4.1.7
+        name: Download Build Artifacts
+        with:
+          name: manywheel-py3_14-xpu
+          path: "${{ runner.temp }}/artifacts/"
       - name: Calculate docker image
         id: calculate-docker-image
-        uses: pytorch/test-infra/.github/actions/calculate-docker-image@main
+        uses: pytorch/test-infra/.github/actions/calculate-docker-image@release/2.12
         with:
           docker-registry: ${{ startsWith(github.event.ref, 'refs/tags/ciflow/') && '308535385114.dkr.ecr.us-east-1.amazonaws.com' || 'docker.io' }}
           docker-image-name: manylinux2_28-builder
           custom-tag-prefix: xpu
           docker-build-dir: .ci/docker
-          working-directory: pytorch
       - name: Pull Docker image
-        uses: pytorch/test-infra/.github/actions/pull-docker-image@main
+        uses: pytorch/test-infra/.github/actions/pull-docker-image@release/2.12
         with:
           docker-image: ${{ steps.calculate-docker-image.outputs.docker-image }}
       - name: Test Pytorch binary
-        uses: ./pytorch/.github/actions/test-pytorch-binary
+        uses: ./.github/actions/test-pytorch-binary
         env:
           DOCKER_IMAGE: ${{ steps.calculate-docker-image.outputs.docker-image }}
       - name: Teardown XPU
@@ -4047,6 +3697,9 @@ jobs:
       build_name: manywheel-py3_14-xpu
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
     uses: ./.github/workflows/_binary-upload.yml
 
   manywheel-py3_14t-cpu-build:
@@ -4110,6 +3763,9 @@ jobs:
       build_name: manywheel-py3_14t-cpu
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
     uses: ./.github/workflows/_binary-upload.yml
 
   manywheel-py3_14t-cuda12_6-build:
@@ -4130,7 +3786,7 @@ jobs:
       runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
       build_name: manywheel-py3_14t-cuda12_6
       build_environment: linux-binary-manywheel
-      PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cublas,cufile,nvjitlink,nvtx]==12.6.3; platform_system == 'Linux' | cuda-bindings==12.9.4; platform_system == 'Linux' | nvidia-cudnn-cu12==9.10.2.21; platform_system == 'Linux' | nvidia-cusparselt-cu12==0.7.1; platform_system == 'Linux' | nvidia-nccl-cu12==2.28.9; platform_system == 'Linux' | nvidia-nvshmem-cu12==3.4.5; platform_system == 'Linux'
+      PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cublas,cufile,nvjitlink,nvtx]==12.6.3; platform_system == 'Linux' | cuda-bindings>=12.9.4,<13; platform_system == 'Linux' | nvidia-cudnn-cu12==9.10.2.21; platform_system == 'Linux' | nvidia-cusparselt-cu12==0.7.1; platform_system == 'Linux' | nvidia-nccl-cu12==2.29.3; platform_system == 'Linux' | nvidia-nvshmem-cu12==3.4.5; platform_system == 'Linux'
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
 
@@ -4154,7 +3810,7 @@ jobs:
       build_name: manywheel-py3_14t-cuda12_6
       build_environment: linux-binary-manywheel
       runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
-      runs_on: linux.4xlarge.nvidia.gpu  # 12.6 build can use maxwell (sm_50) runner
+      runs_on: linux.g4dn.4xlarge.nvidia.gpu
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
   manywheel-py3_14t-cuda12_6-upload:  # Uploading
@@ -4177,9 +3833,12 @@ jobs:
       build_name: manywheel-py3_14t-cuda12_6
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
     uses: ./.github/workflows/_binary-upload.yml
 
-  manywheel-py3_14t-cuda12_8-build:
+  manywheel-py3_14t-cuda13_0-build:
     if: ${{ github.repository_owner == 'pytorch' }}
     uses: ./.github/workflows/_binary-build-linux.yml
     needs: get-label-type
@@ -4188,90 +3847,23 @@ jobs:
       PACKAGE_TYPE: manywheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu128
-      GPU_ARCH_VERSION: "12.8"
+      DESIRED_CUDA: cu130
+      GPU_ARCH_VERSION: "13.0"
       GPU_ARCH_TYPE: cuda
       DOCKER_IMAGE: manylinux2_28-builder
-      DOCKER_IMAGE_TAG_PREFIX: cuda12.8
+      DOCKER_IMAGE_TAG_PREFIX: cuda13.0
       DESIRED_PYTHON: "3.14t"
       runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
-      build_name: manywheel-py3_14t-cuda12_8
+      build_name: manywheel-py3_14t-cuda13_0
       build_environment: linux-binary-manywheel
-      PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cublas,cufile,nvjitlink,nvtx]==12.8.1; platform_system == 'Linux' | cuda-bindings==12.9.4; platform_system == 'Linux' | nvidia-cudnn-cu12==9.17.1.4; platform_system == 'Linux' | nvidia-cusparselt-cu12==0.7.1; platform_system == 'Linux' | nvidia-nccl-cu12==2.28.9; platform_system == 'Linux' | nvidia-nvshmem-cu12==3.4.5; platform_system == 'Linux'
+      PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cufile,nvjitlink,nvtx]==13.0.2; platform_system == 'Linux' | nvidia-cublas>=13.1.0.3,<=13.1.1.3; platform_system == 'Linux' | cuda-bindings>=13.0.3,<14; platform_system == 'Linux' | nvidia-cudnn-cu13==9.20.0.48; platform_system == 'Linux' | nvidia-cusparselt-cu13==0.8.1; platform_system == 'Linux' | nvidia-nccl-cu13==2.29.7; platform_system == 'Linux' | nvidia-nvshmem-cu13==3.4.5; platform_system == 'Linux'
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
 
-  manywheel-py3_14t-cuda12_8-test:  # Testing
+  manywheel-py3_14t-cuda13_0-test:  # Testing
     if: ${{ github.repository_owner == 'pytorch' }}
     needs:
-      - manywheel-py3_14t-cuda12_8-build
-      - get-label-type
-    uses: ./.github/workflows/_binary-test-linux.yml
-    with:
-      PYTORCH_ROOT: /pytorch
-      PACKAGE_TYPE: manywheel
-      # TODO: This is a legacy variable that we eventually want to get rid of in
-      #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu128
-      GPU_ARCH_VERSION: "12.8"
-      GPU_ARCH_TYPE: cuda
-      DOCKER_IMAGE: manylinux2_28-builder
-      DOCKER_IMAGE_TAG_PREFIX: cuda12.8
-      DESIRED_PYTHON: "3.14t"
-      build_name: manywheel-py3_14t-cuda12_8
-      build_environment: linux-binary-manywheel
-      runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
-      runs_on: linux.g4dn.4xlarge.nvidia.gpu # 12.8+ builds need sm_70+ runner
-    secrets:
-      github-token: ${{ secrets.GITHUB_TOKEN }}
-  manywheel-py3_14t-cuda12_8-upload:  # Uploading
-    if: ${{ github.repository_owner == 'pytorch' }}
-    permissions:
-      id-token: write
-      contents: read
-    needs: manywheel-py3_14t-cuda12_8-test
-    with:
-      PYTORCH_ROOT: /pytorch
-      PACKAGE_TYPE: manywheel
-      # TODO: This is a legacy variable that we eventually want to get rid of in
-      #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu128
-      GPU_ARCH_VERSION: "12.8"
-      GPU_ARCH_TYPE: cuda
-      DOCKER_IMAGE: manylinux2_28-builder
-      DOCKER_IMAGE_TAG_PREFIX: cuda12.8
-      DESIRED_PYTHON: "3.14t"
-      build_name: manywheel-py3_14t-cuda12_8
-    secrets:
-      github-token: ${{ secrets.GITHUB_TOKEN }}
-    uses: ./.github/workflows/_binary-upload.yml
-
-  manywheel-py3_14t-cuda12_9-build:
-    if: ${{ github.repository_owner == 'pytorch' }}
-    uses: ./.github/workflows/_binary-build-linux.yml
-    needs: get-label-type
-    with:
-      PYTORCH_ROOT: /pytorch
-      PACKAGE_TYPE: manywheel
-      # TODO: This is a legacy variable that we eventually want to get rid of in
-      #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu129
-      GPU_ARCH_VERSION: "12.9"
-      GPU_ARCH_TYPE: cuda
-      DOCKER_IMAGE: manylinux2_28-builder
-      DOCKER_IMAGE_TAG_PREFIX: cuda12.9
-      DESIRED_PYTHON: "3.14t"
-      runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
-      build_name: manywheel-py3_14t-cuda12_9
-      build_environment: linux-binary-manywheel
-      PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cublas,cufile,nvjitlink,nvtx]==12.9.1; platform_system == 'Linux' | cuda-bindings==12.9.4; platform_system == 'Linux' | nvidia-cudnn-cu12==9.17.1.4; platform_system == 'Linux' | nvidia-cusparselt-cu12==0.7.1; platform_system == 'Linux' | nvidia-nccl-cu12==2.28.9; platform_system == 'Linux' | nvidia-nvshmem-cu12==3.4.5; platform_system == 'Linux'
-    secrets:
-      github-token: ${{ secrets.GITHUB_TOKEN }}
-
-  manywheel-py3_14t-cuda12_9-test:  # Testing
-    if: ${{ github.repository_owner == 'pytorch' }}
-    needs:
-      - manywheel-py3_14t-cuda12_9-build
+      - manywheel-py3_14t-cuda13_0-build
       - get-label-type
     uses: ./.github/workflows/_binary-test-linux.yml
     with:
@@ -4279,41 +3871,44 @@ jobs:
       PACKAGE_TYPE: manywheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu129
-      GPU_ARCH_VERSION: "12.9"
+      DESIRED_CUDA: cu130
+      GPU_ARCH_VERSION: "13.0"
       GPU_ARCH_TYPE: cuda
       DOCKER_IMAGE: manylinux2_28-builder
-      DOCKER_IMAGE_TAG_PREFIX: cuda12.9
+      DOCKER_IMAGE_TAG_PREFIX: cuda13.0
       DESIRED_PYTHON: "3.14t"
-      build_name: manywheel-py3_14t-cuda12_9
+      build_name: manywheel-py3_14t-cuda13_0
       build_environment: linux-binary-manywheel
       runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
-      runs_on: linux.g4dn.4xlarge.nvidia.gpu # 12.8+ builds need sm_70+ runner
+      runs_on: linux.g4dn.4xlarge.nvidia.gpu
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
-  manywheel-py3_14t-cuda12_9-upload:  # Uploading
+  manywheel-py3_14t-cuda13_0-upload:  # Uploading
     if: ${{ github.repository_owner == 'pytorch' }}
     permissions:
       id-token: write
       contents: read
-    needs: manywheel-py3_14t-cuda12_9-test
+    needs: manywheel-py3_14t-cuda13_0-test
     with:
       PYTORCH_ROOT: /pytorch
       PACKAGE_TYPE: manywheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu129
-      GPU_ARCH_VERSION: "12.9"
+      DESIRED_CUDA: cu130
+      GPU_ARCH_VERSION: "13.0"
       GPU_ARCH_TYPE: cuda
       DOCKER_IMAGE: manylinux2_28-builder
-      DOCKER_IMAGE_TAG_PREFIX: cuda12.9
+      DOCKER_IMAGE_TAG_PREFIX: cuda13.0
       DESIRED_PYTHON: "3.14t"
-      build_name: manywheel-py3_14t-cuda12_9
+      build_name: manywheel-py3_14t-cuda13_0
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
     uses: ./.github/workflows/_binary-upload.yml
 
-  manywheel-py3_14t-cuda13_0-build:
+  manywheel-py3_14t-cuda13_2-build:
     if: ${{ github.repository_owner == 'pytorch' }}
     uses: ./.github/workflows/_binary-build-linux.yml
     needs: get-label-type
@@ -4322,23 +3917,23 @@ jobs:
       PACKAGE_TYPE: manywheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu130
-      GPU_ARCH_VERSION: "13.0"
+      DESIRED_CUDA: cu132
+      GPU_ARCH_VERSION: "13.2"
       GPU_ARCH_TYPE: cuda
       DOCKER_IMAGE: manylinux2_28-builder
-      DOCKER_IMAGE_TAG_PREFIX: cuda13.0
+      DOCKER_IMAGE_TAG_PREFIX: cuda13.2
       DESIRED_PYTHON: "3.14t"
       runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
-      build_name: manywheel-py3_14t-cuda13_0
+      build_name: manywheel-py3_14t-cuda13_2
       build_environment: linux-binary-manywheel
-      PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cublas,cufile,nvjitlink,nvtx]==13.0.2; platform_system == 'Linux' | cuda-bindings==13.0.3; platform_system == 'Linux' | nvidia-cudnn-cu13==9.17.1.4; platform_system == 'Linux' | nvidia-cusparselt-cu13==0.8.0; platform_system == 'Linux' | nvidia-nccl-cu13==2.28.9; platform_system == 'Linux' | nvidia-nvshmem-cu13==3.4.5; platform_system == 'Linux'
+      PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cublas,cufile,nvjitlink,nvtx]==13.2.1; platform_system == 'Linux' | cuda-bindings>=13.0.3,<14; platform_system == 'Linux' | nvidia-cudnn-cu13==9.20.0.48; platform_system == 'Linux' | nvidia-cusparselt-cu13==0.8.1; platform_system == 'Linux' | nvidia-nccl-cu13==2.29.7; platform_system == 'Linux' | nvidia-nvshmem-cu13==3.4.5; platform_system == 'Linux'
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
 
-  manywheel-py3_14t-cuda13_0-test:  # Testing
+  manywheel-py3_14t-cuda13_2-test:  # Testing
     if: ${{ github.repository_owner == 'pytorch' }}
     needs:
-      - manywheel-py3_14t-cuda13_0-build
+      - manywheel-py3_14t-cuda13_2-build
       - get-label-type
     uses: ./.github/workflows/_binary-test-linux.yml
     with:
@@ -4346,38 +3941,41 @@ jobs:
       PACKAGE_TYPE: manywheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu130
-      GPU_ARCH_VERSION: "13.0"
+      DESIRED_CUDA: cu132
+      GPU_ARCH_VERSION: "13.2"
       GPU_ARCH_TYPE: cuda
       DOCKER_IMAGE: manylinux2_28-builder
-      DOCKER_IMAGE_TAG_PREFIX: cuda13.0
+      DOCKER_IMAGE_TAG_PREFIX: cuda13.2
       DESIRED_PYTHON: "3.14t"
-      build_name: manywheel-py3_14t-cuda13_0
+      build_name: manywheel-py3_14t-cuda13_2
       build_environment: linux-binary-manywheel
       runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
-      runs_on: linux.g4dn.4xlarge.nvidia.gpu # 12.8+ builds need sm_70+ runner
+      runs_on: linux.g4dn.4xlarge.nvidia.gpu
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
-  manywheel-py3_14t-cuda13_0-upload:  # Uploading
+  manywheel-py3_14t-cuda13_2-upload:  # Uploading
     if: ${{ github.repository_owner == 'pytorch' }}
     permissions:
       id-token: write
       contents: read
-    needs: manywheel-py3_14t-cuda13_0-test
+    needs: manywheel-py3_14t-cuda13_2-test
     with:
       PYTORCH_ROOT: /pytorch
       PACKAGE_TYPE: manywheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu130
-      GPU_ARCH_VERSION: "13.0"
+      DESIRED_CUDA: cu132
+      GPU_ARCH_VERSION: "13.2"
       GPU_ARCH_TYPE: cuda
       DOCKER_IMAGE: manylinux2_28-builder
-      DOCKER_IMAGE_TAG_PREFIX: cuda13.0
+      DOCKER_IMAGE_TAG_PREFIX: cuda13.2
       DESIRED_PYTHON: "3.14t"
-      build_name: manywheel-py3_14t-cuda13_0
+      build_name: manywheel-py3_14t-cuda13_2
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
     uses: ./.github/workflows/_binary-upload.yml
 
   manywheel-py3_14t-rocm7_1-build:
@@ -4396,7 +3994,7 @@ jobs:
       DOCKER_IMAGE_TAG_PREFIX: rocm7.1
       DESIRED_PYTHON: "3.14t"
       runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
-      timeout-minutes: 300
+      timeout-minutes: 420
       build_name: manywheel-py3_14t-rocm7_1
       build_environment: linux-binary-manywheel
     secrets:
@@ -4407,7 +4005,7 @@ jobs:
     needs:
       - manywheel-py3_14t-rocm7_1-build
       - get-label-type
-    runs-on: linux.rocm.gpu.mi250.1
+    runs-on: linux.rocm.gpu.gfx942.1
     timeout-minutes: 240
     env:
       PYTORCH_ROOT: /pytorch
@@ -4425,51 +4023,47 @@ jobs:
       id-token: write
       contents: read
     steps:
-      - name: Setup ROCm
-        uses: ./.github/actions/setup-rocm
-      - uses: actions/download-artifact@v4.1.7
-        name: Download Build Artifacts
-        with:
-          name: manywheel-py3_14t-rocm7_1
-          path: "${{ runner.temp }}/artifacts/"
       - name: Checkout PyTorch
         uses: actions/checkout@v4
         with:
-          ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
           submodules: recursive
-          path: pytorch
           show-progress: false
       - name: Clean PyTorch checkout
         run: |
           # Remove any artifacts from the previous checkouts
           git clean -fxd
-        working-directory: pytorch
+      - name: Setup ROCm
+        uses: ./.github/actions/setup-rocm
+      - uses: actions/download-artifact@v4.1.7
+        name: Download Build Artifacts
+        with:
+          name: manywheel-py3_14t-rocm7_1
+          path: "${{ runner.temp }}/artifacts/"
       - name: ROCm set GPU_FLAG
         run: |
           echo "GPU_FLAG=--device=/dev/mem --device=/dev/kfd --device=/dev/dri --group-add video --group-add daemon" >> "${GITHUB_ENV}"
       - name: configure aws credentials
         id: aws_creds
         if: ${{ startsWith(github.event.ref, 'refs/tags/ciflow/') }}
-        uses: aws-actions/configure-aws-credentials@v4
+        uses: aws-actions/configure-aws-credentials@7474bc4690e29a8392af63c5b98e7449536d5c3a # v4
         with:
           role-to-assume: arn:aws:iam::308535385114:role/gha_workflow_s3_and_ecr_read_only
           aws-region: us-east-1
           role-duration-seconds: 18000
       - name: Calculate docker image
         id: calculate-docker-image
-        uses: pytorch/test-infra/.github/actions/calculate-docker-image@main
+        uses: pytorch/test-infra/.github/actions/calculate-docker-image@release/2.12
         with:
           docker-registry: ${{ startsWith(github.event.ref, 'refs/tags/ciflow/') && '308535385114.dkr.ecr.us-east-1.amazonaws.com' || 'docker.io' }}
           docker-image-name: manylinux2_28-builder
           custom-tag-prefix: rocm7.1
           docker-build-dir: .ci/docker
-          working-directory: pytorch
       - name: Pull Docker image
-        uses: pytorch/test-infra/.github/actions/pull-docker-image@main
+        uses: pytorch/test-infra/.github/actions/pull-docker-image@release/2.12
         with:
           docker-image: ${{ steps.calculate-docker-image.outputs.docker-image }}
       - name: Test Pytorch binary
-        uses: ./pytorch/.github/actions/test-pytorch-binary
+        uses: ./.github/actions/test-pytorch-binary
         env:
           DOCKER_IMAGE: ${{ steps.calculate-docker-image.outputs.docker-image }}
       - name: Teardown ROCm
@@ -4494,6 +4088,9 @@ jobs:
       build_name: manywheel-py3_14t-rocm7_1
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
     uses: ./.github/workflows/_binary-upload.yml
 
   manywheel-py3_14t-rocm7_2-build:
@@ -4512,7 +4109,7 @@ jobs:
       DOCKER_IMAGE_TAG_PREFIX: rocm7.2
       DESIRED_PYTHON: "3.14t"
       runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
-      timeout-minutes: 300
+      timeout-minutes: 420
       build_name: manywheel-py3_14t-rocm7_2
       build_environment: linux-binary-manywheel
     secrets:
@@ -4523,7 +4120,7 @@ jobs:
     needs:
       - manywheel-py3_14t-rocm7_2-build
       - get-label-type
-    runs-on: linux.rocm.gpu.mi250.1
+    runs-on: linux.rocm.gpu.gfx942.1
     timeout-minutes: 240
     env:
       PYTORCH_ROOT: /pytorch
@@ -4541,51 +4138,47 @@ jobs:
       id-token: write
       contents: read
     steps:
-      - name: Setup ROCm
-        uses: ./.github/actions/setup-rocm
-      - uses: actions/download-artifact@v4.1.7
-        name: Download Build Artifacts
-        with:
-          name: manywheel-py3_14t-rocm7_2
-          path: "${{ runner.temp }}/artifacts/"
       - name: Checkout PyTorch
         uses: actions/checkout@v4
         with:
-          ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
           submodules: recursive
-          path: pytorch
           show-progress: false
       - name: Clean PyTorch checkout
         run: |
           # Remove any artifacts from the previous checkouts
           git clean -fxd
-        working-directory: pytorch
+      - name: Setup ROCm
+        uses: ./.github/actions/setup-rocm
+      - uses: actions/download-artifact@v4.1.7
+        name: Download Build Artifacts
+        with:
+          name: manywheel-py3_14t-rocm7_2
+          path: "${{ runner.temp }}/artifacts/"
       - name: ROCm set GPU_FLAG
         run: |
           echo "GPU_FLAG=--device=/dev/mem --device=/dev/kfd --device=/dev/dri --group-add video --group-add daemon" >> "${GITHUB_ENV}"
       - name: configure aws credentials
         id: aws_creds
         if: ${{ startsWith(github.event.ref, 'refs/tags/ciflow/') }}
-        uses: aws-actions/configure-aws-credentials@v4
+        uses: aws-actions/configure-aws-credentials@7474bc4690e29a8392af63c5b98e7449536d5c3a # v4
         with:
           role-to-assume: arn:aws:iam::308535385114:role/gha_workflow_s3_and_ecr_read_only
           aws-region: us-east-1
           role-duration-seconds: 18000
       - name: Calculate docker image
         id: calculate-docker-image
-        uses: pytorch/test-infra/.github/actions/calculate-docker-image@main
+        uses: pytorch/test-infra/.github/actions/calculate-docker-image@release/2.12
         with:
           docker-registry: ${{ startsWith(github.event.ref, 'refs/tags/ciflow/') && '308535385114.dkr.ecr.us-east-1.amazonaws.com' || 'docker.io' }}
           docker-image-name: manylinux2_28-builder
           custom-tag-prefix: rocm7.2
           docker-build-dir: .ci/docker
-          working-directory: pytorch
       - name: Pull Docker image
-        uses: pytorch/test-infra/.github/actions/pull-docker-image@main
+        uses: pytorch/test-infra/.github/actions/pull-docker-image@release/2.12
         with:
           docker-image: ${{ steps.calculate-docker-image.outputs.docker-image }}
       - name: Test Pytorch binary
-        uses: ./pytorch/.github/actions/test-pytorch-binary
+        uses: ./.github/actions/test-pytorch-binary
         env:
           DOCKER_IMAGE: ${{ steps.calculate-docker-image.outputs.docker-image }}
       - name: Teardown ROCm
@@ -4610,6 +4203,9 @@ jobs:
       build_name: manywheel-py3_14t-rocm7_2
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
     uses: ./.github/workflows/_binary-upload.yml
 
   manywheel-py3_14t-xpu-build:
@@ -4655,42 +4251,38 @@ jobs:
       id-token: write
       contents: read
     steps:
-      - name: Setup XPU
-        uses: pytorch/pytorch/.github/actions/setup-xpu@main
-      - name: Login to ECR
-        uses: pytorch/pytorch/.github/actions/ecr-login@main
-      - uses: actions/download-artifact@v4.1.7
-        name: Download Build Artifacts
-        with:
-          name: manywheel-py3_14t-xpu
-          path: "${{ runner.temp }}/artifacts/"
       - name: Checkout PyTorch
         uses: actions/checkout@v4
         with:
-          ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
           submodules: recursive
-          path: pytorch
           show-progress: false
       - name: Clean PyTorch checkout
         run: |
           # Remove any artifacts from the previous checkouts
           git clean -fxd
-        working-directory: pytorch
+      - name: Setup XPU
+        uses: pytorch/pytorch/.github/actions/setup-xpu@release/2.12
+      - name: Login to ECR
+        uses: pytorch/pytorch/.github/actions/ecr-login@release/2.12
+      - uses: actions/download-artifact@v4.1.7
+        name: Download Build Artifacts
+        with:
+          name: manywheel-py3_14t-xpu
+          path: "${{ runner.temp }}/artifacts/"
       - name: Calculate docker image
         id: calculate-docker-image
-        uses: pytorch/test-infra/.github/actions/calculate-docker-image@main
+        uses: pytorch/test-infra/.github/actions/calculate-docker-image@release/2.12
         with:
           docker-registry: ${{ startsWith(github.event.ref, 'refs/tags/ciflow/') && '308535385114.dkr.ecr.us-east-1.amazonaws.com' || 'docker.io' }}
           docker-image-name: manylinux2_28-builder
           custom-tag-prefix: xpu
           docker-build-dir: .ci/docker
-          working-directory: pytorch
       - name: Pull Docker image
-        uses: pytorch/test-infra/.github/actions/pull-docker-image@main
+        uses: pytorch/test-infra/.github/actions/pull-docker-image@release/2.12
         with:
           docker-image: ${{ steps.calculate-docker-image.outputs.docker-image }}
       - name: Test Pytorch binary
-        uses: ./pytorch/.github/actions/test-pytorch-binary
+        uses: ./.github/actions/test-pytorch-binary
         env:
           DOCKER_IMAGE: ${{ steps.calculate-docker-image.outputs.docker-image }}
       - name: Teardown XPU
@@ -4714,4 +4306,372 @@ jobs:
       build_name: manywheel-py3_14t-xpu
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
+    uses: ./.github/workflows/_binary-upload.yml
+
+  libtorch-cpu-shared-with-deps-release-extract:
+    if: ${{ github.repository_owner == 'pytorch' }}
+    needs:
+      - manywheel-py3_10-cpu-build
+      - get-label-type
+    runs-on: "${{ needs.get-label-type.outputs.label-type }}linux.4xlarge"
+    timeout-minutes: 60
+    env:
+      DESIRED_CUDA: cpu
+      LIBTORCH_VARIANT: shared-with-deps
+    steps:
+      - name: Checkout PyTorch
+        uses: actions/checkout@v4
+        with:
+          ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
+          sparse-checkout: .ci/libtorch/
+          show-progress: false
+      - uses: actions/download-artifact@v4.1.7
+        name: Download Wheel Artifact
+        with:
+          name: manywheel-py3_10-cpu
+          path: "${{ runner.temp }}/wheel_artifact/"
+      - name: Extract libtorch from wheel
+        run: |
+          set -eux
+          mkdir -p "${{ runner.temp }}/libtorch_output"
+          python3 .ci/libtorch/extract_libtorch_from_wheel.py \
+            --wheel-dir "${{ runner.temp }}/wheel_artifact" \
+            --output-dir "${{ runner.temp }}/libtorch_output" \
+            --platform linux \
+            --desired-cuda "$DESIRED_CUDA" \
+            --libtorch-variant "$LIBTORCH_VARIANT"
+      - uses: actions/upload-artifact@v4.4.0
+        if: always()
+        with:
+          name: libtorch-cpu-shared-with-deps-release
+          retention-days: 14
+          if-no-files-found: error
+          path: "${{ runner.temp }}/libtorch_output/"
+  libtorch-cpu-shared-with-deps-release-upload:  # Uploading
+    if: ${{ github.repository_owner == 'pytorch' }}
+    permissions:
+      id-token: write
+      contents: read
+    needs: libtorch-cpu-shared-with-deps-release-extract
+    with:
+      PYTORCH_ROOT: /pytorch
+      PACKAGE_TYPE: libtorch
+      DESIRED_CUDA: cpu
+      GPU_ARCH_TYPE: cpu
+      LIBTORCH_CONFIG: release
+      LIBTORCH_VARIANT: shared-with-deps
+      build_name: libtorch-cpu-shared-with-deps-release
+    secrets:
+      github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
+    uses: ./.github/workflows/_binary-upload.yml
+
+  libtorch-cuda12_6-shared-with-deps-release-extract:
+    if: ${{ github.repository_owner == 'pytorch' }}
+    needs:
+      - manywheel-py3_10-cuda12_6-build
+      - get-label-type
+    runs-on: "${{ needs.get-label-type.outputs.label-type }}linux.4xlarge"
+    timeout-minutes: 60
+    env:
+      DESIRED_CUDA: cu126
+      LIBTORCH_VARIANT: shared-with-deps
+    steps:
+      - name: Checkout PyTorch
+        uses: actions/checkout@v4
+        with:
+          ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
+          sparse-checkout: .ci/libtorch/
+          show-progress: false
+      - uses: actions/download-artifact@v4.1.7
+        name: Download Wheel Artifact
+        with:
+          name: manywheel-py3_10-cuda12_6
+          path: "${{ runner.temp }}/wheel_artifact/"
+      - name: Extract libtorch from wheel
+        run: |
+          set -eux
+          mkdir -p "${{ runner.temp }}/libtorch_output"
+          python3 .ci/libtorch/extract_libtorch_from_wheel.py \
+            --wheel-dir "${{ runner.temp }}/wheel_artifact" \
+            --output-dir "${{ runner.temp }}/libtorch_output" \
+            --platform linux \
+            --desired-cuda "$DESIRED_CUDA" \
+            --libtorch-variant "$LIBTORCH_VARIANT"
+      - uses: actions/upload-artifact@v4.4.0
+        if: always()
+        with:
+          name: libtorch-cuda12_6-shared-with-deps-release
+          retention-days: 14
+          if-no-files-found: error
+          path: "${{ runner.temp }}/libtorch_output/"
+  libtorch-cuda12_6-shared-with-deps-release-upload:  # Uploading
+    if: ${{ github.repository_owner == 'pytorch' }}
+    permissions:
+      id-token: write
+      contents: read
+    needs: libtorch-cuda12_6-shared-with-deps-release-extract
+    with:
+      PYTORCH_ROOT: /pytorch
+      PACKAGE_TYPE: libtorch
+      DESIRED_CUDA: cu126
+      GPU_ARCH_VERSION: "12.6"
+      GPU_ARCH_TYPE: cuda
+      LIBTORCH_CONFIG: release
+      LIBTORCH_VARIANT: shared-with-deps
+      build_name: libtorch-cuda12_6-shared-with-deps-release
+    secrets:
+      github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
+    uses: ./.github/workflows/_binary-upload.yml
+
+  libtorch-cuda13_0-shared-with-deps-release-extract:
+    if: ${{ github.repository_owner == 'pytorch' }}
+    needs:
+      - manywheel-py3_10-cuda13_0-build
+      - get-label-type
+    runs-on: "${{ needs.get-label-type.outputs.label-type }}linux.4xlarge"
+    timeout-minutes: 60
+    env:
+      DESIRED_CUDA: cu130
+      LIBTORCH_VARIANT: shared-with-deps
+    steps:
+      - name: Checkout PyTorch
+        uses: actions/checkout@v4
+        with:
+          ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
+          sparse-checkout: .ci/libtorch/
+          show-progress: false
+      - uses: actions/download-artifact@v4.1.7
+        name: Download Wheel Artifact
+        with:
+          name: manywheel-py3_10-cuda13_0
+          path: "${{ runner.temp }}/wheel_artifact/"
+      - name: Extract libtorch from wheel
+        run: |
+          set -eux
+          mkdir -p "${{ runner.temp }}/libtorch_output"
+          python3 .ci/libtorch/extract_libtorch_from_wheel.py \
+            --wheel-dir "${{ runner.temp }}/wheel_artifact" \
+            --output-dir "${{ runner.temp }}/libtorch_output" \
+            --platform linux \
+            --desired-cuda "$DESIRED_CUDA" \
+            --libtorch-variant "$LIBTORCH_VARIANT"
+      - uses: actions/upload-artifact@v4.4.0
+        if: always()
+        with:
+          name: libtorch-cuda13_0-shared-with-deps-release
+          retention-days: 14
+          if-no-files-found: error
+          path: "${{ runner.temp }}/libtorch_output/"
+  libtorch-cuda13_0-shared-with-deps-release-upload:  # Uploading
+    if: ${{ github.repository_owner == 'pytorch' }}
+    permissions:
+      id-token: write
+      contents: read
+    needs: libtorch-cuda13_0-shared-with-deps-release-extract
+    with:
+      PYTORCH_ROOT: /pytorch
+      PACKAGE_TYPE: libtorch
+      DESIRED_CUDA: cu130
+      GPU_ARCH_VERSION: "13.0"
+      GPU_ARCH_TYPE: cuda
+      LIBTORCH_CONFIG: release
+      LIBTORCH_VARIANT: shared-with-deps
+      build_name: libtorch-cuda13_0-shared-with-deps-release
+    secrets:
+      github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
+    uses: ./.github/workflows/_binary-upload.yml
+
+  libtorch-cuda13_2-shared-with-deps-release-extract:
+    if: ${{ github.repository_owner == 'pytorch' }}
+    needs:
+      - manywheel-py3_10-cuda13_2-build
+      - get-label-type
+    runs-on: "${{ needs.get-label-type.outputs.label-type }}linux.4xlarge"
+    timeout-minutes: 60
+    env:
+      DESIRED_CUDA: cu132
+      LIBTORCH_VARIANT: shared-with-deps
+    steps:
+      - name: Checkout PyTorch
+        uses: actions/checkout@v4
+        with:
+          ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
+          sparse-checkout: .ci/libtorch/
+          show-progress: false
+      - uses: actions/download-artifact@v4.1.7
+        name: Download Wheel Artifact
+        with:
+          name: manywheel-py3_10-cuda13_2
+          path: "${{ runner.temp }}/wheel_artifact/"
+      - name: Extract libtorch from wheel
+        run: |
+          set -eux
+          mkdir -p "${{ runner.temp }}/libtorch_output"
+          python3 .ci/libtorch/extract_libtorch_from_wheel.py \
+            --wheel-dir "${{ runner.temp }}/wheel_artifact" \
+            --output-dir "${{ runner.temp }}/libtorch_output" \
+            --platform linux \
+            --desired-cuda "$DESIRED_CUDA" \
+            --libtorch-variant "$LIBTORCH_VARIANT"
+      - uses: actions/upload-artifact@v4.4.0
+        if: always()
+        with:
+          name: libtorch-cuda13_2-shared-with-deps-release
+          retention-days: 14
+          if-no-files-found: error
+          path: "${{ runner.temp }}/libtorch_output/"
+  libtorch-cuda13_2-shared-with-deps-release-upload:  # Uploading
+    if: ${{ github.repository_owner == 'pytorch' }}
+    permissions:
+      id-token: write
+      contents: read
+    needs: libtorch-cuda13_2-shared-with-deps-release-extract
+    with:
+      PYTORCH_ROOT: /pytorch
+      PACKAGE_TYPE: libtorch
+      DESIRED_CUDA: cu132
+      GPU_ARCH_VERSION: "13.2"
+      GPU_ARCH_TYPE: cuda
+      LIBTORCH_CONFIG: release
+      LIBTORCH_VARIANT: shared-with-deps
+      build_name: libtorch-cuda13_2-shared-with-deps-release
+    secrets:
+      github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
+    uses: ./.github/workflows/_binary-upload.yml
+
+  libtorch-rocm7_1-shared-with-deps-release-extract:
+    if: ${{ github.repository_owner == 'pytorch' }}
+    needs:
+      - manywheel-py3_10-rocm7_1-build
+      - get-label-type
+    runs-on: "${{ needs.get-label-type.outputs.label-type }}linux.4xlarge"
+    timeout-minutes: 60
+    env:
+      DESIRED_CUDA: rocm7.1
+      LIBTORCH_VARIANT: shared-with-deps
+    steps:
+      - name: Checkout PyTorch
+        uses: actions/checkout@v4
+        with:
+          ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
+          sparse-checkout: .ci/libtorch/
+          show-progress: false
+      - uses: actions/download-artifact@v4.1.7
+        name: Download Wheel Artifact
+        with:
+          name: manywheel-py3_10-rocm7_1
+          path: "${{ runner.temp }}/wheel_artifact/"
+      - name: Extract libtorch from wheel
+        run: |
+          set -eux
+          mkdir -p "${{ runner.temp }}/libtorch_output"
+          python3 .ci/libtorch/extract_libtorch_from_wheel.py \
+            --wheel-dir "${{ runner.temp }}/wheel_artifact" \
+            --output-dir "${{ runner.temp }}/libtorch_output" \
+            --platform linux \
+            --desired-cuda "$DESIRED_CUDA" \
+            --libtorch-variant "$LIBTORCH_VARIANT"
+      - uses: actions/upload-artifact@v4.4.0
+        if: always()
+        with:
+          name: libtorch-rocm7_1-shared-with-deps-release
+          retention-days: 14
+          if-no-files-found: error
+          path: "${{ runner.temp }}/libtorch_output/"
+  libtorch-rocm7_1-shared-with-deps-release-upload:  # Uploading
+    if: ${{ github.repository_owner == 'pytorch' }}
+    permissions:
+      id-token: write
+      contents: read
+    needs: libtorch-rocm7_1-shared-with-deps-release-extract
+    with:
+      PYTORCH_ROOT: /pytorch
+      PACKAGE_TYPE: libtorch
+      DESIRED_CUDA: rocm7.1
+      GPU_ARCH_VERSION: "7.1"
+      GPU_ARCH_TYPE: rocm
+      LIBTORCH_CONFIG: release
+      LIBTORCH_VARIANT: shared-with-deps
+      build_name: libtorch-rocm7_1-shared-with-deps-release
+    secrets:
+      github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
+    uses: ./.github/workflows/_binary-upload.yml
+
+  libtorch-rocm7_2-shared-with-deps-release-extract:
+    if: ${{ github.repository_owner == 'pytorch' }}
+    needs:
+      - manywheel-py3_10-rocm7_2-build
+      - get-label-type
+    runs-on: "${{ needs.get-label-type.outputs.label-type }}linux.4xlarge"
+    timeout-minutes: 60
+    env:
+      DESIRED_CUDA: rocm7.2
+      LIBTORCH_VARIANT: shared-with-deps
+    steps:
+      - name: Checkout PyTorch
+        uses: actions/checkout@v4
+        with:
+          ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
+          sparse-checkout: .ci/libtorch/
+          show-progress: false
+      - uses: actions/download-artifact@v4.1.7
+        name: Download Wheel Artifact
+        with:
+          name: manywheel-py3_10-rocm7_2
+          path: "${{ runner.temp }}/wheel_artifact/"
+      - name: Extract libtorch from wheel
+        run: |
+          set -eux
+          mkdir -p "${{ runner.temp }}/libtorch_output"
+          python3 .ci/libtorch/extract_libtorch_from_wheel.py \
+            --wheel-dir "${{ runner.temp }}/wheel_artifact" \
+            --output-dir "${{ runner.temp }}/libtorch_output" \
+            --platform linux \
+            --desired-cuda "$DESIRED_CUDA" \
+            --libtorch-variant "$LIBTORCH_VARIANT"
+      - uses: actions/upload-artifact@v4.4.0
+        if: always()
+        with:
+          name: libtorch-rocm7_2-shared-with-deps-release
+          retention-days: 14
+          if-no-files-found: error
+          path: "${{ runner.temp }}/libtorch_output/"
+  libtorch-rocm7_2-shared-with-deps-release-upload:  # Uploading
+    if: ${{ github.repository_owner == 'pytorch' }}
+    permissions:
+      id-token: write
+      contents: read
+    needs: libtorch-rocm7_2-shared-with-deps-release-extract
+    with:
+      PYTORCH_ROOT: /pytorch
+      PACKAGE_TYPE: libtorch
+      DESIRED_CUDA: rocm7.2
+      GPU_ARCH_VERSION: "7.2"
+      GPU_ARCH_TYPE: rocm
+      LIBTORCH_CONFIG: release
+      LIBTORCH_VARIANT: shared-with-deps
+      build_name: libtorch-rocm7_2-shared-with-deps-release
+    secrets:
+      github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
     uses: ./.github/workflows/_binary-upload.yml
diff --git a/.github/workflows/generated-linux-s390x-binary-manywheel-nightly.yml b/.github/workflows/generated-linux-s390x-binary-manywheel-nightly.yml
index 2314dfa31db47..09405e9d2cc8b 100644
--- a/.github/workflows/generated-linux-s390x-binary-manywheel-nightly.yml
+++ b/.github/workflows/generated-linux-s390x-binary-manywheel-nightly.yml
@@ -42,7 +42,7 @@ jobs:
   get-label-type:
     if: github.repository_owner == 'pytorch'
     name: get-label-type
-    uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@main
+    uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@release/2.12
     with:
       triggering_actor: ${{ github.triggering_actor }}
       issue_owner: ${{ github.event.pull_request.user.login || github.event.issue.user.login }}
@@ -111,6 +111,9 @@ jobs:
       build_name: manywheel-py3_10-cpu-s390x
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
     uses: ./.github/workflows/_binary-upload.yml
 
   manywheel-py3_11-cpu-s390x-build:
@@ -176,6 +179,9 @@ jobs:
       build_name: manywheel-py3_11-cpu-s390x
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
     uses: ./.github/workflows/_binary-upload.yml
 
   manywheel-py3_12-cpu-s390x-build:
@@ -241,6 +247,9 @@ jobs:
       build_name: manywheel-py3_12-cpu-s390x
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
     uses: ./.github/workflows/_binary-upload.yml
 
   manywheel-py3_13-cpu-s390x-build:
@@ -306,6 +315,9 @@ jobs:
       build_name: manywheel-py3_13-cpu-s390x
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
     uses: ./.github/workflows/_binary-upload.yml
 
   manywheel-py3_13t-cpu-s390x-build:
@@ -371,6 +383,9 @@ jobs:
       build_name: manywheel-py3_13t-cpu-s390x
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
     uses: ./.github/workflows/_binary-upload.yml
 
   manywheel-py3_14-cpu-s390x-build:
@@ -436,6 +451,9 @@ jobs:
       build_name: manywheel-py3_14-cpu-s390x
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
     uses: ./.github/workflows/_binary-upload.yml
 
   manywheel-py3_14t-cpu-s390x-build:
@@ -501,4 +519,7 @@ jobs:
       build_name: manywheel-py3_14t-cpu-s390x
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
     uses: ./.github/workflows/_binary-upload.yml
diff --git a/.github/workflows/generated-macos-arm64-binary-libtorch-release-nightly.yml b/.github/workflows/generated-macos-arm64-binary-libtorch-release-nightly.yml
deleted file mode 100644
index 2d47b83e5822e..0000000000000
--- a/.github/workflows/generated-macos-arm64-binary-libtorch-release-nightly.yml
+++ /dev/null
@@ -1,128 +0,0 @@
-# @generated DO NOT EDIT MANUALLY
-
-# Template is at:    .github/templates/macos_binary_build_workflow.yml.j2
-# Generation script: .github/scripts/generate_ci_workflows.py
-name: macos-arm64-binary-libtorch-release
-
-on:
-# TODO: Migrate to new ciflow trigger, reference https://github.com/pytorch/pytorch/pull/70321
-  push:
-    # NOTE: Meta Employees can trigger new nightlies using: https://fburl.com/trigger_pytorch_nightly_build
-    branches:
-      - nightly
-    tags:
-      # NOTE: Binary build pipelines should only get triggered on release candidate builds
-      # Release candidate tags look like: v1.11.0-rc1
-      - v[0-9]+.[0-9]+.[0-9]+-rc[0-9]+
-      - 'ciflow/binaries/*'
-      - 'ciflow/binaries_libtorch/*'
-  workflow_dispatch:
-
-env:
-  ALPINE_IMAGE: "308535385114.dkr.ecr.us-east-1.amazonaws.com/tool/alpine"
-  AWS_DEFAULT_REGION: us-east-1
-  BUILD_ENVIRONMENT: macos-arm64-binary-libtorch-release
-  GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
-  PR_NUMBER: ${{ github.event.pull_request.number }}
-  SKIP_ALL_TESTS: 0
-concurrency:
-  group: macos-arm64-binary-libtorch-release-${{ github.event.pull_request.number || github.ref_name }}-${{ github.ref_type == 'branch' && github.sha }}-${{ github.event_name == 'workflow_dispatch' }}
-  cancel-in-progress: true
-
-jobs:
-  libtorch-cpu-shared-with-deps-release-build:
-    if: ${{ github.repository_owner == 'pytorch' }}
-    runs-on: macos-14-xlarge
-    timeout-minutes: 240
-    env:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
-      PACKAGE_TYPE: libtorch
-      # TODO: This is a legacy variable that we eventually want to get rid of in
-      #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cpu
-      GPU_ARCH_TYPE: cpu
-      SKIP_ALL_TESTS: 1
-      LIBTORCH_CONFIG: release
-      LIBTORCH_VARIANT: shared-with-deps
-      # This is a dummy value for libtorch to work correctly with our batch scripts
-      # without this value pip does not get installed for some reason
-      DESIRED_PYTHON: "3.10"
-    steps:
-      # NOTE: These environment variables are put here so that they can be applied on every job equally
-      #       They are also here because setting them at a workflow level doesn't give us access to the
-      #       runner.temp variable, which we need.
-      - name: Populate binary env
-        shell: bash
-        run: |
-          # shellcheck disable=SC2129
-          echo "BINARY_ENV_FILE=${RUNNER_TEMP}/env" >> "${GITHUB_ENV}"
-          # shellcheck disable=SC2129
-          echo "PYTORCH_FINAL_PACKAGE_DIR=${RUNNER_TEMP}/artifacts" >> "${GITHUB_ENV}"
-          # shellcheck disable=SC2129
-          echo "MAC_PACKAGE_WORK_DIR=${RUNNER_TEMP}" >> "${GITHUB_ENV}"
-      - name: Setup Python
-        uses: actions/setup-python@v6
-        with:
-          # .4 version is min minor for 3.10, and also no-gil version of 3.13 needs at least 3.13.3
-          python-version: "3.10.4"
-          freethreaded: false
-      - name: Checkout PyTorch
-        uses: actions/checkout@v4
-        with:
-          ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
-          submodules: recursive
-          path: pytorch
-          show-progress: false
-      - name: Clean PyTorch checkout
-        run: |
-          # Remove any artifacts from the previous checkouts
-          git clean -fxd
-        working-directory: pytorch
-      - name: Populate binary env
-        run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_populate_env.sh"
-      - name: Build PyTorch binary
-        run: |
-          set -eux -o pipefail
-          # shellcheck disable=SC1090
-          source "${BINARY_ENV_FILE:-/Users/distiller/project/env}"
-          mkdir -p "$PYTORCH_FINAL_PACKAGE_DIR"
-
-          # Build
-          USE_PYTORCH_METAL_EXPORT=1
-          USE_COREML_DELEGATE=1
-          TORCH_PACKAGE_NAME="${TORCH_PACKAGE_NAME//-/_}"
-          export USE_PYTORCH_METAL_EXPORT
-          export USE_COREML_DELEGATE
-          export TORCH_PACKAGE_NAME
-          "${PYTORCH_ROOT}/.ci/wheel/install_libomp.sh"
-          "${PYTORCH_ROOT}/.ci/wheel/build_wheel.sh"
-      - uses: actions/upload-artifact@v4.4.0
-        if: always()
-        with:
-          name: libtorch-cpu-shared-with-deps-release
-          retention-days: 14
-          if-no-files-found: error
-          path: "${{ env.PYTORCH_FINAL_PACKAGE_DIR }}"
-  libtorch-cpu-shared-with-deps-release-upload:  # Uploading
-    if: ${{ github.repository_owner == 'pytorch' }}
-    permissions:
-      id-token: write
-      contents: read
-    needs: libtorch-cpu-shared-with-deps-release-build
-    with:
-      PYTORCH_ROOT: /pytorch
-      PACKAGE_TYPE: libtorch
-      # TODO: This is a legacy variable that we eventually want to get rid of in
-      #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cpu
-      GPU_ARCH_TYPE: cpu
-      DOCKER_IMAGE: libtorch-cxx11-builder
-      DOCKER_IMAGE_TAG_PREFIX: cpu
-      LIBTORCH_CONFIG: release
-      LIBTORCH_VARIANT: shared-with-deps
-      build_name: libtorch-cpu-shared-with-deps-release
-      use_s3: False
-    secrets:
-      github-token: ${{ secrets.GITHUB_TOKEN }}
-    uses: ./.github/workflows/_binary-upload.yml
diff --git a/.github/workflows/generated-macos-arm64-binary-wheel-nightly.yml b/.github/workflows/generated-macos-arm64-binary-wheel-nightly.yml
index 70d210c7eabb6..8f1011024185e 100644
--- a/.github/workflows/generated-macos-arm64-binary-wheel-nightly.yml
+++ b/.github/workflows/generated-macos-arm64-binary-wheel-nightly.yml
@@ -15,6 +15,7 @@ on:
       # Release candidate tags look like: v1.11.0-rc1
       - v[0-9]+.[0-9]+.[0-9]+-rc[0-9]+
       - 'ciflow/binaries/*'
+      - 'ciflow/binaries_libtorch/*'
       - 'ciflow/binaries_wheel/*'
   workflow_dispatch:
 
@@ -32,10 +33,10 @@ concurrency:
 jobs:
   wheel-py3_10-cpu-build:
     if: ${{ github.repository_owner == 'pytorch' }}
-    runs-on: macos-14-xlarge
+    runs-on: macos-26-xlarge
     timeout-minutes: 240
     env:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
+      PYTORCH_ROOT: ${{ github.workspace }}
       PACKAGE_TYPE: wheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
@@ -43,6 +44,7 @@ jobs:
       GPU_ARCH_TYPE: cpu
       SKIP_ALL_TESTS: 1
       DESIRED_PYTHON: "3.10"
+      PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cufile,nvjitlink,nvtx]==13.0.2; platform_system == 'Linux' | nvidia-cublas>=13.1.0.3,<=13.1.1.3; platform_system == 'Linux' | cuda-bindings>=13.0.3,<14; platform_system == 'Linux' | nvidia-cudnn-cu13==9.20.0.48; platform_system == 'Linux' | nvidia-cusparselt-cu13==0.8.1; platform_system == 'Linux' | nvidia-nccl-cu13==2.29.7; platform_system == 'Linux' | nvidia-nvshmem-cu13==3.4.5; platform_system == 'Linux'
     steps:
       # NOTE: These environment variables are put here so that they can be applied on every job equally
       #       They are also here because setting them at a workflow level doesn't give us access to the
@@ -65,18 +67,15 @@ jobs:
       - name: Checkout PyTorch
         uses: actions/checkout@v4
         with:
-          ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
           submodules: recursive
-          path: pytorch
           show-progress: false
       - name: Clean PyTorch checkout
         run: |
           # Remove any artifacts from the previous checkouts
           git clean -fxd
-        working-directory: pytorch
       - name: Populate binary env
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_populate_env.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_populate_env.sh"
       - name: Build PyTorch binary
         run: |
           set -eux -o pipefail
@@ -139,13 +138,16 @@ jobs:
       use_s3: False
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
     uses: ./.github/workflows/_binary-upload.yml
   wheel-py3_11-cpu-build:
     if: ${{ github.repository_owner == 'pytorch' }}
-    runs-on: macos-14-xlarge
+    runs-on: macos-26-xlarge
     timeout-minutes: 240
     env:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
+      PYTORCH_ROOT: ${{ github.workspace }}
       PACKAGE_TYPE: wheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
@@ -153,6 +155,7 @@ jobs:
       GPU_ARCH_TYPE: cpu
       SKIP_ALL_TESTS: 1
       DESIRED_PYTHON: "3.11"
+      PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cufile,nvjitlink,nvtx]==13.0.2; platform_system == 'Linux' | nvidia-cublas>=13.1.0.3,<=13.1.1.3; platform_system == 'Linux' | cuda-bindings>=13.0.3,<14; platform_system == 'Linux' | nvidia-cudnn-cu13==9.20.0.48; platform_system == 'Linux' | nvidia-cusparselt-cu13==0.8.1; platform_system == 'Linux' | nvidia-nccl-cu13==2.29.7; platform_system == 'Linux' | nvidia-nvshmem-cu13==3.4.5; platform_system == 'Linux'
     steps:
       # NOTE: These environment variables are put here so that they can be applied on every job equally
       #       They are also here because setting them at a workflow level doesn't give us access to the
@@ -175,18 +178,15 @@ jobs:
       - name: Checkout PyTorch
         uses: actions/checkout@v4
         with:
-          ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
           submodules: recursive
-          path: pytorch
           show-progress: false
       - name: Clean PyTorch checkout
         run: |
           # Remove any artifacts from the previous checkouts
           git clean -fxd
-        working-directory: pytorch
       - name: Populate binary env
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_populate_env.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_populate_env.sh"
       - name: Build PyTorch binary
         run: |
           set -eux -o pipefail
@@ -249,13 +249,16 @@ jobs:
       use_s3: False
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
     uses: ./.github/workflows/_binary-upload.yml
   wheel-py3_12-cpu-build:
     if: ${{ github.repository_owner == 'pytorch' }}
-    runs-on: macos-14-xlarge
+    runs-on: macos-26-xlarge
     timeout-minutes: 240
     env:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
+      PYTORCH_ROOT: ${{ github.workspace }}
       PACKAGE_TYPE: wheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
@@ -263,6 +266,7 @@ jobs:
       GPU_ARCH_TYPE: cpu
       SKIP_ALL_TESTS: 1
       DESIRED_PYTHON: "3.12"
+      PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cufile,nvjitlink,nvtx]==13.0.2; platform_system == 'Linux' | nvidia-cublas>=13.1.0.3,<=13.1.1.3; platform_system == 'Linux' | cuda-bindings>=13.0.3,<14; platform_system == 'Linux' | nvidia-cudnn-cu13==9.20.0.48; platform_system == 'Linux' | nvidia-cusparselt-cu13==0.8.1; platform_system == 'Linux' | nvidia-nccl-cu13==2.29.7; platform_system == 'Linux' | nvidia-nvshmem-cu13==3.4.5; platform_system == 'Linux'
     steps:
       # NOTE: These environment variables are put here so that they can be applied on every job equally
       #       They are also here because setting them at a workflow level doesn't give us access to the
@@ -285,18 +289,15 @@ jobs:
       - name: Checkout PyTorch
         uses: actions/checkout@v4
         with:
-          ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
           submodules: recursive
-          path: pytorch
           show-progress: false
       - name: Clean PyTorch checkout
         run: |
           # Remove any artifacts from the previous checkouts
           git clean -fxd
-        working-directory: pytorch
       - name: Populate binary env
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_populate_env.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_populate_env.sh"
       - name: Build PyTorch binary
         run: |
           set -eux -o pipefail
@@ -359,13 +360,16 @@ jobs:
       use_s3: False
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
     uses: ./.github/workflows/_binary-upload.yml
   wheel-py3_13-cpu-build:
     if: ${{ github.repository_owner == 'pytorch' }}
-    runs-on: macos-14-xlarge
+    runs-on: macos-26-xlarge
     timeout-minutes: 240
     env:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
+      PYTORCH_ROOT: ${{ github.workspace }}
       PACKAGE_TYPE: wheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
@@ -373,6 +377,7 @@ jobs:
       GPU_ARCH_TYPE: cpu
       SKIP_ALL_TESTS: 1
       DESIRED_PYTHON: "3.13"
+      PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cufile,nvjitlink,nvtx]==13.0.2; platform_system == 'Linux' | nvidia-cublas>=13.1.0.3,<=13.1.1.3; platform_system == 'Linux' | cuda-bindings>=13.0.3,<14; platform_system == 'Linux' | nvidia-cudnn-cu13==9.20.0.48; platform_system == 'Linux' | nvidia-cusparselt-cu13==0.8.1; platform_system == 'Linux' | nvidia-nccl-cu13==2.29.7; platform_system == 'Linux' | nvidia-nvshmem-cu13==3.4.5; platform_system == 'Linux'
     steps:
       # NOTE: These environment variables are put here so that they can be applied on every job equally
       #       They are also here because setting them at a workflow level doesn't give us access to the
@@ -395,18 +400,15 @@ jobs:
       - name: Checkout PyTorch
         uses: actions/checkout@v4
         with:
-          ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
           submodules: recursive
-          path: pytorch
           show-progress: false
       - name: Clean PyTorch checkout
         run: |
           # Remove any artifacts from the previous checkouts
           git clean -fxd
-        working-directory: pytorch
       - name: Populate binary env
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_populate_env.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_populate_env.sh"
       - name: Build PyTorch binary
         run: |
           set -eux -o pipefail
@@ -469,13 +471,16 @@ jobs:
       use_s3: False
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
     uses: ./.github/workflows/_binary-upload.yml
   wheel-py3_13t-cpu-build:
     if: ${{ github.repository_owner == 'pytorch' }}
-    runs-on: macos-14-xlarge
+    runs-on: macos-26-xlarge
     timeout-minutes: 240
     env:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
+      PYTORCH_ROOT: ${{ github.workspace }}
       PACKAGE_TYPE: wheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
@@ -483,6 +488,7 @@ jobs:
       GPU_ARCH_TYPE: cpu
       SKIP_ALL_TESTS: 1
       DESIRED_PYTHON: "3.13t"
+      PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cufile,nvjitlink,nvtx]==13.0.2; platform_system == 'Linux' | nvidia-cublas>=13.1.0.3,<=13.1.1.3; platform_system == 'Linux' | cuda-bindings>=13.0.3,<14; platform_system == 'Linux' | nvidia-cudnn-cu13==9.20.0.48; platform_system == 'Linux' | nvidia-cusparselt-cu13==0.8.1; platform_system == 'Linux' | nvidia-nccl-cu13==2.29.7; platform_system == 'Linux' | nvidia-nvshmem-cu13==3.4.5; platform_system == 'Linux'
     steps:
       # NOTE: These environment variables are put here so that they can be applied on every job equally
       #       They are also here because setting them at a workflow level doesn't give us access to the
@@ -505,18 +511,15 @@ jobs:
       - name: Checkout PyTorch
         uses: actions/checkout@v4
         with:
-          ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
           submodules: recursive
-          path: pytorch
           show-progress: false
       - name: Clean PyTorch checkout
         run: |
           # Remove any artifacts from the previous checkouts
           git clean -fxd
-        working-directory: pytorch
       - name: Populate binary env
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_populate_env.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_populate_env.sh"
       - name: Build PyTorch binary
         run: |
           set -eux -o pipefail
@@ -579,13 +582,16 @@ jobs:
       use_s3: False
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
     uses: ./.github/workflows/_binary-upload.yml
   wheel-py3_14-cpu-build:
     if: ${{ github.repository_owner == 'pytorch' }}
-    runs-on: macos-14-xlarge
+    runs-on: macos-26-xlarge
     timeout-minutes: 240
     env:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
+      PYTORCH_ROOT: ${{ github.workspace }}
       PACKAGE_TYPE: wheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
@@ -593,6 +599,7 @@ jobs:
       GPU_ARCH_TYPE: cpu
       SKIP_ALL_TESTS: 1
       DESIRED_PYTHON: "3.14"
+      PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cufile,nvjitlink,nvtx]==13.0.2; platform_system == 'Linux' | nvidia-cublas>=13.1.0.3,<=13.1.1.3; platform_system == 'Linux' | cuda-bindings>=13.0.3,<14; platform_system == 'Linux' | nvidia-cudnn-cu13==9.20.0.48; platform_system == 'Linux' | nvidia-cusparselt-cu13==0.8.1; platform_system == 'Linux' | nvidia-nccl-cu13==2.29.7; platform_system == 'Linux' | nvidia-nvshmem-cu13==3.4.5; platform_system == 'Linux'
     steps:
       # NOTE: These environment variables are put here so that they can be applied on every job equally
       #       They are also here because setting them at a workflow level doesn't give us access to the
@@ -615,18 +622,15 @@ jobs:
       - name: Checkout PyTorch
         uses: actions/checkout@v4
         with:
-          ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
           submodules: recursive
-          path: pytorch
           show-progress: false
       - name: Clean PyTorch checkout
         run: |
           # Remove any artifacts from the previous checkouts
           git clean -fxd
-        working-directory: pytorch
       - name: Populate binary env
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_populate_env.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_populate_env.sh"
       - name: Build PyTorch binary
         run: |
           set -eux -o pipefail
@@ -689,13 +693,16 @@ jobs:
       use_s3: False
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
     uses: ./.github/workflows/_binary-upload.yml
   wheel-py3_14t-cpu-build:
     if: ${{ github.repository_owner == 'pytorch' }}
-    runs-on: macos-14-xlarge
+    runs-on: macos-26-xlarge
     timeout-minutes: 240
     env:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
+      PYTORCH_ROOT: ${{ github.workspace }}
       PACKAGE_TYPE: wheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
@@ -703,6 +710,7 @@ jobs:
       GPU_ARCH_TYPE: cpu
       SKIP_ALL_TESTS: 1
       DESIRED_PYTHON: "3.14t"
+      PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cufile,nvjitlink,nvtx]==13.0.2; platform_system == 'Linux' | nvidia-cublas>=13.1.0.3,<=13.1.1.3; platform_system == 'Linux' | cuda-bindings>=13.0.3,<14; platform_system == 'Linux' | nvidia-cudnn-cu13==9.20.0.48; platform_system == 'Linux' | nvidia-cusparselt-cu13==0.8.1; platform_system == 'Linux' | nvidia-nccl-cu13==2.29.7; platform_system == 'Linux' | nvidia-nvshmem-cu13==3.4.5; platform_system == 'Linux'
     steps:
       # NOTE: These environment variables are put here so that they can be applied on every job equally
       #       They are also here because setting them at a workflow level doesn't give us access to the
@@ -725,18 +733,15 @@ jobs:
       - name: Checkout PyTorch
         uses: actions/checkout@v4
         with:
-          ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
           submodules: recursive
-          path: pytorch
           show-progress: false
       - name: Clean PyTorch checkout
         run: |
           # Remove any artifacts from the previous checkouts
           git clean -fxd
-        working-directory: pytorch
       - name: Populate binary env
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_populate_env.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_populate_env.sh"
       - name: Build PyTorch binary
         run: |
           set -eux -o pipefail
@@ -799,4 +804,71 @@ jobs:
       use_s3: False
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
+    uses: ./.github/workflows/_binary-upload.yml
+  libtorch-cpu-shared-with-deps-release-extract:
+    if: ${{ github.repository_owner == 'pytorch' }}
+    needs: wheel-py3_10-cpu-build
+    runs-on: macos-26-xlarge
+    timeout-minutes: 60
+    env:
+      DESIRED_CUDA: cpu
+      LIBTORCH_VARIANT: shared-with-deps
+    steps:
+      - name: Checkout PyTorch
+        uses: actions/checkout@v4
+        with:
+          ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
+          sparse-checkout: .ci/libtorch/
+          show-progress: false
+      - name: Setup Python
+        uses: actions/setup-python@v6
+        with:
+          # .4 version is min minor for 3.10, and also no-gil version of 3.13 needs at least 3.13.3
+          python-version: "3.10.4"
+          freethreaded: false
+      - uses: actions/download-artifact@v4.1.7
+        name: Download Wheel Artifact
+        with:
+          name: wheel-py3_10-cpu
+          path: "${{ runner.temp }}/wheel_artifact/"
+      - name: Extract libtorch from wheel
+        run: |
+          set -eux
+          mkdir -p "${{ runner.temp }}/libtorch_output"
+          python3 .ci/libtorch/extract_libtorch_from_wheel.py \
+            --wheel-dir "${{ runner.temp }}/wheel_artifact" \
+            --output-dir "${{ runner.temp }}/libtorch_output" \
+            --platform macos \
+            --desired-cuda "$DESIRED_CUDA" \
+            --libtorch-variant "$LIBTORCH_VARIANT"
+      - uses: actions/upload-artifact@v4.4.0
+        if: always()
+        with:
+          name: libtorch-cpu-shared-with-deps-release
+          retention-days: 14
+          if-no-files-found: error
+          path: "${{ runner.temp }}/libtorch_output/"
+  libtorch-cpu-shared-with-deps-release-upload:  # Uploading
+    if: ${{ github.repository_owner == 'pytorch' }}
+    permissions:
+      id-token: write
+      contents: read
+    needs: libtorch-cpu-shared-with-deps-release-extract
+    with:
+      PYTORCH_ROOT: /pytorch
+      PACKAGE_TYPE: libtorch
+      DESIRED_CUDA: cpu
+      GPU_ARCH_TYPE: cpu
+      LIBTORCH_CONFIG: release
+      LIBTORCH_VARIANT: shared-with-deps
+      build_name: libtorch-cpu-shared-with-deps-release
+      use_s3: False
+    secrets:
+      github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
     uses: ./.github/workflows/_binary-upload.yml
diff --git a/.github/workflows/generated-windows-arm64-binary-libtorch-debug-nightly.yml b/.github/workflows/generated-windows-arm64-binary-libtorch-debug-nightly.yml
index 7c26dbc3b9eea..8edab3d40a604 100644
--- a/.github/workflows/generated-windows-arm64-binary-libtorch-debug-nightly.yml
+++ b/.github/workflows/generated-windows-arm64-binary-libtorch-debug-nightly.yml
@@ -41,7 +41,7 @@ jobs:
   get-label-type:
     if: github.repository_owner == 'pytorch'
     name: get-label-type
-    uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@main
+    uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@release/2.12
     with:
       triggering_actor: ${{ github.triggering_actor }}
       issue_owner: ${{ github.event.pull_request.user.login || github.event.issue.user.login }}
@@ -53,7 +53,7 @@ jobs:
     runs-on: "windows-11-arm64-preview"
     timeout-minutes: 360
     env:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
+      PYTORCH_ROOT: ${{ github.workspace }}
       PACKAGE_TYPE: libtorch
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
@@ -85,36 +85,35 @@ jobs:
       - name: Git checkout PyTorch
         uses: actions/checkout@v4
         with:
-          path: "pytorch"
           submodules: recursive
       - name: Bootstrap Python
         shell: cmd
         run: |
-          "pytorch/.ci/pytorch/windows/arm64/bootstrap_python.bat"
+          ".ci/pytorch/windows/arm64/bootstrap_python.bat"
       - name: Bootstrap APL
         shell: cmd
         run: |
-          "pytorch/.ci/pytorch/windows/arm64/bootstrap_apl.bat"
+          ".ci/pytorch/windows/arm64/bootstrap_apl.bat"
       - name: Bootstrap Rust
         shell: cmd
         run: |
-          "pytorch/.ci/pytorch/windows/arm64/bootstrap_rust.bat"
+          ".ci/pytorch/windows/arm64/bootstrap_rust.bat"
       - name: Bootstrap sccache
         shell: cmd
         run: |
-          "pytorch/.ci/pytorch/windows/arm64/bootstrap_sccache.bat"
+          ".ci/pytorch/windows/arm64/bootstrap_sccache.bat"
       - name: Bootstrap Libuv
         shell: cmd
         run: |
-          "pytorch/.ci/pytorch/windows/arm64/bootstrap_libuv.bat"
+          ".ci/pytorch/windows/arm64/bootstrap_libuv.bat"
       - name: Populate binary env
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_populate_env.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_populate_env.sh"
       - name: Build PyTorch binary
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_windows_build.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_windows_build.sh"
       - uses: actions/upload-artifact@v4.4.0
         if: always()
         with:
@@ -130,7 +129,7 @@ jobs:
     runs-on: "windows-11-arm64-preview"
     timeout-minutes: 360
     env:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
+      PYTORCH_ROOT: ${{ github.workspace }}
       PACKAGE_TYPE: libtorch
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
@@ -157,20 +156,19 @@ jobs:
       - name: Git checkout PyTorch
         uses: actions/checkout@v4
         with:
-          path: "pytorch"
           submodules: recursive
       - name: Bootstrap APL
         shell: cmd
         run: |
-          "pytorch/.ci/pytorch/windows/arm64/bootstrap_apl.bat"
+          ".ci/pytorch/windows/arm64/bootstrap_apl.bat"
       - name: Bootstrap Python
         shell: cmd
         run: |
-          "pytorch/.ci/pytorch/windows/arm64/bootstrap_python.bat"
+          ".ci/pytorch/windows/arm64/bootstrap_python.bat"
       - name: Bootstrap Rust
         shell: cmd
         run: |
-          "pytorch/.ci/pytorch/windows/arm64/bootstrap_rust.bat"
+          ".ci/pytorch/windows/arm64/bootstrap_rust.bat"
       - uses: actions/download-artifact@v4.1.7
         name: Download Build Artifacts
         with:
@@ -179,11 +177,11 @@ jobs:
       - name: Populate binary env
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_populate_env.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_populate_env.sh"
       - name: Test PyTorch binary
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_windows_test.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_windows_test.sh"
   libtorch-cpu-shared-with-deps-debug-upload:  # Uploading
     if: ${{ github.repository_owner == 'pytorch' }}
     permissions:
@@ -191,7 +189,7 @@ jobs:
       contents: read
     needs: libtorch-cpu-shared-with-deps-debug-test
     with:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
+      PYTORCH_ROOT: ${{ github.workspace }}
       PACKAGE_TYPE: libtorch
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
@@ -205,4 +203,7 @@ jobs:
       build_name: libtorch-cpu-shared-with-deps-debug
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
     uses: ./.github/workflows/_binary-upload.yml
diff --git a/.github/workflows/generated-windows-arm64-binary-libtorch-release-nightly.yml b/.github/workflows/generated-windows-arm64-binary-libtorch-release-nightly.yml
deleted file mode 100644
index 5e30b66183840..0000000000000
--- a/.github/workflows/generated-windows-arm64-binary-libtorch-release-nightly.yml
+++ /dev/null
@@ -1,208 +0,0 @@
-# @generated DO NOT EDIT MANUALLY
-
-# Template is at:    .github/templates/windows_binary_build_workflow.yml.j2
-# Generation script: .github/scripts/generate_ci_workflows.py
-name: windows-arm64-binary-libtorch-release
-
-on:
-  push:
-    # NOTE: Meta Employees can trigger new nightlies using: https://fburl.com/trigger_pytorch_nightly_build
-    branches:
-      - nightly
-    tags:
-      # NOTE: Binary build pipelines should only get triggered on release candidate builds
-      # Release candidate tags look like: v1.11.0-rc1
-      - v[0-9]+.[0-9]+.[0-9]+-rc[0-9]+
-      - 'ciflow/binaries/*'
-      - 'ciflow/binaries_libtorch/*'
-  workflow_dispatch:
-
-env:
-  # Needed for conda builds
-  ALPINE_IMAGE: "308535385114.dkr.ecr.us-east-1.amazonaws.com/tool/alpine"
-  AWS_DEFAULT_REGION: us-east-1
-  BUILD_ENVIRONMENT: windows-arm64-binary-libtorch-release
-  GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
-  PR_NUMBER: ${{ github.event.pull_request.number }}
-  SHA1: ${{ github.event.pull_request.head.sha || github.sha }}
-  SKIP_ALL_TESTS: 1
-  OS: windows-arm64
-  PYTORCH_ROOT: /pytorch
-  DOWNLOADS_DIR: c:\temp\downloads
-  DEPENDENCIES_DIR: c:\temp\dependencies
-  ENABLE_APL: 1
-  ENABLE_OPENBLAS: 0
-  MSVC_VERSION : 14.42
-concurrency:
-  group: windows-arm64-binary-libtorch-release-${{ github.event.pull_request.number || github.ref_name }}-${{ github.ref_type == 'branch' && github.sha }}-${{ github.event_name == 'workflow_dispatch' }}
-  cancel-in-progress: true
-
-jobs:
-  get-label-type:
-    if: github.repository_owner == 'pytorch'
-    name: get-label-type
-    uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@main
-    with:
-      triggering_actor: ${{ github.triggering_actor }}
-      issue_owner: ${{ github.event.pull_request.user.login || github.event.issue.user.login }}
-      curr_branch: ${{ github.head_ref || github.ref_name }}
-      curr_ref_type: ${{ github.ref_type }}
-  libtorch-cpu-shared-with-deps-release-build:
-    if: ${{ github.repository_owner == 'pytorch' }}
-    needs: get-label-type
-    runs-on: "windows-11-arm64-preview"
-    timeout-minutes: 360
-    env:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
-      PACKAGE_TYPE: libtorch
-      # TODO: This is a legacy variable that we eventually want to get rid of in
-      #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cpu
-      GPU_ARCH_TYPE: cpu
-      SKIP_ALL_TESTS: 1
-      LIBTORCH_CONFIG: release
-      LIBTORCH_VARIANT: shared-with-deps
-      # This is a dummy value for libtorch to work correctly with our batch scripts
-      # without this value pip does not get installed for some reason
-      DESIRED_PYTHON: "3.10"
-    steps:
-      - name: Populate binary env
-        shell: cmd
-        run: |
-          echo BINARY_ENV_FILE=%RUNNER_TEMP%/env>> %GITHUB_ENV%
-          echo PYTORCH_FINAL_PACKAGE_DIR=%RUNNER_TEMP%/artifacts>> %GITHUB_ENV%
-          echo WIN_PACKAGE_WORK_DIR=%RUNNER_TEMP%>> %GITHUB_ENV%
-      - name: Bootstrap folders
-        shell: cmd
-        run: |
-          mkdir "%NIGHTLIES_PYTORCH_ROOT%"
-          mkdir "%PYTORCH_FINAL_PACKAGE_DIR%"
-      - name: Enable long paths
-        shell: cmd
-        run: |
-          git config --system --get core.longpaths || echo "core.longpaths is not set, setting it now"
-          git config --system core.longpaths true
-      - name: Git checkout PyTorch
-        uses: actions/checkout@v4
-        with:
-          path: "pytorch"
-          submodules: recursive
-      - name: Bootstrap Python
-        shell: cmd
-        run: |
-          "pytorch/.ci/pytorch/windows/arm64/bootstrap_python.bat"
-      - name: Bootstrap APL
-        shell: cmd
-        run: |
-          "pytorch/.ci/pytorch/windows/arm64/bootstrap_apl.bat"
-      - name: Bootstrap Rust
-        shell: cmd
-        run: |
-          "pytorch/.ci/pytorch/windows/arm64/bootstrap_rust.bat"
-      - name: Bootstrap sccache
-        shell: cmd
-        run: |
-          "pytorch/.ci/pytorch/windows/arm64/bootstrap_sccache.bat"
-      - name: Bootstrap Libuv
-        shell: cmd
-        run: |
-          "pytorch/.ci/pytorch/windows/arm64/bootstrap_libuv.bat"
-      - name: Populate binary env
-        shell: bash
-        run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_populate_env.sh"
-      - name: Build PyTorch binary
-        shell: bash
-        run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_windows_build.sh"
-      - uses: actions/upload-artifact@v4.4.0
-        if: always()
-        with:
-          name: libtorch-cpu-shared-with-deps-release
-          retention-days: 14
-          if-no-files-found: error
-          path: "${{ env.PYTORCH_FINAL_PACKAGE_DIR }}"
-  libtorch-cpu-shared-with-deps-release-test:  # Testing
-    if: ${{ github.repository_owner == 'pytorch' }}
-    needs:
-      - libtorch-cpu-shared-with-deps-release-build
-      - get-label-type
-    runs-on: "windows-11-arm64-preview"
-    timeout-minutes: 360
-    env:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
-      PACKAGE_TYPE: libtorch
-      # TODO: This is a legacy variable that we eventually want to get rid of in
-      #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cpu
-      GPU_ARCH_TYPE: cpu
-      SKIP_ALL_TESTS: 1
-      LIBTORCH_CONFIG: release
-      LIBTORCH_VARIANT: shared-with-deps
-      # This is a dummy value for libtorch to work correctly with our batch scripts
-      # without this value pip does not get installed for some reason
-      DESIRED_PYTHON: "3.10"
-    steps:
-      - name: Populate binary env
-        shell: cmd
-        run: |
-          echo BINARY_ENV_FILE=%RUNNER_TEMP%/env>> %GITHUB_ENV%
-          echo PYTORCH_FINAL_PACKAGE_DIR=%RUNNER_TEMP%/artifacts>> %GITHUB_ENV%
-          echo WIN_PACKAGE_WORK_DIR=%RUNNER_TEMP%>> %GITHUB_ENV%
-      - name: Enable long paths
-        shell: cmd
-        run: |
-          git config --system --get core.longpaths || echo "core.longpaths is not set, setting it now"
-          git config --system core.longpaths true
-      - name: Git checkout PyTorch
-        uses: actions/checkout@v4
-        with:
-          path: "pytorch"
-          submodules: recursive
-      - name: Bootstrap APL
-        shell: cmd
-        run: |
-          "pytorch/.ci/pytorch/windows/arm64/bootstrap_apl.bat"
-      - name: Bootstrap Python
-        shell: cmd
-        run: |
-          "pytorch/.ci/pytorch/windows/arm64/bootstrap_python.bat"
-      - name: Bootstrap Rust
-        shell: cmd
-        run: |
-          "pytorch/.ci/pytorch/windows/arm64/bootstrap_rust.bat"
-      - uses: actions/download-artifact@v4.1.7
-        name: Download Build Artifacts
-        with:
-          name: libtorch-cpu-shared-with-deps-release
-          path: "${{ env.PYTORCH_FINAL_PACKAGE_DIR }}"
-      - name: Populate binary env
-        shell: bash
-        run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_populate_env.sh"
-      - name: Test PyTorch binary
-        shell: bash
-        run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_windows_test.sh"
-  libtorch-cpu-shared-with-deps-release-upload:  # Uploading
-    if: ${{ github.repository_owner == 'pytorch' }}
-    permissions:
-      id-token: write
-      contents: read
-    needs: libtorch-cpu-shared-with-deps-release-test
-    with:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
-      PACKAGE_TYPE: libtorch
-      # TODO: This is a legacy variable that we eventually want to get rid of in
-      #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cpu
-      GPU_ARCH_TYPE: cpu
-      LIBTORCH_CONFIG: release
-      LIBTORCH_VARIANT: shared-with-deps
-      # This is a dummy value for libtorch to work correctly with our batch scripts
-      # without this value pip does not get installed for some reason
-      DESIRED_PYTHON: "3.10"
-      build_name: libtorch-cpu-shared-with-deps-release
-    secrets:
-      github-token: ${{ secrets.GITHUB_TOKEN }}
-    uses: ./.github/workflows/_binary-upload.yml
diff --git a/.github/workflows/generated-windows-arm64-binary-wheel-nightly.yml b/.github/workflows/generated-windows-arm64-binary-wheel-nightly.yml
index 1368bc942350e..e3825b76cdbcc 100644
--- a/.github/workflows/generated-windows-arm64-binary-wheel-nightly.yml
+++ b/.github/workflows/generated-windows-arm64-binary-wheel-nightly.yml
@@ -14,6 +14,7 @@ on:
       # Release candidate tags look like: v1.11.0-rc1
       - v[0-9]+.[0-9]+.[0-9]+-rc[0-9]+
       - 'ciflow/binaries/*'
+      - 'ciflow/binaries_libtorch/*'
       - 'ciflow/binaries_wheel/*'
   workflow_dispatch:
 
@@ -41,7 +42,7 @@ jobs:
   get-label-type:
     if: github.repository_owner == 'pytorch'
     name: get-label-type
-    uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@main
+    uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@release/2.12
     with:
       triggering_actor: ${{ github.triggering_actor }}
       issue_owner: ${{ github.event.pull_request.user.login || github.event.issue.user.login }}
@@ -53,7 +54,7 @@ jobs:
     runs-on: "windows-11-arm64-preview"
     timeout-minutes: 360
     env:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
+      PYTORCH_ROOT: ${{ github.workspace }}
       PACKAGE_TYPE: wheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
@@ -81,36 +82,35 @@ jobs:
       - name: Git checkout PyTorch
         uses: actions/checkout@v4
         with:
-          path: "pytorch"
           submodules: recursive
       - name: Bootstrap Python
         shell: cmd
         run: |
-          "pytorch/.ci/pytorch/windows/arm64/bootstrap_python.bat"
+          ".ci/pytorch/windows/arm64/bootstrap_python.bat"
       - name: Bootstrap APL
         shell: cmd
         run: |
-          "pytorch/.ci/pytorch/windows/arm64/bootstrap_apl.bat"
+          ".ci/pytorch/windows/arm64/bootstrap_apl.bat"
       - name: Bootstrap Rust
         shell: cmd
         run: |
-          "pytorch/.ci/pytorch/windows/arm64/bootstrap_rust.bat"
+          ".ci/pytorch/windows/arm64/bootstrap_rust.bat"
       - name: Bootstrap sccache
         shell: cmd
         run: |
-          "pytorch/.ci/pytorch/windows/arm64/bootstrap_sccache.bat"
+          ".ci/pytorch/windows/arm64/bootstrap_sccache.bat"
       - name: Bootstrap Libuv
         shell: cmd
         run: |
-          "pytorch/.ci/pytorch/windows/arm64/bootstrap_libuv.bat"
+          ".ci/pytorch/windows/arm64/bootstrap_libuv.bat"
       - name: Populate binary env
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_populate_env.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_populate_env.sh"
       - name: Build PyTorch binary
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_windows_build.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_windows_build.sh"
       - uses: actions/upload-artifact@v4.4.0
         if: always()
         with:
@@ -126,7 +126,7 @@ jobs:
     runs-on: "windows-11-arm64-preview"
     timeout-minutes: 360
     env:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
+      PYTORCH_ROOT: ${{ github.workspace }}
       PACKAGE_TYPE: wheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
@@ -149,20 +149,19 @@ jobs:
       - name: Git checkout PyTorch
         uses: actions/checkout@v4
         with:
-          path: "pytorch"
           submodules: recursive
       - name: Bootstrap APL
         shell: cmd
         run: |
-          "pytorch/.ci/pytorch/windows/arm64/bootstrap_apl.bat"
+          ".ci/pytorch/windows/arm64/bootstrap_apl.bat"
       - name: Bootstrap Python
         shell: cmd
         run: |
-          "pytorch/.ci/pytorch/windows/arm64/bootstrap_python.bat"
+          ".ci/pytorch/windows/arm64/bootstrap_python.bat"
       - name: Bootstrap Rust
         shell: cmd
         run: |
-          "pytorch/.ci/pytorch/windows/arm64/bootstrap_rust.bat"
+          ".ci/pytorch/windows/arm64/bootstrap_rust.bat"
       - uses: actions/download-artifact@v4.1.7
         name: Download Build Artifacts
         with:
@@ -171,11 +170,11 @@ jobs:
       - name: Populate binary env
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_populate_env.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_populate_env.sh"
       - name: Test PyTorch binary
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_windows_test.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_windows_test.sh"
   wheel-py3_11-cpu-upload:  # Uploading
     if: ${{ github.repository_owner == 'pytorch' }}
     permissions:
@@ -183,7 +182,7 @@ jobs:
       contents: read
     needs: wheel-py3_11-cpu-test
     with:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
+      PYTORCH_ROOT: ${{ github.workspace }}
       PACKAGE_TYPE: wheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
@@ -193,6 +192,9 @@ jobs:
       build_name: wheel-py3_11-cpu
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
     uses: ./.github/workflows/_binary-upload.yml
   wheel-py3_12-cpu-build:
     if: ${{ github.repository_owner == 'pytorch' }}
@@ -200,7 +202,7 @@ jobs:
     runs-on: "windows-11-arm64-preview"
     timeout-minutes: 360
     env:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
+      PYTORCH_ROOT: ${{ github.workspace }}
       PACKAGE_TYPE: wheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
@@ -228,36 +230,35 @@ jobs:
       - name: Git checkout PyTorch
         uses: actions/checkout@v4
         with:
-          path: "pytorch"
           submodules: recursive
       - name: Bootstrap Python
         shell: cmd
         run: |
-          "pytorch/.ci/pytorch/windows/arm64/bootstrap_python.bat"
+          ".ci/pytorch/windows/arm64/bootstrap_python.bat"
       - name: Bootstrap APL
         shell: cmd
         run: |
-          "pytorch/.ci/pytorch/windows/arm64/bootstrap_apl.bat"
+          ".ci/pytorch/windows/arm64/bootstrap_apl.bat"
       - name: Bootstrap Rust
         shell: cmd
         run: |
-          "pytorch/.ci/pytorch/windows/arm64/bootstrap_rust.bat"
+          ".ci/pytorch/windows/arm64/bootstrap_rust.bat"
       - name: Bootstrap sccache
         shell: cmd
         run: |
-          "pytorch/.ci/pytorch/windows/arm64/bootstrap_sccache.bat"
+          ".ci/pytorch/windows/arm64/bootstrap_sccache.bat"
       - name: Bootstrap Libuv
         shell: cmd
         run: |
-          "pytorch/.ci/pytorch/windows/arm64/bootstrap_libuv.bat"
+          ".ci/pytorch/windows/arm64/bootstrap_libuv.bat"
       - name: Populate binary env
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_populate_env.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_populate_env.sh"
       - name: Build PyTorch binary
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_windows_build.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_windows_build.sh"
       - uses: actions/upload-artifact@v4.4.0
         if: always()
         with:
@@ -273,7 +274,7 @@ jobs:
     runs-on: "windows-11-arm64-preview"
     timeout-minutes: 360
     env:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
+      PYTORCH_ROOT: ${{ github.workspace }}
       PACKAGE_TYPE: wheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
@@ -296,20 +297,19 @@ jobs:
       - name: Git checkout PyTorch
         uses: actions/checkout@v4
         with:
-          path: "pytorch"
           submodules: recursive
       - name: Bootstrap APL
         shell: cmd
         run: |
-          "pytorch/.ci/pytorch/windows/arm64/bootstrap_apl.bat"
+          ".ci/pytorch/windows/arm64/bootstrap_apl.bat"
       - name: Bootstrap Python
         shell: cmd
         run: |
-          "pytorch/.ci/pytorch/windows/arm64/bootstrap_python.bat"
+          ".ci/pytorch/windows/arm64/bootstrap_python.bat"
       - name: Bootstrap Rust
         shell: cmd
         run: |
-          "pytorch/.ci/pytorch/windows/arm64/bootstrap_rust.bat"
+          ".ci/pytorch/windows/arm64/bootstrap_rust.bat"
       - uses: actions/download-artifact@v4.1.7
         name: Download Build Artifacts
         with:
@@ -318,11 +318,11 @@ jobs:
       - name: Populate binary env
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_populate_env.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_populate_env.sh"
       - name: Test PyTorch binary
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_windows_test.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_windows_test.sh"
   wheel-py3_12-cpu-upload:  # Uploading
     if: ${{ github.repository_owner == 'pytorch' }}
     permissions:
@@ -330,7 +330,7 @@ jobs:
       contents: read
     needs: wheel-py3_12-cpu-test
     with:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
+      PYTORCH_ROOT: ${{ github.workspace }}
       PACKAGE_TYPE: wheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
@@ -340,6 +340,9 @@ jobs:
       build_name: wheel-py3_12-cpu
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
     uses: ./.github/workflows/_binary-upload.yml
   wheel-py3_13-cpu-build:
     if: ${{ github.repository_owner == 'pytorch' }}
@@ -347,7 +350,7 @@ jobs:
     runs-on: "windows-11-arm64-preview"
     timeout-minutes: 360
     env:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
+      PYTORCH_ROOT: ${{ github.workspace }}
       PACKAGE_TYPE: wheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
@@ -375,36 +378,35 @@ jobs:
       - name: Git checkout PyTorch
         uses: actions/checkout@v4
         with:
-          path: "pytorch"
           submodules: recursive
       - name: Bootstrap Python
         shell: cmd
         run: |
-          "pytorch/.ci/pytorch/windows/arm64/bootstrap_python.bat"
+          ".ci/pytorch/windows/arm64/bootstrap_python.bat"
       - name: Bootstrap APL
         shell: cmd
         run: |
-          "pytorch/.ci/pytorch/windows/arm64/bootstrap_apl.bat"
+          ".ci/pytorch/windows/arm64/bootstrap_apl.bat"
       - name: Bootstrap Rust
         shell: cmd
         run: |
-          "pytorch/.ci/pytorch/windows/arm64/bootstrap_rust.bat"
+          ".ci/pytorch/windows/arm64/bootstrap_rust.bat"
       - name: Bootstrap sccache
         shell: cmd
         run: |
-          "pytorch/.ci/pytorch/windows/arm64/bootstrap_sccache.bat"
+          ".ci/pytorch/windows/arm64/bootstrap_sccache.bat"
       - name: Bootstrap Libuv
         shell: cmd
         run: |
-          "pytorch/.ci/pytorch/windows/arm64/bootstrap_libuv.bat"
+          ".ci/pytorch/windows/arm64/bootstrap_libuv.bat"
       - name: Populate binary env
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_populate_env.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_populate_env.sh"
       - name: Build PyTorch binary
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_windows_build.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_windows_build.sh"
       - uses: actions/upload-artifact@v4.4.0
         if: always()
         with:
@@ -420,7 +422,7 @@ jobs:
     runs-on: "windows-11-arm64-preview"
     timeout-minutes: 360
     env:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
+      PYTORCH_ROOT: ${{ github.workspace }}
       PACKAGE_TYPE: wheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
@@ -443,20 +445,19 @@ jobs:
       - name: Git checkout PyTorch
         uses: actions/checkout@v4
         with:
-          path: "pytorch"
           submodules: recursive
       - name: Bootstrap APL
         shell: cmd
         run: |
-          "pytorch/.ci/pytorch/windows/arm64/bootstrap_apl.bat"
+          ".ci/pytorch/windows/arm64/bootstrap_apl.bat"
       - name: Bootstrap Python
         shell: cmd
         run: |
-          "pytorch/.ci/pytorch/windows/arm64/bootstrap_python.bat"
+          ".ci/pytorch/windows/arm64/bootstrap_python.bat"
       - name: Bootstrap Rust
         shell: cmd
         run: |
-          "pytorch/.ci/pytorch/windows/arm64/bootstrap_rust.bat"
+          ".ci/pytorch/windows/arm64/bootstrap_rust.bat"
       - uses: actions/download-artifact@v4.1.7
         name: Download Build Artifacts
         with:
@@ -465,11 +466,11 @@ jobs:
       - name: Populate binary env
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_populate_env.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_populate_env.sh"
       - name: Test PyTorch binary
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_windows_test.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_windows_test.sh"
   wheel-py3_13-cpu-upload:  # Uploading
     if: ${{ github.repository_owner == 'pytorch' }}
     permissions:
@@ -477,7 +478,7 @@ jobs:
       contents: read
     needs: wheel-py3_13-cpu-test
     with:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
+      PYTORCH_ROOT: ${{ github.workspace }}
       PACKAGE_TYPE: wheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
@@ -487,4 +488,67 @@ jobs:
       build_name: wheel-py3_13-cpu
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
+    uses: ./.github/workflows/_binary-upload.yml
+  libtorch-cpu-shared-with-deps-release-extract:
+    if: ${{ github.repository_owner == 'pytorch' }}
+    needs:
+      - wheel-py3_11-cpu-build
+      - get-label-type
+    runs-on: "${{ needs.get-label-type.outputs.label-type }}linux.4xlarge"
+    timeout-minutes: 60
+    env:
+      DESIRED_CUDA: cpu
+      LIBTORCH_VARIANT: shared-with-deps
+    steps:
+      - name: Checkout PyTorch
+        uses: actions/checkout@v4
+        with:
+          ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
+          sparse-checkout: .ci/libtorch/
+          show-progress: false
+      - uses: actions/download-artifact@v4.1.7
+        name: Download Wheel Artifact
+        with:
+          name: wheel-py3_11-cpu
+          path: "${{ runner.temp }}/wheel_artifact/"
+      - name: Extract libtorch from wheel
+        shell: bash
+        run: |
+          set -eux
+          mkdir -p "${{ runner.temp }}/libtorch_output"
+          python3 .ci/libtorch/extract_libtorch_from_wheel.py \
+            --wheel-dir "${{ runner.temp }}/wheel_artifact" \
+            --output-dir "${{ runner.temp }}/libtorch_output" \
+            --platform windows \
+            --desired-cuda "$DESIRED_CUDA" \
+            --libtorch-variant "$LIBTORCH_VARIANT"
+      - uses: actions/upload-artifact@v4.4.0
+        if: always()
+        with:
+          name: libtorch-cpu-shared-with-deps-release
+          retention-days: 14
+          if-no-files-found: error
+          path: "${{ runner.temp }}/libtorch_output/"
+  libtorch-cpu-shared-with-deps-release-upload:  # Uploading
+    if: ${{ github.repository_owner == 'pytorch' }}
+    permissions:
+      id-token: write
+      contents: read
+    needs: libtorch-cpu-shared-with-deps-release-extract
+    with:
+      PYTORCH_ROOT: ${{ github.workspace }}
+      PACKAGE_TYPE: libtorch
+      DESIRED_CUDA: cpu
+      GPU_ARCH_TYPE: cpu
+      LIBTORCH_CONFIG: release
+      LIBTORCH_VARIANT: shared-with-deps
+      build_name: libtorch-cpu-shared-with-deps-release
+    secrets:
+      github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
     uses: ./.github/workflows/_binary-upload.yml
diff --git a/.github/workflows/generated-windows-binary-libtorch-debug-nightly.yml b/.github/workflows/generated-windows-binary-libtorch-debug-nightly.yml
index 3ca3364e5de88..6953a11b64677 100644
--- a/.github/workflows/generated-windows-binary-libtorch-debug-nightly.yml
+++ b/.github/workflows/generated-windows-binary-libtorch-debug-nightly.yml
@@ -35,7 +35,7 @@ jobs:
   get-label-type:
     if: github.repository_owner == 'pytorch'
     name: get-label-type
-    uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@main
+    uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@release/2.12
     with:
       triggering_actor: ${{ github.triggering_actor }}
       issue_owner: ${{ github.event.pull_request.user.login || github.event.issue.user.login }}
@@ -47,7 +47,7 @@ jobs:
     runs-on: "${{ needs.get-label-type.outputs.label-type }}windows.12xlarge"
     timeout-minutes: 360
     env:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
+      PYTORCH_ROOT: ${{ github.workspace }}
       PACKAGE_TYPE: libtorch
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
@@ -84,7 +84,7 @@ jobs:
           echo "instance-type: $(get_ec2_metadata instance-type)"
           echo "system info $(uname -a)"
       - name: "[FB EMPLOYEES] Enable SSH (Click me for login details)"
-        uses: pytorch/test-infra/.github/actions/setup-ssh@main
+        uses: pytorch/test-infra/.github/actions/setup-ssh@release/2.12
         continue-on-error: true
         with:
           github-secret: ${{ secrets.GITHUB_TOKEN }}
@@ -116,23 +116,20 @@ jobs:
       - name: Checkout PyTorch
         uses: actions/checkout@v4
         with:
-          ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
           submodules: recursive
-          path: pytorch
           show-progress: false
       - name: Clean PyTorch checkout
         run: |
           # Remove any artifacts from the previous checkouts
           git clean -fxd
-        working-directory: pytorch
       - name: Populate binary env
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_populate_env.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_populate_env.sh"
       - name: Build PyTorch binary
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_windows_build.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_windows_build.sh"
       - uses: actions/upload-artifact@v4.4.0
         if: always()
         with:
@@ -142,14 +139,12 @@ jobs:
           path: "${{ env.PYTORCH_FINAL_PACKAGE_DIR }}"
       - name: Wait until all sessions have drained
         shell: powershell
-        working-directory: pytorch
         if: always()
         timeout-minutes: 120
         run: |
           .github\scripts\wait_for_ssh_to_drain.ps1
       - name: Kill active ssh sessions if still around (Useful if workflow was cancelled)
         shell: powershell
-        working-directory: pytorch
         if: always()
         run: |
           .github\scripts\kill_active_ssh_sessions.ps1
@@ -162,7 +157,7 @@ jobs:
     runs-on: "${{ needs.get-label-type.outputs.label-type }}windows.4xlarge"
     timeout-minutes: 360
     env:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
+      PYTORCH_ROOT: ${{ github.workspace }}
       PACKAGE_TYPE: libtorch
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
@@ -190,7 +185,7 @@ jobs:
           echo "instance-type: $(get_ec2_metadata instance-type)"
           echo "system info $(uname -a)"
       - name: "[FB EMPLOYEES] Enable SSH (Click me for login details)"
-        uses: pytorch/test-infra/.github/actions/setup-ssh@main
+        uses: pytorch/test-infra/.github/actions/setup-ssh@release/2.12
         continue-on-error: true
         with:
           github-secret: ${{ secrets.GITHUB_TOKEN }}
@@ -222,15 +217,12 @@ jobs:
       - name: Checkout PyTorch
         uses: actions/checkout@v4
         with:
-          ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
           submodules: recursive
-          path: pytorch
           show-progress: false
       - name: Clean PyTorch checkout
         run: |
           # Remove any artifacts from the previous checkouts
           git clean -fxd
-        working-directory: pytorch
       # NOTE: These environment variables are put here so that they can be applied on every job equally
       #       They are also here because setting them at a workflow level doesn't give us access to the
       #       runner.temp variable, which we need.
@@ -248,21 +240,19 @@ jobs:
       - name: Populate binary env
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_populate_env.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_populate_env.sh"
       - name: Test PyTorch binary
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_windows_test.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_windows_test.sh"
       - name: Wait until all sessions have drained
         shell: powershell
-        working-directory: pytorch
         if: always()
         timeout-minutes: 120
         run: |
           .github\scripts\wait_for_ssh_to_drain.ps1
       - name: Kill active ssh sessions if still around (Useful if workflow was cancelled)
         shell: powershell
-        working-directory: pytorch
         if: always()
         run: |
           .github\scripts\kill_active_ssh_sessions.ps1
@@ -273,7 +263,7 @@ jobs:
       contents: read
     needs: libtorch-cpu-shared-with-deps-debug-test
     with:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
+      PYTORCH_ROOT: ${{ github.workspace }}
       PACKAGE_TYPE: libtorch
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
@@ -287,6 +277,9 @@ jobs:
       build_name: libtorch-cpu-shared-with-deps-debug
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
     uses: ./.github/workflows/_binary-upload.yml
   libtorch-cuda12_6-shared-with-deps-debug-build:
     if: ${{ github.repository_owner == 'pytorch' }}
@@ -294,7 +287,7 @@ jobs:
     runs-on: "${{ needs.get-label-type.outputs.label-type }}windows.12xlarge"
     timeout-minutes: 360
     env:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
+      PYTORCH_ROOT: ${{ github.workspace }}
       PACKAGE_TYPE: libtorch
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
@@ -332,7 +325,7 @@ jobs:
           echo "instance-type: $(get_ec2_metadata instance-type)"
           echo "system info $(uname -a)"
       - name: "[FB EMPLOYEES] Enable SSH (Click me for login details)"
-        uses: pytorch/test-infra/.github/actions/setup-ssh@main
+        uses: pytorch/test-infra/.github/actions/setup-ssh@release/2.12
         continue-on-error: true
         with:
           github-secret: ${{ secrets.GITHUB_TOKEN }}
@@ -364,23 +357,20 @@ jobs:
       - name: Checkout PyTorch
         uses: actions/checkout@v4
         with:
-          ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
           submodules: recursive
-          path: pytorch
           show-progress: false
       - name: Clean PyTorch checkout
         run: |
           # Remove any artifacts from the previous checkouts
           git clean -fxd
-        working-directory: pytorch
       - name: Populate binary env
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_populate_env.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_populate_env.sh"
       - name: Build PyTorch binary
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_windows_build.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_windows_build.sh"
       - uses: actions/upload-artifact@v4.4.0
         if: always()
         with:
@@ -390,14 +380,12 @@ jobs:
           path: "${{ env.PYTORCH_FINAL_PACKAGE_DIR }}"
       - name: Wait until all sessions have drained
         shell: powershell
-        working-directory: pytorch
         if: always()
         timeout-minutes: 120
         run: |
           .github\scripts\wait_for_ssh_to_drain.ps1
       - name: Kill active ssh sessions if still around (Useful if workflow was cancelled)
         shell: powershell
-        working-directory: pytorch
         if: always()
         run: |
           .github\scripts\kill_active_ssh_sessions.ps1
@@ -410,7 +398,7 @@ jobs:
     runs-on: "${{ needs.get-label-type.outputs.label-type }}windows.g4dn.xlarge"
     timeout-minutes: 360
     env:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
+      PYTORCH_ROOT: ${{ github.workspace }}
       PACKAGE_TYPE: libtorch
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
@@ -439,7 +427,7 @@ jobs:
           echo "instance-type: $(get_ec2_metadata instance-type)"
           echo "system info $(uname -a)"
       - name: "[FB EMPLOYEES] Enable SSH (Click me for login details)"
-        uses: pytorch/test-infra/.github/actions/setup-ssh@main
+        uses: pytorch/test-infra/.github/actions/setup-ssh@release/2.12
         continue-on-error: true
         with:
           github-secret: ${{ secrets.GITHUB_TOKEN }}
@@ -471,15 +459,12 @@ jobs:
       - name: Checkout PyTorch
         uses: actions/checkout@v4
         with:
-          ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
           submodules: recursive
-          path: pytorch
           show-progress: false
       - name: Clean PyTorch checkout
         run: |
           # Remove any artifacts from the previous checkouts
           git clean -fxd
-        working-directory: pytorch
       # NOTE: These environment variables are put here so that they can be applied on every job equally
       #       They are also here because setting them at a workflow level doesn't give us access to the
       #       runner.temp variable, which we need.
@@ -497,21 +482,19 @@ jobs:
       - name: Populate binary env
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_populate_env.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_populate_env.sh"
       - name: Test PyTorch binary
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_windows_test.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_windows_test.sh"
       - name: Wait until all sessions have drained
         shell: powershell
-        working-directory: pytorch
         if: always()
         timeout-minutes: 120
         run: |
           .github\scripts\wait_for_ssh_to_drain.ps1
       - name: Kill active ssh sessions if still around (Useful if workflow was cancelled)
         shell: powershell
-        working-directory: pytorch
         if: always()
         run: |
           .github\scripts\kill_active_ssh_sessions.ps1
@@ -522,7 +505,7 @@ jobs:
       contents: read
     needs: libtorch-cuda12_6-shared-with-deps-debug-test
     with:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
+      PYTORCH_ROOT: ${{ github.workspace }}
       PACKAGE_TYPE: libtorch
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
@@ -537,19 +520,22 @@ jobs:
       build_name: libtorch-cuda12_6-shared-with-deps-debug
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
     uses: ./.github/workflows/_binary-upload.yml
-  libtorch-cuda12_8-shared-with-deps-debug-build:
+  libtorch-cuda13_0-shared-with-deps-debug-build:
     if: ${{ github.repository_owner == 'pytorch' }}
     needs: get-label-type
     runs-on: "${{ needs.get-label-type.outputs.label-type }}windows.12xlarge"
     timeout-minutes: 360
     env:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
+      PYTORCH_ROOT: ${{ github.workspace }}
       PACKAGE_TYPE: libtorch
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu128
-      GPU_ARCH_VERSION: "12.8"
+      DESIRED_CUDA: cu130
+      GPU_ARCH_VERSION: "13.0"
       GPU_ARCH_TYPE: cuda
       SKIP_ALL_TESTS: 1
       LIBTORCH_CONFIG: debug
@@ -582,7 +568,7 @@ jobs:
           echo "instance-type: $(get_ec2_metadata instance-type)"
           echo "system info $(uname -a)"
       - name: "[FB EMPLOYEES] Enable SSH (Click me for login details)"
-        uses: pytorch/test-infra/.github/actions/setup-ssh@main
+        uses: pytorch/test-infra/.github/actions/setup-ssh@release/2.12
         continue-on-error: true
         with:
           github-secret: ${{ secrets.GITHUB_TOKEN }}
@@ -614,58 +600,53 @@ jobs:
       - name: Checkout PyTorch
         uses: actions/checkout@v4
         with:
-          ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
           submodules: recursive
-          path: pytorch
           show-progress: false
       - name: Clean PyTorch checkout
         run: |
           # Remove any artifacts from the previous checkouts
           git clean -fxd
-        working-directory: pytorch
       - name: Populate binary env
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_populate_env.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_populate_env.sh"
       - name: Build PyTorch binary
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_windows_build.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_windows_build.sh"
       - uses: actions/upload-artifact@v4.4.0
         if: always()
         with:
-          name: libtorch-cuda12_8-shared-with-deps-debug
+          name: libtorch-cuda13_0-shared-with-deps-debug
           retention-days: 14
           if-no-files-found: error
           path: "${{ env.PYTORCH_FINAL_PACKAGE_DIR }}"
       - name: Wait until all sessions have drained
         shell: powershell
-        working-directory: pytorch
         if: always()
         timeout-minutes: 120
         run: |
           .github\scripts\wait_for_ssh_to_drain.ps1
       - name: Kill active ssh sessions if still around (Useful if workflow was cancelled)
         shell: powershell
-        working-directory: pytorch
         if: always()
         run: |
           .github\scripts\kill_active_ssh_sessions.ps1
 
-  libtorch-cuda12_8-shared-with-deps-debug-test:  # Testing
+  libtorch-cuda13_0-shared-with-deps-debug-test:  # Testing
     if: ${{ github.repository_owner == 'pytorch' }}
     needs:
-      - libtorch-cuda12_8-shared-with-deps-debug-build
+      - libtorch-cuda13_0-shared-with-deps-debug-build
       - get-label-type
     runs-on: "${{ needs.get-label-type.outputs.label-type }}windows.g4dn.xlarge"
     timeout-minutes: 360
     env:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
+      PYTORCH_ROOT: ${{ github.workspace }}
       PACKAGE_TYPE: libtorch
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu128
-      GPU_ARCH_VERSION: "12.8"
+      DESIRED_CUDA: cu130
+      GPU_ARCH_VERSION: "13.0"
       GPU_ARCH_TYPE: cuda
       SKIP_ALL_TESTS: 1
       LIBTORCH_CONFIG: debug
@@ -689,7 +670,7 @@ jobs:
           echo "instance-type: $(get_ec2_metadata instance-type)"
           echo "system info $(uname -a)"
       - name: "[FB EMPLOYEES] Enable SSH (Click me for login details)"
-        uses: pytorch/test-infra/.github/actions/setup-ssh@main
+        uses: pytorch/test-infra/.github/actions/setup-ssh@release/2.12
         continue-on-error: true
         with:
           github-secret: ${{ secrets.GITHUB_TOKEN }}
@@ -721,15 +702,12 @@ jobs:
       - name: Checkout PyTorch
         uses: actions/checkout@v4
         with:
-          ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
           submodules: recursive
-          path: pytorch
           show-progress: false
       - name: Clean PyTorch checkout
         run: |
           # Remove any artifacts from the previous checkouts
           git clean -fxd
-        working-directory: pytorch
       # NOTE: These environment variables are put here so that they can be applied on every job equally
       #       They are also here because setting them at a workflow level doesn't give us access to the
       #       runner.temp variable, which we need.
@@ -742,64 +720,65 @@ jobs:
       - uses: actions/download-artifact@v4.1.7
         name: Download Build Artifacts
         with:
-          name: libtorch-cuda12_8-shared-with-deps-debug
+          name: libtorch-cuda13_0-shared-with-deps-debug
           path: "${{ env.PYTORCH_FINAL_PACKAGE_DIR }}"
       - name: Populate binary env
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_populate_env.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_populate_env.sh"
       - name: Test PyTorch binary
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_windows_test.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_windows_test.sh"
       - name: Wait until all sessions have drained
         shell: powershell
-        working-directory: pytorch
         if: always()
         timeout-minutes: 120
         run: |
           .github\scripts\wait_for_ssh_to_drain.ps1
       - name: Kill active ssh sessions if still around (Useful if workflow was cancelled)
         shell: powershell
-        working-directory: pytorch
         if: always()
         run: |
           .github\scripts\kill_active_ssh_sessions.ps1
-  libtorch-cuda12_8-shared-with-deps-debug-upload:  # Uploading
+  libtorch-cuda13_0-shared-with-deps-debug-upload:  # Uploading
     if: ${{ github.repository_owner == 'pytorch' }}
     permissions:
       id-token: write
       contents: read
-    needs: libtorch-cuda12_8-shared-with-deps-debug-test
+    needs: libtorch-cuda13_0-shared-with-deps-debug-test
     with:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
+      PYTORCH_ROOT: ${{ github.workspace }}
       PACKAGE_TYPE: libtorch
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu128
-      GPU_ARCH_VERSION: "12.8"
+      DESIRED_CUDA: cu130
+      GPU_ARCH_VERSION: "13.0"
       GPU_ARCH_TYPE: cuda
       LIBTORCH_CONFIG: debug
       LIBTORCH_VARIANT: shared-with-deps
       # This is a dummy value for libtorch to work correctly with our batch scripts
       # without this value pip does not get installed for some reason
       DESIRED_PYTHON: "3.10"
-      build_name: libtorch-cuda12_8-shared-with-deps-debug
+      build_name: libtorch-cuda13_0-shared-with-deps-debug
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
     uses: ./.github/workflows/_binary-upload.yml
-  libtorch-cuda13_0-shared-with-deps-debug-build:
+  libtorch-cuda13_2-shared-with-deps-debug-build:
     if: ${{ github.repository_owner == 'pytorch' }}
     needs: get-label-type
     runs-on: "${{ needs.get-label-type.outputs.label-type }}windows.12xlarge"
     timeout-minutes: 360
     env:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
+      PYTORCH_ROOT: ${{ github.workspace }}
       PACKAGE_TYPE: libtorch
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu130
-      GPU_ARCH_VERSION: "13.0"
+      DESIRED_CUDA: cu132
+      GPU_ARCH_VERSION: "13.2"
       GPU_ARCH_TYPE: cuda
       SKIP_ALL_TESTS: 1
       LIBTORCH_CONFIG: debug
@@ -832,7 +811,7 @@ jobs:
           echo "instance-type: $(get_ec2_metadata instance-type)"
           echo "system info $(uname -a)"
       - name: "[FB EMPLOYEES] Enable SSH (Click me for login details)"
-        uses: pytorch/test-infra/.github/actions/setup-ssh@main
+        uses: pytorch/test-infra/.github/actions/setup-ssh@release/2.12
         continue-on-error: true
         with:
           github-secret: ${{ secrets.GITHUB_TOKEN }}
@@ -864,58 +843,53 @@ jobs:
       - name: Checkout PyTorch
         uses: actions/checkout@v4
         with:
-          ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
           submodules: recursive
-          path: pytorch
           show-progress: false
       - name: Clean PyTorch checkout
         run: |
           # Remove any artifacts from the previous checkouts
           git clean -fxd
-        working-directory: pytorch
       - name: Populate binary env
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_populate_env.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_populate_env.sh"
       - name: Build PyTorch binary
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_windows_build.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_windows_build.sh"
       - uses: actions/upload-artifact@v4.4.0
         if: always()
         with:
-          name: libtorch-cuda13_0-shared-with-deps-debug
+          name: libtorch-cuda13_2-shared-with-deps-debug
           retention-days: 14
           if-no-files-found: error
           path: "${{ env.PYTORCH_FINAL_PACKAGE_DIR }}"
       - name: Wait until all sessions have drained
         shell: powershell
-        working-directory: pytorch
         if: always()
         timeout-minutes: 120
         run: |
           .github\scripts\wait_for_ssh_to_drain.ps1
       - name: Kill active ssh sessions if still around (Useful if workflow was cancelled)
         shell: powershell
-        working-directory: pytorch
         if: always()
         run: |
           .github\scripts\kill_active_ssh_sessions.ps1
 
-  libtorch-cuda13_0-shared-with-deps-debug-test:  # Testing
+  libtorch-cuda13_2-shared-with-deps-debug-test:  # Testing
     if: ${{ github.repository_owner == 'pytorch' }}
     needs:
-      - libtorch-cuda13_0-shared-with-deps-debug-build
+      - libtorch-cuda13_2-shared-with-deps-debug-build
       - get-label-type
     runs-on: "${{ needs.get-label-type.outputs.label-type }}windows.g4dn.xlarge"
     timeout-minutes: 360
     env:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
+      PYTORCH_ROOT: ${{ github.workspace }}
       PACKAGE_TYPE: libtorch
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu130
-      GPU_ARCH_VERSION: "13.0"
+      DESIRED_CUDA: cu132
+      GPU_ARCH_VERSION: "13.2"
       GPU_ARCH_TYPE: cuda
       SKIP_ALL_TESTS: 1
       LIBTORCH_CONFIG: debug
@@ -939,7 +913,7 @@ jobs:
           echo "instance-type: $(get_ec2_metadata instance-type)"
           echo "system info $(uname -a)"
       - name: "[FB EMPLOYEES] Enable SSH (Click me for login details)"
-        uses: pytorch/test-infra/.github/actions/setup-ssh@main
+        uses: pytorch/test-infra/.github/actions/setup-ssh@release/2.12
         continue-on-error: true
         with:
           github-secret: ${{ secrets.GITHUB_TOKEN }}
@@ -971,15 +945,12 @@ jobs:
       - name: Checkout PyTorch
         uses: actions/checkout@v4
         with:
-          ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
           submodules: recursive
-          path: pytorch
           show-progress: false
       - name: Clean PyTorch checkout
         run: |
           # Remove any artifacts from the previous checkouts
           git clean -fxd
-        working-directory: pytorch
       # NOTE: These environment variables are put here so that they can be applied on every job equally
       #       They are also here because setting them at a workflow level doesn't give us access to the
       #       runner.temp variable, which we need.
@@ -992,49 +963,50 @@ jobs:
       - uses: actions/download-artifact@v4.1.7
         name: Download Build Artifacts
         with:
-          name: libtorch-cuda13_0-shared-with-deps-debug
+          name: libtorch-cuda13_2-shared-with-deps-debug
           path: "${{ env.PYTORCH_FINAL_PACKAGE_DIR }}"
       - name: Populate binary env
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_populate_env.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_populate_env.sh"
       - name: Test PyTorch binary
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_windows_test.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_windows_test.sh"
       - name: Wait until all sessions have drained
         shell: powershell
-        working-directory: pytorch
         if: always()
         timeout-minutes: 120
         run: |
           .github\scripts\wait_for_ssh_to_drain.ps1
       - name: Kill active ssh sessions if still around (Useful if workflow was cancelled)
         shell: powershell
-        working-directory: pytorch
         if: always()
         run: |
           .github\scripts\kill_active_ssh_sessions.ps1
-  libtorch-cuda13_0-shared-with-deps-debug-upload:  # Uploading
+  libtorch-cuda13_2-shared-with-deps-debug-upload:  # Uploading
     if: ${{ github.repository_owner == 'pytorch' }}
     permissions:
       id-token: write
       contents: read
-    needs: libtorch-cuda13_0-shared-with-deps-debug-test
+    needs: libtorch-cuda13_2-shared-with-deps-debug-test
     with:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
+      PYTORCH_ROOT: ${{ github.workspace }}
       PACKAGE_TYPE: libtorch
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu130
-      GPU_ARCH_VERSION: "13.0"
+      DESIRED_CUDA: cu132
+      GPU_ARCH_VERSION: "13.2"
       GPU_ARCH_TYPE: cuda
       LIBTORCH_CONFIG: debug
       LIBTORCH_VARIANT: shared-with-deps
       # This is a dummy value for libtorch to work correctly with our batch scripts
       # without this value pip does not get installed for some reason
       DESIRED_PYTHON: "3.10"
-      build_name: libtorch-cuda13_0-shared-with-deps-debug
+      build_name: libtorch-cuda13_2-shared-with-deps-debug
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
     uses: ./.github/workflows/_binary-upload.yml
diff --git a/.github/workflows/generated-windows-binary-libtorch-release-nightly.yml b/.github/workflows/generated-windows-binary-libtorch-release-nightly.yml
deleted file mode 100644
index c6d1e2cf3b017..0000000000000
--- a/.github/workflows/generated-windows-binary-libtorch-release-nightly.yml
+++ /dev/null
@@ -1,1040 +0,0 @@
-# @generated DO NOT EDIT MANUALLY
-
-# Template is at:    .github/templates/windows_binary_build_workflow.yml.j2
-# Generation script: .github/scripts/generate_ci_workflows.py
-name: windows-binary-libtorch-release
-
-on:
-  push:
-    # NOTE: Meta Employees can trigger new nightlies using: https://fburl.com/trigger_pytorch_nightly_build
-    branches:
-      - nightly
-    tags:
-      # NOTE: Binary build pipelines should only get triggered on release candidate builds
-      # Release candidate tags look like: v1.11.0-rc1
-      - v[0-9]+.[0-9]+.[0-9]+-rc[0-9]+
-      - 'ciflow/binaries/*'
-      - 'ciflow/binaries_libtorch/*'
-  workflow_dispatch:
-
-env:
-  # Needed for conda builds
-  ALPINE_IMAGE: "308535385114.dkr.ecr.us-east-1.amazonaws.com/tool/alpine"
-  AWS_DEFAULT_REGION: us-east-1
-  BUILD_ENVIRONMENT: windows-binary-libtorch-release
-  GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
-  PR_NUMBER: ${{ github.event.pull_request.number }}
-  SHA1: ${{ github.event.pull_request.head.sha || github.sha }}
-  SKIP_ALL_TESTS: 1
-  OS: windows
-concurrency:
-  group: windows-binary-libtorch-release-${{ github.event.pull_request.number || github.ref_name }}-${{ github.ref_type == 'branch' && github.sha }}-${{ github.event_name == 'workflow_dispatch' }}
-  cancel-in-progress: true
-
-jobs:
-  get-label-type:
-    if: github.repository_owner == 'pytorch'
-    name: get-label-type
-    uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@main
-    with:
-      triggering_actor: ${{ github.triggering_actor }}
-      issue_owner: ${{ github.event.pull_request.user.login || github.event.issue.user.login }}
-      curr_branch: ${{ github.head_ref || github.ref_name }}
-      curr_ref_type: ${{ github.ref_type }}
-  libtorch-cpu-shared-with-deps-release-build:
-    if: ${{ github.repository_owner == 'pytorch' }}
-    needs: get-label-type
-    runs-on: "${{ needs.get-label-type.outputs.label-type }}windows.12xlarge"
-    timeout-minutes: 360
-    env:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
-      PACKAGE_TYPE: libtorch
-      # TODO: This is a legacy variable that we eventually want to get rid of in
-      #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cpu
-      GPU_ARCH_TYPE: cpu
-      SKIP_ALL_TESTS: 1
-      LIBTORCH_CONFIG: release
-      LIBTORCH_VARIANT: shared-with-deps
-      # This is a dummy value for libtorch to work correctly with our batch scripts
-      # without this value pip does not get installed for some reason
-      DESIRED_PYTHON: "3.10"
-    steps:
-      # NOTE: These environment variables are put here so that they can be applied on every job equally
-      #       They are also here because setting them at a workflow level doesn't give us access to the
-      #       runner.temp variable, which we need.
-      - name: Populate binary env
-        shell: bash
-        run: |
-          echo "BINARY_ENV_FILE=${RUNNER_TEMP}/env" >> "${GITHUB_ENV}"
-          echo "PYTORCH_FINAL_PACKAGE_DIR=${RUNNER_TEMP}/artifacts" >> "${GITHUB_ENV}"
-          echo "WIN_PACKAGE_WORK_DIR=${RUNNER_TEMP}"
-      - name: Display EC2 information
-        shell: bash
-        run: |
-          set -euo pipefail
-          function get_ec2_metadata() {
-            # Pulled from instance metadata endpoint for EC2
-            # see https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/instancedata-data-retrieval.html
-            category=$1
-            curl -H "X-aws-ec2-metadata-token: $(curl -s -X PUT "http://169.254.169.254/latest/api/token" -H "X-aws-ec2-metadata-token-ttl-seconds: 30")" -fsSL "http://169.254.169.254/latest/meta-data/${category}"
-          }
-          echo "ami-id: $(get_ec2_metadata ami-id)"
-          echo "instance-id: $(get_ec2_metadata instance-id)"
-          echo "instance-type: $(get_ec2_metadata instance-type)"
-          echo "system info $(uname -a)"
-      - name: "[FB EMPLOYEES] Enable SSH (Click me for login details)"
-        uses: pytorch/test-infra/.github/actions/setup-ssh@main
-        continue-on-error: true
-        with:
-          github-secret: ${{ secrets.GITHUB_TOKEN }}
-      - name: Enable git long paths and symlinks on Windows and disable fsmonitor daemon
-        shell: bash
-        run: |
-          git config --global core.longpaths true
-          git config --global core.symlinks true
-
-          # https://git-scm.com/docs/git-fsmonitor--daemon.  The daemon could lock
-          # the directory on Windows and prevent GHA from checking out as reported
-          # in https://github.com/actions/checkout/issues/1018
-          git config --global core.fsmonitor false
-      # Needed for binary builds, see: https://github.com/pytorch/pytorch/issues/73339#issuecomment-1058981560
-      - name: Enable long paths on Windows
-        shell: powershell
-        run: |
-          Set-ItemProperty -Path "HKLM:\\SYSTEM\CurrentControlSet\Control\FileSystem" -Name "LongPathsEnabled" -Value 1
-      # Since it's just a defensive command, the workflow should continue even the command fails. This step can be
-      # removed once Windows Defender is removed from the AMI
-      - name: Disables Windows Defender scheduled and real-time scanning for files in directories used by PyTorch
-        continue-on-error: true
-        shell: powershell
-        run: |
-          Add-MpPreference -ExclusionPath $(Get-Location).tostring(),$Env:TEMP -ErrorAction Ignore
-          # Let's both exclude the path and disable Windows Defender completely just to be sure
-          # that it doesn't interfere
-          Set-MpPreference -DisableRealtimeMonitoring $True -ErrorAction Ignore
-      - name: Checkout PyTorch
-        uses: actions/checkout@v4
-        with:
-          ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
-          submodules: recursive
-          path: pytorch
-          show-progress: false
-      - name: Clean PyTorch checkout
-        run: |
-          # Remove any artifacts from the previous checkouts
-          git clean -fxd
-        working-directory: pytorch
-      - name: Populate binary env
-        shell: bash
-        run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_populate_env.sh"
-      - name: Build PyTorch binary
-        shell: bash
-        run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_windows_build.sh"
-      - uses: actions/upload-artifact@v4.4.0
-        if: always()
-        with:
-          name: libtorch-cpu-shared-with-deps-release
-          retention-days: 14
-          if-no-files-found: error
-          path: "${{ env.PYTORCH_FINAL_PACKAGE_DIR }}"
-      - name: Wait until all sessions have drained
-        shell: powershell
-        working-directory: pytorch
-        if: always()
-        timeout-minutes: 120
-        run: |
-          .github\scripts\wait_for_ssh_to_drain.ps1
-      - name: Kill active ssh sessions if still around (Useful if workflow was cancelled)
-        shell: powershell
-        working-directory: pytorch
-        if: always()
-        run: |
-          .github\scripts\kill_active_ssh_sessions.ps1
-
-  libtorch-cpu-shared-with-deps-release-test:  # Testing
-    if: ${{ github.repository_owner == 'pytorch' }}
-    needs:
-      - libtorch-cpu-shared-with-deps-release-build
-      - get-label-type
-    runs-on: "${{ needs.get-label-type.outputs.label-type }}windows.4xlarge"
-    timeout-minutes: 360
-    env:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
-      PACKAGE_TYPE: libtorch
-      # TODO: This is a legacy variable that we eventually want to get rid of in
-      #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cpu
-      GPU_ARCH_TYPE: cpu
-      SKIP_ALL_TESTS: 1
-      LIBTORCH_CONFIG: release
-      LIBTORCH_VARIANT: shared-with-deps
-      # This is a dummy value for libtorch to work correctly with our batch scripts
-      # without this value pip does not get installed for some reason
-      DESIRED_PYTHON: "3.10"
-    steps:
-      - name: Display EC2 information
-        shell: bash
-        run: |
-          set -euo pipefail
-          function get_ec2_metadata() {
-            # Pulled from instance metadata endpoint for EC2
-            # see https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/instancedata-data-retrieval.html
-            category=$1
-            curl -H "X-aws-ec2-metadata-token: $(curl -s -X PUT "http://169.254.169.254/latest/api/token" -H "X-aws-ec2-metadata-token-ttl-seconds: 30")" -fsSL "http://169.254.169.254/latest/meta-data/${category}"
-          }
-          echo "ami-id: $(get_ec2_metadata ami-id)"
-          echo "instance-id: $(get_ec2_metadata instance-id)"
-          echo "instance-type: $(get_ec2_metadata instance-type)"
-          echo "system info $(uname -a)"
-      - name: "[FB EMPLOYEES] Enable SSH (Click me for login details)"
-        uses: pytorch/test-infra/.github/actions/setup-ssh@main
-        continue-on-error: true
-        with:
-          github-secret: ${{ secrets.GITHUB_TOKEN }}
-      - name: Enable git long paths and symlinks on Windows and disable fsmonitor daemon
-        shell: bash
-        run: |
-          git config --global core.longpaths true
-          git config --global core.symlinks true
-
-          # https://git-scm.com/docs/git-fsmonitor--daemon.  The daemon could lock
-          # the directory on Windows and prevent GHA from checking out as reported
-          # in https://github.com/actions/checkout/issues/1018
-          git config --global core.fsmonitor false
-      # Needed for binary builds, see: https://github.com/pytorch/pytorch/issues/73339#issuecomment-1058981560
-      - name: Enable long paths on Windows
-        shell: powershell
-        run: |
-          Set-ItemProperty -Path "HKLM:\\SYSTEM\CurrentControlSet\Control\FileSystem" -Name "LongPathsEnabled" -Value 1
-      # Since it's just a defensive command, the workflow should continue even the command fails. This step can be
-      # removed once Windows Defender is removed from the AMI
-      - name: Disables Windows Defender scheduled and real-time scanning for files in directories used by PyTorch
-        continue-on-error: true
-        shell: powershell
-        run: |
-          Add-MpPreference -ExclusionPath $(Get-Location).tostring(),$Env:TEMP -ErrorAction Ignore
-          # Let's both exclude the path and disable Windows Defender completely just to be sure
-          # that it doesn't interfere
-          Set-MpPreference -DisableRealtimeMonitoring $True -ErrorAction Ignore
-      - name: Checkout PyTorch
-        uses: actions/checkout@v4
-        with:
-          ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
-          submodules: recursive
-          path: pytorch
-          show-progress: false
-      - name: Clean PyTorch checkout
-        run: |
-          # Remove any artifacts from the previous checkouts
-          git clean -fxd
-        working-directory: pytorch
-      # NOTE: These environment variables are put here so that they can be applied on every job equally
-      #       They are also here because setting them at a workflow level doesn't give us access to the
-      #       runner.temp variable, which we need.
-      - name: Populate binary env
-        shell: bash
-        run: |
-          echo "BINARY_ENV_FILE=${RUNNER_TEMP}/env" >> "${GITHUB_ENV}"
-          echo "PYTORCH_FINAL_PACKAGE_DIR=${RUNNER_TEMP}/artifacts" >> "${GITHUB_ENV}"
-          echo "WIN_PACKAGE_WORK_DIR=${RUNNER_TEMP}"
-      - uses: actions/download-artifact@v4.1.7
-        name: Download Build Artifacts
-        with:
-          name: libtorch-cpu-shared-with-deps-release
-          path: "${{ env.PYTORCH_FINAL_PACKAGE_DIR }}"
-      - name: Populate binary env
-        shell: bash
-        run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_populate_env.sh"
-      - name: Test PyTorch binary
-        shell: bash
-        run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_windows_test.sh"
-      - name: Wait until all sessions have drained
-        shell: powershell
-        working-directory: pytorch
-        if: always()
-        timeout-minutes: 120
-        run: |
-          .github\scripts\wait_for_ssh_to_drain.ps1
-      - name: Kill active ssh sessions if still around (Useful if workflow was cancelled)
-        shell: powershell
-        working-directory: pytorch
-        if: always()
-        run: |
-          .github\scripts\kill_active_ssh_sessions.ps1
-  libtorch-cpu-shared-with-deps-release-upload:  # Uploading
-    if: ${{ github.repository_owner == 'pytorch' }}
-    permissions:
-      id-token: write
-      contents: read
-    needs: libtorch-cpu-shared-with-deps-release-test
-    with:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
-      PACKAGE_TYPE: libtorch
-      # TODO: This is a legacy variable that we eventually want to get rid of in
-      #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cpu
-      GPU_ARCH_TYPE: cpu
-      LIBTORCH_CONFIG: release
-      LIBTORCH_VARIANT: shared-with-deps
-      # This is a dummy value for libtorch to work correctly with our batch scripts
-      # without this value pip does not get installed for some reason
-      DESIRED_PYTHON: "3.10"
-      build_name: libtorch-cpu-shared-with-deps-release
-    secrets:
-      github-token: ${{ secrets.GITHUB_TOKEN }}
-    uses: ./.github/workflows/_binary-upload.yml
-  libtorch-cuda12_6-shared-with-deps-release-build:
-    if: ${{ github.repository_owner == 'pytorch' }}
-    needs: get-label-type
-    runs-on: "${{ needs.get-label-type.outputs.label-type }}windows.12xlarge"
-    timeout-minutes: 360
-    env:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
-      PACKAGE_TYPE: libtorch
-      # TODO: This is a legacy variable that we eventually want to get rid of in
-      #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu126
-      GPU_ARCH_VERSION: "12.6"
-      GPU_ARCH_TYPE: cuda
-      SKIP_ALL_TESTS: 1
-      LIBTORCH_CONFIG: release
-      LIBTORCH_VARIANT: shared-with-deps
-      # This is a dummy value for libtorch to work correctly with our batch scripts
-      # without this value pip does not get installed for some reason
-      DESIRED_PYTHON: "3.10"
-    steps:
-      # NOTE: These environment variables are put here so that they can be applied on every job equally
-      #       They are also here because setting them at a workflow level doesn't give us access to the
-      #       runner.temp variable, which we need.
-      - name: Populate binary env
-        shell: bash
-        run: |
-          echo "BINARY_ENV_FILE=${RUNNER_TEMP}/env" >> "${GITHUB_ENV}"
-          echo "PYTORCH_FINAL_PACKAGE_DIR=${RUNNER_TEMP}/artifacts" >> "${GITHUB_ENV}"
-          echo "WIN_PACKAGE_WORK_DIR=${RUNNER_TEMP}"
-      - name: Display EC2 information
-        shell: bash
-        run: |
-          set -euo pipefail
-          function get_ec2_metadata() {
-            # Pulled from instance metadata endpoint for EC2
-            # see https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/instancedata-data-retrieval.html
-            category=$1
-            curl -H "X-aws-ec2-metadata-token: $(curl -s -X PUT "http://169.254.169.254/latest/api/token" -H "X-aws-ec2-metadata-token-ttl-seconds: 30")" -fsSL "http://169.254.169.254/latest/meta-data/${category}"
-          }
-          echo "ami-id: $(get_ec2_metadata ami-id)"
-          echo "instance-id: $(get_ec2_metadata instance-id)"
-          echo "instance-type: $(get_ec2_metadata instance-type)"
-          echo "system info $(uname -a)"
-      - name: "[FB EMPLOYEES] Enable SSH (Click me for login details)"
-        uses: pytorch/test-infra/.github/actions/setup-ssh@main
-        continue-on-error: true
-        with:
-          github-secret: ${{ secrets.GITHUB_TOKEN }}
-      - name: Enable git long paths and symlinks on Windows and disable fsmonitor daemon
-        shell: bash
-        run: |
-          git config --global core.longpaths true
-          git config --global core.symlinks true
-
-          # https://git-scm.com/docs/git-fsmonitor--daemon.  The daemon could lock
-          # the directory on Windows and prevent GHA from checking out as reported
-          # in https://github.com/actions/checkout/issues/1018
-          git config --global core.fsmonitor false
-      # Needed for binary builds, see: https://github.com/pytorch/pytorch/issues/73339#issuecomment-1058981560
-      - name: Enable long paths on Windows
-        shell: powershell
-        run: |
-          Set-ItemProperty -Path "HKLM:\\SYSTEM\CurrentControlSet\Control\FileSystem" -Name "LongPathsEnabled" -Value 1
-      # Since it's just a defensive command, the workflow should continue even the command fails. This step can be
-      # removed once Windows Defender is removed from the AMI
-      - name: Disables Windows Defender scheduled and real-time scanning for files in directories used by PyTorch
-        continue-on-error: true
-        shell: powershell
-        run: |
-          Add-MpPreference -ExclusionPath $(Get-Location).tostring(),$Env:TEMP -ErrorAction Ignore
-          # Let's both exclude the path and disable Windows Defender completely just to be sure
-          # that it doesn't interfere
-          Set-MpPreference -DisableRealtimeMonitoring $True -ErrorAction Ignore
-      - name: Checkout PyTorch
-        uses: actions/checkout@v4
-        with:
-          ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
-          submodules: recursive
-          path: pytorch
-          show-progress: false
-      - name: Clean PyTorch checkout
-        run: |
-          # Remove any artifacts from the previous checkouts
-          git clean -fxd
-        working-directory: pytorch
-      - name: Populate binary env
-        shell: bash
-        run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_populate_env.sh"
-      - name: Build PyTorch binary
-        shell: bash
-        run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_windows_build.sh"
-      - uses: actions/upload-artifact@v4.4.0
-        if: always()
-        with:
-          name: libtorch-cuda12_6-shared-with-deps-release
-          retention-days: 14
-          if-no-files-found: error
-          path: "${{ env.PYTORCH_FINAL_PACKAGE_DIR }}"
-      - name: Wait until all sessions have drained
-        shell: powershell
-        working-directory: pytorch
-        if: always()
-        timeout-minutes: 120
-        run: |
-          .github\scripts\wait_for_ssh_to_drain.ps1
-      - name: Kill active ssh sessions if still around (Useful if workflow was cancelled)
-        shell: powershell
-        working-directory: pytorch
-        if: always()
-        run: |
-          .github\scripts\kill_active_ssh_sessions.ps1
-
-  libtorch-cuda12_6-shared-with-deps-release-test:  # Testing
-    if: ${{ github.repository_owner == 'pytorch' }}
-    needs:
-      - libtorch-cuda12_6-shared-with-deps-release-build
-      - get-label-type
-    runs-on: "${{ needs.get-label-type.outputs.label-type }}windows.g4dn.xlarge"
-    timeout-minutes: 360
-    env:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
-      PACKAGE_TYPE: libtorch
-      # TODO: This is a legacy variable that we eventually want to get rid of in
-      #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu126
-      GPU_ARCH_VERSION: "12.6"
-      GPU_ARCH_TYPE: cuda
-      SKIP_ALL_TESTS: 1
-      LIBTORCH_CONFIG: release
-      LIBTORCH_VARIANT: shared-with-deps
-      # This is a dummy value for libtorch to work correctly with our batch scripts
-      # without this value pip does not get installed for some reason
-      DESIRED_PYTHON: "3.10"
-    steps:
-      - name: Display EC2 information
-        shell: bash
-        run: |
-          set -euo pipefail
-          function get_ec2_metadata() {
-            # Pulled from instance metadata endpoint for EC2
-            # see https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/instancedata-data-retrieval.html
-            category=$1
-            curl -H "X-aws-ec2-metadata-token: $(curl -s -X PUT "http://169.254.169.254/latest/api/token" -H "X-aws-ec2-metadata-token-ttl-seconds: 30")" -fsSL "http://169.254.169.254/latest/meta-data/${category}"
-          }
-          echo "ami-id: $(get_ec2_metadata ami-id)"
-          echo "instance-id: $(get_ec2_metadata instance-id)"
-          echo "instance-type: $(get_ec2_metadata instance-type)"
-          echo "system info $(uname -a)"
-      - name: "[FB EMPLOYEES] Enable SSH (Click me for login details)"
-        uses: pytorch/test-infra/.github/actions/setup-ssh@main
-        continue-on-error: true
-        with:
-          github-secret: ${{ secrets.GITHUB_TOKEN }}
-      - name: Enable git long paths and symlinks on Windows and disable fsmonitor daemon
-        shell: bash
-        run: |
-          git config --global core.longpaths true
-          git config --global core.symlinks true
-
-          # https://git-scm.com/docs/git-fsmonitor--daemon.  The daemon could lock
-          # the directory on Windows and prevent GHA from checking out as reported
-          # in https://github.com/actions/checkout/issues/1018
-          git config --global core.fsmonitor false
-      # Needed for binary builds, see: https://github.com/pytorch/pytorch/issues/73339#issuecomment-1058981560
-      - name: Enable long paths on Windows
-        shell: powershell
-        run: |
-          Set-ItemProperty -Path "HKLM:\\SYSTEM\CurrentControlSet\Control\FileSystem" -Name "LongPathsEnabled" -Value 1
-      # Since it's just a defensive command, the workflow should continue even the command fails. This step can be
-      # removed once Windows Defender is removed from the AMI
-      - name: Disables Windows Defender scheduled and real-time scanning for files in directories used by PyTorch
-        continue-on-error: true
-        shell: powershell
-        run: |
-          Add-MpPreference -ExclusionPath $(Get-Location).tostring(),$Env:TEMP -ErrorAction Ignore
-          # Let's both exclude the path and disable Windows Defender completely just to be sure
-          # that it doesn't interfere
-          Set-MpPreference -DisableRealtimeMonitoring $True -ErrorAction Ignore
-      - name: Checkout PyTorch
-        uses: actions/checkout@v4
-        with:
-          ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
-          submodules: recursive
-          path: pytorch
-          show-progress: false
-      - name: Clean PyTorch checkout
-        run: |
-          # Remove any artifacts from the previous checkouts
-          git clean -fxd
-        working-directory: pytorch
-      # NOTE: These environment variables are put here so that they can be applied on every job equally
-      #       They are also here because setting them at a workflow level doesn't give us access to the
-      #       runner.temp variable, which we need.
-      - name: Populate binary env
-        shell: bash
-        run: |
-          echo "BINARY_ENV_FILE=${RUNNER_TEMP}/env" >> "${GITHUB_ENV}"
-          echo "PYTORCH_FINAL_PACKAGE_DIR=${RUNNER_TEMP}/artifacts" >> "${GITHUB_ENV}"
-          echo "WIN_PACKAGE_WORK_DIR=${RUNNER_TEMP}"
-      - uses: actions/download-artifact@v4.1.7
-        name: Download Build Artifacts
-        with:
-          name: libtorch-cuda12_6-shared-with-deps-release
-          path: "${{ env.PYTORCH_FINAL_PACKAGE_DIR }}"
-      - name: Populate binary env
-        shell: bash
-        run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_populate_env.sh"
-      - name: Test PyTorch binary
-        shell: bash
-        run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_windows_test.sh"
-      - name: Wait until all sessions have drained
-        shell: powershell
-        working-directory: pytorch
-        if: always()
-        timeout-minutes: 120
-        run: |
-          .github\scripts\wait_for_ssh_to_drain.ps1
-      - name: Kill active ssh sessions if still around (Useful if workflow was cancelled)
-        shell: powershell
-        working-directory: pytorch
-        if: always()
-        run: |
-          .github\scripts\kill_active_ssh_sessions.ps1
-  libtorch-cuda12_6-shared-with-deps-release-upload:  # Uploading
-    if: ${{ github.repository_owner == 'pytorch' }}
-    permissions:
-      id-token: write
-      contents: read
-    needs: libtorch-cuda12_6-shared-with-deps-release-test
-    with:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
-      PACKAGE_TYPE: libtorch
-      # TODO: This is a legacy variable that we eventually want to get rid of in
-      #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu126
-      GPU_ARCH_VERSION: "12.6"
-      GPU_ARCH_TYPE: cuda
-      LIBTORCH_CONFIG: release
-      LIBTORCH_VARIANT: shared-with-deps
-      # This is a dummy value for libtorch to work correctly with our batch scripts
-      # without this value pip does not get installed for some reason
-      DESIRED_PYTHON: "3.10"
-      build_name: libtorch-cuda12_6-shared-with-deps-release
-    secrets:
-      github-token: ${{ secrets.GITHUB_TOKEN }}
-    uses: ./.github/workflows/_binary-upload.yml
-  libtorch-cuda12_8-shared-with-deps-release-build:
-    if: ${{ github.repository_owner == 'pytorch' }}
-    needs: get-label-type
-    runs-on: "${{ needs.get-label-type.outputs.label-type }}windows.12xlarge"
-    timeout-minutes: 360
-    env:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
-      PACKAGE_TYPE: libtorch
-      # TODO: This is a legacy variable that we eventually want to get rid of in
-      #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu128
-      GPU_ARCH_VERSION: "12.8"
-      GPU_ARCH_TYPE: cuda
-      SKIP_ALL_TESTS: 1
-      LIBTORCH_CONFIG: release
-      LIBTORCH_VARIANT: shared-with-deps
-      # This is a dummy value for libtorch to work correctly with our batch scripts
-      # without this value pip does not get installed for some reason
-      DESIRED_PYTHON: "3.10"
-    steps:
-      # NOTE: These environment variables are put here so that they can be applied on every job equally
-      #       They are also here because setting them at a workflow level doesn't give us access to the
-      #       runner.temp variable, which we need.
-      - name: Populate binary env
-        shell: bash
-        run: |
-          echo "BINARY_ENV_FILE=${RUNNER_TEMP}/env" >> "${GITHUB_ENV}"
-          echo "PYTORCH_FINAL_PACKAGE_DIR=${RUNNER_TEMP}/artifacts" >> "${GITHUB_ENV}"
-          echo "WIN_PACKAGE_WORK_DIR=${RUNNER_TEMP}"
-      - name: Display EC2 information
-        shell: bash
-        run: |
-          set -euo pipefail
-          function get_ec2_metadata() {
-            # Pulled from instance metadata endpoint for EC2
-            # see https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/instancedata-data-retrieval.html
-            category=$1
-            curl -H "X-aws-ec2-metadata-token: $(curl -s -X PUT "http://169.254.169.254/latest/api/token" -H "X-aws-ec2-metadata-token-ttl-seconds: 30")" -fsSL "http://169.254.169.254/latest/meta-data/${category}"
-          }
-          echo "ami-id: $(get_ec2_metadata ami-id)"
-          echo "instance-id: $(get_ec2_metadata instance-id)"
-          echo "instance-type: $(get_ec2_metadata instance-type)"
-          echo "system info $(uname -a)"
-      - name: "[FB EMPLOYEES] Enable SSH (Click me for login details)"
-        uses: pytorch/test-infra/.github/actions/setup-ssh@main
-        continue-on-error: true
-        with:
-          github-secret: ${{ secrets.GITHUB_TOKEN }}
-      - name: Enable git long paths and symlinks on Windows and disable fsmonitor daemon
-        shell: bash
-        run: |
-          git config --global core.longpaths true
-          git config --global core.symlinks true
-
-          # https://git-scm.com/docs/git-fsmonitor--daemon.  The daemon could lock
-          # the directory on Windows and prevent GHA from checking out as reported
-          # in https://github.com/actions/checkout/issues/1018
-          git config --global core.fsmonitor false
-      # Needed for binary builds, see: https://github.com/pytorch/pytorch/issues/73339#issuecomment-1058981560
-      - name: Enable long paths on Windows
-        shell: powershell
-        run: |
-          Set-ItemProperty -Path "HKLM:\\SYSTEM\CurrentControlSet\Control\FileSystem" -Name "LongPathsEnabled" -Value 1
-      # Since it's just a defensive command, the workflow should continue even the command fails. This step can be
-      # removed once Windows Defender is removed from the AMI
-      - name: Disables Windows Defender scheduled and real-time scanning for files in directories used by PyTorch
-        continue-on-error: true
-        shell: powershell
-        run: |
-          Add-MpPreference -ExclusionPath $(Get-Location).tostring(),$Env:TEMP -ErrorAction Ignore
-          # Let's both exclude the path and disable Windows Defender completely just to be sure
-          # that it doesn't interfere
-          Set-MpPreference -DisableRealtimeMonitoring $True -ErrorAction Ignore
-      - name: Checkout PyTorch
-        uses: actions/checkout@v4
-        with:
-          ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
-          submodules: recursive
-          path: pytorch
-          show-progress: false
-      - name: Clean PyTorch checkout
-        run: |
-          # Remove any artifacts from the previous checkouts
-          git clean -fxd
-        working-directory: pytorch
-      - name: Populate binary env
-        shell: bash
-        run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_populate_env.sh"
-      - name: Build PyTorch binary
-        shell: bash
-        run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_windows_build.sh"
-      - uses: actions/upload-artifact@v4.4.0
-        if: always()
-        with:
-          name: libtorch-cuda12_8-shared-with-deps-release
-          retention-days: 14
-          if-no-files-found: error
-          path: "${{ env.PYTORCH_FINAL_PACKAGE_DIR }}"
-      - name: Wait until all sessions have drained
-        shell: powershell
-        working-directory: pytorch
-        if: always()
-        timeout-minutes: 120
-        run: |
-          .github\scripts\wait_for_ssh_to_drain.ps1
-      - name: Kill active ssh sessions if still around (Useful if workflow was cancelled)
-        shell: powershell
-        working-directory: pytorch
-        if: always()
-        run: |
-          .github\scripts\kill_active_ssh_sessions.ps1
-
-  libtorch-cuda12_8-shared-with-deps-release-test:  # Testing
-    if: ${{ github.repository_owner == 'pytorch' }}
-    needs:
-      - libtorch-cuda12_8-shared-with-deps-release-build
-      - get-label-type
-    runs-on: "${{ needs.get-label-type.outputs.label-type }}windows.g4dn.xlarge"
-    timeout-minutes: 360
-    env:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
-      PACKAGE_TYPE: libtorch
-      # TODO: This is a legacy variable that we eventually want to get rid of in
-      #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu128
-      GPU_ARCH_VERSION: "12.8"
-      GPU_ARCH_TYPE: cuda
-      SKIP_ALL_TESTS: 1
-      LIBTORCH_CONFIG: release
-      LIBTORCH_VARIANT: shared-with-deps
-      # This is a dummy value for libtorch to work correctly with our batch scripts
-      # without this value pip does not get installed for some reason
-      DESIRED_PYTHON: "3.10"
-    steps:
-      - name: Display EC2 information
-        shell: bash
-        run: |
-          set -euo pipefail
-          function get_ec2_metadata() {
-            # Pulled from instance metadata endpoint for EC2
-            # see https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/instancedata-data-retrieval.html
-            category=$1
-            curl -H "X-aws-ec2-metadata-token: $(curl -s -X PUT "http://169.254.169.254/latest/api/token" -H "X-aws-ec2-metadata-token-ttl-seconds: 30")" -fsSL "http://169.254.169.254/latest/meta-data/${category}"
-          }
-          echo "ami-id: $(get_ec2_metadata ami-id)"
-          echo "instance-id: $(get_ec2_metadata instance-id)"
-          echo "instance-type: $(get_ec2_metadata instance-type)"
-          echo "system info $(uname -a)"
-      - name: "[FB EMPLOYEES] Enable SSH (Click me for login details)"
-        uses: pytorch/test-infra/.github/actions/setup-ssh@main
-        continue-on-error: true
-        with:
-          github-secret: ${{ secrets.GITHUB_TOKEN }}
-      - name: Enable git long paths and symlinks on Windows and disable fsmonitor daemon
-        shell: bash
-        run: |
-          git config --global core.longpaths true
-          git config --global core.symlinks true
-
-          # https://git-scm.com/docs/git-fsmonitor--daemon.  The daemon could lock
-          # the directory on Windows and prevent GHA from checking out as reported
-          # in https://github.com/actions/checkout/issues/1018
-          git config --global core.fsmonitor false
-      # Needed for binary builds, see: https://github.com/pytorch/pytorch/issues/73339#issuecomment-1058981560
-      - name: Enable long paths on Windows
-        shell: powershell
-        run: |
-          Set-ItemProperty -Path "HKLM:\\SYSTEM\CurrentControlSet\Control\FileSystem" -Name "LongPathsEnabled" -Value 1
-      # Since it's just a defensive command, the workflow should continue even the command fails. This step can be
-      # removed once Windows Defender is removed from the AMI
-      - name: Disables Windows Defender scheduled and real-time scanning for files in directories used by PyTorch
-        continue-on-error: true
-        shell: powershell
-        run: |
-          Add-MpPreference -ExclusionPath $(Get-Location).tostring(),$Env:TEMP -ErrorAction Ignore
-          # Let's both exclude the path and disable Windows Defender completely just to be sure
-          # that it doesn't interfere
-          Set-MpPreference -DisableRealtimeMonitoring $True -ErrorAction Ignore
-      - name: Checkout PyTorch
-        uses: actions/checkout@v4
-        with:
-          ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
-          submodules: recursive
-          path: pytorch
-          show-progress: false
-      - name: Clean PyTorch checkout
-        run: |
-          # Remove any artifacts from the previous checkouts
-          git clean -fxd
-        working-directory: pytorch
-      # NOTE: These environment variables are put here so that they can be applied on every job equally
-      #       They are also here because setting them at a workflow level doesn't give us access to the
-      #       runner.temp variable, which we need.
-      - name: Populate binary env
-        shell: bash
-        run: |
-          echo "BINARY_ENV_FILE=${RUNNER_TEMP}/env" >> "${GITHUB_ENV}"
-          echo "PYTORCH_FINAL_PACKAGE_DIR=${RUNNER_TEMP}/artifacts" >> "${GITHUB_ENV}"
-          echo "WIN_PACKAGE_WORK_DIR=${RUNNER_TEMP}"
-      - uses: actions/download-artifact@v4.1.7
-        name: Download Build Artifacts
-        with:
-          name: libtorch-cuda12_8-shared-with-deps-release
-          path: "${{ env.PYTORCH_FINAL_PACKAGE_DIR }}"
-      - name: Populate binary env
-        shell: bash
-        run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_populate_env.sh"
-      - name: Test PyTorch binary
-        shell: bash
-        run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_windows_test.sh"
-      - name: Wait until all sessions have drained
-        shell: powershell
-        working-directory: pytorch
-        if: always()
-        timeout-minutes: 120
-        run: |
-          .github\scripts\wait_for_ssh_to_drain.ps1
-      - name: Kill active ssh sessions if still around (Useful if workflow was cancelled)
-        shell: powershell
-        working-directory: pytorch
-        if: always()
-        run: |
-          .github\scripts\kill_active_ssh_sessions.ps1
-  libtorch-cuda12_8-shared-with-deps-release-upload:  # Uploading
-    if: ${{ github.repository_owner == 'pytorch' }}
-    permissions:
-      id-token: write
-      contents: read
-    needs: libtorch-cuda12_8-shared-with-deps-release-test
-    with:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
-      PACKAGE_TYPE: libtorch
-      # TODO: This is a legacy variable that we eventually want to get rid of in
-      #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu128
-      GPU_ARCH_VERSION: "12.8"
-      GPU_ARCH_TYPE: cuda
-      LIBTORCH_CONFIG: release
-      LIBTORCH_VARIANT: shared-with-deps
-      # This is a dummy value for libtorch to work correctly with our batch scripts
-      # without this value pip does not get installed for some reason
-      DESIRED_PYTHON: "3.10"
-      build_name: libtorch-cuda12_8-shared-with-deps-release
-    secrets:
-      github-token: ${{ secrets.GITHUB_TOKEN }}
-    uses: ./.github/workflows/_binary-upload.yml
-  libtorch-cuda13_0-shared-with-deps-release-build:
-    if: ${{ github.repository_owner == 'pytorch' }}
-    needs: get-label-type
-    runs-on: "${{ needs.get-label-type.outputs.label-type }}windows.12xlarge"
-    timeout-minutes: 360
-    env:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
-      PACKAGE_TYPE: libtorch
-      # TODO: This is a legacy variable that we eventually want to get rid of in
-      #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu130
-      GPU_ARCH_VERSION: "13.0"
-      GPU_ARCH_TYPE: cuda
-      SKIP_ALL_TESTS: 1
-      LIBTORCH_CONFIG: release
-      LIBTORCH_VARIANT: shared-with-deps
-      # This is a dummy value for libtorch to work correctly with our batch scripts
-      # without this value pip does not get installed for some reason
-      DESIRED_PYTHON: "3.10"
-    steps:
-      # NOTE: These environment variables are put here so that they can be applied on every job equally
-      #       They are also here because setting them at a workflow level doesn't give us access to the
-      #       runner.temp variable, which we need.
-      - name: Populate binary env
-        shell: bash
-        run: |
-          echo "BINARY_ENV_FILE=${RUNNER_TEMP}/env" >> "${GITHUB_ENV}"
-          echo "PYTORCH_FINAL_PACKAGE_DIR=${RUNNER_TEMP}/artifacts" >> "${GITHUB_ENV}"
-          echo "WIN_PACKAGE_WORK_DIR=${RUNNER_TEMP}"
-      - name: Display EC2 information
-        shell: bash
-        run: |
-          set -euo pipefail
-          function get_ec2_metadata() {
-            # Pulled from instance metadata endpoint for EC2
-            # see https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/instancedata-data-retrieval.html
-            category=$1
-            curl -H "X-aws-ec2-metadata-token: $(curl -s -X PUT "http://169.254.169.254/latest/api/token" -H "X-aws-ec2-metadata-token-ttl-seconds: 30")" -fsSL "http://169.254.169.254/latest/meta-data/${category}"
-          }
-          echo "ami-id: $(get_ec2_metadata ami-id)"
-          echo "instance-id: $(get_ec2_metadata instance-id)"
-          echo "instance-type: $(get_ec2_metadata instance-type)"
-          echo "system info $(uname -a)"
-      - name: "[FB EMPLOYEES] Enable SSH (Click me for login details)"
-        uses: pytorch/test-infra/.github/actions/setup-ssh@main
-        continue-on-error: true
-        with:
-          github-secret: ${{ secrets.GITHUB_TOKEN }}
-      - name: Enable git long paths and symlinks on Windows and disable fsmonitor daemon
-        shell: bash
-        run: |
-          git config --global core.longpaths true
-          git config --global core.symlinks true
-
-          # https://git-scm.com/docs/git-fsmonitor--daemon.  The daemon could lock
-          # the directory on Windows and prevent GHA from checking out as reported
-          # in https://github.com/actions/checkout/issues/1018
-          git config --global core.fsmonitor false
-      # Needed for binary builds, see: https://github.com/pytorch/pytorch/issues/73339#issuecomment-1058981560
-      - name: Enable long paths on Windows
-        shell: powershell
-        run: |
-          Set-ItemProperty -Path "HKLM:\\SYSTEM\CurrentControlSet\Control\FileSystem" -Name "LongPathsEnabled" -Value 1
-      # Since it's just a defensive command, the workflow should continue even the command fails. This step can be
-      # removed once Windows Defender is removed from the AMI
-      - name: Disables Windows Defender scheduled and real-time scanning for files in directories used by PyTorch
-        continue-on-error: true
-        shell: powershell
-        run: |
-          Add-MpPreference -ExclusionPath $(Get-Location).tostring(),$Env:TEMP -ErrorAction Ignore
-          # Let's both exclude the path and disable Windows Defender completely just to be sure
-          # that it doesn't interfere
-          Set-MpPreference -DisableRealtimeMonitoring $True -ErrorAction Ignore
-      - name: Checkout PyTorch
-        uses: actions/checkout@v4
-        with:
-          ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
-          submodules: recursive
-          path: pytorch
-          show-progress: false
-      - name: Clean PyTorch checkout
-        run: |
-          # Remove any artifacts from the previous checkouts
-          git clean -fxd
-        working-directory: pytorch
-      - name: Populate binary env
-        shell: bash
-        run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_populate_env.sh"
-      - name: Build PyTorch binary
-        shell: bash
-        run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_windows_build.sh"
-      - uses: actions/upload-artifact@v4.4.0
-        if: always()
-        with:
-          name: libtorch-cuda13_0-shared-with-deps-release
-          retention-days: 14
-          if-no-files-found: error
-          path: "${{ env.PYTORCH_FINAL_PACKAGE_DIR }}"
-      - name: Wait until all sessions have drained
-        shell: powershell
-        working-directory: pytorch
-        if: always()
-        timeout-minutes: 120
-        run: |
-          .github\scripts\wait_for_ssh_to_drain.ps1
-      - name: Kill active ssh sessions if still around (Useful if workflow was cancelled)
-        shell: powershell
-        working-directory: pytorch
-        if: always()
-        run: |
-          .github\scripts\kill_active_ssh_sessions.ps1
-
-  libtorch-cuda13_0-shared-with-deps-release-test:  # Testing
-    if: ${{ github.repository_owner == 'pytorch' }}
-    needs:
-      - libtorch-cuda13_0-shared-with-deps-release-build
-      - get-label-type
-    runs-on: "${{ needs.get-label-type.outputs.label-type }}windows.g4dn.xlarge"
-    timeout-minutes: 360
-    env:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
-      PACKAGE_TYPE: libtorch
-      # TODO: This is a legacy variable that we eventually want to get rid of in
-      #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu130
-      GPU_ARCH_VERSION: "13.0"
-      GPU_ARCH_TYPE: cuda
-      SKIP_ALL_TESTS: 1
-      LIBTORCH_CONFIG: release
-      LIBTORCH_VARIANT: shared-with-deps
-      # This is a dummy value for libtorch to work correctly with our batch scripts
-      # without this value pip does not get installed for some reason
-      DESIRED_PYTHON: "3.10"
-    steps:
-      - name: Display EC2 information
-        shell: bash
-        run: |
-          set -euo pipefail
-          function get_ec2_metadata() {
-            # Pulled from instance metadata endpoint for EC2
-            # see https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/instancedata-data-retrieval.html
-            category=$1
-            curl -H "X-aws-ec2-metadata-token: $(curl -s -X PUT "http://169.254.169.254/latest/api/token" -H "X-aws-ec2-metadata-token-ttl-seconds: 30")" -fsSL "http://169.254.169.254/latest/meta-data/${category}"
-          }
-          echo "ami-id: $(get_ec2_metadata ami-id)"
-          echo "instance-id: $(get_ec2_metadata instance-id)"
-          echo "instance-type: $(get_ec2_metadata instance-type)"
-          echo "system info $(uname -a)"
-      - name: "[FB EMPLOYEES] Enable SSH (Click me for login details)"
-        uses: pytorch/test-infra/.github/actions/setup-ssh@main
-        continue-on-error: true
-        with:
-          github-secret: ${{ secrets.GITHUB_TOKEN }}
-      - name: Enable git long paths and symlinks on Windows and disable fsmonitor daemon
-        shell: bash
-        run: |
-          git config --global core.longpaths true
-          git config --global core.symlinks true
-
-          # https://git-scm.com/docs/git-fsmonitor--daemon.  The daemon could lock
-          # the directory on Windows and prevent GHA from checking out as reported
-          # in https://github.com/actions/checkout/issues/1018
-          git config --global core.fsmonitor false
-      # Needed for binary builds, see: https://github.com/pytorch/pytorch/issues/73339#issuecomment-1058981560
-      - name: Enable long paths on Windows
-        shell: powershell
-        run: |
-          Set-ItemProperty -Path "HKLM:\\SYSTEM\CurrentControlSet\Control\FileSystem" -Name "LongPathsEnabled" -Value 1
-      # Since it's just a defensive command, the workflow should continue even the command fails. This step can be
-      # removed once Windows Defender is removed from the AMI
-      - name: Disables Windows Defender scheduled and real-time scanning for files in directories used by PyTorch
-        continue-on-error: true
-        shell: powershell
-        run: |
-          Add-MpPreference -ExclusionPath $(Get-Location).tostring(),$Env:TEMP -ErrorAction Ignore
-          # Let's both exclude the path and disable Windows Defender completely just to be sure
-          # that it doesn't interfere
-          Set-MpPreference -DisableRealtimeMonitoring $True -ErrorAction Ignore
-      - name: Checkout PyTorch
-        uses: actions/checkout@v4
-        with:
-          ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
-          submodules: recursive
-          path: pytorch
-          show-progress: false
-      - name: Clean PyTorch checkout
-        run: |
-          # Remove any artifacts from the previous checkouts
-          git clean -fxd
-        working-directory: pytorch
-      # NOTE: These environment variables are put here so that they can be applied on every job equally
-      #       They are also here because setting them at a workflow level doesn't give us access to the
-      #       runner.temp variable, which we need.
-      - name: Populate binary env
-        shell: bash
-        run: |
-          echo "BINARY_ENV_FILE=${RUNNER_TEMP}/env" >> "${GITHUB_ENV}"
-          echo "PYTORCH_FINAL_PACKAGE_DIR=${RUNNER_TEMP}/artifacts" >> "${GITHUB_ENV}"
-          echo "WIN_PACKAGE_WORK_DIR=${RUNNER_TEMP}"
-      - uses: actions/download-artifact@v4.1.7
-        name: Download Build Artifacts
-        with:
-          name: libtorch-cuda13_0-shared-with-deps-release
-          path: "${{ env.PYTORCH_FINAL_PACKAGE_DIR }}"
-      - name: Populate binary env
-        shell: bash
-        run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_populate_env.sh"
-      - name: Test PyTorch binary
-        shell: bash
-        run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_windows_test.sh"
-      - name: Wait until all sessions have drained
-        shell: powershell
-        working-directory: pytorch
-        if: always()
-        timeout-minutes: 120
-        run: |
-          .github\scripts\wait_for_ssh_to_drain.ps1
-      - name: Kill active ssh sessions if still around (Useful if workflow was cancelled)
-        shell: powershell
-        working-directory: pytorch
-        if: always()
-        run: |
-          .github\scripts\kill_active_ssh_sessions.ps1
-  libtorch-cuda13_0-shared-with-deps-release-upload:  # Uploading
-    if: ${{ github.repository_owner == 'pytorch' }}
-    permissions:
-      id-token: write
-      contents: read
-    needs: libtorch-cuda13_0-shared-with-deps-release-test
-    with:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
-      PACKAGE_TYPE: libtorch
-      # TODO: This is a legacy variable that we eventually want to get rid of in
-      #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu130
-      GPU_ARCH_VERSION: "13.0"
-      GPU_ARCH_TYPE: cuda
-      LIBTORCH_CONFIG: release
-      LIBTORCH_VARIANT: shared-with-deps
-      # This is a dummy value for libtorch to work correctly with our batch scripts
-      # without this value pip does not get installed for some reason
-      DESIRED_PYTHON: "3.10"
-      build_name: libtorch-cuda13_0-shared-with-deps-release
-    secrets:
-      github-token: ${{ secrets.GITHUB_TOKEN }}
-    uses: ./.github/workflows/_binary-upload.yml
diff --git a/.github/workflows/generated-windows-binary-wheel-nightly.yml b/.github/workflows/generated-windows-binary-wheel-nightly.yml
index e23118631d3a9..1aeec381ec19b 100644
--- a/.github/workflows/generated-windows-binary-wheel-nightly.yml
+++ b/.github/workflows/generated-windows-binary-wheel-nightly.yml
@@ -14,6 +14,7 @@ on:
       # Release candidate tags look like: v1.11.0-rc1
       - v[0-9]+.[0-9]+.[0-9]+-rc[0-9]+
       - 'ciflow/binaries/*'
+      - 'ciflow/binaries_libtorch/*'
       - 'ciflow/binaries_wheel/*'
   workflow_dispatch:
 
@@ -35,7 +36,7 @@ jobs:
   get-label-type:
     if: github.repository_owner == 'pytorch'
     name: get-label-type
-    uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@main
+    uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@release/2.12
     with:
       triggering_actor: ${{ github.triggering_actor }}
       issue_owner: ${{ github.event.pull_request.user.login || github.event.issue.user.login }}
@@ -47,7 +48,7 @@ jobs:
     runs-on: "${{ needs.get-label-type.outputs.label-type }}windows.12xlarge"
     timeout-minutes: 360
     env:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
+      PYTORCH_ROOT: ${{ github.workspace }}
       PACKAGE_TYPE: wheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
@@ -55,6 +56,7 @@ jobs:
       GPU_ARCH_TYPE: cpu
       SKIP_ALL_TESTS: 1
       DESIRED_PYTHON: "3.10"
+      PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cufile,nvjitlink,nvtx]==13.0.2; platform_system == 'Linux' | nvidia-cublas>=13.1.0.3,<=13.1.1.3; platform_system == 'Linux' | cuda-bindings>=13.0.3,<14; platform_system == 'Linux' | nvidia-cudnn-cu13==9.20.0.48; platform_system == 'Linux' | nvidia-cusparselt-cu13==0.8.1; platform_system == 'Linux' | nvidia-nccl-cu13==2.29.7; platform_system == 'Linux' | nvidia-nvshmem-cu13==3.4.5; platform_system == 'Linux'
     steps:
       # NOTE: These environment variables are put here so that they can be applied on every job equally
       #       They are also here because setting them at a workflow level doesn't give us access to the
@@ -80,7 +82,7 @@ jobs:
           echo "instance-type: $(get_ec2_metadata instance-type)"
           echo "system info $(uname -a)"
       - name: "[FB EMPLOYEES] Enable SSH (Click me for login details)"
-        uses: pytorch/test-infra/.github/actions/setup-ssh@main
+        uses: pytorch/test-infra/.github/actions/setup-ssh@release/2.12
         continue-on-error: true
         with:
           github-secret: ${{ secrets.GITHUB_TOKEN }}
@@ -112,23 +114,20 @@ jobs:
       - name: Checkout PyTorch
         uses: actions/checkout@v4
         with:
-          ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
           submodules: recursive
-          path: pytorch
           show-progress: false
       - name: Clean PyTorch checkout
         run: |
           # Remove any artifacts from the previous checkouts
           git clean -fxd
-        working-directory: pytorch
       - name: Populate binary env
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_populate_env.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_populate_env.sh"
       - name: Build PyTorch binary
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_windows_build.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_windows_build.sh"
       - uses: actions/upload-artifact@v4.4.0
         if: always()
         with:
@@ -138,14 +137,12 @@ jobs:
           path: "${{ env.PYTORCH_FINAL_PACKAGE_DIR }}"
       - name: Wait until all sessions have drained
         shell: powershell
-        working-directory: pytorch
         if: always()
         timeout-minutes: 120
         run: |
           .github\scripts\wait_for_ssh_to_drain.ps1
       - name: Kill active ssh sessions if still around (Useful if workflow was cancelled)
         shell: powershell
-        working-directory: pytorch
         if: always()
         run: |
           .github\scripts\kill_active_ssh_sessions.ps1
@@ -158,7 +155,7 @@ jobs:
     runs-on: "${{ needs.get-label-type.outputs.label-type }}windows.4xlarge"
     timeout-minutes: 360
     env:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
+      PYTORCH_ROOT: ${{ github.workspace }}
       PACKAGE_TYPE: wheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
@@ -182,7 +179,7 @@ jobs:
           echo "instance-type: $(get_ec2_metadata instance-type)"
           echo "system info $(uname -a)"
       - name: "[FB EMPLOYEES] Enable SSH (Click me for login details)"
-        uses: pytorch/test-infra/.github/actions/setup-ssh@main
+        uses: pytorch/test-infra/.github/actions/setup-ssh@release/2.12
         continue-on-error: true
         with:
           github-secret: ${{ secrets.GITHUB_TOKEN }}
@@ -214,15 +211,12 @@ jobs:
       - name: Checkout PyTorch
         uses: actions/checkout@v4
         with:
-          ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
           submodules: recursive
-          path: pytorch
           show-progress: false
       - name: Clean PyTorch checkout
         run: |
           # Remove any artifacts from the previous checkouts
           git clean -fxd
-        working-directory: pytorch
       # NOTE: These environment variables are put here so that they can be applied on every job equally
       #       They are also here because setting them at a workflow level doesn't give us access to the
       #       runner.temp variable, which we need.
@@ -240,21 +234,19 @@ jobs:
       - name: Populate binary env
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_populate_env.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_populate_env.sh"
       - name: Test PyTorch binary
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_windows_test.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_windows_test.sh"
       - name: Wait until all sessions have drained
         shell: powershell
-        working-directory: pytorch
         if: always()
         timeout-minutes: 120
         run: |
           .github\scripts\wait_for_ssh_to_drain.ps1
       - name: Kill active ssh sessions if still around (Useful if workflow was cancelled)
         shell: powershell
-        working-directory: pytorch
         if: always()
         run: |
           .github\scripts\kill_active_ssh_sessions.ps1
@@ -265,7 +257,7 @@ jobs:
       contents: read
     needs: wheel-py3_10-cpu-test
     with:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
+      PYTORCH_ROOT: ${{ github.workspace }}
       PACKAGE_TYPE: wheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
@@ -275,6 +267,9 @@ jobs:
       build_name: wheel-py3_10-cpu
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
     uses: ./.github/workflows/_binary-upload.yml
   wheel-py3_10-cuda12_6-build:
     if: ${{ github.repository_owner == 'pytorch' }}
@@ -282,7 +277,7 @@ jobs:
     runs-on: "${{ needs.get-label-type.outputs.label-type }}windows.12xlarge"
     timeout-minutes: 360
     env:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
+      PYTORCH_ROOT: ${{ github.workspace }}
       PACKAGE_TYPE: wheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
@@ -316,7 +311,7 @@ jobs:
           echo "instance-type: $(get_ec2_metadata instance-type)"
           echo "system info $(uname -a)"
       - name: "[FB EMPLOYEES] Enable SSH (Click me for login details)"
-        uses: pytorch/test-infra/.github/actions/setup-ssh@main
+        uses: pytorch/test-infra/.github/actions/setup-ssh@release/2.12
         continue-on-error: true
         with:
           github-secret: ${{ secrets.GITHUB_TOKEN }}
@@ -348,23 +343,20 @@ jobs:
       - name: Checkout PyTorch
         uses: actions/checkout@v4
         with:
-          ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
           submodules: recursive
-          path: pytorch
           show-progress: false
       - name: Clean PyTorch checkout
         run: |
           # Remove any artifacts from the previous checkouts
           git clean -fxd
-        working-directory: pytorch
       - name: Populate binary env
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_populate_env.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_populate_env.sh"
       - name: Build PyTorch binary
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_windows_build.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_windows_build.sh"
       - uses: actions/upload-artifact@v4.4.0
         if: always()
         with:
@@ -374,14 +366,12 @@ jobs:
           path: "${{ env.PYTORCH_FINAL_PACKAGE_DIR }}"
       - name: Wait until all sessions have drained
         shell: powershell
-        working-directory: pytorch
         if: always()
         timeout-minutes: 120
         run: |
           .github\scripts\wait_for_ssh_to_drain.ps1
       - name: Kill active ssh sessions if still around (Useful if workflow was cancelled)
         shell: powershell
-        working-directory: pytorch
         if: always()
         run: |
           .github\scripts\kill_active_ssh_sessions.ps1
@@ -394,7 +384,7 @@ jobs:
     runs-on: "${{ needs.get-label-type.outputs.label-type }}windows.g4dn.xlarge"
     timeout-minutes: 360
     env:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
+      PYTORCH_ROOT: ${{ github.workspace }}
       PACKAGE_TYPE: wheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
@@ -419,7 +409,7 @@ jobs:
           echo "instance-type: $(get_ec2_metadata instance-type)"
           echo "system info $(uname -a)"
       - name: "[FB EMPLOYEES] Enable SSH (Click me for login details)"
-        uses: pytorch/test-infra/.github/actions/setup-ssh@main
+        uses: pytorch/test-infra/.github/actions/setup-ssh@release/2.12
         continue-on-error: true
         with:
           github-secret: ${{ secrets.GITHUB_TOKEN }}
@@ -451,15 +441,12 @@ jobs:
       - name: Checkout PyTorch
         uses: actions/checkout@v4
         with:
-          ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
           submodules: recursive
-          path: pytorch
           show-progress: false
       - name: Clean PyTorch checkout
         run: |
           # Remove any artifacts from the previous checkouts
           git clean -fxd
-        working-directory: pytorch
       # NOTE: These environment variables are put here so that they can be applied on every job equally
       #       They are also here because setting them at a workflow level doesn't give us access to the
       #       runner.temp variable, which we need.
@@ -477,21 +464,19 @@ jobs:
       - name: Populate binary env
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_populate_env.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_populate_env.sh"
       - name: Test PyTorch binary
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_windows_test.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_windows_test.sh"
       - name: Wait until all sessions have drained
         shell: powershell
-        working-directory: pytorch
         if: always()
         timeout-minutes: 120
         run: |
           .github\scripts\wait_for_ssh_to_drain.ps1
       - name: Kill active ssh sessions if still around (Useful if workflow was cancelled)
         shell: powershell
-        working-directory: pytorch
         if: always()
         run: |
           .github\scripts\kill_active_ssh_sessions.ps1
@@ -502,7 +487,7 @@ jobs:
       contents: read
     needs: wheel-py3_10-cuda12_6-test
     with:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
+      PYTORCH_ROOT: ${{ github.workspace }}
       PACKAGE_TYPE: wheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
@@ -513,19 +498,22 @@ jobs:
       build_name: wheel-py3_10-cuda12_6
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
     uses: ./.github/workflows/_binary-upload.yml
-  wheel-py3_10-cuda12_8-build:
+  wheel-py3_10-cuda13_0-build:
     if: ${{ github.repository_owner == 'pytorch' }}
     needs: get-label-type
     runs-on: "${{ needs.get-label-type.outputs.label-type }}windows.12xlarge"
     timeout-minutes: 360
     env:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
+      PYTORCH_ROOT: ${{ github.workspace }}
       PACKAGE_TYPE: wheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu128
-      GPU_ARCH_VERSION: "12.8"
+      DESIRED_CUDA: cu130
+      GPU_ARCH_VERSION: "13.0"
       GPU_ARCH_TYPE: cuda
       SKIP_ALL_TESTS: 1
       DESIRED_PYTHON: "3.10"
@@ -554,7 +542,7 @@ jobs:
           echo "instance-type: $(get_ec2_metadata instance-type)"
           echo "system info $(uname -a)"
       - name: "[FB EMPLOYEES] Enable SSH (Click me for login details)"
-        uses: pytorch/test-infra/.github/actions/setup-ssh@main
+        uses: pytorch/test-infra/.github/actions/setup-ssh@release/2.12
         continue-on-error: true
         with:
           github-secret: ${{ secrets.GITHUB_TOKEN }}
@@ -586,58 +574,53 @@ jobs:
       - name: Checkout PyTorch
         uses: actions/checkout@v4
         with:
-          ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
           submodules: recursive
-          path: pytorch
           show-progress: false
       - name: Clean PyTorch checkout
         run: |
           # Remove any artifacts from the previous checkouts
           git clean -fxd
-        working-directory: pytorch
       - name: Populate binary env
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_populate_env.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_populate_env.sh"
       - name: Build PyTorch binary
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_windows_build.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_windows_build.sh"
       - uses: actions/upload-artifact@v4.4.0
         if: always()
         with:
-          name: wheel-py3_10-cuda12_8
+          name: wheel-py3_10-cuda13_0
           retention-days: 14
           if-no-files-found: error
           path: "${{ env.PYTORCH_FINAL_PACKAGE_DIR }}"
       - name: Wait until all sessions have drained
         shell: powershell
-        working-directory: pytorch
         if: always()
         timeout-minutes: 120
         run: |
           .github\scripts\wait_for_ssh_to_drain.ps1
       - name: Kill active ssh sessions if still around (Useful if workflow was cancelled)
         shell: powershell
-        working-directory: pytorch
         if: always()
         run: |
           .github\scripts\kill_active_ssh_sessions.ps1
 
-  wheel-py3_10-cuda12_8-test:  # Testing
+  wheel-py3_10-cuda13_0-test:  # Testing
     if: ${{ github.repository_owner == 'pytorch' }}
     needs:
-      - wheel-py3_10-cuda12_8-build
+      - wheel-py3_10-cuda13_0-build
       - get-label-type
     runs-on: "${{ needs.get-label-type.outputs.label-type }}windows.g4dn.xlarge"
     timeout-minutes: 360
     env:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
+      PYTORCH_ROOT: ${{ github.workspace }}
       PACKAGE_TYPE: wheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu128
-      GPU_ARCH_VERSION: "12.8"
+      DESIRED_CUDA: cu130
+      GPU_ARCH_VERSION: "13.0"
       GPU_ARCH_TYPE: cuda
       SKIP_ALL_TESTS: 1
       DESIRED_PYTHON: "3.10"
@@ -657,7 +640,7 @@ jobs:
           echo "instance-type: $(get_ec2_metadata instance-type)"
           echo "system info $(uname -a)"
       - name: "[FB EMPLOYEES] Enable SSH (Click me for login details)"
-        uses: pytorch/test-infra/.github/actions/setup-ssh@main
+        uses: pytorch/test-infra/.github/actions/setup-ssh@release/2.12
         continue-on-error: true
         with:
           github-secret: ${{ secrets.GITHUB_TOKEN }}
@@ -689,15 +672,12 @@ jobs:
       - name: Checkout PyTorch
         uses: actions/checkout@v4
         with:
-          ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
           submodules: recursive
-          path: pytorch
           show-progress: false
       - name: Clean PyTorch checkout
         run: |
           # Remove any artifacts from the previous checkouts
           git clean -fxd
-        working-directory: pytorch
       # NOTE: These environment variables are put here so that they can be applied on every job equally
       #       They are also here because setting them at a workflow level doesn't give us access to the
       #       runner.temp variable, which we need.
@@ -710,60 +690,61 @@ jobs:
       - uses: actions/download-artifact@v4.1.7
         name: Download Build Artifacts
         with:
-          name: wheel-py3_10-cuda12_8
+          name: wheel-py3_10-cuda13_0
           path: "${{ env.PYTORCH_FINAL_PACKAGE_DIR }}"
       - name: Populate binary env
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_populate_env.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_populate_env.sh"
       - name: Test PyTorch binary
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_windows_test.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_windows_test.sh"
       - name: Wait until all sessions have drained
         shell: powershell
-        working-directory: pytorch
         if: always()
         timeout-minutes: 120
         run: |
           .github\scripts\wait_for_ssh_to_drain.ps1
       - name: Kill active ssh sessions if still around (Useful if workflow was cancelled)
         shell: powershell
-        working-directory: pytorch
         if: always()
         run: |
           .github\scripts\kill_active_ssh_sessions.ps1
-  wheel-py3_10-cuda12_8-upload:  # Uploading
+  wheel-py3_10-cuda13_0-upload:  # Uploading
     if: ${{ github.repository_owner == 'pytorch' }}
     permissions:
       id-token: write
       contents: read
-    needs: wheel-py3_10-cuda12_8-test
+    needs: wheel-py3_10-cuda13_0-test
     with:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
+      PYTORCH_ROOT: ${{ github.workspace }}
       PACKAGE_TYPE: wheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu128
-      GPU_ARCH_VERSION: "12.8"
+      DESIRED_CUDA: cu130
+      GPU_ARCH_VERSION: "13.0"
       GPU_ARCH_TYPE: cuda
       DESIRED_PYTHON: "3.10"
-      build_name: wheel-py3_10-cuda12_8
+      build_name: wheel-py3_10-cuda13_0
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
     uses: ./.github/workflows/_binary-upload.yml
-  wheel-py3_10-cuda13_0-build:
+  wheel-py3_10-cuda13_2-build:
     if: ${{ github.repository_owner == 'pytorch' }}
     needs: get-label-type
     runs-on: "${{ needs.get-label-type.outputs.label-type }}windows.12xlarge"
     timeout-minutes: 360
     env:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
+      PYTORCH_ROOT: ${{ github.workspace }}
       PACKAGE_TYPE: wheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu130
-      GPU_ARCH_VERSION: "13.0"
+      DESIRED_CUDA: cu132
+      GPU_ARCH_VERSION: "13.2"
       GPU_ARCH_TYPE: cuda
       SKIP_ALL_TESTS: 1
       DESIRED_PYTHON: "3.10"
@@ -792,7 +773,7 @@ jobs:
           echo "instance-type: $(get_ec2_metadata instance-type)"
           echo "system info $(uname -a)"
       - name: "[FB EMPLOYEES] Enable SSH (Click me for login details)"
-        uses: pytorch/test-infra/.github/actions/setup-ssh@main
+        uses: pytorch/test-infra/.github/actions/setup-ssh@release/2.12
         continue-on-error: true
         with:
           github-secret: ${{ secrets.GITHUB_TOKEN }}
@@ -824,58 +805,53 @@ jobs:
       - name: Checkout PyTorch
         uses: actions/checkout@v4
         with:
-          ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
           submodules: recursive
-          path: pytorch
           show-progress: false
       - name: Clean PyTorch checkout
         run: |
           # Remove any artifacts from the previous checkouts
           git clean -fxd
-        working-directory: pytorch
       - name: Populate binary env
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_populate_env.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_populate_env.sh"
       - name: Build PyTorch binary
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_windows_build.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_windows_build.sh"
       - uses: actions/upload-artifact@v4.4.0
         if: always()
         with:
-          name: wheel-py3_10-cuda13_0
+          name: wheel-py3_10-cuda13_2
           retention-days: 14
           if-no-files-found: error
           path: "${{ env.PYTORCH_FINAL_PACKAGE_DIR }}"
       - name: Wait until all sessions have drained
         shell: powershell
-        working-directory: pytorch
         if: always()
         timeout-minutes: 120
         run: |
           .github\scripts\wait_for_ssh_to_drain.ps1
       - name: Kill active ssh sessions if still around (Useful if workflow was cancelled)
         shell: powershell
-        working-directory: pytorch
         if: always()
         run: |
           .github\scripts\kill_active_ssh_sessions.ps1
 
-  wheel-py3_10-cuda13_0-test:  # Testing
+  wheel-py3_10-cuda13_2-test:  # Testing
     if: ${{ github.repository_owner == 'pytorch' }}
     needs:
-      - wheel-py3_10-cuda13_0-build
+      - wheel-py3_10-cuda13_2-build
       - get-label-type
     runs-on: "${{ needs.get-label-type.outputs.label-type }}windows.g4dn.xlarge"
     timeout-minutes: 360
     env:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
+      PYTORCH_ROOT: ${{ github.workspace }}
       PACKAGE_TYPE: wheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu130
-      GPU_ARCH_VERSION: "13.0"
+      DESIRED_CUDA: cu132
+      GPU_ARCH_VERSION: "13.2"
       GPU_ARCH_TYPE: cuda
       SKIP_ALL_TESTS: 1
       DESIRED_PYTHON: "3.10"
@@ -895,7 +871,7 @@ jobs:
           echo "instance-type: $(get_ec2_metadata instance-type)"
           echo "system info $(uname -a)"
       - name: "[FB EMPLOYEES] Enable SSH (Click me for login details)"
-        uses: pytorch/test-infra/.github/actions/setup-ssh@main
+        uses: pytorch/test-infra/.github/actions/setup-ssh@release/2.12
         continue-on-error: true
         with:
           github-secret: ${{ secrets.GITHUB_TOKEN }}
@@ -927,15 +903,12 @@ jobs:
       - name: Checkout PyTorch
         uses: actions/checkout@v4
         with:
-          ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
           submodules: recursive
-          path: pytorch
           show-progress: false
       - name: Clean PyTorch checkout
         run: |
           # Remove any artifacts from the previous checkouts
           git clean -fxd
-        working-directory: pytorch
       # NOTE: These environment variables are put here so that they can be applied on every job equally
       #       They are also here because setting them at a workflow level doesn't give us access to the
       #       runner.temp variable, which we need.
@@ -948,47 +921,48 @@ jobs:
       - uses: actions/download-artifact@v4.1.7
         name: Download Build Artifacts
         with:
-          name: wheel-py3_10-cuda13_0
+          name: wheel-py3_10-cuda13_2
           path: "${{ env.PYTORCH_FINAL_PACKAGE_DIR }}"
       - name: Populate binary env
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_populate_env.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_populate_env.sh"
       - name: Test PyTorch binary
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_windows_test.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_windows_test.sh"
       - name: Wait until all sessions have drained
         shell: powershell
-        working-directory: pytorch
         if: always()
         timeout-minutes: 120
         run: |
           .github\scripts\wait_for_ssh_to_drain.ps1
       - name: Kill active ssh sessions if still around (Useful if workflow was cancelled)
         shell: powershell
-        working-directory: pytorch
         if: always()
         run: |
           .github\scripts\kill_active_ssh_sessions.ps1
-  wheel-py3_10-cuda13_0-upload:  # Uploading
+  wheel-py3_10-cuda13_2-upload:  # Uploading
     if: ${{ github.repository_owner == 'pytorch' }}
     permissions:
       id-token: write
       contents: read
-    needs: wheel-py3_10-cuda13_0-test
+    needs: wheel-py3_10-cuda13_2-test
     with:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
+      PYTORCH_ROOT: ${{ github.workspace }}
       PACKAGE_TYPE: wheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu130
-      GPU_ARCH_VERSION: "13.0"
+      DESIRED_CUDA: cu132
+      GPU_ARCH_VERSION: "13.2"
       GPU_ARCH_TYPE: cuda
       DESIRED_PYTHON: "3.10"
-      build_name: wheel-py3_10-cuda13_0
+      build_name: wheel-py3_10-cuda13_2
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
     uses: ./.github/workflows/_binary-upload.yml
   wheel-py3_10-xpu-build:
     if: ${{ github.repository_owner == 'pytorch' }}
@@ -996,7 +970,7 @@ jobs:
     runs-on: "${{ needs.get-label-type.outputs.label-type }}windows.12xlarge"
     timeout-minutes: 360
     env:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
+      PYTORCH_ROOT: ${{ github.workspace }}
       PACKAGE_TYPE: wheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
@@ -1030,7 +1004,7 @@ jobs:
           echo "instance-type: $(get_ec2_metadata instance-type)"
           echo "system info $(uname -a)"
       - name: "[FB EMPLOYEES] Enable SSH (Click me for login details)"
-        uses: pytorch/test-infra/.github/actions/setup-ssh@main
+        uses: pytorch/test-infra/.github/actions/setup-ssh@release/2.12
         continue-on-error: true
         with:
           github-secret: ${{ secrets.GITHUB_TOKEN }}
@@ -1062,23 +1036,20 @@ jobs:
       - name: Checkout PyTorch
         uses: actions/checkout@v4
         with:
-          ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
           submodules: recursive
-          path: pytorch
           show-progress: false
       - name: Clean PyTorch checkout
         run: |
           # Remove any artifacts from the previous checkouts
           git clean -fxd
-        working-directory: pytorch
       - name: Populate binary env
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_populate_env.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_populate_env.sh"
       - name: Build PyTorch binary
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_windows_build.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_windows_build.sh"
       - uses: actions/upload-artifact@v4.4.0
         if: always()
         with:
@@ -1088,14 +1059,12 @@ jobs:
           path: "${{ env.PYTORCH_FINAL_PACKAGE_DIR }}"
       - name: Wait until all sessions have drained
         shell: powershell
-        working-directory: pytorch
         if: always()
         timeout-minutes: 120
         run: |
           .github\scripts\wait_for_ssh_to_drain.ps1
       - name: Kill active ssh sessions if still around (Useful if workflow was cancelled)
         shell: powershell
-        working-directory: pytorch
         if: always()
         run: |
           .github\scripts\kill_active_ssh_sessions.ps1
@@ -1108,7 +1077,7 @@ jobs:
     runs-on: "${{ needs.get-label-type.outputs.label-type }}windows.4xlarge"
     timeout-minutes: 360
     env:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
+      PYTORCH_ROOT: ${{ github.workspace }}
       PACKAGE_TYPE: wheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
@@ -1132,7 +1101,7 @@ jobs:
           echo "instance-type: $(get_ec2_metadata instance-type)"
           echo "system info $(uname -a)"
       - name: "[FB EMPLOYEES] Enable SSH (Click me for login details)"
-        uses: pytorch/test-infra/.github/actions/setup-ssh@main
+        uses: pytorch/test-infra/.github/actions/setup-ssh@release/2.12
         continue-on-error: true
         with:
           github-secret: ${{ secrets.GITHUB_TOKEN }}
@@ -1164,15 +1133,12 @@ jobs:
       - name: Checkout PyTorch
         uses: actions/checkout@v4
         with:
-          ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
           submodules: recursive
-          path: pytorch
           show-progress: false
       - name: Clean PyTorch checkout
         run: |
           # Remove any artifacts from the previous checkouts
           git clean -fxd
-        working-directory: pytorch
       # NOTE: These environment variables are put here so that they can be applied on every job equally
       #       They are also here because setting them at a workflow level doesn't give us access to the
       #       runner.temp variable, which we need.
@@ -1190,21 +1156,19 @@ jobs:
       - name: Populate binary env
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_populate_env.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_populate_env.sh"
       - name: Test PyTorch binary
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_windows_test.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_windows_test.sh"
       - name: Wait until all sessions have drained
         shell: powershell
-        working-directory: pytorch
         if: always()
         timeout-minutes: 120
         run: |
           .github\scripts\wait_for_ssh_to_drain.ps1
       - name: Kill active ssh sessions if still around (Useful if workflow was cancelled)
         shell: powershell
-        working-directory: pytorch
         if: always()
         run: |
           .github\scripts\kill_active_ssh_sessions.ps1
@@ -1215,7 +1179,7 @@ jobs:
       contents: read
     needs: wheel-py3_10-xpu-test
     with:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
+      PYTORCH_ROOT: ${{ github.workspace }}
       PACKAGE_TYPE: wheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
@@ -1225,6 +1189,9 @@ jobs:
       build_name: wheel-py3_10-xpu
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
     uses: ./.github/workflows/_binary-upload.yml
   wheel-py3_11-cpu-build:
     if: ${{ github.repository_owner == 'pytorch' }}
@@ -1232,7 +1199,7 @@ jobs:
     runs-on: "${{ needs.get-label-type.outputs.label-type }}windows.12xlarge"
     timeout-minutes: 360
     env:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
+      PYTORCH_ROOT: ${{ github.workspace }}
       PACKAGE_TYPE: wheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
@@ -1240,6 +1207,7 @@ jobs:
       GPU_ARCH_TYPE: cpu
       SKIP_ALL_TESTS: 1
       DESIRED_PYTHON: "3.11"
+      PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cufile,nvjitlink,nvtx]==13.0.2; platform_system == 'Linux' | nvidia-cublas>=13.1.0.3,<=13.1.1.3; platform_system == 'Linux' | cuda-bindings>=13.0.3,<14; platform_system == 'Linux' | nvidia-cudnn-cu13==9.20.0.48; platform_system == 'Linux' | nvidia-cusparselt-cu13==0.8.1; platform_system == 'Linux' | nvidia-nccl-cu13==2.29.7; platform_system == 'Linux' | nvidia-nvshmem-cu13==3.4.5; platform_system == 'Linux'
     steps:
       # NOTE: These environment variables are put here so that they can be applied on every job equally
       #       They are also here because setting them at a workflow level doesn't give us access to the
@@ -1265,7 +1233,7 @@ jobs:
           echo "instance-type: $(get_ec2_metadata instance-type)"
           echo "system info $(uname -a)"
       - name: "[FB EMPLOYEES] Enable SSH (Click me for login details)"
-        uses: pytorch/test-infra/.github/actions/setup-ssh@main
+        uses: pytorch/test-infra/.github/actions/setup-ssh@release/2.12
         continue-on-error: true
         with:
           github-secret: ${{ secrets.GITHUB_TOKEN }}
@@ -1297,23 +1265,20 @@ jobs:
       - name: Checkout PyTorch
         uses: actions/checkout@v4
         with:
-          ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
           submodules: recursive
-          path: pytorch
           show-progress: false
       - name: Clean PyTorch checkout
         run: |
           # Remove any artifacts from the previous checkouts
           git clean -fxd
-        working-directory: pytorch
       - name: Populate binary env
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_populate_env.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_populate_env.sh"
       - name: Build PyTorch binary
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_windows_build.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_windows_build.sh"
       - uses: actions/upload-artifact@v4.4.0
         if: always()
         with:
@@ -1323,14 +1288,12 @@ jobs:
           path: "${{ env.PYTORCH_FINAL_PACKAGE_DIR }}"
       - name: Wait until all sessions have drained
         shell: powershell
-        working-directory: pytorch
         if: always()
         timeout-minutes: 120
         run: |
           .github\scripts\wait_for_ssh_to_drain.ps1
       - name: Kill active ssh sessions if still around (Useful if workflow was cancelled)
         shell: powershell
-        working-directory: pytorch
         if: always()
         run: |
           .github\scripts\kill_active_ssh_sessions.ps1
@@ -1343,7 +1306,7 @@ jobs:
     runs-on: "${{ needs.get-label-type.outputs.label-type }}windows.4xlarge"
     timeout-minutes: 360
     env:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
+      PYTORCH_ROOT: ${{ github.workspace }}
       PACKAGE_TYPE: wheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
@@ -1367,7 +1330,7 @@ jobs:
           echo "instance-type: $(get_ec2_metadata instance-type)"
           echo "system info $(uname -a)"
       - name: "[FB EMPLOYEES] Enable SSH (Click me for login details)"
-        uses: pytorch/test-infra/.github/actions/setup-ssh@main
+        uses: pytorch/test-infra/.github/actions/setup-ssh@release/2.12
         continue-on-error: true
         with:
           github-secret: ${{ secrets.GITHUB_TOKEN }}
@@ -1399,15 +1362,12 @@ jobs:
       - name: Checkout PyTorch
         uses: actions/checkout@v4
         with:
-          ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
           submodules: recursive
-          path: pytorch
           show-progress: false
       - name: Clean PyTorch checkout
         run: |
           # Remove any artifacts from the previous checkouts
           git clean -fxd
-        working-directory: pytorch
       # NOTE: These environment variables are put here so that they can be applied on every job equally
       #       They are also here because setting them at a workflow level doesn't give us access to the
       #       runner.temp variable, which we need.
@@ -1425,21 +1385,19 @@ jobs:
       - name: Populate binary env
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_populate_env.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_populate_env.sh"
       - name: Test PyTorch binary
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_windows_test.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_windows_test.sh"
       - name: Wait until all sessions have drained
         shell: powershell
-        working-directory: pytorch
         if: always()
         timeout-minutes: 120
         run: |
           .github\scripts\wait_for_ssh_to_drain.ps1
       - name: Kill active ssh sessions if still around (Useful if workflow was cancelled)
         shell: powershell
-        working-directory: pytorch
         if: always()
         run: |
           .github\scripts\kill_active_ssh_sessions.ps1
@@ -1450,7 +1408,7 @@ jobs:
       contents: read
     needs: wheel-py3_11-cpu-test
     with:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
+      PYTORCH_ROOT: ${{ github.workspace }}
       PACKAGE_TYPE: wheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
@@ -1460,6 +1418,9 @@ jobs:
       build_name: wheel-py3_11-cpu
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
     uses: ./.github/workflows/_binary-upload.yml
   wheel-py3_11-cuda12_6-build:
     if: ${{ github.repository_owner == 'pytorch' }}
@@ -1467,7 +1428,7 @@ jobs:
     runs-on: "${{ needs.get-label-type.outputs.label-type }}windows.12xlarge"
     timeout-minutes: 360
     env:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
+      PYTORCH_ROOT: ${{ github.workspace }}
       PACKAGE_TYPE: wheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
@@ -1501,7 +1462,7 @@ jobs:
           echo "instance-type: $(get_ec2_metadata instance-type)"
           echo "system info $(uname -a)"
       - name: "[FB EMPLOYEES] Enable SSH (Click me for login details)"
-        uses: pytorch/test-infra/.github/actions/setup-ssh@main
+        uses: pytorch/test-infra/.github/actions/setup-ssh@release/2.12
         continue-on-error: true
         with:
           github-secret: ${{ secrets.GITHUB_TOKEN }}
@@ -1533,23 +1494,20 @@ jobs:
       - name: Checkout PyTorch
         uses: actions/checkout@v4
         with:
-          ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
           submodules: recursive
-          path: pytorch
           show-progress: false
       - name: Clean PyTorch checkout
         run: |
           # Remove any artifacts from the previous checkouts
           git clean -fxd
-        working-directory: pytorch
       - name: Populate binary env
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_populate_env.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_populate_env.sh"
       - name: Build PyTorch binary
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_windows_build.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_windows_build.sh"
       - uses: actions/upload-artifact@v4.4.0
         if: always()
         with:
@@ -1559,14 +1517,12 @@ jobs:
           path: "${{ env.PYTORCH_FINAL_PACKAGE_DIR }}"
       - name: Wait until all sessions have drained
         shell: powershell
-        working-directory: pytorch
         if: always()
         timeout-minutes: 120
         run: |
           .github\scripts\wait_for_ssh_to_drain.ps1
       - name: Kill active ssh sessions if still around (Useful if workflow was cancelled)
         shell: powershell
-        working-directory: pytorch
         if: always()
         run: |
           .github\scripts\kill_active_ssh_sessions.ps1
@@ -1579,7 +1535,7 @@ jobs:
     runs-on: "${{ needs.get-label-type.outputs.label-type }}windows.g4dn.xlarge"
     timeout-minutes: 360
     env:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
+      PYTORCH_ROOT: ${{ github.workspace }}
       PACKAGE_TYPE: wheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
@@ -1604,7 +1560,7 @@ jobs:
           echo "instance-type: $(get_ec2_metadata instance-type)"
           echo "system info $(uname -a)"
       - name: "[FB EMPLOYEES] Enable SSH (Click me for login details)"
-        uses: pytorch/test-infra/.github/actions/setup-ssh@main
+        uses: pytorch/test-infra/.github/actions/setup-ssh@release/2.12
         continue-on-error: true
         with:
           github-secret: ${{ secrets.GITHUB_TOKEN }}
@@ -1636,15 +1592,12 @@ jobs:
       - name: Checkout PyTorch
         uses: actions/checkout@v4
         with:
-          ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
           submodules: recursive
-          path: pytorch
           show-progress: false
       - name: Clean PyTorch checkout
         run: |
           # Remove any artifacts from the previous checkouts
           git clean -fxd
-        working-directory: pytorch
       # NOTE: These environment variables are put here so that they can be applied on every job equally
       #       They are also here because setting them at a workflow level doesn't give us access to the
       #       runner.temp variable, which we need.
@@ -1662,21 +1615,19 @@ jobs:
       - name: Populate binary env
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_populate_env.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_populate_env.sh"
       - name: Test PyTorch binary
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_windows_test.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_windows_test.sh"
       - name: Wait until all sessions have drained
         shell: powershell
-        working-directory: pytorch
         if: always()
         timeout-minutes: 120
         run: |
           .github\scripts\wait_for_ssh_to_drain.ps1
       - name: Kill active ssh sessions if still around (Useful if workflow was cancelled)
         shell: powershell
-        working-directory: pytorch
         if: always()
         run: |
           .github\scripts\kill_active_ssh_sessions.ps1
@@ -1687,7 +1638,7 @@ jobs:
       contents: read
     needs: wheel-py3_11-cuda12_6-test
     with:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
+      PYTORCH_ROOT: ${{ github.workspace }}
       PACKAGE_TYPE: wheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
@@ -1698,19 +1649,22 @@ jobs:
       build_name: wheel-py3_11-cuda12_6
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
     uses: ./.github/workflows/_binary-upload.yml
-  wheel-py3_11-cuda12_8-build:
+  wheel-py3_11-cuda13_0-build:
     if: ${{ github.repository_owner == 'pytorch' }}
     needs: get-label-type
     runs-on: "${{ needs.get-label-type.outputs.label-type }}windows.12xlarge"
     timeout-minutes: 360
     env:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
+      PYTORCH_ROOT: ${{ github.workspace }}
       PACKAGE_TYPE: wheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu128
-      GPU_ARCH_VERSION: "12.8"
+      DESIRED_CUDA: cu130
+      GPU_ARCH_VERSION: "13.0"
       GPU_ARCH_TYPE: cuda
       SKIP_ALL_TESTS: 1
       DESIRED_PYTHON: "3.11"
@@ -1739,7 +1693,7 @@ jobs:
           echo "instance-type: $(get_ec2_metadata instance-type)"
           echo "system info $(uname -a)"
       - name: "[FB EMPLOYEES] Enable SSH (Click me for login details)"
-        uses: pytorch/test-infra/.github/actions/setup-ssh@main
+        uses: pytorch/test-infra/.github/actions/setup-ssh@release/2.12
         continue-on-error: true
         with:
           github-secret: ${{ secrets.GITHUB_TOKEN }}
@@ -1771,58 +1725,53 @@ jobs:
       - name: Checkout PyTorch
         uses: actions/checkout@v4
         with:
-          ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
           submodules: recursive
-          path: pytorch
           show-progress: false
       - name: Clean PyTorch checkout
         run: |
           # Remove any artifacts from the previous checkouts
           git clean -fxd
-        working-directory: pytorch
       - name: Populate binary env
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_populate_env.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_populate_env.sh"
       - name: Build PyTorch binary
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_windows_build.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_windows_build.sh"
       - uses: actions/upload-artifact@v4.4.0
         if: always()
         with:
-          name: wheel-py3_11-cuda12_8
+          name: wheel-py3_11-cuda13_0
           retention-days: 14
           if-no-files-found: error
           path: "${{ env.PYTORCH_FINAL_PACKAGE_DIR }}"
       - name: Wait until all sessions have drained
         shell: powershell
-        working-directory: pytorch
         if: always()
         timeout-minutes: 120
         run: |
           .github\scripts\wait_for_ssh_to_drain.ps1
       - name: Kill active ssh sessions if still around (Useful if workflow was cancelled)
         shell: powershell
-        working-directory: pytorch
         if: always()
         run: |
           .github\scripts\kill_active_ssh_sessions.ps1
 
-  wheel-py3_11-cuda12_8-test:  # Testing
+  wheel-py3_11-cuda13_0-test:  # Testing
     if: ${{ github.repository_owner == 'pytorch' }}
     needs:
-      - wheel-py3_11-cuda12_8-build
+      - wheel-py3_11-cuda13_0-build
       - get-label-type
     runs-on: "${{ needs.get-label-type.outputs.label-type }}windows.g4dn.xlarge"
     timeout-minutes: 360
     env:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
+      PYTORCH_ROOT: ${{ github.workspace }}
       PACKAGE_TYPE: wheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu128
-      GPU_ARCH_VERSION: "12.8"
+      DESIRED_CUDA: cu130
+      GPU_ARCH_VERSION: "13.0"
       GPU_ARCH_TYPE: cuda
       SKIP_ALL_TESTS: 1
       DESIRED_PYTHON: "3.11"
@@ -1842,7 +1791,7 @@ jobs:
           echo "instance-type: $(get_ec2_metadata instance-type)"
           echo "system info $(uname -a)"
       - name: "[FB EMPLOYEES] Enable SSH (Click me for login details)"
-        uses: pytorch/test-infra/.github/actions/setup-ssh@main
+        uses: pytorch/test-infra/.github/actions/setup-ssh@release/2.12
         continue-on-error: true
         with:
           github-secret: ${{ secrets.GITHUB_TOKEN }}
@@ -1874,15 +1823,12 @@ jobs:
       - name: Checkout PyTorch
         uses: actions/checkout@v4
         with:
-          ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
           submodules: recursive
-          path: pytorch
           show-progress: false
       - name: Clean PyTorch checkout
         run: |
           # Remove any artifacts from the previous checkouts
           git clean -fxd
-        working-directory: pytorch
       # NOTE: These environment variables are put here so that they can be applied on every job equally
       #       They are also here because setting them at a workflow level doesn't give us access to the
       #       runner.temp variable, which we need.
@@ -1895,60 +1841,61 @@ jobs:
       - uses: actions/download-artifact@v4.1.7
         name: Download Build Artifacts
         with:
-          name: wheel-py3_11-cuda12_8
+          name: wheel-py3_11-cuda13_0
           path: "${{ env.PYTORCH_FINAL_PACKAGE_DIR }}"
       - name: Populate binary env
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_populate_env.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_populate_env.sh"
       - name: Test PyTorch binary
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_windows_test.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_windows_test.sh"
       - name: Wait until all sessions have drained
         shell: powershell
-        working-directory: pytorch
         if: always()
         timeout-minutes: 120
         run: |
           .github\scripts\wait_for_ssh_to_drain.ps1
       - name: Kill active ssh sessions if still around (Useful if workflow was cancelled)
         shell: powershell
-        working-directory: pytorch
         if: always()
         run: |
           .github\scripts\kill_active_ssh_sessions.ps1
-  wheel-py3_11-cuda12_8-upload:  # Uploading
+  wheel-py3_11-cuda13_0-upload:  # Uploading
     if: ${{ github.repository_owner == 'pytorch' }}
     permissions:
       id-token: write
       contents: read
-    needs: wheel-py3_11-cuda12_8-test
+    needs: wheel-py3_11-cuda13_0-test
     with:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
+      PYTORCH_ROOT: ${{ github.workspace }}
       PACKAGE_TYPE: wheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu128
-      GPU_ARCH_VERSION: "12.8"
+      DESIRED_CUDA: cu130
+      GPU_ARCH_VERSION: "13.0"
       GPU_ARCH_TYPE: cuda
       DESIRED_PYTHON: "3.11"
-      build_name: wheel-py3_11-cuda12_8
+      build_name: wheel-py3_11-cuda13_0
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
     uses: ./.github/workflows/_binary-upload.yml
-  wheel-py3_11-cuda13_0-build:
+  wheel-py3_11-cuda13_2-build:
     if: ${{ github.repository_owner == 'pytorch' }}
     needs: get-label-type
     runs-on: "${{ needs.get-label-type.outputs.label-type }}windows.12xlarge"
     timeout-minutes: 360
     env:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
+      PYTORCH_ROOT: ${{ github.workspace }}
       PACKAGE_TYPE: wheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu130
-      GPU_ARCH_VERSION: "13.0"
+      DESIRED_CUDA: cu132
+      GPU_ARCH_VERSION: "13.2"
       GPU_ARCH_TYPE: cuda
       SKIP_ALL_TESTS: 1
       DESIRED_PYTHON: "3.11"
@@ -1977,7 +1924,7 @@ jobs:
           echo "instance-type: $(get_ec2_metadata instance-type)"
           echo "system info $(uname -a)"
       - name: "[FB EMPLOYEES] Enable SSH (Click me for login details)"
-        uses: pytorch/test-infra/.github/actions/setup-ssh@main
+        uses: pytorch/test-infra/.github/actions/setup-ssh@release/2.12
         continue-on-error: true
         with:
           github-secret: ${{ secrets.GITHUB_TOKEN }}
@@ -2009,58 +1956,53 @@ jobs:
       - name: Checkout PyTorch
         uses: actions/checkout@v4
         with:
-          ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
           submodules: recursive
-          path: pytorch
           show-progress: false
       - name: Clean PyTorch checkout
         run: |
           # Remove any artifacts from the previous checkouts
           git clean -fxd
-        working-directory: pytorch
       - name: Populate binary env
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_populate_env.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_populate_env.sh"
       - name: Build PyTorch binary
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_windows_build.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_windows_build.sh"
       - uses: actions/upload-artifact@v4.4.0
         if: always()
         with:
-          name: wheel-py3_11-cuda13_0
+          name: wheel-py3_11-cuda13_2
           retention-days: 14
           if-no-files-found: error
           path: "${{ env.PYTORCH_FINAL_PACKAGE_DIR }}"
       - name: Wait until all sessions have drained
         shell: powershell
-        working-directory: pytorch
         if: always()
         timeout-minutes: 120
         run: |
           .github\scripts\wait_for_ssh_to_drain.ps1
       - name: Kill active ssh sessions if still around (Useful if workflow was cancelled)
         shell: powershell
-        working-directory: pytorch
         if: always()
         run: |
           .github\scripts\kill_active_ssh_sessions.ps1
 
-  wheel-py3_11-cuda13_0-test:  # Testing
+  wheel-py3_11-cuda13_2-test:  # Testing
     if: ${{ github.repository_owner == 'pytorch' }}
     needs:
-      - wheel-py3_11-cuda13_0-build
+      - wheel-py3_11-cuda13_2-build
       - get-label-type
     runs-on: "${{ needs.get-label-type.outputs.label-type }}windows.g4dn.xlarge"
     timeout-minutes: 360
     env:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
+      PYTORCH_ROOT: ${{ github.workspace }}
       PACKAGE_TYPE: wheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu130
-      GPU_ARCH_VERSION: "13.0"
+      DESIRED_CUDA: cu132
+      GPU_ARCH_VERSION: "13.2"
       GPU_ARCH_TYPE: cuda
       SKIP_ALL_TESTS: 1
       DESIRED_PYTHON: "3.11"
@@ -2080,7 +2022,7 @@ jobs:
           echo "instance-type: $(get_ec2_metadata instance-type)"
           echo "system info $(uname -a)"
       - name: "[FB EMPLOYEES] Enable SSH (Click me for login details)"
-        uses: pytorch/test-infra/.github/actions/setup-ssh@main
+        uses: pytorch/test-infra/.github/actions/setup-ssh@release/2.12
         continue-on-error: true
         with:
           github-secret: ${{ secrets.GITHUB_TOKEN }}
@@ -2112,15 +2054,12 @@ jobs:
       - name: Checkout PyTorch
         uses: actions/checkout@v4
         with:
-          ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
           submodules: recursive
-          path: pytorch
           show-progress: false
       - name: Clean PyTorch checkout
         run: |
           # Remove any artifacts from the previous checkouts
           git clean -fxd
-        working-directory: pytorch
       # NOTE: These environment variables are put here so that they can be applied on every job equally
       #       They are also here because setting them at a workflow level doesn't give us access to the
       #       runner.temp variable, which we need.
@@ -2133,47 +2072,48 @@ jobs:
       - uses: actions/download-artifact@v4.1.7
         name: Download Build Artifacts
         with:
-          name: wheel-py3_11-cuda13_0
+          name: wheel-py3_11-cuda13_2
           path: "${{ env.PYTORCH_FINAL_PACKAGE_DIR }}"
       - name: Populate binary env
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_populate_env.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_populate_env.sh"
       - name: Test PyTorch binary
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_windows_test.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_windows_test.sh"
       - name: Wait until all sessions have drained
         shell: powershell
-        working-directory: pytorch
         if: always()
         timeout-minutes: 120
         run: |
           .github\scripts\wait_for_ssh_to_drain.ps1
       - name: Kill active ssh sessions if still around (Useful if workflow was cancelled)
         shell: powershell
-        working-directory: pytorch
         if: always()
         run: |
           .github\scripts\kill_active_ssh_sessions.ps1
-  wheel-py3_11-cuda13_0-upload:  # Uploading
+  wheel-py3_11-cuda13_2-upload:  # Uploading
     if: ${{ github.repository_owner == 'pytorch' }}
     permissions:
       id-token: write
       contents: read
-    needs: wheel-py3_11-cuda13_0-test
+    needs: wheel-py3_11-cuda13_2-test
     with:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
+      PYTORCH_ROOT: ${{ github.workspace }}
       PACKAGE_TYPE: wheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu130
-      GPU_ARCH_VERSION: "13.0"
+      DESIRED_CUDA: cu132
+      GPU_ARCH_VERSION: "13.2"
       GPU_ARCH_TYPE: cuda
       DESIRED_PYTHON: "3.11"
-      build_name: wheel-py3_11-cuda13_0
+      build_name: wheel-py3_11-cuda13_2
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
     uses: ./.github/workflows/_binary-upload.yml
   wheel-py3_11-xpu-build:
     if: ${{ github.repository_owner == 'pytorch' }}
@@ -2181,7 +2121,7 @@ jobs:
     runs-on: "${{ needs.get-label-type.outputs.label-type }}windows.12xlarge"
     timeout-minutes: 360
     env:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
+      PYTORCH_ROOT: ${{ github.workspace }}
       PACKAGE_TYPE: wheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
@@ -2215,7 +2155,7 @@ jobs:
           echo "instance-type: $(get_ec2_metadata instance-type)"
           echo "system info $(uname -a)"
       - name: "[FB EMPLOYEES] Enable SSH (Click me for login details)"
-        uses: pytorch/test-infra/.github/actions/setup-ssh@main
+        uses: pytorch/test-infra/.github/actions/setup-ssh@release/2.12
         continue-on-error: true
         with:
           github-secret: ${{ secrets.GITHUB_TOKEN }}
@@ -2247,23 +2187,20 @@ jobs:
       - name: Checkout PyTorch
         uses: actions/checkout@v4
         with:
-          ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
           submodules: recursive
-          path: pytorch
           show-progress: false
       - name: Clean PyTorch checkout
         run: |
           # Remove any artifacts from the previous checkouts
           git clean -fxd
-        working-directory: pytorch
       - name: Populate binary env
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_populate_env.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_populate_env.sh"
       - name: Build PyTorch binary
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_windows_build.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_windows_build.sh"
       - uses: actions/upload-artifact@v4.4.0
         if: always()
         with:
@@ -2273,14 +2210,12 @@ jobs:
           path: "${{ env.PYTORCH_FINAL_PACKAGE_DIR }}"
       - name: Wait until all sessions have drained
         shell: powershell
-        working-directory: pytorch
         if: always()
         timeout-minutes: 120
         run: |
           .github\scripts\wait_for_ssh_to_drain.ps1
       - name: Kill active ssh sessions if still around (Useful if workflow was cancelled)
         shell: powershell
-        working-directory: pytorch
         if: always()
         run: |
           .github\scripts\kill_active_ssh_sessions.ps1
@@ -2293,7 +2228,7 @@ jobs:
     runs-on: "${{ needs.get-label-type.outputs.label-type }}windows.4xlarge"
     timeout-minutes: 360
     env:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
+      PYTORCH_ROOT: ${{ github.workspace }}
       PACKAGE_TYPE: wheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
@@ -2317,7 +2252,7 @@ jobs:
           echo "instance-type: $(get_ec2_metadata instance-type)"
           echo "system info $(uname -a)"
       - name: "[FB EMPLOYEES] Enable SSH (Click me for login details)"
-        uses: pytorch/test-infra/.github/actions/setup-ssh@main
+        uses: pytorch/test-infra/.github/actions/setup-ssh@release/2.12
         continue-on-error: true
         with:
           github-secret: ${{ secrets.GITHUB_TOKEN }}
@@ -2349,15 +2284,12 @@ jobs:
       - name: Checkout PyTorch
         uses: actions/checkout@v4
         with:
-          ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
           submodules: recursive
-          path: pytorch
           show-progress: false
       - name: Clean PyTorch checkout
         run: |
           # Remove any artifacts from the previous checkouts
           git clean -fxd
-        working-directory: pytorch
       # NOTE: These environment variables are put here so that they can be applied on every job equally
       #       They are also here because setting them at a workflow level doesn't give us access to the
       #       runner.temp variable, which we need.
@@ -2375,21 +2307,19 @@ jobs:
       - name: Populate binary env
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_populate_env.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_populate_env.sh"
       - name: Test PyTorch binary
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_windows_test.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_windows_test.sh"
       - name: Wait until all sessions have drained
         shell: powershell
-        working-directory: pytorch
         if: always()
         timeout-minutes: 120
         run: |
           .github\scripts\wait_for_ssh_to_drain.ps1
       - name: Kill active ssh sessions if still around (Useful if workflow was cancelled)
         shell: powershell
-        working-directory: pytorch
         if: always()
         run: |
           .github\scripts\kill_active_ssh_sessions.ps1
@@ -2400,7 +2330,7 @@ jobs:
       contents: read
     needs: wheel-py3_11-xpu-test
     with:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
+      PYTORCH_ROOT: ${{ github.workspace }}
       PACKAGE_TYPE: wheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
@@ -2410,6 +2340,9 @@ jobs:
       build_name: wheel-py3_11-xpu
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
     uses: ./.github/workflows/_binary-upload.yml
   wheel-py3_12-cpu-build:
     if: ${{ github.repository_owner == 'pytorch' }}
@@ -2417,7 +2350,7 @@ jobs:
     runs-on: "${{ needs.get-label-type.outputs.label-type }}windows.12xlarge"
     timeout-minutes: 360
     env:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
+      PYTORCH_ROOT: ${{ github.workspace }}
       PACKAGE_TYPE: wheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
@@ -2425,6 +2358,7 @@ jobs:
       GPU_ARCH_TYPE: cpu
       SKIP_ALL_TESTS: 1
       DESIRED_PYTHON: "3.12"
+      PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cufile,nvjitlink,nvtx]==13.0.2; platform_system == 'Linux' | nvidia-cublas>=13.1.0.3,<=13.1.1.3; platform_system == 'Linux' | cuda-bindings>=13.0.3,<14; platform_system == 'Linux' | nvidia-cudnn-cu13==9.20.0.48; platform_system == 'Linux' | nvidia-cusparselt-cu13==0.8.1; platform_system == 'Linux' | nvidia-nccl-cu13==2.29.7; platform_system == 'Linux' | nvidia-nvshmem-cu13==3.4.5; platform_system == 'Linux'
     steps:
       # NOTE: These environment variables are put here so that they can be applied on every job equally
       #       They are also here because setting them at a workflow level doesn't give us access to the
@@ -2450,7 +2384,7 @@ jobs:
           echo "instance-type: $(get_ec2_metadata instance-type)"
           echo "system info $(uname -a)"
       - name: "[FB EMPLOYEES] Enable SSH (Click me for login details)"
-        uses: pytorch/test-infra/.github/actions/setup-ssh@main
+        uses: pytorch/test-infra/.github/actions/setup-ssh@release/2.12
         continue-on-error: true
         with:
           github-secret: ${{ secrets.GITHUB_TOKEN }}
@@ -2482,23 +2416,20 @@ jobs:
       - name: Checkout PyTorch
         uses: actions/checkout@v4
         with:
-          ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
           submodules: recursive
-          path: pytorch
           show-progress: false
       - name: Clean PyTorch checkout
         run: |
           # Remove any artifacts from the previous checkouts
           git clean -fxd
-        working-directory: pytorch
       - name: Populate binary env
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_populate_env.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_populate_env.sh"
       - name: Build PyTorch binary
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_windows_build.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_windows_build.sh"
       - uses: actions/upload-artifact@v4.4.0
         if: always()
         with:
@@ -2508,14 +2439,12 @@ jobs:
           path: "${{ env.PYTORCH_FINAL_PACKAGE_DIR }}"
       - name: Wait until all sessions have drained
         shell: powershell
-        working-directory: pytorch
         if: always()
         timeout-minutes: 120
         run: |
           .github\scripts\wait_for_ssh_to_drain.ps1
       - name: Kill active ssh sessions if still around (Useful if workflow was cancelled)
         shell: powershell
-        working-directory: pytorch
         if: always()
         run: |
           .github\scripts\kill_active_ssh_sessions.ps1
@@ -2528,7 +2457,7 @@ jobs:
     runs-on: "${{ needs.get-label-type.outputs.label-type }}windows.4xlarge"
     timeout-minutes: 360
     env:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
+      PYTORCH_ROOT: ${{ github.workspace }}
       PACKAGE_TYPE: wheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
@@ -2552,7 +2481,7 @@ jobs:
           echo "instance-type: $(get_ec2_metadata instance-type)"
           echo "system info $(uname -a)"
       - name: "[FB EMPLOYEES] Enable SSH (Click me for login details)"
-        uses: pytorch/test-infra/.github/actions/setup-ssh@main
+        uses: pytorch/test-infra/.github/actions/setup-ssh@release/2.12
         continue-on-error: true
         with:
           github-secret: ${{ secrets.GITHUB_TOKEN }}
@@ -2584,15 +2513,12 @@ jobs:
       - name: Checkout PyTorch
         uses: actions/checkout@v4
         with:
-          ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
           submodules: recursive
-          path: pytorch
           show-progress: false
       - name: Clean PyTorch checkout
         run: |
           # Remove any artifacts from the previous checkouts
           git clean -fxd
-        working-directory: pytorch
       # NOTE: These environment variables are put here so that they can be applied on every job equally
       #       They are also here because setting them at a workflow level doesn't give us access to the
       #       runner.temp variable, which we need.
@@ -2610,21 +2536,19 @@ jobs:
       - name: Populate binary env
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_populate_env.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_populate_env.sh"
       - name: Test PyTorch binary
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_windows_test.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_windows_test.sh"
       - name: Wait until all sessions have drained
         shell: powershell
-        working-directory: pytorch
         if: always()
         timeout-minutes: 120
         run: |
           .github\scripts\wait_for_ssh_to_drain.ps1
       - name: Kill active ssh sessions if still around (Useful if workflow was cancelled)
         shell: powershell
-        working-directory: pytorch
         if: always()
         run: |
           .github\scripts\kill_active_ssh_sessions.ps1
@@ -2635,7 +2559,7 @@ jobs:
       contents: read
     needs: wheel-py3_12-cpu-test
     with:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
+      PYTORCH_ROOT: ${{ github.workspace }}
       PACKAGE_TYPE: wheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
@@ -2645,6 +2569,9 @@ jobs:
       build_name: wheel-py3_12-cpu
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
     uses: ./.github/workflows/_binary-upload.yml
   wheel-py3_12-cuda12_6-build:
     if: ${{ github.repository_owner == 'pytorch' }}
@@ -2652,7 +2579,7 @@ jobs:
     runs-on: "${{ needs.get-label-type.outputs.label-type }}windows.12xlarge"
     timeout-minutes: 360
     env:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
+      PYTORCH_ROOT: ${{ github.workspace }}
       PACKAGE_TYPE: wheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
@@ -2686,7 +2613,7 @@ jobs:
           echo "instance-type: $(get_ec2_metadata instance-type)"
           echo "system info $(uname -a)"
       - name: "[FB EMPLOYEES] Enable SSH (Click me for login details)"
-        uses: pytorch/test-infra/.github/actions/setup-ssh@main
+        uses: pytorch/test-infra/.github/actions/setup-ssh@release/2.12
         continue-on-error: true
         with:
           github-secret: ${{ secrets.GITHUB_TOKEN }}
@@ -2718,23 +2645,20 @@ jobs:
       - name: Checkout PyTorch
         uses: actions/checkout@v4
         with:
-          ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
           submodules: recursive
-          path: pytorch
           show-progress: false
       - name: Clean PyTorch checkout
         run: |
           # Remove any artifacts from the previous checkouts
           git clean -fxd
-        working-directory: pytorch
       - name: Populate binary env
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_populate_env.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_populate_env.sh"
       - name: Build PyTorch binary
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_windows_build.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_windows_build.sh"
       - uses: actions/upload-artifact@v4.4.0
         if: always()
         with:
@@ -2744,14 +2668,12 @@ jobs:
           path: "${{ env.PYTORCH_FINAL_PACKAGE_DIR }}"
       - name: Wait until all sessions have drained
         shell: powershell
-        working-directory: pytorch
         if: always()
         timeout-minutes: 120
         run: |
           .github\scripts\wait_for_ssh_to_drain.ps1
       - name: Kill active ssh sessions if still around (Useful if workflow was cancelled)
         shell: powershell
-        working-directory: pytorch
         if: always()
         run: |
           .github\scripts\kill_active_ssh_sessions.ps1
@@ -2764,7 +2686,7 @@ jobs:
     runs-on: "${{ needs.get-label-type.outputs.label-type }}windows.g4dn.xlarge"
     timeout-minutes: 360
     env:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
+      PYTORCH_ROOT: ${{ github.workspace }}
       PACKAGE_TYPE: wheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
@@ -2789,7 +2711,7 @@ jobs:
           echo "instance-type: $(get_ec2_metadata instance-type)"
           echo "system info $(uname -a)"
       - name: "[FB EMPLOYEES] Enable SSH (Click me for login details)"
-        uses: pytorch/test-infra/.github/actions/setup-ssh@main
+        uses: pytorch/test-infra/.github/actions/setup-ssh@release/2.12
         continue-on-error: true
         with:
           github-secret: ${{ secrets.GITHUB_TOKEN }}
@@ -2821,15 +2743,12 @@ jobs:
       - name: Checkout PyTorch
         uses: actions/checkout@v4
         with:
-          ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
           submodules: recursive
-          path: pytorch
           show-progress: false
       - name: Clean PyTorch checkout
         run: |
           # Remove any artifacts from the previous checkouts
           git clean -fxd
-        working-directory: pytorch
       # NOTE: These environment variables are put here so that they can be applied on every job equally
       #       They are also here because setting them at a workflow level doesn't give us access to the
       #       runner.temp variable, which we need.
@@ -2847,21 +2766,19 @@ jobs:
       - name: Populate binary env
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_populate_env.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_populate_env.sh"
       - name: Test PyTorch binary
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_windows_test.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_windows_test.sh"
       - name: Wait until all sessions have drained
         shell: powershell
-        working-directory: pytorch
         if: always()
         timeout-minutes: 120
         run: |
           .github\scripts\wait_for_ssh_to_drain.ps1
       - name: Kill active ssh sessions if still around (Useful if workflow was cancelled)
         shell: powershell
-        working-directory: pytorch
         if: always()
         run: |
           .github\scripts\kill_active_ssh_sessions.ps1
@@ -2872,7 +2789,7 @@ jobs:
       contents: read
     needs: wheel-py3_12-cuda12_6-test
     with:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
+      PYTORCH_ROOT: ${{ github.workspace }}
       PACKAGE_TYPE: wheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
@@ -2883,19 +2800,22 @@ jobs:
       build_name: wheel-py3_12-cuda12_6
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
     uses: ./.github/workflows/_binary-upload.yml
-  wheel-py3_12-cuda12_8-build:
+  wheel-py3_12-cuda13_0-build:
     if: ${{ github.repository_owner == 'pytorch' }}
     needs: get-label-type
     runs-on: "${{ needs.get-label-type.outputs.label-type }}windows.12xlarge"
     timeout-minutes: 360
     env:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
+      PYTORCH_ROOT: ${{ github.workspace }}
       PACKAGE_TYPE: wheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu128
-      GPU_ARCH_VERSION: "12.8"
+      DESIRED_CUDA: cu130
+      GPU_ARCH_VERSION: "13.0"
       GPU_ARCH_TYPE: cuda
       SKIP_ALL_TESTS: 1
       DESIRED_PYTHON: "3.12"
@@ -2924,7 +2844,7 @@ jobs:
           echo "instance-type: $(get_ec2_metadata instance-type)"
           echo "system info $(uname -a)"
       - name: "[FB EMPLOYEES] Enable SSH (Click me for login details)"
-        uses: pytorch/test-infra/.github/actions/setup-ssh@main
+        uses: pytorch/test-infra/.github/actions/setup-ssh@release/2.12
         continue-on-error: true
         with:
           github-secret: ${{ secrets.GITHUB_TOKEN }}
@@ -2956,58 +2876,53 @@ jobs:
       - name: Checkout PyTorch
         uses: actions/checkout@v4
         with:
-          ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
           submodules: recursive
-          path: pytorch
           show-progress: false
       - name: Clean PyTorch checkout
         run: |
           # Remove any artifacts from the previous checkouts
           git clean -fxd
-        working-directory: pytorch
       - name: Populate binary env
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_populate_env.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_populate_env.sh"
       - name: Build PyTorch binary
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_windows_build.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_windows_build.sh"
       - uses: actions/upload-artifact@v4.4.0
         if: always()
         with:
-          name: wheel-py3_12-cuda12_8
+          name: wheel-py3_12-cuda13_0
           retention-days: 14
           if-no-files-found: error
           path: "${{ env.PYTORCH_FINAL_PACKAGE_DIR }}"
       - name: Wait until all sessions have drained
         shell: powershell
-        working-directory: pytorch
         if: always()
         timeout-minutes: 120
         run: |
           .github\scripts\wait_for_ssh_to_drain.ps1
       - name: Kill active ssh sessions if still around (Useful if workflow was cancelled)
         shell: powershell
-        working-directory: pytorch
         if: always()
         run: |
           .github\scripts\kill_active_ssh_sessions.ps1
 
-  wheel-py3_12-cuda12_8-test:  # Testing
+  wheel-py3_12-cuda13_0-test:  # Testing
     if: ${{ github.repository_owner == 'pytorch' }}
     needs:
-      - wheel-py3_12-cuda12_8-build
+      - wheel-py3_12-cuda13_0-build
       - get-label-type
     runs-on: "${{ needs.get-label-type.outputs.label-type }}windows.g4dn.xlarge"
     timeout-minutes: 360
     env:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
+      PYTORCH_ROOT: ${{ github.workspace }}
       PACKAGE_TYPE: wheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu128
-      GPU_ARCH_VERSION: "12.8"
+      DESIRED_CUDA: cu130
+      GPU_ARCH_VERSION: "13.0"
       GPU_ARCH_TYPE: cuda
       SKIP_ALL_TESTS: 1
       DESIRED_PYTHON: "3.12"
@@ -3027,7 +2942,7 @@ jobs:
           echo "instance-type: $(get_ec2_metadata instance-type)"
           echo "system info $(uname -a)"
       - name: "[FB EMPLOYEES] Enable SSH (Click me for login details)"
-        uses: pytorch/test-infra/.github/actions/setup-ssh@main
+        uses: pytorch/test-infra/.github/actions/setup-ssh@release/2.12
         continue-on-error: true
         with:
           github-secret: ${{ secrets.GITHUB_TOKEN }}
@@ -3059,15 +2974,12 @@ jobs:
       - name: Checkout PyTorch
         uses: actions/checkout@v4
         with:
-          ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
           submodules: recursive
-          path: pytorch
           show-progress: false
       - name: Clean PyTorch checkout
         run: |
           # Remove any artifacts from the previous checkouts
           git clean -fxd
-        working-directory: pytorch
       # NOTE: These environment variables are put here so that they can be applied on every job equally
       #       They are also here because setting them at a workflow level doesn't give us access to the
       #       runner.temp variable, which we need.
@@ -3080,60 +2992,61 @@ jobs:
       - uses: actions/download-artifact@v4.1.7
         name: Download Build Artifacts
         with:
-          name: wheel-py3_12-cuda12_8
+          name: wheel-py3_12-cuda13_0
           path: "${{ env.PYTORCH_FINAL_PACKAGE_DIR }}"
       - name: Populate binary env
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_populate_env.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_populate_env.sh"
       - name: Test PyTorch binary
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_windows_test.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_windows_test.sh"
       - name: Wait until all sessions have drained
         shell: powershell
-        working-directory: pytorch
         if: always()
         timeout-minutes: 120
         run: |
           .github\scripts\wait_for_ssh_to_drain.ps1
       - name: Kill active ssh sessions if still around (Useful if workflow was cancelled)
         shell: powershell
-        working-directory: pytorch
         if: always()
         run: |
           .github\scripts\kill_active_ssh_sessions.ps1
-  wheel-py3_12-cuda12_8-upload:  # Uploading
+  wheel-py3_12-cuda13_0-upload:  # Uploading
     if: ${{ github.repository_owner == 'pytorch' }}
     permissions:
       id-token: write
       contents: read
-    needs: wheel-py3_12-cuda12_8-test
+    needs: wheel-py3_12-cuda13_0-test
     with:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
+      PYTORCH_ROOT: ${{ github.workspace }}
       PACKAGE_TYPE: wheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu128
-      GPU_ARCH_VERSION: "12.8"
+      DESIRED_CUDA: cu130
+      GPU_ARCH_VERSION: "13.0"
       GPU_ARCH_TYPE: cuda
       DESIRED_PYTHON: "3.12"
-      build_name: wheel-py3_12-cuda12_8
+      build_name: wheel-py3_12-cuda13_0
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
     uses: ./.github/workflows/_binary-upload.yml
-  wheel-py3_12-cuda13_0-build:
+  wheel-py3_12-cuda13_2-build:
     if: ${{ github.repository_owner == 'pytorch' }}
     needs: get-label-type
     runs-on: "${{ needs.get-label-type.outputs.label-type }}windows.12xlarge"
     timeout-minutes: 360
     env:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
+      PYTORCH_ROOT: ${{ github.workspace }}
       PACKAGE_TYPE: wheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu130
-      GPU_ARCH_VERSION: "13.0"
+      DESIRED_CUDA: cu132
+      GPU_ARCH_VERSION: "13.2"
       GPU_ARCH_TYPE: cuda
       SKIP_ALL_TESTS: 1
       DESIRED_PYTHON: "3.12"
@@ -3162,7 +3075,7 @@ jobs:
           echo "instance-type: $(get_ec2_metadata instance-type)"
           echo "system info $(uname -a)"
       - name: "[FB EMPLOYEES] Enable SSH (Click me for login details)"
-        uses: pytorch/test-infra/.github/actions/setup-ssh@main
+        uses: pytorch/test-infra/.github/actions/setup-ssh@release/2.12
         continue-on-error: true
         with:
           github-secret: ${{ secrets.GITHUB_TOKEN }}
@@ -3194,58 +3107,53 @@ jobs:
       - name: Checkout PyTorch
         uses: actions/checkout@v4
         with:
-          ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
           submodules: recursive
-          path: pytorch
           show-progress: false
       - name: Clean PyTorch checkout
         run: |
           # Remove any artifacts from the previous checkouts
           git clean -fxd
-        working-directory: pytorch
       - name: Populate binary env
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_populate_env.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_populate_env.sh"
       - name: Build PyTorch binary
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_windows_build.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_windows_build.sh"
       - uses: actions/upload-artifact@v4.4.0
         if: always()
         with:
-          name: wheel-py3_12-cuda13_0
+          name: wheel-py3_12-cuda13_2
           retention-days: 14
           if-no-files-found: error
           path: "${{ env.PYTORCH_FINAL_PACKAGE_DIR }}"
       - name: Wait until all sessions have drained
         shell: powershell
-        working-directory: pytorch
         if: always()
         timeout-minutes: 120
         run: |
           .github\scripts\wait_for_ssh_to_drain.ps1
       - name: Kill active ssh sessions if still around (Useful if workflow was cancelled)
         shell: powershell
-        working-directory: pytorch
         if: always()
         run: |
           .github\scripts\kill_active_ssh_sessions.ps1
 
-  wheel-py3_12-cuda13_0-test:  # Testing
+  wheel-py3_12-cuda13_2-test:  # Testing
     if: ${{ github.repository_owner == 'pytorch' }}
     needs:
-      - wheel-py3_12-cuda13_0-build
+      - wheel-py3_12-cuda13_2-build
       - get-label-type
     runs-on: "${{ needs.get-label-type.outputs.label-type }}windows.g4dn.xlarge"
     timeout-minutes: 360
     env:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
+      PYTORCH_ROOT: ${{ github.workspace }}
       PACKAGE_TYPE: wheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu130
-      GPU_ARCH_VERSION: "13.0"
+      DESIRED_CUDA: cu132
+      GPU_ARCH_VERSION: "13.2"
       GPU_ARCH_TYPE: cuda
       SKIP_ALL_TESTS: 1
       DESIRED_PYTHON: "3.12"
@@ -3265,7 +3173,7 @@ jobs:
           echo "instance-type: $(get_ec2_metadata instance-type)"
           echo "system info $(uname -a)"
       - name: "[FB EMPLOYEES] Enable SSH (Click me for login details)"
-        uses: pytorch/test-infra/.github/actions/setup-ssh@main
+        uses: pytorch/test-infra/.github/actions/setup-ssh@release/2.12
         continue-on-error: true
         with:
           github-secret: ${{ secrets.GITHUB_TOKEN }}
@@ -3297,15 +3205,12 @@ jobs:
       - name: Checkout PyTorch
         uses: actions/checkout@v4
         with:
-          ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
           submodules: recursive
-          path: pytorch
           show-progress: false
       - name: Clean PyTorch checkout
         run: |
           # Remove any artifacts from the previous checkouts
           git clean -fxd
-        working-directory: pytorch
       # NOTE: These environment variables are put here so that they can be applied on every job equally
       #       They are also here because setting them at a workflow level doesn't give us access to the
       #       runner.temp variable, which we need.
@@ -3318,47 +3223,48 @@ jobs:
       - uses: actions/download-artifact@v4.1.7
         name: Download Build Artifacts
         with:
-          name: wheel-py3_12-cuda13_0
+          name: wheel-py3_12-cuda13_2
           path: "${{ env.PYTORCH_FINAL_PACKAGE_DIR }}"
       - name: Populate binary env
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_populate_env.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_populate_env.sh"
       - name: Test PyTorch binary
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_windows_test.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_windows_test.sh"
       - name: Wait until all sessions have drained
         shell: powershell
-        working-directory: pytorch
         if: always()
         timeout-minutes: 120
         run: |
           .github\scripts\wait_for_ssh_to_drain.ps1
       - name: Kill active ssh sessions if still around (Useful if workflow was cancelled)
         shell: powershell
-        working-directory: pytorch
         if: always()
         run: |
           .github\scripts\kill_active_ssh_sessions.ps1
-  wheel-py3_12-cuda13_0-upload:  # Uploading
+  wheel-py3_12-cuda13_2-upload:  # Uploading
     if: ${{ github.repository_owner == 'pytorch' }}
     permissions:
       id-token: write
       contents: read
-    needs: wheel-py3_12-cuda13_0-test
+    needs: wheel-py3_12-cuda13_2-test
     with:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
+      PYTORCH_ROOT: ${{ github.workspace }}
       PACKAGE_TYPE: wheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu130
-      GPU_ARCH_VERSION: "13.0"
+      DESIRED_CUDA: cu132
+      GPU_ARCH_VERSION: "13.2"
       GPU_ARCH_TYPE: cuda
       DESIRED_PYTHON: "3.12"
-      build_name: wheel-py3_12-cuda13_0
+      build_name: wheel-py3_12-cuda13_2
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
     uses: ./.github/workflows/_binary-upload.yml
   wheel-py3_12-xpu-build:
     if: ${{ github.repository_owner == 'pytorch' }}
@@ -3366,7 +3272,7 @@ jobs:
     runs-on: "${{ needs.get-label-type.outputs.label-type }}windows.12xlarge"
     timeout-minutes: 360
     env:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
+      PYTORCH_ROOT: ${{ github.workspace }}
       PACKAGE_TYPE: wheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
@@ -3400,7 +3306,7 @@ jobs:
           echo "instance-type: $(get_ec2_metadata instance-type)"
           echo "system info $(uname -a)"
       - name: "[FB EMPLOYEES] Enable SSH (Click me for login details)"
-        uses: pytorch/test-infra/.github/actions/setup-ssh@main
+        uses: pytorch/test-infra/.github/actions/setup-ssh@release/2.12
         continue-on-error: true
         with:
           github-secret: ${{ secrets.GITHUB_TOKEN }}
@@ -3432,23 +3338,20 @@ jobs:
       - name: Checkout PyTorch
         uses: actions/checkout@v4
         with:
-          ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
           submodules: recursive
-          path: pytorch
           show-progress: false
       - name: Clean PyTorch checkout
         run: |
           # Remove any artifacts from the previous checkouts
           git clean -fxd
-        working-directory: pytorch
       - name: Populate binary env
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_populate_env.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_populate_env.sh"
       - name: Build PyTorch binary
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_windows_build.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_windows_build.sh"
       - uses: actions/upload-artifact@v4.4.0
         if: always()
         with:
@@ -3458,14 +3361,12 @@ jobs:
           path: "${{ env.PYTORCH_FINAL_PACKAGE_DIR }}"
       - name: Wait until all sessions have drained
         shell: powershell
-        working-directory: pytorch
         if: always()
         timeout-minutes: 120
         run: |
           .github\scripts\wait_for_ssh_to_drain.ps1
       - name: Kill active ssh sessions if still around (Useful if workflow was cancelled)
         shell: powershell
-        working-directory: pytorch
         if: always()
         run: |
           .github\scripts\kill_active_ssh_sessions.ps1
@@ -3478,7 +3379,7 @@ jobs:
     runs-on: "${{ needs.get-label-type.outputs.label-type }}windows.4xlarge"
     timeout-minutes: 360
     env:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
+      PYTORCH_ROOT: ${{ github.workspace }}
       PACKAGE_TYPE: wheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
@@ -3502,7 +3403,7 @@ jobs:
           echo "instance-type: $(get_ec2_metadata instance-type)"
           echo "system info $(uname -a)"
       - name: "[FB EMPLOYEES] Enable SSH (Click me for login details)"
-        uses: pytorch/test-infra/.github/actions/setup-ssh@main
+        uses: pytorch/test-infra/.github/actions/setup-ssh@release/2.12
         continue-on-error: true
         with:
           github-secret: ${{ secrets.GITHUB_TOKEN }}
@@ -3534,15 +3435,12 @@ jobs:
       - name: Checkout PyTorch
         uses: actions/checkout@v4
         with:
-          ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
           submodules: recursive
-          path: pytorch
           show-progress: false
       - name: Clean PyTorch checkout
         run: |
           # Remove any artifacts from the previous checkouts
           git clean -fxd
-        working-directory: pytorch
       # NOTE: These environment variables are put here so that they can be applied on every job equally
       #       They are also here because setting them at a workflow level doesn't give us access to the
       #       runner.temp variable, which we need.
@@ -3560,21 +3458,19 @@ jobs:
       - name: Populate binary env
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_populate_env.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_populate_env.sh"
       - name: Test PyTorch binary
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_windows_test.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_windows_test.sh"
       - name: Wait until all sessions have drained
         shell: powershell
-        working-directory: pytorch
         if: always()
         timeout-minutes: 120
         run: |
           .github\scripts\wait_for_ssh_to_drain.ps1
       - name: Kill active ssh sessions if still around (Useful if workflow was cancelled)
         shell: powershell
-        working-directory: pytorch
         if: always()
         run: |
           .github\scripts\kill_active_ssh_sessions.ps1
@@ -3585,7 +3481,7 @@ jobs:
       contents: read
     needs: wheel-py3_12-xpu-test
     with:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
+      PYTORCH_ROOT: ${{ github.workspace }}
       PACKAGE_TYPE: wheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
@@ -3595,6 +3491,9 @@ jobs:
       build_name: wheel-py3_12-xpu
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
     uses: ./.github/workflows/_binary-upload.yml
   wheel-py3_13-cpu-build:
     if: ${{ github.repository_owner == 'pytorch' }}
@@ -3602,7 +3501,7 @@ jobs:
     runs-on: "${{ needs.get-label-type.outputs.label-type }}windows.12xlarge"
     timeout-minutes: 360
     env:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
+      PYTORCH_ROOT: ${{ github.workspace }}
       PACKAGE_TYPE: wheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
@@ -3610,6 +3509,7 @@ jobs:
       GPU_ARCH_TYPE: cpu
       SKIP_ALL_TESTS: 1
       DESIRED_PYTHON: "3.13"
+      PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cufile,nvjitlink,nvtx]==13.0.2; platform_system == 'Linux' | nvidia-cublas>=13.1.0.3,<=13.1.1.3; platform_system == 'Linux' | cuda-bindings>=13.0.3,<14; platform_system == 'Linux' | nvidia-cudnn-cu13==9.20.0.48; platform_system == 'Linux' | nvidia-cusparselt-cu13==0.8.1; platform_system == 'Linux' | nvidia-nccl-cu13==2.29.7; platform_system == 'Linux' | nvidia-nvshmem-cu13==3.4.5; platform_system == 'Linux'
     steps:
       # NOTE: These environment variables are put here so that they can be applied on every job equally
       #       They are also here because setting them at a workflow level doesn't give us access to the
@@ -3635,7 +3535,7 @@ jobs:
           echo "instance-type: $(get_ec2_metadata instance-type)"
           echo "system info $(uname -a)"
       - name: "[FB EMPLOYEES] Enable SSH (Click me for login details)"
-        uses: pytorch/test-infra/.github/actions/setup-ssh@main
+        uses: pytorch/test-infra/.github/actions/setup-ssh@release/2.12
         continue-on-error: true
         with:
           github-secret: ${{ secrets.GITHUB_TOKEN }}
@@ -3667,23 +3567,20 @@ jobs:
       - name: Checkout PyTorch
         uses: actions/checkout@v4
         with:
-          ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
           submodules: recursive
-          path: pytorch
           show-progress: false
       - name: Clean PyTorch checkout
         run: |
           # Remove any artifacts from the previous checkouts
           git clean -fxd
-        working-directory: pytorch
       - name: Populate binary env
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_populate_env.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_populate_env.sh"
       - name: Build PyTorch binary
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_windows_build.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_windows_build.sh"
       - uses: actions/upload-artifact@v4.4.0
         if: always()
         with:
@@ -3693,14 +3590,12 @@ jobs:
           path: "${{ env.PYTORCH_FINAL_PACKAGE_DIR }}"
       - name: Wait until all sessions have drained
         shell: powershell
-        working-directory: pytorch
         if: always()
         timeout-minutes: 120
         run: |
           .github\scripts\wait_for_ssh_to_drain.ps1
       - name: Kill active ssh sessions if still around (Useful if workflow was cancelled)
         shell: powershell
-        working-directory: pytorch
         if: always()
         run: |
           .github\scripts\kill_active_ssh_sessions.ps1
@@ -3713,7 +3608,7 @@ jobs:
     runs-on: "${{ needs.get-label-type.outputs.label-type }}windows.4xlarge"
     timeout-minutes: 360
     env:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
+      PYTORCH_ROOT: ${{ github.workspace }}
       PACKAGE_TYPE: wheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
@@ -3737,7 +3632,7 @@ jobs:
           echo "instance-type: $(get_ec2_metadata instance-type)"
           echo "system info $(uname -a)"
       - name: "[FB EMPLOYEES] Enable SSH (Click me for login details)"
-        uses: pytorch/test-infra/.github/actions/setup-ssh@main
+        uses: pytorch/test-infra/.github/actions/setup-ssh@release/2.12
         continue-on-error: true
         with:
           github-secret: ${{ secrets.GITHUB_TOKEN }}
@@ -3769,15 +3664,12 @@ jobs:
       - name: Checkout PyTorch
         uses: actions/checkout@v4
         with:
-          ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
           submodules: recursive
-          path: pytorch
           show-progress: false
       - name: Clean PyTorch checkout
         run: |
           # Remove any artifacts from the previous checkouts
           git clean -fxd
-        working-directory: pytorch
       # NOTE: These environment variables are put here so that they can be applied on every job equally
       #       They are also here because setting them at a workflow level doesn't give us access to the
       #       runner.temp variable, which we need.
@@ -3795,21 +3687,19 @@ jobs:
       - name: Populate binary env
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_populate_env.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_populate_env.sh"
       - name: Test PyTorch binary
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_windows_test.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_windows_test.sh"
       - name: Wait until all sessions have drained
         shell: powershell
-        working-directory: pytorch
         if: always()
         timeout-minutes: 120
         run: |
           .github\scripts\wait_for_ssh_to_drain.ps1
       - name: Kill active ssh sessions if still around (Useful if workflow was cancelled)
         shell: powershell
-        working-directory: pytorch
         if: always()
         run: |
           .github\scripts\kill_active_ssh_sessions.ps1
@@ -3820,7 +3710,7 @@ jobs:
       contents: read
     needs: wheel-py3_13-cpu-test
     with:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
+      PYTORCH_ROOT: ${{ github.workspace }}
       PACKAGE_TYPE: wheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
@@ -3830,6 +3720,9 @@ jobs:
       build_name: wheel-py3_13-cpu
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
     uses: ./.github/workflows/_binary-upload.yml
   wheel-py3_13-cuda12_6-build:
     if: ${{ github.repository_owner == 'pytorch' }}
@@ -3837,7 +3730,7 @@ jobs:
     runs-on: "${{ needs.get-label-type.outputs.label-type }}windows.12xlarge"
     timeout-minutes: 360
     env:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
+      PYTORCH_ROOT: ${{ github.workspace }}
       PACKAGE_TYPE: wheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
@@ -3871,7 +3764,7 @@ jobs:
           echo "instance-type: $(get_ec2_metadata instance-type)"
           echo "system info $(uname -a)"
       - name: "[FB EMPLOYEES] Enable SSH (Click me for login details)"
-        uses: pytorch/test-infra/.github/actions/setup-ssh@main
+        uses: pytorch/test-infra/.github/actions/setup-ssh@release/2.12
         continue-on-error: true
         with:
           github-secret: ${{ secrets.GITHUB_TOKEN }}
@@ -3903,23 +3796,20 @@ jobs:
       - name: Checkout PyTorch
         uses: actions/checkout@v4
         with:
-          ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
           submodules: recursive
-          path: pytorch
           show-progress: false
       - name: Clean PyTorch checkout
         run: |
           # Remove any artifacts from the previous checkouts
           git clean -fxd
-        working-directory: pytorch
       - name: Populate binary env
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_populate_env.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_populate_env.sh"
       - name: Build PyTorch binary
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_windows_build.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_windows_build.sh"
       - uses: actions/upload-artifact@v4.4.0
         if: always()
         with:
@@ -3929,14 +3819,12 @@ jobs:
           path: "${{ env.PYTORCH_FINAL_PACKAGE_DIR }}"
       - name: Wait until all sessions have drained
         shell: powershell
-        working-directory: pytorch
         if: always()
         timeout-minutes: 120
         run: |
           .github\scripts\wait_for_ssh_to_drain.ps1
       - name: Kill active ssh sessions if still around (Useful if workflow was cancelled)
         shell: powershell
-        working-directory: pytorch
         if: always()
         run: |
           .github\scripts\kill_active_ssh_sessions.ps1
@@ -3949,7 +3837,7 @@ jobs:
     runs-on: "${{ needs.get-label-type.outputs.label-type }}windows.g4dn.xlarge"
     timeout-minutes: 360
     env:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
+      PYTORCH_ROOT: ${{ github.workspace }}
       PACKAGE_TYPE: wheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
@@ -3974,7 +3862,7 @@ jobs:
           echo "instance-type: $(get_ec2_metadata instance-type)"
           echo "system info $(uname -a)"
       - name: "[FB EMPLOYEES] Enable SSH (Click me for login details)"
-        uses: pytorch/test-infra/.github/actions/setup-ssh@main
+        uses: pytorch/test-infra/.github/actions/setup-ssh@release/2.12
         continue-on-error: true
         with:
           github-secret: ${{ secrets.GITHUB_TOKEN }}
@@ -4006,15 +3894,12 @@ jobs:
       - name: Checkout PyTorch
         uses: actions/checkout@v4
         with:
-          ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
           submodules: recursive
-          path: pytorch
           show-progress: false
       - name: Clean PyTorch checkout
         run: |
           # Remove any artifacts from the previous checkouts
           git clean -fxd
-        working-directory: pytorch
       # NOTE: These environment variables are put here so that they can be applied on every job equally
       #       They are also here because setting them at a workflow level doesn't give us access to the
       #       runner.temp variable, which we need.
@@ -4032,21 +3917,19 @@ jobs:
       - name: Populate binary env
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_populate_env.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_populate_env.sh"
       - name: Test PyTorch binary
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_windows_test.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_windows_test.sh"
       - name: Wait until all sessions have drained
         shell: powershell
-        working-directory: pytorch
         if: always()
         timeout-minutes: 120
         run: |
           .github\scripts\wait_for_ssh_to_drain.ps1
       - name: Kill active ssh sessions if still around (Useful if workflow was cancelled)
         shell: powershell
-        working-directory: pytorch
         if: always()
         run: |
           .github\scripts\kill_active_ssh_sessions.ps1
@@ -4057,7 +3940,7 @@ jobs:
       contents: read
     needs: wheel-py3_13-cuda12_6-test
     with:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
+      PYTORCH_ROOT: ${{ github.workspace }}
       PACKAGE_TYPE: wheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
@@ -4068,19 +3951,22 @@ jobs:
       build_name: wheel-py3_13-cuda12_6
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
     uses: ./.github/workflows/_binary-upload.yml
-  wheel-py3_13-cuda12_8-build:
+  wheel-py3_13-cuda13_0-build:
     if: ${{ github.repository_owner == 'pytorch' }}
     needs: get-label-type
     runs-on: "${{ needs.get-label-type.outputs.label-type }}windows.12xlarge"
     timeout-minutes: 360
     env:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
+      PYTORCH_ROOT: ${{ github.workspace }}
       PACKAGE_TYPE: wheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu128
-      GPU_ARCH_VERSION: "12.8"
+      DESIRED_CUDA: cu130
+      GPU_ARCH_VERSION: "13.0"
       GPU_ARCH_TYPE: cuda
       SKIP_ALL_TESTS: 1
       DESIRED_PYTHON: "3.13"
@@ -4109,7 +3995,7 @@ jobs:
           echo "instance-type: $(get_ec2_metadata instance-type)"
           echo "system info $(uname -a)"
       - name: "[FB EMPLOYEES] Enable SSH (Click me for login details)"
-        uses: pytorch/test-infra/.github/actions/setup-ssh@main
+        uses: pytorch/test-infra/.github/actions/setup-ssh@release/2.12
         continue-on-error: true
         with:
           github-secret: ${{ secrets.GITHUB_TOKEN }}
@@ -4141,58 +4027,53 @@ jobs:
       - name: Checkout PyTorch
         uses: actions/checkout@v4
         with:
-          ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
           submodules: recursive
-          path: pytorch
           show-progress: false
       - name: Clean PyTorch checkout
         run: |
           # Remove any artifacts from the previous checkouts
           git clean -fxd
-        working-directory: pytorch
       - name: Populate binary env
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_populate_env.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_populate_env.sh"
       - name: Build PyTorch binary
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_windows_build.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_windows_build.sh"
       - uses: actions/upload-artifact@v4.4.0
         if: always()
         with:
-          name: wheel-py3_13-cuda12_8
+          name: wheel-py3_13-cuda13_0
           retention-days: 14
           if-no-files-found: error
           path: "${{ env.PYTORCH_FINAL_PACKAGE_DIR }}"
       - name: Wait until all sessions have drained
         shell: powershell
-        working-directory: pytorch
         if: always()
         timeout-minutes: 120
         run: |
           .github\scripts\wait_for_ssh_to_drain.ps1
       - name: Kill active ssh sessions if still around (Useful if workflow was cancelled)
         shell: powershell
-        working-directory: pytorch
         if: always()
         run: |
           .github\scripts\kill_active_ssh_sessions.ps1
 
-  wheel-py3_13-cuda12_8-test:  # Testing
+  wheel-py3_13-cuda13_0-test:  # Testing
     if: ${{ github.repository_owner == 'pytorch' }}
     needs:
-      - wheel-py3_13-cuda12_8-build
+      - wheel-py3_13-cuda13_0-build
       - get-label-type
     runs-on: "${{ needs.get-label-type.outputs.label-type }}windows.g4dn.xlarge"
     timeout-minutes: 360
     env:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
+      PYTORCH_ROOT: ${{ github.workspace }}
       PACKAGE_TYPE: wheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu128
-      GPU_ARCH_VERSION: "12.8"
+      DESIRED_CUDA: cu130
+      GPU_ARCH_VERSION: "13.0"
       GPU_ARCH_TYPE: cuda
       SKIP_ALL_TESTS: 1
       DESIRED_PYTHON: "3.13"
@@ -4212,7 +4093,7 @@ jobs:
           echo "instance-type: $(get_ec2_metadata instance-type)"
           echo "system info $(uname -a)"
       - name: "[FB EMPLOYEES] Enable SSH (Click me for login details)"
-        uses: pytorch/test-infra/.github/actions/setup-ssh@main
+        uses: pytorch/test-infra/.github/actions/setup-ssh@release/2.12
         continue-on-error: true
         with:
           github-secret: ${{ secrets.GITHUB_TOKEN }}
@@ -4244,15 +4125,12 @@ jobs:
       - name: Checkout PyTorch
         uses: actions/checkout@v4
         with:
-          ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
           submodules: recursive
-          path: pytorch
           show-progress: false
       - name: Clean PyTorch checkout
         run: |
           # Remove any artifacts from the previous checkouts
           git clean -fxd
-        working-directory: pytorch
       # NOTE: These environment variables are put here so that they can be applied on every job equally
       #       They are also here because setting them at a workflow level doesn't give us access to the
       #       runner.temp variable, which we need.
@@ -4265,60 +4143,61 @@ jobs:
       - uses: actions/download-artifact@v4.1.7
         name: Download Build Artifacts
         with:
-          name: wheel-py3_13-cuda12_8
+          name: wheel-py3_13-cuda13_0
           path: "${{ env.PYTORCH_FINAL_PACKAGE_DIR }}"
       - name: Populate binary env
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_populate_env.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_populate_env.sh"
       - name: Test PyTorch binary
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_windows_test.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_windows_test.sh"
       - name: Wait until all sessions have drained
         shell: powershell
-        working-directory: pytorch
         if: always()
         timeout-minutes: 120
         run: |
           .github\scripts\wait_for_ssh_to_drain.ps1
       - name: Kill active ssh sessions if still around (Useful if workflow was cancelled)
         shell: powershell
-        working-directory: pytorch
         if: always()
         run: |
           .github\scripts\kill_active_ssh_sessions.ps1
-  wheel-py3_13-cuda12_8-upload:  # Uploading
+  wheel-py3_13-cuda13_0-upload:  # Uploading
     if: ${{ github.repository_owner == 'pytorch' }}
     permissions:
       id-token: write
       contents: read
-    needs: wheel-py3_13-cuda12_8-test
+    needs: wheel-py3_13-cuda13_0-test
     with:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
+      PYTORCH_ROOT: ${{ github.workspace }}
       PACKAGE_TYPE: wheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu128
-      GPU_ARCH_VERSION: "12.8"
+      DESIRED_CUDA: cu130
+      GPU_ARCH_VERSION: "13.0"
       GPU_ARCH_TYPE: cuda
       DESIRED_PYTHON: "3.13"
-      build_name: wheel-py3_13-cuda12_8
+      build_name: wheel-py3_13-cuda13_0
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
     uses: ./.github/workflows/_binary-upload.yml
-  wheel-py3_13-cuda13_0-build:
+  wheel-py3_13-cuda13_2-build:
     if: ${{ github.repository_owner == 'pytorch' }}
     needs: get-label-type
     runs-on: "${{ needs.get-label-type.outputs.label-type }}windows.12xlarge"
     timeout-minutes: 360
     env:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
+      PYTORCH_ROOT: ${{ github.workspace }}
       PACKAGE_TYPE: wheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu130
-      GPU_ARCH_VERSION: "13.0"
+      DESIRED_CUDA: cu132
+      GPU_ARCH_VERSION: "13.2"
       GPU_ARCH_TYPE: cuda
       SKIP_ALL_TESTS: 1
       DESIRED_PYTHON: "3.13"
@@ -4347,7 +4226,7 @@ jobs:
           echo "instance-type: $(get_ec2_metadata instance-type)"
           echo "system info $(uname -a)"
       - name: "[FB EMPLOYEES] Enable SSH (Click me for login details)"
-        uses: pytorch/test-infra/.github/actions/setup-ssh@main
+        uses: pytorch/test-infra/.github/actions/setup-ssh@release/2.12
         continue-on-error: true
         with:
           github-secret: ${{ secrets.GITHUB_TOKEN }}
@@ -4379,58 +4258,53 @@ jobs:
       - name: Checkout PyTorch
         uses: actions/checkout@v4
         with:
-          ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
           submodules: recursive
-          path: pytorch
           show-progress: false
       - name: Clean PyTorch checkout
         run: |
           # Remove any artifacts from the previous checkouts
           git clean -fxd
-        working-directory: pytorch
       - name: Populate binary env
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_populate_env.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_populate_env.sh"
       - name: Build PyTorch binary
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_windows_build.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_windows_build.sh"
       - uses: actions/upload-artifact@v4.4.0
         if: always()
         with:
-          name: wheel-py3_13-cuda13_0
+          name: wheel-py3_13-cuda13_2
           retention-days: 14
           if-no-files-found: error
           path: "${{ env.PYTORCH_FINAL_PACKAGE_DIR }}"
       - name: Wait until all sessions have drained
         shell: powershell
-        working-directory: pytorch
         if: always()
         timeout-minutes: 120
         run: |
           .github\scripts\wait_for_ssh_to_drain.ps1
       - name: Kill active ssh sessions if still around (Useful if workflow was cancelled)
         shell: powershell
-        working-directory: pytorch
         if: always()
         run: |
           .github\scripts\kill_active_ssh_sessions.ps1
 
-  wheel-py3_13-cuda13_0-test:  # Testing
+  wheel-py3_13-cuda13_2-test:  # Testing
     if: ${{ github.repository_owner == 'pytorch' }}
     needs:
-      - wheel-py3_13-cuda13_0-build
+      - wheel-py3_13-cuda13_2-build
       - get-label-type
     runs-on: "${{ needs.get-label-type.outputs.label-type }}windows.g4dn.xlarge"
     timeout-minutes: 360
     env:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
+      PYTORCH_ROOT: ${{ github.workspace }}
       PACKAGE_TYPE: wheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu130
-      GPU_ARCH_VERSION: "13.0"
+      DESIRED_CUDA: cu132
+      GPU_ARCH_VERSION: "13.2"
       GPU_ARCH_TYPE: cuda
       SKIP_ALL_TESTS: 1
       DESIRED_PYTHON: "3.13"
@@ -4450,7 +4324,7 @@ jobs:
           echo "instance-type: $(get_ec2_metadata instance-type)"
           echo "system info $(uname -a)"
       - name: "[FB EMPLOYEES] Enable SSH (Click me for login details)"
-        uses: pytorch/test-infra/.github/actions/setup-ssh@main
+        uses: pytorch/test-infra/.github/actions/setup-ssh@release/2.12
         continue-on-error: true
         with:
           github-secret: ${{ secrets.GITHUB_TOKEN }}
@@ -4482,15 +4356,12 @@ jobs:
       - name: Checkout PyTorch
         uses: actions/checkout@v4
         with:
-          ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
           submodules: recursive
-          path: pytorch
           show-progress: false
       - name: Clean PyTorch checkout
         run: |
           # Remove any artifacts from the previous checkouts
           git clean -fxd
-        working-directory: pytorch
       # NOTE: These environment variables are put here so that they can be applied on every job equally
       #       They are also here because setting them at a workflow level doesn't give us access to the
       #       runner.temp variable, which we need.
@@ -4503,47 +4374,48 @@ jobs:
       - uses: actions/download-artifact@v4.1.7
         name: Download Build Artifacts
         with:
-          name: wheel-py3_13-cuda13_0
+          name: wheel-py3_13-cuda13_2
           path: "${{ env.PYTORCH_FINAL_PACKAGE_DIR }}"
       - name: Populate binary env
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_populate_env.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_populate_env.sh"
       - name: Test PyTorch binary
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_windows_test.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_windows_test.sh"
       - name: Wait until all sessions have drained
         shell: powershell
-        working-directory: pytorch
         if: always()
         timeout-minutes: 120
         run: |
           .github\scripts\wait_for_ssh_to_drain.ps1
       - name: Kill active ssh sessions if still around (Useful if workflow was cancelled)
         shell: powershell
-        working-directory: pytorch
         if: always()
         run: |
           .github\scripts\kill_active_ssh_sessions.ps1
-  wheel-py3_13-cuda13_0-upload:  # Uploading
+  wheel-py3_13-cuda13_2-upload:  # Uploading
     if: ${{ github.repository_owner == 'pytorch' }}
     permissions:
       id-token: write
       contents: read
-    needs: wheel-py3_13-cuda13_0-test
+    needs: wheel-py3_13-cuda13_2-test
     with:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
+      PYTORCH_ROOT: ${{ github.workspace }}
       PACKAGE_TYPE: wheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu130
-      GPU_ARCH_VERSION: "13.0"
+      DESIRED_CUDA: cu132
+      GPU_ARCH_VERSION: "13.2"
       GPU_ARCH_TYPE: cuda
       DESIRED_PYTHON: "3.13"
-      build_name: wheel-py3_13-cuda13_0
+      build_name: wheel-py3_13-cuda13_2
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
     uses: ./.github/workflows/_binary-upload.yml
   wheel-py3_13-xpu-build:
     if: ${{ github.repository_owner == 'pytorch' }}
@@ -4551,7 +4423,7 @@ jobs:
     runs-on: "${{ needs.get-label-type.outputs.label-type }}windows.12xlarge"
     timeout-minutes: 360
     env:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
+      PYTORCH_ROOT: ${{ github.workspace }}
       PACKAGE_TYPE: wheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
@@ -4585,7 +4457,7 @@ jobs:
           echo "instance-type: $(get_ec2_metadata instance-type)"
           echo "system info $(uname -a)"
       - name: "[FB EMPLOYEES] Enable SSH (Click me for login details)"
-        uses: pytorch/test-infra/.github/actions/setup-ssh@main
+        uses: pytorch/test-infra/.github/actions/setup-ssh@release/2.12
         continue-on-error: true
         with:
           github-secret: ${{ secrets.GITHUB_TOKEN }}
@@ -4617,23 +4489,20 @@ jobs:
       - name: Checkout PyTorch
         uses: actions/checkout@v4
         with:
-          ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
           submodules: recursive
-          path: pytorch
           show-progress: false
       - name: Clean PyTorch checkout
         run: |
           # Remove any artifacts from the previous checkouts
           git clean -fxd
-        working-directory: pytorch
       - name: Populate binary env
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_populate_env.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_populate_env.sh"
       - name: Build PyTorch binary
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_windows_build.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_windows_build.sh"
       - uses: actions/upload-artifact@v4.4.0
         if: always()
         with:
@@ -4643,14 +4512,12 @@ jobs:
           path: "${{ env.PYTORCH_FINAL_PACKAGE_DIR }}"
       - name: Wait until all sessions have drained
         shell: powershell
-        working-directory: pytorch
         if: always()
         timeout-minutes: 120
         run: |
           .github\scripts\wait_for_ssh_to_drain.ps1
       - name: Kill active ssh sessions if still around (Useful if workflow was cancelled)
         shell: powershell
-        working-directory: pytorch
         if: always()
         run: |
           .github\scripts\kill_active_ssh_sessions.ps1
@@ -4663,7 +4530,7 @@ jobs:
     runs-on: "${{ needs.get-label-type.outputs.label-type }}windows.4xlarge"
     timeout-minutes: 360
     env:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
+      PYTORCH_ROOT: ${{ github.workspace }}
       PACKAGE_TYPE: wheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
@@ -4687,7 +4554,7 @@ jobs:
           echo "instance-type: $(get_ec2_metadata instance-type)"
           echo "system info $(uname -a)"
       - name: "[FB EMPLOYEES] Enable SSH (Click me for login details)"
-        uses: pytorch/test-infra/.github/actions/setup-ssh@main
+        uses: pytorch/test-infra/.github/actions/setup-ssh@release/2.12
         continue-on-error: true
         with:
           github-secret: ${{ secrets.GITHUB_TOKEN }}
@@ -4719,15 +4586,12 @@ jobs:
       - name: Checkout PyTorch
         uses: actions/checkout@v4
         with:
-          ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
           submodules: recursive
-          path: pytorch
           show-progress: false
       - name: Clean PyTorch checkout
         run: |
           # Remove any artifacts from the previous checkouts
           git clean -fxd
-        working-directory: pytorch
       # NOTE: These environment variables are put here so that they can be applied on every job equally
       #       They are also here because setting them at a workflow level doesn't give us access to the
       #       runner.temp variable, which we need.
@@ -4745,21 +4609,19 @@ jobs:
       - name: Populate binary env
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_populate_env.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_populate_env.sh"
       - name: Test PyTorch binary
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_windows_test.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_windows_test.sh"
       - name: Wait until all sessions have drained
         shell: powershell
-        working-directory: pytorch
         if: always()
         timeout-minutes: 120
         run: |
           .github\scripts\wait_for_ssh_to_drain.ps1
       - name: Kill active ssh sessions if still around (Useful if workflow was cancelled)
         shell: powershell
-        working-directory: pytorch
         if: always()
         run: |
           .github\scripts\kill_active_ssh_sessions.ps1
@@ -4770,7 +4632,7 @@ jobs:
       contents: read
     needs: wheel-py3_13-xpu-test
     with:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
+      PYTORCH_ROOT: ${{ github.workspace }}
       PACKAGE_TYPE: wheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
@@ -4780,6 +4642,9 @@ jobs:
       build_name: wheel-py3_13-xpu
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
     uses: ./.github/workflows/_binary-upload.yml
   wheel-py3_13t-cpu-build:
     if: ${{ github.repository_owner == 'pytorch' }}
@@ -4787,7 +4652,7 @@ jobs:
     runs-on: "${{ needs.get-label-type.outputs.label-type }}windows.12xlarge"
     timeout-minutes: 360
     env:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
+      PYTORCH_ROOT: ${{ github.workspace }}
       PACKAGE_TYPE: wheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
@@ -4795,6 +4660,7 @@ jobs:
       GPU_ARCH_TYPE: cpu
       SKIP_ALL_TESTS: 1
       DESIRED_PYTHON: "3.13t"
+      PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cufile,nvjitlink,nvtx]==13.0.2; platform_system == 'Linux' | nvidia-cublas>=13.1.0.3,<=13.1.1.3; platform_system == 'Linux' | cuda-bindings>=13.0.3,<14; platform_system == 'Linux' | nvidia-cudnn-cu13==9.20.0.48; platform_system == 'Linux' | nvidia-cusparselt-cu13==0.8.1; platform_system == 'Linux' | nvidia-nccl-cu13==2.29.7; platform_system == 'Linux' | nvidia-nvshmem-cu13==3.4.5; platform_system == 'Linux'
     steps:
       # NOTE: These environment variables are put here so that they can be applied on every job equally
       #       They are also here because setting them at a workflow level doesn't give us access to the
@@ -4820,7 +4686,7 @@ jobs:
           echo "instance-type: $(get_ec2_metadata instance-type)"
           echo "system info $(uname -a)"
       - name: "[FB EMPLOYEES] Enable SSH (Click me for login details)"
-        uses: pytorch/test-infra/.github/actions/setup-ssh@main
+        uses: pytorch/test-infra/.github/actions/setup-ssh@release/2.12
         continue-on-error: true
         with:
           github-secret: ${{ secrets.GITHUB_TOKEN }}
@@ -4852,23 +4718,20 @@ jobs:
       - name: Checkout PyTorch
         uses: actions/checkout@v4
         with:
-          ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
           submodules: recursive
-          path: pytorch
           show-progress: false
       - name: Clean PyTorch checkout
         run: |
           # Remove any artifacts from the previous checkouts
           git clean -fxd
-        working-directory: pytorch
       - name: Populate binary env
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_populate_env.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_populate_env.sh"
       - name: Build PyTorch binary
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_windows_build.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_windows_build.sh"
       - uses: actions/upload-artifact@v4.4.0
         if: always()
         with:
@@ -4878,14 +4741,12 @@ jobs:
           path: "${{ env.PYTORCH_FINAL_PACKAGE_DIR }}"
       - name: Wait until all sessions have drained
         shell: powershell
-        working-directory: pytorch
         if: always()
         timeout-minutes: 120
         run: |
           .github\scripts\wait_for_ssh_to_drain.ps1
       - name: Kill active ssh sessions if still around (Useful if workflow was cancelled)
         shell: powershell
-        working-directory: pytorch
         if: always()
         run: |
           .github\scripts\kill_active_ssh_sessions.ps1
@@ -4898,7 +4759,7 @@ jobs:
     runs-on: "${{ needs.get-label-type.outputs.label-type }}windows.4xlarge"
     timeout-minutes: 360
     env:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
+      PYTORCH_ROOT: ${{ github.workspace }}
       PACKAGE_TYPE: wheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
@@ -4922,7 +4783,7 @@ jobs:
           echo "instance-type: $(get_ec2_metadata instance-type)"
           echo "system info $(uname -a)"
       - name: "[FB EMPLOYEES] Enable SSH (Click me for login details)"
-        uses: pytorch/test-infra/.github/actions/setup-ssh@main
+        uses: pytorch/test-infra/.github/actions/setup-ssh@release/2.12
         continue-on-error: true
         with:
           github-secret: ${{ secrets.GITHUB_TOKEN }}
@@ -4954,15 +4815,12 @@ jobs:
       - name: Checkout PyTorch
         uses: actions/checkout@v4
         with:
-          ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
           submodules: recursive
-          path: pytorch
           show-progress: false
       - name: Clean PyTorch checkout
         run: |
           # Remove any artifacts from the previous checkouts
           git clean -fxd
-        working-directory: pytorch
       # NOTE: These environment variables are put here so that they can be applied on every job equally
       #       They are also here because setting them at a workflow level doesn't give us access to the
       #       runner.temp variable, which we need.
@@ -4980,21 +4838,19 @@ jobs:
       - name: Populate binary env
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_populate_env.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_populate_env.sh"
       - name: Test PyTorch binary
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_windows_test.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_windows_test.sh"
       - name: Wait until all sessions have drained
         shell: powershell
-        working-directory: pytorch
         if: always()
         timeout-minutes: 120
         run: |
           .github\scripts\wait_for_ssh_to_drain.ps1
       - name: Kill active ssh sessions if still around (Useful if workflow was cancelled)
         shell: powershell
-        working-directory: pytorch
         if: always()
         run: |
           .github\scripts\kill_active_ssh_sessions.ps1
@@ -5005,7 +4861,7 @@ jobs:
       contents: read
     needs: wheel-py3_13t-cpu-test
     with:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
+      PYTORCH_ROOT: ${{ github.workspace }}
       PACKAGE_TYPE: wheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
@@ -5015,6 +4871,9 @@ jobs:
       build_name: wheel-py3_13t-cpu
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
     uses: ./.github/workflows/_binary-upload.yml
   wheel-py3_13t-cuda12_6-build:
     if: ${{ github.repository_owner == 'pytorch' }}
@@ -5022,7 +4881,7 @@ jobs:
     runs-on: "${{ needs.get-label-type.outputs.label-type }}windows.12xlarge"
     timeout-minutes: 360
     env:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
+      PYTORCH_ROOT: ${{ github.workspace }}
       PACKAGE_TYPE: wheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
@@ -5056,7 +4915,7 @@ jobs:
           echo "instance-type: $(get_ec2_metadata instance-type)"
           echo "system info $(uname -a)"
       - name: "[FB EMPLOYEES] Enable SSH (Click me for login details)"
-        uses: pytorch/test-infra/.github/actions/setup-ssh@main
+        uses: pytorch/test-infra/.github/actions/setup-ssh@release/2.12
         continue-on-error: true
         with:
           github-secret: ${{ secrets.GITHUB_TOKEN }}
@@ -5088,23 +4947,20 @@ jobs:
       - name: Checkout PyTorch
         uses: actions/checkout@v4
         with:
-          ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
           submodules: recursive
-          path: pytorch
           show-progress: false
       - name: Clean PyTorch checkout
         run: |
           # Remove any artifacts from the previous checkouts
           git clean -fxd
-        working-directory: pytorch
       - name: Populate binary env
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_populate_env.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_populate_env.sh"
       - name: Build PyTorch binary
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_windows_build.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_windows_build.sh"
       - uses: actions/upload-artifact@v4.4.0
         if: always()
         with:
@@ -5114,14 +4970,12 @@ jobs:
           path: "${{ env.PYTORCH_FINAL_PACKAGE_DIR }}"
       - name: Wait until all sessions have drained
         shell: powershell
-        working-directory: pytorch
         if: always()
         timeout-minutes: 120
         run: |
           .github\scripts\wait_for_ssh_to_drain.ps1
       - name: Kill active ssh sessions if still around (Useful if workflow was cancelled)
         shell: powershell
-        working-directory: pytorch
         if: always()
         run: |
           .github\scripts\kill_active_ssh_sessions.ps1
@@ -5134,7 +4988,7 @@ jobs:
     runs-on: "${{ needs.get-label-type.outputs.label-type }}windows.g4dn.xlarge"
     timeout-minutes: 360
     env:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
+      PYTORCH_ROOT: ${{ github.workspace }}
       PACKAGE_TYPE: wheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
@@ -5159,7 +5013,7 @@ jobs:
           echo "instance-type: $(get_ec2_metadata instance-type)"
           echo "system info $(uname -a)"
       - name: "[FB EMPLOYEES] Enable SSH (Click me for login details)"
-        uses: pytorch/test-infra/.github/actions/setup-ssh@main
+        uses: pytorch/test-infra/.github/actions/setup-ssh@release/2.12
         continue-on-error: true
         with:
           github-secret: ${{ secrets.GITHUB_TOKEN }}
@@ -5191,15 +5045,12 @@ jobs:
       - name: Checkout PyTorch
         uses: actions/checkout@v4
         with:
-          ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
           submodules: recursive
-          path: pytorch
           show-progress: false
       - name: Clean PyTorch checkout
         run: |
           # Remove any artifacts from the previous checkouts
           git clean -fxd
-        working-directory: pytorch
       # NOTE: These environment variables are put here so that they can be applied on every job equally
       #       They are also here because setting them at a workflow level doesn't give us access to the
       #       runner.temp variable, which we need.
@@ -5217,21 +5068,19 @@ jobs:
       - name: Populate binary env
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_populate_env.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_populate_env.sh"
       - name: Test PyTorch binary
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_windows_test.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_windows_test.sh"
       - name: Wait until all sessions have drained
         shell: powershell
-        working-directory: pytorch
         if: always()
         timeout-minutes: 120
         run: |
           .github\scripts\wait_for_ssh_to_drain.ps1
       - name: Kill active ssh sessions if still around (Useful if workflow was cancelled)
         shell: powershell
-        working-directory: pytorch
         if: always()
         run: |
           .github\scripts\kill_active_ssh_sessions.ps1
@@ -5242,7 +5091,7 @@ jobs:
       contents: read
     needs: wheel-py3_13t-cuda12_6-test
     with:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
+      PYTORCH_ROOT: ${{ github.workspace }}
       PACKAGE_TYPE: wheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
@@ -5253,19 +5102,22 @@ jobs:
       build_name: wheel-py3_13t-cuda12_6
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
     uses: ./.github/workflows/_binary-upload.yml
-  wheel-py3_13t-cuda12_8-build:
+  wheel-py3_13t-cuda13_0-build:
     if: ${{ github.repository_owner == 'pytorch' }}
     needs: get-label-type
     runs-on: "${{ needs.get-label-type.outputs.label-type }}windows.12xlarge"
     timeout-minutes: 360
     env:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
+      PYTORCH_ROOT: ${{ github.workspace }}
       PACKAGE_TYPE: wheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu128
-      GPU_ARCH_VERSION: "12.8"
+      DESIRED_CUDA: cu130
+      GPU_ARCH_VERSION: "13.0"
       GPU_ARCH_TYPE: cuda
       SKIP_ALL_TESTS: 1
       DESIRED_PYTHON: "3.13t"
@@ -5294,7 +5146,7 @@ jobs:
           echo "instance-type: $(get_ec2_metadata instance-type)"
           echo "system info $(uname -a)"
       - name: "[FB EMPLOYEES] Enable SSH (Click me for login details)"
-        uses: pytorch/test-infra/.github/actions/setup-ssh@main
+        uses: pytorch/test-infra/.github/actions/setup-ssh@release/2.12
         continue-on-error: true
         with:
           github-secret: ${{ secrets.GITHUB_TOKEN }}
@@ -5326,58 +5178,53 @@ jobs:
       - name: Checkout PyTorch
         uses: actions/checkout@v4
         with:
-          ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
           submodules: recursive
-          path: pytorch
           show-progress: false
       - name: Clean PyTorch checkout
         run: |
           # Remove any artifacts from the previous checkouts
           git clean -fxd
-        working-directory: pytorch
       - name: Populate binary env
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_populate_env.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_populate_env.sh"
       - name: Build PyTorch binary
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_windows_build.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_windows_build.sh"
       - uses: actions/upload-artifact@v4.4.0
         if: always()
         with:
-          name: wheel-py3_13t-cuda12_8
+          name: wheel-py3_13t-cuda13_0
           retention-days: 14
           if-no-files-found: error
           path: "${{ env.PYTORCH_FINAL_PACKAGE_DIR }}"
       - name: Wait until all sessions have drained
         shell: powershell
-        working-directory: pytorch
         if: always()
         timeout-minutes: 120
         run: |
           .github\scripts\wait_for_ssh_to_drain.ps1
       - name: Kill active ssh sessions if still around (Useful if workflow was cancelled)
         shell: powershell
-        working-directory: pytorch
         if: always()
         run: |
           .github\scripts\kill_active_ssh_sessions.ps1
 
-  wheel-py3_13t-cuda12_8-test:  # Testing
+  wheel-py3_13t-cuda13_0-test:  # Testing
     if: ${{ github.repository_owner == 'pytorch' }}
     needs:
-      - wheel-py3_13t-cuda12_8-build
+      - wheel-py3_13t-cuda13_0-build
       - get-label-type
     runs-on: "${{ needs.get-label-type.outputs.label-type }}windows.g4dn.xlarge"
     timeout-minutes: 360
     env:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
+      PYTORCH_ROOT: ${{ github.workspace }}
       PACKAGE_TYPE: wheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu128
-      GPU_ARCH_VERSION: "12.8"
+      DESIRED_CUDA: cu130
+      GPU_ARCH_VERSION: "13.0"
       GPU_ARCH_TYPE: cuda
       SKIP_ALL_TESTS: 1
       DESIRED_PYTHON: "3.13t"
@@ -5397,7 +5244,7 @@ jobs:
           echo "instance-type: $(get_ec2_metadata instance-type)"
           echo "system info $(uname -a)"
       - name: "[FB EMPLOYEES] Enable SSH (Click me for login details)"
-        uses: pytorch/test-infra/.github/actions/setup-ssh@main
+        uses: pytorch/test-infra/.github/actions/setup-ssh@release/2.12
         continue-on-error: true
         with:
           github-secret: ${{ secrets.GITHUB_TOKEN }}
@@ -5429,15 +5276,12 @@ jobs:
       - name: Checkout PyTorch
         uses: actions/checkout@v4
         with:
-          ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
           submodules: recursive
-          path: pytorch
           show-progress: false
       - name: Clean PyTorch checkout
         run: |
           # Remove any artifacts from the previous checkouts
           git clean -fxd
-        working-directory: pytorch
       # NOTE: These environment variables are put here so that they can be applied on every job equally
       #       They are also here because setting them at a workflow level doesn't give us access to the
       #       runner.temp variable, which we need.
@@ -5450,60 +5294,61 @@ jobs:
       - uses: actions/download-artifact@v4.1.7
         name: Download Build Artifacts
         with:
-          name: wheel-py3_13t-cuda12_8
+          name: wheel-py3_13t-cuda13_0
           path: "${{ env.PYTORCH_FINAL_PACKAGE_DIR }}"
       - name: Populate binary env
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_populate_env.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_populate_env.sh"
       - name: Test PyTorch binary
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_windows_test.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_windows_test.sh"
       - name: Wait until all sessions have drained
         shell: powershell
-        working-directory: pytorch
         if: always()
         timeout-minutes: 120
         run: |
           .github\scripts\wait_for_ssh_to_drain.ps1
       - name: Kill active ssh sessions if still around (Useful if workflow was cancelled)
         shell: powershell
-        working-directory: pytorch
         if: always()
         run: |
           .github\scripts\kill_active_ssh_sessions.ps1
-  wheel-py3_13t-cuda12_8-upload:  # Uploading
+  wheel-py3_13t-cuda13_0-upload:  # Uploading
     if: ${{ github.repository_owner == 'pytorch' }}
     permissions:
       id-token: write
       contents: read
-    needs: wheel-py3_13t-cuda12_8-test
+    needs: wheel-py3_13t-cuda13_0-test
     with:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
+      PYTORCH_ROOT: ${{ github.workspace }}
       PACKAGE_TYPE: wheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu128
-      GPU_ARCH_VERSION: "12.8"
+      DESIRED_CUDA: cu130
+      GPU_ARCH_VERSION: "13.0"
       GPU_ARCH_TYPE: cuda
       DESIRED_PYTHON: "3.13t"
-      build_name: wheel-py3_13t-cuda12_8
+      build_name: wheel-py3_13t-cuda13_0
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
     uses: ./.github/workflows/_binary-upload.yml
-  wheel-py3_13t-cuda13_0-build:
+  wheel-py3_13t-cuda13_2-build:
     if: ${{ github.repository_owner == 'pytorch' }}
     needs: get-label-type
     runs-on: "${{ needs.get-label-type.outputs.label-type }}windows.12xlarge"
     timeout-minutes: 360
     env:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
+      PYTORCH_ROOT: ${{ github.workspace }}
       PACKAGE_TYPE: wheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu130
-      GPU_ARCH_VERSION: "13.0"
+      DESIRED_CUDA: cu132
+      GPU_ARCH_VERSION: "13.2"
       GPU_ARCH_TYPE: cuda
       SKIP_ALL_TESTS: 1
       DESIRED_PYTHON: "3.13t"
@@ -5532,7 +5377,7 @@ jobs:
           echo "instance-type: $(get_ec2_metadata instance-type)"
           echo "system info $(uname -a)"
       - name: "[FB EMPLOYEES] Enable SSH (Click me for login details)"
-        uses: pytorch/test-infra/.github/actions/setup-ssh@main
+        uses: pytorch/test-infra/.github/actions/setup-ssh@release/2.12
         continue-on-error: true
         with:
           github-secret: ${{ secrets.GITHUB_TOKEN }}
@@ -5564,58 +5409,53 @@ jobs:
       - name: Checkout PyTorch
         uses: actions/checkout@v4
         with:
-          ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
           submodules: recursive
-          path: pytorch
           show-progress: false
       - name: Clean PyTorch checkout
         run: |
           # Remove any artifacts from the previous checkouts
           git clean -fxd
-        working-directory: pytorch
       - name: Populate binary env
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_populate_env.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_populate_env.sh"
       - name: Build PyTorch binary
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_windows_build.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_windows_build.sh"
       - uses: actions/upload-artifact@v4.4.0
         if: always()
         with:
-          name: wheel-py3_13t-cuda13_0
+          name: wheel-py3_13t-cuda13_2
           retention-days: 14
           if-no-files-found: error
           path: "${{ env.PYTORCH_FINAL_PACKAGE_DIR }}"
       - name: Wait until all sessions have drained
         shell: powershell
-        working-directory: pytorch
         if: always()
         timeout-minutes: 120
         run: |
           .github\scripts\wait_for_ssh_to_drain.ps1
       - name: Kill active ssh sessions if still around (Useful if workflow was cancelled)
         shell: powershell
-        working-directory: pytorch
         if: always()
         run: |
           .github\scripts\kill_active_ssh_sessions.ps1
 
-  wheel-py3_13t-cuda13_0-test:  # Testing
+  wheel-py3_13t-cuda13_2-test:  # Testing
     if: ${{ github.repository_owner == 'pytorch' }}
     needs:
-      - wheel-py3_13t-cuda13_0-build
+      - wheel-py3_13t-cuda13_2-build
       - get-label-type
     runs-on: "${{ needs.get-label-type.outputs.label-type }}windows.g4dn.xlarge"
     timeout-minutes: 360
     env:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
+      PYTORCH_ROOT: ${{ github.workspace }}
       PACKAGE_TYPE: wheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu130
-      GPU_ARCH_VERSION: "13.0"
+      DESIRED_CUDA: cu132
+      GPU_ARCH_VERSION: "13.2"
       GPU_ARCH_TYPE: cuda
       SKIP_ALL_TESTS: 1
       DESIRED_PYTHON: "3.13t"
@@ -5635,7 +5475,7 @@ jobs:
           echo "instance-type: $(get_ec2_metadata instance-type)"
           echo "system info $(uname -a)"
       - name: "[FB EMPLOYEES] Enable SSH (Click me for login details)"
-        uses: pytorch/test-infra/.github/actions/setup-ssh@main
+        uses: pytorch/test-infra/.github/actions/setup-ssh@release/2.12
         continue-on-error: true
         with:
           github-secret: ${{ secrets.GITHUB_TOKEN }}
@@ -5667,15 +5507,12 @@ jobs:
       - name: Checkout PyTorch
         uses: actions/checkout@v4
         with:
-          ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
           submodules: recursive
-          path: pytorch
           show-progress: false
       - name: Clean PyTorch checkout
         run: |
           # Remove any artifacts from the previous checkouts
           git clean -fxd
-        working-directory: pytorch
       # NOTE: These environment variables are put here so that they can be applied on every job equally
       #       They are also here because setting them at a workflow level doesn't give us access to the
       #       runner.temp variable, which we need.
@@ -5688,47 +5525,48 @@ jobs:
       - uses: actions/download-artifact@v4.1.7
         name: Download Build Artifacts
         with:
-          name: wheel-py3_13t-cuda13_0
+          name: wheel-py3_13t-cuda13_2
           path: "${{ env.PYTORCH_FINAL_PACKAGE_DIR }}"
       - name: Populate binary env
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_populate_env.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_populate_env.sh"
       - name: Test PyTorch binary
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_windows_test.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_windows_test.sh"
       - name: Wait until all sessions have drained
         shell: powershell
-        working-directory: pytorch
         if: always()
         timeout-minutes: 120
         run: |
           .github\scripts\wait_for_ssh_to_drain.ps1
       - name: Kill active ssh sessions if still around (Useful if workflow was cancelled)
         shell: powershell
-        working-directory: pytorch
         if: always()
         run: |
           .github\scripts\kill_active_ssh_sessions.ps1
-  wheel-py3_13t-cuda13_0-upload:  # Uploading
+  wheel-py3_13t-cuda13_2-upload:  # Uploading
     if: ${{ github.repository_owner == 'pytorch' }}
     permissions:
       id-token: write
       contents: read
-    needs: wheel-py3_13t-cuda13_0-test
+    needs: wheel-py3_13t-cuda13_2-test
     with:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
+      PYTORCH_ROOT: ${{ github.workspace }}
       PACKAGE_TYPE: wheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu130
-      GPU_ARCH_VERSION: "13.0"
+      DESIRED_CUDA: cu132
+      GPU_ARCH_VERSION: "13.2"
       GPU_ARCH_TYPE: cuda
       DESIRED_PYTHON: "3.13t"
-      build_name: wheel-py3_13t-cuda13_0
+      build_name: wheel-py3_13t-cuda13_2
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
     uses: ./.github/workflows/_binary-upload.yml
   wheel-py3_13t-xpu-build:
     if: ${{ github.repository_owner == 'pytorch' }}
@@ -5736,7 +5574,7 @@ jobs:
     runs-on: "${{ needs.get-label-type.outputs.label-type }}windows.12xlarge"
     timeout-minutes: 360
     env:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
+      PYTORCH_ROOT: ${{ github.workspace }}
       PACKAGE_TYPE: wheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
@@ -5770,7 +5608,7 @@ jobs:
           echo "instance-type: $(get_ec2_metadata instance-type)"
           echo "system info $(uname -a)"
       - name: "[FB EMPLOYEES] Enable SSH (Click me for login details)"
-        uses: pytorch/test-infra/.github/actions/setup-ssh@main
+        uses: pytorch/test-infra/.github/actions/setup-ssh@release/2.12
         continue-on-error: true
         with:
           github-secret: ${{ secrets.GITHUB_TOKEN }}
@@ -5802,23 +5640,20 @@ jobs:
       - name: Checkout PyTorch
         uses: actions/checkout@v4
         with:
-          ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
           submodules: recursive
-          path: pytorch
           show-progress: false
       - name: Clean PyTorch checkout
         run: |
           # Remove any artifacts from the previous checkouts
           git clean -fxd
-        working-directory: pytorch
       - name: Populate binary env
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_populate_env.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_populate_env.sh"
       - name: Build PyTorch binary
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_windows_build.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_windows_build.sh"
       - uses: actions/upload-artifact@v4.4.0
         if: always()
         with:
@@ -5828,14 +5663,12 @@ jobs:
           path: "${{ env.PYTORCH_FINAL_PACKAGE_DIR }}"
       - name: Wait until all sessions have drained
         shell: powershell
-        working-directory: pytorch
         if: always()
         timeout-minutes: 120
         run: |
           .github\scripts\wait_for_ssh_to_drain.ps1
       - name: Kill active ssh sessions if still around (Useful if workflow was cancelled)
         shell: powershell
-        working-directory: pytorch
         if: always()
         run: |
           .github\scripts\kill_active_ssh_sessions.ps1
@@ -5848,7 +5681,7 @@ jobs:
     runs-on: "${{ needs.get-label-type.outputs.label-type }}windows.4xlarge"
     timeout-minutes: 360
     env:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
+      PYTORCH_ROOT: ${{ github.workspace }}
       PACKAGE_TYPE: wheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
@@ -5872,7 +5705,7 @@ jobs:
           echo "instance-type: $(get_ec2_metadata instance-type)"
           echo "system info $(uname -a)"
       - name: "[FB EMPLOYEES] Enable SSH (Click me for login details)"
-        uses: pytorch/test-infra/.github/actions/setup-ssh@main
+        uses: pytorch/test-infra/.github/actions/setup-ssh@release/2.12
         continue-on-error: true
         with:
           github-secret: ${{ secrets.GITHUB_TOKEN }}
@@ -5904,15 +5737,12 @@ jobs:
       - name: Checkout PyTorch
         uses: actions/checkout@v4
         with:
-          ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
           submodules: recursive
-          path: pytorch
           show-progress: false
       - name: Clean PyTorch checkout
         run: |
           # Remove any artifacts from the previous checkouts
           git clean -fxd
-        working-directory: pytorch
       # NOTE: These environment variables are put here so that they can be applied on every job equally
       #       They are also here because setting them at a workflow level doesn't give us access to the
       #       runner.temp variable, which we need.
@@ -5930,21 +5760,19 @@ jobs:
       - name: Populate binary env
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_populate_env.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_populate_env.sh"
       - name: Test PyTorch binary
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_windows_test.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_windows_test.sh"
       - name: Wait until all sessions have drained
         shell: powershell
-        working-directory: pytorch
         if: always()
         timeout-minutes: 120
         run: |
           .github\scripts\wait_for_ssh_to_drain.ps1
       - name: Kill active ssh sessions if still around (Useful if workflow was cancelled)
         shell: powershell
-        working-directory: pytorch
         if: always()
         run: |
           .github\scripts\kill_active_ssh_sessions.ps1
@@ -5955,7 +5783,7 @@ jobs:
       contents: read
     needs: wheel-py3_13t-xpu-test
     with:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
+      PYTORCH_ROOT: ${{ github.workspace }}
       PACKAGE_TYPE: wheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
@@ -5965,6 +5793,9 @@ jobs:
       build_name: wheel-py3_13t-xpu
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
     uses: ./.github/workflows/_binary-upload.yml
   wheel-py3_14-cpu-build:
     if: ${{ github.repository_owner == 'pytorch' }}
@@ -5972,7 +5803,7 @@ jobs:
     runs-on: "${{ needs.get-label-type.outputs.label-type }}windows.12xlarge"
     timeout-minutes: 360
     env:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
+      PYTORCH_ROOT: ${{ github.workspace }}
       PACKAGE_TYPE: wheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
@@ -5980,6 +5811,7 @@ jobs:
       GPU_ARCH_TYPE: cpu
       SKIP_ALL_TESTS: 1
       DESIRED_PYTHON: "3.14"
+      PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cufile,nvjitlink,nvtx]==13.0.2; platform_system == 'Linux' | nvidia-cublas>=13.1.0.3,<=13.1.1.3; platform_system == 'Linux' | cuda-bindings>=13.0.3,<14; platform_system == 'Linux' | nvidia-cudnn-cu13==9.20.0.48; platform_system == 'Linux' | nvidia-cusparselt-cu13==0.8.1; platform_system == 'Linux' | nvidia-nccl-cu13==2.29.7; platform_system == 'Linux' | nvidia-nvshmem-cu13==3.4.5; platform_system == 'Linux'
     steps:
       # NOTE: These environment variables are put here so that they can be applied on every job equally
       #       They are also here because setting them at a workflow level doesn't give us access to the
@@ -6005,7 +5837,7 @@ jobs:
           echo "instance-type: $(get_ec2_metadata instance-type)"
           echo "system info $(uname -a)"
       - name: "[FB EMPLOYEES] Enable SSH (Click me for login details)"
-        uses: pytorch/test-infra/.github/actions/setup-ssh@main
+        uses: pytorch/test-infra/.github/actions/setup-ssh@release/2.12
         continue-on-error: true
         with:
           github-secret: ${{ secrets.GITHUB_TOKEN }}
@@ -6037,23 +5869,20 @@ jobs:
       - name: Checkout PyTorch
         uses: actions/checkout@v4
         with:
-          ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
           submodules: recursive
-          path: pytorch
           show-progress: false
       - name: Clean PyTorch checkout
         run: |
           # Remove any artifacts from the previous checkouts
           git clean -fxd
-        working-directory: pytorch
       - name: Populate binary env
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_populate_env.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_populate_env.sh"
       - name: Build PyTorch binary
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_windows_build.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_windows_build.sh"
       - uses: actions/upload-artifact@v4.4.0
         if: always()
         with:
@@ -6063,14 +5892,12 @@ jobs:
           path: "${{ env.PYTORCH_FINAL_PACKAGE_DIR }}"
       - name: Wait until all sessions have drained
         shell: powershell
-        working-directory: pytorch
         if: always()
         timeout-minutes: 120
         run: |
           .github\scripts\wait_for_ssh_to_drain.ps1
       - name: Kill active ssh sessions if still around (Useful if workflow was cancelled)
         shell: powershell
-        working-directory: pytorch
         if: always()
         run: |
           .github\scripts\kill_active_ssh_sessions.ps1
@@ -6083,7 +5910,7 @@ jobs:
     runs-on: "${{ needs.get-label-type.outputs.label-type }}windows.4xlarge"
     timeout-minutes: 360
     env:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
+      PYTORCH_ROOT: ${{ github.workspace }}
       PACKAGE_TYPE: wheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
@@ -6107,7 +5934,7 @@ jobs:
           echo "instance-type: $(get_ec2_metadata instance-type)"
           echo "system info $(uname -a)"
       - name: "[FB EMPLOYEES] Enable SSH (Click me for login details)"
-        uses: pytorch/test-infra/.github/actions/setup-ssh@main
+        uses: pytorch/test-infra/.github/actions/setup-ssh@release/2.12
         continue-on-error: true
         with:
           github-secret: ${{ secrets.GITHUB_TOKEN }}
@@ -6139,15 +5966,12 @@ jobs:
       - name: Checkout PyTorch
         uses: actions/checkout@v4
         with:
-          ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
           submodules: recursive
-          path: pytorch
           show-progress: false
       - name: Clean PyTorch checkout
         run: |
           # Remove any artifacts from the previous checkouts
           git clean -fxd
-        working-directory: pytorch
       # NOTE: These environment variables are put here so that they can be applied on every job equally
       #       They are also here because setting them at a workflow level doesn't give us access to the
       #       runner.temp variable, which we need.
@@ -6165,21 +5989,19 @@ jobs:
       - name: Populate binary env
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_populate_env.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_populate_env.sh"
       - name: Test PyTorch binary
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_windows_test.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_windows_test.sh"
       - name: Wait until all sessions have drained
         shell: powershell
-        working-directory: pytorch
         if: always()
         timeout-minutes: 120
         run: |
           .github\scripts\wait_for_ssh_to_drain.ps1
       - name: Kill active ssh sessions if still around (Useful if workflow was cancelled)
         shell: powershell
-        working-directory: pytorch
         if: always()
         run: |
           .github\scripts\kill_active_ssh_sessions.ps1
@@ -6190,7 +6012,7 @@ jobs:
       contents: read
     needs: wheel-py3_14-cpu-test
     with:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
+      PYTORCH_ROOT: ${{ github.workspace }}
       PACKAGE_TYPE: wheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
@@ -6200,6 +6022,9 @@ jobs:
       build_name: wheel-py3_14-cpu
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
     uses: ./.github/workflows/_binary-upload.yml
   wheel-py3_14-cuda12_6-build:
     if: ${{ github.repository_owner == 'pytorch' }}
@@ -6207,7 +6032,7 @@ jobs:
     runs-on: "${{ needs.get-label-type.outputs.label-type }}windows.12xlarge"
     timeout-minutes: 360
     env:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
+      PYTORCH_ROOT: ${{ github.workspace }}
       PACKAGE_TYPE: wheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
@@ -6241,7 +6066,7 @@ jobs:
           echo "instance-type: $(get_ec2_metadata instance-type)"
           echo "system info $(uname -a)"
       - name: "[FB EMPLOYEES] Enable SSH (Click me for login details)"
-        uses: pytorch/test-infra/.github/actions/setup-ssh@main
+        uses: pytorch/test-infra/.github/actions/setup-ssh@release/2.12
         continue-on-error: true
         with:
           github-secret: ${{ secrets.GITHUB_TOKEN }}
@@ -6273,23 +6098,20 @@ jobs:
       - name: Checkout PyTorch
         uses: actions/checkout@v4
         with:
-          ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
           submodules: recursive
-          path: pytorch
           show-progress: false
       - name: Clean PyTorch checkout
         run: |
           # Remove any artifacts from the previous checkouts
           git clean -fxd
-        working-directory: pytorch
       - name: Populate binary env
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_populate_env.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_populate_env.sh"
       - name: Build PyTorch binary
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_windows_build.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_windows_build.sh"
       - uses: actions/upload-artifact@v4.4.0
         if: always()
         with:
@@ -6299,14 +6121,12 @@ jobs:
           path: "${{ env.PYTORCH_FINAL_PACKAGE_DIR }}"
       - name: Wait until all sessions have drained
         shell: powershell
-        working-directory: pytorch
         if: always()
         timeout-minutes: 120
         run: |
           .github\scripts\wait_for_ssh_to_drain.ps1
       - name: Kill active ssh sessions if still around (Useful if workflow was cancelled)
         shell: powershell
-        working-directory: pytorch
         if: always()
         run: |
           .github\scripts\kill_active_ssh_sessions.ps1
@@ -6319,7 +6139,7 @@ jobs:
     runs-on: "${{ needs.get-label-type.outputs.label-type }}windows.g4dn.xlarge"
     timeout-minutes: 360
     env:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
+      PYTORCH_ROOT: ${{ github.workspace }}
       PACKAGE_TYPE: wheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
@@ -6344,7 +6164,7 @@ jobs:
           echo "instance-type: $(get_ec2_metadata instance-type)"
           echo "system info $(uname -a)"
       - name: "[FB EMPLOYEES] Enable SSH (Click me for login details)"
-        uses: pytorch/test-infra/.github/actions/setup-ssh@main
+        uses: pytorch/test-infra/.github/actions/setup-ssh@release/2.12
         continue-on-error: true
         with:
           github-secret: ${{ secrets.GITHUB_TOKEN }}
@@ -6376,15 +6196,12 @@ jobs:
       - name: Checkout PyTorch
         uses: actions/checkout@v4
         with:
-          ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
           submodules: recursive
-          path: pytorch
           show-progress: false
       - name: Clean PyTorch checkout
         run: |
           # Remove any artifacts from the previous checkouts
           git clean -fxd
-        working-directory: pytorch
       # NOTE: These environment variables are put here so that they can be applied on every job equally
       #       They are also here because setting them at a workflow level doesn't give us access to the
       #       runner.temp variable, which we need.
@@ -6402,21 +6219,19 @@ jobs:
       - name: Populate binary env
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_populate_env.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_populate_env.sh"
       - name: Test PyTorch binary
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_windows_test.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_windows_test.sh"
       - name: Wait until all sessions have drained
         shell: powershell
-        working-directory: pytorch
         if: always()
         timeout-minutes: 120
         run: |
           .github\scripts\wait_for_ssh_to_drain.ps1
       - name: Kill active ssh sessions if still around (Useful if workflow was cancelled)
         shell: powershell
-        working-directory: pytorch
         if: always()
         run: |
           .github\scripts\kill_active_ssh_sessions.ps1
@@ -6427,7 +6242,7 @@ jobs:
       contents: read
     needs: wheel-py3_14-cuda12_6-test
     with:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
+      PYTORCH_ROOT: ${{ github.workspace }}
       PACKAGE_TYPE: wheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
@@ -6438,19 +6253,22 @@ jobs:
       build_name: wheel-py3_14-cuda12_6
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
     uses: ./.github/workflows/_binary-upload.yml
-  wheel-py3_14-cuda12_8-build:
+  wheel-py3_14-cuda13_0-build:
     if: ${{ github.repository_owner == 'pytorch' }}
     needs: get-label-type
     runs-on: "${{ needs.get-label-type.outputs.label-type }}windows.12xlarge"
     timeout-minutes: 360
     env:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
+      PYTORCH_ROOT: ${{ github.workspace }}
       PACKAGE_TYPE: wheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu128
-      GPU_ARCH_VERSION: "12.8"
+      DESIRED_CUDA: cu130
+      GPU_ARCH_VERSION: "13.0"
       GPU_ARCH_TYPE: cuda
       SKIP_ALL_TESTS: 1
       DESIRED_PYTHON: "3.14"
@@ -6479,7 +6297,7 @@ jobs:
           echo "instance-type: $(get_ec2_metadata instance-type)"
           echo "system info $(uname -a)"
       - name: "[FB EMPLOYEES] Enable SSH (Click me for login details)"
-        uses: pytorch/test-infra/.github/actions/setup-ssh@main
+        uses: pytorch/test-infra/.github/actions/setup-ssh@release/2.12
         continue-on-error: true
         with:
           github-secret: ${{ secrets.GITHUB_TOKEN }}
@@ -6511,58 +6329,53 @@ jobs:
       - name: Checkout PyTorch
         uses: actions/checkout@v4
         with:
-          ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
           submodules: recursive
-          path: pytorch
           show-progress: false
       - name: Clean PyTorch checkout
         run: |
           # Remove any artifacts from the previous checkouts
           git clean -fxd
-        working-directory: pytorch
       - name: Populate binary env
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_populate_env.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_populate_env.sh"
       - name: Build PyTorch binary
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_windows_build.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_windows_build.sh"
       - uses: actions/upload-artifact@v4.4.0
         if: always()
         with:
-          name: wheel-py3_14-cuda12_8
+          name: wheel-py3_14-cuda13_0
           retention-days: 14
           if-no-files-found: error
           path: "${{ env.PYTORCH_FINAL_PACKAGE_DIR }}"
       - name: Wait until all sessions have drained
         shell: powershell
-        working-directory: pytorch
         if: always()
         timeout-minutes: 120
         run: |
           .github\scripts\wait_for_ssh_to_drain.ps1
       - name: Kill active ssh sessions if still around (Useful if workflow was cancelled)
         shell: powershell
-        working-directory: pytorch
         if: always()
         run: |
           .github\scripts\kill_active_ssh_sessions.ps1
 
-  wheel-py3_14-cuda12_8-test:  # Testing
+  wheel-py3_14-cuda13_0-test:  # Testing
     if: ${{ github.repository_owner == 'pytorch' }}
     needs:
-      - wheel-py3_14-cuda12_8-build
+      - wheel-py3_14-cuda13_0-build
       - get-label-type
     runs-on: "${{ needs.get-label-type.outputs.label-type }}windows.g4dn.xlarge"
     timeout-minutes: 360
     env:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
+      PYTORCH_ROOT: ${{ github.workspace }}
       PACKAGE_TYPE: wheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu128
-      GPU_ARCH_VERSION: "12.8"
+      DESIRED_CUDA: cu130
+      GPU_ARCH_VERSION: "13.0"
       GPU_ARCH_TYPE: cuda
       SKIP_ALL_TESTS: 1
       DESIRED_PYTHON: "3.14"
@@ -6582,7 +6395,7 @@ jobs:
           echo "instance-type: $(get_ec2_metadata instance-type)"
           echo "system info $(uname -a)"
       - name: "[FB EMPLOYEES] Enable SSH (Click me for login details)"
-        uses: pytorch/test-infra/.github/actions/setup-ssh@main
+        uses: pytorch/test-infra/.github/actions/setup-ssh@release/2.12
         continue-on-error: true
         with:
           github-secret: ${{ secrets.GITHUB_TOKEN }}
@@ -6614,15 +6427,12 @@ jobs:
       - name: Checkout PyTorch
         uses: actions/checkout@v4
         with:
-          ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
           submodules: recursive
-          path: pytorch
           show-progress: false
       - name: Clean PyTorch checkout
         run: |
           # Remove any artifacts from the previous checkouts
           git clean -fxd
-        working-directory: pytorch
       # NOTE: These environment variables are put here so that they can be applied on every job equally
       #       They are also here because setting them at a workflow level doesn't give us access to the
       #       runner.temp variable, which we need.
@@ -6635,60 +6445,61 @@ jobs:
       - uses: actions/download-artifact@v4.1.7
         name: Download Build Artifacts
         with:
-          name: wheel-py3_14-cuda12_8
+          name: wheel-py3_14-cuda13_0
           path: "${{ env.PYTORCH_FINAL_PACKAGE_DIR }}"
       - name: Populate binary env
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_populate_env.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_populate_env.sh"
       - name: Test PyTorch binary
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_windows_test.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_windows_test.sh"
       - name: Wait until all sessions have drained
         shell: powershell
-        working-directory: pytorch
         if: always()
         timeout-minutes: 120
         run: |
           .github\scripts\wait_for_ssh_to_drain.ps1
       - name: Kill active ssh sessions if still around (Useful if workflow was cancelled)
         shell: powershell
-        working-directory: pytorch
         if: always()
         run: |
           .github\scripts\kill_active_ssh_sessions.ps1
-  wheel-py3_14-cuda12_8-upload:  # Uploading
+  wheel-py3_14-cuda13_0-upload:  # Uploading
     if: ${{ github.repository_owner == 'pytorch' }}
     permissions:
       id-token: write
       contents: read
-    needs: wheel-py3_14-cuda12_8-test
+    needs: wheel-py3_14-cuda13_0-test
     with:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
+      PYTORCH_ROOT: ${{ github.workspace }}
       PACKAGE_TYPE: wheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu128
-      GPU_ARCH_VERSION: "12.8"
+      DESIRED_CUDA: cu130
+      GPU_ARCH_VERSION: "13.0"
       GPU_ARCH_TYPE: cuda
       DESIRED_PYTHON: "3.14"
-      build_name: wheel-py3_14-cuda12_8
+      build_name: wheel-py3_14-cuda13_0
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
     uses: ./.github/workflows/_binary-upload.yml
-  wheel-py3_14-cuda13_0-build:
+  wheel-py3_14-cuda13_2-build:
     if: ${{ github.repository_owner == 'pytorch' }}
     needs: get-label-type
     runs-on: "${{ needs.get-label-type.outputs.label-type }}windows.12xlarge"
     timeout-minutes: 360
     env:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
+      PYTORCH_ROOT: ${{ github.workspace }}
       PACKAGE_TYPE: wheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu130
-      GPU_ARCH_VERSION: "13.0"
+      DESIRED_CUDA: cu132
+      GPU_ARCH_VERSION: "13.2"
       GPU_ARCH_TYPE: cuda
       SKIP_ALL_TESTS: 1
       DESIRED_PYTHON: "3.14"
@@ -6717,7 +6528,7 @@ jobs:
           echo "instance-type: $(get_ec2_metadata instance-type)"
           echo "system info $(uname -a)"
       - name: "[FB EMPLOYEES] Enable SSH (Click me for login details)"
-        uses: pytorch/test-infra/.github/actions/setup-ssh@main
+        uses: pytorch/test-infra/.github/actions/setup-ssh@release/2.12
         continue-on-error: true
         with:
           github-secret: ${{ secrets.GITHUB_TOKEN }}
@@ -6749,58 +6560,53 @@ jobs:
       - name: Checkout PyTorch
         uses: actions/checkout@v4
         with:
-          ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
           submodules: recursive
-          path: pytorch
           show-progress: false
       - name: Clean PyTorch checkout
         run: |
           # Remove any artifacts from the previous checkouts
           git clean -fxd
-        working-directory: pytorch
       - name: Populate binary env
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_populate_env.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_populate_env.sh"
       - name: Build PyTorch binary
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_windows_build.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_windows_build.sh"
       - uses: actions/upload-artifact@v4.4.0
         if: always()
         with:
-          name: wheel-py3_14-cuda13_0
+          name: wheel-py3_14-cuda13_2
           retention-days: 14
           if-no-files-found: error
           path: "${{ env.PYTORCH_FINAL_PACKAGE_DIR }}"
       - name: Wait until all sessions have drained
         shell: powershell
-        working-directory: pytorch
         if: always()
         timeout-minutes: 120
         run: |
           .github\scripts\wait_for_ssh_to_drain.ps1
       - name: Kill active ssh sessions if still around (Useful if workflow was cancelled)
         shell: powershell
-        working-directory: pytorch
         if: always()
         run: |
           .github\scripts\kill_active_ssh_sessions.ps1
 
-  wheel-py3_14-cuda13_0-test:  # Testing
+  wheel-py3_14-cuda13_2-test:  # Testing
     if: ${{ github.repository_owner == 'pytorch' }}
     needs:
-      - wheel-py3_14-cuda13_0-build
+      - wheel-py3_14-cuda13_2-build
       - get-label-type
     runs-on: "${{ needs.get-label-type.outputs.label-type }}windows.g4dn.xlarge"
     timeout-minutes: 360
     env:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
+      PYTORCH_ROOT: ${{ github.workspace }}
       PACKAGE_TYPE: wheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu130
-      GPU_ARCH_VERSION: "13.0"
+      DESIRED_CUDA: cu132
+      GPU_ARCH_VERSION: "13.2"
       GPU_ARCH_TYPE: cuda
       SKIP_ALL_TESTS: 1
       DESIRED_PYTHON: "3.14"
@@ -6820,7 +6626,7 @@ jobs:
           echo "instance-type: $(get_ec2_metadata instance-type)"
           echo "system info $(uname -a)"
       - name: "[FB EMPLOYEES] Enable SSH (Click me for login details)"
-        uses: pytorch/test-infra/.github/actions/setup-ssh@main
+        uses: pytorch/test-infra/.github/actions/setup-ssh@release/2.12
         continue-on-error: true
         with:
           github-secret: ${{ secrets.GITHUB_TOKEN }}
@@ -6852,15 +6658,12 @@ jobs:
       - name: Checkout PyTorch
         uses: actions/checkout@v4
         with:
-          ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
           submodules: recursive
-          path: pytorch
           show-progress: false
       - name: Clean PyTorch checkout
         run: |
           # Remove any artifacts from the previous checkouts
           git clean -fxd
-        working-directory: pytorch
       # NOTE: These environment variables are put here so that they can be applied on every job equally
       #       They are also here because setting them at a workflow level doesn't give us access to the
       #       runner.temp variable, which we need.
@@ -6873,47 +6676,48 @@ jobs:
       - uses: actions/download-artifact@v4.1.7
         name: Download Build Artifacts
         with:
-          name: wheel-py3_14-cuda13_0
+          name: wheel-py3_14-cuda13_2
           path: "${{ env.PYTORCH_FINAL_PACKAGE_DIR }}"
       - name: Populate binary env
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_populate_env.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_populate_env.sh"
       - name: Test PyTorch binary
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_windows_test.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_windows_test.sh"
       - name: Wait until all sessions have drained
         shell: powershell
-        working-directory: pytorch
         if: always()
         timeout-minutes: 120
         run: |
           .github\scripts\wait_for_ssh_to_drain.ps1
       - name: Kill active ssh sessions if still around (Useful if workflow was cancelled)
         shell: powershell
-        working-directory: pytorch
         if: always()
         run: |
           .github\scripts\kill_active_ssh_sessions.ps1
-  wheel-py3_14-cuda13_0-upload:  # Uploading
+  wheel-py3_14-cuda13_2-upload:  # Uploading
     if: ${{ github.repository_owner == 'pytorch' }}
     permissions:
       id-token: write
       contents: read
-    needs: wheel-py3_14-cuda13_0-test
+    needs: wheel-py3_14-cuda13_2-test
     with:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
+      PYTORCH_ROOT: ${{ github.workspace }}
       PACKAGE_TYPE: wheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu130
-      GPU_ARCH_VERSION: "13.0"
+      DESIRED_CUDA: cu132
+      GPU_ARCH_VERSION: "13.2"
       GPU_ARCH_TYPE: cuda
       DESIRED_PYTHON: "3.14"
-      build_name: wheel-py3_14-cuda13_0
+      build_name: wheel-py3_14-cuda13_2
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
     uses: ./.github/workflows/_binary-upload.yml
   wheel-py3_14-xpu-build:
     if: ${{ github.repository_owner == 'pytorch' }}
@@ -6921,7 +6725,7 @@ jobs:
     runs-on: "${{ needs.get-label-type.outputs.label-type }}windows.12xlarge"
     timeout-minutes: 360
     env:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
+      PYTORCH_ROOT: ${{ github.workspace }}
       PACKAGE_TYPE: wheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
@@ -6955,7 +6759,7 @@ jobs:
           echo "instance-type: $(get_ec2_metadata instance-type)"
           echo "system info $(uname -a)"
       - name: "[FB EMPLOYEES] Enable SSH (Click me for login details)"
-        uses: pytorch/test-infra/.github/actions/setup-ssh@main
+        uses: pytorch/test-infra/.github/actions/setup-ssh@release/2.12
         continue-on-error: true
         with:
           github-secret: ${{ secrets.GITHUB_TOKEN }}
@@ -6987,23 +6791,20 @@ jobs:
       - name: Checkout PyTorch
         uses: actions/checkout@v4
         with:
-          ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
           submodules: recursive
-          path: pytorch
           show-progress: false
       - name: Clean PyTorch checkout
         run: |
           # Remove any artifacts from the previous checkouts
           git clean -fxd
-        working-directory: pytorch
       - name: Populate binary env
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_populate_env.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_populate_env.sh"
       - name: Build PyTorch binary
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_windows_build.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_windows_build.sh"
       - uses: actions/upload-artifact@v4.4.0
         if: always()
         with:
@@ -7013,14 +6814,12 @@ jobs:
           path: "${{ env.PYTORCH_FINAL_PACKAGE_DIR }}"
       - name: Wait until all sessions have drained
         shell: powershell
-        working-directory: pytorch
         if: always()
         timeout-minutes: 120
         run: |
           .github\scripts\wait_for_ssh_to_drain.ps1
       - name: Kill active ssh sessions if still around (Useful if workflow was cancelled)
         shell: powershell
-        working-directory: pytorch
         if: always()
         run: |
           .github\scripts\kill_active_ssh_sessions.ps1
@@ -7033,7 +6832,7 @@ jobs:
     runs-on: "${{ needs.get-label-type.outputs.label-type }}windows.4xlarge"
     timeout-minutes: 360
     env:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
+      PYTORCH_ROOT: ${{ github.workspace }}
       PACKAGE_TYPE: wheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
@@ -7057,7 +6856,7 @@ jobs:
           echo "instance-type: $(get_ec2_metadata instance-type)"
           echo "system info $(uname -a)"
       - name: "[FB EMPLOYEES] Enable SSH (Click me for login details)"
-        uses: pytorch/test-infra/.github/actions/setup-ssh@main
+        uses: pytorch/test-infra/.github/actions/setup-ssh@release/2.12
         continue-on-error: true
         with:
           github-secret: ${{ secrets.GITHUB_TOKEN }}
@@ -7089,15 +6888,12 @@ jobs:
       - name: Checkout PyTorch
         uses: actions/checkout@v4
         with:
-          ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
           submodules: recursive
-          path: pytorch
           show-progress: false
       - name: Clean PyTorch checkout
         run: |
           # Remove any artifacts from the previous checkouts
           git clean -fxd
-        working-directory: pytorch
       # NOTE: These environment variables are put here so that they can be applied on every job equally
       #       They are also here because setting them at a workflow level doesn't give us access to the
       #       runner.temp variable, which we need.
@@ -7115,21 +6911,19 @@ jobs:
       - name: Populate binary env
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_populate_env.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_populate_env.sh"
       - name: Test PyTorch binary
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_windows_test.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_windows_test.sh"
       - name: Wait until all sessions have drained
         shell: powershell
-        working-directory: pytorch
         if: always()
         timeout-minutes: 120
         run: |
           .github\scripts\wait_for_ssh_to_drain.ps1
       - name: Kill active ssh sessions if still around (Useful if workflow was cancelled)
         shell: powershell
-        working-directory: pytorch
         if: always()
         run: |
           .github\scripts\kill_active_ssh_sessions.ps1
@@ -7140,7 +6934,7 @@ jobs:
       contents: read
     needs: wheel-py3_14-xpu-test
     with:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
+      PYTORCH_ROOT: ${{ github.workspace }}
       PACKAGE_TYPE: wheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
@@ -7150,6 +6944,9 @@ jobs:
       build_name: wheel-py3_14-xpu
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
     uses: ./.github/workflows/_binary-upload.yml
   wheel-py3_14t-cpu-build:
     if: ${{ github.repository_owner == 'pytorch' }}
@@ -7157,7 +6954,7 @@ jobs:
     runs-on: "${{ needs.get-label-type.outputs.label-type }}windows.12xlarge"
     timeout-minutes: 360
     env:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
+      PYTORCH_ROOT: ${{ github.workspace }}
       PACKAGE_TYPE: wheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
@@ -7165,6 +6962,7 @@ jobs:
       GPU_ARCH_TYPE: cpu
       SKIP_ALL_TESTS: 1
       DESIRED_PYTHON: "3.14t"
+      PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cufile,nvjitlink,nvtx]==13.0.2; platform_system == 'Linux' | nvidia-cublas>=13.1.0.3,<=13.1.1.3; platform_system == 'Linux' | cuda-bindings>=13.0.3,<14; platform_system == 'Linux' | nvidia-cudnn-cu13==9.20.0.48; platform_system == 'Linux' | nvidia-cusparselt-cu13==0.8.1; platform_system == 'Linux' | nvidia-nccl-cu13==2.29.7; platform_system == 'Linux' | nvidia-nvshmem-cu13==3.4.5; platform_system == 'Linux'
     steps:
       # NOTE: These environment variables are put here so that they can be applied on every job equally
       #       They are also here because setting them at a workflow level doesn't give us access to the
@@ -7190,7 +6988,7 @@ jobs:
           echo "instance-type: $(get_ec2_metadata instance-type)"
           echo "system info $(uname -a)"
       - name: "[FB EMPLOYEES] Enable SSH (Click me for login details)"
-        uses: pytorch/test-infra/.github/actions/setup-ssh@main
+        uses: pytorch/test-infra/.github/actions/setup-ssh@release/2.12
         continue-on-error: true
         with:
           github-secret: ${{ secrets.GITHUB_TOKEN }}
@@ -7222,23 +7020,20 @@ jobs:
       - name: Checkout PyTorch
         uses: actions/checkout@v4
         with:
-          ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
           submodules: recursive
-          path: pytorch
           show-progress: false
       - name: Clean PyTorch checkout
         run: |
           # Remove any artifacts from the previous checkouts
           git clean -fxd
-        working-directory: pytorch
       - name: Populate binary env
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_populate_env.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_populate_env.sh"
       - name: Build PyTorch binary
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_windows_build.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_windows_build.sh"
       - uses: actions/upload-artifact@v4.4.0
         if: always()
         with:
@@ -7248,14 +7043,12 @@ jobs:
           path: "${{ env.PYTORCH_FINAL_PACKAGE_DIR }}"
       - name: Wait until all sessions have drained
         shell: powershell
-        working-directory: pytorch
         if: always()
         timeout-minutes: 120
         run: |
           .github\scripts\wait_for_ssh_to_drain.ps1
       - name: Kill active ssh sessions if still around (Useful if workflow was cancelled)
         shell: powershell
-        working-directory: pytorch
         if: always()
         run: |
           .github\scripts\kill_active_ssh_sessions.ps1
@@ -7268,7 +7061,7 @@ jobs:
     runs-on: "${{ needs.get-label-type.outputs.label-type }}windows.4xlarge"
     timeout-minutes: 360
     env:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
+      PYTORCH_ROOT: ${{ github.workspace }}
       PACKAGE_TYPE: wheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
@@ -7292,7 +7085,7 @@ jobs:
           echo "instance-type: $(get_ec2_metadata instance-type)"
           echo "system info $(uname -a)"
       - name: "[FB EMPLOYEES] Enable SSH (Click me for login details)"
-        uses: pytorch/test-infra/.github/actions/setup-ssh@main
+        uses: pytorch/test-infra/.github/actions/setup-ssh@release/2.12
         continue-on-error: true
         with:
           github-secret: ${{ secrets.GITHUB_TOKEN }}
@@ -7324,15 +7117,12 @@ jobs:
       - name: Checkout PyTorch
         uses: actions/checkout@v4
         with:
-          ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
           submodules: recursive
-          path: pytorch
           show-progress: false
       - name: Clean PyTorch checkout
         run: |
           # Remove any artifacts from the previous checkouts
           git clean -fxd
-        working-directory: pytorch
       # NOTE: These environment variables are put here so that they can be applied on every job equally
       #       They are also here because setting them at a workflow level doesn't give us access to the
       #       runner.temp variable, which we need.
@@ -7350,21 +7140,19 @@ jobs:
       - name: Populate binary env
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_populate_env.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_populate_env.sh"
       - name: Test PyTorch binary
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_windows_test.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_windows_test.sh"
       - name: Wait until all sessions have drained
         shell: powershell
-        working-directory: pytorch
         if: always()
         timeout-minutes: 120
         run: |
           .github\scripts\wait_for_ssh_to_drain.ps1
       - name: Kill active ssh sessions if still around (Useful if workflow was cancelled)
         shell: powershell
-        working-directory: pytorch
         if: always()
         run: |
           .github\scripts\kill_active_ssh_sessions.ps1
@@ -7375,7 +7163,7 @@ jobs:
       contents: read
     needs: wheel-py3_14t-cpu-test
     with:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
+      PYTORCH_ROOT: ${{ github.workspace }}
       PACKAGE_TYPE: wheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
@@ -7385,6 +7173,9 @@ jobs:
       build_name: wheel-py3_14t-cpu
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
     uses: ./.github/workflows/_binary-upload.yml
   wheel-py3_14t-cuda12_6-build:
     if: ${{ github.repository_owner == 'pytorch' }}
@@ -7392,7 +7183,7 @@ jobs:
     runs-on: "${{ needs.get-label-type.outputs.label-type }}windows.12xlarge"
     timeout-minutes: 360
     env:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
+      PYTORCH_ROOT: ${{ github.workspace }}
       PACKAGE_TYPE: wheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
@@ -7426,7 +7217,7 @@ jobs:
           echo "instance-type: $(get_ec2_metadata instance-type)"
           echo "system info $(uname -a)"
       - name: "[FB EMPLOYEES] Enable SSH (Click me for login details)"
-        uses: pytorch/test-infra/.github/actions/setup-ssh@main
+        uses: pytorch/test-infra/.github/actions/setup-ssh@release/2.12
         continue-on-error: true
         with:
           github-secret: ${{ secrets.GITHUB_TOKEN }}
@@ -7458,23 +7249,20 @@ jobs:
       - name: Checkout PyTorch
         uses: actions/checkout@v4
         with:
-          ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
           submodules: recursive
-          path: pytorch
           show-progress: false
       - name: Clean PyTorch checkout
         run: |
           # Remove any artifacts from the previous checkouts
           git clean -fxd
-        working-directory: pytorch
       - name: Populate binary env
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_populate_env.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_populate_env.sh"
       - name: Build PyTorch binary
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_windows_build.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_windows_build.sh"
       - uses: actions/upload-artifact@v4.4.0
         if: always()
         with:
@@ -7484,14 +7272,12 @@ jobs:
           path: "${{ env.PYTORCH_FINAL_PACKAGE_DIR }}"
       - name: Wait until all sessions have drained
         shell: powershell
-        working-directory: pytorch
         if: always()
         timeout-minutes: 120
         run: |
           .github\scripts\wait_for_ssh_to_drain.ps1
       - name: Kill active ssh sessions if still around (Useful if workflow was cancelled)
         shell: powershell
-        working-directory: pytorch
         if: always()
         run: |
           .github\scripts\kill_active_ssh_sessions.ps1
@@ -7504,7 +7290,7 @@ jobs:
     runs-on: "${{ needs.get-label-type.outputs.label-type }}windows.g4dn.xlarge"
     timeout-minutes: 360
     env:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
+      PYTORCH_ROOT: ${{ github.workspace }}
       PACKAGE_TYPE: wheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
@@ -7529,7 +7315,7 @@ jobs:
           echo "instance-type: $(get_ec2_metadata instance-type)"
           echo "system info $(uname -a)"
       - name: "[FB EMPLOYEES] Enable SSH (Click me for login details)"
-        uses: pytorch/test-infra/.github/actions/setup-ssh@main
+        uses: pytorch/test-infra/.github/actions/setup-ssh@release/2.12
         continue-on-error: true
         with:
           github-secret: ${{ secrets.GITHUB_TOKEN }}
@@ -7561,15 +7347,12 @@ jobs:
       - name: Checkout PyTorch
         uses: actions/checkout@v4
         with:
-          ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
           submodules: recursive
-          path: pytorch
           show-progress: false
       - name: Clean PyTorch checkout
         run: |
           # Remove any artifacts from the previous checkouts
           git clean -fxd
-        working-directory: pytorch
       # NOTE: These environment variables are put here so that they can be applied on every job equally
       #       They are also here because setting them at a workflow level doesn't give us access to the
       #       runner.temp variable, which we need.
@@ -7587,21 +7370,19 @@ jobs:
       - name: Populate binary env
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_populate_env.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_populate_env.sh"
       - name: Test PyTorch binary
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_windows_test.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_windows_test.sh"
       - name: Wait until all sessions have drained
         shell: powershell
-        working-directory: pytorch
         if: always()
         timeout-minutes: 120
         run: |
           .github\scripts\wait_for_ssh_to_drain.ps1
       - name: Kill active ssh sessions if still around (Useful if workflow was cancelled)
         shell: powershell
-        working-directory: pytorch
         if: always()
         run: |
           .github\scripts\kill_active_ssh_sessions.ps1
@@ -7612,7 +7393,7 @@ jobs:
       contents: read
     needs: wheel-py3_14t-cuda12_6-test
     with:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
+      PYTORCH_ROOT: ${{ github.workspace }}
       PACKAGE_TYPE: wheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
@@ -7623,19 +7404,22 @@ jobs:
       build_name: wheel-py3_14t-cuda12_6
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
     uses: ./.github/workflows/_binary-upload.yml
-  wheel-py3_14t-cuda12_8-build:
+  wheel-py3_14t-cuda13_0-build:
     if: ${{ github.repository_owner == 'pytorch' }}
     needs: get-label-type
     runs-on: "${{ needs.get-label-type.outputs.label-type }}windows.12xlarge"
     timeout-minutes: 360
     env:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
+      PYTORCH_ROOT: ${{ github.workspace }}
       PACKAGE_TYPE: wheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu128
-      GPU_ARCH_VERSION: "12.8"
+      DESIRED_CUDA: cu130
+      GPU_ARCH_VERSION: "13.0"
       GPU_ARCH_TYPE: cuda
       SKIP_ALL_TESTS: 1
       DESIRED_PYTHON: "3.14t"
@@ -7664,7 +7448,7 @@ jobs:
           echo "instance-type: $(get_ec2_metadata instance-type)"
           echo "system info $(uname -a)"
       - name: "[FB EMPLOYEES] Enable SSH (Click me for login details)"
-        uses: pytorch/test-infra/.github/actions/setup-ssh@main
+        uses: pytorch/test-infra/.github/actions/setup-ssh@release/2.12
         continue-on-error: true
         with:
           github-secret: ${{ secrets.GITHUB_TOKEN }}
@@ -7696,58 +7480,53 @@ jobs:
       - name: Checkout PyTorch
         uses: actions/checkout@v4
         with:
-          ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
           submodules: recursive
-          path: pytorch
           show-progress: false
       - name: Clean PyTorch checkout
         run: |
           # Remove any artifacts from the previous checkouts
           git clean -fxd
-        working-directory: pytorch
       - name: Populate binary env
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_populate_env.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_populate_env.sh"
       - name: Build PyTorch binary
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_windows_build.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_windows_build.sh"
       - uses: actions/upload-artifact@v4.4.0
         if: always()
         with:
-          name: wheel-py3_14t-cuda12_8
+          name: wheel-py3_14t-cuda13_0
           retention-days: 14
           if-no-files-found: error
           path: "${{ env.PYTORCH_FINAL_PACKAGE_DIR }}"
       - name: Wait until all sessions have drained
         shell: powershell
-        working-directory: pytorch
         if: always()
         timeout-minutes: 120
         run: |
           .github\scripts\wait_for_ssh_to_drain.ps1
       - name: Kill active ssh sessions if still around (Useful if workflow was cancelled)
         shell: powershell
-        working-directory: pytorch
         if: always()
         run: |
           .github\scripts\kill_active_ssh_sessions.ps1
 
-  wheel-py3_14t-cuda12_8-test:  # Testing
+  wheel-py3_14t-cuda13_0-test:  # Testing
     if: ${{ github.repository_owner == 'pytorch' }}
     needs:
-      - wheel-py3_14t-cuda12_8-build
+      - wheel-py3_14t-cuda13_0-build
       - get-label-type
     runs-on: "${{ needs.get-label-type.outputs.label-type }}windows.g4dn.xlarge"
     timeout-minutes: 360
     env:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
+      PYTORCH_ROOT: ${{ github.workspace }}
       PACKAGE_TYPE: wheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu128
-      GPU_ARCH_VERSION: "12.8"
+      DESIRED_CUDA: cu130
+      GPU_ARCH_VERSION: "13.0"
       GPU_ARCH_TYPE: cuda
       SKIP_ALL_TESTS: 1
       DESIRED_PYTHON: "3.14t"
@@ -7767,7 +7546,7 @@ jobs:
           echo "instance-type: $(get_ec2_metadata instance-type)"
           echo "system info $(uname -a)"
       - name: "[FB EMPLOYEES] Enable SSH (Click me for login details)"
-        uses: pytorch/test-infra/.github/actions/setup-ssh@main
+        uses: pytorch/test-infra/.github/actions/setup-ssh@release/2.12
         continue-on-error: true
         with:
           github-secret: ${{ secrets.GITHUB_TOKEN }}
@@ -7799,15 +7578,12 @@ jobs:
       - name: Checkout PyTorch
         uses: actions/checkout@v4
         with:
-          ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
           submodules: recursive
-          path: pytorch
           show-progress: false
       - name: Clean PyTorch checkout
         run: |
           # Remove any artifacts from the previous checkouts
           git clean -fxd
-        working-directory: pytorch
       # NOTE: These environment variables are put here so that they can be applied on every job equally
       #       They are also here because setting them at a workflow level doesn't give us access to the
       #       runner.temp variable, which we need.
@@ -7820,60 +7596,61 @@ jobs:
       - uses: actions/download-artifact@v4.1.7
         name: Download Build Artifacts
         with:
-          name: wheel-py3_14t-cuda12_8
+          name: wheel-py3_14t-cuda13_0
           path: "${{ env.PYTORCH_FINAL_PACKAGE_DIR }}"
       - name: Populate binary env
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_populate_env.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_populate_env.sh"
       - name: Test PyTorch binary
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_windows_test.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_windows_test.sh"
       - name: Wait until all sessions have drained
         shell: powershell
-        working-directory: pytorch
         if: always()
         timeout-minutes: 120
         run: |
           .github\scripts\wait_for_ssh_to_drain.ps1
       - name: Kill active ssh sessions if still around (Useful if workflow was cancelled)
         shell: powershell
-        working-directory: pytorch
         if: always()
         run: |
           .github\scripts\kill_active_ssh_sessions.ps1
-  wheel-py3_14t-cuda12_8-upload:  # Uploading
+  wheel-py3_14t-cuda13_0-upload:  # Uploading
     if: ${{ github.repository_owner == 'pytorch' }}
     permissions:
       id-token: write
       contents: read
-    needs: wheel-py3_14t-cuda12_8-test
+    needs: wheel-py3_14t-cuda13_0-test
     with:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
+      PYTORCH_ROOT: ${{ github.workspace }}
       PACKAGE_TYPE: wheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu128
-      GPU_ARCH_VERSION: "12.8"
+      DESIRED_CUDA: cu130
+      GPU_ARCH_VERSION: "13.0"
       GPU_ARCH_TYPE: cuda
       DESIRED_PYTHON: "3.14t"
-      build_name: wheel-py3_14t-cuda12_8
+      build_name: wheel-py3_14t-cuda13_0
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
     uses: ./.github/workflows/_binary-upload.yml
-  wheel-py3_14t-cuda13_0-build:
+  wheel-py3_14t-cuda13_2-build:
     if: ${{ github.repository_owner == 'pytorch' }}
     needs: get-label-type
     runs-on: "${{ needs.get-label-type.outputs.label-type }}windows.12xlarge"
     timeout-minutes: 360
     env:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
+      PYTORCH_ROOT: ${{ github.workspace }}
       PACKAGE_TYPE: wheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu130
-      GPU_ARCH_VERSION: "13.0"
+      DESIRED_CUDA: cu132
+      GPU_ARCH_VERSION: "13.2"
       GPU_ARCH_TYPE: cuda
       SKIP_ALL_TESTS: 1
       DESIRED_PYTHON: "3.14t"
@@ -7902,7 +7679,7 @@ jobs:
           echo "instance-type: $(get_ec2_metadata instance-type)"
           echo "system info $(uname -a)"
       - name: "[FB EMPLOYEES] Enable SSH (Click me for login details)"
-        uses: pytorch/test-infra/.github/actions/setup-ssh@main
+        uses: pytorch/test-infra/.github/actions/setup-ssh@release/2.12
         continue-on-error: true
         with:
           github-secret: ${{ secrets.GITHUB_TOKEN }}
@@ -7934,58 +7711,53 @@ jobs:
       - name: Checkout PyTorch
         uses: actions/checkout@v4
         with:
-          ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
           submodules: recursive
-          path: pytorch
           show-progress: false
       - name: Clean PyTorch checkout
         run: |
           # Remove any artifacts from the previous checkouts
           git clean -fxd
-        working-directory: pytorch
       - name: Populate binary env
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_populate_env.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_populate_env.sh"
       - name: Build PyTorch binary
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_windows_build.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_windows_build.sh"
       - uses: actions/upload-artifact@v4.4.0
         if: always()
         with:
-          name: wheel-py3_14t-cuda13_0
+          name: wheel-py3_14t-cuda13_2
           retention-days: 14
           if-no-files-found: error
           path: "${{ env.PYTORCH_FINAL_PACKAGE_DIR }}"
       - name: Wait until all sessions have drained
         shell: powershell
-        working-directory: pytorch
         if: always()
         timeout-minutes: 120
         run: |
           .github\scripts\wait_for_ssh_to_drain.ps1
       - name: Kill active ssh sessions if still around (Useful if workflow was cancelled)
         shell: powershell
-        working-directory: pytorch
         if: always()
         run: |
           .github\scripts\kill_active_ssh_sessions.ps1
 
-  wheel-py3_14t-cuda13_0-test:  # Testing
+  wheel-py3_14t-cuda13_2-test:  # Testing
     if: ${{ github.repository_owner == 'pytorch' }}
     needs:
-      - wheel-py3_14t-cuda13_0-build
+      - wheel-py3_14t-cuda13_2-build
       - get-label-type
     runs-on: "${{ needs.get-label-type.outputs.label-type }}windows.g4dn.xlarge"
     timeout-minutes: 360
     env:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
+      PYTORCH_ROOT: ${{ github.workspace }}
       PACKAGE_TYPE: wheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu130
-      GPU_ARCH_VERSION: "13.0"
+      DESIRED_CUDA: cu132
+      GPU_ARCH_VERSION: "13.2"
       GPU_ARCH_TYPE: cuda
       SKIP_ALL_TESTS: 1
       DESIRED_PYTHON: "3.14t"
@@ -8005,7 +7777,7 @@ jobs:
           echo "instance-type: $(get_ec2_metadata instance-type)"
           echo "system info $(uname -a)"
       - name: "[FB EMPLOYEES] Enable SSH (Click me for login details)"
-        uses: pytorch/test-infra/.github/actions/setup-ssh@main
+        uses: pytorch/test-infra/.github/actions/setup-ssh@release/2.12
         continue-on-error: true
         with:
           github-secret: ${{ secrets.GITHUB_TOKEN }}
@@ -8037,15 +7809,12 @@ jobs:
       - name: Checkout PyTorch
         uses: actions/checkout@v4
         with:
-          ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
           submodules: recursive
-          path: pytorch
           show-progress: false
       - name: Clean PyTorch checkout
         run: |
           # Remove any artifacts from the previous checkouts
           git clean -fxd
-        working-directory: pytorch
       # NOTE: These environment variables are put here so that they can be applied on every job equally
       #       They are also here because setting them at a workflow level doesn't give us access to the
       #       runner.temp variable, which we need.
@@ -8058,47 +7827,48 @@ jobs:
       - uses: actions/download-artifact@v4.1.7
         name: Download Build Artifacts
         with:
-          name: wheel-py3_14t-cuda13_0
+          name: wheel-py3_14t-cuda13_2
           path: "${{ env.PYTORCH_FINAL_PACKAGE_DIR }}"
       - name: Populate binary env
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_populate_env.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_populate_env.sh"
       - name: Test PyTorch binary
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_windows_test.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_windows_test.sh"
       - name: Wait until all sessions have drained
         shell: powershell
-        working-directory: pytorch
         if: always()
         timeout-minutes: 120
         run: |
           .github\scripts\wait_for_ssh_to_drain.ps1
       - name: Kill active ssh sessions if still around (Useful if workflow was cancelled)
         shell: powershell
-        working-directory: pytorch
         if: always()
         run: |
           .github\scripts\kill_active_ssh_sessions.ps1
-  wheel-py3_14t-cuda13_0-upload:  # Uploading
+  wheel-py3_14t-cuda13_2-upload:  # Uploading
     if: ${{ github.repository_owner == 'pytorch' }}
     permissions:
       id-token: write
       contents: read
-    needs: wheel-py3_14t-cuda13_0-test
+    needs: wheel-py3_14t-cuda13_2-test
     with:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
+      PYTORCH_ROOT: ${{ github.workspace }}
       PACKAGE_TYPE: wheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
-      DESIRED_CUDA: cu130
-      GPU_ARCH_VERSION: "13.0"
+      DESIRED_CUDA: cu132
+      GPU_ARCH_VERSION: "13.2"
       GPU_ARCH_TYPE: cuda
       DESIRED_PYTHON: "3.14t"
-      build_name: wheel-py3_14t-cuda13_0
+      build_name: wheel-py3_14t-cuda13_2
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
     uses: ./.github/workflows/_binary-upload.yml
   wheel-py3_14t-xpu-build:
     if: ${{ github.repository_owner == 'pytorch' }}
@@ -8106,7 +7876,7 @@ jobs:
     runs-on: "${{ needs.get-label-type.outputs.label-type }}windows.12xlarge"
     timeout-minutes: 360
     env:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
+      PYTORCH_ROOT: ${{ github.workspace }}
       PACKAGE_TYPE: wheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
@@ -8140,7 +7910,7 @@ jobs:
           echo "instance-type: $(get_ec2_metadata instance-type)"
           echo "system info $(uname -a)"
       - name: "[FB EMPLOYEES] Enable SSH (Click me for login details)"
-        uses: pytorch/test-infra/.github/actions/setup-ssh@main
+        uses: pytorch/test-infra/.github/actions/setup-ssh@release/2.12
         continue-on-error: true
         with:
           github-secret: ${{ secrets.GITHUB_TOKEN }}
@@ -8172,23 +7942,20 @@ jobs:
       - name: Checkout PyTorch
         uses: actions/checkout@v4
         with:
-          ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
           submodules: recursive
-          path: pytorch
           show-progress: false
       - name: Clean PyTorch checkout
         run: |
           # Remove any artifacts from the previous checkouts
           git clean -fxd
-        working-directory: pytorch
       - name: Populate binary env
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_populate_env.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_populate_env.sh"
       - name: Build PyTorch binary
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_windows_build.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_windows_build.sh"
       - uses: actions/upload-artifact@v4.4.0
         if: always()
         with:
@@ -8198,14 +7965,12 @@ jobs:
           path: "${{ env.PYTORCH_FINAL_PACKAGE_DIR }}"
       - name: Wait until all sessions have drained
         shell: powershell
-        working-directory: pytorch
         if: always()
         timeout-minutes: 120
         run: |
           .github\scripts\wait_for_ssh_to_drain.ps1
       - name: Kill active ssh sessions if still around (Useful if workflow was cancelled)
         shell: powershell
-        working-directory: pytorch
         if: always()
         run: |
           .github\scripts\kill_active_ssh_sessions.ps1
@@ -8218,7 +7983,7 @@ jobs:
     runs-on: "${{ needs.get-label-type.outputs.label-type }}windows.4xlarge"
     timeout-minutes: 360
     env:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
+      PYTORCH_ROOT: ${{ github.workspace }}
       PACKAGE_TYPE: wheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
@@ -8242,7 +8007,7 @@ jobs:
           echo "instance-type: $(get_ec2_metadata instance-type)"
           echo "system info $(uname -a)"
       - name: "[FB EMPLOYEES] Enable SSH (Click me for login details)"
-        uses: pytorch/test-infra/.github/actions/setup-ssh@main
+        uses: pytorch/test-infra/.github/actions/setup-ssh@release/2.12
         continue-on-error: true
         with:
           github-secret: ${{ secrets.GITHUB_TOKEN }}
@@ -8274,15 +8039,12 @@ jobs:
       - name: Checkout PyTorch
         uses: actions/checkout@v4
         with:
-          ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
           submodules: recursive
-          path: pytorch
           show-progress: false
       - name: Clean PyTorch checkout
         run: |
           # Remove any artifacts from the previous checkouts
           git clean -fxd
-        working-directory: pytorch
       # NOTE: These environment variables are put here so that they can be applied on every job equally
       #       They are also here because setting them at a workflow level doesn't give us access to the
       #       runner.temp variable, which we need.
@@ -8300,21 +8062,19 @@ jobs:
       - name: Populate binary env
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_populate_env.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_populate_env.sh"
       - name: Test PyTorch binary
         shell: bash
         run: |
-          "${PYTORCH_ROOT}/.circleci/scripts/binary_windows_test.sh"
+          "${PYTORCH_ROOT}/.ci/pytorch/binary_windows_test.sh"
       - name: Wait until all sessions have drained
         shell: powershell
-        working-directory: pytorch
         if: always()
         timeout-minutes: 120
         run: |
           .github\scripts\wait_for_ssh_to_drain.ps1
       - name: Kill active ssh sessions if still around (Useful if workflow was cancelled)
         shell: powershell
-        working-directory: pytorch
         if: always()
         run: |
           .github\scripts\kill_active_ssh_sessions.ps1
@@ -8325,7 +8085,7 @@ jobs:
       contents: read
     needs: wheel-py3_14t-xpu-test
     with:
-      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
+      PYTORCH_ROOT: ${{ github.workspace }}
       PACKAGE_TYPE: wheel
       # TODO: This is a legacy variable that we eventually want to get rid of in
       #       favor of GPU_ARCH_VERSION
@@ -8335,4 +8095,250 @@ jobs:
       build_name: wheel-py3_14t-xpu
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
+    uses: ./.github/workflows/_binary-upload.yml
+  libtorch-cpu-shared-with-deps-release-extract:
+    if: ${{ github.repository_owner == 'pytorch' }}
+    needs:
+      - wheel-py3_10-cpu-build
+      - get-label-type
+    runs-on: "${{ needs.get-label-type.outputs.label-type }}linux.4xlarge"
+    timeout-minutes: 60
+    env:
+      DESIRED_CUDA: cpu
+      LIBTORCH_VARIANT: shared-with-deps
+    steps:
+      - name: Checkout PyTorch
+        uses: actions/checkout@v4
+        with:
+          ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
+          sparse-checkout: .ci/libtorch/
+          show-progress: false
+      - uses: actions/download-artifact@v4.1.7
+        name: Download Wheel Artifact
+        with:
+          name: wheel-py3_10-cpu
+          path: "${{ runner.temp }}/wheel_artifact/"
+      - name: Extract libtorch from wheel
+        shell: bash
+        run: |
+          set -eux
+          mkdir -p "${{ runner.temp }}/libtorch_output"
+          python3 .ci/libtorch/extract_libtorch_from_wheel.py \
+            --wheel-dir "${{ runner.temp }}/wheel_artifact" \
+            --output-dir "${{ runner.temp }}/libtorch_output" \
+            --platform windows \
+            --desired-cuda "$DESIRED_CUDA" \
+            --libtorch-variant "$LIBTORCH_VARIANT"
+      - uses: actions/upload-artifact@v4.4.0
+        if: always()
+        with:
+          name: libtorch-cpu-shared-with-deps-release
+          retention-days: 14
+          if-no-files-found: error
+          path: "${{ runner.temp }}/libtorch_output/"
+  libtorch-cpu-shared-with-deps-release-upload:  # Uploading
+    if: ${{ github.repository_owner == 'pytorch' }}
+    permissions:
+      id-token: write
+      contents: read
+    needs: libtorch-cpu-shared-with-deps-release-extract
+    with:
+      PYTORCH_ROOT: ${{ github.workspace }}
+      PACKAGE_TYPE: libtorch
+      DESIRED_CUDA: cpu
+      GPU_ARCH_TYPE: cpu
+      LIBTORCH_CONFIG: release
+      LIBTORCH_VARIANT: shared-with-deps
+      build_name: libtorch-cpu-shared-with-deps-release
+    secrets:
+      github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
+    uses: ./.github/workflows/_binary-upload.yml
+  libtorch-cuda12_6-shared-with-deps-release-extract:
+    if: ${{ github.repository_owner == 'pytorch' }}
+    needs:
+      - wheel-py3_10-cuda12_6-build
+      - get-label-type
+    runs-on: "${{ needs.get-label-type.outputs.label-type }}linux.4xlarge"
+    timeout-minutes: 60
+    env:
+      DESIRED_CUDA: cu126
+      LIBTORCH_VARIANT: shared-with-deps
+    steps:
+      - name: Checkout PyTorch
+        uses: actions/checkout@v4
+        with:
+          ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
+          sparse-checkout: .ci/libtorch/
+          show-progress: false
+      - uses: actions/download-artifact@v4.1.7
+        name: Download Wheel Artifact
+        with:
+          name: wheel-py3_10-cuda12_6
+          path: "${{ runner.temp }}/wheel_artifact/"
+      - name: Extract libtorch from wheel
+        shell: bash
+        run: |
+          set -eux
+          mkdir -p "${{ runner.temp }}/libtorch_output"
+          python3 .ci/libtorch/extract_libtorch_from_wheel.py \
+            --wheel-dir "${{ runner.temp }}/wheel_artifact" \
+            --output-dir "${{ runner.temp }}/libtorch_output" \
+            --platform windows \
+            --desired-cuda "$DESIRED_CUDA" \
+            --libtorch-variant "$LIBTORCH_VARIANT"
+      - uses: actions/upload-artifact@v4.4.0
+        if: always()
+        with:
+          name: libtorch-cuda12_6-shared-with-deps-release
+          retention-days: 14
+          if-no-files-found: error
+          path: "${{ runner.temp }}/libtorch_output/"
+  libtorch-cuda12_6-shared-with-deps-release-upload:  # Uploading
+    if: ${{ github.repository_owner == 'pytorch' }}
+    permissions:
+      id-token: write
+      contents: read
+    needs: libtorch-cuda12_6-shared-with-deps-release-extract
+    with:
+      PYTORCH_ROOT: ${{ github.workspace }}
+      PACKAGE_TYPE: libtorch
+      DESIRED_CUDA: cu126
+      GPU_ARCH_VERSION: "12.6"
+      GPU_ARCH_TYPE: cuda
+      LIBTORCH_CONFIG: release
+      LIBTORCH_VARIANT: shared-with-deps
+      build_name: libtorch-cuda12_6-shared-with-deps-release
+    secrets:
+      github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
+    uses: ./.github/workflows/_binary-upload.yml
+  libtorch-cuda13_0-shared-with-deps-release-extract:
+    if: ${{ github.repository_owner == 'pytorch' }}
+    needs:
+      - wheel-py3_10-cuda13_0-build
+      - get-label-type
+    runs-on: "${{ needs.get-label-type.outputs.label-type }}linux.4xlarge"
+    timeout-minutes: 60
+    env:
+      DESIRED_CUDA: cu130
+      LIBTORCH_VARIANT: shared-with-deps
+    steps:
+      - name: Checkout PyTorch
+        uses: actions/checkout@v4
+        with:
+          ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
+          sparse-checkout: .ci/libtorch/
+          show-progress: false
+      - uses: actions/download-artifact@v4.1.7
+        name: Download Wheel Artifact
+        with:
+          name: wheel-py3_10-cuda13_0
+          path: "${{ runner.temp }}/wheel_artifact/"
+      - name: Extract libtorch from wheel
+        shell: bash
+        run: |
+          set -eux
+          mkdir -p "${{ runner.temp }}/libtorch_output"
+          python3 .ci/libtorch/extract_libtorch_from_wheel.py \
+            --wheel-dir "${{ runner.temp }}/wheel_artifact" \
+            --output-dir "${{ runner.temp }}/libtorch_output" \
+            --platform windows \
+            --desired-cuda "$DESIRED_CUDA" \
+            --libtorch-variant "$LIBTORCH_VARIANT"
+      - uses: actions/upload-artifact@v4.4.0
+        if: always()
+        with:
+          name: libtorch-cuda13_0-shared-with-deps-release
+          retention-days: 14
+          if-no-files-found: error
+          path: "${{ runner.temp }}/libtorch_output/"
+  libtorch-cuda13_0-shared-with-deps-release-upload:  # Uploading
+    if: ${{ github.repository_owner == 'pytorch' }}
+    permissions:
+      id-token: write
+      contents: read
+    needs: libtorch-cuda13_0-shared-with-deps-release-extract
+    with:
+      PYTORCH_ROOT: ${{ github.workspace }}
+      PACKAGE_TYPE: libtorch
+      DESIRED_CUDA: cu130
+      GPU_ARCH_VERSION: "13.0"
+      GPU_ARCH_TYPE: cuda
+      LIBTORCH_CONFIG: release
+      LIBTORCH_VARIANT: shared-with-deps
+      build_name: libtorch-cuda13_0-shared-with-deps-release
+    secrets:
+      github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
+    uses: ./.github/workflows/_binary-upload.yml
+  libtorch-cuda13_2-shared-with-deps-release-extract:
+    if: ${{ github.repository_owner == 'pytorch' }}
+    needs:
+      - wheel-py3_10-cuda13_2-build
+      - get-label-type
+    runs-on: "${{ needs.get-label-type.outputs.label-type }}linux.4xlarge"
+    timeout-minutes: 60
+    env:
+      DESIRED_CUDA: cu132
+      LIBTORCH_VARIANT: shared-with-deps
+    steps:
+      - name: Checkout PyTorch
+        uses: actions/checkout@v4
+        with:
+          ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
+          sparse-checkout: .ci/libtorch/
+          show-progress: false
+      - uses: actions/download-artifact@v4.1.7
+        name: Download Wheel Artifact
+        with:
+          name: wheel-py3_10-cuda13_2
+          path: "${{ runner.temp }}/wheel_artifact/"
+      - name: Extract libtorch from wheel
+        shell: bash
+        run: |
+          set -eux
+          mkdir -p "${{ runner.temp }}/libtorch_output"
+          python3 .ci/libtorch/extract_libtorch_from_wheel.py \
+            --wheel-dir "${{ runner.temp }}/wheel_artifact" \
+            --output-dir "${{ runner.temp }}/libtorch_output" \
+            --platform windows \
+            --desired-cuda "$DESIRED_CUDA" \
+            --libtorch-variant "$LIBTORCH_VARIANT"
+      - uses: actions/upload-artifact@v4.4.0
+        if: always()
+        with:
+          name: libtorch-cuda13_2-shared-with-deps-release
+          retention-days: 14
+          if-no-files-found: error
+          path: "${{ runner.temp }}/libtorch_output/"
+  libtorch-cuda13_2-shared-with-deps-release-upload:  # Uploading
+    if: ${{ github.repository_owner == 'pytorch' }}
+    permissions:
+      id-token: write
+      contents: read
+    needs: libtorch-cuda13_2-shared-with-deps-release-extract
+    with:
+      PYTORCH_ROOT: ${{ github.workspace }}
+      PACKAGE_TYPE: libtorch
+      DESIRED_CUDA: cu132
+      GPU_ARCH_VERSION: "13.2"
+      GPU_ARCH_TYPE: cuda
+      LIBTORCH_CONFIG: release
+      LIBTORCH_VARIANT: shared-with-deps
+      build_name: libtorch-cuda13_2-shared-with-deps-release
+    secrets:
+      github-token: ${{ secrets.GITHUB_TOKEN }}
+      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
+      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
+      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
     uses: ./.github/workflows/_binary-upload.yml
diff --git a/.github/workflows/h100-cutlass-backend.yml b/.github/workflows/h100-cutlass-backend.yml
index e5406f7600133..c3778fcc9c8ef 100644
--- a/.github/workflows/h100-cutlass-backend.yml
+++ b/.github/workflows/h100-cutlass-backend.yml
@@ -21,13 +21,14 @@ concurrency:
 permissions:
   id-token: write
   contents: read
+  actions: read
 
 jobs:
 
   get-label-type:
     if: github.repository_owner == 'pytorch'
     name: get-label-type
-    uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@main
+    uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@release/2.12
     with:
       triggering_actor: ${{ github.triggering_actor }}
       issue_owner: ${{ github.event.pull_request.user.login || github.event.issue.user.login }}
diff --git a/.github/workflows/h100-distributed.yml b/.github/workflows/h100-distributed.yml
index 0e5370a51c160..b004a4b38a481 100644
--- a/.github/workflows/h100-distributed.yml
+++ b/.github/workflows/h100-distributed.yml
@@ -18,13 +18,14 @@ concurrency:
 permissions:
   id-token: write
   contents: read
+  actions: read
 
 jobs:
 
   get-label-type:
     if: github.repository_owner == 'pytorch'
     name: get-label-type
-    uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@main
+    uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@release/2.12
     with:
       triggering_actor: ${{ github.triggering_actor }}
       issue_owner: ${{ github.event.pull_request.user.login || github.event.issue.user.login }}
diff --git a/.github/workflows/h100-symm-mem.yml b/.github/workflows/h100-symm-mem.yml
index 09c362a546024..f00f1d0583ff2 100644
--- a/.github/workflows/h100-symm-mem.yml
+++ b/.github/workflows/h100-symm-mem.yml
@@ -18,13 +18,14 @@ concurrency:
 permissions:
   id-token: write
   contents: read
+  actions: read
 
 jobs:
 
   get-label-type:
     if: github.repository_owner == 'pytorch'
     name: get-label-type
-    uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@main
+    uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@release/2.12
     with:
       triggering_actor: ${{ github.triggering_actor }}
       issue_owner: ${{ github.event.pull_request.user.login || github.event.issue.user.login }}
diff --git a/.github/workflows/inductor-micro-benchmark-x86.yml b/.github/workflows/inductor-micro-benchmark-x86.yml
index b0e9c869cdcf0..744ef7d7f75ce 100644
--- a/.github/workflows/inductor-micro-benchmark-x86.yml
+++ b/.github/workflows/inductor-micro-benchmark-x86.yml
@@ -16,6 +16,7 @@ concurrency:
 permissions:
   id-token: write
   contents: read
+  actions: read
 
 jobs:
   inductor-build:
diff --git a/.github/workflows/inductor-micro-benchmark.yml b/.github/workflows/inductor-micro-benchmark.yml
index 5813aa28365e7..0bc84e9a26b39 100644
--- a/.github/workflows/inductor-micro-benchmark.yml
+++ b/.github/workflows/inductor-micro-benchmark.yml
@@ -16,11 +16,12 @@ concurrency:
 permissions:
   id-token: write
   contents: read
+  actions: read
 
 jobs:
   get-default-label-prefix:
     name: get-default-label-prefix
-    uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@main
+    uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@release/2.12
     if: ${{ (github.event_name != 'schedule' || github.repository == 'pytorch/pytorch') && github.repository_owner == 'pytorch' }}
     with:
       triggering_actor: ${{ github.triggering_actor }}
@@ -30,14 +31,14 @@ jobs:
       opt_out_experiments: lf
 
   build:
-    name: cuda12.8-py3.10-gcc11-sm80
+    name: cuda13.0-py3.10-gcc11-sm80
     uses: ./.github/workflows/_linux-build.yml
     needs:
       - get-default-label-prefix
     with:
       runner_prefix: "${{ needs.get-default-label-prefix.outputs.label-type }}"
-      build-environment: linux-jammy-cuda12.8-py3.10-gcc11-sm80
-      docker-image-name: ci-image:pytorch-linux-jammy-cuda12.8-cudnn9-py3-gcc11-inductor-benchmarks
+      build-environment: linux-jammy-cuda13.0-py3.10-gcc11-sm80
+      docker-image-name: ci-image:pytorch-linux-jammy-cuda13.0-cudnn9-py3-gcc11-inductor-benchmarks
       cuda-arch-list: '8.0'
       test-matrix: |
         { include: [
@@ -46,7 +47,7 @@ jobs:
     secrets: inherit
 
   test:
-    name: cuda12.8-py3.10-gcc11-sm80
+    name: cuda13.0-py3.10-gcc11-sm80
     uses: ./.github/workflows/_linux-test.yml
     needs: build
     with:
@@ -55,30 +56,3 @@ jobs:
       test-matrix: ${{ needs.build.outputs.test-matrix }}
       timeout-minutes: 720
     secrets: inherit
-
-  build-cuda13:
-    name: cuda13.0-py3.10-gcc11-sm80
-    uses: ./.github/workflows/_linux-build.yml
-    needs:
-      - get-default-label-prefix
-    with:
-      runner_prefix: "${{ needs.get-default-label-prefix.outputs.label-type }}"
-      build-environment: linux-jammy-cuda13.0-py3.10-gcc11-sm80
-      docker-image-name: ci-image:pytorch-linux-jammy-cuda13.0-cudnn9-py3-gcc11-inductor-benchmarks
-      cuda-arch-list: '8.0'
-      test-matrix: |
-        { include: [
-          { config: "inductor-micro-benchmark", shard: 1, num_shards: 1, runner: "linux.aws.a100", owners: ["oncall:pt2"] },
-        ]}
-    secrets: inherit
-
-  test-cuda13:
-    name: cuda13.0-py3.10-gcc11-sm80
-    uses: ./.github/workflows/_linux-test.yml
-    needs: build-cuda13
-    with:
-      build-environment: linux-jammy-cuda13.0-py3.10-gcc11-sm80
-      docker-image: ${{ needs.build-cuda13.outputs.docker-image }}
-      test-matrix: ${{ needs.build-cuda13.outputs.test-matrix }}
-      timeout-minutes: 720
-    secrets: inherit
diff --git a/.github/workflows/inductor-nightly.yml b/.github/workflows/inductor-nightly.yml
index 4258e8fdb0c84..dd950d4b1b099 100644
--- a/.github/workflows/inductor-nightly.yml
+++ b/.github/workflows/inductor-nightly.yml
@@ -19,11 +19,12 @@ concurrency:
 permissions:
   id-token: write
   contents: read
+  actions: read
 
 jobs:
   get-default-label-prefix:
     name: get-default-label-prefix
-    uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@main
+    uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@release/2.12
     if: ${{ (github.event_name != 'schedule' || github.repository == 'pytorch/pytorch') && github.repository_owner == 'pytorch' }}
     with:
       triggering_actor: ${{ github.triggering_actor }}
diff --git a/.github/workflows/inductor-pallas.yml b/.github/workflows/inductor-pallas.yml
index 8676434d0e580..88e43044b5af0 100644
--- a/.github/workflows/inductor-pallas.yml
+++ b/.github/workflows/inductor-pallas.yml
@@ -16,11 +16,12 @@ concurrency:
 permissions:
   id-token: write
   contents: read
+  actions: read
 
 jobs:
   get-label-type:
     name: get-label-type
-    uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@main
+    uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@release/2.12
     if: ${{ (github.event_name != 'schedule' || github.repository == 'pytorch/pytorch') && github.repository_owner == 'pytorch' }}
     with:
       triggering_actor: ${{ github.triggering_actor }}
@@ -46,9 +47,6 @@ jobs:
     secrets: inherit
 
   linux-jammy-py3_12-inductor-pallas-gpu-test:
-    permissions:
-      id-token: write
-      contents: read
     name: pallas-gpu-py3.12-inductor
     uses: ./.github/workflows/_linux-test.yml
     needs: linux-jammy-py3_12-inductor-pallas-gpu-build
@@ -73,9 +71,6 @@ jobs:
     secrets: inherit
 
   linux-jammy-py3_12-inductor-pallas-tpu-test:
-    permissions:
-      id-token: write
-      contents: read
     name: pallas-tpu-py3.12-inductor
     uses: ./.github/workflows/_linux-test.yml
     needs: linux-jammy-py3_12-inductor-pallas-tpu-build
diff --git a/.github/workflows/inductor-perf-compare.yml b/.github/workflows/inductor-perf-compare.yml
index 17b265500d47e..93f3d30fea56e 100644
--- a/.github/workflows/inductor-perf-compare.yml
+++ b/.github/workflows/inductor-perf-compare.yml
@@ -13,12 +13,13 @@ concurrency:
 permissions:
   id-token: write
   contents: read
+  actions: read
 
 jobs:
   get-default-label-prefix:
     if: github.repository_owner == 'pytorch'
     name: get-default-label-prefix
-    uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@main
+    uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@release/2.12
     with:
       triggering_actor: ${{ github.triggering_actor }}
       issue_owner: ${{ github.event.pull_request.user.login || github.event.issue.user.login }}
@@ -27,15 +28,15 @@ jobs:
       opt_out_experiments: lf
 
   build:
-    name: cuda12.8-py3.10-gcc11-sm80
+    name: cuda13.0-py3.10-gcc11-sm80
     uses: ./.github/workflows/_linux-build.yml
     needs:
       - get-default-label-prefix
     with:
       runner_prefix: "${{ needs.get-default-label-prefix.outputs.label-type }}"
       runner: linux.4xlarge.memory
-      build-environment: linux-jammy-cuda12.8-py3.10-gcc11-sm80
-      docker-image-name: ci-image:pytorch-linux-jammy-cuda12.8-cudnn9-py3-gcc11-inductor-benchmarks
+      build-environment: linux-jammy-cuda13.0-py3.10-gcc11-sm80
+      docker-image-name: ci-image:pytorch-linux-jammy-cuda13.0-cudnn9-py3-gcc11-inductor-benchmarks
       cuda-arch-list: '8.0'
       test-matrix: |
         { include: [
@@ -48,7 +49,7 @@ jobs:
     secrets: inherit
 
   test:
-    name: cuda12.8-py3.10-gcc11-sm80
+    name: cuda13.0-py3.10-gcc11-sm80
     uses: ./.github/workflows/_linux-test.yml
     needs: build
     with:
@@ -59,39 +60,5 @@ jobs:
       disable-monitor: false
       monitor-log-interval: 15
       monitor-data-collect-interval: 4
-    secrets: inherit
-
-  build-cuda13:
-    name: cuda13.0-py3.10-gcc11-sm80
-    uses: ./.github/workflows/_linux-build.yml
-    needs:
-      - get-default-label-prefix
-    with:
-      runner_prefix: "${{ needs.get-default-label-prefix.outputs.label-type }}"
-      runner: linux.4xlarge.memory
-      build-environment: linux-jammy-cuda13.0-py3.10-gcc11-sm80
-      docker-image-name: ci-image:pytorch-linux-jammy-cuda13.0-cudnn9-py3-gcc11-inductor-benchmarks
-      cuda-arch-list: '8.0'
-      test-matrix: |
-        { include: [
-          { config: "inductor_huggingface_perf_compare", shard: 1, num_shards: 1, runner: "linux.aws.a100" },
-          { config: "inductor_timm_perf_compare", shard: 1, num_shards: 2, runner: "linux.aws.a100" },
-          { config: "inductor_timm_perf_compare", shard: 2, num_shards: 2, runner: "linux.aws.a100" },
-          { config: "inductor_torchbench_perf_compare", shard: 1, num_shards: 1, runner: "linux.aws.a100" },
-        ]}
-      build-additional-packages: "vision audio torchao"
-    secrets: inherit
-
-  test-cuda13:
-    name: cuda13.0-py3.10-gcc11-sm80
-    uses: ./.github/workflows/_linux-test.yml
-    needs: build-cuda13
-    with:
-      build-environment: linux-jammy-cuda13.0-py3.10-gcc11-sm80
-      docker-image: ${{ needs.build-cuda13.outputs.docker-image }}
-      test-matrix: ${{ needs.build-cuda13.outputs.test-matrix }}
-      # disable monitor in perf tests for more investigation
-      disable-monitor: false
-      monitor-log-interval: 15
-      monitor-data-collect-interval: 4
+      export-profiler-trace: "1"
     secrets: inherit
diff --git a/.github/workflows/inductor-perf-test-b200.yml b/.github/workflows/inductor-perf-test-b200.yml
index 07c7a0cf987d7..029ad1fe8a019 100644
--- a/.github/workflows/inductor-perf-test-b200.yml
+++ b/.github/workflows/inductor-perf-test-b200.yml
@@ -64,11 +64,12 @@ concurrency:
 permissions:
   id-token: write
   contents: read
+  actions: read
 
 jobs:
   get-label-type:
     name: get-label-type
-    uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@main
+    uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@release/2.12
     if: ${{ (github.event_name != 'schedule' || github.repository == 'pytorch/pytorch') && github.repository_owner == 'pytorch' }}
     with:
       triggering_actor: ${{ github.triggering_actor }}
@@ -78,7 +79,7 @@ jobs:
       opt_out_experiments: lf
 
   build:
-    name: cuda12.8-py3.10-gcc11-sm100
+    name: cuda13.0-py3.10-gcc11-sm100
     uses: ./.github/workflows/_linux-build.yml
     needs: get-label-type
     with:
@@ -88,8 +89,8 @@ jobs:
       # from trunk. Also use a memory-intensive runner here because memory is
       # usually the bottleneck
       runner: linux.12xlarge.memory
-      build-environment: linux-jammy-cuda12.8-py3.10-gcc11-sm100
-      docker-image-name: ci-image:pytorch-linux-jammy-cuda12.8-cudnn9-py3-gcc11-inductor-benchmarks
+      build-environment: linux-jammy-cuda13.0-py3.10-gcc11-sm100
+      docker-image-name: ci-image:pytorch-linux-jammy-cuda13.0-cudnn9-py3-gcc11-inductor-benchmarks
       cuda-arch-list: '10.0'
       test-matrix: |
         { include: [
@@ -102,7 +103,7 @@ jobs:
     secrets: inherit
 
   test-periodically:
-    name: cuda12.8-py3.10-gcc11-sm100
+    name: cuda13.0-py3.10-gcc11-sm100
     uses: ./.github/workflows/_linux-test.yml
     needs: build
     if: github.event.schedule == '0 7 * * 1-6'
@@ -119,7 +120,7 @@ jobs:
     secrets: inherit
 
   test-weekly:
-    name: cuda12.8-py3.10-gcc11-sm100
+    name: cuda13.0-py3.10-gcc11-sm100
     uses: ./.github/workflows/_linux-test.yml
     needs: build
     if: github.event.schedule == '0 7 * * 0'
@@ -136,7 +137,7 @@ jobs:
     secrets: inherit
 
   test:
-    name: cuda12.8-py3.10-gcc11-sm100
+    name: cuda13.0-py3.10-gcc11-sm100
     uses: ./.github/workflows/_linux-test.yml
     needs: build
     with:
diff --git a/.github/workflows/inductor-perf-test-nightly-aarch64.yml b/.github/workflows/inductor-perf-test-nightly-aarch64.yml
index f7b3517dccc06..1e1524495a402 100644
--- a/.github/workflows/inductor-perf-test-nightly-aarch64.yml
+++ b/.github/workflows/inductor-perf-test-nightly-aarch64.yml
@@ -51,11 +51,12 @@ concurrency:
 permissions:
   id-token: write
   contents: read
+  actions: read
 
 jobs:
   get-label-type:
     name: get-label-type
-    uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@main
+    uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@release/2.12
     if: ${{ (github.event_name != 'schedule' || github.repository == 'pytorch/pytorch') && github.repository_owner == 'pytorch' }}
     with:
       triggering_actor: ${{ github.triggering_actor }}
diff --git a/.github/workflows/inductor-perf-test-nightly-h100.yml b/.github/workflows/inductor-perf-test-nightly-h100.yml
index 155b995e3c8a2..5b5f89bee0073 100644
--- a/.github/workflows/inductor-perf-test-nightly-h100.yml
+++ b/.github/workflows/inductor-perf-test-nightly-h100.yml
@@ -51,6 +51,11 @@ on:
         required: false
         type: boolean
         default: false
+      deterministic_perf:
+        description: Run benchmarks with deterministic mode enabled?
+        required: false
+        type: boolean
+        default: false
       benchmark_configs:
         description: The list of configs used the benchmark
         required: false
@@ -69,11 +74,12 @@ concurrency:
 permissions:
   id-token: write
   contents: read
+  actions: read
 
 jobs:
   get-label-type:
     name: get-label-type
-    uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@main
+    uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@release/2.12
     if: ${{ (github.event_name != 'schedule' || github.repository == 'pytorch/pytorch') && github.repository_owner == 'pytorch' }}
     with:
       triggering_actor: ${{ github.triggering_actor }}
@@ -93,8 +99,8 @@ jobs:
       # from trunk. Also use a memory-intensive runner here because memory is
       # usually the bottleneck
       runner: linux.12xlarge.memory
-      build-environment: linux-jammy-cuda12.8-py3.10-gcc11-sm90
-      docker-image-name: ci-image:pytorch-linux-jammy-cuda12.8-cudnn9-py3-gcc11-inductor-benchmarks
+      build-environment: linux-jammy-cuda13.0-py3.10-gcc11-sm90
+      docker-image-name: ci-image:pytorch-linux-jammy-cuda13.0-cudnn9-py3-gcc11-inductor-benchmarks
       cuda-arch-list: '9.0'
       test-matrix: |
         { include: [
@@ -131,7 +137,7 @@ jobs:
     if: github.event.schedule == '15 0 * * 1-6'
     with:
       build-environment: ${{ needs.build.outputs.build-environment }}
-      dashboard-tag: training-true-inference-true-default-true-dynamic-true-cudagraphs-true-cppwrapper-true-aotinductor-true-freezing_cudagraphs-true
+      dashboard-tag: training-true-inference-true-default-true-dynamic-true-cudagraphs-true-cppwrapper-true-aotinductor-true-freezing_cudagraphs-true-deterministic_perf-true
       docker-image: ${{ needs.build.outputs.docker-image }}
       test-matrix: ${{ needs.build.outputs.test-matrix }}
       timeout-minutes: 720
@@ -139,6 +145,8 @@ jobs:
       disable-monitor: false
       monitor-log-interval: 15
       monitor-data-collect-interval: 4
+      export-profiler-trace: "1"
+      enable-torch-trace: "1"
     secrets: inherit
 
   test-weekly:
@@ -148,7 +156,7 @@ jobs:
     if: github.event.schedule == '0 7 * * 0'
     with:
       build-environment: ${{ needs.build.outputs.build-environment }}
-      dashboard-tag: training-true-inference-true-default-true-dynamic-true-cudagraphs-true-cppwrapper-true-aotinductor-true-freezing_cudagraphs-true-maxautotune-true-freeze_autotune_cudagraphs-true
+      dashboard-tag: training-true-inference-true-default-true-dynamic-true-cudagraphs-true-cppwrapper-true-aotinductor-true-freezing_cudagraphs-true-maxautotune-true-freeze_autotune_cudagraphs-true-deterministic_perf-true
       docker-image: ${{ needs.build.outputs.docker-image }}
       test-matrix: ${{ needs.build.outputs.test-matrix }}
       timeout-minutes: 1440
@@ -156,6 +164,8 @@ jobs:
       disable-monitor: false
       monitor-log-interval: 15
       monitor-data-collect-interval: 4
+      export-profiler-trace: "1"
+      enable-torch-trace: "1"
     secrets: inherit
 
   test:
@@ -167,7 +177,7 @@ jobs:
     if: ${{ github.event_name == 'workflow_dispatch' || github.event_name == 'pull_request' }}
     with:
       build-environment: ${{ needs.build.outputs.build-environment }}
-      dashboard-tag: training-${{ inputs.training || 'true' }}-inference-${{ inputs.inference || 'true' }}-default-${{ inputs.default || 'true' }}-dynamic-${{ inputs.dynamic || 'true' }}-cudagraphs-${{ inputs.cudagraphs || 'true' }}-cppwrapper-${{ inputs.cppwrapper || 'false' }}-aotinductor-${{ inputs.aotinductor || 'false' }}-maxautotune-${{ inputs.maxautotune || 'false' }}-freezing_cudagraphs-${{ inputs.freezing_cudagraphs || 'false' }}
+      dashboard-tag: training-${{ inputs.training || 'true' }}-inference-${{ inputs.inference || 'true' }}-default-${{ inputs.default || 'true' }}-dynamic-${{ inputs.dynamic || 'true' }}-cudagraphs-${{ inputs.cudagraphs || 'true' }}-cppwrapper-${{ inputs.cppwrapper || 'false' }}-aotinductor-${{ inputs.aotinductor || 'false' }}-maxautotune-${{ inputs.maxautotune || 'false' }}-freezing_cudagraphs-${{ inputs.freezing_cudagraphs || 'false' }}-deterministic_perf-${{ inputs.deterministic_perf || 'false' }}
       docker-image: ${{ needs.build.outputs.docker-image }}
       test-matrix: ${{ needs.build.outputs.test-matrix }}
       timeout-minutes: 720
@@ -175,4 +185,6 @@ jobs:
       disable-monitor: false
       monitor-log-interval: 15
       monitor-data-collect-interval: 4
+      export-profiler-trace: "1"
+      enable-torch-trace: "1"
     secrets: inherit
diff --git a/.github/workflows/inductor-perf-test-nightly-rocm-mi300.yml b/.github/workflows/inductor-perf-test-nightly-rocm-mi300.yml
index c556c6b455783..f904bd7ab7b7f 100644
--- a/.github/workflows/inductor-perf-test-nightly-rocm-mi300.yml
+++ b/.github/workflows/inductor-perf-test-nightly-rocm-mi300.yml
@@ -68,7 +68,7 @@ permissions: read-all
 jobs:
   get-label-type:
     name: get-label-type
-    uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@main
+    uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@release/2.12
     if: ${{ (github.event_name != 'schedule' || github.repository == 'pytorch/pytorch') && github.repository_owner == 'pytorch' }}
     with:
       triggering_actor: ${{ github.triggering_actor }}
@@ -79,10 +79,10 @@ jobs:
 
   linux-jammy-rocm-py3_10-inductor-benchmark-build:
     if: github.repository_owner == 'pytorch'
-    name: rocm-py3_10-inductor-benchmark-build
+    name: linux-jammy-rocm-py3.10-mi300
     uses: ./.github/workflows/_linux-build.yml
     with:
-      build-environment: linux-jammy-rocm-py3_10
+      build-environment: linux-jammy-rocm-py3.10-mi300
       docker-image-name: ci-image:pytorch-linux-jammy-rocm-n-py3-benchmarks
       test-matrix: |
         { include: [
@@ -114,7 +114,7 @@ jobs:
     permissions:
       id-token: write
       contents: read
-    name: rocm-py3_10-inductor-benchmark-test
+    name: linux-jammy-rocm-py3.10-mi300
     uses: ./.github/workflows/_rocm-test.yml
     needs: linux-jammy-rocm-py3_10-inductor-benchmark-build
     with:
diff --git a/.github/workflows/inductor-perf-test-nightly-rocm-mi355.yml b/.github/workflows/inductor-perf-test-nightly-rocm-mi355.yml
index e6fd83193202c..7e06a39b4522b 100644
--- a/.github/workflows/inductor-perf-test-nightly-rocm-mi355.yml
+++ b/.github/workflows/inductor-perf-test-nightly-rocm-mi355.yml
@@ -68,7 +68,7 @@ permissions: read-all
 jobs:
   get-label-type:
     name: get-label-type
-    uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@main
+    uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@release/2.12
     if: ${{ (github.event_name != 'schedule' || github.repository == 'pytorch/pytorch') && github.repository_owner == 'pytorch' }}
     with:
       triggering_actor: ${{ github.triggering_actor }}
@@ -79,10 +79,10 @@ jobs:
 
   linux-jammy-rocm-py3_10-inductor-benchmark-build:
     if: github.repository_owner == 'pytorch'
-    name: rocm-py3_10-inductor-benchmark-build
+    name: linux-jammy-rocm-py3.10-mi355
     uses: ./.github/workflows/_linux-build.yml
     with:
-      build-environment: linux-jammy-rocm-py3_10
+      build-environment: linux-jammy-rocm-py3.10-mi355
       docker-image-name: ci-image:pytorch-linux-jammy-rocm-n-py3-benchmarks
       test-matrix: |
         { include: [
@@ -114,7 +114,7 @@ jobs:
     permissions:
       id-token: write
       contents: read
-    name: rocm-py3_10-inductor-benchmark-test
+    name: linux-jammy-rocm-py3.10-mi355
     uses: ./.github/workflows/_rocm-test.yml
     needs: linux-jammy-rocm-py3_10-inductor-benchmark-build
     with:
diff --git a/.github/workflows/inductor-perf-test-nightly-x86-zen.yml b/.github/workflows/inductor-perf-test-nightly-x86-zen.yml
index eee51b7ff8889..7f78e1706923e 100644
--- a/.github/workflows/inductor-perf-test-nightly-x86-zen.yml
+++ b/.github/workflows/inductor-perf-test-nightly-x86-zen.yml
@@ -61,11 +61,12 @@ concurrency:
 permissions:
   id-token: write
   contents: read
+  actions: read
 
 jobs:
   get-label-type:
     name: get-label-type
-    uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@main
+    uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@release/2.12
     if: ${{ (github.event_name != 'schedule' || github.repository == 'pytorch/pytorch') && github.repository_owner == 'pytorch' }}
     with:
       triggering_actor: ${{ github.triggering_actor }}
diff --git a/.github/workflows/inductor-perf-test-nightly-x86.yml b/.github/workflows/inductor-perf-test-nightly-x86.yml
index 87875831e2a0b..ef1a7224c9098 100644
--- a/.github/workflows/inductor-perf-test-nightly-x86.yml
+++ b/.github/workflows/inductor-perf-test-nightly-x86.yml
@@ -61,11 +61,12 @@ concurrency:
 permissions:
   id-token: write
   contents: read
+  actions: read
 
 jobs:
   get-label-type:
     name: get-label-type
-    uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@main
+    uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@release/2.12
     if: ${{ (github.event_name != 'schedule' || github.repository == 'pytorch/pytorch') && github.repository_owner == 'pytorch' }}
     with:
       triggering_actor: ${{ github.triggering_actor }}
diff --git a/.github/workflows/inductor-perf-test-nightly-xpu.yml b/.github/workflows/inductor-perf-test-nightly-xpu.yml
index b51795c663957..404c6db4fc381 100644
--- a/.github/workflows/inductor-perf-test-nightly-xpu.yml
+++ b/.github/workflows/inductor-perf-test-nightly-xpu.yml
@@ -68,7 +68,7 @@ permissions: read-all
 jobs:
   get-label-type:
     name: get-label-type
-    uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@main
+    uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@release/2.12
     if: ${{ (github.event_name != 'schedule' || github.repository == 'pytorch/pytorch') && github.repository_owner == 'pytorch' }}
     with:
       triggering_actor: ${{ github.triggering_actor }}
diff --git a/.github/workflows/inductor-perf-test-nightly.yml b/.github/workflows/inductor-perf-test-nightly.yml
index 2a8e29278b8bc..ddcdc85e20b2f 100644
--- a/.github/workflows/inductor-perf-test-nightly.yml
+++ b/.github/workflows/inductor-perf-test-nightly.yml
@@ -64,11 +64,12 @@ concurrency:
 permissions:
   id-token: write
   contents: read
+  actions: read
 
 jobs:
   get-label-type:
     name: get-label-type
-    uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@main
+    uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@release/2.12
     if: ${{ (github.event_name != 'schedule' || github.repository == 'pytorch/pytorch') && github.repository_owner == 'pytorch' }}
     with:
       triggering_actor: ${{ github.triggering_actor }}
@@ -78,15 +79,15 @@ jobs:
       opt_out_experiments: lf
 
   build:
-    name: cuda12.8-py3.10-gcc11-sm80
+    name: cuda13.0-py3.10-gcc11-sm80
     uses: ./.github/workflows/_linux-build.yml
     needs: get-label-type
     with:
       runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
       # Every bit to make perf run faster helps
       runner: linux.12xlarge.memory
-      build-environment: linux-jammy-cuda12.8-py3.10-gcc11-sm80
-      docker-image-name: ci-image:pytorch-linux-jammy-cuda12.8-cudnn9-py3-gcc11-inductor-benchmarks
+      build-environment: linux-jammy-cuda13.0-py3.10-gcc11-sm80
+      docker-image-name: ci-image:pytorch-linux-jammy-cuda13.0-cudnn9-py3-gcc11-inductor-benchmarks
       cuda-arch-list: '8.0'
       test-matrix: |
         { include: [
@@ -115,7 +116,7 @@ jobs:
     secrets: inherit
 
   test-nightly:
-    name: cuda12.8-py3.10-gcc11-sm80
+    name: cuda13.0-py3.10-gcc11-sm80
     uses: ./.github/workflows/_linux-test.yml
     needs: build
     if: github.event.schedule == '0 7 * * 1-6'
@@ -131,7 +132,7 @@ jobs:
     secrets: inherit
 
   test-weekly:
-    name: cuda12.8-py3.10-gcc11-sm80
+    name: cuda13.0-py3.10-gcc11-sm80
     uses: ./.github/workflows/_linux-test.yml
     needs: build
     if: github.event.schedule == '0 7 * * 0'
@@ -148,7 +149,7 @@ jobs:
     secrets: inherit
 
   test:
-    name: cuda12.8-py3.10-gcc11-sm80
+    name: cuda13.0-py3.10-gcc11-sm80
     uses: ./.github/workflows/_linux-test.yml
     needs: build
     if: github.event_name == 'workflow_dispatch'
@@ -162,89 +163,3 @@ jobs:
       monitor-log-interval: 15
       monitor-data-collect-interval: 4
     secrets: inherit
-
-  build-cuda13:
-    name: cuda13.0-py3.10-gcc11-sm80
-    uses: ./.github/workflows/_linux-build.yml
-    needs: get-label-type
-    with:
-      runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
-      # Every bit to make perf run faster helps
-      runner: linux.12xlarge.memory
-      build-environment: linux-jammy-cuda13.0-py3.10-gcc11-sm80
-      docker-image-name: ci-image:pytorch-linux-jammy-cuda13.0-cudnn9-py3-gcc11-inductor-benchmarks
-      cuda-arch-list: '8.0'
-      test-matrix: |
-        { include: [
-          { config: "inductor_huggingface_perf", shard: 1, num_shards: 5, runner: "linux.aws.a100" },
-          { config: "inductor_huggingface_perf", shard: 2, num_shards: 5, runner: "linux.aws.a100" },
-          { config: "inductor_huggingface_perf", shard: 3, num_shards: 5, runner: "linux.aws.a100" },
-          { config: "inductor_huggingface_perf", shard: 4, num_shards: 5, runner: "linux.aws.a100" },
-          { config: "inductor_huggingface_perf", shard: 5, num_shards: 5, runner: "linux.aws.a100" },
-          { config: "inductor_timm_perf", shard: 1, num_shards: 6, runner: "linux.aws.a100" },
-          { config: "inductor_timm_perf", shard: 2, num_shards: 6, runner: "linux.aws.a100" },
-          { config: "inductor_timm_perf", shard: 3, num_shards: 6, runner: "linux.aws.a100" },
-          { config: "inductor_timm_perf", shard: 4, num_shards: 6, runner: "linux.aws.a100" },
-          { config: "inductor_timm_perf", shard: 5, num_shards: 6, runner: "linux.aws.a100" },
-          { config: "inductor_timm_perf", shard: 6, num_shards: 6, runner: "linux.aws.a100" },
-          { config: "inductor_torchbench_perf", shard: 1, num_shards: 6, runner: "linux.aws.a100" },
-          { config: "inductor_torchbench_perf", shard: 2, num_shards: 6, runner: "linux.aws.a100" },
-          { config: "inductor_torchbench_perf", shard: 3, num_shards: 6, runner: "linux.aws.a100" },
-          { config: "inductor_torchbench_perf", shard: 4, num_shards: 6, runner: "linux.aws.a100" },
-          { config: "inductor_torchbench_perf", shard: 5, num_shards: 6, runner: "linux.aws.a100" },
-          { config: "inductor_torchbench_perf", shard: 6, num_shards: 6, runner: "linux.aws.a100" },
-          { config: "cachebench", shard: 1, num_shards: 2, runner: "linux.aws.a100" },
-          { config: "cachebench", shard: 2, num_shards: 2, runner: "linux.aws.a100" },
-        ]}
-      selected-test-configs: ${{ inputs.benchmark_configs }}
-      build-additional-packages: "vision audio torchao"
-    secrets: inherit
-
-  test-nightly-cuda13:
-    name: cuda13.0-py3.10-gcc11-sm80
-    uses: ./.github/workflows/_linux-test.yml
-    needs: build-cuda13
-    if: github.event.schedule == '0 7 * * 1-6'
-    with:
-      build-environment: linux-jammy-cuda13.0-py3.10-gcc11-sm80
-      dashboard-tag: training-true-inference-true-default-true-dynamic-true-cudagraphs-true-cppwrapper-true-aotinductor-true-freezing_cudagraphs-true
-      docker-image: ${{ needs.build-cuda13.outputs.docker-image }}
-      test-matrix: ${{ needs.build-cuda13.outputs.test-matrix }}
-      timeout-minutes: 720
-      disable-monitor: false
-      monitor-log-interval: 15
-      monitor-data-collect-interval: 4
-    secrets: inherit
-
-  test-weekly-cuda13:
-    name: cuda13.0-py3.10-gcc11-sm80
-    uses: ./.github/workflows/_linux-test.yml
-    needs: build-cuda13
-    if: github.event.schedule == '0 7 * * 0'
-    with:
-      build-environment: linux-jammy-cuda13.0-py3.10-gcc11-sm80
-      dashboard-tag: training-true-inference-true-default-true-dynamic-true-cudagraphs-true-cppwrapper-true-aotinductor-true-freezing_cudagraphs-true-maxautotune-true-freeze_autotune_cudagraphs-true
-      docker-image: ${{ needs.build-cuda13.outputs.docker-image }}
-      test-matrix: ${{ needs.build-cuda13.outputs.test-matrix }}
-      timeout-minutes: 1440
-      # disable monitor in perf tests, next step is to enable it
-      disable-monitor: false
-      monitor-log-interval: 15
-      monitor-data-collect-interval: 4
-    secrets: inherit
-
-  test-cuda13:
-    name: cuda13.0-py3.10-gcc11-sm80
-    uses: ./.github/workflows/_linux-test.yml
-    needs: build-cuda13
-    if: github.event_name == 'workflow_dispatch'
-    with:
-      build-environment: linux-jammy-cuda13.0-py3.10-gcc11-sm80
-      dashboard-tag: training-${{ inputs.training }}-inference-${{ inputs.inference }}-default-${{ inputs.default }}-dynamic-${{ inputs.dynamic }}-cudagraphs-${{ inputs.cudagraphs }}-cppwrapper-${{ inputs.cppwrapper }}-aotinductor-${{ inputs.aotinductor }}-maxautotune-${{ inputs.maxautotune }}-freezing_cudagraphs-${{ inputs.freezing_cudagraphs }}
-      docker-image: ${{ needs.build-cuda13.outputs.docker-image }}
-      test-matrix: ${{ needs.build-cuda13.outputs.test-matrix }}
-      timeout-minutes: 720
-      disable-monitor: false
-      monitor-log-interval: 15
-      monitor-data-collect-interval: 4
-    secrets: inherit
diff --git a/.github/workflows/inductor-periodic.yml b/.github/workflows/inductor-periodic.yml
index 18d7c7189f38e..64e79ee3ab941 100644
--- a/.github/workflows/inductor-periodic.yml
+++ b/.github/workflows/inductor-periodic.yml
@@ -18,11 +18,12 @@ concurrency:
 permissions:
   id-token: write
   contents: read
+  actions: read
 
 jobs:
   get-default-label-prefix:
     name: get-default-label-prefix
-    uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@main
+    uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@release/2.12
     if: ${{ (github.event_name != 'schedule' || github.repository == 'pytorch/pytorch') && github.repository_owner == 'pytorch' }}
     with:
       triggering_actor: ${{ github.triggering_actor }}
@@ -38,8 +39,8 @@ jobs:
     with:
       runner_prefix: "${{ needs.get-default-label-prefix.outputs.label-type }}"
       runner: linux.4xlarge.memory
-      build-environment: linux-jammy-cuda12.8-py3.10-gcc11-sm86
-      docker-image-name: ci-image:pytorch-linux-jammy-cuda12.8-cudnn9-py3-gcc11-inductor-benchmarks
+      build-environment: linux-jammy-cuda13.0-py3.10-gcc11-sm86
+      docker-image-name: ci-image:pytorch-linux-jammy-cuda13.0-cudnn9-py3-gcc11-inductor-benchmarks
       cuda-arch-list: '8.0;8.6'
       test-matrix: |
         { include: [
@@ -59,6 +60,9 @@ jobs:
           { config: "dynamic_aot_eager_timm", shard: 1, num_shards: 2, runner: "linux.g5.4xlarge.nvidia.gpu" },
           { config: "dynamic_aot_eager_timm", shard: 2, num_shards: 2, runner: "linux.g5.4xlarge.nvidia.gpu" },
           { config: "dynamic_inductor_huggingface", shard: 1, num_shards: 1, runner: "linux.g5.4xlarge.nvidia.gpu" },
+          { config: "inductor_huggingface_unbacked_parity", shard: 1, num_shards: 1, runner: "linux.g5.4xlarge.nvidia.gpu" },
+          { config: "inductor_huggingface_unbacked_parity", shard: 1, num_shards: 1, runner: "linux.aws.a100" },
+          { config: "inductor_huggingface_unbacked_parity", shard: 1, num_shards: 1, runner: "linux.aws.h100" },
           { config: "dynamic_inductor_timm", shard: 1, num_shards: 2, runner: "linux.g5.4xlarge.nvidia.gpu" },
           { config: "dynamic_inductor_timm", shard: 2, num_shards: 2, runner: "linux.g5.4xlarge.nvidia.gpu" },
           { config: "dynamic_inductor_torchbench", shard: 1, num_shards: 2, runner: "linux.g5.4xlarge.nvidia.gpu" },
@@ -82,63 +86,12 @@ jobs:
       test-matrix: ${{ needs.periodic-dynamo-benchmarks-build.outputs.test-matrix }}
     secrets: inherit
 
-  periodic-dynamo-benchmarks-build-cuda13:
-    name: periodic-dynamo-benchmarks-build-cuda13
-    uses: ./.github/workflows/_linux-build.yml
-    needs: get-default-label-prefix
-    with:
-      runner_prefix: "${{ needs.get-default-label-prefix.outputs.label-type }}"
-      runner: linux.4xlarge.memory
-      build-environment: linux-jammy-cuda13.0-py3.10-gcc11-sm86
-      docker-image-name: ci-image:pytorch-linux-jammy-cuda13.0-cudnn9-py3-gcc11-inductor-benchmarks
-      cuda-arch-list: '8.0;8.6'
-      test-matrix: |
-        { include: [
-          { config: "dynamo_eager_torchbench", shard: 1, num_shards: 2, runner: "linux.g5.4xlarge.nvidia.gpu" },
-          { config: "dynamo_eager_torchbench", shard: 2, num_shards: 2, runner: "linux.g5.4xlarge.nvidia.gpu" },
-          { config: "dynamo_eager_huggingface", shard: 1, num_shards: 1, runner: "linux.g5.4xlarge.nvidia.gpu" },
-          { config: "dynamo_eager_timm", shard: 1, num_shards: 2, runner: "linux.g5.4xlarge.nvidia.gpu" },
-          { config: "dynamo_eager_timm", shard: 2, num_shards: 2, runner: "linux.g5.4xlarge.nvidia.gpu" },
-          { config: "aot_eager_torchbench", shard: 1, num_shards: 2, runner: "linux.g5.4xlarge.nvidia.gpu" },
-          { config: "aot_eager_torchbench", shard: 2, num_shards: 2, runner: "linux.g5.4xlarge.nvidia.gpu" },
-          { config: "aot_eager_huggingface", shard: 1, num_shards: 1, runner: "linux.g5.4xlarge.nvidia.gpu" },
-          { config: "aot_eager_timm", shard: 1, num_shards: 2, runner: "linux.g5.4xlarge.nvidia.gpu" },
-          { config: "aot_eager_timm", shard: 2, num_shards: 2, runner: "linux.g5.4xlarge.nvidia.gpu" },
-          { config: "dynamic_aot_eager_torchbench", shard: 1, num_shards: 2, runner: "linux.g5.4xlarge.nvidia.gpu" },
-          { config: "dynamic_aot_eager_torchbench", shard: 2, num_shards: 2, runner: "linux.g5.4xlarge.nvidia.gpu" },
-          { config: "dynamic_aot_eager_huggingface", shard: 1, num_shards: 1, runner: "linux.g5.4xlarge.nvidia.gpu" },
-          { config: "dynamic_aot_eager_timm", shard: 1, num_shards: 2, runner: "linux.g5.4xlarge.nvidia.gpu" },
-          { config: "dynamic_aot_eager_timm", shard: 2, num_shards: 2, runner: "linux.g5.4xlarge.nvidia.gpu" },
-          { config: "dynamic_inductor_huggingface", shard: 1, num_shards: 1, runner: "linux.g5.4xlarge.nvidia.gpu" },
-          { config: "dynamic_inductor_timm", shard: 1, num_shards: 2, runner: "linux.g5.4xlarge.nvidia.gpu" },
-          { config: "dynamic_inductor_timm", shard: 2, num_shards: 2, runner: "linux.g5.4xlarge.nvidia.gpu" },
-          { config: "dynamic_inductor_torchbench", shard: 1, num_shards: 2, runner: "linux.g5.4xlarge.nvidia.gpu" },
-          { config: "dynamic_inductor_torchbench", shard: 2, num_shards: 2, runner: "linux.g5.4xlarge.nvidia.gpu" },
-          { config: "aot_inductor_huggingface", shard: 1, num_shards: 1, runner: "linux.aws.a100" },
-          { config: "aot_inductor_timm", shard: 1, num_shards: 2, runner: "linux.g5.4xlarge.nvidia.gpu" },
-          { config: "aot_inductor_timm", shard: 2, num_shards: 2, runner: "linux.g5.4xlarge.nvidia.gpu" },
-          { config: "aot_inductor_torchbench", shard: 1, num_shards: 2, runner: "linux.g5.4xlarge.nvidia.gpu" },
-          { config: "aot_inductor_torchbench", shard: 2, num_shards: 2, runner: "linux.g5.4xlarge.nvidia.gpu" },
-        ]}
-      build-additional-packages: "vision audio torchao"
-    secrets: inherit
-
-  periodic-dynamo-benchmarks-test-cuda13:
-    name: periodic-dynamo-benchmarks-test-cuda13
-    uses: ./.github/workflows/_linux-test.yml
-    needs: periodic-dynamo-benchmarks-build-cuda13
-    with:
-      build-environment: linux-jammy-cuda13.0-py3.10-gcc11-sm86
-      docker-image: ${{ needs.periodic-dynamo-benchmarks-build-cuda13.outputs.docker-image }}
-      test-matrix: ${{ needs.periodic-dynamo-benchmarks-build-cuda13.outputs.test-matrix }}
-    secrets: inherit
-
   rocm-periodic-dynamo-benchmarks-build:
     if: github.repository_owner == 'pytorch'
     name: rocm-periodic-dynamo-benchmarks-build
     uses: ./.github/workflows/_linux-build.yml
     with:
-      build-environment: linux-jammy-rocm-py3_10
+      build-environment: linux-jammy-rocm-py3.10-mi355
       docker-image-name: ci-image:pytorch-linux-jammy-rocm-n-py3-benchmarks
       test-matrix: |
         { include: [
@@ -171,9 +124,6 @@ jobs:
     secrets: inherit
 
   rocm-periodic-dynamo-benchmarks-test:
-    permissions:
-      id-token: write
-      contents: read
     name: rocm-periodic-dynamo-benchmarks-test
     uses: ./.github/workflows/_rocm-test.yml
     needs: rocm-periodic-dynamo-benchmarks-build
@@ -191,8 +141,8 @@ jobs:
     with:
       runner_prefix: "${{ needs.get-default-label-prefix.outputs.label-type }}"
       runner: linux.4xlarge.memory
-      build-environment: linux-jammy-cuda12.8-py3.10-gcc11-sm80
-      docker-image-name: ci-image:pytorch-linux-jammy-cuda12.8-cudnn9-py3-gcc11-inductor-benchmarks
+      build-environment: linux-jammy-cuda13.0-py3.10-gcc11-sm80
+      docker-image-name: ci-image:pytorch-linux-jammy-cuda13.0-cudnn9-py3-gcc11-inductor-benchmarks
       cuda-arch-list: '8.0'
       test-matrix: |
         { include: [
diff --git a/.github/workflows/inductor-rocm-mi200.yml b/.github/workflows/inductor-rocm-mi200.yml
index 1698eb4fc85fb..b874352a55701 100644
--- a/.github/workflows/inductor-rocm-mi200.yml
+++ b/.github/workflows/inductor-rocm-mi200.yml
@@ -1,11 +1,11 @@
 name: inductor-rocm-mi200
 
 on:
-  schedule:
-    - cron: 0 0,3,6,9,12,15,18,21 * * * # run every 3 hours for regression coverage
+  #schedule:
+    #- cron: 0 0,3,6,9,12,15,18,21 * * * # run every 3 hours for regression coverage
   push:
-    branches:
-      - release/*
+    #branches:
+      #- release/*
     tags:
       - ciflow/inductor-rocm-mi200/*
   workflow_dispatch:
@@ -17,11 +17,12 @@ concurrency:
 permissions:
   id-token: write
   contents: read
+  actions: read
 
 jobs:
   get-label-type:
     name: get-label-type
-    uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@main
+    uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@release/2.12
     if: ${{ (github.event_name != 'schedule' || github.repository == 'pytorch/pytorch') && github.repository_owner == 'pytorch' }}
     with:
       triggering_actor: ${{ github.triggering_actor }}
@@ -31,12 +32,12 @@ jobs:
       opt_out_experiments: lf
 
   linux-jammy-rocm-py3_10-inductor-build:
-    name: rocm-py3.10-inductor
+    name: linux-jammy-rocm-py3.10-mi200
     uses: ./.github/workflows/_linux-build.yml
     needs: get-label-type
     with:
       runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
-      build-environment: linux-jammy-rocm-py3.10
+      build-environment: linux-jammy-rocm-py3.10-mi200
       docker-image-name: ci-image:pytorch-linux-jammy-rocm-n-py3
       test-matrix: |
         { include: [
@@ -46,10 +47,7 @@ jobs:
     secrets: inherit
 
   linux-jammy-rocm-py3_10-inductor-test:
-    permissions:
-      id-token: write
-      contents: read
-    name: rocm-py3.10-inductor
+    name: linux-jammy-rocm-py3.10-mi200
     uses: ./.github/workflows/_rocm-test.yml
     needs: linux-jammy-rocm-py3_10-inductor-build
     with:
diff --git a/.github/workflows/inductor-rocm-mi300.yml b/.github/workflows/inductor-rocm-mi300.yml
index 633386aba487b..44801016478da 100644
--- a/.github/workflows/inductor-rocm-mi300.yml
+++ b/.github/workflows/inductor-rocm-mi300.yml
@@ -17,19 +17,17 @@ concurrency:
 permissions:
   id-token: write
   contents: read
+  actions: read
 
 jobs:
   target-determination:
     if: github.repository_owner == 'pytorch'
     name: before-test
     uses: ./.github/workflows/target_determination.yml
-    permissions:
-      id-token: write
-      contents: read
 
   get-label-type:
     name: get-label-type
-    uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@main
+    uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@release/2.12
     if: ${{ (github.event_name != 'schedule' || github.repository == 'pytorch/pytorch') && github.repository_owner == 'pytorch' }}
     with:
       triggering_actor: ${{ github.triggering_actor }}
@@ -39,7 +37,7 @@ jobs:
       opt_out_experiments: lf
 
   linux-noble-rocm-py3_12-inductor-build:
-    name: rocm-py3.12-inductor-mi300
+    name: linux-noble-rocm-py3.12-mi300
     uses: ./.github/workflows/_linux-build.yml
     needs: get-label-type
     with:
@@ -54,10 +52,7 @@ jobs:
     secrets: inherit
 
   linux-noble-rocm-py3_12-inductor-test:
-    permissions:
-      id-token: write
-      contents: read
-    name: rocm-py3.12-inductor-mi300
+    name: linux-noble-rocm-py3.12-mi300
     uses: ./.github/workflows/_rocm-test.yml
     needs: linux-noble-rocm-py3_12-inductor-build
     with:
diff --git a/.github/workflows/inductor-rocm-mi355.yml b/.github/workflows/inductor-rocm-mi355.yml
index 70ea41a6da698..0380206235c43 100644
--- a/.github/workflows/inductor-rocm-mi355.yml
+++ b/.github/workflows/inductor-rocm-mi355.yml
@@ -16,19 +16,17 @@ concurrency:
 permissions:
   id-token: write
   contents: read
+  actions: read
 
 jobs:
   target-determination:
     if: github.repository_owner == 'pytorch'
     name: before-test
     uses: ./.github/workflows/target_determination.yml
-    permissions:
-      id-token: write
-      contents: read
 
   get-label-type:
     name: get-label-type
-    uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@main
+    uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@release/2.12
     if: ${{ (github.event_name != 'schedule' || github.repository == 'pytorch/pytorch') && github.repository_owner == 'pytorch' }}
     with:
       triggering_actor: ${{ github.triggering_actor }}
@@ -38,7 +36,7 @@ jobs:
       opt_out_experiments: lf
 
   linux-noble-rocm-py3_12-inductor-build:
-    name: rocm-py3.12-inductor-mi355
+    name: linux-noble-rocm-py3.12-mi355
     uses: ./.github/workflows/_linux-build.yml
     needs: get-label-type
     with:
@@ -53,10 +51,7 @@ jobs:
     secrets: inherit
 
   linux-noble-rocm-py3_12-inductor-test:
-    permissions:
-      id-token: write
-      contents: read
-    name: rocm-py3.12-inductor-mi355
+    name: linux-noble-rocm-py3.12-mi355
     uses: ./.github/workflows/_rocm-test.yml
     needs: linux-noble-rocm-py3_12-inductor-build
     with:
diff --git a/.github/workflows/inductor-unittest.yml b/.github/workflows/inductor-unittest.yml
index b5cea8ceb5265..a2eb23c7b68ce 100644
--- a/.github/workflows/inductor-unittest.yml
+++ b/.github/workflows/inductor-unittest.yml
@@ -18,11 +18,12 @@ concurrency:
 permissions:
   id-token: write
   contents: read
+  actions: read
 
 jobs:
   get-label-type:
     name: get-label-type
-    uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@main
+    uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@release/2.12
     if: ${{ (github.event_name != 'schedule' || github.repository == 'pytorch/pytorch') && github.repository_owner == 'pytorch' }}
     with:
       triggering_actor: ${{ github.triggering_actor }}
@@ -36,8 +37,8 @@ jobs:
     uses: ./.github/workflows/_linux-build.yml
     needs: get-label-type
     with:
-      build-environment: linux-jammy-cuda12.8-py3.10-gcc11-sm86
-      docker-image-name: ci-image:pytorch-linux-jammy-cuda12.8-cudnn9-py3-gcc11-inductor-benchmarks
+      build-environment: linux-jammy-cuda13.0-py3.10-gcc11-sm86
+      docker-image-name: ci-image:pytorch-linux-jammy-cuda13.0-cudnn9-py3-gcc11-inductor-benchmarks
       cuda-arch-list: '8.6'
       runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
       test-matrix: |
@@ -166,8 +167,8 @@ jobs:
         python-version: ['3.11', '3.12', '3.13']
     with:
       runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
-      build-environment: linux-jammy-py${{ matrix.python-version }}-clang15
-      docker-image-name: ci-image:pytorch-linux-jammy-py${{ matrix.python-version }}-clang15
+      build-environment: linux-jammy-py${{ matrix.python-version }}-clang18
+      docker-image-name: ci-image:pytorch-linux-jammy-py${{ matrix.python-version }}-clang18
       test-matrix: |
         { include: [
           { config: "inductor_core", shard: 1, num_shards: 2, runner: "${{ needs.get-label-type.outputs.label-type }}linux.c7i.2xlarge" },
@@ -183,8 +184,8 @@ jobs:
       matrix:
         python-version: ['3.11', '3.12', '3.13']
     with:
-      build-environment: linux-jammy-py${{ matrix.python-version }}-clang15
-      docker-image: ci-image:pytorch-linux-jammy-py${{ matrix.python-version }}-clang15
+      build-environment: linux-jammy-py${{ matrix.python-version }}-clang18
+      docker-image: ci-image:pytorch-linux-jammy-py${{ matrix.python-version }}-clang18
       test-matrix: |
         { include: [
           { config: "inductor_core", shard: 1, num_shards: 2, runner: "${{ needs.get-label-type.outputs.label-type }}linux.c7i.2xlarge" },
diff --git a/.github/workflows/inductor.yml b/.github/workflows/inductor.yml
index acfa832c88ebb..66441b8fbba02 100644
--- a/.github/workflows/inductor.yml
+++ b/.github/workflows/inductor.yml
@@ -25,6 +25,7 @@ concurrency:
 permissions:
   id-token: write
   contents: read
+  actions: read
 
 jobs:
   unit-test:
@@ -35,7 +36,7 @@ jobs:
 
   get-label-type:
     name: get-label-type
-    uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@main
+    uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@release/2.12
     if: ${{ (github.event_name != 'schedule' || github.repository == 'pytorch/pytorch') && github.repository_owner == 'pytorch' }}
     with:
       triggering_actor: ${{ github.triggering_actor }}
@@ -49,8 +50,8 @@ jobs:
     uses: ./.github/workflows/_linux-build.yml
     needs: get-label-type
     with:
-      build-environment: linux-jammy-cuda12.8-py3.10-gcc11-sm86
-      docker-image-name: ci-image:pytorch-linux-jammy-cuda12.8-cudnn9-py3-gcc11-inductor-benchmarks
+      build-environment: linux-jammy-cuda13.0-py3.10-gcc11-sm86
+      docker-image-name: ci-image:pytorch-linux-jammy-cuda13.0-cudnn9-py3-gcc11-inductor-benchmarks
       cuda-arch-list: '8.6'
       runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
       runner: linux.4xlarge.memory
@@ -73,37 +74,7 @@ jobs:
       build-environment: ${{ needs.inductor-build.outputs.build-environment }}
       docker-image: ${{ needs.inductor-build.outputs.docker-image }}
       test-matrix: ${{ needs.inductor-build.outputs.test-matrix }}
-    secrets: inherit
-
-  inductor-build-cuda13:
-    name: inductor-build-cuda13
-    uses: ./.github/workflows/_linux-build.yml
-    needs: get-label-type
-    with:
-      build-environment: linux-jammy-cuda13.0-py3.10-gcc11-sm86
-      docker-image-name: ci-image:pytorch-linux-jammy-cuda13.0-cudnn9-py3-gcc11-inductor-benchmarks
-      cuda-arch-list: '8.6'
-      runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
-      runner: linux.4xlarge.memory
-      test-matrix: |
-        { include: [
-          { config: "inductor_huggingface", shard: 1, num_shards: 1, runner: "${{ needs.get-label-type.outputs.label-type }}linux.g5.4xlarge.nvidia.gpu" },
-          { config: "inductor_timm", shard: 1, num_shards: 2, runner: "${{ needs.get-label-type.outputs.label-type }}linux.g5.4xlarge.nvidia.gpu" },
-          { config: "inductor_timm", shard: 2, num_shards: 2, runner: "${{ needs.get-label-type.outputs.label-type }}linux.g5.4xlarge.nvidia.gpu" },
-          { config: "inductor_torchbench", shard: 1, num_shards: 2, runner: "${{ needs.get-label-type.outputs.label-type }}linux.g5.4xlarge.nvidia.gpu" },
-          { config: "inductor_torchbench", shard: 2, num_shards: 2, runner: "${{ needs.get-label-type.outputs.label-type }}linux.g5.4xlarge.nvidia.gpu" },
-        ]}
-      build-additional-packages: "vision audio torchao"
-    secrets: inherit
-
-  inductor-test-cuda13:
-    name: inductor-test-cuda13
-    uses: ./.github/workflows/_linux-test.yml
-    needs: inductor-build-cuda13
-    with:
-      build-environment: linux-jammy-cuda13.0-py3.10-gcc11-sm86
-      docker-image: ${{ needs.inductor-build-cuda13.outputs.docker-image }}
-      test-matrix: ${{ needs.inductor-build-cuda13.outputs.test-matrix }}
+      enable-torch-trace: "1"
     secrets: inherit
 
   inductor-cpu-build:
diff --git a/.github/workflows/lint-autoformat.yml b/.github/workflows/lint-autoformat.yml
deleted file mode 100644
index b962970dc5b78..0000000000000
--- a/.github/workflows/lint-autoformat.yml
+++ /dev/null
@@ -1,42 +0,0 @@
-name: Apply lint suggestions
-
-on:
-  pull_request:
-    types: [opened, synchronize, reopened, labeled, unlabeled]
-
-jobs:
-  lintrunner-autoformat:
-    permissions:
-      contents: read
-      pull-requests: write
-    runs-on: lf.linux.2xlarge
-    if: ${{ github.repository_owner == 'pytorch' && contains(github.event.pull_request.labels.*.name, 'autoformat') }}
-    steps:
-      - name: Checkout pytorch
-        uses: pytorch/pytorch/.github/actions/checkout-pytorch@main
-        with:
-          submodules: true
-          fetch-depth: 0
-      - name: Run lintrunner (nonretryable)
-        continue-on-error: true
-        run: |
-          set -ex
-          python3 -m venv /tmp/venv
-          source /tmp/venv/bin/activate
-          export ADDITIONAL_LINTRUNNER_ARGS="format --all-files"
-          bash .github/scripts/lintrunner.sh
-      - name: Check for changes
-        id: git-check
-        continue-on-error: true
-        run: |
-          git diff --exit-code || echo "changes=true" >> "$GITHUB_OUTPUT"
-      - name: Suggest changes
-        if: steps.git-check.outputs.changes == 'true'
-        continue-on-error: true
-        uses: parkerbxyz/suggest-changes@a2ec1653b0c4cc8287d682f0066dba4a173cc7f3 # v1.0.8
-        with:
-          comment: "Please commit the suggested changes from pytorch's linter."
-
-concurrency:
-  group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}-${{ github.event_name == 'workflow_dispatch' }}
-  cancel-in-progress: true
diff --git a/.github/workflows/lint-bc.yml b/.github/workflows/lint-bc.yml
index e0de9ede35084..bdba4fb224f26 100644
--- a/.github/workflows/lint-bc.yml
+++ b/.github/workflows/lint-bc.yml
@@ -20,7 +20,7 @@ jobs:
     runs-on: ubuntu-latest
     steps:
       - name: Run BC Lint Action
-        uses: pytorch/test-infra/.github/actions/bc-lint@main
+        uses: pytorch/test-infra/.github/actions/bc-lint@release/2.12
         with:
           repo: ${{ github.event.pull_request.head.repo.full_name }}
           base_sha: ${{ github.event.pull_request.base.sha }}
diff --git a/.github/workflows/lint.yml b/.github/workflows/lint.yml
index e5d998e100a94..f1c1bede5004e 100644
--- a/.github/workflows/lint.yml
+++ b/.github/workflows/lint.yml
@@ -22,7 +22,7 @@ jobs:
   get-label-type:
     if: github.repository_owner == 'pytorch'
     name: get-label-type
-    uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@main
+    uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@release/2.12
     with:
       triggering_actor: ${{ github.triggering_actor }}
       issue_owner: ${{ github.event.pull_request.user.login || github.event.issue.user.login }}
@@ -36,7 +36,7 @@ jobs:
       all_files: ${{ contains(github.event.pull_request.labels.*.name, 'lint-all-files') || contains(github.event.pull_request.labels.*.name, 'Reverted') || github.event_name == 'push' }}
 
   lintrunner-clang:
-    uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main
+    uses: ./.github/workflows/_lint.yml
     # Needed to prevent deduping on HUD
     name: lintrunner-clang-${{ needs.get-changed-files.outputs.changed-files == '*' && 'all' || 'partial' }}
     needs: [get-label-type, get-changed-files]
@@ -56,14 +56,8 @@ jobs:
         contains(needs.get-changed-files.outputs.changed-files, '.metal')
       )
     with:
-      timeout: 120
-      runner: "${{ needs.get-label-type.outputs.label-type }}linux.2xlarge"
-      docker-image: ci-image:pytorch-linux-jammy-cuda12.8-cudnn9-py3.10-linter
-      # NB: A shallow checkout won't work here because calculate-docker-image requires a full checkout
-      # to run git rev-parse HEAD~:.ci/docker when a new image is needed
-      fetch-depth: 0
-      submodules: true
-      ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
+      runner: mt-l-x86iamx-8-16
+      docker-image: ghcr.io/pytorch/test-infra:cuda-x86_64-810d48d
       script: |
         CHANGED_FILES="${{ needs.get-changed-files.outputs.changed-files }}"
         if [ "$CHANGED_FILES" = "*" ]; then
@@ -78,7 +72,7 @@ jobs:
   #       fails to find types when it should
   # NOTE: We should be able to disable this and consolidate with Pyrefly
   lintrunner-pyrefly:
-    uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main
+    uses: ./.github/workflows/_lint.yml
     name: lintrunner-pyrefly-${{ needs.get-changed-files.outputs.changed-files == '*' && 'all' || 'partial' }}
     needs: [get-label-type, get-changed-files]
     # Only run if there are changed files relevant to pyrefly
@@ -89,32 +83,20 @@ jobs:
         contains(needs.get-changed-files.outputs.changed-files, '.pyi')
       )
     with:
-      timeout: 120
-      runner: "${{ needs.get-label-type.outputs.label-type }}linux.2xlarge"
-      docker-image: ci-image:pytorch-linux-jammy-linter
-      # NB: A shallow checkout won't work here because calculate-docker-image requires a full checkout
-      # to run git rev-parse HEAD~:.ci/docker when a new image is needed
-      fetch-depth: 0
-      submodules: true
-      ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
+      runner: mt-l-x86iamx-8-16
+      docker-image: ghcr.io/pytorch/test-infra:cpu-x86_64-810d48d
       script: |
         CHANGED_FILES="${{ needs.get-changed-files.outputs.changed-files }}"
         echo "Running pyrefly"
         ADDITIONAL_LINTRUNNER_ARGS="--take PYREFLY --all-files" .github/scripts/lintrunner.sh
 
   lintrunner-noclang:
-    uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main
+    uses: ./.github/workflows/_lint.yml
     name: lintrunner-noclang-${{ needs.get-changed-files.outputs.changed-files == '*' && 'all' || 'partial' }}
     needs: [get-label-type, get-changed-files]
     with:
-      timeout: 120
-      runner: "${{ needs.get-label-type.outputs.label-type }}linux.2xlarge"
-      docker-image: ci-image:pytorch-linux-jammy-linter
-      # NB: A shallow checkout won't work here because calculate-docker-image requires a full checkout
-      # to run git rev-parse HEAD~:.ci/docker when a new image is needed
-      fetch-depth: 0
-      submodules: true
-      ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
+      runner: mt-l-x86iamx-8-16
+      docker-image: ghcr.io/pytorch/test-infra:cpu-x86_64-810d48d
       script: |
         CHANGED_FILES="${{ needs.get-changed-files.outputs.changed-files }}"
         echo "Running all other linters"
@@ -125,14 +107,12 @@ jobs:
         fi
 
   quick-checks:
-    uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main
+    if: github.repository_owner == 'pytorch'
     needs: get-label-type
+    uses: ./.github/workflows/_lint.yml
     with:
-      timeout: 120
-      runner: "${{ needs.get-label-type.outputs.label-type }}linux.2xlarge"
-      docker-image: ci-image:pytorch-linux-jammy-linter
-      fetch-depth: 0
-      ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
+      runner: mt-l-x86iamx-8-16
+      docker-image: ghcr.io/pytorch/test-infra:cpu-x86_64-810d48d
       script: |
         # Ensure no non-breaking spaces
         # NB: We use 'printf' below rather than '\u000a' since bash pre-4.2
@@ -159,13 +139,12 @@ jobs:
 
   pr-sanity-checks:
     name: pr-sanity-checks
-    needs: get-label-type
-    runs-on: [self-hosted, "${{ needs.get-label-type.outputs.label-type }}linux.large"]
+    runs-on: linux.24_04.4x
     # Only run this on pull requests. This check is simple enough to be done without a Docker image
-    if: github.event_name == 'pull_request' && !contains(github.event.pull_request.labels.*.name, 'skip-pr-sanity-checks')
+    if: ${{ github.event_name == 'pull_request' && !contains(github.event.pull_request.labels.*.name, 'skip-pr-sanity-checks') && github.repository_owner == 'pytorch' }}
     steps:
       - name: Checkout PyTorch
-        uses: pytorch/pytorch/.github/actions/checkout-pytorch@main
+        uses: pytorch/pytorch/.github/actions/checkout-pytorch@release/2.12
         with:
           submodules: false
           fetch-depth: -1
@@ -178,18 +157,15 @@ jobs:
           bash .github/scripts/pr-sanity-check.sh
 
   workflow-checks:
-    uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main
+    if: github.repository_owner == 'pytorch'
     needs: get-label-type
+    uses: ./.github/workflows/_lint.yml
     with:
-      timeout: 120
-      runner: "${{ needs.get-label-type.outputs.label-type }}linux.2xlarge"
-      docker-image: ci-image:pytorch-linux-jammy-linter
-      fetch-depth: -1
-      submodules: true
-      ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
+      runner: mt-l-x86iamx-8-16
+      docker-image: ghcr.io/pytorch/test-infra:cpu-x86_64-810d48d
       script: |
         # Regenerate workflows
-        .github/scripts/generate_ci_workflows.py
+        RELEASE_VERSION_TAG=2.12 .github/scripts/generate_ci_workflows.py
 
         RC=0
         # Assert that regenerating the workflows didn't change them
@@ -212,14 +188,12 @@ jobs:
         exit $RC
 
   toc:
-    uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main
+    if: github.repository_owner == 'pytorch'
     needs: get-label-type
+    uses: ./.github/workflows/_lint.yml
     with:
-      timeout: 120
-      runner: "${{ needs.get-label-type.outputs.label-type }}linux.2xlarge"
-      docker-image: ci-image:pytorch-linux-jammy-linter
-      fetch-depth: 0
-      ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
+      runner: mt-l-x86iamx-8-16
+      docker-image: ghcr.io/pytorch/test-infra:cpu-x86_64-810d48d
       script: |
         # Regenerate ToCs and check that they didn't change
         set -eu
@@ -248,14 +222,11 @@ jobs:
   test-tools:
     name: Test tools
     if: ${{ github.repository == 'pytorch/pytorch' }}
-    uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main
     needs: get-label-type
+    uses: ./.github/workflows/_lint.yml
     with:
-      timeout: 120
-      runner: "${{ needs.get-label-type.outputs.label-type }}linux.2xlarge"
-      docker-image: ci-image:pytorch-linux-jammy-linter
-      fetch-depth: 0
-      ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
+      runner: mt-l-x86iamx-8-16
+      docker-image: ghcr.io/pytorch/test-infra:cpu-x86_64-810d48d
       script: |
         # Test tools
         PYTHONPATH=$(pwd) pytest tools/stats
@@ -268,7 +239,7 @@ jobs:
     runs-on: linux.24_04.4x
     steps:
       - name: Checkout PyTorch
-        uses: pytorch/pytorch/.github/actions/checkout-pytorch@main
+        uses: pytorch/pytorch/.github/actions/checkout-pytorch@release/2.12
         with:
           submodules: false
           fetch-depth: 1
@@ -305,7 +276,7 @@ jobs:
       # [see note: pytorch repo ref]
       # deep clone (fetch-depth 0) required, to allow us to use git log
       - name: Checkout PyTorch
-        uses: pytorch/pytorch/.github/actions/checkout-pytorch@main
+        uses: pytorch/pytorch/.github/actions/checkout-pytorch@release/2.12
         with:
           submodules: false
           fetch-depth: 1
@@ -350,13 +321,51 @@ jobs:
           python3 torch/utils/collect_env.py
 
   link-check:
-    name: Link checks
+    if: github.repository_owner == 'pytorch'
     needs: get-label-type
+    name: Link checks
     uses: ./.github/workflows/_link_check.yml
     with:
       runner: ${{ needs.get-label-type.outputs.label-type }}
-      ref:    ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
-    secrets: inherit
+      ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
+
+  doc-redirects-check:
+    name: doc-redirects-check
+    runs-on: linux.24_04.4x
+    if: github.event_name == 'pull_request' && github.repository_owner == 'pytorch'
+    steps:
+      - name: Checkout PyTorch
+        uses: pytorch/pytorch/.github/actions/checkout-pytorch@release/2.12
+        with:
+          submodules: false
+          fetch-depth: 0
+
+      - name: Doc redirects check (nonretryable)
+        env:
+          BASE_REF: ${{ github.event.pull_request.base.ref }}
+        run: |
+          set -euo pipefail
+
+          # Run the check with auto-fix to generate suggestions
+          python3 .github/scripts/check_doc_redirects.py \
+            --base-ref "origin/${BASE_REF}" \
+            --auto-fix
+
+          # If redirects.py was modified, show the diff and fail
+          if ! git diff --quiet docs/source/redirects.py 2>/dev/null; then
+            echo ""
+            echo "📋 The following redirects were auto-generated:"
+            echo ""
+            git diff docs/source/redirects.py
+            echo ""
+            echo "Please add these changes to your PR by running locally:"
+            echo ""
+            echo "    python3 .github/scripts/check_doc_redirects.py --base-ref origin/main --auto-fix"
+            echo "    git add docs/source/redirects.py"
+            echo "    git commit --amend --no-edit  # or create a new commit"
+            echo ""
+            exit 1
+          fi
 
 concurrency:
   group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}-${{ github.event_name == 'workflow_dispatch' && github.run_id }}
diff --git a/.github/workflows/linux-aarch64.yml b/.github/workflows/linux-aarch64.yml
deleted file mode 100644
index 2c6e6b6dac39c..0000000000000
--- a/.github/workflows/linux-aarch64.yml
+++ /dev/null
@@ -1,62 +0,0 @@
-name: linux-aarch64
-
-on:
-  push:
-    branches:
-      - main
-      - release/*
-    tags:
-      - ciflow/linux-aarch64/*
-      - ciflow/trunk/*
-  workflow_dispatch:
-
-concurrency:
-  group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref_name }}-${{ github.ref_type == 'branch' && github.sha }}-${{ github.event_name == 'workflow_dispatch' && github.run_id }}-${{ github.event_name == 'schedule' }}
-  cancel-in-progress: true
-
-jobs:
-
-  get-label-type:
-    if: github.repository_owner == 'pytorch'
-    name: get-label-type
-    uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@main
-    with:
-      triggering_actor: ${{ github.triggering_actor }}
-      issue_owner: ${{ github.event.pull_request.user.login || github.event.issue.user.login }}
-      curr_branch: ${{ github.head_ref || github.ref_name }}
-      curr_ref_type: ${{ github.ref_type }}
-
-  linux-jammy-aarch64-py3_10-build:
-    name: linux-jammy-aarch64-py3.10
-    uses: ./.github/workflows/_linux-build.yml
-    needs: get-label-type
-    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
-      runner: linux.arm64.m7g.4xlarge
-      test-matrix: |
-        { include: [
-          { config: "default", shard: 1, num_shards: 3, runner: "${{ needs.get-label-type.outputs.label-type }}linux.arm64.m7g.4xlarge" },
-          { config: "default", shard: 2, num_shards: 3, runner: "${{ needs.get-label-type.outputs.label-type }}linux.arm64.m7g.4xlarge" },
-          { config: "default", shard: 3, num_shards: 3, runner: "${{ needs.get-label-type.outputs.label-type }}linux.arm64.m7g.4xlarge" },
-          { config: "openreg", shard: 1, num_shards: 1, runner: "${{ needs.get-label-type.outputs.label-type }}linux.arm64.m7g.4xlarge" },
-          { config: "default", shard: 1, num_shards: 3, runner: "${{ needs.get-label-type.outputs.label-type }}linux.arm64.m8g.4xlarge" },
-          { config: "default", shard: 2, num_shards: 3, runner: "${{ needs.get-label-type.outputs.label-type }}linux.arm64.m8g.4xlarge" },
-          { config: "default", shard: 3, num_shards: 3, runner: "${{ needs.get-label-type.outputs.label-type }}linux.arm64.m8g.4xlarge" },
-          { config: "openreg", shard: 1, num_shards: 1, runner: "${{ needs.get-label-type.outputs.label-type }}linux.arm64.m8g.4xlarge" },
-        ]}
-    secrets: inherit
-
-  linux-jammy-aarch64-py3_10-test:
-    name: linux-jammy-aarch64-py3.10
-    uses: ./.github/workflows/_linux-test.yml
-    needs: linux-jammy-aarch64-py3_10-build
-    permissions:
-      id-token: write
-      contents: read
-    with:
-      build-environment: ${{ needs.linux-jammy-aarch64-py3_10-build.outputs.build-environment }}
-      docker-image: ${{ needs.linux-jammy-aarch64-py3_10-build.outputs.docker-image }}
-      test-matrix: ${{ needs.linux-jammy-aarch64-py3_10-build.outputs.test-matrix }}
-    secrets: inherit
diff --git a/.github/workflows/llm_td_retrieval.yml b/.github/workflows/llm_td_retrieval.yml
index 565a9b25df50f..534cf7aae2807 100644
--- a/.github/workflows/llm_td_retrieval.yml
+++ b/.github/workflows/llm_td_retrieval.yml
@@ -12,7 +12,7 @@ jobs:
     name: get-label-type
     # Don't run on forked repos
     if: github.repository_owner == 'pytorch'
-    uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@main
+    uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@release/2.12
     with:
       triggering_actor: ${{ github.triggering_actor }}
       issue_owner: ${{ github.event.pull_request.user.login || github.event.issue.user.login }}
@@ -26,15 +26,8 @@ jobs:
     continue-on-error: true
     needs: get-label-type
     steps:
-      - name: Clone PyTorch
-        uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
-        with:
-          repository: pytorch/pytorch
-          fetch-depth: 0
-          path: pytorch
-
       - name: Setup Linux
-        uses: ./pytorch/.github/actions/setup-linux
+        uses: pytorch/pytorch/.github/actions/setup-linux@release/2.12
 
       - name: Clone CodeLlama
         uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
@@ -116,5 +109,5 @@ jobs:
           AWS_REGION: ""
 
       - name: Teardown Linux
-        uses: pytorch/test-infra/.github/actions/teardown-linux@main
+        uses: pytorch/test-infra/.github/actions/teardown-linux@release/2.12
         if: always()
diff --git a/.github/workflows/nightly-s3-uploads.yml b/.github/workflows/nightly-s3-uploads.yml
index acf3504dec9ca..5ceb50923be56 100644
--- a/.github/workflows/nightly-s3-uploads.yml
+++ b/.github/workflows/nightly-s3-uploads.yml
@@ -23,7 +23,7 @@ jobs:
     environment: upload-stats
     steps:
       - name: Checkout PyTorch
-        uses: pytorch/pytorch/.github/actions/checkout-pytorch@main
+        uses: pytorch/pytorch/.github/actions/checkout-pytorch@release/2.12
         with:
           fetch-depth: 1
           submodules: false
diff --git a/.github/workflows/nightly.yml b/.github/workflows/nightly.yml
index c47b0c5763078..cdb4931097fa6 100644
--- a/.github/workflows/nightly.yml
+++ b/.github/workflows/nightly.yml
@@ -18,10 +18,15 @@ concurrency:
   group: ${{ github.workflow }}--${{ github.event.pull_request.number || github.sha }}-${{ github.event_name == 'workflow_dispatch' }}
   cancel-in-progress: true
 
+permissions:
+  id-token: write
+  contents: read
+  actions: read
+
 jobs:
   get-label-type:
     name: get-label-type
-    uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@main
+    uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@release/2.12
     if: ${{ (github.event_name != 'schedule' || github.repository == 'pytorch/pytorch') && github.repository_owner == 'pytorch' }}
     with:
       triggering_actor: ${{ github.triggering_actor }}
@@ -45,7 +50,7 @@ jobs:
     with:
       runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
       build-environment: linux-jammy-py3.10-gcc11
-      docker-image-name: ci-image:pytorch-linux-jammy-py3.10-gcc11
+      docker-image-name: ci-image:pytorch-linux-jammy-py3.10-clang18
     secrets: inherit
 
   docs-push:
@@ -90,11 +95,15 @@ jobs:
             repo-owner: vllm-project
             branch: main
             pin-folder: .github/ci_commit_pins
+          - repo-name: torchtitan
+            repo-owner: pytorch
+            branch: main
+            pin-folder: .github/ci_commit_pins
     # Allow this to be triggered on either a schedule or on workflow_dispatch to allow for easier testing
     if: github.repository_owner == 'pytorch' && (github.event_name == 'schedule' || github.event_name == 'workflow_dispatch')
     steps:
       - name: "${{ matrix.repo-owner }}/${{ matrix.repo-name }} update-commit-hash"
-        uses: pytorch/test-infra/.github/actions/update-commit-hash@main
+        uses: pytorch/test-infra/.github/actions/update-commit-hash@release/2.12
         with:
           repo-owner: ${{ matrix.repo-owner }}
           repo-name: ${{ matrix.repo-name }}
diff --git a/.github/workflows/nitpicker.yml b/.github/workflows/nitpicker.yml
index 40bd245ce913f..cb41646745892 100644
--- a/.github/workflows/nitpicker.yml
+++ b/.github/workflows/nitpicker.yml
@@ -19,8 +19,8 @@ jobs:
     if: ${{ github.event.pull_request.number != 26921 && github.repository_owner == 'pytorch' }}
     steps:
     - name: Checkout PyTorch
-      uses: pytorch/pytorch/.github/actions/checkout-pytorch@main
-    - uses: ethanis/nitpicker@v1
+      uses: pytorch/pytorch/.github/actions/checkout-pytorch@release/2.12
+    - uses: ethanis/nitpicker@c102a39683a80c7db9065f8eab7de8b58871f946 # v1
       with:
         nitpicks: '.github/nitpicks.yml'
         token: "${{ secrets.GITHUB_TOKEN }}"
diff --git a/.github/workflows/operator_benchmark.yml b/.github/workflows/operator_benchmark.yml
index e682e1eb06c24..7d6cbe58b726b 100644
--- a/.github/workflows/operator_benchmark.yml
+++ b/.github/workflows/operator_benchmark.yml
@@ -28,6 +28,7 @@ concurrency:
 permissions:
   id-token: write
   contents: read
+  actions: read
 
 jobs:
   x86-opbenchmark-build:
diff --git a/.github/workflows/operator_microbenchmark.yml b/.github/workflows/operator_microbenchmark.yml
index 445cdcc4be04a..1a7c2d8294c29 100644
--- a/.github/workflows/operator_microbenchmark.yml
+++ b/.github/workflows/operator_microbenchmark.yml
@@ -17,11 +17,12 @@ concurrency:
 permissions:
   id-token: write
   contents: read
+  actions: read
 
 jobs:
   get-label-type:
     name: get-label-type
-    uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@main
+    uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@release/2.12
     if: ${{ (github.event_name != 'schedule' || github.repository == 'pytorch/pytorch') && github.repository_owner == 'pytorch' }}
     with:
       triggering_actor: ${{ github.triggering_actor }}
diff --git a/.github/workflows/periodic-rocm-mi200.yml b/.github/workflows/periodic-rocm-mi200.yml
index 865d999623cbb..e7d4bee98d4fa 100644
--- a/.github/workflows/periodic-rocm-mi200.yml
+++ b/.github/workflows/periodic-rocm-mi200.yml
@@ -1,13 +1,13 @@
 name: periodic-rocm-mi200
 
 on:
-  schedule:
-    - cron: 0 0,3,6,9,12,15,18,21 * * * # run every 3 hours for regression coverage
+  #schedule:
+    #- cron: 0 0,3,6,9,12,15,18,21 * * * # run every 3 hours for regression coverage
   push:
     tags:
       - ciflow/periodic-rocm-mi200/*
-    branches:
-      - release/*
+    #branches:
+      #- release/*
   workflow_dispatch:
 
 concurrency:
@@ -17,27 +17,22 @@ concurrency:
 permissions:
   id-token: write
   contents: read
+  actions: read
 
 jobs:
   llm-td:
     if: github.repository_owner == 'pytorch'
     name: before-test
     uses: ./.github/workflows/llm_td_retrieval.yml
-    permissions:
-      id-token: write
-      contents: read
 
   target-determination:
     name: before-test
     uses: ./.github/workflows/target_determination.yml
     needs: llm-td
-    permissions:
-      id-token: write
-      contents: read
 
   get-label-type:
     name: get-label-type
-    uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@main
+    uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@release/2.12
     if: (github.event_name != 'schedule' || github.repository == 'pytorch/pytorch') && github.repository_owner == 'pytorch'
     with:
       triggering_actor: ${{ github.triggering_actor }}
@@ -46,26 +41,23 @@ jobs:
       curr_ref_type: ${{ github.ref_type }}
 
   linux-jammy-rocm-py3_10-build:
-    name: linux-jammy-rocm-py3.10
+    name: linux-jammy-rocm-py3.10-mi200
     uses: ./.github/workflows/_linux-build.yml
     needs: get-label-type
     with:
       runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
-      build-environment: linux-jammy-rocm-py3.10
+      build-environment: linux-jammy-rocm-py3.10-mi200
       docker-image-name: ci-image:pytorch-linux-jammy-rocm-n-py3
       test-matrix: |
         { include: [
-          { config: "distributed", shard: 1, num_shards: 3, runner: "linux.rocm.gpu.4", owners: ["module:rocm", "oncall:distributed"] },
-          { config: "distributed", shard: 2, num_shards: 3, runner: "linux.rocm.gpu.4", owners: ["module:rocm", "oncall:distributed"] },
-          { config: "distributed", shard: 3, num_shards: 3, runner: "linux.rocm.gpu.4", owners: ["module:rocm", "oncall:distributed"] },
+          { config: "distributed", shard: 1, num_shards: 3, runner: "linux.rocm.gpu.2", owners: ["module:rocm", "oncall:distributed"] },
+          { config: "distributed", shard: 2, num_shards: 3, runner: "linux.rocm.gpu.2", owners: ["module:rocm", "oncall:distributed"] },
+          { config: "distributed", shard: 3, num_shards: 3, runner: "linux.rocm.gpu.2", owners: ["module:rocm", "oncall:distributed"] },
         ]}
     secrets: inherit
 
   linux-jammy-rocm-py3_10-test:
-    permissions:
-      id-token: write
-      contents: read
-    name: linux-jammy-rocm-py3.10
+    name: linux-jammy-rocm-py3.10-mi200
     uses: ./.github/workflows/_rocm-test.yml
     needs:
       - linux-jammy-rocm-py3_10-build
diff --git a/.github/workflows/periodic-rocm-mi300.yml b/.github/workflows/periodic-rocm-mi300.yml
index 88da168926444..459528ef6442f 100644
--- a/.github/workflows/periodic-rocm-mi300.yml
+++ b/.github/workflows/periodic-rocm-mi300.yml
@@ -14,28 +14,25 @@ concurrency:
   group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref_name }}-${{ github.ref_type == 'branch' && github.sha }}-${{ github.event_name == 'workflow_dispatch' && github.run_id }}-${{ github.event_name == 'schedule' }}-${{ github.event.schedule }}
   cancel-in-progress: true
 
-permissions: read-all
+permissions:
+  id-token: write
+  contents: read
+  actions: read
 
 jobs:
   llm-td:
     if: github.repository_owner == 'pytorch'
     name: before-test
     uses: ./.github/workflows/llm_td_retrieval.yml
-    permissions:
-      id-token: write
-      contents: read
 
   target-determination:
     name: before-test
     uses: ./.github/workflows/target_determination.yml
     needs: llm-td
-    permissions:
-      id-token: write
-      contents: read
 
   get-label-type:
     name: get-label-type
-    uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@main
+    uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@release/2.12
     if: (github.event_name != 'schedule' || github.repository == 'pytorch/pytorch') && github.repository_owner == 'pytorch'
     with:
       triggering_actor: ${{ github.triggering_actor }}
@@ -60,9 +57,6 @@ jobs:
     secrets: inherit
 
   linux-noble-rocm-py3_12-test:
-    permissions:
-      id-token: write
-      contents: read
     name: linux-noble-rocm-py3.12-mi300
     uses: ./.github/workflows/_rocm-test.yml
     needs:
diff --git a/.github/workflows/periodic-rocm-mi355.yml b/.github/workflows/periodic-rocm-mi355.yml
index 9885ffb2d3832..190041a7a6c3e 100644
--- a/.github/workflows/periodic-rocm-mi355.yml
+++ b/.github/workflows/periodic-rocm-mi355.yml
@@ -2,6 +2,7 @@ name: periodic-rocm-mi355
 
 on:
   schedule:
+    - cron: 0 0,3,6,9,12,15,18,21 * * *
     - cron: 29 8 * * *  # about 1:29am PDT, for mem leak check and rerun disabled tests
   push:
     tags:
@@ -14,28 +15,25 @@ concurrency:
   group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref_name }}-${{ github.ref_type == 'branch' && github.sha }}-${{ github.event_name == 'workflow_dispatch' && github.run_id }}-${{ github.event_name == 'schedule' }}-${{ github.event.schedule }}
   cancel-in-progress: true
 
-permissions: read-all
+permissions:
+  id-token: write
+  contents: read
+  actions: read
 
 jobs:
   llm-td:
     if: github.repository_owner == 'pytorch'
     name: before-test
     uses: ./.github/workflows/llm_td_retrieval.yml
-    permissions:
-      id-token: write
-      contents: read
 
   target-determination:
     name: before-test
     uses: ./.github/workflows/target_determination.yml
     needs: llm-td
-    permissions:
-      id-token: write
-      contents: read
 
   get-label-type:
     name: get-label-type
-    uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@main
+    uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@release/2.12
     if: (github.event_name != 'schedule' || github.repository == 'pytorch/pytorch') && github.repository_owner == 'pytorch'
     with:
       triggering_actor: ${{ github.triggering_actor }}
@@ -60,9 +58,6 @@ jobs:
     secrets: inherit
 
   linux-noble-rocm-py3_12-test:
-    permissions:
-      id-token: write
-      contents: read
     name: linux-noble-rocm-py3.12-mi355
     uses: ./.github/workflows/_rocm-test.yml
     needs:
diff --git a/.github/workflows/periodic.yml b/.github/workflows/periodic.yml
index 16a706a7be6be..b44a3bff003dd 100644
--- a/.github/workflows/periodic.yml
+++ b/.github/workflows/periodic.yml
@@ -23,27 +23,22 @@ concurrency:
 permissions:
   id-token: write
   contents: read
+  actions: read
 
 jobs:
   llm-td:
     if: github.repository_owner == 'pytorch'
     name: before-test
     uses: ./.github/workflows/llm_td_retrieval.yml
-    permissions:
-      id-token: write
-      contents: read
 
   target-determination:
     name: before-test
     uses: ./.github/workflows/target_determination.yml
     needs: llm-td
-    permissions:
-      id-token: write
-      contents: read
 
   get-label-type:
     name: get-label-type
-    uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@main
+    uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@release/2.12
     if: (github.event_name != 'schedule' || github.repository == 'pytorch/pytorch') && github.repository_owner == 'pytorch'
     with:
       triggering_actor: ${{ github.triggering_actor }}
@@ -51,48 +46,25 @@ jobs:
       curr_branch: ${{ github.head_ref || github.ref_name }}
       curr_ref_type: ${{ github.ref_type }}
 
-  linux-jammy-cuda12_4-py3_10-gcc11-build:
-    name: linux-jammy-cuda12.4-py3.10-gcc11
-    uses: ./.github/workflows/_linux-build.yml
-    needs: get-label-type
-    with:
-      runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
-      build-environment: linux-jammy-cuda12.4-py3.10-gcc11
-      docker-image-name: ci-image:pytorch-linux-jammy-cuda12.4-cudnn9-py3-gcc11
-      cuda-arch-list: 7.5
-      test-matrix: |
-        { include: [
-          { config: "legacy_nvidia_driver", shard: 1, num_shards: 5, runner: "${{ needs.get-label-type.outputs.label-type }}linux.g4dn.4xlarge.nvidia.gpu" },
-          { config: "legacy_nvidia_driver", shard: 2, num_shards: 5, runner: "${{ needs.get-label-type.outputs.label-type }}linux.g4dn.4xlarge.nvidia.gpu" },
-          { config: "legacy_nvidia_driver", shard: 3, num_shards: 5, runner: "${{ needs.get-label-type.outputs.label-type }}linux.g4dn.4xlarge.nvidia.gpu" },
-          { config: "legacy_nvidia_driver", shard: 4, num_shards: 5, runner: "${{ needs.get-label-type.outputs.label-type }}linux.g4dn.4xlarge.nvidia.gpu" },
-          { config: "legacy_nvidia_driver", shard: 5, num_shards: 5, runner: "${{ needs.get-label-type.outputs.label-type }}linux.g4dn.4xlarge.nvidia.gpu" },
-        ]}
-    secrets: inherit
-
-  linux-jammy-cuda12_4-py3_10-gcc11-test:
-    name: linux-jammy-cuda12.4-py3.10-gcc11
-    uses: ./.github/workflows/_linux-test.yml
-    needs:
-      - linux-jammy-cuda12_4-py3_10-gcc11-build
-      - target-determination
-    with:
-      build-environment: ${{ needs.linux-jammy-cuda12_4-py3_10-gcc11-build.outputs.build-environment }}
-      docker-image: ${{ needs.linux-jammy-cuda12_4-py3_10-gcc11-build.outputs.docker-image }}
-      test-matrix: ${{ needs.linux-jammy-cuda12_4-py3_10-gcc11-build.outputs.test-matrix }}
-    secrets: inherit
-
-  linux-jammy-cuda12_8-py3_10-gcc11-build:
-    name: linux-jammy-cuda12.8-py3.10-gcc11
+  linux-jammy-cuda13_0-py3_10-gcc11-build:
+    name: linux-jammy-cuda13.0-py3.10-gcc11
     uses: ./.github/workflows/_linux-build.yml
     needs: get-label-type
     with:
       runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
-      build-environment: linux-jammy-cuda12.8-py3.10-gcc11
-      docker-image-name: ci-image:pytorch-linux-jammy-cuda12.8-cudnn9-py3-gcc11
-      cuda-arch-list: 8.6
+      build-environment: linux-jammy-cuda13.0-py3.10-gcc11
+      docker-image-name: ci-image:pytorch-linux-jammy-cuda13.0-cudnn9-py3-gcc11
+      cuda-arch-list: '7.5 8.6 8.9'
       test-matrix: |
         { include: [
+          { config: "default", shard: 1, num_shards: 5, runner: "${{ needs.get-label-type.outputs.label-type }}linux.g6.4xlarge.experimental.nvidia.gpu" },
+          { config: "default", shard: 2, num_shards: 5, runner: "${{ needs.get-label-type.outputs.label-type }}linux.g6.4xlarge.experimental.nvidia.gpu" },
+          { config: "default", shard: 3, num_shards: 5, runner: "${{ needs.get-label-type.outputs.label-type }}linux.g6.4xlarge.experimental.nvidia.gpu" },
+          { config: "default", shard: 4, num_shards: 5, runner: "${{ needs.get-label-type.outputs.label-type }}linux.g6.4xlarge.experimental.nvidia.gpu" },
+          { config: "default", shard: 5, num_shards: 5, runner: "${{ needs.get-label-type.outputs.label-type }}linux.g6.4xlarge.experimental.nvidia.gpu" },
+          { config: "distributed", shard: 1, num_shards: 3, runner: "${{ needs.get-label-type.outputs.label-type }}linux.g4dn.12xlarge.nvidia.gpu" },
+          { config: "distributed", shard: 2, num_shards: 3, runner: "${{ needs.get-label-type.outputs.label-type }}linux.g4dn.12xlarge.nvidia.gpu" },
+          { config: "distributed", shard: 3, num_shards: 3, runner: "${{ needs.get-label-type.outputs.label-type }}linux.g4dn.12xlarge.nvidia.gpu" },
           { config: "nogpu_AVX512", shard: 1, num_shards: 3, runner: "${{ needs.get-label-type.outputs.label-type }}linux.4xlarge" },
           { config: "nogpu_AVX512", shard: 2, num_shards: 3, runner: "${{ needs.get-label-type.outputs.label-type }}linux.4xlarge" },
           { config: "nogpu_AVX512", shard: 3, num_shards: 3, runner: "${{ needs.get-label-type.outputs.label-type }}linux.4xlarge" },
@@ -104,26 +76,26 @@ jobs:
         ]}
     secrets: inherit
 
-  linux-jammy-cuda12_8-py3_10-gcc11-test:
-    name: linux-jammy-cuda12.8-py3.10-gcc11
+  linux-jammy-cuda13_0-py3_10-gcc11-test:
+    name: linux-jammy-cuda13.0-py3.10-gcc11
     uses: ./.github/workflows/_linux-test.yml
     needs:
-      - linux-jammy-cuda12_8-py3_10-gcc11-build
+      - linux-jammy-cuda13_0-py3_10-gcc11-build
       - target-determination
     with:
-      build-environment: ${{ needs.linux-jammy-cuda12_8-py3_10-gcc11-build.outputs.build-environment }}
-      docker-image: ${{ needs.linux-jammy-cuda12_8-py3_10-gcc11-build.outputs.docker-image }}
-      test-matrix: ${{ needs.linux-jammy-cuda12_8-py3_10-gcc11-build.outputs.test-matrix }}
+      build-environment: ${{ needs.linux-jammy-cuda13_0-py3_10-gcc11-build.outputs.build-environment }}
+      docker-image: ${{ needs.linux-jammy-cuda13_0-py3_10-gcc11-build.outputs.docker-image }}
+      test-matrix: ${{ needs.linux-jammy-cuda13_0-py3_10-gcc11-build.outputs.test-matrix }}
     secrets: inherit
 
-  linux-jammy-cuda12_8-py3_10-gcc11-debug-build:
-    name: linux-jammy-cuda12.8-py3.10-gcc11-debug
+  linux-jammy-cuda13_0-py3_10-gcc11-debug-build:
+    name: linux-jammy-cuda13.0-py3.10-gcc11-debug
     uses: ./.github/workflows/_linux-build.yml
     needs: get-label-type
     with:
       runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
-      build-environment: linux-jammy-cuda12.8-py3.10-gcc11-debug
-      docker-image-name: ci-image:pytorch-linux-jammy-cuda12.8-cudnn9-py3-gcc11
+      build-environment: linux-jammy-cuda13.0-py3.10-gcc11-debug
+      docker-image-name: ci-image:pytorch-linux-jammy-cuda13.0-cudnn9-py3-gcc11
       cuda-arch-list: 8.9
       test-matrix: |
         { include: [
@@ -137,58 +109,26 @@ jobs:
         ]}
     secrets: inherit
 
-  linux-jammy-cuda12_8-py3_10-gcc11-debug-test:
-    name: linux-jammy-cuda12.8-py3.10-gcc11-debug
+  linux-jammy-cuda13_0-py3_10-gcc11-debug-test:
+    name: linux-jammy-cuda13.0-py3.10-gcc11-debug
     uses: ./.github/workflows/_linux-test.yml
     needs:
-      - linux-jammy-cuda12_8-py3_10-gcc11-debug-build
+      - linux-jammy-cuda13_0-py3_10-gcc11-debug-build
       - target-determination
     with:
-      build-environment: ${{ needs.linux-jammy-cuda12_8-py3_10-gcc11-debug-build.outputs.build-environment }}
-      docker-image: ${{ needs.linux-jammy-cuda12_8-py3_10-gcc11-debug-build.outputs.docker-image }}
-      test-matrix: ${{ needs.linux-jammy-cuda12_8-py3_10-gcc11-debug-build.outputs.test-matrix }}
+      build-environment: ${{ needs.linux-jammy-cuda13_0-py3_10-gcc11-debug-build.outputs.build-environment }}
+      docker-image: ${{ needs.linux-jammy-cuda13_0-py3_10-gcc11-debug-build.outputs.docker-image }}
+      test-matrix: ${{ needs.linux-jammy-cuda13_0-py3_10-gcc11-debug-build.outputs.test-matrix }}
     secrets: inherit
 
-  linux-jammy-cuda13_0-py3_10-gcc11-build:
-    name: linux-jammy-cuda13.0-py3.10-gcc11
+  linux-jammy-cuda13_0-py3-gcc11-slow-gradcheck-build:
+    name: linux-jammy-cuda13.0-py3-gcc11-slow-gradcheck
     uses: ./.github/workflows/_linux-build.yml
     needs: get-label-type
     with:
       runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
-      cuda-arch-list: 7.5
-      build-environment: linux-jammy-cuda13.0-py3.10-gcc11
+      build-environment: linux-jammy-cuda13.0-py3-gcc11-slow-gradcheck
       docker-image-name: ci-image:pytorch-linux-jammy-cuda13.0-cudnn9-py3-gcc11
-      test-matrix: |
-        { include: [
-          { config: "nogpu_AVX512", shard: 1, num_shards: 3, runner: "${{ needs.get-label-type.outputs.label-type }}linux.4xlarge" },
-          { config: "nogpu_AVX512", shard: 2, num_shards: 3, runner: "${{ needs.get-label-type.outputs.label-type }}linux.4xlarge" },
-          { config: "nogpu_AVX512", shard: 3, num_shards: 3, runner: "${{ needs.get-label-type.outputs.label-type }}linux.4xlarge" },
-          { config: "nogpu_NO_AVX2", shard: 1, num_shards: 2, runner: "${{ needs.get-label-type.outputs.label-type }}linux.4xlarge" },
-          { config: "nogpu_NO_AVX2", shard: 2, num_shards: 2, runner: "${{ needs.get-label-type.outputs.label-type }}linux.4xlarge" },
-          { config: "jit_legacy", shard: 1, num_shards: 1, runner: "${{ needs.get-label-type.outputs.label-type }}linux.g4dn.4xlarge.nvidia.gpu" },
-        ]}
-    secrets: inherit
-
-  linux-jammy-cuda13_0-py3_10-gcc11-test:
-    name: linux-jammy-cuda13.0-py3.10-gcc11
-    uses: ./.github/workflows/_linux-test.yml
-    needs:
-      - linux-jammy-cuda13_0-py3_10-gcc11-build
-      - target-determination
-    with:
-      build-environment: ${{ needs.linux-jammy-cuda13_0-py3_10-gcc11-build.outputs.build-environment }}
-      docker-image: ${{ needs.linux-jammy-cuda13_0-py3_10-gcc11-build.outputs.docker-image }}
-      test-matrix: ${{ needs.linux-jammy-cuda13_0-py3_10-gcc11-build.outputs.test-matrix }}
-    secrets: inherit
-
-  linux-jammy-cuda12_8-py3-gcc11-slow-gradcheck-build:
-    name: linux-jammy-cuda12.8-py3-gcc11-slow-gradcheck
-    uses: ./.github/workflows/_linux-build.yml
-    needs: get-label-type
-    with:
-      runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
-      build-environment: linux-jammy-cuda12.8-py3-gcc11-slow-gradcheck
-      docker-image-name: ci-image:pytorch-linux-jammy-cuda12.8-cudnn9-py3-gcc11
       cuda-arch-list: 8.6
       test-matrix: |
         { include: [
@@ -203,15 +143,15 @@ jobs:
         ]}
     secrets: inherit
 
-  linux-jammy-cuda12_8-py3-gcc11-slow-gradcheck-test:
-    name: linux-jammy-cuda12.8-py3-gcc11-slow-gradcheck
+  linux-jammy-cuda13_0-py3-gcc11-slow-gradcheck-test:
+    name: linux-jammy-cuda13.0-py3-gcc11-slow-gradcheck
     uses: ./.github/workflows/_linux-test.yml
     needs:
-      - linux-jammy-cuda12_8-py3-gcc11-slow-gradcheck-build
+      - linux-jammy-cuda13_0-py3-gcc11-slow-gradcheck-build
       - target-determination
     with:
-      build-environment: ${{ needs.linux-jammy-cuda12_8-py3-gcc11-slow-gradcheck-build.outputs.build-environment }}
-      docker-image: ${{ needs.linux-jammy-cuda12_8-py3-gcc11-slow-gradcheck-build.outputs.docker-image }}
-      test-matrix: ${{ needs.linux-jammy-cuda12_8-py3-gcc11-slow-gradcheck-build.outputs.test-matrix }}
+      build-environment: ${{ needs.linux-jammy-cuda13_0-py3-gcc11-slow-gradcheck-build.outputs.build-environment }}
+      docker-image: ${{ needs.linux-jammy-cuda13_0-py3-gcc11-slow-gradcheck-build.outputs.docker-image }}
+      test-matrix: ${{ needs.linux-jammy-cuda13_0-py3-gcc11-slow-gradcheck-build.outputs.test-matrix }}
       timeout-minutes: 300
     secrets: inherit
diff --git a/.github/workflows/pull.yml b/.github/workflows/pull.yml
index 538e163fb84e1..70b0218d79f4b 100644
--- a/.github/workflows/pull.yml
+++ b/.github/workflows/pull.yml
@@ -33,6 +33,7 @@ concurrency:
 permissions:
   id-token: write
   contents: read
+  actions: read
 
 jobs:
   # See job-filter.yml for rules on adding job filter conditions
@@ -47,26 +48,25 @@ jobs:
     if: github.repository_owner == 'pytorch'
     name: before-test
     uses: ./.github/workflows/llm_td_retrieval.yml
-    permissions:
-      id-token: write
-      contents: read
 
   target-determination:
     name: before-test
     uses: ./.github/workflows/target_determination.yml
     needs: llm-td
-    permissions:
-      id-token: write
-      contents: read
 
   get-label-type:
     name: get-label-type
-    uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@main
+    uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@release/2.12
     if: ${{ (github.event_name != 'schedule' || github.repository == 'pytorch/pytorch') && github.repository_owner == 'pytorch' }}
     with:
       triggering_actor: ${{ github.triggering_actor }}
       issue_owner: ${{ github.event.pull_request.user.login || github.event.issue.user.login }}
       curr_branch: ${{ github.head_ref || github.ref_name }}
+      check_experiments: arc,lf
+
+  # ╠══════════════════════════════════════════════════════════════════════╣
+  # ║ linux-jammy-py3.10-gcc11 (build + test)                              ║
+  # ╠══════════════════════════════════════════════════════════════════════╣
 
   linux-jammy-py3_10-gcc11-build:
     if: ${{ needs.job-filter.outputs.jobs == '' || contains(needs.job-filter.outputs.jobs, ' linux-jammy-py3.10-gcc11 ') || contains(needs.job-filter.outputs.jobs, ' linux-docs ') }}
@@ -78,23 +78,29 @@ jobs:
     with:
       runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
       build-environment: linux-jammy-py3.10-gcc11
-      docker-image-name: ci-image:pytorch-linux-jammy-py3.10-gcc11
+      docker-image-name: ci-image:pytorch-linux-jammy-py3.10-clang18
       test-matrix: |
         { include: [
-          { config: "default", shard: 1, num_shards: 5, runner: "${{ needs.get-label-type.outputs.label-type }}linux.2xlarge" },
-          { config: "default", shard: 2, num_shards: 5, runner: "${{ needs.get-label-type.outputs.label-type }}linux.2xlarge" },
-          { config: "default", shard: 3, num_shards: 5, runner: "${{ needs.get-label-type.outputs.label-type }}linux.2xlarge" },
-          { config: "default", shard: 4, num_shards: 5, runner: "${{ needs.get-label-type.outputs.label-type }}linux.2xlarge" },
-          { config: "default", shard: 5, num_shards: 5, runner: "${{ needs.get-label-type.outputs.label-type }}linux.2xlarge" },
+          { config: "default", shard: 1, num_shards: 5, runner: "${{ needs.get-label-type.outputs.label-type }}linux.4xlarge" },
+          { config: "default", shard: 2, num_shards: 5, runner: "${{ needs.get-label-type.outputs.label-type }}linux.4xlarge" },
+          { config: "default", shard: 3, num_shards: 5, runner: "${{ needs.get-label-type.outputs.label-type }}linux.4xlarge" },
+          { config: "default", shard: 4, num_shards: 5, runner: "${{ needs.get-label-type.outputs.label-type }}linux.4xlarge" },
+          { config: "default", shard: 5, num_shards: 5, runner: "${{ needs.get-label-type.outputs.label-type }}linux.4xlarge" },
           { config: "docs_test", shard: 1, num_shards: 1,  runner: "${{ needs.get-label-type.outputs.label-type }}linux.2xlarge" },
           { config: "jit_legacy", shard: 1, num_shards: 1, runner: "${{ needs.get-label-type.outputs.label-type }}linux.2xlarge" },
           { config: "backwards_compat", shard: 1, num_shards: 1, runner: "${{ needs.get-label-type.outputs.label-type }}linux.c7i.2xlarge" },
-          { config: "distributed", shard: 1, num_shards: 2, runner: "${{ needs.get-label-type.outputs.label-type }}linux.2xlarge" },
-          { config: "distributed", shard: 2, num_shards: 2, runner: "${{ needs.get-label-type.outputs.label-type }}linux.2xlarge" },
+          # Use m7i instead of m7a to match the existing (Intel Xeon) numerics. DTensor crossref tests like linalg.multi_dot have tight
+          # float32 tolerances sensitive to different FMA/reduction order across CPU vendors.
+          { config: "distributed", shard: 1, num_shards: 3, runner: "${{ needs.get-label-type.outputs.label-type }}linux.2xlarge.amx" },
+          { config: "distributed", shard: 2, num_shards: 3, runner: "${{ needs.get-label-type.outputs.label-type }}linux.2xlarge.amx" },
+          { config: "distributed", shard: 3, num_shards: 3, runner: "${{ needs.get-label-type.outputs.label-type }}linux.2xlarge.amx" },
           { config: "numpy_2_x", shard: 1, num_shards: 1, runner: "${{ needs.get-label-type.outputs.label-type }}linux.c7i.2xlarge" },
           { config: "libtorch_agnostic_targetting", shard: 1, num_shards: 1, runner: "${{ needs.get-label-type.outputs.label-type }}linux.2xlarge" },
           { config: "openreg", shard: 1, num_shards: 1, runner: "${{ needs.get-label-type.outputs.label-type }}linux.2xlarge" },
         ]}
+      use-arc: ${{ needs.get-label-type.outputs.use-arc == 'true' }}
+      python-version: "3.10"
+      compiler: gcc11
     secrets: inherit
 
   linux-jammy-py3_10-gcc11-test:
@@ -105,22 +111,74 @@ jobs:
       - linux-jammy-py3_10-gcc11-build
       - target-determination
       - job-filter
+      - get-label-type
     with:
       build-environment: ${{ needs.linux-jammy-py3_10-gcc11-build.outputs.build-environment }}
       docker-image: ${{ needs.linux-jammy-py3_10-gcc11-build.outputs.docker-image }}
       test-matrix: ${{ needs.linux-jammy-py3_10-gcc11-build.outputs.test-matrix }}
       tests-to-include: ${{ github.event.inputs.tests-to-include || '' }}
+      use-arc: ${{ needs.get-label-type.outputs.use-arc == 'true' }}
+      python-version: "3.10"
+      compiler: gcc11
+    secrets: inherit
+
+  # ╠══════════════════════════════════════════════════════════════════════╣
+  # ║ linux-jammy-aarch64-py3.10-gcc11 (build + test)                      ║
+  # ╠══════════════════════════════════════════════════════════════════════╣
+  linux-jammy-aarch64-py3_10-gcc13-build:
+    name: linux-jammy-aarch64-py3.10
+    uses: ./.github/workflows/_linux-build.yml
+    needs:
+      - get-label-type
+      - job-filter
+    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
+      runner: linux.arm64.m8g.4xlarge
+      test-matrix: |
+        { include: [
+          { config: "default", shard: 1, num_shards: 5, runner: "${{ needs.get-label-type.outputs.label-type }}linux.arm64.m8g.4xlarge" },
+          { config: "default", shard: 2, num_shards: 5, runner: "${{ needs.get-label-type.outputs.label-type }}linux.arm64.m8g.4xlarge" },
+          { config: "default", shard: 3, num_shards: 5, runner: "${{ needs.get-label-type.outputs.label-type }}linux.arm64.m8g.4xlarge" },
+          { config: "default", shard: 4, num_shards: 5, runner: "${{ needs.get-label-type.outputs.label-type }}linux.arm64.m8g.4xlarge" },
+          { config: "default", shard: 5, num_shards: 5, runner: "${{ needs.get-label-type.outputs.label-type }}linux.arm64.m8g.4xlarge" },
+          { config: "openreg", shard: 1, num_shards: 1, runner: "${{ needs.get-label-type.outputs.label-type }}linux.arm64.m8g.4xlarge" },
+        ]}
+      use-arc: ${{ needs.get-label-type.outputs.use-arc == 'true' }}
+      python-version: "3.10"
+      compiler: gcc13
     secrets: inherit
 
+  linux-jammy-aarch64-py3_10-gcc13-test:
+    name: linux-jammy-aarch64-py3.10
+    uses: ./.github/workflows/_linux-test.yml
+    needs:
+      - linux-jammy-aarch64-py3_10-gcc13-build
+      - target-determination
+      - job-filter
+      - get-label-type
+    with:
+      build-environment: ${{ needs.linux-jammy-aarch64-py3_10-gcc13-build.outputs.build-environment }}
+      docker-image: ${{ needs.linux-jammy-aarch64-py3_10-gcc13-build.outputs.docker-image }}
+      test-matrix: ${{ needs.linux-jammy-aarch64-py3_10-gcc13-build.outputs.test-matrix }}
+      tests-to-include: ${{ github.event.inputs.tests-to-include || '' }}
+      use-arc: ${{ needs.get-label-type.outputs.use-arc == 'true' }}
+      python-version: "3.10"
+      compiler: gcc13
+    secrets: inherit
 
-  linux-jammy-py3_14t-clang15-build:
-    name: linux-jammy-py3.14t-clang15
+  # ╠══════════════════════════════════════════════════════════════════════╣
+  # ║ linux-jammy-py3.14t-clang18 (build + test)                           ║
+  # ╠══════════════════════════════════════════════════════════════════════╣
+  linux-jammy-py3_14t-clang18-build:
+    name: linux-jammy-py3.14t-clang18
     uses: ./.github/workflows/_linux-build.yml
     needs: get-label-type
     with:
       runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
-      build-environment: linux-jammy-py3.14t-clang15
-      docker-image-name: ci-image:pytorch-linux-jammy-py3.14t-clang15
+      build-environment: linux-jammy-py3.14t-clang18
+      docker-image-name: ci-image:pytorch-linux-jammy-py3.14t-clang18
       test-matrix: |
         { include: [
           { config: "default", shard: 1, num_shards: 5, runner: "${{ needs.get-label-type.outputs.label-type }}linux.4xlarge" },
@@ -130,26 +188,39 @@ jobs:
           { config: "default", shard: 5, num_shards: 5, runner: "${{ needs.get-label-type.outputs.label-type }}linux.4xlarge" },
           { config: "crossref", shard: 1, num_shards: 2, runner: "${{ needs.get-label-type.outputs.label-type }}linux.2xlarge" },
           { config: "crossref", shard: 2, num_shards: 2, runner: "${{ needs.get-label-type.outputs.label-type }}linux.2xlarge" },
-          { config: "dynamo_wrapped", shard: 1, num_shards: 3, runner: "${{ needs.get-label-type.outputs.label-type }}linux.2xlarge" },
-          { config: "dynamo_wrapped", shard: 2, num_shards: 3, runner: "${{ needs.get-label-type.outputs.label-type }}linux.2xlarge" },
-          { config: "dynamo_wrapped", shard: 3, num_shards: 3, runner: "${{ needs.get-label-type.outputs.label-type }}linux.2xlarge" },
+          # Free-threaded Python 3.14t has ~20-50% higher per-object memory overhead from biased reference counting and
+          # per-object locks, and test_nn (531 tests) under dynamo wrapping with compiled autograd consistently OOMs at 64GB.
+          { config: "dynamo_wrapped", shard: 1, num_shards: 3, runner: "${{ needs.get-label-type.outputs.label-type }}linux.4xlarge" },
+          { config: "dynamo_wrapped", shard: 2, num_shards: 3, runner: "${{ needs.get-label-type.outputs.label-type }}linux.4xlarge" },
+          { config: "dynamo_wrapped", shard: 3, num_shards: 3, runner: "${{ needs.get-label-type.outputs.label-type }}linux.4xlarge" },
           { config: "einops", shard: 1, num_shards: 1, runner: "${{ needs.get-label-type.outputs.label-type }}linux.2xlarge" },
           { config: "openreg", shard: 1, num_shards: 1, runner: "${{ needs.get-label-type.outputs.label-type }}linux.2xlarge" },
         ]}
+      use-arc: ${{ needs.get-label-type.outputs.use-arc == 'true' }}
+      python-version: "3.14t"
+      compiler: clang18
     secrets: inherit
 
-  linux-jammy-py3_14t-clang15-test:
-    name: linux-jammy-py3.14t-clang15
+  linux-jammy-py3_14t-clang18-test:
+    name: linux-jammy-py3.14t-clang18
     uses: ./.github/workflows/_linux-test.yml
     needs:
-      - linux-jammy-py3_14t-clang15-build
+      - linux-jammy-py3_14t-clang18-build
       - target-determination
+      - get-label-type
     with:
-      build-environment: linux-jammy-py3.14t-clang15
-      docker-image: ${{ needs.linux-jammy-py3_14t-clang15-build.outputs.docker-image }}
-      test-matrix: ${{ needs.linux-jammy-py3_14t-clang15-build.outputs.test-matrix }}
+      build-environment: linux-jammy-py3.14t-clang18
+      docker-image: ${{ needs.linux-jammy-py3_14t-clang18-build.outputs.docker-image }}
+      test-matrix: ${{ needs.linux-jammy-py3_14t-clang18-build.outputs.test-matrix }}
+      use-arc: ${{ needs.get-label-type.outputs.use-arc == 'true' }}
+      python-version: "3.14t"
+      compiler: clang18
     secrets: inherit
 
+  # ╠══════════════════════════════════════════════════════════════════════╣
+  # ║ linux-docs                                                           ║
+  # ╠══════════════════════════════════════════════════════════════════════╣
+
   linux-docs:
     if: ${{ needs.job-filter.outputs.jobs == '' || contains(needs.job-filter.outputs.jobs, ' linux-docs ') }}
     name: linux-docs
@@ -157,11 +228,21 @@ jobs:
     needs:
       - linux-jammy-py3_10-gcc11-build
       - job-filter
+      - get-label-type
     with:
       build-environment: ${{ needs.linux-jammy-py3_10-gcc11-build.outputs.build-environment }}
       docker-image: ${{ needs.linux-jammy-py3_10-gcc11-build.outputs.docker-image }}
+      run-doxygen: true
+      use-arc: ${{ needs.get-label-type.outputs.use-arc == 'true' }}
+      runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
+      python-version: "3.10"
+      compiler: gcc11
     secrets: inherit
 
+  # ╠══════════════════════════════════════════════════════════════════════╣
+  # ║ linux-jammy-py3.10-gcc11-no-ops (build only)                         ║
+  # ╠══════════════════════════════════════════════════════════════════════╣
+
   linux-jammy-py3_10-gcc11-no-ops:
     if: ${{ needs.job-filter.outputs.jobs == '' || contains(needs.job-filter.outputs.jobs, ' linux-jammy-py3.10-gcc11-no-ops ') }}
     name: linux-jammy-py3.10-gcc11-no-ops
@@ -172,29 +253,19 @@ jobs:
     with:
       runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
       build-environment: linux-jammy-py3.10-gcc11-no-ops
-      docker-image-name: ci-image:pytorch-linux-jammy-py3.10-gcc11
+      docker-image-name: ci-image:pytorch-linux-jammy-py3.10-clang18
       test-matrix: |
         { include: [
           { config: "default", shard: 1, num_shards: 1 },
         ]}
+      use-arc: ${{ needs.get-label-type.outputs.use-arc == 'true' }}
+      python-version: "3.10"
+      compiler: gcc11
     secrets: inherit
 
-  linux-jammy-py3_10-gcc11-pch:
-    if: ${{ needs.job-filter.outputs.jobs == '' || contains(needs.job-filter.outputs.jobs, ' linux-jammy-py3.10-gcc11-pch ') }}
-    name: linux-jammy-py3.10-gcc11-pch
-    uses: ./.github/workflows/_linux-build.yml
-    needs:
-      - get-label-type
-      - job-filter
-    with:
-      runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
-      build-environment: linux-jammy-py3.10-gcc11-pch
-      docker-image-name: ci-image:pytorch-linux-jammy-py3.10-gcc11
-      test-matrix: |
-        { include: [
-          { config: "default", shard: 1, num_shards: 1 },
-        ]}
-    secrets: inherit
+  # ╠══════════════════════════════════════════════════════════════════════╣
+  # ║ linux-jammy-py3.10-clang18-asan (build + test)                       ║
+  # ╠══════════════════════════════════════════════════════════════════════╣
 
   linux-jammy-py3_10-clang18-asan-build:
     if: ${{ needs.job-filter.outputs.jobs == '' || contains(needs.job-filter.outputs.jobs, ' linux-jammy-py3.10-clang18-asan ') }}
@@ -207,7 +278,7 @@ jobs:
       runner: linux.c7i.4xlarge
       runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
       build-environment: linux-jammy-py3.10-clang18-asan
-      docker-image-name: ci-image:pytorch-linux-jammy-py3-clang18-asan
+      docker-image-name: ci-image:pytorch-linux-jammy-py3.10-clang18
       test-matrix: |
         { include: [
           { config: "default", shard: 1, num_shards: 7, runner: "${{ needs.get-label-type.outputs.label-type }}linux.4xlarge" },
@@ -219,7 +290,11 @@ jobs:
           { config: "default", shard: 7, num_shards: 7, runner: "${{ needs.get-label-type.outputs.label-type }}linux.4xlarge" },
           { config: "openreg", shard: 1, num_shards: 1, runner: "${{ needs.get-label-type.outputs.label-type }}linux.4xlarge" },
         ]}
-      sync-tag: asan-build
+      use-arc: ${{ needs.get-label-type.outputs.use-arc == 'true' }}
+      python-version: "3.10"
+      compiler: clang18
+      # TODO (huydhn): Add this back once other workflow migrates
+      # sync-tag: asan-build
     secrets: inherit
 
   linux-jammy-py3_10-clang18-asan-test:
@@ -230,58 +305,34 @@ jobs:
       - linux-jammy-py3_10-clang18-asan-build
       - target-determination
       - job-filter
+      - get-label-type
     with:
       build-environment: ${{ needs.linux-jammy-py3_10-clang18-asan-build.outputs.build-environment }}
       docker-image: ${{ needs.linux-jammy-py3_10-clang18-asan-build.outputs.docker-image }}
       test-matrix: ${{ needs.linux-jammy-py3_10-clang18-asan-build.outputs.test-matrix }}
-      sync-tag: asan-test
       tests-to-include: ${{ github.event.inputs.tests-to-include || '' }}
+      use-arc: ${{ needs.get-label-type.outputs.use-arc == 'true' }}
+      python-version: "3.10"
+      compiler: clang18
+      # TODO (huydhn): Add this back once other workflow migrates
+      # sync-tag: asan-test
     secrets: inherit
 
-  linux-jammy-py3_10-clang15-onnx-build:
-    if: ${{ needs.job-filter.outputs.jobs == '' || contains(needs.job-filter.outputs.jobs, ' linux-jammy-py3.10-clang15-onnx ') }}
-    name: linux-jammy-py3.10-clang15-onnx
-    uses: ./.github/workflows/_linux-build.yml
-    needs:
-      - get-label-type
-      - job-filter
-    with:
-      runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
-      build-environment: linux-jammy-py3.10-clang15-onnx
-      docker-image-name: ci-image:pytorch-linux-jammy-py3-clang15-onnx
-      test-matrix: |
-        { include: [
-          { config: "default", shard: 1, num_shards: 2, runner: "${{ needs.get-label-type.outputs.label-type }}linux.c7i.2xlarge" },
-          { config: "default", shard: 2, num_shards: 2, runner: "${{ needs.get-label-type.outputs.label-type }}linux.c7i.2xlarge" },
-        ]}
-    secrets: inherit
+  # ╠══════════════════════════════════════════════════════════════════════╣
+  # ║ linux-jammy-py3.10-clang18 (build + test)                            ║
+  # ╠══════════════════════════════════════════════════════════════════════╣
 
-  linux-jammy-py3_10-clang15-onnx-test:
-    if: ${{ needs.job-filter.outputs.jobs == '' || contains(needs.job-filter.outputs.jobs, ' linux-jammy-py3.10-clang15-onnx ') }}
-    name: linux-jammy-py3.10-clang15-onnx
-    uses: ./.github/workflows/_linux-test.yml
-    needs:
-      - linux-jammy-py3_10-clang15-onnx-build
-      - target-determination
-      - job-filter
-    with:
-      build-environment: ${{ needs.linux-jammy-py3_10-clang15-onnx-build.outputs.build-environment }}
-      docker-image: ${{ needs.linux-jammy-py3_10-clang15-onnx-build.outputs.docker-image }}
-      test-matrix: ${{ needs.linux-jammy-py3_10-clang15-onnx-build.outputs.test-matrix }}
-      tests-to-include: ${{ github.event.inputs.tests-to-include || '' }}
-    secrets: inherit
-
-  linux-jammy-py3_10-clang15-build:
-    if: ${{ needs.job-filter.outputs.jobs == '' || contains(needs.job-filter.outputs.jobs, ' linux-jammy-py3.10-clang15 ') }}
-    name: linux-jammy-py3.10-clang15
+  linux-jammy-py3_10-clang18-build:
+    if: ${{ needs.job-filter.outputs.jobs == '' || contains(needs.job-filter.outputs.jobs, ' linux-jammy-py3.10-clang18 ') }}
+    name: linux-jammy-py3.10-clang18
     uses: ./.github/workflows/_linux-build.yml
     needs:
       - get-label-type
       - job-filter
     with:
       runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
-      build-environment: linux-jammy-py3.10-clang15
-      docker-image-name: ci-image:pytorch-linux-jammy-py3.10-clang15
+      build-environment: linux-jammy-py3.10-clang18
+      docker-image-name: ci-image:pytorch-linux-jammy-py3.10-clang18
       test-matrix: |
         { include: [
           { config: "default", shard: 1, num_shards: 5, runner: "${{ needs.get-label-type.outputs.label-type }}linux.4xlarge" },
@@ -296,35 +347,47 @@ jobs:
           { config: "dynamo_wrapped", shard: 3, num_shards: 3, runner: "${{ needs.get-label-type.outputs.label-type }}linux.2xlarge" },
           { config: "einops", shard: 1, num_shards: 1, runner: "${{ needs.get-label-type.outputs.label-type }}linux.2xlarge" },
           { config: "openreg", shard: 1, num_shards: 1, runner: "${{ needs.get-label-type.outputs.label-type }}linux.2xlarge" },
+          { config: "onnx", shard: 1, num_shards: 1, runner: "${{ needs.get-label-type.outputs.label-type }}linux.4xlarge" },
         ]}
+      use-arc: ${{ needs.get-label-type.outputs.use-arc == 'true' }}
+      python-version: "3.10"
+      compiler: clang18
     secrets: inherit
 
-  linux-jammy-py3_10-clang15-test:
-    if: ${{ needs.job-filter.outputs.jobs == '' || contains(needs.job-filter.outputs.jobs, ' linux-jammy-py3.10-clang15 ') }}
-    name: linux-jammy-py3.10-clang15
+  linux-jammy-py3_10-clang18-test:
+    if: ${{ needs.job-filter.outputs.jobs == '' || contains(needs.job-filter.outputs.jobs, ' linux-jammy-py3.10-clang18 ') }}
+    name: linux-jammy-py3.10-clang18
     uses: ./.github/workflows/_linux-test.yml
     needs:
-      - linux-jammy-py3_10-clang15-build
+      - linux-jammy-py3_10-clang18-build
       - target-determination
       - job-filter
+      - get-label-type
     with:
-      build-environment: ${{ needs.linux-jammy-py3_10-clang15-build.outputs.build-environment }}
-      docker-image: ${{ needs.linux-jammy-py3_10-clang15-build.outputs.docker-image }}
-      test-matrix: ${{ needs.linux-jammy-py3_10-clang15-build.outputs.test-matrix }}
+      build-environment: ${{ needs.linux-jammy-py3_10-clang18-build.outputs.build-environment }}
+      docker-image: ${{ needs.linux-jammy-py3_10-clang18-build.outputs.docker-image }}
+      test-matrix: ${{ needs.linux-jammy-py3_10-clang18-build.outputs.test-matrix }}
       tests-to-include: ${{ github.event.inputs.tests-to-include || '' }}
+      use-arc: ${{ needs.get-label-type.outputs.use-arc == 'true' }}
+      python-version: "3.10"
+      compiler: clang18
     secrets: inherit
 
-  linux-jammy-py3_14-clang15-build:
-    if: ${{ needs.job-filter.outputs.jobs == '' || contains(needs.job-filter.outputs.jobs, ' linux-jammy-py3.14-clang15 ') }}
-    name: linux-jammy-py3.14-clang15
+  # ╠══════════════════════════════════════════════════════════════════════╣
+  # ║ linux-jammy-py3.14-clang18 (build + test)                            ║
+  # ╠══════════════════════════════════════════════════════════════════════╣
+
+  linux-jammy-py3_14-clang18-build:
+    if: ${{ needs.job-filter.outputs.jobs == '' || contains(needs.job-filter.outputs.jobs, ' linux-jammy-py3.14-clang18 ') }}
+    name: linux-jammy-py3.14-clang18
     uses: ./.github/workflows/_linux-build.yml
     needs:
       - get-label-type
       - job-filter
     with:
       runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
-      build-environment: linux-jammy-py3.14-clang15
-      docker-image-name: ci-image:pytorch-linux-jammy-py3.14-clang15
+      build-environment: linux-jammy-py3.14-clang18
+      docker-image-name: ci-image:pytorch-linux-jammy-py3.14-clang18
       test-matrix: |
         { include: [
           { config: "default", shard: 1, num_shards: 5, runner: "${{ needs.get-label-type.outputs.label-type }}linux.4xlarge" },
@@ -340,56 +403,57 @@ jobs:
           { config: "einops", shard: 1, num_shards: 1, runner: "${{ needs.get-label-type.outputs.label-type }}linux.2xlarge" },
           { config: "openreg", shard: 1, num_shards: 1, runner: "${{ needs.get-label-type.outputs.label-type }}linux.2xlarge" },
         ]}
+      use-arc: ${{ needs.get-label-type.outputs.use-arc == 'true' }}
+      python-version: "3.14"
+      compiler: clang18
     secrets: inherit
 
-  linux-jammy-py3_14-clang15-test:
-    if: ${{ needs.job-filter.outputs.jobs == '' || contains(needs.job-filter.outputs.jobs, ' linux-jammy-py3.14-clang15 ') }}
-    name: linux-jammy-py3.14-clang15
+  linux-jammy-py3_14-clang18-test:
+    if: ${{ needs.job-filter.outputs.jobs == '' || contains(needs.job-filter.outputs.jobs, ' linux-jammy-py3.14-clang18 ') }}
+    name: linux-jammy-py3.14-clang18
     uses: ./.github/workflows/_linux-test.yml
     needs:
-      - linux-jammy-py3_14-clang15-build
+      - linux-jammy-py3_14-clang18-build
       - job-filter
+      - get-label-type
     with:
-      build-environment: ${{ needs.linux-jammy-py3_14-clang15-build.outputs.build-environment }}
-      docker-image: ${{ needs.linux-jammy-py3_14-clang15-build.outputs.docker-image }}
-      test-matrix: ${{ needs.linux-jammy-py3_14-clang15-build.outputs.test-matrix }}
+      build-environment: ${{ needs.linux-jammy-py3_14-clang18-build.outputs.build-environment }}
+      docker-image: ${{ needs.linux-jammy-py3_14-clang18-build.outputs.docker-image }}
+      test-matrix: ${{ needs.linux-jammy-py3_14-clang18-build.outputs.test-matrix }}
       tests-to-include: ${{ github.event.inputs.tests-to-include || '' }}
+      use-arc: ${{ needs.get-label-type.outputs.use-arc == 'true' }}
+      python-version: "3.14"
+      compiler: clang18
     secrets: inherit
 
-  linux-jammy-cuda12_8-cudnn9-py3_10-clang15-build:
-    if: ${{ needs.job-filter.outputs.jobs == '' || contains(needs.job-filter.outputs.jobs, ' linux-jammy-cuda12.8-cudnn9-py3.10-clang15 ') }}
-    name: linux-jammy-cuda12.8-cudnn9-py3.10-clang15
+  # ╠══════════════════════════════════════════════════════════════════════╣
+  # ║ linux-jammy-cuda12.8-cudnn9-py3.10-clang18 (build only)              ║
+  # ╠══════════════════════════════════════════════════════════════════════╣
+
+  linux-jammy-cuda12_8-cudnn9-py3_10-clang18-build:
+    if: ${{ needs.job-filter.outputs.jobs == '' || contains(needs.job-filter.outputs.jobs, ' linux-jammy-cuda12.8-cudnn9-py3.10-clang18 ') }}
+    name: linux-jammy-cuda12.8-cudnn9-py3.10-clang18
     uses: ./.github/workflows/_linux-build.yml
     needs:
       - get-label-type
       - job-filter
     with:
       runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
-      build-environment: linux-jammy-cuda12.8-cudnn9-py3.10-clang15
-      docker-image-name: ci-image:pytorch-linux-jammy-cuda12.8-cudnn9-py3.10-clang15
+      build-environment: linux-jammy-cuda12.8-cudnn9-py3.10-clang18
+      docker-image-name: ci-image:pytorch-linux-jammy-cuda12.8-cudnn9-py3.10-clang18
       test-matrix: |
         { include: [
           { config: "default", shard: 1, num_shards: 1 },
         ]}
+      use-arc: ${{ needs.get-label-type.outputs.use-arc == 'true' }}
+      python-version: "3.10"
+      compiler: clang18
+      cuda-version: "12.8"
     secrets: inherit
 
-  linux-jammy-cpu-py3_10-gcc11-bazel-test:
-    if: ${{ needs.job-filter.outputs.jobs == '' || contains(needs.job-filter.outputs.jobs, ' linux-jammy-cpu-py3.10-gcc11-bazel-test ') }}
-    name: linux-jammy-cpu-py3.10-gcc11-bazel-test
-    uses: ./.github/workflows/_bazel-build-test.yml
-    needs:
-      - get-label-type
-      - job-filter
-    with:
-      runner: "${{ needs.get-label-type.outputs.label-type }}linux.large"
-      build-environment: linux-jammy-cuda12.8-py3.10-gcc11-bazel-test
-      docker-image-name: ci-image:pytorch-linux-jammy-cuda12.8-cudnn9-py3-gcc11
-      cuda-version: cpu
-      test-matrix: |
-        { include: [
-          { config: "default", shard: 1, num_shards: 1, runner: "${{ needs.get-label-type.outputs.label-type }}linux.4xlarge" },
-        ]}
-    secrets: inherit
+  # ╠══════════════════════════════════════════════════════════════════════╣
+  # ║ linux-jammy-py3.10-gcc11-mobile-lightweight-dispatch (build only)    ║
+  # ╠══════════════════════════════════════════════════════════════════════╣
 
   linux-jammy-py3_10-gcc11-mobile-lightweight-dispatch-build:
     if: ${{ needs.job-filter.outputs.jobs == '' || contains(needs.job-filter.outputs.jobs, ' linux-jammy-py3.10-gcc11-mobile-lightweight-dispatch-build ') }}
@@ -401,16 +465,23 @@ jobs:
     with:
       runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
       build-environment: linux-jammy-py3.10-gcc11-mobile-lightweight-dispatch-build
-      docker-image-name: ci-image:pytorch-linux-jammy-py3.10-gcc11
+      docker-image-name: ci-image:pytorch-linux-jammy-py3.10-clang18
       build-generates-artifacts: false
       test-matrix: |
         { include: [
           { config: "default", shard: 1, num_shards: 1 },
         ]}
+      use-arc: ${{ needs.get-label-type.outputs.use-arc == 'true' }}
+      python-version: "3.10"
+      compiler: gcc11
     secrets: inherit
 
+  # ╠══════════════════════════════════════════════════════════════════════╣
+  # ║ linux-jammy-rocm-py3.10 (build)                                      ║
+  # ╠══════════════════════════════════════════════════════════════════════╣
+
   linux-jammy-rocm-py3_10-build:
-    if: github.event_name == 'pull_request' || (needs.job-filter.outputs.jobs != '' && contains(needs.job-filter.outputs.jobs, ' linux-jammy-rocm-py3.10 '))
+    if: ${{ github.event_name == 'pull_request' || (needs.job-filter.outputs.jobs != '' && contains(needs.job-filter.outputs.jobs, ' linux-jammy-rocm-py3.10 ')) }}
     # don't run build twice on main
     name: linux-jammy-rocm-py3.10
     uses: ./.github/workflows/_linux-build.yml
@@ -421,78 +492,21 @@ jobs:
       runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
       build-environment: linux-jammy-rocm-py3.10
       docker-image-name: ci-image:pytorch-linux-jammy-rocm-n-py3
-      sync-tag: rocm-build
+      # TODO (huydhn): Add this back once other workflow migrates
+      # sync-tag: rocm-build
       test-matrix: |
         { include: [
           { config: "default", shard: 1, num_shards: 3, runner: "linux.rocm.gpu.2" },
           { config: "default", shard: 2, num_shards: 3, runner: "linux.rocm.gpu.2" },
           { config: "default", shard: 3, num_shards: 3, runner: "linux.rocm.gpu.2" },
         ]}
+      use-arc: ${{ needs.get-label-type.outputs.use-arc == 'true' }}
+      python-version: "3.10"
     secrets: inherit
 
-  linux-jammy-cuda12_8-py3_10-gcc11-inductor-build:
-    if: ${{ needs.job-filter.outputs.jobs == '' || contains(needs.job-filter.outputs.jobs, ' cuda12.8-py3.10-gcc11-sm75 ') }}
-    name: cuda12.8-py3.10-gcc11-sm75
-    uses: ./.github/workflows/_linux-build.yml
-    needs:
-      - get-label-type
-      - job-filter
-    with:
-      runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
-      build-environment: linux-jammy-cuda12.8-py3.10-gcc11-sm75
-      docker-image-name: ci-image:pytorch-linux-jammy-cuda12.8-cudnn9-py3-gcc11-inductor-benchmarks
-      cuda-arch-list: '7.5'
-      test-matrix: |
-        { include: [
-          { config: "pr_time_benchmarks", shard: 1, num_shards: 1, runner: "linux.g4dn.metal.nvidia.gpu" },
-        ]}
-    secrets: inherit
-
-  linux-jammy-cuda12_8-py3_10-gcc11-inductor-test:
-    if: ${{ needs.job-filter.outputs.jobs == '' || contains(needs.job-filter.outputs.jobs, ' cuda12.8-py3.10-gcc11-sm75 ') }}
-    name: cuda12.8-py3.10-gcc11-sm75
-    uses: ./.github/workflows/_linux-test.yml
-    needs:
-      - linux-jammy-cuda12_8-py3_10-gcc11-inductor-build
-      - job-filter
-    with:
-      build-environment: ${{ needs.linux-jammy-cuda12_8-py3_10-gcc11-inductor-build.outputs.build-environment }}
-      docker-image: ${{ needs.linux-jammy-cuda12_8-py3_10-gcc11-inductor-build.outputs.docker-image }}
-      test-matrix: ${{ needs.linux-jammy-cuda12_8-py3_10-gcc11-inductor-build.outputs.test-matrix }}
-      tests-to-include: ${{ github.event.inputs.tests-to-include || '' }}
-    secrets: inherit
-
-  linux-jammy-cuda13_0-py3_10-gcc11-inductor-build:
-    if: ${{ needs.job-filter.outputs.jobs == '' || contains(needs.job-filter.outputs.jobs, ' cuda13.0-py3.10-gcc11-sm75 ') }}
-    name: cuda13.0-py3.10-gcc11-sm75
-    uses: ./.github/workflows/_linux-build.yml
-    needs:
-      - get-label-type
-      - job-filter
-    with:
-      runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
-      build-environment: linux-jammy-cuda13.0-py3.10-gcc11-sm75
-      docker-image-name: ci-image:pytorch-linux-jammy-cuda13.0-cudnn9-py3-gcc11-inductor-benchmarks
-      cuda-arch-list: '7.5'
-      test-matrix: |
-        { include: [
-          { config: "pr_time_benchmarks", shard: 1, num_shards: 1, runner: "linux.g4dn.metal.nvidia.gpu" },
-        ]}
-    secrets: inherit
-
-  linux-jammy-cuda13_0-py3_10-gcc11-inductor-test:
-    if: ${{ needs.job-filter.outputs.jobs == '' || contains(needs.job-filter.outputs.jobs, ' cuda13.0-py3.10-gcc11-sm75 ') }}
-    name: cuda13.0-py3.10-gcc11-sm75
-    uses: ./.github/workflows/_linux-test.yml
-    needs:
-      - linux-jammy-cuda13_0-py3_10-gcc11-inductor-build
-      - job-filter
-    with:
-      build-environment: linux-jammy-cuda13.0-py3.10-gcc11-sm75
-      docker-image: ${{ needs.linux-jammy-cuda13_0-py3_10-gcc11-inductor-build.outputs.docker-image }}
-      test-matrix: ${{ needs.linux-jammy-cuda13_0-py3_10-gcc11-inductor-build.outputs.test-matrix }}
-      tests-to-include: ${{ github.event.inputs.tests-to-include || '' }}
-    secrets: inherit
+  # ╠══════════════════════════════════════════════════════════════════════╣
+  # ║ linux-jammy-xpu-n-py3.10 (build)                                     ║
+  # ╠══════════════════════════════════════════════════════════════════════╣
 
   linux-jammy-xpu-n-py3_10-build:
     if: ${{ needs.job-filter.outputs.jobs == '' || contains(needs.job-filter.outputs.jobs, ' linux-jammy-xpu-n-py3.10 ') }}
@@ -502,8 +516,6 @@ jobs:
       - get-label-type
       - job-filter
     with:
-      # This should sync with the build in xpu.yml but xpu uses a larger runner
-      # sync-tag: linux-xpu-n-build
       runner_prefix: ${{ needs.get-label-type.outputs.label-type }}
       runner: linux.c7i.4xlarge
       build-environment: linux-noble-xpu-n-py3.10
@@ -515,20 +527,29 @@ jobs:
           { config: "default", shard: 3, num_shards: 4, runner: "linux.idc.xpu" },
           { config: "default", shard: 4, num_shards: 4, runner: "linux.idc.xpu" },
         ]}
+      use-arc: ${{ needs.get-label-type.outputs.use-arc == 'true' }}
+      python-version: "3.10"
     secrets: inherit
 
+  # ╠══════════════════════════════════════════════════════════════════════╣
+  # ║ dynamo-cpython (build + test)                                        ║
+  # ╠══════════════════════════════════════════════════════════════════════╣
+
   dynamo-cpython-build:
     name: dynamo-cpython-build
     uses: ./.github/workflows/_linux-build.yml
     needs: get-label-type
     with:
       runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
-      build-environment: linux-jammy-py3.13-clang15
-      docker-image-name: ci-image:pytorch-linux-jammy-py3.13-clang15
+      build-environment: linux-jammy-py3.13-clang18
+      docker-image-name: ci-image:pytorch-linux-jammy-py3.13-clang18
       test-matrix: |
         { include: [
           { config: "dynamo_cpython", shard: 1, num_shards: 1, runner: "${{ needs.get-label-type.outputs.label-type }}linux.c7i.2xlarge" },
         ]}
+      use-arc: ${{ needs.get-label-type.outputs.use-arc == 'true' }}
+      python-version: "3.13"
+      compiler: clang18
     secrets: inherit
 
   dynamo-cpython-test:
@@ -536,10 +557,10 @@ jobs:
     uses: ./.github/workflows/_linux-test.yml
     needs: [get-label-type, dynamo-cpython-build]
     with:
-      build-environment: linux-jammy-py3.13-clang15
-      docker-image: ci-image:pytorch-linux-jammy-py3.13-clang15
-      test-matrix: |
-        { include: [
-          { config: "dynamo_cpython", shard: 1, num_shards: 1, runner: "${{ needs.get-label-type.outputs.label-type }}linux.c7i.2xlarge" },
-        ]}
+      build-environment: linux-jammy-py3.13-clang18
+      docker-image: ${{ needs.dynamo-cpython-build.outputs.docker-image }}
+      test-matrix: ${{ needs.dynamo-cpython-build.outputs.test-matrix }}
+      use-arc: ${{ needs.get-label-type.outputs.use-arc == 'true' }}
+      python-version: "3.13"
+      compiler: clang18
     secrets: inherit
diff --git a/.github/workflows/quantization-periodic.yml b/.github/workflows/quantization-periodic.yml
index 8dd97ff9308db..3302ebc351c28 100644
--- a/.github/workflows/quantization-periodic.yml
+++ b/.github/workflows/quantization-periodic.yml
@@ -16,11 +16,12 @@ concurrency:
 permissions:
   id-token: write
   contents: read
+  actions: read
 
 jobs:
   get-default-label-prefix:
     name: get-default-label-prefix
-    uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@main
+    uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@release/2.12
     if: ${{ (github.event_name != 'schedule' || github.repository == 'pytorch/pytorch') && github.repository_owner == 'pytorch' }}
     with:
       triggering_actor: ${{ github.triggering_actor }}
diff --git a/.github/workflows/rocm-mi200.yml b/.github/workflows/rocm-mi200.yml
index 1bbb538527f13..426089e8774ef 100644
--- a/.github/workflows/rocm-mi200.yml
+++ b/.github/workflows/rocm-mi200.yml
@@ -2,13 +2,13 @@ name: rocm-mi200
 
 on:
   push:
-    branches:
-      - release/*
+    #branches:
+      #- release/*
     tags:
       - ciflow/rocm-mi200/*
   workflow_dispatch:
-  schedule:
-    - cron: 0 0,3,6,9,12,15,18,21 * * * # run every 3 hours for regression coverage
+  #schedule:
+    #- cron: 0 0,3,6,9,12,15,18,21 * * * # run every 3 hours for regression coverage
 
 
 concurrency:
@@ -28,7 +28,7 @@ jobs:
 
   get-label-type:
     name: get-label-type
-    uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@main
+    uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@release/2.12
     if: ${{ (github.event_name != 'schedule' || github.repository == 'pytorch/pytorch') && github.repository_owner == 'pytorch' }}
     with:
       triggering_actor: ${{ github.triggering_actor }}
@@ -38,12 +38,12 @@ jobs:
 
   linux-jammy-rocm-py3_10-build:
     if: ${{ (github.event_name != 'schedule' || github.repository == 'pytorch/pytorch') && github.repository_owner == 'pytorch' }}
-    name: linux-jammy-rocm-py3.10
+    name: linux-jammy-rocm-py3.10-mi200
     uses: ./.github/workflows/_linux-build.yml
     needs: get-label-type
     with:
       runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
-      build-environment: linux-jammy-rocm-py3.10
+      build-environment: linux-jammy-rocm-py3.10-mi200
       docker-image-name: ci-image:pytorch-linux-jammy-rocm-n-py3
       sync-tag: rocm-build
       test-matrix: |
@@ -61,7 +61,7 @@ jobs:
     permissions:
       id-token: write
       contents: read
-    name: linux-jammy-rocm-py3.10
+    name: linux-jammy-rocm-py3.10-mi200
     uses: ./.github/workflows/_rocm-test.yml
     needs:
       - linux-jammy-rocm-py3_10-build
diff --git a/.github/workflows/rocm-mi300.yml b/.github/workflows/rocm-mi300.yml
index 9c2bae06f32bd..2bc8e68aaf9c1 100644
--- a/.github/workflows/rocm-mi300.yml
+++ b/.github/workflows/rocm-mi300.yml
@@ -14,20 +14,20 @@ 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' }}
   cancel-in-progress: true
 
-permissions: read-all
+permissions:
+  id-token: write
+  contents: read
+  actions: read
 
 jobs:
   target-determination:
     if: github.repository_owner == 'pytorch'
     name: before-test
     uses: ./.github/workflows/target_determination.yml
-    permissions:
-      id-token: write
-      contents: read
 
   get-label-type:
     name: get-label-type
-    uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@main
+    uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@release/2.12
     if: ${{ (github.event_name != 'schedule' || github.repository == 'pytorch/pytorch') && github.repository_owner == 'pytorch' }}
     with:
       triggering_actor: ${{ github.triggering_actor }}
@@ -56,9 +56,6 @@ jobs:
     secrets: inherit
 
   linux-noble-rocm-py3_12-test:
-    permissions:
-      id-token: write
-      contents: read
     name: linux-noble-rocm-py3.12-mi300
     uses: ./.github/workflows/_rocm-test.yml
     needs:
diff --git a/.github/workflows/rocm-mi355.yml b/.github/workflows/rocm-mi355.yml
index 5a77695011f3c..452d29b378834 100644
--- a/.github/workflows/rocm-mi355.yml
+++ b/.github/workflows/rocm-mi355.yml
@@ -12,20 +12,20 @@ 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' }}
   cancel-in-progress: true
 
-permissions: read-all
+permissions:
+  id-token: write
+  contents: read
+  actions: read
 
 jobs:
   target-determination:
     if: github.repository_owner == 'pytorch'
     name: before-test
     uses: ./.github/workflows/target_determination.yml
-    permissions:
-      id-token: write
-      contents: read
 
   get-label-type:
     name: get-label-type
-    uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@main
+    uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@release/2.12
     if: ${{ (github.event_name != 'schedule' || github.repository == 'pytorch/pytorch') && github.repository_owner == 'pytorch' }}
     with:
       triggering_actor: ${{ github.triggering_actor }}
@@ -54,9 +54,6 @@ jobs:
     secrets: inherit
 
   linux-noble-rocm-py3_12-test:
-    permissions:
-      id-token: write
-      contents: read
     name: linux-noble-rocm-py3.12-mi355
     uses: ./.github/workflows/_rocm-test.yml
     needs:
diff --git a/.github/workflows/rocm-navi31.yml b/.github/workflows/rocm-navi31.yml
index bf1661b35e210..2c78e25a1785a 100644
--- a/.github/workflows/rocm-navi31.yml
+++ b/.github/workflows/rocm-navi31.yml
@@ -15,20 +15,20 @@ 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' }}
   cancel-in-progress: true
 
-permissions: read-all
+permissions:
+  id-token: write
+  contents: read
+  actions: read
 
 jobs:
   target-determination:
     if: github.repository_owner == 'pytorch'
     name: before-test
     uses: ./.github/workflows/target_determination.yml
-    permissions:
-      id-token: write
-      contents: read
 
   get-label-type:
     name: get-label-type
-    uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@main
+    uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@release/2.12
     if: ${{ (github.event_name != 'schedule' || github.repository == 'pytorch/pytorch') && github.repository_owner == 'pytorch' }}
     with:
       triggering_actor: ${{ github.triggering_actor }}
@@ -38,12 +38,12 @@ jobs:
 
   linux-jammy-rocm-py3_10-build:
     if: ${{ (github.event_name != 'schedule' || github.repository == 'pytorch/pytorch') && github.repository_owner == 'pytorch' }}
-    name: linux-jammy-rocm-py3.10
+    name: linux-jammy-rocm-py3.10-navi31
     uses: ./.github/workflows/_linux-build.yml
     needs: get-label-type
     with:
       runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
-      build-environment: linux-jammy-rocm-py3.10
+      build-environment: linux-jammy-rocm-py3.10-navi31
       docker-image-name: ci-image:pytorch-linux-jammy-rocm-n-py3
       sync-tag: rocm-build
       test-matrix: |
@@ -54,10 +54,7 @@ jobs:
     secrets: inherit
 
   linux-jammy-rocm-py3_10-test:
-    permissions:
-      id-token: write
-      contents: read
-    name: linux-jammy-rocm-py3_10
+    name: linux-jammy-rocm-py3.10-navi31
     uses: ./.github/workflows/_rocm-test.yml
     needs:
       - linux-jammy-rocm-py3_10-build
@@ -67,9 +64,8 @@ jobs:
       docker-image: ${{ needs.linux-jammy-rocm-py3_10-build.outputs.docker-image }}
       test-matrix: ${{ needs.linux-jammy-rocm-py3_10-build.outputs.test-matrix }}
       tests-to-include: >-
-         ${{ github.event_name == 'schedule' && 'test_nn test_torch test_cuda test_ops test_unary_ufuncs test_binary_ufuncs
          test_autograd inductor/test_torchinductor inductor/test_kernel_benchmark
          inductor/test_pad_mm inductor/test_benchmark_fusion inductor/test_aot_inductor
          inductor/test_torchinductor inductor/test_decompose_mem_bound_mm
-         inductor/test_flex_attention inductor/test_max_autotune' || '' }}
+         inductor/test_flex_attention inductor/test_max_autotune
     secrets: inherit
diff --git a/.github/workflows/rocm-nightly.yml b/.github/workflows/rocm-nightly.yml
index 649de4ab2f689..5dc1ba4d3a410 100644
--- a/.github/workflows/rocm-nightly.yml
+++ b/.github/workflows/rocm-nightly.yml
@@ -2,20 +2,33 @@ name: rocm-nightly
 
 on:
   push:
-    branches:
-      - nightly  # Trigger when nightly branch is updated
-  workflow_dispatch:  # Allow manual triggering
+    tags:
+      - ciflow/rocm-nightly/*
+  workflow_dispatch:
+  schedule:
+    - cron: 0 0 * * *  # midnight UTC
 
 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' }}
   cancel-in-progress: true
 
-permissions: read-all
+permissions:
+  id-token: write
+  contents: read
+  actions: read
 
 jobs:
+  target-determination:
+    if: github.repository_owner == 'pytorch'
+    name: before-test
+    uses: ./.github/workflows/target_determination.yml
+    permissions:
+      id-token: write
+      contents: read
+
   get-label-type:
     name: get-label-type
-    uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@main
+    uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@release/2.12
     if: ${{ (github.event_name != 'schedule' || github.repository == 'pytorch/pytorch') && github.repository_owner == 'pytorch' }}
     with:
       triggering_actor: ${{ github.triggering_actor }}
@@ -32,4 +45,28 @@ jobs:
       runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
       build-environment: linux-noble-rocm-nightly-py3.12-gfx942
       docker-image-name: ci-image:pytorch-linux-noble-rocm-nightly-py3
+      test-matrix: |
+        { include: [
+          { config: "default", shard: 1, num_shards: 6, runner: "linux.rocm.gpu.gfx942.1" },
+          { config: "default", shard: 2, num_shards: 6, runner: "linux.rocm.gpu.gfx942.1" },
+          { config: "default", shard: 3, num_shards: 6, runner: "linux.rocm.gpu.gfx942.1" },
+          { config: "default", shard: 4, num_shards: 6, runner: "linux.rocm.gpu.gfx942.1" },
+          { config: "default", shard: 5, num_shards: 6, runner: "linux.rocm.gpu.gfx942.1" },
+          { config: "default", shard: 6, num_shards: 6, runner: "linux.rocm.gpu.gfx942.1" },
+        ]}
+    secrets: inherit
+
+  linux-noble-rocm-nightly-py3_12-test:
+    permissions:
+      id-token: write
+      contents: read
+    name: linux-noble-rocm-nightly-py3.12-gfx942
+    uses: ./.github/workflows/_rocm-test.yml
+    needs:
+      - linux-noble-rocm-nightly-py3_12-build
+      - target-determination
+    with:
+      build-environment: ${{ needs.linux-noble-rocm-nightly-py3_12-build.outputs.build-environment }}
+      docker-image: ${{ needs.linux-noble-rocm-nightly-py3_12-build.outputs.docker-image }}
+      test-matrix: ${{ needs.linux-noble-rocm-nightly-py3_12-build.outputs.test-matrix }}
     secrets: inherit
diff --git a/.github/workflows/runner-determinator-validator.yml b/.github/workflows/runner-determinator-validator.yml
deleted file mode 100644
index 0710229a7b9ff..0000000000000
--- a/.github/workflows/runner-determinator-validator.yml
+++ /dev/null
@@ -1,41 +0,0 @@
-name: Validate Runner Determinator Script is in Sync
-
-on:
-  # Run on PRs when the runner-determinator script is updated to ensure it's copies are kept in sync
-  pull_request:
-    paths:
-      - .github/workflows/_runner-determinator.yml
-      - .github/workflows/runner-determinator-validator.yml
-      - .github/scripts/runner_determinator.py
-  workflow_dispatch:
-
-concurrency:
-  group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}-${{ github.event_name == 'workflow_dispatch' }}-${{ github.event_name == 'schedule' }}
-  cancel-in-progress: true
-
-jobs:
-  check-runner-determinator:
-    if: github.repository_owner == 'pytorch'
-    runs-on: ubuntu-latest
-
-    steps:
-    - name: Checkout repository
-      uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
-
-    - name: Run Hardcode runner-determinator script
-      id: hardcode-script
-      run: |
-        # Extract the script content from _runner-determinator.yml and skip the first 10 spaces of each line
-        script_content=$(awk '/cat <<EOF > runner_determinator.py/{flag=1;next}/EOF$/{flag=0}flag{print substr($0, 11)}' .github/workflows/_runner-determinator.yml)
-
-        # Write the extracted script content to runner_determinator.py
-        echo "$script_content" > runner_determinator_workflow.py
-
-    - name: Compare runner-determinator script embedded in workflow with checked in script
-      run: |
-        # Compare the extracted runner_determinator script with the existing one
-        # If this check fails, then make sure the contents of .github/scripts/runner_determinator.py is in sync with the
-        # version embedded into .github/workflows/_runner-determinator.yml
-        diff runner_determinator_workflow.py .github/scripts/runner_determinator.py
-      # Fail the job if the scripts are not identical
-      continue-on-error: false
diff --git a/.github/workflows/runner_determinator_script_sync.yaml b/.github/workflows/runner_determinator_script_sync.yaml
deleted file mode 100644
index a5f52f6980f7c..0000000000000
--- a/.github/workflows/runner_determinator_script_sync.yaml
+++ /dev/null
@@ -1,43 +0,0 @@
-name: runner-determinator
-
-on:
-  workflow_dispatch:
-  pull_request:
-    branches: [main]
-    paths:
-      - .github/workflows/_runner-determinator.yaml
-      - .github/workflows/_runner_determinator_script_sync.yaml
-      - .github/workflows/scripts/runner_determinator.py
-
-jobs:
-  python-script-sync-check:
-    if: github.repository_owner == 'pytorch'
-    runs-on: ubuntu-latest
-
-    steps:
-      - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
-        with:
-          sparse-checkout: |
-            .github
-
-      - name: Extract the script from runner_determinator
-        run: |
-            # Runner determinator files
-            RUNNER_DETERMINATOR_WORKFLOW_FILE=.github/workflows/_runner-determinator.yml
-            RUNNER_DETERMINATOR_PYTHON_SCRIPT_FILE=.github/scripts/runner_determinator.py
-
-            # Parse the job file, extract the script and run it, up to the final EOF,
-            # to generate the python file in the local folder
-            yq '.jobs.runner-determinator.steps[] | select(.id == "hardcode-script") | .run' \
-                "${RUNNER_DETERMINATOR_WORKFLOW_FILE}" | sed '/^EOF$/q' | bash
-
-            set +e
-            DIFF="$(diff "$(basename ${RUNNER_DETERMINATOR_PYTHON_SCRIPT_FILE})" ${RUNNER_DETERMINATOR_PYTHON_SCRIPT_FILE})"
-            IS_DIFF=$?
-            set -e
-            if [ $IS_DIFF -eq 0 ]; then
-                echo "Scripts are in sync! ^_^";
-            else
-                echo -e "Scripts are *NOT* in sync:\n ${DIFF}";
-                exit 1
-            fi
diff --git a/.github/workflows/s390x-periodic.yml b/.github/workflows/s390x-periodic.yml
index bf94ee3bb00f1..f583a5b425a44 100644
--- a/.github/workflows/s390x-periodic.yml
+++ b/.github/workflows/s390x-periodic.yml
@@ -18,23 +18,18 @@ concurrency:
 permissions:
   id-token: write
   contents: read
+  actions: read
 
 jobs:
   llm-td:
     if: github.repository_owner == 'pytorch'
     name: before-test
     uses: ./.github/workflows/llm_td_retrieval.yml
-    permissions:
-      id-token: write
-      contents: read
 
   target-determination:
     name: before-test
     uses: ./.github/workflows/target_determination.yml
     needs: llm-td
-    permissions:
-      id-token: write
-      contents: read
 
   linux-manylinux-2_28-py3-cpu-s390x-build:
     if: github.repository_owner == 'pytorch'
@@ -60,9 +55,6 @@ jobs:
     secrets: inherit
 
   linux-manylinux-2_28-py3-cpu-s390x-test:
-    permissions:
-      id-token: write
-      contents: read
     name: linux-manylinux-2_28-py3-cpu-s390x
     uses: ./.github/workflows/_linux-test.yml
     needs:
diff --git a/.github/workflows/slow-rocm-mi200.yml b/.github/workflows/slow-rocm-mi200.yml
index 937f04980522e..fe17280b5b2e8 100644
--- a/.github/workflows/slow-rocm-mi200.yml
+++ b/.github/workflows/slow-rocm-mi200.yml
@@ -5,13 +5,13 @@ name: slow-rocm-mi200
 
 on:
   push:
-    branches:
-      - release/*
+    #branches:
+      #- release/*
     tags:
       - ciflow/slow/*
       - ciflow/slow-rocm-mi200/*
-  schedule:
-    - cron: 0 */3 * * *
+  #schedule:
+    #- cron: 0 */3 * * *
   workflow_dispatch:
 
 concurrency:
@@ -21,27 +21,22 @@ concurrency:
 permissions:
   id-token: write
   contents: read
+  actions: read
 
 jobs:
   llm-td:
     if: github.repository_owner == 'pytorch'
     name: before-test
     uses: ./.github/workflows/llm_td_retrieval.yml
-    permissions:
-      id-token: write
-      contents: read
 
   target-determination:
     name: before-test
     uses: ./.github/workflows/target_determination.yml
     needs: llm-td
-    permissions:
-      id-token: write
-      contents: read
 
   get-label-type:
     name: get-label-type
-    uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@main
+    uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@release/2.12
     if: ${{ (github.event_name != 'schedule' || github.repository == 'pytorch/pytorch') && github.repository_owner == 'pytorch' }}
     with:
       triggering_actor: ${{ github.triggering_actor }}
@@ -50,12 +45,12 @@ jobs:
       curr_ref_type: ${{ github.ref_type }}
 
   linux-jammy-rocm-py3_10-build:
-    name: linux-jammy-rocm-py3.10
+    name: linux-jammy-rocm-py3.10-mi200
     uses: ./.github/workflows/_linux-build.yml
     needs: get-label-type
     with:
       runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
-      build-environment: linux-jammy-rocm-py3.10
+      build-environment: linux-jammy-rocm-py3.10-mi200
       docker-image-name: ci-image:pytorch-linux-jammy-rocm-n-py3
       sync-tag: rocm-build
       test-matrix: |
@@ -66,10 +61,7 @@ jobs:
     secrets: inherit
 
   linux-jammy-rocm-py3_10-test:
-    permissions:
-      id-token: write
-      contents: read
-    name: linux-jammy-rocm-py3.10
+    name: linux-jammy-rocm-py3.10-mi200
     uses: ./.github/workflows/_rocm-test.yml
     needs:
       - linux-jammy-rocm-py3_10-build
diff --git a/.github/workflows/slow.yml b/.github/workflows/slow.yml
index 8da9c9bd219d5..d5cb9d02c75fa 100644
--- a/.github/workflows/slow.yml
+++ b/.github/workflows/slow.yml
@@ -21,27 +21,22 @@ concurrency:
 permissions:
   id-token: write
   contents: read
+  actions: read
 
 jobs:
   llm-td:
     if: github.repository_owner == 'pytorch'
     name: before-test
     uses: ./.github/workflows/llm_td_retrieval.yml
-    permissions:
-      id-token: write
-      contents: read
 
   target-determination:
     name: before-test
     uses: ./.github/workflows/target_determination.yml
     needs: llm-td
-    permissions:
-      id-token: write
-      contents: read
 
   get-label-type:
     name: get-label-type
-    uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@main
+    uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@release/2.12
     if: ${{ (github.event_name != 'schedule' || github.repository == 'pytorch/pytorch') && github.repository_owner == 'pytorch' }}
     with:
       triggering_actor: ${{ github.triggering_actor }}
@@ -107,14 +102,14 @@ jobs:
       test-matrix: ${{ needs.linux-jammy-cuda13_0-py3_10-gcc11-sm86-build.outputs.test-matrix }}
     secrets: inherit
 
-  linux-jammy-py3_10-clang15-build:
-    name: linux-jammy-py3.10-clang15
+  linux-jammy-py3_10-clang18-build:
+    name: linux-jammy-py3.10-clang18
     uses: ./.github/workflows/_linux-build.yml
     needs: get-label-type
     with:
       runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
-      build-environment: linux-jammy-py3.10-clang15
-      docker-image-name: ci-image:pytorch-linux-jammy-py3.10-clang15
+      build-environment: linux-jammy-py3.10-clang18
+      docker-image-name: ci-image:pytorch-linux-jammy-py3.10-clang18
       test-matrix: |
         { include: [
           { config: "slow", shard: 1, num_shards: 2, runner: "linux.2xlarge" },
@@ -122,16 +117,16 @@ jobs:
         ]}
     secrets: inherit
 
-  linux-jammy-py3_10-clang15-test:
-    name: linux-jammy-py3.10-clang15
+  linux-jammy-py3_10-clang18-test:
+    name: linux-jammy-py3.10-clang18
     uses: ./.github/workflows/_linux-test.yml
     needs:
-      - linux-jammy-py3_10-clang15-build
+      - linux-jammy-py3_10-clang18-build
       - target-determination
     with:
-      build-environment: ${{ needs.linux-jammy-py3_10-clang15-build.outputs.build-environment }}
-      docker-image: ${{ needs.linux-jammy-py3_10-clang15-build.outputs.docker-image }}
-      test-matrix: ${{ needs.linux-jammy-py3_10-clang15-build.outputs.test-matrix }}
+      build-environment: ${{ needs.linux-jammy-py3_10-clang18-build.outputs.build-environment }}
+      docker-image: ${{ needs.linux-jammy-py3_10-clang18-build.outputs.docker-image }}
+      test-matrix: ${{ needs.linux-jammy-py3_10-clang18-build.outputs.test-matrix }}
     secrets: inherit
 
   linux-jammy-py3_10-clang18-asan-build:
@@ -142,7 +137,7 @@ jobs:
       runner: linux.c7i.4xlarge
       runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
       build-environment: linux-jammy-py3.10-clang18-asan
-      docker-image-name: ci-image:pytorch-linux-jammy-py3-clang18-asan
+      docker-image-name: ci-image:pytorch-linux-jammy-py3.10-clang18
       test-matrix: |
         { include: [
           { config: "slow", shard: 1, num_shards: 3, runner: "${{ needs.get-label-type.outputs.label-type }}linux.4xlarge" },
diff --git a/.github/workflows/target-determination-indexer.yml b/.github/workflows/target-determination-indexer.yml
index 3438b1dd5ac57..01ab07ece6a67 100644
--- a/.github/workflows/target-determination-indexer.yml
+++ b/.github/workflows/target-determination-indexer.yml
@@ -13,7 +13,7 @@ jobs:
   get-label-type:
     if: github.repository_owner == 'pytorch'
     name: get-label-type
-    uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@main
+    uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@release/2.12
     with:
       triggering_actor: ${{ github.triggering_actor }}
       issue_owner: ${{ github.event.pull_request.user.login || github.event.issue.user.login }}
@@ -31,14 +31,14 @@ jobs:
           path: pytorch
 
       - name: Setup Linux
-        uses: ./pytorch/.github/actions/setup-linux
+        uses: pytorch/pytorch/.github/actions/setup-linux@release/2.12
 
       - name: Login to ECR
         uses: ./pytorch/.github/actions/ecr-login
 
       - name: Calculate docker image
         id: calculate-docker-image
-        uses: pytorch/test-infra/.github/actions/calculate-docker-image@main
+        uses: pytorch/test-infra/.github/actions/calculate-docker-image@release/2.12
         with:
           docker-image-name: ci-image:pytorch-linux-jammy-cuda12.8-cudnn9-py3-gcc11
           working-directory: pytorch
@@ -53,13 +53,13 @@ jobs:
           echo "docker pull ghcr.io/pytorch/ci-image:${tag/:/-}"
 
       - name: Pull docker image
-        uses: pytorch/test-infra/.github/actions/pull-docker-image@main
+        uses: pytorch/test-infra/.github/actions/pull-docker-image@release/2.12
         with:
           docker-image: ${{ steps.calculate-docker-image.outputs.docker-image }}
 
       - name: Install nvidia driver, nvidia-docker runtime, set GPU_FLAG
         id: install-nvidia-driver
-        uses: pytorch/test-infra/.github/actions/setup-nvidia@main
+        uses: pytorch/test-infra/.github/actions/setup-nvidia@release/2.12
 
       - name: Clone CodeLlama
         uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
@@ -152,7 +152,7 @@ jobs:
             "s3://target-determinator-assets/indexes/latest/${ZIP_NAME}"
 
       - name: Teardown Linux
-        uses: pytorch/test-infra/.github/actions/teardown-linux@main
+        uses: pytorch/test-infra/.github/actions/teardown-linux@release/2.12
         if: always()
 
 concurrency:
diff --git a/.github/workflows/target_determination.yml b/.github/workflows/target_determination.yml
index c712b11185a76..ded830557d4e9 100644
--- a/.github/workflows/target_determination.yml
+++ b/.github/workflows/target_determination.yml
@@ -9,7 +9,7 @@ jobs:
     name: get-label-type
     # Don't run on forked repos
     if: github.repository_owner == 'pytorch'
-    uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@main
+    uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@release/2.12
     with:
       triggering_actor: ${{ github.triggering_actor }}
       issue_owner: ${{ github.event.pull_request.user.login || github.event.issue.user.login }}
@@ -22,18 +22,11 @@ jobs:
     runs-on: "${{ needs.get-label-type.outputs.label-type }}linux.2xlarge"
     needs: get-label-type
     steps:
-      # [pytorch repo ref]
-      # Use a pytorch/pytorch reference instead of a reference to the local
-      # checkout because when we run this action we don't *have* a local
-      # checkout. In other cases you should prefer a local checkout.
-      - name: Checkout PyTorch
-        uses: pytorch/pytorch/.github/actions/checkout-pytorch@main
+      - name: Setup Linux
+        uses: pytorch/pytorch/.github/actions/setup-linux@release/2.12
         with:
           submodules: false
 
-      - name: Setup Linux
-        uses: ./.github/actions/setup-linux
-
       - name: Get workflow job id
         id: get-job-id
         uses: ./.github/actions/get-workflow-job-id
diff --git a/.github/workflows/test-b200.yml b/.github/workflows/test-b200.yml
index 19dcb07c29844..1a1dba622da74 100644
--- a/.github/workflows/test-b200.yml
+++ b/.github/workflows/test-b200.yml
@@ -35,13 +35,14 @@ concurrency:
 permissions:
   id-token: write
   contents: read
+  actions: read
 
 jobs:
 
   get-label-type:
     if: github.repository_owner == 'pytorch'
     name: get-label-type
-    uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@main
+    uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@release/2.12
     with:
       triggering_actor: ${{ github.triggering_actor }}
       issue_owner: ${{ github.event.pull_request.user.login || github.event.issue.user.login }}
diff --git a/.github/workflows/test-check-binary.yml b/.github/workflows/test-check-binary.yml
index 883b2d253aa8f..d60f816d95b08 100644
--- a/.github/workflows/test-check-binary.yml
+++ b/.github/workflows/test-check-binary.yml
@@ -15,7 +15,7 @@ jobs:
   check_binary_linux_cpu:
     if: github.repository_owner == 'pytorch'
     name: Test check_binary.sh for Linux CPU
-    uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main
+    uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@release/2.12
     with:
       docker-image: python:3.11
       docker-build-dir: "skip-docker-build"
@@ -30,7 +30,7 @@ jobs:
   check_binary_linux_cuda:
     if: github.repository_owner == 'pytorch'
     name: Test check_binary.sh for Linux CUDA
-    uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main
+    uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@release/2.12
     with:
       runner: linux.g4dn.4xlarge.nvidia.gpu
       docker-image: python:3.11
diff --git a/.github/workflows/test-h100.yml b/.github/workflows/test-h100.yml
index 4351b427b0b8a..5964585f1a839 100644
--- a/.github/workflows/test-h100.yml
+++ b/.github/workflows/test-h100.yml
@@ -22,13 +22,14 @@ concurrency:
 permissions:
   id-token: write
   contents: read
+  actions: read
 
 jobs:
 
   get-label-type:
     if: github.repository_owner == 'pytorch'
     name: get-label-type
-    uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@main
+    uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@release/2.12
     with:
       triggering_actor: ${{ github.triggering_actor }}
       issue_owner: ${{ github.event.pull_request.user.login || github.event.issue.user.login }}
diff --git a/.github/workflows/tools-unit-tests.yml b/.github/workflows/tools-unit-tests.yml
index 4f87992eb5d72..cb58f534eb0da 100644
--- a/.github/workflows/tools-unit-tests.yml
+++ b/.github/workflows/tools-unit-tests.yml
@@ -25,7 +25,7 @@ jobs:
     runs-on: ubuntu-latest
     steps:
       - name: Checkout pytorch
-        uses: pytorch/pytorch/.github/actions/checkout-pytorch@main
+        uses: pytorch/pytorch/.github/actions/checkout-pytorch@release/2.12
         with:
           submodules: true
           fetch-depth: 0
@@ -52,7 +52,7 @@ jobs:
     runs-on: ubuntu-latest
     steps:
       - name: Checkout pytorch
-        uses: pytorch/pytorch/.github/actions/checkout-pytorch@main
+        uses: pytorch/pytorch/.github/actions/checkout-pytorch@release/2.12
         with:
           submodules: true
           fetch-depth: 0
diff --git a/.github/workflows/torchbench.yml b/.github/workflows/torchbench.yml
index 508c39a653600..8195a87720a1e 100644
--- a/.github/workflows/torchbench.yml
+++ b/.github/workflows/torchbench.yml
@@ -13,12 +13,13 @@ concurrency:
 permissions:
   id-token: write
   contents: read
+  actions: read
 
 jobs:
   get-default-label-prefix:
     if: github.repository_owner == 'pytorch'
     name: get-default-label-prefix
-    uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@main
+    uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@release/2.12
     with:
       triggering_actor: ${{ github.triggering_actor }}
       issue_owner: ${{ github.event.pull_request.user.login || github.event.issue.user.login }}
@@ -26,14 +27,14 @@ jobs:
       curr_ref_type: ${{ github.ref_type }}
 
   build:
-    name: cuda12.8-py3.10-gcc11-sm80
+    name: cuda13.0-py3.10-gcc11-sm80
     uses: ./.github/workflows/_linux-build.yml
     needs:
       - get-default-label-prefix
     with:
       runner_prefix: "${{ needs.get-default-label-prefix.outputs.label-type }}"
-      build-environment: linux-jammy-cuda12.8-py3.10-gcc11-sm80
-      docker-image-name: ci-image:pytorch-linux-jammy-cuda12.8-cudnn9-py3-gcc11-inductor-benchmarks
+      build-environment: linux-jammy-cuda13.0-py3.10-gcc11-sm80
+      docker-image-name: ci-image:pytorch-linux-jammy-cuda13.0-cudnn9-py3-gcc11-inductor-benchmarks
       cuda-arch-list: '8.0'
       test-matrix: |
         { include: [
@@ -42,7 +43,7 @@ jobs:
     secrets: inherit
 
   test:
-    name: cuda12.8-py3.10-gcc11-sm80
+    name: cuda13.0-py3.10-gcc11-sm80
     uses: ./.github/workflows/_linux-test.yml
     needs: build
     with:
diff --git a/.github/workflows/torchtitan.yml b/.github/workflows/torchtitan.yml
new file mode 100644
index 0000000000000..b12f980ec4202
--- /dev/null
+++ b/.github/workflows/torchtitan.yml
@@ -0,0 +1,47 @@
+name: torchtitan-test
+
+on:
+  push:
+    branches:
+      - main
+      - release/*
+    tags:
+      - ciflow/torchtitan/*
+  workflow_dispatch:
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref_name }}-${{ github.ref_type == 'branch' && github.sha }}-${{ github.event_name == 'workflow_dispatch' }}
+  cancel-in-progress: true
+
+permissions:
+  id-token: write
+  contents: read
+  actions: read
+
+jobs:
+  build:
+    name: torchtitan-x-pytorch-build
+    if: github.repository_owner == 'pytorch'
+    uses: ./.github/workflows/_linux-build.yml
+    with:
+      build-environment: linux-jammy-cuda12.8-py3-gcc11
+      docker-image-name: ci-image:pytorch-linux-jammy-cuda12.8-cudnn9-py3-gcc11
+      cuda-arch-list: '8.0 8.9 9.0 10.0'
+      runner: linux.24xlarge.memory
+      # torchtitan_h100_integration disabled pending runner availability
+      test-matrix: |
+        { include: [
+          { config: "torchtitan_features_integration", shard: 1, num_shards: 1, runner: "linux.g5.48xlarge.nvidia.gpu" },
+          { config: "torchtitan_models_integration", shard: 1, num_shards: 1, runner: "linux.g5.48xlarge.nvidia.gpu" }
+        ]}
+    secrets: inherit
+
+  test:
+    name: torchtitan-x-pytorch-test
+    uses: ./.github/workflows/_linux-test.yml
+    needs: build
+    with:
+      build-environment: linux-jammy-cuda12.8-py3-gcc11
+      docker-image: ${{ needs.build.outputs.docker-image }}
+      test-matrix: ${{ needs.build.outputs.test-matrix }}
+    secrets: inherit
diff --git a/.github/workflows/trunk-rocm-sandbox.yml b/.github/workflows/trunk-rocm-sandbox.yml
index aee6a5d87df09..9e80cd282769d 100644
--- a/.github/workflows/trunk-rocm-sandbox.yml
+++ b/.github/workflows/trunk-rocm-sandbox.yml
@@ -1,13 +1,13 @@
 name: trunk-rocm-sandbox
 
 on:
-  push:
-    branches:
-      - main
-      - release/*
+  #push:
+  #  branches:
+  #    - main
+  #    - release/*
   workflow_dispatch:
   schedule:
-    - cron: 29 8 * * *  # about 1:29am PDT
+    - cron: 0 0,3,6,9,12,15,18,21 * * * # run every 3 hours for regression coverage
 
 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' }}
@@ -16,27 +16,22 @@ concurrency:
 permissions:
   id-token: write
   contents: read
+  actions: read
 
 jobs:
   llm-td:
     if: github.repository_owner == 'pytorch'
     name: before-test
     uses: ./.github/workflows/llm_td_retrieval.yml
-    permissions:
-      id-token: write
-      contents: read
 
   target-determination:
     name: before-test
     uses: ./.github/workflows/target_determination.yml
     needs: llm-td
-    permissions:
-      id-token: write
-      contents: read
 
   get-label-type:
     name: get-label-type
-    uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@main
+    uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@release/2.12
     if: ${{ (github.event_name != 'schedule' || github.repository == 'pytorch/pytorch') && github.repository_owner == 'pytorch' }}
     with:
       triggering_actor: ${{ github.triggering_actor }}
@@ -55,22 +50,21 @@ jobs:
       sync-tag: rocm-build
       test-matrix: |
         { include: [
-          { config: "default", shard: 1, num_shards: 6, runner: "linux.rocm.gpu.gfx950.1.test" },
-          { config: "default", shard: 2, num_shards: 6, runner: "linux.rocm.gpu.gfx950.1.test" },
-          { config: "default", shard: 3, num_shards: 6, runner: "linux.rocm.gpu.gfx950.1.test" },
-          { config: "default", shard: 4, num_shards: 6, runner: "linux.rocm.gpu.gfx950.1.test" },
-          { config: "default", shard: 5, num_shards: 6, runner: "linux.rocm.gpu.gfx950.1.test" },
-          { config: "default", shard: 6, num_shards: 6, runner: "linux.rocm.gpu.gfx950.1.test" },
-          { config: "distributed", shard: 1, num_shards: 3, runner: "linux.rocm.gpu.gfx950.4.test" },
-          { config: "distributed", shard: 2, num_shards: 3, runner: "linux.rocm.gpu.gfx950.4.test" },
-          { config: "distributed", shard: 3, num_shards: 3, runner: "linux.rocm.gpu.gfx950.4.test" },
+          { config: "default", shard: 1, num_shards: 6, runner: "linux.rocm.gpu.mi210.1.test" },
+          { config: "default", shard: 2, num_shards: 6, runner: "linux.rocm.gpu.mi210.1.test" },
+          { config: "default", shard: 3, num_shards: 6, runner: "linux.rocm.gpu.mi210.1.test" },
+          { config: "default", shard: 4, num_shards: 6, runner: "linux.rocm.gpu.mi210.1.test" },
+          { config: "default", shard: 5, num_shards: 6, runner: "linux.rocm.gpu.mi210.1.test" },
+          { config: "default", shard: 6, num_shards: 6, runner: "linux.rocm.gpu.mi210.1.test" },
+          { config: "inductor", shard: 1, num_shards: 2, runner: "linux.rocm.gpu.mi210.1.test" },
+          { config: "inductor", shard: 2, num_shards: 2, runner: "linux.rocm.gpu.mi210.1.test" },
+          { config: "distributed", shard: 1, num_shards: 3, runner: "linux.rocm.gpu.mi210.2.test" },
+          { config: "distributed", shard: 2, num_shards: 3, runner: "linux.rocm.gpu.mi210.2.test" },
+          { config: "distributed", shard: 3, num_shards: 3, runner: "linux.rocm.gpu.mi210.2.test" },
         ]}
     secrets: inherit
 
   linux-jammy-rocm-py3_10-test:
-    permissions:
-      id-token: write
-      contents: read
     name: linux-jammy-rocm-py3.10
     uses: ./.github/workflows/_rocm-test.yml
     needs:
diff --git a/.github/workflows/trunk-tagging.yml b/.github/workflows/trunk-tagging.yml
index d96f2de8366aa..27472f786794c 100644
--- a/.github/workflows/trunk-tagging.yml
+++ b/.github/workflows/trunk-tagging.yml
@@ -28,8 +28,8 @@ jobs:
       - name: Pre-checkout validation
         run: |
           # For workflow_dispatch, validate SHA format before checkout
-          if [ -n "${{ github.event.inputs.commit_sha }}" ]; then
-            COMMIT_SHA="${{ github.event.inputs.commit_sha }}"
+          if [ -n "${INPUT_COMMIT_SHA}" ]; then
+            COMMIT_SHA="${INPUT_COMMIT_SHA}"
 
             # Verify it's a well-formed SHA (40 hex characters)
             if ! echo "${COMMIT_SHA}" | grep -qE '^[a-f0-9]{40}$'; then
@@ -42,6 +42,8 @@ jobs:
             echo "✅ Using current commit SHA - no pre-checkout validation needed"
           fi
 
+        env:
+          INPUT_COMMIT_SHA: ${{ github.event.inputs.commit_sha }}
       - name: Checkout repository
         uses: actions/checkout@v4
         with:
@@ -53,8 +55,8 @@ jobs:
       - name: Set commit SHA
         id: commit
         run: |
-          if [ -n "${{ github.event.inputs.commit_sha }}" ]; then
-            COMMIT_SHA="${{ github.event.inputs.commit_sha }}"
+          if [ -n "${INPUT_COMMIT_SHA}" ]; then
+            COMMIT_SHA="${INPUT_COMMIT_SHA}"
           else
             COMMIT_SHA="${{ github.sha }}"
           fi
@@ -63,6 +65,8 @@ jobs:
             echo "tag_name=trunk/${COMMIT_SHA}"
           } >> "${GITHUB_OUTPUT}"
 
+        env:
+          INPUT_COMMIT_SHA: ${{ github.event.inputs.commit_sha }}
       - name: Validate commit SHA
         run: |
           COMMIT_SHA="${{ steps.commit.outputs.sha }}"
@@ -74,7 +78,7 @@ jobs:
           fi
 
           # For workflow_dispatch, verify the commit exists on main branch
-          if [ -n "${{ github.event.inputs.commit_sha }}" ]; then
+          if [ -n "${INPUT_COMMIT_SHA}" ]; then
             echo "Manual dispatch detected - validating commit is on main branch..."
 
             # Get all commits reachable from main branch
@@ -89,6 +93,8 @@ jobs:
             echo "✅ Commit ${COMMIT_SHA} is valid (automatic push trigger)"
           fi
 
+        env:
+          INPUT_COMMIT_SHA: ${{ github.event.inputs.commit_sha }}
       - name: Create and push tag(s) with retry
         id: check_tag
         env:
@@ -309,7 +315,10 @@ jobs:
             echo "  Name: ${{ steps.commit.outputs.tag_name }}"
             echo "  Commit: ${{ steps.commit.outputs.sha }}"
             echo "  Trigger: ${{ github.event_name }}"
-            if [ -n "${{ github.event.inputs.commit_sha }}" ]; then
-              echo "  Manual commit: ${{ github.event.inputs.commit_sha }}"
+            if [ -n "${INPUT_COMMIT_SHA}" ]; then
+              echo "  Manual commit: ${INPUT_COMMIT_SHA}"
             fi
           fi
+
+        env:
+          INPUT_COMMIT_SHA: ${{ github.event.inputs.commit_sha }}
\ No newline at end of file
diff --git a/.github/workflows/trunk.yml b/.github/workflows/trunk.yml
index 65bc108c842af..d7e1d0b1cab24 100644
--- a/.github/workflows/trunk.yml
+++ b/.github/workflows/trunk.yml
@@ -30,6 +30,7 @@ concurrency:
 permissions:
   id-token: write
   contents: read
+  actions: read
 
 jobs:
   # See job-filter.yml for rules on adding job filter conditions
@@ -44,21 +45,15 @@ jobs:
     if: github.repository_owner == 'pytorch'
     name: before-test
     uses: ./.github/workflows/llm_td_retrieval.yml
-    permissions:
-      id-token: write
-      contents: read
 
   target-determination:
     name: before-test
     uses: ./.github/workflows/target_determination.yml
     needs: llm-td
-    permissions:
-      id-token: write
-      contents: read
 
   get-label-type:
     name: get-label-type
-    uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@main
+    uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@release/2.12
     if: ${{ (github.event_name != 'schedule' || github.repository == 'pytorch/pytorch') && github.repository_owner == 'pytorch' }}
     with:
       triggering_actor: ${{ github.triggering_actor }}
@@ -78,7 +73,7 @@ jobs:
       docker-image-name: ci-image:pytorch-linux-jammy-cuda12.8-cudnn9-py3-gcc11
       build-generates-artifacts: false
       runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
-      runner: "linux.c7i.4xlarge"
+      runner: "linux.r7i.4xlarge"
       test-matrix: |
         { include: [
           { config: "default", shard: 1, num_shards: 1 },
@@ -86,7 +81,7 @@ jobs:
     secrets: inherit
 
   linux-jammy-cuda12_8-py3_10-gcc11-build:
-    if: ${{ needs.job-filter.outputs.jobs == '' || contains(needs.job-filter.outputs.jobs, ' linux-jammy-cuda12.8-py3.10-gcc11 ') || contains(needs.job-filter.outputs.jobs, ' cross-compile-linux-test ') }}
+    if: ${{ needs.job-filter.outputs.jobs == '' || contains(needs.job-filter.outputs.jobs, ' linux-jammy-cuda12.8-py3.10-gcc11 ') }}
     name: linux-jammy-cuda12.8-py3.10-gcc11
     uses: ./.github/workflows/_linux-build.yml
     needs:
@@ -107,29 +102,17 @@ jobs:
           { config: "distributed", shard: 1, num_shards: 3, runner: "${{ needs.get-label-type.outputs.label-type }}linux.g4dn.12xlarge.nvidia.gpu" },
           { config: "distributed", shard: 2, num_shards: 3, runner: "${{ needs.get-label-type.outputs.label-type }}linux.g4dn.12xlarge.nvidia.gpu" },
           { config: "distributed", shard: 3, num_shards: 3, runner: "${{ needs.get-label-type.outputs.label-type }}linux.g4dn.12xlarge.nvidia.gpu" },
-          { config: "pr_time_benchmarks", shard: 1, num_shards: 1, runner: "linux.g4dn.metal.nvidia.gpu" },
           { config: "libtorch_agnostic_targetting", shard: 1, num_shards: 1, runner: "linux.g4dn.metal.nvidia.gpu" },
         ]}
     secrets: inherit
 
-  linux-jammy-cuda12_8-py3_10-gcc11-test:
-    if: ${{ needs.job-filter.outputs.jobs == '' || contains(needs.job-filter.outputs.jobs, ' linux-jammy-cuda12.8-py3.10-gcc11 ') }}
-    name: linux-jammy-cuda12.8-py3.10-gcc11
-    uses: ./.github/workflows/_linux-test.yml
-    needs:
-      - linux-jammy-cuda12_8-py3_10-gcc11-build
-      - target-determination
-      - job-filter
-    with:
-      timeout-minutes: 360
-      build-environment: ${{ needs.linux-jammy-cuda12_8-py3_10-gcc11-build.outputs.build-environment }}
-      docker-image: ${{ needs.linux-jammy-cuda12_8-py3_10-gcc11-build.outputs.docker-image }}
-      test-matrix: ${{ needs.linux-jammy-cuda12_8-py3_10-gcc11-build.outputs.test-matrix }}
-      tests-to-include: ${{ github.event.inputs.tests-to-include || '' }}
-    secrets: inherit
+  # CUDA 12.8 GPU tests moved to periodic.yml to reduce per-commit compute.
+  # CUDA 13.0 (the newer shipping version) remains per-commit for forward-looking coverage.
+  # The 12.8 build job is kept because cross-compile-linux-test depends on it.
+  # See P2188981399 for the full CI workflow analysis.
 
   linux-jammy-cuda13_0-py3_10-gcc11-build:
-    if: ${{ needs.job-filter.outputs.jobs == '' || contains(needs.job-filter.outputs.jobs, ' linux-jammy-cuda13.0-py3.10-gcc11 ') }}
+    if: ${{ needs.job-filter.outputs.jobs == '' || contains(needs.job-filter.outputs.jobs, ' linux-jammy-cuda13.0-py3.10-gcc11 ') || contains(needs.job-filter.outputs.jobs, ' cross-compile-linux-test-cuda13 ') }}
     name: linux-jammy-cuda13.0-py3.10-gcc11
     uses: ./.github/workflows/_linux-build.yml
     needs:
@@ -171,23 +154,7 @@ jobs:
     secrets: inherit
 
   # no-ops builds test USE_PER_OPERATOR_HEADERS=0 where ATen/ops is not generated
-  linux-jammy-cuda12_8-py3_10-gcc11-no-ops-build:
-    if: ${{ needs.job-filter.outputs.jobs == '' || contains(needs.job-filter.outputs.jobs, ' linux-jammy-cuda12.8-py3.10-gcc11-no-ops ') }}
-    name: linux-jammy-cuda12.8-py3.10-gcc11-no-ops
-    uses: ./.github/workflows/_linux-build.yml
-    needs:
-      - get-label-type
-      - job-filter
-    with:
-      runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
-      build-environment: linux-jammy-cuda12.8-py3.10-gcc11-no-ops
-      docker-image-name: ci-image:pytorch-linux-jammy-cuda12.8-cudnn9-py3-gcc11
-      test-matrix: |
-        { include: [
-          { config: "default", shard: 1, num_shards: 1 },
-        ]}
-    secrets: inherit
-
+  # CUDA 12.8 no-ops build moved to periodic (only 13.0 remains per-commit)
   linux-jammy-cuda13_0-py3_10-gcc11-no-ops-build:
     if: ${{ needs.job-filter.outputs.jobs == '' || contains(needs.job-filter.outputs.jobs, ' linux-jammy-cuda13.0-py3.10-gcc11-no-ops ') }}
     name: linux-jammy-cuda13.0-py3.10-gcc11-no-ops
@@ -281,28 +248,29 @@ jobs:
       disable-monitor: false
     secrets: inherit
 
-  win-vs2022-cuda12_8-py3-build:
-    if: ${{ needs.job-filter.outputs.jobs == '' || contains(needs.job-filter.outputs.jobs, ' win-vs2022-cuda12.8-py3 ') || contains(needs.job-filter.outputs.jobs, ' cross-compile-linux-test ') }}
-    name: win-vs2022-cuda12.8-py3
+
+  win-vs2022-cuda13_0-py3-build:
+    if: ${{ needs.job-filter.outputs.jobs == '' || contains(needs.job-filter.outputs.jobs, ' win-vs2022-cuda13.0-py3 ') || contains(needs.job-filter.outputs.jobs, ' cross-compile-linux-test-cuda13 ') }}
+    name: win-vs2022-cuda13.0-py3
     uses: ./.github/workflows/_win-build.yml
     needs:
       - get-label-type
       - job-filter
     with:
-      build-environment: win-vs2022-cuda12.8-py3
-      cuda-version: "12.8"
+      build-environment: win-vs2022-cuda13.0-py3
+      cuda-version: "13.0"
       runner: "${{ needs.get-label-type.outputs.label-type }}windows.4xlarge.nonephemeral"
     secrets: inherit
 
   linux-jammy-rocm-py3_10-build:
-    name: linux-jammy-rocm-py3.10
+    name: linux-jammy-rocm-py3.10-mi355
     uses: ./.github/workflows/_linux-build.yml
     needs:
       - get-label-type
       - job-filter
     with:
       runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
-      build-environment: linux-jammy-rocm-py3.10
+      build-environment: linux-jammy-rocm-py3.10-mi355
       docker-image-name: ci-image:pytorch-linux-jammy-rocm-n-py3
       sync-tag: rocm-build
       test-matrix: |
@@ -313,17 +281,14 @@ jobs:
           { config: "default", shard: 4, num_shards: 6, runner: "linux.rocm.gpu.gfx950.1" },
           { config: "default", shard: 5, num_shards: 6, runner: "linux.rocm.gpu.gfx950.1" },
           { config: "default", shard: 6, num_shards: 6, runner: "linux.rocm.gpu.gfx950.1" },
-          { config: "distributed", shard: 1, num_shards: 3, runner: "linux.rocm.gpu.gfx950.4" },
-          { config: "distributed", shard: 2, num_shards: 3, runner: "linux.rocm.gpu.gfx950.4" },
-          { config: "distributed", shard: 3, num_shards: 3, runner: "linux.rocm.gpu.gfx950.4" },
+          { config: "distributed", shard: 1, num_shards: 3, runner: "linux.rocm.gpu.gfx950.2" },
+          { config: "distributed", shard: 2, num_shards: 3, runner: "linux.rocm.gpu.gfx950.2" },
+          { config: "distributed", shard: 3, num_shards: 3, runner: "linux.rocm.gpu.gfx950.2" },
         ]}
     secrets: inherit
 
   linux-jammy-rocm-py3_10-test:
-    permissions:
-      id-token: write
-      contents: read
-    name: linux-jammy-rocm-py3.10
+    name: linux-jammy-rocm-py3.10-mi355
     uses: ./.github/workflows/_rocm-test.yml
     needs:
       - linux-jammy-rocm-py3_10-build
@@ -340,19 +305,6 @@ jobs:
     if: ${{ needs.job-filter.outputs.jobs == '' || contains(needs.job-filter.outputs.jobs, ' inductor-build ') }}
     name: inductor-build
     uses: ./.github/workflows/_linux-build.yml
-    needs:
-      - get-label-type
-      - job-filter
-    with:
-      build-environment: linux-jammy-cuda12.8-py3.12-gcc11-sm80
-      docker-image-name: ci-image:pytorch-linux-jammy-cuda12.8-cudnn9-py3-gcc11-inductor-benchmarks
-      cuda-arch-list: '8.0'
-    secrets: inherit
-
-  inductor-build-cuda13:
-    if: ${{ needs.job-filter.outputs.jobs == '' || contains(needs.job-filter.outputs.jobs, ' inductor-build-cuda13 ') }}
-    name: inductor-build-cuda13
-    uses: ./.github/workflows/_linux-build.yml
     needs:
       - get-label-type
       - job-filter
@@ -362,22 +314,23 @@ jobs:
       cuda-arch-list: '8.0'
     secrets: inherit
 
-  # Test cross-compiled models with Windows libs extracted from wheel
-  cross-compile-linux-test:
-    if: ${{ needs.job-filter.outputs.jobs == '' || contains(needs.job-filter.outputs.jobs, ' cross-compile-linux-test ') }}
-    name: cross-compile-linux-test
+
+  # Test cross-compiled models with Windows libs extracted from wheel (CUDA 13.0)
+  cross-compile-linux-test-cuda13:
+    if: ${{ needs.job-filter.outputs.jobs == '' || contains(needs.job-filter.outputs.jobs, ' cross-compile-linux-test-cuda13 ') }}
+    name: cross-compile-linux-test-cuda13
     uses: ./.github/workflows/_linux-test.yml
     needs:
-      - linux-jammy-cuda12_8-py3_10-gcc11-build
+      - linux-jammy-cuda13_0-py3_10-gcc11-build
       - get-label-type
-      - win-vs2022-cuda12_8-py3-build
+      - win-vs2022-cuda13_0-py3-build
       - job-filter
     with:
-      build-environment: ${{ needs.linux-jammy-cuda12_8-py3_10-gcc11-build.outputs.build-environment }}
-      docker-image: ${{ needs.linux-jammy-cuda12_8-py3_10-gcc11-build.outputs.docker-image }}
+      build-environment: ${{ needs.linux-jammy-cuda13_0-py3_10-gcc11-build.outputs.build-environment }}
+      docker-image: ${{ needs.linux-jammy-cuda13_0-py3_10-gcc11-build.outputs.docker-image }}
       test-matrix: |
         { include: [
-          { config: "aoti_cross_compile_for_windows", shard: 1, num_shards: 1, runner: "${{ needs.get-label-type.outputs.label-type }}linux.g6.4xlarge.experimental.nvidia.gpu", win_torch_wheel_artifact: "win-vs2022-cuda12.8-py3" },
+          { config: "aoti_cross_compile_for_windows", shard: 1, num_shards: 1, runner: "${{ needs.get-label-type.outputs.label-type }}linux.g6.4xlarge.experimental.nvidia.gpu", win_torch_wheel_artifact: "win-vs2022-cuda13.0-py3" },
         ]}
       tests-to-include: ${{ github.event.inputs.tests-to-include || '' }}
     secrets: inherit
@@ -414,9 +367,9 @@ jobs:
       tests-to-include: ${{ github.event.inputs.tests-to-include || '' }}
     secrets: inherit
 
-  linux-jammy-py3-clang15-executorch-build:
-#    if: ${{ needs.job-filter.outputs.jobs == '' || contains(needs.job-filter.outputs.jobs, ' linux-jammy-py3-clang15-executorch ') }}
-    name: linux-jammy-py3-clang15-executorch
+  linux-jammy-py3-clang18-executorch-build:
+#    if: ${{ needs.job-filter.outputs.jobs == '' || contains(needs.job-filter.outputs.jobs, ' linux-jammy-py3-clang18-executorch ') }}
+    name: linux-jammy-py3-clang18-executorch
     uses: ./.github/workflows/_linux-build.yml
     needs:
       - get-label-type
@@ -424,25 +377,25 @@ jobs:
     if: false # Has been broken for a while
     with:
       runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
-      build-environment: linux-jammy-py3-clang15-executorch
-      docker-image-name: ci-image:pytorch-linux-jammy-py3-clang15-executorch
+      build-environment: linux-jammy-py3-clang18-executorch
+      docker-image-name: ci-image:pytorch-linux-jammy-py3-clang18-executorch
       test-matrix: |
         { include: [
           { config: "executorch", shard: 1, num_shards: 1, runner: "${{ needs.get-label-type.outputs.label-type }}linux.2xlarge" },
         ]}
     secrets: inherit
 
-  linux-jammy-py3-clang15-executorch-test:
-    if: ${{ needs.job-filter.outputs.jobs == '' || contains(needs.job-filter.outputs.jobs, ' linux-jammy-py3-clang15-executorch ') }}
-    name: linux-jammy-py3-clang15-executorch
+  linux-jammy-py3-clang18-executorch-test:
+    if: ${{ needs.job-filter.outputs.jobs == '' || contains(needs.job-filter.outputs.jobs, ' linux-jammy-py3-clang18-executorch ') }}
+    name: linux-jammy-py3-clang18-executorch
     uses: ./.github/workflows/_linux-test.yml
     needs:
-      - linux-jammy-py3-clang15-executorch-build
+      - linux-jammy-py3-clang18-executorch-build
       - job-filter
     with:
-      build-environment: ${{ needs.linux-jammy-py3-clang15-executorch-build.outputs.build-environment }}
-      docker-image: ${{ needs.linux-jammy-py3-clang15-executorch-build.outputs.docker-image }}
-      test-matrix: ${{ needs.linux-jammy-py3-clang15-executorch-build.outputs.test-matrix }}
+      build-environment: ${{ needs.linux-jammy-py3-clang18-executorch-build.outputs.build-environment }}
+      docker-image: ${{ needs.linux-jammy-py3-clang18-executorch-build.outputs.docker-image }}
+      test-matrix: ${{ needs.linux-jammy-py3-clang18-executorch-build.outputs.test-matrix }}
       tests-to-include: ${{ github.event.inputs.tests-to-include || '' }}
     secrets: inherit
 
@@ -457,5 +410,35 @@ jobs:
       runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
       runner: linux.r7i.2xlarge
       build-environment: linux-jammy-py3.10-gcc11-full-debug-build-only
-      docker-image-name: ci-image:pytorch-linux-jammy-py3.10-gcc11
+      docker-image-name: ci-image:pytorch-linux-jammy-py3.10-clang18
+    secrets: inherit
+
+  linux-jammy-aarch64-py3_10-gcc13-build:
+    name: linux-jammy-aarch64-py3.10
+    uses: ./.github/workflows/_linux-build.yml
+    needs: get-label-type
+    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
+      runner: linux.arm64.m7g.4xlarge
+      # Periodic AArch64 tests for SVE256 coverage
+      test-matrix: |
+        { include: [
+          { config: "default", shard: 1, num_shards: 5, runner: "${{ needs.get-label-type.outputs.label-type }}linux.arm64.m7g.4xlarge" },
+          { config: "default", shard: 2, num_shards: 5, runner: "${{ needs.get-label-type.outputs.label-type }}linux.arm64.m7g.4xlarge" },
+          { config: "default", shard: 3, num_shards: 5, runner: "${{ needs.get-label-type.outputs.label-type }}linux.arm64.m7g.4xlarge" },
+          { config: "default", shard: 4, num_shards: 5, runner: "${{ needs.get-label-type.outputs.label-type }}linux.arm64.m7g.4xlarge" },
+          { config: "default", shard: 5, num_shards: 5, runner: "${{ needs.get-label-type.outputs.label-type }}linux.arm64.m7g.4xlarge" },
+        ]}
+    secrets: inherit
+
+  linux-jammy-aarch64-py3_10-gcc13-test:
+    name: linux-jammy-aarch64-py3.10
+    uses: ./.github/workflows/_linux-test.yml
+    needs: linux-jammy-aarch64-py3_10-gcc13-build
+    with:
+      build-environment: ${{ needs.linux-jammy-aarch64-py3_10-gcc13-build.outputs.build-environment }}
+      docker-image: ${{ needs.linux-jammy-aarch64-py3_10-gcc13-build.outputs.docker-image }}
+      test-matrix: ${{ needs.linux-jammy-aarch64-py3_10-gcc13-build.outputs.test-matrix }}
     secrets: inherit
diff --git a/.github/workflows/unstable.yml b/.github/workflows/unstable.yml
index b5955127d9fb3..e346eb1d23386 100644
--- a/.github/workflows/unstable.yml
+++ b/.github/workflows/unstable.yml
@@ -46,7 +46,7 @@ jobs:
 
   get-label-type:
     name: get-label-type
-    uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@main
+    uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@release/2.12
     if: ${{ (github.event_name != 'schedule' || github.repository == 'pytorch/pytorch') && github.repository_owner == 'pytorch' }}
     with:
       triggering_actor: ${{ github.triggering_actor }}
diff --git a/.github/workflows/update-viablestrict.yml b/.github/workflows/update-viablestrict.yml
index 1b4af0f274913..84b6a073ffd8f 100644
--- a/.github/workflows/update-viablestrict.yml
+++ b/.github/workflows/update-viablestrict.yml
@@ -18,12 +18,12 @@ jobs:
     environment: ${{ (github.event_name == 'schedule') && 'mergebot' || '' }}
     steps:
       - name: Update viable/strict
-        uses: pytorch/test-infra/.github/actions/update-viablestrict@main
+        uses: pytorch/test-infra/.github/actions/update-viablestrict@release/2.12
         id: update_viablestrict
         with:
           repository: pytorch/pytorch
           stable-branch: viable/strict
-          requires: '[\"pull\", \"trunk\", \"lint\",  \"linux-aarch64\"]'
+          requires: '[\"pull\", \"trunk\", \"lint\"]'
           secret-bot-token: ${{ secrets.MERGEBOT_TOKEN }}
           clickhouse-url: ${{ secrets.CLICKHOUSE_URL }}
           clickhouse-username: ${{ secrets.CLICKHOUSE_VIABLESTRICT_USERNAME }}
diff --git a/.github/workflows/update_pytorch_labels.yml b/.github/workflows/update_pytorch_labels.yml
index a1b8c38141ae8..066d331ddbbac 100644
--- a/.github/workflows/update_pytorch_labels.yml
+++ b/.github/workflows/update_pytorch_labels.yml
@@ -17,7 +17,7 @@ jobs:
       contents: read
     steps:
       - name: Checkout PyTorch
-        uses: pytorch/pytorch/.github/actions/checkout-pytorch@main
+        uses: pytorch/pytorch/.github/actions/checkout-pytorch@release/2.12
         with:
           fetch-depth: 1
           submodules: false
diff --git a/.github/workflows/upload-test-stats-while-running.yml b/.github/workflows/upload-test-stats-while-running.yml
index 9aecaad0e068f..43811a3bf4bf2 100644
--- a/.github/workflows/upload-test-stats-while-running.yml
+++ b/.github/workflows/upload-test-stats-while-running.yml
@@ -15,15 +15,11 @@ jobs:
     name: Upload test stats while running
     runs-on: linux.2xlarge
     steps:
-      - name: Checkout PyTorch
-        uses: pytorch/pytorch/.github/actions/checkout-pytorch@main
+      - name: Setup Linux
+        uses: pytorch/pytorch/.github/actions/setup-linux@release/2.12
         with:
-          fetch-depth: 1
           submodules: false
 
-      - name: Setup Linux
-        uses: ./.github/actions/setup-linux
-
       - name: Install requirements
         run: |
           python3 -m pip install requests==2.32.2 boto3==1.35.42
diff --git a/.github/workflows/upload-test-stats.yml b/.github/workflows/upload-test-stats.yml
index 5dcbfe7fd65fa..21e49a2b0db17 100644
--- a/.github/workflows/upload-test-stats.yml
+++ b/.github/workflows/upload-test-stats.yml
@@ -20,6 +20,7 @@ on:
       - rocm-mi300
       - rocm-mi355
       - rocm-navi31
+      - rocm-nightly
       - inductor-micro-benchmark
       - inductor-micro-benchmark-x86
       - inductor-cu124
@@ -29,7 +30,6 @@ on:
       - inductor-perf-nightly-rocm-mi300
       - inductor-perf-nightly-rocm-mi355
       - mac-mps
-      - linux-aarch64
     types:
       - completed
 
@@ -66,7 +66,7 @@ jobs:
         run: echo "${TRIGGERING_WORKFLOW}"
 
       - name: Checkout PyTorch
-        uses: pytorch/pytorch/.github/actions/checkout-pytorch@main
+        uses: pytorch/pytorch/.github/actions/checkout-pytorch@release/2.12
 
       - name: Configure aws credentials
         uses: aws-actions/configure-aws-credentials@ececac1a45f3b08a01d2dd070d28d111c5fe6722 # v4.1.0
diff --git a/.github/workflows/upload-torch-dynamo-perf-stats.yml b/.github/workflows/upload-torch-dynamo-perf-stats.yml
index 07471619437a2..2d0ee2684d0de 100644
--- a/.github/workflows/upload-torch-dynamo-perf-stats.yml
+++ b/.github/workflows/upload-torch-dynamo-perf-stats.yml
@@ -32,7 +32,7 @@ jobs:
     name: Upload dynamo performance stats for ${{ github.event.workflow_run.id }}, attempt ${{ github.event.workflow_run.run_attempt }}
     steps:
       - name: Checkout PyTorch
-        uses: pytorch/pytorch/.github/actions/checkout-pytorch@main
+        uses: pytorch/pytorch/.github/actions/checkout-pytorch@release/2.12
         with:
           submodules: false
           fetch-depth: 1
diff --git a/.github/workflows/upload_test_stats_intermediate.yml b/.github/workflows/upload_test_stats_intermediate.yml
index 5702562006055..2247bba923384 100644
--- a/.github/workflows/upload_test_stats_intermediate.yml
+++ b/.github/workflows/upload_test_stats_intermediate.yml
@@ -17,7 +17,7 @@ jobs:
     environment: upload-stats
     steps:
       - name: Checkout PyTorch
-        uses: pytorch/pytorch/.github/actions/checkout-pytorch@main
+        uses: pytorch/pytorch/.github/actions/checkout-pytorch@release/2.12
         with:
           fetch-depth: 1
           submodules: false
diff --git a/.github/workflows/vllm-benchmark.yml b/.github/workflows/vllm-benchmark.yml
index 24c57119cae85..74f576ae1dadb 100644
--- a/.github/workflows/vllm-benchmark.yml
+++ b/.github/workflows/vllm-benchmark.yml
@@ -20,6 +20,12 @@ on:
         required: false
         type: string
         default: ''
+      compilation_config:
+        description: |
+          A JSON string for vLLM --compilation-config, applied to all benchmark tests, e.g. {"cudagraph_mode": "PIECEWISE", "use_inductor_graph_partition": true}
+        required: false
+        type: string
+        default: ''
       runners:
         description: |
           A comma-separated list of runners from .github/scripts/generate_vllm_benchmark_matrix.py to run the benchmark (optional, default to run everything)
@@ -41,10 +47,10 @@ jobs:
     outputs:
       benchmark_matrix: ${{ steps.set-parameters.outputs.benchmark_matrix }}
       docker_image: ${{ steps.calculate-docker-image.outputs.docker-image }}
-      torch_cuda_arch_list: '8.0 8.9 9.0 10.0'
-      build_environment: linux-jammy-cuda12.9-py3.12-gcc11
+      torch_cuda_arch_list: '8.0 8.9 9.0 10.0 12.0'
+      build_environment: linux-jammy-cuda13.0-py3.12-gcc11
     steps:
-      - uses: pytorch/test-infra/.github/actions/setup-uv@main
+      - uses: pytorch/test-infra/.github/actions/setup-uv@release/2.12
         with:
           python-version: "3.12"
           activate-environment: "true"
@@ -82,10 +88,10 @@ jobs:
 
       - name: Calculate docker image
         id: calculate-docker-image
-        uses: pytorch/test-infra/.github/actions/calculate-docker-image@main
+        uses: pytorch/test-infra/.github/actions/calculate-docker-image@release/2.12
         with:
           working-directory: pytorch/pytorch
-          docker-image-name: ci-image:pytorch-linux-jammy-cuda12.9-cudnn9-py3.12-gcc11-vllm
+          docker-image-name: ci-image:pytorch-linux-jammy-cuda13.0-cudnn9-py3.12-gcc11-vllm
 
   build:
     name: Build PyTorch and vLLM
@@ -120,4 +126,6 @@ jobs:
       pytorch_branch: ${{ inputs.pytorch_branch }}
       pytorch_commit: ${{ inputs.pytorch_commit }}
       models: ${{ matrix.models }}
+      compilation_config: ${{ inputs.compilation_config }}
+      is_nightly: ${{ github.event_name == 'schedule' }}
     secrets: inherit
diff --git a/.github/workflows/vllm.yml b/.github/workflows/vllm.yml
index df657d057bbd0..f65e73042c245 100644
--- a/.github/workflows/vllm.yml
+++ b/.github/workflows/vllm.yml
@@ -16,6 +16,7 @@ concurrency:
 permissions:
   id-token: write
   contents: read
+  actions: read
 
 jobs:
   build:
@@ -27,9 +28,9 @@ jobs:
       allow-reuse-old-whl: false
       build-additional-packages: "vision audio"
       build-external-packages: "vllm"
-      build-environment: linux-jammy-cuda12.9-py3.12-gcc11
-      docker-image-name: ci-image:pytorch-linux-jammy-cuda12.9-cudnn9-py3.12-gcc11-vllm
-      cuda-arch-list: '8.0 8.9 9.0 10.0'
+      build-environment: linux-jammy-cuda13.0-py3.12-gcc11
+      docker-image-name: ci-image:pytorch-linux-jammy-cuda13.0-cudnn9-py3.12-gcc11-vllm
+      cuda-arch-list: '8.0 8.9 9.0 10.0 12.0'
       runner: linux.24xlarge.memory
       test-matrix: |
         { include: [
@@ -58,7 +59,7 @@ jobs:
     uses: ./.github/workflows/_linux-test.yml
     needs: build
     with:
-      build-environment: linux-jammy-cuda12.9-py3.12-gcc11
+      build-environment: linux-jammy-cuda13.0-py3.12-gcc11
       docker-image: ${{ needs.build.outputs.docker-image }}
       test-matrix: ${{ needs.build.outputs.test-matrix }}
     secrets: inherit
diff --git a/.github/workflows/weekly.yml b/.github/workflows/weekly.yml
index 7bed6c785d4db..c5a3311857327 100644
--- a/.github/workflows/weekly.yml
+++ b/.github/workflows/weekly.yml
@@ -22,7 +22,7 @@ jobs:
           fetch-depth: 0
       - name: update-xla-commit-hash
         continue-on-error: true
-        uses: pytorch/test-infra/.github/actions/update-commit-hash@main
+        uses: pytorch/test-infra/.github/actions/update-commit-hash@release/2.12
         with:
           repo-name: xla
           branch: master
diff --git a/.github/workflows/win-arm64-build-test.yml b/.github/workflows/win-arm64-build-test.yml
index 95b4e2f027f60..5c31628b96871 100644
--- a/.github/workflows/win-arm64-build-test.yml
+++ b/.github/workflows/win-arm64-build-test.yml
@@ -31,7 +31,7 @@ jobs:
     steps:
       - name: configure aws credentials
         id: aws_creds
-        uses: aws-actions/configure-aws-credentials@v4
+        uses: aws-actions/configure-aws-credentials@7474bc4690e29a8392af63c5b98e7449536d5c3a # v4
         with:
           role-to-assume: arn:aws:iam::308535385114:role/gha_workflow_sscache
           aws-region: us-east-1
diff --git a/.github/workflows/xpu.yml b/.github/workflows/xpu.yml
index 84b8aa3cd91d4..32853bb956425 100644
--- a/.github/workflows/xpu.yml
+++ b/.github/workflows/xpu.yml
@@ -14,12 +14,17 @@ 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' }}
   cancel-in-progress: true
 
+permissions:
+  id-token: write
+  contents: read
+  actions: read
+
 jobs:
 
   get-label-type:
     if: github.repository_owner == 'pytorch'
     name: get-label-type
-    uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@main
+    uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@release/2.12
     with:
       triggering_actor: ${{ github.triggering_actor }}
       issue_owner: ${{ github.event.pull_request.user.login || github.event.issue.user.login }}
@@ -78,15 +83,38 @@ jobs:
     name: linux-noble-xpu-n-py3.10
     uses: ./.github/workflows/_xpu-test.yml
     needs: linux-noble-xpu-n-py3_10-build
-    permissions:
-      id-token: write
-      contents: read
     with:
       build-environment: ${{ needs.linux-noble-xpu-n-py3_10-build.outputs.build-environment }}
       docker-image: ${{ needs.linux-noble-xpu-n-py3_10-build.outputs.docker-image }}
       test-matrix: ${{ needs.linux-noble-xpu-n-py3_10-build.outputs.test-matrix }}
     secrets: inherit
 
+  linux-noble-xpu-n-py3_10-client-build:
+    name: linux-noble-xpu-n-py3.10-client
+    uses: ./.github/workflows/_linux-build.yml
+    needs: get-label-type
+    with:
+      sync-tag: linux-xpu-n-client-build
+      runner_prefix: ${{ needs.get-label-type.outputs.label-type }}
+      build-environment: linux-noble-xpu-n-py3.10-client
+      docker-image-name: ci-image:pytorch-linux-noble-xpu-n-py3-client
+      runner: linux.c7i.12xlarge
+      test-matrix: |
+        { include: [
+          { config: "smoke_xpu", shard: 1, num_shards: 1, runner: "linux.client.xpu" },
+        ]}
+    secrets: inherit
+
+  linux-noble-xpu-n-py3_10-client-test:
+    name: linux-noble-xpu-n-py3.10-client
+    uses: ./.github/workflows/_xpu-test.yml
+    needs: linux-noble-xpu-n-py3_10-client-build
+    with:
+      build-environment: ${{ needs.linux-noble-xpu-n-py3_10-client-build.outputs.build-environment }}
+      docker-image: ${{ needs.linux-noble-xpu-n-py3_10-client-build.outputs.docker-image }}
+      test-matrix: ${{ needs.linux-noble-xpu-n-py3_10-client-build.outputs.test-matrix }}
+    secrets: inherit
+
   windows-xpu-n-1-build:
     if: github.repository_owner == 'pytorch'
     name: win-vs2022-xpu-n-1-py3
diff --git a/.gitignore b/.gitignore
index 095b9cf8f9456..5105977ae52c0 100644
--- a/.gitignore
+++ b/.gitignore
@@ -8,8 +8,7 @@
 
 ## PyTorch
 
-.coverage
-coverage.xml
+agent_space/
 .dmypy.json
 .gradle
 .hypothesis
@@ -40,7 +39,6 @@ docs/build/
 docs/cpp/src
 docs/src/**/*
 docs/cpp/build
-docs/cpp/source/api
 docs/cpp/source/html/
 docs/cpp/source/latex/
 docs/source/compile/generated/
@@ -52,7 +50,6 @@ usage_log*
 test-reports/
 test/*.bak
 test/**/*.bak
-test/.coverage
 test/.hypothesis/
 test/cpp/api/mnist
 test/custom_operator/model.pt
@@ -127,7 +124,8 @@ 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/*
+torch/_inductor/kernel/vendored_templates/cutedsl/kernels/*
+torch/_inductor/kernel/vendored_templates/cutedsl_grouped_gemm.py
 test/inductor/test_tlx*
 minifier_launcher.py
 aten/src/ATen/native/transformers/hip/flash_attn/ck/fav_v3/asm_fmha_v3_bwd_configs.hpp
@@ -247,8 +245,8 @@ docs/source/scripts/lr_scheduler_images/
 
 # build, distribute, and bins (+ python proto bindings)
 build/
-# Allow tools/build/ for build support.
-!tools/build/
+# Allow tools/build/bazel for build support.
+!tools/build/bazel
 build_host_protoc
 build_android
 build_ios
@@ -366,7 +364,13 @@ xla/
 pr.diff
 
 # coverage files
-*/**/.coverage.*
+.coverage.*
+.coverage
+coverage.xml
+*.profdata
+*.profraw
+coverage-html
+coverage.lcov
 
 # buck generated files
 .buckd/
diff --git a/.lintrunner.toml b/.lintrunner.toml
index 43874512fc85e..cd4f0a3dce756 100644
--- a/.lintrunner.toml
+++ b/.lintrunner.toml
@@ -16,6 +16,8 @@ exclude_patterns = [
     'functorch/docs/source/tutorials/**',
     'torch/_inductor/fx_passes/serialized_patterns/**',
     'torch/_inductor/autoheuristic/artifacts/**',
+    'torch/_inductor/kernel/vendored_templates/cutedsl/kernels/**',
+    'torch/_inductor/kernel/vendored_templates/cutedsl/dense_blockscaled_gemm_persistent.py',
     'scripts/**',
     'test/generated_type_hints_smoketest.py',
     'test/test_torchfuzz_repros.py',
@@ -46,6 +48,8 @@ code = 'CLANGFORMAT'
 include_patterns = [
     'aten/src/ATen/*.h',
     'aten/src/ATen/cpu/vec/**/*.h',
+    'aten/src/ATen/accelerator/**/*.h',
+    'aten/src/ATen/accelerator/**/*.cpp',
     'aten/src/ATen/mps/**/*.mm',
     'aten/src/ATen/mps/**/*.h',
     'aten/src/ATen/xpu/**/*.h',
@@ -117,8 +121,13 @@ include_patterns = [
     'torchgen/**/*.pyi',
     'functorch/**/*.py',
     'functorch/**/*.pyi',
+    'tools/**/*.py',
+    'tools/**/*.pyi',
+]
+exclude_patterns = [
+    'torch/_inductor/kernel/vendored_templates/cutedsl/kernels/**',
+    'torch/_inductor/kernel/vendored_templates/cutedsl/dense_blockscaled_gemm_persistent.py',
 ]
-exclude_patterns = []
 command = [
     'uv',
     'run',
@@ -137,6 +146,8 @@ include_patterns = [
     'aten/src/ATen/cuda/*.cpp',
     'aten/src/ATen/cpu/*.h',
     'aten/src/ATen/cpu/*.cpp',
+    'aten/src/ATen/accelerator/**/*.h',
+    'aten/src/ATen/accelerator/**/*.cpp',
     'aten/src/ATen/core/*.h',
     'aten/src/ATen/core/*.cpp',
     'aten/src/ATen/cudnn/*.h',
@@ -180,6 +191,11 @@ exclude_patterns = [
     'aten/src/ATen/cpu/FlushDenormal.cpp',
     'aten/src/ATen/cpu/vml.h',
     'aten/src/ATen/CPUFixedAllocator.h',
+    # interned_strings.h includes the generated aten_interned_strings.h which
+    # #errors under TORCH_ASSERT_ONLY_METHOD_OPERATORS. The .cpp file #undefs
+    # the macro before including the header, but clang-tidy lints the .h
+    # standalone using the .cpp's compile flags, so the #error always fires.
+    'aten/src/ATen/core/interned_strings.h',
     'aten/src/ATen/Parallel*.h',
     'c10/xpu/**/*.h',
     'c10/xpu/**/*.cpp',
@@ -505,6 +521,44 @@ command = [
     '@{{PATHSFILE}}'
 ]
 
+[[linter]]
+code = 'RAWTHROW'
+include_patterns = [
+    '**/*.cpp',
+    '**/*.h',
+]
+exclude_patterns = [
+    # Vendored code (do not modify)
+    'c10/util/flat_hash_map.h',
+    'c10/util/order_preserving_flat_hash_map.h',
+    'aten/src/ATen/native/quantized/cpu/qnnpack/**',
+    'aten/src/ATen/native/transformers/cuda/mem_eff_attention/**',
+    # Standalone code that cannot depend on c10
+    'torch/csrc/inductor/aoti_runtime/**',
+    # Test files use EXPECT_THROW which is a gtest macro
+    'test/cpp/**/*.cpp',
+    # JIT frontend uses throw(ErrorReport(...)) pervasively
+    'torch/csrc/jit/frontend/schema_type_parser.cpp',
+]
+command = [
+    'python3',
+    'tools/linter/adapters/grep_linter.py',
+    '--pattern=\bthrow\b',
+    '--allowlist-pattern=@allow-raw-throw',
+    '--linter-name=RAWTHROW',
+    '--error-name=raw throw statement',
+    """--error-description=\
+        Do not use raw `throw` in C++ code. Use TORCH_CHECK, \
+        TORCH_CHECK_WITH, C10_THROW_ERROR, or other error-reporting macros \
+        instead. See c10/util/Exception.h for available macros. If this is a \
+        re-throw (throw;), restructure the code to avoid try/catch/throw or \
+        add the file to the RAWTHROW exclude list in .lintrunner.toml with \
+        justification.\
+    """,
+    '--',
+    '@{{PATHSFILE}}'
+]
+
 [[linter]]
 code = 'INCLUDE'
 include_patterns = [
@@ -1108,6 +1162,8 @@ exclude_patterns = [
     'torch/_vendor/**',
     'torch/_inductor/fx_passes/serialized_patterns/**',
     'torch/_inductor/autoheuristic/artifacts/**',
+    'torch/_inductor/kernel/vendored_templates/cutedsl/kernels/**',
+    'torch/_inductor/kernel/vendored_templates/cutedsl/dense_blockscaled_gemm_persistent.py',
     'torch/utils/model_dump/preact.mjs',
 ]
 is_formatter = true
@@ -1141,6 +1197,8 @@ exclude_patterns = [
     'torch/_vendor/**',
     'torch/_inductor/fx_passes/serialized_patterns/**',
     'torch/_inductor/autoheuristic/artifacts/**',
+    # Needs Python 3.12+
+    'torch/testing/_internal/py312_intrinsics.py',
     # These files are all grandfathered in, feel free to remove from this list
     # as necessary
     'test/quantization/__init__.py',
@@ -1445,27 +1503,6 @@ command = [
     '@{{PATHSFILE}}'
 ]
 
-[[linter]]
-code = 'BAZEL_LINTER'
-include_patterns = ['WORKSPACE']
-command = [
-    'python3',
-    'tools/linter/adapters/bazel_linter.py',
-    '--binary=.lintbin/bazel',
-    '--',
-    '@{{PATHSFILE}}'
-]
-init_command = [
-    'python3',
-    'tools/linter/adapters/s3_init.py',
-    '--config-json=tools/linter/adapters/s3_init_config.json',
-    '--linter=bazel',
-    '--dry-run={{DRYRUN}}',
-    '--output-dir=.lintbin',
-    '--output-name=bazel',
-]
-is_formatter = true
-
 [[linter]]
 code = 'LINTRUNNER_VERSION'
 include_patterns = ['**']
@@ -1491,6 +1528,8 @@ exclude_patterns = [
     'functorch/docs/**',
     'torch/_inductor/fx_passes/serialized_patterns/**',
     'torch/_inductor/autoheuristic/artifacts/**',
+    'torch/_inductor/kernel/vendored_templates/cutedsl/kernels/**',
+    'torch/_inductor/kernel/vendored_templates/cutedsl/dense_blockscaled_gemm_persistent.py',
     'test/dynamo/cpython/**',
     'test/test_torchfuzz_repros.py',
     'scripts/**',
@@ -1670,6 +1709,10 @@ command = [
 include_patterns = [
    'torch/_inductor/**/*.py'
 ]
+exclude_patterns = [
+    'torch/_inductor/kernel/vendored_templates/cutedsl/kernels/**',
+    'torch/_inductor/kernel/vendored_templates/cutedsl/dense_blockscaled_gemm_persistent.py',
+]
 is_formatter = false
 
 # `import_linter` reports on importing disallowed third party libraries.
diff --git a/.spin/cmds.py b/.spin/cmds.py
index f31acfa977424..f1aceb77edd25 100644
--- a/.spin/cmds.py
+++ b/.spin/cmds.py
@@ -8,6 +8,11 @@
 import spin
 
 
+CWD = Path(__file__).absolute().parent.parent
+sys.path.insert(0, str(CWD))  # this only affects the current process
+from tools.clean import clean as _clean
+
+
 def file_digest(file, algorithm: str):
     try:
         return hashlib.file_digest(file, algorithm)
@@ -133,7 +138,6 @@ def regenerate_clangtidy_files():
 #: These linters are expected to need less than 3s cpu time total
 VERY_FAST_LINTERS = {
     "ATEN_CPU_GPU_AGNOSTIC",
-    "BAZEL_LINTER",
     "C10_NODISCARD",
     "C10_UNUSED",
     "CALL_ONCE",
@@ -441,6 +445,12 @@ def quickfix(ctx, *, lintrunner_args, **kwargs):
     ctx.invoke(quicklint, apply_patches=True)
 
 
+@click.command()
+def clean():
+    """Clean, that is remove all files in .gitignore except in the NOT-CLEAN-FILES section."""
+    _clean()
+
+
 @click.command()
 def regenerate_github_workflows():
     """Regenerate GitHub workflows from templates."""
diff --git a/AGENTS.md b/AGENTS.md
deleted file mode 100644
index 718217d3e663d..0000000000000
--- a/AGENTS.md
+++ /dev/null
@@ -1,18 +0,0 @@
-- This is the only AGENTS.md, there are no recursive AGENTS.md
-- When you are working on a bug, first create a standalone file that
-  reproduces the bug and verify it fails in the expected way.  Use this to
-  test if your changes work.  Once the change is passing, find an appropriate
-  test file to add the test to and make sure to follow local conventions on
-  the test file.
-- If you are running the real test suite, DO NOT run the entire test suite.
-  Instead run only a single test case, e.g., 'python test/test_torch.py TestTorch.test_dir'
-- Do NOT run setup.py, you do not have a working build environment
-- Do NOT run pre-commit, it is not setup
-- To run lint, run 'lintrunner -a' (which will autoapply changes)
-- Do NOT attempt to install dependencies, you do not have Internet access
-- Do NOT create summary files unless explicitly asked
-- When you are ready to make a PR, do exactly these steps:
-  - git stash -u
-  - git reset --hard $(cat /tmp/orig_work.txt) # NB: reset to the LOCAL branch, do NOT fetch
-  - git stash pop
-  - Resolve conflicts if necessary
diff --git a/AGENTS.md b/AGENTS.md
new file mode 120000
index 0000000000000..681311eb9cf45
--- /dev/null
+++ b/AGENTS.md
@@ -0,0 +1 @@
+CLAUDE.md
\ No newline at end of file
diff --git a/BUILD.bazel b/BUILD.bazel
deleted file mode 100644
index 4737a2a0c486c..0000000000000
--- a/BUILD.bazel
+++ /dev/null
@@ -1,1104 +0,0 @@
-load("@bazel_skylib//lib:paths.bzl", "paths")
-load("@com_github_google_flatbuffers//:build_defs.bzl", "flatbuffer_cc_library")
-load("@pybind11_bazel//:build_defs.bzl", "pybind_extension")
-load("@rules_cc//cc:defs.bzl", "cc_binary", "cc_library", "cc_test")
-load("@rules_python//python:defs.bzl", "py_library", "py_test")
-load("@pytorch//third_party:substitution.bzl", "header_template_rule", "template_rule")
-load("@pytorch//:tools/bazel.bzl", "rules")
-load("@pytorch//tools/rules:cu.bzl", "cu_library")
-load("@pytorch//tools/config:defs.bzl", "if_cuda")
-load("@pytorch//:aten.bzl", "generate_aten", "intern_build_aten_ops")
-load(":build.bzl", "GENERATED_AUTOGRAD_CPP", "GENERATED_AUTOGRAD_PYTHON", "define_targets")
-load(":build_variables.bzl", "jit_core_sources", "lazy_tensor_ts_sources", "libtorch_core_sources", "libtorch_cuda_sources", "libtorch_distributed_sources", "libtorch_extra_sources", "libtorch_python_core_sources", "torch_cpp_srcs", "libtorch_python_cuda_sources", "libtorch_python_distributed_sources")
-load(":ufunc_defs.bzl", "aten_ufunc_generated_cpu_kernel_sources", "aten_ufunc_generated_cpu_sources", "aten_ufunc_generated_cuda_sources")
-load("//:tools/bazel.bzl", "rules")
-
-# Export files for use by torch/headeronly (where version.h generation now lives)
-exports_files(["version.txt"])
-
-define_targets(rules = rules)
-
-COMMON_COPTS = [
-    "-DHAVE_MALLOC_USABLE_SIZE=1",
-    "-DHAVE_MMAP=1",
-    "-DHAVE_SHM_OPEN=1",
-    "-DHAVE_SHM_UNLINK=1",
-    "-D_FILE_OFFSET_BITS=64",
-    "-DUSE_FBGEMM",
-    "-DUSE_DISTRIBUTED",
-    "-DAT_PER_OPERATOR_HEADERS",
-    "-DATEN_THREADING=NATIVE",
-    "-DNO_CUDNN_DESTROY_HANDLE",
-] + if_cuda([
-    "-DUSE_CUDA",
-    "-DUSE_CUDNN",
-    # TODO: This should be passed only when building for CUDA-11.5 or newer
-    # use cub in a safe manner, see:
-    # https://github.com/pytorch/pytorch/pull/55292
-    "-DCUB_WRAPPED_NAMESPACE=at_cuda_detail",
-])
-
-aten_generation_srcs = ["aten/src/ATen/native/native_functions.yaml"] + ["aten/src/ATen/native/tags.yaml"] + glob(["aten/src/ATen/templates/**"])
-
-generated_cpu_cpp = [
-    "aten/src/ATen/RegisterBackendSelect.cpp",
-    "aten/src/ATen/RegisterCPU_0.cpp",
-    "aten/src/ATen/RegisterCPU_1.cpp",
-    "aten/src/ATen/RegisterCPU_2.cpp",
-    "aten/src/ATen/RegisterCPU_3.cpp",
-    "aten/src/ATen/RegisterFunctionalization_0.cpp",
-    "aten/src/ATen/RegisterFunctionalization_1.cpp",
-    "aten/src/ATen/RegisterFunctionalization_2.cpp",
-    "aten/src/ATen/RegisterFunctionalization_3.cpp",
-    # "aten/src/ATen/RegisterFunctionalizationEverything.cpp",
-    "aten/src/ATen/RegisterMkldnnCPU_0.cpp",
-    "aten/src/ATen/RegisterNestedTensorCPU_0.cpp",
-    "aten/src/ATen/RegisterQuantizedCPU_0.cpp",
-    "aten/src/ATen/RegisterSparseCPU_0.cpp",
-    "aten/src/ATen/RegisterSparseCsrCPU_0.cpp",
-    "aten/src/ATen/RegisterZeroTensor_0.cpp",
-    "aten/src/ATen/RegisterCompositeImplicitAutograd_0.cpp",
-    "aten/src/ATen/RegisterCompositeImplicitAutogradNestedTensor_0.cpp",
-    "aten/src/ATen/RegisterCompositeExplicitAutograd_0.cpp",
-    "aten/src/ATen/RegisterCompositeExplicitAutogradNonFunctional_0.cpp",
-    "aten/src/ATen/RegisterMeta_0.cpp",
-    "aten/src/ATen/RegisterSparseMeta_0.cpp",
-    "aten/src/ATen/RegisterQuantizedMeta_0.cpp",
-    "aten/src/ATen/RegisterNestedTensorMeta_0.cpp",
-    "aten/src/ATen/RegisterSchema.cpp",
-    "aten/src/ATen/CPUFunctions.h",
-    "aten/src/ATen/CPUFunctions_inl.h",
-    "aten/src/ATen/CompositeExplicitAutogradFunctions.h",
-    "aten/src/ATen/CompositeExplicitAutogradFunctions_inl.h",
-    "aten/src/ATen/CompositeExplicitAutogradNonFunctionalFunctions.h",
-    "aten/src/ATen/CompositeExplicitAutogradNonFunctionalFunctions_inl.h",
-    "aten/src/ATen/CompositeImplicitAutogradFunctions.h",
-    "aten/src/ATen/CompositeImplicitAutogradFunctions_inl.h",
-    "aten/src/ATen/CompositeImplicitAutogradNestedTensorFunctions.h",
-    "aten/src/ATen/CompositeImplicitAutogradNestedTensorFunctions_inl.h",
-    "aten/src/ATen/CompositeViewCopyKernels.cpp",
-    "aten/src/ATen/FunctionalInverses.h",
-    "aten/src/ATen/Functions.h",
-    "aten/src/ATen/Functions.cpp",
-    "aten/src/ATen/RedispatchFunctions.h",
-    "aten/src/ATen/Operators.h",
-    "aten/src/ATen/Operators_0.cpp",
-    "aten/src/ATen/Operators_1.cpp",
-    "aten/src/ATen/Operators_2.cpp",
-    "aten/src/ATen/Operators_3.cpp",
-    "aten/src/ATen/Operators_4.cpp",
-    "aten/src/ATen/NativeFunctions.h",
-    "aten/src/ATen/MetaFunctions.h",
-    "aten/src/ATen/MetaFunctions_inl.h",
-    "aten/src/ATen/MethodOperators.h",
-    "aten/src/ATen/NativeMetaFunctions.h",
-    "aten/src/ATen/RegistrationDeclarations.h",
-    "aten/src/ATen/VmapGeneratedPlumbing.h",
-    "aten/src/ATen/ViewMetaClasses.h",
-    "aten/src/ATen/ViewMetaClasses.cpp",
-    "aten/src/ATen/core/aten_interned_strings.h",
-    "aten/src/ATen/core/enum_tag.h",
-    "aten/src/ATen/core/TensorBody.h",
-    "aten/src/ATen/core/TensorMethods.cpp",
-    "aten/src/ATen/core/ATenOpList.cpp",
-]
-
-generated_cuda_cpp = [
-    "aten/src/ATen/CUDAFunctions.h",
-    "aten/src/ATen/CUDAFunctions_inl.h",
-    "aten/src/ATen/RegisterCUDA_0.cpp",
-    "aten/src/ATen/RegisterNestedTensorCUDA_0.cpp",
-    "aten/src/ATen/RegisterQuantizedCUDA_0.cpp",
-    "aten/src/ATen/RegisterSparseCUDA_0.cpp",
-    "aten/src/ATen/RegisterSparseCsrCUDA_0.cpp",
-]
-
-generate_aten(
-    name = "generated_aten_cpp",
-    srcs = aten_generation_srcs,
-    outs = (
-        generated_cpu_cpp +
-        generated_cuda_cpp +
-        aten_ufunc_generated_cpu_sources("aten/src/ATen/{}") +
-        aten_ufunc_generated_cpu_kernel_sources("aten/src/ATen/{}") +
-        aten_ufunc_generated_cuda_sources("aten/src/ATen/{}") + [
-            "aten/src/ATen/Declarations.yaml",
-        ]
-    ),
-    generator = "//torchgen:gen",
-)
-
-filegroup(
-    name = "cpp_generated_code",
-    srcs = GENERATED_AUTOGRAD_CPP,
-    data = [":generate-code"],
-)
-
-# ATen
-filegroup(
-    name = "aten_base_cpp",
-    srcs = glob([
-        "aten/src/ATen/*.cpp",
-        "aten/src/ATen/functorch/*.cpp",
-        "aten/src/ATen/detail/*.cpp",
-        "aten/src/ATen/cpu/*.cpp",
-    ]),
-)
-
-filegroup(
-    name = "ATen_CORE_SRCS",
-    srcs = glob(
-        [
-            "aten/src/ATen/core/**/*.cpp",
-        ],
-        exclude = [
-            "aten/src/ATen/core/**/*_test.cpp",
-        ],
-    ),
-)
-
-filegroup(
-    name = "aten_native_cpp",
-    srcs = glob(["aten/src/ATen/native/*.cpp"]),
-)
-
-filegroup(
-    name = "aten_native_sparse_cpp",
-    srcs = glob(["aten/src/ATen/native/sparse/*.cpp"]),
-)
-
-filegroup(
-    name = "aten_native_nested_cpp",
-    srcs = glob(["aten/src/ATen/native/nested/*.cpp"]),
-)
-
-filegroup(
-    name = "aten_native_quantized_cpp",
-    srcs = glob(
-        [
-            "aten/src/ATen/native/quantized/*.cpp",
-            "aten/src/ATen/native/quantized/cpu/*.cpp",
-        ],
-    ),
-)
-
-filegroup(
-    name = "aten_native_transformers_cpp",
-    srcs = glob(["aten/src/ATen/native/transformers/*.cpp"]),
-)
-
-filegroup(
-    name = "aten_native_mkl_cpp",
-    srcs = glob([
-        "aten/src/ATen/native/mkl/*.cpp",
-        "aten/src/ATen/mkl/*.cpp",
-    ]),
-)
-
-filegroup(
-    name = "aten_native_mkldnn_cpp",
-    srcs = glob(["aten/src/ATen/native/mkldnn/*.cpp"]),
-)
-
-filegroup(
-    name = "aten_native_xnnpack",
-    srcs = glob(["aten/src/ATen/native/xnnpack/*.cpp"]),
-)
-
-filegroup(
-    name = "aten_base_vulkan",
-    srcs = glob(["aten/src/ATen/vulkan/*.cpp"]),
-)
-
-filegroup(
-    name = "aten_base_metal",
-    srcs = glob(["aten/src/ATen/metal/*.cpp"]),
-)
-
-filegroup(
-    name = "ATen_QUANTIZED_SRCS",
-    srcs = glob(
-        [
-            "aten/src/ATen/quantized/**/*.cpp",
-        ],
-        exclude = [
-            "aten/src/ATen/quantized/**/*_test.cpp",
-        ],
-    ),
-)
-
-filegroup(
-    name = "aten_cuda_cpp_srcs",
-    srcs = glob(
-        [
-            "aten/src/ATen/cuda/*.cpp",
-            "aten/src/ATen/cuda/detail/*.cpp",
-            "aten/src/ATen/cuda/tunable/*.cpp",
-            "aten/src/ATen/cudnn/*.cpp",
-            "aten/src/ATen/native/cuda/*.cpp",
-            "aten/src/ATen/native/cuda/linalg/*.cpp",
-            "aten/src/ATen/native/cudnn/*.cpp",
-            "aten/src/ATen/native/miopen/*.cpp",
-            "aten/src/ATen/native/nested/cuda/*.cpp",
-            "aten/src/ATen/native/quantized/cuda/*.cpp",
-            "aten/src/ATen/native/quantized/cudnn/*.cpp",
-            "aten/src/ATen/native/sparse/cuda/*.cpp",
-            "aten/src/ATen/native/transformers/cuda/*.cpp",
-        ],
-    ),
-)
-
-filegroup(
-    name = "aten_cu_srcs",
-    srcs = glob([
-        "aten/src/ATen/cuda/*.cu",
-        "aten/src/ATen/cuda/detail/*.cu",
-        "aten/src/ATen/native/cuda/*.cu",
-        "aten/src/ATen/native/nested/cuda/*.cu",
-        "aten/src/ATen/native/quantized/cuda/*.cu",
-        "aten/src/ATen/native/sparse/cuda/*.cu",
-        "aten/src/ATen/native/transformers/cuda/*.cu",
-    ]) + aten_ufunc_generated_cuda_sources("aten/src/ATen/{}"),
-    # It's a bit puzzling to me why it's not necessary to declare the
-    # target that generates these sources...
-)
-
-# TODO: Enable support for KleidiAI bazel build
-header_template_rule(
-    name = "aten_src_ATen_config",
-    src = "aten/src/ATen/Config.h.in",
-    out = "aten/src/ATen/Config.h",
-    include = "aten/src",
-    substitutions = {
-        "@AT_MKLDNN_ENABLED@": "1",
-        "@AT_MKLDNN_ACL_ENABLED@": "0",
-        "@AT_MKL_ENABLED@": "1",
-        "@AT_MKL_SEQUENTIAL@": "0",
-        "@AT_POCKETFFT_ENABLED@": "0",
-        "@AT_NNPACK_ENABLED@": "0",
-        "@CAFFE2_STATIC_LINK_CUDA_INT@": "0",
-        "@AT_BUILD_WITH_BLAS@": "1",
-        "@AT_BUILD_WITH_LAPACK@": "1",
-        "@AT_PARALLEL_OPENMP@": "0",
-        "@AT_PARALLEL_NATIVE@": "1",
-        "@AT_BLAS_F2C@": "0",
-        "@AT_BLAS_USE_CBLAS_DOT@": "1",
-        "@AT_KLEIDIAI_ENABLED@": "0",
-        "@AT_USE_EIGEN_SPARSE@": "0",
-    },
-)
-
-header_template_rule(
-    name = "aten_src_ATen_cuda_config",
-    src = "aten/src/ATen/cuda/CUDAConfig.h.in",
-    out = "aten/src/ATen/cuda/CUDAConfig.h",
-    include = "aten/src",
-    substitutions = {
-        "@AT_CUDNN_ENABLED@": "1",
-        "@AT_CUSPARSELT_ENABLED@": "0",
-        "@AT_HIPSPARSELT_ENABLED@": "0",
-        "@AT_ROCM_ENABLED@": "0",
-        "@AT_MAGMA_ENABLED@": "0",
-        "@NVCC_FLAGS_EXTRA@": "",
-    },
-)
-
-cc_library(
-    name = "aten_headers",
-    hdrs = [
-        "torch/csrc/Export.h",
-        "torch/csrc/jit/frontend/function_schema_parser.h",
-    ] + glob(
-        [
-            "aten/src/**/*.h",
-            "aten/src/**/*.hpp",
-            "aten/src/ATen/cuda/**/*.cuh",
-            "aten/src/ATen/native/**/*.cuh",
-            "aten/src/THC/*.cuh",
-        ],
-    ) + [
-        ":aten_src_ATen_config",
-        ":generated_aten_cpp",
-    ],
-    includes = [
-        "aten/src",
-    ],
-    deps = [
-        "//c10",
-    ],
-)
-
-ATEN_COPTS = COMMON_COPTS + [
-    "-DCAFFE2_BUILD_MAIN_LIBS",
-    "-DHAVE_AVX_CPU_DEFINITION",
-    "-DHAVE_AVX2_CPU_DEFINITION",
-    "-fvisibility-inlines-hidden",
-    "-fno-math-errno",
-    "-fno-trapping-math",
-]
-
-intern_build_aten_ops(
-    copts = ATEN_COPTS,
-    extra_impls = aten_ufunc_generated_cpu_kernel_sources("aten/src/ATen/{}"),
-    deps = [
-        ":aten_headers",
-        "@fbgemm",
-        "@mkl",
-        "@sleef",
-        "@mkl_dnn//:mkl-dnn",
-    ],
-)
-
-cc_library(
-    name = "aten",
-    srcs = [
-        ":ATen_CORE_SRCS",
-        ":ATen_QUANTIZED_SRCS",
-        ":aten_base_cpp",
-        ":aten_base_metal",
-        ":aten_base_vulkan",
-        ":aten_native_cpp",
-        ":aten_native_mkl_cpp",
-        ":aten_native_mkldnn_cpp",
-        ":aten_native_nested_cpp",
-        ":aten_native_quantized_cpp",
-        ":aten_native_sparse_cpp",
-        ":aten_native_transformers_cpp",
-        ":aten_native_xnnpack",
-        ":aten_src_ATen_config",
-    ] + generated_cpu_cpp + aten_ufunc_generated_cpu_sources("aten/src/ATen/{}"),
-    copts = ATEN_COPTS,
-    linkopts = [
-      "-ldl",
-    ],
-    data = if_cuda(
-        [":libcaffe2_nvrtc.so"],
-        [],
-    ),
-    visibility = ["//visibility:public"],
-    deps = [
-        ":ATen_CPU",
-        ":aten_headers",
-        ":caffe2_for_aten_headers",
-        ":torch_headers",
-        "@fbgemm",
-        "@ideep",
-        "@nlohmann",
-    ],
-    alwayslink = True,
-)
-
-cc_library(
-    name = "aten_nvrtc",
-    srcs = glob([
-        "aten/src/ATen/cuda/nvrtc_stub/*.cpp",
-    ]),
-    copts = ATEN_COPTS,
-    linkstatic = True,
-    visibility = ["//visibility:public"],
-    deps = [
-        ":aten_headers",
-        "//c10",
-        "@cuda",
-        "@cuda//:cuda_driver",
-        "@cuda//:nvrtc",
-    ],
-    alwayslink = True,
-)
-
-cc_binary(
-    name = "libcaffe2_nvrtc.so",
-    linkshared = True,
-    visibility = ["//visibility:public"],
-    deps = [
-        ":aten_nvrtc",
-    ],
-)
-
-cc_library(
-    name = "aten_cuda_cpp",
-    srcs = [":aten_cuda_cpp_srcs"] + generated_cuda_cpp,
-    hdrs = [":aten_src_ATen_cuda_config"],
-    copts = ATEN_COPTS,
-    visibility = ["//visibility:public"],
-    deps = [
-        ":aten",
-        "@cuda",
-        "@cuda//:cusolver",
-        "@cuda//:nvrtc",
-        "@cudnn",
-        "@cudnn_frontend",
-    ],
-    alwayslink = True,
-)
-
-torch_cuda_half_options = [
-    "-DCUDA_HAS_FP16=1",
-    "-D__CUDA_NO_HALF_OPERATORS__",
-    "-D__CUDA_NO_HALF_CONVERSIONS__",
-    "-D__CUDA_NO_BFLOAT16_CONVERSIONS__",
-    "-D__CUDA_NO_HALF2_OPERATORS__",
-]
-
-cu_library(
-    name = "aten_cuda",
-    srcs = [":aten_cu_srcs"],
-    copts = ATEN_COPTS + torch_cuda_half_options,
-    visibility = ["//visibility:public"],
-    deps = [
-        ":aten_cuda_cpp",
-        "//c10/util:bit_cast",
-        "@cuda//:cublas",
-        "@cuda//:cufft",
-        "@cuda//:cusparse",
-        "@cutlass",
-    ],
-    alwayslink = True,
-)
-
-# caffe2
-CAFFE2_COPTS = COMMON_COPTS + [
-    "-Dcaffe2_EXPORTS",
-    "-DCAFFE2_USE_CUDNN",
-    "-DCAFFE2_BUILD_MAIN_LIB",
-    "-fvisibility-inlines-hidden",
-    "-fno-math-errno",
-    "-fno-trapping-math",
-]
-
-filegroup(
-    name = "caffe2_core_srcs",
-    srcs = [
-        "caffe2/core/common.cc",
-    ],
-)
-
-filegroup(
-    name = "caffe2_perfkernels_srcs",
-    srcs = [
-        "caffe2/perfkernels/embedding_lookup_idx.cc",
-    ],
-)
-
-
-filegroup(
-    name = "caffe2_serialize_srcs",
-    srcs = [
-        "caffe2/serialize/file_adapter.cc",
-        "caffe2/serialize/inline_container.cc",
-        "caffe2/serialize/istream_adapter.cc",
-        "caffe2/serialize/read_adapter_interface.cc",
-    ],
-)
-
-filegroup(
-    name = "caffe2_utils_srcs",
-    srcs = [
-        "caffe2/utils/proto_wrap.cc",
-        "caffe2/utils/string_utils.cc",
-        "caffe2/utils/threadpool/ThreadPool.cc",
-        "caffe2/utils/threadpool/pthreadpool.cc",
-        "caffe2/utils/threadpool/pthreadpool_impl.cc",
-        "caffe2/utils/threadpool/thread_pool_guard.cpp",
-    ],
-)
-
-# To achieve finer granularity and make debug easier, caffe2 is split into three libraries:
-# ATen, caffe2 and caffe2_for_aten_headers. ATen lib group up source codes under
-# aten/ directory and caffe2 contains most files under `caffe2/` directory. Since the
-# ATen lib and the caffe2 lib would depend on each other, `caffe2_for_aten_headers` is split
-# out from `caffe2` to avoid dependency cycle.
-cc_library(
-    name = "caffe2_for_aten_headers",
-    hdrs = [
-        "caffe2/core/common.h",
-        "caffe2/perfkernels/common.h",
-        "caffe2/perfkernels/embedding_lookup_idx.h",
-        "caffe2/utils/fixed_divisor.h",
-    ] + glob([
-        "caffe2/utils/threadpool/*.h",
-    ]),
-    copts = CAFFE2_COPTS,
-    visibility = ["//visibility:public"],
-    deps = [
-        ":caffe2_core_macros",
-        "//c10",
-    ],
-)
-
-cc_library(
-    name = "caffe2_headers",
-    hdrs = glob(
-        [
-            "caffe2/perfkernels/*.h",
-            "caffe2/serialize/*.h",
-            "caffe2/utils/*.h",
-            "caffe2/utils/threadpool/*.h",
-            "modules/**/*.h",
-        ],
-        exclude = [
-            "caffe2/core/macros.h",
-        ],
-    ) + if_cuda(glob([
-        "caffe2/**/*.cuh",
-    ])),
-    copts = CAFFE2_COPTS,
-    visibility = ["//visibility:public"],
-    deps = [
-        ":caffe2_core_macros",
-        ":caffe2_for_aten_headers",
-    ],
-)
-
-cc_library(
-    name = "caffe2",
-    srcs = [
-        ":caffe2_core_srcs",
-        ":caffe2_perfkernels_srcs",
-        ":caffe2_serialize_srcs",
-        ":caffe2_utils_srcs",
-    ],
-    copts = CAFFE2_COPTS + ["-mf16c"],
-    linkstatic = 1,
-    visibility = ["//visibility:public"],
-    deps = [
-        ":caffe2_core_macros",
-        ":caffe2_headers",
-        ":caffe2_perfkernels_avx",
-        ":caffe2_perfkernels_avx2",
-        "//third_party/miniz-3.0.2:miniz",
-        "@com_google_protobuf//:protobuf",
-        "@eigen",
-        "@fbgemm//:fbgemm_src_headers",
-        "@fmt",
-        "@onnx",
-    ] + if_cuda(
-        [
-            ":aten_cuda",
-            "@tensorpipe//:tensorpipe_cuda",
-        ],
-        [
-            ":aten",
-            "@tensorpipe//:tensorpipe_cpu",
-        ],
-    ),
-    alwayslink = True,
-)
-
-cu_library(
-    name = "torch_cuda",
-    srcs = [
-        "torch/csrc/distributed/c10d/NanCheck.cu",
-        "torch/csrc/distributed/c10d/quantization/quantization_gpu.cu",
-        "torch/csrc/distributed/c10d/symm_mem/intra_node_comm.cu",
-    ],
-    copts = torch_cuda_half_options,
-    visibility = ["//visibility:public"],
-    deps = [
-        ":aten",
-        "@cuda//:cublas",
-        "@cuda//:curand",
-        "@cudnn",
-        "@eigen",
-        "@tensorpipe//:tensorpipe_cuda",
-    ],
-    alwayslink = True,
-)
-
-PERF_COPTS = [
-    "-DHAVE_AVX_CPU_DEFINITION",
-    "-DHAVE_AVX2_CPU_DEFINITION",
-    "-DENABLE_ALIAS=1",
-    "-DHAVE_MALLOC_USABLE_SIZE=1",
-    "-DHAVE_MMAP=1",
-    "-DHAVE_SHM_OPEN=1",
-    "-DHAVE_SHM_UNLINK=1",
-    "-DSLEEF_STATIC_LIBS=1",
-    "-DTH_BALS_MKL",
-    "-D_FILE_OFFSET_BITS=64",
-    "-DUSE_FBGEMM",
-    "-fvisibility-inlines-hidden",
-    "-Wunused-parameter",
-    "-fno-math-errno",
-    "-fno-trapping-math",
-    "-mf16c",
-]
-
-PERF_HEADERS = glob([
-    "caffe2/perfkernels/*.h",
-    "caffe2/core/*.h",
-])
-
-cc_library(
-    name = "caffe2_perfkernels_avx",
-    srcs = glob([
-        "caffe2/perfkernels/*_avx.cc",
-    ]),
-    hdrs = PERF_HEADERS,
-    copts = PERF_COPTS + [
-        "-mavx",
-    ],
-    visibility = ["//visibility:public"],
-    deps = [
-        ":caffe2_headers",
-        "//c10",
-    ],
-    alwayslink = True,
-)
-
-cc_library(
-    name = "caffe2_perfkernels_avx2",
-    srcs = glob([
-        "caffe2/perfkernels/*_avx2.cc",
-    ]),
-    hdrs = PERF_HEADERS,
-    copts = PERF_COPTS + [
-        "-mavx2",
-        "-mfma",
-        "-mavx",
-    ],
-    visibility = ["//visibility:public"],
-    deps = [
-        ":caffe2_headers",
-        "//c10",
-    ],
-    alwayslink = True,
-)
-
-# torch
-torch_cuda_headers = glob(["torch/csrc/cuda/*.h"])
-
-flatbuffer_cc_library(
-    name = "torch_flatbuffers",
-    srcs = [
-        "torch/csrc/jit/serialization/mobile_bytecode.fbs",
-    ],
-    flatc_args = ["--cpp", "--gen-mutable", "--scoped-enums"],
-    out_prefix = "torch/csrc/jit/serialization/",
-)
-
-cc_library(
-    name = "torch_headers",
-    hdrs = if_cuda(
-        torch_cuda_headers,
-    ) + glob(
-        [
-            "torch/*.h",
-            "torch/csrc/**/*.h",
-            "torch/nativert/**/*.h",
-            "torch/csrc/distributed/c10d/**/*.hpp",
-            "torch/lib/libshm/*.h",
-        ],
-        exclude = [
-            "torch/csrc/*/generated/*.h",
-            "torch/csrc/jit/serialization/mobile_bytecode_generated.h",
-        ] + torch_cuda_headers,
-    ) + GENERATED_AUTOGRAD_CPP + [
-        "//torch/headeronly:version_h",
-    ],
-    includes = [
-        "third_party/kineto/libkineto/include",
-        "torch/csrc",
-        "torch/csrc/api/include",
-        "torch/csrc/distributed",
-        "torch/lib",
-        "torch/lib/libshm",
-    ],
-    visibility = ["//visibility:public"],
-    deps = [
-        ":aten_headers",
-        ":caffe2_headers",
-        ":torch_flatbuffers",
-        "//c10",
-        "@com_github_google_flatbuffers//:flatbuffers",
-        "@local_config_python//:python_headers",
-        "@onnx",
-    ],
-    alwayslink = True,
-)
-
-TORCH_COPTS = COMMON_COPTS + [
-    "-Dtorch_EXPORTS",
-    "-DHAVE_AVX_CPU_DEFINITION",
-    "-DHAVE_AVX2_CPU_DEFINITION",
-    "-DCAFFE2_USE_GLOO",
-    "-fvisibility-inlines-hidden",
-    "-fno-math-errno ",
-    "-fno-trapping-math",
-    "-Wno-error=unused-function",
-]
-
-torch_sources = {
-    k: ""
-    for k in (
-        libtorch_core_sources +
-        libtorch_distributed_sources +
-        torch_cpp_srcs +
-        libtorch_extra_sources +
-        jit_core_sources +
-        lazy_tensor_ts_sources +
-        GENERATED_AUTOGRAD_CPP
-    )
-}.keys()
-
-cc_library(
-    name = "torch",
-    srcs = if_cuda(glob(
-        libtorch_cuda_sources,
-        exclude = [
-            "torch/csrc/cuda/nccl.cpp",
-            "torch/csrc/cuda/python_nccl.cpp",
-            "torch/csrc/distributed/c10d/NanCheck.cu",
-            "torch/csrc/distributed/c10d/cuda/AsyncMM.cu",
-            "torch/csrc/distributed/c10d/quantization/quantization_gpu.cu",
-            "torch/csrc/distributed/c10d/symm_mem/CUDASymmetricMemory.cu",
-            "torch/csrc/distributed/c10d/symm_mem/CUDASymmetricMemoryOps.cu",
-            "torch/csrc/distributed/c10d/symm_mem/CUDASymmetricMemoryUtils.cpp",
-            "torch/csrc/distributed/c10d/symm_mem/cuda_mem_pool.cpp",
-            "torch/csrc/distributed/c10d/symm_mem/intra_node_comm.cu",
-        ],
-    )) + torch_sources,
-    copts = TORCH_COPTS,
-    linkopts = [
-      "-lrt",
-    ],
-    defines = [
-        "CAFFE2_NIGHTLY_VERSION=20200115",
-    ],
-    visibility = ["//visibility:public"],
-    deps = [
-        ":caffe2",
-        ":torch_headers",
-        "@kineto",
-        "@cpp-httplib",
-        "@nlohmann",
-    ] + if_cuda([
-        "@cuda//:nvToolsExt",
-        "@cutlass",
-        ":torch_cuda",
-    ]),
-    alwayslink = True,
-)
-
-cc_library(
-    name = "shm",
-    srcs = glob(["torch/lib/libshm/*.cpp"]),
-    linkopts = [
-      "-lrt",
-    ],
-    deps = [
-        ":torch",
-    ],
-)
-
-cc_library(
-    name = "libtorch_headers",
-    hdrs = glob([
-        "**/*.h",
-        "**/*.cuh",
-    ]) + [
-        # We need the filegroup here because the raw list causes Bazel
-        # to see duplicate files. It knows how to deduplicate with the
-        # filegroup.
-        ":cpp_generated_code",
-    ],
-    includes = [
-        "torch/csrc/api/include",
-        "torch/csrc/distributed",
-        "torch/lib",
-        "torch/lib/libshm",
-    ],
-    visibility = ["//visibility:public"],
-    deps = [
-        ":torch_headers",
-    ],
-)
-
-cc_library(
-    name = "torch_python",
-    srcs = libtorch_python_core_sources
-        + if_cuda(libtorch_python_cuda_sources)
-        + if_cuda(libtorch_python_distributed_sources)
-        + GENERATED_AUTOGRAD_PYTHON,
-    hdrs = glob([
-        "torch/csrc/generic/*.cpp",
-    ]),
-    copts = COMMON_COPTS + if_cuda(["-DUSE_CUDA=1"]),
-    deps = [
-        ":torch",
-        ":shm",
-        "@pybind11",
-    ],
-)
-
-pybind_extension(
-    name = "torch/_C",
-    srcs = ["torch/csrc/stub.c"],
-    deps = [
-        ":torch_python",
-        ":aten_nvrtc",
-    ],
-)
-
-cc_binary(
-    name = "torch/bin/torch_shm_manager",
-    srcs = [
-        "torch/lib/libshm/manager.cpp",
-    ],
-    deps = [
-        ":shm",
-    ],
-    linkstatic = False,
-)
-
-template_rule(
-    name = "gen_version_py",
-    src = ":torch/version.py.tpl",
-    out = "torch/version.py",
-    substitutions = if_cuda({
-        # Set default to 11.2. Otherwise Torchvision complains about incompatibility.
-        "{{CUDA_VERSION}}": "11.2",
-        "{{VERSION}}": "2.0.0",
-    }, {
-        "{{CUDA_VERSION}}": "None",
-        "{{VERSION}}": "2.0.0",
-    }),
-)
-
-py_library(
-    name = "pytorch_py",
-    visibility = ["//visibility:public"],
-    srcs = glob(["torch/**/*.py"], exclude = ["torch/version.py"]) + [":torch/version.py"] + glob(["functorch/**/*.py"]),
-    deps = [
-        rules.requirement("numpy"),
-        rules.requirement("pyyaml"),
-        rules.requirement("requests"),
-        rules.requirement("setuptools"),
-        rules.requirement("sympy"),
-        rules.requirement("typing_extensions"),
-        "//torchgen",
-    ],
-    data = [
-        ":torch/_C.so",
-        ":torch/bin/torch_shm_manager",
-    ],
-)
-
-# cpp api tests
-cc_library(
-    name = "test_support",
-    testonly = True,
-    srcs = [
-        "test/cpp/api/support.cpp",
-    ],
-    hdrs = [
-        "test/cpp/api/init_baseline.h",
-        "test/cpp/api/optim_baseline.h",
-        "test/cpp/api/support.h",
-        "test/cpp/common/support.h",
-    ],
-    deps = [
-        ":torch",
-        "@com_google_googletest//:gtest_main",
-    ],
-)
-
-# Torch integration tests rely on a labeled data set from the MNIST database.
-# http://yann.lecun.com/exdb/mnist/
-
-cpp_api_tests = glob(
-    ["test/cpp/api/*.cpp"],
-    exclude = [
-        "test/cpp/api/imethod.cpp",
-        "test/cpp/api/integration.cpp",
-    ],
-)
-
-cc_test(
-    name = "integration_test",
-    size = "medium",
-    srcs = ["test/cpp/api/integration.cpp"],
-    data = [
-        ":download_mnist",
-    ],
-    tags = [
-        "gpu-required",
-    ],
-    deps = [
-        ":test_support",
-        "@com_google_googletest//:gtest_main",
-    ],
-)
-
-[
-    cc_test(
-        name = paths.split_extension(paths.basename(filename))[0].replace("-", "_") + "_test",
-        size = "medium",
-        srcs = [filename],
-        deps = [
-            ":test_support",
-            "@com_google_googletest//:gtest_main",
-        ],
-    )
-    for filename in cpp_api_tests
-]
-
-test_suite(
-    name = "api_tests",
-    tests = [
-        "any_test",
-        "autograd_test",
-        "dataloader_test",
-        "enum_test",
-        "expanding_array_test",
-        "functional_test",
-        "init_test",
-        "integration_test",
-        "jit_test",
-        "memory_test",
-        "misc_test",
-        "module_test",
-        "modulelist_test",
-        "modules_test",
-        "nn_utils_test",
-        "optim_test",
-        "ordered_dict_test",
-        "rnn_test",
-        "sequential_test",
-        "serialize_test",
-        "static_test",
-        "tensor_options_test",
-        "tensor_test",
-        "torch_include_test",
-    ],
-)
-
-# dist autograd tests
-cc_test(
-    name = "torch_dist_autograd_test",
-    size = "small",
-    srcs = ["test/cpp/dist_autograd/test_dist_autograd.cpp"],
-    tags = [
-        "exclusive",
-        "gpu-required",
-    ],
-    deps = [
-        ":torch",
-        "@com_google_googletest//:gtest_main",
-    ],
-)
-
-# jit tests
-# Because these individual unit tests require custom registering,
-# it is easier to mimic the cmake build by globing together a single test.
-cc_test(
-    name = "jit_tests",
-    size = "small",
-    srcs = glob(
-        [
-            "test/cpp/jit/*.cpp",
-            "test/cpp/jit/*.h",
-            "test/cpp/tensorexpr/*.cpp",
-            "test/cpp/tensorexpr/*.h",
-        ],
-        exclude = [
-            # skip this since <pybind11/embed.h> is not found in OSS build
-            "test/cpp/jit/test_exception.cpp",
-        ],
-    ),
-    linkstatic = True,
-    tags = [
-        "exclusive",
-        "gpu-required",
-    ],
-    deps = [
-        ":torch",
-        "@com_google_googletest//:gtest_main",
-    ],
-)
-
-cc_test(
-    name = "lazy_tests",
-    size = "small",
-    srcs = glob(
-        [
-            "test/cpp/lazy/*.cpp",
-            "test/cpp/lazy/*.h",
-        ],
-        exclude = [
-            # skip these since they depend on generated LazyIr.h which isn't available in bazel yet
-            "test/cpp/lazy/test_ir.cpp",
-            "test/cpp/lazy/test_lazy_ops.cpp",
-            "test/cpp/lazy/test_lazy_ops_util.cpp",
-            "test/cpp/lazy/test_lazy_graph_executor.cpp",
-        ],
-    ),
-    linkstatic = True,
-    tags = [
-        "exclusive",
-    ],
-    deps = [
-        ":torch",
-        "@com_google_googletest//:gtest_main",
-    ],
-)
-
-# python api tests
-
-py_test(
-    name = "test_bazel",
-    srcs = ["test/_test_bazel.py"],
-    main = "test/_test_bazel.py",
-    deps = [
-        ":pytorch_py",
-        rules.requirement("networkx"),
-    ],
-)
-
-# all tests
-test_suite(
-    name = "all_tests",
-    tests = [
-        "api_tests",
-        "jit_tests",
-        "torch_dist_autograd_test",
-        "//c10/test:tests",
-    ],
-)
-
-# An internal genrule that we are converging with refers to these file
-# as if they are from this package, so we alias them for
-# compatibility.
-
-[
-    alias(
-        name = paths.basename(path),
-        actual = path,
-    )
-    for path in [
-        "aten/src/ATen/templates/DispatchKeyNativeFunctions.cpp",
-        "aten/src/ATen/templates/DispatchKeyNativeFunctions.h",
-        "aten/src/ATen/templates/LazyIr.h",
-        "aten/src/ATen/templates/LazyNonNativeIr.h",
-        "aten/src/ATen/templates/RegisterDispatchKey.cpp",
-        "aten/src/ATen/templates/RegisterDispatchDefinitions.ini",
-        "aten/src/ATen/templates/ViewMetaClassesPythonBinding.cpp",
-        "aten/src/ATen/native/native_functions.yaml",
-        "aten/src/ATen/native/tags.yaml",
-        "aten/src/ATen/native/ts_native_functions.yaml",
-        "torch/csrc/lazy/core/shape_inference.h",
-        "torch/csrc/lazy/ts_backend/ts_native_functions.cpp",
-    ]
-]
-
-genrule(
-    name = "download_mnist",
-    srcs = ["//:tools/download_mnist.py"],
-    outs = [
-        "mnist/train-images-idx3-ubyte",
-        "mnist/train-labels-idx1-ubyte",
-        "mnist/t10k-images-idx3-ubyte",
-        "mnist/t10k-labels-idx1-ubyte",
-    ],
-    cmd = "python3 tools/download_mnist.py -d $(RULEDIR)/mnist",
-)
diff --git a/CLAUDE.md b/CLAUDE.md
index 245820413f52a..640a24bbcb410 100644
--- a/CLAUDE.md
+++ b/CLAUDE.md
@@ -1,3 +1,11 @@
+# Scratch Space
+
+Use `agent_space/` (git-ignored, at repo root) for temporary scripts, scratch files, and throwaway experiments. Do not commit files from this directory.
+
+# PR Review
+
+When asked to review a PR, always use the /pr-review skill.
+
 # Environment
 
 If any tool you're trying to use (pip, python, spin, etc) is missing, always stop and ask the user if an environment is needed. Do NOT try to find alternatives or install these tools.
@@ -23,6 +31,8 @@ if __name__ == "__main__":
 ```
 
 To test Tensor equality, use assertEqual.
+For tests over multiple inputs, use the `@parametrize` decorator.
+For any test that checks numerics of the on-device implementation, use `instantiate_device_type_tests` to write device-generic tests.
 
 # Linting
 
@@ -41,6 +51,10 @@ entirely.
 
 Disclose that the PR was authored with Claude.
 
+If a commit message contains `ghstack-source-id` or `Pull-Request` trailers,
+you MUST preserve them when rewriting or splitting commit messages. ghstack
+will update the source id automatically when needed.
+
 # Coding Style Guidelines
 
 Follow these rules for all code changes in this repository:
diff --git a/CMakeLists.txt b/CMakeLists.txt
index 188f08bb272f5..d5941243d66db 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -27,6 +27,17 @@ endif()
 # ---[ Project and semantic versioning.
 project(Torch CXX C)
 
+# When built via scikit-build-core, Python_EXECUTABLE is already set.
+# For standalone cmake invocations, find it via find_package.
+find_package(Python COMPONENTS Interpreter REQUIRED)
+
+# Forward environment variables to CMake variables, replicating the behavior
+# previously handled by setup.py / tools/setup_helpers/cmake.py.
+# When built via setup.py, these are already set as -D flags; the module
+# skips variables that are already defined, so it is safe to include
+# unconditionally.
+include(cmake/EnvVarForwarding.cmake)
+
 if(${CMAKE_SYSTEM_NAME} STREQUAL "Linux")
   set(LINUX TRUE)
 else()
@@ -40,12 +51,12 @@ string(FIND "${CMAKE_CXX_FLAGS}" "-std=c++" env_cxx_standard)
 if(env_cxx_standard GREATER -1)
   message(
     WARNING
-      "C++ standard version definition detected in environment variable."
-      "PyTorch requires -std=c++17. Please remove -std=c++ settings in your environment."
+      "C++ standard version definition detected in environment variable. "
+      "PyTorch requires -std=c++20. Please remove -std=c++ settings in your environment."
   )
 endif()
 set(CMAKE_CXX_STANDARD
-    17
+    20
     CACHE STRING
           "The C++ standard whose features are requested to build this target.")
 set(CMAKE_C_STANDARD
@@ -56,11 +67,11 @@ set(CMAKE_C_STANDARD
 # ---[ Utils
 include(cmake/public/utils.cmake)
 
-# --- [ Check that minimal gcc version is 9.3+
-if(CMAKE_COMPILER_IS_GNUCXX AND CMAKE_CXX_COMPILER_VERSION VERSION_LESS 9.3)
+# --- [ Check that minimal gcc version is 11.3+
+if(CMAKE_COMPILER_IS_GNUCXX AND CMAKE_CXX_COMPILER_VERSION VERSION_LESS 11.3)
   message(
     FATAL_ERROR
-      "GCC-9.3 or newer is required to compile PyTorch, but found ${CMAKE_CXX_COMPILER_VERSION}"
+      "GCC-11.3 or newer is required to compile PyTorch, but found ${CMAKE_CXX_COMPILER_VERSION}"
   )
 endif()
 
@@ -249,6 +260,7 @@ endif()
 option(USE_CUDA "Use CUDA" ON)
 
 option(USE_XPU "Use XPU" ON)
+option(USE_MTIA "Use MTIA" OFF)
 cmake_dependent_option(
   BUILD_LAZY_CUDA_LINALG "Build cuda linalg ops as separate library" ON
   "USE_CUDA AND LINUX AND BUILD_PYTHON" OFF)
@@ -291,7 +303,7 @@ option(USE_NNPACK "Use NNPACK" ON)
 cmake_dependent_option(USE_NUMA "Use NUMA. Only available on Linux." ON "LINUX"
                        OFF)
 cmake_dependent_option(USE_NVRTC "Use NVRTC. Only available if USE_CUDA is on."
-                       OFF "USE_CUDA" OFF)
+                       ON "USE_CUDA" OFF)
 option(USE_NUMPY "Use NumPy" ON)
 option(USE_OBSERVERS "Use observers module." OFF)
 option(USE_OPENCL "Use OpenCL" OFF)
@@ -373,7 +385,7 @@ option(HAVE_SOVERSION "Whether to add SOVERSION to the shared objects" OFF)
 option(BUILD_LIBTORCH_CPU_WITH_DEBUG
        "Enable RelWithDebInfo for libtorch_cpu target only" OFF)
 cmake_dependent_option(
-  USE_CCACHE "Attempt using CCache to wrap the compilation" ON "UNIX" OFF)
+  USE_CCACHE "Attempt using [S]CCache to wrap the compilation" ON "UNIX" OFF)
 option(WERROR "Build with -Werror supported by the compiler" OFF)
 option(
   DEBUG_CUDA
@@ -420,20 +432,23 @@ endif()
 
 if(USE_CCACHE)
   find_program(CCACHE_PROGRAM ccache)
+  find_program(SCCACHE_EXECUTABLE sccache)
   if(CCACHE_PROGRAM)
-    set(CMAKE_C_COMPILER_LAUNCHER
-        "${CCACHE_PROGRAM}"
-        CACHE STRING "C compiler launcher")
-    set(CMAKE_CXX_COMPILER_LAUNCHER
-        "${CCACHE_PROGRAM}"
-        CACHE STRING "CXX compiler launcher")
-    set(CMAKE_CUDA_COMPILER_LAUNCHER
-        "${CCACHE_PROGRAM}"
-        CACHE STRING "CUDA compiler launcher")
+    foreach(LANG CXX C CUDA)
+      set(CMAKE_${LANG}_COMPILER_LAUNCHER
+          "${CCACHE_PROGRAM}"
+          CACHE STRING "${LANG} compiler launcher")
+    endforeach()
+  elseif(SCCACHE_EXECUTABLE)
+    foreach(LANG CXX C CUDA)
+      set(CMAKE_${LANG}_COMPILER_LAUNCHER
+          "${SCCACHE_EXECUTABLE}"
+          CACHE STRING "${LANG} compiler launcher")
+    endforeach()
   else()
     message(
       STATUS
-        "Could not find ccache. Consider installing ccache to speed up compilation."
+        "Could not find neither ccache nor sccache. Consider installing ccache to speed up compilation."
     )
   endif()
 endif()
@@ -587,7 +602,7 @@ if(MSVC)
   # https://learn.microsoft.com/en-us/cpp/build/reference/zc-cplusplus?view=msvc-170
   # * https://en.cppreference.com/w/cpp/preprocessor/replace#Predefined_macros
   set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} /Zc:__cplusplus")
-  set(CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS} -Xcompiler  /Zc:__cplusplus")
+  set(CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS} -Xcompiler /Zc:__cplusplus -Xcompiler /Zc:preprocessor")
 
   set(CMAKE_NINJA_CMCLDEPS_RC OFF)
   if(MSVC_Z7_OVERRIDE)
@@ -810,6 +825,7 @@ if(INTERN_BUILD_MOBILE)
 endif()
 
 # ---[ Version numbers for generated libraries
+set_property(DIRECTORY APPEND PROPERTY CMAKE_CONFIGURE_DEPENDS version.txt)
 file(READ version.txt TORCH_DEFAULT_VERSION)
 # Strip trailing newline
 string(REGEX REPLACE "\n$" "" TORCH_DEFAULT_VERSION "${TORCH_DEFAULT_VERSION}")
@@ -821,18 +837,12 @@ if("${TORCH_DEFAULT_VERSION} " STREQUAL " ")
 endif()
 set(TORCH_BUILD_VERSION
     "${TORCH_DEFAULT_VERSION}"
-    CACHE STRING "Torch build version")
+    CACHE STRING "Torch build version" FORCE)
 if(DEFINED ENV{PYTORCH_BUILD_VERSION})
   set(TORCH_BUILD_VERSION
       "$ENV{PYTORCH_BUILD_VERSION}"
       CACHE STRING "Torch build version" FORCE)
 endif()
-if(NOT TORCH_BUILD_VERSION)
-  # An empty string was specified so force version to the default
-  set(TORCH_BUILD_VERSION
-      "${TORCH_DEFAULT_VERSION}"
-      CACHE STRING "Torch build version" FORCE)
-endif()
 caffe2_parse_version_str(TORCH ${TORCH_BUILD_VERSION})
 set(TORCH_SOVERSION "${TORCH_VERSION_MAJOR}.${TORCH_VERSION_MINOR}")
 
@@ -1183,6 +1193,12 @@ if(NOT MSVC)
     string(APPEND CMAKE_CXX_FLAGS_DEBUG " -fno-omit-frame-pointer -O0")
     string(APPEND CMAKE_LINKER_FLAGS_DEBUG " -fno-omit-frame-pointer -O0")
   endif()
+  # aarch64 C++ stack unwinding uses frame-pointer chain walking, so frame
+  # pointers must be present in all build types.  The cost is negligible on
+  # aarch64 (31 GPRs vs x86-64's 16, so dedicating x29 rarely spills).
+  if(CMAKE_SYSTEM_PROCESSOR MATCHES "aarch64")
+    append_cxx_flag_if_supported("-fno-omit-frame-pointer" CMAKE_CXX_FLAGS)
+  endif()
   append_cxx_flag_if_supported("-fno-math-errno" CMAKE_CXX_FLAGS)
   append_cxx_flag_if_supported("-fno-trapping-math" CMAKE_CXX_FLAGS)
   append_cxx_flag_if_supported("-Werror=format" CMAKE_CXX_FLAGS)
@@ -1190,6 +1206,10 @@ if(NOT MSVC)
     append_cxx_flag_if_supported("-Wno-dangling-reference" CMAKE_CXX_FLAGS)
     append_cxx_flag_if_supported("-Wno-error=dangling-reference" CMAKE_CXX_FLAGS)
   endif()
+  # needed for compat with newer versions of clang that use C++20 mangling rules
+  if(CMAKE_CXX_COMPILER_ID MATCHES "Clang" AND CMAKE_CXX_COMPILER_VERSION VERSION_GREATER_EQUAL 18)
+    append_cxx_flag_if_supported("-fclang-abi-compat=17" CMAKE_CXX_FLAGS)
+  endif()
 else()
   # Define export functions for AOTI.
   add_compile_definitions(EXPORT_AOTI_FUNCTIONS)
@@ -1227,10 +1247,11 @@ if(USE_CPP_CODE_COVERAGE)
   if("${CMAKE_CXX_COMPILER_ID}" STREQUAL "GNU")
     string(APPEND CMAKE_C_FLAGS " --coverage -fprofile-abs-path")
     string(APPEND CMAKE_CXX_FLAGS " --coverage -fprofile-abs-path")
+    string(APPEND CMAKE_OBJCXX_FLAGS " --coverage -fprofile-abs-path")
   elseif("${CMAKE_CXX_COMPILER_ID}" MATCHES "Clang")
     string(APPEND CMAKE_C_FLAGS " -fprofile-instr-generate -fcoverage-mapping")
-    string(APPEND CMAKE_CXX_FLAGS
-           " -fprofile-instr-generate -fcoverage-mapping")
+    string(APPEND CMAKE_CXX_FLAGS " -fprofile-instr-generate -fcoverage-mapping")
+    string(APPEND CMAKE_OBJCXX_FLAGS " -fprofile-instr-generate -fcoverage-mapping")
   else()
     message(
       ERROR
diff --git a/CODEOWNERS b/CODEOWNERS
index b869b69cbfae7..604ad3c694d9a 100644
--- a/CODEOWNERS
+++ b/CODEOWNERS
@@ -48,6 +48,17 @@ nn/qat/ @jerryzh168
 # c10d backend APIs
 /torch/csrc/distributed/c10d/Backend.* @kwen2501
 /torch/csrc/distributed/c10d/Ops.* @kwen2501
+# FSDP, DDP
+/torch/distributed/_composable/fsdp/ @weifengpy
+/torch/distributed/_composable/replicate.py @weifengpy
+/torch/distributed/_composable/replicate_with_fsdp.py @weifengpy
+/torch/distributed/fsdp/ @weifengpy
+/torch/distributed/optim/ @weifengpy
+/torch/nn/parallel/ @weifengpy
+# DTensor view, matmul, _StridedShard
+/torch/distributed/tensor/placement_types.py @weifengpy
+/torch/distributed/tensor/_ops/_view_ops.py @weifengpy
+/torch/distributed/tensor/_ops/_matrix_ops.py @weifengpy
 
 # ONNX Export
 /torch/_dynamo/backends/onnxrt.py @titaiwangms @xadupre @justinchuby
@@ -121,19 +132,20 @@ torch/_subclasses/fake_tensor.py @aorenste
 torch/_subclasses/fake_utils.py @aorenste
 
 # torch MPS
-test/test_mps.py @kulinseth @malfet
-aten/src/ATen/mps/ @kulinseth @malfet
-aten/src/ATen/native/mps/ @kulinseth @malfet
+test/test_mps.py @malfet
+c10/metal/ @malfet
+aten/src/ATen/mps/ @malfet
+aten/src/ATen/native/mps/ @malfet
 
 # MTIA
 aten/src/ATen/detail/MTIAHooksInterface.h @egienvalue
 torch/csrc/mtia/ @egienvalue
 
 # Profiler
-torch/csrc/autograd/profiler* @sraikund16
-torch/autograd/profiler* @sraikund16
-torch/csrc/profiler/ @sraikund16
-torch/profiler/ @sraikund16
+torch/csrc/autograd/profiler* @scotts
+torch/autograd/profiler* @scotts
+torch/csrc/profiler/ @scotts
+torch/profiler/ @scotts
 
 # AOTDispatch tests
 test/functorch/test_aotdispatch.py @ezyang @Chillee
@@ -172,9 +184,9 @@ caffe2/utils/hip @jeffdaily @jithunnair-amd
 /torch/_export/serde/schema.py @SherlockNoMad @zhxchen17
 
 # Dynamic Shapes
-/torch/fx/experimental/symbolic_shapes.py @bobrenjc93 @laithsakka
-/torch/fx/experimental/sym_node.py @bobrenjc93 @laithsakka
-/torch/fx/experimental/recording.py @bobrenjc93 @laithsakka
+/torch/fx/experimental/symbolic_shapes.py @laithsakka
+/torch/fx/experimental/sym_node.py @laithsakka
+/torch/fx/experimental/recording.py @laithsakka
 
 # ProxyTorchDispatchMode
 torch/fx/experimental/proxy_tensor.py @aorenste
@@ -235,3 +247,6 @@ torch/backends/cudnn/ @eqy @syed-ahmed @Aidyn-A
 /torch/sparse/ @amjames @pearu @nikitaved
 /test/test_sparse.py @amjames @pearu @nikitaved
 /test/test_sparse_csr.py @amjames @pearu @nikitaved
+
+# Python/DSL native ops
+/torch/_native/ @slayton58 @drisspg
diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md
index bf18749ae4329..e92aec464946a 100644
--- a/CONTRIBUTING.md
+++ b/CONTRIBUTING.md
@@ -16,6 +16,7 @@ aspects of contributing to PyTorch.
 - [Codebase structure](#codebase-structure)
 - [AI-Assisted Development](#ai-assisted-development)
 - [Spin](#spin)
+  - [Building](#building)
   - [Linting](#linting)
     - [default lint](#default-lint)
   - [Regenerating](#regenerating)
@@ -277,7 +278,6 @@ dependencies as well as the nightly binaries into the repo directory.
   * [core](caffe2/core) - Core files of Caffe2, e.g., tensor, workspace,
     blobs, etc.
   * ...
-* [.circleci](.circleci) - CircleCI configuration management. [README](.circleci/README.md)
 
 ## AI-Assisted Development
 
@@ -301,6 +301,14 @@ helps running common tasks.
 To list the available tasks, run `spin --help`.
 Currently, we support the following tasks with Spin:
 
+### Building
+
+To support building and general development, the following commands exist.
+
+|command||
+|-|-|
+|`clean`|clean, that is remove files and directories listed in .gitignore before the NOT-CLEAN-FILES marker|
+
 ### Linting
 
 Spin helps with linting by making sure that lintrunner is installed correctly
diff --git a/Dockerfile b/Dockerfile
index 7b28e9b056e62..1de43ad526687 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -44,10 +44,17 @@ ARG INSTALL_CHANNEL=whl/nightly
 # Automatically set by buildx
 ARG TARGETPLATFORM
 
-# INSTALL_CHANNEL whl - release, whl/nightly - nightly, whl/test - test channels
+# torchaudio does not publish wheels against the CUDA 13.2 index yet,
+# so skip it when CUDA_PATH=cu132 to keep this Dockerfile building.
 RUN case ${TARGETPLATFORM} in \
          "linux/arm64")  pip3 install --extra-index-url https://download.pytorch.org/whl/cpu/ torch torchvision torchaudio ;; \
-         *)              pip3 install --index-url https://download.pytorch.org/${INSTALL_CHANNEL}/${CUDA_PATH#.}/ torch torchvision torchaudio ;; \
+         *) \
+             if [ "${CUDA_PATH}" = "cu132" ]; then \
+                 pip3 install --index-url https://download.pytorch.org/${INSTALL_CHANNEL}/${CUDA_PATH}/ torch torchvision; \
+             else \
+                 pip3 install --index-url https://download.pytorch.org/${INSTALL_CHANNEL}/${CUDA_PATH}/ torch torchvision torchaudio; \
+             fi \
+             ;; \
     esac
 RUN pip3 install torchelastic
 RUN IS_CUDA=$(python3 -c 'import torch ; print(torch.cuda._is_compiled())'); \
diff --git a/README.md b/README.md
index 4be858bea6b08..eddc034f7b9b1 100644
--- a/README.md
+++ b/README.md
@@ -35,6 +35,7 @@ Our trunk health (Continuous Integration signals) can be found at [hud.pytorch.o
     - [Using pre-built images](#using-pre-built-images)
     - [Building the image yourself](#building-the-image-yourself)
   - [Building the Documentation](#building-the-documentation)
+    - [Troubleshooting CI Errors](#troubleshooting-ci-errors)
     - [Building a PDF](#building-a-pdf)
   - [Previous Versions](#previous-versions)
 - [Getting Started](#getting-started)
@@ -134,7 +135,7 @@ This enables you to train bigger deep learning models than before.
 
 ### Extensions Without Pain
 
-Writing new neural network modules, or interfacing with PyTorch's Tensor API was designed to be straightforward
+Writing new neural network modules, or interfacing with PyTorch's Tensor API, was designed to be straightforward
 and with minimal abstractions.
 
 You can write new neural network layers in Python using the torch API
@@ -162,8 +163,10 @@ They require JetPack 4.2 and above, and [@dusty-nv](https://github.com/dusty-nv)
 #### Prerequisites
 If you are installing from source, you will need:
 - Python 3.10 or later
-- A compiler that fully supports C++17, such as clang or gcc (gcc 9.4.0 or newer is required, on Linux)
+- A compiler that fully supports C++20, such as clang or gcc (gcc 11.3.0 or newer is required, on Linux)
 - Visual Studio or Visual Studio Build Tool (Windows only)
+- At least 10 GB of free disk space
+- 30-60 minutes for the initial build (subsequent rebuilds are much faster)
 
 \* PyTorch CI uses Visual C++ BuildTools, which come with Visual Studio Enterprise,
 Professional, or Community Editions. You can also install the build tools from
@@ -179,7 +182,6 @@ $ source <CONDA_INSTALL_DIR>/bin/activate
 $ conda create -y -n <CONDA_NAME>
 $ conda activate <CONDA_NAME>
 ```
-
 * Windows:
 
 ```bash
@@ -197,7 +199,7 @@ packages (e.g., CUDA, MKL.)
 ##### NVIDIA CUDA Support
 If you want to compile with CUDA support, [select a supported version of CUDA from our support matrix](https://pytorch.org/get-started/locally/), then install the following:
 - [NVIDIA CUDA](https://developer.nvidia.com/cuda-downloads)
-- [NVIDIA cuDNN](https://developer.nvidia.com/cudnn) v8.5 or above
+- [NVIDIA cuDNN](https://developer.nvidia.com/cudnn) v9.0 or above
 - [Compiler](https://gist.github.com/ax3l/9489132) compatible with CUDA
 
 Note: You could refer to the [cuDNN Support Matrix](https://docs.nvidia.com/deeplearning/cudnn/backend/latest/reference/support-matrix.html) for cuDNN versions with the various supported CUDA, CUDA driver, and NVIDIA hardware.
@@ -260,15 +262,6 @@ pip install mkl-static mkl-include
 make triton
 ```
 
-**On MacOS**
-
-```bash
-# Add this package on intel x86 processor machines only
-pip install mkl-static mkl-include
-# Add these packages if torch.distributed is needed
-conda install pkg-config libuv
-```
-
 **On Windows**
 
 ```bash
@@ -407,7 +400,7 @@ ccmake build  # or cmake-gui build
 
 #### Using pre-built images
 
-You can also pull a pre-built docker image from Docker Hub and run with docker v19.03+
+You can also pull a pre-built docker image from Docker Hub and run with docker v23.0+
 
 ```bash
 docker run --gpus all --rm -ti --ipc=host pytorch/pytorch:latest
@@ -419,11 +412,11 @@ should increase shared memory size either with `--ipc=host` or `--shm-size` comm
 
 #### Building the image yourself
 
-**NOTE:** Must be built with a docker version > 18.06
+**NOTE:** Must be built with a Docker version >= 23.0
 
-The `Dockerfile` is supplied to build images with CUDA 11.1 support and cuDNN v8.
+The Dockerfile is supplied to build images with CUDA 12.1 support and cuDNN v9.
 You can pass `PYTHON_VERSION=x.y` make variable to specify which Python version is to be used by Miniconda, or leave it
-unset to use the default.
+unset to use the default, as the Dockerfile uses system Python.
 
 ```bash
 make -f docker.Makefile
@@ -440,7 +433,7 @@ make -f docker.Makefile
 ### Building the Documentation
 
 To build documentation in various formats, you will need [Sphinx](http://www.sphinx-doc.org)
-and the pytorch_sphinx_theme2.
+and the `pytorch_sphinx_theme2`.
 
 Before you build the documentation locally, ensure `torch` is
 installed in your environment. For small fixes, you can install the
@@ -463,22 +456,59 @@ Run `make` to get a list of all available output formats.
 If you get a katex error run `npm install katex`.  If it persists, try
 `npm install -g katex`
 
-> [!NOTE]
-> If you installed `nodejs` with a different package manager (e.g.,
-> `conda`) then `npm` will probably install a version of `katex` that is not
-> compatible with your version of `nodejs` and doc builds will fail.
-> A combination of versions that is known to work is `node@6.13.1` and
-> `katex@0.13.18`. To install the latter with `npm` you can run
-> ```npm install -g katex@0.13.18```
-
 > [!NOTE]
 > If you see a numpy incompatibility error, run:
 > ```
 > pip install 'numpy<2'
 > ```
 
-When you make changes to the dependencies run by CI, edit the
-`.ci/docker/requirements-docs.txt` file.
+
+#### Troubleshooting CI Errors
+Your build may show errors you didn't have locally - here's how to find the errors relevant to the docs.
+
+If the build has any errors, you will see something like this on the PR:
+
+<img width="781" height="400" alt="Monosnap Update installation instructions for doc build · Pull Request #169534 · pytorch:pytorch 2025-12-18 18-22-53" src="https://github.com/user-attachments/assets/49a3dfe7-81c2-4246-852b-bc3f807e95af" />
+
+Any doc-related errors will occur in jobs that include "doc" somewhere in the title. It doesn't look like any of these jobs are relevant to our docs.
+
+
+Let's take a look anyway. Click on the job to see the logs:
+
+<img width="1187" height="668" alt="Monosnap Update installation instructions for doc build · pytorch:pytorch@7380336 2025-12-18 18-24-15" src="https://github.com/user-attachments/assets/117df543-8356-4323-8e1c-ef02a95554ba" />
+
+And we can be sure that this job does not involve docs.
+
+Looking at this build, we can see these jobs are relevant to our docs - and they didn't have any errors:
+
+<img width="777" height="395" alt="Check the docs jobs" src="https://github.com/user-attachments/assets/5d7c196b-2d40-49ad-87e3-f57de6e14a5b" />
+
+You might also see a comment on the PR like this:
+
+<img width="651" height="246" alt="PR Comment" src="https://github.com/user-attachments/assets/27e0120a-ba33-4b1c-b4a5-bf3064520586" />
+
+We can see that some of these issues are relevant to our docs.
+
+Open the logs by clicking on the `gh` link:
+
+<img width="873" height="360" alt="View Logs" src="https://github.com/user-attachments/assets/ab5b862f-8026-489c-b95e-a6cd4257e4b7" />
+
+And here we can see there is a doc-related error:
+
+<img width="1117" height="433" alt="Doc Error" src="https://github.com/user-attachments/assets/0a275921-736d-43a7-ab0f-3e8854d43280" />
+
+You can always find the relevant doc builds by going to the `Checks` tab on your PR, and scrolling down to `pull`.
+
+<img width="481" height="561" alt="checks" src="https://github.com/user-attachments/assets/eef18f2b-7134-4e2e-bd90-bcdc12800132" />
+
+You can either click through or toggle the accordion to see all of the jobs here, where you can see the docs jobs highlighted:
+
+<img width="570" height="611" alt="jobs" src="https://github.com/user-attachments/assets/f62812ca-caee-421b-863c-54f38fd28d46" />
+
+If you click through, you'll see the doc jobs at the bottom, like this:
+
+<img width="354" height="312" alt="View Docs jobs" src="https://github.com/user-attachments/assets/8fadb935-5314-4c4b-a1b5-133781754f03" />
+
 
 #### Building a PDF
 
@@ -522,7 +552,7 @@ on [our website](https://pytorch.org/get-started/previous-versions).
 
 ## Getting Started
 
-Three pointers to get you started:
+Pointers to get you started:
 - [Tutorials: get you started with understanding and using PyTorch](https://pytorch.org/tutorials/)
 - [Examples: easy to understand PyTorch code across all domains](https://github.com/pytorch/examples)
 - [The API Reference](https://pytorch.org/docs/)
diff --git a/RELEASE.md b/RELEASE.md
index 0e83a889dc239..3120931b38557 100644
--- a/RELEASE.md
+++ b/RELEASE.md
@@ -51,6 +51,7 @@ Following is the Release Compatibility Matrix for PyTorch releases:
 
 | PyTorch version | Python | C++ | Stable CUDA | Experimental CUDA | Stable ROCm |
 | --- | --- | --- | --- | --- | --- |
+| 2.11 | >=3.10, <=(3.14, 3.14t experimental) | C++17 | CUDA 12.6 (CUDNN 9.10.2.21), CUDA 12.8 (CUDNN 9.17.1.4), CUDA 13.0 (CUDNN 9.17.1.4) | -- | ROCm 7.2 |
 | 2.10 | >=3.10, <=(3.14, 3.14t experimental) | C++17 | CUDA 12.6 (CUDNN 9.10.2.21), CUDA 12.8 (CUDNN 9.10.2.21) | CUDA 13.0 (CUDNN 9.15.1.9) | ROCm 7.1 |
 | 2.9 | >=3.10, <=(3.14, 3.14t experimental) | C++17 | CUDA 12.6 (CUDNN 9.10.2.21), CUDA 12.8 (CUDNN 9.10.2.21) | CUDA 13.0 (CUDNN 9.13.0.50) | ROCm 6.4 |
 | 2.8 | >=3.9, <=3.13, (3.13t experimental) | C++17 | CUDA 12.6 (CUDNN 9.10.2.21), CUDA 12.8 (CUDNN 9.10.2.21) | CUDA 12.9 (CUDNN 9.10.2.21) | ROCm 6.4 |
@@ -67,19 +68,19 @@ Following is the Release Compatibility Matrix for PyTorch releases:
 
 ### PyTorch CUDA Support Matrix
 
-For Release 2.9 and 2.10 PyTorch Supports following CUDA Architectures:
+For Release 2.11 PyTorch Supports following CUDA Architectures:
 
 | CUDA | architectures supported for Linux x86 and Windows builds | notes |
 | --- | --- | --- |
 | 12.6.3 | Maxwell(5.0), Pascal(6.0), Volta(7.0), Turing(7.5), Ampere(8.0, 8.6), Hopper(9.0) | |
-| 12.8.1 | Volta(7.0), Turing(7.5), Ampere(8.0, 8.6), Hopper(9.0), Blackwell(10.0, 12.0)  | |
-| 13.0.0 | Turing(7.5), Ampere(8.0, 8.6), Hopper(9.0), Blackwell(10.0, 12.0+PTX) | +PTX available on linux builds only |
+| 12.8.1 | Turing(7.5), Ampere(8.0, 8.6), Hopper(9.0), Blackwell(10.0, 12.0)  | |
+| 13.0.2 | Turing(7.5), Ampere(8.0, 8.6), Hopper(9.0), Blackwell(10.0, 12.0+PTX) | +PTX available on linux builds only |
 
 | CUDA | architectures supported for Linux aarch64 builds |
 | --- | --- |
 | 12.6.3 | Ampere(8.0), Hopper(9.0) |
 | 12.8.1 | Ampere(8.0), Hopper(9.0), Blackwell(10.0, 12.0)  |
-| 13.0.0 | Ampere(8.0), Hopper(9.0), Blackwell(10.0, 11.0, 12.0+PTX) |
+| 13.0.2 | Ampere(8.0), Hopper(9.0), Blackwell(10.0, 11.0, 12.0+PTX) |
 
 ## Release Cadence
 
@@ -128,8 +129,8 @@ Following requirements need to be met prior to cutting a release branch:
 
 * Triton release branch must be created (e.g., [release/3.6.x](https://github.com/triton-lang/triton/tree/release/3.6.x)) and the Triton pin update PR must be landed (e.g., [#168096](https://github.com/pytorch/pytorch/pull/168096)) at least 1 week before the branch cut
 * Resolve all outstanding issues in the milestones that are feature work and release blocking (for example [release 2.10 milestone](https://github.com/pytorch/pytorch/milestone/57)). A report of outstanding cherry-picks can be produced by running the [github-analytics-daily workflow](https://github.com/pytorch/test-infra/blob/main/.github/workflows/github-analytics-daily.yml)
-* Validate that all new workflows have been created in the PyTorch and domain libraries included in the release. Validate it against all dimensions of release matrix, including operating systems (Linux, MacOS, Windows), Python versions as well as CPU architectures (x86 and arm) and accelerator versions (CUDA, ROCm, XPU).
-* All [viable/strict](.github/workflows/update-viablestrict.yml) jobs are green, which requires the following jobs to pass: `pull`, `trunk`, `lint`, `linux-aarch64`
+* Validate that all new workflows have been created in the PyTorch and domain libraries included in the release. Validate it against all dimensions of release matrix, including operating systems (Linux, macOS, Windows), Python versions as well as CPU architectures (x86 and arm) and accelerator versions (CUDA, ROCm, XPU).
+* All [viable/strict](.github/workflows/update-viablestrict.yml) jobs are green, which requires the following jobs to pass: `pull`, `trunk`, `lint`
 * All the nightly jobs for pytorch and domain libraries should be green. Validate this using the following HUD links:
   * [PyTorch](https://hud.pytorch.org/hud/pytorch/pytorch/nightly)
   * [TorchVision](https://hud.pytorch.org/hud/pytorch/vision/nightly)
@@ -474,7 +475,7 @@ Supported OS flavors are summarized in the table below:
 | Operating System family | Architecture | Notes |
 | --- | --- | --- |
 | Linux | aarch64, x86_64 | Wheels are manylinux2014 compatible, i.e. they should be runnable on any Linux system with glibc-2.17 or above. |
-| MacOS | arm64 | Builds should be compatible with MacOS 11 (Big Sur) or newer, but are actively tested against MacOS 14 (Sonoma). MPS support is enabled on MacOS 13 (Ventura) or later. |
+| macOS | arm64 | Builds should be compatible with macOS 11 (Big Sur) or newer, but are actively tested against macOS 14 (Sonoma). MPS support is enabled on macOS 14 (Sonoma) or later. |
 | Windows | x86_64 | Builds are compatible with Windows-10 or newer. |
 
 # Submitting Tutorials
diff --git a/SECURITY.md b/SECURITY.md
index e6ac059a4067f..5d8f845b87c99 100644
--- a/SECURITY.md
+++ b/SECURITY.md
@@ -1,12 +1,14 @@
 # Security Policy
 
  - [**Reporting a Vulnerability**](#reporting-a-vulnerability)
+ - [**Issues That Are Not Security Vulnerabilities**](#issues-that-are-not-security-vulnerabilities)
  - [**Using PyTorch Securely**](#using-pytorch-securely)
    - [Untrusted models](#untrusted-models)
    - [TorchScript models](#torchscript-models)
    - [Untrusted inputs](#untrusted-inputs)
    - [Data privacy](#data-privacy)
    - [Using distributed features](#using-distributed-features)
+- [**Backporting Security Fixes**](#security-fixes-and-old-releases)
 - [**CI/CD security principles**](#cicd-security-principles)
 ## Reporting Security Issues
 
@@ -18,14 +20,25 @@ Please report security issues using https://github.com/pytorch/pytorch/security/
 
 All reports submitted through the security advisories mechanism would **either be made public or dismissed by the team within 90 days of the submission**. If advisory has been closed on the grounds that it is not a security issue, please do not hesitate to create an [new issue](https://github.com/pytorch/pytorch/issues/new?template=bug-report.yml) as it is still likely a valid issue within the framework.
 
-**Note on crashes and out of bounds access**: PyTorch is a computational framework that performs operations on behalf of the caller. Like many low-level libraries, PyTorch generally does not validate all inputs to every function—the responsibility for providing valid arguments lies with the calling code. While crashes and out of bounds memory access should be reported as bugs, they are generally not considered security vulnerabilities in PyTorch's threat model.
-
-**Note on numerical stability**: Issues related to incorrect calculations, numerical instability, or precision loss (CWE-682: Incorrect Calculation) should be reported as regular bugs, not as security vulnerabilities.
-
 Please refer to the following page for our responsible disclosure policy, reward guidelines, and those things that should not be reported:
 
 https://www.facebook.com/whitehat
 
+## Issues That Are Not Security Vulnerabilities
+
+PyTorch is a framework that executes user-provided code, including model definitions, custom operators, and training scripts. Like many low-level computational libraries, PyTorch generally does not validate all inputs to every function - the responsibility for providing valid arguments lies with the calling code. An attacker who already has the ability to execute arbitrary code locally, or to modify files on the system, does not gain any additional capability by exploiting PyTorch. The following categories of reports should be filed as regular bugs, **not** as security vulnerabilities:
+
+- **Crashes and out-of-bounds access**: Passing invalid arguments (wrong shapes, dtypes, device placement, etc.) may result in crashes, segfaults, or out-of-bounds memory access. These are often caused by missing input validation and are bugs worth reporting as regular issues, but not security vulnerabilities — the caller is already running native code and could achieve the same effect without PyTorch.
+
+- **Numerical accuracy and stability**: Differences in numerical results between devices (CPU vs. CUDA), dtypes, platforms, or library versions are expected behavior in floating-point computing. This includes precision loss (CWE-682), non-determinism across runs, and results that diverge from a reference implementation.
+
+- **Internal or "unsafe" functions**: Some PyTorch functions are intentionally low-level or unchecked (e.g., functions prefixed with `_`, internal C++ APIs, or APIs documented as unsafe). These exist for performance or flexibility and are not intended to be hardened against adversarial callers. If an attacker can invoke such functions with arbitrary arguments, they likely have local execution access already.
+
+- **Denial of service via resource consumption**: PyTorch is designed to run computationally intensive workloads that fully utilize CPU, GPU, and memory. Crafting inputs that cause high resource consumption or long-running operations is trivial and expected — this is not a security vulnerability.
+
+- **Local filesystem and cache trust**: PyTorch heavily uses local caching for performance (e.g., compiled kernel artifacts, Triton cache). These cache should be located in local user-only accessible locations and those are trusted by design. If an attacker has write access to the local filesystem, they can already execute arbitrary code — poisoning a PyTorch cache does not grant any additional capability.
+
+If your security advisory is closed because it falls into one of these categories, please don't be discouraged — these are still valuable reports. We encourage you to re-file them as a [regular issue](https://github.com/pytorch/pytorch/issues/new?template=bug-report.yml) so they can be tracked and fixed as bugs.
 
 ## Using PyTorch Securely
 **PyTorch models are programs**, so treat its security seriously -- running untrusted models is equivalent to running untrusted code. In general we recommend that model weights and the python code for the model are distributed independently. That said, be careful about where you get the python code from and who wrote it (preferentially check for a provenance or checksums, do not run any pip installed package).
@@ -75,6 +88,10 @@ PyTorch can be used for distributed computing, and as such there is a `torch.dis
 
 For performance reasons, none of the PyTorch Distributed primitives (including c10d, RPC, and TCPStore) include any authorization protocol and will send messages unencrypted. They accept connections from anywhere, and execute the workload sent without performing any checks. Therefore, if you run a PyTorch Distributed program on your network, anybody with access to the network can execute arbitrary code with the privileges of the user running PyTorch.
 
+## Backporting Security Fixes
+
+Security fixes are only applied to the current release. Fixes are never backported to older versions of PyTorch. The only security concern applicable to previously published releases is tampering with the published binaries themselves (e.g., unauthorized modification of packages on PyPI or download indexes). If you believe a published binary has been compromised, please report it through the [security advisory process](#reporting-security-issues).
+
 ## CI/CD security principles
 _Audience_: Contributors and reviewers, especially if modifying the workflow files/build system.
 
diff --git a/WORKSPACE b/WORKSPACE
deleted file mode 100644
index c7d6307a9a394..0000000000000
--- a/WORKSPACE
+++ /dev/null
@@ -1,369 +0,0 @@
-workspace(name = "pytorch")
-
-load("@bazel_tools//tools/build_defs/repo:http.bzl", "http_archive")
-load("//tools/rules:workspace.bzl", "new_patched_local_repository")
-
-http_archive(
-    name = "rules_cc",
-    patches = [
-        "//:tools/rules_cc/cuda_support.patch",
-    ],
-    strip_prefix = "rules_cc-40548a2974f1aea06215272d9c2b47a14a24e556",
-    urls = [
-        "https://mirror.bazel.build/github.com/bazelbuild/rules_cc/archive/40548a2974f1aea06215272d9c2b47a14a24e556.tar.gz",
-        "https://github.com/bazelbuild/rules_cc/archive/40548a2974f1aea06215272d9c2b47a14a24e556.tar.gz",
-    ],
-)
-
-http_archive(
-    name = "rules_cuda",
-    strip_prefix = "runtime-b1c7cce21ba4661c17ac72421c6a0e2015e7bef3/third_party/rules_cuda",
-    urls = ["https://github.com/tensorflow/runtime/archive/b1c7cce21ba4661c17ac72421c6a0e2015e7bef3.tar.gz"],
-)
-
-http_archive(
-    name = "platforms",
-    urls = [
-        "https://mirror.bazel.build/github.com/bazelbuild/platforms/releases/download/0.0.10/platforms-0.0.10.tar.gz",
-        # TODO Fix bazel linter to support hashes for release tarballs.
-        # "https://github.com/bazelbuild/platforms/releases/download/0.0.10/platforms-0.0.10.tar.gz",
-    ],
-    # sha256 = "218efe8ee736d26a3572663b374a253c012b716d8af0c07e842e82f238a0a7ee",
-)
-
-load("@rules_cuda//cuda:dependencies.bzl", "rules_cuda_dependencies")
-
-rules_cuda_dependencies(with_rules_cc = False)
-
-load("@rules_cc//cc:repositories.bzl", "rules_cc_toolchains")
-
-rules_cc_toolchains()
-
-http_archive(
-    name = "bazel_skylib",
-    urls = [
-        "https://github.com/bazelbuild/bazel-skylib/releases/download/1.0.2/bazel-skylib-1.0.2.tar.gz",
-    ],
-)
-
-http_archive(
-    name = "pybind11_bazel",
-    strip_prefix = "pybind11_bazel-b162c7c88a253e3f6b673df0c621aca27596ce6b",
-    urls = ["https://github.com/pybind/pybind11_bazel/archive/b162c7c88a253e3f6b673df0c621aca27596ce6b.zip"],
-)
-
-new_local_repository(
-    name = "pybind11",
-    build_file = "@pybind11_bazel//:pybind11.BUILD",
-    path = "third_party/pybind11",
-)
-
-http_archive(
-    name = "com_github_glog",
-    build_file_content = """
-licenses(['notice'])
-
-load(':bazel/glog.bzl', 'glog_library')
-# TODO: figure out why enabling gflags leads to SIGSEV on the logging init
-glog_library(with_gflags=0)
-    """,
-    strip_prefix = "glog-0.4.0",
-    urls = [
-        "https://github.com/google/glog/archive/v0.4.0.tar.gz",
-    ],
-)
-
-http_archive(
-    name = "com_github_gflags_gflags",
-    strip_prefix = "gflags-2.2.2",
-    urls = [
-        "https://github.com/gflags/gflags/archive/v2.2.2.tar.gz",
-    ],
-)
-
-new_local_repository(
-    name = "gloo",
-    build_file = "//third_party:gloo.BUILD",
-    path = "third_party/gloo",
-)
-
-new_local_repository(
-    name = "onnx",
-    build_file = "//third_party:onnx.BUILD",
-    path = "third_party/onnx",
-)
-
-local_repository(
-    name = "com_google_protobuf",
-    path = "third_party/protobuf",
-)
-
-new_local_repository(
-    name = "eigen",
-    build_file = "//third_party:eigen.BUILD",
-    path = "third_party/eigen",
-)
-
-new_local_repository(
-    name = "cutlass",
-    build_file = "//third_party:cutlass.BUILD",
-    path = "third_party/cutlass",
-)
-
-new_local_repository(
-    name = "fbgemm",
-    build_file = "//third_party:fbgemm/BUILD.bazel",
-    path = "third_party/fbgemm",
-    repo_mapping = {"@cpuinfo": "@org_pytorch_cpuinfo"},
-)
-
-new_local_repository(
-    name = "ideep",
-    build_file = "//third_party:ideep.BUILD",
-    path = "third_party/ideep",
-)
-
-new_local_repository(
-    name = "mkl_dnn",
-    build_file = "//third_party:mkl-dnn.BUILD",
-    path = "third_party/ideep/mkl-dnn",
-)
-
-new_local_repository(
-    name = "org_pytorch_cpuinfo",
-    build_file = "//third_party:cpuinfo/BUILD.bazel",
-    path = "third_party/cpuinfo",
-)
-
-new_local_repository(
-    name = "asmjit",
-    build_file = "//third_party:fbgemm/external/asmjit.BUILD",
-    path = "third_party/fbgemm/external/asmjit",
-)
-
-new_local_repository(
-    name = "sleef",
-    build_file = "//third_party:sleef.BUILD",
-    path = "third_party/sleef",
-)
-
-new_local_repository(
-    name = "fmt",
-    build_file = "//third_party:fmt.BUILD",
-    path = "third_party/fmt",
-)
-
-new_local_repository(
-    name = "kineto",
-    build_file = "//third_party:kineto.BUILD",
-    path = "third_party/kineto",
-)
-
-new_local_repository(
-    name = "opentelemetry-cpp",
-    build_file = "//third_party::opentelemetry-cpp.BUILD",
-    path = "third_party/opentelemetry-cpp",
-)
-
-new_local_repository(
-    name = "cpp-httplib",
-    build_file = "//third_party:cpp-httplib.BUILD",
-    path = "third_party/cpp-httplib",
-)
-
-new_local_repository(
-    name = "nlohmann",
-    build_file = "//third_party:nlohmann.BUILD",
-    path = "third_party/nlohmann",
-)
-
-new_local_repository(
-    name = "moodycamel",
-    build_file = "//third_party:moodycamel.BUILD",
-    path = "third_party/concurrentqueue",
-)
-
-new_local_repository(
-    name = "tensorpipe",
-    build_file = "//third_party:tensorpipe.BUILD",
-    path = "third_party/tensorpipe",
-)
-
-http_archive(
-    name = "mkl",
-    build_file = "//third_party:mkl.BUILD",
-    sha256 = "59154b30dd74561e90d547f9a3af26c75b6f4546210888f09c9d4db8f4bf9d4c",
-    strip_prefix = "lib",
-    urls = [
-        "https://anaconda.org/anaconda/mkl/2020.0/download/linux-64/mkl-2020.0-166.tar.bz2",
-    ],
-)
-
-http_archive(
-    name = "mkl_headers",
-    build_file = "//third_party:mkl_headers.BUILD",
-    sha256 = "2af3494a4bebe5ddccfdc43bacc80fcd78d14c1954b81d2c8e3d73b55527af90",
-    urls = [
-        "https://anaconda.org/anaconda/mkl-include/2020.0/download/linux-64/mkl-include-2020.0-166.tar.bz2",
-    ],
-)
-
-http_archive(
-    name = "rules_python",
-    # TODO Fix bazel linter to support hashes for release tarballs.
-    #
-    # sha256 = "94750828b18044533e98a129003b6a68001204038dc4749f40b195b24c38f49f",
-    strip_prefix = "rules_python-0.21.0",
-    url = "https://github.com/bazelbuild/rules_python/releases/download/0.21.0/rules_python-0.21.0.tar.gz",
-)
-
-load("@rules_python//python:repositories.bzl", "py_repositories")
-
-py_repositories()
-
-load("@rules_python//python:repositories.bzl", "python_register_toolchains")
-
-python_register_toolchains(
-    name = "python3_10",
-    python_version = "3.10",
-)
-
-load("@python3_10//:defs.bzl", "interpreter")
-load("@rules_python//python:pip.bzl", "pip_parse")
-
-pip_parse(
-    name = "pip_deps",
-    python_interpreter_target = interpreter,
-    requirements_lock = "//:tools/build/bazel/requirements.txt",
-)
-
-load("@pip_deps//:requirements.bzl", "install_deps")
-
-install_deps()
-
-load("@pybind11_bazel//:python_configure.bzl", "python_configure")
-
-python_configure(
-    name = "local_config_python",
-    python_interpreter_target = interpreter,
-)
-
-load("@com_google_protobuf//:protobuf_deps.bzl", "protobuf_deps")
-
-protobuf_deps()
-
-new_local_repository(
-    name = "cuda",
-    build_file = "@//third_party:cuda.BUILD",
-    path = "/usr/local/cuda",
-)
-
-new_local_repository(
-    name = "cudnn",
-    build_file = "@//third_party:cudnn.BUILD",
-    path = "/usr/local/cuda",
-)
-
-new_local_repository(
-    name = "cudnn_frontend",
-    build_file = "@//third_party:cudnn_frontend.BUILD",
-    path = "third_party/cudnn_frontend/",
-)
-
-local_repository(
-    name = "com_github_google_flatbuffers",
-    path = "third_party/flatbuffers",
-)
-
-local_repository(
-    name = "google_benchmark",
-    path = "third_party/benchmark",
-)
-
-local_repository(
-    name = "com_google_googletest",
-    path = "third_party/googletest",
-)
-
-local_repository(
-    name = "pthreadpool",
-    path = "third_party/pthreadpool",
-    repo_mapping = {"@com_google_benchmark": "@google_benchmark"},
-)
-
-local_repository(
-    name = "FXdiv",
-    path = "third_party/FXdiv",
-    repo_mapping = {"@com_google_benchmark": "@google_benchmark"},
-)
-
-local_repository(
-    name = "XNNPACK",
-    path = "third_party/XNNPACK",
-    repo_mapping = {"@com_google_benchmark": "@google_benchmark"},
-)
-
-local_repository(
-    name = "gemmlowp",
-    path = "third_party/gemmlowp/gemmlowp",
-)
-
-local_repository(
-    name = "kleidiai",
-    path = "third_party/kleidiai",
-    repo_mapping = {"@com_google_googletest": "@com_google_benchmark"},
-)
-
-### Unused repos start
-
-# `unused` repos are defined to hide bazel files from submodules of submodules.
-# This allows us to run `bazel build //...` and not worry about the submodules madness.
-# Otherwise everything traverses recursively and a lot of submodules of submodules have
-# they own bazel build files.
-
-local_repository(
-    name = "unused_tensorpipe_googletest",
-    path = "third_party/tensorpipe/third_party/googletest",
-)
-
-local_repository(
-    name = "unused_fbgemm",
-    path = "third_party/fbgemm",
-)
-
-local_repository(
-    name = "unused_ftm_bazel",
-    path = "third_party/fmt/support/bazel",
-)
-
-local_repository(
-    name = "unused_kineto_fmt_bazel",
-    path = "third_party/kineto/libkineto/third_party/fmt/support/bazel",
-)
-
-local_repository(
-    name = "unused_kineto_dynolog_googletest",
-    path = "third_party/kineto/libkineto/third_party/dynolog/third_party/googletest",
-)
-
-local_repository(
-    name = "unused_kineto_dynolog_gflags",
-    path = "third_party/kineto/libkineto/third_party/dynolog/third_party/gflags",
-)
-
-local_repository(
-    name = "unused_kineto_dynolog_glog",
-    path = "third_party/kineto/libkineto/third_party/dynolog/third_party/glog",
-)
-
-local_repository(
-    name = "unused_kineto_googletest",
-    path = "third_party/kineto/libkineto/third_party/googletest",
-)
-
-local_repository(
-    name = "unused_onnx_benchmark",
-    path = "third_party/onnx/third_party/benchmark",
-)
-
-### Unused repos end
diff --git a/android/pytorch_android/CMakeLists.txt b/android/pytorch_android/CMakeLists.txt
index 653256c256c25..4ee3011683011 100644
--- a/android/pytorch_android/CMakeLists.txt
+++ b/android/pytorch_android/CMakeLists.txt
@@ -1,4 +1,4 @@
-cmake_minimum_required(VERSION 3.5)
+cmake_minimum_required(VERSION 3.10)
 option(BUILD_LITE_INTERPRETER "Master flag to build pytorch_jni_lite" ON)
 message(
   STATUS
@@ -14,7 +14,7 @@ endif()
 
 include(GNUInstallDirs)
 
-set(CMAKE_CXX_STANDARD 17 CACHE STRING "The C++ standard whose features are requested to build this target.")
+set(CMAKE_CXX_STANDARD 20 CACHE STRING "The C++ standard whose features are requested to build this target.")
 set(CMAKE_VERBOSE_MAKEFILE ON)
 message(STATUS "ANDROID_STL:${ANDROID_STL}")
 
diff --git a/android/pytorch_android/generate_test_torchscripts.py b/android/pytorch_android/generate_test_torchscripts.py
index 507222ec17b03..13507ce1afac4 100644
--- a/android/pytorch_android/generate_test_torchscripts.py
+++ b/android/pytorch_android/generate_test_torchscripts.py
@@ -1,5 +1,3 @@
-from typing import Optional
-
 import torch
 from torch import Tensor
 
@@ -86,11 +84,11 @@ def tupleIntSumReturnTuple(
         return (input, sum)
 
     @torch.jit.script_method
-    def optionalIntIsNone(self, input: Optional[int]) -> bool:
+    def optionalIntIsNone(self, input: int | None) -> bool:
         return input is None
 
     @torch.jit.script_method
-    def intEq0None(self, input: int) -> Optional[int]:
+    def intEq0None(self, input: int) -> int | None:
         if input == 0:
             return None
         return input
diff --git a/android/pytorch_android_torchvision/CMakeLists.txt b/android/pytorch_android_torchvision/CMakeLists.txt
index 3e23fdc1a1e5b..b49b779a18b1d 100644
--- a/android/pytorch_android_torchvision/CMakeLists.txt
+++ b/android/pytorch_android_torchvision/CMakeLists.txt
@@ -1,6 +1,6 @@
-cmake_minimum_required(VERSION 3.5)
+cmake_minimum_required(VERSION 3.10)
 project(pytorch_vision_jni CXX)
-set(CMAKE_CXX_STANDARD 17 CACHE STRING "The C++ standard whose features are requested to build this target.")
+set(CMAKE_CXX_STANDARD 20 CACHE STRING "The C++ standard whose features are requested to build this target.")
 set(CMAKE_VERBOSE_MAKEFILE ON)
 
 set(pytorch_vision_cpp_DIR ${CMAKE_CURRENT_LIST_DIR}/src/main/cpp)
diff --git a/aten/src/ATen/CMakeLists.txt b/aten/src/ATen/CMakeLists.txt
index f1ac6246f4d56..fc5d69381e648 100644
--- a/aten/src/ATen/CMakeLists.txt
+++ b/aten/src/ATen/CMakeLists.txt
@@ -74,8 +74,8 @@ if(NOT BUILD_LITE_INTERPRETER)
 endif()
 EXCLUDE(ATen_CORE_SRCS "${ATen_CORE_SRCS}" ${ATen_CORE_TEST_SRCS})
 
-file(GLOB base_h "*.h" "detail/*.h" "cpu/*.h" "cpu/vec/vec512/*.h" "cpu/vec/vec128/*.h" "cpu/vec/vec256/*.h" "cpu/vec/vec256/vsx/*.h" "cpu/vec/vec256/zarch/*.h" "cpu/vec/sve/*.h" "cpu/vec/*.h" "quantized/*.h" "functorch/*.h")
-file(GLOB base_cpp "*.cpp" "detail/*.cpp" "cpu/*.cpp" "functorch/*.cpp")
+file(GLOB base_h "*.h" "detail/*.h" "cpu/*.h" "cpu/vec/vec512/*.h" "cpu/vec/vec128/*.h" "cpu/vec/vec256/*.h" "cpu/vec/vec256/vsx/*.h" "cpu/vec/vec256/zarch/*.h" "cpu/vec/sve/*.h" "cpu/vec/*.h" "quantized/*.h" "functorch/*.h" "accelerator/*.h")
+file(GLOB base_cpp "*.cpp" "detail/*.cpp" "cpu/*.cpp" "functorch/*.cpp" "accelerator/*.cpp")
 file(GLOB cuda_h "cuda/*.h" "cuda/detail/*.h" "cuda/*.cuh" "cuda/detail/*.cuh" "cuda/tunable/*.cuh" "cuda/tunable/*.h")
 file(GLOB cuda_cpp "cuda/*.cpp" "cuda/detail/*.cpp" "cuda/tunable/*.cpp")
 file(GLOB cuda_nvrtc_stub_h "cuda/nvrtc_stub/*.h")
@@ -202,9 +202,36 @@ file(GLOB native_flash_attn_api_cpp "native/transformers/cuda/flash_attn/flash_a
 file(GLOB flash_attention_hip_hip "native/transformers/hip/flash_attn/*.hip")
 # if USE_FLASH_ATTENTION is set, ensure CK instances get generated
 if(USE_FLASH_ATTENTION)
-  if("$ENV{USE_CK_FLASH_ATTENTION}" STREQUAL "1")
-    message(STATUS "USE_CK_FLASH_ATTENTION is being deprecated. Please use USE_ROCM_CK_SDPA instead")
-    caffe2_update_option(USE_ROCM_CK_SDPA ON)
+  # Now building CK by default
+  if(USE_ROCM)
+    if(DEFINED ENV{USE_ROCM_CK_SDPA})
+        if(ENV{USE_ROCM_CK_SDPA} EQUAL 1)
+          caffe2_update_option(USE_ROCM_CK_SDPA ON)
+        endif()
+    else()
+      caffe2_update_option(USE_ROCM_CK_SDPA ON)
+    endif()
+    # CK SDPA only supports gfx942/gfx950. Compute the intersection of the
+    # requested PYTORCH_ROCM_ARCH with the supported archs and build the
+    # ck_sdpa target for that subset only. This keeps CK SDPA working when
+    # PYTORCH_ROCM_ARCH mixes supported and unsupported archs: e.g.
+    # "gfx942;gfx1010" builds CK SDPA for gfx942 and skips gfx1010. If none of
+    # the supported archs are requested, disable the target entirely so it does
+    # not end up with an empty arch list and error out.
+    if(USE_ROCM_CK_SDPA)
+      set(CK_SDPA_ARCHS "")
+      foreach(ARCH ${PYTORCH_ROCM_ARCH})
+        if("${ARCH}" STREQUAL "gfx942" OR "${ARCH}" STREQUAL "gfx950")
+          list(APPEND CK_SDPA_ARCHS "${ARCH}")
+        endif()
+      endforeach()
+      if(CK_SDPA_ARCHS)
+        message(STATUS "USE_ROCM_CK_SDPA enabled for supported archs: ${CK_SDPA_ARCHS} (from PYTORCH_ROCM_ARCH: ${PYTORCH_ROCM_ARCH})")
+      else()
+        message(STATUS "USE_ROCM_CK_SDPA disabled: PYTORCH_ROCM_ARCH (${PYTORCH_ROCM_ARCH}) has no supported arch (gfx942/gfx950)")
+        caffe2_update_option(USE_ROCM_CK_SDPA OFF)
+      endif()
+    endif()
   endif()
   if(USE_ROCM_CK_SDPA)
     if(DEFINED ENV{PYTORCH_ROCM_ARCH})
@@ -289,9 +316,31 @@ if(USE_FLASH_ATTENTION)
          "native/transformers/hip/flash_attn/ck/fav_v3/*.hip")
 
     set_source_files_properties(${ck_sdpa_sources_hip} PROPERTIES HIP_SOURCE_PROPERTY_FORMAT 1)
+
+    # hip_add_library() inherits the global HIP_CLANG_FLAGS, which carry an
+    # --offload-arch=<arch> entry for every arch in PYTORCH_ROCM_ARCH (see
+    # cmake/Dependencies.cmake), including archs CK SDPA does not support. Strip
+    # those inherited offload-arch flags for this target and re-add only the
+    # supported archs (CK_SDPA_ARCHS), so ck_sdpa compiles for gfx942/gfx950
+    # only, regardless of what else PYTORCH_ROCM_ARCH requests. Restore the
+    # flags afterwards so other HIP targets keep building for the full arch list.
+    set(_ck_sdpa_saved_hip_clang_flags "${HIP_CLANG_FLAGS}")
+    set(_ck_sdpa_hip_clang_flags "")
+    foreach(_flag ${HIP_CLANG_FLAGS})
+      if(NOT "${_flag}" MATCHES "^--offload-arch=")
+        list(APPEND _ck_sdpa_hip_clang_flags "${_flag}")
+      endif()
+    endforeach()
+    foreach(_arch ${CK_SDPA_ARCHS})
+      list(APPEND _ck_sdpa_hip_clang_flags "--offload-arch=${_arch}")
+    endforeach()
+    set(HIP_CLANG_FLAGS "${_ck_sdpa_hip_clang_flags}")
+
     hip_add_library(ck_sdpa STATIC
                     ${ck_sdpa_sources_hip}
                     HIPCC_OPTIONS ${HIP_HCC_FLAGS} ${CK_SDPA_EXTRA_HIPCC_FLAGS})
+
+    set(HIP_CLANG_FLAGS "${_ck_sdpa_saved_hip_clang_flags}")
     set_target_properties(ck_sdpa PROPERTIES POSITION_INDEPENDENT_CODE ON)
     target_compile_definitions(ck_sdpa PUBLIC ${CK_SDPA_EXTRA_HIPCC_OPTIONS})
     target_compile_definitions(ck_sdpa PRIVATE AITER_EMBEDDED_HSA_HEADER="aiter_embedded_hsa.h")
@@ -407,6 +456,21 @@ IF(USE_MSLK)
     # Add MSLK include directories for gemm_torch.h
     list(APPEND ATen_CUDA_INCLUDE ${MSLK_INCLUDE})
   elseif(USE_ROCM)
+    # MSLK only supports gfx942/gfx950. Auto-disable when none of those archs
+    # are in PYTORCH_ROCM_ARCH; otherwise the mslk target ends up with an empty
+    # HIP_ARCHITECTURES and CMake errors out.
+    set(_mslk_hip_arches "")
+    foreach(ARCH gfx942 gfx950)
+      if(${ARCH} IN_LIST PYTORCH_ROCM_ARCH)
+        list(APPEND _mslk_hip_arches ${ARCH})
+      endif()
+    endforeach()
+  endif()
+  if(USE_MSLK AND USE_ROCM AND _mslk_hip_arches STREQUAL "")
+    message(STATUS "USE_MSLK disabled: PYTORCH_ROCM_ARCH (${PYTORCH_ROCM_ARCH}) has no MI3xx arch (gfx942/gfx950)")
+    caffe2_update_option(USE_MSLK OFF)
+  endif()
+  if(USE_MSLK AND USE_ROCM)
     # Path to vendored CK
     set(MSLK_CK_DIR "${MSLK_THIRD_PARTY}/composable_kernel")
     rocm_generate_ck_conf(
@@ -430,22 +494,11 @@ IF(USE_MSLK)
         list(PREPEND MSLK_EXTRA_HIPCC_FLAGS -mllvm -amdgpu-coerce-illegal-types=1)
       endif()
 
-    # Only compile for gfx942 and gfx950.
-    set(HIP_CLANG_FLAGS_ORIGINAL ${HIP_CLANG_FLAGS})
-    string(REGEX REPLACE "--offload-arch=[^ ]*" "" FILTERED_HIP_CLANG_FLAGS "${HIP_CLANG_FLAGS}")
-    foreach(ARCH gfx942 gfx950)
-      if(${ARCH} IN_LIST PYTORCH_ROCM_ARCH)
-        list(APPEND FILTERED_HIP_CLANG_FLAGS --offload-arch=${ARCH})
-      endif()
-    endforeach()
-    set(HIP_CLANG_FLAGS ${FILTERED_HIP_CLANG_FLAGS})
-
-    hip_add_library(
-      mslk STATIC
-      ${mslk_native_rocm_hip}
-      HIPCC_OPTIONS ${HIP_HCC_FLAGS} ${MSLK_EXTRA_HIPCC_FLAGS})
-    set(HIP_CLANG_FLAGS ${HIP_CLANG_FLAGS_ORIGINAL})
-    set_target_properties(mslk PROPERTIES POSITION_INDEPENDENT_CODE ON)
+    add_library(mslk STATIC ${mslk_native_rocm_hip})
+    set_target_properties(mslk PROPERTIES
+      POSITION_INDEPENDENT_CODE ON
+      HIP_ARCHITECTURES "${_mslk_hip_arches}")
+    target_compile_options(mslk PRIVATE $<$<COMPILE_LANGUAGE:HIP>:${MSLK_EXTRA_HIP_FLAGS}> ${HIP_CXX_FLAGS})
     target_compile_definitions(mslk PRIVATE MSLK_NO_EXTENDED_SHAPES)
 
     target_include_directories(mslk PRIVATE
@@ -462,7 +515,9 @@ IF(USE_MSLK)
   endif()
 
   # Ensure torch_cpu is built before mslk to avoid race conditions
-  add_dependencies(mslk torch_cpu)
+  if(TARGET mslk)
+    add_dependencies(mslk torch_cpu)
+  endif()
 endif()
 
 # XNNPACK
@@ -627,11 +682,33 @@ if(USE_ROCM)
     ${native_quantized_hip_hip}
     ${native_transformers_hip_hip} ${native_transformers_src_hip_hip}
   )
+  file(GLOB native_hip_bgemm CONFIGURE_DEPENDS "native/hip/bgemm_kernels/*.hip")
+  file(GLOB native_hip_ck CONFIGURE_DEPENDS "native/hip/ck*.hip")
   if(NOT USE_ROCM_CK_GEMM)
-    file(GLOB native_hip_bgemm "native/hip/bgemm_kernels/*.hip")
-    file(GLOB native_hip_ck "native/hip/ck*.hip")
     exclude(ATen_HIP_SRCS "${ATen_HIP_SRCS}"
       ${native_hip_bgemm} ${native_hip_ck})
+  else()
+    # composable_kernel has no gfx1250 support yet, so its CK GEMM kernels fail to
+    # compile for that arch. Build CK GEMM as a separate library with gfx1250
+    # filtered out of HIP_ARCHITECTURES so every other requested arch still gets
+    # CK GEMM. Remove this filtering once the CK submodule supports gfx1250.
+    exclude(ATen_HIP_SRCS "${ATen_HIP_SRCS}"
+      ${native_hip_bgemm} ${native_hip_ck})
+    set(_ck_gemm_hip_arches ${PYTORCH_ROCM_ARCH})
+    list(REMOVE_ITEM _ck_gemm_hip_arches "gfx1250")
+    set_source_files_properties(${native_hip_bgemm} ${native_hip_ck} PROPERTIES LANGUAGE HIP)
+    add_library(ck_gemm STATIC ${native_hip_bgemm} ${native_hip_ck})
+    set_target_properties(ck_gemm PROPERTIES
+      POSITION_INDEPENDENT_CODE ON
+      HIP_ARCHITECTURES "${_ck_gemm_hip_arches}")
+    target_compile_options(ck_gemm PRIVATE ${HIP_CXX_FLAGS})
+    target_include_directories(ck_gemm PRIVATE
+      ${CMAKE_CURRENT_SOURCE_DIR}/hip
+      ${CMAKE_CURRENT_SOURCE_DIR}/../../../third_party/composable_kernel/include
+      ${CMAKE_CURRENT_SOURCE_DIR}/../../../third_party/composable_kernel/library/include
+      ${CMAKE_CURRENT_BINARY_DIR}/composable_kernel)
+    # Ensure torch_cpu (and its codegen) is built before ck_gemm to avoid race conditions
+    add_dependencies(ck_gemm torch_cpu)
   endif()
 
   # TODO: Codegen separate files for HIP and use those (s/cuda_generated_sources/hip_generated_sources)
@@ -764,6 +841,7 @@ if(NOT EMSCRIPTEN AND NOT INTERN_BUILD_MOBILE)
     set(SLEEF_BUILD_GNUABI_LIBS OFF CACHE BOOL "Don't build sleef gnuabi libs" FORCE)
     set(SLEEF_BUILD_TESTS OFF CACHE BOOL "Don't build sleef tests" FORCE)
     set(SLEEF_BUILD_SCALAR_LIB OFF CACHE BOOL "libsleefscalar will be built." FORCE)
+    set(SLEEF_DISABLE_OPENMP ON CACHE BOOL "OpenMP only used by sleef DFT and tests, both disabled" FORCE)
     if(WIN32)
       set(SLEEF_BUILD_WITH_LIBM OFF CACHE BOOL "Don't build sleef with libm for Windows." FORCE)
     endif()
@@ -891,12 +969,23 @@ if(USE_MPS)
             metal_to_air(${SHADER} ${TGT_BASIC} "-std=metal3.1")
         endforeach()
         air_to_metallib(kernels_basic.metallib ${AIR_BASIC})
+        set(METALLIB_DEPS kernels_basic.metallib)
+        if(CAN_COMPILE_METAL_40)
+            foreach(SHADER ${native_mps_metal} ${native_mps_sparse_metal})
+                cmake_path(GET SHADER STEM TGT_STEM)
+                string(CONCAT TGT_40 ${TGT_STEM} "_40.air")
+                list(APPEND AIR_40 ${TGT_40})
+                metal_to_air(${SHADER} ${TGT_40} "-std=metal4.0")
+            endforeach()
+            air_to_metallib(kernels_40.metallib ${AIR_40})
+            list(APPEND METALLIB_DEPS kernels_40.metallib)
+        endif()
         add_custom_command(
                           COMMAND echo "// $$(date)" > metallib_dummy.cpp
-                          DEPENDS kernels_basic.metallib
+                          DEPENDS ${METALLIB_DEPS}
                           OUTPUT metallib_dummy.cpp
                           COMMENT "Updating metallibs timestamp")
-        add_custom_target(metallibs DEPENDS kernels_basic.metallib metallib_dummy.cpp)
+        add_custom_target(metallibs DEPENDS ${METALLIB_DEPS} metallib_dummy.cpp)
     else()
         file(MAKE_DIRECTORY "${CMAKE_CURRENT_BINARY_DIR}/native/mps")
         foreach(SHADER ${native_mps_metal} ${native_mps_sparse_metal})
@@ -924,7 +1013,7 @@ install(FILES "${CMAKE_CURRENT_BINARY_DIR}/cmake-exports/ATenConfig.cmake"
 
 set(INSTALL_HEADERS ${base_h} ${ATen_CORE_HEADERS} ${native_nested_h} ${ATen_TRANSFORMER_HEADERS})
 if(NOT INTERN_BUILD_MOBILE)
-  list(APPEND INSTALL_HEADERS ${native_h} ${native_cpu_h} ${native_ao_sparse_h} ${native_quantized_h} ${cuda_h} ${native_cuda_h} ${native_hip_h} ${native_mtia_h} ${cudnn_h} ${hip_h} ${mtia_h} ${xpu_h} ${mps_h} ${native_kleidiai_h} ${native_mps_h} ${native_utils_h} ${miopen_h} ${mkldnn_xpu_h})
+  list(APPEND INSTALL_HEADERS ${native_h} ${native_cpu_h} ${native_ao_sparse_h} ${native_quantized_h} ${cuda_h} ${native_cuda_h} ${native_hip_h} ${native_mtia_h} ${cudnn_h} ${hip_h} ${mtia_h} ${xpu_h} ${mps_h} ${native_kleidiai_h} ${native_mps_h} ${native_utils_h} ${miopen_h})
   # Metal
   if(USE_PYTORCH_METAL_EXPORT)
     # Add files needed from exporting metal models(optimized_for_mobile)
diff --git a/aten/src/ATen/Context.cpp b/aten/src/ATen/Context.cpp
index c342590b58c42..d0c1d6eb88887 100644
--- a/aten/src/ATen/Context.cpp
+++ b/aten/src/ATen/Context.cpp
@@ -505,7 +505,7 @@ at::BlasBackend Context::blasPreferredBackend() {
 bool Context::ckSupported() {
 #ifdef USE_ROCM
   static const std::vector<std::string> supported_archs = {
-    "gfx90a", "gfx942", "gfx950"
+    "gfx942", "gfx950"
   };
   for (auto index : c10::irange(detail::getCUDAHooks().deviceCount())) {
     if(!detail::getCUDAHooks().isGPUArch(supported_archs, index)) {
diff --git a/aten/src/ATen/Context.h b/aten/src/ATen/Context.h
index 9f24ea32245ec..de6a7dda66d73 100644
--- a/aten/src/ATen/Context.h
+++ b/aten/src/ATen/Context.h
@@ -502,7 +502,11 @@ class TORCH_API Context {
       (c10::utils::check_env("TORCH_BLAS_PREFER_CUBLASLT") == true ||
        c10::utils::check_env("TORCH_BLAS_PREFER_HIPBLASLT") == true) // alias
       ? at::BlasBackend::Cublaslt
-      : at::BlasBackend::Default;
+      : ((c10::utils::check_env("TORCH_BLAS_PREFER_CUBLASLT") == false ||
+          c10::utils::check_env("TORCH_BLAS_PREFER_HIPBLASLT") ==
+              false) // alias
+             ? at::BlasBackend::Cublas
+             : at::BlasBackend::Default);
   at::ROCmFABackend rocm_fa_preferred_backend =
       c10::utils::check_env("TORCH_ROCM_FA_PREFER_CK") == true
       ? at::ROCmFABackend::Ck
diff --git a/aten/src/ATen/DeviceAccelerator.cpp b/aten/src/ATen/DeviceAccelerator.cpp
index efab9ec9c5927..fed2e6c4febe8 100644
--- a/aten/src/ATen/DeviceAccelerator.cpp
+++ b/aten/src/ATen/DeviceAccelerator.cpp
@@ -1,4 +1,5 @@
 #include <ATen/Context.h>
+#include <ATen/core/CachingHostAllocator.h>
 #include <ATen/DeviceAccelerator.h>
 #include <c10/core/impl/VirtualGuardImpl.h>
 
@@ -136,6 +137,11 @@ c10::DeviceCapability getDeviceCapability(c10::DeviceIndex device_index) {
   c10::impl::VirtualGuardImpl impl(device_type);
   return impl.getDeviceCapability({device_type, device_index});
 }
+
+void emptyHostCache() {
+  const auto device_type = getAccelerator(true).value();
+  at::getHostAllocator(device_type)->empty_cache();
+}
 // NOLINTEND(bugprone-unchecked-optional-access)
 
 } // namespace at::accelerator
diff --git a/aten/src/ATen/DeviceAccelerator.h b/aten/src/ATen/DeviceAccelerator.h
index d24b42ca459e7..90fad38a5faca 100644
--- a/aten/src/ATen/DeviceAccelerator.h
+++ b/aten/src/ATen/DeviceAccelerator.h
@@ -5,7 +5,7 @@
 #include <c10/core/DeviceType.h>
 #include <c10/macros/Macros.h>
 
-#include <ATen/detail/MTIAHooksInterface.h>
+#include <ATen/accelerator/Graph.h>
 #include <optional>
 
 namespace at::accelerator {
@@ -78,11 +78,19 @@ TORCH_API c10::DeviceIndex maybeExchangeDevice(c10::DeviceIndex device_index);
 TORCH_API c10::DeviceCapability getDeviceCapability(
     c10::DeviceIndex device_index);
 
+// Releases all unused device memory currently held by the accelerator's
+// device-side caching allocator. The freed memory becomes available for reuse
+// by other applications or processes.
 TORCH_API inline void emptyCache() {
   const auto device_type = getAccelerator(true).value();
   at::getDeviceAllocator(device_type)->emptyCache();
 }
 
+// Releases all unused host (pinned) memory currently held by the accelerator's
+// host-side caching allocator. The freed memory becomes available for reuse by
+// other applications or processes.
+TORCH_API void emptyHostCache();
+
 TORCH_API inline at::CachingDeviceAllocator::DeviceStats getDeviceStats(
     c10::DeviceIndex device_index) {
   const auto device_type = getAccelerator(true).value();
diff --git a/aten/src/ATen/ExpandUtils.cpp b/aten/src/ATen/ExpandUtils.cpp
index 36c580695e34e..ff9b4573c9568 100644
--- a/aten/src/ATen/ExpandUtils.cpp
+++ b/aten/src/ATen/ExpandUtils.cpp
@@ -28,14 +28,17 @@ Container infer_size_impl(ArrayType a, ArrayType b) {
     auto sizeA = (dimA >= 0) ? a[dimA] : 1;
     auto sizeB = (dimB >= 0) ? b[dimB] : 1;
 
-    TORCH_CHECK(
-        sizeA == sizeB || sizeA == 1 || sizeB == 1,
-        "The size of tensor a (", sizeA,
-        ") must match the size of tensor b (", sizeB,
-        ") at non-singleton dimension ", i);
-
-      // 1s map to the other size (even 0).
-      expandedSizes[i] = sizeA == 1 ? sizeB : sizeA;
+    // If we know for sure one of them is 1, we are done.
+    // otherwise for unbacked we assume non-broadcasting path.
+    // Unbacked semantics should match python infer_size fake_impl.
+    if (!TORCH_GUARD_OR_FALSE(sym_eq(sizeA, 1)) && !TORCH_GUARD_OR_FALSE(sym_eq(sizeB, 1))) {
+      TORCH_MAYBE_SYM_CHECK(sym_eq(sizeA, sizeB),
+          "The size of tensor a (", sizeA,
+          ") must match the size of tensor b (", sizeB,
+          ") at non-singleton dimension ", i);
+    }
+    // 1s map to the other size (even 0).
+    expandedSizes[i] = TORCH_GUARD_OR_FALSE(sym_eq(sizeA, 1)) ? sizeB : sizeA;
   }
 
   return expandedSizes;
diff --git a/aten/src/ATen/FunctionalInverses.cpp b/aten/src/ATen/FunctionalInverses.cpp
index 123d87b304148..c372ae0ad339f 100644
--- a/aten/src/ATen/FunctionalInverses.cpp
+++ b/aten/src/ATen/FunctionalInverses.cpp
@@ -105,12 +105,10 @@ static Tensor unsqueeze_copy_to(const Tensor & self, IntArrayRef dim, c10::SymIn
 
 Tensor FunctionalInverses::_fw_primal_inverse(const at::Tensor& base, const at::Tensor& mutated_view, InverseReturnMode inverse_return_mode, int64_t level) {
     TORCH_INTERNAL_ASSERT(false, "Attempted to call _fw_primal() during the functionalization pass. For now, this is not supported.");
-    return Tensor();
 }
 
 Tensor FunctionalInverses::_make_dual_inverse(const at::Tensor& base, const at::Tensor& mutated_view, InverseReturnMode inverse_return_mode, const at::Tensor& tangent, int64_t level) {
     TORCH_INTERNAL_ASSERT(false, "Attempted to call _make_dual() during the functionalization pass. For now, this is not supported.");
-    return Tensor();
 }
 
 Tensor FunctionalInverses::view_as_real_inverse(const Tensor& base, const Tensor& mutated_view, InverseReturnMode inverse_return_mode) {
@@ -301,7 +299,6 @@ Tensor FunctionalInverses::transpose_int_inverse(const Tensor& base, const Tenso
 
 Tensor FunctionalInverses::_nested_view_from_buffer_inverse(const Tensor& base, const Tensor& mutated_view, InverseReturnMode inverse_return_mode, const Tensor& nested_sizes, const Tensor& nested_strides, const Tensor& storage_offsets) {
     TORCH_INTERNAL_ASSERT(false, "Attempted to call _nested_view_from_buffer() during the functionalization pass. For now, nested tensors aren't supported during functionalization");
-    return Tensor();
 }
 
 Tensor FunctionalInverses::_nested_view_from_jagged_inverse(const Tensor& base, const Tensor& mutated_view, InverseReturnMode inverse_return_mode, const Tensor& offsets, const Tensor& dummy, const std::optional<Tensor>& lengths, int64_t ragged_idx, const std::optional<Tensor>& min_seqlen, const std::optional<Tensor>& max_seqlen) {
@@ -342,47 +339,38 @@ Tensor FunctionalInverses::unsqueeze_inverse(const Tensor& base, const Tensor& m
 
 Tensor FunctionalInverses::_indices_inverse(const Tensor& base, const Tensor& mutated_view, InverseReturnMode inverse_return_mode) {
     TORCH_INTERNAL_ASSERT(false, "Attempted to call _indices() during the functionalization pass. For now, sparse tensors aren't supported during functionalization");
-    return Tensor();
 }
 
 Tensor FunctionalInverses::_values_inverse(const Tensor& base, const Tensor& mutated_view, InverseReturnMode inverse_return_mode) {
     TORCH_INTERNAL_ASSERT(false, "Attempted to call _values() during the functionalization pass. For now, sparse tensors aren't supported during functionalization");
-    return Tensor();
 }
 
 Tensor FunctionalInverses::indices_inverse(const Tensor& base, const Tensor& mutated_view, InverseReturnMode inverse_return_mode) {
     TORCH_INTERNAL_ASSERT(false, "Attempted to call indices() during the functionalization pass. For now, sparse tensors aren't supported during functionalization");
-    return Tensor();
 }
 
 Tensor FunctionalInverses::values_inverse(const Tensor& base, const Tensor& mutated_view, InverseReturnMode inverse_return_mode) {
     TORCH_INTERNAL_ASSERT(false, "Attempted to call values() during the functionalization pass. For now, sparse tensors aren't supported during functionalization");
-    return Tensor();
 }
 
 Tensor FunctionalInverses::_sparse_broadcast_to_inverse(const Tensor& base, const Tensor& mutated_view, InverseReturnMode inverse_return_mode, at::IntArrayRef size) {
     TORCH_INTERNAL_ASSERT(false, "Attempted to call _sparse_broadcast_to() during the functionalization pass. For now, sparse tensors aren't supported during functionalization");
-    return Tensor();
 }
 
 Tensor FunctionalInverses::crow_indices_inverse(const at::Tensor& base, const at::Tensor& mutated_view, InverseReturnMode inverse_return_mode) {
     TORCH_INTERNAL_ASSERT(false, "Attempted to call crow_indices() during the functionalization pass. For now, sparse tensors aren't supported during functionalization");
-    return Tensor();
 }
 
 Tensor FunctionalInverses::col_indices_inverse(const at::Tensor& base, const at::Tensor& mutated_view, InverseReturnMode inverse_return_mode) {
     TORCH_INTERNAL_ASSERT(false, "Attempted to call col_indices() during the functionalization pass. For now, sparse tensors aren't supported during functionalization");
-    return Tensor();
 }
 
 Tensor FunctionalInverses::ccol_indices_inverse(const at::Tensor& base, const at::Tensor& mutated_view, InverseReturnMode inverse_return_mode) {
     TORCH_INTERNAL_ASSERT(false, "Attempted to call ccol_indices() during the functionalization pass. For now, sparse tensors aren't supported during functionalization");
-    return Tensor();
 }
 
 Tensor FunctionalInverses::row_indices_inverse(const at::Tensor& base, const at::Tensor& mutated_view, InverseReturnMode inverse_return_mode) {
     TORCH_INTERNAL_ASSERT(false, "Attempted to call row_indices() during the functionalization pass. For now, sparse tensors aren't supported during functionalization");
-    return Tensor();
 }
 
 Tensor FunctionalInverses::unbind_int_inverse(const Tensor& base, const Tensor& mutated_view, InverseReturnMode inverse_return_mode, int64_t mutated_view_idx, int64_t dim) {
diff --git a/aten/src/ATen/FunctionalTensorWrapper.cpp b/aten/src/ATen/FunctionalTensorWrapper.cpp
index 9610360be4dd9..760a2a38aba39 100644
--- a/aten/src/ATen/FunctionalTensorWrapper.cpp
+++ b/aten/src/ATen/FunctionalTensorWrapper.cpp
@@ -864,6 +864,18 @@ void functionalize_op_helper(const c10::OperatorHandle& op, torch::jit::Stack* s
         "The composite op functionalization fallback expects its inputs all not to be functional tensors");
       auto t_new = c10::IValue(at::functionalization::impl::to_functional_tensor(opt_tensors));
       (*stack)[arguments_begin + idx] = t_new;
+    } else if (ivalue.isList()) {
+      // Handle nested lists containing tensor lists (e.g., Tensor[][]).
+      auto list = ivalue.toList();
+      for (const auto i : c10::irange(list.size())) {
+        const auto& elem = list.get(i);
+        if (elem.isTensorList()) {
+          auto tensors = elem.toTensorList();
+          TORCH_INTERNAL_ASSERT(!at::functionalization::impl::isFunctionalTensor(tensors),
+            "The composite op functionalization fallback expects its inputs all not to be functional tensors");
+          list.set(i, c10::IValue(at::functionalization::impl::to_functional_tensor(tensors)));
+        }
+      }
     }
   }
 
@@ -907,6 +919,17 @@ void functionalize_op_helper(const c10::OperatorHandle& op, torch::jit::Stack* s
       at::functionalization::impl::sync(opt_tensors);
       auto t_new = c10::IValue(at::functionalization::impl::from_functional_tensor(opt_tensors));
       (*stack)[returns_begin + idx] = t_new;
+    } else if (ivalue.isList()) {
+      // Handle nested lists containing tensor lists (e.g., Tensor[][]).
+      auto list = ivalue.toList();
+      for (const auto i : c10::irange(list.size())) {
+        const auto& elem = list.get(i);
+        if (elem.isTensorList()) {
+          auto tensors = elem.toTensorList();
+          at::functionalization::impl::sync(tensors);
+          list.set(i, c10::IValue(at::functionalization::impl::from_functional_tensor(tensors)));
+        }
+      }
     }
   }
 }
diff --git a/aten/src/ATen/FunctionalizeFallbackKernel.cpp b/aten/src/ATen/FunctionalizeFallbackKernel.cpp
index 100b4efe90b67..6d8134b01efa1 100644
--- a/aten/src/ATen/FunctionalizeFallbackKernel.cpp
+++ b/aten/src/ATen/FunctionalizeFallbackKernel.cpp
@@ -112,6 +112,21 @@ namespace {
           auto t_new = c10::IValue(at::functionalization::impl::from_functional_tensor(opt_tensors));
           (*stack)[arguments_begin + idx] = t_new;
         }
+      } else if (ivalue.isList()) {
+        // Handle nested lists containing tensor lists (e.g., Tensor[][]).
+        auto list = ivalue.toList();
+        for (const auto i : c10::irange(list.size())) {
+          const auto& elem = list.get(i);
+          if (elem.isTensorList()) {
+            any_tensor_inputs = true;
+            auto tensors = elem.toTensorList();
+            if (at::functionalization::impl::isFunctionalTensor(tensors)) {
+              any_functional_inputs = true;
+              at::functionalization::impl::sync(tensors);
+              list.set(i, c10::IValue(at::functionalization::impl::from_functional_tensor(tensors)));
+            }
+          }
+        }
       }
     }
     // we should wrap the output if any inputs were wrapped,
@@ -142,6 +157,16 @@ namespace {
         auto opt_tensors = ivalue.toOptionalTensorList();
         auto t_new = c10::IValue(at::functionalization::impl::to_functional_tensor(opt_tensors));
         (*stack)[returns_begin + idx] = t_new;
+      } else if (ivalue.isList() && should_wrap_outputs) {
+        // Handle nested lists containing tensor lists (e.g., Tensor[][]).
+        auto list = ivalue.toList();
+        for (const auto i : c10::irange(list.size())) {
+          const auto& elem = list.get(i);
+          if (elem.isTensorList()) {
+            auto tensors = elem.toTensorList();
+            list.set(i, c10::IValue(at::functionalization::impl::to_functional_tensor(tensors)));
+          }
+        }
       }
     }
   }
diff --git a/aten/src/ATen/MapAllocator.cpp b/aten/src/ATen/MapAllocator.cpp
index f2f0545410794..448f7738de279 100644
--- a/aten/src/ATen/MapAllocator.cpp
+++ b/aten/src/ATen/MapAllocator.cpp
@@ -295,20 +295,21 @@ MapAllocator::MapAllocator(WithFd /*unused*/, std::string_view filename, int fd,
 #ifdef HAVE_POSIX_FALLOCATE
           if (flags_ & ALLOCATOR_MAPPED_SHAREDMEM) {
             for (;;) {
-              if (posix_fallocate(fd, 0, static_cast<off_t>(size)) == 0) {
+              int err = posix_fallocate(fd, 0, static_cast<off_t>(size));
+              if (err == 0) {
                 break;
               }
 
-              if (errno == EINTR) {
+              if (err == EINTR) {
                 continue;
               }
 
-              if (errno == EINVAL || errno == EOPNOTSUPP) {
+              if (err == EINVAL || err == EOPNOTSUPP) {
                 // the underlying filesystem does not support the operation
                 break;
               }
 
-              TORCH_CHECK(false, "unable to allocate shared memory(shm) for file <", filename_, ">: ", c10::utils::str_error(errno), " (", errno, ")");
+              TORCH_CHECK(false, "unable to allocate shared memory(shm) for file <", filename_, ">: ", c10::utils::str_error(err), " (", err, ")");
             }
           }
 #endif
diff --git a/aten/src/ATen/Version.cpp b/aten/src/ATen/Version.cpp
index e59aab44ffd23..4a1c1525157f3 100644
--- a/aten/src/ATen/Version.cpp
+++ b/aten/src/ATen/Version.cpp
@@ -102,7 +102,7 @@ std::string get_cpu_capability() {
 #elif defined(HAVE_ZVECTOR_CPU_DEFINITION)
     case native::CPUCapability::ZVECTOR:
       return "Z VECTOR";
-#elif defined(HAVE_SVE256_CPU_DEFINITION) && defined(HAVE_ARM_BF16_CPU_DEFINITION)
+#elif defined(HAVE_SVE_CPU_DEFINITION)
     case native::CPUCapability::SVE256:
       return "SVE256";
 #else
@@ -218,8 +218,9 @@ std::string get_cxx_flags() {
     "Buck does not populate the `CXX_FLAGS` field of Caffe2 build options. "
     "As a result, `get_cxx_flags` is OSS only."
   );
-  #endif
+  #else
   return caffe2::GetBuildOptions().at("CXX_FLAGS");
+  #endif
 }
 
 }
diff --git a/aten/src/ATen/accelerator/Graph.cpp b/aten/src/ATen/accelerator/Graph.cpp
new file mode 100644
index 0000000000000..674a6e5a662cb
--- /dev/null
+++ b/aten/src/ATen/accelerator/Graph.cpp
@@ -0,0 +1,16 @@
+#include <ATen/DeviceAccelerator.h>
+#include <ATen/accelerator/Graph.h>
+
+namespace at::accelerator {
+
+Graph::Graph(bool keep_graph) {
+  c10::DeviceType device_type = at::accelerator::getAccelerator(true).value();
+  TORCH_CHECK(
+      has_graph_impl(device_type),
+      "Graph is not supported on device type: ",
+      device_type);
+  GraphImplArgs args{keep_graph};
+  impl_ = create_graph_impl(device_type, args);
+}
+
+} // namespace at::accelerator
diff --git a/aten/src/ATen/accelerator/Graph.h b/aten/src/ATen/accelerator/Graph.h
new file mode 100644
index 0000000000000..a3981be3a98c2
--- /dev/null
+++ b/aten/src/ATen/accelerator/Graph.h
@@ -0,0 +1,64 @@
+#pragma once
+#include <ATen/core/GraphImplInterface.h>
+
+namespace at::accelerator {
+
+struct TORCH_API Graph {
+  Graph(bool keep_graph = false);
+  ~Graph() = default;
+
+  C10_DISABLE_COPY_AND_ASSIGN(Graph);
+
+  // Begin graph capture on the current device and stream. Different accelerator
+  // backends may support different capture modes. `GraphCaptureMode::Default`
+  // lets the backend choose an appropriate capture strategy. If the requested
+  // capture mode is not supported, behavior is backend-specific (e.g., warn,
+  // raise an error, or fall back to `GraphCaptureMode::Default` or a
+  // backend-specific mode).
+  void capture_begin(
+      MempoolId_t pool = {0, 0},
+      GraphCaptureMode capture_mode = GraphCaptureMode::Default) {
+    impl_->capture_begin(pool, capture_mode);
+  }
+
+  // End graph capture and finalize the captured graph if `keep_graph` is false.
+  void capture_end() {
+    impl_->capture_end();
+  }
+
+  // Instantiate the captured graph for execution.
+  void instantiate() {
+    impl_->instantiate();
+  }
+
+  // Replay the previously captured graph.
+  void replay() {
+    impl_->replay();
+  }
+
+  // After reset(), the instance may be reused for a new capture.
+  void reset() {
+    impl_->reset();
+  }
+
+  // Return the memory pool associated with the captured graph.
+  MempoolId_t pool() const {
+    return impl_->pool();
+  }
+
+  // Enable backend-specific debug behavior for graph capture/replay.
+  void enable_debug_mode() {
+    impl_->enable_debug_mode();
+  }
+
+  // Dump the captured graph to a file for debugging purposes. The file format
+  // and content are backend-specific.
+  void debug_dump(const std::string& path) {
+    impl_->debug_dump(path);
+  }
+
+ private:
+  std::unique_ptr<at::GraphImplInterface> impl_;
+};
+
+} // namespace at::accelerator
diff --git a/aten/src/ATen/autocast_mode.h b/aten/src/ATen/autocast_mode.h
index 655b2343d5d5c..c8cc4f1a49e98 100644
--- a/aten/src/ATen/autocast_mode.h
+++ b/aten/src/ATen/autocast_mode.h
@@ -872,6 +872,7 @@ copy pasted in from VariableTypeEverything.cpp with appropriate substitutions.
   _(softplus)                         \
   _(layer_norm)                       \
   _(native_layer_norm)                \
+  _(rms_norm)                         \
   _(group_norm)                       \
   _(frobenius_norm, dim)              \
   _(nuclear_norm)                     \
diff --git a/aten/src/ATen/core/ATen_pch.h b/aten/src/ATen/core/ATen_pch.h
index 9a097845d003b..dfc643aa953da 100644
--- a/aten/src/ATen/core/ATen_pch.h
+++ b/aten/src/ATen/core/ATen_pch.h
@@ -97,7 +97,6 @@
 #include <c10/util/AlignOf.h>
 #include <c10/util/ArrayRef.h>
 #include <c10/util/BFloat16.h>
-#include <c10/util/C++17.h>
 #include <c10/util/ConstexprCrc.h>
 #include <c10/util/Deprecated.h>
 #include <c10/util/DimVector.h>
diff --git a/aten/src/ATen/core/Dimname.cpp b/aten/src/ATen/core/Dimname.cpp
index 66aa8cb69e1ed..1a2fcdfe2674a 100644
--- a/aten/src/ATen/core/Dimname.cpp
+++ b/aten/src/ATen/core/Dimname.cpp
@@ -4,7 +4,10 @@
 
 namespace at {
 
-static Symbol kWildcard = Symbol::dimname("*");
+static Symbol kWildcard() {
+  static Symbol singleton = Symbol::dimname("*");
+  return singleton;
+}
 
 std::ostream& operator<<(std::ostream& out, const Dimname& dimname) {
   if (dimname.type() == NameType::WILDCARD) {
@@ -46,7 +49,7 @@ static void check_valid_identifier(const std::string& name) {
 
 Dimname Dimname::fromSymbol(Symbol name) {
   TORCH_INTERNAL_ASSERT(name.is_dimname());
-  if (name == kWildcard) {
+  if (name == kWildcard()) {
     return Dimname::wildcard();
   }
   check_valid_identifier(name.toUnqualString());
@@ -54,7 +57,7 @@ Dimname Dimname::fromSymbol(Symbol name) {
 }
 
 Dimname Dimname::wildcard() {
-  static Dimname result(kWildcard, NameType::WILDCARD);
+  static Dimname result(kWildcard(), NameType::WILDCARD);
   return result;
 }
 
diff --git a/aten/src/ATen/core/Formatting.cpp b/aten/src/ATen/core/Formatting.cpp
index 62b16a83e523b..a1a1ee0acc6ec 100644
--- a/aten/src/ATen/core/Formatting.cpp
+++ b/aten/src/ATen/core/Formatting.cpp
@@ -303,6 +303,9 @@ std::ostream& print(
   } else if (tensor_.is_mps()) {
     // MPS does not support double tensors, so first copy then convert
     tensor = tensor_.to(kCPU).to(kDouble).contiguous();
+  } else if (tensor_.is_privateuseone()) {
+    // PrivateUseOne backends may not support double tensors
+    tensor = tensor_.to(kCPU).to(kDouble).contiguous();
   } else {
     tensor = tensor_.to(kCPU, kDouble).contiguous();
   }
diff --git a/aten/src/ATen/core/GraphImplInterface.cpp b/aten/src/ATen/core/GraphImplInterface.cpp
new file mode 100644
index 0000000000000..d1c8ea30bc14e
--- /dev/null
+++ b/aten/src/ATen/core/GraphImplInterface.cpp
@@ -0,0 +1,7 @@
+#include <ATen/core/GraphImplInterface.h>
+
+namespace at {
+
+C10_DEFINE_REGISTRY(GraphImplRegistry, GraphImplInterface, GraphImplArgs)
+
+} // namespace at
diff --git a/aten/src/ATen/core/GraphImplInterface.h b/aten/src/ATen/core/GraphImplInterface.h
new file mode 100644
index 0000000000000..c72482345c4d9
--- /dev/null
+++ b/aten/src/ATen/core/GraphImplInterface.h
@@ -0,0 +1,100 @@
+#pragma once
+#include <c10/core/Allocator.h>
+#include <c10/util/Registry.h>
+
+namespace at {
+
+enum class GraphCaptureMode : int8_t {
+  // Backend-defined default capture behavior.
+  Default = 0,
+  // Potentially unsafe API calls are prohibited. Errors may occur if capture in
+  // the current thread affects other threads.
+  Global,
+  // Potentially unsafe API calls are prohibited. Errors occur only if capture
+  // in the current thread affects itself.
+  ThreadLocal,
+  // The current thread is allowed to make potentially unsafe API calls, except
+  // for calls that inherently conflict with stream capture.
+  Relaxed,
+};
+
+// Arguments used to construct a GraphImplInterface instance.
+//
+// This struct is intentionally lightweight and extensible so that new options
+// can be added in a backward-compatible way without breaking existing or
+// out-of-tree backends.
+struct TORCH_API GraphImplArgs {
+  // Whether to keep the underlying raw graph after capture is complete.
+  bool keep_graph = false;
+};
+
+// A lightweight, backend-agnostic interface that provides a unified API for
+// graph capture and replay.
+//
+// Each backend (e.g. CUDA, XPU, etc.) implements this interface and registers
+// its implementation via GraphImplRegistry. Implementations are required to
+// provide a constructor that accepts `GraphImplArgs`.
+// While the concrete semantics and detailed behavior of capture and replay may
+// vary across backends, the API contract exposed here is consistent.
+struct TORCH_API GraphImplInterface {
+  virtual ~GraphImplInterface() = default;
+  // Begin graph capture on the current device and stream.
+  // `pool` specifies the memory pool to be used during capture.
+  // `capture_mode` controls how capture interacts with other concurrent work.
+  // Its exact semantics are backend-specific. If a backend does not support the
+  // requested capture mode, it may choose to emit a warning, raise an error, or
+  // fall back to `GraphCaptureMode::Default` or a backend-specific mode.
+  virtual void capture_begin(
+      MempoolId_t pool = {0, 0},
+      GraphCaptureMode capture_mode = GraphCaptureMode::Default) = 0;
+
+  // End graph capture and maybe finalize the captured graph.
+  virtual void capture_end() = 0;
+
+  // Instantiate the captured graph for execution.
+  virtual void instantiate() = 0;
+
+  // Replay the previously captured graph.
+  virtual void replay() = 0;
+
+  // Reset internal state and release any backend-specific resources.
+  // After reset(), the instance may be reused for a new capture.
+  virtual void reset() = 0;
+
+  // Return the memory pool associated with the captured graph.
+  virtual MempoolId_t pool() const = 0;
+
+  // Enable backend-specific debug behavior for graph capture/replay.
+  // Implementations may enable extra validation and/or logging to help diagnose
+  // issues. Backends that do not support debug mode could implement this as a
+  // no-op.
+  virtual void enable_debug_mode() = 0;
+
+  // Dump the captured graph to a file for debugging purposes. The file format
+  // and content are backend-specific.
+  virtual void debug_dump(const std::string& path) = 0;
+};
+
+TORCH_DECLARE_REGISTRY(GraphImplRegistry, GraphImplInterface, GraphImplArgs);
+
+// Registry mapping DeviceType -> GraphImplInterface implementation.
+// The key is the string returned by c10::DeviceTypeName(device_type, false).
+#define REGISTER_GRAPH_IMPL(key, impl) \
+C10_REGISTER_CLASS(GraphImplRegistry, key, impl)
+
+// Check whether a graph implementation is registered for the given device type.
+inline bool has_graph_impl(const c10::DeviceType device_type) {
+  auto key = c10::DeviceTypeName(device_type, /*lowercase=*/false);
+  return GraphImplRegistry()->Has(key);
+}
+
+// Factory function to create a graph implementation for the given device.
+// Returns nullptr if no implementation is registered for the device.
+inline std::unique_ptr<GraphImplInterface> create_graph_impl(
+    const c10::DeviceType device_type,
+    const GraphImplArgs& args = {}) {
+  auto key = c10::DeviceTypeName(device_type, /*lowercase=*/false);
+  return GraphImplRegistry()->Create(key, args);
+}
+
+} // namespace at
diff --git a/aten/src/ATen/core/List_test.cpp b/aten/src/ATen/core/List_test.cpp
index 8dd677790cd0b..48118463bd2c6 100644
--- a/aten/src/ATen/core/List_test.cpp
+++ b/aten/src/ATen/core/List_test.cpp
@@ -4,6 +4,13 @@
 using namespace c10;
 using std::string;
 
+// TODO(NS): Remove me
+// For some reason emplace tests in this file fails to compile with C++20 with the following warning
+// aten/src/ATen/core/ivalue.h:240:3: error: array subscript 0 is outside array bounds of 'c10::IValue [0]'
+#if defined(__GNUC__) && __GNUC__ == 13
+#pragma GCC diagnostic ignored "-Warray-bounds"
+#endif
+
 // NOLINTBEGIN(performance-move-const-arg, bugprone-use-after-move, *analyzer*Move)
 TEST(ListTestIValueBasedList, givenEmptyList_whenCallingEmpty_thenReturnsTrue) {
     List<string> list;
diff --git a/aten/src/ATen/core/TensorBase.h b/aten/src/ATen/core/TensorBase.h
index 1f4c23e7c87cf..7715ece7c6187 100644
--- a/aten/src/ATen/core/TensorBase.h
+++ b/aten/src/ATen/core/TensorBase.h
@@ -24,7 +24,6 @@
 #include <c10/core/TensorOptions.h>
 #include <c10/core/UndefinedTensorImpl.h>
 #include <c10/core/WrapDimMinimal.h>
-#include <c10/util/C++17.h>
 #include <c10/util/Exception.h>
 #include <c10/util/ExclusivelyOwned.h>
 #include <c10/util/ExclusivelyOwnedTensorTraits.h>
@@ -628,12 +627,10 @@ class TORCH_API TensorBase {
     return mutable_data_ptr();
   }
 
-  template <typename T, std::enable_if_t<!std::is_const_v<T>, int> = 0>
+  // Implemented in aten/src/ATen/templates/TensorMethods.cpp
+  template <typename T>
   const T* const_data_ptr() const;
 
-  template <typename T, std::enable_if_t<std::is_const_v<T>, int> = 0>
-  const std::remove_const_t<T>* const_data_ptr() const;
-
   template <typename T>
   T* mutable_data_ptr() const;
 
diff --git a/aten/src/ATen/core/Vitals.cpp b/aten/src/ATen/core/Vitals.cpp
deleted file mode 100644
index db58c03830539..0000000000000
--- a/aten/src/ATen/core/Vitals.cpp
+++ /dev/null
@@ -1,96 +0,0 @@
-#include <ATen/core/Vitals.h>
-#include <c10/util/env.h>
-#include <iostream>
-
-namespace at::vitals {
-
-APIVitals VitalsAPI;
-
-std::ostream& operator<<(std::ostream& os, TorchVital const& tv) {
-  for (const auto& m : tv.attrs) {
-    os << "[TORCH_VITAL] " << tv.name << '.' << m.first << "\t\t "
-       << m.second.value << '\n';
-  }
-  return os;
-}
-
-TorchVital::~TorchVital() {
-  if (torchVitalEnabled()) {
-    std::cout << *this;
-  }
-}
-
-TorchVitalAttr& TorchVital::create(const std::string& attr) {
-  return create(attr, /* force = */ false);
-}
-
-TorchVitalAttr& TorchVital::create(const std::string& attr, bool force) {
-  if (!(torchVitalEnabled() || force)) {
-    static TorchVitalAttr disabled;
-    return disabled;
-  }
-  auto iter = attrs.find(attr);
-  if (iter == attrs.end()) {
-    auto r = attrs.emplace(attr, TorchVitalAttr());
-    return r.first->second;
-  }
-  return iter->second;
-}
-
-bool torchVitalEnabled() {
-  // If this is a performance hit, make `enabled` variable static
-  // and return `const bool&` instead
-  bool enabled = []() {
-    auto const e = c10::utils::get_env("TORCH_VITAL");
-    if (e.has_value()) {
-      return !e.value().empty();
-    }
-    return false;
-  }();
-  if (enabled) {
-    VitalsAPI.vitals_enabled = true;
-  }
-  return VitalsAPI.vitals_enabled;
-}
-
-std::string APIVitals::readVitals() {
-  if (!torchVitalEnabled()) {
-    return "";
-  }
-
-  std::stringstream buf;
-  for (const auto& x : name_map_) {
-    buf << x.second;
-  }
-  return buf.str();
-}
-
-bool APIVitals::setVital(
-    const std::string& vital_name,
-    const std::string& attr_name,
-    const std::string& value,
-    bool force) {
-  if (!(torchVitalEnabled() || force)) {
-    return false;
-  }
-
-  auto iter = name_map_.find(vital_name);
-  TorchVital* vital = nullptr;
-  if (iter == name_map_.end()) {
-    auto r = name_map_.emplace(vital_name, TorchVital(vital_name));
-    vital = &r.first->second;
-  } else {
-    vital = &iter->second;
-  }
-
-  vital->create(attr_name, force).write(value, force);
-  return true;
-}
-
-APIVitals::APIVitals() : vitals_enabled(false) {
-  // Set default values, force is necessary because in unit tests the env
-  // variable may not be set when global APIVitals are constructed.
-  setVital("CUDA", "used", "False", /* force = */ true);
-}
-
-} // namespace at::vitals
diff --git a/aten/src/ATen/core/Vitals.h b/aten/src/ATen/core/Vitals.h
deleted file mode 100644
index 2fd7729744a10..0000000000000
--- a/aten/src/ATen/core/Vitals.h
+++ /dev/null
@@ -1,94 +0,0 @@
-#pragma once
-#include <ostream>
-#include <sstream>
-#include <unordered_map>
-
-#include <c10/core/impl/LocalDispatchKeySet.h>
-
-namespace at::vitals {
-
-TORCH_API bool torchVitalEnabled();
-
-struct TORCH_API TorchVitalAttr {
-  // always initialized to empty
-  std::string value;
-  template <typename T>
-  TorchVitalAttr& operator<<(const T& t) {
-    if (torchVitalEnabled()) {
-      std::stringstream ss;
-      ss << t;
-      value += ss.str();
-    }
-    return *this;
-  }
-
-  template <typename T>
-  void write(const T& t, bool force) {
-    if (force || torchVitalEnabled()) {
-      std::stringstream ss;
-      ss << t;
-      value = ss.str();
-    }
-  }
-};
-
-struct TORCH_API TorchVital {
-  std::string name;
-  std::unordered_map<std::string, TorchVitalAttr> attrs;
-
-  explicit TorchVital(std::string n) : name(std::move(n)) {}
-  TorchVital(const TorchVital&) = default;
-  TorchVital(TorchVital&&) = default;
-  TorchVital& operator=(const TorchVital&) = default;
-  TorchVital& operator=(TorchVital&&) = default;
-  TorchVital() = delete;
-
-  TorchVitalAttr& create(const std::string& attr);
-  TorchVitalAttr& create(const std::string& attr, bool force);
-  friend std::ostream& operator<<(std::ostream& os, const TorchVital& dt);
-
-  ~TorchVital();
-};
-
-std::ostream& operator<<(std::ostream& os, TorchVital const& tv);
-
-// A way to access vitals by string names instead of by global reference.
-// This enables access to vitals from the PythonAPI.
-class TORCH_API APIVitals {
- public:
-  bool vitals_enabled;
-
-  // Set any vital sign that was added to the map.
-  bool setVital(
-      const std::string& vital_name,
-      const std::string& attr_name,
-      const std::string& value,
-      bool force = false);
-  std::string readVitals();
-
-  APIVitals();
-
-  // Ensure this stays a singleton
-  APIVitals(APIVitals const& other) = delete;
-  APIVitals(APIVitals&& other) = delete;
-  APIVitals& operator=(const APIVitals&) = delete;
-  APIVitals& operator=(APIVitals&&) = delete;
-  ~APIVitals() = default;
-
- private:
-  std::unordered_map<std::string, TorchVital> name_map_;
-};
-
-extern TORCH_API APIVitals VitalsAPI;
-
-} // namespace at::vitals
-
-#define TORCH_VITAL_DECLARE(name) \
-  TORCH_API at::vitals::TorchVital TorchVital_##name;
-
-#define TORCH_VITAL_DEFINE(name) \
-  TORCH_API at::vitals::TorchVital TorchVital_##name(#name);
-
-#define TORCH_VITAL_BASE(name) TorchVital_##name
-
-#define TORCH_VITAL(name, attr) TORCH_VITAL_BASE(name).create(#attr)
diff --git a/aten/src/ATen/core/boxing/KernelFunction_impl.h b/aten/src/ATen/core/boxing/KernelFunction_impl.h
index 0ce1b77dece0e..fe9e1d2ea961c 100644
--- a/aten/src/ATen/core/boxing/KernelFunction_impl.h
+++ b/aten/src/ATen/core/boxing/KernelFunction_impl.h
@@ -3,7 +3,6 @@
 #include <ATen/core/boxing/impl/boxing.h>
 #include <ATen/core/boxing/impl/make_boxed_from_unboxed_functor.h>
 
-#include <c10/util/C++17.h>
 #include <type_traits>
 
 namespace c10 {
diff --git a/aten/src/ATen/core/boxing/impl/make_boxed_from_unboxed_functor.h b/aten/src/ATen/core/boxing/impl/make_boxed_from_unboxed_functor.h
index 5f11d1715f34a..ac6fd7b09931f 100644
--- a/aten/src/ATen/core/boxing/impl/make_boxed_from_unboxed_functor.h
+++ b/aten/src/ATen/core/boxing/impl/make_boxed_from_unboxed_functor.h
@@ -731,6 +731,22 @@ struct push_outputs<void, AllowDeprecatedTypes> final {
   static void copy(int /*dummy*/, Stack* /*stack*/) {}
 };
 
+// decay_if_tuple ensures that if T is a tuple, all of its elements are decayed.
+// This is useful for kernels that return a tuple of references (e.g., Tensor&),
+// to avoid dangling references after the stack is dropped.
+template <typename T>
+struct decay_if_tuple {
+  using type = std::decay_t<T>;
+};
+
+template <typename... Args>
+struct decay_if_tuple<std::tuple<Args...>> {
+  using type = std::tuple<std::decay_t<Args>...>;
+};
+
+template <typename T>
+using decay_if_tuple_t = typename decay_if_tuple<T>::type;
+
 // make_boxed_from_unboxed_functor
 
 template <class KernelFunctor, bool AllowDeprecatedTypes>
@@ -761,7 +777,7 @@ struct make_boxed_from_unboxed_functor final {
       // we actually store it by value and don't get a dangling reference. This
       // is only required because some kernels still return `Tensor&`. [Note:
       // VC++ and 'std': ambiguous symbol]
-      using ReturnType_ = ::std::decay_t<ReturnType>;
+      using ReturnType_ = decay_if_tuple_t<ReturnType>;
       ReturnType_ output = call_functor_with_args_from_stack<
           KernelFunctor,
           AllowDeprecatedTypes>(functor, dispatchKeySet, stack);
diff --git a/aten/src/ATen/core/enum_tag.h b/aten/src/ATen/core/enum_tag.h
new file mode 100644
index 0000000000000..13f9c65a63488
--- /dev/null
+++ b/aten/src/ATen/core/enum_tag.h
@@ -0,0 +1,4 @@
+#pragma once
+
+// at::Tag is now generated into torch/headeronly
+#include <torch/headeronly/core/enum_tag.h>
diff --git a/aten/src/ATen/core/interned_strings.h b/aten/src/ATen/core/interned_strings.h
index 38942031befcd..532c921ca4258 100644
--- a/aten/src/ATen/core/interned_strings.h
+++ b/aten/src/ATen/core/interned_strings.h
@@ -187,7 +187,6 @@ namespace c10 {
   _(aten, append)                    \
   _(aten, as_tensor)                 \
   _(aten, adaptive_avg_pool2d_backward) \
-  _(aten, dim)                       \
   _(aten, format)                    \
   _(aten, percentFormat)             \
   _(aten, __not__)                   \
diff --git a/aten/src/ATen/cpu/vec/functional_base.h b/aten/src/ATen/cpu/vec/functional_base.h
index 112121b297055..bca8a719818db 100644
--- a/aten/src/ATen/cpu/vec/functional_base.h
+++ b/aten/src/ATen/cpu/vec/functional_base.h
@@ -103,7 +103,7 @@ struct VecReduceAllSIMD<float, Op> {
        // !defined(C10_MOBILE)
 
 #if defined(__aarch64__) && !defined(C10_MOBILE) && !defined(__CUDACC__) && \
-    !defined(CPU_CAPABILITY_SVE)
+    !defined(CPU_CAPABILITY_SVE256)
 template <typename Op>
 struct VecReduceAllSIMD<float, Op> {
   static inline float apply(
@@ -141,7 +141,7 @@ struct VecReduceAllSIMD<float, std::plus<Vectorized<float>>> {
   }
 };
 #endif // defined(__aarch64__) && !defined(C10_MOBILE) && !defined(__CUDACC__)
-       // && !defined(CPU_CAPABILITY_SVE)
+       // && !defined(CPU_CAPABILITY_SVE256)
 
 #if defined(__aarch64__) && !defined(C10_MOBILE) && !defined(__CUDACC__) && \
     defined(CPU_CAPABILITY_SVE256)
diff --git a/aten/src/ATen/cpu/vec/sve/sve_helper.h b/aten/src/ATen/cpu/vec/sve/sve_helper.h
index f3786019064c1..2f66395024f40 100644
--- a/aten/src/ATen/cpu/vec/sve/sve_helper.h
+++ b/aten/src/ATen/cpu/vec/sve/sve_helper.h
@@ -4,7 +4,7 @@
 
 #include <ATen/cpu/vec/vec_base.h>
 
-#if defined(CPU_CAPABILITY_SVE)
+#if defined(CPU_CAPABILITY_SVE256)
 
 // Define the data type of VLS(vector-length specific).
 typedef svbool_t vls_pred_t
@@ -77,4 +77,4 @@ typedef svfloat64_t vls_float64_t
 #define ALL_F64_TRUE_MASK svreinterpret_f64_s64(ALL_S64_TRUE_MASK)
 #define ALL_F64_FALSE_MASK svreinterpret_f64_s64(ALL_S64_FALSE_MASK)
 
-#endif // defined(CPU_CAPABILITY_SVE)
+#endif // defined(CPU_CAPABILITY_SVE256)
diff --git a/aten/src/ATen/cpu/vec/sve/vec_bfloat16.h b/aten/src/ATen/cpu/vec/sve/vec_bfloat16.h
index 757ef839f965a..07438e6b48140 100644
--- a/aten/src/ATen/cpu/vec/sve/vec_bfloat16.h
+++ b/aten/src/ATen/cpu/vec/sve/vec_bfloat16.h
@@ -586,7 +586,7 @@ Vectorized<BFloat16> inline fmadd(
   return a * b + c;
 }
 
-#endif // defined(CPU_CAPABILITY_SVE) && defined(__ARM_FEATURE_BF16)
+#endif // defined(CPU_CAPABILITY_SVE256) && defined(__ARM_FEATURE_BF16)
 
 } // namespace CPU_CAPABILITY
 } // namespace vec
diff --git a/aten/src/ATen/cpu/vec/sve/vec_common_sve.h b/aten/src/ATen/cpu/vec/sve/vec_common_sve.h
index 69ed5d061bd88..7c033c45945c3 100644
--- a/aten/src/ATen/cpu/vec/sve/vec_common_sve.h
+++ b/aten/src/ATen/cpu/vec/sve/vec_common_sve.h
@@ -8,7 +8,7 @@
 #include <ATen/cpu/vec/sve/sve_helper.h>
 #include <ATen/cpu/vec/vec_base.h>
 
-#if defined(CPU_CAPABILITY_SVE)
+#if defined(CPU_CAPABILITY_SVE256)
 #include <ATen/cpu/vec/sve/vec_bfloat16.h>
 #include <ATen/cpu/vec/sve/vec_double.h>
 #include <ATen/cpu/vec/sve/vec_float.h>
@@ -27,7 +27,7 @@ namespace at::vec {
 // accessed as `at::vec`.
 inline namespace CPU_CAPABILITY {
 
-#if defined(CPU_CAPABILITY_SVE)
+#if defined(CPU_CAPABILITY_SVE256)
 
 // ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ CAST ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 #define DEFINE_SVE_CAST(t1_t, t1_prefix, t2_t, t2_prefix)                 \
@@ -230,7 +230,7 @@ std::pair<
 }
 #endif // __ARM_FEATURE_BF16
 
-#endif // defined(CPU_CAPABILITY_SVE)
+#endif // defined(CPU_CAPABILITY_SVE256)
 
 } // namespace CPU_CAPABILITY
 } // namespace at::vec
diff --git a/aten/src/ATen/cpu/vec/sve/vec_double.h b/aten/src/ATen/cpu/vec/sve/vec_double.h
index c2aaeda9d9d00..8d93a3b730e73 100644
--- a/aten/src/ATen/cpu/vec/sve/vec_double.h
+++ b/aten/src/ATen/cpu/vec/sve/vec_double.h
@@ -22,7 +22,7 @@ namespace at::vec {
 // accessed as `at::vec`.
 inline namespace CPU_CAPABILITY {
 
-#if defined(CPU_CAPABILITY_SVE)
+#if defined(CPU_CAPABILITY_SVE256)
 
 template <>
 struct is_vec_specialized_for<double> : std::bool_constant<true> {};
@@ -627,7 +627,7 @@ Vectorized<double> inline fnmsub(
   return svnmad_f64_x(ptrue, a, b, c);
 }
 
-#endif // defined(CPU_CAPABILITY_SVE)
+#endif // defined(CPU_CAPABILITY_SVE256)
 
 } // namespace CPU_CAPABILITY
 } // namespace at::vec
diff --git a/aten/src/ATen/cpu/vec/sve/vec_float.h b/aten/src/ATen/cpu/vec/sve/vec_float.h
index 379aa95584c97..8136e6b8c6ee9 100644
--- a/aten/src/ATen/cpu/vec/sve/vec_float.h
+++ b/aten/src/ATen/cpu/vec/sve/vec_float.h
@@ -22,7 +22,7 @@ namespace at::vec {
 // accessed as `at::vec`.
 inline namespace CPU_CAPABILITY {
 
-#if defined(CPU_CAPABILITY_SVE)
+#if defined(CPU_CAPABILITY_SVE256)
 
 template <>
 struct is_vec_specialized_for<float> : std::bool_constant<true> {};
@@ -285,8 +285,56 @@ class Vectorized<float> {
     svfloat32_t poly = svmla_x(svptrue_b32(), r, r2, c1);
     return svmla_x(svptrue_b32(), scale, scale, poly);
   }
+  // Implementation from Arm Optimized Routines:
+  // https://github.com/ARM-software/optimized-routines/blob/v26.01/math/aarch64/experimental/sve/sv_expf_inline.h
   Vectorized<float> fexp_u20() const {
-    return exp_u20();
+    // fast exponential intended for cases where outputs will be downcasted to
+    // FP16 / BF16 (e.g. attention softmax).
+    // Accurate within 1 ULP for FP16
+    // Accurate within 1 ULP for BF16 for inputs in [-87.346, max_float] &
+    // clamps
+    //  inputs < -87.346 to zero.
+    // Implementation is similar to exp_u20, but:
+    // - approximates exp(r) - 1 as r instead of r + 0.5 r^2
+    // - does not split natural log (ln) into high / low parts
+    // - avoids special case code by clamping exp(x) to 0 for x < -87.346 and
+    // inf for x > 88.717
+
+    constexpr float upper_bound = 0x1.62dea4p+6f;
+    constexpr float lower_bound = -0x1.5d619ap+6f;
+
+    const svfloat32_t ln2 = svdup_n_f32(0x1.62e43p-1f);
+    const svfloat32_t inv_ln2 = svdup_n_f32(0x1.715476p+0f);
+
+    constexpr float shift = 0x1.803f8p17f;
+
+    // exp(x) = 2^n (1 + poly(r)), with 1 + poly(r) in [1/sqrt(2),sqrt(2)]
+    // x = ln2*n + r, with r in [-ln2/2, ln2/2] and poly(r) ~= r.
+
+    // n = round(x/(ln2/N))
+    svfloat32_t z = svmad_x(svptrue_b32(), inv_ln2, values, shift);
+    svfloat32_t n = svsub_x(svptrue_b32(), z, shift);
+
+    // n = round(x/(ln2/N))
+    svfloat32_t r = svmls_x(svptrue_b32(), values, n, ln2);
+
+    // scale = 2^(n/N)
+    svfloat32_t scale = svexpa(svreinterpret_u32(z));
+
+    // poly(r) = exp(r) - 1 ~= r
+    svfloat32_t y = svmla_x(svptrue_b32(), scale, scale, r);
+
+    // clamp to 0, inf
+    y = svsel_f32(
+        svcmplt_f32(svptrue_b32(), values, svdup_n_f32(lower_bound)),
+        svdup_n_f32(0.0f),
+        y);
+    y = svsel_f32(
+        svcmpgt_f32(svptrue_b32(), values, svdup_n_f32(upper_bound)),
+        svdup_n_f32(INFINITY),
+        y);
+
+    return y;
   }
   Vectorized<float> fmod(const Vectorized<float>& q) const {
     USE_SLEEF(
@@ -764,7 +812,7 @@ Vectorized<float> inline fnmsub(
   return svnmad_f32_x(ptrue, a, b, c);
 }
 
-#endif // defined(CPU_CAPABILITY_SVE)
+#endif // defined(CPU_CAPABILITY_SVE256)
 
 } // namespace CPU_CAPABILITY
 } // namespace at::vec
diff --git a/aten/src/ATen/cpu/vec/sve/vec_int.h b/aten/src/ATen/cpu/vec/sve/vec_int.h
index f0bc42caa9502..e9cad3251308d 100644
--- a/aten/src/ATen/cpu/vec/sve/vec_int.h
+++ b/aten/src/ATen/cpu/vec/sve/vec_int.h
@@ -15,7 +15,7 @@ namespace at::vec {
 // accessed as `at::vec`.
 inline namespace CPU_CAPABILITY {
 
-#if defined(CPU_CAPABILITY_SVE)
+#if defined(CPU_CAPABILITY_SVE256)
 
 #define VEC_INT_SVE_TEMPLATE(vl, bit)                                         \
   template <>                                                                 \
@@ -493,7 +493,7 @@ Vectorized<int8_t> inline operator>>(
   return svasr_s8_x(ptrue, a, svreinterpret_u8_s8(b));
 }
 
-#endif // defined(CPU_CAPABILITY_SVE)
+#endif // defined(CPU_CAPABILITY_SVE256)
 
 } // namespace CPU_CAPABILITY
 } // namespace at::vec
diff --git a/aten/src/ATen/cpu/vec/sve/vec_qint.h b/aten/src/ATen/cpu/vec/sve/vec_qint.h
index 61cb63cb1e12a..c3107720dc8d5 100644
--- a/aten/src/ATen/cpu/vec/sve/vec_qint.h
+++ b/aten/src/ATen/cpu/vec/sve/vec_qint.h
@@ -46,7 +46,7 @@ namespace at::vec {
 // accessed as `at::vec`.
 inline namespace CPU_CAPABILITY {
 
-#if defined(CPU_CAPABILITY_SVE)
+#if defined(CPU_CAPABILITY_SVE256)
 
 // NOTE: These are low-performance implementations that we fall back on
 // if we are not building with SVE. This may not be an issue, because
@@ -600,7 +600,7 @@ Vectorized<c10::quint8> inline maximum(
   return a.maximum(b);
 }
 
-#endif // defined(CPU_CAPABILITY_SVE)
+#endif // defined(CPU_CAPABILITY_SVE256)
 
 } // namespace CPU_CAPABILITY
 } // namespace at::vec
diff --git a/aten/src/ATen/cpu/vec/vec128/vec128.h b/aten/src/ATen/cpu/vec/vec128/vec128.h
index 9f9079d475a8f..ba740dbf43878 100644
--- a/aten/src/ATen/cpu/vec/vec128/vec128.h
+++ b/aten/src/ATen/cpu/vec/vec128/vec128.h
@@ -4,7 +4,7 @@
 #include <ATen/cpu/vec/intrinsics.h>
 
 #ifdef __aarch64__
-#if !defined(CPU_CAPABILITY_SVE)
+#if !defined(CPU_CAPABILITY_SVE256)
 #include <ATen/cpu/vec/vec128/vec128_bfloat16_neon.h>
 #include <ATen/cpu/vec/vec128/vec128_double_neon.h>
 #include <ATen/cpu/vec/vec128/vec128_float_neon.h>
diff --git a/aten/src/ATen/cpu/vec/vec128/vec128_float_neon.h b/aten/src/ATen/cpu/vec/vec128/vec128_float_neon.h
index 6a64226475cf3..3ee50822458b5 100644
--- a/aten/src/ATen/cpu/vec/vec128/vec128_float_neon.h
+++ b/aten/src/ATen/cpu/vec/vec128/vec128_float_neon.h
@@ -321,10 +321,10 @@ class Vectorized<float> {
     }
 
     const float32x4_t inv_ln2 = vdupq_n_f32(0x1.715476p+0f);
-    const float ln2_hi = 0x1.62e4p-1f;
-    const float ln2_lo = 0x1.7f7d1cp-20f;
-    const float c0 = 0x1.0e4020p-7f;
-    const float c2 = 0x1.555e66p-3f;
+    constexpr float ln2_hi = 0x1.62e4p-1f;
+    constexpr float ln2_lo = 0x1.7f7d1cp-20f;
+    constexpr float c0 = 0x1.0e4020p-7f;
+    constexpr float c2 = 0x1.555e66p-3f;
     const float32x4_t ln2_c02 = {ln2_hi, ln2_lo, c0, c2};
 
     const uint32x4_t exponent_bias = vdupq_n_u32(0x3f800000);
@@ -354,7 +354,47 @@ class Vectorized<float> {
     return vexpq_f32_u20();
   }
   Vectorized<float> fexp_u20() const {
-    return exp_u20();
+    // fast exponential intended for cases where outputs will be downcasted to
+    // FP16 / BF16 (e.g. attention softmax). Accurate within 1 ULP for FP16
+    // Accurate within 1 ULP for BF16 for inputs in [-87.683, 88.376] & clamps
+    // inputs outside this range to 0 / inf. Implementation is similar to
+    // exp_u20, but:
+    // - uses a third degree polynomial approximation for exp(r) instead of a
+    // fifth degree one, with coefficients re-tuned.
+    // - does not split natural log (ln) into high / low parts
+    // - clamps exp(x) to 0 for x < -87.346351f and inf for x > 88.3762589f
+
+    const float32x4_t lower_bound = vdupq_n_f32(-0x1.5ebb82p+6f);
+    const float32x4_t upper_bound = vdupq_n_f32(0x1.61814ap+6f);
+    const float32x4_t inv_ln2 = vdupq_n_f32(0x1.715476p+0f);
+    constexpr float ln2 = 0x1.62e43p-1f;
+    constexpr float c2 = 0x1.5592ecp-3f;
+    const float32x4_t c3 = vdupq_n_f32(0x1.017d34p-1f);
+    const uint32x4_t lt_lower = vcltq_f32(values, lower_bound);
+    const uint32x4_t gt_upper = vcgtq_f32(values, upper_bound);
+
+    // exp(x) = 2^n (1 + exp(r))
+    // r = x - n*ln2, with n = round(x/ln2)
+    // exp(r) ~ poly(r) = r + r^2 * (c3 + c2 * r)
+
+    // n = round(x / ln2), r = x - n*ln2
+    float32x4_t n = vrndaq_f32(vmulq_f32(values, inv_ln2));
+    float32x4_t r = vfmsq_n_f32(values, n, ln2);
+    // e = n << 23
+    uint32x4_t e = vshlq_n_u32(vreinterpretq_u32_s32(vcvtq_s32_f32(n)), 23);
+
+    float32x4_t r2 = vmulq_f32(r, r);
+    float32x4_t q = vfmaq_n_f32(c3, r, c2);
+    float32x4_t s = vaddq_f32(vdupq_n_f32(1.0f), r);
+    float32x4_t p = vfmaq_f32(s, q, r2);
+
+    // 2^n * p
+    float32x4_t y =
+        vreinterpretq_f32_u32(vaddq_u32(vreinterpretq_u32_f32(p), e));
+
+    y = vbslq_f32(lt_lower, vdupq_n_f32(0.0f), y);
+    y = vbslq_f32(gt_upper, vdupq_n_f32(INFINITY), y);
+    return y;
   }
   DEFINE_SLEEF_COMPATIBLE_BINARY_ELEMENTWISE_FUNC_WITH_SLEEF_NAME(
       fmod,
diff --git a/aten/src/ATen/cpu/vec/vec256/vec256_qint.h b/aten/src/ATen/cpu/vec/vec256/vec256_qint.h
index 353d3a1219cd2..9e12589b58590 100644
--- a/aten/src/ATen/cpu/vec/vec256/vec256_qint.h
+++ b/aten/src/ATen/cpu/vec/vec256/vec256_qint.h
@@ -169,6 +169,53 @@ at::vec::Vectorized<uint8_t> inline convert_float_to_int8(
   return _mm256_castsi128_si256(result);
 }
 
+template <typename T>
+at::vec::Vectorized<T> inline round_convert_float_to_int8(
+    at::vec::Vectorized<float> src);
+
+template <>
+at::vec::Vectorized<int8_t> inline round_convert_float_to_int8(
+    at::vec::Vectorized<float> src) {
+  // Convert from float32 to int32 with round nearest
+  __m256i x_values_int32 = _mm256_cvtps_epi32(src);
+
+  // Convert from int32 to int16 using signed saturation
+  __m256i xy_packed_v = _mm256_packs_epi32(x_values_int32, x_values_int32);
+
+  constexpr auto min_val = std::numeric_limits<int8_t>::min();
+  constexpr auto max_val = std::numeric_limits<int8_t>::max();
+
+  // Convert from int16 to int8 using unsigned saturation
+  __m256i xyzw_clamped_v = pack_saturate_and_clamp<int8_t>(
+      xy_packed_v, xy_packed_v, min_val, max_val);
+  __m256i permute_mask_v =
+      _mm256_set_epi32(0x07, 0x03, 0x06, 0x02, 0x05, 0x01, 0x04, 0x00);
+  return _mm256_permutevar8x32_epi32(xyzw_clamped_v, permute_mask_v);
+}
+
+template <>
+at::vec::Vectorized<uint8_t> inline round_convert_float_to_int8(
+    at::vec::Vectorized<float> src) {
+  // The type of *_val should be int32_t to ensure correct clamping behavior.
+  constexpr auto min_val = std::numeric_limits<int32_t>::min();
+  constexpr auto max_val = std::numeric_limits<int32_t>::max();
+  __m256 float32_min_val = _mm256_set1_ps(float(min_val));
+  __m256 float32_max_val = _mm256_set1_ps(float(max_val));
+  __m256 float32_src = _mm256_max_ps(src, float32_min_val);
+  float32_src = _mm256_min_ps(float32_src, float32_max_val);
+  __m256i round_src = _mm256_cvtps_epi32(float32_src);
+
+  __m128i r1 = _mm256_castsi256_si128(round_src);
+  __m128i mask = _mm_setr_epi8(
+      0, 4, 8, 12, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1);
+  __m128i r1_shuffled = _mm_shuffle_epi8(r1, mask);
+  __m128i r2 = _mm256_extractf128_si256(round_src, 1);
+  __m128i r2_shuffled = _mm_shuffle_epi8(r2, mask);
+  __m128i result = _mm_unpacklo_epi32(r1_shuffled, r2_shuffled);
+
+  return _mm256_castsi128_si256(result);
+}
+
 template <typename T>
 __FORCE_INLINE void QuantizeAvx2(
     const float* src,
diff --git a/aten/src/ATen/cpu/vec/vec512/vec512_convert.h b/aten/src/ATen/cpu/vec/vec512/vec512_convert.h
index a4adc222fa188..920c42df6b202 100644
--- a/aten/src/ATen/cpu/vec/vec512/vec512_convert.h
+++ b/aten/src/ATen/cpu/vec/vec512/vec512_convert.h
@@ -214,6 +214,24 @@ struct VecConvert<
   }
 };
 
+template <typename dst_t>
+struct VecRoundConvert<
+    dst_t,
+    1,
+    float,
+    2,
+    typename std::enable_if_t<is_8bit_integer_v<dst_t>, void>> {
+  static inline VectorizedN<dst_t, 1> apply(const VectorizedN<float, 2>& src) {
+    at::vec::Vectorized<dst_t> vec1 =
+        round_convert_float_to_int8<dst_t>(src[0]);
+    at::vec::Vectorized<dst_t> vec2 =
+        round_convert_float_to_int8<dst_t>(src[1]);
+    __m128i vec2_lo = _mm512_castsi512_si128(vec2);
+    __m512i out = _mm512_inserti32x4(vec1, vec2_lo, 1);
+    return VectorizedN<dst_t, 1>(at::vec::Vectorized<dst_t>(out));
+  }
+};
+
 template <typename dst_t>
 struct VecConvert<
     dst_t,
@@ -245,6 +263,18 @@ struct VecConvert<
   }
 };
 
+template <typename dst_t>
+struct VecRoundConvert<
+    dst_t,
+    1,
+    float,
+    1,
+    typename std::enable_if_t<is_8bit_integer_v<dst_t>, void>> {
+  static inline VectorizedN<dst_t, 1> apply(const VectorizedN<float, 1>& src) {
+    return round_convert_float_to_int8<dst_t>(src[0]);
+  }
+};
+
 template <typename src_t>
 struct VecConvert<
     float,
@@ -334,6 +364,36 @@ struct VecConvert<float, 1, Float8_e5m2, 1> {
   }
 };
 
+template <>
+struct VecConvert<float, 2, Float8_e4m3fn, 1> {
+  static inline VectorizedN<float, 2> apply(
+      const VectorizedN<Float8_e4m3fn, 1>& src_n) {
+    at::vec::Vectorized<Float8_e4m3fn> src = src_n[0];
+    __m512i values = src;
+    __m512 result0, result1;
+    cvtfp8e4m3_fp32(_mm512_castsi512_si128(values), result0);
+    cvtfp8e4m3_fp32(_mm512_extracti32x4_epi32(values, 1), result1);
+    return VectorizedN<float, 2>(
+        at::vec::Vectorized<float>(result0),
+        at::vec::Vectorized<float>(result1));
+  }
+};
+
+template <>
+struct VecConvert<Float8_e4m3fn, 1, float, 2> {
+  static inline VectorizedN<Float8_e4m3fn, 1> apply(
+      const VectorizedN<float, 2>& src_n) {
+    at::vec::Vectorized<float> src0 = src_n[0];
+    at::vec::Vectorized<float> src1 = src_n[1];
+    __m128i lane0 = cvtfp32_fp8e4m3(src0);
+    __m128i lane1 = cvtfp32_fp8e4m3(src1);
+    __m512i result = _mm512_setzero_si512();
+    result = _mm512_inserti32x4(result, lane0, 0);
+    result = _mm512_inserti32x4(result, lane1, 1);
+    return at::vec::Vectorized<Float8_e4m3fn>(result);
+  }
+};
+
 #endif
 
 } // namespace CPU_CAPABILITY
diff --git a/aten/src/ATen/cpu/vec/vec512/vec512_float8.h b/aten/src/ATen/cpu/vec/vec512/vec512_float8.h
index 0a54986d82b78..60035b7bee75c 100644
--- a/aten/src/ATen/cpu/vec/vec512/vec512_float8.h
+++ b/aten/src/ATen/cpu/vec/vec512/vec512_float8.h
@@ -17,6 +17,12 @@ inline namespace CPU_CAPABILITY {
 #if defined(CPU_CAPABILITY_AVX512) && !defined(_MSC_VER)
 
 static inline void cvtfp8e4m3_fp32(const __m128i& a, __m512& o) {
+#ifdef __AVX10_2__
+  // It takes effect only with compatible compiler and hardware in Inductor CPP
+  // backend. No effect on eager mode.
+  __m256h f16_vec = _mm256_cvthf8_ph(a);
+  o = _mm512_cvtph_ps(_mm256_castph_si256(f16_vec));
+#else
   // Zero Extend
   __m512i x = _mm512_cvtepu8_epi32(a);
   __m512i val = _mm512_and_epi32(
@@ -92,9 +98,17 @@ static inline void cvtfp8e4m3_fp32(const __m128i& a, __m512& o) {
   val = _mm512_or_si512(val, _mm512_slli_epi32(sign, 24));
 
   o = _mm512_castsi512_ps(val);
+#endif
 }
 
 static inline __m128i cvtfp32_fp8e4m3(const __m512& src) {
+#ifdef __AVX10_2__
+  // It takes effect only with compatible compiler and hardware in Inductor CPP
+  // backend. No effect on eager mode.
+  __m256i f16_vec =
+      _mm512_cvt_roundps_ph(src, _MM_FROUND_TO_NEAREST_INT | _MM_FROUND_NO_EXC);
+  return _mm256_cvtph_hf8(_mm256_castsi256_ph(f16_vec));
+#else
   // cvt 16x32 from fp32 to fp8 e4m3
   const __m512i sign_mask = _mm512_set1_epi32(0x80000000);
   const __m512i fp8_max = _mm512_set1_epi32(UINT32_C(1087) << 20);
@@ -156,6 +170,7 @@ static inline __m128i cvtfp32_fp8e4m3(const __m512& src) {
 
   // Narrow 32-bit integers to 8-bit
   return _mm512_cvtepi32_epi8(packed);
+#endif
 }
 
 static inline float fp8e4m3_to_fp32_scalar(uint8_t val) {
@@ -362,6 +377,15 @@ class Vectorizedf8 {
     }
   }
 
+  static Vectorized<T> blendv(
+      const Vectorized<T>& a,
+      const Vectorized<T>& b,
+      const Vectorized<T>& mask) {
+    auto msb_one = _mm512_set1_epi8(0xFF);
+    auto mask_ = _mm512_cmp_epu8_mask((__m512i)mask, msb_one, _MM_CMPINT_EQ);
+    return _mm512_mask_blend_epi8(mask_, (__m512i)a, (__m512i)b);
+  }
+
   Vectorized<T> abs() const {
     return _mm512_andnot_si512(_mm512_set1_epi8(0x80), values);
   }
diff --git a/aten/src/ATen/cpu/vec/vec512/vec512_qint.h b/aten/src/ATen/cpu/vec/vec512/vec512_qint.h
index e723cba61bf2d..d5899b4b4f948 100644
--- a/aten/src/ATen/cpu/vec/vec512/vec512_qint.h
+++ b/aten/src/ATen/cpu/vec/vec512/vec512_qint.h
@@ -178,6 +178,36 @@ at::vec::Vectorized<uint8_t> inline convert_float_to_int8(
   return _mm512_castsi128_si512(int8_src);
 }
 
+template <typename T>
+at::vec::Vectorized<T> inline round_convert_float_to_int8(
+    at::vec::Vectorized<float> src);
+
+template <>
+at::vec::Vectorized<int8_t> inline round_convert_float_to_int8(
+    at::vec::Vectorized<float> src) {
+  // Convert from float32 to int32 with round nearest
+  __m512i int32_src_round = _mm512_cvt_roundps_epi32(
+      src, (_MM_FROUND_TO_NEAREST_INT | _MM_FROUND_NO_EXC));
+  // Convert from int32 to int8 with saturate
+  __m128i int8_src = _mm512_cvtsepi32_epi8(int32_src_round);
+  return _mm512_castsi128_si512(int8_src);
+}
+
+template <>
+at::vec::Vectorized<uint8_t> inline round_convert_float_to_int8(
+    at::vec::Vectorized<float> src) {
+  // Clamp float32 to unsigned int range
+  constexpr auto min_val = std::numeric_limits<uint8_t>::min();
+  __m512 float32_min_val = _mm512_set1_ps(float(min_val));
+  __m512 float32_src = _mm512_max_ps(src, float32_min_val);
+  // Convert from float32 to int32 with round nearest
+  __m512i int32_src_round = _mm512_cvt_roundps_epi32(
+      float32_src, (_MM_FROUND_TO_NEAREST_INT | _MM_FROUND_NO_EXC));
+  // Convert from int32 to int8 with saturate
+  __m128i int8_src = _mm512_cvtusepi32_epi8(int32_src_round);
+  return _mm512_castsi128_si512(int8_src);
+}
+
 template <typename T>
 __FORCE_INLINE void QuantizeAvx512(
     const float* src,
diff --git a/aten/src/ATen/cpu/vec/vec_base.h b/aten/src/ATen/cpu/vec/vec_base.h
index 8741d099ea914..736005c07c584 100644
--- a/aten/src/ATen/cpu/vec/vec_base.h
+++ b/aten/src/ATen/cpu/vec/vec_base.h
@@ -68,7 +68,7 @@ Windows llvm will not have this definition.
 #define VECTOR_WIDTH 64
 #define int_vector __m512i
 #elif defined(__aarch64__) && \
-    !defined(CPU_CAPABILITY_SVE) // CPU_CAPABILITY_AVX512
+    !defined(CPU_CAPABILITY_SVE256) // CPU_CAPABILITY_AVX512
 // SVE code expects 256-vectors; leave that set for SVE?
 #if defined(__GNUC__)
 #define __at_align__ __attribute__((aligned(16)))
@@ -104,6 +104,9 @@ struct is_floating_point
 template <typename T>
 constexpr bool is_floating_point_v = is_floating_point<T>::value;
 
+template <typename T>
+constexpr bool is_integer_v = std::is_integral_v<T>;
+
 template <typename T>
 struct is_reduced_floating_point
     : std::integral_constant<
diff --git a/aten/src/ATen/cpu/vec/vec_convert.h b/aten/src/ATen/cpu/vec/vec_convert.h
index c91858641cf2f..cce9d46910e78 100644
--- a/aten/src/ATen/cpu/vec/vec_convert.h
+++ b/aten/src/ATen/cpu/vec/vec_convert.h
@@ -27,18 +27,60 @@ struct VecConvert {
   }
 };
 
+// Specialized for float-to-int vector conversion.
+// It ensures the result staying within the destination
+// type's limits without any undefined behavior.
+template <
+    typename dst_t,
+    int dst_n,
+    typename src_t,
+    int src_n,
+    typename =
+        std::enable_if_t<is_integer_v<dst_t> && is_floating_point_v<src_t>>>
+struct VecRoundConvert {
+  static inline VectorizedN<dst_t, dst_n> apply(
+      const VectorizedN<src_t, src_n>& src) {
+    constexpr int count = std::min(
+        VectorizedN<src_t, src_n>::size(), VectorizedN<dst_t, dst_n>::size());
+    __at_align__ src_t src_buf[VectorizedN<src_t, src_n>::size()];
+    src.store(src_buf);
+    __at_align__ dst_t dst_buf[VectorizedN<dst_t, dst_n>::size()];
+    constexpr auto min_val =
+        static_cast<src_t>(std::numeric_limits<dst_t>::min());
+    constexpr auto max_val =
+        static_cast<src_t>(std::numeric_limits<dst_t>::max());
+    for (int i = 0; i < count; i++) {
+      dst_buf[i] = static_cast<dst_t>(
+          std::min(std::max(std::round(src_buf[i]), min_val), max_val));
+    }
+    return VectorizedN<dst_t, dst_n>::loadu(dst_buf, count);
+  }
+};
+
 template <typename dst_t, typename src_t>
 inline std::enable_if_t<std::is_same_v<dst_t, src_t>, Vectorized<src_t>> convert(
     const Vectorized<src_t>& src) {
   return src;
 }
 
+template <typename dst_t, typename src_t>
+inline std::enable_if_t<std::is_same_v<dst_t, src_t>, Vectorized<src_t>>
+round_convert(const Vectorized<src_t>& src) {
+  return src;
+}
+
 template <typename dst_t, typename src_t>
 inline std::enable_if_t<!std::is_same_v<dst_t, src_t>, Vectorized<dst_t>>
 convert(const Vectorized<src_t>& src) {
   return VecConvert<dst_t, 1, src_t, 1>::apply(src);
 }
 
+template <typename dst_t, typename src_t>
+inline std::enable_if_t<!std::is_same_v<dst_t, src_t>, Vectorized<dst_t>>
+round_convert(const Vectorized<src_t>& src) {
+  return VecRoundConvert<dst_t, 1, src_t, 1>::apply(src);
+}
+
 template <
     typename dst_t,
     int dst_n,
@@ -49,6 +91,17 @@ inline VectorizedN<dst_t, dst_n> convert(const VectorizedN<src_t, src_n>& src) {
   return VecConvert<dst_t, dst_n, src_t, src_n>::apply(src);
 }
 
+template <
+    typename dst_t,
+    int dst_n,
+    typename src_t,
+    int src_n,
+    std::enable_if_t<dst_n != 1, int> = 0>
+inline VectorizedN<dst_t, dst_n> round_convert(
+    const VectorizedN<src_t, src_n>& src) {
+  return VecRoundConvert<dst_t, dst_n, src_t, src_n>::apply(src);
+}
+
 template <
     typename dst_t,
     int dst_n,
@@ -61,6 +114,18 @@ convert(const VectorizedN<src_t, src_n>& src) {
   return VecConvert<dst_t, dst_n, src_t, src_n>::apply(src);
 }
 
+template <
+    typename dst_t,
+    int dst_n,
+    typename src_t,
+    int src_n,
+    bool keep = false,
+    std::enable_if_t<dst_n == 1, int> = 0>
+inline std::conditional_t<keep, VectorizedN<dst_t, 1>, Vectorized<dst_t>>
+round_convert(const VectorizedN<src_t, src_n>& src) {
+  return VecRoundConvert<dst_t, dst_n, src_t, src_n>::apply(src);
+}
+
 } // namespace CPU_CAPABILITY
 
 template <
diff --git a/aten/src/ATen/cuda/Atomic.cuh b/aten/src/ATen/cuda/Atomic.cuh
index f16be30f8b71b..d38bcf9f83576 100644
--- a/aten/src/ATen/cuda/Atomic.cuh
+++ b/aten/src/ATen/cuda/Atomic.cuh
@@ -408,14 +408,7 @@ inline __device__ float gpuAtomicMul (float * address, float val) {
 
 template <typename T>
 __host__ __device__ T safe_max(T a, T b) {
-  #if defined(__HIPCC__)
-  // TODO: remove this special case for HIP when issue is fixed:
-  //       https://github.com/ROCm/hip/issues/2209
-    T max = at::_isnan(a) ? a : (at::_isnan(b) ? b : std::max<T>(a, b));
-  #else
-    T max = at::_isnan(b) ? b : std::max<T>(a, b);
-  #endif
-
+  T max = at::_isnan(b) ? b : std::max<T>(a, b);
   return max;
 }
 
@@ -468,14 +461,7 @@ inline __device__ float gpuAtomicMax(float * address, float val) {
 
 template <typename T>
 __host__ __device__ T safe_min(T a, T b) {
-  #if defined(__HIPCC__)
-  // TODO: remove this special case for HIP when issue is fixed:
-  //       https://github.com/ROCm/hip/issues/2209
-    T min = at::_isnan(a) ? a : (at::_isnan(b) ? b : std::min<T>(a, b));
-  #else
-    T min = at::_isnan(b) ? b : std::min<T>(a, b);
-  #endif
-
+  T min = at::_isnan(b) ? b : std::min<T>(a, b);
   return min;
 }
 
diff --git a/aten/src/ATen/cuda/CUDAApplyUtils.cuh b/aten/src/ATen/cuda/CUDAApplyUtils.cuh
index 1ae122f6643a6..9429c168c4119 100644
--- a/aten/src/ATen/cuda/CUDAApplyUtils.cuh
+++ b/aten/src/ATen/cuda/CUDAApplyUtils.cuh
@@ -209,10 +209,10 @@ inline void rearrangeDims(detail::TensorInfo<T1, IndexType>* aInfo,
 //     bound), and recursively call `ApplyOpN` with `remaining_steps - 1`.
 //  3. At `remaining_steps = 0`,
 //       if `step = 1`, call `op(tensor1_val, tensor2_val, ...)`;
-//       if `step > 1`, call `op(n, tensor1_val1, tensor1_val2, ..., tesor1_valstep,
-//                                  tensor2_val1, tensor2_val2, ..., tesor2_valstep,
+//       if `step > 1`, call `op(n, tensor1_val1, tensor1_val2, ..., tensor1_valstep,
+//                                  tensor2_val1, tensor2_val2, ..., tensor2_valstep,
 //                                       ...
-//                                  tensorN_val1, tensorN_val2, ..., tesorN_valstep);`
+//                                  tensorN_val1, tensorN_val2, ..., tensorN_valstep);`
 //
 // See NOTE [ CUDA_tensor_applyN helpers ] above for how Op may look like.
 
diff --git a/aten/src/ATen/cuda/CUDABlas.cpp b/aten/src/ATen/cuda/CUDABlas.cpp
index b9e66fea5ebdb..f36c87ce108bd 100644
--- a/aten/src/ATen/cuda/CUDABlas.cpp
+++ b/aten/src/ATen/cuda/CUDABlas.cpp
@@ -1956,7 +1956,6 @@ case ScalingType::TensorWise:
 
     default:
       TORCH_CHECK(false);
-      return -1;
   }
 }
 
@@ -2013,7 +2012,11 @@ void scaled_gemm(
   }
   else if (mat1_scale_dtype == kFloat8_e8m0fnu && mat2_scale_dtype == kFloat8_e8m0fnu) {
   #if ROCM_VERSION >= 70000
-            if (at::detail::getCUDAHooks().isGPUArch({"gfx950"})) {
+            std::vector<std::string> mx_archs{"gfx950"};
+  #if ROCM_VERSION >= 71400
+            mx_archs.push_back("gfx1250");
+  #endif
+            if (at::detail::getCUDAHooks().isGPUArch(mx_archs)) {
                 // TODO: add constraints based on hipblaslt internals
                 TORCH_CHECK((m % 16 == 0) && (n % 16 == 0) && (k % 128 == 0),
                            "M, N must be multiples of 16 and K should be multiple of 128 for MX format. "
diff --git a/aten/src/ATen/cuda/CUDAContextLight.h b/aten/src/ATen/cuda/CUDAContextLight.h
index b3737a823a1b3..e93d5b19d56bf 100644
--- a/aten/src/ATen/cuda/CUDAContextLight.h
+++ b/aten/src/ATen/cuda/CUDAContextLight.h
@@ -85,10 +85,11 @@ TORCH_CUDA_CPP_API c10::Allocator* getCUDADeviceAllocator();
 
 /* Handles */
 TORCH_CUDA_CPP_API cusparseHandle_t getCurrentCUDASparseHandle();
-TORCH_CUDA_CPP_API cublasHandle_t getCurrentCUDABlasHandle();
+TORCH_CUDA_CPP_API cublasHandle_t getCurrentCUDABlasHandle(bool setup = true);
 TORCH_CUDA_CPP_API cublasLtHandle_t getCurrentCUDABlasLtHandle();
 
 TORCH_CUDA_CPP_API void clearCublasWorkspaces();
+TORCH_CUDA_CPP_API void clearCublasWorkspacesForStream(cudaStream_t stream);
 struct WorkspaceMapWithMutex {
   std::map<std::tuple<void*, void*>, at::DataPtr> map;
   std::shared_mutex mutex;
diff --git a/aten/src/ATen/cuda/CUDAGeneratorImpl.cpp b/aten/src/ATen/cuda/CUDAGeneratorImpl.cpp
index 536020ffa1c86..9b690502bf7a5 100644
--- a/aten/src/ATen/cuda/CUDAGeneratorImpl.cpp
+++ b/aten/src/ATen/cuda/CUDAGeneratorImpl.cpp
@@ -82,151 +82,144 @@ Generator createCUDAGenerator(DeviceIndex device_index) {
 } // namespace cuda::detail
 
 /**
- * Creates a clone of this CUDA Generator State.
+ * Allocate GPU tensors for this capture state.
+ *
+ * We allocate on the default stream so that the caching allocator routes
+ * these tensors to the default memory pool, not the graph's capture pool.
  */
+void CUDAGeneratorCaptureState::initialize(uint64_t seed) {
+  if (is_initialized()) {
+    return;
+  }
+
+  auto options = at::TensorOptions().device(at::kCUDA).dtype(at::kLong);
+  c10::InferenceMode inference_guard(false);
+
+  // Allocate on the default stream so that the caching allocator routes
+  // these tensors to the default memory pool, not the graph's capture pool.
+  // The relaxed capture mode guard is needed because the thread-local capture
+  // mode may be Global (set by cudaStreamBeginCapture), which would block
+  // cudaMalloc even on a non-capturing stream.
+  c10::cuda::CUDAStreamCaptureModeGuard capture_mode_guard(
+      cudaStreamCaptureModeRelaxed);
+  c10::cuda::CUDAStreamGuard stream_guard(c10::cuda::getDefaultCUDAStream());
+
+  rng_state_seed_extragraph_ = at::empty({1}, options);
+  rng_state_offset_extragraph_ = at::empty({1}, options);
+
+  // Synchronize the default stream so that any prior work completes before
+  // a different stream writes to this memory.
+  c10::cuda::getDefaultCUDAStream().synchronize();
+
+  offset_intragraph_ = 0;
+}
+
+void CUDAGeneratorCaptureState::increase(uint64_t increment) {
+  // see Note [Why enforce RNG offset % 4 == 0?]
+  TORCH_INTERNAL_ASSERT(
+      offset_intragraph_ % 4 == 0, "RNG offset must be a multiple of 4.");
+  TORCH_INTERNAL_ASSERT(
+      offset_intragraph_ <= std::numeric_limits<uint64_t>::max() - increment,
+      "Increment causes overflow in the offset value.");
+  offset_intragraph_ += increment;
+}
+
+uint64_t CUDAGeneratorCaptureState::finalize() {
+  uint64_t result = offset_intragraph_;
+  offset_intragraph_ = 0;
+  return result;
+}
+
+void CUDAGeneratorCaptureState::setup_for_replay(uint64_t seed, uint64_t philox_offset) {
+  TORCH_INTERNAL_ASSERT(is_initialized(),
+      "Capture state not initialized");
+  rng_state_seed_extragraph_.fill_(static_cast<int64_t>(seed));
+  rng_state_offset_extragraph_.fill_(static_cast<int64_t>(philox_offset));
+}
+
 c10::intrusive_ptr<CUDAGeneratorState> CUDAGeneratorState::clone() {
-  return make_intrusive<CUDAGeneratorState>(
-      seed_, philox_offset_per_thread_, offset_intragraph_);
+  return make_intrusive<CUDAGeneratorState>(seed_, philox_offset_per_thread_);
 }
 
 /**
- * Function to increase the internal offset based on the specified increment.
+ * Lookup capture state for a capture ID. Returns nullptr if not found.
  */
-void CUDAGeneratorState::increase(uint64_t increment) {
-  // Rounds increment up to the nearest multiple of 4 to meet alignment
-  // requirements.
-  // see Note [Why enforce RNG offset % 4 == 0?]
-  increment = ((increment + 3) / 4) * 4;
-  // Handling different behaviors based on whether capturing is active.
-  if (at::cuda::currentStreamCaptureStatus() != at::cuda::CaptureStatus::None) {
-    // Ensures that the state is actually capturing.
-    TORCH_CHECK(
-        capturing_,
-        "Attempt to increase offset for a CUDA generator not in capture mode.");
-    // Ensures the offset is a multiple of 4
-    // see Note [Why enforce RNG offset % 4 == 0?]
-    TORCH_INTERNAL_ASSERT(
-        offset_intragraph_ % 4 == 0, "RNG offset must be a multiple of 4.");
-    // Ensures the increment does not cause overflow.
-    TORCH_INTERNAL_ASSERT(
-        offset_intragraph_ <= std::numeric_limits<uint64_t>::max() - increment,
-        "Increment causes overflow in the offset value.");
-    offset_intragraph_ += increment;
-  } else {
-    // Checks that the increment is expected outside graph capturing.
-    TORCH_CHECK(
-        !capturing_,
-        "Offset increment outside graph capture encountered unexpectedly.");
-    // Ensures the offset is a multiple of 4
-    // see Note [Why enforce RNG offset % 4 == 0?]
-    TORCH_INTERNAL_ASSERT(
-        philox_offset_per_thread_ % 4 == 0,
-        "RNG offset must be a multiple of 4.");
-    philox_offset_per_thread_ += increment;
+CUDAGeneratorCaptureState* CUDAGeneratorState::get_capture_state(CaptureId_t capture_id) {
+  std::lock_guard<std::mutex> lock(capture_states_mutex_);
+  auto it = capture_states_.find(capture_id);
+  if (it != capture_states_.end()) {
+    return it->second.get();
   }
+  return nullptr;
 }
 
 /**
- * Registers this state to a CUDA graph to manage within the graph.
+ * Create and initialize capture state for a given capture ID.
+ * Called during capture_begin for each registered generator.
  */
-void CUDAGeneratorState::register_graph(cuda::CUDAGraph* graph) {
-  // Ensures that the RNG state is not currently being captured.
-  at::cuda::assertNotCapturing(
-      "Cannot register the state during capturing stage.");
-
-  // If this is the first graph to be registered, allocate memory for the seed
-  // and offset on the GPU.
-  if (registered_graphs_.empty()) {
-    auto options = at::TensorOptions().device(at::kCUDA).dtype(at::kLong);
-    // Create these tensors outside of inference mode to ensure they can be
-    // modified in-place later. If we create them as inference tensors,
-    // subsequent fill_() calls outside inference mode
-    // will fail with "Inplace update to inference tensor outside InferenceMode".
-    c10::InferenceMode guard(false);
-    seed_extragraph_ = at::empty({1}, options);
-    offset_extragraph_ = at::empty({1}, options);
+void CUDAGeneratorState::init_capture_state(CaptureId_t capture_id) {
+  {
+    std::lock_guard<std::mutex> lock(capture_states_mutex_);
+    if (capture_states_.count(capture_id)) {
+      return;
+    }
   }
 
-  // Insert the graph into the set of registered graphs if it's not already
-  // registered.
-  if (registered_graphs_.find(graph) == registered_graphs_.end()) {
-    registered_graphs_.insert(graph);
+  auto capture_state = make_intrusive<CUDAGeneratorCaptureState>();
+  capture_state->initialize(seed_);
+
+  std::lock_guard<std::mutex> lock(capture_states_mutex_);
+  if (!capture_states_.count(capture_id)) {
+    capture_states_[capture_id] = std::move(capture_state);
   }
 }
 
 /**
- * Unregisters a CUDA graph from the RNG state.
+ * Function to increase the internal offset based on the specified increment.
  */
-void CUDAGeneratorState::unregister_graph(cuda::CUDAGraph* graph) {
-  // Verify the graph was previously registered.
-  TORCH_CHECK(
-      registered_graphs_.find(graph) != registered_graphs_.end(),
-      "The graph should be registered to the state");
-
-  // Remove the graph from the set of registered graphs.
-  registered_graphs_.erase(graph);
-
-  // If no more graphs are registered, deallocate the GPU memory for the seed
-  // and offset.
-  if (registered_graphs_.empty()) {
-    seed_extragraph_.reset();
-    offset_extragraph_.reset();
+void CUDAGeneratorState::increase(uint64_t increment) {
+  // see Note [Why enforce RNG offset % 4 == 0?]
+  increment = ((increment + 3) / 4) * 4;
+
+  auto capture_id = at::cuda::currentStreamCaptureId();
+  if (capture_id.has_value()) {
+    auto* capture_state = get_capture_state(capture_id.value());
+    TORCH_CHECK(capture_state != nullptr,
+        "RNG op during graph capture but generator is not registered with "
+        "the capturing graph. Call graph.register_generator_state() before "
+        "capture_begin().");
+    capture_state->increase(increment);
+  } else {
+    TORCH_INTERNAL_ASSERT(
+        philox_offset_per_thread_ % 4 == 0,
+        "RNG offset must be a multiple of 4.");
+    philox_offset_per_thread_ += increment;
   }
 }
 
-/**
- * Note [Explicit Registration of Generators to the CUDA Graph]
- * ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
- *
- * Ideally, it would be more user-friendly if the state could be exchanged and generators
- * could be registered with the CUDA graph implicitly. However, resetting GPU tensors during
- * the capture stage causes these reset operations to be recorded within the CUDA graph.
- * This behavior is undesirable because we do not want these tensors to be reset during
- * the replay stage of the graph.
- *
- * As of now, there is no available method to perform a CUDA operation during the graph's
- * recording phase without having that operation be included in the CUDA graph.
- * This limitation necessitates explicit user action to register generators with the graph.
- * By requiring users to manually register their generators, we can ensure that state resets
- * (capture_prologue) only occur before the graph capture begins, thus avoiding unintended
- * resets during the replay of the graph. See https://github.com/pytorch/pytorch/pull/114068.
- */
-
-/**
- * Performs the prologue steps for capturing a CUDA graph state.
- * This method is intended to reset graph-related state variables before capturing begins.
- */
-void CUDAGeneratorState::capture_prologue() {
-  capturing_ = true;
-  offset_intragraph_ = 0;
-  seed_extragraph_.fill_(static_cast<int64_t>(seed_));
-  offset_extragraph_.fill_(0);
+uint64_t CUDAGeneratorState::capture_epilogue(CaptureId_t capture_id) {
+  auto* capture_state = get_capture_state(capture_id);
+  if (capture_state) {
+    return capture_state->finalize();
+  }
+  return 0;
 }
 
-/**
- * Ends the capturing phase and resets related variables, returning the whole
- * graph increment.
- */
-uint64_t CUDAGeneratorState::capture_epilogue() {
-  capturing_ = false;
-  return offset_intragraph_;
+void CUDAGeneratorState::remove_capture_state(CaptureId_t capture_id) {
+  std::lock_guard<std::mutex> lock(capture_states_mutex_);
+  capture_states_.erase(capture_id);
 }
 
-/**
- * Prepares the state for replay by setting initial state tensors and applying
- * total increment.
- */
-void CUDAGeneratorState::replay_prologue(uint64_t wholegraph_increment) {
-  // Ensures the generator is not in capturing mode.
-  at::cuda::assertNotCapturing(
-      "Cannot prepare for replay during capturing stage.");
-  if (wholegraph_increment) {
-      seed_extragraph_.fill_(static_cast<int64_t>(seed_));
-      offset_extragraph_.fill_(static_cast<int64_t>(philox_offset_per_thread_));
-      // Applies the total increment achieved during previous captures to update the
-      // offset.
-      increase(wholegraph_increment);
+void CUDAGeneratorState::replay_prologue(CaptureId_t capture_id, uint64_t wholegraph_increment) {
+  if (wholegraph_increment == 0) {
+    return;
   }
+
+  auto* capture_state = get_capture_state(capture_id);
+  TORCH_INTERNAL_ASSERT(capture_state != nullptr,
+      "replay_prologue called but no capture state found for this capture_id");
+  capture_state->setup_for_replay(seed_, philox_offset_per_thread_);
+  philox_offset_per_thread_ += wholegraph_increment;
 }
 
 /**
@@ -271,7 +264,8 @@ CUDAGeneratorImpl::CUDAGeneratorImpl(
  * See Note [Acquire lock when using random generators]
  */
 void CUDAGeneratorImpl::set_current_seed(uint64_t seed) {
-  if (C10_LIKELY(at::cuda::currentStreamCaptureStatus() == at::cuda::CaptureStatus::None)) {
+  auto capture_id = at::cuda::currentStreamCaptureId();
+  if (C10_LIKELY(!capture_id.has_value())) {
     state_->seed_ = seed;
     state_->philox_offset_per_thread_ = 0;
     no_reset_rnn_state_.clear();
@@ -410,10 +404,14 @@ void CUDAGeneratorImpl::set_philox_offset_per_thread(uint64_t offset) {
   // set_philox_offset_per_thread instead of set_offset will cause the
   // cudnn RNN rng state to become stale.
   TORCH_CHECK(offset % 4 == 0, "offset must be a multiple of 4");
-  if (C10_LIKELY(at::cuda::currentStreamCaptureStatus() == at::cuda::CaptureStatus::None)) {
+  auto capture_id = at::cuda::currentStreamCaptureId();
+  if (C10_LIKELY(!capture_id.has_value())) {
     state_->philox_offset_per_thread_ = offset;
   } else {
-    state_->offset_intragraph_ = offset;
+    auto* capture_state = state_->get_capture_state(capture_id.value());
+    TORCH_CHECK(capture_state != nullptr,
+        "Generator not registered with the capturing graph.");
+    capture_state->offset_intragraph_ = offset;
   }
 }
 
@@ -421,26 +419,19 @@ void CUDAGeneratorImpl::set_philox_offset_per_thread(uint64_t offset) {
  * Gets the current philox_offset_per_thread_ of CUDAGeneratorImpl.
  */
 uint64_t CUDAGeneratorImpl::philox_offset_per_thread() const {
-  if (C10_LIKELY(at::cuda::currentStreamCaptureStatus() == at::cuda::CaptureStatus::None)) {
+  auto capture_id = at::cuda::currentStreamCaptureId();
+  if (C10_LIKELY(!capture_id.has_value())) {
     return state_->philox_offset_per_thread_;
   } else {
-    return state_->offset_intragraph_;
+    auto* capture_state = state_->get_capture_state(capture_id.value());
+    TORCH_CHECK(capture_state != nullptr,
+        "Generator not registered with the capturing graph.");
+    return capture_state->offset_intragraph_;
   }
 }
 
-/**
- * Registers this state to a CUDA graph to manage within the graph.
- */
 void CUDAGeneratorImpl::register_graph(cuda::CUDAGraph* graph) {
   graph->register_generator_state(state_);
-  state_->register_graph(graph);
-}
-
-/**
- * Unregisters a CUDA graph from the RNG state.
- */
-void CUDAGeneratorImpl::unregister_graph(cuda::CUDAGraph* graph) {
-  state_->unregister_graph(graph);
 }
 
 /**
@@ -465,12 +456,20 @@ void CUDAGeneratorImpl::unregister_graph(cuda::CUDAGraph* graph) {
  * See Note [Acquire lock when using random generators]
  */
 PhiloxCudaState CUDAGeneratorImpl::philox_cuda_state(uint64_t increment) {
-  if (at::cuda::currentStreamCaptureStatus() != at::cuda::CaptureStatus::None) {
-    uint64_t offset = state_->offset_intragraph_;
+  auto capture_id = at::cuda::currentStreamCaptureId();
+  if (capture_id.has_value()) {
+    auto* capture_state = state_->get_capture_state(capture_id.value());
+    TORCH_CHECK(capture_state != nullptr,
+        "RNG op during graph capture but generator is not registered with "
+        "the capturing graph. Call graph.register_generator_state() before "
+        "capture_begin().");
+
+    uint64_t offset = capture_state->offset_intragraph_;
     state_->increase(increment);
+
     return PhiloxCudaState(
-        state_->seed_extragraph_.data_ptr<int64_t>(),
-        state_->offset_extragraph_.data_ptr<int64_t>(),
+        capture_state->rng_state_seed_extragraph_.data_ptr<int64_t>(),
+        capture_state->rng_state_offset_extragraph_.data_ptr<int64_t>(),
         offset);
   } else {
     uint64_t offset = state_->philox_offset_per_thread_;
diff --git a/aten/src/ATen/cuda/CUDAGeneratorImpl.h b/aten/src/ATen/cuda/CUDAGeneratorImpl.h
index d4ab49382e7ff..fe62fc53eef07 100644
--- a/aten/src/ATen/cuda/CUDAGeneratorImpl.h
+++ b/aten/src/ATen/cuda/CUDAGeneratorImpl.h
@@ -4,15 +4,20 @@
 #include <ATen/core/Generator.h>
 #include <ATen/core/TensorBase.h>
 #include <ATen/cuda/PhiloxCudaState.h>
+#include <c10/core/Allocator.h>
+#include <c10/util/flat_hash_map.h>
 #include <atomic>
 #include <memory>
-#include <unordered_set>
+#include <mutex>
+
 namespace at {
 
 namespace cuda {
 struct CUDAGraph;
 }
 
+using CaptureId_t = c10::CaptureId_t;
+
 /**
  * Note [CUDA Graph-safe RNG states]
  * ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
@@ -93,32 +98,51 @@ struct CUDAGraph;
  *
  */
 
+/**
+ * Per-capture state for a generator.
+ * Each (generator, capture_id) pair gets its own CUDAGeneratorCaptureState.
+ * This holds the GPU tensors and offset tracking for a specific graph capture.
+ */
+struct CUDAGeneratorCaptureState : public c10::intrusive_ptr_target {
+  uint64_t offset_intragraph_{0};
+  at::TensorBase rng_state_seed_extragraph_;
+  at::TensorBase rng_state_offset_extragraph_;
+
+  CUDAGeneratorCaptureState() = default;
+
+  bool is_initialized() const { return rng_state_seed_extragraph_.defined(); }
+  void initialize(uint64_t seed);
+  void increase(uint64_t increment);
+  uint64_t finalize();
+  void setup_for_replay(uint64_t seed, uint64_t philox_offset);
+};
+
+/**
+ * Generator state that supports multiple concurrent graph captures.
+ * Each capture gets its own CUDAGeneratorCaptureState keyed by CaptureId_t.
+ */
 struct CUDAGeneratorState : public c10::intrusive_ptr_target {
   uint64_t seed_;
   uint64_t philox_offset_per_thread_;
-  uint64_t offset_intragraph_;
-  bool capturing_{};
-  std::unordered_set<cuda::CUDAGraph*> registered_graphs_;
-  at::TensorBase seed_extragraph_;
-  at::TensorBase offset_extragraph_;
+
+  // Map from capture ID to per-capture state
+  ska::flat_hash_map<CaptureId_t, c10::intrusive_ptr<CUDAGeneratorCaptureState>> capture_states_;
+  mutable std::mutex capture_states_mutex_;
 
   CUDAGeneratorState(
       uint64_t seed = default_rng_seed_val,
-      uint64_t philox_offset_per_thread = 0,
-      uint64_t offset_intragraph = 0)
+      uint64_t philox_offset_per_thread = 0)
       : seed_(seed),
-        philox_offset_per_thread_(philox_offset_per_thread),
-        offset_intragraph_(offset_intragraph) {}
+        philox_offset_per_thread_(philox_offset_per_thread) {}
 
   void increase(uint64_t increment);
 
-  void register_graph(cuda::CUDAGraph* graph);
-  void unregister_graph(cuda::CUDAGraph* graph);
+  CUDAGeneratorCaptureState* get_capture_state(CaptureId_t capture_id);
+  void init_capture_state(CaptureId_t capture_id);
+  uint64_t capture_epilogue(CaptureId_t capture_id);
+  void replay_prologue(CaptureId_t capture_id, uint64_t wholegraph_increment);
+  void remove_capture_state(CaptureId_t capture_id);
 
-  void capture_prologue();
-  // capture_epilogue returns the wholegraph_increment
-  uint64_t capture_epilogue();
-  void replay_prologue(uint64_t wholegraph_increment);
   c10::intrusive_ptr<CUDAGeneratorState> clone();
 };
 
@@ -147,7 +171,6 @@ struct TORCH_CUDA_CPP_API CUDAGeneratorImpl : public c10::GeneratorImpl {
   uint64_t philox_offset_per_thread() const;
 
   void register_graph(cuda::CUDAGraph* graph);
-  void unregister_graph(cuda::CUDAGraph* graph);
 
   // Generates a PhiloxCudaState with a specified increment, and increment
   // current state
diff --git a/aten/src/ATen/cuda/CUDAGraph.cpp b/aten/src/ATen/cuda/CUDAGraph.cpp
index 3de3cc3f5259c..82ddb68566651 100644
--- a/aten/src/ATen/cuda/CUDAGraph.cpp
+++ b/aten/src/ATen/cuda/CUDAGraph.cpp
@@ -1,17 +1,40 @@
 #include <ATen/core/CachingHostAllocator.h>
+#include <ATen/cuda/CUDAContextLight.h>
 #include <ATen/cuda/CUDAGeneratorImpl.h>
 #include <ATen/cuda/CUDAGraph.h>
+#include <ATen/cuda/CUDAGraphsUtils.cuh>
 #include <ATen/cuda/Exceptions.h>
 #include <ATen/cuda/MemPool.h>
 #include <ATen/Functions.h>
+#include <c10/cuda/CUDAAllocatorConfig.h>
 #include <c10/cuda/CUDAFunctions.h>
 
 #include <cstddef>
+#include <optional>
 
 namespace at::cuda {
 
 static bool _cuda_graphs_debug = false;
 
+// To support stream capture across multiple threads, we use a global
+// hashmap mapping cuda stream capture IDs to CUDAGraph objects. This
+// was originally a thread_local std::stack<CUDAGraph*>, but that was
+// not acceptable since stream capture does span threads in certain
+// circumstances (in particular, during autograd).
+static std::mutex _currently_capturing_graphs_mutex;
+static ska::flat_hash_map<CaptureId_t, CUDAGraph*> _currently_capturing_graphs;
+
+
+#if defined(USE_ROCM)
+// Returns true when at least one CUDAGraph capture is currently active in this
+// process. Uses the same mutex-protected capture map as capture lifecycle
+// bookkeeping.
+bool is_graph_capture_active() {
+  std::unique_lock<std::mutex> lock(_currently_capturing_graphs_mutex);
+  return !_currently_capturing_graphs.empty();
+}
+#endif // defined(USE_ROCM)
+
 MempoolId_t graph_pool_handle() {
   // Sets just the second value, to distinguish it from MempoolId_ts created from
   // cudaStreamGetCaptureInfo id_s in capture_begin.
@@ -57,21 +80,36 @@ void CUDAGraph::register_generator_state(const at::Generator& generator) {
   cuda_gen->register_graph(this);
 }
 
-void CUDAGraph::capture_begin(MempoolId_t pool/*=0*/, cudaStreamCaptureMode capture_mode) {
+
+template <>
+std::function<bool(cudaStream_t)> CUDAGraph::create_allocate_filter<cudaStream_t>() const {
+  return [this](cudaStream_t stream) {
+    auto capture_id_opt = c10::cuda::captureIdMayInitCtx(stream);
+    return capture_id_opt.has_value() && capture_id_opt.value() == capture_id_;
+  };
+}
+
+template <>
+std::function<bool(c10::Stream)> CUDAGraph::create_allocate_filter<c10::Stream>() const {
+  return [this](c10::Stream stream) {
+    cudaStream_t cuda_stream = CUDAStream(CUDAStream::UNCHECKED, stream);
+    auto capture_id_opt = c10::cuda::captureIdMayInitCtx(cuda_stream);
+    return capture_id_opt.has_value() && capture_id_opt.value() == capture_id_;
+  };
+}
+
+void CUDAGraph::capture_begin(MempoolId_t pool/*={0,0}*/, cudaStreamCaptureMode capture_mode) {
   TORCH_CHECK(!has_graph_exec_,
               "This CUDAGraph instance already owns a captured graph. "
               "To capture a new graph, create a new instance.");
 
+  capture_mode_ = capture_mode;
+
   // default generator is always registered
   auto* gen = get_generator_or_default<CUDAGeneratorImpl>(
       std::nullopt, cuda::detail::getDefaultCUDAGenerator());
   gen->register_graph(this);
 
-  for (auto& [generator_state, wholegraph_increments] :
-       captured_generator_states_) {
-    generator_state->capture_prologue();
-  }
-
   auto stream = at::cuda::getCurrentCUDAStream();
 
   TORCH_CHECK(stream != at::cuda::getDefaultCUDAStream(),
@@ -82,6 +120,16 @@ void CUDAGraph::capture_begin(MempoolId_t pool/*=0*/, cudaStreamCaptureMode capt
   capture_stream_ = stream;
   capture_dev_ = c10::cuda::current_device();
 
+#if defined(USE_ROCM)
+  // hipBLASLt handles are per-(device, stream) on ROCm and lazily created.
+  // Ensure the handle for the intended capture stream exists before
+  // capture begins, because hipblasLtCreate performs internal allocations
+  // that are not allowed once stream capture is active.
+  if (at::globalContext().blasPreferredBackend() == at::BlasBackend::Cublaslt) {
+    (void)at::cuda::getCurrentCUDABlasLtHandle();
+  }
+#endif
+
   if (pool.first != 0 || pool.second != 0) {
     // Either value being nonzero means the user supplied a pool to share.
     // But only one should be nonzero.
@@ -96,31 +144,37 @@ void CUDAGraph::capture_begin(MempoolId_t pool/*=0*/, cudaStreamCaptureMode capt
     TORCH_INTERNAL_ASSERT(mempool_id_.first > 0);
   }
 
-  auto filter = [this](cudaStream_t stream) {
-    cudaStreamCaptureStatus status{};
-    CaptureId_t stream_capture_id = 0;
-    AT_CUDA_CHECK(cudaStreamGetCaptureInfo(stream, &status, &stream_capture_id));
-    return status == cudaStreamCaptureStatus::cudaStreamCaptureStatusActive && stream_capture_id == capture_id_;
-  };
-
   // Addendum: beginAllocateStreamToPool is now called before cudaStreamBeginCapture to prevent an
   // autograd thread's free() call triggering an invalid cudaEventRecord in the caching allocator
   // due to the capture status being updated _after_ a capture had already started.
-  c10::cuda::CUDACachingAllocator::beginAllocateToPool(capture_dev_, mempool_id_, filter);
+  c10::cuda::CUDACachingAllocator::beginAllocateToPool(capture_dev_, mempool_id_, create_allocate_filter<cudaStream_t>());
 
-  at::getHostAllocator(at::kCUDA)->begin_allocate_to_pool(mempool_id_, [filter](c10::Stream stream) {
-    return filter(CUDAStream(CUDAStream::UNCHECKED, stream));
-  });
+  at::getHostAllocator(at::kCUDA)->begin_allocate_to_pool(mempool_id_, create_allocate_filter<c10::Stream>());
+
+  // The pool is now acquired and being recorded to. Track this so reset() can
+  // release it even if the capture fails before capture_end() completes.
+  allocated_pool_ = true;
+  capturing_to_pool_ = true;
 
   // cudaStreamCaptureModeGlobal is the most conservative option to
   // prevent potentially unsafe CUDA API calls during capture.  See
   // https://docs.nvidia.com/cuda/cuda-runtime-api/group__CUDART__STREAM.html#group__CUDART__STREAM_1g9d0535d93a214cbf126835257b16ba85
   AT_CUDA_CHECK(cudaStreamBeginCapture(capture_stream_, capture_mode));
 
-  cudaStreamCaptureStatus status{};
-  AT_CUDA_CHECK(cudaStreamGetCaptureInfo(stream, &status, &capture_id_));
-  TORCH_INTERNAL_ASSERT(status == cudaStreamCaptureStatus::cudaStreamCaptureStatusActive);
+  auto capture_id_opt = c10::cuda::captureIdMayInitCtx(stream);
+  TORCH_INTERNAL_ASSERT(capture_id_opt.has_value(),
+      "Stream should be actively capturing after cudaStreamBeginCapture");
+  capture_id_ = capture_id_opt.value();
 
+  {
+    std::lock_guard<std::mutex> lock(_currently_capturing_graphs_mutex);
+    _currently_capturing_graphs.emplace(capture_id_, this);
+  }
+
+  for (auto& [generator_state, wholegraph_increment] :
+       captured_generator_states_) {
+    generator_state->init_capture_state(capture_id_);
+  }
 }
 
 void CUDAGraph::capture_end() {
@@ -129,16 +183,30 @@ void CUDAGraph::capture_end() {
   TORCH_CHECK(stream.stream() == capture_stream_.stream(),
               "Capture must end on the same stream it began on.");
 
-  AT_CUDA_CHECK(cudaStreamEndCapture(capture_stream_, &graph_));
+  // Capture is over once cudaStreamEndCapture returns (success or failure).
+  // Clear bookkeeping before propagating the return status so watchdog-side
+  // checks cannot observe stale "capture active" state on error paths.
+  cudaError_t endCaptureErr = cudaStreamEndCapture(capture_stream_, &graph_);
+  {
+    std::unique_lock<std::mutex> lock(_currently_capturing_graphs_mutex);
+    TORCH_CHECK(
+        _currently_capturing_graphs.count(capture_id_),
+        "capture_end() called before capture_begin().");
+    _currently_capturing_graphs.erase(capture_id_);
+  }
 
   c10::cuda::CUDACachingAllocator::endAllocateToPool(capture_dev_, mempool_id_);
   at::getHostAllocator(at::kCUDA)->end_allocate_to_pool(mempool_id_);
+  // Allocation recording has stopped (even if endCaptureErr is a failure), so
+  // reset() must not end the pool again.
+  capturing_to_pool_ = false;
+  AT_CUDA_CHECK(endCaptureErr);
 
   TORCH_CHECK(graph_ != nullptr, "Invalid capture.");
 
-  for (auto& [generator_state, wholegraph_increments] :
+  for (auto& [generator_state, wholegraph_increment] :
        captured_generator_states_) {
-    wholegraph_increments = generator_state->capture_epilogue();
+    wholegraph_increment = generator_state->capture_epilogue(capture_id_);
   }
 
   size_t numCUDAGraphNodes = 0;
@@ -200,9 +268,9 @@ void CUDAGraph::replay() {
 
   c10::OptionalDeviceGuard device_guard{capture_stream_.device()};
 
-  for (auto& [generator_state, wholegraph_increments] :
+  for (auto& [generator_state, wholegraph_increment] :
        captured_generator_states_) {
-    generator_state->replay_prologue(wholegraph_increments);
+    generator_state->replay_prologue(capture_id_, wholegraph_increment);
   }
   // graph_exec_ may be replayed in any stream.
   AT_CUDA_CHECK(cudaGraphLaunch(graph_exec_, at::cuda::getCurrentCUDAStream()));
@@ -261,11 +329,39 @@ void CUDAGraph::reset() {
   // and the allocator could end up in all kinds of weird states depending where failure occurred.
   // If the user catches the failure exception in a script, or is running in REPL or (god forbid)
   // a Jupyter notebook, I don't see an easy way for reset() to gracefully fix all such possible error states.
-  if (capture_ended_) {
+
+  if (capture_id_ != 0) {
+    for (auto& [generator_state, wholegraph_increment] : captured_generator_states_) {
+      generator_state->remove_capture_state(capture_id_);
+    }
+  }
+  captured_generator_states_.clear();
+
+  if (capture_id_ != 0) {
+    std::lock_guard<std::mutex> lock(_currently_capturing_graphs_mutex);
+    _currently_capturing_graphs.erase(capture_id_);
+    capture_id_ = 0;
+  }
+
+  if (allocated_pool_) {
+    if (capturing_to_pool_) {
+      // Capture was abandoned before capture_end() ran, so the allocator is
+      // still routing allocations to this pool. Stop that before releasing so
+      // the pool is left in a consistent, freeable state.
+      c10::cuda::CUDACachingAllocator::endAllocateToPool(capture_dev_, mempool_id_);
+      at::getHostAllocator(at::kCUDA)->end_allocate_to_pool(mempool_id_);
+      capturing_to_pool_ = false;
+    }
+
+    // Clean up cuBLAS workspaces allocated on the capture stream, otherwise live allocations prevent
+    // private pool cleanup
+    clearCublasWorkspacesForStream(capture_stream_.stream());
+
     // notifyCaptureDestroy may throw. How should we handle this?
     c10::cuda::CUDACachingAllocator::releasePool(capture_dev_, mempool_id_);
     at::getHostAllocator(at::kCUDA)->release_pool(mempool_id_);
     capture_ended_ = false;
+    allocated_pool_ = false;
   }
   if (has_graph_) {
     C10_CUDA_CHECK_WARN(cudaGraphDestroy(graph_));
@@ -279,16 +375,12 @@ void CUDAGraph::reset() {
 
 // Returns an id another graph's capture_begin can use to share the same memory pool as this graph.
 MempoolId_t CUDAGraph::pool() {
-TORCH_CHECK(capture_ended_,
+  TORCH_CHECK(capture_ended_,
               "Called CUDAGraph::pool() without a preceding successful capture.");
   return mempool_id_;
 }
 
 CUDAGraph::~CUDAGraph() {
-  for (auto& [generator_state, wholegraph_increments] :
-       captured_generator_states_) {
-    generator_state->unregister_graph(this);
-  }
   reset();
 
 // There are recent HIP changes where hipGraphExecDestroy doesn't immediately free memory.
@@ -304,4 +396,202 @@ CUDAGraph::~CUDAGraph() {
 #endif
 }
 
+CUDAGraph* CUDAGraph::get_currently_capturing_graph() {
+  std::unique_lock<std::mutex> lock(_currently_capturing_graphs_mutex);
+  auto capture_id_opt = c10::cuda::currentStreamCaptureIdMayInitCtx();
+  TORCH_CHECK(
+      capture_id_opt.has_value(),
+      "The current stream is not currently capturing.");
+  TORCH_CHECK(
+      _currently_capturing_graphs.count(capture_id_opt.value()),
+      "get_currently_capturing_graph() can be used only between capture_begin() and capture_end(). Did you use a stream without making it depend upon the original stream used for capture?");
+  return _currently_capturing_graphs.at(capture_id_opt.value());
+}
+
+void CUDAGraph::begin_capture_to_if_node(
+    const at::Tensor& scalar_cuda_pred_tensor) {
+#if !defined(USE_ROCM) && (defined(CUDA_VERSION) && CUDA_VERSION >= 12040)
+  TORCH_CHECK(
+      !has_graph_exec_,
+      "This CUDAGraph instance already owns a captured graph.");
+
+  TORCH_CHECK(!c10::cuda::CUDACachingAllocator::CUDAAllocatorConfig::graph_capture_record_stream_reuse(), "'graph_capture_record_stream_reuse:True' allocator config does not work with conditional control flow in a cuda graph today. See issue #175001 for updates");
+
+  cudaStreamCaptureStatus status{};
+  cudaGraph_t currently_capturing_graph{};
+  AT_CUDA_CHECK(cudaStreamGetCaptureInfo(
+      getCurrentCUDAStream(), &status, nullptr, &currently_capturing_graph));
+  TORCH_CHECK(
+      status == cudaStreamCaptureStatusActive,
+      "capture_begin() must be called before begin_capture_to_if_node()");
+  cudaGraphConditionalHandle handle{};
+  AT_CUDA_CHECK(cudaGraphConditionalHandleCreate(
+      &handle, currently_capturing_graph, 0, 0));
+
+  set_conditional_handle(handle, scalar_cuda_pred_tensor);
+
+  const cudaGraphNode_t* dependencies{};
+  const cudaGraphEdgeData* dependency_edges{};
+  size_t num_dependencies = 0;
+#if CUDA_VERSION >= 13000
+  AT_CUDA_CHECK(cudaStreamGetCaptureInfo(
+      getCurrentCUDAStream(),
+      &status,
+      nullptr,
+      &currently_capturing_graph,
+      &dependencies,
+      &dependency_edges,
+      &num_dependencies));
+#else
+  AT_CUDA_CHECK(cudaStreamGetCaptureInfo_v3(
+      getCurrentCUDAStream(),
+      &status,
+      nullptr,
+      &currently_capturing_graph,
+      &dependencies,
+      &dependency_edges,
+      &num_dependencies
+  ));
+#endif
+  TORCH_CHECK(status == cudaStreamCaptureStatusActive);
+
+  cudaGraphNodeParams params{};
+  params.type = cudaGraphNodeTypeConditional;
+  params.conditional.handle = handle;
+  params.conditional.type = cudaGraphCondTypeIf;
+  params.conditional.size = 1;
+
+  cudaGraphNode_t cond_node{};
+#if CUDA_VERSION >= 13000
+  AT_CUDA_CHECK(cudaGraphAddNode(
+      &cond_node,
+      currently_capturing_graph,
+      dependencies,
+      dependency_edges,
+      num_dependencies,
+      &params));
+#else
+  AT_CUDA_CHECK(cudaGraphAddNode_v2(
+      &cond_node,
+      currently_capturing_graph,
+      dependencies,
+      dependency_edges,
+      num_dependencies,
+      &params));
+#endif
+  cudaGraph_t if_node_child_graph = params.conditional.phGraph_out[0];
+
+#if CUDA_VERSION >= 13000
+  AT_CUDA_CHECK(cudaStreamUpdateCaptureDependencies(
+getCurrentCUDAStream(), &cond_node, nullptr, 1, cudaStreamSetCaptureDependencies));
+#else
+  AT_CUDA_CHECK(cudaStreamUpdateCaptureDependencies_v2(
+getCurrentCUDAStream(), &cond_node, nullptr, 1, cudaStreamSetCaptureDependencies));
+#endif
+
+  CUDAStream child_stream = getStreamFromPool();
+  conditional_graph_capture_ids_.push(0);
+
+  c10::cuda::CUDACachingAllocator::endAllocateToPool(capture_dev_, mempool_id_);
+  at::getHostAllocator(at::kCUDA)->end_allocate_to_pool(mempool_id_);
+  c10::cuda::CUDACachingAllocator::beginAllocateToPool(
+      capture_dev_, mempool_id_, create_child_allocate_filter());
+  auto filter = create_child_allocate_filter();
+  at::getHostAllocator(at::kCUDA)->begin_allocate_to_pool(mempool_id_, [filter](c10::Stream stream) {
+    return filter(CUDAStream(CUDAStream::UNCHECKED, stream));
+  });
+
+  AT_CUDA_CHECK(cudaStreamBeginCaptureToGraph(
+      child_stream, if_node_child_graph, nullptr, nullptr, 0, capture_mode_));
+
+  auto child_capture_id_opt = c10::cuda::captureIdMayInitCtx(child_stream);
+  TORCH_INTERNAL_ASSERT(child_capture_id_opt.has_value(),
+      "Child stream should be actively capturing after cudaStreamBeginCaptureToGraph");
+  conditional_graph_capture_ids_.top() = child_capture_id_opt.value();
+
+  conditional_node_streams_.emplace(child_stream);
+
+  {
+    std::unique_lock<std::mutex> lock(_currently_capturing_graphs_mutex);
+    _currently_capturing_graphs.emplace(
+        conditional_graph_capture_ids_.top(), this);
+  }
+
+#else // !defined(USE_ROCM) && (defined(CUDA_VERSION) && CUDA_VERSION >= 12040)
+  AT_ERROR(
+      __func__,
+      " CUDA Graphs conditional nodes are not supported for cuda version < 12.4");
+  return;
+#endif
+}
+
+void CUDAGraph::end_capture_to_conditional_node() {
+#if !defined(USE_ROCM) && (defined(CUDA_VERSION) && CUDA_VERSION >= 12040)
+  TORCH_INTERNAL_ASSERT(
+      !conditional_graph_capture_ids_.empty(),
+      "Missing capture ID for conditional node.");
+
+  CaptureId_t child_capture_id = conditional_graph_capture_ids_.top();
+  bool rng_or_generators_changed = false;
+  for (const auto& [generator_state, wholegraph_increment] :
+       captured_generator_states_) {
+    if (generator_state->get_capture_state(child_capture_id) != nullptr) {
+      rng_or_generators_changed = true;
+      break;
+    }
+  }
+
+  {
+    std::unique_lock<std::mutex> lock(_currently_capturing_graphs_mutex);
+    TORCH_CHECK(
+        _currently_capturing_graphs.count(child_capture_id),
+        "capture_end() called before capture_begin().");
+    _currently_capturing_graphs.erase(child_capture_id);
+  }
+
+  CUDAStream stream = conditional_node_streams_.top().current_stream();
+  AT_CUDA_CHECK(cudaStreamEndCapture(stream.stream(), nullptr));
+  conditional_node_streams_.pop();
+  conditional_graph_capture_ids_.pop();
+
+  c10::cuda::CUDACachingAllocator::endAllocateToPool(capture_dev_, mempool_id_);
+  at::getHostAllocator(at::kCUDA)->end_allocate_to_pool(mempool_id_);
+  if (conditional_graph_capture_ids_.empty()) {
+    c10::cuda::CUDACachingAllocator::beginAllocateToPool(
+        capture_dev_, mempool_id_, create_allocate_filter<cudaStream_t>());
+    at::getHostAllocator(at::kCUDA)->begin_allocate_to_pool(mempool_id_, create_allocate_filter<c10::Stream>());
+  } else {
+    c10::cuda::CUDACachingAllocator::beginAllocateToPool(
+        capture_dev_, mempool_id_, create_child_allocate_filter());
+    auto filter = create_child_allocate_filter();
+    at::getHostAllocator(at::kCUDA)->begin_allocate_to_pool(mempool_id_, [filter](c10::Stream stream) {
+      return filter(CUDAStream(CUDAStream::UNCHECKED, stream));
+    });
+  }
+  constexpr const char* rng_with_conditional_nodes_error =
+      "RNG within data-dependent conditional nodes is not supported yet.";
+  TORCH_CHECK(!rng_or_generators_changed, rng_with_conditional_nodes_error);
+
+#else // !defined(USE_ROCM) && (defined(CUDA_VERSION) && CUDA_VERSION >= 12040)
+  AT_ERROR(
+      __func__,
+      " CUDA Graphs conditional nodes are not supported for cuda version < 12.4");
+#endif
+}
+
+std::function<bool(cudaStream_t)> CUDAGraph::create_child_allocate_filter() {
+#if !defined(USE_ROCM) && (defined(CUDA_VERSION) && CUDA_VERSION >= 12040)
+  return [&current_capture_id = conditional_graph_capture_ids_.top()](cudaStream_t stream) {
+      auto capture_id_opt = c10::cuda::captureIdMayInitCtx(stream);
+      return capture_id_opt.has_value() && capture_id_opt.value() == current_capture_id;
+  };
+#else // !defined(USE_ROCM) && (defined(CUDA_VERSION) && CUDA_VERSION >= 12040)
+  AT_ERROR(
+      __func__,
+      " CUDA Graphs conditional nodes are not supported for cuda version < 12.4");
+  return std::function<bool(cudaStream_t)>();
+#endif
+}
+
+
 } // namespace at::cuda
diff --git a/aten/src/ATen/cuda/CUDAGraph.cu b/aten/src/ATen/cuda/CUDAGraph.cu
new file mode 100644
index 0000000000000..1fe6a1e98e471
--- /dev/null
+++ b/aten/src/ATen/cuda/CUDAGraph.cu
@@ -0,0 +1,30 @@
+#include <ATen/cuda/CUDAGraph.h>
+#include <ATen/cuda/Exceptions.h>
+
+namespace at::cuda {
+
+namespace {
+
+#if !(defined(USE_ROCM)) && (defined(CUDA_VERSION) && CUDA_VERSION >= 12040)
+__global__ void set_conditional_handle_kernel(
+    cudaGraphConditionalHandle handle,
+    const bool* value) {
+  cudaGraphSetConditional(handle, *value);
+}
+#endif
+}
+
+void CUDAGraph::set_conditional_handle(
+    cudaGraphConditionalHandle handle,
+    const Tensor& scalar_cuda_pred_tensor) {
+#if !(defined(USE_ROCM)) && (defined(CUDA_VERSION) && CUDA_VERSION >= 12040)
+  set_conditional_handle_kernel<<<1, 1, 0, getCurrentCUDAStream()>>>(
+      handle, scalar_cuda_pred_tensor.const_data_ptr<bool>());
+  C10_CUDA_KERNEL_LAUNCH_CHECK();
+#else
+  AT_ERROR("not allowed");
+  return;
+#endif
+}
+
+} // namespace at::cuda
diff --git a/aten/src/ATen/cuda/CUDAGraph.h b/aten/src/ATen/cuda/CUDAGraph.h
index a32e7b4b86f07..4f14f6e50986c 100644
--- a/aten/src/ATen/cuda/CUDAGraph.h
+++ b/aten/src/ATen/cuda/CUDAGraph.h
@@ -4,9 +4,21 @@
 #include <c10/core/Device.h>
 #include <c10/cuda/CUDACachingAllocator.h>
 #include <c10/cuda/CUDAGraphsC10Utils.h>
+#include <c10/cuda/CUDAGuard.h>
 #include <c10/cuda/CUDAStream.h>
 #include <c10/util/flat_hash_map.h>
 
+#include <limits>
+#include <optional>
+#include <stack>
+
+#if defined(USE_ROCM) || !(defined(CUDA_VERSION) && CUDA_VERSION >= 12040)
+// this type is not defined until CUDA 12.4, but we use it as a
+// parameter type and return type in some below functions, so we give
+// it the same definition as in CUDA 12.4.
+typedef unsigned long long cudaGraphConditionalHandle;
+#endif // defined(USE_ROCM) || !(defined(CUDA_VERSION) && CUDA_VERSION >= 12040)
+
 namespace at {
 
 struct Generator;
@@ -19,11 +31,40 @@ namespace cuda {
 // to CUDAGraph::capture_begin
 TORCH_CUDA_CPP_API MempoolId_t graph_pool_handle();
 
+// Returns true if any CUDAGraph capture is currently active in this process.
+// Used by ProcessGroupNCCL's ROCm watchdog workaround to avoid calling
+// hipEventQuery during active capture on HIP runtimes without the
+// event-query capture-mode fix (https://github.com/ROCm/clr/pull/3176).
+// Not needed on CUDA/NVIDIA where cross-thread event query does not have this
+// restriction.
+#if defined(USE_ROCM)
+TORCH_CUDA_CPP_API bool is_graph_capture_active();
+#endif // defined(USE_ROCM)
+
 struct TORCH_CUDA_CPP_API CUDAGraph {
   CUDAGraph(bool keep_graph=false);
   ~CUDAGraph();
 
-  // See Note [Explicit Registration of Generators to the CUDA Graph]
+  // Copy and move constructors and assignments are disabled. These
+  // were disabled because pybind11 believed that CUDAGraph was copy
+  // constructable because
+  // pybind11::is_copy_constructible<CUDAGraph>::value originally
+  // evaluated to true. However, it cannot generate a copy constructor
+  // because CUDAGeneratorState, one of CUDAGraph's members, is an
+  // incomplete type unless CUDAGeneratorImpl.h is included. However,
+  // that would create a circular dependency between
+  // CUDAGeneratorImpl.h and CUDAGraph.h. Disabling the copy and move
+  // constructors is the most straightforward way to prevent pybind11
+  // from trying to generate default implementations of them.
+  //
+  // We needed pybind11 to return a reference to a CUDAGraph as part
+  // of wrapping CUDAGraph::get_currently_capturing_graph, which
+  // unearthed the above problem.
+  CUDAGraph(const CUDAGraph&) = delete;
+  CUDAGraph& operator=(const CUDAGraph&) = delete;
+  CUDAGraph(CUDAGraph&& other) = delete;
+  CUDAGraph& operator=(CUDAGraph&& other) = delete;
+
   void register_generator_state(c10::intrusive_ptr<at::CUDAGeneratorState> state);
   void register_generator_state(const at::Generator& generator);
   void capture_begin(
@@ -39,6 +80,18 @@ struct TORCH_CUDA_CPP_API CUDAGraph {
   cudaGraph_t raw_cuda_graph();
   cudaGraphExec_t raw_cuda_graph_exec();
 
+  static CUDAGraph* get_currently_capturing_graph();
+  void begin_capture_to_if_node(const Tensor& scalar_cuda_pred_tensor);
+  void end_capture_to_conditional_node();
+  static void set_conditional_handle(
+      cudaGraphConditionalHandle handle,
+      const Tensor& scalar_cuda_pred_tensor);
+
+ private:
+  template <typename StreamType>
+  std::function<bool(StreamType)> create_allocate_filter() const;
+  std::function<bool(cudaStream_t)> create_child_allocate_filter();
+
  protected:
   cudaGraph_t graph_ = nullptr;
   cudaGraphExec_t graph_exec_ = nullptr;
@@ -54,6 +107,18 @@ struct TORCH_CUDA_CPP_API CUDAGraph {
   // Set to true in capture_end if cudaGraphInstantiate succeeded
   bool has_graph_exec_ = false;
 
+  // Set to true in capture_begin once a private pool has been acquired
+  // (beginAllocateToPool). Tells reset() it must release the pool, even if the
+  // capture failed before capture_end() completed. Otherwise a failed capture
+  // leaks the pool: its use_count never returns to zero, so empty_cache can
+  // never reclaim its segments for the rest of the process.
+  bool allocated_pool_ = false;
+  // Set to true in capture_begin after beginAllocateToPool and cleared in
+  // capture_end after endAllocateToPool. Tells reset() whether the allocator is
+  // still routing allocations to the pool (capture abandoned before capture_end
+  // ran) and must be ended before the pool can be released.
+  bool capturing_to_pool_ = false;
+
   // the ID assigned by cuda during graph capture,
   // used to identify when a stream is participating in capture
   CaptureId_t capture_id_ = 0;
@@ -89,7 +154,18 @@ struct TORCH_CUDA_CPP_API CUDAGraph {
   c10::DeviceIndex capture_dev_{UNDEFINED_DEVICE};
 
   bool keep_graph_;
+  cudaStreamCaptureMode capture_mode_{};
+
+#if !defined(USE_ROCM) && (defined(CUDA_VERSION) && CUDA_VERSION >= 12040)
+  std::stack<at::cuda::CUDAStreamGuard> conditional_node_streams_;
+  std::stack<CaptureId_t> conditional_graph_capture_ids_;
+#endif // !defined(USE_ROCM) && defined(CUDA_VERSION) && CUDA_VERSION >= 12040
 };
 
+template <>
+std::function<bool(cudaStream_t)> CUDAGraph::create_allocate_filter<cudaStream_t>() const;
+template <>
+std::function<bool(c10::Stream)> CUDAGraph::create_allocate_filter<c10::Stream>() const;
+
 } // namespace cuda
 } // namespace at
diff --git a/aten/src/ATen/cuda/CUDAGraphsUtils.cuh b/aten/src/ATen/cuda/CUDAGraphsUtils.cuh
index d3a5b306eeea4..64ab0140793bb 100644
--- a/aten/src/ATen/cuda/CUDAGraphsUtils.cuh
+++ b/aten/src/ATen/cuda/CUDAGraphsUtils.cuh
@@ -19,12 +19,17 @@ using CaptureStatus = c10::cuda::CaptureStatus;
 
 // Use this version where you don't want to create a CUDA context if none exists.
 inline CaptureStatus currentStreamCaptureStatus() {
-  // don't create a context if we don't have to
   if (c10::cuda::hasPrimaryContext(c10::cuda::current_device())) {
     return c10::cuda::currentStreamCaptureStatusMayInitCtx();
-  } else {
-    return CaptureStatus::None;
   }
+  return CaptureStatus::None;
+}
+
+inline std::optional<CaptureId_t> currentStreamCaptureId() {
+  if (c10::cuda::hasPrimaryContext(c10::cuda::current_device())) {
+    return c10::cuda::currentStreamCaptureIdMayInitCtx();
+  }
+  return std::nullopt;
 }
 
 inline void assertNotCapturing(const std::string& attempt) {
diff --git a/aten/src/ATen/cuda/CUDAGreenContext.cpp b/aten/src/ATen/cuda/CUDAGreenContext.cpp
index 7a0eb81bfe3c2..07861d7d78a11 100644
--- a/aten/src/ATen/cuda/CUDAGreenContext.cpp
+++ b/aten/src/ATen/cuda/CUDAGreenContext.cpp
@@ -11,10 +11,27 @@
 C10_DIAGNOSTIC_PUSH_AND_IGNORED_IF_DEFINED("-Wunused-private-field")
 #endif
 
+#if defined(CUDA_VERSION) && CUDA_VERSION >= 13010 && HAS_CUDA_GREEN_CONTEXT()
+#define HAS_CUDA_WORKQUEUE_SUPPORT() 1
+#else
+#define HAS_CUDA_WORKQUEUE_SUPPORT() 0
+#endif
+
 namespace at::cuda {
 
-GreenContext::GreenContext(uint32_t device_id, uint32_t num_sms) {
+GreenContext::GreenContext(
+    uint32_t device_id,
+    std::optional<uint32_t> num_sms,
+    std::optional<int32_t> workqueue_scope,
+    std::optional<uint32_t> workqueue_concurrency_limit) {
 #if HAS_CUDA_GREEN_CONTEXT()
+  TORCH_CHECK(
+      num_sms.has_value() || workqueue_scope.has_value(),
+      "At least one of num_sms or workqueue_scope must be specified");
+  TORCH_CHECK(
+      !workqueue_concurrency_limit.has_value() || workqueue_scope.has_value(),
+      "workqueue_concurrency_limit requires workqueue_scope to be set");
+
   int driver_version;
   C10_CUDA_CHECK(cudaDriverGetVersion(&driver_version));
   TORCH_CHECK(
@@ -29,46 +46,73 @@ GreenContext::GreenContext(uint32_t device_id, uint32_t num_sms) {
     cudaFree(nullptr);
   }
 
-   CUdevice device;
+  CUdevice device;
   device_id_ = device_id;
   C10_CUDA_DRIVER_CHECK(
       c10::cuda::DriverAPI::get()->cuDeviceGet_(&device, device_id));
 
-  // Get device resources
-  CUdevResource device_resource;
-  C10_CUDA_DRIVER_CHECK(c10::cuda::DriverAPI::get()->cuDeviceGetDevResource_(
-      device, &device_resource, CU_DEV_RESOURCE_TYPE_SM));
+  std::vector<CUdevResource> resources;
 
-  TORCH_CHECK(
-      num_sms > 0 && num_sms <= device_resource.sm.smCount,
-      "Invalid number of SMs requested for green context: ",
-      num_sms,
-      " (device has ",
-      device_resource.sm.smCount,
-      " SMs)");
-
-  // Split resources
-  std::vector<CUdevResource> result(1);
-  auto result_data = result.data();
-  unsigned int nb_groups = 1;
-  CUdevResource remaining;
+  // --- SM resource ---
+  if (num_sms.has_value()) {
+    CUdevResource sm_resource;
+    C10_CUDA_DRIVER_CHECK(c10::cuda::DriverAPI::get()->cuDeviceGetDevResource_(
+        device, &sm_resource, CU_DEV_RESOURCE_TYPE_SM));
 
-  C10_CUDA_DRIVER_CHECK(
-      c10::cuda::DriverAPI::get()->cuDevSmResourceSplitByCount_(
-          result_data,
-          &nb_groups,
-          &device_resource,
-          &remaining,
-          0, // default flags
-          num_sms));
+    TORCH_CHECK(
+        *num_sms > 0 && *num_sms <= sm_resource.sm.smCount,
+        "Invalid number of SMs requested for green context: ",
+        *num_sms,
+        " (device has ",
+        sm_resource.sm.smCount,
+        " SMs)");
+
+    // Split resources
+    std::vector<CUdevResource> split_result(1);
+    unsigned int nb_groups = 1;
+    CUdevResource remaining;
+
+    C10_CUDA_DRIVER_CHECK(
+        c10::cuda::DriverAPI::get()->cuDevSmResourceSplitByCount_(
+            split_result.data(),
+            &nb_groups,
+            &sm_resource,
+            &remaining,
+            0, // default flags
+            *num_sms));
+    TORCH_CHECK(nb_groups == 1, "Failed to create single SM resource group");
+    resources.push_back(split_result[0]);
+  }
 
-  TORCH_CHECK(nb_groups == 1, "Failed to create single resource group");
+  // --- Workqueue config resource ---
+  if (workqueue_scope.has_value()) {
+#if HAS_CUDA_WORKQUEUE_SUPPORT()
+    TORCH_CHECK(
+        driver_version >= 13010, "cuda driver too old to use workqueue configuration!");
+    CUdevResource wq_resource{};
+    C10_CUDA_DRIVER_CHECK(c10::cuda::DriverAPI::get()->cuDeviceGetDevResource_(
+        device, &wq_resource, CU_DEV_RESOURCE_TYPE_WORKQUEUE_CONFIG));
+
+    wq_resource.wqConfig.sharingScope =
+        static_cast<CUdevWorkqueueConfigScope>(*workqueue_scope);
+    if (workqueue_concurrency_limit.has_value()) {
+      wq_resource.wqConfig.wqConcurrencyLimit = *workqueue_concurrency_limit;
+    }
+    resources.push_back(wq_resource);
+#else
+    TORCH_CHECK(
+        false,
+        "Workqueue configuration for green contexts requires CUDA 13.1+!");
+#endif
+  }
 
   // Generate resource descriptor
   CUdevResourceDesc desc;
   C10_CUDA_DRIVER_CHECK(
       c10::cuda::DriverAPI::get()->cuDevResourceGenerateDesc_(
-          &desc, result_data, 1));
+          &desc,
+          resources.data(),
+          static_cast<unsigned int>(resources.size())));
 
   // Create green context
   // CU_GREEN_CTX_DEFAULT_STREAM is required per docs:
@@ -83,20 +127,45 @@ GreenContext::GreenContext(uint32_t device_id, uint32_t num_sms) {
 #else
   TORCH_CHECK(false, "Green Context is only supported on CUDA 12.8+!");
 #endif
-  }
+}
 
-  std::unique_ptr<GreenContext> GreenContext::create(
-      uint32_t num_sms,
-      std::optional<uint32_t> device_id) {
+std::unique_ptr<GreenContext> GreenContext::create(
+    std::optional<uint32_t> device_id,
+    std::optional<uint32_t> num_sms,
+    std::optional<int32_t> workqueue_scope,
+    std::optional<uint32_t> workqueue_concurrency_limit) {
 #if HAS_CUDA_GREEN_CONTEXT()
-    if (!device_id.has_value()) {
-      device_id = at::cuda::current_device();
-    }
-    return std::unique_ptr<GreenContext>(new GreenContext(device_id.value(), num_sms));
+  if (!device_id.has_value()) {
+    device_id = at::cuda::current_device();
+  }
+  return std::unique_ptr<GreenContext>(new GreenContext(
+      device_id.value(), num_sms, workqueue_scope, workqueue_concurrency_limit));
 #else
-    TORCH_CHECK(false, "Green Context is only supported on CUDA 12.8+!");
+  TORCH_CHECK(false, "Green Context is only supported on CUDA 12.8+!");
 #endif
+}
+
+uint32_t GreenContext::max_workqueue_concurrency(
+    std::optional<uint32_t> device_id) {
+#if HAS_CUDA_WORKQUEUE_SUPPORT()
+  int driver_version;
+  C10_CUDA_CHECK(cudaDriverGetVersion(&driver_version));
+  TORCH_CHECK(
+      driver_version >= 13010, "cuda driver too old to use workqueue configuration!");
+  if (!device_id.has_value()) {
+    device_id = at::cuda::current_device();
   }
+  CUdevice device;
+  C10_CUDA_DRIVER_CHECK(
+      c10::cuda::DriverAPI::get()->cuDeviceGet_(&device, device_id.value()));
+  CUdevResource wq_resource;
+  C10_CUDA_DRIVER_CHECK(c10::cuda::DriverAPI::get()->cuDeviceGetDevResource_(
+      device, &wq_resource, CU_DEV_RESOURCE_TYPE_WORKQUEUE_CONFIG));
+  return wq_resource.wqConfig.wqConcurrencyLimit;
+#else
+  TORCH_CHECK(false, "Workqueue configuration requires CUDA 13.1+!");
+#endif
+}
 
   // Implement move operations
 #if HAS_CUDA_GREEN_CONTEXT()
diff --git a/aten/src/ATen/cuda/CUDAGreenContext.h b/aten/src/ATen/cuda/CUDAGreenContext.h
index b1eda3f48a45a..a93eedfaeedfa 100644
--- a/aten/src/ATen/cuda/CUDAGreenContext.h
+++ b/aten/src/ATen/cuda/CUDAGreenContext.h
@@ -11,12 +11,24 @@ namespace {
   constexpr int kStreamPerGreenContextPool = 32;
 }
 
+// Workqueue sharing scope for green contexts.
+// Values match the CUDA driver API's CUdevWorkqueueConfigScope enum.
+enum class WorkqueueScope : int32_t {
+  DeviceCtx = 0,
+  Balanced = 1,
+};
+
 class TORCH_CUDA_CPP_API GreenContext {
  public:
-  // Green context creation
   static std::unique_ptr<GreenContext> create(
-      uint32_t num_sms,
-      std::optional<uint32_t> device_id);
+    std::optional<uint32_t> device_id,
+    std::optional<uint32_t> num_sms,
+    std::optional<int32_t> workqueue_scope = std::nullopt,
+    std::optional<uint32_t> workqueue_concurrency_limit = std::nullopt);
+
+  static uint32_t max_workqueue_concurrency(
+      std::optional<uint32_t> device_id = std::nullopt);
+
   ~GreenContext() noexcept;
 
   // Delete copy constructor and assignment
@@ -31,7 +43,12 @@ class TORCH_CUDA_CPP_API GreenContext {
   CUDAStream Stream();
 
  private:
-  GreenContext(uint32_t device_id, uint32_t num_sms);
+  GreenContext(
+    uint32_t device_id,
+    std::optional<uint32_t> num_sms,
+    std::optional<int32_t> workqueue_scope,
+    std::optional<uint32_t> workqueue_concurrency_limit);
+
   // Implement move operations
   GreenContext(GreenContext&& other) noexcept;
   GreenContext& operator=(GreenContext&& other) noexcept;
diff --git a/aten/src/ATen/cuda/CUDAScaledBlas.h b/aten/src/ATen/cuda/CUDAScaledBlas.h
deleted file mode 100644
index 3fd0e2a6a3aae..0000000000000
--- a/aten/src/ATen/cuda/CUDAScaledBlas.h
+++ /dev/null
@@ -1,173 +0,0 @@
-#include <cstdint>
-#include <c10/util/typeid.h>
-#include <c10/util/Exception.h>
-#include <c10/util/SmallVector.h>
-#include <c10/core/Scalar.h>
-#include <c10/core/ScalarType.h>
-#define TORCH_ASSERT_ONLY_METHOD_OPERATORS
-#include <ATen/core/Tensor.h>
-#include <ATen/core/NamedTensor.h>
-#include <ATen/Dispatch.h>
-#include <ATen/ExpandUtils.h>
-#include <ATen/OpMathType.h>
-#include <ATen/TensorUtils.h>
-#include <ATen/cuda/CUDABlas.h>
-#include <ATen/cuda/tunable/Tunable.h>
-#include <ATen/cuda/tunable/TunableGemm.h>
-#include <ATen/native/Resize.h>
-#include <c10/util/MaybeOwned.h>
-#include <ATen/native/GroupedMMUtils.h>
-#include <ATen/native/cuda/RowwiseScaledMM.h>
-#include <ATen/native/cuda/ScaledGroupMM.h>
-#include <ATen/native/cuda/GroupMM.h>
-#include <ATen/ceil_div.h>
-
-#ifdef USE_MSLK
-#include <mslk/gemm/gemm_torch.h>
-#endif
-
-#ifndef AT_PER_OPERATOR_HEADERS
-#include <ATen/Functions.h>
-#include <ATen/NativeFunctions.h>
-#else
-#include <ATen/ops/_addmm_activation_native.h>
-#include <ATen/ops/_efficientzerotensor.h>
-#include <ATen/ops/_scaled_mm_native.h>
-#include <ATen/ops/_unsafe_view_native.h>
-#include <ATen/ops/abs.h>
-#include <ATen/ops/addmm_native.h>
-#include <ATen/ops/addmv_native.h>
-#include <ATen/ops/baddbmm_native.h>
-#include <ATen/ops/bmm_native.h>
-#include <ATen/ops/copy_native.h>
-#include <ATen/ops/dot_native.h>
-#include <ATen/ops/empty.h>
-#include <ATen/ops/empty_strided.h>
-#include <ATen/ops/gelu.h>
-#include <ATen/ops/max.h>
-#include <ATen/ops/mm_native.h>
-#include <ATen/ops/mul.h>
-#include <ATen/ops/relu.h>
-#include <ATen/ops/ones.h>
-#include <ATen/ops/scalar_tensor_native.h>
-#include <ATen/ops/vdot_native.h>
-#endif
-
-using at::blas::ScalingType;
-using at::blas::SwizzleType;
-
-namespace at::cuda::scaled {
-
-static bool _scaled_mm_allowed_device(bool sm90_only=false, bool sm100_only=false) {
-#ifdef USE_ROCM
-    static const std::vector<std::string> archs = {
-        "gfx942",
-#if ROCM_VERSION >= 60300
-        "gfx1200", "gfx1201",
-#endif
-#if ROCM_VERSION >= 60500
-        "gfx950"
-#endif
-    };
-    return at::detail::getCUDAHooks().isGPUArch(archs);
-#else
-    auto dprops = at::cuda::getCurrentDeviceProperties();
-
-    if (sm90_only || sm100_only) {
-      return (sm90_only && dprops->major == 9) || (sm100_only && dprops->major == 10);
-    } else {
-      return dprops->major >= 9 || (dprops->major == 8 && dprops->minor == 9);
-    }
-#endif
-}
-
-#ifdef USE_ROCM
-static bool _scaled_mm_is_fnuz() {
-    return at::detail::getCUDAHooks().isGPUArch({"gfx942"});
-}
-#endif
-/**
- * Track concrete implementations available
- */
-enum class ScaledGemmImplementation {
-  NONE = 0,
-  TENSORWISE_TENSORWISE = 1,
-  ROWWISE_ROWWISE = 2,
-  BLOCK_128x128_1x128 = 3,
-  BLOCK_1x128_128x128 = 4,
-  BLOCK_1x128_1x128 = 5,
-  MXFP8_MXFP8 = 6,
-  NVFP4_NVFP4 = 7,
-  NVFP4_NVFP4_SINGLE_SCALE = 8,
-  MXFP4_MXFP4 = 9,
-};
-
-/**
- * Convert passed int (enum) from python back into a
- * strictly-typed enum
- */
-template <class EnumType, class ArrayType>
-std::vector<EnumType> convert_int_to_enum(ArrayType& v) {
-  std::vector<EnumType> converted;
-  converted.reserve(v.size());
-
-  for (auto vi : v) {
-    converted.push_back(static_cast<EnumType>(vi));
-  }
-  return converted;
-}
-
-bool check_tensorwise_recipe(c10::ScalarType,
-                             std::vector<ScalingType>&,
-                             ArrayRef<Tensor>&,
-                             c10::ScalarType,
-                             std::vector<ScalingType>&,
-                             ArrayRef<Tensor>&);
-
-
-bool check_rowwise_recipe(c10::ScalarType,
-                             std::vector<ScalingType>&,
-                             ArrayRef<Tensor>&,
-                             c10::ScalarType,
-                             std::vector<ScalingType>&,
-                             ArrayRef<Tensor>&);
-
-bool check_nvfp4_recipe(c10::ScalarType,
-                        std::vector<ScalingType>&,
-                        ArrayRef<Tensor>&,
-                        c10::ScalarType,
-                        std::vector<ScalingType>&,
-                        ArrayRef<Tensor>&);
-
-bool check_nvfp4_recipe_single_scale
-                       (c10::ScalarType,
-                        std::vector<ScalingType>&,
-                        ArrayRef<Tensor>&,
-                        c10::ScalarType,
-                        std::vector<ScalingType>&,
-                        ArrayRef<Tensor>&);
-
-bool check_deepseek_recipe(ScalingType,
-                           ScalingType,
-                           c10::ScalarType,
-                           std::vector<ScalingType>&,
-                           ArrayRef<Tensor>&,
-                           c10::ScalarType,
-                           std::vector<ScalingType>&,
-                           ArrayRef<Tensor>&);
-
-bool check_mxfp8_recipe(c10::ScalarType,
-                        std::vector<ScalingType>&,
-                        ArrayRef<Tensor>&,
-                        c10::ScalarType,
-                        std::vector<ScalingType>&,
-                        ArrayRef<Tensor>&);
-
-bool check_mxfp4_recipe(c10::ScalarType,
-                        std::vector<ScalingType>&,
-                        ArrayRef<Tensor>&,
-                        c10::ScalarType,
-                        std::vector<ScalingType>&,
-                        ArrayRef<Tensor>&);
-
-} // namespace at::native::cuda::blas::scaled
diff --git a/aten/src/ATen/cuda/CachingHostAllocator.cpp b/aten/src/ATen/cuda/CachingHostAllocator.cpp
index dcff81333cb2f..c376eb3a0a4d5 100644
--- a/aten/src/ATen/cuda/CachingHostAllocator.cpp
+++ b/aten/src/ATen/cuda/CachingHostAllocator.cpp
@@ -3,6 +3,8 @@
 #include <ATen/cuda/CUDAEvent.h>
 #include <c10/core/thread_pool.h>
 #include <c10/cuda/CUDAAllocatorConfig.h>
+#include <c10/cuda/CUDAGraphsC10Utils.h>
+#include <c10/util/Gauge.h>
 
 #include <cuda_runtime_api.h>
 #include <future>
@@ -14,6 +16,25 @@ using Block = HostBlock<CUDAStream>;
 
 struct CUDACachingHostAllocatorImpl
     : public CachingHostAllocatorImpl<CUDAStream, CUDAEventPool::Event> {
+  void free(void* ctx) override {
+    using Base = CachingHostAllocatorImpl<CUDAStream, CUDAEventPool::Event>;
+    try {
+      Base::free(ctx);
+    } catch (...) {
+      if (!c10::cuda::CUDACachingAllocator::CUDAAllocatorConfig::
+              pinned_free_catch_all()) {
+        TORCH_WARN("Exception in pinned allocator free(), rethrowing");
+        throw;
+      }
+      // pinned_free_catch_all is enabled: suppress the exception to prevent
+      // it from escaping through ~StorageImpl() (implicitly noexcept), which
+      // would cause std::terminate. Allows graceful shutdown to proceed.
+      STATIC_GAUGE(pytorch.CUDACachingHostAllocator.free_fail_catch_all)
+          .record(1);
+      TORCH_WARN("Suppressed exception in pinned allocator free()");
+    }
+  }
+
  private:
   ska::flat_hash_map<void*, bool> use_host_register;
 
@@ -245,9 +266,7 @@ struct CUDACachingHostAllocatorImpl
   }
 
   bool stream_is_capturing(CUDAStream s) const override {
-    cudaStreamCaptureStatus status{cudaStreamCaptureStatusNone};
-    C10_CUDA_CHECK(cudaStreamIsCapturing(s, &status));
-    return status != cudaStreamCaptureStatusNone;
+    return s.is_capturing();
   }
 };
 
diff --git a/aten/src/ATen/cuda/CublasHandlePool.cpp b/aten/src/ATen/cuda/CublasHandlePool.cpp
index b5008bad832d0..6a6d0fe1929cf 100644
--- a/aten/src/ATen/cuda/CublasHandlePool.cpp
+++ b/aten/src/ATen/cuda/CublasHandlePool.cpp
@@ -7,9 +7,14 @@
 #include <map>
 #include <memory>
 #include <regex>
+#include <shared_mutex>
 #include <string>
 #include <tuple>
 
+#if defined(USE_ROCM)
+#include <rocblas/rocblas.h>
+#endif
+
 /**
  * Note [hipblaslt handles]
  * ~~~~~~~~~~~~~~~~~~~~~~~~
@@ -24,6 +29,10 @@
  * For CUDA builds, getCurrentCUDABlasLtHandle will alias for getCurrentCUDABlasHandle,
  * whereas for ROCm builds, it is a distinct function.
  *
+ * Additionally, hipblaslt cannot share a single handle across multiple streams.
+ * On ROCm, getCurrentCUDABlasLtHandle returns a handle unique to each (device, stream)
+ * pair, rather than just per-device like the cublas handle pool.
+ *
  * The workspace pools are separate for ROCm. On CUDA, the env var
  * TORCH_CUBLASLT_UNIFIED_WORKSPACE can be used to opt-in to unifying the workspace pools.
  */
@@ -53,8 +62,6 @@ void destroyCublasLtHandle(cublasLtHandle_t handle) {
 using CuBlasLtPoolType = DeviceThreadHandlePool<cublasLtHandle_t, createCublasLtHandle, destroyCublasLtHandle>;
 
 // ugly hack until hipblasSetWorkspace exists
-#include <rocblas/rocblas.h>
-
 static hipblasStatus_t rocBLASStatusToHIPStatus(rocblas_status error) {
     switch(error) {
     case rocblas_status_size_unchanged:
@@ -130,6 +137,24 @@ void clearCublasWorkspaces() {
   }
 }
 
+void clearCublasWorkspacesForStream(cudaStream_t stream) {
+  void* stream_ptr = static_cast<void*>(stream);
+  {
+    auto& workspace = cublas_handle_stream_to_workspace();
+    std::unique_lock<std::shared_mutex> lock(workspace.mutex);
+    std::erase_if(workspace.map, [stream_ptr](const auto& entry) {
+      return std::get<1>(entry.first) == stream_ptr;
+    });
+  }
+  {
+    auto& workspace = cublaslt_handle_stream_to_workspace();
+    std::unique_lock<std::shared_mutex> lock(workspace.mutex);
+    std::erase_if(workspace.map, [stream_ptr](const auto& entry) {
+      return std::get<1>(entry.first) == stream_ptr;
+    });
+  }
+}
+
 size_t parseChosenWorkspaceSize() {
   auto val = c10::utils::get_env("CUBLAS_WORKSPACE_CONFIG");
 #ifdef USE_ROCM
@@ -140,14 +165,14 @@ size_t parseChosenWorkspaceSize() {
     // for extra convenience
     val = c10::utils::get_env("ROCBLAS_WORKSPACE_CONFIG");
   }
-  /* 32MiB default, 128MiB for gfx94x/gfx95x */
-  const bool gfx94_95 = at::detail::getCUDAHooks().isGPUArch({"gfx94", "gfx95"});
-  const size_t default_size = gfx94_95 ? 1024 * 128 * 1024 : 1024 * 32 * 1024;
+  /* 32MiB default, 128MiB for gfx942/gfx950/gfx1250 */
+  const bool gfx942_950_1250 = at::detail::getCUDAHooks().isGPUArch({"gfx942", "gfx950", "gfx1250"});
+  const size_t default_size = gfx942_950_1250 ? 1024 * 128 * 1024 : 1024 * 32 * 1024;
 #else
-  /* :4096:2:16:8 default, 32MiB for Hopper */
+  /* :4096:2:16:8 default, 32MiB for Hopper and Blackwell */
   cudaDeviceProp* properties = at::cuda::getCurrentDeviceProperties();
-  const bool sm90 = properties != nullptr && properties->major == 9 && properties->minor == 0;
-  const size_t default_size = sm90 ? 4096 * 8 * 1024 : 4096 * 1024 * 2 + 16 * 1024 * 8;
+  const bool use32mb = properties != nullptr && (properties->major == 9 || properties->major == 10 || properties->major == 12);
+  const size_t default_size = use32mb ? 4096 * 8 * 1024 : 4096 * 1024 * 2 + 16 * 1024 * 8;
 #endif
 
   if (val) {
@@ -174,6 +199,19 @@ size_t parseChosenWorkspaceSize() {
   }
 }
 
+#define TORCH_CUBLASLT_UNIFIED_WORKSPACE "TORCH_CUBLASLT_UNIFIED_WORKSPACE"
+#ifndef USE_ROCM
+inline bool unified_cublas_and_lt_workspaces() {
+  static auto unified_env_var = c10::utils::check_env(TORCH_CUBLASLT_UNIFIED_WORKSPACE);
+#if !defined(FBCODE)
+  static bool unified = (unified_env_var == std::nullopt) || (unified_env_var == true);
+#else
+  static bool unified = unified_env_var == true;
+#endif
+  return unified;
+}
+#endif
+
 size_t parseCUDABlasLtWorkspaceSize() {
   auto val = c10::utils::get_env("CUBLASLT_WORKSPACE_SIZE");
 #ifdef USE_ROCM
@@ -183,7 +221,9 @@ size_t parseCUDABlasLtWorkspaceSize() {
   }
   size_t workspace_size = 76*1024; /* Use 76 MB for hipBLASLt */
 #else
-  size_t workspace_size = 1024; /* default size in KiB according to #73328 */
+  /* use CUDABlas default workspace size if unified */
+  /* otherwise, use default size in KiB according to #73328 */
+  size_t workspace_size = unified_cublas_and_lt_workspaces() ? parseChosenWorkspaceSize() / 1024 : 1024;
 #endif
 
   if (val.has_value()) {
@@ -206,24 +246,11 @@ size_t parseCUDABlasLtWorkspaceSize() {
   return workspace_size * 1024;
 }
 
-size_t getChosenWorkspaceSize() {
-  size_t pool_size = parseChosenWorkspaceSize();
-  return pool_size;
-}
-
-#define TORCH_CUBLASLT_UNIFIED_WORKSPACE "TORCH_CUBLASLT_UNIFIED_WORKSPACE"
-
 size_t getCUDABlasLtWorkspaceSize() {
-  size_t pool_size = parseCUDABlasLtWorkspaceSize();
+  static size_t pool_size = parseCUDABlasLtWorkspaceSize();
 #ifndef USE_ROCM
-  static auto unified_env_var = c10::utils::check_env(TORCH_CUBLASLT_UNIFIED_WORKSPACE);
-#if !defined(FBCODE)
-  static bool unified = (unified_env_var == std::nullopt) || (unified_env_var == true);
-#else
-  static bool unified = unified_env_var == true;
-#endif
-  if (unified) {
-    auto cublasWorkspaceSize = getChosenWorkspaceSize();
+  if (unified_cublas_and_lt_workspaces()) {
+    static size_t cublasWorkspaceSize = parseChosenWorkspaceSize();
     if (cublasWorkspaceSize < pool_size) {
       TORCH_WARN_ONCE("Requested unified CUBLASLT workspace size of ", pool_size,
                       " bytes exceeds CUBLAS workspace size of ", cublasWorkspaceSize,
@@ -239,7 +266,7 @@ size_t getCUDABlasLtWorkspaceSize() {
 }
 
 at::DataPtr getNewWorkspace() {
-  return c10::cuda::CUDACachingAllocator::get()->allocate(getChosenWorkspaceSize());
+  return c10::cuda::CUDACachingAllocator::get()->allocate(parseChosenWorkspaceSize());
 }
 
 at::DataPtr getNewCUDABlasLtWorkspace() {
@@ -252,7 +279,7 @@ void setWorkspaceForHandle(cublasHandle_t handle, c10::cuda::CUDAStream stream)
 
   auto& workspace = cublas_handle_stream_to_workspace();
 
-  size_t workspace_size = getChosenWorkspaceSize();
+  size_t workspace_size = parseChosenWorkspaceSize();
 
   // Fast path: check if workspace already exists
   {
@@ -280,17 +307,27 @@ void setWorkspaceForHandle(cublasHandle_t handle, c10::cuda::CUDAStream stream)
 
 void* getCUDABlasLtWorkspace() {
 #ifndef USE_ROCM
-  static bool unified = c10::utils::check_env(TORCH_CUBLASLT_UNIFIED_WORKSPACE) == true;
-  if (unified) {
-    cublasHandle_t handle = at::cuda::getCurrentCUDABlasHandle();
+  if (unified_cublas_and_lt_workspaces()) {
+    cublasHandle_t handle = at::cuda::getCurrentCUDABlasHandle(/*setup=*/false);
     auto stream = c10::cuda::getCurrentCUDAStream();
     cudaStream_t _stream = stream;
     auto key = std::make_tuple(static_cast<void *>(handle), static_cast<void *>(_stream));
     auto& workspace = at::cuda::cublas_handle_stream_to_workspace();
-    std::shared_lock<std::shared_mutex> lock(workspace.mutex);
-    auto workspace_it = workspace.map.find(key);
-    TORCH_INTERNAL_ASSERT(workspace_it != workspace.map.end());
-    return workspace_it->second.mutable_get();
+    {
+      std::shared_lock<std::shared_mutex> lock(workspace.mutex);
+      auto workspace_it = workspace.map.find(key);
+      if (workspace_it != workspace.map.end()) {
+        return workspace_it->second.mutable_get();
+      }
+    }
+    // First use for this handle+stream pair — allocate and insert directly.
+    // No need to call cublasSetWorkspace; Lt passes workspace explicitly.
+    auto new_workspace = getNewWorkspace();
+    {
+      std::unique_lock<std::shared_mutex> lock(workspace.mutex);
+      auto workspace_it = workspace.map.try_emplace(key, std::move(new_workspace)).first;
+      return workspace_it->second.mutable_get();
+    }
   }
 #endif
   cublasLtHandle_t handle = getCurrentCUDABlasLtHandle();
@@ -322,7 +359,7 @@ void* getCUDABlasLtWorkspace() {
   }
 }
 
-cublasHandle_t getCurrentCUDABlasHandle() {
+cublasHandle_t getCurrentCUDABlasHandle(bool setup) {
   c10::DeviceIndex device = 0;
   AT_CUDA_CHECK(c10::cuda::GetDevice(&device));
 
@@ -330,8 +367,6 @@ cublasHandle_t getCurrentCUDABlasHandle() {
   CUcontext pctx = nullptr;
   at::globalContext().getNVRTC().cuCtxGetCurrent(&pctx);
   if (C10_UNLIKELY(!pctx)) {
-    // workaround for corner case where a primary context exists but is not
-    // the current context, seen in multithreaded use-cases
     TORCH_WARN_ONCE("Attempting to run cuBLAS, but there was no current CUDA context! Attempting to set the primary context...");
     at::globalContext().getNVRTC().cuDevicePrimaryCtxRetain(&pctx, device);
     at::globalContext().getNVRTC().cuCtxSetCurrent(pctx);
@@ -354,6 +389,11 @@ cublasHandle_t getCurrentCUDABlasHandle() {
       pool->newPoolWindow());
 
   auto handle = myPoolWindow->reserve(device);
+
+  if (!setup) {
+    return handle;
+  }
+
   auto stream = c10::cuda::getCurrentCUDAStream();
   TORCH_CUDABLAS_CHECK(cublasSetStream(handle, stream));
   // We explicitly set the cublas workspace even though CUDA 12.2+ fixed the
@@ -407,10 +447,14 @@ cublasLtHandle_t getCurrentCUDABlasLtHandle() {
   thread_local std::unique_ptr<CuBlasLtPoolType::PoolWindow> myPoolWindow(
       pool->newPoolWindow());
 
-  auto handle = myPoolWindow->reserve(device);
+  // hipblaslt cannot share a single handle across multiple streams,
+  // so reserve a handle unique to each (device, stream) pair.
+  auto stream = c10::cuda::getCurrentCUDAStream();
+  cudaStream_t _stream = stream;
+  auto handle = myPoolWindow->reserve(device, static_cast<void*>(_stream));
   return handle;
 #else
-  return reinterpret_cast<cublasLtHandle_t>(getCurrentCUDABlasHandle());
+  return reinterpret_cast<cublasLtHandle_t>(getCurrentCUDABlasHandle(/*setup=*/false));
 #endif
 }
 
diff --git a/aten/src/ATen/cuda/StatelessPhilox4x32.cuh b/aten/src/ATen/cuda/StatelessPhilox4x32.cuh
new file mode 100644
index 0000000000000..7d21af6b33480
--- /dev/null
+++ b/aten/src/ATen/cuda/StatelessPhilox4x32.cuh
@@ -0,0 +1,61 @@
+// Stateless Philox-4x32 PRNG implementation.
+//
+// Unlike PhiloxRNGEngine (PhiloxUtils.cuh), this is a pure function: given
+// (seed, offset) it returns 4 pseudo-random uint32 values with no mutable
+// state. This makes it suitable for use in stateless random APIs.
+//
+// The Philox-4x32 cipher operates on a 128-bit counter. The full counter
+// is (offset_lo, offset_hi, subsequence_lo, subsequence_hi), but we fix
+// subsequence=0 so that the entire 128-bit counter space is addressed by
+// the 64-bit offset alone. This keeps the API simple and maintains
+// cross-device consistency. For example, utilizing thread ID-based subsequence
+// numbers and SM-based thread count causes different random values to
+// be generated across GPU types. We avoid this situation by always setting
+// subsequence=0.
+
+#pragma once
+
+#include <cstdint>
+
+namespace at::cuda {
+
+__device__ __forceinline__ uint2 mulhilo32(uint32_t a, uint32_t b) {
+  return {a * b, __umulhi(a, b)};
+}
+
+__device__ __forceinline__ uint4 philox_round(uint4 ctr, uint2 key) {
+  constexpr uint32_t kPhiloxSA = 0xD2511F53;
+  constexpr uint32_t kPhiloxSB = 0xCD9E8D57;
+  uint2 r0 = mulhilo32(kPhiloxSA, ctr.x);
+  uint2 r1 = mulhilo32(kPhiloxSB, ctr.z);
+  return {r1.y ^ ctr.y ^ key.x, r1.x, r0.y ^ ctr.w ^ key.y, r0.x};
+}
+
+// Stateless Philox-4x32. Returns 4 pseudo-random uint32 values (128 bits)
+// determined entirely by (seed, offset). Each unique offset produces a
+// distinct 128-bit output.
+template <int N_ROUNDS = 10>
+__device__ __forceinline__ uint4 philox_4x32(
+    uint64_t seed, uint64_t offset) {
+  uint2 key = {
+      static_cast<uint32_t>(seed),
+      static_cast<uint32_t>(seed >> 32)};
+  uint4 ctr = {
+      static_cast<uint32_t>(offset),
+      static_cast<uint32_t>(offset >> 32),
+      // restrict subsequence=0
+      0, 0};
+
+  constexpr uint32_t kPhilox10A = 0x9E3779B9;
+  constexpr uint32_t kPhilox10B = 0xBB67AE85;
+
+  #pragma unroll
+  for (int i = 0; i < N_ROUNDS - 1; i++) {
+    ctr = philox_round(ctr, key);
+    key.x += kPhilox10A;
+    key.y += kPhilox10B;
+  }
+  return philox_round(ctr, key);
+}
+
+} // namespace at::cuda
diff --git a/aten/src/ATen/cuda/cub.cuh b/aten/src/ATen/cuda/cub.cuh
index ba097dafef791..5f4af0f42a065 100644
--- a/aten/src/ATen/cuda/cub.cuh
+++ b/aten/src/ATen/cuda/cub.cuh
@@ -70,7 +70,11 @@
 #define ATEN_CUB_MAXIMUM() NO_ROCM(at_cuda_detail)ROCM_HIPCUB(::cub)::Max()
 #endif
 
-#if defined(USE_ROCM)
+// Pre ROCm 8.0 we require backporting of FpLimits specialization
+// for 'c10::BFloat16'. However, ROCm 8.0 lifts the compatible
+// CUB version to 3.x which uses libhipcxx and derives numerical
+// limits from there.
+#if defined(USE_ROCM) && CUB_VERSION < 300000
 
 // backport https://github.com/NVIDIA/cub/pull/306 for c10::BFloat16
 
diff --git a/aten/src/ATen/cuda/cub_definitions.cuh b/aten/src/ATen/cuda/cub_definitions.cuh
index e0d7455d4c22b..7d9eb4096133f 100644
--- a/aten/src/ATen/cuda/cub_definitions.cuh
+++ b/aten/src/ATen/cuda/cub_definitions.cuh
@@ -1,14 +1,20 @@
 #pragma once
 
 #if !defined(USE_ROCM)
-#include <cuda.h>  // for CUDA_VERSION
-#endif
-
-#if !defined(USE_ROCM)
+#include <cuda.h> // for CUDA_VERSION
 #include <cub/version.cuh>
 #else
+// Check if we can find HIPCUB_CCCL_VERSION.
+#include <hipcub/hipcub_version.hpp>
+// Older versions of hipCUB do not support the CUB V3 API.
+#if defined(HIPCUB_CCCL_VERSION) && HIPCUB_CCCL_VERSION >= 300000
+#define CUB_VERSION HIPCUB_CCCL_VERSION
+#else
+// If we cannot find a compatible CCCL version, fallback to
+// a very old version.
 #define CUB_VERSION 200001
 #endif
+#endif
 
 // cub support for CUB_WRAPPED_NAMESPACE is added to cub 1.13.1 in:
 // https://github.com/NVIDIA/cub/pull/326
diff --git a/aten/src/ATen/cuda/detail/CUDAHooks.cpp b/aten/src/ATen/cuda/detail/CUDAHooks.cpp
index 03a3a97525a43..47e83d03e6b36 100644
--- a/aten/src/ATen/cuda/detail/CUDAHooks.cpp
+++ b/aten/src/ATen/cuda/detail/CUDAHooks.cpp
@@ -4,7 +4,6 @@
 #include <ATen/Context.h>
 #include <ATen/DeviceGuard.h>
 #include <ATen/DynamicLibrary.h>
-#include <ATen/core/Vitals.h>
 #include <ATen/cuda/CUDAConfig.h>
 #include <ATen/cuda/CUDADevice.h>
 #include <ATen/cuda/Exceptions.h>
@@ -86,9 +85,6 @@ struct _Initializer {
 // let's not if we don't need to!)
 void CUDAHooks::init() const {
   C10_LOG_API_USAGE_ONCE("aten.init.cuda");
-  // Force the update to enable unit testing. This code get executed before unit tests
-  // have a chance to enable vitals.
-  at::vitals::VitalsAPI.setVital("CUDA", "used", "true", /* force = */ true);
 
   const auto num_devices = c10::cuda::device_count_ensure_non_zero();
   c10::cuda::CUDACachingAllocator::init(num_devices);
@@ -558,7 +554,13 @@ const std::vector<std::string>& CUDAHooks::getHipblasltPreferredArchs() const {
     "gfx1200", "gfx1201",
 #endif
 #if ROCM_VERSION >= 70000
-    "gfx950"
+    "gfx950",
+#endif
+#if ROCM_VERSION >= 71300
+    "gfx1100", "gfx1101", "gfx1151",
+#endif
+#if ROCM_VERSION >= 71400
+    "gfx1250",
 #endif
   };
   return archs;
@@ -573,7 +575,7 @@ const std::vector<std::string>& CUDAHooks::getHipblasltSupportedArchs() const {
 #if ROCM_VERSION >= 70000
     "gfx950", "gfx1150", "gfx1151",
 #endif
-#if ROCM_VERSION >= 70200
+#if ROCM_VERSION >= 71400
     "gfx1250"
 #endif
   };
diff --git a/aten/src/ATen/cuda/detail/DeviceThreadHandles.h b/aten/src/ATen/cuda/detail/DeviceThreadHandles.h
index 71a344d281d2a..9ec194fc878b6 100644
--- a/aten/src/ATen/cuda/detail/DeviceThreadHandles.h
+++ b/aten/src/ATen/cuda/detail/DeviceThreadHandles.h
@@ -15,6 +15,7 @@
 
 #pragma once
 
+#include <map>
 #include <unordered_map>
 #include <vector>
 #include <utility>
@@ -114,9 +115,45 @@ struct DeviceThreadHandlePool : public std::enable_shared_from_this<DeviceThread
         return my_handles[device];
     }
 
+#ifdef USE_ROCM
+    // hipblaslt cannot share a single handle across multiple streams, so this
+    // overload returns a handle unique to each (device, stream) pair.
+    Handle_t reserve(int device, void* stream)
+    {
+        auto key = std::make_pair(device, stream);
+        // If this thread already has a handle for this (device, stream), return it
+        if(my_stream_handles.find(key) != my_stream_handles.end())
+        return my_stream_handles[key];
+
+        // otherwise, either grab a handle from the pool if one is available,
+        // or if not, create a new one.
+        auto parent = weak_parent.lock();
+        TORCH_CHECK(parent, "Cannot create handle during program termination");
+        std::lock_guard<std::mutex> guard(parent->mutex);
+
+        if(parent->available_handles[device].size() > 0)
+        {
+        my_stream_handles[key] = parent->available_handles[device].back();
+        parent->available_handles[device].pop_back();
+        }
+        else
+        {
+        parent->created_handles[device].emplace_back(true /*create*/);
+        my_stream_handles[key] = parent->created_handles[device].back().handle;
+        }
+
+        return my_stream_handles[key];
+    }
+#endif
+
     private:
     // Stores the per-device handles currently owned by this thread
     std::unordered_map<int, Handle_t> my_handles;
+#ifdef USE_ROCM
+    // Stores per-(device, stream) handles for ROCm, where hipblaslt
+    // requires a unique handle per stream.
+    std::map<std::pair<int, void*>, Handle_t> my_stream_handles;
+#endif
 
     std::weak_ptr<DeviceThreadHandlePool> weak_parent;
 
@@ -134,6 +171,18 @@ struct DeviceThreadHandlePool : public std::enable_shared_from_this<DeviceThread
             for(auto d_h : my_handles)
                 parent->available_handles[d_h.first].push_back(d_h.second);
         }
+#ifdef USE_ROCM
+        if(!my_stream_handles.empty()) {
+            auto parent = weak_parent.lock();
+            if (!parent) {
+                return;
+            }
+
+            std::lock_guard<std::mutex> guard(parent->mutex);
+            for(auto& [key, handle] : my_stream_handles)
+                parent->available_handles[key.first].push_back(handle);
+        }
+#endif
     }
     };
 
diff --git a/aten/src/ATen/cuda/detail/IndexUtils.cuh b/aten/src/ATen/cuda/detail/IndexUtils.cuh
index 367ab10d3d3bb..b8408f8fe2c61 100644
--- a/aten/src/ATen/cuda/detail/IndexUtils.cuh
+++ b/aten/src/ATen/cuda/detail/IndexUtils.cuh
@@ -33,4 +33,31 @@ getTensorInfo(const at::TensorBase &t) {
     data_ptr, dims, sz, st);
 }
 
+// ForwardIt: only legacy random access iterator is supported.
+template<class ForwardIt, class T, bool is_lower = true>
+static __host__ __device__ __forceinline__
+ForwardIt find_bound(ForwardIt first, ForwardIt last, const T& value) {
+    ForwardIt it;
+    typename std::iterator_traits<ForwardIt>::difference_type count, step;
+    // NOTE: std::distance(first, last) compiles but produces wrong results here,
+    // so only legacy random access iterators are safe in this code.
+    count = last - first;
+
+    while (count > 0) {
+      it = first;
+      step = count / 2;
+      // avoiding std::advance(it, step),
+      // although it does work unlike std::distance
+      it += step;
+      if (is_lower ? *it < value : value >= *it) {
+        first = ++it;
+        count -= step + 1;
+      }
+      else {
+        count = step;
+      }
+    }
+    return first;
+}
+
 } // namespace at::cuda::detail
diff --git a/aten/src/ATen/cuda/tunable/GemmHipblaslt.h b/aten/src/ATen/cuda/tunable/GemmHipblaslt.h
index 29affa2d21ff1..29c15720f4a66 100644
--- a/aten/src/ATen/cuda/tunable/GemmHipblaslt.h
+++ b/aten/src/ATen/cuda/tunable/GemmHipblaslt.h
@@ -94,6 +94,16 @@ constexpr hipDataType HipDataTypeFor<c10::Float4_e2m1fn_x2>() {
 #endif
 }
 
+template <typename T>
+constexpr hipblasComputeType_t HipBlasComputeTypeFor() {
+  return HIPBLAS_COMPUTE_32F;
+}
+
+template <>
+constexpr hipblasComputeType_t HipBlasComputeTypeFor<double>() {
+  return HIPBLAS_COMPUTE_64F;
+}
+
 template <typename T>
 int GetBatchFromParams(const GemmParams<T>* params) {
   return 1;
@@ -175,43 +185,43 @@ int GetStrideCFromParams(const ScaledGemmParams<T>* params) {
 }
 
 template <typename T>
-float GetAlphaFromParams(const GemmParams<T>* params) {
+at::opmath_type<T> GetAlphaFromParams(const GemmParams<T>* params) {
   return params->alpha;
 }
 
 template <typename T>
-float GetAlphaFromParams(const GemmAndBiasParams<T>* params) {
+at::opmath_type<T> GetAlphaFromParams(const GemmAndBiasParams<T>* params) {
   return params->alpha;
 }
 
 template <typename T>
-float GetAlphaFromParams(const GemmStridedBatchedParams<T>* params) {
+at::opmath_type<T> GetAlphaFromParams(const GemmStridedBatchedParams<T>* params) {
   return params->alpha;
 }
 
 template <typename T>
-float GetAlphaFromParams(const ScaledGemmParams<T>* params) {
-  return 1.0;
+at::opmath_type<T> GetAlphaFromParams(const ScaledGemmParams<T>* params) {
+  return at::opmath_type<T>{1.0};
 }
 
 template <typename T>
-float GetBetaFromParams(const GemmParams<T>* params) {
+at::opmath_type<T> GetBetaFromParams(const GemmParams<T>* params) {
   return params->beta;
 }
 
 template <typename T>
-float GetBetaFromParams(const GemmAndBiasParams<T>* params) {
-  return 0.0;
+at::opmath_type<T> GetBetaFromParams(const GemmAndBiasParams<T>* params) {
+  return at::opmath_type<T>{0.0};
 }
 
 template <typename T>
-float GetBetaFromParams(const GemmStridedBatchedParams<T>* params) {
+at::opmath_type<T> GetBetaFromParams(const GemmStridedBatchedParams<T>* params) {
   return params->beta;
 }
 
 template <typename T>
-float GetBetaFromParams(const ScaledGemmParams<T>* params) {
-  return 0.0;
+at::opmath_type<T> GetBetaFromParams(const ScaledGemmParams<T>* params) {
+  return at::opmath_type<T>{0.0};
 }
 
 template <typename T>
@@ -467,8 +477,9 @@ class HipblasltGemmOp : public Callable<ParamsT> {
 
       TORCH_CHECK(transa_outer == opa && transb_outer == opb, "trans mismatch, shouldn't happen");
 
-      float alpha = GetAlphaFromParams<CT>(params);
-      float beta = GetBetaFromParams<CT>(params);
+      using opmath_t = at::opmath_type<CT>;
+      opmath_t alpha = GetAlphaFromParams<CT>(params);
+      opmath_t beta = GetBetaFromParams<CT>(params);
 
       hipblasLtMatrixLayout_t mat_a, mat_b, mat_c;
       if (opa == HIPBLAS_OP_N) {
@@ -505,11 +516,14 @@ class HipblasltGemmOp : public Callable<ParamsT> {
             mat_c, HIPBLASLT_MATRIX_LAYOUT_STRIDED_BATCH_OFFSET, &stride_c, sizeof(stride_c)));
       }
 
-      hipblasComputeType_t computeType = HIPBLAS_COMPUTE_32F;
-      if (at::globalContext().float32Precision(at::Float32Backend::CUDA, at::Float32Op::MATMUL) == at::Float32Precision::TF32) {
-        computeType = HIPBLAS_COMPUTE_32F_FAST_TF32;
+      hipblasComputeType_t computeType = HipBlasComputeTypeFor<CT>();
+      if constexpr (std::is_same_v<CT, float>) {
+        if (at::globalContext().float32Precision(at::Float32Backend::CUDA, at::Float32Op::MATMUL) == at::Float32Precision::TF32) {
+          computeType = HIPBLAS_COMPUTE_32F_FAST_TF32;
+        }
       }
-      HipBlasLtMatmulDescriptor matmul(computeType, HIP_R_32F);
+      auto scale_type = HipDataTypeFor<opmath_t>();
+      HipBlasLtMatmulDescriptor matmul(computeType, scale_type);
       matmul.setAttribute(HIPBLASLT_MATMUL_DESC_TRANSA, opa);
       matmul.setAttribute(HIPBLASLT_MATMUL_DESC_TRANSB, opb);
 
@@ -630,9 +644,11 @@ auto GetHipBlasLtTypeStringAndOps() {
   }
 #endif
 
-  hipblasComputeType_t computeType = HIPBLAS_COMPUTE_32F;
-  if (at::globalContext().allowTF32CuBLAS()) {
-    computeType = HIPBLAS_COMPUTE_32F_FAST_TF32;
+  hipblasComputeType_t computeType = HipBlasComputeTypeFor<CT>();
+  if constexpr (std::is_same_v<CT, float>) {
+    if (at::globalContext().float32Precision(at::Float32Backend::CUDA, at::Float32Op::MATMUL) == at::Float32Precision::TF32) {
+      computeType = HIPBLAS_COMPUTE_32F_FAST_TF32;
+    }
   }
 
   hipblasLtHandle_t handle;
diff --git a/aten/src/ATen/functorch/BatchRulesActivation.cpp b/aten/src/ATen/functorch/BatchRulesActivation.cpp
index 92b5527db77c5..47fe73e4212c3 100644
--- a/aten/src/ATen/functorch/BatchRulesActivation.cpp
+++ b/aten/src/ATen/functorch/BatchRulesActivation.cpp
@@ -21,8 +21,8 @@ glu_batch_rule(const Tensor& self, std::optional<int64_t> self_bdim, int64_t dim
 
   const auto self_ = moveBatchDimToFront(self, self_bdim);
 
-  const auto res = at::glu(self_, dim_);
-  return std::make_tuple(res, 0);
+  auto res = at::glu(self_, dim_);
+  return std::make_tuple(std::move(res), 0);
 }
 
 static std::tuple<Tensor, std::optional<int64_t>> glu_backward_batch_rule(
@@ -42,8 +42,8 @@ static std::tuple<Tensor, std::optional<int64_t>> glu_backward_batch_rule(
   const auto grad_output_ = ensure_has_bdim(moveBatchDimToFront(grad_output, grad_output_bdim), grad_output_bdim.has_value(), batch_size);
   const auto self_ = ensure_has_bdim(moveBatchDimToFront(self, self_bdim), self_bdim.has_value(), batch_size);
 
-  const auto res = at::glu_backward(grad_output_, self_, dim_);
-  return std::make_tuple(res, 0);
+  auto res = at::glu_backward(grad_output_, self_, dim_);
+  return std::make_tuple(std::move(res), 0);
 }
 
 
diff --git a/aten/src/ATen/functorch/BatchRulesBinaryOps.cpp b/aten/src/ATen/functorch/BatchRulesBinaryOps.cpp
index 25fe96c6cda15..0acdd148e3eef 100644
--- a/aten/src/ATen/functorch/BatchRulesBinaryOps.cpp
+++ b/aten/src/ATen/functorch/BatchRulesBinaryOps.cpp
@@ -188,8 +188,8 @@ static std::tuple<Tensor, std::optional<int64_t>> masked_select_batch_rule(
   self_ = maybePadToLogicalRank(self_, 0, max_logical_rank);
 
   // masked_select returns a 1D tensor, so we have to reshape it into 2D
-  const auto result = at::masked_select(self_, mask).view({ batch_size, -1 });
-  return std::make_tuple(result, 0);
+  auto result = at::masked_select(self_, mask).view({ batch_size, -1 });
+  return std::make_tuple(std::move(result), 0);
 }
 
 static std::tuple<Tensor, std::optional<int64_t>> masked_select_backward_batch_rule(
@@ -213,8 +213,8 @@ static std::tuple<Tensor, std::optional<int64_t>> masked_select_backward_batch_r
   self_ = ensure_has_bdim(self_, self_bdim.has_value(), batch_size);
   grad_ = ensure_has_bdim(grad_, grad_bdim.has_value(), batch_size);
 
-  const auto result = at::masked_select_backward(grad_, self_.contiguous(), mask);
-  return std::make_tuple(result, 0);
+  auto result = at::masked_select_backward(grad_, self_.contiguous(), mask);
+  return std::make_tuple(std::move(result), 0);
 }
 
 static std::tuple<Tensor, std::optional<int64_t>> cdist_backward_batch_rule(
@@ -294,7 +294,7 @@ rrelu_with_noise_batch_rule(
 
   auto ret = at::rrelu_with_noise(self_, noise_, lower, upper, training, std::move(generator));
 
-  return std::make_tuple(ret, 0, noise_, 0);
+  return std::make_tuple(std::move(ret), 0, std::move(noise_), 0);
 }
 
 static Tensor rrelu_with_noise_batch(
diff --git a/aten/src/ATen/functorch/BatchRulesConvolution.cpp b/aten/src/ATen/functorch/BatchRulesConvolution.cpp
index 748d5b1687a3c..602517256cef1 100644
--- a/aten/src/ATen/functorch/BatchRulesConvolution.cpp
+++ b/aten/src/ATen/functorch/BatchRulesConvolution.cpp
@@ -329,7 +329,7 @@ convolution_backward_weight_batch_rule(
             dilation, transposed, output_padding, groups, mask);
         auto& grad_weight = std::get<1>(result);
         grad_weight = reshape_dim_outof_symint(1, batch_size, grad_weight);
-        return std::make_tuple(grad_weight, 1);
+        return std::make_tuple(std::move(grad_weight), 1);
       } else {
         // transposed: N(GO), BN(GI) -> N(GO), N(GBI) -> (GBI)O
         const auto dummy_weight = make_dummy(weight, weight_bdim, 0, batch_size);
diff --git a/aten/src/ATen/functorch/BatchRulesFactory.cpp b/aten/src/ATen/functorch/BatchRulesFactory.cpp
index 34a537a9edb40..d5450ed972115 100644
--- a/aten/src/ATen/functorch/BatchRulesFactory.cpp
+++ b/aten/src/ATen/functorch/BatchRulesFactory.cpp
@@ -94,14 +94,14 @@ static std::tuple<Tensor, std::optional<int64_t>> _new_zeros_with_same_feature_m
       // [K0, K1, B, 6], [B, 5], 2 -> [K0, K1, B, 5]
       tangent_ = tangent.movedim(*tangent_bdim, self_num_batch_dims);
     }
-    const auto result = at::_new_zeros_with_same_feature_meta(tangent_, base_, self_num_batch_dims);
-    return std::make_tuple(result, self_num_batch_dims);
+    auto result = at::_new_zeros_with_same_feature_meta(tangent_, base_, self_num_batch_dims);
+    return std::make_tuple(std::move(result), self_num_batch_dims);
   }
 
   // Case 1:
   auto tangent_ = moveBatchDimToFront(tangent, tangent_bdim);
   auto result = at::_new_zeros_with_same_feature_meta(tangent_, base, self_num_batch_dims + 1);
-  return std::make_tuple(result, 0);
+  return std::make_tuple(std::move(result), 0);
 }
 
 static std::tuple<Tensor, std::optional<int64_t>> linspace_logspace_batch_rule_helper(
@@ -139,7 +139,7 @@ static std::tuple<Tensor, std::optional<int64_t>> linspace_logspace_batch_rule_h
     result = result.to(*dtype);
   }
 
-  return std::make_tuple(result, 0);
+  return std::make_tuple(std::move(result), 0);
 }
 
 static std::tuple<Tensor, std::optional<int64_t>> linspace_Tensor_Tensor_batch_rule(
diff --git a/aten/src/ATen/functorch/BatchRulesHelper.h b/aten/src/ATen/functorch/BatchRulesHelper.h
index 0d2f075d0c540..f4583ac32a4a0 100644
--- a/aten/src/ATen/functorch/BatchRulesHelper.h
+++ b/aten/src/ATen/functorch/BatchRulesHelper.h
@@ -141,6 +141,8 @@ void boxed_tensor_inputs_batch_rule(const c10::OperatorHandle& op, torch::jit::S
   auto arguments = torch::jit::pop(*stack, num_arguments);
   std::vector<std::pair<Tensor, std::optional<int64_t>>> tensor_inputs;
   std::vector<int64_t> tensor_pos;
+  tensor_inputs.reserve(num_arguments);
+  tensor_pos.reserve(num_arguments);
   for (const auto idx : c10::irange(0, num_arguments)) {
     const auto& ivalue = arguments[idx];
     if (ivalue.isTensor()) {
diff --git a/aten/src/ATen/functorch/BatchRulesModules.cpp b/aten/src/ATen/functorch/BatchRulesModules.cpp
index 4e0b50c4e3fe7..51cb1d4751d6f 100644
--- a/aten/src/ATen/functorch/BatchRulesModules.cpp
+++ b/aten/src/ATen/functorch/BatchRulesModules.cpp
@@ -168,13 +168,20 @@ static std::tuple<Tensor, std::optional<int64_t>, Tensor, std::optional<int64_t>
 grid_sample_backward_helper_out(
     // NOLINTNEXTLINE(performance-unnecessary-value-param)
     std::tuple<Tensor, Tensor> bw_out,
-    int64_t grad_input_out_bdim,
-    int64_t grad_grid_out_bdim,
+    int64_t grad_input_bdim,
+    int64_t grad_grid_bdim,
     int64_t bdim_size) {
   auto& [grad_input, grad_grid] = bw_out;
-  grad_input = reshape_dim_outof(grad_input_out_bdim, bdim_size, grad_input);
-  grad_grid = reshape_dim_outof(grad_grid_out_bdim, bdim_size, grad_grid);
-  return std::make_tuple(std::move(grad_input), grad_input_out_bdim, std::move(grad_grid), grad_grid_out_bdim);
+  std::optional<int64_t> grad_input_bdim_out, grad_grid_bdim_out;
+  if (grad_input.defined()) {
+    grad_input = reshape_dim_outof(grad_input_bdim, bdim_size, grad_input);
+    grad_input_bdim_out = grad_input_bdim;
+  }
+  if (grad_grid.defined()) {
+    grad_grid = reshape_dim_outof(grad_grid_bdim, bdim_size, grad_grid);
+    grad_grid_bdim_out = grad_grid_bdim;
+  }
+  return std::make_tuple(std::move(grad_input), grad_input_bdim_out, std::move(grad_grid), grad_grid_bdim_out);
 }
 
 
diff --git a/aten/src/ATen/functorch/BatchRulesNorm.cpp b/aten/src/ATen/functorch/BatchRulesNorm.cpp
index 51dae00e6b7ed..affb6ce369f2b 100644
--- a/aten/src/ATen/functorch/BatchRulesNorm.cpp
+++ b/aten/src/ATen/functorch/BatchRulesNorm.cpp
@@ -284,7 +284,7 @@ std::tuple<at::Tensor,at::Tensor,at::Tensor> batch_norm_backward_plumbing(
         training, eps);
     grad_input = makeBatched(std::move(std::get<0>(results)), std::get<1>(results), cur_level);
   }
-  return std::make_tuple(grad_input, grad_weight, grad_bias);
+  return std::make_tuple(std::move(grad_input), std::move(grad_weight), std::move(grad_bias));
 }
 
 static std::tuple<Tensor,Tensor,Tensor> native_group_norm_plumbing(
@@ -351,7 +351,7 @@ static at::Tensor group_norm_backward_no_weight_bias_batch_rule(
   auto rstd_ = moveBatchDimToFront(rstd, rstd_bdim);
 
   const auto bdim_size = get_bdim_size2(grad_out, grad_out_bdim, input, input_bdim);
-  grad_out_ = ensure_has_bdim(grad_out, grad_out_bdim.has_value(), bdim_size);
+  grad_out_ = ensure_has_bdim(grad_out_, grad_out_bdim.has_value(), bdim_size);
   input_ = ensure_has_bdim(input_, input_bdim.has_value(), bdim_size);
   mean_ = ensure_has_bdim(mean_, mean_bdim.has_value(), bdim_size);
   rstd_ = ensure_has_bdim(rstd_, rstd_bdim.has_value(), bdim_size);
@@ -432,7 +432,7 @@ static std::tuple<Tensor,Tensor,Tensor> native_group_norm_backward_plumbing(
     );
     grad_input = makeBatched(std::move(tensor), 0, cur_level);
   }
-  return std::make_tuple(grad_input, grad_weight, grad_bias);
+  return std::make_tuple(std::move(grad_input), std::move(grad_weight), std::move(grad_bias));
 }
 
 static bool has_same_shape(
@@ -524,7 +524,7 @@ native_layer_norm_batch_rule(
     bias_ = maybePadToLogicalRank(bias_, /*has_bdim*/bias_bdim, result_logical_rank);
     result0 = result0 + bias_;
   }
-  return std::make_tuple(result0, 0, mean, stats_bdim, rstd, stats_bdim);
+  return std::make_tuple(std::move(result0), 0, std::move(mean), stats_bdim, std::move(rstd), stats_bdim);
 }
 
 static std::tuple<at::Tensor, std::optional<int64_t>> native_layer_norm_backward_no_weight_bias_batch_rule(
@@ -536,9 +536,9 @@ static std::tuple<at::Tensor, std::optional<int64_t>> native_layer_norm_backward
 
   if (!grad_out_bdim.has_value() && !input_bdim.has_value() &&
       !mean_bdim.has_value() && !rstd_bdim.has_value()) {
-    const auto result = at::native_layer_norm_backward(
+    auto result = at::native_layer_norm_backward(
         grad_out, input, normalized_shape, mean, rstd, std::nullopt, std::nullopt, {true, false, false});
-    return std::make_tuple(std::get<0>(result), std::nullopt);
+    return std::make_tuple(std::get<0>(std::move(result)), std::nullopt);
   }
 
   auto grad_out_ = moveBatchDimToFront(grad_out, grad_out_bdim);
@@ -644,7 +644,7 @@ static std::tuple<at::Tensor,at::Tensor,at::Tensor> native_layer_norm_backward_p
         rstd_value, rstd_bdim);
     grad_input = makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
   }
-  return std::make_tuple(grad_input, grad_weight, grad_bias);
+  return std::make_tuple(std::move(grad_input), std::move(grad_weight), std::move(grad_bias));
 }
 
 template <typename F, F Func>
diff --git a/aten/src/ATen/functorch/BatchRulesScatterOps.cpp b/aten/src/ATen/functorch/BatchRulesScatterOps.cpp
index 80034ff95ca3c..eab17f23ce2bf 100644
--- a/aten/src/ATen/functorch/BatchRulesScatterOps.cpp
+++ b/aten/src/ATen/functorch/BatchRulesScatterOps.cpp
@@ -647,7 +647,7 @@ std::tuple<Tensor, std::optional<int64_t>> index_put_batch_rule(
   values_ = maybe_permute_values(values_, indices, indices_bdims);
 
   auto result = at::index_put(self_, List<std::optional<Tensor>>(indices_), values_, accumulate);
-  return std::make_tuple(result, 0);
+  return std::make_tuple(std::move(result), 0);
 }
 
 // plumbing done since we don't support List<std::optional<Tensor>> in codegen
@@ -704,7 +704,7 @@ std::tuple<Tensor, std::optional<int64_t>> scatter_batch_rule(
   if (self_logical_rank == 0) {
     result = result.squeeze(-1);
   }
-  return std::make_tuple(result, 0);
+  return std::make_tuple(std::move(result), 0);
 }
 
 template <typename Func, typename ...Args>
@@ -742,7 +742,7 @@ inline std::tuple<Tensor, std::optional<int64_t>> scatter_batch_rule(
   if (self_logical_rank == 0) {
     result = result.squeeze(-1);
   }
-  return std::make_tuple(result, 0);
+  return std::make_tuple(std::move(result), 0);
 }
 
 } // namespace
@@ -774,11 +774,22 @@ std::tuple<Tensor, std::optional<int64_t>> scatter_add_batch_rule(
                             self, self_bdim, dim, index, index_bdim, src, src_bdim);
 }
 
+static void check_scatter_inplace_bdim(
+    std::optional<int64_t> self_bdim,
+    std::optional<int64_t> index_bdim,
+    std::optional<int64_t> src_bdim,
+    const char* schema_name) {
+  if (!self_bdim.has_value() && (index_bdim.has_value() || src_bdim.has_value())) {
+    vmapIncompatibleInplaceError(schema_name);
+  }
+}
+
 std::tuple<Tensor, std::optional<int64_t>> scatter_add__batch_rule(
     const Tensor& self, std::optional<int64_t> self_bdim,
     int64_t dim,
     const Tensor& index, std::optional<int64_t> index_bdim,
     const Tensor& src, std::optional<int64_t> src_bdim) {
+  check_scatter_inplace_bdim(self_bdim, index_bdim, src_bdim, "scatter_add_");
   return scatter_batch_rule(ATEN_FN(scatter_add_),
                             self, self_bdim, dim, index, index_bdim, src, src_bdim);
 }
@@ -811,6 +822,8 @@ std::tuple<Tensor, std::optional<int64_t>> scatter_reduce__two_batch_rule(
     const Tensor& src, std::optional<int64_t> src_bdim,
     const std::string_view reduce,
     bool include_self) {
+  check_scatter_inplace_bdim(
+      self_bdim, index_bdim, src_bdim, "scatter_reduce_");
   return scatter_batch_rule(ATEN_FN2(scatter_reduce_, two),
                             self, self_bdim, dim, index, index_bdim, src, src_bdim, reduce, include_self);
 }
@@ -852,7 +865,7 @@ std::tuple<Tensor, std::optional<int64_t>> gather_batch_rule(
   if (index_logical_rank == 0) {
     result = result.squeeze(-1);
   }
-  return std::make_tuple(result, 0);
+  return std::make_tuple(std::move(result), 0);
 }
 
 Tensor get_expanded_index(const Tensor& index, SymIntArrayRef self_size, int64_t dim) {
@@ -1000,7 +1013,7 @@ std::tuple<Tensor, std::optional<int64_t>> index_add_batch_rule_impl(
     if (self_logical_rank == 0) {
       result = result.squeeze(-1);
     }
-    return std::make_tuple(result, 0);
+    return std::make_tuple(std::move(result), 0);
   }
 
   // Index is batched. For-loop and stack is the best thing I can come up with
@@ -1071,7 +1084,7 @@ std::tuple<Tensor,Tensor> binary_pointwise_align(
   tensor_ = maybePadToLogicalRank(tensor_, self_bdim, max_logical_rank);
   other_ = maybePadToLogicalRank(other_, mask_bdim, max_logical_rank);
 
-  return std::make_tuple(tensor_, other_);
+  return std::make_tuple(std::move(tensor_), std::move(other_));
 }
 
 std::tuple<Tensor, std::optional<int64_t>> masked_fill_scalar_batch_rule(
@@ -1082,7 +1095,7 @@ std::tuple<Tensor, std::optional<int64_t>> masked_fill_scalar_batch_rule(
     const Scalar& source) {
   auto tensors = binary_pointwise_align(self, self_bdim, mask, mask_bdim);
   auto result = at::masked_fill(std::get<0>(tensors), std::get<1>(tensors), source);
-  return std::make_tuple(result, 0);
+  return std::make_tuple(std::move(result), 0);
 }
 
 std::tuple<Tensor, std::optional<int64_t>> index_fill_batch_rule_helper(
diff --git a/aten/src/ATen/functorch/BatchRulesUnaryOps.cpp b/aten/src/ATen/functorch/BatchRulesUnaryOps.cpp
index 3cabdd251480f..ef06451a17314 100644
--- a/aten/src/ATen/functorch/BatchRulesUnaryOps.cpp
+++ b/aten/src/ATen/functorch/BatchRulesUnaryOps.cpp
@@ -39,12 +39,12 @@ clone_batch_rule(
     // philosophically vmap hides the batch dims and operates on a per-sample level.
     auto self_ = moveBatchDimToFront(self, self_bdim);
     auto result = at::clone(self_, memory_format);
-    return std::make_tuple(result, 0);
+    return std::make_tuple(std::move(result), 0);
   }
 
   TORCH_INTERNAL_ASSERT(!memory_format.has_value() || memory_format == MemoryFormat::Preserve);
   auto result = at::clone(self, memory_format);
-  return std::make_tuple(result, self_bdim);
+  return std::make_tuple(std::move(result), self_bdim);
 }
 
 std::tuple<Tensor, std::optional<int64_t>>
@@ -55,7 +55,7 @@ view_as_complex_batch_rule(const Tensor& self, std::optional<int64_t> self_bdim)
 
   auto self_ = moveBatchDimToFront(self, self_bdim);
   auto result = at::view_as_complex(self_);
-  return std::make_tuple(result, 0);
+  return std::make_tuple(std::move(result), 0);
 }
 
 }
diff --git a/aten/src/ATen/functorch/BatchRulesViews.cpp b/aten/src/ATen/functorch/BatchRulesViews.cpp
index 08724d4fc1243..c451bb6861998 100644
--- a/aten/src/ATen/functorch/BatchRulesViews.cpp
+++ b/aten/src/ATen/functorch/BatchRulesViews.cpp
@@ -300,7 +300,7 @@ std::tuple<Tensor, std::optional<int64_t>> roll_batch_rule(const Tensor& self, s
   // NOTE: For scalar tensor, we don't need to unsqueeze as reshape
   // with `old_shape` takes care of it.
   output = output.reshape_symint(old_shape);
-  return std::make_tuple(output, 0);
+  return std::make_tuple(std::move(output), 0);
 }
 
 std::tuple<Tensor, std::optional<int64_t>> diagonal_batching_rule(
@@ -498,7 +498,6 @@ std::tuple<Tensor, std::optional<int64_t>> narrow_copy_batch_rule(
   auto logical_rank = rankWithoutBatchDim(self, self_bdim);
   dim = maybe_wrap_dim(dim, logical_rank) + 1;
   auto result = self_.narrow_copy_symint(dim, std::move(start), std::move(length));
-
   return std::make_tuple(std::move(result), 0);
 }
 
diff --git a/aten/src/ATen/functorch/BatchedFallback.cpp b/aten/src/ATen/functorch/BatchedFallback.cpp
index b479639f1c1a5..ee7b9b69eafb6 100644
--- a/aten/src/ATen/functorch/BatchedFallback.cpp
+++ b/aten/src/ATen/functorch/BatchedFallback.cpp
@@ -412,10 +412,8 @@ void batchedNestedTensorForLoopFallback(const c10::OperatorHandle& op, torch::ji
     return;
   }
 
-  if (isInplaceOp(schema)) {
-    TORCH_INTERNAL_ASSERT(false, "vmap fallback not supported for in-place ops on nested tensors");
-    return;
-  }
+  TORCH_INTERNAL_ASSERT(!isInplaceOp(schema), "vmap fallback not supported for in-place ops on nested tensors");
+
   TORCH_CHECK(!schema.is_mutable() && !schema.hasAnyAliasInfo(),
               "Nested batching rule not implemented for ", schema.operator_name(), "; ",
               "the fallback path doesn't work on out= or view ops.");
diff --git a/aten/src/ATen/functorch/BatchedTensorImpl.cpp b/aten/src/ATen/functorch/BatchedTensorImpl.cpp
index 895770fc69921..72af353064661 100644
--- a/aten/src/ATen/functorch/BatchedTensorImpl.cpp
+++ b/aten/src/ATen/functorch/BatchedTensorImpl.cpp
@@ -156,7 +156,6 @@ c10::intrusive_ptr<TensorImpl> BatchedTensorImpl::shallow_copy_and_detach(
     const c10::VariableVersion& version_counter,
     bool allow_tensor_metadata_change) const {
   TORCH_CHECK(false, "accessing `data` under vmap transform is not allowed");
-  return nullptr;
 }
 
 c10::intrusive_ptr<TensorImpl> BatchedTensorImpl::shallow_copy_and_detach(
@@ -164,7 +163,6 @@ c10::intrusive_ptr<TensorImpl> BatchedTensorImpl::shallow_copy_and_detach(
     c10::VariableVersion&& version_counter,
     bool allow_tensor_metadata_change) const {
   TORCH_CHECK(false, "accessing `data` under vmap transform is not allowed");
-  return nullptr;
 }
 
 void BatchedTensorImpl::shallow_copy_from(const c10::intrusive_ptr<TensorImpl>& impl) {
diff --git a/aten/src/ATen/functorch/DynamicLayer.cpp b/aten/src/ATen/functorch/DynamicLayer.cpp
index 1420aaf0ab943..2ef7c6d23ea2c 100644
--- a/aten/src/ATen/functorch/DynamicLayer.cpp
+++ b/aten/src/ATen/functorch/DynamicLayer.cpp
@@ -29,7 +29,8 @@ DynamicLayer::DynamicLayer(
     std::optional<RandomnessType> randomness,
     std::optional<bool> prev_grad_mode,
     std::optional<bool> prev_fwd_grad_mode,
-    std::optional<bool> functionalize_add_back_views)
+    std::optional<bool> functionalize_add_back_views,
+    std::optional<bool> prev_inference_mode)
 {
   if (transform_type == TransformType::Grad) {
     TORCH_INTERNAL_ASSERT(prev_grad_mode.has_value());
@@ -43,10 +44,10 @@ DynamicLayer::DynamicLayer(
       interpreter_ = Interpreter::Vmap(layerId, std::move(batchSize.value()), randomness.value());
       break;
     case TransformType::Grad:
-      interpreter_ = Interpreter::Grad(layerId, prev_grad_mode.value());
+      interpreter_ = Interpreter::Grad(layerId, prev_grad_mode.value(), prev_inference_mode.value_or(false));
       break;
     case TransformType::Jvp:
-      interpreter_ = Interpreter::Jvp(layerId, prev_fwd_grad_mode.value());
+      interpreter_ = Interpreter::Jvp(layerId, prev_fwd_grad_mode.value(), prev_inference_mode.value_or(false));
       break;
     case TransformType::Functionalize:
       // NOLINTNEXTLINE(bugprone-unchecked-optional-access)
@@ -246,10 +247,11 @@ int64_t initAndPushDynamicLayer(
     std::optional<RandomnessType> randomness,
     std::optional<bool> prev_grad_mode,
     std::optional<bool> prev_fwd_grad_mode,
-    std::optional<bool> functionalize_add_back_views) {
+    std::optional<bool> functionalize_add_back_views,
+    std::optional<bool> prev_inference_mode) {
   const auto& dynamicLayerStack = dynamicLayerStackAccessor();
   const int64_t layerId = static_cast<int64_t>(1 + dynamicLayerStack.size());
-  DynamicLayer new_layer(transform_type, layerId, std::move(batch_size), randomness, prev_grad_mode, prev_fwd_grad_mode, functionalize_add_back_views);
+  DynamicLayer new_layer(transform_type, layerId, std::move(batch_size), randomness, prev_grad_mode, prev_fwd_grad_mode, functionalize_add_back_views, prev_inference_mode);
   // NB: this function should be called while holding the GIL to avoid races
   new_layer.interpreter().set_is_alive(true);
   pushDynamicLayer(std::move(new_layer));
diff --git a/aten/src/ATen/functorch/DynamicLayer.h b/aten/src/ATen/functorch/DynamicLayer.h
index 672a33fda0016..bb203ea8ab8ea 100644
--- a/aten/src/ATen/functorch/DynamicLayer.h
+++ b/aten/src/ATen/functorch/DynamicLayer.h
@@ -47,7 +47,8 @@ struct TORCH_API DynamicLayer {
       std::optional<RandomnessType> randomness = std::nullopt,
       std::optional<bool> prev_grad_mode = std::nullopt,
       std::optional<bool> pre_fwd_grad_mode = std::nullopt,
-      std::optional<bool> functionalize_add_back_views = std::nullopt);
+      std::optional<bool> functionalize_add_back_views = std::nullopt,
+      std::optional<bool> prev_inference_mode = std::nullopt);
 
   TransformType key() const;
   int64_t layerId() const;
@@ -69,7 +70,8 @@ TORCH_API int64_t initAndPushDynamicLayer(
     std::optional<RandomnessType> randomness = std::nullopt,
     std::optional<bool> prev_grad_mode = std::nullopt,
     std::optional<bool> prev_fwd_grad_mode = std::nullopt,
-    std::optional<bool> functionalize_add_back_views = std::nullopt);
+    std::optional<bool> functionalize_add_back_views = std::nullopt,
+    std::optional<bool> prev_inference_mode = std::nullopt);
 TORCH_API DynamicLayer popDynamicLayerAndDeleteMetadata();
 TORCH_API std::optional<DynamicLayer> maybeCurrentDynamicLayer();
 TORCH_API const std::vector<DynamicLayer>& getDynamicLayerStack();
diff --git a/aten/src/ATen/functorch/Interpreter.h b/aten/src/ATen/functorch/Interpreter.h
index 3d3b2069387d7..02681f830fb7a 100644
--- a/aten/src/ATen/functorch/Interpreter.h
+++ b/aten/src/ATen/functorch/Interpreter.h
@@ -123,7 +123,8 @@ struct VmapInterpreterMeta {
 };
 
 struct GradInterpreterMeta {
-  explicit GradInterpreterMeta(bool prevGradMode): prevGradMode_(prevGradMode) {}
+  explicit GradInterpreterMeta(bool prevGradMode, bool prevInferenceMode = false)
+    : prevGradMode_(prevGradMode), prevInferenceMode_(prevInferenceMode) {}
   GradInterpreterMeta() = default;
   GradInterpreterMeta(const GradInterpreterMeta&) = default;
   GradInterpreterMeta(GradInterpreterMeta&&) = default;
@@ -132,19 +133,23 @@ struct GradInterpreterMeta {
   ~GradInterpreterMeta() = default;
 
   bool prevGradMode_;
+  bool prevInferenceMode_;
   template <typename T>
   friend void to_json(T& json_j, const GradInterpreterMeta& json_t) {
     json_j["prevGradMode"] = json_t.prevGradMode_;
+    json_j["prevInferenceMode"] = json_t.prevInferenceMode_;
   }
 
   template <typename T>
   friend void from_json(const T& json_j, GradInterpreterMeta& json_t) {
     json_t.prevGradMode_ = json_j["prevGradMode"];
+    json_t.prevInferenceMode_ = json_j.value("prevInferenceMode", false);
   }
 };
 
 struct JvpInterpreterMeta {
-  explicit JvpInterpreterMeta(bool prevFwdGradMode) : prevFwdGradMode_(prevFwdGradMode) {}
+  explicit JvpInterpreterMeta(bool prevFwdGradMode, bool prevInferenceMode = false)
+    : prevFwdGradMode_(prevFwdGradMode), prevInferenceMode_(prevInferenceMode) {}
   JvpInterpreterMeta() = default;
   JvpInterpreterMeta(const JvpInterpreterMeta&) = default;
   JvpInterpreterMeta(JvpInterpreterMeta&&) = default;
@@ -153,14 +158,17 @@ struct JvpInterpreterMeta {
   ~JvpInterpreterMeta() = default;
 
   bool prevFwdGradMode_;
+  bool prevInferenceMode_;
   template <typename T>
   friend void to_json(T& json_j, const JvpInterpreterMeta& json_t) {
     json_j["prevFwdGradMode"] = json_t.prevFwdGradMode_;
+    json_j["prevInferenceMode"] = json_t.prevInferenceMode_;
   }
 
   template <typename T>
   friend void from_json(const T& json_j, JvpInterpreterMeta& json_t) {
     json_t.prevFwdGradMode_ = json_j["prevFwdGradMode"];
+    json_t.prevInferenceMode_ = json_j.value("prevInferenceMode", false);
   }
 };
 
@@ -200,11 +208,11 @@ struct Interpreter {
   static Interpreter Vmap(int64_t level, c10::SymInt batchSize, RandomnessType randomness) {
     return Interpreter(TransformType::Vmap, level, VmapInterpreterMeta(std::move(batchSize), randomness));
   }
-  static Interpreter Grad(int64_t level, bool prevGradMode) {
-    return Interpreter(TransformType::Grad, level, GradInterpreterMeta(prevGradMode));
+  static Interpreter Grad(int64_t level, bool prevGradMode, bool prevInferenceMode = false) {
+    return Interpreter(TransformType::Grad, level, GradInterpreterMeta(prevGradMode, prevInferenceMode));
   }
-  static Interpreter Jvp(int64_t level, bool prevFwdGradMode) {
-    return Interpreter(TransformType::Jvp, level, JvpInterpreterMeta(prevFwdGradMode));
+  static Interpreter Jvp(int64_t level, bool prevFwdGradMode, bool prevInferenceMode = false) {
+    return Interpreter(TransformType::Jvp, level, JvpInterpreterMeta(prevFwdGradMode, prevInferenceMode));
   }
   static Interpreter Functionalize(int64_t level, bool functionalizeAddBackViews) {
     return Interpreter(TransformType::Functionalize, level, FunctionalizeInterpreterMeta(functionalizeAddBackViews));
diff --git a/aten/src/ATen/functorch/LegacyBatchingRegistrations.cpp b/aten/src/ATen/functorch/LegacyBatchingRegistrations.cpp
index 1df4c8938183a..e02f20b102bc7 100644
--- a/aten/src/ATen/functorch/LegacyBatchingRegistrations.cpp
+++ b/aten/src/ATen/functorch/LegacyBatchingRegistrations.cpp
@@ -709,8 +709,10 @@ Tensor nested_cat_batching_rule(const ITensorListRef& tensors, int64_t dim) {
   // Do a cat for each set of zipped unbound components
   const auto num_components = unbound.front().size();
   std::vector<Tensor> outputs;
+  outputs.reserve(num_components);
   for (auto i : c10::irange(num_components)) {
     std::vector<Tensor> arg_list;
+    arg_list.reserve(unbound.size());
     for (auto j : c10::irange(unbound.size())) {
       arg_list.push_back(unbound[j][i]);
     }
diff --git a/aten/src/ATen/functorch/PyTorchOperatorHacks.cpp b/aten/src/ATen/functorch/PyTorchOperatorHacks.cpp
index 667e92970033c..e41a080b1d714 100644
--- a/aten/src/ATen/functorch/PyTorchOperatorHacks.cpp
+++ b/aten/src/ATen/functorch/PyTorchOperatorHacks.cpp
@@ -111,7 +111,9 @@ Tensor binary_cross_entropy_with_logits_hack(
 Tensor trace_backward_decomp(const Tensor& grad, IntArrayRef sizes) {
   TORCH_CHECK(sizes.size() == 2, "expected matrix input");
   auto grad_input = at::zeros(sizes[0] * sizes[1], grad.options());
-  auto indices = at::arange(0, grad_input.numel(), sizes[1] + 1, grad.options().dtype(at::kLong));
+  auto diag_size = std::min(sizes[0], sizes[1]);
+  auto step = sizes[1] + 1;
+  auto indices = at::arange(0, diag_size * step, step, grad.options().dtype(at::kLong));
   // Workaround using index_put instead of yet unsupported index_fill_
   grad_input = grad_input.index_put({indices}, grad);
   return grad_input.view(sizes);
diff --git a/aten/src/ATen/hip/impl/HIPStreamMasqueradingAsCUDA.h b/aten/src/ATen/hip/impl/HIPStreamMasqueradingAsCUDA.h
index a6d24a55c7b0d..1c0047371d7b6 100644
--- a/aten/src/ATen/hip/impl/HIPStreamMasqueradingAsCUDA.h
+++ b/aten/src/ATen/hip/impl/HIPStreamMasqueradingAsCUDA.h
@@ -24,33 +24,6 @@ class HIPStreamMasqueradingAsCUDA final : public c10::cuda::CUDAStream {
   c10::cuda::CUDAStream hip_stream() const { return *this; }
 };
 
-HIPStreamMasqueradingAsCUDA
-inline getStreamFromPoolMasqueradingAsCUDA(const bool isHighPriority = false, DeviceIndex device = -1) {
-  return HIPStreamMasqueradingAsCUDA(c10::cuda::getStreamFromPool(isHighPriority, device));
-}
-
-HIPStreamMasqueradingAsCUDA
-inline getStreamFromPoolMasqueradingAsCUDA(const int priority, DeviceIndex device = -1) {
-  return HIPStreamMasqueradingAsCUDA(c10::cuda::getStreamFromPool(priority, device));
-}
-
-HIPStreamMasqueradingAsCUDA
-inline getStreamFromExternalMasqueradingAsCUDA(hipStream_t ext_stream, DeviceIndex device) {
-  return HIPStreamMasqueradingAsCUDA(c10::cuda::getStreamFromExternal(ext_stream, device));
-}
-
-inline HIPStreamMasqueradingAsCUDA getDefaultHIPStreamMasqueradingAsCUDA(DeviceIndex device_index = -1) {
-  return HIPStreamMasqueradingAsCUDA(c10::cuda::getDefaultCUDAStream(device_index));
-}
-
-inline HIPStreamMasqueradingAsCUDA getCurrentHIPStreamMasqueradingAsCUDA(DeviceIndex device_index = -1) {
-  return HIPStreamMasqueradingAsCUDA(c10::cuda::getCurrentCUDAStream(device_index));
-}
-
-inline void setCurrentHIPStreamMasqueradingAsCUDA(HIPStreamMasqueradingAsCUDA stream) {
-  c10::cuda::setCurrentCUDAStream(stream.hip_stream());
-}
-
 inline std::ostream& operator<<(std::ostream& stream, const HIPStreamMasqueradingAsCUDA& s) {
   stream << s.hip_stream() << " (masquerading as CUDA)";
   return stream;
diff --git a/aten/src/ATen/miopen/Exceptions.h b/aten/src/ATen/miopen/Exceptions.h
index f5f0a4785b1cf..5db684e362a31 100644
--- a/aten/src/ATen/miopen/Exceptions.h
+++ b/aten/src/ATen/miopen/Exceptions.h
@@ -1,3 +1,4 @@
+// @allow-raw-throw
 #pragma once
 
 #include <ATen/miopen/miopen-wrapper.h>
diff --git a/aten/src/ATen/mkl/SparseBlas.cpp b/aten/src/ATen/mkl/SparseBlas.cpp
index d2a8d902b56ef..99e2174f7d13e 100644
--- a/aten/src/ATen/mkl/SparseBlas.cpp
+++ b/aten/src/ATen/mkl/SparseBlas.cpp
@@ -2,6 +2,9 @@
   Provides the implementations of MKL Sparse BLAS function templates.
 */
 #define TORCH_ASSERT_NO_OPERATORS
+#include <ATen/mkl/Sparse.h>
+
+#if AT_USE_MKL_SPARSE()
 #include <ATen/mkl/Exceptions.h>
 #include <ATen/mkl/SparseBlas.h>
 
@@ -364,3 +367,4 @@ sparse_status_t trsm<c10::complex<double>>(
 }
 
 } // namespace at::mkl::sparse
+#endif
diff --git a/aten/src/ATen/mps/EmptyTensor.cpp b/aten/src/ATen/mps/EmptyTensor.cpp
index 6c58de099648d..8c722a255ea36 100644
--- a/aten/src/ATen/mps/EmptyTensor.cpp
+++ b/aten/src/ATen/mps/EmptyTensor.cpp
@@ -12,7 +12,7 @@
 
 #define MPS_ERROR_NOT_COMPILED "PyTorch code is not compiled with MPS enabled"
 #define MPS_ERROR_RUNTIME_TOO_LOW \
-  "The MPS backend is supported on MacOS 14.0+. ", \
+  "The MPS backend is supported on macOS 14.0+. ", \
   "Current OS version can be queried using `sw_vers`"
 #define MPS_ERROR_DOUBLE_NOT_SUPPORTED "Cannot convert a MPS Tensor to float64 dtype " \
   "as the MPS framework doesn't support float64. Please use float32 instead."
@@ -42,7 +42,7 @@ TensorBase empty_mps(
     auto* allocator = at::mps::GetMPSAllocator();
     int64_t nelements = c10::multiply_integers(size);
     auto dtype = dtype_or_default(dtype_opt);
-    TORCH_CHECK_TYPE(dtype != ScalarType::Double, MPS_ERROR_DOUBLE_NOT_SUPPORTED);
+    TORCH_CHECK_TYPE(dtype != kDouble && dtype != kComplexDouble, MPS_ERROR_DOUBLE_NOT_SUPPORTED);
 
 
     auto dtype_meta = scalarTypeToTypeMeta(dtype);
diff --git a/aten/src/ATen/mps/MPSAllocator.h b/aten/src/ATen/mps/MPSAllocator.h
index 1132ca6219cdf..9708fea095008 100644
--- a/aten/src/ATen/mps/MPSAllocator.h
+++ b/aten/src/ATen/mps/MPSAllocator.h
@@ -227,8 +227,6 @@ typedef bool (*HeapComparison)(const HeapBlock*, const HeapBlock*);
 
 struct BufferPool {
   enum class Kind {
-    PRIVATE_SMALL,
-    PRIVATE_LARGE,
     SHARED_SMALL,
     SHARED_LARGE,
     SCALAR,
diff --git a/aten/src/ATen/mps/MPSAllocator.mm b/aten/src/ATen/mps/MPSAllocator.mm
index dfdd67c8f4458..62bacd0a41f61 100644
--- a/aten/src/ATen/mps/MPSAllocator.mm
+++ b/aten/src/ATen/mps/MPSAllocator.mm
@@ -40,15 +40,9 @@
 
 void MPSHeapAllocatorImpl::init_buffer_pools() {
   // using a container for pools to simplify iterating over them
-  // Pool of large buffers with private storage mode
-  m_pools.emplace(BufferPool::Kind::PRIVATE_LARGE,
-                  std::make_unique<BufferPool>(m_device, UsageFlags::PRIVATE | UsageFlags::HAZARD));
   // Pool of large buffers with shared storage mode
   m_pools.emplace(BufferPool::Kind::SHARED_LARGE,
                   std::make_unique<BufferPool>(m_device, UsageFlags::SHARED | UsageFlags::HAZARD));
-  // Pool of small buffers with private storage mode
-  m_pools.emplace(BufferPool::Kind::PRIVATE_SMALL,
-                  std::make_unique<BufferPool>(m_device, UsageFlags::SMALL | UsageFlags::PRIVATE | UsageFlags::HAZARD));
   // Pool of small buffers with shared storage mode
   m_pools.emplace(BufferPool::Kind::SHARED_SMALL,
                   std::make_unique<BufferPool>(m_device, UsageFlags::SMALL | UsageFlags::SHARED | UsageFlags::HAZARD));
@@ -64,12 +58,10 @@
 
   if (usage & UsageFlags::SCALAR) {
     poolKind = BufferPool::Kind::SCALAR;
-  } else if (requested_size <= kMaxScalarAlloc && m_device.hasUnifiedMemory) {
-    poolKind = BufferPool::Kind::SHARED_SMALL;
   } else if (aligned_size <= kMaxSmallAlloc) {
-    poolKind = (usage & UsageFlags::SHARED) ? BufferPool::Kind::SHARED_SMALL : BufferPool::Kind::PRIVATE_SMALL;
+    poolKind = BufferPool::Kind::SHARED_SMALL;
   } else {
-    poolKind = (usage & UsageFlags::SHARED) ? BufferPool::Kind::SHARED_LARGE : BufferPool::Kind::PRIVATE_LARGE;
+    poolKind = BufferPool::Kind::SHARED_LARGE;
   }
   return *m_pools[poolKind];
 }
@@ -830,16 +822,9 @@ static void Delete(void* ptr) {
   return s_mps_shared_alloc;
 }
 
-MPSAllocator& _getPrivateAllocator() {
-  static MPSAllocator s_mps_private_alloc(HeapAllocator::UsageFlags::PRIVATE);
-  return s_mps_private_alloc;
-}
 } // anonymous namespace
 
-IMPSAllocator* getIMPSAllocator(bool sharedAllocator) {
-  if (!sharedAllocator) {
-    return &_getPrivateAllocator();
-  }
+IMPSAllocator* getIMPSAllocator() {
   auto& sa = _getSharedAllocator();
   if (sa.isSharedStorageSupported()) {
     return &sa;
diff --git a/aten/src/ATen/mps/MPSAllocatorInterface.h b/aten/src/ATen/mps/MPSAllocatorInterface.h
index 996c84a8b87f5..d8480408f3b71 100644
--- a/aten/src/ATen/mps/MPSAllocatorInterface.h
+++ b/aten/src/ATen/mps/MPSAllocatorInterface.h
@@ -61,7 +61,7 @@ TORCH_DECLARE_REGISTRY(MPSAllocatorCallbacksRegistry, IMpsAllocatorCallback);
 #define REGISTER_MPS_ALLOCATOR_CALLBACK(name, ...) \
   C10_REGISTER_CLASS(MPSAllocatorCallbacksRegistry, name, __VA_ARGS__)
 
-IMPSAllocator* getIMPSAllocator(bool sharedAllocator = false);
+IMPSAllocator* getIMPSAllocator();
 
 bool isMPSPinnedPtr(const void* data);
 
diff --git a/aten/src/ATen/mps/MPSDevice.h b/aten/src/ATen/mps/MPSDevice.h
index a6ef221d099f0..67f65b99c79e0 100644
--- a/aten/src/ATen/mps/MPSDevice.h
+++ b/aten/src/ATen/mps/MPSDevice.h
@@ -23,6 +23,7 @@ enum class MacOSVersion : uint32_t {
   MACOS_VER_15_1_PLUS,
   MACOS_VER_15_2_PLUS,
   MACOS_VER_26_0_PLUS,
+  MACOS_VER_26_4_PLUS,
 };
 
 //-----------------------------------------------------------------
@@ -77,7 +78,7 @@ class TORCH_API MPSDevice {
 
 TORCH_API bool is_available();
 TORCH_API bool is_macos_13_or_newer(MacOSVersion version);
-TORCH_API at::Allocator* GetMPSAllocator(bool useSharedAllocator = false);
+TORCH_API at::Allocator* GetMPSAllocator();
 
 inline Device getDeviceFromPtr(void* ptr) {
   return {c10::DeviceType::MPS, 0};
diff --git a/aten/src/ATen/mps/MPSDevice.mm b/aten/src/ATen/mps/MPSDevice.mm
index 5b30d2902e184..4e864aeaa73f8 100644
--- a/aten/src/ATen/mps/MPSDevice.mm
+++ b/aten/src/ATen/mps/MPSDevice.mm
@@ -66,6 +66,7 @@ static inline MTLLanguageVersion getMetalLanguageVersion(const id<MTLDevice>& de
   static bool _macos_15_1_plus = is_os_version_at_least(15, 1);
   static bool _macos_15_2_plus = is_os_version_at_least(15, 2);
   static bool _macos_26_0_plus = is_os_version_at_least(26, 0);
+  static bool _macos_26_4_plus = is_os_version_at_least(26, 4);
 
   switch (version) {
     case MacOSVersion::MACOS_VER_14_4_PLUS:
@@ -78,6 +79,8 @@ static inline MTLLanguageVersion getMetalLanguageVersion(const id<MTLDevice>& de
       return _macos_15_2_plus;
     case MacOSVersion::MACOS_VER_26_0_PLUS:
       return _macos_26_0_plus;
+    case MacOSVersion::MACOS_VER_26_4_PLUS:
+      return _macos_26_4_plus;
     default:
       return false;
   }
@@ -110,8 +113,8 @@ static inline MTLLanguageVersion getMetalLanguageVersion(const id<MTLDevice>& de
   return core_count;
 }
 
-at::Allocator* GetMPSAllocator(bool useSharedAllocator) {
-  return getIMPSAllocator(useSharedAllocator);
+at::Allocator* GetMPSAllocator() {
+  return getIMPSAllocator();
 }
 bool is_available() {
   return MPSDevice::getInstance()->device() != nil;
diff --git a/aten/src/ATen/mps/MPSGuardImpl.h b/aten/src/ATen/mps/MPSGuardImpl.h
index 008a8d57f3df6..368515f233e9b 100644
--- a/aten/src/ATen/mps/MPSGuardImpl.h
+++ b/aten/src/ATen/mps/MPSGuardImpl.h
@@ -84,6 +84,18 @@ struct TORCH_API MPSGuardImpl final
   Stream exchangeStream(Stream s) const override {
     return Stream(Stream::DEFAULT, Device(c10::DeviceType::MPS, 0));
   }
+  DeviceCapability getDeviceCapability(Device /* unused */) const override {
+    DeviceCapability cap;
+    cap.capability_data.capability_bits = (1ULL << kIndex_Byte) |
+        (1ULL << kIndex_Char) | (1ULL << kIndex_Short) | (1ULL << kIndex_Int) |
+        (1ULL << kIndex_Long) | (1ULL << kIndex_Half) | (1ULL << kIndex_Float) |
+        (1ULL << kIndex_ComplexHalf) | (1ULL << kIndex_ComplexFloat) |
+        (1ULL << kIndex_Bool) | (1ULL << kIndex_BFloat16) |
+        (1ULL << kIndex_UInt32) | (1ULL << kIndex_UInt16) |
+        (1ULL << kIndex_UInt64);
+    return cap;
+  }
+
   DeviceIndex deviceCount() const noexcept override {
     if (at::hasMPS()) {
       // TODO: extend it for multi-device case
diff --git a/aten/src/ATen/mps/MPSHooks.mm b/aten/src/ATen/mps/MPSHooks.mm
index 662b77c225b48..2471fcb31cbe4 100644
--- a/aten/src/ATen/mps/MPSHooks.mm
+++ b/aten/src/ATen/mps/MPSHooks.mm
@@ -153,7 +153,7 @@
 }
 
 Allocator* MPSHooks::getPinnedMemoryAllocator() const {
-  return at::mps::getIMPSAllocator(true);
+  return at::mps::getIMPSAllocator();
 }
 
 using at::MPSHooksRegistry;
diff --git a/aten/src/ATen/mps/MPSStream.h b/aten/src/ATen/mps/MPSStream.h
index b00890b9f5901..f83a387ee846e 100644
--- a/aten/src/ATen/mps/MPSStream.h
+++ b/aten/src/ATen/mps/MPSStream.h
@@ -73,7 +73,6 @@ class TORCH_API MPSStream {
   MTLComputeCommandEncoder_t commandEncoder();
   void endKernelCoalescing();
   void synchronize(SyncType syncType);
-  void fill(MTLBuffer_t buffer, uint8_t value, size_t length, size_t offset, SyncType syncType = SyncType::NONE);
   void copy(MTLBuffer_t srcBuffer,
             MTLBuffer_t dstBuffer,
             size_t length,
diff --git a/aten/src/ATen/mps/MPSStream.mm b/aten/src/ATen/mps/MPSStream.mm
index b02d2447ce17f..8cecc94f5926b 100644
--- a/aten/src/ATen/mps/MPSStream.mm
+++ b/aten/src/ATen/mps/MPSStream.mm
@@ -157,33 +157,6 @@ @interface MPSGraphExecutionDescriptor ()
   });
 }
 
-void MPSStream::fill(id<MTLBuffer> buffer, uint8_t value, size_t length, size_t offset, SyncType syncType) {
-  if (length == 0) {
-    return;
-  }
-  dispatch_sync_with_rethrow(_serialQueue, ^() {
-    @autoreleasepool {
-      endKernelCoalescing();
-      id<MTLBlitCommandEncoder> blitEncoder = [commandBuffer() blitCommandEncoder];
-
-      // For some reason fillBufferfor stopped working for length > 4Gb on MacOS 26
-      // See https://github.com/pytorch/pytorch/issues/163962
-      // Workaround by batching copy commands into 4Gb chunks
-      constexpr size_t max_copy_size = 0x100000000; // 4GB
-      size_t bytes_filled = 0;
-      size_t bytes_remains = length;
-      while (bytes_remains > 0) {
-        NSUInteger bytes_to_copy = std::min(max_copy_size, bytes_remains);
-        [blitEncoder fillBuffer:buffer range:NSMakeRange(offset + bytes_filled, bytes_to_copy) value:value];
-        bytes_filled += bytes_to_copy;
-        bytes_remains -= bytes_to_copy;
-      }
-      [blitEncoder endEncoding];
-      synchronize(syncType);
-    }
-  });
-}
-
 void MPSStream::copy(id<MTLBuffer> srcBuffer,
                      id<MTLBuffer> dstBuffer,
                      size_t length,
diff --git a/aten/src/ATen/native/BatchLinearAlgebra.cpp b/aten/src/ATen/native/BatchLinearAlgebra.cpp
index 77d4ed32c2f7a..2de3e24b584b8 100644
--- a/aten/src/ATen/native/BatchLinearAlgebra.cpp
+++ b/aten/src/ATen/native/BatchLinearAlgebra.cpp
@@ -1843,7 +1843,21 @@ TORCH_IMPL_FUNC(linalg_cholesky_ex_out)(const Tensor& A,
 
   cholesky_stub(L.device().type(), L, info, upper);
 
-  if (!cpu) {
+  // On non-CPU devices (MAGMA) the pre-copy doesn't zero the unused triangle,
+  // so we must clean up after. On macOS, Accelerate's LAPACK writes into the
+  // unreferenced triangle for matrices larger than its internal block size
+  // (e.g. n > 64), violating the LAPACK spec which says "not referenced"
+  // elements are "never read, written to, or otherwise accessed"
+  // (see https://www.netlib.org/lapack/lug/node121.html).
+  // We work around this by applying the same cleanup on macOS.
+  // TODO(https://github.com/pytorch/pytorch/issues/179152): always
+  // clean up the unused triangle on all platforms.
+#if defined(__APPLE__)
+  constexpr bool needs_triangle_cleanup = true;
+#else
+  const bool needs_triangle_cleanup = !cpu;
+#endif
+  if (needs_triangle_cleanup) {
     if (upper) {
       L.triu_();
     } else {
@@ -3065,7 +3079,7 @@ std::tuple<Tensor, Tensor> linalg_eig(const Tensor& input) {
 
   at::linalg_eig_outf(input, values, vectors);
 
-  return std::tuple<Tensor, Tensor>(values, vectors);
+  return std::tuple<Tensor, Tensor>(std::move(values), std::move(vectors));
 }
 
 Tensor& linalg_eigvals_out(const Tensor& input, Tensor& values) {
diff --git a/aten/src/ATen/native/BatchLinearAlgebraKernel.cpp b/aten/src/ATen/native/BatchLinearAlgebraKernel.cpp
index c4d7e1fdfe9d1..03aa41a8a33f2 100644
--- a/aten/src/ATen/native/BatchLinearAlgebraKernel.cpp
+++ b/aten/src/ATen/native/BatchLinearAlgebraKernel.cpp
@@ -1208,115 +1208,20 @@ void unpack_pivots_cpu_kernel(TensorIterator& iter, const int64_t dim_size, cons
 }
 } // anonymous namespace
 
-REGISTER_ARCH_DISPATCH(cholesky_stub, DEFAULT, &cholesky_kernel)
-REGISTER_AVX512_DISPATCH(cholesky_stub, &cholesky_kernel)
-REGISTER_AVX2_DISPATCH(cholesky_stub, &cholesky_kernel)
-REGISTER_VSX_DISPATCH(cholesky_stub, &cholesky_kernel)
-REGISTER_ZVECTOR_DISPATCH(cholesky_stub, &cholesky_kernel)
-REGISTER_SVE256_DISPATCH(cholesky_stub, &cholesky_kernel)
-
-REGISTER_ARCH_DISPATCH(cholesky_inverse_stub, DEFAULT, &cholesky_inverse_kernel_impl)
-REGISTER_AVX512_DISPATCH(cholesky_inverse_stub, &cholesky_inverse_kernel_impl)
-REGISTER_AVX2_DISPATCH(cholesky_inverse_stub, &cholesky_inverse_kernel_impl)
-REGISTER_VSX_DISPATCH(cholesky_inverse_stub, &cholesky_inverse_kernel_impl)
-REGISTER_ZVECTOR_DISPATCH(cholesky_inverse_stub, &cholesky_inverse_kernel_impl)
-REGISTER_SVE256_DISPATCH(cholesky_inverse_stub, &cholesky_inverse_kernel_impl)
-
-REGISTER_ARCH_DISPATCH(linalg_eig_make_complex_eigenvectors_stub, DEFAULT, &linalg_eig_make_complex_eigenvectors_cpu)
-REGISTER_AVX512_DISPATCH(linalg_eig_make_complex_eigenvectors_stub, &linalg_eig_make_complex_eigenvectors_cpu)
-REGISTER_AVX2_DISPATCH(linalg_eig_make_complex_eigenvectors_stub, &linalg_eig_make_complex_eigenvectors_cpu)
-REGISTER_VSX_DISPATCH(linalg_eig_make_complex_eigenvectors_stub, &linalg_eig_make_complex_eigenvectors_cpu)
-REGISTER_ZVECTOR_DISPATCH(linalg_eig_make_complex_eigenvectors_stub, &linalg_eig_make_complex_eigenvectors_cpu)
-REGISTER_SVE256_DISPATCH(linalg_eig_make_complex_eigenvectors_stub, &linalg_eig_make_complex_eigenvectors_cpu)
-
-REGISTER_ARCH_DISPATCH(linalg_eig_stub, DEFAULT, &linalg_eig_kernel)
-REGISTER_AVX512_DISPATCH(linalg_eig_stub, &linalg_eig_kernel)
-REGISTER_AVX2_DISPATCH(linalg_eig_stub, &linalg_eig_kernel)
-REGISTER_VSX_DISPATCH(linalg_eig_stub, &linalg_eig_kernel)
-REGISTER_ZVECTOR_DISPATCH(linalg_eig_stub, &linalg_eig_kernel)
-REGISTER_SVE256_DISPATCH(linalg_eig_stub, &linalg_eig_kernel)
-
-REGISTER_ARCH_DISPATCH(linalg_eigh_stub, DEFAULT, &linalg_eigh_kernel)
-REGISTER_AVX512_DISPATCH(linalg_eigh_stub, &linalg_eigh_kernel)
-REGISTER_AVX2_DISPATCH(linalg_eigh_stub, &linalg_eigh_kernel)
-REGISTER_VSX_DISPATCH(linalg_eigh_stub, &linalg_eigh_kernel)
-REGISTER_ZVECTOR_DISPATCH(linalg_eigh_stub, &linalg_eigh_kernel)
-REGISTER_SVE256_DISPATCH(linalg_eigh_stub, &linalg_eigh_kernel)
-
-REGISTER_ARCH_DISPATCH(geqrf_stub, DEFAULT, &geqrf_kernel)
-REGISTER_AVX512_DISPATCH(geqrf_stub, &geqrf_kernel)
-REGISTER_AVX2_DISPATCH(geqrf_stub, &geqrf_kernel)
-REGISTER_VSX_DISPATCH(geqrf_stub, &geqrf_kernel)
-REGISTER_ZVECTOR_DISPATCH(geqrf_stub, &geqrf_kernel)
-REGISTER_SVE256_DISPATCH(geqrf_stub, &geqrf_kernel)
-
-REGISTER_ARCH_DISPATCH(orgqr_stub, DEFAULT, &orgqr_kernel_impl)
-REGISTER_AVX512_DISPATCH(orgqr_stub, &orgqr_kernel_impl)
-REGISTER_AVX2_DISPATCH(orgqr_stub, &orgqr_kernel_impl)
-REGISTER_VSX_DISPATCH(orgqr_stub, &orgqr_kernel_impl)
-REGISTER_ZVECTOR_DISPATCH(orgqr_stub, &orgqr_kernel_impl)
-REGISTER_SVE256_DISPATCH(orgqr_stub, &orgqr_kernel_impl)
-
-REGISTER_ARCH_DISPATCH(ormqr_stub, DEFAULT, &ormqr_kernel)
-REGISTER_AVX512_DISPATCH(ormqr_stub, &ormqr_kernel)
-REGISTER_AVX2_DISPATCH(ormqr_stub, &ormqr_kernel)
-REGISTER_VSX_DISPATCH(ormqr_stub, &ormqr_kernel)
-REGISTER_ZVECTOR_DISPATCH(ormqr_stub, &ormqr_kernel)
-REGISTER_SVE256_DISPATCH(ormqr_stub, &ormqr_kernel)
-
-REGISTER_ARCH_DISPATCH(lstsq_stub, DEFAULT, &lstsq_kernel)
-REGISTER_AVX512_DISPATCH(lstsq_stub, &lstsq_kernel)
-REGISTER_AVX2_DISPATCH(lstsq_stub, &lstsq_kernel)
-REGISTER_VSX_DISPATCH(lstsq_stub, &lstsq_kernel)
-REGISTER_ZVECTOR_DISPATCH(lstsq_stub, &lstsq_kernel)
-REGISTER_SVE256_DISPATCH(lstsq_stub, &lstsq_kernel)
-
-REGISTER_ARCH_DISPATCH(triangular_solve_stub, DEFAULT, &triangular_solve_kernel)
-REGISTER_AVX512_DISPATCH(triangular_solve_stub, &triangular_solve_kernel)
-REGISTER_AVX2_DISPATCH(triangular_solve_stub, &triangular_solve_kernel)
-REGISTER_VSX_DISPATCH(triangular_solve_stub, &triangular_solve_kernel)
-REGISTER_ZVECTOR_DISPATCH(triangular_solve_stub, &triangular_solve_kernel)
-REGISTER_SVE256_DISPATCH(triangular_solve_stub, &triangular_solve_kernel)
-
-REGISTER_ARCH_DISPATCH(lu_factor_stub, DEFAULT, &lu_factor_kernel)
-REGISTER_AVX512_DISPATCH(lu_factor_stub, &lu_factor_kernel)
-REGISTER_AVX2_DISPATCH(lu_factor_stub, &lu_factor_kernel)
-REGISTER_VSX_DISPATCH(lu_factor_stub, &lu_factor_kernel)
-REGISTER_ZVECTOR_DISPATCH(lu_factor_stub, &lu_factor_kernel)
-REGISTER_SVE256_DISPATCH(lu_factor_stub, &lu_factor_kernel)
-
-REGISTER_ARCH_DISPATCH(ldl_factor_stub, DEFAULT, &ldl_factor_kernel)
-REGISTER_AVX512_DISPATCH(ldl_factor_stub, &ldl_factor_kernel)
-REGISTER_AVX2_DISPATCH(ldl_factor_stub, &ldl_factor_kernel)
-REGISTER_VSX_DISPATCH(ldl_factor_stub, &ldl_factor_kernel)
-REGISTER_ZVECTOR_DISPATCH(ldl_factor_stub, &ldl_factor_kernel)
-REGISTER_SVE256_DISPATCH(ldl_factor_stub, &ldl_factor_kernel)
-
-REGISTER_ARCH_DISPATCH(ldl_solve_stub, DEFAULT, &ldl_solve_kernel)
-REGISTER_AVX512_DISPATCH(ldl_solve_stub, &ldl_solve_kernel)
-REGISTER_AVX2_DISPATCH(ldl_solve_stub, &ldl_solve_kernel)
-REGISTER_VSX_DISPATCH(ldl_solve_stub, &ldl_solve_kernel)
-REGISTER_ZVECTOR_DISPATCH(ldl_solve_stub, &ldl_solve_kernel)
-REGISTER_SVE256_DISPATCH(ldl_solve_stub, &ldl_solve_kernel)
-
-REGISTER_ARCH_DISPATCH(lu_solve_stub, DEFAULT, &lu_solve_kernel)
-REGISTER_AVX512_DISPATCH(lu_solve_stub, &lu_solve_kernel)
-REGISTER_AVX2_DISPATCH(lu_solve_stub, &lu_solve_kernel)
-REGISTER_VSX_DISPATCH(lu_solve_stub, &lu_solve_kernel)
-REGISTER_ZVECTOR_DISPATCH(lu_solve_stub, &lu_solve_kernel)
-REGISTER_SVE256_DISPATCH(lu_solve_stub, &lu_solve_kernel)
-
-REGISTER_ARCH_DISPATCH(svd_stub, DEFAULT, &svd_kernel)
-REGISTER_AVX512_DISPATCH(svd_stub, &svd_kernel)
-REGISTER_AVX2_DISPATCH(svd_stub, &svd_kernel)
-REGISTER_VSX_DISPATCH(svd_stub, &svd_kernel)
-REGISTER_ZVECTOR_DISPATCH(svd_stub, &svd_kernel)
-REGISTER_SVE256_DISPATCH(svd_stub, &svd_kernel)
-
-REGISTER_ARCH_DISPATCH(unpack_pivots_stub, DEFAULT, &unpack_pivots_cpu_kernel)
-REGISTER_AVX512_DISPATCH(unpack_pivots_stub, &unpack_pivots_cpu_kernel)
-REGISTER_AVX2_DISPATCH(unpack_pivots_stub, &unpack_pivots_cpu_kernel)
-REGISTER_VSX_DISPATCH(unpack_pivots_stub, &unpack_pivots_cpu_kernel)
-REGISTER_ZVECTOR_DISPATCH(unpack_pivots_stub, &unpack_pivots_cpu_kernel)
-REGISTER_SVE256_DISPATCH(unpack_pivots_stub, &unpack_pivots_cpu_kernel)
+REGISTER_ALL_CPU_DISPATCH(cholesky_stub, &cholesky_kernel)
+REGISTER_ALL_CPU_DISPATCH(cholesky_inverse_stub, &cholesky_inverse_kernel_impl)
+REGISTER_ALL_CPU_DISPATCH(linalg_eig_make_complex_eigenvectors_stub, &linalg_eig_make_complex_eigenvectors_cpu)
+REGISTER_ALL_CPU_DISPATCH(linalg_eig_stub, &linalg_eig_kernel)
+REGISTER_ALL_CPU_DISPATCH(linalg_eigh_stub, &linalg_eigh_kernel)
+REGISTER_ALL_CPU_DISPATCH(geqrf_stub, &geqrf_kernel)
+REGISTER_ALL_CPU_DISPATCH(orgqr_stub, &orgqr_kernel_impl)
+REGISTER_ALL_CPU_DISPATCH(ormqr_stub, &ormqr_kernel)
+REGISTER_ALL_CPU_DISPATCH(lstsq_stub, &lstsq_kernel)
+REGISTER_ALL_CPU_DISPATCH(triangular_solve_stub, &triangular_solve_kernel)
+REGISTER_ALL_CPU_DISPATCH(lu_factor_stub, &lu_factor_kernel)
+REGISTER_ALL_CPU_DISPATCH(ldl_factor_stub, &ldl_factor_kernel)
+REGISTER_ALL_CPU_DISPATCH(ldl_solve_stub, &ldl_solve_kernel)
+REGISTER_ALL_CPU_DISPATCH(lu_solve_stub, &lu_solve_kernel)
+REGISTER_ALL_CPU_DISPATCH(svd_stub, &svd_kernel)
+REGISTER_ALL_CPU_DISPATCH(unpack_pivots_stub, &unpack_pivots_cpu_kernel)
 } // namespace at::native
diff --git a/aten/src/ATen/native/Blas.cpp b/aten/src/ATen/native/Blas.cpp
index a943af30741f4..35557b02027de 100644
--- a/aten/src/ATen/native/Blas.cpp
+++ b/aten/src/ATen/native/Blas.cpp
@@ -10,6 +10,7 @@
 #include <ATen/native/mkldnn/Linear.h>
 #include <ATen/native/Resize.h>
 #include <ATen/native/GroupedMMUtils.h>
+#include <ATen/BlasBackend.h>
 #if !defined(__s390x__) && !defined(__powerpc__)
 #include <cpuinfo.h>
 #endif
@@ -42,6 +43,9 @@ TORCH_META_FUNC(addmv)(const Tensor &self, const Tensor &mat, const Tensor &vec,
 
   TORCH_CHECK(mat.size(1) == vec.size(0) && (mat.size(0) == self.numel() || self.numel() == 1),
     "size mismatch, got input (", self.size(0), "), mat (", mat.size(0), "x", mat.size(1), "), vec (", vec.size(0), ")");
+
+  TORCH_CHECK(self.scalar_type() == mat.scalar_type() && mat.scalar_type() == vec.scalar_type(),
+    "addmv input tensors must have the same dtype, but got ", self.scalar_type(), ", ", mat.scalar_type(), ", and ", vec.scalar_type());
   auto names = at::namedinference::propagate_names_for_addmv(mat, vec, self);
   set_output_raw_strided(0, IntArrayRef(mat.sizes().data(), 1), {}, vec.options(), names);
 }
@@ -231,125 +235,4 @@ Tensor vdot(const Tensor &self, const Tensor &other){
 
 }
 
-static Tensor&
-_scaled_mm_out_cpu_emulated(const Tensor& mat1, const Tensor& mat2,
-          const Tensor& scale_a,
-          const Tensor& scale_b,
-          const std::optional<at::Tensor>& bias,
-          const std::optional<at::Tensor>& scale_result,
-          std::optional<c10::ScalarType> out_dtype,
-          bool use_fast_accum,
-          Tensor& out) {
-  TORCH_CHECK(mat1.dim() == 2, "mat1 must be a matrix");
-  TORCH_CHECK(mat2.dim() == 2, "mat2 must be a matrix");
-  TORCH_CHECK(
-      mat1.sizes()[1] == mat2.sizes()[0], "mat1 and mat2 shapes cannot be multiplied (",
-      mat1.sizes()[0], "x", mat1.sizes()[1], " and ", mat2.sizes()[0], "x", mat2.sizes()[1], ")");
-
-  TORCH_INTERNAL_ASSERT((scale_a.numel() == 1 && scale_b.numel() == 1), "Now _scaled_mm only supports per-tensor scaling for CPU backend.");
-  TORCH_CHECK(
-      !scale_result ||
-          (scale_result->numel() == 1 && scale_result->scalar_type() == kFloat),
-      "scale_result must be a float scalar");
-  TORCH_CHECK(!bias || bias->numel() == mat2.sizes()[1], "Bias must be size ", mat2.sizes()[1],
-       " but got ", bias->numel());
-
-  // Check types
-  TORCH_CHECK(!out_dtype || *out_dtype == out.scalar_type(), "out_dtype must match output matrix type");
-  TORCH_CHECK(isFloat8Type(mat1.scalar_type()), "Expected mat1 to be Float8 matrix got ", mat1.scalar_type());
-  TORCH_CHECK(isFloat8Type(mat2.scalar_type()), "Expected mat2 to be Float8 matrix got ", mat2.scalar_type());
-
-  auto mat1_c = mat1.contiguous();
-  auto mat2_c = mat2.contiguous();
-  IntArrayRef mat1_sizes = mat1_c.sizes();
-  IntArrayRef mat2_sizes = mat2_c.sizes();
-  at::native::resize_output(out, {mat1_sizes[0], mat2_sizes[1]});
-
-  float input_scale = scale_a.item<float>();
-  float weight_scale = scale_b.item<float>();
-  float output_scale = 1.0f;
-  if (scale_result.has_value() &&
-      (*out_dtype == ScalarType::Float8_e4m3fn ||
-       *out_dtype == ScalarType::Float8_e5m2)) {
-    output_scale = scale_result.value().item<float>();
-  }
-  auto fp32_mat1 = at::mul(mat1.to(kFloat), input_scale);
-  auto fp32_mat2 = at::mul(mat2_c.to(kFloat), weight_scale);
-  auto out_tmp = at::matmul(fp32_mat1, fp32_mat2);
-  if (bias) {
-    out_tmp.add_(bias.value());
-  }
-  if (*out_dtype == ScalarType::Float8_e4m3fn ||
-      *out_dtype == ScalarType::Float8_e5m2) {
-    out_tmp = at::mul(out_tmp, 1 / output_scale);
-  }
-  out_tmp = out_tmp.to(out.scalar_type());
-  out.copy_(out_tmp);
-  return out;
-}
-
-Tensor&
-_scaled_mm_out_cpu(const Tensor& mat1, const Tensor& mat2,
-          const Tensor& scale_a,
-          const Tensor& scale_b,
-          const std::optional<at::Tensor>& bias,
-          const std::optional<at::Tensor>& scale_result,
-          std::optional<c10::ScalarType> out_dtype,
-          bool use_fast_accum,
-          Tensor& out) {
-#if AT_MKLDNN_ENABLED() && !defined(__powerpc__)
-  if (at::globalContext().userEnabledMkldnn()) {
-    bool mixed_dtype = mat1.scalar_type() != mat2.scalar_type();
-    if ((!mixed_dtype && cpuinfo_has_x86_amx_int8()) ||
-        (mixed_dtype && cpuinfo_has_x86_amx_fp16())) {
-      return mkldnn_scaled_mm(
-          mat1,
-          mat2,
-          scale_a,
-          scale_b,
-          bias,
-          scale_result,
-          out_dtype,
-          use_fast_accum,
-          out);
-    }
-  }
-#endif
-  {
-  return _scaled_mm_out_cpu_emulated(mat1, mat2, scale_a, scale_b, bias, scale_result, out_dtype, use_fast_accum, out);
-  }
-}
-
-Tensor
-_scaled_mm_cpu(const Tensor& mat_a, const Tensor& mat_b,
-          const Tensor& scale_a,
-          const Tensor& scale_b,
-          const std::optional<at::Tensor>& bias,
-          const std::optional<at::Tensor>& scale_result,
-          std::optional<c10::ScalarType> out_dtype,
-          bool use_fast_accum) {
-  const auto out_dtype_ = out_dtype.value_or(mat_a.scalar_type());
-  Tensor out = at::empty({0}, mat_a.options().dtype(out_dtype_));
-  return _scaled_mm_out_cpu(mat_a, mat_b, scale_a, scale_b, bias, scale_result, out_dtype, use_fast_accum, out);
-}
-
-// TODO(vasiliy, future PR): figure out why we need to declare this function, when
-// other functions that live in ATen/native/*.cpp without declarations
-// or headers work just fine.
-Tensor _grouped_mm(const Tensor& mat_a, const Tensor& mat_b,
-const std::optional<at::Tensor>& offs,
-const std::optional<at::Tensor>& bias,
-std::optional<c10::ScalarType> out_dtype);
-
-Tensor _grouped_mm(const Tensor& mat_a, const Tensor& mat_b,
-const std::optional<at::Tensor>& offs,
-const std::optional<at::Tensor>& bias,
-std::optional<c10::ScalarType> out_dtype) {
-  _grouped_mm_validate_inputs(mat_a, mat_b, offs, bias, out_dtype);
-  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_);
-  _grouped_mm_fallback(mat_a, mat_b, offs, bias, out_dtype, out);
-  return out;
-}
-
 }  // namespace at::native
diff --git a/aten/src/ATen/native/CPUBlas.cpp b/aten/src/ATen/native/CPUBlas.cpp
index 87351cdb98717..164e709536135 100644
--- a/aten/src/ATen/native/CPUBlas.cpp
+++ b/aten/src/ATen/native/CPUBlas.cpp
@@ -2,7 +2,6 @@
 #include <ATen/native/CPUBlas.h>
 #include <ATen/native/mkl/LinearAlgebra.h>
 #include <ATen/native/mkldnn/Matmul.h>
-#include <ATen/Config.h>
 
 #include <c10/util/SmallBuffer.h>
 #include <c10/util/irange.h>
@@ -65,6 +64,11 @@ extern "C" void zaxpy_(int *n, void *a, const void *x, int *incx, void *y, int *
 #if ((defined(ONEDNN_UKERNEL_1) || defined(ONEDNN_UKERNEL_2)) && (defined(__x86_64__) || (defined(_M_X64) && !defined(_M_ARM64EC))))
 #define ONEDNN_UKERNEL_ENABLED
 #endif
+
+#if IDEEP_PREREQ(3, 9, 0, 0)
+#define ONEDNN_FP8_UKERNEL_AVAILABLE
+#endif
+
 #endif  // AT_MKLDNN_ENABLED()
 
 #if defined(ONEDNN_UKERNEL_ENABLED)
@@ -994,20 +998,19 @@ struct PackKey {
 };
 
 static inline dnnl::memory::data_type get_dnnl_dtype(ScalarType dtype) {
-  if (dtype == ScalarType::Float) {
-    return dnnl::memory::data_type::f32;
-  } else if (dtype == ScalarType::BFloat16) {
-    return dnnl::memory::data_type::bf16;
-  } else if (dtype == ScalarType::Half) {
-    return dnnl::memory::data_type::f16;
-  } else if (dtype == ScalarType::Int) {
-    return dnnl::memory::data_type::s32;
-  } else if (dtype == ScalarType::Byte) {
-    return dnnl::memory::data_type::u8;
-  } else if (dtype == ScalarType::Char) {
-    return dnnl::memory::data_type::s8;
-  } else {
-    TORCH_CHECK(false, "get_dnnl_dtype expects float/bfloat16/half/int8 tensor input");
+  using at::ScalarType;
+  using data_type = dnnl::memory::data_type;
+  switch (dtype) {
+    case ScalarType::Float:          return data_type::f32;
+    case ScalarType::BFloat16:       return data_type::bf16;
+    case ScalarType::Half:           return data_type::f16;
+    case ScalarType::Int:            return data_type::s32;
+    case ScalarType::Byte:           return data_type::u8;
+    case ScalarType::Char:           return data_type::s8;
+    case ScalarType::Float8_e4m3fn:  return data_type::f8_e4m3;
+    case ScalarType::Float8_e5m2:    return data_type::f8_e5m2;
+    default:
+      TORCH_CHECK(false, "Unsupported dtype for oneDNN");
   }
 }
 
@@ -1107,7 +1110,7 @@ struct GemmHelper {
     // Create a scratchpad buffer for the brgemm execution
     scratchpad = std::vector<uint8_t>(brg.get_scratchpad_size());
     // Prepare default vector of pairs of tensors A and B offsets for each batch.
-    A_B_offsets.reserve(1);
+    A_B_offsets.resize(1);
     A_B_offsets[0] = std::make_pair(0, 0);
   }
   dnnl::ukernel::brgemm brg;
@@ -1178,23 +1181,26 @@ struct Brgemm : public KernelCache <BrgemmKey, GemmHelper> {
     if (!at::globalContext().userEnabledMkldnn()) {
       return false;
     }
-    if (dtype == ScalarType::Half) {
-      static bool fp16_support = dnnl::get_effective_cpu_isa() >= dnnl::cpu_isa::avx512_core_fp16;
-      return fp16_support;
-    } else if (dtype == ScalarType::Float) {
-      static bool fp32_support = dnnl::get_effective_cpu_isa() >= dnnl::cpu_isa::avx2;
-      return fp32_support;
-    } else if (dtype == ScalarType::BFloat16) {
-      static bool bf16_support = dnnl::get_effective_cpu_isa() >= dnnl::cpu_isa::avx512_core;
-      return bf16_support;
-    } else if (dtype == ScalarType::Byte) {
-      static bool u8_support = dnnl::get_effective_cpu_isa() >= dnnl::cpu_isa::avx512_core_amx;
-      return u8_support;
-    } else if (dtype == ScalarType::Char) {
-      static bool s8_support = dnnl::get_effective_cpu_isa() >= dnnl::cpu_isa::avx512_core_vnni;
-      return s8_support;
+    static bool fp16_support = dnnl::get_effective_cpu_isa() >= dnnl::cpu_isa::avx512_core_fp16;
+    static bool fp32_support = dnnl::get_effective_cpu_isa() >= dnnl::cpu_isa::avx2;
+    static bool bf16_support = dnnl::get_effective_cpu_isa() >= dnnl::cpu_isa::avx512_core;
+    static bool u8_support = dnnl::get_effective_cpu_isa() >= dnnl::cpu_isa::avx512_core_amx;
+    static bool s8_support = dnnl::get_effective_cpu_isa() >= dnnl::cpu_isa::avx512_core_vnni;
+#ifdef ONEDNN_FP8_UKERNEL_AVAILABLE
+    static bool f8_support = dnnl::get_effective_cpu_isa() >= dnnl::cpu_isa::avx512_core_amx;
+#else
+    static bool f8_support = false;
+#endif
+    switch (dtype) {
+      case ScalarType::Half:          return fp16_support;
+      case ScalarType::Float:         return fp32_support;
+      case ScalarType::BFloat16:      return bf16_support;
+      case ScalarType::Byte:          return u8_support;
+      case ScalarType::Char:          return s8_support;
+      case ScalarType::Float8_e4m3fn: return f8_support;
+      case ScalarType::Float8_e5m2:   return f8_support;
+      default:                        return false;
     }
-    return false;
   }
 };
 
@@ -1237,17 +1243,23 @@ struct Pack : public KernelCache <PackKey, pack_t> {
     if (!at::globalContext().userEnabledMkldnn()) {
       return false;
     }
-    if (dtype == ScalarType::Half) {
-      static bool fp16_pack = dnnl::get_effective_cpu_isa() >= dnnl::cpu_isa::avx512_core_amx_fp16;
-      return fp16_pack;
-    } else if (dtype == ScalarType::BFloat16) {
-      static bool bf16_pack = dnnl::get_effective_cpu_isa() >= dnnl::cpu_isa::avx512_core_amx;
-      return bf16_pack;
-    } else if (dtype == ScalarType::Byte || dtype == ScalarType::Char) {
-      static bool bit8_pack = dnnl::get_effective_cpu_isa() >= dnnl::cpu_isa::avx512_core_amx;
-      return bit8_pack;
+    static bool fp16_pack = dnnl::get_effective_cpu_isa() >= dnnl::cpu_isa::avx512_core_amx_fp16;
+    static bool bf16_pack = dnnl::get_effective_cpu_isa() >= dnnl::cpu_isa::avx512_core_amx;
+    static bool bit8_pack = dnnl::get_effective_cpu_isa() >= dnnl::cpu_isa::avx512_core_amx;
+#ifdef ONEDNN_FP8_UKERNEL_AVAILABLE
+    static bool fp8_pack = dnnl::get_effective_cpu_isa() >= dnnl::cpu_isa::avx512_core_amx;
+#else
+    static bool fp8_pack = false;
+#endif
+    switch (dtype) {
+      case ScalarType::Half:          return fp16_pack;
+      case ScalarType::BFloat16:      return bf16_pack;
+      case ScalarType::Byte:          return bit8_pack;
+      case ScalarType::Char:          return bit8_pack;
+      case ScalarType::Float8_e4m3fn: return fp8_pack;
+      case ScalarType::Float8_e5m2:   return fp8_pack;
+      default:                        return false;
     }
-    return false;
   }
 };
 #endif
@@ -1419,6 +1431,54 @@ void brgemm(
     "I8 Brgemm is only supported on X64 when oneDNN ukernel is enabled and `amx` is supported");
 }
 
+void brgemm(
+    int64_t M,
+    int64_t N,
+    int64_t K,
+    int64_t ld_a,
+    int64_t ld_b,
+    int64_t ld_c,
+    const bool add_C,
+    const at::Float8_e4m3fn* A,
+    const at::Float8_e4m3fn* B,
+    float* C,
+    bool is_vnni) {
+#if defined(ONEDNN_UKERNEL_ENABLED) && defined(ONEDNN_FP8_UKERNEL_AVAILABLE)
+  if (is_vnni && Brgemm::device_check(ScalarType::Float8_e4m3fn)) {
+    Brgemm::call<at::Float8_e4m3fn, at::Float8_e4m3fn, float>(
+      M, N, K, ld_a, ld_b, ld_c, add_C, A, B, C);
+    return;
+  }
+#endif
+  // raise an error if the path is not supported
+  TORCH_CHECK(false,
+    "F8 Brgemm is only supported on X64 when oneDNN ukernel is enabled and `amx` is supported");
+}
+
+void brgemm(
+    int64_t M,
+    int64_t N,
+    int64_t K,
+    int64_t ld_a,
+    int64_t ld_b,
+    int64_t ld_c,
+    const bool add_C,
+    const at::Float8_e5m2* A,
+    const at::Float8_e5m2* B,
+    float* C,
+    bool is_vnni) {
+#if defined(ONEDNN_UKERNEL_ENABLED) && defined(ONEDNN_FP8_UKERNEL_AVAILABLE)
+  if (is_vnni && Brgemm::device_check(ScalarType::Float8_e5m2)) {
+    Brgemm::call<at::Float8_e5m2, at::Float8_e5m2, float>(
+      M, N, K, ld_a, ld_b, ld_c, add_C, A, B, C);
+    return;
+  }
+#endif
+  // raise an error if the path is not supported
+  TORCH_CHECK(false,
+    "F8 Brgemm is only supported on X64 when oneDNN ukernel is enabled and `amx` is supported");
+}
+
 void brgemm_release(bool is_vnni) {
 #if defined(ONEDNN_UKERNEL_ENABLED)
   if (is_vnni) {
diff --git a/aten/src/ATen/native/CPUBlas.h b/aten/src/ATen/native/CPUBlas.h
index 8b75f12ebaf21..caf166fff6aa2 100644
--- a/aten/src/ATen/native/CPUBlas.h
+++ b/aten/src/ATen/native/CPUBlas.h
@@ -6,6 +6,7 @@
 #include <c10/util/complex.h>
 #include <c10/core/ScalarType.h>
 #include <c10/core/Scalar.h>
+#include <ATen/Config.h>
 
 
 namespace at::native::cpublas {
@@ -215,6 +216,7 @@ void copy(int64_t n, const c10::complex<float> *x, int64_t incx, c10::complex<fl
 #define CPUBLAS_BRGEMM_U8U8I32 // unsigned char * unsigned char -> int32
 #define CPUBLAS_BRGEMM_U8I8I32 // unsigned char * signed char -> int32
 #define CPUBLAS_BRGEMM_I8I8I32 // signed char * signed char -> int32
+#define CPUBLAS_BRGEMM_F8F8F32 // float8 * float8 -> float (e4m3 & e5m2)
 
 TORCH_API void brgemm(
     int64_t M,
@@ -294,6 +296,34 @@ TORCH_API void brgemm(
     int32_t* C,
     bool is_vnni = true);
 
+#ifdef CPUBLAS_BRGEMM_F8F8F32
+TORCH_API void brgemm(
+    int64_t M,
+    int64_t N,
+    int64_t K,
+    int64_t ld_a,
+    int64_t ld_b,
+    int64_t ld_c,
+    const bool add_C,
+    const at::Float8_e4m3fn* A,
+    const at::Float8_e4m3fn* B,
+    float* C,
+    bool is_vnni = true);
+
+TORCH_API void brgemm(
+    int64_t M,
+    int64_t N,
+    int64_t K,
+    int64_t ld_a,
+    int64_t ld_b,
+    int64_t ld_c,
+    const bool add_C,
+    const at::Float8_e5m2* A,
+    const at::Float8_e5m2* B,
+    float* C,
+    bool is_vnni = true);
+#endif
+
 // Release brgemm hardware context
 TORCH_API void brgemm_release(bool is_vnni = true);
 
diff --git a/aten/src/ATen/native/ConvUtils.h b/aten/src/ATen/native/ConvUtils.h
index 8c0771ddfdcc6..c96e0522b67d0 100644
--- a/aten/src/ATen/native/ConvUtils.h
+++ b/aten/src/ATen/native/ConvUtils.h
@@ -471,8 +471,14 @@ inline bool mps_conv_use_channels_last(const at::Tensor& input, const at::Tensor
     return false;
   }
 
+  // Use exact-match so a tensor whose strides merely *look* like channels-last
+  // (e.g. a channel-slice of a channels-last tensor) is not misclassified --
+  // MPS reads raw buffers assuming packed NHWC, which would be incorrect for
+  // such views. exact_match also correctly excludes degenerate 1x1 weights
+  // whose NCHW-contiguous strides happen to also satisfy is_contiguous(CL).
+  // See https://github.com/pytorch/pytorch/issues/180984
   auto is_channel_last = [](const at::Tensor& t) {
-    auto fmt = t.suggest_memory_format();
+    auto fmt = t.suggest_memory_format(/*channels_last_strides_exact_match=*/true);
     return fmt == at::MemoryFormat::ChannelsLast || fmt == at::MemoryFormat::ChannelsLast3d;
   };
   return is_channel_last(input) || is_channel_last(weight);
diff --git a/aten/src/ATen/native/Convolution.cpp b/aten/src/ATen/native/Convolution.cpp
index 5f14d08df5c62..51d83ba16779e 100644
--- a/aten/src/ATen/native/Convolution.cpp
+++ b/aten/src/ATen/native/Convolution.cpp
@@ -11,6 +11,7 @@
 #include <ATen/native/utils/ParamUtils.h>
 #include <ATen/native/xnnpack/Engine.h>
 #include <c10/core/GradMode.h>
+#include <c10/core/SymBool.h>
 #include <c10/util/accumulate.h>
 #include <c10/util/irange.h>
 #include <c10/macros/Macros.h>
@@ -509,12 +510,35 @@ struct ConvParams {
        input.scalar_type() == kFloat && // only on CPU Float Tensors
        // For 1x1 filters, MKLDNN is faster than THNN when multi-threaded,
        // but THNN is faster when single-threaded.
-       (is_strided() || is_dilated() || at::symint::size<T>(input, 0) >= 16 ||
-        at::symint::size<T>(weight, -1) != 1 || at::symint::size<T>(weight, -2) != 1 || at::get_num_threads() > 1) &&
-       (groups > 1
-        || (at::symint::size<T>(weight, -1) > 3 && at::symint::size<T>(weight, -2) > 3)
-        || at::symint::size<T>(input, 0) > 1
-        || at::symint::size<T>(input, 0)*at::symint::size<T>(input, 1)*at::symint::size<T>(input, 2)*at::symint::size<T>(input, 3) > 20480) // for some case, native is faster
+       [&]() {
+         if constexpr (std::is_same_v<T, c10::SymInt>) {
+           return is_strided() || is_dilated() ||
+             TORCH_GUARD_OR_FALSE(at::symint::size<T>(input, 0).sym_ge(16)) ||
+             TORCH_GUARD_OR_FALSE(at::symint::size<T>(weight, -1).sym_ne(1)) ||
+             TORCH_GUARD_OR_FALSE(at::symint::size<T>(weight, -2).sym_ne(1)) ||
+             at::get_num_threads() > 1;
+         } else {
+           return is_strided() || is_dilated() ||
+             at::symint::size<T>(input, 0) >= 16 ||
+             at::symint::size<T>(weight, -1) != 1 ||
+             at::symint::size<T>(weight, -2) != 1 ||
+             at::get_num_threads() > 1;
+         }
+       }() &&
+       [&]() {
+         if constexpr (std::is_same_v<T, c10::SymInt>) {
+           return groups > 1
+             || (TORCH_GUARD_OR_FALSE(at::symint::size<T>(weight, -1).sym_gt(3)) &&
+                 TORCH_GUARD_OR_FALSE(at::symint::size<T>(weight, -2).sym_gt(3)))
+             || TORCH_GUARD_OR_FALSE(at::symint::size<T>(input, 0).sym_gt(1))
+             || TORCH_GUARD_OR_FALSE((at::symint::size<T>(input, 0)*at::symint::size<T>(input, 1)*at::symint::size<T>(input, 2)*at::symint::size<T>(input, 3)).sym_gt(20480));
+         } else {
+           return groups > 1
+             || (at::symint::size<T>(weight, -1) > 3 && at::symint::size<T>(weight, -2) > 3)
+             || at::symint::size<T>(input, 0) > 1
+             || at::symint::size<T>(input, 0)*at::symint::size<T>(input, 1)*at::symint::size<T>(input, 2)*at::symint::size<T>(input, 3) > 20480;
+         }
+       }() // for some case, native is faster
         );
 
 #endif
@@ -651,29 +675,53 @@ static void check_shape_forward(const at::Tensor& input,
            "Expected ", weight_dim, "-dimensional input for ", weight_dim,
            "-dimensional weight ", weight_sizes, ", but got ", k, "-dimensional input of size ",
            at::symint::sizes<T>(input), " instead");
-  TORCH_CHECK(weight_sizes[0] >= groups,
-           "Given groups=", groups, ", expected weight to be at least ", groups,
-           " at dimension 0, but got weight of size ", weight_sizes, " instead");
-  TORCH_CHECK(weight_sizes[0] % groups == 0,
-           "Given groups=", groups, ", expected weight to be divisible by ",
-           groups, " at dimension 0, but got weight of size [", weight_sizes,
-           "] instead");
+  if constexpr (std::is_same_v<T, c10::SymInt>) {
+    TORCH_SYM_CHECK(weight_sizes[0].sym_ge(groups),
+             "Given groups=", groups, ", expected weight to be at least ", groups,
+             " at dimension 0, but got weight of size ", weight_sizes, " instead");
+    TORCH_SYM_CHECK((weight_sizes[0] % groups).sym_eq(0),
+             "Given groups=", groups, ", expected weight to be divisible by ",
+             groups, " at dimension 0, but got weight of size [", weight_sizes,
+             "] instead");
+  } else {
+    TORCH_CHECK(weight_sizes[0] >= groups,
+             "Given groups=", groups, ", expected weight to be at least ", groups,
+             " at dimension 0, but got weight of size ", weight_sizes, " instead");
+    TORCH_CHECK(weight_sizes[0] % groups == 0,
+             "Given groups=", groups, ", expected weight to be divisible by ",
+             groups, " at dimension 0, but got weight of size [", weight_sizes,
+             "] instead");
+  }
 
   if (!transposed) {
     std::vector<T> input_shape;
     std::vector<T> kernel_shape;
     bool kernel_size_correct = true;
 
-    TORCH_CHECK(at::symint::size<T>(input, 1) == (weight_sizes[1] * groups),
-                "Given groups=", groups, ", weight of size ", weight_sizes,
-                ", expected input", at::symint::sizes<T>(input), " to have ",
-                (weight_sizes[1] * groups), " channels, but got ", at::symint::size<T>(input, 1),
-                " channels instead");
-
-    TORCH_CHECK(!bias.defined() || (bias.ndimension() == 1 && at::symint::size<T>(bias, 0) == weight_sizes[0]),
-             "Given weight of size ", weight_sizes,
-             ", expected bias to be 1-dimensional with ", weight_sizes[0], " elements",
-             ", but got bias of size ", at::symint::sizes<T>(bias), " instead");
+    if constexpr (std::is_same_v<T, c10::SymInt>) {
+      TORCH_SYM_CHECK(at::symint::size<T>(input, 1).sym_eq(weight_sizes[1] * groups),
+                  "Given groups=", groups, ", weight of size ", weight_sizes,
+                  ", expected input", at::symint::sizes<T>(input), " to have ",
+                  (weight_sizes[1] * groups), " channels, but got ", at::symint::size<T>(input, 1),
+                  " channels instead");
+
+      TORCH_CHECK(!bias.defined() || (bias.ndimension() == 1 &&
+                  at::symint::size<T>(bias, 0).sym_eq(weight_sizes[0]).expect_true(__FILE__, __LINE__)),
+               "Given weight of size ", weight_sizes,
+               ", expected bias to be 1-dimensional with ", weight_sizes[0], " elements",
+               ", but got bias of size ", at::symint::sizes<T>(bias), " instead");
+    } else {
+      TORCH_CHECK(at::symint::size<T>(input, 1) == (weight_sizes[1] * groups),
+                  "Given groups=", groups, ", weight of size ", weight_sizes,
+                  ", expected input", at::symint::sizes<T>(input), " to have ",
+                  (weight_sizes[1] * groups), " channels, but got ", at::symint::size<T>(input, 1),
+                  " channels instead");
+
+      TORCH_CHECK(!bias.defined() || (bias.ndimension() == 1 && at::symint::size<T>(bias, 0) == weight_sizes[0]),
+               "Given weight of size ", weight_sizes,
+               ", expected bias to be 1-dimensional with ", weight_sizes[0], " elements",
+               ", but got bias of size ", at::symint::sizes<T>(bias), " instead");
+    }
 
     for (const auto i : c10::irange(2, k)) {
       // T could be int64_t or SymInt, Specialized numeric_limts<SymInt> in c10/core/SymInt.h
@@ -683,8 +731,14 @@ static void check_shape_forward(const at::Tensor& input,
       input_shape.push_back(at::symint::size<T>(input, i) + 2 * padding[i-2]);
       // log new kernel size considering dilation
       kernel_shape.push_back(dilation[i-2] * (weight_sizes[i]-1) + 1);
-      if (input_shape.back() < kernel_shape.back()) {
-        kernel_size_correct = false;
+      if constexpr (std::is_same_v<T, c10::SymInt>) {
+        if (TORCH_GUARD_OR_FALSE(input_shape.back().sym_lt(kernel_shape.back()))) {
+          kernel_size_correct = false;
+        }
+      } else {
+        if (input_shape.back() < kernel_shape.back()) {
+          kernel_size_correct = false;
+        }
       }
     }
 
@@ -711,14 +765,26 @@ static void check_shape_forward(const at::Tensor& input,
                   "Given padding=", padding[i-2], " at dimension ", i-2, " , expected padding to be at most ",
                   (std::numeric_limits<T>::max() / 2));
     }
-    TORCH_CHECK(at::symint::size<T>(input, 1) == weight_sizes[0],
-             "Given transposed=", transposed, ", weight of size ", weight_sizes,
-             ", expected input", at::symint::sizes<T>(input), " to have ", weight_sizes[0],
-             " channels, but got ", at::symint::size<T>(input, 1), " channels instead");
-    TORCH_CHECK(!bias.defined() || (bias.ndimension() == 1 && at::symint::size<T>(bias, 0) == weight_sizes[1] * groups),
-             "Given transposed=", transposed, ", weight of size ", weight_sizes,
-             ", expected bias to be 1-dimensional with ", weight_sizes[1] * groups, " elements",
-             ", but got bias of size ", at::symint::sizes<T>(bias), " instead");
+    if constexpr (std::is_same_v<T, c10::SymInt>) {
+      TORCH_SYM_CHECK(at::symint::size<T>(input, 1).sym_eq(weight_sizes[0]),
+               "Given transposed=", transposed, ", weight of size ", weight_sizes,
+               ", expected input", at::symint::sizes<T>(input), " to have ", weight_sizes[0],
+               " channels, but got ", at::symint::size<T>(input, 1), " channels instead");
+      TORCH_CHECK(!bias.defined() || (bias.ndimension() == 1 &&
+                  at::symint::size<T>(bias, 0).sym_eq(weight_sizes[1] * groups).expect_true(__FILE__, __LINE__)),
+               "Given transposed=", transposed, ", weight of size ", weight_sizes,
+               ", expected bias to be 1-dimensional with ", weight_sizes[1] * groups, " elements",
+               ", but got bias of size ", at::symint::sizes<T>(bias), " instead");
+    } else {
+      TORCH_CHECK(at::symint::size<T>(input, 1) == weight_sizes[0],
+               "Given transposed=", transposed, ", weight of size ", weight_sizes,
+               ", expected input", at::symint::sizes<T>(input), " to have ", weight_sizes[0],
+               " channels, but got ", at::symint::size<T>(input, 1), " channels instead");
+      TORCH_CHECK(!bias.defined() || (bias.ndimension() == 1 && at::symint::size<T>(bias, 0) == weight_sizes[1] * groups),
+               "Given transposed=", transposed, ", weight of size ", weight_sizes,
+               ", expected bias to be 1-dimensional with ", weight_sizes[1] * groups, " elements",
+               ", but got bias of size ", at::symint::sizes<T>(bias), " instead");
+    }
   }
 }
 
@@ -1009,7 +1075,7 @@ static Tensor convolution_same(
         input_sizes[i + 2], weight_sizes[i + 2], s, d);
     padding_l.push_back(pad.first);
     padding_r.push_back(pad.second);
-    if (pad.first != pad.second) {
+    if (!TORCH_GUARD_OR_FALSE(pad.first.sym_eq(pad.second))) {
       symmetric_padding = false;
     }
   }
@@ -1025,15 +1091,10 @@ static Tensor convolution_same(
                   " require a zero-padded copy of the input be created");
   SmallVector<c10::SymInt, kDimVectorStaticSize * 2> pad_nd(static_cast<size_t>(2 * dim));
   for (auto i: c10::irange(dim)) {
-    // Apply padding by the difference, leaving only a symmetric padding
-    auto delta_pad = padding_r[i] - padding_l[i];
+    // pad_r >= pad_l always (floor division), so pad the right side by the
+    // difference and use pad_l as the symmetric base for convolution.
     auto pad_idx = 2 * (dim - 1 - i);  // F.pad goes from last dim to first
-    if (delta_pad > 0) {
-      pad_nd[pad_idx + 1] = delta_pad;
-    } else {
-      pad_nd[pad_idx] = delta_pad;
-      padding_l[i] = padding_r[i];
-    }
+    pad_nd[pad_idx + 1] = padding_r[i] - padding_l[i];
   }
   auto padded_input = at::constant_pad_nd_symint(input, pad_nd, 0);
   SymDimVector output_padding(dim);
@@ -1196,10 +1257,19 @@ static ConvBackend _select_conv_backend(
     const ConvParams<T>& params) {
 
   // don't send empty inputs through backends
-  if (at::symint::size<T>(input, 0) == 0 || at::symint::size<T>(input, 1) == 0) {
-    return input.is_mkldnn() ? ConvBackend::MkldnnEmpty : ConvBackend::Empty;
-  } else if (at::symint::numel<T>(input) == 0) {
-    TORCH_CHECK(false, "Only zero batch or zero channel inputs are supported, but got input shape: ", at::symint::sizes<T>(input));
+  if constexpr (std::is_same_v<T, c10::SymInt>) {
+    if (TORCH_GUARD_OR_FALSE(at::symint::size<T>(input, 0).sym_eq(0)) ||
+        TORCH_GUARD_OR_FALSE(at::symint::size<T>(input, 1).sym_eq(0))) {
+      return input.is_mkldnn() ? ConvBackend::MkldnnEmpty : ConvBackend::Empty;
+    } else if (TORCH_GUARD_OR_FALSE(at::symint::numel<T>(input).sym_eq(0))) {
+      TORCH_CHECK(false, "Only zero batch or zero channel inputs are supported, but got input shape: ", at::symint::sizes<T>(input));
+    }
+  } else {
+    if (at::symint::size<T>(input, 0) == 0 || at::symint::size<T>(input, 1) == 0) {
+      return input.is_mkldnn() ? ConvBackend::MkldnnEmpty : ConvBackend::Empty;
+    } else if (at::symint::numel<T>(input) == 0) {
+      TORCH_CHECK(false, "Only zero batch or zero channel inputs are supported, but got input shape: ", at::symint::sizes<T>(input));
+    }
   }
 
   if (params.is_depthwise(input, weight)) {
@@ -1704,10 +1774,6 @@ std::tuple<Tensor, Tensor, Tensor> convolution_backward_overrideable(
         IntArrayRef stride, IntArrayRef padding, IntArrayRef dilation,
         bool transposed, IntArrayRef output_padding, int64_t groups, std::array<bool, 3> output_mask) {
    TORCH_CHECK_NOT_IMPLEMENTED(false, "convolution_backward_overrideable: You are likely triggering this with tensor backend other than CPU/CUDA/MKLDNN, if this is intended, please use TORCH_LIBRARY_IMPL to override this function ");
-  return std::tuple<Tensor, Tensor, Tensor>(
-          at::empty_like(input, LEGACY_CONTIGUOUS_MEMORY_FORMAT),
-          at::empty_like(weight, LEGACY_CONTIGUOUS_MEMORY_FORMAT),
-          at::empty({}));
 }
 
 static Tensor subvariable(const Tensor& var, int64_t dim, int64_t groups, int64_t g) {
@@ -1910,7 +1976,7 @@ std::tuple<Tensor,Tensor,Tensor> _convolution_double_backward( const std::option
     }
   }
 
-  return std::tuple<Tensor,Tensor,Tensor>{ggO, gI, gW};
+  return std::tuple<Tensor,Tensor,Tensor>{std::move(ggO), std::move(gI), std::move(gW)};
 }
 
 static std::tuple<at::Tensor, at::Tensor, at::Tensor> _convolution_backward_nogroup_backend(
@@ -2255,7 +2321,7 @@ std::tuple<Tensor, Tensor, Tensor> convolution_backward(
     }
   }
 
-  return std::make_tuple(backend_grad_input, backend_grad_weight, backend_grad_bias);
+  return std::make_tuple(std::move(backend_grad_input), std::move(backend_grad_weight), std::move(backend_grad_bias));
 }
 
 void _cudnn_set_conv_benchmark_empty_cache(bool enable) {
diff --git a/aten/src/ATen/native/DispatchStub.cpp b/aten/src/ATen/native/DispatchStub.cpp
index 12ad011493210..cff843e0ee5c8 100644
--- a/aten/src/ATen/native/DispatchStub.cpp
+++ b/aten/src/ATen/native/DispatchStub.cpp
@@ -39,14 +39,12 @@ static CPUCapability compute_cpu_capability() {
     }
 #elif defined(HAVE_SVE_CPU_DEFINITION)
     int sve_vl = cpuinfo_get_max_arm_sve_length(); //Returns maximum SVE VL supported by your HW.
-#ifdef HAVE_SVE256_CPU_DEFINITION
+#ifdef HAVE_SVE_CPU_DEFINITION
     if (envar == "sve256") {
       if (sve_vl == 256) {
-#ifdef HAVE_ARM_BF16_CPU_DEFINITION
         if (cpuinfo_has_arm_bf16()) {
           return CPUCapability::SVE256;
         }
-#endif
       }
       TORCH_WARN("SVE256 capability not available on hardware. Falling back to DEFAULT");
       return CPUCapability::DEFAULT;
@@ -105,13 +103,9 @@ static CPUCapability compute_cpu_capability() {
       // Return the default CPU capability.
       return CPUCapability::DEFAULT;
     }
-    #ifdef HAVE_SVE256_CPU_DEFINITION
-        if (sve_vl == 256) { // Check for SVE256
-        #ifdef HAVE_ARM_BF16_CPU_DEFINITION
-          if (cpuinfo_has_arm_bf16())
+    #ifdef HAVE_SVE_CPU_DEFINITION
+        if (sve_vl == 256 && cpuinfo_has_arm_bf16())
             return CPUCapability::SVE256;
-        #endif
-        }
     #endif
     // Return the default CPU capability.
     return CPUCapability::DEFAULT;
@@ -144,7 +138,7 @@ DispatchResult DispatchStubImpl::try_get_call_ptr(
 #ifdef HAVE_ZVECTOR_CPU_DEFINITION
   , void *ZVECTOR
 #endif
-#ifdef HAVE_SVE256_CPU_DEFINITION
+#ifdef HAVE_SVE_CPU_DEFINITION
   , void *SVE256
 #endif
 ) {
@@ -182,7 +176,7 @@ DispatchResult DispatchStubImpl::try_get_call_ptr(
 #ifdef HAVE_ZVECTOR_CPU_DEFINITION
           , ZVECTOR
 #endif
-#ifdef HAVE_SVE256_CPU_DEFINITION
+#ifdef HAVE_SVE_CPU_DEFINITION
           , SVE256
 #endif
         );
@@ -220,7 +214,6 @@ DispatchResult DispatchStubImpl::try_get_call_ptr(
 
     default:
       TORCH_INTERNAL_ASSERT(false, "An unexpected device type was provided ", device_type);
-      return ErrorType::DeviceNotSupported;
   }
 }
 
@@ -239,7 +232,7 @@ void* DispatchStubImpl::get_call_ptr(
 #ifdef HAVE_ZVECTOR_CPU_DEFINITION
   , void *ZVECTOR
 #endif
-#ifdef HAVE_SVE256_CPU_DEFINITION
+#ifdef HAVE_SVE_CPU_DEFINITION
   , void *SVE256
 #endif
 ) {
@@ -263,7 +256,7 @@ void* DispatchStubImpl::get_call_ptr(
       ,
       ZVECTOR
 #endif
-#ifdef HAVE_SVE256_CPU_DEFINITION
+#ifdef HAVE_SVE_CPU_DEFINITION
       ,
       SVE256
 #endif
@@ -274,7 +267,6 @@ void* DispatchStubImpl::get_call_ptr(
       case ErrorType::MissingDeviceKernel:
         TORCH_INTERNAL_ASSERT(
             false, "DispatchStub: missing kernel for ", device_type);
-        return nullptr;
       case ErrorType::DeviceNotSupported:
         TORCH_CHECK(false, "DispatchStub: unsupported device type", device_type);
     }
@@ -298,7 +290,7 @@ DispatchResult DispatchStubImpl::try_choose_cpu_impl(
 #ifdef HAVE_ZVECTOR_CPU_DEFINITION
     , void *ZVECTOR
 #endif
-#ifdef HAVE_SVE256_CPU_DEFINITION
+#ifdef HAVE_SVE_CPU_DEFINITION
     , void *SVE256
 #endif
   ){
@@ -333,7 +325,7 @@ DispatchResult DispatchStubImpl::try_choose_cpu_impl(
     return ZVECTOR != nullptr ? DispatchResult(ZVECTOR) : ErrorType::MissingDeviceKernel;
   }
 #endif
-#ifdef HAVE_SVE256_CPU_DEFINITION
+#ifdef HAVE_SVE_CPU_DEFINITION
   if (capability >= static_cast<int>(CPUCapability::SVE256)) {
     if (C10_UNLIKELY(!SVE256)) {
       // dispatch to DEFAULT, since the SVE kernel is missing
@@ -360,7 +352,7 @@ void* DispatchStubImpl::choose_cpu_impl(
 #ifdef HAVE_ZVECTOR_CPU_DEFINITION
   , void *ZVECTOR
 #endif
-#ifdef HAVE_SVE256_CPU_DEFINITION
+#ifdef HAVE_SVE_CPU_DEFINITION
   , void *SVE256
 #endif
 ) {
@@ -398,7 +390,7 @@ void* DispatchStubImpl::choose_cpu_impl(
     return ZVECTOR;
   }
 #endif
-#ifdef HAVE_SVE256_CPU_DEFINITION
+#ifdef HAVE_SVE_CPU_DEFINITION
   if (capability >= static_cast<int>(CPUCapability::SVE256)) {
     if (C10_UNLIKELY(!SVE256)) {
       // dispatch to DEFAULT, since the SVE kernel is missing
diff --git a/aten/src/ATen/native/DispatchStub.h b/aten/src/ATen/native/DispatchStub.h
index 4c1e403973d15..1c636453c8006 100644
--- a/aten/src/ATen/native/DispatchStub.h
+++ b/aten/src/ATen/native/DispatchStub.h
@@ -64,7 +64,7 @@ enum class CPUCapability {
   VSX = 1,
 #elif defined(HAVE_ZVECTOR_CPU_DEFINITION)
   ZVECTOR = 1,
-#elif defined(HAVE_SVE256_CPU_DEFINITION) && defined(HAVE_ARM_BF16_CPU_DEFINITION)
+#elif defined(HAVE_SVE_CPU_DEFINITION)
   SVE256 = 1,
 #else
   AVX2 = 1,
@@ -115,7 +115,7 @@ struct TORCH_API DispatchStubImpl {
 #ifdef HAVE_ZVECTOR_CPU_DEFINITION
       , void *ZVECTOR
 #endif
-#ifdef HAVE_SVE256_CPU_DEFINITION
+#ifdef HAVE_SVE_CPU_DEFINITION
       , void *SVE256
 #endif
   );
@@ -136,7 +136,7 @@ struct TORCH_API DispatchStubImpl {
 #ifdef HAVE_ZVECTOR_CPU_DEFINITION
     , void *ZVECTOR
 #endif
-#ifdef HAVE_SVE256_CPU_DEFINITION
+#ifdef HAVE_SVE_CPU_DEFINITION
     , void *SVE256
 #endif
   );
@@ -157,7 +157,7 @@ struct TORCH_API DispatchStubImpl {
 #ifdef HAVE_ZVECTOR_CPU_DEFINITION
       , void *ZVECTOR
 #endif
-#ifdef HAVE_SVE256_CPU_DEFINITION
+#ifdef HAVE_SVE_CPU_DEFINITION
       , void *SVE256
 #endif
   );
@@ -181,7 +181,7 @@ struct TORCH_API DispatchStubImpl {
 #ifdef HAVE_ZVECTOR_CPU_DEFINITION
     , void *ZVECTOR
 #endif
-#ifdef HAVE_SVE256_CPU_DEFINITION
+#ifdef HAVE_SVE_CPU_DEFINITION
     , void *SVE256
 #endif
   );
@@ -238,7 +238,7 @@ struct DispatchStub<rT (*)(Args...), T> {
 #ifdef HAVE_ZVECTOR_CPU_DEFINITION
       , reinterpret_cast<void*>(ZVECTOR)
 #endif
-#ifdef HAVE_SVE256_CPU_DEFINITION
+#ifdef HAVE_SVE_CPU_DEFINITION
       , reinterpret_cast<void*>(SVE256)
 #endif
       )
@@ -299,7 +299,7 @@ struct DispatchStub<rT (*)(Args...), T> {
 #ifdef HAVE_ZVECTOR_CPU_DEFINITION
       , reinterpret_cast<void*>(ZVECTOR)
 #endif
-#ifdef HAVE_SVE256_CPU_DEFINITION
+#ifdef HAVE_SVE_CPU_DEFINITION
       , reinterpret_cast<void*>(SVE256)
 #endif
       );
@@ -322,7 +322,7 @@ struct DispatchStub<rT (*)(Args...), T> {
 #ifdef HAVE_ZVECTOR_CPU_DEFINITION
   static TORCH_API FnPtr ZVECTOR;
 #endif
-#ifdef HAVE_SVE256_CPU_DEFINITION
+#ifdef HAVE_SVE_CPU_DEFINITION
   static TORCH_API FnPtr SVE256;
 #endif
 private:
@@ -426,7 +426,7 @@ struct RegisterPRIVATEUSE1Dispatch {
 #define REGISTER_ZVECTOR_DISPATCH(name, fn)
 #endif
 
-#ifdef HAVE_SVE256_CPU_DEFINITION
+#ifdef HAVE_SVE_CPU_DEFINITION
 #define REGISTER_SVE256_DISPATCH(name, fn) REGISTER_ARCH_DISPATCH(name, SVE256, fn)
 #else
 #define REGISTER_SVE256_DISPATCH(name, fn)
diff --git a/aten/src/ATen/native/Distance.cpp b/aten/src/ATen/native/Distance.cpp
index b7dfef7779ca7..af8d6185d4c26 100644
--- a/aten/src/ATen/native/Distance.cpp
+++ b/aten/src/ATen/native/Distance.cpp
@@ -103,8 +103,8 @@ static Tensor cdist_impl(const Tensor& x1, const Tensor& x2, const double p, std
   // See Note [cdist relies on cdist_impl redispatching]
   // Keep this condition in sync with the condition at the Note
   if (!(p == 2 && (mode == 1 || (mode == 0 && (r1 > 25 || r2 > 25))))) {
-    TORCH_CHECK(device1 == kCPU || device1 == kCUDA || device1 == kXPU, "cdist only supports CPU, XPU and CUDA devices, X1 got: ", device1);
-    TORCH_CHECK(device2 == kCPU || device2 == kCUDA || device2 == kXPU, "cdist only supports CPU, XPU and CUDA devices, X2 got: ", device2);
+    TORCH_CHECK(device1 == kCPU || device1 == kCUDA || device1 == kXPU || device1 == kPrivateUse1, "cdist only supports CPU, XPU, CUDA and PrivateUse1 devices, X1 got: ", device1);
+    TORCH_CHECK(device2 == kCPU || device2 == kCUDA || device2 == kXPU || device2 == kPrivateUse1, "cdist only supports CPU, XPU, CUDA and PrivateUse1 devices, X2 got: ", device2);
   }
 
   auto dim1 = x1.dim();
@@ -229,9 +229,9 @@ Tensor _cdist_backward(const Tensor& _grad, const Tensor& _x1, const Tensor& _x2
   int64_t n = x1.size(-2);
   int64_t m = x1.size(-1);
   auto device1 = x1.device().type();
-  TORCH_CHECK(device1 == kCPU || device1 == kCUDA || device1 == kXPU, "_cdist_backward only supports CPU, XPU and CUDA devices, X1 got: ", device1);
+  TORCH_CHECK(device1 == kCPU || device1 == kCUDA || device1 == kXPU || device1 == kPrivateUse1, "_cdist_backward only supports CPU, XPU, CUDA and PrivateUse1 devices, X1 got: ", device1);
   auto device2 = x2.device().type();
-  TORCH_CHECK(device2 == kCPU || device2 == kCUDA || device2 == kXPU, "_cdist_backward only supports CPU, XPU and CUDA devices, X2 got: ", device2);
+  TORCH_CHECK(device2 == kCPU || device2 == kCUDA || device2 == kXPU || device2 == kPrivateUse1, "_cdist_backward only supports CPU, XPU, CUDA and PrivateUse1 devices, X2 got: ", device2);
 
   Tensor grad_x1 =
       at::empty({batch_product, n, m}, x1.options(), LEGACY_CONTIGUOUS_MEMORY_FORMAT);
@@ -245,7 +245,7 @@ Tensor _cdist_backward(const Tensor& _grad, const Tensor& _x1, const Tensor& _x2
 Tensor _pdist_forward(const Tensor& self, const double p) {
   TORCH_CHECK(self.is_contiguous(), "_pdist_forward requires contiguous input");
   auto device = self.device().type();
-  TORCH_CHECK(device == kCPU || device == kCUDA || device == kXPU, "_pdist_forward only supports CPU, XPU and CUDA devices, got: ", device);
+  TORCH_CHECK(device == kCPU || device == kCUDA || device == kXPU || device == kPrivateUse1, "_pdist_forward only supports CPU, XPU, CUDA and PrivateUse1 devices, got: ", device);
   Tensor result = at::empty({0}, self.options(), LEGACY_CONTIGUOUS_MEMORY_FORMAT);
   if (self.size(0) <= 1) {
     result.resize_({0});
@@ -266,7 +266,7 @@ Tensor _pdist_backward(const Tensor& grad, const Tensor& self, const double p, c
   TORCH_CHECK(self.is_contiguous(), "_pdist_backward requires self to be contiguous");
   TORCH_CHECK(pdist.is_contiguous(), "_pdist_backward requires pdist to be contiguous");
   auto device = self.device().type();
-  TORCH_CHECK(device == kCPU || device == kCUDA || device == kXPU, "_pdist_backward only supports CPU, XPU and CUDA devices, got: ", device);
+  TORCH_CHECK(device == kCPU || device == kCUDA || device == kXPU || device == kPrivateUse1, "_pdist_backward only supports CPU, XPU, CUDA and PrivateUse1 devices, got: ", device);
   Tensor result = at::empty_like(self, LEGACY_CONTIGUOUS_MEMORY_FORMAT);
   pdist_backward_stub(device, result, grad, self, p, pdist);
   return result;
diff --git a/aten/src/ATen/native/Distributions.cpp b/aten/src/ATen/native/Distributions.cpp
index 5f34ed9d24c17..6382b1ecaace1 100644
--- a/aten/src/ATen/native/Distributions.cpp
+++ b/aten/src/ATen/native/Distributions.cpp
@@ -594,7 +594,8 @@ Tensor& multinomial_out(const Tensor& self,
   // https://github.com/pytorch/pytorch/issues/11931#issuecomment-625882503
   if (!with_replacement || n_sample == 1) {
     // Sanity checks on `self`.
-    auto is_valid = ((self.max() < INFINITY) & (self.min() >= 0));
+    auto [self_min, self_max] = self.aminmax();
+    auto is_valid = ((self_max < INFINITY) & (self_min >= 0));
     at::_assert_async(is_valid, "probability tensor contains either `inf`, `nan` or element < 0");
     at::Tensor zero_prob_condition;
     if (self.dim() == 1){
diff --git a/aten/src/ATen/native/Dropout.cpp b/aten/src/ATen/native/Dropout.cpp
index f7d32579165b4..421bd83d8d6ff 100644
--- a/aten/src/ATen/native/Dropout.cpp
+++ b/aten/src/ATen/native/Dropout.cpp
@@ -121,7 +121,7 @@ native_dropout_cpu(const Tensor& input, double p, std::optional<bool> train) {
     mask = at::ones_like(input, input.options().dtype(c10::CppTypeToScalarType<bool>::value));
     output = input.clone();
   }
-  return std::make_tuple(output, mask);
+  return std::make_tuple(std::move(output), std::move(mask));
 }
 
 Tensor native_dropout_backward(const Tensor& grad, const Tensor& mask, double scale) {
diff --git a/aten/src/ATen/native/ForeachOpsKernels.cpp b/aten/src/ATen/native/ForeachOpsKernels.cpp
index d411a68b037e0..2fc3c81ce0eb1 100644
--- a/aten/src/ATen/native/ForeachOpsKernels.cpp
+++ b/aten/src/ATen/native/ForeachOpsKernels.cpp
@@ -20,6 +20,7 @@
 #include <ATen/ops/_foreach_ceil_native.h>
 #include <ATen/ops/_foreach_clamp_max_native.h>
 #include <ATen/ops/_foreach_clamp_min_native.h>
+#include <ATen/ops/_foreach_clone_native.h>
 #include <ATen/ops/_foreach_copy_native.h>
 #include <ATen/ops/_foreach_cos_native.h>
 #include <ATen/ops/_foreach_cosh_native.h>
@@ -364,6 +365,21 @@ FOREACH_BINARY_OP_LIST(div)
 FOREACH_BINARY_OP_LIST(clamp_min)
 FOREACH_BINARY_OP_LIST(clamp_max)
 FOREACH_BINARY_OP_LIST(pow)
+
+// _foreach_clone
+std::vector<Tensor> foreach_tensor_clone_slow(
+    TensorList self,
+    std::optional<MemoryFormat> memory_format) {
+  check_foreach_api_restrictions(self);
+
+  std::vector<Tensor> ret{};
+  ret.reserve(self.size());
+  for (const auto& t : self) {
+    ret.emplace_back(t.clone(memory_format));
+  }
+  return ret;
+}
+
 // _foreach_copy_
 void foreach_tensor_copy_list_kernel_slow_(
     TensorList self,
diff --git a/aten/src/ATen/native/ForeachUtils.h b/aten/src/ATen/native/ForeachUtils.h
index f0dce20a6eff4..d68d770f215c0 100644
--- a/aten/src/ATen/native/ForeachUtils.h
+++ b/aten/src/ATen/native/ForeachUtils.h
@@ -93,22 +93,27 @@ inline void check_foreach_api_restrictions(
 // same device and dtype.
 inline bool _check_tensors_share_device_and_dtype(
     ArrayRef<TensorList> tensorLists,
-    const bool skip_dtype_check = false) {
+    const bool skip_cross_list_dtype_check = false) {
   const auto expected_dtype = tensorLists[0][0].dtype();
   const auto expected_device = tensorLists[0][0].device();
 
-  auto is_tensor_okay = [&](const Tensor& tensor) {
-    return (skip_dtype_check || tensor.dtype() == expected_dtype) &&
-        tensor.device() == expected_device && tensor.layout() == at::kStrided &&
-        tensor.is_non_overlapping_and_dense();
-  };
-
   return std::all_of(
       tensorLists.cbegin(),
       tensorLists.cend(),
       [&](const TensorList& tensorList) {
+        if (tensorList.empty()) {
+          return true;
+        }
+        const auto list_dtype = tensorList[0].dtype();
         return std::all_of(
-            tensorList.cbegin(), tensorList.cend(), is_tensor_okay);
+            tensorList.cbegin(), tensorList.cend(), [&](const Tensor& tensor) {
+              return tensor.device() == expected_device &&
+                  tensor.layout() == at::kStrided &&
+                  tensor.is_non_overlapping_and_dense() &&
+                  tensor.dtype() == list_dtype &&
+                  (skip_cross_list_dtype_check ||
+                   tensor.dtype() == expected_dtype);
+            });
       });
 }
 
@@ -192,8 +197,10 @@ inline bool _check_tensors_do_type_promotion_with_scalars(
 inline bool check_fast_path_restrictions(
     ArrayRef<TensorList> tensorLists,
     ArrayRef<Scalar> scalarList = {},
-    bool does_op_promote_integer_inputs_to_float = false) {
-  return _check_tensors_share_device_and_dtype(tensorLists) &&
+    bool does_op_promote_integer_inputs_to_float = false,
+    bool skip_cross_list_dtype_check = false) {
+  return _check_tensors_share_device_and_dtype(
+             tensorLists, skip_cross_list_dtype_check) &&
       _check_tensors_share_sizes_and_strides(tensorLists) &&
       _check_tensors_do_type_promotion_with_scalars(
              tensorLists[0],
@@ -318,9 +325,17 @@ inline FlatMap _group_tensors_by_first_tensors_device_and_dtype(
                 const auto s = tensor->scalar_type();
                 const auto d = tensor->device();
                 // Note: `step` or `state_step` is float32 by default.
+                // BFloat16 is allowed here for mixed-precision optimizer
+                // states (e.g. fp32 params with bf16 exp_avg/exp_avg_sq).
+                // Currently only BFloat16 is permitted because it is the
+                // only low-precision state dtype validated end-to-end in
+                // large-scale training (e.g. DeepSeek-V3 671B).
+                // TBD: make the set of allowed extra dtypes configurable
+                // per optimizer so this function stays dtype-agnostic.
                 if (key.first == d) {
                   return key.second == s || s == at::ScalarType::Float ||
-                      s == at::ScalarType::Double;
+                      s == at::ScalarType::Double ||
+                      s == at::ScalarType::BFloat16;
                 } else if (d.is_cpu()) {
                   // note(crcrpar): There are some test cases (e.g.
                   // TestOptim::test_adam) where state_steps are on CPU and the
@@ -333,7 +348,9 @@ inline FlatMap _group_tensors_by_first_tensors_device_and_dtype(
                 }
               }
             }),
-        "Tensors of the same index must be on the same device and the same dtype except `step` tensors that can be CPU and float32/64 notwithstanding");
+        "Tensors of the same index must be on the same device and the same dtype "
+        "except `step` tensors that can be CPU and float32/64, and optimizer "
+        "states that can be bfloat16 for mixed-precision training");
     grouped_tensors_with_indices.try_emplace(
         key,
         TensorsAndIndicesT{
diff --git a/aten/src/ATen/native/LinearAlgebra.cpp b/aten/src/ATen/native/LinearAlgebra.cpp
index 017b2d4956214..b5896575580f2 100644
--- a/aten/src/ATen/native/LinearAlgebra.cpp
+++ b/aten/src/ATen/native/LinearAlgebra.cpp
@@ -476,7 +476,7 @@ std::tuple<Tensor, Tensor> get_atol_rtol(
            ? at::where(atol_opt.value() > 0, at::zeros({}, options), default_rtol)
            : std::move(default_rtol);
   }
-  return std::make_tuple(atol, rtol);
+  return std::make_tuple(std::move(atol), std::move(rtol));
 }
 
 std::tuple<Tensor, Tensor> get_atol_rtol(
@@ -502,7 +502,7 @@ std::tuple<Tensor, Tensor> get_atol_rtol(
   }
   auto atol_tensor = at::full({}, atol, options);
   auto rtol_tensor = at::full({}, rtol, options);
-  return std::make_tuple(atol_tensor, rtol_tensor);
+  return std::make_tuple(std::move(atol_tensor), std::move(rtol_tensor));
 }
 
 } // anonymous namespace
@@ -2623,6 +2623,7 @@ Tensor compute_T18_scale_square(
   // 3. Multiply remain matrices by 3 times and slice to acc[1:]
   // All processed matrices will be stored in `output_pieces`.
   std::vector<Tensor> output_pieces;
+  output_pieces.reserve(section_numel);
   auto acc = mexp_scaled.index_select(0, sorted_s_inds);
   for (int64_t i = 0; i < section_numel; ++i) {
     for (int64_t j = 0; j < pts[i]; j++) {
@@ -3702,7 +3703,8 @@ Tensor& _int_mm_out_cpu(const Tensor& self, const Tensor& mat2, Tensor& result)
   TORCH_CHECK(self.dim() == 2, func_name, ": Expected self to be of dimension 2 but got ", self.dim());
   TORCH_CHECK(mat2.dim() == 2, func_name, ": Expected mat2 to be of dimension 2 but got ", mat2.dim());
   TORCH_CHECK(self.size(1) == mat2.size(0), func_name, ": self.size(1) needs to match mat2.size(0) but got ", self.size(1), " and ", mat2.size(0));
-  TORCH_CHECK(self.dtype() == at::kChar, func_name, ": Expected self dtype to be of type int8 but got ", self.dtype());
+  TORCH_CHECK(self.dtype() == at::kChar || self.dtype() == at::kByte,
+    func_name, ": Expected self dtype to be int8 or uint8 but got ", self.dtype());
   TORCH_CHECK(mat2.dtype() == at::kChar, func_name, ": Expected mat2 dtype to be of type int8 but got ", mat2.dtype());
   TORCH_CHECK(result.dtype() == at::kInt, func_name, ": Expected result dtype to be of type kInt but got ", result.dtype());
   TORCH_CHECK(result.size(0) == self.size(0), func_name, ": Expected result.size(0) to be ", self.size(0), " but got ", result.size(0));
@@ -3710,6 +3712,7 @@ Tensor& _int_mm_out_cpu(const Tensor& self, const Tensor& mat2, Tensor& result)
   TORCH_CHECK(result.dim() == 2, func_name, ": Expected result to be of dimension 2 but got ", result.dim());
   TORCH_CHECK(result.is_contiguous(), func_name, ": Expected result to be contiguous.");
 
+  // Outer or inner dimension is 0
   if (result.numel() == 0 || self.size(1) == 0) {
     return result.zero_();
   }
@@ -3724,7 +3727,6 @@ Tensor& _int_mm_out_cpu(const Tensor& self, const Tensor& mat2, Tensor& result)
     }
   }
   if (!dispatched) {
-    auto a = reinterpret_cast<int8_t*>(self.data_ptr());
     auto b = reinterpret_cast<int8_t*>(mat2.data_ptr());
     auto c = reinterpret_cast<int32_t*>(result.data_ptr());
     const int64_t m = result.size(0);
@@ -3735,18 +3737,26 @@ Tensor& _int_mm_out_cpu(const Tensor& self, const Tensor& mat2, Tensor& result)
     const int64_t ldb_0 = mat2.strides()[0];
     const int64_t ldb_1 = mat2.strides()[1];
     const int64_t ldc = result.strides()[0];
-    parallel_for(0, m * n, 1, [&](int64_t start, int64_t end) {
-      for (const auto i : c10::irange(start, end)) {
-        auto row = i / n;
-        auto col = i % n;
-        c[row * ldc + col] = 0;
-        for (const auto k : c10::irange(k)) {
-          c[row * ldc + col] = c[row * ldc + col] +
-              static_cast<int32_t>(a[row * lda_0 + k * lda_1]) *
-                  static_cast<int32_t>(b[k * ldb_0 + col * ldb_1]);
-        }
-      }
+    #define COMPUTE_WITH_A_TYPE(a_type)                             \
+    auto a = reinterpret_cast<a_type*>(self.data_ptr());            \
+    at::parallel_for(0, m * n, 1, [&](int64_t start, int64_t end) { \
+      for (const auto i : c10::irange(start, end)) {                \
+        auto row = i / n;                                           \
+        auto col = i % n;                                           \
+        c[row * ldc + col] = 0;                                     \
+        for (const auto k : c10::irange(k)) {                       \
+          c[row * ldc + col] = c[row * ldc + col] +                 \
+              static_cast<int32_t>(a[row * lda_0 + k * lda_1]) *    \
+                  static_cast<int32_t>(b[k * ldb_0 + col * ldb_1]); \
+        }                                                           \
+      }                                                             \
     });
+
+    if (self.scalar_type() == at::kByte) {
+      COMPUTE_WITH_A_TYPE(uint8_t);
+    } else {
+      COMPUTE_WITH_A_TYPE(int8_t);
+    }
   }
   return result;
 }
diff --git a/aten/src/ATen/native/LinearAlgebraUtils.h b/aten/src/ATen/native/LinearAlgebraUtils.h
index 257863573d3a8..7b5c0758ad9c1 100644
--- a/aten/src/ATen/native/LinearAlgebraUtils.h
+++ b/aten/src/ATen/native/LinearAlgebraUtils.h
@@ -336,7 +336,7 @@ inline std::tuple<Tensor,Tensor> _linalg_broadcast_batch_dims(const Tensor& arg1
 
   auto arg1_broadcasted  = arg1_expand_size == arg1.sizes() ? arg1 : arg1.expand(arg1_expand_size);
   auto arg2_broadcasted  = arg2_expand_size == arg2.sizes() ? arg2 : arg2.expand(arg2_expand_size);
-  return std::make_tuple(arg1_broadcasted, arg2_broadcasted);
+  return std::make_tuple(std::move(arg1_broadcasted), std::move(arg2_broadcasted));
 }
 
 inline std::vector<int64_t> broadcast_batch_size(const Tensor& t1, const Tensor& t2, int64_t n_batch_dims) {
@@ -351,6 +351,7 @@ inline Tensor _move_to_end(const Tensor& self, IntArrayRef axes) {
   const std::vector<int64_t> a = axes.vec();
   const int64_t ndim = self.ndimension();
   std::vector<int64_t> perm;
+  perm.reserve(static_cast<size_t>(std::max<int64_t>(0, ndim)));
 
   for (const auto i : c10::irange(ndim)) {
     auto it = std::find(a.begin(), a.end(), i);
@@ -405,7 +406,7 @@ inline std::tuple<DimVector, DimVector, int64_t> _compute_geometry_for_Q(
     n_columns_q = std::min(m, n);
   }
   auto q_strides = batched_matrix_contiguous_strides(q_sizes, /*f-contig*/true);
-  return std::make_tuple(q_sizes, q_strides, n_columns_q);
+  return std::make_tuple(std::move(q_sizes), std::move(q_strides), n_columns_q);
 }
 
 inline bool svd_uses_cusolver(const Tensor& A) {
diff --git a/aten/src/ATen/native/Loss.cpp b/aten/src/ATen/native/Loss.cpp
index 40d79d97c0cdf..d12ead20e2fb3 100644
--- a/aten/src/ATen/native/Loss.cpp
+++ b/aten/src/ATen/native/Loss.cpp
@@ -489,6 +489,9 @@ Tensor& mse_loss_backward_out(const Tensor& grad_output,
     .add_const_input(input)
     .add_const_input(target)
     .add_const_input(grad_output)
+    .promote_inputs_to_common_dtype(true)
+    .cast_common_dtype_to_outputs(true)
+    .enforce_safe_casting_to_output(true)
     .build();
   mse_backward_stub(iter.device_type(), iter, norm);
   return grad_input;
diff --git a/aten/src/ATen/native/LossMulti.h b/aten/src/ATen/native/LossMulti.h
index 8877b05a54cc3..cc339b1ce43c2 100644
--- a/aten/src/ATen/native/LossMulti.h
+++ b/aten/src/ATen/native/LossMulti.h
@@ -56,8 +56,10 @@ namespace at::native {
 
     TORCH_CHECK(
         target.dim() <= 1 && target.numel() == nframe,
-        "inconsistent target size, expected ", nframe, " but got ",
-        target.sizes());
+        "multi_margin_loss: target tensor should be 1-D with size equal to "
+        "the number of input samples (batch size). Expected target size [",
+        nframe, "], but got ", target.sizes(),
+        ". Input has shape ", input.sizes(), ".");
     if (weight && weight->defined()) {
       TORCH_CHECK(
           weight->dim() <= 1 && weight->numel() == dim,
diff --git a/aten/src/ATen/native/LossNLL.cpp b/aten/src/ATen/native/LossNLL.cpp
index 94bf75be2ee8f..0f7de004af8e8 100644
--- a/aten/src/ATen/native/LossNLL.cpp
+++ b/aten/src/ATen/native/LossNLL.cpp
@@ -47,6 +47,10 @@ TORCH_META_FUNC(nll_loss_forward)
   TORCH_CHECK(
       target.dim() <= 1,
       "0D or 1D target tensor expected, multi-target not supported");
+  TORCH_CHECK(
+      target.scalar_type() == kLong || target.scalar_type() == kByte,
+      "expected target dtype to be Long or Byte, but got ",
+      target.scalar_type());
   if (self.dim() == 1 && target.dim() == 1) {
       TORCH_CHECK_VALUE(
           target.size(0) == 1,
@@ -682,9 +686,21 @@ Tensor nll_loss_nd_symint(
         false, "Expected 1 or more dimensions (got ", self.dim(), ")");
   }
 
-  if (self.dim() != 1 && self.sym_sizes()[0] != target.sym_sizes()[0]) {
-    TORCH_CHECK_VALUE(
-        false,
+  if (self.dim() != 1) {
+    auto sizes_match = self.sym_sizes()[0].sym_eq(target.sym_sizes()[0]);
+    if (TORCH_GUARD_OR_FALSE(sizes_match.sym_not())) {
+      // Statically known mismatch - raise ValueError for eager mode
+      TORCH_CHECK_VALUE(
+          false,
+          "Expected input batch_size (",
+          self.sym_sizes()[0],
+          ") to match target batch_size (",
+          target.sym_sizes()[0],
+          ").");
+    }
+    // For unbacked symbolic shapes, emit runtime check.
+    TORCH_SYM_CHECK(
+        sizes_match,
         "Expected input batch_size (",
         self.sym_sizes()[0],
         ") to match target batch_size (",
diff --git a/aten/src/ATen/native/LossNLL2d.cpp b/aten/src/ATen/native/LossNLL2d.cpp
index 6f7fd8ad1e136..f43938a1f6aab 100644
--- a/aten/src/ATen/native/LossNLL2d.cpp
+++ b/aten/src/ATen/native/LossNLL2d.cpp
@@ -56,6 +56,10 @@ inline void check_inputs_nll_loss2d(
       "only batches of spatial inputs supported (4D tensors), "
       "but got input of dimension: ",
       input.dim());
+  TORCH_CHECK(
+      target.scalar_type() == kLong || target.scalar_type() == kByte,
+      "expected target dtype to be Long or Byte, but got ",
+      target.scalar_type());
   TORCH_CHECK(
       !weight.defined() || weight.numel() == input.size(1),
       "weight tensor should be defined either for all or no classes");
@@ -428,7 +432,7 @@ std::tuple<Tensor, Tensor> nll_loss2d_forward_cpu(
   auto total_weight = at::empty({0}, self.options());
   at::native::nll_loss2d_forward_out_cpu(
       self, target, weight_opt, reduction, ignore_index, output, total_weight);
-  return std::make_tuple(output, total_weight);
+  return std::make_tuple(std::move(output), std::move(total_weight));
 }
 
 Tensor& nll_loss2d_backward_out_cpu(const Tensor& grad_output,
diff --git a/aten/src/ATen/native/Math.h b/aten/src/ATen/native/Math.h
index de0a40ae412d6..7d9e73779a4bd 100644
--- a/aten/src/ATen/native/Math.h
+++ b/aten/src/ATen/native/Math.h
@@ -2281,7 +2281,7 @@ inline C10_HOST_DEVICE T airy_ai_forward(T x) {
 
     int domain_flag = 0;
 
-    T ai;
+    T ai = T(0.0);
 
     if (std::isinf(x)) {
         return std::numeric_limits<T>::quiet_NaN();
@@ -3267,7 +3267,7 @@ inline C10_HOST_DEVICE T modified_bessel_i0_forward(T x) {
             +8.04490411014108831608e-01,
     };
 
-    T p;
+    T p = T{0};
     T q = 0.0;
 
     if (std::abs(x) <= T(8.0)) {
@@ -3355,7 +3355,7 @@ inline C10_HOST_DEVICE T modified_bessel_i1_forward(T x) {
             +7.78576235018280120474e-01,
     };
 
-    T p;
+    T p = T{0};
     T q = 0.0;
 
     if (std::abs(x) <= T(8.0)) {
@@ -3440,7 +3440,7 @@ inline C10_HOST_DEVICE T modified_bessel_k0_forward(T x) {
         return std::numeric_limits<T>::quiet_NaN();
     }
 
-    T p;
+    T p = T{0};
     T q = 0.0;
 
     if (x <= T(2.0)) {
@@ -3518,7 +3518,7 @@ inline C10_HOST_DEVICE T modified_bessel_k1_forward(T x) {
         return std::numeric_limits<T>::quiet_NaN();
     }
 
-    T p;
+    T p = T{0};
     T q = 0.0;
 
     if (x <= T(2.0)) {
@@ -3595,7 +3595,7 @@ inline C10_HOST_DEVICE T scaled_modified_bessel_k0_forward(T x) {
         return std::numeric_limits<T>::quiet_NaN();
     }
 
-    T p;
+    T p = T{0};
     T q = 0.0;
 
     if (x <= T(2.0)) {
@@ -3673,7 +3673,7 @@ inline C10_HOST_DEVICE T scaled_modified_bessel_k1_forward(T x) {
         return std::numeric_limits<T>::quiet_NaN();
     }
 
-    T p;
+    T p = T{0};
     T q = 0.0;
 
     if (x <= T(2.0)) {
diff --git a/aten/src/ATen/native/NaiveConvolutionTranspose2d.cpp b/aten/src/ATen/native/NaiveConvolutionTranspose2d.cpp
index 08c42a0d470c7..2042f7f5d6a62 100644
--- a/aten/src/ATen/native/NaiveConvolutionTranspose2d.cpp
+++ b/aten/src/ATen/native/NaiveConvolutionTranspose2d.cpp
@@ -868,7 +868,7 @@ static std::tuple<Tensor, Tensor, Tensor> slow_conv_transpose2d_backward_cpu(
         1);
   }
 
-  return std::tuple<Tensor, Tensor, Tensor>(grad_input, grad_weight, grad_bias);
+  return std::tuple<Tensor, Tensor, Tensor>(std::move(grad_input), std::move(grad_weight), std::move(grad_bias));
 }
 
 REGISTER_ALL_CPU_DISPATCH(slow_conv_transpose2d_backward_stub, &slow_conv_transpose2d_backward_cpu)
diff --git a/aten/src/ATen/native/Normalization.cpp b/aten/src/ATen/native/Normalization.cpp
index 76ce2802f05b0..71dab6aed8955 100644
--- a/aten/src/ATen/native/Normalization.cpp
+++ b/aten/src/ATen/native/Normalization.cpp
@@ -118,14 +118,14 @@ struct Var {
   }
 };
 
-static inline bool is_contiguous(const Tensor& t) {
+static bool is_contiguous_in_any_format(const Tensor& t) {
   return t.is_contiguous() || t.is_contiguous(at::MemoryFormat::ChannelsLast) || t.is_contiguous(at::MemoryFormat::ChannelsLast3d);
 }
 
 // For some ambiguous cases, it is possible a channels last contiguous Tensor has
 //   `suggest_memory_format` of Contiguous.
 // See https://github.com/pytorch/pytorch/issues/63224 for details.
-static inline MemoryFormat suggest_memory_format_contig(const Tensor& t) {
+static MemoryFormat suggest_memory_format_contig(const Tensor& t) {
   return t.is_contiguous() ?
     at::MemoryFormat::Contiguous : (t.is_contiguous(at::MemoryFormat::ChannelsLast3d) ?
     at::MemoryFormat::ChannelsLast3d : at::MemoryFormat::ChannelsLast);
@@ -138,8 +138,8 @@ static std::tuple<Tensor,Tensor,Tensor> batch_norm_cpu_transform_input_template(
     const Tensor& running_mean /* optional */, const Tensor& running_var /* optional */,
     bool train, double eps, Tensor& output) {
 
-  bool all_contiguous = is_contiguous(input)
-    && is_contiguous(output)
+  bool all_contiguous = is_contiguous_in_any_format(input)
+    && is_contiguous_in_any_format(output)
     && (!weight.defined() || weight.is_contiguous())
     && (!bias.defined() || bias.is_contiguous())
     && running_mean.is_contiguous()
@@ -207,7 +207,7 @@ static std::tuple<Tensor,Tensor> batch_norm_cpu_update_stats_template(
   TORCH_CHECK(input.numel() != 0, "input tensor must have at least one element, but got input_sizes = ", input.sizes());
   int64_t n = input.numel() / n_input;
 
-  bool all_contiguous = is_contiguous(input);
+  bool all_contiguous = is_contiguous_in_any_format(input);
   constexpr bool mixed_type = !std::is_same_v<scalar_t, param_t>;
   // Using float data type for Half _var_sum in batchnorm stats updating on CPU
   // to avoid _var_sum overflow since the representation range of Half is small.
@@ -300,7 +300,7 @@ static std::tuple<Tensor,Tensor> batch_norm_cpu_update_stats_template(
 
   constexpr bool mixed_type = !std::is_same_v<scalar_t, param_t>;
   const auto dtype = mixed_type ? kFloat : input.scalar_type();
-  Tensor save_mean = is_contiguous(input) ? at::empty({n_input}, input.options().dtype(dtype)) : at::mean(input, /*dim=*/reduce_dims, /*keepdim=*/false, dtype);
+  Tensor save_mean = is_contiguous_in_any_format(input) ? at::empty({n_input}, input.options().dtype(dtype)) : at::mean(input, /*dim=*/reduce_dims, /*keepdim=*/false, dtype);
   Tensor save_var_transform = at::empty({n_input}, input.options().dtype(dtype));
   return batch_norm_cpu_update_stats_template<scalar_t, param_t, VarTransform>(input, running_mean, running_var, momentum, eps, save_mean, save_var_transform);
 }
@@ -331,8 +331,8 @@ static std::tuple<Tensor, Tensor, Tensor> batch_norm_backward_cpu_template(
 
   // since we are directly manipulating pointers in contiguous path,
   // need to make sure input and grad_out have the same memory format.
-  bool all_contiguous = is_contiguous(input)
-      && is_contiguous(grad_out_)
+  bool all_contiguous = is_contiguous_in_any_format(input)
+      && is_contiguous_in_any_format(grad_out_)
       && input.suggest_memory_format() == grad_out_.suggest_memory_format();
 
   if (all_contiguous) {
@@ -341,7 +341,7 @@ static std::tuple<Tensor, Tensor, Tensor> batch_norm_backward_cpu_template(
     }
     batch_norm_cpu_backward_stub(kCPU, grad_input, grad_weight, grad_bias,
         grad_out_, input, weight, running_mean, running_var, save_mean, save_invstd, train, eps);
-    return std::make_tuple(grad_input, grad_weight, grad_bias);
+    return std::make_tuple(std::move(grad_input), std::move(grad_weight), std::move(grad_bias));
   }
 
   auto weight_a = conditional_accessor_1d<const param_t>(weight);
@@ -491,7 +491,7 @@ static std::tuple<Tensor, Tensor, Tensor> batch_norm_backward_cpu_template(
         }
       }
     });
-  return std::make_tuple(grad_input, grad_weight, grad_bias);
+  return std::make_tuple(std::move(grad_input), std::move(grad_weight), std::move(grad_bias));
 }
 
 BatchNormBackend _select_batch_norm_backend(
@@ -671,7 +671,7 @@ std::tuple<Tensor, Tensor, Tensor> _batch_norm_impl_index_backward(
     if (output_mask[0] && weight.defined()) {
       grad_input = grad_output * weight[0];
     }
-    return std::make_tuple(grad_input, grad_weight, grad_bias);
+    return std::make_tuple(std::move(grad_input), std::move(grad_weight), std::move(grad_bias));
   }
 
   // backward in inference mode is not supported in cudnn, fallback to native
@@ -863,7 +863,7 @@ std::tuple<Tensor, Tensor, Tensor> batch_norm_cpu(const Tensor& self, const std:
   checkBackend("batch_norm_cpu", {self, weight, bias, running_mean, running_var}, Backend::CPU);
 
   // Prepare output tensor
-  const bool all_contiguous = is_contiguous(self)
+  const bool all_contiguous = is_contiguous_in_any_format(self)
     && (!weight.defined() || weight.is_contiguous())
     && (!bias.defined() || bias.is_contiguous())
     && running_mean.is_contiguous()
@@ -885,7 +885,7 @@ std::tuple<Tensor, Tensor, Tensor> batch_norm_cpu(const Tensor& self, const std:
       save_mean = at::empty({0}, self.options().dtype(kFloat));
       save_var = at::empty({0}, self.options().dtype(kFloat));
     } else {
-      save_mean = is_contiguous(self) ? at::empty({self.size(1)}, self.options().dtype(kFloat)) : at::mean(self, /*dim=*/reduce_dims, /*keepdim=*/false, kFloat);
+      save_mean = is_contiguous_in_any_format(self) ? at::empty({self.size(1)}, self.options().dtype(kFloat)) : at::mean(self, /*dim=*/reduce_dims, /*keepdim=*/false, kFloat);
       save_var = at::empty({self.size(1)}, self.options().dtype(kFloat));
     }
   } else {
@@ -893,7 +893,7 @@ std::tuple<Tensor, Tensor, Tensor> batch_norm_cpu(const Tensor& self, const std:
       save_mean = at::empty({0}, self.options());
       save_var = at::empty({0}, self.options());
     } else {
-      save_mean = is_contiguous(self) ? at::empty({self.size(1)}, self.options()) : at::mean(self, /*dim=*/reduce_dims, /*keepdim=*/false);
+      save_mean = is_contiguous_in_any_format(self) ? at::empty({self.size(1)}, self.options()) : at::mean(self, /*dim=*/reduce_dims, /*keepdim=*/false);
       save_var = at::empty({self.size(1)}, self.options());
     }
   }
diff --git a/aten/src/ATen/native/Onehot.cpp b/aten/src/ATen/native/Onehot.cpp
index 2a20f95f10c20..195b58eac6a0f 100644
--- a/aten/src/ATen/native/Onehot.cpp
+++ b/aten/src/ATen/native/Onehot.cpp
@@ -48,16 +48,14 @@ Tensor one_hot(const Tensor &self, int64_t num_classes) {
     }
 
     // non-empty tensor
-    if (self.device().type() != at::kCUDA && self.device().type() != at::kMPS &&
-        self.device().type() != at::kPrivateUse1 && self.device().type() != at::kXLA) {
+    if (self.device().is_cpu()) {
       // for cuda, rely on device assert thrown by scatter
       TORCH_CHECK(self.min().item().toLong() >= 0, "Class values must be non-negative.");
     }
     if (num_classes == -1) {
         num_classes = self.max().item().toLong() + 1;
     } else {
-        if (self.device().type() != at::kCUDA && self.device().type() != at::kMPS &&
-            self.device().type() != at::kPrivateUse1 && self.device().type() != at::kXLA) {
+        if (self.device().is_cpu()) {
           // rely on device asserts from scatter to avoid sync here
           TORCH_CHECK(num_classes > self.max().item().toLong(), "Class values must be smaller than num_classes.");
         } else {
diff --git a/aten/src/ATen/native/PackedSequence.cpp b/aten/src/ATen/native/PackedSequence.cpp
index be7961b2a2452..d757e8ba41c8d 100644
--- a/aten/src/ATen/native/PackedSequence.cpp
+++ b/aten/src/ATen/native/PackedSequence.cpp
@@ -106,7 +106,7 @@ std::tuple<Tensor, Tensor> _pack_padded_sequence(const Tensor& _input, const Ten
     TORCH_CHECK(l >= prev_l);
   }
 
-  return std::make_tuple(at::cat(steps), batch_sizes_t);
+  return std::make_tuple(at::cat(steps), std::move(batch_sizes_t));
 }
 
 // `grad` could be on arbitrary device and of arbitrary dtype, but `_batch_sizes`
@@ -200,7 +200,7 @@ std::tuple<Tensor, Tensor> _pad_packed_sequence(const Tensor& data, const Tensor
     output = output.transpose(0, 1);
   }
 
-  return std::make_tuple(output, lengths_t);
+  return std::make_tuple(std::move(output), std::move(lengths_t));
 }
 
 Tensor pad_sequence(TensorList sequences, bool batch_first, double padding_value, const std::string_view padding_side) {
diff --git a/aten/src/ATen/native/README.md b/aten/src/ATen/native/README.md
index ab836593dc297..87512d6058a4b 100644
--- a/aten/src/ATen/native/README.md
+++ b/aten/src/ATen/native/README.md
@@ -5,7 +5,7 @@ in one of the `cpp` files in this directory.
 
 Like all ATen methods/functions, native functions are made available
 from both ATen's C++ and Python APIs.  In C++, they are made available
-either as methods on `Tensor` (`t.mymeth()`) and functions in the ATen
+either as methods on `Tensor` (`t.mymeth()`) or as functions in the ATen
 namespace (`at::myfunc()`).  In PyTorch, they are made available as
 methods on `Variable` or as functions on `torch._C._FunctionBase`.
 (It is the user's responsibility to re-export these functions in
diff --git a/aten/src/ATen/native/RNN.cpp b/aten/src/ATen/native/RNN.cpp
index a18053cc9a3f6..80c466a6b5815 100644
--- a/aten/src/ATen/native/RNN.cpp
+++ b/aten/src/ATen/native/RNN.cpp
@@ -104,8 +104,9 @@ bool use_mkldnn(const Tensor& input, TensorList params, TensorList hx) {
        (input.scalar_type() == kHalf && !at::GradMode::is_enabled() &&
         mkldnn_fp16_device_check())) &&
       input.numel() != 0;
-#endif
+#else
   return false;
+#endif
 }
 
 bool use_cudnn(const Tensor& t) {
@@ -852,7 +853,7 @@ struct FullLayer : Layer<Tensor, hidden_type, cell_params> {
       hidden = cell_(input, hidden, params, pre_compute_input);
       step_outputs.emplace_back(hidden_as_output(hidden));
     }
-    return {step_outputs, hidden};
+    return {std::move(step_outputs), std::move(hidden)};
   }
 
   output_type operator()(
@@ -1111,7 +1112,7 @@ apply_layer_stack(const Layer<io_type, hidden_type, weight_type>& layer, const i
     }
   }
 
-  return {layer_input, final_hiddens};
+  return {std::move(layer_input), std::move(final_hiddens)};
 }
 
 ////////////////////////////////////////////////////////////////////////////////
diff --git a/aten/src/ATen/native/RNN.h b/aten/src/ATen/native/RNN.h
index afebf06a0fba3..efae0682f6813 100644
--- a/aten/src/ATen/native/RNN.h
+++ b/aten/src/ATen/native/RNN.h
@@ -28,6 +28,13 @@ DECLARE_DISPATCH(rnn_packed_fn, rnn_tanh_packed_miopen_stub)
 DECLARE_DISPATCH(rnn_packed_fn, rnn_relu_packed_cudnn_stub)
 DECLARE_DISPATCH(rnn_packed_fn, rnn_relu_packed_miopen_stub)
 
+// Releases cached RNN dropout state buffers, which are long-lived allocations
+// (process-lifetime for cuDNN, thread-lifetime for MIOpen) that empty_cache()
+// cannot reclaim on its own. Only the corresponding backend's function is
+// defined for a given build.
+TORCH_CUDA_CPP_API void _cudnn_clear_dropout_state();
+TORCH_CUDA_CPP_API void _miopen_clear_dropout_state();
+
 inline void check_attributes(const Tensor& input, const TensorList& params, const TensorList& hiddens, bool check_dtype=false) {
   auto input_device = input.device();
   auto input_dtype = input.scalar_type();
diff --git a/aten/src/ATen/native/ReduceOps.cpp b/aten/src/ATen/native/ReduceOps.cpp
index 6be4d72b88ed6..d5dc9ba1eeb46 100644
--- a/aten/src/ATen/native/ReduceOps.cpp
+++ b/aten/src/ATen/native/ReduceOps.cpp
@@ -117,6 +117,7 @@
 #include <ATen/ops/var_mean.h>
 #include <ATen/ops/var_mean_native.h>
 #include <ATen/ops/var_native.h>
+#include <ATen/ops/where.h>
 #include <ATen/ops/zeros.h>
 #include <ATen/ops/zeros_like.h>
 #endif
@@ -645,7 +646,12 @@ Tensor cumprod_backward(const Tensor& grad, const Tensor& input, int64_t dim, co
   // O(n) implementation. The derivative of this implementation is _not_
   // the second derivative of cumprod. As such, we fallback to a less efficient
   // O(n^2) implementation when at::GradMode::is_enabled().
-  if (!at::GradMode::is_enabled() && !are_inputs_tensors_sublcass) {
+  //
+  // NOTE: We use at::where instead of masked_scatter_/masked_select to make
+  // this path composite compliant for tensor subclasses (e.g., FakeTensors
+  // used by torch.compile). masked_select has dynamic output shape which
+  // causes issues with tracing. See https://github.com/pytorch/pytorch/issues/136263
+  if (!at::GradMode::is_enabled()) {
     // n.b. This could probably be implemented much faster with a kernel
 
     // From here on we need to use some mask gymnastics to
@@ -669,9 +675,16 @@ Tensor cumprod_backward(const Tensor& grad, const Tensor& input, int64_t dim, co
     // case k < z1
     // select everything before the first zero [0, z1)
     auto mask = cumsum == 0;
-    // equiv to grad_input[mask] = deriv[grad]
-    grad_input.masked_scatter_(mask,
-        reversed_cumsum(w.masked_fill(~mask, 0.), dim).div_(input_conj).masked_select(mask));
+    // Compute gradient for positions before the first zero
+    // Using at::where instead of masked_scatter_ for composite compliance
+    auto grad_before_first_zero = reversed_cumsum(w.masked_fill(~mask, 0.), dim);
+    if (!are_inputs_tensors_sublcass) {
+      grad_before_first_zero = grad_before_first_zero.div_(input_conj);
+    } else {
+      grad_before_first_zero = grad_before_first_zero.div(input_conj);
+    }
+    grad_input = at::where(mask, grad_before_first_zero, grad_input);
+
     // select everything from the first zero to the second zero [z1, z2)
     mask = cumsum == 1;
 
@@ -693,13 +706,21 @@ Tensor cumprod_backward(const Tensor& grad, const Tensor& input, int64_t dim, co
     // dy_j / dx_z1 = sum(cumprod(input[z1+1:z2] * grad[z1+1:z2])) * prod(output[z1-1])
     // relu_() necessary as gather does not support negative indices
     // finally, we do grad_input[z1] = dy_j / dx_z1
-    grad_input.masked_scatter_(first_zero_mask,
-                               input_conj.masked_fill(~mask, 1.).cumprod(dim)
-                                    .mul_(grad.masked_fill(cumsum != 1, 0.))
-                                    .sum(dim, /*keepdim*/true)
-                                    .mul_(at::gather(output_conj, dim, (first_zero_index - 1).relu_())
-                                          .masked_fill_(first_zero_index == 0, 1.))
-                                    .masked_select(first_zero_mask));
+    // Using at::where instead of masked_scatter_ for composite compliance
+    auto grad_at_first_zero = input_conj.masked_fill(~mask, 1.).cumprod(dim);
+    const auto grad_masked = grad.masked_fill(cumsum != 1, 0.);
+    const auto output_before_zero = at::gather(output_conj, dim, (first_zero_index - 1).relu_())
+                                      .masked_fill_(first_zero_index == 0, 1.);
+    if (!are_inputs_tensors_sublcass) {
+      grad_at_first_zero = grad_at_first_zero.mul_(grad_masked)
+                             .sum(dim, /*keepdim*/true)
+                             .mul_(output_before_zero);
+    } else {
+      grad_at_first_zero = grad_at_first_zero.mul(grad_masked)
+                             .sum(dim, /*keepdim*/true)
+                             .mul(output_before_zero);
+    }
+    grad_input = at::where(first_zero_mask, grad_at_first_zero, grad_input);
     return grad_input;
   } else { // GradMode::enabled()
     /*
@@ -840,7 +861,7 @@ std::tuple<Tensor, Tensor> cummax(const Tensor& self, int64_t dim) {
   auto values = at::empty(self.sizes(), self.options());
   auto indices = at::empty(self.sizes(), self.options().dtype(at::kLong));
   at::cummax_out(values, indices, self, dim);
-  return std::make_tuple(values, indices);
+  return std::make_tuple(std::move(values), std::move(indices));
 }
 
 void cummin_helper_cpu(const Tensor& self, Tensor& values, Tensor& indices, int64_t dim) {
@@ -879,7 +900,7 @@ std::tuple<Tensor, Tensor> cummin(const Tensor& self, int64_t dim) {
   auto values = at::empty(self.sizes(), self.options());
   auto indices = at::empty(self.sizes(), self.options().dtype(at::kLong));
   at::cummin_out(values, indices, self, dim);
-  return std::make_tuple(values, indices);
+  return std::make_tuple(std::move(values), std::move(indices));
 }
 
 Tensor cummaxmin_backward(const Tensor& grad, const Tensor& input, const Tensor& indices, int64_t dim) {
@@ -1864,8 +1885,8 @@ static Tensor& std_var_out(
     const char* fname, Tensor& result, const Tensor& self,
     at::OptionalIntArrayRef dim, const std::optional<Scalar>& correction_opt,
     bool keepdim, bool take_sqrt) {
-  TORCH_CHECK(self.device().is_cpu() || self.device().is_cuda() || self.device().is_xpu(),
-              "std and var supports tensors on a CPU, CUDA, or XPU device only, but got: ",
+  TORCH_CHECK(self.device().is_cpu() || self.device().is_cuda() || self.device().is_xpu() || self.device().is_privateuseone(),
+              "std and var supports tensors on a CPU, CUDA, XPU or PrivateUse1 device only, but got: ",
               self.device().type());
   TORCH_CHECK(self.layout() == Layout::Strided,
               "std and var only supports strided layout, got: ", self.layout());
@@ -1937,8 +1958,8 @@ static std::tuple<Tensor&, Tensor&> std_var_mean_out(
     at::OptionalIntArrayRef dim, const std::optional<Scalar>& correction_opt,
     bool keepdim, bool take_sqrt) {
   AT_ASSERT(result1.defined() && result2.defined());
-  TORCH_CHECK(self.device().is_cpu() || self.is_cuda() || self.is_xpu(),
-              fname, " supports tensors on a CPU, CUDA, or XPU device only, got: ",
+  TORCH_CHECK(self.device().is_cpu() || self.is_cuda() || self.is_xpu() || self.is_privateuseone(),
+              fname, " supports tensors on a CPU, CUDA, XPU or PrivateUse1 device only, got: ",
               self.device().type());
   TORCH_CHECK(self.layout() == Layout::Strided,
               fname, " only supports strided layout, got: ", self.layout());
@@ -2351,10 +2372,14 @@ bool cpu_equal(const Tensor& self, const Tensor& other) {
 Tensor value_selecting_reduction_backward_symint(const Tensor& grad, int64_t dim, const Tensor& indices, c10::SymIntArrayRef sizes, bool keepdim) {
   auto inplace_scatter_if_not_tensor_subclass =
       [&](const Tensor& grad_out, const Tensor& indices_) {
-        auto grad_in = at::zeros_symint(sizes, grad_out.options());
         if (areAnyTensorSubclassLike({grad, indices})) {
+          // Use new_zeros_symint so that tensor subclasses (e.g. DTensor)
+          // can intercept the zeros creation through dispatch, ensuring
+          // the result has matching subclass type for subsequent scatter.
+          auto grad_in = grad_out.new_zeros_symint(sizes);
           return grad_in.scatter(dim, indices_, grad_out);
         }
+        auto grad_in = at::zeros_symint(sizes, grad_out.options());
         return grad_in.scatter_(dim, indices_, grad_out);
       };
 
diff --git a/aten/src/ATen/native/Repeat.cpp b/aten/src/ATen/native/Repeat.cpp
index 11b528b445ed4..62c66e2be1bea 100644
--- a/aten/src/ATen/native/Repeat.cpp
+++ b/aten/src/ATen/native/Repeat.cpp
@@ -105,7 +105,7 @@ Tensor repeat_interleave_symint(
     std::optional<SymInt> output_size) {
   Tensor input = dim_opt ? self : self.flatten();
   int64_t dim = c10::maybe_wrap_dim(dim_opt.value_or(0), self.dim());
-  TORCH_CHECK(repeats >= 0, "Repeats must be non-negative");
+  TORCH_SYM_CHECK(repeats.sym_ge(0), "Repeats must be non-negative");
 
   input = input.unsqueeze(dim + 1);
   auto expand_shape = input.sym_sizes().vec();
@@ -115,9 +115,13 @@ Tensor repeat_interleave_symint(
   // This argument doesn't really make sense for the scalar overload, but exists
   // for consistency with the tensor overload
   if (output_size) {
-    auto calculated_size = (repeats * expand_shape[dim]).guard_int(__FILE__, __LINE__);
-    TORCH_CHECK(*output_size == calculated_size, "repeat_interleave: Invalid output_size, expected ",
-                calculated_size, " but got ", *output_size);
+    auto calculated_size = repeats * expand_shape[dim];
+    TORCH_SYM_CHECK(
+        output_size->sym_eq(calculated_size),
+        "repeat_interleave: Invalid output_size, expected ",
+        calculated_size,
+        " but got ",
+        *output_size);
   }
 
   return input.clone(at::MemoryFormat::Contiguous).flatten(dim, dim + 1);
diff --git a/aten/src/ATen/native/Resize.h b/aten/src/ATen/native/Resize.h
index 3346cd2cb220e..d4bc55e30e701 100644
--- a/aten/src/ATen/native/Resize.h
+++ b/aten/src/ATen/native/Resize.h
@@ -181,24 +181,76 @@ inline void checkSetStorage(Tensor& result, Storage storage, T storage_offset,
  * (size, stride, storage_offset) must be in bounds for self's storage.
  */
 template <typename T>
-inline void setStrided(
-    const Tensor& self,
+void checkAsStridedArgs(
     ArrayRef<T> size,
     ArrayRef<T> stride,
     T storage_offset) {
-  TORCH_CHECK(size.size() == stride.size(), "mismatch in length of strides and shape");
+  TORCH_CHECK(
+      size.size() == stride.size(), "mismatch in length of strides and shape");
   for (const auto& val : stride) {
-    TORCH_CHECK(val >= 0,
-                "as_strided: Negative strides are not supported at the moment, "
-                "got strides: ", stride);
+    TORCH_CHECK(
+        val >= 0,
+        "as_strided: Negative strides are not supported at the moment, "
+        "got strides: ",
+        stride);
   }
+  TORCH_CHECK(storage_offset >= 0, "Tensor: invalid storage offset ", storage_offset);
+}
+
+template <typename T>
+void checkAsStridedArgsAllowUnbackedSymInts(
+    ArrayRef<T> size,
+    ArrayRef<T> stride,
+    T storage_offset) {
+  TORCH_CHECK(
+      size.size() == stride.size(), "mismatch in length of strides and shape");
+  if constexpr (std::is_same_v<T, c10::SymInt>) {
+    // FakeTensor/Meta view replay can pass ephemeral symbolic metadata here,
+    // so only validate values once the SymInts become concrete.
+    for (const auto& val : stride) {
+      if (auto maybe_val = val.maybe_as_int()) {
+        TORCH_CHECK(
+            *maybe_val >= 0,
+            "as_strided: Negative strides are not supported at the moment, "
+            "got strides: ",
+            stride);
+      }
+    }
+
+    if (auto maybe_storage_offset = storage_offset.maybe_as_int()) {
+      TORCH_CHECK(
+          *maybe_storage_offset >= 0,
+          "Tensor: invalid storage offset ",
+          storage_offset);
+    }
+  } else {
+    for (const auto& val : stride) {
+      TORCH_CHECK(
+          val >= 0,
+          "as_strided: Negative strides are not supported at the moment, "
+          "got strides: ",
+          stride);
+    }
+
+    TORCH_CHECK(
+        storage_offset >= 0,
+        "Tensor: invalid storage offset ",
+        storage_offset);
+  }
+}
+
+template <typename T>
+inline void setStrided(
+    const Tensor& self,
+    ArrayRef<T> size,
+    ArrayRef<T> stride,
+    T storage_offset) {
+  checkAsStridedArgs(size, stride, storage_offset);
 
   auto* self_ = self.unsafeGetTensorImpl();
   checkInBoundsForStorage(
       size, stride, storage_offset, self_->dtype(), self_->storage());
 
-  /* storage offset */
-  TORCH_CHECK(storage_offset >= 0, "Tensor: invalid storage offset ", storage_offset);
   self_->set_sizes_and_strides(size, stride, storage_offset);
 }
 
diff --git a/aten/src/ATen/native/RowwisePrune.cpp b/aten/src/ATen/native/RowwisePrune.cpp
index 26062a071ff70..ec4e5fff10a01 100644
--- a/aten/src/ATen/native/RowwisePrune.cpp
+++ b/aten/src/ATen/native/RowwisePrune.cpp
@@ -53,8 +53,8 @@ std::tuple<Tensor, Tensor> _rowwise_prune_helper(
       }
     }
   });
-  return std::tuple<Tensor, Tensor>(pruned_2d_tensor,
-      compressed_indices_mapping);
+  return std::tuple<Tensor, Tensor>(std::move(pruned_2d_tensor),
+      std::move(compressed_indices_mapping));
 }
 
 } // namespace
diff --git a/aten/src/ATen/native/ScaledBlas.cpp b/aten/src/ATen/native/ScaledBlas.cpp
new file mode 100644
index 0000000000000..0b384d7226b73
--- /dev/null
+++ b/aten/src/ATen/native/ScaledBlas.cpp
@@ -0,0 +1,536 @@
+#define TORCH_ASSERT_ONLY_METHOD_OPERATORS
+#include <ATen/core/Tensor.h>
+#include <ATen/core/NamedTensor.h>
+#include <ATen/Dispatch.h>
+#include <ATen/ExpandUtils.h>
+#include <ATen/NamedTensorUtils.h>
+#include <ATen/Config.h>
+
+#include <ATen/native/mkldnn/Matmul.h>
+#include <ATen/native/mkldnn/Linear.h>
+#include <ATen/native/Resize.h>
+#include <ATen/native/GroupedMMUtils.h>
+#include <ATen/BlasBackend.h>
+#include <ATen/native/ScaledBlasUtils.h>
+#if !defined(__s390x__) && !defined(__powerpc__)
+#include <cpuinfo.h>
+#endif
+
+#ifndef AT_PER_OPERATOR_HEADERS
+#include <ATen/CPUFunctions.h>
+#include <ATen/Functions.h>
+#include <ATen/NativeFunctions.h>
+#else
+#include <ATen/ops/_efficientzerotensor.h>
+#include <ATen/ops/addmv.h>
+#include <ATen/ops/addmv_native.h>
+#include <ATen/ops/copy_native.h>
+#include <ATen/ops/dot.h>
+#include <ATen/ops/dot_native.h>
+#include <ATen/ops/empty.h>
+#include <ATen/ops/mul_cpu_dispatch.h>
+#include <ATen/ops/mv_native.h>
+#include <ATen/ops/scalar_tensor_native.h>
+#include <ATen/ops/vdot_native.h>
+#include <ATen/ops/_scaled_mm_native.h>
+#include <ATen/ops/_scaled_mm_v2_native.h>
+#include <ATen/ops/mul.h>
+#include <ATen/ops/matmul.h>
+#endif
+
+namespace at::native {
+
+using at::blas::ScalingType;
+using at::blas::SwizzleType;
+
+namespace {
+
+void invalid_scaling_config(
+    const at::Tensor& mat_a,
+    const at::Tensor& mat_b,
+    const std::optional<at::Tensor>& scale_a,
+    const std::optional<at::Tensor>& scale_b) {
+  std::stringstream exception_ss;
+  exception_ss << "Invalid scaling configuration.\n";
+  exception_ss << "- For TensorWise scaling, a and b should be float8, scales should be float and singletons.\n";
+  exception_ss << "- For RowWise scaling, a and b should be float8, scales should be float, scale_a should be (";
+  exception_ss << mat_a.size(0) << ", 1) and scale_b should be (1, " << mat_b.size(1) << "), and both should be contiguous.\n";
+  exception_ss << "Got mat_a.dtype()=" << mat_a.scalar_type();
+  if (scale_a.has_value()) {
+    exception_ss << ", scale_a.dtype()=" << scale_a.value().scalar_type() << ", scale_a.size()=" << scale_a.value().sizes() << ", scale_a.stride()=" << scale_a.value().strides();
+  }
+  else {
+    exception_ss << ", scale_a=None";
+  }
+  exception_ss << ", mat_b.dtype()=" << mat_b.scalar_type();
+  if (scale_b.has_value()) {
+    exception_ss << ", scale_b.dtype()=" << scale_b.value().scalar_type() << ", scale_b.size()=" << scale_b.value().sizes() << " and scale_b.stride()=" << scale_b.value().strides();
+  }
+  else {
+    exception_ss << " and scale_b=None";
+  }
+
+  TORCH_CHECK_VALUE(false, exception_ss.str());
+}
+
+/*
+ * Scaling Type Determination:
+ * ---------------------------
+ * Conditions and corresponding Scaling Types:
+ *
+ * - If scale.numel() == 1:
+ *   - Returns TensorWise.
+ *
+ * - Else if scale.dim() == 2 && scale.size(0) == outer_dim && scale.size(1) ==
+ * 1:
+ *   - Returns RowWise.
+ *
+ * - Otherwise:
+ *   - Returns Error.
+ */
+
+bool is_tensorwise_scaling(const at::Tensor& t, const at::Tensor& scale) {
+  return at::isFloat8Type(t.scalar_type()) &&
+      scale.scalar_type() == at::kFloat && scale.numel() == 1;
+}
+
+bool is_rowwise_scaling(const at::Tensor& t, const at::Tensor& scale) {
+  return (
+      at::isFloat8Type(t.scalar_type()) && scale.scalar_type() == at::kFloat &&
+      scale.dim() == 2 && scale.size(0) == t.size(0) && scale.size(1) == 1 &&
+      scale.is_contiguous());
+}
+
+bool is_desired_scaling(
+    const at::Tensor& t,
+    const at::Tensor& scale,
+    ScalingType desired_scaling) {
+  auto result = desired_scaling == ScalingType::TensorWise
+      ? is_tensorwise_scaling(t, scale)
+      : is_rowwise_scaling(t, scale);
+  return result;
+}
+
+std::pair<ScalingType, ScalingType> get_joint_scaling(
+    std::initializer_list<std::pair<ScalingType, ScalingType>> options,
+    const at::Tensor& mat_a,
+    const at::Tensor& mat_b,
+    const at::Tensor& scale_a,
+    const at::Tensor& scale_b) {
+  for (auto [lhs, rhs] : options) {
+    if (is_desired_scaling(mat_a, scale_a, lhs) &&
+        is_desired_scaling(mat_b.t(), scale_b.t(), rhs)) {
+      return {lhs, rhs};
+    }
+  }
+  invalid_scaling_config(mat_a, mat_b, scale_a, scale_b);
+  return {};
+}
+
+} // namespace
+
+static Tensor&
+_scaled_mm_out_cpu_emulated(const Tensor& mat1, const Tensor& mat2,
+          const Tensor& scale_a,
+          const Tensor& scale_b,
+          const std::optional<at::Tensor>& bias,
+          const std::optional<at::Tensor>& scale_result,
+          std::optional<c10::ScalarType> out_dtype,
+          bool use_fast_accum,
+          const scaled::ScaledGemmImplementation gemm_impl,
+          Tensor& out) {
+  TORCH_CHECK(mat1.dim() == 2, "mat1 must be a matrix");
+  TORCH_CHECK(mat2.dim() == 2, "mat2 must be a matrix");
+  TORCH_CHECK(
+      mat1.sizes()[1] == mat2.sizes()[0], "mat_a and mat_b shapes cannot be multiplied (",
+      mat1.sizes()[0], "x", mat1.sizes()[1], " and ", mat2.sizes()[0], "x", mat2.sizes()[1], ")");
+
+  TORCH_CHECK_VALUE(gemm_impl == scaled::ScaledGemmImplementation::TENSORWISE_TENSORWISE ||
+    gemm_impl == scaled::ScaledGemmImplementation::ROWWISE_ROWWISE, "Unsupported scaling implementation");
+
+  if (gemm_impl == scaled::ScaledGemmImplementation::TENSORWISE_TENSORWISE) {
+    // Restrictions:
+    // A, B are FP8, scales are fp32
+    TORCH_CHECK_VALUE(
+    isFloat8Type(mat1.scalar_type()) && isFloat8Type(mat2.scalar_type()),
+    "mat1 and mat2 must be fp8 types, got: ",
+    mat1.scalar_type(),
+    mat2.scalar_type());
+    TORCH_CHECK_VALUE(
+        scale_a.numel() == 1 && scale_a.scalar_type() == kFloat,
+        "scale_a must have 1 Float element");
+    TORCH_CHECK_VALUE(
+        scale_b.numel() == 1 && scale_b.scalar_type() == kFloat,
+        "scale_b must have 1 Float element");
+  }
+  else {
+    // Restrictions:
+    // A, B are FP8, scales are fp32, shape M/N for A/B
+    TORCH_CHECK_VALUE(
+        isFloat8Type(mat1.scalar_type()) && isFloat8Type(mat2.scalar_type()),
+        "mat1 and mat_b must be fp8 types, got: ",
+        mat1.scalar_type(),
+        mat2.scalar_type());
+    TORCH_CHECK_VALUE(
+        scale_a.size(0) == mat1.size(0) && scale_a.size(1) == 1,
+        "scale_a must have shape [",
+        mat1.size(0),
+        ", 1], got [",
+        scale_a.sizes(),
+        "]");
+    TORCH_CHECK_VALUE(
+        scale_a.numel() == mat1.size(0) && scale_a.scalar_type() == kFloat,
+        "scale_a must have ",
+        mat1.size(0),
+        " Float elements, got ",
+        scale_a.numel());
+    TORCH_CHECK_VALUE(
+        scale_b.numel() == mat2.size(1) && scale_b.scalar_type() == kFloat,
+        "scale_b must have ",
+        mat2.size(1),
+        " Float elements, got ",
+        scale_b.numel());
+
+    TORCH_CHECK_VALUE(
+        scale_a.stride(1) == 1,
+        "expected scale_a.stride(1) to be 1, but got ",
+        scale_a.stride(1));
+    TORCH_CHECK_VALUE(
+        scale_b.stride(1) == 1,
+        "expected scale_b.stride(1) to be 1, but got ",
+        scale_b.stride(1));
+  }
+
+  TORCH_CHECK(
+      !scale_result ||
+          (scale_result->numel() == 1 && scale_result->scalar_type() == kFloat),
+      "scale_result must be a float scalar");
+  TORCH_CHECK(!bias || bias->numel() == mat2.sizes()[1], "Bias must be size ", mat2.sizes()[1],
+       " but got ", bias->numel());
+
+  // Check types
+  TORCH_CHECK(isFloat8Type(mat1.scalar_type()), "Expected mat1 to be Float8 matrix got ", mat1.scalar_type());
+  TORCH_CHECK(isFloat8Type(mat2.scalar_type()), "Expected mat2 to be Float8 matrix got ", mat2.scalar_type());
+
+  auto mat1_cont = mat1.contiguous();
+  auto mat2_cont = mat2.contiguous();
+  IntArrayRef mat1_sizes = mat1_cont.sizes();
+  IntArrayRef mat2_sizes = mat2_cont.sizes();
+  at::native::resize_output(out, {mat1_sizes[0], mat2_sizes[1]});
+
+  const auto out_dtype_ = out_dtype.value_or(c10::ScalarType::BFloat16);
+
+  float output_scale = 1.0f;
+  if (scale_result.has_value() &&
+      (out_dtype_ == ScalarType::Float8_e4m3fn ||
+       out_dtype_ == ScalarType::Float8_e5m2)) {
+    output_scale = scale_result.value().item<float>();
+  }
+
+  at::Tensor fp32_mat_a;
+  at::Tensor fp32_mat_b;
+  if (gemm_impl == scaled::ScaledGemmImplementation::TENSORWISE_TENSORWISE) {
+    fp32_mat_a = at::mul(mat1_cont.to(kFloat), scale_a.item<float>());
+    fp32_mat_b = at::mul(mat2_cont.to(kFloat), scale_b.item<float>());
+  }
+  else {
+    fp32_mat_a = at::mul(mat1_cont.to(kFloat), scale_a);
+    fp32_mat_b = at::mul(mat2_cont.to(kFloat), scale_b);
+  }
+
+  auto out_tmp = at::matmul(fp32_mat_a, fp32_mat_b);
+  if (bias) {
+    out_tmp.add_(bias.value());
+  }
+  if (out_dtype_ == ScalarType::Float8_e4m3fn ||
+      out_dtype_ == ScalarType::Float8_e5m2) {
+    out_tmp = at::mul(out_tmp, 1 / output_scale);
+  }
+  out_tmp = out_tmp.to(out.scalar_type());
+  out.copy_(out_tmp);
+  return out;
+}
+
+Tensor&
+_scaled_mm_out_cpu(const Tensor& mat1, const Tensor& mat2,
+          const Tensor& scale_a,
+          const Tensor& scale_b,
+          const std::optional<at::Tensor>& bias,
+          const std::optional<at::Tensor>& scale_result,
+          std::optional<c10::ScalarType> out_dtype,
+          bool use_fast_accum,
+          Tensor& out) {
+#if AT_MKLDNN_ENABLED() && !defined(__powerpc__)
+  if (at::globalContext().userEnabledMkldnn() && scale_a.numel() == 1 && scale_b.numel() == 1) {
+    bool mixed_dtype = mat1.scalar_type() != mat2.scalar_type();
+    if ((!mixed_dtype && cpuinfo_has_x86_amx_int8()) ||
+        (mixed_dtype && cpuinfo_has_x86_amx_fp16())) {
+      return mkldnn_scaled_mm(
+          mat1,
+          mat2,
+          scale_a,
+          scale_b,
+          bias,
+          scale_result,
+          out_dtype,
+          use_fast_accum,
+          out);
+    }
+  }
+#endif
+  {
+    TORCH_CHECK_VALUE(
+      !out_dtype || *out_dtype == out.scalar_type(),
+      "out_dtype must match output matrix type");
+
+    const auto [scaling_choice_a, scaling_choice_b] = get_joint_scaling(
+        {
+            std::make_pair(ScalingType::TensorWise, ScalingType::TensorWise),
+            std::make_pair(ScalingType::RowWise, ScalingType::RowWise),
+        },
+        mat1,
+        mat2,
+        scale_a,
+        scale_b);
+
+    scaled::ScaledGemmImplementation gemm_impl{scaled::ScaledGemmImplementation::NONE};
+
+    if (scaling_choice_a == ScalingType::TensorWise && scaling_choice_b == ScalingType::TensorWise) {
+      gemm_impl = scaled::ScaledGemmImplementation::TENSORWISE_TENSORWISE;
+    }
+    else if (scaling_choice_a == ScalingType::RowWise && scaling_choice_b == ScalingType::RowWise) {
+      gemm_impl = scaled::ScaledGemmImplementation::ROWWISE_ROWWISE;
+    }
+
+  return _scaled_mm_out_cpu_emulated(mat1, mat2, scale_a, scale_b, bias, scale_result, out_dtype, use_fast_accum, gemm_impl, out);
+  }
+}
+
+Tensor
+_scaled_mm_cpu(const Tensor& mat_a, const Tensor& mat_b,
+          const Tensor& scale_a,
+          const Tensor& scale_b,
+          const std::optional<at::Tensor>& bias,
+          const std::optional<at::Tensor>& scale_result,
+          std::optional<c10::ScalarType> out_dtype,
+          bool use_fast_accum) {
+  const auto out_dtype_ = out_dtype.value_or(mat_a.scalar_type());
+  Tensor out = at::empty({0}, mat_a.options().dtype(out_dtype_));
+  return _scaled_mm_out_cpu(mat_a, mat_b, scale_a, scale_b, bias, scale_result, out_dtype, use_fast_accum, out);
+}
+
+using acceptance_fn = std::function<bool(
+    c10::ScalarType,
+    std::vector<ScalingType>&,
+    ArrayRef<Tensor>&,
+    c10::ScalarType,
+    std::vector<ScalingType>&,
+    ArrayRef<Tensor>&)>;
+
+namespace scaled_blas = at::native::scaled;
+using scaled_blas::convert_int_to_enum;
+using scaled_blas::ScaledGemmImplementation;
+
+std::array<std::tuple<std::string, acceptance_fn, ScaledGemmImplementation>, 2>
+    scale_kernel_dispatch = {{
+      {"tensorwise_tensorwise",
+       scaled_blas::check_tensorwise_recipe,
+       ScaledGemmImplementation::TENSORWISE_TENSORWISE},
+      {"rowwise_rowwise",
+       scaled_blas::check_rowwise_recipe,
+       ScaledGemmImplementation::ROWWISE_ROWWISE},
+
+  }};
+
+Tensor& _scaled_mm_cpu_v2_out(
+    const Tensor& mat_a,
+    const Tensor& mat_b,
+    ArrayRef<Tensor> scale_a,
+    IntArrayRef scale_recipe_a,
+    IntArrayRef swizzle_a,
+    ArrayRef<Tensor> scale_b,
+    IntArrayRef scale_recipe_b,
+    IntArrayRef swizzle_b,
+    const std::optional<Tensor>& bias,
+    const std::optional<c10::ScalarType> out_dtype,
+    IntArrayRef contraction_dim,
+    bool use_fast_accum,
+    Tensor& out) {
+  TORCH_CHECK_VALUE(mat_a.dim() == 2, "mat_a must be a matrix");
+  TORCH_CHECK_VALUE(mat_b.dim() == 2, "mat_b must be a matrix");
+
+  // If any of M, K, N is 0 - return early (the tensorwise/rowwise float8 gemm kernels
+  // do not support this case).
+  if (mat_a.size(0) == 0 || mat_a.size(1) == 0 || mat_b.size(1) == 0) {
+    // `out` was created with `at::empty`. In the case where we are multiplying
+    // MxK by KxN and K is the zero dim, we need to initialize here to properly
+    // return a tensor of zeros.
+    at::native::resize_output(out, {mat_a.size(0), mat_b.size(1)});
+    if (mat_a.size(1) == 0) {
+      out.zero_();
+    }
+
+    return out;
+  }
+
+  // Check if the input matrix sizes can be multiplied
+  // - if optional contraction dims are provided, use those
+  //   -- mostly for < 1B formats (i.e. nvfp4x2) where cheap .t() is not available.
+  if (contraction_dim.size() > 0) {
+    TORCH_CHECK_VALUE(contraction_dim.size() == 2, "contraction_dim must have exactly 2 elements");
+    auto mat_a_dim = contraction_dim[0];
+    auto mat_b_dim = contraction_dim[1];
+    TORCH_CHECK_VALUE(
+        mat_a.size(mat_a_dim) == mat_b.size(mat_b_dim), "mat_a and mat_b shapes cannot be multiplied (",
+        mat_a.size(0), "x", mat_a.size(1), " and ", mat_b.size(0), "x", mat_b.size(1), ") ",
+        "with contraction dims mat_a: ", mat_a_dim, ", mat_b: ", mat_b_dim);
+  } else {
+    TORCH_CHECK_VALUE(
+        mat_a.size(1) == mat_b.size(0), "mat_a and mat_b shapes cannot be multiplied (",
+        mat_a.size(0), "x", mat_a.size(1), " and ", mat_b.size(0), "x", mat_b.size(1), ")");
+  }
+
+  TORCH_CHECK_VALUE(
+    !bias || bias->numel() == mat_b.sizes()[1],
+    "Bias must be size ",
+    mat_b.sizes()[1],
+    " but got ",
+    bias->numel());
+
+  TORCH_CHECK_VALUE(
+    !out_dtype || *out_dtype == out.scalar_type(),
+    "out_dtype must match output matrix type");
+
+  if (bias) {
+    TORCH_CHECK_VALUE(
+        bias->scalar_type() == kFloat ||
+            bias->scalar_type() == c10::ScalarType::BFloat16 ||
+            bias->scalar_type() == c10::ScalarType::Half,
+        "Bias must be Float32 or BFloat16 or Half, but got ",
+        bias->scalar_type());
+  }
+
+  // Align with CUDA's default out to be bf16
+  const auto out_dtype_ = out_dtype.value_or(c10::ScalarType::BFloat16);
+
+  // Conversion of implicitly-defined enums to explicit
+  auto scale_recipe_a_enum = convert_int_to_enum<ScalingType>(scale_recipe_a);
+  auto swizzle_a_enum = convert_int_to_enum<SwizzleType>(swizzle_a);
+  auto scale_recipe_b_enum = convert_int_to_enum<ScalingType>(scale_recipe_b);
+  auto swizzle_b_enum = convert_int_to_enum<SwizzleType>(swizzle_b);
+
+  if (!swizzle_a_enum.empty() && !swizzle_b_enum.empty()) {
+    TORCH_CHECK_VALUE(
+      swizzle_a_enum[0] == at::blas::SwizzleType::NO_SWIZZLE &&
+          swizzle_b_enum[0] == at::blas::SwizzleType::NO_SWIZZLE,
+      "CPU does not support swizzle.");
+  }
+
+  // at this point we can start working out what we want to be doing
+  // Try to do as few steps as possible.
+  // NOTE: support is deliberately sparse, can explicitly enumerate all
+  // combinations allowed. Do this via a list of defined (name, acceptance,
+  // concrete_impl) tuples.
+  bool found_impl = false;
+  ScaledGemmImplementation gemm_impl = ScaledGemmImplementation::NONE;
+
+  for (const auto& fn_entry : scale_kernel_dispatch) {
+    const auto [name, accept_fn, scaled_gemm_impl] = fn_entry;
+    const bool ok = accept_fn(
+        mat_a.scalar_type(),
+        scale_recipe_a_enum,
+        scale_a,
+        mat_b.scalar_type(),
+        scale_recipe_b_enum,
+        scale_b);
+
+    if (ok) {
+      gemm_impl = scaled_gemm_impl;
+      found_impl = true;
+      break;
+    }
+  }
+
+  if (!found_impl) {
+    const std::optional<at::Tensor> scale_a_opt = scale_a.empty() ? std::optional<at::Tensor>{std::nullopt} : std::optional<at::Tensor>{scale_a[0]};
+    const std::optional<at::Tensor> scale_b_opt = scale_b.empty() ? std::optional<at::Tensor>{std::nullopt} : std::optional<at::Tensor>{scale_b[0]};
+
+    invalid_scaling_config(mat_a, mat_b, scale_a_opt, scale_b_opt);
+  }
+
+  at::native::resize_output(out, {mat_a.size(0), mat_b.size(1)});
+
+  auto bias_ = bias.value_or(Tensor());
+
+  if (gemm_impl == ScaledGemmImplementation::TENSORWISE_TENSORWISE ||
+      gemm_impl == ScaledGemmImplementation::ROWWISE_ROWWISE) {
+    _scaled_mm_out_cpu_emulated(
+      mat_a,
+      mat_b,
+      scale_a[0],
+      scale_b[0],
+      bias,
+      std::nullopt,  // scale-result
+      out_dtype_,
+      use_fast_accum,
+      gemm_impl,
+      out);
+  } else {
+    TORCH_CHECK_VALUE(false, "Invalid state - found an implementation, but not really");
+  }
+
+  return out;
+}
+
+Tensor _scaled_mm_cpu_v2(
+    const Tensor& mat_a,
+    const Tensor& mat_b,
+    ArrayRef<Tensor> scale_a,
+    IntArrayRef scale_recipe_a,
+    IntArrayRef swizzle_a,
+    ArrayRef<Tensor> scale_b,
+    IntArrayRef scale_recipe_b,
+    IntArrayRef swizzle_b,
+    const std::optional<Tensor>& bias,
+    const std::optional<c10::ScalarType> out_dtype,
+    IntArrayRef contraction_dim,
+    bool use_fast_accum) {
+  const auto out_dtype_ = out_dtype.value_or(mat_a.scalar_type());
+  Tensor out = at::empty({0}, mat_a.options().dtype(out_dtype_));
+
+  return _scaled_mm_cpu_v2_out(
+    mat_a,
+    mat_b,
+    scale_a,
+    scale_recipe_a,
+    swizzle_a,
+    scale_b,
+    scale_recipe_b,
+    swizzle_b,
+    bias,
+    out_dtype,
+    contraction_dim,
+    use_fast_accum,
+    out);
+}
+
+// TODO(vasiliy, future PR): figure out why we need to declare this function, when
+// other functions that live in ATen/native/*.cpp without declarations
+// or headers work just fine.
+Tensor _grouped_mm(const Tensor& mat_a, const Tensor& mat_b,
+const std::optional<at::Tensor>& offs,
+const std::optional<at::Tensor>& bias,
+std::optional<c10::ScalarType> out_dtype);
+
+Tensor _grouped_mm(const Tensor& mat_a, const Tensor& mat_b,
+const std::optional<at::Tensor>& offs,
+const std::optional<at::Tensor>& bias,
+std::optional<c10::ScalarType> out_dtype) {
+  _grouped_mm_validate_inputs(mat_a, mat_b, offs, bias, out_dtype);
+  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_);
+  _grouped_mm_fallback(mat_a, mat_b, offs, bias, out_dtype, out);
+  return out;
+}
+
+}  // namespace at::native
diff --git a/aten/src/ATen/cuda/CUDAScaledBlas.cpp b/aten/src/ATen/native/ScaledBlasUtils.cpp
similarity index 70%
rename from aten/src/ATen/cuda/CUDAScaledBlas.cpp
rename to aten/src/ATen/native/ScaledBlasUtils.cpp
index e7903780457b8..4cdd796638d9c 100644
--- a/aten/src/ATen/cuda/CUDAScaledBlas.cpp
+++ b/aten/src/ATen/native/ScaledBlasUtils.cpp
@@ -5,27 +5,20 @@
 #include <c10/core/Scalar.h>
 #include <c10/core/ScalarType.h>
 #define TORCH_ASSERT_ONLY_METHOD_OPERATORS
+
+#include <ATen/BlasBackend.h>
 #include <ATen/core/Tensor.h>
 #include <ATen/core/NamedTensor.h>
 #include <ATen/Dispatch.h>
 #include <ATen/ExpandUtils.h>
 #include <ATen/OpMathType.h>
 #include <ATen/TensorUtils.h>
-#include <ATen/cuda/CUDABlas.h>
-#include <ATen/cuda/tunable/Tunable.h>
-#include <ATen/cuda/tunable/TunableGemm.h>
 #include <ATen/native/Resize.h>
 #include <c10/util/MaybeOwned.h>
 #include <ATen/native/GroupedMMUtils.h>
-#include <ATen/native/cuda/RowwiseScaledMM.h>
-#include <ATen/native/cuda/ScaledGroupMM.h>
-#include <ATen/native/cuda/GroupMM.h>
+#include <ATen/native/ScaledBlasUtils.h>
 #include <ATen/ceil_div.h>
 
-#ifdef USE_MSLK
-#include <mslk/gemm/gemm_torch.h>
-#endif
-
 #ifndef AT_PER_OPERATOR_HEADERS
 #include <ATen/Functions.h>
 #include <ATen/NativeFunctions.h>
@@ -56,18 +49,19 @@
 using at::blas::ScalingType;
 using at::blas::SwizzleType;
 
-namespace at::cuda::scaled {
+namespace at::native::scaled {
 
 /**
  * Both inputs must be fp8,
  * Each needs a single scale, {Tensorwise (float)}
  */
-bool check_tensorwise_recipe(c10::ScalarType type_a,
-                             std::vector<ScalingType>& recipe_a,
-                             ArrayRef<Tensor>& scales_a,
-                             c10::ScalarType type_b,
-                             std::vector<ScalingType>& recipe_b,
-                             ArrayRef<Tensor>& scales_b) {
+bool check_tensorwise_recipe(
+    c10::ScalarType type_a,
+    std::vector<ScalingType>& recipe_a,
+    c10::ArrayRef<Tensor>& scales_a,
+    c10::ScalarType type_b,
+    std::vector<ScalingType>& recipe_b,
+    c10::ArrayRef<Tensor>& scales_b) {
   // both types must be fp8
   if (!isFloat8Type(type_a) || !isFloat8Type(type_b)) {
     return false;
@@ -90,12 +84,13 @@ bool check_tensorwise_recipe(c10::ScalarType type_a,
  * Both inputs must be fp8,
  * Each needs scales, {Rowwise (float)}
  */
-bool check_rowwise_recipe(c10::ScalarType type_a,
-                             std::vector<ScalingType>& recipe_a,
-                             ArrayRef<Tensor>& scales_a,
-                             c10::ScalarType type_b,
-                             std::vector<ScalingType>& recipe_b,
-                             ArrayRef<Tensor>& scales_b) {
+bool check_rowwise_recipe(
+    c10::ScalarType type_a,
+    std::vector<ScalingType>& recipe_a,
+    ArrayRef<Tensor>& scales_a,
+    c10::ScalarType type_b,
+    std::vector<ScalingType>& recipe_b,
+    ArrayRef<Tensor>& scales_b) {
   // both types must be fp8
   if (!isFloat8Type(type_a) || !isFloat8Type(type_b)) {
     return false;
@@ -115,18 +110,18 @@ bool check_rowwise_recipe(c10::ScalarType type_a,
   return true;
 }
 
-
 /**
  * Two-level scaling, canonical NVFP4
  * Both inputs must be fp4
  * A, B need 2 scales, {Blockwise_1x16 (e4m3), Tensorwise (fp32)}
  */
-bool check_nvfp4_recipe(c10::ScalarType type_a,
-                        std::vector<ScalingType>& recipe_a,
-                        ArrayRef<Tensor>& scales_a,
-                        c10::ScalarType type_b,
-                        std::vector<ScalingType>& recipe_b,
-                        ArrayRef<Tensor>& scales_b) {
+bool check_nvfp4_recipe(
+    c10::ScalarType type_a,
+    std::vector<ScalingType>& recipe_a,
+    ArrayRef<Tensor>& scales_a,
+    c10::ScalarType type_b,
+    std::vector<ScalingType>& recipe_b,
+    ArrayRef<Tensor>& scales_b) {
   // both types must be fp4
   if (type_a != ScalarType::Float4_e2m1fn_x2 || type_b != ScalarType::Float4_e2m1fn_x2) {
     return false;
@@ -151,13 +146,13 @@ bool check_nvfp4_recipe(c10::ScalarType type_a,
  * Both inputs must be fp4
  * A, B need 1 scale, {Blockwise_1x16 (e4m3)}
  */
-bool check_nvfp4_recipe_single_scale
-                       (c10::ScalarType type_a,
-                        std::vector<ScalingType>& recipe_a,
-                        ArrayRef<Tensor>& scales_a,
-                        c10::ScalarType type_b,
-                        std::vector<ScalingType>& recipe_b,
-                        ArrayRef<Tensor>& scales_b) {
+bool check_nvfp4_recipe_single_scale(
+    c10::ScalarType type_a,
+    std::vector<ScalingType>& recipe_a,
+    ArrayRef<Tensor>& scales_a,
+    c10::ScalarType type_b,
+    std::vector<ScalingType>& recipe_b,
+    ArrayRef<Tensor>& scales_b) {
   // both types must be fp4
   if (type_a != ScalarType::Float4_e2m1fn_x2 || type_b != ScalarType::Float4_e2m1fn_x2) {
     return false;
@@ -179,16 +174,18 @@ bool check_nvfp4_recipe_single_scale
 
 /**
  * Both inputs must be fp8
- * A, B must only have 1 scale each, A: {Blockwise_1x128 (float), B: {Blockwise_128x128 (float)
+ * A, B must only have 1 scale each, A: {Blockwise_1x128 (float), B:
+ * {Blockwise_128x128 (float)
  */
-bool check_deepseek_recipe(ScalingType expected_recipe_a,
-                           ScalingType expected_recipe_b,
-                           c10::ScalarType type_a,
-                           std::vector<ScalingType>& recipe_a,
-                           ArrayRef<Tensor>& scales_a,
-                           c10::ScalarType type_b,
-                           std::vector<ScalingType>& recipe_b,
-                           ArrayRef<Tensor>& scales_b) {
+bool check_deepseek_recipe(
+    ScalingType expected_recipe_a,
+    ScalingType expected_recipe_b,
+    c10::ScalarType type_a,
+    std::vector<ScalingType>& recipe_a,
+    ArrayRef<Tensor>& scales_a,
+    c10::ScalarType type_b,
+    std::vector<ScalingType>& recipe_b,
+    ArrayRef<Tensor>& scales_b) {
   // both types must be fp8
   if (type_a != ScalarType::Float8_e4m3fn || type_b != ScalarType::Float8_e4m3fn) {
     return false;
@@ -212,12 +209,13 @@ bool check_deepseek_recipe(ScalingType expected_recipe_a,
  * Both inputs must be fp8
  * A, B must have 1 scale each, {Blockwise_1x32, e8m0}
  */
-bool check_mxfp8_recipe(c10::ScalarType type_a,
-                        std::vector<ScalingType>& recipe_a,
-                        ArrayRef<Tensor>& scales_a,
-                        c10::ScalarType type_b,
-                        std::vector<ScalingType>& recipe_b,
-                        ArrayRef<Tensor>& scales_b) {
+bool check_mxfp8_recipe(
+    c10::ScalarType type_a,
+    std::vector<ScalingType>& recipe_a,
+    ArrayRef<Tensor>& scales_a,
+    c10::ScalarType type_b,
+    std::vector<ScalingType>& recipe_b,
+    ArrayRef<Tensor>& scales_b) {
   // both types must be fp8
   if (type_a != ScalarType::Float8_e4m3fn || type_b != ScalarType::Float8_e4m3fn) {
     return false;
@@ -241,12 +239,13 @@ bool check_mxfp8_recipe(c10::ScalarType type_a,
  * Both inputs must be fp4
  * A, B must have 1 scale each, {Blockwise_1x32, e8m0}
  */
-bool check_mxfp4_recipe(c10::ScalarType type_a,
-                        std::vector<ScalingType>& recipe_a,
-                        ArrayRef<Tensor>& scales_a,
-                        c10::ScalarType type_b,
-                        std::vector<ScalingType>& recipe_b,
-                        ArrayRef<Tensor>& scales_b) {
+bool check_mxfp4_recipe(
+    c10::ScalarType type_a,
+    std::vector<ScalingType>& recipe_a,
+    ArrayRef<Tensor>& scales_a,
+    c10::ScalarType type_b,
+    std::vector<ScalingType>& recipe_b,
+    ArrayRef<Tensor>& scales_b) {
   // both types must be fp4
   if (type_a != ScalarType::Float4_e2m1fn_x2 || type_b != ScalarType::Float4_e2m1fn_x2) {
     return false;
@@ -266,4 +265,4 @@ bool check_mxfp4_recipe(c10::ScalarType type_a,
   return true;
 }
 
-} // namespace at::native::cuda::blas::scaled
+}  // at::scaled
diff --git a/aten/src/ATen/native/ScaledBlasUtils.h b/aten/src/ATen/native/ScaledBlasUtils.h
new file mode 100644
index 0000000000000..7ab0c357c8660
--- /dev/null
+++ b/aten/src/ATen/native/ScaledBlasUtils.h
@@ -0,0 +1,138 @@
+#pragma once
+
+#include <ATen/BlasBackend.h>
+#include <ATen/core/Tensor.h>
+
+using at::blas::ScalingType;
+using at::blas::SwizzleType;
+
+namespace at::native::scaled {
+
+/**
+ * Track concrete implementations available
+ */
+enum class ScaledGemmImplementation {
+  NONE = 0,
+  TENSORWISE_TENSORWISE = 1,
+  ROWWISE_ROWWISE = 2,
+  BLOCK_128x128_1x128 = 3,
+  BLOCK_1x128_128x128 = 4,
+  BLOCK_1x128_1x128 = 5,
+  MXFP8_MXFP8 = 6,
+  NVFP4_NVFP4 = 7,
+  NVFP4_NVFP4_SINGLE_SCALE = 8,
+  MXFP4_MXFP4 = 9,
+};
+
+/**
+ * Convert passed int (enum) from python back into a
+ * strictly-typed enum
+ */
+template <class EnumType, class ArrayType>
+std::vector<EnumType> convert_int_to_enum(ArrayType& v) {
+  std::vector<EnumType> converted;
+  converted.reserve(v.size());
+
+  for (auto vi : v) {
+    converted.push_back(static_cast<EnumType>(vi));
+  }
+  return converted;
+}
+
+/**
+ * Both inputs must be fp8,
+ * Each needs a single scale, {Tensorwise (float)}
+ */
+TORCH_API
+bool check_tensorwise_recipe(
+    c10::ScalarType type_a,
+    std::vector<ScalingType>& recipe_a,
+    ArrayRef<Tensor>& scales_a,
+    c10::ScalarType type_b,
+    std::vector<ScalingType>& recipe_b,
+    ArrayRef<Tensor>& scales_b);
+
+/**
+ * Both inputs must be fp8,
+ * Each needs scales, {Rowwise (float)}
+ */
+TORCH_API
+bool check_rowwise_recipe(
+    c10::ScalarType type_a,
+    std::vector<ScalingType>& recipe_a,
+    ArrayRef<Tensor>& scales_a,
+    c10::ScalarType type_b,
+    std::vector<ScalingType>& recipe_b,
+    ArrayRef<Tensor>& scales_b);
+
+/**
+ * Two-level scaling, canonical NVFP4
+ * Both inputs must be fp4
+ * A, B need 2 scales, {Blockwise_1x16 (e4m3), Tensorwise (fp32)}
+ */
+TORCH_API
+bool check_nvfp4_recipe(
+    c10::ScalarType type_a,
+    std::vector<ScalingType>& recipe_a,
+    ArrayRef<Tensor>& scales_a,
+    c10::ScalarType type_b,
+    std::vector<ScalingType>& recipe_b,
+    ArrayRef<Tensor>& scales_b);
+
+/**
+ * Single-level scaling, what PyT currently understands
+ * Both inputs must be fp4
+ * A, B need 1 scale, {Blockwise_1x16 (e4m3)}
+ */
+TORCH_API
+bool check_nvfp4_recipe_single_scale(
+    c10::ScalarType type_a,
+    std::vector<ScalingType>& recipe_a,
+    ArrayRef<Tensor>& scales_a,
+    c10::ScalarType type_b,
+    std::vector<ScalingType>& recipe_b,
+    ArrayRef<Tensor>& scales_b);
+
+/**
+ * Both inputs must be fp8
+ * A, B must only have 1 scale each, A: {Blockwise_1x128 (float), B:
+ * {Blockwise_128x128 (float)
+ */
+TORCH_API
+bool check_deepseek_recipe(
+    ScalingType expected_recipe_a,
+    ScalingType expected_recipe_b,
+    c10::ScalarType type_a,
+    std::vector<ScalingType>& recipe_a,
+    ArrayRef<Tensor>& scales_a,
+    c10::ScalarType type_b,
+    std::vector<ScalingType>& recipe_b,
+    ArrayRef<Tensor>& scales_b);
+
+/**
+ * Both inputs must be fp8
+ * A, B must have 1 scale each, {Blockwise_1x32, e8m0}
+ */
+TORCH_API
+bool check_mxfp8_recipe(
+    c10::ScalarType type_a,
+    std::vector<ScalingType>& recipe_a,
+    ArrayRef<Tensor>& scales_a,
+    c10::ScalarType type_b,
+    std::vector<ScalingType>& recipe_b,
+    ArrayRef<Tensor>& scales_b);
+
+/**
+ * Both inputs must be fp4
+ * A, B must have 1 scale each, {Blockwise_1x32, e8m0}
+ */
+TORCH_API
+bool check_mxfp4_recipe(
+    c10::ScalarType type_a,
+    std::vector<ScalingType>& recipe_a,
+    ArrayRef<Tensor>& scales_a,
+    c10::ScalarType type_b,
+    std::vector<ScalingType>& recipe_b,
+    ArrayRef<Tensor>& scales_b);
+
+} // namespace at::scaled
diff --git a/aten/src/ATen/native/SegmentReduce.cpp b/aten/src/ATen/native/SegmentReduce.cpp
index 2b61bcec6a828..7254c581e2947 100644
--- a/aten/src/ATen/native/SegmentReduce.cpp
+++ b/aten/src/ATen/native/SegmentReduce.cpp
@@ -458,26 +458,8 @@ Tensor segment_reduce_kernel(
   }
 }
 
-REGISTER_ARCH_DISPATCH(
-    _segment_reduce_lengths_stub,
-    DEFAULT,
-    &_segment_reduce_lengths_cpu_kernel)
-REGISTER_AVX2_DISPATCH(_segment_reduce_lengths_stub, &_segment_reduce_lengths_cpu_kernel)
-REGISTER_AVX512_DISPATCH(_segment_reduce_lengths_stub, &_segment_reduce_lengths_cpu_kernel)
-REGISTER_VSX_DISPATCH(_segment_reduce_lengths_stub, &_segment_reduce_lengths_cpu_kernel)
-REGISTER_ZVECTOR_DISPATCH(_segment_reduce_lengths_stub, &_segment_reduce_lengths_cpu_kernel)
-REGISTER_SVE256_DISPATCH(_segment_reduce_lengths_stub, &_segment_reduce_lengths_cpu_kernel)
-
-// offsets dispatches
-REGISTER_ARCH_DISPATCH(
-    _segment_reduce_offsets_stub,
-    DEFAULT,
-    &_segment_reduce_offsets_cpu_kernel)
-REGISTER_AVX2_DISPATCH(_segment_reduce_offsets_stub, &_segment_reduce_offsets_cpu_kernel)
-REGISTER_AVX512_DISPATCH(_segment_reduce_offsets_stub, &_segment_reduce_offsets_cpu_kernel)
-REGISTER_VSX_DISPATCH(_segment_reduce_offsets_stub, &_segment_reduce_offsets_cpu_kernel)
-REGISTER_ZVECTOR_DISPATCH(_segment_reduce_offsets_stub, &_segment_reduce_offsets_cpu_kernel)
-REGISTER_SVE256_DISPATCH(_segment_reduce_offsets_stub, &_segment_reduce_offsets_cpu_kernel)
+REGISTER_ALL_CPU_DISPATCH(_segment_reduce_lengths_stub, &_segment_reduce_lengths_cpu_kernel)
+REGISTER_ALL_CPU_DISPATCH(_segment_reduce_offsets_stub, &_segment_reduce_offsets_cpu_kernel)
 
 // Currently some computation is being duplicated across forward and backward.
 // TODO: Cache indices in forward pass to reuse in backward
@@ -532,44 +514,7 @@ Tensor _segment_reduce_backward_kernel(
   }
 }
 
-REGISTER_ARCH_DISPATCH(
-    _segment_reduce_lengths_backward_stub,
-    DEFAULT,
-    &_segment_reduce_cpu_lengths_backward_kernel)
-REGISTER_AVX512_DISPATCH(
-    _segment_reduce_lengths_backward_stub,
-    &_segment_reduce_cpu_lengths_backward_kernel)
-REGISTER_AVX2_DISPATCH(
-    _segment_reduce_lengths_backward_stub,
-    &_segment_reduce_cpu_lengths_backward_kernel)
-REGISTER_VSX_DISPATCH(
-    _segment_reduce_lengths_backward_stub,
-    &_segment_reduce_cpu_lengths_backward_kernel)
-REGISTER_ZVECTOR_DISPATCH(
-    _segment_reduce_lengths_backward_stub,
-    &_segment_reduce_cpu_lengths_backward_kernel)
-REGISTER_SVE256_DISPATCH(
-    _segment_reduce_lengths_backward_stub,
-    &_segment_reduce_cpu_lengths_backward_kernel)
-
-REGISTER_ARCH_DISPATCH(
-    _segment_reduce_offsets_backward_stub,
-    DEFAULT,
-    &_segment_reduce_cpu_offsets_backward_kernel)
-REGISTER_AVX512_DISPATCH(
-    _segment_reduce_offsets_backward_stub,
-    &_segment_reduce_cpu_offsets_backward_kernel)
-REGISTER_AVX2_DISPATCH(
-    _segment_reduce_offsets_backward_stub,
-    &_segment_reduce_cpu_offsets_backward_kernel)
-REGISTER_VSX_DISPATCH(
-    _segment_reduce_offsets_backward_stub,
-    &_segment_reduce_cpu_offsets_backward_kernel)
-REGISTER_ZVECTOR_DISPATCH(
-    _segment_reduce_offsets_backward_stub,
-    &_segment_reduce_cpu_offsets_backward_kernel)
-REGISTER_SVE256_DISPATCH(
-    _segment_reduce_offsets_backward_stub,
-    &_segment_reduce_cpu_offsets_backward_kernel)
+REGISTER_ALL_CPU_DISPATCH(_segment_reduce_lengths_backward_stub, &_segment_reduce_cpu_lengths_backward_kernel)
+REGISTER_ALL_CPU_DISPATCH(_segment_reduce_offsets_backward_stub, &_segment_reduce_cpu_offsets_backward_kernel)
 
 } // namespace at::native
diff --git a/aten/src/ATen/native/SharedReduceOps.h b/aten/src/ATen/native/SharedReduceOps.h
index 1575d770b5368..ec1eaa80f36fe 100644
--- a/aten/src/ATen/native/SharedReduceOps.h
+++ b/aten/src/ATen/native/SharedReduceOps.h
@@ -24,24 +24,12 @@
 #if defined(__CUDACC__) || defined(__HIPCC__)
 template <typename scalar_t>
 inline C10_DEVICE scalar_t max_propagate_nan(scalar_t a, scalar_t b) {
-#if defined(__HIPCC__)
-  // TODO: remove this special case for HIP when issue is fixed:
-  //       https://github.com/ROCm/hip/issues/2209
-  scalar_t max = at::_isnan(a) ? a : (at::_isnan(b) ? b : std::max(a, b));
-#else
   scalar_t max = at::_isnan(b) ? b : std::max(a, b);
-#endif
   return max;
 }
 template <typename scalar_t>
 inline C10_DEVICE scalar_t min_propagate_nan(scalar_t a, scalar_t b) {
-#if defined(__HIPCC__)
-  // TODO: remove this special case for HIP when issue is fixed:
-  //       https://github.com/ROCm/hip/issues/2209
-  scalar_t min = at::_isnan(a) ? a : (at::_isnan(b) ? b : std::min(a, b));
-#else
   scalar_t min = at::_isnan(b) ? b : std::min(a, b);
-#endif
   return min;
 }
 #define MAX(X, Y) max_propagate_nan(X,Y)
diff --git a/aten/src/ATen/native/SobolEngineOps.cpp b/aten/src/ATen/native/SobolEngineOps.cpp
index 27fc833ce6570..c061adc475859 100644
--- a/aten/src/ATen/native/SobolEngineOps.cpp
+++ b/aten/src/ATen/native/SobolEngineOps.cpp
@@ -58,7 +58,7 @@ std::tuple<Tensor, Tensor> _sobol_engine_draw(const Tensor& quasi, int64_t n, co
   });
 
   result.mul_(RECIPD);
-  return std::tuple<Tensor, Tensor>(result, wquasi);
+  return std::tuple<Tensor, Tensor>(std::move(result), std::move(wquasi));
 }
 
 /// This is the core function to fast-forward a `SobolEngine` given
diff --git a/aten/src/ATen/native/Sorting.cpp b/aten/src/ATen/native/Sorting.cpp
index 44215a26018f0..9c1d6c4c4e0bc 100644
--- a/aten/src/ATen/native/Sorting.cpp
+++ b/aten/src/ATen/native/Sorting.cpp
@@ -206,7 +206,7 @@ QUANTILE_INTERPOLATION_MODE get_quantile_interpolation_mode(
 }
 
 void quantile_checks(const Tensor& self, const Tensor& q) {
-  TORCH_CHECK(self.numel() > 0, "quantile() input tensor must be non-empty");
+  TORCH_SYM_CHECK(self.sym_numel().sym_gt(0), "quantile() input tensor must be non-empty");
   TORCH_CHECK(q.dim() <= 1, "quantile() q must be a scalar or 1D tensor");
   TORCH_CHECK(
       self.scalar_type() == kFloat || self.scalar_type() == kDouble,
@@ -219,26 +219,26 @@ void quantile_checks(const Tensor& self, const Tensor& q) {
       "quantile() q tensor must be on the same device as the input tensor");
 }
 
-std::vector<int64_t> quantile_output_shape(
+std::vector<c10::SymInt> quantile_output_shape(
     const std::optional<int64_t> original_dim,
     const Tensor& self,
     const Tensor& q,
     const bool keepdim,
     int64_t wrapped_dim) {
   // Compute output shape: q_size + reduced_size
-  std::vector<int64_t> out_shape;
+  std::vector<c10::SymInt> out_shape;
   if (original_dim && self.dim() > 0) {
-    out_shape = self.sizes().vec();
+    out_shape = self.sym_sizes().vec();
     if (keepdim) {
       out_shape[wrapped_dim] = 1;
     } else {
       out_shape.erase(out_shape.begin() + wrapped_dim);
     }
   } else if (keepdim) {
-    out_shape = std::vector<int64_t>(self.dim(), 1);
+    out_shape = std::vector<c10::SymInt>(self.dim(), 1);
   }
   if (q.dim() > 0) {
-    out_shape.insert(out_shape.begin(), q.numel());
+    out_shape.insert(out_shape.begin(), q.sym_numel());
   }
 
   return out_shape;
@@ -252,11 +252,13 @@ Tensor quantile_compute(
     const QUANTILE_INTERPOLATION_MODE& interpolation,
     const bool ignore_nan,
     int64_t wrapped_dim,
-    std::vector<int64_t> out_shape) {
+    std::vector<c10::SymInt> out_shape) {
   // Checks that all q values are between 0 and 1, inclusive
   // NOTE: this check is only performed when running on the CPU to avoid
   // synchronizing an accelerator with the CPU
-  if (self.device().is_cpu()) {
+  // The check is also skipped when the actual q values are not available yet
+  // e.g. with symbolic shapes or during export
+  if (self.device().is_cpu() && !isTensorSubclassLike(q)) {
     auto all_q_in_range = q.ge(0).logical_and_(q.le(1)).all();
     TORCH_CHECK(at::is_scalar_tensor_true(all_q_in_range),
                 "quantile() q values must be in the range [0, 1]");
@@ -275,18 +277,18 @@ Tensor quantile_compute(
 
   // Treat q as a 1D tensor for the following computations
   if (q.dim() == 0) {
-    out_shape.insert(out_shape.begin(), q.numel());
+    out_shape.insert(out_shape.begin(), 1);
   }
 
   // View input as reduced_size + size of dim to reduce
-  std::vector<int64_t> in_shape(out_shape.size());
+  std::vector<c10::SymInt> in_shape(out_shape.size());
   std::copy(out_shape.begin() + 1, out_shape.end(), in_shape.begin());
-  in_shape[in_shape.size() - 1] = sorted.size(-1);
-  sorted = sorted.view(in_shape);
+  in_shape[in_shape.size() - 1] = sorted.sym_size(-1);
+  sorted = sorted.view_symint(in_shape);
 
   // Ensure converting from int64_t to double won't overflow
-  TORCH_CHECK(
-      sorted.size(-1) <= std::pow(2, 24),
+  TORCH_SYM_CHECK(
+      sorted.sym_size(-1).sym_le(1 << 24),
       "quantile() input tensor is too large");
 
   // Convert q in [0, 1] to ranks in [0, reduction_size)
@@ -308,7 +310,7 @@ Tensor quantile_compute(
   } else {
     // For quantile, compute ranks based on reduction size. If there is nan
     // set rank to last index so the quantile computed will be nan.
-    int64_t last_index = sorted.size(-1) - 1;
+    auto last_index = sorted.sym_size(-1) - 1;
     std::vector<Tensor> tl =
         at::broadcast_tensors({q * last_index, sorted.isnan().any(-1, true)});
     ranks = at::masked_fill(tl[0], tl[1], last_index);
@@ -388,7 +390,7 @@ void quantile_out_impl(
   int64_t wrapped_dim = at::maybe_wrap_dim(original_dim.value_or(0), self.dim());
 
   auto out_shape = quantile_output_shape(original_dim, self, q, keepdim, wrapped_dim);
-  resize_output(out, out_shape);
+  resize_output_symint(out, out_shape);
 
   auto quantile = quantile_compute(
       self, q, original_dim, keepdim, interpolation, ignore_nan, wrapped_dim, std::move(out_shape));
@@ -812,7 +814,7 @@ std::tuple<Tensor, Tensor> kthvalue(
   Tensor values = at::empty({0}, self.options());
   Tensor indices = at::empty({0}, self.options().dtype(kLong));
   at::kthvalue_out(values, indices, self, k, dim, keepdim);
-  return std::make_tuple(values, indices);
+  return std::make_tuple(std::move(values), std::move(indices));
 }
 
 std::tuple<Tensor, Tensor> kthvalue(
@@ -877,7 +879,7 @@ std::tuple<Tensor, Tensor> median(
   Tensor values = at::empty({0}, self.options());
   Tensor indices = at::empty({0}, self.options().dtype(kLong));
   at::median_out(values, indices, self, dim, keepdim);
-  return std::make_tuple(values, indices);
+  return std::make_tuple(std::move(values), std::move(indices));
 }
 
 std::tuple<Tensor, Tensor> median(
@@ -924,7 +926,7 @@ std::tuple<Tensor, Tensor> nanmedian(
   Tensor values = at::empty({0}, self.options());
   Tensor indices = at::empty({0}, self.options().dtype(kLong));
   at::nanmedian_out(values, indices, self, dim, keepdim);
-  return std::make_tuple(values, indices);
+  return std::make_tuple(std::move(values), std::move(indices));
 }
 
 std::tuple<Tensor, Tensor> nanmedian(
diff --git a/aten/src/ATen/native/SpectralOps.cpp b/aten/src/ATen/native/SpectralOps.cpp
index 91a0c3ff8cf93..6a8dd6c75769b 100644
--- a/aten/src/ATen/native/SpectralOps.cpp
+++ b/aten/src/ATen/native/SpectralOps.cpp
@@ -80,9 +80,9 @@ ScalarType promote_type_fft(ScalarType type, bool require_complex, Device device
   }
 
   const bool maybe_support_half = (
-    // Only CUDA supports half precision, but since meta tensors don't have a
+    // CUDA and XPU support half precision, but since meta tensors don't have a
     // device we err on the side of accepting it
-    device.is_cuda() || device.is_meta()
+    device.is_cuda() || device.is_meta() || device.is_xpu()
   );
   if (maybe_support_half) {
     TORCH_CHECK(type == kHalf || type == kFloat || type == kDouble, "Unsupported dtype ", type);
diff --git a/aten/src/ATen/native/TensorAdvancedIndexing.cpp b/aten/src/ATen/native/TensorAdvancedIndexing.cpp
index ecb67c9ef3799..4f91a78c0c17c 100644
--- a/aten/src/ATen/native/TensorAdvancedIndexing.cpp
+++ b/aten/src/ATen/native/TensorAdvancedIndexing.cpp
@@ -988,7 +988,8 @@ Tensor& _index_put_impl_(
     }
   }
   if ((self.device().type() == DeviceType::CUDA ||
-       self.device().type() == DeviceType::XPU) &&
+       self.device().type() == DeviceType::XPU ||
+       self.device().type() == DeviceType::PrivateUse1) &&
       (accumulate ||
        (globalContext().deterministicAlgorithms() && value_.numel() > 1))) {
     TORCH_CHECK(
diff --git a/aten/src/ATen/native/TensorAdvancedIndexingUtils.h b/aten/src/ATen/native/TensorAdvancedIndexingUtils.h
index 6f127b711d3e8..b7b8ab17e9a4e 100644
--- a/aten/src/ATen/native/TensorAdvancedIndexingUtils.h
+++ b/aten/src/ATen/native/TensorAdvancedIndexingUtils.h
@@ -66,7 +66,7 @@ inline std::tuple<bool, Tensor> canDispatchToMaskedFill(
        c10::irange(num_ind, self.ndimension())) {
     mask = mask.unsqueeze(-1);
   }
-  return std::make_tuple(true, mask);
+  return std::make_tuple(true, std::move(mask));
 }
 
 inline AdvancedIndex make_info(Tensor self, IOptTensorListRef orig) {
diff --git a/aten/src/ATen/native/TensorCompare.cpp b/aten/src/ATen/native/TensorCompare.cpp
index 1a3843e9cdca8..68492256cc802 100644
--- a/aten/src/ATen/native/TensorCompare.cpp
+++ b/aten/src/ATen/native/TensorCompare.cpp
@@ -1,3 +1,4 @@
+#include <limits>
 #define TORCH_ASSERT_ONLY_METHOD_OPERATORS
 #include <ATen/Dispatch.h>
 #include <ATen/NamedTensorUtils.h>
@@ -380,7 +381,14 @@ Tensor isclose(
     if (isTensorSubclassLike(other)) {
       close.__ior__(self.isnan().bitwise_and(other.isnan()));
     } else {
-      close.__ior__(self.isnan().__iand__(other.isnan()));
+      // In-place __iand__ requires the target shape to equal the broadcast
+      // result. Check shapes to pick the right in-place target, or fall back
+      // to out-of-place for mutual broadcast (e.g. [3,1] vs [1,4] -> [3,4]).
+      if (self.sizes() == other.sizes()) {
+        close.__ior__(self.isnan().__iand__(other.isnan()));
+      } else {
+        close.__ior__(self.isnan().bitwise_and(other.isnan()));
+      }
     }
   }
 
@@ -826,15 +834,18 @@ TORCH_IMPL_FUNC(clamp_out)
  const OptionalScalarRef max,
  const Tensor& result) {
   using at::native::detail::ClampLimits;
-  if (min && max) {
-    if (min.get().toDouble() != min.get().toDouble() ||
-        max.get().toDouble() != max.get().toDouble()) {
-      at::fill_(
-          const_cast<Tensor&>(result),
-          std::numeric_limits<double>::quiet_NaN());
-    } else {
-      clamp_scalar_stub(device_type(), *this, min.get(), max.get());
-    }
+
+  // If either scalar bound is NaN, the result should be all NaNs
+  const bool min_is_nan =
+      min.has_value() && (min.get().toDouble() != min.get().toDouble());
+  const bool max_is_nan =
+      max.has_value() && (max.get().toDouble() != max.get().toDouble());
+
+  if (min_is_nan || max_is_nan) {
+    at::fill_(
+        const_cast<Tensor&>(result), std::numeric_limits<double>::quiet_NaN());
+  } else if (min && max) {
+    clamp_scalar_stub(device_type(), *this, min.get(), max.get());
   } else if (max) {
     clamp_max_scalar_stub(device_type(), *this, max.get());
   } else if (min) {
diff --git a/aten/src/ATen/native/TensorConversions.cpp b/aten/src/ATen/native/TensorConversions.cpp
index 909381ce5ed30..697150b0fbb3a 100644
--- a/aten/src/ATen/native/TensorConversions.cpp
+++ b/aten/src/ATen/native/TensorConversions.cpp
@@ -652,7 +652,6 @@ Tensor to_dense_backward(
     default:
       TORCH_CHECK(
           false, "to_dense_backward: Unsupported input layout: ", input_layout);
-      return Tensor{};
   }
 }
 
@@ -1399,7 +1398,6 @@ Tensor dense_to_sparse_with_mask(
       " to ",
       layout_to,
       " conversion not supported");
-  return Tensor{};
 }
 
 Tensor dense_to_sparse_csr(
@@ -1482,7 +1480,6 @@ Tensor dense_to_sparse(
       " to ",
       layout_to,
       " conversion not supported");
-  return Tensor{};
 }
 
 Tensor dense_to_sparse(const Tensor& self, int64_t sparse_dim) {
@@ -1766,7 +1763,6 @@ Tensor sparse_compressed_to_sparse_csr(
       false,
       "sparse_compressed_to_sparse_csr: expected SparseCsr or SparseCsc layout but got ",
       self.layout());
-  return Tensor{};
 }
 
 Tensor sparse_compressed_to_sparse_csc(
@@ -1787,7 +1783,6 @@ Tensor sparse_compressed_to_sparse_csc(
       false,
       "sparse_compressed_to_sparse_csc: expected SparseCsr or SparseCsc layout but got ",
       self.layout());
-  return Tensor{};
 }
 
 Tensor coo_to_sparse_csr(
@@ -2222,7 +2217,6 @@ Tensor sparse_compressed_to_sparse_bsr(
       false,
       "sparse_compressed_to_sparse_bsr: expected SparseCsr, SparseCsc, SparseBsr or SparseBsc layout but got ",
       self.layout());
-  return Tensor{};
 }
 
 Tensor sparse_compressed_to_sparse_bsc(
@@ -2260,7 +2254,6 @@ Tensor sparse_compressed_to_sparse_bsc(
       false,
       "sparse_compressed_to_sparse_bsc: expected SparseCsr, SparseCsc, SparseBsr or SparseBsc layout but got ",
       self.layout());
-  return Tensor{};
 }
 
 Tensor sparse_coo_to_sparse(const Tensor& self, const int64_t sparse_dim) {
@@ -2273,7 +2266,6 @@ Tensor sparse_coo_to_sparse(const Tensor& self, const int64_t sparse_dim) {
       " to ",
       kSparse,
       " conversion not supported");
-  return Tensor{};
 }
 
 Tensor sparse_compressed_to_sparse(
@@ -2377,7 +2369,6 @@ Tensor sparse_compressed_to_sparse(
       " to ",
       layout_to,
       " conversion not supported");
-  return Tensor{};
 }
 
 Tensor sparse_coo_to_sparse(
@@ -2414,7 +2405,6 @@ Tensor sparse_coo_to_sparse(
       " to ",
       layout_to,
       " conversion not supported");
-  return Tensor{};
 }
 
 Tensor to_sparse(const Tensor& self, const int64_t sparse_dim) {
diff --git a/aten/src/ATen/native/TensorProperties.cpp b/aten/src/ATen/native/TensorProperties.cpp
index c15b082f107b2..26cafbb4b2585 100644
--- a/aten/src/ATen/native/TensorProperties.cpp
+++ b/aten/src/ATen/native/TensorProperties.cpp
@@ -13,10 +13,15 @@
 #include <ATen/ops/contiguous_native.h>
 #include <ATen/ops/cudnn_is_acceptable_native.h>
 #include <ATen/ops/detach_native.h>
+#include <ATen/ops/dim_native.h>
 #include <ATen/ops/equal.h>
+#include <ATen/ops/get_device_native.h>
+#include <ATen/ops/is_contiguous_native.h>
 #include <ATen/ops/is_same_size_native.h>
 #include <ATen/ops/is_set_to_native.h>
+#include <ATen/ops/numel_native.h>
 #include <ATen/ops/size_native.h>
+#include <ATen/ops/storage_offset_native.h>
 #include <ATen/ops/stride_native.h>
 #include <ATen/ops/sym_is_contiguous_native.h>
 #include <ATen/ops/sym_numel_native.h>
@@ -76,6 +81,30 @@ c10::SymInt sym_storage_offset(const Tensor& self) {
   return self.sym_storage_offset();
 }
 
+int64_t numel(const Tensor& self) {
+  return self.numel();
+}
+
+int64_t dim(const Tensor& self) {
+  return self.dim();
+}
+
+int64_t get_device(const Tensor& self) {
+  return self.get_device();
+}
+
+int64_t storage_offset(const Tensor& self) {
+  return self.storage_offset();
+}
+
+bool is_contiguous(const Tensor& self) {
+  return self.is_contiguous();
+}
+
+bool is_contiguous(const Tensor& self, at::MemoryFormat memory_format) {
+  return self.is_contiguous(memory_format);
+}
+
 int64_t size(const Tensor& self, Dimname dim) {
   size_t pos_dim = dimname_to_position(self, dim);
   return self.sizes()[pos_dim];
diff --git a/aten/src/ATen/native/TensorShape.cpp b/aten/src/ATen/native/TensorShape.cpp
index d52072cec6a10..5a97b3cf8386e 100644
--- a/aten/src/ATen/native/TensorShape.cpp
+++ b/aten/src/ATen/native/TensorShape.cpp
@@ -1422,11 +1422,12 @@ Tensor as_strided_tensorimpl(
 }
 
 template <typename T>
-static inline void setStridedUnchecked(
+static void setStridedUnchecked(
     const Tensor& self,
     ArrayRef<T> size,
     ArrayRef<T> stride,
     T&& storage_offset) {
+  checkAsStridedArgsAllowUnbackedSymInts(size, stride, storage_offset);
   auto* self_ = self.unsafeGetTensorImpl();
   self_->set_sizes_and_strides(size, stride, std::forward<T>(storage_offset));
 }
diff --git a/aten/src/ATen/native/TestOps.cpp b/aten/src/ATen/native/TestOps.cpp
index 90dbf97075093..a1dbde708157b 100644
--- a/aten/src/ATen/native/TestOps.cpp
+++ b/aten/src/ATen/native/TestOps.cpp
@@ -137,7 +137,6 @@ Tensor FunctionalInverses::_test_autograd_multiple_dispatch_view_inverse(const a
     TORCH_INTERNAL_ASSERT(false,
     "Attempted to call _test_autograd_multiple_dispatch_view_inverse() during the functionalization pass. ",
     "This function is for testing only and should never be called.");
-    return Tensor();
 }
 
 } // namespace at::functionalization
diff --git a/aten/src/ATen/native/TriangularOps.cpp b/aten/src/ATen/native/TriangularOps.cpp
index 5560f3e79f273..009c1eab673a5 100644
--- a/aten/src/ATen/native/TriangularOps.cpp
+++ b/aten/src/ATen/native/TriangularOps.cpp
@@ -185,7 +185,10 @@ Tensor trace_backward_symint(const Tensor& grad, c10::SymIntArrayRef sizes) {
   TORCH_CHECK(sizes.size() == 2, "expected matrix input");
 
   auto grad_input = at::zeros_symint(sizes[0] * sizes[1], grad.options());
-  auto indices = at::arange(0, grad_input.numel(), sizes[1] + 1, grad.options().dtype(at::kLong));
+  // Only min(M, N) diagonal elements exist for an (M, N) matrix
+  auto diag_size = std::min(sizes[0], sizes[1]);
+  auto step = sizes[1] + 1;
+  auto indices = at::arange(0, diag_size * step, step, grad.options().dtype(at::kLong));
   // for composite compliance, use out-of-place variant of
   // `index_fill` if grad tensor is a Tensor Subclass.
   if (isTensorSubclassLike(grad)) {
diff --git a/aten/src/ATen/native/UnaryOps.cpp b/aten/src/ATen/native/UnaryOps.cpp
index acf14f3dfcdd5..ef534208e61e8 100644
--- a/aten/src/ATen/native/UnaryOps.cpp
+++ b/aten/src/ATen/native/UnaryOps.cpp
@@ -912,6 +912,10 @@ Tensor& mvlgamma_(Tensor& self, int64_t p) {
 }
 
 Tensor& mvlgamma_out(const Tensor& self, int64_t p, Tensor& result) {
+  TORCH_CHECK(
+    self.device() == result.device(),
+    "Expected tensors to be on the same device, but found ", self.device(), " and ", result.device()
+  );
   auto out = self.mvlgamma(p);
   TORCH_CHECK(
       at::can_cast(out.scalar_type(), result.scalar_type()),
@@ -936,7 +940,7 @@ std::tuple<Tensor, Tensor> frexp(const Tensor& self) {
   Tensor exponent = at::empty_like(self, self.options().dtype(at::kInt));
 
   at::frexp_out(mantissa, exponent, self);
-  return std::tuple<Tensor, Tensor>(mantissa, exponent);
+  return std::tuple<Tensor, Tensor>(std::move(mantissa), std::move(exponent));
 }
 
 std::tuple<Tensor&, Tensor&> frexp_out(const Tensor& self,
diff --git a/aten/src/ATen/native/Unique.cpp b/aten/src/ATen/native/Unique.cpp
index b14079e7ea19c..a805053ef5357 100644
--- a/aten/src/ATen/native/Unique.cpp
+++ b/aten/src/ATen/native/Unique.cpp
@@ -433,7 +433,7 @@ std::tuple<Tensor, Tensor, Tensor> _unique_dim_cpu_template(
   output = output.view(new_sizes);
   output = output.moveaxis(0, dim);
 
-  return std::make_tuple(output, inverse_indices, counts);
+  return std::make_tuple(std::move(output), std::move(inverse_indices), std::move(counts));
 }
 
 } // namespace
diff --git a/aten/src/ATen/native/UpSample.h b/aten/src/ATen/native/UpSample.h
index cf6727c2207c7..6483ca256414e 100644
--- a/aten/src/ATen/native/UpSample.h
+++ b/aten/src/ATen/native/UpSample.h
@@ -78,6 +78,7 @@ using _upsampling_bilinear2d_aa = void(*)(const Tensor& output, const Tensor& in
 using upsampling_trilinear3d = void(*)(const Tensor& output, const Tensor& input, bool align_corners, scale_t scales_d, scale_t scales_h, scale_t scales_w);
 using upsampling_bicubic2d = void(*)(const Tensor& output, const Tensor& input, bool align_corners, scale_t scales_h, scale_t scales_w);
 using _upsampling_bicubic2d_aa = void(*)(const Tensor& output, const Tensor& input, bool align_corners, scale_t scales_h, scale_t scales_w);
+using _upsampling_lanczos2d_aa = void(*)(const Tensor& output, const Tensor& input, bool align_corners, scale_t scales_h, scale_t scales_w);
 DECLARE_DISPATCH(upsampling_nearest1d, upsample_nearest1d_kernel)
 DECLARE_DISPATCH(_upsampling_nearest_exact1d, _upsample_nearest_exact1d_kernel)
 DECLARE_DISPATCH(upsampling_nearest2d, upsample_nearest2d_kernel)
@@ -101,6 +102,8 @@ DECLARE_DISPATCH(upsampling_trilinear3d, upsample_trilinear3d_backward_kernel)
 DECLARE_DISPATCH(upsampling_bicubic2d, upsample_bicubic2d_kernel)
 DECLARE_DISPATCH(_upsampling_bicubic2d_aa, _upsample_bicubic2d_aa_kernel)
 DECLARE_DISPATCH(_upsampling_bicubic2d_aa, _upsample_bicubic2d_aa_backward_kernel)
+DECLARE_DISPATCH(_upsampling_lanczos2d_aa, _upsample_lanczos2d_aa_kernel)
+DECLARE_DISPATCH(_upsampling_lanczos2d_aa, _upsample_lanczos2d_aa_backward_kernel)
 
 [[maybe_unused]] inline std::array<int64_t, 3> upsample_1d_common_check(
     IntArrayRef input_size,
diff --git a/aten/src/ATen/native/UpSampleLanczos2d.cpp b/aten/src/ATen/native/UpSampleLanczos2d.cpp
new file mode 100644
index 0000000000000..22ae7d972029f
--- /dev/null
+++ b/aten/src/ATen/native/UpSampleLanczos2d.cpp
@@ -0,0 +1,102 @@
+#define TORCH_ASSERT_ONLY_METHOD_OPERATORS
+#include <ATen/core/Tensor.h>
+#include <ATen/TensorMeta.h>
+#include <ATen/native/UpSample.h>
+
+#ifndef AT_PER_OPERATOR_HEADERS
+#include <ATen/Functions.h>
+#include <ATen/NativeFunctions.h>
+#else
+#include <ATen/ops/_upsample_lanczos2d_aa.h>
+#include <ATen/ops/_upsample_lanczos2d_aa_backward.h>
+#include <ATen/ops/_upsample_lanczos2d_aa_backward_native.h>
+#include <ATen/ops/_upsample_lanczos2d_aa_native.h>
+#endif
+
+namespace at::meta {
+
+TORCH_META_FUNC(_upsample_lanczos2d_aa) (
+  const Tensor& input, IntArrayRef output_size, bool align_corners, std::optional<double> scales_h, std::optional<double> scales_w
+) {
+  auto full_output_size = native::upsample_2d_common_check(input.sizes(), output_size);
+
+  TORCH_CHECK(
+      input.numel() != 0 || c10::multiply_integers(input.sizes().begin() + 1, input.sizes().end()),
+      "Non-empty 4D data tensor expected but got a tensor with sizes ",
+      input.sizes());
+
+  set_output_raw_strided(0, full_output_size, {}, input.options().memory_format(input.suggest_memory_format()));
+}
+
+TORCH_META_FUNC(_upsample_lanczos2d_aa_backward) (
+  const Tensor& grad_output,
+  IntArrayRef output_size,
+  IntArrayRef input_size,
+  bool align_corners,
+  std::optional<double> scales_h,
+  std::optional<double> scales_w
+) {
+  auto full_output_size = native::upsample_2d_common_check(input_size, output_size);
+
+  TORCH_CHECK(
+      grad_output.dim() == 4,
+      "Expected grad_output to be a tensor of dimension 4 but got: dimension ", grad_output.dim());
+
+  for (const auto i : c10::irange(4)) {
+    TORCH_CHECK(
+        grad_output.size(i) == full_output_size[i],
+        "Expected grad_output to have the same shape as output;",
+        " output.size(", i, ") = ", full_output_size[i],
+        " but got grad_output.size(", i, ") = ", grad_output.size(i));
+  }
+
+  set_output_raw_strided(0, input_size, {}, grad_output.options());
+}
+
+} // namespace at::meta
+namespace at::native {
+
+TORCH_IMPL_FUNC(_upsample_lanczos2d_aa_out_cpu) (
+    const Tensor& input,
+    IntArrayRef output_size,
+    bool align_corners,
+    std::optional<double> scales_h,
+    std::optional<double> scales_w,
+    const Tensor& output
+) {
+  _upsample_lanczos2d_aa_kernel(kCPU, output, input, align_corners, scales_h, scales_w);
+}
+
+TORCH_IMPL_FUNC(_upsample_lanczos2d_aa_backward_out_cpu) (
+    const Tensor& grad_output,
+    IntArrayRef output_size,
+    IntArrayRef input_size,
+    bool align_corners,
+    std::optional<double> scales_h,
+    std::optional<double> scales_w,
+    const Tensor& grad_input
+) {
+  grad_input.zero_();
+  _upsample_lanczos2d_aa_backward_kernel(kCPU, grad_input, grad_output, align_corners, scales_h, scales_w);
+}
+
+// vec variant
+
+using at::native::upsample::compute_output_size;
+using at::native::upsample::get_scale_value;
+
+Tensor _upsample_lanczos2d_aa(
+    const Tensor& input,
+    at::OptionalIntArrayRef output_size,
+    bool align_corners,
+    std::optional<ArrayRef<double>> scale_factors) {
+  auto osize = compute_output_size(input.sizes(), output_size, scale_factors);
+  auto scale_h = get_scale_value(scale_factors, 0);
+  auto scale_w = get_scale_value(scale_factors, 1);
+  return at::_upsample_lanczos2d_aa(input, osize, align_corners, scale_h, scale_w);
+}
+
+DEFINE_DISPATCH(_upsample_lanczos2d_aa_kernel);
+DEFINE_DISPATCH(_upsample_lanczos2d_aa_backward_kernel);
+
+} // namespace at::native
diff --git a/aten/src/ATen/native/WeightNorm.cpp b/aten/src/ATen/native/WeightNorm.cpp
index bbd39809085ab..e893b06658649 100644
--- a/aten/src/ATen/native/WeightNorm.cpp
+++ b/aten/src/ATen/native/WeightNorm.cpp
@@ -59,7 +59,7 @@ std::tuple<Tensor,Tensor> weight_norm_cpu(
   auto norm = at::empty_strided(g.sizes(), g.strides(), g.options().dtype(dtype));
   weight_norm_stub(kCPU, w, norm, v, g, dim);
 
-  return std::tuple<Tensor, Tensor>{w, norm};
+  return std::tuple<Tensor, Tensor>{std::move(w), std::move(norm)};
 }
 
 std::tuple<Tensor, Tensor> weight_norm_backward_cpu(
@@ -76,7 +76,7 @@ std::tuple<Tensor, Tensor> weight_norm_backward_cpu(
   auto grad_g = at::empty_like(saved_g, at::MemoryFormat::Contiguous);
   weight_norm_backward_stub(kCPU, grad_v, grad_g, grad_w, saved_v, saved_g, saved_norm, dim);
 
-  return std::tuple<Tensor, Tensor>{grad_v, grad_g};
+  return std::tuple<Tensor, Tensor>{std::move(grad_v), std::move(grad_g)};
 }
 
 Tensor _weight_norm
@@ -146,13 +146,13 @@ std::tuple<Tensor, Tensor> _weight_norm_differentiable_backward
     auto per_dim_sums = (grad_w*saved_v).view({saved_v.size(0), -1}).sum(1).view(bcast_size);
     auto grad_v = (saved_g/norms)*(grad_w - saved_v*(per_dim_sums/(norms*norms)));
     auto grad_g = per_dim_sums/norms;
-    return std::tuple<Tensor, Tensor>{grad_v, grad_g};
+    return std::tuple<Tensor, Tensor>{std::move(grad_v), std::move(grad_g)};
   } else { // dim == last_dim
     bcast_size[last_dim] = last_size;
     auto per_dim_sums = (grad_w*saved_v).view({-1, last_size}).sum(0).view(bcast_size);
     auto grad_v = (saved_g/norms)*(grad_w - saved_v*(per_dim_sums/(norms*norms)));
     auto grad_g = per_dim_sums/norms;
-    return std::tuple<Tensor, Tensor>{grad_v, grad_g};
+    return std::tuple<Tensor, Tensor>{std::move(grad_v), std::move(grad_g)};
   }
 }
 
diff --git a/aten/src/ATen/native/ao_sparse/quantized/cpu/qlinear_dynamic.cpp b/aten/src/ATen/native/ao_sparse/quantized/cpu/qlinear_dynamic.cpp
index 12841ad8e7391..8ecf4a2324074 100644
--- a/aten/src/ATen/native/ao_sparse/quantized/cpu/qlinear_dynamic.cpp
+++ b/aten/src/ATen/native/ao_sparse/quantized/cpu/qlinear_dynamic.cpp
@@ -29,7 +29,6 @@ at::Tensor PackedLinearWeightQnnp::apply_dynamic_impl<true>(
       false,
       "Sparse quantized dynamic linear with fused relu is not yet "
       "supported on qnnpack backend.");
-  return at::Tensor();
 }
 
 template <>
diff --git a/aten/src/ATen/native/cpu/DistributionTemplates.h b/aten/src/ATen/native/cpu/DistributionTemplates.h
index 1f8693902a32b..0a30fcb2aab26 100644
--- a/aten/src/ATen/native/cpu/DistributionTemplates.h
+++ b/aten/src/ATen/native/cpu/DistributionTemplates.h
@@ -4,6 +4,7 @@
 #include <ATen/Dispatch.h>
 #include <ATen/Dispatch_v2.h>
 #include <ATen/ExpandBase.h>
+#include <ATen/OpMathType.h>
 #include <ATen/core/DistributionsHelper.h>
 #include <ATen/native/TensorIterator.h>
 #include <ATen/native/cpu/Loops.h>
@@ -205,25 +206,53 @@ void normal_fill_VSX(const TensorBase &self, const scalar_t mean, const scalar_t
 
 template <typename scalar_t, typename RNG>
 void normal_fill(const TensorBase &self, const scalar_t mean, const scalar_t std, RNG generator) {
+  using opmath_t = at::opmath_type<scalar_t>;
   scalar_t *data = self.data_ptr<scalar_t>();
   auto size = self.numel();
   std::lock_guard<std::mutex> lock(generator->mutex_);
-  for (const auto i : c10::irange(size)) {
-    at::uniform_real_distribution<scalar_t> uniform(0, 1);
-    data[i] = uniform(generator);
-  }
-
-  for (int64_t i = 0; i < size - 15; i += 16) {
-    normal_fill_16<scalar_t>(data + i, mean, std);
-  }
-  if (size % 16 != 0) {
-    // Recompute the last 16 values.
-    data = data + size - 16;
-    for (const auto i : c10::irange(16)) {
-      at::uniform_real_distribution<scalar_t> uniform(0, 1);
+  auto omean = static_cast<opmath_t>(mean);
+  auto ostd = static_cast<opmath_t>(std);
+  at::uniform_real_distribution<opmath_t> uniform(0, 1);
+
+  if constexpr (std::is_same_v<scalar_t, opmath_t>) {
+    // float/double: generate uniform samples directly into the output buffer,
+    // then apply Box-Muller in-place.
+    for (const auto i : c10::irange(size)) {
       data[i] = uniform(generator);
     }
-    normal_fill_16<scalar_t>(data, mean, std);
+    for (int64_t i = 0; i < size - 15; i += 16) {
+      normal_fill_16<scalar_t>(data + i, omean, ostd);
+    }
+    if (size % 16 != 0) {
+      data = data + size - 16;
+      for (const auto i : c10::irange(16)) {
+        data[i] = uniform(generator);
+      }
+      normal_fill_16<scalar_t>(data, omean, ostd);
+    }
+  } else {
+    // bf16/fp16: generate in opmath_t precision using a stack buffer,
+    // apply Box-Muller, then cast down to scalar_t.
+    opmath_t buf[16];
+    for (int64_t i = 0; i < size - 15; i += 16) {
+      for (int j = 0; j < 16; j++) {
+        buf[j] = uniform(generator);
+      }
+      normal_fill_16<opmath_t>(buf, omean, ostd);
+      for (int j = 0; j < 16; j++) {
+        data[i + j] = static_cast<scalar_t>(buf[j]);
+      }
+    }
+    if (size % 16 != 0) {
+      int64_t offset = size - 16;
+      for (int j = 0; j < 16; j++) {
+        buf[j] = uniform(generator);
+      }
+      normal_fill_16<opmath_t>(buf, omean, ostd);
+      for (int j = 0; j < 16; j++) {
+        data[offset + j] = static_cast<scalar_t>(buf[j]);
+      }
+    }
   }
 }
 
@@ -267,11 +296,16 @@ template<typename RNG>
 void uniform_kernel(TensorIteratorBase& iter, double from_, double to_, RNG generator) {
   AT_DISPATCH_FLOATING_TYPES_AND2(kHalf, kBFloat16, iter.dtype(), "uniform_kernel_cpu", [&]() {
     std::lock_guard<std::mutex> lock(generator->mutex_);
-    auto from = static_cast<scalar_t>(from_);
-    auto to = static_cast<scalar_t>(to_);
-    at::uniform_real_distribution<scalar_t> uniform(from, to);
-    cpu_serial_kernel(iter, [&uniform, generator]() -> scalar_t {
-      return static_cast<scalar_t>(uniform(generator));
+    using opmath_t = at::opmath_type<scalar_t>;
+    auto from = static_cast<opmath_t>(from_);
+    auto to = static_cast<opmath_t>(to_);
+    auto to_scalar = static_cast<scalar_t>(to_);
+    auto from_scalar = static_cast<scalar_t>(from_);
+    at::uniform_real_distribution<opmath_t> uniform(from, to);
+    cpu_serial_kernel(iter, [&uniform, generator, to_scalar, from_scalar]() -> scalar_t {
+      auto value = static_cast<scalar_t>(uniform(generator));
+      // Clamp if the float→scalar_t cast rounded up to the upper bound
+      return value == to_scalar ? from_scalar : value;
     });
   });
 }
diff --git a/aten/src/ATen/native/cpu/LerpKernel.cpp b/aten/src/ATen/native/cpu/LerpKernel.cpp
index 2cbb204d32ae6..fdf19b8075863 100644
--- a/aten/src/ATen/native/cpu/LerpKernel.cpp
+++ b/aten/src/ATen/native/cpu/LerpKernel.cpp
@@ -19,7 +19,7 @@ Vectorized<scalar_t> is_lerp_weight_small(Vectorized<scalar_t> weight) {
 // is_lerp_weight_small doesn't work for complex because z.abs() returns a
 // complex vector which can't be compared. Either implement it with z.abs_2_(),
 // or fallback to the scalar function.
-#if !(defined(CPU_CAPABILITY_DEFAULT) || defined(_MSC_VER) || defined(CPU_CAPABILITY_SVE))
+#if !(defined(CPU_CAPABILITY_DEFAULT) || defined(_MSC_VER) || defined(CPU_CAPABILITY_SVE256))
 template <typename value_t>
 Vectorized<c10::complex<value_t>> is_lerp_weight_small(Vectorized<c10::complex<value_t>> weight) {
   using vec_reg_t = decltype(weight.abs_2_());
diff --git a/aten/src/ATen/native/cpu/Loops.h b/aten/src/ATen/native/cpu/Loops.h
index aad618a258a37..adadb636e7f9b 100644
--- a/aten/src/ATen/native/cpu/Loops.h
+++ b/aten/src/ATen/native/cpu/Loops.h
@@ -27,7 +27,6 @@
 //
 
 #include <cstdint>
-#include <c10/util/C++17.h>
 #include <c10/util/Load.h>
 #include <c10/util/irange.h>
 #include <ATen/detail/FunctionTraits.h>
@@ -302,13 +301,15 @@ VectorizedLoop2d<op_t, vop_t> make_vectorized_loop2d(
 }
 
 template <typename func_t>
-void cpu_kernel(TensorIteratorBase& iter, func_t&& op, int64_t grain_size = at::internal::GRAIN_SIZE) {
+void cpu_kernel(TensorIteratorBase& iter, func_t&& op, int64_t grain_size = at::internal::GRAIN_SIZE, bool check_dynamic_casting = true) {
   using traits = function_traits<func_t>;
   // this could be extended to work with void return types
   TORCH_INTERNAL_ASSERT(iter.ninputs() == traits::arity);
   TORCH_INTERNAL_ASSERT(iter.noutputs() == 1);
   // dynamic casting not currently supported on CPU
-  TORCH_INTERNAL_ASSERT(!needs_dynamic_casting<func_t>::check(iter));
+  if (check_dynamic_casting) {
+    TORCH_INTERNAL_ASSERT(!needs_dynamic_casting<func_t>::check(iter));
+  }
 
   iter.for_each([&](char** data, const int64_t* strides, int64_t n) {
     // basic loop can handle 1d slices with arbitrary strides, and 1d slices is all that
@@ -318,6 +319,11 @@ void cpu_kernel(TensorIteratorBase& iter, func_t&& op, int64_t grain_size = at::
   iter.cast_outputs();
 }
 
+template <typename func_t>
+void cpu_kernel_opaque(TensorIteratorBase& iter, func_t&& op, int64_t grain_size = at::internal::GRAIN_SIZE) {
+  return cpu_kernel(iter, op, grain_size, false);
+}
+
 // This function helps write elementwise kernels that requires multiple outputs.
 // It follows the similar structure of cpu_kernel.
 // Instead of `basic_loop` function, a new `multiple_outputs_loop` function is
diff --git a/aten/src/ATen/native/cpu/ReducedPrecisionFloatGemvFastPathKernel.cpp b/aten/src/ATen/native/cpu/ReducedPrecisionFloatGemvFastPathKernel.cpp
index 8d22201ed63c4..3d2d1b336c6f9 100644
--- a/aten/src/ATen/native/cpu/ReducedPrecisionFloatGemvFastPathKernel.cpp
+++ b/aten/src/ATen/native/cpu/ReducedPrecisionFloatGemvFastPathKernel.cpp
@@ -150,16 +150,16 @@ float reduce(vec::VectorizedN<float, kF32RegistersPerIteration>& x) {
 // BFDOT. Deferring that for now to get the NEON/ASIMD BFDOT path
 // working.
 #if __ARM_FEATURE_BF16_VECTOR_ARITHMETIC
-#if defined(__aarch64__) && !defined(CPU_CAPABILITY_SVE) && defined(__clang__) && __clang_major__ > 15
+#if defined(__aarch64__) && !defined(CPU_CAPABILITY_SVE256) && defined(__clang__) && __clang_major__ > 15
 // https://godbolt.org/z/z8P4Yncra
 #define COMPILER_SUPPORTS_BF16_TARGET 1
-#elif defined(__aarch64__) && !defined(CPU_CAPABILITY_SVE) && !defined(__clang__) && defined(__GNUC__) && __GNUC__ >= 10
+#elif defined(__aarch64__) && !defined(CPU_CAPABILITY_SVE256) && !defined(__clang__) && defined(__GNUC__) && __GNUC__ >= 10
 // https://gcc.gnu.org/gcc-10/changes.html
 // https://godbolt.org/z/cdGG7vn8o
 #define COMPILER_SUPPORTS_BF16_TARGET 1
-#else // defined(__aarch64__) && !defined(CPU_CAPABILITY_SVE) && defined(__clang__) && __clang_major__ > 15
+#else // defined(__aarch64__) && !defined(CPU_CAPABILITY_SVE256) && defined(__clang__) && __clang_major__ > 15
 #define COMPILER_SUPPORTS_BF16_TARGET 0
-#endif // defined(__aarch64__) && !defined(CPU_CAPABILITY_SVE) && defined(__clang__) && __clang_major__ > 15
+#endif // defined(__aarch64__) && !defined(CPU_CAPABILITY_SVE256) && defined(__clang__) && __clang_major__ > 15
 #else // __ARM_FEATURE_BF16_VECTOR_ARITHMETIC
 #define COMPILER_SUPPORTS_BF16_TARGET 0
 #endif // __ARM_FEATURE_BF16_VECTOR_ARITHMETIC
@@ -212,7 +212,7 @@ std::pair<vec::Vectorized<float>, vec::Vectorized<float>> fmadd(
     const vec::Vectorized<c10::Half>& b,
     const vec::Vectorized<float>& acc_low,
     const vec::Vectorized<float>& acc_high) {
-#if defined(__ARM_FEATURE_FP16_FML) && !defined(CPU_CAPABILITY_SVE)
+#if defined(__ARM_FEATURE_FP16_FML) && !defined(CPU_CAPABILITY_SVE256)
   return std::make_pair(vfmlalq_low_f16(acc_low, a, b), vfmlalq_high_f16(acc_high, a, b));
 #else
   const auto [a_float_low, a_float_high] = convert_half_float(a);
@@ -233,7 +233,7 @@ std::pair<vec::Vectorized<float>, vec::Vectorized<float>> fmadd(
 
 // Return a + b_low * c_low + b_high * c_high
 vec::Vectorized<float> fmadd(vec::Vectorized<float> a, vec::Vectorized<Half> b, vec::Vectorized<Half> c) {
-#if defined(__aarch64__) && defined(__ARM_FEATURE_FP16_FML) && !defined(__ARM_FEATURE_SVE)
+#if defined(__aarch64__) && defined(__ARM_FEATURE_FP16_FML) && !defined(__ARM_FEATURE_SVE256)
   // NOTE: this instruction is an optional instruction in ARM v8.2 and
   // v8.3, but mandatory in v8.4 per
   // https://developer.arm.com/documentation/ddi0596/2021-03/SIMD-FP-Instructions/FMLAL--FMLAL2--vector---Floating-point-fused-Multiply-Add-Long-to-accumulator--vector--?lang=en
diff --git a/aten/src/ATen/native/cpu/UpSampleKernel.cpp b/aten/src/ATen/native/cpu/UpSampleKernel.cpp
index 79583b59edaf1..6274f22ecad7c 100644
--- a/aten/src/ATen/native/cpu/UpSampleKernel.cpp
+++ b/aten/src/ATen/native/cpu/UpSampleKernel.cpp
@@ -9,6 +9,7 @@
 #include <ATen/native/cpu/utils.h>
 #include <c10/util/irange.h>
 #include <ATen/native/cpu/UpSampleKernelAVXAntialias.h>
+#include <ATen/native/cpu/UpSampleKernelNEONAntialias.h>
 
 #ifndef AT_PER_OPERATOR_HEADERS
 #include <ATen/Functions.h>
@@ -23,67 +24,54 @@ namespace {
 
 using scale_t = std::vector<std::optional<double>>;
 
-// TODO: this file could benefit from a global renaming of its functions /
-// classes and terms, as well as from adding more comments. In particular:
-// - It's not obvious that despite their names (and the file name), all these
-//   kernels don't just do upsampling: they do general interpolation, i.e. they
-//   also all support downscaling.
-// - the term "horizontal" or "within dims" or "contiguous dim" refers to the
-//   last dimension.
-//   It's not specific to 2D images and applies to 3D (and 1D??) inputs as well.
-//   Similarly "vertical" or "across dims" refers to all dims that aren't the
-//   last one. In other kernels these are also referred to as "zero-stride" and
-//   "non-zero-stride" - we should unify all this.
-// - the terms "zero-stride" and "non-zero strides" refer to the weights and
-//   indices, not to the contiguity of input or output
-// - It's not always clear which kernel is vectorized and which one isn't.
-// - The functions like _use_vectorized_kernel_cond() should be renamed and
-//   their description updated, because they're not the only "fork" in the
-//   code-path where a choice is made between a vectorized kernel vs a
-//   non-vectorized one. See e.g. upsample_bilinear2d_kernel_impl() where we
-//   already make a similar check, before the one in
-//   _use_vectorized_kernel_cond().
-// - It's not always clear which code is part of a "separable interpolation"
-//   code-path.
-// - Some names need to be more specific. For example
-//   "cpu_upsample_generic_aa()" looks like a super generic name, but the function
-//   is instead fairly specific - we need to make that clearer.
-// - Some functions have a "aa" suffix but it doesn't mean that they only
-//   support antialias. Some of them also support antialias=False now.
-// - Various comments are outdated. Case in point: the one just below about the
-//   `Interpolate` struct being used for cpu_upsample_linear:
-//   cpu_upsample_linear doesn't exist anymore, and these structs are used for
-//   various modes, *not* just linear.
-// - It'd be useful to document how interpolation works in general, and in particular state explicitly:
-//   - that the weights and indices across a given dimension are the same for
-//     all pixels (hence the benefit of pre-computing them)
-//   - that it can be "separated", i.e. we can do the horizontal pass and the
-//     vertical pass independently (and that some kernels are written this way,
-//     while some aren't.)
-// - we can probably remove the template over index_t, because it's always
-//   hard-coded as int64_t
-
-
-// Helper structs and methods for cpu_upsample_linear
+// Naming conventions used in this file:
 //
-// Interpolation methods that used below are separable, and as such we can compute the interpolation
-// independently per dimension in a recursive way. Please, refer to #10482 for more context.
+// - "non_separable": All spatial dimensions are interpolated in a single
+//   TensorIterator pass, using the recursive InterpolateNonSeparable struct.
+//   Non-separable is always multi-dimensional (Nd).
+//
+// - "separable": The multi-dimensional interpolation is decomposed into a
+//   sequence of 1d passes (one per spatial dimension). The entry point is Nd
+//   (it loops over dims), and inner functions are 1d.
+//
+//   - "1d": Refers to processing a single spatial dimension within the
+//     separable approach. Not to be confused with 1d interpolation (e.g.
+//     linear); it means one dimension of a multi-dimensional separable
+//     decomposition.
+//
+//   - "Nd": Templated on out_ndims (1, 2, or 3) for the separable approach.
+
+
+// ---- Non-separable interpolation ----
+//
+// Used by: nearest, linear, bilinear (float), cubic (float), trilinear.
+// Processes all spatial dims in a single TensorIterator pass via the recursive
+// InterpolateNonSeparable struct.
+//
+// Call chain:
+//   upsample_non_separable_Nd_kernel_impl
+//     -> upsample_non_separable
+//       -> basic_loop_non_separable
+//         -> interpolate_non_separable
+//           -> InterpolateNonSeparable (recursive struct)
+//
+// Helper structs and methods for upsample_non_separable
 //
 // Interpolation structure to compute output value in n-dimensional case.
 // - recursively compute interpolated output for each dimension
 // - we rely a lot on compiler's code optimization such that implemented operations
 //   can be automatically factorized and vectorized using SSE and AVX2
 template <int n, typename scalar_t, typename opmath_t, typename index_t, int interp_size>
-struct Interpolate {
+struct InterpolateNonSeparable {
     static inline opmath_t eval(char* src, char** data, const int64_t* strides, int64_t i) {
       index_t ids = *(index_t*)&data[0][i * strides[0]];
       opmath_t wts = *(scalar_t*)&data[1][i * strides[1]];
-      opmath_t t = Interpolate<n - 1, scalar_t, opmath_t, index_t, interp_size>::eval(src + ids, &data[2 * interp_size], &strides[2 * interp_size], i);
+      opmath_t t = InterpolateNonSeparable<n - 1, scalar_t, opmath_t, index_t, interp_size>::eval(src + ids, &data[2 * interp_size], &strides[2 * interp_size], i);
       opmath_t output = t * wts;
       for (const auto j : c10::irange(1, interp_size)) {
         ids = *(index_t*)&data[2 * j + 0][i * strides[2 * j + 0]];
         wts = *(scalar_t*)&data[2 * j + 1][i * strides[2 * j + 1]];
-        t = Interpolate<n - 1, scalar_t, opmath_t, index_t, interp_size>::eval(src + ids, &data[2 * interp_size], &strides[2 * interp_size], i);
+        t = InterpolateNonSeparable<n - 1, scalar_t, opmath_t, index_t, interp_size>::eval(src + ids, &data[2 * interp_size], &strides[2 * interp_size], i);
         output += t * wts;
       }
       return output;
@@ -91,7 +79,7 @@ struct Interpolate {
 };
 
 template <typename scalar_t, typename opmath_t, typename index_t, int interp_size>
-struct Interpolate<1, scalar_t, opmath_t, index_t, interp_size> {
+struct InterpolateNonSeparable<1, scalar_t, opmath_t, index_t, interp_size> {
     static inline opmath_t eval(char* src, char** data, const int64_t* strides, int64_t i) {
       index_t ids = *(index_t*)&data[0][i * strides[0]];
       opmath_t wts = *(scalar_t*)&data[1][i * strides[1]];
@@ -108,15 +96,15 @@ struct Interpolate<1, scalar_t, opmath_t, index_t, interp_size> {
 };
 
 template <int n, typename scalar_t, typename opmath_t, typename index_t>
-struct Interpolate<n, scalar_t, opmath_t, index_t, 1> {
+struct InterpolateNonSeparable<n, scalar_t, opmath_t, index_t, 1> {
     static inline opmath_t eval(char* src, char** data, const int64_t* strides, int64_t i) {
       index_t ids = *(index_t*)&data[0][i * strides[0]];
-      return Interpolate<n - 1, scalar_t, opmath_t, index_t, 1>::eval(src + ids, &data[2], &strides[2], i);
+      return InterpolateNonSeparable<n - 1, scalar_t, opmath_t, index_t, 1>::eval(src + ids, &data[2], &strides[2], i);
   }
 };
 
 template <typename scalar_t, typename opmath_t, typename index_t>
-struct Interpolate<1, scalar_t, opmath_t, index_t, 1> {
+struct InterpolateNonSeparable<1, scalar_t, opmath_t, index_t, 1> {
     static inline opmath_t eval(char* src, char** data, const int64_t* strides, int64_t i) {
       index_t ids = *(index_t*)&data[0][i * strides[0]];
       return *(scalar_t *)&src[ids];
@@ -126,25 +114,25 @@ struct Interpolate<1, scalar_t, opmath_t, index_t, 1> {
 // There is an unexpected 2x slowdown for upsample_trilinear3d channels_first
 // for both 1 and 6 threads. We have to specialize this case as below:
 // Once the issue is fixed we can keep generic implementation and remove:
-// struct Interpolate<n, scalar_t, index_t, 2> and
-// struct Interpolate<1, scalar_t, index_t, 2>
+// struct InterpolateNonSeparable<n, scalar_t, index_t, 2> and
+// struct InterpolateNonSeparable<1, scalar_t, index_t, 2>
 template <int n, typename scalar_t, typename opmath_t, typename index_t>
-struct Interpolate<n, scalar_t, opmath_t, index_t, 2> {
+struct InterpolateNonSeparable<n, scalar_t, opmath_t, index_t, 2> {
     static inline opmath_t eval(char* src, char** data, const int64_t* strides, int64_t i) {
         index_t i0 = *(index_t*)&data[0][i * strides[0]];
         index_t i1 = *(index_t*)&data[2][i * strides[2]];
         opmath_t w0 = *(scalar_t *)&data[1][i * strides[1]];
         opmath_t w1 = *(scalar_t *)&data[3][i * strides[3]];
 
-        opmath_t t0 = Interpolate<n - 1, scalar_t, opmath_t, index_t, 2>::eval(src + i0, &data[4], &strides[4], i);
-        opmath_t t1 = Interpolate<n - 1, scalar_t, opmath_t, index_t, 2>::eval(src + i1, &data[4], &strides[4], i);
+        opmath_t t0 = InterpolateNonSeparable<n - 1, scalar_t, opmath_t, index_t, 2>::eval(src + i0, &data[4], &strides[4], i);
+        opmath_t t1 = InterpolateNonSeparable<n - 1, scalar_t, opmath_t, index_t, 2>::eval(src + i1, &data[4], &strides[4], i);
 
         return t0 * w0 + t1 * w1;
   }
 };
 
 template <typename scalar_t, typename opmath_t, typename index_t>
-struct Interpolate<1, scalar_t, opmath_t, index_t, 2> {
+struct InterpolateNonSeparable<1, scalar_t, opmath_t, index_t, 2> {
     static inline opmath_t eval(char* src, char** data, const int64_t* strides, int64_t i) {
         index_t i0 = *(index_t*)&data[0][i * strides[0]];
         index_t i1 = *(index_t*)&data[2][i * strides[2]];
@@ -157,13 +145,28 @@ struct Interpolate<1, scalar_t, opmath_t, index_t, 2> {
 };
 
 template <int n, typename scalar_t, typename index_t, int interp_size>
-inline scalar_t interpolate(char* src, char** data, const int64_t* strides, int64_t i) {
+inline scalar_t interpolate_non_separable(char* src, char** data, const int64_t* strides, int64_t i) {
   using opmath_t = at::opmath_type<scalar_t>;
-  return Interpolate<n, scalar_t, opmath_t, index_t, interp_size>::eval(src, data, strides, i);
+  return InterpolateNonSeparable<n, scalar_t, opmath_t, index_t, interp_size>::eval(src, data, strides, i);
 }
 
+// ---- Separable interpolation ----
+//
+// Used by: bilinear/bicubic with antialias=True, and bilinear/bicubic uint8
+// (as fallback when AVX isn't supported).
+// Processes one spatial dimension at a time. The outer loop over dimensions
+// is in upsample_separable_Nd_kernel_impl.
+//
+// Call chain:
+//   upsample_separable_Nd_kernel_impl  (loops over dims)
+//     -> upsample_separable_1d
+//       -> basic_loop_separable_1d_horizontal  (for last spatial dim, i.e. W)
+//         -> interpolate_separable_1d
+//       -> basic_loop_separable_1d_vertical  (for other spatial dims, e.g. H, D)
+//         -> interpolate_separable_1d_zero_strides
+
 template <typename scalar_t, typename index_t>
-inline scalar_t interpolate_aa_single_dim_zero_strides(
+inline scalar_t interpolate_separable_1d_zero_strides(
     char* src,
     char** data,
     const index_t ids_stride) {
@@ -187,7 +190,7 @@ inline scalar_t interpolate_aa_single_dim_zero_strides(
 }
 
 template <typename scalar_t, typename index_t>
-inline scalar_t interpolate_aa_single_dim(
+inline scalar_t interpolate_separable_1d(
     char* src,
     char** data,
     const int64_t* strides,
@@ -233,6 +236,10 @@ inline bool is_contiguous_stride(const int64_t* strides) {
 // Helper class to recursively check if all input strides corresponding to interpolated dimensions
 // are equal zero except on a single dimension.
 //
+// Note: "zero-stride" and "non-zero stride" here refer to the strides of the
+// pre-computed indices and weights tensors (extra TensorIterator inputs), NOT
+// to the contiguity of the input or output data tensors.
+//
 // Inputs: array of strides of size N, non_zero_stride_dim which can be -1, 0, 1, 2, ...
 //   if non_zero_stride_dim, we check that all strides are equal zero, otherwise
 //   4 strides corresponding to the strides for index_0, weight_0, index_1 and weight_1 for non_zero_stride_dim
@@ -256,7 +263,7 @@ inline bool is_contiguous_stride(const int64_t* strides) {
 // strides=(0,  0,  0,  0,  0,  0,  0,  0, ..., 0,  0,  0,  0)
 //
 // Using these methods we can hint the compiler to factorize constant indices and weights
-// in cpu_upsample_linear method
+// in upsample_non_separable
 template <int N, int non_zero_stride_dim, typename scalar_t, typename index_t, int interp_size>
 struct CheckAlmostAllZeroStrides {
   static inline bool eval(const int64_t* strides) {
@@ -288,17 +295,17 @@ inline bool check_almost_all_zero_stride(const int64_t* strides) {
 
 // Helper method to compute interpolation for nearest, linear, cubic modes
 template <typename scalar_t, typename index_t, int out_ndims, int interp_size>
-inline void basic_loop(char** data, const int64_t* strides, int64_t n) {
+inline void basic_loop_non_separable(char** data, const int64_t* strides, int64_t n) {
   char* dst = data[0];
   char* src = data[1];
   for (const auto i : c10::irange(n)) {
-    *(scalar_t*)&dst[i * strides[0]] = interpolate<out_ndims, scalar_t, index_t, interp_size>(
+    *(scalar_t*)&dst[i * strides[0]] = interpolate_non_separable<out_ndims, scalar_t, index_t, interp_size>(
         src + i * strides[1], &data[2], &strides[2], i);
   }
 }
 
 template <typename scalar_t>
-inline void basic_loop_aa_vertical(
+inline void basic_loop_separable_1d_vertical(
     char** data,
     const int64_t* strides,
     int64_t n,
@@ -310,13 +317,13 @@ inline void basic_loop_aa_vertical(
 
   for (const auto i : c10::irange(n)) {
     *(scalar_t*)&dst[i * strides[0]] =
-        interpolate_aa_single_dim_zero_strides<scalar_t, int64_t>(
+        interpolate_separable_1d_zero_strides<scalar_t, int64_t>(
             src + i * strides[1], &data[2], ids_stride);
   }
 }
 
 template <>
-inline void basic_loop_aa_vertical<uint8_t>(
+inline void basic_loop_separable_1d_vertical<uint8_t>(
     char** data,
     const int64_t* strides,
     int64_t n,
@@ -354,7 +361,7 @@ inline void basic_loop_aa_vertical<uint8_t>(
 }
 
 template <typename scalar_t>
-inline void basic_loop_aa_horizontal(
+inline void basic_loop_separable_1d_horizontal(
     char** data,
     const int64_t* strides,
     int64_t n,
@@ -367,20 +374,20 @@ inline void basic_loop_aa_horizontal(
   if (strides[1] == 0) {
     for (const auto i : c10::irange(n)) {
       *(scalar_t*)&dst[i * strides[0]] =
-          interpolate_aa_single_dim<scalar_t, int64_t>(
+          interpolate_separable_1d<scalar_t, int64_t>(
               src, &data[2], &strides[2], i, ids_stride);
     }
   } else {
     for (const auto i : c10::irange(n)) {
       *(scalar_t*)&dst[i * strides[0]] =
-          interpolate_aa_single_dim<scalar_t, int64_t>(
+          interpolate_separable_1d<scalar_t, int64_t>(
               src + i * strides[1], &data[2], &strides[2], i, ids_stride);
     }
   }
 }
 
 template <>
-inline void basic_loop_aa_horizontal<uint8_t>(
+inline void basic_loop_separable_1d_horizontal<uint8_t>(
     char** data,
     const int64_t* strides,
     int64_t n,
@@ -431,10 +438,10 @@ inline void basic_loop_aa_horizontal<uint8_t>(
 // output_DN[a] = interpolate(input_DN[a], w_DN[a], input_DN[a+1], w_DN[a+1], ...)
 // and i - dimension index and a - linear index for spatial coordinates
 //
-// The recursive call is implemented with InterpLinear struct using template for
+// The recursive call is implemented with the InterpolateNonSeparable struct using template for
 // the loop unrolling on compile time.
 template <typename scalar_t, int out_ndims, int interp_size>
-void cpu_upsample_generic(at::TensorIterator& iter)
+void upsample_non_separable(at::TensorIterator& iter)
 {
   auto loop = [&](char** data, const int64_t* strides, int64_t n) {
     // special-cases to let the compiler apply compile-time input-specific optimizations
@@ -442,21 +449,21 @@ void cpu_upsample_generic(at::TensorIterator& iter)
         // NOLINTNEXTLINE(bugprone-branch-clone)
         check_almost_all_zero_stride<out_ndims, 1, scalar_t, int64_t, interp_size>(&strides[2]))) {
       // contiguous channels-first case
-      basic_loop<scalar_t, int64_t, out_ndims, interp_size>(data, strides, n);
+      basic_loop_non_separable<scalar_t, int64_t, out_ndims, interp_size>(data, strides, n);
     } else if ((strides[0] == sizeof(scalar_t) && (strides[1] == sizeof(scalar_t)) &&
                check_almost_all_zero_stride<out_ndims, -1, scalar_t, int64_t, interp_size>(&strides[2]))) {
       // contiguous channels-last case
-      basic_loop<scalar_t, int64_t, out_ndims, interp_size>(data, strides, n);
+      basic_loop_non_separable<scalar_t, int64_t, out_ndims, interp_size>(data, strides, n);
     } else {
       // fallback
-      basic_loop<scalar_t, int64_t, out_ndims, interp_size>(data, strides, n);
+      basic_loop_non_separable<scalar_t, int64_t, out_ndims, interp_size>(data, strides, n);
     }
   };
   iter.for_each(loop);
 }
 
 template <typename scalar_t, typename scale_type, nearest_idx_fn_t nearest_idx_fn>
-void cpu_upsample_nearest_channels_last(
+void upsample_nearest_channels_last(
     const Tensor& output_,
     const Tensor& input_,
     const scale_type& scales) {
@@ -561,7 +568,7 @@ inline VecType<scalar_t> interpolate(const scalar_t* t, accscalar_t w, Args... a
 }
 
 template <typename scalar_t, typename scale_type>
-void cpu_upsample_linear_channels_last(
+void upsample_linear_channels_last(
     const Tensor& output_,
     const Tensor& input_,
     bool align_corners,
@@ -708,10 +715,10 @@ void cpu_upsample_linear_channels_last(
   };
 
   if (ndim == 4) {
-    // upsample nearest 2d
+    // bilinear 2d
     at::parallel_for(0, num_batches, at::internal::GRAIN_SIZE / output_slice_size / 4, loop2d);
   } else {
-    // upsample nearest 3d
+    // trilinear 3d
     TORCH_INTERNAL_ASSERT(ndim == 5);
     at::parallel_for(0, num_batches, at::internal::GRAIN_SIZE / output_slice_size / 8, loop3d);
   }
@@ -721,7 +728,7 @@ void cpu_upsample_linear_channels_last(
   }
 }
 
-// Helper structs to use with upsample_generic_Nd_kernel_impl
+// Helper structs to use with upsample_non_separable_Nd_kernel_impl
 struct HelperInterpBase {
 
   static inline void init_indices_weights(
@@ -787,7 +794,7 @@ struct HelperInterpBase {
   // for interpolation with antialiasing=false mode. It returns the maximal weights value.
   // This function is templated with scalar_t for type of scale and weights but is only used for
   // bilinear/bicubic modes on uint8 input and antialiasing=false (in this case scalar_t is double).
-  // For float input types we are using upsample_generic_Nd_kernel_impl and compute_indices_weights methods
+  // For float input types we are using upsample_non_separable_Nd_kernel_impl and compute_indices_weights methods
   template <typename scalar_t, typename aa_filter_fn_t>
   static inline scalar_t _compute_indices_min_size_weights(
     const int64_t i, const int64_t input_size, const scalar_t scale,
@@ -938,7 +945,7 @@ struct HelperInterpBase {
       idx_ptr_stride[i] = stride;
       wt_idx_ptr[i] = i * max_interp_size * weight_index_stride;
     }
-    return {output, max_interp_size, wt_max};
+    return {std::move(output), max_interp_size, wt_max};
   }
 
   /*
@@ -947,7 +954,7 @@ struct HelperInterpBase {
   weights as double, but then convert them to int16 via some conversion logic
   detailed below. This allows us to compute all interpolation operation (sum of
   multiplications) as ints instead of floats. The result is converted back into
-  uint8 in basic_loop_aa_horizontal<uint8_t> (and vertical)
+  uint8 in basic_loop_separable_1d_horizontal<uint8_t> (and vertical)
 
   In essence the idea is to avoid a multiplication between a float (the
   weight) and an int (the pixel value) and instead run a multiplication between
@@ -1028,7 +1035,7 @@ struct HelperInterpBase {
       }
     }
 
-    return {indices_weights, aligned_interp_size, weights_precision};
+    return {std::move(indices_weights), aligned_interp_size, weights_precision};
   }
 };
 
@@ -1409,6 +1416,89 @@ struct HelperInterpCubic : public HelperInterpBase {
 
 };
 
+struct HelperInterpLanczos : public HelperInterpBase {
+
+  static constexpr int interp_size = 6;
+
+  // Taken from
+  // https://github.com/python-pillow/Pillow/blob/8004234d879254cc354935ad42fbb51b1700925e/src/libImaging/Resample.c#L64-L86
+  template<typename scalar_t>
+  static inline scalar_t sinc_filter(scalar_t x) {
+    if (x == 0.0) {
+      return 1.0;
+    }
+    x *= c10::pi<scalar_t>;
+    return std::sin(x) / x;
+  }
+
+  template<typename scalar_t>
+  static inline scalar_t aa_filter(scalar_t x) {
+    // Lanczos-3 filter: sinc(x) * sinc(x/3) for |x| < 3
+    x = std::abs(x);
+    if (x < 3.0) {
+      return sinc_filter(x) * sinc_filter(x / 3.0);
+    }
+    return 0.0;
+  }
+
+  static inline std::vector<Tensor> compute_index_ranges_weights(
+    at::ScalarType scalar_type,
+    int64_t input_size,
+    int64_t output_size,
+    int64_t stride,
+    int64_t ndims,
+    int64_t reshape_dim,
+    bool align_corners,
+    const std::optional<double>& opt_scale,
+    bool antialias
+  ) {
+
+    std::vector<Tensor> indices_weights;
+    AT_DISPATCH_FLOATING_TYPES(
+      scalar_type, "compute_index_ranges_weights", [&] {
+
+        scalar_t scale = area_pixel_compute_scale<scalar_t>(
+            input_size, output_size, align_corners, opt_scale);
+
+        auto interp_size = HelperInterpLanczos::interp_size;
+
+        indices_weights = std::get<0>(HelperInterpLanczos::_compute_index_ranges_weights<scalar_t>(
+            input_size,
+            output_size,
+            stride,
+            ndims,
+            reshape_dim,
+            scale,
+            interp_size,
+            &HelperInterpLanczos::aa_filter<scalar_t>,
+            /*antialias=*/antialias,
+            /*align_corners=*/align_corners));
+      }
+    );
+    return indices_weights;
+  }
+
+  static inline std::tuple<std::vector<Tensor>, int, unsigned int> compute_index_ranges_int16_weights(
+    int64_t input_size,
+    int64_t output_size,
+    int64_t stride,
+    int64_t ndims,
+    int64_t reshape_dim,
+    bool align_corners,
+    const std::optional<double>& opt_scale,
+    bool antialias,
+    bool align_i32=false
+  ) {
+
+    auto interp_size = HelperInterpLanczos::interp_size;
+    auto fn = HelperInterpLanczos::aa_filter<double>;
+    return HelperInterpLanczos::_compute_index_ranges_int16_weights(
+        input_size, output_size, stride, ndims, reshape_dim,
+        align_corners, opt_scale, interp_size, fn, antialias, align_i32);
+  }
+
+};
+
 // Generic upsampling interpolation kernel for N-d case.
 // Input is assumed to be like NCHW, NCL, NCKHW - interpolated spatial dimension
 // are those from the end up to batch size N and number of channels C.
@@ -1418,7 +1508,7 @@ struct HelperInterpCubic : public HelperInterpBase {
 // - scale_type is template type for scales, typically std::optional<double>
 // - template<typename> class F is one of the above structs to compute indices and weights
 template <int out_ndims, typename scale_type, class F>
-void upsample_generic_Nd_kernel_impl(
+void upsample_non_separable_Nd_kernel_impl(
     const Tensor& output,
     const Tensor& input,
     bool align_corners,
@@ -1478,43 +1568,27 @@ void upsample_generic_Nd_kernel_impl(
 
   if (interp_size > 1) {
     // Nearest also supports uint8 tensor, so need to handle it separately
+    // Dispatch name should be "upsample_non_separable" but we keep the old
+    // name for internal BC.
     AT_DISPATCH_FLOATING_TYPES_AND2(
         kBFloat16, kHalf, iter.dtype(), "upsample_generic_Nd", [&] {
         // MSVC can not catch constexpr int interp_size here
         constexpr int mode = F::interp_size;
-        cpu_upsample_generic<scalar_t, out_ndims, mode>(iter);
+        upsample_non_separable<scalar_t, out_ndims, mode>(iter);
     });
   } else {
+    // Dispatch name should be "upsample_non_separable" but we keep the old
+    // name for internal BC.
     AT_DISPATCH_FLOATING_TYPES_AND3(kByte, kBFloat16, kHalf,
         iter.dtype(), "upsample_generic_Nd", [&] {
         constexpr int mode = F::interp_size;
-        cpu_upsample_generic<scalar_t, out_ndims, mode>(iter);
+        upsample_non_separable<scalar_t, out_ndims, mode>(iter);
     });
   }
 }
 
-template <typename scalar_t, bool is_horizontal>
-void cpu_upsample_generic_aa(at::TensorIterator& iter, unsigned int weights_precision) {
-
-  auto loop = [&](char** data, const int64_t* strides, int64_t n) {
-    if constexpr (is_horizontal) {
-
-      // Strides are : X 0 | 8 8 8 0 8  (Channels first)
-      // Strides are : X X | 0 0 0 0 0  (Channels last)
-      basic_loop_aa_horizontal<scalar_t>(data, strides, n, weights_precision);
-    } else {
-      // Strides are : X Y | 0 0 0 0 0 (Channels first)
-      // Strides are : X X | 0 0 0 0 0 (Channels last)
-      // upsampling data between contiguous dimensions (aka vertical resampling)
-      basic_loop_aa_vertical<scalar_t>(data, strides, n, weights_precision);
-    }
-  };
-
-  iter.for_each(loop);
-}
-
 template <int out_ndims, typename scale_type, class F, bool is_horizontal>
-void _separable_upsample_generic_Nd_kernel_impl_single_dim(
+void upsample_separable_1d(
     const Tensor& output,
     const Tensor& input,
     int interp_dim,
@@ -1543,8 +1617,8 @@ void _separable_upsample_generic_Nd_kernel_impl_single_dim(
   unsigned int weights_precision = 0;
 
   if (input_scalar_type == at::kByte) {
-    // This is a special branch to provide uint8 dtype support for bilinear and bicubic modes only
-    TORCH_INTERNAL_ASSERT(F::interp_size == 2 || F::interp_size == 4);
+    // This is a special branch to provide uint8 dtype support for bilinear, bicubic and lanczos modes only
+    TORCH_INTERNAL_ASSERT(F::interp_size == 2 || F::interp_size == 4 || F::interp_size == 6);
     int unused = 0;
     std::tie(indices_weights, unused, weights_precision) =
       F::compute_index_ranges_int16_weights(
@@ -1574,9 +1648,22 @@ void _separable_upsample_generic_Nd_kernel_impl_single_dim(
 
   auto iter = config.build();
 
+  // Dispatch name should be "upsample_separable_1d" but we keep the old
+  // name for internal BC.
   AT_DISPATCH_FLOATING_TYPES_AND(
       at::ScalarType::Byte, iter.dtype(), "upsample_generic_Nd_aa", [&] {
-        cpu_upsample_generic_aa<scalar_t, is_horizontal>(iter, weights_precision);
+        auto loop = [&](char** data, const int64_t* strides, int64_t n) {
+          if constexpr (is_horizontal) {
+            // Strides are : X 0 | 8 8 8 0 8  (Channels first)
+            // Strides are : X X | 0 0 0 0 0  (Channels last)
+            basic_loop_separable_1d_horizontal<scalar_t>(data, strides, n, weights_precision);
+          } else {
+            // Strides are : X Y | 0 0 0 0 0 (Channels first)
+            // Strides are : X X | 0 0 0 0 0 (Channels last)
+            basic_loop_separable_1d_vertical<scalar_t>(data, strides, n, weights_precision);
+          }
+        };
+        iter.for_each(loop);
       });
 }
 
@@ -1585,7 +1672,7 @@ void _separable_upsample_generic_Nd_kernel_impl_single_dim(
 // (dtype == uint8 and mode in ("bilinear", "bicubic")): this is used as
 // fallback in these settings when AVX isn't supported.
 template <int out_ndims, typename scale_type, class F>
-void separable_upsample_generic_Nd_kernel_impl(
+void upsample_separable_Nd_kernel_impl(
     const Tensor& output,
     const Tensor& input,
     bool align_corners,
@@ -1604,29 +1691,29 @@ void separable_upsample_generic_Nd_kernel_impl(
   at::Tensor temp_output, temp_input = input;
 
   int interp_dim = 0;
-  // Precompute the number of single dim resize method invocations
+  // Precompute the number of 1d resize ops
   // to avoid copying temporary buffer to output
-  int num_single_dim_ops = 0;
+  int num_1d_ops = 0;
   for (const auto i : c10::irange(out_ndims)) {
     interp_dim = 2 + out_ndims - 1 - i;
     if (output_shape[interp_dim] != input_shape[interp_dim]) {
-      num_single_dim_ops += 1;
+      num_1d_ops += 1;
     }
   }
 
-  // upsampling data within the contiguous dimension (aka horizontal resampling)
+  // Horizontal resampling (last spatial dim, i.e. W)
   interp_dim = 2 + out_ndims - 1;
   if (output_shape[interp_dim] != input_shape[interp_dim]) {
 
-    num_single_dim_ops -= 1;
-    if (num_single_dim_ops > 0) {
+    num_1d_ops -= 1;
+    if (num_1d_ops > 0) {
       temp_oshape[interp_dim] = output_shape[interp_dim];
       temp_output = at::empty(temp_oshape, input.options());
     } else {
       temp_output = output;
     }
 
-    _separable_upsample_generic_Nd_kernel_impl_single_dim<
+    upsample_separable_1d<
         out_ndims,
         scale_t,
         F,
@@ -1635,20 +1722,20 @@ void separable_upsample_generic_Nd_kernel_impl(
     temp_input = temp_output;
   }
 
-  // upsampling data between contiguous dimensions (aka vertical resampling)
+  // Vertical resampling (remaining spatial dims, e.g. H, D)
   for (const auto i : c10::irange(1, out_ndims)) {
     interp_dim = 2 + out_ndims - 1 - i;
     if (output_shape[interp_dim] != input_shape[interp_dim]) {
 
-      num_single_dim_ops -= 1;
-      if (num_single_dim_ops > 0) {
+      num_1d_ops -= 1;
+      if (num_1d_ops > 0) {
         temp_oshape[interp_dim] = output_shape[interp_dim];
         temp_output = at::empty(temp_oshape, input.options());
       } else {
         temp_output = output;
       }
 
-      _separable_upsample_generic_Nd_kernel_impl_single_dim<
+      upsample_separable_1d<
           out_ndims,
           scale_t,
           F,
@@ -1663,7 +1750,7 @@ void upsample_nearest1d_kernel_impl(
     const Tensor& output,
     const Tensor& input,
     std::optional<double> scales_w) {
-  upsample_generic_Nd_kernel_impl<1, scale_t, HelperInterpNearest>(
+  upsample_non_separable_Nd_kernel_impl<1, scale_t, HelperInterpNearest>(
       output, input, false, {scales_w});
 }
 
@@ -1671,26 +1758,28 @@ void _upsample_nearest_exact1d_kernel_impl(
     const Tensor& output,
     const Tensor& input,
     std::optional<double> scales_w) {
-  upsample_generic_Nd_kernel_impl<1, scale_t, HelperInterpNearestExact>(
+  upsample_non_separable_Nd_kernel_impl<1, scale_t, HelperInterpNearestExact>(
     output, input, false, {scales_w});
 }
 
-int _use_vectorized_kernel_cond_2d(
+int _use_channels_last_kernel_2d(
     const Tensor& output,
     const Tensor& input) {
-      // This condition is used to know whether we should dispatch to a vectorized
-      // kernel, or to the more general upsample_generic_Nd_kernel_impl(). For now,
-      // the vectorized kernels are only optimized for channels_last and when C >= 4
-      // (shape = NCHW). For a very wide range of use-cases (typically image or mask
-      // resizing where we have C < 4), using upsample_generic_Nd_kernel_impl() is
-      // actually faster. On top of that, benchmarks showed that this also depends on
-      // the *output* size (output_H + output_W), for both upsampling and
-      // downsampling. The current 128 threshold was determined through benchmarks.
+      // This condition is used to know whether we should dispatch to a
+      // channels-last-optimized kernel, or to the more general
+      // upsample_non_separable_Nd_kernel_impl(). For now, the channels-last kernels
+      // are only optimized for channels_last and when C >= 4 (shape = NCHW).
+      // For a very wide range of use-cases (typically image or mask resizing
+      // where we have C < 4), using upsample_non_separable_Nd_kernel_impl() is
+      // actually faster. On top of that, benchmarks showed that this also
+      // depends on the *output* size (output_H + output_W), for both
+      // upsampling and downsampling. The current 128 threshold was determined
+      // through benchmarks.
       return ((input.is_contiguous(at::MemoryFormat::ChannelsLast)) && (input.size(1) > 3)) || ((output.size(-2) + output.size(-1)) <= 128);
 }
 
-int _use_vectorized_kernel_cond_3d(
-    // Similar to _use_vectorized_kernel_cond_2d() but for 3d resampling (e.g. videos)
+int _use_channels_last_kernel_3d(
+    // Similar to _use_channels_last_kernel_2d() but for 3d resampling (e.g. videos)
     // Note that unlike the 2d case, this is not subject to small output size
     // overhead - hence the absence of the 128 threshold in the condition.
     const Tensor& output,
@@ -1704,13 +1793,13 @@ void upsample_nearest2d_kernel_impl(
     const Tensor& input,
     std::optional<double> scales_h,
     std::optional<double> scales_w) {
-  if (_use_vectorized_kernel_cond_2d(output, input)) {
+  if (_use_channels_last_kernel_2d(output, input)) {
     AT_DISPATCH_FLOATING_TYPES_AND3(kByte, kBFloat16, kHalf,
         input.scalar_type(), "upsample_nearest2d_channels_last", [&] {
-      cpu_upsample_nearest_channels_last<scalar_t, scale_t, nearest_idx>(output, input, {scales_h, scales_w});
+      upsample_nearest_channels_last<scalar_t, scale_t, nearest_idx>(output, input, {scales_h, scales_w});
     });
   } else {
-    upsample_generic_Nd_kernel_impl<2, scale_t, HelperInterpNearest>(
+    upsample_non_separable_Nd_kernel_impl<2, scale_t, HelperInterpNearest>(
       output, input, false, {scales_h, scales_w});
   }
 }
@@ -1720,12 +1809,12 @@ void _upsample_nearest_exact2d_kernel_impl(
     const Tensor& input,
     std::optional<double> scales_h,
     std::optional<double> scales_w) {
-  if (_use_vectorized_kernel_cond_2d(output, input)) {
+  if (_use_channels_last_kernel_2d(output, input)) {
     AT_DISPATCH_FLOATING_TYPES_AND3(kByte, kBFloat16, kHalf, input.scalar_type(), "upsample_nearest2d_channels_last", [&] {
-      cpu_upsample_nearest_channels_last<scalar_t, scale_t, nearest_exact_idx>(output, input, {scales_h, scales_w});
+      upsample_nearest_channels_last<scalar_t, scale_t, nearest_exact_idx>(output, input, {scales_h, scales_w});
     });
   } else {
-    upsample_generic_Nd_kernel_impl<2, scale_t, HelperInterpNearestExact>(
+    upsample_non_separable_Nd_kernel_impl<2, scale_t, HelperInterpNearestExact>(
       output, input, false, {scales_h, scales_w});
   }
 }
@@ -1736,13 +1825,13 @@ void upsample_nearest3d_kernel_impl(
     std::optional<double> scales_d,
     std::optional<double> scales_h,
     std::optional<double> scales_w) {
-  if (_use_vectorized_kernel_cond_3d(output, input)) {
+  if (_use_channels_last_kernel_3d(output, input)) {
     AT_DISPATCH_FLOATING_TYPES_AND3(kByte, kBFloat16, kHalf,
         input.scalar_type(), "upsample_nearest3d_channels_last", [&] {
-      cpu_upsample_nearest_channels_last<scalar_t, scale_t, nearest_idx>(output, input, {scales_d, scales_h, scales_w});
+      upsample_nearest_channels_last<scalar_t, scale_t, nearest_idx>(output, input, {scales_d, scales_h, scales_w});
     });
   } else {
-    upsample_generic_Nd_kernel_impl<3, scale_t, HelperInterpNearest>(
+    upsample_non_separable_Nd_kernel_impl<3, scale_t, HelperInterpNearest>(
       output, input, false, {scales_d, scales_h, scales_w});
   }
 }
@@ -1753,12 +1842,12 @@ void _upsample_nearest_exact3d_kernel_impl(
     std::optional<double> scales_d,
     std::optional<double> scales_h,
     std::optional<double> scales_w) {
-  if (_use_vectorized_kernel_cond_3d(output, input)) {
+  if (_use_channels_last_kernel_3d(output, input)) {
     AT_DISPATCH_FLOATING_TYPES_AND3(kByte, kBFloat16, kHalf, input.scalar_type(), "upsample_nearest3d_channels_last", [&] {
-      cpu_upsample_nearest_channels_last<scalar_t, scale_t, nearest_exact_idx>(output, input, {scales_d, scales_h, scales_w});
+      upsample_nearest_channels_last<scalar_t, scale_t, nearest_exact_idx>(output, input, {scales_d, scales_h, scales_w});
     });
   } else {
-    upsample_generic_Nd_kernel_impl<3, scale_t, HelperInterpNearestExact>(
+    upsample_non_separable_Nd_kernel_impl<3, scale_t, HelperInterpNearestExact>(
       output, input, false, {scales_d, scales_h, scales_w});
   }
 }
@@ -1768,7 +1857,7 @@ void upsample_linear1d_kernel_impl(
     const Tensor& input,
     bool align_corners,
     std::optional<double> scales_w) {
-  upsample_generic_Nd_kernel_impl<1, scale_t, HelperInterpLinear>(
+  upsample_non_separable_Nd_kernel_impl<1, scale_t, HelperInterpLinear>(
     output, input, align_corners, {scales_w});
 }
 
@@ -1780,17 +1869,17 @@ void upsample_bilinear2d_kernel_impl_float(
     std::optional<double> scales_h,
     std::optional<double> scales_w) {
 
-  // See note above about _use_vectorized_kernel_cond_2d(output, input). The extra cond is present
+  // See note above about _use_channels_last_kernel_2d(output, input). The extra cond is present
   // because benchmarks showed that with only 1 thread, images (C == 3) were
-  // slightly faster with the vectorized kernel than with the generic one.
+  // slightly faster with the channels-last kernel than with the generic one.
   // That's not the case for masks though (C == 1), which strongly benefit from
   // using the generic kernel.
-  if ((_use_vectorized_kernel_cond_2d(output, input)) || (at::get_num_threads() == 1 && input.size(1) == 3)) {
+  if ((_use_channels_last_kernel_2d(output, input)) || (at::get_num_threads() == 1 && input.size(1) == 3)) {
     AT_DISPATCH_FLOATING_TYPES_AND2(kBFloat16, kHalf, input.scalar_type(), "upsample_bilinear2d_channels_last", [&] {
-      cpu_upsample_linear_channels_last<scalar_t, scale_t>(output, input, align_corners, {scales_h, scales_w});
+      upsample_linear_channels_last<scalar_t, scale_t>(output, input, align_corners, {scales_h, scales_w});
     });
   } else {
-    upsample_generic_Nd_kernel_impl<2, scale_t, HelperInterpLinear>(
+    upsample_non_separable_Nd_kernel_impl<2, scale_t, HelperInterpLinear>(
       output, input, align_corners, {scales_h, scales_w});
   }
 }
@@ -1802,25 +1891,27 @@ void upsample_bilinear2d_kernel_impl(
     std::optional<double> scales_h,
     std::optional<double> scales_w) {
 
-  if (input.dtype() == at::kByte){
+  if (input.dtype() == at::kByte) {
     #ifdef CPU_CAPABILITY_AVX2
       if (input.size(1) <= 4) {
-        upsample_avx_bilinear_bicubic_uint8<scale_t, HelperInterpLinear>(input,
+        return upsample_avx_bilinear_bicubic_uint8<scale_t, HelperInterpLinear>(input,
           output, align_corners, {scales_h, scales_w},
           /*antialias=*/false);
-      } else {
-        separable_upsample_generic_Nd_kernel_impl<2, scale_t, HelperInterpLinear>(
-          output, input, align_corners, {scales_h, scales_w},
+      }
+    #elif defined(__aarch64__)
+      if (input.size(1) == 3
+          && input.is_contiguous(at::MemoryFormat::ChannelsLast)
+          && output.is_contiguous(at::MemoryFormat::ChannelsLast)) {
+        return upsample_neon_bilinear_bicubic_uint8<scale_t, HelperInterpLinear>(
+          input, output, align_corners, {scales_h, scales_w},
           /*antialias=*/false);
       }
-    #else  // CPU_CAPABILITY_AVX2
-      separable_upsample_generic_Nd_kernel_impl<2, scale_t, HelperInterpLinear>(
-        output, input, align_corners, {scales_h, scales_w},
-        /*antialias=*/false);
     #endif  // CPU_CAPABILITY_AVX2
-  } else {
-    upsample_bilinear2d_kernel_impl_float(output, input, align_corners, scales_h, scales_w);
+    return upsample_separable_Nd_kernel_impl<2, scale_t, HelperInterpLinear>(
+      output, input, align_corners, {scales_h, scales_w},
+      /*antialias=*/false);
   }
+  return upsample_bilinear2d_kernel_impl_float(output, input, align_corners, scales_h, scales_w);
 }
 
 
@@ -1830,21 +1921,26 @@ void upsample_bilinear2d_aa_kernel_impl(
     bool align_corners,
     std::optional<double> scales_h,
     std::optional<double> scales_w) {
-#ifdef CPU_CAPABILITY_AVX2
-  if (input.dtype() == at::kByte && input.size(1) <= 4) {
-    upsample_avx_bilinear_bicubic_uint8<scale_t, HelperInterpLinear>(
-      input, output, align_corners, {scales_h, scales_w},
-      /*antialias=*/true);
-  } else {
-    separable_upsample_generic_Nd_kernel_impl<2, scale_t, HelperInterpLinear>(
-        output, input, align_corners, {scales_h, scales_w},
-        /*antialias=*/true);
+  if (input.dtype() == at::kByte) {
+    #ifdef CPU_CAPABILITY_AVX2
+      if (input.size(1) <= 4) {
+        return upsample_avx_bilinear_bicubic_uint8<scale_t, HelperInterpLinear>(
+          input, output, align_corners, {scales_h, scales_w},
+          /*antialias=*/true);
+      }
+    #elif defined(__aarch64__)
+      if (input.size(1) == 3
+          && input.is_contiguous(at::MemoryFormat::ChannelsLast)
+          && output.is_contiguous(at::MemoryFormat::ChannelsLast)) {
+        return upsample_neon_bilinear_bicubic_uint8<scale_t, HelperInterpLinear>(
+          input, output, align_corners, {scales_h, scales_w},
+          /*antialias=*/true);
+      }
+    #endif  // CPU_CAPABILITY_AVX2
   }
-#else // CPU_CAPABILITY_AVX2
-  separable_upsample_generic_Nd_kernel_impl<2, scale_t, HelperInterpLinear>(
-      output, input, align_corners, {scales_h, scales_w},
-      /*antialias=*/true);
-#endif // CPU_CAPABILITY_AVX2
+  return upsample_separable_Nd_kernel_impl<2, scale_t, HelperInterpLinear>(
+    output, input, align_corners, {scales_h, scales_w},
+    /*antialias=*/true);
 }
 
 void upsample_trilinear3d_kernel_impl(
@@ -1854,12 +1950,12 @@ void upsample_trilinear3d_kernel_impl(
     std::optional<double> scales_d,
     std::optional<double> scales_h,
     std::optional<double> scales_w) {
-  if ((_use_vectorized_kernel_cond_3d(output, input))) {
+  if ((_use_channels_last_kernel_3d(output, input))) {
     AT_DISPATCH_FLOATING_TYPES_AND2(kBFloat16, kHalf, input.scalar_type(), "upsample_trilinear3d_channels_last", [&] {
-      cpu_upsample_linear_channels_last<scalar_t, scale_t>(output, input, align_corners, {scales_d, scales_h, scales_w});
+      upsample_linear_channels_last<scalar_t, scale_t>(output, input, align_corners, {scales_d, scales_h, scales_w});
     });
   } else {
-    upsample_generic_Nd_kernel_impl<3, scale_t, HelperInterpLinear>(
+    upsample_non_separable_Nd_kernel_impl<3, scale_t, HelperInterpLinear>(
       output, input, align_corners, {scales_d, scales_h, scales_w});
   }
 }
@@ -1871,27 +1967,28 @@ void upsample_bicubic2d_kernel_impl(
     std::optional<double> scales_h,
     std::optional<double> scales_w) {
 
-  if (input.dtype() == at::kByte){
+  if (input.dtype() == at::kByte) {
     #ifdef CPU_CAPABILITY_AVX2
       if (input.size(1) <= 4) {
-        upsample_avx_bilinear_bicubic_uint8<scale_t, HelperInterpCubic>(input,
+        return upsample_avx_bilinear_bicubic_uint8<scale_t, HelperInterpCubic>(input,
           output, align_corners, {scales_h, scales_w},
           /*antialias=*/false);
-      } else {
-        separable_upsample_generic_Nd_kernel_impl<2, scale_t, HelperInterpCubic>(
-          output, input, align_corners, {scales_h, scales_w},
+      }
+    #elif defined(__aarch64__)
+      if (input.size(1) == 3
+          && input.is_contiguous(at::MemoryFormat::ChannelsLast)
+          && output.is_contiguous(at::MemoryFormat::ChannelsLast)) {
+        return upsample_neon_bilinear_bicubic_uint8<scale_t, HelperInterpCubic>(
+          input, output, align_corners, {scales_h, scales_w},
           /*antialias=*/false);
       }
-    #else  // CPU_CAPABILITY_AVX2
-      separable_upsample_generic_Nd_kernel_impl<2, scale_t, HelperInterpCubic>(
-        output, input, align_corners, {scales_h, scales_w},
-        /*antialias=*/false);
     #endif  // CPU_CAPABILITY_AVX2
+    return upsample_separable_Nd_kernel_impl<2, scale_t, HelperInterpCubic>(
+      output, input, align_corners, {scales_h, scales_w},
+      /*antialias=*/false);
   }
-  else {
-    upsample_generic_Nd_kernel_impl<2, scale_t, HelperInterpCubic>(
-      output, input, align_corners, {scales_h, scales_w});
-  }
+  return upsample_non_separable_Nd_kernel_impl<2, scale_t, HelperInterpCubic>(
+    output, input, align_corners, {scales_h, scales_w});
 }
 
 void upsample_bicubic2d_aa_kernel_impl(
@@ -1901,28 +1998,62 @@ void upsample_bicubic2d_aa_kernel_impl(
     std::optional<double> scales_h,
     std::optional<double> scales_w) {
 
-#ifdef CPU_CAPABILITY_AVX2
-  if (input.dtype() == at::kByte && input.size(1) <= 4) {
-    upsample_avx_bilinear_bicubic_uint8<scale_t, HelperInterpCubic>(
-      input, output, align_corners, {scales_h, scales_w},
-      /*antialias=*/true);
-  } else {
-    separable_upsample_generic_Nd_kernel_impl<2, scale_t, HelperInterpCubic>(
-        output, input, align_corners, {scales_h, scales_w},
-        /*antialias=*/true);
+  if (input.dtype() == at::kByte) {
+    #ifdef CPU_CAPABILITY_AVX2
+      if (input.size(1) <= 4) {
+        return upsample_avx_bilinear_bicubic_uint8<scale_t, HelperInterpCubic>(
+          input, output, align_corners, {scales_h, scales_w},
+          /*antialias=*/true);
+      }
+    #elif defined(__aarch64__)
+      if (input.size(1) == 3
+          && input.is_contiguous(at::MemoryFormat::ChannelsLast)
+          && output.is_contiguous(at::MemoryFormat::ChannelsLast)) {
+        return upsample_neon_bilinear_bicubic_uint8<scale_t, HelperInterpCubic>(
+          input, output, align_corners, {scales_h, scales_w},
+          /*antialias=*/true);
+      }
+    #endif  // CPU_CAPABILITY_AVX2
   }
-#else // CPU_CAPABILITY_AVX2
-  separable_upsample_generic_Nd_kernel_impl<2, scale_t, HelperInterpCubic>(
-      output, input, align_corners, {scales_h, scales_w},
-      /*antialias=*/true);
-#endif // CPU_CAPABILITY_AVX2
+  return upsample_separable_Nd_kernel_impl<2, scale_t, HelperInterpCubic>(
+    output, input, align_corners, {scales_h, scales_w},
+    /*antialias=*/true);
+}
+
+void upsample_lanczos2d_aa_kernel_impl(
+    const Tensor& output,
+    const Tensor& input,
+    bool align_corners,
+    std::optional<double> scales_h,
+    std::optional<double> scales_w) {
+
+  if (input.dtype() == at::kByte) {
+    #ifdef CPU_CAPABILITY_AVX2
+      if (input.size(1) <= 4) {
+        return upsample_avx_bilinear_bicubic_uint8<scale_t, HelperInterpLanczos>(
+          input, output, align_corners, {scales_h, scales_w},
+          /*antialias=*/true);
+      }
+    #elif defined(__aarch64__)
+      if (input.size(1) == 3
+          && input.is_contiguous(at::MemoryFormat::ChannelsLast)
+          && output.is_contiguous(at::MemoryFormat::ChannelsLast)) {
+        return upsample_neon_bilinear_bicubic_uint8<scale_t, HelperInterpLanczos>(
+          input, output, align_corners, {scales_h, scales_w},
+          /*antialias=*/true);
+      }
+    #endif  // CPU_CAPABILITY_AVX2
+  }
+  return upsample_separable_Nd_kernel_impl<2, scale_t, HelperInterpLanczos>(
+    output, input, align_corners, {scales_h, scales_w},
+    /*antialias=*/true);
 }
 
 template <
     typename scalar_t,
     typename scale_type,
     class F>
-void cpu_upsample_genNd_backward_aa(
+void upsample_separable_Nd_backward_aa(
     const Tensor& grad_input_,
     const Tensor& grad_output_,
     bool align_corners,
@@ -2040,7 +2171,7 @@ void upsample_bilinear2d_aa_backward_kernel_impl(
     std::optional<double> scales_w) {
   AT_DISPATCH_FLOATING_TYPES(
       grad_output.scalar_type(), "upsample_bilinear2d_aa_backward_cpu", [&] {
-        cpu_upsample_genNd_backward_aa<scalar_t, scale_t, HelperInterpLinear>(
+        upsample_separable_Nd_backward_aa<scalar_t, scale_t, HelperInterpLinear>(
             grad_input, grad_output, align_corners, {scales_h, scales_w});
       });
 }
@@ -2053,7 +2184,20 @@ void upsample_bicubic2d_aa_backward_kernel_impl(
     std::optional<double> scales_w) {
   AT_DISPATCH_FLOATING_TYPES(
       grad_output.scalar_type(), "upsample_bicubic2d_aa_backward_cpu", [&] {
-        cpu_upsample_genNd_backward_aa<scalar_t, scale_t, HelperInterpCubic>(
+        upsample_separable_Nd_backward_aa<scalar_t, scale_t, HelperInterpCubic>(
+            grad_input, grad_output, align_corners, {scales_h, scales_w});
+      });
+}
+
+void upsample_lanczos2d_aa_backward_kernel_impl(
+    const Tensor& grad_input,
+    const Tensor& grad_output,
+    bool align_corners,
+    std::optional<double> scales_h,
+    std::optional<double> scales_w) {
+  AT_DISPATCH_FLOATING_TYPES(
+      grad_output.scalar_type(), "upsample_lanczos2d_aa_backward_cpu", [&] {
+        upsample_separable_Nd_backward_aa<scalar_t, scale_t, HelperInterpLanczos>(
             grad_input, grad_output, align_corners, {scales_h, scales_w});
       });
 }
@@ -2076,4 +2220,7 @@ REGISTER_DISPATCH(upsample_trilinear3d_kernel, &upsample_trilinear3d_kernel_impl
 REGISTER_DISPATCH(upsample_bicubic2d_kernel, &upsample_bicubic2d_kernel_impl)
 REGISTER_DISPATCH(_upsample_bicubic2d_aa_kernel, &upsample_bicubic2d_aa_kernel_impl)
 REGISTER_DISPATCH(_upsample_bicubic2d_aa_backward_kernel, &upsample_bicubic2d_aa_backward_kernel_impl)
+
+REGISTER_DISPATCH(_upsample_lanczos2d_aa_kernel, &upsample_lanczos2d_aa_kernel_impl)
+REGISTER_DISPATCH(_upsample_lanczos2d_aa_backward_kernel, &upsample_lanczos2d_aa_backward_kernel_impl)
 } // namespace at::native
diff --git a/aten/src/ATen/native/cpu/UpSampleKernelAVXAntialias.h b/aten/src/ATen/native/cpu/UpSampleKernelAVXAntialias.h
index c1bf79dfa44e6..3c351b05dc92f 100644
--- a/aten/src/ATen/native/cpu/UpSampleKernelAVXAntialias.h
+++ b/aten/src/ATen/native/cpu/UpSampleKernelAVXAntialias.h
@@ -176,7 +176,7 @@ void ImagingResampleHorizontal(
   //
 
   // TODO: we may want to merge that into the fallback code (currently called
-  // basic_loop_aa_horizontal<uint8_t>)
+  // basic_loop_separable_1d_horizontal<uint8_t>)
   // Although this may not be needed if / when we port all this code to use
   // Vec.h since this would potentially give us another fall-back implem
 
@@ -252,7 +252,7 @@ void ImagingResampleVertical(
   //   oB[xoffset + i] = b[xoffset + ymin[i]] * w[i, 0] + ... + b[xoffset + ymin[i] + (K-1) * xsize] * w[i, K-1]
 
   // TODO: we may want to merge that into the fallback code (currently called
-  // basic_loop_aa_vertical<uint8_t>)
+  // basic_loop_separable_1d_vertical<uint8_t>)
   // Although this may not be needed if / when we port all this code to use
   // Vec.h since this would potentially give us another fall-back implem
   const int16_t* kk = (int16_t*)(vert_indices_weights[3].const_data_ptr<double>());
@@ -289,7 +289,7 @@ void ImagingResampleVertical(
 // mode for uint8 dtype when C <= 4, with or without antialias. The
 // implem is based on PIL-SIMD.
 // Its equivalent implementation (fallback) for when AVX isn't supported or when
-// C > 4 is separable_upsample_generic_Nd_kernel_impl()  There are a bunch of
+// C > 4 is upsample_separable_Nd_kernel_impl()  There are a bunch of
 // future improvement that can be done: look for the TODOs in this file.
 // For details on how the weights are computed and how the multiplications are
 // run on int (instead of float weights), see
diff --git a/aten/src/ATen/native/cpu/UpSampleKernelNEONAntialias.h b/aten/src/ATen/native/cpu/UpSampleKernelNEONAntialias.h
new file mode 100644
index 0000000000000..c8d3a1a59a508
--- /dev/null
+++ b/aten/src/ATen/native/cpu/UpSampleKernelNEONAntialias.h
@@ -0,0 +1,381 @@
+// NEON-optimized uint8 bilinear resize with antialiasing for aarch64.
+//
+// This is the NEON counterpart of UpSampleKernelAVXAntialias.h.
+// It only supports num_channels == 3 with channels-last input.
+
+#pragma once
+#if defined(__aarch64__)
+
+#include <ATen/core/Tensor.h>
+#include <arm_neon.h>
+#include <c10/util/irange.h>
+
+#ifndef AT_PER_OPERATOR_HEADERS
+#include <ATen/Functions.h>
+#else
+#include <ATen/ops/empty.h>
+#endif
+
+namespace {
+
+// Emulate SSE _mm_madd_epi16: multiply 8 int16 pairs element-wise, then add
+// adjacent 32-bit products pairwise, producing 4 int32 results.
+// [a0*b0+a1*b1, a2*b2+a3*b3, a4*b4+a5*b5, a6*b6+a7*b7]
+inline int32x4_t neon_madd_s16(int16x8_t a, int16x8_t b) {
+  int32x4_t prod_low = vmull_s16(vget_low_s16(a), vget_low_s16(b));
+  int32x4_t prod_high = vmull_s16(vget_high_s16(a), vget_high_s16(b));
+  return vpaddq_s32(prod_low, prod_high);
+}
+
+// Interpolation horizontal pass: compute x-axis interpolation outputs.
+//
+// Input data is interleaved RGB:
+//   input = [r0, g0, b0, r1, g1, b1, r2, g2, b2, ...]
+// Weights are float values rescaled to int16:
+//   weights[i] = [w[i, 0], w[i, 1], ..., w[i, K-1]]
+// For each output pixel i, we compute:
+//   oR[i] = r[xmin[i]] * w[i,0] + ... + r[xmin[i]+K-1] * w[i,K-1]
+//   oG[i] = g[xmin[i]] * w[i,0] + ... + g[xmin[i]+K-1] * w[i,K-1]
+//   oB[i] = b[xmin[i]] * w[i,0] + ... + b[xmin[i]+K-1] * w[i,K-1]
+//
+// Processes one row at a time.
+void NeonResampleHorizontal(const at::Tensor& unpacked_output,
+                            const at::Tensor& unpacked_input,
+                            int ksize,
+                            const std::vector<at::Tensor>& horiz_indices_weights,
+                            unsigned int horiz_weights_precision) {
+  const auto* kk = (const int16_t*)(horiz_indices_weights[3].const_data_ptr<double>());
+
+  auto xout = unpacked_output.size(2);
+  auto yin = unpacked_output.size(1);
+  auto xin = unpacked_input.size(2);
+  const auto num_channels = unpacked_input.size(0);
+  TORCH_INTERNAL_ASSERT(num_channels == 3);
+
+  const int64_t* idx_ptr_xmin = horiz_indices_weights[0].const_data_ptr<int64_t>();
+  const int64_t* idx_ptr_size = horiz_indices_weights[1].const_data_ptr<int64_t>();
+
+  uint8_t* output_p = unpacked_output.data_ptr<uint8_t>();
+  const uint8_t* input_p = unpacked_input.const_data_ptr<uint8_t>();
+
+  auto xout_stride = xout * num_channels;
+  auto xin_stride = xin * num_channels;
+
+  for (int64_t yy = 0; yy < yin; yy++) {
+    uint8_t* C10_RESTRICT lineOut = output_p + yy * xout_stride;
+    const uint8_t* C10_RESTRICT lineIn = input_p + yy * xin_stride;
+    const int32_t initial_val = 1 << (horiz_weights_precision - 1);
+
+    // Horizontal convolution for a single row.
+    //
+    // For each output pixel out_x, computes the weighted sum of ids_size input
+    // pixels using the weight vector k[]. Uses vld3_u8 to load 8 interleaved RGB
+    // pixels and deinterleave them into separate R, G, B vectors. This avoids
+    // the shuffle masks needed by the AVX2 path.
+    //
+    // The weight vector is split as: ids_size = num_blocks_8 * 8 + remainder.
+    // We process 8 weights in a vectorized loop, then handle remaining weights
+    // in a scalar cleanup loop. Accumulators are kept per-channel in int32 to
+    // avoid overflow.
+    for (int64_t out_x = 0; out_x < xout; out_x++) {
+      // ids_min is a byte offset (pre-multiplied by num_channels)
+      const int64_t ids_min = idx_ptr_xmin[out_x];
+      const int64_t ids_size = idx_ptr_size[out_x];
+      const int16_t* k = &kk[out_x * ksize];
+      const uint8_t* lineIn_min = lineIn + ids_min;
+
+      int32x4_t acc_r = vdupq_n_s32(0);
+      int32x4_t acc_g = vdupq_n_s32(0);
+      int32x4_t acc_b = vdupq_n_s32(0);
+
+      int64_t i = 0;
+
+      // Block 8: load 8 RGB pixels with vld3, deinterleave, widen to int16,
+      // multiply-accumulate with 8 weights per channel.
+      // Each vmlal_s16 accumulates 4 products; calling it on the low and high
+      // halves puts all 8 products into 4 int32 lanes, reduced later by
+      // vaddvq_s32.
+      for (; i + 8 <= ids_size; i += 8) {
+        uint8x8x3_t rgb = vld3_u8(lineIn_min + num_channels * i);
+        int16x8_t weights = vld1q_s16(&k[i]);
+
+        int16x8_t r16 = vreinterpretq_s16_u16(vmovl_u8(rgb.val[0]));
+        int16x8_t g16 = vreinterpretq_s16_u16(vmovl_u8(rgb.val[1]));
+        int16x8_t b16 = vreinterpretq_s16_u16(vmovl_u8(rgb.val[2]));
+
+        acc_r = vmlal_s16(acc_r, vget_low_s16(r16), vget_low_s16(weights));
+        acc_r = vmlal_s16(acc_r, vget_high_s16(r16), vget_high_s16(weights));
+        acc_g = vmlal_s16(acc_g, vget_low_s16(g16), vget_low_s16(weights));
+        acc_g = vmlal_s16(acc_g, vget_high_s16(g16), vget_high_s16(weights));
+        acc_b = vmlal_s16(acc_b, vget_low_s16(b16), vget_low_s16(weights));
+        acc_b = vmlal_s16(acc_b, vget_high_s16(b16), vget_high_s16(weights));
+      }
+
+      // Block 4: handle 4 pixels that didn't fit in a block of 8
+      // We also use vld3_u8 here, which still loads 8 pixels, but we only use
+      // the lower half - so the computation is correct.
+      // On all rows except the last one, reading 8 pixels is safe (the tensors
+      // are channels-last). But on the last row, we have to be careful not to
+      // read past the buffer, hence the extra boundary check.
+      const uint8_t* block_of_4_safe_load_end = input_p + yin * xin_stride - 24;
+      for (; i + 4 <= ids_size && lineIn_min + num_channels * i <= block_of_4_safe_load_end; i += 4) {
+        uint8x8x3_t rgb = vld3_u8(lineIn_min + num_channels * i);
+        int16x4_t weights4 = vld1_s16(&k[i]);
+
+        int16x4_t r16 = vget_low_s16(vreinterpretq_s16_u16(vmovl_u8(rgb.val[0])));
+        int16x4_t g16 = vget_low_s16(vreinterpretq_s16_u16(vmovl_u8(rgb.val[1])));
+        int16x4_t b16 = vget_low_s16(vreinterpretq_s16_u16(vmovl_u8(rgb.val[2])));
+
+        acc_r = vmlal_s16(acc_r, r16, weights4);
+        acc_g = vmlal_s16(acc_g, g16, weights4);
+        acc_b = vmlal_s16(acc_b, b16, weights4);
+      }
+
+      // Horizontal reduction + rounding bias
+      int32_t sum_r = vaddvq_s32(acc_r) + initial_val;
+      int32_t sum_g = vaddvq_s32(acc_g) + initial_val;
+      int32_t sum_b = vaddvq_s32(acc_b) + initial_val;
+
+      // Scalar cleanup for remaining pixels
+      for (; i < ids_size; i++) {
+        int16_t w = k[i];
+        const uint8_t* p = lineIn_min + num_channels * i;
+        sum_r += w * p[0];
+        sum_g += w * p[1];
+        sum_b += w * p[2];
+      }
+
+      // Right-shift and clamp to [0, 255]
+      uint8_t* out = lineOut + num_channels * out_x;
+      out[0] = static_cast<uint8_t>(std::clamp(sum_r >> horiz_weights_precision, 0, 255));
+      out[1] = static_cast<uint8_t>(std::clamp(sum_g >> horiz_weights_precision, 0, 255));
+      out[2] = static_cast<uint8_t>(std::clamp(sum_b >> horiz_weights_precision, 0, 255));
+    }
+  }
+}
+
+// Interpolation vertical pass: compute y-axis interpolation outputs.
+//
+// Input data is interleaved RGB:
+//   input = [r0, g0, b0, r1, g1, b1, ...]
+// For each output row yy, we compute a weighted sum of ids_size input rows:
+//   oC[x] = C[ymin[yy] * xstride + x] * w[yy,0] + ... + C[(ymin[yy]+K-1) * xstride + x] * w[yy,K-1]
+// for each channel C in {R, G, B} and each horizontal position x.
+void NeonResampleVertical(const at::Tensor& unpacked_output,
+                          const at::Tensor& unpacked_input,
+                          int ksize,
+                          const std::vector<at::Tensor>& vert_indices_weights,
+                          unsigned int vert_weights_precision) {
+  const auto* kk = (const int16_t*)(vert_indices_weights[3].const_data_ptr<double>());
+
+  const int64_t* idx_ptr_xmin = vert_indices_weights[0].const_data_ptr<int64_t>();
+  const int64_t* idx_ptr_size = vert_indices_weights[1].const_data_ptr<int64_t>();
+
+  uint8_t* output_p = unpacked_output.data_ptr<uint8_t>();
+  const uint8_t* input_p = unpacked_input.const_data_ptr<uint8_t>();
+
+  auto xout = unpacked_output.size(2);
+  auto yout = unpacked_output.size(1);
+  const auto num_channels = unpacked_input.size(0);
+  TORCH_INTERNAL_ASSERT(num_channels == unpacked_output.size(0));
+
+  const int64_t data_size = xout * num_channels;
+
+  for (const auto yy : c10::irange(yout)) {
+    const auto* k = &kk[yy * ksize];
+    auto ids_min = idx_ptr_xmin[yy];
+    auto ids_size = idx_ptr_size[yy];
+
+    uint8_t* C10_RESTRICT lineOut = output_p + yy * data_size;
+    const int32_t initial_val = 1 << (vert_weights_precision - 1);
+    const int32x4_t initial = vdupq_n_s32(initial_val);
+
+    // Vertical convolution for one output row.
+    // Computes a weighted sum of ids_size input rows for all x positions.
+    //
+    // The data is treated as a flat byte array of size xout * num_channels.
+    // We process 16 bytes at a time in the NEON path, then fall back to scalar
+    // for the remainder. The vertical pass doesn't need channel deinterleaving
+    // because the same weight applies to all channels at a given (x, y) position.
+    //
+    // To process 2 weights at a time, we interleave pixels from two input rows
+    // using vzipq_u8, then use neon_madd_s16 (which pairs adjacent int16 values,
+    // multiplies, and adds) so that each pair [pixel_row_i, pixel_row_i+1] is
+    // multiplied by [w_i, w_i+1] and summed in one operation.
+    int64_t j = 0;
+
+    // Process 16 bytes at a time using 4 accumulators (sss0..sss3),
+    // each holding 4 int32 partial sums.
+    for (; j + 16 <= data_size; j += 16) {
+      int32x4_t sss0 = initial;
+      int32x4_t sss1 = initial;
+      int32x4_t sss2 = initial;
+      int32x4_t sss3 = initial;
+
+      const uint8_t* lineIn_min = input_p + j + ids_min;
+
+      // Process 2 weights at a time by interleaving rows.
+      // For weights w0, w1 and pixel bytes from row_i and row_{i+1}:
+      //   mmk = [w0, w1, w0, w1, w0, w1, w0, w1]
+      //   vzipq interleaves: [r0_0, r1_0, r0_1, r1_1, ...] (from 2 rows)
+      //   neon_madd_s16 computes: r0_0*w0 + r1_0*w1, r0_1*w0 + r1_1*w1, ...
+      int64_t i = 0;
+      for (; i + 1 < ids_size; i += 2) {
+        int16_t w0 = k[i];
+        int16_t w1 = k[i + 1];
+        int16x8_t mmk = {w0, w1, w0, w1, w0, w1, w0, w1};
+
+        uint8x16_t src1 = vld1q_u8(lineIn_min + i * data_size);
+        uint8x16_t src2 = vld1q_u8(lineIn_min + (i + 1) * data_size);
+
+        uint8x16x2_t interleaved = vzipq_u8(src1, src2);
+
+        int16x8_t pix0 = vreinterpretq_s16_u16(vmovl_u8(vget_low_u8(interleaved.val[0])));
+        int16x8_t pix1 = vreinterpretq_s16_u16(vmovl_u8(vget_high_u8(interleaved.val[0])));
+        int16x8_t pix2 = vreinterpretq_s16_u16(vmovl_u8(vget_low_u8(interleaved.val[1])));
+        int16x8_t pix3 = vreinterpretq_s16_u16(vmovl_u8(vget_high_u8(interleaved.val[1])));
+
+        sss0 = vaddq_s32(sss0, neon_madd_s16(pix0, mmk));
+        sss1 = vaddq_s32(sss1, neon_madd_s16(pix1, mmk));
+        sss2 = vaddq_s32(sss2, neon_madd_s16(pix2, mmk));
+        sss3 = vaddq_s32(sss3, neon_madd_s16(pix3, mmk));
+      }
+
+      // Handle remaining single weight (when ids_size is odd)
+      for (; i < ids_size; i++) {
+        int16x8_t mmk = vdupq_n_s16(k[i]);
+
+        uint8x16_t src = vld1q_u8(lineIn_min + i * data_size);
+
+        int16x8_t pix_lo = vreinterpretq_s16_u16(vmovl_u8(vget_low_u8(src)));
+        int16x8_t pix_hi = vreinterpretq_s16_u16(vmovl_u8(vget_high_u8(src)));
+
+        sss0 = vaddq_s32(sss0, vmull_s16(vget_low_s16(pix_lo), vget_low_s16(mmk)));
+        sss1 = vaddq_s32(sss1, vmull_s16(vget_high_s16(pix_lo), vget_high_s16(mmk)));
+        sss2 = vaddq_s32(sss2, vmull_s16(vget_low_s16(pix_hi), vget_low_s16(mmk)));
+        sss3 = vaddq_s32(sss3, vmull_s16(vget_high_s16(pix_hi), vget_high_s16(mmk)));
+      }
+
+      // Right-shift accumulators by coefs_precision to convert fixed-point -> int
+      int32x4_t shift = vdupq_n_s32(-static_cast<int32_t>(vert_weights_precision));
+      sss0 = vshlq_s32(sss0, shift);
+      sss1 = vshlq_s32(sss1, shift);
+      sss2 = vshlq_s32(sss2, shift);
+      sss3 = vshlq_s32(sss3, shift);
+
+      // Narrow int32 -> int16 (with saturation), then int16 -> uint8 (with
+      // saturation), clamping to [0, 255]
+      int16x8_t narrow_lo = vcombine_s16(vqmovn_s32(sss0), vqmovn_s32(sss1));
+      int16x8_t narrow_hi = vcombine_s16(vqmovn_s32(sss2), vqmovn_s32(sss3));
+
+      // Store 16 output bytes
+      vst1_u8(lineOut + j, vqmovun_s16(narrow_lo));
+      vst1_u8(lineOut + j + 8, vqmovun_s16(narrow_hi));
+    }
+
+    // Scalar fallback for remaining bytes
+    for (; j < data_size; j++) {
+      int32_t sss = initial_val;
+      const uint8_t* lineIn_min = input_p + j + ids_min;
+
+      for (int64_t i = 0; i < ids_size; i++) {
+        sss += k[i] * static_cast<int32_t>(lineIn_min[i * data_size]);
+      }
+
+      sss >>= vert_weights_precision;
+      lineOut[j] = static_cast<uint8_t>(std::clamp(sss, 0, 255));
+    }
+  }
+}
+
+// Main entry point for NEON-accelerated uint8 bilinear resize.
+//
+// Only supports num_channels == 3 with channels-last memory format.
+// Mirrors upsample_avx_bilinear_bicubic_uint8 but uses NEON intrinsics and
+// works directly on interleaved RGB data (no unpack/pack step needed).
+//
+// Weights are computed as int16 values by compute_index_ranges_int16_weights
+// with align_i32=false (no int32 alignment padding needed, unlike AVX2 which
+// processes pixels as 4-byte RGBA units).
+template <typename scale_type, class F>
+void upsample_neon_bilinear_bicubic_uint8(const at::Tensor& input_,
+                                          const at::Tensor& output,
+                                          bool align_corners,
+                                          const scale_type& scales,
+                                          bool antialias) {
+  auto batch_size = input_.size(0);
+  auto num_channels = input_.size(1);
+  auto xin = input_.size(3);
+  auto yin = input_.size(2);
+  auto xout = output.size(3);
+  auto yout = output.size(2);
+
+  if (xin == xout && yin == yout) {
+    output.copy_(input_);
+    return;
+  }
+
+  TORCH_INTERNAL_ASSERT(num_channels == 3);
+  TORCH_INTERNAL_ASSERT(output.is_contiguous(at::MemoryFormat::ChannelsLast));
+
+  auto input = input_.contiguous(at::MemoryFormat::ChannelsLast);
+
+  auto need_horizontal = xout != xin;
+  auto need_vertical = yout != yin;
+
+  int ksize_horiz, ksize_vert;
+  std::vector<at::Tensor> horiz_indices_weights, vert_indices_weights;
+  unsigned int horiz_weights_precision, vert_weights_precision;
+
+  if (need_horizontal) {
+    int interp_dim = 3;
+    std::tie(horiz_indices_weights, ksize_horiz, horiz_weights_precision) = F::compute_index_ranges_int16_weights(
+        /*input_size=*/xin,
+        /*output_size=*/xout,
+        /*stride=*/num_channels,
+        /*ndims=*/4,
+        /*reshape_dim=*/interp_dim,
+        /*align_corners=*/align_corners,
+        /*opt_scale=*/scales[interp_dim - 2],
+        /*antialias=*/antialias,
+        /*align_i32=*/false);
+  }
+
+  if (need_vertical) {
+    int interp_dim = 2;
+    std::tie(vert_indices_weights, ksize_vert, vert_weights_precision) = F::compute_index_ranges_int16_weights(
+        /*input_size=*/yin,
+        /*output_size=*/yout,
+        /*stride=*/num_channels * xout,
+        /*ndims=*/4,
+        /*reshape_dim=*/interp_dim,
+        /*align_corners=*/align_corners,
+        /*opt_scale=*/scales[interp_dim - 2],
+        /*antialias=*/antialias,
+        /*align_i32=*/false);
+  }
+
+  at::Tensor buffer_horiz;
+  if (need_horizontal && need_vertical) {
+    buffer_horiz = at::empty({num_channels, yin, xout}, input.options());
+  }
+
+  for (const auto i : c10::irange(batch_size)) {
+    at::Tensor input_slice = input[i];
+
+    if (need_horizontal) {
+      at::Tensor horiz_output = need_vertical ? buffer_horiz : output[i];
+      NeonResampleHorizontal(horiz_output, input_slice, ksize_horiz, horiz_indices_weights, horiz_weights_precision);
+      if (need_vertical) {
+        input_slice = horiz_output;
+      }
+    }
+    if (need_vertical) {
+      NeonResampleVertical(output[i], input_slice, ksize_vert, vert_indices_weights, vert_weights_precision);
+    }
+  }
+}
+
+} // anonymous namespace
+
+#endif // __aarch64__
diff --git a/aten/src/ATen/native/cpu/int4mm_kernel.cpp b/aten/src/ATen/native/cpu/int4mm_kernel.cpp
index 9a9897b88fba2..72911aea5f5bb 100644
--- a/aten/src/ATen/native/cpu/int4mm_kernel.cpp
+++ b/aten/src/ATen/native/cpu/int4mm_kernel.cpp
@@ -40,7 +40,7 @@ inline bool is_block_start(int index, int BLOCK_SIZE) {
 #if (defined(CPU_CAPABILITY_AVX512) || defined(CPU_CAPABILITY_AVX2)) && !defined(_MSC_VER)
 // convert 16x int4 to int8, handle 64 bits at a time
 // used in avx2 and avx512
-inline __m128i conver_int4_to_int8(const uint8_t* data) {
+inline __m128i convert_int4_to_int8(const uint8_t* data) {
   __m128i tmp = _mm_loadu_si64((const __m128i*)data);
   __m128i bytes = _mm_cvtepu8_epi16(tmp);
   const __m128i lowMask = _mm_set1_epi8(0xF);
@@ -169,7 +169,7 @@ inline void tinygemm_kernel(
           vb[3] = _mm512_fmadd_ps(vb[3], scale[3], zero[3]);
         }
       } else {
-        __m128i b8 = conver_int4_to_int8(B + k * ldb + col * 8);
+        __m128i b8 = convert_int4_to_int8(B + k * ldb + col * 8);
         __m512i b32 = _mm512_cvtepu8_epi32(b8);
         vb[col] = _mm512_permutexvar_ps(b32, lut);
         vb[col] = _mm512_fmadd_ps(vb[col], scale[col], zero[col]);
@@ -312,7 +312,7 @@ inline void tinygemm_kernel(
       } else {
         if constexpr (col % 2 == 0) {
           // de-quantize per 64 bits (16x int4)
-          __m128i b8 = conver_int4_to_int8(B + k * ldb + col * 4);
+          __m128i b8 = convert_int4_to_int8(B + k * ldb + col * 4);
           __m128i b8_val0 = _mm_set1_epi64x(_mm_extract_epi64(b8, 0));
           __m128i b8_val1 = _mm_set1_epi64x(_mm_extract_epi64(b8, 1));
           if (k + PREFETCH_SIZE_K < K) {
diff --git a/aten/src/ATen/native/cuda/AdaptiveMaxPooling2d.cu b/aten/src/ATen/native/cuda/AdaptiveMaxPooling2d.cu
index 2030dbb904505..30745cb8f0a0a 100644
--- a/aten/src/ATen/native/cuda/AdaptiveMaxPooling2d.cu
+++ b/aten/src/ATen/native/cuda/AdaptiveMaxPooling2d.cu
@@ -393,7 +393,7 @@ TORCH_IMPL_FUNC(adaptive_max_pool2d_backward_out_cuda)
             C10_CUDA_KERNEL_LAUNCH_CHECK();
           } else {
             // run updateGradInput kernel
-            atomicadaptivemaxgradinput<<<
+            adaptivemaxgradinput<<<
                 blocks,
                 threads,
                 0,
diff --git a/aten/src/ATen/native/cuda/AveragePool3d.cu b/aten/src/ATen/native/cuda/AveragePool3d.cu
index dabcf5b63be99..f290864a9320b 100644
--- a/aten/src/ATen/native/cuda/AveragePool3d.cu
+++ b/aten/src/ATen/native/cuda/AveragePool3d.cu
@@ -500,11 +500,6 @@ TORCH_IMPL_FUNC(avg_pool3d_backward_out_cuda) (
   const int64_t oheight = gradOutput.size(-2);
   const int64_t owidth = gradOutput.size(-1);
 
-  /* XXX shape check behavior from TH */
-  const int64_t otime_for_shape_check = pooling_output_shape<int64_t>(itime, kT, padT, dT, 1, ceil_mode);
-  const int64_t oheight_for_shape_check = pooling_output_shape<int64_t>(iheight, kH, padH, dH, 1, ceil_mode);
-  const int64_t owidth_for_chape_check = pooling_output_shape<int64_t>(iwidth, kW, padW, dW, 1, ceil_mode);
-
   const bool kernelsOverlap = (dT < kT) || (dH < kH) || (dW < kW);
 
   Tensor work_grad_input = gradInput;
diff --git a/aten/src/ATen/native/cuda/BinaryRemainderKernel.cu b/aten/src/ATen/native/cuda/BinaryRemainderKernel.cu
index d05db3dc5823c..d981a176c274b 100644
--- a/aten/src/ATen/native/cuda/BinaryRemainderKernel.cu
+++ b/aten/src/ATen/native/cuda/BinaryRemainderKernel.cu
@@ -6,6 +6,7 @@
 #include <ATen/native/TensorIterator.h>
 #include <c10/util/TypeSafeSignMath.h>
 
+#include <limits>
 #include <type_traits>
 
 // NOTE: CUDA on Windows requires that the enclosing function
@@ -15,6 +16,25 @@ namespace at::native {
 
 void remainder_kernel_cuda(TensorIteratorBase& iter) {
   if (isIntegralType(iter.common_dtype(), /*includeBool*/ false)) {
+#ifndef USE_ROCM
+    // Guard uint8 div-by-zero: NVIDIA hardware returns all-1s (255) for
+    // integer mod by zero, while ROCm returns the dividend unchanged, so
+    // we only need this explicit guard on the CUDA path.
+    AT_DISPATCH_INTEGRAL_TYPES(iter.common_dtype(), "remainder_cuda", [&]() {
+      gpu_kernel_with_scalars(iter, []GPU_LAMBDA(scalar_t a, scalar_t b) -> scalar_t {
+        if constexpr (std::is_same_v<scalar_t, uint8_t>) {
+          if (b == 0) {
+            return std::numeric_limits<uint8_t>::max();
+          }
+        }
+        scalar_t r = a % b;
+        if (r != 0 && c10::signs_differ(r, b)) {
+          r += b;
+        }
+        return r;
+      });
+    });
+#else
     AT_DISPATCH_INTEGRAL_TYPES(iter.common_dtype(), "remainder_cuda", [&]() {
       gpu_kernel_with_scalars(iter, []GPU_LAMBDA(scalar_t a, scalar_t b) -> scalar_t {
         scalar_t r = a % b;
@@ -24,6 +44,7 @@ void remainder_kernel_cuda(TensorIteratorBase& iter) {
         return r;
       });
     });
+#endif
   } else {
     AT_DISPATCH_FLOATING_TYPES_AND2(kHalf, kBFloat16, iter.common_dtype(), "remainder_cuda", [&]() {
       gpu_kernel_with_scalars(iter,
diff --git a/aten/src/ATen/native/cuda/Blas.cpp b/aten/src/ATen/native/cuda/Blas.cpp
index 4a03faa02ef56..c37e28b0c96bb 100644
--- a/aten/src/ATen/native/cuda/Blas.cpp
+++ b/aten/src/ATen/native/cuda/Blas.cpp
@@ -12,7 +12,7 @@
 #include <ATen/OpMathType.h>
 #include <ATen/TensorUtils.h>
 #include <ATen/cuda/CUDABlas.h>
-#include <ATen/cuda/CUDAScaledBlas.h>
+#include <ATen/native/ScaledBlasUtils.h>
 #include <ATen/cuda/tunable/Tunable.h>
 #include <ATen/cuda/tunable/TunableGemm.h>
 #include <ATen/native/Resize.h>
@@ -34,6 +34,7 @@
 #else
 #include <ATen/ops/_addmm_activation_native.h>
 #include <ATen/ops/_efficientzerotensor.h>
+#include <ATen/ops/_int_mm_native.h>
 #include <ATen/ops/_scaled_mm_native.h>
 #include <ATen/ops/_unsafe_view_native.h>
 #include <ATen/ops/abs.h>
@@ -108,8 +109,7 @@ cuda::blas::GEMMAndBiasActivationEpilogue activation_to_gemm_and_blas_arg(Activa
     case Activation::GELU:
       return cuda::blas::GEMMAndBiasActivationEpilogue::GELU;
     default:
-      TORCH_CHECK(false);
-      return cuda::blas::GEMMAndBiasActivationEpilogue::None;
+      TORCH_CHECK(false, "Unknown activation epologue type");
   }
 }
 
@@ -228,9 +228,6 @@ static bool isInputCompliesAddmmCudaLt(
       mat2_sizes[0] > 1 && mat2_sizes[1] > 1
     )
   );
-
-  // no compliance by default
-  return false;
 }
 
 template <typename scalar_t>
diff --git a/aten/src/ATen/native/cuda/CUDALoops.cuh b/aten/src/ATen/native/cuda/CUDALoops.cuh
index b1031a95a3d0e..00f0aedf4578c 100644
--- a/aten/src/ATen/native/cuda/CUDALoops.cuh
+++ b/aten/src/ATen/native/cuda/CUDALoops.cuh
@@ -171,8 +171,15 @@ __global__ void vectorized_elementwise_kernel(int N, func_t f, array_t data) {
     constexpr auto io_size = calc_io_size<func_t>();
     int remaining = N - io_block_work_size<io_size>() * blockIdx.x;
 
-    if (remaining < io_block_work_size<io_size>()) { // if this block handles the reminder,
-                  // just do a naive unrolled loop
+    // note: unless the compiler has a good reason to move code, it won't.
+    // Thus, the if-condition typically comes first in SASS, so the "hot" path
+    // should go first, improving instruction cache use.
+    if (remaining >= io_block_work_size<io_size>()) { // if this block has a full `block_work_size` data to handle, use
+      // vectorized memory access
+      elementwise_kernel_helper(
+        f, memory::policies::vectorized<vec_size, array_t, elems_per_thread<io_size>()>(data));
+    } else { // if this block handles the reminder,
+      // just do a naive unrolled loop
       auto input_calc = TrivialOffsetCalculator<traits::arity>();
       auto output_calc = TrivialOffsetCalculator<1>();
       auto loader = memory::LoadWithoutCast();
@@ -186,10 +193,6 @@ __global__ void vectorized_elementwise_kernel(int N, func_t f, array_t data) {
       elems_per_thread<io_size>()>(
       data, remaining, input_calc, output_calc, loader, storer);
       elementwise_kernel_helper(f, policy);
-    } else { // if this block has a full `block_work_size` data to handle, use
-        // vectorized memory access
-      elementwise_kernel_helper(
-      f, memory::policies::vectorized<vec_size, array_t, elems_per_thread<io_size>()>(data));
     }
 #else
     CUDA_KERNEL_ASSERT(false && "Fatal! vectorized_elementwise_kernel<8,...> supports only sm_90 and sm_10x. Please report an issue on GitHub.");
@@ -199,8 +202,12 @@ __global__ void vectorized_elementwise_kernel(int N, func_t f, array_t data) {
     constexpr auto io_size = calc_io_size<func_t>();
     int remaining = N - io_block_work_size<io_size>() * blockIdx.x;
 
-    if (remaining < io_block_work_size<io_size>()) { // if this block handles the reminder,
-                   // just do a naive unrolled loop
+    if (remaining >= io_block_work_size<io_size>()) { // if this block has a full `block_work_size` data to handle, use
+      // vectorized memory access
+      elementwise_kernel_helper(
+        f, memory::policies::vectorized<vec_size, array_t, elems_per_thread<io_size>()>(data));
+    } else { // if this block handles the reminder,
+      // just do a naive unrolled loop
       auto input_calc = TrivialOffsetCalculator<traits::arity>();
       auto output_calc = TrivialOffsetCalculator<1>();
       auto loader = memory::LoadWithoutCast();
@@ -214,10 +221,6 @@ __global__ void vectorized_elementwise_kernel(int N, func_t f, array_t data) {
       elems_per_thread<io_size>()>(
       data, remaining, input_calc, output_calc, loader, storer);
       elementwise_kernel_helper(f, policy);
-    } else { // if this block has a full `block_work_size` data to handle, use
-         // vectorized memory access
-      elementwise_kernel_helper(
-      f, memory::policies::vectorized<vec_size, array_t, elems_per_thread<io_size>()>(data));
     }
   }
 }
@@ -313,12 +316,11 @@ static inline void launch_vectorized_kernel(
   if (p->major != 9 && p->major != 10) {
     vec_size = std::min<uint16_t>(vec_size, 4);
   }
-  // Here we purposely omit vec8 for 1-byte data because of a bug in NVCC
-  // that causes some numerical mismatches with uint8 on sm80 and sm90.
-  // TODO: Revisit this after CUDA 12.8 update.
+#if !defined(CUDA_VERSION) || CUDA_VERSION < 12080
   if constexpr (sizeof(cpp_type) < 2) {
     vec_size = std::min<uint16_t>(vec_size, 4);
   }
+#endif
   int tws = elems_per_thread<io_size>();
 #endif
   int bws = tws * num_threads();
@@ -646,7 +648,6 @@ void gpu_kernel_impl_nocast(TensorIteratorBase& iter, const func_t& f) {
   TORCH_INTERNAL_ASSERT(iter.can_use_32bit_indexing());
   TORCH_INTERNAL_ASSERT(iter.ninputs() == traits::arity);
   TORCH_INTERNAL_ASSERT(iter.noutputs() == 1);
-  TORCH_INTERNAL_ASSERT(!needs_dynamic_casting<func_t>::check(iter));
 
   std::array<char*, ntensors> data;
   for (int i = 0; i < ntensors; i++) {
diff --git a/aten/src/ATen/native/cuda/CUDAScalar.cu b/aten/src/ATen/native/cuda/CUDAScalar.cu
index 169a2ab92615f..3f2bec46117a1 100644
--- a/aten/src/ATen/native/cuda/CUDAScalar.cu
+++ b/aten/src/ATen/native/cuda/CUDAScalar.cu
@@ -10,28 +10,74 @@
 #endif
 
 #include <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDACachingAllocator.h>
 
 namespace at::native {
 
+namespace {
+
+bool is_cuda_caching_allocator_tensor(const Tensor& self) {
+  auto* cuda_allocator = c10::cuda::CUDACachingAllocator::get();
+  if (cuda_allocator == nullptr) {
+    return false;
+  }
+  // SymmMem/NVSHMEM/rocSHMEM tensors are typically backed by custom
+  // from_blob-style deleters, so this check filters them out and keeps
+  // the direct dereference path limited to allocator-managed CUDA memory.
+  return self.storage().data_ptr().get_deleter() == cuda_allocator->raw_deleter();
+}
+
+template <typename scalar_t>
+void _local_scalar_dense_cuda_impl(const Tensor& self, Scalar& r) {
+#if defined(USE_ROCM) && (ROCM_VERSION >= 70200)
+  // If this is a large BAR device, we can just read directly from VRAM
+  if (
+      at::cuda::getCurrentDeviceProperties()->isLargeBar &&
+      is_cuda_caching_allocator_tensor(self)) {
+    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+#if ROCM_VERSION < 71400
+    hipStreamCaptureStatus captureStatus;
+    C10_CUDA_CHECK(hipStreamGetCaptureInfo(stream, &captureStatus, nullptr));
+    if (C10_LIKELY(captureStatus == hipStreamCaptureStatusNone)) {
+      at::cuda::stream_synchronize(stream);
+      r = Scalar(*self.template const_data_ptr<scalar_t>());
+    } else {
+      C10_CUDA_CHECK(hipErrorStreamCaptureUnsupported);
+    }
+#else // ROCM_VERSION
+    // HIP reports the errors during graph capture in ROCm 7.14+
+    // no StreamCaptureStatus check required
+    at::cuda::stream_synchronize(stream);
+    r = Scalar(*self.template const_data_ptr<scalar_t>());
+#endif // ROCM_VERSION
+    return;
+  }
+#endif
+
+  // Create pinned memory for the scalar value to avoid implicit
+  // locking/sync in cuda library due to pageable memory
+  auto value = at::detail::empty_cpu(
+    {1}, /* size */
+    c10::CppTypeToScalarType<scalar_t>(), /* dtype */
+    std::nullopt, /* layout */
+    std::nullopt, /* device */
+    true, /* pin_memory */
+    std::nullopt /* memory format */
+  );
+  cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+  at::cuda::memcpy_and_sync(value.template mutable_data_ptr<scalar_t>(), self.template const_data_ptr<scalar_t>(), sizeof(scalar_t), cudaMemcpyDeviceToHost, stream);
+  r = Scalar(*value.template const_data_ptr<scalar_t>());
+}
+
+} // anonymous namespace
+
 Scalar _local_scalar_dense_cuda(const Tensor& self) {
   Scalar r;
   TORCH_CHECK(self.numel() > 0, "_local_scalar_dense: Empty tensor not supported");
     AT_DISPATCH_V2(
       self.scalar_type(), "_local_scalar_dense_cuda", AT_WRAP([&] {
-          // Create pinned memory for the scalar value to avoid implicit
-          // locking/sync in cuda library due to pageable memory
-          auto value = at::detail::empty_cpu(
-            {1}, /* size */
-            c10::CppTypeToScalarType<scalar_t>(), /* dtype */
-            std::nullopt, /* layout */
-            std::nullopt, /* device */
-            true, /* pin_memory */
-            std::nullopt /* memory format */
-          );
-          cudaStream_t stream = at::cuda::getCurrentCUDAStream();
-          at::cuda::memcpy_and_sync(value.mutable_data_ptr<scalar_t>(), self.const_data_ptr<scalar_t>(), sizeof(scalar_t), cudaMemcpyDeviceToHost, stream);
-          r = Scalar(*value.const_data_ptr<scalar_t>());
-        }), AT_EXPAND(AT_ALL_TYPES_AND_COMPLEX), kComplexHalf, kHalf, kBool, kBFloat16, AT_EXPAND(AT_FLOAT8_TYPES), AT_EXPAND(AT_BAREBONES_UNSIGNED_TYPES));
+        _local_scalar_dense_cuda_impl<scalar_t>(self, r);
+      }), AT_EXPAND(AT_ALL_TYPES_AND_COMPLEX), kComplexHalf, kHalf, kBool, kBFloat16, AT_EXPAND(AT_FLOAT8_TYPES), AT_EXPAND(AT_BAREBONES_UNSIGNED_TYPES));
   return r;
 }
 
diff --git a/aten/src/ATen/native/cuda/CompositeRandomAccessor.h b/aten/src/ATen/native/cuda/CompositeRandomAccessor.h
index eb8587d1f9337..d519e333b1e9d 100644
--- a/aten/src/ATen/native/cuda/CompositeRandomAccessor.h
+++ b/aten/src/ATen/native/cuda/CompositeRandomAccessor.h
@@ -25,7 +25,7 @@ void swap(
   references_holder<Values, References> rh1,
   references_holder<Values, References> rh2
 ) {
-  return thrust::swap(rh1.data(), rh2.data());
+  thrust::swap(rh1.data(), rh2.data());
 }
 
 template <int N, typename Values, typename References>
diff --git a/aten/src/ATen/native/cuda/Copy.cu b/aten/src/ATen/native/cuda/Copy.cu
index 9e78d2e71f043..68741f90ef85e 100644
--- a/aten/src/ATen/native/cuda/Copy.cu
+++ b/aten/src/ATen/native/cuda/Copy.cu
@@ -21,8 +21,7 @@
 #include <c10/cuda/CUDAStream.h>
 #include <ATen/cuda/CUDAGraphsUtils.cuh>
 
-// TODO(NS): Investigate why FP8 conversion intrinsics end up being slower
-#ifdef AT_USE_NV_CVT_INTRINSICS
+#if defined(CUDA_VERSION) && CUDA_VERSION >= 13000
 #include <cuda_fp8.h>
 #endif
 
@@ -69,25 +68,53 @@ void float16tofloat32_copy_kernel_cuda(TensorIteratorBase &iter) {
 }
 #endif
 
+template <typename SrcT>
+struct ConvertToFloat8E4M3fnOp {
+  __device__ __forceinline__ Float8_e4m3fn operator()(SrcT value) const {
+#if defined(CUDA_VERSION) && CUDA_VERSION >= 13000 && defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 890
+    __nv_fp8_storage_t x;
+    if constexpr (std::is_same_v<SrcT, float>) {
+      x = __nv_cvt_float_to_fp8(value, __NV_SATFINITE, __NV_E4M3);
+    } else if constexpr (std::is_same_v<SrcT, Half>) {
+      x = __nv_cvt_halfraw_to_fp8(static_cast<__half>(value), __NV_SATFINITE, __NV_E4M3);
+    } else if constexpr (std::is_same_v<SrcT, BFloat16>) {
+      x = __nv_cvt_bfloat16raw_to_fp8(static_cast<__nv_bfloat16>(value), __NV_SATFINITE, __NV_E4M3);
+    } else {
+      x = __nv_cvt_float_to_fp8(static_cast<float>(value), __NV_SATFINITE, __NV_E4M3);
+    }
+    return Float8_e4m3fn(x, Float8_e4m3fn::from_bits());
+#else
+    return Float8_e4m3fn(value);
+#endif
+  }
+};
+
+// e5m2 intrinsics are correct but slower; only used for float on Blackwell
+// to work around the ptxas subnormal codegen bug.
+struct ConvertFloatToFloat8E5M2Op {
+  __device__ __forceinline__ Float8_e5m2 operator()(float value) const {
+#if defined(CUDA_VERSION) && CUDA_VERSION >= 13020 && defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1000
+    auto x = __nv_cvt_float_to_fp8(value, __NV_NOSAT, __NV_E5M2);
+    return Float8_e5m2(x, Float8_e5m2::from_bits());
+#else
+    return Float8_e5m2(value);
+#endif
+  }
+};
+
 void float8_copy_kernel_cuda(TensorIteratorBase &iter) {
   ScalarType dtype = iter.dtype(0);
   ScalarType other_dtype = iter.dtype(1);
   if (dtype == kFloat8_e4m3fn) {
     switch (other_dtype) {
       case kFloat:
-         gpu_kernel_nocast(iter, [] GPU_LAMBDA(float value) {
-             return Float8_e4m3fn(value);
-         });
+         gpu_kernel_nocast(iter, ConvertToFloat8E4M3fnOp<float>{});
          break;
       case kHalf:
-         gpu_kernel_nocast(iter, [] GPU_LAMBDA(Half value) {
-             return Float8_e4m3fn(value);
-         });
+         gpu_kernel_nocast(iter, ConvertToFloat8E4M3fnOp<Half>{});
          break;
       case kBFloat16:
-         gpu_kernel_nocast(iter, [] GPU_LAMBDA(BFloat16 value) {
-             return Float8_e4m3fn(value);
-         });
+         gpu_kernel_nocast(iter, ConvertToFloat8E4M3fnOp<BFloat16>{});
          break;
       default:
         gpu_kernel(iter, [] GPU_LAMBDA(Float8_e4m3fn x) { return x; });
@@ -96,33 +123,16 @@ void float8_copy_kernel_cuda(TensorIteratorBase &iter) {
   } else if (dtype == kFloat8_e5m2) {
     switch (other_dtype) {
       case kFloat:
-         gpu_kernel_nocast(iter, [] GPU_LAMBDA(float value) {
-#ifdef AT_USE_NV_CVT_INTRINSICS
-             const auto x =  __nv_cvt_float_to_fp8(value, __NV_NOSAT, __NV_E5M2);
-             return Float8_e5m2(x, Float8_e5m2::from_bits());
-#else
-             return Float8_e5m2(value);
-#endif
-         });
+         gpu_kernel_nocast(iter, ConvertFloatToFloat8E5M2Op{});
          break;
       case kHalf:
          gpu_kernel_nocast(iter, [] GPU_LAMBDA(Half value) {
-#ifdef AT_USE_NV_CVT_INTRINSICS
-             const auto x =  __nv_cvt_halfraw_to_fp8(static_cast<__half>(value), __NV_NOSAT, __NV_E5M2);
-             return Float8_e5m2(x, Float8_e5m2::from_bits());
-#else
              return Float8_e5m2(value);
-#endif
          });
          break;
       case kBFloat16:
          gpu_kernel_nocast(iter, [] GPU_LAMBDA(BFloat16 value) {
-#ifdef AT_USE_NV_CVT_INTRINSICS
-             const auto x =  __nv_cvt_bfloat16raw_to_fp8(static_cast<__nv_bfloat16>(value), __NV_NOSAT, __NV_E5M2);
-             return Float8_e5m2(x, Float8_e5m2::from_bits());
-#else
              return Float8_e5m2(value);
-#endif
          });
          break;
       default:
@@ -194,7 +204,7 @@ void float8_copy_kernel_cuda(TensorIteratorBase &iter) {
          break;
     }
   } else {
-    TORCH_CHECK(false, "This supposed ot be called only for Float8 types");
+    TORCH_CHECK(false, "This supposed to be called only for Float8 types");
   }
 }
 
diff --git a/aten/src/ATen/native/cuda/DistanceKernel.cu b/aten/src/ATen/native/cuda/DistanceKernel.cu
index e6f4dadb23df6..7da0055a0e3ef 100644
--- a/aten/src/ATen/native/cuda/DistanceKernel.cu
+++ b/aten/src/ATen/native/cuda/DistanceKernel.cu
@@ -9,6 +9,7 @@
 #include <ATen/native/cuda/block_reduce.cuh>
 #include <ATen/native/cuda/DeviceSqrt.cuh>
 #include <ATen/native/Distance.h>
+#include <ATen/NumericUtils.h>
 
 #ifndef AT_PER_OPERATOR_HEADERS
 #include <ATen/Functions.h>
@@ -34,7 +35,13 @@ struct dists {
 
   // Zero norm
   struct zero {
-    static __forceinline__ __device__ void inc(scalar_t& agg, const scalar_t diff, const scalar_t /*p*/) { agg += diff != 0.0; }
+    static __forceinline__ __device__ void inc(scalar_t& agg, const scalar_t diff, const scalar_t /*p*/) {
+      if (diff != diff) { // NaN
+        agg = diff;
+      } else if (diff != 0.0) {
+        agg += 1.0;
+      }
+    }
     static __forceinline__ __device__ scalar_t finish(const scalar_t agg, const scalar_t /*p*/) { return agg; }
     static __forceinline__ __device__ void agg(scalar_t& update, const scalar_t other) { update += other; }
   };
diff --git a/aten/src/ATen/native/cuda/Distributions.cpp b/aten/src/ATen/native/cuda/Distributions.cpp
index be397f4bc217f..f270ea66d3755 100644
--- a/aten/src/ATen/native/cuda/Distributions.cpp
+++ b/aten/src/ATen/native/cuda/Distributions.cpp
@@ -28,6 +28,14 @@ Tensor _s_poisson_cuda(const Tensor& lambda, std::optional<Generator> gen_) {
 
 // NOLINTNEXTLINE(performance-unnecessary-value-param)
 Tensor _s_binomial_cuda(const Tensor& count, const Tensor& prob, std::optional<Generator> gen_) {
+  TORCH_CHECK_VALUE(
+      at::isFloatingType(count.scalar_type()),
+      "binomial only supports floating-point dtypes for count, got: ",
+      count.scalar_type());
+  TORCH_CHECK_VALUE(
+      at::isFloatingType(prob.scalar_type()),
+      "binomial only supports floating-point dtypes for prob, got: ",
+      prob.scalar_type());
   auto gen = get_generator_or_default<CUDAGeneratorImpl>(gen_, cuda::detail::getDefaultCUDAGenerator());
   Tensor ret = at::empty(count.sizes(), count.options());
   at::TensorIterator iter = at::TensorIteratorConfig()
diff --git a/aten/src/ATen/native/cuda/ForeachFunctors.cuh b/aten/src/ATen/native/cuda/ForeachFunctors.cuh
index 37175776097df..1e18e42c34565 100644
--- a/aten/src/ATen/native/cuda/ForeachFunctors.cuh
+++ b/aten/src/ATen/native/cuda/ForeachFunctors.cuh
@@ -77,6 +77,34 @@ __device__ bool init_args(
   return all_aligned;
 }
 
+template <
+    int depth,
+    typename param_type,
+    typename grad_type,
+    typename exp_avg_type,
+    typename exp_avg_sq_type>
+__device__ bool init_args_mixed_prec(
+    param_type** param_args,
+    grad_type** grad_args,
+    exp_avg_type** exp_avg_args,
+    exp_avg_sq_type** exp_avg_sq_args,
+    FusedOptimizerTensorListMetadata<depth>& tl,
+    const int64_t chunk_idx,
+    const int64_t chunk_size,
+    const int64_t tensor_loc) {
+  *param_args =
+      (param_type*)tl.addresses[0][tensor_loc] + chunk_idx * chunk_size;
+  *grad_args = (grad_type*)tl.addresses[1][tensor_loc] + chunk_idx * chunk_size;
+  *exp_avg_args =
+      (exp_avg_type*)tl.addresses[2][tensor_loc] + chunk_idx * chunk_size;
+  *exp_avg_sq_args =
+      (exp_avg_sq_type*)tl.addresses[3][tensor_loc] + chunk_idx * chunk_size;
+
+  bool all_aligned = is_aligned(*param_args) && is_aligned(*grad_args) &&
+      is_aligned(*exp_avg_args) && is_aligned(*exp_avg_sq_args);
+  return all_aligned;
+}
+
 template <int depth, typename T>
 __device__ void load_args(
     T r_args[][kILP],
@@ -96,6 +124,43 @@ __device__ void load_args(
   }
 }
 
+template <
+    typename T,
+    typename param_type,
+    typename grad_type,
+    typename exp_avg_type,
+    typename exp_avg_sq_type>
+__device__ void load_args(
+    T r_args[][kILP],
+    const param_type* param_args,
+    const grad_type* grad_args,
+    const exp_avg_type* exp_avg_args,
+    const exp_avg_sq_type* exp_avg_sq_args,
+    const int64_t i_start,
+    const int64_t chunk_size,
+    const int64_t n) {
+#pragma unroll
+  for (int ii = 0; ii < kILP; ii++) {
+    const auto i = i_start + threadIdx.x + ii * blockDim.x;
+    r_args[0][ii] = 0;
+    if (i < n && i < chunk_size) {
+      r_args[0][ii] = static_cast<T>(param_args[i]);
+    }
+    r_args[1][ii] = 0;
+    if (i < n && i < chunk_size) {
+      r_args[1][ii] = static_cast<T>(grad_args[i]);
+    }
+    r_args[2][ii] = 0;
+    if (i < n && i < chunk_size) {
+      r_args[2][ii] = static_cast<T>(exp_avg_args[i]);
+    }
+    r_args[3][ii] = 0;
+    if (i < n && i < chunk_size) {
+      r_args[3][ii] = static_cast<T>(exp_avg_sq_args[i]);
+    }
+  }
+}
+
 template <typename T>
 __device__ void store_args(
     T* dst,
@@ -111,6 +176,21 @@ __device__ void store_args(
   }
 }
 
+template <typename dT, typename sT>
+__device__ void store_args(
+    dT* dst,
+    sT* src,
+    const int64_t i_start,
+    const int64_t chunk_size,
+    const int64_t n) {
+#pragma unroll
+  for (int ii = 0; ii < kILP; ii++) {
+    const int64_t i = i_start + threadIdx.x + ii * blockDim.x;
+    if (i < n && i < chunk_size)
+      dst[i] = static_cast<dT>(src[ii]);
+  }
+}
+
 template <int res_arg_index, typename Op, typename T, typename opmath_t>
 __device__ __forceinline__ void binary_op_scalar(
     T r_args[][kILP],
diff --git a/aten/src/ATen/native/cuda/ForeachReduceOp.cu b/aten/src/ATen/native/cuda/ForeachReduceOp.cu
index 5bdc11af2d55f..e7a073fe75a55 100644
--- a/aten/src/ATen/native/cuda/ForeachReduceOp.cu
+++ b/aten/src/ATen/native/cuda/ForeachReduceOp.cu
@@ -28,22 +28,25 @@
 
 namespace at::native {
 
-// _foreach_norm supports only L1, L2, and inf norm
-enum class NormType { L1, L2, LInf };
+// _foreach_norm supports L0, L1, L2, and inf norm
+enum class NormType { L0, L1, L2, LInf };
 
 // NOTE: This is a simple variant of TensorListMetadata in MultiTensorApply.cuh
 // as we only need to track addresses for the lpnorm_cleanup function below.
 // Why is this struct necessary? For the same reason the TensorListMetadata
 // struct is necessary--which is to ferry static metadata to the CUDA kernel
-// while complying with the 4kb size constraint. Since we only need to track
-// addresses, we introduce this struct to be able to fit more Tensor pointers at
-// a time, currently 400 empirically, compared to the much smaller values in
-// depth_to_max_tensors. This way, we can launch fewer kernels for better
-// performance.
+// while complying with the kernel arg size constraint. Since we only need to
+// track addresses, we introduce this struct to be able to fit more Tensor
+// pointers at a time compared to depth_to_max_tensors. This way, we can
+// launch fewer kernels for better performance.
 //
 // IF YOU USE THIS STRUCT, PLEASE ADD A ONE-OFF TEST IN test_foreach.py AS THIS
 // IS CURRENTLY ONLY TESTED FOR _foreach_norm.
-const size_t MAX_TENSORS_PER_KERNEL = 400;
+#if defined(CUDART_VERSION) && CUDART_VERSION >= 13000 && !defined(USE_ROCM)
+static constexpr size_t MAX_TENSORS_PER_KERNEL = 3200;
+#else
+static constexpr size_t MAX_TENSORS_PER_KERNEL = 400;
+#endif
 struct TensorListAddresses {
   const void* addresses[MAX_TENSORS_PER_KERNEL];
 };
@@ -281,7 +284,10 @@ struct LpNormFunctor {
 #pragma unroll
         for (int ii = 0; ii < kILP; ii++) {
           const auto next = static_cast<out_opmath_t>(r_x[ii]);
-          if constexpr (norm_type == NormType::LInf) {
+          if constexpr (norm_type == NormType::L0) {
+            vals[ii] +=
+                next != out_opmath_t(0) ? out_opmath_t(1) : out_opmath_t(0);
+          } else if constexpr (norm_type == NormType::LInf) {
             vals[ii] = max_propagate_nan(vals[ii], ::abs(next));
           } else {
             vals[ii] += norm_type == NormType::L1 ? ::abs(next) : next * next;
@@ -296,7 +302,10 @@ struct LpNormFunctor {
           int i = i_start + threadIdx.x + ii * blockDim.x;
           if (i < n && i < chunk_size) {
             const auto next = static_cast<out_opmath_t>(x[i]);
-            if constexpr (norm_type == NormType::LInf) {
+            if constexpr (norm_type == NormType::L0) {
+              vals[ii] +=
+                  next != out_opmath_t(0) ? out_opmath_t(1) : out_opmath_t(0);
+            } else if constexpr (norm_type == NormType::LInf) {
               vals[ii] = max_propagate_nan(vals[ii], ::abs(next));
             } else {
               vals[ii] += norm_type == NormType::L1 ? ::abs(next) : next * next;
@@ -314,7 +323,8 @@ struct LpNormFunctor {
         val += vals[i];
       }
     }
-    auto final_val = norm_type == NormType::L1 || norm_type == NormType::L2
+    auto final_val = norm_type == NormType::L0 || norm_type == NormType::L1 ||
+            norm_type == NormType::L2
         ? at::native::cuda_utils::BlockReduceSum(val, s_vals)
         : at::native::cuda_utils::BlockReduceMax(val, s_vals);
 
@@ -348,8 +358,8 @@ __global__ void lpnorm_cleanup(
       val += output_this_tensor[i];
     }
   }
-  out_opmath_t final_val =
-      norm_type == NormType::L1 || norm_type == NormType::L2
+  out_opmath_t final_val = norm_type == NormType::L0 ||
+          norm_type == NormType::L1 || norm_type == NormType::L2
       ? at::native::cuda_utils::BlockReduceSum<out_opmath_t>(val, vals)
       : at::native::cuda_utils::BlockReduceMax(val, vals);
   if (threadIdx.x == 0) {
@@ -476,24 +486,31 @@ std::vector<Tensor> foreach_tensor_norm_cuda_internal(
         AT_DISPATCH_OUT_DTYPES(
             output_dtype, ForeachNormDispatchName<apply_root>::value, [&]() {
               using out_opmath_t = typename at::opmath_type<out_t>;
-              if (p == static_cast<double>(1)) {
+              if (p == static_cast<double>(0)) {
+                multi_tensor_apply<1>(
+                    tensor_lists,
+                    LpNormFunctor<scalar_t, NormType::L0, out_t>(),
+                    output_per_tensor.template mutable_data_ptr<out_opmath_t>(),
+                    max_chunks_per_tensor);
+              } else if (p == static_cast<double>(1)) {
                 multi_tensor_apply<1>(
                     tensor_lists,
                     LpNormFunctor<scalar_t, NormType::L1, out_t>(),
-                    output_per_tensor.mutable_data_ptr<out_opmath_t>(),
+                    output_per_tensor.template mutable_data_ptr<out_opmath_t>(),
                     max_chunks_per_tensor);
               } else if (p == static_cast<double>(2)) {
                 multi_tensor_apply<1>(
                     tensor_lists,
                     LpNormFunctor<scalar_t, NormType::L2, out_t>(),
-                    output_per_tensor.mutable_data_ptr<out_opmath_t>(),
+                    output_per_tensor.template mutable_data_ptr<out_opmath_t>(),
                     max_chunks_per_tensor);
               } else if constexpr (support_infinity) {
                 if (p == std::numeric_limits<double>::infinity()) {
                   multi_tensor_apply<1>(
                       tensor_lists,
                       LpNormFunctor<scalar_t, NormType::LInf, out_t>(),
-                      output_per_tensor.mutable_data_ptr<out_opmath_t>(),
+                      output_per_tensor
+                          .template mutable_data_ptr<out_opmath_t>(),
                       max_chunks_per_tensor);
                 }
               }
@@ -515,13 +532,23 @@ std::vector<Tensor> foreach_tensor_norm_cuda_internal(
                 for (const auto j : c10::irange(num_tensors_this_kernel)) {
                   addr_struct.addresses[j] =
                       vec_res[i * MAX_TENSORS_PER_KERNEL + j]
-                          .mutable_data_ptr<out_t>();
+                          .template mutable_data_ptr<out_t>();
                 }
 
-                if (p == static_cast<double>(1)) {
+                if (p == static_cast<double>(0)) {
+                  lpnorm_cleanup<scalar_t, NormType::L0, out_t, apply_root>
+                      <<<num_tensors_this_kernel, 512, 0, stream>>>(
+                          output_per_tensor
+                                  .template const_data_ptr<out_opmath_t>() +
+                              i * MAX_TENSORS_PER_KERNEL *
+                                  max_chunks_per_tensor,
+                          addr_struct,
+                          max_chunks_per_tensor);
+                } else if (p == static_cast<double>(1)) {
                   lpnorm_cleanup<scalar_t, NormType::L1, out_t, apply_root>
                       <<<num_tensors_this_kernel, 512, 0, stream>>>(
-                          output_per_tensor.const_data_ptr<out_opmath_t>() +
+                          output_per_tensor
+                                  .template const_data_ptr<out_opmath_t>() +
                               i * MAX_TENSORS_PER_KERNEL *
                                   max_chunks_per_tensor,
                           addr_struct,
@@ -529,7 +556,8 @@ std::vector<Tensor> foreach_tensor_norm_cuda_internal(
                 } else if (p == static_cast<double>(2)) {
                   lpnorm_cleanup<scalar_t, NormType::L2, out_t, apply_root>
                       <<<num_tensors_this_kernel, 512, 0, stream>>>(
-                          output_per_tensor.const_data_ptr<out_opmath_t>() +
+                          output_per_tensor
+                                  .template const_data_ptr<out_opmath_t>() +
                               i * MAX_TENSORS_PER_KERNEL *
                                   max_chunks_per_tensor,
                           addr_struct,
@@ -538,7 +566,8 @@ std::vector<Tensor> foreach_tensor_norm_cuda_internal(
                   if (p == std::numeric_limits<double>::infinity()) {
                     lpnorm_cleanup<scalar_t, NormType::LInf, out_t, apply_root>
                         <<<num_tensors_this_kernel, 512, 0, stream>>>(
-                            output_per_tensor.const_data_ptr<out_opmath_t>() +
+                            output_per_tensor
+                                    .template const_data_ptr<out_opmath_t>() +
                                 i * MAX_TENSORS_PER_KERNEL *
                                     max_chunks_per_tensor,
                             addr_struct,
@@ -609,7 +638,8 @@ std::vector<Tensor> foreach_tensor_norm_cuda(
         });
   }
   if (!can_use_fast_route(tensors) || has_int_or_complex ||
-      !(p == static_cast<double>(1) || p == static_cast<double>(2) ||
+      !(p == static_cast<double>(0) || p == static_cast<double>(1) ||
+        p == static_cast<double>(2) ||
         p == std::numeric_limits<double>::infinity())) {
     return foreach_tensor_norm_slow(tensors, ord, dtype);
   }
diff --git a/aten/src/ATen/native/cuda/ForeachUnaryOp.cu b/aten/src/ATen/native/cuda/ForeachUnaryOp.cu
index bb070f9d97616..9a36a879615a0 100644
--- a/aten/src/ATen/native/cuda/ForeachUnaryOp.cu
+++ b/aten/src/ATen/native/cuda/ForeachUnaryOp.cu
@@ -13,6 +13,7 @@
 #include <ATen/ops/_foreach_asin_native.h>
 #include <ATen/ops/_foreach_atan_native.h>
 #include <ATen/ops/_foreach_ceil_native.h>
+#include <ATen/ops/_foreach_clone_native.h>
 #include <ATen/ops/_foreach_cos_native.h>
 #include <ATen/ops/_foreach_cosh_native.h>
 #include <ATen/ops/_foreach_erf_native.h>
@@ -40,7 +41,9 @@
 #include <ATen/ops/_foreach_trunc_native.h>
 #include <ATen/ops/_foreach_zero_native.h>
 
+#include <ATen/ops/_foreach_copy_native.h>
 #include <ATen/ops/empty_like_native.h>
+#include <ATen/ops/empty_strided_native.h>
 #endif
 
 namespace at::native {
@@ -431,4 +434,48 @@ void foreach_tensor_zero_cuda_(TensorList tensors) {
       });
 }
 
+std::vector<Tensor> foreach_tensor_clone_cuda(
+    TensorList self,
+    std::optional<MemoryFormat> memory_format) {
+  check_foreach_api_restrictions(self);
+  if (!_check_tensors_share_device_and_dtype({self})) {
+    return at::native::foreach_tensor_clone_slow(self, memory_format);
+  }
+
+  std::vector<Tensor> ret{};
+  ret.reserve(self.size());
+
+  auto realized_memory_format = memory_format.value_or(MemoryFormat::Preserve);
+  for (const auto& s : self) {
+    // This logic modified from at::native::clone.
+    if (realized_memory_format == MemoryFormat::Preserve) {
+      if (s.is_non_overlapping_and_dense()) {
+        // Copy all strides, this is marginally faster than calling empty_like
+        auto options = s.options();
+        ret.emplace_back(at::native::empty_strided_cuda(
+            s.sizes(),
+            s.strides(),
+            c10::optTypeMetaToScalarType(options.dtype_opt()),
+            options.layout_opt(),
+            options.device_opt(),
+            options.pinned_memory_opt()));
+      } else {
+        ret.emplace_back(at::native::empty_like(s));
+      }
+    } else {
+      auto options = s.options();
+      ret.emplace_back(at::native::empty_like(
+          s,
+          c10::optTypeMetaToScalarType(options.dtype_opt()),
+          options.layout_opt(),
+          options.device_opt(),
+          options.pinned_memory_opt(),
+          realized_memory_format));
+    }
+  }
+
+  at::native::foreach_tensor_copy_list_kernel_cuda_(ret, self);
+  return ret;
+}
+
 } // namespace at::native
diff --git a/aten/src/ATen/native/cuda/FusedAdamKernel.cu b/aten/src/ATen/native/cuda/FusedAdamKernel.cu
index 0858f24e17c6e..498dc634fab2b 100644
--- a/aten/src/ATen/native/cuda/FusedAdamKernel.cu
+++ b/aten/src/ATen/native/cuda/FusedAdamKernel.cu
@@ -29,10 +29,15 @@ void _fused_adam_kernel_cuda_(
     const bool maximize,
     const std::optional<at::Tensor>& grad_scale,
     const std::optional<at::Tensor>& found_inf) {
+  const bool is_mixed_precision =
+      params[0].scalar_type() != exp_avgs[0].scalar_type();
   if (amsgrad) {
     TORCH_CHECK(
         at::native::check_fast_path_restrictions(
-            {params, grads, exp_avgs, exp_avg_sqs, max_exp_avg_sqs}),
+            {params, grads, exp_avgs, exp_avg_sqs, max_exp_avg_sqs},
+            /*scalarList=*/{},
+            /*does_op_promote_integer_inputs_to_float=*/false,
+            /*skip_cross_list_dtype_check=*/is_mixed_precision),
         "params, grads, exp_avgs, exp_avg_sqs, and max_exp_avg_sqs must have same dtype, device, and layout");
     _fused_adam_amsgrad_cuda_impl_(
         params,
@@ -52,7 +57,10 @@ void _fused_adam_kernel_cuda_(
   } else {
     TORCH_CHECK(
         at::native::check_fast_path_restrictions(
-            {params, grads, exp_avgs, exp_avg_sqs}),
+            {params, grads, exp_avgs, exp_avg_sqs},
+            /*scalarList=*/{},
+            /*does_op_promote_integer_inputs_to_float=*/false,
+            /*skip_cross_list_dtype_check=*/is_mixed_precision),
         "params, grads, exp_avgs, and exp_avg_sqs must have same dtype, device, and layout");
     _fused_adam_cuda_impl_(
         params,
@@ -125,10 +133,15 @@ void _fused_adam_kernel_cuda_(
       lr.device() == param_device,
       "lr must be on the same GPU device as the params");
 
+  const bool is_mixed_precision =
+      params[0].scalar_type() != exp_avgs[0].scalar_type();
   if (amsgrad) {
     TORCH_CHECK(
         at::native::check_fast_path_restrictions(
-            {params, grads, exp_avgs, exp_avg_sqs, max_exp_avg_sqs}),
+            {params, grads, exp_avgs, exp_avg_sqs, max_exp_avg_sqs},
+            /*scalarList=*/{},
+            /*does_op_promote_integer_inputs_to_float=*/false,
+            /*skip_cross_list_dtype_check=*/is_mixed_precision),
         "params, grads, exp_avgs, exp_avg_sqs, and max_exp_avg_sqs must have same dtype, device, and layout");
     _fused_adam_amsgrad_cuda_impl_(
         params,
@@ -148,7 +161,10 @@ void _fused_adam_kernel_cuda_(
   } else {
     TORCH_CHECK(
         at::native::check_fast_path_restrictions(
-            {params, grads, exp_avgs, exp_avg_sqs}),
+            {params, grads, exp_avgs, exp_avg_sqs},
+            /*scalarList=*/{},
+            /*does_op_promote_integer_inputs_to_float=*/false,
+            /*skip_cross_list_dtype_check=*/is_mixed_precision),
         "params, grads, exp_avgs, and exp_avg_sqs must have same dtype, device, and layout");
     _fused_adam_cuda_impl_(
         params,
diff --git a/aten/src/ATen/native/cuda/FusedAdamWKernel.cu b/aten/src/ATen/native/cuda/FusedAdamWKernel.cu
index 4b758aa574c66..13e11df2a24e9 100644
--- a/aten/src/ATen/native/cuda/FusedAdamWKernel.cu
+++ b/aten/src/ATen/native/cuda/FusedAdamWKernel.cu
@@ -29,10 +29,15 @@ void _fused_adamw_kernel_cuda_(
     const bool maximize,
     const std::optional<at::Tensor>& grad_scale,
     const std::optional<at::Tensor>& found_inf) {
+  const bool is_mixed_precision =
+      params[0].scalar_type() != exp_avgs[0].scalar_type();
   if (amsgrad) {
     TORCH_CHECK(
         at::native::check_fast_path_restrictions(
-            {params, grads, exp_avgs, exp_avg_sqs, max_exp_avg_sqs}),
+            {params, grads, exp_avgs, exp_avg_sqs, max_exp_avg_sqs},
+            /*scalarList=*/{},
+            /*does_op_promote_integer_inputs_to_float=*/false,
+            /*skip_cross_list_dtype_check=*/is_mixed_precision),
         "params, grads, exp_avgs, exp_avg_sqs, and max_exp_avg_sqs must have same dtype, device, and layout");
     _fused_adamw_amsgrad_cuda_impl_(
         params,
@@ -52,7 +57,10 @@ void _fused_adamw_kernel_cuda_(
   } else {
     TORCH_CHECK(
         at::native::check_fast_path_restrictions(
-            {params, grads, exp_avgs, exp_avg_sqs}),
+            {params, grads, exp_avgs, exp_avg_sqs},
+            /*scalarList=*/{},
+            /*does_op_promote_integer_inputs_to_float=*/false,
+            /*skip_cross_list_dtype_check=*/is_mixed_precision),
         "params, grads, exp_avgs, and exp_avg_sqs must have same dtype, device, and layout");
     _fused_adamw_cuda_impl_(
         params,
@@ -125,10 +133,15 @@ void _fused_adamw_kernel_cuda_(
       lr.device() == param_device,
       "lr must be on the same GPU device as the params");
 
+  const bool is_mixed_precision =
+      params[0].scalar_type() != exp_avgs[0].scalar_type();
   if (amsgrad) {
     TORCH_CHECK(
         at::native::check_fast_path_restrictions(
-            {params, grads, exp_avgs, exp_avg_sqs, max_exp_avg_sqs}),
+            {params, grads, exp_avgs, exp_avg_sqs, max_exp_avg_sqs},
+            /*scalarList=*/{},
+            /*does_op_promote_integer_inputs_to_float=*/false,
+            /*skip_cross_list_dtype_check=*/is_mixed_precision),
         "params, grads, exp_avgs, exp_avg_sqs, and max_exp_avg_sqs must have same dtype, device, and layout");
     _fused_adamw_amsgrad_cuda_impl_(
         params,
@@ -148,7 +161,10 @@ void _fused_adamw_kernel_cuda_(
   } else {
     TORCH_CHECK(
         at::native::check_fast_path_restrictions(
-            {params, grads, exp_avgs, exp_avg_sqs}),
+            {params, grads, exp_avgs, exp_avg_sqs},
+            /*scalarList=*/{},
+            /*does_op_promote_integer_inputs_to_float=*/false,
+            /*skip_cross_list_dtype_check=*/is_mixed_precision),
         "params, grads, exp_avgs, and exp_avg_sqs must have same dtype, device, and layout");
     _fused_adamw_cuda_impl_(
         params,
diff --git a/aten/src/ATen/native/cuda/GroupedBlas.cpp b/aten/src/ATen/native/cuda/GroupedBlas.cpp
index 5875c9a805724..5b393401a29b4 100644
--- a/aten/src/ATen/native/cuda/GroupedBlas.cpp
+++ b/aten/src/ATen/native/cuda/GroupedBlas.cpp
@@ -13,7 +13,7 @@
 #include <ATen/OpMathType.h>
 #include <ATen/TensorUtils.h>
 #include <ATen/cuda/CUDABlas.h>
-#include <ATen/cuda/CUDAScaledBlas.h>
+#include <ATen/native/ScaledBlasUtils.h>
 #include <ATen/cuda/tunable/Tunable.h>
 #include <ATen/cuda/tunable/TunableGemm.h>
 #include <ATen/native/Resize.h>
@@ -37,6 +37,7 @@
 #else
 #include <ATen/ops/_addmm_activation_native.h>
 #include <ATen/ops/_efficientzerotensor.h>
+#include <ATen/ops/_grouped_mm_native.h>
 #include <ATen/ops/_scaled_mm_native.h>
 #include <ATen/ops/_unsafe_view_native.h>
 #include <ATen/ops/abs.h>
@@ -61,15 +62,37 @@
 using at::blas::ScalingType;
 using at::blas::SwizzleType;
 
-namespace scaled_blas = at::cuda::scaled;
+namespace scaled_blas = at::native::scaled;
 using scaled_blas::ScaledGemmImplementation;
 using scaled_blas::convert_int_to_enum;
-using scaled_blas::_scaled_mm_allowed_device;
 
 namespace at::native {
 
 namespace {
 
+bool _scaled_mm_allowed_device(bool sm90_only=false, bool sm100_only=false) {
+#ifdef USE_ROCM
+  static const std::vector<std::string> archs = {
+    "gfx942",
+#if ROCM_VERSION >= 60300
+    "gfx1200", "gfx1201",
+#endif
+#if ROCM_VERSION >= 60500
+    "gfx950"
+#endif
+};
+  return at::detail::getCUDAHooks().isGPUArch(archs);
+#else
+  auto dprops = at::cuda::getCurrentDeviceProperties();
+
+  if (sm90_only || sm100_only) {
+    return (sm90_only && dprops->major == 9) || (sm100_only && dprops->major == 10);
+  } else {
+    return dprops->major >= 9 || (dprops->major == 8 && dprops->minor == 9);
+  }
+#endif
+}
+
 // 2d-2d and 2d-3d
 // scaling=MXFP8
 // CUDA-only
@@ -110,10 +133,10 @@ _mx8_mx8_bf16_grouped_mm_mslk(
         scale_b,
         offs.value(),
         out);
+    return out;
 #else
     TORCH_CHECK_NOT_IMPLEMENTED(false, "mxfp8_mxfp8 grouped gemm requires compile with USE_MSLK");
 #endif
-    return out;
 }
 
 // 2d-2d and 2d-3d cases
@@ -130,7 +153,7 @@ _f8_f8_bf16_rowwise_grouped_mm_cuda(
           const bool use_fast_accum,
           Tensor& out) {
   TORCH_CHECK_VALUE(mat_a.dtype() == at::kFloat8_e4m3fn, "Expected mat_a to be Float8_e4m3 matrix got ", mat_a.scalar_type());
-  TORCH_CHECK_VALUE(mat_b.dtype() == at::kFloat8_e4m3fn, "Expected mat_a to be Float8_e4m3 matrix got ", mat_b.scalar_type());
+  TORCH_CHECK_VALUE(mat_b.dtype() == at::kFloat8_e4m3fn, "Expected mat_b to be Float8_e4m3 matrix got ", mat_b.scalar_type());
 
   at::cuda::detail::f8f8bf16_grouped_mm(
       mat_a,
@@ -175,10 +198,10 @@ _f8_f8_bf16_rowwise_grouped_mm_rocm(
       scale_b,
       offs,
       out);
+  return out;
 #else
   TORCH_CHECK_NOT_IMPLEMENTED(false, "grouped gemm is not supported without USE_MSLK on ROCM")
 #endif
-  return out;
 
 }
 #endif // USE_ROCM
@@ -269,11 +292,11 @@ _f4_f4_bf16_grouped_mm_mslk(
       out,
       combined_global_scale
   );
+
+  return out;
 #else
   TORCH_CHECK_NOT_IMPLEMENTED(false, "nvfp4 grouped gemm is not supported without USE_MSLK, and only for CUDA")
 #endif
-
-  return out;
 }
 
 void _check_scales_fp8_rowwise(const Tensor& mat, const Tensor& scale, const int dim, const int arg_idx, const int scale_multiplier=1) {
diff --git a/aten/src/ATen/native/cuda/IndexKernelUtils.cu b/aten/src/ATen/native/cuda/IndexKernelUtils.cu
index e44450c542e83..d9bf529f38db1 100644
--- a/aten/src/ATen/native/cuda/IndexKernelUtils.cu
+++ b/aten/src/ATen/native/cuda/IndexKernelUtils.cu
@@ -30,7 +30,7 @@ void vectorized_gather_kernel_launch(char * out, char * inp, index_t * idx, int
   constexpr int64_t max_num_threads=256;
   auto num_threads = at::round_up(
       at::ceil_div(slice_size_in_bytes, Alignment),
-      static_cast<int64_t>(C10_WARP_SIZE));
+      static_cast<int64_t>(at::cuda::warp_size()));
   uint32_t grid_y = at::cuda::getCurrentDeviceProperties()->maxGridSize[1];
   grid_y = std::min(static_cast<uint32_t>(at::ceil_div(slice_size_in_bytes, max_num_threads * Alignment)), grid_y);
   dim3 grid = {static_cast<uint32_t>(num_ind), grid_y, 1};
diff --git a/aten/src/ATen/native/cuda/Indexing.cu b/aten/src/ATen/native/cuda/Indexing.cu
index 44f4a552e758b..a052943c597ec 100644
--- a/aten/src/ATen/native/cuda/Indexing.cu
+++ b/aten/src/ATen/native/cuda/Indexing.cu
@@ -38,6 +38,7 @@
 
 #include <ATen/cuda/CUDAContext.h>
 #include <ATen/cuda/cub.h>
+#include <ATen/cuda/detail/IntegerDivider.cuh>
 #include <c10/util/irange.h>
 #include <c10/core/QScheme.h>
 #include <ATen/native/quantized/AffineQuantizerBase.h>
@@ -1018,7 +1019,7 @@ __global__ void indexFuncSmallIndex(cuda::detail::TensorInfo<T, IndexType> dst,
     // Lua indices begin at 1
     IndexType dstIndex =
         indices.data[cuda::detail::IndexToOffset<const IndicesType, IndexType, IdxDim>::get(srcIndex, indices)];
-    CUDA_KERNEL_ASSERT(dstIndex < dstAddDimSize);
+    CUDA_KERNEL_ASSERT(dstIndex < static_cast<IndexType>(dstAddDimSize));
 
     // We stride over the output ignoring the indexed dimension
     // (innerSize), whose offset calculation is handled differently
@@ -1046,8 +1047,18 @@ __global__ void indexFuncSmallIndex(cuda::detail::TensorInfo<T, IndexType> dst,
 // the number of indices chosen is small, then the
 // indexFuncSmallIndex kernel is a better choice to reduce memory
 // accesses.
+//
+// When VecSize > 1, each thread processes VecSize consecutive elements,
+// amortizing the divmod and index lookup cost.  Boundary cases (where
+// VecSize elements would cross a slice boundary) are handled inline so
+// the dispatch site does not need to check sliceSize alignment or stride.
+// When VecSize == 1 (default), this is the original scalar kernel.
+//
+// Uses IntDivider for fast integer divmod (multiply+shift instead of
+// hardware division).  For uint32_t this is a significant win on
+// non-power-of-2 innerSize; for power-of-2 it matches compiler shifts.
 template <typename T, typename IndicesType, typename IndexType, int DstDim, int SrcDim, int IdxDim,
-          bool IndexIsMajor, typename func_t>
+          bool IndexIsMajor, int VecSize = 1, typename func_t>
 __global__ void indexFuncLargeIndex(cuda::detail::TensorInfo<T, IndexType> dst,
                                     cuda::detail::TensorInfo<const T, IndexType> src,
                                     cuda::detail::TensorInfo<const IndicesType, IndexType> indices,
@@ -1058,37 +1069,86 @@ __global__ void indexFuncLargeIndex(cuda::detail::TensorInfo<T, IndexType> dst,
                                     int64_t dstAddDimSize,
                                     int64_t dstNumel,
                                     const func_t& op,
-                                    T alpha) {
-  // We stride over the output including the indexed dimension
-  // (totalSize), and calculate the destination index point based on that
-  for (IndexType linearIndex = blockIdx.x * blockDim.x + threadIdx.x;
-       linearIndex < totalSize;
-       linearIndex += gridDim.x * blockDim.x) {
+                                    T alpha,
+                                    cuda::detail::IntDivider<IndexType> innerSizeDivider) {
+  // Helper to process a single element at linearIndex.
+  auto processElement = [&](IndexType linearIndex) {
+    auto dm = innerSizeDivider.divmod(linearIndex);
     IndexType srcIndex, elementInSlice;
-    if (IndexIsMajor) {
-      srcIndex = linearIndex / innerSize;
-      elementInSlice = linearIndex % innerSize;
-    }
-    else {
-      elementInSlice = linearIndex / innerSize;
-      srcIndex = linearIndex % innerSize;
+    if constexpr (IndexIsMajor) {
+      srcIndex = dm.div;
+      elementInSlice = dm.mod;
+    } else {
+      elementInSlice = dm.div;
+      srcIndex = dm.mod;
     }
 
-    // Lua indices begin at 1
     IndexType dstIndex =
         indices.data[cuda::detail::IndexToOffset<const IndicesType, IndexType, IdxDim>::get(srcIndex, indices)];
-    CUDA_KERNEL_ASSERT(dstIndex < dstAddDimSize);
+    CUDA_KERNEL_ASSERT(dstIndex < static_cast<IndexType>(dstAddDimSize));
 
     IndexType dstOffset =
-      cuda::detail::IndexToOffset<T, IndexType, DstDim>::get(elementInSlice, dst);
+        cuda::detail::IndexToOffset<T, IndexType, DstDim>::get(elementInSlice, dst);
     dstOffset += dstIndex * dst.strides[dstAddDim];
 
     IndexType srcOffset =
-      cuda::detail::IndexToOffset<const T, IndexType, SrcDim>::get(elementInSlice, src);
+        cuda::detail::IndexToOffset<const T, IndexType, SrcDim>::get(elementInSlice, src);
     srcOffset += srcIndex * src.strides[srcAddDim];
 
     T val = src.data[srcOffset] * alpha;
     op(dst.data, dstOffset, dstNumel, &val);
+  };
+
+  if constexpr (VecSize > 1) {
+    // Round up to include tail elements in the vectorized loop.
+    const IndexType totalVecs = at::ceil_div(totalSize, (IndexType)VecSize);
+    for (IndexType vecIdx = blockIdx.x * blockDim.x + threadIdx.x;
+         vecIdx < totalVecs;
+         vecIdx += gridDim.x * blockDim.x) {
+      IndexType baseLinear = vecIdx * VecSize;
+      auto dm = innerSizeDivider.divmod(baseLinear);
+      IndexType srcIndex = dm.div;
+      IndexType elementBase = dm.mod;
+
+      if (elementBase + VecSize <= innerSize && baseLinear + VecSize <= totalSize) {
+        // Fast path: all VecSize elements are in the same slice and within bounds.
+        // Hoist index lookup outside the unrolled loop.
+        IndexType dstIndex =
+            indices.data[cuda::detail::IndexToOffset<const IndicesType, IndexType, IdxDim>::get(srcIndex, indices)];
+        CUDA_KERNEL_ASSERT(dstIndex < static_cast<IndexType>(dstAddDimSize));
+
+        #pragma unroll
+        for (int v = 0; v < VecSize; ++v) {
+          IndexType elementInSlice = elementBase + v;
+          IndexType dstOffset =
+              cuda::detail::IndexToOffset<T, IndexType, DstDim>::get(elementInSlice, dst);
+          dstOffset += dstIndex * dst.strides[dstAddDim];
+
+          IndexType srcOffset =
+              cuda::detail::IndexToOffset<const T, IndexType, SrcDim>::get(elementInSlice, src);
+          srcOffset += srcIndex * src.strides[srcAddDim];
+
+          T val = src.data[srcOffset] * alpha;
+          op(dst.data, dstOffset, dstNumel, &val);
+        }
+      } else {
+        // Slow path: elements cross a slice boundary or are in the tail.
+        #pragma unroll
+        for (int v = 0; v < VecSize; ++v) {
+          IndexType li = baseLinear + v;
+          if (li < totalSize) {
+            processElement(li);
+          }
+        }
+      }
+    }
+  } else {
+    // Scalar path: one element per thread.
+    for (IndexType linearIndex = blockIdx.x * blockDim.x + threadIdx.x;
+         linearIndex < totalSize;
+         linearIndex += gridDim.x * blockDim.x) {
+      processElement(linearIndex);
+    }
   }
 }
 
@@ -1180,13 +1240,19 @@ void index_add_cuda_impl(const Tensor& self, int64_t dim, const Tensor& index, c
 
 #define LARGE_INDEX(TENSOR_TYPE, INDICES_TYPE, TYPE,                        \
                     SELF_DIM, SOURCE_DIM, IDX_DIM, IDX_IS_MAJOR)            \
-  indexFuncLargeIndex<TENSOR_TYPE, INDICES_TYPE, TYPE,                      \
-                      SELF_DIM, SOURCE_DIM, IDX_DIM, IDX_IS_MAJOR>          \
-    <<<largeIndexGrid, largeIndexBlock, 0, stream>>>(                       \
-      selfInfo, sourceInfo, indexInfo,                                      \
-      selfAddDim, sourceAddDim, sourceTotalSize,                            \
-      (IDX_IS_MAJOR) ? sliceSize : numIndex,                                \
-      selfAddDimSize, selfNumel, reduce_add, alpha_value);                  \
+  {                                                                          \
+    const TYPE innerSizeVal =                                                \
+        static_cast<TYPE>((IDX_IS_MAJOR) ? sliceSize : numIndex);            \
+    cuda::detail::IntDivider<TYPE> innerSizeDivider(innerSizeVal);           \
+    indexFuncLargeIndex<TENSOR_TYPE, INDICES_TYPE, TYPE,                     \
+                        SELF_DIM, SOURCE_DIM, IDX_DIM,                       \
+                        IDX_IS_MAJOR, (IDX_IS_MAJOR) ? 4 : 1>               \
+      <<<largeIndexGrid, largeIndexBlock, 0, stream>>>(                     \
+        selfInfo, sourceInfo, indexInfo,                                     \
+        selfAddDim, sourceAddDim, sourceTotalSize,                          \
+        innerSizeVal, selfAddDimSize, selfNumel,                            \
+        reduce_add, alpha_value, innerSizeDivider);                         \
+  }                                                                          \
   C10_CUDA_KERNEL_LAUNCH_CHECK();
 
   uint64_t defaultMaxBlockThreads = getDefaultMaxThreadsPerBlock();
@@ -1355,13 +1421,19 @@ void index_reduce_func_cuda_impl(
 
 #define LARGE_INDEX(TENSOR_TYPE, INDICES_TYPE, TYPE,                                     \
                     SELF_DIM, SOURCE_DIM, IDX_DIM, IDX_IS_MAJOR)                         \
-  indexFuncLargeIndex<TENSOR_TYPE, INDICES_TYPE, TYPE,                                   \
-                     SELF_DIM, SOURCE_DIM, IDX_DIM, IDX_IS_MAJOR>                        \
-    <<<largeIndexGrid, largeIndexBlock, 0, stream>>>(                                    \
-      selfInfo, sourceInfo, indexInfo,                                                   \
-      selfReduceDim, sourceReduceDim, sourceTotalSize,                                   \
-      (IDX_IS_MAJOR) ? sliceSize : numIndex,                                             \
-      selfReduceDimSize, selfNumel, reduce_func, alpha_value);                           \
+  {                                                                                       \
+    const TYPE innerSizeVal =                                                             \
+        static_cast<TYPE>((IDX_IS_MAJOR) ? sliceSize : numIndex);                         \
+    cuda::detail::IntDivider<TYPE> innerSizeDivider(innerSizeVal);                        \
+    indexFuncLargeIndex<TENSOR_TYPE, INDICES_TYPE, TYPE,                                  \
+                        SELF_DIM, SOURCE_DIM, IDX_DIM,                                     \
+                        IDX_IS_MAJOR, (IDX_IS_MAJOR) ? 4 : 1>                              \
+      <<<largeIndexGrid, largeIndexBlock, 0, stream>>>(                                   \
+        selfInfo, sourceInfo, indexInfo,                                                   \
+        selfReduceDim, sourceReduceDim, sourceTotalSize,                                  \
+        innerSizeVal, selfReduceDimSize, selfNumel,                                       \
+        reduce_func, alpha_value, innerSizeDivider);                                      \
+  }                                                                                       \
   C10_CUDA_KERNEL_LAUNCH_CHECK();
 
   uint64_t defaultMaxBlockThreads = getDefaultMaxThreadsPerBlock();
@@ -1795,36 +1867,6 @@ Tensor & masked_fill__cuda(Tensor& self, const Tensor & mask, const Tensor & val
   return masked_fill__cuda(self, mask, value.item());
 }
 
-namespace {
-
-// ForwardIt: only legacy random access iterator is supported.
-template<class ForwardIt, class T, bool is_lower = true>
-static __host__ __device__ __forceinline__
-ForwardIt find_bound(ForwardIt first, ForwardIt last, const T& value) {
-    ForwardIt it;
-    typename std::iterator_traits<ForwardIt>::difference_type count, step;
-    // NOTE: std::distance(first, last) compiles but produces wrong results here,
-    // so only legacy random access iterators are safe in this code.
-    count = last - first;
-
-    while (count > 0) {
-      it = first;
-      step = count / 2;
-      // avoiding std::advance(it, step),
-      // although it does work unlike std::distance
-      it += step;
-      if (is_lower ? *it < value : value >= *it) {
-        first = ++it;
-        count -= step + 1;
-      }
-      else {
-        count = step;
-      }
-    }
-    return first;
-}
-
-}
 
 Tensor index_select_sparse_cuda(const Tensor& self, int64_t dim, const Tensor& index) {
   const auto ndim = self.dim();
@@ -1916,12 +1958,12 @@ Tensor index_select_sparse_cuda(const Tensor& self, int64_t dim, const Tensor& i
               [ptr_intrsc_counts_nneg_index, ptr_sorted_dim_indices, nnz] GPU_LAMBDA (
                 index_t idx_val, index_t idx_idx
               ) -> index_t {
-                auto* lb = find_bound<const index_t*, index_t, true>(
+                auto* lb = at::cuda::detail::find_bound<const index_t*, index_t, true>(
                   ptr_sorted_dim_indices,
                   ptr_sorted_dim_indices + nnz,
                   idx_val
                 );
-                auto* ub = find_bound<const index_t*, index_t, false>(
+                auto* ub = at::cuda::detail::find_bound<const index_t*, index_t, false>(
                   ptr_sorted_dim_indices,
                   ptr_sorted_dim_indices + nnz,
                   idx_val
diff --git a/aten/src/ATen/native/cuda/KernelUtils.cuh b/aten/src/ATen/native/cuda/KernelUtils.cuh
index 9c100ca206adf..12feeb6d63af3 100644
--- a/aten/src/ATen/native/cuda/KernelUtils.cuh
+++ b/aten/src/ATen/native/cuda/KernelUtils.cuh
@@ -228,7 +228,9 @@ __device__ __forceinline__ void fastAtomicAdd(
 // This function implements a committed store.
 // Upon returning, the store is committed to global memory.
 // This is useful in avoiding the need for fences.
-template <typename T>
+// If multiple stores are done in a row there is option to skip
+// waiting for commit for all but the last store.
+template <typename T, bool wait_for_commit = true>
 __device__ inline void cmtdStore(void* address, T value) {
       int constexpr num_long_per_val = sizeof(value)/sizeof(long);
       int constexpr num_int_per_val = sizeof(value)/sizeof(int);
@@ -252,13 +254,16 @@ __device__ inline void cmtdStore(void* address, T value) {
       else if constexpr (num_char_per_val*sizeof(char) == sizeof(value))
         for (int i=0; i<num_char_per_val; i++)
           __hip_atomic_store(reinterpret_cast<char *>(address)+i, _pnr.c[i], __ATOMIC_RELAXED, __HIP_MEMORY_SCOPE_AGENT);
-      __atomic_signal_fence(__ATOMIC_SEQ_CST);
+      if constexpr (wait_for_commit)
+      {
+        __atomic_signal_fence(__ATOMIC_SEQ_CST);
 #ifdef __gfx1250__
-      asm volatile("s_wait_loadcnt(0)" ::: "memory");
+        asm volatile("s_wait_loadcnt(0)" ::: "memory");
 #else
-      asm volatile("s_waitcnt vmcnt(0)" ::: "memory");
+        asm volatile("s_waitcnt vmcnt(0)" ::: "memory");
 #endif
-      __atomic_signal_fence(__ATOMIC_SEQ_CST);
+        __atomic_signal_fence(__ATOMIC_SEQ_CST);
+      }
 }
 #endif
 
diff --git a/aten/src/ATen/native/cuda/Loops.cuh b/aten/src/ATen/native/cuda/Loops.cuh
index a80c51fa6a9cb..063db29d38279 100644
--- a/aten/src/ATen/native/cuda/Loops.cuh
+++ b/aten/src/ATen/native/cuda/Loops.cuh
@@ -8,6 +8,7 @@
 #include <ATen/native/cuda/thread_constants.h>
 #include <ATen/native/cuda/MemoryAccess.cuh>
 
+#include <c10/util/C++17.h>
 #include <tuple>
 
 
@@ -80,7 +81,7 @@ __device__ inline void elementwise_kernel_helper(func_t f, policy_t policy) {
 namespace at:: native {
 
 template <typename func_t>
-void gpu_kernel_nocast(TensorIteratorBase& iter, const func_t& f) {
+void gpu_kernel_nocast(TensorIteratorBase& iter, const func_t& f, bool check_cast = true) {
 
   for (int arg = 0; arg < iter.ntensors(); arg++) {
     TORCH_INTERNAL_ASSERT(
@@ -94,14 +95,22 @@ void gpu_kernel_nocast(TensorIteratorBase& iter, const func_t& f) {
 
   if (!iter.can_use_32bit_indexing()) {
     for (auto& sub_iter : iter.with_32bit_indexing()) {
-      gpu_kernel_nocast(sub_iter, f);
+      gpu_kernel_nocast(sub_iter, f, check_cast);
     }
     return;
   }
 
+  if (check_cast) {
+    TORCH_INTERNAL_ASSERT(!needs_dynamic_casting<func_t>::check(iter));
+  }
   gpu_kernel_impl_nocast(iter, f);
 }
 
+template <typename func_t>
+void gpu_kernel_opaque(TensorIteratorBase& iter, const func_t& f) {
+  gpu_kernel_nocast(iter, f, false);
+}
+
 template <typename func_t>
 void gpu_kernel(TensorIteratorBase& iter, const func_t& f) {
 
diff --git a/aten/src/ATen/native/cuda/MemoryAccess.cuh b/aten/src/ATen/native/cuda/MemoryAccess.cuh
index 373b44cca7901..ac248d8cc00c8 100644
--- a/aten/src/ATen/native/cuda/MemoryAccess.cuh
+++ b/aten/src/ATen/native/cuda/MemoryAccess.cuh
@@ -187,28 +187,35 @@ template <int vec_size, typename scalar_t>
 __device__ aligned_vector<scalar_t, vec_size> load_vector(const scalar_t *base_ptr, uint32_t offset) {
   using vec_t = aligned_vector<scalar_t, vec_size>;
   auto *from = reinterpret_cast<const vec_t *>(base_ptr);
-#if defined(USE_ROCM) && defined(__gfx942__)
-  using longx2 = __attribute__((__vector_size__(4*sizeof(int)))) int;
-  if constexpr (sizeof(vec_t) == sizeof(int)) {
-   union {
-     vec_t v;
-     int   i;
-   } tmpt = { .i = __builtin_nontemporal_load(reinterpret_cast<const int *>(&(from[offset]))) };
-   return tmpt.v;
-  }
-  else if constexpr (sizeof(vec_t) == sizeof(long)) {
-   union {
-     vec_t v;
-     long   i;
-   } tmpt = { .i = __builtin_nontemporal_load(reinterpret_cast<const long *>(&(from[offset]))) };
-   return tmpt.v;
-  }
-  else if constexpr (sizeof(vec_t) == sizeof(longx2)) {
-   union {
-     vec_t v;
-     longx2  i;
-   } tmpt = { .i = __builtin_nontemporal_load(reinterpret_cast<const longx2 *>(&(from[offset]))) };
-   return tmpt.v;
+#if defined(USE_ROCM)
+  // NOTE: This nontemporal vectorized load is tuned for gfx942 (CDNA3, Wave64).
+  // Do NOT extend to gfx1250 (GFX12.5, Wave32) without first validating
+  // intrinsic codegen and re-benchmarking — the register layout and wave
+  // width differ. gfx1250 falls through to the generic path below, which
+  // is correctness-safe.
+  if(__builtin_amdgcn_processor_is("gfx942")) {
+    using longx2 = __attribute__((__vector_size__(4*sizeof(int)))) int;
+    if constexpr (sizeof(vec_t) == sizeof(int)) {
+    union {
+      vec_t v;
+      int   i;
+    } tmpt = { .i = __builtin_nontemporal_load(reinterpret_cast<const int *>(&(from[offset]))) };
+    return tmpt.v;
+    }
+    else if constexpr (sizeof(vec_t) == sizeof(long)) {
+    union {
+      vec_t v;
+      long   i;
+    } tmpt = { .i = __builtin_nontemporal_load(reinterpret_cast<const long *>(&(from[offset]))) };
+    return tmpt.v;
+    }
+    else if constexpr (sizeof(vec_t) == sizeof(longx2)) {
+    union {
+      vec_t v;
+      longx2  i;
+    } tmpt = { .i = __builtin_nontemporal_load(reinterpret_cast<const longx2 *>(&(from[offset]))) };
+    return tmpt.v;
+    }
   }
 #endif
   return from[offset];
diff --git a/aten/src/ATen/native/cuda/MultiMarginLoss.cu b/aten/src/ATen/native/cuda/MultiMarginLoss.cu
index ad7b3638b489d..3de86ff4bfc64 100644
--- a/aten/src/ATen/native/cuda/MultiMarginLoss.cu
+++ b/aten/src/ATen/native/cuda/MultiMarginLoss.cu
@@ -2,6 +2,7 @@
 #include <ATen/core/Tensor.h>
 #include <ATen/AccumulateType.h>
 #include <ATen/Dispatch.h>
+#include <ATen/native/LossMulti.h>
 #include <ATen/native/Resize.h>
 #include <c10/cuda/CUDAStream.h>
 #include <c10/cuda/CUDAException.h>
@@ -127,38 +128,6 @@ __global__ void MultiMarginLoss_backward_kernel(
   }
 }
 
-void multi_margin_loss_shape_check(
-    int64_t& nframe,
-    int64_t& dim,
-    const int64_t& ndims,
-    const Tensor& input,
-    const Tensor& target,
-    const std::optional<Tensor>& weight) {
-    TORCH_CHECK(
-        (ndims == 2 && input.size(1) != 0) || (ndims == 1 && input.size(0) != 0) || ndims == 0,
-        "Expected non-empty vector or matrix with optional 0-dim batch size, but got: ",
-        input.sizes());
-
-    if (ndims <= 1) {
-      nframe = 1;
-      dim = ndims == 0 ? 1 : input.size(0);
-    } else {
-      nframe = input.size(0);
-      dim = input.size(1);
-    }
-
-    TORCH_CHECK(
-        target.dim() <= 1 && target.numel() == nframe,
-        "inconsistent target size, expected ", nframe, " but got ",
-        target.sizes());
-    if (weight && weight->defined()) {
-      TORCH_CHECK(
-          weight->dim() <= 1 && weight->numel() == dim,
-          "inconsistent weight size, expected ", dim, " but got ",
-          weight->sizes());
-    }
-}
-
 }  // namespace (anonymous)
 
 Tensor& multi_margin_loss_cuda_out(
@@ -196,7 +165,6 @@ Tensor& multi_margin_loss_cuda_out(
   AT_DISPATCH_FLOATING_TYPES_AND2(kHalf, kBFloat16, input.scalar_type(), "multi_margin_loss_cuda", [&] {
     const scalar_t margin = margin_.to<scalar_t>();
     if (input.dim() <= 1) {
-      TORCH_CHECK(target.dim() <= 1 && target.numel() == nframe, "inconsistent target size");
       dim3 blocks(1);
       dim3 threads(MULTIMARGIN_THREADS);
       if (p == 1) {
@@ -225,9 +193,6 @@ Tensor& multi_margin_loss_cuda_out(
     } else {
       auto in_sizes = input.sizes();
       TORCH_INTERNAL_ASSERT(in_sizes.size() == 2);
-      // allow zero-dim target for 2D input.
-      TORCH_CHECK(in_sizes[1] != 0 && target.dim() <= 1 && target.numel() == nframe,
-                "inconsistent target size");
       dim3 blocks(nframe);
       dim3 threads(MULTIMARGIN_THREADS);
 
@@ -363,8 +328,6 @@ Tensor& multi_margin_loss_cuda_backward_out(
     } else {
       auto in_sizes = input.sizes();
       TORCH_INTERNAL_ASSERT(in_sizes.size() == 2);
-      TORCH_CHECK((in_sizes[1] != 0) && (target.dim() <= 1) && (target.numel() == nframe),
-                  "inconsistent target size");
       dim3 blocks(in_sizes[0]);
       dim3 threads(MULTIMARGIN_THREADS);
 
diff --git a/aten/src/ATen/native/cuda/MultiTensorApply.cuh b/aten/src/ATen/native/cuda/MultiTensorApply.cuh
index 2fe431f778b1a..d11cfc19c5856 100644
--- a/aten/src/ATen/native/cuda/MultiTensorApply.cuh
+++ b/aten/src/ATen/native/cuda/MultiTensorApply.cuh
@@ -15,13 +15,35 @@ static constexpr int64_t kChunkSize = 65536;
 static constexpr int64_t kBlockSize = 512;
 
 // TODO(crcrpar): Add `n>5` for `low prec params & their higher prec copy`
-// TensorListMetadata has to be < 4KB - the limit for kernel launch argument
+// TensorListMetadata has to fit within the CUDA kernel launch argument limit.
+// While CUDA 12.1, driver version R530+ and Volta+ would work with 32KB, we
+// decide to be safe and only swap for CUDA 13+ during compile time. This saves
+// binary size and will guarantees 32KB kernel arg space; older versions are
+// still limited to 4KB. We adopt naive values for 32KB from
+// https://github.com/pytorch/pytorch/pull/134373.
+// TODO: The values for 32KB can very much be optimized further.
+#if defined(CUDART_VERSION) && CUDART_VERSION >= 13000 && !defined(USE_ROCM)
+
+static constexpr int depth_to_max_tensors[5] = {770, 448, 336, 252, 210};
+static constexpr int depth_to_max_blocks[5] = {2240, 2240, 2240, 2240, 2240};
+static constexpr int depth_to_max_tensors_scalarlist[5] =
+    {672, 448, 336, 252, 210};
+static constexpr int depth_to_max_tensors_scalarlist_of_complex_double[2] = {
+    504,
+    420};
+using block_index_t = uint16_t;
+
+#else
+
 static constexpr int depth_to_max_tensors[5] = {110, 64, 48, 36, 30};
 static constexpr int depth_to_max_blocks[5] = {320, 320, 320, 320, 320};
 static constexpr int depth_to_max_tensors_scalarlist[5] = {96, 64, 48, 36, 30};
 static constexpr int depth_to_max_tensors_scalarlist_of_complex_double[2] = {
     72,
     60};
+using block_index_t = unsigned char;
+
+#endif
 
 template <typename T>
 __device__ __forceinline__ bool is_aligned(T* p) {
@@ -42,7 +64,7 @@ template <int n>
 struct TensorListMetadata {
   const void* addresses[n][depth_to_max_tensors[n - 1]];
   int64_t numel_for_tensor[depth_to_max_tensors[n - 1]];
-  unsigned char block_to_tensor[depth_to_max_blocks[n - 1]];
+  block_index_t block_to_tensor[depth_to_max_blocks[n - 1]];
   int block_to_chunk[depth_to_max_blocks[n - 1]];
   int start_tensor_this_launch;
 };
@@ -52,12 +74,12 @@ struct TensorListScalarListMetadata {
   const void* addresses[n][depth_to_max_tensors_scalarlist[n - 1]];
   int64_t numel_for_tensor[depth_to_max_tensors_scalarlist[n - 1]];
   scalar_vals_t scalar_vals[depth_to_max_tensors_scalarlist[n - 1]];
-  unsigned char block_to_tensor[depth_to_max_blocks[n - 1]];
+  block_index_t block_to_tensor[depth_to_max_blocks[n - 1]];
   int block_to_chunk[depth_to_max_blocks[n - 1]];
 };
 
-// note(mkozuki): `n` of 1&2 violate the limit of cuda kernel argument size of
-// 4kb with `c10::complex<double>`
+// note(mkozuki): `n` of 1&2 violate the limit of cuda kernel argument size
+// with `c10::complex<double>`
 template <>
 struct TensorListScalarListMetadata<c10::complex<double>, 1> {
   const void* addresses[1]
@@ -66,7 +88,7 @@ struct TensorListScalarListMetadata<c10::complex<double>, 1> {
       numel_for_tensor[depth_to_max_tensors_scalarlist_of_complex_double[0]];
   c10::complex<double>
       scalar_vals[depth_to_max_tensors_scalarlist_of_complex_double[0]];
-  unsigned char block_to_tensor[depth_to_max_blocks[1 - 1]];
+  block_index_t block_to_tensor[depth_to_max_blocks[1 - 1]];
   int block_to_chunk[depth_to_max_blocks[1 - 1]];
 };
 
@@ -78,19 +100,22 @@ struct TensorListScalarListMetadata<c10::complex<double>, 2> {
       numel_for_tensor[depth_to_max_tensors_scalarlist_of_complex_double[1]];
   c10::complex<double>
       scalar_vals[depth_to_max_tensors_scalarlist_of_complex_double[1]];
-  unsigned char block_to_tensor[depth_to_max_blocks[2 - 1]];
+  block_index_t block_to_tensor[depth_to_max_blocks[2 - 1]];
   int block_to_chunk[depth_to_max_blocks[2 - 1]];
 };
 
 // NOTE(crcrpar): This is a conservative resolution to handle `state_steps`
 // whose each element is `at::Tensor` of 1 element representing the number of
 // `step`s called so far.
+// We're aware this struct overflows the kernel arg limit at n=1 (4244 bytes),
+// but our current fused optimizers only instantiate at n>=4 so it's not a
+// concern (yet).
 template <int n>
 struct FusedOptimizerTensorListMetadata {
   const void* addresses[n][depth_to_max_tensors[n - 1]];
   int64_t numel_for_tensor[depth_to_max_tensors[n - 1]];
-  const void* state_steps_addresses[depth_to_max_tensors_scalarlist[n - 1]];
-  unsigned char block_to_tensor[depth_to_max_blocks[n - 1]];
+  const void* state_steps_addresses[depth_to_max_tensors[n - 1]];
+  block_index_t block_to_tensor[depth_to_max_blocks[n - 1]];
   int block_to_chunk[depth_to_max_blocks[n - 1]];
   int start_tensor_this_launch;
 };
diff --git a/aten/src/ATen/native/cuda/Normalization.cuh b/aten/src/ATen/native/cuda/Normalization.cuh
index bbd65419bbb92..d9ceeea53ceec 100644
--- a/aten/src/ATen/native/cuda/Normalization.cuh
+++ b/aten/src/ATen/native/cuda/Normalization.cuh
@@ -8,6 +8,7 @@
 #include <ATen/cuda/DeviceUtils.cuh>
 #include <ATen/native/cuda/block_reduce.cuh>
 #include <ATen/native/cuda/DeviceSqrt.cuh>
+#include <ATen/native/cuda/KernelUtils.cuh>
 #include <ATen/native/cuda/LaunchUtils.h>
 #include <c10/macros/Macros.h>
 
@@ -133,7 +134,7 @@ __device__ scalar_t reduce(Op op, PTA tensor, int plane) {
     }
 #endif
   }
-  __shared__ scalar_t shared[C10_WARP_SIZE];
+  __shared__ scalar_t shared[C10_WARP_SIZE_UPPER_BOUND];
   SumReduceOp<scalar_t> reduce_op;
   sum = cuda_utils::BlockReduce<scalar_t, SumReduceOp<scalar_t>, cuda_utils::Block2D>(sum, reduce_op, 0, shared);
   if (threadIdx.x == 0 && threadIdx.y == 0) {
@@ -287,7 +288,7 @@ __global__ void batch_norm_collect_statistics_kernel(
     GenericPackedTensorAccessor<stat_accscalar_t, 1, RestrictPtrTraits, index_t> save_mean,
     GenericPackedTensorAccessor<stat_accscalar_t, 1, RestrictPtrTraits, index_t> save_transformed_var) {
 
-  __shared__ int shared_n[2 * 2 * C10_WARP_SIZE + C10_WARP_SIZE];
+  __shared__ int shared_n[2 * 2 * C10_WARP_SIZE_UPPER_BOUND + C10_WARP_SIZE_UPPER_BOUND];
 
   int plane = blockIdx.x;
   int N = input.size(0) * input.size(2);
@@ -749,7 +750,7 @@ void batch_norm_elemt_cuda_template(const Tensor& output_, const Tensor& input_,
 
   // The input_transform kernel is pointwise, but we need to balance reading parameters (save_var/mean,
   // weight/bias) - which we only do once and have a for loop afterwards - with having many threads and blocks
-  // and good occupancy. Quiet likely, we could go with even more blocks than 1024.
+  // and good occupancy. Quite likely, we could go with even more blocks than 1024.
   // The various planes are independent, so we use blocks for them.
   int tf = std::max<int>(getNumThreads(input.size(2)/4),
                          std::min<int>(getNumThreads(input.size(2)), 64));
@@ -1063,12 +1064,22 @@ batch_norm_collect_statistics_channels_last_kernel(
     address_base = c_offset + blockIdx.y * stride;
     // write data to staging_data;
     if (threadIdx.y == 0 && c_offset < stride) {
+#ifndef USE_ROCM
       staging_mean[address_base] = mean_th;
       staging_m2n[address_base] = m2_th;
       staging_count[address_base] = count_th;
+#else
+      // In architectures with split caches, global fences are costly.
+      // Here we preempt need for fences by committing stores to global memory.
+      cmtdStore<accscalar_t, false>((void*)&staging_mean[address_base], mean_th);
+      cmtdStore<accscalar_t, false>((void*)&staging_m2n[address_base], m2_th);
+      cmtdStore((void*)&staging_count[address_base], count_th);
+#endif
     }
 
+#ifndef USE_ROCM
     __threadfence();
+#endif
     __syncthreads(); // ensuring writes to staging_ is visible to all blocks
 
     __shared__ bool is_last_block_done;
@@ -1288,11 +1299,20 @@ __global__ void batch_norm_backward_reduce_channels_last_kernel(
     address_base = c_offset + blockIdx.y * stride;
     // write data to staging_data;
     if (threadIdx.y == 0 && c_offset < stride) {
+#ifndef USE_ROCM
       staging_sum_dy[address_base] = sum_dy_th;
       staging_sum_dy_xmu[address_base] = sum_dy_xmu_th;
+#else
+      // In architectures with split caches, global fences are costly.
+      // Here we preempt need for fences by committing stores to global memory.
+      cmtdStore<accscalar_t, false>((void*)&staging_sum_dy[address_base], sum_dy_th);
+      cmtdStore((void*)&staging_sum_dy_xmu[address_base], sum_dy_xmu_th);
+#endif
     }
 
+#ifndef USE_ROCM
     __threadfence();
+#endif
     __syncthreads(); // ensuring writes to staging_ is visible to all blocks
 
     __shared__ bool is_last_block_done;
diff --git a/aten/src/ATen/native/cuda/PersistentSoftmax.cuh b/aten/src/ATen/native/cuda/PersistentSoftmax.cuh
index f0871fa0ead6f..fee3d0a5aba1d 100644
--- a/aten/src/ATen/native/cuda/PersistentSoftmax.cuh
+++ b/aten/src/ATen/native/cuda/PersistentSoftmax.cuh
@@ -64,12 +64,12 @@ __device__ __forceinline__ void warp_reduce(acc_t* sum) {
 // input_t=half,  acc_t=float, output_t=float => read half tensor, float accumulators, write float tensor.
 // input_t_float, acc_t=float, output_t=half  => read float tensor, float accumulators, write half tensor.
 
-template <typename input_t, typename output_t, typename acc_t, int log2_elements, bool is_log_softmax, bool is_masked>
+template <typename input_t, typename output_t, typename acc_t, int log2_elements, bool is_log_softmax, bool is_masked, int WARP_SIZE_PARAM>
 __global__ void softmax_warp_forward(output_t *dst, const input_t *src, int batch_size, int stride, int element_count, const bool *mask = nullptr, const int head_chunk_size = -1, bool is_transformer_mask = false)
 {
     // WARP_SIZE and WARP_BATCH must match the return values batches_per_warp and warp_size of method warp_softmax_forward_kernel.
     constexpr int next_power_of_two = 1 << log2_elements;
-    constexpr int WARP_SIZE = (next_power_of_two < C10_WARP_SIZE) ? next_power_of_two : C10_WARP_SIZE;
+    constexpr int WARP_SIZE = (next_power_of_two < WARP_SIZE_PARAM) ? next_power_of_two : WARP_SIZE_PARAM;
     constexpr int WARP_ITERATIONS = next_power_of_two / WARP_SIZE;
     constexpr int WARP_BATCH = (next_power_of_two <= 128) ? 2 : 1;
 
@@ -211,12 +211,12 @@ __global__ void softmax_warp_forward(output_t *dst, const input_t *src, int batc
     }
 }
 
-template <typename input_t, typename output_t, typename acc_t, int log2_elements, bool is_log_softmax, bool is_masked>
+template <typename input_t, typename output_t, typename acc_t, int log2_elements, bool is_log_softmax, bool is_masked, int WARP_SIZE_PARAM>
 __global__ void softmax_warp_backward(output_t *gradInput, const input_t *grad, const input_t *output, int batch_size, int stride, int element_count, const bool *mask = nullptr)
 {
     // WARP_SIZE and WARP_BATCH must match the return values batches_per_warp and warp_size of method warp_softmax_backward_kernel.
     constexpr int next_power_of_two = 1 << log2_elements;
-    constexpr int WARP_SIZE = (next_power_of_two < C10_WARP_SIZE) ? next_power_of_two : C10_WARP_SIZE;
+    constexpr int WARP_SIZE = (next_power_of_two < WARP_SIZE_PARAM) ? next_power_of_two : WARP_SIZE_PARAM;
     constexpr int WARP_ITERATIONS = next_power_of_two / WARP_SIZE;
     constexpr int WARP_BATCH = (next_power_of_two <= 128) ? 2 : 1;
 
@@ -325,12 +325,29 @@ void dispatch_softmax_forward(output_t *dst, const input_t *src, int softmax_ele
         dim3 threads(warp_size, warps_per_block, 1);
         // Launch code would be more elegant if C++ supported FOR CONSTEXPR
         switch (log2_elements) {
+#ifdef USE_ROCM
+            // To support ROCm amdgcnspirv target, we must compile both a 32 and 64 warpSize version of each kernel
             #define LAUNCH_SOFTMAX_WARP_FORWARD(L2E) case L2E:                    \
-            softmax_warp_forward<input_t, output_t, acc_t, L2E, is_log_softmax, is_masked>   \
+            if (warp_size == 64) { \
+              softmax_warp_forward<input_t, output_t, acc_t, L2E, is_log_softmax, is_masked, 64>   \
+                  <<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(dst,   \
+                      src, batch_count, softmax_elements_stride, softmax_elements, mask, chunk_size, is_transformer_mask); \
+            } \
+            else { \
+              softmax_warp_forward<input_t, output_t, acc_t, L2E, is_log_softmax, is_masked, 32>   \
+                  <<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(dst,   \
+                      src, batch_count, softmax_elements_stride, softmax_elements, mask, chunk_size, is_transformer_mask); \
+            } \
+            C10_CUDA_KERNEL_LAUNCH_CHECK();                                       \
+            break;
+#else
+            #define LAUNCH_SOFTMAX_WARP_FORWARD(L2E) case L2E:                    \
+            softmax_warp_forward<input_t, output_t, acc_t, L2E, is_log_softmax, is_masked, 32>   \
                 <<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(dst,   \
                     src, batch_count, softmax_elements_stride, softmax_elements, mask, chunk_size, is_transformer_mask); \
             C10_CUDA_KERNEL_LAUNCH_CHECK();                                       \
             break;
+#endif
 
             LAUNCH_SOFTMAX_WARP_FORWARD(0);  // 1
             LAUNCH_SOFTMAX_WARP_FORWARD(1);  // 2
@@ -376,13 +393,32 @@ void dispatch_softmax_backward(output_t *grad_input, const input_t *grad, const
         dim3 threads(warp_size, warps_per_block, 1);
         // Launch code would be more elegant if C++ supported FOR CONSTEXPR
         switch (log2_elements) {
+#ifdef USE_ROCM
+            // To support ROCm amdgcnspirv target, we must compile both a 32 and 64 warpSize version of each kernel
+            #define LAUNCH_SOFTMAX_WARP_BACKWARD(L2E) case L2E:                      \
+            if (warp_size == 64) { \
+              softmax_warp_backward<input_t, output_t, acc_t, L2E, is_log_softmax, is_masked, 64> \
+                  <<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>       \
+                  (grad_input, grad, output, batch_count, softmax_elements_stride, \
+                  softmax_elements, mask);                                              \
+            } \
+            else { \
+              softmax_warp_backward<input_t, output_t, acc_t, L2E, is_log_softmax, is_masked, 32> \
+                  <<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>       \
+                  (grad_input, grad, output, batch_count, softmax_elements_stride, \
+                  softmax_elements, mask);                                              \
+            } \
+            C10_CUDA_KERNEL_LAUNCH_CHECK();                                      \
+            break;
+#else
             #define LAUNCH_SOFTMAX_WARP_BACKWARD(L2E) case L2E:                      \
-            softmax_warp_backward<input_t, output_t, acc_t, L2E, is_log_softmax, is_masked> \
+            softmax_warp_backward<input_t, output_t, acc_t, L2E, is_log_softmax, is_masked, 32> \
                 <<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>       \
                 (grad_input, grad, output, batch_count, softmax_elements_stride, \
                 softmax_elements, mask);                                              \
             C10_CUDA_KERNEL_LAUNCH_CHECK();                                      \
             break;
+#endif
 
             LAUNCH_SOFTMAX_WARP_BACKWARD(0); // 1
             LAUNCH_SOFTMAX_WARP_BACKWARD(1); // 2
diff --git a/aten/src/ATen/native/cuda/PhiloxDistribution.cu b/aten/src/ATen/native/cuda/PhiloxDistribution.cu
new file mode 100644
index 0000000000000..193ea9618e748
--- /dev/null
+++ b/aten/src/ATen/native/cuda/PhiloxDistribution.cu
@@ -0,0 +1,322 @@
+#define TORCH_ASSERT_ONLY_METHOD_OPERATORS
+
+#include <ATen/core/Tensor.h>
+#include <ATen/cuda/CUDAContext.h>
+#include <ATen/cuda/StatelessPhilox4x32.cuh>
+#include <ATen/Dispatch.h>
+#include <ATen/ExpandUtils.h>
+#include <ATen/cuda/detail/OffsetCalculator.cuh>
+#include <ATen/native/cuda/MemoryAccess.cuh>
+#include <ATen/core/TransformationHelper.h>
+#include <type_traits>
+
+#ifndef AT_PER_OPERATOR_HEADERS
+#include <ATen/NativeFunctions.h>
+#else
+#include <ATen/ops/_philox_normal_native.h>
+#include <ATen/ops/_philox_uniform_native.h>
+#endif
+
+namespace at::native {
+
+namespace {
+
+using at::cuda::philox_4x32;
+
+// Elements produced per Philox 4x32 call: 4 for float/half/bfloat16, 2 for double.
+// Note that we use a full float for each generated half/bfloat16 for better numerics.
+template <typename scalar_t>
+constexpr int elems_per_call = std::is_same_v<scalar_t, double> ? 2 : 4;
+
+// Box-Muller: convert 4 uniform uint32 values into 4 standard normal floats.
+__device__ __forceinline__ float4 box_muller_float(uint4 r) {
+  constexpr float M = 2.3283064365386963e-10f; // 1/2^32
+  constexpr float TWO_PI = 6.2831853071795864f;
+  // Map to (0, 1] to avoid log(0).
+  float u1 = fmaf(r.x, M, M * 0.5f);
+  float u2 = fmaf(r.y, M, M * 0.5f);
+  float u3 = fmaf(r.z, M, M * 0.5f);
+  float u4 = fmaf(r.w, M, M * 0.5f);
+
+  float radius1 = sqrtf(-2.0f * __logf(u1));
+  float radius2 = sqrtf(-2.0f * __logf(u3));
+  float s1, c1, s2, c2;
+  __sincosf(TWO_PI * u2, &s1, &c1);
+  __sincosf(TWO_PI * u4, &s2, &c2);
+  return {radius1 * c1, radius1 * s1, radius2 * c2, radius2 * s2};
+}
+
+// Box-Muller: convert 4 uint32 values (packed into 2 uint64) into 2 standard
+// normal doubles.
+__device__ __forceinline__ double2 box_muller_double(uint4 r) {
+  constexpr double M = 2.3283064365386963e-10; // 1/2^32
+  constexpr double TWO_PI = 6.2831853071795864;
+  // Pack pairs of uint32 for ~64 bits of uniform randomness.
+  double u1 = fma(static_cast<double>(r.x), M,
+                  static_cast<double>(r.y) * M * M + M * M * 0.5);
+  double u2 = fma(static_cast<double>(r.z), M,
+                  static_cast<double>(r.w) * M * M + M * M * 0.5);
+
+  double radius = ::sqrt(-2.0 * ::log(u1));
+  double s, c;
+  ::sincos(TWO_PI * u2, &s, &c);
+  return {radius * c, radius * s};
+}
+
+// Single-key kernel: one thread per chunk of elements, where each chunk
+// comes from a single Philox 4x32 call. Uses vectorized stores for full
+// chunks and scalar writes for the tail.
+template <typename scalar_t, typename sample_t, typename param_t>
+__global__ void philox_single_key_kernel(
+    scalar_t* __restrict__ output,
+    const uint64_t* __restrict__ key,
+    int64_t num_elems,
+    sample_t sample_func,
+    param_t param_func) {
+
+  // Use vectorized load to get (seed, offset)
+  auto key_vec = memory::ld_vec<16>(key);
+  auto* key_vals = reinterpret_cast<const uint64_t*>(&key_vec);
+  uint64_t seed = key_vals[0];
+  uint64_t offset = key_vals[1];
+
+  // Use vectorized stores for full chunks since they're aligned.
+  constexpr int epc = elems_per_call<scalar_t>;
+  int64_t num_full_chunks = num_elems / epc;
+  int64_t chunk = static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x;
+  if (chunk < num_full_chunks) {
+    auto sample = sample_func(seed, offset + static_cast<uint64_t>(chunk));
+    constexpr int vec_bytes = epc * sizeof(scalar_t);
+    memory::Vec<vec_bytes> v;
+    auto* vals = reinterpret_cast<scalar_t*>(&v);
+    #pragma unroll
+    for (int j = 0; j < epc; j++) {
+      vals[j] = param_func((&sample.x)[j]);
+    }
+    memory::st_vec<vec_bytes>(output + chunk * epc, v);
+  }
+
+  // Scalar tail for remaining elements.
+  if (chunk == num_full_chunks) {
+    int64_t tail_start = num_full_chunks * epc;
+    auto sample = sample_func(seed, offset + static_cast<uint64_t>(num_full_chunks));
+    for (int j = 0; j < num_elems - tail_start; j++) {
+      output[tail_start + j] = param_func((&sample.x)[j]);
+    }
+  }
+}
+
+// Multi-key kernel: one thread per (key_idx, chunk) pair, where each chunk
+// comes from a single Philox 4x32 call. Uses vectorized stores for full
+// chunks and scalar writes for the tail.
+template <typename scalar_t, typename sample_t, typename param_t>
+__global__ void philox_multi_key_kernel(
+    scalar_t* __restrict__ output,
+    const uint64_t* __restrict__ keys,
+    int64_t num_keys,
+    int64_t elems_per_key,
+    sample_t sample_func,
+    param_t param_func,
+    OffsetCalculator<1> key_offset_calc) {
+  constexpr int epc = elems_per_call<scalar_t>;
+  int64_t chunks_per_key = (elems_per_key + epc - 1) / epc;
+  int64_t total_threads = num_keys * chunks_per_key;
+  int64_t tid = static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x;
+  if (tid >= total_threads) return;
+
+  // Determine correct (seed, offset) to use and sample.
+  int64_t key_idx = tid / chunks_per_key;
+  int64_t chunk = tid % chunks_per_key;
+  auto elem_offset = key_offset_calc.get(key_idx)[0];
+  uint64_t seed = keys[elem_offset];
+  uint64_t offset = keys[elem_offset + 1];
+  auto sample = sample_func(seed, offset + static_cast<uint64_t>(chunk));
+
+  // Vectorized writes require aligned base addresses. This is guaranteed
+  // when elems_per_key is a multiple of epc, since
+  // base = key_idx * elems_per_key + chunk * epc.
+  int64_t full_chunks_per_key = elems_per_key / epc;
+  bool aligned = elems_per_key % epc == 0;
+  int64_t base = key_idx * elems_per_key + chunk * epc;
+  if (aligned && chunk < full_chunks_per_key) {
+    constexpr int vec_bytes = epc * sizeof(scalar_t);
+    memory::Vec<vec_bytes> v;
+    auto* vals = reinterpret_cast<scalar_t*>(&v);
+    #pragma unroll
+    for (int j = 0; j < epc; j++) {
+      vals[j] = param_func((&sample.x)[j]);
+    }
+    memory::st_vec<vec_bytes>(output + base, v);
+  } else {
+    for (int j = 0; j < epc && chunk * epc + j < elems_per_key; j++) {
+      output[base + j] = param_func((&sample.x)[j]);
+    }
+  }
+}
+
+// Dispatches to single-key or multi-key kernels as needed.
+template <typename scalar_t, typename sample_t, typename param_t>
+void philox_distribution_kernel(
+    const char* op_name,
+    Tensor& self, const Tensor& key,
+    const sample_t& sample_func, const param_t& param_func) {
+  TORCH_CHECK(self.is_floating_point(),
+      op_name, ": self must be a floating point tensor, got ",
+      self.scalar_type());
+  TORCH_CHECK(key.scalar_type() == kUInt64,
+      op_name, ": key must have dtype uint64, got ",
+      key.scalar_type());
+  TORCH_CHECK(self.device() == key.device(),
+      op_name, ": self and key must be on the same device, got ",
+      self.device(), " and ", key.device());
+  TORCH_CHECK(key.dim() >= 1 && key.size(-1) == 2,
+      op_name, ": key must have shape (2,) or (*batch, 2), got shape ",
+      key.sizes());
+  if (key.dim() > 1) {
+    TORCH_CHECK(key.dim() == self.dim() + 1,
+        op_name, ": batched key must have ndim == output ndim + 1, "
+        "got key shape ", key.sizes(), " with output shape ", self.sizes());
+    auto key_batch = key.sizes().slice(0, self.dim());
+    TORCH_CHECK(is_expandable_to(key_batch, self.sizes()),
+        op_name, ": key batch shape ", key_batch,
+        " is not broadcastable with output shape ", self.sizes());
+  }
+
+  if (self.numel() == 0) {
+    return;
+  }
+
+  // Ensure contiguous, aligned output for vectorized stores. Clone if needed
+  // to ensure alignment; the result is copied back into self afterwards.
+  constexpr int vec_bytes = elems_per_call<scalar_t> * sizeof(scalar_t);
+  auto output = self.contiguous();
+  if (reinterpret_cast<uintptr_t>(output.data_ptr()) % vec_bytes != 0) {
+    output = output.clone();
+  }
+
+  constexpr int block_size = 256;
+
+  if (key.dim() == 1) {
+    // === Launch single key kernel ===
+    constexpr int epc = elems_per_call<scalar_t>;
+    int64_t num_chunks = (self.numel() + epc - 1) / epc;
+    int num_blocks = static_cast<int>((num_chunks + block_size - 1) / block_size);
+
+    auto key_contig = key.contiguous();
+    philox_single_key_kernel<scalar_t>
+        <<<num_blocks, block_size, 0, at::cuda::getCurrentCUDAStream()>>>(
+        output.mutable_data_ptr<scalar_t>(),
+        key_contig.data_ptr<uint64_t>(),
+        self.numel(), sample_func, param_func);
+  } else {
+    // === Launch batched (multiple) key kernel ===
+    // The kernel writes each key's output as a contiguous block of
+    // elems_per_key elements. We determine elems_per_key by counting
+    // trailing size-1 key dims; these are the output dimensions that a
+    // single key generates over. For example, with key shape (4, 1, 1, 2)
+    // and output shape (4, 10, 100): key_dims=1, elems_per_key=1000.
+    int64_t elems_per_key = 1;
+    int64_t key_dims = self.dim();
+    for (int64_t i = self.dim() - 1; i >= 0; i--) {
+      if (key.size(i) != 1) break;
+      elems_per_key *= self.size(i);
+      key_dims--;
+    }
+    int64_t num_keys = self.numel() / elems_per_key;
+
+    // Handle key, self broadcasting via OffsetCalculator.
+    c10::SmallVector<int64_t, MAX_DIMS> oc_sizes(key_dims);
+    c10::SmallVector<int64_t, MAX_DIMS> oc_strides(key_dims);
+    for (int64_t i = 0; i < key_dims; i++) {
+      int64_t dim = key_dims - 1 - i;
+      oc_sizes[i] = self.size(dim);
+      oc_strides[i] = key.size(dim) > 1 ? key.stride(dim) : 0;
+    }
+    const int64_t* oc_strides_ptr = oc_strides.data();
+    auto key_offset_calc = OffsetCalculator<1>(
+        key_dims, oc_sizes.data(), &oc_strides_ptr);
+
+    int64_t chunks_per_key =
+        (elems_per_key + elems_per_call<scalar_t> - 1) / elems_per_call<scalar_t>;
+    int64_t total_threads = num_keys * chunks_per_key;
+    int num_blocks = static_cast<int>((total_threads + block_size - 1) / block_size);
+
+    philox_multi_key_kernel<scalar_t>
+        <<<num_blocks, block_size, 0, at::cuda::getCurrentCUDAStream()>>>(
+        output.mutable_data_ptr<scalar_t>(),
+        key.data_ptr<uint64_t>(),
+        num_keys, elems_per_key,
+        sample_func, param_func, key_offset_calc);
+  }
+  C10_CUDA_KERNEL_LAUNCH_CHECK();
+
+  if (output.data_ptr() != self.data_ptr()) {
+    self.copy_(output);
+  }
+}
+
+} // anonymous namespace
+
+Tensor& _philox_uniform_cuda_(
+    Tensor& self, const Tensor& key, double low, double high) {
+  AT_DISPATCH_FLOATING_TYPES_AND2(
+      kHalf, kBFloat16, self.scalar_type(), "_philox_uniform_", [&] {
+    auto sample_func = []() {
+      if constexpr (std::is_same_v<scalar_t, double>) {
+        return [] __device__ (uint64_t seed, uint64_t offset) {
+          uint4 r = philox_4x32(seed, offset);
+          ulonglong2 packed;
+          packed.x = (static_cast<unsigned long long>(r.x) << 32) | r.y;
+          packed.y = (static_cast<unsigned long long>(r.z) << 32) | r.w;
+          return packed;
+        };
+      } else {
+        return [] __device__ (uint64_t seed, uint64_t offset) {
+          return philox_4x32(seed, offset);
+        };
+      }
+    }();
+
+    auto lo = static_cast<scalar_t>(low);
+    auto hi = static_cast<scalar_t>(high);
+    auto param_func = [lo, hi] __device__ (auto rand) {
+      return static_cast<scalar_t>(
+          at::transformation::uniform_real(rand, lo, hi));
+    };
+
+    philox_distribution_kernel<scalar_t>(
+        "_philox_uniform_", self, key, sample_func, param_func);
+  });
+  return self;
+}
+
+Tensor& _philox_normal_cuda_(
+    Tensor& self, const Tensor& key, double mean, double stddev) {
+  AT_DISPATCH_FLOATING_TYPES_AND2(
+      kHalf, kBFloat16, self.scalar_type(), "_philox_normal_", [&] {
+    using compute_t = std::conditional_t<std::is_same_v<scalar_t, double>, double, float>;
+    auto sample_func = []() {
+      if constexpr (std::is_same_v<scalar_t, double>) {
+        return [] __device__ (uint64_t seed, uint64_t offset) {
+          return box_muller_double(philox_4x32(seed, offset));
+        };
+      } else {
+        return [] __device__ (uint64_t seed, uint64_t offset) {
+          return box_muller_float(philox_4x32(seed, offset));
+        };
+      }
+    }();
+
+    auto mu = static_cast<compute_t>(mean);
+    auto sigma = static_cast<compute_t>(stddev);
+    auto param_func = [mu, sigma] __device__ (compute_t rand) {
+      return static_cast<scalar_t>(rand * sigma + mu);
+    };
+
+    philox_distribution_kernel<scalar_t>(
+        "_philox_normal_", self, key, sample_func, param_func);
+  });
+  return self;
+}
+
+} // namespace at::native
diff --git a/aten/src/ATen/native/cuda/PhiloxKeySplit.cu b/aten/src/ATen/native/cuda/PhiloxKeySplit.cu
new file mode 100644
index 0000000000000..07893c061ee0d
--- /dev/null
+++ b/aten/src/ATen/native/cuda/PhiloxKeySplit.cu
@@ -0,0 +1,140 @@
+#define TORCH_ASSERT_ONLY_METHOD_OPERATORS
+
+#include <ATen/core/Tensor.h>
+#include <ATen/cuda/CUDAContext.h>
+#include <ATen/cuda/StatelessPhilox4x32.cuh>
+
+#ifndef AT_PER_OPERATOR_HEADERS
+#include <ATen/Functions.h>
+#include <ATen/NativeFunctions.h>
+#else
+#include <ATen/ops/_philox_key_fold_in_native.h>
+#include <ATen/ops/_philox_key_split_native.h>
+#include <ATen/ops/empty.h>
+#include <ATen/ops/empty_like.h>
+#endif
+
+namespace at::native {
+
+namespace {
+
+using at::cuda::philox_4x32;
+
+// Derive a new (seed, offset) key from 4 random uint32 values.
+// Use 2 uint32s for the 64-bit seed and 2 for the 64-bit offset.
+__device__ __forceinline__ void philox_derive_key(
+    uint4 r,
+    uint64_t* out_seed,
+    uint64_t* out_offset) {
+  *out_seed = static_cast<uint64_t>(r.x) | (static_cast<uint64_t>(r.y) << 32);
+  *out_offset = static_cast<uint64_t>(r.z) | (static_cast<uint64_t>(r.w) << 32);
+}
+
+// Grid-stride loop over (split_idx, key_idx) pairs.
+__global__ void philox_key_split_kernel(
+    const uint64_t* __restrict__ input,
+    uint64_t* __restrict__ output,
+    int64_t num_keys,
+    int64_t num_splits) {
+  int64_t total = num_keys * num_splits;
+  int64_t tid = static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x;
+  for (; tid < total; tid += static_cast<int64_t>(gridDim.x) * blockDim.x) {
+    int64_t split_idx = tid / num_keys;
+    int64_t key_idx = tid % num_keys;
+
+    uint64_t seed = input[key_idx * 2];
+    uint64_t offset = input[key_idx * 2 + 1];
+
+    // Sample randomness to get the next (seed, offset pair).
+    uint4 r = philox_4x32(seed, offset + static_cast<uint64_t>(split_idx));
+    int64_t out = (split_idx * num_keys + key_idx) * 2;
+    philox_derive_key(r, &output[out], &output[out + 1]);
+  }
+}
+
+__global__ void philox_key_fold_in_kernel(
+    const uint64_t* __restrict__ input,
+    uint64_t* __restrict__ output,
+    int64_t num_keys,
+    int64_t data) {
+  int64_t idx = static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x;
+  for (; idx < num_keys; idx += static_cast<int64_t>(gridDim.x) * blockDim.x) {
+    uint64_t seed = input[idx * 2];
+    uint64_t offset = input[idx * 2 + 1];
+
+    // Sample randomness to get the next (seed, offset pair).
+    uint4 r = philox_4x32(seed, offset + static_cast<uint64_t>(data));
+    philox_derive_key(r, &output[idx * 2], &output[idx * 2 + 1]);
+  }
+}
+
+} // anonymous namespace
+
+Tensor _philox_key_split_cuda(const Tensor& key, int64_t num_splits) {
+  TORCH_CHECK(key.dim() >= 1 && key.size(-1) == 2,
+      "_philox_key_split: key must have shape (*batch, 2), got shape ",
+      key.sizes());
+  TORCH_CHECK(key.scalar_type() == kUInt64,
+      "_philox_key_split: key must have dtype uint64, got ",
+      key.scalar_type());
+  TORCH_CHECK(num_splits > 0,
+      "_philox_key_split: num_splits must be positive, got ",
+      num_splits);
+
+  // Output shape: (num_splits, *key.shape)
+  auto output_sizes = key.sizes().vec();
+  output_sizes.insert(output_sizes.begin(), num_splits);
+  Tensor output = at::empty(output_sizes, key.options());
+  int64_t num_keys = key.numel() / 2;
+  if (num_keys == 0) {
+    return output;
+  }
+
+  int64_t total_threads = num_keys * num_splits;
+  constexpr int block_size = 256;
+  int num_blocks = std::min(
+      static_cast<int>((total_threads + block_size - 1) / block_size),
+      at::cuda::getCurrentDeviceProperties()->multiProcessorCount * 4);
+
+  auto key_contig = key.contiguous();
+  philox_key_split_kernel<<<num_blocks, block_size, 0,
+      at::cuda::getCurrentCUDAStream()>>>(
+      key_contig.data_ptr<uint64_t>(),
+      output.data_ptr<uint64_t>(),
+      num_keys, num_splits);
+  C10_CUDA_KERNEL_LAUNCH_CHECK();
+
+  return output;
+}
+
+Tensor _philox_key_fold_in_cuda(const Tensor& key, int64_t data) {
+  TORCH_CHECK(key.dim() >= 1 && key.size(-1) == 2,
+      "_philox_key_fold_in: key must have shape (*batch, 2), got shape ",
+      key.sizes());
+  TORCH_CHECK(key.scalar_type() == kUInt64,
+      "_philox_key_fold_in: key must have dtype uint64, got ",
+      key.scalar_type());
+
+  Tensor output = at::empty_like(key);
+  int64_t num_keys = key.numel() / 2;
+  if (num_keys == 0) {
+    return output;
+  }
+
+  constexpr int block_size = 256;
+  int num_blocks = std::min(
+      static_cast<int>((num_keys + block_size - 1) / block_size),
+      at::cuda::getCurrentDeviceProperties()->multiProcessorCount * 4);
+
+  auto key_contig = key.contiguous();
+  philox_key_fold_in_kernel<<<num_blocks, block_size, 0,
+      at::cuda::getCurrentCUDAStream()>>>(
+      key_contig.data_ptr<uint64_t>(),
+      output.data_ptr<uint64_t>(),
+      num_keys, data);
+  C10_CUDA_KERNEL_LAUNCH_CHECK();
+
+  return output;
+}
+
+} // namespace at::native
diff --git a/aten/src/ATen/native/cuda/Reduce.cuh b/aten/src/ATen/native/cuda/Reduce.cuh
index ee3a05c854540..71084941d5e1d 100644
--- a/aten/src/ATen/native/cuda/Reduce.cuh
+++ b/aten/src/ATen/native/cuda/Reduce.cuh
@@ -635,10 +635,10 @@ struct ReduceOp {
     using args_vec_t = std::array<arg_t, output_vec_size>;
     int dim_x = blockDim.x;
     args_vec_t* shared = (args_vec_t*)shared_memory;
-    if (dim_x > warpSize) {
+    if (dim_x > C10_WARP_SIZE) {
       int address_base = threadIdx.x + threadIdx.y*blockDim.x;
       shared[address_base] = value;
-      for (int offset = dim_x/2; offset >= warpSize; offset >>= 1) {
+      for (int offset = dim_x/2; offset >= C10_WARP_SIZE; offset >>= 1) {
         __syncthreads();
         if (threadIdx.x < offset && threadIdx.x + offset < blockDim.x) {
           args_vec_t other = shared[address_base + offset];
@@ -649,7 +649,7 @@ struct ReduceOp {
           shared[address_base] = value;
         }
       }
-      dim_x = warpSize;
+      dim_x = C10_WARP_SIZE;
     }
 
     __syncthreads();
@@ -836,9 +836,9 @@ struct ReduceOp {
           }
         }
       } else {
+#if defined(USE_ROCM) && ROCM_VERSION <= 71300
         index_t input_offset = threadIdx.y;
         index_t step = blockDim.y;
-#ifdef USE_ROCM // Prefetch loads to better hide their latency
         #define PRFCH 4
         for (; input_offset < config.ctas_per_output; input_offset += step*PRFCH) {
          arg_vec_t next[PRFCH];
@@ -855,6 +855,14 @@ struct ReduceOp {
          }
         }
 #else
+#if defined(USE_ROCM)
+        int input_offset = threadIdx.y;
+        int step = blockDim.y;
+        #pragma unroll
+#else
+        index_t input_offset = threadIdx.y;
+        index_t step = blockDim.y;
+#endif
         for (; input_offset < config.ctas_per_output; input_offset += step) {
           index_t idx = config.staging_memory_offset(input_offset);
           arg_vec_t next = reduce_buffer[idx];
@@ -1117,7 +1125,7 @@ ReduceConfig setReduceConfig(const TensorIterator& iter){
   if (iter.ndim() == 0 || reduction_on_fastest_striding_dimension) {
     // Split the input across lanes if the input is contiguous in the reduced
     // dimension. This will require reduction between threads using warp
-    // shuffle instructions and shared memory (if block_width > warpSize).
+    // shuffle instructions and shared memory (if block_width > C10_WARP_SIZE).
     config.input_mult[0] = config.split_input(block_width);
   } else {
     // Otherwise split the output across lanes in a warp.
diff --git a/aten/src/ATen/native/cuda/ReduceAMinMaxKernel.cu b/aten/src/ATen/native/cuda/ReduceAMinMaxKernel.cu
index cdd5daab2d983..dcc2fcfac607f 100644
--- a/aten/src/ATen/native/cuda/ReduceAMinMaxKernel.cu
+++ b/aten/src/ATen/native/cuda/ReduceAMinMaxKernel.cu
@@ -15,6 +15,8 @@
 #include <ATen/NumericUtils.h>
 #include <ATen/cuda/NumericLimits.cuh>
 
+#include <thrust/pair.h>
+
 namespace at::native {
 
 template <typename scalar_t>
diff --git a/aten/src/ATen/native/cuda/ReduceArgMaxKernel.cu b/aten/src/ATen/native/cuda/ReduceArgMaxKernel.cu
index d88ed2ac0f5cf..818f2403ec593 100644
--- a/aten/src/ATen/native/cuda/ReduceArgMaxKernel.cu
+++ b/aten/src/ATen/native/cuda/ReduceArgMaxKernel.cu
@@ -15,6 +15,8 @@
 #include <ATen/NumericUtils.h>
 #include <ATen/cuda/NumericLimits.cuh>
 
+#include <thrust/pair.h>
+
 namespace at::native {
 
 template <typename scalar_t, typename acc_t = scalar_t>
diff --git a/aten/src/ATen/native/cuda/ReduceArgMinKernel.cu b/aten/src/ATen/native/cuda/ReduceArgMinKernel.cu
index 44c44f6bf0300..ad7bee87694ed 100644
--- a/aten/src/ATen/native/cuda/ReduceArgMinKernel.cu
+++ b/aten/src/ATen/native/cuda/ReduceArgMinKernel.cu
@@ -15,6 +15,8 @@
 #include <ATen/NumericUtils.h>
 #include <ATen/cuda/NumericLimits.cuh>
 
+#include <thrust/pair.h>
+
 namespace at::native {
 
 template <typename scalar_t, typename acc_t = scalar_t>
diff --git a/aten/src/ATen/native/cuda/ReduceLogicKernel.cu b/aten/src/ATen/native/cuda/ReduceLogicKernel.cu
index 9e806b2aa7183..9724819fd681e 100644
--- a/aten/src/ATen/native/cuda/ReduceLogicKernel.cu
+++ b/aten/src/ATen/native/cuda/ReduceLogicKernel.cu
@@ -13,8 +13,8 @@ void and_kernel_cuda(TensorIterator& iter) {
       kHalf, kBFloat16, kBool, iter.common_dtype(), "and_cuda", [&]() {
         gpu_reduce_kernel<scalar_t, bool>(
             iter,
-            func_wrapper<bool>([] GPU_LAMBDA(scalar_t a, scalar_t b) -> bool {
-              return (static_cast<bool>(a) && static_cast<bool>(b));
+            func_wrapper<bool>([] GPU_LAMBDA(bool acc, scalar_t val) -> bool {
+              return (acc && static_cast<bool>(val));
             }),
             true);
       });
@@ -25,8 +25,8 @@ void or_kernel_cuda(TensorIterator& iter) {
       kHalf, kBFloat16, kBool, iter.common_dtype(), "or_cuda", [&]() {
         gpu_reduce_kernel<scalar_t, bool>(
             iter,
-            func_wrapper<bool>([] GPU_LAMBDA(scalar_t a, scalar_t b) -> bool {
-              return (static_cast<bool>(a) || static_cast<bool>(b));
+            func_wrapper<bool>([] GPU_LAMBDA(bool acc, scalar_t val) -> bool {
+              return (acc || static_cast<bool>(val));
             }),
             false);
       });
diff --git a/aten/src/ATen/native/cuda/ReduceMaxValuesKernel.cu b/aten/src/ATen/native/cuda/ReduceMaxValuesKernel.cu
index e8d1e88ebb3ec..8345e5dc77b8b 100644
--- a/aten/src/ATen/native/cuda/ReduceMaxValuesKernel.cu
+++ b/aten/src/ATen/native/cuda/ReduceMaxValuesKernel.cu
@@ -15,6 +15,8 @@
 #include <ATen/NumericUtils.h>
 #include <ATen/cuda/NumericLimits.cuh>
 
+#include <thrust/pair.h>
+
 namespace at::native {
 
 template <typename acc_t>
diff --git a/aten/src/ATen/native/cuda/ReduceMinValuesKernel.cu b/aten/src/ATen/native/cuda/ReduceMinValuesKernel.cu
index e01ca6c88ebc8..2ce0e0d20453c 100644
--- a/aten/src/ATen/native/cuda/ReduceMinValuesKernel.cu
+++ b/aten/src/ATen/native/cuda/ReduceMinValuesKernel.cu
@@ -15,6 +15,7 @@
 #include <ATen/NumericUtils.h>
 #include <ATen/cuda/NumericLimits.cuh>
 
+#include <thrust/pair.h>
 
 namespace at::native {
 
diff --git a/aten/src/ATen/native/cuda/ReduceMomentKernel.cu b/aten/src/ATen/native/cuda/ReduceMomentKernel.cu
index cabe86b313e94..df9e34ae7b38f 100644
--- a/aten/src/ATen/native/cuda/ReduceMomentKernel.cu
+++ b/aten/src/ATen/native/cuda/ReduceMomentKernel.cu
@@ -8,6 +8,8 @@
 #include <ATen/Dispatch.h>
 #include <ATen/native/ReduceOps.h>
 
+#include <thrust/pair.h>
+
 namespace at::native {
 
 template <typename scalar_t, typename out_t=scalar_t>
diff --git a/aten/src/ATen/native/cuda/RowwiseScaledMM.cu b/aten/src/ATen/native/cuda/RowwiseScaledMM.cu
index 967f522c2eb8f..4bafe771c7e61 100644
--- a/aten/src/ATen/native/cuda/RowwiseScaledMM.cu
+++ b/aten/src/ATen/native/cuda/RowwiseScaledMM.cu
@@ -198,7 +198,7 @@ void f8f8bf16_rowwise_impl(
           cutlass::epilogue::collective::EpilogueTileAuto,
           DtypeAccum,
           DtypeEpilogue,
-          DtypeOutput,
+          void, // Indicate there is no beta scaling to save register
           LayoutOutput,
           AlignmentOutput,
           DtypeOutput,
@@ -255,7 +255,7 @@ void f8f8bf16_rowwise_impl(
                            : nullptr},
          {{reinterpret_cast<DtypeScale*>(w_scale.data_ptr())},
           {{reinterpret_cast<DtypeScale*>(x_scale.data_ptr())}}}}},
-       reinterpret_cast<DtypeOutput*>(out.data_ptr()),
+       nullptr,
        stride_output,
        reinterpret_cast<DtypeOutput*>(out.data_ptr()),
        stride_output}};
@@ -288,7 +288,7 @@ void f8f8bf16_rowwise_impl(
   }
 
   // Initialize CUTLASS kernel with arguments and workspace pointer
-  status = gemm.initialize(arguments, workspace.data_ptr());
+  status = gemm.initialize(arguments, workspace.data_ptr(), at::cuda::getCurrentCUDAStream());
   if (status != cutlass::Status::kSuccess) {
     throw std::runtime_error("cutlass cannot initialize");
   }
@@ -390,7 +390,7 @@ void f8f8bf16_rowwise_impl_sm100_sm120(
       TileShape, ClusterShape,
       cutlass::epilogue::collective::EpilogueTileAuto,
       DtypeAccum, DtypeEpilogue,
-      DtypeOutput, LayoutOutput, AlignmentOutput,
+      void, LayoutOutput, AlignmentOutput,
       DtypeOutput, LayoutOutput, AlignmentOutput,
       EpilogueScheduleType,
       EpilogueEVT>::CollectiveOp;
@@ -448,7 +448,7 @@ void f8f8bf16_rowwise_impl_sm100_sm120(
                            : nullptr},
          {{reinterpret_cast<DtypeScale*>(w_scale.data_ptr())},
           {{reinterpret_cast<DtypeScale*>(x_scale.data_ptr())}}}}},
-       reinterpret_cast<DtypeOutput*>(out.data_ptr()),
+       nullptr,
        stride_output,
        reinterpret_cast<DtypeOutput*>(out.data_ptr()),
        stride_output}};
@@ -481,7 +481,7 @@ void f8f8bf16_rowwise_impl_sm100_sm120(
   }
 
   // Initialize CUTLASS kernel with arguments and workspace pointer
-  status = gemm.initialize(arguments, workspace.data_ptr());
+  status = gemm.initialize(arguments, workspace.data_ptr(), at::cuda::getCurrentCUDAStream());
   if (status != cutlass::Status::kSuccess) {
     throw std::runtime_error("cutlass cannot initialize");
   }
@@ -697,7 +697,7 @@ void f8f8bf16_rowwise_impl_sm89(
   }
 
   // Initialize CUTLASS kernel with arguments and workspace pointer
-  status = gemm.initialize(arguments, workspace.data_ptr());
+  status = gemm.initialize(arguments, workspace.data_ptr(), at::cuda::getCurrentCUDAStream());
   if (status != cutlass::Status::kSuccess) {
     throw std::runtime_error("cutlass cannot initialize");
   }
diff --git a/aten/src/ATen/native/cuda/ScaledBlas.cpp b/aten/src/ATen/native/cuda/ScaledBlas.cpp
index cafcc28b3d2c1..e1a1db85d6c1c 100644
--- a/aten/src/ATen/native/cuda/ScaledBlas.cpp
+++ b/aten/src/ATen/native/cuda/ScaledBlas.cpp
@@ -12,11 +12,12 @@
 #include <ATen/OpMathType.h>
 #include <ATen/TensorUtils.h>
 #include <ATen/cuda/CUDABlas.h>
-#include <ATen/cuda/CUDAScaledBlas.h>
+#include <ATen/native/ScaledBlasUtils.h>
 #include <ATen/cuda/tunable/Tunable.h>
 #include <ATen/cuda/tunable/TunableGemm.h>
 #include <ATen/native/Resize.h>
 #include <c10/util/MaybeOwned.h>
+#include <c10/util/StringUtil.h>
 #include <ATen/native/GroupedMMUtils.h>
 #include <ATen/native/cuda/RowwiseScaledMM.h>
 #include <ATen/native/cuda/ScaledGroupMM.h>
@@ -35,6 +36,7 @@
 #include <ATen/ops/_addmm_activation_native.h>
 #include <ATen/ops/_efficientzerotensor.h>
 #include <ATen/ops/_scaled_mm_native.h>
+#include <ATen/ops/_scaled_mm_v2_native.h>
 #include <ATen/ops/_unsafe_view_native.h>
 #include <ATen/ops/abs.h>
 #include <ATen/ops/addmm_native.h>
@@ -61,14 +63,15 @@ class cublasCommonArgs;
 using at::blas::ScalingType;
 using at::blas::SwizzleType;
 
-namespace scaled_blas = at::cuda::scaled;
+namespace scaled_blas = at::native::scaled;
 using scaled_blas::ScaledGemmImplementation;
 using scaled_blas::convert_int_to_enum;
-using scaled_blas::_scaled_mm_allowed_device;
 
 namespace at::native {
 
-static bool _scaled_mm_allowed_device(bool sm90_only=false, bool sm100_only=false) {
+namespace{
+
+bool _scaled_mm_allowed_device(bool sm90_only=false, bool sm100_only=false) {
 #ifdef USE_ROCM
     static const std::vector<std::string> archs = {
         "gfx942",
@@ -76,7 +79,10 @@ static bool _scaled_mm_allowed_device(bool sm90_only=false, bool sm100_only=fals
         "gfx1200", "gfx1201",
 #endif
 #if ROCM_VERSION >= 60500
-        "gfx950"
+        "gfx950",
+#endif
+#if ROCM_VERSION >= 71400
+        "gfx1250",
 #endif
     };
     return at::detail::getCUDAHooks().isGPUArch(archs);
@@ -92,12 +98,23 @@ static bool _scaled_mm_allowed_device(bool sm90_only=false, bool sm100_only=fals
 }
 
 #ifdef USE_ROCM
-static bool _scaled_mm_is_fnuz() {
+bool _scaled_mm_is_fnuz() {
     return at::detail::getCUDAHooks().isGPUArch({"gfx942"});
 }
-#endif
 
-namespace{
+#if ROCM_VERSION >= 70000
+static void check_blockwise_e8m0fnu_arch_supported() {
+  std::vector<std::string> mx_archs{"gfx950"};
+#if ROCM_VERSION >= 71400
+  mx_archs.push_back("gfx1250");
+#endif
+  TORCH_CHECK_NOT_IMPLEMENTED(
+      at::detail::getCUDAHooks().isGPUArch(mx_archs),
+      "Block-wise scaling for Float8_e8m0fnu is only supported on ",
+      c10::Join(",", mx_archs));
+}
+#endif
+#endif
 
 /*
  * Scaling Type Determination:
@@ -206,8 +223,7 @@ bool is_desired_scaling(const at::Tensor& t, const at::Tensor& scale, ScalingTyp
     case ScalingType::BlockWise128x128:
       return is_blockwise_128x128_scaling(t, scale);
     default:
-      TORCH_CHECK(false);
-      return false;
+      TORCH_CHECK(false, "Unknown scaling type");
   }
 }
 
@@ -428,7 +444,6 @@ _scaled_gemm(
   }
 }
 
-} // namespace
 
 // NOTE(slayton58): This is defined as part of the _v2 code (way) below - declare the signature here
 //                  to help cleanup v1 call structure.
@@ -441,6 +456,8 @@ _scaled_rowwise_rowwise(
           bool /*use_fast_accum*/,
           Tensor& /*out*/);
 
+} // namespace
+
 
 // Computes matrix multiply + bias while applying scaling to input and output matrices
 // Scales are only applicable when matrices are of Float8 type and assumed to be equal to 1.0 by default.
@@ -621,9 +638,8 @@ _scaled_mm_out_cuda(const Tensor& mat1, const Tensor& mat2,
   }
   else if (scaling_choice_a == ScalingType::BlockWise1x32 && scaling_choice_b == ScalingType::BlockWise1x32) {
 #ifdef USE_ROCM
-    #if ROCM_VERSION >= 70000
-    TORCH_CHECK_NOT_IMPLEMENTED(at::detail::getCUDAHooks().isGPUArch({"gfx950"}),
-                "Block-wise scaling for Float8_e8m0fnu is only supported on gfx950");
+#if ROCM_VERSION >= 70000
+    check_blockwise_e8m0fnu_arch_supported();
 
     int packed_factor = 1;
     if (mat1.scalar_type() == ScalarType::Float4_e2m1fn_x2) {
@@ -661,13 +677,11 @@ _scaled_mm_cuda(const Tensor& mat_a, const Tensor& mat_b,
   return _scaled_mm_out_cuda(mat_a, mat_b, scale_a, scale_b, bias, scale_result, out_dtype, use_fast_accum, out);
 }
 
+namespace {
+
 using acceptance_fn = std::function<bool(c10::ScalarType, std::vector<ScalingType>&, ArrayRef<Tensor>&, c10::ScalarType, std::vector<ScalingType>&, ArrayRef<Tensor>&)>;
 using namespace std::placeholders;
 
-namespace scaled_blas = at::cuda::scaled;
-using scaled_blas::ScaledGemmImplementation;
-using scaled_blas::convert_int_to_enum;
-
 std::array<std::tuple<std::string, acceptance_fn, ScaledGemmImplementation>, 9> scale_kernel_dispatch = {{
   { "tensorwise_tensorwise", scaled_blas::check_tensorwise_recipe, ScaledGemmImplementation::TENSORWISE_TENSORWISE },
   { "rowwise_rowwise", scaled_blas::check_rowwise_recipe, ScaledGemmImplementation::ROWWISE_ROWWISE},
@@ -829,7 +843,7 @@ _scaled_block1x128_block1x128(
     scale_a.stride(0) == 1 &&
     (
       scale_a.stride(1) == M ||
-      (scale_a.size(1) == 1 && scale_b.stride(1) == 1)
+      (scale_a.size(1) == 1 && scale_a.stride(1) == 1)
     ),
     "scale_a strides must be (", 1, ", ", M, "); got: ", scale_a.strides()
   );
@@ -851,7 +865,7 @@ _scaled_block1x128_block1x128(
         scale_b.stride(1) == 1
       )
     ),
-    "scale_b strides must be (", 1, ", ", N, "); got: ", scale_a.strides()
+    "scale_b strides must be (", 1, ", ", N, "); got: ", scale_b.strides()
   );
 
   auto scaling_choice_a = ScalingType::BlockWise1x128;
@@ -989,7 +1003,7 @@ _scaled_block1x128_block128x128(
         scale_a.stride(1) == 1
       )
     ),
-    "scale_a must have strides (1, ", M, "); got ", scale_b.strides()
+    "scale_a must have strides (1, ", M, "); got ", scale_a.strides()
   );
   // scale_b shape
   TORCH_CHECK_VALUE(
@@ -1066,8 +1080,7 @@ _scaled_mxfp8_mxfp8(
 
 #ifdef USE_ROCM
 #if ROCM_VERSION >= 70000
-  TORCH_CHECK_NOT_IMPLEMENTED(at::detail::getCUDAHooks().isGPUArch({"gfx950"}),
-              "Block-wise scaling for Float8_e8m0fnu is only supported on gfx950");
+  check_blockwise_e8m0fnu_arch_supported();
 
   TORCH_CHECK_VALUE(mat_a.size(0) % 32 == 0 && mat_a.size(1) % 32 == 0 &&
               mat_b.size(0) % 32 == 0 && mat_b.size(1) % 32 == 0,
@@ -1152,8 +1165,7 @@ _scaled_mxfp4_mxfp4(
   auto scaling_choice_b = ScalingType::BlockWise1x32;
 
 #if ROCM_VERSION >= 70000
-  TORCH_CHECK_NOT_IMPLEMENTED(at::detail::getCUDAHooks().isGPUArch({"gfx950"}),
-              "Block-wise scaling for Float8_e8m0fnu is only supported on gfx950");
+  check_blockwise_e8m0fnu_arch_supported();
 
   TORCH_CHECK_VALUE(mat_a.size(0) % 32 == 0 && mat_a.size(1) % 32 == 0 &&
               mat_b.size(0) % 32 == 0 && mat_b.size(1) % 32 == 0,
@@ -1228,6 +1240,67 @@ _scaled_nvfp4_nvfp4(
   return _scaled_gemm(mat_a, mat_b, scale_a, scale_b, scaling_choice_a, scaling_choice_b, bias, false /* use_fast_accum */, out, alpha);
 }
 
+void check_swizzle_lengths(ScaledGemmImplementation impl,
+                           std::vector<SwizzleType>& swizzle_a,
+                           std::vector<SwizzleType>& swizzle_b) {
+  // Store implementations that care about swizzling, and how many swizzle arguments
+  // they have to have
+  // NOTE(slayton): auto here is unable to deduce the correct type..
+  std::array<std::tuple<ScaledGemmImplementation, unsigned int>, 4> swizzled_impl = {{
+    // {implementation, # required arguments}
+    {ScaledGemmImplementation::MXFP8_MXFP8, 1},
+    {ScaledGemmImplementation::NVFP4_NVFP4, 2},
+    {ScaledGemmImplementation::NVFP4_NVFP4_SINGLE_SCALE, 1},
+    {ScaledGemmImplementation::MXFP4_MXFP4, 1}
+  }};
+
+  // Only check MX/NVFP formats on NVIDIA
+  for (auto [check_impl, num_args] : swizzled_impl) {
+    if (impl != check_impl) {
+      continue;
+    }
+#ifdef USE_ROCM
+    if (
+        check_impl != ScaledGemmImplementation::MXFP8_MXFP8 &&
+        check_impl != ScaledGemmImplementation::MXFP4_MXFP4) {
+      // ROCm currently does not support NVFP4 paths.
+      break;
+    }
+    TORCH_CHECK_VALUE(
+        swizzle_a.size() == 1 && swizzle_b.size() == 1,
+        "For ROCM MX gemm, swizzle_a and swizzle_b must each have 1 value, got ",
+        swizzle_a.size(),
+        " and ",
+        swizzle_b.size());
+    TORCH_CHECK_VALUE(
+        swizzle_a[0] == SwizzleType::NO_SWIZZLE &&
+            swizzle_b[0] == SwizzleType::NO_SWIZZLE,
+        "For ROCM MX gemm, swizzle_a and swizzle_b must both be NO_SWIZZLE");
+#else
+    TORCH_CHECK_VALUE(
+        swizzle_a.size() == num_args,
+        "swizzle_a must have ",
+        num_args,
+        " value",
+        num_args == 1 ? "" : "s",
+        ", got ",
+        swizzle_a.size());
+    TORCH_CHECK_VALUE(
+        swizzle_b.size() == num_args,
+        "swizzle_b must have ",
+        num_args,
+        " value",
+        num_args == 1 ? "" : "s",
+        ", got ",
+        swizzle_b.size());
+#endif
+
+    // No need to check anything else
+    break;
+  }
+}
+
+};  // anonymous namespace
 
 // V2: Computes matrix multiply + bias while applying scaling to input and output matrices
 // Scales are only applicable when matrices are of Float8 type and assumed to be equal to 1.0 by default.
@@ -1420,6 +1493,8 @@ _scaled_mm_cuda_v2_out(
 
   auto bias_ = bias.value_or(Tensor());
 
+  check_swizzle_lengths(gemm_impl, swizzle_a_enum, swizzle_b_enum);
+
   // dispatch to appropriate lower-level calls for error checking & execution
   if (gemm_impl == ScaledGemmImplementation::TENSORWISE_TENSORWISE) {
     return _scaled_tensorwise_tensorwise(mat_a, mat_b, scale_a[0], scale_b[0], bias, out_dtype_, use_fast_accum, out);
diff --git a/aten/src/ATen/native/cuda/ScaledGroupMM.cu b/aten/src/ATen/native/cuda/ScaledGroupMM.cu
index a276579b15fd2..69a1a27c76bb1 100644
--- a/aten/src/ATen/native/cuda/ScaledGroupMM.cu
+++ b/aten/src/ATen/native/cuda/ScaledGroupMM.cu
@@ -192,7 +192,7 @@ void f8f8bf16_grouped_gemm_impl_sm90(
           cutlass::epilogue::collective::EpilogueTileAuto,
           DtypeAccum,
           DtypeAccum,
-          DtypeOutput,
+          void, // Indicate there is no beta scaling to save register
           LayoutOutput*,
           AlignmentOutput,
           DtypeOutput,
@@ -390,7 +390,7 @@ void f8f8bf16_grouped_gemm_impl_sm90(
        (const DtypeB**)inputB_ptrs,
        stride_B},
       {{{{inputB_scale_ptrs}, {{inputA_scale_ptrs}, {}, {}}, {}}, {}},
-       (const DtypeOutput**)output_ptrs,
+       nullptr,
        stride_output,
        output_ptrs,
        stride_output}};
diff --git a/aten/src/ATen/native/cuda/Sort.cu b/aten/src/ATen/native/cuda/Sort.cu
index bae17a6a01236..a05ded2e2bb0f 100644
--- a/aten/src/ATen/native/cuda/Sort.cu
+++ b/aten/src/ATen/native/cuda/Sort.cu
@@ -107,7 +107,7 @@ struct SmallBitonicSort {
 // For small sorts (n <= 128) we use warpMergeSortKVInPlace which
 // sorts one slice per warp and potentially multiple slices in the
 // same block for improved occupancy with large batch sizes.
-template <int sort_size>
+template <int sort_size, int WARP_SIZE>
 struct WarpMergeSort {
 
   template <int A, typename K, typename V, typename IndexType>
@@ -128,7 +128,7 @@ struct WarpMergeSort {
     const auto min_grid = minimum_grid_for_occupancy(
         warpMergeSortKVInPlace<
             A, -1, sort_size, max_block_y,
-            K, V, LTOp<K, true>, IndexType>,
+            K, V, LTOp<K, true>, IndexType, WARP_SIZE>,
         block_x * max_block_y);
     const auto max_batch = std::max(IndexType{1}, keySlices / min_grid);
     const int block_y = std::min(IndexType(max_block_y), max_batch);
@@ -142,7 +142,9 @@ struct WarpMergeSort {
 
     if (descending) {
       const K invalid_key = at::numeric_limits<K>::lower_bound();
-      warpMergeSortKVInPlace<A, -1, sort_size, max_block_y>
+      warpMergeSortKVInPlace<
+          A, -1, sort_size, max_block_y,
+          K, V, GTOp<K, true>, IndexType, WARP_SIZE>
         <<<grid, block, 0, stream>>>(
           keyInfo,
           keySlices,
@@ -161,7 +163,9 @@ struct WarpMergeSort {
         }
         return at::numeric_limits<K>::upper_bound();
       }();
-      warpMergeSortKVInPlace<A, -1, sort_size, max_block_y>
+      warpMergeSortKVInPlace<
+          A, -1, sort_size, max_block_y,
+          K, V, LTOp<K, true>, IndexType, WARP_SIZE>
         <<<grid, block, 0, stream>>>(
           keyInfo,
           keySlices,
@@ -375,7 +379,16 @@ void sortKeyValueInplace(
     sortCommon(SmallBitonicSort{}, key, value, dim, descending);
 #if HAS_WARP_MERGE_SORT()
   } else if (sort_size <= 128) {
-    sortCommon(WarpMergeSort<128>{}, key, value, dim, descending);
+#ifdef USE_ROCM
+    if (at::cuda::warp_size() == 32) {
+      sortCommon(WarpMergeSort<128, 32>{}, key, value, dim, descending);
+    }
+    else {
+      sortCommon(WarpMergeSort<128, 64>{}, key, value, dim, descending);
+    }
+#else
+    sortCommon(WarpMergeSort<128, C10_WARP_SIZE>{}, key, value, dim, descending);
+#endif
 #endif
   } else {
     sortCommon(MediumRadixSort{}, key, value, dim, descending);
diff --git a/aten/src/ATen/native/cuda/SortStable.cu b/aten/src/ATen/native/cuda/SortStable.cu
index 8117eeeec558e..33f3189ffbe44 100644
--- a/aten/src/ATen/native/cuda/SortStable.cu
+++ b/aten/src/ATen/native/cuda/SortStable.cu
@@ -21,11 +21,11 @@ namespace {
 struct offset_t {
   int stride;
   int begin;
-  __device__ int operator[](int i) const {
+  __host__ __device__ int operator[](int i) const {
     return stride * (begin + i);
   }
 #if CCCL_VERSION >= 3001000
-  __device__ offset_t& operator+=(int i) {
+  __host__ __device__ offset_t& operator+=(int i) {
     begin += i;
     return *this;
   }
diff --git a/aten/src/ATen/native/cuda/SortUtils.cuh b/aten/src/ATen/native/cuda/SortUtils.cuh
index 8e424707c7ef6..8bcb0ae323cfc 100644
--- a/aten/src/ATen/native/cuda/SortUtils.cuh
+++ b/aten/src/ATen/native/cuda/SortUtils.cuh
@@ -165,15 +165,18 @@ bitonicSortKVInPlace(at::cuda::detail::TensorInfo<K, IndexType> keys,
 
 #if HAS_WARP_MERGE_SORT()
 
+// Note [warp merge sort WARP_SIZE template param]
+// ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+// warpMergeSortKVInPlace was written assuming C10_WARP_SIZE is a constexpr.
+// In torch/headeronly/macros/Macros.h, C10_WARP_SIZE is 32 for CUDA, and on
+// ROCm it will be 32 or 64 based on the current compile-time gfx target.
+// Ideally, warpSize should be used instead of C10_WARP_SIZE in device code, but
+// C10_WARP_SIZE within this kernel has been used as a template parameter for
+// some device functions. Therefore, a template param for WARP_SIZE was added.
 template <int KeyDims, int ValueDims, int sort_size, int max_block_dim_y,
-          typename K, typename V, typename Comparator, typename IndexType>
-#if !defined(USE_ROCM)
-// On CUDA, use explicit launch bounds for better occupancy
-C10_LAUNCH_BOUNDS_1(C10_WARP_SIZE * max_block_dim_y)
-#endif
-// Note: ROCm doesn't use launch bounds here because C10_WARP_SIZE is not
-// a true compile-time constant in device code (it's a constexpr function).
-// The compiler infers good launch bounds from the kernel code automatically.
+          typename K, typename V, typename Comparator, typename IndexType,
+          int WARP_SIZE>
+C10_LAUNCH_BOUNDS_1(WARP_SIZE * max_block_dim_y)
 __global__ void
 warpMergeSortKVInPlace(
     at::cuda::detail::TensorInfo<K, IndexType> keys,
@@ -207,17 +210,17 @@ warpMergeSortKVInPlace(
 
   namespace cub = ROCM_HIPCUB(at_cuda_detail::cub);
 
-  CUDA_KERNEL_ASSERT(blockDim.x == C10_WARP_SIZE);
+  CUDA_KERNEL_ASSERT(blockDim.x == WARP_SIZE);
   CUDA_KERNEL_ASSERT(blockDim.y <= max_block_dim_y);
-  constexpr int items_per_thread = sort_size / C10_WARP_SIZE;
+  constexpr int items_per_thread = sort_size / WARP_SIZE;
   static_assert(
-      items_per_thread * C10_WARP_SIZE == sort_size,
-      "sort_size must be a multiple of C10_WARP_SIZE");
+      items_per_thread * WARP_SIZE == sort_size,
+      "sort_size must be a multiple of WARP_SIZE template param");
 
 
   using LoadKeys = cub::WarpLoad<K, items_per_thread, cub::WARP_LOAD_TRANSPOSE>;
   using LoadValues = cub::WarpLoad<V, items_per_thread, cub::WARP_LOAD_TRANSPOSE>;
-  using Sort = cub::WarpMergeSort<K, items_per_thread, C10_WARP_SIZE, V>;
+  using Sort = cub::WarpMergeSort<K, items_per_thread, WARP_SIZE, V>;
   using StoreKeys = cub::WarpStore<K, items_per_thread, cub::WARP_STORE_TRANSPOSE>;
   using StoreValues = cub::WarpStore<V, items_per_thread, cub::WARP_STORE_TRANSPOSE>;
 
diff --git a/aten/src/ATen/native/cuda/Sorting.cu b/aten/src/ATen/native/cuda/Sorting.cu
index 1c4f06fe262e6..dcf535aa0f016 100644
--- a/aten/src/ATen/native/cuda/Sorting.cu
+++ b/aten/src/ATen/native/cuda/Sorting.cu
@@ -29,9 +29,16 @@ __global__ void gatherKthValue(
     index_t inputWithinSliceStride,
     cuda::detail::TensorInfo<scalar_t, index_t> kthValue,
     cuda::detail::TensorInfo<int64_t, index_t> indices) {
-  // Indices are limited to integer fp precision, so counts can fit in
-  // int32, regardless of index_t
-  __shared__ int smem[C10_WARP_SIZE]; // one per each warp, up to warp limit
+  // smem is used by radixSelect for radix bin counts. Type must be index_t to
+  // handle sliceSize > INT_MAX.
+#ifndef USE_ROCM
+  __shared__ index_t smem[C10_WARP_SIZE]; // one per each warp, up to warp limit
+#else
+  // Maximum shared memory size for radix select (used in countRadixAggregateCounts): NUM_BUFFERS * MAX_WARPS * RADIX_SIZE.
+  // HIP workgroups have at most 1024 threads. Warp size is at least 32 (can be 64 on some
+  // architectures), so we use 32 for safety: 2 buffers * (1024/32) warps * 4 radix bins = 256.
+  __shared__ index_t smem[256];
+#endif
 
   index_t slice = getLinearBlockId<index_t>();
   if (slice >= numInputSlices) {
@@ -106,9 +113,16 @@ __global__ void gatherMedian(
     index_t numInputSlices,
     index_t inputWithinSliceStride,
     bool ignore_nan) {
-  // Shared memory for the subroutine RadixSelect. Note that RadixSelect converts the
-  // floating point type to int with the same relative ordering.
-  __shared__ int smem[C10_WARP_SIZE]; // one per each warp, up to warp limit
+  // smem is used by radixSelect for radix bin counts. Type must be index_t to
+  // handle sliceSize > INT_MAX.
+#ifndef USE_ROCM
+  __shared__ index_t smem[C10_WARP_SIZE]; // one per each warp, up to warp limit
+#else
+  // Maximum shared memory size for radix select (used in countRadixAggregateCounts): NUM_BUFFERS * MAX_WARPS * RADIX_SIZE.
+  // HIP workgroups have at most 1024 threads. Warp size is at least 32 (can be 64 on some
+  // architectures), so we use 32 for safety: 2 buffers * (1024/32) warps * 4 radix bins = 256.
+  __shared__ index_t smem[256];
+#endif
 
   index_t slice = getLinearBlockId<index_t>();
   if (slice >= numInputSlices) {
diff --git a/aten/src/ATen/native/cuda/SortingRadixSelect.cuh b/aten/src/ATen/native/cuda/SortingRadixSelect.cuh
index 977fc76b295be..67c15e0dd3385 100644
--- a/aten/src/ATen/native/cuda/SortingRadixSelect.cuh
+++ b/aten/src/ATen/native/cuda/SortingRadixSelect.cuh
@@ -3,6 +3,7 @@
 #include <ATen/cuda/AsmUtils.cuh>
 #include <ATen/cuda/Atomic.cuh>
 #include <ATen/cuda/DeviceUtils.cuh>
+#include <type_traits>
 
 namespace at::native {
 
@@ -215,10 +216,17 @@ __device__ void countRadixUsingMask(
   }
 
   // Now, for each warp, sum values
+  // Note: uint64_t on Linux is unsigned long, but CUDA atomicAdd expects
+  // unsigned long long. We use reinterpret_cast for compatibility.
   if (at::cuda::getLaneId() == 0) {
 #pragma unroll
     for (uint32_t i = 0; i < RadixSize; ++i) {
-      gpuAtomicAddNoReturn(&smem[i], counts[i]);
+      if constexpr (std::is_same_v<CountType, uint64_t>) {
+        atomicAdd(reinterpret_cast<unsigned long long*>(smem) + i,
+                  static_cast<unsigned long long>(counts[i]));
+      } else {
+        atomicAdd(&smem[i], counts[i]);
+      }
     }
   }
 
@@ -318,6 +326,7 @@ template <
     typename CountType,
     int RadixSize,
     int RadixBits,
+    bool prefetch,
     typename DataAccessor>
 __device__ __forceinline__ void countRadixLoop(
     CountType counts[RadixSize], // counts[i] will be the number of matching
@@ -395,32 +404,38 @@ __device__ __forceinline__ void countRadixLoop(
 
   // phase 2: processing 1 element at an iteration.
 
-  // prefetching. This is specifically useful for global memory access.
-  scalar_t v = unroll_segment + threadIdx.x < loopBound
-      ? getData(unroll_segment + threadIdx.x)
-      : static_cast<scalar_t>(0);
-  // we pad loopbound to round_up(loopbound, warpSize) to make sure all threads
-  // in the warp participate in the ballot.
+  // prefetching pattern if prefetch is true.
+  // prefetching pattern is only useful for global memory access.
+  scalar_t v_curr;
+  if constexpr (prefetch) {
+    v_curr = unroll_segment + threadIdx.x < loopBound
+        ? getData(unroll_segment + threadIdx.x)
+        : static_cast<scalar_t>(0);
+  }
   for (index_t i = unroll_segment + threadIdx.x;
-       i < round_up(
-               static_cast<index_t>(loopBound), static_cast<index_t>(warpSize));
+       i < loopBound;
        i += blockDim.x) {
-    // prefetch the next element.
-    scalar_t v_next = i + blockDim.x < loopBound ? getData(i + blockDim.x)
-                                                 : static_cast<scalar_t>(0);
-
-    bool hasVal = false;
-    bitwise_t digitInRadix = static_cast<bitwise_t>(0);
-    if (i < loopBound) {
-      bitwise_t val = TopKTypeConfig<scalar_t>::convert(v);
-      // check if bit pattern matches the pattern we have already discovered for
-      // topk value v.
-      hasVal = ((val & desiredMask) == desired);
-      // get the bits [radixDigitPos, radixDigitPos+RADIX_BITS-1] of the value
-      // v.
-      digitInRadix = at::cuda::Bitfield<bitwise_t>::getBitfield(
-          val, radixDigitPos, RadixBits);
-    }
+        scalar_t v_local; // the current element.
+        scalar_t v_next; // the next element. Used for prefetching.
+
+        if constexpr (prefetch) {
+          // prefetch the next element.
+          v_local = v_curr;
+          v_next = i + blockDim.x < loopBound ? getData(i + blockDim.x)
+                                              : static_cast<scalar_t>(0);
+        }
+        else {
+          v_local = getData(i); // if no prefetching, just get the current element.
+        }
+
+        bitwise_t val = TopKTypeConfig<scalar_t>::convert(v_local);
+        // check if bit pattern matches the pattern we have already discovered for
+        // topk value v.
+        bool hasVal = ((val & desiredMask) == desired);
+        // get the bits [radixDigitPos, radixDigitPos+RADIX_BITS-1] of the value
+        // v.
+        bitwise_t digitInRadix = at::cuda::Bitfield<bitwise_t>::getBitfield(
+            val, radixDigitPos, RadixBits);
 
 // counting across the warp.
 #pragma unroll
@@ -432,7 +447,79 @@ __device__ __forceinline__ void countRadixLoop(
       counts[j] += __popcll(WARP_BALLOT(vote));
     }
 
-    v = v_next; // closing the prefetching loop.
+    if constexpr (prefetch) {
+      v_curr = v_next; // closing the prefetching loop.
+    }
+  }
+}
+
+// Aggregates radix matches across all warps and distributes results back to all threads.
+// Uses double-buffering via buffer_index (0 or 1) to alternate between two smem segments,
+// preventing race conditions between concurrent iterations. Since countRadixUsingMaskDataSmem
+// performs __syncthreads() internally, at most two loop iterations can be in flight
+// simultaneously, so two buffers are sufficient. buffer_index is toggled after each
+// countRadixUsingMaskDataSmem invocation.
+template <
+    typename CountType,
+    int RadixSize,
+    int RadixBits>
+__device__ __forceinline__ void countRadixAggregateCounts(
+    CountType counts[RadixSize], // counts[i] will be the number of matching
+                                 // elements ((val & desiredMask) == desired)
+                                 // that have the digits [radixDigitPos,
+                                 // radixDigitPos+RADIX_BITS-1] set to i.
+    CountType* smem, // shared memory for inter-warp reduction of counts.
+    int buffer_index){ // buffer index for smem.
+
+  // Maximum number of warps per workgroup. HIP workgroups have at most 1024 threads.
+  // Warp size is at least 32 (can be 64 on some architectures), so we use 32 for safety.
+  // This sizes shared memory buffers to accommodate all possible warps: 1024/32 = 32.
+  constexpr uint MAX_WARPS = 1024/C10_WARP_SIZE_LOWER_BOUND;
+  const int buffer_offset = buffer_index * MAX_WARPS * RadixSize; // offset of the buffer in smem.
+  const uint WARP_BITS = __builtin_ctz(C10_WARP_SIZE);
+
+  const uint num_warps = blockDim.x >> WARP_BITS;  // Actual number of warps in this block
+  const uint warp_id = threadIdx.x >> WARP_BITS; // = threadIdx.x / C10_WARP_SIZE
+  const int lane_id = at::cuda::getLaneId(); // = threadIdx.x % C10_WARP_SIZE
+
+  // Stage 1: Each warp's lane 0 stores its counts in smem.
+  // Layout after Stage 1: [warp0: all radix bins], [warp1: all radix bins], ...
+  // this layout starts from index buffer_offset.
+  if (lane_id == 0) {
+#pragma unroll
+    for (uint32_t i = 0; i < RadixSize; ++i) {
+      smem[
+            buffer_offset
+          + warp_id * RadixSize
+          + i
+          ] = counts[i];
+    }
+  }
+
+  __syncthreads(); // wait for all warps to finish storing their counts to smem.
+
+  // Stage 2: Warp0 performs reduction for all bins.
+  // Layout after Stage 2: [final radix0 sum], [final radix1 sum], ..., [final radix(RadixSize-1) sum]
+  // this layout starts from index buffer_offset.
+  if (warp_id == 0 && lane_id < RadixSize) {
+    CountType sum = 0;
+#pragma unroll
+    for (int w = 0; w < num_warps; ++w) {
+      sum += smem[
+                    buffer_offset
+                  + w * RadixSize
+                  + lane_id
+                  ];
+    }
+    smem[buffer_offset + lane_id] = sum;
+  }
+
+  __syncthreads(); // Wait for warp 0 to finish reduction.
+
+  // Stage 3: Each thread reads the final counts from smem.
+#pragma unroll
+  for (uint32_t i = 0; i < RadixSize; ++i) {
+    counts[i] = smem[buffer_offset + i];
   }
 }
 
@@ -457,6 +544,7 @@ __device__ void countRadixUsingMaskDataSmem(
                            // digits [radixDigitPos, radixDigitPos+RADIX_BITS-1]
                            // set to i in the warp.
     CountType* smem, // shared memory for inter-warp reduction of counts.
+    int buffer_index, // buffer index for smem.
     bitwise_t
         desired, // combined with desiredMask to filter relevant elements. An
                  // element is relevant if ((val & desiredMask) == desired).
@@ -479,21 +567,13 @@ __device__ void countRadixUsingMaskDataSmem(
     counts[i] = 0; // initialize counts to 0.
   }
 
-  // initialize smem to 0. This is for reduction of counts across all warps.
-  if (threadIdx.x < RadixSize) {
-    smem[threadIdx.x] = 0;
-  }
-
-  __syncthreads(); // wait for all threads in the block to finish initializing
-                   // smem.
-
   // count the distribution of the bits in the radix digit at `radixDigitPos` to
   // `radixDigitPos`+RADIX_BITS-1 for values that match the desired pattern
   // ((val & desiredMask) == desired). counts[] will hold the results for the
   // current warp.
   if (dataSmemSize >
       0) { // if shared memory is filled, use dataSmem as the input data.
-    countRadixLoop<scalar_t, bitwise_t, index_t, int, RadixSize, RadixBits>(
+    countRadixLoop<scalar_t, bitwise_t, index_t, CountType, RadixSize, RadixBits, /*prefetch =*/ false>(
         counts,
         desired,
         desiredMask,
@@ -501,7 +581,7 @@ __device__ void countRadixUsingMaskDataSmem(
         dataSmemSize,
         [&](index_t i) -> scalar_t { return dataSmem[i]; });
   } else { // if shared memory is not filled, fall back to global memory.
-    countRadixLoop<scalar_t, bitwise_t, index_t, int, RadixSize, RadixBits>(
+    countRadixLoop<scalar_t, bitwise_t, index_t, CountType, RadixSize, RadixBits, /*prefetch =*/ true>(
         counts,
         desired,
         desiredMask,
@@ -512,27 +592,11 @@ __device__ void countRadixUsingMaskDataSmem(
         });
   }
 
-  // accumulate the counts across all warps.
-  // sum for each warp is added to smem by thread 0 in the warp.
-  if (at::cuda::getLaneId() == 0) {
-#pragma unroll
-    for (uint32_t i = 0; i < RadixSize; ++i) {
-      gpuAtomicAddNoReturn(
-          &smem[i],
-          counts[i]); // thread0 in warp atomically adds the counts to smem.
-    }
-  }
-
-  __syncthreads(); // wait for all warps to finish adding their counts to smem.
-
-// each thread reads the final counts from smem.
-#pragma unroll
-  for (uint32_t i = 0; i < RadixSize; ++i) {
-    counts[i] = smem[i];
-  }
-
-  __syncthreads(); // wait for all threads in the block to finish reading the
-                   // counts.
+  // aggregate counts across all warps and distribute results back to all threads.
+  countRadixAggregateCounts<CountType, RadixSize, RadixBits>(
+    counts,
+    smem,
+    buffer_index);
 }
 
 // This is the main loop of the findPattern function that finds the unique value
@@ -627,6 +691,15 @@ __device__ scalar_t findPatternDataSmem(
     const scalar_t* dataSmem, // input data stored in shared memory.
     index_t dataSmemSize) { // input data size stored in shared memory.
 
+  // Ensure all threads have finished reading from smem before overwriting it.
+  // countRadixAggregateCounts Stage 3 reads from smem[buffer_offset + i];
+  // when buffer_offset == 0, those locations overlap with smem[0]/smem[1]
+  // written below. Warp 0 (which writes smem[0]/smem[1]) may get ahead of
+  // lagging warps still in Stage 3. Syncing here (rather than at the end of
+  // Stage 3) is cheaper because findPatternDataSmem is called at most once per
+  // radixSelect invocation, only when a unique element is found (count == 1).
+  __syncthreads();
+
   // initialize smem to 0.
   // smem[0] is a flag to indicate if a value has been found.
   // smem[1] is the found value.
@@ -740,21 +813,15 @@ __device__ __forceinline__ void fillDataSmem(
     scalar_t v = threadIdx.x < sliceSize
         ? doLdg(&data[threadIdx.x * withinSliceStride])
         : static_cast<scalar_t>(0);
-    // we pad sliceSize to round_up(sliceSize, warpSize) to make sure all
-    // threads in the warp participate in the ballot.
-    for (index_t i = threadIdx.x; i <
-         round_up(static_cast<index_t>(sliceSize),
-                  static_cast<index_t>(warpSize));
+
+    for (index_t i = threadIdx.x; i < sliceSize;
          i += blockDim.x) {
       scalar_t v_next = (i + blockDim.x) < sliceSize
           ? doLdg(&data[(i + blockDim.x) * withinSliceStride])
           : static_cast<scalar_t>(0);
 
-      bool match = false;
-      if (i < sliceSize) {
-        match =
-            ((TopKTypeConfig<scalar_t>::convert(v) & desiredMask) == desired);
-      }
+      bool match =
+          (TopKTypeConfig<scalar_t>::convert(v) & desiredMask) == desired;
 
       // Warp-level ballot
       uint64_t ballot = WARP_BALLOT(
@@ -813,11 +880,13 @@ __device__ void radixSelect(
     bool largest,
     index_t sliceSize,
     index_t withinSliceStride,
-    int* smem,
+    index_t* smem,
     scalar_t* topK) {
   // Per-thread buckets into which we accumulate digit counts in our
   // radix
-  int counts[RADIX_SIZE];
+  //
+  // counts must be index_t to safely handle sliceSize > INT_MAX.
+  index_t counts[RADIX_SIZE];
 
 #ifdef USE_ROCM
 
@@ -849,6 +918,14 @@ __device__ void radixSelect(
   __syncthreads(); // so the initialization is visible to all threads in the
                    // blocks.
 
+  // buffer index for smem. We use two segments of smem for inter-warp communication of counts.
+  // Given the counting operation in countRadixUsingMaskDataSmem performs __syncthreads() internally,
+  // we need to alternate between the at most two segments of smem to avoid race conditions.
+  // No more than two iterations of the loop will be "in flight" at any given time because
+  // of the __syncthreads() in countRadixUsingMaskDataSmem.
+  // buffer_index is either 0 or 1. It is toggled after each countRadixUsingMaskDataSmem invocation.
+  int buffer_index = 0;
+
 #endif
 
   // We only consider elements x such that (x & desiredMask) == desired
@@ -860,7 +937,7 @@ __device__ void radixSelect(
   // We are looking for the top kToFind-th element when iterating over
   // digits; this count gets reduced by elimination when counting
   // successive digits
-  int kToFind = k;
+  index_t kToFind = k;
 
   // We start at the most significant digit in our radix, scanning
   // through to the least significant digit
@@ -890,11 +967,12 @@ __device__ void radixSelect(
         scalar_t,
         bitwise_t,
         index_t,
-        int,
+        index_t,
         RADIX_SIZE,
         RADIX_BITS>(
         counts,
         smem,
+        buffer_index,
         desired,
         desiredMask,
         digitPos,
@@ -903,12 +981,15 @@ __device__ void radixSelect(
         data,
         dataSmem,
         dataSmemSize);
+
+    buffer_index ^= 1; // toggle buffer index.
+
 #else
     countRadixUsingMask<
         scalar_t,
         bitwise_t,
         index_t,
-        int,
+        index_t,
         RADIX_SIZE,
         RADIX_BITS>(
         counts,
@@ -921,7 +1002,7 @@ __device__ void radixSelect(
         data);
 
 #endif
-    auto found_unique = [&](int i, int count) -> bool {
+    auto found_unique = [&](int i, index_t count) -> bool {
       /* All threads have the same value in counts here, so all */
       /* threads will return from the function. */
       if (count == 1 && kToFind == 1) {
@@ -963,7 +1044,7 @@ __device__ void radixSelect(
       }
       return false;
     };
-    auto found_non_unique = [&](int i, int count) -> bool {
+    auto found_non_unique = [&](int i, index_t count) -> bool {
       if (count >= kToFind) {
         desired = at::cuda::Bitfield<bitwise_t>::setBitfield(
             desired, i, digitPos, RADIX_BITS);
@@ -999,7 +1080,7 @@ __device__ void radixSelect(
       // Process in descending order
 #pragma unroll
       for (int i = RADIX_SIZE - 1; i >= 0; --i) {
-        int count = counts[i];
+        index_t count = counts[i];
         if (found_unique(i, count)) {
           return;
         }
@@ -1011,7 +1092,7 @@ __device__ void radixSelect(
       // Process in ascending order
 #pragma unroll
       for (int i = 0; i < RADIX_SIZE; ++i) {
-        int count = counts[i];
+        index_t count = counts[i];
         if (found_unique(i, count)) {
           return;
         }
diff --git a/aten/src/ATen/native/cuda/SummaryOps.cu b/aten/src/ATen/native/cuda/SummaryOps.cu
index 00916903c3eb1..dd192bf91444e 100644
--- a/aten/src/ATen/native/cuda/SummaryOps.cu
+++ b/aten/src/ATen/native/cuda/SummaryOps.cu
@@ -326,8 +326,9 @@ Tensor _histc_cuda_template(
   bounds_t maxvalue = max;
 
   if (min == max && self.numel() > 0) {
-    minvalue = *self.min().cpu().const_data_ptr<input_t>();
-    maxvalue = *self.max().cpu().const_data_ptr<input_t>();
+    auto [min_tensor, max_tensor] = self.aminmax();
+    minvalue = min_tensor.item<input_t>();
+    maxvalue = max_tensor.item<input_t>();
   }
   if (minvalue == maxvalue) {
     minvalue = minvalue - 1;
diff --git a/aten/src/ATen/native/cuda/TensorTopK.cu b/aten/src/ATen/native/cuda/TensorTopK.cu
index e3025bace508e..2d21944e92ceb 100644
--- a/aten/src/ATen/native/cuda/TensorTopK.cu
+++ b/aten/src/ATen/native/cuda/TensorTopK.cu
@@ -53,9 +53,15 @@ __global__ void gatherTopK(at::cuda::detail::TensorInfo<const T, IndexType> inpu
                            at::cuda::detail::TensorInfo<int64_t, IndexType> indices,
                            IndexType indicesWithinSliceStride,
                            T* kthValues) {
-  // Indices are limited to integer fp precision, so counts can fit in
-  // int32, regardless of IndexType
-  __shared__ int smem[32]; // one per each warp, up to warp limit
+  // smem is used for:
+  // 1. radixSelect: radix bin counts (RADIX_SIZE=4 elements)
+  // 2. exclusiveBinaryPrefixScan: warp prefix sums (≤32 elements)
+  // 3. findPattern: flag and value (2 elements, cast to scalar_t*)
+  //
+  // Type must be IndexType to safely handle sliceSize > INT_MAX.
+  // In radix selection, counts can exceed INT_MAX when billions of
+  // elements fall into a single radix bin.
+  __shared__ IndexType smem[32];
   IndexType slice = getLinearBlockId<IndexType>();
   if (slice >= numInputSlices) {
     return;
@@ -115,13 +121,13 @@ __global__ void gatherTopK(at::cuda::detail::TensorInfo<const T, IndexType> inpu
       hasTopK = inRange && (convertedV < topKConverted);
     }
 
-    int index;
-    int carry;
-    at::cuda::exclusiveBinaryPrefixScan<int, true>(
-        smem, hasTopK, &index, &carry, AddOp<int>());
+    IndexType index;
+    IndexType carry;
+    at::cuda::exclusiveBinaryPrefixScan<IndexType, true>(
+        smem, hasTopK, &index, &carry, AddOp<IndexType>());
 
     if (hasTopK) {
-      int writeIndex = writeIndexStart + index;
+      IndexType writeIndex = writeIndexStart + index;
       CUDA_KERNEL_ASSERT(writeIndex < outputSliceSize);
 
       IndexType topKOffset = writeIndex * topKWithinSliceStride;
@@ -149,13 +155,13 @@ __global__ void gatherTopK(at::cuda::detail::TensorInfo<const T, IndexType> inpu
     const auto convertedV = at::native::TopKTypeConfig<T>::convert(v);
     bool hasTopK = inRange && (convertedV == topKConverted);
 
-    int index;
-    int carry;
-    at::cuda::exclusiveBinaryPrefixScan<int, true>(
-        smem, hasTopK, &index, &carry, AddOp<int>());
+    IndexType index;
+    IndexType carry;
+    at::cuda::exclusiveBinaryPrefixScan<IndexType, true>(
+        smem, hasTopK, &index, &carry, AddOp<IndexType>());
 
     if (hasTopK && index < topKRemaining) {
-      int writeIndex = writeIndexStart + index;
+      IndexType writeIndex = writeIndexStart + index;
       CUDA_KERNEL_ASSERT(writeIndex < outputSliceSize);
 
       IndexType topKOffset = writeIndex * topKWithinSliceStride;
@@ -200,9 +206,10 @@ smaller) than the k-th element but phase 2 tops off the output as long as there
 // start_index: index to write the result to. (output of function)
 // my_offset: offset to write the result to. (output of function)
 // warp_count: number of threads that have values to add to the output. (output of function)
+template <typename IndexType>
 __device__ __forceinline__ void reserveWarpSpace(bool hasTopK,
-                                                int& writeIndexStart,
-                                                int& start_index,
+                                                IndexType& writeIndexStart,
+                                                IndexType& start_index,
                                                 int& my_offset,
                                                 int& warp_count) {
   auto ballot = WARP_BALLOT(hasTopK); // a bitmask of threads that have hasTopK == true within the warp.
@@ -213,7 +220,7 @@ __device__ __forceinline__ void reserveWarpSpace(bool hasTopK,
   // if > 0 threads have hasTopK == true within the warp,
   // reserve space for them by incrementing writeIndexStart atomically + saving the old value  as start index.
   if (warp_count > 0 && lane_id == 0) {
-    start_index = atomicAdd(&writeIndexStart, warp_count);
+    start_index = atomicAdd(&writeIndexStart, (IndexType)warp_count);
   }
   start_index = __shfl(start_index, 0); // broadcast the start index to all threads in the warp.
 
@@ -228,7 +235,7 @@ __device__ __forceinline__ void writeResult(T* topKSliceStart,
                                             IndexType topKWithinSliceStride,
                                             IndexType indicesWithinSliceStride,
                                             IndexType outputSliceSize,
-                                            int writeIndex,
+                                            IndexType writeIndex,
                                             T v,
                                             IndexType i){
   CUDA_KERNEL_ASSERT(writeIndex < outputSliceSize); // assert that the write index is within the output slice size.
@@ -255,10 +262,15 @@ __global__ void gatherTopK(at::cuda::detail::TensorInfo<const T, IndexType> inpu
                             IndexType indicesWithinSliceStride,
                             T* kthValues) {
 
-  // Indices are limited to integer fp precision, so counts can fit in
-  // int32, regardless of IndexType
-  __shared__ int smem[64];
-  __shared__ int writeIndexStart; // index to track where to write results. This is shared by all threads in the block. Increases atomically.
+  // smem and counts must use IndexType to safely handle sliceSize > INT_MAX.
+  // In radix selection, counts tracks elements matching a radix pattern,
+  // which can exceed INT_MAX when billions of elements fall into one bin.
+
+  // Maximum shared memory size for radix select (used in countRadixAggregateCounts): NUM_BUFFERS * MAX_WARPS * RADIX_SIZE.
+  // HIP workgroups have at most 1024 threads. Warp size is at least 32 (can be 64 on some
+  // architectures), so we use 32 for safety: 2 buffers * (1024/32) warps * 4 radix bins = 256.
+  __shared__ IndexType smem[256];
+  __shared__ IndexType writeIndexStart; // index to track where to write results. This is shared by all threads in the block. Increases atomically.
 
   IndexType slice = getLinearBlockId<IndexType>();
   if (slice >= numInputSlices) {
@@ -307,7 +319,7 @@ __global__ void gatherTopK(at::cuda::detail::TensorInfo<const T, IndexType> inpu
   }
   __syncthreads();
   // All threads within the warp need to participate in the loop, so rounding up to a multiple of the warp size.
-  IndexType numIterations = round_up(inputSliceSize, (IndexType) warpSize);
+  IndexType numIterations = round_up(inputSliceSize, (IndexType) C10_WARP_SIZE);
 
   // phase 1: write actual > `pattern` (or < `pattern`, depending on the sort direction) values to the output.
   // prefetching data from global memory.
@@ -321,7 +333,8 @@ __global__ void gatherTopK(at::cuda::detail::TensorInfo<const T, IndexType> inpu
       hasTopK = (largest) ? (convertedV > topKConverted) : (convertedV < topKConverted);
     }
 
-    int start_index, my_offset, warp_count;
+    IndexType start_index;
+    int my_offset, warp_count;
     reserveWarpSpace(hasTopK, writeIndexStart, start_index, my_offset, warp_count);
 
     // now warp has reserved space for itself. If hasTopK == true, we need to find the index to write the result to.
@@ -360,7 +373,8 @@ __global__ void gatherTopK(at::cuda::detail::TensorInfo<const T, IndexType> inpu
       hasTopK = convertedV == topKConverted;
     }
 
-    int start_index, my_offset, warp_count;
+    IndexType start_index;
+    int my_offset, warp_count;
     reserveWarpSpace(hasTopK, writeIndexStart, start_index, my_offset, warp_count);
 
     if ((warp_count > 0) && (outputSliceSize <= start_index)){
@@ -368,7 +382,7 @@ __global__ void gatherTopK(at::cuda::detail::TensorInfo<const T, IndexType> inpu
     }
 
     if (hasTopK){
-      int slots_available = outputSliceSize - start_index;
+      IndexType slots_available = outputSliceSize - start_index;
       if (my_offset < slots_available){
         writeResult(topKSliceStart,
           indicesSliceStart,
@@ -433,8 +447,9 @@ constexpr int MAX_WARP_TOPK_SLICE = 512;
 // GTOp/LTOp instead of bitwise conversion. Bitwise conversion is only needed for radix sorting.
 
 // Kernel using WarpMergeSort for small topK operations
+// See Note [warp merge sort WARP_SIZE template param]
 template <int KeyDims, int ValueDims, int sort_size, int max_block_dim_y,
-          typename scalar_t, typename IndexType, bool is_descending>
+          typename scalar_t, typename IndexType, bool is_descending, int WARP_SIZE>
 __global__ void warpMergeSortTopK(
     at::cuda::detail::TensorInfo<const scalar_t, IndexType> input,
     IndexType inputSliceSize,
@@ -472,14 +487,14 @@ __global__ void warpMergeSortTopK(
 
   namespace cub = ROCM_HIPCUB(at_cuda_detail::cub);
 
-  CUDA_KERNEL_ASSERT(blockDim.x == C10_WARP_SIZE);
+  CUDA_KERNEL_ASSERT(blockDim.x == WARP_SIZE);
   CUDA_KERNEL_ASSERT(blockDim.y <= max_block_dim_y);
-  constexpr int items_per_thread = sort_size / C10_WARP_SIZE;
-  static_assert(items_per_thread * C10_WARP_SIZE == sort_size,
-                "sort_size must be a multiple of C10_WARP_SIZE");
+  constexpr int items_per_thread = sort_size / WARP_SIZE;
+  static_assert(items_per_thread * WARP_SIZE == sort_size,
+                "sort_size must be a multiple of WARP_SIZE template param");
 
   using LoadKeys = cub::WarpLoad<scalar_t, items_per_thread, cub::WARP_LOAD_TRANSPOSE>;
-  using Sort = cub::WarpMergeSort<scalar_t, items_per_thread, C10_WARP_SIZE, int64_t>;
+  using Sort = cub::WarpMergeSort<scalar_t, items_per_thread, WARP_SIZE, int64_t>;
   using StoreKeys = cub::WarpStore<scalar_t, items_per_thread, cub::WARP_STORE_TRANSPOSE>;
   using StoreIndices = cub::WarpStore<int64_t, items_per_thread, cub::WARP_STORE_TRANSPOSE>;
 
@@ -555,12 +570,30 @@ void launch(
                         "Too many slices for warp topk");
 
   // Dispatch based on sort size and sort direction
+  // See Note [warp merge sort WARP_SIZE template param]
+#ifdef USE_ROCM
   #define LAUNCH_KERNEL(SORT_SIZE, IS_DESCENDING) \
-    warpMergeSortTopK<Dim, Dim, SORT_SIZE, max_block_dim_y, scalar_t, IndexType, IS_DESCENDING> \
+  if (at::cuda::warp_size() == 32) { \
+    warpMergeSortTopK<Dim, Dim, SORT_SIZE, max_block_dim_y, scalar_t, IndexType, IS_DESCENDING, 32> \
+      <<<grid, block, 0, stream>>>( \
+          input, inputSliceSize, k, numInputSlices, inputWithinSliceStride, \
+          topK, topKWithinSliceStride, indices, indicesWithinSliceStride); \
+  } \
+  else { \
+    warpMergeSortTopK<Dim, Dim, SORT_SIZE, max_block_dim_y, scalar_t, IndexType, IS_DESCENDING, 64> \
+      <<<grid, block, 0, stream>>>( \
+          input, inputSliceSize, k, numInputSlices, inputWithinSliceStride, \
+          topK, topKWithinSliceStride, indices, indicesWithinSliceStride); \
+  } \
+  C10_CUDA_KERNEL_LAUNCH_CHECK()
+#else
+  #define LAUNCH_KERNEL(SORT_SIZE, IS_DESCENDING) \
+    warpMergeSortTopK<Dim, Dim, SORT_SIZE, max_block_dim_y, scalar_t, IndexType, IS_DESCENDING, 32> \
       <<<grid, block, 0, stream>>>( \
           input, inputSliceSize, k, numInputSlices, inputWithinSliceStride, \
           topK, topKWithinSliceStride, indices, indicesWithinSliceStride); \
     C10_CUDA_KERNEL_LAUNCH_CHECK()
+#endif
 
   // We have specialized launches for different sizes, as sort_size affects
   // shared memory, registers per thread and occupancy. We can use 'LAUNCH_KERNEL(512, false);'
@@ -835,8 +868,9 @@ __global__ void computeBlockwiseWithinKCounts(
     }
   }
 
-  constexpr int num_warps = RADIX_DIGITS / C10_WARP_SIZE;
-  __shared__ uint32_t warp_counts[num_warps];
+  constexpr int SHMEM_SIZE = RADIX_DIGITS / C10_WARP_SIZE_LOWER_BOUND;  // max shmem size on ROCm
+  const int num_warps = RADIX_DIGITS / C10_WARP_SIZE;
+  __shared__ uint32_t warp_counts[SHMEM_SIZE];
   if (tidx % C10_WARP_SIZE == 0) {
     warp_counts[warp] = count;
   }
@@ -1009,8 +1043,9 @@ __global__ void computeBlockwiseWithinKCounts(
     }
   }
 
-  constexpr int num_warps = RADIX_DIGITS / C10_WARP_SIZE;
-  __shared__ uint32_t warp_counts[num_warps];
+  constexpr int SHMEM_SIZE = RADIX_DIGITS / C10_WARP_SIZE_LOWER_BOUND; // max shmem size on ROCm
+  const int num_warps = RADIX_DIGITS / C10_WARP_SIZE;
+  __shared__ uint32_t warp_counts[SHMEM_SIZE];
   if (tidx % C10_WARP_SIZE == 0) {
     warp_counts[warp] = count;
   }
diff --git a/aten/src/ATen/native/cuda/TriangularOps.cu b/aten/src/ATen/native/cuda/TriangularOps.cu
index dd8084b48ebf4..024a55b56d3dc 100644
--- a/aten/src/ATen/native/cuda/TriangularOps.cu
+++ b/aten/src/ATen/native/cuda/TriangularOps.cu
@@ -44,67 +44,138 @@ __global__ void triu_tril_kernel(
     const int64_t k,
     const int64_t N_padded,
     const IndexType last_dim_padded) {
+#if !defined(USE_ROCM)
   int64_t linear_idx = (((int64_t)blockIdx.x) * blockDim.x + threadIdx.x) * elements_per_thread;
   if (linear_idx >= N_padded) {
     return;
   }
+#else
+  // ROCm limits the total number of threads at 2^32 - use a strided loop to stay within this limit
+  for (int64_t linear_idx_retained = (((int64_t) blockIdx.x) * blockDim.x + threadIdx.x) * elements_per_thread;
+       linear_idx_retained < N_padded;
+       linear_idx_retained += blockDim.x * gridDim.x * elements_per_thread)
+  {
+    int64_t linear_idx { linear_idx_retained }; // linear_idx_retained persists for the next iteration
+#endif // !defined(USE_ROCM)
+
+    auto dims = self_info.dims;
+
+    // Compute column index and row index
+    IndexType col = linear_idx % last_dim_padded;
+    linear_idx /= last_dim_padded;
+    IndexType row = linear_idx % self_info.sizes[dims - 2];
+
+    if constexpr (inplace) {
+      bool mask_all_true = upper ? (col - row >= k) : (col + elements_per_thread - row <= k);
+      if (mask_all_true)
+#if !defined(USE_ROCM)
+        return;
+#else
+        // The strided loop must proceed on ROCm
+        continue;
+#endif
+    }
 
-  auto dims = self_info.dims;
-
-  // Compute column index amd row index
-  IndexType col = linear_idx % last_dim_padded;
-  linear_idx /= last_dim_padded;
-  IndexType row = linear_idx % self_info.sizes[dims - 2];
-
-  if constexpr (inplace) {
-    bool mask_all_true = upper ? (col - row >= k) : (col + elements_per_thread - row <= k);
-    if (mask_all_true)
-      return;
-  }
-
-  // Compute offset
-  IndexType self_offset = 0, result_offset = 0;
-  self_offset += self_info.strides[dims - 1] * col;
-  result_offset += result_info.strides[dims - 1] * col;
-  linear_idx /= self_info.sizes[dims - 2];
-  self_offset += self_info.strides[dims - 2] * row;
-  result_offset += result_info.strides[dims - 2] * row;
-
-  // Compute remaining offsets
-  IndexType running_index;
-  #pragma unroll
-  for (IndexType i = dims - 3; i >= 0; --i) {
-    running_index = linear_idx % self_info.sizes[i];
-    linear_idx /= self_info.sizes[i];
-    self_offset += running_index * self_info.strides[i];
-    result_offset += running_index * result_info.strides[i];
-  }
+    // Compute offset
+    IndexType self_offset = 0, result_offset = 0;
+    self_offset += self_info.strides[dims - 1] * col;
+    result_offset += result_info.strides[dims - 1] * col;
+    linear_idx /= self_info.sizes[dims - 2];
+    self_offset += self_info.strides[dims - 2] * row;
+    result_offset += result_info.strides[dims - 2] * row;
 
-  if constexpr (inplace) {
+    // Compute remaining offsets
+    IndexType running_index;
     #pragma unroll
-    for (int i = 0; i < elements_per_thread && col + i < self_info.sizes[dims - 1]; i++) {
-      bool mask = upper ? (col + i - row >= k) : (col + i - row <= k);
-      if (!mask)
-        result_info.data[result_offset + i * result_info.strides[dims - 1]] = scalar_t(0);
+    for (IndexType i = dims - 3; i >= 0; --i) {
+      running_index = linear_idx % self_info.sizes[i];
+      linear_idx /= self_info.sizes[i];
+      self_offset += running_index * self_info.strides[i];
+      result_offset += running_index * result_info.strides[i];
     }
-  } else {
-    scalar_t frag[elements_per_thread] = {};
-    bool has_mask = (upper && col + elements_per_thread - row >= k) || (!upper && col - row <= k);
-    if (has_mask) {
-      #pragma unroll
-      for (int i = 0; i < elements_per_thread && col + i < self_info.sizes[dims - 1]; i++)
-        frag[i] = self_info.data[self_offset + i * self_info.strides[dims - 1]];
 
+    if constexpr (inplace) {
       #pragma unroll
-      for (int i = 0; i < elements_per_thread; i++) {
+      for (int i = 0; i < elements_per_thread && col + i < self_info.sizes[dims - 1]; i++) {
         bool mask = upper ? (col + i - row >= k) : (col + i - row <= k);
-        frag[i] = mask ? frag[i] : scalar_t(0);
+        if (!mask)
+          result_info.data[result_offset + i * result_info.strides[dims - 1]] = scalar_t(0);
+      }
+    } else {
+      scalar_t frag[elements_per_thread] = {};
+      bool has_mask = (upper && col + elements_per_thread - row >= k) || (!upper && col - row <= k);
+      if (has_mask) {
+        #pragma unroll
+        for (int i = 0; i < elements_per_thread && col + i < self_info.sizes[dims - 1]; i++)
+          frag[i] = self_info.data[self_offset + i * self_info.strides[dims - 1]];
+
+        #pragma unroll
+        for (int i = 0; i < elements_per_thread; i++) {
+          bool mask = upper ? (col + i - row >= k) : (col + i - row <= k);
+          frag[i] = mask ? frag[i] : scalar_t(0);
+        }
       }
+
+      #pragma unroll
+      for (int i = 0; i < elements_per_thread && col + i < self_info.sizes[dims - 1]; i++)
+        result_info.data[result_offset + i * result_info.strides[dims - 1]] = frag[i];
     }
+#if defined(USE_ROCM)
+  } // end of the strided loop on ROCm
+#endif // !defined(USE_ROCM)
+}
 
-    #pragma unroll
-    for (int i = 0; i < elements_per_thread && col + i < self_info.sizes[dims - 1]; i++)
-      result_info.data[result_offset + i * result_info.strides[dims - 1]] = frag[i];
+// Helper extracted from the AT_DISPATCH lambda below. The USE_ROCM
+// preprocessor branches must live outside the AT_DISPATCH macro expansion;
+// MSVC's traditional preprocessor does not accept #if/#else/#endif inside a
+// macro argument and the Windows build would otherwise fail.
+template <bool upper, typename scalar_t>
+void launch_triu_tril_kernel(const Tensor& result, const Tensor& self, int64_t k) {
+#if !defined(USE_ROCM)
+  constexpr int elements_per_thread = sizeof(scalar_t) < 8 ? 8 / sizeof(scalar_t) : 1;
+#else
+  // Pair smaller scalars with more elements per thread so the bounded grid
+  // (see dim_grid below) still covers the full output via the kernel's
+  // strided loop.
+  constexpr int elements_per_thread =
+    sizeof(scalar_t) <= 2 ? 4 :
+    sizeof(scalar_t) <= 8 ? 2 : 1;
+#endif
+  auto sizes = self.sizes();
+  int64_t last_dim_padded = round_up<int64_t>(sizes.back(), elements_per_thread);
+  int64_t N_padded = c10::multiply_integers(sizes.begin(), sizes.end() - 1) * last_dim_padded;
+  dim3 dim_block = block_size;
+
+#if !defined(USE_ROCM)
+  dim3 dim_grid((N_padded / elements_per_thread + dim_block.x - 1) / dim_block.x);
+#else
+  // ROCm caps total threads (grid * block) at 2^32; bound the grid here and
+  // let the kernel walk the remainder via a strided loop so large matrices
+  // with 64-bit indexing produce correct results instead of silently
+  // truncating.
+  const int num_mp = at::cuda::getCurrentDeviceProperties()->multiProcessorCount;
+  constexpr int grid_multiplier = sizeof(scalar_t) <= 2 ? 16 : 32;
+  dim3 dim_grid(num_mp * grid_multiplier);
+#endif
+
+  if (cuda::detail::canUse32BitIndexMath(result) && cuda::detail::canUse32BitIndexMath(self)) {
+    auto result_info = cuda::detail::getTensorInfo<scalar_t, int32_t>(result);
+    auto self_info = cuda::detail::getTensorInfo<const scalar_t, int32_t>(self);
+    BOOL_SWITCH(self.is_same(result), inplace, [&] {
+      triu_tril_kernel<scalar_t, int32_t, upper, elements_per_thread, inplace>
+        <<<dim_grid, dim_block, 0, at::cuda::getCurrentCUDAStream()>>>(
+          result_info, self_info, k, N_padded, last_dim_padded);
+    });
+    C10_CUDA_KERNEL_LAUNCH_CHECK();
+  } else {
+    auto result_info = cuda::detail::getTensorInfo<scalar_t, int64_t>(result);
+    auto self_info = cuda::detail::getTensorInfo<const scalar_t, int64_t>(self);
+    BOOL_SWITCH(self.is_same(result), inplace, [&] {
+      triu_tril_kernel<scalar_t, int64_t, upper, elements_per_thread, inplace>
+        <<<dim_grid, dim_block, 0, at::cuda::getCurrentCUDAStream()>>>(
+          result_info, self_info, k, N_padded, last_dim_padded);
+    });
+    C10_CUDA_KERNEL_LAUNCH_CHECK();
   }
 }
 
@@ -116,31 +187,7 @@ void triu_tril_cuda_template(const Tensor& result, const Tensor& self, int64_t k
       at::ScalarType::BFloat16,
       at::ScalarType::Bool,
       self.scalar_type(), "triu_tril_cuda_template", [&] {
-    constexpr int elements_per_thread = sizeof(scalar_t) < 8 ? 8 / sizeof(scalar_t) : 1;
-    auto sizes = self.sizes();
-    int64_t last_dim_padded = round_up<int64_t>(sizes.back(), elements_per_thread);
-    int64_t N_padded = c10::multiply_integers(sizes.begin(), sizes.end() - 1) * last_dim_padded;
-    dim3 dim_block = block_size;
-    dim3 dim_grid((N_padded / elements_per_thread + dim_block.x - 1) / dim_block.x);
-    if (cuda::detail::canUse32BitIndexMath(result) && cuda::detail::canUse32BitIndexMath(self)) {
-      auto result_info = cuda::detail::getTensorInfo<scalar_t, int32_t>(result);
-      auto self_info = cuda::detail::getTensorInfo<const scalar_t, int32_t>(self);
-      BOOL_SWITCH(self.is_same(result), inplace, [&] {
-        triu_tril_kernel<scalar_t, int32_t, upper, elements_per_thread, inplace>
-          <<<dim_grid, dim_block, 0, at::cuda::getCurrentCUDAStream()>>>(
-            result_info, self_info, k, N_padded, last_dim_padded);
-      });
-      C10_CUDA_KERNEL_LAUNCH_CHECK();
-    } else {
-      auto result_info = cuda::detail::getTensorInfo<scalar_t, int64_t>(result);
-      auto self_info = cuda::detail::getTensorInfo<const scalar_t, int64_t>(self);
-      BOOL_SWITCH(self.is_same(result), inplace, [&] {
-        triu_tril_kernel<scalar_t, int64_t, upper, elements_per_thread, inplace>
-          <<<dim_grid, dim_block, 0, at::cuda::getCurrentCUDAStream()>>>(
-            result_info, self_info, k, N_padded, last_dim_padded);
-      });
-      C10_CUDA_KERNEL_LAUNCH_CHECK();
-    }
+    launch_triu_tril_kernel<upper, scalar_t>(result, self, k);
   });
 }
 
diff --git a/aten/src/ATen/native/cuda/UpSampleBicubic2d.cu b/aten/src/ATen/native/cuda/UpSampleBicubic2d.cu
index 31cdf0a5688b7..53993d1bd9b83 100644
--- a/aten/src/ATen/native/cuda/UpSampleBicubic2d.cu
+++ b/aten/src/ATen/native/cuda/UpSampleBicubic2d.cu
@@ -7,6 +7,7 @@
 #include <ATen/Utils.h>
 #include <ATen/cuda/CUDAContext.h>
 #include <ATen/native/cuda/UpSample.cuh>
+#include <c10/util/irange.h>
 
 #ifndef AT_PER_OPERATOR_HEADERS
 #include <ATen/Functions.h>
@@ -93,6 +94,81 @@ __global__ void upsample_bicubic2d_out_frame(
   }
 }
 
+// Parallelized across batch*channels via blockIdx.z.
+// Faster than upsample_bicubic2d_out_frame when the output spatial size is
+// small (e.g. Kimi K2.5 position embeddings: 64x64 to 74x74 with 1152 channels).
+template <typename scalar_t, typename accscalar_t>
+C10_LAUNCH_BOUNDS_1(1024)
+__global__ void upsample_bicubic2d_out_frame_parallel(
+    const int num_elements,
+    const accscalar_t height_scale,
+    const accscalar_t width_scale,
+    const bool align_corners,
+    const PackedTensorAccessor64<const scalar_t, 4> idata,
+    PackedTensorAccessor64<scalar_t, 4> odata) {
+  int index = threadIdx.x + blockIdx.x * blockDim.x;
+
+  const int batchsize = idata.size(0);
+  const int channels = idata.size(1);
+  const int input_height = idata.size(2);
+  const int input_width = idata.size(3);
+  const int output_height = odata.size(2);
+  const int output_width = odata.size(3);
+
+  if (index >= num_elements) {
+    return;
+  }
+
+  const int output_x = index % output_width;
+  const int output_y = index / output_width;
+
+  if (input_height == output_height && input_width == output_width) {
+    for (int i = blockIdx.z; i < batchsize * channels; i += gridDim.z) {
+      int n = i / channels;
+      int c = i % channels;
+      const scalar_t val = idata[n][c][output_y][output_x];
+      odata[n][c][output_y][output_x] = val;
+    }
+    return;
+  }
+
+  accscalar_t real_x = area_pixel_compute_source_index(
+      width_scale, output_x, align_corners, /*cubic=*/true);
+  int in_x = floorf(real_x);
+  accscalar_t t_x = real_x - in_x;
+
+  accscalar_t real_y = area_pixel_compute_source_index(
+      height_scale, output_y, align_corners, /*cubic=*/true);
+  int in_y = floorf(real_y);
+  accscalar_t t_y = real_y - in_y;
+
+  for (int i = blockIdx.z; i < batchsize * channels; i += gridDim.z) {
+    int n = i / channels;
+    int c = i % channels;
+    accscalar_t coefficients[4];
+
+    for (const auto k : c10::irange(4)) {
+      coefficients[k] = cubic_interp1d(
+          upsample_get_value_bounded<scalar_t>(
+              idata, n, c, input_height, input_width, in_y - 1 + k, in_x - 1),
+          upsample_get_value_bounded<scalar_t>(
+              idata, n, c, input_height, input_width, in_y - 1 + k, in_x + 0),
+          upsample_get_value_bounded<scalar_t>(
+              idata, n, c, input_height, input_width, in_y - 1 + k, in_x + 1),
+          upsample_get_value_bounded<scalar_t>(
+              idata, n, c, input_height, input_width, in_y - 1 + k, in_x + 2),
+          t_x);
+    }
+
+    odata[n][c][output_y][output_x] = static_cast<scalar_t>(cubic_interp1d(
+        coefficients[0],
+        coefficients[1],
+        coefficients[2],
+        coefficients[3],
+        t_y));
+  }
+}
+
 // Backward (adjoint) operation 1 <- 2 (accumulates)
 template <typename scalar_t, typename accscalar_t>
 C10_LAUNCH_BOUNDS_1(1024)
@@ -204,17 +280,38 @@ static void upsample_bicubic2d_out_cuda_template(
         const accscalar_t rwidth = area_pixel_compute_scale<accscalar_t>(
             input_width, output_width, align_corners, scales_w);
 
-        upsample_bicubic2d_out_frame<scalar_t, accscalar_t>
-            <<<ceil_div(num_output_elements, max_threads),
-               max_threads,
-               0,
-               stream>>>(
-                num_output_elements,
-                rheight,
-                rwidth,
-                align_corners,
-                idata,
-                odata);
+        const int num_blocks = ceil_div(num_output_elements, max_threads);
+
+        // For small output spatial sizes the original kernel underutilizes
+        // the GPU because it loops over batch*channels sequentially.
+        // The parallel variant spreads that work across blockIdx.z instead.
+        // Threshold of 18432 covers all VLM position embedding resizing
+        // shapes (up to ~130x130 grids from 2048x2048 images) while
+        // staying well below the crossover point where the original
+        // kernel catches up (~65k+ output elements).
+        if (num_output_elements <= 18432 && input.size(0) * input.size(1) > 0) {
+          int* maxGridSize = at::cuda::getCurrentDeviceProperties()->maxGridSize;
+          int grid_z = std::min<int>(maxGridSize[2],
+                                     input.size(0) * input.size(1));
+          dim3 grid(num_blocks, 1, grid_z);
+          upsample_bicubic2d_out_frame_parallel<scalar_t, accscalar_t>
+              <<<grid, max_threads, 0, stream>>>(
+                  num_output_elements,
+                  rheight,
+                  rwidth,
+                  align_corners,
+                  idata,
+                  odata);
+        } else {
+          upsample_bicubic2d_out_frame<scalar_t, accscalar_t>
+              <<<num_blocks, max_threads, 0, stream>>>(
+                  num_output_elements,
+                  rheight,
+                  rwidth,
+                  align_corners,
+                  idata,
+                  odata);
+        }
         C10_CUDA_KERNEL_LAUNCH_CHECK();
       });
 }
diff --git a/aten/src/ATen/native/cuda/UpSampleNearest2d.cu b/aten/src/ATen/native/cuda/UpSampleNearest2d.cu
index e9470b3a6b473..0bf40e7b19bec 100644
--- a/aten/src/ATen/native/cuda/UpSampleNearest2d.cu
+++ b/aten/src/ATen/native/cuda/UpSampleNearest2d.cu
@@ -79,7 +79,18 @@ __global__ void upsample_nearest2d_out_frame(
   }
 }
 
-template <typename scalar_t, nn_compute_source_index_fn_t nn_compute_source_index_fn>
+// NHWC (channels-last) forward kernel.
+//
+// `kGridStride` selects the launch contract:
+//   - false: one element per thread; the launch grid covers `out_numel` exactly.
+//            The loop body executes once and the trailing `break` lets the
+//            compiler drop the loop entirely, preserving the original fast path.
+//   - true : grid is host-clamped (currently only on ROCm/HIP, see the comment
+//            on upsample_nearest2d_out_cuda_template) and threads stride over
+//            the output to cover the full range.
+template <typename scalar_t,
+          nn_compute_source_index_fn_t nn_compute_source_index_fn,
+          bool kGridStride>
 C10_LAUNCH_BOUNDS_1(1024)
 __global__ void upsample_nearest2d_nhwc_out_frame(
     const scalar_t* idata,
@@ -92,19 +103,28 @@ __global__ void upsample_nearest2d_nhwc_out_frame(
     float height_scale,
     float width_scale,
     const size_t out_numel) {
+  const int64_t start = static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x;
+  const int64_t step  = static_cast<int64_t>(gridDim.x) * blockDim.x;
+  const int64_t end   = static_cast<int64_t>(out_numel);
 
-    const int64_t index = ((int64_t) blockIdx.x) * blockDim.x + threadIdx.x;
-
-    if (index < out_numel) {
-    const auto c = index % channels;
+  for (int64_t index = start; index < end; index += step) {
+    const auto c  = index % channels;
     const auto w2 = (index / channels) % width2;
     const auto h2 = (index / channels / width2) % height2;
-    const auto n = index / channels / width2 / height2;
+    const auto n  = index / channels / width2 / height2;
 
-    const size_t h1 = height1 == height2 ? h2 : nn_compute_source_index_fn(height_scale, h2, height1);
-    const size_t w1 = width1 == width2 ? w2 : nn_compute_source_index_fn(width_scale, w2, width1);
+    const size_t h1 = height1 == height2
+        ? h2
+        : nn_compute_source_index_fn(height_scale, h2, height1);
+    const size_t w1 = width1 == width2
+        ? w2
+        : nn_compute_source_index_fn(width_scale, w2, width1);
 
     odata[index] = idata[idx_cl(n, h1, w1, c, height1, width1, channels)];
+
+    if constexpr (!kGridStride) {
+      break;
+    }
   }
 }
 
@@ -202,6 +222,13 @@ __global__ void upsample_nearest2d_backward_nhwc_out_frame(
   }
 }
 
+// On ROCm/HIP, gridDim.x * blockDim.x must fit in uint32_t per dimension (the
+// HSA AQL dispatch packet stores grid_size_{x,y,z} as uint32_t); violating it
+// returns hipErrorInvalidConfiguration.  Only the NHWC 1D-grid forward launch
+// can hit this in practice: with output.numel() near 2^32, a one-element-per-
+// thread grid overflows that field.  The contiguous (NCHW) 3D launch derives
+// grid_x/grid_y from int output_width/output_height, so per-dimension grids
+// are bounded well below the HIP limit.  See PR pytorch/pytorch#180310.
 template<nn_compute_source_index_fn_t nn_compute_source_index_fn>
 static void upsample_nearest2d_out_cuda_template(
     const Tensor& output,
@@ -234,6 +261,9 @@ static void upsample_nearest2d_out_cuda_template(
     return;
   }
 
+  // =================================================
+  // PATH 1: NHWC (channels-last) layout -- 1D grid
+  // =================================================
   // heuristic: only use channels_last path when it's faster than the contiguous path
   if (memory_format == at::MemoryFormat::ChannelsLast && channels >= 4 && \
         output.is_contiguous(memory_format)) {
@@ -247,25 +277,47 @@ static void upsample_nearest2d_out_cuda_template(
     const int64_t num_kernels = output.numel();
     const int64_t num_threads = std::min(at::cuda::getCurrentDeviceProperties()->maxThreadsPerBlock, 1024);
 
+    const int64_t grid = ceil_div(num_kernels, num_threads);
+#ifdef USE_ROCM
+    constexpr int64_t kHipMaxGlobalWorkSize = 4294967295LL;  // UINT32_MAX
+    const int64_t safe_max_grid = kHipMaxGlobalWorkSize / num_threads;
+    const bool use_grid_stride = grid > safe_max_grid;
+    const unsigned int launch_grid =
+        static_cast<unsigned int>(use_grid_stride ? safe_max_grid : grid);
+#else
+    constexpr bool use_grid_stride = false;
+    const int64_t launch_grid = grid;
+#endif
+
     AT_DISPATCH_FLOATING_TYPES_AND3(ScalarType::Half, ScalarType::BFloat16, ScalarType::Byte, input.scalar_type(), "upsample_nearest2d_nhwc_out_frame", [&] {
       const scalar_t* idata = input.const_data_ptr<scalar_t>();
       scalar_t* odata = output.mutable_data_ptr<scalar_t>();
-      upsample_nearest2d_nhwc_out_frame<scalar_t, nn_compute_source_index_fn>
-        <<<ceil_div(num_kernels, num_threads), num_threads, 0, at::cuda::getCurrentCUDAStream()>>>(
-          idata,
-          odata,
-          channels,
-          input_height,
-          input_width,
-          output_height,
-          output_width,
-          height_scale,
-          width_scale,
-          output.numel()
-      );
+      auto launch = [&](auto kgrid_stride) {
+        constexpr bool kGridStride = decltype(kgrid_stride)::value;
+        upsample_nearest2d_nhwc_out_frame<scalar_t, nn_compute_source_index_fn, kGridStride>
+          <<<launch_grid, num_threads, 0, at::cuda::getCurrentCUDAStream()>>>(
+            idata,
+            odata,
+            channels,
+            input_height,
+            input_width,
+            output_height,
+            output_width,
+            height_scale,
+            width_scale,
+            output.numel());
+      };
+      if (use_grid_stride) {
+        launch(std::true_type{});
+      } else {
+        launch(std::false_type{});
+      }
       C10_CUDA_KERNEL_LAUNCH_CHECK();
     });
   }
+  // ==============================================
+  // PATH 2: Contiguous (NCHW) layout -- 3D grid
+  // ==============================================
   else {
     // This is needed for non-contiguous tensors.
     Tensor output_c = output.is_contiguous() ? output : at::empty(output.sizes(), output.options());
@@ -293,15 +345,17 @@ static void upsample_nearest2d_out_cuda_template(
     int grid_y = ceil_div(output_height, block_y);
     int grid_z = std::min<int>(
         maxGridSize[2], ceil_div(nc, (int64_t) block_z * 4));
-    const dim3 grid(grid_x, grid_y, grid_z);
-    // Error out on cases where grid_x & grid_y exceeds limit of launch config, as
-    // the current kernel implementation doesn't loop over the two dimensions.
-    // This is unlikely to happen.
+
+    // output_height / output_width are int, so grid_x and grid_y are bounded
+    // by INT_MAX/block_dim - well below both the CUDA maxGridSize and the
+    // ROCm/HIP uint32_t per-dimension limit.  No HIP-specific clamp needed.
     // TODO: kernel implementation could stride on spatial dimension. We probably
     //       need to overhaul the kernel.
     TORCH_CHECK(
         grid_x <= maxGridSize[0] && grid_y <= maxGridSize[1],
         "input tensor has spatial dimension larger than the kernel capacity");
+
+    const dim3 grid(grid_x, grid_y, grid_z);
     cudaStream_t stream = at::cuda::getCurrentCUDAStream();
     AT_DISPATCH_FLOATING_TYPES_AND3(ScalarType::Half, ScalarType::BFloat16, ScalarType::Byte, input.scalar_type(), "upsample_nearest2d_out_frame", [&] {
           using accscalar_t = at::acc_type<scalar_t, true>;
@@ -329,6 +383,12 @@ static void upsample_nearest2d_out_cuda_template(
   }
 }
 
+// TODO(pytorch/pytorch#180310): Apply the same ROCm/HIP grid clamp to the
+// backward path.  The backward NHWC launch is bounded by
+// TORCH_CHECK(grad_input.numel() < INT_MAX), so its grid stays under 2^32.
+// The backward contiguous (NCHW) launch derives `n = grad_input.numel() /
+// nbatch`, which can exceed UINT32_MAX / bdim.x and therefore still needs a
+// host-side clamp + grid-stride fallback for completeness.
 template<nn_bw_compute_source_index_fn_t nn_bw_compute_source_index_fn>
 static void upsample_nearest2d_backward_out_cuda_template(
     const Tensor& grad_input,
diff --git a/aten/src/ATen/native/cuda/cutlass_extensions/gemm/kernel/fpA_intB_gemm.h b/aten/src/ATen/native/cuda/cutlass_extensions/gemm/kernel/fpA_intB_gemm.h
index dc07a53543ad7..f75606bdc9eae 100644
--- a/aten/src/ATen/native/cuda/cutlass_extensions/gemm/kernel/fpA_intB_gemm.h
+++ b/aten/src/ATen/native/cuda/cutlass_extensions/gemm/kernel/fpA_intB_gemm.h
@@ -68,13 +68,13 @@ struct GemmFpAIntB {
     using ElementA     = typename Mma::IteratorA::Element;
     using LayoutA      = typename Mma::IteratorA::Layout;
     using ElementB     = typename Mma::IteratorB::Element;
-    using LayoutB      = typename Mma::IteratorB::Element;
+    using LayoutB      = typename Mma::IteratorB::Layout;
     using ElementC     = typename Epilogue::OutputTileIterator::Element;
     using LayoutC      = typename Mma::LayoutC;
     using ElementScale = ElementC;
 
     static ComplexTransform const kTransformA = Mma::kTransformA;
-    static ComplexTransform const kTransformB = Mma::kTransformA;
+    static ComplexTransform const kTransformB = Mma::kTransformB;
 
     // Type definitions about the mainloop.
     using Operator         = typename Mma::Operator;
diff --git a/aten/src/ATen/native/cuda/fused_adagrad_utils.cuh b/aten/src/ATen/native/cuda/fused_adagrad_utils.cuh
index cb82c30fae198..fd70a33affdcf 100644
--- a/aten/src/ATen/native/cuda/fused_adagrad_utils.cuh
+++ b/aten/src/ATen/native/cuda/fused_adagrad_utils.cuh
@@ -104,6 +104,9 @@ struct FusedAdagradMathFunctor {
             found_inf_ptr);
 
         load_store(args[kParamIdx], r_args[kParamIdx], i_start, 0);
+        if (grad_scale_ptr) {
+          load_store(args[kGradIdx], r_args[kGradIdx], i_start, 0);
+        }
         load_store(args[kStateSumIdx], r_args[kStateSumIdx], i_start, 0);
       }
     } else {
diff --git a/aten/src/ATen/native/cuda/fused_adam_amsgrad_impl.cu b/aten/src/ATen/native/cuda/fused_adam_amsgrad_impl.cu
index cef07de1b41f9..f4844a79b2d70 100644
--- a/aten/src/ATen/native/cuda/fused_adam_amsgrad_impl.cu
+++ b/aten/src/ATen/native/cuda/fused_adam_amsgrad_impl.cu
@@ -1,6 +1,7 @@
 #include <ATen/native/cuda/fused_adam_amsgrad_impl.cuh>
 
 #include <ATen/Dispatch.h>
+#include <ATen/Dispatch_v2.h>
 #include <ATen/native/ForeachUtils.h>
 #include <ATen/native/cuda/MultiTensorApply.cuh>
 #include <ATen/native/cuda/fused_adam_utils.cuh>
@@ -36,26 +37,64 @@ void _fused_adam_amsgrad_cuda_impl_(
       found_inf.has_value() ? found_inf->data_ptr<float>() : nullptr;
   const float* lr_ptr = nullptr;
 
-  AT_DISPATCH_FLOATING_TYPES_AND2(
-      kHalf,
-      kBFloat16,
-      params[0].scalar_type(),
-      "fused_adam_kernel_cuda",
-      [&]() {
-        multi_tensor_apply_for_fused_optimizer<5>(
-            tensor_lists,
-            state_steps,
-            FusedAdamMathFunctor<scalar_t, 5, ADAM_MODE::ORIGINAL, true>(),
-            lr_ptr, // unused
-            lr,
-            beta1,
-            beta2,
-            weight_decay,
-            eps,
-            maximize,
-            grad_scale_ptr,
-            found_inf_ptr);
-      });
+  if (params[0].scalar_type() != exp_avgs[0].scalar_type()) {
+    validate_mixed_precision_dtypes(
+        params,
+        grads,
+        exp_avgs,
+        exp_avg_sqs,
+        max_exp_avg_sqs,
+        "Mixed-precision fused Adam");
+    AT_DISPATCH_V2(
+        exp_avgs[0].scalar_type(),
+        "fused_adam_amsgrad_mp_kernel_cuda",
+        AT_WRAP([&]() {
+          multi_tensor_apply_for_fused_optimizer<5>(
+              tensor_lists,
+              state_steps,
+              FusedAdamMathFunctorMP<
+                  float,
+                  float,
+                  float,
+                  scalar_t,
+                  scalar_t,
+                  scalar_t,
+                  5,
+                  ADAM_MODE::ORIGINAL,
+                  true>(),
+              lr_ptr, // unused
+              lr,
+              beta1,
+              beta2,
+              weight_decay,
+              eps,
+              maximize,
+              grad_scale_ptr,
+              found_inf_ptr);
+        }),
+        kBFloat16);
+  } else {
+    AT_DISPATCH_FLOATING_TYPES_AND2(
+        kHalf,
+        kBFloat16,
+        params[0].scalar_type(),
+        "fused_adam_kernel_cuda",
+        [&]() {
+          multi_tensor_apply_for_fused_optimizer<5>(
+              tensor_lists,
+              state_steps,
+              FusedAdamMathFunctor<scalar_t, 5, ADAM_MODE::ORIGINAL, true>(),
+              lr_ptr, // unused
+              lr,
+              beta1,
+              beta2,
+              weight_decay,
+              eps,
+              maximize,
+              grad_scale_ptr,
+              found_inf_ptr);
+        });
+  }
 }
 
 // The following overload simply has a Tensor lr
@@ -87,26 +126,64 @@ void _fused_adam_amsgrad_cuda_impl_(
       found_inf.has_value() ? found_inf->data_ptr<float>() : nullptr;
   const float* lr_ptr = lr.const_data_ptr<float>();
 
-  AT_DISPATCH_FLOATING_TYPES_AND2(
-      kHalf,
-      kBFloat16,
-      params[0].scalar_type(),
-      "fused_adam_kernel_cuda",
-      [&]() {
-        multi_tensor_apply_for_fused_optimizer<5>(
-            tensor_lists,
-            state_steps,
-            FusedAdamMathFunctor<scalar_t, 5, ADAM_MODE::ORIGINAL, true>(),
-            lr_ptr,
-            1.0, // unused
-            beta1,
-            beta2,
-            weight_decay,
-            eps,
-            maximize,
-            grad_scale_ptr,
-            found_inf_ptr);
-      });
+  if (params[0].scalar_type() != exp_avgs[0].scalar_type()) {
+    validate_mixed_precision_dtypes(
+        params,
+        grads,
+        exp_avgs,
+        exp_avg_sqs,
+        max_exp_avg_sqs,
+        "Mixed-precision fused Adam");
+    AT_DISPATCH_V2(
+        exp_avgs[0].scalar_type(),
+        "fused_adam_amsgrad_mp_kernel_cuda",
+        AT_WRAP([&]() {
+          multi_tensor_apply_for_fused_optimizer<5>(
+              tensor_lists,
+              state_steps,
+              FusedAdamMathFunctorMP<
+                  float,
+                  float,
+                  float,
+                  scalar_t,
+                  scalar_t,
+                  scalar_t,
+                  5,
+                  ADAM_MODE::ORIGINAL,
+                  true>(),
+              lr_ptr,
+              1.0, // unused
+              beta1,
+              beta2,
+              weight_decay,
+              eps,
+              maximize,
+              grad_scale_ptr,
+              found_inf_ptr);
+        }),
+        kBFloat16);
+  } else {
+    AT_DISPATCH_FLOATING_TYPES_AND2(
+        kHalf,
+        kBFloat16,
+        params[0].scalar_type(),
+        "fused_adam_kernel_cuda",
+        [&]() {
+          multi_tensor_apply_for_fused_optimizer<5>(
+              tensor_lists,
+              state_steps,
+              FusedAdamMathFunctor<scalar_t, 5, ADAM_MODE::ORIGINAL, true>(),
+              lr_ptr,
+              1.0, // unused
+              beta1,
+              beta2,
+              weight_decay,
+              eps,
+              maximize,
+              grad_scale_ptr,
+              found_inf_ptr);
+        });
+  }
 }
 
 } // namespace at::native
diff --git a/aten/src/ATen/native/cuda/fused_adam_impl.cu b/aten/src/ATen/native/cuda/fused_adam_impl.cu
index 2c1f5ce0d6d57..e51e446ea5e74 100644
--- a/aten/src/ATen/native/cuda/fused_adam_impl.cu
+++ b/aten/src/ATen/native/cuda/fused_adam_impl.cu
@@ -1,6 +1,7 @@
 #include <ATen/native/cuda/fused_adam_impl.cuh>
 
 #include <ATen/Dispatch.h>
+#include <ATen/Dispatch_v2.h>
 #include <ATen/native/ForeachUtils.h>
 #include <ATen/native/cuda/MultiTensorApply.cuh>
 #include <ATen/native/cuda/fused_adam_utils.cuh>
@@ -31,26 +32,59 @@ void _fused_adam_cuda_impl_(
       found_inf.has_value() ? found_inf->data_ptr<float>() : nullptr;
   const float* lr_ptr = nullptr;
 
-  AT_DISPATCH_FLOATING_TYPES_AND2(
-      kHalf,
-      kBFloat16,
-      params[0].scalar_type(),
-      "fused_adam_kernel_cuda",
-      [&]() {
-        multi_tensor_apply_for_fused_optimizer<4>(
-            tensor_lists,
-            state_steps,
-            FusedAdamMathFunctor<scalar_t, 4, ADAM_MODE::ORIGINAL, false>(),
-            lr_ptr, // unused
-            lr,
-            beta1,
-            beta2,
-            weight_decay,
-            eps,
-            maximize,
-            grad_scale_ptr,
-            found_inf_ptr);
-      });
+  if (params[0].scalar_type() != exp_avgs[0].scalar_type()) {
+    validate_mixed_precision_dtypes(
+        params, grads, exp_avgs, exp_avg_sqs, "Mixed-precision fused Adam");
+    AT_DISPATCH_V2(
+        exp_avgs[0].scalar_type(),
+        "fused_adam_mp_kernel_cuda",
+        AT_WRAP([&]() {
+          multi_tensor_apply_for_fused_optimizer<4>(
+              tensor_lists,
+              state_steps,
+              FusedAdamMathFunctorMP<
+                  float,
+                  float,
+                  float,
+                  scalar_t,
+                  scalar_t,
+                  float,
+                  4,
+                  ADAM_MODE::ORIGINAL,
+                  false>(),
+              lr_ptr, // unused
+              lr,
+              beta1,
+              beta2,
+              weight_decay,
+              eps,
+              maximize,
+              grad_scale_ptr,
+              found_inf_ptr);
+        }),
+        kBFloat16);
+  } else {
+    AT_DISPATCH_FLOATING_TYPES_AND2(
+        kHalf,
+        kBFloat16,
+        params[0].scalar_type(),
+        "fused_adam_kernel_cuda",
+        [&]() {
+          multi_tensor_apply_for_fused_optimizer<4>(
+              tensor_lists,
+              state_steps,
+              FusedAdamMathFunctor<scalar_t, 4, ADAM_MODE::ORIGINAL, false>(),
+              lr_ptr, // unused
+              lr,
+              beta1,
+              beta2,
+              weight_decay,
+              eps,
+              maximize,
+              grad_scale_ptr,
+              found_inf_ptr);
+        });
+  }
 }
 
 // The following overload simply has a Tensor lr
@@ -77,26 +111,59 @@ void _fused_adam_cuda_impl_(
       found_inf.has_value() ? found_inf->data_ptr<float>() : nullptr;
   const float* lr_ptr = lr.const_data_ptr<float>();
 
-  AT_DISPATCH_FLOATING_TYPES_AND2(
-      kHalf,
-      kBFloat16,
-      params[0].scalar_type(),
-      "fused_adam_kernel_cuda",
-      [&]() {
-        multi_tensor_apply_for_fused_optimizer<4>(
-            tensor_lists,
-            state_steps,
-            FusedAdamMathFunctor<scalar_t, 4, ADAM_MODE::ORIGINAL, false>(),
-            lr_ptr,
-            1.0, // unused
-            beta1,
-            beta2,
-            weight_decay,
-            eps,
-            maximize,
-            grad_scale_ptr,
-            found_inf_ptr);
-      });
+  if (params[0].scalar_type() != exp_avgs[0].scalar_type()) {
+    validate_mixed_precision_dtypes(
+        params, grads, exp_avgs, exp_avg_sqs, "Mixed-precision fused Adam");
+    AT_DISPATCH_V2(
+        exp_avgs[0].scalar_type(),
+        "fused_adam_mp_kernel_cuda",
+        AT_WRAP([&]() {
+          multi_tensor_apply_for_fused_optimizer<4>(
+              tensor_lists,
+              state_steps,
+              FusedAdamMathFunctorMP<
+                  float,
+                  float,
+                  float,
+                  scalar_t,
+                  scalar_t,
+                  float,
+                  4,
+                  ADAM_MODE::ORIGINAL,
+                  false>(),
+              lr_ptr,
+              1.0, // unused
+              beta1,
+              beta2,
+              weight_decay,
+              eps,
+              maximize,
+              grad_scale_ptr,
+              found_inf_ptr);
+        }),
+        kBFloat16);
+  } else {
+    AT_DISPATCH_FLOATING_TYPES_AND2(
+        kHalf,
+        kBFloat16,
+        params[0].scalar_type(),
+        "fused_adam_kernel_cuda",
+        [&]() {
+          multi_tensor_apply_for_fused_optimizer<4>(
+              tensor_lists,
+              state_steps,
+              FusedAdamMathFunctor<scalar_t, 4, ADAM_MODE::ORIGINAL, false>(),
+              lr_ptr,
+              1.0, // unused
+              beta1,
+              beta2,
+              weight_decay,
+              eps,
+              maximize,
+              grad_scale_ptr,
+              found_inf_ptr);
+        });
+  }
 }
 
 } // namespace at::native
diff --git a/aten/src/ATen/native/cuda/fused_adam_utils.cuh b/aten/src/ATen/native/cuda/fused_adam_utils.cuh
index 0533db28558b7..5440577ded0ad 100644
--- a/aten/src/ATen/native/cuda/fused_adam_utils.cuh
+++ b/aten/src/ATen/native/cuda/fused_adam_utils.cuh
@@ -3,12 +3,73 @@
 #include <ATen/native/cuda/ForeachFunctors.cuh>
 #include <ATen/native/cuda/MultiTensorApply.cuh>
 #include <ATen/native/cuda/Pow.cuh>
+#include <type_traits>
 #include <utility>
 
 namespace at::native {
 
 enum class ADAM_MODE : uint8_t { ORIGINAL = 0, ADAMW = 1 };
 
+// Validates the dtype configuration for mixed-precision fused Adam/AdamW.
+//
+// Currently the only supported configuration is:
+//   params/grads: float32
+//   optimizer states (exp_avg, exp_avg_sq, ...): bfloat16
+//
+// This specific configuration (fp32 params + bf16 optimizer states) has been
+// validated end-to-end in large-scale training runs (e.g. DeepSeek-V3 671B)
+// and is the only one for which training convergence has been demonstrated.
+// Additional mixed-precision configurations (e.g. float16 states) can be
+// enabled here once convergence is verified for those as well.
+//
+// Only [0] is checked because within-list dtype homogeneity is guaranteed by
+// _check_tensors_share_device_and_dtype (with skip_cross_list_dtype_check)
+// and the Python-side grouping in
+// _group_tensors_by_first_tensors_device_and_dtype.
+inline void validate_mixed_precision_dtypes(
+    at::TensorList params,
+    at::TensorList grads,
+    at::TensorList exp_avgs,
+    at::TensorList exp_avg_sqs,
+    const char* op_name) {
+  TORCH_CHECK(
+      params[0].scalar_type() == at::kFloat,
+      op_name,
+      " requires float32 params, got ",
+      params[0].scalar_type());
+  TORCH_CHECK(
+      grads[0].scalar_type() == at::kFloat,
+      op_name,
+      " requires float32 grads, got ",
+      grads[0].scalar_type());
+  TORCH_CHECK(
+      exp_avgs[0].scalar_type() == at::kBFloat16,
+      op_name,
+      " requires bfloat16 optimizer states, got ",
+      exp_avgs[0].scalar_type());
+  TORCH_CHECK(
+      exp_avg_sqs[0].scalar_type() == at::kBFloat16,
+      op_name,
+      " requires bfloat16 optimizer states, got ",
+      exp_avg_sqs[0].scalar_type());
+}
+
+inline void validate_mixed_precision_dtypes(
+    at::TensorList params,
+    at::TensorList grads,
+    at::TensorList exp_avgs,
+    at::TensorList exp_avg_sqs,
+    at::TensorList max_exp_avg_sqs,
+    const char* op_name) {
+  validate_mixed_precision_dtypes(
+      params, grads, exp_avgs, exp_avg_sqs, op_name);
+  TORCH_CHECK(
+      max_exp_avg_sqs[0].scalar_type() == at::kBFloat16,
+      op_name,
+      " requires bfloat16 max_exp_avg_sqs, got ",
+      max_exp_avg_sqs[0].scalar_type());
+}
+
 namespace {
 
 constexpr uint8_t kParamIdx = 0;
@@ -211,6 +272,282 @@ struct FusedAdamMathFunctor {
     }
   }
 };
+
+template <
+    typename scalar_type,
+    typename param_type,
+    typename grad_type,
+    typename exp_avg_type,
+    typename exp_avg_sq_type,
+    typename max_exp_avg_sq_type,
+    int depth,
+    ADAM_MODE adam_mode,
+    bool amsgrad>
+struct FusedAdamMathFunctorMP {
+  static_assert(
+      depth == 4 || depth == 5,
+      "depth of 4 for Adam, depth of 5 for Adam with AMSGrad.");
+  using opmath_t = at::opmath_type<scalar_type>;
+  C10_DEVICE __forceinline__ void operator()(
+      int64_t chunk_size,
+      FusedOptimizerTensorListMetadata<depth>& tl,
+      const float* lr_ptr,
+      const double& lr,
+      const double& beta1,
+      const double& beta2,
+      const double& weight_decay,
+      const double& eps,
+      const bool& maximize,
+      const float* grad_scale_ptr,
+      const float* found_inf_ptr) {
+    const auto tensor_loc = tl.block_to_tensor[blockIdx.x];
+    const auto chunk_idx = tl.block_to_chunk[blockIdx.x];
+    const double lr_double = lr_ptr ? *lr_ptr : lr;
+
+    if (found_inf_ptr && *found_inf_ptr == 1) {
+      return;
+    }
+    const auto [bias_correction1, bias_correction2_sqrt] =
+        [&]() -> std::pair<double, double> {
+      auto* step_count =
+          reinterpret_cast<const float*>(tl.state_steps_addresses[tensor_loc]);
+      const auto bias_correction1 = 1 - at::native::pow_(beta1, *step_count);
+      const auto bias_correction2 = 1 - at::native::pow_(beta2, *step_count);
+      const auto bias_correction2_sqrt = std::sqrt(bias_correction2);
+      return {bias_correction1, bias_correction2_sqrt};
+    }();
+
+    param_type* param_args;
+    grad_type* grad_args;
+    exp_avg_type* exp_avg_args;
+    exp_avg_sq_type* exp_avg_sq_args;
+    [[maybe_unused]] max_exp_avg_sq_type* max_exp_avg_sq_args;
+
+    // r_args represents the state when everything is casted to scalar_type
+    // to be passed into the adam_math function. scalar_type is our operation
+    // math type.
+    scalar_type r_args[depth][kILP];
+
+    // n = total numel of tensor - what's already been processed
+    // so n = numel in current tensor not yet processed
+    const auto n = tl.numel_for_tensor[tensor_loc] - chunk_idx * chunk_size;
+
+    bool all_aligned = init_args_mixed_prec<
+        depth,
+        param_type,
+        grad_type,
+        exp_avg_type,
+        exp_avg_sq_type>(
+        &param_args,
+        &grad_args,
+        &exp_avg_args,
+        &exp_avg_sq_args,
+        tl,
+        chunk_idx,
+        chunk_size,
+        tensor_loc);
+    if constexpr (amsgrad) {
+      max_exp_avg_sq_args =
+          (max_exp_avg_sq_type*)tl.addresses[kMaxExpAvgSqIdx][tensor_loc] +
+          chunk_idx * chunk_size;
+      all_aligned = all_aligned && is_aligned(max_exp_avg_sq_args);
+    }
+    if ((n % kILP == 0) && (chunk_size % kILP == 0) && all_aligned) {
+      for (int64_t i_start = threadIdx.x;
+           i_start * kILP < n && i_start * kILP < chunk_size;
+           i_start += blockDim.x) {
+        if constexpr (!std::is_same_v<scalar_type, param_type>) {
+          scalar_type casted_param_args[kILP];
+          for (int ii = 0; ii < kILP; ii++) {
+            casted_param_args[ii] =
+                static_cast<scalar_type>(param_args[ii + i_start * kILP]);
+          }
+          load_store(r_args[kParamIdx], casted_param_args, 0, 0);
+        } else {
+          load_store(r_args[kParamIdx], (scalar_type*)param_args, 0, i_start);
+        }
+        if constexpr (!std::is_same_v<scalar_type, grad_type>) {
+          scalar_type casted_grad_args[kILP];
+          for (int ii = 0; ii < kILP; ii++) {
+            casted_grad_args[ii] =
+                static_cast<scalar_type>(grad_args[ii + i_start * kILP]);
+          }
+          load_store(r_args[kGradIdx], casted_grad_args, 0, 0);
+        } else {
+          load_store(r_args[kGradIdx], (scalar_type*)grad_args, 0, i_start);
+        }
+        if constexpr (!std::is_same_v<scalar_type, exp_avg_type>) {
+          scalar_type casted_exp_avg_args[kILP];
+          for (int ii = 0; ii < kILP; ii++) {
+            casted_exp_avg_args[ii] =
+                static_cast<scalar_type>(exp_avg_args[ii + i_start * kILP]);
+          }
+          load_store(r_args[kExpAvgIdx], casted_exp_avg_args, 0, 0);
+        } else {
+          load_store(
+              r_args[kExpAvgIdx], (scalar_type*)exp_avg_args, 0, i_start);
+        }
+        if constexpr (!std::is_same_v<scalar_type, exp_avg_sq_type>) {
+          scalar_type casted_exp_avg_sq_args[kILP];
+          for (int ii = 0; ii < kILP; ii++) {
+            casted_exp_avg_sq_args[ii] =
+                static_cast<scalar_type>(exp_avg_sq_args[ii + i_start * kILP]);
+          }
+          load_store(r_args[kExpAvgSqIdx], casted_exp_avg_sq_args, 0, 0);
+        } else {
+          load_store(
+              r_args[kExpAvgSqIdx], (scalar_type*)exp_avg_sq_args, 0, i_start);
+        }
+        if constexpr (amsgrad) {
+          if constexpr (!std::is_same_v<scalar_type, max_exp_avg_sq_type>) {
+            scalar_type casted[kILP];
+            for (int ii = 0; ii < kILP; ii++) {
+              casted[ii] = static_cast<scalar_type>(
+                  max_exp_avg_sq_args[ii + i_start * kILP]);
+            }
+            load_store(r_args[kMaxExpAvgSqIdx], casted, 0, 0);
+          } else {
+            load_store(
+                r_args[kMaxExpAvgSqIdx],
+                (scalar_type*)max_exp_avg_sq_args,
+                0,
+                i_start);
+          }
+        }
+        adam_math<scalar_type, opmath_t, depth, adam_mode, amsgrad>(
+            r_args,
+            lr_double,
+            beta1,
+            beta2,
+            weight_decay,
+            eps,
+            maximize,
+            grad_scale_ptr,
+            found_inf_ptr,
+            bias_correction1,
+            bias_correction2_sqrt);
+        if constexpr (!std::is_same_v<scalar_type, param_type>) {
+          param_type casted_r_args[kILP];
+          for (int ii = 0; ii < kILP; ii++) {
+            casted_r_args[ii] = static_cast<param_type>(r_args[kParamIdx][ii]);
+          }
+          load_store(param_args, casted_r_args, i_start, 0);
+        } else {
+          load_store(param_args, (param_type*)r_args[kParamIdx], i_start, 0);
+        }
+        if constexpr (!std::is_same_v<scalar_type, exp_avg_type>) {
+          exp_avg_type casted_r_args[kILP];
+          for (int ii = 0; ii < kILP; ii++) {
+            casted_r_args[ii] =
+                static_cast<exp_avg_type>(r_args[kExpAvgIdx][ii]);
+          }
+          load_store(exp_avg_args, casted_r_args, i_start, 0);
+        } else {
+          load_store(
+              exp_avg_args, (exp_avg_type*)r_args[kExpAvgIdx], i_start, 0);
+        }
+        if constexpr (!std::is_same_v<scalar_type, exp_avg_sq_type>) {
+          exp_avg_sq_type casted_r_args[kILP];
+          for (int ii = 0; ii < kILP; ii++) {
+            casted_r_args[ii] =
+                static_cast<exp_avg_sq_type>(r_args[kExpAvgSqIdx][ii]);
+          }
+          load_store(exp_avg_sq_args, casted_r_args, i_start, 0);
+        } else {
+          load_store(
+              exp_avg_sq_args,
+              (exp_avg_sq_type*)r_args[kExpAvgSqIdx],
+              i_start,
+              0);
+        }
+        if constexpr (amsgrad) {
+          if constexpr (!std::is_same_v<scalar_type, max_exp_avg_sq_type>) {
+            max_exp_avg_sq_type casted[kILP];
+            for (int ii = 0; ii < kILP; ii++) {
+              casted[ii] =
+                  static_cast<max_exp_avg_sq_type>(r_args[kMaxExpAvgSqIdx][ii]);
+            }
+            load_store(max_exp_avg_sq_args, casted, i_start, 0);
+          } else {
+            load_store(
+                max_exp_avg_sq_args,
+                (max_exp_avg_sq_type*)r_args[kMaxExpAvgSqIdx],
+                i_start,
+                0);
+          }
+        }
+        if (grad_scale_ptr) {
+          if constexpr (!std::is_same_v<scalar_type, grad_type>) {
+            grad_type casted_r_args[kILP];
+            for (int ii = 0; ii < kILP; ii++) {
+              casted_r_args[ii] = static_cast<grad_type>(r_args[kGradIdx][ii]);
+            }
+            load_store(grad_args, casted_r_args, i_start, 0);
+          } else {
+            load_store(grad_args, (grad_type*)r_args[kGradIdx], i_start, 0);
+          }
+        }
+      }
+    } else {
+      for (int64_t i_start = 0; i_start < n && i_start < chunk_size;
+           i_start += blockDim.x * kILP) {
+        load_args<
+            scalar_type,
+            param_type,
+            grad_type,
+            exp_avg_type,
+            exp_avg_sq_type>(
+            r_args,
+            param_args,
+            grad_args,
+            exp_avg_args,
+            exp_avg_sq_args,
+            i_start,
+            chunk_size,
+            n);
+        if constexpr (amsgrad) {
+#pragma unroll
+          for (int ii = 0; ii < kILP; ii++) {
+            const auto i = i_start + threadIdx.x + ii * blockDim.x;
+            r_args[kMaxExpAvgSqIdx][ii] = 0;
+            if (i < n && i < chunk_size) {
+              r_args[kMaxExpAvgSqIdx][ii] =
+                  static_cast<scalar_type>(max_exp_avg_sq_args[i]);
+            }
+          }
+        }
+        adam_math<scalar_type, opmath_t, depth, adam_mode, amsgrad>(
+            r_args,
+            lr_double,
+            beta1,
+            beta2,
+            weight_decay,
+            eps,
+            maximize,
+            grad_scale_ptr,
+            found_inf_ptr,
+            bias_correction1,
+            bias_correction2_sqrt);
+        store_args(param_args, r_args[kParamIdx], i_start, chunk_size, n);
+        store_args(exp_avg_args, r_args[kExpAvgIdx], i_start, chunk_size, n);
+        store_args(
+            exp_avg_sq_args, r_args[kExpAvgSqIdx], i_start, chunk_size, n);
+        if constexpr (amsgrad) {
+          store_args(
+              max_exp_avg_sq_args,
+              r_args[kMaxExpAvgSqIdx],
+              i_start,
+              chunk_size,
+              n);
+        }
+        if (grad_scale_ptr) {
+          store_args(grad_args, r_args[kGradIdx], i_start, chunk_size, n);
+        }
+      }
+    }
+  }
+};
+
 } // namespace
 
 } // namespace at::native
diff --git a/aten/src/ATen/native/cuda/fused_adamw_amsgrad_impl.cu b/aten/src/ATen/native/cuda/fused_adamw_amsgrad_impl.cu
index b2eff4839133f..dba2b2c2af830 100644
--- a/aten/src/ATen/native/cuda/fused_adamw_amsgrad_impl.cu
+++ b/aten/src/ATen/native/cuda/fused_adamw_amsgrad_impl.cu
@@ -1,6 +1,7 @@
 #include <ATen/native/cuda/fused_adamw_amsgrad_impl.cuh>
 
 #include <ATen/Dispatch.h>
+#include <ATen/Dispatch_v2.h>
 #include <ATen/native/ForeachUtils.h>
 #include <ATen/native/cuda/MultiTensorApply.cuh>
 #include <ATen/native/cuda/fused_adam_utils.cuh>
@@ -36,26 +37,64 @@ void _fused_adamw_amsgrad_cuda_impl_(
       found_inf.has_value() ? found_inf->data_ptr<float>() : nullptr;
   const float* lr_ptr = nullptr;
 
-  AT_DISPATCH_FLOATING_TYPES_AND2(
-      kHalf,
-      kBFloat16,
-      params[0].scalar_type(),
-      "fused_adamw_kernel_cuda",
-      [&]() {
-        multi_tensor_apply_for_fused_optimizer<5>(
-            tensor_lists,
-            state_steps,
-            FusedAdamMathFunctor<scalar_t, 5, ADAM_MODE::ADAMW, true>(),
-            lr_ptr, // unused
-            lr,
-            beta1,
-            beta2,
-            weight_decay,
-            eps,
-            maximize,
-            grad_scale_ptr,
-            found_inf_ptr);
-      });
+  if (params[0].scalar_type() != exp_avgs[0].scalar_type()) {
+    validate_mixed_precision_dtypes(
+        params,
+        grads,
+        exp_avgs,
+        exp_avg_sqs,
+        max_exp_avg_sqs,
+        "Mixed-precision fused AdamW");
+    AT_DISPATCH_V2(
+        exp_avgs[0].scalar_type(),
+        "fused_adamw_amsgrad_mp_kernel_cuda",
+        AT_WRAP([&]() {
+          multi_tensor_apply_for_fused_optimizer<5>(
+              tensor_lists,
+              state_steps,
+              FusedAdamMathFunctorMP<
+                  float,
+                  float,
+                  float,
+                  scalar_t,
+                  scalar_t,
+                  scalar_t,
+                  5,
+                  ADAM_MODE::ADAMW,
+                  true>(),
+              lr_ptr, // unused
+              lr,
+              beta1,
+              beta2,
+              weight_decay,
+              eps,
+              maximize,
+              grad_scale_ptr,
+              found_inf_ptr);
+        }),
+        kBFloat16);
+  } else {
+    AT_DISPATCH_FLOATING_TYPES_AND2(
+        kHalf,
+        kBFloat16,
+        params[0].scalar_type(),
+        "fused_adamw_kernel_cuda",
+        [&]() {
+          multi_tensor_apply_for_fused_optimizer<5>(
+              tensor_lists,
+              state_steps,
+              FusedAdamMathFunctor<scalar_t, 5, ADAM_MODE::ADAMW, true>(),
+              lr_ptr, // unused
+              lr,
+              beta1,
+              beta2,
+              weight_decay,
+              eps,
+              maximize,
+              grad_scale_ptr,
+              found_inf_ptr);
+        });
+  }
 }
 
 // The following overload simply has a Tensor lr
@@ -87,26 +126,64 @@ void _fused_adamw_amsgrad_cuda_impl_(
       found_inf.has_value() ? found_inf->data_ptr<float>() : nullptr;
   const float* lr_ptr = lr.const_data_ptr<float>();
 
-  AT_DISPATCH_FLOATING_TYPES_AND2(
-      kHalf,
-      kBFloat16,
-      params[0].scalar_type(),
-      "fused_adamw_kernel_cuda",
-      [&]() {
-        multi_tensor_apply_for_fused_optimizer<5>(
-            tensor_lists,
-            state_steps,
-            FusedAdamMathFunctor<scalar_t, 5, ADAM_MODE::ADAMW, true>(),
-            lr_ptr,
-            1.0, // unused
-            beta1,
-            beta2,
-            weight_decay,
-            eps,
-            maximize,
-            grad_scale_ptr,
-            found_inf_ptr);
-      });
+  if (params[0].scalar_type() != exp_avgs[0].scalar_type()) {
+    validate_mixed_precision_dtypes(
+        params,
+        grads,
+        exp_avgs,
+        exp_avg_sqs,
+        max_exp_avg_sqs,
+        "Mixed-precision fused AdamW");
+    AT_DISPATCH_V2(
+        exp_avgs[0].scalar_type(),
+        "fused_adamw_amsgrad_mp_kernel_cuda",
+        AT_WRAP([&]() {
+          multi_tensor_apply_for_fused_optimizer<5>(
+              tensor_lists,
+              state_steps,
+              FusedAdamMathFunctorMP<
+                  float,
+                  float,
+                  float,
+                  scalar_t,
+                  scalar_t,
+                  scalar_t,
+                  5,
+                  ADAM_MODE::ADAMW,
+                  true>(),
+              lr_ptr,
+              1.0, // unused
+              beta1,
+              beta2,
+              weight_decay,
+              eps,
+              maximize,
+              grad_scale_ptr,
+              found_inf_ptr);
+        }),
+        kBFloat16);
+  } else {
+    AT_DISPATCH_FLOATING_TYPES_AND2(
+        kHalf,
+        kBFloat16,
+        params[0].scalar_type(),
+        "fused_adamw_kernel_cuda",
+        [&]() {
+          multi_tensor_apply_for_fused_optimizer<5>(
+              tensor_lists,
+              state_steps,
+              FusedAdamMathFunctor<scalar_t, 5, ADAM_MODE::ADAMW, true>(),
+              lr_ptr,
+              1.0, // unused
+              beta1,
+              beta2,
+              weight_decay,
+              eps,
+              maximize,
+              grad_scale_ptr,
+              found_inf_ptr);
+        });
+  }
 }
 
 } // namespace at::native
diff --git a/aten/src/ATen/native/cuda/fused_adamw_impl.cu b/aten/src/ATen/native/cuda/fused_adamw_impl.cu
index 90318854bec4c..0a65fee7ba84d 100644
--- a/aten/src/ATen/native/cuda/fused_adamw_impl.cu
+++ b/aten/src/ATen/native/cuda/fused_adamw_impl.cu
@@ -1,6 +1,7 @@
 #include <ATen/native/cuda/fused_adamw_impl.cuh>
 
 #include <ATen/Dispatch.h>
+#include <ATen/Dispatch_v2.h>
 #include <ATen/native/ForeachUtils.h>
 #include <ATen/native/cuda/MultiTensorApply.cuh>
 #include <ATen/native/cuda/fused_adam_utils.cuh>
@@ -31,26 +32,59 @@ void _fused_adamw_cuda_impl_(
       found_inf.has_value() ? found_inf->data_ptr<float>() : nullptr;
   const float* lr_ptr = nullptr;
 
-  AT_DISPATCH_FLOATING_TYPES_AND2(
-      kHalf,
-      kBFloat16,
-      params[0].scalar_type(),
-      "fused_adamw_kernel_cuda",
-      [&]() {
-        multi_tensor_apply_for_fused_optimizer<4>(
-            tensor_lists,
-            state_steps,
-            FusedAdamMathFunctor<scalar_t, 4, ADAM_MODE::ADAMW, false>(),
-            lr_ptr, // unused
-            lr,
-            beta1,
-            beta2,
-            weight_decay,
-            eps,
-            maximize,
-            grad_scale_ptr,
-            found_inf_ptr);
-      });
+  if (params[0].scalar_type() != exp_avgs[0].scalar_type()) {
+    validate_mixed_precision_dtypes(
+        params, grads, exp_avgs, exp_avg_sqs, "Mixed-precision fused AdamW");
+    AT_DISPATCH_V2(
+        exp_avgs[0].scalar_type(),
+        "fused_adamw_mp_kernel_cuda",
+        AT_WRAP([&]() {
+          multi_tensor_apply_for_fused_optimizer<4>(
+              tensor_lists,
+              state_steps,
+              FusedAdamMathFunctorMP<
+                  float,
+                  float,
+                  float,
+                  scalar_t,
+                  scalar_t,
+                  float,
+                  4,
+                  ADAM_MODE::ADAMW,
+                  false>(),
+              lr_ptr, // unused
+              lr,
+              beta1,
+              beta2,
+              weight_decay,
+              eps,
+              maximize,
+              grad_scale_ptr,
+              found_inf_ptr);
+        }),
+        kBFloat16);
+  } else {
+    AT_DISPATCH_FLOATING_TYPES_AND2(
+        kHalf,
+        kBFloat16,
+        params[0].scalar_type(),
+        "fused_adamw_kernel_cuda",
+        [&]() {
+          multi_tensor_apply_for_fused_optimizer<4>(
+              tensor_lists,
+              state_steps,
+              FusedAdamMathFunctor<scalar_t, 4, ADAM_MODE::ADAMW, false>(),
+              lr_ptr, // unused
+              lr,
+              beta1,
+              beta2,
+              weight_decay,
+              eps,
+              maximize,
+              grad_scale_ptr,
+              found_inf_ptr);
+        });
+  }
 }
 
 // The following overload simply has a Tensor lr
@@ -77,26 +111,59 @@ void _fused_adamw_cuda_impl_(
       found_inf.has_value() ? found_inf->data_ptr<float>() : nullptr;
   const float* lr_ptr = lr.const_data_ptr<float>();
 
-  AT_DISPATCH_FLOATING_TYPES_AND2(
-      kHalf,
-      kBFloat16,
-      params[0].scalar_type(),
-      "fused_adamw_kernel_cuda",
-      [&]() {
-        multi_tensor_apply_for_fused_optimizer<4>(
-            tensor_lists,
-            state_steps,
-            FusedAdamMathFunctor<scalar_t, 4, ADAM_MODE::ADAMW, false>(),
-            lr_ptr,
-            1.0, // unused
-            beta1,
-            beta2,
-            weight_decay,
-            eps,
-            maximize,
-            grad_scale_ptr,
-            found_inf_ptr);
-      });
+  if (params[0].scalar_type() != exp_avgs[0].scalar_type()) {
+    validate_mixed_precision_dtypes(
+        params, grads, exp_avgs, exp_avg_sqs, "Mixed-precision fused AdamW");
+    AT_DISPATCH_V2(
+        exp_avgs[0].scalar_type(),
+        "fused_adamw_mp_kernel_cuda",
+        AT_WRAP([&]() {
+          multi_tensor_apply_for_fused_optimizer<4>(
+              tensor_lists,
+              state_steps,
+              FusedAdamMathFunctorMP<
+                  float,
+                  float,
+                  float,
+                  scalar_t,
+                  scalar_t,
+                  float,
+                  4,
+                  ADAM_MODE::ADAMW,
+                  false>(),
+              lr_ptr,
+              1.0, // unused
+              beta1,
+              beta2,
+              weight_decay,
+              eps,
+              maximize,
+              grad_scale_ptr,
+              found_inf_ptr);
+        }),
+        kBFloat16);
+  } else {
+    AT_DISPATCH_FLOATING_TYPES_AND2(
+        kHalf,
+        kBFloat16,
+        params[0].scalar_type(),
+        "fused_adamw_kernel_cuda",
+        [&]() {
+          multi_tensor_apply_for_fused_optimizer<4>(
+              tensor_lists,
+              state_steps,
+              FusedAdamMathFunctor<scalar_t, 4, ADAM_MODE::ADAMW, false>(),
+              lr_ptr,
+              1.0, // unused
+              beta1,
+              beta2,
+              weight_decay,
+              eps,
+              maximize,
+              grad_scale_ptr,
+              found_inf_ptr);
+        });
+  }
 }
 
 } // namespace at::native
diff --git a/aten/src/ATen/native/cuda/group_norm_kernel.cu b/aten/src/ATen/native/cuda/group_norm_kernel.cu
index 254cd69466e8d..23c25554b3a0e 100644
--- a/aten/src/ATen/native/cuda/group_norm_kernel.cu
+++ b/aten/src/ATen/native/cuda/group_norm_kernel.cu
@@ -3,8 +3,6 @@
 
 #include <type_traits>
 
-#include <thrust/tuple.h>
-
 #include <ATen/core/Tensor.h>
 #include <ATen/AccumulateType.h>
 #include <ATen/Dispatch.h>
@@ -28,6 +26,19 @@ namespace {
 constexpr int kCUDANumThreads = 256;
 constexpr int kReduceTileSize = 32;
 
+// Reduce across exactly 32 lanes (offsets 16, 8, 4, 2, 1).
+// On NVIDIA (warp=32) this is identical to WarpReduceSum.
+// On AMD (wavefront=64) this avoids summing across two tile columns
+// when the block is (32, 16) and consecutive y-rows share a wavefront.
+template <typename T>
+__inline__ __device__ T ReduceSum32(T val) {
+#pragma unroll
+  for (int offset = 16; offset > 0; offset >>= 1) {
+    val += WARP_SHFL_DOWN(val, offset);
+  }
+  return val;
+}
+
 template <typename T>
 __global__ void RowwiseMomentsCUDAKernel(
     int64_t N,
@@ -54,7 +65,7 @@ __global__ void RowwiseMomentsCUDAKernel(
     // https://github.com/pytorch/pytorch/pull/13967
     __shared__ typename std::aligned_storage<
         sizeof(WelfordType),
-        alignof(WelfordType)>::type val_shared[C10_WARP_SIZE];
+        alignof(WelfordType)>::type val_shared[C10_WARP_SIZE_UPPER_BOUND];
     WelfordType* val_shared_ptr = reinterpret_cast<WelfordType*>(val_shared);
     val = cuda_utils::BlockReduce(
         val,
@@ -125,8 +136,8 @@ __global__ void Compute1dBackwardFusedParamsCUDAKernel(
     sum1 = cuda_utils::WarpReduceSum<T_ACC>(sum1);
     sum2 = cuda_utils::WarpReduceSum<T_ACC>(sum2);
   } else {
-    __shared__ T_ACC ds_shared[C10_WARP_SIZE];
-    __shared__ T_ACC db_shared[C10_WARP_SIZE];
+    __shared__ T_ACC ds_shared[C10_WARP_SIZE_UPPER_BOUND];
+    __shared__ T_ACC db_shared[C10_WARP_SIZE_UPPER_BOUND];
     sum1 = cuda_utils::BlockReduceSum<T_ACC>(sum1, ds_shared);
     sum2 = cuda_utils::BlockReduceSum<T_ACC>(sum2, db_shared);
   }
@@ -240,8 +251,11 @@ __global__ void GammaBeta1dBackwardCUDAKernel2(
   // Do warp reduce for the 1st 16 cols in the tile.
   T_ACC sum1 = g_shared[threadIdx.x][threadIdx.y];
   T_ACC sum2 = b_shared[threadIdx.x][threadIdx.y];
-  sum1 = cuda_utils::WarpReduceSum<T_ACC>(sum1);
-  sum2 = cuda_utils::WarpReduceSum<T_ACC>(sum2);
+  // Use ReduceSum32 (not WarpReduceSum) to reduce exactly 32 lanes.
+  // On AMD wavefront-64, WarpReduceSum would incorrectly sum across two
+  // tile columns since consecutive y-rows share a wavefront.
+  sum1 = ReduceSum32<T_ACC>(sum1);
+  sum2 = ReduceSum32<T_ACC>(sum2);
   if (threadIdx.x == 0) {
     const int64_t c = blockIdx.x * blockDim.x + threadIdx.y;
     if (c < C) {
@@ -257,8 +271,8 @@ __global__ void GammaBeta1dBackwardCUDAKernel2(
   // Do warp reduce for the 2nd 16 cols in the tile.
   sum1 = g_shared[threadIdx.x][threadIdx.y + blockDim.y];
   sum2 = b_shared[threadIdx.x][threadIdx.y + blockDim.y];
-  sum1 = cuda_utils::WarpReduceSum<T_ACC>(sum1);
-  sum2 = cuda_utils::WarpReduceSum<T_ACC>(sum2);
+  sum1 = ReduceSum32<T_ACC>(sum1);
+  sum2 = ReduceSum32<T_ACC>(sum2);
   if (threadIdx.x == 0) {
     const int64_t c = blockIdx.x * blockDim.x + threadIdx.y + blockDim.y;
     if (c < C) {
@@ -292,8 +306,8 @@ __global__ void ComputeInternalGradientsCUDAKernel(
     sum1 = cuda_utils::WarpReduceSum<T_ACC>(sum1);
     sum2 = cuda_utils::WarpReduceSum<T_ACC>(sum2);
   } else {
-    __shared__ T_ACC ds_shared[C10_WARP_SIZE];
-    __shared__ T_ACC db_shared[C10_WARP_SIZE];
+    __shared__ T_ACC ds_shared[C10_WARP_SIZE_UPPER_BOUND];
+    __shared__ T_ACC db_shared[C10_WARP_SIZE_UPPER_BOUND];
     sum1 = cuda_utils::BlockReduceSum<T_ACC>(sum1, ds_shared);
     sum2 = cuda_utils::BlockReduceSum<T_ACC>(sum2, db_shared);
   }
@@ -335,8 +349,8 @@ __global__ void ComputeBackwardFusedParamsCUDAKernel(
     sum1 = cuda_utils::WarpReduceSum<T_ACC>(sum1);
     sum2 = cuda_utils::WarpReduceSum<T_ACC>(sum2);
   } else {
-    __shared__ T_ACC ds_shared[C10_WARP_SIZE];
-    __shared__ T_ACC db_shared[C10_WARP_SIZE];
+    __shared__ T_ACC ds_shared[C10_WARP_SIZE_UPPER_BOUND];
+    __shared__ T_ACC db_shared[C10_WARP_SIZE_UPPER_BOUND];
     sum1 = cuda_utils::BlockReduceSum<T_ACC>(sum1, ds_shared);
     sum2 = cuda_utils::BlockReduceSum<T_ACC>(sum2, db_shared);
   }
@@ -442,10 +456,11 @@ __global__ void GammaBetaBackwardCUDAKernel2(
   __syncthreads();
 
   // Do warp reduce for the 1st 16 cols in the tile.
+  // Use ReduceSum32 for correctness on AMD wavefront-64 (see above).
   T_ACC sum1 = g_shared[threadIdx.x][threadIdx.y];
   T_ACC sum2 = b_shared[threadIdx.x][threadIdx.y];
-  sum1 = cuda_utils::WarpReduceSum<T_ACC>(sum1);
-  sum2 = cuda_utils::WarpReduceSum<T_ACC>(sum2);
+  sum1 = ReduceSum32<T_ACC>(sum1);
+  sum2 = ReduceSum32<T_ACC>(sum2);
   if (threadIdx.x == 0) {
     const int64_t c = blockIdx.x * blockDim.x + threadIdx.y;
     if (c < C) {
@@ -461,8 +476,8 @@ __global__ void GammaBetaBackwardCUDAKernel2(
   // Do warp reduce for the 2nd 16 cols in the tile.
   sum1 = g_shared[threadIdx.x][threadIdx.y + blockDim.y];
   sum2 = b_shared[threadIdx.x][threadIdx.y + blockDim.y];
-  sum1 = cuda_utils::WarpReduceSum<T_ACC>(sum1);
-  sum2 = cuda_utils::WarpReduceSum<T_ACC>(sum2);
+  sum1 = ReduceSum32<T_ACC>(sum1);
+  sum2 = ReduceSum32<T_ACC>(sum2);
   if (threadIdx.x == 0) {
     const int64_t c = blockIdx.x * blockDim.x + threadIdx.y + blockDim.y;
     if (c < C) {
diff --git a/aten/src/ATen/native/cuda/int4mm.cu b/aten/src/ATen/native/cuda/int4mm.cu
index ca00c944b3259..21893ea5622f4 100644
--- a/aten/src/ATen/native/cuda/int4mm.cu
+++ b/aten/src/ATen/native/cuda/int4mm.cu
@@ -146,6 +146,10 @@ static bool isCDNA2orLater(int index) {
     return at::detail::getCUDAHooks().isGPUArch({"gfx90a", "gfx942", "gfx950"}, index);
 }
 
+static bool isCDNA5orLater(int index) {
+  return at::detail::getCUDAHooks().isGPUArch({"gfx1250"}, index);
+}
+
 #else
 constexpr int32_t kWarpSize = 32;
 #endif
@@ -576,7 +580,14 @@ struct BLayout_TC_int4 {
           // type pun, the __nv_bfloat162 value in bf16x2x4 is a struct and
           // can't be used as a 32-bit asm register argument for `mma`
           static_assert(sizeof(bf16x2x4) == sizeof(out[0][0]), "");
+          // On Windows with ROCm, std::memcpy resolves to a __host__-only
+          // function and cannot be called from __device__ code. Use the raw
+          // memcpy which the HIP compiler provides as a __device__ builtin.
+#if defined(_WIN32) && defined(USE_ROCM)
+          memcpy(&out[i][j], &v, sizeof(bf16x2x4_u32));
+#else
           std::memcpy(&out[i][j], &v, sizeof(bf16x2x4_u32));
+#endif
         }
       }
     }
@@ -1091,6 +1102,11 @@ at::Tensor _weight_int4pack_mm_cuda(
       A.device() == B.device() && A.device() == qScaleAndZeros.device());
 
 #if defined(USE_ROCM)
+  if (isCDNA5orLater(A.device().index())) {
+    TORCH_CHECK(false,
+                "_weight_int4pack_mm_cuda is not yet supported on gfx1250. "
+                "A WMMA-based implementation is required for gfx1250.");
+  }
   if (!isCDNA2orLater(A.device().index())) {
     TORCH_CHECK(false, "_weight_int4pack_mm_cuda is only supported on AMD gpu arch greater than or equal to CDNA2");
   }
@@ -1286,6 +1302,11 @@ at::Tensor _convert_weight_to_int4pack_cuda(
   TORCH_CHECK(innerKTiles == 2 || innerKTiles == 4 || innerKTiles == 8);
 
 #if defined(USE_ROCM)
+  if (isCDNA5orLater(in.device().index())) {
+    TORCH_CHECK(false,
+                "_convert_weight_to_int4pack_cuda is not yet supported on gfx1250. "
+                "A WMMA-based implementation is required for gfx1250.");
+  }
   if (!isCDNA2orLater(in.device().index())) {
     TORCH_CHECK(false, "_convert_weight_to_int4pack_cuda is only supported on AMD gpu arch greater than or equal to CDNA2");
   }
diff --git a/aten/src/ATen/native/cuda/jit_utils.cpp b/aten/src/ATen/native/cuda/jit_utils.cpp
index f74c86f435fbc..74739b348db1f 100644
--- a/aten/src/ATen/native/cuda/jit_utils.cpp
+++ b/aten/src/ATen/native/cuda/jit_utils.cpp
@@ -981,7 +981,6 @@ int calc_thread_work_size(
     } else {
         return 4;
     }
-    return io_size;
 #else
     auto io_size = at::cuda::jit::calc_io_size(nInputs, nOutputs, inputs_type, result_type);
     TORCH_INTERNAL_ASSERT(io_size > 0);
@@ -990,7 +989,6 @@ int calc_thread_work_size(
     } else {
         return 8;
     }
-    return io_size;
 #endif
 }
 
@@ -1599,7 +1597,7 @@ NvrtcFunction jit_pwise_function(
       &program, code.c_str(), nullptr, 0, nullptr, nullptr));
 
 #ifdef USE_ROCM
-  std::vector<const char*> args = {"--std=c++17"};
+  std::vector<const char*> args = {"--std=c++20"};
 #else
   // Constructs nvrtc build arguments
   // CUDA 11.1 allows going directly to SASS (sm_) instead of PTX (compute_)
@@ -1613,7 +1611,7 @@ NvrtcFunction jit_pwise_function(
       (compile_to_sass ? "sm_" : "compute_") + std::to_string(cuda_major) +
       std::to_string(cuda_minor);
   std::vector<const char*> args = {
-      "--std=c++17", compute.c_str(), "-default-device"};
+      "--std=c++20", compute.c_str(), "-default-device"};
 #endif
 
   #ifndef NDEBUG
diff --git a/aten/src/ATen/native/cuda/layer_norm_kernel.cu b/aten/src/ATen/native/cuda/layer_norm_kernel.cu
index 0a4b58cbdbd85..59f5e4a041117 100644
--- a/aten/src/ATen/native/cuda/layer_norm_kernel.cu
+++ b/aten/src/ATen/native/cuda/layer_norm_kernel.cu
@@ -1,10 +1,9 @@
 #define TORCH_ASSERT_ONLY_METHOD_OPERATORS
 #include <ATen/native/layer_norm.h>
 
+#include <tuple>
 #include <type_traits>
 
-#include <thrust/tuple.h>
-
 #include <ATen/core/Tensor.h>
 #include <ATen/AccumulateType.h>
 #include <ATen/Dispatch.h>
@@ -33,12 +32,6 @@ namespace at::native {
 namespace {
 
 constexpr int kCUDANumThreads = 256;
-#ifdef USE_ROCM
-// C10_WARP_SIZE is not constexpr for host code.
-#define kWarpSize C10_WARP_SIZE
-#else
-constexpr unsigned int kWarpSize = C10_WARP_SIZE;
-#endif
 constexpr int vec_size = 4; //we could make it dependent on dtype, but that would lead to different results between float and low-p types
 
 // aligned vector generates vectorized load/store on CUDA (copy-pasted from MemoryAccess.cuh)
@@ -68,7 +61,7 @@ __global__ void RowwiseMomentsCUDAKernel(
 
   __shared__
       typename std::aligned_storage<sizeof(WelfordType), alignof(WelfordType)>::
-          type val_shared[C10_WARP_SIZE];
+          type val_shared[C10_WARP_SIZE_UPPER_BOUND];
   WelfordType* val_shared_ptr = reinterpret_cast<WelfordType*>(val_shared);
 
   const int64_t i = blockIdx.x;
@@ -663,7 +656,7 @@ blockReduceGammaBetaBackwardsHelper(
   constexpr int rows_per_thread_y = rows_per_block_y / block_dim_y;
   int64_t thread_x = blockIdx.x * block_dim_x + threadIdx.x;
 
-    int lane_id = (threadIdx.y * blockDim.x + threadIdx.x) & (kWarpSize - 1);
+    int lane_id = (threadIdx.y * blockDim.x + threadIdx.x) & (C10_WARP_SIZE - 1);
     int64_t mean_index = M_start + threadIdx.y * rows_per_thread_y;
     T_ACC warp_mean = 0, warp_rstd = 0;
     if (lane_id < rows_per_thread_y && mean_index + lane_id < M) {
@@ -696,9 +689,9 @@ blockReduceGammaBetaBackwardsHelper(
 
     #pragma unroll
     for (int i = 0; i < rows_per_thread_y; ++i) {
-      T_ACC rstd_reg = WARP_SHFL(warp_rstd, i, kWarpSize);
+      T_ACC rstd_reg = WARP_SHFL(warp_rstd, i, C10_WARP_SIZE);
       if constexpr (!rms_norm){
-        T_ACC mean_reg = WARP_SHFL(warp_mean, i, kWarpSize);
+        T_ACC mean_reg = WARP_SHFL(warp_mean, i, C10_WARP_SIZE);
         dg_sum += dY_regs[i] * (X_regs[i] - mean_reg) * rstd_reg;
         db_sum += dY_regs[i];
       } else{
@@ -776,7 +769,7 @@ __launch_bounds__(block_dim_x * block_dim_y)
     T* __restrict__ db) {
   // This assert is a compile-time check only.
   constexpr int rows_per_thread_y = rows_per_block_y / block_dim_y;
-  static_assert(rows_per_thread_y <= kWarpSize);
+  static_assert(rows_per_thread_y <= C10_WARP_SIZE_LOWER_BOUND);
 
   T_ACC dg_sum = 0;
   T_ACC db_sum = 0;
@@ -819,7 +812,7 @@ __launch_bounds__(block_dim_x * block_dim_y)
   } else {
     // The caller requested a full reduction so we must reduce across
     // warps using shared memory and warp shuffles.
-    static_assert(rows_per_thread_y <= C10_WARP_SIZE);
+    static_assert(rows_per_thread_y <= C10_WARP_SIZE_LOWER_BOUND);
     alignas(sizeof(double)) extern __shared__ char s_data1[];
     T_ACC* s_data_typed = reinterpret_cast<T_ACC*>(&s_data1);
     T_ACC* s_dg;
@@ -835,8 +828,8 @@ __launch_bounds__(block_dim_x * block_dim_y)
     // Load transposed so that a warp holds an entire column
     // Because block_dim_x != block_dim_y in the general case, we need
     // some code to handle the general case.
-    static_assert(block_dim_x * block_dim_y % C10_WARP_SIZE == 0);
-    constexpr int warps_available_to_reduce = block_dim_x * block_dim_y / C10_WARP_SIZE;
+    static_assert(block_dim_x * block_dim_y % C10_WARP_SIZE_LOWER_BOUND == 0);
+    const int warps_available_to_reduce = block_dim_x * block_dim_y / C10_WARP_SIZE;
     int thread_id = threadIdx.y * block_dim_x + threadIdx.x;
     int warp_id = thread_id / C10_WARP_SIZE;
     int lane_id = thread_id & (C10_WARP_SIZE - 1);
@@ -849,8 +842,8 @@ __launch_bounds__(block_dim_x * block_dim_y)
       }
       #pragma unroll
       for (unsigned delta = block_dim_y >> 1; delta >= 1; delta >>= 1) {
-        reg_dg += WARP_SHFL_XOR(reg_dg, delta, kWarpSize);
-        reg_db += WARP_SHFL_XOR(reg_db, delta, kWarpSize);
+        reg_dg += WARP_SHFL_XOR(reg_dg, delta, C10_WARP_SIZE);
+        reg_db += WARP_SHFL_XOR(reg_db, delta, C10_WARP_SIZE);
       }
       // Reduce is done. Now write it out to global memory.
       int64_t out_index = ((int64_t)blockIdx.x) * block_dim_x + i;
@@ -956,7 +949,11 @@ void LaunchGammaBetaBackwardCUDAKernel(
     Tensor* dgamma,
     Tensor* dbeta,
     cudaStream_t cuda_stream) {
+#ifdef USE_ROCM
+  constexpr int block_dim_x = 64;
+#else
   constexpr int block_dim_x = 32;
+#endif
   const int sm_count = at::cuda::getCurrentDeviceProperties()->multiProcessorCount;
   if (M > 64 * 1024 && N / block_dim_x < sm_count / 2) {
     // We have a situation where M >> N and N is small.
@@ -1011,7 +1008,13 @@ void LaunchGammaBetaBackwardCUDAKernel(
     } else if (M < 256) {
       ConfigureAndLaunchGammaBetaBackwardKernel<T, T_ACC, block_dim_x, 16, 128, rms_norm>(dY_data, X_data, mean_data, rstd_data, M, N, dgamma, dbeta, cuda_stream);
     } else {
+#ifdef USE_ROCM
+      // Cap block_dim_y at 16 to keep total threads (64*16=1024) within GPU limits.
+      // rows_per_thread_y = 256/16 = 16, still within warp size constraint.
+      ConfigureAndLaunchGammaBetaBackwardKernel<T, T_ACC, block_dim_x, 16, 256, rms_norm>(dY_data, X_data, mean_data, rstd_data, M, N, dgamma, dbeta, cuda_stream);
+#else
       ConfigureAndLaunchGammaBetaBackwardKernel<T, T_ACC, block_dim_x, 32, 256, rms_norm>(dY_data, X_data, mean_data, rstd_data, M, N, dgamma, dbeta, cuda_stream);
+#endif
     }
   }
 }
@@ -1618,40 +1621,12 @@ void LayerNormBackwardKernelImplInternal(
               dbeta_data);
       C10_CUDA_KERNEL_LAUNCH_CHECK();
     } else {
-      // For small batch size, do colwise reduce directly.
-      const int part_size = warp_size;
-      const dim3 threads2(warp_size, 4, 1);
-      const dim3 blocks2((N + threads2.x - 1) / threads2.x, part_size, 1);
-      const int nshared2_a = 2 * sizeof(T_ACC) * threads2.y * threads2.y * (threads2.x + 1);
-      const int nshared2_b = threads2.x * threads2.y * sizeof(T_ACC);
-      const int nshared2 = nshared2_a > nshared2_b ? nshared2_a : nshared2_b;
-
-      const auto part_grad_dtype = at::toAccumulateType(X.scalar_type(), true);
-      Tensor part_grad_gamma = at::empty({part_size,N}, gamma.options().dtype(part_grad_dtype));
-      Tensor part_grad_beta = at::native::empty_like(part_grad_gamma);
-
-      cuComputePartGradGammaBeta<T, T_ACC, rms_norm><<<blocks2, threads2, nshared2, cuda_stream>>>(
-                      dY_data,
-                      X_data,
-                      M,N,
-                      mean_data,
-                      rstd_data,
-                      part_grad_gamma.template data_ptr<T_ACC>(),
-                      part_grad_beta.template data_ptr<T_ACC>());
-      C10_CUDA_KERNEL_LAUNCH_CHECK();
-
-      const dim3 threads3(warp_size, 8, 1); // Optimization for ROCm
-      const dim3 blocks3((N + threads3.x - 1) / threads3.x, 1, 1);
-      const int nshared3 = threads3.x * threads3.y * sizeof(T_ACC);
-
-      cuComputeGradGammaBeta<T, T_ACC, rms_norm><<<blocks3, threads3, nshared3, cuda_stream>>>(
-                      part_grad_gamma.template data_ptr<T_ACC>(),
-                      part_grad_beta.template data_ptr<T_ACC>(),
-                      part_size,
-                      M,N,
-                      dgamma_data,
-                      dbeta_data);
-      C10_CUDA_KERNEL_LAUNCH_CHECK();
+      // Use the optimized tiled kernel adapted for wavefront-64.
+      // This replaces the legacy two-pass cuComputePartGradGammaBeta +
+      // cuComputeGradGammaBeta approach with a single-pass tiled reduction
+      // that has coalesced memory access and adaptive tile sizing.
+      LaunchGammaBetaBackwardCUDAKernel<T, T_ACC, rms_norm>(
+        dY_data, X_data, mean_data, rstd_data, M, N, dgamma, dbeta, cuda_stream);
     }
 #else
     LaunchGammaBetaBackwardCUDAKernel<T, T_ACC, rms_norm>(
diff --git a/aten/src/ATen/native/cuda/linalg/BatchLinearAlgebra.cpp b/aten/src/ATen/native/cuda/linalg/BatchLinearAlgebra.cpp
index 83e957cc6729a..da21a3fed9c00 100644
--- a/aten/src/ATen/native/cuda/linalg/BatchLinearAlgebra.cpp
+++ b/aten/src/ATen/native/cuda/linalg/BatchLinearAlgebra.cpp
@@ -39,8 +39,6 @@
 #include <magma_v2.h>
 #include <ATen/cuda/detail/CUDAHooks.h>
 
-const bool use_magma_ = true;
-
 namespace {
 struct MagmaInitializer {
   MagmaInitializer() {
@@ -60,9 +58,6 @@ struct MagmaInitializer {
 #error "MAGMA release minor or micro version >= 10, please correct AT_MAGMA_VERSION"
 #endif
 
-#else
-const bool use_magma_ = false;
-
 #endif
 
 namespace at::native {
@@ -110,16 +105,6 @@ void magmaLuNoPivBatched(
     magma_int_t m, magma_int_t n, scalar_t** dA_array, magma_int_t ldda,
     magma_int_t* info_array, magma_int_t batchsize, const MAGMAQueue& magma_queue);
 
-template<class scalar_t>
-void magmaCholeskySolve(
-    magma_uplo_t uplo, magma_int_t n, magma_int_t nrhs, scalar_t* dA, magma_int_t ldda,
-    scalar_t* dB, magma_int_t lddb, magma_int_t* info);
-
-template<class scalar_t>
-void magmaCholeskySolveBatched(
-    magma_uplo_t uplo, magma_int_t n, magma_int_t nrhs, scalar_t** dA_array, magma_int_t ldda,
-    scalar_t** dB_array, magma_int_t lddb, magma_int_t& info, magma_int_t batchsize, const MAGMAQueue& magma_queue);
-
 #if defined(USE_ROCM)
 template<class scalar_t>
 void magmaTriangularSolveBatched(
@@ -128,12 +113,7 @@ void magmaTriangularSolveBatched(
     const MAGMAQueue& magma_queue);
 #endif // defined(USE_ROCM)
 
-template<class scalar_t, class value_t=scalar_t>
-void magmaSyevd(
-    magma_vec_t jobz, magma_uplo_t uplo, magma_int_t n, scalar_t* dA, magma_int_t ldda,
-    value_t* w, scalar_t* wA, magma_int_t ldwa, scalar_t* work, magma_int_t lwork, value_t* rwork,
-    magma_int_t lrwork, magma_int_t* iwork, magma_int_t liwork, magma_int_t* info);
-
+#if defined(USE_ROCM) || !(defined(CUSOLVER_VERSION) && (CUSOLVER_VERSION >= 11702))
 template<class scalar_t, class value_t=scalar_t>
 void magmaEig(
     magma_vec_t jobvl, magma_vec_t jobvr, magma_int_t n, scalar_t *A, magma_int_t lda,
@@ -141,17 +121,7 @@ void magmaEig(
     scalar_t *VR, magma_int_t ldvr, scalar_t *work, magma_int_t lwork,
     value_t *rwork,
     magma_int_t *info);
-
-template<class scalar_t>
-void magmaLuSolve(
-    magma_int_t n, magma_int_t nrhs, scalar_t* dA, magma_int_t ldda, magma_int_t* ipiv,
-    scalar_t* dB, magma_int_t lddb, magma_int_t* info, magma_trans_t trans);
-
-template<class scalar_t>
-void magmaLuSolveBatched(
-    magma_int_t n, magma_int_t nrhs, scalar_t** dA_array, magma_int_t ldda, magma_int_t** dipiv_array,
-    scalar_t** dB_array, magma_int_t lddb, magma_int_t& info,
-    magma_int_t batchsize, const MAGMAQueue& magma_queue, magma_trans_t trans);
+#endif
 
 #if AT_MAGMA_VERSION >= 254
 
@@ -351,82 +321,6 @@ void magmaLuNoPivBatched<c10::complex<float>>(
   AT_CUDA_CHECK(cudaGetLastError());
 }
 
-template<>
-void magmaCholeskySolve<double>(
-    magma_uplo_t uplo, magma_int_t n, magma_int_t nrhs, double* dA, magma_int_t ldda,
-    double* dB, magma_int_t lddb, magma_int_t* info) {
-  MagmaStreamSyncGuard guard;
-  magma_dpotrs_gpu(uplo, n, nrhs, dA, ldda, dB, lddb, info);
-  AT_CUDA_CHECK(cudaGetLastError());
-}
-
-template<>
-void magmaCholeskySolve<float>(
-    magma_uplo_t uplo, magma_int_t n, magma_int_t nrhs, float* dA, magma_int_t ldda,
-    float* dB, magma_int_t lddb, magma_int_t* info) {
-  MagmaStreamSyncGuard guard;
-  magma_spotrs_gpu(uplo, n, nrhs, dA, ldda, dB, lddb, info);
-  AT_CUDA_CHECK(cudaGetLastError());
-}
-
-template<>
-void magmaCholeskySolve<c10::complex<double>>(
-    magma_uplo_t uplo, magma_int_t n, magma_int_t nrhs, c10::complex<double>* dA, magma_int_t ldda,
-    c10::complex<double>* dB, magma_int_t lddb, magma_int_t* info) {
-  MagmaStreamSyncGuard guard;
-  magma_zpotrs_gpu(uplo, n, nrhs,
-    reinterpret_cast<magmaDoubleComplex*>(dA), ldda,
-    reinterpret_cast<magmaDoubleComplex*>(dB), lddb, info);
-  AT_CUDA_CHECK(cudaGetLastError());
-}
-
-template<>
-void magmaCholeskySolve<c10::complex<float>>(
-    magma_uplo_t uplo, magma_int_t n, magma_int_t nrhs, c10::complex<float>* dA, magma_int_t ldda,
-    c10::complex<float>* dB, magma_int_t lddb, magma_int_t* info) {
-  MagmaStreamSyncGuard guard;
-  magma_cpotrs_gpu(uplo, n, nrhs,
-    reinterpret_cast<magmaFloatComplex*>(dA), ldda,
-    reinterpret_cast<magmaFloatComplex*>(dB), lddb, info);
-  AT_CUDA_CHECK(cudaGetLastError());
-}
-
-template<>
-void magmaCholeskySolveBatched<double>(
-    magma_uplo_t uplo, magma_int_t n, magma_int_t nrhs, double** dA_array, magma_int_t ldda,
-    double** dB_array, magma_int_t lddb, magma_int_t& info, magma_int_t batchsize, const MAGMAQueue& magma_queue) {
-  info = magma_dpotrs_batched(uplo, n, nrhs, dA_array, ldda, dB_array, lddb, batchsize, magma_queue.get_queue());
-  AT_CUDA_CHECK(cudaGetLastError());
-}
-
-template<>
-void magmaCholeskySolveBatched<float>(
-    magma_uplo_t uplo, magma_int_t n, magma_int_t nrhs, float** dA_array, magma_int_t ldda,
-    float** dB_array, magma_int_t lddb, magma_int_t& info, magma_int_t batchsize, const MAGMAQueue& magma_queue) {
-  info = magma_spotrs_batched(uplo, n, nrhs, dA_array, ldda, dB_array, lddb, batchsize, magma_queue.get_queue());
-  AT_CUDA_CHECK(cudaGetLastError());
-}
-
-template<>
-void magmaCholeskySolveBatched<c10::complex<double>>(
-    magma_uplo_t uplo, magma_int_t n, magma_int_t nrhs, c10::complex<double>** dA_array, magma_int_t ldda,
-    c10::complex<double>** dB_array, magma_int_t lddb, magma_int_t& info, magma_int_t batchsize, const MAGMAQueue& magma_queue) {
-  info = magma_zpotrs_batched(uplo, n, nrhs,
-    reinterpret_cast<magmaDoubleComplex**>(dA_array), ldda,
-    reinterpret_cast<magmaDoubleComplex**>(dB_array), lddb, batchsize, magma_queue.get_queue());
-  AT_CUDA_CHECK(cudaGetLastError());
-}
-
-template<>
-void magmaCholeskySolveBatched<c10::complex<float>>(
-    magma_uplo_t uplo, magma_int_t n, magma_int_t nrhs, c10::complex<float>** dA_array, magma_int_t ldda,
-    c10::complex<float>** dB_array, magma_int_t lddb, magma_int_t& info, magma_int_t batchsize, const MAGMAQueue& magma_queue) {
-  info = magma_cpotrs_batched(uplo, n, nrhs,
-    reinterpret_cast<magmaFloatComplex**>(dA_array), ldda,
-    reinterpret_cast<magmaFloatComplex**>(dB_array), lddb, batchsize, magma_queue.get_queue());
-  AT_CUDA_CHECK(cudaGetLastError());
-}
-
 #if defined(USE_ROCM)
 template<>
 void magmaTriangularSolveBatched<double>(
@@ -471,54 +365,7 @@ void magmaTriangularSolveBatched<c10::complex<float>>(
 }
 #endif // defined(USE_ROCM)
 
-template<>
-void magmaSyevd<double>(
-    magma_vec_t jobz, magma_uplo_t uplo, magma_int_t n, double* dA, magma_int_t ldda,
-    double* w, double* wA, magma_int_t ldwa, double* work, magma_int_t lwork, double* rwork,
-    magma_int_t lrwork, magma_int_t* iwork, magma_int_t liwork, magma_int_t* info) {
-  (void)rwork;  // unused
-  (void)lrwork;  // unused
-  MagmaStreamSyncGuard guard;
-  magma_dsyevd_gpu(jobz, uplo, n, dA, ldda, w, wA, ldwa, work, lwork, iwork, liwork, info);
-  AT_CUDA_CHECK(cudaGetLastError());
-}
-
-template<>
-void magmaSyevd<float>(
-    magma_vec_t jobz, magma_uplo_t uplo, magma_int_t n, float* dA, magma_int_t ldda,
-    float* w, float* wA, magma_int_t ldwa, float* work, magma_int_t lwork, float* rwork,
-    magma_int_t lrwork, magma_int_t* iwork, magma_int_t liwork, magma_int_t* info) {
-  (void)rwork;  // unused
-  (void)lrwork;  // unused
-  MagmaStreamSyncGuard guard;
-  magma_ssyevd_gpu(jobz, uplo, n, dA, ldda, w, wA, ldwa, work, lwork, iwork, liwork, info);
-  AT_CUDA_CHECK(cudaGetLastError());
-}
-
-template<>
-void magmaSyevd<c10::complex<double>, double>(
-    magma_vec_t jobz, magma_uplo_t uplo, magma_int_t n, c10::complex<double>* dA, magma_int_t ldda,
-    double* w, c10::complex<double>* wA, magma_int_t ldwa, c10::complex<double>* work, magma_int_t lwork, double* rwork,
-    magma_int_t lrwork, magma_int_t* iwork, magma_int_t liwork, magma_int_t* info) {
-  MagmaStreamSyncGuard guard;
-  magma_zheevd_gpu(
-      jobz, uplo, n, reinterpret_cast<magmaDoubleComplex*>(dA), ldda, w, reinterpret_cast<magmaDoubleComplex*>(wA),
-      ldwa, reinterpret_cast<magmaDoubleComplex*>(work), lwork, rwork, lrwork, iwork, liwork, info);
-  AT_CUDA_CHECK(cudaGetLastError());
-}
-
-template<>
-void magmaSyevd<c10::complex<float>, float>(
-    magma_vec_t jobz, magma_uplo_t uplo, magma_int_t n, c10::complex<float>* dA, magma_int_t ldda,
-    float* w, c10::complex<float>* wA, magma_int_t ldwa, c10::complex<float>* work, magma_int_t lwork, float* rwork,
-    magma_int_t lrwork, magma_int_t* iwork, magma_int_t liwork, magma_int_t* info) {
-  MagmaStreamSyncGuard guard;
-  magma_cheevd_gpu(
-      jobz, uplo, n, reinterpret_cast<magmaFloatComplex*>(dA), ldda, w, reinterpret_cast<magmaFloatComplex*>(wA),
-      ldwa, reinterpret_cast<magmaFloatComplex*>(work), lwork, rwork, lrwork, iwork, liwork, info);
-  AT_CUDA_CHECK(cudaGetLastError());
-}
-
+#if defined(USE_ROCM) || !(defined(CUSOLVER_VERSION) && (CUSOLVER_VERSION >= 11702))
 template<>
 void magmaEig<double>(
     magma_vec_t jobvl, magma_vec_t jobvr, magma_int_t n,
@@ -598,78 +445,7 @@ void magmaEig<c10::complex<float>, float>(
          rwork, info);
   AT_CUDA_CHECK(cudaGetLastError());
 }
-
-template<>
-void magmaLuSolve<double>(
-    magma_int_t n, magma_int_t nrhs, double* dA, magma_int_t ldda, magma_int_t* ipiv,
-    double* dB, magma_int_t lddb, magma_int_t* info, magma_trans_t trans) {
-  MagmaStreamSyncGuard guard;
-  magma_dgetrs_gpu(trans, n, nrhs, dA, ldda, ipiv, dB, lddb, info);
-  AT_CUDA_CHECK(cudaGetLastError());
-}
-
-template<>
-void magmaLuSolve<float>(
-    magma_int_t n, magma_int_t nrhs, float* dA, magma_int_t ldda, magma_int_t* ipiv,
-    float* dB, magma_int_t lddb, magma_int_t* info, magma_trans_t trans) {
-  MagmaStreamSyncGuard guard;
-  magma_sgetrs_gpu(trans, n, nrhs, dA, ldda, ipiv, dB, lddb, info);
-  AT_CUDA_CHECK(cudaGetLastError());
-}
-
-template<>
-void magmaLuSolve<c10::complex<double>>(
-    magma_int_t n, magma_int_t nrhs, c10::complex<double>* dA, magma_int_t ldda, magma_int_t* ipiv,
-    c10::complex<double>* dB, magma_int_t lddb, magma_int_t* info, magma_trans_t trans) {
-  MagmaStreamSyncGuard guard;
-  magma_zgetrs_gpu(trans, n, nrhs, reinterpret_cast<magmaDoubleComplex*>(dA), ldda, ipiv, reinterpret_cast<magmaDoubleComplex*>(dB), lddb, info);
-  AT_CUDA_CHECK(cudaGetLastError());
-}
-
-template<>
-void magmaLuSolve<c10::complex<float>>(
-    magma_int_t n, magma_int_t nrhs, c10::complex<float>* dA, magma_int_t ldda, magma_int_t* ipiv,
-    c10::complex<float>* dB, magma_int_t lddb, magma_int_t* info, magma_trans_t trans) {
-  MagmaStreamSyncGuard guard;
-  magma_cgetrs_gpu(trans, n, nrhs, reinterpret_cast<magmaFloatComplex*>(dA), ldda, ipiv, reinterpret_cast<magmaFloatComplex*>(dB), lddb, info);
-  AT_CUDA_CHECK(cudaGetLastError());
-}
-
-template<>
-void magmaLuSolveBatched<double>(
-    magma_int_t n, magma_int_t nrhs, double** dA_array, magma_int_t ldda, magma_int_t** dipiv_array,
-    double** dB_array, magma_int_t lddb, magma_int_t& info,
-    magma_int_t batchsize, const MAGMAQueue& magma_queue, magma_trans_t trans) {
-  info = magma_dgetrs_batched(trans, n, nrhs, dA_array, ldda, dipiv_array, dB_array, lddb, batchsize, magma_queue.get_queue());
-  AT_CUDA_CHECK(cudaGetLastError());
-}
-
-template<>
-void magmaLuSolveBatched<float>(
-    magma_int_t n, magma_int_t nrhs, float** dA_array, magma_int_t ldda, magma_int_t** dipiv_array,
-    float** dB_array, magma_int_t lddb, magma_int_t& info,
-    magma_int_t batchsize, const MAGMAQueue& magma_queue, magma_trans_t trans) {
- info = magma_sgetrs_batched(trans, n, nrhs, dA_array, ldda, dipiv_array, dB_array, lddb, batchsize, magma_queue.get_queue());
- AT_CUDA_CHECK(cudaGetLastError());
-}
-
-template<>
-void magmaLuSolveBatched<c10::complex<double>>(
-    magma_int_t n, magma_int_t nrhs, c10::complex<double>** dA_array, magma_int_t ldda, magma_int_t** dipiv_array,
-    c10::complex<double>** dB_array, magma_int_t lddb, magma_int_t& info,
-    magma_int_t batchsize, const MAGMAQueue& magma_queue, magma_trans_t trans) {
-  info = magma_zgetrs_batched(trans, n, nrhs, reinterpret_cast<magmaDoubleComplex**>(dA_array), ldda, dipiv_array, reinterpret_cast<magmaDoubleComplex**>(dB_array), lddb, batchsize, magma_queue.get_queue());
-  AT_CUDA_CHECK(cudaGetLastError());
-}
-
-template<>
-void magmaLuSolveBatched<c10::complex<float>>(
-    magma_int_t n, magma_int_t nrhs, c10::complex<float>** dA_array, magma_int_t ldda, magma_int_t** dipiv_array,
-    c10::complex<float>** dB_array, magma_int_t lddb, magma_int_t& info,
-    magma_int_t batchsize, const MAGMAQueue& magma_queue, magma_trans_t trans) {
- info = magma_cgetrs_batched(trans, n, nrhs, reinterpret_cast<magmaFloatComplex**>(dA_array), ldda, dipiv_array, reinterpret_cast<magmaFloatComplex**>(dB_array), lddb, batchsize, magma_queue.get_queue());
- AT_CUDA_CHECK(cudaGetLastError());
-}
+#endif
 
 namespace {
 
@@ -719,14 +495,17 @@ magma_trans_t to_magma(TransposeType trans) {
 
 namespace {
 
-void _warn_once_magma_deprecation(const std::string& op_name) {
+void _warn_once_magma_deprecation(const std::string& op_name, bool force_cusolver = true) {
   if (at::globalContext().linalgPreferredBackend() == at::LinalgBackend::Magma) {
+    std::string warn_force_cusolver = force_cusolver
+      ? " " + op_name + " will try dispatching to cuSOLVER instead. " +
+        "If you see any error messages, please, file an issue on GitHub."
+      : "";
     TORCH_WARN_ONCE(
-      op_name, ": ",
-      "MAGMA, as a linear algebra backend, is deprecated and will be removed ",
-      "in future releases. ",
-      op_name, " will try dispatching to cuSOLVER instead. "
-      "If you see any error messages, please, file an issue on GitHub."
+      op_name, ": "
+      "MAGMA, as a linear algebra backend, is deprecated and will be removed "
+      "in future releases.",
+      warn_force_cusolver
     );
   }
 }
@@ -858,106 +637,58 @@ REGISTER_CUDA_DISPATCH(ldl_solve_stub, &ldl_solve_kernel)
 
 // ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ cholesky_solve ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
-template <typename scalar_t>
-static void apply_cholesky_solve(Tensor& b, Tensor& A, bool upper, int64_t& info) {
-#if !AT_MAGMA_ENABLED()
-TORCH_CHECK(false, "cholesky_solve: MAGMA library not found in "
-    "compilation. Please rebuild with MAGMA.");
-#else
-  magma_uplo_t uplo = upper ? MagmaUpper : MagmaLower;
-
-  auto A_data = A.data_ptr<scalar_t>();
-  auto b_data = b.data_ptr<scalar_t>();
-  magma_int_t n = magma_int_cast(A.size(-2), "A.size(-2)");
-  magma_int_t lda = std::max<magma_int_t>(1, n);
-  magma_int_t nrhs = magma_int_cast(b.size(-1), "b.size(-1)");
-
-  int info_tmp = 0;
-  if (b.dim() == 2) {
-    magmaCholeskySolve<scalar_t>(uplo, n, nrhs, A_data, lda,
-                                 b_data, lda, &info_tmp);
-    info = info_tmp;
+namespace {
+// At the time of writing, the unconditional dispatch
+// to the native cholesky_solve method in cuSOLVER is slow
+// with batched inputs.
+template <bool use_dedicated_kernel_unconditionally = false>
+inline void _cholesky_solve_helper_cuda_cusolver_algo_selector(
+  Tensor& self,
+  const Tensor& A,
+  bool upper) {
+  if constexpr (use_dedicated_kernel_unconditionally) {
+    _cholesky_solve_helper_cuda_cusolver(self, A, upper);
   } else {
-    auto A_mat_stride = matrixStride(A);
-    auto b_mat_stride = matrixStride(b);
-    magma_int_t batch_size = magma_int_cast(batchCount(A), "batchCount");
-
-    scalar_t** A_array;
-    scalar_t** b_array;
-
-    ALLOCATE_ARRAY(A_array, scalar_t*, batch_size);
-    ALLOCATE_ARRAY(b_array, scalar_t*, batch_size);
-
-    // Set up the created arrays
-    for (int64_t i = 0; i < batch_size; i++) {
-      A_array[i] = &A_data[i * A_mat_stride];
-      b_array[i] = &b_data[i * b_mat_stride];
-    }
-
-    MAGMAQueue magma_queue(b.get_device());
-
-    constexpr int64_t batch_limit = 65535;
-    // Compute as many batches of 65535 possible
-    // The number of "mini"-batches are floor(batch_size / batch_limit)
-    // and these cover floor(batch_size / batch_limit) * batch_limit matrix solves
-    int64_t mini_batches = batch_size / batch_limit, mini_idx;
-    for (mini_idx = 0; mini_idx < mini_batches * batch_limit; mini_idx += batch_limit) {
-      scalar_t** A_array_cur = &A_array[mini_idx];
-      scalar_t** b_array_cur = &b_array[mini_idx];
-
-      magmaCholeskySolveBatched<scalar_t>(
-          uplo, n, nrhs, A_array_cur, lda, b_array_cur, lda,
-          info_tmp, batch_limit, magma_queue);
-
-      if (info_tmp != 0) {
-        break;
-      }
-    }
-
-    // Compute whatever is left = batch_size - floor(batch_size / batch_limit) * batch_limit
-    // which concisely is equal to batch_size % batch_limit
-    if (batch_size % batch_limit != 0 && info_tmp == 0) {
-      magmaCholeskySolveBatched<scalar_t>(
-          uplo, n, nrhs, &A_array[mini_idx], lda, &b_array[mini_idx], lda,
-          info_tmp, batch_size % batch_limit, magma_queue);
+    // TODO: cusolverDn<T>potrsBatched only supports nrhs == 1 and does not have good performance.
+    // TODO: Non-batched potrs is too slow in the batched setting compared to two triangular solves.
+    // Non-batched input -> non-batched potrs.
+    // Batched input -> two triangular solves.
+    if (batchCount(self) == 1) {
+      _cholesky_solve_helper_cuda_cusolver(self, A, upper);
+    } else {
+      const auto L = upper
+        ? c10::MaybeOwned<Tensor>::owned(A.mH())
+        : c10::MaybeOwned<Tensor>::borrowed(A);
+      // NOTE: we tolerate redispatch with at::triangular_solve_triangular
+      // because it handles memory layout optimization and conj/neg flags.
+      // IMPORTANT NOTE: `self` and `A` are not processed for kernel calls yet!
+      // Step 1: Solve for Y: L Y = B or U^H Y = B.
+      at::linalg_solve_triangular_out(self, *L, self, /*upper=*/false);
+      // Step 2: Solve for X: L^H X = Y or U X = Y.
+      at::linalg_solve_triangular_out(self, L->mH(), self, /*upper=*/true);
     }
-
-    info = info_tmp;
   }
-#endif
 }
 
-Tensor _cholesky_solve_helper_cuda_magma(const Tensor& self, const Tensor& A, bool upper) {
-  int64_t info = 0;
-  auto self_working_copy = cloneBatchedColumnMajor(self);
-  auto A_working_copy = cloneBatchedColumnMajor(A);
-  AT_DISPATCH_FLOATING_AND_COMPLEX_TYPES(self.scalar_type(), "cholesky_solve_cuda", [&]{
-    apply_cholesky_solve<scalar_t>(self_working_copy, A_working_copy, upper, info);
-  });
-  TORCH_CHECK(info == 0, "MAGMA cholesky_solve : invalid argument: ", -info);
-  return self_working_copy;
+inline void _cholesky_solve_helper_cuda_cusolver_dispatcher(
+    Tensor& self,
+    const Tensor& A,
+    bool upper) {
+  // For now, unconditional dispatch to the dedicated cholesky solve
+  // kernel in cuSOLVER is slow for batched inputs.
+  // TODO: switch once resolved.
+  _cholesky_solve_helper_cuda_cusolver_algo_selector<
+    /*use_dedicated_kernel_unconditionally=*/false
+  >(self, A, upper);
 }
 
-// Todo: cusolverDn<T>potrsBatched only supports nrhs == 1 and does not have good performance.
-//     Batched cholesky_solve is dispatched to magma.
+} // namespace (anonymous)
+
 Tensor _cholesky_solve_helper_cuda(const Tensor& self, const Tensor& A, bool upper) {
-#if defined(USE_LINALG_SOLVER)
-  auto preferred_backend = at::globalContext().linalgPreferredBackend();
-  switch (preferred_backend) {
-    case at::LinalgBackend::Cusolver:
-      return _cholesky_solve_helper_cuda_cusolver(self, A, upper);
-    case at::LinalgBackend::Magma:
-      return _cholesky_solve_helper_cuda_magma(self, A, upper);
-    default:
-      if (batchCount(self) == 1 || !use_magma_) {
-        return _cholesky_solve_helper_cuda_cusolver(self, A, upper);
-      } else {
-        return _cholesky_solve_helper_cuda_magma(self, A, upper);
-      }
-  }
-#else
-  return _cholesky_solve_helper_cuda_magma(self, A, upper);
-#endif
+  _warn_once_magma_deprecation("linalg.cholesky_solve");
+  at::Tensor self_working_copy = cloneBatchedColumnMajor(self);
+  _cholesky_solve_helper_cuda_cusolver_dispatcher(self_working_copy, A, upper);
+  return self_working_copy;
 }
 
 // ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ cholesky ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
@@ -971,81 +702,16 @@ REGISTER_CUDA_DISPATCH(cholesky_stub, &cholesky_kernel)
 
 // ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ cholesky_inverse ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
-/*
-Computes the inverse of a symmetric (Hermitian) positive-definite matrix n-by-n matrix 'input' using the Cholesky solver
-This is an in-place routine, content of 'input' is overwritten.
-'infos' is an int Tensor containing error codes for each matrix in the batched input.
-MAGMA requires 'infos' to reside in CPU memory.
-For more information see MAGMA's documentation for POTRS routine.
-*/
-template <typename scalar_t>
-static void apply_cholesky_inverse(Tensor& input, Tensor& infos, bool upper) {
-#if !AT_MAGMA_ENABLED()
-  TORCH_CHECK(false, "cholesky_inverse: MAGMA library not found in compilation. Please rebuild with MAGMA.");
-#else
-  // magmaCholeskyInverse (magma_dpotri_gpu) is slow because internally
-  // it transfers data several times between GPU and CPU and calls lapack routine on CPU
-  // using magmaCholeskySolveBatched is a lot faster
-  // note that magmaCholeskySolve is also slow
-
-  // 'input' is modified in-place we need to clone it and replace with a diagonal matrix
-  // for apply_cholesky_solve
-  auto input_working_copy = cloneBatchedColumnMajor(input);
-
-  // 'input' tensor has to be a batch of diagonal matrix
-  input.fill_(0);
-  input.diagonal(/*offset=*/0, /*dim1=*/-2, /*dim2=*/-1).fill_(1);
-
-  Tensor result_u, input_u;
-  if (input.dim() == 2) {
-    // unsqueezing here so that the batched version is used
-    result_u = input.unsqueeze(0);
-    input_u = input_working_copy.unsqueeze(0);
-  } else {
-    result_u = input;
-    input_u = input_working_copy;
-  }
-
-  // magma's potrs_batched doesn't take matrix-wise array of ints as an 'info' argument
-  // it returns a single 'magma_int_t'
-  // if info = 0 the operation is successful, if info = -i, the i-th parameter had an illegal value.
-  int64_t info_tmp = 0;
-  apply_cholesky_solve<scalar_t>(result_u, input_u, upper, info_tmp);
-  infos.fill_(info_tmp);
-#endif
-}
-
-// This is a type dispatching helper function for 'apply_cholesky_inverse'
-Tensor& cholesky_inverse_kernel_impl_magma(Tensor &result, Tensor& infos, bool upper) {
-  AT_DISPATCH_FLOATING_AND_COMPLEX_TYPES(result.scalar_type(), "cholesky_inverse_out_cuda", [&]{
-    apply_cholesky_inverse<scalar_t>(result, infos, upper);
-  });
-  return result;
-}
-
-Tensor& cholesky_inverse_kernel_impl(Tensor &result, Tensor& infos, bool upper) {
+Tensor& cholesky_inverse_kernel_impl(Tensor &result, [[maybe_unused]] Tensor& infos, bool upper) {
   // This function calculates the inverse matrix in-place
   // result should be in column major order and contain matrices to invert
-  // the content of result is overwritten by 'apply_cholesky_inverse'
-#if defined(USE_LINALG_SOLVER)
-  auto preferred_backend = at::globalContext().linalgPreferredBackend();
-  switch (preferred_backend) {
-    case at::LinalgBackend::Cusolver:
-      return cholesky_inverse_kernel_impl_cusolver(result, infos, upper);
-    case at::LinalgBackend::Magma:
-      return cholesky_inverse_kernel_impl_magma(result, infos, upper);
-    default:
-      if (batchCount(result) == 1 ||
-          !use_magma_) {
-        return cholesky_inverse_kernel_impl_cusolver(result, infos, upper);
-      } else {
-        return cholesky_inverse_kernel_impl_magma(result, infos, upper);
-      }
-  }
-#else
-  return cholesky_inverse_kernel_impl_magma(result, infos, upper);
-#endif
-
+  // the content of result is overwritten
+  _warn_once_magma_deprecation("linalg.cholesky_inverse");
+  at::Tensor A = cloneBatchedColumnMajor(result);
+  result.fill_(0);
+  result.diagonal(/*offset=*/0, /*dim1=*/-2, /*dim2=*/-1).fill_(1);
+  _cholesky_solve_helper_cuda_cusolver_dispatcher(result, A, upper);
+  return result;
 }
 
 REGISTER_CUDA_DISPATCH(cholesky_inverse_stub, &cholesky_inverse_kernel_impl)
@@ -1408,129 +1074,9 @@ REGISTER_CUDA_DISPATCH(geqrf_stub, &geqrf_kernel)
 
 // ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ linalg_eigh ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
-template <typename scalar_t>
-static void apply_magma_eigh(const Tensor& values, const Tensor& vectors, const Tensor& infos, bool upper, bool compute_eigenvectors) {
-#if !AT_MAGMA_ENABLED()
-  TORCH_CHECK(
-    false,
-    "Calling torch.linalg.eigh/eigvalsh on a CUDA tensor requires compiling ",
-    "PyTorch with MAGMA. Please use PyTorch built with MAGMA support.");
-#else
-  TORCH_INTERNAL_ASSERT_DEBUG_ONLY(values.device() == kCPU);
-  TORCH_INTERNAL_ASSERT_DEBUG_ONLY(infos.device() == kCPU);
-
-  using value_t = typename c10::scalar_value_type<scalar_t>::type;
-
-  magma_uplo_t uplo = upper ? MagmaUpper : MagmaLower;
-  magma_vec_t jobz = compute_eigenvectors ? MagmaVec : MagmaNoVec;
-
-  magma_int_t n = magma_int_cast(vectors.size(-1), "n");
-  auto lda = std::max<magma_int_t>(1, n);
-  auto batch_size = batchCount(vectors);
-
-  auto vectors_stride = matrixStride(vectors);
-  auto values_stride = values.size(-1);
-
-  auto vectors_data = vectors.data_ptr<scalar_t>();
-  auto values_data = values.data_ptr<value_t>();
-  auto infos_data = infos.data_ptr<magma_int_t>();
-
-  scalar_t* wA;
-  ALLOCATE_ARRAY(wA, scalar_t, lda * lda);
-
-  // Run once, first to get the optimum work sizes.
-  // Since we deal with batches of matrices with the same dimensions, doing this outside
-  // the loop saves (batch_size - 1) workspace queries which would provide the same result
-  // and (batch_size - 1) calls to allocate and deallocate workspace using at::empty()
-  magma_int_t lwork = -1;
-  scalar_t wkopt;
-  magma_int_t liwork = -1;
-  magma_int_t iwkopt;
-  magma_int_t lrwork = -1;
-  value_t rwkopt;
-  magmaSyevd<scalar_t, value_t>(jobz, uplo, n, vectors_data, lda, values_data,
-    wA, lda, &wkopt, lwork, &rwkopt, lrwork, &iwkopt, liwork, infos_data);
-
-  scalar_t* work;
-  magma_int_t* iwork;
-  lwork = magma_int_cast(std::max<int64_t>(1, real_impl<scalar_t, value_t>(wkopt)), "work_size");
-  liwork = magma_int_cast(std::max<int64_t>(1, iwkopt), "iwork_size");
-  ALLOCATE_ARRAY(work, scalar_t, lwork);
-  ALLOCATE_ARRAY(iwork, magma_int_t, liwork);
-
-  value_t* rwork = nullptr;
-  c10::Storage storage_rwork;
-  if (vectors.is_complex()) {
-    lrwork = magma_int_cast(std::max<int64_t>(1, rwkopt), "rwork_size");
-    storage_rwork = pin_memory<value_t>(lrwork);
-    rwork = static_cast<value_t*>(storage_rwork.mutable_data());
-  }
-
-  for (decltype(batch_size) i = 0; i < batch_size; i++) {
-    scalar_t* vectors_working_ptr = &vectors_data[i * vectors_stride];
-    value_t* values_working_ptr = &values_data[i * values_stride];
-    magma_int_t* info_working_ptr = &infos_data[i];
-    magmaSyevd<scalar_t, value_t>(jobz, uplo, n, vectors_working_ptr, lda, values_working_ptr,
-      wA, lda, work, lwork, rwork, lrwork, iwork, liwork, info_working_ptr);
-    // The current behaviour for Linear Algebra functions to raise an error if something goes wrong
-    // or input doesn't satisfy some requirement
-    // therefore return early since further computations will be wasted anyway
-    if (*info_working_ptr != 0) {
-      return;
-    }
-  }
-#endif
-}
-
-// This is a type dispatch function for 'apply_magma_eigh'
-// For small inputs result is computed on CPU
-void linalg_eigh_magma(const Tensor& eigenvalues, const Tensor& eigenvectors, const Tensor& infos, bool upper, bool compute_eigenvectors) {
-  // MAGMA just calls LAPACK for eigenvectors.size(-1) <= 128
-  // See https://bitbucket.org/icl/magma/src/e6fdca447bd402693e8b0b950a898b6879bbcc41/src/zheevd_gpu.cpp?at=master#lines-258
-  // in addition lda is ignored breaking 0x0 inputs
-  if (eigenvectors.size(-1) > 128) {
-    // MAGMA requires eigenvalues and infos tensors to reside on CPU
-    Tensor eigenvalues_cpu = eigenvalues.to(kCPU);
-    Tensor infos_cpu = infos.to(kCPU);
-
-    AT_DISPATCH_FLOATING_AND_COMPLEX_TYPES(
-      eigenvectors.scalar_type(), "linalg_eigh_magma", [&] {
-        apply_magma_eigh<scalar_t>(
-            eigenvalues_cpu, eigenvectors, infos_cpu, upper, compute_eigenvectors);
-      });
-
-    // Transfer computed by MAGMA results from CPU to GPU
-    eigenvalues.copy_(eigenvalues_cpu);
-    infos.copy_(infos_cpu);
-  } else { // eigenvectors.size(-1) <= 128
-    // transfer to CPU, compute the result and copy back to GPU
-    // this is faster than going through MAGMA that does the same
-    Tensor eigenvalues_cpu = at::empty_like(eigenvalues, eigenvalues.options().device(kCPU));
-    if (compute_eigenvectors) {
-      Tensor eigenvectors_cpu = at::empty_like(eigenvectors, eigenvectors.options().device(kCPU));
-      at::linalg_eigh_out(eigenvalues_cpu, eigenvectors_cpu, eigenvectors.to(kCPU), upper ? "U" : "L");
-      eigenvectors.copy_(eigenvectors_cpu);
-    } else {
-      at::linalg_eigvalsh_out(eigenvalues_cpu, eigenvectors.to(kCPU), upper ? "U" : "L");
-    }
-    eigenvalues.copy_(eigenvalues_cpu);
-  }
-}
-
 void linalg_eigh_kernel(const Tensor& eigenvalues, const Tensor& eigenvectors, const Tensor& infos, bool upper, bool compute_eigenvectors) {
-#if defined(USE_LINALG_SOLVER)
-  auto preferred_backend = at::globalContext().linalgPreferredBackend();
-  switch (preferred_backend) {
-    case at::LinalgBackend::Magma:
-      linalg_eigh_magma(eigenvalues, eigenvectors, infos, upper, compute_eigenvectors);
-      break;
-    case at::LinalgBackend::Cusolver:
-    default:
-      linalg_eigh_cusolver(eigenvalues, eigenvectors, infos, upper, compute_eigenvectors);
-  }
-#else
-  linalg_eigh_magma(eigenvalues, eigenvectors, infos, upper, compute_eigenvectors);
-#endif
+  _warn_once_magma_deprecation("linalg.eigh");
+  linalg_eigh_cusolver(eigenvalues, eigenvectors, infos, upper, compute_eigenvectors);
 }
 
 REGISTER_CUDA_DISPATCH(linalg_eigh_stub, &linalg_eigh_kernel)
@@ -1543,6 +1089,7 @@ This is an in-place routine, content of 'input', 'values', 'vectors' is overwrit
 'infos' is an int Tensor containing error codes for each matrix in the batched input.
 For more information see MAGMA's documentation for GEEV routine.
 */
+#if defined(USE_ROCM) || !(defined(CUSOLVER_VERSION) && (CUSOLVER_VERSION >= 11702))
 template <typename scalar_t>
 void apply_magma_eig(Tensor& values, Tensor& vectors, Tensor& input, Tensor& infos, bool compute_eigenvectors) {
 #if !AT_MAGMA_ENABLED()
@@ -1620,25 +1167,22 @@ void linalg_eig_magma(Tensor& eigenvalues, Tensor& eigenvectors, Tensor& infos,
   eigenvectors.copy_(eigenvectors_cpu);
   infos.copy_(infos_cpu);
 }
+#endif
+
 void linalg_eig_kernel(Tensor& eigenvalues, Tensor& eigenvectors, Tensor& infos, const Tensor& input, bool compute_eigenvectors) {
+  TORCH_INTERNAL_ASSERT_DEBUG_ONLY(input.is_cuda());
   // This function calculates the non-symmetric eigendecomposition in-place
   // tensors should be in batched column major memory format
   // the content of eigenvalues, eigenvectors and infos is overwritten by 'linalg_eig_magma' or
   // 'linalg_eig_cusolver_xgeev' both geev routines modify the provided input matrix in-place, therefore we need a copy
-
-  TORCH_INTERNAL_ASSERT_DEBUG_ONLY(input.is_cuda());
-#if defined(CUSOLVER_VERSION) && (CUSOLVER_VERSION >= 11702)
-  auto preferred_backend = at::globalContext().linalgPreferredBackend();
-  switch (preferred_backend) {
-    case at::LinalgBackend::Cusolver:
-    default:
-      linalg_eig_cusolver_xgeev(eigenvalues, eigenvectors, input, infos, compute_eigenvectors);
-      return;
-    case at::LinalgBackend::Magma:
-      break; // MAGMA path handled below
-  }
-#endif
+#if !defined(USE_ROCM) && defined(CUSOLVER_VERSION) && (CUSOLVER_VERSION >= 11702)
+  _warn_once_magma_deprecation("linalg.eig");
+  linalg_eig_cusolver_xgeev(eigenvalues, eigenvectors, input, infos, compute_eigenvectors);
+#else
+  // hipSolver does not have `geev`
+  _warn_once_magma_deprecation("linalg.eig", /*force_cusolver=*/false);
   linalg_eig_magma(eigenvalues, eigenvectors, infos, input, compute_eigenvectors);
+#endif
 }
 
 REGISTER_CUDA_DISPATCH(linalg_eig_stub, &linalg_eig_kernel)
@@ -1666,160 +1210,6 @@ REGISTER_CUDA_DISPATCH(svd_stub, &svd_kernel)
 
 // ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ lu_solve ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
-/*
-  Solves the matrix equation A X = B
-  X and B are n-by-nrhs matrices, A is represented using the LU factorization.
-  This is an in-place routine, content of `B` is overwritten.
-  This is a "looped" variant for calling single input MAGMA function on batched input.
-
-  Args:
-  * `LU` - [in] the LU factorization of matrix A (see at::linalg_lu_factor)
-  * `pivots` - [in] the pivot indices (see at::linalg_lu_factor)
-  * `B` -  [in] the right hand side matrix B
-           [out] the solution matrix X
-
-  For further details, please see the MAGMA documentation for magma_dgetrs_gpu.
-*/
-template <typename scalar_t>
-static void apply_lu_solve_looped_magma(const Tensor& LU, const Tensor& pivots, const Tensor& B, TransposeType transpose) {
-#if !AT_MAGMA_ENABLED()
-  TORCH_CHECK(
-      false,
-      "Calling linalg.lu_solve on a CUDA tensor requires compiling ",
-      "PyTorch with MAGMA. Please rebuild with MAGMA.");
-#else
-  auto trans = to_magma(transpose);
-  auto b_data = B.data_ptr<scalar_t>();
-  auto lu_data = LU.data_ptr<scalar_t>();
-
-  // MAGMA requires pivots to be a CPU tensor
-  Tensor pivots_cpu = pivots.cpu();
-  auto pivots_data = pivots_cpu.data_ptr<magma_int_t>();
-
-  auto b_stride = matrixStride(B);
-  auto lu_stride = LU.dim() > 2 ? LU.stride(-3) : 0;
-  auto pivots_stride = pivots_cpu.dim() > 1 ? pivots_cpu.stride(-2) : 0;
-  auto batch_size = batchCount(B);
-
-  magma_int_t n = magma_int_cast(LU.size(-2), "n");
-  magma_int_t nrhs = magma_int_cast(B.size(-1), "nrhs");
-  auto leading_dimension = std::max<magma_int_t>(1, n);
-
-  // LU and pivots tensors can be broadcasted to B
-  // here we construct a helper indexing tensor to linearly index into lu and pivots
-  IntArrayRef lu_batch_shape(LU.sizes().data(), LU.dim() - 2);
-  IntArrayRef b_batch_shape(B.sizes().data(), B.dim() - 2);
-  BroadcastLinearIndices lu_index(
-      batchCount(LU), lu_batch_shape, b_batch_shape);
-
-  int info = 0;
-  for (decltype(batch_size) i = 0; i < batch_size; i++) {
-    int64_t lu_index_i = lu_index(i);
-    scalar_t* b_working_ptr = &b_data[i * b_stride];
-    scalar_t* lu_working_ptr = &lu_data[lu_index_i * lu_stride];
-    int* pivots_working_ptr = &pivots_data[lu_index_i * pivots_stride];
-
-    magmaLuSolve<scalar_t>(n, nrhs, lu_working_ptr, leading_dimension, pivots_working_ptr, b_working_ptr, leading_dimension, &info, trans);
-
-    // info from magmaLuSolve only reports if the i-th parameter is wrong
-    // so we don't need to check it all the time
-    TORCH_INTERNAL_ASSERT_DEBUG_ONLY(info == 0);
-  }
-#endif
-}
-
-/*
-  Solves the matrix equation A X = B
-  X and B are n-by-nrhs matrices, A is represented using the LU factorization.
-  This is an in-place routine, content of `B` is overwritten.
-  This is a specialized batched variant, it is expected to be faster than the "looped" version only for small inputs.
-
-  Args:
-  * `lu` - [in] the LU factorization of matrix A (see at::linalg_lu_factor)
-  * `pivots` - [in] the pivot indices (see at::linalg_lu_factor)
-  * `B` -  [in] the right hand side matrix B
-           [out] the solution matrix X
-
-  For further details, please see the MAGMA documentation for magma_dgetrs_batched.
-*/
-template <typename scalar_t>
-static void apply_lu_solve_batched_magma(const Tensor& LU, const Tensor& pivots, const Tensor& B, TransposeType transpose) {
-#if !AT_MAGMA_ENABLED()
-  TORCH_CHECK(
-      false,
-      "Calling linalg.lu_solve on a CUDA tensor requires compiling ",
-      "PyTorch with MAGMA. Please rebuild with MAGMA.");
-#else
-  TORCH_INTERNAL_ASSERT(batchCount(B) == batchCount(LU), "batch_size of LU and B must be the same");
-  TORCH_INTERNAL_ASSERT(batchCount(LU) == batchCount(pivots.unsqueeze(-1)), "batch_size of LU and pivots must be the same");
-  auto trans = to_magma(transpose);
-  auto b_data = B.data_ptr<scalar_t>();
-  auto lu_data = LU.data_ptr<scalar_t>();
-
-  magma_int_t n = magma_int_cast(LU.size(-2), "n");
-  magma_int_t nrhs = magma_int_cast(B.size(-1), "nrhs");
-  auto leading_dimension = std::max<magma_int_t>(1, n);
-
-  auto pivots_data = pivots.data_ptr<magma_int_t>();
-
-  auto b_stride = matrixStride(B);
-  auto lu_stride = matrixStride(LU);
-  auto pivots_stride = pivots.size(-1);
-  magma_int_t batch_size = magma_int_cast(batchCount(B), "batchCount");
-
-  magma_int_t** pivots_array;
-  scalar_t** lu_array;
-  scalar_t** b_array;
-
-  ALLOCATE_ARRAY(pivots_array, magma_int_t*, batch_size);
-  ALLOCATE_ARRAY(lu_array, scalar_t*, batch_size);
-  ALLOCATE_ARRAY(b_array, scalar_t*, batch_size);
-
-  for (int64_t i = 0; i < batch_size; i++) {
-    pivots_array[i] = &pivots_data[i * pivots_stride];
-    b_array[i] = &b_data[i * b_stride];
-    lu_array[i] = &lu_data[i * lu_stride];
-  }
-
-  MAGMAQueue magma_queue(B.get_device());
-
-  // Compute the result in batches of 65535
-  // that is the maximum allowed number for batch_size in MAGMA
-  constexpr int64_t batch_limit = 65535;
-
-  for (int64_t mini_idx = 0; mini_idx < batch_size; mini_idx += batch_limit) {
-    int64_t nbatches = std::min(batch_limit, batch_size - mini_idx);
-    scalar_t** lu_array_cur = &lu_array[mini_idx];
-    scalar_t** b_array_cur = &b_array[mini_idx];
-    magma_int_t** pivots_array_cur = &pivots_array[mini_idx];
-
-    int info;
-    magmaLuSolveBatched<scalar_t>(
-        n, nrhs, lu_array_cur, leading_dimension,
-        pivots_array_cur, b_array_cur, leading_dimension,
-        info, nbatches, magma_queue, trans);
-
-    // info from magmaLuSolveBatched only reports if the i-th parameter is wrong
-    // so we don't need to check it all the time
-    TORCH_INTERNAL_ASSERT_DEBUG_ONLY(info == 0);
-  }
-#endif
-}
-
-static void lu_solve_batched_magma(const Tensor& LU, const Tensor& pivots, const Tensor& B, TransposeType trans) {
-  // There is a bug in MAGMA when TransposeType is transpose or conj-transpose.
-  TORCH_INTERNAL_ASSERT(trans == TransposeType::NoTranspose);
-  AT_DISPATCH_FLOATING_AND_COMPLEX_TYPES(LU.scalar_type(), "lu_solve_batched_magma", [&]{
-    apply_lu_solve_batched_magma<scalar_t>(LU, pivots, B, trans);
-  });
-}
-
-static void lu_solve_looped_magma(const Tensor& LU, const Tensor& pivots, const Tensor& B, TransposeType trans) {
-  AT_DISPATCH_FLOATING_AND_COMPLEX_TYPES(LU.scalar_type(), "lu_solve_looped_magma", [&]{
-    apply_lu_solve_looped_magma<scalar_t>(LU, pivots, B, trans);
-  });
-}
-
 c10::MaybeOwned<Tensor> maybe_expand_lu(const Tensor& B, const Tensor& LU) {
   // B and LU have the same number of dimensions
   if (batchCount(B) != batchCount(LU)) {
@@ -1853,9 +1243,6 @@ static void lu_solve_kernel(const Tensor& LU, const Tensor& pivots, const Tensor
   auto b = batchCount(B);
   auto n = LU.size(-2);
   auto k = B.size(-1);
-  // magma implementation of LU solve cannot handle a b tensor with last dim > 1024
-  // See https://bitbucket.org/icl/magma/issues/19/dgesv_batched-dgetrs_batched-fails-for
-  bool over_batched_magma_dim_limit = k > 1024;
   // heuristics determined from tests discussed in https://github.com/pytorch/pytorch/pull/72935
 
   // Computes X = U^{-1}L^{-1}P^T B via triangular solves
@@ -1900,52 +1287,35 @@ static void lu_solve_kernel(const Tensor& LU, const Tensor& pivots, const Tensor
     }
   };
 
+  _warn_once_magma_deprecation("linalg.lu_solve");
+
 #ifdef USE_LINALG_SOLVER
   auto lu_solve_batched_cublas_fn = [](const Tensor& LU, const Tensor& pivots, const Tensor& B, TransposeType trans) {
     auto LU_ = maybe_expand_lu(B, LU);
     auto pivots_ = maybe_expand_pivots(B, pivots);
     lu_solve_batched_cublas(*LU_, *pivots_, B, trans);
   };
-#endif
-
-  auto lu_solve_batched_magma_fn = [](const Tensor& LU, const Tensor& pivots, const Tensor& B, TransposeType trans) {
-    auto LU_ = maybe_expand_lu(B, LU);
-    auto pivots_ = maybe_expand_pivots(B, pivots);
-    lu_solve_batched_magma(*LU_, *pivots_, B, trans);
-  };
-
 
   // Preferred Backend
   auto preferred_backend = at::globalContext().linalgPreferredBackend();
-#ifdef USE_LINALG_SOLVER
   if (preferred_backend == at::LinalgBackend::Cusolver) {
+    // TODO: Re-eval this condition
     if (b <= 2 && n >= 64) {
       lu_solve_looped_cusolver(LU, pivots, B, trans);
     } else {
       lu_solve_batched_cublas_fn(LU, pivots, B, trans);
     }
     return;
-  } else
-#endif // ifdef USE_LINALG_SOLVER
-  if (preferred_backend == at::LinalgBackend::Magma) {
-    // Looped magma is very slow, but batched magma is buggy in these two cases
-    if (!over_batched_magma_dim_limit && trans == TransposeType::NoTranspose) {
-      lu_solve_batched_magma_fn(LU, pivots, B, trans);
-    }
-    else {
-      lu_solve_looped_magma(LU, pivots, B, trans);
-    }
-    return;
   }
 
+  // TODO: Re-eval this heuristic
   // Heuristic
   //if (n == k) {
   // if (k <= 16) batched_cublas
   // else solve_triag
   //} else {
   //if (n <= 8) {
-  // if (k >= 256 && NoTranspose) batched_magma
-  // else batched_cusolver
+  // batched_cublas
   //} else if (n <= 32) {
   //  b <= 2 looped_cusolver
   //  k <= 8 batched_cusolver
@@ -1963,8 +1333,6 @@ static void lu_solve_kernel(const Tensor& LU, const Tensor& pivots, const Tensor
   //}
   //}
 
-
-#ifdef USE_LINALG_SOLVER
   // Particular case when multiplying A^{-1}B where B is square
   // In this case doing two triangular solves is almost always fastest
   if (n == k) {
@@ -1976,31 +1344,27 @@ static void lu_solve_kernel(const Tensor& LU, const Tensor& pivots, const Tensor
     return;
   }
 
-if (n <= 8) {
-  if (use_magma_ && !over_batched_magma_dim_limit && trans == TransposeType::NoTranspose && k >= 256) {
-    lu_solve_batched_magma_fn(LU, pivots, B, trans);
-  } else {
-    lu_solve_batched_cublas_fn(LU, pivots, B, trans);
-  }
-} else if (n <= 64) {
-  if (b <= 2 && (k <= 64 || trans != TransposeType::NoTranspose || n <= 32)) {
-    lu_solve_looped_cusolver(LU, pivots, B, trans);
-  } else if (k <= 8) {
+  if (n <= 8) {
     lu_solve_batched_cublas_fn(LU, pivots, B, trans);
-  } else {
-    lu_solve_triangular(LU, pivots, B, trans);
-  }
-} else if (n <= 128) {
-  if (b <= 2 && k <= 2)  {
-    lu_solve_looped_cusolver(LU, pivots, B, trans);
-  } else if (k <= 2)  {
-    lu_solve_batched_cublas_fn(LU, pivots, B, trans);
-  } else {
+  } else if (n <= 64) {
+    if (b <= 2 && (k <= 64 || trans != TransposeType::NoTranspose || n <= 32)) {
+      lu_solve_looped_cusolver(LU, pivots, B, trans);
+    } else if (k <= 8) {
+      lu_solve_batched_cublas_fn(LU, pivots, B, trans);
+    } else {
+      lu_solve_triangular(LU, pivots, B, trans);
+    }
+  } else if (n <= 128) {
+    if (b <= 2 && k <= 2)  {
+      lu_solve_looped_cusolver(LU, pivots, B, trans);
+    } else if (k <= 2)  {
+      lu_solve_batched_cublas_fn(LU, pivots, B, trans);
+    } else {
+      lu_solve_triangular(LU, pivots, B, trans);
+    }
+  } else { // n > 128
     lu_solve_triangular(LU, pivots, B, trans);
   }
-} else { // n > 128
-  lu_solve_triangular(LU, pivots, B, trans);
-}
 #else
   // No cublas or cusolver
   // lu_solve_triangular is almost always best
diff --git a/aten/src/ATen/native/cuda/linalg/BatchLinearAlgebraLib.cpp b/aten/src/ATen/native/cuda/linalg/BatchLinearAlgebraLib.cpp
index 05111e2be9932..65c89cb709790 100644
--- a/aten/src/ATen/native/cuda/linalg/BatchLinearAlgebraLib.cpp
+++ b/aten/src/ATen/native/cuda/linalg/BatchLinearAlgebraLib.cpp
@@ -899,53 +899,29 @@ static void apply_cholesky_cusolver_potrsBatched(Tensor& self_working_copy, cons
   );
 }
 
-Tensor _cholesky_solve_helper_cuda_cusolver(const Tensor& self, const Tensor& A, bool upper) {
+void _cholesky_solve_helper_cuda_cusolver(Tensor& self, const Tensor& A, bool upper) {
   const int64_t batch_size = batchCount(self);
   at::Tensor infos = at::zeros({1}, self.options().dtype(at::kInt));
-  at::Tensor self_working_copy = cloneBatchedColumnMajor(self);
   at::Tensor A_column_major_copy = cloneBatchedColumnMajor(A);
 
-  const int64_t nrhs = self_working_copy.size(-1);
+  const int64_t nrhs = self.size(-1);
 
   // cusolverDn<t>potrsBatched only supports nrhs == 1
   if (batch_size > 1 && nrhs == 1) {
     AT_DISPATCH_FLOATING_AND_COMPLEX_TYPES(self.scalar_type(), "cholesky_cuda_potrs_batched", [&] {
-      apply_cholesky_cusolver_potrsBatched<scalar_t>(self_working_copy, A_column_major_copy, upper, infos);
+      apply_cholesky_cusolver_potrsBatched<scalar_t>(self, A_column_major_copy, upper, infos);
     });
   } else {
     AT_DISPATCH_FLOATING_AND_COMPLEX_TYPES(self.scalar_type(), "cholesky_cuda_potrs", [&] {
-      apply_cholesky_cusolver_potrs<scalar_t>(self_working_copy, A_column_major_copy, upper, infos);
+      apply_cholesky_cusolver_potrs<scalar_t>(self, A_column_major_copy, upper, infos);
     });
   }
 
   // info from potrs and potrsBatched only report if the i-th parameter is wrong, not about the matrix singularity, etc.
   // So we don't need to check it all the time.
   TORCH_INTERNAL_ASSERT_DEBUG_ONLY(infos.item().toInt() == 0);
-
-  return self_working_copy;
-}
-
-
-void _cholesky_inverse_cusolver_potrs_based(Tensor& result, Tensor& infos, bool upper) {
-  at::Tensor input_working_copy = cloneBatchedColumnMajor(result);
-  at::Tensor infos_gpu = at::zeros({1}, result.options().dtype(at::kInt));
-  result.fill_(0);
-  result.diagonal(/*offset=*/0, /*dim1=*/-2, /*dim2=*/-1).fill_(1);
-  AT_DISPATCH_FLOATING_AND_COMPLEX_TYPES(result.scalar_type(), "cholesky_cuda_potri", [&] {
-    apply_cholesky_cusolver_potrs<scalar_t>(result, input_working_copy, upper, infos_gpu);
-  });
-
-  // Debug only: info of cusolver potrs only check if the i-th parameter is wrong
-  // Function argument `infos` is a CPU tensor, the following copy will cause a device-host sync.
-  // infos.copy_(infos_gpu);
-}
-
-Tensor& cholesky_inverse_kernel_impl_cusolver(Tensor &result, Tensor& infos, bool upper) {
-  _cholesky_inverse_cusolver_potrs_based(result, infos, upper);
-  return result;
 }
 
-
 /*
   The geqrf function computes the QR decomposition of a m x n matrix A.
 
@@ -1402,63 +1378,6 @@ static void apply_syevd(const Tensor& values, const Tensor& vectors, const Tenso
   }
 }
 
-template <typename scalar_t>
-static void apply_syevj(const Tensor& values, const Tensor& vectors, const Tensor& infos, bool upper, bool compute_eigenvectors) {
-  using value_t = typename c10::scalar_value_type<scalar_t>::type;
-
-  cublasFillMode_t uplo = upper ? CUBLAS_FILL_MODE_UPPER : CUBLAS_FILL_MODE_LOWER;
-  cusolverEigMode_t jobz = compute_eigenvectors ? CUSOLVER_EIG_MODE_VECTOR : CUSOLVER_EIG_MODE_NOVECTOR;
-
-  int n = cuda_int_cast(vectors.size(-1), "n");
-  int lda = std::max<int>(1, n);
-  auto batch_size = batchCount(vectors);
-
-  auto vectors_stride = matrixStride(vectors);
-  auto values_stride = values.size(-1);
-
-  auto vectors_data = vectors.data_ptr<scalar_t>();
-  auto values_data = values.data_ptr<value_t>();
-  auto infos_data = infos.data_ptr<int>();
-
-  // syevj_params controls the numerical accuracy of syevj
-  // by default the tolerance is set to machine accuracy
-  // the maximum number of iteration of Jacobi method by default is 100
-  // cuSOLVER documentations says: "15 sweeps are good enough to converge to machine accuracy"
-  // LAPACK has SVD routine based on similar Jacobi algorithm (gesvj) and there a maximum of 30 iterations is set
-  // Let's use the default values for now
-  syevjInfo_t syevj_params;
-  TORCH_CUSOLVER_CHECK(cusolverDnCreateSyevjInfo(&syevj_params));
-
-  // get the optimal work size and allocate workspace tensor
-  int lwork;
-  at::cuda::solver::syevj_bufferSize<scalar_t>(
-      at::cuda::getCurrentCUDASolverDnHandle(), jobz, uplo, n, vectors_data, lda, values_data, &lwork, syevj_params);
-
-  for (decltype(batch_size) i = 0; i < batch_size; i++) {
-    scalar_t* vectors_working_ptr = &vectors_data[i * vectors_stride];
-    value_t* values_working_ptr = &values_data[i * values_stride];
-    int* info_working_ptr = &infos_data[i];
-    auto handle = at::cuda::getCurrentCUDASolverDnHandle();
-
-    // allocate workspace storage on device
-    auto& allocator = *at::cuda::getCUDADeviceAllocator();
-    auto work_data = allocator.allocate(sizeof(scalar_t) * lwork);
-    at::cuda::solver::syevj<scalar_t>(
-        handle,
-        jobz,
-        uplo,
-        n,
-        vectors_working_ptr,
-        lda,
-        values_working_ptr,
-        static_cast<scalar_t*>(work_data.get()),
-        lwork,
-        info_working_ptr,
-        syevj_params);
-  }
-  TORCH_CUSOLVER_CHECK(cusolverDnDestroySyevjInfo(syevj_params));
-}
-
 template <typename scalar_t>
 static void apply_syevj_batched(const Tensor& values, const Tensor& vectors, const Tensor& infos, bool upper, bool compute_eigenvectors) {
   using value_t = typename c10::scalar_value_type<scalar_t>::type;
@@ -1574,12 +1493,6 @@ static void linalg_eigh_cusolver_syevd(const Tensor& eigenvalues, const Tensor&
   });
 }
 
-static void linalg_eigh_cusolver_syevj(const Tensor& eigenvalues, const Tensor& eigenvectors, const Tensor& infos, bool upper, bool compute_eigenvectors) {
-  AT_DISPATCH_FLOATING_AND_COMPLEX_TYPES(eigenvectors.scalar_type(), "linalg_eigh_cuda", [&] {
-    apply_syevj<scalar_t>(eigenvalues, eigenvectors, infos, upper, compute_eigenvectors);
-  });
-}
-
 static void linalg_eigh_cusolver_syevj_batched(const Tensor& eigenvalues, const Tensor& eigenvectors, const Tensor& infos, bool upper, bool compute_eigenvectors) {
   AT_DISPATCH_FLOATING_AND_COMPLEX_TYPES(eigenvectors.scalar_type(), "linalg_eigh_cuda", [&] {
     apply_syevj_batched<scalar_t>(eigenvalues, eigenvectors, infos, upper, compute_eigenvectors);
@@ -1602,17 +1515,9 @@ void linalg_eigh_cusolver(const Tensor& eigenvalues, const Tensor& eigenvectors,
 #endif // ROCSOLVER_SYEVD_BATCHED_ENABLED
     linalg_eigh_cusolver_syevd(eigenvalues, eigenvectors, infos, upper, compute_eigenvectors);
 #else // not USE_ROCM
-  if (batchCount(eigenvectors) > 1 && eigenvectors.size(-1) <= 32) {
-    // Use syevjBatched for batched matrix operation when matrix size <= 32
-    // See https://github.com/pytorch/pytorch/pull/53040#issuecomment-788264724
-    linalg_eigh_cusolver_syevj_batched(eigenvalues, eigenvectors, infos, upper, compute_eigenvectors);
-  } else if (eigenvectors.scalar_type() == at::kFloat && eigenvectors.size(-1) >= 32 && eigenvectors.size(-1) <= 512) {
-    // syevj is better than syevd for float32 dtype and matrix sizes 32x32 - 512x512
-    // See https://github.com/pytorch/pytorch/pull/53040#issuecomment-788264724
-    linalg_eigh_cusolver_syevj(eigenvalues, eigenvectors, infos, upper, compute_eigenvectors);
-  } else {
-    linalg_eigh_cusolver_syevd(eigenvalues, eigenvectors, infos, upper, compute_eigenvectors);
-  }
+
+  linalg_eigh_cusolver_syevj_batched(eigenvalues, eigenvectors, infos, upper, compute_eigenvectors);
+
 #endif
 }
 
diff --git a/aten/src/ATen/native/cuda/linalg/BatchLinearAlgebraLib.h b/aten/src/ATen/native/cuda/linalg/BatchLinearAlgebraLib.h
index c59649d246b70..2c1937fb27e9a 100644
--- a/aten/src/ATen/native/cuda/linalg/BatchLinearAlgebraLib.h
+++ b/aten/src/ATen/native/cuda/linalg/BatchLinearAlgebraLib.h
@@ -59,8 +59,7 @@ void svd_cusolver(const Tensor& A,
 
 // entrance of calculations of `cholesky` using cusolver potrf and potrfBatched
 void cholesky_helper_cusolver(const Tensor& input, bool upper, const Tensor& info);
-Tensor _cholesky_solve_helper_cuda_cusolver(const Tensor& self, const Tensor& A, bool upper);
-Tensor& cholesky_inverse_kernel_impl_cusolver(Tensor& result, Tensor& infos, bool upper);
+void _cholesky_solve_helper_cuda_cusolver(Tensor& self, const Tensor& A, bool upper);
 
 void geqrf_cusolver(const Tensor& input, const Tensor& tau);
 void ormqr_cusolver(const Tensor& input, const Tensor& tau, const Tensor& other, bool left, bool transpose);
diff --git a/aten/src/ATen/native/cuda/linalg/CUDASolver.cpp b/aten/src/ATen/native/cuda/linalg/CUDASolver.cpp
index 44fcd17c99841..3f7a13424294b 100644
--- a/aten/src/ATen/native/cuda/linalg/CUDASolver.cpp
+++ b/aten/src/ATen/native/cuda/linalg/CUDASolver.cpp
@@ -1208,170 +1208,6 @@ void syevd<c10::complex<double>, double>(
       info));
 }
 
-template <>
-void syevj_bufferSize<float>(
-    cusolverDnHandle_t handle,
-    cusolverEigMode_t jobz,
-    cublasFillMode_t uplo,
-    int n,
-    const float* A,
-    int lda,
-    const float* W,
-    int* lwork,
-    syevjInfo_t params) {
-  TORCH_CUSOLVER_CHECK(cusolverDnSsyevj_bufferSize(
-      handle, jobz, uplo, n, A, lda, W, lwork, params));
-}
-
-template <>
-void syevj_bufferSize<double>(
-    cusolverDnHandle_t handle,
-    cusolverEigMode_t jobz,
-    cublasFillMode_t uplo,
-    int n,
-    const double* A,
-    int lda,
-    const double* W,
-    int* lwork,
-    syevjInfo_t params) {
-  TORCH_CUSOLVER_CHECK(cusolverDnDsyevj_bufferSize(
-      handle, jobz, uplo, n, A, lda, W, lwork, params));
-}
-
-template <>
-void syevj_bufferSize<c10::complex<float>, float>(
-    cusolverDnHandle_t handle,
-    cusolverEigMode_t jobz,
-    cublasFillMode_t uplo,
-    int n,
-    const c10::complex<float>* A,
-    int lda,
-    const float* W,
-    int* lwork,
-    syevjInfo_t params) {
-  TORCH_CUSOLVER_CHECK(cusolverDnCheevj_bufferSize(
-      handle,
-      jobz,
-      uplo,
-      n,
-      reinterpret_cast<const cuComplex*>(A),
-      lda,
-      W,
-      lwork,
-      params));
-}
-
-template <>
-void syevj_bufferSize<c10::complex<double>, double>(
-    cusolverDnHandle_t handle,
-    cusolverEigMode_t jobz,
-    cublasFillMode_t uplo,
-    int n,
-    const c10::complex<double>* A,
-    int lda,
-    const double* W,
-    int* lwork,
-    syevjInfo_t params) {
-  TORCH_CUSOLVER_CHECK(cusolverDnZheevj_bufferSize(
-      handle,
-      jobz,
-      uplo,
-      n,
-      reinterpret_cast<const cuDoubleComplex*>(A),
-      lda,
-      W,
-      lwork,
-      params));
-}
-
-template <>
-void syevj<float>(
-    cusolverDnHandle_t handle,
-    cusolverEigMode_t jobz,
-    cublasFillMode_t uplo,
-    int n,
-    float* A,
-    int lda,
-    float* W,
-    float* work,
-    int lwork,
-    int* info,
-    syevjInfo_t params) {
-  TORCH_CUSOLVER_CHECK(cusolverDnSsyevj(
-      handle, jobz, uplo, n, A, lda, W, work, lwork, info, params));
-}
-
-template <>
-void syevj<double>(
-    cusolverDnHandle_t handle,
-    cusolverEigMode_t jobz,
-    cublasFillMode_t uplo,
-    int n,
-    double* A,
-    int lda,
-    double* W,
-    double* work,
-    int lwork,
-    int* info,
-    syevjInfo_t params) {
-  TORCH_CUSOLVER_CHECK(cusolverDnDsyevj(
-      handle, jobz, uplo, n, A, lda, W, work, lwork, info, params));
-}
-
-template <>
-void syevj<c10::complex<float>, float>(
-    cusolverDnHandle_t handle,
-    cusolverEigMode_t jobz,
-    cublasFillMode_t uplo,
-    int n,
-    c10::complex<float>* A,
-    int lda,
-    float* W,
-    c10::complex<float>* work,
-    int lwork,
-    int* info,
-    syevjInfo_t params) {
-  TORCH_CUSOLVER_CHECK(cusolverDnCheevj(
-      handle,
-      jobz,
-      uplo,
-      n,
-      reinterpret_cast<cuComplex*>(A),
-      lda,
-      W,
-      reinterpret_cast<cuComplex*>(work),
-      lwork,
-      info,
-      params));
-}
-
-template <>
-void syevj<c10::complex<double>, double>(
-    cusolverDnHandle_t handle,
-    cusolverEigMode_t jobz,
-    cublasFillMode_t uplo,
-    int n,
-    c10::complex<double>* A,
-    int lda,
-    double* W,
-    c10::complex<double>* work,
-    int lwork,
-    int* info,
-    syevjInfo_t params) {
-  TORCH_CUSOLVER_CHECK(cusolverDnZheevj(
-      handle,
-      jobz,
-      uplo,
-      n,
-      reinterpret_cast<cuDoubleComplex*>(A),
-      lda,
-      W,
-      reinterpret_cast<cuDoubleComplex*>(work),
-      lwork,
-      info,
-      params));
-}
-
 template <>
 void syevjBatched_bufferSize<float>(
     cusolverDnHandle_t handle,
diff --git a/aten/src/ATen/native/cuda/linalg/CUDASolver.h b/aten/src/ATen/native/cuda/linalg/CUDASolver.h
index 2e72853953ed4..b2f31f768cb90 100644
--- a/aten/src/ATen/native/cuda/linalg/CUDASolver.h
+++ b/aten/src/ATen/native/cuda/linalg/CUDASolver.h
@@ -484,51 +484,6 @@ template <>
 void syevd<c10::complex<double>, double>(
     CUDASOLVER_SYEVD_ARGTYPES(c10::complex<double>, double));
 
-#define CUDASOLVER_SYEVJ_BUFFERSIZE_ARGTYPES(scalar_t, value_t)             \
-  cusolverDnHandle_t handle, cusolverEigMode_t jobz, cublasFillMode_t uplo, \
-      int n, const scalar_t *A, int lda, const value_t *W, int *lwork,      \
-      syevjInfo_t params
-
-template <class scalar_t, class value_t = scalar_t>
-void syevj_bufferSize(CUDASOLVER_SYEVJ_BUFFERSIZE_ARGTYPES(scalar_t, value_t)) {
-  static_assert(false&&sizeof(scalar_t),
-      "at::cuda::solver::syevj_bufferSize: not implemented");
-}
-
-template <>
-void syevj_bufferSize<float>(
-    CUDASOLVER_SYEVJ_BUFFERSIZE_ARGTYPES(float, float));
-template <>
-void syevj_bufferSize<double>(
-    CUDASOLVER_SYEVJ_BUFFERSIZE_ARGTYPES(double, double));
-template <>
-void syevj_bufferSize<c10::complex<float>, float>(
-    CUDASOLVER_SYEVJ_BUFFERSIZE_ARGTYPES(c10::complex<float>, float));
-template <>
-void syevj_bufferSize<c10::complex<double>, double>(
-    CUDASOLVER_SYEVJ_BUFFERSIZE_ARGTYPES(c10::complex<double>, double));
-
-#define CUDASOLVER_SYEVJ_ARGTYPES(scalar_t, value_t)                        \
-  cusolverDnHandle_t handle, cusolverEigMode_t jobz, cublasFillMode_t uplo, \
-      int n, scalar_t *A, int lda, value_t *W, scalar_t *work, int lwork,   \
-      int *info, syevjInfo_t params
-
-template <class scalar_t, class value_t = scalar_t>
-void syevj(CUDASOLVER_SYEVJ_ARGTYPES(scalar_t, value_t)) {
-  static_assert(false&&sizeof(scalar_t), "at::cuda::solver::syevj: not implemented");
-}
-
-template <>
-void syevj<float>(CUDASOLVER_SYEVJ_ARGTYPES(float, float));
-template <>
-void syevj<double>(CUDASOLVER_SYEVJ_ARGTYPES(double, double));
-template <>
-void syevj<c10::complex<float>, float>(
-    CUDASOLVER_SYEVJ_ARGTYPES(c10::complex<float>, float));
-template <>
-void syevj<c10::complex<double>, double>(
-    CUDASOLVER_SYEVJ_ARGTYPES(c10::complex<double>, double));
-
 #define CUDASOLVER_SYEVJ_BATCHED_BUFFERSIZE_ARGTYPES(scalar_t, value_t)     \
   cusolverDnHandle_t handle, cusolverEigMode_t jobz, cublasFillMode_t uplo, \
       int n, const scalar_t *A, int lda, const value_t *W, int *lwork,      \
diff --git a/aten/src/ATen/native/cudnn/BatchNorm.h b/aten/src/ATen/native/cudnn/BatchNorm.h
index 3da76c0c16e41..7bd0a18b94de2 100644
--- a/aten/src/ATen/native/cudnn/BatchNorm.h
+++ b/aten/src/ATen/native/cudnn/BatchNorm.h
@@ -1,6 +1,6 @@
 namespace at::native {
 
-TORCH_API size_t
+TORCH_CUDA_CU_API size_t
 _get_cudnn_batch_norm_reserve_space_size(const Tensor& input_t, bool training);
 
 } // namespace at::native
diff --git a/aten/src/ATen/native/cudnn/ConvShared.cpp b/aten/src/ATen/native/cudnn/ConvShared.cpp
index 1adf440b26916..ff4a02346810a 100644
--- a/aten/src/ATen/native/cudnn/ConvShared.cpp
+++ b/aten/src/ATen/native/cudnn/ConvShared.cpp
@@ -632,7 +632,8 @@ std::tuple<at::Tensor, at::Tensor> cudnn_convolution_backward(
     }
   }
 
-  return std::tuple<Tensor, Tensor>{grad_input, grad_weight};
+  return std::tuple<Tensor, Tensor>{
+      std::move(grad_input), std::move(grad_weight)};
 }
 
 Tensor cudnn_convolution_transpose_backward_weight(
@@ -702,7 +703,8 @@ std::tuple<at::Tensor, at::Tensor> cudnn_convolution_transpose_backward(
         allow_tf32);
   }
 
-  return std::tuple<Tensor, Tensor>{grad_input, grad_weight};
+  return std::tuple<Tensor, Tensor>{
+      std::move(grad_input), std::move(grad_weight)};
 }
 
 Tensor cudnn_convolution_relu(
diff --git a/aten/src/ATen/native/cudnn/GridSampler.cpp b/aten/src/ATen/native/cudnn/GridSampler.cpp
index 21ac35ff7238e..1b1e3d19699b7 100644
--- a/aten/src/ATen/native/cudnn/GridSampler.cpp
+++ b/aten/src/ATen/native/cudnn/GridSampler.cpp
@@ -183,7 +183,8 @@ std::tuple<Tensor, Tensor> cudnn_grid_sampler_backward(
       &zero,
       grad_grid_t.data_ptr()));
 
-  return std::tuple<Tensor, Tensor>{grad_input_t, grad_grid_t};
+  return std::tuple<Tensor, Tensor>{
+      std::move(grad_input_t), std::move(grad_grid_t)};
 }
 
 } // namespace at::native
diff --git a/aten/src/ATen/native/cudnn/RNN.cpp b/aten/src/ATen/native/cudnn/RNN.cpp
index dfeee0d7d3879..8c346de54396b 100644
--- a/aten/src/ATen/native/cudnn/RNN.cpp
+++ b/aten/src/ATen/native/cudnn/RNN.cpp
@@ -108,6 +108,8 @@ Tensor _cudnn_init_dropout_state(
       false, "_cudnn_init_dropout_state: ATen not compiled with cuDNN support");
 }
 
+void _cudnn_clear_dropout_state() {}
+
 } // namespace native
 } // namespace at
 
@@ -1960,7 +1962,7 @@ std::tuple<Tensor, Tensor, Tensor> _cudnn_rnn_backward_input(
     dx = dx.transpose_(0, 1);
   }
 
-  return std::make_tuple(dx, dhx, dcx);
+  return std::make_tuple(std::move(dx), std::move(dhx), std::move(dcx));
 }
 
 // NB: This MUST BE CALLED AFTER _cudnn_rnn_backward_input.
@@ -2249,7 +2251,7 @@ std::tuple<Tensor, Tensor, Tensor, std::vector<Tensor>> _cudnn_rnn_backward(
         reserve);
   }
   return std::tuple<Tensor, Tensor, Tensor, std::vector<Tensor>>{
-      dx, dhx, dcx, dw};
+      std::move(dx), std::move(dhx), std::move(dcx), std::move(dw)};
 }
 
 // TODO: I am not sure if we actually need the 'dropout' and 'train' parameters
@@ -2376,12 +2378,7 @@ struct DropoutState {
     if (event) {
 #if !defined(USE_ROCM)
       // See Note [DropoutState and CUDA graph capture]
-      cudaStreamCaptureStatus status;
-      AT_CUDA_CHECK(cudaStreamGetCaptureInfo(
-          cuda::getCurrentCUDAStream(), &status, &capture_id_last_lock));
-      if (status == cudaStreamCaptureStatus::cudaStreamCaptureStatusNone) {
-        capture_id_last_lock = 0;
-      }
+      capture_id_last_lock = at::cuda::currentStreamCaptureId().value_or(0);
       if (capture_id_last_lock == capture_id_last_unlock) {
         event->block(cuda::getCurrentCUDAStream());
       }
@@ -2396,12 +2393,7 @@ struct DropoutState {
       event->record();
 #if !defined(USE_ROCM)
       // See Note [DropoutState and CUDA graph capture]
-      cudaStreamCaptureStatus status;
-      AT_CUDA_CHECK(cudaStreamGetCaptureInfo(
-          cuda::getCurrentCUDAStream(), &status, &capture_id_last_unlock));
-      if (status == cudaStreamCaptureStatus::cudaStreamCaptureStatusNone) {
-        capture_id_last_unlock = 0;
-      }
+      capture_id_last_unlock = at::cuda::currentStreamCaptureId().value_or(0);
       TORCH_INTERNAL_ASSERT(capture_id_last_unlock == capture_id_last_lock);
 #endif
     }
@@ -2409,21 +2401,29 @@ struct DropoutState {
   }
 };
 
+// Each state is slightly over 2MB and initialized lazily, so it's fine to
+// cache them. The cache is process-lifetime state; it can be released via
+// clear_dropout_state() (exposed as torch._C._cudnn_clear_dropout_state()).
+std::vector<DropoutState>& dropout_state_cache() {
+  static std::vector<DropoutState> cache{
+      static_cast<size_t>(cuda::getNumGPUs())};
+  return cache;
+}
+
+std::mutex& dropout_state_cache_mutex() {
+  static std::mutex mut;
+  return mut;
+}
+
 DropoutState& get_dropout_state(
     double dropout_p,
     bool train,
     TensorOptions options) {
-  // Each state is slightly over 2MB and initialized lazily, so it's fine to
-  // cache them.
-  static std::vector<DropoutState> dropout_state_cache{
-      static_cast<size_t>(cuda::getNumGPUs())};
-  static std::mutex state_cache_mut;
-
   AT_ASSERT(options.device().is_cuda());
   auto device = options.device().index();
 
-  std::unique_lock<std::mutex> lock{state_cache_mut};
-  auto& state = dropout_state_cache.at(device);
+  std::unique_lock<std::mutex> lock{dropout_state_cache_mutex()};
+  auto& state = dropout_state_cache().at(device);
   if (train && dropout_p > 0) {
     const auto& gen = at::detail::getCUDAHooks().getDefaultGenerator(device);
     auto gen_impl = gen.get<at::CUDAGeneratorImpl>();
@@ -2822,6 +2822,15 @@ TORCH_LIBRARY_IMPL(aten, Meta, m) {
 
 } // namespace
 
+void _cudnn_clear_dropout_state() {
+  std::lock_guard<std::mutex> lock{dropout_state_cache_mutex()};
+  for (auto& state : dropout_state_cache()) {
+    std::lock_guard<std::mutex> state_lock{state.mutex};
+    state.buffer = Tensor();
+    state.event.reset();
+  }
+}
+
 } // namespace at::native
 
 #endif // AT_CUDNN_ENABLED()
diff --git a/aten/src/ATen/native/hip/ck_group_gemm.hip b/aten/src/ATen/native/hip/ck_group_gemm.hip
index 7b6014fbddb4d..aa53441af2156 100644
--- a/aten/src/ATen/native/hip/ck_group_gemm.hip
+++ b/aten/src/ATen/native/hip/ck_group_gemm.hip
@@ -13,6 +13,7 @@
 #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/device/impl/device_grouped_gemm_xdl_splitk_cshuffle.hpp>
 #include <ck/tensor_operation/gpu/element/element_wise_operation.hpp>
 #include <ck/utility/tuple.hpp>
 
@@ -31,7 +32,54 @@ namespace CkTypes {
 }
 
 template <typename ALayout, typename BLayout, typename DataType>
-using GroupedGemmKernel = ck::tensor_operation::device::DeviceGroupedGemmMultipleDSplitKXdlCShuffleTwoStage<
+using GroupedGemmXdlSplit = ck::tensor_operation::device::DeviceGroupedGemmXdlSplitKCShuffle<
+    ALayout,                                 // Layout for A
+    BLayout,                                 // Layout for B
+    ck::Tuple<>,                             // DsLayout
+    ck::tensor_layout::gemm::RowMajor,       // ELayout/CLayout
+    DataType,                                // A type
+    DataType,                                // B type
+    CkTypes::F32,                            // Accumulator type
+    DataType,                                // CShuffle type
+    ck::Tuple<>,                             // Ds type
+    DataType,                                // E type
+    CkTypes::PassThrough,                    // Elementwise functors
+    CkTypes::PassThrough,
+    CkTypes::PassThrough,
+    ck::tensor_operation::device::GemmSpecialization::MNKPadding,
+    1,                                       // NumGemmKPrefetchStage
+    256,                                     // BlockSize
+    256,                                     // MPerBlock
+    128,                                     // NPerBlock
+    32,                                      // KPerBlock
+    8,                                       // AK1
+    8,                                       // BK1
+    32,                                      // MPerXDL
+    32,                                      // NPerXDL
+    4,                                       // MXdlPerWave
+    2,                                       // NXdlPerWave
+    S<1,4,64,1>,                             // ABlockTransferThreadClusterLengths_K0_M_K1
+    S<0,2,1,3>,                              // ABlockTransferThreadClusterArrangeOrder
+    S<0,2,1,3>,                              // ABlockTransferSrcAccessOrder
+    3,                                       // ABlockTransferSrcVectorDim
+    8,                                       // ABlockTransferSrcScalarPerVector
+    8,                                       // ABlockTransferDstScalarPerVector_K1
+    1,                                       // ABlockLdsExtraM
+    S<1,4,64,1>,                             // BBlockTransferThreadClusterLengths_K0_N_K1
+    S<0,2,1,3>,                              // BBlockTransferThreadClusterArrangeOrder
+    S<0,2,1,3>,                              // BBlockTransferSrcAccessOrder
+    3,                                       // BBlockTransferSrcVectorDim
+    8,                                       // BBlockTransferSrcScalarPerVector
+    8,                                       // BBlockTransferDstScalarPerVector_K1
+    1,                                       // BBlockLdsExtraN
+    1,                                       // CShuffleMXdlPerWavePerShuffle
+    1,                                       // CShuffleNXdlPerWavePerShuffle
+    S<1,32,1,8>,                             // CDEBlockTransferClusterLengths_MBlock_MPerBlock_NBlock_NPerBlock
+    4                                        // CDEBlockTransferScalarPerVector_NPerBlock
+>;
+
+template <typename ALayout, typename BLayout, typename DataType>
+using GroupedGemmMultipleDSplit = 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,
@@ -52,7 +100,6 @@ void launch_grouped_bgemm_ck_impl_dispatch(
     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;
@@ -381,27 +428,51 @@ void launch_grouped_bgemm_ck_impl_dispatch(
     // 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 = c10::cuda::CUDACachingAllocator::raw_alloc(arg_buf_size);
-    void* ws_buf = c10::cuda::CUDACachingAllocator::raw_alloc(ws_size);
-
-    gemm_instance.SetDeviceKernelArgs(&argument, gemm_arg_buf);
-    gemm_instance.SetWorkSpacePointer(&argument, ws_buf);
-
-    auto invoker = gemm_instance.MakeInvoker();
-    hipStream_t stream = c10::cuda::getCurrentCUDAStream();
-    invoker.Run(argument, {stream});
-    c10::cuda::CUDACachingAllocator::raw_delete(gemm_arg_buf);
-    c10::cuda::CUDACachingAllocator::raw_delete(ws_buf);
+    bool all_K_equal = true;
+    ck::index_t base_K = gemm_descs.empty() ? 0 : gemm_descs[0].K_;
+    for (const auto& desc : gemm_descs) {
+        if (desc.K_ != base_K) {
+           all_K_equal = false;
+           break;
+       }
+    }
+
+    auto run_gemm = [](auto& gemm_instance, auto& argument) {
+        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 = c10::cuda::CUDACachingAllocator::raw_alloc(arg_buf_size);
+        void* ws_buf = c10::cuda::CUDACachingAllocator::raw_alloc(ws_size);
+
+        gemm_instance.SetDeviceKernelArgs(&argument, gemm_arg_buf);
+        gemm_instance.SetWorkSpacePointer(&argument, ws_buf);
+
+        auto invoker = gemm_instance.MakeInvoker();
+        hipStream_t stream = c10::cuda::getCurrentCUDAStream();
+        invoker.Run(argument, {stream});
+        c10::cuda::CUDACachingAllocator::raw_delete(gemm_arg_buf);
+        c10::cuda::CUDACachingAllocator::raw_delete(ws_buf);
+    };
+
+    if (all_K_equal) {
+        using DeviceOp = GroupedGemmXdlSplit<ALayout, BLayout, DataType>;
+        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{}
+        );
+        run_gemm(gemm_instance, argument);
+    } else {
+        using DeviceOp = GroupedGemmMultipleDSplit<ALayout, BLayout, DataType>;
+        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{}
+        );
+        run_gemm(gemm_instance, argument);
+    }
 }
 
 void group_gemm_ck(
diff --git a/aten/src/ATen/native/layer_norm.cpp b/aten/src/ATen/native/layer_norm.cpp
index dadfe8aef5fd6..3fab14edfb433 100644
--- a/aten/src/ATen/native/layer_norm.cpp
+++ b/aten/src/ATen/native/layer_norm.cpp
@@ -250,6 +250,7 @@ std::tuple<Tensor, Tensor, Tensor> math_native_layer_norm(
     out = out.add(bias);
   }
   std::vector<int64_t> stat_shape;
+  stat_shape.reserve(input.dim());
   for (const auto idx : c10::irange(axis)) {
     stat_shape.push_back(input_shape[idx]);
   }
@@ -268,6 +269,7 @@ std::tuple<Tensor, Tensor> rms_norm_composite(
     std::optional<double> eps) {
 
   std::vector<int64_t> dims_to_reduce;
+  dims_to_reduce.reserve(normalized_shape.size());
   for (const auto i : c10::irange(normalized_shape.size())) {
     dims_to_reduce.push_back(input.dim() - i - 1);
   }
diff --git a/aten/src/ATen/native/miopen/BatchNorm_miopen.cpp b/aten/src/ATen/native/miopen/BatchNorm_miopen.cpp
index 0c122c9e13d4d..1bbd2a99ce70a 100644
--- a/aten/src/ATen/native/miopen/BatchNorm_miopen.cpp
+++ b/aten/src/ATen/native/miopen/BatchNorm_miopen.cpp
@@ -156,7 +156,7 @@ std::tuple<Tensor, Tensor, Tensor> miopen_batch_norm(
   // save_mean and save_var can be undefined
   // If this causes problems, we can initialize them to empty tensors
   // of the correct type
-  return std::tuple<Tensor, Tensor, Tensor>{output_t, save_mean, save_var};
+  return std::tuple<Tensor, Tensor, Tensor>{std::move(output_t), std::move(save_mean), std::move(save_var)};
 }
 
 std::tuple<Tensor, Tensor, Tensor> miopen_batch_norm_backward(
@@ -235,7 +235,7 @@ std::tuple<Tensor, Tensor, Tensor> miopen_batch_norm_backward(
     save_mean->const_data_ptr(),
     save_var->const_data_ptr()));
 
-  return std::tuple<Tensor,Tensor,Tensor>{grad_input_t, grad_weight_t, grad_bias_t};
+  return std::tuple<Tensor,Tensor,Tensor>{std::move(grad_input_t), std::move(grad_weight_t), std::move(grad_bias_t)};
 }
 
 }}  // namespace native
diff --git a/aten/src/ATen/native/miopen/Conv_miopen.cpp b/aten/src/ATen/native/miopen/Conv_miopen.cpp
index 0d61c2700f118..1748ae7bec0ec 100644
--- a/aten/src/ATen/native/miopen/Conv_miopen.cpp
+++ b/aten/src/ATen/native/miopen/Conv_miopen.cpp
@@ -1088,7 +1088,7 @@ std::tuple<at::Tensor, at::Tensor, at::Tensor> miopen_convolution_transpose_back
     grad_bias = miopen_convolution_backward_bias(grad_output);
   }
 
-  return std::tuple<Tensor,Tensor,Tensor>{grad_input, grad_weight, grad_bias};
+  return std::tuple<Tensor,Tensor,Tensor>{std::move(grad_input), std::move(grad_weight), std::move(grad_bias)};
 }
 
 // ---------------------------------------------------------------------
@@ -1572,7 +1572,7 @@ std::tuple<at::Tensor, at::Tensor, at::Tensor> miopen_convolution_backward(
     }
   }
 
-  return std::tuple<Tensor, Tensor, Tensor>{grad_input, grad_weight, grad_bias};
+  return std::tuple<Tensor, Tensor, Tensor>{std::move(grad_input), std::move(grad_weight), std::move(grad_bias)};
 }
 
 Tensor miopen_convolution_transpose_forward(
@@ -1761,7 +1761,7 @@ std::tuple<at::Tensor, at::Tensor, at::Tensor> miopen_depthwise_convolution_back
     grad_bias = miopen_convolution_backward_bias(grad_output);
   }
 
-  return std::tuple<Tensor,Tensor,Tensor>{grad_input, grad_weight, grad_bias};
+  return std::tuple<Tensor,Tensor,Tensor>{std::move(grad_input), std::move(grad_weight), std::move(grad_bias)};
 }
 
 // ---------------------------------------------------------------------
diff --git a/aten/src/ATen/native/miopen/LossCTC_miopen.cpp b/aten/src/ATen/native/miopen/LossCTC_miopen.cpp
index 21797e7537d59..6200f7ede7df5 100644
--- a/aten/src/ATen/native/miopen/LossCTC_miopen.cpp
+++ b/aten/src/ATen/native/miopen/LossCTC_miopen.cpp
@@ -12,6 +12,7 @@
 #include <ATen/ops/empty_like.h>
 #include <ATen/ops/miopen_ctc_loss.h>
 #include <ATen/ops/miopen_ctc_loss_native.h>
+#include <ATen/ops/_use_miopen_ctc_loss_native.h>
 #endif
 
 // TODO: Remove the condition on AT_ROCM_ENABLED entirely,
@@ -206,35 +207,16 @@ std::tuple<Tensor, Tensor> miopen_ctc_loss(
   Tensor costs = at::empty({batch_size}, log_probs->options());
   Tensor grad = at::empty_like(log_probs_t, LEGACY_CONTIGUOUS_MEMORY_FORMAT);
 
-  // MIOpen requires labels and lengths on GPU
-  Tensor labels_gpu = targets_t.to(Device(at::kCUDA), at::kInt);
-  Tensor label_lengths_gpu = at::empty(
-      {static_cast<int64_t>(target_lengths.size())},
-      at::TensorOptions().dtype(at::kInt).device(at::kCUDA));
-  Tensor input_lengths_gpu = at::empty(
-      {static_cast<int64_t>(input_lengths.size())},
-      at::TensorOptions().dtype(at::kInt).device(at::kCUDA));
-
-  C10_CUDA_CHECK(hipMemcpy(
-      label_lengths_gpu.data_ptr<int>(),
-      target_lengths.data(),
-      target_lengths.size() * sizeof(int),
-      hipMemcpyHostToDevice));
-  C10_CUDA_CHECK(hipMemcpy(
-      input_lengths_gpu.data_ptr<int>(),
-      input_lengths.data(),
-      input_lengths.size() * sizeof(int),
-      hipMemcpyHostToDevice));
-
+  // MIOpen reads labels/lengths on the host.
   size_t workspace_size;
   (void)deterministic; // MIOpen only supports deterministic algorithm
   MIOPEN_CHECK(miopenGetCTCLossWorkspaceSize(
       handle,
       probs_desc,
       grads_desc,
-      labels_gpu.data_ptr<int>(),
-      label_lengths_gpu.data_ptr<int>(),
-      input_lengths_gpu.data_ptr<int>(),
+      targets_t.data_ptr<int>(),
+      target_lengths.data(),
+      input_lengths.data(),
       MIOPEN_CTC_LOSS_ALGO_DETERMINISTIC,
       ctc_desc,
       &workspace_size));
@@ -245,9 +227,9 @@ std::tuple<Tensor, Tensor> miopen_ctc_loss(
       handle,
       probs_desc,
       log_probs_t.data_ptr(),
-      labels_gpu.data_ptr<int>(),
-      label_lengths_gpu.data_ptr<int>(),
-      input_lengths_gpu.data_ptr<int>(),
+      targets_t.data_ptr<int>(),
+      target_lengths.data(),
+      input_lengths.data(),
       costs.data_ptr(),
       grads_desc,
       grad.data_ptr(),
diff --git a/aten/src/ATen/native/miopen/RNN_miopen.cpp b/aten/src/ATen/native/miopen/RNN_miopen.cpp
index 2a9568884d959..b301d23258c21 100644
--- a/aten/src/ATen/native/miopen/RNN_miopen.cpp
+++ b/aten/src/ATen/native/miopen/RNN_miopen.cpp
@@ -316,6 +316,12 @@ thread_local DropoutDescriptor RNNDescriptors::dropout_desc {};
 // Each state is 0.75 MB so there is no problem in caching all of them for each thread
 thread_local std::unique_ptr<DropoutState> RNNDescriptors::dropout_states { nullptr };
 
+// Releases the calling thread's cached MIOpen RNN dropout state buffer. This is
+// thread-lifetime state that empty_cache() cannot reclaim on its own.
+void _miopen_clear_dropout_state() {
+  RNNDescriptors::dropout_states.reset();
+}
+
 Tensor permute_wei_for_miopen(Tensor wei, int64_t mode)
 {
     if (mode < 2)
@@ -743,7 +749,7 @@ std::tuple<Tensor, Tensor, Tensor, Tensor> miopen_rnn_backward_input(
         dx = dx.transpose_(0, 1);
     }
 
-    return std::make_tuple(dx, dhx, dcx, workspace);
+    return std::make_tuple(std::move(dx), std::move(dhx), std::move(dcx), std::move(workspace));
 }
 
 std::vector<Tensor> miopen_rnn_backward_weight(
@@ -866,7 +872,7 @@ std::tuple<Tensor, Tensor, Tensor, std::vector<Tensor>> miopen_rnn_backward(
             }
         }
     }
-    return std::tuple<Tensor, Tensor, Tensor, std::vector<Tensor>>{dx, dhx, dcx, dw};
+    return std::tuple<Tensor, Tensor, Tensor, std::vector<Tensor>>{std::move(dx), std::move(dhx), std::move(dcx), std::move(dw)};
 }
 
 namespace {
diff --git a/aten/src/ATen/native/mkl/SparseCsrLinearAlgebra.cpp b/aten/src/ATen/native/mkl/SparseCsrLinearAlgebra.cpp
index cd87850fe9eb9..6b34c0c7d8b4a 100644
--- a/aten/src/ATen/native/mkl/SparseCsrLinearAlgebra.cpp
+++ b/aten/src/ATen/native/mkl/SparseCsrLinearAlgebra.cpp
@@ -23,7 +23,6 @@ Tensor& _sparse_mm_mkl_(
 #else
   TORCH_CHECK(false, "sparse_mm_mkl: ATen not compiled with MKL support");
 #endif
-  return self; // for stopping compiler warnings.
 }
 } // namespace native
 
diff --git a/aten/src/ATen/native/mkl/SpectralOps.cpp b/aten/src/ATen/native/mkl/SpectralOps.cpp
index da4fc3fb6d079..21720d9e2698d 100644
--- a/aten/src/ATen/native/mkl/SpectralOps.cpp
+++ b/aten/src/ATen/native/mkl/SpectralOps.cpp
@@ -160,12 +160,7 @@ static void _fft_fill_with_conjugate_symmetry_cpu_(
 }
 
 // Register this one implementation for all cpu types instead of compiling multiple times
-REGISTER_ARCH_DISPATCH(fft_fill_with_conjugate_symmetry_stub, DEFAULT, &_fft_fill_with_conjugate_symmetry_cpu_)
-REGISTER_AVX2_DISPATCH(fft_fill_with_conjugate_symmetry_stub, &_fft_fill_with_conjugate_symmetry_cpu_)
-REGISTER_AVX512_DISPATCH(fft_fill_with_conjugate_symmetry_stub, &_fft_fill_with_conjugate_symmetry_cpu_)
-REGISTER_ZVECTOR_DISPATCH(fft_fill_with_conjugate_symmetry_stub, &_fft_fill_with_conjugate_symmetry_cpu_)
-REGISTER_VSX_DISPATCH(fft_fill_with_conjugate_symmetry_stub, &_fft_fill_with_conjugate_symmetry_cpu_)
-REGISTER_SVE256_DISPATCH(fft_fill_with_conjugate_symmetry_stub, &_fft_fill_with_conjugate_symmetry_cpu_)
+REGISTER_ALL_CPU_DISPATCH(fft_fill_with_conjugate_symmetry_stub, &_fft_fill_with_conjugate_symmetry_cpu_)
 
 // _out variants can be shared between PocketFFT and MKL
 Tensor& _fft_r2c_mkl_out(const Tensor& self, IntArrayRef dim, int64_t normalization,
diff --git a/aten/src/ATen/native/mkldnn/Conv.cpp b/aten/src/ATen/native/mkldnn/Conv.cpp
index 64943484dceb0..d01b93a4a634e 100644
--- a/aten/src/ATen/native/mkldnn/Conv.cpp
+++ b/aten/src/ATen/native/mkldnn/Conv.cpp
@@ -87,6 +87,10 @@ static void check_shape_forward(const Tensor& input,
 
   std::vector<int64_t> input_shape;
   std::vector<int64_t> kernel_shape;
+  if (k > 2) {
+    input_shape.reserve(k - 2);
+    kernel_shape.reserve(k - 2);
+  }
   bool kernel_size_correct = true;
 
   for (const auto i : c10::irange(2, k)) {
@@ -442,7 +446,9 @@ Tensor mkldnn_convolution_pointwise_binary(
     auto weight =
         weight_t.is_mkldnn() ? weight_t : weight_t.contiguous(memory_format);
     auto other = other_t.contiguous(memory_format);
-    auto output = at::empty(output_sizes, input_t.options()).contiguous(memory_format);
+    auto output = at::empty(
+        output_sizes,
+        input_t.options().memory_format(memory_format));
     const ideep::tensor x = itensor_from_tensor(input);
     const ideep::tensor w = itensor_from_tensor(weight);
     const ideep::tensor z = itensor_from_tensor(other);
diff --git a/aten/src/ATen/native/mkldnn/Linear.cpp b/aten/src/ATen/native/mkldnn/Linear.cpp
index 719865063e20d..d01b2cdd356ad 100644
--- a/aten/src/ATen/native/mkldnn/Linear.cpp
+++ b/aten/src/ATen/native/mkldnn/Linear.cpp
@@ -203,7 +203,7 @@ std::tuple<Tensor, Tensor, Tensor> mkldnn_linear_backward(
   if (output_mask[1] || output_mask[2]) {
     std::tie(grad_weight, grad_bias) = at::mkldnn_linear_backward_weights(grad_output, input, weight, output_mask[2]);
   }
-  return std::tuple<Tensor, Tensor, Tensor>{grad_input, grad_weight, grad_bias};
+  return std::tuple<Tensor, Tensor, Tensor>{std::move(grad_input), std::move(grad_weight), std::move(grad_bias)};
 }
 
 Tensor mkldnn_linear_pointwise(
diff --git a/aten/src/ATen/native/mkldnn/Matmul.cpp b/aten/src/ATen/native/mkldnn/Matmul.cpp
index 92eb5a9c83870..e1cc99f75ab26 100644
--- a/aten/src/ATen/native/mkldnn/Matmul.cpp
+++ b/aten/src/ATen/native/mkldnn/Matmul.cpp
@@ -519,9 +519,15 @@ static void _mkldnn_matmul_i8i8i32_with_primitive(
     const Tensor &mat2,
     const Tensor &result) {
   // Create ideep tensors for oneDNN computation
+  ideep::tensor::data_type src_dtype;
+  if (mat1.scalar_type() == at::kByte) {
+    src_dtype = ideep::tensor::data_type::u8;
+  } else {
+    src_dtype = ideep::tensor::data_type::s8;
+  }
   auto src = ideep::tensor(
       {mat1.sizes().vec(),
-       ideep::tensor::data_type::s8,
+       src_dtype,
        mat1.strides().vec()},
       mat1.data_ptr());
   auto wei = ideep::tensor(
@@ -557,6 +563,7 @@ static void _mkldnn_matmul_i8i8i32_with_primitive(
   primitive.execute(ideep::stream::default_stream(), args);
 }
 
+template <typename a_type>
 static void _mkldnn_gemm_i8i8i32_with_blas(
   const Tensor& self,
   const Tensor& mat2,
@@ -576,42 +583,42 @@ static void _mkldnn_gemm_i8i8i32_with_blas(
     const float alpha = 1;
     const float beta = 0;
 
-    int8_t ao = 0;
+    a_type ao = 0;
     int8_t bo = 0;
     int32_t co = 0;
 
-    dnnl::gemm_s8s8s32(
-        transa,
-        transb,
-        offsetc,
-        m,
-        n,
-        k,
-        alpha,
-        static_cast<int8_t*>(self.data_ptr()),
-        lda,
-        ao,
-        static_cast<int8_t*>(mat2.data_ptr()),
-        ldb,
-        bo,
-        beta,
-        static_cast<int32_t*>(result.data_ptr()),
-        ldc,
-        &co);
+    #define CALL_DNNL_GEMM(gemm_func, a_ptr_type, a_data_ptr) \
+      gemm_func(                                              \
+        transa, transb, offsetc, m, n, k,                     \
+        alpha,                                                \
+        static_cast<a_ptr_type>(a_data_ptr), lda, ao,         \
+        static_cast<int8_t*>(mat2.data_ptr()), ldb, bo,       \
+        beta,                                                 \
+        static_cast<int32_t*>(result.data_ptr()), ldc, &co)
+
+    if (self.scalar_type() == at::kByte) { //uint8
+      CALL_DNNL_GEMM(dnnl::gemm_u8s8s32, uint8_t*, self.data_ptr());
+    } else { //int8
+      CALL_DNNL_GEMM(dnnl::gemm_s8s8s32, int8_t*, self.data_ptr());
+    }
   }
 
 void mkldnn_matmul_i8i8i32(
     const Tensor &mat1,
     const Tensor &mat2,
     const Tensor &result) {
-  // x:s8 * w:s8 -> y:s32
+  // x:u8 or s8 * w:s8 -> y:s32
   // both inputs should be 2d
   // In most cases, using DNNL blas API is faster but it requires a/b contiguous along one dimentsion
   bool a_is_contigous = (mat1.stride(0) == 1 || mat1.stride(1) == 1);
   bool b_is_contigous = (mat2.stride(0) == 1 || mat2.stride(1) == 1);
 
   if (a_is_contigous && b_is_contigous) {
-    _mkldnn_gemm_i8i8i32_with_blas(mat1, mat2, result);
+    if (mat1.scalar_type() == at::kByte) {
+      _mkldnn_gemm_i8i8i32_with_blas<uint8_t>(mat1, mat2, result);
+    } else {
+      _mkldnn_gemm_i8i8i32_with_blas<int8_t>(mat1, mat2, result);
+    }
   } else {
     _mkldnn_matmul_i8i8i32_with_primitive(mat1, mat2, result);
   }
diff --git a/aten/src/ATen/native/mkldnn/Matmul.h b/aten/src/ATen/native/mkldnn/Matmul.h
index 80247497d58f0..49709b6acb7a5 100644
--- a/aten/src/ATen/native/mkldnn/Matmul.h
+++ b/aten/src/ATen/native/mkldnn/Matmul.h
@@ -81,7 +81,7 @@ bool use_mkldnn_matmul(
     const Tensor& mat2,
     const Tensor& result);
 
-// x:s8 * w:s8 -> y:s32
+// x:u8 or s8 * w:s8 -> y:s32
 TORCH_API void mkldnn_matmul_i8i8i32(
     const Tensor &mat1,
     const Tensor &mat2,
diff --git a/aten/src/ATen/native/mkldnn/RNN.cpp b/aten/src/ATen/native/mkldnn/RNN.cpp
index c93198aec33a2..63e1f5a9e1242 100644
--- a/aten/src/ATen/native/mkldnn/RNN.cpp
+++ b/aten/src/ATen/native/mkldnn/RNN.cpp
@@ -432,7 +432,7 @@ std::tuple<Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor> mkldnn_rnn_la
       forward_hint.workspace_desc(), workspace.template data_ptr<uint8_t>());
   ideep::lstm_backward::compute(forward_hint, x, hx, cx, w1, w2, b, y, hy, cy, diff_y, diff_hy, diff_cy, mkldnn_workspace, diff_x, diff_hx, diff_cx, diff_w1, diff_w2, diff_b, reverse);
   auto diff_b2_ = at::clone(diff_b_);
-  return std::make_tuple(diff_x_, diff_w1_, diff_w2_, diff_b_, diff_b2_, diff_hx_, diff_cx_);
+  return std::make_tuple(std::move(diff_x_), std::move(diff_w1_), std::move(diff_w2_), std::move(diff_b_), std::move(diff_b2_), std::move(diff_hx_), std::move(diff_cx_));
 }
 
 // MKLDNN RNN integration notes:
@@ -530,7 +530,7 @@ static std::tuple<Tensor, Tensor, Tensor> mkldnn_rnn(
   if (batch_first) {
     output = output.transpose(0, 1);
   }
-  return std::make_tuple(output, hy, cy);
+  return std::make_tuple(std::move(output), std::move(hy), std::move(cy));
 }
 
 ////////////////////////////////////////////////////////////////////////////////
diff --git a/aten/src/ATen/native/mkldnn/xpu/Blas.cpp b/aten/src/ATen/native/mkldnn/xpu/Blas.cpp
index 7272fab02261a..05ab70cdaa2c3 100644
--- a/aten/src/ATen/native/mkldnn/xpu/Blas.cpp
+++ b/aten/src/ATen/native/mkldnn/xpu/Blas.cpp
@@ -85,6 +85,30 @@ Tensor& addmm_out(
       " but got:",
       self.sizes());
 
+  // Bypass OneDNN optimization path for float64 due to lack of full double
+  // precision support.
+  if (mat1.scalar_type() == at::kDouble) {
+    bool is_inplace = self.is_same(result);
+    bool is_beta_ne_zero = beta.to<double>() != 0.0;
+
+    Tensor self_copy;
+    if (is_inplace && is_beta_ne_zero) {
+      self_copy = self.clone();
+    }
+
+    onednn::matmul(result, mat1, mat2, Tensor(), true, onednn::Attr());
+
+    if (alpha.to<double>() != 1.0) {
+      result.mul_(alpha);
+    }
+
+    if (is_beta_ne_zero) {
+      result.add_(is_inplace ? self_copy : self, beta);
+    }
+
+    return result;
+  }
+
   // general case
   Tensor bias = Tensor();
   onednn::Attr attr;
@@ -223,6 +247,31 @@ Tensor& baddbmm_out(
     return result;
   }
 
+  // Bypass OneDNN optimization path for float64 due to lack of full double
+  // precision support.
+  if (batch1.scalar_type() == at::kDouble ||
+      batch2.scalar_type() == at::kDouble) {
+    bool is_inplace = input.is_same(result);
+    bool is_beta_ne_zero = beta.to<double>() != 0.0;
+
+    Tensor input_copy;
+    if (is_inplace && is_beta_ne_zero) {
+      input_copy = input.clone();
+    }
+
+    onednn::matmul(result, batch1, batch2, Tensor(), true, onednn::Attr());
+
+    if (alpha.to<double>() != 1.0) {
+      result.mul_(alpha);
+    }
+
+    if (is_beta_ne_zero) {
+      result.add_(is_inplace ? input_copy : input, beta);
+    }
+
+    return result;
+  }
+
   // general case
   onednn::Attr attr;
   float beta_ = beta.to<float>();
@@ -326,6 +375,17 @@ Tensor& addmv_out(
   }
 
   Tensor vec_v = vec.view({vec.size(0), 1});
+
+  bool is_float64 =
+      mat.scalar_type() == at::kDouble || vec.scalar_type() == at::kDouble;
+  bool is_inplace = self.is_same(out);
+  if (is_float64 && is_inplace) {
+    Tensor self_v_copy = self_v.clone();
+    at::native::xpu::addmm_out(self_v_copy, mat, vec_v, beta, alpha, out);
+    out.resize_({mat.size(0)});
+    return out;
+  }
+
   at::native::xpu::addmm_out(self_v, mat, vec_v, beta, alpha, out);
   out.resize_({mat.size(0)});
   return out;
diff --git a/aten/src/ATen/native/mkldnn/xpu/Conv.cpp b/aten/src/ATen/native/mkldnn/xpu/Conv.cpp
index eb8f154dedaf0..d6736750fb000 100644
--- a/aten/src/ATen/native/mkldnn/xpu/Conv.cpp
+++ b/aten/src/ATen/native/mkldnn/xpu/Conv.cpp
@@ -384,9 +384,10 @@ Tensor _convolution_out(
   at::MemoryFormat mfmt = is_channels_last_suggested
       ? get_cl_tag_by_ndim(input.ndimension())
       : at::MemoryFormat::Contiguous;
-  auto bias = bias_r.defined() ? bias_r.contiguous() : bias_r;
-  input = input.contiguous(mfmt);
-  weight = weight.contiguous(mfmt);
+
+  auto bias = bias_r.defined() ? make_contiguous_and_aligned(bias_r) : bias_r;
+  input = make_contiguous_and_aligned(input, mfmt);
+  weight = make_contiguous_and_aligned(weight, mfmt);
   check_shape_forward(input, weight, bias, params, true);
 
   Tensor output;
@@ -591,9 +592,9 @@ std::tuple<Tensor, Tensor, Tensor> convolution_backward_overrideable(
   auto mfmt = is_channels_last_suggested
       ? get_cl_tag_by_ndim(input_.ndimension())
       : at::MemoryFormat::Contiguous;
-  grad_output_ = grad_output_.contiguous(mfmt);
-  weight_ = weight_.contiguous(mfmt);
-  input_ = input_.contiguous(mfmt);
+  grad_output_ = make_contiguous_and_aligned(grad_output_, mfmt);
+  weight_ = make_contiguous_and_aligned(weight_, mfmt);
+  input_ = make_contiguous_and_aligned(input_, mfmt);
 
   auto opt = grad_output_.options();
   Tensor grad_input;
diff --git a/aten/src/ATen/native/mkldnn/xpu/Linear.cpp b/aten/src/ATen/native/mkldnn/xpu/Linear.cpp
index 472c726f9420e..f6bd18112a0c0 100644
--- a/aten/src/ATen/native/mkldnn/xpu/Linear.cpp
+++ b/aten/src/ATen/native/mkldnn/xpu/Linear.cpp
@@ -22,7 +22,10 @@ collapse_in_out_dim(at::Tensor input, int64_t dim, at::Tensor weight) {
   // [BM, N]
   std::vector<int64_t> output_reshaped_size{
       input_reshaped_size[0], weight.size(0)};
-  return {input_reshaped_size, output_size, output_reshaped_size};
+  return {
+      std::move(input_reshaped_size),
+      std::move(output_size),
+      std::move(output_reshaped_size)};
 }
 
 Tensor linear_pointwise(
diff --git a/aten/src/ATen/native/mkldnn/xpu/ScaledBlas.cpp b/aten/src/ATen/native/mkldnn/xpu/ScaledBlas.cpp
index 654cb296cda57..0854fe5940079 100644
--- a/aten/src/ATen/native/mkldnn/xpu/ScaledBlas.cpp
+++ b/aten/src/ATen/native/mkldnn/xpu/ScaledBlas.cpp
@@ -3,6 +3,7 @@
 #include <ATen/WrapDimUtilsMulti.h>
 #include <ATen/ceil_div.h>
 #include <ATen/native/Resize.h>
+#include <ATen/native/ScaledBlasUtils.h>
 #include <ATen/native/mkldnn/xpu/detail/oneDNN.h>
 #include <ATen/native/xpu/Blas.h>
 #include <ATen/xpu/XPUScaledBlas.h>
@@ -359,7 +360,7 @@ using acceptance_fn = std::function<bool(
     ArrayRef<Tensor>&)>;
 using namespace std::placeholders;
 
-namespace scaled_blas = at::native::onednn::scaled;
+namespace scaled_blas = at::native::scaled;
 using scaled_blas::convert_int_to_enum;
 using scaled_blas::ScaledGemmImplementation;
 
diff --git a/aten/src/ATen/native/mkldnn/xpu/detail/QMatmul.cpp b/aten/src/ATen/native/mkldnn/xpu/detail/QMatmul.cpp
index f79dfadd65454..953b3b5a1b2e4 100644
--- a/aten/src/ATen/native/mkldnn/xpu/detail/QMatmul.cpp
+++ b/aten/src/ATen/native/mkldnn/xpu/detail/QMatmul.cpp
@@ -345,8 +345,14 @@ struct ScaleSpec {
       int64_t outer_dim,
       int64_t inner_dim,
       const std::string& arg_type) const {
-    if (groups == dnnl::memory::dims{1, 1})
+    // TensorWise: mask=0, groups={} -> single scale
+    if (groups.empty()) {
+      TORCH_INTERNAL_ASSERT(
+          mask == 0,
+          "Empty groups only valid for TensorWise (mask=0), got mask=",
+          mask);
       return 1; // tensorwise scaling
+    }
 
     TORCH_CHECK(
         arg_type == "src" || arg_type == "wei",
@@ -399,8 +405,8 @@ inline ScaleSpec make_scale_spec(
   if (scaling_type == at::blas::ScalingType::TensorWise) {
     // Scale tensorwise. The same as `--attr-scales=common`.
     // mask=0 : scale whole tensor
-    // groups={1, 1}: indicates that there is only one group for scaling
-    return {0, {1, 1}, dnnl::memory::data_type::f32};
+    // groups={}: indicates that there is only one group for scaling
+    return {0, {}, dnnl::memory::data_type::f32};
   } else {
     // (scaling_type == at::blas::ScalingType::RowWise)
     // Scale RowWise. The same as `--attr-scales=per_dim_01`.
@@ -479,7 +485,7 @@ sycl::event scaled_matmul(
   // scale_result tensor currently only supports scalar(TensorWise Scaling).
   bool with_dst_scale = scale_result && scale_result->defined();
   if (with_dst_scale) {
-    op_attr.set_scales(DNNL_ARG_DST, 0, {1}, dnnl::memory::data_type::f32);
+    op_attr.set_scales(DNNL_ARG_DST, 0, {}, dnnl::memory::data_type::f32);
   }
 
   op_attr.set_scratchpad_mode(dnnl::scratchpad_mode::user);
diff --git a/aten/src/ATen/native/mkldnn/xpu/detail/Utils.cpp b/aten/src/ATen/native/mkldnn/xpu/detail/Utils.cpp
index a8a6b870ff6b6..ae26798d11640 100644
--- a/aten/src/ATen/native/mkldnn/xpu/detail/Utils.cpp
+++ b/aten/src/ATen/native/mkldnn/xpu/detail/Utils.cpp
@@ -287,6 +287,27 @@ void undo_broadcast(at::Tensor& tensor) {
   return;
 }
 
+bool is_64_bytes_aligned(const at::Tensor& tensor) {
+  constexpr uintptr_t alignment_byte = 64;
+  auto data_ptr = reinterpret_cast<uintptr_t>(tensor.const_data_ptr());
+  return (data_ptr % alignment_byte) == 0;
+}
+
+at::Tensor make_contiguous_and_aligned(
+    const at::Tensor& tensor,
+    std::optional<at::MemoryFormat> memory_format) {
+  at::Tensor out = memory_format.has_value() ? tensor.contiguous(*memory_format)
+                                             : tensor.contiguous();
+  if (!is_64_bytes_aligned(out)) {
+    TORCH_WARN(
+        "Tensor is not 64-byte aligned. Cloning to ensure alignment for oneDNN "
+        "operations, which incurs a device-to-device copy.");
+    out = out.clone();
+  }
+
+  return out;
+}
+
 bool is_onednn_matmul_strides(const at::Tensor& tensor) {
   // https://oneapi-src.github.io/oneDNN/dev_guide_matmul.html
   // oneDNN matmul only support 2-dim and 3-dim
@@ -300,11 +321,8 @@ bool is_onednn_matmul_strides(const at::Tensor& tensor) {
   if (tensor.is_contiguous())
     return true;
 
-  if (tensor.storage_offset() > 0) {
-    // currently onednn asks 64 byte alignment
-    constexpr int alignment_byte = 64;
-    if (reinterpret_cast<uintptr_t>(tensor.data_ptr()) % alignment_byte > 0)
-      return false;
+  if (tensor.storage_offset() > 0 && !is_64_bytes_aligned(tensor)) {
+    return false;
   }
 
   // the overlapped cases are not supported
diff --git a/aten/src/ATen/native/mkldnn/xpu/detail/Utils.h b/aten/src/ATen/native/mkldnn/xpu/detail/Utils.h
index 0055fc2f296ad..e841ab8de6583 100644
--- a/aten/src/ATen/native/mkldnn/xpu/detail/Utils.h
+++ b/aten/src/ATen/native/mkldnn/xpu/detail/Utils.h
@@ -3,6 +3,7 @@
 #include <ATen/Tensor.h>
 #include <ATen/core/Tensor.h>
 #include <iostream>
+#include <optional>
 
 #include <ATen/core/grad_mode.h>
 #include <c10/core/MemoryFormat.h>
@@ -46,6 +47,12 @@ void undo_broadcast(at::Tensor& tensor);
 
 bool is_onednn_matmul_strides(const at::Tensor& tensor);
 
+bool is_64_bytes_aligned(const at::Tensor& tensor);
+
+at::Tensor make_contiguous_and_aligned(
+    const at::Tensor& tensor,
+    std::optional<at::MemoryFormat> memory_format = std::nullopt);
+
 bool is_broadcast_from_other_to_self(
     const at::Tensor& self,
     const at::Tensor& other);
diff --git a/aten/src/ATen/native/mkldnn/xpu/detail/oneDNNContext.h b/aten/src/ATen/native/mkldnn/xpu/detail/oneDNNContext.h
index d919e87bfa586..f28aaa131e51e 100644
--- a/aten/src/ATen/native/mkldnn/xpu/detail/oneDNNContext.h
+++ b/aten/src/ATen/native/mkldnn/xpu/detail/oneDNNContext.h
@@ -13,7 +13,7 @@
 
 namespace at::native::onednn {
 
-TORCH_XPU_API dnnl::memory make_onednn_memory(
+dnnl::memory make_onednn_memory(
     dnnl::memory::desc md,
     dnnl::engine& engine,
     void* ptr);
@@ -22,7 +22,7 @@ TORCH_XPU_API dnnl::memory make_onednn_memory(
 bool set_onednn_verbose(int level);
 
 // GpuEngineManager singleton
-struct TORCH_XPU_API GpuEngineManager {
+struct GpuEngineManager {
   static GpuEngineManager& Instance(); // Singleton
 
   dnnl::engine& get_engine(
@@ -50,7 +50,7 @@ struct TORCH_XPU_API GpuEngineManager {
 };
 
 // GpuStreamManager singleton
-struct TORCH_XPU_API GpuStreamManager {
+struct GpuStreamManager {
   static GpuStreamManager& Instance(); // Singleton
 
   dnnl::stream& get_stream(
diff --git a/aten/src/ATen/native/mps/MetalShaderLibrary.h b/aten/src/ATen/native/mps/MetalShaderLibrary.h
index c39e1e5928df0..defaadf5b9e58 100644
--- a/aten/src/ATen/native/mps/MetalShaderLibrary.h
+++ b/aten/src/ATen/native/mps/MetalShaderLibrary.h
@@ -141,7 +141,8 @@ class MetalShaderLibrary {
       TensorIteratorBase& iter,
       const std::string& name,
       const std::optional<c10::Scalar> alpha = std::nullopt,
-      const std::optional<c10::ScalarType> scalar_arg_type = std::nullopt);
+      const std::optional<c10::ScalarType> scalar_arg_type = std::nullopt,
+      bool supports_vec4 = false);
   void exec_binary_kernel(
       TensorIteratorBase& iter,
       const std::string& name,
diff --git a/aten/src/ATen/native/mps/OperationUtils.h b/aten/src/ATen/native/mps/OperationUtils.h
index 7b76da64200c9..112f22742ac13 100644
--- a/aten/src/ATen/native/mps/OperationUtils.h
+++ b/aten/src/ATen/native/mps/OperationUtils.h
@@ -28,16 +28,6 @@
 
 #include <MetalPerformanceShaders/MetalPerformanceShaders.h>
 
-@interface MPSGraph (PyTorchFixups)
-- (MPSGraphTensor*)minimumWithNaNPropagationAndIntFallbackWithPrimaryTensor:(MPSGraphTensor*)primaryTensor
-                                                            secondaryTensor:(MPSGraphTensor*)secondaryTensor
-                                                                       name:(NSString*)name;
-
-- (MPSGraphTensor*)maximumWithNaNPropagationAndIntFallbackWithPrimaryTensor:(MPSGraphTensor*)primaryTensor
-                                                            secondaryTensor:(MPSGraphTensor*)secondaryTensor
-                                                                       name:(NSString*)name;
-@end
-
 using namespace at::mps;
 
 namespace at::native::mps {
@@ -54,6 +44,7 @@ struct MPSScalar {
     float f; // MPS doesn't support 'double'
     at::Half h;
     int64_t i;
+    uint64_t u;
     bool b;
     c10::complex<float> cf;
     c10::complex<at::Half> ch;
@@ -260,7 +251,7 @@ struct MPSKernelCache {
     __block MPSCachedKernel* cachedKernel = nil;
     MPSCacheKey hash = std::hash<std::string>{}(key);
     dispatch_sync_with_rethrow(serialQueue_, ^() {
-      if (cache_.count(hash) != 0) {
+      if (cache_.contains(hash)) {
         auto& entry = cache_.at(hash);
         TORCH_INTERNAL_ASSERT_DEBUG_ONLY(key == entry.key_, "Key collision in the MPS cached kernel!\n");
         cachedKernel = entry.cachedKernel_;
diff --git a/aten/src/ATen/native/mps/OperationUtils.mm b/aten/src/ATen/native/mps/OperationUtils.mm
index a826c0caf95cd..65067ab796cfe 100644
--- a/aten/src/ATen/native/mps/OperationUtils.mm
+++ b/aten/src/ATen/native/mps/OperationUtils.mm
@@ -25,34 +25,6 @@
 #include <mach-o/dyld.h>
 #include <mach-o/getsect.h>
 
-@implementation MPSGraph (PyTorchFixups)
-- (MPSGraphTensor*)minimumWithNaNPropagationAndIntFallbackWithPrimaryTensor:(MPSGraphTensor*)primaryTensor
-                                                            secondaryTensor:(MPSGraphTensor*)secondaryTensor
-                                                                       name:(NSString*)name {
-  // As of MacOS-15.1 m..imumWithNanPropagation is only defined for floating types and calling it with integral
-  // arguments results in
-  //  /AppleInternal/Library/BuildRoots/c7c74b64-74b4-11ef-aeda-9635a580fe0d/Library/Caches/com.apple.xbs/Sources/MetalPerformanceShaders/MPSCore/Utility/MPSKernelDAG.mm:805:
-  //  failed assertion `Error getting visible function: (null) Function isNaN_u8_i8 was not found in the library'
-  if (([primaryTensor dataType] & MPSDataTypeFloatBit) == 0) {
-    return [self minimumWithPrimaryTensor:primaryTensor secondaryTensor:secondaryTensor name:name];
-  }
-  return [self minimumWithNaNPropagationWithPrimaryTensor:primaryTensor secondaryTensor:secondaryTensor name:name];
-}
-
-- (MPSGraphTensor*)maximumWithNaNPropagationAndIntFallbackWithPrimaryTensor:(MPSGraphTensor*)primaryTensor
-                                                            secondaryTensor:(MPSGraphTensor*)secondaryTensor
-                                                                       name:(NSString*)name {
-  // As of MacOS-15.1 m..imumWithNanPropagation is only defined for floating types and calling it with integral
-  // arguments results in
-  //  /AppleInternal/Library/BuildRoots/c7c74b64-74b4-11ef-aeda-9635a580fe0d/Library/Caches/com.apple.xbs/Sources/MetalPerformanceShaders/MPSCore/Utility/MPSKernelDAG.mm:805:
-  //  failed assertion `Error getting visible function: (null) Function isNaN_u8_i8 was not found in the library'
-  if (([primaryTensor dataType] & MPSDataTypeFloatBit) == 0) {
-    return [self maximumWithPrimaryTensor:primaryTensor secondaryTensor:secondaryTensor name:name];
-  }
-  return [self maximumWithNaNPropagationWithPrimaryTensor:primaryTensor secondaryTensor:secondaryTensor name:name];
-}
-@end
-
 namespace at::native::mps {
 /**
  * Computes distance from lowest to highest element offset in given tensor.
@@ -470,8 +442,6 @@ bool isTooLargeForMPSGraph(const Tensor& tensor, bool useMPSStridedAPI) {
   auto sizes = src.sizes();
   auto nStrides = strides.size();
   auto nonZeroStrides = src.strides();
-  int64_t crtNonZeroStride = 1;
-  bool hasZeroStrides = false;
   auto sortedStridesIndices = getSortedStrides(nonZeroStrides);
 
   NSMutableArray<NSNumber*>* sortedStridesShape = [NSMutableArray arrayWithCapacity:nStrides];
@@ -484,8 +454,8 @@ bool isTooLargeForMPSGraph(const Tensor& tensor, bool useMPSStridedAPI) {
   MPSShape* originalSortedStridesShape = sortedStridesShape;
   bool hasNonZeroStrides = nStrides == 0 ? false : nonZeroStrides[sortedStridesIndices[nStrides - 1]] != 1;
   if (hasNonZeroStrides) {
-    originalSortedMPSShape = [sortedMPSShape copy];
-    originalSortedStridesShape = [sortedStridesShape copy];
+    originalSortedMPSShape = [[sortedMPSShape copy] autorelease];
+    originalSortedStridesShape = [[sortedStridesShape copy] autorelease];
     [sortedStridesShape addObject:[NSNumber numberWithInteger:1]];
     [sortedMPSShape addObject:[NSNumber numberWithInteger:1]];
   }
@@ -563,7 +533,19 @@ bool isTooLargeForMPSGraph(const Tensor& tensor, bool useMPSStridedAPI) {
   if ((_tensor.is_contiguous() && !_tensor.storage_offset()) || !useMPSStridedAPI || !is_macOS_15_0_or_newer) {
     auto shape = mpsShape_ ? mpsShape_ : getMPSShape(_tensor);
     check_mps_shape(shape);
-    _value = [[[MPSGraphTensorData alloc] initWithMTLBuffer:srcBuf shape:shape dataType:dataType] autorelease];
+    if (!_tensor.storage_offset()) {
+      _value = [[[MPSGraphTensorData alloc] initWithMTLBuffer:srcBuf shape:shape dataType:dataType] autorelease];
+    } else {
+      // initWithMTLBuffer:shape:dataType: has no offset parameter; use MPSNDArray to apply storage_offset.
+      // This arises e.g. for channels_last tensors with a non-zero storage offset (gatherTensorData=false path).
+      MPSNDArrayDescriptor* desc = [MPSNDArrayDescriptor descriptorWithDataType:dataType shape:shape];
+      desc.preferPackedRows = YES;
+      MPSNDArray* ndArray = [[[MPSNDArray alloc] initWithBuffer:srcBuf
+                                                         offset:_tensor.storage_offset() * _tensor.element_size()
+                                                     descriptor:desc] autorelease];
+      TORCH_INTERNAL_ASSERT(ndArray);
+      _value = [[[MPSGraphTensorData alloc] initWithMPSNDArray:ndArray] autorelease];
+    }
   } else {
     IntArrayRef view_shape;
     if (mpsShape_) {
@@ -661,6 +643,8 @@ MPSScalar getMPSScalar(const Scalar& scalar, ScalarType type) {
     case ScalarType::ComplexDouble:
       return {.size = sizeof(int64_t), .type = type, .value.cf = scalar.to<c10::complex<float>>()};
     // Unsigned types
+    case ScalarType::UInt64:
+      return {.size = sizeof(uint64_t), .type = type, .value.u = scalar.to<uint64_t>()};
     case ScalarType::UInt32:
       return {.size = sizeof(uint32_t), .type = type, .value.i = scalar.to<uint32_t>()};
     case ScalarType::UInt16:
@@ -947,16 +931,21 @@ void executeMPSAllocatorCallback(void* ptr, EventType event) override {}
   return raw_ptr;
 }
 
-class BundledShaderLibary : public MetalShaderLibrary {
+class BundledShaderLibrary : public MetalShaderLibrary {
  public:
-  BundledShaderLibary() : MetalShaderLibrary("") {}
+  BundledShaderLibrary() : MetalShaderLibrary("") {}
 
  protected:
   id<MTLLibrary> getLibrary() override {
     if (C10_UNLIKELY(!library)) {
       auto device = MPSDevice::getInstance()->device();
       NSError* error = nil;
-      library = [device newLibraryWithData:getSectionData("metal_basic") error:&error];
+#ifdef CAN_BUILD_METAL_4
+      const auto section_name = is_macos_13_or_newer(MacOSVersion::MACOS_VER_26_0_PLUS) ? "metal_40" : "metal_basic";
+#else
+      const auto section_name = "metal_basic";
+#endif
+      library = [device newLibraryWithData:getSectionData(section_name) error:&error];
       TORCH_CHECK(library, "Failed to create metal library, error: ", [[error description] UTF8String]);
     }
     return library;
@@ -992,11 +981,12 @@ static dispatch_data_t getSectionData(const std::string& name) {
 void MetalShaderLibrary::exec_unary_kernel(TensorIteratorBase& iter,
                                            const std::string& name,
                                            std::optional<c10::Scalar> alpha,
-                                           std::optional<c10::ScalarType> scalar_arg_type) {
+                                           std::optional<c10::ScalarType> scalar_arg_type,
+                                           bool supports_vec4) {
   // Decompose 64-bit tensor into 32-bit ones
   if (!iter.can_use_32bit_indexing()) {
     for (auto&& sub_iter : iter.with_32bit_indexing()) {
-      exec_unary_kernel(sub_iter, name, alpha, scalar_arg_type);
+      exec_unary_kernel(sub_iter, name, alpha, scalar_arg_type, supports_vec4);
     }
     return;
   }
@@ -1008,10 +998,12 @@ static dispatch_data_t getSectionData(const std::string& name) {
     return;
   }
   using namespace mps;
+  bool use_vec4 =
+      supports_vec4 && iter.is_contiguous() && !alpha.has_value() && at::isFloatingType(iter.common_dtype());
   const auto alpha_type = scalar_arg_type.has_value() ? scalar_arg_type.value() : iter.common_dtype();
   auto kernel_name = fmt::format("{}_{}_{}_{}{}",
                                  name,
-                                 iter.is_contiguous() ? "dense" : "strided",
+                                 use_vec4 ? "dense_vec4" : (iter.is_contiguous() ? "dense" : "strided"),
                                  scalarToMetalTypeString(outputTensor),
                                  scalarToMetalTypeString(inputTensor),
                                  alpha.has_value() ? fmt::format("_{}", scalarToMetalTypeString(alpha_type)) : "");
@@ -1026,17 +1018,22 @@ static dispatch_data_t getSectionData(const std::string& name) {
 
       [computeEncoder setComputePipelineState:cplState];
       bind_iter_tensors(computeEncoder, iter);
-      if (!iter.is_contiguous()) {
-        mtl_setArgs<2>(computeEncoder,
-                       outputTensor.sizes(),
-                       inputTensor.strides(),
-                       outputTensor.strides(),
-                       inputTensor.ndimension());
-      }
-      if (alpha) {
-        mtl_setBytes(computeEncoder, getMPSScalar(*alpha, alpha_type), iter.is_contiguous() ? 2 : 6);
+      if (use_vec4) {
+        mtl_setBytes(computeEncoder, length, 2);
+        mtl_dispatch1DJob(computeEncoder, cplState, (length + 3) / 4);
+      } else {
+        if (!iter.is_contiguous()) {
+          mtl_setArgs<2>(computeEncoder,
+                         outputTensor.sizes(),
+                         inputTensor.strides(),
+                         outputTensor.strides(),
+                         inputTensor.ndimension());
+        }
+        if (alpha) {
+          mtl_setBytes(computeEncoder, getMPSScalar(*alpha, alpha_type), iter.is_contiguous() ? 2 : 6);
+        }
+        mtl_dispatch1DJob(computeEncoder, cplState, length);
       }
-      mtl_dispatch1DJob(computeEncoder, cplState, length);
 
       getMPSProfiler().endProfileKernel(cplState);
     });
@@ -1117,7 +1114,8 @@ static dispatch_data_t getSectionData(const std::string& name) {
     }
   }
 
-  const auto alpha_type = scalar_arg_type.has_value() ? scalar_arg_type.value() : iter.common_dtype();
+  const auto alpha_type =
+      scalar_arg_type.has_value() ? scalar_arg_type.value() : (cast_needed ? out.scalar_type() : iter.common_dtype());
   const auto alpha_suffix = alpha.has_value() ? fmt::format("_{}", scalarToMetalTypeString(alpha_type)) : "";
 
   std::string kernel_name;
@@ -1330,7 +1328,7 @@ static dispatch_data_t getSectionData(const std::string& name) {
 }
 
 MetalShaderLibrary& MetalShaderLibrary::getBundledLibrary() {
-  static BundledShaderLibary l;
+  static BundledShaderLibrary l;
   return l;
 }
 
diff --git a/aten/src/ATen/native/mps/kernels/ActivationKernel.metal b/aten/src/ATen/native/mps/kernels/ActivationKernel.metal
index 7d1f3aa5bacf6..dc892dceea2d0 100644
--- a/aten/src/ATen/native/mps/kernels/ActivationKernel.metal
+++ b/aten/src/ATen/native/mps/kernels/ActivationKernel.metal
@@ -44,6 +44,27 @@ REGISTER_BINARY_ALPHA_OP(shrink_backward, float, float, float);
 REGISTER_BINARY_ALPHA_OP(shrink_backward, half, half, half);
 REGISTER_BINARY_ALPHA_OP(shrink_backward, bfloat, bfloat, bfloat);
 
+struct relu_functor {
+  template <typename T>
+  inline T operator()(const T x) {
+    return x > T(0) ? x : T(0);
+  }
+};
+
+REGISTER_UNARY_OP(relu, float, float);
+REGISTER_UNARY_OP(relu, half, half);
+REGISTER_UNARY_OP(relu, bfloat, bfloat);
+REGISTER_UNARY_OP(relu, long, long);
+REGISTER_UNARY_OP(relu, int, int);
+REGISTER_UNARY_OP(relu, short, short);
+REGISTER_UNARY_OP(relu, char, char);
+REGISTER_UNARY_OP(relu, uchar, uchar);
+REGISTER_UNARY_OP(relu, bool, bool);
+
+REGISTER_UNARY_VEC4_OP(relu, float, float);
+REGISTER_UNARY_VEC4_OP(relu, half, half);
+REGISTER_UNARY_VEC4_OP(relu, bfloat, bfloat);
+
 struct hardsigmoid_functor {
   template <typename T>
   inline T operator()(const T x) {
@@ -180,3 +201,37 @@ REGISTER_UNARY_ALPHA_OP(leaky_relu, bfloat, bfloat, bfloat);
 REGISTER_BINARY_ALPHA_OP(leaky_relu_backward, float, float, float);
 REGISTER_BINARY_ALPHA_OP(leaky_relu_backward, half, half, half);
 REGISTER_BINARY_ALPHA_OP(leaky_relu_backward, bfloat, bfloat, bfloat);
+
+struct silu_functor {
+  template <typename T>
+  inline T operator()(const T x) {
+    float xf = float(x);
+    return static_cast<T>(xf / (1.0f + ::metal::precise::exp(-xf)));
+  }
+};
+
+REGISTER_UNARY_OP(silu, float, float);
+REGISTER_UNARY_OP(silu, half, half);
+REGISTER_UNARY_OP(silu, bfloat, bfloat);
+REGISTER_UNARY_OP(silu, int, int);
+REGISTER_UNARY_OP(silu, short, short);
+REGISTER_UNARY_OP(silu, char, char);
+REGISTER_UNARY_OP(silu, uchar, uchar);
+REGISTER_UNARY_OP(silu, bool, bool);
+
+REGISTER_UNARY_VEC4_OP(silu, float, float);
+REGISTER_UNARY_VEC4_OP(silu, half, half);
+REGISTER_UNARY_VEC4_OP(silu, bfloat, bfloat);
+
+struct silu_backward_functor {
+  template <typename T>
+  inline T operator()(const T grad_output, const T self) {
+    float sf = float(self);
+    float sig = 1.0f / (1.0f + ::metal::precise::exp(-sf));
+    return static_cast<T>(float(grad_output) * sig * (1.0f + sf - sf * sig));
+  }
+};
+
+REGISTER_BINARY_OP(silu_backward, float, float);
+REGISTER_BINARY_OP(silu_backward, half, half);
+REGISTER_BINARY_OP(silu_backward, bfloat, bfloat);
diff --git a/aten/src/ATen/native/mps/kernels/Attention.metal b/aten/src/ATen/native/mps/kernels/Attention.metal
index 5a317895f508e..fb8e578dd4ed4 100644
--- a/aten/src/ATen/native/mps/kernels/Attention.metal
+++ b/aten/src/ATen/native/mps/kernels/Attention.metal
@@ -20,6 +20,7 @@ template <typename T, int D, int V = D>
     const device bool* mask [[buffer(9)]],
     const constant uint3& mask_strides [[buffer(10)]],
     const constant bool& has_mask [[buffer(11)]],
+    const constant uint4& qkv_batch_strides_heads [[buffer(12)]],
     uint3 tid [[threadgroup_position_in_grid]],
     uint3 tpg [[threadgroups_per_grid]],
     uint simd_gid [[simdgroup_index_in_threadgroup]],
@@ -30,13 +31,17 @@ template <typename T, int D, int V = D>
   constexpr uint v_per_thread = V / BD;
   const uint q_head_stride = qkv_head_strides.x;
   const uint q_seq_stride = qkv_seq_strides.x;
+  const uint q_batch_stride = qkv_batch_strides_heads.x;
   const uint k_head_stride = qkv_head_strides.y;
   const uint k_seq_stride = qkv_seq_strides.y;
+  const uint k_batch_stride = qkv_batch_strides_heads.y;
   const uint v_head_stride = qkv_head_strides.z;
   const uint v_seq_stride = qkv_seq_strides.z;
-  const uint mask_head_stride = mask_strides.x;
-  const uint mask_kv_seq_stride = mask_strides.y;
-  const uint mask_q_seq_stride = mask_strides.z;
+  const uint v_batch_stride = qkv_batch_strides_heads.z;
+  const uint num_heads = qkv_batch_strides_heads.w;
+  const uint mask_kv_seq_stride = mask_strides.x;
+  const uint mask_q_seq_stride = mask_strides.y;
+  const uint mask_head_stride = mask_strides.z;
   uint inner_k_stride = BN * int(k_seq_stride);
   uint inner_v_stride = BN * int(v_seq_stride);
 
@@ -50,21 +55,22 @@ template <typename T, int D, int V = D>
   threadgroup U max_scores[BN];
   threadgroup U sum_exp_scores[BN];
 
-  // Adjust positions
-  const int head_idx = tid.x;
+  const int bh_idx = tid.x;
+  const int batch_idx = bh_idx / int(num_heads);
+  const int head_idx = bh_idx - batch_idx * int(num_heads);
   const int q_seq_idx = tid.y;
   const int kv_head_idx = head_idx / gqa_factor;
   const int Q = tpg.y;
-  const int group_offset = head_idx * Q + q_seq_idx;
+  const int group_offset = bh_idx * Q + q_seq_idx;
   const int o_offset = group_offset;
-  queries += head_idx * q_head_stride + q_seq_idx * q_seq_stride +
-      simd_lid * qk_per_thread;
-  keys += kv_head_idx * k_head_stride + simd_gid * k_seq_stride +
-      simd_lid * qk_per_thread;
-  values += kv_head_idx * v_head_stride + simd_gid * v_seq_stride +
-      simd_lid * v_per_thread;
+  queries += batch_idx * q_batch_stride + head_idx * q_head_stride +
+      q_seq_idx * q_seq_stride + simd_lid * qk_per_thread;
+  keys += batch_idx * k_batch_stride + kv_head_idx * k_head_stride +
+      simd_gid * k_seq_stride + simd_lid * qk_per_thread;
+  values += batch_idx * v_batch_stride + kv_head_idx * v_head_stride +
+      simd_gid * v_seq_stride + simd_lid * v_per_thread;
   if (has_mask) {
-    mask += head_idx * mask_head_stride + simd_gid * mask_kv_seq_stride +
+    mask += bh_idx * mask_head_stride + simd_gid * mask_kv_seq_stride +
         q_seq_idx * mask_q_seq_stride;
   }
 
@@ -164,6 +170,7 @@ template <typename T, int D, int V = D>
     const device bool* mask [[buffer(11)]],
     const constant uint3& mask_strides [[buffer(12)]],
     const constant bool& has_mask [[buffer(13)]],
+    const constant uint4& qkv_batch_strides_heads [[buffer(14)]],
     uint3 tid [[threadgroup_position_in_grid]],
     uint3 tpg [[threadgroups_per_grid]],
     uint simd_gid [[simdgroup_index_in_threadgroup]],
@@ -174,10 +181,14 @@ template <typename T, int D, int V = D>
   constexpr int v_per_thread = V / BD;
   const int q_head_stride = qkv_head_strides.x;
   const int q_seq_stride = qkv_seq_strides.x;
+  const int q_batch_stride = qkv_batch_strides_heads.x;
   const int k_head_stride = qkv_head_strides.y;
   const int k_seq_stride = qkv_seq_strides.y;
+  const int k_batch_stride = qkv_batch_strides_heads.y;
   const int v_head_stride = qkv_head_strides.z;
   const int v_seq_stride = qkv_seq_strides.z;
+  const int v_batch_stride = qkv_batch_strides_heads.z;
+  const int num_heads = qkv_batch_strides_heads.w;
   const int mask_kv_seq_stride = mask_strides.x;
   const int mask_q_seq_stride = mask_strides.y;
   const int mask_head_stride = mask_strides.z;
@@ -195,22 +206,23 @@ template <typename T, int D, int V = D>
   threadgroup U max_scores[BN];
   threadgroup U sum_exp_scores[BN];
 
-  // Adjust positions
   const int block_idx = tid.z;
-  const int head_idx = tid.x;
+  const int bh_idx = tid.x;
+  const int batch_idx = bh_idx / int(num_heads);
+  const int head_idx = bh_idx - batch_idx * int(num_heads);
   const int q_seq_idx = tid.y;
-  const int o_offset = head_idx * tpg.y + q_seq_idx;
+  const int o_offset = bh_idx * tpg.y + q_seq_idx;
   const int kv_head_idx = head_idx / gqa_factor;
 
-  queries += head_idx * q_head_stride + q_seq_idx * q_seq_stride +
-      simd_lid * qk_per_thread;
-  keys += kv_head_idx * k_head_stride +
+  queries += batch_idx * q_batch_stride + head_idx * q_head_stride +
+      q_seq_idx * q_seq_stride + simd_lid * qk_per_thread;
+  keys += batch_idx * k_batch_stride + kv_head_idx * k_head_stride +
       (block_idx * BN + simd_gid) * k_seq_stride + simd_lid * qk_per_thread;
-  values += kv_head_idx * v_head_stride +
+  values += batch_idx * v_batch_stride + kv_head_idx * v_head_stride +
       (block_idx * BN + simd_gid) * v_seq_stride + simd_lid * v_per_thread;
   out += o_offset * blocks * V + block_idx * V + simd_lid * v_per_thread;
   if (has_mask) {
-    mask += head_idx * mask_head_stride +
+    mask += bh_idx * mask_head_stride +
         (block_idx * BN + simd_gid) * mask_kv_seq_stride +
         q_seq_idx * mask_q_seq_stride;
   }
@@ -524,48 +536,50 @@ kernel void attention(
   }
 }
 
-#define INSTANTIATE_SDPA_VECTOR(DTYPE, QK_DIM, VALUE_DIM)   \
-  template [[host_name("sdpa_vector_" #DTYPE "_" #QK_DIM    \
-                       "_" #VALUE_DIM)]] kernel void        \
-  sdpa_vector<DTYPE, QK_DIM, VALUE_DIM>(                    \
-      const device DTYPE* queries [[buffer(0)]],            \
-      const device DTYPE* keys [[buffer(1)]],               \
-      const device DTYPE* values [[buffer(2)]],             \
-      device DTYPE* out [[buffer(3)]],                      \
-      const constant uint& gqa_factor [[buffer(4)]],        \
-      const constant uint& N [[buffer(5)]],                 \
-      const constant uint3& qkv_head_strides [[buffer(6)]], \
-      const constant uint3& qkv_seq_strides [[buffer(7)]],  \
-      const constant float& scale [[buffer(8)]],            \
-      const device bool* mask [[buffer(9)]],                \
-      const constant uint3& mask_strides [[buffer(10)]],    \
-      const constant bool& has_mask [[buffer(11)]],         \
-      uint3 tid [[threadgroup_position_in_grid]],           \
-      uint3 tpg [[threadgroups_per_grid]],                  \
-      uint simd_gid [[simdgroup_index_in_threadgroup]],     \
+#define INSTANTIATE_SDPA_VECTOR(DTYPE, QK_DIM, VALUE_DIM)           \
+  template [[host_name("sdpa_vector_" #DTYPE "_" #QK_DIM            \
+                       "_" #VALUE_DIM)]] kernel void                \
+  sdpa_vector<DTYPE, QK_DIM, VALUE_DIM>(                            \
+      const device DTYPE* queries [[buffer(0)]],                    \
+      const device DTYPE* keys [[buffer(1)]],                       \
+      const device DTYPE* values [[buffer(2)]],                     \
+      device DTYPE* out [[buffer(3)]],                              \
+      const constant uint& gqa_factor [[buffer(4)]],                \
+      const constant uint& N [[buffer(5)]],                         \
+      const constant uint3& qkv_head_strides [[buffer(6)]],         \
+      const constant uint3& qkv_seq_strides [[buffer(7)]],          \
+      const constant float& scale [[buffer(8)]],                    \
+      const device bool* mask [[buffer(9)]],                        \
+      const constant uint3& mask_strides [[buffer(10)]],            \
+      const constant bool& has_mask [[buffer(11)]],                 \
+      const constant uint4& qkv_batch_strides_heads [[buffer(12)]], \
+      uint3 tid [[threadgroup_position_in_grid]],                   \
+      uint3 tpg [[threadgroups_per_grid]],                          \
+      uint simd_gid [[simdgroup_index_in_threadgroup]],             \
       uint simd_lid [[thread_index_in_simdgroup]]);
 
-#define INSTANTIATE_SDPA_VECTOR_2PASS_1(DTYPE, QK_DIM, VALUE_DIM) \
-  template [[host_name("sdpa_vector_2pass_1_" #DTYPE "_" #QK_DIM  \
-                       "_" #VALUE_DIM)]] kernel void              \
-  sdpa_vector_2pass_1<DTYPE, QK_DIM, VALUE_DIM>(                  \
-      const device DTYPE* queries [[buffer(0)]],                  \
-      const device DTYPE* keys [[buffer(1)]],                     \
-      const device DTYPE* values [[buffer(2)]],                   \
-      device DTYPE* out [[buffer(3)]],                            \
-      device float* sums [[buffer(4)]],                           \
-      device float* maxs [[buffer(5)]],                           \
-      const constant uint& gqa_factor [[buffer(6)]],              \
-      const constant uint& N [[buffer(7)]],                       \
-      const constant uint3& qkv_head_strides [[buffer(8)]],       \
-      const constant uint3& qkv_seq_strides [[buffer(9)]],        \
-      const constant float& scale [[buffer(10)]],                 \
-      const device bool* mask [[buffer(11)]],                     \
-      const constant uint3& mask_strides [[buffer(12)]],          \
-      const constant bool& has_mask [[buffer(13)]],               \
-      uint3 tid [[threadgroup_position_in_grid]],                 \
-      uint3 tpg [[threadgroups_per_grid]],                        \
-      uint simd_gid [[simdgroup_index_in_threadgroup]],           \
+#define INSTANTIATE_SDPA_VECTOR_2PASS_1(DTYPE, QK_DIM, VALUE_DIM)   \
+  template [[host_name("sdpa_vector_2pass_1_" #DTYPE "_" #QK_DIM    \
+                       "_" #VALUE_DIM)]] kernel void                \
+  sdpa_vector_2pass_1<DTYPE, QK_DIM, VALUE_DIM>(                    \
+      const device DTYPE* queries [[buffer(0)]],                    \
+      const device DTYPE* keys [[buffer(1)]],                       \
+      const device DTYPE* values [[buffer(2)]],                     \
+      device DTYPE* out [[buffer(3)]],                              \
+      device float* sums [[buffer(4)]],                             \
+      device float* maxs [[buffer(5)]],                             \
+      const constant uint& gqa_factor [[buffer(6)]],                \
+      const constant uint& N [[buffer(7)]],                         \
+      const constant uint3& qkv_head_strides [[buffer(8)]],         \
+      const constant uint3& qkv_seq_strides [[buffer(9)]],          \
+      const constant float& scale [[buffer(10)]],                   \
+      const device bool* mask [[buffer(11)]],                       \
+      const constant uint3& mask_strides [[buffer(12)]],            \
+      const constant bool& has_mask [[buffer(13)]],                 \
+      const constant uint4& qkv_batch_strides_heads [[buffer(14)]], \
+      uint3 tid [[threadgroup_position_in_grid]],                   \
+      uint3 tpg [[threadgroups_per_grid]],                          \
+      uint simd_gid [[simdgroup_index_in_threadgroup]],             \
       uint simd_lid [[thread_index_in_simdgroup]]);
 
 #define INSTANTIATE_SDPA_VECTOR_AGGREGATION(DTYPE, VALUE_DIM)                 \
diff --git a/aten/src/ATen/native/mps/kernels/BinaryKernel.metal b/aten/src/ATen/native/mps/kernels/BinaryKernel.metal
index 15d259cc9b27a..1def63fff16ae 100644
--- a/aten/src/ATen/native/mps/kernels/BinaryKernel.metal
+++ b/aten/src/ATen/native/mps/kernels/BinaryKernel.metal
@@ -63,14 +63,14 @@ struct fmin_functor {
 struct maximum_functor {
   template <typename T>
   inline T operator()(const T a, const T b) {
-    return max(a, b);
+    return c10::metal::max(a, b);
   }
 };
 
 struct minimum_functor {
   template <typename T>
   inline T operator()(const T a, const T b) {
-    return min(a, b);
+    return c10::metal::min(a, b);
   }
 };
 
@@ -122,6 +122,20 @@ struct logaddexp2_functor {
   }
 };
 
+struct xlogy_functor {
+  template <typename T, enable_if_t<is_floating_point_v<T>, bool> = true>
+  inline T operator()(const T a, const T b) {
+    return static_cast<T>(c10::metal::xlogy(a, b));
+  }
+  template <typename T, enable_if_t<is_integral_v<T>, bool> = true>
+  inline float operator()(const T a, const T b) {
+    return c10::metal::xlogy(float(a), float(b));
+  }
+  inline float operator()(const bool a, const bool b) {
+    return (a && !b) ? -INFINITY : 0;
+  }
+};
+
 struct xlog1py_functor {
   template <typename T, enable_if_t<is_floating_point_v<T>, bool> = true>
   inline T operator()(const T a, const T b) {
@@ -449,6 +463,8 @@ REGISTER_FLOAT_BINARY_OP(logaddexp);
 REGISTER_INT2FLOAT_BINARY_OP(logaddexp);
 REGISTER_FLOAT_BINARY_OP(logaddexp2);
 REGISTER_INT2FLOAT_BINARY_OP(logaddexp2);
+REGISTER_FLOAT_BINARY_OP(xlogy);
+REGISTER_INT2FLOAT_BINARY_OP(xlogy);
 REGISTER_FLOAT_BINARY_OP(xlog1py);
 REGISTER_INT2FLOAT_BINARY_OP(xlog1py);
 REGISTER_FLOAT_BINARY_OP(chebyshev_polynomial_t);
@@ -544,3 +560,176 @@ REGISTER_BINARY_ALPHA_OP(sub_alpha, float2, float2, float2);
 REGISTER_BINARY_ALPHA_OP(sub_alpha, half2, half2, half2);
 REGISTER_BINARY_ALPHA_OP(lerp_alpha, float2, float2, float2);
 REGISTER_BINARY_ALPHA_OP(lerp_alpha, half2, half2, half2);
+
+// lerp with tensor weight: lerp(s, e, w) = fma(w, e - s, s)
+template <typename T>
+inline T lerp_op(T s, T e, T w) {
+  return fma(w, e - s, s);
+}
+
+inline bfloat lerp_op(bfloat s, bfloat e, bfloat w) {
+  return static_cast<bfloat>(fma(float(w), float(e) - float(s), float(s)));
+}
+
+inline long lerp_op(long s, long e, long w) {
+  return s + w * (e - s);
+}
+
+inline float2 lerp_op(float2 s, float2 e, float2 w) {
+  return s + mul(w, e - s);
+}
+
+template <typename T>
+kernel void lerp_tensor_dense(
+    device T* out [[buffer(0)]],
+    device const T* self [[buffer(1)]],
+    device const T* end [[buffer(2)]],
+    device const T* weight [[buffer(3)]],
+    uint tid [[thread_position_in_grid]]) {
+  out[tid] = lerp_op(self[tid], end[tid], weight[tid]);
+}
+
+// Scalar weight broadcast: self/end/out contiguous, weight is a single element
+template <typename T>
+kernel void lerp_tensor_scalar_weight(
+    device T* out [[buffer(0)]],
+    device const T* self [[buffer(1)]],
+    device const T* end [[buffer(2)]],
+    device const T& weight [[buffer(3)]],
+    uint tid [[thread_position_in_grid]]) {
+  out[tid] = lerp_op(self[tid], end[tid], weight);
+}
+
+// 2D strided: coordinates from 2D dispatch, no integer division
+template <typename T>
+kernel void lerp_tensor_strided_2d(
+    device void* out_ptr [[buffer(0)]],
+    constant void* self_ptr [[buffer(1)]],
+    constant void* end_ptr [[buffer(2)]],
+    constant void* weight_ptr [[buffer(3)]],
+    constant long* out_strides [[buffer(4)]],
+    constant long* self_strides [[buffer(5)]],
+    constant long* end_strides [[buffer(6)]],
+    constant long* weight_strides [[buffer(7)]],
+    uint2 tid [[thread_position_in_grid]]) {
+  int out_off =
+      int(tid.x) * int(out_strides[0]) + int(tid.y) * int(out_strides[1]);
+  int self_off =
+      int(tid.x) * int(self_strides[0]) + int(tid.y) * int(self_strides[1]);
+  int end_off =
+      int(tid.x) * int(end_strides[0]) + int(tid.y) * int(end_strides[1]);
+  int wt_off =
+      int(tid.x) * int(weight_strides[0]) + int(tid.y) * int(weight_strides[1]);
+  ref_at_offs<T>(out_ptr, long(out_off)) = lerp_op(
+      val_at_offs<T>(self_ptr, long(self_off)),
+      val_at_offs<T>(end_ptr, long(end_off)),
+      val_at_offs<T>(weight_ptr, long(wt_off)));
+}
+
+// 3D strided: coordinates from 3D dispatch, no integer division
+template <typename T>
+kernel void lerp_tensor_strided_3d(
+    device void* out_ptr [[buffer(0)]],
+    constant void* self_ptr [[buffer(1)]],
+    constant void* end_ptr [[buffer(2)]],
+    constant void* weight_ptr [[buffer(3)]],
+    constant long* out_strides [[buffer(4)]],
+    constant long* self_strides [[buffer(5)]],
+    constant long* end_strides [[buffer(6)]],
+    constant long* weight_strides [[buffer(7)]],
+    uint3 tid [[thread_position_in_grid]]) {
+  int out_off = int(tid.x) * int(out_strides[0]) +
+      int(tid.y) * int(out_strides[1]) + int(tid.z) * int(out_strides[2]);
+  int self_off = int(tid.x) * int(self_strides[0]) +
+      int(tid.y) * int(self_strides[1]) + int(tid.z) * int(self_strides[2]);
+  int end_off = int(tid.x) * int(end_strides[0]) +
+      int(tid.y) * int(end_strides[1]) + int(tid.z) * int(end_strides[2]);
+  int wt_off = int(tid.x) * int(weight_strides[0]) +
+      int(tid.y) * int(weight_strides[1]) + int(tid.z) * int(weight_strides[2]);
+  ref_at_offs<T>(out_ptr, long(out_off)) = lerp_op(
+      val_at_offs<T>(self_ptr, long(self_off)),
+      val_at_offs<T>(end_ptr, long(end_off)),
+      val_at_offs<T>(weight_ptr, long(wt_off)));
+}
+
+template <typename T>
+kernel void lerp_tensor_strided(
+    device void* out_ptr [[buffer(0)]],
+    constant void* self_ptr [[buffer(1)]],
+    constant void* end_ptr [[buffer(2)]],
+    constant void* weight_ptr [[buffer(3)]],
+    constant long* sizes [[buffer(4)]],
+    constant long* out_strides [[buffer(5)]],
+    constant long* self_strides [[buffer(6)]],
+    constant long* end_strides [[buffer(7)]],
+    constant long* weight_strides [[buffer(8)]],
+    constant uint& ndim [[buffer(9)]],
+    uint tid [[thread_position_in_grid]]) {
+  int pos[max_ndim];
+  pos_from_thread_index(int(tid), pos, sizes, ndim);
+  auto self_off = offset_from_coord(pos, self_strides, ndim);
+  auto end_off = offset_from_coord(pos, end_strides, ndim);
+  auto weight_off = offset_from_coord(pos, weight_strides, ndim);
+  auto out_off = offset_from_coord(pos, out_strides, ndim);
+  ref_at_offs<T>(out_ptr, out_off) = lerp_op(
+      val_at_offs<T>(self_ptr, self_off),
+      val_at_offs<T>(end_ptr, end_off),
+      val_at_offs<T>(weight_ptr, weight_off));
+}
+
+#define INSTANTIATE_LERP(DTYPE)                                           \
+  template [[host_name("lerp_tensor_dense_" #DTYPE)]] kernel void         \
+  lerp_tensor_dense<DTYPE>(                                               \
+      device DTYPE * out [[buffer(0)]],                                   \
+      device const DTYPE* self [[buffer(1)]],                             \
+      device const DTYPE* end [[buffer(2)]],                              \
+      device const DTYPE* weight [[buffer(3)]],                           \
+      uint tid [[thread_position_in_grid]]);                              \
+  template [[host_name("lerp_tensor_scalar_weight_" #DTYPE)]] kernel void \
+  lerp_tensor_scalar_weight<DTYPE>(                                       \
+      device DTYPE * out [[buffer(0)]],                                   \
+      device const DTYPE* self [[buffer(1)]],                             \
+      device const DTYPE* end [[buffer(2)]],                              \
+      device const DTYPE& weight [[buffer(3)]],                           \
+      uint tid [[thread_position_in_grid]]);                              \
+  template [[host_name("lerp_tensor_strided_2d_" #DTYPE)]] kernel void    \
+  lerp_tensor_strided_2d<DTYPE>(                                          \
+      device void* out_ptr [[buffer(0)]],                                 \
+      constant void* self_ptr [[buffer(1)]],                              \
+      constant void* end_ptr [[buffer(2)]],                               \
+      constant void* weight_ptr [[buffer(3)]],                            \
+      constant long* out_strides [[buffer(4)]],                           \
+      constant long* self_strides [[buffer(5)]],                          \
+      constant long* end_strides [[buffer(6)]],                           \
+      constant long* weight_strides [[buffer(7)]],                        \
+      uint2 tid [[thread_position_in_grid]]);                             \
+  template [[host_name("lerp_tensor_strided_3d_" #DTYPE)]] kernel void    \
+  lerp_tensor_strided_3d<DTYPE>(                                          \
+      device void* out_ptr [[buffer(0)]],                                 \
+      constant void* self_ptr [[buffer(1)]],                              \
+      constant void* end_ptr [[buffer(2)]],                               \
+      constant void* weight_ptr [[buffer(3)]],                            \
+      constant long* out_strides [[buffer(4)]],                           \
+      constant long* self_strides [[buffer(5)]],                          \
+      constant long* end_strides [[buffer(6)]],                           \
+      constant long* weight_strides [[buffer(7)]],                        \
+      uint3 tid [[thread_position_in_grid]]);                             \
+  template [[host_name("lerp_tensor_strided_" #DTYPE)]] kernel void       \
+  lerp_tensor_strided<DTYPE>(                                             \
+      device void* out_ptr [[buffer(0)]],                                 \
+      constant void* self_ptr [[buffer(1)]],                              \
+      constant void* end_ptr [[buffer(2)]],                               \
+      constant void* weight_ptr [[buffer(3)]],                            \
+      constant long* sizes [[buffer(4)]],                                 \
+      constant long* out_strides [[buffer(5)]],                           \
+      constant long* self_strides [[buffer(6)]],                          \
+      constant long* end_strides [[buffer(7)]],                           \
+      constant long* weight_strides [[buffer(8)]],                        \
+      constant uint& ndim [[buffer(9)]],                                  \
+      uint tid [[thread_position_in_grid]]);
+
+INSTANTIATE_LERP(float);
+INSTANTIATE_LERP(half);
+INSTANTIATE_LERP(bfloat);
+INSTANTIATE_LERP(float2);
+INSTANTIATE_LERP(long);
diff --git a/aten/src/ATen/native/mps/kernels/ConstantKernel.metal b/aten/src/ATen/native/mps/kernels/ConstantKernel.metal
new file mode 100644
index 0000000000000..1b4c1277b6448
--- /dev/null
+++ b/aten/src/ATen/native/mps/kernels/ConstantKernel.metal
@@ -0,0 +1,102 @@
+#include <c10/metal/indexing.h>
+#include <metal_stdlib>
+
+using namespace metal;
+
+template <typename T>
+kernel void fill_scalar_dense(
+    device T* out [[buffer(0)]],
+    constant T& fill_val [[buffer(1)]],
+    constant uint& numel [[buffer(2)]],
+    uint index [[thread_position_in_grid]]) {
+  if (index < numel)
+    out[index] = fill_val;
+}
+
+// Single-byte types: each thread fills 4 elements via vec<T,4> writes.
+template <typename T>
+kernel void fill_scalar_dense_vec4(
+    device T* out [[buffer(0)]],
+    constant T& fill_val [[buffer(1)]],
+    constant uint& numel [[buffer(2)]],
+    uint index [[thread_position_in_grid]]) {
+  uint base = index * 4;
+  if (base + 4 <= numel) {
+    *(device vec<T, 4>*)(out + base) = vec<T, 4>(fill_val);
+  } else {
+    for (uint i = base; i < numel; i++)
+      out[i] = fill_val;
+  }
+}
+
+#define REGISTER_FILL_OP(T)                                   \
+  template [[host_name("fill_scalar_dense_" #T)]] kernel void \
+  fill_scalar_dense<T>(device T*, constant T&, constant uint&, uint)
+
+REGISTER_FILL_OP(float);
+REGISTER_FILL_OP(half);
+REGISTER_FILL_OP(bfloat);
+REGISTER_FILL_OP(long);
+REGISTER_FILL_OP(ulong);
+REGISTER_FILL_OP(int);
+REGISTER_FILL_OP(uint);
+REGISTER_FILL_OP(short);
+REGISTER_FILL_OP(ushort);
+REGISTER_FILL_OP(float2);
+REGISTER_FILL_OP(half2);
+
+// Byte types use the vec4 variant under the same naming convention.
+#define REGISTER_FILL_BYTE_OP(T)                              \
+  template [[host_name("fill_scalar_dense_" #T)]] kernel void \
+  fill_scalar_dense_vec4<T>(device T*, constant T&, constant uint&, uint)
+
+REGISTER_FILL_BYTE_OP(char);
+REGISTER_FILL_BYTE_OP(uchar);
+REGISTER_FILL_BYTE_OP(bool);
+
+// 2D dispatch: tid.y = dim-0 index (no division), tid.x = linear index for
+// dims 1..ndim-1. For an N-dim tensor this requires N-1 divisions instead of N,
+// and consecutive threads in x access consecutive addresses in the innermost
+// dimension (coalesced writes).
+// Strides from TensorIterator are in bytes, use ref_at_offs<T>(out+offs)
+template <typename T>
+kernel void fill_scalar_strided(
+    device void* out [[buffer(0)]],
+    constant T& fill_val [[buffer(1)]],
+    constant long* sizes [[buffer(2)]],
+    constant long* strides [[buffer(3)]],
+    constant uint& ndim [[buffer(4)]],
+    uint2 tid [[thread_position_in_grid]]) {
+  long offset = long(tid.y) * strides[0];
+  uint inner = tid.x;
+  for (uint i = 1; i < ndim; i++) {
+    offset += long(inner % uint(sizes[i])) * strides[i];
+    inner /= uint(sizes[i]);
+  }
+  c10::metal::ref_at_offs<T>(out, offset) = fill_val;
+}
+
+#define REGISTER_FILL_STRIDED_OP(T)                             \
+  template [[host_name("fill_scalar_strided_" #T)]] kernel void \
+  fill_scalar_strided<T>(                                       \
+      device void*,                                             \
+      constant T&,                                              \
+      constant long*,                                           \
+      constant long*,                                           \
+      constant uint&,                                           \
+      uint2)
+
+REGISTER_FILL_STRIDED_OP(float);
+REGISTER_FILL_STRIDED_OP(half);
+REGISTER_FILL_STRIDED_OP(bfloat);
+REGISTER_FILL_STRIDED_OP(long);
+REGISTER_FILL_STRIDED_OP(ulong);
+REGISTER_FILL_STRIDED_OP(int);
+REGISTER_FILL_STRIDED_OP(uint);
+REGISTER_FILL_STRIDED_OP(short);
+REGISTER_FILL_STRIDED_OP(ushort);
+REGISTER_FILL_STRIDED_OP(char);
+REGISTER_FILL_STRIDED_OP(uchar);
+REGISTER_FILL_STRIDED_OP(bool);
+REGISTER_FILL_STRIDED_OP(float2);
+REGISTER_FILL_STRIDED_OP(half2);
diff --git a/aten/src/ATen/native/mps/kernels/CrossKernel.metal b/aten/src/ATen/native/mps/kernels/CrossKernel.metal
index bceae51c02db4..c8b756b9e7b70 100644
--- a/aten/src/ATen/native/mps/kernels/CrossKernel.metal
+++ b/aten/src/ATen/native/mps/kernels/CrossKernel.metal
@@ -1,69 +1,104 @@
-#include <metal_array>
-
+#include <c10/metal/indexing.h>
+#include <c10/metal/utils.h>
 #include <metal_stdlib>
 using namespace metal;
+using namespace c10::metal;
 
-#define REGISTER_CROSS_FUNC(DTYPE)                       \
-  static inline DTYPE##3 cross(DTYPE##3 x, DTYPE##3 y) { \
-    DTYPE##3 out;                                        \
-    out.x = x.y * y.z - x.z * y.y;                       \
-    out.y = x.z * y.x - x.x * y.z;                       \
-    out.z = x.x * y.y - x.y * y.x;                       \
-    return out;                                          \
-  }
-
-// Metal only supports half and float for native cross implementation.
-// For all the other data types, implement cross manually.
-REGISTER_CROSS_FUNC(int);
-REGISTER_CROSS_FUNC(long);
-REGISTER_CROSS_FUNC(short);
-REGISTER_CROSS_FUNC(char);
-REGISTER_CROSS_FUNC(uchar);
-REGISTER_CROSS_FUNC(bool);
-REGISTER_CROSS_FUNC(bfloat);
+template <typename T>
+inline void cross_product(
+    device T* out,
+    long os,
+    constant T* x,
+    long xs,
+    constant T* y,
+    long ys) {
+  out[0] = mul(x[xs], y[2 * ys]) - mul(x[2 * xs], y[ys]);
+  out[os] = mul(x[2 * xs], y[0]) - mul(x[0], y[2 * ys]);
+  out[2 * os] = mul(x[0], y[ys]) - mul(x[xs], y[0]);
+}
 
-template <typename T, typename U>
-kernel void cross(
-    constant void* input_ [[buffer(0)]],
-    constant void* other_ [[buffer(1)]],
-    device void* out_ [[buffer(2)]],
-    constant uint3* offsets [[buffer(3)]],
-    constant int64_t& outStride [[buffer(4)]],
-    constant int64_t& inputStride [[buffer(5)]],
-    constant int64_t& otherStride [[buffer(6)]],
+// Handles any contiguous layout regardless of which dim is the cross dim.
+// For cross-dim element stride B, thread tid covers the triple whose base is
+//   base = tid + 2*(tid/B)*B
+// i.e. B=numThreads (outermost) → base=tid; B=1 (innermost) → base=3*tid.
+template <typename T>
+kernel void cross_dense(
+    device T* out [[buffer(0)]],
+    constant T* input [[buffer(1)]],
+    constant T* other [[buffer(2)]],
+    constant uint& dim_stride [[buffer(3)]],
     uint tid [[thread_position_in_grid]]) {
-  device T* out = (device T*)((device uint8_t*)out_ + offsets[tid].x);
-  constant T* input = (constant T*)((constant uint8_t*)input_ + offsets[tid].y);
-  constant T* other = (constant T*)((constant uint8_t*)other_ + offsets[tid].z);
-
-  const U x = {
-      input[0 * inputStride], input[1 * inputStride], input[2 * inputStride]};
-  const U y = {
-      other[0 * otherStride], other[1 * otherStride], other[2 * otherStride]};
-  const U res = cross(x, y);
+  uint base = tid + 2 * (tid / dim_stride) * dim_stride;
+  cross_product(
+      out + base,
+      long(dim_stride),
+      input + base,
+      long(dim_stride),
+      other + base,
+      long(dim_stride));
+}
 
-  out[0 * outStride] = res.x;
-  out[1 * outStride] = res.y;
-  out[2 * outStride] = res.z;
+template <typename T>
+kernel void cross_strided(
+    device T* out [[buffer(0)]],
+    constant T* input [[buffer(1)]],
+    constant T* other [[buffer(2)]],
+    constant long* squashed_sizes [[buffer(3)]],
+    constant long* out_strides [[buffer(4)]],
+    constant long* input_strides [[buffer(5)]],
+    constant long* other_strides [[buffer(6)]],
+    constant uint& ndim [[buffer(7)]],
+    constant uint& dim [[buffer(8)]],
+    uint tid [[thread_position_in_grid]]) {
+  int pos[max_ndim];
+  pos_from_thread_index(int(tid), pos, squashed_sizes, ndim - 1);
+  long out_offs = 0, in_offs = 0, oth_offs = 0;
+  uint p = 0;
+  for (uint d = 0; d < ndim; d++) {
+    if (d == dim)
+      continue;
+    out_offs += pos[p] * out_strides[d];
+    in_offs += pos[p] * input_strides[d];
+    oth_offs += pos[p] * other_strides[d];
+    p++;
+  }
+  cross_product(
+      out + out_offs,
+      out_strides[dim],
+      input + in_offs,
+      input_strides[dim],
+      other + oth_offs,
+      other_strides[dim]);
 }
 
-#define REGISTER_CROSS_OP(DTYPE)                                              \
-  template [[host_name("cross_" #DTYPE)]] kernel void cross<DTYPE, DTYPE##3>( \
-      constant void* input_ [[buffer(0)]],                                    \
-      constant void* other_ [[buffer(1)]],                                    \
-      device void* out_ [[buffer(2)]],                                        \
-      constant uint3* offsets [[buffer(3)]],                                  \
-      constant int64_t& outStride [[buffer(4)]],                              \
-      constant int64_t& inputStride [[buffer(5)]],                            \
-      constant int64_t& otherStride [[buffer(6)]],                            \
-      uint tid [[thread_position_in_grid]]);
+#define REGISTER_CROSS_OP(DTYPE)                              \
+  template [[host_name("cross_dense_" #DTYPE)]] kernel void   \
+  cross_dense<DTYPE>(                                         \
+      device DTYPE * out [[buffer(0)]],                       \
+      constant DTYPE * input [[buffer(1)]],                   \
+      constant DTYPE * other [[buffer(2)]],                   \
+      constant uint & dim_stride [[buffer(3)]],               \
+      uint tid [[thread_position_in_grid]]);                  \
+  template [[host_name("cross_strided_" #DTYPE)]] kernel void \
+  cross_strided<DTYPE>(                                       \
+      device DTYPE * out [[buffer(0)]],                       \
+      constant DTYPE * input [[buffer(1)]],                   \
+      constant DTYPE * other [[buffer(2)]],                   \
+      constant long* squashed_sizes [[buffer(3)]],            \
+      constant long* out_strides [[buffer(4)]],               \
+      constant long* input_strides [[buffer(5)]],             \
+      constant long* other_strides [[buffer(6)]],             \
+      constant uint& ndim [[buffer(7)]],                      \
+      constant uint& dim [[buffer(8)]],                       \
+      uint tid [[thread_position_in_grid]])
 
 REGISTER_CROSS_OP(float);
 REGISTER_CROSS_OP(half);
-REGISTER_CROSS_OP(int);
+REGISTER_CROSS_OP(bfloat);
+REGISTER_CROSS_OP(float2);
+REGISTER_CROSS_OP(half2);
 REGISTER_CROSS_OP(long);
+REGISTER_CROSS_OP(int);
 REGISTER_CROSS_OP(short);
 REGISTER_CROSS_OP(char);
 REGISTER_CROSS_OP(uchar);
-REGISTER_CROSS_OP(bool);
-REGISTER_CROSS_OP(bfloat);
diff --git a/aten/src/ATen/native/mps/kernels/Distributions.metal b/aten/src/ATen/native/mps/kernels/Distributions.metal
index 6c9a2722f378d..2c8a9404ab533 100644
--- a/aten/src/ATen/native/mps/kernels/Distributions.metal
+++ b/aten/src/ATen/native/mps/kernels/Distributions.metal
@@ -1,4 +1,5 @@
 #include <c10/metal/random.h>
+#include <c10/metal/special_math.h>
 #include <metal_stdlib>
 
 using namespace metal;
@@ -49,6 +50,26 @@ kernel void geometric(
   output[tid] = static_cast<T>(result);
 }
 
+template <typename T>
+kernel void exponential(
+    device T* output [[buffer(0)]],
+    constant float2& params [[buffer(1)]],
+    constant long2& seed_base_offset [[buffer(2)]],
+    constant uint& numel [[buffer(3)]],
+    uint tid [[thread_position_in_grid]]) {
+  uint base = tid * 4;
+  uint4 raw =
+      c10::metal::philox4::rand(seed_base_offset.x, seed_base_offset.y + tid);
+  float lambda = params.x;
+  uint count = min(4u, numel - base);
+  for (uint i = 0; i < count; ++i) {
+    float u = clamp(
+        c10::metal::detail::uint32_to_uniform_float(raw[i]), eps, 1.0f - eps);
+    output[base + i] =
+        static_cast<T>(-::metal::precise::log(1.0f - u) / lambda);
+  }
+}
+
 #define REGISTER_OP(NAME, DTYPE)                                    \
   template [[host_name(#NAME "_" #DTYPE)]] kernel void NAME<DTYPE>( \
       device DTYPE*, constant float2&, constant long2&, uint)
@@ -69,3 +90,190 @@ REGISTER_OP(geometric, long);
 REGISTER_OP(geometric, short);
 REGISTER_OP(geometric, char);
 REGISTER_OP(geometric, uchar);
+
+#define REGISTER_EXPONENTIAL(DTYPE)                         \
+  template [[host_name("exponential_" #DTYPE)]] kernel void \
+  exponential<DTYPE>(                                       \
+      device DTYPE*, constant float2&, constant long2&, constant uint&, uint)
+
+REGISTER_EXPONENTIAL(float);
+REGISTER_EXPONENTIAL(half);
+REGISTER_EXPONENTIAL(bfloat);
+
+// Marsaglia & Tsang (2000) acceptance-rejection method for Gamma distribution.
+// Adapted from aten/src/ATen/native/Distributions.h sample_gamma(),
+// which originates from NumPy's random module (Copyright 2005 Robert Kern).
+// Each thread uses a per-thread RNG offset stride to allow variable-length
+// rejection loops without colliding with other threads' random streams.
+constant constexpr int GAMMA_RANDOMS_STRIDE = 32;
+
+template <typename T>
+kernel void standard_gamma(
+    device T* output [[buffer(0)]],
+    device const T* alpha_in [[buffer(1)]],
+    constant long2& seed_base_offset [[buffer(2)]],
+    uint tid [[thread_position_in_grid]]) {
+  float alpha = static_cast<float>(alpha_in[tid]);
+  float scale = 1.0f;
+  long base =
+      seed_base_offset.y + static_cast<long>(tid) * GAMMA_RANDOMS_STRIDE;
+  long seed = seed_base_offset.x;
+  int rng_idx = 0;
+
+  // Boost alpha < 1 for higher acceptance probability
+  if (alpha < 1.0f) {
+    if (alpha == 0.0f) {
+      output[tid] = static_cast<T>(0.0f);
+      return;
+    }
+    float u = c10::metal::rand(seed, base + rng_idx++);
+    scale = ::metal::precise::pow(1.0f - u, 1.0f / alpha);
+    alpha += 1.0f;
+  }
+
+  // Marsaglia & Tsang acceptance-rejection
+  float d = alpha - 1.0f / 3.0f;
+  float c = 1.0f / ::metal::precise::sqrt(9.0f * d);
+  for (;;) {
+    float x, y;
+    do {
+      x = c10::metal::randn(seed, base + rng_idx++);
+      y = 1.0f + c * x;
+    } while (y <= 0.0f);
+    float v = y * y * y;
+    float u = 1.0f - c10::metal::rand(seed, base + rng_idx++);
+    float xx = x * x;
+    if (u < 1.0f - 0.0331f * xx * xx) {
+      float result = scale * d * v;
+      output[tid] = static_cast<T>(max(result, FLT_MIN));
+      return;
+    }
+    if (::metal::precise::log(u) <
+        0.5f * xx + d * (1.0f - v + ::metal::precise::log(v))) {
+      float result = scale * d * v;
+      output[tid] = static_cast<T>(max(result, FLT_MIN));
+      return;
+    }
+  }
+}
+
+#define REGISTER_GAMMA(DTYPE)                      \
+  template [[host_name("standard_gamma_" #DTYPE)]] \
+  kernel void standard_gamma<DTYPE>(               \
+      device DTYPE*, device const DTYPE*, constant long2&, uint)
+
+REGISTER_GAMMA(float);
+REGISTER_GAMMA(half);
+REGISTER_GAMMA(bfloat);
+
+// Reparameterized gradient for Gamma distribution.
+// Computes -(d/dalpha cdf(x;alpha)) / pdf(x;alpha).
+// Adapted from aten/src/ATen/native/Distributions.h standard_gamma_grad_one().
+
+constant constexpr float GAMMA_GRAD_COEF_UV[3][8] = {
+    {0.16009398f,
+     -0.094634809f,
+     0.025146376f,
+     -0.0030648343f,
+     1.0f,
+     0.32668115f,
+     0.10406089f,
+     0.0014179084f},
+    {0.53487893f,
+     0.1298071f,
+     0.065735949f,
+     -0.0015649758f,
+     0.16639465f,
+     0.020070113f,
+     -0.0035938915f,
+     -0.00058392623f},
+    {0.040121004f,
+     -0.0065914022f,
+     -0.0026286047f,
+     -0.0013441777f,
+     0.017050642f,
+     -0.0021309326f,
+     0.00085092367f,
+     -1.5247877e-07f},
+};
+
+template <typename T>
+kernel void standard_gamma_grad(
+    device T* output [[buffer(0)]],
+    device const T* self_data [[buffer(1)]],
+    device const T* output_data [[buffer(2)]],
+    uint tid [[thread_position_in_grid]]) {
+  float alpha = static_cast<float>(self_data[tid]);
+  float x = static_cast<float>(output_data[tid]);
+
+  // Region 1: Small x - Taylor series expansion
+  if (x < 0.8f) {
+    float numer = 1.0f;
+    float denom = alpha;
+    float series1 = numer / denom;
+    float series2 = numer / (denom * denom);
+    for (int i = 1; i <= 5; ++i) {
+      numer *= -x / static_cast<float>(i);
+      denom += 1.0f;
+      series1 += numer / denom;
+      series2 += numer / (denom * denom);
+    }
+    float pow_x_alpha = ::metal::precise::pow(x, alpha);
+    float gamma_pdf =
+        ::metal::precise::pow(x, alpha - 1.0f) * ::metal::precise::exp(-x);
+    float gamma_cdf = pow_x_alpha * series1;
+    float gamma_cdf_alpha =
+        (::metal::precise::log(x) - c10::metal::digamma(alpha)) * gamma_cdf -
+        pow_x_alpha * series2;
+    float result = -gamma_cdf_alpha / gamma_pdf;
+    output[tid] = static_cast<T>(isnan(result) ? 0.0f : result);
+    return;
+  }
+
+  // Region 2: Large alpha - Rice saddle point expansion
+  if (alpha > 8.0f) {
+    if (0.9f * alpha <= x && x <= 1.1f * alpha) {
+      float numer_1 = 1.0f + 24.0f * alpha * (1.0f + 12.0f * alpha);
+      float numer_2 = 1440.0f * (alpha * alpha) +
+          6.0f * x * (53.0f - 120.0f * x) - 65.0f * x * x / alpha +
+          alpha * (107.0f + 3600.0f * x);
+      float denom = 1244160.0f * (alpha * alpha) * (alpha * alpha);
+      output[tid] = static_cast<T>(numer_1 * numer_2 / denom);
+      return;
+    }
+    float denom = ::metal::precise::sqrt(8.0f * alpha);
+    float term2 = denom / (alpha - x);
+    float term3 = ::metal::precise::pow(
+        x - alpha - alpha * ::metal::precise::log(x / alpha), -1.5f);
+    float term23 = (x < alpha) ? term2 - term3 : term2 + term3;
+    float term1 = ::metal::precise::log(x / alpha) * term23 -
+        ::metal::precise::sqrt(2.0f / alpha) * (alpha + x) /
+            ((alpha - x) * (alpha - x));
+    float stirling =
+        1.0f + 1.0f / (12.0f * alpha) * (1.0f + 1.0f / (24.0f * alpha));
+    float numer = x * term1;
+    output[tid] = static_cast<T>(-stirling * numer / denom);
+    return;
+  }
+
+  // Region 3: Moderate alpha - bivariate rational approximation
+  float u = ::metal::precise::log(x / alpha);
+  float v = ::metal::precise::log(alpha);
+  float coef_v[8];
+  for (int i = 0; i < 8; ++i) {
+    coef_v[i] = GAMMA_GRAD_COEF_UV[0][i] +
+        u * (GAMMA_GRAD_COEF_UV[1][i] + u * GAMMA_GRAD_COEF_UV[2][i]);
+  }
+  float p = coef_v[0] + v * (coef_v[1] + v * (coef_v[2] + v * coef_v[3]));
+  float q = coef_v[4] + v * (coef_v[5] + v * (coef_v[6] + v * coef_v[7]));
+  output[tid] = static_cast<T>(::metal::precise::exp(p / q));
+}
+
+#define REGISTER_GAMMA_GRAD(DTYPE)                      \
+  template [[host_name("standard_gamma_grad_" #DTYPE)]] \
+  kernel void standard_gamma_grad<DTYPE>(               \
+      device DTYPE*, device const DTYPE*, device const DTYPE*, uint)
+
+REGISTER_GAMMA_GRAD(float);
+REGISTER_GAMMA_GRAD(half);
+REGISTER_GAMMA_GRAD(bfloat);
diff --git a/aten/src/ATen/native/mps/kernels/Eye.metal b/aten/src/ATen/native/mps/kernels/Eye.metal
new file mode 100644
index 0000000000000..ea389866241c4
--- /dev/null
+++ b/aten/src/ATen/native/mps/kernels/Eye.metal
@@ -0,0 +1,59 @@
+#include <metal_stdlib>
+using namespace metal;
+
+// For complex types (float2/half2), "1" is (1, 0), not (1, 1)
+template <typename T>
+constexpr T eye_one() {
+  return static_cast<T>(1);
+}
+template <>
+constexpr float2 eye_one<float2>() {
+  return float2(1.0, 0.0);
+}
+template <>
+constexpr half2 eye_one<half2>() {
+  return half2(half(1.0), half(0.0));
+}
+
+// Single-pass: writes both 0s and 1s in one dispatch (better for small tensors)
+template <typename T>
+kernel void eye(
+    device T* output [[buffer(0)]],
+    constant long& stride0 [[buffer(1)]],
+    constant long& stride1 [[buffer(2)]],
+    uint2 pos [[thread_position_in_grid]]) {
+  output[pos.y * stride0 + pos.x * stride1] =
+      (pos.x == pos.y) ? eye_one<T>() : static_cast<T>(0);
+}
+
+// Diagonal-only: writes 1s to pre-zeroed tensor (better for large tensors)
+template <typename T>
+kernel void eye_diag(
+    device T* output [[buffer(0)]],
+    constant long& diag_stride [[buffer(1)]],
+    uint index [[thread_position_in_grid]]) {
+  output[index * diag_stride] = eye_one<T>();
+}
+
+#define REGISTER_EYE_OP(DTYPE)                                            \
+  template [[host_name("eye_" #DTYPE)]] kernel void eye<DTYPE>(           \
+      device DTYPE * output [[buffer(0)]],                                \
+      constant long& stride0 [[buffer(1)]],                               \
+      constant long& stride1 [[buffer(2)]],                               \
+      uint2 pos [[thread_position_in_grid]]);                             \
+  template [[host_name("eye_diag_" #DTYPE)]] kernel void eye_diag<DTYPE>( \
+      device DTYPE * output [[buffer(0)]],                                \
+      constant long& diag_stride [[buffer(1)]],                           \
+      uint index [[thread_position_in_grid]]);
+
+REGISTER_EYE_OP(float);
+REGISTER_EYE_OP(half);
+REGISTER_EYE_OP(bfloat);
+REGISTER_EYE_OP(float2);
+REGISTER_EYE_OP(half2);
+REGISTER_EYE_OP(long);
+REGISTER_EYE_OP(int);
+REGISTER_EYE_OP(short);
+REGISTER_EYE_OP(char);
+REGISTER_EYE_OP(uchar);
+REGISTER_EYE_OP(bool);
diff --git a/aten/src/ATen/native/mps/kernels/Gamma.metal b/aten/src/ATen/native/mps/kernels/Gamma.metal
deleted file mode 100644
index 1c150a726edb1..0000000000000
--- a/aten/src/ATen/native/mps/kernels/Gamma.metal
+++ /dev/null
@@ -1,117 +0,0 @@
-/*
- * The gamma function approximations follow John D Cook's
- * c++ implementation:  https://www.johndcook.com/Gamma.cpp.
- * (BSD License)
- *
- *
- * The digamma kernel and helper function is derived from the pytorch cpu
- * of this function, which is itself derived from the implementation
- * of the digamma function in the Cephes Math Library.
- * See note [3-Clause BSD License for the Cephes Math Library].
- */
-
-#include <c10/metal/special_math.h>
-#include <metal_stdlib>
-using namespace metal;
-
-float calc_trigamma(float x) {
-  float sign = 1.0f;
-  float result = 0.0f;
-
-  if (x < 0.0f) {
-    sign = -1.0f;
-    auto sin_pi_x = sin(M_PI_F * x);
-    result -= (M_PI_F * M_PI_F) / (sin_pi_x * sin_pi_x);
-    x = 1.0f - x;
-  }
-
-  else if (x == 0.0) {
-    return INFINITY;
-  }
-
-  else if (x < 1.0) {
-    result += 1.0 / (x * x);
-    x += 1.0f;
-  }
-
-  for (int i = 0; i < 6; ++i) {
-    result += 1.0f / (x * x);
-    x += 1.0f;
-  }
-
-  const float ixx = 1.0f / (x * x);
-  result +=
-      (1.0f + 1.0f / (2.0f * x) +
-       ixx * ((1.0f / 6.0f) - ixx * ((1.0f / 30.0f) - ixx * (1.0f / 42.0f)))) /
-      x;
-  return sign * result;
-}
-
-template <typename T0, typename T1>
-kernel void lgamma(
-    constant T0* input [[buffer(0)]],
-    device T1* output [[buffer(1)]],
-    uint id [[thread_position_in_grid]]) {
-  output[id] =
-      static_cast<T1>(c10::metal::log_gamma(static_cast<float>(input[id])));
-}
-
-template <typename T0, typename T1>
-kernel void digamma(
-    constant T0* input [[buffer(0)]],
-    device T1* output [[buffer(1)]],
-    uint id [[thread_position_in_grid]]) {
-  float x = input[id];
-  output[id] = static_cast<T1>(c10::metal::digamma(x));
-}
-
-template <typename T0, typename T1>
-kernel void trigamma(
-    constant T0* input [[buffer(0)]],
-    device T1* output [[buffer(1)]],
-    uint id [[thread_position_in_grid]]) {
-  float x = input[id];
-  output[id] = static_cast<T1>(calc_trigamma(x));
-}
-
-template <typename T0, typename T1>
-kernel void polygamma(
-    constant T0* input [[buffer(0)]],
-    device T1* output [[buffer(1)]],
-    constant int64_t& order [[buffer(2)]],
-    uint id [[thread_position_in_grid]]) {
-  // already blocked if n <= 1
-  output[id] = static_cast<T1>(
-      c10::metal::polygamma(order, static_cast<float>(input[id])));
-}
-
-#define INSTANTIATE_GAMMA_KERNELS(DTYPE0, DTYPE1)                             \
-  template [[host_name("lgamma_" #DTYPE0 "_" #DTYPE1)]] kernel void lgamma(   \
-      constant DTYPE0* input [[buffer(0)]],                                   \
-      device DTYPE1* output [[buffer(1)]],                                    \
-      uint id [[thread_position_in_grid]]);                                   \
-  template [[host_name("digamma_" #DTYPE0 "_" #DTYPE1)]] kernel void digamma( \
-      constant DTYPE0* input [[buffer(0)]],                                   \
-      device DTYPE1* output [[buffer(1)]],                                    \
-      uint id [[thread_position_in_grid]]);                                   \
-  template [[host_name("trigamma_" #DTYPE0 "_" #DTYPE1)]] kernel void         \
-  trigamma(                                                                   \
-      constant DTYPE0* input [[buffer(0)]],                                   \
-      device DTYPE1* output [[buffer(1)]],                                    \
-      uint id [[thread_position_in_grid]]);                                   \
-  template [[host_name("polygamma_" #DTYPE0 "_" #DTYPE1)]] kernel void        \
-  polygamma(                                                                  \
-      constant DTYPE0* input [[buffer(0)]],                                   \
-      device DTYPE1* output [[buffer(1)]],                                    \
-      constant int64_t& order [[buffer(2)]],                                  \
-      uint id [[thread_position_in_grid]]);
-
-INSTANTIATE_GAMMA_KERNELS(bfloat, bfloat);
-INSTANTIATE_GAMMA_KERNELS(half, half);
-INSTANTIATE_GAMMA_KERNELS(float, float);
-INSTANTIATE_GAMMA_KERNELS(bool, float);
-INSTANTIATE_GAMMA_KERNELS(uchar, float);
-INSTANTIATE_GAMMA_KERNELS(char, float);
-INSTANTIATE_GAMMA_KERNELS(short, float);
-INSTANTIATE_GAMMA_KERNELS(int, float);
-INSTANTIATE_GAMMA_KERNELS(long, float);
diff --git a/aten/src/ATen/native/mps/kernels/GridSampler.h b/aten/src/ATen/native/mps/kernels/GridSampler.h
index c2b3cad3cd47d..c9d4112508613 100644
--- a/aten/src/ATen/native/mps/kernels/GridSampler.h
+++ b/aten/src/ATen/native/mps/kernels/GridSampler.h
@@ -1,15 +1,6 @@
 #pragma once
 #include <c10/metal/common.h>
 
-#ifdef __METAL__
-enum class GridSamplerInterpolation { Bilinear, Nearest, Bicubic };
-enum class GridSamplerPadding { Zeros, Border, Reflection };
-#else
-#include <ATen/native/GridSamplerUtils.h>
-using at::native::GridSamplerInterpolation;
-using at::native::GridSamplerPadding;
-#endif
-
 template <unsigned N = 5, typename idx_type_t = int32_t>
 struct GridSamplerParams {
   int32_t sampler_dims;
@@ -19,7 +10,5 @@ struct GridSamplerParams {
   ::c10::metal::array<idx_type_t, N> input_strides;
   ::c10::metal::array<idx_type_t, N> grid_sizes;
   ::c10::metal::array<idx_type_t, N> grid_strides;
-  GridSamplerInterpolation interpolation_mode;
-  GridSamplerPadding padding_mode;
   bool align_corners;
 };
diff --git a/aten/src/ATen/native/mps/kernels/GridSampler.metal b/aten/src/ATen/native/mps/kernels/GridSampler.metal
index ceb0cb3d55ab8..8314f63bdbce9 100644
--- a/aten/src/ATen/native/mps/kernels/GridSampler.metal
+++ b/aten/src/ATen/native/mps/kernels/GridSampler.metal
@@ -1,4 +1,5 @@
 #include <ATen/native/mps/kernels/GridSampler.h>
+#include <ATen/native/mps/kernels/SamplingHelpers.h>
 #include <c10/metal/utils.h>
 #include <metal_array>
 #include <metal_stdlib>
@@ -68,35 +69,11 @@ static int32_t mod(int32_t a, int32_t b) {
 // Sentinel index value to indicate zero padding
 constant int32_t IDX_ZERO = -1;
 
-// Apply padding to an index into the input
-static int32_t pad_input_index(
-    int32_t idx,
-    int32_t input_size,
-    GridSamplerPadding padding_mode,
-    bool align_corners) {
-  int32_t idx_padded = idx;
-
-  if (padding_mode == GridSamplerPadding::Zeros) {
-    idx_padded = (idx < 0) ? IDX_ZERO : idx_padded;
-    idx_padded = (idx >= input_size) ? IDX_ZERO : idx_padded;
-
-  } else if (padding_mode == GridSamplerPadding::Border) {
-    idx_padded = (idx < 0) ? 0 : idx_padded;
-    idx_padded = (idx >= input_size) ? input_size - 1 : idx_padded;
-
-  } else if (padding_mode == GridSamplerPadding::Reflection) {
-    auto scale_length = align_corners ? (input_size - 1) : input_size;
-    auto idx_mod = mod(idx, scale_length);
-    auto idx_mod_reverse = (input_size - 1) - idx_mod;
-    bool is_reverse = (abs(idx - idx_mod) / scale_length) % 2 == 1;
-    idx_padded = is_reverse ? idx_mod_reverse : idx_mod;
-  }
-  return idx_padded;
-}
-
 // Unnormalize grid coordinate from [-1, 1] to pixel space
-template <typename T>
-static T grid_sampler_unnormalize(T coord, int32_t size, bool align_corners) {
+static float grid_sampler_unnormalize(
+    float coord,
+    int32_t size,
+    bool align_corners) {
   if (align_corners) {
     return ((coord + 1) / 2) * (size - 1);
   } else {
@@ -105,9 +82,8 @@ static T grid_sampler_unnormalize(T coord, int32_t size, bool align_corners) {
 }
 
 // Clip coordinates for border padding
-template <typename T>
-static T clip_coordinates(T in, int32_t clip_limit) {
-  return ::metal::clamp(in, static_cast<T>(0), static_cast<T>(clip_limit - 1));
+static float clip_coordinates(float in, int32_t clip_limit) {
+  return ::metal::clamp(in, 0.0, clip_limit - 1.0);
 }
 
 // Reflect coordinates for reflection padding
@@ -124,68 +100,65 @@ static T reflect_coordinates(T in, int32_t twice_low, int32_t twice_high) {
   return (flips % 2 == 0) ? (extra + min_val) : (span - extra + min_val);
 }
 
-// Compute source index with padding mode
-template <typename T>
-static T compute_source_index(
-    T coord,
-    int32_t size,
-    GridSamplerPadding padding_mode,
-    bool align_corners) {
-  coord = grid_sampler_unnormalize(coord, size, align_corners);
-  if (padding_mode == GridSamplerPadding::Border) {
-    coord = clip_coordinates(coord, size);
-  } else if (padding_mode == GridSamplerPadding::Reflection) {
+// Padding functors: each encapsulates the padding logic for integer indices
+// (pad) and float source coordinates (compute_source).
+struct PadZeros {
+  static constant constexpr bool checks_bounds = true;
+
+  static int32_t pad(int32_t idx, int32_t input_size, bool) {
+    return (idx < 0 || idx >= input_size) ? IDX_ZERO : idx;
+  }
+
+  static float compute_source(float coord, int32_t size, bool align_corners) {
+    return grid_sampler_unnormalize(coord, size, align_corners);
+  }
+};
+
+struct PadBorder {
+  static constant constexpr bool checks_bounds = false;
+
+  static int32_t pad(int32_t idx, int32_t input_size, bool) {
+    return clamp(idx, 0, input_size - 1);
+  }
+
+  static float compute_source(float coord, int32_t size, bool align_corners) {
+    coord = grid_sampler_unnormalize(coord, size, align_corners);
+    return clip_coordinates(coord, size);
+  }
+};
+
+struct PadReflection {
+  static constant constexpr bool checks_bounds = false;
+
+  static int32_t pad(int32_t idx, int32_t input_size, bool align_corners) {
+    auto scale_length = align_corners ? (input_size - 1) : input_size;
+    auto idx_mod = mod(idx, scale_length);
+    auto idx_mod_reverse = (input_size - 1) - idx_mod;
+    bool is_reverse = (abs(idx - idx_mod) / scale_length) % 2 == 1;
+    return is_reverse ? idx_mod_reverse : idx_mod;
+  }
+
+  static float compute_source(float coord, int32_t size, bool align_corners) {
+    coord = grid_sampler_unnormalize(coord, size, align_corners);
     if (align_corners) {
       coord = reflect_coordinates(coord, 0, 2 * (size - 1));
     } else {
       coord = reflect_coordinates(coord, -1, 2 * size - 1);
     }
-    coord = clip_coordinates(coord, size);
+    return clip_coordinates(coord, size);
   }
-  return coord;
-}
-
-// Cubic convolution helper 1: for |x| < 1
-template <typename T>
-static T cubic_convolution1(T x, T A) {
-  return ((A + 2) * x - (A + 3)) * x * x + 1;
-}
-
-// Cubic convolution helper 2: for 1 <= |x| < 2
-template <typename T>
-static T cubic_convolution2(T x, T A) {
-  return ((A * x - 5 * A) * x + 8 * A) * x - 4 * A;
-}
-
-// Get cubic upsampling coefficients (Catmull-Rom spline with A=-0.75)
-template <typename T>
-static void get_cubic_coefficients(T coeffs[4], T t) {
-  T A = static_cast<T>(-0.75);
-  coeffs[0] = cubic_convolution2(t + 1, A);
-  coeffs[1] = cubic_convolution1(t, A);
-  coeffs[2] = cubic_convolution1(1 - t, A);
-  coeffs[3] = cubic_convolution2(2 - t, A);
-}
-
-// 1D cubic interpolation
-template <typename T>
-static T cubic_interp1d(T x0, T x1, T x2, T x3, T t) {
-  T coeffs[4];
-  get_cubic_coefficients(coeffs, t);
-  return x0 * coeffs[0] + x1 * coeffs[1] + x2 * coeffs[2] + x3 * coeffs[3];
-}
+};
 
 // 2D Bilinear interpolation
-template <typename T>
+template <typename Pad, typename T>
 static T interpolate_bilinear_2d(
     constant T* input,
-    opmath_t<T> ix,
-    opmath_t<T> iy,
+    float ix,
+    float iy,
     int32_t inp_H,
     int32_t inp_W,
     int32_t inp_sH,
     int32_t inp_sW,
-    GridSamplerPadding padding_mode,
     bool align_corners) {
   ix = grid_sampler_unnormalize(ix, inp_W, align_corners);
   iy = grid_sampler_unnormalize(iy, inp_H, align_corners);
@@ -199,51 +172,39 @@ static T interpolate_bilinear_2d(
   int32_t ix_se = ix_nw + 1;
   int32_t iy_se = iy_nw + 1;
 
-  opmath_t<T> nw = (static_cast<opmath_t<T>>(ix_se) - ix) *
-      (static_cast<opmath_t<T>>(iy_se) - iy);
-  opmath_t<T> ne = (ix - static_cast<opmath_t<T>>(ix_sw)) *
-      (static_cast<opmath_t<T>>(iy_sw) - iy);
-  opmath_t<T> sw = (static_cast<opmath_t<T>>(ix_ne) - ix) *
-      (iy - static_cast<opmath_t<T>>(iy_ne));
-  opmath_t<T> se = (ix - static_cast<opmath_t<T>>(ix_nw)) *
-      (iy - static_cast<opmath_t<T>>(iy_nw));
-
-  int32_t iy_nw_p = pad_input_index(iy_nw, inp_H, padding_mode, align_corners);
-  int32_t ix_nw_p = pad_input_index(ix_nw, inp_W, padding_mode, align_corners);
-  int32_t iy_ne_p = pad_input_index(iy_ne, inp_H, padding_mode, align_corners);
-  int32_t ix_ne_p = pad_input_index(ix_ne, inp_W, padding_mode, align_corners);
-  int32_t iy_sw_p = pad_input_index(iy_sw, inp_H, padding_mode, align_corners);
-  int32_t ix_sw_p = pad_input_index(ix_sw, inp_W, padding_mode, align_corners);
-  int32_t iy_se_p = pad_input_index(iy_se, inp_H, padding_mode, align_corners);
-  int32_t ix_se_p = pad_input_index(ix_se, inp_W, padding_mode, align_corners);
-
-  opmath_t<T> out_acc = static_cast<opmath_t<T>>(0);
+  const auto nw = (ix_se - ix) * (iy_se - iy);
+  const auto ne = (ix - ix_sw) * (iy_sw - iy);
+  const auto sw = (ix_ne - ix) * (iy - iy_ne);
+  const auto se = (ix - ix_nw) * (iy - iy_nw);
+
+  int32_t iy_nw_p = Pad::pad(iy_nw, inp_H, align_corners);
+  int32_t ix_nw_p = Pad::pad(ix_nw, inp_W, align_corners);
+  int32_t iy_ne_p = Pad::pad(iy_ne, inp_H, align_corners);
+  int32_t ix_ne_p = Pad::pad(ix_ne, inp_W, align_corners);
+  int32_t iy_sw_p = Pad::pad(iy_sw, inp_H, align_corners);
+  int32_t ix_sw_p = Pad::pad(ix_sw, inp_W, align_corners);
+  int32_t iy_se_p = Pad::pad(iy_se, inp_H, align_corners);
+  int32_t ix_se_p = Pad::pad(ix_se, inp_W, align_corners);
+
+  opmath_t<T> out_acc = 0;
   if (iy_nw_p != IDX_ZERO && ix_nw_p != IDX_ZERO) {
-    out_acc +=
-        static_cast<opmath_t<T>>(input[iy_nw_p * inp_sH + ix_nw_p * inp_sW]) *
-        nw;
+    out_acc += input[iy_nw_p * inp_sH + ix_nw_p * inp_sW] * nw;
   }
   if (iy_ne_p != IDX_ZERO && ix_ne_p != IDX_ZERO) {
-    out_acc +=
-        static_cast<opmath_t<T>>(input[iy_ne_p * inp_sH + ix_ne_p * inp_sW]) *
-        ne;
+    out_acc += input[iy_ne_p * inp_sH + ix_ne_p * inp_sW] * ne;
   }
   if (iy_sw_p != IDX_ZERO && ix_sw_p != IDX_ZERO) {
-    out_acc +=
-        static_cast<opmath_t<T>>(input[iy_sw_p * inp_sH + ix_sw_p * inp_sW]) *
-        sw;
+    out_acc += input[iy_sw_p * inp_sH + ix_sw_p * inp_sW] * sw;
   }
   if (iy_se_p != IDX_ZERO && ix_se_p != IDX_ZERO) {
-    out_acc +=
-        static_cast<opmath_t<T>>(input[iy_se_p * inp_sH + ix_se_p * inp_sW]) *
-        se;
+    out_acc += input[iy_se_p * inp_sH + ix_se_p * inp_sW] * se;
   }
 
   return static_cast<T>(out_acc);
 }
 
 // 2D Nearest neighbor interpolation
-template <typename T>
+template <typename Pad, typename T>
 static T interpolate_nearest_2d(
     constant T* input,
     opmath_t<T> ix,
@@ -252,18 +213,14 @@ static T interpolate_nearest_2d(
     int32_t inp_W,
     int32_t inp_sH,
     int32_t inp_sW,
-    GridSamplerPadding padding_mode,
     bool align_corners) {
-  // For Border and Reflection modes, apply padding to float coordinates
-  // before rounding (matches CPU behavior)
-  ix = compute_source_index(ix, inp_W, padding_mode, align_corners);
-  iy = compute_source_index(iy, inp_H, padding_mode, align_corners);
+  ix = Pad::compute_source(ix, inp_W, align_corners);
+  iy = Pad::compute_source(iy, inp_H, align_corners);
 
   int32_t ix_nearest = static_cast<int32_t>(rint(ix));
   int32_t iy_nearest = static_cast<int32_t>(rint(iy));
 
-  // For Zeros padding, check bounds after rounding
-  if (padding_mode == GridSamplerPadding::Zeros) {
+  if (Pad::checks_bounds) {
     if (ix_nearest < 0 || ix_nearest >= inp_W || iy_nearest < 0 ||
         iy_nearest >= inp_H) {
       return static_cast<T>(0);
@@ -274,7 +231,7 @@ static T interpolate_nearest_2d(
 }
 
 // Helper to get bounded value for bicubic interpolation
-template <typename T>
+template <typename Pad, typename T>
 static opmath_t<T> get_bicubic_value(
     constant T* input,
     int32_t y,
@@ -283,35 +240,33 @@ static opmath_t<T> get_bicubic_value(
     int32_t inp_W,
     int32_t inp_sH,
     int32_t inp_sW,
-    GridSamplerPadding padding_mode,
     bool align_corners) {
-  int32_t y_p = pad_input_index(y, inp_H, padding_mode, align_corners);
-  int32_t x_p = pad_input_index(x, inp_W, padding_mode, align_corners);
+  int32_t y_p = Pad::pad(y, inp_H, align_corners);
+  int32_t x_p = Pad::pad(x, inp_W, align_corners);
   if (y_p == IDX_ZERO || x_p == IDX_ZERO) {
-    return static_cast<opmath_t<T>>(0);
+    return 0;
   }
-  return static_cast<opmath_t<T>>(input[y_p * inp_sH + x_p * inp_sW]);
+  return input[y_p * inp_sH + x_p * inp_sW];
 }
 
 // 2D Bicubic interpolation
-template <typename T>
+template <typename Pad, typename T>
 static T interpolate_bicubic_2d(
     constant T* input,
-    opmath_t<T> ix,
-    opmath_t<T> iy,
+    float ix,
+    float iy,
     int32_t inp_H,
     int32_t inp_W,
     int32_t inp_sH,
     int32_t inp_sW,
-    GridSamplerPadding padding_mode,
     bool align_corners) {
   ix = grid_sampler_unnormalize(ix, inp_W, align_corners);
   iy = grid_sampler_unnormalize(iy, inp_H, align_corners);
 
-  opmath_t<T> ix_nw = floor(ix);
-  opmath_t<T> iy_nw = floor(iy);
-  opmath_t<T> tx = ix - ix_nw;
-  opmath_t<T> ty = iy - iy_nw;
+  auto ix_nw = floor(ix);
+  auto iy_nw = floor(iy);
+  auto tx = ix - ix_nw;
+  auto ty = iy - iy_nw;
 
   opmath_t<T> coefficients[4];
   int32_t ix_nw_i = static_cast<int32_t>(ix_nw);
@@ -319,7 +274,7 @@ static T interpolate_bicubic_2d(
 
   for (int32_t i = 0; i < 4; ++i) {
     coefficients[i] = cubic_interp1d(
-        get_bicubic_value<T>(
+        get_bicubic_value<Pad, T>(
             input,
             iy_nw_i - 1 + i,
             ix_nw_i - 1,
@@ -327,9 +282,8 @@ static T interpolate_bicubic_2d(
             inp_W,
             inp_sH,
             inp_sW,
-            padding_mode,
             align_corners),
-        get_bicubic_value<T>(
+        get_bicubic_value<Pad, T>(
             input,
             iy_nw_i - 1 + i,
             ix_nw_i + 0,
@@ -337,9 +291,8 @@ static T interpolate_bicubic_2d(
             inp_W,
             inp_sH,
             inp_sW,
-            padding_mode,
             align_corners),
-        get_bicubic_value<T>(
+        get_bicubic_value<Pad, T>(
             input,
             iy_nw_i - 1 + i,
             ix_nw_i + 1,
@@ -347,9 +300,8 @@ static T interpolate_bicubic_2d(
             inp_W,
             inp_sH,
             inp_sW,
-            padding_mode,
             align_corners),
-        get_bicubic_value<T>(
+        get_bicubic_value<Pad, T>(
             input,
             iy_nw_i - 1 + i,
             ix_nw_i + 2,
@@ -357,7 +309,6 @@ static T interpolate_bicubic_2d(
             inp_W,
             inp_sH,
             inp_sW,
-            padding_mode,
             align_corners),
         tx);
   }
@@ -366,15 +317,68 @@ static T interpolate_bicubic_2d(
       coefficients[0], coefficients[1], coefficients[2], coefficients[3], ty));
 }
 
+// Interpolation functors for the 2D kernel.
+// Each wraps the respective interpolation function with the padding already
+// baked in via the Pad template parameter.
+template <typename Pad>
+struct Bilinear2D {
+  template <typename T>
+  static T interpolate(
+      constant T* input,
+      opmath_t<T> ix,
+      opmath_t<T> iy,
+      int32_t inp_H,
+      int32_t inp_W,
+      int32_t inp_sH,
+      int32_t inp_sW,
+      bool align_corners) {
+    return interpolate_bilinear_2d<Pad, T>(
+        input, ix, iy, inp_H, inp_W, inp_sH, inp_sW, align_corners);
+  }
+};
+
+template <typename Pad>
+struct Nearest2D {
+  template <typename T>
+  static T interpolate(
+      constant T* input,
+      opmath_t<T> ix,
+      opmath_t<T> iy,
+      int32_t inp_H,
+      int32_t inp_W,
+      int32_t inp_sH,
+      int32_t inp_sW,
+      bool align_corners) {
+    return interpolate_nearest_2d<Pad, T>(
+        input, ix, iy, inp_H, inp_W, inp_sH, inp_sW, align_corners);
+  }
+};
+
+template <typename Pad>
+struct Bicubic2D {
+  template <typename T>
+  static T interpolate(
+      constant T* input,
+      float ix,
+      float iy,
+      int32_t inp_H,
+      int32_t inp_W,
+      int32_t inp_sH,
+      int32_t inp_sW,
+      bool align_corners) {
+    return interpolate_bicubic_2d<Pad, T>(
+        input, ix, iy, inp_H, inp_W, inp_sH, inp_sW, align_corners);
+  }
+};
+
 // 2D grid sampler kernel
-template <typename T>
+template <typename Interp, typename T>
 kernel void grid_sampler_2d(
     device T* output [[buffer(0)]],
     constant T* input [[buffer(1)]],
     constant T* grid [[buffer(2)]],
-    constant GridSamplerParams<5>& params [[buffer(3)]],
+    constant GridSamplerParams<4>& params [[buffer(3)]],
     uint tid [[thread_position_in_grid]]) {
-  auto N = params.output_sizes[0];
   auto C = params.output_sizes[1];
   auto out_H = params.output_sizes[2];
   auto out_W = params.output_sizes[3];
@@ -394,15 +398,8 @@ kernel void grid_sampler_2d(
   auto grid_sW = params.grid_strides[2];
   auto grid_sCoor = params.grid_strides[3];
 
-  auto interpolation_mode = params.interpolation_mode;
-  auto padding_mode = params.padding_mode;
   auto align_corners = params.align_corners;
 
-  auto num_elements = N * out_H * out_W;
-  if (tid >= static_cast<uint>(num_elements)) {
-    return;
-  }
-
   int32_t w = tid % out_W;
   int32_t h = (tid / out_W) % out_H;
   int32_t n = tid / (out_H * out_W);
@@ -415,41 +412,8 @@ kernel void grid_sampler_2d(
   auto out_ptr_NCHW = output + n * out_sN + h * out_sH + w * out_sW;
 
   for (int32_t c = 0; c < C; ++c) {
-    T result;
-    if (interpolation_mode == GridSamplerInterpolation::Bilinear) {
-      result = interpolate_bilinear_2d<T>(
-          inp_ptr_N,
-          ix,
-          iy,
-          inp_H,
-          inp_W,
-          inp_sH,
-          inp_sW,
-          padding_mode,
-          align_corners);
-    } else if (interpolation_mode == GridSamplerInterpolation::Nearest) {
-      result = interpolate_nearest_2d<T>(
-          inp_ptr_N,
-          ix,
-          iy,
-          inp_H,
-          inp_W,
-          inp_sH,
-          inp_sW,
-          padding_mode,
-          align_corners);
-    } else {
-      result = interpolate_bicubic_2d<T>(
-          inp_ptr_N,
-          ix,
-          iy,
-          inp_H,
-          inp_W,
-          inp_sH,
-          inp_sW,
-          padding_mode,
-          align_corners);
-    }
+    auto result = Interp::template interpolate<T>(
+        inp_ptr_N, ix, iy, inp_H, inp_W, inp_sH, inp_sW, align_corners);
     out_ptr_NCHW[c * out_sC] = result;
     inp_ptr_N += inp_sC;
   }
@@ -530,15 +494,9 @@ T interpolate_linear_3d(
       scale0_right * scale1_right * scale2_right * h);
 }
 
-// Calculates a single output element.
-// `input` shape:
-//    2 sampler dims: (Hin, Win)
-//    3 sampler dims: (Din, Hin, Win)
-// `coords` values:
-//    2 sampler dims: (Wcoord, Hcoord)
-//    3 sampler dims: (Wcoord, Hcoord, Dcoord)
-template <typename T>
-void grid_sampler_single_element(
+// 3D bilinear sampling for a single output element.
+template <typename Pad, typename T>
+void grid_sampler_3d_single_element(
     device T* output,
     constant T* input,
     constant T* coords,
@@ -546,89 +504,37 @@ void grid_sampler_single_element(
     constant int32_t* input_sizes,
     constant int32_t* input_strides,
     int32_t coord_stride,
-    GridSamplerInterpolation interpolation_mode,
-    GridSamplerPadding padding_mode,
     bool align_corners) {
   int32_t left_indices[3];
   int32_t right_indices[3];
   opmath_t<T> scales[3];
 
-  // For each dimension, find the pair of indices in the corresponding dimension
-  // of `input` which surround the grid coordinate in that dimension. We'll do
-  // this by mapping different coordinate spaces onto each other. There are
-  // basically three different coordinate spaces to keep in mind:
-  //
-  //  * aligned grid space
-  //    - `-1` refers to the leftmost input value.
-  //    - `1` refers to the rightmost input value.
-  //
-  //  * unaligned grid space
-  //    - `-1` refers to the midpoint between the leftmost input value and
-  //      a padding value to the left of that.
-  //    - `1` refers to the midpoint between the rightmost input value and
-  //      a padding value to the right of that.
-  //
-  //  * input index space
-  //    - `n` refers to the n-th value of the input.
-  //    - `0` refers to the leftmost input value.
-  //    - `N-1` refers to the rightmost input value.
-  //
-  // If `align_corners == False`, then the coordinates are is in unaligned grid
-  // space, and we will map it onto aligned grid space. If `align_corners ==
-  // True`, then coordinates are already in aligned grid space.
-  //
-  // Then we will map unaligned grid space onto input index space, making it
-  // relatively simple to find the two input indices that surround the
-  // coordinate.
   for (auto coord_dim = 0; coord_dim < dims; coord_dim++) {
     auto input_dim = dims - coord_dim - 1;
     auto input_size = input_sizes[input_dim];
     auto coord = static_cast<opmath_t<T>>(coords[coord_dim * coord_stride]);
 
     if (!align_corners) {
-      // Map unaligned grid space to aligned grid space
       auto corner_alignment_factor = static_cast<opmath_t<T>>(input_size) /
           static_cast<opmath_t<T>>(input_size - 1);
       coord = coord * corner_alignment_factor;
     }
 
-    // Map aligned grid space to input index space
     coord = (coord + 1) * (static_cast<opmath_t<T>>(input_size - 1) / 2);
 
-    // Get the input indices surrounding the coordinate, apply padding to them,
-    // and obtain the scaling factor between the two for interpolation.
     auto left_idx = static_cast<int32_t>(floor(coord));
     auto right_idx = static_cast<int32_t>(ceil(coord));
-    left_indices[input_dim] =
-        pad_input_index(left_idx, input_size, padding_mode, align_corners);
-    right_indices[input_dim] =
-        pad_input_index(right_idx, input_size, padding_mode, align_corners);
-
-    auto scale = coord - left_idx;
-
-    if (interpolation_mode == GridSamplerInterpolation::Nearest) {
-      // Round to nearest even (banker's rounding) - matches CPU behavior
-      // For halfway cases (scale == 0.5), pick the even-indexed neighbor
-      if (scale == 0.5) {
-        // if left_idx is even -> round to 0, if odd round to 1
-        scale = (left_idx % 2 == 0) ? 0 : 1;
-      } else {
-        scale = (scale < 0.5) ? 0 : 1;
-      }
-    }
-    scales[input_dim] = scale;
+    left_indices[input_dim] = Pad::pad(left_idx, input_size, align_corners);
+    right_indices[input_dim] = Pad::pad(right_idx, input_size, align_corners);
+    scales[input_dim] = coord - left_idx;
   }
 
-  // Now that we have the bounding indices and scale factor for each dimension
-  // of the input, we can interpolate.
-  if (dims == 3) {
-    *output = interpolate_linear_3d(
-        input, input_strides, left_indices, right_indices, scales);
-  }
+  *output = interpolate_linear_3d(
+      input, input_strides, left_indices, right_indices, scales);
 }
 
-template <typename T>
-kernel void grid_sampler(
+template <typename Pad, typename T>
+kernel void grid_sampler_3d(
     device T* output [[buffer(0)]],
     constant T* input [[buffer(1)]],
     constant T* grid [[buffer(2)]],
@@ -657,11 +563,7 @@ kernel void grid_sampler(
   input_strides += 2;
   auto coord_stride = grid_strides[sampler_dims + 1];
 
-  auto interpolation_mode = params.interpolation_mode;
-  auto padding_mode = params.padding_mode;
-  auto align_corners = params.align_corners;
-
-  grid_sampler_single_element(
+  grid_sampler_3d_single_element<Pad>(
       output,
       input,
       coords,
@@ -669,27 +571,40 @@ kernel void grid_sampler(
       input_sizes,
       input_strides,
       coord_stride,
-      interpolation_mode,
-      padding_mode,
-      align_corners);
+      params.align_corners);
 }
 
-#define REGISTER_GRID_SAMPLER_OPS(DTYPE)                    \
-  template [[host_name("grid_sampler_" #DTYPE)]]            \
-  kernel void grid_sampler<DTYPE>(                          \
-      device DTYPE * output [[buffer(0)]],                  \
-      constant DTYPE * input [[buffer(1)]],                 \
-      constant DTYPE * grid [[buffer(2)]],                  \
-      constant GridSamplerParams<5> & params [[buffer(3)]], \
-      uint tid [[thread_position_in_grid]]);                \
-  template [[host_name("grid_sampler_2d_" #DTYPE)]]         \
-  kernel void grid_sampler_2d<DTYPE>(                       \
-      device DTYPE * output [[buffer(0)]],                  \
-      constant DTYPE * input [[buffer(1)]],                 \
-      constant DTYPE * grid [[buffer(2)]],                  \
-      constant GridSamplerParams<5> & params [[buffer(3)]], \
+#define REGISTER_GRID_SAMPLER_2D(DTYPE, INTERP, INAME, PAD, PNAME)      \
+  template [[host_name("grid_sampler_2d_" INAME "_" PNAME "_" #DTYPE)]] \
+  kernel void grid_sampler_2d<INTERP<PAD>, DTYPE>(                      \
+      device DTYPE * output [[buffer(0)]],                              \
+      constant DTYPE * input [[buffer(1)]],                             \
+      constant DTYPE * grid [[buffer(2)]],                              \
+      constant GridSamplerParams<4> & params [[buffer(3)]],             \
       uint tid [[thread_position_in_grid]]);
 
+#define REGISTER_GRID_SAMPLER_2D_INTERP(DTYPE, INTERP, INAME)         \
+  REGISTER_GRID_SAMPLER_2D(DTYPE, INTERP, INAME, PadZeros, "zeros")   \
+  REGISTER_GRID_SAMPLER_2D(DTYPE, INTERP, INAME, PadBorder, "border") \
+  REGISTER_GRID_SAMPLER_2D(DTYPE, INTERP, INAME, PadReflection, "reflection")
+
+#define REGISTER_GRID_SAMPLER_3D(DTYPE, PAD, PNAME)           \
+  template [[host_name("grid_sampler_3d_" PNAME "_" #DTYPE)]] \
+  kernel void grid_sampler_3d<PAD, DTYPE>(                    \
+      device DTYPE * output [[buffer(0)]],                    \
+      constant DTYPE * input [[buffer(1)]],                   \
+      constant DTYPE * grid [[buffer(2)]],                    \
+      constant GridSamplerParams<5> & params [[buffer(3)]],   \
+      uint tid [[thread_position_in_grid]]);
+
+#define REGISTER_GRID_SAMPLER_OPS(DTYPE)                         \
+  REGISTER_GRID_SAMPLER_2D_INTERP(DTYPE, Bilinear2D, "bilinear") \
+  REGISTER_GRID_SAMPLER_2D_INTERP(DTYPE, Nearest2D, "nearest")   \
+  REGISTER_GRID_SAMPLER_2D_INTERP(DTYPE, Bicubic2D, "bicubic")   \
+  REGISTER_GRID_SAMPLER_3D(DTYPE, PadZeros, "zeros")             \
+  REGISTER_GRID_SAMPLER_3D(DTYPE, PadBorder, "border")           \
+  REGISTER_GRID_SAMPLER_3D(DTYPE, PadReflection, "reflection")
+
 REGISTER_GRID_SAMPLER_OPS(float);
 REGISTER_GRID_SAMPLER_OPS(half);
 REGISTER_GRID_SAMPLER_OPS(bfloat);
diff --git a/aten/src/ATen/native/mps/kernels/HistogramKernel.metal b/aten/src/ATen/native/mps/kernels/HistogramKernel.metal
index 39922e03918d6..689eee6e47f11 100644
--- a/aten/src/ATen/native/mps/kernels/HistogramKernel.metal
+++ b/aten/src/ATen/native/mps/kernels/HistogramKernel.metal
@@ -40,9 +40,9 @@ kernel void histogramdd(
     constant T* rightmost_edge [[buffer(8)]],
     constant int64_t* local_out_strides [[buffer(9)]],
     constant uint8_t& algorithm [[buffer(10)]],
-    constant uint8_t& has_weight [[buffer(11)]],
+    constant int64_t& weight_stride [[buffer(11)]],
     uint tid [[thread_position_in_grid]]) {
-  constexpr T eps = 4e-6;
+  constexpr auto eps = T(4e-6);
   bool skip_element = false;
   int64_t hist_index = 0;
   int64_t bin_seq_offset = 0;
@@ -90,7 +90,7 @@ kernel void histogramdd(
   if (!skip_element) {
     // In the unweighted case, the default weight is 1
     local_out[local_out_strides[0] * tid + hist_index] +=
-        has_weight ? weight[tid] : 1;
+        (weight_stride >= 0) ? weight[tid * weight_stride] : 1;
   }
 }
 
@@ -108,11 +108,12 @@ kernel void histogramdd(
       constant DTYPE* rightmost_edge [[buffer(8)]],             \
       constant int64_t* local_out_strides [[buffer(9)]],        \
       constant uint8_t& bin_selection_algorithm [[buffer(10)]], \
-      constant uint8_t& has_weight [[buffer(11)]],              \
+      constant int64_t& weight_stride [[buffer(11)]],           \
       uint tid [[thread_position_in_grid]]);
 
 REGISTER_HISTOGRAMDD_OP(float);
 REGISTER_HISTOGRAMDD_OP(half);
+REGISTER_HISTOGRAMDD_OP(bfloat);
 
 kernel void kernel_index_offset(
     constant uint* strides [[buffer(0)]],
diff --git a/aten/src/ATen/native/mps/kernels/Indexing.h b/aten/src/ATen/native/mps/kernels/Indexing.h
new file mode 100644
index 0000000000000..7c50b1f0f191e
--- /dev/null
+++ b/aten/src/ATen/native/mps/kernels/Indexing.h
@@ -0,0 +1,13 @@
+#pragma once
+#include <c10/metal/common.h>
+
+template <unsigned N = c10::metal::max_ndim>
+struct IndexReduceParams {
+  int32_t ndim;
+  int32_t reduce_dim;
+  ::c10::metal::array<uint32_t, N> self_strides;
+  ::c10::metal::array<uint32_t, N> self_sizes;
+  ::c10::metal::array<uint32_t, N> source_strides;
+  ::c10::metal::array<uint32_t, N> source_sizes;
+  uint32_t index_stride;
+};
diff --git a/aten/src/ATen/native/mps/kernels/Indexing.metal b/aten/src/ATen/native/mps/kernels/Indexing.metal
index 09fe380b4c2b3..dcc0c4ff3cb82 100644
--- a/aten/src/ATen/native/mps/kernels/Indexing.metal
+++ b/aten/src/ATen/native/mps/kernels/Indexing.metal
@@ -1,3 +1,4 @@
+#include <ATen/native/mps/kernels/Indexing.h>
 #include <c10/metal/atomic.h>
 #include <c10/metal/error.h>
 #include <c10/metal/indexing.h>
@@ -275,6 +276,88 @@ REGISTER_INDEX_OP(put_accumulate, bool, bool);
 REGISTER_INDEX_OP(put_accumulate, float2, float2);
 REGISTER_INDEX_OP(put_accumulate, half2, half2);
 
+struct IndexReduceOp {
+  template <typename T>
+  static T prod(T a, T b) {
+    return c10::metal::mul(a, b);
+  }
+
+  template <typename T>
+  static T mean(T a, T b) {
+    return a + b;
+  }
+
+  template <typename T>
+  static T amin(T a, T b) {
+    return min(a, b);
+  }
+
+  template <typename T>
+  static T amax(T a, T b) {
+    return max(a, b);
+  }
+};
+
+template <typename T, typename IT, T (*ReduceOp)(T, T)>
+kernel void index_reduce(
+    device AtomicType_t<T>* self [[buffer(0)]],
+    device IT* index [[buffer(1)]],
+    device T* source [[buffer(2)]],
+    constant IndexReduceParams<>& params [[buffer(3)]],
+    uint tid [[thread_position_in_grid]]) {
+  uint32_t tid_ = tid;
+  long source_offset = 0;
+  long self_offset = 0;
+
+  for (int32_t dim = params.ndim - 1; dim >= 0; dim--) {
+    auto source_size = params.source_sizes[dim];
+    auto dim_idx = tid_ % source_size;
+
+    source_offset += dim_idx * params.source_strides[dim];
+
+    if (dim == params.reduce_dim) {
+      uint32_t self_dim_idx =
+          static_cast<uint32_t>(index[dim_idx * params.index_stride]);
+      self_offset += self_dim_idx * params.self_strides[dim];
+    } else {
+      self_offset += dim_idx * params.self_strides[dim];
+    }
+
+    tid_ /= source_size;
+  }
+
+  T source_elem = source[source_offset];
+
+  AtomicType<T>::atomic_binary_op(self, self_offset, source_elem, ReduceOp);
+}
+
+#define REGISTER_INDEX_REDUCE_OP(ReduceOp, T, IT)                  \
+  template [[host_name("index_reduce_" #ReduceOp "_" #T "_" #IT)]] \
+  kernel void index_reduce<T, IT, IndexReduceOp::ReduceOp<T>>(     \
+      device AtomicType_t<T> * self [[buffer(0)]],                 \
+      device IT * index [[buffer(1)]],                             \
+      device T * source [[buffer(2)]],                             \
+      constant IndexReduceParams<> & params [[buffer(3)]],         \
+      uint tid [[thread_position_in_grid]]);
+
+#define REGISTER_INDEX_REDUCE_OP_ALL_REDUCE_TYPES(T, IT) \
+  REGISTER_INDEX_REDUCE_OP(amax, T, IT);                 \
+  REGISTER_INDEX_REDUCE_OP(mean, T, IT);                 \
+  REGISTER_INDEX_REDUCE_OP(amin, T, IT);                 \
+  REGISTER_INDEX_REDUCE_OP(prod, T, IT);
+
+#define REGISTER_INDEX_REDUCE_OP_ALL_INDEX_TYPES(T)  \
+  REGISTER_INDEX_REDUCE_OP_ALL_REDUCE_TYPES(T, int); \
+  REGISTER_INDEX_REDUCE_OP_ALL_REDUCE_TYPES(T, long);
+
+REGISTER_INDEX_REDUCE_OP_ALL_INDEX_TYPES(float);
+REGISTER_INDEX_REDUCE_OP_ALL_INDEX_TYPES(half);
+REGISTER_INDEX_REDUCE_OP_ALL_INDEX_TYPES(bfloat);
+REGISTER_INDEX_REDUCE_OP_ALL_INDEX_TYPES(int);
+REGISTER_INDEX_REDUCE_OP_ALL_INDEX_TYPES(short);
+REGISTER_INDEX_REDUCE_OP_ALL_INDEX_TYPES(char);
+REGISTER_INDEX_REDUCE_OP_ALL_INDEX_TYPES(uchar);
+
 template <typename StridesT, typename DataT>
 kernel void kernel_index_offsets(
     constant StridesT* strides [[buffer(0)]],
@@ -442,6 +525,154 @@ kernel void index_copy_strided(
   }
 }
 
+// Scatter-based index_fill: each thread writes exactly one element.
+// Launches indices_numel * slice_numel threads, where
+// slice_numel = numel / self.size(dim).
+
+// Dense (contiguous output): offset computed analytically.
+template <typename T>
+kernel void index_fill_dense(
+    device T* output,
+    constant long* indices,
+    constant T& fill_val,
+    constant long& dim_size,
+    constant long&
+        dim_stride, // output.stride(dim); trailing product for contiguous
+    constant long& slice_numel,
+    uint thread_index [[thread_position_in_grid]]) {
+  long j =
+      thread_index % slice_numel; // position within the slice (non-dim dims)
+  long i = thread_index / slice_numel; // which index
+  long idx = indices[i];
+  if (idx < 0) {
+    idx += dim_size;
+  }
+  long before = j / dim_stride;
+  long after = j % dim_stride;
+  output[before * (dim_size * dim_stride) + idx * dim_stride + after] =
+      fill_val;
+}
+
+// Strided (non-contiguous output or index): offset computed via slice strides.
+template <typename T>
+kernel void index_fill_strided(
+    device T* output,
+    constant long* indices,
+    constant T& fill_val,
+    constant long& dim_size,
+    constant long& dim_out_stride, // output.stride(dim)
+    constant long*
+        slice_sizes, // output sizes with dim removed (length = ndim-1)
+    constant long*
+        slice_out_strides, // output strides with dim removed (length = ndim-1)
+    constant uint& slice_ndim, // ndim - 1
+    constant long& slice_numel,
+    constant long& indices_stride, // index.stride(0)
+    uint thread_index [[thread_position_in_grid]]) {
+  long j = thread_index % slice_numel;
+  long i = thread_index / slice_numel;
+  long idx = indices[i * indices_stride];
+  if (idx < 0) {
+    idx += dim_size;
+  }
+  int slice_pos[max_ndim];
+  pos_from_thread_index(int(j), slice_pos, slice_sizes, slice_ndim);
+  long out_offset = offset_from_coord(slice_pos, slice_out_strides, slice_ndim);
+  output[out_offset + idx * dim_out_stride] = fill_val;
+}
+
+// Two-pass mask-based index_fill for large index counts.
+// Avoids write conflicts from duplicate indices and improves cache locality.
+//
+// Pass 0 (zero_mask): clear the mask buffer in the same compute encoder to
+//   avoid a blit/compute encoder switch from using at::zeros().
+// Pass 1 (set_mask): mark which dim-positions to fill.
+// Pass 2 (from_mask): fill all elements whose dim-coordinate is marked.
+
+kernel void index_fill_zero_mask(
+    device bool* mask,
+    uint i [[thread_position_in_grid]]) {
+  mask[i] = false;
+}
+
+kernel void index_fill_set_mask(
+    device bool* mask,
+    constant long* indices,
+    constant long& dim_size,
+    constant long& indices_stride,
+    uint thread_index [[thread_position_in_grid]]) {
+  long idx = indices[thread_index * indices_stride];
+  if (idx < 0)
+    idx += dim_size;
+  mask[idx] = true;
+}
+
+// Dense: fills output elements whose dim-coordinate is marked in mask.
+// Uses a 3D thread grid (inner, dim, outer) to avoid any integer division:
+//   x = inner index  (0..inner_size-1 where inner_size = stride(dim))
+//   y = dim coord    (0..dim_size-1)           <-- no division needed
+//   z = outer index  (0..outer_size-1)
+// Output offset = (outer * dim_size + dim) * inner_size + inner.
+template <typename T>
+kernel void index_fill_dense_from_mask(
+    device T* output,
+    constant bool* mask,
+    constant T& fill_val,
+    constant uint& dim_size,
+    constant uint& inner_size,
+    uint3 thread_pos [[thread_position_in_grid]]) {
+  uint inner_idx = thread_pos.x;
+  uint dim_idx = thread_pos.y;
+  uint outer_idx = thread_pos.z;
+  if (mask[dim_idx]) {
+    output[(outer_idx * dim_size + dim_idx) * inner_size + inner_idx] =
+        fill_val;
+  }
+}
+
+// Strided: thread i fills its element if its dim-coordinate is marked.
+template <typename T>
+kernel void index_fill_strided_from_mask(
+    device T* output,
+    constant bool* mask,
+    constant T& fill_val,
+    constant long* sizes,
+    constant long* out_strides,
+    constant uint& dim,
+    constant uint& ndim,
+    uint thread_index [[thread_position_in_grid]]) {
+  int pos[max_ndim];
+  pos_from_thread_index(int(thread_index), pos, sizes, ndim);
+  if (mask[pos[dim]]) {
+    long out_offset = offset_from_coord(pos, out_strides, ndim);
+    output[out_offset] = fill_val;
+  }
+}
+
+#define INSTANTIATE_INDEX_FILL(T)                  \
+  template [[host_name("index_fill_dense_" #T)]]   \
+  kernel void index_fill_dense<T>(                 \
+      device T*,                                   \
+      constant long*,                              \
+      constant T&,                                 \
+      constant long&,                              \
+      constant long&,                              \
+      constant long&,                              \
+      uint);                                       \
+  template [[host_name("index_fill_strided_" #T)]] \
+  kernel void index_fill_strided<T>(               \
+      device T*,                                   \
+      constant long*,                              \
+      constant T&,                                 \
+      constant long&,                              \
+      constant long&,                              \
+      constant long*,                              \
+      constant long*,                              \
+      constant uint&,                              \
+      constant long&,                              \
+      constant long&,                              \
+      uint);
+
 #define INSTANTIATE_INDEX_COPY(T, index_t)                      \
   template [[host_name("index_copy_dense_" #T "_" #index_t)]]   \
   kernel void index_copy_dense<T, index_t>(                     \
@@ -516,3 +747,281 @@ INSTANTIATE_INDEX_COPY(float2, int);
 INSTANTIATE_INDEX_COPY(float2, long);
 INSTANTIATE_INDEX_COPY(half2, int);
 INSTANTIATE_INDEX_COPY(half2, long);
+
+INSTANTIATE_INDEX_FILL(float);
+INSTANTIATE_INDEX_FILL(bool);
+INSTANTIATE_INDEX_FILL(half);
+INSTANTIATE_INDEX_FILL(bfloat);
+INSTANTIATE_INDEX_FILL(int);
+INSTANTIATE_INDEX_FILL(long);
+INSTANTIATE_INDEX_FILL(short);
+INSTANTIATE_INDEX_FILL(char);
+INSTANTIATE_INDEX_FILL(uchar);
+INSTANTIATE_INDEX_FILL(float2);
+INSTANTIATE_INDEX_FILL(half2);
+
+#define INSTANTIATE_INDEX_FILL_FROM_MASK(T)                  \
+  template [[host_name("index_fill_dense_from_mask_" #T)]]   \
+  kernel void index_fill_dense_from_mask<T>(                 \
+      device T*,                                             \
+      constant bool*,                                        \
+      constant T&,                                           \
+      constant uint&,                                        \
+      constant uint&,                                        \
+      uint3);                                                \
+  template [[host_name("index_fill_strided_from_mask_" #T)]] \
+  kernel void index_fill_strided_from_mask<T>(               \
+      device T*,                                             \
+      constant bool*,                                        \
+      constant T&,                                           \
+      constant long*,                                        \
+      constant long*,                                        \
+      constant uint&,                                        \
+      constant uint&,                                        \
+      uint);
+
+INSTANTIATE_INDEX_FILL_FROM_MASK(float)
+INSTANTIATE_INDEX_FILL_FROM_MASK(half)
+INSTANTIATE_INDEX_FILL_FROM_MASK(bfloat)
+INSTANTIATE_INDEX_FILL_FROM_MASK(int)
+INSTANTIATE_INDEX_FILL_FROM_MASK(long)
+INSTANTIATE_INDEX_FILL_FROM_MASK(short)
+INSTANTIATE_INDEX_FILL_FROM_MASK(char)
+INSTANTIATE_INDEX_FILL_FROM_MASK(uchar)
+INSTANTIATE_INDEX_FILL_FROM_MASK(bool)
+INSTANTIATE_INDEX_FILL_FROM_MASK(float2)
+INSTANTIATE_INDEX_FILL_FROM_MASK(half2)
+
+// Nonzero kernel implementation using prefix-sum + scatter approach.
+//
+// Step 1 (count_nonzero_prefix_sum): Each threadgroup computes an exclusive
+// prefix sum of the nonzero flags over its chunk. Per-threadgroup totals are
+// written to block_sums.
+//
+// Step 2 (prefix_sum_blocks): A single threadgroup computes the exclusive
+// prefix sum of block_sums → block_offsets and writes the total nonzero count
+// to a 1-element buffer. The host reads back only that single int, then
+// allocates the output tensor.
+//
+// Step 3 (scatter_nonzero_indices): Each thread with a nonzero element writes
+// its multi-dimensional indices into the output at the position determined by
+// block_offsets[tgid] + prefix[tid].
+
+template <typename T, enable_if_t<!is_complex_v<T>, bool> = true>
+inline bool is_nonzero(T val) {
+  return val != T(0);
+}
+
+template <typename T, enable_if_t<is_complex_v<T>, bool> = true>
+inline bool is_nonzero(T val) {
+  return val.x != 0 || val.y != 0;
+}
+
+template <typename T>
+[[max_total_threads_per_threadgroup(1024)]]
+kernel void count_nonzero_prefix_sum(
+    const device T* input [[buffer(0)]],
+    device int* prefix [[buffer(1)]],
+    device int* block_sums [[buffer(2)]],
+    uint tid [[thread_position_in_grid]],
+    uint lid [[thread_position_in_threadgroup]],
+    uint tgsize [[threads_per_threadgroup]],
+    uint tgid [[threadgroup_position_in_grid]],
+    uint simd_lane_id [[thread_index_in_simdgroup]],
+    uint simd_group_id [[simdgroup_index_in_threadgroup]]) {
+  uint num_simds = (tgsize + simdgroup_size - 1) / simdgroup_size;
+
+  int flag = is_nonzero(input[tid]) ? 1 : 0;
+
+  // Inclusive prefix sum within SIMD group using shuffle
+  int val = flag;
+  for (uint offset = 1; offset < simdgroup_size; offset <<= 1) {
+    int other = simd_shuffle_and_fill_up(val, 0, static_cast<ushort>(offset));
+    val += other;
+  }
+
+  // The last lane in each simd group writes its total.
+  // For full groups this is lane 31; for the last (partial) group we compute
+  // which lane is actually last.
+  threadgroup int simdgroup_totals[32];
+  bool is_last_lane_in_simd;
+  if (simd_group_id < num_simds - 1) {
+    is_last_lane_in_simd = (simd_lane_id == simdgroup_size - 1);
+  } else {
+    uint lanes_in_last = tgsize - simd_group_id * simdgroup_size;
+    is_last_lane_in_simd = (simd_lane_id == lanes_in_last - 1);
+  }
+  if (is_last_lane_in_simd) {
+    simdgroup_totals[simd_group_id] = val;
+  }
+  threadgroup_barrier(mem_flags::mem_threadgroup);
+
+  // First simd group computes exclusive prefix sum of simd group totals
+  threadgroup int simdgroup_offsets[32];
+  if (simd_group_id == 0) {
+    int sg_val =
+        (simd_lane_id < num_simds) ? simdgroup_totals[simd_lane_id] : 0;
+    for (uint offset = 1; offset < simdgroup_size; offset <<= 1) {
+      int other =
+          simd_shuffle_and_fill_up(sg_val, 0, static_cast<ushort>(offset));
+      sg_val += other;
+    }
+    int exclusive = simd_shuffle_and_fill_up(sg_val, 0, static_cast<ushort>(1));
+    simdgroup_offsets[simd_lane_id] = exclusive;
+  }
+  threadgroup_barrier(mem_flags::mem_threadgroup);
+
+  int exclusive_val = val - flag + simdgroup_offsets[simd_group_id];
+
+  prefix[tid] = exclusive_val;
+
+  if (lid == tgsize - 1) {
+    block_sums[tgid] =
+        simdgroup_offsets[num_simds - 1] + simdgroup_totals[num_simds - 1];
+  }
+}
+
+// Step 2: exclusive prefix sum of block_sums → block_offsets, and write
+// total nonzero count to a 1-element buffer.  Runs in a single threadgroup.
+// Each thread handles ceil(num_blocks / tgsize) consecutive blocks via a
+// serial loop, then the per-thread totals are scanned in parallel.
+[[max_total_threads_per_threadgroup(1024)]]
+kernel void prefix_sum_blocks(
+    const device int* block_sums [[buffer(0)]],
+    device int* block_offsets [[buffer(1)]],
+    device int* total_nonzero [[buffer(2)]],
+    constant uint& num_blocks [[buffer(3)]],
+    uint lid [[thread_position_in_threadgroup]],
+    uint tgsize [[threads_per_threadgroup]],
+    uint simd_lane_id [[thread_index_in_simdgroup]],
+    uint simd_group_id [[simdgroup_index_in_threadgroup]]) {
+  uint num_simds = (tgsize + simdgroup_size - 1) / simdgroup_size;
+
+  // Each thread handles a contiguous chunk of blocks
+  uint chunk_size = (num_blocks + tgsize - 1) / tgsize;
+  uint start = lid * chunk_size;
+  uint end = min(start + chunk_size, num_blocks);
+
+  // Serial sum over this thread's chunk
+  int chunk_total = 0;
+  for (uint i = start; i < end; i++) {
+    chunk_total += block_sums[i];
+  }
+
+  // Parallel inclusive prefix sum of chunk_totals across threads
+  int val = chunk_total;
+  for (uint offset = 1; offset < simdgroup_size; offset <<= 1) {
+    int other = simd_shuffle_and_fill_up(val, 0, static_cast<ushort>(offset));
+    val += other;
+  }
+
+  threadgroup int simdgroup_totals[32];
+  bool is_last_lane_in_simd;
+  if (simd_group_id < num_simds - 1) {
+    is_last_lane_in_simd = (simd_lane_id == simdgroup_size - 1);
+  } else {
+    uint lanes_in_last = tgsize - simd_group_id * simdgroup_size;
+    is_last_lane_in_simd = (simd_lane_id == lanes_in_last - 1);
+  }
+  if (is_last_lane_in_simd) {
+    simdgroup_totals[simd_group_id] = val;
+  }
+  threadgroup_barrier(mem_flags::mem_threadgroup);
+
+  threadgroup int simdgroup_offsets[32];
+  if (simd_group_id == 0) {
+    int sg_val =
+        (simd_lane_id < num_simds) ? simdgroup_totals[simd_lane_id] : 0;
+    for (uint offset = 1; offset < simdgroup_size; offset <<= 1) {
+      int other =
+          simd_shuffle_and_fill_up(sg_val, 0, static_cast<ushort>(offset));
+      sg_val += other;
+    }
+    int exclusive = simd_shuffle_and_fill_up(sg_val, 0, static_cast<ushort>(1));
+    simdgroup_offsets[simd_lane_id] = exclusive;
+  }
+  threadgroup_barrier(mem_flags::mem_threadgroup);
+
+  // This thread's exclusive offset = inclusive_scan - chunk_total +
+  // simdgroup_offset
+  int thread_offset = val - chunk_total + simdgroup_offsets[simd_group_id];
+
+  // Write block_offsets for this thread's chunk using a serial exclusive scan
+  int running = thread_offset;
+  for (uint i = start; i < end; i++) {
+    block_offsets[i] = running;
+    running += block_sums[i];
+  }
+
+  if (lid == tgsize - 1) {
+    *total_nonzero =
+        simdgroup_offsets[num_simds - 1] + simdgroup_totals[num_simds - 1];
+  }
+}
+
+// Scatter the multi-dimensional indices of nonzero elements.
+// Output layout: out[position * ndim + d] = index along dimension d.
+template <typename T>
+[[max_total_threads_per_threadgroup(1024)]]
+kernel void scatter_nonzero_indices(
+    const device T* input [[buffer(0)]],
+    const device int* prefix [[buffer(1)]],
+    device int64_t* output [[buffer(2)]],
+    constant int& ndim [[buffer(3)]],
+    constant int64_t* sizes [[buffer(4)]],
+    constant int* block_offsets [[buffer(5)]],
+    constant int& max_entries [[buffer(6)]],
+    uint tid [[thread_position_in_grid]],
+    uint tgid [[threadgroup_position_in_grid]]) {
+  if (!is_nonzero(input[tid]))
+    return;
+
+  int pos = block_offsets[tgid] + prefix[tid];
+  if (pos >= max_entries)
+    return;
+
+  uint flat = tid;
+  for (int d = ndim - 1; d >= 0; d--) {
+    int64_t dim_size = sizes[d];
+    output[pos * ndim + d] =
+        static_cast<int64_t>(flat % static_cast<uint>(dim_size));
+    flat /= static_cast<uint>(dim_size);
+  }
+}
+
+#define REGISTER_NONZERO_KERNELS(DTYPE)                                      \
+  template [[host_name("count_nonzero_prefix_sum_" #DTYPE)]] [[kernel]] void \
+  count_nonzero_prefix_sum<DTYPE>(                                           \
+      const device DTYPE* input [[buffer(0)]],                               \
+      device int* prefix [[buffer(1)]],                                      \
+      device int* block_sums [[buffer(2)]],                                  \
+      uint tid [[thread_position_in_grid]],                                  \
+      uint lid [[thread_position_in_threadgroup]],                           \
+      uint tgsize [[threads_per_threadgroup]],                               \
+      uint tgid [[threadgroup_position_in_grid]],                            \
+      uint simd_lane_id [[thread_index_in_simdgroup]],                       \
+      uint simd_group_id [[simdgroup_index_in_threadgroup]]);                \
+                                                                             \
+  template [[host_name("scatter_nonzero_indices_" #DTYPE)]] [[kernel]] void  \
+  scatter_nonzero_indices<DTYPE>(                                            \
+      const device DTYPE* input [[buffer(0)]],                               \
+      const device int* prefix [[buffer(1)]],                                \
+      device int64_t* output [[buffer(2)]],                                  \
+      constant int& ndim [[buffer(3)]],                                      \
+      constant int64_t* sizes [[buffer(4)]],                                 \
+      constant int* block_offsets [[buffer(5)]],                             \
+      constant int& max_entries [[buffer(6)]],                               \
+      uint tid [[thread_position_in_grid]],                                  \
+      uint tgid [[threadgroup_position_in_grid]])
+
+REGISTER_NONZERO_KERNELS(float);
+REGISTER_NONZERO_KERNELS(half);
+REGISTER_NONZERO_KERNELS(bfloat);
+REGISTER_NONZERO_KERNELS(long);
+REGISTER_NONZERO_KERNELS(int);
+REGISTER_NONZERO_KERNELS(short);
+REGISTER_NONZERO_KERNELS(char);
+REGISTER_NONZERO_KERNELS(uchar);
+REGISTER_NONZERO_KERNELS(bool);
+REGISTER_NONZERO_KERNELS(float2);
+REGISTER_NONZERO_KERNELS(half2);
diff --git a/aten/src/ATen/native/mps/kernels/LinearAlgebra.h b/aten/src/ATen/native/mps/kernels/LinearAlgebra.h
index a77122ec55b09..556e9cda92718 100644
--- a/aten/src/ATen/native/mps/kernels/LinearAlgebra.h
+++ b/aten/src/ATen/native/mps/kernels/LinearAlgebra.h
@@ -21,3 +21,8 @@ struct UnpackPivotsParams {
   uint32_t pivots_batch_stride;
   uint32_t dim_size;
 };
+
+struct QrParams {
+  uint32_t m;
+  uint32_t n;
+};
diff --git a/aten/src/ATen/native/mps/kernels/LinearAlgebra.metal b/aten/src/ATen/native/mps/kernels/LinearAlgebra.metal
index 6c093475dde1d..ad68c9549fa00 100644
--- a/aten/src/ATen/native/mps/kernels/LinearAlgebra.metal
+++ b/aten/src/ATen/native/mps/kernels/LinearAlgebra.metal
@@ -1,4 +1,6 @@
 #include <ATen/native/mps/kernels/LinearAlgebra.h>
+#include <c10/metal/common.h>
+#include <c10/metal/reduction_utils.h>
 #include <c10/metal/utils.h>
 #include <metal_array>
 #include <metal_simdgroup>
@@ -931,6 +933,182 @@ kernel void unpack_pivots(
   }
 }
 
+template <typename T>
+kernel void linalg_qr_householder(
+    device T* A [[buffer(0)]],
+    device T* Q [[buffer(1)]],
+    device T* R [[buffer(2)]],
+    device int* info [[buffer(3)]],
+    constant QrParams& params [[buffer(4)]],
+    device T* v_work [[buffer(5)]],
+    uint3 thread_pos [[thread_position_in_threadgroup]],
+    uint3 tpg [[threads_per_threadgroup]],
+    uint3 tg_pos [[threadgroup_position_in_grid]]) {
+  using opmath_t = c10::metal::opmath_t<T>;
+
+  const uint32_t tid = thread_pos.x;
+  const uint32_t group_size = tpg.x;
+  const uint32_t m = params.m;
+  const uint32_t n = params.n;
+
+  // Batch indexing
+  const uint32_t batch_idx = tg_pos.x;
+  const uint32_t A_stride = m * n;
+  const uint32_t Q_stride = m * m;
+  const uint32_t R_stride = m * n;
+  const uint32_t v_stride = m;
+
+  device T* A_batch = A + batch_idx * A_stride;
+  device T* Q_batch = Q + batch_idx * Q_stride;
+  device T* R_batch = R + batch_idx * R_stride;
+  device T* v_batch = v_work + batch_idx * v_stride;
+
+  constexpr auto kMaxThreadsPerThreadgroup = 1024;
+  constexpr auto kMaxSIMDGroups =
+      kMaxThreadsPerThreadgroup / c10::metal::simdgroup_size;
+
+  threadgroup opmath_t scratch[kMaxSIMDGroups];
+  threadgroup opmath_t tau_shared;
+
+  // initialize Q = Identity (m x m)
+  for (uint32_t i = tid; i < m * m; i += group_size) {
+    Q_batch[i] = static_cast<T>((i / m == i % m) ? 1.0 : 0.0);
+  }
+
+  // initialize R = A (m x n)
+  for (uint32_t i = tid; i < m * n; i += group_size) {
+    R_batch[i] = A_batch[i];
+  }
+  threadgroup_barrier(mem_flags::mem_device);
+
+  for (uint32_t k = 0; k < min(m, n); k++) {
+    // Step 1: compute norm of R[k:m, k] and copy to v_batch
+    opmath_t norm_sq = 0.0;
+    for (uint32_t i = k + tid; i < m; i += group_size) {
+      T r_ik = R_batch[i * n + k];
+      v_batch[i] = r_ik;
+      const auto val = static_cast<opmath_t>(r_ik);
+      norm_sq = fma(val, val, norm_sq);
+    }
+    const auto norm = ::metal::precise::sqrt(
+        c10::metal::threadgroup_sum(scratch, norm_sq, tid, group_size));
+
+    // scale norm_eps by matrix dimension to handle accumulated error
+    const auto norm_eps = ::metal::numeric_limits<opmath_t>::epsilon() * m;
+    const auto tau_eps = ::metal::numeric_limits<opmath_t>::epsilon();
+
+    // Step 2: compute Householder vector and tau
+    if (tid == 0) {
+      // LAPACK convention: skip reflection for last row to preserve natural
+      // sign When k == m - 1, there's only one element in the column, so
+      // reflection would just flip its sign. Instead, preserve whatever value
+      // emerged from prior transformations to match LAPACK's behavior.
+      if (fabs(norm) < norm_eps || k == m - 1) {
+        tau_shared = 0.0;
+      } else {
+        opmath_t alpha = static_cast<opmath_t>(v_batch[k]);
+        opmath_t sign_alpha = (alpha >= 0.0) ? 1.0 : -1.0;
+        opmath_t beta = sign_alpha * norm;
+        opmath_t u1 = alpha + beta;
+
+        tau_shared = 1.0 + fabs(alpha) / norm;
+
+        v_batch[k] = static_cast<T>(1.0); // always 1 by construction
+        for (uint32_t i = k + 1; i < m; i++) {
+          v_batch[i] = static_cast<T>(static_cast<opmath_t>(v_batch[i]) / u1);
+        }
+
+        R_batch[k * n + k] = static_cast<T>(-beta);
+      }
+    }
+    threadgroup_barrier(mem_flags::mem_device);
+
+    const auto tau = tau_shared;
+    if (tau < tau_eps)
+      continue;
+
+    // (zero out column k below diagonal)
+    for (uint32_t i = k + 1 + tid; i < m; i += group_size) {
+      R_batch[i * n + k] = static_cast<T>(0.0);
+    }
+
+    // Step 3: apply reflection to trailing columns of R
+    // Parallelize across columns: each SIMD group (32 threads) handles one
+    // column
+    uint32_t simd_lane = tid % c10::metal::simdgroup_size;
+    uint32_t simd_group_id = tid / c10::metal::simdgroup_size;
+    uint32_t num_simd_groups = group_size / c10::metal::simdgroup_size;
+
+    for (uint32_t j_base = k + 1; j_base < n; j_base += num_simd_groups) {
+      uint32_t j = j_base + simd_group_id;
+      if (j < n) {
+        // Each SIMD group computes dot product for its column
+        // Use SIMD reduction within the group
+        opmath_t dot = 0.0;
+        for (uint32_t i = k + simd_lane; i < m; i += 32) {
+          opmath_t v_i = static_cast<opmath_t>(v_batch[i]);
+          opmath_t r_ij = static_cast<opmath_t>(R_batch[i * n + j]);
+          dot = fma(v_i, r_ij, dot);
+        }
+        opmath_t vt_col = simd_sum(dot);
+        opmath_t factor = tau * vt_col;
+
+        // Update column
+        for (uint32_t i = k + simd_lane; i < m; i += 32) {
+          opmath_t v_i = static_cast<opmath_t>(v_batch[i]);
+          opmath_t r_ij = static_cast<opmath_t>(R_batch[i * n + j]);
+          R_batch[i * n + j] = static_cast<T>(r_ij - v_i * factor);
+        }
+      }
+    }
+    threadgroup_barrier(mem_flags::mem_device);
+
+    // Step 4: accumulate Q = Q * H_k
+    // each SIMD group handles one row
+    for (uint32_t i_base = 0; i_base < m; i_base += num_simd_groups) {
+      uint32_t i = i_base + simd_group_id;
+      if (i < m) {
+        opmath_t dot = 0.0;
+        for (uint32_t j = k + simd_lane; j < m; j += 32) {
+          opmath_t v_j = static_cast<opmath_t>(v_batch[j]);
+          opmath_t q_ij = static_cast<opmath_t>(Q_batch[i * m + j]);
+          dot = fma(q_ij, v_j, dot);
+        }
+        opmath_t row_v = simd_sum(dot);
+        opmath_t factor = tau * row_v;
+
+        // Update row
+        for (uint32_t j = k + simd_lane; j < m; j += 32) {
+          opmath_t v_j = static_cast<opmath_t>(v_batch[j]);
+          opmath_t q_ij = static_cast<opmath_t>(Q_batch[i * m + j]);
+          Q_batch[i * m + j] = static_cast<T>(q_ij - v_j * factor);
+        }
+      }
+    }
+
+    threadgroup_barrier(mem_flags::mem_device);
+  }
+
+  if (tid == 0) {
+    info[0] = 0;
+  }
+}
+
+#define REGISTER_QR(T)                                \
+  template [[host_name("linalg_qr_householder_" #T)]] \
+  kernel void linalg_qr_householder<T>(               \
+      device T * A [[buffer(0)]],                     \
+      device T * Q [[buffer(1)]],                     \
+      device T * R [[buffer(2)]],                     \
+      device int* info [[buffer(3)]],                 \
+      constant QrParams& params [[buffer(4)]],        \
+      device T* v_work [[buffer(5)]],                 \
+      uint3 tid [[thread_position_in_threadgroup]],   \
+      uint3 tpg [[threads_per_threadgroup]],          \
+      uint3 tg_pos [[threadgroup_position_in_grid]]);
+
+REGISTER_QR(float);
+
 #define INSTANTIATE_MM_OPS(DTYPE)                                           \
   template [[host_name("matmul_" #DTYPE)]] kernel void matmul<DTYPE>(       \
       constant DTYPE * mat1Data [[buffer(0)]],                              \
diff --git a/aten/src/ATen/native/mps/kernels/ReduceOps.h b/aten/src/ATen/native/mps/kernels/ReduceOps.h
new file mode 100644
index 0000000000000..e978c356efa22
--- /dev/null
+++ b/aten/src/ATen/native/mps/kernels/ReduceOps.h
@@ -0,0 +1,17 @@
+#pragma once
+#include <c10/metal/common.h>
+
+#define MAX_THREADGROUP_SIZE static_cast<uint32_t>(1024)
+
+template <unsigned N = c10::metal::max_ndim>
+struct NormParams {
+  float p;
+  uint32_t reduction_size;
+  uint32_t ndim;
+
+  ::c10::metal::array<uint32_t, N> input_sizes;
+  ::c10::metal::array<uint32_t, N> input_strides;
+
+  ::c10::metal::array<uint32_t, N> output_sizes;
+  ::c10::metal::array<uint32_t, N> output_strides;
+};
diff --git a/aten/src/ATen/native/mps/kernels/ReduceOps.metal b/aten/src/ATen/native/mps/kernels/ReduceOps.metal
new file mode 100644
index 0000000000000..79ad71de98db3
--- /dev/null
+++ b/aten/src/ATen/native/mps/kernels/ReduceOps.metal
@@ -0,0 +1,177 @@
+#include <ATen/native/mps/kernels/ReduceOps.h>
+#include <c10/metal/atomic.h>
+#include <c10/metal/utils.h>
+#include <metal_array>
+#include <metal_stdlib>
+
+using namespace metal;
+using namespace c10::metal;
+
+struct norm_abs_functor {
+  template <typename T, enable_if_t<!is_complex_v<T>, bool> = true>
+  inline T operator()(const T x) {
+    return static_cast<T>(::precise::abs(x));
+  }
+
+  template <typename T, enable_if_t<is_complex_v<T>, bool> = true>
+  inline float operator()(const T x) {
+    const auto abs_2 = ::precise::abs(float2(x));
+    return c10::metal::hypot(abs_2.x, abs_2.y);
+  }
+};
+
+// `reduction_idx` is the index of a particular batch of input elements that all
+// get reduced to one output element. `reduction_element_idx` is the index of
+// just one input element within its batch.
+static uint32_t get_input_offset(
+    uint32_t reduction_element_idx,
+    uint32_t reduction_idx,
+    constant NormParams<>& params) {
+  uint32_t input_offset = 0;
+
+  for (int32_t dim = params.ndim - 1; dim >= 0; dim--) {
+    auto input_dim_size = params.input_sizes[dim];
+    auto output_dim_size = params.output_sizes[dim];
+
+    // If the the input and output have the same size for this dim, then this
+    // dim is not being reduced, so we index by `reduction_idx`
+    if (input_dim_size == output_dim_size) {
+      auto index_in_dim = reduction_idx % input_dim_size;
+      reduction_idx /= input_dim_size;
+      input_offset += index_in_dim * params.input_strides[dim];
+
+      // Otherwise, this dim is being reduced, so we index by
+      // `reduction_element_idx`
+    } else {
+      auto index_in_dim = reduction_element_idx % input_dim_size;
+      reduction_element_idx /= input_dim_size;
+      input_offset += index_in_dim * params.input_strides[dim];
+    }
+  }
+  return input_offset;
+}
+
+// In this kernel, each threadgroup is responsible for calculating one element
+// of the output.
+// TI - dtype of the input tensor.
+// TO - dtype of the output tensor.
+template <typename TI, typename TO>
+kernel void norm(
+    constant TI* input [[buffer(0)]],
+    device TO* output [[buffer(1)]],
+    constant NormParams<>& params [[buffer(2)]],
+    uint tid [[thread_position_in_threadgroup]],
+    uint tptg [[threads_per_threadgroup]],
+    uint tgid [[threadgroup_position_in_grid]],
+    uint simd_lane_id [[thread_index_in_simdgroup]],
+    uint simdgroup_id [[simdgroup_index_in_threadgroup]],
+    uint simdgroup_size [[threads_per_simdgroup]]) {
+  using TA = opmath_t<TO>;
+  TA output_val = 0;
+  const auto p = static_cast<TA>(params.p);
+
+  if (p == INFINITY) {
+    output_val = -INFINITY;
+  } else if (p == -INFINITY) {
+    output_val = INFINITY;
+  }
+
+  // First, all the input elements assigned to the threadgroup are divided
+  // between all the threads in the threadgroup, and each thread reduces those
+  // elements down to one partial `output_val`.
+  for (uint32_t reduction_element_idx = tid;
+       reduction_element_idx < params.reduction_size;
+       reduction_element_idx += tptg) {
+    auto input_elem =
+        input[get_input_offset(reduction_element_idx, tgid, params)];
+    auto input_abs = static_cast<TA>(norm_abs_functor()(input_elem));
+
+    if (p == INFINITY) {
+      output_val = max(input_abs, output_val);
+
+    } else if (p == -INFINITY) {
+      output_val = min(input_abs, output_val);
+
+    } else if (p == 0) {
+      output_val += (input_abs == 0) ? 0 : 1;
+
+    } else {
+      output_val += static_cast<TA>(::precise::pow(input_abs, p));
+    }
+  }
+
+  // Next, all the threads in a threadgroup reduce their `output_val`s together
+  // with a series of SIMD group reductions.
+  auto threads_remaining = tptg;
+  threadgroup TA shared_outputs[MAX_THREADGROUP_SIZE];
+
+  while (threads_remaining > 1) {
+    if (p == INFINITY) {
+      output_val = simd_max(output_val);
+    } else if (p == -INFINITY) {
+      output_val = simd_min(output_val);
+    } else {
+      output_val = simd_sum(output_val);
+    }
+
+    threads_remaining =
+        (threads_remaining + simdgroup_size - 1) / simdgroup_size;
+
+    if (threads_remaining > 1) {
+      // One thread from each SIMD group writes to a shared buffer
+      if (simd_lane_id == 0) {
+        shared_outputs[simdgroup_id] = output_val;
+      }
+
+      threadgroup_barrier(mem_flags::mem_threadgroup);
+
+      // The remaining threads each read one of the partial outputs from the
+      // shared buffer
+      if (tid < threads_remaining) {
+        output_val = shared_outputs[tid];
+      } else {
+        return;
+      }
+    }
+  }
+
+  // Finally, one thread in the threadgroup writes the final output
+  if (tid == 0) {
+    uint32_t output_offset = 0;
+    uint32_t reduction_idx = tgid;
+
+    for (int32_t dim = params.ndim - 1; dim >= 0; dim--) {
+      auto output_dim_size = params.output_sizes[dim];
+
+      if (output_dim_size > 1) {
+        auto index_in_dim = reduction_idx % output_dim_size;
+        reduction_idx /= output_dim_size;
+        output_offset += index_in_dim * params.output_strides[dim];
+      }
+    }
+
+    if (p != 0 && p != 1 && p != INFINITY && p != -INFINITY) {
+      output_val = static_cast<TA>(::precise::pow(output_val, 1 / p));
+    }
+    output[output_offset] = static_cast<TO>(output_val);
+  }
+}
+
+#define REGISTER_NORM(TI, TO)                               \
+  template [[host_name("norm_" #TI "_" #TO)]]               \
+  kernel void norm<TI, TO>(                                 \
+      constant TI * input [[buffer(0)]],                    \
+      device TO * output [[buffer(1)]],                     \
+      constant NormParams<> & params [[buffer(2)]],         \
+      uint tid [[thread_position_in_threadgroup]],          \
+      uint tptg [[threads_per_threadgroup]],                \
+      uint tgid [[threadgroup_position_in_grid]],           \
+      uint simd_lane_id [[thread_index_in_simdgroup]],      \
+      uint simdgroup_id [[simdgroup_index_in_threadgroup]], \
+      uint simdgroup_size [[threads_per_simdgroup]]);
+
+REGISTER_NORM(float, float);
+REGISTER_NORM(half, half);
+REGISTER_NORM(bfloat, bfloat);
+REGISTER_NORM(float2, float);
+REGISTER_NORM(half2, half);
diff --git a/aten/src/ATen/native/mps/kernels/SamplingHelpers.h b/aten/src/ATen/native/mps/kernels/SamplingHelpers.h
new file mode 100644
index 0000000000000..2fd06dfc26247
--- /dev/null
+++ b/aten/src/ATen/native/mps/kernels/SamplingHelpers.h
@@ -0,0 +1,42 @@
+#pragma once
+
+// Shared cubic interpolation helpers used by both GridSampler and UpSample
+// kernels. Based on the Catmull-Rom spline with A=-0.75.
+// See
+// https://en.wikipedia.org/wiki/Bicubic_interpolation#Bicubic_convolution_algorithm
+
+template <typename accscalar_t>
+accscalar_t cubic_convolution1(accscalar_t x, accscalar_t A) {
+  return ((A + 2) * x - (A + 3)) * x * x + 1;
+}
+
+template <typename accscalar_t>
+accscalar_t cubic_convolution2(accscalar_t x, accscalar_t A) {
+  return ((A * x - 5 * A) * x + 8 * A) * x - 4 * A;
+}
+
+template <typename accscalar_t>
+void get_cubic_coefficients(accscalar_t coeffs[4], accscalar_t t) {
+  accscalar_t A = -0.75;
+
+  accscalar_t x1 = t;
+  coeffs[0] = cubic_convolution2<accscalar_t>(x1 + 1.0, A);
+  coeffs[1] = cubic_convolution1<accscalar_t>(x1, A);
+
+  accscalar_t x2 = 1.0 - t;
+  coeffs[2] = cubic_convolution1<accscalar_t>(x2, A);
+  coeffs[3] = cubic_convolution2<accscalar_t>(x2 + 1.0, A);
+}
+
+template <typename scalar_t, typename accscalar_t>
+accscalar_t cubic_interp1d(
+    scalar_t x0,
+    scalar_t x1,
+    scalar_t x2,
+    scalar_t x3,
+    accscalar_t t) {
+  accscalar_t coeffs[4];
+  get_cubic_coefficients<accscalar_t>(coeffs, t);
+
+  return x0 * coeffs[0] + x1 * coeffs[1] + x2 * coeffs[2] + x3 * coeffs[3];
+}
diff --git a/aten/src/ATen/native/mps/kernels/ScanKernel.metal b/aten/src/ATen/native/mps/kernels/ScanKernel.metal
index de493af7aaa05..1fde656f46e69 100644
--- a/aten/src/ATen/native/mps/kernels/ScanKernel.metal
+++ b/aten/src/ATen/native/mps/kernels/ScanKernel.metal
@@ -3,26 +3,16 @@
 using namespace metal;
 
 #include <c10/metal/common.h>
+#include <c10/metal/special_math.h>
 #include <c10/metal/utils.h>
 
-using c10::metal::accum_t;
+using namespace c10::metal;
 
 struct LogAddExp {
   template <typename T>
   T operator()(T x, T y) {
-    // Reference:
-    // https://www.tensorflow.org/api_docs/python/tf/math/cumulative_logsumexp
-    T min_val = c10::metal::min(x, y);
-    T max_val = c10::metal::max(x, y);
-
-    if (min_val != max_val || metal::isfinite(min_val)) {
-      // nan will be propagated here
-      return c10::metal::log1p(metal::exp(min_val - max_val)) + max_val;
-    } else {
-      // special case to correctly handle infinite cases
-      return x;
-    }
-  };
+    return c10::metal::logaddexp(x, y);
+  }
 };
 
 C10_METAL_CONSTEXPR auto simd_size = c10::metal::simdgroup_size;
@@ -123,6 +113,29 @@ struct LogCumSumExpOp {
   }
 };
 
+template <typename T, typename acc_t = accum_t<T>>
+struct CumProdOp {
+  static constexpr constant acc_t init =
+      is_complex_v<acc_t> ? acc_t(1, 0) : static_cast<acc_t>(1);
+
+  acc_t operator()(acc_t a, acc_t b) {
+    return c10::metal::mul(a, b);
+  }
+
+  acc_t simd_scan(acc_t x) {
+    for (int i = 1; i <= 16; i *= 2) {
+      acc_t other = simd_shuffle_and_fill_up(x, init, i);
+      x = this->operator()(x, other);
+    }
+    return x;
+  }
+
+  acc_t simd_exclusive_scan(acc_t x) {
+    x = simd_scan(x);
+    return simd_shuffle_and_fill_up(x, init, 1);
+  }
+};
+
 // Pair structure to hold value and index for cummin/cummax operations
 template <typename T, typename acc_t = accum_t<T>>
 struct ValueIndexPair {
@@ -765,6 +778,11 @@ kernel void scan_with_indices_outer_dim(
 REGISTER_SCAN_OP(logcumsumexp, LogCumSumExpOp, float, 4);
 REGISTER_SCAN_OP(logcumsumexp, LogCumSumExpOp, half, 4);
 REGISTER_SCAN_OP(logcumsumexp, LogCumSumExpOp, bfloat, 4);
+REGISTER_SCAN_OP(logcumsumexp, LogCumSumExpOp, float2, 2);
+REGISTER_SCAN_OP(logcumsumexp, LogCumSumExpOp, half2, 4);
+
+REGISTER_SCAN_OP(cumprod, CumProdOp, float2, 2);
+REGISTER_SCAN_OP(cumprod, CumProdOp, half2, 4);
 
 // Scan with indices operations for cummin/cummax
 REGISTER_SCAN_WITH_INDICES_OP(cummin, CumMinOp, float, 4);
diff --git a/aten/src/ATen/native/mps/kernels/UnaryKernel.metal b/aten/src/ATen/native/mps/kernels/UnaryKernel.metal
index 66dad7440aa6f..81d333337ea55 100644
--- a/aten/src/ATen/native/mps/kernels/UnaryKernel.metal
+++ b/aten/src/ATen/native/mps/kernels/UnaryKernel.metal
@@ -91,7 +91,8 @@ struct abs_functor {
   }
   template <typename T, enable_if_t<is_complex_v<T>, bool> = true>
   inline T operator()(const T x) {
-    return T(::precise::sqrt(dot(x, x)), 0);
+    const auto abs_2 = ::precise::abs(float2(x));
+    return T(c10::metal::hypot(abs_2.x, abs_2.y));
   }
 };
 
@@ -420,6 +421,66 @@ struct log2_functor {
   }
 };
 
+struct lgamma_functor {
+  template <typename T>
+  inline enable_if_t<is_scalar_floating_point_v<T>, T> operator()(const T x) {
+    return static_cast<T>(c10::metal::log_gamma(static_cast<float>(x)));
+  }
+  template <typename T>
+  inline enable_if_t<is_scalar_integral_v<T>, float> operator()(const T x) {
+    return c10::metal::log_gamma(static_cast<float>(x));
+  }
+  inline float operator()(const bool x) {
+    return x ? 0 : INFINITY;
+  }
+};
+
+struct digamma_functor {
+  template <typename T>
+  inline enable_if_t<is_scalar_floating_point_v<T>, T> operator()(const T x) {
+    return static_cast<T>(c10::metal::digamma(static_cast<float>(x)));
+  }
+  template <typename T>
+  inline enable_if_t<is_scalar_integral_v<T>, float> operator()(const T x) {
+    return c10::metal::digamma(static_cast<float>(x));
+  }
+  inline float operator()(const bool x) {
+    return c10::metal::digamma(static_cast<float>(x));
+  }
+};
+
+struct trigamma_functor {
+  template <typename T>
+  inline enable_if_t<is_scalar_floating_point_v<T>, T> operator()(const T x) {
+    return static_cast<T>(c10::metal::trigamma(static_cast<float>(x)));
+  }
+  template <typename T>
+  inline enable_if_t<is_scalar_integral_v<T>, float> operator()(const T x) {
+    return c10::metal::trigamma(static_cast<float>(x));
+  }
+  inline float operator()(const bool x) {
+    return c10::metal::trigamma(static_cast<float>(x));
+  }
+};
+
+struct polygamma_functor {
+  template <typename T>
+  inline enable_if_t<is_scalar_floating_point_v<T>, T> operator()(
+      const T x,
+      const int order) {
+    return static_cast<T>(c10::metal::polygamma(order, static_cast<float>(x)));
+  }
+  template <typename T>
+  inline enable_if_t<is_scalar_integral_v<T>, float> operator()(
+      const T x,
+      const int order) {
+    return c10::metal::polygamma(order, static_cast<float>(x));
+  }
+  inline float operator()(const bool x, const int order) {
+    return c10::metal::polygamma(order, static_cast<float>(x));
+  }
+};
+
 struct exp2_functor {
   template <typename T>
   inline enable_if_t<is_scalar_floating_point_v<T>, T> operator()(const T x) {
@@ -614,6 +675,9 @@ REGISTER_UNARY_OP(abs, half, half);
   REGISTER_UNARY_OP(log10, DTYPE1, DTYPE0);        \
   REGISTER_UNARY_OP(log1p, DTYPE1, DTYPE0);        \
   REGISTER_UNARY_OP(log2, DTYPE1, DTYPE0);         \
+  REGISTER_UNARY_OP(lgamma, DTYPE1, DTYPE0);       \
+  REGISTER_UNARY_OP(digamma, DTYPE1, DTYPE0);      \
+  REGISTER_UNARY_OP(trigamma, DTYPE1, DTYPE0);     \
   REGISTER_UNARY_OP(sinc, DTYPE1, DTYPE0);         \
   REGISTER_UNARY_OP(sqrt, DTYPE1, DTYPE0);         \
   REGISTER_UNARY_OP(reciprocal, DTYPE1, DTYPE0);   \
@@ -685,3 +749,13 @@ REGISTER_UNARY_ALPHA_OP(pow_scalar, char, int, char);
 REGISTER_UNARY_ALPHA_OP(pow_scalar, short, int, short);
 REGISTER_UNARY_ALPHA_OP(pow_scalar, int, int, int);
 REGISTER_UNARY_ALPHA_OP(pow_scalar, long, int, long);
+
+REGISTER_UNARY_ALPHA_OP(polygamma, bfloat, int, bfloat);
+REGISTER_UNARY_ALPHA_OP(polygamma, half, int, half);
+REGISTER_UNARY_ALPHA_OP(polygamma, float, int, float);
+REGISTER_UNARY_ALPHA_OP(polygamma, bool, int, float);
+REGISTER_UNARY_ALPHA_OP(polygamma, uchar, int, float);
+REGISTER_UNARY_ALPHA_OP(polygamma, char, int, float);
+REGISTER_UNARY_ALPHA_OP(polygamma, short, int, float);
+REGISTER_UNARY_ALPHA_OP(polygamma, int, int, float);
+REGISTER_UNARY_ALPHA_OP(polygamma, long, int, float);
diff --git a/aten/src/ATen/native/mps/kernels/UnfoldBackward.metal b/aten/src/ATen/native/mps/kernels/UnfoldBackward.metal
index 8369258a30a6a..d62c4159e5c4d 100644
--- a/aten/src/ATen/native/mps/kernels/UnfoldBackward.metal
+++ b/aten/src/ATen/native/mps/kernels/UnfoldBackward.metal
@@ -42,7 +42,7 @@ kernel void unfold_backward(
   // Shift grad_in to start of unfold windows
   pos[dim] = 0;
   grad_in += offset_from_coord(pos, input_strides, ndim);
-  float rc = 0;
+  opmath_t<T> rc = 0;
   const auto in_dim_stride = input_strides[dim];
   const auto in_last_dim_stride = input_strides[ndim];
   for (auto in_dim_idx = left_fold_idx; in_dim_idx <= right_fold_idx;
@@ -71,3 +71,4 @@ kernel void unfold_backward(
 INSTANTIATE_UNFOLD_BACKWARD(float);
 INSTANTIATE_UNFOLD_BACKWARD(half);
 INSTANTIATE_UNFOLD_BACKWARD(bfloat);
+INSTANTIATE_UNFOLD_BACKWARD(float2);
diff --git a/aten/src/ATen/native/mps/kernels/UpSample.metal b/aten/src/ATen/native/mps/kernels/UpSample.metal
index fa9b5a1bb107d..cd4be4d0fc5bc 100644
--- a/aten/src/ATen/native/mps/kernels/UpSample.metal
+++ b/aten/src/ATen/native/mps/kernels/UpSample.metal
@@ -1,3 +1,4 @@
+#include <ATen/native/mps/kernels/SamplingHelpers.h>
 #include <ATen/native/mps/kernels/UpSample.h>
 #include <c10/metal/atomic.h>
 #include <metal_stdlib>
@@ -5,45 +6,6 @@
 using namespace metal;
 using namespace c10::metal;
 
-// Based on
-// https://en.wikipedia.org/wiki/Bicubic_interpolation#Bicubic_convolution_algorithm
-template <typename accscalar_t>
-accscalar_t cubic_convolution1(accscalar_t x, accscalar_t A) {
-  return ((A + 2) * x - (A + 3)) * x * x + 1;
-}
-
-template <typename accscalar_t>
-accscalar_t cubic_convolution2(accscalar_t x, accscalar_t A) {
-  return ((A * x - 5 * A) * x + 8 * A) * x - 4 * A;
-}
-
-template <typename accscalar_t>
-void get_cubic_upsampling_coefficients(accscalar_t coeffs[4], accscalar_t t) {
-  accscalar_t A = -0.75;
-
-  accscalar_t x1 = t;
-  coeffs[0] = cubic_convolution2<accscalar_t>(x1 + 1.0, A);
-  coeffs[1] = cubic_convolution1<accscalar_t>(x1, A);
-
-  // opposite coefficients
-  accscalar_t x2 = 1.0 - t;
-  coeffs[2] = cubic_convolution1<accscalar_t>(x2, A);
-  coeffs[3] = cubic_convolution2<accscalar_t>(x2 + 1.0, A);
-}
-
-template <typename scalar_t, typename accscalar_t>
-accscalar_t cubic_interp1d(
-    scalar_t x0,
-    scalar_t x1,
-    scalar_t x2,
-    scalar_t x3,
-    accscalar_t t) {
-  accscalar_t coeffs[4];
-  get_cubic_upsampling_coefficients<accscalar_t>(coeffs, t);
-
-  return x0 * coeffs[0] + x1 * coeffs[1] + x2 * coeffs[2] + x3 * coeffs[3];
-}
-
 template <typename accscalar_t>
 accscalar_t area_pixel_compute_source_index(
     accscalar_t scale,
@@ -720,8 +682,8 @@ kernel void upsample_bicubic2d_backward(
   float x_coeffs[4];
   float y_coeffs[4];
 
-  get_cubic_upsampling_coefficients(x_coeffs, t_x);
-  get_cubic_upsampling_coefficients(y_coeffs, t_y);
+  get_cubic_coefficients(x_coeffs, t_x);
+  get_cubic_coefficients(y_coeffs, t_y);
 
   for (int n = 0; n < output_sizes.x; n++) {
     for (int c = 0; c < output_sizes.y; ++c) {
diff --git a/aten/src/ATen/native/mps/operations/Activation.mm b/aten/src/ATen/native/mps/operations/Activation.mm
index 802c648c888d5..205b6bb3872c2 100644
--- a/aten/src/ATen/native/mps/operations/Activation.mm
+++ b/aten/src/ATen/native/mps/operations/Activation.mm
@@ -22,8 +22,6 @@
 #include <ATen/ops/mish_native.h>
 #include <ATen/ops/relu_native.h>
 #include <ATen/ops/sigmoid_backward_native.h>
-#include <ATen/ops/silu_backward_native.h>
-#include <ATen/ops/silu_native.h>
 #include <ATen/ops/softplus_backward_native.h>
 #include <ATen/ops/softplus_native.h>
 #include <ATen/ops/tanh_backward_native.h>
@@ -35,86 +33,6 @@
 
 namespace at::native {
 
-Tensor relu_mps(const Tensor& self) {
-  using namespace mps;
-  using CachedGraph = MPSUnaryCachedGraph;
-
-  bool executeGatherOp =
-      !(self.is_contiguous(MemoryFormat::Contiguous) || self.is_contiguous(MemoryFormat::ChannelsLast) ||
-        self.is_contiguous(MemoryFormat::ChannelsLast3d));
-  Tensor output = at::empty_like(self, executeGatherOp ? MemoryFormat::Contiguous : MemoryFormat::Preserve);
-
-  if (output.numel() == 0) {
-    return output;
-  }
-
-  MPSStream* stream = getCurrentMPSStream();
-  @autoreleasepool {
-    std::string key = "relu" + getTensorsStringKey({self});
-    auto cachedGraph = LookUpOrCreateCachedGraph<CachedGraph>(key, [&](auto mpsGraph, auto newCachedGraph) {
-      MPSGraphTensor* inputTensor = mpsGraphRankedPlaceHolder(mpsGraph, self);
-      // passing selector of reLUWithTensor on the mpsGraph object
-      MPSGraphTensor* outputTensor = [mpsGraph reLUWithTensor:inputTensor name:nil];
-
-      newCachedGraph->inputTensor_ = inputTensor;
-      newCachedGraph->outputTensor_ = outputTensor;
-    });
-
-    Placeholder selfPlaceholder = Placeholder(cachedGraph->inputTensor_, self, nil, executeGatherOp);
-    Placeholder outputPlaceholder = Placeholder(cachedGraph->outputTensor_, output, nil, false);
-
-    // Create dictionary of inputs and outputs
-    auto feeds = dictionaryFromPlaceholders(selfPlaceholder);
-    runMPSGraph(stream, cachedGraph->graph(), feeds, outputPlaceholder);
-  }
-
-  return output;
-}
-
-Tensor& relu_mps_(Tensor& self) {
-  using namespace mps;
-  using CachedGraph = MPSUnaryCachedGraph;
-
-  if (self.numel() == 0) {
-    return self;
-  }
-  // Inplace relu
-  Tensor& output = self;
-  bool executeGatherOp =
-      !(self.is_contiguous(MemoryFormat::Contiguous) || self.is_contiguous(MemoryFormat::ChannelsLast) ||
-        self.is_contiguous(MemoryFormat::ChannelsLast3d));
-  Tensor out;
-  if (executeGatherOp) {
-    out = at::empty_like(self, MemoryFormat::Contiguous);
-  }
-
-  MPSStream* stream = getCurrentMPSStream();
-
-  @autoreleasepool {
-    std::string key = "relu_" + getTensorsStringKey({self});
-    auto cachedGraph = LookUpOrCreateCachedGraph<CachedGraph>(key, [&](auto mpsGraph, auto newCachedGraph) {
-      MPSGraphTensor* inputTensor = mpsGraphRankedPlaceHolder(mpsGraph, self);
-      // passing selector of reLUWithTensor on the mpsGraph object
-      MPSGraphTensor* outputTensor = [mpsGraph reLUWithTensor:inputTensor name:nil];
-
-      newCachedGraph->inputTensor_ = inputTensor;
-      newCachedGraph->outputTensor_ = outputTensor;
-    });
-
-    Placeholder selfPlaceholder = Placeholder(cachedGraph->inputTensor_, self, nil, executeGatherOp);
-    Placeholder outputPlaceholder = Placeholder(cachedGraph->outputTensor_, executeGatherOp ? out : output, nil, false);
-
-    // Create dictionary of inputs and outputs
-    auto feeds = dictionaryFromPlaceholders(selfPlaceholder);
-    runMPSGraph(stream, cachedGraph->graph(), feeds, outputPlaceholder);
-    if (executeGatherOp) {
-      output.copy_(out);
-    }
-  }
-
-  return output;
-}
-
 TORCH_IMPL_FUNC(log_softmax_mps_out)
 (const Tensor& self, const int64_t dim, const bool half_to_float, const Tensor& out) {
   TORCH_CHECK_NOT_IMPLEMENTED(self.scalar_type() != kLong, "MPS doesn't know how to do exponent_i64");
@@ -1216,112 +1134,6 @@ Tensor prelu_mps(const Tensor& self, const Tensor& weight_) {
   return std::tuple<Tensor, Tensor>{grad_input, weight_grad};
 }
 
-TORCH_IMPL_FUNC(silu_out_mps)(const Tensor& self, const Tensor& result) {
-  using namespace mps;
-  using CachedGraph = MPSUnaryCachedGraph;
-
-  TORCH_CHECK(self.is_mps());
-  TORCH_CHECK_NOT_IMPLEMENTED(self.scalar_type() != kLong, "MPS doesn't know how to do exponent_i64");
-
-  // Empty output
-  if (result.numel() == 0)
-    return;
-
-  MPSStream* stream = getCurrentMPSStream();
-
-  bool executeGatherOp =
-      !(self.is_contiguous(MemoryFormat::Contiguous) || self.is_contiguous(MemoryFormat::ChannelsLast) ||
-        self.is_contiguous(MemoryFormat::ChannelsLast3d));
-  Tensor result_ = at::empty_like(self, executeGatherOp ? MemoryFormat::Contiguous : MemoryFormat::Preserve);
-
-  @autoreleasepool {
-    std::string key = "silu_out_mps:" + getTensorsStringKey({self});
-
-    auto cachedGraph = LookUpOrCreateCachedGraph<CachedGraph>(key, [&](auto mpsGraph, auto newCachedGraph) {
-      MPSGraphTensor* inputTensor = mpsGraphRankedPlaceHolder(mpsGraph, self);
-
-      MPSGraphTensor* unitTensor = [mpsGraph constantWithScalar:1.0 shape:@[ @1 ] dataType:getMPSDataType(self)];
-      MPSGraphTensor* negativeInput = [mpsGraph negativeWithTensor:inputTensor name:nil];
-      MPSGraphTensor* expNegativeTensor = [mpsGraph exponentWithTensor:negativeInput name:nil];
-      MPSGraphTensor* expPlusOneTensor = [mpsGraph additionWithPrimaryTensor:expNegativeTensor
-                                                             secondaryTensor:unitTensor
-                                                                        name:nil];
-      MPSGraphTensor* outputTensor = [mpsGraph divisionWithPrimaryTensor:inputTensor
-                                                         secondaryTensor:expPlusOneTensor
-                                                                    name:nil];
-
-      newCachedGraph->inputTensor_ = inputTensor;
-      newCachedGraph->outputTensor_ = outputTensor;
-    });
-
-    Placeholder selfPlaceholder = Placeholder(cachedGraph->inputTensor_, self, nil, executeGatherOp);
-    Placeholder outputPlaceholder =
-        Placeholder(cachedGraph->outputTensor_, executeGatherOp ? result_ : result, nil, false);
-
-    auto feeds = dictionaryFromPlaceholders(selfPlaceholder);
-    runMPSGraph(stream, cachedGraph->graph(), feeds, outputPlaceholder);
-  }
-  if (executeGatherOp) {
-    result.copy_(result_);
-  }
-}
-
-TORCH_IMPL_FUNC(silu_backward_out_mps)
-(const Tensor& grad_output, const Tensor& self, const Tensor& grad_input) {
-  using namespace mps;
-  using CachedGraph = MPSUnaryGradCachedGraph;
-  TORCH_CHECK(grad_output.is_mps());
-
-  // Empty output
-  if (grad_input.numel() == 0)
-    return;
-
-  MPSStream* stream = getCurrentMPSStream();
-
-  @autoreleasepool {
-    std::string key = "silu_out_backward_mps:" + getTensorsStringKey({grad_output});
-
-    auto cachedGraph = LookUpOrCreateCachedGraph<CachedGraph>(key, [&](auto mpsGraph, auto newCachedGraph) {
-      MPSGraphTensor* inputTensor = mpsGraphRankedPlaceHolder(mpsGraph, self);
-      MPSGraphTensor* gradOutputTensor = mpsGraphRankedPlaceHolder(mpsGraph, grad_output);
-
-      MPSGraphTensor* unitTensor = [mpsGraph constantWithScalar:1.0 shape:@[ @1 ] dataType:getMPSDataType(grad_output)];
-      MPSGraphTensor* negativeInput = [mpsGraph negativeWithTensor:inputTensor name:nil];
-      MPSGraphTensor* expNegativeTensor = [mpsGraph exponentWithTensor:negativeInput name:nil];
-      MPSGraphTensor* expPlusOneTensor = [mpsGraph additionWithPrimaryTensor:expNegativeTensor
-                                                             secondaryTensor:unitTensor
-                                                                        name:nil];
-      MPSGraphTensor* sigmoidTensor = [mpsGraph reciprocalWithTensor:expPlusOneTensor name:nil];
-      MPSGraphTensor* oneMinusSigmoid = [mpsGraph subtractionWithPrimaryTensor:unitTensor
-                                                               secondaryTensor:sigmoidTensor
-                                                                          name:nil];
-      MPSGraphTensor* inputTimesDiff = [mpsGraph multiplicationWithPrimaryTensor:inputTensor
-                                                                 secondaryTensor:oneMinusSigmoid
-                                                                            name:nil];
-      MPSGraphTensor* onePlusTensor = [mpsGraph additionWithPrimaryTensor:unitTensor
-                                                          secondaryTensor:inputTimesDiff
-                                                                     name:nil];
-      MPSGraphTensor* gradTensor = [mpsGraph multiplicationWithPrimaryTensor:sigmoidTensor
-                                                             secondaryTensor:onePlusTensor
-                                                                        name:nil];
-      MPSGraphTensor* gradInputTensor = [mpsGraph multiplicationWithPrimaryTensor:gradTensor
-                                                                  secondaryTensor:gradOutputTensor
-                                                                             name:nil];
-
-      newCachedGraph->gradOutputTensor_ = gradOutputTensor;
-      newCachedGraph->inputTensor_ = inputTensor;
-      newCachedGraph->gradInputTensor_ = gradInputTensor;
-    });
-
-    Placeholder selfPlaceholder = Placeholder(cachedGraph->inputTensor_, self);
-    Placeholder gradOutputPlaceholder = Placeholder(cachedGraph->gradOutputTensor_, grad_output);
-    Placeholder gradInputPlaceholder = Placeholder(cachedGraph->gradInputTensor_, grad_input);
-
-    auto feeds = dictionaryFromPlaceholders(selfPlaceholder, gradOutputPlaceholder);
-    runMPSGraph(stream, cachedGraph->graph(), feeds, gradInputPlaceholder);
-  }
-}
-
 // -------------------------------------------------
 // Hardtanh backward
 
diff --git a/aten/src/ATen/native/mps/operations/ActivationKernel.mm b/aten/src/ATen/native/mps/operations/ActivationKernel.mm
index f6d3ad986ade0..2138b0ddb6a96 100644
--- a/aten/src/ATen/native/mps/operations/ActivationKernel.mm
+++ b/aten/src/ATen/native/mps/operations/ActivationKernel.mm
@@ -4,6 +4,18 @@
 #include <ATen/mps/MPSProfiler.h>
 #include <ATen/native/Activation.h>
 #include <ATen/native/mps/OperationUtils.h>
+
+#ifndef AT_PER_OPERATOR_HEADERS
+#include <ATen/NativeFunctions.h>
+#else
+#include <ATen/ops/add.h>
+#include <ATen/ops/mul.h>
+#include <ATen/ops/mul_native.h>
+#include <ATen/ops/relu_native.h>
+#include <ATen/ops/rsub.h>
+#include <ATen/ops/sigmoid.h>
+#include <ATen/ops/sigmoid_native.h>
+#endif
 #include <ATen/native/mps/kernels/Activation.h>
 #include <fmt/format.h>
 
@@ -15,6 +27,25 @@
 #include <ATen/native/mps/ActivationKernel_metallib.h>
 #endif
 
+Tensor relu_mps(const Tensor& self) {
+  TORCH_CHECK(!self.is_complex(), "relu is not supported for complex types");
+  auto output = at::empty_like(self);
+  if (output.numel() == 0)
+    return output;
+  auto iter = at::TensorIteratorConfig().add_output(output).add_input(self).build();
+  lib.exec_unary_kernel(iter, "relu", /*alpha=*/std::nullopt, /*scalar_arg_type=*/std::nullopt, /*supports_vec4=*/true);
+  return output;
+}
+
+Tensor& relu_mps_(Tensor& self) {
+  TORCH_CHECK(!self.is_complex(), "relu is not supported for complex types");
+  if (self.numel() == 0)
+    return self;
+  auto iter = at::TensorIteratorConfig().add_output(self).add_input(self).set_check_mem_overlap(false).build();
+  lib.exec_unary_kernel(iter, "relu", /*alpha=*/std::nullopt, /*scalar_arg_type=*/std::nullopt, /*supports_vec4=*/true);
+  return self;
+}
+
 static void hardshrink_kernel(TensorIteratorBase& iter, const Scalar& lambda = 0.5) {
   lib.exec_unary_kernel(iter, "hardshrink", lambda);
 }
@@ -67,6 +98,30 @@ static void elu_backward_kernel(TensorIteratorBase& iter,
   });
 }
 
+static void silu_kernel(TensorIteratorBase& iter) {
+  if (isComplexType(iter.common_dtype())) {
+    auto out = iter.output(0);
+    auto self = iter.input(0);
+    at::mul_out(out, self, at::sigmoid(self));
+    return;
+  }
+  lib.exec_unary_kernel(iter, "silu", /*alpha=*/std::nullopt, /*scalar_arg_type=*/std::nullopt, /*supports_vec4=*/true);
+}
+
+static void silu_backward_kernel(TensorIteratorBase& iter) {
+  if (isComplexType(iter.common_dtype())) {
+    auto grad_input = iter.output(0);
+    auto grad_output = iter.input(0);
+    auto self = iter.input(1);
+    auto sig = at::sigmoid(self);
+    auto one_minus_sig = at::rsub(sig, 1);
+    auto inner = at::add(at::mul(self, one_minus_sig), 1);
+    grad_input.copy_(at::mul(grad_output, at::mul(sig, inner)));
+    return;
+  }
+  lib.exec_binary_kernel(iter, "silu_backward");
+}
+
 static void leaky_relu_kernel(TensorIteratorBase& iter, const Scalar& negative_slope) {
   lib.exec_unary_kernel(iter, "leaky_relu", negative_slope);
 }
@@ -86,5 +141,7 @@ static void leaky_relu_backward_kernel(TensorIteratorBase& iter, const Scalar& n
 REGISTER_DISPATCH(elu_backward_stub, elu_backward_kernel);
 REGISTER_DISPATCH(leaky_relu_stub, leaky_relu_kernel);
 REGISTER_DISPATCH(leaky_relu_backward_stub, leaky_relu_backward_kernel);
+REGISTER_DISPATCH(silu_stub, silu_kernel);
+REGISTER_DISPATCH(silu_backward_stub, silu_backward_kernel);
 
 } // namespace at::native
diff --git a/aten/src/ATen/native/mps/operations/Amp.mm b/aten/src/ATen/native/mps/operations/Amp.mm
index e410d434ec7ad..2b1e11575cd50 100644
--- a/aten/src/ATen/native/mps/operations/Amp.mm
+++ b/aten/src/ATen/native/mps/operations/Amp.mm
@@ -29,7 +29,6 @@ static void _amp_non_finite_check_and_unscale_mps_single_impl(const Tensor& scal
   }
   TORCH_CHECK(scaled_grad.is_mps(), "Tensor is not on the MPS device.");
   TORCH_CHECK(scaled_grad.numel() <= std::numeric_limits<uint32_t>::max(), "scaled_grad is too large");
-  float inv_scale_val = inv_scale.item<float>();
   auto stream = getCurrentMPSStream();
   auto device = MPSDevice::getInstance()->device();
   auto ampPipelineState =
@@ -43,7 +42,7 @@ static void _amp_non_finite_check_and_unscale_mps_single_impl(const Tensor& scal
   dispatch_sync_with_rethrow(stream->queue(), ^() {
     auto computeEncoder = stream->commandEncoder();
     [computeEncoder setComputePipelineState:ampPipelineState];
-    mtl_setArgs(computeEncoder, scaled_grad, found_inf, inv_scale_val);
+    mtl_setArgs(computeEncoder, scaled_grad, found_inf, inv_scale);
     [computeEncoder dispatchThreads:gridSize threadsPerThreadgroup:threadGroupSize];
   });
 }
diff --git a/aten/src/ATen/native/mps/operations/Attention.mm b/aten/src/ATen/native/mps/operations/Attention.mm
index ce57174177885..6cce569d5e14d 100644
--- a/aten/src/ATen/native/mps/operations/Attention.mm
+++ b/aten/src/ATen/native/mps/operations/Attention.mm
@@ -64,9 +64,9 @@
   int64_t batchSize = query.size(0);
   int64_t num_head = query.size(1);
   int64_t qSize = query.size(2);
-  int64_t headSize = query.size(3);
+  int64_t valueHeadSize = value.size(3);
   int64_t maxSeqLength = key.size(2);
-  auto out = at::empty({batchSize, num_head, qSize, headSize}, query.options());
+  auto out = at::empty({batchSize, num_head, qSize, valueHeadSize}, query.options());
   auto attn = at::empty({batchSize, num_head, qSize, maxSeqLength}, query.options());
   auto scale_factor = sdp::calculate_scale(query, scale).expect_float();
   @autoreleasepool {
@@ -149,12 +149,25 @@
     runMPSGraph(getCurrentMPSStream(), cachedGraph->graph(), feeds, outs);
   }
 
-  auto final_out = unsqueezed ? out.view_as(orig_query) : out;
-  auto final_attn = unsqueezed ? (orig_query.dim() == 3 ? attn.squeeze(0) : [&]{
-    std::vector<int64_t> shape(orig_query.sizes().begin(), orig_query.sizes().end() - 3);
-    shape.insert(shape.end(), {attn.size(1), attn.size(2), attn.size(3)});
-    return attn.view(shape);
-  }()) : attn;
+  auto final_out = out;
+  auto final_attn = attn;
+  if (unsqueezed) {
+    if (orig_query.dim() == 3) {
+      final_out = out.squeeze(0);
+      final_attn = attn.squeeze(0);
+    } else {
+      std::vector<int64_t> prefix_shape(orig_query.sizes().begin(), orig_query.sizes().end() - 3);
+
+      auto out_shape = prefix_shape;
+      auto attn_shape = prefix_shape;
+
+      out_shape.insert(out_shape.end(), {out.size(1), out.size(2), out.size(3)});
+      attn_shape.insert(attn_shape.end(), {attn.size(1), attn.size(2), attn.size(3)});
+
+      final_out = out.view(out_shape);
+      final_attn = attn.view(attn_shape);
+    }
+  }
 
   return {std::move(final_out), std::move(final_attn)};
 }
@@ -170,6 +183,8 @@
                                                        std::optional<double> scale,
                                                        const Tensor& orig_query,
                                                        bool unsqueezed) {
+  TORCH_CHECK(q_.size(3) == k_.size(3) && q_.size(3) == v_.size(3),
+              "sdpa_vector_fast_mps expects query, key, and value to have the same head dimension");
   const auto macOS15_0_plus = is_macos_13_or_newer(MacOSVersion::MACOS_VER_15_0_PLUS);
   using namespace mps;
   uint batchSize = q_.size(0);
@@ -179,10 +194,13 @@
   uint maxSeqLength = k_.size(2);
   uint N = k_.size(2);
   uint B = q_.size(0) * q_.size(1);
+  uint q_batch_stride = q_.stride(0);
   uint q_head_stride = q_.stride(1);
   uint q_seq_stride = q_.stride(2);
+  uint k_batch_stride = k_.stride(0);
   uint k_head_stride = k_.stride(1);
   uint k_seq_stride = k_.stride(2);
+  uint v_batch_stride = v_.stride(0);
   uint v_head_stride = v_.stride(1);
   uint v_seq_stride = v_.stride(2);
 
@@ -220,7 +238,8 @@
         mtl_setArgs<9>(
             computeEncoder, mask_.value(), std::array<uint32_t, 3>{kv_seq_stride, q_seq_stride, head_stride});
       }
-      mtl_setArgs<11>(computeEncoder, has_mask);
+      mtl_setArgs<11>(
+          computeEncoder, has_mask, std::array<uint32_t, 4>{q_batch_stride, k_batch_stride, v_batch_stride, num_head});
       [computeEncoder dispatchThreadgroups:grid_dims threadsPerThreadgroup:group_dims];
     }
   });
@@ -246,6 +265,8 @@
                                                         std::optional<double> scale,
                                                         const Tensor& orig_query,
                                                         bool unsqueezed) {
+  TORCH_CHECK(q_.size(3) == k_.size(3) && q_.size(3) == v_.size(3),
+              "sdpa_vector_2pass_mps expects query, key, and value to have the same head dimension");
   using namespace mps;
   uint batchSize = q_.size(0);
   uint num_heads = q_.size(1);
@@ -256,17 +277,20 @@
   uint B = batchSize * num_heads;
   uint gqa_factor = q_.size(1) / k_.size(1);
 
+  uint q_batch_stride = q_.stride(0);
   uint q_head_stride = q_.stride(1);
   uint q_seq_stride = q_.stride(2);
+  uint k_batch_stride = k_.stride(0);
   uint k_head_stride = k_.stride(1);
   uint k_seq_stride = k_.stride(2);
+  uint v_batch_stride = v_.stride(0);
   uint v_head_stride = v_.stride(1);
   uint v_seq_stride = v_.stride(2);
 
   auto out = at::empty({batchSize, num_heads, seq_len_q, headSize}, q_.options());
   auto intermediate = at::empty({batchSize, num_heads, seq_len_q, blocks, headSize}, q_.options());
-  auto sums = at::empty({batchSize, num_heads, seq_len_q, blocks}, q_.options());
-  auto maxs = at::empty({batchSize, num_heads, seq_len_q, blocks}, q_.options());
+  auto sums = at::empty({batchSize, num_heads, seq_len_q, blocks}, q_.options().dtype(kFloat));
+  auto maxs = at::empty({batchSize, num_heads, seq_len_q, blocks}, q_.options().dtype(kFloat));
 
   auto scale_factor = sdp::calculate_scale(orig_query, scale).expect_float();
   bool has_mask = mask_.has_value();
@@ -307,7 +331,8 @@
         uint head_stride = (nd >= 3 && mask.size(nd - 3) > 1) ? mask.stride(nd - 3) : 0;
         mtl_setArgs<11>(computeEncoder, mask, std::array<uint32_t, 3>{kv_seq_stride, q_seq_stride, head_stride});
       }
-      mtl_setArgs<13>(computeEncoder, has_mask);
+      mtl_setArgs<13>(
+          computeEncoder, has_mask, std::array<uint32_t, 4>{q_batch_stride, k_batch_stride, v_batch_stride, num_heads});
       [computeEncoder dispatchThreadgroups:grid_dims threadsPerThreadgroup:group_dims];
       // 2nd pass
       [computeEncoder setComputePipelineState:sdpa_vector_pass2PSO];
@@ -422,13 +447,42 @@
 }
 
 std::tuple<Tensor, Tensor> _scaled_dot_product_attention_math_mps(const Tensor& query,
-                                                                  const Tensor& key,
-                                                                  const Tensor& value,
+                                                                  const Tensor& key_,
+                                                                  const Tensor& value_,
                                                                   const std::optional<Tensor>& attn_mask,
                                                                   double dropout_p,
                                                                   bool is_causal,
                                                                   const std::optional<Tensor>& dropout_mask,
-                                                                  std::optional<double> scale) {
+                                                                  std::optional<double> scale,
+                                                                  bool enable_gqa) {
+  TORCH_CHECK_NOT_IMPLEMENTED(c10::isFloatingType(query.scalar_type()),
+                              "scaled_dot_product_attention for MPS does not support dtype ",
+                              query.scalar_type());
+  TORCH_CHECK_NOT_IMPLEMENTED(c10::isFloatingType(key_.scalar_type()),
+                              "scaled_dot_product_attention for MPS does not support dtype ",
+                              key_.scalar_type());
+  TORCH_CHECK_NOT_IMPLEMENTED(c10::isFloatingType(value_.scalar_type()),
+                              "scaled_dot_product_attention for MPS does not support dtype ",
+                              value_.scalar_type());
+  const auto any_nested = query.is_nested() || key_.is_nested() || value_.is_nested();
+  const auto all_contiguous =
+      query.is_contiguous_or_false() && key_.is_contiguous_or_false() && value_.is_contiguous_or_false();
+  auto key = key_;
+  auto value = value_;
+  if (enable_gqa) {
+    int64_t q_heads = query.size(-3);
+    int64_t k_heads = key_.size(-3);
+    int64_t repeat_factor = q_heads / k_heads;
+
+    if (repeat_factor > 1) {
+      TORCH_CHECK(q_heads % k_heads == 0,
+                  "For GQA, the query tensor's head dimension (" + std::to_string(q_heads) +
+                      ") must be divisible by the key tensor's head dimension (" + std::to_string(k_heads) + ").");
+      key = key_.repeat_interleave(repeat_factor, /*dim=*/-3);
+      value = value_.repeat_interleave(repeat_factor, /*dim=*/-3);
+    }
+  }
+
   auto query_tuple = ensure_4d(query);
   Tensor q_ = std::get<0>(query_tuple);
   bool unsqueezed = std::get<1>(query_tuple);
@@ -469,19 +523,12 @@
     return sdpa_general_mps(q_, k_, v_, mask_, dropout_p, is_causal, dropout_mask, scale, query, unsqueezed);
   }
 
-  // dispatch to the fast SDPA implementation
-  auto is_contiguous_or_head_seq_transposed = [](const Tensor& t) -> bool {
-    if (t.is_contiguous())
-      return true;
-    auto sizes = t.sizes();
-    auto strides = t.strides();
-    return (strides[3] == 1) && (strides[2] == sizes[3] * sizes[1]) && (strides[1] == sizes[3]) &&
-        (strides[0] == strides[2] * sizes[2]);
-  };
+  // Kernels load head-dim elements linearly, so stride(-1) == 1 is the only hard requirement
+  auto can_use_kernel_strides = [](const Tensor& t) -> bool { return t.stride(-1) == 1; };
 
-  Tensor q_contig = is_contiguous_or_head_seq_transposed(q_) ? q_ : q_.contiguous();
-  Tensor k_contig = k_.contiguous();
-  Tensor v_contig = v_.contiguous();
+  Tensor q_contig = can_use_kernel_strides(q_) ? q_ : q_.contiguous();
+  Tensor k_contig = can_use_kernel_strides(k_) ? k_ : k_.contiguous();
+  Tensor v_contig = can_use_kernel_strides(v_) ? v_ : v_.contiguous();
 
   // for short sequences, differentiate based on key sequence length
   if ((k_.size(2) >= 1024) || (k_.size(1) < q_.size(1) && k_.size(2) >= 4096)) {
diff --git a/aten/src/ATen/native/mps/operations/BinaryKernel.mm b/aten/src/ATen/native/mps/operations/BinaryKernel.mm
index 03ecee79449c0..9f78ab11dfab1 100644
--- a/aten/src/ATen/native/mps/operations/BinaryKernel.mm
+++ b/aten/src/ATen/native/mps/operations/BinaryKernel.mm
@@ -109,6 +109,10 @@ static void logaddexp2_mps_kernel(TensorIteratorBase& iter) {
   lib.exec_binary_kernel(iter, "logaddexp2");
 }
 
+static void xlogy_mps_kernel(TensorIteratorBase& iter) {
+  lib.exec_binary_kernel(iter, "xlogy");
+}
+
 static void xlog1py_mps_kernel(TensorIteratorBase& iter) {
   TORCH_CHECK_TYPE(isFloatingType(iter.common_dtype()), "xlog1py_mps not implemented for non-floating types");
   lib.exec_binary_kernel(iter, "xlog1py");
@@ -186,6 +190,68 @@ static void lerp_scalar_mps_kernel(at::TensorIteratorBase& iter, const Scalar& w
   lib.exec_binary_kernel(iter, "lerp_alpha", weight);
 }
 
+static void lerp_tensor_mps_kernel(at::TensorIteratorBase& iter) {
+  using namespace mps;
+  auto type_str = scalarToMetalTypeString(iter.common_dtype());
+  auto numel = static_cast<uint32_t>(iter.numel());
+  auto ndim = static_cast<uint32_t>(iter.ndim());
+
+  // simple elementwise kernel for dense tensors
+  if (iter.is_contiguous()) {
+    auto pso = lib.getPipelineStateForFunc("lerp_tensor_dense_" + type_str);
+    dispatch_sync_with_rethrow(getCurrentMPSStream()->queue(), ^() {
+      auto computeEncoder = getCurrentMPSStream()->commandEncoder();
+      [computeEncoder setComputePipelineState:pso];
+      bind_iter_tensors(computeEncoder, iter);
+      mtl_dispatch1DJob(computeEncoder, pso, numel);
+    });
+    return;
+  }
+
+  // Scalar weight broadcast path
+  if (ndim == 1 && iter.strides(3)[0] == 0) {
+    auto pso = lib.getPipelineStateForFunc("lerp_tensor_scalar_weight_" + type_str);
+    dispatch_sync_with_rethrow(getCurrentMPSStream()->queue(), ^() {
+      auto computeEncoder = getCurrentMPSStream()->commandEncoder();
+      [computeEncoder setComputePipelineState:pso];
+      bind_iter_tensors(computeEncoder, iter);
+      mtl_dispatch1DJob(computeEncoder, pso, numel);
+    });
+    return;
+  }
+
+  // 2D/3D: multi-dimensional dispatch, to avoid integer division for coordinates
+  if (ndim >= 2 && ndim <= 3) {
+    auto pso = lib.getPipelineStateForFunc(fmt::format("lerp_tensor_strided_{}d_{}", ndim, type_str));
+    dispatch_sync_with_rethrow(getCurrentMPSStream()->queue(), ^() {
+      auto computeEncoder = getCurrentMPSStream()->commandEncoder();
+      [computeEncoder setComputePipelineState:pso];
+      bind_iter_tensors(computeEncoder, iter);
+      mtl_setArgs<4>(computeEncoder, iter.strides(0), iter.strides(1), iter.strides(2), iter.strides(3));
+      auto sizes = iter.shape();
+      auto maxTg = [pso maxTotalThreadsPerThreadgroup];
+      auto tg_x = std::min(static_cast<NSUInteger>(sizes[0]), maxTg);
+      auto tg_y = std::min(static_cast<NSUInteger>(sizes[1]), maxTg / tg_x);
+      auto grid_z = ndim > 2 ? static_cast<NSUInteger>(sizes[2]) : 1;
+      auto tg_z = std::min(grid_z, std::max(maxTg / (tg_x * tg_y), (NSUInteger)1));
+      [computeEncoder dispatchThreads:MTLSizeMake(sizes[0], sizes[1], grid_z)
+                threadsPerThreadgroup:MTLSizeMake(tg_x, tg_y, tg_z)];
+    });
+    return;
+  }
+
+  // General strided fallback
+  auto pso = lib.getPipelineStateForFunc("lerp_tensor_strided_" + type_str);
+  dispatch_sync_with_rethrow(getCurrentMPSStream()->queue(), ^() {
+    auto computeEncoder = getCurrentMPSStream()->commandEncoder();
+    [computeEncoder setComputePipelineState:pso];
+    bind_iter_tensors(computeEncoder, iter);
+    mtl_setArgs<4>(
+        computeEncoder, iter.shape(), iter.strides(0), iter.strides(1), iter.strides(2), iter.strides(3), ndim);
+    mtl_dispatch1DJob(computeEncoder, pso, numel);
+  });
+}
+
 static void native_dropout_mask_and_scale_mps_kernel(at::TensorIteratorBase& iter, const Scalar& scale) {
   lib.exec_binary_kernel(iter, "native_dropout_mask_and_scale", scale);
 }
@@ -242,6 +308,7 @@ static void gcd_mps_kernel(TensorIteratorBase& iter) {
 REGISTER_DISPATCH(zeta_stub, &zeta_mps_kernel)
 REGISTER_DISPATCH(logaddexp_stub, &logaddexp_mps_kernel);
 REGISTER_DISPATCH(logaddexp2_stub, &logaddexp2_mps_kernel);
+REGISTER_DISPATCH(xlogy_stub, &xlogy_mps_kernel)
 REGISTER_DISPATCH(xlog1py_stub, &xlog1py_mps_kernel)
 REGISTER_DISPATCH(chebyshev_polynomial_t_stub, &chebyshev_polynomial_t_mps_kernel)
 REGISTER_DISPATCH(chebyshev_polynomial_u_stub, &chebyshev_polynomial_u_mps_kernel)
@@ -256,6 +323,7 @@ static void gcd_mps_kernel(TensorIteratorBase& iter) {
 REGISTER_DISPATCH(polar_stub, &polar_mps_kernel);
 REGISTER_DISPATCH(complex_stub, &complex_mps_kernel);
 REGISTER_DISPATCH(lerp_kernel_scalar_weight, &lerp_scalar_mps_kernel)
+REGISTER_DISPATCH(lerp_kernel_tensor_weight, &lerp_tensor_mps_kernel)
 REGISTER_DISPATCH(mul_stub, &mul_mps_kernel)
 REGISTER_DISPATCH(div_true_stub, &div_true_mps_kernel)
 REGISTER_DISPATCH(div_floor_stub, &div_floor_mps_kernel)
diff --git a/aten/src/ATen/native/mps/operations/BinaryOps.mm b/aten/src/ATen/native/mps/operations/BinaryOps.mm
index cfbb194916b83..e56dd92679507 100644
--- a/aten/src/ATen/native/mps/operations/BinaryOps.mm
+++ b/aten/src/ATen/native/mps/operations/BinaryOps.mm
@@ -20,15 +20,12 @@
 #include <ATen/ops/logical_or_native.h>
 #include <ATen/ops/logical_xor_native.h>
 #include <ATen/ops/lt_native.h>
-#include <ATen/ops/maximum_native.h>
-#include <ATen/ops/minimum_native.h>
 #include <ATen/ops/ne_native.h>
 #include <ATen/ops/pow.h>
 #include <ATen/ops/pow_native.h>
 #include <ATen/ops/result_type.h>
 #include <ATen/ops/sub_native.h>
 #include <ATen/ops/view_as_real.h>
-#include <ATen/ops/xlogy_native.h>
 #endif
 
 namespace at::native {
@@ -243,8 +240,6 @@ static void add_sub_lerp_template(const Tensor& self,
 CREATE_MPS_STRUCTURED_BOOLEAN_OP_FUNC(gt_tensor_out_mps, greaterThan, Tensor);
 
 // Arithmetic Binary Ops
-CREATE_MPS_STRUCTURED_BINARY_OP_FUNC(minimum_out_mps, minimumWithNaNPropagationAndIntFallback, Tensor);
-CREATE_MPS_STRUCTURED_BINARY_OP_FUNC(maximum_out_mps, maximumWithNaNPropagationAndIntFallback, Tensor);
 CREATE_MPS_STRUCTURED_BINARY_OP_FUNC(pow_tensor_tensor_out_mps, power, Tensor);
 CREATE_MPS_BINARY_COMPARISON_OP_FUNC(logical_and_out_mps, logicalAND, Tensor);
 CREATE_MPS_BINARY_COMPARISON_OP_FUNC(logical_or_out_mps, logicalOR, Tensor);
@@ -266,29 +261,4 @@ static void add_sub_lerp_template(const Tensor& self,
   }
 }
 
-TORCH_IMPL_FUNC(xlogy_out_mps)(const Tensor& self, const Tensor& other, const Tensor& output) {
-  mps::BinaryOpBlock xlogy_op_block = ^BinaryOpFn(cachedGraph, primaryCastTensor, secondaryCastTensor) {
-    MPSGraph* mpsGraph = cachedGraph->graph();
-    MPSGraphTensor* zeroTensor = [mpsGraph constantWithScalar:0.0 shape:@[ @1 ] dataType:primaryCastTensor.dataType];
-    MPSGraphTensor* yIsNaNPredicateTensor = [mpsGraph isNaNWithTensor:secondaryCastTensor name:nil];
-    MPSGraphTensor* logyTensor = [mpsGraph logarithmWithTensor:secondaryCastTensor name:nil];
-    MPSGraphTensor* xlogyTensor = [mpsGraph multiplicationWithPrimaryTensor:primaryCastTensor
-                                                            secondaryTensor:logyTensor
-                                                                       name:nil];
-    MPSGraphTensor* xEqualZeroPredicateTensor = [mpsGraph equalWithPrimaryTensor:primaryCastTensor
-                                                                 secondaryTensor:zeroTensor
-                                                                            name:nil];
-    MPSGraphTensor* outputTensor = [mpsGraph selectWithPredicateTensor:xEqualZeroPredicateTensor
-                                                   truePredicateTensor:zeroTensor
-                                                  falsePredicateTensor:xlogyTensor
-                                                                  name:nil];
-    outputTensor = [mpsGraph selectWithPredicateTensor:yIsNaNPredicateTensor
-                                   truePredicateTensor:secondaryCastTensor
-                                  falsePredicateTensor:outputTensor
-                                                  name:nil];
-    return outputTensor;
-  };
-  mps::binaryOpTensor(self, other, output, "xlogy_out_mps", xlogy_op_block);
-}
-
 } // namespace at::native
diff --git a/aten/src/ATen/native/mps/operations/ConstantOps.mm b/aten/src/ATen/native/mps/operations/ConstantOps.mm
index e36ac4dc45246..c0d3b795ded9f 100644
--- a/aten/src/ATen/native/mps/operations/ConstantOps.mm
+++ b/aten/src/ATen/native/mps/operations/ConstantOps.mm
@@ -1,120 +1,89 @@
 //  Copyright © 2022 Apple Inc.
 #define TORCH_ASSERT_ONLY_METHOD_OPERATORS
+#include <ATen/native/Fill.h>
 #include <ATen/native/mps/OperationUtils.h>
+#include <fmt/format.h>
 
 #ifndef AT_PER_OPERATOR_HEADERS
 #include <ATen/Functions.h>
 #include <ATen/NativeFunctions.h>
 #else
-#include <ATen/ops/fill_native.h>
-#include <ATen/ops/view_as_real.h>
-#include <ATen/ops/zero_native.h>
+#include <ATen/ops/empty_native.h>
 #endif
 
 namespace at::native {
 
-static Tensor& fill_scalar_mps_impl(Tensor& self, const Scalar& value) {
-  using namespace mps;
+#ifndef PYTORCH_JIT_COMPILE_SHADERS
+static auto& lib = mps::MetalShaderLibrary::getBundledLibrary();
+#else
+#include <ATen/native/mps/ConstantKernel_metallib.h>
+#endif
 
-  if (self.numel() == 0) {
-    return self;
-  }
-  Tensor output = self;
-  bool needsCopyToOutput = false;
-  if (needsGather(self)) {
-    output = at::empty(self.sizes(), self.scalar_type(), std::nullopt, kMPS, std::nullopt, std::nullopt);
-    needsCopyToOutput = true;
+static void fill_mps_kernel(TensorIterator& iter, const Scalar& value) {
+  using namespace mps;
+  if (iter.numel() == 0) {
+    return;
   }
 
-  struct CachedGraph : public MPSCachedGraph {
-    CachedGraph(MPSGraph* graph) : MPSCachedGraph(graph) {}
-    MPSGraphTensor* inputTensor_ = nil;
-    MPSGraphTensor* outputTensor_ = nil;
-  };
-
-  @autoreleasepool {
-    std::string key = "fill_scalar_mps_impl" + getTensorsStringKey(self) + ":" + std::to_string(value.toDouble());
-
-    auto cachedGraph = LookUpOrCreateCachedGraph<CachedGraph>(key, [&](auto mpsGraph, auto newCachedGraph) {
-      MPSGraphTensor* inputTensor = mpsGraphScalarPlaceHolder(mpsGraph, getMPSDataType(self.scalar_type()));
-      MPSGraphTensor* outputTensor = [mpsGraph identityWithTensor:inputTensor name:nil];
-      newCachedGraph->inputTensor_ = inputTensor;
-      newCachedGraph->outputTensor_ = outputTensor;
-    });
-
-    auto mpsScalar = getMPSScalar(value, self.scalar_type());
-    auto mpsScalarData = getMPSGraphTensorFromScalar(getCurrentMPSStream(), mpsScalar);
-    NSDictionary<MPSGraphTensor*, MPSGraphTensorData*>* feeds = @{cachedGraph->inputTensor_ : mpsScalarData};
-
-    Placeholder outputPlaceholder =
-        Placeholder(cachedGraph->outputTensor_, needsCopyToOutput ? output : self, nullptr, !needsCopyToOutput);
-
-    NSDictionary<MPSGraphTensor*, MPSGraphTensorData*>* results =
-        @{outputPlaceholder.getMPSGraphTensor() : outputPlaceholder.getMPSGraphTensorData()};
-
-    runMPSGraph(getCurrentMPSStream(), cachedGraph->graph(), feeds, results);
-
-    if (needsCopyToOutput) {
-      self.copy_(output);
+  // Metal compute kernels use uint (32-bit) thread indices; decompose large
+  // tensors into chunks that fit in 32-bit indexing.
+  if (!iter.can_use_32bit_indexing()) {
+    for (auto&& sub_iter : iter.with_32bit_indexing()) {
+      fill_mps_kernel(sub_iter, value);
     }
+    return;
   }
 
-  return self;
-}
-
-static Tensor& fill_mps_tensor_(Tensor& self, uint8_t value) {
-  TORCH_INTERNAL_ASSERT(self.is_contiguous());
+  const Tensor& self = iter.tensor(0);
+  const auto dtype = self.scalar_type();
   const auto stream = getCurrentMPSStream();
-  auto storage_byte_offset = self.storage_offset() * self.itemsize();
-  stream->fill(mps::getMTLBufferStorage(self), value, self.nbytes(), storage_byte_offset);
-  return self;
-}
-
-Tensor& fill_scalar_mps(Tensor& self, const Scalar& value) {
-  if (isComplexType(self.scalar_type())) {
-    auto self_as_real = at::view_as_real(self);
-    auto self_as_real_real = self_as_real.select(self.dim(), 0);
-    auto self_as_real_imag = self_as_real.select(self.dim(), 1);
-    if (value.isComplex()) {
-      auto value_cdouble = value.to<c10::complex<double>>();
-      fill_scalar_mps_impl(self_as_real_real, value_cdouble.real());
-      fill_scalar_mps_impl(self_as_real_imag, value_cdouble.imag());
-      return self;
-    }
-    fill_scalar_mps_impl(self_as_real_real, value);
-    fill_scalar_mps_impl(self_as_real_imag, 0.0f);
-    return self;
-  }
-  // check if it's possible to use fillBuffer() to fill the Tensor's storage
-  if (self.is_contiguous()) {
-    if (value.toDouble() == 0.0) {
-      return fill_mps_tensor_(self, 0);
-    }
-    if (self.scalar_type() == kBool) {
-      return fill_mps_tensor_(self, value.toBool());
-    }
-    if (self.scalar_type() == kByte) {
-      return fill_mps_tensor_(self, value.toByte());
-    }
-    if (self.scalar_type() == kChar) {
-      return fill_mps_tensor_(self, value.toChar());
-    }
+  const auto type_str = scalarToMetalTypeString(dtype);
+  const bool can_fill_linearly = self.is_non_overlapping_and_dense();
+
+  // For tensors with gaps or overlaps (e.g. stride-2 slices) use a 2D strided
+  // kernel: tid.y indexes dim 0 directly (no division), tid.x is the linear
+  // index for the remaining dims.  Consecutive threads in x write consecutive
+  // addresses in the innermost dimension, giving coalesced writes.
+  if (!can_fill_linearly) {
+    auto fillPSO = lib.getPipelineStateForFunc(fmt::format("fill_scalar_strided_{}", type_str));
+    const int64_t dim0_size = iter.ndim() > 0 ? iter.shape()[0] : 1;
+    const int64_t inner_numel = iter.numel() / dim0_size;
+    const uint32_t ndim = static_cast<uint32_t>(iter.ndim());
+    dispatch_sync_with_rethrow(stream->queue(), ^() {
+      @autoreleasepool {
+        auto computeEncoder = stream->commandEncoder();
+        auto mpsScalar = getMPSScalar(value, dtype);
+        [computeEncoder setComputePipelineState:fillPSO];
+        bind_iter_tensors(computeEncoder, iter);
+        mtl_setArgs<1>(computeEncoder, mpsScalar, iter.shape(), iter.strides(0), ndim);
+        const NSUInteger maxTG = fillPSO.maxTotalThreadsPerThreadgroup;
+        const MTLSize tgSize = MTLSizeMake(std::min(maxTG, (NSUInteger)inner_numel), 1, 1);
+        const MTLSize gridSize = MTLSizeMake(inner_numel, dim0_size, 1);
+        [computeEncoder dispatchThreads:gridSize threadsPerThreadgroup:tgSize];
+      }
+    });
+    return;
   }
 
-  return fill_scalar_mps_impl(self, value);
-}
-
-Tensor& fill_tensor_mps_(Tensor& self, const Tensor& value) {
-  TORCH_CHECK(value.dim() == 0,
-              "fill_ only supports 0-dimension value tensor but got tensor with ",
-              value.dim(),
-              " dimensions.");
-  Scalar scalar_value = value.item();
-  return fill_scalar_mps(self, scalar_value);
+  // Single-byte dtypes (bool, uint8, int8) use vec4 kernels that fill
+  // 4 elements per thread; all others fill 1 element per thread.
+  const bool is_byte_type = self.element_size() == 1;
+  const uint32_t numel = static_cast<uint32_t>(iter.numel());
+  const int64_t threads = is_byte_type ? (numel + 3) / 4 : numel;
+
+  auto fillPSO = lib.getPipelineStateForFunc(fmt::format("fill_scalar_dense_{}", type_str));
+  dispatch_sync_with_rethrow(stream->queue(), ^() {
+    @autoreleasepool {
+      auto computeEncoder = stream->commandEncoder();
+      auto mpsScalar = getMPSScalar(value, dtype);
+      [computeEncoder setComputePipelineState:fillPSO];
+      bind_iter_tensors(computeEncoder, iter);
+      mtl_setArgs<1>(computeEncoder, mpsScalar, numel);
+      mtl_dispatch1DJob(computeEncoder, fillPSO, threads);
+    }
+  });
 }
 
-Tensor& zero_mps_(Tensor& self) {
-  return fill_scalar_mps(self, 0.0f);
-}
+REGISTER_DISPATCH(fill_stub, &fill_mps_kernel);
 
 } // namespace at::native
diff --git a/aten/src/ATen/native/mps/operations/Convolution.mm b/aten/src/ATen/native/mps/operations/Convolution.mm
index a457267a9d850..cf46afdd92b9d 100644
--- a/aten/src/ATen/native/mps/operations/Convolution.mm
+++ b/aten/src/ATen/native/mps/operations/Convolution.mm
@@ -114,7 +114,10 @@ static Tensor _mps_convolution_impl(const Tensor& input_t,
   const bool is_macos_15_plus = is_macos_13_or_newer(MacOSVersion::MACOS_VER_15_0_PLUS);
 
   const bool is3DConv = input_t.dim() == 5;
-  const auto memory_format = input_t.suggest_memory_format();
+  // Use exact-match: a channel-slice of a channels-last tensor has CL-like
+  // strides but is not NHWC-packed, so the raw-buffer NHWC path would misread
+  // it. See https://github.com/pytorch/pytorch/issues/180984
+  const auto memory_format = input_t.suggest_memory_format(/*channels_last_strides_exact_match=*/true);
   const auto input_suggested_layout = memory_format == kChannelsLast && is_macos_15_plus ? kChannelsLast : kContiguous;
   const bool is_channels_last = mps_conv_use_channels_last(input_t, weight_t) && !is3DConv;
   const bool bias_defined = bias_opt ? bias_opt->defined() : false;
diff --git a/aten/src/ATen/native/mps/operations/Copy.mm b/aten/src/ATen/native/mps/operations/Copy.mm
index a3cba05c975cf..fedf6d6cfe837 100644
--- a/aten/src/ATen/native/mps/operations/Copy.mm
+++ b/aten/src/ATen/native/mps/operations/Copy.mm
@@ -333,6 +333,7 @@ void copy_blit_mps(void* dst, const void* src, size_t size) {
 } // namespace mps
 
 Tensor _copy_from_and_resize_mps(const at::Tensor& self, const at::Tensor& dst) {
+  const_cast<Tensor&>(dst).resize_as_(self);
   return mps::mps_copy_(const_cast<Tensor&>(dst), self, false);
 }
 
diff --git a/aten/src/ATen/native/mps/operations/CrossKernel.mm b/aten/src/ATen/native/mps/operations/CrossKernel.mm
index 35073a94ac5be..8bd4a0af920dc 100644
--- a/aten/src/ATen/native/mps/operations/CrossKernel.mm
+++ b/aten/src/ATen/native/mps/operations/CrossKernel.mm
@@ -1,6 +1,5 @@
 //  Copyright © 2022 Apple Inc.
 #define TORCH_ASSERT_ONLY_METHOD_OPERATORS
-#include <ATen/TensorIterator.h>
 #include <ATen/mps/MPSProfiler.h>
 #include <ATen/native/Cross.h>
 #include <ATen/native/mps/OperationUtils.h>
@@ -17,39 +16,53 @@
 #endif
 
 void cross_mps_impl(const Tensor& out, const Tensor& input, const Tensor& other, int64_t dim) {
-  TORCH_CHECK(input.dtype() != at::kDouble, "float64 is not supported on MPS");
+  // numThreads = number of cross triplets = numel / 3
+  const auto numThreads = out.numel() / 3;
+  // The dense kernel handles any contiguous layout: it uses dim_stride and a
+  // closed-form index mapping that works for outermost, innermost, and any dim.
+  const bool is_dense = out.is_contiguous() && input.is_contiguous() && other.is_contiguous();
+  const std::string suffix = is_dense ? "dense" : "strided";
 
-  auto iter = TensorIteratorConfig()
-                  .add_output(out)
-                  .add_input(input)
-                  .add_input(other)
-                  .resize_outputs(false)
-                  .declare_static_shape(out.sizes(), /*squash_dims=*/dim)
-                  .build();
+  // For the strided kernel, build squashed sizes (cross dim removed) for the
+  // thread-index → position mapping, plus full element strides (all ndim dims,
+  // including the cross dim) so the kernel can read strides[dim] directly.
+  const int64_t ndim = out.dim();
+  std::vector<int64_t> squashed_sizes, out_strides, input_strides, other_strides;
+  if (!is_dense) {
+    squashed_sizes.reserve(ndim - 1);
+    out_strides.resize(ndim);
+    input_strides.resize(ndim);
+    other_strides.resize(ndim);
+    for (int64_t d = 0; d < ndim; ++d) {
+      out_strides[d] = out.stride(d);
+      input_strides[d] = input.stride(d);
+      other_strides[d] = other.stride(d);
+      if (d != dim) {
+        squashed_sizes.push_back(out.size(d));
+      }
+    }
+  }
 
-  id<MTLDevice> device = MPSDevice::getInstance()->device();
-  MPSStream* mpsStream = getCurrentMPSStream();
-  const int64_t out_dim_stride = out.stride(dim);
-  const int64_t input_dim_stride = input.stride(dim);
-  const int64_t other_dim_stride = other.stride(dim);
-  const uint32_t nDim = iter.ndim();
-  constexpr uint32_t nOffsets = 3;
-  const uint32_t numThreads = iter.numel();
+  const auto mpsStream = getCurrentMPSStream();
   dispatch_sync_with_rethrow(mpsStream->queue(), ^() {
     @autoreleasepool {
       id<MTLComputeCommandEncoder> computeEncoder = mpsStream->commandEncoder();
-      auto kernelDataOffsets = generateKernelDataOffsets(computeEncoder, iter);
-
-      auto crossPSO = lib.getPipelineStateForFunc("cross_" + scalarToMetalTypeString(out));
-
-      // this function call is a no-op if MPS Profiler is not enabled
+      auto crossPSO = lib.getPipelineStateForFunc(fmt::format("cross_{}_{}", suffix, scalarToMetalTypeString(out)));
       getMPSProfiler().beginProfileKernel(crossPSO, "cross", {input, other});
-
       [computeEncoder setComputePipelineState:crossPSO];
-      mtl_setArgs(
-          computeEncoder, input, other, out, kernelDataOffsets, out_dim_stride, input_dim_stride, other_dim_stride);
+      mtl_setArgs(computeEncoder, out, input, other);
+      if (is_dense) {
+        mtl_setArgs<3>(computeEncoder, static_cast<uint32_t>(out.stride(dim)));
+      } else {
+        mtl_setArgs<3>(computeEncoder,
+                       squashed_sizes,
+                       out_strides,
+                       input_strides,
+                       other_strides,
+                       static_cast<uint32_t>(ndim),
+                       static_cast<uint32_t>(dim));
+      }
       mtl_dispatch1DJob(computeEncoder, crossPSO, numThreads);
-
       getMPSProfiler().endProfileKernel(crossPSO);
     }
   });
diff --git a/aten/src/ATen/native/mps/operations/Distributions.mm b/aten/src/ATen/native/mps/operations/Distributions.mm
index d1d7bd60f7a16..a1016fe5cdc90 100644
--- a/aten/src/ATen/native/mps/operations/Distributions.mm
+++ b/aten/src/ATen/native/mps/operations/Distributions.mm
@@ -11,6 +11,8 @@
 #include <ATen/Functions.h>
 #include <ATen/NativeFunctions.h>
 #else
+#include <ATen/ops/_standard_gamma_grad_native.h>
+#include <ATen/ops/_standard_gamma_native.h>
 #include <ATen/ops/argmax.h>
 #include <ATen/ops/bernoulli_native.h>
 #include <ATen/ops/cauchy_native.h>
@@ -419,38 +421,13 @@ Tensor normal_mps(const Tensor& mean, const Tensor& std, std::optional<Generator
   return random_mps_(self, 0, std::nullopt, gen);
 }
 
-// Exponential distribution
-Tensor& exponential_mps_(Tensor& self, double lambda, std::optional<Generator> gen) {
-  TORCH_CHECK(lambda > 0.0, "exponential_ expects lambda > 0.0, but found lambda=", lambda);
-
-  mps::RandomOpBlock random_op_block = ^RandomOpFn(cachedGraph, randomTensor) {
-    MPSGraph* mpsGraph = cachedGraph->graph();
-    MPSGraphTensor* unitTensor = [mpsGraph constantWithScalar:1.0f dataType:randomTensor.dataType];
-    MPSGraphTensor* minusLambdaTensor = [mpsGraph constantWithScalar:-lambda dataType:randomTensor.dataType];
-    MPSGraphTensor* subtractTensor = [mpsGraph subtractionWithPrimaryTensor:unitTensor
-                                                            secondaryTensor:randomTensor
-                                                                       name:nil];
-    MPSGraphTensor* logTensor = [mpsGraph logarithmWithTensor:subtractTensor name:nil];
-    return [mpsGraph divisionWithPrimaryTensor:logTensor secondaryTensor:minusLambdaTensor name:nil];
-  };
-  auto eps = std::numeric_limits<float>::epsilon();
-  return mps::random_mps_impl<double>(self,
-                                      eps,
-                                      1.0,
-                                      std::nullopt,
-                                      std::nullopt,
-                                      MPSGraphRandomDistributionUniform,
-                                      gen,
-                                      "exponential_mps_:" + std::to_string(lambda),
-                                      random_op_block);
-}
-
 static Tensor& distribution_kernel_mps_impl(Tensor& self,
                                             double param1,
                                             double param2,
                                             const std::string& kernel_name,
                                             int64_t randoms_per_element,
-                                            std::optional<Generator> gen) {
+                                            std::optional<Generator> gen,
+                                            int64_t elements_per_thread = 1) {
   if (self.numel() == 0) {
     return self;
   }
@@ -483,7 +460,13 @@ Tensor normal_mps(const Tensor& mean, const Tensor& std, std::optional<Generator
                     output,
                     std::array<float, 2>{static_cast<float>(param1), static_cast<float>(param2)},
                     std::array<long, 2>{seed, base_offset});
-        mtl_dispatch1DJob(computeEncoder, pso, output.numel());
+        if (elements_per_thread > 1) {
+          auto numel = static_cast<uint32_t>(output.numel());
+          mtl_setBytes(computeEncoder, numel, 3);
+          mtl_dispatch1DJob(computeEncoder, pso, (numel + elements_per_thread - 1) / elements_per_thread);
+        } else {
+          mtl_dispatch1DJob(computeEncoder, pso, output.numel());
+        }
       }
     });
   }
@@ -495,6 +478,11 @@ Tensor normal_mps(const Tensor& mean, const Tensor& std, std::optional<Generator
   return self;
 }
 
+Tensor& exponential_mps_(Tensor& self, double lambda, std::optional<Generator> gen) {
+  TORCH_CHECK(lambda > 0.0, "exponential_ expects lambda > 0.0, but found lambda=", lambda);
+  return distribution_kernel_mps_impl(self, lambda, 0.0, "exponential", 1, gen, /*elements_per_thread=*/4);
+}
+
 Tensor& cauchy_mps_(Tensor& self, double median, double sigma, std::optional<Generator> gen) {
   TORCH_CHECK(sigma > 0.0, "cauchy_ expects sigma > 0.0, but found sigma=", sigma);
   return distribution_kernel_mps_impl(self, median, sigma, "cauchy", 1, gen);
@@ -511,6 +499,74 @@ Tensor normal_mps(const Tensor& mean, const Tensor& std, std::optional<Generator
   return distribution_kernel_mps_impl(self, log_one_minus_p, 0.0, "geometric", 1, gen);
 }
 
+Tensor _s_gamma_mps(const Tensor& alpha, std::optional<Generator> gen) {
+  if (alpha.numel() == 0) {
+    return at::empty_like(alpha);
+  }
+
+  using namespace mps;
+
+  auto mps_gen = get_generator_or_default<MPSGeneratorImpl>(gen, at::mps::detail::getDefaultMPSGenerator());
+  auto stream = getCurrentMPSStream();
+  Tensor ret = at::empty_like(alpha, alpha.options(), at::MemoryFormat::Contiguous);
+  auto alpha_contig = alpha.contiguous();
+
+  @autoreleasepool {
+    auto pso = lib.getPipelineStateForFunc("standard_gamma_" + scalarToMetalTypeString(ret));
+
+    int64_t seed;
+    int64_t base_offset;
+    // Each thread may consume up to GAMMA_RANDOMS_STRIDE random numbers
+    constexpr int64_t GAMMA_RANDOMS_STRIDE = 32;
+    {
+      // See Note [Acquire lock when using random generators]
+      std::lock_guard<std::mutex> lock(mps_gen->mutex_);
+      seed = static_cast<int64_t>(mps_gen->current_seed());
+      base_offset = static_cast<int64_t>(mps_gen->get_offset());
+      mps_gen->set_offset(base_offset + GAMMA_RANDOMS_STRIDE * ret.numel());
+    }
+
+    dispatch_sync_with_rethrow(stream->queue(), ^() {
+      @autoreleasepool {
+        auto computeEncoder = stream->commandEncoder();
+        [computeEncoder setComputePipelineState:pso];
+        mtl_setArgs(computeEncoder, ret, alpha_contig, std::array<long, 2>{seed, base_offset});
+        mtl_dispatch1DJob(computeEncoder, pso, ret.numel());
+      }
+    });
+  }
+
+  return ret;
+}
+
+Tensor _standard_gamma_grad_mps(const Tensor& self, const Tensor& output) {
+  if (self.numel() == 0) {
+    return at::empty_like(self);
+  }
+
+  using namespace mps;
+
+  auto stream = getCurrentMPSStream();
+  Tensor ret = at::empty_like(self, self.options(), at::MemoryFormat::Contiguous);
+  const auto self_contig = self.contiguous();
+  const auto output_contig = output.contiguous();
+
+  @autoreleasepool {
+    auto pso = lib.getPipelineStateForFunc("standard_gamma_grad_" + scalarToMetalTypeString(ret));
+
+    dispatch_sync_with_rethrow(stream->queue(), ^() {
+      @autoreleasepool {
+        auto computeEncoder = stream->commandEncoder();
+        [computeEncoder setComputePipelineState:pso];
+        mtl_setArgs(computeEncoder, ret, self_contig, output_contig);
+        mtl_dispatch1DJob(computeEncoder, pso, ret.numel());
+      }
+    });
+  }
+
+  return ret;
+}
+
 Tensor& randperm_out_mps(int64_t n, std::optional<Generator> generator, Tensor& result) {
   TORCH_CHECK(n >= 0, "n must be non-negative, got", n);
   TORCH_CHECK(!generator.has_value() || (generator.has_value() && result.device() == generator->device()),
diff --git a/aten/src/ATen/native/mps/operations/Eye.mm b/aten/src/ATen/native/mps/operations/Eye.mm
index 592654dd57740..1b603304f565e 100644
--- a/aten/src/ATen/native/mps/operations/Eye.mm
+++ b/aten/src/ATen/native/mps/operations/Eye.mm
@@ -1,5 +1,7 @@
 #define TORCH_ASSERT_ONLY_METHOD_OPERATORS
+#include <ATen/core/Tensor.h>
 #include <ATen/native/mps/OperationUtils.h>
+
 #ifndef AT_PER_OPERATOR_HEADERS
 #include <ATen/Functions.h>
 #include <ATen/NativeFunctions.h>
@@ -7,95 +9,76 @@
 #include <ATen/ops/eye_native.h>
 #endif
 
-// Steps to add op for MPS backend:
-// 1. Register the op in aten/src/ATen/native/native_functions.yaml with the "MPS" dispatch key
-// 2. Define the function interface for the MPS backend similar to other
-//    backends depending on whether its structured or non-structured
-// 3. Add boiler-plate error checking code as expected for the Op
-// 4. The code structure roughly follows the pattern
-//    a) get the MPS stream handle to encode work onto
-//    b) get an instance of MPSGraphCache and create a key unique to the Graph
-//       needed for implementing this Op. Any shape, dataType or parameter
-//       passed to the MPSGraph during its construction will need to be included
-//       here.
-//    c) Create the graph using make_mps_graph() and add operations to the
-//       instance of MPSGraph. This is if the Cache->lookup() fails.
-//    d) Store the MPSGraphTensors for inputs and output which are needed at
-//       runtime.
-//    e) Use the CachedGraph instance's inputs and output to create Placeholders
-//       You will need to pass in Tensor to create MPSGraphTensorData objects.
-//    f) Using MPSGraphTensor and MPSGraphTensorData instances create a feeds
-//       dictionary.
-//    g) Then call runMPSGraph() with input params and return the result.
-//
-
 namespace at::native {
+using namespace mps;
+
+#ifndef PYTORCH_JIT_COMPILE_SHADERS
+static auto& lib = MetalShaderLibrary::getBundledLibrary();
+#else
+#include <ATen/native/mps/Eye_metallib.h>
+#endif
 
 Tensor& eye_out_mps(int64_t n, Tensor& result) {
   // the default value of `m` equals to `n`
   return eye_out_mps(n, n, result);
 }
 
-using namespace mps;
-
 Tensor& eye_out_mps(int64_t n, int64_t m, Tensor& result) {
-  // This is one example of boiler-plate error checking, taking after CPU/CUDA counterparts
   TORCH_CHECK(n >= 0, "n must be greater or equal to 0, got ", n);
   TORCH_CHECK(m >= 0, "m must be greater or equal to 0, got ", m);
 
   result.resize_({n, m});
-  result.zero_();
 
-  // Handle empty outputs
   if (result.numel() == 0)
     return result;
 
-  // Get MPS stream
-  MPSStream* stream = getCurrentMPSStream();
-
-  auto outputDataType = result.scalar_type();
-  // Derive from MPSCachedGraph
-  // This structure is used to cache an MPSGraph with certain keys, so that we don't have to compile the same MPSGraph
-  // time and time again for the same operation The keys of this structure are based on the inputs and outputs needed
-  // for the operation here, we don't have any input tensors, just an output tensor.
-  // If the operator to be added is unary or binary, instead of creating a new CachedGraph struct yourself, please
-  // consider using `MPSUnaryCachedGraph` or `MPSBinaryCachedGraph` and their corresponding Grad versions in
-  // `OperationUtils.h`.
-  struct CachedGraph : public MPSCachedGraph {
-    CachedGraph(MPSGraph* graph) : MPSCachedGraph(graph) {}
-    MPSGraphTensor* outputTensor_ = nil;
-  };
-
-  @autoreleasepool {
-    // A key is used to identify the MPSGraph which was created once, and can be reused if the parameters, data types
-    // etc match the earlier created MPSGraph
-    std::string key = "eye_out_mps:" + getTensorsStringKey({result});
-    auto cachedGraph = LookUpOrCreateCachedGraph<CachedGraph>(key, [&](auto* mpsGraph, auto* newCachedGraph) {
-      MPSGraphTensor* onesTensor = [mpsGraph constantWithScalar:1.0f
-                                                          shape:getMPSShape(result)
-                                                       dataType:getMPSDataType(outputDataType)];
-
-      // Here we can call the MPSGraph API needed to execute the operation.
-      // The API details can be found here:
-      // https://developer.apple.com/documentation/metalperformanceshadersgraph/mpsgraph
-      MPSGraphTensor* outputTensor = [mpsGraph bandPartWithTensor:onesTensor numLower:0 numUpper:0 name:nil];
-
-      if ([outputTensor dataType] != getMPSDataType(outputDataType)) {
-        outputTensor = castMPSTensor(mpsGraph, outputTensor, outputDataType);
+  MPSStream* mpsStream = getCurrentMPSStream();
+  auto stride0 = result.stride(0);
+  auto stride1 = result.stride(1);
+
+  // Small tensors: single-pass 2D kernel (one dispatch, no zero_() overhead).
+  // Large tensors: zero_() + diagonal fill (memset is faster than n*m branching writes).
+  constexpr int64_t kSinglePassThreshold = 1 << 22;
+
+  if (n * m <= kSinglePassThreshold) {
+    auto key = "eye_" + scalarToMetalTypeString(result);
+    id<MTLComputeCommandEncoder> computeEncoder = mpsStream->commandEncoder();
+    id<MTLComputePipelineState> pso = lib.getPipelineStateForFunc(key);
+
+    // Map x to the smaller stride for coalesced writes
+    bool swap = stride0 < stride1;
+    auto x_stride = swap ? stride0 : stride1;
+    auto y_stride = swap ? stride1 : stride0;
+    auto grid_x = swap ? static_cast<NSUInteger>(n) : static_cast<NSUInteger>(m);
+    auto grid_y = swap ? static_cast<NSUInteger>(m) : static_cast<NSUInteger>(n);
+
+    dispatch_sync(mpsStream->queue(), ^() {
+      @autoreleasepool {
+        [computeEncoder setComputePipelineState:pso];
+        mtl_setArgs(computeEncoder, result, y_stride, x_stride);
+
+        auto maxTg = [pso maxTotalThreadsPerThreadgroup];
+        auto tg_x = std::min(grid_x, maxTg);
+        auto tg_y = std::min(grid_y, std::max(maxTg / tg_x, static_cast<NSUInteger>(1)));
+        [computeEncoder dispatchThreads:MTLSizeMake(grid_x, grid_y, 1)
+                  threadsPerThreadgroup:MTLSizeMake(tg_x, tg_y, 1)];
+      }
+    });
+  } else {
+    result.zero_();
+    int64_t sz = std::min(n, m);
+    int64_t diag_stride = stride0 + stride1;
+    auto key = "eye_diag_" + scalarToMetalTypeString(result);
+    id<MTLComputeCommandEncoder> computeEncoder = mpsStream->commandEncoder();
+    id<MTLComputePipelineState> pso = lib.getPipelineStateForFunc(key);
+
+    dispatch_sync(mpsStream->queue(), ^() {
+      @autoreleasepool {
+        [computeEncoder setComputePipelineState:pso];
+        mtl_setArgs(computeEncoder, result, diag_stride);
+        mtl_dispatch1DJob(computeEncoder, pso, sz);
       }
-      newCachedGraph->outputTensor_ = outputTensor;
     });
-
-    // Create placeholders which use the keys of the CachedGraph to create inputs and outputs of the operation
-    Placeholder outputPlaceholder = Placeholder(cachedGraph->outputTensor_, result);
-
-    // Create dictionary of inputs/feeds and outputs/results
-    // In this case, there are no inputs, so the feeds are nil
-    NSDictionary<MPSGraphTensor*, MPSGraphTensorData*>* feeds = nil;
-    auto results = dictionaryFromPlaceholders(outputPlaceholder);
-
-    // Run the graph
-    runMPSGraph(stream, cachedGraph->graph(), feeds, results);
   }
 
   return result;
diff --git a/aten/src/ATen/native/mps/operations/Gamma.mm b/aten/src/ATen/native/mps/operations/Gamma.mm
deleted file mode 100644
index 9feb5eba84992..0000000000000
--- a/aten/src/ATen/native/mps/operations/Gamma.mm
+++ /dev/null
@@ -1,158 +0,0 @@
-#define TORCH_ASSERT_ONLY_METHOD_OPERATORS
-#include <ATen/ExpandUtils.h>
-#include <ATen/mps/MPSProfiler.h>
-#include <ATen/native/Resize.h>
-#include <ATen/native/mps/OperationUtils.h>
-#include <fmt/format.h>
-
-#include <ATen/ops/digamma_native.h>
-#include <ATen/ops/lgamma_native.h>
-#include <ATen/ops/polygamma_native.h>
-
-namespace at::native {
-namespace mps {
-#ifndef PYTORCH_JIT_COMPILE_SHADERS
-static auto& lib = MetalShaderLibrary::getBundledLibrary();
-#else
-#include <ATen/native/mps/Gamma_metallib.h>
-#endif
-
-static id<MTLComputePipelineState> getCPLState(const Tensor& t1, const Tensor& t2, const std::string& fname) {
-  return lib.getPipelineStateForFunc(
-      fmt::format("{}_{}_{}", fname, scalarToMetalTypeString(t1), scalarToMetalTypeString(t2)));
-}
-
-} // namespace mps
-
-TORCH_IMPL_FUNC(lgamma_out_mps)(const Tensor& self, const Tensor& output_) {
-  TORCH_CHECK(self.scalar_type() != ScalarType::Double, "MPS does not support lgamma_out op with scalar type: Double");
-
-  Tensor output = output_;
-  bool needs_output_copy = false;
-  uint32_t length = output.numel();
-  if (length == 0) {
-    return;
-  }
-
-  if (mps::needsGather(output_)) {
-    output = output.contiguous();
-    needs_output_copy = true;
-  }
-
-  using namespace mps;
-
-  @autoreleasepool {
-    auto cplState = getCPLState(self, output, "lgamma");
-
-    MPSStream* mpsStream = getCurrentMPSStream();
-    dispatch_sync(mpsStream->queue(), ^() {
-      id<MTLComputeCommandEncoder> computeEncoder = mpsStream->commandEncoder();
-
-      getMPSProfiler().beginProfileKernel(cplState, "lgamma_out", {self});
-
-      [computeEncoder setComputePipelineState:cplState];
-      mtl_setArgs(computeEncoder, self, output);
-      mtl_dispatch1DJob(computeEncoder, cplState, length);
-
-      getMPSProfiler().endProfileKernel(cplState);
-    });
-  }
-  if (needs_output_copy) {
-    output_.copy_(output);
-  }
-}
-
-TORCH_IMPL_FUNC(digamma_out_mps)(const Tensor& self, const Tensor& output_) {
-  TORCH_CHECK(self.scalar_type() != ScalarType::Double, "MPS does not support digamma_out op with scalar type: Double");
-
-  Tensor output = output_;
-  bool needs_output_copy = false;
-  uint32_t length = output.numel();
-  if (length == 0) {
-    return;
-  }
-
-  if (mps::needsGather(output_)) {
-    output = output.contiguous();
-    needs_output_copy = true;
-  }
-
-  using namespace mps;
-
-  @autoreleasepool {
-    id<MTLComputePipelineState> cplState = getCPLState(self, output, "digamma");
-
-    MPSStream* mpsStream = getCurrentMPSStream();
-    dispatch_sync(mpsStream->queue(), ^() {
-      id<MTLComputeCommandEncoder> computeEncoder = mpsStream->commandEncoder();
-
-      getMPSProfiler().beginProfileKernel(cplState, "digamma_out", {self});
-
-      [computeEncoder setComputePipelineState:cplState];
-      mtl_setArgs(computeEncoder, self, output);
-      mtl_dispatch1DJob(computeEncoder, cplState, length);
-
-      getMPSProfiler().endProfileKernel(cplState);
-    });
-  }
-  if (needs_output_copy) {
-    output_.copy_(output);
-  }
-}
-
-TORCH_IMPL_FUNC(polygamma_out_mps)(const int64_t order, const Tensor& self, const Tensor& output_) {
-  TORCH_CHECK(self.scalar_type() != ScalarType::Double,
-              "MPS does not support polygamma_out op with scalar type: Double");
-  TORCH_CHECK(order >= 0, "Polygamma is implemented only for nonnegative real numbers");
-
-  Tensor output = output_;
-  bool needs_output_copy = false;
-  uint32_t length = output.numel();
-  if (length == 0) {
-    return;
-  }
-
-  if (mps::needsGather(output_)) {
-    output = output.contiguous();
-    needs_output_copy = true;
-  }
-
-  using namespace mps;
-
-  std::string func_name;
-
-  if (order == 0) {
-    func_name = "digamma";
-  } else if (order == 1) {
-    func_name = "trigamma";
-  } else {
-    func_name = "polygamma";
-  }
-
-  @autoreleasepool {
-    id<MTLComputePipelineState> cplState = getCPLState(self, output, func_name);
-
-    MPSStream* mpsStream = getCurrentMPSStream();
-    dispatch_sync(mpsStream->queue(), ^() {
-      id<MTLComputeCommandEncoder> computeEncoder = mpsStream->commandEncoder();
-
-      getMPSProfiler().beginProfileKernel(cplState, func_name, {self});
-
-      [computeEncoder setComputePipelineState:cplState];
-      mtl_setArgs(computeEncoder, self, output);
-
-      if (func_name == "polygamma") {
-        mtl_setBytes(computeEncoder, order, 2);
-      }
-
-      mtl_dispatch1DJob(computeEncoder, cplState, length);
-
-      getMPSProfiler().endProfileKernel(cplState);
-    });
-  }
-  if (needs_output_copy) {
-    output_.copy_(output);
-  }
-}
-
-} // namespace at::native
diff --git a/aten/src/ATen/native/mps/operations/GridSampler.mm b/aten/src/ATen/native/mps/operations/GridSampler.mm
index 5d53b9d6dfad5..2a232decc66ed 100644
--- a/aten/src/ATen/native/mps/operations/GridSampler.mm
+++ b/aten/src/ATen/native/mps/operations/GridSampler.mm
@@ -24,6 +24,32 @@
 
 namespace mps {
 
+static const char* interp_to_string(GridSamplerInterpolation mode) {
+  switch (mode) {
+    case GridSamplerInterpolation::Bilinear:
+      return "bilinear";
+    case GridSamplerInterpolation::Nearest:
+      return "nearest";
+    case GridSamplerInterpolation::Bicubic:
+      return "bicubic";
+  }
+  TORCH_CHECK(false, "Unrecognised interpolation mode: ", mode);
+  return "";
+}
+
+static const char* padding_to_string(GridSamplerPadding mode) {
+  switch (mode) {
+    case GridSamplerPadding::Zeros:
+      return "zeros";
+    case GridSamplerPadding::Border:
+      return "border";
+    case GridSamplerPadding::Reflection:
+      return "reflection";
+  }
+  TORCH_CHECK(false, "Unrecognised padding mode: ", mode);
+  return "";
+}
+
 static void grid_sampler_2d_mps_impl(Tensor& output,
                                      const Tensor& input,
                                      const Tensor& grid,
@@ -46,30 +72,10 @@ static void grid_sampler_2d_mps_impl(Tensor& output,
   auto interpolation_mode = static_cast<GridSamplerInterpolation>(_interpolation_mode);
   auto padding_mode = static_cast<GridSamplerPadding>(_padding_mode);
 
-  switch (interpolation_mode) {
-    case GridSamplerInterpolation::Bilinear:
-    case GridSamplerInterpolation::Nearest:
-    case GridSamplerInterpolation::Bicubic:
-      break;
-    default:
-      TORCH_CHECK(false, "grid_sampler_2d: Unrecognised interpolation mode: ", _interpolation_mode);
-  }
-
-  switch (padding_mode) {
-    case GridSamplerPadding::Zeros:
-    case GridSamplerPadding::Border:
-    case GridSamplerPadding::Reflection:
-      break;
-    default:
-      TORCH_CHECK(false, "grid_sampler_2d: Unrecognised Padding Mode: ", _padding_mode);
-  }
-
   auto dims = input.dim();
 
-  GridSamplerParams<5> params;
+  GridSamplerParams<4> params;
   params.sampler_dims = 2;
-  params.padding_mode = padding_mode;
-  params.interpolation_mode = interpolation_mode;
   params.align_corners = align_corners;
 
   for (const auto dim : c10::irange(dims)) {
@@ -91,7 +97,10 @@ static void grid_sampler_2d_mps_impl(Tensor& output,
   dispatch_sync_with_rethrow(mpsStream->queue(), ^() {
     @autoreleasepool {
       id<MTLComputeCommandEncoder> computeEncoder = mpsStream->commandEncoder();
-      auto pso = lib.getPipelineStateForFunc("grid_sampler_2d_" + scalarToMetalTypeString(input));
+      auto pso = lib.getPipelineStateForFunc(fmt::format("grid_sampler_2d_{}_{}_{}",
+                                                         interp_to_string(interpolation_mode),
+                                                         padding_to_string(padding_mode),
+                                                         scalarToMetalTypeString(input)));
 
       getMPSProfiler().beginProfileKernel(pso, "grid_sampler_2d", {input, grid});
       [computeEncoder setComputePipelineState:pso];
@@ -144,15 +153,6 @@ static void grid_sampler_3d_mps_impl(Tensor& output,
       TORCH_CHECK(false, op_name, ": Unrecognised interpolation mode: ", _interpolation_mode);
   }
 
-  switch (padding_mode) {
-    case GridSamplerPadding::Zeros:
-    case GridSamplerPadding::Border:
-    case GridSamplerPadding::Reflection:
-      break;
-    default:
-      TORCH_CHECK(false, op_name, ": Unrecognised Padding Mode: ", _padding_mode);
-  }
-
   auto input_size = input.sizes();
   auto grid_size = grid.sizes();
   output.resize_({input_size[0], input_size[1], grid_size[1], grid_size[2], grid_size[3]}, MemoryFormat::Contiguous);
@@ -161,8 +161,6 @@ static void grid_sampler_3d_mps_impl(Tensor& output,
 
   GridSamplerParams<5> params;
   params.sampler_dims = sampler_dims;
-  params.padding_mode = padding_mode;
-  params.interpolation_mode = interpolation_mode;
   params.align_corners = align_corners;
 
   for (const auto dim : c10::irange(dims)) {
@@ -180,7 +178,8 @@ static void grid_sampler_3d_mps_impl(Tensor& output,
   dispatch_sync_with_rethrow(mpsStream->queue(), ^() {
     @autoreleasepool {
       id<MTLComputeCommandEncoder> computeEncoder = mpsStream->commandEncoder();
-      auto pso = lib.getPipelineStateForFunc("grid_sampler_" + scalarToMetalTypeString(input));
+      auto pso = lib.getPipelineStateForFunc(
+          fmt::format("grid_sampler_3d_{}_{}", padding_to_string(padding_mode), scalarToMetalTypeString(input)));
 
       getMPSProfiler().beginProfileKernel(pso, op_name, {input, grid});
       [computeEncoder setComputePipelineState:pso];
diff --git a/aten/src/ATen/native/mps/operations/HistogramKernel.mm b/aten/src/ATen/native/mps/operations/HistogramKernel.mm
index 2f12d4eef95a7..0e0dcc53b04e6 100644
--- a/aten/src/ATen/native/mps/operations/HistogramKernel.mm
+++ b/aten/src/ATen/native/mps/operations/HistogramKernel.mm
@@ -1,5 +1,5 @@
 #define TORCH_ASSERT_ONLY_METHOD_OPERATORS
-#include <ATen/Dispatch.h>
+#include <ATen/Dispatch_v2.h>
 #include <ATen/mps/MPSProfiler.h>
 #include <ATen/native/Histogram.h>
 #include <ATen/native/mps/OperationUtils.h>
@@ -33,7 +33,6 @@ void histogramdd_kernel_impl(Tensor& hist_output,
                              const TensorList& bin_edges,
                              const Tensor& input,
                              const std::optional<Tensor>& weight) {
-  TORCH_CHECK(input.dtype() != at::kDouble, "float64 is not supported on MPS");
   TORCH_INTERNAL_ASSERT(input.dim() == 2);
 
   constexpr uint8_t bin_selection_algorithm = algorithm;
@@ -41,11 +40,11 @@ void histogramdd_kernel_impl(Tensor& hist_output,
   const bool has_weight = weight.has_value();
 
   if (has_weight) {
-    TORCH_CHECK(weight.value().is_contiguous(), "histogramdd(): weight should be contiguous on MPS");
     TORCH_INTERNAL_ASSERT(weight.value().dim() == 1 && weight.value().numel() == N);
     TORCH_INTERNAL_ASSERT(weight.value().scalar_type() == input.scalar_type());
   }
 
+  const int64_t weight_stride = has_weight ? weight.value().stride(0) : -1;
   const int64_t D = input.size(1);
   size_t bin_edges_numel = 0;
   TORCH_INTERNAL_ASSERT(int64_t(bin_edges.size()) == D);
@@ -135,7 +134,7 @@ void histogramdd_kernel_impl(Tensor& hist_output,
                      rightmost_edge,
                      thread_histograms.strides(),
                      bin_selection_algorithm,
-                     has_weight);
+                     weight_stride);
 
       mtl_dispatch1DJob(computeEncoder, histogramPSO, numThreads);
 
@@ -167,9 +166,15 @@ static void histogramdd_out_mps_template(const Tensor& self,
     bin_edges_contig[dim] = bin_edges[dim].contiguous();
   }
 
-  AT_DISPATCH_FLOATING_TYPES(self.scalar_type(), "histogram_mps", [&]() {
-    mps::histogramdd_kernel_impl<scalar_t, bin_algorithm>(hist, bin_edges_contig, reshaped_input, reshaped_weight);
-  });
+  AT_DISPATCH_V2(self.scalar_type(),
+                 "histogram_mps",
+                 AT_WRAP([&]() {
+                   mps::histogramdd_kernel_impl<scalar_t, bin_algorithm>(
+                       hist, bin_edges_contig, reshaped_input, reshaped_weight);
+                 }),
+                 kFloat,
+                 kHalf,
+                 kBFloat16);
 
   /* Divides each bin's value by the total count/weight in all bins,
    * and by the bin's volume.
@@ -211,6 +216,7 @@ static void histogramdd_linear_kernel(const Tensor& self,
   if (local_search) {
     // histogramdd codepath: both hist and bin_edges are eventually returned as output,
     // so we'll keep them consistent
+    TORCH_CHECK(self.scalar_type() == kFloat, "histogram is only supported for float32");
     mps::histogramdd_out_mps_template<mps::LINEAR_INTERPOLATION_WITH_LOCAL_SEARCH>(
         self, weight, density, hist, bin_edges);
   } else {
diff --git a/aten/src/ATen/native/mps/operations/Indexing.mm b/aten/src/ATen/native/mps/operations/Indexing.mm
index 30e36887611e4..1433f4ad9844c 100644
--- a/aten/src/ATen/native/mps/operations/Indexing.mm
+++ b/aten/src/ATen/native/mps/operations/Indexing.mm
@@ -3,6 +3,7 @@
 #define TORCH_ASSERT_ONLY_METHOD_OPERATORS
 #include <ATen/Dispatch_v2.h>
 #include <ATen/native/mps/OperationUtils.h>
+#include <ATen/native/mps/kernels/Indexing.h>
 
 #include <ATen/AccumulateType.h>
 #include <ATen/ExpandUtils.h>
@@ -25,26 +26,25 @@
 #include <ATen/Functions.h>
 #include <ATen/NativeFunctions.h>
 #else
-#include <ATen/ops/count_nonzero.h>
-#include <ATen/ops/count_nonzero_native.h>
+#include <ATen/native/IndexKernel.h>
 #include <ATen/ops/embedding_dense_backward_native.h>
 #include <ATen/ops/flip_native.h>
 #include <ATen/ops/index.h>
 #include <ATen/ops/index_add_native.h>
 #include <ATen/ops/index_copy_native.h>
-#include <ATen/ops/index_fill_native.h>
 #include <ATen/ops/index_put.h>
+#include <ATen/ops/index_reduce_native.h>
 #include <ATen/ops/index_select_native.h>
 #include <ATen/ops/masked_fill_native.h>
 #include <ATen/ops/masked_scatter_native.h>
 #include <ATen/ops/masked_select_native.h>
 #include <ATen/ops/nonzero.h>
 #include <ATen/ops/nonzero_native.h>
+#include <ATen/ops/nonzero_static_native.h>
+#include <ATen/ops/ones_like.h>
 #include <ATen/ops/view_as_real.h>
 #endif
 
-constexpr auto nonZeroMaxSize = 1UL << 24;
-
 namespace at::native {
 namespace mps {
 
@@ -289,150 +289,139 @@ static void index_put_kernel_mps(TensorIterator& iter,
   });
 }
 
-static Tensor nonzero_fallback(const Tensor& self) {
-  return at::nonzero(self.to("cpu")).to("mps");
-}
-
-static Tensor& nonzero_out_native_mps(const Tensor& self, Tensor& out_) {
+// Metal kernel-based nonzero using prefix-sum + scatter.
+// Step 1: Per-element exclusive prefix sum of nonzero flags + block totals.
+// Step 2: GPU prefix sum of block totals → block offsets + total count.
+// Host (optional):   Read back total count, allocate output, unless max_element is provided
+// Step 3: Scatter multi-dimensional indices into the output.
+static void nonzero_impl_mps(const Tensor& self, Tensor& out_, std::optional<int64_t> max_elements) {
   using namespace mps;
 
-  int64_t nDim = self.dim();
+  TORCH_CHECK(self.numel() < std::numeric_limits<int>::max(),
+              "nonzero is not supported for tensors with more than INT_MAX elements, "
+              "See https://github.com/pytorch/pytorch/issues/51871");
+  TORCH_CHECK(out_.dtype() == at::kLong, "Expected output type to be Long, but got ", out_.dtype());
+  TORCH_CHECK(self.device() == out_.device(),
+              "expected self and out to be on the same device, but got out on ",
+              out_.device(),
+              " and self on ",
+              self.device());
+  TORCH_CHECK(out_.is_mps());
+
+  Tensor input = self.contiguous();
+  const int64_t nDim = self.dim();
+  const auto numel = static_cast<uint32_t>(input.numel());
+  const auto type_str = scalarToMetalTypeString(input);
   MPSStream* stream = getCurrentMPSStream();
-  using CachedGraph = MPSUnaryCachedGraph;
 
+  auto pso_step1 = lib.getPipelineStateForFunc(fmt::format("count_nonzero_prefix_sum_{}", type_str));
+  auto pso_step2 = lib.getPipelineStateForFunc("prefix_sum_blocks");
+  auto pso_step3 = lib.getPipelineStateForFunc(fmt::format("scatter_nonzero_indices_{}", type_str));
+  TORCH_INTERNAL_ASSERT([pso_step1 maxTotalThreadsPerThreadgroup] == [pso_step3 maxTotalThreadsPerThreadgroup],
+                        "nonzero: step 1 and step 3 threadgroup sizes must match");
+
+  uint32_t threads_per_group = static_cast<uint32_t>([pso_step1 maxTotalThreadsPerThreadgroup]);
+  uint32_t num_blocks = (numel + threads_per_group - 1) / threads_per_group;
+
+  auto tmp = at::empty({input.numel() + 2 * num_blocks + 1}, input.options().dtype(kInt));
+  Tensor prefix_buf = tmp.slice(0, 0, numel);
+  Tensor block_sums_buf = tmp.slice(0, numel, numel + num_blocks);
+  Tensor block_offsets_buf = tmp.slice(0, numel + num_blocks, numel + 2 * num_blocks);
+  Tensor total_nonzero_buf = tmp.slice(0, numel + 2 * num_blocks, numel + 2 * num_blocks + 1);
+
+  // Steps 1+2: compute prefix sums and block offsets entirely on GPU
   dispatch_sync_with_rethrow(stream->queue(), ^() {
-    stream->synchronize(SyncType::COMMIT_AND_WAIT);
+    @autoreleasepool {
+      auto computeEncoder = stream->commandEncoder();
+
+      [computeEncoder setComputePipelineState:pso_step1];
+      mtl_setArgs(computeEncoder, input, prefix_buf, block_sums_buf);
+      mtl_dispatch1DJob(computeEncoder, pso_step1, numel);
+
+      [computeEncoder setComputePipelineState:pso_step2];
+      mtl_setArgs(computeEncoder, block_sums_buf, block_offsets_buf, total_nonzero_buf, num_blocks);
+      uint32_t tg_size_blocks = std::min(1024u, ((num_blocks + 31) / 32) * 32);
+      [computeEncoder dispatchThreads:MTLSizeMake(tg_size_blocks, 1, 1)
+                threadsPerThreadgroup:MTLSizeMake(tg_size_blocks, 1, 1)];
+    }
   });
-  int64_t total_nonzero = at::count_nonzero(self).item<int64_t>();
-  at::native::resize_output(out_, {total_nonzero, nDim});
-  if (out_.numel() == 0) {
-    return out_;
-  }
 
-  bool contiguous_output = !needsGather(out_);
-  Tensor out = out_;
-  if (!contiguous_output) {
-    out = at::empty_like(out_, MemoryFormat::Contiguous);
+  if (!max_elements) {
+    // Dynamic path: sync to learn output size
+    const int64_t total_nonzero = total_nonzero_buf.item<int>();
+    at::native::resize_output(out_, {total_nonzero, nDim});
+    max_elements = total_nonzero;
   }
 
-  @autoreleasepool {
-    std::string key = "nonzero_out_native_mps" + getTensorsStringKey(self);
-    auto cachedGraph = LookUpOrCreateCachedGraph<CachedGraph>(key, [&](auto mpsGraph, auto newCachedGraph) {
-      MPSGraphTensor* inputTensor = mpsGraphRankedPlaceHolder(mpsGraph, self);
+  if (out_.numel() == 0) {
+    return;
+  }
 
-      MPSGraphTensor* outputTensor = [mpsGraph nonZeroIndicesOfTensor:inputTensor name:nil];
+  bool contiguous_output = out_.is_contiguous();
+  Tensor out = contiguous_output ? out_ : at::empty_like(out_, MemoryFormat::Contiguous);
 
-      newCachedGraph->inputTensor_ = inputTensor;
-      newCachedGraph->outputTensor_ = outputTensor;
-    });
+  int ndim_int = static_cast<int>(nDim);
+  int max_entries = static_cast<int>(*max_elements);
 
-    Placeholder selfPlaceholder = Placeholder(cachedGraph->inputTensor_, self);
-    Placeholder outputPlaceholder = Placeholder(cachedGraph->outputTensor_, out);
-    auto feeds = dictionaryFromPlaceholders(selfPlaceholder);
-    runMPSGraph(stream, cachedGraph->graph(), feeds, outputPlaceholder);
-  }
+  // Step 3: scatter indices, capped at max_entries
+  dispatch_sync_with_rethrow(stream->queue(), ^() {
+    @autoreleasepool {
+      auto computeEncoder = stream->commandEncoder();
+      [computeEncoder setComputePipelineState:pso_step3];
+      mtl_setArgs(computeEncoder, input, prefix_buf, out, ndim_int, input.sizes(), block_offsets_buf, max_entries);
+      mtl_dispatch1DJob(computeEncoder, pso_step3, numel);
+    }
+  });
 
   if (!contiguous_output) {
     out_.copy_(out);
   }
-
-  return out_;
 }
 
 Tensor& nonzero_out_mps(const Tensor& self, Tensor& out_) {
-  if (self.is_complex()) {
-    TORCH_WARN_ONCE("MPS: nonzero op is not supported for complex datatypes. ",
-                    "Falling back on CPU. This may have performance implications.");
-    Tensor out_fallback = nonzero_fallback(self);
-    at::native::resize_output(out_, out_fallback.sizes());
-    out_.copy_(out_fallback);
-    return out_;
-  }
-
   int64_t nDim = self.dim();
   if (self.numel() == 0) {
     at::native::resize_output(out_, {0, nDim});
     return out_;
   }
 
-  using namespace mps;
-  const uint32_t maxDimensions = 16;
-
-  TORCH_CHECK(self.numel() < std::numeric_limits<int>::max(),
-              "nonzero is not supported for tensors with more than INT_MAX elements, \
-  See https://github.com/pytorch/pytorch/issues/51871");
-  TORCH_CHECK(
-      out_.dtype() == at::kLong, "Expected object of scalar type ", at::kLong, " as out, but got ", out_.dtype());
-  TORCH_CHECK(self.device() == out_.device(),
-              "expected self and out to be on the same device, but got out on ",
-              out_.device(),
-              " and self on ",
-              self.device());
-  TORCH_CHECK(self.dim() <= maxDimensions, "nonzero is not supported for tensor with more than ", 16, " dimensions");
-  TORCH_CHECK(out_.is_mps());
+  nonzero_impl_mps(self, out_, std::nullopt);
+  return out_;
+}
 
-  if (!is_macos_13_or_newer(MacOSVersion::MACOS_VER_15_0_PLUS) &&
-      (self.numel() >= nonZeroMaxSize || self.is_complex())) {
-    TORCH_WARN_ONCE("MPS: nonzero op is not natively supported for the provided input on MacOS14",
-                    "Falling back on CPU. This may have performance implications.",
-                    "See github.com/pytorch/pytorch/issues/122916 for further info");
-    Tensor out_fallback = nonzero_fallback(self);
-    at::native::resize_output(out_, out_fallback.sizes());
-    out_.copy_(out_fallback);
-    return out_;
-  }
+Tensor nonzero_mps(const Tensor& self) {
+  Tensor out = at::empty({0}, self.options().dtype(kLong));
+  return nonzero_out_mps(self, out);
+}
 
-  MPSStream* stream = getCurrentMPSStream();
-  using CachedGraph = MPSUnaryCachedGraph;
+Tensor& nonzero_static_out_mps(const Tensor& self, int64_t size, int64_t fill_value, Tensor& result) {
+  TORCH_CHECK(size >= 0, "nonzero_static: 'size' must be an non-negative integer");
 
-  dispatch_sync_with_rethrow(stream->queue(), ^() {
-    stream->synchronize(SyncType::COMMIT_AND_WAIT);
-  });
-  int64_t total_nonzero = at::count_nonzero(self).item<int64_t>();
-  at::native::resize_output(out_, {total_nonzero, nDim});
-  if (out_.numel() == 0) {
-    return out_;
+  int64_t nDim = self.dim();
+  if (result.dim() != 2 || result.size(0) != size || result.size(1) != nDim) {
+    at::native::resize_output(result, {size, nDim});
   }
 
-  bool contiguous_output = !needsGather(out_);
-  Tensor out = out_;
-  if (!contiguous_output) {
-    out = at::empty_like(out_, MemoryFormat::Contiguous);
+  if (result.size(0) == 0 || result.size(1) == 0) {
+    return result;
   }
 
-  @autoreleasepool {
-    std::string key = "nonzero_out_native_mps" + getTensorsStringKey(self);
-    auto cachedGraph = LookUpOrCreateCachedGraph<CachedGraph>(key, [&](auto mpsGraph, auto newCachedGraph) {
-      MPSGraphTensor* inputTensor = mpsGraphRankedPlaceHolder(mpsGraph, self);
-
-      MPSGraphTensor* outputTensor = [mpsGraph nonZeroIndicesOfTensor:inputTensor name:nil];
+  result.fill_(fill_value);
 
-      newCachedGraph->inputTensor_ = inputTensor;
-      newCachedGraph->outputTensor_ = outputTensor;
-    });
-
-    Placeholder selfPlaceholder = Placeholder(cachedGraph->inputTensor_, self);
-    Placeholder outputPlaceholder = Placeholder(cachedGraph->outputTensor_, out);
-    auto feeds = dictionaryFromPlaceholders(selfPlaceholder);
-    runMPSGraph(stream, cachedGraph->graph(), feeds, outputPlaceholder);
-  }
-
-  if (!contiguous_output) {
-    out_.copy_(out);
+  if (self.numel() == 0) {
+    return result;
   }
 
-  return out_;
+  nonzero_impl_mps(self, result, size);
+  return result;
 }
 
-Tensor nonzero_mps(const Tensor& self) {
-  if (self.is_complex()) {
-    TORCH_WARN_ONCE("MPS: nonzero op is not supported for complex datatypes ",
-                    "Falling back on CPU. This may have performance implications.");
-    return nonzero_fallback(self);
-  }
-
-  Tensor out = at::empty({0}, self.options().dtype(kLong));
-  return nonzero_out_mps(self, out);
+Tensor nonzero_static_mps(const Tensor& self, int64_t size, int64_t fill_value) {
+  TORCH_CHECK(size >= 0, "nonzero_static: 'size' must be an non-negative integer");
+  int64_t nDim = self.dim();
+  auto result = at::empty({size, nDim}, at::TensorOptions().dtype(kLong).device(kMPS));
+  nonzero_static_out_mps(self, size, fill_value, result);
+  return result;
 }
 
 Tensor masked_select_mps(const Tensor& self, const Tensor& mask) {
@@ -537,6 +526,13 @@ Tensor flip_mps(const Tensor& self, IntArrayRef dims) {
 
   auto casted_type = isFloatingType(source.scalar_type()) ? ScalarType::Float : ScalarType::Int;
 
+  bool needs_gather = needsGather(result);
+  Tensor output;
+  if (needs_gather) {
+    output = at::empty_like(result, MemoryFormat::Contiguous);
+    output.copy_(result);
+  }
+
   struct CachedGraph : public MPSCachedGraph {
     CachedGraph(MPSGraph* graph) : MPSCachedGraph(graph) {}
     MPSGraphTensor* inputTensor_ = nil;
@@ -582,7 +578,7 @@ Tensor flip_mps(const Tensor& self, IntArrayRef dims) {
     Placeholder selfPlaceholder = Placeholder(cachedGraph->inputTensor_, self);
     Placeholder indexPlaceholder = Placeholder(cachedGraph->indexTensor_, index);
     Placeholder sourcePlaceholder = Placeholder(cachedGraph->sourceTensor_, source);
-    Placeholder outputPlaceholder = Placeholder(cachedGraph->outputTensor_, result);
+    Placeholder outputPlaceholder = Placeholder(cachedGraph->outputTensor_, needs_gather ? output : result);
     MPSScalar alpha_scalar = getMPSScalar(alpha, casted_type);
 
     NSDictionary<MPSGraphTensor*, MPSGraphTensorData*>* feeds = @{
@@ -593,6 +589,10 @@ Tensor flip_mps(const Tensor& self, IntArrayRef dims) {
     };
     runMPSGraph(stream, cachedGraph->graph(), feeds, outputPlaceholder);
   }
+
+  if (needs_gather) {
+    result.copy_(output);
+  }
 }
 
 Tensor index_select_mps(const Tensor& self, int64_t dim, const Tensor& index) {
@@ -702,6 +702,123 @@ Tensor index_select_mps(const Tensor& self, int64_t dim, const Tensor& index) {
   return output;
 }
 
+static inline ReductionType index_reduce_type(const std::string_view& reduce) {
+  if (reduce == "prod") {
+    return ReductionType::PROD;
+  } else if (reduce == "mean") {
+    return ReductionType::MEAN;
+  } else if (reduce == "amax") {
+    return ReductionType::MAX;
+  } else if (reduce == "amin") {
+    return ReductionType::MIN;
+  } else {
+    TORCH_CHECK(false, "reduce argument must be either prod, mean, amax or amin, got ", reduce, ".");
+  }
+}
+
+template <typename scalar_t>
+static inline scalar_t highest_value() {
+  if constexpr (std::numeric_limits<scalar_t>::has_infinity) {
+    return std::numeric_limits<scalar_t>::infinity();
+  } else {
+    return std::numeric_limits<scalar_t>::max();
+  }
+}
+
+template <typename scalar_t>
+static inline scalar_t lowest_value() {
+  if constexpr (std::numeric_limits<scalar_t>::has_infinity) {
+    return -std::numeric_limits<scalar_t>::infinity();
+  } else {
+    return std::numeric_limits<scalar_t>::lowest();
+  }
+}
+
+template <typename scalar_t>
+static inline scalar_t index_reduce_init_value(ReductionType reduction_type) {
+  if (reduction_type == ReductionType::PROD) {
+    return 1;
+  } else if (reduction_type == ReductionType::MEAN) {
+    return 0;
+  } else if (reduction_type == ReductionType::MAX) {
+    return lowest_value<scalar_t>();
+  } else if (reduction_type == ReductionType::MIN) {
+    return highest_value<scalar_t>();
+  } else {
+    TORCH_INTERNAL_ASSERT(false, "reduction type not supported");
+  }
+}
+
+TORCH_IMPL_FUNC(index_reduce_mps_out)
+(const Tensor& self,
+ int64_t dim,
+ const Tensor& index,
+ const Tensor& source,
+ const std::string_view reduce,
+ bool include_self,
+ const Tensor& result) {
+  TORCH_WARN_ONCE("index_reduce() is in beta and the API may change at any time.");
+  TORCH_CHECK(self.scalar_type() != c10::kLong, "index_reduce for MPS does not support torch.long dtype");
+  TORCH_CHECK(self.scalar_type() != c10::kComplexFloat, "index_reduce for MPS does not support torch.cfloat dtype");
+
+  auto reduction_type = index_reduce_type(reduce);
+
+  if (!result.is_same(self)) {
+    result.copy_(self);
+  }
+
+  if (!include_self) {
+    AT_DISPATCH_ALL_TYPES_AND2(at::ScalarType::Half,
+                               at::ScalarType::BFloat16,
+                               result.scalar_type(),
+                               "index_reduce_func_mps_exclude_input_init",
+                               [&] {
+                                 scalar_t init_val = index_reduce_init_value<scalar_t>(reduction_type);
+                                 result.index_fill_(dim, index.to(at::ScalarType::Long), init_val);
+                               });
+  }
+
+  IndexReduceParams params;
+  params.index_stride = index.stride(0);
+  params.reduce_dim = dim;
+  params.ndim = result.dim();
+
+  for (const auto dim : c10::irange(result.dim())) {
+    params.self_strides[dim] = result.stride(dim);
+    params.self_sizes[dim] = result.size(dim);
+    params.source_strides[dim] = source.stride(dim);
+    params.source_sizes[dim] = source.size(dim);
+  }
+
+  MPSStream* stream = getCurrentMPSStream();
+
+  auto num_threads = source.numel();
+
+  dispatch_sync_with_rethrow(stream->queue(), ^() {
+    @autoreleasepool {
+      id<MTLComputeCommandEncoder> compute_encoder = stream->commandEncoder();
+      auto pipeline_state = mps::lib.getPipelineStateForFunc(fmt::format(
+          "index_reduce_{}_{}_{}", reduce, mps::scalarToMetalTypeString(result), mps::scalarToMetalTypeString(index)));
+      getMPSProfiler().beginProfileKernel(pipeline_state, "index_reduce", {result, index, source});
+      [compute_encoder setComputePipelineState:pipeline_state];
+      mps::mtl_setArgs(compute_encoder, result, index, source, params);
+      mps::mtl_dispatch1DJob(compute_encoder, pipeline_state, num_threads);
+      getMPSProfiler().endProfileKernel(pipeline_state);
+    }
+  });
+
+  if (reduction_type == ReductionType::MEAN) {
+    auto counts = include_self ? at::ones_like(result) : at::zeros_like(result);
+    counts.index_add_(dim, index, at::ones_like(source));
+    counts.masked_fill_(counts.eq(0), 1);
+    if (result.is_floating_point() || result.is_complex()) {
+      result.div_(counts);
+    } else {
+      result.div_(counts, "floor");
+    }
+  }
+}
+
 Tensor& masked_fill__mps(Tensor& self, const Tensor& mask, const Scalar& value) {
   using namespace mps;
 
@@ -850,6 +967,8 @@ Tensor embedding_dense_backward_mps(const Tensor& grad_,
   TORCH_CHECK(mask.scalar_type() == ScalarType::Byte || mask.scalar_type() == ScalarType::Bool,
               "masked_scatter: expected BoolTensor or ByteTensor for mask");
 
+  bool was_scalar = self.dim() == 0;
+
   auto mask_temp =
       (mask.dim() == 0) ? c10::MaybeOwned<Tensor>::owned(mask.unsqueeze(0)) : c10::MaybeOwned<Tensor>::borrowed(mask);
   auto self_temp =
@@ -862,6 +981,9 @@ Tensor embedding_dense_backward_mps(const Tensor& grad_,
   auto indices =
       at::native::expandTensors(*std::get<1>(mask_self_expanded),
                                 c10::List<std::optional<at::Tensor>>({*std::move(std::get<0>(mask_self_expanded))}));
+
+  TORCH_CHECK(indices[0].numel() <= source.numel(), "Number of elements of source < number of ones in mask");
+
   // next broadcast all index tensors together
   try {
     indices = at::expand_outplace(indices);
@@ -879,104 +1001,165 @@ Tensor embedding_dense_backward_mps(const Tensor& grad_,
   for (const auto index : indices) {
     final_indices.push_back(index);
   }
-  return at::index_put_out(self, *std::get<1>(mask_self_expanded), final_indices, source.resize_(indices[0].numel()));
-}
-
-Tensor& index_fill_mps_(Tensor& self, int64_t dim, const Tensor& index, const Tensor& source) {
-  using namespace mps;
-  MPSStream* stream = getCurrentMPSStream();
-  auto num_indices = index.numel();
-  dim = maybe_wrap_dim(dim, self.dim());
-
-  // Checks
-  TORCH_CHECK_INDEX(index.dim() <= 1, "index_fill_(): Index is supposed to be a vector");
-  TORCH_CHECK(index.scalar_type() == ScalarType::Long || index.scalar_type() == ScalarType::Int,
-              "index_fill_(): Expected dtype int32 or int64 for index");
-  TORCH_CHECK(dim == 0 || dim < self.dim(), "index_fill_(): Indexing dim ", dim, " is out of bounds of tensor");
-  TORCH_CHECK(self.is_complex() || !source.is_complex(),
-              "index_fill_(): Converting complex Scalar to non-complex type is not supported");
-  // MPS.scatter crashes if used with complex dtypes
-  TORCH_CHECK(!c10::isComplexType(self.scalar_type()), "index_fill_(): Complex types are yet not supported");
 
-  // Empty index
-  if (num_indices == 0) {
-    return self;
+  at::index_put_out(
+      self, *std::get<1>(mask_self_expanded), final_indices, source.flatten().narrow(0, 0, indices[0].numel()));
+  if (was_scalar) {
+    self.squeeze_();
   }
+  return self;
+}
 
-  // Scalar input
-  if (self.dim() == 0 && self.numel() == 1) {
-    return self.copy_(source);
+static void index_fill_mps_kernel(TensorIterator& iter,
+                                  int64_t dim,
+                                  int64_t self_dim_size,
+                                  int64_t self_dim_stride,
+                                  const Scalar& source) {
+  if (iter.numel() == 0) {
+    return;
   }
+  using namespace mps;
 
-  // Derive from MPSCachedGraph
-  struct CachedGraph : public MPSCachedGraph {
-    CachedGraph(MPSGraph* graph) : MPSCachedGraph(graph) {}
-    MPSGraphTensor* inputTensor_ = nil;
-    MPSGraphTensor* indexTensor_ = nil;
-    MPSGraphTensor* updateTensor_ = nil;
-    MPSGraphTensor* outputTensor_ = nil;
-  };
-
-  auto inputType = getMPSDataType(self);
-  auto sourceType = getMPSDataType(source);
-  if (inputType == MPSDataTypeUInt8 || inputType == MPSDataTypeBool) {
-    inputType = MPSDataTypeInt8;
-  }
+  // Reconstruct original self from the restrided iterator output.
+  // iter.tensor(0) has stride[dim]=0 and size[dim]=index_numel; restore originals.
+  const Tensor& self_rs = iter.tensor(0);
+  auto self_sizes = self_rs.sizes().vec();
+  auto self_strides = self_rs.strides().vec();
+  self_sizes[dim] = self_dim_size;
+  self_strides[dim] = self_dim_stride;
+  Tensor self = self_rs.as_strided(self_sizes, self_strides, self_rs.storage_offset());
+
+  // Reconstruct original 1D index from the restrided iterator input.
+  // iter.tensor(1) has shape [..., index_numel, ...] with stride[dim]=original stride.
+  const Tensor& idx_rs = iter.tensor(1);
+  Tensor index = idx_rs.as_strided({idx_rs.size(dim)}, {idx_rs.stride(dim)}, idx_rs.storage_offset());
+
+  const bool is_dense = self.is_contiguous() && index.is_contiguous();
+  const auto type_str = scalarToMetalTypeString(self);
+  const int64_t dim_size = self_dim_size;
+  const int64_t indices_numel = index.numel();
+
+  // For large index counts, use a two-pass mask approach: mark which dim-positions
+  // to fill, then iterate over all elements checking the mask. This avoids write
+  // conflicts from duplicate indices and produces cache-friendly sequential writes.
+  // Threshold: ≥6% fill rate (indices_numel × 16 ≥ dim_size).  Empirically this
+  // is where the mask's O(numel) sequential scan beats scatter's growing write-
+  // conflict overhead.
+  const bool use_mask = (indices_numel * 16 >= dim_size);
 
-  std::vector<int64_t> source_shape(self.sizes().begin(), self.sizes().end());
-  source_shape[dim] = index.numel();
-  auto expanded_source = source.expand(source_shape);
+  auto stream = getCurrentMPSStream();
+  if (use_mask) {
+    // Use at::empty (not at::zeros) to avoid a blit encoder fill that would
+    // force a blit→compute encoder switch, adding latency per call.
+    // The mask is zeroed via a compute kernel inside the dispatch block instead.
+    // TODO: Remove me after https://github.com/pytorch/pytorch/issues/175859 is fixed
+    auto mask = at::empty({dim_size}, at::TensorOptions().dtype(at::kBool).device(self.device()));
+    auto zeroMaskPSO = lib.getPipelineStateForFunc("index_fill_zero_mask");
+    auto setMaskPSO = lib.getPipelineStateForFunc("index_fill_set_mask");
+    auto fillMaskPSO = lib.getPipelineStateForFunc(
+        fmt::format("index_fill_{}_from_mask_{}", is_dense ? "dense" : "strided", type_str));
+
+    c10::SmallVector<int64_t> all_sizes, all_strides;
+    if (!is_dense) {
+      for (int64_t d = 0; d < self.dim(); d++) {
+        all_sizes.push_back(self.size(d));
+        all_strides.push_back(self.stride(d));
+      }
+    }
 
-  @autoreleasepool {
-    std::string key =
-        "index_fill_mps_" + getTensorsStringKey({self, index, expanded_source}) + ":" + std::to_string(dim);
-    auto cachedGraph = LookUpOrCreateCachedGraph<CachedGraph>(key, [&](auto mpsGraph, auto newCachedGraph) {
-      MPSGraphTensor* inputTensor = mpsGraphRankedPlaceHolder(mpsGraph, inputType, getMPSShape(self));
-      MPSGraphTensor* indexTensor = mpsGraphRankedPlaceHolder(mpsGraph, index);
-      MPSGraphTensor* updateTensor = mpsGraphRankedPlaceHolder(mpsGraph, expanded_source);
-      MPSGraphTensor* castedUpdateTensor = updateTensor;
-      if (inputType != sourceType) {
-        castedUpdateTensor = castMPSTensor(mpsGraph, updateTensor, inputType);
+    dispatch_sync_with_rethrow(stream->queue(), ^() {
+      @autoreleasepool {
+        auto computeEncoder = stream->commandEncoder();
+        auto mpsScalar = getMPSScalar(source, self.scalar_type());
+        long indices_stride = index.stride(0);
+
+        // Pass 0: zero the mask entirely within the compute encoder.
+        [computeEncoder setComputePipelineState:zeroMaskPSO];
+        mtl_setArgs(computeEncoder, mask);
+        mtl_dispatch1DJob(computeEncoder, zeroMaskPSO, dim_size);
+
+        [computeEncoder setComputePipelineState:setMaskPSO];
+        mtl_setArgs(computeEncoder, mask, index);
+        mtl_setArgs<2>(computeEncoder, dim_size, indices_stride);
+        mtl_dispatch1DJob(computeEncoder, setMaskPSO, indices_numel);
+
+        [computeEncoder setComputePipelineState:fillMaskPSO];
+        mtl_setArgs(computeEncoder, self, mask);
+        mtl_setBytes(computeEncoder, mpsScalar, 2);
+        if (is_dense) {
+          // 3D dispatch: (inner, dim, outer) avoids integer division in the kernel.
+          // Compute inner_size from size products, not stride(dim): for size-1 dims,
+          // is_contiguous() returns true regardless of stride, so stride(dim) can
+          // exceed the actual inner element count, causing outer_size=0 (no threads).
+          uint32_t dim_size_u = static_cast<uint32_t>(dim_size);
+          uint32_t inner_size = 1;
+          for (int64_t d = dim + 1; d < self.dim(); d++) {
+            inner_size *= static_cast<uint32_t>(self.size(d));
+          }
+          uint32_t outer_size = static_cast<uint32_t>(self.numel() / ((int64_t)dim_size * inner_size));
+          mtl_setArgs<3>(computeEncoder, dim_size_u, inner_size);
+          NSUInteger maxTG = [fillMaskPSO maxTotalThreadsPerThreadgroup];
+          // Fill threadgroup across dimensions to ensure adequate occupancy.
+          // When inner_size is small (e.g. 1 for dim=last), use y-dim threads.
+          NSUInteger tgX = std::min((NSUInteger)inner_size, maxTG);
+          NSUInteger tgY = std::min((NSUInteger)dim_size_u, maxTG / tgX);
+          NSUInteger tgZ = std::min((NSUInteger)outer_size, maxTG / (tgX * tgY));
+          [computeEncoder dispatchThreads:MTLSizeMake(inner_size, dim_size_u, outer_size)
+                    threadsPerThreadgroup:MTLSizeMake(tgX, tgY, tgZ)];
+        } else {
+          uint32_t dim_u = static_cast<uint32_t>(dim);
+          uint32_t ndim_u = static_cast<uint32_t>(self.dim());
+          mtl_setArgs<3>(computeEncoder, all_sizes, all_strides, dim_u, ndim_u);
+          mtl_dispatch1DJob(computeEncoder, fillMaskPSO, self.numel());
+        }
       }
-      MPSGraphTensor* outputTensor = [mpsGraph scatterWithDataTensor:inputTensor
-                                                       updatesTensor:castedUpdateTensor
-                                                       indicesTensor:indexTensor
-                                                                axis:(NSInteger)dim
-                                                                mode:MPSGraphScatterModeSet
-                                                                name:nil];
-      newCachedGraph->inputTensor_ = inputTensor;
-      newCachedGraph->indexTensor_ = indexTensor;
-      newCachedGraph->updateTensor_ = updateTensor;
-      newCachedGraph->outputTensor_ = outputTensor;
     });
+  } else {
+    // Scatter: one thread per (index, slice-element) pair.
+    const int64_t slice_numel = self.numel() / dim_size;
+    auto indexFillPSO =
+        lib.getPipelineStateForFunc(fmt::format("index_fill_{}_{}", is_dense ? "dense" : "strided", type_str));
+
+    c10::SmallVector<int64_t> slice_sizes, slice_out_strides;
+    if (!is_dense) {
+      for (int64_t d = 0; d < self.dim(); d++) {
+        if (d != dim) {
+          slice_sizes.push_back(self.size(d));
+          slice_out_strides.push_back(self.stride(d));
+        }
+      }
+    }
 
-    Placeholder selfPlaceholder = Placeholder(cachedGraph->inputTensor_,
-                                              self,
-                                              /*mpsShape=*/nullptr,
-                                              /*gatherTensorData=*/true,
-                                              /*dataType=*/inputType);
-    Placeholder indexPlaceholder = Placeholder(cachedGraph->indexTensor_, index);
-    Placeholder updatePlaceholder = Placeholder(cachedGraph->updateTensor_,
-                                                expanded_source,
-                                                /*mpsShape=*/nullptr,
-                                                /*gatherTensorData=*/true,
-                                                /*dataType=*/sourceType);
-    Placeholder outputPlaceholder = Placeholder(cachedGraph->outputTensor_,
-                                                self,
-                                                /*mpsShape=*/nullptr,
-                                                /*gatherTensorData=*/false,
-                                                /*dataType=*/inputType);
-
-    auto feeds = dictionaryFromPlaceholders(selfPlaceholder, indexPlaceholder, updatePlaceholder);
-    runMPSGraph(stream, cachedGraph->graph(), feeds, outputPlaceholder);
+    dispatch_sync_with_rethrow(stream->queue(), ^() {
+      @autoreleasepool {
+        auto computeEncoder = stream->commandEncoder();
+        auto mpsScalar = getMPSScalar(source, self.scalar_type());
+        [computeEncoder setComputePipelineState:indexFillPSO];
+        mtl_setArgs(computeEncoder, self, index);
+        mtl_setBytes(computeEncoder, mpsScalar, 2);
+        if (is_dense) {
+          long dim_stride = self.stride(dim);
+          mtl_setArgs<3>(computeEncoder, dim_size, dim_stride, slice_numel);
+        } else {
+          long dim_out_stride = self.stride(dim);
+          long indices_stride = index.stride(0);
+          uint32_t slice_ndim = static_cast<uint32_t>(self.dim() - 1);
+          mtl_setArgs<3>(computeEncoder,
+                         dim_size,
+                         dim_out_stride,
+                         slice_sizes,
+                         slice_out_strides,
+                         slice_ndim,
+                         slice_numel,
+                         indices_stride);
+        }
+        mtl_dispatch1DJob(computeEncoder, indexFillPSO, indices_numel * slice_numel);
+      }
+    });
   }
-  return self;
-}
-
-Tensor& index_fill_mps_(Tensor& self, int64_t dim, const Tensor& index, const Scalar& source) {
-  return self.index_fill_(dim, index, mps::wrapped_scalar_tensor_mps(source, self.device()));
 }
 
 REGISTER_DISPATCH(index_stub, &mps::index_kernel_mps)
+REGISTER_DISPATCH(index_fill_stub, &index_fill_mps_kernel)
 REGISTER_DISPATCH(index_put_stub, &mps::index_put_kernel_mps)
 } // namespace at::native
diff --git a/aten/src/ATen/native/mps/operations/Lerp.mm b/aten/src/ATen/native/mps/operations/Lerp.mm
deleted file mode 100644
index 4d13b8b2466bd..0000000000000
--- a/aten/src/ATen/native/mps/operations/Lerp.mm
+++ /dev/null
@@ -1,49 +0,0 @@
-#define TORCH_ASSERT_ONLY_METHOD_OPERATORS
-#include <ATen/core/Tensor.h>
-#include <ATen/native/mps/OperationUtils.h>
-
-#ifndef AT_PER_OPERATOR_HEADERS
-#include <ATen/Functions.h>
-#include <ATen/NativeFunctions.h>
-#else
-#include <ATen/ops/add.h>
-#include <ATen/ops/lerp_native.h>
-#endif
-
-namespace at::native {
-TORCH_IMPL_FUNC(lerp_Tensor_mps)(const Tensor& self, const Tensor& end, const Tensor& weight, const Tensor& out) {
-  TORCH_CHECK(out.is_mps());
-  std::array<TensorArg, 4> args{{{out, "out", 0}, {self, "self", 1}, {end, "end", 2}, {weight, "weight", 3}}};
-  checkAllSameGPU(__func__, args);
-  using namespace mps;
-  struct CachedGraph : public MPSCachedGraph {
-    CachedGraph(MPSGraph* graph) : MPSCachedGraph(graph) {}
-    MPSGraphTensor* selfTensor_ = nil;
-    MPSGraphTensor* endTensor_ = nil;
-    MPSGraphTensor* weightTensor_ = nil;
-    MPSGraphTensor* outputTensor_ = nil;
-  };
-  @autoreleasepool {
-    std::string key = "lerp_Tensor_mps" + getTensorsStringKey({self, end, weight});
-    auto cachedGraph = LookUpOrCreateCachedGraph<CachedGraph>(key, [&](auto mpsGraph, auto graph) {
-      auto selfTensor = mpsGraphRankedPlaceHolder(mpsGraph, self);
-      auto endTensor = mpsGraphRankedPlaceHolder(mpsGraph, end);
-      auto weightTensor = mpsGraphRankedPlaceHolder(mpsGraph, weight);
-      auto distance = [mpsGraph subtractionWithPrimaryTensor:endTensor secondaryTensor:selfTensor name:nil];
-      auto weighedDistance = [mpsGraph multiplicationWithPrimaryTensor:weightTensor secondaryTensor:distance name:nil];
-      auto output = [mpsGraph additionWithPrimaryTensor:selfTensor secondaryTensor:weighedDistance name:nil];
-      graph->selfTensor_ = selfTensor;
-      graph->endTensor_ = endTensor;
-      graph->weightTensor_ = weightTensor;
-      graph->outputTensor_ = output;
-    });
-    auto selfPlaceholder = Placeholder(cachedGraph->selfTensor_, self);
-    auto endPlaceholder = Placeholder(cachedGraph->endTensor_, end);
-    auto weightPlaceholder = Placeholder(cachedGraph->weightTensor_, weight);
-    auto outputPlaceholder = Placeholder(cachedGraph->outputTensor_, out);
-    auto feeds = dictionaryFromPlaceholders(selfPlaceholder, endPlaceholder, weightPlaceholder);
-    runMPSGraph(getCurrentMPSStream(), cachedGraph->graph(), feeds, outputPlaceholder);
-  }
-}
-
-} // namespace at::native
diff --git a/aten/src/ATen/native/mps/operations/Linear.mm b/aten/src/ATen/native/mps/operations/Linear.mm
index 219086edd8e37..2ccb89bc4d702 100644
--- a/aten/src/ATen/native/mps/operations/Linear.mm
+++ b/aten/src/ATen/native/mps/operations/Linear.mm
@@ -76,8 +76,6 @@ Tensor _mps_linear(const Tensor& input, const Tensor& weight_arg, const std::opt
   TORCH_CHECK(input.is_mps(), "Tensor for argument input is on ", input.device(), " but expected on mps");
   TORCH_CHECK(supportedFloatingOrComplexType(weight_arg), "MPS device does not support linear for non-float weights");
   TORCH_CHECK(weight_arg.is_mps(), "Tensor for argument weight is on ", weight_arg.device(), " but expected on mps");
-  TORCH_CHECK((input.scalar_type() != kComplexFloat && input.scalar_type() != kComplexHalf),
-              "mps linear does not support complex types");
 
   const Tensor& bias = *(at::borrow_from_optional_tensor(bias_opt));
   const bool is_bias_defined = bias.defined();
@@ -117,8 +115,9 @@ Tensor _mps_linear(const Tensor& input, const Tensor& weight_arg, const std::opt
 
   // No-graph execution causes nonsense if these are non-contiguous.
   const bool is_contiguous = input.is_contiguous() && weight.is_contiguous() && bias.is_contiguous();
+  const bool is_complex = input.is_complex() || weight.is_complex() || (is_bias_defined && bias.is_complex());
 
-  if (is_macos_13_or_newer(MacOSVersion::MACOS_VER_15_0_PLUS) && is_contiguous) {
+  if (is_macos_13_or_newer(MacOSVersion::MACOS_VER_15_0_PLUS) && is_contiguous && !is_complex) {
     _mps_linear_nograph(input, weight, bias, output);
     // Squeeze last dim of 1D linear
     return weight_arg.dim() != 1 ? output : output.squeeze(-1);
@@ -206,8 +205,7 @@ static Tensor _mps_linear_backward_input(IntArrayRef input_size, const Tensor& g
     MPSGraphTensor* outputTensor_ = nil;
   };
 
-  Tensor output = at::empty(
-      input_size, grad_output.scalar_type(), std::nullopt, kMPS, std::nullopt, grad_output.suggest_memory_format());
+  Tensor output = at::empty(input_size, grad_output.options());
   TORCH_CHECK(output.is_mps());
   if (grad_output.numel() == 0) {
     return output;
@@ -275,18 +273,8 @@ static Tensor _mps_linear_backward_input(IntArrayRef input_size, const Tensor& g
   TORCH_CHECK(grad_output_reshaped.is_mps());
   TORCH_CHECK(input_reshaped.is_mps());
 
-  Tensor output = at::empty({grad_output_reshaped.size(1), input_reshaped.size(1)},
-                            grad_output.scalar_type(),
-                            std::nullopt,
-                            kMPS,
-                            std::nullopt,
-                            grad_output.suggest_memory_format());
-  Tensor bias = at::empty({grad_output_reshaped.size(1)},
-                          grad_output.scalar_type(),
-                          std::nullopt,
-                          kMPS,
-                          std::nullopt,
-                          grad_output.suggest_memory_format());
+  Tensor output = at::empty({grad_output_reshaped.size(1), input_reshaped.size(1)}, grad_output.options());
+  Tensor bias = at::empty({grad_output_reshaped.size(1)}, grad_output.options());
   TORCH_CHECK(output.is_mps());
   TORCH_CHECK(bias.is_mps());
 
diff --git a/aten/src/ATen/native/mps/operations/LinearAlgebra.mm b/aten/src/ATen/native/mps/operations/LinearAlgebra.mm
index d98134469ec04..6bfd7d1049c4b 100644
--- a/aten/src/ATen/native/mps/operations/LinearAlgebra.mm
+++ b/aten/src/ATen/native/mps/operations/LinearAlgebra.mm
@@ -17,6 +17,7 @@
 #include <ATen/Functions.h>
 #include <ATen/NativeFunctions.h>
 #else
+#include <ATen/ops/_cholesky_solve_helper_native.h>
 #include <ATen/ops/_linalg_solve_ex_native.h>
 #include <ATen/ops/addbmm_native.h>
 #include <ATen/ops/addmm_native.h>
@@ -31,6 +32,8 @@
 #include <ATen/ops/linalg_lu_factor_ex_native.h>
 #include <ATen/ops/linalg_lu_factor_native.h>
 #include <ATen/ops/linalg_lu_native.h>
+#include <ATen/ops/linalg_qr.h>
+#include <ATen/ops/linalg_qr_native.h>
 #include <ATen/ops/linalg_solve_triangular_native.h>
 #include <ATen/ops/lu_unpack.h>
 #include <ATen/ops/lu_unpack_native.h>
@@ -307,6 +310,20 @@ bool use_metal_mm(const Tensor& self, const Tensor& other, const Tensor& output)
   if (c10::isComplexType(self.scalar_type()) && self.size(1) > max_complex_inner_size) {
     return true;
   }
+  // Detect conditions that would trigger LORADOWN GEMV kernel with potential padding overflow
+  // See https://github.com/pytorch/pytorch/issues/178056
+  if (self.scalar_type() == at::ScalarType::Half && (self.size(0) <= 16 || other.size(1) <= 16) &&
+      self.stride(1) == 1 && other.stride(0) == 1) {
+    int64_t self_padding = self.stride(0) - self.size(1);
+    int64_t other_padding = other.stride(1) - other.size(0);
+
+    if (self_padding > 15 || other_padding > 15 || self_padding % 4 != 0 || other_padding % 4 != 0) {
+      TORCH_WARN_ONCE(
+          "MPS mm implementation has a known issue with this shape, dtype and slice. Dispatching to metal implementation instead. This may impact performance.");
+      return true;
+    }
+  }
+
   return !is_macos_14_4_or_newer &&
       (self.stride(0) > max_stride_size || self.stride(1) > max_stride_size || self.size(0) > max_stride_size ||
        self.size(1) > max_stride_size || other.stride(0) > max_stride_size || other.stride(1) > max_stride_size ||
@@ -395,34 +412,32 @@ static void linalg_lu_factor_ex_out_mps_impl(const Tensor& A,
   dispatch_sync_with_rethrow(mpsStream->queue(), ^() {
     @autoreleasepool {
       id<MTLCommandBuffer> commandBuffer = mpsStream->commandBuffer();
-      MPSMatrixDecompositionLU* filter = [[[MPSMatrixDecompositionLU alloc] initWithDevice:device
-                                                                                      rows:aRows
-                                                                                   columns:aCols] autorelease];
-
-      MPSMatrixDescriptor* sourceMatrixDesc = [MPSMatrixDescriptor matrixDescriptorWithRows:aRows
-                                                                                    columns:aCols
-                                                                                   matrices:1
-                                                                                   rowBytes:aCols * aElemSize
-                                                                                matrixBytes:aRows * aCols * aElemSize
-                                                                                   dataType:getMPSDataType(A_)];
-      MPSMatrixDescriptor* pivotsMatrixDesc = [MPSMatrixDescriptor matrixDescriptorWithRows:1
-                                                                                    columns:numPivots
-                                                                                   matrices:1
-                                                                                   rowBytes:numPivots * sizeof(uint32_t)
-                                                                                matrixBytes:numPivots * sizeof(uint32_t)
-                                                                                   dataType:MPSDataTypeUInt32];
+      auto filter = [[[MPSMatrixDecompositionLU alloc] initWithDevice:device rows:aRows columns:aCols] autorelease];
+
+      auto sourceMatrixDesc = [MPSMatrixDescriptor matrixDescriptorWithRows:aRows
+                                                                    columns:aCols
+                                                                   matrices:1
+                                                                   rowBytes:aCols * aElemSize
+                                                                matrixBytes:aRows * aCols * aElemSize
+                                                                   dataType:getMPSDataType(A_)];
+      auto pivotsMatrixDesc = [MPSMatrixDescriptor matrixDescriptorWithRows:1
+                                                                    columns:numPivots
+                                                                   matrices:1
+                                                                   rowBytes:numPivots * sizeof(uint32_t)
+                                                                matrixBytes:numPivots * sizeof(uint32_t)
+                                                                   dataType:MPSDataTypeUInt32];
 
       for (const auto i : c10::irange(batchSize)) {
         const uint64_t aBatchOffset = i * aRows * aCols;
-        MPSMatrix* sourceMatrix = [[[MPSMatrix alloc] initWithBuffer:aBuffer
-                                                              offset:(A_.storage_offset() + aBatchOffset) * aElemSize
-                                                          descriptor:sourceMatrixDesc] autorelease];
-        MPSMatrix* pivotIndices = [[[MPSMatrix alloc] initWithBuffer:getMTLBufferStorage(pivots_list[i])
-                                                              offset:0
-                                                          descriptor:pivotsMatrixDesc] autorelease];
-        MPSMatrix* solutionMatrix = [[[MPSMatrix alloc] initWithBuffer:aBuffer
-                                                                offset:(A_.storage_offset() + aBatchOffset) * aElemSize
-                                                            descriptor:sourceMatrixDesc] autorelease];
+        auto sourceMatrix = [[[MPSMatrix alloc] initWithBuffer:aBuffer
+                                                        offset:(A_.storage_offset() + aBatchOffset) * aElemSize
+                                                    descriptor:sourceMatrixDesc] autorelease];
+        auto pivotIndices = [[[MPSMatrix alloc] initWithBuffer:getMTLBufferStorage(pivots_list[i])
+                                                        offset:0
+                                                    descriptor:pivotsMatrixDesc] autorelease];
+        auto solutionMatrix = [[[MPSMatrix alloc] initWithBuffer:aBuffer
+                                                          offset:(A_.storage_offset() + aBatchOffset) * aElemSize
+                                                      descriptor:sourceMatrixDesc] autorelease];
         id<MTLBuffer> statusBuffer = getMTLBufferStorage(status_tensors[i]);
         [filter encodeToCommandBuffer:commandBuffer
                          sourceMatrix:sourceMatrix
@@ -490,8 +505,8 @@ static void linalg_solve_out_mps_impl(const Tensor& A,
 
   uint64_t numPivots = std::min(aRows, aCols);
   std::vector<int64_t> pivot_sizes(A_t.sizes().begin(), A_t.sizes().end() - 2);
-  info.fill_(0); // will be set to 1 during kernel if something fails
   resize_output(info, pivot_sizes);
+  info.fill_(0); // will be set by kernel on failure
   pivot_sizes.push_back(numPivots);
   resize_output(pivots, pivot_sizes);
 
@@ -499,23 +514,23 @@ static void linalg_solve_out_mps_impl(const Tensor& A,
     return;
   }
 
+  // Save original shape before flattening for the LU output
+  auto A_original_sizes = A_t.sizes().vec();
+
   if (A_t.dim() > 3) {
     A_t = A_t.flatten(0, -3);
   }
 
   uint64_t batchSize = (A_t.dim() > 2) ? A_t.size(0) : 1;
   std::vector<Tensor> status_tensors;
-  std::vector<Tensor> pivots_list;
-
   status_tensors.reserve(batchSize);
-  pivots_list.reserve(batchSize);
   for ([[maybe_unused]] const auto i : c10::irange(batchSize)) {
     status_tensors.push_back(at::zeros(1, kInt, std::nullopt, kMPS, std::nullopt));
-    pivots_list.push_back(at::zeros(numPivots, kInt, std::nullopt, kMPS, std::nullopt));
   }
 
-  resize_output(LU, A_t.sizes());
-  Tensor LU_ = LU;
+  // LU must keep the original (unflattened) shape for the backward pass
+  resize_output(LU, A_original_sizes);
+  Tensor LU_ = (LU.dim() > 3) ? LU.flatten(0, -3) : LU;
   if (!LU_.is_same(A_t)) {
     A_t = LU_.copy_(A_t);
   } else {
@@ -534,6 +549,7 @@ static void linalg_solve_out_mps_impl(const Tensor& A,
   id<MTLBuffer> luBuffer = getMTLBufferStorage(LU_);
   id<MTLBuffer> bBuffer = getMTLBufferStorage(B_t);
   id<MTLBuffer> resultBuffer = getMTLBufferStorage(result_t);
+  id<MTLBuffer> pivotsBuffer = getMTLBufferStorage(pivots);
 
   MPSStream* mpsStream = getCurrentMPSStream();
   id<MTLDevice> device = MPSDevice::getInstance()->device();
@@ -542,41 +558,37 @@ static void linalg_solve_out_mps_impl(const Tensor& A,
     @autoreleasepool {
       id<MTLCommandBuffer> commandBuffer = mpsStream->commandBuffer();
 
-      MPSMatrixDecompositionLU* lu_decomp = [[[MPSMatrixDecompositionLU alloc] initWithDevice:device
-                                                                                         rows:aRows
-                                                                                      columns:aCols] autorelease];
-
-      MPSMatrixSolveLU* solver = [[[MPSMatrixSolveLU alloc] initWithDevice:device
-                                                                 transpose:false
-                                                                     order:aRows
-                                                    numberOfRightHandSides:numberOfRightHandSides] autorelease];
-
-      MPSMatrixDescriptor* luMatrixDesc = [MPSMatrixDescriptor matrixDescriptorWithRows:aRows
-                                                                                columns:aCols
-                                                                               matrices:1
-                                                                               rowBytes:aCols * aElemSize
-                                                                            matrixBytes:aRows * aCols * aElemSize
-                                                                               dataType:getMPSDataType(LU_)];
-      MPSMatrixDescriptor* rhsMatrixDesc =
-          [MPSMatrixDescriptor matrixDescriptorWithRows:aRows
-                                                columns:numberOfRightHandSides
-                                               matrices:1
-                                               rowBytes:numberOfRightHandSides * aElemSize
-                                            matrixBytes:aRows * numberOfRightHandSides * aElemSize
-                                               dataType:getMPSDataType(B_t)];
-      MPSMatrixDescriptor* resultMatrixDesc =
-          [MPSMatrixDescriptor matrixDescriptorWithRows:aRows
-                                                columns:numberOfRightHandSides
-                                               matrices:1
-                                               rowBytes:numberOfRightHandSides * aElemSize
-                                            matrixBytes:aRows * numberOfRightHandSides * aElemSize
-                                               dataType:getMPSDataType(result_t)];
-      MPSMatrixDescriptor* pivotsMatrixDesc = [MPSMatrixDescriptor matrixDescriptorWithRows:1
-                                                                                    columns:numPivots
-                                                                                   matrices:1
-                                                                                   rowBytes:numPivots * sizeof(uint32_t)
-                                                                                matrixBytes:numPivots * sizeof(uint32_t)
-                                                                                   dataType:MPSDataTypeUInt32];
+      auto lu_decomp = [[[MPSMatrixDecompositionLU alloc] initWithDevice:device rows:aRows columns:aCols] autorelease];
+
+      auto solver = [[[MPSMatrixSolveLU alloc] initWithDevice:device
+                                                    transpose:false
+                                                        order:aRows
+                                       numberOfRightHandSides:numberOfRightHandSides] autorelease];
+
+      auto luMatrixDesc = [MPSMatrixDescriptor matrixDescriptorWithRows:aRows
+                                                                columns:aCols
+                                                               matrices:1
+                                                               rowBytes:aCols * aElemSize
+                                                            matrixBytes:aRows * aCols * aElemSize
+                                                               dataType:getMPSDataType(LU_)];
+      auto rhsMatrixDesc = [MPSMatrixDescriptor matrixDescriptorWithRows:aRows
+                                                                 columns:numberOfRightHandSides
+                                                                matrices:1
+                                                                rowBytes:numberOfRightHandSides * aElemSize
+                                                             matrixBytes:aRows * numberOfRightHandSides * aElemSize
+                                                                dataType:getMPSDataType(B_t)];
+      auto resultMatrixDesc = [MPSMatrixDescriptor matrixDescriptorWithRows:aRows
+                                                                    columns:numberOfRightHandSides
+                                                                   matrices:1
+                                                                   rowBytes:numberOfRightHandSides * aElemSize
+                                                                matrixBytes:aRows * numberOfRightHandSides * aElemSize
+                                                                   dataType:getMPSDataType(result_t)];
+      auto pivotsMatrixDesc = [MPSMatrixDescriptor matrixDescriptorWithRows:1
+                                                                    columns:numPivots
+                                                                   matrices:1
+                                                                   rowBytes:numPivots * sizeof(uint32_t)
+                                                                matrixBytes:numPivots * sizeof(uint32_t)
+                                                                   dataType:MPSDataTypeUInt32];
 
       for (const auto i : c10::irange(batchSize)) {
         const uint64_t batchOffsetA = i * aRows * aCols;
@@ -593,15 +605,14 @@ static void linalg_solve_out_mps_impl(const Tensor& A,
                                                            offset:(result_t.storage_offset() + batchOffsetB) * aElemSize
                                                        descriptor:resultMatrixDesc] autorelease];
 
-        MPSMatrix* mpsPivots = [[[MPSMatrix alloc] initWithBuffer:getMTLBufferStorage(pivots_list[i])
-                                                           offset:0
-                                                       descriptor:pivotsMatrixDesc] autorelease];
-        id<MTLBuffer> statusBuffer = getMTLBufferStorage(status_tensors[i]);
+        auto mpsPivots = [[[MPSMatrix alloc] initWithBuffer:pivotsBuffer
+                                                     offset:(pivots.storage_offset() + i * numPivots) * sizeof(uint32_t)
+                                                 descriptor:pivotsMatrixDesc] autorelease];
         [lu_decomp encodeToCommandBuffer:commandBuffer
                             sourceMatrix:mpsLU
                             resultMatrix:mpsLU
                             pivotIndices:mpsPivots
-                                  status:statusBuffer];
+                                  status:getMTLBufferStorage(status_tensors[i])];
         [solver encodeToCommandBuffer:commandBuffer
                          sourceMatrix:mpsLU
                   rightHandSideMatrix:mpsRHS
@@ -611,12 +622,13 @@ static void linalg_solve_out_mps_impl(const Tensor& A,
     }
   });
 
-  auto stacked_status = A.dim() > 2 ? at::stack(status_tensors) : status_tensors[0];
-  std::vector<int64_t> info_sizes(A.sizes().begin(), A.sizes().end() - 2);
-  info.copy_(stacked_status.view(info_sizes));
+  pivots.add_(1); // MPS is 0-based, PyTorch/LAPACK is 1-based
+
+  auto stacked_status = batchSize > 1 ? at::stack(status_tensors) : status_tensors[0];
+  info.copy_(stacked_status.view(info.sizes()));
 
   if (check_errors) {
-    for (const auto i : c10::irange(status_tensors.size())) {
+    for (const auto i : c10::irange(batchSize)) {
       int status = status_tensors[i].item<int>();
       TORCH_CHECK(status == 0,
                   "solve(): Linear solve failed at the ",
@@ -676,9 +688,11 @@ static void linalg_inv_ex_out_mps_impl(const Tensor& A, bool check_errors, const
 
   TORCH_CHECK(output.is_mps());
 
+  // Edge case behaviors must match _int_mm_out_cpu CPU implementation
   // Transpose inputs if needed
-  if (output.numel() == 0) {
-    return output;
+  // Outer or inner dimension is 0
+  if (output.numel() == 0 || self.size(1) == 0) {
+    return output.zero_();
   }
 
   // MPS matmul returns silently incorrect results if one of the matrix dimensions is greater than 2**15
@@ -698,13 +712,21 @@ static void linalg_inv_ex_out_mps_impl(const Tensor& A, bool check_errors, const
     });
     // MPS TODO:
     // Strided API doesn't play nice with complex data types (at least not in case of matmul).
+    // MPSGraph's matrixMultiplication produces incorrect results with stride-0 NDArray
+    // inputs on macOS < 26.4 (only every 16th row is computed). Contiguify such tensors
+    // by disabling the strided API so they go through the gather/clone path first.
+    // See https://github.com/pytorch/pytorch/issues/180201
+    static const bool is_macOS_26_4_or_newer = is_macos_13_or_newer(MacOSVersion::MACOS_VER_26_4_PLUS);
+    auto hasZeroStride = [](const Tensor& t) {
+      return std::ranges::any_of(t.strides(), [](auto s) { return s == 0; });
+    };
+    auto useStridedSelf = !isComplexType(self.scalar_type()) && (is_macOS_26_4_or_newer || !hasZeroStride(self));
+    auto useStridedOther = !isComplexType(other.scalar_type()) && (is_macOS_26_4_or_newer || !hasZeroStride(other));
     auto selfPlaceholder = self.numel() != 0
-        ? Placeholder(
-              cachedGraph->inputTensor_, self, nil, true, MPSDataTypeInvalid, !isComplexType(self.scalar_type()))
+        ? Placeholder(cachedGraph->inputTensor_, self, nil, true, MPSDataTypeInvalid, useStridedSelf)
         : Placeholder();
     auto otherPlaceholder = other.numel() != 0
-        ? Placeholder(
-              cachedGraph->otherTensor_, other, nil, true, MPSDataTypeInvalid, !isComplexType(other.scalar_type()))
+        ? Placeholder(cachedGraph->otherTensor_, other, nil, true, MPSDataTypeInvalid, useStridedOther)
         : Placeholder();
     auto outputPlaceholder = Placeholder(cachedGraph->outputTensor_, output);
 
@@ -869,6 +891,17 @@ static void linalg_inv_ex_out_mps_impl(const Tensor& A, bool check_errors, const
   if (output.numel() == 0) {
     return output;
   }
+  // Inner dimension is 0
+  // Early out as some paths in the code below do not handle this case correctly
+  if (self.size(1) == 0) {
+    if (beta.toDouble() == 0.0) {
+      output.zero_();
+    } else {
+      output.copy_(*bias_);
+      output.mul_(beta);
+    }
+    return output;
+  }
 
   if (use_metal_mm(self, other, output)) {
     return do_metal_addmm(self, other, output, alpha, beta, *bias_);
@@ -888,7 +921,8 @@ static void linalg_inv_ex_out_mps_impl(const Tensor& A, bool check_errors, const
     std::string key = "addmm_out_mps_impl" + getTensorsStringKey({self, other, *bias_}) + ":" +
         std::to_string(beta.toDouble()) + ":" + std::to_string(alpha.toDouble());
     auto cachedGraph = LookUpOrCreateCachedGraph<CachedGraph>(key, [&](auto mpsGraph, auto newCachedGraph) {
-      MPSGraphTensor* biasTensor = mpsGraphRankedPlaceHolder(mpsGraph, *bias_);
+      auto biasTensor = mpsGraphRankedPlaceHolder(mpsGraph, *bias_);
+      auto biasTensor_ = bias_->is_conj() ? [mpsGraph conjugateWithTensor:biasTensor name:nil] : biasTensor;
 
       // TODO: Use alpha and beta here with fill_.Scalar and mul
       auto [selfTensor, otherTensor, productTensor] = do_mm(mpsGraph, self, other);
@@ -901,11 +935,11 @@ static void linalg_inv_ex_out_mps_impl(const Tensor& A, bool check_errors, const
                                                             secondaryTensor:alphaTensor
                                                                        name:@"MM/alpha*(mat1@mat2)"];
       }
-      auto biasTimesBetaTensor = biasTensor;
+      auto biasTimesBetaTensor = biasTensor_;
       if (is_beta_non_zero && beta.toDouble() != 1.0) {
         auto betaTensor = [mpsGraph constantWithScalar:beta.toDouble()
                                               dataType:getMPSScalarType((*bias_).scalar_type())];
-        biasTimesBetaTensor = [mpsGraph multiplicationWithPrimaryTensor:biasTensor
+        biasTimesBetaTensor = [mpsGraph multiplicationWithPrimaryTensor:biasTensor_
                                                         secondaryTensor:betaTensor
                                                                    name:@"MM/beta*input"];
       }
@@ -1068,6 +1102,11 @@ static void linalg_inv_ex_out_mps_impl(const Tensor& A, bool check_errors, const
 }
 
 static Tensor& bmm_out_mps_impl(const Tensor& batch1, const Tensor& batch2, Tensor& result) {
+  TORCH_CHECK(batch1.scalar_type() == batch2.scalar_type(),
+              "Expected arguments of same type but got ",
+              batch1.scalar_type(),
+              " and ",
+              batch2.scalar_type());
   using namespace mps;
 
   // Matmul not supported if any output dimension size is larger than 2**32
@@ -1112,7 +1151,8 @@ static void linalg_inv_ex_out_mps_impl(const Tensor& A, bool check_errors, const
   // Call tiled implementation if the number of elements exceeds 2^32
   uint64_t resultSize = batch1.size(0) * batch1.size(1) * batch2.size(2);
   if (resultSize > pow(2, 32)) {
-    result = tiled_bmm_out_mps_impl(batch1, batch2, result);
+    // Tiled path uses MPSNDArray directly, so resolve conjugate views upfront
+    result = tiled_bmm_out_mps_impl(batch1.resolve_conj(), batch2.resolve_conj(), result);
     return result;
   }
 
@@ -1130,16 +1170,18 @@ static void linalg_inv_ex_out_mps_impl(const Tensor& A, bool check_errors, const
         std::to_string(doTranspose);
 
     auto cachedGraph = LookUpOrCreateCachedGraph<CachedGraph>(key, [&](auto mpsGraph, auto newCachedGraph) {
-      MPSGraphTensor* batch1Tensor = mps::mpsGraphUnrankedPlaceHolder(mpsGraph, getMPSDataType(batch1.scalar_type()));
-      MPSGraphTensor* batch2Tensor = mps::mpsGraphUnrankedPlaceHolder(mpsGraph, getMPSDataType(batch2.scalar_type()));
-      MPSGraphTensor* batch2TensorTranspose = batch2Tensor;
+      auto batch1Tensor = mps::mpsGraphUnrankedPlaceHolder(mpsGraph, getMPSDataType(batch1.scalar_type()));
+      auto batch2Tensor = mps::mpsGraphUnrankedPlaceHolder(mpsGraph, getMPSDataType(batch2.scalar_type()));
+
+      auto batch1TensorOp = batch1.is_conj() ? [mpsGraph conjugateWithTensor:batch1Tensor name:nil] : batch1Tensor;
+      auto batch2TensorOp = batch2.is_conj() ? [mpsGraph conjugateWithTensor:batch2Tensor name:nil] : batch2Tensor;
 
       if (doTranspose) {
-        batch2TensorTranspose = [mpsGraph transposeTensor:batch2Tensor dimension:-1 withDimension:-2 name:nil];
+        batch2TensorOp = [mpsGraph transposeTensor:batch2TensorOp dimension:-1 withDimension:-2 name:nil];
       }
 
-      MPSGraphTensor* productTensor = [mpsGraph matrixMultiplicationWithPrimaryTensor:batch1Tensor
-                                                                      secondaryTensor:batch2TensorTranspose
+      MPSGraphTensor* productTensor = [mpsGraph matrixMultiplicationWithPrimaryTensor:batch1TensorOp
+                                                                      secondaryTensor:batch2TensorOp
                                                                                  name:@"MM/(batch1@batch2)"];
 
       newCachedGraph->batch1Tensor_ = batch1Tensor;
@@ -1232,11 +1274,11 @@ static void linalg_inv_ex_out_mps_impl(const Tensor& A, bool check_errors, const
         const uint64_t aBatchOffset = i * aRows * aCols;
         const uint64_t bBatchOffset = i * bRows * bCols;
         MPSMatrix* sourceMatrix = [[[MPSMatrix alloc] initWithBuffer:aBuffer
-                                                              offset:(A_t.storage_offset() + aBatchOffset) * aElemSize
+                                                              offset:(A_.storage_offset() + aBatchOffset) * aElemSize
                                                           descriptor:sourceMatrixDesc] autorelease];
         MPSMatrix* rightHandSideMatrix =
             [[[MPSMatrix alloc] initWithBuffer:bBuffer
-                                        offset:(B_t.storage_offset() + bBatchOffset) * bElemSize
+                                        offset:(B_.storage_offset() + bBatchOffset) * bElemSize
                                     descriptor:rightHandSideMatrixDesc] autorelease];
         MPSMatrix* solutionMatrix = [[[MPSMatrix alloc] initWithBuffer:outBuffer
                                                                 offset:(out.storage_offset() + bBatchOffset) * bElemSize
@@ -1424,8 +1466,8 @@ static void cholesky_stub_impl(const Tensor& out, const Tensor& info, bool upper
   if (result.numel() == 0) {
     return result;
   }
-  auto cholesky = upper ? result.triu().clone() : result.tril().clone();
-  cholesky = cholesky.contiguous();
+  auto cholesky =
+      upper ? result.triu().clone(at::MemoryFormat::Contiguous) : result.tril().clone(at::MemoryFormat::Contiguous);
 
   auto n = result.size(-1);
   auto identity = at::eye(n, result.options()).expand_as(result).contiguous();
@@ -1445,6 +1487,107 @@ static void cholesky_stub_impl(const Tensor& out, const Tensor& info, bool upper
   return result;
 }
 
+static void metal_qr_kernel_impl(const Tensor& A, const Tensor& Q, const Tensor& R, bool reduced_mode) {
+  using namespace mps;
+
+  auto m = A.size(-2);
+  auto n = A.size(-1);
+
+  int64_t batch_size = 1;
+  for (int64_t i = 0; i < A.dim() - 2; i++) {
+    batch_size *= A.size(i);
+  }
+
+  auto A_work = A.reshape({batch_size, m, n}).contiguous();
+
+  QrParams params;
+  params.m = m;
+  params.n = n;
+
+  auto info = at::zeros({1}, A.options().dtype(kInt));
+  MPSStream* stream = getCurrentMPSStream();
+
+  Tensor Q_work = at::empty({batch_size, m, m}, A.options());
+  Tensor R_work = at::empty({batch_size, m, n}, A.options());
+  Tensor v_work = at::empty({batch_size, m}, A.options());
+
+  dispatch_sync_with_rethrow(stream->queue(), ^() {
+    @autoreleasepool {
+      auto compute_encoder = stream->commandEncoder();
+      auto pso = lib.getPipelineStateForFunc(fmt::format("linalg_qr_householder_{}", scalarToMetalTypeString(A)));
+
+      getMPSProfiler().beginProfileKernel(pso, "linalg_qr", {A});
+      [compute_encoder setComputePipelineState:pso];
+
+      MTLSize threadGroupSize = MTLSizeMake(1024, 1, 1);
+      // one threadgroup per matrix in batch
+      MTLSize gridSize = MTLSizeMake(batch_size, 1, 1);
+
+      mtl_setArgs(compute_encoder, A_work, Q_work, R_work, info, params, v_work);
+      [compute_encoder dispatchThreadgroups:gridSize threadsPerThreadgroup:threadGroupSize];
+
+      getMPSProfiler().endProfileKernel(pso);
+    }
+  });
+
+  bool is_batched = A.dim() > 2;
+
+  if (reduced_mode) {
+    auto k = std::min(m, n);
+    auto Q_reduced = Q_work.narrow(-1, 0, k); // [batch, m, k]
+    auto R_reduced = R_work.narrow(-2, 0, k); // [batch, k, n]
+
+    if (is_batched) {
+      Q.copy_(Q_reduced.reshape(Q.sizes()));
+      R.copy_(R_reduced.reshape(R.sizes()));
+    } else {
+      Q.copy_(Q_reduced.squeeze(0));
+      R.copy_(R_reduced.squeeze(0));
+    }
+  } else {
+    // Q=mxm, R=mxn
+    if (is_batched) {
+      Q.copy_(Q_work.reshape(Q.sizes()));
+      R.copy_(R_work.reshape(R.sizes()));
+    } else {
+      Q.copy_(Q_work.squeeze(0));
+      R.copy_(R_work.squeeze(0));
+    }
+  }
+
+  if (info.item<int>() != 0) {
+    TORCH_CHECK(false, "linalg_qr: MPS kernel failed with error code ", info.item<int>());
+  }
+}
+
+static void linalg_qr_out_impl_mps(const Tensor& A, const Tensor& Q, const Tensor& R, const c10::string_view mode) {
+  using namespace mps;
+
+  TORCH_CHECK(A.scalar_type() == kFloat, "linalg_qr: MPS currently supports float32 only");
+
+  if (A.numel() == 0) {
+    return;
+  }
+
+  auto m = A.size(-2);
+  auto n = A.size(-1);
+
+  if (std::min(m, n) > 512) {
+    TORCH_WARN_ONCE(
+        "linalg_qr: MPS implementation is currently limited to min(m,n) <= 512, "
+        "falling back to CPU.");
+    auto A_cpu = A.to(at::kCPU);
+    auto [Q_cpu, R_cpu] = at::linalg_qr(A_cpu, mode);
+    const_cast<Tensor&>(Q).copy_(Q_cpu.to(at::kMPS));
+    const_cast<Tensor&>(R).copy_(R_cpu.to(at::kMPS));
+    return;
+  }
+
+  bool reduced_mode = (mode != "complete");
+
+  metal_qr_kernel_impl(A, Q, R, reduced_mode);
+}
+
 } // namespace mps
 
 Tensor addr_mps(const Tensor& self, const Tensor& vec1, const Tensor& vec2, const Scalar& beta, const Scalar& alpha) {
@@ -1633,6 +1776,28 @@ Tensor linalg_solve_triangular_mps(const Tensor& A, const Tensor& B, bool upper,
   return out;
 }
 
+Tensor _cholesky_solve_helper_mps(const Tensor& self, const Tensor& A, bool upper) {
+  auto out = at::empty({0}, self.options().memory_format(MemoryFormat::Contiguous));
+  const bool first_transpose = upper;
+  const bool second_transpose = !upper;
+
+  mps::linalg_solve_triangular_mps_impl(A,
+                                        self,
+                                        upper,
+                                        first_transpose,
+                                        /*left=*/true,
+                                        /*unitriangular=*/false,
+                                        out);
+  mps::linalg_solve_triangular_mps_impl(A,
+                                        out,
+                                        upper,
+                                        second_transpose,
+                                        /*left=*/true,
+                                        /*unitriangular=*/false,
+                                        out);
+  return out;
+}
+
 TORCH_IMPL_FUNC(triangular_solve_mps_out)
 (const Tensor& self,
  const Tensor& A,
@@ -1669,6 +1834,10 @@ Tensor linalg_solve_triangular_mps(const Tensor& A, const Tensor& B, bool upper,
   mps::linalg_inv_ex_out_mps_impl(A, check_errors, result, info);
 }
 
+TORCH_IMPL_FUNC(linalg_qr_out_mps)(const Tensor& A, c10::string_view mode, const Tensor& Q, const Tensor& R) {
+  mps::linalg_qr_out_impl_mps(A, Q, R, mode);
+}
+
 REGISTER_DISPATCH(cholesky_stub, mps::cholesky_stub_impl)
 REGISTER_DISPATCH(unpack_pivots_stub, mps::unpack_pivots_stub_impl)
 REGISTER_DISPATCH(orgqr_stub, mps::orgqr_stub_impl);
diff --git a/aten/src/ATen/native/mps/operations/Normalization.mm b/aten/src/ATen/native/mps/operations/Normalization.mm
index 7441692b6c291..4e7b37f9d8883 100644
--- a/aten/src/ATen/native/mps/operations/Normalization.mm
+++ b/aten/src/ATen/native/mps/operations/Normalization.mm
@@ -70,6 +70,18 @@ static void get_shapes(MPSShape* input_shape_readonly,
       axes[i] = [NSNumber numberWithInt:i];
   }
 }
+
+static std::string get_mem_string(c10::MemoryFormat memory_format) {
+  switch (memory_format) {
+    case at::MemoryFormat::Contiguous:
+      return "Contiguous";
+    case at::MemoryFormat::ChannelsLast:
+      return "ChannelsLast";
+    default:
+      break;
+  }
+  TORCH_INTERNAL_ASSERT(false, "Unexpected memory format", memory_format);
+}
 } // namespace mps
 
 // Inverse standard deviation now becomes variance (without epsilon)
@@ -112,7 +124,10 @@ static void get_shapes(MPSShape* input_shape_readonly,
   const bool has_weight = (weight_opt.has_value() && weight_opt->defined());
   const bool has_bias = (bias_opt.has_value() && bias_opt->defined());
 
-  auto memory_format = self.suggest_memory_format();
+  // Use exact-match: a channel-slice of a channels-last tensor has CL-like
+  // strides but is not NHWC-packed.
+  // See https://github.com/pytorch/pytorch/issues/180984
+  auto memory_format = self.suggest_memory_format(/*channels_last_strides_exact_match=*/true);
 
   if (output.numel() == 0) {
     return std::tuple<Tensor&, Tensor&, Tensor&>(output, save_mean, save_var);
@@ -120,18 +135,6 @@ static void get_shapes(MPSShape* input_shape_readonly,
   }
 
   @autoreleasepool {
-    std::string mem_format_key;
-    switch (memory_format) {
-      case at::MemoryFormat::Contiguous:
-        mem_format_key = "Contiguous";
-        break;
-      case at::MemoryFormat::ChannelsLast:
-        mem_format_key = "ChannelsLast";
-        break;
-      default:
-        assert(0 && "Check should have been done earlier\n");
-    }
-
     // Number of elements in one channel, needed for bessel correction term
     const int64_t N = self.numel() / save_mean.numel();
     MPSShape* input_shape_readonly = mps::getMPSShape(self);
@@ -147,14 +150,20 @@ static void get_shapes(MPSShape* input_shape_readonly,
 
     NSString* ns_shape_key = [[input_shape valueForKey:@"description"] componentsJoinedByString:@","];
 
-    std::string key = "batch_norm_mps_out:" + mem_format_key + ":" + std::to_string(epsilon) + ":" +
-        std::to_string(momentum) + ":" + std::to_string(train) + ":" + std::to_string(has_running_mean) + ":" +
-        std::to_string(has_weight) + ":" + std::to_string(has_bias) + ":" + [ns_shape_key UTF8String] + ":" +
-        getTensorsStringKey({self,
-                             weight_opt.value_or(Tensor()),
-                             bias_opt.value_or(Tensor()),
-                             running_mean_opt.value_or(Tensor()),
-                             running_var_opt.value_or(Tensor())});
+    std::string key = fmt::format("batch_norm_mps_out:{}:{}:{}:{}:{}:{}:{}:{}:{}",
+                                  get_mem_string(memory_format),
+                                  epsilon,
+                                  momentum,
+                                  train,
+                                  has_running_mean,
+                                  has_weight,
+                                  has_bias,
+                                  [ns_shape_key UTF8String],
+                                  getTensorsStringKey({self,
+                                                       weight_opt.value_or(Tensor()),
+                                                       bias_opt.value_or(Tensor()),
+                                                       running_mean_opt.value_or(Tensor()),
+                                                       running_var_opt.value_or(Tensor())}));
     auto input_mps_dtype = getMPSDataType(self);
 
     // Dim where channels are located
@@ -215,25 +224,29 @@ Check if running mean exists (maybe do this check before making graph)
         MPSGraphTensor* batchVarianceTensor = [mpsGraph varianceOfTensor:inputTensor axes:axes name:nil];
         varTensor = batchVarianceTensor;
         if (has_running_mean) {
+          // Running stats may have a different dtype (e.g. float32 with float16 input)
+          auto running_mean_dtype = getMPSDataType(running_mean_opt.value());
           // TODO: This is not the formula used in PyTorch, is this OK? Seems more robust
           // float besselCorrectionTerm = float(N) / std::max(N - 1.0f, 1.0f);
           float besselCorrectionTerm = float(N) / float(N - 1.0f);
           MPSGraphTensor* besselConstantTensor = [mpsGraph constantWithScalar:(double)besselCorrectionTerm
                                                                         shape:@[ @1 ]
-                                                                     dataType:input_mps_dtype];
-          MPSGraphTensor* unbiasedVarianceTensor = [mpsGraph multiplicationWithPrimaryTensor:batchVarianceTensor
-                                                                             secondaryTensor:besselConstantTensor
-                                                                                        name:nil];
+                                                                     dataType:running_mean_dtype];
+          MPSGraphTensor* unbiasedVarianceTensor =
+              [mpsGraph multiplicationWithPrimaryTensor:castMPSTensor(mpsGraph, batchVarianceTensor, running_mean_dtype)
+                                        secondaryTensor:besselConstantTensor
+                                                   name:nil];
           MPSGraphTensor* momentumTensor = [mpsGraph constantWithScalar:(double)momentum
                                                                   shape:@[ @1 ]
-                                                               dataType:input_mps_dtype];
+                                                               dataType:running_mean_dtype];
           MPSGraphTensor* oneMinusMomentum = [mpsGraph constantWithScalar:(double)(1.0 - momentum)
                                                                     shape:@[ @1 ]
-                                                                 dataType:input_mps_dtype];
+                                                                 dataType:running_mean_dtype];
           // Compute updated running mean
-          MPSGraphTensor* scaledBatchMean = [mpsGraph multiplicationWithPrimaryTensor:batchMeanTensor
-                                                                      secondaryTensor:momentumTensor
-                                                                                 name:nil];
+          MPSGraphTensor* scaledBatchMean =
+              [mpsGraph multiplicationWithPrimaryTensor:castMPSTensor(mpsGraph, batchMeanTensor, running_mean_dtype)
+                                        secondaryTensor:momentumTensor
+                                                   name:nil];
           MPSGraphTensor* scaledRunningMean = [mpsGraph multiplicationWithPrimaryTensor:runningMeanTensor
                                                                         secondaryTensor:oneMinusMomentum
                                                                                    name:nil];
@@ -272,12 +285,16 @@ Check if running mean exists (maybe do this check before making graph)
         varTensor = saveVarTensor;
       }
 
+      // Cast weight and bias to input dtype if needed (mixed-precision support)
+      MPSGraphTensor* gammaTensor = has_weight ? castMPSTensor(mpsGraph, weightTensor, input_mps_dtype) : nil;
+      MPSGraphTensor* betaTensor = has_bias ? castMPSTensor(mpsGraph, biasTensor, input_mps_dtype) : nil;
+
       // Compute output of batch norm
       MPSGraphTensor* outputTensor = [mpsGraph normalizationWithTensor:inputTensor
                                                             meanTensor:saveMeanTensor
                                                         varianceTensor:varTensor
-                                                           gammaTensor:weightTensor
-                                                            betaTensor:biasTensor
+                                                           gammaTensor:gammaTensor
+                                                            betaTensor:betaTensor
                                                                epsilon:(float)epsilon
                                                                   name:nil];
 
@@ -379,7 +396,8 @@ Check if running mean exists (maybe do this check before making graph)
                                                   bool train,
                                                   double momentum,
                                                   double epsilon) {
-  const auto memory_format = self.suggest_memory_format();
+  // See https://github.com/pytorch/pytorch/issues/180984
+  const auto memory_format = self.suggest_memory_format(/*channels_last_strides_exact_match=*/true);
 
   auto output = at::empty(self.sizes(), self.scalar_type(), std::nullopt, kMPS, std::nullopt, memory_format);
 
@@ -493,22 +511,6 @@ Check if running mean exists (maybe do this check before making graph)
       self, weight_opt, bias_opt, Tensor(), Tensor(), train, momentum, epsilon, output, save_mean, save_var);
 }
 
-static std::string get_mem_string(c10::MemoryFormat memory_format) {
-  std::string mem_format_key;
-  switch (memory_format) {
-    case at::MemoryFormat::Contiguous:
-      mem_format_key = "Contiguous";
-      break;
-    case at::MemoryFormat::ChannelsLast:
-      mem_format_key = "ChannelsLast";
-      break;
-    default:
-      assert(0 && "Invalid memory format\n");
-  }
-
-  return mem_format_key;
-}
-
 // Batch norm backward
 std::tuple<Tensor, Tensor, Tensor> _new_batch_norm_backward_mps(const Tensor& grad_output,
                                                                 const Tensor& input,
@@ -547,7 +549,8 @@ Check if running mean exists (maybe do this check before making graph)
   Tensor grad_weight;
   Tensor grad_bias;
 
-  const auto memory_format = input.suggest_memory_format();
+  // See https://github.com/pytorch/pytorch/issues/180984
+  const auto memory_format = input.suggest_memory_format(/*channels_last_strides_exact_match=*/true);
 
   if (grad_input_mask[0]) {
     grad_input = at::empty(input.sizes(), input.scalar_type(), std::nullopt, kMPS, std::nullopt, memory_format);
@@ -608,18 +611,6 @@ Check if running mean exists (maybe do this check before making graph)
   }
 
   @autoreleasepool {
-    std::string mem_format_key;
-    switch (memory_format) {
-      case at::MemoryFormat::Contiguous:
-        mem_format_key = "Contiguous";
-        break;
-      case at::MemoryFormat::ChannelsLast:
-        mem_format_key = "ChannelsLast";
-        break;
-      default:
-        assert(0 && "Check should have been done earlier\n");
-    }
-
     MPSShape* input_shape_readonly = mps::getMPSShape(input);
     int num_input_dims = [input_shape_readonly count];
     NSMutableArray<NSNumber*>* input_shape = [NSMutableArray<NSNumber*> arrayWithCapacity:num_input_dims];
@@ -632,10 +623,18 @@ Check if running mean exists (maybe do this check before making graph)
 
     NSString* ns_shape_key = [[input_shape valueForKey:@"description"] componentsJoinedByString:@","];
 
-    std::string key = "batch_norm_backward_mps:" + mem_format_key + ":" + std::to_string(epsilon) + ":" +
-        std::to_string(train) + ":" + std::to_string(has_running_mean) + ":" + std::to_string(has_weight) + ":" +
-        [ns_shape_key UTF8String] + ":" + c10::Join(",", grad_input_mask) + ":" + getMPSTypeString(input);
     auto input_mps_dtype = getMPSDataType(input);
+    auto weight_mps_dtype = has_weight ? getMPSDataType(weight_opt.value()) : input_mps_dtype;
+    std::string key = fmt::format("batch_norm_backward_mps:{}:{}:{}:{}:{}:{}:{}:{}:{}",
+                                  get_mem_string(memory_format),
+                                  epsilon,
+                                  train,
+                                  has_running_mean,
+                                  has_weight,
+                                  [ns_shape_key UTF8String],
+                                  c10::Join(",", grad_input_mask),
+                                  getMPSTypeString(input),
+                                  has_weight ? getMPSTypeString(weight_opt.value()) : "none");
     auto cachedGraph = LookUpOrCreateCachedGraph<CachedGraph>(key, [&](auto mpsGraph, auto newCachedGraph) {
       // NCHW - Channels dim is 1
       int channelsDim = 1;
@@ -644,8 +643,11 @@ Check if running mean exists (maybe do this check before making graph)
       // Shape is the ORIGINAL NCHW shape
       auto gradOutputTensor = mpsGraphRankedPlaceHolder(mpsGraph, getMPSDataType(grad_out), input_shape_readonly);
       MPSGraphTensor* weightTensor = nil;
-      if (has_weight)
-        weightTensor = mpsGraphRankedPlaceHolder(mpsGraph, getMPSDataType(weight_opt.value()), new_mean_shape);
+      MPSGraphTensor* weightTensorCasted = nil;
+      if (has_weight) {
+        weightTensor = mpsGraphRankedPlaceHolder(mpsGraph, weight_mps_dtype, new_mean_shape);
+        weightTensorCasted = castMPSTensor(mpsGraph, weightTensor, input_mps_dtype);
+      }
       MPSGraphTensor* runningMeanTensor = nil;
       MPSGraphTensor* runningVarTensor = nil;
       if (has_running_mean) {
@@ -714,7 +716,7 @@ Check if running mean exists (maybe do this check before making graph)
                                                                          sourceTensor:inputTensor
                                                                            meanTensor:saveMeanTensor
                                                                        varianceTensor:revertSaveVarTensor
-                                                                          gammaTensor:weightTensor
+                                                                          gammaTensor:weightTensorCasted
                                                                   gammaGradientTensor:gradWeightTensor
                                                                    betaGradientTensor:gradBiasTensor
                                                                         reductionAxes:axes
@@ -782,7 +784,7 @@ Check if running mean exists (maybe do this check before making graph)
           gradInputTensor = [mpsGraph multiplicationWithPrimaryTensor:unitTensor secondaryTensor:rsqrtTensor name:nil];
           if (has_weight)
             gradInputTensor = [mpsGraph multiplicationWithPrimaryTensor:gradInputTensor
-                                                        secondaryTensor:weightTensor
+                                                        secondaryTensor:weightTensorCasted
                                                                    name:nil];
           gradInputTensor = [mpsGraph multiplicationWithPrimaryTensor:gradInputTensor
                                                       secondaryTensor:gradOutputTensor
@@ -794,11 +796,13 @@ Check if running mean exists (maybe do this check before making graph)
         gradWeightTensor = [mpsGraph reshapeTensor:gradWeightTensor
                                          withShape:@[ input_shape_readonly[channelsDim] ]
                                               name:nil];
+        gradWeightTensor = castMPSTensor(mpsGraph, gradWeightTensor, weight_mps_dtype);
       }
       if (grad_input_mask[2]) {
         gradBiasTensor = [mpsGraph reshapeTensor:gradBiasTensor
                                        withShape:@[ input_shape_readonly[channelsDim] ]
                                             name:nil];
+        gradBiasTensor = castMPSTensor(mpsGraph, gradBiasTensor, weight_mps_dtype);
       }
 
       MPSGraphTensor* gradInputTensorFinal = nil;
@@ -902,6 +906,7 @@ Check if running mean exists (maybe do this check before making graph)
   const Tensor& weight = *weight_maybe_owned;
   c10::MaybeOwned<Tensor> bias_maybe_owned = at::borrow_from_optional_tensor(bias_opt);
   const Tensor& bias = *bias_maybe_owned;
+  auto bias_contig = bias.expect_contiguous();
 
   auto M_N = _check_layer_norm_inputs(input, normalized_shape, weight, bias);
   auto M = M_N.first;
@@ -938,11 +943,11 @@ Check if running mean exists (maybe do this check before making graph)
 
       mps::mtl_setArgs(computeEncoder, *X, out, mean, rstd, axis_size, epsilon_buf, use_weight_buf, use_bias_buf);
       if (use_weight_and_bias_buf) {
-        mps::mtl_setArgs<8>(computeEncoder, *gamma, bias);
+        mps::mtl_setArgs<8>(computeEncoder, *gamma, *bias_contig);
       } else if (use_weight_buf) {
         mps::mtl_setArgs<8>(computeEncoder, *gamma);
       } else if (use_bias_buf) {
-        mps::mtl_setArgs<9>(computeEncoder, bias);
+        mps::mtl_setArgs<9>(computeEncoder, *bias_contig);
       }
       MTLSize numThreads = MTLSizeMake(std::min((axis_size + N_READS - 1) / N_READS, 1024), 1, 1);
       MTLSize numThreadgroups = MTLSizeMake(M, 1, 1);
@@ -959,7 +964,7 @@ Check if running mean exists (maybe do this check before making graph)
   }
   mean = mean.view(stat_shape);
   rstd = rstd.view(stat_shape);
-  return std::make_tuple(out, mean, rstd);
+  return std::make_tuple(std::move(out), std::move(mean), std::move(rstd));
 }
 
 std::tuple<Tensor, Tensor, Tensor> layer_norm_backward_mps(const Tensor& grad_out,
@@ -1092,9 +1097,12 @@ Check if running mean exists (maybe do this check before making graph)
       for (const auto i : c10::irange(num_normalized_dims))
         bn_gamma_shape[i + 2] = input_shape[i + num_channel_dims];
 
-      std::string key = "layer_norm_backward_mps:" + std::to_string(has_weight) + ":" +
-          getArrayRefString(normalized_shape) + ":" + getArrayRefString((*X).sizes()) + ":" +
-          c10::Join(",", grad_input_mask) + ":" + getMPSTypeString(*X);
+      std::string key = fmt::format("layer_norm_backward_mps:{}:{}:{}:{}:{}",
+                                    has_weight,
+                                    getArrayRefString(normalized_shape),
+                                    getArrayRefString((*X).sizes()),
+                                    c10::Join(",", grad_input_mask),
+                                    getMPSTypeString(*X));
       auto cachedGraph = LookUpOrCreateCachedGraph<CachedGraph>(key, [&](auto mpsGraph, auto newCachedGraph) {
         MPSGraphTensor* inputTensor = mpsGraphRankedPlaceHolder(mpsGraph, *X);
         MPSGraphTensor* gradOutputTensor = mpsGraphRankedPlaceHolder(mpsGraph, *dOut);
diff --git a/aten/src/ATen/native/mps/operations/Pooling.mm b/aten/src/ATen/native/mps/operations/Pooling.mm
index f9c541019a024..96fd799889473 100644
--- a/aten/src/ATen/native/mps/operations/Pooling.mm
+++ b/aten/src/ATen/native/mps/operations/Pooling.mm
@@ -69,7 +69,10 @@ static void pool2d_template(const Tensor& input,
   const bool is_backward_pass = grad_output.defined();
   const bool has_indices = indices.defined();
   const bool has_divisor = divisor_override.has_value() && divisor_override.value() != 0;
-  const auto suggested_memory_format = input.suggest_memory_format();
+  // Use exact-match: a channel-slice of a channels-last tensor has CL-like
+  // strides but is not NHWC-packed, so the raw-buffer NHWC path would misread
+  // it. See https://github.com/pytorch/pytorch/issues/180984
+  const auto suggested_memory_format = input.suggest_memory_format(/*channels_last_strides_exact_match=*/true);
   // for max_pool2d_with_indices() we cannot pass ChannelsLast (i.e., NHWC) to 'desc.dataLayout' in MPSGraph.
   // Because the returned indices will be selected based on NHWC memory layout which will
   // be incompatible with the PyTorch's global NCHW layout.
@@ -157,11 +160,17 @@ static void pool2d_template(const Tensor& input,
     padH = padW = 0;
   }
   @autoreleasepool {
-    std::string key = op_name + getTensorsStringKey({input, indices, grad_output}) + ":K[" +
-        getArrayRefString(kernel_size) + "]:S[" + getArrayRefString(stride) + "]:P[" + getArrayRefString(padding) +
-        "]:D[" + getArrayRefString(dilation) + "]" + (ceil_mode ? ":ceil" : "") +
-        (count_include_pad ? ":include_pad" : "") + (has_divisor ? ":divisor" : "") + ":" +
-        (suggested_memory_format == MemoryFormat::ChannelsLast ? "NHWC" : "NCHW");
+    const auto key = fmt::format("{}{}:K[{}]:S[{}]:P[{}]:D[{}]{}{}{}:{}",
+                                 op_name,
+                                 getTensorsStringKey({input, indices, grad_output}),
+                                 getArrayRefString(kernel_size),
+                                 getArrayRefString(stride),
+                                 getArrayRefString(padding),
+                                 getArrayRefString(dilation),
+                                 ceil_mode ? ":ceil" : "",
+                                 count_include_pad ? ":include_pad" : "",
+                                 has_divisor ? ":divisor" : "",
+                                 suggested_memory_format == MemoryFormat::ChannelsLast ? "NHWC" : "NCHW");
 
     MPSShape* inputShape = getMPSShape(input, memory_format);
     MPSShape* gradOutputShape = is_backward_pass ? getMPSShape(grad_output, memory_format) : nullptr;
diff --git a/aten/src/ATen/native/mps/operations/RMSNorm.mm b/aten/src/ATen/native/mps/operations/RMSNorm.mm
index 7948b5acd8e93..33348c6cbff8e 100644
--- a/aten/src/ATen/native/mps/operations/RMSNorm.mm
+++ b/aten/src/ATen/native/mps/operations/RMSNorm.mm
@@ -55,7 +55,7 @@
       if (N <= LOOPED_LIMIT) {
         size_t threadgroup_needed = (N + N_READS - 1) / N_READS;
         size_t simds_needed = (threadgroup_needed + SIMD_SIZE - 1) / SIMD_SIZE;
-        size_t threadgroup_size = SIMD_SIZE * simds_needed;
+        threadgroup_size = SIMD_SIZE * simds_needed;
         assert(threadgroup_size <= maxThreadsPerGroup);
       }
       size_t n_threads = M * threadgroup_size;
diff --git a/aten/src/ATen/native/mps/operations/ReduceOps.mm b/aten/src/ATen/native/mps/operations/ReduceOps.mm
index e634eefee2058..2ec070c84be7e 100644
--- a/aten/src/ATen/native/mps/operations/ReduceOps.mm
+++ b/aten/src/ATen/native/mps/operations/ReduceOps.mm
@@ -3,8 +3,10 @@
 #include <ATen/ExpandUtils.h>
 #include <ATen/TensorUtils.h>
 #include <ATen/native/Pool.h>
+#include <ATen/native/ReduceOps.h>
 #include <ATen/native/ReduceOpsUtils.h>
 #include <ATen/native/mps/OperationUtils.h>
+#include <ATen/native/mps/kernels/ReduceOps.h>
 #include <c10/util/irange.h>
 
 #ifndef AT_PER_OPERATOR_HEADERS
@@ -20,7 +22,7 @@
 #include <ATen/ops/argmax_native.h>
 #include <ATen/ops/argmin_native.h>
 #include <ATen/ops/count_nonzero_native.h>
-#include <ATen/ops/linalg_vector_norm_native.h>
+#include <ATen/ops/linalg_vector_norm.h>
 #include <ATen/ops/max_native.h>
 #include <ATen/ops/mean_native.h>
 #include <ATen/ops/median.h>
@@ -28,7 +30,6 @@
 #include <ATen/ops/min_native.h>
 #include <ATen/ops/nanmedian_native.h>
 #include <ATen/ops/nansum_native.h>
-#include <ATen/ops/norm_native.h>
 #include <ATen/ops/prod_native.h>
 #include <ATen/ops/std_mean_native.h>
 #include <ATen/ops/std_native.h>
@@ -41,9 +42,12 @@
 
 namespace at::native {
 namespace mps {
-typedef MPSGraphTensor* (^NormOpBlock)(MPSBinaryCachedGraph*, MPSGraphTensor*, MPSGraphTensor*);
-#define NormOpFn(graph, primary, secondary) \
-  MPSGraphTensor*(MPSBinaryCachedGraph * graph, MPSGraphTensor * primary, MPSGraphTensor * secondary)
+
+#ifndef PYTORCH_JIT_COMPILE_SHADERS
+static auto& lib = MetalShaderLibrary::getBundledLibrary();
+#else
+#include <ATen/native/mps/ReduceOps_metallib.h>
+#endif
 
 enum StdVarType { STANDARD_VARIANCE, STANDARD_DEVIATION };
 
@@ -303,143 +307,58 @@ static void reduction_out_mps(const Tensor& input_t,
   }
 }
 
-static void impl_func_norm_mps(const Tensor& input_tensor,
-                               const Tensor& other_tensor,
-                               const OptionalScalarRef& opt_p,
-                               IntArrayRef dim,
-                               bool keepdim,
-                               std::optional<ScalarType> opt_dtype,
-                               const Tensor& output_t,
-                               bool cdist = false,
-                               std::optional<IntArrayRef> input_broadcasted_shape = std::nullopt,
-                               NormOpBlock normOpBlock = nullptr) {
-  auto p = opt_p.has_value() ? opt_p.get().to<double>() : Scalar(2.0).to<double>();
-  if (input_tensor.numel() == 0) {
-    output_t.fill_((p < 0) ? INFINITY : 0);
-    return;
-  }
-
-  auto input_t = (input_tensor.sizes().size() == 0) ? input_tensor.view({1}) : input_tensor;
-  auto in_dtype = opt_dtype.value_or(input_tensor.scalar_type());
-  auto mps_input_dtype = getMPSDataType(in_dtype);
-  TORCH_CHECK(!input_tensor.is_complex(), "norm ops are not supported for complex yet");
-
-  IntArrayRef input_shape = cdist ? input_broadcasted_shape.value() : input_t.sizes();
-
-  for (const auto dim_val : dim) {
-    auto wrap_dim = maybe_wrap_dim(dim_val, input_shape.size());
-    TORCH_CHECK(wrap_dim < static_cast<decltype(wrap_dim)>(input_shape.size()),
-                "norm_out_mps: reduction dim must be in the range of input shape")
-  }
-
-  auto reciprocal_p = 1 / p;
-  bool pIsZero = (p == 0.0);
-  bool pIsPosInf = (p == std::numeric_limits<double>::infinity());
-  bool pIsNegInf = (p == -std::numeric_limits<double>::infinity());
-
-  int64_t num_input_dims = input_shape.size();
-  int64_t num_reduce_dims = dim.size();
-  int64_t num_output_dims;
-
-  // For output shape calculation, assume that keepdim is true
-  num_output_dims = num_input_dims;
-  NSMutableArray<NSNumber*>* apparent_output_shape = nil;
-  NSMutableArray<NSNumber*>* apparent_input_shape = nil;
-
-  // Reduction axes
-  NSMutableArray<NSNumber*>* axes;
-  set_axes(axes, num_reduce_dims, dim, input_shape.size());
-
-  set_apparent_shapes(apparent_output_shape, apparent_input_shape, num_reduce_dims, num_output_dims, input_shape, axes);
+static void norm_kernel_mps(TensorIterator& iter, const Scalar& p_scalar) {
+  const Tensor& output = iter.output(0);
+  const Tensor& input = iter.input(0);
+  auto p = p_scalar.to<double>();
 
-  NSArray<NSNumber*>* wrappedAxes = getTensorAxes(input_shape, dim);
-  if (cdist) {
-    apparent_input_shape = [getMPSShape(input_tensor.sizes()) mutableCopy];
-    apparent_output_shape = [getMPSShape(output_t.sizes()) mutableCopy];
+  if (input.numel() == 0) {
+    output.fill_((p < 0) ? INFINITY : 0);
+    return;
   }
 
-  if (output_t.numel() == 0) {
+  if (output.numel() == 0) {
     return;
   }
 
-  auto stream = getCurrentMPSStream();
-  @autoreleasepool {
-    NSString* ns_key = [[wrappedAxes valueForKey:@"description"] componentsJoinedByString:@","];
-    std::string keepdim_info = (keepdim) ? "keepdim=1" : "keepdim=0";
-    std::string tensor_key = cdist ? getTensorsStringKey({input_tensor, other_tensor}) : getTensorsStringKey({input_t});
-    std::string key = std::string("norm_out_mps:") + [ns_key UTF8String] + ":" + tensor_key + ":p" + std::to_string(p) +
-        ":" + keepdim_info + ":" + toString(in_dtype);
-
-    auto cachedGraph = LookUpOrCreateCachedGraph<MPSBinaryCachedGraph>(key, [&](auto mpsGraph, auto newCachedGraph) {
-      newCachedGraph->inputTensor_ = mpsGraphRankedPlaceHolder(mpsGraph, input_tensor);
-
-      if (cdist) {
-        newCachedGraph->otherTensor_ = mpsGraphRankedPlaceHolder(mpsGraph, other_tensor);
-      }
-
-      MPSGraphTensor* inputTensor = cdist
-          ? normOpBlock(newCachedGraph, newCachedGraph->inputTensor_, newCachedGraph->otherTensor_)
-          : newCachedGraph->inputTensor_;
-
-      if (opt_dtype.has_value()) {
-        inputTensor = castMPSTensor(mpsGraph, inputTensor, mps_input_dtype);
-      }
-
-      MPSGraphTensor* outputTensor;
-
-      if (pIsZero) {
-        MPSGraphTensor* zeros = [mpsGraph constantWithScalar:0.0 dataType:mps_input_dtype];
-        MPSGraphTensor* ones = [mpsGraph constantWithScalar:1.0 dataType:mps_input_dtype];
-        MPSGraphTensor* nonZeros = [mpsGraph selectWithPredicateTensor:inputTensor
-                                                   truePredicateTensor:ones
-                                                  falsePredicateTensor:zeros
-                                                                  name:nil];
-        outputTensor = [mpsGraph reductionSumWithTensor:nonZeros axes:wrappedAxes name:nil];
-      } else if (pIsPosInf) {
-        MPSGraphTensor* absoluteTensor = [mpsGraph absoluteWithTensor:inputTensor name:nil];
-        outputTensor = [mpsGraph reductionMaximumWithTensor:absoluteTensor axes:wrappedAxes name:nil];
-      } else if (pIsNegInf) {
-        MPSGraphTensor* absoluteTensor = [mpsGraph absoluteWithTensor:inputTensor name:nil];
-        outputTensor = [mpsGraph reductionMinimumWithTensor:absoluteTensor axes:wrappedAxes name:nil];
-      } else {
-        MPSGraphTensor* absoluteTensor = [mpsGraph absoluteWithTensor:inputTensor name:nil];
+  // Number of input elements that are reduced into one output element
+  uint32_t reduction_size = input.numel() / output.numel();
 
-        MPSGraphTensor* powerValTensor = [mpsGraph constantWithScalar:p dataType:mps_input_dtype];
+  TORCH_INTERNAL_ASSERT(output.dim() == input.dim());
 
-        MPSGraphTensor* reciprocalPowerValTensor = [mpsGraph constantWithScalar:reciprocal_p dataType:mps_input_dtype];
+  NormParams params;
 
-        MPSGraphTensor* powerTensor = [mpsGraph powerWithPrimaryTensor:absoluteTensor
-                                                       secondaryTensor:powerValTensor
-                                                                  name:nil];
+  params.ndim = input.dim();
+  params.p = static_cast<float>(p);
+  params.reduction_size = reduction_size;
 
-        MPSGraphTensor* reductionSumTensor = [mpsGraph reductionSumWithTensor:powerTensor axes:wrappedAxes name:nil];
+  for (const auto dim_idx : c10::irange(input.dim())) {
+    params.input_sizes[dim_idx] = input.size(dim_idx);
+    params.input_strides[dim_idx] = input.stride(dim_idx);
+    params.output_sizes[dim_idx] = output.size(dim_idx);
+    params.output_strides[dim_idx] = output.stride(dim_idx);
+  }
 
-        outputTensor = [mpsGraph powerWithPrimaryTensor:reductionSumTensor
-                                        secondaryTensor:reciprocalPowerValTensor
-                                                   name:nil];
-      }
+  MPSStream* stream = getCurrentMPSStream();
 
-      if (cdist) {
-        outputTensor = [mpsGraph reshapeTensor:outputTensor withShape:getMPSShape(output_t) name:nil];
-      }
+  dispatch_sync_with_rethrow(stream->queue(), ^() {
+    @autoreleasepool {
+      id<MTLComputeCommandEncoder> compute_encoder = stream->commandEncoder();
+      auto pipeline_state = lib.getPipelineStateForFunc(
+          fmt::format("norm_{}_{}", scalarToMetalTypeString(input), scalarToMetalTypeString(output)));
+      getMPSProfiler().beginProfileKernel(pipeline_state, "norm", {input});
+      [compute_encoder setComputePipelineState:pipeline_state];
+      mtl_setArgs(compute_encoder, input, output, params);
 
-      newCachedGraph->outputTensor_ = outputTensor;
-    });
+      auto threads_per_group = std::min(MAX_THREADGROUP_SIZE, reduction_size);
+      uint32_t num_threads = output.numel() * threads_per_group;
 
-    auto otherPlaceholder = Placeholder();
-    auto inputPlaceholder = Placeholder(cachedGraph->inputTensor_, input_t);
-    auto outputPlaceholder = Placeholder(cachedGraph->outputTensor_, output_t, apparent_output_shape);
+      [compute_encoder dispatchThreads:MTLSizeMake(num_threads, 1, 1)
+                 threadsPerThreadgroup:MTLSizeMake(threads_per_group, 1, 1)];
 
-    NSMutableDictionary<MPSGraphTensor*, MPSGraphTensorData*>* feeds = [NSMutableDictionary dictionary];
-    feeds[inputPlaceholder.getMPSGraphTensor()] = inputPlaceholder.getMPSGraphTensorData();
-
-    if (cdist) {
-      otherPlaceholder = Placeholder(cachedGraph->otherTensor_, other_tensor);
-      feeds[otherPlaceholder.getMPSGraphTensor()] = otherPlaceholder.getMPSGraphTensorData();
+      getMPSProfiler().endProfileKernel(pipeline_state);
     }
-
-    runMPSGraph(stream, cachedGraph->graph(), feeds, outputPlaceholder);
-  }
+  });
 }
 
 static Tensor std_var_common_impl_mps(const Tensor& input_t,
@@ -1113,32 +1032,6 @@ Tensor count_nonzero_mps(const Tensor& self, IntArrayRef dims) {
   reduction_out_mps(input_t, opt_dim, keepdim, dtype, output_t, MPSReductionType::MEAN, "mean_out_mps");
 }
 
-TORCH_IMPL_FUNC(norm_out_mps)
-(const Tensor& self, const OptionalScalarRef opt_p, IntArrayRef dim, bool keepdim, const Tensor& result) {
-  impl_func_norm_mps(self, self, opt_p, dim, keepdim, std::nullopt, result, /*cdist=*/false);
-}
-
-TORCH_IMPL_FUNC(norm_dtype_out_mps)
-(const Tensor& self,
- const OptionalScalarRef opt_p,
- IntArrayRef dim,
- bool keepdim,
- ScalarType dtype,
- const Tensor& result) {
-  impl_func_norm_mps(self, self, opt_p, dim, keepdim, dtype, result, /*cdist=*/false);
-}
-
-TORCH_IMPL_FUNC(linalg_vector_norm_out_mps)
-(const Tensor& self,
- const Scalar& scalar_ord,
- OptionalIntArrayRef opt_dim,
- bool keepdim,
- std::optional<ScalarType> opt_dtype,
- const Tensor& result) {
-  impl_func_norm_mps(
-      self, self, scalar_ord, opt_dim.value_or(IntArrayRef{}), keepdim, opt_dtype, result, /*cdist=*/false);
-}
-
 Tensor _cdist_forward_mps(const Tensor& x1, const Tensor& x2, const double p, std::optional<int64_t> compute_mode) {
   TORCH_CHECK(x1.dim() >= 2, "cdist only supports at least 2D tensors, X1 got: ", x1.dim(), "D");
   TORCH_CHECK(x2.dim() >= 2, "cdist only supports at least 2D tensors, X2 got: ", x2.dim(), "D");
@@ -1149,96 +1042,20 @@ Tensor _cdist_forward_mps(const Tensor& x1, const Tensor& x2, const double p, st
               x2.size(-1));
   TORCH_CHECK(
       at::isFloatingType(x1.scalar_type()), "cdist only supports floating-point dtypes, X1 got: ", x1.scalar_type());
-  auto device1 = x1.device().type();
   TORCH_CHECK(
       at::isFloatingType(x2.scalar_type()), "cdist only supports floating-point dtypes, X2 got: ", x2.scalar_type());
-  auto device2 = x2.device().type();
   TORCH_CHECK(p >= 0, "cdist only supports non-negative p values");
-  TORCH_CHECK(device1 == device2, "X1 and X2 must have the same device type. X1: ", device1, " X2: ", device2);
-  TORCH_CHECK(x1.is_mps() && (x1.get_device() == x2.get_device()),
-              "device of X1 (",
-              x1.get_device(),
-              ") must match device of X2 (",
-              x2.get_device(),
-              ")");
-
-  int64_t c1 = x1.size(-1);
-  int64_t c2 = x2.size(-1);
-
-  auto dim1 = x1.dim();
-  auto dim2 = x2.dim();
+
   int64_t mode = compute_mode.value_or(0);
   TORCH_CHECK(mode >= 0 && mode <= 2, "possible modes: 0, 1, 2, but was: ", mode);
 
-  int64_t r1 = x1.size(-2);
-  int64_t r2 = x2.size(-2);
-
-  // For batch calculation we expand all dimensions(except the last two) to one, with size that equals to product of
-  // them. The last two dimensions will stay the same
-  IntArrayRef batch_tensor1(x1.sizes().data(), dim1 - 2);
-  IntArrayRef batch_tensor2(x2.sizes().data(), dim2 - 2);
-  std::vector<int64_t> expand_batch_portion = infer_size(batch_tensor1, batch_tensor2);
-  std::vector<int64_t> tensor1_expand_size(expand_batch_portion);
-  tensor1_expand_size.insert(tensor1_expand_size.end(), {r1, c1});
-  std::vector<int64_t> tensor2_expand_size(expand_batch_portion);
-  tensor2_expand_size.insert(tensor2_expand_size.end(), {r2, c2});
-
-  const int64_t expand_batch_product = c10::multiply_integers(expand_batch_portion);
-  std::vector<int64_t> tensor1_view{expand_batch_product, r1, c1};
-  std::vector<int64_t> tensor2_view{expand_batch_product, r2, c2};
-
-  std::vector<int64_t> output_shape(expand_batch_portion);
-  output_shape.insert(output_shape.end(), {r1, r2});
+  Tensor x1_ = x1.unsqueeze(-2);
+  Tensor x2_ = x2.unsqueeze(-3);
+  Tensor diff = x1_.sub(x2_);
+  IntArrayRef output_shape(diff.sizes().data(), diff.dim() - 1);
   Tensor result = at::empty(output_shape, x1.options());
+  linalg_vector_norm_out(result, diff, p, makeArrayRef<int64_t>(-1), /*keepdim=*/false, /*dtype=*/std::nullopt);
 
-  NormOpBlock norm_op_block = ^NormOpFn(cachedGraph, x1Tensor, x2Tensor) {
-    MPSGraph* mpsGraph = cachedGraph->graph();
-
-    MPSGraphTensor* inputBroadcast = [mpsGraph broadcastTensor:x1Tensor
-                                                       toShape:getMPSShape(tensor1_expand_size)
-                                                          name:nil];
-    MPSGraphTensor* inputBroadcastReshape = [mpsGraph reshapeTensor:inputBroadcast
-                                                          withShape:getMPSShape(tensor1_view)
-                                                               name:nil];
-
-    MPSGraphTensor* otherBroadcast = [mpsGraph broadcastTensor:x2Tensor
-                                                       toShape:getMPSShape(tensor2_expand_size)
-                                                          name:nil];
-    MPSGraphTensor* otherBroadcastReshape = [mpsGraph reshapeTensor:otherBroadcast
-                                                          withShape:getMPSShape(tensor2_view)
-                                                               name:nil];
-
-    NSMutableArray<MPSGraphTensor*>* inputArray = [NSMutableArray arrayWithCapacity:tensor1_view[1]];
-    NSMutableArray<MPSGraphTensor*>* otherArray = [NSMutableArray arrayWithCapacity:tensor2_view[1]];
-
-    for (const auto i : c10::irange(tensor2_view[1])) {
-      inputArray[i] = inputBroadcastReshape;
-    }
-
-    for (const auto i : c10::irange(tensor1_view[1])) {
-      otherArray[i] = otherBroadcastReshape;
-    }
-
-    MPSGraphTensor* inputTensorReshaped = [mpsGraph concatTensors:inputArray dimension:1 interleave:YES name:nil];
-    MPSGraphTensor* otherTensorReshaped = [mpsGraph concatTensors:otherArray dimension:1 interleave:NO name:nil];
-
-    MPSGraphTensor* inputTensorPNorm = [mpsGraph subtractionWithPrimaryTensor:inputTensorReshaped
-                                                              secondaryTensor:otherTensorReshaped
-                                                                         name:nil];
-    return inputTensorPNorm;
-  };
-
-  IntArrayRef inputBroadcastSize = makeArrayRef(tensor1_view.data(), tensor1_view.size());
-  impl_func_norm_mps(x1,
-                     x2,
-                     OptionalScalarRef(p),
-                     makeArrayRef<int64_t>(2),
-                     false,
-                     std::nullopt,
-                     result,
-                     /*cdist=*/true,
-                     inputBroadcastSize,
-                     norm_op_block);
   return result;
 }
 
@@ -1701,4 +1518,6 @@ Tensor nanmedian_mps(const Tensor& self) {
   return {var, mean};
 }
 
+REGISTER_DISPATCH(norm_stub, &mps::norm_kernel_mps)
+
 } // namespace at::native
diff --git a/aten/src/ATen/native/mps/operations/Repeat.mm b/aten/src/ATen/native/mps/operations/Repeat.mm
index f350b0137b05e..7abfb0a2e4500 100644
--- a/aten/src/ATen/native/mps/operations/Repeat.mm
+++ b/aten/src/ATen/native/mps/operations/Repeat.mm
@@ -10,6 +10,14 @@
 #include <ATen/ops/repeat_native.h>
 #include <fmt/format.h>
 
+#ifndef AT_PER_OPERATOR_HEADERS
+#include <ATen/Functions.h>
+#include <ATen/NativeFunctions.h>
+#else
+#include <ATen/ops/view_as_complex.h>
+#include <ATen/ops/view_as_real.h>
+#endif
+
 namespace at::native {
 
 Tensor permute_mps(const Tensor& self, IntArrayRef dims) {
@@ -36,7 +44,13 @@ Tensor repeat_mps(const Tensor& self, IntArrayRef repeats) {
 
   TORCH_CHECK(repeats.size() >= (size_t)self.dim(),
               "Number of dimensions of repeat dims can not be smaller than number of dimensions of tensor");
-  TORCH_CHECK(!self.is_complex(), "repeat(): Not supported for complex yet!");
+
+  if (self.is_complex()) {
+    std::vector<int64_t> repeats_real = repeats.vec();
+    repeats_real.push_back(1);
+    auto self_real = at::view_as_real(self);
+    return at::view_as_complex(repeat_mps(self_real, repeats_real));
+  }
 
   // Add new leading dimensions to the tensor if the
   // number of target dimensions is larger than the
diff --git a/aten/src/ATen/native/mps/operations/ScanKernel.h b/aten/src/ATen/native/mps/operations/ScanKernel.h
new file mode 100644
index 0000000000000..06bedfbdd215e
--- /dev/null
+++ b/aten/src/ATen/native/mps/operations/ScanKernel.h
@@ -0,0 +1,11 @@
+#pragma once
+
+namespace at::native::mps {
+
+void scan_simple_mps_impl(
+    const Tensor& self,
+    const Tensor& output,
+    int64_t dim,
+    const std::string& op_name);
+
+} // namespace at::native::mps
diff --git a/aten/src/ATen/native/mps/operations/ScanKernel.mm b/aten/src/ATen/native/mps/operations/ScanKernel.mm
index 80495ba9d501d..9b8a8a5d485f7 100644
--- a/aten/src/ATen/native/mps/operations/ScanKernel.mm
+++ b/aten/src/ATen/native/mps/operations/ScanKernel.mm
@@ -3,6 +3,7 @@
 #include <ATen/mps/MPSProfiler.h>
 #include <ATen/native/ReduceOps.h>
 #include <ATen/native/mps/OperationUtils.h>
+#include <ATen/native/mps/operations/ScanKernel.h>
 
 #ifndef AT_PER_OPERATOR_HEADERS
 #include <ATen/Functions.h>
@@ -45,7 +46,7 @@
   return {static_cast<uint32_t>(grid_x), static_cast<uint32_t>(grid_y)};
 }
 
-static void scan_simple_mps_impl(const Tensor& self, const Tensor& output, int64_t dim, const std::string& op_name) {
+void scan_simple_mps_impl(const Tensor& self, const Tensor& output, int64_t dim, const std::string& op_name) {
   if (output.numel() == 0) {
     return;
   }
@@ -60,12 +61,10 @@ static void scan_simple_mps_impl(const Tensor& self, const Tensor& output, int64
   // Preprocess output tensor - ensure it's contiguous for Metal shaders
   Tensor output_tensor = output;
   bool output_needs_copy = !output.is_contiguous();
-  Tensor temp_output;
 
   if (output_needs_copy) {
     // Create a temporary contiguous tensor with the same shape and type
-    temp_output = at::empty_like(output, output.options()).contiguous();
-    output_tensor = temp_output;
+    output_tensor = at::empty_like(output, output.options().memory_format(c10::MemoryFormat::Contiguous));
   }
 
   // Determine which kernel to use based on scan dimension position
diff --git a/aten/src/ATen/native/mps/operations/ScatterGather.mm b/aten/src/ATen/native/mps/operations/ScatterGather.mm
index ce65421c71c9d..b9b055b4fee6f 100644
--- a/aten/src/ATen/native/mps/operations/ScatterGather.mm
+++ b/aten/src/ATen/native/mps/operations/ScatterGather.mm
@@ -10,10 +10,23 @@
 #include <ATen/ops/gather_native.h>
 #include <ATen/ops/scatter_add_native.h>
 #include <ATen/ops/scatter_native.h>
+#include <ATen/ops/view_as_real.h>
 #endif
 
 namespace at::native {
 
+static Tensor maybe_expand_0_dim(const Tensor& t) {
+  return t.dim() == 0 ? t.view({1}) : t;
+}
+
+static Tensor expand_index_as_real(const Tensor& index) {
+  auto index_view = maybe_expand_0_dim(index);
+  std::vector<int64_t> index_expanded_sizes = index_view.sizes().vec();
+  index_expanded_sizes.push_back(2);
+  auto index_expanded = index_view.unsqueeze(-1).expand(index_expanded_sizes);
+  return index_expanded;
+}
+
 TORCH_IMPL_FUNC(gather_out_mps)
 (const Tensor& self_arg, int64_t dim, const Tensor& index, bool sparse_grad, const Tensor& output) {
   using namespace mps;
@@ -27,7 +40,14 @@
   TORCH_CHECK(!sparse_grad, "sparse_grad not supported in MPS yet")
   TORCH_CHECK(self.scalar_type() == output.scalar_type(), "gather(): self and output must have the same scalar type");
   TORCH_CHECK(dim >= 0 && dim < self.dim(), "gather(): Indexing dim ", dim, " is out of bounds of tensor");
-  TORCH_CHECK(!self.is_complex(), "gather(): Yet not supported for complex");
+
+  if (self.is_complex()) {
+    auto self_real = at::view_as_real(self);
+    auto index_expanded = expand_index_as_real(index);
+    auto output_real = at::view_as_real(maybe_expand_0_dim(output));
+    structured_gather_out_mps::impl(self_real, dim, index_expanded, sparse_grad, output_real);
+    return;
+  }
 
   struct CachedGraph : public MPSCachedGraph {
     CachedGraph(MPSGraph* graph) : MPSCachedGraph(graph) {}
@@ -145,7 +165,15 @@ static void scatter_mps_general(const Tensor& self_arg,
   TORCH_CHECK(self.scalar_type() == output.scalar_type() && output.scalar_type() == src.scalar_type(),
               "scatter(): self, src and output must have the same scalar type");
   TORCH_CHECK(dim >= 0 && dim < self.dim(), "scatter(): Indexing dim ", dim, " is out of bounds of tensor");
-  TORCH_CHECK(!self.is_complex(), "scatter(): Yet not supported for complex");
+
+  if (self.is_complex()) {
+    auto self_real = at::view_as_real(self);
+    auto index_expanded = expand_index_as_real(index);
+    auto src_real = at::view_as_real(maybe_expand_0_dim(src));
+    auto output_real = at::view_as_real(maybe_expand_0_dim(output));
+    scatter_mps_general(self_real, dim, index_expanded, src_real, output_real, func_name, reduce);
+    return;
+  }
 
   struct CachedGraph : public MPSCachedGraph {
     CachedGraph(MPSGraph* graph) : MPSCachedGraph(graph) {}
diff --git a/aten/src/ATen/native/mps/operations/UnaryKernel.mm b/aten/src/ATen/native/mps/operations/UnaryKernel.mm
index 474f7a3f67500..b0c73043a313d 100644
--- a/aten/src/ATen/native/mps/operations/UnaryKernel.mm
+++ b/aten/src/ATen/native/mps/operations/UnaryKernel.mm
@@ -51,6 +51,16 @@ static void erfcx_kernel(TensorIteratorBase& iter) {
   lib.exec_unary_kernel(iter, "erfcx");
 }
 
+static void polygamma_kernel(TensorIteratorBase& iter, int64_t order) {
+  if (order == 0) {
+    return lib.exec_unary_kernel(iter, "digamma");
+  } else if (order == 1) {
+    return lib.exec_unary_kernel(iter, "trigamma");
+  } else {
+    return lib.exec_unary_kernel(iter, "polygamma", order, kInt);
+  }
+}
+
 REGISTER_UNARY_TI_DISPATCH(exp);
 REGISTER_UNARY_TI_DISPATCH(expm1);
 REGISTER_UNARY_TI_DISPATCH(erf);
@@ -77,10 +87,13 @@ static void erfcx_kernel(TensorIteratorBase& iter) {
 REGISTER_UNARY_TI_DISPATCH(log2);
 REGISTER_UNARY_TI_DISPATCH(log);
 REGISTER_UNARY_TI_DISPATCH(log1p);
+REGISTER_UNARY_TI_DISPATCH(lgamma);
+REGISTER_UNARY_TI_DISPATCH(digamma);
 REGISTER_UNARY_TI_DISPATCH(bitwise_not);
 REGISTER_UNARY_TI_DISPATCH(round);
 REGISTER_UNARY_TI_DISPATCH(sigmoid);
 REGISTER_DISPATCH(special_erfcx_stub, erfcx_kernel);
 REGISTER_DISPATCH(round_decimals_stub, round_decimals_kernel);
 REGISTER_DISPATCH(pow_tensor_scalar_stub, pow_tensor_scalar_kernel);
+REGISTER_DISPATCH(polygamma_stub, polygamma_kernel);
 } // namespace at::native
diff --git a/aten/src/ATen/native/mps/operations/UnaryOps.mm b/aten/src/ATen/native/mps/operations/UnaryOps.mm
index eb827d9f44c26..9dc318a393a5c 100644
--- a/aten/src/ATen/native/mps/operations/UnaryOps.mm
+++ b/aten/src/ATen/native/mps/operations/UnaryOps.mm
@@ -3,6 +3,7 @@
 #include <ATen/native/UnaryOps.h>
 #include <ATen/native/mps/Copy.h>
 #include <ATen/native/mps/OperationUtils.h>
+#include <ATen/native/mps/operations/ScanKernel.h>
 
 #ifndef AT_PER_OPERATOR_HEADERS
 #include <ATen/Functions.h>
@@ -333,8 +334,21 @@ static void cumulative_op_impl(const Tensor& self,
               "(original dim is ",
               dim,
               ")");
-  TORCH_CHECK(!self.is_complex(), "cumulative ops are not yet supported for complex");
   auto input = dtype.has_value() ? self.to(dtype.value()) : self;
+  if (input.is_complex()) {
+    if (cumulativeOpType == MPSCumulativeOpType::CUMSUM) {
+      auto input_real = at::view_as_real(input.dim() == 0 ? input.view({1}) : input);
+      auto result_real = at::view_as_real(result.dim() == 0 ? result.view({1}) : result);
+      return cumulative_op_impl(
+          input_real, wrapped_dim, std::nullopt, result_real, MPSCumulativeOpType::CUMSUM, "cumsum_out_mps");
+    } else if (cumulativeOpType == MPSCumulativeOpType::CUMPROD) {
+      auto input_view = input.dim() == 0 ? input.view({1}) : input;
+      auto result_view = result.dim() == 0 ? result.view({1}) : result;
+      return mps::scan_simple_mps_impl(input_view, result_view, wrapped_dim, "cumprod");
+    } else {
+      TORCH_INTERNAL_ASSERT(false);
+    }
+  }
 
   // issue #103810551: cumsum / cumprod are broken for int8, int16 and as chances for overflow are pretty high, cast to
   // int32 fixed in macOS 13.3
diff --git a/aten/src/ATen/native/mtia/EmptyTensor.cpp b/aten/src/ATen/native/mtia/EmptyTensor.cpp
deleted file mode 100644
index e7ff8b719aa21..0000000000000
--- a/aten/src/ATen/native/mtia/EmptyTensor.cpp
+++ /dev/null
@@ -1,86 +0,0 @@
-#include <ATen/Context.h>
-#include <ATen/EmptyTensor.h>
-#include <ATen/native/mtia/EmptyTensor.h>
-#include <c10/core/Allocator.h>
-#include <c10/core/DeviceGuard.h>
-
-namespace at::detail {
-
-at::Allocator* GetMTIAAllocator() {
-  return GetAllocator(DeviceType::MTIA);
-}
-
-TensorBase empty_mtia(
-    IntArrayRef size,
-    ScalarType dtype,
-    std::optional<Device> device_opt,
-    std::optional<c10::MemoryFormat> memory_format_opt) {
-  at::globalContext().lazyInitDevice(c10::DeviceType::MTIA);
-  const auto device = device_or_default(device_opt);
-  TORCH_INTERNAL_ASSERT(device.is_mtia());
-  const DeviceGuard device_guard(device);
-  auto* allocator = GetMTIAAllocator();
-  constexpr c10::DispatchKeySet mtia_dks(c10::DispatchKey::MTIA);
-  return at::detail::empty_generic(
-      size, allocator, mtia_dks, dtype, memory_format_opt);
-}
-
-TensorBase empty_mtia(
-    IntArrayRef size,
-    std::optional<ScalarType> dtype_opt,
-    std::optional<Layout> layout_opt,
-    std::optional<Device> device_opt,
-    std::optional<bool> pin_memory_opt,
-    std::optional<c10::MemoryFormat> memory_format_opt) {
-  const auto dtype = dtype_or_default(dtype_opt);
-  return at::detail::empty_mtia(size, dtype, device_opt, memory_format_opt);
-}
-
-TensorBase empty_mtia(IntArrayRef size, const TensorOptions& options) {
-  return at::detail::empty_mtia(
-      size,
-      optTypeMetaToScalarType(options.dtype_opt()),
-      options.layout_opt(),
-      options.device_opt(),
-      options.pinned_memory_opt(),
-      options.memory_format_opt());
-}
-
-TensorBase empty_strided_mtia(
-    IntArrayRef size,
-    IntArrayRef stride,
-    ScalarType dtype,
-    std::optional<Device> device_opt) {
-  at::globalContext().lazyInitDevice(c10::DeviceType::MTIA);
-  const auto device = device_or_default(device_opt);
-  const DeviceGuard device_guard(device);
-  auto* allocator = GetMTIAAllocator();
-  constexpr c10::DispatchKeySet mtia_dks(c10::DispatchKey::MTIA);
-  return at::detail::empty_strided_generic(
-      size, stride, allocator, mtia_dks, dtype);
-}
-
-TensorBase empty_strided_mtia(
-    IntArrayRef size,
-    IntArrayRef stride,
-    std::optional<ScalarType> dtype_opt,
-    std::optional<Layout> layout_opt,
-    std::optional<Device> device_opt,
-    std::optional<bool> pin_memory_opt) {
-  const auto dtype = dtype_or_default(dtype_opt);
-  return at::detail::empty_strided_mtia(size, stride, dtype, device_opt);
-}
-
-TensorBase empty_strided_mtia(
-    IntArrayRef size,
-    IntArrayRef stride,
-    const TensorOptions& options) {
-  return at::detail::empty_strided_mtia(
-      size,
-      stride,
-      optTypeMetaToScalarType(options.dtype_opt()),
-      options.layout_opt(),
-      options.device_opt(),
-      options.pinned_memory_opt());
-}
-} // namespace at::detail
diff --git a/aten/src/ATen/native/mtia/EmptyTensor.h b/aten/src/ATen/native/mtia/EmptyTensor.h
deleted file mode 100644
index afd1e58d40f07..0000000000000
--- a/aten/src/ATen/native/mtia/EmptyTensor.h
+++ /dev/null
@@ -1,42 +0,0 @@
-
-#pragma once
-#include <ATen/core/TensorBase.h>
-
-namespace at::detail {
-
-TensorBase empty_mtia(
-    IntArrayRef size,
-    ScalarType dtype,
-    std::optional<Device> device_opt,
-    std::optional<c10::MemoryFormat> memory_format_opt);
-
-TensorBase empty_mtia(
-    IntArrayRef size,
-    std::optional<ScalarType> dtype_opt,
-    std::optional<Layout> layout_opt,
-    std::optional<Device> device_opt,
-    std::optional<bool> pin_memory_opt,
-    std::optional<c10::MemoryFormat> memory_format_opt);
-
-TensorBase empty_mtia(IntArrayRef size, const TensorOptions& options);
-
-TensorBase empty_strided_mtia(
-    IntArrayRef size,
-    IntArrayRef stride,
-    ScalarType dtype,
-    std::optional<Device> device_opt);
-
-TensorBase empty_strided_mtia(
-    IntArrayRef size,
-    IntArrayRef stride,
-    std::optional<ScalarType> dtype_opt,
-    std::optional<Layout> layout_opt,
-    std::optional<Device> device_opt,
-    std::optional<bool> pin_memory_opt);
-
-TensorBase empty_strided_mtia(
-    IntArrayRef size,
-    IntArrayRef stride,
-    const TensorOptions& options);
-
-} // namespace at::detail
diff --git a/aten/src/ATen/native/native_functions.yaml b/aten/src/ATen/native/native_functions.yaml
index 43a818b782754..9e39eda0368fe 100644
--- a/aten/src/ATen/native/native_functions.yaml
+++ b/aten/src/ATen/native/native_functions.yaml
@@ -480,6 +480,7 @@
     CompositeExplicitAutograd: _conj_physical
     SparseCsrCPU, SparseCsrCUDA, SparseCsrMPS, SparseCsrMeta: conj_physical_sparse_csr
   autogen: _conj_physical.out
+  tags: pointwise
 
 - func: conj_physical(Tensor self) -> Tensor
   variants: function, method
@@ -1322,7 +1323,7 @@
 - func: logical_xor.out(Tensor self, Tensor other, *, Tensor(a!) out) -> Tensor(a!)
   device_check: NoCheck   # TensorIterator
   dispatch:
-    CPU, CUDA: logical_xor_out
+    CPU, CUDA, MTIA: logical_xor_out
     MPS: logical_xor_out_mps
   tags: pointwise
 
@@ -2735,8 +2736,7 @@
   device_check: NoCheck   # TensorIterator
   variants: function, method
   dispatch:
-    CPU, CUDA: fill_
-    MPS: fill_scalar_mps
+    CPU, CUDA, MPS: fill_
     QuantizedCPU, QuantizedCUDA: fill_quantized_
     Meta: fill_meta_
     SparseCsrCPU, SparseCsrCUDA, SparseCsrMeta: fill_sparse_csr_
@@ -2747,8 +2747,7 @@
   device_check: NoCheck   # TensorIterator
   variants: function, method
   dispatch:
-    CPU, CUDA: fill_
-    MPS: fill_tensor_mps_
+    CPU, CUDA, MPS: fill_
     QuantizedCPU, QuantizedCUDA: fill_quantized_
     Meta: fill_meta_
     NestedTensorCPU, NestedTensorHPU, NestedTensorCUDA: fill_nested_
@@ -2948,6 +2947,8 @@
   autogen: _grid_sampler_2d_cpu_fallback.out
 
 - func: _grid_sampler_2d_cpu_fallback_backward(Tensor grad_output, Tensor input, Tensor grid, int interpolation_mode, int padding_mode, bool align_corners) -> (Tensor, Tensor)
+  dispatch:
+    CompositeExplicitAutograd: _grid_sampler_2d_cpu_fallback_backward
 
 - func: grid_sampler_3d(Tensor input, Tensor grid, int interpolation_mode, int padding_mode, bool align_corners) -> Tensor
   dispatch:
@@ -3296,6 +3297,42 @@
   device_guard: False
   manual_cpp_binding: True
 
+- func: numel(Tensor self) -> int
+  variants: method
+  device_check: NoCheck
+  device_guard: False
+  manual_cpp_binding: True
+
+- func: dim(Tensor self) -> int
+  variants: method
+  device_check: NoCheck
+  device_guard: False
+  manual_cpp_binding: True
+
+- func: get_device(Tensor self) -> int
+  variants: method
+  device_check: NoCheck
+  device_guard: False
+  manual_cpp_binding: True
+
+- func: storage_offset(Tensor self) -> int
+  variants: method
+  device_check: NoCheck
+  device_guard: False
+  manual_cpp_binding: True
+
+- func: is_contiguous(Tensor self) -> bool
+  variants: method
+  device_check: NoCheck
+  device_guard: False
+  manual_cpp_binding: True
+
+- func: is_contiguous.memory_format(Tensor self, MemoryFormat memory_format) -> bool
+  variants: method
+  device_check: NoCheck
+  device_guard: False
+  manual_cpp_binding: True
+
 - func: kl_div(Tensor self, Tensor target, int reduction=Mean, *, bool log_target=False) -> Tensor
 
 - func: kron(Tensor self, Tensor other) -> Tensor
@@ -3350,11 +3387,13 @@
   dispatch:
     CUDA: _fused_rms_norm_cuda
     MPS: _fused_rms_norm_mps
+    XPU: _fused_rms_norm_xpu
     CompositeImplicitAutograd: rms_norm_composite
 
 - func: _fused_rms_norm_backward(Tensor grad_out, Tensor input, int[] normalized_shape, Tensor rstd, Tensor? weight, bool[2] output_mask) -> (Tensor, Tensor)
   dispatch:
     CUDA: _fused_rms_norm_backward_cuda
+    XPU: _fused_rms_norm_backward_xpu
 
 - func: nan_to_num(Tensor self, float? nan=None, float? posinf=None, float? neginf=None) -> Tensor
   variants: function, method
@@ -3691,8 +3730,7 @@
   structured_inherits: TensorIteratorBase
   variants: function
   dispatch:
-    CPU, CUDA: xlogy_out
-    MPS: xlogy_out_mps
+    CPU, CUDA, MPS: xlogy_out
   tags: pointwise
 
 - func: xlogy.OutScalar_Self(Scalar self, Tensor other, *, Tensor(a!) out) -> Tensor(a!)
@@ -3886,7 +3924,7 @@
   device_check: NoCheck   # TensorIterator
   structured: True
   dispatch:
-    CPU, CUDA, MTIA: aminmax_out
+    CPU, CUDA: aminmax_out
     MPS: aminmax_out_mps
   tags: reduction
 
@@ -3941,7 +3979,7 @@
 - func: amax.out(Tensor self, int[1] dim=[], bool keepdim=False, *, Tensor(a!) out) -> Tensor(a!)
   structured: True
   dispatch:
-    CPU, CUDA, MTIA: amax_out
+    CPU, CUDA: amax_out
     MPS: amax_out_mps
   tags: reduction
 
@@ -4130,7 +4168,7 @@
 - func: amin.out(Tensor self, int[1] dim=[], bool keepdim=False, *, Tensor(a!) out) -> Tensor(a!)
   structured: True
   dispatch:
-    CPU, CUDA, MTIA: amin_out
+    CPU, CUDA: amin_out
     MPS: amin_out_mps
   tags: reduction
 
@@ -4419,7 +4457,7 @@
 
 - func: mvlgamma.out(Tensor self, int p, *, Tensor(a!) out) -> Tensor(a!)
   dispatch:
-    CPU, CUDA: mvlgamma_out
+    CPU, CUDA, MPS: mvlgamma_out
   tags: pointwise
 
 - func: mvlgamma(Tensor self, int p) -> Tensor
@@ -5381,6 +5419,7 @@
 
 - func: selu_(Tensor(a!) self) -> Tensor(a!)
   device_check: NoCheck   # TensorIterator
+  tags: pointwise
 
 - func: celu(Tensor self, Scalar alpha=1.0) -> Tensor
   device_check: NoCheck   # TensorIterator
@@ -5393,6 +5432,7 @@
   dispatch:
     CompositeExplicitAutograd: celu_
   autogen: celu.out
+  tags: pointwise
 
 - func: silu(Tensor self) -> Tensor
   structured_delegate: silu.out
@@ -5413,8 +5453,7 @@
   structured_inherits: TensorIteratorBase
   python_module: nn
   dispatch:
-    CPU, CUDA, MTIA: silu_out
-    MPS: silu_out_mps
+    CPU, CUDA, MPS, MTIA: silu_out
   tags: pointwise
 
 - func: silu_backward.grad_input(Tensor grad_output, Tensor self, *, Tensor(a!) grad_input) -> Tensor(a!)
@@ -5422,8 +5461,7 @@
   structured_inherits: TensorIteratorBase
   python_module: nn
   dispatch:
-    CPU, CUDA: silu_backward_out
-    MPS: silu_backward_out_mps
+    CPU, CUDA, MPS: silu_backward_out
   tags: pointwise
 
 - func: silu_backward(Tensor grad_output, Tensor self) -> Tensor
@@ -6656,7 +6694,6 @@
   dispatch:
     CPU, CUDA: var
     MPS: var_mps
-    MTIA: var_mtia
   tags: [core, reduction]
 
 - func: var.out(Tensor self, int[1]? dim, bool unbiased=True, bool keepdim=False, *, Tensor(a!) out) -> Tensor(a!)
@@ -6818,6 +6855,7 @@
   dispatch:
     CPU: _standard_gamma_grad_cpu
     CUDA: _standard_gamma_grad_cuda
+    MPS: _standard_gamma_grad_mps
   autogen: _standard_gamma_grad.out
 
 - func: _standard_gamma(Tensor self, Generator? generator=None) -> Tensor
@@ -6825,9 +6863,32 @@
   dispatch:
     CPU: _s_gamma_cpu
     CUDA: _s_gamma_cuda
+    MPS: _s_gamma_mps
   tags: nondeterministic_seeded
   autogen: _standard_gamma.out
 
+- func: _philox_key_split(Tensor key, int num_splits) -> Tensor
+  variants: function
+  dispatch:
+    CUDA: _philox_key_split_cuda
+
+- func: _philox_key_fold_in(Tensor key, int data) -> Tensor
+  variants: function
+  dispatch:
+    CUDA: _philox_key_fold_in_cuda
+
+- func: _philox_normal_(Tensor(a!) self, Tensor key, float mean=0, float std=1) -> Tensor(a!)
+  variants: function, method
+  dispatch:
+    CUDA: _philox_normal_cuda_
+  autogen: _philox_normal, _philox_normal.out
+
+- func: _philox_uniform_(Tensor(a!) self, Tensor key, float low=0, float high=1) -> Tensor(a!)
+  variants: function, method
+  dispatch:
+    CUDA: _philox_uniform_cuda_
+  autogen: _philox_uniform, _philox_uniform.out
+
 - func: _dirichlet_grad(Tensor x, Tensor alpha, Tensor total) -> Tensor
   dispatch:
     CPU: _dirichlet_grad_cpu
@@ -7016,16 +7077,14 @@
   structured: True
   device_check: NoCheck   # TensorIterator
   dispatch:
-    CPU, CUDA: norm_dtype_out
-    MPS: norm_dtype_out_mps
+    CPU, CUDA, MPS: norm_dtype_out
   tags: reduction
 
 - func: norm.out(Tensor self, Scalar? p, int[1] dim, bool keepdim=False, *, Tensor(a!) out) -> Tensor(a!)
   structured: True
   device_check: NoCheck   # TensorIterator
   dispatch:
-    CPU, CUDA: norm_out
-    MPS: norm_out_mps
+    CPU, CUDA, MPS: norm_out
   tags: reduction
 
 # These four redispatch in their implementation, so OK to be CompositeImplicitAutograd
@@ -7118,8 +7177,7 @@
   device_check: NoCheck   # TensorIterator
   variants: method, function
   dispatch:
-    CPU, CUDA: zero_
-    MPS: zero_mps_
+    CPU, CUDA, MPS: zero_
     Meta: zero_meta_
     SparseCPU, SparseCUDA, SparseMPS, SparseMeta: zero_sparse_
     SparseCsrCPU, SparseCsrCUDA, SparseCsrMeta: zero_sparse_csr_
@@ -7327,6 +7385,7 @@
 - func: _scaled_mm_v2(Tensor self, Tensor mat2, Tensor[] scale_a, int[] recipe_a, int[] swizzle_a, Tensor[] scale_b, int[] recipe_b, int[] swizzle_b, Tensor? bias, ScalarType? out_dtype, int[] contraction_dim=[], bool use_fast_accum=False) -> Tensor
   variants: function
   dispatch:
+    CPU: _scaled_mm_cpu_v2
     CUDA: _scaled_mm_cuda_v2
     XPU: _scaled_mm_xpu_v2
   tags: needs_exact_strides
@@ -7334,6 +7393,7 @@
 - func: _scaled_mm_v2.out(Tensor self, Tensor mat2, Tensor[] scale_a, int[] recipe_a, int[] swizzle_a, Tensor[] scale_b, int[] recipe_b, int[] swizzle_b, Tensor? bias, ScalarType? out_dtype, int[] contraction_dim=[], bool use_fast_accum=False, *, Tensor(a!) out) -> Tensor(a!)
   variants: function
   dispatch:
+    CPU: _scaled_mm_cpu_v2_out
     CUDA: _scaled_mm_cuda_v2_out
     XPU: _scaled_mm_xpu_v2_out
   tags: needs_exact_strides
@@ -8422,6 +8482,7 @@
   dispatch:
     CPU: index_reduce_cpu_out
     CUDA: index_reduce_cuda_out
+    MPS: index_reduce_mps_out
 
 - func: index_reduce_(Tensor(a!) self, int dim, Tensor index, Tensor source, str reduce, *, bool include_self=True) -> Tensor(a!)
   structured_delegate: index_reduce.out
@@ -8435,9 +8496,7 @@
   device_check: NoCheck   # TensorIterator
   variants: method
   dispatch:
-    CPU: index_fill_
-    CUDA: index_fill_
-    MPS: index_fill_mps_
+    CPU, CUDA, MPS: index_fill_
   autogen: index_fill.int_Scalar_out
 
 - func: index_fill.int_Scalar(Tensor self, int dim, Tensor index, Scalar value) -> Tensor
@@ -8450,8 +8509,7 @@
   device_check: NoCheck   # TensorIterator
   variants: method
   dispatch:
-    CPU, CUDA: index_fill_
-    MPS: index_fill_mps_
+    CPU, CUDA, MPS: index_fill_
   autogen: index_fill.int_Tensor_out
 
 - func: index_fill.int_Tensor(Tensor self, int dim, Tensor index, Tensor value) -> Tensor
@@ -9570,12 +9628,14 @@
   dispatch:
     CPU: nonzero_static_out_cpu
     CUDA: nonzero_static_out_cuda
+    MPS: nonzero_static_out_mps
 
 - func: nonzero_static(Tensor self, *, SymInt size, int fill_value=-1) -> Tensor
   variants: method, function
   dispatch:
     CPU: nonzero_static_cpu
     CUDA: nonzero_static_cuda
+    MPS: nonzero_static_mps
 
 - func: nonzero_numpy(Tensor self) -> Tensor[]
   variants: method, function
@@ -9740,6 +9800,7 @@
   dispatch:
     CPU: _cholesky_solve_helper_cpu
     CUDA: _cholesky_solve_helper_cuda
+    MPS: _cholesky_solve_helper_mps
   autogen: _cholesky_solve_helper.out
 
 - func: cholesky_inverse(Tensor self, bool upper=False) -> Tensor
@@ -9818,8 +9879,7 @@
   structured: True
   structured_inherits: TensorIteratorBase
   dispatch:
-    CPU, CUDA: lgamma_out
-    MPS: lgamma_out_mps
+    CPU, CUDA, MPS: lgamma_out
   tags: pointwise
 
 - func: lgamma_(Tensor(a!) self) -> Tensor(a!)
@@ -9839,8 +9899,7 @@
   structured: True
   structured_inherits: TensorIteratorBase
   dispatch:
-    CPU, CUDA: digamma_out
-    MPS: digamma_out_mps
+    CPU, CUDA, MPS: digamma_out
   tags: pointwise
 
 - func: digamma(Tensor self) -> Tensor
@@ -9854,8 +9913,7 @@
   structured: True
   structured_inherits: TensorIteratorBase
   dispatch:
-    CPU, CUDA: polygamma_out
-    MPS: polygamma_out_mps
+    CPU, CUDA, MPS: polygamma_out
   tags: pointwise
 
 - func: polygamma(int n, Tensor self) -> Tensor
@@ -10014,8 +10072,7 @@
   structured: True
   structured_inherits: TensorIteratorBase
   dispatch:
-    CPU, CUDA: lerp_Tensor
-    MPS: lerp_Tensor_mps
+    CPU, CUDA, MPS: lerp_Tensor
   tags: pointwise
 
 - func: lerp.Scalar(Tensor self, Tensor end, Scalar weight) -> Tensor
@@ -10300,8 +10357,7 @@
   structured_inherits: TensorIteratorBase
   device_check: NoCheck   # TensorIterator
   dispatch:
-    CPU, CUDA, MTIA: maximum_out
-    MPS: maximum_out_mps
+    CPU, CUDA, MTIA, MPS: maximum_out
   tags: pointwise
 
 # binary max, alias of maximum
@@ -10333,8 +10389,7 @@
   structured_inherits: TensorIteratorBase
   device_check: NoCheck   # TensorIterator
   dispatch:
-    CPU, CUDA, MTIA: minimum_out
-    MPS: minimum_out_mps
+    CPU, CUDA, MTIA, MPS: minimum_out
   tags: pointwise
 
 # binary min, alias for minimum
@@ -11802,6 +11857,14 @@
     CUDA: foreach_tensor_zero_cuda_
   autogen: _foreach_zero, _foreach_zero.out
 
+- func: _foreach_clone(Tensor[] self, *, MemoryFormat? memory_format=None) -> Tensor[]
+  device_check: NoCheck   # foreach kernels fall back to slow path when tensor are on different devices
+  variants: function
+  dispatch:
+    CompositeExplicitAutograd: foreach_tensor_clone_slow
+    CUDA: foreach_tensor_clone_cuda
+  autogen: _foreach_clone.out
+
 - func: _foreach_copy_(Tensor(a!)[] self, Tensor[] src, bool non_blocking=False) -> ()
   device_check: NoCheck   # foreach kernels fall back to slow path when tensor are on different devices
   variants: function
@@ -12135,6 +12198,7 @@
   structured_delegate: elu.out
   device_check: NoCheck   # TensorIterator
   python_module: nn
+  tags: pointwise
 
 - func: glu.out(Tensor self, int dim=-1, *, Tensor(a!) out) -> Tensor(a!)
   structured: True
@@ -12196,6 +12260,7 @@
   structured_delegate: hardsigmoid.out
   device_check: NoCheck   # TensorIterator
   python_module: nn
+  tags: pointwise
 
 - func: hardsigmoid_backward.grad_input(Tensor grad_output, Tensor self, *, Tensor(a!) grad_input) -> Tensor(a!)
   structured: True
@@ -12241,6 +12306,7 @@
   dispatch:
     CPU, CUDA, MPS: hardtanh_
     QuantizedCPU: hardtanh_quantized_cpu_
+  tags: pointwise
 
 - func: hardswish.out(Tensor self, *, Tensor(a!) out) -> Tensor(a!)
   device_check: NoCheck   # TensorIterator
@@ -12281,7 +12347,7 @@
   python_module: nn
   dispatch:
     QuantizedCPU: leaky_relu_quantized_cpu
-  tags: core
+  tags: [core, pointwise]
 
 - func: leaky_relu_backward.grad_input(Tensor grad_output, Tensor self, Scalar negative_slope, bool self_is_result, *, Tensor(a!) grad_input) -> Tensor(a!)
   structured: True
@@ -12300,6 +12366,7 @@
   python_module: nn
   dispatch:
     QuantizedCPU: leaky_relu_quantized_cpu_
+  tags: pointwise
 
 - func: log_sigmoid.out(Tensor self, *, Tensor(a!) out) -> Tensor(a!)
   device_check: NoCheck   # TensorIterator
@@ -12956,6 +13023,10 @@
   python_module: nn
   autogen: _upsample_bicubic2d_aa.vec_out
 
+- func: _upsample_lanczos2d_aa.vec(Tensor input, SymInt[]? output_size, bool align_corners, float[]? scale_factors) -> Tensor
+  python_module: nn
+  autogen: _upsample_lanczos2d_aa.vec_out
+
 - func: upsample_nearest1d.vec(Tensor input, SymInt[]? output_size, float[]? scale_factors) -> Tensor
   python_module: nn
   autogen: upsample_nearest1d.vec_out
@@ -13102,6 +13173,26 @@
   python_module: nn
   structured_delegate: _upsample_bicubic2d_aa_backward.grad_input
 
+- func: _upsample_lanczos2d_aa.out(Tensor self, SymInt[2] output_size, bool align_corners, float? scales_h=None, float? scales_w=None, *, Tensor(a!) out) -> Tensor(a!)
+  python_module: nn
+  structured: True
+  dispatch:
+    CPU: _upsample_lanczos2d_aa_out_cpu
+
+- func: _upsample_lanczos2d_aa(Tensor self, SymInt[2] output_size, bool align_corners, float? scales_h=None, float? scales_w=None) -> Tensor
+  python_module: nn
+  structured_delegate: _upsample_lanczos2d_aa.out
+
+- func: _upsample_lanczos2d_aa_backward.grad_input(Tensor grad_output, SymInt[2] output_size, SymInt[4] input_size, bool align_corners, float? scales_h=None, float? scales_w=None, *, Tensor(a!) grad_input) -> Tensor(a!)
+  python_module: nn
+  structured: True
+  dispatch:
+    CPU: _upsample_lanczos2d_aa_backward_out_cpu
+
+- func: _upsample_lanczos2d_aa_backward(Tensor grad_output, SymInt[2] output_size, SymInt[4] input_size, bool align_corners, float? scales_h=None, float? scales_w=None) -> Tensor
+  python_module: nn
+  structured_delegate: _upsample_lanczos2d_aa_backward.grad_input
+
 - func: upsample_trilinear3d.out(Tensor self, SymInt[3] output_size, bool align_corners, float? scales_d=None, float? scales_h=None, float? scales_w=None, *, Tensor(a!) out) -> Tensor(a!)
   python_module: nn
   structured: True
@@ -14523,8 +14614,7 @@
   python_module: linalg
   structured: True
   dispatch:
-    CPU, CUDA: linalg_vector_norm_out
-    MPS: linalg_vector_norm_out_mps
+    CPU, CUDA, MPS: linalg_vector_norm_out
   tags: reduction
 
 # Computes sum(|x|^ord) - the "power sum" without the final root.
@@ -14685,6 +14775,7 @@
   structured: True
   dispatch:
     CPU, CUDA: linalg_qr_out
+    MPS: linalg_qr_out_mps
 
 - func: linalg_matrix_power(Tensor self, int n) -> Tensor
   python_module: linalg
@@ -15185,7 +15276,7 @@
   variants: function
   tags: nondeterministic_seeded
 
-- func: _scaled_dot_product_attention_math_for_mps(Tensor query, Tensor key, Tensor value, Tensor? attn_mask=None, float dropout_p=0.0, bool is_causal=False, Tensor? dropout_mask=None, *, float? scale=None) -> (Tensor, Tensor)
+- func: _scaled_dot_product_attention_math_for_mps(Tensor query, Tensor key, Tensor value, Tensor? attn_mask=None, float dropout_p=0.0, bool is_causal=False, Tensor? dropout_mask=None, *, float? scale=None, bool enable_gqa=False) -> (Tensor, Tensor)
   dispatch:
     MPS: _scaled_dot_product_attention_math_mps
   tags: nondeterministic_seeded
@@ -15257,12 +15348,18 @@
     NestedTensorCUDA: _scaled_dot_product_cudnn_attention_nestedtensor_backward_cuda
   tags: nondeterministic_seeded
 
-- func: _flash_attention_forward(Tensor query, Tensor key, Tensor value, Tensor? cum_seq_q, Tensor? cum_seq_k, SymInt max_q, SymInt max_k, float dropout_p, bool is_causal, bool return_debug_mask, *, float? scale=None, SymInt? window_size_left=None, SymInt? window_size_right=None, Tensor? seqused_k=None, Tensor? alibi_slopes=None) -> (Tensor output, Tensor softmax_logsumexp, Tensor rng_state, Tensor unused, Tensor debug_attn_mask)
+- func: _flash_attention_forward(Tensor query, Tensor key, Tensor value, Tensor? cum_seq_q, Tensor? cum_seq_k, SymInt max_q, SymInt max_k, float dropout_p, bool is_causal, bool return_debug_mask, *, float? scale=None, SymInt? window_size_left=None, SymInt? window_size_right=None, Tensor? seqused_k=None, Tensor? alibi_slopes=None, Tensor? block_table=None, int? num_splits=None) -> (Tensor output, Tensor softmax_logsumexp, Tensor rng_state, Tensor unused, Tensor debug_attn_mask)
   variants: function
   dispatch:
     CUDA: _flash_attention_forward
   tags: nondeterministic_seeded
 
+- func: _flash_attention_forward_no_dropout_inplace(Tensor(a!) out, Tensor query, Tensor key, Tensor value, Tensor? cum_seq_q, Tensor? cum_seq_k, SymInt max_q, SymInt max_k, float dropout_p, bool is_causal, bool return_debug_mask, *, float? scale=None, SymInt? window_size_left=None, SymInt? window_size_right=None, Tensor? seqused_k=None, Tensor? alibi_slopes=None, Tensor? block_table=None, int? num_splits=None) -> Tensor softmax_logsumexp
+  variants: function
+  dispatch:
+    CUDA: _flash_attention_forward_no_dropout_inplace
+  tags: nondeterministic_seeded
+
 - func: _flash_attention_forward.quantized(Tensor query, Tensor key, Tensor value, Tensor? cum_seq_q, Tensor? cum_seq_k, SymInt max_q, SymInt max_k, float dropout_p, bool is_causal, bool return_debug_mask, Tensor? q_descale, Tensor? k_descale, Tensor? v_descale, *, float? scale=None, SymInt? window_size_left=None, SymInt? window_size_right=None, Tensor? seqused_k=None, Tensor? alibi_slopes=None) -> (Tensor output, Tensor softmax_logsumexp, Tensor rng_state, Tensor unused, Tensor debug_attn_mask)
   variants: function
   dispatch:
diff --git a/aten/src/ATen/native/nested/NestedTensorBackward.cpp b/aten/src/ATen/native/nested/NestedTensorBackward.cpp
index 841b67addd1bd..80c9119fa1717 100644
--- a/aten/src/ATen/native/nested/NestedTensorBackward.cpp
+++ b/aten/src/ATen/native/nested/NestedTensorBackward.cpp
@@ -30,7 +30,7 @@ std::tuple<Tensor, Tensor> matmul_backward_nested(
   if (grad_input_mask[1]) {
     grad_other = at::matmul(self.transpose(-1, -2), grad);
   }
-  return std::make_tuple(grad_self, grad_other);
+  return std::make_tuple(std::move(grad_self), std::move(grad_other));
 }
 
 std::tuple<Tensor, Tensor, Tensor> nested_linear_backward(
diff --git a/aten/src/ATen/native/nested/NestedTensorMath.cpp b/aten/src/ATen/native/nested/NestedTensorMath.cpp
index 318bbb3728a85..120af06ca840a 100644
--- a/aten/src/ATen/native/nested/NestedTensorMath.cpp
+++ b/aten/src/ATen/native/nested/NestedTensorMath.cpp
@@ -615,7 +615,6 @@ Tensor squeeze_nested(const Tensor& self) {
   "squeeze(): For nested tensors, squeeze without the dim argument is not supported ",
   "at the moment, however you can use squeeze(Tensor self, int dim) instead ",
   "if you need this feature, please open an issue on github describing your use case.");
-  return self;
 }
 
 Tensor squeeze_dim_nested(const Tensor& self, IntArrayRef dims) {
@@ -1022,6 +1021,7 @@ static Tensor cat_nested_as_jagged(
   const auto first_item_dim = first_item.dim();
   const auto first_item_batch_size = first_item.size(0);
   std::vector<Tensor> jagged_views;
+  jagged_views.reserve(tensors.size());
   for (auto i : c10::irange(tensors.size())) {
     auto t = tensors[i].get();
     TORCH_CHECK(t.is_nested(),
@@ -1073,6 +1073,8 @@ static Tensor cat_nested_impl(
     // handle simple case of dim=0: concat NT components
     std::vector<at::Tensor> buffers;
     std::vector<at::Tensor> sizes;
+    buffers.reserve(tensors.size());
+    sizes.reserve(tensors.size());
     for (const auto i : c10::irange(tensors.size())) {
       const Tensor& t = tensors[i];
       TORCH_CHECK(
diff --git a/aten/src/ATen/native/nested/NestedTensorUtils.h b/aten/src/ATen/native/nested/NestedTensorUtils.h
index 47b8e2ad8086f..43556def235ec 100644
--- a/aten/src/ATen/native/nested/NestedTensorUtils.h
+++ b/aten/src/ATen/native/nested/NestedTensorUtils.h
@@ -433,6 +433,8 @@ inline Tensor wrap_tensor_node(
   } else { // Slow path
     std::vector<Tensor> flat_tensors;
     std::vector<Tensor> sizes;
+    flat_tensors.reserve(tensor_node.degree());
+    sizes.reserve(tensor_node.degree());
     for (const auto i : c10::irange(tensor_node.degree())) {
       flat_tensors.push_back(tensor_node.children(i).reshape(-1).contiguous());
       sizes.push_back(
diff --git a/aten/src/ATen/native/nested/cuda/NestedTensorTransformerFunctions.cpp b/aten/src/ATen/native/nested/cuda/NestedTensorTransformerFunctions.cpp
index 96c6ab8310f80..ff2213deeab44 100644
--- a/aten/src/ATen/native/nested/cuda/NestedTensorTransformerFunctions.cpp
+++ b/aten/src/ATen/native/nested/cuda/NestedTensorTransformerFunctions.cpp
@@ -463,7 +463,7 @@ std::tuple<at::Tensor, at::Tensor, at::Tensor> _scaled_dot_product_flash_attenti
   grad_k = wrap_buffer(grad_k.view(-1), key.transpose(1,2)._nested_tensor_size()).transpose(1,2);
   grad_v = wrap_buffer(grad_v.view(-1), value.transpose(1,2)._nested_tensor_size()).transpose(1,2);
 
-  return std::make_tuple(grad_q, grad_k, grad_v);
+  return std::make_tuple(std::move(grad_q), std::move(grad_k), std::move(grad_v));
 }
 
 } // namespace at::native
diff --git a/aten/src/ATen/native/quantized/FakeQuantPerChannelAffine.cpp b/aten/src/ATen/native/quantized/FakeQuantPerChannelAffine.cpp
index ffe6f4c31829d..af4d543648180 100644
--- a/aten/src/ATen/native/quantized/FakeQuantPerChannelAffine.cpp
+++ b/aten/src/ATen/native/quantized/FakeQuantPerChannelAffine.cpp
@@ -260,6 +260,6 @@ std::tuple<Tensor, Tensor, Tensor> _fake_quantize_learnable_per_channel_affine_b
   auto dScale = dScale_vec.sum(at::IntArrayRef(axis_for_reduction.data(), numElements));
   auto dZeroPoint = dZeroPoint_vec.sum(at::IntArrayRef(axis_for_reduction.data(), numElements));
 
-  return std::make_tuple(dX, dScale, dZeroPoint);
+  return std::make_tuple(std::move(dX), std::move(dScale), std::move(dZeroPoint));
 }
 } // namespace at::native
diff --git a/aten/src/ATen/native/quantized/FakeQuantPerTensorAffine.cpp b/aten/src/ATen/native/quantized/FakeQuantPerTensorAffine.cpp
index 88ac05cffe9e2..c491d0d9fe64e 100644
--- a/aten/src/ATen/native/quantized/FakeQuantPerTensorAffine.cpp
+++ b/aten/src/ATen/native/quantized/FakeQuantPerTensorAffine.cpp
@@ -86,7 +86,7 @@ std::tuple<Tensor, Tensor> fake_quantize_per_tensor_affine_cachemask(
       self.device().type(), Y, mask, self, scale, zero_point, quant_min, quant_max);
   // TODO(future, optional): look into packing the mask further (BoolTensor uses
   //   1 byte per element, we only need 1 bit per element).
-  return std::make_tuple(Y, mask);
+  return std::make_tuple(std::move(Y), std::move(mask));
 }
 
 std::tuple<Tensor, Tensor> _fake_quantize_per_tensor_affine_cachemask_tensor_qparams(
@@ -106,7 +106,7 @@ std::tuple<Tensor, Tensor> _fake_quantize_per_tensor_affine_cachemask_tensor_qpa
       self.device().type(), Y, mask, self, scale, zero_point, fake_quant_enabled, quant_min, quant_max);
   // TODO(future, optional): look into packing the mask further (BoolTensor uses
   //   1 byte per element, we only need 1 bit per element).
-  return std::make_tuple(Y, mask);
+  return std::make_tuple(std::move(Y), std::move(mask));
 }
 
 /* Backward path to fake-quantize the 'inputs' tensor, with mask.
@@ -231,7 +231,7 @@ std::tuple<Tensor, Tensor, Tensor> _fake_quantize_learnable_per_tensor_affine_ba
   auto dScale = dScale_vec.sum().unsqueeze(0).to(scale_.device());
   auto dZeroPoint = dZeroPoint_vec.sum().unsqueeze(0).to(zero_point_.device());
 
-  return std::make_tuple(dX, dScale, dZeroPoint);
+  return std::make_tuple(std::move(dX), std::move(dScale), std::move(dZeroPoint));
 }
 
 } // namespace at::native
diff --git a/aten/src/ATen/native/quantized/cpu/ChannelShuffle.cpp b/aten/src/ATen/native/quantized/cpu/ChannelShuffle.cpp
index 5df69d01b2549..36baffc36c1ed 100644
--- a/aten/src/ATen/native/quantized/cpu/ChannelShuffle.cpp
+++ b/aten/src/ATen/native/quantized/cpu/ChannelShuffle.cpp
@@ -102,10 +102,11 @@ Tensor channel_shuffle_quantized_cpu(
     int64_t groups) {
 #ifdef USE_PYTORCH_QNNPACK
   return quantized_channel_shuffle_impl(self, groups);
-#endif
+#else
   // If QNNPACK is not available then fall back to the
   // non quantized path.
   return at::native::channel_shuffle(self, groups);
+#endif
 }
 
 // Keep the registry in the anonymous namespace.
diff --git a/aten/src/ATen/native/quantized/cpu/qembeddingbag_unpack.cpp b/aten/src/ATen/native/quantized/cpu/qembeddingbag_unpack.cpp
index b115b25c42784..333796e8a24ce 100644
--- a/aten/src/ATen/native/quantized/cpu/qembeddingbag_unpack.cpp
+++ b/aten/src/ATen/native/quantized/cpu/qembeddingbag_unpack.cpp
@@ -94,7 +94,6 @@ at::Tensor PackedEmbeddingBagWeight::unpack() {
   TORCH_INTERNAL_ASSERT(
       false,
       "We currently only support 8-bit and 4-bit quantization of embedding_bag.");
-  return weight_origin;
 }
 
 namespace at::native {
diff --git a/aten/src/ATen/native/quantized/cpu/qnnpack/CMakeLists.txt b/aten/src/ATen/native/quantized/cpu/qnnpack/CMakeLists.txt
index 799ec5e174083..19b7851b3e7ba 100644
--- a/aten/src/ATen/native/quantized/cpu/qnnpack/CMakeLists.txt
+++ b/aten/src/ATen/native/quantized/cpu/qnnpack/CMakeLists.txt
@@ -4,7 +4,7 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
-cmake_minimum_required(VERSION 3.5 FATAL_ERROR)
+cmake_minimum_required(VERSION 3.10 FATAL_ERROR)
 
 include(GNUInstallDirs)
 
diff --git a/aten/src/ATen/native/quantized/cudnn/Pooling.cpp b/aten/src/ATen/native/quantized/cudnn/Pooling.cpp
index 7fe44de11e54c..857d0c52140e9 100644
--- a/aten/src/ATen/native/quantized/cudnn/Pooling.cpp
+++ b/aten/src/ATen/native/quantized/cudnn/Pooling.cpp
@@ -65,7 +65,6 @@ Tensor adaptive_avg_pool2d_quantized_cuda(
     return at::quantize_per_tensor(result_fp32, input.q_scale(), input.q_zero_point(), input.scalar_type());
 #else // USE_CUDA
   TORCH_CHECK(false, "at::native::adaptive_avg_pool2d_quantized_cuda: ATen not compiled with USE_CUDA support");
-  return Tensor{}; // never reached, placates the compiler
 #endif
 }
 
@@ -214,7 +213,6 @@ Tensor quantized_max_pool2d_cudnn(
 #endif // AT_CUDNN_ENABLED()
 #else // USE_CUDA
   TORCH_CHECK(false, "at::native::quantized_max_pool2d_cudnn: ATen not compiled with USE_CUDA support");
-  return Tensor{}; // never reached, placates the compiler
 #endif
 }
 
diff --git a/aten/src/ATen/native/sparse/FlattenIndicesKernel.cpp b/aten/src/ATen/native/sparse/FlattenIndicesKernel.cpp
index 35e3ebaa9f8b6..ee41472a74118 100644
--- a/aten/src/ATen/native/sparse/FlattenIndicesKernel.cpp
+++ b/aten/src/ATen/native/sparse/FlattenIndicesKernel.cpp
@@ -22,11 +22,6 @@ Tensor flatten_indices_cpu_kernel(const Tensor& indices, IntArrayRef size) {
 
 }
 
-REGISTER_ARCH_DISPATCH(flatten_indices_stub, DEFAULT, &flatten_indices_cpu_kernel)
-REGISTER_AVX512_DISPATCH(flatten_indices_stub, &flatten_indices_cpu_kernel)
-REGISTER_AVX2_DISPATCH(flatten_indices_stub, &flatten_indices_cpu_kernel)
-REGISTER_VSX_DISPATCH(flatten_indices_stub, &flatten_indices_cpu_kernel)
-REGISTER_ZVECTOR_DISPATCH(flatten_indices_stub, &flatten_indices_cpu_kernel)
-REGISTER_SVE256_DISPATCH(flatten_indices_stub, &flatten_indices_cpu_kernel)
+REGISTER_ALL_CPU_DISPATCH(flatten_indices_stub, &flatten_indices_cpu_kernel)
 
 } // namespace at::native
diff --git a/aten/src/ATen/native/sparse/SparseBinaryOpIntersectionCommon.h b/aten/src/ATen/native/sparse/SparseBinaryOpIntersectionCommon.h
index 7cec767d44660..e50171f48f4c6 100644
--- a/aten/src/ATen/native/sparse/SparseBinaryOpIntersectionCommon.h
+++ b/aten/src/ATen/native/sparse/SparseBinaryOpIntersectionCommon.h
@@ -297,8 +297,8 @@ void _sparse_binary_op_intersection_kernel_impl(
   std::tie(sorted_hash, argsort_hash) = [&]() -> std::tuple<Tensor, Tensor> {
     if (probably_coalesced.is_coalesced()) {
       // NOTE: argsort.dtype == nnz_arange.dtype
-      const auto argsort = nnz_arange.narrow(-1, 0, probably_coalesced._nnz());
-      return std::make_tuple(probably_coalesced_indices_hash, argsort);
+      auto argsort = nnz_arange.narrow(-1, 0, probably_coalesced._nnz());
+      return std::make_tuple(probably_coalesced_indices_hash, std::move(argsort));
     } else {
       // NOTE: we want argsort.dtype == nnz_arange.dtype,
       // but sort() produces indices of type int64_t,
diff --git a/aten/src/ATen/native/sparse/SparseBinaryOpIntersectionKernel.cpp b/aten/src/ATen/native/sparse/SparseBinaryOpIntersectionKernel.cpp
index 979dbdd033ac3..5fc5df5abca24 100644
--- a/aten/src/ATen/native/sparse/SparseBinaryOpIntersectionKernel.cpp
+++ b/aten/src/ATen/native/sparse/SparseBinaryOpIntersectionKernel.cpp
@@ -155,24 +155,7 @@ void sparse_mask_projection_out_cpu_kernel(
 
 }
 
-REGISTER_ARCH_DISPATCH(mul_sparse_sparse_out_stub, DEFAULT, &mul_sparse_sparse_out_cpu_kernel)
-REGISTER_AVX512_DISPATCH(mul_sparse_sparse_out_stub, &mul_sparse_sparse_out_cpu_kernel)
-REGISTER_AVX2_DISPATCH(mul_sparse_sparse_out_stub, &mul_sparse_sparse_out_cpu_kernel)
-REGISTER_VSX_DISPATCH(mul_sparse_sparse_out_stub, &mul_sparse_sparse_out_cpu_kernel)
-REGISTER_ZVECTOR_DISPATCH(mul_sparse_sparse_out_stub, &mul_sparse_sparse_out_cpu_kernel)
-REGISTER_SVE256_DISPATCH(mul_sparse_sparse_out_stub, &mul_sparse_sparse_out_cpu_kernel)
-
-REGISTER_ARCH_DISPATCH(sparse_mask_intersection_out_stub, DEFAULT, &sparse_mask_intersection_out_cpu_kernel)
-REGISTER_AVX512_DISPATCH(sparse_mask_intersection_out_stub, &sparse_mask_intersection_out_cpu_kernel)
-REGISTER_AVX2_DISPATCH(sparse_mask_intersection_out_stub, &sparse_mask_intersection_out_cpu_kernel)
-REGISTER_VSX_DISPATCH(sparse_mask_intersection_out_stub, &sparse_mask_intersection_out_cpu_kernel)
-REGISTER_ZVECTOR_DISPATCH(sparse_mask_intersection_out_stub, &sparse_mask_intersection_out_cpu_kernel)
-REGISTER_SVE256_DISPATCH(sparse_mask_intersection_out_stub, &sparse_mask_intersection_out_cpu_kernel)
-
-REGISTER_ARCH_DISPATCH(sparse_mask_projection_out_stub, DEFAULT, &sparse_mask_projection_out_cpu_kernel)
-REGISTER_AVX512_DISPATCH(sparse_mask_projection_out_stub, &sparse_mask_projection_out_cpu_kernel)
-REGISTER_AVX2_DISPATCH(sparse_mask_projection_out_stub, &sparse_mask_projection_out_cpu_kernel)
-REGISTER_VSX_DISPATCH(sparse_mask_projection_out_stub, &sparse_mask_projection_out_cpu_kernel)
-REGISTER_ZVECTOR_DISPATCH(sparse_mask_projection_out_stub, &sparse_mask_projection_out_cpu_kernel)
-REGISTER_SVE256_DISPATCH(sparse_mask_projection_out_stub, &sparse_mask_projection_out_cpu_kernel)
+REGISTER_ALL_CPU_DISPATCH(mul_sparse_sparse_out_stub, &mul_sparse_sparse_out_cpu_kernel)
+REGISTER_ALL_CPU_DISPATCH(sparse_mask_intersection_out_stub, &sparse_mask_intersection_out_cpu_kernel)
+REGISTER_ALL_CPU_DISPATCH(sparse_mask_projection_out_stub, &sparse_mask_projection_out_cpu_kernel)
 }
diff --git a/aten/src/ATen/native/sparse/SparseCsrTensor.cpp b/aten/src/ATen/native/sparse/SparseCsrTensor.cpp
index 5ecfbd589ca39..79ce7fb17bfb3 100644
--- a/aten/src/ATen/native/sparse/SparseCsrTensor.cpp
+++ b/aten/src/ATen/native/sparse/SparseCsrTensor.cpp
@@ -189,8 +189,10 @@ static void _validate_sparse_compressed_tensor_args_worker(const Tensor& compres
               batch_ndim, " + ", block_ndim, ") but got ", values.dim());
 
   // 3.5
-  TORCH_CHECK(plain_indices.stride(-1) == 1,
-              "expected ", plain_indices_name, " to be a contiguous tensor per batch");
+  if (plain_indices.numel() != 0) {
+    TORCH_CHECK(plain_indices.stride(-1) == 1,
+                "expected ", plain_indices_name, " to be a contiguous tensor per batch");
+  }
 
   // 3.6
   TORCH_CHECK(compressed_indices.stride(-1) == 1,
diff --git a/aten/src/ATen/native/sparse/SparseCsrTensorMath.cpp b/aten/src/ATen/native/sparse/SparseCsrTensorMath.cpp
index edb95c23b98ba..cb8baf89a86e2 100644
--- a/aten/src/ATen/native/sparse/SparseCsrTensorMath.cpp
+++ b/aten/src/ATen/native/sparse/SparseCsrTensorMath.cpp
@@ -731,8 +731,8 @@ Tensor& addmm_out_sparse_compressed_cpu(
       " without MKL. PyTorch built with MKL has better support for addmm with sparse CPU tensors.");
 #else
   sparse::impl::mkl::addmm_out_sparse_csr(mat1, mat2, beta, alpha, result);
-#endif
   return result;
+#endif
 }
 
 Tensor addmm_sparse_compressed_dense(
diff --git a/aten/src/ATen/native/sparse/SparseTensor.cpp b/aten/src/ATen/native/sparse/SparseTensor.cpp
index 5dec3746eaa88..b18b4dfccabac 100644
--- a/aten/src/ATen/native/sparse/SparseTensor.cpp
+++ b/aten/src/ATen/native/sparse/SparseTensor.cpp
@@ -95,7 +95,6 @@ bool is_coalesced_sparse(const SparseTensor& self) {
 
 bool is_coalesced_default(const Tensor& self) {
   TORCH_CHECK(false, "is_coalesced expected sparse coordinate tensor layout but got ", self.layout());
-  return false;
 }
 
 int64_t _nnz_sparse(const SparseTensor& self) {
diff --git a/aten/src/ATen/native/sparse/cuda/SoftMax.cu b/aten/src/ATen/native/sparse/cuda/SoftMax.cu
index 19b8c377f796d..11653c28894f4 100644
--- a/aten/src/ATen/native/sparse/cuda/SoftMax.cu
+++ b/aten/src/ATen/native/sparse/cuda/SoftMax.cu
@@ -316,8 +316,13 @@ std::tuple<Tensor, Tensor, Tensor, Tensor> compute_pool_max(
       [offsets_ptr] __device__(int64_t x, int64_t y) {
         return offsets_ptr[x] == offsets_ptr[y];
       });
+#if !defined(USE_ROCM)
+  auto new_sz = ::cuda::std::distance(
+      thrust_ptr(pool_sizes.template data_ptr<int64_t>()), new_end.second);
+#else
   auto new_sz = thrust::distance(
       thrust_ptr(pool_sizes.template data_ptr<int64_t>()), new_end.second);
+#endif
   pool_sizes.resize_({new_sz});
 
   auto pool_offsets = pool_sizes.clone();
diff --git a/aten/src/ATen/native/sparse/cuda/SparseCUDAApplyUtils.cuh b/aten/src/ATen/native/sparse/cuda/SparseCUDAApplyUtils.cuh
index 078b354659d35..7b9e61ceb4cc7 100644
--- a/aten/src/ATen/native/sparse/cuda/SparseCUDAApplyUtils.cuh
+++ b/aten/src/ATen/native/sparse/cuda/SparseCUDAApplyUtils.cuh
@@ -246,11 +246,7 @@ __global__ void coalesceValuesKernel(
 // `if constexpr` when CUDA codes will be compiled under C++-17, see
 // gh-56055 for blockers.
 template<typename Dtype>
-#ifdef USE_ROCM
-C10_LAUNCH_BOUNDS_1(C10_WARP_SIZE_STATIC*4)
-#else
-C10_LAUNCH_BOUNDS_1(C10_WARP_SIZE*4)
-#endif
+C10_LAUNCH_BOUNDS_1(C10_WARP_SIZE_UPPER_BOUND*4)
 __global__ void coalesceValuesKernel(
   int64_t *segment_offsets, int64_t *value_indices,
   bool *values, bool *newValues,
diff --git a/aten/src/ATen/native/sparse/cuda/SparseCUDATensor.cu b/aten/src/ATen/native/sparse/cuda/SparseCUDATensor.cu
index b8afef136d98f..9726f391129f9 100644
--- a/aten/src/ATen/native/sparse/cuda/SparseCUDATensor.cu
+++ b/aten/src/ATen/native/sparse/cuda/SparseCUDATensor.cu
@@ -26,7 +26,6 @@
 #include <thrust/device_vector.h>
 #include <thrust/gather.h>
 #include <thrust/generate.h>
-#include <thrust/pair.h>
 #include <thrust/scan.h>
 #include <thrust/sequence.h>
 #include <thrust/sort.h>
diff --git a/aten/src/ATen/native/sparse/cuda/SparseCUDATensorMath.cu b/aten/src/ATen/native/sparse/cuda/SparseCUDATensorMath.cu
index dd39b67363b45..6489b2b46256e 100644
--- a/aten/src/ATen/native/sparse/cuda/SparseCUDATensorMath.cu
+++ b/aten/src/ATen/native/sparse/cuda/SparseCUDATensorMath.cu
@@ -16,6 +16,7 @@
 #include <ATen/WrapDimUtilsMulti.h>
 #include <ATen/ExpandUtils.h>
 #include <c10/cuda/CUDACachingAllocator.h>
+#include <c10/macros/Macros.h>
 
 #ifndef AT_PER_OPERATOR_HEADERS
 #include <ATen/Functions.h>
@@ -670,34 +671,22 @@ Tensor bmm_sparse_cuda(const SparseTensor& self, const Tensor& mat2) {
 }
 
 __global__ void search_end_matrix_indices_cuda_kernel(
-  int64_t* mat_el_end_indices,
+  int64_t* __restrict__ mat_el_end_indices,
   int64_t num_matrices,
-  const TensorInfo<int64_t, int64_t> indices_1D_ti,
+  const int64_t* __restrict__ indices_1D,
   const int64_t num_elements
 ){
-  const int64_t target_mat_num = ((int64_t) blockIdx.x) * blockDim.x + threadIdx.x;
+  const int64_t target_mat_num = static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x;
   if (target_mat_num >= num_matrices) return;
 
-  const int64_t* indices_1D = indices_1D_ti.data;
-  const int64_t indices_1D_stride = indices_1D_ti.strides[0];
-  int64_t start_idx = 0;
-  int64_t end_idx = num_elements - 1;
-
-  while (start_idx < end_idx) {
-    int64_t mid_idx = (start_idx + end_idx + 1) >> 1;
-    int64_t mat_num = indices_1D[mid_idx*indices_1D_stride];
-    if (mat_num > target_mat_num) {
-      end_idx = mid_idx - 1;
-    } else {
-      start_idx = mid_idx;
-    }
-  }
-
-  if (indices_1D[start_idx*indices_1D_stride] == target_mat_num) {
-    mat_el_end_indices[target_mat_num] = start_idx;
-  } else {
-    mat_el_end_indices[target_mat_num] = -1;
-  }
+  // Points to the upper-bound of `target_mat_num`
+  const auto target_mat_num_ub = at::cuda::detail::find_bound(
+    indices_1D,
+    indices_1D + num_elements,
+    target_mat_num + 1
+  );
+  // Store the "end" offset of each matrix into the "batch" dim
+  mat_el_end_indices[target_mat_num] = target_mat_num_ub - indices_1D;
 }
 
 // Search through a 1D tensor of sorted sparse matrix
@@ -707,20 +696,17 @@ void search_end_matrix_indices(int64_t* mat_el_end_indices, int64_t num_matrices
   c10::cuda::GetDevice(&curDevice);
   cudaStream_t stream = at::cuda::getCurrentCUDAStream(curDevice);
 
-  auto indices_1D_ti = getTensorInfo<int64_t, int64_t>(indices_1D);
-  int64_t grid_size = (num_matrices / 64)+1;
   int64_t block_size = 64;
+  int64_t grid_size = (num_matrices + block_size - 1) / block_size;
   int64_t num_elements = indices_1D.size(0);
 
   search_end_matrix_indices_cuda_kernel<<<grid_size, block_size, 0, stream>>>(
     mat_el_end_indices,
     num_matrices,
-    indices_1D_ti,
+    indices_1D.data_ptr<int64_t>(),
     num_elements
   );
   C10_CUDA_KERNEL_LAUNCH_CHECK();
-
-  cudaDeviceSynchronize();
 }
 
 cudaDataType getTensorCudaDataType(Tensor self) {
@@ -739,6 +725,12 @@ cudaDataType getTensorCudaDataType(Tensor self) {
   return cuda_data_type;
 }
 
+// cusparseSpMV bug on CUDA < 13.1 -> COO row index array needs to be 16-byte aligned.
+// See https://github.com/pytorch/pytorch/issues/167901
+#if defined(USE_CUDA) && CUDA_VERSION < 13010
+#define CUSPARSE_SPMV_ALIGNMENT_BUG_PRESENT
+#endif
+
 Tensor& bmm_out_sparse_cuda(const SparseTensor& self, const Tensor& mat2, Tensor& result) {
   TORCH_CHECK(!mat2.is_sparse(), "bmm_sparse: Tensor 'mat2' must be dense");
   TORCH_CHECK(self.dense_dim() == 0, "bmm_sparse: Tensor 'self' must have 0 dense dims, but has ", self.dense_dim());
@@ -752,10 +744,10 @@ Tensor& bmm_out_sparse_cuda(const SparseTensor& self, const Tensor& mat2, Tensor
   int64_t dim_j = self.size(2);
   int64_t dim_k = mat2.size(2);
 
-  result.resize_({num_matrices, dim_k, dim_i});
+  result.resize_({num_matrices, dim_i, dim_k});
 
   if ((self._nnz() == 0) || (dim_j == 0) || (dim_k == 0)) {
-    result.zero_().transpose_(1, 2);
+    result.zero_();
     return result;
   }
 
@@ -768,7 +760,7 @@ Tensor& bmm_out_sparse_cuda(const SparseTensor& self, const Tensor& mat2, Tensor
     tmp_result = result;
     need_copy_result = false;
   } else {
-    tmp_result = at::empty({num_matrices, dim_k, dim_i}, result.options(), at::MemoryFormat::Contiguous);
+    tmp_result = at::empty({num_matrices, dim_i, dim_k}, result.options(), at::MemoryFormat::Contiguous);
     need_copy_result = true;
   }
 
@@ -785,29 +777,56 @@ Tensor& bmm_out_sparse_cuda(const SparseTensor& self, const Tensor& mat2, Tensor
   Tensor values =      self_coalesced._values();
 
   Tensor indices_dim0 = indices[0];
+  Tensor indices_dim1 = indices[1];
+  Tensor indices_dim2 = indices[2];
 
-  // Need to convert dim1 and dim2 indices to 32-bit since cusparseSpMM
-  // only supports 32-bit indices
-  Tensor indices_dim1 = indices[1].to(ScalarType::Int);
-  Tensor indices_dim2 = indices[2].to(ScalarType::Int);
-
-  std::vector<int64_t> mat_el_end_indices_host(num_matrices);
-
-  {
-    auto& allocator = *::c10::cuda::CUDACachingAllocator::get();
-    auto dataPtr = allocator.allocate(num_matrices*sizeof(int64_t));
-    int64_t* mat_el_end_indices_device = static_cast<int64_t*>(dataPtr.get());
+  auto mat_el_end_indices_host = at::empty({num_matrices}, indices.options().device(at::kCPU).pinned_memory(true));
+  search_end_matrix_indices(mat_el_end_indices_host.data_ptr<int64_t>(), num_matrices, indices_dim0);
+  cudaDeviceSynchronize();
+  auto* mat_el_end_indices = mat_el_end_indices_host.data_ptr<int64_t>();
+
+  // cusparseSpMV bug on CUDA < 13.1 -> COO row index array needs to be 16-byte aligned,
+  // so we use a buffer for misaglined sub-arrays to copy into.
+#ifdef CUSPARSE_SPMV_ALIGNMENT_BUG_PRESENT
+  auto aligned_row_indices_buffer = [&]() -> Tensor {
+    if (dim_k == 1) { // implies cusparseSpMV
+      const auto* row_indices_start_ptr = indices_dim1.data_ptr<int64_t>();
+      const auto* mat_end_offsets_ptr = mat_el_end_indices_host.data_ptr<int64_t>();
+      int64_t max_nnz = 0;
+      int64_t start_offset = 0;
+      for (const auto i : c10::irange(num_matrices)) {
+        const auto* row_indices_ptr = row_indices_start_ptr + start_offset;
+        const int64_t row_indices_ptr_not_aligned = (
+            (reinterpret_cast<uintptr_t>(row_indices_ptr) % 16) / 8
+        );
+        const auto end_offset = mat_end_offsets_ptr[i];
+        const auto nnz = end_offset - start_offset;
+        max_nnz = std::max(row_indices_ptr_not_aligned * nnz, max_nnz);
+        start_offset = end_offset;
+      }
+      return max_nnz ? at::empty({max_nnz}, indices.options()) : Tensor{};
+    }
+    return Tensor{};
+  }();
+#endif
 
-    search_end_matrix_indices(mat_el_end_indices_device, num_matrices, indices_dim0);
-    AT_CUDA_CHECK(cudaMemcpy(
-      mat_el_end_indices_host.data(),
-      mat_el_end_indices_device,
-      num_matrices*sizeof(int64_t),
-      cudaMemcpyDeviceToHost
-    ));
-  }
-  // Need a pointer to an array to access within a lambda
-  int64_t* mat_el_end_indices = mat_el_end_indices_host.data();
+  // MSVC is not happy with having macros in AT_DISPATCH,
+  // so we are using a lambda which we try to force-inline
+  const auto maybe_provide_aligned_buffer_idx_ptr
+    = [&](int64_t* idx_ptr, int64_t start_offset, int64_t len) C10_ALWAYS_INLINE_ATTRIBUTE -> int64_t* {
+#ifdef CUSPARSE_SPMV_ALIGNMENT_BUG_PRESENT
+      auto* start = idx_ptr + start_offset;
+      const auto is_misaligned_start = (reinterpret_cast<uintptr_t>(start) % 16) != 0;
+      if (is_misaligned_start && aligned_row_indices_buffer.defined()) {
+        aligned_row_indices_buffer.narrow(0, 0, len)
+          .copy_(indices_dim1.narrow(0, start_offset, len));
+        return aligned_row_indices_buffer.data_ptr<int64_t>();
+      }
+      return start;
+#else
+      return idx_ptr + start_offset;
+#endif
+  };
 
   Scalar beta = 0;
   Scalar alpha = 1;
@@ -828,112 +847,106 @@ Tensor& bmm_out_sparse_cuda(const SparseTensor& self, const Tensor& mat2, Tensor
     values.scalar_type(), "bmm_sparse_cuda", [&] {
       scalar_t alpha_val = alpha.to<scalar_t>();
       scalar_t beta_val = beta.to<scalar_t>();
-      uint32_t* row_indices_start_ptr = reinterpret_cast<uint32_t*>(indices_dim1.data_ptr());
-      uint32_t* col_indices_start_ptr = reinterpret_cast<uint32_t*>(indices_dim2.data_ptr());
-      scalar_t* values_start_ptr = reinterpret_cast<scalar_t*>(values.data_ptr());
-      scalar_t* mat2_start_ptr = reinterpret_cast<scalar_t*>(mat2_contig.data_ptr());
-      scalar_t* result_start_ptr = reinterpret_cast<scalar_t*>(tmp_result.data_ptr());
+      auto* row_indices_start_ptr = indices_dim1.data_ptr<int64_t>();
+      auto* col_indices_start_ptr = indices_dim2.data_ptr<int64_t>();
+      auto* values_start_ptr = values.data_ptr<scalar_t>();
+      auto* mat2_start_ptr = mat2_contig.data_ptr<scalar_t>();
+      auto* result_start_ptr = tmp_result.data_ptr<scalar_t>();
       for (
         int64_t cur_mat_num = 0;
-        (cur_mat_num < num_matrices);
-        cur_mat_num++
+        cur_mat_num < num_matrices;
+        ++cur_mat_num
       ) {
         int64_t mat_el_end_idx = mat_el_end_indices[cur_mat_num];
 
-        if (mat_el_end_idx != -1) {
-          mat_el_end_idx++;
-
-          // Create tensors to view just the current set of matrices
-          int64_t sparse_nnz = mat_el_end_idx - mat_el_begin_idx;
-
-          cudaDataType cuda_data_type = getTensorCudaDataType(mat2_contig);
-          uint32_t* row_indices_ptr = &row_indices_start_ptr[mat_el_begin_idx];
-          uint32_t* col_indices_ptr = &col_indices_start_ptr[mat_el_begin_idx];
-          scalar_t* values_ptr = &values_start_ptr[mat_el_begin_idx];
-
-          cusparseSpMatDescr_t sparse_descr;
-          TORCH_CUDASPARSE_CHECK(cusparseCreateCoo(
-            &sparse_descr,
-            dim_i,
-            dim_j,
-            sparse_nnz,
-            reinterpret_cast<void*>(row_indices_ptr),
-            reinterpret_cast<void*>(col_indices_ptr),
-            reinterpret_cast<void*>(values_ptr),
-            CUSPARSE_INDEX_32I,
-            CUSPARSE_INDEX_BASE_ZERO,
-            cuda_data_type
-          ));
-          scalar_t* mat2_ptr = &mat2_start_ptr[dim_k*dim_j*cur_mat_num];
-          cusparseDnMatDescr_t dense_descr;
-          TORCH_CUDASPARSE_CHECK(cusparseCreateDnMat(
-            &dense_descr,
-            dim_k,
-            dim_j,
-            dim_k,
-            reinterpret_cast<void*>(mat2_ptr),
-            cuda_data_type,
-            CUSPARSE_ORDER_COL
-          ));
-          scalar_t* result_ptr = &result_start_ptr[dim_i*dim_k*cur_mat_num];
-          cusparseDnMatDescr_t result_descr;
-          TORCH_CUDASPARSE_CHECK(cusparseCreateDnMat(
-            &result_descr,
-            dim_i,
-            dim_k,
-            dim_i,
-            reinterpret_cast<void*>(result_ptr),
-            cuda_data_type,
-            CUSPARSE_ORDER_COL
-          ));
-          size_t required_workspace_buffer_size = 0;
-          TORCH_CUDASPARSE_CHECK(cusparseSpMM_bufferSize(
-            cusparse_handle,
-            CUSPARSE_OPERATION_NON_TRANSPOSE,
-            CUSPARSE_OPERATION_TRANSPOSE,
-            (void*)&alpha_val,
-            sparse_descr,
-            dense_descr,
-            (void*)&beta_val,
-            result_descr,
-            cuda_data_type,
-            mm_alg,
-            &required_workspace_buffer_size
-          ));
-          if (required_workspace_buffer_size > workspace_buffer_size) {
-            workspace_buffer_size = required_workspace_buffer_size;
-            dataPtr = allocator.allocate(workspace_buffer_size);
-            workspace_buffer = dataPtr.get();
-          }
-          TORCH_CUDASPARSE_CHECK(cusparseSpMM(
-            cusparse_handle,
-            CUSPARSE_OPERATION_NON_TRANSPOSE,
-            CUSPARSE_OPERATION_TRANSPOSE,
-            (void*)&alpha_val,
-            sparse_descr,
-            dense_descr,
-            (void*)&beta_val,
-            result_descr,
-            cuda_data_type,
-            mm_alg,
-            workspace_buffer
-          ));
-          TORCH_CUDASPARSE_CHECK(cusparseDestroySpMat(sparse_descr));
-          TORCH_CUDASPARSE_CHECK(cusparseDestroyDnMat(dense_descr));
-          TORCH_CUDASPARSE_CHECK(cusparseDestroyDnMat(result_descr));
-          mat_el_begin_idx = mat_el_end_idx;
-        } else {
-          tmp_result[cur_mat_num].zero_();
+        // Create variables to view just the current set of matrices
+        int64_t sparse_nnz = mat_el_end_idx - mat_el_begin_idx;
+        cudaDataType cuda_data_type = getTensorCudaDataType(mat2_contig);
+        auto* row_indices_ptr = maybe_provide_aligned_buffer_idx_ptr(
+          row_indices_start_ptr,
+          mat_el_begin_idx,
+          sparse_nnz
+        );
+        auto* col_indices_ptr = &col_indices_start_ptr[mat_el_begin_idx];
+        scalar_t* values_ptr = &values_start_ptr[mat_el_begin_idx];
+
+        cusparseSpMatDescr_t sparse_descr;
+        TORCH_CUDASPARSE_CHECK(cusparseCreateCoo(
+          &sparse_descr,
+          dim_i,
+          dim_j,
+          sparse_nnz,
+          row_indices_ptr,
+          col_indices_ptr,
+          values_ptr,
+          CUSPARSE_INDEX_64I,
+          CUSPARSE_INDEX_BASE_ZERO,
+          cuda_data_type
+        ));
+        scalar_t* mat2_ptr = &mat2_start_ptr[dim_k*dim_j*cur_mat_num];
+        cusparseDnMatDescr_t dense_descr;
+        TORCH_CUDASPARSE_CHECK(cusparseCreateDnMat(
+          &dense_descr,
+          dim_j,
+          dim_k,
+          dim_k,
+          mat2_ptr,
+          cuda_data_type,
+          CUSPARSE_ORDER_ROW
+        ));
+        scalar_t* result_ptr = &result_start_ptr[dim_i*dim_k*cur_mat_num];
+        cusparseDnMatDescr_t result_descr;
+        TORCH_CUDASPARSE_CHECK(cusparseCreateDnMat(
+          &result_descr,
+          dim_i,
+          dim_k,
+          dim_k,
+          result_ptr,
+          cuda_data_type,
+          CUSPARSE_ORDER_ROW
+        ));
+        size_t required_workspace_buffer_size = 0;
+        TORCH_CUDASPARSE_CHECK(cusparseSpMM_bufferSize(
+          cusparse_handle,
+          CUSPARSE_OPERATION_NON_TRANSPOSE,
+          CUSPARSE_OPERATION_NON_TRANSPOSE,
+          &alpha_val,
+          sparse_descr,
+          dense_descr,
+          &beta_val,
+          result_descr,
+          cuda_data_type,
+          mm_alg,
+          &required_workspace_buffer_size
+        ));
+        if (required_workspace_buffer_size > workspace_buffer_size) {
+          workspace_buffer_size = required_workspace_buffer_size;
+          dataPtr = allocator.allocate(workspace_buffer_size);
+          workspace_buffer = dataPtr.get();
         }
+        TORCH_CUDASPARSE_CHECK(cusparseSpMM(
+          cusparse_handle,
+          CUSPARSE_OPERATION_NON_TRANSPOSE,
+          CUSPARSE_OPERATION_NON_TRANSPOSE,
+          &alpha_val,
+          sparse_descr,
+          dense_descr,
+          &beta_val,
+          result_descr,
+          cuda_data_type,
+          mm_alg,
+          workspace_buffer
+        ));
+        TORCH_CUDASPARSE_CHECK(cusparseDestroySpMat(sparse_descr));
+        TORCH_CUDASPARSE_CHECK(cusparseDestroyDnMat(dense_descr));
+        TORCH_CUDASPARSE_CHECK(cusparseDestroyDnMat(result_descr));
+        mat_el_begin_idx = mat_el_end_idx;
       }
     }
   );
   if (need_copy_result) {
     result.copy_(tmp_result);
   }
-  // Need to transpose the result matrices since cusparse stores
-  // them in column-major order in memory
-  result.transpose_(1,2);
 
   return result;
 }
diff --git a/aten/src/ATen/native/sparse/cuda/SparseSemiStructuredTile.cu b/aten/src/ATen/native/sparse/cuda/SparseSemiStructuredTile.cu
index 599a59a18312c..cfe91ecc2e5e8 100644
--- a/aten/src/ATen/native/sparse/cuda/SparseSemiStructuredTile.cu
+++ b/aten/src/ATen/native/sparse/cuda/SparseSemiStructuredTile.cu
@@ -59,7 +59,7 @@ struct MetadataCuSparseLt {
     // TODO: Cast metadata to Short
     static_assert(kBytesPerScalar == 2, "or modify the last dim below");
     metadata = metadata.view({rows / 128, cols / 32, 256});
-    return std::make_tuple(storage, packed, metadata);
+    return std::make_tuple(std::move(storage), std::move(packed), std::move(metadata));
   }
 
   MetadataCuSparseLt(at::Tensor metaN, at::Tensor metaT, int rows, int cols) {
diff --git a/aten/src/ATen/native/sparse/cuda/cuSPARSELtOps.cpp b/aten/src/ATen/native/sparse/cuda/cuSPARSELtOps.cpp
index 9d735ac0f2c88..f0f4172ccf6f4 100644
--- a/aten/src/ATen/native/sparse/cuda/cuSPARSELtOps.cpp
+++ b/aten/src/ATen/native/sparse/cuda/cuSPARSELtOps.cpp
@@ -2,6 +2,9 @@
 #include <unordered_map>
 #include <mutex>
 #include <string_view>
+#include <unordered_map>
+#include <vector>
+#include <c10/util/StringUtil.h>
 #if AT_CUSPARSELT_ENABLED()
 
 namespace at::native {
@@ -22,37 +25,51 @@ thread_local bool handle_initialized = false;
 c10::once_flag g_hipSparseLtSupportInitFlag;
 static bool g_hipSparseLtSupported = false;
 
+static const std::vector<std::string>& hipSparseLtSupportedArchs() {
+#if ROCM_VERSION >= 71400
+  static const std::vector<std::string> archs = {"gfx950", "gfx942", "gfx1250"};
+#elif ROCM_VERSION >= 71200
+  static const std::vector<std::string> archs = {"gfx950", "gfx942"};
+#else
+  static const std::vector<std::string> archs = {};
+#endif
+  return archs;
+}
+
 // Initialize the hipSparseLt support status once for the platform
 static void initHipSparseLtSupport() {
-    // Default to not supported
-    g_hipSparseLtSupported = false;
-
-    // Check only the first available device
-    try {
-        if (at::cuda::device_count() > 0) {
-            g_hipSparseLtSupported = at::detail::getCUDAHooks().isGPUArch({"gfx950", "gfx942"}, 0);
-        }
-    } catch (const std::exception&) {
-        // If an exception occurs during device property check, we assume hipSparseLt is not supported
-        // This could happen due to driver issues, device access problems, or other runtime errors
-        g_hipSparseLtSupported = false;
-        TORCH_WARN("Exception occurred while checking hipSparseLt support. Assuming not supported.");
+  // Default to not supported
+  g_hipSparseLtSupported = false;
+
+  // Check only the first available device
+  try {
+    if (at::cuda::device_count() > 0) {
+      g_hipSparseLtSupported = at::detail::getCUDAHooks().isGPUArch(
+          hipSparseLtSupportedArchs(), 0);
     }
+  } catch (const std::exception&) {
+    // If an exception occurs during device property check, we assume
+    // hipSparseLt is not supported This could happen due to driver issues,
+    // device access problems, or other runtime errors
+     g_hipSparseLtSupported = false;
+    TORCH_WARN(
+        "Exception occurred while checking hipSparseLt support. Assuming not supported.");
+  }
 }
 
 static bool isHipSparseLtSupported() {
     // Initialize support check only once
     c10::call_once(g_hipSparseLtSupportInitFlag, initHipSparseLtSupport);
 
-    // Return cached result (platform-wide)
-    if (!g_hipSparseLtSupported) {
-        TORCH_CHECK(
-            false,
-            "hipSparseLt not supported on this device, supported architectures: "
-            "gfx950, gfx942. "
-            "required ROCM version: 6.4.0 or later.");
-    }
-    return g_hipSparseLtSupported;
+  // Return cached result (platform-wide)
+  if (!g_hipSparseLtSupported) {
+    TORCH_CHECK(
+        false,
+        "hipSparseLt not supported on this device. Supported architectures: ",
+        c10::Join(", ", hipSparseLtSupportedArchs()),
+        ". hipSparseLt on ROCm requires ROCm 7.12 or newer.");
+  }
+  return g_hipSparseLtSupported;
 }
 #endif
 
@@ -84,7 +101,7 @@ at::Tensor _cslt_compress(const Tensor& sparse_input) {
       type = CUDA_R_32F;
       break;
 #endif
-#if defined(CUSPARSELT_VERSION) && CUSPARSELT_VERSION >= 602 && !defined(USE_ROCM)
+#if defined(CUSPARSELT_VERSION) && CUSPARSELT_VERSION >= 602 || defined(USE_ROCM)
     case at::ScalarType::Float8_e4m3fn:
       type = CUDA_R_8F_E4M3;
       break;
@@ -203,12 +220,18 @@ std::tuple<at::Tensor, int64_t, int64_t, int64_t, int64_t> _cslt_sparse_mm_impl(
       compute_type = CUSPARSE_COMPUTE_32F;
       break;
 #endif
-// if cuSPARSELt >= 6.2.3, we can add Float8 support
-#if defined(CUSPARSELT_VERSION) && CUSPARSELT_VERSION >= 602 && !defined(USE_ROCM)
+// cuSPARSELt >= 0.6.2 or hipSparseLt: add Float8 support
+#if defined(CUSPARSELT_VERSION) && CUSPARSELT_VERSION >= 602 || defined(USE_ROCM)
     case at::ScalarType::Float8_e4m3fn:
       input_type = CUDA_R_8F_E4M3;
+#ifdef USE_ROCM
+      // hipSparseLt 0.2.7: FP8 input only supports FP32 output
+      output_type = CUDA_R_32F;
+      C_type = CUDA_R_32F;
+#else
       output_type = CUDA_R_8F_E4M3;
       C_type = CUDA_R_16F;
+#endif
       compute_type = CUSPARSE_COMPUTE_32F;
       break;
 #endif
@@ -265,10 +288,11 @@ std::tuple<at::Tensor, int64_t, int64_t, int64_t, int64_t> _cslt_sparse_mm_impl(
           break;
       }
     }
-// cslt 0.6.2+: fp8 fp8 -> {fp8, fp16, bf16, fp32} support
-#if defined(CUSPARSELT_VERSION) && CUSPARSELT_VERSION >= 602 && !defined(USE_ROCM)
+// cslt 0.6.2+ or hipSparseLt: fp8 output dtype support
+#if defined(CUSPARSELT_VERSION) && CUSPARSELT_VERSION >= 602 || defined(USE_ROCM)
     else if (input_type == CUDA_R_8F_E4M3) {
       switch (out_dtype) {
+#ifndef USE_ROCM
         case at::ScalarType::Float8_e4m3fn:
           output_type = CUDA_R_8F_E4M3;
           C_type = CUDA_R_16F;
@@ -281,6 +305,7 @@ std::tuple<at::Tensor, int64_t, int64_t, int64_t, int64_t> _cslt_sparse_mm_impl(
           output_type = CUDA_R_16BF;
           C_type = CUDA_R_16BF;
           break;
+#endif
         case at::ScalarType::Float:
           output_type = CUDA_R_32F;
           C_type = CUDA_R_32F;
@@ -288,7 +313,11 @@ std::tuple<at::Tensor, int64_t, int64_t, int64_t, int64_t> _cslt_sparse_mm_impl(
         default:
           TORCH_CHECK(
               false,
-              "Unsupported out_dtype passed, must be one of {fp16, bf16, float32} for fp8 inputs");
+#ifdef USE_ROCM
+              "Unsupported out_dtype passed, must be float32 for fp8 inputs on ROCm");
+#else
+              "Unsupported out_dtype passed, must be one of {fp8, fp16, bf16, float32} for fp8 inputs");
+#endif
           break;
       }
     }
@@ -468,6 +497,8 @@ std::tuple<at::Tensor, int64_t, int64_t, int64_t, int64_t> _cslt_sparse_mm_impl(
         &alg_id,
         sizeof(alg_id)));
 
+#ifndef USE_ROCM
+    // hipSPARSELt does not support querying SPLIT_K attributes
     TORCH_CUDASPARSE_CHECK(cusparseLtMatmulAlgGetAttribute(
         &handle,
         &alg_sel,
@@ -481,6 +512,7 @@ std::tuple<at::Tensor, int64_t, int64_t, int64_t, int64_t> _cslt_sparse_mm_impl(
         CUSPARSELT_MATMUL_SPLIT_K_MODE,
         &splitKMode,
         sizeof(splitKMode)));
+#endif
 
     TORCH_CUDASPARSE_CHECK(cusparseLtMatmulAlgGetAttribute(
         &handle,
diff --git a/aten/src/ATen/native/sparse/cuda/cuSPARSELtOps.h b/aten/src/ATen/native/sparse/cuda/cuSPARSELtOps.h
index 20b6c663f99f8..f06bda807c7cd 100644
--- a/aten/src/ATen/native/sparse/cuda/cuSPARSELtOps.h
+++ b/aten/src/ATen/native/sparse/cuda/cuSPARSELtOps.h
@@ -14,6 +14,16 @@
 #include <cstdint>
 
 #if AT_CUSPARSELT_ENABLED()
+// ROCm 7.0.2's amd_hip_bf16.h (pulled in via hipsparselt.h -> hip_fp8.h ->
+// amd_hip_fp8.h -> amd_hip_bf16.h) contains device-only builtins (warpSize,
+// __shfl_*_sync) that fail during host compilation. Pre-define the hip_fp8.h
+// include guard to prevent this chain. The hipsparselt C API uses opaque
+// handles and enum types, not C++ fp8 types directly.
+#if defined(USE_ROCM) && !defined(__HIP_DEVICE_COMPILE__)
+#ifndef HIP_INCLUDE_HIP_HIP_FP8_H
+#define HIP_INCLUDE_HIP_HIP_FP8_H
+#endif
+#endif
 #include <cusparseLt.h>
 #endif
 
diff --git a/aten/src/ATen/native/sparse/mps/SparseMPSTensorMath.mm b/aten/src/ATen/native/sparse/mps/SparseMPSTensorMath.mm
index d9aadcf51fb94..31037b888eba2 100644
--- a/aten/src/ATen/native/sparse/mps/SparseMPSTensorMath.mm
+++ b/aten/src/ATen/native/sparse/mps/SparseMPSTensorMath.mm
@@ -1103,7 +1103,9 @@ Tensor sparse_sparse_matmul_mps(const Tensor& mat1_, const Tensor& mat2_) {
   auto v_out = vA_out.mul(vB_out);
 
   // build (2, P) indices
-  auto out_indices = at::empty({2, P}, at::device(device).dtype(at::kLong)).contiguous();
+  auto out_indices = at::empty(
+      {2, P},
+      at::device(device).dtype(at::kLong).memory_format(c10::MemoryFormat::Contiguous));
   out_indices.select(0, 0).copy_(i_out);
   out_indices.select(0, 1).copy_(j_out);
 
@@ -1150,6 +1152,8 @@ Tensor _sparse_sum_backward_mps(const Tensor& grad_, const SparseTensor& input_,
   int64_t sparse_dims_to_sum_size = 0;
   std::vector<int64_t> sparse_dims_to_keep_v;
   std::vector<int64_t> dense_dims_to_sum_v;
+  sparse_dims_to_keep_v.reserve(input_sparse_dim);
+  dense_dims_to_sum_v.reserve(input_dense_dim);
 
   for (auto d = 0; d < input_dim; d++) {
     if (dims_to_sum_b[static_cast<size_t>(d)]) {
diff --git a/aten/src/ATen/native/transformers/attention.cpp b/aten/src/ATen/native/transformers/attention.cpp
index a83793a1640f4..eecfe1058c2c4 100644
--- a/aten/src/ATen/native/transformers/attention.cpp
+++ b/aten/src/ATen/native/transformers/attention.cpp
@@ -264,7 +264,7 @@ std::tuple<Tensor, Tensor, Tensor> transform_bias_rescale_qkv_cpu(
   auto q_k_v_s =
       at::native::split(q_k_v.view({3 * B, num_head, T, dim_per_head}), B, 0);
   TORCH_INTERNAL_ASSERT_DEBUG_ONLY(q_k_v_s.size() == 3);
-  return std::make_tuple(q_k_v_s[0], q_k_v_s[1], q_k_v_s[2]);
+  return std::make_tuple(std::move(q_k_v_s[0]), std::move(q_k_v_s[1]), std::move(q_k_v_s[2]));
 }
 
 std::tuple<Tensor, Tensor> native_multi_head_attention_cpu(
@@ -443,12 +443,7 @@ int64_t _fused_sdp_choice_cpp(const Tensor& query_, const Tensor& key, const Ten
   return static_cast<int64_t>(backend);
 }
 
-REGISTER_ARCH_DISPATCH(_fused_sdp_choice_stub, DEFAULT, &_fused_sdp_choice_cpp)
-REGISTER_AVX2_DISPATCH(_fused_sdp_choice_stub, &_fused_sdp_choice_cpp)
-REGISTER_AVX512_DISPATCH(_fused_sdp_choice_stub, &_fused_sdp_choice_cpp)
-REGISTER_VSX_DISPATCH(_fused_sdp_choice_stub, &_fused_sdp_choice_cpp)
-REGISTER_ZVECTOR_DISPATCH(_fused_sdp_choice_stub, &_fused_sdp_choice_cpp)
-REGISTER_SVE256_DISPATCH(_fused_sdp_choice_stub, &_fused_sdp_choice_cpp)
+REGISTER_ALL_CPU_DISPATCH(_fused_sdp_choice_stub, &_fused_sdp_choice_cpp)
 REGISTER_HPU_DISPATCH(_fused_sdp_choice_stub, &_fused_sdp_choice_meta)
 
 int64_t _fused_sdp_choice_meta(
@@ -497,6 +492,20 @@ int64_t _fused_sdp_choice_meta(
     return choice_int;
   }
 #endif
+  bool has_xpu = query_key_set.has(c10::DispatchKey::XPU);
+  if (has_xpu) {
+    auto choice_int = _fused_sdp_choice_stub(
+        at::kXPU,
+        query_,
+        key,
+        value,
+        attn_mask_,
+        dropout_p,
+        is_causal,
+        scale,
+        enable_gqa);
+    return choice_int;
+  }
   return static_cast<int64_t>(sdp::SDPBackend::math);
 }
 namespace {
@@ -538,7 +547,7 @@ inline void validate_sdpa_input(
 // the math and memory efficient attn_mask implementation
 //  Args:
 //    attn_mask: attn_mask of shape (B, L, S) or (L, S) or (B, N_heads, L, S)
-std::optional<Tensor> convert_boolean_attn_mask_(const std::optional<Tensor>& attn_mask, caffe2::TypeMeta dtype, double neg_inf) {
+std::optional<Tensor> convert_boolean_attn_mask(const std::optional<Tensor>& attn_mask, caffe2::TypeMeta dtype) {
   // Pass through
   if (!attn_mask.has_value()) {
     return std::nullopt;
@@ -546,23 +555,13 @@ std::optional<Tensor> convert_boolean_attn_mask_(const std::optional<Tensor>& at
   // Convert boolean mask to additive mask; need to invert mask to indicate what
   // to mask *out*.
   if (attn_mask->dtype() == at::kBool) {
+    constexpr double neg_inf = -std::numeric_limits<double>::infinity();
     return at::where(*attn_mask, 0.0, at::scalar_tensor(neg_inf, at::TensorOptions().dtype(dtype).device(attn_mask->device())));
   }
   // Otherwise, attn_mask represents an additive attention tensor
   return attn_mask;
 }
 
-std::optional<Tensor> convert_boolean_attn_mask(const std::optional<Tensor>& attn_mask, caffe2::TypeMeta dtype) {
-  return convert_boolean_attn_mask_(attn_mask, dtype, -std::numeric_limits<double>::infinity());
-}
-
-// alternate version to workaround -inf issue with cuDNN
-// TODO(eqy): delete this when cuDNN -inf issue is resolved
-std::optional<Tensor> convert_boolean_attn_mask_cudnn(const std::optional<Tensor>& attn_mask, caffe2::TypeMeta dtype) {
-  // TODO Use the max type of the input and output
-  return convert_boolean_attn_mask_(attn_mask, dtype, -65504.0);
-}
-
 // Memory Efficient Attention requires a padded attn mask bias
 // This function pads the attn_mask bias to be a multiple of 16
 // Then slices the padded bias to the original size
@@ -691,11 +690,11 @@ Tensor _safe_softmax(
 // greater than 0.0 is specified.
 //
 // Args:
-//     query (Tensor): Query tensor; shape (N, ..., L, E)
-//     key (Tensor): Key tensor; shape (N, ..., S, E)
-//     value (Tensor): Value tensor; shape (N, ..., S, E)
+//     query (Tensor): Query tensor; shape (N, ..., Hq, L, E)
+//     key (Tensor): Key tensor; shape (N, ..., H, S, E)
+//     value (Tensor): Value tensor; shape (N, ..., H, S, Ev)
 //     attn_mask (optional Tensor): Attention mask; shape must be broadcastable to the shape of attention weights,
-//         which is (N,..., L, S). Two types of masks are supported.
+//         which is (N,..., Hq, L, S). Two types of masks are supported.
 //         A boolean mask where a value of True indicates that the element *should* take part in attention.
 //         A float mask of the same type as query, key, value that is added to the attention score.
 //     dropout_p (float): Dropout probability; if greater than 0.0, dropout is applied
@@ -707,14 +706,17 @@ Tensor _safe_softmax(
 //         sparse masks) via tensor subclassing, allowing for a leaner API.
 //
 // Returns a tensor:
-//     output (Tensor): Attention output; shape (N, ..., L, E)
+//     output (Tensor): Attention output; shape (N, ..., Hq, L, Ev)
 //
 // Shape legend:
 //     N: Batch size
 //     ...: Any number of other batch dimensions (optional)
 //     S: Source sequence length
 //     L: Target sequence length
-//     E: Embedding dimension
+//     E: Embedding dimension of the query and key
+//     Ev: Embedding dimension of the value
+//     Hq: Number of heads of query
+//     H: Number of heads of key and value
 Tensor scaled_dot_product_attention(
     const Tensor& query_,
     const Tensor& key,
@@ -726,6 +728,22 @@ Tensor scaled_dot_product_attention(
     bool enable_gqa) {
   using sdp::SDPBackend;
   validate_sdpa_input(query_, key, value, attn_mask_, dropout_p, is_causal, scale);
+  // NB: This op is CompositeImplicitAutograd — autograd traces through the
+  // implementation rather than using an explicit backward formula. We must
+  // return early here because the scale computation (1/sqrt(head_dim)) is
+  // undefined when head_dim is 0, and backends don't uniformly handle
+  // zero-element tensors. We use the sum()*0 trick (same pattern as
+  // _batch_norm_impl_index in Normalization.cpp) to make the output depend
+  // on all inputs so that backward() produces correctly-shaped zero
+  // gradients instead of None.
+  if (TORCH_GUARD_OR_FALSE(query_.sym_numel().sym_eq(0)) ||
+      TORCH_GUARD_OR_FALSE(key.sym_numel().sym_eq(0)) ||
+      TORCH_GUARD_OR_FALSE(value.sym_numel().sym_eq(0))) {
+    auto output_shape = query_.sym_sizes().vec();
+    output_shape[output_shape.size() - 1] = value.sym_size(-1);
+    auto out = at::zeros_symint(output_shape, query_.options());
+    return out + (query_.sum() + key.sum() + value.sum()) * 0;
+  }
   int64_t choice_int = static_cast<int64_t>(sdp::SDPBackend::math);
   if (_fused_sdp_choice_stub.is_device_supported(query_.device().type())) {
     choice_int = _fused_sdp_choice_stub(query_.device().type(),
@@ -733,8 +751,7 @@ Tensor scaled_dot_product_attention(
   }
   const auto query_device_type = query_.device().type();
   const auto backend = static_cast<SDPBackend>(choice_int);
-  const auto convert_attn_func = backend != SDPBackend::cudnn_attention ? convert_boolean_attn_mask : convert_boolean_attn_mask_cudnn;
-  auto attn_mask = convert_attn_func(attn_mask_, query_.dtype());
+  auto attn_mask = convert_boolean_attn_mask(attn_mask_, query_.dtype());
   switch (backend) {
     case SDPBackend::cudnn_attention: {
       bool compute_logsumexp = should_compute_logsumexp(query_, key, value);
@@ -775,48 +792,8 @@ Tensor scaled_dot_product_attention(
       return std::get<0>(out_lse_softmax);
     }
     case SDPBackend::math: {
-#ifdef USE_MPS
-      TORCH_CHECK_NOT_IMPLEMENTED(
-        c10::isFloatingType(query_.scalar_type()),
-        "scaled_dot_product_attention for MPS does not support dtype ",
-        query_.scalar_type());
-      TORCH_CHECK_NOT_IMPLEMENTED(
-        c10::isFloatingType(key.scalar_type()),
-        "scaled_dot_product_attention for MPS does not support dtype ",
-        key.scalar_type());
-      TORCH_CHECK_NOT_IMPLEMENTED(
-        c10::isFloatingType(value.scalar_type()),
-        "scaled_dot_product_attention for MPS does not support dtype ",
-        value.scalar_type());
-      const auto any_nested = query_.is_nested() || key.is_nested() || value.is_nested();
       const bool any_inputs_require_grad = query_.requires_grad() || key.requires_grad() || value.requires_grad();
-      const auto all_contiguous = query_.is_contiguous_or_false() && key.is_contiguous_or_false() && value.is_contiguous_or_false();
-      if (query_device_type == DeviceType::MPS && dropout_p == 0.0
-          && !(GradMode::is_enabled() && any_inputs_require_grad)
-          && (all_contiguous || mps::is_macos_13_or_newer(mps::MacOSVersion::MACOS_VER_15_0_PLUS))
-          && !any_nested) {
-        if (enable_gqa) {
-          int64_t q_heads = query_.size(-3);
-          int64_t k_heads = key.size(-3);
-          int64_t repeat_factor = q_heads / k_heads;
-
-          if (repeat_factor > 1) {
-            TORCH_CHECK(q_heads % k_heads == 0,
-                          "For GQA, the query tensor's head dimension (" + std::to_string(q_heads) +
-                                    ") must be divisible by the key tensor's head dimension (" + std::to_string(k_heads) + ").");
-            auto repeated_key = key.repeat_interleave(repeat_factor, /*dim=*/-3);
-            auto repeated_value = value.repeat_interleave(repeat_factor, /*dim=*/-3);
-            return std::get<0>(at::_scaled_dot_product_attention_math_for_mps(
-              query_,
-              repeated_key,
-              repeated_value,
-              attn_mask,
-              dropout_p,
-              is_causal,
-              std::nullopt, /*dropout_mask*/
-              scale));
-          }
-        }
+      if (query_device_type == c10::kMPS && !(at::GradMode::is_enabled() && any_inputs_require_grad)) {
         return std::get<0>(at::_scaled_dot_product_attention_math_for_mps(
             query_,
             key,
@@ -825,9 +802,9 @@ Tensor scaled_dot_product_attention(
             dropout_p,
             is_causal,
             std::nullopt, /*dropout_mask*/
-            scale));
+            scale,
+            enable_gqa));
       }
-#endif
       return std::get<0>(at::_scaled_dot_product_attention_math(
           query_,
           key,
@@ -843,7 +820,6 @@ Tensor scaled_dot_product_attention(
       TORCH_CHECK(
           false,
           "No viable backend for scaled_dot_product_attention was found.");
-      return Tensor();
   }
 }
 
@@ -1053,11 +1029,6 @@ _scaled_dot_product_fused_attention_overrideable_backward(
     const at::Tensor & philox_offset,
     std::optional<double> scale) {
   TORCH_CHECK_NOT_IMPLEMENTED(false, "_scaled_dot_product_fused_attention_overrideable_backward not implemented: This is an operator for privateuse1 backends, please use TORCH_LIBRARY_IMPL to override this function ");
-  return std::tuple<Tensor, Tensor, Tensor, Tensor>(
-          at::empty_like(query),
-          at::empty_like(key),
-          at::empty_like(value),
-          at::empty_like(attn_bias));
 }
 
 Tensor triton_multi_head_attention(
diff --git a/aten/src/ATen/native/transformers/cuda/attention.cu b/aten/src/ATen/native/transformers/cuda/attention.cu
index f2df246521fb5..93466a28064f4 100644
--- a/aten/src/ATen/native/transformers/cuda/attention.cu
+++ b/aten/src/ATen/native/transformers/cuda/attention.cu
@@ -86,6 +86,7 @@
 #include <ATen/native/transformers/cuda/mem_eff_attention/pytorch_utils.h>
 #else
 // MemoryEfficient Attention Specific Imports for ROCM
+#include <ATen/native/transformers/hip/gemm_kernel_utils.h>
 #ifndef DISABLE_AOTRITON
 #include <ATen/native/transformers/hip/aotriton_adapter.h>
 #include <aotriton/flash.h>
@@ -95,7 +96,7 @@
 #endif
 #endif
 
-#if defined(USE_ROCM) && (defined(USE_FLASH_ATTENTION) || defined(USE_MEM_EFF_ATTENTION))
+#if defined(USE_ROCM) && defined(USE_FLASH_ATTENTION)
 namespace pytorch_flash
 {
 std::tuple<
@@ -466,6 +467,138 @@ Tensor collapse_dims_1_and_2(const Tensor& sizes) {
   return (sizes_dim1 * sizes_dim2).contiguous();
 }
 
+std::tuple<Tensor, Tensor, Tensor, Tensor, Tensor>
+_flash_attention_forward_impl(
+    const Tensor& query,
+    const Tensor& key,
+    const Tensor& value,
+    const std::optional<Tensor>& cumulative_sequence_length_q,
+    const std::optional<Tensor>& cumulative_sequence_length_k,
+    int64_t max_seqlen_batch_q,
+    int64_t max_seqlen_batch_k,
+    double dropout_p,
+    bool is_causal,
+    bool return_debug_mask,
+    std::optional<double> scale,
+    std::optional<int64_t> window_size_left,
+    std::optional<int64_t> window_size_right,
+    const std::optional<Tensor>& _seqused_k,
+    const std::optional<Tensor>& _alibi_slopes,
+    const std::optional<Tensor>& _block_table,
+    std::optional<Tensor> out,
+    std::optional<int64_t> num_splits
+    ) {
+#if defined(USE_FLASH_ATTENTION)
+  TORCH_CHECK(
+      !num_splits.has_value(),
+      "num_splits requires FA3. Register FA3 with `register_flash_attention_fa3()` to set num_splits.");
+  const auto softmax_scale =
+      sdp::calculate_scale(query, scale).expect_float();
+
+  std::optional<Tensor> seqused_k = _seqused_k;
+  std::optional<at::Tensor> block_table = _block_table;
+  std::optional<Tensor> alibi_slopes = _alibi_slopes;
+  const float softcap = 0.0;
+
+#ifdef USE_ROCM  // ROCM backend accepts std::optional for window_size_left/right directly.
+#ifdef DISABLE_AOTRITON  // CK backend, Passing window_size as it is
+  const auto window_left = window_size_left;
+  const auto window_right = window_size_right;
+#else  // AOTriton implements "generalized" SWA and negative size means negative shifting.
+  // aotriton_adapter::parse_window_size tries to match the behavior of CUTLASS backend
+  using sdp::aotriton_adapter::parse_window_size;
+  const auto [window_left, window_right] = parse_window_size(window_size_left,
+                                                             window_size_right);
+#endif
+#else  // USE_ROCM
+  const int window_left = window_size_left.value_or(-1);
+  const int window_right = window_size_right.value_or(-1);
+#endif  // USE_ROCM
+
+  // We are going to have two paths:
+  // 1. The standard MHA path for dense tensors
+  // 2. The Varseqlen path
+  TORCH_CHECK(
+      cumulative_sequence_length_q.has_value() ==
+          cumulative_sequence_length_k.has_value(),
+      "cumulative_sequence_length_q and cumulative_sequence_length_k must be both set or both not set");
+  Tensor output, q_padded, k_padded, v_padded, logsumexp, output_shape,
+      philox_seed, philox_offset, debug_attn_mask;
+  if (cumulative_sequence_length_q.has_value()) {
+    std::tie(
+        output,
+        q_padded,
+        k_padded,
+        v_padded,
+        logsumexp,
+        philox_seed,
+        philox_offset,
+        debug_attn_mask) =
+        FLASH_NAMESPACE::mha_varlen_fwd(
+            query,
+            key,
+            value,
+            out,
+            cumulative_sequence_length_q.value(),
+            cumulative_sequence_length_k.value(),
+            seqused_k, /*seqused_k*/
+            block_table, /*block_table*/
+            alibi_slopes, /*alibi_slopes*/
+            max_seqlen_batch_q,
+            max_seqlen_batch_k,
+            dropout_p,
+            softmax_scale,
+            false /*zero_tensors*/,
+            is_causal,
+            window_left,
+            window_right,
+            softcap,
+            return_debug_mask,
+            std::nullopt /*gen_*/);
+  } else {
+    std::tie(
+        output,
+        q_padded,
+        k_padded,
+        v_padded,
+        logsumexp,
+        philox_seed,
+        philox_offset,
+        debug_attn_mask) =
+        FLASH_NAMESPACE::mha_fwd(
+            query,
+            key,
+            value,
+            out,
+            alibi_slopes,
+            dropout_p,
+            softmax_scale,
+            is_causal,
+            window_left,
+            window_right,
+            softcap,
+            return_debug_mask, /*return_softmax (this is used for testing)*/
+            std::nullopt);
+  }
+  debug_attn_mask =
+      return_debug_mask ? debug_attn_mask : at::empty({0}, query.options());
+  return std::make_tuple(
+      std::move(output),
+      std::move(logsumexp),
+      std::move(philox_seed),
+      std::move(philox_offset),
+      std::move(debug_attn_mask));
+
+#endif
+  TORCH_CHECK(false, "USE_FLASH_ATTENTION was not enabled for build.")
+  return std::make_tuple(
+      Tensor(),
+      Tensor(),
+      Tensor(),
+      Tensor(),
+      Tensor());
+}
+
 } // namespace
 // compute q = (q + q_bias) / sqrt(dim_per_head), k = k + k_bias, v = v + v_bias
 __host__ std::tuple<Tensor, Tensor, Tensor> transform_bias_rescale_qkv_cuda(
@@ -568,7 +701,7 @@ __host__ std::tuple<Tensor, Tensor, Tensor> transform_bias_rescale_qkv_cuda(
 #undef CALL_KERNEL
   auto q_k_v_s =
       at::native::split(q_k_v.view({3 * B, num_head, T, dim_per_head}), B, 0);
-  return std::make_tuple(q_k_v_s[0], q_k_v_s[1], q_k_v_s[2]);
+  return std::make_tuple(std::move(q_k_v_s[0]), std::move(q_k_v_s[1]), std::move(q_k_v_s[2]));
 }
 
 std::tuple<Tensor, Tensor> native_multi_head_attention_cuda(
@@ -646,7 +779,7 @@ std::tuple<Tensor, Tensor> native_multi_head_attention_cuda(
 
 #endif
   const auto dim_per_head = D / num_head;
-  if ((query.is_same(key) && key.is_same(value)) && dim_per_head % 8 == 0 && !need_weights) {
+  if ((query.is_same(key) && key.is_same(value)) && !need_weights) {
 
     // We have not done linear projection yet but the input for SDP
     // Is expected to be 4 dimensional. We "cheaply" create view tensors
@@ -824,7 +957,7 @@ std::tuple<Tensor, Tensor, Tensor, Tensor, c10::SymInt, c10::SymInt, Tensor, Ten
   // Reshape output to convert nnz to batch_size and seq_len
   Tensor attention = output.transpose(1,2);
 
-  return std::make_tuple(attention, logsumexp, Tensor(), Tensor(), max_seqlen_batch_q, max_seqlen_batch_k, philox_seed, philox_offset, debug_attn_mask);
+  return std::make_tuple(std::move(attention), std::move(logsumexp), Tensor(), Tensor(), max_seqlen_batch_q, max_seqlen_batch_k, std::move(philox_seed), std::move(philox_offset), std::move(debug_attn_mask));
 }
 
 std::tuple<Tensor, Tensor, Tensor, Tensor, c10::SymInt, c10::SymInt, Tensor, Tensor, Tensor> _scaled_dot_product_flash_attention_cuda_quantized(
@@ -1174,115 +1307,52 @@ _flash_attention_forward(
     std::optional<int64_t> window_size_left,
     std::optional<int64_t> window_size_right,
     const std::optional<Tensor>& _seqused_k,
-    const std::optional<Tensor>& _alibi_slopes
+    const std::optional<Tensor>& _alibi_slopes,
+    const std::optional<Tensor>& _block_table,
+    std::optional<int64_t> num_splits
     ) {
-#if defined(USE_FLASH_ATTENTION)
-  const auto softmax_scale =
-      sdp::calculate_scale(query, scale).expect_float();
-  std::optional<Tensor> out = std::nullopt;
-
-  std::optional<Tensor> seqused_k = _seqused_k;
-  std::optional<at::Tensor> block_table = std::nullopt;  // we are not using the block table yet
-  std::optional<Tensor> alibi_slopes = _alibi_slopes;
-  const float softcap = 0.0;
-
-#ifdef USE_ROCM  // ROCM backend accepts std::optional for window_size_left/right directly.
-#ifdef DISABLE_AOTRITON  // CK backend, Passing window_size as it is
-  const auto window_left = window_size_left;
-  const auto window_right = window_size_right;
-#else  // AOTriton implements "generalized" SWA and negative size means negative shifting.
-  // aotriton_adapter::parse_window_size tries to match the behavior of CUTLASS backend
-  using sdp::aotriton_adapter::parse_window_size;
-  const auto [window_left, window_right] = parse_window_size(window_size_left,
-                                                             window_size_right);
-#endif
-#else  // USE_ROCM
-  const int window_left = window_size_left.value_or(-1);
-  const int window_right = window_size_right.value_or(-1);
-#endif  // USE_ROCM
-
-  // We are going to have two paths:
-  // 1. The standard MHA path for dense tensors
-  // 2. The Varseqlen path
-  TORCH_CHECK(
-      cumulative_sequence_length_q.has_value() ==
-          cumulative_sequence_length_k.has_value(),
-      "cumulative_sequence_length_q and cumulative_sequence_length_k must be both set or both not set");
-  Tensor output, q_padded, k_padded, v_padded, logsumexp, output_shape,
-      philox_seed, philox_offset, debug_attn_mask;
-  if (cumulative_sequence_length_q.has_value()) {
-    std::tie(
-        output,
-        q_padded,
-        k_padded,
-        v_padded,
-        logsumexp,
-        philox_seed,
-        philox_offset,
-        debug_attn_mask) =
-        FLASH_NAMESPACE::mha_varlen_fwd(
-            query,
-            key,
-            value,
-            out,
-            cumulative_sequence_length_q.value(),
-            cumulative_sequence_length_k.value(),
-            seqused_k, /*seqused_k*/
-            block_table, /*block_table*/
-            alibi_slopes, /*alibi_slopes*/
-            max_seqlen_batch_q,
-            max_seqlen_batch_k,
-            dropout_p,
-            softmax_scale,
-            false /*zero_tensors*/,
-            is_causal,
-            window_left,
-            window_right,
-            softcap,
-            return_debug_mask,
-            std::nullopt /*gen_*/);
-  } else {
-    std::tie(
-        output,
-        q_padded,
-        k_padded,
-        v_padded,
-        logsumexp,
-        philox_seed,
-        philox_offset,
-        debug_attn_mask) =
-        FLASH_NAMESPACE::mha_fwd(
-            query,
-            key,
-            value,
-            out,
-            alibi_slopes,
-            dropout_p,
-            softmax_scale,
-            is_causal,
-            window_left,
-            window_right,
-            softcap,
-            return_debug_mask, /*return_softmax (this is used for testing)*/
-            std::nullopt);
-  }
-  debug_attn_mask =
-      return_debug_mask ? debug_attn_mask : at::empty({0}, query.options());
-  return std::make_tuple(
-      std::move(output),
-      std::move(logsumexp),
-      std::move(philox_seed),
-      std::move(philox_offset),
-      std::move(debug_attn_mask));
+  return _flash_attention_forward_impl(
+      query, key, value,
+      cumulative_sequence_length_q, cumulative_sequence_length_k,
+      max_seqlen_batch_q, max_seqlen_batch_k,
+      dropout_p, is_causal, return_debug_mask,
+      scale, window_size_left, window_size_right,
+      _seqused_k, _alibi_slopes, _block_table,
+      /*out=*/std::nullopt, num_splits);
+}
 
-#endif
-  TORCH_CHECK(false, "USE_FLASH_ATTENTION was not enabled for build.")
-  return std::make_tuple(
-      Tensor(),
-      Tensor(),
-      Tensor(),
-      Tensor(),
-      Tensor());
+Tensor
+_flash_attention_forward_no_dropout_inplace(
+    Tensor& out,
+    const Tensor& query,
+    const Tensor& key,
+    const Tensor& value,
+    const std::optional<Tensor>& cumulative_sequence_length_q,
+    const std::optional<Tensor>& cumulative_sequence_length_k,
+    int64_t max_seqlen_batch_q,
+    int64_t max_seqlen_batch_k,
+    double dropout_p,
+    bool is_causal,
+    bool return_debug_mask,
+    std::optional<double> scale,
+    std::optional<int64_t> window_size_left,
+    std::optional<int64_t> window_size_right,
+    const std::optional<Tensor>& _seqused_k,
+    const std::optional<Tensor>& _alibi_slopes,
+    const std::optional<Tensor>& _block_table,
+    std::optional<int64_t> num_splits
+    ) {
+  TORCH_CHECK(dropout_p == 0.0);
+  auto [output, logsumexp, philox_seed, philox_offset, debug_attn_mask] =
+      _flash_attention_forward_impl(
+          query, key, value,
+          cumulative_sequence_length_q, cumulative_sequence_length_k,
+          max_seqlen_batch_q, max_seqlen_batch_k,
+          dropout_p, is_causal, return_debug_mask,
+          scale, window_size_left, window_size_right,
+          _seqused_k, _alibi_slopes, _block_table,
+          /*out=*/std::make_optional(out), num_splits);
+  return logsumexp;
 }
 
 std::tuple<Tensor, Tensor, Tensor, Tensor, Tensor>
@@ -1493,6 +1563,10 @@ std::tuple<Tensor, Tensor, Tensor, Tensor, c10::SymInt, c10::SymInt> _efficient_
     // compute_logsumexp is false
     constexpr int kAlignLSE = 1;
     res = at::empty({B, M, num_heads, Kv}, query.options());
+    // TODO: Use Compact Varlen LSE
+    //       The current memory allocation is strictly larger than necessary
+    //       (total_q <= max_seqlen_q * B)
+    //       The problem is total_q is not available here.
     at::Tensor softmax_lse;
     logsumexp = at::empty(
       { B, num_heads, compute_logsumexp ? max_seqlen_q : 0},
@@ -1521,8 +1595,6 @@ std::tuple<Tensor, Tensor, Tensor, Tensor, c10::SymInt, c10::SymInt> _efficient_
       atomic_counter = at::zeros({1}, query.options().dtype(at::kInt));
     }
 
-    using aotriton::v2::flash::attn_fwd;
-    using aotriton::v2::flash::attn_fwd_compact_varlen;
     using sdp::aotriton_adapter::mk_aotensor;
     using sdp::aotriton_adapter::mk_aoscalartensor;
     using sdp::aotriton_adapter::mk_philoxtensor;
@@ -1538,92 +1610,47 @@ std::tuple<Tensor, Tensor, Tensor, Tensor, c10::SymInt, c10::SymInt> _efficient_
     auto offset_output = mk_philoxtensor(use_philox_state ? offset_t.data_ptr<int64_t>() : nullptr);
     auto persistent_counter = mk_atomictensor(is_causal ? atomic_counter.data_ptr<int32_t>() : nullptr);
     hipError_t err; // TODO: Error handling
-    if constexpr (AOTRITON_ALWAYS_V3_API) {  // Better readability than nesting ifdef
-#if AOTRITON_V3_API  // if constexpr does not stop errors from undefined functions
-      using aotriton::v3::flash::CausalType;
-      using aotriton::v3::flash::VarlenType;
-      using aotriton::v3::flash::WindowValue;
-      aotriton::v3::flash::attn_fwd_params params;
-      params.Q = mk_aotensor(q_t, "q");
-      params.K = mk_aotensor(k_t, "k");
-      params.V = mk_aotensor(v_t, "v");
-      params.Sm_scale = softmax_scale;
-      params.L = compute_logsumexp ? mk_aotensor<2>(softmax_lse, "M") : empty_t2;
-      params.Out = mk_aotensor(output_t, "Out");
-      params.Max_seqlen_q = max_seqlen_q;    // Unused if cu_seqlens_q is empty
-      params.Max_seqlen_k = max_seqlen_k;    // Unused if cu_seqlens_k is empty
-      params.dropout_p = dropout_p;
-      params.philox_seed_ptr = seed;
-      params.philox_offset1 = offset1;
-      params.philox_offset2 = offset2;
-      params.philox_seed_output = seed_output;
-      params.philox_offset_output = offset_output;
-      params.encoded_softmax = mk_aotensor(softmax_fa_t, "encoded_softmax");
-      params.persistent_atomic_counter = persistent_counter;
-      params.causal_type = is_causal ? CausalType::WindowedAttention : CausalType::None;
-      if (static_cast<int64_t>(sdp::CustomMaskType::CausalFromTopLeft) == custom_mask_type) {
-        params.window_left = WindowValue::TopLeftAligned;
-        params.window_right = WindowValue::TopLeftAligned;
-      } else if (static_cast<int64_t>(sdp::CustomMaskType::CausalFromBottomRight) == custom_mask_type) {
-        params.window_left = WindowValue::BottomRightAligned;
-        params.window_right = WindowValue::BottomRightAligned;
-      }
-      if (bias.has_value()) {
-        params.B = mk_aotensor(bias.value(), "bias");
-      }
-      if (seqstart_q.has_value()) {
-        params.varlen_type = VarlenType::CompactVarlen;
-        params.cu_seqlens_q = mk_aotensor<1>(seqstart_q.value(), "cu_seqlens_q");
-        params.cu_seqlens_k = mk_aotensor<1>(seqstart_k.value(), "cu_seqlens_k");
-      } else {
-        params.varlen_type = VarlenType::None;
-      }
-      err = aotriton::v3::flash::attn_fwd(params,
-                                          aotriton::v3::flash::attn_fwd_params::kVersion,
-                                          stream);
-#endif  // AOTRITON_V3_API
-    } else if (seqstart_q.has_value()) {
-      // varlen aka nested tensor
-      err = attn_fwd_compact_varlen(mk_aotensor(q_t, "q"),
-                                    mk_aotensor(k_t, "k"),
-                                    mk_aotensor(v_t, "v"),
-                                    bias.has_value() ? mk_aotensor(bias.value(), "bias"): empty_t4,
-                                    mk_aotensor<1>(seqstart_q.value(), "cu_seqlens_q"),
-                                    mk_aotensor<1>(seqstart_k.value(), "cu_seqlens_k"),
-                                    max_seqlen_q,
-                                    max_seqlen_k,
-                                    softmax_scale,
-                                    compute_logsumexp ? mk_aotensor<2>(softmax_lse, "M") : empty_t2,
-                                    mk_aotensor(output_t, "Out"),
-                                    dropout_p,
-                                    seed,
-                                    offset1,
-                                    offset2,
-                                    seed_output,
-                                    offset_output,
-                                    mk_aotensor(softmax_fa_t, "encoded_softmax"),
-                                    is_causal,
-                                    persistent_counter,
-                                    stream);
+    using aotriton::v3::flash::CausalType;
+    using aotriton::v3::flash::VarlenType;
+    using aotriton::v3::flash::WindowValue;
+    aotriton::v3::flash::attn_fwd_params params;
+    params.Q = mk_aotensor(q_t, "q");
+    params.K = mk_aotensor(k_t, "k");
+    params.V = mk_aotensor(v_t, "v");
+    params.Sm_scale = softmax_scale;
+    params.L = compute_logsumexp ? mk_aotensor<2>(softmax_lse, "M") : empty_t2;
+    params.Out = mk_aotensor(output_t, "Out");
+    params.Max_seqlen_q = max_seqlen_q;    // Unused if cu_seqlens_q is empty
+    params.Max_seqlen_k = max_seqlen_k;    // Unused if cu_seqlens_k is empty
+    params.dropout_p = dropout_p;
+    params.philox_seed_ptr = seed;
+    params.philox_offset1 = offset1;
+    params.philox_offset2 = offset2;
+    params.philox_seed_output = seed_output;
+    params.philox_offset_output = offset_output;
+    params.encoded_softmax = mk_aotensor(softmax_fa_t, "encoded_softmax");
+    params.persistent_atomic_counter = persistent_counter;
+    params.causal_type = is_causal ? CausalType::WindowedAttention : CausalType::None;
+    if (static_cast<int64_t>(sdp::CustomMaskType::CausalFromTopLeft) == custom_mask_type) {
+      params.window_left = WindowValue::TopLeftAligned;
+      params.window_right = WindowValue::TopLeftAligned;
+    } else if (static_cast<int64_t>(sdp::CustomMaskType::CausalFromBottomRight) == custom_mask_type) {
+      params.window_left = WindowValue::BottomRightAligned;
+      params.window_right = WindowValue::BottomRightAligned;
+    }
+    if (bias.has_value()) {
+      params.B = mk_aotensor(bias.value(), "bias");
+    }
+    if (seqstart_q.has_value()) {
+      params.varlen_type = VarlenType::CompactVarlen;
+      params.cu_seqlens_q = mk_aotensor<1>(seqstart_q.value(), "cu_seqlens_q");
+      params.cu_seqlens_k = mk_aotensor<1>(seqstart_k.value(), "cu_seqlens_k");
     } else {
-      err = attn_fwd(mk_aotensor(q_t, "q"),
-                     mk_aotensor(k_t, "k"),
-                     mk_aotensor(v_t, "v"),
-                     bias.has_value() ? mk_aotensor(bias.value(), "bias"): empty_t4,
-                     softmax_scale,
-                     compute_logsumexp ? mk_aotensor<2>(softmax_lse, "M") : empty_t2,
-                     mk_aotensor(output_t, "Out"),
-                     dropout_p,
-                     seed,
-                     offset1,
-                     offset2,
-                     seed_output,
-                     offset_output,
-                     mk_aotensor(softmax_fa_t, "encoded_softmax"),
-                     is_causal,
-                     persistent_counter,
-                     stream);
+      params.varlen_type = VarlenType::None;
     }
+    err = aotriton::v3::flash::attn_fwd(params,
+                                        aotriton::v3::flash::attn_fwd_params::kVersion,
+                                        stream);
 #else
     TORCH_CHECK(false, "Attempting to use AOTriton mem_eff_forward backend in a build that has not built AOTriton");
 #endif
diff --git a/aten/src/ATen/native/transformers/cuda/attention_backward.cu b/aten/src/ATen/native/transformers/cuda/attention_backward.cu
index 183f99e975cda..b7e52617a0ec6 100644
--- a/aten/src/ATen/native/transformers/cuda/attention_backward.cu
+++ b/aten/src/ATen/native/transformers/cuda/attention_backward.cu
@@ -547,6 +547,17 @@ _efficient_attention_backward(
                 "[AOTriton] Accelerated SDPA only supports MI200/MI300X/7900XTX/9070XT GPUs"
                 " (gfx90a/gfx942/gfx1100/gfx1201)")
     }
+    bool deterministic{false};
+    auto& ctx = at::globalContext();
+    if (ctx.deterministicAlgorithms()) {
+      if (ctx.deterministicAlgorithmsWarnOnly()) {
+        TORCH_WARN_ONCE(
+            "Memory Efficient attention defaults to a non-deterministic algorithm. ",
+            "To explicitly enable determinism call torch.use_deterministic_algorithms(True, warn_only=False).");
+      } else {
+        deterministic = true;
+      }
+    }
     const auto softmax_scale = sdp::calculate_scale(query, scale).expect_float();
     bool is_causal;
     if (static_cast<int64_t>(sdp::CustomMaskType::NoCustomMask) == custom_mask_type) {
@@ -569,139 +580,60 @@ _efficient_attention_backward(
     at::Tensor dout_t = grad_out.permute({0,2,1,3});
     at::Tensor softmax_lse = logsumexp.view({B * nH, max_seqlen_q});
     hipError_t err;
-    using aotriton::v2::flash::attn_bwd;
-    using aotriton::v2::flash::attn_bwd_fused;
-    using aotriton::v2::flash::attn_bwd_compact_varlen;
     using sdp::aotriton_adapter::mk_aotensor;
     using sdp::aotriton_adapter::mk_aoscalartensor;
     using sdp::aotriton_adapter::cast_dtype;
     aotriton::TensorView<4> empty_t4(0, {0, 0, 0, 0}, {0, 0, 0, 0}, cast_dtype(query.dtype()));
-    if constexpr (AOTRITON_ALWAYS_V3_API) {  // Better readability than nesting ifdef
-#if AOTRITON_V3_API  // if constexpr does not stop errors from undefined functions
-      using aotriton::v3::flash::CausalType;
-      using aotriton::v3::flash::VarlenType;
-      using aotriton::v3::flash::WindowValue;
-      aotriton::v3::flash::attn_bwd_params params;
-      params.Q = mk_aotensor(q_t, "q");
-      params.K = mk_aotensor(k_t, "k");
-      params.V = mk_aotensor(v_t, "v");
-      params.B = bias.has_value() ? mk_aotensor(bias.value(), "bias") : empty_t4;
-      params.Sm_scale = softmax_scale;
-      params.Out = mk_aotensor(out_t, "out");
-      params.DO = mk_aotensor(dout_t, "dout");
-      params.DK = mk_aotensor(dk_t, "dk");
-      params.DV = mk_aotensor(dv_t, "dv");
-      params.DQ = mk_aotensor(dq_t, "dq");
-      params.DB = bias_requires_grad ? mk_aotensor(grad_bias, "db") : empty_t4;
-      params.L = mk_aotensor<2>(softmax_lse, "L");
-      params.Max_seqlen_q = max_seqlen_q;        // Unused if cu_seqlens_q is empty
-      params.Max_seqlen_k = max_seqlen_k;        // Unused if cu_seqlens_k is empty
-      params.dropout_p = float(dropout_p);
-      params.philox_seed_ptr =  mk_aoscalartensor(philox_seed);
-      params.philox_offset1 = mk_aoscalartensor(philox_offset);
-      params.philox_offset2 = 0;
-      params.causal_type = is_causal ? CausalType::WindowedAttention : CausalType::None;
-      if (static_cast<int64_t>(sdp::CustomMaskType::CausalFromTopLeft) == custom_mask_type) {
-        params.window_left = WindowValue::TopLeftAligned;
-        params.window_right = WindowValue::TopLeftAligned;
-      } else if (static_cast<int64_t>(sdp::CustomMaskType::CausalFromBottomRight) == custom_mask_type) {
-        params.window_left = WindowValue::BottomRightAligned;
-        params.window_right = WindowValue::BottomRightAligned;
-      }
-#if AOTRITON_ALWAYS_V3_API
-      using sdp::aotriton_adapter::mklazy_empty_like;
-      using sdp::aotriton_adapter::mklazy_fp32zeros;
-      using sdp::aotriton_adapter::LazyTensorContext;
-      LazyTensorContext lazy_delta { .like_tensor = softmax_lse, .tensor_name = "delta" };
-      LazyTensorContext lazy_dq_acc { .like_tensor = dq_t, .tensor_name = "dq_acc" };
-      params.D = mklazy_empty_like<2>(&lazy_delta);
-      params.DQ_ACC = mklazy_fp32zeros<4>(&lazy_dq_acc);
-#else
-      at::Tensor delta = at::empty_like(softmax_lse).contiguous();
-      params.D = mk_aotensor<2>(delta, "delta");
-#endif
-      if (cu_seqlens_q.has_value()) {
-        params.varlen_type = VarlenType::CompactVarlen;
-        params.cu_seqlens_q = mk_aotensor<1>(cu_seqlens_q.value(), "cu_seqlens_q");
-        params.cu_seqlens_k = mk_aotensor<1>(cu_seqlens_k.value(), "cu_seqlens_k");
-      } else {
-        params.varlen_type = VarlenType::None;
-      }
-      err = aotriton::v3::flash::attn_bwd(params,
-                                          aotriton::v3::flash::attn_bwd_params::kVersion,
-                                          stream);
-#endif  // AOTRITON_V3_API
-    } else if (cu_seqlens_q.has_value()) {
-      at::Tensor delta = at::empty_like(softmax_lse).contiguous();
-      // varlen aka Nested tensor
-      err = attn_bwd_compact_varlen(mk_aotensor(q_t, "q"),
-                                    mk_aotensor(k_t, "k"),
-                                    mk_aotensor(v_t, "v"),
-                                    mk_aotensor<1>(cu_seqlens_q.value(), "cu_seqlens_q"),
-                                    mk_aotensor<1>(cu_seqlens_k.value(), "cu_seqlens_k"),
-                                    max_seqlen_q,
-                                    max_seqlen_k,
-                                    bias.has_value() ? mk_aotensor(bias.value(), "bias") : empty_t4,
-                                    softmax_scale,
-                                    mk_aotensor(out_t, "out"),
-                                    mk_aotensor(dout_t, "dout"),
-                                    mk_aotensor(dq_t, "dq"),
-                                    mk_aotensor(dk_t, "dk"),
-                                    mk_aotensor(dv_t, "dv"),
-                                    bias_requires_grad ? mk_aotensor(grad_bias, "db") : empty_t4,
-                                    mk_aotensor<2>(softmax_lse, "L"),
-                                    mk_aotensor<2>(delta, "delta"),
-                                    float(dropout_p),
-                                    mk_aoscalartensor(philox_seed),
-                                    mk_aoscalartensor(philox_offset),
-                                    0,
-                                    is_causal,
-                                    stream);
-    } else { // cu_seqlens.has_value
-      auto d_head = Kv;
-      bool use_fused_bwd = d_head <= 192 && d_head * max_seqlen_q < 64 * 512;
-      if (use_fused_bwd) {
-        err = attn_bwd_fused(mk_aotensor(q_t, "q"),
-                             mk_aotensor(k_t, "k"),
-                             mk_aotensor(v_t, "v"),
-                             bias.has_value() ? mk_aotensor(bias.value(), "bias") : empty_t4,
-                             softmax_scale,
-                             mk_aotensor(out_t, "out"),
-                             mk_aotensor(dout_t, "dout"),
-                             mk_aotensor(dq_t, "dq"),
-                             mk_aotensor(dk_t, "dk"),
-                             mk_aotensor(dv_t, "dv"),
-                             bias_requires_grad ? mk_aotensor(grad_bias, "db") : empty_t4,
-                             mk_aotensor<2>(softmax_lse, "L"),
-                             float(dropout_p),
-                             mk_aoscalartensor(philox_seed),
-                             mk_aoscalartensor(philox_offset),
-                             0,
-                             is_causal,
-                             stream);
-      } else {
-        at::Tensor delta = at::empty_like(softmax_lse).contiguous();
-        err = attn_bwd(mk_aotensor(q_t, "q"),
-                     mk_aotensor(k_t, "k"),
-                     mk_aotensor(v_t, "v"),
-                     bias.has_value() ? mk_aotensor(bias.value(), "bias") : empty_t4,
-                     softmax_scale,
-                     mk_aotensor(out_t, "out"),
-                     mk_aotensor(dout_t, "dout"),
-                     mk_aotensor(dq_t, "dq"),
-                     mk_aotensor(dk_t, "dk"),
-                     mk_aotensor(dv_t, "dv"),
-                     bias_requires_grad ? mk_aotensor(grad_bias, "db") : empty_t4,
-                     mk_aotensor<2>(softmax_lse, "L"),
-                     mk_aotensor<2>(delta, "delta"),
-                     float(dropout_p),
-                     mk_aoscalartensor(philox_seed),
-                     mk_aoscalartensor(philox_offset),
-                     0,
-                     is_causal,
-                     stream);
-      } //used_fused_bwd
-    } // cuseqlen.has_value
+    using aotriton::v3::flash::CausalType;
+    using aotriton::v3::flash::VarlenType;
+    using aotriton::v3::flash::WindowValue;
+    aotriton::v3::flash::attn_bwd_params params;
+    params.Q = mk_aotensor(q_t, "q");
+    params.K = mk_aotensor(k_t, "k");
+    params.V = mk_aotensor(v_t, "v");
+    params.B = bias.has_value() ? mk_aotensor(bias.value(), "bias") : empty_t4;
+    params.Sm_scale = softmax_scale;
+    params.Out = mk_aotensor(out_t, "out");
+    params.DO = mk_aotensor(dout_t, "dout");
+    params.DK = mk_aotensor(dk_t, "dk");
+    params.DV = mk_aotensor(dv_t, "dv");
+    params.DQ = mk_aotensor(dq_t, "dq");
+    params.DB = bias_requires_grad ? mk_aotensor(grad_bias, "db") : empty_t4;
+    params.L = mk_aotensor<2>(softmax_lse, "L");
+    params.Max_seqlen_q = max_seqlen_q;        // Unused if cu_seqlens_q is empty
+    params.Max_seqlen_k = max_seqlen_k;        // Unused if cu_seqlens_k is empty
+    params.dropout_p = float(dropout_p);
+    params.philox_seed_ptr =  mk_aoscalartensor(philox_seed);
+    params.philox_offset1 = mk_aoscalartensor(philox_offset);
+    params.philox_offset2 = 0;
+    params.causal_type = is_causal ? CausalType::WindowedAttention : CausalType::None;
+    if (static_cast<int64_t>(sdp::CustomMaskType::CausalFromTopLeft) == custom_mask_type) {
+      params.window_left = WindowValue::TopLeftAligned;
+      params.window_right = WindowValue::TopLeftAligned;
+    } else if (static_cast<int64_t>(sdp::CustomMaskType::CausalFromBottomRight) == custom_mask_type) {
+      params.window_left = WindowValue::BottomRightAligned;
+      params.window_right = WindowValue::BottomRightAligned;
+    }
+    using sdp::aotriton_adapter::mklazy_empty_like;
+    using sdp::aotriton_adapter::mklazy_fp32zeros;
+    using sdp::aotriton_adapter::LazyTensorContext;
+    LazyTensorContext lazy_delta { .like_tensor = softmax_lse, .tensor_name = "delta" };
+    LazyTensorContext lazy_dq_acc { .like_tensor = dq_t, .tensor_name = "dq_acc" };
+    params.D = mklazy_empty_like<2>(&lazy_delta);
+    params.DQ_ACC = mklazy_fp32zeros<4>(&lazy_dq_acc);
+    if (cu_seqlens_q.has_value()) {
+      params.varlen_type = VarlenType::CompactVarlen;
+      params.cu_seqlens_q = mk_aotensor<1>(cu_seqlens_q.value(), "cu_seqlens_q");
+      params.cu_seqlens_k = mk_aotensor<1>(cu_seqlens_k.value(), "cu_seqlens_k");
+    } else {
+      params.varlen_type = VarlenType::None;
+    }
+    aotriton::v3::flash::attn_options opts;
+    opts.deterministic = deterministic;
+    err = aotriton::v3::flash::attn_bwd(params,
+                                        aotriton::v3::flash::attn_bwd_params::kVersion,
+                                        stream,
+                                        &opts);
 #else  // DISABLE_AOTRITON
     TORCH_CHECK(false, "Attempting to use aotriton mem_eff_backward backend in a build that has not built AOTriton");
 #endif
diff --git a/aten/src/ATen/native/transformers/cuda/mem_eff_attention/epilogue/epilogue_thread_apply_logsumexp.h b/aten/src/ATen/native/transformers/cuda/mem_eff_attention/epilogue/epilogue_thread_apply_logsumexp.h
index 156034954d9e6..58e5234f53526 100644
--- a/aten/src/ATen/native/transformers/cuda/mem_eff_attention/epilogue/epilogue_thread_apply_logsumexp.h
+++ b/aten/src/ATen/native/transformers/cuda/mem_eff_attention/epilogue/epilogue_thread_apply_logsumexp.h
@@ -35,8 +35,6 @@
 
 #pragma once
 
-#include <cuda_fp16.h>
-
 #include <cutlass/array.h>
 #include <cutlass/cutlass.h>
 #include <cutlass/epilogue/thread/activation.h>
@@ -70,26 +68,21 @@ struct ArrayExponential {
   }
 };
 
+template <int N>
+struct dependent_false : std::false_type {};
+
 template <int ElementsPerAccess>
 struct ArrayExponential<half_t, ElementsPerAccess> {
-  CUTLASS_DEVICE
-  Array<half_t, ElementsPerAccess> operator()(
-      Array<half_t, ElementsPerAccess> const& input) const {
-    Array<half_t, ElementsPerAccess> result;
-
-    int const kVectorCount = ElementsPerAccess / 2;
-
-    __half2 const* input_ptr =
-        reinterpret_cast<__half2 const*>(input.raw_data());
-    __half2* res_ptr = reinterpret_cast<__half2*>(result.raw_data());
-
-    CUTLASS_PRAGMA_UNROLL
-    for (int i = 0; i < kVectorCount; ++i) {
-      res_ptr[i] = h2exp(input_ptr[i]);
-    }
+  static_assert(
+      dependent_false<ElementsPerAccess>::value,
+      "ArrayExponential is not implemented for half_t");
+};
 
-    return result;
-  }
+template <int ElementsPerAccess>
+struct ArrayExponential<bfloat16_t, ElementsPerAccess> {
+  static_assert(
+      dependent_false<ElementsPerAccess>::value,
+      "ArrayExponential is not implemented for bfloat16_t");
 };
 } // namespace detail
 
diff --git a/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB.h b/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB.h
index 384fc9130ef21..5a8709f8579c6 100644
--- a/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB.h
+++ b/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB.h
@@ -436,16 +436,6 @@ __global__ void __launch_bounds__(
     AttentionBackwardKernel<cutlass::arch::Sm50, float, true, false, false, 64, 64, 65536>::kNumThreads,
     AttentionBackwardKernel<cutlass::arch::Sm50, float, true, false, false, 64, 64, 65536>::kMinBlocksPerSm)
 fmha_cutlassB_f32_aligned_64x64_k65536_sm50(typename AttentionBackwardKernel<cutlass::arch::Sm50, float, true, false, false, 64, 64, 65536>::Params p);
-#if defined(CUDA_VERSION) && CUDA_VERSION == 12040 && !defined(USE_ROCM)
-__global__ void __launch_bounds__(
-    AttentionBackwardKernel<cutlass::arch::Sm50, float, true, true, false, 32, 32, 32>::kNumThreads,
-    AttentionBackwardKernel<cutlass::arch::Sm50, float, true, true, false, 32, 32, 32>::kMinBlocksPerSm)
-fmha_cutlassB_f32_aligned_32x32_k32_dropout_sm50(typename AttentionBackwardKernel<cutlass::arch::Sm50, float, true, true, false, 32, 32, 32>::Params p);
-__global__ void __launch_bounds__(
-    AttentionBackwardKernel<cutlass::arch::Sm50, float, true, true, false, 32, 32, 64>::kNumThreads,
-    AttentionBackwardKernel<cutlass::arch::Sm50, float, true, true, false, 32, 32, 64>::kMinBlocksPerSm)
-fmha_cutlassB_f32_aligned_32x32_k64_dropout_sm50(typename AttentionBackwardKernel<cutlass::arch::Sm50, float, true, true, false, 32, 32, 64>::Params p);
-#else
 __global__ void __launch_bounds__(
     AttentionBackwardKernel<cutlass::arch::Sm50, float, true, true, false, 64, 64, 32>::kNumThreads,
     AttentionBackwardKernel<cutlass::arch::Sm50, float, true, true, false, 64, 64, 32>::kMinBlocksPerSm)
@@ -454,7 +444,6 @@ __global__ void __launch_bounds__(
     AttentionBackwardKernel<cutlass::arch::Sm50, float, true, true, false, 64, 64, 64>::kNumThreads,
     AttentionBackwardKernel<cutlass::arch::Sm50, float, true, true, false, 64, 64, 64>::kMinBlocksPerSm)
 fmha_cutlassB_f32_aligned_64x64_k64_dropout_sm50(typename AttentionBackwardKernel<cutlass::arch::Sm50, float, true, true, false, 64, 64, 64>::Params p);
-#endif
 __global__ void __launch_bounds__(
     AttentionBackwardKernel<cutlass::arch::Sm50, float, true, true, false, 64, 64, 128>::kNumThreads,
     AttentionBackwardKernel<cutlass::arch::Sm50, float, true, true, false, 64, 64, 128>::kMinBlocksPerSm)
@@ -501,13 +490,8 @@ template <typename T> void dispatch_cutlassB_f32_sm50(T cb, int cc) {
     cb(AttentionBackwardKernel<cutlass::arch::Sm50, float, true, false, false, 64, 64, 64>(), fmha_cutlassB_f32_aligned_64x64_k64_sm50);
     cb(AttentionBackwardKernel<cutlass::arch::Sm50, float, true, false, false, 64, 64, 128>(), fmha_cutlassB_f32_aligned_64x64_k128_sm50);
     cb(AttentionBackwardKernel<cutlass::arch::Sm50, float, true, false, false, 64, 64, 65536>(), fmha_cutlassB_f32_aligned_64x64_k65536_sm50);
-#if defined(CUDA_VERSION) && CUDA_VERSION == 12040 && !defined(USE_ROCM)
-    cb(AttentionBackwardKernel<cutlass::arch::Sm50, float, true, true, false, 32, 32, 32>(), fmha_cutlassB_f32_aligned_32x32_k32_dropout_sm50);
-    cb(AttentionBackwardKernel<cutlass::arch::Sm50, float, true, true, false, 32, 32, 64>(), fmha_cutlassB_f32_aligned_32x32_k64_dropout_sm50);
-#else
     cb(AttentionBackwardKernel<cutlass::arch::Sm50, float, true, true, false, 64, 64, 32>(), fmha_cutlassB_f32_aligned_64x64_k32_dropout_sm50);
     cb(AttentionBackwardKernel<cutlass::arch::Sm50, float, true, true, false, 64, 64, 64>(), fmha_cutlassB_f32_aligned_64x64_k64_dropout_sm50);
-#endif
     cb(AttentionBackwardKernel<cutlass::arch::Sm50, float, true, true, false, 64, 64, 128>(), fmha_cutlassB_f32_aligned_64x64_k128_dropout_sm50);
     cb(AttentionBackwardKernel<cutlass::arch::Sm50, float, true, true, false, 64, 64, 65536>(), fmha_cutlassB_f32_aligned_64x64_k65536_dropout_sm50);
     cb(AttentionBackwardKernel<cutlass::arch::Sm50, float, false, false, false, 64, 64, 32>(), fmha_cutlassB_f32_notaligned_64x64_k32_sm50);
@@ -884,31 +868,31 @@ template <typename T> void dispatch_cutlassB_f32_sm80(T cb, int cc) {
 template <typename DT, typename T>
 void dispatch_cutlassB(T cb, int cc = 0) {
 
-    if (std::is_same_v<DT, cutlass::half_t> && 70 <= cc && cc < 75) {
+    if (std::is_same_v<DT, cutlass::half_t> && 70 <= cc && cc <= 74) {
         dispatch_cutlassB_f16_sm70(cb, cc);
     }
-    if (std::is_same_v<DT, cutlass::bfloat16_t> && 80 <= cc && cc <= 120) {
+    if (std::is_same_v<DT, cutlass::bfloat16_t> && 80 <= cc && cc <= 121) {
         dispatch_cutlassB_bf16_sm80(cb, cc);
     }
-    if (std::is_same_v<DT, cutlass::half_t> && 80 <= cc && cc <= 120) {
+    if (std::is_same_v<DT, cutlass::half_t> && 80 <= cc && cc <= 121) {
         dispatch_cutlassB_f16_sm80(cb, cc);
     }
-    if (std::is_same_v<DT, cutlass::half_t> && 50 <= cc && cc < 70) {
+    if (std::is_same_v<DT, cutlass::half_t> && 50 <= cc && cc <= 69) {
         dispatch_cutlassB_f16_sm50(cb, cc);
     }
-    if (std::is_same_v<DT, float> && 50 <= cc && cc < 70) {
+    if (std::is_same_v<DT, float> && 50 <= cc && cc <= 69) {
         dispatch_cutlassB_f32_sm50(cb, cc);
     }
-    if (std::is_same_v<DT, float> && 70 <= cc && cc < 75) {
+    if (std::is_same_v<DT, float> && 70 <= cc && cc <= 74) {
         dispatch_cutlassB_f32_sm70(cb, cc);
     }
-    if (std::is_same_v<DT, cutlass::half_t> && 75 <= cc && cc < 80) {
+    if (std::is_same_v<DT, cutlass::half_t> && 75 <= cc && cc <= 79) {
         dispatch_cutlassB_f16_sm75(cb, cc);
     }
-    if (std::is_same_v<DT, float> && 75 <= cc && cc < 80) {
+    if (std::is_same_v<DT, float> && 75 <= cc && cc <= 79) {
         dispatch_cutlassB_f32_sm75(cb, cc);
     }
-    if (std::is_same_v<DT, float> && 80 <= cc && cc <= 120) {
+    if (std::is_same_v<DT, float> && 80 <= cc && cc <= 121) {
         dispatch_cutlassB_f32_sm80(cb, cc);
     }
 }
diff --git a/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_bf16_aligned_k128.cu b/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_bf16_aligned_k128.cu
index 6b0bfdc62dcb7..1502d527cc4c6 100644
--- a/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_bf16_aligned_k128.cu
+++ b/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_bf16_aligned_k128.cu
@@ -14,7 +14,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_bf16_aligned_128x128_k128_seqaligned_sm80(typename AttentionBackwardKernel<cutlass::arch::Sm80, cutlass::bfloat16_t, true, false, true, 128, 128, 128, true>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 800
-#if __CUDA_ARCH__ <= 1200
+#if __CUDA_ARCH__ <= 1210
   if (!p.advance_to_block()) {
     return;
   }
@@ -23,7 +23,7 @@ fmha_cutlassB_bf16_aligned_128x128_k128_seqaligned_sm80(typename AttentionBackwa
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_bf16_aligned_128x128_k128_seqaligned_sm80` is for sm80-sm100, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_bf16_aligned_128x128_k128_seqaligned_sm80` is for sm80-sm121, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -33,7 +33,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_bf16_aligned_128x128_k128_sm80(typename AttentionBackwardKernel<cutlass::arch::Sm80, cutlass::bfloat16_t, true, false, true, 128, 128, 128>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 800
-#if __CUDA_ARCH__ <= 1200
+#if __CUDA_ARCH__ <= 1210
   if (!p.advance_to_block()) {
     return;
   }
@@ -42,7 +42,7 @@ fmha_cutlassB_bf16_aligned_128x128_k128_sm80(typename AttentionBackwardKernel<cu
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_bf16_aligned_128x128_k128_sm80` is for sm80-sm100, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_bf16_aligned_128x128_k128_sm80` is for sm80-sm121, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -52,7 +52,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_bf16_aligned_64x64_k128_seqaligned_sm80(typename AttentionBackwardKernel<cutlass::arch::Sm80, cutlass::bfloat16_t, true, false, false, 64, 64, 128, true>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 800
-#if __CUDA_ARCH__ <= 1200
+#if __CUDA_ARCH__ <= 1210
   if (!p.advance_to_block()) {
     return;
   }
@@ -61,7 +61,7 @@ fmha_cutlassB_bf16_aligned_64x64_k128_seqaligned_sm80(typename AttentionBackward
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_bf16_aligned_64x64_k128_seqaligned_sm80` is for sm80-sm100, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_bf16_aligned_64x64_k128_seqaligned_sm80` is for sm80-sm121, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -71,7 +71,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_bf16_aligned_64x64_k128_sm80(typename AttentionBackwardKernel<cutlass::arch::Sm80, cutlass::bfloat16_t, true, false, false, 64, 64, 128>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 800
-#if __CUDA_ARCH__ <= 1200
+#if __CUDA_ARCH__ <= 1210
   if (!p.advance_to_block()) {
     return;
   }
@@ -80,7 +80,7 @@ fmha_cutlassB_bf16_aligned_64x64_k128_sm80(typename AttentionBackwardKernel<cutl
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_bf16_aligned_64x64_k128_sm80` is for sm80-sm100, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_bf16_aligned_64x64_k128_sm80` is for sm80-sm121, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
diff --git a/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_bf16_aligned_k128_dropout.cu b/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_bf16_aligned_k128_dropout.cu
index 76c87287a6c80..dee3b893144c3 100644
--- a/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_bf16_aligned_k128_dropout.cu
+++ b/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_bf16_aligned_k128_dropout.cu
@@ -14,7 +14,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_bf16_aligned_128x128_k128_dropout_sm80(typename AttentionBackwardKernel<cutlass::arch::Sm80, cutlass::bfloat16_t, true, true, true, 128, 128, 128>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 800
-#if __CUDA_ARCH__ <= 1200
+#if __CUDA_ARCH__ <= 1210
   if (!p.advance_to_block()) {
     return;
   }
@@ -23,7 +23,7 @@ fmha_cutlassB_bf16_aligned_128x128_k128_dropout_sm80(typename AttentionBackwardK
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_bf16_aligned_128x128_k128_dropout_sm80` is for sm80-sm100, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_bf16_aligned_128x128_k128_dropout_sm80` is for sm80-sm121, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -33,7 +33,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_bf16_aligned_64x64_k128_dropout_sm80(typename AttentionBackwardKernel<cutlass::arch::Sm80, cutlass::bfloat16_t, true, true, false, 64, 64, 128>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 800
-#if __CUDA_ARCH__ <= 1200
+#if __CUDA_ARCH__ <= 1210
   if (!p.advance_to_block()) {
     return;
   }
@@ -42,7 +42,7 @@ fmha_cutlassB_bf16_aligned_64x64_k128_dropout_sm80(typename AttentionBackwardKer
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_bf16_aligned_64x64_k128_dropout_sm80` is for sm80-sm100, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_bf16_aligned_64x64_k128_dropout_sm80` is for sm80-sm121, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
diff --git a/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_bf16_aligned_k32.cu b/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_bf16_aligned_k32.cu
index 79b731c55b129..f235218f135dd 100644
--- a/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_bf16_aligned_k32.cu
+++ b/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_bf16_aligned_k32.cu
@@ -14,7 +14,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_bf16_aligned_64x64_k32_seqaligned_sm80(typename AttentionBackwardKernel<cutlass::arch::Sm80, cutlass::bfloat16_t, true, false, true, 64, 64, 32, true>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 800
-#if __CUDA_ARCH__ <= 1200
+#if __CUDA_ARCH__ <= 1210
   if (!p.advance_to_block()) {
     return;
   }
@@ -23,7 +23,7 @@ fmha_cutlassB_bf16_aligned_64x64_k32_seqaligned_sm80(typename AttentionBackwardK
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_bf16_aligned_64x64_k32_seqaligned_sm80` is for sm80-sm100, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_bf16_aligned_64x64_k32_seqaligned_sm80` is for sm80-sm121, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -33,7 +33,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_bf16_aligned_64x64_k32_sm80(typename AttentionBackwardKernel<cutlass::arch::Sm80, cutlass::bfloat16_t, true, false, true, 64, 64, 32>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 800
-#if __CUDA_ARCH__ <= 1200
+#if __CUDA_ARCH__ <= 1210
   if (!p.advance_to_block()) {
     return;
   }
@@ -42,7 +42,7 @@ fmha_cutlassB_bf16_aligned_64x64_k32_sm80(typename AttentionBackwardKernel<cutla
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_bf16_aligned_64x64_k32_sm80` is for sm80-sm100, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_bf16_aligned_64x64_k32_sm80` is for sm80-sm121, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
diff --git a/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_bf16_aligned_k32_dropout.cu b/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_bf16_aligned_k32_dropout.cu
index b99b516477df3..55cc359174c81 100644
--- a/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_bf16_aligned_k32_dropout.cu
+++ b/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_bf16_aligned_k32_dropout.cu
@@ -14,7 +14,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_bf16_aligned_64x64_k32_dropout_sm80(typename AttentionBackwardKernel<cutlass::arch::Sm80, cutlass::bfloat16_t, true, true, true, 64, 64, 32>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 800
-#if __CUDA_ARCH__ <= 1200
+#if __CUDA_ARCH__ <= 1210
   if (!p.advance_to_block()) {
     return;
   }
@@ -23,7 +23,7 @@ fmha_cutlassB_bf16_aligned_64x64_k32_dropout_sm80(typename AttentionBackwardKern
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_bf16_aligned_64x64_k32_dropout_sm80` is for sm80-sm100, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_bf16_aligned_64x64_k32_dropout_sm80` is for sm80-sm121, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
diff --git a/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_bf16_aligned_k64.cu b/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_bf16_aligned_k64.cu
index 08c544561d958..17926cf09d6b4 100644
--- a/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_bf16_aligned_k64.cu
+++ b/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_bf16_aligned_k64.cu
@@ -14,7 +14,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_bf16_aligned_64x64_k64_seqaligned_sm80(typename AttentionBackwardKernel<cutlass::arch::Sm80, cutlass::bfloat16_t, true, false, true, 64, 64, 64, true>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 800
-#if __CUDA_ARCH__ <= 1200
+#if __CUDA_ARCH__ <= 1210
   if (!p.advance_to_block()) {
     return;
   }
@@ -23,7 +23,7 @@ fmha_cutlassB_bf16_aligned_64x64_k64_seqaligned_sm80(typename AttentionBackwardK
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_bf16_aligned_64x64_k64_seqaligned_sm80` is for sm80-sm100, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_bf16_aligned_64x64_k64_seqaligned_sm80` is for sm80-sm121, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -33,7 +33,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_bf16_aligned_64x64_k64_sm80(typename AttentionBackwardKernel<cutlass::arch::Sm80, cutlass::bfloat16_t, true, false, true, 64, 64, 64>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 800
-#if __CUDA_ARCH__ <= 1200
+#if __CUDA_ARCH__ <= 1210
   if (!p.advance_to_block()) {
     return;
   }
@@ -42,7 +42,7 @@ fmha_cutlassB_bf16_aligned_64x64_k64_sm80(typename AttentionBackwardKernel<cutla
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_bf16_aligned_64x64_k64_sm80` is for sm80-sm100, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_bf16_aligned_64x64_k64_sm80` is for sm80-sm121, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
diff --git a/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_bf16_aligned_k64_dropout.cu b/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_bf16_aligned_k64_dropout.cu
index 04308a418d1a2..c2fecd50a9a8e 100644
--- a/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_bf16_aligned_k64_dropout.cu
+++ b/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_bf16_aligned_k64_dropout.cu
@@ -14,7 +14,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_bf16_aligned_64x64_k64_dropout_sm80(typename AttentionBackwardKernel<cutlass::arch::Sm80, cutlass::bfloat16_t, true, true, true, 64, 64, 64>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 800
-#if __CUDA_ARCH__ <= 1200
+#if __CUDA_ARCH__ <= 1210
   if (!p.advance_to_block()) {
     return;
   }
@@ -23,7 +23,7 @@ fmha_cutlassB_bf16_aligned_64x64_k64_dropout_sm80(typename AttentionBackwardKern
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_bf16_aligned_64x64_k64_dropout_sm80` is for sm80-sm100, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_bf16_aligned_64x64_k64_dropout_sm80` is for sm80-sm121, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
diff --git a/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_bf16_aligned_k65536.cu b/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_bf16_aligned_k65536.cu
index 31546f71e8641..c3dcb8a597cbe 100644
--- a/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_bf16_aligned_k65536.cu
+++ b/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_bf16_aligned_k65536.cu
@@ -14,7 +14,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_bf16_aligned_128x64_k65536_sm80(typename AttentionBackwardKernel<cutlass::arch::Sm80, cutlass::bfloat16_t, true, false, false, 128, 64, 65536>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 800
-#if __CUDA_ARCH__ <= 1200
+#if __CUDA_ARCH__ <= 1210
   if (!p.advance_to_block()) {
     return;
   }
@@ -23,7 +23,7 @@ fmha_cutlassB_bf16_aligned_128x64_k65536_sm80(typename AttentionBackwardKernel<c
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_bf16_aligned_128x64_k65536_sm80` is for sm80-sm100, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_bf16_aligned_128x64_k65536_sm80` is for sm80-sm121, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -33,7 +33,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_bf16_aligned_64x64_k65536_sm80(typename AttentionBackwardKernel<cutlass::arch::Sm80, cutlass::bfloat16_t, true, false, false, 64, 64, 65536>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 800
-#if __CUDA_ARCH__ <= 1200
+#if __CUDA_ARCH__ <= 1210
   if (!p.advance_to_block()) {
     return;
   }
@@ -42,7 +42,7 @@ fmha_cutlassB_bf16_aligned_64x64_k65536_sm80(typename AttentionBackwardKernel<cu
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_bf16_aligned_64x64_k65536_sm80` is for sm80-sm100, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_bf16_aligned_64x64_k65536_sm80` is for sm80-sm121, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
diff --git a/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_bf16_aligned_k65536_dropout.cu b/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_bf16_aligned_k65536_dropout.cu
index 04c8b728bc629..225c97172d6d4 100644
--- a/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_bf16_aligned_k65536_dropout.cu
+++ b/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_bf16_aligned_k65536_dropout.cu
@@ -14,7 +14,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_bf16_aligned_128x64_k65536_dropout_sm80(typename AttentionBackwardKernel<cutlass::arch::Sm80, cutlass::bfloat16_t, true, true, false, 128, 64, 65536>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 800
-#if __CUDA_ARCH__ <= 1200
+#if __CUDA_ARCH__ <= 1210
   if (!p.advance_to_block()) {
     return;
   }
@@ -23,7 +23,7 @@ fmha_cutlassB_bf16_aligned_128x64_k65536_dropout_sm80(typename AttentionBackward
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_bf16_aligned_128x64_k65536_dropout_sm80` is for sm80-sm100, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_bf16_aligned_128x64_k65536_dropout_sm80` is for sm80-sm121, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -33,7 +33,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_bf16_aligned_64x64_k65536_dropout_sm80(typename AttentionBackwardKernel<cutlass::arch::Sm80, cutlass::bfloat16_t, true, true, false, 64, 64, 65536>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 800
-#if __CUDA_ARCH__ <= 1200
+#if __CUDA_ARCH__ <= 1210
   if (!p.advance_to_block()) {
     return;
   }
@@ -42,7 +42,7 @@ fmha_cutlassB_bf16_aligned_64x64_k65536_dropout_sm80(typename AttentionBackwardK
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_bf16_aligned_64x64_k65536_dropout_sm80` is for sm80-sm100, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_bf16_aligned_64x64_k65536_dropout_sm80` is for sm80-sm121, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
diff --git a/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_bf16_aligned_k96.cu b/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_bf16_aligned_k96.cu
index c3bc1ec8b6495..da1a33f7f0b21 100644
--- a/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_bf16_aligned_k96.cu
+++ b/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_bf16_aligned_k96.cu
@@ -14,7 +14,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_bf16_aligned_128x64_k96_sm80(typename AttentionBackwardKernel<cutlass::arch::Sm80, cutlass::bfloat16_t, true, false, true, 128, 64, 96>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 800
-#if __CUDA_ARCH__ <= 1200
+#if __CUDA_ARCH__ <= 1210
   if (!p.advance_to_block()) {
     return;
   }
@@ -23,7 +23,7 @@ fmha_cutlassB_bf16_aligned_128x64_k96_sm80(typename AttentionBackwardKernel<cutl
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_bf16_aligned_128x64_k96_sm80` is for sm80-sm100, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_bf16_aligned_128x64_k96_sm80` is for sm80-sm121, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
diff --git a/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_f16_aligned_k128.cu b/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_f16_aligned_k128.cu
index cad26eaf1914b..ff30477445ad7 100644
--- a/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_f16_aligned_k128.cu
+++ b/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_f16_aligned_k128.cu
@@ -14,7 +14,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_f16_aligned_128x64_k128_seqaligned_sm70(typename AttentionBackwardKernel<cutlass::arch::Sm70, cutlass::half_t, true, false, false, 128, 64, 128, true>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 700
-#if __CUDA_ARCH__ < 750
+#if __CUDA_ARCH__ <= 740
   if (!p.advance_to_block()) {
     return;
   }
@@ -23,7 +23,7 @@ fmha_cutlassB_f16_aligned_128x64_k128_seqaligned_sm70(typename AttentionBackward
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_f16_aligned_128x64_k128_seqaligned_sm70` is for sm70-sm75, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_f16_aligned_128x64_k128_seqaligned_sm70` is for sm70-sm74, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -33,7 +33,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_f16_aligned_128x128_k128_seqaligned_sm80(typename AttentionBackwardKernel<cutlass::arch::Sm80, cutlass::half_t, true, false, true, 128, 128, 128, true>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 800
-#if __CUDA_ARCH__ <= 1200
+#if __CUDA_ARCH__ <= 1210
   if (!p.advance_to_block()) {
     return;
   }
@@ -42,7 +42,7 @@ fmha_cutlassB_f16_aligned_128x128_k128_seqaligned_sm80(typename AttentionBackwar
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_f16_aligned_128x128_k128_seqaligned_sm80` is for sm80-sm100, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_f16_aligned_128x128_k128_seqaligned_sm80` is for sm80-sm121, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -52,7 +52,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_f16_aligned_128x64_k128_sm70(typename AttentionBackwardKernel<cutlass::arch::Sm70, cutlass::half_t, true, false, false, 128, 64, 128>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 700
-#if __CUDA_ARCH__ < 750
+#if __CUDA_ARCH__ <= 740
   if (!p.advance_to_block()) {
     return;
   }
@@ -61,7 +61,7 @@ fmha_cutlassB_f16_aligned_128x64_k128_sm70(typename AttentionBackwardKernel<cutl
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_f16_aligned_128x64_k128_sm70` is for sm70-sm75, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_f16_aligned_128x64_k128_sm70` is for sm70-sm74, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -71,7 +71,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_f16_aligned_128x64_k128_sm75(typename AttentionBackwardKernel<cutlass::arch::Sm75, cutlass::half_t, true, false, false, 128, 64, 128>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 750
-#if __CUDA_ARCH__ < 800
+#if __CUDA_ARCH__ <= 790
   if (!p.advance_to_block()) {
     return;
   }
@@ -80,7 +80,7 @@ fmha_cutlassB_f16_aligned_128x64_k128_sm75(typename AttentionBackwardKernel<cutl
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_f16_aligned_128x64_k128_sm75` is for sm75-sm80, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_f16_aligned_128x64_k128_sm75` is for sm75-sm79, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -90,7 +90,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_f16_aligned_128x128_k128_sm80(typename AttentionBackwardKernel<cutlass::arch::Sm80, cutlass::half_t, true, false, true, 128, 128, 128>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 800
-#if __CUDA_ARCH__ <= 1200
+#if __CUDA_ARCH__ <= 1210
   if (!p.advance_to_block()) {
     return;
   }
@@ -99,7 +99,7 @@ fmha_cutlassB_f16_aligned_128x128_k128_sm80(typename AttentionBackwardKernel<cut
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_f16_aligned_128x128_k128_sm80` is for sm80-sm100, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_f16_aligned_128x128_k128_sm80` is for sm80-sm121, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -109,7 +109,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_f16_aligned_64x64_k128_seqaligned_sm70(typename AttentionBackwardKernel<cutlass::arch::Sm70, cutlass::half_t, true, false, false, 64, 64, 128, true>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 700
-#if __CUDA_ARCH__ < 750
+#if __CUDA_ARCH__ <= 740
   if (!p.advance_to_block()) {
     return;
   }
@@ -118,7 +118,7 @@ fmha_cutlassB_f16_aligned_64x64_k128_seqaligned_sm70(typename AttentionBackwardK
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_f16_aligned_64x64_k128_seqaligned_sm70` is for sm70-sm75, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_f16_aligned_64x64_k128_seqaligned_sm70` is for sm70-sm74, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -128,7 +128,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_f16_aligned_64x64_k128_seqaligned_sm80(typename AttentionBackwardKernel<cutlass::arch::Sm80, cutlass::half_t, true, false, false, 64, 64, 128, true>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 800
-#if __CUDA_ARCH__ <= 1200
+#if __CUDA_ARCH__ <= 1210
   if (!p.advance_to_block()) {
     return;
   }
@@ -137,7 +137,7 @@ fmha_cutlassB_f16_aligned_64x64_k128_seqaligned_sm80(typename AttentionBackwardK
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_f16_aligned_64x64_k128_seqaligned_sm80` is for sm80-sm100, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_f16_aligned_64x64_k128_seqaligned_sm80` is for sm80-sm121, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -147,7 +147,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_f16_aligned_64x64_k128_sm50(typename AttentionBackwardKernel<cutlass::arch::Sm50, cutlass::half_t, true, false, false, 64, 64, 128>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 500
-#if __CUDA_ARCH__ < 700
+#if __CUDA_ARCH__ <= 690
   if (!p.advance_to_block()) {
     return;
   }
@@ -156,7 +156,7 @@ fmha_cutlassB_f16_aligned_64x64_k128_sm50(typename AttentionBackwardKernel<cutla
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_f16_aligned_64x64_k128_sm50` is for sm50-sm70, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_f16_aligned_64x64_k128_sm50` is for sm50-sm69, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -166,7 +166,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_f16_aligned_64x64_k128_sm70(typename AttentionBackwardKernel<cutlass::arch::Sm70, cutlass::half_t, true, false, false, 64, 64, 128>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 700
-#if __CUDA_ARCH__ < 750
+#if __CUDA_ARCH__ <= 740
   if (!p.advance_to_block()) {
     return;
   }
@@ -175,7 +175,7 @@ fmha_cutlassB_f16_aligned_64x64_k128_sm70(typename AttentionBackwardKernel<cutla
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_f16_aligned_64x64_k128_sm70` is for sm70-sm75, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_f16_aligned_64x64_k128_sm70` is for sm70-sm74, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -185,7 +185,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_f16_aligned_64x64_k128_sm75(typename AttentionBackwardKernel<cutlass::arch::Sm75, cutlass::half_t, true, false, false, 64, 64, 128>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 750
-#if __CUDA_ARCH__ < 800
+#if __CUDA_ARCH__ <= 790
   if (!p.advance_to_block()) {
     return;
   }
@@ -194,7 +194,7 @@ fmha_cutlassB_f16_aligned_64x64_k128_sm75(typename AttentionBackwardKernel<cutla
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_f16_aligned_64x64_k128_sm75` is for sm75-sm80, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_f16_aligned_64x64_k128_sm75` is for sm75-sm79, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -204,7 +204,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_f16_aligned_64x64_k128_sm80(typename AttentionBackwardKernel<cutlass::arch::Sm80, cutlass::half_t, true, false, false, 64, 64, 128>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 800
-#if __CUDA_ARCH__ <= 1200
+#if __CUDA_ARCH__ <= 1210
   if (!p.advance_to_block()) {
     return;
   }
@@ -213,7 +213,7 @@ fmha_cutlassB_f16_aligned_64x64_k128_sm80(typename AttentionBackwardKernel<cutla
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_f16_aligned_64x64_k128_sm80` is for sm80-sm100, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_f16_aligned_64x64_k128_sm80` is for sm80-sm121, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
diff --git a/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_f16_aligned_k128_dropout.cu b/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_f16_aligned_k128_dropout.cu
index cc840d3dc52fb..82c31ab786825 100644
--- a/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_f16_aligned_k128_dropout.cu
+++ b/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_f16_aligned_k128_dropout.cu
@@ -14,7 +14,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_f16_aligned_128x64_k128_dropout_sm70(typename AttentionBackwardKernel<cutlass::arch::Sm70, cutlass::half_t, true, true, false, 128, 64, 128>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 700
-#if __CUDA_ARCH__ < 750
+#if __CUDA_ARCH__ <= 740
   if (!p.advance_to_block()) {
     return;
   }
@@ -23,7 +23,7 @@ fmha_cutlassB_f16_aligned_128x64_k128_dropout_sm70(typename AttentionBackwardKer
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_f16_aligned_128x64_k128_dropout_sm70` is for sm70-sm75, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_f16_aligned_128x64_k128_dropout_sm70` is for sm70-sm74, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -33,7 +33,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_f16_aligned_128x64_k128_dropout_sm75(typename AttentionBackwardKernel<cutlass::arch::Sm75, cutlass::half_t, true, true, false, 128, 64, 128>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 750
-#if __CUDA_ARCH__ < 800
+#if __CUDA_ARCH__ <= 790
   if (!p.advance_to_block()) {
     return;
   }
@@ -42,7 +42,7 @@ fmha_cutlassB_f16_aligned_128x64_k128_dropout_sm75(typename AttentionBackwardKer
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_f16_aligned_128x64_k128_dropout_sm75` is for sm75-sm80, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_f16_aligned_128x64_k128_dropout_sm75` is for sm75-sm79, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -52,7 +52,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_f16_aligned_128x128_k128_dropout_sm80(typename AttentionBackwardKernel<cutlass::arch::Sm80, cutlass::half_t, true, true, true, 128, 128, 128>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 800
-#if __CUDA_ARCH__ <= 1200
+#if __CUDA_ARCH__ <= 1210
   if (!p.advance_to_block()) {
     return;
   }
@@ -61,7 +61,7 @@ fmha_cutlassB_f16_aligned_128x128_k128_dropout_sm80(typename AttentionBackwardKe
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_f16_aligned_128x128_k128_dropout_sm80` is for sm80-sm100, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_f16_aligned_128x128_k128_dropout_sm80` is for sm80-sm121, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -71,7 +71,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_f16_aligned_64x64_k128_dropout_sm50(typename AttentionBackwardKernel<cutlass::arch::Sm50, cutlass::half_t, true, true, false, 64, 64, 128>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 500
-#if __CUDA_ARCH__ < 700
+#if __CUDA_ARCH__ <= 690
   if (!p.advance_to_block()) {
     return;
   }
@@ -80,7 +80,7 @@ fmha_cutlassB_f16_aligned_64x64_k128_dropout_sm50(typename AttentionBackwardKern
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_f16_aligned_64x64_k128_dropout_sm50` is for sm50-sm70, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_f16_aligned_64x64_k128_dropout_sm50` is for sm50-sm69, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -90,7 +90,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_f16_aligned_64x64_k128_dropout_sm70(typename AttentionBackwardKernel<cutlass::arch::Sm70, cutlass::half_t, true, true, false, 64, 64, 128>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 700
-#if __CUDA_ARCH__ < 750
+#if __CUDA_ARCH__ <= 740
   if (!p.advance_to_block()) {
     return;
   }
@@ -99,7 +99,7 @@ fmha_cutlassB_f16_aligned_64x64_k128_dropout_sm70(typename AttentionBackwardKern
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_f16_aligned_64x64_k128_dropout_sm70` is for sm70-sm75, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_f16_aligned_64x64_k128_dropout_sm70` is for sm70-sm74, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -109,7 +109,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_f16_aligned_64x64_k128_dropout_sm75(typename AttentionBackwardKernel<cutlass::arch::Sm75, cutlass::half_t, true, true, false, 64, 64, 128>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 750
-#if __CUDA_ARCH__ < 800
+#if __CUDA_ARCH__ <= 790
   if (!p.advance_to_block()) {
     return;
   }
@@ -118,7 +118,7 @@ fmha_cutlassB_f16_aligned_64x64_k128_dropout_sm75(typename AttentionBackwardKern
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_f16_aligned_64x64_k128_dropout_sm75` is for sm75-sm80, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_f16_aligned_64x64_k128_dropout_sm75` is for sm75-sm79, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -128,7 +128,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_f16_aligned_64x64_k128_dropout_sm80(typename AttentionBackwardKernel<cutlass::arch::Sm80, cutlass::half_t, true, true, false, 64, 64, 128>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 800
-#if __CUDA_ARCH__ <= 1200
+#if __CUDA_ARCH__ <= 1210
   if (!p.advance_to_block()) {
     return;
   }
@@ -137,7 +137,7 @@ fmha_cutlassB_f16_aligned_64x64_k128_dropout_sm80(typename AttentionBackwardKern
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_f16_aligned_64x64_k128_dropout_sm80` is for sm80-sm100, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_f16_aligned_64x64_k128_dropout_sm80` is for sm80-sm121, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
diff --git a/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_f16_aligned_k32.cu b/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_f16_aligned_k32.cu
index cf4b7ff192f10..a30b2234e5659 100644
--- a/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_f16_aligned_k32.cu
+++ b/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_f16_aligned_k32.cu
@@ -14,7 +14,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_f16_aligned_64x64_k32_seqaligned_sm70(typename AttentionBackwardKernel<cutlass::arch::Sm70, cutlass::half_t, true, false, false, 64, 64, 32, true>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 700
-#if __CUDA_ARCH__ < 750
+#if __CUDA_ARCH__ <= 740
   if (!p.advance_to_block()) {
     return;
   }
@@ -23,7 +23,7 @@ fmha_cutlassB_f16_aligned_64x64_k32_seqaligned_sm70(typename AttentionBackwardKe
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_f16_aligned_64x64_k32_seqaligned_sm70` is for sm70-sm75, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_f16_aligned_64x64_k32_seqaligned_sm70` is for sm70-sm74, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -33,7 +33,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_f16_aligned_64x64_k32_seqaligned_sm80(typename AttentionBackwardKernel<cutlass::arch::Sm80, cutlass::half_t, true, false, true, 64, 64, 32, true>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 800
-#if __CUDA_ARCH__ <= 1200
+#if __CUDA_ARCH__ <= 1210
   if (!p.advance_to_block()) {
     return;
   }
@@ -42,7 +42,7 @@ fmha_cutlassB_f16_aligned_64x64_k32_seqaligned_sm80(typename AttentionBackwardKe
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_f16_aligned_64x64_k32_seqaligned_sm80` is for sm80-sm100, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_f16_aligned_64x64_k32_seqaligned_sm80` is for sm80-sm121, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -52,7 +52,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_f16_aligned_64x64_k32_sm50(typename AttentionBackwardKernel<cutlass::arch::Sm50, cutlass::half_t, true, false, false, 64, 64, 32>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 500
-#if __CUDA_ARCH__ < 700
+#if __CUDA_ARCH__ <= 690
   if (!p.advance_to_block()) {
     return;
   }
@@ -61,7 +61,7 @@ fmha_cutlassB_f16_aligned_64x64_k32_sm50(typename AttentionBackwardKernel<cutlas
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_f16_aligned_64x64_k32_sm50` is for sm50-sm70, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_f16_aligned_64x64_k32_sm50` is for sm50-sm69, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -71,7 +71,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_f16_aligned_64x64_k32_sm70(typename AttentionBackwardKernel<cutlass::arch::Sm70, cutlass::half_t, true, false, false, 64, 64, 32>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 700
-#if __CUDA_ARCH__ < 750
+#if __CUDA_ARCH__ <= 740
   if (!p.advance_to_block()) {
     return;
   }
@@ -80,7 +80,7 @@ fmha_cutlassB_f16_aligned_64x64_k32_sm70(typename AttentionBackwardKernel<cutlas
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_f16_aligned_64x64_k32_sm70` is for sm70-sm75, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_f16_aligned_64x64_k32_sm70` is for sm70-sm74, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -90,7 +90,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_f16_aligned_64x64_k32_sm75(typename AttentionBackwardKernel<cutlass::arch::Sm75, cutlass::half_t, true, false, false, 64, 64, 32>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 750
-#if __CUDA_ARCH__ < 800
+#if __CUDA_ARCH__ <= 790
   if (!p.advance_to_block()) {
     return;
   }
@@ -99,7 +99,7 @@ fmha_cutlassB_f16_aligned_64x64_k32_sm75(typename AttentionBackwardKernel<cutlas
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_f16_aligned_64x64_k32_sm75` is for sm75-sm80, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_f16_aligned_64x64_k32_sm75` is for sm75-sm79, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -109,7 +109,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_f16_aligned_64x64_k32_sm80(typename AttentionBackwardKernel<cutlass::arch::Sm80, cutlass::half_t, true, false, true, 64, 64, 32>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 800
-#if __CUDA_ARCH__ <= 1200
+#if __CUDA_ARCH__ <= 1210
   if (!p.advance_to_block()) {
     return;
   }
@@ -118,7 +118,7 @@ fmha_cutlassB_f16_aligned_64x64_k32_sm80(typename AttentionBackwardKernel<cutlas
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_f16_aligned_64x64_k32_sm80` is for sm80-sm100, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_f16_aligned_64x64_k32_sm80` is for sm80-sm121, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
diff --git a/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_f16_aligned_k32_dropout.cu b/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_f16_aligned_k32_dropout.cu
index 47c25b413f983..6323374a988e3 100644
--- a/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_f16_aligned_k32_dropout.cu
+++ b/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_f16_aligned_k32_dropout.cu
@@ -14,7 +14,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_f16_aligned_64x64_k32_dropout_sm50(typename AttentionBackwardKernel<cutlass::arch::Sm50, cutlass::half_t, true, true, false, 64, 64, 32>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 500
-#if __CUDA_ARCH__ < 700
+#if __CUDA_ARCH__ <= 690
   if (!p.advance_to_block()) {
     return;
   }
@@ -23,7 +23,7 @@ fmha_cutlassB_f16_aligned_64x64_k32_dropout_sm50(typename AttentionBackwardKerne
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_f16_aligned_64x64_k32_dropout_sm50` is for sm50-sm70, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_f16_aligned_64x64_k32_dropout_sm50` is for sm50-sm69, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -33,7 +33,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_f16_aligned_64x64_k32_dropout_sm70(typename AttentionBackwardKernel<cutlass::arch::Sm70, cutlass::half_t, true, true, false, 64, 64, 32>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 700
-#if __CUDA_ARCH__ < 750
+#if __CUDA_ARCH__ <= 740
   if (!p.advance_to_block()) {
     return;
   }
@@ -42,7 +42,7 @@ fmha_cutlassB_f16_aligned_64x64_k32_dropout_sm70(typename AttentionBackwardKerne
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_f16_aligned_64x64_k32_dropout_sm70` is for sm70-sm75, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_f16_aligned_64x64_k32_dropout_sm70` is for sm70-sm74, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -52,7 +52,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_f16_aligned_64x64_k32_dropout_sm75(typename AttentionBackwardKernel<cutlass::arch::Sm75, cutlass::half_t, true, true, false, 64, 64, 32>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 750
-#if __CUDA_ARCH__ < 800
+#if __CUDA_ARCH__ <= 790
   if (!p.advance_to_block()) {
     return;
   }
@@ -61,7 +61,7 @@ fmha_cutlassB_f16_aligned_64x64_k32_dropout_sm75(typename AttentionBackwardKerne
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_f16_aligned_64x64_k32_dropout_sm75` is for sm75-sm80, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_f16_aligned_64x64_k32_dropout_sm75` is for sm75-sm79, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -71,7 +71,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_f16_aligned_64x64_k32_dropout_sm80(typename AttentionBackwardKernel<cutlass::arch::Sm80, cutlass::half_t, true, true, true, 64, 64, 32>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 800
-#if __CUDA_ARCH__ <= 1200
+#if __CUDA_ARCH__ <= 1210
   if (!p.advance_to_block()) {
     return;
   }
@@ -80,7 +80,7 @@ fmha_cutlassB_f16_aligned_64x64_k32_dropout_sm80(typename AttentionBackwardKerne
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_f16_aligned_64x64_k32_dropout_sm80` is for sm80-sm100, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_f16_aligned_64x64_k32_dropout_sm80` is for sm80-sm121, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
diff --git a/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_f16_aligned_k64.cu b/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_f16_aligned_k64.cu
index 745b4b802a6bc..5e3fdbbc45ac0 100644
--- a/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_f16_aligned_k64.cu
+++ b/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_f16_aligned_k64.cu
@@ -14,7 +14,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_f16_aligned_64x64_k64_seqaligned_sm70(typename AttentionBackwardKernel<cutlass::arch::Sm70, cutlass::half_t, true, false, false, 64, 64, 64, true>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 700
-#if __CUDA_ARCH__ < 750
+#if __CUDA_ARCH__ <= 740
   if (!p.advance_to_block()) {
     return;
   }
@@ -23,7 +23,7 @@ fmha_cutlassB_f16_aligned_64x64_k64_seqaligned_sm70(typename AttentionBackwardKe
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_f16_aligned_64x64_k64_seqaligned_sm70` is for sm70-sm75, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_f16_aligned_64x64_k64_seqaligned_sm70` is for sm70-sm74, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -33,7 +33,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_f16_aligned_64x64_k64_seqaligned_sm80(typename AttentionBackwardKernel<cutlass::arch::Sm80, cutlass::half_t, true, false, true, 64, 64, 64, true>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 800
-#if __CUDA_ARCH__ <= 1200
+#if __CUDA_ARCH__ <= 1210
   if (!p.advance_to_block()) {
     return;
   }
@@ -42,7 +42,7 @@ fmha_cutlassB_f16_aligned_64x64_k64_seqaligned_sm80(typename AttentionBackwardKe
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_f16_aligned_64x64_k64_seqaligned_sm80` is for sm80-sm100, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_f16_aligned_64x64_k64_seqaligned_sm80` is for sm80-sm121, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -52,7 +52,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_f16_aligned_64x64_k64_sm50(typename AttentionBackwardKernel<cutlass::arch::Sm50, cutlass::half_t, true, false, false, 64, 64, 64>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 500
-#if __CUDA_ARCH__ < 700
+#if __CUDA_ARCH__ <= 690
   if (!p.advance_to_block()) {
     return;
   }
@@ -61,7 +61,7 @@ fmha_cutlassB_f16_aligned_64x64_k64_sm50(typename AttentionBackwardKernel<cutlas
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_f16_aligned_64x64_k64_sm50` is for sm50-sm70, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_f16_aligned_64x64_k64_sm50` is for sm50-sm69, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -71,7 +71,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_f16_aligned_64x64_k64_sm70(typename AttentionBackwardKernel<cutlass::arch::Sm70, cutlass::half_t, true, false, false, 64, 64, 64>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 700
-#if __CUDA_ARCH__ < 750
+#if __CUDA_ARCH__ <= 740
   if (!p.advance_to_block()) {
     return;
   }
@@ -80,7 +80,7 @@ fmha_cutlassB_f16_aligned_64x64_k64_sm70(typename AttentionBackwardKernel<cutlas
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_f16_aligned_64x64_k64_sm70` is for sm70-sm75, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_f16_aligned_64x64_k64_sm70` is for sm70-sm74, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -90,7 +90,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_f16_aligned_64x64_k64_sm75(typename AttentionBackwardKernel<cutlass::arch::Sm75, cutlass::half_t, true, false, false, 64, 64, 64>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 750
-#if __CUDA_ARCH__ < 800
+#if __CUDA_ARCH__ <= 790
   if (!p.advance_to_block()) {
     return;
   }
@@ -99,7 +99,7 @@ fmha_cutlassB_f16_aligned_64x64_k64_sm75(typename AttentionBackwardKernel<cutlas
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_f16_aligned_64x64_k64_sm75` is for sm75-sm80, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_f16_aligned_64x64_k64_sm75` is for sm75-sm79, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -109,7 +109,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_f16_aligned_64x64_k64_sm80(typename AttentionBackwardKernel<cutlass::arch::Sm80, cutlass::half_t, true, false, true, 64, 64, 64>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 800
-#if __CUDA_ARCH__ <= 1200
+#if __CUDA_ARCH__ <= 1210
   if (!p.advance_to_block()) {
     return;
   }
@@ -118,7 +118,7 @@ fmha_cutlassB_f16_aligned_64x64_k64_sm80(typename AttentionBackwardKernel<cutlas
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_f16_aligned_64x64_k64_sm80` is for sm80-sm100, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_f16_aligned_64x64_k64_sm80` is for sm80-sm121, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
diff --git a/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_f16_aligned_k64_dropout.cu b/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_f16_aligned_k64_dropout.cu
index c143627fa3a1e..863dab8a6aee2 100644
--- a/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_f16_aligned_k64_dropout.cu
+++ b/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_f16_aligned_k64_dropout.cu
@@ -14,7 +14,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_f16_aligned_64x64_k64_dropout_sm50(typename AttentionBackwardKernel<cutlass::arch::Sm50, cutlass::half_t, true, true, false, 64, 64, 64>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 500
-#if __CUDA_ARCH__ < 700
+#if __CUDA_ARCH__ <= 690
   if (!p.advance_to_block()) {
     return;
   }
@@ -23,7 +23,7 @@ fmha_cutlassB_f16_aligned_64x64_k64_dropout_sm50(typename AttentionBackwardKerne
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_f16_aligned_64x64_k64_dropout_sm50` is for sm50-sm70, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_f16_aligned_64x64_k64_dropout_sm50` is for sm50-sm69, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -33,7 +33,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_f16_aligned_64x64_k64_dropout_sm70(typename AttentionBackwardKernel<cutlass::arch::Sm70, cutlass::half_t, true, true, false, 64, 64, 64>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 700
-#if __CUDA_ARCH__ < 750
+#if __CUDA_ARCH__ <= 740
   if (!p.advance_to_block()) {
     return;
   }
@@ -42,7 +42,7 @@ fmha_cutlassB_f16_aligned_64x64_k64_dropout_sm70(typename AttentionBackwardKerne
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_f16_aligned_64x64_k64_dropout_sm70` is for sm70-sm75, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_f16_aligned_64x64_k64_dropout_sm70` is for sm70-sm74, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -52,7 +52,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_f16_aligned_64x64_k64_dropout_sm75(typename AttentionBackwardKernel<cutlass::arch::Sm75, cutlass::half_t, true, true, false, 64, 64, 64>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 750
-#if __CUDA_ARCH__ < 800
+#if __CUDA_ARCH__ <= 790
   if (!p.advance_to_block()) {
     return;
   }
@@ -61,7 +61,7 @@ fmha_cutlassB_f16_aligned_64x64_k64_dropout_sm75(typename AttentionBackwardKerne
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_f16_aligned_64x64_k64_dropout_sm75` is for sm75-sm80, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_f16_aligned_64x64_k64_dropout_sm75` is for sm75-sm79, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -71,7 +71,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_f16_aligned_64x64_k64_dropout_sm80(typename AttentionBackwardKernel<cutlass::arch::Sm80, cutlass::half_t, true, true, true, 64, 64, 64>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 800
-#if __CUDA_ARCH__ <= 1200
+#if __CUDA_ARCH__ <= 1210
   if (!p.advance_to_block()) {
     return;
   }
@@ -80,7 +80,7 @@ fmha_cutlassB_f16_aligned_64x64_k64_dropout_sm80(typename AttentionBackwardKerne
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_f16_aligned_64x64_k64_dropout_sm80` is for sm80-sm100, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_f16_aligned_64x64_k64_dropout_sm80` is for sm80-sm121, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
diff --git a/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_f16_aligned_k65536.cu b/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_f16_aligned_k65536.cu
index 33fbd2dd05888..da087d15a07f2 100644
--- a/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_f16_aligned_k65536.cu
+++ b/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_f16_aligned_k65536.cu
@@ -14,7 +14,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_f16_aligned_128x64_k65536_sm70(typename AttentionBackwardKernel<cutlass::arch::Sm70, cutlass::half_t, true, false, false, 128, 64, 65536>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 700
-#if __CUDA_ARCH__ < 750
+#if __CUDA_ARCH__ <= 740
   if (!p.advance_to_block()) {
     return;
   }
@@ -23,7 +23,7 @@ fmha_cutlassB_f16_aligned_128x64_k65536_sm70(typename AttentionBackwardKernel<cu
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_f16_aligned_128x64_k65536_sm70` is for sm70-sm75, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_f16_aligned_128x64_k65536_sm70` is for sm70-sm74, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -33,7 +33,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_f16_aligned_128x64_k65536_sm75(typename AttentionBackwardKernel<cutlass::arch::Sm75, cutlass::half_t, true, false, false, 128, 64, 65536>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 750
-#if __CUDA_ARCH__ < 800
+#if __CUDA_ARCH__ <= 790
   if (!p.advance_to_block()) {
     return;
   }
@@ -42,7 +42,7 @@ fmha_cutlassB_f16_aligned_128x64_k65536_sm75(typename AttentionBackwardKernel<cu
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_f16_aligned_128x64_k65536_sm75` is for sm75-sm80, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_f16_aligned_128x64_k65536_sm75` is for sm75-sm79, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -52,7 +52,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_f16_aligned_128x64_k65536_sm80(typename AttentionBackwardKernel<cutlass::arch::Sm80, cutlass::half_t, true, false, false, 128, 64, 65536>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 800
-#if __CUDA_ARCH__ <= 1200
+#if __CUDA_ARCH__ <= 1210
   if (!p.advance_to_block()) {
     return;
   }
@@ -61,7 +61,7 @@ fmha_cutlassB_f16_aligned_128x64_k65536_sm80(typename AttentionBackwardKernel<cu
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_f16_aligned_128x64_k65536_sm80` is for sm80-sm100, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_f16_aligned_128x64_k65536_sm80` is for sm80-sm121, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -71,7 +71,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_f16_aligned_64x64_k65536_sm50(typename AttentionBackwardKernel<cutlass::arch::Sm50, cutlass::half_t, true, false, false, 64, 64, 65536>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 500
-#if __CUDA_ARCH__ < 700
+#if __CUDA_ARCH__ <= 690
   if (!p.advance_to_block()) {
     return;
   }
@@ -80,7 +80,7 @@ fmha_cutlassB_f16_aligned_64x64_k65536_sm50(typename AttentionBackwardKernel<cut
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_f16_aligned_64x64_k65536_sm50` is for sm50-sm70, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_f16_aligned_64x64_k65536_sm50` is for sm50-sm69, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -90,7 +90,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_f16_aligned_64x64_k65536_sm70(typename AttentionBackwardKernel<cutlass::arch::Sm70, cutlass::half_t, true, false, false, 64, 64, 65536>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 700
-#if __CUDA_ARCH__ < 750
+#if __CUDA_ARCH__ <= 740
   if (!p.advance_to_block()) {
     return;
   }
@@ -99,7 +99,7 @@ fmha_cutlassB_f16_aligned_64x64_k65536_sm70(typename AttentionBackwardKernel<cut
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_f16_aligned_64x64_k65536_sm70` is for sm70-sm75, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_f16_aligned_64x64_k65536_sm70` is for sm70-sm74, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -109,7 +109,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_f16_aligned_64x64_k65536_sm75(typename AttentionBackwardKernel<cutlass::arch::Sm75, cutlass::half_t, true, false, false, 64, 64, 65536>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 750
-#if __CUDA_ARCH__ < 800
+#if __CUDA_ARCH__ <= 790
   if (!p.advance_to_block()) {
     return;
   }
@@ -118,7 +118,7 @@ fmha_cutlassB_f16_aligned_64x64_k65536_sm75(typename AttentionBackwardKernel<cut
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_f16_aligned_64x64_k65536_sm75` is for sm75-sm80, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_f16_aligned_64x64_k65536_sm75` is for sm75-sm79, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -128,7 +128,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_f16_aligned_64x64_k65536_sm80(typename AttentionBackwardKernel<cutlass::arch::Sm80, cutlass::half_t, true, false, false, 64, 64, 65536>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 800
-#if __CUDA_ARCH__ <= 1200
+#if __CUDA_ARCH__ <= 1210
   if (!p.advance_to_block()) {
     return;
   }
@@ -137,7 +137,7 @@ fmha_cutlassB_f16_aligned_64x64_k65536_sm80(typename AttentionBackwardKernel<cut
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_f16_aligned_64x64_k65536_sm80` is for sm80-sm100, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_f16_aligned_64x64_k65536_sm80` is for sm80-sm121, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
diff --git a/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_f16_aligned_k65536_dropout.cu b/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_f16_aligned_k65536_dropout.cu
index 3c307c221d802..c32d576a8a790 100644
--- a/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_f16_aligned_k65536_dropout.cu
+++ b/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_f16_aligned_k65536_dropout.cu
@@ -14,7 +14,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_f16_aligned_128x64_k65536_dropout_sm70(typename AttentionBackwardKernel<cutlass::arch::Sm70, cutlass::half_t, true, true, false, 128, 64, 65536>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 700
-#if __CUDA_ARCH__ < 750
+#if __CUDA_ARCH__ <= 740
   if (!p.advance_to_block()) {
     return;
   }
@@ -23,7 +23,7 @@ fmha_cutlassB_f16_aligned_128x64_k65536_dropout_sm70(typename AttentionBackwardK
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_f16_aligned_128x64_k65536_dropout_sm70` is for sm70-sm75, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_f16_aligned_128x64_k65536_dropout_sm70` is for sm70-sm74, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -33,7 +33,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_f16_aligned_128x64_k65536_dropout_sm75(typename AttentionBackwardKernel<cutlass::arch::Sm75, cutlass::half_t, true, true, false, 128, 64, 65536>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 750
-#if __CUDA_ARCH__ < 800
+#if __CUDA_ARCH__ <= 790
   if (!p.advance_to_block()) {
     return;
   }
@@ -42,7 +42,7 @@ fmha_cutlassB_f16_aligned_128x64_k65536_dropout_sm75(typename AttentionBackwardK
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_f16_aligned_128x64_k65536_dropout_sm75` is for sm75-sm80, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_f16_aligned_128x64_k65536_dropout_sm75` is for sm75-sm79, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -52,7 +52,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_f16_aligned_128x64_k65536_dropout_sm80(typename AttentionBackwardKernel<cutlass::arch::Sm80, cutlass::half_t, true, true, false, 128, 64, 65536>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 800
-#if __CUDA_ARCH__ <= 1200
+#if __CUDA_ARCH__ <= 1210
   if (!p.advance_to_block()) {
     return;
   }
@@ -61,7 +61,7 @@ fmha_cutlassB_f16_aligned_128x64_k65536_dropout_sm80(typename AttentionBackwardK
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_f16_aligned_128x64_k65536_dropout_sm80` is for sm80-sm100, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_f16_aligned_128x64_k65536_dropout_sm80` is for sm80-sm121, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -71,7 +71,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_f16_aligned_64x64_k65536_dropout_sm50(typename AttentionBackwardKernel<cutlass::arch::Sm50, cutlass::half_t, true, true, false, 64, 64, 65536>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 500
-#if __CUDA_ARCH__ < 700
+#if __CUDA_ARCH__ <= 690
   if (!p.advance_to_block()) {
     return;
   }
@@ -80,7 +80,7 @@ fmha_cutlassB_f16_aligned_64x64_k65536_dropout_sm50(typename AttentionBackwardKe
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_f16_aligned_64x64_k65536_dropout_sm50` is for sm50-sm70, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_f16_aligned_64x64_k65536_dropout_sm50` is for sm50-sm69, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -90,7 +90,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_f16_aligned_64x64_k65536_dropout_sm70(typename AttentionBackwardKernel<cutlass::arch::Sm70, cutlass::half_t, true, true, false, 64, 64, 65536>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 700
-#if __CUDA_ARCH__ < 750
+#if __CUDA_ARCH__ <= 740
   if (!p.advance_to_block()) {
     return;
   }
@@ -99,7 +99,7 @@ fmha_cutlassB_f16_aligned_64x64_k65536_dropout_sm70(typename AttentionBackwardKe
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_f16_aligned_64x64_k65536_dropout_sm70` is for sm70-sm75, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_f16_aligned_64x64_k65536_dropout_sm70` is for sm70-sm74, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -109,7 +109,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_f16_aligned_64x64_k65536_dropout_sm75(typename AttentionBackwardKernel<cutlass::arch::Sm75, cutlass::half_t, true, true, false, 64, 64, 65536>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 750
-#if __CUDA_ARCH__ < 800
+#if __CUDA_ARCH__ <= 790
   if (!p.advance_to_block()) {
     return;
   }
@@ -118,7 +118,7 @@ fmha_cutlassB_f16_aligned_64x64_k65536_dropout_sm75(typename AttentionBackwardKe
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_f16_aligned_64x64_k65536_dropout_sm75` is for sm75-sm80, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_f16_aligned_64x64_k65536_dropout_sm75` is for sm75-sm79, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -128,7 +128,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_f16_aligned_64x64_k65536_dropout_sm80(typename AttentionBackwardKernel<cutlass::arch::Sm80, cutlass::half_t, true, true, false, 64, 64, 65536>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 800
-#if __CUDA_ARCH__ <= 1200
+#if __CUDA_ARCH__ <= 1210
   if (!p.advance_to_block()) {
     return;
   }
@@ -137,7 +137,7 @@ fmha_cutlassB_f16_aligned_64x64_k65536_dropout_sm80(typename AttentionBackwardKe
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_f16_aligned_64x64_k65536_dropout_sm80` is for sm80-sm100, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_f16_aligned_64x64_k65536_dropout_sm80` is for sm80-sm121, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
diff --git a/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_f16_aligned_k96.cu b/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_f16_aligned_k96.cu
index c0cc22ae370fb..371609b89e72c 100644
--- a/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_f16_aligned_k96.cu
+++ b/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_f16_aligned_k96.cu
@@ -14,7 +14,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_f16_aligned_128x64_k96_sm80(typename AttentionBackwardKernel<cutlass::arch::Sm80, cutlass::half_t, true, false, true, 128, 64, 96>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 800
-#if __CUDA_ARCH__ <= 1200
+#if __CUDA_ARCH__ <= 1210
   if (!p.advance_to_block()) {
     return;
   }
@@ -23,7 +23,7 @@ fmha_cutlassB_f16_aligned_128x64_k96_sm80(typename AttentionBackwardKernel<cutla
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_f16_aligned_128x64_k96_sm80` is for sm80-sm100, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_f16_aligned_128x64_k96_sm80` is for sm80-sm121, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
diff --git a/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_f16_notaligned_k128.cu b/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_f16_notaligned_k128.cu
index 4e513594d87e7..cec7a64e5c57f 100644
--- a/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_f16_notaligned_k128.cu
+++ b/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_f16_notaligned_k128.cu
@@ -14,7 +14,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_f16_notaligned_128x64_k128_sm70(typename AttentionBackwardKernel<cutlass::arch::Sm70, cutlass::half_t, false, false, false, 128, 64, 128>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 700
-#if __CUDA_ARCH__ < 750
+#if __CUDA_ARCH__ <= 740
   if (!p.advance_to_block()) {
     return;
   }
@@ -23,7 +23,7 @@ fmha_cutlassB_f16_notaligned_128x64_k128_sm70(typename AttentionBackwardKernel<c
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_f16_notaligned_128x64_k128_sm70` is for sm70-sm75, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_f16_notaligned_128x64_k128_sm70` is for sm70-sm74, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -33,7 +33,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_f16_notaligned_128x64_k128_sm75(typename AttentionBackwardKernel<cutlass::arch::Sm75, cutlass::half_t, false, false, false, 128, 64, 128>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 750
-#if __CUDA_ARCH__ < 800
+#if __CUDA_ARCH__ <= 790
   if (!p.advance_to_block()) {
     return;
   }
@@ -42,7 +42,7 @@ fmha_cutlassB_f16_notaligned_128x64_k128_sm75(typename AttentionBackwardKernel<c
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_f16_notaligned_128x64_k128_sm75` is for sm75-sm80, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_f16_notaligned_128x64_k128_sm75` is for sm75-sm79, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -52,7 +52,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_f16_notaligned_64x64_k128_sm50(typename AttentionBackwardKernel<cutlass::arch::Sm50, cutlass::half_t, false, false, false, 64, 64, 128>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 500
-#if __CUDA_ARCH__ < 700
+#if __CUDA_ARCH__ <= 690
   if (!p.advance_to_block()) {
     return;
   }
@@ -61,7 +61,7 @@ fmha_cutlassB_f16_notaligned_64x64_k128_sm50(typename AttentionBackwardKernel<cu
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_f16_notaligned_64x64_k128_sm50` is for sm50-sm70, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_f16_notaligned_64x64_k128_sm50` is for sm50-sm69, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -71,7 +71,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_f16_notaligned_64x64_k128_sm70(typename AttentionBackwardKernel<cutlass::arch::Sm70, cutlass::half_t, false, false, false, 64, 64, 128>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 700
-#if __CUDA_ARCH__ < 750
+#if __CUDA_ARCH__ <= 740
   if (!p.advance_to_block()) {
     return;
   }
@@ -80,7 +80,7 @@ fmha_cutlassB_f16_notaligned_64x64_k128_sm70(typename AttentionBackwardKernel<cu
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_f16_notaligned_64x64_k128_sm70` is for sm70-sm75, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_f16_notaligned_64x64_k128_sm70` is for sm70-sm74, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -90,7 +90,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_f16_notaligned_64x64_k128_sm75(typename AttentionBackwardKernel<cutlass::arch::Sm75, cutlass::half_t, false, false, false, 64, 64, 128>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 750
-#if __CUDA_ARCH__ < 800
+#if __CUDA_ARCH__ <= 790
   if (!p.advance_to_block()) {
     return;
   }
@@ -99,7 +99,7 @@ fmha_cutlassB_f16_notaligned_64x64_k128_sm75(typename AttentionBackwardKernel<cu
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_f16_notaligned_64x64_k128_sm75` is for sm75-sm80, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_f16_notaligned_64x64_k128_sm75` is for sm75-sm79, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
diff --git a/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_f16_notaligned_k128_dropout.cu b/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_f16_notaligned_k128_dropout.cu
index f82db06a789a6..cecc99f707c8b 100644
--- a/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_f16_notaligned_k128_dropout.cu
+++ b/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_f16_notaligned_k128_dropout.cu
@@ -14,7 +14,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_f16_notaligned_128x64_k128_dropout_sm70(typename AttentionBackwardKernel<cutlass::arch::Sm70, cutlass::half_t, false, true, false, 128, 64, 128>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 700
-#if __CUDA_ARCH__ < 750
+#if __CUDA_ARCH__ <= 740
   if (!p.advance_to_block()) {
     return;
   }
@@ -23,7 +23,7 @@ fmha_cutlassB_f16_notaligned_128x64_k128_dropout_sm70(typename AttentionBackward
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_f16_notaligned_128x64_k128_dropout_sm70` is for sm70-sm75, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_f16_notaligned_128x64_k128_dropout_sm70` is for sm70-sm74, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -33,7 +33,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_f16_notaligned_128x64_k128_dropout_sm75(typename AttentionBackwardKernel<cutlass::arch::Sm75, cutlass::half_t, false, true, false, 128, 64, 128>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 750
-#if __CUDA_ARCH__ < 800
+#if __CUDA_ARCH__ <= 790
   if (!p.advance_to_block()) {
     return;
   }
@@ -42,7 +42,7 @@ fmha_cutlassB_f16_notaligned_128x64_k128_dropout_sm75(typename AttentionBackward
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_f16_notaligned_128x64_k128_dropout_sm75` is for sm75-sm80, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_f16_notaligned_128x64_k128_dropout_sm75` is for sm75-sm79, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -52,7 +52,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_f16_notaligned_64x64_k128_dropout_sm50(typename AttentionBackwardKernel<cutlass::arch::Sm50, cutlass::half_t, false, true, false, 64, 64, 128>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 500
-#if __CUDA_ARCH__ < 700
+#if __CUDA_ARCH__ <= 690
   if (!p.advance_to_block()) {
     return;
   }
@@ -61,7 +61,7 @@ fmha_cutlassB_f16_notaligned_64x64_k128_dropout_sm50(typename AttentionBackwardK
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_f16_notaligned_64x64_k128_dropout_sm50` is for sm50-sm70, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_f16_notaligned_64x64_k128_dropout_sm50` is for sm50-sm69, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -71,7 +71,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_f16_notaligned_64x64_k128_dropout_sm70(typename AttentionBackwardKernel<cutlass::arch::Sm70, cutlass::half_t, false, true, false, 64, 64, 128>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 700
-#if __CUDA_ARCH__ < 750
+#if __CUDA_ARCH__ <= 740
   if (!p.advance_to_block()) {
     return;
   }
@@ -80,7 +80,7 @@ fmha_cutlassB_f16_notaligned_64x64_k128_dropout_sm70(typename AttentionBackwardK
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_f16_notaligned_64x64_k128_dropout_sm70` is for sm70-sm75, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_f16_notaligned_64x64_k128_dropout_sm70` is for sm70-sm74, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -90,7 +90,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_f16_notaligned_64x64_k128_dropout_sm75(typename AttentionBackwardKernel<cutlass::arch::Sm75, cutlass::half_t, false, true, false, 64, 64, 128>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 750
-#if __CUDA_ARCH__ < 800
+#if __CUDA_ARCH__ <= 790
   if (!p.advance_to_block()) {
     return;
   }
@@ -99,7 +99,7 @@ fmha_cutlassB_f16_notaligned_64x64_k128_dropout_sm75(typename AttentionBackwardK
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_f16_notaligned_64x64_k128_dropout_sm75` is for sm75-sm80, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_f16_notaligned_64x64_k128_dropout_sm75` is for sm75-sm79, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
diff --git a/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_f16_notaligned_k32.cu b/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_f16_notaligned_k32.cu
index 4b51b045fa161..6b31b66035a3e 100644
--- a/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_f16_notaligned_k32.cu
+++ b/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_f16_notaligned_k32.cu
@@ -14,7 +14,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_f16_notaligned_64x64_k32_sm50(typename AttentionBackwardKernel<cutlass::arch::Sm50, cutlass::half_t, false, false, false, 64, 64, 32>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 500
-#if __CUDA_ARCH__ < 700
+#if __CUDA_ARCH__ <= 690
   if (!p.advance_to_block()) {
     return;
   }
@@ -23,7 +23,7 @@ fmha_cutlassB_f16_notaligned_64x64_k32_sm50(typename AttentionBackwardKernel<cut
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_f16_notaligned_64x64_k32_sm50` is for sm50-sm70, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_f16_notaligned_64x64_k32_sm50` is for sm50-sm69, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -33,7 +33,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_f16_notaligned_64x64_k32_sm70(typename AttentionBackwardKernel<cutlass::arch::Sm70, cutlass::half_t, false, false, false, 64, 64, 32>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 700
-#if __CUDA_ARCH__ < 750
+#if __CUDA_ARCH__ <= 740
   if (!p.advance_to_block()) {
     return;
   }
@@ -42,7 +42,7 @@ fmha_cutlassB_f16_notaligned_64x64_k32_sm70(typename AttentionBackwardKernel<cut
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_f16_notaligned_64x64_k32_sm70` is for sm70-sm75, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_f16_notaligned_64x64_k32_sm70` is for sm70-sm74, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -52,7 +52,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_f16_notaligned_64x64_k32_sm75(typename AttentionBackwardKernel<cutlass::arch::Sm75, cutlass::half_t, false, false, false, 64, 64, 32>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 750
-#if __CUDA_ARCH__ < 800
+#if __CUDA_ARCH__ <= 790
   if (!p.advance_to_block()) {
     return;
   }
@@ -61,7 +61,7 @@ fmha_cutlassB_f16_notaligned_64x64_k32_sm75(typename AttentionBackwardKernel<cut
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_f16_notaligned_64x64_k32_sm75` is for sm75-sm80, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_f16_notaligned_64x64_k32_sm75` is for sm75-sm79, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
diff --git a/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_f16_notaligned_k32_dropout.cu b/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_f16_notaligned_k32_dropout.cu
index 1cff4d6ee1023..382936de41cfd 100644
--- a/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_f16_notaligned_k32_dropout.cu
+++ b/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_f16_notaligned_k32_dropout.cu
@@ -14,7 +14,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_f16_notaligned_64x64_k32_dropout_sm50(typename AttentionBackwardKernel<cutlass::arch::Sm50, cutlass::half_t, false, true, false, 64, 64, 32>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 500
-#if __CUDA_ARCH__ < 700
+#if __CUDA_ARCH__ <= 690
   if (!p.advance_to_block()) {
     return;
   }
@@ -23,7 +23,7 @@ fmha_cutlassB_f16_notaligned_64x64_k32_dropout_sm50(typename AttentionBackwardKe
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_f16_notaligned_64x64_k32_dropout_sm50` is for sm50-sm70, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_f16_notaligned_64x64_k32_dropout_sm50` is for sm50-sm69, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -33,7 +33,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_f16_notaligned_64x64_k32_dropout_sm70(typename AttentionBackwardKernel<cutlass::arch::Sm70, cutlass::half_t, false, true, false, 64, 64, 32>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 700
-#if __CUDA_ARCH__ < 750
+#if __CUDA_ARCH__ <= 740
   if (!p.advance_to_block()) {
     return;
   }
@@ -42,7 +42,7 @@ fmha_cutlassB_f16_notaligned_64x64_k32_dropout_sm70(typename AttentionBackwardKe
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_f16_notaligned_64x64_k32_dropout_sm70` is for sm70-sm75, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_f16_notaligned_64x64_k32_dropout_sm70` is for sm70-sm74, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -52,7 +52,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_f16_notaligned_64x64_k32_dropout_sm75(typename AttentionBackwardKernel<cutlass::arch::Sm75, cutlass::half_t, false, true, false, 64, 64, 32>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 750
-#if __CUDA_ARCH__ < 800
+#if __CUDA_ARCH__ <= 790
   if (!p.advance_to_block()) {
     return;
   }
@@ -61,7 +61,7 @@ fmha_cutlassB_f16_notaligned_64x64_k32_dropout_sm75(typename AttentionBackwardKe
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_f16_notaligned_64x64_k32_dropout_sm75` is for sm75-sm80, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_f16_notaligned_64x64_k32_dropout_sm75` is for sm75-sm79, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
diff --git a/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_f16_notaligned_k64.cu b/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_f16_notaligned_k64.cu
index 45dd860b6d747..27d3992e2f60d 100644
--- a/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_f16_notaligned_k64.cu
+++ b/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_f16_notaligned_k64.cu
@@ -14,7 +14,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_f16_notaligned_64x64_k64_sm50(typename AttentionBackwardKernel<cutlass::arch::Sm50, cutlass::half_t, false, false, false, 64, 64, 64>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 500
-#if __CUDA_ARCH__ < 700
+#if __CUDA_ARCH__ <= 690
   if (!p.advance_to_block()) {
     return;
   }
@@ -23,7 +23,7 @@ fmha_cutlassB_f16_notaligned_64x64_k64_sm50(typename AttentionBackwardKernel<cut
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_f16_notaligned_64x64_k64_sm50` is for sm50-sm70, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_f16_notaligned_64x64_k64_sm50` is for sm50-sm69, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -33,7 +33,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_f16_notaligned_64x64_k64_sm70(typename AttentionBackwardKernel<cutlass::arch::Sm70, cutlass::half_t, false, false, false, 64, 64, 64>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 700
-#if __CUDA_ARCH__ < 750
+#if __CUDA_ARCH__ <= 740
   if (!p.advance_to_block()) {
     return;
   }
@@ -42,7 +42,7 @@ fmha_cutlassB_f16_notaligned_64x64_k64_sm70(typename AttentionBackwardKernel<cut
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_f16_notaligned_64x64_k64_sm70` is for sm70-sm75, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_f16_notaligned_64x64_k64_sm70` is for sm70-sm74, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -52,7 +52,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_f16_notaligned_64x64_k64_sm75(typename AttentionBackwardKernel<cutlass::arch::Sm75, cutlass::half_t, false, false, false, 64, 64, 64>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 750
-#if __CUDA_ARCH__ < 800
+#if __CUDA_ARCH__ <= 790
   if (!p.advance_to_block()) {
     return;
   }
@@ -61,7 +61,7 @@ fmha_cutlassB_f16_notaligned_64x64_k64_sm75(typename AttentionBackwardKernel<cut
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_f16_notaligned_64x64_k64_sm75` is for sm75-sm80, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_f16_notaligned_64x64_k64_sm75` is for sm75-sm79, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
diff --git a/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_f16_notaligned_k64_dropout.cu b/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_f16_notaligned_k64_dropout.cu
index 40a10272e1e65..67dd9b2f8b0ec 100644
--- a/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_f16_notaligned_k64_dropout.cu
+++ b/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_f16_notaligned_k64_dropout.cu
@@ -14,7 +14,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_f16_notaligned_64x64_k64_dropout_sm50(typename AttentionBackwardKernel<cutlass::arch::Sm50, cutlass::half_t, false, true, false, 64, 64, 64>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 500
-#if __CUDA_ARCH__ < 700
+#if __CUDA_ARCH__ <= 690
   if (!p.advance_to_block()) {
     return;
   }
@@ -23,7 +23,7 @@ fmha_cutlassB_f16_notaligned_64x64_k64_dropout_sm50(typename AttentionBackwardKe
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_f16_notaligned_64x64_k64_dropout_sm50` is for sm50-sm70, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_f16_notaligned_64x64_k64_dropout_sm50` is for sm50-sm69, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -33,7 +33,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_f16_notaligned_64x64_k64_dropout_sm70(typename AttentionBackwardKernel<cutlass::arch::Sm70, cutlass::half_t, false, true, false, 64, 64, 64>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 700
-#if __CUDA_ARCH__ < 750
+#if __CUDA_ARCH__ <= 740
   if (!p.advance_to_block()) {
     return;
   }
@@ -42,7 +42,7 @@ fmha_cutlassB_f16_notaligned_64x64_k64_dropout_sm70(typename AttentionBackwardKe
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_f16_notaligned_64x64_k64_dropout_sm70` is for sm70-sm75, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_f16_notaligned_64x64_k64_dropout_sm70` is for sm70-sm74, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -52,7 +52,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_f16_notaligned_64x64_k64_dropout_sm75(typename AttentionBackwardKernel<cutlass::arch::Sm75, cutlass::half_t, false, true, false, 64, 64, 64>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 750
-#if __CUDA_ARCH__ < 800
+#if __CUDA_ARCH__ <= 790
   if (!p.advance_to_block()) {
     return;
   }
@@ -61,7 +61,7 @@ fmha_cutlassB_f16_notaligned_64x64_k64_dropout_sm75(typename AttentionBackwardKe
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_f16_notaligned_64x64_k64_dropout_sm75` is for sm75-sm80, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_f16_notaligned_64x64_k64_dropout_sm75` is for sm75-sm79, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
diff --git a/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_f16_notaligned_k65536.cu b/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_f16_notaligned_k65536.cu
index 87a5e9e59e42d..b1c42c2ff3ee1 100644
--- a/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_f16_notaligned_k65536.cu
+++ b/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_f16_notaligned_k65536.cu
@@ -14,7 +14,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_f16_notaligned_128x64_k65536_sm70(typename AttentionBackwardKernel<cutlass::arch::Sm70, cutlass::half_t, false, false, false, 128, 64, 65536>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 700
-#if __CUDA_ARCH__ < 750
+#if __CUDA_ARCH__ <= 740
   if (!p.advance_to_block()) {
     return;
   }
@@ -23,7 +23,7 @@ fmha_cutlassB_f16_notaligned_128x64_k65536_sm70(typename AttentionBackwardKernel
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_f16_notaligned_128x64_k65536_sm70` is for sm70-sm75, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_f16_notaligned_128x64_k65536_sm70` is for sm70-sm74, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -33,7 +33,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_f16_notaligned_128x64_k65536_sm75(typename AttentionBackwardKernel<cutlass::arch::Sm75, cutlass::half_t, false, false, false, 128, 64, 65536>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 750
-#if __CUDA_ARCH__ < 800
+#if __CUDA_ARCH__ <= 790
   if (!p.advance_to_block()) {
     return;
   }
@@ -42,7 +42,7 @@ fmha_cutlassB_f16_notaligned_128x64_k65536_sm75(typename AttentionBackwardKernel
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_f16_notaligned_128x64_k65536_sm75` is for sm75-sm80, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_f16_notaligned_128x64_k65536_sm75` is for sm75-sm79, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -52,7 +52,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_f16_notaligned_64x64_k65536_sm50(typename AttentionBackwardKernel<cutlass::arch::Sm50, cutlass::half_t, false, false, false, 64, 64, 65536>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 500
-#if __CUDA_ARCH__ < 700
+#if __CUDA_ARCH__ <= 690
   if (!p.advance_to_block()) {
     return;
   }
@@ -61,7 +61,7 @@ fmha_cutlassB_f16_notaligned_64x64_k65536_sm50(typename AttentionBackwardKernel<
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_f16_notaligned_64x64_k65536_sm50` is for sm50-sm70, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_f16_notaligned_64x64_k65536_sm50` is for sm50-sm69, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -71,7 +71,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_f16_notaligned_64x64_k65536_sm70(typename AttentionBackwardKernel<cutlass::arch::Sm70, cutlass::half_t, false, false, false, 64, 64, 65536>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 700
-#if __CUDA_ARCH__ < 750
+#if __CUDA_ARCH__ <= 740
   if (!p.advance_to_block()) {
     return;
   }
@@ -80,7 +80,7 @@ fmha_cutlassB_f16_notaligned_64x64_k65536_sm70(typename AttentionBackwardKernel<
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_f16_notaligned_64x64_k65536_sm70` is for sm70-sm75, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_f16_notaligned_64x64_k65536_sm70` is for sm70-sm74, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -90,7 +90,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_f16_notaligned_64x64_k65536_sm75(typename AttentionBackwardKernel<cutlass::arch::Sm75, cutlass::half_t, false, false, false, 64, 64, 65536>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 750
-#if __CUDA_ARCH__ < 800
+#if __CUDA_ARCH__ <= 790
   if (!p.advance_to_block()) {
     return;
   }
@@ -99,7 +99,7 @@ fmha_cutlassB_f16_notaligned_64x64_k65536_sm75(typename AttentionBackwardKernel<
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_f16_notaligned_64x64_k65536_sm75` is for sm75-sm80, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_f16_notaligned_64x64_k65536_sm75` is for sm75-sm79, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
diff --git a/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_f16_notaligned_k65536_dropout.cu b/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_f16_notaligned_k65536_dropout.cu
index 5aafa412aa266..6fbd82ea97171 100644
--- a/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_f16_notaligned_k65536_dropout.cu
+++ b/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_f16_notaligned_k65536_dropout.cu
@@ -14,7 +14,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_f16_notaligned_128x64_k65536_dropout_sm70(typename AttentionBackwardKernel<cutlass::arch::Sm70, cutlass::half_t, false, true, false, 128, 64, 65536>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 700
-#if __CUDA_ARCH__ < 750
+#if __CUDA_ARCH__ <= 740
   if (!p.advance_to_block()) {
     return;
   }
@@ -23,7 +23,7 @@ fmha_cutlassB_f16_notaligned_128x64_k65536_dropout_sm70(typename AttentionBackwa
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_f16_notaligned_128x64_k65536_dropout_sm70` is for sm70-sm75, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_f16_notaligned_128x64_k65536_dropout_sm70` is for sm70-sm74, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -33,7 +33,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_f16_notaligned_128x64_k65536_dropout_sm75(typename AttentionBackwardKernel<cutlass::arch::Sm75, cutlass::half_t, false, true, false, 128, 64, 65536>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 750
-#if __CUDA_ARCH__ < 800
+#if __CUDA_ARCH__ <= 790
   if (!p.advance_to_block()) {
     return;
   }
@@ -42,7 +42,7 @@ fmha_cutlassB_f16_notaligned_128x64_k65536_dropout_sm75(typename AttentionBackwa
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_f16_notaligned_128x64_k65536_dropout_sm75` is for sm75-sm80, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_f16_notaligned_128x64_k65536_dropout_sm75` is for sm75-sm79, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -52,7 +52,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_f16_notaligned_64x64_k65536_dropout_sm50(typename AttentionBackwardKernel<cutlass::arch::Sm50, cutlass::half_t, false, true, false, 64, 64, 65536>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 500
-#if __CUDA_ARCH__ < 700
+#if __CUDA_ARCH__ <= 690
   if (!p.advance_to_block()) {
     return;
   }
@@ -61,7 +61,7 @@ fmha_cutlassB_f16_notaligned_64x64_k65536_dropout_sm50(typename AttentionBackwar
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_f16_notaligned_64x64_k65536_dropout_sm50` is for sm50-sm70, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_f16_notaligned_64x64_k65536_dropout_sm50` is for sm50-sm69, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -71,7 +71,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_f16_notaligned_64x64_k65536_dropout_sm70(typename AttentionBackwardKernel<cutlass::arch::Sm70, cutlass::half_t, false, true, false, 64, 64, 65536>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 700
-#if __CUDA_ARCH__ < 750
+#if __CUDA_ARCH__ <= 740
   if (!p.advance_to_block()) {
     return;
   }
@@ -80,7 +80,7 @@ fmha_cutlassB_f16_notaligned_64x64_k65536_dropout_sm70(typename AttentionBackwar
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_f16_notaligned_64x64_k65536_dropout_sm70` is for sm70-sm75, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_f16_notaligned_64x64_k65536_dropout_sm70` is for sm70-sm74, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -90,7 +90,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_f16_notaligned_64x64_k65536_dropout_sm75(typename AttentionBackwardKernel<cutlass::arch::Sm75, cutlass::half_t, false, true, false, 64, 64, 65536>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 750
-#if __CUDA_ARCH__ < 800
+#if __CUDA_ARCH__ <= 790
   if (!p.advance_to_block()) {
     return;
   }
@@ -99,7 +99,7 @@ fmha_cutlassB_f16_notaligned_64x64_k65536_dropout_sm75(typename AttentionBackwar
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_f16_notaligned_64x64_k65536_dropout_sm75` is for sm75-sm80, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_f16_notaligned_64x64_k65536_dropout_sm75` is for sm75-sm79, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
diff --git a/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_f32_aligned_k128.cu b/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_f32_aligned_k128.cu
index 7ca636a909e58..39aa4598eb870 100644
--- a/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_f32_aligned_k128.cu
+++ b/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_f32_aligned_k128.cu
@@ -14,7 +14,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_f32_aligned_128x64_k128_sm80(typename AttentionBackwardKernel<cutlass::arch::Sm80, float, true, false, false, 128, 64, 128>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 800
-#if __CUDA_ARCH__ <= 1200
+#if __CUDA_ARCH__ <= 1210
   if (!p.advance_to_block()) {
     return;
   }
@@ -23,7 +23,7 @@ fmha_cutlassB_f32_aligned_128x64_k128_sm80(typename AttentionBackwardKernel<cutl
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_f32_aligned_128x64_k128_sm80` is for sm80-sm100, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_f32_aligned_128x64_k128_sm80` is for sm80-sm121, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -33,7 +33,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_f32_aligned_64x64_k128_sm50(typename AttentionBackwardKernel<cutlass::arch::Sm50, float, true, false, false, 64, 64, 128>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 500
-#if __CUDA_ARCH__ < 700
+#if __CUDA_ARCH__ <= 690
   if (!p.advance_to_block()) {
     return;
   }
@@ -42,7 +42,7 @@ fmha_cutlassB_f32_aligned_64x64_k128_sm50(typename AttentionBackwardKernel<cutla
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_f32_aligned_64x64_k128_sm50` is for sm50-sm70, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_f32_aligned_64x64_k128_sm50` is for sm50-sm69, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -52,7 +52,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_f32_aligned_64x64_k128_sm70(typename AttentionBackwardKernel<cutlass::arch::Sm70, float, true, false, false, 64, 64, 128>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 700
-#if __CUDA_ARCH__ < 750
+#if __CUDA_ARCH__ <= 740
   if (!p.advance_to_block()) {
     return;
   }
@@ -61,7 +61,7 @@ fmha_cutlassB_f32_aligned_64x64_k128_sm70(typename AttentionBackwardKernel<cutla
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_f32_aligned_64x64_k128_sm70` is for sm70-sm75, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_f32_aligned_64x64_k128_sm70` is for sm70-sm74, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -71,7 +71,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_f32_aligned_64x64_k128_sm75(typename AttentionBackwardKernel<cutlass::arch::Sm75, float, true, false, false, 64, 64, 128>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 750
-#if __CUDA_ARCH__ < 800
+#if __CUDA_ARCH__ <= 790
   if (!p.advance_to_block()) {
     return;
   }
@@ -80,7 +80,7 @@ fmha_cutlassB_f32_aligned_64x64_k128_sm75(typename AttentionBackwardKernel<cutla
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_f32_aligned_64x64_k128_sm75` is for sm75-sm80, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_f32_aligned_64x64_k128_sm75` is for sm75-sm79, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -90,7 +90,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_f32_aligned_64x64_k128_sm80(typename AttentionBackwardKernel<cutlass::arch::Sm80, float, true, false, false, 64, 64, 128>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 800
-#if __CUDA_ARCH__ <= 1200
+#if __CUDA_ARCH__ <= 1210
   if (!p.advance_to_block()) {
     return;
   }
@@ -99,7 +99,7 @@ fmha_cutlassB_f32_aligned_64x64_k128_sm80(typename AttentionBackwardKernel<cutla
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_f32_aligned_64x64_k128_sm80` is for sm80-sm100, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_f32_aligned_64x64_k128_sm80` is for sm80-sm121, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
diff --git a/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_f32_aligned_k128_dropout.cu b/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_f32_aligned_k128_dropout.cu
index c279e4c0040b8..a8985a446e626 100644
--- a/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_f32_aligned_k128_dropout.cu
+++ b/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_f32_aligned_k128_dropout.cu
@@ -14,7 +14,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_f32_aligned_128x64_k128_dropout_sm80(typename AttentionBackwardKernel<cutlass::arch::Sm80, float, true, true, false, 128, 64, 128>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 800
-#if __CUDA_ARCH__ <= 1200
+#if __CUDA_ARCH__ <= 1210
   if (!p.advance_to_block()) {
     return;
   }
@@ -23,7 +23,7 @@ fmha_cutlassB_f32_aligned_128x64_k128_dropout_sm80(typename AttentionBackwardKer
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_f32_aligned_128x64_k128_dropout_sm80` is for sm80-sm100, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_f32_aligned_128x64_k128_dropout_sm80` is for sm80-sm121, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -33,7 +33,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_f32_aligned_64x64_k128_dropout_sm50(typename AttentionBackwardKernel<cutlass::arch::Sm50, float, true, true, false, 64, 64, 128>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 500
-#if __CUDA_ARCH__ < 700
+#if __CUDA_ARCH__ <= 690
   if (!p.advance_to_block()) {
     return;
   }
@@ -42,7 +42,7 @@ fmha_cutlassB_f32_aligned_64x64_k128_dropout_sm50(typename AttentionBackwardKern
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_f32_aligned_64x64_k128_dropout_sm50` is for sm50-sm70, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_f32_aligned_64x64_k128_dropout_sm50` is for sm50-sm69, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -52,7 +52,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_f32_aligned_64x64_k128_dropout_sm70(typename AttentionBackwardKernel<cutlass::arch::Sm70, float, true, true, false, 64, 64, 128>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 700
-#if __CUDA_ARCH__ < 750
+#if __CUDA_ARCH__ <= 740
   if (!p.advance_to_block()) {
     return;
   }
@@ -61,7 +61,7 @@ fmha_cutlassB_f32_aligned_64x64_k128_dropout_sm70(typename AttentionBackwardKern
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_f32_aligned_64x64_k128_dropout_sm70` is for sm70-sm75, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_f32_aligned_64x64_k128_dropout_sm70` is for sm70-sm74, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -71,7 +71,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_f32_aligned_64x64_k128_dropout_sm75(typename AttentionBackwardKernel<cutlass::arch::Sm75, float, true, true, false, 64, 64, 128>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 750
-#if __CUDA_ARCH__ < 800
+#if __CUDA_ARCH__ <= 790
   if (!p.advance_to_block()) {
     return;
   }
@@ -80,7 +80,7 @@ fmha_cutlassB_f32_aligned_64x64_k128_dropout_sm75(typename AttentionBackwardKern
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_f32_aligned_64x64_k128_dropout_sm75` is for sm75-sm80, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_f32_aligned_64x64_k128_dropout_sm75` is for sm75-sm79, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -90,7 +90,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_f32_aligned_64x64_k128_dropout_sm80(typename AttentionBackwardKernel<cutlass::arch::Sm80, float, true, true, false, 64, 64, 128>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 800
-#if __CUDA_ARCH__ <= 1200
+#if __CUDA_ARCH__ <= 1210
   if (!p.advance_to_block()) {
     return;
   }
@@ -99,7 +99,7 @@ fmha_cutlassB_f32_aligned_64x64_k128_dropout_sm80(typename AttentionBackwardKern
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_f32_aligned_64x64_k128_dropout_sm80` is for sm80-sm100, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_f32_aligned_64x64_k128_dropout_sm80` is for sm80-sm121, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
diff --git a/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_f32_aligned_k32.cu b/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_f32_aligned_k32.cu
index fe238a798ea91..acc3f63155aae 100644
--- a/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_f32_aligned_k32.cu
+++ b/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_f32_aligned_k32.cu
@@ -14,7 +14,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_f32_aligned_64x64_k32_sm50(typename AttentionBackwardKernel<cutlass::arch::Sm50, float, true, false, false, 64, 64, 32>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 500
-#if __CUDA_ARCH__ < 700
+#if __CUDA_ARCH__ <= 690
   if (!p.advance_to_block()) {
     return;
   }
@@ -23,7 +23,7 @@ fmha_cutlassB_f32_aligned_64x64_k32_sm50(typename AttentionBackwardKernel<cutlas
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_f32_aligned_64x64_k32_sm50` is for sm50-sm70, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_f32_aligned_64x64_k32_sm50` is for sm50-sm69, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -33,7 +33,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_f32_aligned_64x64_k32_sm70(typename AttentionBackwardKernel<cutlass::arch::Sm70, float, true, false, false, 64, 64, 32>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 700
-#if __CUDA_ARCH__ < 750
+#if __CUDA_ARCH__ <= 740
   if (!p.advance_to_block()) {
     return;
   }
@@ -42,7 +42,7 @@ fmha_cutlassB_f32_aligned_64x64_k32_sm70(typename AttentionBackwardKernel<cutlas
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_f32_aligned_64x64_k32_sm70` is for sm70-sm75, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_f32_aligned_64x64_k32_sm70` is for sm70-sm74, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -52,7 +52,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_f32_aligned_64x64_k32_sm75(typename AttentionBackwardKernel<cutlass::arch::Sm75, float, true, false, false, 64, 64, 32>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 750
-#if __CUDA_ARCH__ < 800
+#if __CUDA_ARCH__ <= 790
   if (!p.advance_to_block()) {
     return;
   }
@@ -61,7 +61,7 @@ fmha_cutlassB_f32_aligned_64x64_k32_sm75(typename AttentionBackwardKernel<cutlas
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_f32_aligned_64x64_k32_sm75` is for sm75-sm80, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_f32_aligned_64x64_k32_sm75` is for sm75-sm79, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -71,7 +71,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_f32_aligned_64x64_k32_sm80(typename AttentionBackwardKernel<cutlass::arch::Sm80, float, true, false, false, 64, 64, 32>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 800
-#if __CUDA_ARCH__ <= 1200
+#if __CUDA_ARCH__ <= 1210
   if (!p.advance_to_block()) {
     return;
   }
@@ -80,7 +80,7 @@ fmha_cutlassB_f32_aligned_64x64_k32_sm80(typename AttentionBackwardKernel<cutlas
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_f32_aligned_64x64_k32_sm80` is for sm80-sm100, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_f32_aligned_64x64_k32_sm80` is for sm80-sm121, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
diff --git a/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_f32_aligned_k32_dropout.cu b/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_f32_aligned_k32_dropout.cu
index 425cb293471fa..d33d5aa595b26 100644
--- a/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_f32_aligned_k32_dropout.cu
+++ b/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_f32_aligned_k32_dropout.cu
@@ -8,34 +8,13 @@
 // This file is auto-generated. See "generate_kernels.py"
 #include <ATen/native/transformers/cuda/mem_eff_attention/kernel_backward.h>
 using namespace PyTorchMemEffAttention;
-#if defined(CUDA_VERSION) && CUDA_VERSION == 12040 && !defined(USE_ROCM)
-__global__ void __launch_bounds__(
-    AttentionBackwardKernel<cutlass::arch::Sm50, float, true, true, false, 32, 32, 32>::kNumThreads,
-    AttentionBackwardKernel<cutlass::arch::Sm50, float, true, true, false, 32, 32, 32>::kMinBlocksPerSm)
-fmha_cutlassB_f32_aligned_32x32_k32_dropout_sm50(typename AttentionBackwardKernel<cutlass::arch::Sm50, float, true, true, false, 32, 32, 32>::Params p) {
-#ifdef __CUDA_ARCH__
-#if __CUDA_ARCH__ >= 500
-#if __CUDA_ARCH__ < 700
-  if (!p.advance_to_block()) {
-    return;
-  }
-  AttentionBackwardKernel<cutlass::arch::Sm50, float, true, true, false, 32, 32, 32>::attention_kernel(p);
-  return;
-#endif
-#endif
-    printf(
-        "FATAL: kernel `fmha_cutlassB_f32_aligned_32x32_k32_dropout_sm50` is for sm50-sm70, but was built for sm%d\n",
-        int(__CUDA_ARCH__ + 0) / 10);
-#endif
-}
-#else
 __global__ void __launch_bounds__(
     AttentionBackwardKernel<cutlass::arch::Sm50, float, true, true, false, 64, 64, 32>::kNumThreads,
     AttentionBackwardKernel<cutlass::arch::Sm50, float, true, true, false, 64, 64, 32>::kMinBlocksPerSm)
 fmha_cutlassB_f32_aligned_64x64_k32_dropout_sm50(typename AttentionBackwardKernel<cutlass::arch::Sm50, float, true, true, false, 64, 64, 32>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 500
-#if __CUDA_ARCH__ < 700
+#if __CUDA_ARCH__ <= 690
   if (!p.advance_to_block()) {
     return;
   }
@@ -44,18 +23,17 @@ fmha_cutlassB_f32_aligned_64x64_k32_dropout_sm50(typename AttentionBackwardKerne
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_f32_aligned_64x64_k32_dropout_sm50` is for sm50-sm70, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_f32_aligned_64x64_k32_dropout_sm50` is for sm50-sm69, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
-#endif
 __global__ void __launch_bounds__(
     AttentionBackwardKernel<cutlass::arch::Sm70, float, true, true, false, 64, 64, 32>::kNumThreads,
     AttentionBackwardKernel<cutlass::arch::Sm70, float, true, true, false, 64, 64, 32>::kMinBlocksPerSm)
 fmha_cutlassB_f32_aligned_64x64_k32_dropout_sm70(typename AttentionBackwardKernel<cutlass::arch::Sm70, float, true, true, false, 64, 64, 32>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 700
-#if __CUDA_ARCH__ < 750
+#if __CUDA_ARCH__ <= 740
   if (!p.advance_to_block()) {
     return;
   }
@@ -64,7 +42,7 @@ fmha_cutlassB_f32_aligned_64x64_k32_dropout_sm70(typename AttentionBackwardKerne
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_f32_aligned_64x64_k32_dropout_sm70` is for sm70-sm75, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_f32_aligned_64x64_k32_dropout_sm70` is for sm70-sm74, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -74,7 +52,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_f32_aligned_64x64_k32_dropout_sm75(typename AttentionBackwardKernel<cutlass::arch::Sm75, float, true, true, false, 64, 64, 32>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 750
-#if __CUDA_ARCH__ < 800
+#if __CUDA_ARCH__ <= 790
   if (!p.advance_to_block()) {
     return;
   }
@@ -83,7 +61,7 @@ fmha_cutlassB_f32_aligned_64x64_k32_dropout_sm75(typename AttentionBackwardKerne
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_f32_aligned_64x64_k32_dropout_sm75` is for sm75-sm80, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_f32_aligned_64x64_k32_dropout_sm75` is for sm75-sm79, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -93,7 +71,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_f32_aligned_64x64_k32_dropout_sm80(typename AttentionBackwardKernel<cutlass::arch::Sm80, float, true, true, false, 64, 64, 32>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 800
-#if __CUDA_ARCH__ <= 1200
+#if __CUDA_ARCH__ <= 1210
   if (!p.advance_to_block()) {
     return;
   }
@@ -102,7 +80,7 @@ fmha_cutlassB_f32_aligned_64x64_k32_dropout_sm80(typename AttentionBackwardKerne
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_f32_aligned_64x64_k32_dropout_sm80` is for sm80-sm100, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_f32_aligned_64x64_k32_dropout_sm80` is for sm80-sm121, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
diff --git a/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_f32_aligned_k64.cu b/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_f32_aligned_k64.cu
index da5a0bea92690..e31ef4db4c57e 100644
--- a/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_f32_aligned_k64.cu
+++ b/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_f32_aligned_k64.cu
@@ -14,7 +14,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_f32_aligned_64x64_k64_sm50(typename AttentionBackwardKernel<cutlass::arch::Sm50, float, true, false, false, 64, 64, 64>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 500
-#if __CUDA_ARCH__ < 700
+#if __CUDA_ARCH__ <= 690
   if (!p.advance_to_block()) {
     return;
   }
@@ -23,7 +23,7 @@ fmha_cutlassB_f32_aligned_64x64_k64_sm50(typename AttentionBackwardKernel<cutlas
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_f32_aligned_64x64_k64_sm50` is for sm50-sm70, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_f32_aligned_64x64_k64_sm50` is for sm50-sm69, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -33,7 +33,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_f32_aligned_64x64_k64_sm70(typename AttentionBackwardKernel<cutlass::arch::Sm70, float, true, false, false, 64, 64, 64>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 700
-#if __CUDA_ARCH__ < 750
+#if __CUDA_ARCH__ <= 740
   if (!p.advance_to_block()) {
     return;
   }
@@ -42,7 +42,7 @@ fmha_cutlassB_f32_aligned_64x64_k64_sm70(typename AttentionBackwardKernel<cutlas
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_f32_aligned_64x64_k64_sm70` is for sm70-sm75, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_f32_aligned_64x64_k64_sm70` is for sm70-sm74, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -52,7 +52,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_f32_aligned_64x64_k64_sm75(typename AttentionBackwardKernel<cutlass::arch::Sm75, float, true, false, false, 64, 64, 64>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 750
-#if __CUDA_ARCH__ < 800
+#if __CUDA_ARCH__ <= 790
   if (!p.advance_to_block()) {
     return;
   }
@@ -61,7 +61,7 @@ fmha_cutlassB_f32_aligned_64x64_k64_sm75(typename AttentionBackwardKernel<cutlas
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_f32_aligned_64x64_k64_sm75` is for sm75-sm80, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_f32_aligned_64x64_k64_sm75` is for sm75-sm79, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -71,7 +71,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_f32_aligned_64x64_k64_sm80(typename AttentionBackwardKernel<cutlass::arch::Sm80, float, true, false, false, 64, 64, 64>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 800
-#if __CUDA_ARCH__ <= 1200
+#if __CUDA_ARCH__ <= 1210
   if (!p.advance_to_block()) {
     return;
   }
@@ -80,7 +80,7 @@ fmha_cutlassB_f32_aligned_64x64_k64_sm80(typename AttentionBackwardKernel<cutlas
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_f32_aligned_64x64_k64_sm80` is for sm80-sm100, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_f32_aligned_64x64_k64_sm80` is for sm80-sm121, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
diff --git a/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_f32_aligned_k64_dropout.cu b/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_f32_aligned_k64_dropout.cu
index 8dd2ac4c436fb..d1d5fc0a19353 100644
--- a/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_f32_aligned_k64_dropout.cu
+++ b/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_f32_aligned_k64_dropout.cu
@@ -8,34 +8,13 @@
 // This file is auto-generated. See "generate_kernels.py"
 #include <ATen/native/transformers/cuda/mem_eff_attention/kernel_backward.h>
 using namespace PyTorchMemEffAttention;
-#if defined(CUDA_VERSION) && CUDA_VERSION == 12040 && !defined(USE_ROCM)
-__global__ void __launch_bounds__(
-    AttentionBackwardKernel<cutlass::arch::Sm50, float, true, true, false, 32, 32, 64>::kNumThreads,
-    AttentionBackwardKernel<cutlass::arch::Sm50, float, true, true, false, 32, 32, 64>::kMinBlocksPerSm)
-fmha_cutlassB_f32_aligned_32x32_k64_dropout_sm50(typename AttentionBackwardKernel<cutlass::arch::Sm50, float, true, true, false, 32, 32, 64>::Params p) {
-#ifdef __CUDA_ARCH__
-#if __CUDA_ARCH__ >= 500
-#if __CUDA_ARCH__ < 700
-  if (!p.advance_to_block()) {
-    return;
-  }
-  AttentionBackwardKernel<cutlass::arch::Sm50, float, true, true, false, 32, 32, 64>::attention_kernel(p);
-  return;
-#endif
-#endif
-    printf(
-        "FATAL: kernel `fmha_cutlassB_f32_aligned_64x64_k64_dropout_sm50` is for sm50-sm70, but was built for sm%d\n",
-        int(__CUDA_ARCH__ + 0) / 10);
-#endif
-}
-#else
 __global__ void __launch_bounds__(
     AttentionBackwardKernel<cutlass::arch::Sm50, float, true, true, false, 64, 64, 64>::kNumThreads,
     AttentionBackwardKernel<cutlass::arch::Sm50, float, true, true, false, 64, 64, 64>::kMinBlocksPerSm)
 fmha_cutlassB_f32_aligned_64x64_k64_dropout_sm50(typename AttentionBackwardKernel<cutlass::arch::Sm50, float, true, true, false, 64, 64, 64>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 500
-#if __CUDA_ARCH__ < 700
+#if __CUDA_ARCH__ <= 690
   if (!p.advance_to_block()) {
     return;
   }
@@ -44,18 +23,17 @@ fmha_cutlassB_f32_aligned_64x64_k64_dropout_sm50(typename AttentionBackwardKerne
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_f32_aligned_64x64_k64_dropout_sm50` is for sm50-sm70, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_f32_aligned_64x64_k64_dropout_sm50` is for sm50-sm69, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
-#endif
 __global__ void __launch_bounds__(
     AttentionBackwardKernel<cutlass::arch::Sm70, float, true, true, false, 64, 64, 64>::kNumThreads,
     AttentionBackwardKernel<cutlass::arch::Sm70, float, true, true, false, 64, 64, 64>::kMinBlocksPerSm)
 fmha_cutlassB_f32_aligned_64x64_k64_dropout_sm70(typename AttentionBackwardKernel<cutlass::arch::Sm70, float, true, true, false, 64, 64, 64>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 700
-#if __CUDA_ARCH__ < 750
+#if __CUDA_ARCH__ <= 740
   if (!p.advance_to_block()) {
     return;
   }
@@ -64,7 +42,7 @@ fmha_cutlassB_f32_aligned_64x64_k64_dropout_sm70(typename AttentionBackwardKerne
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_f32_aligned_64x64_k64_dropout_sm70` is for sm70-sm75, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_f32_aligned_64x64_k64_dropout_sm70` is for sm70-sm74, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -74,7 +52,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_f32_aligned_64x64_k64_dropout_sm75(typename AttentionBackwardKernel<cutlass::arch::Sm75, float, true, true, false, 64, 64, 64>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 750
-#if __CUDA_ARCH__ < 800
+#if __CUDA_ARCH__ <= 790
   if (!p.advance_to_block()) {
     return;
   }
@@ -83,7 +61,7 @@ fmha_cutlassB_f32_aligned_64x64_k64_dropout_sm75(typename AttentionBackwardKerne
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_f32_aligned_64x64_k64_dropout_sm75` is for sm75-sm80, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_f32_aligned_64x64_k64_dropout_sm75` is for sm75-sm79, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -93,7 +71,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_f32_aligned_64x64_k64_dropout_sm80(typename AttentionBackwardKernel<cutlass::arch::Sm80, float, true, true, false, 64, 64, 64>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 800
-#if __CUDA_ARCH__ <= 1200
+#if __CUDA_ARCH__ <= 1210
   if (!p.advance_to_block()) {
     return;
   }
@@ -102,7 +80,7 @@ fmha_cutlassB_f32_aligned_64x64_k64_dropout_sm80(typename AttentionBackwardKerne
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_f32_aligned_64x64_k64_dropout_sm80` is for sm80-sm100, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_f32_aligned_64x64_k64_dropout_sm80` is for sm80-sm121, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
diff --git a/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_f32_aligned_k65536.cu b/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_f32_aligned_k65536.cu
index 5a6a56898a770..5af3d48a1a839 100644
--- a/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_f32_aligned_k65536.cu
+++ b/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_f32_aligned_k65536.cu
@@ -14,7 +14,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_f32_aligned_128x64_k65536_sm80(typename AttentionBackwardKernel<cutlass::arch::Sm80, float, true, false, false, 128, 64, 65536>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 800
-#if __CUDA_ARCH__ <= 1200
+#if __CUDA_ARCH__ <= 1210
   if (!p.advance_to_block()) {
     return;
   }
@@ -23,7 +23,7 @@ fmha_cutlassB_f32_aligned_128x64_k65536_sm80(typename AttentionBackwardKernel<cu
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_f32_aligned_128x64_k65536_sm80` is for sm80-sm100, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_f32_aligned_128x64_k65536_sm80` is for sm80-sm121, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -33,7 +33,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_f32_aligned_64x64_k65536_sm50(typename AttentionBackwardKernel<cutlass::arch::Sm50, float, true, false, false, 64, 64, 65536>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 500
-#if __CUDA_ARCH__ < 700
+#if __CUDA_ARCH__ <= 690
   if (!p.advance_to_block()) {
     return;
   }
@@ -42,7 +42,7 @@ fmha_cutlassB_f32_aligned_64x64_k65536_sm50(typename AttentionBackwardKernel<cut
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_f32_aligned_64x64_k65536_sm50` is for sm50-sm70, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_f32_aligned_64x64_k65536_sm50` is for sm50-sm69, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -52,7 +52,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_f32_aligned_64x64_k65536_sm70(typename AttentionBackwardKernel<cutlass::arch::Sm70, float, true, false, false, 64, 64, 65536>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 700
-#if __CUDA_ARCH__ < 750
+#if __CUDA_ARCH__ <= 740
   if (!p.advance_to_block()) {
     return;
   }
@@ -61,7 +61,7 @@ fmha_cutlassB_f32_aligned_64x64_k65536_sm70(typename AttentionBackwardKernel<cut
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_f32_aligned_64x64_k65536_sm70` is for sm70-sm75, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_f32_aligned_64x64_k65536_sm70` is for sm70-sm74, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -71,7 +71,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_f32_aligned_64x64_k65536_sm75(typename AttentionBackwardKernel<cutlass::arch::Sm75, float, true, false, false, 64, 64, 65536>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 750
-#if __CUDA_ARCH__ < 800
+#if __CUDA_ARCH__ <= 790
   if (!p.advance_to_block()) {
     return;
   }
@@ -80,7 +80,7 @@ fmha_cutlassB_f32_aligned_64x64_k65536_sm75(typename AttentionBackwardKernel<cut
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_f32_aligned_64x64_k65536_sm75` is for sm75-sm80, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_f32_aligned_64x64_k65536_sm75` is for sm75-sm79, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -90,7 +90,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_f32_aligned_64x64_k65536_sm80(typename AttentionBackwardKernel<cutlass::arch::Sm80, float, true, false, false, 64, 64, 65536>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 800
-#if __CUDA_ARCH__ <= 1200
+#if __CUDA_ARCH__ <= 1210
   if (!p.advance_to_block()) {
     return;
   }
@@ -99,7 +99,7 @@ fmha_cutlassB_f32_aligned_64x64_k65536_sm80(typename AttentionBackwardKernel<cut
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_f32_aligned_64x64_k65536_sm80` is for sm80-sm100, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_f32_aligned_64x64_k65536_sm80` is for sm80-sm121, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
diff --git a/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_f32_aligned_k65536_dropout.cu b/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_f32_aligned_k65536_dropout.cu
index 3bc89c5614ee7..4a11a86c8e90a 100644
--- a/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_f32_aligned_k65536_dropout.cu
+++ b/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_f32_aligned_k65536_dropout.cu
@@ -14,7 +14,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_f32_aligned_128x64_k65536_dropout_sm80(typename AttentionBackwardKernel<cutlass::arch::Sm80, float, true, true, false, 128, 64, 65536>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 800
-#if __CUDA_ARCH__ <= 1200
+#if __CUDA_ARCH__ <= 1210
   if (!p.advance_to_block()) {
     return;
   }
@@ -23,7 +23,7 @@ fmha_cutlassB_f32_aligned_128x64_k65536_dropout_sm80(typename AttentionBackwardK
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_f32_aligned_128x64_k65536_dropout_sm80` is for sm80-sm100, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_f32_aligned_128x64_k65536_dropout_sm80` is for sm80-sm121, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -33,7 +33,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_f32_aligned_64x64_k65536_dropout_sm50(typename AttentionBackwardKernel<cutlass::arch::Sm50, float, true, true, false, 64, 64, 65536>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 500
-#if __CUDA_ARCH__ < 700
+#if __CUDA_ARCH__ <= 690
   if (!p.advance_to_block()) {
     return;
   }
@@ -42,7 +42,7 @@ fmha_cutlassB_f32_aligned_64x64_k65536_dropout_sm50(typename AttentionBackwardKe
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_f32_aligned_64x64_k65536_dropout_sm50` is for sm50-sm70, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_f32_aligned_64x64_k65536_dropout_sm50` is for sm50-sm69, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -52,7 +52,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_f32_aligned_64x64_k65536_dropout_sm70(typename AttentionBackwardKernel<cutlass::arch::Sm70, float, true, true, false, 64, 64, 65536>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 700
-#if __CUDA_ARCH__ < 750
+#if __CUDA_ARCH__ <= 740
   if (!p.advance_to_block()) {
     return;
   }
@@ -61,7 +61,7 @@ fmha_cutlassB_f32_aligned_64x64_k65536_dropout_sm70(typename AttentionBackwardKe
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_f32_aligned_64x64_k65536_dropout_sm70` is for sm70-sm75, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_f32_aligned_64x64_k65536_dropout_sm70` is for sm70-sm74, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -71,7 +71,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_f32_aligned_64x64_k65536_dropout_sm75(typename AttentionBackwardKernel<cutlass::arch::Sm75, float, true, true, false, 64, 64, 65536>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 750
-#if __CUDA_ARCH__ < 800
+#if __CUDA_ARCH__ <= 790
   if (!p.advance_to_block()) {
     return;
   }
@@ -80,7 +80,7 @@ fmha_cutlassB_f32_aligned_64x64_k65536_dropout_sm75(typename AttentionBackwardKe
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_f32_aligned_64x64_k65536_dropout_sm75` is for sm75-sm80, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_f32_aligned_64x64_k65536_dropout_sm75` is for sm75-sm79, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -90,7 +90,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_f32_aligned_64x64_k65536_dropout_sm80(typename AttentionBackwardKernel<cutlass::arch::Sm80, float, true, true, false, 64, 64, 65536>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 800
-#if __CUDA_ARCH__ <= 1200
+#if __CUDA_ARCH__ <= 1210
   if (!p.advance_to_block()) {
     return;
   }
@@ -99,7 +99,7 @@ fmha_cutlassB_f32_aligned_64x64_k65536_dropout_sm80(typename AttentionBackwardKe
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_f32_aligned_64x64_k65536_dropout_sm80` is for sm80-sm100, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_f32_aligned_64x64_k65536_dropout_sm80` is for sm80-sm121, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
diff --git a/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_f32_notaligned_k128.cu b/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_f32_notaligned_k128.cu
index 90d3c1831e005..3b9ba3450ec5b 100644
--- a/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_f32_notaligned_k128.cu
+++ b/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_f32_notaligned_k128.cu
@@ -14,7 +14,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_f32_notaligned_64x64_k128_sm50(typename AttentionBackwardKernel<cutlass::arch::Sm50, float, false, false, false, 64, 64, 128>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 500
-#if __CUDA_ARCH__ < 700
+#if __CUDA_ARCH__ <= 690
   if (!p.advance_to_block()) {
     return;
   }
@@ -23,7 +23,7 @@ fmha_cutlassB_f32_notaligned_64x64_k128_sm50(typename AttentionBackwardKernel<cu
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_f32_notaligned_64x64_k128_sm50` is for sm50-sm70, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_f32_notaligned_64x64_k128_sm50` is for sm50-sm69, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -33,7 +33,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_f32_notaligned_64x64_k128_sm70(typename AttentionBackwardKernel<cutlass::arch::Sm70, float, false, false, false, 64, 64, 128>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 700
-#if __CUDA_ARCH__ < 750
+#if __CUDA_ARCH__ <= 740
   if (!p.advance_to_block()) {
     return;
   }
@@ -42,7 +42,7 @@ fmha_cutlassB_f32_notaligned_64x64_k128_sm70(typename AttentionBackwardKernel<cu
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_f32_notaligned_64x64_k128_sm70` is for sm70-sm75, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_f32_notaligned_64x64_k128_sm70` is for sm70-sm74, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -52,7 +52,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_f32_notaligned_64x64_k128_sm75(typename AttentionBackwardKernel<cutlass::arch::Sm75, float, false, false, false, 64, 64, 128>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 750
-#if __CUDA_ARCH__ < 800
+#if __CUDA_ARCH__ <= 790
   if (!p.advance_to_block()) {
     return;
   }
@@ -61,7 +61,7 @@ fmha_cutlassB_f32_notaligned_64x64_k128_sm75(typename AttentionBackwardKernel<cu
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_f32_notaligned_64x64_k128_sm75` is for sm75-sm80, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_f32_notaligned_64x64_k128_sm75` is for sm75-sm79, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
diff --git a/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_f32_notaligned_k128_dropout.cu b/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_f32_notaligned_k128_dropout.cu
index 9bb8a6a6c7deb..ad0db51b268b2 100644
--- a/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_f32_notaligned_k128_dropout.cu
+++ b/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_f32_notaligned_k128_dropout.cu
@@ -14,7 +14,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_f32_notaligned_64x64_k128_dropout_sm50(typename AttentionBackwardKernel<cutlass::arch::Sm50, float, false, true, false, 64, 64, 128>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 500
-#if __CUDA_ARCH__ < 700
+#if __CUDA_ARCH__ <= 690
   if (!p.advance_to_block()) {
     return;
   }
@@ -23,7 +23,7 @@ fmha_cutlassB_f32_notaligned_64x64_k128_dropout_sm50(typename AttentionBackwardK
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_f32_notaligned_64x64_k128_dropout_sm50` is for sm50-sm70, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_f32_notaligned_64x64_k128_dropout_sm50` is for sm50-sm69, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -33,7 +33,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_f32_notaligned_64x64_k128_dropout_sm70(typename AttentionBackwardKernel<cutlass::arch::Sm70, float, false, true, false, 64, 64, 128>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 700
-#if __CUDA_ARCH__ < 750
+#if __CUDA_ARCH__ <= 740
   if (!p.advance_to_block()) {
     return;
   }
@@ -42,7 +42,7 @@ fmha_cutlassB_f32_notaligned_64x64_k128_dropout_sm70(typename AttentionBackwardK
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_f32_notaligned_64x64_k128_dropout_sm70` is for sm70-sm75, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_f32_notaligned_64x64_k128_dropout_sm70` is for sm70-sm74, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -52,7 +52,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_f32_notaligned_64x64_k128_dropout_sm75(typename AttentionBackwardKernel<cutlass::arch::Sm75, float, false, true, false, 64, 64, 128>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 750
-#if __CUDA_ARCH__ < 800
+#if __CUDA_ARCH__ <= 790
   if (!p.advance_to_block()) {
     return;
   }
@@ -61,7 +61,7 @@ fmha_cutlassB_f32_notaligned_64x64_k128_dropout_sm75(typename AttentionBackwardK
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_f32_notaligned_64x64_k128_dropout_sm75` is for sm75-sm80, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_f32_notaligned_64x64_k128_dropout_sm75` is for sm75-sm79, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
diff --git a/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_f32_notaligned_k32.cu b/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_f32_notaligned_k32.cu
index db1dc62a95f94..cc1621b21dfb4 100644
--- a/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_f32_notaligned_k32.cu
+++ b/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_f32_notaligned_k32.cu
@@ -14,7 +14,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_f32_notaligned_64x64_k32_sm50(typename AttentionBackwardKernel<cutlass::arch::Sm50, float, false, false, false, 64, 64, 32>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 500
-#if __CUDA_ARCH__ < 700
+#if __CUDA_ARCH__ <= 690
   if (!p.advance_to_block()) {
     return;
   }
@@ -23,7 +23,7 @@ fmha_cutlassB_f32_notaligned_64x64_k32_sm50(typename AttentionBackwardKernel<cut
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_f32_notaligned_64x64_k32_sm50` is for sm50-sm70, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_f32_notaligned_64x64_k32_sm50` is for sm50-sm69, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -33,7 +33,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_f32_notaligned_64x64_k32_sm70(typename AttentionBackwardKernel<cutlass::arch::Sm70, float, false, false, false, 64, 64, 32>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 700
-#if __CUDA_ARCH__ < 750
+#if __CUDA_ARCH__ <= 740
   if (!p.advance_to_block()) {
     return;
   }
@@ -42,7 +42,7 @@ fmha_cutlassB_f32_notaligned_64x64_k32_sm70(typename AttentionBackwardKernel<cut
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_f32_notaligned_64x64_k32_sm70` is for sm70-sm75, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_f32_notaligned_64x64_k32_sm70` is for sm70-sm74, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -52,7 +52,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_f32_notaligned_64x64_k32_sm75(typename AttentionBackwardKernel<cutlass::arch::Sm75, float, false, false, false, 64, 64, 32>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 750
-#if __CUDA_ARCH__ < 800
+#if __CUDA_ARCH__ <= 790
   if (!p.advance_to_block()) {
     return;
   }
@@ -61,7 +61,7 @@ fmha_cutlassB_f32_notaligned_64x64_k32_sm75(typename AttentionBackwardKernel<cut
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_f32_notaligned_64x64_k32_sm75` is for sm75-sm80, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_f32_notaligned_64x64_k32_sm75` is for sm75-sm79, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
diff --git a/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_f32_notaligned_k32_dropout.cu b/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_f32_notaligned_k32_dropout.cu
index d50ba8c5e81d0..ad23f87b49218 100644
--- a/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_f32_notaligned_k32_dropout.cu
+++ b/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_f32_notaligned_k32_dropout.cu
@@ -14,7 +14,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_f32_notaligned_64x64_k32_dropout_sm50(typename AttentionBackwardKernel<cutlass::arch::Sm50, float, false, true, false, 64, 64, 32>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 500
-#if __CUDA_ARCH__ < 700
+#if __CUDA_ARCH__ <= 690
   if (!p.advance_to_block()) {
     return;
   }
@@ -23,7 +23,7 @@ fmha_cutlassB_f32_notaligned_64x64_k32_dropout_sm50(typename AttentionBackwardKe
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_f32_notaligned_64x64_k32_dropout_sm50` is for sm50-sm70, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_f32_notaligned_64x64_k32_dropout_sm50` is for sm50-sm69, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -33,7 +33,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_f32_notaligned_64x64_k32_dropout_sm70(typename AttentionBackwardKernel<cutlass::arch::Sm70, float, false, true, false, 64, 64, 32>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 700
-#if __CUDA_ARCH__ < 750
+#if __CUDA_ARCH__ <= 740
   if (!p.advance_to_block()) {
     return;
   }
@@ -42,7 +42,7 @@ fmha_cutlassB_f32_notaligned_64x64_k32_dropout_sm70(typename AttentionBackwardKe
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_f32_notaligned_64x64_k32_dropout_sm70` is for sm70-sm75, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_f32_notaligned_64x64_k32_dropout_sm70` is for sm70-sm74, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -52,7 +52,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_f32_notaligned_64x64_k32_dropout_sm75(typename AttentionBackwardKernel<cutlass::arch::Sm75, float, false, true, false, 64, 64, 32>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 750
-#if __CUDA_ARCH__ < 800
+#if __CUDA_ARCH__ <= 790
   if (!p.advance_to_block()) {
     return;
   }
@@ -61,7 +61,7 @@ fmha_cutlassB_f32_notaligned_64x64_k32_dropout_sm75(typename AttentionBackwardKe
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_f32_notaligned_64x64_k32_dropout_sm75` is for sm75-sm80, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_f32_notaligned_64x64_k32_dropout_sm75` is for sm75-sm79, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
diff --git a/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_f32_notaligned_k64.cu b/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_f32_notaligned_k64.cu
index c2b996e2b9741..43d0430498a8f 100644
--- a/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_f32_notaligned_k64.cu
+++ b/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_f32_notaligned_k64.cu
@@ -14,7 +14,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_f32_notaligned_64x64_k64_sm50(typename AttentionBackwardKernel<cutlass::arch::Sm50, float, false, false, false, 64, 64, 64>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 500
-#if __CUDA_ARCH__ < 700
+#if __CUDA_ARCH__ <= 690
   if (!p.advance_to_block()) {
     return;
   }
@@ -23,7 +23,7 @@ fmha_cutlassB_f32_notaligned_64x64_k64_sm50(typename AttentionBackwardKernel<cut
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_f32_notaligned_64x64_k64_sm50` is for sm50-sm70, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_f32_notaligned_64x64_k64_sm50` is for sm50-sm69, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -33,7 +33,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_f32_notaligned_64x64_k64_sm70(typename AttentionBackwardKernel<cutlass::arch::Sm70, float, false, false, false, 64, 64, 64>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 700
-#if __CUDA_ARCH__ < 750
+#if __CUDA_ARCH__ <= 740
   if (!p.advance_to_block()) {
     return;
   }
@@ -42,7 +42,7 @@ fmha_cutlassB_f32_notaligned_64x64_k64_sm70(typename AttentionBackwardKernel<cut
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_f32_notaligned_64x64_k64_sm70` is for sm70-sm75, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_f32_notaligned_64x64_k64_sm70` is for sm70-sm74, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -52,7 +52,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_f32_notaligned_64x64_k64_sm75(typename AttentionBackwardKernel<cutlass::arch::Sm75, float, false, false, false, 64, 64, 64>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 750
-#if __CUDA_ARCH__ < 800
+#if __CUDA_ARCH__ <= 790
   if (!p.advance_to_block()) {
     return;
   }
@@ -61,7 +61,7 @@ fmha_cutlassB_f32_notaligned_64x64_k64_sm75(typename AttentionBackwardKernel<cut
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_f32_notaligned_64x64_k64_sm75` is for sm75-sm80, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_f32_notaligned_64x64_k64_sm75` is for sm75-sm79, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
diff --git a/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_f32_notaligned_k64_dropout.cu b/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_f32_notaligned_k64_dropout.cu
index 890cca641e5f2..28200fc2925d3 100644
--- a/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_f32_notaligned_k64_dropout.cu
+++ b/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_f32_notaligned_k64_dropout.cu
@@ -14,7 +14,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_f32_notaligned_64x64_k64_dropout_sm50(typename AttentionBackwardKernel<cutlass::arch::Sm50, float, false, true, false, 64, 64, 64>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 500
-#if __CUDA_ARCH__ < 700
+#if __CUDA_ARCH__ <= 690
   if (!p.advance_to_block()) {
     return;
   }
@@ -23,7 +23,7 @@ fmha_cutlassB_f32_notaligned_64x64_k64_dropout_sm50(typename AttentionBackwardKe
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_f32_notaligned_64x64_k64_dropout_sm50` is for sm50-sm70, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_f32_notaligned_64x64_k64_dropout_sm50` is for sm50-sm69, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -33,7 +33,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_f32_notaligned_64x64_k64_dropout_sm70(typename AttentionBackwardKernel<cutlass::arch::Sm70, float, false, true, false, 64, 64, 64>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 700
-#if __CUDA_ARCH__ < 750
+#if __CUDA_ARCH__ <= 740
   if (!p.advance_to_block()) {
     return;
   }
@@ -42,7 +42,7 @@ fmha_cutlassB_f32_notaligned_64x64_k64_dropout_sm70(typename AttentionBackwardKe
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_f32_notaligned_64x64_k64_dropout_sm70` is for sm70-sm75, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_f32_notaligned_64x64_k64_dropout_sm70` is for sm70-sm74, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -52,7 +52,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_f32_notaligned_64x64_k64_dropout_sm75(typename AttentionBackwardKernel<cutlass::arch::Sm75, float, false, true, false, 64, 64, 64>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 750
-#if __CUDA_ARCH__ < 800
+#if __CUDA_ARCH__ <= 790
   if (!p.advance_to_block()) {
     return;
   }
@@ -61,7 +61,7 @@ fmha_cutlassB_f32_notaligned_64x64_k64_dropout_sm75(typename AttentionBackwardKe
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_f32_notaligned_64x64_k64_dropout_sm75` is for sm75-sm80, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_f32_notaligned_64x64_k64_dropout_sm75` is for sm75-sm79, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
diff --git a/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_f32_notaligned_k65536.cu b/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_f32_notaligned_k65536.cu
index eb6979dc7e933..6f2608306cb9c 100644
--- a/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_f32_notaligned_k65536.cu
+++ b/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_f32_notaligned_k65536.cu
@@ -14,7 +14,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_f32_notaligned_64x64_k65536_sm50(typename AttentionBackwardKernel<cutlass::arch::Sm50, float, false, false, false, 64, 64, 65536>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 500
-#if __CUDA_ARCH__ < 700
+#if __CUDA_ARCH__ <= 690
   if (!p.advance_to_block()) {
     return;
   }
@@ -23,7 +23,7 @@ fmha_cutlassB_f32_notaligned_64x64_k65536_sm50(typename AttentionBackwardKernel<
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_f32_notaligned_64x64_k65536_sm50` is for sm50-sm70, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_f32_notaligned_64x64_k65536_sm50` is for sm50-sm69, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -33,7 +33,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_f32_notaligned_64x64_k65536_sm70(typename AttentionBackwardKernel<cutlass::arch::Sm70, float, false, false, false, 64, 64, 65536>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 700
-#if __CUDA_ARCH__ < 750
+#if __CUDA_ARCH__ <= 740
   if (!p.advance_to_block()) {
     return;
   }
@@ -42,7 +42,7 @@ fmha_cutlassB_f32_notaligned_64x64_k65536_sm70(typename AttentionBackwardKernel<
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_f32_notaligned_64x64_k65536_sm70` is for sm70-sm75, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_f32_notaligned_64x64_k65536_sm70` is for sm70-sm74, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -52,7 +52,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_f32_notaligned_64x64_k65536_sm75(typename AttentionBackwardKernel<cutlass::arch::Sm75, float, false, false, false, 64, 64, 65536>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 750
-#if __CUDA_ARCH__ < 800
+#if __CUDA_ARCH__ <= 790
   if (!p.advance_to_block()) {
     return;
   }
@@ -61,7 +61,7 @@ fmha_cutlassB_f32_notaligned_64x64_k65536_sm75(typename AttentionBackwardKernel<
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_f32_notaligned_64x64_k65536_sm75` is for sm75-sm80, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_f32_notaligned_64x64_k65536_sm75` is for sm75-sm79, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
diff --git a/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_f32_notaligned_k65536_dropout.cu b/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_f32_notaligned_k65536_dropout.cu
index eddb78c1b1738..5e3a349d5639c 100644
--- a/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_f32_notaligned_k65536_dropout.cu
+++ b/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB_f32_notaligned_k65536_dropout.cu
@@ -14,7 +14,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_f32_notaligned_64x64_k65536_dropout_sm50(typename AttentionBackwardKernel<cutlass::arch::Sm50, float, false, true, false, 64, 64, 65536>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 500
-#if __CUDA_ARCH__ < 700
+#if __CUDA_ARCH__ <= 690
   if (!p.advance_to_block()) {
     return;
   }
@@ -23,7 +23,7 @@ fmha_cutlassB_f32_notaligned_64x64_k65536_dropout_sm50(typename AttentionBackwar
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_f32_notaligned_64x64_k65536_dropout_sm50` is for sm50-sm70, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_f32_notaligned_64x64_k65536_dropout_sm50` is for sm50-sm69, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -33,7 +33,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_f32_notaligned_64x64_k65536_dropout_sm70(typename AttentionBackwardKernel<cutlass::arch::Sm70, float, false, true, false, 64, 64, 65536>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 700
-#if __CUDA_ARCH__ < 750
+#if __CUDA_ARCH__ <= 740
   if (!p.advance_to_block()) {
     return;
   }
@@ -42,7 +42,7 @@ fmha_cutlassB_f32_notaligned_64x64_k65536_dropout_sm70(typename AttentionBackwar
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_f32_notaligned_64x64_k65536_dropout_sm70` is for sm70-sm75, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_f32_notaligned_64x64_k65536_dropout_sm70` is for sm70-sm74, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -52,7 +52,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassB_f32_notaligned_64x64_k65536_dropout_sm75(typename AttentionBackwardKernel<cutlass::arch::Sm75, float, false, true, false, 64, 64, 65536>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 750
-#if __CUDA_ARCH__ < 800
+#if __CUDA_ARCH__ <= 790
   if (!p.advance_to_block()) {
     return;
   }
@@ -61,7 +61,7 @@ fmha_cutlassB_f32_notaligned_64x64_k65536_dropout_sm75(typename AttentionBackwar
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassB_f32_notaligned_64x64_k65536_dropout_sm75` is for sm75-sm80, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassB_f32_notaligned_64x64_k65536_dropout_sm75` is for sm75-sm79, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
diff --git a/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassF.h b/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassF.h
index 1d2191990fd74..53a6a3b5316fd 100644
--- a/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassF.h
+++ b/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassF.h
@@ -283,31 +283,31 @@ template <typename T> void dispatch_cutlassF_f32_sm80(T cb, int cc) {
 template <typename DT, typename T>
 void dispatch_cutlassF(T cb, int cc = 0) {
 
-    if (std::is_same_v<DT, cutlass::bfloat16_t> && 80 <= cc && cc <= 120) {
+    if (std::is_same_v<DT, cutlass::bfloat16_t> && 80 <= cc && cc <= 121) {
         dispatch_cutlassF_bf16_sm80(cb, cc);
     }
-    if (std::is_same_v<DT, cutlass::half_t> && 50 <= cc && cc < 70) {
+    if (std::is_same_v<DT, cutlass::half_t> && 50 <= cc && cc <= 69) {
         dispatch_cutlassF_f16_sm50(cb, cc);
     }
-    if (std::is_same_v<DT, cutlass::half_t> && 70 <= cc && cc < 75) {
+    if (std::is_same_v<DT, cutlass::half_t> && 70 <= cc && cc <= 74) {
         dispatch_cutlassF_f16_sm70(cb, cc);
     }
-    if (std::is_same_v<DT, cutlass::half_t> && 75 <= cc && cc < 80) {
+    if (std::is_same_v<DT, cutlass::half_t> && 75 <= cc && cc <= 79) {
         dispatch_cutlassF_f16_sm75(cb, cc);
     }
-    if (std::is_same_v<DT, cutlass::half_t> && 80 <= cc && cc <= 120) {
+    if (std::is_same_v<DT, cutlass::half_t> && 80 <= cc && cc <= 121) {
         dispatch_cutlassF_f16_sm80(cb, cc);
     }
-    if (std::is_same_v<DT, float> && 50 <= cc && cc < 70) {
+    if (std::is_same_v<DT, float> && 50 <= cc && cc <= 69) {
         dispatch_cutlassF_f32_sm50(cb, cc);
     }
-    if (std::is_same_v<DT, float> && 70 <= cc && cc < 75) {
+    if (std::is_same_v<DT, float> && 70 <= cc && cc <= 74) {
         dispatch_cutlassF_f32_sm70(cb, cc);
     }
-    if (std::is_same_v<DT, float> && 75 <= cc && cc < 80) {
+    if (std::is_same_v<DT, float> && 75 <= cc && cc <= 79) {
         dispatch_cutlassF_f32_sm75(cb, cc);
     }
-    if (std::is_same_v<DT, float> && 80 <= cc && cc <= 120) {
+    if (std::is_same_v<DT, float> && 80 <= cc && cc <= 121) {
         dispatch_cutlassF_f32_sm80(cb, cc);
     }
 }
diff --git a/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassF_bf16_aligned.cu b/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassF_bf16_aligned.cu
index 7b78fc53d2522..e7faeac9252cd 100644
--- a/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassF_bf16_aligned.cu
+++ b/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassF_bf16_aligned.cu
@@ -14,7 +14,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassF_bf16_aligned_64x64_rf_sm80(typename AttentionKernel<cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64, true, true>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 800
-#if __CUDA_ARCH__ <= 1200
+#if __CUDA_ARCH__ <= 1210
   if (!p.advance_to_block()) {
     return;
   }
@@ -23,7 +23,7 @@ fmha_cutlassF_bf16_aligned_64x64_rf_sm80(typename AttentionKernel<cutlass::bfloa
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassF_bf16_aligned_64x64_rf_sm80` is for sm80-sm100, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassF_bf16_aligned_64x64_rf_sm80` is for sm80-sm121, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -33,7 +33,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassF_bf16_aligned_64x128_rf_sm80(typename AttentionKernel<cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 128, true, true>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 800
-#if __CUDA_ARCH__ <= 1200
+#if __CUDA_ARCH__ <= 1210
   if (!p.advance_to_block()) {
     return;
   }
@@ -42,7 +42,7 @@ fmha_cutlassF_bf16_aligned_64x128_rf_sm80(typename AttentionKernel<cutlass::bflo
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassF_bf16_aligned_64x128_rf_sm80` is for sm80-sm100, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassF_bf16_aligned_64x128_rf_sm80` is for sm80-sm121, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -52,7 +52,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassF_bf16_aligned_32x128_gmem_sm80(typename AttentionKernel<cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 65536, true, true>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 800
-#if __CUDA_ARCH__ <= 1200
+#if __CUDA_ARCH__ <= 1210
   if (!p.advance_to_block()) {
     return;
   }
@@ -61,7 +61,7 @@ fmha_cutlassF_bf16_aligned_32x128_gmem_sm80(typename AttentionKernel<cutlass::bf
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassF_bf16_aligned_32x128_gmem_sm80` is for sm80-sm100, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassF_bf16_aligned_32x128_gmem_sm80` is for sm80-sm121, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
diff --git a/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassF_f16_aligned.cu b/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassF_f16_aligned.cu
index 92e5fe5a96d37..31a193da19efd 100644
--- a/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassF_f16_aligned.cu
+++ b/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassF_f16_aligned.cu
@@ -14,7 +14,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassF_f16_aligned_64x64_rf_sm50(typename AttentionKernel<cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64, true, true>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 500
-#if __CUDA_ARCH__ < 700
+#if __CUDA_ARCH__ <= 690
   if (!p.advance_to_block()) {
     return;
   }
@@ -23,7 +23,7 @@ fmha_cutlassF_f16_aligned_64x64_rf_sm50(typename AttentionKernel<cutlass::half_t
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassF_f16_aligned_64x64_rf_sm50` is for sm50-sm70, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassF_f16_aligned_64x64_rf_sm50` is for sm50-sm69, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -33,7 +33,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassF_f16_aligned_64x64_rf_sm70(typename AttentionKernel<cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64, true, true>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 700
-#if __CUDA_ARCH__ < 750
+#if __CUDA_ARCH__ <= 740
   if (!p.advance_to_block()) {
     return;
   }
@@ -42,7 +42,7 @@ fmha_cutlassF_f16_aligned_64x64_rf_sm70(typename AttentionKernel<cutlass::half_t
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassF_f16_aligned_64x64_rf_sm70` is for sm70-sm75, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassF_f16_aligned_64x64_rf_sm70` is for sm70-sm74, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -52,7 +52,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassF_f16_aligned_64x64_rf_sm75(typename AttentionKernel<cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64, true, true>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 750
-#if __CUDA_ARCH__ < 800
+#if __CUDA_ARCH__ <= 790
   if (!p.advance_to_block()) {
     return;
   }
@@ -61,7 +61,7 @@ fmha_cutlassF_f16_aligned_64x64_rf_sm75(typename AttentionKernel<cutlass::half_t
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassF_f16_aligned_64x64_rf_sm75` is for sm75-sm80, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassF_f16_aligned_64x64_rf_sm75` is for sm75-sm79, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -71,7 +71,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassF_f16_aligned_64x64_rf_sm80(typename AttentionKernel<cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64, true, true>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 800
-#if __CUDA_ARCH__ <= 1200
+#if __CUDA_ARCH__ <= 1210
   if (!p.advance_to_block()) {
     return;
   }
@@ -80,7 +80,7 @@ fmha_cutlassF_f16_aligned_64x64_rf_sm80(typename AttentionKernel<cutlass::half_t
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassF_f16_aligned_64x64_rf_sm80` is for sm80-sm100, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassF_f16_aligned_64x64_rf_sm80` is for sm80-sm121, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -90,7 +90,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassF_f16_aligned_32x128_rf_sm50(typename AttentionKernel<cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 128, true, true>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 500
-#if __CUDA_ARCH__ < 700
+#if __CUDA_ARCH__ <= 690
   if (!p.advance_to_block()) {
     return;
   }
@@ -99,7 +99,7 @@ fmha_cutlassF_f16_aligned_32x128_rf_sm50(typename AttentionKernel<cutlass::half_
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassF_f16_aligned_32x128_rf_sm50` is for sm50-sm70, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassF_f16_aligned_32x128_rf_sm50` is for sm50-sm69, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -109,7 +109,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassF_f16_aligned_32x128_rf_sm70(typename AttentionKernel<cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 128, true, true>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 700
-#if __CUDA_ARCH__ < 750
+#if __CUDA_ARCH__ <= 740
   if (!p.advance_to_block()) {
     return;
   }
@@ -118,7 +118,7 @@ fmha_cutlassF_f16_aligned_32x128_rf_sm70(typename AttentionKernel<cutlass::half_
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassF_f16_aligned_32x128_rf_sm70` is for sm70-sm75, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassF_f16_aligned_32x128_rf_sm70` is for sm70-sm74, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -128,7 +128,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassF_f16_aligned_32x128_rf_sm75(typename AttentionKernel<cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 128, true, true>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 750
-#if __CUDA_ARCH__ < 800
+#if __CUDA_ARCH__ <= 790
   if (!p.advance_to_block()) {
     return;
   }
@@ -137,7 +137,7 @@ fmha_cutlassF_f16_aligned_32x128_rf_sm75(typename AttentionKernel<cutlass::half_
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassF_f16_aligned_32x128_rf_sm75` is for sm75-sm80, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassF_f16_aligned_32x128_rf_sm75` is for sm75-sm79, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -147,7 +147,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassF_f16_aligned_64x128_rf_sm80(typename AttentionKernel<cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 128, true, true>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 800
-#if __CUDA_ARCH__ <= 1200
+#if __CUDA_ARCH__ <= 1210
   if (!p.advance_to_block()) {
     return;
   }
@@ -156,7 +156,7 @@ fmha_cutlassF_f16_aligned_64x128_rf_sm80(typename AttentionKernel<cutlass::half_
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassF_f16_aligned_64x128_rf_sm80` is for sm80-sm100, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassF_f16_aligned_64x128_rf_sm80` is for sm80-sm121, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -166,7 +166,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassF_f16_aligned_32x128_gmem_sm50(typename AttentionKernel<cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 65536, true, true>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 500
-#if __CUDA_ARCH__ < 700
+#if __CUDA_ARCH__ <= 690
   if (!p.advance_to_block()) {
     return;
   }
@@ -175,7 +175,7 @@ fmha_cutlassF_f16_aligned_32x128_gmem_sm50(typename AttentionKernel<cutlass::hal
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassF_f16_aligned_32x128_gmem_sm50` is for sm50-sm70, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassF_f16_aligned_32x128_gmem_sm50` is for sm50-sm69, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -185,7 +185,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassF_f16_aligned_32x128_gmem_sm70(typename AttentionKernel<cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 65536, true, true>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 700
-#if __CUDA_ARCH__ < 750
+#if __CUDA_ARCH__ <= 740
   if (!p.advance_to_block()) {
     return;
   }
@@ -194,7 +194,7 @@ fmha_cutlassF_f16_aligned_32x128_gmem_sm70(typename AttentionKernel<cutlass::hal
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassF_f16_aligned_32x128_gmem_sm70` is for sm70-sm75, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassF_f16_aligned_32x128_gmem_sm70` is for sm70-sm74, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -204,7 +204,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassF_f16_aligned_32x128_gmem_sm75(typename AttentionKernel<cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 65536, true, true>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 750
-#if __CUDA_ARCH__ < 800
+#if __CUDA_ARCH__ <= 790
   if (!p.advance_to_block()) {
     return;
   }
@@ -213,7 +213,7 @@ fmha_cutlassF_f16_aligned_32x128_gmem_sm75(typename AttentionKernel<cutlass::hal
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassF_f16_aligned_32x128_gmem_sm75` is for sm75-sm80, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassF_f16_aligned_32x128_gmem_sm75` is for sm75-sm79, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -223,7 +223,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassF_f16_aligned_32x128_gmem_sm80(typename AttentionKernel<cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 65536, true, true>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 800
-#if __CUDA_ARCH__ <= 1200
+#if __CUDA_ARCH__ <= 1210
   if (!p.advance_to_block()) {
     return;
   }
@@ -232,7 +232,7 @@ fmha_cutlassF_f16_aligned_32x128_gmem_sm80(typename AttentionKernel<cutlass::hal
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassF_f16_aligned_32x128_gmem_sm80` is for sm80-sm100, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassF_f16_aligned_32x128_gmem_sm80` is for sm80-sm121, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
diff --git a/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassF_f16_notaligned.cu b/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassF_f16_notaligned.cu
index a93fda7d76b2d..7531491bafd88 100644
--- a/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassF_f16_notaligned.cu
+++ b/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassF_f16_notaligned.cu
@@ -14,7 +14,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassF_f16_notaligned_64x64_rf_sm50(typename AttentionKernel<cutlass::half_t, cutlass::arch::Sm50, false, 64, 64, 64, true, true>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 500
-#if __CUDA_ARCH__ < 700
+#if __CUDA_ARCH__ <= 690
   if (!p.advance_to_block()) {
     return;
   }
@@ -23,7 +23,7 @@ fmha_cutlassF_f16_notaligned_64x64_rf_sm50(typename AttentionKernel<cutlass::hal
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassF_f16_notaligned_64x64_rf_sm50` is for sm50-sm70, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassF_f16_notaligned_64x64_rf_sm50` is for sm50-sm69, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -33,7 +33,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassF_f16_notaligned_64x64_rf_sm70(typename AttentionKernel<cutlass::half_t, cutlass::arch::Sm70, false, 64, 64, 64, true, true>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 700
-#if __CUDA_ARCH__ < 750
+#if __CUDA_ARCH__ <= 740
   if (!p.advance_to_block()) {
     return;
   }
@@ -42,7 +42,7 @@ fmha_cutlassF_f16_notaligned_64x64_rf_sm70(typename AttentionKernel<cutlass::hal
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassF_f16_notaligned_64x64_rf_sm70` is for sm70-sm75, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassF_f16_notaligned_64x64_rf_sm70` is for sm70-sm74, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -52,7 +52,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassF_f16_notaligned_64x64_rf_sm75(typename AttentionKernel<cutlass::half_t, cutlass::arch::Sm75, false, 64, 64, 64, true, true>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 750
-#if __CUDA_ARCH__ < 800
+#if __CUDA_ARCH__ <= 790
   if (!p.advance_to_block()) {
     return;
   }
@@ -61,7 +61,7 @@ fmha_cutlassF_f16_notaligned_64x64_rf_sm75(typename AttentionKernel<cutlass::hal
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassF_f16_notaligned_64x64_rf_sm75` is for sm75-sm80, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassF_f16_notaligned_64x64_rf_sm75` is for sm75-sm79, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -71,7 +71,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassF_f16_notaligned_32x128_rf_sm50(typename AttentionKernel<cutlass::half_t, cutlass::arch::Sm50, false, 32, 128, 128, true, true>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 500
-#if __CUDA_ARCH__ < 700
+#if __CUDA_ARCH__ <= 690
   if (!p.advance_to_block()) {
     return;
   }
@@ -80,7 +80,7 @@ fmha_cutlassF_f16_notaligned_32x128_rf_sm50(typename AttentionKernel<cutlass::ha
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassF_f16_notaligned_32x128_rf_sm50` is for sm50-sm70, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassF_f16_notaligned_32x128_rf_sm50` is for sm50-sm69, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -90,7 +90,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassF_f16_notaligned_32x128_rf_sm70(typename AttentionKernel<cutlass::half_t, cutlass::arch::Sm70, false, 32, 128, 128, true, true>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 700
-#if __CUDA_ARCH__ < 750
+#if __CUDA_ARCH__ <= 740
   if (!p.advance_to_block()) {
     return;
   }
@@ -99,7 +99,7 @@ fmha_cutlassF_f16_notaligned_32x128_rf_sm70(typename AttentionKernel<cutlass::ha
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassF_f16_notaligned_32x128_rf_sm70` is for sm70-sm75, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassF_f16_notaligned_32x128_rf_sm70` is for sm70-sm74, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -109,7 +109,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassF_f16_notaligned_32x128_rf_sm75(typename AttentionKernel<cutlass::half_t, cutlass::arch::Sm75, false, 32, 128, 128, true, true>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 750
-#if __CUDA_ARCH__ < 800
+#if __CUDA_ARCH__ <= 790
   if (!p.advance_to_block()) {
     return;
   }
@@ -118,7 +118,7 @@ fmha_cutlassF_f16_notaligned_32x128_rf_sm75(typename AttentionKernel<cutlass::ha
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassF_f16_notaligned_32x128_rf_sm75` is for sm75-sm80, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassF_f16_notaligned_32x128_rf_sm75` is for sm75-sm79, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -128,7 +128,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassF_f16_notaligned_32x128_gmem_sm50(typename AttentionKernel<cutlass::half_t, cutlass::arch::Sm50, false, 32, 128, 65536, true, true>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 500
-#if __CUDA_ARCH__ < 700
+#if __CUDA_ARCH__ <= 690
   if (!p.advance_to_block()) {
     return;
   }
@@ -137,7 +137,7 @@ fmha_cutlassF_f16_notaligned_32x128_gmem_sm50(typename AttentionKernel<cutlass::
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassF_f16_notaligned_32x128_gmem_sm50` is for sm50-sm70, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassF_f16_notaligned_32x128_gmem_sm50` is for sm50-sm69, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -147,7 +147,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassF_f16_notaligned_32x128_gmem_sm70(typename AttentionKernel<cutlass::half_t, cutlass::arch::Sm70, false, 32, 128, 65536, true, true>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 700
-#if __CUDA_ARCH__ < 750
+#if __CUDA_ARCH__ <= 740
   if (!p.advance_to_block()) {
     return;
   }
@@ -156,7 +156,7 @@ fmha_cutlassF_f16_notaligned_32x128_gmem_sm70(typename AttentionKernel<cutlass::
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassF_f16_notaligned_32x128_gmem_sm70` is for sm70-sm75, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassF_f16_notaligned_32x128_gmem_sm70` is for sm70-sm74, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -166,7 +166,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassF_f16_notaligned_32x128_gmem_sm75(typename AttentionKernel<cutlass::half_t, cutlass::arch::Sm75, false, 32, 128, 65536, true, true>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 750
-#if __CUDA_ARCH__ < 800
+#if __CUDA_ARCH__ <= 790
   if (!p.advance_to_block()) {
     return;
   }
@@ -175,7 +175,7 @@ fmha_cutlassF_f16_notaligned_32x128_gmem_sm75(typename AttentionKernel<cutlass::
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassF_f16_notaligned_32x128_gmem_sm75` is for sm75-sm80, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassF_f16_notaligned_32x128_gmem_sm75` is for sm75-sm79, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
diff --git a/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassF_f32_aligned.cu b/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassF_f32_aligned.cu
index 5857194f76c9e..fcd006d8c8499 100644
--- a/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassF_f32_aligned.cu
+++ b/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassF_f32_aligned.cu
@@ -14,7 +14,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassF_f32_aligned_64x64_rf_sm50(typename AttentionKernel<float, cutlass::arch::Sm50, true, 64, 64, 64, true, true>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 500
-#if __CUDA_ARCH__ < 700
+#if __CUDA_ARCH__ <= 690
   if (!p.advance_to_block()) {
     return;
   }
@@ -23,7 +23,7 @@ fmha_cutlassF_f32_aligned_64x64_rf_sm50(typename AttentionKernel<float, cutlass:
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassF_f32_aligned_64x64_rf_sm50` is for sm50-sm70, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassF_f32_aligned_64x64_rf_sm50` is for sm50-sm69, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -33,7 +33,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassF_f32_aligned_64x64_rf_sm70(typename AttentionKernel<float, cutlass::arch::Sm70, true, 64, 64, 64, true, true>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 700
-#if __CUDA_ARCH__ < 750
+#if __CUDA_ARCH__ <= 740
   if (!p.advance_to_block()) {
     return;
   }
@@ -42,7 +42,7 @@ fmha_cutlassF_f32_aligned_64x64_rf_sm70(typename AttentionKernel<float, cutlass:
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassF_f32_aligned_64x64_rf_sm70` is for sm70-sm75, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassF_f32_aligned_64x64_rf_sm70` is for sm70-sm74, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -52,7 +52,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassF_f32_aligned_64x64_rf_sm75(typename AttentionKernel<float, cutlass::arch::Sm75, true, 64, 64, 64, true, true>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 750
-#if __CUDA_ARCH__ < 800
+#if __CUDA_ARCH__ <= 790
   if (!p.advance_to_block()) {
     return;
   }
@@ -61,7 +61,7 @@ fmha_cutlassF_f32_aligned_64x64_rf_sm75(typename AttentionKernel<float, cutlass:
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassF_f32_aligned_64x64_rf_sm75` is for sm75-sm80, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassF_f32_aligned_64x64_rf_sm75` is for sm75-sm79, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -71,7 +71,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassF_f32_aligned_64x64_rf_sm80(typename AttentionKernel<float, cutlass::arch::Sm80, true, 64, 64, 64, true, true>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 800
-#if __CUDA_ARCH__ <= 1200
+#if __CUDA_ARCH__ <= 1210
   if (!p.advance_to_block()) {
     return;
   }
@@ -80,7 +80,7 @@ fmha_cutlassF_f32_aligned_64x64_rf_sm80(typename AttentionKernel<float, cutlass:
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassF_f32_aligned_64x64_rf_sm80` is for sm80-sm100, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassF_f32_aligned_64x64_rf_sm80` is for sm80-sm121, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -90,7 +90,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassF_f32_aligned_32x128_rf_sm50(typename AttentionKernel<float, cutlass::arch::Sm50, true, 32, 128, 128, true, true>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 500
-#if __CUDA_ARCH__ < 700
+#if __CUDA_ARCH__ <= 690
   if (!p.advance_to_block()) {
     return;
   }
@@ -99,7 +99,7 @@ fmha_cutlassF_f32_aligned_32x128_rf_sm50(typename AttentionKernel<float, cutlass
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassF_f32_aligned_32x128_rf_sm50` is for sm50-sm70, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassF_f32_aligned_32x128_rf_sm50` is for sm50-sm69, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -109,7 +109,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassF_f32_aligned_32x128_rf_sm70(typename AttentionKernel<float, cutlass::arch::Sm70, true, 32, 128, 128, true, true>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 700
-#if __CUDA_ARCH__ < 750
+#if __CUDA_ARCH__ <= 740
   if (!p.advance_to_block()) {
     return;
   }
@@ -118,7 +118,7 @@ fmha_cutlassF_f32_aligned_32x128_rf_sm70(typename AttentionKernel<float, cutlass
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassF_f32_aligned_32x128_rf_sm70` is for sm70-sm75, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassF_f32_aligned_32x128_rf_sm70` is for sm70-sm74, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -128,7 +128,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassF_f32_aligned_32x128_rf_sm75(typename AttentionKernel<float, cutlass::arch::Sm75, true, 32, 128, 128, true, true>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 750
-#if __CUDA_ARCH__ < 800
+#if __CUDA_ARCH__ <= 790
   if (!p.advance_to_block()) {
     return;
   }
@@ -137,7 +137,7 @@ fmha_cutlassF_f32_aligned_32x128_rf_sm75(typename AttentionKernel<float, cutlass
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassF_f32_aligned_32x128_rf_sm75` is for sm75-sm80, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassF_f32_aligned_32x128_rf_sm75` is for sm75-sm79, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -147,7 +147,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassF_f32_aligned_64x128_rf_sm80(typename AttentionKernel<float, cutlass::arch::Sm80, true, 64, 128, 128, true, true>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 800
-#if __CUDA_ARCH__ <= 1200
+#if __CUDA_ARCH__ <= 1210
   if (!p.advance_to_block()) {
     return;
   }
@@ -156,7 +156,7 @@ fmha_cutlassF_f32_aligned_64x128_rf_sm80(typename AttentionKernel<float, cutlass
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassF_f32_aligned_64x128_rf_sm80` is for sm80-sm100, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassF_f32_aligned_64x128_rf_sm80` is for sm80-sm121, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -166,7 +166,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassF_f32_aligned_32x128_gmem_sm50(typename AttentionKernel<float, cutlass::arch::Sm50, true, 32, 128, 65536, true, true>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 500
-#if __CUDA_ARCH__ < 700
+#if __CUDA_ARCH__ <= 690
   if (!p.advance_to_block()) {
     return;
   }
@@ -175,7 +175,7 @@ fmha_cutlassF_f32_aligned_32x128_gmem_sm50(typename AttentionKernel<float, cutla
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassF_f32_aligned_32x128_gmem_sm50` is for sm50-sm70, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassF_f32_aligned_32x128_gmem_sm50` is for sm50-sm69, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -185,7 +185,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassF_f32_aligned_32x128_gmem_sm70(typename AttentionKernel<float, cutlass::arch::Sm70, true, 32, 128, 65536, true, true>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 700
-#if __CUDA_ARCH__ < 750
+#if __CUDA_ARCH__ <= 740
   if (!p.advance_to_block()) {
     return;
   }
@@ -194,7 +194,7 @@ fmha_cutlassF_f32_aligned_32x128_gmem_sm70(typename AttentionKernel<float, cutla
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassF_f32_aligned_32x128_gmem_sm70` is for sm70-sm75, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassF_f32_aligned_32x128_gmem_sm70` is for sm70-sm74, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -204,7 +204,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassF_f32_aligned_32x128_gmem_sm75(typename AttentionKernel<float, cutlass::arch::Sm75, true, 32, 128, 65536, true, true>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 750
-#if __CUDA_ARCH__ < 800
+#if __CUDA_ARCH__ <= 790
   if (!p.advance_to_block()) {
     return;
   }
@@ -213,7 +213,7 @@ fmha_cutlassF_f32_aligned_32x128_gmem_sm75(typename AttentionKernel<float, cutla
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassF_f32_aligned_32x128_gmem_sm75` is for sm75-sm80, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassF_f32_aligned_32x128_gmem_sm75` is for sm75-sm79, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -223,7 +223,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassF_f32_aligned_32x128_gmem_sm80(typename AttentionKernel<float, cutlass::arch::Sm80, true, 32, 128, 65536, true, true>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 800
-#if __CUDA_ARCH__ <= 1200
+#if __CUDA_ARCH__ <= 1210
   if (!p.advance_to_block()) {
     return;
   }
@@ -232,7 +232,7 @@ fmha_cutlassF_f32_aligned_32x128_gmem_sm80(typename AttentionKernel<float, cutla
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassF_f32_aligned_32x128_gmem_sm80` is for sm80-sm100, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassF_f32_aligned_32x128_gmem_sm80` is for sm80-sm121, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
diff --git a/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassF_f32_notaligned.cu b/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassF_f32_notaligned.cu
index 8c13e7560247d..630027a777d53 100644
--- a/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassF_f32_notaligned.cu
+++ b/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassF_f32_notaligned.cu
@@ -14,7 +14,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassF_f32_notaligned_64x64_rf_sm50(typename AttentionKernel<float, cutlass::arch::Sm50, false, 64, 64, 64, true, true>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 500
-#if __CUDA_ARCH__ < 700
+#if __CUDA_ARCH__ <= 690
   if (!p.advance_to_block()) {
     return;
   }
@@ -23,7 +23,7 @@ fmha_cutlassF_f32_notaligned_64x64_rf_sm50(typename AttentionKernel<float, cutla
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassF_f32_notaligned_64x64_rf_sm50` is for sm50-sm70, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassF_f32_notaligned_64x64_rf_sm50` is for sm50-sm69, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -33,7 +33,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassF_f32_notaligned_64x64_rf_sm70(typename AttentionKernel<float, cutlass::arch::Sm70, false, 64, 64, 64, true, true>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 700
-#if __CUDA_ARCH__ < 750
+#if __CUDA_ARCH__ <= 740
   if (!p.advance_to_block()) {
     return;
   }
@@ -42,7 +42,7 @@ fmha_cutlassF_f32_notaligned_64x64_rf_sm70(typename AttentionKernel<float, cutla
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassF_f32_notaligned_64x64_rf_sm70` is for sm70-sm75, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassF_f32_notaligned_64x64_rf_sm70` is for sm70-sm74, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -52,7 +52,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassF_f32_notaligned_64x64_rf_sm75(typename AttentionKernel<float, cutlass::arch::Sm75, false, 64, 64, 64, true, true>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 750
-#if __CUDA_ARCH__ < 800
+#if __CUDA_ARCH__ <= 790
   if (!p.advance_to_block()) {
     return;
   }
@@ -61,7 +61,7 @@ fmha_cutlassF_f32_notaligned_64x64_rf_sm75(typename AttentionKernel<float, cutla
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassF_f32_notaligned_64x64_rf_sm75` is for sm75-sm80, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassF_f32_notaligned_64x64_rf_sm75` is for sm75-sm79, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -71,7 +71,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassF_f32_notaligned_32x128_rf_sm50(typename AttentionKernel<float, cutlass::arch::Sm50, false, 32, 128, 128, true, true>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 500
-#if __CUDA_ARCH__ < 700
+#if __CUDA_ARCH__ <= 690
   if (!p.advance_to_block()) {
     return;
   }
@@ -80,7 +80,7 @@ fmha_cutlassF_f32_notaligned_32x128_rf_sm50(typename AttentionKernel<float, cutl
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassF_f32_notaligned_32x128_rf_sm50` is for sm50-sm70, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassF_f32_notaligned_32x128_rf_sm50` is for sm50-sm69, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -90,7 +90,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassF_f32_notaligned_32x128_rf_sm70(typename AttentionKernel<float, cutlass::arch::Sm70, false, 32, 128, 128, true, true>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 700
-#if __CUDA_ARCH__ < 750
+#if __CUDA_ARCH__ <= 740
   if (!p.advance_to_block()) {
     return;
   }
@@ -99,7 +99,7 @@ fmha_cutlassF_f32_notaligned_32x128_rf_sm70(typename AttentionKernel<float, cutl
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassF_f32_notaligned_32x128_rf_sm70` is for sm70-sm75, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassF_f32_notaligned_32x128_rf_sm70` is for sm70-sm74, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -109,7 +109,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassF_f32_notaligned_32x128_rf_sm75(typename AttentionKernel<float, cutlass::arch::Sm75, false, 32, 128, 128, true, true>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 750
-#if __CUDA_ARCH__ < 800
+#if __CUDA_ARCH__ <= 790
   if (!p.advance_to_block()) {
     return;
   }
@@ -118,7 +118,7 @@ fmha_cutlassF_f32_notaligned_32x128_rf_sm75(typename AttentionKernel<float, cutl
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassF_f32_notaligned_32x128_rf_sm75` is for sm75-sm80, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassF_f32_notaligned_32x128_rf_sm75` is for sm75-sm79, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -128,7 +128,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassF_f32_notaligned_32x128_gmem_sm50(typename AttentionKernel<float, cutlass::arch::Sm50, false, 32, 128, 65536, true, true>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 500
-#if __CUDA_ARCH__ < 700
+#if __CUDA_ARCH__ <= 690
   if (!p.advance_to_block()) {
     return;
   }
@@ -137,7 +137,7 @@ fmha_cutlassF_f32_notaligned_32x128_gmem_sm50(typename AttentionKernel<float, cu
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassF_f32_notaligned_32x128_gmem_sm50` is for sm50-sm70, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassF_f32_notaligned_32x128_gmem_sm50` is for sm50-sm69, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -147,7 +147,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassF_f32_notaligned_32x128_gmem_sm70(typename AttentionKernel<float, cutlass::arch::Sm70, false, 32, 128, 65536, true, true>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 700
-#if __CUDA_ARCH__ < 750
+#if __CUDA_ARCH__ <= 740
   if (!p.advance_to_block()) {
     return;
   }
@@ -156,7 +156,7 @@ fmha_cutlassF_f32_notaligned_32x128_gmem_sm70(typename AttentionKernel<float, cu
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassF_f32_notaligned_32x128_gmem_sm70` is for sm70-sm75, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassF_f32_notaligned_32x128_gmem_sm70` is for sm70-sm74, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
@@ -166,7 +166,7 @@ __global__ void __launch_bounds__(
 fmha_cutlassF_f32_notaligned_32x128_gmem_sm75(typename AttentionKernel<float, cutlass::arch::Sm75, false, 32, 128, 65536, true, true>::Params p) {
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= 750
-#if __CUDA_ARCH__ < 800
+#if __CUDA_ARCH__ <= 790
   if (!p.advance_to_block()) {
     return;
   }
@@ -175,7 +175,7 @@ fmha_cutlassF_f32_notaligned_32x128_gmem_sm75(typename AttentionKernel<float, cu
 #endif
 #endif
     printf(
-        "FATAL: kernel `fmha_cutlassF_f32_notaligned_32x128_gmem_sm75` is for sm75-sm80, but was built for sm%d\n",
+        "FATAL: kernel `fmha_cutlassF_f32_notaligned_32x128_gmem_sm75` is for sm75-sm79, but was built for sm%d\n",
         int(__CUDA_ARCH__ + 0) / 10);
 #endif
 }
diff --git a/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/generate_kernels.py b/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/generate_kernels.py
index 7b83617f643d2..bf249985c6cf0 100644
--- a/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/generate_kernels.py
+++ b/aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/generate_kernels.py
@@ -13,7 +13,7 @@
 import itertools
 from dataclasses import dataclass, field
 from pathlib import Path
-from typing import Optional, TypeVar
+from typing import TypeVar
 
 
 DTYPES = {
@@ -22,7 +22,7 @@
     "bf16": "cutlass::bfloat16_t",
 }
 
-SM = [50, 70, 75, 80, 100]  # Sm80 kernels support up to Sm100
+SM_RANGES: list[tuple[int, int]] = [(50, 69), (70, 74), (75, 79), (80, 121)]
 
 KERNEL_IMPL_TEMPLATE = """__global__ void __launch_bounds__(
     {CPP_CLASS}::kNumThreads,
@@ -30,7 +30,7 @@
 {NAME}(typename {CPP_CLASS}::Params p) {{
 #ifdef __CUDA_ARCH__
 #if __CUDA_ARCH__ >= {SM}0
-#if __CUDA_ARCH__ < {SM_MAX}0
+#if __CUDA_ARCH__ <= {SM_MAX}0
   if (!p.advance_to_block()) {{
     return;
   }}
@@ -57,7 +57,7 @@ class FwdKernel:
     max_k: int
     supports_dropout: bool = True
     supports_bias: bool = True
-    dispatch_cond: Optional[str] = None
+    dispatch_cond: str | None = None
 
     def __post_init__(self) -> None:
         # Set kernel selection priority
@@ -117,7 +117,7 @@ def cpp_impl(self) -> str:
     def get_all(cls) -> list["FwdKernel"]:
         kernels: list[FwdKernel] = []
         for aligned, dtype, (sm, sm_max) in itertools.product(
-            [True, False], DTYPES.keys(), itertools.pairwise(SM)
+            [True, False], DTYPES.keys(), SM_RANGES
         ):
             # Remove some kernels we don't use
             if dtype == "bf16" and sm < 80:
@@ -154,7 +154,7 @@ class BwdKernel:
     block_i: int
     block_j: int
     max_k: int
-    dispatch_cond: Optional[str] = None
+    dispatch_cond: str | None = None
     keys_queries_aligned_to_blocksizes: bool = False
 
     def __post_init__(self) -> None:
@@ -228,7 +228,7 @@ def get_all(cls) -> list["BwdKernel"]:
         for aligned, dtype, (sm, sm_max), apply_dropout, max_k in itertools.product(
             [True, False],
             DTYPES.keys(),
-            itertools.pairwise(SM),
+            SM_RANGES,
             [True, False],
             [32, 64, 128, 2**16],
         ):
@@ -282,12 +282,13 @@ def get_all(cls) -> list["BwdKernel"]:
         # Add some specialized kernels for stable diffusion BW (K=80)
         # This is the only kernel that can keep the outputs on RF on
         # Sm86/Sm89, so it's much faster than the 64x64 one
+        sm80_range = next(sm_range for sm_range in SM_RANGES if sm_range[0] == 80)
         for dtype in ["f16", "bf16"]:
             kernels.append(
                 cls(
                     aligned=True,
                     dtype=dtype,
-                    sm_range=(80, SM[SM.index(80) + 1]),
+                    sm_range=sm80_range,
                     apply_dropout=False,
                     preload_mmas=True,
                     block_i=128,
@@ -308,7 +309,7 @@ def write_decl_impl(
     family_name: str,
     impl_file: str,
     autogen_dir: Path,
-    disable_def: Optional[str] = None,
+    disable_def: str | None = None,
 ) -> None:
     cpp_file_header = """/*
  * Copyright (c) Meta Platforms, Inc. and affiliates.
@@ -352,7 +353,7 @@ def write_decl_impl(
             declarations += f"    {_call}"
         declarations += "}\n\n"
         dispatch_all += f"""
-    if (std::is_same_v<DT, {DTYPES[cat_dt]}> && {cat_sm} <= cc && cc < {cat_sm_max}) {{
+    if (std::is_same_v<DT, {DTYPES[cat_dt]}> && {cat_sm} <= cc && cc <= {cat_sm_max}) {{
         {dispatch_category_fn}(cb, cc);
     }}"""
 
@@ -378,7 +379,7 @@ def write_decl_impl(
         (autogen_dir / f"{family_name}_{f}.cu").write_text(impl_cu)
 
 
-def main(output_dir: Optional[str]) -> None:
+def main(output_dir: str | None) -> None:
     if output_dir is None:
         output_dir = Path(__file__).parent
     else:
diff --git a/aten/src/ATen/native/transformers/cuda/sdp_utils.cpp b/aten/src/ATen/native/transformers/cuda/sdp_utils.cpp
index 67a5a296e2afe..0b5159cca5a37 100644
--- a/aten/src/ATen/native/transformers/cuda/sdp_utils.cpp
+++ b/aten/src/ATen/native/transformers/cuda/sdp_utils.cpp
@@ -140,24 +140,19 @@ int64_t minimum_gemm_alignment(sdp_params const& params) {
   return matmul_alignment_mn;
 }
 
-// On ROCM, ME and FA share the backend, and hence they share the checking
-// function for fundamental limitations by the GPU kernel
-// caller_is_meff is added to make the TORCH_WARN message showing the correct result
-template<bool caller_is_meff = false>
-bool check_head_dim_size_flash(sdp_params const& params, bool debug) {
 #if USE_ROCM_ATTENTION
-  if (at::cuda::device_count() == 0) {
-    return false;
-  }
-  // AOTriton 0.9+ supports head_dim up to 512
-  const static auto max_hdim = []() {
+inline int aotriton_max_hdim() {
+  static const int max_hdim = []() {
 #if AOTRITON_VERSION_CURRENT == AOTRITON_VERSION_INT(0, 11)
-    // gfx11xx only support hdim <= 256 on AOTriton 0.11
+    // gfx11xx only support hdim <= 256 on AOTriton 0.11/0.12
     auto dprops = at::cuda::getCurrentDeviceProperties();
     const c10::basic_string_view<char> arch(dprops->gcnArchName);
     if (arch.starts_with("gfx11")) {
       return 256;
     }
+    if (arch.starts_with("gfx1250")) {
+      return 256;
+    }
 #endif // AOTriton 0.11
 #if AOTRITON_VERSION_CURRENT >= AOTRITON_VERSION_INT(0, 9)
     return 512;
@@ -165,7 +160,28 @@ bool check_head_dim_size_flash(sdp_params const& params, bool debug) {
     return 256;
 #endif
   }();
-  const auto max_size = c10::SymInt(max_hdim);
+  return max_hdim;
+}
+#endif // USE_ROCM_ATTENTION
+
+// For AOTriton <= 0.11:
+// On ROCM, ME and FA share the backend, and hence they share the checking
+// function for fundamental limitations by the GPU kernel
+// caller_is_meff is added to make the TORCH_WARN message showing the correct result
+//
+// FIXME: revert this reuse when removing AOTriton <= 0.11 support
+//
+// AOTriton 0.12 supports hdim_qk != hdim_vo, but we cannot enable this in
+// check_head_dim_size_flash because it changes the backend selection logic for
+// FA, which can break certain workloads that rely on the behavior of rejecting
+// FA for hdim_qk != hdim_vo
+template<bool caller_is_meff = false>
+bool check_head_dim_size_flash(sdp_params const& params, bool debug) {
+#if USE_ROCM_ATTENTION
+  if (at::cuda::device_count() == 0) {
+    return false;
+  }
+  const auto max_size = c10::SymInt(aotriton_max_hdim());
 #else
   // All head_dim sizes must be equal and less than 256
   const auto max_size = c10::SymInt(256);
@@ -245,9 +261,20 @@ bool check_head_dim_size_flash_nested(sdp_params const& params, bool debug) {
 }
 
 bool check_head_dim_size_mem_efficient(sdp_params const& params, bool debug) {
+#if USE_ROCM_ATTENTION
+#if AOTRITON_VERSION_CURRENT < AOTRITON_VERSION_INT(0, 12)
+  return check_head_dim_size_flash_nested<true /* caller_is_meff */>(params, debug);
+#endif
+#endif
   const auto query_size_last = params.query.sym_size(-1);
   const auto value_size_last = params.value.sym_size(-1);
+#ifdef USE_ROCM
+  bool is_half = (params.query.dtype() == at::kHalf) ||
+      (params.query.dtype() == at::kBFloat16);
+  const int64_t alignment = is_half ? 8 : 4;
+#else
   const int64_t alignment = minimum_gemm_alignment(params);
+#endif
   if (!(query_size_last == params.key.sym_size(-1) &&
         query_size_last % alignment == 0 && query_size_last > 0 &&
         value_size_last % alignment == 0 && value_size_last > 0)) {
@@ -266,6 +293,27 @@ bool check_head_dim_size_mem_efficient(sdp_params const& params, bool debug) {
     }
     return false;
   }
+#if USE_ROCM_ATTENTION
+#if AOTRITON_VERSION_CURRENT >= AOTRITON_VERSION_INT(0, 12)
+  const auto max_size = c10::SymInt(aotriton_max_hdim());
+  if (!(query_size_last <= max_size && value_size_last <= max_size)) {
+    if (debug) {
+      TORCH_WARN(
+          "Mem efficient attention on ROCM requires last dimension of inputs to less or equal than ",
+          max_size,
+          ". ",
+          "Got Query.size(-1): ",
+          query_size_last,
+          ", Key.size(-1): ",
+          params.key.sym_size(-1),
+          ", Value.size(-1): ",
+          params.value.sym_size(-1),
+          " instead. (Note this limit differs among architectures)");
+    }
+    return false;
+  }
+#endif // AOTRITON_VERSION_CURRENT >= AOTRITON_VERSION_INT(0, 12)
+#endif // USE_ROCM_ATTENTION
   return true;
 }
 
@@ -483,6 +531,18 @@ bool check_all_tensors_on_device(sdp_params const& params, bool debug) {
   return true;
 }
 
+bool check_cudnn_dropout(sdp_params const& params, bool debug) {
+  if (params.dropout * 16.0 != std::floor(params.dropout * 16.0)) {
+    if (debug) {
+      TORCH_WARN("cuDNN dropout probability resolution is limited to 1/16."
+                 "Use a dropout probability that is a multiple of 1/16 to "
+                 "select the cuDNN backend");
+    }
+    return false;
+  }
+  return true;
+}
+
 bool check_cudnn_tensor_shapes(sdp_params const& params, bool debug) {
   const auto s_q = params.query.sym_size(2);
   const auto s_k = params.key.sym_size(2);
@@ -758,7 +818,8 @@ bool can_use_cudnn_attention(const sdp_params& params, bool debug) {
           check_cudnn_deterministic,
           check_dtypes_low_precision,
           check_attn_mask_shape,
-          check_cudnn_hardware_support
+          check_cudnn_hardware_support,
+          check_cudnn_dropout
           );
   for (auto& constraint : general_constraints) {
     if (!constraint(params, debug)) {
@@ -867,11 +928,7 @@ bool can_use_mem_efficient_attention(sdp_params const& params, bool debug) {
       check_all_tensors_on_device,
       check_mem_efficient_hardware_support,
       check_tensor_shapes,
-#ifdef USE_ROCM
-      check_head_dim_size_flash<true /* caller_is_meff */>
-#else
       check_head_dim_size_mem_efficient
-#endif
   );
   for (auto& constraint : general_constraints) {
     if (!constraint(params, debug)) {
@@ -881,10 +938,8 @@ bool can_use_mem_efficient_attention(sdp_params const& params, bool debug) {
 
   if (has_for_nested_inputs(params)) {
     constexpr auto nested_constraints = c10::array_of<bool (*)(sdp_params const&, bool)>(
-#ifndef USE_ROCM  // ME and FA shares backend on ROCM and thus supports training
+#ifndef USE_ROCM  // ME and FA share the backend on ROCM and thus support training
         check_requires_grad_and_nested,
-#else // Meanwhile ME on ROCM share the limits of FA about head dimensions
-        check_head_dim_size_flash_nested<true /* caller_is_meff */>,
 #endif
         check_batch_size_nested,
         check_for_seq_len_0_nested_tensor);
@@ -924,8 +979,8 @@ bool can_use_mem_efficient_attention(sdp_params const& params, bool debug) {
   if (dprop->major >= 8) {
     return check_tensor_dtype(params, greater_than_or_equal_sm80_mem_efficient_dtypes, debug);
   }
-#endif
   return check_tensor_dtype(params, less_than_sm80_mem_efficient_dtypes, debug);
+#endif
 }
 
 SDPBackend select_sdp_backend(sdp_params const& kernel_params) {
diff --git a/aten/src/ATen/native/transformers/hip/aotriton_adapter.h b/aten/src/ATen/native/transformers/hip/aotriton_adapter.h
index e40376ae0c3a7..43e9c616e35fa 100644
--- a/aten/src/ATen/native/transformers/hip/aotriton_adapter.h
+++ b/aten/src/ATen/native/transformers/hip/aotriton_adapter.h
@@ -11,6 +11,10 @@
 #include <tuple>
 #include <optional>
 
+#if AOTRITON_VERSION_CURRENT >= AOTRITON_VERSION_INT(0, 12)
+#define AOTRITON_V2_API_FLASH_ATTN_H  // Suppress the include of deprecated flash/v2.h
+#endif
+
 ////////////////////////////////////////////////////////////////////////////////
 // Common macros copied from cuda/mem_eff_attention/gemm_kernel_utils.h
 ////////////////////////////////////////////////////////////////////////////////
@@ -127,8 +131,17 @@ struct LazyTensorContext {
 
 template<int kRank, bool kRequireZeros>
 struct LazyTensorFunctions : public LazyTensorContext {
-  static aotriton::TensorView<kRank> acquire(void* cookie) {
+#if AOTRITON_VERSION_CURRENT >= AOTRITON_VERSION_INT(0, 12)
+  using HolderType = aotriton::LazyTensor<kRank>;
+#else
+  using HolderType = void;
+#endif
+  static aotriton::TensorView<kRank> acquire(HolderType* self) {
+#if AOTRITON_VERSION_CURRENT >= AOTRITON_VERSION_INT(0, 12)
+    auto ctx = (LazyTensorContext*)self->cookie;
+#else
     auto ctx = (LazyTensorContext*)cookie;
+#endif
     if (!ctx->tensor.defined()) {
       auto q = ctx->like_tensor;
       if constexpr (kRequireZeros) {
@@ -141,7 +154,7 @@ struct LazyTensorFunctions : public LazyTensorContext {
     return mk_aotensor<kRank>(ctx->tensor, ctx->tensor_name);
   }
 
-  static void dispose(void* cookie) {
+  static void dispose(HolderType* cookie) {
   }
 };
 
diff --git a/aten/src/ATen/native/transformers/hip/aotriton_versions.h b/aten/src/ATen/native/transformers/hip/aotriton_versions.h
index 2f5d3f0e12228..b95731119e848 100644
--- a/aten/src/ATen/native/transformers/hip/aotriton_versions.h
+++ b/aten/src/ATen/native/transformers/hip/aotriton_versions.h
@@ -17,4 +17,10 @@
 #define AOTRITON_V3_API 0
 #endif
 
+#if AOTRITON_VERSION_CURRENT >= AOTRITON_VERSION_INT(0, 12)
+#define AOTRITON_COMPACT_VARLEN_LSE 1
+#else
+#define AOTRITON_COMPACT_VARLEN_LSE 1
+#endif
+
 #endif
diff --git a/aten/src/ATen/native/transformers/hip/flash_attn/aot/mha_all_aot.hip b/aten/src/ATen/native/transformers/hip/flash_attn/aot/mha_all_aot.hip
index b96e80d5e5a9e..c16f7d1aad233 100644
--- a/aten/src/ATen/native/transformers/hip/flash_attn/aot/mha_all_aot.hip
+++ b/aten/src/ATen/native/transformers/hip/flash_attn/aot/mha_all_aot.hip
@@ -257,7 +257,6 @@ mha_fwd_aot(const at::Tensor &q,         // batch_size x seqlen_q x num_heads x
   }
 
   hipError_t err; // TODO: Error handling
-  using aotriton::v2::flash::attn_fwd;
   using sdp::aotriton_adapter::mk_aotensor;
   using sdp::aotriton_adapter::mk_aoscalartensor;
   using sdp::aotriton_adapter::mk_philoxtensor;
@@ -270,56 +269,42 @@ mha_fwd_aot(const at::Tensor &q,         // batch_size x seqlen_q x num_heads x
   auto seed_output = mk_philoxtensor(use_philox_state ? seed_t.data_ptr<int64_t>() : nullptr);
   auto offset_output = mk_philoxtensor(use_philox_state ? offset_t.data_ptr<int64_t>() : nullptr);
   auto persistent_counter = mk_atomictensor(is_causal ? atomic_counter.data_ptr<int32_t>() : nullptr);
-  if (uses_swa || AOTRITON_ALWAYS_V3_API) {
-#if AOTRITON_V3_API
-    using aotriton::v3::flash::CausalType;
-    using aotriton::v3::flash::VarlenType;
-    aotriton::v3::flash::attn_fwd_params params;
-    params.Q = mk_aotensor(q_t, "q");
-    params.K = mk_aotensor(k_t, "k");
-    params.V = mk_aotensor(v_t, "v");
-    params.Sm_scale = softmax_scale;
-    params.L = mk_aotensor<2>(M, "M");
-    params.Out = mk_aotensor(output_t, "Out");
-    params.Max_seqlen_q = seqlen_q;    // Unused if cu_seqlens_q is empty
-    params.Max_seqlen_k = seqlen_k;    // Unused if cu_seqlens_k is empty
-    params.dropout_p = p_dropout;
-    params.philox_seed_ptr = seed;
-    params.philox_offset1 = offset1;
-    params.philox_offset2 = offset2;
-    params.philox_seed_output = seed_output;
-    params.philox_offset_output = offset_output;
-    params.encoded_softmax = mk_aotensor(softmax_fa_t, "encoded_softmax");
-    params.persistent_atomic_counter = persistent_counter;
-    params.causal_type = is_causal ? CausalType::WindowedAttention : CausalType::None;
-    params.varlen_type = VarlenType::None;
-    params.window_left = window_left;
-    params.window_right = window_right;
-    err = aotriton::v3::flash::attn_fwd(params,
-                                        aotriton::v3::flash::attn_fwd_params::kVersion,
-                                        stream);
-#endif
-  } else {
-    err = attn_fwd(mk_aotensor(q_t, "q"),
-                   mk_aotensor(k_t, "k"),
-                   mk_aotensor(v_t, "v"),
-                   empty_bias,
-                   softmax_scale,
-                   mk_aotensor<2>(M, "M"),
-                   mk_aotensor(output_t, "Out"),
-                   p_dropout,
-                   seed,
-                   offset1,
-                   offset2,
-                   seed_output,
-                   offset_output,
-                   mk_aotensor(softmax_fa_t, "encoded_softmax"),
-                   is_causal,
-                   persistent_counter,
-                   stream);
-  }
-
-  return {out, q_padded, k_padded, v_padded, M.view({batch_size, num_heads, seqlen_q}), seed_t, offset_t, softmax_fa_t};
+  using aotriton::v3::flash::CausalType;
+  using aotriton::v3::flash::VarlenType;
+  aotriton::v3::flash::attn_fwd_params params;
+  params.Q = mk_aotensor(q_t, "q");
+  params.K = mk_aotensor(k_t, "k");
+  params.V = mk_aotensor(v_t, "v");
+  params.Sm_scale = softmax_scale;
+  params.L = mk_aotensor<2>(M, "M");
+  params.Out = mk_aotensor(output_t, "Out");
+  params.Max_seqlen_q = seqlen_q;    // Unused if cu_seqlens_q is empty
+  params.Max_seqlen_k = seqlen_k;    // Unused if cu_seqlens_k is empty
+  params.dropout_p = p_dropout;
+  params.philox_seed_ptr = seed;
+  params.philox_offset1 = offset1;
+  params.philox_offset2 = offset2;
+  params.philox_seed_output = seed_output;
+  params.philox_offset_output = offset_output;
+  params.encoded_softmax = mk_aotensor(softmax_fa_t, "encoded_softmax");
+  params.persistent_atomic_counter = persistent_counter;
+  params.causal_type = is_causal ? CausalType::WindowedAttention : CausalType::None;
+  params.varlen_type = VarlenType::None;
+  params.window_left = window_left;
+  params.window_right = window_right;
+  err = aotriton::v3::flash::attn_fwd(params,
+                                      aotriton::v3::flash::attn_fwd_params::kVersion,
+                                      stream);
+  // Note: These are propagated up to the return of mha_fwd(). comments
+  //       represent the assignments at that level
+  return {out,          // output
+          q_padded,     // q_padded
+          k_padded,     // k_padded
+          v_padded,     // v_padded
+          M.view({batch_size, num_heads, seqlen_q}),   // logsumexp
+          seed_t,       // philox_seed
+          offset_t,     // philox_offset
+          softmax_fa_t};// debug_attn_mask
 }
 
 std::tuple<at::Tensor, at::Tensor, at::Tensor, at::Tensor, at::Tensor, at::Tensor, at::Tensor, at::Tensor>
@@ -342,7 +327,7 @@ mha_varlen_fwd_aot(const at::Tensor &q,  // total_q x num_heads x head_size, tot
                std::optional<int64_t> window_size_right,
                const bool return_softmax,
                const std::optional<at::Generator>& gen_) {
-  TORCH_CHECK(!seqused_k.has_value(), "[ROCm] mha_varlen_fwd: seqused_k must be nullopt");
+  bool strided_varlen = seqused_k.has_value();
   const bool paged_KV = block_table_.has_value();
   TORCH_CHECK(!paged_KV, "[ROCm] mha_varlen_fwd: block_table_ must be nullopt");
   TORCH_CHECK(!alibi_slopes_.has_value(), "[ROCm] mha_varlen_fwd: alibi_slopes_ must be nullopt");
@@ -417,8 +402,13 @@ mha_varlen_fwd_aot(const at::Tensor &q,  // total_q x num_heads x head_size, tot
 
   auto opts = q.options();
 
+#if AOTRITON_COMPACT_VARLEN_LSE
+  auto softmax_lse = at::empty({num_heads, total_q}, opts.dtype(at::kFloat));
+  at::Tensor M = softmax_lse;
+#else
   auto softmax_lse = at::empty({batch_size, num_heads, max_seqlen_q}, opts.dtype(at::kFloat));
   at::Tensor M = softmax_lse.view({batch_size * num_heads, max_seqlen_q});
+#endif
   at::Tensor softmax_fa_t;
   // Only return softmax if there's dropout to reduce compilation time
   if (return_softmax) {
@@ -457,7 +447,6 @@ mha_varlen_fwd_aot(const at::Tensor &q,  // total_q x num_heads x head_size, tot
 
   if (max_seqlen_k > 0) {
     hipError_t err; // TODO: Error handling
-    using aotriton::v2::flash::attn_fwd_compact_varlen;
     using sdp::aotriton_adapter::mk_aotensor;
     using sdp::aotriton_adapter::mk_aoscalartensor;
     using sdp::aotriton_adapter::mk_philoxtensor;
@@ -475,60 +464,48 @@ mha_varlen_fwd_aot(const at::Tensor &q,  // total_q x num_heads x head_size, tot
     auto seed_output = use_philox_state ? mk_philoxtensor(seed_t.data_ptr<int64_t>()) : nullscalar;
     auto offset_output = use_philox_state ? mk_philoxtensor(offset_t.data_ptr<int64_t>()) : nullscalar;
     auto persistent_counter = mk_atomictensor(is_causal ? atomic_counter.data_ptr<int32_t>() : nullptr);
-    if (uses_swa || AOTRITON_ALWAYS_V3_API) {
-#if AOTRITON_V3_API
-      using aotriton::v3::flash::CausalType;
-      using aotriton::v3::flash::VarlenType;
-      aotriton::v3::flash::attn_fwd_params params;
-      params.Q = mk_aotensor(q_padded, "q");
-      params.K = mk_aotensor(k_padded, "k");
-      params.V = mk_aotensor(v_padded, "v");
-      params.Sm_scale = softmax_scale;
-      params.L = mk_aotensor<2>(M, "M");
-      params.Out = mk_aotensor(out_padded, "Out");
+    using aotriton::v3::flash::CausalType;
+    using aotriton::v3::flash::VarlenType;
+    aotriton::v3::flash::attn_fwd_params params;
+    params.Q = mk_aotensor(q_padded, "q");
+    params.K = mk_aotensor(k_padded, "k");
+    params.V = mk_aotensor(v_padded, "v");
+    params.Sm_scale = softmax_scale;
+    params.L = mk_aotensor<2>(M, "logsumexp");
+    params.Out = mk_aotensor(out_padded, "Out");
+    params.Max_seqlen_q = max_seqlen_q;    // Unused if cu_seqlens_q is empty
+    params.Max_seqlen_k = max_seqlen_k;    // Unused if cu_seqlens_k is empty
+    params.dropout_p = p_dropout;
+    params.philox_seed_ptr = seed;
+    params.philox_offset1 = offset1;
+    params.philox_offset2 = offset2;
+    params.philox_seed_output = seed_output;
+    params.philox_offset_output = offset_output;
+    params.encoded_softmax = mk_aotensor(softmax_fa_t, "encoded_softmax");
+    params.persistent_atomic_counter = persistent_counter;
+    params.causal_type = is_causal ? CausalType::WindowedAttention : CausalType::None;
+    params.varlen_type = strided_varlen ? VarlenType::StridedVarlen : VarlenType::CompactVarlen;
+    if (strided_varlen) {
+      // seqused_k holds per-batch actual kv lengths; the kernel expects cumulative
+      // offsets for cu_seqlens_k so it can compute seqlen_k via differencing.
+      // seq_strides_k carries the real memory offsets into the KV cache.
+      const auto& seqused = seqused_k.value();
+      const int num_seqs = seqused.size(0);
+      at::Tensor cu_seqlens_k_from_seqused = at::zeros({num_seqs + 1}, seqused.options());
+      cu_seqlens_k_from_seqused.slice(0, 1).copy_(seqused.cumsum(0));
       params.cu_seqlens_q = mk_aotensor<1>(cu_seqlens_q, "cu_seqlens_q");
-      params.cu_seqlens_k = mk_aotensor<1>(cu_seqlens_k, "cu_seqlens_k");
-      params.Max_seqlen_q = max_seqlen_q;    // Unused if cu_seqlens_q is empty
-      params.Max_seqlen_k = max_seqlen_k;    // Unused if cu_seqlens_k is empty
-      params.dropout_p = p_dropout;
-      params.philox_seed_ptr = seed;
-      params.philox_offset1 = offset1;
-      params.philox_offset2 = offset2;
-      params.philox_seed_output = seed_output;
-      params.philox_offset_output = offset_output;
-      params.encoded_softmax = mk_aotensor(softmax_fa_t, "encoded_softmax");
-      params.persistent_atomic_counter = persistent_counter;
-      params.causal_type = is_causal ? CausalType::WindowedAttention : CausalType::None;
-      params.varlen_type = VarlenType::CompactVarlen;
-      params.window_left = window_left;
-      params.window_right = window_right;
-      err = aotriton::v3::flash::attn_fwd(params,
-                                          aotriton::v3::flash::attn_fwd_params::kVersion,
-                                          stream);
-#endif
+      params.cu_seqlens_k = mk_aotensor<1>(cu_seqlens_k_from_seqused, "cu_seqlens_k");
+      params.seq_strides_q = mk_aotensor<1>(cu_seqlens_q, "seq_strides_q");
+      params.seq_strides_k = mk_aotensor<1>(cu_seqlens_k, "seq_strides_k");
     } else {
-      err = attn_fwd_compact_varlen(mk_aotensor(q_padded, "q"),
-                                    mk_aotensor(k_padded, "k"),
-                                    mk_aotensor(v_padded, "v"),
-                                    empty_bias,
-                                    mk_aotensor<1>(cu_seqlens_q, "cu_seqlens_q"),
-                                    mk_aotensor<1>(cu_seqlens_k, "cu_seqlens_k"),
-                                    max_seqlen_q,
-                                    max_seqlen_k,
-                                    softmax_scale,
-                                    mk_aotensor<2>(M, "M"),
-                                    mk_aotensor(out_padded, "Out"),
-                                    p_dropout,
-                                    seed,
-                                    offset1,
-                                    offset2,
-                                    seed_output,
-                                    offset_output,
-                                    mk_aotensor(softmax_fa_t, "encoded_softmax"),
-                                    is_causal,
-                                    persistent_counter,
-                                    stream);
+      params.cu_seqlens_q = mk_aotensor<1>(cu_seqlens_q, "cu_seqlens_q");
+      params.cu_seqlens_k = mk_aotensor<1>(cu_seqlens_k, "cu_seqlens_k");
     }
+    params.window_left = window_left;
+    params.window_right = window_right;
+    err = aotriton::v3::flash::attn_fwd(params,
+                                        aotriton::v3::flash::attn_fwd_params::kVersion,
+                                        stream);
   } else {
     // If seqlen_k == 0, then we have an empty tensor. We need to set the output to 0.
     out.zero_();
@@ -635,7 +612,7 @@ mha_bwd_aot(const at::Tensor &dout,  // batch_size x seqlen_q x num_heads, x hea
     TORCH_CHECK(dv.stride(-1) == 1, "dv must have contiguous last dimension");
     CHECK_SHAPE(dv, batch_size, seqlen_k, num_heads_k, head_size);
   } else {
-    dv = at::empty_like(k);
+    dv = at::empty_like(v);
   }
 
   auto [needs_swa, window_left, window_right] = calculate_swa(window_size_left,
@@ -670,96 +647,44 @@ mha_bwd_aot(const at::Tensor &dout,  // batch_size x seqlen_q x num_heads, x hea
   hipError_t err; // TODO: Error handling
   using sdp::aotriton_adapter::mk_aotensor;
   using sdp::aotriton_adapter::mk_aoscalartensor;
-  if (uses_swa || AOTRITON_ALWAYS_V3_API) {
-#if AOTRITON_V3_API
-    // Fused BWD does not support SWA
-    using aotriton::v3::flash::CausalType;
-    using aotriton::v3::flash::VarlenType;
-    aotriton::v3::flash::attn_bwd_params params;
-    params.Q = mk_aotensor(q_t, "q");
-    params.K = mk_aotensor(k_t, "k");
-    params.V = mk_aotensor(v_t, "v");
-    params.Sm_scale = softmax_scale;
-    params.Out = mk_aotensor(out_t, "out");
-    params.DO = mk_aotensor(dout_t, "dout");
-    params.DQ = mk_aotensor(dq_t, "dq");
-    params.DK = mk_aotensor(dk_t, "dk");
-    params.DV = mk_aotensor(dv_t, "dv");
-    params.L = mk_aotensor<2>(softmax_lse_cont, "L");
-    params.Max_seqlen_q = seqlen_q;        // Unused if cu_seqlens_q is empty
-    params.Max_seqlen_k = seqlen_k;        // Unused if cu_seqlens_k is empty
-    params.dropout_p = p_dropout;
-    params.philox_seed_ptr =  mk_aoscalartensor(philox_seed);
-    params.philox_offset1 = mk_aoscalartensor(philox_offset);
-    params.philox_offset2 = 0;
-    // SWA in AOTriton Kernels is treated as "Generalized Causal masks"
-    params.causal_type = is_causal || uses_swa ? CausalType::WindowedAttention : CausalType::None;
-    params.window_left = window_left;
-    params.window_right = window_right;
-    params.varlen_type = VarlenType::None;
-#if AOTRITON_ALWAYS_V3_API
-    using sdp::aotriton_adapter::mklazy_empty_like;
-    using sdp::aotriton_adapter::mklazy_fp32zeros;
-    using sdp::aotriton_adapter::LazyTensorContext;
-    LazyTensorContext lazy_delta { .like_tensor = softmax_lse_cont, .tensor_name = "delta" };
-    LazyTensorContext lazy_dq_acc { .like_tensor = dq_t, .tensor_name = "dq_acc" };
-    params.D = mklazy_empty_like<2>(&lazy_delta);
-    params.DQ_ACC = mklazy_fp32zeros<4>(&lazy_dq_acc);
-#else
-    at::Tensor delta = at::empty_like(softmax_lse_cont).contiguous();
-    params.D = mk_aotensor<2>(delta, "delta");
-#endif
-    err = aotriton::v3::flash::attn_bwd(params,
-                                        aotriton::v3::flash::attn_bwd_params::kVersion,
-                                        stream);
-#endif
-  } else if (use_fused_bwd) {
-    using aotriton::v2::flash::attn_bwd_fused;
-    using sdp::aotriton_adapter::cast_dtype;
-    aotriton::TensorView<4> empty_bias(0, {0,0,0,0}, {0,0,0,0}, cast_dtype(q.dtype()));
-    err = attn_bwd_fused(mk_aotensor(q_t, "q"),
-                         mk_aotensor(k_t, "k"),
-                         mk_aotensor(v_t, "v"),
-                         empty_bias,
-                         softmax_scale,
-                         mk_aotensor(out_t, "out"),
-                         mk_aotensor(dout_t, "dout"),
-                         mk_aotensor(dq_t, "dq"),
-                         mk_aotensor(dk_t, "dk"),
-                         mk_aotensor(dv_t, "dv"),
-                         empty_bias,  // dbb
-                         mk_aotensor<2>(softmax_lse_cont, "L"),
-                         p_dropout,
-                         mk_aoscalartensor(philox_seed),
-                         mk_aoscalartensor(philox_offset),
-                         0,
-                         is_causal,
-                         stream);
-  } else {
-    at::Tensor delta = at::empty_like(softmax_lse_cont).contiguous();
-    using aotriton::v2::flash::attn_bwd;
-    using sdp::aotriton_adapter::cast_dtype;
-    aotriton::TensorView<4> empty_bias(0, {0,0,0,0}, {0,0,0,0}, cast_dtype(q.dtype()));
-    err = attn_bwd(mk_aotensor(q_t, "q"),
-                   mk_aotensor(k_t, "k"),
-                   mk_aotensor(v_t, "v"),
-                   empty_bias,
-                   softmax_scale,
-                   mk_aotensor(out_t, "out"),
-                   mk_aotensor(dout_t, "dout"),
-                   mk_aotensor(dq_t, "dq"),
-                   mk_aotensor(dk_t, "dk"),
-                   mk_aotensor(dv_t, "dv"),
-                   empty_bias,  // db
-                   mk_aotensor<2>(softmax_lse_cont, "L"),
-                   mk_aotensor<2>(delta, "delta"),
-                   p_dropout,
-                   mk_aoscalartensor(philox_seed),
-                   mk_aoscalartensor(philox_offset),
-                   0,
-                   is_causal,
-                   stream);
-  }
+  // Fused BWD does not support SWA
+  using aotriton::v3::flash::CausalType;
+  using aotriton::v3::flash::VarlenType;
+  aotriton::v3::flash::attn_bwd_params params;
+  params.Q = mk_aotensor(q_t, "q");
+  params.K = mk_aotensor(k_t, "k");
+  params.V = mk_aotensor(v_t, "v");
+  params.Sm_scale = softmax_scale;
+  params.Out = mk_aotensor(out_t, "out");
+  params.DO = mk_aotensor(dout_t, "dout");
+  params.DQ = mk_aotensor(dq_t, "dq");
+  params.DK = mk_aotensor(dk_t, "dk");
+  params.DV = mk_aotensor(dv_t, "dv");
+  params.L = mk_aotensor<2>(softmax_lse_cont, "L");
+  params.Max_seqlen_q = seqlen_q;        // Unused if cu_seqlens_q is empty
+  params.Max_seqlen_k = seqlen_k;        // Unused if cu_seqlens_k is empty
+  params.dropout_p = p_dropout;
+  params.philox_seed_ptr =  mk_aoscalartensor(philox_seed);
+  params.philox_offset1 = mk_aoscalartensor(philox_offset);
+  params.philox_offset2 = 0;
+  // SWA in AOTriton Kernels is treated as "Generalized Causal masks"
+  params.causal_type = is_causal || uses_swa ? CausalType::WindowedAttention : CausalType::None;
+  params.window_left = window_left;
+  params.window_right = window_right;
+  params.varlen_type = VarlenType::None;
+  using sdp::aotriton_adapter::mklazy_empty_like;
+  using sdp::aotriton_adapter::mklazy_fp32zeros;
+  using sdp::aotriton_adapter::LazyTensorContext;
+  LazyTensorContext lazy_delta { .like_tensor = softmax_lse_cont, .tensor_name = "delta" };
+  LazyTensorContext lazy_dq_acc { .like_tensor = dq_t, .tensor_name = "dq_acc" };
+  params.D = mklazy_empty_like<2>(&lazy_delta);
+  params.DQ_ACC = mklazy_fp32zeros<4>(&lazy_dq_acc);
+  aotriton::v3::flash::attn_options options;
+  options.deterministic = deterministic;
+  err = aotriton::v3::flash::attn_bwd(params,
+                                      aotriton::v3::flash::attn_bwd_params::kVersion,
+                                      stream,
+                                      &options);
 
   return { dq, dk, dv, softmax_d };
 }
@@ -842,7 +767,11 @@ mha_varlen_bwd_aot(const at::Tensor &dout,  // total_q x num_heads, x head_size
   CHECK_SHAPE(cu_seqlens_q, batch_size + 1);
   CHECK_SHAPE(cu_seqlens_k, batch_size + 1);
 
+#if AOTRITON_COMPACT_VARLEN_LSE
+  at::Tensor softmax_lse_cont = softmax_lse.view({num_heads, total_q}).contiguous();
+#else
   at::Tensor softmax_lse_cont = softmax_lse.view({batch_size * num_heads, max_seqlen_q}).contiguous();
+#endif
 
   at::Tensor q_padded, k_padded, v_padded;
   q_padded = q.unsqueeze(0).transpose(1, 2);
@@ -923,79 +852,45 @@ mha_varlen_bwd_aot(const at::Tensor &dout,  // total_q x num_heads, x head_size
     hipError_t err; // TODO: Error handling
     using sdp::aotriton_adapter::mk_aotensor;
     using sdp::aotriton_adapter::mk_aoscalartensor;
-    if (uses_swa || AOTRITON_ALWAYS_V3_API) {
-#if AOTRITON_V3_API
-      using aotriton::v3::flash::CausalType;
-      using aotriton::v3::flash::VarlenType;
-      aotriton::v3::flash::attn_bwd_params params;
-      params.Q = mk_aotensor(q_padded, "q");
-      params.K = mk_aotensor(k_padded, "k");
-      params.V = mk_aotensor(v_padded, "v");
-      params.Sm_scale = softmax_scale;
-      params.Out = mk_aotensor(out_t, "out");
-      params.DO = mk_aotensor(dout_t, "dout");
-      params.DK = mk_aotensor(dk_padded, "dk");
-      params.DV = mk_aotensor(dv_padded, "dv");
-      params.DQ = mk_aotensor(dq_padded, "dq");
-      params.L = mk_aotensor<2>(softmax_lse_cont, "L");
-      params.cu_seqlens_q = mk_aotensor<1>(cu_seqlens_q, "cu_seqlens_q");
-      params.cu_seqlens_k = mk_aotensor<1>(cu_seqlens_k, "cu_seqlens_k");
-      params.Max_seqlen_q = max_seqlen_q;        // Unused if cu_seqlens_q is empty
-      params.Max_seqlen_k = max_seqlen_k;        // Unused if cu_seqlens_k is empty
-      params.dropout_p = p_dropout;
-      params.philox_seed_ptr =  mk_aoscalartensor(philox_seed);
-      params.philox_offset1 = mk_aoscalartensor(philox_offset);
-      params.philox_offset2 = 0;
-      // SWA in AOTriton Kernels is treated as "Generalized Causal masks"
-      params.causal_type = is_causal || uses_swa ? CausalType::WindowedAttention : CausalType::None;
-      params.varlen_type = VarlenType::CompactVarlen;
-      params.window_left = window_left;
-      params.window_right = window_right;
-#if AOTRITON_ALWAYS_V3_API
-      using sdp::aotriton_adapter::mklazy_empty_like;
-      using sdp::aotriton_adapter::mklazy_fp32zeros;
-      using sdp::aotriton_adapter::LazyTensorContext;
-      LazyTensorContext lazy_delta { .like_tensor = softmax_lse_cont, .tensor_name = "delta" };
-      LazyTensorContext lazy_dq_acc { .like_tensor = dq_padded, .tensor_name = "dq_acc" };
-      params.D = mklazy_empty_like<2>(&lazy_delta);
-      params.DQ_ACC = mklazy_fp32zeros<4>(&lazy_dq_acc);
-#else
-      at::Tensor delta = at::empty_like(softmax_lse_cont).contiguous();
-      params.D = mk_aotensor<2>(delta, "delta");
-#endif
-      err = aotriton::v3::flash::attn_bwd(params,
-                                          aotriton::v3::flash::attn_bwd_params::kVersion,
-                                          stream);
-#endif  // AOTRITON_ALWAYS_V3_API
-    } else {
-      using aotriton::v2::flash::attn_bwd_compact_varlen;
-      using sdp::aotriton_adapter::cast_dtype;
-      at::Tensor delta = at::empty_like(softmax_lse_cont).contiguous();
-      aotriton::TensorView<4> empty_bias(0, {0,0,0,0}, {0,0,0,0}, cast_dtype(q.dtype()));
-      err = attn_bwd_compact_varlen(mk_aotensor(q_padded, "q"),
-                                    mk_aotensor(k_padded, "k"),
-                                    mk_aotensor(v_padded, "v"),
-                                    mk_aotensor<1>(cu_seqlens_q, "cu_seqlens_q"),
-                                    mk_aotensor<1>(cu_seqlens_k, "cu_seqlens_k"),
-                                    max_seqlen_q,
-                                    max_seqlen_k,
-                                    empty_bias,
-                                    softmax_scale,
-                                    mk_aotensor(out_t, "out"),
-                                    mk_aotensor(dout_t, "dout"),
-                                    mk_aotensor(dq_padded, "dq"),
-                                    mk_aotensor(dk_padded, "dk"),
-                                    mk_aotensor(dv_padded, "dv"),
-                                    empty_bias,
-                                    mk_aotensor<2>(softmax_lse_cont, "L"),
-                                    mk_aotensor<2>(delta, "delta"),
-                                    p_dropout,
-                                    mk_aoscalartensor(philox_seed),
-                                    mk_aoscalartensor(philox_offset),
-                                    0,
-                                    is_causal,
-                                    stream);
-    }
+    using aotriton::v3::flash::CausalType;
+    using aotriton::v3::flash::VarlenType;
+    aotriton::v3::flash::attn_bwd_params params;
+    params.Q = mk_aotensor(q_padded, "q");
+    params.K = mk_aotensor(k_padded, "k");
+    params.V = mk_aotensor(v_padded, "v");
+    params.Sm_scale = softmax_scale;
+    params.Out = mk_aotensor(out_t, "out");
+    params.DO = mk_aotensor(dout_t, "dout");
+    params.DK = mk_aotensor(dk_padded, "dk");
+    params.DV = mk_aotensor(dv_padded, "dv");
+    params.DQ = mk_aotensor(dq_padded, "dq");
+    params.L = mk_aotensor<2>(softmax_lse_cont, "L");
+    params.cu_seqlens_q = mk_aotensor<1>(cu_seqlens_q, "cu_seqlens_q");
+    params.cu_seqlens_k = mk_aotensor<1>(cu_seqlens_k, "cu_seqlens_k");
+    params.Max_seqlen_q = max_seqlen_q;        // Unused if cu_seqlens_q is empty
+    params.Max_seqlen_k = max_seqlen_k;        // Unused if cu_seqlens_k is empty
+    params.dropout_p = p_dropout;
+    params.philox_seed_ptr =  mk_aoscalartensor(philox_seed);
+    params.philox_offset1 = mk_aoscalartensor(philox_offset);
+    params.philox_offset2 = 0;
+    // SWA in AOTriton Kernels is treated as "Generalized Causal masks"
+    params.causal_type = is_causal || uses_swa ? CausalType::WindowedAttention : CausalType::None;
+    params.varlen_type = VarlenType::CompactVarlen;
+    params.window_left = window_left;
+    params.window_right = window_right;
+    using sdp::aotriton_adapter::mklazy_empty_like;
+    using sdp::aotriton_adapter::mklazy_fp32zeros;
+    using sdp::aotriton_adapter::LazyTensorContext;
+    LazyTensorContext lazy_delta { .like_tensor = softmax_lse_cont, .tensor_name = "delta" };
+    LazyTensorContext lazy_dq_acc { .like_tensor = dq_padded, .tensor_name = "dq_acc" };
+    params.D = mklazy_empty_like<2>(&lazy_delta);
+    params.DQ_ACC = mklazy_fp32zeros<4>(&lazy_dq_acc);
+    aotriton::v3::flash::attn_options options;
+    options.deterministic = deterministic;
+    err = aotriton::v3::flash::attn_bwd(params,
+                                        aotriton::v3::flash::attn_bwd_params::kVersion,
+                                        stream,
+                                        &options);
   } else {
     // If seqlen_q == 0, then we have an empty tensor. We need to set the output to 0.
     dq.zero_();
diff --git a/aten/src/ATen/native/transformers/hip/flash_attn/ck/CMakeLists.txt b/aten/src/ATen/native/transformers/hip/flash_attn/ck/CMakeLists.txt
index 819880cf3bc5c..8e3ecc2dabc1e 100644
--- a/aten/src/ATen/native/transformers/hip/flash_attn/ck/CMakeLists.txt
+++ b/aten/src/ATen/native/transformers/hip/flash_attn/ck/CMakeLists.txt
@@ -1,7 +1,7 @@
 # generate a list of kernels, but not actually emit files at config stage
 execute_process(
-  COMMAND python3 ${CMAKE_SOURCE_DIR}/third_party/composable_kernel/example/ck_tile/01_fmha/generate.py
-  --api fwd --receipt 4 --list_blobs ${CMAKE_CURRENT_LIST_DIR}/fwd_blob_list.txt
+  COMMAND python3 ${CMAKE_SOURCE_DIR}/third_party/composable_kernel/example/ck_tile/01_fmha/generate.py --optdim=32,64,128,256
+  --api fwd --receipt 4 --filter "*_lse*ntrload*nsink*" --list_blobs ${CMAKE_CURRENT_LIST_DIR}/fwd_blob_list.txt
   RESULT_VARIABLE ret
 )
 
@@ -10,8 +10,8 @@ if(ret AND NOT ret EQUAL 0)
 endif()
 
 execute_process(
-  COMMAND python3 ${CMAKE_SOURCE_DIR}/third_party/composable_kernel/example/ck_tile/01_fmha/generate.py
-  --api fwd_splitkv --receipt 4 --list_blobs ${CMAKE_CURRENT_LIST_DIR}/fwd_splitkv_blob_list.txt
+  COMMAND python3 ${CMAKE_SOURCE_DIR}/third_party/composable_kernel/example/ck_tile/01_fmha/generate.py --optdim=32,64,128,256
+  --api fwd_splitkv --receipt 4 --filter "*psdv*_lse*_nsquant*" --list_blobs ${CMAKE_CURRENT_LIST_DIR}/fwd_splitkv_blob_list.txt
   RESULT_VARIABLE ret
 )
 
@@ -20,8 +20,8 @@ if(ret AND NOT ret EQUAL 0)
 endif()
 
 execute_process(
-  COMMAND python3 ${CMAKE_SOURCE_DIR}/third_party/composable_kernel/example/ck_tile/01_fmha/generate.py
-  --api fwd_appendkv --receipt 4 --list_blobs ${CMAKE_CURRENT_LIST_DIR}/fwd_appendkv_blob_list.txt
+  COMMAND python3 ${CMAKE_SOURCE_DIR}/third_party/composable_kernel/example/ck_tile/01_fmha/generate.py --optdim=32,64,128,256
+  --api fwd_appendkv --receipt 4 --filter "*psskddv_*" --list_blobs ${CMAKE_CURRENT_LIST_DIR}/fwd_appendkv_blob_list.txt
   RESULT_VARIABLE ret
 )
 
@@ -31,7 +31,7 @@ endif()
 
 execute_process(
   COMMAND python3 ${CMAKE_SOURCE_DIR}/third_party/composable_kernel/example/ck_tile/01_fmha/generate.py
-  --api bwd --receipt 4 --list_blobs ${CMAKE_CURRENT_LIST_DIR}/bwd_blob_list.txt
+  --api bwd --optdim=32,64,128,256 --receipt 4 --filter "*psdv*@*psd*@*_pd1dv1*_ntrload*" --list_blobs ${CMAKE_CURRENT_LIST_DIR}/bwd_blob_list.txt
   RESULT_VARIABLE ret
 )
 
@@ -40,28 +40,28 @@ if(ret AND NOT ret EQUAL 0)
 endif()
 
 # Generate the files for both fwd, fwd_splitkv, fwd_appendkv, and bwd
-execute_process(COMMAND python3 ${CMAKE_SOURCE_DIR}/third_party/composable_kernel/example/ck_tile/01_fmha/generate.py --api fwd --receipt 4 --output_dir ${CMAKE_CURRENT_LIST_DIR}
+execute_process(COMMAND python3 ${CMAKE_SOURCE_DIR}/third_party/composable_kernel/example/ck_tile/01_fmha/generate.py --api fwd  --optdim=32,64,128,256 --receipt 4 --filter "*_lse*ntrload*nsink*" --output_dir ${CMAKE_CURRENT_LIST_DIR}
 )
 
 if(ret AND NOT ret EQUAL 0)
   message( FATAL_ERROR "CK Tile FMHA FAILED to generate FWD kernels.")
 endif()
 
-execute_process(COMMAND python3 ${CMAKE_SOURCE_DIR}/third_party/composable_kernel/example/ck_tile/01_fmha/generate.py --api fwd_splitkv --receipt 4 --output_dir ${CMAKE_CURRENT_LIST_DIR}
+execute_process(COMMAND python3 ${CMAKE_SOURCE_DIR}/third_party/composable_kernel/example/ck_tile/01_fmha/generate.py --api fwd_splitkv --optdim=32,64,128,256 --receipt 4 --filter "*psdv*_lse*_nsquant*" --output_dir ${CMAKE_CURRENT_LIST_DIR}
 )
 
 if(ret AND NOT ret EQUAL 0)
     message( FATAL_ERROR "CK Tile FMHA FAILED to generate FWD_SPLITKV kernels.")
 endif()
 
-execute_process(COMMAND python3 ${CMAKE_SOURCE_DIR}/third_party/composable_kernel/example/ck_tile/01_fmha/generate.py --api fwd_appendkv --receipt 4 --output_dir ${CMAKE_CURRENT_LIST_DIR}
+execute_process(COMMAND python3 ${CMAKE_SOURCE_DIR}/third_party/composable_kernel/example/ck_tile/01_fmha/generate.py --api fwd_appendkv --optdim=32,64,128,256 --receipt 4 --filter "*psskddv_*" --output_dir ${CMAKE_CURRENT_LIST_DIR}
 )
 
 if(ret AND NOT ret EQUAL 0)
     message( FATAL_ERROR "CK Tile FMHA FAILED to generate FWD_APPENDKV kernels.")
 endif()
 
-execute_process(COMMAND python3 ${CMAKE_SOURCE_DIR}/third_party/composable_kernel/example/ck_tile/01_fmha/generate.py --api bwd --receipt 4 --output_dir ${CMAKE_CURRENT_LIST_DIR}
+execute_process(COMMAND python3 ${CMAKE_SOURCE_DIR}/third_party/composable_kernel/example/ck_tile/01_fmha/generate.py --api bwd --optdim=32,64,128,256 --receipt 4 --filter "*psdv*@*psd*@*_pd1dv1*_ntrload*" --output_dir ${CMAKE_CURRENT_LIST_DIR}
   RESULT_VARIABLE ret
 )
 
diff --git a/aten/src/ATen/native/transformers/hip/flash_attn/ck/fav_v3/CMakeLists.txt b/aten/src/ATen/native/transformers/hip/flash_attn/ck/fav_v3/CMakeLists.txt
index e0eb28652c1e6..e0e96cc8d5b32 100644
--- a/aten/src/ATen/native/transformers/hip/flash_attn/ck/fav_v3/CMakeLists.txt
+++ b/aten/src/ATen/native/transformers/hip/flash_attn/ck/fav_v3/CMakeLists.txt
@@ -1,14 +1,38 @@
 include(CMakePrintHelpers)
 
-# Generate AITER/CK Asm code
+# Create an isolated venv for codegen.py which requires pandas>=2.1.0 (for
+# the "future.no_silent_downcasting" option).  This avoids polluting the CI /
+# host Python environment while still satisfying the dependency.
+set(CODEGEN_VENV_DIR "${CMAKE_CURRENT_BINARY_DIR}/_codegen_venv")
+
+if(NOT EXISTS "${CODEGEN_VENV_DIR}/bin/python3")
+    message(STATUS "Creating isolated codegen venv at ${CODEGEN_VENV_DIR}")
+    execute_process(
+        COMMAND python3 -m venv "${CODEGEN_VENV_DIR}"
+        RESULT_VARIABLE ret
+    )
+    if(ret AND NOT ret EQUAL 0)
+        message(FATAL_ERROR "Failed to create codegen venv")
+    endif()
+
+    execute_process(
+        COMMAND "${CODEGEN_VENV_DIR}/bin/pip" install --quiet "pandas>=2.1.0,<=2.3.3" numpy
+        RESULT_VARIABLE ret
+    )
+    if(ret AND NOT ret EQUAL 0)
+        message(FATAL_ERROR "Failed to install pandas/numpy into codegen venv")
+    endif()
+endif()
+
+# Generate AITER/CK Asm code (using the venv python so pandas is new enough)
 execute_process(
     COMMAND ${CMAKE_COMMAND} -E env "AITER_GPU_ARCHS=gfx942;gfx950"
-            python3 ${CMAKE_SOURCE_DIR}/third_party/aiter/hsa/codegen.py -m fmha_v3_bwd --output_dir ${CMAKE_CURRENT_LIST_DIR}
+            "${CODEGEN_VENV_DIR}/bin/python3" ${CMAKE_SOURCE_DIR}/third_party/aiter/hsa/codegen.py -m fmha_v3_bwd --output_dir ${CMAKE_CURRENT_LIST_DIR}
     RESULT_VARIABLE ret
 )
 
 if(ret AND NOT ret EQUAL 0)
-    message( FATAL_ERROR "Failed to generate FAv3 CK Kernels")
+    message(FATAL_ERROR "Failed to generate FAv3 CK Kernels")
 endif()
 
 execute_process(COMMAND bash -c "cp ${CMAKE_SOURCE_DIR}/third_party/aiter/csrc/cpp_itfs/mha_bwd.cu ${CMAKE_CURRENT_LIST_DIR}/mha_bwd.hip")
diff --git a/aten/src/ATen/native/transformers/hip/flash_attn/ck/mha_bwd_ck.hip b/aten/src/ATen/native/transformers/hip/flash_attn/ck/mha_bwd_ck.hip
index 2d3692d9f98df..d7633521d8eef 100644
--- a/aten/src/ATen/native/transformers/hip/flash_attn/ck/mha_bwd_ck.hip
+++ b/aten/src/ATen/native/transformers/hip/flash_attn/ck/mha_bwd_ck.hip
@@ -127,7 +127,7 @@ aiter::mha_bwd_args get_ck_fmha_bwd_args(const mask_info &mask,
     return aiter::mha_bwd_args{
         // aiter args
         static_cast<int>(mask.type),
-        true,   // use_asm_v3
+        hdim <= 192,   // use_asm_v3: ASM v3 only supports head dim <= 192
         true,   // v3_atomic_fp32
         1,      // v3_bf16_cvt
         false,  // v3_api_check
@@ -375,11 +375,16 @@ mha_bwd_ck(const at::Tensor &dout,                   // batch_size x seqlen_q x
         dv_expanded = dv;
     }
 
-    uint64_t drop_seed = 1, drop_offset = 0;
-    drop_seed = *philox_seed.data_ptr<int64_t>();
-    drop_offset = *philox_offset.data_ptr<int64_t>();
-    auto drop_seed_offset = std::make_pair(&drop_seed, &drop_offset);
+    auto gen = at::get_generator_or_default<at::CUDAGeneratorImpl>(
+            std::nullopt, at::cuda::detail::getDefaultCUDAGenerator());
 
+    uint64_t* drop_seed, drop_offset;
+    int64_t counter_offset = batch_size * num_heads * ck_tile::get_warp_size();
+    std::pair<uint64_t*, uint64_t*> drop_seed_offset = {nullptr,nullptr};
+    if(is_dropout) {
+        drop_seed_offset.first = philox_seed.data_ptr<uint64_t>();
+        drop_seed_offset.second = philox_offset.data_ptr<uint64_t>();
+    }
 
     if (seqlen_q > 0) {
         ck_tile::stream_config stream_config{stream};
diff --git a/aten/src/ATen/native/transformers/hip/flash_attn/ck/mha_fwd_ck.hip b/aten/src/ATen/native/transformers/hip/flash_attn/ck/mha_fwd_ck.hip
index 441589e70d763..90ae0e7aa6afc 100644
--- a/aten/src/ATen/native/transformers/hip/flash_attn/ck/mha_fwd_ck.hip
+++ b/aten/src/ATen/native/transformers/hip/flash_attn/ck/mha_fwd_ck.hip
@@ -6,9 +6,100 @@
 #include <fmha_fwd.hpp>
 #include <mask.hpp>
 
+#include <cstdio>
+
+namespace {
+/* Debug traces. Comment in to aid in debug of 'Invalid argument to fmha_fwd' error
+const char* mask_enum_str(mask_enum m) {
+    switch (m) {
+        case mask_enum::no_mask:          return "no_mask";
+        case mask_enum::mask_top_left:    return "mask_top_left";
+        case mask_enum::mask_bottom_right:return "mask_bottom_right";
+        case mask_enum::window_generic:   return "window_generic";
+        default:                          return "unknown";
+    }
+}
+
+const char* bias_enum_str(bias_enum b) {
+    switch (b) {
+        case bias_enum::no_bias:          return "no_bias";
+        case bias_enum::elementwise_bias: return "elementwise_bias";
+        case bias_enum::alibi:            return "alibi";
+        default:                          return "unknown";
+    }
+}
+
+const char* qscale_enum_str(quant_scale_enum q) {
+    switch (q) {
+        case quant_scale_enum::no_scale:      return "no_scale";
+        case quant_scale_enum::pertensor:      return "pertensor";
+        case quant_scale_enum::blockscale:     return "blockscale";
+        case quant_scale_enum::kv_blockscale:  return "kv_blockscale";
+        default:                               return "unknown";
+    }
+}
+
+void dump_fmha_fwd_kernel_selection_params(const fmha_fwd_traits& t,
+                                           const fmha_fwd_args& a) {
+    std::printf("=== fmha_fwd kernel selection params ===\n");
+    std::printf("--- traits ---\n");
+    std::printf("  data_type          : %s\n",  t.data_type.c_str());
+    std::printf("  hdim_q             : %d\n",  t.hdim_q);
+    std::printf("  hdim_v             : %d\n",  t.hdim_v);
+    std::printf("  is_group_mode      : %s\n",  t.is_group_mode      ? "true" : "false");
+    std::printf("  is_v_rowmajor      : %s\n",  t.is_v_rowmajor      ? "true" : "false");
+    std::printf("  has_logits_soft_cap: %s\n",  t.has_logits_soft_cap ? "true" : "false");
+    std::printf("  mask_type          : %s (%d)\n", mask_enum_str(t.mask_type),
+                static_cast<int>(t.mask_type));
+    std::printf("  bias_type          : %s (%d)\n", bias_enum_str(t.bias_type),
+                static_cast<int>(t.bias_type));
+    std::printf("  has_lse            : %s\n",  t.has_lse             ? "true" : "false");
+    std::printf("  has_dropout        : %s\n",  t.has_dropout         ? "true" : "false");
+    std::printf("  qscale_type        : %s (%d)\n", qscale_enum_str(t.qscale_type),
+                static_cast<int>(t.qscale_type));
+    std::printf("  skip_min_seqlen_q  : %s\n",  t.skip_min_seqlen_q  ? "true" : "false");
+    std::printf("  has_sink           : %s\n",  t.has_sink            ? "true" : "false");
+    std::printf("--- args ---\n");
+    std::printf("  hdim_q             : %d\n",  static_cast<int>(a.hdim_q));
+    std::printf("  hdim_v             : %d\n",  static_cast<int>(a.hdim_v));
+    std::printf("  seqlen_k           : %d\n",  static_cast<int>(a.seqlen_k));
+    std::printf("  cu_seqlen_k_ptr    : %p\n",  a.cu_seqlen_k_ptr);
+    std::printf("  batch              : %d\n",  static_cast<int>(a.batch));
+    std::printf("  nhead_q            : %d\n",  static_cast<int>(a.nhead_q));
+    std::printf("  max_seqlen_q       : %d\n",  static_cast<int>(a.max_seqlen_q));
+    std::printf("========================================\n");
+    std::fflush(stdout);
+}
+*/
+} // anonymous namespace
 
 namespace pytorch_flash {
 
+// SFINAE trait to detect block-scale quantization fields in fmha_fwd_args.
+// Newer versions of composable_kernel add these fields; older versions don't.
+// This lets the PyTorch integration layer work with either version.
+template <typename T, typename = void>
+struct has_fmha_block_scale_fields : std::false_type {};
+
+template <typename T>
+struct has_fmha_block_scale_fields<T,
+    std::void_t<decltype(std::declval<T&>().block_scale_size_q)>> : std::true_type {};
+
+template <typename Args>
+void set_fmha_fwd_block_scale_fields([[maybe_unused]] Args& args) {
+    if constexpr (has_fmha_block_scale_fields<Args>::value) {
+        args.block_scale_seqstart_q_ptr = nullptr;
+        args.block_scale_seqstart_k_ptr = nullptr;
+        args.nhead_stride_q_descale = 0;
+        args.nhead_stride_k_descale = 0;
+        args.nhead_stride_v_descale = 0;
+        args.batch_stride_q_descale = 0;
+        args.batch_stride_k_descale = 0;
+        args.batch_stride_v_descale = 0;
+        args.block_scale_size_q = 0;
+        args.block_scale_size_kv = 0;
+    }
+}
 
 fmha_fwd_traits get_ck_fmha_fwd_traits(const mask_info &mask,
                                        std::string dtype,
@@ -96,61 +187,58 @@ fmha_fwd_args get_ck_fmha_fwd_args(bool has_lse,
         nhead_stride_bias = a_b.stride(1);
         batch_stride_bias = a_b.stride(0);
     }
-    return fmha_fwd_args{q.data_ptr(),
-                         k.data_ptr(),
-                         v.data_ptr(),
-                         attn_bias_ptr, // bias
-                         nullptr, // q_descale_ptr
-                         nullptr, // k_descale_ptr
-                         nullptr, // v_descale_ptr
-                         has_dropout_randval ? dropout_randval.data_ptr() : nullptr,
-                         has_lse ? softmax_lse.data_ptr() : nullptr,
-                         out.data_ptr(),
-                         nullptr, // seqstart_q
-                         nullptr, // seqstart_k
-                         nullptr, // seqlen_q_ptr
-                         nullptr, // seqlen_k_ptr
-                         nullptr, // cu_seqlen_q_ptr
-                         nullptr, // cu_seqlen_k_ptr
-                         nullptr, // sink_ptr
-                         seqlen_q,
-                         seqlen_k,
-                         b,
-                         seqlen_q,                          // max_seqlen_q
-                         d,                                 // hdim_q
-                         d,                                 // hdim_v
-                         h,                                 // nhead
-                         h_k,                               // nhead_k
-                         softmax_scale,                     // scale_s
-                         0.0f,                              // logits_soft_cap
-                         stride_q,
-                         stride_k,
-                         stride_v,
-                         stride_attn_bias,
-                         stride_randval,
-                         stride_o,
-                         nhead_stride_q,
-                         nhead_stride_k,
-                         nhead_stride_v,
-                         nhead_stride_bias,                 // nhead_stride_bias
-                         nhead_stride_randval,
-                         nhead_stride_lse,
-                         nhead_stride_o,
-                         batch_stride_q,
-                         batch_stride_k,
-                         batch_stride_v,
-                         batch_stride_bias,                 // batch_stride_bias
-                         batch_stride_randval,
-                         batch_stride_lse,
-                         batch_stride_o,
-                         mask.left,
-                         mask.right,
-                         0,                                 // sink_size
-                         static_cast<ck_tile::index_t>(mask.type),
-                         -1,                                // min_seqlen_q
-                         p_dropout,
-                         has_dropout_randval,
-                         drop_seed_offset};
+    fmha_fwd_args args{};
+    args.q_ptr = q.data_ptr();
+    args.k_ptr = k.data_ptr();
+    args.v_ptr = v.data_ptr();
+    args.bias_ptr = attn_bias_ptr;
+    args.q_descale_ptr = nullptr;
+    args.k_descale_ptr = nullptr;
+    args.v_descale_ptr = nullptr;
+    args.rand_val_ptr = has_dropout_randval ? dropout_randval.data_ptr() : nullptr;
+    args.lse_ptr = has_lse ? softmax_lse.data_ptr() : nullptr;
+    args.o_ptr = out.data_ptr();
+    args.sink_ptr = nullptr;
+    args.seqlen_q = seqlen_q;
+    args.seqlen_k = seqlen_k;
+    args.batch = b;
+    args.max_seqlen_q = seqlen_q;
+    args.hdim_q = d;
+    args.hdim_v = d;
+    args.nhead_q = h;
+    args.nhead_k = h_k;
+    args.scale_s = softmax_scale;
+    args.logits_soft_cap = 0.0f;
+    args.stride_q = stride_q;
+    args.stride_k = stride_k;
+    args.stride_v = stride_v;
+    args.stride_bias = stride_attn_bias;
+    args.stride_randval = stride_randval;
+    args.stride_o = stride_o;
+    args.nhead_stride_q = nhead_stride_q;
+    args.nhead_stride_k = nhead_stride_k;
+    args.nhead_stride_v = nhead_stride_v;
+    args.nhead_stride_bias = nhead_stride_bias;
+    args.nhead_stride_randval = nhead_stride_randval;
+    args.nhead_stride_lse = nhead_stride_lse;
+    args.nhead_stride_o = nhead_stride_o;
+    args.batch_stride_q = batch_stride_q;
+    args.batch_stride_k = batch_stride_k;
+    args.batch_stride_v = batch_stride_v;
+    args.batch_stride_bias = batch_stride_bias;
+    args.batch_stride_randval = batch_stride_randval;
+    args.batch_stride_lse = batch_stride_lse;
+    args.batch_stride_o = batch_stride_o;
+    args.window_size_left = mask.left;
+    args.window_size_right = mask.right;
+    args.sink_size = 0;
+    args.mask_type = static_cast<ck_tile::index_t>(mask.type);
+    args.min_seqlen_q = -1;
+    args.p_drop = p_dropout;
+    args.s_randval = has_dropout_randval;
+    args.drop_seed_offset = drop_seed_offset;
+    set_fmha_fwd_block_scale_fields(args);
+    return args;
 }
 
 std::tuple<at::Tensor, at::Tensor, at::Tensor, at::Tensor, at::Tensor, at::Tensor, at::Tensor, at::Tensor>
@@ -183,7 +271,6 @@ mha_fwd_ck(const at::Tensor &q,                      // batch_size x seqlen_q x
     TORCH_CHECK(v.stride(-1) == 1, "Input tensor must have contiguous last dimension");
 
     const auto sizes = q.sizes();
-
     const int batch_size = sizes[0];
     int seqlen_q = sizes[1];
     int num_heads = sizes[2];
@@ -232,7 +319,6 @@ mha_fwd_ck(const at::Tensor &q,                      // batch_size x seqlen_q x
     CHECK_SHAPE(k, batch_size, seqlen_k, num_heads_k, head_size);
     CHECK_SHAPE(v, batch_size, seqlen_k, num_heads_k, head_size);
 
-
     at::Tensor q_padded, k_padded, v_padded;
     if (head_size % 8 != 0) {
         q_padded = at::pad(temp_q, {0, 8 - head_size % 8});
@@ -245,7 +331,6 @@ mha_fwd_ck(const at::Tensor &q,                      // batch_size x seqlen_q x
         v_padded = v;
     }
 
-
     at::Tensor out;
     if (out_.has_value()) {
         out = out_.value();
@@ -272,7 +357,6 @@ mha_fwd_ck(const at::Tensor &q,                      // batch_size x seqlen_q x
     auto opts = q.options();
     bool has_lse = true;
     bool has_dropout = p_dropout > 0.0f;
-
     at::Tensor softmax_lse;
     // TODO - check gradient, only training require lse
     softmax_lse = at::empty({batch_size, num_heads, seqlen_q}, opts.dtype(at::kFloat));
@@ -283,47 +367,46 @@ mha_fwd_ck(const at::Tensor &q,                      // batch_size x seqlen_q x
         p = at::empty({batch_size, num_heads, seqlen_q, seqlen_k}, opts.dtype(at::kByte));
     }
     else {
-        p = at::empty({ 0 }, opts);
+        p = at::empty({ 0, 0, 0, 0 }, opts.dtype(at::kByte));
     }
 
+    uint64_t drop_seed = 1, drop_offset = 0;
+    at::Tensor seed_t, offset_t;
     int64_t counter_offset = batch_size * num_heads * ck_tile::get_warp_size();
-    auto rng_state = at::empty({2}, opts.dtype(at::kLong));
-    auto rng_state_ptr = reinterpret_cast<uint64_t*>(rng_state.data_ptr());
 
 
-
-    at::Tensor seed_t, offset_t;
+    // rng_state is used to pass philox params to CK in a type it likes i.e. uint64
+    auto rng_state_options = at::TensorOptions().dtype(at::kUInt64).device(at::kCUDA);
+    auto rng_state = at::zeros({2}, rng_state_options.dtype(at::kUInt64));
+    auto _unused = at::empty({}, at::dtype(c10::kUInt64).device(at::kCUDA));
 
     if (p_dropout > 0.0)  {
+
         auto gen = at::get_generator_or_default<at::CUDAGeneratorImpl>(
             gen_, at::cuda::detail::getDefaultCUDAGenerator());
+
         // See Note [Acquire lock when using random generators]
         std::lock_guard<std::mutex> lock(gen->mutex_);
-
         auto philox_args = gen->philox_cuda_state(counter_offset);
 
+        std::tie(drop_seed, drop_offset) = at::cuda::philox::unpack(philox_args);
+        rng_state[0] = *(reinterpret_cast<int64_t*>(&drop_seed));
+        rng_state[1] = *(reinterpret_cast<int64_t*>(&drop_offset));
 
-
-        hipLaunchKernelGGL(
-            flash::ParsePhiloxCudaState, dim3(1), dim3(64), 0, at::cuda::getCurrentCUDAStream(), philox_args, rng_state_ptr);
-        seed_t = at::scalar_tensor(at::Scalar(static_cast<uint64_t>(rng_state_ptr[0])), at::dtype(at::kLong));
-        offset_t = at::scalar_tensor(at::Scalar(static_cast<uint64_t>(rng_state_ptr[1])), at::dtype(at::kLong));
-    }
-    else
-    {
+    } else {
         seed_t = at::empty({}, at::dtype(at::kLong).device(at::kCUDA));
         offset_t = at::empty({}, at::dtype(at::kLong).device(at::kCUDA));
     }
-
     std::optional<at::Tensor> attn_bias;
+    std::pair<uint64_t*, uint64_t*> drop_seed_offset;
     if( attn_bias_.has_value())
     {
       attn_bias = attn_bias_;
     }
-
     if (seqlen_k > 0) {
-        auto drop_seed_offset = std::make_pair(rng_state_ptr, rng_state_ptr + 1);
-        auto stream = at::cuda::getCurrentCUDAStream().stream();
+        drop_seed_offset.first = rng_state[0].data_ptr<uint64_t>();
+        drop_seed_offset.second = rng_state[1].data_ptr<uint64_t>();
+        auto stream = at::cuda::getCurrentHIPStream().stream();
         ck_tile::stream_config stream_config{stream};
 
         auto traits =
@@ -364,12 +447,28 @@ mha_fwd_ck(const at::Tensor &q,                      // batch_size x seqlen_q x
         out.zero_();
         softmax_lse.fill_(std::numeric_limits<float>::infinity());
     }
-
     if (seqlenq_ngroups_swapped) {
         out = out.transpose(1, 2).reshape({batch_size, 1, num_heads_k * seqlen_q, head_size});
         q_padded = q_padded.transpose(1, 2).reshape({batch_size, 1, num_heads_k * seqlen_q, head_size});
         softmax_lse = softmax_lse.reshape({batch_size, num_heads_k * seqlen_q, 1});
     }
-    return {out, q_padded, k_padded, v_padded, softmax_lse, seed_t, offset_t, p};
+    // Before returning populate seed_t and offset_t if dropout is used
+    if(p_dropout > 0.0)
+    {
+        seed_t = at::scalar_tensor(at::Scalar(reinterpret_cast<uint64_t>(*(drop_seed_offset.first))),
+                                   at::dtype(at::kUInt64).device(at::kCUDA));
+        offset_t = at::scalar_tensor(at::Scalar(reinterpret_cast<uint64_t>(*(drop_seed_offset.second))),
+                                   at::dtype(at::kUInt64).device(at::kCUDA));
+    }
+    // Note: These are propagated up to the return of mha_fwd(). comments
+    //       represent the assignments at that level
+    return {out,            // output
+            q_padded,       // q_padded
+            k_padded,       // k_padded
+            v_padded,       // v_padded
+            softmax_lse,    // logsumexp
+            seed_t,         // philox_seed
+            offset_t,       // philox_offset
+            p};             // debug_attn_mask
 }
 } //namespace pytorch_flash
diff --git a/aten/src/ATen/native/transformers/hip/flash_attn/ck/mha_varlen_bwd_ck.hip b/aten/src/ATen/native/transformers/hip/flash_attn/ck/mha_varlen_bwd_ck.hip
index 3031965833608..896afc7320e96 100644
--- a/aten/src/ATen/native/transformers/hip/flash_attn/ck/mha_varlen_bwd_ck.hip
+++ b/aten/src/ATen/native/transformers/hip/flash_attn/ck/mha_varlen_bwd_ck.hip
@@ -401,8 +401,10 @@ mha_varlen_bwd_ck(const at::Tensor &dout,                   // total_q x num_hea
 
     uint64_t drop_seed = 1, drop_offset = 0;
 
-    drop_seed = *philox_seed.data_ptr<uint64_t>();
-    drop_offset = *philox_offset.data_ptr<uint64_t>();
+    if (is_dropout) {
+        drop_seed = *reinterpret_cast<uint64_t*>(philox_seed.data_ptr());
+        drop_offset = *reinterpret_cast<uint64_t*>(philox_offset.data_ptr());
+    }
     auto drop_seed_offset = std::make_pair(&drop_seed, &drop_offset);
 
     if (max_seqlen_q > 0) {
diff --git a/aten/src/ATen/native/transformers/hip/flash_attn/ck/mha_varlen_fwd_ck.hip b/aten/src/ATen/native/transformers/hip/flash_attn/ck/mha_varlen_fwd_ck.hip
index 2a9d4899e8236..eb4a3b841fa3e 100644
--- a/aten/src/ATen/native/transformers/hip/flash_attn/ck/mha_varlen_fwd_ck.hip
+++ b/aten/src/ATen/native/transformers/hip/flash_attn/ck/mha_varlen_fwd_ck.hip
@@ -8,6 +8,32 @@
 
 namespace pytorch_flash {
 
+// SFINAE trait to detect block-scale quantization fields in fmha_fwd_args.
+// Newer versions of composable_kernel add these fields; older versions don't.
+// This lets the PyTorch integration layer work with either version.
+template <typename T, typename = void>
+struct has_fmha_block_scale_fields : std::false_type {};
+
+template <typename T>
+struct has_fmha_block_scale_fields<T,
+    std::void_t<decltype(std::declval<T&>().block_scale_size_q)>> : std::true_type {};
+
+template <typename Args>
+void set_fmha_fwd_block_scale_fields([[maybe_unused]] Args& args) {
+    if constexpr (has_fmha_block_scale_fields<Args>::value) {
+        args.block_scale_seqstart_q_ptr = nullptr;
+        args.block_scale_seqstart_k_ptr = nullptr;
+        args.nhead_stride_q_descale = 0;
+        args.nhead_stride_k_descale = 0;
+        args.nhead_stride_v_descale = 0;
+        args.batch_stride_q_descale = 0;
+        args.batch_stride_k_descale = 0;
+        args.batch_stride_v_descale = 0;
+        args.block_scale_size_q = 0;
+        args.block_scale_size_kv = 0;
+    }
+}
+
 fmha_fwd_traits get_ck_fmha_varlen_fwd_traits(const mask_info &mask,
                                               std::string dtype,
                                               int head_size,
@@ -73,7 +99,7 @@ fmha_fwd_args get_ck_fmha_varlen_fwd_args(bool has_lse,
     ck_tile::index_t nhead_stride_k = k.stride(1);
     ck_tile::index_t nhead_stride_v = v.stride(1);
     ck_tile::index_t nhead_stride_o = out.stride(1);
-    ck_tile::index_t nhead_stride_lse = has_lse ? softmax_lse.stride(1) : 0;
+    ck_tile::index_t nhead_stride_lse = has_lse ? softmax_lse.stride(0) : 0;
     ck_tile::index_t nhead_stride_randval = has_dropout_randval ? dropout_randval.stride(0) : 0;
 
     ck_tile::index_t batch_stride_q = 0;
@@ -81,7 +107,7 @@ fmha_fwd_args get_ck_fmha_varlen_fwd_args(bool has_lse,
     ck_tile::index_t batch_stride_v = 0;
     ck_tile::index_t batch_stride_o = 0;
 
-    ck_tile::index_t batch_stride_lse = has_lse ? softmax_lse.stride(0) : 0;
+    ck_tile::index_t batch_stride_lse = 0;
     ck_tile::index_t batch_stride_randval = 0;
 
     void *attn_bias_ptr = nullptr;
@@ -97,61 +123,60 @@ fmha_fwd_args get_ck_fmha_varlen_fwd_args(bool has_lse,
         stride_attn_bias = a_b.stride(0);
     }
 
-    return fmha_fwd_args{q.data_ptr(),
-                         k.data_ptr(),
-                         v.data_ptr(),
-                         attn_bias_ptr, // bias
-                         nullptr, // q_descale_ptr
-                         nullptr, // k_descale_ptr
-                         nullptr, // v_descale_ptr
-                         has_dropout_randval ? dropout_randval.data_ptr() : nullptr,
-                         has_lse ? softmax_lse.data_ptr() : nullptr,
-                         out.data_ptr(),
-                         seqlens_q.data_ptr(), // seqstart_q
-                         seqlens_k.data_ptr(), // seqstart_k
-                         nullptr, // seqlen_q_ptr
-                         nullptr, // seqlen_k_ptr
-                         nullptr, // cu_seqlen_q_ptr
-                         nullptr, // cu_seqlen_k_ptr
-                         nullptr, // sink_ptr
-                         total_q,
-                         total_k,
-                         b,
-                         max_seqlen_q,
-                         d,             // hdim_q
-                         d,             // hdim_v
-                         h,             // nhead
-                         h_k,           // nhead_k
-                         softmax_scale, // scale_s
-                         0.0f,          // logits_soft_cap
-                         stride_q,
-                         stride_k,
-                         stride_v,
-                         stride_attn_bias,
-                         stride_randval,
-                         stride_o,
-                         nhead_stride_q,
-                         nhead_stride_k,
-                         nhead_stride_v,
-                         0, // nhead_stride_bias, FA without bias
-                         nhead_stride_randval,
-                         nhead_stride_lse,
-                         nhead_stride_o,
-                         batch_stride_q,
-                         batch_stride_k,
-                         batch_stride_v,
-                         0, // batch_stride_bias, FA without bias
-                         batch_stride_randval,
-                         batch_stride_lse,
-                         batch_stride_o,
-                         mask.left,
-                         mask.right,
-                         0,             // sink_size
-                         static_cast<ck_tile::index_t>(mask.type),
-                         -1,                                // min_seqlen_q
-                         p_dropout,
-                         has_dropout_randval,
-                         drop_seed_offset};
+    fmha_fwd_args args{};
+    args.q_ptr = q.data_ptr();
+    args.k_ptr = k.data_ptr();
+    args.v_ptr = v.data_ptr();
+    args.bias_ptr = attn_bias_ptr;
+    args.q_descale_ptr = nullptr;
+    args.k_descale_ptr = nullptr;
+    args.v_descale_ptr = nullptr;
+    args.rand_val_ptr = has_dropout_randval ? dropout_randval.data_ptr() : nullptr;
+    args.lse_ptr = has_lse ? softmax_lse.data_ptr() : nullptr;
+    args.o_ptr = out.data_ptr();
+    args.seqstart_q_ptr = seqlens_q.data_ptr();
+    args.seqstart_k_ptr = seqlens_k.data_ptr();
+    args.sink_ptr = nullptr;
+    args.seqlen_q = total_q;
+    args.seqlen_k = total_k;
+    args.batch = b;
+    args.max_seqlen_q = max_seqlen_q;
+    args.hdim_q = d;
+    args.hdim_v = d;
+    args.nhead_q = h;
+    args.nhead_k = h_k;
+    args.scale_s = softmax_scale;
+    args.logits_soft_cap = 0.0f;
+    args.stride_q = stride_q;
+    args.stride_k = stride_k;
+    args.stride_v = stride_v;
+    args.stride_bias = stride_attn_bias;
+    args.stride_randval = stride_randval;
+    args.stride_o = stride_o;
+    args.nhead_stride_q = nhead_stride_q;
+    args.nhead_stride_k = nhead_stride_k;
+    args.nhead_stride_v = nhead_stride_v;
+    args.nhead_stride_bias = 0;
+    args.nhead_stride_randval = nhead_stride_randval;
+    args.nhead_stride_lse = nhead_stride_lse;
+    args.nhead_stride_o = nhead_stride_o;
+    args.batch_stride_q = batch_stride_q;
+    args.batch_stride_k = batch_stride_k;
+    args.batch_stride_v = batch_stride_v;
+    args.batch_stride_bias = 0;
+    args.batch_stride_randval = batch_stride_randval;
+    args.batch_stride_lse = batch_stride_lse;
+    args.batch_stride_o = batch_stride_o;
+    args.window_size_left = mask.left;
+    args.window_size_right = mask.right;
+    args.sink_size = 0;
+    args.mask_type = static_cast<ck_tile::index_t>(mask.type);
+    args.min_seqlen_q = -1;
+    args.p_drop = p_dropout;
+    args.s_randval = has_dropout_randval;
+    args.drop_seed_offset = drop_seed_offset;
+    set_fmha_fwd_block_scale_fields(args);
+    return args;
 }
 
 std::tuple<at::Tensor, at::Tensor, at::Tensor, at::Tensor, at::Tensor, at::Tensor, at::Tensor, at::Tensor>
@@ -287,7 +312,7 @@ mha_varlen_fwd_ck(const at::Tensor &q,                   // total_q x num_heads
 
     at::Tensor softmax_lse;
     // TODO - check gradient, only training require lse
-    softmax_lse = at::empty({batch_size, num_heads, max_seqlen_q}, opts.dtype(at::kFloat));
+    softmax_lse = at::empty({num_heads, total_q}, opts.dtype(at::kFloat));
 
     at::Tensor p;
     if (return_dropout_randval) {
diff --git a/aten/src/ATen/native/transformers/sdp_utils.h b/aten/src/ATen/native/transformers/sdp_utils.h
index 89a150f359cf9..18f13781944f8 100644
--- a/aten/src/ATen/native/transformers/sdp_utils.h
+++ b/aten/src/ATen/native/transformers/sdp_utils.h
@@ -21,10 +21,14 @@ void alloc_with_matching_layout(
   std::vector<int> fill_order(shape.size());
   std::iota(fill_order.begin(), fill_order.end(), 0);
   const auto q_strides = q.strides();
+  // note: why INT64_MAX instead of 1.
+  // When Q's strides include 0, e.g. (0, 0, 128, 1), mapping stride 0 to 1 leads to
+  // fill_order of [0, 1, 3, 2], i.e. the output strides are [1, 8, 1024, 16].
+  // To match output strides with Q, use INT64_MAx so that broadcast dims come last in fill_order.
   std::stable_sort(
       fill_order.begin(), fill_order.end(), [&q_strides](int idx1, int idx2) {
-        int64_t s1 = q_strides[idx1] ? q_strides[idx1] : 1;
-        int64_t s2 = q_strides[idx2] ? q_strides[idx2] : 1;
+        int64_t s1 = q_strides[idx1] ? q_strides[idx1] : INT64_MAX;
+        int64_t s2 = q_strides[idx2] ? q_strides[idx2] : INT64_MAX;
         return s1 < s2;
       });
   std::vector<int64_t> ordered_strides(shape.size());
@@ -33,9 +37,7 @@ void alloc_with_matching_layout(
     ordered_strides[dim_idx] = current_stride;
     current_stride *= shape[dim_idx];
   }
-  output = at::empty(at::IntArrayRef(shape), q.options())
-               .as_strided(
-                   at::IntArrayRef(shape), at::IntArrayRef(ordered_strides), 0);
+  output = at::empty_strided(at::IntArrayRef(shape), at::IntArrayRef(ordered_strides), q.options());
 }
 
 void permute_to_matching_layout(const Tensor& output, Tensor& grad_output) {
diff --git a/aten/src/ATen/native/transformers/xpu/sdp_utils.cpp b/aten/src/ATen/native/transformers/xpu/sdp_utils.cpp
index a69ecf3274d2a..b5abfa27e8baa 100644
--- a/aten/src/ATen/native/transformers/xpu/sdp_utils.cpp
+++ b/aten/src/ATen/native/transformers/xpu/sdp_utils.cpp
@@ -29,7 +29,8 @@ bool check_flash_attention_hardware_support(
       c10::array_of<sycl::ext::oneapi::experimental::architecture>(
           sycl::ext::oneapi::experimental::architecture::intel_gpu_pvc,
           sycl::ext::oneapi::experimental::architecture::intel_gpu_pvc_vg,
-          sycl::ext::oneapi::experimental::architecture::intel_gpu_bmg_g21);
+          sycl::ext::oneapi::experimental::architecture::intel_gpu_bmg_g21,
+          sycl::ext::oneapi::experimental::architecture::intel_gpu_bmg_g31);
   auto* device_prop = at::xpu::getCurrentDeviceProperties();
   auto device_architecture = device_prop->architecture;
 
@@ -39,7 +40,7 @@ bool check_flash_attention_hardware_support(
           device_architecture) == supported_architectures.end()) {
     if (debug) {
       TORCH_WARN(
-          "XPU device architecture does not support flash attention. Supported architectures are: intel_gpu_pvc, intel_gpu_pvc_vg, intel_gpu_bmg_g21.");
+          "XPU device architecture does not support flash attention. Supported architectures are: intel_gpu_pvc, intel_gpu_pvc_vg, intel_gpu_bmg_g21, intel_gpu_bmg_g31.");
     }
     return false;
   }
@@ -80,9 +81,12 @@ inline bool check_flash_attention_datatype(
 inline bool check_flash_attention_head_dim_size(
     sdp_params const& params,
     bool debug) {
-  const int query_size_last = params.query.size(3);
-  const int key_size_last = params.key.size(3);
-  const int value_size_last = params.value.size(3);
+  // Use sym_size to preserve symbolic shapes during tracing.
+  // Using concrete .size() would materialize symbolic dimensions into static
+  // guards, preventing dynamic shape generalization across recompilations.
+  const auto query_size_last = params.query.sym_size(-1);
+  const auto key_size_last = params.key.sym_size(-1);
+  const auto value_size_last = params.value.sym_size(-1);
 
   const bool head_dims_equal = (query_size_last == key_size_last) &&
       (query_size_last == value_size_last);
@@ -101,7 +105,7 @@ inline bool check_flash_attention_head_dim_size(
     return false;
   }
 
-  constexpr auto max_supported_headdim = 192;
+  const auto max_supported_headdim = c10::SymInt(192);
   if (query_size_last > max_supported_headdim) {
     if (debug) {
       TORCH_WARN(
diff --git a/aten/src/ATen/native/utils/ParamsHash.h b/aten/src/ATen/native/utils/ParamsHash.h
index 49e9e784d9426..e5a8df6cef240 100644
--- a/aten/src/ATen/native/utils/ParamsHash.h
+++ b/aten/src/ATen/native/utils/ParamsHash.h
@@ -35,9 +35,7 @@ struct ParamsEqual {
   static_assert(std::is_standard_layout_v<Params>, "Params is not POD");
 
   bool operator()(const Params& a, const Params& b) const noexcept {
-    auto ptr1 = reinterpret_cast<const uint8_t*>(&a);
-    auto ptr2 = reinterpret_cast<const uint8_t*>(&b);
-    return memcmp(ptr1, ptr2, sizeof(Params)) == 0;
+    return memcmp(&a, &b, sizeof(Params)) == 0;
   }
 };
 
@@ -60,26 +58,18 @@ struct ParamsWrapper {
     memcpy(&(this->pod), &(other.pod), sizeof(this->pod));
   }
 
-  ParamsWrapper(ParamsWrapper&& other) noexcept {
-    memcpy(&(this->pod), &(other.pod), sizeof(this->pod));
-  }
-
   ParamsWrapper& operator=(const ParamsWrapper& other) noexcept {
     memcpy(&(this->pod), &(other.pod), sizeof(this->pod));
     return *this;
   }
 
-  ParamsWrapper& operator=(ParamsWrapper&& other) noexcept {
-    memcpy(&(this->pod), &(other.pod), sizeof(this->pod));
-    return *this;
-  }
+  ParamsWrapper(ParamsWrapper&& other) = delete;
+  ParamsWrapper& operator=(ParamsWrapper&& other) = delete;
 
   inline friend bool operator==(
       const ParamsWrapper& lhs,
       const ParamsWrapper& rhs) noexcept {
-    auto ptr1 = reinterpret_cast<const uint8_t*>(&(lhs.pod));
-    auto ptr2 = reinterpret_cast<const uint8_t*>(&(rhs.pod));
-    return memcmp(ptr1, ptr2, sizeof(lhs.pod)) == 0;
+    return memcmp(&lhs.pod, &rhs.pod, sizeof(T)) == 0;
   }
 };
 
@@ -87,20 +77,9 @@ struct ParamsWrapper {
 // constructors for additional safety
 template <typename ParamsWrapper>
 struct ParamsWrapperHash {
-  // Params must be a POD because we read out its memory
-  // contents as char* when hashing
-  static_assert(
-      std::is_standard_layout_v<decltype(ParamsWrapper::pod)>,
-      "ParamsWrapper cannot wrap non-POD data");
-
   size_t operator()(const ParamsWrapper& params_wrapper) const noexcept {
-    auto ptr = reinterpret_cast<const uint8_t*>(&(params_wrapper.pod));
-    uint32_t value = 0x811C9DC5;
-    for (const auto i : c10::irange(sizeof(params_wrapper.pod))) {
-      value ^= ptr[i];
-      value *= 0x01000193;
-    }
-    return (size_t)value;
+    ParamsHash<decltype(ParamsWrapper::pod)> hasher;
+    return hasher(params_wrapper.pod);
   }
 };
 
diff --git a/aten/src/ATen/native/vulkan/api/Exception.h b/aten/src/ATen/native/vulkan/api/Exception.h
index 328ff67894e89..b9db4fea4e575 100644
--- a/aten/src/ATen/native/vulkan/api/Exception.h
+++ b/aten/src/ATen/native/vulkan/api/Exception.h
@@ -1,3 +1,4 @@
+// @allow-raw-throw
 #pragma once
 // @lint-ignore-every CLANGTIDY facebook-hte-BadMemberName
 #ifdef USE_VULKAN_API
diff --git a/aten/src/ATen/native/vulkan/ops/Lstm.cpp b/aten/src/ATen/native/vulkan/ops/Lstm.cpp
index 9d7209f030bcb..3e76cab598e25 100644
--- a/aten/src/ATen/native/vulkan/ops/Lstm.cpp
+++ b/aten/src/ATen/native/vulkan/ops/Lstm.cpp
@@ -158,7 +158,8 @@ std::tuple<Tensor, Tensor, Tensor> lstm_input(
   x = x.reshape({batch_size, seq_length, x.size(1)});
   h_n = h_n.reshape({h_n.size(0) * h_n.size(1), h_n.size(2), h_n.size(3)});
   c_n = c_n.reshape({c_n.size(0) * c_n.size(1), c_n.size(2), c_n.size(3)});
-  return std::tuple<Tensor, Tensor, Tensor>(x, h_n, c_n);
+  return std::tuple<Tensor, Tensor, Tensor>(
+      std::move(x), std::move(h_n), std::move(c_n));
 }
 
 #ifdef USE_VULKAN_API
@@ -419,7 +420,8 @@ std::tuple<Tensor, Tensor, Tensor> run_lstm_context(
   x = x.reshape({batch_size, seq_length, x.size(1)});
   h_n = h_n.reshape({h_n.size(0) * h_n.size(1), h_n.size(2), h_n.size(3)});
   c_n = c_n.reshape({c_n.size(0) * c_n.size(1), c_n.size(2), c_n.size(3)});
-  return std::tuple<Tensor, Tensor, Tensor>(x, h_n, c_n);
+  return std::tuple<Tensor, Tensor, Tensor>(
+      std::move(x), std::move(h_n), std::move(c_n));
 }
 
 } // namespace ops
diff --git a/aten/src/ATen/native/vulkan/ops/Mm.cpp b/aten/src/ATen/native/vulkan/ops/Mm.cpp
index aeb25b56e60e9..a5a0d93cccc3d 100644
--- a/aten/src/ATen/native/vulkan/ops/Mm.cpp
+++ b/aten/src/ATen/native/vulkan/ops/Mm.cpp
@@ -631,6 +631,7 @@ Tensor run_quantized_addmm_context(
     return output;
   } else {
     std::vector<int64_t> shape;
+    shape.reserve(static_cast<size_t>(std::max<int64_t>(0, input_arg.dim())));
     for (const auto i : c10::irange(input_arg.dim() - 1)) {
       shape.emplace_back(input_arg.size(i));
     }
@@ -751,6 +752,7 @@ Tensor run_addmm_context(
     return output;
   } else {
     std::vector<int64_t> shape;
+    shape.reserve(static_cast<size_t>(std::max<int64_t>(0, input_arg.dim())));
     for (const auto i : c10::irange(input_arg.dim() - 1)) {
       shape.emplace_back(input_arg.size(i));
     }
diff --git a/aten/src/ATen/native/vulkan/ops/NativeLayerNorm.cpp b/aten/src/ATen/native/vulkan/ops/NativeLayerNorm.cpp
index 1ec6957162cbb..338363c49cdbe 100644
--- a/aten/src/ATen/native/vulkan/ops/NativeLayerNorm.cpp
+++ b/aten/src/ATen/native/vulkan/ops/NativeLayerNorm.cpp
@@ -76,6 +76,7 @@ std::tuple<Tensor, Tensor, Tensor> native_layer_norm(
   const Tensor bias = bias_opt->is_vulkan() ? *bias_opt : bias_opt->vulkan();
 
   std::vector<int64_t> dims_to_reduce;
+  dims_to_reduce.reserve(normalized_shape.size());
   for (const auto i : c10::irange(normalized_shape.size())) {
     dims_to_reduce.push_back(input_arg.dim() - i - 1);
   }
diff --git a/aten/src/ATen/native/vulkan/ops/Repeat.cpp b/aten/src/ATen/native/vulkan/ops/Repeat.cpp
index 94c29b55d0571..0c08d7f7c26dc 100644
--- a/aten/src/ATen/native/vulkan/ops/Repeat.cpp
+++ b/aten/src/ATen/native/vulkan/ops/Repeat.cpp
@@ -37,6 +37,8 @@ Tensor repeat(const Tensor& self, const IntArrayRef repeats) {
 
   std::vector<at::Tensor> tensor_seq_to_concat;
   for (const auto i : c10::irange(out_ndims)) {
+    tensor_seq_to_concat.reserve(
+        static_cast<size_t>(std::max<int64_t>(0, repeats[i])));
     for (const auto k : c10::irange(repeats[i])) {
       (void)k;
       tensor_seq_to_concat.emplace_back(tensor_to_repeat.clone());
diff --git a/aten/src/ATen/nnapi/CMakeLists.txt b/aten/src/ATen/nnapi/CMakeLists.txt
index 065729e5c887e..4ba4048220370 100644
--- a/aten/src/ATen/nnapi/CMakeLists.txt
+++ b/aten/src/ATen/nnapi/CMakeLists.txt
@@ -1,9 +1,9 @@
 # Define this to build the NNAPI binding out of tree.
 if(PYTORCH_NNAPI_STANDALONE)
-  cmake_minimum_required(VERSION 3.5 FATAL_ERROR)
+  cmake_minimum_required(VERSION 3.10 FATAL_ERROR)
   project(pytorch_nnapi)
 
-  set(CMAKE_CXX_STANDARD 14 CACHE STRING "The C++ standard whose features are requested to build this target.")
+  set(CMAKE_CXX_STANDARD 20 CACHE STRING "The C++ standard whose features are requested to build this target.")
   find_package(Torch REQUIRED)
 
   set(NNAPI_SRCS
diff --git a/aten/src/ATen/templates/TensorMethods.cpp b/aten/src/ATen/templates/TensorMethods.cpp
index 0504dccc385c9..c1440abc4cb0e 100644
--- a/aten/src/ATen/templates/TensorMethods.cpp
+++ b/aten/src/ATen/templates/TensorMethods.cpp
@@ -8,39 +8,39 @@ namespace at {
 namespace {
 
 // Verifies the requested type is the same as the Tensor's type.
-void check_type(const TensorBase& tensor, ScalarType type, std::string_view type_name) {
+void check_type(const TensorBase& tensor, ScalarType type) {
   TORCH_CHECK(
       tensor.scalar_type() == type
       || (isQIntType(tensor.scalar_type())
           && toUnderlying(tensor.scalar_type()) == type),
-      "expected scalar type ", type_name, " but found ", tensor.scalar_type());
+      "expected scalar type ", type, " but found ", tensor.scalar_type());
 }
 
 } // namespace
 
-#define DEFINE_CAST(T, name)                                         \
-   template <>                                                       \
-   TORCH_API const T* TensorBase::const_data_ptr() const {           \
-     check_type(*this, ScalarType::name, #name);                     \
-     return this->unsafeGetTensorImpl()->data_ptr_impl<T>();         \
-   }                                                                 \
-                                                                     \
-   template <>                                                       \
-   TORCH_API const T* TensorBase::const_data_ptr<const T>() const {  \
-     check_type(*this, ScalarType::name, #name);                     \
-     return this->unsafeGetTensorImpl()->data_ptr_impl<std::remove_const_t<T>>(); \
-   }                                                                 \
-                                                                     \
-   template <>                                                       \
-   TORCH_API T* TensorBase::mutable_data_ptr() const {               \
-     check_type(*this, ScalarType::name, #name);                     \
-     return this->unsafeGetTensorImpl()->mutable_data_ptr_impl<T>(); \
-   }                                                                 \
-                                                                     \
-   template <>                                                       \
-   TORCH_API T* TensorBase::data_ptr() const {                       \
-     return mutable_data_ptr<T>();                                   \
-   }                                                                 \
+template <typename T>
+const T* TensorBase::const_data_ptr() const {
+  using NonConstT = std::remove_const_t<T>;
+  check_type(*this, c10::CppTypeToScalarType<NonConstT>());
+  return this->unsafeGetTensorImpl()->data_ptr_impl<NonConstT>();
+}
+
+template <typename T>
+T* TensorBase::mutable_data_ptr() const {
+  check_type(*this, c10::CppTypeToScalarType<T>());
+  return this->unsafeGetTensorImpl()->mutable_data_ptr_impl<T>();
+}
+
+template <typename T>
+T* TensorBase::data_ptr() const {
+  return this->mutable_data_ptr<T>();
+}
+
+#define DEFINE_CAST(T, name)                                                \
+   template TORCH_API const T* TensorBase::const_data_ptr<T>() const;       \
+   template TORCH_API const T* TensorBase::const_data_ptr<const T>() const; \
+   template TORCH_API T* TensorBase::mutable_data_ptr() const;              \
+   template TORCH_API T* TensorBase::data_ptr() const;
 
  AT_FORALL_SCALAR_TYPES_WITH_COMPLEX(DEFINE_CAST)
  AT_FORALL_QINT_TYPES(DEFINE_CAST)
diff --git a/aten/src/ATen/templates/enum_tag.h b/aten/src/ATen/templates/enum_tag.h
index 1320fbc28ab8f..49c174b721153 100644
--- a/aten/src/ATen/templates/enum_tag.h
+++ b/aten/src/ATen/templates/enum_tag.h
@@ -2,9 +2,18 @@
 
 // ${generated_comment}
 
+#include <torch/headeronly/macros/Macros.h>
+
+HIDDEN_NAMESPACE_BEGIN(torch, headeronly)
+
+// Enum of valid tags obtained from the entries in tags.yaml
+enum class Tag {
+    ${enum_of_valid_tags}
+};
+
+HIDDEN_NAMESPACE_END(torch, headeronly)
+
+// Re-expose in the at:: namespace for backward compatibility
 namespace at {
-    // Enum of valid tags obtained from the entries in tags.yaml
-    enum class Tag {
-        ${enum_of_valid_tags}
-    };
+    using torch::headeronly::Tag;
 }
diff --git a/aten/src/ATen/test/CMakeLists.txt b/aten/src/ATen/test/CMakeLists.txt
index 923b7119a42fc..62c6c354d79b4 100644
--- a/aten/src/ATen/test/CMakeLists.txt
+++ b/aten/src/ATen/test/CMakeLists.txt
@@ -9,6 +9,7 @@ list(APPEND ATen_CPU_TEST_SRCS
   ${CMAKE_CURRENT_SOURCE_DIR}/Dimname_test.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/MaybeOwned_test.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/NamedTensor_test.cpp
+  ${CMAKE_CURRENT_SOURCE_DIR}/accelerator_graph_test.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/apply_utils_test.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/atest.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/basic.cpp
diff --git a/aten/src/ATen/test/accelerator_graph_test.cpp b/aten/src/ATen/test/accelerator_graph_test.cpp
new file mode 100644
index 0000000000000..a1a8eed56f7b7
--- /dev/null
+++ b/aten/src/ATen/test/accelerator_graph_test.cpp
@@ -0,0 +1,75 @@
+#include <gtest/gtest.h>
+
+#include <ATen/accelerator/Graph.h>
+
+static bool is_capture_begin_called = false;
+static bool is_capture_end_called = false;
+static bool keep_raw_graph = false;
+static bool graph_debug = false;
+
+struct DummyGraphImpl : public at::GraphImplInterface {
+  DummyGraphImpl(const at::GraphImplArgs& args = {}) {
+    keep_raw_graph = args.keep_graph;
+  }
+
+  void capture_begin(
+      [[maybe_unused]] at::MempoolId_t pool = {0, 0},
+      [[maybe_unused]] at::GraphCaptureMode capture_mode =
+          at::GraphCaptureMode::Default) override {
+    is_capture_begin_called = true;
+  }
+
+  void capture_end() override {
+    is_capture_end_called = true;
+  }
+
+  void instantiate() override {
+    TORCH_CHECK_NOT_IMPLEMENTED(false, "Not implemented");
+  }
+
+  void replay() override {
+    TORCH_CHECK_NOT_IMPLEMENTED(false, "Not implemented");
+  }
+
+  void reset() override {
+    TORCH_CHECK_NOT_IMPLEMENTED(false, "Not implemented");
+  }
+
+  at::MempoolId_t pool() const override {
+    return {10, 0};
+  }
+
+  void enable_debug_mode() override {
+    graph_debug = true;
+  };
+
+  void debug_dump(const std::string& path) override {
+    TORCH_CHECK_NOT_IMPLEMENTED(false, "Not implemented");
+  }
+};
+
+namespace at {
+REGISTER_GRAPH_IMPL(DUMMY, DummyGraphImpl)
+}
+
+TEST(AcceleratorGraphTest, graphRegistrationAndCapture) {
+  EXPECT_EQ(at::has_graph_impl(at::kPrivateUse1), false);
+  at::register_privateuse1_backend("DUMMY");
+  EXPECT_EQ(at::has_graph_impl(at::kPrivateUse1), true);
+  EXPECT_EQ(is_capture_begin_called, false);
+  EXPECT_EQ(is_capture_end_called, false);
+  EXPECT_EQ(keep_raw_graph, false);
+  auto graph = at::accelerator::Graph(true);
+  EXPECT_EQ(keep_raw_graph, true);
+  graph.capture_begin();
+  EXPECT_EQ(is_capture_begin_called, true);
+  graph.capture_end();
+  EXPECT_EQ(is_capture_end_called, true);
+  EXPECT_EQ(graph.pool(), (at::MempoolId_t{10, 0}));
+  auto graph1 = at::accelerator::Graph();
+  EXPECT_EQ(keep_raw_graph, false);
+  EXPECT_EQ(graph_debug, false);
+  graph1.enable_debug_mode();
+  EXPECT_EQ(graph_debug, true);
+  ASSERT_THROW(graph1.debug_dump("abc"), c10::Error);
+}
diff --git a/aten/src/ATen/test/cuda_distributions_test.cu b/aten/src/ATen/test/cuda_distributions_test.cu
index 30de06081ee03..9771a269e7608 100644
--- a/aten/src/ATen/test/cuda_distributions_test.cu
+++ b/aten/src/ATen/test/cuda_distributions_test.cu
@@ -46,7 +46,7 @@ void assert_with_expected_uniforms(uint64_t counter_offset) {
   C10_CUDA_KERNEL_LAUNCH_CHECK();
 
   // Wait for GPU to finish before accessing on host
-  cudaDeviceSynchronize();
+  ASSERT_EQ(cudaSuccess, cudaDeviceSynchronize());
 
   // get 4 new float from uniform_()
   auto self = at::empty({4}, at::TensorOptions(at::kCUDA));
@@ -172,10 +172,10 @@ TEST(RandomPermutationTest, TestIslandShuffle) {
   bool shuffled1 = false;
   bool shuffled2 = false;
   for (int i = 0; i < 100; i++) {
-    cudaDeviceSynchronize();
+    ASSERT_EQ(cudaSuccess, cudaDeviceSynchronize());
     std::optional<at::Generator> gen = std::nullopt;
     randperm_handle_duplicate_keys(keys, values, 8, 5, gen);
-    cudaDeviceSynchronize();
+    ASSERT_EQ(cudaSuccess, cudaDeviceSynchronize());
     std::vector<int> slice1 = {values[0], values[1], values[2]};
     std::vector<int> slice2 = {values[3], values[4]};
     if (slice1 != valid_perms1[0]) {
diff --git a/aten/src/ATen/test/cuda_generator_test.cu b/aten/src/ATen/test/cuda_generator_test.cu
index f82db6de1d5b8..8da5feddc4000 100644
--- a/aten/src/ATen/test/cuda_generator_test.cu
+++ b/aten/src/ATen/test/cuda_generator_test.cu
@@ -103,7 +103,7 @@ TEST(CUDAGeneratorImpl, TestPhiloxEngineOffset2) {
   //   Assert that engine2 should be increment_val+1 steps behind engine1.
   if (!at::cuda::is_available()) return;
   test_engine_offset2();
-  cudaDeviceSynchronize();
+  ASSERT_EQ(cudaSuccess, cudaDeviceSynchronize());
   bool isEQ = cudaGetLastError() == cudaSuccess;
   ASSERT_TRUE(isEQ);
 }
@@ -130,7 +130,7 @@ TEST(CUDAGeneratorImpl, TestPhiloxEngineOffset3) {
   //   Assert that engine1 is 1 step behind engine2.
   if (!at::cuda::is_available()) return;
   test_engine_offset3();
-  cudaDeviceSynchronize();
+  ASSERT_EQ(cudaSuccess, cudaDeviceSynchronize());
   bool isEQ = cudaGetLastError() == cudaSuccess;
   ASSERT_TRUE(isEQ);
 }
@@ -154,7 +154,7 @@ TEST(CUDAGeneratorImpl, TestPhiloxEngineIndex) {
   //   Assert that the engines have different sequences.
   if (!at::cuda::is_available()) return;
   test_engine_thread_index();
-  cudaDeviceSynchronize();
+  ASSERT_EQ(cudaSuccess, cudaDeviceSynchronize());
   bool isEQ = cudaGetLastError() == cudaSuccess;
   ASSERT_TRUE(isEQ);
 }
diff --git a/aten/src/ATen/test/cuda_vectorized_test.cu b/aten/src/ATen/test/cuda_vectorized_test.cu
index 1b3ed4dc4ac42..2af8fc93a046d 100644
--- a/aten/src/ATen/test/cuda_vectorized_test.cu
+++ b/aten/src/ATen/test/cuda_vectorized_test.cu
@@ -94,7 +94,7 @@ TEST(TestVectorizedMemoryAccess, CopyKernel) {
 
   // vec4 copy
   reset_buffers();
-  cudaDeviceSynchronize();
+  ASSERT_EQ(cudaSuccess, cudaDeviceSynchronize());
   constexpr int total_work_size = buffer_size * 4;
   vectorized_copy<double, 4><<<total_work_size / block_work_size() , num_threads()>>>(b2, b1);
   C10_CUDA_KERNEL_LAUNCH_CHECK();
@@ -109,7 +109,7 @@ TEST(TestVectorizedMemoryAccess, CopyKernel) {
 
   // vec2 copy
   reset_buffers();
-  cudaDeviceSynchronize();
+  ASSERT_EQ(cudaSuccess, cudaDeviceSynchronize());
   vectorized_copy<double, 2><<<total_work_size / block_work_size() , num_threads()>>>(b2, b1);
   C10_CUDA_KERNEL_LAUNCH_CHECK();
 
@@ -123,7 +123,7 @@ TEST(TestVectorizedMemoryAccess, CopyKernel) {
 
   // vec1 copy
   reset_buffers();
-  cudaDeviceSynchronize();
+  ASSERT_EQ(cudaSuccess, cudaDeviceSynchronize());
   vectorized_copy<double, 1><<<total_work_size / block_work_size() , num_threads()>>>(b2, b1);
   C10_CUDA_KERNEL_LAUNCH_CHECK();
 
@@ -143,9 +143,9 @@ TEST(TestVectorizedMemoryAccess, CopyKernel) {
       b1 = reinterpret_cast<double *>(reinterpret_cast<char *>(buffer1) + i);
       b2 = reinterpret_cast<double *>(reinterpret_cast<char *>(buffer2) + j);
       (void)cudaGetLastError();
-      cudaDeviceSynchronize();
+      ASSERT_EQ(cudaSuccess, cudaDeviceSynchronize());
       vectorized_copy<double, 4><<<1, num_threads()>>>(b2, b1);
-      cudaDeviceSynchronize();
+      ASSERT_EQ(cudaSuccess, cudaDeviceSynchronize());
       auto err = cudaGetLastError();
       if (i % 16 == 0 && j % 16 == 0) {
         ASSERT_EQ(err, cudaSuccess);
diff --git a/aten/src/ATen/test/native_test.cpp b/aten/src/ATen/test/native_test.cpp
index 5f27ce4886e47..ff4507941a648 100644
--- a/aten/src/ATen/test/native_test.cpp
+++ b/aten/src/ATen/test/native_test.cpp
@@ -253,6 +253,9 @@ void test(TensorOptions T, TensorOptions AccT) {
 }
 
 TEST(TestNative, NativeTestCPU) {
+#if defined(__aarch64__)
+  GTEST_SKIP() << "Known failure on AArch64 (label is too far / stack test mismatch)";
+#endif
   manual_seed(123);
 
   test(at::device(kCPU).dtype(kFloat),
diff --git a/aten/src/ATen/test/test_install/CMakeLists.txt b/aten/src/ATen/test/test_install/CMakeLists.txt
index a07044a9ce3bb..10d6805a092cb 100644
--- a/aten/src/ATen/test/test_install/CMakeLists.txt
+++ b/aten/src/ATen/test/test_install/CMakeLists.txt
@@ -1,10 +1,7 @@
-cmake_minimum_required(VERSION 3.5)
+cmake_minimum_required(VERSION 3.10)
 find_package(ATen REQUIRED)
 include_directories(${ATEN_INCLUDE_DIR})
 
-# C++17
-if(not MSVC)
-    set(CMAKE_CXX_FLAGS "--std=c++17 ${CMAKE_CXX_FLAGS}")
-endif()
+set(CMAKE_CXX_STANDARD 20 CACHE STRING "The C++ standard whose features are requested to build this target.")
 add_executable(main main.cpp)
 target_link_libraries(main ${ATEN_LIBRARIES})
diff --git a/aten/src/ATen/test/vec_test_all_types.cpp b/aten/src/ATen/test/vec_test_all_types.cpp
index 69cafabb51a57..08436ae8f8c7e 100644
--- a/aten/src/ATen/test/vec_test_all_types.cpp
+++ b/aten/src/ATen/test/vec_test_all_types.cpp
@@ -381,7 +381,7 @@ namespace {
     TYPED_TEST(Hyperbolic, Tanh) {
         using vec = TypeParam;
 // NOTE: Because SVE uses ACL logic, the precision changes, hence the adjusted tolerance.
-#if defined(CPU_CAPABILITY_SVE)
+#if defined(CPU_CAPABILITY_SVE256)
         using UVT = UvalueType<vec>;
         UVT tolerance = getDefaultTolerance<UVT>();
         test_unary<vec>(
@@ -621,14 +621,14 @@ namespace {
         }
       }
     }
-#if defined(CPU_CAPABILITY_SVE) && defined(__ARM_FEATURE_BF16)
+#if defined(CPU_CAPABILITY_SVE256) && defined(__ARM_FEATURE_BF16)
     TEST(NanBfloat16, IsNan) {
       for (unsigned int ii = 0; ii < 0xFFFF; ++ii) {
         c10::BFloat16 val(ii, c10::BFloat16::from_bits());
         bool expected = std::isnan(val);
-        CACHE_ALIGN c10::BFloat16 actual_vals[at::vec::SVE256::Vectorized<c10::BFloat16>::size()];
-        at::vec::SVE256::Vectorized<c10::BFloat16>(val).isnan().store(actual_vals);
-        for (int jj = 0; jj < at::vec::SVE256::Vectorized<c10::BFloat16>::size(); ++jj) {
+        CACHE_ALIGN c10::BFloat16 actual_vals[at::vec::CPU_CAPABILITY::Vectorized<c10::BFloat16>::size()];
+        at::vec::CPU_CAPABILITY::Vectorized<c10::BFloat16>(val).isnan().store(actual_vals);
+        for (int jj = 0; jj < at::vec::CPU_CAPABILITY::Vectorized<c10::BFloat16>::size(); ++jj) {
           EXPECT_EQ(expected, c10::bit_cast<uint16_t>(actual_vals[jj]) != 0) << "bf16 isnan failure for bit pattern " << std::hex << ii << std::dec;
         }
       }
@@ -2196,7 +2196,12 @@ namespace {
       ASSERT_TRUE(vec_pinf.has_inf_nan()) << "Test failed for positive Infinity\n";
       ASSERT_TRUE(vec_ninf.has_inf_nan()) << "Test failed for negative Infinity\n";
     }
-#if !defined(CPU_CAPABILITY_SVE)
+// Building below for SVE with gcc-13 fails with following internal error
+// during GIMPLE pass: sink
+// In member function ‘virtual void {anonymous}::VecConvertBFloat16_ExhaustiveToFloat_Test::TestBody()’:
+// vec_test_all_types.cpp:2265:10: internal compiler error: Segmentation fault
+// 2265 |     TEST(VecConvertBFloat16, ExhaustiveToFloat) {
+#if !defined(CPU_CAPABILITY_SVE256)
     template <typename vec, typename dst_t>
     void test_convert_to(const char* dst_t_name) {
       using src_t = ValueType<vec>;
@@ -2335,7 +2340,7 @@ namespace {
     #undef TEST_MASK_LOAD
     #undef TEST_MASK_LOAD_N
     }
-#if !defined(CPU_CAPABILITY_SVE)
+#if !defined(CPU_CAPABILITY_SVE256)
     TYPED_TEST(VecMaskTests, MaskedCheck) {
       using VT = ValueType<TypeParam>;
       using vec = TypeParam;
@@ -2360,7 +2365,7 @@ namespace {
     #undef TEST_MASK_CHECK_N
     }
 #endif
-#if !defined(CPU_CAPABILITY_SVE)
+#if !defined(CPU_CAPABILITY_SVE256)
     TYPED_TEST(VecMaskTests, ToFrom) {
       using vec = TypeParam;
       using VT = ValueType<TypeParam>;
@@ -2387,7 +2392,7 @@ namespace {
       }
     }
 #endif
-#if !defined(CPU_CAPABILITY_SVE)
+#if !defined(CPU_CAPABILITY_SVE256)
     TYPED_TEST(VecMaskTests, Cast) {
       using vec = TypeParam;
       using src_t = ValueType<TypeParam>;
diff --git a/aten/src/ATen/test/vec_test_all_types.h b/aten/src/ATen/test/vec_test_all_types.h
index 7ee77f53d5377..dd2d68e77c66b 100644
--- a/aten/src/ATen/test/vec_test_all_types.h
+++ b/aten/src/ATen/test/vec_test_all_types.h
@@ -56,7 +56,7 @@ CACHE_ALIGN #define
   defined(CPU_CAPABILITY_AVX512) && (defined(__GNUC__) || defined(__GNUG__))
 #undef CHECK_DEQUANT_WITH_LOW_PRECISION
 #define CHECK_WITH_FMA 1
-#elif defined(CPU_CAPABILITY_SVE)
+#elif defined(CPU_CAPABILITY_SVE256)
 #define CHECK_DEQUANT_WITH_LOW_PRECISION 1
 #define CHECK_WITH_FMA 1
 #elif !defined(CPU_CAPABILITY_VSX) && !defined(CPU_CAPABILITY_AVX2)
diff --git a/aten/src/ATen/test/vitals.cpp b/aten/src/ATen/test/vitals.cpp
deleted file mode 100644
index eaf1cc152bc37..0000000000000
--- a/aten/src/ATen/test/vitals.cpp
+++ /dev/null
@@ -1,96 +0,0 @@
-#include <gmock/gmock.h>
-#include <gtest/gtest.h>
-
-#include <ATen/ATen.h>
-#include <ATen/core/Vitals.h>
-#include <c10/util/env.h>
-#include <c10/util/irange.h>
-#include <cstdlib>
-
-using namespace at::vitals;
-using ::testing::HasSubstr;
-
-TEST(Vitals, Basic) {
-  std::stringstream buffer;
-
-  std::streambuf* sbuf = std::cout.rdbuf();
-  std::cout.rdbuf(buffer.rdbuf());
-  {
-    c10::utils::set_env("TORCH_VITAL", "1");
-    TORCH_VITAL_DEFINE(Testing);
-    TORCH_VITAL(Testing, Attribute0) << 1;
-    TORCH_VITAL(Testing, Attribute1) << '1';
-    TORCH_VITAL(Testing, Attribute2) << 1.0f;
-    TORCH_VITAL(Testing, Attribute3) << 1.0;
-    auto t = at::ones({1, 1});
-    TORCH_VITAL(Testing, Attribute4) << t;
-  }
-  std::cout.rdbuf(sbuf);
-
-  auto s = buffer.str();
-  ASSERT_THAT(s, HasSubstr("Testing.Attribute0\t\t 1"));
-  ASSERT_THAT(s, HasSubstr("Testing.Attribute1\t\t 1"));
-  ASSERT_THAT(s, HasSubstr("Testing.Attribute2\t\t 1"));
-  ASSERT_THAT(s, HasSubstr("Testing.Attribute3\t\t 1"));
-  ASSERT_THAT(s, HasSubstr("Testing.Attribute4\t\t  1"));
-}
-
-TEST(Vitals, MultiString) {
-  std::stringstream buffer;
-
-  std::streambuf* sbuf = std::cout.rdbuf();
-  std::cout.rdbuf(buffer.rdbuf());
-  {
-    c10::utils::set_env("TORCH_VITAL", "1");
-    TORCH_VITAL_DEFINE(Testing);
-    TORCH_VITAL(Testing, Attribute0) << 1 << " of " << 2;
-    TORCH_VITAL(Testing, Attribute1) << 1;
-    TORCH_VITAL(Testing, Attribute1) << " of ";
-    TORCH_VITAL(Testing, Attribute1) << 2;
-  }
-  std::cout.rdbuf(sbuf);
-
-  auto s = buffer.str();
-  ASSERT_THAT(s, HasSubstr("Testing.Attribute0\t\t 1 of 2"));
-  ASSERT_THAT(s, HasSubstr("Testing.Attribute1\t\t 1 of 2"));
-}
-
-TEST(Vitals, OnAndOff) {
-  for (const auto i : c10::irange(2)) {
-    std::stringstream buffer;
-
-    std::streambuf* sbuf = std::cout.rdbuf();
-    std::cout.rdbuf(buffer.rdbuf());
-    {
-      c10::utils::set_env("TORCH_VITAL", i ? "1" : "0");
-      TORCH_VITAL_DEFINE(Testing);
-      TORCH_VITAL(Testing, Attribute0) << 1;
-    }
-    std::cout.rdbuf(sbuf);
-
-    auto s = buffer.str();
-    auto f = s.find("Testing.Attribute0\t\t 1");
-    if (i) {
-      ASSERT_TRUE(f != std::string::npos);
-    } else {
-      ASSERT_TRUE(f == std::string::npos);
-    }
-  }
-}
-
-TEST(Vitals, APIVitals) {
-  std::stringstream buffer;
-  bool rvalue = false;
-  std::streambuf* sbuf = std::cout.rdbuf();
-  std::cout.rdbuf(buffer.rdbuf());
-  {
-    c10::utils::set_env("TORCH_VITAL", "1");
-    APIVitals api_vitals;
-    rvalue = api_vitals.setVital("TestingSetVital", "TestAttr", "TestValue");
-  }
-  std::cout.rdbuf(sbuf);
-
-  auto s = buffer.str();
-  ASSERT_TRUE(rvalue);
-  ASSERT_THAT(s, HasSubstr("TestingSetVital.TestAttr\t\t TestValue"));
-}
diff --git a/aten/src/ATen/xpu/CachingHostAllocator.cpp b/aten/src/ATen/xpu/CachingHostAllocator.cpp
index 2be5eaf81bb03..c6a1a807801e2 100644
--- a/aten/src/ATen/xpu/CachingHostAllocator.cpp
+++ b/aten/src/ATen/xpu/CachingHostAllocator.cpp
@@ -43,8 +43,7 @@ struct XPUCachingHostAllocatorImpl
   }
 
   bool stream_is_capturing(XPUStream s) const override {
-    return c10::xpu::CaptureStatus(s.queue().ext_oneapi_get_state()) ==
-        c10::xpu::CaptureStatus::Recording;
+    return s.is_capturing();
   }
 };
 
diff --git a/aten/src/ATen/xpu/PeerToPeerAccess.h b/aten/src/ATen/xpu/PeerToPeerAccess.h
index 01812c8f6306c..ae7457e09c712 100644
--- a/aten/src/ATen/xpu/PeerToPeerAccess.h
+++ b/aten/src/ATen/xpu/PeerToPeerAccess.h
@@ -1,15 +1,24 @@
 #pragma once
 
-#include <c10/core/Device.h>
-#include <c10/macros/Macros.h>
+#include <c10/xpu/PeerToPeerAccess.h>
+
+#include <ATen/Context.h>
 
 namespace at::xpu {
 namespace detail {
-void init_p2p_access_cache(c10::DeviceIndex num_devices);
+// Initialize the peer-to-peer access cache for XPU devices.
+inline void init_p2p_access_cache(c10::DeviceIndex num_devices) {
+  c10::xpu::detail::init_p2p_access_cache(num_devices);
+}
 } // namespace detail
 
-TORCH_XPU_API bool get_p2p_access(
+// Query if peer-to-peer access is available between two devices.
+// This wrapper ensures XPU lazy initialization before forwarding to c10.
+inline bool get_p2p_access(
     c10::DeviceIndex dev,
-    c10::DeviceIndex dev_to_access);
+    c10::DeviceIndex dev_to_access) {
+  at::globalContext().lazyInitDevice(c10::DeviceType::XPU);
+  return c10::xpu::get_p2p_access(dev, dev_to_access);
+}
 
 } // namespace at::xpu
diff --git a/aten/src/ATen/xpu/XPUGeneratorImpl.cpp b/aten/src/ATen/xpu/XPUGeneratorImpl.cpp
index 09c84eeef9cea..133a9e4728dd0 100644
--- a/aten/src/ATen/xpu/XPUGeneratorImpl.cpp
+++ b/aten/src/ATen/xpu/XPUGeneratorImpl.cpp
@@ -97,7 +97,7 @@ void XPUGeneratorState::increase(uint64_t increment) {
 }
 
 // State can be used by multiple graph
-void XPUGeneratorState::register_graph(xpu::XPUGraph* graph) {
+void XPUGeneratorState::register_graph(xpu::XPUGraphImpl* graph) {
   // Ensures that the RNG state is not currently being captured.
   at::xpu::assertNotCapturing(
       "Cannot register the state during capturing stage.");
@@ -113,7 +113,7 @@ void XPUGeneratorState::register_graph(xpu::XPUGraph* graph) {
   }
 }
 
-void XPUGeneratorState::unregister_graph(xpu::XPUGraph* graph) {
+void XPUGeneratorState::unregister_graph(xpu::XPUGraphImpl* graph) {
   TORCH_CHECK(
       registered_graphs_.find(graph) != registered_graphs_.end(),
       "The graph should be registered to the state");
@@ -280,12 +280,12 @@ uint64_t XPUGeneratorImpl::philox_offset_per_thread() const {
   }
 }
 
-void XPUGeneratorImpl::register_graph(xpu::XPUGraph* graph) {
+void XPUGeneratorImpl::register_graph(xpu::XPUGraphImpl* graph) {
   graph->register_generator_state(state_);
   state_->register_graph(graph);
 }
 
-void XPUGeneratorImpl::unregister_graph(xpu::XPUGraph* graph) {
+void XPUGeneratorImpl::unregister_graph(xpu::XPUGraphImpl* graph) {
   state_->unregister_graph(graph);
 }
 
diff --git a/aten/src/ATen/xpu/XPUGeneratorImpl.h b/aten/src/ATen/xpu/XPUGeneratorImpl.h
index 8ee4967a4f13f..853d1f28f006a 100644
--- a/aten/src/ATen/xpu/XPUGeneratorImpl.h
+++ b/aten/src/ATen/xpu/XPUGeneratorImpl.h
@@ -8,7 +8,7 @@
 namespace at {
 
 namespace xpu {
-struct XPUGraph;
+struct XPUGraphImpl;
 }
 
 struct XPUGeneratorState : public c10::intrusive_ptr_target {
@@ -16,7 +16,7 @@ struct XPUGeneratorState : public c10::intrusive_ptr_target {
   uint64_t philox_offset_per_thread_;
   uint32_t offset_intragraph_;
   bool capturing_{};
-  std::unordered_set<xpu::XPUGraph*> registered_graphs_;
+  std::unordered_set<xpu::XPUGraphImpl*> registered_graphs_;
   at::TensorBase seed_extragraph_{};
   at::TensorBase offset_extragraph_{};
 
@@ -29,8 +29,8 @@ struct XPUGeneratorState : public c10::intrusive_ptr_target {
         offset_intragraph_(offset_intragraph) {}
 
   void increase(uint64_t increment);
-  void register_graph(xpu::XPUGraph* graph);
-  void unregister_graph(xpu::XPUGraph* graph);
+  void register_graph(xpu::XPUGraphImpl* graph);
+  void unregister_graph(xpu::XPUGraphImpl* graph);
   void capture_prologue();
   uint64_t capture_epilogue();
   void replay_prologue(uint64_t wholegraph_increment);
@@ -62,8 +62,8 @@ struct TORCH_XPU_API XPUGeneratorImpl : public GeneratorImpl {
   void set_philox_offset_per_thread(uint64_t offset);
   uint64_t philox_offset_per_thread() const;
 
-  void register_graph(xpu::XPUGraph* graph);
-  void unregister_graph(xpu::XPUGraph* graph);
+  void register_graph(xpu::XPUGraphImpl* graph);
+  void unregister_graph(xpu::XPUGraphImpl* graph);
   PhiloxXpuState philox_xpu_state(uint64_t increment);
   std::pair<uint64_t, uint64_t> philox_engine_inputs(uint64_t increment);
   static c10::DeviceType device_type();
diff --git a/aten/src/ATen/xpu/XPUGraph.cpp b/aten/src/ATen/xpu/XPUGraph.cpp
index 3d3a22e5b7219..d6d3cad494a04 100644
--- a/aten/src/ATen/xpu/XPUGraph.cpp
+++ b/aten/src/ATen/xpu/XPUGraph.cpp
@@ -1,6 +1,5 @@
 #include <ATen/Functions.h>
 #include <ATen/core/CachingHostAllocator.h>
-#include <ATen/xpu/XPUGeneratorImpl.h>
 #include <ATen/xpu/XPUGraph.h>
 #include <c10/xpu/XPUFunctions.h>
 
@@ -16,23 +15,43 @@ MempoolId_t graph_pool_handle() {
   return c10::xpu::MemPool::graph_pool_handle();
 }
 
-XPUGraph::XPUGraph(bool keep_graph)
+XPUGraphImpl::XPUGraphImpl(const GraphImplArgs& args)
     : capture_stream_(at::xpu::getCurrentXPUStream()),
-      keep_graph_(keep_graph) {}
+      keep_graph_(args.keep_graph) {}
 
-void XPUGraph::register_generator_state(
+void XPUGraphImpl::register_generator_state(
     c10::intrusive_ptr<at::XPUGeneratorState> state) {
   captured_generator_states_[std::move(state)] = 0;
 }
 
-void XPUGraph::register_generator_state(const at::Generator& generator) {
+void XPUGraphImpl::register_generator_state(const at::Generator& generator) {
   c10::intrusive_ptr<XPUGeneratorImpl> xpu_gen =
       dynamic_intrusive_pointer_cast<XPUGeneratorImpl>(
           generator.getIntrusivePtr());
   xpu_gen->register_graph(this);
 }
 
-void XPUGraph::capture_begin(MempoolId_t pool) {
+void XPUGraphImpl::capture_begin(
+    MempoolId_t pool /*={0,0}*/,
+    GraphCaptureMode capture_mode) {
+  switch (capture_mode) {
+    case GraphCaptureMode::Default:
+      break;
+
+    case GraphCaptureMode::Global:
+    case GraphCaptureMode::ThreadLocal:
+    case GraphCaptureMode::Relaxed:
+      TORCH_WARN(
+          "XPUGraph currently only support default GraphCaptureMode. "
+          "Falling back to default capture behavior.");
+      break;
+
+    default:
+      TORCH_CHECK(
+          false,
+          "Invalid GraphCaptureMode value: ",
+          static_cast<int>(capture_mode));
+  }
   TORCH_CHECK(
       !has_graph_exec_,
       "This XPUGraph instance already owns a captured graph. "
@@ -90,7 +109,7 @@ void XPUGraph::capture_begin(MempoolId_t pool) {
       capture_stream_.queue().ext_oneapi_get_state() == queue_state::recording);
 }
 
-void XPUGraph::capture_end() {
+void XPUGraphImpl::capture_end() {
   auto stream = at::xpu::getCurrentXPUStream();
 
   TORCH_CHECK(
@@ -125,7 +144,7 @@ void XPUGraph::capture_end() {
   }
 }
 
-void XPUGraph::instantiate() {
+void XPUGraphImpl::instantiate() {
   TORCH_CHECK(
       capture_ended_,
       "capture_end() must have been called before calling instantiate");
@@ -141,7 +160,7 @@ void XPUGraph::instantiate() {
   has_graph_exec_ = true;
 }
 
-void XPUGraph::replay() {
+void XPUGraphImpl::replay() {
   TORCH_CHECK(
       capture_ended_,
       "Called XPUGraph::replay without a preceding successful capture.");
@@ -164,7 +183,7 @@ void XPUGraph::replay() {
   queue.ext_oneapi_graph(*graph_exec_);
 }
 
-void XPUGraph::reset() {
+void XPUGraphImpl::reset() {
   if (capture_ended_) {
     c10::xpu::XPUCachingAllocator::releasePool(capture_dev_, mempool_id_);
     at::getHostAllocator(at::kXPU)->release_pool(mempool_id_);
@@ -180,11 +199,11 @@ void XPUGraph::reset() {
   }
 }
 
-void XPUGraph::enable_debug_mode() {
+void XPUGraphImpl::enable_debug_mode() {
   _xpu_graphs_debug = true;
 }
 
-void XPUGraph::debug_dump(const std::string& debug_path) {
+void XPUGraphImpl::debug_dump(const std::string& debug_path) {
   TORCH_CHECK(
       debug_path.size() >= 4 &&
           debug_path.substr(debug_path.size() - 4) == ".dot",
@@ -207,7 +226,7 @@ void XPUGraph::debug_dump(const std::string& debug_path) {
   }
 }
 
-xpuGraph_t* XPUGraph::raw_xpu_graph() {
+xpuGraph_t* XPUGraphImpl::raw_xpu_graph() {
   TORCH_CHECK(
       keep_graph_,
       "You cannot access the raw xpuGraph_t instance unless XPUGraph was initialized with keep_graph=true");
@@ -217,7 +236,7 @@ xpuGraph_t* XPUGraph::raw_xpu_graph() {
   return graph_.get();
 }
 
-xpuGraphExec_t* XPUGraph::raw_xpu_graph_exec() {
+xpuGraphExec_t* XPUGraphImpl::raw_xpu_graph_exec() {
   TORCH_CHECK(
       has_graph_exec_,
       "You cannot access the raw xpuGraphExec_t instance until instantiate() has been called");
@@ -226,14 +245,14 @@ xpuGraphExec_t* XPUGraph::raw_xpu_graph_exec() {
 
 // Returns an id another graph's capture_begin can use to share the same memory
 // pool as this graph.
-MempoolId_t XPUGraph::pool() {
+MempoolId_t XPUGraphImpl::pool() const {
   TORCH_CHECK(
       capture_ended_,
       "Called XPUGraph::pool() without a preceding successful capture.");
   return mempool_id_;
 }
 
-XPUGraph::~XPUGraph() {
+XPUGraphImpl::~XPUGraphImpl() {
   for (auto& [generator_state, wholegraph_increments] :
        captured_generator_states_) {
     generator_state->unregister_graph(this);
@@ -241,4 +260,6 @@ XPUGraph::~XPUGraph() {
   reset();
 }
 
+REGISTER_GRAPH_IMPL(XPU, XPUGraphImpl)
+
 } // namespace at::xpu
diff --git a/aten/src/ATen/xpu/XPUGraph.h b/aten/src/ATen/xpu/XPUGraph.h
index 0f23a1e093ee0..fd60bb5b6ba9e 100644
--- a/aten/src/ATen/xpu/XPUGraph.h
+++ b/aten/src/ATen/xpu/XPUGraph.h
@@ -1,18 +1,15 @@
 #pragma once
 
 #include <ATen/Tensor.h>
+#include <ATen/core/GraphImplInterface.h>
+#include <ATen/xpu/XPUGeneratorImpl.h>
 #include <c10/core/Device.h>
 #include <c10/util/flat_hash_map.h>
 #include <c10/xpu/XPUCachingAllocator.h>
 #include <c10/xpu/XPUGraphsC10Utils.h>
 #include <c10/xpu/XPUStream.h>
 
-namespace at {
-
-struct Generator;
-struct XPUGeneratorState;
-
-namespace xpu {
+namespace at::xpu {
 
 TORCH_XPU_API MempoolId_t graph_pool_handle();
 
@@ -21,21 +18,26 @@ using xpuGraph_t = sycl::ext::oneapi::experimental::command_graph<
 using xpuGraphExec_t = sycl::ext::oneapi::experimental::command_graph<
     sycl::ext::oneapi::experimental::graph_state::executable>;
 
-struct TORCH_XPU_API XPUGraph {
-  XPUGraph(bool keep_graph = false);
-  ~XPUGraph();
+struct TORCH_XPU_API XPUGraphImpl : public at::GraphImplInterface {
+  XPUGraphImpl(const GraphImplArgs& args = {});
+  ~XPUGraphImpl() override;
+
+  C10_DISABLE_COPY_AND_ASSIGN(XPUGraphImpl);
 
   void register_generator_state(
       c10::intrusive_ptr<at::XPUGeneratorState> state);
   void register_generator_state(const at::Generator& generator);
-  void capture_begin(MempoolId_t pool = {0, 0});
-  void capture_end();
-  void instantiate();
-  void replay();
-  void reset();
-  MempoolId_t pool();
-  void enable_debug_mode();
-  void debug_dump(const std::string& debug_path);
+
+  void capture_begin(
+      MempoolId_t pool = {0, 0},
+      GraphCaptureMode capture_mode = GraphCaptureMode::Default) override;
+  void capture_end() override;
+  void instantiate() override;
+  void replay() override;
+  void reset() override;
+  MempoolId_t pool() const override;
+  void enable_debug_mode() override;
+  void debug_dump(const std::string& debug_path) override;
   xpuGraph_t* raw_xpu_graph();
   xpuGraphExec_t* raw_xpu_graph_exec();
 
@@ -59,5 +61,58 @@ struct TORCH_XPU_API XPUGraph {
   bool keep_graph_;
 };
 
-} // namespace xpu
-} // namespace at
+struct TORCH_XPU_API XPUGraph {
+  XPUGraph(bool keep_graph = false) {
+    GraphImplArgs args;
+    args.keep_graph = keep_graph;
+    impl_ = std::make_unique<XPUGraphImpl>(args);
+  }
+  ~XPUGraph() = default;
+
+  C10_DISABLE_COPY_AND_ASSIGN(XPUGraph);
+  XPUGraph(XPUGraph&& other) = delete;
+  XPUGraph& operator=(XPUGraph&& other) = delete;
+
+  void register_generator_state(
+      c10::intrusive_ptr<at::XPUGeneratorState> state) {
+    impl_->register_generator_state(state);
+  }
+  void register_generator_state(const at::Generator& generator) {
+    impl_->register_generator_state(generator);
+  }
+  void capture_begin(MempoolId_t pool = {0, 0}) {
+    impl_->capture_begin(pool);
+  }
+  void capture_end() {
+    impl_->capture_end();
+  }
+  void instantiate() {
+    impl_->instantiate();
+  }
+  void replay() {
+    impl_->replay();
+  }
+  void reset() {
+    impl_->reset();
+  }
+  MempoolId_t pool() const {
+    return impl_->pool();
+  }
+  void enable_debug_mode() {
+    impl_->enable_debug_mode();
+  }
+  void debug_dump(const std::string& debug_path) {
+    impl_->debug_dump(debug_path);
+  }
+  xpuGraph_t* raw_xpu_graph() {
+    return impl_->raw_xpu_graph();
+  }
+  xpuGraphExec_t* raw_xpu_graph_exec() {
+    return impl_->raw_xpu_graph_exec();
+  }
+
+ private:
+  std::unique_ptr<XPUGraphImpl> impl_;
+};
+
+} // namespace at::xpu
diff --git a/aten/tools/run_tests.sh b/aten/tools/run_tests.sh
index 93e48aec90851..2c57d428f58ea 100755
--- a/aten/tools/run_tests.sh
+++ b/aten/tools/run_tests.sh
@@ -55,6 +55,7 @@ run_if_exists cuda_allocator_test
 if [ "$VALGRIND" == "ON" ]; then
   # NB: As these tests are invoked by valgrind, let's leave them for now as it's
   # unclear if valgrind -> python -> gtest would work
+  export LD_LIBRARY_PATH="${CPP_TESTS_DIR}:${LD_LIBRARY_PATH}"
   valgrind --suppressions="$VALGRIND_SUP" --error-exitcode=1 "${CPP_TESTS_DIR}/basic" --gtest_filter='-*CUDA'
   if [[ -x ${CPP_TESTS_DIR}/tensor_interop_test ]]; then
     valgrind --suppressions="$VALGRIND_SUP" --error-exitcode=1 "${CPP_TESTS_DIR}/tensor_interop_test"
diff --git a/aten/tools/valgrind.sup b/aten/tools/valgrind.sup
index 585487c4d2be2..c86c374bd1307 100644
--- a/aten/tools/valgrind.sup
+++ b/aten/tools/valgrind.sup
@@ -10,6 +10,16 @@
    ...
 }
 
+{
+   <ld_strncmp_rpath_addr8>
+   Memcheck:Addr8
+   fun:strncmp
+   fun:is_dst
+   ...
+   fun:decompose_rpath
+   ...
+}
+
 {
    ignore_empty_generic_uninitialised_conditional_jump
    Memcheck:Cond
diff --git a/benchmarks/data/samplers_benchmark.py b/benchmarks/data/samplers_benchmark.py
index 6cdd0c77c6519..6dde875d218c8 100644
--- a/benchmarks/data/samplers_benchmark.py
+++ b/benchmarks/data/samplers_benchmark.py
@@ -2,7 +2,6 @@
 
 import time
 from collections.abc import Iterable, Iterator
-from typing import Union
 
 import numpy as np
 from tabulate import tabulate
@@ -15,7 +14,7 @@ class NewBatchSampler(Sampler[list[int]]):
 
     def __init__(
         self,
-        sampler: Union[Sampler[int], Iterable[int]],
+        sampler: Sampler[int] | Iterable[int],
         batch_size: int,
         drop_last: bool,
     ) -> None:
diff --git a/benchmarks/diffusion/compile_benchmark.py b/benchmarks/diffusion/compile_benchmark.py
new file mode 100644
index 0000000000000..934636da28840
--- /dev/null
+++ b/benchmarks/diffusion/compile_benchmark.py
@@ -0,0 +1,210 @@
+"""
+Compile benchmarks for diffusion model transformers.
+
+These models have large numbers of repeated transformer blocks, leading to high
+cold-start compile times under full-model compilation. This makes them good
+candidates for regional/hierarchical compilation strategies that reduce compile
+time by compiling individual blocks instead of the full model.
+
+But if you are looking to reduce the cold start time for full models, these are
+excellent candidates.
+
+Example:
+    python compile_benchmark.py --model auroflow --mode full
+    python compile_benchmark.py --model flux --mode full --backend eager
+
+If you see issues downloading models, try: pip uninstall hf_xet
+"""
+
+import argparse
+import time
+
+import diffusers
+from diffusers import (
+    AuraFlowPipeline,
+    AuraFlowTransformer2DModel,
+    GGUFQuantizationConfig,
+)
+
+import torch
+from torch._higher_order_ops.invoke_subgraph import mark_compile_region
+
+
+def compile_model(model, mode, block_types, backend="inductor"):
+    if mode == "full":
+        model.compile(backend=backend, fullgraph=True, mode="reduce-overhead")
+    elif mode == "regional":
+        for submod in model.modules():
+            if isinstance(submod, block_types):
+                print("Compiling", submod.__class__)
+                submod.compile(backend=backend, fullgraph=True)
+    elif mode == "hierarchical":
+        for submod in model.modules():
+            if isinstance(submod, block_types):
+                submod.__class__.forward = mark_compile_region(submod.__class__.forward)
+        model.compile(backend=backend, fullgraph=True)
+
+
+def bench(run_fn, warmup_steps=1, bench_steps=50):
+    run_fn(warmup_steps)
+    t0 = time.perf_counter()
+    run_fn(bench_steps)
+    t1 = time.perf_counter()
+    print(f"{t1 - t0:.3f}s")
+
+
+def auroflow_benchmark(mode, backend="inductor"):
+    transformer = AuraFlowTransformer2DModel.from_single_file(
+        "https://huggingface.co/city96/AuraFlow-v0.3-gguf/blob/main/aura_flow_0.3-Q2_K.gguf",
+        quantization_config=GGUFQuantizationConfig(compute_dtype=torch.bfloat16),
+        torch_dtype=torch.bfloat16,
+    )
+    pipe = AuraFlowPipeline.from_pretrained(
+        "fal/AuraFlow-v0.3",
+        torch_dtype=torch.bfloat16,
+        transformer=transformer,
+    ).to("cuda")
+
+    block_types = (
+        diffusers.models.transformers.auraflow_transformer_2d.AuraFlowSingleTransformerBlock,
+        diffusers.models.transformers.auraflow_transformer_2d.AuraFlowJointTransformerBlock,
+    )
+    compile_model(pipe.transformer, mode, block_types, backend)
+
+    def run(steps):
+        pipe("A cute pony", width=512, height=512, num_inference_steps=steps)
+
+    bench(run)
+
+
+def wan_benchmark(mode, backend="inductor"):
+    import numpy as np
+    from diffusers import AutoencoderKLWan, WanImageToVideoPipeline
+    from diffusers.utils import load_image
+    from transformers import CLIPVisionModel
+
+    model_id = "Wan-AI/Wan2.1-I2V-14B-480P-Diffusers"
+    image_encoder = CLIPVisionModel.from_pretrained(
+        model_id, subfolder="image_encoder", torch_dtype=torch.float32
+    )
+    vae = AutoencoderKLWan.from_pretrained(
+        model_id, subfolder="vae", torch_dtype=torch.float32
+    )
+    pipe = WanImageToVideoPipeline.from_pretrained(
+        model_id, vae=vae, image_encoder=image_encoder, torch_dtype=torch.bfloat16
+    ).to("cuda")
+
+    image = load_image(
+        "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/astronaut.jpg"
+    )
+    max_area = 480 * 832
+    aspect_ratio = image.height / image.width
+    mod_value = pipe.vae_scale_factor_spatial * pipe.transformer.config.patch_size[1]
+    height = round(np.sqrt(max_area * aspect_ratio)) // mod_value * mod_value
+    width = round(np.sqrt(max_area / aspect_ratio)) // mod_value * mod_value
+    image = image.resize((width, height))
+
+    prompt = (
+        "An astronaut hatching from an egg, on the surface of the moon, the darkness and depth of space realised in "
+        "the background. High quality, ultrarealistic detail and breath-taking movie-like camera shot."
+    )  # noqa: B950
+    negative_prompt = "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards"  # noqa: B950
+
+    block_types = (diffusers.models.transformers.transformer_wan.WanTransformerBlock,)
+    compile_model(pipe.transformer, mode, block_types, backend)
+
+    def run(steps):
+        pipe(
+            image=image,
+            prompt=prompt,
+            negative_prompt=negative_prompt,
+            height=height,
+            width=width,
+            num_frames=33,
+            num_inference_steps=steps,
+            guidance_scale=5.0,
+        )
+
+    bench(run)
+
+
+def ltx_benchmark(mode, backend="inductor"):
+    from diffusers import LTXConditionPipeline
+
+    pipe = LTXConditionPipeline.from_pretrained(
+        "Lightricks/LTX-Video-0.9.7-dev", torch_dtype=torch.bfloat16
+    ).to("cuda")
+    pipe.vae.enable_tiling()
+
+    height = 512 - (512 % pipe.vae_spatial_compression_ratio)
+    width = 704 - (704 % pipe.vae_spatial_compression_ratio)
+
+    prompt = "The video depicts a winding mountain road covered in snow, with a single vehicle traveling along it. The road is flanked by steep, rocky cliffs and sparse vegetation. The landscape is characterized by rugged terrain and a river visible in the distance. The scene captures the solitude and beauty of a winter drive through a mountainous region."  # noqa: B950
+    negative_prompt = "worst quality, inconsistent motion, blurry, jittery, distorted"
+
+    block_types = (
+        diffusers.models.transformers.transformer_ltx.LTXVideoTransformerBlock,
+    )
+    compile_model(pipe.transformer, mode, block_types, backend)
+
+    def run(steps):
+        pipe(
+            conditions=None,
+            prompt=prompt,
+            negative_prompt=negative_prompt,
+            width=width,
+            height=height,
+            num_frames=121,
+            num_inference_steps=steps,
+            generator=torch.Generator().manual_seed(0),
+            output_type="latent",
+        )
+
+    bench(run)
+
+
+def flux_benchmark(mode, backend="inductor"):
+    from diffusers import FluxPipeline
+
+    pipe = FluxPipeline.from_pretrained(
+        "black-forest-labs/FLUX.1-dev",
+        torch_dtype=torch.bfloat16,
+    ).to("cuda")
+
+    prompt = "A cat holding a sign that says hello world"
+
+    block_types = (
+        diffusers.models.transformers.transformer_flux.FluxTransformerBlock,
+        diffusers.models.transformers.transformer_flux.FluxSingleTransformerBlock,
+    )
+    compile_model(pipe.transformer, mode, block_types, backend)
+
+    def run(steps):
+        pipe(
+            prompt,
+            height=1024,
+            width=1024,
+            guidance_scale=3.5,
+            num_inference_steps=steps,
+            max_sequence_length=512,
+        )
+
+    bench(run)
+
+
+BENCHMARKS = {
+    "auroflow": auroflow_benchmark,
+    "wan": wan_benchmark,
+    "ltx": ltx_benchmark,
+    "flux": flux_benchmark,
+}
+
+MODES = ("eager", "full", "regional", "hierarchical")
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--model", choices=BENCHMARKS, required=True)
+    parser.add_argument("--mode", choices=MODES, required=True)
+    parser.add_argument("--backend", default="inductor")
+    args = parser.parse_args()
+    BENCHMARKS[args.model](args.mode, args.backend)
diff --git a/benchmarks/dynamo/ci_expected_accuracy/aot_eager_huggingface_inference.csv b/benchmarks/dynamo/ci_expected_accuracy/aot_eager_huggingface_inference.csv
index d3f0fbba71826..87dd88078f222 100644
--- a/benchmarks/dynamo/ci_expected_accuracy/aot_eager_huggingface_inference.csv
+++ b/benchmarks/dynamo/ci_expected_accuracy/aot_eager_huggingface_inference.csv
@@ -50,7 +50,7 @@ LayoutLMForMaskedLM,pass,0
 
 
 
-M2M100ForConditionalGeneration,pass,0
+M2M100ForConditionalGeneration,pass,7
 
 
 
diff --git a/benchmarks/dynamo/ci_expected_accuracy/aot_eager_huggingface_training.csv b/benchmarks/dynamo/ci_expected_accuracy/aot_eager_huggingface_training.csv
index 82d21ea3cb298..5ca03b5ecf9fb 100644
--- a/benchmarks/dynamo/ci_expected_accuracy/aot_eager_huggingface_training.csv
+++ b/benchmarks/dynamo/ci_expected_accuracy/aot_eager_huggingface_training.csv
@@ -50,7 +50,7 @@ LayoutLMForMaskedLM,pass,5
 
 
 
-M2M100ForConditionalGeneration,pass,4
+M2M100ForConditionalGeneration,pass,11
 
 
 
diff --git a/benchmarks/dynamo/ci_expected_accuracy/aot_eager_torchbench_inference.csv b/benchmarks/dynamo/ci_expected_accuracy/aot_eager_torchbench_inference.csv
index bdf3313659b66..b0daa9a947ad4 100644
--- a/benchmarks/dynamo/ci_expected_accuracy/aot_eager_torchbench_inference.csv
+++ b/benchmarks/dynamo/ci_expected_accuracy/aot_eager_torchbench_inference.csv
@@ -202,10 +202,6 @@ pyhpc_turbulent_kinetic_energy,pass,0
 
 
 
-pytorch_CycleGAN_and_pix2pix,pass,0
-
-
-
 pytorch_stargan,pass,0
 
 
diff --git a/benchmarks/dynamo/ci_expected_accuracy/aot_eager_torchbench_training.csv b/benchmarks/dynamo/ci_expected_accuracy/aot_eager_torchbench_training.csv
index ee7838505d67c..9ff11382f67ad 100644
--- a/benchmarks/dynamo/ci_expected_accuracy/aot_eager_torchbench_training.csv
+++ b/benchmarks/dynamo/ci_expected_accuracy/aot_eager_torchbench_training.csv
@@ -130,10 +130,6 @@ phlippe_resnet,pass,6
 
 
 
-pytorch_CycleGAN_and_pix2pix,pass,6
-
-
-
 pytorch_stargan,pass,6
 
 
diff --git a/benchmarks/dynamo/ci_expected_accuracy/aot_inductor_torchbench_inference.csv b/benchmarks/dynamo/ci_expected_accuracy/aot_inductor_torchbench_inference.csv
index 6fdfc6e72bda3..5d9d432e644a7 100644
--- a/benchmarks/dynamo/ci_expected_accuracy/aot_inductor_torchbench_inference.csv
+++ b/benchmarks/dynamo/ci_expected_accuracy/aot_inductor_torchbench_inference.csv
@@ -178,10 +178,6 @@ pyhpc_turbulent_kinetic_energy,pass,0
 
 
 
-pytorch_CycleGAN_and_pix2pix,pass,0
-
-
-
 pytorch_stargan,pass,0
 
 
diff --git a/benchmarks/dynamo/ci_expected_accuracy/cpu_aot_inductor_amp_freezing_torchbench_inference.csv b/benchmarks/dynamo/ci_expected_accuracy/cpu_aot_inductor_amp_freezing_torchbench_inference.csv
index 061727b22329c..c9a93b51d10c7 100644
--- a/benchmarks/dynamo/ci_expected_accuracy/cpu_aot_inductor_amp_freezing_torchbench_inference.csv
+++ b/benchmarks/dynamo/ci_expected_accuracy/cpu_aot_inductor_amp_freezing_torchbench_inference.csv
@@ -162,10 +162,6 @@ pyhpc_turbulent_kinetic_energy,pass,0
 
 
 
-pytorch_CycleGAN_and_pix2pix,pass,0
-
-
-
 pytorch_stargan,pass,0
 
 
diff --git a/benchmarks/dynamo/ci_expected_accuracy/cpu_aot_inductor_freezing_torchbench_inference.csv b/benchmarks/dynamo/ci_expected_accuracy/cpu_aot_inductor_freezing_torchbench_inference.csv
index 061727b22329c..c9a93b51d10c7 100644
--- a/benchmarks/dynamo/ci_expected_accuracy/cpu_aot_inductor_freezing_torchbench_inference.csv
+++ b/benchmarks/dynamo/ci_expected_accuracy/cpu_aot_inductor_freezing_torchbench_inference.csv
@@ -162,10 +162,6 @@ pyhpc_turbulent_kinetic_energy,pass,0
 
 
 
-pytorch_CycleGAN_and_pix2pix,pass,0
-
-
-
 pytorch_stargan,pass,0
 
 
diff --git a/benchmarks/dynamo/ci_expected_accuracy/cpu_inductor_amp_freezing_huggingface_inference.csv b/benchmarks/dynamo/ci_expected_accuracy/cpu_inductor_amp_freezing_huggingface_inference.csv
index 439c9bf530468..e4aabce10466d 100644
--- a/benchmarks/dynamo/ci_expected_accuracy/cpu_inductor_amp_freezing_huggingface_inference.csv
+++ b/benchmarks/dynamo/ci_expected_accuracy/cpu_inductor_amp_freezing_huggingface_inference.csv
@@ -50,7 +50,7 @@ LayoutLMForMaskedLM,pass,0
 
 
 
-M2M100ForConditionalGeneration,pass,0
+M2M100ForConditionalGeneration,pass,7
 
 
 
diff --git a/benchmarks/dynamo/ci_expected_accuracy/cpu_inductor_amp_freezing_torchbench_inference.csv b/benchmarks/dynamo/ci_expected_accuracy/cpu_inductor_amp_freezing_torchbench_inference.csv
index fa64abf3c8246..8fd420c017d12 100644
--- a/benchmarks/dynamo/ci_expected_accuracy/cpu_inductor_amp_freezing_torchbench_inference.csv
+++ b/benchmarks/dynamo/ci_expected_accuracy/cpu_inductor_amp_freezing_torchbench_inference.csv
@@ -50,47 +50,47 @@ densenet121,pass,0
 
 
 
-detectron2_fasterrcnn_r_101_c4,fail_accuracy,42
+detectron2_fasterrcnn_r_101_c4,fail_accuracy,43
 
 
 
-detectron2_fasterrcnn_r_101_dc5,fail_accuracy,42
+detectron2_fasterrcnn_r_101_dc5,fail_accuracy,43
 
 
 
-detectron2_fasterrcnn_r_101_fpn,fail_accuracy,46
+detectron2_fasterrcnn_r_101_fpn,fail_accuracy,47
 
 
 
-detectron2_fasterrcnn_r_50_c4,fail_accuracy,42
+detectron2_fasterrcnn_r_50_c4,fail_accuracy,43
 
 
 
-detectron2_fasterrcnn_r_50_dc5,fail_accuracy,42
+detectron2_fasterrcnn_r_50_dc5,fail_accuracy,43
 
 
 
-detectron2_fasterrcnn_r_50_fpn,fail_accuracy,46
+detectron2_fasterrcnn_r_50_fpn,fail_accuracy,47
 
 
 
-detectron2_fcos_r_50_fpn,fail_accuracy,24
+detectron2_fcos_r_50_fpn,fail_accuracy,25
 
 
 
-detectron2_maskrcnn_r_101_c4,fail_accuracy,56
+detectron2_maskrcnn_r_101_c4,fail_accuracy,57
 
 
 
-detectron2_maskrcnn_r_101_fpn,fail_accuracy,62
+detectron2_maskrcnn_r_101_fpn,fail_accuracy,63
 
 
 
-detectron2_maskrcnn_r_50_c4,fail_accuracy,56
+detectron2_maskrcnn_r_50_c4,fail_accuracy,57
 
 
 
-detectron2_maskrcnn_r_50_fpn,fail_accuracy,62
+detectron2_maskrcnn_r_50_fpn,fail_accuracy,63
 
 
 
@@ -186,10 +186,6 @@ pyhpc_turbulent_kinetic_energy,pass,0
 
 
 
-pytorch_CycleGAN_and_pix2pix,pass,0
-
-
-
 pytorch_stargan,pass,0
 
 
@@ -246,7 +242,7 @@ vgg16,pass,0
 
 
 
-vision_maskrcnn,fail_accuracy,30
+vision_maskrcnn,fail_accuracy,31
 
 
 
diff --git a/benchmarks/dynamo/ci_expected_accuracy/cpu_inductor_freezing_huggingface_inference.csv b/benchmarks/dynamo/ci_expected_accuracy/cpu_inductor_freezing_huggingface_inference.csv
index 439c9bf530468..e4aabce10466d 100644
--- a/benchmarks/dynamo/ci_expected_accuracy/cpu_inductor_freezing_huggingface_inference.csv
+++ b/benchmarks/dynamo/ci_expected_accuracy/cpu_inductor_freezing_huggingface_inference.csv
@@ -50,7 +50,7 @@ LayoutLMForMaskedLM,pass,0
 
 
 
-M2M100ForConditionalGeneration,pass,0
+M2M100ForConditionalGeneration,pass,7
 
 
 
diff --git a/benchmarks/dynamo/ci_expected_accuracy/cpu_inductor_freezing_torchbench_inference.csv b/benchmarks/dynamo/ci_expected_accuracy/cpu_inductor_freezing_torchbench_inference.csv
index 1d5a70739dd79..813520528261e 100644
--- a/benchmarks/dynamo/ci_expected_accuracy/cpu_inductor_freezing_torchbench_inference.csv
+++ b/benchmarks/dynamo/ci_expected_accuracy/cpu_inductor_freezing_torchbench_inference.csv
@@ -186,10 +186,6 @@ pyhpc_turbulent_kinetic_energy,pass,0
 
 
 
-pytorch_CycleGAN_and_pix2pix,pass,0
-
-
-
 pytorch_stargan,pass,0
 
 
diff --git a/benchmarks/dynamo/ci_expected_accuracy/cpu_inductor_huggingface_inference.csv b/benchmarks/dynamo/ci_expected_accuracy/cpu_inductor_huggingface_inference.csv
index 439c9bf530468..2712e2d1a5012 100644
--- a/benchmarks/dynamo/ci_expected_accuracy/cpu_inductor_huggingface_inference.csv
+++ b/benchmarks/dynamo/ci_expected_accuracy/cpu_inductor_huggingface_inference.csv
@@ -30,7 +30,7 @@ DistilBertForMaskedLM,pass,0
 
 
 
-DistillGPT2,pass,2
+DistillGPT2,pass,3
 
 
 
diff --git a/benchmarks/dynamo/ci_expected_accuracy/cpu_inductor_torchbench_inference.csv b/benchmarks/dynamo/ci_expected_accuracy/cpu_inductor_torchbench_inference.csv
index 0157149df8cb8..0a51409c04ac3 100644
--- a/benchmarks/dynamo/ci_expected_accuracy/cpu_inductor_torchbench_inference.csv
+++ b/benchmarks/dynamo/ci_expected_accuracy/cpu_inductor_torchbench_inference.csv
@@ -186,10 +186,6 @@ pyhpc_turbulent_kinetic_energy,pass,0
 
 
 
-pytorch_CycleGAN_and_pix2pix,pass,0
-
-
-
 pytorch_stargan,pass,0
 
 
diff --git a/benchmarks/dynamo/ci_expected_accuracy/dynamic_aot_eager_huggingface_inference.csv b/benchmarks/dynamo/ci_expected_accuracy/dynamic_aot_eager_huggingface_inference.csv
index d3f0fbba71826..87dd88078f222 100644
--- a/benchmarks/dynamo/ci_expected_accuracy/dynamic_aot_eager_huggingface_inference.csv
+++ b/benchmarks/dynamo/ci_expected_accuracy/dynamic_aot_eager_huggingface_inference.csv
@@ -50,7 +50,7 @@ LayoutLMForMaskedLM,pass,0
 
 
 
-M2M100ForConditionalGeneration,pass,0
+M2M100ForConditionalGeneration,pass,7
 
 
 
diff --git a/benchmarks/dynamo/ci_expected_accuracy/dynamic_aot_eager_huggingface_training.csv b/benchmarks/dynamo/ci_expected_accuracy/dynamic_aot_eager_huggingface_training.csv
index 82d21ea3cb298..5ca03b5ecf9fb 100644
--- a/benchmarks/dynamo/ci_expected_accuracy/dynamic_aot_eager_huggingface_training.csv
+++ b/benchmarks/dynamo/ci_expected_accuracy/dynamic_aot_eager_huggingface_training.csv
@@ -50,7 +50,7 @@ LayoutLMForMaskedLM,pass,5
 
 
 
-M2M100ForConditionalGeneration,pass,4
+M2M100ForConditionalGeneration,pass,11
 
 
 
diff --git a/benchmarks/dynamo/ci_expected_accuracy/dynamic_aot_eager_torchbench_inference.csv b/benchmarks/dynamo/ci_expected_accuracy/dynamic_aot_eager_torchbench_inference.csv
index 16034d6bdfe72..8d4895e4a1ccd 100644
--- a/benchmarks/dynamo/ci_expected_accuracy/dynamic_aot_eager_torchbench_inference.csv
+++ b/benchmarks/dynamo/ci_expected_accuracy/dynamic_aot_eager_torchbench_inference.csv
@@ -202,10 +202,6 @@ pyhpc_turbulent_kinetic_energy,pass,0
 
 
 
-pytorch_CycleGAN_and_pix2pix,pass,0
-
-
-
 pytorch_stargan,pass,0
 
 
diff --git a/benchmarks/dynamo/ci_expected_accuracy/dynamic_aot_eager_torchbench_training.csv b/benchmarks/dynamo/ci_expected_accuracy/dynamic_aot_eager_torchbench_training.csv
index 54486aa44d557..6743f7c739ef1 100644
--- a/benchmarks/dynamo/ci_expected_accuracy/dynamic_aot_eager_torchbench_training.csv
+++ b/benchmarks/dynamo/ci_expected_accuracy/dynamic_aot_eager_torchbench_training.csv
@@ -130,10 +130,6 @@ phlippe_resnet,pass,6
 
 
 
-pytorch_CycleGAN_and_pix2pix,pass,6
-
-
-
 pytorch_stargan,pass,6
 
 
diff --git a/benchmarks/dynamo/ci_expected_accuracy/dynamic_cpu_aot_inductor_amp_freezing_torchbench_inference.csv b/benchmarks/dynamo/ci_expected_accuracy/dynamic_cpu_aot_inductor_amp_freezing_torchbench_inference.csv
index 7deb0fbba56b5..f7df046d04cce 100644
--- a/benchmarks/dynamo/ci_expected_accuracy/dynamic_cpu_aot_inductor_amp_freezing_torchbench_inference.csv
+++ b/benchmarks/dynamo/ci_expected_accuracy/dynamic_cpu_aot_inductor_amp_freezing_torchbench_inference.csv
@@ -146,10 +146,6 @@ pyhpc_turbulent_kinetic_energy,pass,0
 
 
 
-pytorch_CycleGAN_and_pix2pix,pass,0
-
-
-
 pytorch_stargan,pass,0
 
 
diff --git a/benchmarks/dynamo/ci_expected_accuracy/dynamic_cpu_aot_inductor_freezing_torchbench_inference.csv b/benchmarks/dynamo/ci_expected_accuracy/dynamic_cpu_aot_inductor_freezing_torchbench_inference.csv
index 7deb0fbba56b5..f7df046d04cce 100644
--- a/benchmarks/dynamo/ci_expected_accuracy/dynamic_cpu_aot_inductor_freezing_torchbench_inference.csv
+++ b/benchmarks/dynamo/ci_expected_accuracy/dynamic_cpu_aot_inductor_freezing_torchbench_inference.csv
@@ -146,10 +146,6 @@ pyhpc_turbulent_kinetic_energy,pass,0
 
 
 
-pytorch_CycleGAN_and_pix2pix,pass,0
-
-
-
 pytorch_stargan,pass,0
 
 
diff --git a/benchmarks/dynamo/ci_expected_accuracy/dynamic_cpu_inductor_huggingface_inference.csv b/benchmarks/dynamo/ci_expected_accuracy/dynamic_cpu_inductor_huggingface_inference.csv
index 439c9bf530468..2712e2d1a5012 100644
--- a/benchmarks/dynamo/ci_expected_accuracy/dynamic_cpu_inductor_huggingface_inference.csv
+++ b/benchmarks/dynamo/ci_expected_accuracy/dynamic_cpu_inductor_huggingface_inference.csv
@@ -30,7 +30,7 @@ DistilBertForMaskedLM,pass,0
 
 
 
-DistillGPT2,pass,2
+DistillGPT2,pass,3
 
 
 
diff --git a/benchmarks/dynamo/ci_expected_accuracy/dynamic_cpu_inductor_torchbench_inference.csv b/benchmarks/dynamo/ci_expected_accuracy/dynamic_cpu_inductor_torchbench_inference.csv
index 57ba90bd512bb..46d453de48c7d 100644
--- a/benchmarks/dynamo/ci_expected_accuracy/dynamic_cpu_inductor_torchbench_inference.csv
+++ b/benchmarks/dynamo/ci_expected_accuracy/dynamic_cpu_inductor_torchbench_inference.csv
@@ -170,10 +170,6 @@ pyhpc_turbulent_kinetic_energy,pass,0
 
 
 
-pytorch_CycleGAN_and_pix2pix,pass,0
-
-
-
 pytorch_stargan,pass,0
 
 
diff --git a/benchmarks/dynamo/ci_expected_accuracy/dynamic_cpu_max_autotune_inductor_amp_freezing_huggingface_inference.csv b/benchmarks/dynamo/ci_expected_accuracy/dynamic_cpu_max_autotune_inductor_amp_freezing_huggingface_inference.csv
index 439c9bf530468..e4aabce10466d 100644
--- a/benchmarks/dynamo/ci_expected_accuracy/dynamic_cpu_max_autotune_inductor_amp_freezing_huggingface_inference.csv
+++ b/benchmarks/dynamo/ci_expected_accuracy/dynamic_cpu_max_autotune_inductor_amp_freezing_huggingface_inference.csv
@@ -50,7 +50,7 @@ LayoutLMForMaskedLM,pass,0
 
 
 
-M2M100ForConditionalGeneration,pass,0
+M2M100ForConditionalGeneration,pass,7
 
 
 
diff --git a/benchmarks/dynamo/ci_expected_accuracy/dynamic_cpu_max_autotune_inductor_amp_freezing_torchbench_inference.csv b/benchmarks/dynamo/ci_expected_accuracy/dynamic_cpu_max_autotune_inductor_amp_freezing_torchbench_inference.csv
index 723ef7a272ea1..dfa70e34afc53 100644
--- a/benchmarks/dynamo/ci_expected_accuracy/dynamic_cpu_max_autotune_inductor_amp_freezing_torchbench_inference.csv
+++ b/benchmarks/dynamo/ci_expected_accuracy/dynamic_cpu_max_autotune_inductor_amp_freezing_torchbench_inference.csv
@@ -190,10 +190,6 @@ pyhpc_turbulent_kinetic_energy,pass,0
 
 
 
-pytorch_CycleGAN_and_pix2pix,pass,0
-
-
-
 pytorch_stargan,pass,0
 
 
diff --git a/benchmarks/dynamo/ci_expected_accuracy/dynamic_inductor_huggingface_inference.csv b/benchmarks/dynamo/ci_expected_accuracy/dynamic_inductor_huggingface_inference.csv
index d3f0fbba71826..87dd88078f222 100644
--- a/benchmarks/dynamo/ci_expected_accuracy/dynamic_inductor_huggingface_inference.csv
+++ b/benchmarks/dynamo/ci_expected_accuracy/dynamic_inductor_huggingface_inference.csv
@@ -50,7 +50,7 @@ LayoutLMForMaskedLM,pass,0
 
 
 
-M2M100ForConditionalGeneration,pass,0
+M2M100ForConditionalGeneration,pass,7
 
 
 
diff --git a/benchmarks/dynamo/ci_expected_accuracy/dynamic_inductor_huggingface_training.csv b/benchmarks/dynamo/ci_expected_accuracy/dynamic_inductor_huggingface_training.csv
index 82d21ea3cb298..5ca03b5ecf9fb 100644
--- a/benchmarks/dynamo/ci_expected_accuracy/dynamic_inductor_huggingface_training.csv
+++ b/benchmarks/dynamo/ci_expected_accuracy/dynamic_inductor_huggingface_training.csv
@@ -50,7 +50,7 @@ LayoutLMForMaskedLM,pass,5
 
 
 
-M2M100ForConditionalGeneration,pass,4
+M2M100ForConditionalGeneration,pass,11
 
 
 
diff --git a/benchmarks/dynamo/ci_expected_accuracy/dynamic_inductor_torchbench_inference.csv b/benchmarks/dynamo/ci_expected_accuracy/dynamic_inductor_torchbench_inference.csv
index 1ea87d2648875..3ef5716e9306a 100644
--- a/benchmarks/dynamo/ci_expected_accuracy/dynamic_inductor_torchbench_inference.csv
+++ b/benchmarks/dynamo/ci_expected_accuracy/dynamic_inductor_torchbench_inference.csv
@@ -74,7 +74,7 @@ detectron2_fasterrcnn_r_50_fpn,eager_fail_to_run,0
 
 
 
-detectron2_fcos_r_50_fpn,pass,22
+detectron2_fcos_r_50_fpn,pass,23
 
 
 
@@ -202,10 +202,6 @@ pyhpc_turbulent_kinetic_energy,pass,0
 
 
 
-pytorch_CycleGAN_and_pix2pix,pass,0
-
-
-
 pytorch_stargan,pass,0
 
 
diff --git a/benchmarks/dynamo/ci_expected_accuracy/dynamic_inductor_torchbench_training.csv b/benchmarks/dynamo/ci_expected_accuracy/dynamic_inductor_torchbench_training.csv
index 24ad3a397e2cc..9dd6d40930389 100644
--- a/benchmarks/dynamo/ci_expected_accuracy/dynamic_inductor_torchbench_training.csv
+++ b/benchmarks/dynamo/ci_expected_accuracy/dynamic_inductor_torchbench_training.csv
@@ -130,10 +130,6 @@ phlippe_resnet,pass,6
 
 
 
-pytorch_CycleGAN_and_pix2pix,pass,6
-
-
-
 pytorch_stargan,pass,6
 
 
@@ -194,7 +190,7 @@ vgg16,pass,6
 
 
 
-vision_maskrcnn,fail_to_run,37
+vision_maskrcnn,fail_to_run,38
 
 
 
diff --git a/benchmarks/dynamo/ci_expected_accuracy/dynamo_eager_huggingface_inference.csv b/benchmarks/dynamo/ci_expected_accuracy/dynamo_eager_huggingface_inference.csv
index d3f0fbba71826..87dd88078f222 100644
--- a/benchmarks/dynamo/ci_expected_accuracy/dynamo_eager_huggingface_inference.csv
+++ b/benchmarks/dynamo/ci_expected_accuracy/dynamo_eager_huggingface_inference.csv
@@ -50,7 +50,7 @@ LayoutLMForMaskedLM,pass,0
 
 
 
-M2M100ForConditionalGeneration,pass,0
+M2M100ForConditionalGeneration,pass,7
 
 
 
diff --git a/benchmarks/dynamo/ci_expected_accuracy/dynamo_eager_huggingface_training.csv b/benchmarks/dynamo/ci_expected_accuracy/dynamo_eager_huggingface_training.csv
index 82d21ea3cb298..5ca03b5ecf9fb 100644
--- a/benchmarks/dynamo/ci_expected_accuracy/dynamo_eager_huggingface_training.csv
+++ b/benchmarks/dynamo/ci_expected_accuracy/dynamo_eager_huggingface_training.csv
@@ -50,7 +50,7 @@ LayoutLMForMaskedLM,pass,5
 
 
 
-M2M100ForConditionalGeneration,pass,4
+M2M100ForConditionalGeneration,pass,11
 
 
 
diff --git a/benchmarks/dynamo/ci_expected_accuracy/dynamo_eager_torchbench_inference.csv b/benchmarks/dynamo/ci_expected_accuracy/dynamo_eager_torchbench_inference.csv
index bdf3313659b66..b0daa9a947ad4 100644
--- a/benchmarks/dynamo/ci_expected_accuracy/dynamo_eager_torchbench_inference.csv
+++ b/benchmarks/dynamo/ci_expected_accuracy/dynamo_eager_torchbench_inference.csv
@@ -202,10 +202,6 @@ pyhpc_turbulent_kinetic_energy,pass,0
 
 
 
-pytorch_CycleGAN_and_pix2pix,pass,0
-
-
-
 pytorch_stargan,pass,0
 
 
diff --git a/benchmarks/dynamo/ci_expected_accuracy/dynamo_eager_torchbench_training.csv b/benchmarks/dynamo/ci_expected_accuracy/dynamo_eager_torchbench_training.csv
index ee7838505d67c..9ff11382f67ad 100644
--- a/benchmarks/dynamo/ci_expected_accuracy/dynamo_eager_torchbench_training.csv
+++ b/benchmarks/dynamo/ci_expected_accuracy/dynamo_eager_torchbench_training.csv
@@ -130,10 +130,6 @@ phlippe_resnet,pass,6
 
 
 
-pytorch_CycleGAN_and_pix2pix,pass,6
-
-
-
 pytorch_stargan,pass,6
 
 
diff --git a/benchmarks/dynamo/ci_expected_accuracy/inductor_huggingface_inference.csv b/benchmarks/dynamo/ci_expected_accuracy/inductor_huggingface_inference.csv
index d3f0fbba71826..1333ff866469e 100644
--- a/benchmarks/dynamo/ci_expected_accuracy/inductor_huggingface_inference.csv
+++ b/benchmarks/dynamo/ci_expected_accuracy/inductor_huggingface_inference.csv
@@ -30,7 +30,7 @@ DistilBertForMaskedLM,pass,0
 
 
 
-DistillGPT2,pass,2
+DistillGPT2,pass,3
 
 
 
@@ -122,4 +122,4 @@ meta-llama/Llama-3.2-1B,pass,0
 
 
 
-openai/whisper-tiny,pass,5
+openai/whisper-tiny,pass,7
diff --git a/benchmarks/dynamo/ci_expected_accuracy/inductor_huggingface_training.csv b/benchmarks/dynamo/ci_expected_accuracy/inductor_huggingface_training.csv
index 82d21ea3cb298..452bf8707fbfb 100644
--- a/benchmarks/dynamo/ci_expected_accuracy/inductor_huggingface_training.csv
+++ b/benchmarks/dynamo/ci_expected_accuracy/inductor_huggingface_training.csv
@@ -10,7 +10,7 @@ AllenaiLongformerBase,pass,9
 
 
 
-BartForCausalLM,pass,6
+BartForCausalLM,pass,7
 
 
 
@@ -30,7 +30,7 @@ DistilBertForMaskedLM,pass,5
 
 
 
-DistillGPT2,pass,7
+DistillGPT2,pass,8
 
 
 
@@ -54,7 +54,7 @@ M2M100ForConditionalGeneration,pass,4
 
 
 
-MBartForCausalLM,pass,6
+MBartForCausalLM,pass,7
 
 
 
@@ -70,15 +70,15 @@ MobileBertForMaskedLM,pass,3
 
 
 
-OPTForCausalLM,pass,8
+OPTForCausalLM,pass,7
 
 
 
-PLBartForCausalLM,pass,6
+PLBartForCausalLM,pass,7
 
 
 
-PegasusForCausalLM,pass,6
+PegasusForCausalLM,pass,7
 
 
 
diff --git a/benchmarks/dynamo/ci_expected_accuracy/inductor_torchbench_inference.csv b/benchmarks/dynamo/ci_expected_accuracy/inductor_torchbench_inference.csv
index d5395066bbb42..d04a273049dbf 100644
--- a/benchmarks/dynamo/ci_expected_accuracy/inductor_torchbench_inference.csv
+++ b/benchmarks/dynamo/ci_expected_accuracy/inductor_torchbench_inference.csv
@@ -74,7 +74,7 @@ detectron2_fasterrcnn_r_50_fpn,eager_fail_to_run,0
 
 
 
-detectron2_fcos_r_50_fpn,pass,22
+detectron2_fcos_r_50_fpn,pass,23
 
 
 
@@ -202,10 +202,6 @@ pyhpc_turbulent_kinetic_energy,pass,0
 
 
 
-pytorch_CycleGAN_and_pix2pix,pass,0
-
-
-
 pytorch_stargan,pass,0
 
 
diff --git a/benchmarks/dynamo/ci_expected_accuracy/inductor_torchbench_training.csv b/benchmarks/dynamo/ci_expected_accuracy/inductor_torchbench_training.csv
index 24ad3a397e2cc..90c1f0972a3b2 100644
--- a/benchmarks/dynamo/ci_expected_accuracy/inductor_torchbench_training.csv
+++ b/benchmarks/dynamo/ci_expected_accuracy/inductor_torchbench_training.csv
@@ -130,10 +130,6 @@ phlippe_resnet,pass,6
 
 
 
-pytorch_CycleGAN_and_pix2pix,pass,6
-
-
-
 pytorch_stargan,pass,6
 
 
@@ -174,6 +170,10 @@ soft_actor_critic,pass,6
 
 
 
+speech_transformer,pass,16
+
+
+
 squeezenet1_1,pass,6
 
 
@@ -194,7 +194,7 @@ vgg16,pass,6
 
 
 
-vision_maskrcnn,fail_to_run,37
+vision_maskrcnn,fail_to_run,38
 
 
 
diff --git a/benchmarks/dynamo/ci_expected_accuracy/rocm/aot_eager_huggingface_inference.csv b/benchmarks/dynamo/ci_expected_accuracy/rocm/aot_eager_huggingface_inference.csv
index d3f0fbba71826..87dd88078f222 100644
--- a/benchmarks/dynamo/ci_expected_accuracy/rocm/aot_eager_huggingface_inference.csv
+++ b/benchmarks/dynamo/ci_expected_accuracy/rocm/aot_eager_huggingface_inference.csv
@@ -50,7 +50,7 @@ LayoutLMForMaskedLM,pass,0
 
 
 
-M2M100ForConditionalGeneration,pass,0
+M2M100ForConditionalGeneration,pass,7
 
 
 
diff --git a/benchmarks/dynamo/ci_expected_accuracy/rocm/aot_eager_huggingface_training.csv b/benchmarks/dynamo/ci_expected_accuracy/rocm/aot_eager_huggingface_training.csv
index deea60d3203e2..af120d3d31b3b 100644
--- a/benchmarks/dynamo/ci_expected_accuracy/rocm/aot_eager_huggingface_training.csv
+++ b/benchmarks/dynamo/ci_expected_accuracy/rocm/aot_eager_huggingface_training.csv
@@ -30,7 +30,7 @@ DistilBertForMaskedLM,pass,5
 
 
 
-DistillGPT2,pass,7
+DistillGPT2,fail_accuracy,7
 
 
 
@@ -50,7 +50,7 @@ LayoutLMForMaskedLM,pass,5
 
 
 
-M2M100ForConditionalGeneration,pass,4
+M2M100ForConditionalGeneration,pass,11
 
 
 
diff --git a/benchmarks/dynamo/ci_expected_accuracy/rocm/aot_eager_torchbench_inference.csv b/benchmarks/dynamo/ci_expected_accuracy/rocm/aot_eager_torchbench_inference.csv
index 9e1d982030b13..5f486d36c45b0 100644
--- a/benchmarks/dynamo/ci_expected_accuracy/rocm/aot_eager_torchbench_inference.csv
+++ b/benchmarks/dynamo/ci_expected_accuracy/rocm/aot_eager_torchbench_inference.csv
@@ -202,7 +202,6 @@ pyhpc_turbulent_kinetic_energy,pass,0
 
 
 
-pytorch_CycleGAN_and_pix2pix,pass,0
 
 
 
diff --git a/benchmarks/dynamo/ci_expected_accuracy/rocm/aot_eager_torchbench_training.csv b/benchmarks/dynamo/ci_expected_accuracy/rocm/aot_eager_torchbench_training.csv
index fa52782561be2..d171c2e06bffa 100644
--- a/benchmarks/dynamo/ci_expected_accuracy/rocm/aot_eager_torchbench_training.csv
+++ b/benchmarks/dynamo/ci_expected_accuracy/rocm/aot_eager_torchbench_training.csv
@@ -130,7 +130,6 @@ phlippe_resnet,pass,6
 
 
 
-pytorch_CycleGAN_and_pix2pix,pass,6
 
 
 
diff --git a/benchmarks/dynamo/ci_expected_accuracy/rocm/aot_inductor_torchbench_inference.csv b/benchmarks/dynamo/ci_expected_accuracy/rocm/aot_inductor_torchbench_inference.csv
index 6d4c4c0359a50..fb70c29d79fda 100644
--- a/benchmarks/dynamo/ci_expected_accuracy/rocm/aot_inductor_torchbench_inference.csv
+++ b/benchmarks/dynamo/ci_expected_accuracy/rocm/aot_inductor_torchbench_inference.csv
@@ -178,7 +178,6 @@ pyhpc_turbulent_kinetic_energy,pass,0
 
 
 
-pytorch_CycleGAN_and_pix2pix,pass,0
 
 
 
diff --git a/benchmarks/dynamo/ci_expected_accuracy/rocm/dynamic_aot_eager_huggingface_inference.csv b/benchmarks/dynamo/ci_expected_accuracy/rocm/dynamic_aot_eager_huggingface_inference.csv
index d3f0fbba71826..87dd88078f222 100644
--- a/benchmarks/dynamo/ci_expected_accuracy/rocm/dynamic_aot_eager_huggingface_inference.csv
+++ b/benchmarks/dynamo/ci_expected_accuracy/rocm/dynamic_aot_eager_huggingface_inference.csv
@@ -50,7 +50,7 @@ LayoutLMForMaskedLM,pass,0
 
 
 
-M2M100ForConditionalGeneration,pass,0
+M2M100ForConditionalGeneration,pass,7
 
 
 
diff --git a/benchmarks/dynamo/ci_expected_accuracy/rocm/dynamic_aot_eager_huggingface_training.csv b/benchmarks/dynamo/ci_expected_accuracy/rocm/dynamic_aot_eager_huggingface_training.csv
index deea60d3203e2..af120d3d31b3b 100644
--- a/benchmarks/dynamo/ci_expected_accuracy/rocm/dynamic_aot_eager_huggingface_training.csv
+++ b/benchmarks/dynamo/ci_expected_accuracy/rocm/dynamic_aot_eager_huggingface_training.csv
@@ -30,7 +30,7 @@ DistilBertForMaskedLM,pass,5
 
 
 
-DistillGPT2,pass,7
+DistillGPT2,fail_accuracy,7
 
 
 
@@ -50,7 +50,7 @@ LayoutLMForMaskedLM,pass,5
 
 
 
-M2M100ForConditionalGeneration,pass,4
+M2M100ForConditionalGeneration,pass,11
 
 
 
diff --git a/benchmarks/dynamo/ci_expected_accuracy/rocm/dynamic_aot_eager_torchbench_inference.csv b/benchmarks/dynamo/ci_expected_accuracy/rocm/dynamic_aot_eager_torchbench_inference.csv
index 61ffb7197f676..2dd7684e5bd4f 100644
--- a/benchmarks/dynamo/ci_expected_accuracy/rocm/dynamic_aot_eager_torchbench_inference.csv
+++ b/benchmarks/dynamo/ci_expected_accuracy/rocm/dynamic_aot_eager_torchbench_inference.csv
@@ -202,7 +202,6 @@ pyhpc_turbulent_kinetic_energy,pass,0
 
 
 
-pytorch_CycleGAN_and_pix2pix,pass,0
 
 
 
diff --git a/benchmarks/dynamo/ci_expected_accuracy/rocm/dynamic_aot_eager_torchbench_training.csv b/benchmarks/dynamo/ci_expected_accuracy/rocm/dynamic_aot_eager_torchbench_training.csv
index c9d8189dc8b74..4bb772b41d8f0 100644
--- a/benchmarks/dynamo/ci_expected_accuracy/rocm/dynamic_aot_eager_torchbench_training.csv
+++ b/benchmarks/dynamo/ci_expected_accuracy/rocm/dynamic_aot_eager_torchbench_training.csv
@@ -130,7 +130,6 @@ phlippe_resnet,pass,6
 
 
 
-pytorch_CycleGAN_and_pix2pix,pass,6
 
 
 
diff --git a/benchmarks/dynamo/ci_expected_accuracy/rocm/dynamic_inductor_huggingface_inference.csv b/benchmarks/dynamo/ci_expected_accuracy/rocm/dynamic_inductor_huggingface_inference.csv
index d3f0fbba71826..87dd88078f222 100644
--- a/benchmarks/dynamo/ci_expected_accuracy/rocm/dynamic_inductor_huggingface_inference.csv
+++ b/benchmarks/dynamo/ci_expected_accuracy/rocm/dynamic_inductor_huggingface_inference.csv
@@ -50,7 +50,7 @@ LayoutLMForMaskedLM,pass,0
 
 
 
-M2M100ForConditionalGeneration,pass,0
+M2M100ForConditionalGeneration,pass,7
 
 
 
diff --git a/benchmarks/dynamo/ci_expected_accuracy/rocm/dynamic_inductor_huggingface_training.csv b/benchmarks/dynamo/ci_expected_accuracy/rocm/dynamic_inductor_huggingface_training.csv
index deea60d3203e2..af120d3d31b3b 100644
--- a/benchmarks/dynamo/ci_expected_accuracy/rocm/dynamic_inductor_huggingface_training.csv
+++ b/benchmarks/dynamo/ci_expected_accuracy/rocm/dynamic_inductor_huggingface_training.csv
@@ -30,7 +30,7 @@ DistilBertForMaskedLM,pass,5
 
 
 
-DistillGPT2,pass,7
+DistillGPT2,fail_accuracy,7
 
 
 
@@ -50,7 +50,7 @@ LayoutLMForMaskedLM,pass,5
 
 
 
-M2M100ForConditionalGeneration,pass,4
+M2M100ForConditionalGeneration,pass,11
 
 
 
diff --git a/benchmarks/dynamo/ci_expected_accuracy/rocm/dynamic_inductor_timm_training.csv b/benchmarks/dynamo/ci_expected_accuracy/rocm/dynamic_inductor_timm_training.csv
index b544ebd80d0a1..a07866a587a1e 100644
--- a/benchmarks/dynamo/ci_expected_accuracy/rocm/dynamic_inductor_timm_training.csv
+++ b/benchmarks/dynamo/ci_expected_accuracy/rocm/dynamic_inductor_timm_training.csv
@@ -70,4 +70,4 @@ vit_base_patch14_dinov2.lvd142m,pass,7
 
 
 
-vit_base_patch16_siglip_256,fail_accuracy,7
+vit_base_patch16_siglip_256,pass,7
diff --git a/benchmarks/dynamo/ci_expected_accuracy/rocm/dynamic_inductor_torchbench_inference.csv b/benchmarks/dynamo/ci_expected_accuracy/rocm/dynamic_inductor_torchbench_inference.csv
index a678afd25c32d..0c22866fffccb 100644
--- a/benchmarks/dynamo/ci_expected_accuracy/rocm/dynamic_inductor_torchbench_inference.csv
+++ b/benchmarks/dynamo/ci_expected_accuracy/rocm/dynamic_inductor_torchbench_inference.csv
@@ -74,7 +74,7 @@ detectron2_fasterrcnn_r_50_fpn,eager_fail_to_run,0
 
 
 
-detectron2_fcos_r_50_fpn,pass,22
+detectron2_fcos_r_50_fpn,pass,23
 
 
 
@@ -202,7 +202,6 @@ pyhpc_turbulent_kinetic_energy,pass,0
 
 
 
-pytorch_CycleGAN_and_pix2pix,pass,0
 
 
 
diff --git a/benchmarks/dynamo/ci_expected_accuracy/rocm/dynamic_inductor_torchbench_training.csv b/benchmarks/dynamo/ci_expected_accuracy/rocm/dynamic_inductor_torchbench_training.csv
index e059c66ca3d04..a4d42c7d7723a 100644
--- a/benchmarks/dynamo/ci_expected_accuracy/rocm/dynamic_inductor_torchbench_training.csv
+++ b/benchmarks/dynamo/ci_expected_accuracy/rocm/dynamic_inductor_torchbench_training.csv
@@ -90,7 +90,7 @@ microbench_unbacked_tolist_sum,pass,9
 
 
 
-mnasnet1_0,pass,0
+mnasnet1_0,pass,7
 
 
 
@@ -102,7 +102,7 @@ mobilenet_v2_quantized_qat,eager_fail_to_run,0
 
 
 
-mobilenet_v3_large,pass,0
+mobilenet_v3_large,pass,7
 
 
 
@@ -130,7 +130,6 @@ phlippe_resnet,pass,6
 
 
 
-pytorch_CycleGAN_and_pix2pix,pass,6
 
 
 
@@ -142,7 +141,7 @@ pytorch_unet,pass_due_to_skip,7
 
 
 
-resnet152,eager_two_runs_differ,0
+resnet152,eager_two_runs_differ,7
 
 
 
@@ -194,7 +193,7 @@ vgg16,pass,6
 
 
 
-vision_maskrcnn,fail_to_run,37
+vision_maskrcnn,fail_to_run,38
 
 
 
diff --git a/benchmarks/dynamo/ci_expected_accuracy/rocm/dynamo_eager_huggingface_inference.csv b/benchmarks/dynamo/ci_expected_accuracy/rocm/dynamo_eager_huggingface_inference.csv
index d3f0fbba71826..87dd88078f222 100644
--- a/benchmarks/dynamo/ci_expected_accuracy/rocm/dynamo_eager_huggingface_inference.csv
+++ b/benchmarks/dynamo/ci_expected_accuracy/rocm/dynamo_eager_huggingface_inference.csv
@@ -50,7 +50,7 @@ LayoutLMForMaskedLM,pass,0
 
 
 
-M2M100ForConditionalGeneration,pass,0
+M2M100ForConditionalGeneration,pass,7
 
 
 
diff --git a/benchmarks/dynamo/ci_expected_accuracy/rocm/dynamo_eager_huggingface_training.csv b/benchmarks/dynamo/ci_expected_accuracy/rocm/dynamo_eager_huggingface_training.csv
index deea60d3203e2..af120d3d31b3b 100644
--- a/benchmarks/dynamo/ci_expected_accuracy/rocm/dynamo_eager_huggingface_training.csv
+++ b/benchmarks/dynamo/ci_expected_accuracy/rocm/dynamo_eager_huggingface_training.csv
@@ -30,7 +30,7 @@ DistilBertForMaskedLM,pass,5
 
 
 
-DistillGPT2,pass,7
+DistillGPT2,fail_accuracy,7
 
 
 
@@ -50,7 +50,7 @@ LayoutLMForMaskedLM,pass,5
 
 
 
-M2M100ForConditionalGeneration,pass,4
+M2M100ForConditionalGeneration,pass,11
 
 
 
diff --git a/benchmarks/dynamo/ci_expected_accuracy/rocm/dynamo_eager_torchbench_inference.csv b/benchmarks/dynamo/ci_expected_accuracy/rocm/dynamo_eager_torchbench_inference.csv
index 8dbc90cc8d2c7..a18f3b215ecc8 100644
--- a/benchmarks/dynamo/ci_expected_accuracy/rocm/dynamo_eager_torchbench_inference.csv
+++ b/benchmarks/dynamo/ci_expected_accuracy/rocm/dynamo_eager_torchbench_inference.csv
@@ -202,7 +202,6 @@ pyhpc_turbulent_kinetic_energy,pass,0
 
 
 
-pytorch_CycleGAN_and_pix2pix,pass,0
 
 
 
diff --git a/benchmarks/dynamo/ci_expected_accuracy/rocm/dynamo_eager_torchbench_training.csv b/benchmarks/dynamo/ci_expected_accuracy/rocm/dynamo_eager_torchbench_training.csv
index 12d994100dfe3..d7d0a9b0b3292 100644
--- a/benchmarks/dynamo/ci_expected_accuracy/rocm/dynamo_eager_torchbench_training.csv
+++ b/benchmarks/dynamo/ci_expected_accuracy/rocm/dynamo_eager_torchbench_training.csv
@@ -130,7 +130,6 @@ phlippe_resnet,pass,6
 
 
 
-pytorch_CycleGAN_and_pix2pix,pass,6
 
 
 
diff --git a/benchmarks/dynamo/ci_expected_accuracy/rocm/inductor_huggingface_inference.csv b/benchmarks/dynamo/ci_expected_accuracy/rocm/inductor_huggingface_inference.csv
index 46f1e5adf4ec9..6f65795e3f04e 100644
--- a/benchmarks/dynamo/ci_expected_accuracy/rocm/inductor_huggingface_inference.csv
+++ b/benchmarks/dynamo/ci_expected_accuracy/rocm/inductor_huggingface_inference.csv
@@ -50,7 +50,7 @@ LayoutLMForMaskedLM,pass,0
 
 
 
-M2M100ForConditionalGeneration,pass,0
+M2M100ForConditionalGeneration,pass,7
 
 
 
diff --git a/benchmarks/dynamo/ci_expected_accuracy/rocm/inductor_huggingface_training.csv b/benchmarks/dynamo/ci_expected_accuracy/rocm/inductor_huggingface_training.csv
index e06f3bde8af13..07ec2bb634b39 100644
--- a/benchmarks/dynamo/ci_expected_accuracy/rocm/inductor_huggingface_training.csv
+++ b/benchmarks/dynamo/ci_expected_accuracy/rocm/inductor_huggingface_training.csv
@@ -50,7 +50,7 @@ LayoutLMForMaskedLM,pass,5
 
 
 
-M2M100ForConditionalGeneration,pass,4
+M2M100ForConditionalGeneration,pass,11
 
 
 
diff --git a/benchmarks/dynamo/ci_expected_accuracy/rocm/inductor_torchbench_inference.csv b/benchmarks/dynamo/ci_expected_accuracy/rocm/inductor_torchbench_inference.csv
index ef7c4f2a01e9f..55457b2f60695 100644
--- a/benchmarks/dynamo/ci_expected_accuracy/rocm/inductor_torchbench_inference.csv
+++ b/benchmarks/dynamo/ci_expected_accuracy/rocm/inductor_torchbench_inference.csv
@@ -210,7 +210,6 @@ pyhpc_turbulent_kinetic_energy,pass,0
 
 
 
-pytorch_CycleGAN_and_pix2pix,pass,0
 
 
 
diff --git a/benchmarks/dynamo/ci_expected_accuracy/rocm/inductor_torchbench_training.csv b/benchmarks/dynamo/ci_expected_accuracy/rocm/inductor_torchbench_training.csv
index 91e6df19ff02a..529710bbc21c1 100644
--- a/benchmarks/dynamo/ci_expected_accuracy/rocm/inductor_torchbench_training.csv
+++ b/benchmarks/dynamo/ci_expected_accuracy/rocm/inductor_torchbench_training.csv
@@ -134,7 +134,6 @@ phlippe_resnet,pass,6
 
 
 
-pytorch_CycleGAN_and_pix2pix,pass,6
 
 
 
diff --git a/benchmarks/dynamo/common.py b/benchmarks/dynamo/common.py
index ee8a95f5c8459..eb3b44dfdc918 100644
--- a/benchmarks/dynamo/common.py
+++ b/benchmarks/dynamo/common.py
@@ -25,7 +25,7 @@
 import time
 import weakref
 from contextlib import contextmanager
-from typing import Any, NamedTuple, Optional, overload, TYPE_CHECKING, TypeVar
+from typing import Any, NamedTuple, overload, TYPE_CHECKING, TypeVar
 from unittest.mock import MagicMock
 
 import numpy as np
@@ -475,7 +475,6 @@ def output_signpost(data, args, suite, error=None):
         "log_conv_args",
         "recompile_profiler",
         "find_batch_sizes",
-        "download_only",
         # Redundant
         "batch_size",
         "batch_size_file",
@@ -967,7 +966,9 @@ def latency_experiment_summary(suite_name, args, model, timings, **kwargs):
         first_fields.append(kwargs["tag"])
     headers = first_headers + ["speedup", "abs_latency"]
     row = first_fields + [float(speedup), median[1] * 1000]
-    msg = f"{median[0] * 1000} ms, {median[1] * 1000} ms, {speedup:.3f}x"
+    msg = f"{speedup:.3f}x"
+    if getattr(args, "_print_latency_ms", False):
+        msg = f"{median[0] * 1000:.4f} ms, {median[1] * 1000:.4f} ms, {msg}"
     if args.baseline:
         headers.extend(
             [
@@ -1154,6 +1155,8 @@ def maybe_mark_profile(*args, **kwargs):
     headers = first_headers + ["speedup", "abs_latency"]
     row = first_fields + [float(speedup), median[1] * 1000]
     msg = f"{speedup:.3f}x"
+    if getattr(args, "_print_latency_ms", False):
+        msg = f"{median[0] * 1000:.4f} ms, {median[1] * 1000:.4f} ms, {msg}"
     if args.baseline:
         headers.extend(
             [
@@ -1785,7 +1788,7 @@ def __init__(self):
         self.grad_scaler = DummyGradScaler()
         self.autocast = contextlib.nullcontext
         self.autocast_arg = {}
-        self.optimizer: Optional[torch.optim.Optimizer] = None
+        self.optimizer: torch.optim.Optimizer | None = None
         self._args = None
 
     def setup_amp(self, current_device=None):
@@ -1950,10 +1953,6 @@ def disable_cudagraph_models(self):
     def guard_on_nn_module_models(self):
         return set()
 
-    @property
-    def inline_inbuilt_nn_modules_models(self):
-        return set()
-
     def get_tolerance_and_cosine_flag(self, is_training, current_device, name):
         raise NotImplementedError
 
@@ -2025,17 +2024,18 @@ def cast_based_on_args(self, model, example_inputs):
 
         return model, example_inputs
 
-    def validate_model(self, model, example_inputs):
+    def validate_model(self, name, model, example_inputs):
         """
         Runs the eager model with example inputs to ensure that eager passes.
         """
         model = self.deepcopy_model(model)
         example_inputs = clone_inputs(example_inputs)
         model, example_inputs = self.cast_based_on_args(model, example_inputs)
-        try:
-            self.model_iter_fn(model, example_inputs)
-        except Exception as e:
-            raise RuntimeError("Eager run failed") from e
+        with self.pick_grad(name, self.args.training):
+            try:
+                self.model_iter_fn(model, example_inputs)
+            except Exception as e:
+                raise RuntimeError("Eager run failed") from e
 
     def maybe_cast(self, model, example_inputs):
         model, example_inputs = self.cast_based_on_args(model, example_inputs)
@@ -3205,6 +3205,11 @@ def parse_args(args=None):
     parser.add_argument(
         "--freezing", action="store_true", help="turn on freezing", default=False
     )
+    parser.add_argument(
+        "--deterministic",
+        action="store_true",
+        help="Enable deterministic mode (torch.use_deterministic_algorithms, cudnn.deterministic, etc.)",
+    )
     parser.add_argument(
         "--inductor-config",
         "-c",
@@ -3296,6 +3301,16 @@ def get_example_inputs(self):
         action="store_true",
         help="Mark batch dimension as unbacked using mark_unbacked. Implies --dynamic-shapes",
     )
+    parser.add_argument(
+        "--compare-backed-unbacked",
+        action="store_true",
+        help="Run both dynamic-batch-only (backed) and unbacked-batch-only, then compare results side by side",
+    )
+    parser.add_argument(
+        "--_print-latency-ms",
+        action="store_true",
+        help=argparse.SUPPRESS,
+    )
     parser.add_argument(
         "--specialize-int", action="store_true", help="Run with specialize_int=True."
     )
@@ -3627,7 +3642,6 @@ def get_example_inputs(self):
             "int8dynamic",
             "int8weightonly",
             "int4weightonly",
-            "autoquant",
             "noquant",
         ],
         default=None,
@@ -3683,13 +3697,8 @@ def get_example_inputs(self):
         action="store_true",
         help="finds the largest batch size that could fit on GPUs",
     )
-    group.add_argument(
-        "--download-only",
-        action="store_true",
-        help="Download all models and exit without running benchmarks.",
-    )
 
-    mode_group = parser.add_mutually_exclusive_group(required=False)
+    mode_group = parser.add_mutually_exclusive_group(required=True)
     mode_group.add_argument(
         "--accuracy",
         action="store_true",
@@ -3703,7 +3712,7 @@ def get_example_inputs(self):
         action="store_true",
         help="extracts the tolerance for each model with small batch size and eval mode",
     )
-    run_mode_group = parser.add_mutually_exclusive_group(required=False)
+    run_mode_group = parser.add_mutually_exclusive_group(required=True)
     run_mode_group.add_argument(
         "--training",
         action="store_true",
@@ -3712,13 +3721,7 @@ def get_example_inputs(self):
     run_mode_group.add_argument(
         "--inference", action="store_true", help="Performs inference"
     )
-    args = parser.parse_args(args)
-    if not args.download_only:
-        if not any([args.accuracy, args.performance, args.tolerance]):
-            parser.error("one of --accuracy/--performance/--tolerance is required")
-        if not any([args.training, args.inference]):
-            parser.error("one of --training/--inference is required")
-    return args
+    return parser.parse_args(args)
 
 
 def process_caching_precompile():
@@ -3820,7 +3823,159 @@ def main(runner, original_dir=None, args=None):
         else:
             # single process path just uses the main process
             args.world_size = 1
-            process_entry(0, runner, original_dir, args)
+            if args.compare_backed_unbacked:
+                _run_compare_backed_unbacked(runner, args)
+            else:
+                process_entry(0, runner, original_dir, args)
+
+
+def _run_compare_backed_unbacked(runner, args):
+    """Run backed and unbacked per-model, alternating, and compare speedup."""
+    import re
+    import subprocess
+
+    def print_comparison(all_results):
+        print(f"\n{'=' * 80}", flush=True)
+        print("COMPARISON", flush=True)
+        print(f"{'=' * 80}", flush=True)
+        print(
+            f"  {'model':<40s} {'backed_ms':>10s} {'unbacked_ms':>11s} {'diff':>8s}",
+            flush=True,
+        )
+        print(f"  {'-' * 40} {'-' * 10} {'-' * 11} {'-' * 8}", flush=True)
+        for name, modes in all_results.items():
+            b_ms = modes.get("backed_ms")
+            u_ms = modes.get("unbacked_ms")
+            if b_ms is not None and u_ms is not None:
+                ms_diff_pct = (u_ms - b_ms) / b_ms * 100
+                print(
+                    f"  {name:<40s} {b_ms:>10.3f} {u_ms:>11.3f} {ms_diff_pct:>+7.1f}%",
+                    flush=True,
+                )
+            elif b_ms is not None:
+                print(
+                    f"  {name:<40s} {b_ms:>10.3f} {'N/A':>11s} {'N/A':>8s}", flush=True
+                )
+            elif u_ms is not None:
+                print(
+                    f"  {name:<40s} {'N/A':>10s} {u_ms:>11.3f} {'N/A':>8s}", flush=True
+                )
+            else:
+                backed = (
+                    "FAILED" if "backed" not in modes else f"{modes['backed']:.3f}x"
+                )
+                unbacked = (
+                    "FAILED" if "unbacked" not in modes else f"{modes['unbacked']:.3f}x"
+                )
+                print(
+                    f"  {name:<40s} {backed:>10s} {unbacked:>11s} {'N/A':>8s}",
+                    flush=True,
+                )
+        print(f"{'=' * 80}", flush=True)
+
+    # Build base command, stripping --compare-backed-unbacked, --only, --filter and their values
+    # Handles both space-separated (--filter VALUE) and equals-separated (--filter=VALUE) forms
+    filtered = []
+    skip_next = False
+    for a in sys.argv:
+        if a == "--compare-backed-unbacked":
+            continue
+        if skip_next:
+            skip_next = False
+            continue
+        if a == "--only" or a.startswith("--only="):
+            if "=" not in a:
+                skip_next = True
+            continue
+        if a == "--filter" or a.startswith("--filter="):
+            if "=" not in a:
+                skip_next = True
+            continue
+        filtered.append(a)
+    base_cmd = [sys.executable, "-B"] + filtered
+
+    # Get model list from runner
+    runner.args = args
+    args.filter = args.filter or [r"."]
+    args.exclude = args.exclude or [r"^$"]
+    args.exclude_exact = args.exclude_exact or []
+    models = list(runner.iter_model_names(args))
+
+    if args.only:
+        models = [args.only]
+
+    all_results = {}
+    for model in models:
+        print(f"\n--- {model} ---", flush=True)
+        for mode, flag in [
+            ("backed", "--dynamic-batch-only"),
+            ("unbacked", "--unbacked-batch-only"),
+        ]:
+            cmd = base_cmd + ["--only", model, flag, "--_print-latency-ms"]
+            print(f"  {mode}...", end=" ", flush=True)
+            try:
+                proc = subprocess.Popen(
+                    cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True
+                )
+                stdout, stderr = proc.communicate(timeout=600)
+            except subprocess.TimeoutExpired:
+                proc.kill()
+                print("TIMEOUT", flush=True)
+                continue
+            except Exception as e:
+                print(f"ERROR ({e})", flush=True)
+                continue
+
+            speedup_match = re.search(r"(\d+\.\d+)x", stdout)
+            latency_match = re.search(r"([\d.]+) ms, ([\d.]+) ms,", stdout)
+            if speedup_match:
+                speedup = float(speedup_match.group(1))
+                eager_ms = float(latency_match.group(1)) if latency_match else None
+                compiled_ms = float(latency_match.group(2)) if latency_match else None
+                extra = ""
+                if eager_ms and compiled_ms:
+                    extra = f" (eager={eager_ms:.3f} ms, compiled={compiled_ms:.3f} ms)"
+                print(f"{speedup:.3f}x{extra}", flush=True)
+                if model not in all_results:
+                    all_results[model] = {}
+                all_results[model][mode] = speedup
+                if eager_ms is not None:
+                    all_results[model][f"{mode}_eager_ms"] = eager_ms
+                if compiled_ms is not None:
+                    all_results[model][f"{mode}_ms"] = compiled_ms
+            else:
+                err_match = re.search(
+                    r"(Error|Exception|Traceback).*", stdout + stderr, re.IGNORECASE
+                )
+                if err_match:
+                    print("FAILED", flush=True)
+                else:
+                    print("SKIP", flush=True)
+
+        # Print running diff for this model
+        if (
+            model in all_results
+            and "backed_ms" in all_results[model]
+            and "unbacked_ms" in all_results[model]
+        ):
+            b_ms = all_results[model]["backed_ms"]
+            u_ms = all_results[model]["unbacked_ms"]
+            ms_diff_pct = (u_ms - b_ms) / b_ms * 100
+            print(
+                f"  => diff: {ms_diff_pct:+.1f}% ({b_ms:.3f} ms vs {u_ms:.3f} ms)",
+                flush=True,
+            )
+        elif (
+            model in all_results
+            and "backed" in all_results[model]
+            and "unbacked" in all_results[model]
+        ):
+            b = all_results[model]["backed"]
+            u = all_results[model]["unbacked"]
+            diff_pct = (u - b) / b * 100
+            print(f"  => diff: {diff_pct:+.1f}% (ratio-based, no ms data)", flush=True)
+
+    print_comparison(all_results)
 
 
 def write_csv_when_exception(args, name: str, status: str, device=None):
@@ -3841,7 +3996,7 @@ def write_csv_when_exception(args, name: str, status: str, device=None):
         write_outputs(output_filename, headers, row)
 
 
-def setup_determinism_for_accuracy_test(args):
+def setup_determinism(args):
     if args.only is not None and args.only not in {
         "alexnet",
         "Background_Matting",
@@ -3877,6 +4032,11 @@ def setup_determinism_for_accuracy_test(args):
     torch.backends.cudnn.deterministic = True
     torch.backends.cudnn.benchmark = False
     torch.backends.mkldnn.deterministic = True
+    os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"
+    torch.backends.cudnn.allow_tf32 = False
+    torch.backends.cuda.matmul.allow_tf32 = False
+    torch.backends.cuda.allow_fp16_bf16_reduction_math_sdp(False)
+    patch_torch_manual_seed()
 
 
 def run(runner, args, original_dir=None):
@@ -3937,6 +4097,9 @@ def run(runner, args, original_dir=None):
                 "DLRM+DDP is unsupported as it requires sharding the embedding layer separately from DDP"
             )
             return sys.exit(-1)
+    if args.deterministic and not args.accuracy:
+        setup_determinism(args)
+
     if args.accuracy:
         # Use small batch size. We use >1 batch size to ensure we test
         # batch_norm type of operators that work on batch dims.
@@ -3960,9 +4123,8 @@ def run(runner, args, original_dir=None):
             args.use_eval_mode = True
         inductor_config.fallback_random = True
 
-        setup_determinism_for_accuracy_test(args)
+        setup_determinism(args)
 
-        os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"
         if args.only is not None and args.only in {
             "nvidia_deeprecommender",
         }:
@@ -3970,13 +4132,6 @@ def run(runner, args, original_dir=None):
             torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction = False
             torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction = False
 
-        torch.backends.cudnn.allow_tf32 = False
-        torch.backends.cuda.matmul.allow_tf32 = False
-        torch.backends.cuda.allow_fp16_bf16_reduction_math_sdp(False)
-
-        # Remove randomness when torch manual seed is called
-        patch_torch_manual_seed()
-
         # Some models e.g. yolov3 assert batch size on n_gpus
         if "CUDA_VISIBLE_DEVICES" not in os.environ and not args.multiprocess:
             args.device_index = "0"
@@ -4188,8 +4343,8 @@ def run(runner, args, original_dir=None):
             )
             model_iter_fn = baseline_ctx(runner.model_iter_fn)
 
-            # needed to avoid error that causes inconsistent timing due to:
-            # Unable to hit fast path of CUDAGraphs because of pending, uninvoked backwards
+            # needed to avoid CUDAGraph fast-path warning / inconsistent timing when prior
+            # outputs still require backward (see torch._inductor.cudagraph_trees)
             def model_iter_fn_and_mark_step(*args, **kwargs):
                 torch.compiler.cudagraph_mark_step_begin()
                 model_iter_fn(*args, **kwargs)
@@ -4282,20 +4437,6 @@ def model_iter_fn_and_mark_step(*args, **kwargs):
             write_outputs(output_filename, [], [args.only, batch_size])
         return
 
-    if args.download_only:
-        model_names = list(runner.iter_model_names(args))
-        failed = []
-        for name in model_names:
-            try:
-                runner._download_model(name)
-                print(f"Downloaded: {name}")
-            except Exception as e:
-                print(f"Failed: {name}: {e}")
-                failed.append(name)
-        if failed:
-            sys.exit(1)
-        return
-
     should_profile_details = args.profile_details
     args.profile_details = {}
     if args.export_profiler_trace:
@@ -4483,7 +4624,7 @@ def detect_and_mark_batch(t, use_unbacked=False):
                             # Use duck_shape_id="batch" so all batch dimensions
                             # share the same unbacked symbol
                             torch._dynamo.decorators.mark_unbacked(
-                                t, i, shape_id="batch", hint_override=batch_size
+                                t, i, shape_id="batch", hint_override=batch_size, min=1
                             )
                         else:
                             torch._dynamo.maybe_mark_dynamic(t, i)
@@ -4524,22 +4665,17 @@ def detect_and_mark_batch(t, use_unbacked=False):
             if name in runner.guard_on_nn_module_models:
                 guard_ctx = torch._dynamo.config.patch(guard_nn_modules=True)
 
-            inline_ctx = contextlib.nullcontext()
-            if name in runner.inline_inbuilt_nn_modules_models:
-                inline_ctx = torch._dynamo.config.patch(inline_inbuilt_nn_modules=True)
-
             with guard_ctx:
-                with inline_ctx:
-                    runner.run_one_model(
-                        name,
-                        model,
-                        example_inputs,
-                        optimize_ctx,
-                        experiment,
-                        explain=args.explain,
-                        tag=args.tag,
-                        batch_size=batch_size if args.dynamic_batch_only else None,
-                    )
+                runner.run_one_model(
+                    name,
+                    model,
+                    example_inputs,
+                    optimize_ctx,
+                    experiment,
+                    explain=args.explain,
+                    tag=args.tag,
+                    batch_size=batch_size if args.dynamic_batch_only else None,
+                )
         if args.generate_aot_autograd_stats:
             stats_file = output_filename.split(".csv")[0] + "_stats.csv"
             write_outputs(
diff --git a/benchmarks/dynamo/genai_layers/utils.py b/benchmarks/dynamo/genai_layers/utils.py
index 090e58676f4ed..02e4643f096c3 100644
--- a/benchmarks/dynamo/genai_layers/utils.py
+++ b/benchmarks/dynamo/genai_layers/utils.py
@@ -3,7 +3,7 @@
 from collections import defaultdict
 from collections.abc import Callable
 from dataclasses import dataclass
-from typing import Any, Optional
+from typing import Any
 
 import matplotlib.pyplot as plt
 from scipy.stats import gmean
@@ -242,8 +242,8 @@ def get_backend_colors() -> dict[str, str]:
 
 def visualize_comparison(
     profiling_results: dict[str, list[Performance]],
-    title: Optional[str] = None,
-    output_path: Optional[str] = None,
+    title: str | None = None,
+    output_path: str | None = None,
 ) -> None:
     """
     Create a single memory_bandwidth comparison plot from profiling results.
diff --git a/benchmarks/dynamo/huggingface.py b/benchmarks/dynamo/huggingface.py
index b91777f70ed98..d62d99bfc4ffb 100755
--- a/benchmarks/dynamo/huggingface.py
+++ b/benchmarks/dynamo/huggingface.py
@@ -489,7 +489,7 @@ def generate(self, _, example_inputs, collect_outputs=True):
         else:
             model.eval()
 
-        self.validate_model(model, example_inputs)
+        self.validate_model(model_name, model, example_inputs)
         return device, model_name, model, example_inputs, batch_size
 
     def iter_model_names(self, args):
@@ -552,7 +552,7 @@ def compute_loss(self, pred):
         return pred[0]
 
     def forward_pass(self, mod, inputs, collect_outputs=True):
-        with self.autocast(**self.autocast_arg):
+        with torch.no_grad(), self.autocast(**self.autocast_arg):
             res = mod(**inputs)
         return res.logits if self.hf_llm else res
 
diff --git a/benchmarks/dynamo/launch_compile_op_numerics.py b/benchmarks/dynamo/launch_compile_op_numerics.py
new file mode 100644
index 0000000000000..0b6871b1c08bf
--- /dev/null
+++ b/benchmarks/dynamo/launch_compile_op_numerics.py
@@ -0,0 +1,974 @@
+import argparse
+import asyncio
+import datetime
+import functools
+import hashlib
+import itertools
+import json
+import logging
+import os
+import re
+import shlex
+import sys
+from dataclasses import dataclass, field
+from pathlib import Path
+
+
+log = logging.getLogger(__name__)
+
+try:
+    import torch
+except ImportError:
+    torch = None
+
+PYTORCH_NIGHTLY_CUDA_VERSIONS = ["12.6", "12.8", "13.0"]
+PYTORCH_CUDA_VERSIONS = {
+    "2.9.1": ["12.6", "12.8", "13.0"],
+    "2.9.0": ["12.6", "12.8", "13.0"],
+    "2.8.0": ["12.6", "12.8", "12.9"],
+    "2.7.1": ["11.8", "12.6", "12.8"],
+    "2.7.0": ["11.8", "12.6", "12.8"],
+    "2.6.0": ["11.8", "12.4", "12.6"],
+    "2.5.1": ["11.8", "12.1", "12.4"],
+    "2.5.0": ["11.8", "12.1", "12.4"],
+    "2.4.1": ["11.8", "12.1", "12.4"],
+    "2.4.0": ["11.8", "12.1", "12.4"],
+    "2.3.1": ["11.8", "12.1"],
+    "2.3.0": ["11.8", "12.1"],
+    "2.2.2": ["11.8", "12.1"],
+    "2.2.1": ["11.8", "12.1"],
+    "2.2.0": ["11.8", "12.1"],
+    "2.1.2": ["11.8", "12.1"],
+    "2.1.1": ["11.8", "12.1"],
+    "2.1.0": ["11.8", "12.1"],
+    "2.0.1": ["11.8"],
+    "2.0.0": ["11.8"],
+}
+
+ENABLED_CONFIGS = [
+    #    ("git:5f09e6a6c93e0b5bf75b635cddc03b85bbe85938", "12.8"),
+    ("nightly", "13.0"),
+    # ("2.9.1", "12.8"),
+    # ("2.7.1", "12.8"),
+]
+
+PYTHON_VERSION = "3.11"
+
+
+@dataclass
+class Mode:
+    compile: bool
+    backend: str | None = None
+    mode: str | None = None
+    options: dict | None = None
+    env: dict = field(default_factory=dict)
+
+    def __call__(self, fn):
+        if self.compile:
+            kwargs = {}
+            if self.backend is not None:
+                kwargs["backend"] = self.backend
+            if self.mode is not None:
+                kwargs["mode"] = self.mode
+            if self.options is not None:
+                options = {
+                    k: v
+                    for k, v in self.options.items()
+                    if k in torch._inductor.list_options()
+                }
+                kwargs["options"] = options
+            return torch.compile(fn, **kwargs)
+        assert self.backend is None  # noqa: S101
+        assert self.mode is None  # noqa: S101
+        assert self.options is None  # noqa: S101
+        return fn
+
+
+MODES = {
+    "eager": Mode(compile=False),
+    "decomp": Mode(compile=True, backend="aot_eager_decomp_partition"),
+    "compile_numerics": Mode(
+        compile=True,
+        options={
+            "emulate_precision_casts": True,
+            "use_fast_math": False,
+            "emulate_division_rounding": True,
+            "eager_numerics.division_rounding": True,
+            "eager_numerics.disable_ftz": True,
+        },
+    ),
+    "compile": Mode(compile=True),
+}
+
+DTYPES = ["float32", "float16", "bfloat16"]
+
+
+async def run(command, capture_output=True, capture_stderr=True, stderr=None, **kwargs):
+    log.info("Running command: %s", command)
+    proc = await asyncio.create_subprocess_exec(
+        *shlex.split(command),
+        stdout=asyncio.subprocess.PIPE if capture_output else None,
+        stderr=asyncio.subprocess.STDOUT if capture_stderr else stderr,
+        stdin=asyncio.subprocess.DEVNULL,
+        **kwargs,
+    )
+    if capture_output:
+        result = await proc.communicate()
+        result = result[0].decode("utf-8")
+        log.info("Command result: %s", result)
+        log.info("Return code: %s", proc.returncode)
+        return result
+    else:
+        await proc.wait()
+        log.info("Return code: %s", proc.returncode)
+        return
+
+
+async def copy_file_to_remote(hostname, source, destination):
+    await run(f"scp {source} {hostname}:{destination}")
+
+
+async def execute_on_remote(hostname, command):
+    await run(f"ssh {hostname} {command}", capture_output=False)
+
+
+async def copy_results_from_remote(hostname, run_id):
+    await run(f"mkdir -p results_{run_id}/logs")
+    await run(f"scp {hostname}:/workspace/result_*.jsonl results_{run_id}/")
+    await run(
+        f"scp {hostname}:/workspace/logs.tar.gz results_{run_id}/logs_{hostname}.tar.gz"
+    )
+    await run(f"mkdir -p results_{run_id}/logs")
+    await run(
+        f"tar xzf results_{run_id}/logs_{hostname}.tar.gz -C results_{run_id}/logs --strip 2"
+    )
+    await run(
+        "bash -c 'for f in *.trace; do cd $f; uvx tlparse --no-browser *.log & cd ..; done; wait'",
+        cwd=f"results_{run_id}/logs",
+    )
+    await run(f"rm results_{run_id}/logs_{hostname}.tar.gz")
+    await run(f"bash -c 'rm results_{run_id}/logs/*.trace/*.log'")
+    await run(f"bash -c 'rm results_{run_id}/logs/*.trace/tl_out/raw*'")
+
+
+async def do_numerics_test(args, hostname, gpu, run_id):
+    await copy_file_to_remote(hostname, __file__, "/workspace/run.py")
+    await execute_on_remote(hostname, f"python3 /workspace/run.py --worker --gpu {gpu}")
+    await copy_results_from_remote(hostname, run_id)
+
+
+async def reserve_gpu(gpu):
+    stdout = await run(
+        f"gpu-dev reserve -g 1 -h 24 -t {gpu} --ignore-no-persist --disk none"
+    )
+    # the output has a string like this: SSH Command: ssh gpu-dev-1db65ec7
+    match = re.search(r"SSH Command: ssh (.*)", stdout)
+    hostname = match.group(1)
+    log.info(
+        "Hostname: %s, Reservation ID: %s",
+        hostname,
+        hostname.removeprefix("gpu-dev-"),
+    )
+    return hostname, hostname.removeprefix("gpu-dev-")
+
+
+async def cancel_reservation(reservation_id):
+    await run(f"gpu-dev cancel {reservation_id}")
+
+
+async def run_on_gpu(gpu, command):
+    log.info("Reserving GPU %s", gpu)
+    hostname, reservation_id = await reserve_gpu(gpu)
+    try:
+        log.info("Running command on %s", hostname)
+        await command(hostname=hostname)
+    except Exception:
+        log.exception("Error running command on %s", gpu)
+    finally:
+        log.info("Cancelling reservation %s", reservation_id)
+        await cancel_reservation(reservation_id)
+
+
+async def launcher(args):
+    run_id = datetime.datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
+    await asyncio.gather(
+        *[
+            run_on_gpu(
+                gpu,
+                functools.partial(do_numerics_test, args=args, gpu=gpu, run_id=run_id),
+            )
+            for gpu in args.gpu
+        ]
+    )
+    await run(f"{sys.executable} {__file__} --html {run_id}")
+
+
+def format_results_to_html(run_id):
+    results_dir = Path(f"results_{run_id}")
+    for f in results_dir.glob("*.jsonl"):
+        html_file = f.with_suffix(".html")
+        lines = f.read_text().splitlines(keepends=True)
+        data = [json.loads(line) for line in lines]
+        KEYS = data[0].keys()
+        print(KEYS)
+        FILTER_KEYS = [
+            "gpu",
+            "pytorch_version",
+            "cuda_version",
+            "pytorch_mode",
+            "is_golden",
+            "data_type",
+            "function",
+            "pass_type",
+            "category",
+            "match_full",
+            "match_normal",
+        ]
+        KEY_VALUES = {k: sorted({line[k] for line in data}) for k in FILTER_KEYS}
+        with html_file.open("w") as html_f:
+            html_f.write("<html><body style='font-family: monospace;'>\n")
+            html_f.write(f"<h1>Results for {f.stem}</h1>\n")
+            for key, values in KEY_VALUES.items():
+                html_f.write(
+                    f"<details id='section-{key}'><summary>{key} ({len(values)})</summary>\n"
+                )
+                html_f.write(
+                    f"<a href='#' onclick='document.querySelectorAll(\"#section-{key} input\").forEach(checkbox => checkbox.checked = true);'>Enable all</a>\n"  # noqa: B950
+                )
+                html_f.write(
+                    f"<a href='#' onclick='document.querySelectorAll(\"#section-{key} input\").forEach(checkbox => checkbox.checked = false);'>Disable all</a>\n"  # noqa: B950
+                )
+                for value in values:
+                    safe_value = str(value).replace(".", "_")
+                    html_f.write(
+                        f"<label>\n"
+                        f"<input type='checkbox' id='filter-{key}-{safe_value}' checked>"
+                        f"{value}</label>\n"
+                    )
+                    html_f.write(
+                        f"<style>body:has(#filter-{key}-{safe_value}:not(:checked)) .visible-{key}-{safe_value} {{ display: none; }}</style>\n"  # noqa: B950
+                    )
+                html_f.write("</details>")
+            # one more filter: show "first" mode only
+            html_f.write("<details><summary>First mode that fails only</summary>\n")
+            html_f.write(
+                "<label><input type='checkbox' id='filter-mode-first' checked>"  # noqa: B950
+                "Exclude items where a previous mode also fails</label>\n"
+            )
+            html_f.write("</details>\n")
+            # column filter:
+            html_f.write(
+                "<details id='section-columns'><summary>Column filter</summary>\n"
+            )
+            html_f.write(
+                "<a href='#' onclick='document.querySelectorAll(\"#section-columns input\").forEach(checkbox => checkbox.checked = true);'>Enable all</a>\n"  # noqa: B950
+            )
+            html_f.write(
+                "<a href='#' onclick='document.querySelectorAll(\"#section-columns input\").forEach(checkbox => checkbox.checked = false);'>Disable all</a>\n"  # noqa: B950
+            )
+            DEFAULT_HIDDEN_KEYS = [
+                "gpu",
+                "pytorch_version",
+                "cuda_version",
+                "is_golden",
+            ]
+            for key in KEYS:
+                html_f.write(
+                    f"<label><input type='checkbox' id='filter-column-{key}' {'' if key in DEFAULT_HIDDEN_KEYS else 'checked'}>{key}</label>\n"  # noqa: B950
+                )
+                html_f.write(
+                    f"<style>body:has(#filter-column-{key}:not(:checked)) .visible-column-{key} {{ display: none; }}</style>\n"
+                )
+            html_f.write("</details>\n")
+            html_f.write("<table>\n")
+            html_f.write("<tr>\n")
+            for key in KEYS:
+                html_f.write(f"<th class='visible-column-{key}'>{key}</th>\n")
+
+            html_f.write("<th>log</th>\n")
+            html_f.write("<th>tlparse</th>\n")
+            html_f.write("</tr>\n")
+            for line in lines:
+                data = json.loads(line)
+                classes = [
+                    f"visible-{k}-{str(data[k]).replace('.', '_')}" for k in KEY_VALUES
+                ]
+                html_f.write(f"<tr class='{' '.join(classes)}'>\n")
+                for key in KEYS:
+                    html_f.write(f"<td class='visible-column-{key}'>\n")
+                    if key == "mismatch_sample":
+                        html_f.write(
+                            f"<input type='checkbox' id='filter-mismatch-sample-{data['identifier']}'>\n"
+                        )
+                        html_f.write(
+                            f"<style>body:has(#filter-mismatch-sample-{data['identifier']}:not(:checked)) .visible-mismatch-sample-{data['identifier']} {{ display: none; }}</style>\n"  # noqa: B950
+                        )
+                    else:
+                        html_f.write(f"{data[key]}\n")
+                    html_f.write("</td>\n")
+                html_f.write(
+                    f"<td><a href='logs/{data['identifier']}.log'>log</a></td>"
+                )
+                html_f.write(
+                    f"<td><a href='logs/{data['identifier']}.trace/tl_out/index.html'>trace</a></td>\n"
+                )
+                html_f.write("</tr>\n")
+                if data["mismatch_sample"]:
+                    html_f.write(
+                        f"<tr class='visible-mismatch-sample-{data['identifier']} {' '.join(classes)}'>\n"
+                    )
+                    html_f.write(f"<td colspan='{len(KEYS) + 2}'><pre>\n")
+                    html_f.write("<table>\n")
+                    html_f.write("<tr>\n")
+                    html_f.write("<th>pos</th>\n")
+                    html_f.write("<th>input</th>\n")
+                    html_f.write("<th>output</th>\n")
+                    html_f.write("<th>golden</th>\n")
+                    html_f.write("<th>rel_err</th>\n")
+                    html_f.write("</tr>\n")
+                    for sample in data["mismatch_sample"]:
+                        html_f.write("<tr>\n")
+                        html_f.write(f"<td>{sample['pos']}</td>\n")
+                        html_f.write(
+                            f"<td>{', '.join(map(str, sample['input']))}</td>\n"
+                        )
+                        html_f.write(f"<td>{sample['output']}</td>\n")
+                        html_f.write(f"<td>{sample['golden']}</td>\n")
+                        html_f.write(f"<td>{sample['rel_err']}</td>\n")
+                        html_f.write("</tr>\n")
+                    html_f.write("</table>\n")
+                    html_f.write("</td>\n")
+                    html_f.write("</tr>\n")
+
+            html_f.write("</table>\n")
+            html_f.write("<style>\n")
+            # implement the various filters
+            html_f.write("</style>\n")
+            html_f.write("</body></html>\n")
+
+
+global_utils_lock = asyncio.Lock()
+cuda_download_lock = asyncio.Lock()
+venv_creation_lock = asyncio.Lock()
+pytorch_build_lock = asyncio.Lock()
+numerics_test_lock = asyncio.Lock()
+
+
+async def create_results(args, config):
+    pytorch_version, cuda_version = config
+    async with global_utils_lock:
+        await ensure_global_utils()
+    async with cuda_download_lock:
+        await ensure_cuda(args, cuda_version)
+    async with venv_creation_lock:
+        venv = await create_venv(args, config)
+    async with pytorch_build_lock:
+        await maybe_build_pytorch(args, config, venv)
+    async with numerics_test_lock:
+        await run_numerics_test(args, config, venv)
+
+
+async def ensure_global_utils():
+    if os.environ.get("HAS_GLOBAL_UTILS", "0") == "1":
+        return
+    await run("mkdir global_bin")
+    os.environ["PATH"] = f"/workspace/global_bin:{os.environ.get('PATH', '')}"
+
+    os.environ["UV_INSTALL_DIR"] = "/workspace/global_bin"
+    os.environ["UV_CACHE_DIR"] = "/workspace/.cache/uv"
+    await run("wget --no-verbose https://astral.sh/uv/install.sh")
+    await run("bash install.sh")
+    await run("rm install.sh")
+
+    await run(
+        "wget --no-verbose https://github.com/seeraven/gitcache/releases/download/v1.0.29/gitcache_v1.0.29_Ubuntu22.04_x86_64"  # noqa: B950
+    )
+    await run("mv gitcache_v1.0.29_Ubuntu22.04_x86_64 global_bin/gitcache")
+    await run("chmod +x global_bin/gitcache")
+    which_gitcache = await run("which gitcache")
+    await run(f"ln -s {which_gitcache.strip()} global_bin/git")
+    os.environ["GITCACHE_DIR"] = "/workspace/.cache/gitcache"
+
+    await run("sudo apt-get update")
+    await run("sudo apt-get install -y ccache git-lfs")
+    os.environ["CMAKE_C_COMPILER_LAUNCHER"] = "ccache"
+    os.environ["CMAKE_CXX_COMPILER_LAUNCHER"] = "ccache"
+    os.environ["CMAKE_CUDA_COMPILER_LAUNCHER"] = "ccache"
+    os.environ["CCACHE_DIR"] = "/workspace/.cache/ccache"
+    os.environ["CCACHE_NOHASHDIR"] = "1"
+    os.environ["CCACHE_BASEDIR"] = "/workspace"
+    os.environ["CMAKE_BUILD_PARALLEL_LEVEL"] = "20"
+    os.environ["MAX_JOBS"] = "8"
+    os.environ["MAKEFLAGS"] = "-j20"
+    os.environ["NINJAFLAGS"] = "-j20"
+
+    os.environ["HAS_GLOBAL_UTILS"] = "1"
+
+
+async def ensure_cuda(args, cuda_version):
+    if os.path.exists(f"/workspace/cuda-{cuda_version}"):
+        return
+    await run("sudo apt-get update")
+    await run(
+        f"sudo apt-get install -y cuda-toolkit-{cuda_version.split('.')[0]}-{cuda_version.split('.')[1]}"
+    )
+
+
+async def create_venv(args, config):
+    pytorch_version, cuda_version = config
+    env = os.environ.copy()
+    # cuda_home = f"/usr/local/cuda-{cuda_version}"
+    cuda_home = (
+        f"/usr/local/cuda-{cuda_version.split('.')[0]}.{cuda_version.split('.')[1]}"
+    )
+    env["CUDA_HOME"] = cuda_home
+    env["CUDA_ROOT"] = cuda_home
+    env["PATH"] = f"{cuda_home}/bin:{env.get('PATH', '')}"
+    env["LD_LIBRARY_PATH"] = f"{cuda_home}/lib64:{env.get('LD_LIBRARY_PATH', '')}"
+    env["CUDA_NVCC_EXECUTABLE"] = f"{cuda_home}/bin/nvcc"
+    env["PYTORCH_NVCC"] = f"{cuda_home}/bin/nvcc"
+    venv_dir = f"/workspace/venv-pytorch_{pytorch_version}-cuda_{cuda_version}"
+    await run(f"mkdir -p {venv_dir}")
+    await run(f"uv venv -p {PYTHON_VERSION} --managed-python", cwd=venv_dir, env=env)
+    # update env with venv paths
+    env["VIRTUAL_ENV"] = f"{venv_dir}/.venv"
+    env["PATH"] = f"{venv_dir}/.venv/bin:{env.get('PATH', '')}"
+    await run("uv pip install pip numpy", env=env, cwd=venv_dir)
+
+    index_url = f"https://download.pytorch.org/whl/cu{cuda_version.replace('.', '')}"
+    if pytorch_version == "nightly":
+        await run(
+            f"uv pip install --pre torch --index-url {index_url}", env=env, cwd=venv_dir
+        )
+    elif not pytorch_version.startswith("git:"):
+        await run(
+            f"uv pip install torch=={pytorch_version} --index-url {index_url}",
+            env=env,
+            cwd=venv_dir,
+        )
+    return {"env": env, "cwd": venv_dir}
+
+
+async def maybe_build_pytorch(args, config, venv):
+    pytorch_version, cuda_version = config
+    if pytorch_version.startswith("git:"):
+        raise NotImplementedError("Building PyTorch from git is not supported yet")
+    return
+
+
+async def run_numerics_test(args, config, venv):
+    pytorch_version, cuda_version = config
+    golden_flag = ""
+    if args.create_golden:
+        golden_flag = " --create-golden"
+    for mode in MODES:
+        cmd = (
+            f"python {__file__} --runner --gpu {args.gpu[0]}"
+            f" --pytorch-version {pytorch_version}"
+            f" --cuda-version {cuda_version}"
+            f" --mode {mode}"
+            f" --golden /workspace/golden{golden_flag}"
+        )
+        await run(cmd, capture_output=False, **venv)
+        golden_flag = ""
+
+
+def sortable_config_key(config):
+    pytorch_version, cuda_version = config
+    cuda_version = tuple(map(int, cuda_version.split(".")))
+
+    if pytorch_version == "nightly":
+        pytorch_version = (-1,)
+    elif pytorch_version.startswith("git:"):
+        pytorch_version = (-2,)
+    else:
+        pytorch_version = tuple(map(int, pytorch_version.split(".")))
+
+    return (pytorch_version, cuda_version)
+
+
+async def worker(args):
+    assert len(args.gpu) == 1  # noqa: S101
+    # find golden task
+    # a config is a tuple (pytorch_version, cuda_version)
+    # where pytorch_version can either be a version string or nightly or a git hash
+    golden_config = sorted(ENABLED_CONFIGS, key=sortable_config_key)[-1]
+    await run("mkdir -p /workspace/logs")
+    args.create_golden = True
+    await create_results(args, golden_config)
+    args.create_golden = False
+    await asyncio.gather(
+        *[
+            create_results(args, config)
+            for config in ENABLED_CONFIGS
+            if config != golden_config
+        ]
+    )
+    await run("tar czf /workspace/logs.tar.gz /workspace/logs")
+
+
+CATEGORIES = {
+    "reduction": ["torch.sum", "torch.mean", "torch.softmax"],
+    "matrix": ["torch.matmul"],
+    "normalization": [
+        "torch.nn.functional.layer_norm",
+        "torch.nn.functional.rms_norm",
+    ],
+    "activation": [
+        "torch.nn.functional.relu",
+        "torch.nn.functional.sigmoid",
+        "torch.nn.functional.tanh",
+        "torch.nn.functional.gelu",
+        "torch.nn.functional.silu",
+    ],
+    "elementary": [
+        "torch.sin",
+        "torch.cos",
+        "torch.tan",
+        "torch.sigmoid",
+        "torch.exp",
+        "torch.exp2",
+        "torch.log",
+        "torch.log2",
+        "torch.sqrt",
+        "torch.erf",
+        "torch.reciprocal",
+        "torch.rsqrt",
+    ],
+    "binary": [
+        "torch.add",
+        "torch.sub",
+        "torch.mul",
+        "torch.div",
+        "torch.pow",
+    ],
+}
+
+PASS_TYPES = {
+    "reduction": ["fwd", "bwd_0"],
+    "matrix": ["fwd", "bwd_0", "bwd_1"],
+    "normalization": ["fwd", "bwd_0", "bwd_1"],
+    "activation": ["fwd", "bwd_0"],
+    "elementary": ["fwd", "bwd_0"],
+    "binary": ["fwd", "bwd_0", "bwd_1"],
+}
+
+PASSES = {}
+
+
+def register_pass(pass_type):
+    def decorator(func):
+        PASSES[pass_type] = func
+        return func
+
+    return decorator
+
+
+@register_pass("fwd")
+def pass_fwd(callable):
+    def wrapper(*args):
+        return callable(*args)
+
+    return wrapper
+
+
+@register_pass("bwd_0")
+def pass_bwd_0(callable):
+    def wrapper(*args):
+        args = [arg.detach().requires_grad_(idx == 0) for idx, arg in enumerate(args)]
+        output = callable(*args)
+        output.sum().backward()
+        return args[0].grad.detach()
+
+    return wrapper
+
+
+@register_pass("bwd_1")
+def pass_bwd_1(callable):
+    def wrapper(*args):
+        args = [arg.detach().requires_grad_(idx == 1) for idx, arg in enumerate(args)]
+        output = callable(*args)
+        output.sum().backward()
+        return args[1].grad.detach()
+
+    return wrapper
+
+
+def filter_nan(a):
+    return torch.where(torch.isnan(a), 0, a)
+
+
+def mantissa_bits(dtype) -> int:
+    return {torch.float32: 23, torch.bfloat16: 5, torch.float16: 10}[dtype]
+
+
+def float_to_int_type(dtype):
+    return {
+        torch.float32: torch.uint32,
+        torch.bfloat16: torch.uint16,
+        torch.float16: torch.uint16,
+    }[dtype]
+
+
+def generate_test_tensor(dtype, slow=False):
+    if dtype in [torch.float16, torch.bfloat16]:
+        return filter_nan(
+            torch.arange(0, 2**16, dtype=torch.int32, device="cuda")
+            .to(torch.uint16)
+            .view(dtype)
+        )
+    if dtype == torch.float32:
+        if slow:
+            return filter_nan(
+                torch.arange(0, 2**32, dtype=torch.int64, device="cuda")
+                .to(torch.uint32)
+                .view(dtype)
+            )
+        else:
+            # E8M5, E5M10, E8M23
+            result = []
+            for t in [
+                generate_test_tensor(torch.bfloat16),
+                generate_test_tensor(torch.float16),
+            ]:
+                orig_dtype = t.dtype
+                t = t.to(dtype).view(torch.uint32)
+                missing_mantissa_bits = mantissa_bits(dtype) - mantissa_bits(orig_dtype)
+                r = torch.randint(
+                    0, 2**missing_mantissa_bits, t.shape, dtype=t.dtype, device="cuda"
+                )
+                result.append(t.view(dtype))
+                result.append((t.view(torch.int32) | r.view(torch.int32)).view(dtype))
+            return filter_nan(torch.cat(result))
+
+
+def make_input(dtype, category):
+    dtype = getattr(torch, dtype)
+    from torch.testing import make_tensor
+
+    if category in ["matrix"]:
+        return make_tensor(1024, 1024, dtype=dtype, device="cuda"), make_tensor(
+            1024, 1024, dtype=dtype, device="cuda"
+        )
+    if category == "normalization":
+        return make_tensor(1024, 1024, dtype=dtype, device="cuda"), make_tensor(
+            1024, dtype=dtype, device="cuda"
+        )
+    if category == "reduction":
+        return (make_tensor(1024, 1024, dtype=dtype, device="cuda"),)
+    if category == "binary":
+        test_tensor = generate_test_tensor(dtype)
+        # randomly permute twice:
+        return (
+            test_tensor[torch.randperm(test_tensor.shape[0])],
+            test_tensor[torch.randperm(test_tensor.shape[0])],
+        )
+    return (generate_test_tensor(dtype),)
+
+
+def make_function(function, category):
+    if category == "reduction":
+        return eval(f"lambda x: {function}(x, dim=-1)")
+    elif category == "matrix":
+        return eval(f"lambda x, y: {function}(x, y)")
+    elif category == "normalization":
+        return eval(f"lambda x, y: {function}(x, [x.shape[-1]], y, eps=1e-5)")
+    elif category == "binary":
+        return eval(f"lambda x, y: {function}(x, y)")
+    elif category == "activation":
+        return eval(f"lambda x: {function}(x)")
+    elif category == "elementary":
+        return eval(f"lambda x: {function}(x)")
+    else:
+        raise ValueError(f"Unknown category: {category}")
+
+
+def rel_err_ulp(a, b, dtype):
+    return (a - b).abs() / (b.abs() * torch.finfo(dtype).eps + torch.finfo(dtype).tiny)
+
+
+def evaluate_output(input, output, golden_output):
+    # todo: how to make sure this handles denormals in the input well?
+    # for now, only handle it if the inputs have the same shape, otherwise assume there are none
+    # handle more than one input
+
+    input = [i.flatten().float() for i in input if i.shape == output.shape]
+
+    dtype = golden_output.dtype
+
+    output = output.flatten().float()
+    golden_output = golden_output.flatten().float()
+
+    # we are checking subnormals separate from the rest of the numbers
+    # we also need to check NaNs and Infs carefully
+
+    subnormal_mask = golden_output.abs() < torch.finfo(dtype).smallest_normal
+    for i in input:
+        subnormal_mask |= i.abs() < torch.finfo(dtype).smallest_normal
+    nan_mask_golden = torch.isnan(golden_output)
+    golden_output = torch.where(nan_mask_golden, float("nan"), golden_output)
+    nan_mask_output = torch.isnan(output)
+    output = torch.where(nan_mask_output, float("nan"), output)
+
+    equal_mask = (output == golden_output) | (nan_mask_golden & nan_mask_output)
+    equal_subnormal_mask = (
+        torch.where(subnormal_mask, 0, output)
+        == torch.where(subnormal_mask, 0, golden_output)
+    ) | equal_mask
+    output_flushed = torch.where(subnormal_mask, 0.0, output)
+    golden_flushed = torch.where(subnormal_mask, 0.0, golden_output)
+    output_flushed = torch.where(~torch.isfinite(output_flushed), 0.0, output_flushed)
+    golden_flushed = torch.where(~torch.isfinite(golden_flushed), 0.0, golden_flushed)
+
+    num_nonequal = (~equal_mask).sum().item()
+    num_nonequal_subnormal = (~equal_subnormal_mask).sum().item()
+
+    err = rel_err_ulp(output_flushed, golden_flushed, dtype)
+    max_ulp_to_golden = err.max().item()
+    avg_ulp_to_golden = err.mean().item()
+    mismatch_sample = []
+    pos = (~equal_mask).nonzero().squeeze(1)
+    log.info(
+        "pos.shape: %s, input: %s, golden_output.shape: %s, output.shape: %s",
+        pos.shape,
+        [i.shape for i in input],
+        golden_output.shape,
+        output.shape,
+    )
+    ordered = torch.argsort(err[pos], descending=True)
+    random = torch.randperm(pos.shape[0])
+    for sampling in [ordered, random]:
+        for i in range(min(pos.shape[0], 5)):
+            sample_idx = pos[sampling[i]]
+            mismatch_sample.append(
+                {
+                    "pos": sample_idx.item(),
+                    "input": [inp[sample_idx].item() for inp in input],
+                    "output": output[sample_idx].item(),
+                    "golden": golden_output[sample_idx].item(),
+                    "rel_err": err[sample_idx].item(),
+                }
+            )
+
+    # use hashlib.sha256 to hash the tensors
+    # this i
+    return {
+        "normal_hash": hashlib.sha256(
+            output_flushed.cpu().numpy().tobytes()
+        ).hexdigest()[:8],
+        "full_hash": hashlib.sha256(output.cpu().numpy().tobytes()).hexdigest()[:8],
+        "max_ulp_to_golden": max_ulp_to_golden,
+        "avg_ulp_to_golden": avg_ulp_to_golden,
+        "num_nonequal": num_nonequal,
+        "num_nonequal_subnormal": num_nonequal_subnormal,
+        "num_total": output.shape[0],
+        "match_full": num_nonequal == 0,
+        "match_normal": num_nonequal_subnormal == 0,
+        "mismatch_sample": mismatch_sample,
+    }
+
+
+def create_golden(args):
+    golden_inputs = {
+        dtype: {category: make_input(dtype, category) for category in CATEGORIES}
+        for dtype in DTYPES
+    }
+    golden_outputs = {
+        dtype: {
+            category: {
+                function: {
+                    pass_type: MODES[args.mode](
+                        PASSES[pass_type](make_function(function, category))
+                    )(*golden_inputs[dtype][category])
+                    for pass_type in PASS_TYPES[category]
+                }
+                for function in CATEGORIES[category]
+            }
+            for category in CATEGORIES
+        }
+        for dtype in DTYPES
+    }
+    return golden_inputs, golden_outputs
+
+
+def get_metadata(args, dtype, category, function, pass_type):
+    metadata = {
+        "gpu": args.gpu[0],
+        "pytorch_version": args.pytorch_version,
+        "cuda_version": args.cuda_version,
+        "pytorch_mode": args.mode,
+        "is_golden": args.create_golden,
+        "data_type": dtype,
+        "function": function,
+        "pass_type": pass_type,
+        "category": category,
+    }
+    identifier = hashlib.sha256(json.dumps(metadata).encode()).hexdigest()[:8]
+    metadata["identifier"] = identifier
+    return metadata
+
+
+def run_test_case(args):
+    golden = torch.load(args.golden)
+    golden_inputs, golden_outputs = golden
+    metadata = get_metadata(
+        args, args.dtype, args.category, args.function, args.pass_type
+    )
+    input = golden_inputs[args.dtype][args.category]
+    callable = make_function(args.function, args.category)
+    callable = PASSES[args.pass_type](callable)
+    callable = MODES[args.mode](callable)
+    output = callable(*input)
+    evaluation = evaluate_output(
+        input,
+        output,
+        golden_outputs[args.dtype][args.category][args.function][args.pass_type],
+    )
+    data = json.dumps(metadata | evaluation)
+    log.info(data)
+    print(data)
+
+
+test_concurrency = asyncio.Semaphore(16)
+
+
+async def launch_test_case(args, dtype, category, function, pass_type):
+    async with test_concurrency:
+        env = os.environ | MODES[args.mode].env
+        metadata = get_metadata(args, dtype, category, function, pass_type)
+        identifier = metadata["identifier"]
+        env["TORCH_TRACE"] = f"/workspace/logs/{identifier}.trace"
+        Path(env["TORCH_TRACE"]).mkdir(parents=True, exist_ok=True)
+        env["TORCH_LOGS"] = "+all"
+        env["TORCHINDUCTOR_FORCE_DISABLE_CACHES"] = "1"
+        env["TRITON_ALWAYS_COMPILE"] = "1"
+        cmd_args = [
+            "--test",
+            f"--gpu {args.gpu[0]}",
+            f"--pytorch-version {args.pytorch_version}",
+            f"--cuda-version {args.cuda_version}",
+            f"--mode {args.mode}",
+            f"--golden {args.golden}",
+            f"--dtype {dtype}",
+            f"--category {category}",
+            f"--function {function}",
+            f"--pass-type {pass_type}",
+        ]
+        if args.create_golden:
+            cmd_args.append("--create-golden")
+        return await run(
+            f"{sys.executable} {__file__} {' '.join(cmd_args)}",
+            env=env,
+            capture_stderr=False,
+            stderr=open(f"/workspace/logs/{identifier}.log", "w"),
+        )
+
+
+async def runner(args):
+    log.info("Runner called! %s", args)
+
+    if args.create_golden:
+        golden = create_golden(args)
+        torch.save(golden, args.golden)
+
+    results = await asyncio.gather(
+        *[
+            launch_test_case(
+                args,
+                dtype,
+                category,
+                function,
+                pass_type,
+            )
+            for dtype, category in itertools.product(DTYPES, CATEGORIES)
+            for function, pass_type in itertools.product(
+                CATEGORIES[category], PASS_TYPES[category]
+            )
+        ]
+    )
+
+    with open(
+        f"/workspace/result_{args.gpu[0]}_{args.pytorch_version}_{args.cuda_version}.jsonl",
+        "a+",
+    ) as result:
+        for r in results:
+            result.write(r)
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--launcher",
+        default=False,
+        action="store_true",
+        help="launch workers across machines",
+    )
+    parser.add_argument(
+        "--worker",
+        default=False,
+        action="store_true",
+        help="launches jobs on a single machine",
+    )
+    parser.add_argument(
+        "--runner",
+        default=False,
+        action="store_true",
+        help="runs tests in a specific environment",
+    )
+    parser.add_argument(
+        "--test",
+        action="store_true",
+        help="runs a single test case",
+    )
+    parser.add_argument("--gpu", nargs="+", type=str, default=["t4"])
+    parser.add_argument("--html", type=str, default=None)
+    parser.add_argument("--golden", type=str, default=None)
+    parser.add_argument("--create-golden", action="store_true")
+    parser.add_argument("--dtype", type=str)
+    parser.add_argument("--category", type=str)
+    parser.add_argument("--function", type=str)
+    parser.add_argument("--pass-type", type=str)
+    parser.add_argument("--pytorch-version", type=str)
+    parser.add_argument("--cuda-version", type=str)
+    parser.add_argument("--mode")
+    args = parser.parse_args()
+    if args.runner:
+        assert torch is not None  # noqa: S101
+        torch.set_default_device("cuda")
+        (gpu,) = args.gpu
+        logging.basicConfig(
+            level=logging.INFO,
+            format=f"%(asctime)s - runner:{gpu}/{args.pytorch_version}/{args.cuda_version}/{args.mode} - %(message)s",  # noqa: B950
+        )
+        asyncio.run(runner(args))
+    if args.worker:
+        (gpu,) = args.gpu
+        logging.basicConfig(
+            level=logging.INFO, format=f"%(asctime)s - worker:{gpu} - %(message)s"
+        )
+        asyncio.run(worker(args))
+    if args.launcher:
+        logging.basicConfig(
+            level=logging.INFO, format="%(asctime)s - launcher - %(message)s"
+        )
+        try:
+            asyncio.run(launcher(args))
+        except KeyboardInterrupt:
+            log.error("Cancelled by user")
+    if args.test:
+        test_id = (
+            f"{args.gpu[0]}/{args.pytorch_version}/{args.cuda_version}"
+            f"/{args.mode}/{args.dtype}/{args.category}"
+            f"/{args.function}/{args.pass_type}"
+        )
+        logging.basicConfig(
+            level=logging.INFO,
+            format=f"%(asctime)s - test:{test_id} - %(message)s",
+        )
+        run_test_case(args)
+    if args.html:
+        format_results_to_html(args.html)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/benchmarks/dynamo/perf_cli.py b/benchmarks/dynamo/perf_cli.py
new file mode 100644
index 0000000000000..f4c85a5e4a2a8
--- /dev/null
+++ b/benchmarks/dynamo/perf_cli.py
@@ -0,0 +1,1894 @@
+#!/usr/bin/env python3
+"""
+CLI tool for launching, summarizing, and reproducing inductor perf regression runs.
+
+Usage:
+    python benchmarks/dynamo/perf_cli.py launch [--device a100 h100 ...] [--ref BRANCH] [--wait]
+    python benchmarks/dynamo/perf_cli.py summary <run-id|branch> [--top 5] [--config PATTERN]
+    python benchmarks/dynamo/perf_cli.py repro <run-id> [--model MODEL] [--suite SUITE] [--print-only]
+
+Requires: gh CLI (authenticated), internet access to S3 (gha-artifacts bucket).
+"""
+
+from __future__ import annotations
+
+import argparse
+import csv
+import io
+import json
+import os
+import re
+import shutil
+import subprocess
+import sys
+import time
+import urllib.request
+import zipfile
+from collections import defaultdict
+from dataclasses import dataclass, field
+from datetime import datetime
+from math import exp, log
+from pathlib import Path
+
+
+WORKFLOWS = {
+    "a100": {
+        "name": "inductor-A100-perf-nightly",
+        "id": 42513231,
+    },
+    "a100-compare": {
+        "name": "inductor-A100-perf-compare",
+        "id": 50531883,
+    },
+    "h100": {
+        "name": "inductor-perf-nightly-h100",
+        "id": 144201955,
+    },
+    "b200": {
+        "name": "inductor-perf-b200",
+        "id": 173716622,
+    },
+    "rocm-mi300": {
+        "name": "inductor-perf-nightly-rocm-mi300",
+        "id": 197925166,
+    },
+    "rocm-mi355": {
+        "name": "inductor-perf-nightly-rocm-mi355",
+        "id": 197925165,
+    },
+    "x86": {
+        "name": "inductor-perf-nightly-x86",
+        "id": 108782874,
+    },
+    "x86-zen": {
+        "name": "inductor-perf-nightly-x86-zen",
+        "id": 167573808,
+    },
+    "aarch64": {
+        "name": "inductor-perf-nightly-aarch64",
+        "id": 109196799,
+    },
+    "macos": {
+        "name": "inductor-perf-nightly-macos",
+        "id": 117199085,
+    },
+    "xpu": {
+        "name": "inductor-perf-nightly-xpu",
+        "id": 201149053,
+    },
+}
+
+DEVICE_CHOICES = sorted(k for k in WORKFLOWS if k != "a100-compare")
+
+S3_BUCKET = "gha-artifacts"
+S3_URL = f"https://{S3_BUCKET}.s3.amazonaws.com"
+REPO = "pytorch/pytorch"
+
+# Regex to parse test job names like:
+# "cuda13.0-py3.10-gcc11-sm80 / test (inductor_huggingface_perf, 1, 5, linux.aws.a100)"
+JOB_RE = re.compile(
+    r"test \((?P<config>[^,]+),\s*(?P<shard>\d+),\s*(?P<num_shards>\d+),\s*(?P<runner>[^)]+)\)"
+)
+
+PERF_CONFIGS = re.compile(r"inductor_(huggingface|timm|torchbench)_perf")
+
+SUITE_ALIASES = {
+    "hf": "huggingface",
+    "huggingface": "huggingface",
+    "timm": "timm_models",
+    "timm_models": "timm_models",
+    "tb": "torchbench",
+    "torchbench": "torchbench",
+}
+
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+
+def gh(*args: str, json_output: bool = False) -> str | dict | list:
+    cmd = ["gh"] + list(args)
+    result = subprocess.run(cmd, capture_output=True, text=True)
+    if result.returncode != 0:
+        print(f"gh error: {result.stderr.strip()}", file=sys.stderr)
+        sys.exit(1)
+    if json_output:
+        return json.loads(result.stdout)
+    return result.stdout.strip()
+
+
+def git(*args: str) -> str:
+    result = subprocess.run(["git"] + list(args), capture_output=True, text=True)
+    return result.stdout.strip()
+
+
+def gmean(values: list[float]) -> float:
+    if not values:
+        return 0.0
+    return exp(sum(log(v) for v in values if v > 0) / max(len(values), 1))
+
+
+@dataclass
+class Metric:
+    name: str
+    field: str  # attribute on ModelResult
+    unit: str  # display suffix
+    higher_is_better: bool
+    aggregate: str  # "gmean" or "mean"
+
+
+METRICS = {
+    "speedup": Metric("speedup", "speedup", "x", True, "gmean"),
+    "compilation_latency": Metric(
+        "compilation latency", "compilation_latency", "s", False, "mean"
+    ),
+    "compression_ratio": Metric(
+        "memory compression", "compression_ratio", "x", True, "gmean"
+    ),
+    "abs_latency": Metric("absolute latency", "abs_latency", "ms", False, "mean"),
+}
+
+METRIC_CHOICES = list(METRICS.keys())
+
+# HUD uses 5% relative threshold for flagging regressions
+RELATIVE_THRESHOLD = 0.05
+
+WORKFLOW_NAME_TO_NIGHTLY_ID = {
+    v["name"]: v["id"] for k, v in WORKFLOWS.items() if k != "a100-compare"
+}
+# compare's baseline is A100 nightly
+WORKFLOW_NAME_TO_NIGHTLY_ID["inductor-A100-perf-compare"] = WORKFLOWS["a100"]["id"]
+
+
+def _short_config(config: str, device: str = "") -> str:
+    """Compact label: e.g. 'a100 cudagraphs huggingface training'."""
+    c = config
+    backend = c
+    for s in ("_huggingface_", "_timm_models_", "_torchbench_"):
+        if s in c:
+            backend = c.split(s)[0]
+            break
+    backend = backend.removeprefix("inductor_")
+    mode = "training" if "training" in c else "inference" if "inference" in c else ""
+    suite = ""
+    for s in ("huggingface", "timm_models", "torchbench"):
+        if s in c:
+            suite = s
+            break
+    parts = [p for p in (device, backend, suite, mode) if p]
+    return " ".join(parts)
+
+
+@dataclass
+class ModelResult:
+    name: str
+    speedup: float
+    abs_latency: float = 0.0
+    compilation_latency: float = 0.0
+    compression_ratio: float = 0.0
+    eager_peak_mem: float = 0.0
+    dynamo_peak_mem: float = 0.0
+    config: str = ""
+    device: str = ""
+
+    @property
+    def short_config(self) -> str:
+        return _short_config(self.config, self.device)
+
+
+@dataclass
+class PerfData:
+    config: str  # e.g. "inductor_with_cudagraphs_huggingface_amp_training_cuda"
+    models: list[ModelResult] = field(default_factory=list)
+    device: str = ""
+
+    @property
+    def suite(self) -> str:
+        for s in ("huggingface", "timm_models", "torchbench"):
+            if s in self.config:
+                return s
+        return "unknown"
+
+    @property
+    def mode(self) -> str:
+        if "training" in self.config:
+            return "training"
+        if "inference" in self.config:
+            return "inference"
+        return "unknown"
+
+    @property
+    def dtype(self) -> str:
+        # Config format: {backend}_{suite}_{dtype}_{mode}_{device}
+        # e.g. inductor_with_cudagraphs_huggingface_amp_training_cuda
+        for s in ("_huggingface_", "_timm_models_", "_torchbench_"):
+            if s in self.config:
+                tail = self.config.split(s, 1)[1]
+                # tail is e.g. "amp_training_cuda"
+                parts = tail.split("_")
+                if parts:
+                    return parts[0]
+        return "unknown"
+
+    @property
+    def short_name(self) -> str:
+        name = self.config
+        for s in ("_huggingface_", "_timm_models_", "_torchbench_"):
+            if s in name:
+                name = name.split(s)[0]
+                break
+        return name
+
+    @property
+    def runtime(self) -> str:
+        # Last token in config: e.g. "..._training_cuda" → "cuda"
+        parts = self.config.rsplit("_", 1)
+        if len(parts) == 2:
+            return parts[1]
+        return "unknown"
+
+    @property
+    def qualified_config(self) -> str:
+        if self.device:
+            return f"{self.device}/{self.config}"
+        return self.config
+
+    def gmean_speedup(self) -> float:
+        vals = [m.speedup for m in self.models if m.speedup > 0]
+        return gmean(vals)
+
+    def aggregate_metric(self, metric: Metric) -> float:
+        vals = [
+            getattr(m, metric.field)
+            for m in self.models
+            if getattr(m, metric.field) > 0
+        ]
+        if not vals:
+            return 0.0
+        if metric.aggregate == "gmean":
+            return gmean(vals)
+        return sum(vals) / len(vals)
+
+
+# ---------------------------------------------------------------------------
+# Artifact downloading
+# ---------------------------------------------------------------------------
+
+
+def get_run_jobs(run_id: int) -> list[dict]:
+    data = gh(
+        "run",
+        "view",
+        str(run_id),
+        "--repo",
+        REPO,
+        "--json",
+        "jobs",
+        json_output=True,
+    )
+    return data["jobs"]
+
+
+def get_perf_jobs(jobs: list[dict]) -> list[dict]:
+    perf_jobs = []
+    for job in jobs:
+        m = JOB_RE.search(job["name"])
+        if not m:
+            continue
+        config = m.group("config")
+        if not PERF_CONFIGS.match(config):
+            continue
+        if job.get("conclusion") != "success":
+            continue
+        perf_jobs.append(
+            {
+                "config": config,
+                "shard": m.group("shard"),
+                "num_shards": m.group("num_shards"),
+                "runner": m.group("runner"),
+                "job_id": job["databaseId"],
+                "name": job["name"],
+            }
+        )
+    return perf_jobs
+
+
+def s3_artifact_url(run_id: int, attempt: int, job: dict) -> str:
+    config = job["config"]
+    shard = job["shard"]
+    num_shards = job["num_shards"]
+    runner = job["runner"]
+    job_id = job["job_id"]
+    filename = f"test-reports-test-{config}-{shard}-{num_shards}-{runner}_{job_id}.zip"
+    return f"{S3_URL}/{REPO}/{run_id}/{attempt}/artifact/{filename}"
+
+
+CACHE_DIR = Path.home() / ".cache" / "perf_cli"
+
+
+def get_cache_dir(run_id: int, attempt: int) -> Path:
+    d = CACHE_DIR / f"{run_id}" / f"{attempt}"
+    d.mkdir(parents=True, exist_ok=True)
+    return d
+
+
+def download_and_extract_csvs(
+    run_id: int,
+    jobs: list[dict],
+    attempt: int = 1,
+    no_cache: bool = False,
+) -> list[tuple[str, str]]:
+    """Download artifacts and return list of (csv_filename, csv_content) pairs."""
+    cache = get_cache_dir(run_id, attempt)
+    results = []
+    fetched = 0
+
+    for job in jobs:
+        # Check cache first
+        cache_key = (
+            f"{job['config']}-{job['shard']}-{job['num_shards']}-{job['job_id']}"
+        )
+        cache_marker = cache / f"{cache_key}.done"
+
+        if not no_cache and cache_marker.exists():
+            # Read cached CSVs
+            for csv_file in cache.glob(f"{cache_key}__*.csv"):
+                csv_name = csv_file.name.split("__", 1)[1]
+                results.append((csv_name, csv_file.read_text()))
+            continue
+
+        url = s3_artifact_url(run_id, attempt, job)
+        zip_path = cache / f"{cache_key}.zip"
+        try:
+            urllib.request.urlretrieve(url, str(zip_path))
+            fetched += 1
+        except urllib.error.HTTPError as e:
+            print(
+                f"  warning: failed to download shard {job['config']} "
+                f"shard {job['shard']}: {e}",
+                file=sys.stderr,
+            )
+            continue
+
+        with zipfile.ZipFile(zip_path) as zf:
+            for name in zf.namelist():
+                if name.endswith("_performance.csv"):
+                    csv_name = os.path.basename(name)
+                    with zf.open(name) as f:
+                        content = f.read().decode("utf-8")
+                        results.append((csv_name, content))
+                        # Write to cache
+                        (cache / f"{cache_key}__{csv_name}").write_text(content)
+
+        # Mark this shard as cached and remove the zip
+        cache_marker.touch()
+        zip_path.unlink(missing_ok=True)
+
+    if fetched > 0:
+        print(f"  downloaded {fetched} shards (cached at {cache})")
+    elif results:
+        print(f"  using cached data from {cache}")
+
+    return results
+
+
+def parse_csvs(csv_pairs: list[tuple[str, str]], device: str = "") -> list[PerfData]:
+    grouped: dict[str, list[ModelResult]] = defaultdict(list)
+
+    for csv_name, content in csv_pairs:
+        # csv_name like: inductor_with_cudagraphs_huggingface_amp_training_cuda_performance.csv
+        config = csv_name.replace("_performance.csv", "")
+        reader = csv.DictReader(io.StringIO(content))
+        for row in reader:
+            try:
+                speedup = float(row.get("speedup", 0))
+            except (ValueError, TypeError):
+                continue
+            grouped[config].append(
+                ModelResult(
+                    name=row.get("name", "?"),
+                    speedup=speedup,
+                    abs_latency=float(row.get("abs_latency", 0) or 0),
+                    compilation_latency=float(row.get("compilation_latency", 0) or 0),
+                    compression_ratio=float(row.get("compression_ratio", 0) or 0),
+                    eager_peak_mem=float(row.get("eager_peak_mem", 0) or 0),
+                    dynamo_peak_mem=float(row.get("dynamo_peak_mem", 0) or 0),
+                    config=config,
+                    device=device,
+                )
+            )
+
+    return [
+        PerfData(config=k, models=v, device=device) for k, v in sorted(grouped.items())
+    ]
+
+
+# ---------------------------------------------------------------------------
+# S-curve rendering
+# ---------------------------------------------------------------------------
+
+
+def subsample(items: list, max_rows: int) -> list:
+    """Evenly subsample a sorted list, always keeping first and last."""
+    n = len(items)
+    if n <= max_rows:
+        return items
+    # Always include first and last; evenly space the rest
+    indices = {0, n - 1}
+    for i in range(1, max_rows - 1):
+        indices.add(round(i * (n - 1) / (max_rows - 1)))
+    return [items[i] for i in sorted(indices)]
+
+
+def render_scurve(
+    perf: PerfData,
+    metric: Metric,
+    top_n: int = 5,
+    term_width: int | None = None,
+    term_height: int | None = None,
+):
+    if not perf.models:
+        return
+
+    if term_width is None or term_height is None:
+        sz = shutil.get_terminal_size((100, 50))
+        term_width = term_width or sz.columns
+        term_height = term_height or sz.lines
+
+    live = [m for m in perf.models if getattr(m, metric.field) > 0]
+    if not live:
+        return
+
+    sorted_models = sorted(live, key=lambda m: getattr(m, metric.field))
+    agg = perf.aggregate_metric(metric)
+    n = len(sorted_models)
+
+    # Reserve lines for header (2) + axis label (1) + padding (2)
+    max_rows = max(term_height - 5, 15)
+    display = subsample(sorted_models, max_rows)
+    skipped = n - len(display)
+
+    agg_label = metric.aggregate
+    header = f"{perf.config} ({n} data points, {agg_label}={agg:.2f}{metric.unit})"
+    if skipped > 0:
+        header += f" [showing {len(display)}/{n}]"
+    print(f"\n  {header}")
+    print(f"  {'─' * min(len(header), term_width - 4)}")
+
+    def fmt_val(v: float) -> str:
+        if metric.unit == "s" or metric.unit == "ms":
+            return f"{v:7.1f}{metric.unit}"
+        return f"{v:5.2f}{metric.unit}"
+
+    # Layout: "  {name:<max_name}  {val:>8}  {dots}"
+    sample_val = fmt_val(display[0] and getattr(display[0], metric.field))
+    max_name = min(max(len(m.name) for m in display), 30)
+    val_width = len(sample_val)
+    prefix_len = 2 + max_name + 2 + val_width + 2
+    plot_width = max(term_width - prefix_len - 1, 20)
+
+    def get_val(m):
+        return getattr(m, metric.field)
+
+    min_val = get_val(sorted_models[0])
+    p95_idx = max(0, int(n * 0.95) - 1)
+    p95_val = get_val(sorted_models[p95_idx])
+
+    # For ratio metrics (speedup, compression_ratio), anchor at 1.0
+    # For absolute metrics (latency), anchor at 0
+    if metric.unit == "x":
+        plot_min = min(min_val, 0.5)
+        plot_max = max(p95_val * 1.1, 1.5)
+        marker_val = 1.0
+        marker_label = "1.0x"
+    else:
+        plot_min = 0
+        plot_max = p95_val * 1.1
+        marker_val = None
+        marker_label = None
+
+    span = plot_max - plot_min
+    if span == 0:
+        span = 1
+
+    def val_to_col(v: float) -> int:
+        return max(
+            0, min(plot_width - 1, int((v - plot_min) / span * (plot_width - 1)))
+        )
+
+    marker_col = val_to_col(marker_val) if marker_val is not None else None
+
+    for m in display:
+        name = m.name[:max_name].ljust(max_name)
+        v = get_val(m)
+        col = val_to_col(v)
+        bar = [" "] * plot_width
+        if marker_col is not None:
+            bar[marker_col] = "|"
+        for i in range(col + 1):
+            if marker_col is not None and i == marker_col:
+                bar[i] = "|"
+            else:
+                bar[i] = "·"
+        print(f"  {name}  {fmt_val(v)}  {''.join(bar)}")
+
+    pad = " " * prefix_len
+    if marker_label and marker_col is not None:
+        print(f"{pad}{' ' * marker_col}{marker_label}")
+    else:
+        print()
+
+
+def print_worst_offenders(perf: PerfData, metric: Metric, top_n: int = 5):
+    def get_val(m):
+        return getattr(m, metric.field)
+
+    live = [m for m in perf.models if get_val(m) > 0]
+    if not live:
+        return
+    # "worst" depends on metric direction
+    if metric.higher_is_better:
+        worst = sorted(live, key=get_val)[:top_n]
+    else:
+        worst = sorted(live, key=get_val, reverse=True)[:top_n]
+    print(f"\n  Worst offenders ({metric.name}):")
+    for i, m in enumerate(worst, 1):
+        v = get_val(m)
+        parts = [f"{v:.3f}{metric.unit}"]
+        if m.config:
+            parts.append(m.short_config)
+        detail = "  ".join(parts)
+        print(f"    {i}. {m.name:<30}  {detail}")
+
+
+# ---------------------------------------------------------------------------
+# Subcommands
+# ---------------------------------------------------------------------------
+
+
+def build_dispatch_inputs(args) -> list[str]:
+    """Build -f flags for workflow dispatch inputs from CLI args."""
+    flags = []
+    bool_inputs = [
+        "training",
+        "inference",
+        "default",
+        "dynamic",
+        "cppwrapper",
+        "cudagraphs",
+        "freezing_cudagraphs",
+        "aotinductor",
+        "maxautotune",
+    ]
+    for name in bool_inputs:
+        val = getattr(args, name, None)
+        if val is not None:
+            flags.extend(["-f", f"{name}={'true' if val else 'false'}"])
+    if args.benchmark_configs:
+        flags.extend(["-f", f"benchmark_configs={args.benchmark_configs}"])
+    return flags
+
+
+def dispatch_one(device: str, ref: str, extra_flags: list[str]) -> int | None:
+    wf = WORKFLOWS[device]
+    print(f"\nLaunching {wf['name']} on ref: {ref}")
+
+    dispatch_args = [
+        "workflow",
+        "run",
+        str(wf["id"]),
+        "--repo",
+        REPO,
+        "--ref",
+        ref,
+    ] + extra_flags
+
+    gh(*dispatch_args)
+    print("Dispatched. Waiting a few seconds for the run to appear...")
+    time.sleep(5)
+
+    runs = gh(
+        "run",
+        "list",
+        "--repo",
+        REPO,
+        "--workflow",
+        str(wf["id"]),
+        "--branch",
+        ref,
+        "--limit",
+        "1",
+        "--json",
+        "databaseId,status,url,createdAt",
+        json_output=True,
+    )
+    if not runs:
+        print("Could not find the dispatched run. Check the Actions tab manually.")
+        return None
+
+    run = runs[0]
+    run_id = run["databaseId"]
+    url = f"https://github.com/{REPO}/actions/runs/{run_id}"
+    print(f"Run ID:  {run_id}")
+    print(f"URL:     {url}")
+    print(f"Status:  {run['status']}")
+    return run_id
+
+
+def wait_for_runs(pending: dict[str, int]) -> dict[str, int]:
+    """Poll all runs until they complete. Returns dict of successful runs."""
+    print(f"\nWaiting for {len(pending)} run(s)...", flush=True)
+    remaining = dict(pending)
+    succeeded: dict[str, int] = {}
+    while remaining:
+        time.sleep(30)
+        done = []
+        for device, run_id in remaining.items():
+            data = gh(
+                "run",
+                "view",
+                str(run_id),
+                "--repo",
+                REPO,
+                "--json",
+                "status,conclusion",
+                json_output=True,
+            )
+            status = data.get("status", "unknown")
+            if status == "completed":
+                conclusion = data.get("conclusion", "unknown")
+                print(f"  {device} (run {run_id}): {conclusion}")
+                if conclusion == "success":
+                    succeeded[device] = run_id
+                done.append(device)
+        for d in done:
+            del remaining[d]
+        if remaining:
+            ts = datetime.now().strftime("%H:%M:%S")
+            print(f"  [{ts}] still waiting: {', '.join(remaining)}...", flush=True)
+    return succeeded
+
+
+def cmd_launch(args):
+    ref = args.ref or git("rev-parse", "--abbrev-ref", "HEAD")
+    if ref == "HEAD":
+        ref = git("rev-parse", "HEAD")
+
+    extra_flags = build_dispatch_inputs(args)
+    launched: dict[str, int] = {}
+    for device in args.device:
+        if device not in WORKFLOWS:
+            print(f"Unknown device: {device}", file=sys.stderr)
+            sys.exit(1)
+        run_id = dispatch_one(device, ref, extra_flags)
+        if run_id:
+            launched[device] = run_id
+
+    if (args.wait or args.wait_and_summarize) and launched:
+        succeeded = wait_for_runs(launched)
+        if args.wait_and_summarize and succeeded:
+            # Use the first device's run ID as the positional arg; pass all
+            # device→run_id pairs via _run_ids so cmd_summary skips resolution.
+            first_run = next(iter(succeeded.values()))
+            summary_args = argparse.Namespace(
+                run_id=str(first_run),
+                device=list(succeeded.keys()),
+                _run_ids=succeeded,
+                baseline="latest",
+                metric="speedup",
+                top=5,
+                config=None,
+                suite=None,
+                mode=None,
+                group_by=None,
+                attempt=1,
+                no_cache=False,
+            )
+            print(f"\n{'=' * 70}")
+            print("Summary")
+            print(f"{'=' * 70}")
+            cmd_summary(summary_args)
+
+
+def filter_perf(all_perf: list[PerfData], args) -> list[PerfData]:
+    result = all_perf
+    if getattr(args, "config", None):
+        pattern = re.compile(args.config, re.IGNORECASE)
+        result = [p for p in result if pattern.search(p.config)]
+    if getattr(args, "suite", None):
+        suite = SUITE_ALIASES.get(args.suite, args.suite)
+        result = [p for p in result if p.suite == suite]
+    if getattr(args, "mode", None):
+        result = [p for p in result if p.mode == args.mode]
+    if getattr(args, "backend", None):
+        pattern = re.compile(args.backend, re.IGNORECASE)
+        result = [p for p in result if pattern.search(p.short_name)]
+    if getattr(args, "dtype", None):
+        result = [p for p in result if p.dtype == args.dtype]
+    if getattr(args, "runtime", None):
+        result = [p for p in result if p.runtime == args.runtime]
+    return result
+
+
+GROUP_KEY_FNS: dict[str, callable] = {
+    "config": lambda p: p.qualified_config,
+    "suite": lambda p: p.suite,
+    "mode": lambda p: p.mode,
+    "backend": lambda p: p.short_name,
+    "device": lambda p: p.device or "unknown",
+    "dtype": lambda p: p.dtype,
+    "runtime": lambda p: p.runtime,
+}
+
+GROUP_CHOICES = sorted(GROUP_KEY_FNS.keys())
+
+
+def group_perf(all_perf: list[PerfData], group_by: list[str] | None) -> list[PerfData]:
+    if not group_by:
+        all_models = []
+        for p in all_perf:
+            all_models.extend(p.models)
+        return [PerfData(config="all", models=all_models)]
+
+    fns = []
+    for key in group_by:
+        fn = GROUP_KEY_FNS.get(key)
+        if fn is None:
+            print(f"Unknown group-by: {key}", file=sys.stderr)
+            sys.exit(1)
+        fns.append(fn)
+
+    def composite_key(p: PerfData) -> str:
+        return " / ".join(fn(p) for fn in fns)
+
+    groups: dict[str, list[ModelResult]] = defaultdict(list)
+    for p in all_perf:
+        groups[composite_key(p)].extend(p.models)
+
+    return [PerfData(config=k, models=v) for k, v in sorted(groups.items())]
+
+
+# ---------------------------------------------------------------------------
+# Baseline comparison
+# ---------------------------------------------------------------------------
+
+
+NIGHTLY_WORKFLOW_IDS = {k: v["id"] for k, v in WORKFLOWS.items() if k != "a100-compare"}
+
+
+def _find_latest_run(
+    branch: str,
+    device: str,
+) -> dict | None:
+    """Find the latest successful perf nightly run for a branch + device.
+
+    Returns {databaseId, createdAt, headSha} or None.
+    """
+    wf_id = NIGHTLY_WORKFLOW_IDS.get(device)
+    if wf_id is None:
+        return None
+
+    runs = gh(
+        "run",
+        "list",
+        "--repo",
+        REPO,
+        "--workflow",
+        str(wf_id),
+        "--branch",
+        branch,
+        "--status",
+        "success",
+        "--limit",
+        "1",
+        "--json",
+        "databaseId,createdAt,headSha",
+        json_output=True,
+    )
+    if not runs:
+        return None
+    return runs[0]
+
+
+def resolve_run(branch: str, device: str) -> int:
+    """Find the latest successful perf nightly run for a branch + device."""
+    run = _find_latest_run(branch, device)
+    if not run:
+        print(
+            f"No successful perf run found for branch '{branch}' on {device}.",
+            file=sys.stderr,
+        )
+        sys.exit(1)
+    run_id = run["databaseId"]
+    created = run["createdAt"][:10]
+    print(f"Resolved '{branch}' → run {run_id} ({device}, {created})")
+    return run_id
+
+
+def discover_runs(branch: str) -> dict[str, int]:
+    """Auto-discover all devices with successful runs for a branch.
+
+    Finds the latest commit SHA that has runs, then returns all runs matching
+    that commit.
+    """
+    candidates: list[tuple[str, dict]] = []
+    for device in DEVICE_CHOICES:
+        run = _find_latest_run(branch, device)
+        if run:
+            candidates.append((device, run))
+
+    if not candidates:
+        print(
+            f"No successful perf runs found for branch '{branch}' on any device.",
+            file=sys.stderr,
+        )
+        sys.exit(1)
+
+    # Pick the most recent commit (by createdAt) and collect all runs on it
+    candidates.sort(key=lambda x: x[1]["createdAt"], reverse=True)
+    target_sha = candidates[0][1]["headSha"]
+
+    result = {}
+    for device, run in candidates:
+        if run["headSha"] == target_sha:
+            run_id = run["databaseId"]
+            created = run["createdAt"][:10]
+            print(
+                f"Discovered '{branch}' → run {run_id} ({device}, {created}, {target_sha[:10]})"
+            )
+            result[device] = run_id
+
+    return result
+
+
+def resolve_runs(branch: str, devices: list[str]) -> dict[str, int]:
+    """Resolve the latest successful run for each device. Returns {device: run_id}."""
+    result = {}
+    for device in devices:
+        result[device] = resolve_run(branch, device)
+    return result
+
+
+def resolve_baseline(head_run_id: int) -> int:
+    """Find the latest successful nightly on main, skipping the head run itself."""
+    run_data = gh(
+        "run",
+        "view",
+        str(head_run_id),
+        "--repo",
+        REPO,
+        "--json",
+        "workflowName",
+        json_output=True,
+    )
+    wf_name = run_data.get("workflowName", "")
+    nightly_id = WORKFLOW_NAME_TO_NIGHTLY_ID.get(wf_name)
+    if nightly_id is None:
+        print(
+            f"Don't know which nightly corresponds to workflow '{wf_name}'",
+            file=sys.stderr,
+        )
+        sys.exit(1)
+
+    runs = gh(
+        "run",
+        "list",
+        "--repo",
+        REPO,
+        "--workflow",
+        str(nightly_id),
+        "--branch",
+        "main",
+        "--status",
+        "success",
+        "--limit",
+        "5",
+        "--json",
+        "databaseId,createdAt,headBranch",
+        json_output=True,
+    )
+    # Skip the head run itself to avoid comparing a run to itself
+    for run in runs:
+        if run["databaseId"] != head_run_id:
+            baseline_id = run["databaseId"]
+            created = run["createdAt"][:10]
+            print(f"Baseline: run {baseline_id} (main, {created})")
+            return baseline_id
+
+    print("No suitable baseline nightly found on main.", file=sys.stderr)
+    sys.exit(1)
+
+
+def device_for_workflow(workflow_name: str) -> str:
+    """Reverse-lookup device key from workflow name."""
+    for k, v in WORKFLOWS.items():
+        if v["name"] == workflow_name:
+            return k
+    return ""
+
+
+@dataclass
+class RunMeta:
+    run_id: int
+    head_sha: str
+    head_branch: str
+    workflow_name: str
+    created_at: str
+    event: str
+
+    @property
+    def short_sha(self) -> str:
+        return self.head_sha[:10]
+
+    @property
+    def date(self) -> str:
+        return self.created_at[:10]
+
+
+def fetch_run_meta(run_id: int) -> RunMeta:
+    data = gh(
+        "api",
+        f"repos/{REPO}/actions/runs/{run_id}",
+        "-q",
+        "{headSha: .head_sha, headBranch: .head_branch, workflowName: .name, createdAt: .created_at, event: .event}",
+        json_output=True,
+    )
+    return RunMeta(
+        run_id=run_id,
+        head_sha=data.get("headSha", "unknown"),
+        head_branch=data.get("headBranch", "unknown"),
+        workflow_name=data.get("workflowName", "unknown"),
+        created_at=data.get("createdAt", "unknown"),
+        event=data.get("event", "unknown"),
+    )
+
+
+def print_run_header(
+    label: str,
+    metas: list[RunMeta],
+    configs: list[str] | None = None,
+):
+    print(f"\n  {label}")
+    print(f"  {'─' * len(label)}")
+    if len(metas) == 1:
+        m = metas[0]
+        print(f"  Run:      {m.run_id}  ({m.workflow_name})")
+        print(f"  Commit:   {m.short_sha}  ({m.head_branch}, {m.date})")
+    else:
+        # Show commit from first (should all match for multi-device)
+        print(
+            f"  Commit:   {metas[0].short_sha}  ({metas[0].head_branch}, {metas[0].date})"
+        )
+        print("  Runs:")
+        for m in metas:
+            print(f"    {m.run_id}  ({m.workflow_name})")
+    if configs:
+        print(f"  Configs:  {len(configs)} — {', '.join(sorted(configs))}")
+
+
+def fetch_run_perf(
+    run_id: int,
+    attempt: int,
+    no_cache: bool,
+    device: str = "",
+    allow_empty: bool = False,
+) -> list[PerfData]:
+    """Fetch and parse perf data for a run."""
+    jobs = get_run_jobs(run_id)
+    perf_jobs = get_perf_jobs(jobs)
+    if not perf_jobs:
+        if allow_empty:
+            print(f"  no perf jobs in run {run_id}, skipping", file=sys.stderr)
+            return []
+        print(f"No successful perf jobs in run {run_id}.", file=sys.stderr)
+        sys.exit(1)
+    csv_pairs = download_and_extract_csvs(run_id, perf_jobs, attempt, no_cache=no_cache)
+    if not csv_pairs:
+        if allow_empty:
+            print(f"  no CSVs in run {run_id}, skipping", file=sys.stderr)
+            return []
+        print(f"No performance CSVs in run {run_id}.", file=sys.stderr)
+        sys.exit(1)
+    return parse_csvs(csv_pairs, device=device)
+
+
+@dataclass
+class ModelDelta:
+    name: str
+    base_val: float
+    head_val: float
+    config: str = ""
+    device: str = ""
+
+    @property
+    def delta_pct(self) -> float:
+        if self.base_val == 0:
+            return 0.0
+        return (self.head_val - self.base_val) / self.base_val * 100
+
+    @property
+    def short_config(self) -> str:
+        return _short_config(self.config, self.device)
+
+
+@dataclass
+class ConfigAgg:
+    base_agg: float
+    base_count: int
+    head_agg: float
+    head_count: int
+    paired_ratio: float  # gmean(head_val / base_val) over paired models
+    paired_count: int
+
+
+def compute_deltas(
+    head_perf: list[PerfData], base_perf: list[PerfData], metric: Metric
+) -> tuple[list[ModelDelta], dict[str, ConfigAgg]]:
+    """Join head and base on (device, config, model_name) and compute deltas.
+
+    Returns (per_model_deltas, per_config_aggregates).
+    per_config_aggregates maps qualified_config -> ConfigAgg.
+    """
+    # Build base lookup: (device, config, model_name) -> metric value
+    base_lookup: dict[tuple[str, str, str], float] = {}
+    for perf in base_perf:
+        for m in perf.models:
+            v = getattr(m, metric.field)
+            if v > 0:
+                base_lookup[(perf.device, perf.config, m.name)] = v
+
+    deltas = []
+    for perf in head_perf:
+        for m in perf.models:
+            head_val = getattr(m, metric.field)
+            if head_val <= 0:
+                continue
+            key = (perf.device, perf.config, m.name)
+            if key not in base_lookup:
+                continue
+            base_val = base_lookup[key]
+            deltas.append(
+                ModelDelta(
+                    name=m.name,
+                    base_val=base_val,
+                    head_val=head_val,
+                    config=perf.config,
+                    device=perf.device,
+                )
+            )
+
+    # Group deltas by qualified_config for paired aggregates
+    deltas_by_qconfig: dict[str, list[ModelDelta]] = defaultdict(list)
+    for d in deltas:
+        qc = f"{d.device}/{d.config}" if d.device else d.config
+        deltas_by_qconfig[qc].append(d)
+
+    # Per-config aggregates (keyed by qualified_config for display)
+    config_aggs: dict[str, ConfigAgg] = {}
+    base_by_qconfig: dict[str, PerfData] = {p.qualified_config: p for p in base_perf}
+    for perf in head_perf:
+        qc = perf.qualified_config
+        if qc not in base_by_qconfig:
+            continue
+        base_perf_data = base_by_qconfig[qc]
+        head_agg = perf.aggregate_metric(metric)
+        base_agg = base_perf_data.aggregate_metric(metric)
+        head_count = len([m for m in perf.models if getattr(m, metric.field) > 0])
+        base_count = len(
+            [m for m in base_perf_data.models if getattr(m, metric.field) > 0]
+        )
+
+        paired = deltas_by_qconfig.get(qc, [])
+        ratios = [d.head_val / d.base_val for d in paired if d.base_val > 0]
+        config_aggs[qc] = ConfigAgg(
+            base_agg=base_agg,
+            base_count=base_count,
+            head_agg=head_agg,
+            head_count=head_count,
+            paired_ratio=gmean(ratios) if ratios else 0.0,
+            paired_count=len(ratios),
+        )
+
+    return deltas, config_aggs
+
+
+def print_comparison_table(config_aggs: dict[str, ConfigAgg], metric: Metric):
+    u = metric.unit
+    print(f"\n{'Config':<55} {'base':>16} {'new':>16} {'head/base':>18}")
+    print("─" * 108)
+    for config in sorted(config_aggs):
+        agg = config_aggs[config]
+        flag = ""
+        if agg.paired_ratio > 0:
+            delta_pct = (agg.paired_ratio - 1.0) * 100
+            if abs(delta_pct) > RELATIVE_THRESHOLD * 100:
+                if metric.higher_is_better:
+                    flag = " !!" if delta_pct < 0 else " ++"
+                else:
+                    flag = " !!" if delta_pct > 0 else " ++"
+        print(
+            f"  {config:<53} "
+            f"{agg.base_agg:>5.2f}{u} (n={agg.base_count}) "
+            f"{agg.head_agg:>5.2f}{u} (n={agg.head_count}) "
+            f"{agg.paired_ratio:>5.3f}x (n={agg.paired_count}){flag}"
+        )
+
+
+def print_regressions(deltas: list[ModelDelta], metric: Metric, top_n: int):
+    # For higher_is_better metrics, regression = negative delta
+    # For lower_is_better metrics, regression = positive delta
+    if metric.higher_is_better:
+        bad = [d for d in deltas if d.delta_pct < -RELATIVE_THRESHOLD * 100]
+        bad.sort(key=lambda d: d.delta_pct)
+    else:
+        bad = [d for d in deltas if d.delta_pct > RELATIVE_THRESHOLD * 100]
+        bad.sort(key=lambda d: d.delta_pct, reverse=True)
+
+    if not bad:
+        print(f"\n  No regressions (>{RELATIVE_THRESHOLD * 100:.0f}% change).")
+        return
+
+    print(f"\n  Regressions ({len(bad)} models, showing top {min(top_n, len(bad))}):")
+    for i, d in enumerate(bad[:top_n], 1):
+        print(
+            f"    {i}. {d.name:<30}  "
+            f"{d.base_val:.2f}{metric.unit} → {d.head_val:.2f}{metric.unit}  "
+            f"{d.delta_pct:>+6.1f}%  {d.short_config}"
+        )
+
+
+def render_delta_scurve(
+    deltas: list[ModelDelta],
+    metric: Metric,
+    term_width: int | None = None,
+    term_height: int | None = None,
+):
+    if not deltas:
+        return
+
+    if term_width is None or term_height is None:
+        sz = shutil.get_terminal_size((100, 50))
+        term_width = term_width or sz.columns
+        term_height = term_height or sz.lines
+
+    sorted_deltas = sorted(deltas, key=lambda d: d.delta_pct)
+    n = len(sorted_deltas)
+    max_rows = max(term_height - 5, 15)
+    display = subsample(sorted_deltas, max_rows)
+    skipped = n - len(display)
+
+    header = f"Delta S-curve ({n} models)"
+    if skipped > 0:
+        header += f" [showing {len(display)}/{n}]"
+    print(f"\n  {header}")
+    print(f"  {'─' * min(len(header), term_width - 4)}")
+
+    max_name = min(max(len(d.name) for d in display), 28)
+    # "  name  +12.3%  {bar}"
+    prefix_len = 2 + max_name + 2 + 7 + 2
+    plot_width = max(term_width - prefix_len - 1, 20)
+
+    # Range: cap at p5/p95 to avoid outlier squishing
+    p5_idx = max(0, int(n * 0.05))
+    p95_idx = min(n - 1, int(n * 0.95))
+    range_lo = min(sorted_deltas[p5_idx].delta_pct, -10)
+    range_hi = max(sorted_deltas[p95_idx].delta_pct, 10)
+    # Ensure symmetric-ish around 0
+    abs_max = max(abs(range_lo), abs(range_hi))
+    range_lo = -abs_max
+    range_hi = abs_max
+    span = range_hi - range_lo
+    if span == 0:
+        span = 1
+
+    def pct_to_col(pct: float) -> int:
+        return max(
+            0, min(plot_width - 1, int((pct - range_lo) / span * (plot_width - 1)))
+        )
+
+    zero_col = pct_to_col(0)
+
+    for d in display:
+        name = d.name[:max_name].ljust(max_name)
+        col = pct_to_col(d.delta_pct)
+        bar = [" "] * plot_width
+        bar[zero_col] = "|"
+        if col <= zero_col:
+            for i in range(col, zero_col):
+                bar[i] = "·"
+            bar[zero_col] = "|"
+        else:
+            bar[zero_col] = "|"
+            for i in range(zero_col + 1, col + 1):
+                bar[i] = "·"
+        print(f"  {name}  {d.delta_pct:>+6.1f}%  {''.join(bar)}")
+
+    pad = " " * prefix_len
+    print(f"{pad}{' ' * zero_col}0%")
+
+
+def print_summary_table(all_perf: list[PerfData], metric: Metric):
+    agg_label = metric.aggregate
+    print(f"\n{'Config':<65} {'Models':>6} {agg_label:>10}")
+    print("─" * 85)
+    for perf in all_perf:
+        n = len([m for m in perf.models if getattr(m, metric.field) > 0])
+        agg = perf.aggregate_metric(metric)
+        print(f"  {perf.qualified_config:<63} {n:>6} {agg:>8.2f}{metric.unit}")
+
+
+def _resolve_head_runs(args) -> dict[str, int]:
+    """Parse run_id arg into {device: run_id} mapping.
+
+    run_id can be:
+      - A single numeric run ID (device inferred from workflow)
+      - A branch name (resolved across all --device values, or auto-discovered)
+    Pre-resolved IDs can be passed via args._run_ids.
+    """
+    # Pre-resolved (from --wait-and-summarize)
+    if hasattr(args, "_run_ids") and args._run_ids:
+        return args._run_ids
+
+    raw = args.run_id
+    try:
+        run_id = int(raw)
+        # Single run ID — infer device from the workflow
+        meta = fetch_run_meta(run_id)
+        device = device_for_workflow(meta.workflow_name)
+        return {device: run_id}
+    except ValueError:
+        # Branch name
+        if args.device:
+            return resolve_runs(raw, args.device)
+        return discover_runs(raw)
+
+
+def cmd_summary(args):
+    attempt = args.attempt
+    metric = METRICS[args.metric]
+
+    head_run_ids = _resolve_head_runs(args)
+
+    # Fetch head runs
+    auto_discovered = not hasattr(args, "_run_ids") and not args.device
+    head_metas: list[RunMeta] = []
+    head_perf: list[PerfData] = []
+    for device, run_id in list(head_run_ids.items()):
+        print(f"Fetching head run {run_id} ({device})...")
+        perf = fetch_run_perf(
+            run_id,
+            attempt,
+            args.no_cache,
+            device=device,
+            allow_empty=auto_discovered,
+        )
+        if not perf:
+            del head_run_ids[device]
+            continue
+        head_metas.append(fetch_run_meta(run_id))
+        head_perf.extend(perf)
+
+    head_perf = filter_perf(head_perf, args)
+    if not head_perf:
+        print("No configs matched filters.")
+        sys.exit(1)
+
+    head_configs = [p.qualified_config for p in head_perf]
+
+    # Baseline comparison mode
+    if args.baseline and args.baseline.lower() != "none":
+        base_metas: list[RunMeta] = []
+        base_perf: list[PerfData] = []
+
+        for device, head_run_id in head_run_ids.items():
+            if args.baseline == "latest":
+                baseline_id = resolve_baseline(head_run_id)
+            else:
+                try:
+                    baseline_id = int(args.baseline)
+                except ValueError:
+                    baseline_id = resolve_run(args.baseline, device)
+
+            print(f"Fetching baseline run {baseline_id} ({device})...")
+            base_data = fetch_run_perf(
+                baseline_id,
+                attempt,
+                args.no_cache,
+                device=device,
+                allow_empty=auto_discovered,
+            )
+            if not base_data:
+                continue
+            base_metas.append(fetch_run_meta(baseline_id))
+            base_perf.extend(base_data)
+
+        base_perf = filter_perf(base_perf, args)
+        if not base_perf:
+            print("No baseline configs matched filters.")
+            sys.exit(1)
+
+        print_run_header("HEAD", head_metas, head_configs)
+        print_run_header("BASE", base_metas, [p.qualified_config for p in base_perf])
+        print()
+
+        deltas, config_aggs = compute_deltas(head_perf, base_perf, metric)
+        if not deltas:
+            print("No matching models between head and baseline.")
+            sys.exit(1)
+
+        print_comparison_table(config_aggs, metric)
+        print_regressions(deltas, metric, args.top)
+        render_delta_scurve(deltas, metric)
+        return
+
+    # Absolute mode (no baseline)
+    print_run_header("RUN", head_metas, head_configs)
+    print()
+
+    print_summary_table(head_perf, metric)
+    grouped = group_perf(head_perf, args.group_by)
+    for perf in grouped:
+        print_worst_offenders(perf, metric, args.top)
+        render_scurve(perf, metric, args.top)
+
+
+CONFIG_RE = re.compile(
+    r"(?P<backend>inductor_[a-z_]+?)_"
+    r"(?P<suite>huggingface|timm_models|torchbench)_"
+    r"(?P<dtype>\w+)_"
+    r"(?P<mode>training|inference)_"
+    r"(?P<device>\w+)"
+)
+
+
+def config_to_command(
+    config: str,
+    suite: str,
+    model: str | None = None,
+) -> str | None:
+    """Turn a config name into a runnable benchmark command."""
+    m = CONFIG_RE.match(config)
+    if not m:
+        return None
+
+    backend_variant = m.group("backend")
+    dtype = m.group("dtype")
+    mode = m.group("mode")
+    runtime = m.group("device")
+
+    # Runtime → --device flag (strip platform suffix like _x86_zen)
+    device_flag = runtime.split("_")[0]  # "cpu_x86_zen" → "cpu"
+
+    cmd_parts = [
+        "python",
+        f"benchmarks/dynamo/{suite}.py",
+        f"--{mode}",
+        f"--{dtype}",
+        "--backend",
+        "inductor",
+        "--device",
+        device_flag,
+    ]
+
+    if "no_cudagraphs" in backend_variant:
+        cmd_parts.append("--disable-cudagraphs")
+    if "dynamic" in backend_variant:
+        cmd_parts.extend(["--dynamic-shapes", "--dynamic-batch-only"])
+    if "cpp_wrapper" in backend_variant:
+        cmd_parts.insert(0, "TORCHINDUCTOR_CPP_WRAPPER=1")
+        cmd_parts.append("--disable-cudagraphs")
+    if "freezing" in backend_variant:
+        cmd_parts.append("--freezing")
+    if "max_autotune" in backend_variant:
+        cmd_parts.insert(0, "TORCHINDUCTOR_MAX_AUTOTUNE=1")
+    if "aot_inductor" in backend_variant:
+        cmd_parts.append("--export-aot-inductor")
+        cmd_parts.append("--disable-cudagraphs")
+
+    cmd_parts.extend(["--performance", "--cold-start-latency"])
+
+    if model:
+        cmd_parts.extend(["--only", model])
+
+    cmd_parts.extend(["--output", f"{config}_performance.csv"])
+
+    return " ".join(cmd_parts)
+
+
+def cmd_repro(args):
+    run_ids = _resolve_head_runs(args)
+
+    # Fetch perf data to discover configs
+    auto_discovered = not hasattr(args, "_run_ids") and not args.device
+    all_perf: list[PerfData] = []
+    metas: list[RunMeta] = []
+    for device, run_id in list(run_ids.items()):
+        print(f"Fetching run {run_id} ({device})...")
+        metas.append(fetch_run_meta(run_id))
+        perf = fetch_run_perf(
+            run_id,
+            args.attempt,
+            no_cache=False,
+            device=device,
+            allow_empty=auto_discovered,
+        )
+        all_perf.extend(perf)
+
+    all_perf = filter_perf(all_perf, args)
+    if not all_perf:
+        print("No configs matched filters.")
+        sys.exit(1)
+
+    print_run_header("REPRO", metas)
+
+    configs_seen: dict[str, str] = {}  # config_name → suite
+    for perf in all_perf:
+        configs_seen[perf.config] = perf.suite
+
+    count = 0
+    commands: list[str] = []
+    for config in sorted(configs_seen):
+        suite = configs_seen[config]
+        cmd = config_to_command(config, suite, args.model)
+        if not cmd:
+            continue
+        commands.append(f"# {config}\n{cmd}")
+        count += 1
+
+    print(f"\nReproducible commands ({count} configs):\n")
+    for cmd in commands:
+        print(f"{cmd}\n")
+
+
+PIN_DIR = Path(".ci/docker/ci_commit_pins")
+
+
+def read_pin(name: str) -> str:
+    """Read a pinned commit or requirements file."""
+    path = PIN_DIR / name
+    if not path.exists():
+        return f"<{name} not found>"
+    return path.read_text().strip()
+
+
+def cmd_prepare_repro(args):
+    suites = set()
+    if args.suite:
+        suite = SUITE_ALIASES.get(args.suite, args.suite)
+        suites = {suite}
+    else:
+        suites = {"huggingface", "timm_models", "torchbench"}
+
+    torchbench_pin = read_pin("torchbench.txt")
+    timm_pin = read_pin("timm.txt")
+    hf_reqs = read_pin("huggingface-requirements.txt")
+
+    print("# Setup commands for inductor perf benchmark suites")
+    print("# These mirror what CI does in the inductor-benchmarks Docker image.")
+    print("#")
+    print("# Pinned versions (commits, package versions) are read live from")
+    print("#   .ci/docker/ci_commit_pins/")
+    print("# Install steps are based on:")
+    print("#   .ci/docker/common/install_inductor_benchmark_deps.sh  (build-time)")
+    print("#   .ci/pytorch/test.sh                                   (runtime)")
+    print("# If the setup process changes, check those files.")
+    print()
+
+    if "huggingface" in suites:
+        print("# ── HuggingFace ──")
+        for line in hf_reqs.splitlines():
+            line = line.strip()
+            if line and not line.startswith("#"):
+                print(f"pip install {line}")
+        print()
+
+    if "timm_models" in suites:
+        print("# ── Timm ──")
+        print(
+            f"pip install git+https://github.com/huggingface/pytorch-image-models@{timm_pin}"
+        )
+        print()
+
+    if "torchbench" in suites:
+        print("# ── TorchBench ──")
+        print("git clone https://github.com/pytorch/benchmark torchbench")
+        print(f"cd torchbench && git checkout {torchbench_pin}")
+        print("python install.py --continue_on_fail")
+        print("cd ..")
+        print()
+        print("# Set PYTHONPATH so benchmark scripts find torchbench")
+        print("export PYTHONPATH=$(pwd)/torchbench")
+        print()
+
+    print("# ── Runtime dependencies ──")
+    print("pip install torchvision torchaudio")
+    if "torchbench" in suites:
+        print("pip install opencv-python==4.8.0.74")
+    print()
+
+    print("# ── Environment variables ──")
+    print("export TORCHINDUCTOR_FX_GRAPH_CACHE=True")
+    print("export TORCHINDUCTOR_AUTOGRAD_CACHE=True")
+    print()
+
+    if not args.no_repro:
+        # Also print repro commands if we have a run_id
+        if args.run_id:
+            # Reuse the repro logic
+            repro_args = argparse.Namespace(
+                run_id=args.run_id,
+                device=args.device,
+                model=args.model,
+                suite=args.suite,
+                mode=args.mode,
+                backend=None,
+                dtype=None,
+                runtime=None,
+                attempt=args.attempt,
+            )
+            print("# ── Benchmark commands ──")
+            cmd_repro(repro_args)
+
+
+# ---------------------------------------------------------------------------
+# Main
+# ---------------------------------------------------------------------------
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="CLI tool for inductor perf regression runs",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        epilog="""
+Examples:
+  # Launch a perf run on your current branch (A100)
+  python benchmarks/dynamo/perf_cli.py launch
+
+  # Launch on both A100 and H100
+  python benchmarks/dynamo/perf_cli.py launch --device a100 h100
+
+  # Launch on H100 and wait for completion
+  python benchmarks/dynamo/perf_cli.py launch --device h100 --wait
+
+  # Launch with only inference, no training
+  python benchmarks/dynamo/perf_cli.py launch --no-training --inference
+
+  # Launch with dynamic shapes enabled
+  python benchmarks/dynamo/perf_cli.py launch --dynamic
+
+  # Launch on ROCm MI300
+  python benchmarks/dynamo/perf_cli.py launch --device rocm-mi300
+
+  # Check your branch against latest main nightly
+  python benchmarks/dynamo/perf_cli.py summary my-feature-branch
+
+  # Same but on H100
+  python benchmarks/dynamo/perf_cli.py summary my-feature-branch --device h100
+
+  # Compare across A100 and H100 in one summary
+  python benchmarks/dynamo/perf_cli.py summary my-feature-branch --device a100 h100
+
+  # Use a specific run ID instead of branch name
+  python benchmarks/dynamo/perf_cli.py summary 22842783236
+
+  # Compare against a specific baseline (run ID or branch)
+  python benchmarks/dynamo/perf_cli.py summary 22842783236 --baseline 22816292132
+
+  # Absolute metrics only (no comparison)
+  python benchmarks/dynamo/perf_cli.py summary 22842783236 --baseline none
+
+  # Filter to cudagraphs training
+  python benchmarks/dynamo/perf_cli.py summary main --config cudagraphs --mode training
+
+  # Show commands to reproduce a run locally for a single model
+  python benchmarks/dynamo/perf_cli.py repro 22842783236 --model BERT_pytorch --suite tb
+""",
+    )
+    sub = parser.add_subparsers(dest="command", required=True)
+
+    # -- launch --
+    p_launch = sub.add_parser(
+        "launch", help="Launch a perf regression run on your branch"
+    )
+    p_launch.add_argument(
+        "--device",
+        nargs="+",
+        default=["a100"],
+        choices=DEVICE_CHOICES,
+        metavar="DEVICE",
+        help=f"Devices to launch (default: a100). Choices: {', '.join(DEVICE_CHOICES)}",
+    )
+    p_launch.add_argument(
+        "--ref",
+        type=str,
+        default=None,
+        help="Git ref to benchmark (default: current branch)",
+    )
+    p_launch.add_argument(
+        "--wait", action="store_true", help="Wait for all runs to complete"
+    )
+    p_launch.add_argument(
+        "--wait-and-summarize",
+        dest="wait_and_summarize",
+        action="store_true",
+        help="Wait for all runs, then print summary vs. latest main nightly",
+    )
+    # Workflow dispatch inputs
+    launch_opts = p_launch.add_argument_group("workflow options")
+    launch_opts.add_argument(
+        "--training",
+        action="store_true",
+        default=None,
+        help="Enable training benchmarks",
+    )
+    launch_opts.add_argument(
+        "--no-training",
+        dest="training",
+        action="store_false",
+        help="Disable training benchmarks",
+    )
+    launch_opts.add_argument(
+        "--inference",
+        action="store_true",
+        default=None,
+        help="Enable inference benchmarks",
+    )
+    launch_opts.add_argument(
+        "--no-inference",
+        dest="inference",
+        action="store_false",
+        help="Disable inference benchmarks",
+    )
+    launch_opts.add_argument(
+        "--cudagraphs", action="store_true", default=None, help="Enable cudagraphs"
+    )
+    launch_opts.add_argument(
+        "--no-cudagraphs",
+        dest="cudagraphs",
+        action="store_false",
+        help="Disable cudagraphs",
+    )
+    launch_opts.add_argument(
+        "--dynamic", action="store_true", default=None, help="Enable dynamic shapes"
+    )
+    launch_opts.add_argument(
+        "--no-dynamic",
+        dest="dynamic",
+        action="store_false",
+        help="Disable dynamic shapes",
+    )
+    launch_opts.add_argument(
+        "--cppwrapper", action="store_true", default=None, help="Enable cpp wrapper"
+    )
+    launch_opts.add_argument(
+        "--freezing-cudagraphs",
+        dest="freezing_cudagraphs",
+        action="store_true",
+        default=None,
+    )
+    launch_opts.add_argument("--aotinductor", action="store_true", default=None)
+    launch_opts.add_argument("--maxautotune", action="store_true", default=None)
+    launch_opts.add_argument(
+        "--default", dest="default", action="store_true", default=None
+    )
+    launch_opts.add_argument(
+        "--benchmark-configs",
+        dest="benchmark_configs",
+        type=str,
+        default=None,
+        help="Override benchmark_configs input",
+    )
+
+    # -- summary --
+    p_summary = sub.add_parser("summary", help="Summarize results of a perf run")
+    p_summary.add_argument(
+        "run_id", type=str, help="GitHub Actions run ID or branch name"
+    )
+    p_summary.add_argument(
+        "--device",
+        nargs="+",
+        default=None,
+        choices=DEVICE_CHOICES,
+        metavar="DEVICE",
+        help=f"Device(s) to summarize (default: auto-discover from branch). Choices: {', '.join(DEVICE_CHOICES)}",
+    )
+    p_summary.add_argument(
+        "--baseline",
+        type=str,
+        default="latest",
+        help="Baseline run ID, 'latest' for most recent main nightly (default), or 'none' to disable",
+    )
+    p_summary.add_argument(
+        "--metric",
+        type=str,
+        default="speedup",
+        choices=METRIC_CHOICES,
+        help="Metric to display (default: speedup)",
+    )
+    p_summary.add_argument(
+        "--top",
+        type=int,
+        default=5,
+        help="Number of worst offenders to show (default: 5)",
+    )
+    # Filters
+    filters = p_summary.add_argument_group("filters")
+    filters.add_argument(
+        "--config",
+        type=str,
+        default=None,
+        help="Regex to filter config names (e.g. 'cudagraphs', 'dynamic')",
+    )
+    filters.add_argument(
+        "--suite", type=str, default=None, help="Filter to suite: hf, timm, tb"
+    )
+    filters.add_argument(
+        "--mode",
+        type=str,
+        choices=["training", "inference"],
+        default=None,
+        help="Filter to training or inference",
+    )
+    filters.add_argument(
+        "--backend",
+        type=str,
+        default=None,
+        help="Regex to filter backend (e.g. 'cudagraphs', 'dynamic')",
+    )
+    filters.add_argument(
+        "--dtype",
+        type=str,
+        default=None,
+        help="Filter to dtype (e.g. amp, float16, bfloat16)",
+    )
+    filters.add_argument(
+        "--runtime",
+        type=str,
+        default=None,
+        help="Filter to runtime (e.g. cuda, cpu, xpu)",
+    )
+    p_summary.add_argument(
+        "--group-by",
+        dest="group_by",
+        nargs="+",
+        default=None,
+        choices=GROUP_CHOICES,
+        metavar="KEY",
+        help=f"Group S-curves by key(s) (default: single combined). Choices: {', '.join(GROUP_CHOICES)}",
+    )
+    p_summary.add_argument(
+        "--attempt", type=int, default=1, help="Run attempt number (default: 1)"
+    )
+    p_summary.add_argument(
+        "--no-cache",
+        dest="no_cache",
+        action="store_true",
+        default=False,
+        help="Re-download artifacts even if cached",
+    )
+
+    # -- repro --
+    p_repro = sub.add_parser("repro", help="Reproduce a remote perf run locally")
+    p_repro.add_argument(
+        "run_id", type=str, help="GitHub Actions run ID or branch name"
+    )
+    p_repro.add_argument(
+        "--device",
+        nargs="+",
+        default=None,
+        choices=DEVICE_CHOICES,
+        metavar="DEVICE",
+        help="Device(s) (default: auto-discover from branch)",
+    )
+    p_repro.add_argument(
+        "--model",
+        type=str,
+        default=None,
+        help="Run only this model (e.g. BERT_pytorch)",
+    )
+    p_repro.add_argument(
+        "--suite", type=str, default=None, help="Filter to suite (hf, timm, tb)"
+    )
+    p_repro.add_argument(
+        "--mode",
+        type=str,
+        choices=["training", "inference"],
+        default=None,
+        help="Filter to mode",
+    )
+    p_repro.add_argument(
+        "--backend",
+        type=str,
+        default=None,
+        help="Regex to filter backend (e.g. 'cudagraphs', 'dynamic')",
+    )
+    p_repro.add_argument(
+        "--dtype", type=str, default=None, help="Filter to dtype (e.g. amp, bfloat16)"
+    )
+    p_repro.add_argument(
+        "--runtime", type=str, default=None, help="Filter to runtime (e.g. cuda, cpu)"
+    )
+    p_repro.add_argument(
+        "--attempt", type=int, default=1, help="Run attempt number (default: 1)"
+    )
+
+    # -- prepare-repro --
+    p_prep = sub.add_parser(
+        "prepare-repro", help="Show setup commands to prepare benchmark suites locally"
+    )
+    p_prep.add_argument(
+        "run_id",
+        nargs="?",
+        type=str,
+        default=None,
+        help="Optional: run ID or branch to also show benchmark commands",
+    )
+    p_prep.add_argument(
+        "--device",
+        nargs="+",
+        default=None,
+        choices=DEVICE_CHOICES,
+        metavar="DEVICE",
+        help="Device(s) for benchmark commands",
+    )
+    p_prep.add_argument(
+        "--suite",
+        type=str,
+        default=None,
+        help="Only show setup for this suite (hf, timm, tb)",
+    )
+    p_prep.add_argument(
+        "--mode",
+        type=str,
+        choices=["training", "inference"],
+        default=None,
+        help="Filter benchmark commands to mode",
+    )
+    p_prep.add_argument(
+        "--model", type=str, default=None, help="Filter benchmark commands to model"
+    )
+    p_prep.add_argument(
+        "--no-repro",
+        dest="no_repro",
+        action="store_true",
+        help="Only show setup, skip benchmark commands",
+    )
+    p_prep.add_argument(
+        "--attempt", type=int, default=1, help="Run attempt number (default: 1)"
+    )
+
+    args = parser.parse_args()
+
+    if args.command == "launch":
+        cmd_launch(args)
+    elif args.command == "summary":
+        cmd_summary(args)
+    elif args.command == "repro":
+        cmd_repro(args)
+    elif args.command == "prepare-repro":
+        cmd_prepare_repro(args)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/benchmarks/dynamo/pr_time_benchmarks/benchmarks/benchmark_base.py b/benchmarks/dynamo/pr_time_benchmarks/benchmarks/benchmark_base.py
index 17848a4a01e26..179a4e780ce59 100644
--- a/benchmarks/dynamo/pr_time_benchmarks/benchmarks/benchmark_base.py
+++ b/benchmarks/dynamo/pr_time_benchmarks/benchmarks/benchmark_base.py
@@ -3,7 +3,6 @@
 import json
 import os
 from abc import ABC, abstractmethod
-from typing import Optional
 from typing_extensions import Self
 
 import torch._C._instruction_counter as i_counter
@@ -114,7 +113,7 @@ def category(self) -> str:
     def device(self) -> str:
         return self._device
 
-    def is_dynamic(self) -> Optional[bool]:
+    def is_dynamic(self) -> bool | None:
         return self._dynamic
 
     def description(self) -> str:
diff --git a/benchmarks/dynamo/pr_time_benchmarks/benchmarks/unbind_split_flip_cat.py b/benchmarks/dynamo/pr_time_benchmarks/benchmarks/unbind_split_flip_cat.py
index 67bfbfb9cc22d..e3c8a04c86dda 100644
--- a/benchmarks/dynamo/pr_time_benchmarks/benchmarks/unbind_split_flip_cat.py
+++ b/benchmarks/dynamo/pr_time_benchmarks/benchmarks/unbind_split_flip_cat.py
@@ -4,6 +4,7 @@
 
 import torch
 import torch.export
+from torch.fx.experimental._config import AggressiveGuardFreeMode
 
 
 class UnbindSplitFlipCatModel(torch.nn.Module):
@@ -69,6 +70,10 @@ def description(self):
 
     def _prepare_once(self):
         torch._dynamo.config.capture_scalar_outputs = True
+        torch.fx.experimental._config.aggressive_guard_free_semantics = (
+            AggressiveGuardFreeMode.SKIP_RANGE_ANALYSIS
+        )
+        torch.fx.config.do_not_emit_stack_traces = True
         torch.manual_seed(0)
 
         self.model = UnbindSplitFlipCatModel(
diff --git a/benchmarks/dynamo/pr_time_benchmarks/expected_results.csv b/benchmarks/dynamo/pr_time_benchmarks/expected_results.csv
index 35529ae86b71a..375c73eb6e917 100644
--- a/benchmarks/dynamo/pr_time_benchmarks/expected_results.csv
+++ b/benchmarks/dynamo/pr_time_benchmarks/expected_results.csv
@@ -1,8 +1,8 @@
-add_loop_eager,compile_time_instruction_count,3347000000,0.1
+add_loop_eager,compile_time_instruction_count,3479000000,0.1
 
 
 
-add_loop_eager_dynamic,compile_time_instruction_count,4724000000,0.1
+add_loop_eager_dynamic,compile_time_instruction_count,4779000000,0.1
 
 
 
@@ -10,7 +10,7 @@ add_loop_inductor,compile_time_instruction_count,29660000000,0.1
 
 
 
-add_loop_inductor_dynamic_gpu,compile_time_instruction_count,38190000000,0.1
+add_loop_inductor_dynamic_gpu,compile_time_instruction_count,38150000000,0.1
 
 
 
@@ -18,15 +18,15 @@ add_loop_inductor_gpu,compile_time_instruction_count,26140000000,0.1
 
 
 
-basic_modules_ListOfLinears_eager,compile_time_instruction_count,898000000,0.1
+basic_modules_ListOfLinears_eager,compile_time_instruction_count,932400000,0.1
 
 
 
-basic_modules_ListOfLinears_inductor,compile_time_instruction_count,15240000000,0.1
+basic_modules_ListOfLinears_inductor,compile_time_instruction_count,15640000000,0.1
 
 
 
-basic_modules_ListOfLinears_inductor_gpu_force_shape_pad,compile_time_instruction_count,17720000000,0.1
+basic_modules_ListOfLinears_inductor_gpu_force_shape_pad,compile_time_instruction_count,17270000000,0.1
 
 
 
@@ -34,19 +34,19 @@ basic_modules_ListOfLinears_inductor_gpu,compile_time_instruction_count,11090000
 
 
 
-update_hint_regression,compile_time_instruction_count,1645000000,0.1
+update_hint_regression,compile_time_instruction_count,1609000000,0.1
 
 
 
-sum_floordiv_regression,compile_time_instruction_count,3581000000,0.1
+sum_floordiv_regression,compile_time_instruction_count,3641000000,0.1
 
 
 
-symint_sum,compile_time_instruction_count,3237000000,0.1
+symint_sum,compile_time_instruction_count,3210000000,0.1
 
 
 
-symint_sum_loop,compile_time_instruction_count,4583000000,0.1
+symint_sum_loop,compile_time_instruction_count,4448000000,0.1
 
 
 
@@ -54,15 +54,15 @@ aotdispatcher_inference_nosubclass_cpu,compile_time_instruction_count,1793000000
 
 
 
-aotdispatcher_inference_subclass_cpu,compile_time_instruction_count,4988000000,0.1
+aotdispatcher_inference_subclass_cpu,compile_time_instruction_count,5109000000,0.1
 
 
 
-aotdispatcher_partitioner_cpu,compile_time_instruction_count,8217000000,0.1
+aotdispatcher_partitioner_cpu,compile_time_instruction_count,8318000000,0.1
 
 
 
-aotdispatcher_partitioner_cpu2,compile_time_instruction_count,1918000000,0.1
+aotdispatcher_partitioner_cpu2,compile_time_instruction_count,1937000000,0.1
 
 
 
@@ -70,23 +70,23 @@ aotdispatcher_training_nosubclass_cpu,compile_time_instruction_count,3152000000,
 
 
 
-aotdispatcher_training_subclass_cpu,compile_time_instruction_count,8301000000,0.1
+aotdispatcher_training_subclass_cpu,compile_time_instruction_count,8317000000,0.1
 
 
 
-mm_loop_inductor_gpu,compile_time_instruction_count,4597000000,0.1
+mm_loop_inductor_gpu,compile_time_instruction_count,4707000000,0.1
 
 
 
-mm_loop_inductor_dynamic_gpu,compile_time_instruction_count,9051000000,0.1
+mm_loop_inductor_dynamic_gpu,compile_time_instruction_count,9270000000,0.1
 
 
 
-basic_NestedModule_eager,compile_time_instruction_count,5927000000,0.1
+basic_NestedModule_eager,compile_time_instruction_count,5866000000,0.1
 
 
 
-basic_InlineMod_eager,compile_time_instruction_count,7480000000,0.1
+basic_InlineMod_eager,compile_time_instruction_count,7863000000,0.1
 
 
 
@@ -94,19 +94,19 @@ dtensor_dispatch_detach,instruction_count,56080,0.1
 
 
 
-dtensor_dispatch_to_from_local,instruction_count,200600,0.1
+dtensor_dispatch_to_from_local,instruction_count,228200,0.1
 
 
 
-dtensor_dispatch_collectives,instruction_count,26720000,0.1
+dtensor_dispatch_collectives,instruction_count,28890000,0.1
 
 
 
-dtensor_dispatch_add_backward,instruction_count,346201,0.1
+dtensor_dispatch_add_backward,instruction_count,360100,0.1
 
 
 
-dtensor_dispatch_inplace,instruction_count,56530,0.1
+dtensor_dispatch_inplace,instruction_count,61490,0.1
 
 
 
@@ -114,8 +114,8 @@ dtensor_dispatch_view,instruction_count,66300,0.1
 
 
 
-dtensor_dispatch_random,instruction_count,557400,0.1
+dtensor_dispatch_random,instruction_count,592800,0.1
 
 
 
-dtensor_dispatch_custom_handler,instruction_count,36120,0.1
+dtensor_dispatch_custom_handler,instruction_count,36050,0.1
diff --git a/benchmarks/dynamo/timm_models.py b/benchmarks/dynamo/timm_models.py
index 5d845e4fcee00..09e7647537a4a 100755
--- a/benchmarks/dynamo/timm_models.py
+++ b/benchmarks/dynamo/timm_models.py
@@ -56,53 +56,6 @@ def pip_install(package):
         TIMM_MODELS[model_name] = int(batch_size)
 
 
-# TODO - Figure out the reason of cold start memory spike
-
-BATCH_SIZE_DIVISORS = {
-    "beit_base_patch16_224": 2,
-    "deit_base_distilled_patch16_224": 2,
-    "gluon_xception65": 2,
-    "mobilevit_s": 2,
-    "swin_base_patch4_window7_224": 2,
-}
-
-REQUIRE_HIGHER_TOLERANCE = {
-    "inception_v3",
-    "mobilenetv3_large_100",
-}
-
-REQUIRE_HIGHER_TOLERANCE_FP16_XPU = {
-    "botnet26t_256",
-}
-
-REQUIRE_HIGHER_TOLERANCE_AMP = {}
-
-REQUIRE_EVEN_HIGHER_TOLERANCE = {
-    "deit_base_distilled_patch16_224",
-    "vit_base_patch16_siglip_256",
-}
-
-# These models need higher tolerance in MaxAutotune mode
-REQUIRE_EVEN_HIGHER_TOLERANCE_MAX_AUTOTUNE = {}
-
-REQUIRE_HIGHER_TOLERANCE_FOR_FREEZING = {
-    "adv_inception_v3",
-}
-
-SCALED_COMPUTE_LOSS = {
-    "mobilevit_s",
-}
-
-FORCE_AMP_FOR_FP16_BF16_MODELS = {}
-
-SKIP_ACCURACY_CHECK_AS_EAGER_NON_DETERMINISTIC_MODELS = {}
-
-REQUIRE_LARGER_MULTIPLIER_FOR_SMALLER_TENSOR = {
-    "inception_v3",
-    "mobilenetv3_large_100",
-}
-
-
 def refresh_model_names():
     import glob
 
@@ -199,6 +152,22 @@ def _config(self):
     def _skip(self):
         return self._config["skip"]
 
+    @property
+    def _batch_size(self):
+        return self._config["batch_size"]
+
+    @property
+    def _tolerance(self):
+        return self._config["tolerance"]
+
+    @property
+    def _accuracy(self):
+        return self._config["accuracy"]
+
+    @property
+    def _require_larger_multiplier_for_smaller_tensor(self):
+        return self._config["require_larger_multiplier_for_smaller_tensor"]
+
     @property
     def skip_models_for_cpu(self):
         return self._skip["device"]["cpu"]
@@ -213,7 +182,7 @@ def skip_models(self):
 
     @property
     def force_amp_for_fp16_bf16_models(self):
-        return FORCE_AMP_FOR_FP16_BF16_MODELS
+        return self._config["dtype"]["force_amp_for_fp16_bf16_models"]
 
     @property
     def force_fp16_for_bf16_models(self):
@@ -226,17 +195,13 @@ def get_output_amp_train_process_func(self):
     @property
     def skip_accuracy_check_as_eager_non_deterministic(self):
         if self.args.accuracy and self.args.training:
-            return SKIP_ACCURACY_CHECK_AS_EAGER_NON_DETERMINISTIC_MODELS
+            return self._accuracy["skip"]["eager_not_deterministic"]
         return set()
 
     @property
     def guard_on_nn_module_models(self):
         return {}
 
-    @property
-    def inline_inbuilt_nn_modules_models(self):
-        return {}
-
     @download_retry_decorator
     def _download_model(self, model_name):
         model = create_model(
@@ -284,9 +249,10 @@ def load_model(
         input_size = data_config["input_size"]
         recorded_batch_size = TIMM_MODELS[model_name]
 
-        if model_name in BATCH_SIZE_DIVISORS:
+        batch_size_divisors = self._batch_size["divisors"]
+        if model_name in batch_size_divisors:
             recorded_batch_size = max(
-                int(recorded_batch_size / BATCH_SIZE_DIVISORS[model_name]), 1
+                int(recorded_batch_size / batch_size_divisors[model_name]), 1
             )
         batch_size = batch_size or recorded_batch_size
 
@@ -308,7 +274,7 @@ def load_model(
 
         self.loss = torch.nn.CrossEntropyLoss().to(device)
 
-        if model_name in SCALED_COMPUTE_LOSS:
+        if model_name in self._config["scaled_compute_loss"]:
             self.compute_loss = self.scaled_compute_loss
 
         if is_training and not use_eval_mode:
@@ -316,7 +282,7 @@ def load_model(
         else:
             model.eval()
 
-        self.validate_model(model, example_inputs)
+        self.validate_model(model_name, model, example_inputs)
 
         return device, model_name, model, example_inputs, batch_size
 
@@ -344,34 +310,34 @@ def pick_grad(self, name, is_training):
             return torch.no_grad()
 
     def use_larger_multiplier_for_smaller_tensor(self, name):
-        return name in REQUIRE_LARGER_MULTIPLIER_FOR_SMALLER_TENSOR
+        return name in self._require_larger_multiplier_for_smaller_tensor
 
     def get_tolerance_and_cosine_flag(self, is_training, current_device, name):
         cosine = self.args.cosine
         tolerance = 1e-3
 
-        if self.args.freezing and name in REQUIRE_HIGHER_TOLERANCE_FOR_FREEZING:
+        if self.args.freezing and name in self._tolerance["freezing"]:
             # the conv-batchnorm fusion used under freezing may cause relatively
-            # large numerical difference. We need are larger tolerance.
+            # large numerical difference. We need a larger tolerance.
             # Check https://github.com/pytorch/pytorch/issues/120545 for context
             tolerance = 8 * 1e-2
 
         if is_training:
             from torch._inductor import config as inductor_config
 
-            if name == "beit_base_patch16_224":
+            if name in self._tolerance["highest_training"]:
                 tolerance = 16 * 1e-2
-            elif name in REQUIRE_EVEN_HIGHER_TOLERANCE or (
+            elif name in self._tolerance["even_higher"] or (
                 inductor_config.max_autotune
-                and name in REQUIRE_EVEN_HIGHER_TOLERANCE_MAX_AUTOTUNE
+                and name in self._tolerance["even_higher_max_autotune"]
             ):
                 tolerance = 8 * 1e-2
-            elif name in REQUIRE_HIGHER_TOLERANCE or (
-                self.args.amp and name in REQUIRE_HIGHER_TOLERANCE_AMP
+            elif name in self._tolerance["higher_training"] or (
+                self.args.amp and name in self._tolerance["higher_amp"]
             ):
                 tolerance = 4 * 1e-2
             elif (
-                name in REQUIRE_HIGHER_TOLERANCE_FP16_XPU
+                name in self._tolerance["higher_fp16_xpu"]
                 and self.args.float16
                 and current_device == "xpu"
             ):
diff --git a/benchmarks/dynamo/timm_models.yaml b/benchmarks/dynamo/timm_models.yaml
index 1650a87500537..3e94903791f93 100644
--- a/benchmarks/dynamo/timm_models.yaml
+++ b/benchmarks/dynamo/timm_models.yaml
@@ -11,3 +11,62 @@ skip:
       - dm_nfnet_f0
       - nfnet_l0
       - visformer_small
+
+
+# TODO - Figure out the reason of cold start memory spike
+batch_size:
+  divisors:
+    beit_base_patch16_224: 2
+    deit_base_distilled_patch16_224: 2
+    gluon_xception65: 2
+    mobilevit_s: 2
+    swin_base_patch4_window7_224: 2
+
+
+tolerance:
+  higher_training:
+    - inception_v3
+    - mobilenetv3_large_100
+
+  higher_fp16_xpu:
+    - botnet26t_256
+
+  higher_amp: []
+
+  even_higher:
+    - deit_base_distilled_patch16_224
+    - vit_base_patch16_siglip_256
+
+  # These models need higher tolerance in MaxAutotune mode
+  even_higher_max_autotune: []
+
+  # beit_base_patch16_224 needs 16 * 1e-2 tolerance in training
+  highest_training:
+    - beit_base_patch16_224
+
+  freezing:
+    # the conv-batchnorm fusion used under freezing may cause relatively
+    # large numerical difference. We need a larger tolerance.
+    # Check https://github.com/pytorch/pytorch/issues/120545 for context
+    - adv_inception_v3
+
+  cosine: []
+
+
+scaled_compute_loss:
+  - mobilevit_s
+
+
+require_larger_multiplier_for_smaller_tensor:
+  - inception_v3
+  - mobilenetv3_large_100
+  - vit_base_patch14_dinov2.lvd142m
+
+
+dtype:
+  force_amp_for_fp16_bf16_models: []
+
+
+accuracy:
+  skip:
+    eager_not_deterministic: []
diff --git a/benchmarks/dynamo/torchao_backend.py b/benchmarks/dynamo/torchao_backend.py
index 6b4204db7b361..a32baf1bc7eb2 100644
--- a/benchmarks/dynamo/torchao_backend.py
+++ b/benchmarks/dynamo/torchao_backend.py
@@ -14,7 +14,6 @@ def setup_baseline():
 
 def torchao_optimize_ctx(quantization: str):
     from torchao.quantization.quant_api import (
-        autoquant,
         int4_weight_only,
         int8_dynamic_activation_int8_weight,
         int8_weight_only,
@@ -35,21 +34,7 @@ def _torchao_apply(module: torch.nn.Module, example_inputs: Any):
                     quantize_(module, int8_weight_only(), set_inductor_config=False)
                 elif quantization == "int4weightonly":
                     quantize_(module, int4_weight_only(), set_inductor_config=False)
-                if quantization == "autoquant":
-                    autoquant(module, error_on_unseen=False, set_inductor_config=False)
-                    if isinstance(example_inputs, dict):
-                        module(**example_inputs)
-                    else:
-                        module(*example_inputs)
-                    from torchao.quantization.autoquant import AUTOQUANT_CACHE
-
-                    if len(AUTOQUANT_CACHE) == 0:
-                        raise Exception(  # noqa: TRY002
-                            "NotAutoquantizable"
-                            f"Found no autoquantizable layers in model {type(module)}, stopping autoquantized run"
-                        )
-                else:
-                    unwrap_tensor_subclass(module)
+                unwrap_tensor_subclass(module)
                 setattr(module, "_quantized", True)  # noqa: B010
             model_iter_fn(module, example_inputs)
 
diff --git a/benchmarks/dynamo/torchbench.py b/benchmarks/dynamo/torchbench.py
index f44f7923adf1d..cbe575f926d22 100755
--- a/benchmarks/dynamo/torchbench.py
+++ b/benchmarks/dynamo/torchbench.py
@@ -9,6 +9,7 @@
 import tempfile
 import warnings
 from collections import namedtuple
+from functools import cached_property
 from os.path import abspath, exists
 
 import torch
@@ -207,18 +208,11 @@ def guard_on_nn_module_models(self):
             "vision_maskrcnn",
         }
 
-    @property
-    def inline_inbuilt_nn_modules_models(self):
-        return {
-            "basic_gnn_edgecnn",
-            "drq",
-            "hf_Reformer",
-            "DALLE2_pytorch",
-            "detectron2_maskrcnn_r_50_fpn",
-            "detectron2_maskrcnn_r_101_fpn",
-            "vision_maskrcnn",
-            "doctr_reco_predictor",
-        }
+    @cached_property
+    def _fb_models_available(self):
+        """This property exists because importing IS_FBCODE causes some models to be
+        frozen out of setting certain config flags."""
+        return importlib.util.find_spec("torchbenchmark.models.fb") is not None
 
     def load_model(
         self,
@@ -234,11 +228,14 @@ def load_model(
             )
         is_training = self.args.training
         use_eval_mode = self.args.use_eval_mode
+
         candidates = [
             f"torchbenchmark.models.{model_name}",
             f"torchbenchmark.canary_models.{model_name}",
-            f"torchbenchmark.models.fb.{model_name}",
         ]
+        if self._fb_models_available:
+            candidates.append(f"torchbenchmark.models.fb.{model_name}")
+
         for c in candidates:
             try:
                 module = importlib.import_module(c)
@@ -384,7 +381,7 @@ def load_model(
         ):
             model.config.use_cache = False
 
-        self.validate_model(model, example_inputs)
+        self.validate_model(benchmark.name, model, example_inputs)
         return device, benchmark.name, model, example_inputs, batch_size
 
     def iter_model_names(self, args):
@@ -415,7 +412,7 @@ def iter_model_names(self, args):
             yield model_name
 
     def pick_grad(self, name, is_training):
-        if is_training or name in ("maml",):
+        if is_training or name == "maml":
             return torch.enable_grad()
         else:
             return torch.no_grad()
diff --git a/benchmarks/dynamo/torchbench.yaml b/benchmarks/dynamo/torchbench.yaml
index b7f6229bed1cd..6ab39de217ba8 100644
--- a/benchmarks/dynamo/torchbench.yaml
+++ b/benchmarks/dynamo/torchbench.yaml
@@ -38,12 +38,12 @@ tolerance:
     - tacotron2
     - yolov3
     - squeezenet1_1
+    - shufflenet_v2_x1_0
 
   higher_fp16:
     - doctr_reco_predictor
     - drq
     - phlippe_resnet
-    - pytorch_CycleGAN_and_pix2pix
 
   higher_bf16:
     - doctr_reco_predictor
@@ -219,6 +219,9 @@ skip:
     # Has never been working correctly
     # https://github.com/pytorch/pytorch/issues/172015#issuecomment-3730509098
     - modded_nanogpt
+    # Broken since mid-2025, eager_fail_to_run on all platforms.
+    # Legacy 2017 model providing zero CI signal.
+    - pytorch_CycleGAN_and_pix2pix
 
   device:
     cpu:
diff --git a/benchmarks/functional_autograd_benchmark/torchaudio_models.py b/benchmarks/functional_autograd_benchmark/torchaudio_models.py
index 23e494e82aea7..b23eb320694df 100644
--- a/benchmarks/functional_autograd_benchmark/torchaudio_models.py
+++ b/benchmarks/functional_autograd_benchmark/torchaudio_models.py
@@ -3,7 +3,6 @@
 
 import math
 from collections import OrderedDict
-from typing import Optional
 
 import torch
 import torch.nn.functional as F
@@ -510,9 +509,9 @@ def forward(
         query: torch.Tensor,
         key: torch.Tensor,
         value: torch.Tensor,
-        attn_mask: Optional[torch.Tensor] = None,
-        bias_k: Optional[torch.Tensor] = None,
-        bias_v: Optional[torch.Tensor] = None,
+        attn_mask: torch.Tensor | None = None,
+        bias_k: torch.Tensor | None = None,
+        bias_v: torch.Tensor | None = None,
     ) -> tuple[torch.Tensor, torch.Tensor]:
         r"""
         Args:
@@ -590,9 +589,9 @@ def forward(
         query: torch.Tensor,
         key: torch.Tensor,
         value: torch.Tensor,
-        attn_mask: Optional[torch.Tensor] = None,
-        bias_k: Optional[torch.Tensor] = None,
-        bias_v: Optional[torch.Tensor] = None,
+        attn_mask: torch.Tensor | None = None,
+        bias_k: torch.Tensor | None = None,
+        bias_v: torch.Tensor | None = None,
     ) -> tuple[torch.Tensor, torch.Tensor]:
         r"""Uses a scaled dot product with the projected key-value pair to update
         the projected query.
diff --git a/benchmarks/functional_autograd_benchmark/utils.py b/benchmarks/functional_autograd_benchmark/utils.py
index 8efc0bdcddd12..e67e5d48eb119 100644
--- a/benchmarks/functional_autograd_benchmark/utils.py
+++ b/benchmarks/functional_autograd_benchmark/utils.py
@@ -1,18 +1,17 @@
 from collections import defaultdict
 from collections.abc import Callable
-from typing import Optional, Union
 
 import torch
 from torch import nn, Tensor
 
 
 # Type helpers
-InputsType = Union[Tensor, tuple[Tensor, ...]]
+InputsType = Tensor | tuple[Tensor, ...]
 # A Getter takes in a device and returns a callable and the inputs to that callable
 GetterReturnType = tuple[Callable[..., Tensor], InputsType]
 GetterType = Callable[[torch.device], GetterReturnType]
 # V here refers to the v in either vjp, jvp, vhp or hvp
-VType = Union[None, Tensor, tuple[Tensor, ...]]
+VType = None | Tensor | tuple[Tensor, ...]
 # Type used to store timing results. The first key is the model name, the second key
 # is the task name, the result is a Tuple of: speedup, mean_before, var_before, mean_after, var_after.
 TimingResultType = dict[str, dict[str, tuple[float, ...]]]
@@ -78,7 +77,7 @@ def load_weights(mod: nn.Module, names: list[str], params: tuple[Tensor, ...]) -
 
 # Utilities to read/write markdown table-like content.
 def to_markdown_table(
-    res: TimingResultType, header: Optional[tuple[str, ...]] = None
+    res: TimingResultType, header: tuple[str, ...] | None = None
 ) -> str:
     if header is None:
         header = ("model", "task", "mean", "var")
diff --git a/benchmarks/gpt_fast/common.py b/benchmarks/gpt_fast/common.py
index 4cbd0bd0f2dc6..211d5f6c5e301 100644
--- a/benchmarks/gpt_fast/common.py
+++ b/benchmarks/gpt_fast/common.py
@@ -1,6 +1,5 @@
 import dataclasses
 from collections.abc import Callable
-from typing import Optional
 
 
 all_experiments: dict[str, Callable] = {}
@@ -18,7 +17,7 @@ class Experiment:
     is_model: bool = False
 
 
-def register_experiment(name: Optional[str] = None):
+def register_experiment(name: str | None = None):
     def decorator(func):
         key = name or func.__name__
         all_experiments[key] = func
diff --git a/benchmarks/gpt_fast/generate.py b/benchmarks/gpt_fast/generate.py
index 42faefc6d58a9..a54e6df059e3b 100644
--- a/benchmarks/gpt_fast/generate.py
+++ b/benchmarks/gpt_fast/generate.py
@@ -2,9 +2,7 @@
 import itertools
 import platform
 import time
-from typing import Optional
 
-import torchao
 from common import Experiment, register_experiment
 from mixtral_moe_model import ConditionalFeedForward, Transformer as MixtralMoE
 from mixtral_moe_quantize import (
@@ -30,12 +28,12 @@
 class GPTModelConfig:
     name: str
     module: type
-    mode: Optional[str]
+    mode: str | None
     quantizer: type
     token_per_sec: float
     memory_bandwidth: float
     compilation_time: float
-    batch_size: Optional[int] = None
+    batch_size: int | None = None
 
 
 def device_sync(device):
@@ -62,7 +60,7 @@ def multinomial_sample_one_no_sync(
     return torch.argmax(probs_sort / q, dim=-1, keepdim=True).to(dtype=torch.int)
 
 
-def logits_to_probs(logits, temperature: float = 1.0, top_k: Optional[int] = None):
+def logits_to_probs(logits, temperature: float = 1.0, top_k: int | None = None):
     logits = logits / max(temperature, 1e-5)
 
     if top_k is not None:
@@ -73,7 +71,7 @@ def logits_to_probs(logits, temperature: float = 1.0, top_k: Optional[int] = Non
     return probs
 
 
-def sample(logits, temperature: float = 1.0, top_k: Optional[int] = None):
+def sample(logits, temperature: float = 1.0, top_k: int | None = None):
     probs = logits_to_probs(logits[0, -1], temperature, top_k)
     idx_next = multinomial_sample_one_no_sync(probs)
     return idx_next, probs
@@ -229,36 +227,6 @@ def run_experiment(
     start = -1
     compilation_time = None
 
-    if x.mode == "autoquant":
-        print("Using autoquant")
-        model = torchao.autoquant(model, manual=True, error_on_unseen=False)
-        generate(model, prompt, max_new_tokens, temperature=temperature, top_k=top_k)
-        model.finalize_autoquant()
-
-    if x.mode == "autoquant_v2":
-        print("Using autoquant_v2")
-        from torchao.prototype.quantization.autoquant_v2 import autoquant_v2
-
-        p = prompt.view(1, -1)
-        T = prompt.size(0)
-        T_new = T + max_new_tokens
-        max_seq_length = min(T_new, model.config.block_size)
-        input_pos = torch.arange(0, T, device=device)
-        example_input = (p, input_pos)
-
-        with torch.device(device):
-            model.setup_caches(max_batch_size=1, max_seq_length=max_seq_length)
-        model = autoquant_v2(
-            model,
-            manual=True,
-            error_on_unseen=False,
-            example_input=example_input,
-            batch_size=x.batch_size,
-        )
-        torch.compiler.cudagraph_mark_step_begin()
-        generate(model, prompt, max_new_tokens, temperature=temperature, top_k=top_k)
-        model.finalize_autoquant()
-
     global decode_one_token, prefill, compiled
     if not compiled:
         compiled = True
@@ -444,199 +412,3 @@ def run_mixtral_8x7b_int8(device: str = "cuda"):
             True,
         ),
     ]
-
-
-# token_per_sec and memory_bandwidth target numbers are for A100-40GB, which are different from the typical A100-80GB.
-def run_llama2_7b_autoquant(device: str = "cuda"):
-    model = GPTModelConfig(
-        "Llama-2-7b-chat-hf",
-        LLaMA,
-        "autoquant",
-        None,
-        144,
-        957,
-        136,
-    )
-    token_per_sec, memory_bandwidth, compilation_time = run_experiment(
-        model, device=device
-    )
-    return [
-        Experiment(
-            model.name,
-            "token_per_sec",
-            model.token_per_sec,
-            f"{token_per_sec:.02f}",
-            model.mode,
-            device,
-            get_arch_name(),
-            True,
-        ),
-        Experiment(
-            model.name,
-            "memory_bandwidth(GB/s)",
-            model.memory_bandwidth,
-            f"{memory_bandwidth:.02f}",
-            model.mode,
-            device,
-            get_arch_name(),
-            True,
-        ),
-        Experiment(
-            model.name,
-            "compilation_time(s)",
-            model.compilation_time,
-            f"{compilation_time:.02f}",
-            model.mode,
-            device,
-            get_arch_name(),
-            True,
-        ),
-    ]
-
-
-# token_per_sec and memory_bandwidth target numbers are for A100-40GB, which are different from the typical A100-80GB.
-def run_mixtral_8x7b_autoquant(device: str = "cuda"):
-    # We reduced the original number of layers from 32 to 16 to adapt CI memory limitation.
-    model = GPTModelConfig(
-        "Mixtral-8x7B-v0.1",
-        MixtralMoE,
-        "autoquant",
-        None,
-        175,
-        1130,
-        133,
-    )
-    token_per_sec, memory_bandwidth, compilation_time = run_experiment(
-        model, device=device
-    )
-    return [
-        Experiment(
-            model.name,
-            "token_per_sec",
-            model.token_per_sec,
-            f"{token_per_sec:.02f}",
-            model.mode,
-            device,
-            get_arch_name(),
-            True,
-        ),
-        Experiment(
-            model.name,
-            "memory_bandwidth(GB/s)",
-            model.memory_bandwidth,
-            f"{memory_bandwidth:.02f}",
-            model.mode,
-            device,
-            get_arch_name(),
-            True,
-        ),
-        Experiment(
-            model.name,
-            "compilation_time(s)",
-            model.compilation_time,
-            f"{compilation_time:.02f}",
-            model.mode,
-            device,
-            get_arch_name(),
-            True,
-        ),
-    ]
-
-
-# token_per_sec and memory_bandwidth target numbers are for A100-40GB, which are different from the typical A100-80GB.
-def run_llama2_7b_autoquant_v2(device: str = "cuda"):
-    model = GPTModelConfig(
-        "Llama-2-7b-chat-hf",
-        LLaMA,
-        "autoquant_v2",
-        None,
-        144,
-        957,
-        136,
-        6,  # batch_size
-    )
-    token_per_sec, memory_bandwidth, compilation_time = run_experiment(
-        model, device=device
-    )
-    return [
-        Experiment(
-            model.name,
-            "token_per_sec",
-            model.token_per_sec,
-            f"{token_per_sec:.02f}",
-            model.mode,
-            device,
-            get_arch_name(),
-            True,
-        ),
-        Experiment(
-            model.name,
-            "memory_bandwidth(GB/s)",
-            model.memory_bandwidth,
-            f"{memory_bandwidth:.02f}",
-            model.mode,
-            device,
-            get_arch_name(),
-            True,
-        ),
-        Experiment(
-            model.name,
-            "compilation_time(s)",
-            model.compilation_time,
-            f"{compilation_time:.02f}",
-            model.mode,
-            device,
-            get_arch_name(),
-            True,
-        ),
-    ]
-
-
-# token_per_sec and memory_bandwidth target numbers are for A100-40GB, which are different from the typical A100-80GB.
-def run_mixtral_8x7b_autoquant_v2(device: str = "cuda"):
-    # We reduced the original number of layers from 32 to 16 to adapt CI memory limitation.
-    model = GPTModelConfig(
-        "Mixtral-8x7B-v0.1",
-        MixtralMoE,
-        "autoquant_v2",
-        None,
-        175,
-        1130,
-        133,
-        6,  # batch_size
-    )
-    token_per_sec, memory_bandwidth, compilation_time = run_experiment(
-        model, device=device
-    )
-    return [
-        Experiment(
-            model.name,
-            "token_per_sec",
-            model.token_per_sec,
-            f"{token_per_sec:.02f}",
-            model.mode,
-            device,
-            get_arch_name(),
-            True,
-        ),
-        Experiment(
-            model.name,
-            "memory_bandwidth(GB/s)",
-            model.memory_bandwidth,
-            f"{memory_bandwidth:.02f}",
-            model.mode,
-            device,
-            get_arch_name(),
-            True,
-        ),
-        Experiment(
-            model.name,
-            "compilation_time(s)",
-            model.compilation_time,
-            f"{compilation_time:.02f}",
-            model.mode,
-            device,
-            get_arch_name(),
-            True,
-        ),
-    ]
diff --git a/benchmarks/gpt_fast/mixtral_moe_model.py b/benchmarks/gpt_fast/mixtral_moe_model.py
index 5b1ac3945d044..a0733588cab48 100644
--- a/benchmarks/gpt_fast/mixtral_moe_model.py
+++ b/benchmarks/gpt_fast/mixtral_moe_model.py
@@ -1,6 +1,5 @@
 # flake8: noqa: E266, C417, B950
 from dataclasses import dataclass
-from typing import Optional
 
 import torch
 import torch.nn as nn
@@ -106,8 +105,8 @@ def __init__(self, config: ModelArgs) -> None:
         self.norm = RMSNorm(config.dim, eps=config.norm_eps)
         self.output = nn.Linear(config.dim, config.vocab_size, bias=False)
 
-        self.freqs_cis: Optional[Tensor] = None
-        self.mask_cache: Optional[Tensor] = None
+        self.freqs_cis: Tensor | None = None
+        self.mask_cache: Tensor | None = None
         self.max_batch_size = -1
         self.max_seq_length = -1
 
@@ -135,7 +134,7 @@ def setup_caches(self, max_batch_size, max_seq_length):
             torch.ones(self.max_seq_length, self.max_seq_length, dtype=torch.bool)
         )
 
-    def forward(self, idx: Tensor, input_pos: Optional[Tensor] = None) -> Tensor:
+    def forward(self, idx: Tensor, input_pos: Tensor | None = None) -> Tensor:
         if self.freqs_cis is None:
             raise AssertionError("Caches must be initialized first")
         mask = self.causal_mask[None, None, input_pos]
@@ -201,7 +200,7 @@ def forward(
         x: Tensor,
         freqs_cis: Tensor,
         mask: Tensor,
-        input_pos: Optional[Tensor] = None,
+        input_pos: Tensor | None = None,
     ) -> Tensor:
         bsz, seqlen, _ = x.shape
 
diff --git a/benchmarks/gpt_fast/model.py b/benchmarks/gpt_fast/model.py
index 5f2fd6e20d926..5a1805bb1c32a 100644
--- a/benchmarks/gpt_fast/model.py
+++ b/benchmarks/gpt_fast/model.py
@@ -1,6 +1,5 @@
 # flake8: noqa: E266, C417, B950
 from dataclasses import dataclass
-from typing import Optional
 
 import torch
 import torch.nn as nn
@@ -125,8 +124,8 @@ def __init__(self, config: ModelArgs) -> None:
         self.norm = RMSNorm(config.dim, eps=config.norm_eps)
         self.output = nn.Linear(config.dim, config.vocab_size, bias=False)
 
-        self.freqs_cis: Optional[Tensor] = None
-        self.mask_cache: Optional[Tensor] = None
+        self.freqs_cis: Tensor | None = None
+        self.mask_cache: Tensor | None = None
         self.max_batch_size = -1
         self.max_seq_length = -1
 
@@ -154,7 +153,7 @@ def setup_caches(self, max_batch_size, max_seq_length):
             torch.ones(self.max_seq_length, self.max_seq_length, dtype=torch.bool)
         )
 
-    def forward(self, idx: Tensor, input_pos: Optional[Tensor] = None) -> Tensor:
+    def forward(self, idx: Tensor, input_pos: Tensor | None = None) -> Tensor:
         if self.freqs_cis is None:
             raise AssertionError("Caches must be initialized first")
         mask = self.causal_mask[None, None, input_pos]
@@ -220,7 +219,7 @@ def forward(
         x: Tensor,
         freqs_cis: Tensor,
         mask: Tensor,
-        input_pos: Optional[Tensor] = None,
+        input_pos: Tensor | None = None,
     ) -> Tensor:
         bsz, seqlen, _ = x.shape
 
diff --git a/benchmarks/inductor_backends/cutlass.py b/benchmarks/inductor_backends/cutlass.py
index eb5c3963340a7..c231e9ddcda63 100644
--- a/benchmarks/inductor_backends/cutlass.py
+++ b/benchmarks/inductor_backends/cutlass.py
@@ -11,7 +11,7 @@
 from collections import defaultdict
 from collections.abc import Callable
 from dataclasses import asdict, dataclass, field
-from typing import Any, Optional
+from typing import Any
 
 from tabulate import tabulate
 from tqdm import tqdm
@@ -363,7 +363,7 @@ def generate_experiment_configs(
 
 def calculate_table_data(results: list[ExperimentResults]) -> dict:
     table_data = defaultdict(list)
-    aten_perf: Optional[float] = None
+    aten_perf: float | None = None
 
     for experiment_result in results:
         for key, value in experiment_result.asdict().items():
diff --git a/benchmarks/instruction_counts/applications/ci.py b/benchmarks/instruction_counts/applications/ci.py
index 7a4f39b093d98..abd00d70dd410 100644
--- a/benchmarks/instruction_counts/applications/ci.py
+++ b/benchmarks/instruction_counts/applications/ci.py
@@ -6,7 +6,6 @@
 import hashlib
 import json
 import time
-from typing import Union
 
 from core.expand import materialize
 from definitions.standard import BENCHMARKS
@@ -56,7 +55,7 @@ def main(argv: list[str]) -> None:
     results = Runner(work_orders, cadence=30.0).run()
 
     # TODO: Annotate with TypedDict when 3.8 is the minimum supported version.
-    grouped_results: dict[str, dict[str, list[Union[float, int]]]] = {
+    grouped_results: dict[str, dict[str, list[float | int]]] = {
         key: {"times": [], "counts": []} for key in keys
     }
 
diff --git a/benchmarks/instruction_counts/core/api.py b/benchmarks/instruction_counts/core/api.py
index a41f9c226dc07..b2aaa5ec843c0 100644
--- a/benchmarks/instruction_counts/core/api.py
+++ b/benchmarks/instruction_counts/core/api.py
@@ -7,7 +7,7 @@
 import itertools as it
 import re
 import textwrap
-from typing import Optional, TYPE_CHECKING, Union
+from typing import TYPE_CHECKING
 
 from worker.main import WorkerTimerArgs
 
@@ -148,22 +148,22 @@ class GroupedBenchmark:
     # `GroupedStmts` (init_from_stmts) they are passed through from user args.
     # In the case of `GroupedModules` (init_from_model) they are generated
     # using `signature`. (e.g. `f(x, y) -> z` generates `z = model(x, y)`)
-    py_fwd_stmt: Optional[str]
-    cpp_fwd_stmt: Optional[str]
+    py_fwd_stmt: str | None
+    cpp_fwd_stmt: str | None
 
     # Code block used to define a model. `init_from_stmts` will never populate
     # `cpp_model_setup`, but if TorchScript is requested it will generate
     # `py_model_setup` using `torch.jit.script`.
-    py_model_setup: Optional[str]
-    cpp_model_setup: Optional[str]
+    py_model_setup: str | None
+    cpp_model_setup: str | None
 
     # True if this benchmark used `init_from_stmts`, otherwise False.
     inferred_model_setup: bool
 
     # Described above
     setup: GroupedSetup
-    signature_args: Optional[tuple[str, ...]]
-    signature_output: Optional[str]
+    signature_args: tuple[str, ...] | None
+    signature_output: str | None
     torchscript: bool
     autograd: bool
     num_threads: tuple[int, ...]
@@ -171,14 +171,14 @@ class GroupedBenchmark:
     @classmethod
     def init_from_stmts(
         cls,
-        py_stmt: Optional[str] = None,
-        cpp_stmt: Optional[str] = None,
+        py_stmt: str | None = None,
+        cpp_stmt: str | None = None,
         # Generic constructor arguments
         setup: GroupedSetup = GroupedSetup(),
-        signature: Optional[str] = None,
+        signature: str | None = None,
         torchscript: bool = False,
         autograd: bool = False,
-        num_threads: Union[int, tuple[int, ...]] = 1,
+        num_threads: int | tuple[int, ...] = 1,
     ) -> "GroupedBenchmark":
         """Create a set of benchmarks from free-form statements.
 
@@ -219,14 +219,14 @@ def init_from_stmts(
     @classmethod
     def init_from_model(
         cls,
-        py_model_setup: Optional[str] = None,
-        cpp_model_setup: Optional[str] = None,
+        py_model_setup: str | None = None,
+        cpp_model_setup: str | None = None,
         # Generic constructor arguments
         setup: GroupedSetup = GroupedSetup(),
-        signature: Optional[str] = None,
+        signature: str | None = None,
         torchscript: bool = False,
         autograd: bool = False,
-        num_threads: Union[int, tuple[int, ...]] = 1,
+        num_threads: int | tuple[int, ...] = 1,
     ) -> "GroupedBenchmark":
         """Create a set of benchmarks using torch.nn Modules.
 
@@ -263,8 +263,8 @@ def init_from_variants(
         cls,
         py_block: str = "",
         cpp_block: str = "",
-        num_threads: Union[int, tuple[int, ...]] = 1,
-    ) -> dict[Union[tuple[str, ...], Optional[str]], "GroupedBenchmark"]:
+        num_threads: int | tuple[int, ...] = 1,
+    ) -> dict[tuple[str, ...] | str | None, "GroupedBenchmark"]:
         py_cases, py_setup, py_global_setup = cls._parse_variants(
             py_block, Language.PYTHON
         )
@@ -285,7 +285,7 @@ def init_from_variants(
         # NB: The key is actually `Tuple[str, ...]`, however MyPy gets confused
         #     and we use the superset `Union[Tuple[str, ...], Optional[str]` to
         #     match the expected signature.
-        variants: dict[Union[tuple[str, ...], Optional[str]], GroupedBenchmark] = {}
+        variants: dict[tuple[str, ...] | str | None, GroupedBenchmark] = {}
 
         seen_labels: set[str] = set()
         for label in it.chain(py_cases.keys(), cpp_cases.keys()):
@@ -338,8 +338,8 @@ def __post_init__(self) -> None:
 
     @staticmethod
     def _parse_signature(
-        signature: Optional[str],
-    ) -> tuple[Optional[tuple[str, ...]], Optional[str]]:
+        signature: str | None,
+    ) -> tuple[tuple[str, ...] | None, str | None]:
         if signature is None:
             return None, None
 
@@ -362,9 +362,9 @@ def _parse_signature(
 
     @staticmethod
     def _model_from_py_stmt(
-        py_stmt: Optional[str],
-        signature_args: Optional[tuple[str, ...]],
-        signature_output: Optional[str],
+        py_stmt: str | None,
+        signature_args: tuple[str, ...] | None,
+        signature_output: str | None,
     ) -> str:
         if py_stmt is None:
             raise ValueError("`py_stmt` must be defined in order to derive a model.")
@@ -383,7 +383,7 @@ def model({", ".join(signature_args)}):
     @staticmethod
     def _make_model_invocation(
         signature_args: tuple[str, ...],
-        signature_output: Optional[str],
+        signature_output: str | None,
         runtime: RuntimeMode,
     ) -> tuple[str, str]:
         py_prefix, cpp_prefix = "", ""
diff --git a/benchmarks/instruction_counts/core/expand.py b/benchmarks/instruction_counts/core/expand.py
index 5fb4eb76a8cb2..58eae4ba16915 100644
--- a/benchmarks/instruction_counts/core/expand.py
+++ b/benchmarks/instruction_counts/core/expand.py
@@ -12,7 +12,7 @@
 import re
 import textwrap
 import uuid
-from typing import Optional, TYPE_CHECKING
+from typing import TYPE_CHECKING
 
 import torch
 
@@ -37,7 +37,7 @@
 )
 
 
-def _generate_torchscript_file(model_src: str, name: str) -> Optional[str]:
+def _generate_torchscript_file(model_src: str, name: str) -> str | None:
     """Returns the path a saved model if one can be constructed from `spec`.
 
     Because TorchScript requires actual source code in order to script a
@@ -99,7 +99,7 @@ def _get_stmt(
     runtime: RuntimeMode,
     autograd: AutogradMode,
     language: Language,
-) -> Optional[str]:
+) -> str | None:
     """Specialize a GroupedBenchmark for a particular configuration."""
     is_python = language == Language.PYTHON
 
@@ -143,7 +143,7 @@ def _get_setup(
     runtime: RuntimeMode,
     language: Language,
     stmt: str,
-    model_path: Optional[str],
+    model_path: str | None,
 ) -> str:
     """Specialize a GroupedBenchmark for a particular configuration.
 
@@ -234,7 +234,7 @@ def materialize(benchmarks: FlatIntermediateDefinition) -> FlatDefinition:
             if not isinstance(args, GroupedBenchmark):
                 raise AssertionError(f"Expected GroupedBenchmark, but got {type(args)}")
 
-            model_path: Optional[str] = None
+            model_path: str | None = None
             if args.py_model_setup and args.torchscript:
                 model_setup = (
                     f"{args.py_model_setup}\njit_model = torch.jit.script(model)"
diff --git a/benchmarks/instruction_counts/core/types.py b/benchmarks/instruction_counts/core/types.py
index d509ffaa7deb3..346d37b9fc90f 100644
--- a/benchmarks/instruction_counts/core/types.py
+++ b/benchmarks/instruction_counts/core/types.py
@@ -2,8 +2,6 @@
 
 # mypy: ignore-errors
 
-from typing import Optional, Union
-
 from core.api import AutoLabels, GroupedBenchmark, TimerArgs
 
 
@@ -67,19 +65,16 @@
 # Allow strings in definition for convenience, and None to signify a base
 # case. (No subsequent entry needed. See the "add" example above.)
 Label = tuple[str, ...]
-_Label = Union[Label, Optional[str]]
+_Label = Label | str | None
 
-_Value = Union[
-    Union[TimerArgs, GroupedBenchmark],
-    dict[_Label, "_Value"],
-]
+_Value = TimerArgs | GroupedBenchmark | dict[_Label, "_Value"]
 
 Definition = dict[_Label, _Value]
 
 # We initially have to parse (flatten) to an intermediate state in order to
 # build TorchScript models since multiple entries will share the same model
 # artifact.
-FlatIntermediateDefinition = dict[Label, Union[TimerArgs, GroupedBenchmark]]
+FlatIntermediateDefinition = dict[Label, TimerArgs | GroupedBenchmark]
 
 # Final parsed schema.
 FlatDefinition = tuple[tuple[Label, AutoLabels, TimerArgs], ...]
diff --git a/benchmarks/instruction_counts/core/utils.py b/benchmarks/instruction_counts/core/utils.py
index f960cfd0fbffb..6c6e2b3253bcc 100644
--- a/benchmarks/instruction_counts/core/utils.py
+++ b/benchmarks/instruction_counts/core/utils.py
@@ -3,7 +3,6 @@
 import re
 import shutil
 import textwrap
-from typing import Optional
 
 from core.api import GroupedBenchmark, TimerArgs
 from core.types import Definition, FlatIntermediateDefinition, Label
@@ -11,7 +10,7 @@
 from torch.utils.benchmark.utils.common import _make_temp_dir
 
 
-_TEMPDIR: Optional[str] = None
+_TEMPDIR: str | None = None
 
 
 def get_temp_dir() -> str:
diff --git a/benchmarks/instruction_counts/execution/runner.py b/benchmarks/instruction_counts/execution/runner.py
index 74e099d816606..0ebe837aaee59 100644
--- a/benchmarks/instruction_counts/execution/runner.py
+++ b/benchmarks/instruction_counts/execution/runner.py
@@ -8,7 +8,6 @@
 import textwrap
 import threading
 import time
-from typing import Optional, Union
 
 from worker.main import WorkerFailure, WorkerOutput
 
@@ -21,9 +20,9 @@
 class WorkerFailed(Exception):
     """Raised in the main process when a worker failure is detected."""
 
-    def __init__(self, cmd: str, wrapped_trace: Optional[str] = None) -> None:
+    def __init__(self, cmd: str, wrapped_trace: str | None = None) -> None:
         self.cmd: str = cmd
-        self.wrapped_trace: Optional[str] = wrapped_trace
+        self.wrapped_trace: str | None = wrapped_trace
         super().__init__()
 
 
@@ -65,7 +64,7 @@ def __init__(self, min_core_id: int, max_core_id: int) -> None:
         self._reservations: dict[str, tuple[int, ...]] = {}
         self._lock = threading.Lock()
 
-    def reserve(self, n: int) -> Optional[str]:
+    def reserve(self, n: int) -> str | None:
         """Simple first-fit policy.
 
         If successful, return a string for `taskset`. Otherwise, return None.
@@ -95,7 +94,7 @@ class Runner:
     def __init__(
         self,
         work_items: tuple[WorkOrder, ...],
-        core_pool: Optional[CorePool] = None,
+        core_pool: CorePool | None = None,
         cadence: float = 1.0,
     ) -> None:
         self._work_items: tuple[WorkOrder, ...] = work_items
@@ -110,7 +109,7 @@ def __init__(
         # Debug information for ETA and error messages.
         self._start_time: float = -1
         self._durations: dict[WorkOrder, float] = {}
-        self._currently_processed: Optional[WorkOrder] = None
+        self._currently_processed: WorkOrder | None = None
 
         if len(work_items) != len(set(work_items)):
             raise ValueError("Duplicate work items.")
@@ -164,7 +163,7 @@ def _update_active_jobs(self) -> None:
                 active_jobs.append(job)
                 continue
 
-            result: Union[WorkerOutput, WorkerFailure] = job.result
+            result: WorkerOutput | WorkerFailure = job.result
             if isinstance(result, WorkerOutput):
                 self._results[job.work_order] = result
                 if job.cpu_list is None:
diff --git a/benchmarks/instruction_counts/execution/work.py b/benchmarks/instruction_counts/execution/work.py
index 3c1f33344c8ac..deac1db48d7ff 100644
--- a/benchmarks/instruction_counts/execution/work.py
+++ b/benchmarks/instruction_counts/execution/work.py
@@ -10,7 +10,7 @@
 import subprocess
 import time
 import uuid
-from typing import Optional, TYPE_CHECKING, Union
+from typing import TYPE_CHECKING
 
 from core.api import AutoLabels
 from core.types import Label
@@ -46,8 +46,8 @@ class WorkOrder:
     label: Label
     autolabels: AutoLabels
     timer_args: WorkerTimerArgs
-    source_cmd: Optional[str] = None
-    timeout: Optional[float] = None
+    source_cmd: str | None = None
+    timeout: float | None = None
     retries: int = 0
 
     def __hash__(self) -> int:
@@ -67,17 +67,17 @@ class _BenchmarkProcess:
     """Wraps subprocess.Popen for a given WorkOrder."""
 
     _work_order: WorkOrder
-    _cpu_list: Optional[str]
+    _cpu_list: str | None
     _proc: PopenType
 
     # Internal bookkeeping
     _communication_file: str
     _start_time: float
-    _end_time: Optional[float] = None
-    _retcode: Optional[int]
-    _result: Optional[Union[WorkerOutput, WorkerFailure]] = None
+    _end_time: float | None = None
+    _retcode: int | None
+    _result: WorkerOutput | WorkerFailure | None = None
 
-    def __init__(self, work_order: WorkOrder, cpu_list: Optional[str]) -> None:
+    def __init__(self, work_order: WorkOrder, cpu_list: str | None) -> None:
         self._work_order = work_order
         self._cpu_list = cpu_list
         self._start_time = time.time()
@@ -129,13 +129,13 @@ def duration(self) -> float:
         return (self._end_time or time.time()) - self._start_time
 
     @property
-    def result(self) -> Union[WorkerOutput, WorkerFailure]:
+    def result(self) -> WorkerOutput | WorkerFailure:
         self._maybe_collect()
         if self._result is None:
             raise AssertionError("result is None after collection")
         return self._result
 
-    def poll(self) -> Optional[int]:
+    def poll(self) -> int | None:
         self._maybe_collect()
         return self._retcode
 
@@ -189,7 +189,7 @@ class InProgress:
     _proc: _BenchmarkProcess
     _timeouts: int = 0
 
-    def __init__(self, work_order: WorkOrder, cpu_list: Optional[str]):
+    def __init__(self, work_order: WorkOrder, cpu_list: str | None):
         self._work_order = work_order
         self._proc = _BenchmarkProcess(work_order, cpu_list)
 
@@ -198,7 +198,7 @@ def work_order(self) -> WorkOrder:
         return self._proc._work_order
 
     @property
-    def cpu_list(self) -> Optional[str]:
+    def cpu_list(self) -> str | None:
         return self._proc._cpu_list
 
     @property
@@ -232,7 +232,7 @@ def check_finished(self) -> bool:
         raise subprocess.TimeoutExpired(cmd=self._proc.cmd, timeout=timeout)
 
     @property
-    def result(self) -> Union[WorkerOutput, WorkerFailure]:
+    def result(self) -> WorkerOutput | WorkerFailure:
         return self._proc.result
 
     def __hash__(self) -> int:
diff --git a/benchmarks/instruction_counts/worker/main.py b/benchmarks/instruction_counts/worker/main.py
index e2f513741ce20..e08e741029844 100644
--- a/benchmarks/instruction_counts/worker/main.py
+++ b/benchmarks/instruction_counts/worker/main.py
@@ -24,7 +24,7 @@
 import sys
 import timeit
 import traceback
-from typing import Any, TYPE_CHECKING, Union
+from typing import Any, TYPE_CHECKING
 
 
 if TYPE_CHECKING:
@@ -121,7 +121,7 @@ def load_input(self) -> WorkerTimerArgs:
             raise AssertionError(f"expected WorkerTimerArgs, got {type(result)}")
         return result
 
-    def load_output(self) -> Union[WorkerTimerArgs, WorkerOutput, WorkerFailure]:
+    def load_output(self) -> WorkerTimerArgs | WorkerOutput | WorkerFailure:
         """Convenience method for type safe loading."""
         result = self.load()
         if not isinstance(result, (WorkerTimerArgs, WorkerOutput, WorkerFailure)):
@@ -163,7 +163,7 @@ def _run(timer_args: WorkerTimerArgs) -> WorkerOutput:
 
 
 def main(communication_file: str) -> None:
-    result: Union[WorkerOutput, WorkerFailure]
+    result: WorkerOutput | WorkerFailure
     try:
         with open(communication_file, "rb") as f:
             timer_args: WorkerTimerArgs = WorkerUnpickler(f).load_input()
diff --git a/benchmarks/operator_benchmark/benchmark_core.py b/benchmarks/operator_benchmark/benchmark_core.py
index b183d65539737..679aea36cd479 100644
--- a/benchmarks/operator_benchmark/benchmark_core.py
+++ b/benchmarks/operator_benchmark/benchmark_core.py
@@ -8,7 +8,7 @@
 import timeit
 from collections import namedtuple
 from dataclasses import asdict, dataclass
-from typing import Any, Optional
+from typing import Any
 
 import benchmark_utils
 
@@ -598,7 +598,7 @@ def _output_json(
             @dataclass
             class BenchmarkInfo:
                 name: str
-                mode: Optional[str]
+                mode: str | None
                 dtype: str
                 extra_info: dict[str, Any]
 
@@ -614,7 +614,7 @@ class MetricInfo:
                 name: str
                 unit: str
                 benchmark_values: list[float]
-                target_value: Optional[float]
+                target_value: float | None
 
             @dataclass
             class BenchmarkRecord:
diff --git a/benchmarks/operator_benchmark/benchmark_runner.py b/benchmarks/operator_benchmark/benchmark_runner.py
index 6568cf9bf3ee6..4c96b40b8db47 100644
--- a/benchmarks/operator_benchmark/benchmark_runner.py
+++ b/benchmarks/operator_benchmark/benchmark_runner.py
@@ -16,6 +16,7 @@
     description="Run microbenchmarks.",
     formatter_class=argparse.ArgumentDefaultsHelpFormatter,
     conflict_handler="resolve",
+    allow_abbrev=False,
 )
 
 
diff --git a/benchmarks/operator_benchmark/pt/interpolate_test.py b/benchmarks/operator_benchmark/pt/interpolate_test.py
index c65f0b9c241e9..7eb435a640eae 100644
--- a/benchmarks/operator_benchmark/pt/interpolate_test.py
+++ b/benchmarks/operator_benchmark/pt/interpolate_test.py
@@ -14,13 +14,14 @@ def init(
         channels_last=False,
         mode="linear",
         dtype=torch.float,
+        device="cpu",
     ):
         input_image = torch.randint(
             0,
             256,
             size=input_size,
             dtype=dtype,
-            device="cpu",
+            device=device,
             requires_grad=self.auto_set(),
         )
         if channels_last:
@@ -151,7 +152,29 @@ def forward(self, input_image, output_size, mode, align_corners):
 )
 
 
-for config in (config_short, config_long, config_3d, config_5d):
+# CUDA bicubic with high channel counts (VLM position embedding workloads).
+# These exercise the parallel batch/channel kernel for small output spatial sizes.
+config_cuda_bicubic = op_bench.config_list(
+    attr_names=["input_size", "output_size"],
+    attrs=[
+        [(64, 768, 16, 16), (6, 6)],
+        [(64, 768, 16, 16), (14, 14)],
+        [(64, 1152, 16, 16), (6, 6)],
+        [(64, 1152, 32, 32), (14, 14)],
+        [(4, 256, 64, 64), (32, 32)],
+        # Regression check: large output with few channels should not regress
+        [(1, 3, 320, 320), (256, 256)],
+        [(1, 3, 500, 500), (800, 800)],
+    ],
+    cross_product_configs={
+        "mode": ["bicubic"],
+        "device": ["cuda"],
+    },
+    tags=["short"],
+)
+
+
+for config in (config_short, config_long, config_3d, config_5d, config_cuda_bicubic):
     op_bench.generate_pt_test(config, InterpolateBenchmark)
 
 
diff --git a/benchmarks/operator_benchmark/pt/qembedding_bag_lookups_test.py b/benchmarks/operator_benchmark/pt/qembedding_bag_lookups_test.py
index f5fd9dc4ff345..7ea8da20a4dcb 100644
--- a/benchmarks/operator_benchmark/pt/qembedding_bag_lookups_test.py
+++ b/benchmarks/operator_benchmark/pt/qembedding_bag_lookups_test.py
@@ -1,5 +1,3 @@
-from typing import Optional
-
 import numpy as np
 
 import operator_benchmark as op_bench
@@ -180,10 +178,10 @@ def forward(
         indices,
         offsets,
         mode: int,
-        per_sample_weights: Optional[torch.Tensor],
+        per_sample_weights: torch.Tensor | None,
         include_last_offset: bool,
         is_pruned_weights: bool,
-        compressed_indices: Optional[torch.Tensor],
+        compressed_indices: torch.Tensor | None,
     ):
         return self.op_func(
             prepacked_weights,
@@ -292,10 +290,10 @@ def forward(
         indices,
         offsets,
         mode: int,
-        per_sample_weights: Optional[torch.Tensor],
+        per_sample_weights: torch.Tensor | None,
         include_last_offset: bool,
         is_pruned_weights: bool,
-        compressed_indices: Optional[torch.Tensor],
+        compressed_indices: torch.Tensor | None,
     ):
         return self.op_func(
             prepacked_weights,
diff --git a/benchmarks/operator_benchmark/pt/scaled_mm_common.py b/benchmarks/operator_benchmark/pt/scaled_mm_common.py
index 7440023480a6f..2f4657889942f 100644
--- a/benchmarks/operator_benchmark/pt/scaled_mm_common.py
+++ b/benchmarks/operator_benchmark/pt/scaled_mm_common.py
@@ -1,7 +1,6 @@
 import importlib.util
 import os
 from types import ModuleType
-from typing import Optional
 
 import torch
 from torch.torch_version import TorchVersion
@@ -15,7 +14,7 @@
 and utilities to reduce code duplication.
 """
 
-_TEST_SCALED_MATMUL_CUDA_MOD: Optional[ModuleType] = None
+_TEST_SCALED_MATMUL_CUDA_MOD: ModuleType | None = None
 
 
 # Shared benchmark shapes for scaled matmul operations
diff --git a/benchmarks/serialization/export_save_linear_bench.py b/benchmarks/serialization/export_save_linear_bench.py
new file mode 100644
index 0000000000000..6f3f9360e9896
--- /dev/null
+++ b/benchmarks/serialization/export_save_linear_bench.py
@@ -0,0 +1,287 @@
+import argparse
+import gc
+import statistics
+import tempfile
+import time
+from dataclasses import dataclass
+from pathlib import Path
+
+import torch
+
+
+DEFAULT_NUM_PARAMS = (
+    1_000_000,
+    10_000_000,
+    100_000_000,
+    1_000_000_000,
+    2_000_000_000,
+    3_000_000_000,
+    4_000_000_000,
+)
+NUM_LAYERS = 5
+DEFAULT_REPEATS = 5
+
+
+class LinearModel(torch.nn.Module):
+    def __init__(self, hidden_size: int, *, dtype: torch.dtype) -> None:
+        super().__init__()
+        self.layers = torch.nn.ModuleList(
+            [
+                torch.nn.Linear(
+                    hidden_size,
+                    hidden_size,
+                    bias=True,
+                    dtype=dtype,
+                )
+                for _ in range(NUM_LAYERS)
+            ]
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        for layer in self.layers:
+            x = layer(x)
+        return x
+
+
+@dataclass
+class BenchmarkResult:
+    num_params: int
+    median_save_ms: float | None
+    status: str
+    error_detail: str | None = None
+
+
+def _parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(
+        description=(
+            "Benchmark torch.export.save for a five-layer LinearModel across "
+            "a range of parameter counts."
+        )
+    )
+    parser.add_argument(
+        "--num-params",
+        type=int,
+        nargs="+",
+        default=list(DEFAULT_NUM_PARAMS),
+        help=(
+            "Target total parameter counts to benchmark. The script derives the "
+            "nearest hidden size for each target."
+        ),
+    )
+    parser.add_argument(
+        "--repeats",
+        type=int,
+        default=DEFAULT_REPEATS,
+        help="Number of torch.export.save timings to collect per parameter-count case.",
+    )
+    parser.add_argument(
+        "--batch-size",
+        type=int,
+        default=0,
+        help=(
+            "Batch size for the example input passed to torch.export.export. "
+            "The default of 0 keeps the benchmark focused on serialization cost."
+        ),
+    )
+    parser.add_argument(
+        "--dtype",
+        choices=("float32", "float16", "bfloat16"),
+        default="float32",
+        help="Parameter and example-input dtype.",
+    )
+    return parser.parse_args()
+
+
+def _dtype_from_name(dtype_name: str) -> torch.dtype:
+    return {
+        "float32": torch.float32,
+        "float16": torch.float16,
+        "bfloat16": torch.bfloat16,
+    }[dtype_name]
+
+
+def _parameter_count(hidden_size: int) -> int:
+    return NUM_LAYERS * (hidden_size * hidden_size + hidden_size)
+
+
+def _hidden_size_from_num_params(num_params: int) -> int:
+    if num_params <= 0:
+        raise ValueError("--num-params values must be positive")
+
+    root = (-1.0 + (1.0 + 4.0 * num_params / NUM_LAYERS) ** 0.5) / 2.0
+    base_hidden_size = max(1, int(root))
+    candidates = [
+        candidate
+        for candidate in (
+            base_hidden_size - 1,
+            base_hidden_size,
+            base_hidden_size + 1,
+            base_hidden_size + 2,
+        )
+        if candidate >= 1
+    ]
+    return min(
+        candidates,
+        key=lambda hidden_size: abs(_parameter_count(hidden_size) - num_params),
+    )
+
+
+def _format_num_params(num_params: int) -> str:
+    for threshold, suffix in (
+        (1_000_000_000, "B"),
+        (1_000_000, "M"),
+        (1_000, "K"),
+    ):
+        if num_params >= threshold:
+            value = num_params / threshold
+            return f"{value:.1f}".rstrip("0").rstrip(".") + suffix
+    return str(num_params)
+
+
+def _measure_save_times_ms(
+    exported_program: torch.export.ExportedProgram,
+    repeats: int,
+) -> list[float]:
+    times_ms: list[float] = []
+    with tempfile.TemporaryDirectory(prefix="export_save_bench_") as temp_dir:
+        temp_path = Path(temp_dir)
+        for iteration in range(repeats):
+            save_path = temp_path / f"exported_program_{iteration}.pt2"
+            start_time = time.perf_counter()
+            torch.export.save(exported_program, save_path)
+            elapsed_ms = (time.perf_counter() - start_time) * 1000.0
+            times_ms.append(elapsed_ms)
+            save_path.unlink(missing_ok=True)
+    return times_ms
+
+
+def _status_for_exception(exc: BaseException) -> tuple[str, str]:
+    message = (
+        str(exc).strip().splitlines()[0] if str(exc).strip() else type(exc).__name__
+    )
+    lowered = message.lower()
+    if (
+        isinstance(exc, MemoryError)
+        or "out of memory" in lowered
+        or "can't allocate memory" in lowered
+    ):
+        return "oom", message
+    return "error", message
+
+
+def _run_case(
+    target_num_params: int,
+    *,
+    repeats: int,
+    batch_size: int,
+    dtype: torch.dtype,
+) -> BenchmarkResult:
+    exported_program = None
+    hidden_size = _hidden_size_from_num_params(target_num_params)
+    num_params = _parameter_count(hidden_size)
+    try:
+        with torch.no_grad():
+            model = LinearModel(hidden_size, dtype=dtype).eval()
+            model.requires_grad_(False)
+            example_input = torch.zeros(batch_size, hidden_size, dtype=dtype)
+            exported_program = torch.export.export(
+                model,
+                (example_input,),
+                strict=True,
+            )
+
+        save_times_ms = _measure_save_times_ms(exported_program, repeats)
+        return BenchmarkResult(
+            num_params=num_params,
+            median_save_ms=statistics.median(save_times_ms),
+            status="ok",
+        )
+    except (MemoryError, RuntimeError) as exc:
+        status, detail = _status_for_exception(exc)
+        return BenchmarkResult(
+            num_params=num_params,
+            median_save_ms=None,
+            status=status,
+            error_detail=detail,
+        )
+    finally:
+        del exported_program
+        del example_input
+        del model
+        gc.collect()
+
+
+def _format_table(results: list[BenchmarkResult]) -> str:
+    headers = ("num_params", "median_save_ms", "status")
+    rows = [
+        (
+            _format_num_params(result.num_params),
+            (
+                f"{result.median_save_ms:.3f}"
+                if result.median_save_ms is not None
+                else "n/a"
+            ),
+            result.status,
+        )
+        for result in results
+    ]
+    widths = [
+        max(len(header), *(len(row[column]) for row in rows))
+        for column, header in enumerate(headers)
+    ]
+
+    def _format_row(row: tuple[str, ...]) -> str:
+        return " | ".join(
+            [
+                row[0].rjust(widths[0]),
+                row[1].rjust(widths[1]),
+                row[2].ljust(widths[2]),
+            ]
+        )
+
+    lines = [
+        _format_row(headers),
+        "-+-".join("-" * width for width in widths),
+    ]
+    lines.extend(_format_row(row) for row in rows)
+    return "\n".join(lines)
+
+
+def main() -> None:
+    args = _parse_args()
+    dtype = _dtype_from_name(args.dtype)
+
+    print("Benchmarking torch.export.save for a 5-layer LinearModel")
+    print(f"dtype={args.dtype}, batch_size={args.batch_size}, repeats={args.repeats}")
+
+    results = [
+        _run_case(
+            num_params,
+            repeats=args.repeats,
+            batch_size=args.batch_size,
+            dtype=dtype,
+        )
+        for num_params in args.num_params
+    ]
+
+    print()
+    print(
+        "Note: This benchmark is highly sensitive to disk performance and OS "
+        "behavior. Results for smaller parameter counts can be noisy or flaky."
+    )
+    print()
+    print(_format_table(results))
+
+    failures = [result for result in results if result.error_detail is not None]
+    if failures:
+        print()
+        print("failed cases:")
+        for result in failures:
+            print(
+                f"- num_params={_format_num_params(result.num_params)}: "
+                f"{result.error_detail}"
+            )
+
+
+if __name__ == "__main__":
+    main()
diff --git a/benchmarks/static_runtime/test_static_runtime.cc b/benchmarks/static_runtime/test_static_runtime.cc
index 69d925675d478..cf8f333bc2c25 100644
--- a/benchmarks/static_runtime/test_static_runtime.cc
+++ b/benchmarks/static_runtime/test_static_runtime.cc
@@ -1074,6 +1074,10 @@ TEST(StaticRuntime, NanToNum) {
 }
 
 TEST(StaticRuntime, Stack) {
+#if defined(__aarch64__) || defined(_M_ARM64)
+  // See https://github.com/pytorch/pytorch/issues/178522.
+  GTEST_SKIP() << "Skipping StaticRuntime.Stack on AArch64.";
+#endif
   const auto stack_dim = R"JIT(
     def forward(self, a: Tensor, b: Tensor, dim: int):
         inputs = [a]
@@ -2307,6 +2311,10 @@ TEST(StaticRuntime, Append) {
 }
 
 TEST(StaticRuntime, QuantizedLinear) {
+#if defined(__aarch64__) || defined(_M_ARM64)
+  // See https://github.com/pytorch/pytorch/issues/178522.
+  GTEST_SKIP() << "Skipping QuantizedLinear on AArch64.";
+#endif
   const std::string quantize_script = R"IR(
     graph(%input: Tensor, %weights: Tensor):
         %scale: float = prim::Constant[value=1.]()
diff --git a/benchmarks/transformer/attention_bias_benchmarks.py b/benchmarks/transformer/attention_bias_benchmarks.py
index dad17e4191a12..27e749518b87e 100644
--- a/benchmarks/transformer/attention_bias_benchmarks.py
+++ b/benchmarks/transformer/attention_bias_benchmarks.py
@@ -2,7 +2,6 @@
 from collections.abc import Callable
 from dataclasses import asdict, dataclass
 from functools import partial
-from typing import Union
 
 import numpy as np
 from tabulate import tabulate
@@ -107,7 +106,7 @@ def forward(
         query: torch.Tensor,
         key: torch.Tensor,
         value: torch.Tensor,
-        mask: Union[torch.Tensor, CausalBias],
+        mask: torch.Tensor | CausalBias,
     ):
         query_projected = F.linear(query, self.q_proj_weight)
         key_projected = F.linear(key, self.k_proj_weight)
diff --git a/benchmarks/transformer/better_transformer_vs_mha_functional.py b/benchmarks/transformer/better_transformer_vs_mha_functional.py
index b5e0a6b75d31b..fa2e2ee151ecd 100644
--- a/benchmarks/transformer/better_transformer_vs_mha_functional.py
+++ b/benchmarks/transformer/better_transformer_vs_mha_functional.py
@@ -16,7 +16,6 @@
 from collections import defaultdict, OrderedDict
 from pathlib import Path
 from pprint import pprint
-from typing import Optional
 
 import numpy as np
 from prettytable import PrettyTable
@@ -159,7 +158,7 @@ def run(
     return result_entry
 
 
-def main(save_path: Optional[Path], error_path: Optional[Path]):
+def main(save_path: Path | None, error_path: Path | None):
     table = PrettyTable()
     entries = defaultdict(list)
 
diff --git a/benchmarks/transformer/score_mod.py b/benchmarks/transformer/score_mod.py
index f9fb6a925ed40..e196b7de791c0 100644
--- a/benchmarks/transformer/score_mod.py
+++ b/benchmarks/transformer/score_mod.py
@@ -10,7 +10,7 @@
 from contextlib import nullcontext
 from dataclasses import asdict, dataclass
 from functools import partial, wraps
-from typing import Literal, Optional, Union
+from typing import Literal
 
 import numpy as np
 from config_utils import heads_input_type, load_config_file, print_default_config
@@ -182,8 +182,8 @@ class Times:
 @dataclass(frozen=True)
 class ExperimentResults:
     fwd_time: float
-    bwd_time: Optional[float]
-    sparsity: Optional[float] = None
+    bwd_time: float | None
+    sparsity: float | None = None
 
 
 @dataclass(frozen=True)
@@ -270,7 +270,7 @@ def generate_jagged_inputs(
     B, Hq, M, Hkv, N, D = shape
 
     def offsets_to_lengths(
-        offsets: torch.Tensor, device: Union[str, torch.device]
+        offsets: torch.Tensor, device: str | torch.device
     ) -> torch.tensor:
         """Converts a list of offsets to a list of lengths. Reverse op of attn_gym.masks.document_mask.length_to_offsets
 
@@ -650,7 +650,7 @@ def get_average_speedups(results: list[Experiment], type: str, backend: str):
     return table_data
 
 
-def print_results(results: list[Experiment], save_path: Optional[str] = None):
+def print_results(results: list[Experiment], save_path: str | None = None):
     table_data = defaultdict(list)
     for experiment in results:
         backends = experiment.config.backends + ["flex"]
@@ -964,7 +964,7 @@ def generate_FA_callable(
         FA_kwargs["cu_seqlens_k"] = kwargs["offsets"].to(torch.int32)
 
         def offsets_to_lengths(
-            offsets: torch.Tensor, device: Union[str, torch.device]
+            offsets: torch.Tensor, device: str | torch.device
         ) -> torch.tensor:
             lengths = offsets[1:] - offsets[:-1]
             return lengths
@@ -1223,7 +1223,7 @@ def _output_json_for_dashboard(
     import math
     import platform
     from dataclasses import asdict, dataclass
-    from typing import Any, Optional
+    from typing import Any
 
     # Prepare headers and records for JSON output
     records = []
@@ -1269,7 +1269,7 @@ def _output_json_for_dashboard(
             @dataclass
             class BenchmarkInfo:
                 name: str
-                mode: Optional[str]
+                mode: str | None
                 dtype: str
                 extra_info: dict[str, Any]
 
@@ -1285,7 +1285,7 @@ class MetricInfo:
                 name: str
                 unit: str
                 benchmark_values: list[float]
-                target_value: Optional[float]
+                target_value: float | None
 
             @dataclass
             class BenchmarkRecord:
@@ -1430,10 +1430,10 @@ def main(
     backend: list[Backend] | None = None,
     max_autotune: bool = False,
     decoding: bool = False,
-    kv_size: Optional[list[int]] = None,
+    kv_size: list[int] | None = None,
     throughput: bool = True,
-    save_path: Optional[str] = None,
-    output_json_for_dashboard: Optional[str] = None,
+    save_path: str | None = None,
+    output_json_for_dashboard: str | None = None,
     benchmark_name: str = "PyTorch operator microbenchmark",
 ) -> None:
     """Run sweep over sizes and score mods for flex attention.
diff --git a/benchmarks/transformer/sdp.py b/benchmarks/transformer/sdp.py
index bc23aeb8e4a09..7ea088fbd1d11 100644
--- a/benchmarks/transformer/sdp.py
+++ b/benchmarks/transformer/sdp.py
@@ -5,7 +5,6 @@
 from dataclasses import dataclass
 from pathlib import Path
 from pprint import pprint
-from typing import Optional
 
 import numpy as np
 from prettytable import PrettyTable
@@ -26,7 +25,7 @@ class ExperimentConfig:
     max_sequence_len: int
     embed_dimension: int
     dtype: torch.dtype
-    pad_percentage: Optional[float]
+    pad_percentage: float | None
     enable_math: bool
     enable_flash: bool
     enable_mem_efficient: bool
@@ -65,9 +64,9 @@ def get_entry_names(cls) -> list[str]:
 @dataclass(frozen=True)
 class ExperimentResults:
     nn_mha_time: float
-    compiled_nn_mha_time: Optional[float]
+    compiled_nn_mha_time: float | None
     composite_mha_time: float
-    compiled_composite_mha_time: Optional[float]
+    compiled_composite_mha_time: float | None
 
     def get_entries(self) -> list:
         return [
@@ -306,7 +305,7 @@ def generate_experiments(
     return configs
 
 
-def main(save_path: Optional[Path]):
+def main(save_path: Path | None):
     seed = 123
     np.random.seed(seed)
     torch.manual_seed(seed)
diff --git a/buckbuild.bzl b/buckbuild.bzl
index 1745af1b06dd4..c34ab922dda33 100644
--- a/buckbuild.bzl
+++ b/buckbuild.bzl
@@ -71,15 +71,15 @@ def read_bool(section, field, default, required = True):
     else:
         fail("`{}:{}`: no value set".format(section, field))
 
-def _is_build_mode_dev():
-    if is_production_build_android():
-        # Android Prod builds
-        return False
+def _select_if_build_mode_dev(dev_value, default = []):
     if is_production_build_ios() or is_profile_build_ios():
-        # iOS Prod builds
-        return False
+        return default
+
+    return select({
+        "DEFAULT": default,
+        "ovr_config//build_mode:optimization[dev]": dev_value,
+    })
 
-    return True
 
 def _get_enable_lightweight_dispatch():
     return read_bool("pt", "enable_lightweight_dispatch", False)
@@ -93,10 +93,24 @@ def get_enable_mobile_dispatch_keys_trimming():
 def get_disable_per_op_profiling():
     return read_bool("pt", "disable_per_op_profiling", True)
 
-def get_strip_error_messages():
+def strip_error_messages_select(value, default = []):
     if IS_OSS:
-        return True  # always strip in OSS CI to expose potential issues
-    return read_bool("pt", "strip_error_messages", not _is_build_mode_dev())
+        return value  # always strip in OSS CI to expose potential issues
+    strip_error = read_bool("pt", "strip_error_messages", default = None, required = False)
+
+    if strip_error == None:
+
+        if is_production_build_ios() or is_profile_build_ios():
+            return value
+
+        return select({
+            "DEFAULT": default,
+            "ovr_config//build_mode:optimization[opt]": value,
+        })
+
+    if strip_error:
+        return value
+    return default
 
 def get_disable_warn():
     return read_bool("pt", "disable_warn", False)
@@ -249,9 +263,7 @@ _COMMON_PREPROCESSOR_FLAGS = [
     "-DNO_EXPORT",
 ] + (
     ["-DC10_MOBILE_TRIM_DISPATCH_KEYS"] if get_enable_mobile_dispatch_keys_trimming() else []
-) + (
-    ["-DSTRIP_ERROR_MESSAGES"] if get_strip_error_messages() else []
-) + (
+) + strip_error_messages_select(["-DSTRIP_ERROR_MESSAGES"]) + (
     ["-DDISABLE_WARN"] if get_disable_warn() else []
 )
 
@@ -282,9 +294,9 @@ def get_aten_preprocessor_flags():
         "-DUSE_RUY_QMATMUL",
     ]
     if get_disable_per_op_profiling():
-        ATEN_PREPROCESSOR_FLAGS.append("-DPYTORCH_DISABLE_PER_OP_PROFILING")
+        ATEN_PREPROCESSOR_FLAGS += ["-DPYTORCH_DISABLE_PER_OP_PROFILING"]
     if _get_enable_record_kernel_dtype():
-        ATEN_PREPROCESSOR_FLAGS.append("-DENABLE_RECORD_KERNEL_FUNCTION_DTYPE")
+        ATEN_PREPROCESSOR_FLAGS += ["-DENABLE_RECORD_KERNEL_FUNCTION_DTYPE"]
     return ATEN_PREPROCESSOR_FLAGS
 
 def get_pt_preprocessor_flags():
@@ -295,8 +307,7 @@ def get_pt_preprocessor_flags():
         "-DNO_CUDNN_DESTROY_HANDLE",
     ]
 
-    if _is_build_mode_dev():
-        PT_PREPROCESSOR_FLAGS.append("-DENABLE_PYTORCH_NON_PRODUCTION_BUILDS")
+    PT_PREPROCESSOR_FLAGS += _select_if_build_mode_dev(["-DENABLE_PYTORCH_NON_PRODUCTION_BUILDS"])
     return PT_PREPROCESSOR_FLAGS
 
 # This needs to be kept in sync with https://github.com/pytorch/pytorch/blob/release/1.9/torchgen/gen.py#L892  @lint-ignore
@@ -399,7 +410,6 @@ def get_aten_generated_files(enabled_backends):
         "core/TensorBody.h",
         "core/TensorMethods.cpp",
         "core/aten_interned_strings.h",
-        "core/enum_tag.h",
         "torch/csrc/inductor/aoti_torch/generated/c_shim_cpu.cpp",
     ] + get_aten_derived_type_srcs(enabled_backends)
 
@@ -905,6 +915,7 @@ def define_buck_targets(
             ("aten/src", "ATen/ops/*.h"),
             # ATen Base
             ("aten/src", "ATen/*.h"),
+            ("aten/src", "ATen/accelerator/*.h"),
             ("aten/src", "ATen/cpu/**/*.h"),
             ("aten/src", "ATen/detail/*.h"),
             ("aten/src", "ATen/functorch/**/*.h"),
@@ -932,6 +943,11 @@ def define_buck_targets(
             ("aten/src", "ATen/native/mkl/*.h"),
             ("aten/src", "ATen/native/mkldnn/*.h"),
         ]),
+        # ATen/core/enum_tag.h is a forwarding header that includes from
+        # torch/headeronly, so we need to export that dependency.
+        exported_deps = [
+            "//xplat/caffe2/torch/headeronly:torch_headeronly",
+        ],
         visibility = ["PUBLIC"],
         labels = labels,
     )
@@ -1219,7 +1235,6 @@ def define_buck_targets(
             "ViewMetaClasses.h": ":gen_aten[ViewMetaClasses.h]",
             "core/TensorBody.h": ":gen_aten[core/TensorBody.h]",
             "core/aten_interned_strings.h": ":gen_aten[core/aten_interned_strings.h]",
-            "core/enum_tag.h": ":gen_aten[core/enum_tag.h]",
         }),
         labels = labels,
     )
@@ -1519,7 +1534,12 @@ def define_buck_targets(
         srcs = [
             "torch/csrc/api/src/data/samplers/random.cpp",
             "torch/csrc/api/src/data/samplers/sequential.cpp",
+            "torch/csrc/api/src/optim/adagrad.cpp",
+            "torch/csrc/api/src/optim/adam.cpp",
+            "torch/csrc/api/src/optim/adamw.cpp",
+            "torch/csrc/api/src/optim/lbfgs.cpp",
             "torch/csrc/api/src/optim/optimizer.cpp",
+            "torch/csrc/api/src/optim/rmsprop.cpp",
             "torch/csrc/api/src/optim/serialize.cpp",
             "torch/csrc/api/src/optim/sgd.cpp",
             "torch/csrc/api/src/serialize/input-archive.cpp",
diff --git a/build.bzl b/build.bzl
index 0ce5b63e7b6ca..e2eadae6af78b 100644
--- a/build.bzl
+++ b/build.bzl
@@ -186,7 +186,6 @@ GENERATED_H_CORE = [
     "core/TensorBody.h",
     "MethodOperators.h",
     "core/aten_interned_strings.h",
-    "core/enum_tag.h",
 ]
 
 GENERATED_H_CUDA = [
@@ -327,6 +326,7 @@ GENERATED_AUTOGRAD_CPP = [
 
 GENERATED_AOTI_CPP = [
     "torch/csrc/inductor/aoti_torch/generated/c_shim_cpu.cpp",
+    "torch/csrc/inductor/aoti_torch/generated/c_shim_aten.cpp",
 ]
 
 GENERATED_AOTI_CUDA_CPP = [
diff --git a/build_variables.bzl b/build_variables.bzl
index 2db2ce477464d..91c5923b6cd76 100644
--- a/build_variables.bzl
+++ b/build_variables.bzl
@@ -112,6 +112,7 @@ libtorch_profiler_sources = [
     "torch/csrc/profiler/standalone/itt_observer.cpp",
     "torch/csrc/profiler/standalone/nvtx_observer.cpp",
     "torch/csrc/profiler/standalone/privateuse1_observer.cpp",
+    "torch/csrc/profiler/standalone/privateuse1_profiler.cpp",
     "torch/csrc/profiler/stubs/base.cpp",
     "torch/csrc/profiler/orchestration/vulkan.cpp",
     "torch/csrc/profiler/perf.cpp",
@@ -504,6 +505,7 @@ libtorch_distributed_base_sources = [
     "torch/csrc/distributed/c10d/Functional.cpp",
     "torch/csrc/distributed/c10d/GlooDeviceFactory.cpp",
     "torch/csrc/distributed/c10d/GroupRegistry.cpp",
+    "torch/csrc/distributed/c10d/NanCheck.cpp",
     "torch/csrc/distributed/c10d/Ops.cpp",
     "torch/csrc/distributed/c10d/ParamCommsUtils.cpp",
     "torch/csrc/distributed/c10d/PrefixStore.cpp",
@@ -689,7 +691,7 @@ libtorch_lite_cmake_sources = sorted(
     torch_mobile_core,
 )
 
-libtorch_cmake_sources = libtorch_core_sources + libtorch_core_jit_sources + libtorch_nativert_sources
+libtorch_cmake_sources = libtorch_core_sources + libtorch_core_jit_sources
 
 libtorch_extra_sources = libtorch_core_jit_sources + [
     "torch/csrc/autograd/TraceTypeManual.cpp",
@@ -740,6 +742,7 @@ libtorch_cuda_core_sources = [
     "torch/csrc/cuda/shim_common.cpp",
     "torch/csrc/inductor/aoti_runner/model_container_runner_cuda.cpp",
     "torch/csrc/inductor/aoti_torch/shim_cuda.cpp",
+    "torch/csrc/inductor/inductor_ops_gpu.cpp",
     "torch/csrc/jit/codegen/fuser/cuda/fused_kernel.cpp",
     "torch/csrc/profiler/stubs/cuda.cpp",
     "torch/csrc/autograd/functions/comm.cpp",
@@ -774,6 +777,7 @@ libtorch_cuda_distributed_extra_sources = [
     "torch/csrc/distributed/c10d/symm_mem/CudaDMAConnectivity.cpp",
     "torch/csrc/distributed/c10d/symm_mem/NCCLSymmetricMemory.cu",
     "torch/csrc/distributed/c10d/symm_mem/nccl_extension.cu",
+    "torch/csrc/distributed/c10d/symm_mem/ops/nccl_reduce_scatter_offset.cu",
     "torch/csrc/distributed/c10d/symm_mem/intra_node_comm.cpp",
     "torch/csrc/distributed/c10d/symm_mem/intra_node_comm.cu",
     "torch/csrc/distributed/c10d/symm_mem/cuda_mem_pool.cpp",
@@ -784,7 +788,7 @@ libtorch_nvshmem_sources = [
     "torch/csrc/distributed/c10d/cuda/utils.cpp",
     "torch/csrc/distributed/c10d/symm_mem/CUDASymmetricMemoryUtils.cpp",
     "torch/csrc/distributed/c10d/symm_mem/nvshmem_extension.cu",
-    "torch/csrc/distributed/c10d/symm_mem/NVSHMEMSymmetricMemory.cu",
+    "torch/csrc/distributed/c10d/symm_mem/NVSHMEMSymmetricMemory.cpp",
 ]
 
 libtorch_cuda_distributed_sources = libtorch_cuda_distributed_base_sources + libtorch_cuda_distributed_extra_sources
@@ -1157,6 +1161,7 @@ aten_cpu_source_non_codegen_list = [
     "aten/src/ATen/LegacyVmapMode.cpp",
     "aten/src/ATen/LegacyVmapTransforms.cpp",
     "aten/src/ATen/core/BackendSelectFallbackKernel.cpp",
+    "aten/src/ATen/core/CachingHostAllocator.cpp",
     "aten/src/ATen/core/DeprecatedTypeProperties.cpp",
     "aten/src/ATen/core/DeprecatedTypePropertiesRegistry.cpp",
     "aten/src/ATen/core/Dict.cpp",
@@ -1170,7 +1175,6 @@ aten_cpu_source_non_codegen_list = [
     "aten/src/ATen/core/Tensor.cpp",
     "aten/src/ATen/core/VariableFallbackKernel.cpp",
     "aten/src/ATen/core/VariableHooksInterface.cpp",
-    "aten/src/ATen/core/Vitals.cpp",
     "aten/src/ATen/core/boxing/KernelFunction.cpp",
     "aten/src/ATen/core/custom_class.cpp",
     "aten/src/ATen/core/dispatch/DispatchKeyExtractor.cpp",
@@ -1472,6 +1476,8 @@ aten_native_source_non_codegen_list = [
     "aten/src/ATen/native/RowwisePrune.cpp",
     "aten/src/ATen/native/SegmentReduce.cpp",
     "aten/src/ATen/native/Scalar.cpp",
+    "aten/src/ATen/native/ScaledBlas.cpp",
+    "aten/src/ATen/native/ScaledBlasUtils.cpp",
     "aten/src/ATen/native/SobolEngineOps.cpp",
     "aten/src/ATen/native/SobolEngineOpsUtils.cpp",
     "aten/src/ATen/native/SoftMax.cpp",
@@ -1499,6 +1505,7 @@ aten_native_source_non_codegen_list = [
     # "aten/src/ATen/native/UpSample.cpp",
     "aten/src/ATen/native/UpSampleBicubic2d.cpp",
     "aten/src/ATen/native/UpSampleBilinear2d.cpp",
+    "aten/src/ATen/native/UpSampleLanczos2d.cpp",
     "aten/src/ATen/native/UpSampleLinear1d.cpp",
     "aten/src/ATen/native/UpSampleNearest1d.cpp",
     "aten/src/ATen/native/UpSampleNearest2d.cpp",
diff --git a/c10/CMakeLists.txt b/c10/CMakeLists.txt
index f82e460cafc31..51dfc9f91f6bc 100644
--- a/c10/CMakeLists.txt
+++ b/c10/CMakeLists.txt
@@ -1,7 +1,7 @@
 cmake_minimum_required(VERSION 3.27 FATAL_ERROR)
 project(c10 CXX)
 
-set(CMAKE_CXX_STANDARD 17 CACHE STRING "The C++ standard whose features are requested to build this target.")
+set(CMAKE_CXX_STANDARD 20 CACHE STRING "The C++ standard whose features are requested to build this target.")
 set(CMAKE_EXPORT_COMPILE_COMMANDS ON)
 
 # Main build file for the C10 library.
diff --git a/c10/core/CachingDeviceAllocator.h b/c10/core/CachingDeviceAllocator.h
index 6800c79ab09ff..2e26e4091abe7 100644
--- a/c10/core/CachingDeviceAllocator.h
+++ b/c10/core/CachingDeviceAllocator.h
@@ -56,6 +56,9 @@ struct DeviceStats {
   // un-mapped and free memory.
   int64_t num_device_free = 0;
 
+  // COUNT: total number of allocations rejected by OOM preemption policy
+  int64_t num_oom_rejections = 0;
+
   // SIZE: maximum block size that is allowed to be split.
   int64_t max_split_size = 0;
 };
diff --git a/c10/core/DeviceType.cpp b/c10/core/DeviceType.cpp
index 907493981e117..4d5ea626a8800 100644
--- a/c10/core/DeviceType.cpp
+++ b/c10/core/DeviceType.cpp
@@ -62,9 +62,6 @@ std::string DeviceTypeName(DeviceType d, bool lower_case) {
           ". If you have recently updated the caffe2.proto file to add a new "
           "device type, did you forget to update the DeviceTypeName() "
           "function to reflect such recent changes?");
-      // The below code won't run but is needed to suppress some compiler
-      // warnings.
-      return "";
   }
 }
 
diff --git a/c10/core/ScalarType.h b/c10/core/ScalarType.h
index 040c6abb7d8e2..dc9d168f053e7 100644
--- a/c10/core/ScalarType.h
+++ b/c10/core/ScalarType.h
@@ -54,6 +54,31 @@ inline size_t elementSize(ScalarType t) {
 #undef CASE_ELEMENTSIZE_CASE
 }
 
+inline ScalarType opaqueScalarType(ScalarType t) {
+  auto esize = elementSize(t);
+  ScalarType result;
+  switch (esize) {
+    case 1:
+      result = kByte;
+      break;
+    case 2:
+      result = kUInt16;
+      break;
+    case 4:
+      result = kUInt32;
+      break;
+    case 8:
+      result = kUInt64;
+      break;
+    case 16:
+      result = kComplexDouble;
+      break;
+    default:
+      TORCH_CHECK(false, "Unknown ScalarType");
+  }
+  return result;
+}
+
 inline bool isIntegralType(ScalarType t, bool includeBool) {
   bool isIntegral =
       (t == ScalarType::Byte || t == ScalarType::Char || t == ScalarType::Int ||
diff --git a/c10/core/Storage.h b/c10/core/Storage.h
index e061375b8887b..efd4cc2e9cde3 100644
--- a/c10/core/Storage.h
+++ b/c10/core/Storage.h
@@ -152,6 +152,10 @@ struct C10_API Storage {
     storage_impl_->set_data_ptr_noswap(std::move(data_ptr));
   }
 
+  void swap_data_ptr(Storage& other) const {
+    storage_impl_->swap_data_ptr(*other.storage_impl_);
+  }
+
   DeviceType device_type() const {
     return storage_impl_->device_type();
   }
diff --git a/c10/core/StorageImpl.h b/c10/core/StorageImpl.h
index 7dfb8d310c33c..8479256c7ca86 100644
--- a/c10/core/StorageImpl.h
+++ b/c10/core/StorageImpl.h
@@ -181,6 +181,24 @@ struct C10_API StorageImpl : public c10::intrusive_ptr_target {
     refresh_has_data_ptr_check();
   }
 
+  void swap_data_ptr(StorageImpl& other) {
+    maybe_materialize_cow();
+    other.maybe_materialize_cow();
+    std::swap(data_ptr_, other.data_ptr_);
+    std::swap(size_bytes_, other.size_bytes_);
+    std::swap(
+        size_bytes_is_heap_allocated_, other.size_bytes_is_heap_allocated_);
+    std::swap(resizable_, other.resizable_);
+    std::swap(allocator_, other.allocator_);
+    std::swap(throw_on_immutable_data_ptr_, other.throw_on_immutable_data_ptr_);
+    std::swap(throw_on_mutable_data_ptr_, other.throw_on_mutable_data_ptr_);
+    std::swap(
+        warn_deprecated_on_mutable_data_ptr_,
+        other.warn_deprecated_on_mutable_data_ptr_);
+    refresh_has_data_ptr_check();
+    other.refresh_has_data_ptr_check();
+  }
+
   const void* data() const {
     if (C10_UNLIKELY(throw_on_immutable_data_ptr_)) {
       throw_data_ptr_access_error();
@@ -294,6 +312,14 @@ struct C10_API StorageImpl : public c10::intrusive_ptr_target {
     refresh_has_data_ptr_check();
   }
 
+  void clear_data_ptr_access_error_msg_() {
+    throw_on_immutable_data_ptr_ = false;
+    if (extra_meta_) {
+      extra_meta_->custom_data_ptr_error_msg_ = std::nullopt;
+    }
+    refresh_has_data_ptr_check();
+  }
+
   void set_throw_on_mutable_data_ptr() {
     throw_on_mutable_data_ptr_ = true;
     refresh_has_data_ptr_check();
diff --git a/c10/core/Stream.cpp b/c10/core/Stream.cpp
index 45339719a91b1..afb8f525dabe4 100644
--- a/c10/core/Stream.cpp
+++ b/c10/core/Stream.cpp
@@ -22,6 +22,12 @@ void Stream::synchronize() const {
   impl.synchronizeStream(*this);
 }
 
+// Return whether this stream is currently under graph capturing mode.
+bool Stream::is_capturing() const {
+  impl::VirtualGuardImpl impl{device_.type()};
+  return impl.isStreamCapturing(*this);
+}
+
 // Not very parsable, but I don't know a good compact syntax for streams.
 // Feel free to change this into something more compact if needed.
 std::ostream& operator<<(std::ostream& stream, const Stream& s) {
diff --git a/c10/core/Stream.h b/c10/core/Stream.h
index 3890faa05d9c4..08eb0eed3c6e1 100644
--- a/c10/core/Stream.h
+++ b/c10/core/Stream.h
@@ -134,6 +134,10 @@ class C10_API Stream final {
   // on this stream has completed running on the device.
   void synchronize() const;
 
+  // Return the stream is currently recording work for graph capture. True while
+  // the stream is in capture mode, false otherwise.
+  bool is_capturing() const;
+
   // The purpose of this function is to more conveniently permit binding
   // of Stream to and from Python.  Without packing, I have to setup a whole
   // class with two fields (device and stream id); with packing I can just
diff --git a/c10/core/SymInt.h b/c10/core/SymInt.h
index 0ee218c8cf54e..ad004a5b40352 100644
--- a/c10/core/SymInt.h
+++ b/c10/core/SymInt.h
@@ -54,13 +54,9 @@ class C10_API SymInt {
   // the negative value is -1; i.e., not user controlled)
   SymInt(Unchecked /*unused*/, int64_t d) : data_(d) {}
 
-  // TODO: these implementations are not optimal because they allocate a
-  // temporary and then use the move constructor/assignment
-  SymInt(const SymInt& s) : data_(0) {
+  SymInt(const SymInt& s) : data_(s.data_) {
     if (s.is_heap_allocated()) {
-      *this = SymInt(s.toSymNode());
-    } else {
-      data_ = s.data_;
+      c10::raw::intrusive_ptr::incref(s.toSymNodeImplUnowned());
     }
   }
   SymInt(SymInt&& s) noexcept : data_(s.data_) {
@@ -69,10 +65,10 @@ class C10_API SymInt {
 
   SymInt& operator=(const SymInt& s) {
     if (this != &s) {
+      release_();
+      data_ = s.data_;
       if (s.is_heap_allocated()) {
-        *this = SymInt(s.toSymNode());
-      } else {
-        data_ = s.data_;
+        c10::raw::intrusive_ptr::incref(s.toSymNodeImplUnowned());
       }
     }
     return *this;
diff --git a/c10/core/SymbolicShapeMeta.h b/c10/core/SymbolicShapeMeta.h
index 0820038968a8e..59c3147c878f7 100644
--- a/c10/core/SymbolicShapeMeta.h
+++ b/c10/core/SymbolicShapeMeta.h
@@ -141,7 +141,7 @@ class C10_API SymbolicShapeMeta {
     available_.fetch_or(is_contiguous_avail);
   }
   void assume_channels_last_contiguous(SymBool val = true) {
-    is_contiguous_ = std::move(val);
+    is_channels_last_contiguous_ = std::move(val);
     available_.fetch_or(is_channels_last_contiguous_avail);
   }
   void assume_channels_last_3d_contiguous(SymBool val = true) {
diff --git a/c10/core/impl/DeviceGuardImplInterface.cpp b/c10/core/impl/DeviceGuardImplInterface.cpp
index 428ea63c04151..6d69c31e3e1e6 100644
--- a/c10/core/impl/DeviceGuardImplInterface.cpp
+++ b/c10/core/impl/DeviceGuardImplInterface.cpp
@@ -21,11 +21,6 @@ DeviceGuardImplRegistrar::DeviceGuardImplRegistrar(
   registerDeviceGuard(type, impl);
 }
 
-namespace {
-thread_local std::unique_ptr<DeviceGuardImplInterface> tls_fake_device_guard =
-    nullptr;
-} // namespace
-
 void ensureCUDADeviceGuardSet() {
   constexpr auto cuda_idx = static_cast<std::size_t>(DeviceType::CUDA);
 
@@ -38,8 +33,8 @@ void ensureCUDADeviceGuardSet() {
     // In following cases, we override CUDA guard interface with a no-op
     // device guard. When p->deviceCount() == 0, cuda build is enabled, but no
     // cuda devices available.
-    tls_fake_device_guard = std::make_unique<FakeGuardImpl<DeviceType::CUDA>>();
-    device_guard_impl_registry[cuda_idx].store(tls_fake_device_guard.get());
+    static FakeGuardImpl<DeviceType::CUDA> fake_cuda_guard;
+    device_guard_impl_registry[cuda_idx].store(&fake_cuda_guard);
   }
 }
 
diff --git a/c10/core/impl/DeviceGuardImplInterface.h b/c10/core/impl/DeviceGuardImplInterface.h
index 01ef32163bb5c..ef7a628835d97 100644
--- a/c10/core/impl/DeviceGuardImplInterface.h
+++ b/c10/core/impl/DeviceGuardImplInterface.h
@@ -228,6 +228,13 @@ struct C10_API DeviceGuardImplInterface {
     TORCH_CHECK(false, "Backend doesn't support synchronizing streams.");
   }
 
+  /**
+   * Return true if this stream is currently recording work for graph capture.
+   */
+  virtual bool isStreamCapturing(const Stream& /*stream*/) const {
+    TORCH_CHECK(false, "Backend doesn't support stream capture query.");
+  }
+
   /**
    * Wait (by blocking the calling thread) until all the work previously
    * recorded on the event has completed running on the device.
diff --git a/c10/core/impl/LocalDispatchKeySet.cpp b/c10/core/impl/LocalDispatchKeySet.cpp
index 0b3d87fce2410..c8277f19acc7c 100644
--- a/c10/core/impl/LocalDispatchKeySet.cpp
+++ b/c10/core/impl/LocalDispatchKeySet.cpp
@@ -46,29 +46,21 @@ void _force_tls_local_dispatch_key_set(LocalDispatchKeySet key_set) {
 // RAII API
 
 IncludeDispatchKeyGuard::IncludeDispatchKeyGuard(DispatchKeySet include)
-    : tls_(&raw_local_dispatch_key_set), include_(include - tls_->included()) {
-  if (!include_.empty()) {
-    tls_->set_included(tls_->included() | include_);
-  }
+    : tls_(&raw_local_dispatch_key_set), saved_state_(tls_->included()) {
+  tls_->set_included(saved_state_ | include);
 }
 
 IncludeDispatchKeyGuard::~IncludeDispatchKeyGuard() {
-  if (!include_.empty()) {
-    tls_->set_included(tls_->included() - include_);
-  }
+  tls_->set_included(saved_state_);
 }
 
 ExcludeDispatchKeyGuard::ExcludeDispatchKeyGuard(DispatchKeySet exclude)
-    : tls_(&raw_local_dispatch_key_set), exclude_(exclude - tls_->excluded()) {
-  if (!exclude_.empty()) {
-    tls_->set_excluded(tls_->excluded() | exclude_);
-  }
+    : tls_(&raw_local_dispatch_key_set), saved_state_(tls_->excluded()) {
+  tls_->set_excluded(saved_state_ | exclude);
 }
 
 ExcludeDispatchKeyGuard::~ExcludeDispatchKeyGuard() {
-  if (!exclude_.empty()) {
-    tls_->set_excluded(tls_->excluded() - exclude_);
-  }
+  tls_->set_excluded(saved_state_);
 }
 
 // Non-RAII API
diff --git a/c10/core/impl/LocalDispatchKeySet.h b/c10/core/impl/LocalDispatchKeySet.h
index bba089bb2ad11..bd8fe811b2c06 100644
--- a/c10/core/impl/LocalDispatchKeySet.h
+++ b/c10/core/impl/LocalDispatchKeySet.h
@@ -94,7 +94,7 @@ class C10_API IncludeDispatchKeyGuard {
   // A little micro-optimization to save us from tls_get_addr call
   // on destruction
   PODLocalDispatchKeySet* tls_;
-  DispatchKeySet include_;
+  DispatchKeySet saved_state_;
 };
 
 class C10_API ExcludeDispatchKeyGuard {
@@ -112,7 +112,7 @@ class C10_API ExcludeDispatchKeyGuard {
   // A little micro-optimization to save us from tls_get_addr call
   // on destruction
   PODLocalDispatchKeySet* tls_;
-  DispatchKeySet exclude_;
+  DispatchKeySet saved_state_;
 };
 
 struct C10_API ForceDispatchKeyGuard {
diff --git a/c10/core/impl/VirtualGuardImpl.h b/c10/core/impl/VirtualGuardImpl.h
index 6c86a71676e55..2444e4e03b3ee 100644
--- a/c10/core/impl/VirtualGuardImpl.h
+++ b/c10/core/impl/VirtualGuardImpl.h
@@ -89,6 +89,9 @@ class VirtualGuardImpl final : public DeviceGuardImplInterface {
   void synchronizeStream(const Stream& stream) const override {
     impl_->synchronizeStream(stream);
   }
+  bool isStreamCapturing(const Stream& stream) const override {
+    return impl_->isStreamCapturing(stream);
+  }
 
   void recordDataPtrOnStream(const c10::DataPtr& data_ptr, const Stream& stream)
       const override {
diff --git a/c10/cuda/CUDAAllocatorConfig.cpp b/c10/cuda/CUDAAllocatorConfig.cpp
index 5414d838cd8c4..f6fb7c9d6d9a8 100644
--- a/c10/cuda/CUDAAllocatorConfig.cpp
+++ b/c10/cuda/CUDAAllocatorConfig.cpp
@@ -109,6 +109,12 @@ void CUDAAllocatorConfig::parseArgs(const std::string& env) {
     } else if (key == "per_process_memory_fraction") {
       i = parsePerProcessMemoryFraction(tokenizer, i);
       used_native_specific_option = true;
+    } else if (key == "pinned_free_catch_all") {
+      i = parsePinnedFreeCatchAll(tokenizer, i);
+      used_native_specific_option = true;
+    } else if (key == "throw_on_cudamalloc_oom") {
+      i = parseThrowOnCudaMallocOom(tokenizer, i);
+      used_native_specific_option = true;
     } else {
       const auto& keys =
           c10::CachingAllocator::AcceleratorAllocatorConfig::getKeys();
@@ -192,6 +198,25 @@ size_t CUDAAllocatorConfig::parsePinnedReserveSegmentSize(
   return i;
 }
 
+size_t CUDAAllocatorConfig::parsePinnedFreeCatchAll(
+    const c10::CachingAllocator::ConfigTokenizer& tokenizer,
+    size_t i) {
+  tokenizer.checkToken(++i, ":");
+  m_pinned_free_catch_all = tokenizer.toBool(++i);
+  return i;
+}
+
+size_t CUDAAllocatorConfig::parseThrowOnCudaMallocOom(
+    const c10::CachingAllocator::ConfigTokenizer& tokenizer,
+    size_t i) {
+  // Format: throw_on_cudamalloc_oom:true or throw_on_cudamalloc_oom:false
+  // When enabled, throws OOM error before calling cudaMalloc if the allocation
+  // would likely fail due to insufficient memory.
+  tokenizer.checkToken(++i, ":");
+  m_throw_on_cudamalloc_oom = tokenizer.toBool(++i);
+  return i;
+}
+
 REGISTER_ALLOCATOR_CONFIG_PARSE_HOOK(CUDAAllocatorConfig)
 
 } // namespace c10::cuda::CUDACachingAllocator
diff --git a/c10/cuda/CUDAAllocatorConfig.h b/c10/cuda/CUDAAllocatorConfig.h
index 4e6097a406bc2..928384a874394 100644
--- a/c10/cuda/CUDAAllocatorConfig.h
+++ b/c10/cuda/CUDAAllocatorConfig.h
@@ -34,7 +34,8 @@ class C10_CUDA_API CUDAAllocatorConfig {
   static bool expandable_segments() {
     bool enabled = c10::CachingAllocator::AcceleratorAllocatorConfig::
         use_expandable_segments();
-#ifndef PYTORCH_C10_DRIVER_API_SUPPORTED
+#if !defined(PYTORCH_C10_DRIVER_API_SUPPORTED) && \
+    (!defined(USE_ROCM) || (ROCM_VERSION < 70000))
     if (enabled) {
       TORCH_WARN_ONCE("expandable_segments not supported on this platform")
     }
@@ -65,6 +66,13 @@ class C10_CUDA_API CUDAAllocatorConfig {
     return instance().m_per_process_memory_fraction;
   }
 
+  // When enabled, throws OOM error before calling cudaMalloc if the allocation
+  // would likely fail due to insufficient memory. This provides early failure
+  // with clear error messages instead of letting cudaMalloc fail.
+  static bool throw_on_cudamalloc_oom() {
+    return instance().m_throw_on_cudamalloc_oom;
+  }
+
   /** Pinned memory allocator settings */
   static bool pinned_use_cuda_host_register() {
     return instance().m_pinned_use_cuda_host_register;
@@ -92,6 +100,10 @@ class C10_CUDA_API CUDAAllocatorConfig {
     return 128;
   }
 
+  static bool pinned_free_catch_all() {
+    return instance().m_pinned_free_catch_all;
+  }
+
   C10_DEPRECATED_MESSAGE(
       "c10::cuda::CUDACachingAllocator::CUDAAllocatorConfig::roundup_power2_divisions() is deprecated. Please use c10::CachingAllocator::AcceleratorAllocatorConfig::roundup_power2_divisions() instead.")
   static size_t roundup_power2_divisions(size_t size) {
@@ -157,7 +169,9 @@ class C10_CUDA_API CUDAAllocatorConfig {
         "graph_capture_record_stream_reuse",
         "pinned_reserve_segment_size_mb",
         "pinned_num_register_threads",
-        "per_process_memory_fraction"};
+        "per_process_memory_fraction",
+        "pinned_free_catch_all",
+        "throw_on_cudamalloc_oom"};
     return keys;
   }
 
@@ -185,6 +199,12 @@ class C10_CUDA_API CUDAAllocatorConfig {
   double parsePerProcessMemoryFraction(
       const c10::CachingAllocator::ConfigTokenizer& tokenizer,
       size_t i);
+  size_t parsePinnedFreeCatchAll(
+      const c10::CachingAllocator::ConfigTokenizer& tokenizer,
+      size_t i);
+  size_t parseThrowOnCudaMallocOom(
+      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};
@@ -198,6 +218,10 @@ class C10_CUDA_API CUDAAllocatorConfig {
   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};
+  std::atomic<bool> m_pinned_free_catch_all{false};
+  // When true, throw OOM error before calling cudaMalloc if allocation would
+  // fail
+  std::atomic<bool> m_throw_on_cudamalloc_oom{false};
 };
 
 // Keep this for backwards compatibility
diff --git a/c10/cuda/CUDACachingAllocator.cpp b/c10/cuda/CUDACachingAllocator.cpp
index 2ab4effc7853d..774c79102b8d6 100644
--- a/c10/cuda/CUDACachingAllocator.cpp
+++ b/c10/cuda/CUDACachingAllocator.cpp
@@ -5,6 +5,7 @@
 #include <c10/cuda/CUDAException.h>
 #include <c10/cuda/CUDAFunctions.h>
 #include <c10/cuda/CUDAGuard.h>
+#include <c10/cuda/PeerToPeerAccess.h>
 #include <c10/util/Gauge.h>
 #include <c10/util/Logging.h>
 #include <c10/util/ScopeExit.h>
@@ -17,16 +18,23 @@
 #include <c10/util/llvmMathExtras.h>
 #include <c10/util/static_tracepoint.h>
 
-#if !defined(USE_ROCM) && defined(PYTORCH_C10_DRIVER_API_SUPPORTED)
+#if defined(PYTORCH_C10_DRIVER_API_SUPPORTED) || defined(USE_ROCM)
+#if defined(PYTORCH_C10_DRIVER_API_SUPPORTED)
 #include <c10/cuda/driver_api.h>
+#endif
+#ifndef _WIN32
 #include <sys/syscall.h>
 #include <sys/types.h>
 #include <unistd.h>
+#else
+#include <process.h>
+#endif
 #endif
 
 #include <c10/util/Exception.h>
 #include <cuda_runtime_api.h>
 #include <algorithm>
+#include <atomic>
 #include <cstddef>
 #include <cstdint>
 #include <deque>
@@ -133,7 +141,7 @@ namespace Native {
  */
 
 // counter to track order for Mempool Registration
-thread_local int32_t registration_counter_global = -1;
+std::atomic<int32_t> registration_counter_global{-1};
 
 static char SHAREABLE_HANDLE_VERSION = 2;
 enum ShareableHandleType : char {
@@ -158,12 +166,12 @@ void decrease_stat_array(
 struct Block;
 struct PrivatePool;
 typedef bool (*Comparison)(const Block*, const Block*);
-static bool BlockComparatorSize(const Block* a, const Block* b);
+static bool BlockComparatorRegistrationCounter(const Block* a, const Block* b);
 static bool BlockComparatorAddress(const Block* a, const Block* b);
 
 struct BlockPool {
   BlockPool(bool small, PrivatePool* private_pool = nullptr)
-      : blocks(BlockComparatorSize),
+      : blocks(BlockComparatorRegistrationCounter),
         unmapped(BlockComparatorAddress),
         is_small(small),
         owner_PrivatePool(private_pool) {}
@@ -226,14 +234,16 @@ struct Block {
         requested_size(0),
         pool(pool),
         ptr(ptr) {
-    registration_counter = ++registration_counter_global;
+    registration_counter =
+        registration_counter_global.fetch_add(1, std::memory_order_relaxed) + 1;
   }
 
   // constructor for search key
+  // Use the default value for registration_counter and not modify
+  // registration_counter_global, because the search key is just a
+  // dummy placeholder.
   Block(c10::DeviceIndex device, cudaStream_t stream, size_t size)
-      : device(device), stream(stream), size(size), requested_size(0) {
-    registration_counter = ++registration_counter_global;
-  }
+      : device(device), stream(stream), size(size), requested_size(0) {}
 
   size_t gc_count() {
     TORCH_INTERNAL_ASSERT(pool);
@@ -269,7 +279,8 @@ struct SegmentRange {
   SegmentRange(void* p, size_t s) : ptr(static_cast<char*>(p)), size(s) {}
 };
 
-#if !defined(USE_ROCM) && defined(PYTORCH_C10_DRIVER_API_SUPPORTED)
+#if !defined(USE_ROCM) && defined(PYTORCH_C10_DRIVER_API_SUPPORTED) || \
+    defined(USE_ROCM)
 
 /*
 Note [Expandable Segments]
@@ -383,8 +394,13 @@ struct ExpandableSegment {
     // This allows for some cases where we have to unmap pages earlier in the
     // segment to put them at the end.
     max_handles_ = numSegments(prop.totalGlobalMem + prop.totalGlobalMem / 8);
+#ifdef USE_ROCM
+    C10_CUDA_CHECK(hipMemAddressReserve(
+        &ptr_, segment_size_ * max_handles_, 0ULL, 0, 0ULL));
+#else
     C10_CUDA_DRIVER_CHECK(DriverAPI::get()->cuMemAddressReserve_(
         &ptr_, segment_size_ * max_handles_, 0ULL, 0, 0ULL));
+#endif
   }
   ExpandableSegment(const ExpandableSegment&) = delete;
   ExpandableSegment(ExpandableSegment&&) = delete;
@@ -408,20 +424,22 @@ struct ExpandableSegment {
     // if it fails, use posix file handle
     if (CUDAAllocatorConfig::expandable_segments_handle_type() ==
         Expandable_Segments_Handle_Type::UNSPECIFIED) {
+#ifndef USE_ROCM
       CUDAAllocatorConfig::set_expandable_segments_handle_type(
           Expandable_Segments_Handle_Type::FABRIC_HANDLE);
       auto output = map(range);
       if (output.ptr != nullptr) {
         return output;
       }
+#endif
       // if fabric handle is not supported, use posix file handle.
       CUDAAllocatorConfig::set_expandable_segments_handle_type(
           Expandable_Segments_Handle_Type::POSIX_FD);
       return map(range);
     }
 
-    while (end > handles_.size()) {
-      handles_.emplace_back(std::nullopt);
+    if (end > handles_.size()) {
+      handles_.resize(end);
     }
     for (auto i : c10::irange(begin, end)) {
       TORCH_INTERNAL_ASSERT(!handles_.at(i));
@@ -445,33 +463,60 @@ struct ExpandableSegment {
       if (enable_ipc_handles) {
         if (CUDAAllocatorConfig::expandable_segments_handle_type() !=
             Expandable_Segments_Handle_Type::FABRIC_HANDLE) {
+#ifdef USE_ROCM
+          prop.requestedHandleType = hipMemHandleTypePosixFileDescriptor;
+#else
           prop.requestedHandleTypes = CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR;
+#endif
         } else {
+#ifndef USE_ROCM
           prop.requestedHandleTypes = CU_MEM_HANDLE_TYPE_FABRIC;
+#endif
         }
       }
       int flag = 0;
+#ifndef USE_ROCM
       C10_CUDA_DRIVER_CHECK(DriverAPI::get()->cuDeviceGetAttribute_(
           &flag,
           CU_DEVICE_ATTRIBUTE_GPU_DIRECT_RDMA_WITH_CUDA_VMM_SUPPORTED,
           device_));
+#endif
       if (flag)
         prop.allocFlags.gpuDirectRDMACapable = 1;
       prop.location.type = CU_MEM_LOCATION_TYPE_DEVICE;
       // NOLINTNEXTLINE(bugprone-signed-char-misuse)
       prop.location.id = static_cast<int>(device_);
+#ifdef USE_ROCM
+      auto status = hipMemCreate(&handle, segment_size_, &prop, 0);
+#else
       auto status =
           DriverAPI::get()->cuMemCreate_(&handle, segment_size_, &prop, 0);
+#endif
       if (status != CUDA_SUCCESS) {
         if (status == CUDA_ERROR_OUT_OF_MEMORY) {
+#ifdef USE_ROCM
+          // hipMemCreate above returned hipErrorOutOfMemory and treated it
+          // like a sticky runtime error. Which means we need to clear it.
+          // Unlike the corresponding CUDA Driver API.
+          (void)hipGetLastError();
+#endif
           for (auto j : c10::irange(begin, i)) {
             // NOLINTNEXTLINE(bugprone-unchecked-optional-access)
             auto h = handles_.at(j).value();
             handles_.at(j) = std::nullopt;
+#ifdef USE_ROCM
+            C10_CUDA_CHECK(hipMemRelease(h.handle));
+#else
             C10_CUDA_DRIVER_CHECK(DriverAPI::get()->cuMemRelease_(h.handle));
+#endif
           }
           trimHandles();
           return rangeFromHandles(begin, begin);
+#ifdef USE_ROCM
+        } else {
+          C10_CUDA_CHECK(status);
+        }
+#else
         } else if (
             CUDAAllocatorConfig::expandable_segments_handle_type() ==
             Expandable_Segments_Handle_Type::FABRIC_HANDLE) {
@@ -487,6 +532,7 @@ struct ExpandableSegment {
         } else {
           C10_CUDA_DRIVER_CHECK(status);
         }
+#endif
       }
       handles_.at(i) = Handle{handle, std::nullopt};
     }
@@ -522,7 +568,11 @@ struct ExpandableSegment {
     // thereby ensuring that the handle can be correctly matched in
     // ipcMemHandle_to_devptr.
     ShareHeader header{};
+#ifdef _WIN32
+    header.pid = _getpid();
+#else
     header.pid = getpid();
+#endif
     header.segment_size = segment_size_;
     header.num_handles = end - begin;
 
@@ -534,8 +584,13 @@ struct ExpandableSegment {
           Expandable_Segments_Handle_Type::FABRIC_HANDLE) {
         if (!handle.shareable_handle) {
           int fd = 0;
+#ifdef USE_ROCM
+          C10_CUDA_CHECK(hipMemExportToShareableHandle(
+              &fd, handle.handle, hipMemHandleTypePosixFileDescriptor, 0));
+#else
           C10_CUDA_DRIVER_CHECK(DriverAPI::get()->cuMemExportToShareableHandle_(
               &fd, handle.handle, CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR, 0));
+#endif
           handle.shareable_handle = fd;
           LOG(INFO) << "use posix fd to share expandable segments.";
         }
@@ -546,6 +601,10 @@ struct ExpandableSegment {
             reinterpret_cast<const char*>(&*handle.shareable_handle),
             sizeof(int));
       } else {
+#ifdef USE_ROCM
+        TORCH_INTERNAL_ASSERT(
+            false, "expandable segment with fabric handle not supported");
+#else
         if (!handle.shareable_handle) {
           CUmemFabricHandle fabric_handle;
           C10_CUDA_DRIVER_CHECK(DriverAPI::get()->cuMemExportToShareableHandle_(
@@ -559,6 +618,7 @@ struct ExpandableSegment {
         buf.write(
             reinterpret_cast<const char*>(&*handle.shareable_handle),
             sizeof(CUmemFabricHandle));
+#endif
       }
     }
     return rangeFromHandles(begin, end);
@@ -574,14 +634,20 @@ struct ExpandableSegment {
         device, std::nullopt, header.segment_size, std::move(peers));
 // older build setups (e.g. multiwheels) do not have this syscall, added 2020
 // but the kernel on the system might still support it.
+#ifndef _WIN32
 #ifndef SYS_pidfd_open
 #define SYS_pidfd_open 434
 #endif
 #ifndef SYS_pidfd_getfd
 #define SYS_pidfd_getfd 438
 #endif
+#endif // !_WIN32
     if (CUDAAllocatorConfig::expandable_segments_handle_type() !=
         Expandable_Segments_Handle_Type::FABRIC_HANDLE) {
+#ifdef _WIN32
+      TORCH_CHECK(
+          false, "IPC expandable segments are not supported on Windows");
+#else
       auto pidfd = syscall(SYS_pidfd_open, header.pid, 0);
       TORCH_CHECK(
           pidfd != -1 || errno != ENOSYS,
@@ -597,9 +663,13 @@ struct ExpandableSegment {
           auto err = errno;
           close(static_cast<int>(pidfd));
           for (auto& h : segment->handles_) {
+#ifdef USE_ROCM
+            C10_CUDA_CHECK(hipMemRelease(h.value().handle));
+#else
             C10_CUDA_DRIVER_CHECK(
                 // NOLINTNEXTLINE(bugprone-unchecked-optional-access)
                 DriverAPI::get()->cuMemRelease_(h.value().handle));
+#endif
             h = std::nullopt;
           }
           TORCH_CHECK(
@@ -609,31 +679,56 @@ struct ExpandableSegment {
           TORCH_CHECK(false, "pidfd_getfd: ", c10::utils::str_error(err));
         }
         CUmemGenericAllocationHandle handle = 0;
-        C10_CUDA_DRIVER_CHECK(DriverAPI::get()->cuMemImportFromShareableHandle_(
-            &handle,
-            // NOLINTNEXTLINE(performance-no-int-to-ptr)
-            (void*)(uintptr_t)myfd,
-            CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR));
+#ifdef USE_ROCM
+#if ROCM_VERSION >= 70100
+        void* myfd_handle =
+            reinterpret_cast<void*>(static_cast<uintptr_t>(myfd));
+#else
+        void* myfd_handle = (void*)(uintptr_t)&myfd;
+#endif
+        C10_CUDA_CHECK(hipMemImportFromShareableHandle(
+            &handle, myfd_handle, hipMemHandleTypePosixFileDescriptor));
+#else
+        C10_CUDA_DRIVER_CHECK_MSG(
+            DriverAPI::get()->cuMemImportFromShareableHandle_(
+                &handle,
+                // NOLINTNEXTLINE(performance-no-int-to-ptr)
+                (void*)(uintptr_t)myfd,
+                CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR),
+            " fabric_info: {",
+            get_nvml_fabric_info(device),
+            "}");
+#endif
         LOG(INFO) << "use posix fd to import expandable segments.";
         close(static_cast<int>(myfd));
         segment->handles_.emplace_back(Handle{handle, std::nullopt});
       }
       close(static_cast<int>(pidfd));
+#endif // !_WIN32
     } else {
+#ifdef USE_ROCM
+      TORCH_INTERNAL_ASSERT(
+          false, "expandable segment with fabric handle not supported");
+#else
       for (auto i : c10::irange(header.num_handles)) {
         (void)i;
         CUmemFabricHandle fabric_handle;
         buf.read(
             reinterpret_cast<char*>(&fabric_handle), sizeof(CUmemFabricHandle));
         CUmemGenericAllocationHandle handle = 0;
-        C10_CUDA_DRIVER_CHECK(DriverAPI::get()->cuMemImportFromShareableHandle_(
-            &handle,
-            // NOLINTNEXTLINE(performance-no-int-to-ptr)
-            (void*)&fabric_handle,
-            CU_MEM_HANDLE_TYPE_FABRIC));
+        C10_CUDA_DRIVER_CHECK_MSG(
+            DriverAPI::get()->cuMemImportFromShareableHandle_(
+                &handle,
+                // NOLINTNEXTLINE(performance-no-int-to-ptr)
+                (void*)&fabric_handle,
+                CU_MEM_HANDLE_TYPE_FABRIC),
+            " fabric_info: {",
+            get_nvml_fabric_info(device),
+            "}");
         LOG(INFO) << "use fabric handle to import expandable segments.";
         segment->handles_.emplace_back(Handle{handle, std::nullopt});
       }
+#endif
     }
     segment->mapAndSetAccess(0, header.num_handles);
     return segment;
@@ -669,23 +764,55 @@ struct ExpandableSegment {
   ~ExpandableSegment() {
     forEachAllocatedRange(
         [&](size_t begin, size_t end) { unmapHandles(begin, end); });
+#ifdef USE_ROCM
+    C10_CUDA_CHECK(hipMemAddressFree(ptr_, segment_size_ * max_handles_));
+#else
     C10_CUDA_DRIVER_CHECK(DriverAPI::get()->cuMemAddressFree_(
         ptr_, segment_size_ * max_handles_));
+#endif
   }
 
  private:
   void setAccess(c10::DeviceIndex device, size_t begin, size_t end) {
-    CUmemAccessDesc desc;
-    desc.location.type = CU_MEM_LOCATION_TYPE_DEVICE;
+#if defined(USE_ROCM) && (ROCM_VERSION >= 70200)
+    constexpr int num_desc = 2;
+    CUmemAccessDesc desc[num_desc];
+    desc[1].location.type = CU_MEM_LOCATION_TYPE_HOST;
+    desc[1].location.id = 0; // ignored
+    desc[1].flags = CU_MEM_ACCESS_FLAGS_PROT_READWRITE;
+#else
+    constexpr int num_desc = 1;
+    CUmemAccessDesc desc[num_desc];
+#endif
+    desc[0].location.type = CU_MEM_LOCATION_TYPE_DEVICE;
     // NOLINTNEXTLINE(bugprone-signed-char-misuse)
-    desc.location.id = static_cast<int>(device);
-    desc.flags = CU_MEM_ACCESS_FLAGS_PROT_READWRITE;
+    desc[0].location.id = static_cast<int>(device);
+    desc[0].flags = CU_MEM_ACCESS_FLAGS_PROT_READWRITE;
+#ifdef USE_ROCM
+    C10_CUDA_CHECK(hipMemSetAccess(
+        ptr() + begin * segment_size_,
+        (end - begin) * segment_size_,
+        &desc[0],
+        num_desc));
+#else
     C10_CUDA_DRIVER_CHECK(DriverAPI::get()->cuMemSetAccess_(
-        ptr_ + begin * segment_size_, (end - begin) * segment_size_, &desc, 1));
+        ptr_ + begin * segment_size_,
+        (end - begin) * segment_size_,
+        &desc[0],
+        num_desc));
+#endif
   }
 
   void mapAndSetAccess(size_t begin, size_t end) {
     for (auto i : c10::irange(begin, end)) {
+#ifdef USE_ROCM
+      C10_CUDA_CHECK(hipMemMap(
+          ptr() + i * segment_size_,
+          segment_size_,
+          0,
+          handles_.at(i).value().handle,
+          0ULL));
+#else
       C10_CUDA_DRIVER_CHECK(DriverAPI::get()->cuMemMap_(
           ptr_ + i * segment_size_,
           segment_size_,
@@ -693,6 +820,7 @@ struct ExpandableSegment {
           // NOLINTNEXTLINE(bugprone-unchecked-optional-access)
           handles_.at(i).value().handle,
           0ULL));
+#endif
     }
     mapped_size_ += (end - begin) * segment_size_;
     setAccess(device_, begin, end);
@@ -719,19 +847,31 @@ struct ExpandableSegment {
       // NOLINTNEXTLINE(bugprone-unchecked-optional-access)
       Handle h = handles_.at(i).value();
       handles_.at(i) = std::nullopt;
+#ifdef USE_ROCM
+      C10_CUDA_CHECK(hipMemUnmap(ptr() + segment_size_ * i, segment_size_));
+#else
       C10_CUDA_DRIVER_CHECK(DriverAPI::get()->cuMemUnmap_(
           ptr_ + segment_size_ * i, segment_size_));
+#endif
       if (h.shareable_handle) {
+#ifndef _WIN32
         close(std::get<int>(*h.shareable_handle));
+#endif
       }
+#ifdef USE_ROCM
+      C10_CUDA_CHECK(hipMemRelease(h.handle));
+#else
       C10_CUDA_DRIVER_CHECK(DriverAPI::get()->cuMemRelease_(h.handle));
+#endif
     }
     trimHandles();
   }
   void trimHandles() {
-    while (!handles_.empty() && !handles_.back()) {
-      handles_.pop_back();
-    }
+    auto it = std::find_if(
+        handles_.rbegin(),
+        handles_.rend(),
+        [](const std::optional<Handle>& opt) { return opt.has_value(); });
+    handles_.erase(it.base(), handles_.end());
   }
   void forEachAllocatedRange(const std::function<void(size_t, size_t)>& fn) {
     size_t start = 0;
@@ -770,7 +910,11 @@ struct ExpandableSegment {
     std::optional<std::variant<int, CUmemFabricHandle>> shareable_handle;
   };
   struct ShareHeader {
+#ifdef _WIN32
+    int pid;
+#else
     pid_t pid;
+#endif
     size_t segment_size;
     size_t num_handles;
   };
@@ -864,15 +1008,16 @@ struct RestoreResult {
   std::vector<Block*> allocations_created;
 };
 
-bool BlockComparatorSize(const Block* a, const Block* b) {
+bool BlockComparatorRegistrationCounter(const Block* a, const Block* b) {
   if (a->stream != b->stream) {
     return (uintptr_t)a->stream < (uintptr_t)b->stream;
   }
   if (a->size != b->size) {
     return a->size < b->size;
   }
-  return (uintptr_t)a->ptr < (uintptr_t)b->ptr;
+  return a->registration_counter < b->registration_counter;
 }
+
 bool BlockComparatorAddress(const Block* a, const Block* b) {
   if (a->stream != b->stream) {
     return (uintptr_t)a->stream < (uintptr_t)b->stream;
@@ -880,6 +1025,14 @@ bool BlockComparatorAddress(const Block* a, const Block* b) {
   return (uintptr_t)a->ptr < (uintptr_t)b->ptr;
 }
 
+// Info about OOM rejection, used to defer observer callbacks outside of lock
+struct OomRejectionInfo {
+  bool rejected{false};
+  size_t alloc_size{0};
+  size_t total_allocated{0};
+  size_t device_total{0};
+};
+
 struct AllocParams {
   AllocParams(
       c10::DeviceIndex device,
@@ -910,6 +1063,7 @@ struct AllocParams {
   Block* block{nullptr};
   StatTypes stat_types = {false};
   cudaError_t err{cudaSuccess};
+  OomRejectionInfo oom_rejection_info;
 };
 
 // Note: cudaEventCreate when concurrently invoked from multiple threads can be
@@ -1273,6 +1427,7 @@ class DeviceCachingAllocator {
 
   // XXX - maybe we should generalize and have multiple events
   std::vector<OutOfMemoryObserver> oom_observers_;
+  std::vector<OomRejectionObserver> oom_rejection_observers_;
 
   std::vector<AllocatorTraceTracker> trace_trackers_;
 
@@ -1381,6 +1536,10 @@ class DeviceCachingAllocator {
     oom_observers_.emplace_back(std::move(observer));
   }
 
+  void attachOomRejectionObserver(OomRejectionObserver observer) {
+    oom_rejection_observers_.emplace_back(std::move(observer));
+  }
+
   void attachAllocatorTraceTracker(AllocatorTraceTracker tracker) {
     std::unique_lock<std::recursive_mutex> lock(mutex);
     trace_trackers_.emplace_back(std::move(tracker));
@@ -1461,19 +1620,28 @@ class DeviceCachingAllocator {
       // cudaMalloc. So far this function has not modified allocator state, but
       // keep in mind that any observed allocator state may change across calls
       // to alloc_block since it may release the lock.
-      block_found = alloc_block(params, false, context, lock)
-          // Try to use memory pools that have opted in as overflow before
-          // expensive memory freeing operations.
-          || try_mempool_fallback(
-                        params, size, stream, device_id, alloc_size, stats)
-          // Free enough available cached blocks to satisfy alloc and retry
-          // alloc.
-          || (release_available_cached_blocks(params, context) &&
-              alloc_block(params, false, context, lock))
-          // Free all non-split cached blocks and retry alloc.
-          || (C10_LIKELY(captures_underway.empty()) &&
-              release_cached_blocks(context, {0, 0}) &&
-              alloc_block(params, true, context, lock));
+      block_found = alloc_block(params, false, context, lock);
+
+      // If allocation was rejected by OOM policy, skip retry chain and fail
+      // immediately
+      if (!block_found && params.oom_rejection_info.rejected) {
+        // Skip retry chain - will be handled below in the !block_found path
+      } else if (!block_found) {
+        // Normal retry chain: try various strategies to free memory and retry
+        block_found =
+            // Try to use memory pools that have opted in as overflow before
+            // expensive memory freeing operations.
+            try_mempool_fallback(
+                params, size, stream, device_id, alloc_size, stats)
+            // Free enough available cached blocks to satisfy alloc and retry
+            // alloc.
+            || (release_available_cached_blocks(params, context) &&
+                alloc_block(params, false, context, lock))
+            // Free all non-split cached blocks and retry alloc.
+            || (C10_LIKELY(captures_underway.empty()) &&
+                release_cached_blocks(context, {0, 0}) &&
+                alloc_block(params, true, context, lock));
+      }
     }
 
     if (!block_found) {
@@ -1481,6 +1649,43 @@ class DeviceCachingAllocator {
       // alloc_block should have thrown an exception already.
       TORCH_INTERNAL_ASSERT(params.err == cudaErrorMemoryAllocation);
 
+      // Handle OOM rejection separately - return nullptr instead of throwing
+      // This allows callers to handle rejection gracefully without crashing
+      if (params.oom_rejection_info.rejected) {
+        if (!oom_rejection_observers_.empty()) {
+          auto observers_local = oom_rejection_observers_;
+          auto rejection_info = params.oom_rejection_info;
+          auto device = device_id;
+
+          // Release lock before dispatching observers
+          lock.unlock();
+
+          // Dispatch observers asynchronously to minimize latency impact on
+          // rejection path
+          std::thread([observers_local = std::move(observers_local),
+                       rejection_info,
+                       device]() {
+            try {
+              for (const auto& observer : observers_local) {
+                observer(
+                    device,
+                    rejection_info.alloc_size,
+                    rejection_info.total_allocated,
+                    rejection_info.device_total);
+              }
+            } catch (const std::exception& e) {
+              LOG(ERROR) << "Exception in OOM rejection observer: " << e.what();
+            } catch (...) {
+              LOG(ERROR) << "Unknown exception in OOM rejection observer";
+            }
+          }).detach();
+        } else {
+          lock.unlock();
+        }
+
+        return nullptr;
+      }
+
       size_t device_free = 0;
       size_t device_total = 0;
       C10_CUDA_CHECK(cudaMemGetInfo(&device_free, &device_total));
@@ -2180,6 +2385,7 @@ class DeviceCachingAllocator {
     stats.num_sync_all_streams = 0;
     stats.num_device_alloc = 0;
     stats.num_device_free = 0;
+    stats.num_oom_rejections = 0;
     stats.oversize_allocations.reset_accumulated();
     stats.oversize_segments.reset_accumulated();
   }
@@ -2656,18 +2862,62 @@ class DeviceCachingAllocator {
   void endAllocateToPool(MempoolId_t mempool_id) {
     std::lock_guard<std::recursive_mutex> lock(mutex);
 
-    if (CUDAAllocatorConfig::graph_capture_record_stream_reuse() &&
-        !graph_reuse_context.empty()) {
-      auto capture_id = mempool_to_capture_id[mempool_id];
-      auto graph_context = graph_reuse_context[capture_id];
-      for (auto& [stream, _] : graph_context.visited) {
-        TORCH_INTERNAL_ASSERT(
-            stream_get_capture_info(stream).status ==
-                cudaStreamCaptureStatusNone,
-            "This stream should not be capturing when the capture is ended");
+    if (CUDAAllocatorConfig::graph_capture_record_stream_reuse()) {
+      // Remove stream reuse context and mapping for this capture, if present.
+      if (!graph_reuse_context.empty()) {
+        auto capture_id = mempool_to_capture_id[mempool_id];
+        auto graph_context = graph_reuse_context[capture_id];
+        for (auto& [stream, _] : graph_context.visited) {
+          TORCH_INTERNAL_ASSERT(
+              stream_get_capture_info(stream).status ==
+                  cudaStreamCaptureStatusNone,
+              "This stream should not be capturing when the capture is ended");
+        }
+        graph_reuse_context.erase(capture_id);
+        mempool_to_capture_id.erase(mempool_id);
+      }
+
+      // Free deferred blocks associated with the ended pool, if any.
+      // cudaStreamEndCapture would have failed if any stream used during
+      // capture hadn't been joined back, so all stream uses on these
+      // blocks are known to be complete and we can safely clear them.
+      if (!deferred_blocks.empty()) {
+        auto pool_it = graph_pools.find(mempool_id);
+        if (pool_it != graph_pools.end()) {
+          auto* private_pool = pool_it->second.get();
+          auto context = maybeGatherContext(RecordContext::ALL);
+          std::vector<Block*> blocks_to_erase;
+          for (auto& [block, markers] : deferred_blocks) {
+            if (block->pool->owner_PrivatePool == private_pool) {
+              // At capture end, handle blocks associated with non-capturing
+              // streams. Remove only stream uses introduced during capture
+              // (guaranteed complete), and for any leftover pre-capture uses,
+              // insert events to track their completion. This aligns with
+              // insert_events_deferred_until_no_capture semantics.
+              remove_cudagraph_stream_uses(block);
+              if (block->stream_uses.empty()) {
+                free_block(block, context);
+              } else {
+                // Pre-capture stream uses remain; record events so
+                // process_events can free the block once they complete.
+                insert_events(block);
+                // block->event_count should likely be non-zero here since
+                // block->stream_uses is not empty. Defensive: still free if
+                // event_count is zero, but this should be rare.
+                if (block->event_count == 0) {
+                  free_block(block, context);
+                }
+              }
+              // Must erase from deferred_blocks regardless of which branch we
+              // took.
+              blocks_to_erase.push_back(block);
+            }
+          }
+          for (auto* b : blocks_to_erase) {
+            deferred_blocks.erase(b);
+          }
+        }
       }
-      graph_reuse_context.erase(capture_id);
-      mempool_to_capture_id.erase(mempool_id);
     }
 
     for (auto it = captures_underway.begin(); it != captures_underway.end();
@@ -3294,8 +3544,37 @@ class DeviceCachingAllocator {
         total_allocated_memory + size > allowed_memory_maximum.value()) {
       p.err = cudaErrorMemoryAllocation;
       return false;
-      // Temporarily disable checkpointing & cudagraphs internally
-    } else if (
+    }
+
+    // When throw_on_cudamalloc_oom is enabled, check
+    // per_process_memory_fraction to reject allocations that would exceed the
+    // configured limit. This prevents fatal HSA/driver aborts by throwing
+    // OutOfMemoryError instead.
+    if (CUDAAllocatorConfig::throw_on_cudamalloc_oom()) {
+      size_t device_total = static_cast<size_t>(device_prop.totalGlobalMem);
+      size_t max_allowed = static_cast<size_t>(
+          CUDAAllocatorConfig::per_process_memory_fraction() *
+          static_cast<double>(device_total));
+      if (total_allocated_memory + size > max_allowed) {
+        stats.num_oom_rejections++;
+        p.oom_rejection_info = {
+            true, size, total_allocated_memory, device_total};
+        C10_LOG_EVERY_MS(WARNING, 60000)
+            << "Preemptively rejecting allocation: requested "
+            << format_size(size) << " but total_allocated_memory ("
+            << format_size(total_allocated_memory)
+            << ") + requested size would exceed memory limit ("
+            << format_size(max_allowed) << " = "
+            << CUDAAllocatorConfig::per_process_memory_fraction() * 100
+            << "% of " << format_size(device_total)
+            << "). Set throw_on_cudamalloc_oom:false to disable.";
+        p.err = cudaErrorMemoryAllocation;
+        return false;
+      }
+    }
+
+    if (
+        // Temporarily disable checkpointing & cudagraphs internally
         p.is_expandable_segments_active &&
         !(in_fbcode && p.pool->owner_PrivatePool)) {
       p.block = try_allocate_expandable_block(
@@ -3837,7 +4116,7 @@ class DeviceCachingAllocator {
 
     if (record_history) {
       // Skip if action is in the skip_actions set
-      bool should_skip = skip_actions_list.count(action) > 0;
+      bool should_skip = skip_actions_list.contains(action);
       if (!should_skip) {
         alloc_buffer.insertEntries(te);
       }
@@ -3967,6 +4246,18 @@ class NativeCachingAllocator : public CUDAAllocator {
         device,
         ": did you call init?");
     Block* block = device_allocator[device]->malloc(size, stream);
+    // block can be nullptr if allocation was rejected by
+    // throw_on_cudamalloc_oom
+    // + per_process_memory_fraction policy. Throw OutOfMemoryError so callers
+    // with OOM error handlers can catch it gracefully.
+    if (C10_UNLIKELY(!block)) {
+      TORCH_CHECK_WITH(
+          OutOfMemoryError,
+          false,
+          "CUDA out of memory. Allocation was preemptively rejected because "
+          "it would exceed per_process_memory_fraction limit. Requested size: ",
+          format_size(size));
+    }
     add_allocated_block(block);
     *devPtr = block->ptr;
     const c10::impl::PyInterpreter* interp = c10::impl::GPUTrace::get_trace();
@@ -4095,6 +4386,15 @@ class NativeCachingAllocator : public CUDAAllocator {
     return device_allocator[device]->isHistoryEnabled();
   }
 
+  std::shared_ptr<GatheredContext> getContextForPointer(
+      const void* ptr) override {
+    Block* block = get_allocated_block(const_cast<void*>(ptr));
+    if (!block) {
+      return nullptr;
+    }
+    return block->context_when_allocated;
+  }
+
   bool checkPoolLiveAllocations(
       c10::DeviceIndex device,
       MempoolId_t mempool_id,
@@ -4109,6 +4409,12 @@ class NativeCachingAllocator : public CUDAAllocator {
     }
   }
 
+  void attachOomRejectionObserver(OomRejectionObserver observer) override {
+    for (auto& allocator : device_allocator) {
+      allocator->attachOomRejectionObserver(observer);
+    }
+  }
+
   void attachAllocatorTraceTracker(AllocatorTraceTracker tracker) override {
     for (auto& allocator : device_allocator) {
       allocator->attachAllocatorTraceTracker(tracker);
@@ -4655,7 +4961,7 @@ struct BackendStaticInitializer {
 // initialization, and doing so there may introduce static initialization
 // order (SIOF) issues.
 #define HIP_MASQUERADING_AS_CUDA "cuda"
-    at::SetAllocator(c10::Device(HIP_MASQUERADING_AS_CUDA).type(), r, 0);
+    at::SetAllocator(c10::Device(HIP_MASQUERADING_AS_CUDA).type(), r, 1);
     allocator.store(r);
 #undef HIP_MASQUERADING_AS_CUDA
   }
diff --git a/c10/cuda/CUDACachingAllocator.h b/c10/cuda/CUDACachingAllocator.h
index 88b6d8905becc..c2d502976dda6 100644
--- a/c10/cuda/CUDACachingAllocator.h
+++ b/c10/cuda/CUDACachingAllocator.h
@@ -95,14 +95,26 @@ using OutOfMemoryObserver = std::function<void(
     size_t device_total,
     size_t device_free)>;
 
+// Observer called when an allocation is preemptively rejected due to
+// throw_on_cudamalloc_oom policy. Parameters:
+//   - device: GPU device index
+//   - alloc_size: size of the rejected allocation request
+//   - total_allocated: total memory allocated before the request
+//   - device_total: total GPU memory
+using OomRejectionObserver = std::function<void(
+    int64_t device,
+    size_t alloc_size,
+    size_t total_allocated,
+    size_t device_total)>;
+
 struct ShareableHandle {
   ptrdiff_t offset;
   std::string handle;
 };
 
 struct StreamSegmentSize {
-  StreamSegmentSize(cudaStream_t s, bool small, size_t sz)
-      : stream(s), is_small_pool(small), total_size(sz) {}
+  StreamSegmentSize(cudaStream_t s, bool small_, size_t sz)
+      : stream(s), is_small_pool(small_), total_size(sz) {}
   cudaStream_t stream;
   bool is_small_pool;
   size_t total_size;
@@ -196,6 +208,10 @@ class CUDAAllocator : public DeviceAllocator {
         " does not yet support recordHistory. "
         "If you need it, please file an issue describing your use case.");
   }
+  virtual std::shared_ptr<GatheredContext> getContextForPointer(
+      const void* ptr) {
+    return nullptr;
+  }
   virtual void recordHistory(
       bool enabled,
       CreateContextFn context_recorder,
@@ -212,6 +228,7 @@ class CUDAAllocator : public DeviceAllocator {
     return "";
   }
   virtual void attachOutOfMemoryObserver(OutOfMemoryObserver observer) = 0;
+  virtual void attachOomRejectionObserver(OomRejectionObserver observer) = 0;
 
   // Attached AllocatorTraceTracker callbacks will be called while the
   // per-device allocator lock is held. Any additional locks taken from within
@@ -401,6 +418,10 @@ inline bool isHistoryEnabled() {
   return get()->isHistoryEnabled();
 }
 
+inline std::shared_ptr<GatheredContext> getContextForPointer(const void* ptr) {
+  return get()->getContextForPointer(ptr);
+}
+
 inline bool checkPoolLiveAllocations(
     c10::DeviceIndex device,
     MempoolId_t mempool_id,
@@ -413,6 +434,10 @@ inline void attachOutOfMemoryObserver(OutOfMemoryObserver observer) {
   get()->attachOutOfMemoryObserver(std::move(observer));
 }
 
+inline void attachOomRejectionObserver(OomRejectionObserver observer) {
+  get()->attachOomRejectionObserver(std::move(observer));
+}
+
 inline void attachAllocatorTraceTracker(AllocatorTraceTracker tracker) {
   get()->attachAllocatorTraceTracker(std::move(tracker));
 }
diff --git a/c10/cuda/CUDAEvent.h b/c10/cuda/CUDAEvent.h
index 98e0b56819e43..e6931ebb0a081 100644
--- a/c10/cuda/CUDAEvent.h
+++ b/c10/cuda/CUDAEvent.h
@@ -141,9 +141,9 @@ struct CUDAEvent {
         ".");
     CUDAGuard guard(device_index_);
 
-#ifndef USE_ROCM
+#if !defined(USE_ROCM) || ROCM_VERSION >= 70000
     // it is an error to use cudaEventRecordExternal when not doing stream
-    // capture
+    // capture (same applies to hipEventRecordExternal on ROCm 7.0+)
     unsigned int flags = (c10::cuda::currentStreamCaptureStatusMayInitCtx() !=
                               c10::cuda::CaptureStatus::None &&
                           external_)
@@ -168,9 +168,9 @@ struct CUDAEvent {
   void block(const CUDAStream& stream) {
     if (is_created_) {
       CUDAGuard guard(stream.device_index());
-#ifndef USE_ROCM
+#if !defined(USE_ROCM) || ROCM_VERSION >= 70000
       // it is an error to use cudaEventWaitExternal when not doing stream
-      // capture
+      // capture (same applies to hipEventWaitExternal on ROCm 7.0+)
       unsigned int flags = (c10::cuda::currentStreamCaptureStatusMayInitCtx() !=
                                 c10::cuda::CaptureStatus::None &&
                             external_)
@@ -252,8 +252,9 @@ struct CUDAEvent {
 
   void createEvent(DeviceIndex device_index) {
     external_ = (flags_ & cudaEventExternal) != 0;
-#ifdef USE_ROCM
-    TORCH_CHECK(!external_, "External events are disallowed in rocm");
+#if defined(USE_ROCM) && ROCM_VERSION < 70000
+    TORCH_CHECK(
+        !external_, "External CUDA-graph events on ROCm require ROCm 7.0+");
 #endif
     flags_ &= ~cudaEventExternal;
     device_index_ = device_index;
diff --git a/c10/cuda/CUDAFunctions.h b/c10/cuda/CUDAFunctions.h
index 16db498a8f9d6..d9f048b129b30 100644
--- a/c10/cuda/CUDAFunctions.h
+++ b/c10/cuda/CUDAFunctions.h
@@ -91,8 +91,9 @@ C10_CUDA_API void __inline__ memcpy_and_sync(
         c10::kCUDA, reinterpret_cast<uintptr_t>(stream));
   }
 #if defined(USE_ROCM) && USE_ROCM
-  // As of ROCm 6.4.1, HIP runtime does not raise an error during capture of
-  // hipMemcpyWithStream which is a synchronous call. Thus, we add a check
+#if ROCM_VERSION < 71400
+  // As of ROCm 6.4.1-7.13.x, HIP runtime does not raise an error during capture
+  // of hipMemcpyWithStream which is a synchronous call. Thus, we add a check
   // here explicitly.
   hipStreamCaptureStatus captureStatus;
   C10_CUDA_CHECK(hipStreamGetCaptureInfo(stream, &captureStatus, nullptr));
@@ -101,6 +102,11 @@ C10_CUDA_API void __inline__ memcpy_and_sync(
   } else {
     C10_CUDA_CHECK(hipErrorStreamCaptureUnsupported);
   }
+#else // ROCM_VERSION
+  // HIP reports the errors during graph capture in ROCm 7.14+
+  // no StreamCaptureStatus check required
+  C10_CUDA_CHECK(hipMemcpyWithStream(dst, src, nbytes, kind, stream));
+#endif // ROCM_VERSION
 #else
   C10_CUDA_CHECK(cudaMemcpyAsync(dst, src, nbytes, kind, stream));
   C10_CUDA_CHECK(cudaStreamSynchronize(stream));
diff --git a/c10/cuda/CUDAGraphsC10Utils.h b/c10/cuda/CUDAGraphsC10Utils.h
index 936875fd71d5c..3709aa64c7845 100644
--- a/c10/cuda/CUDAGraphsC10Utils.h
+++ b/c10/cuda/CUDAGraphsC10Utils.h
@@ -1,8 +1,10 @@
 #pragma once
 
+#include <c10/core/Allocator.h>
 #include <c10/cuda/CUDAStream.h>
+
 #include <iostream>
-#include <utility>
+#include <optional>
 
 // CUDA Graphs utils used by c10 and aten.
 // aten/cuda/CUDAGraphsUtils.cuh adds utils used by aten only.
@@ -67,10 +69,41 @@ inline std::ostream& operator<<(std::ostream& os, CaptureStatus status) {
 
 // Use this version where you're sure a CUDA context exists already.
 inline CaptureStatus currentStreamCaptureStatusMayInitCtx() {
-  cudaStreamCaptureStatus is_capturing{cudaStreamCaptureStatusNone};
+  cudaStreamCaptureStatus status{cudaStreamCaptureStatusNone};
   C10_CUDA_CHECK(
-      cudaStreamIsCapturing(c10::cuda::getCurrentCUDAStream(), &is_capturing));
-  return CaptureStatus(is_capturing);
+      cudaStreamIsCapturing(c10::cuda::getCurrentCUDAStream(), &status));
+  return CaptureStatus(status);
+}
+
+inline CaptureStatus captureStatusMayInitCtx(cudaStream_t stream) {
+  cudaStreamCaptureStatus status{cudaStreamCaptureStatusNone};
+  C10_CUDA_CHECK(cudaStreamIsCapturing(stream, &status));
+  return CaptureStatus(status);
+}
+
+inline bool isStreamCapturingMayInitCtx(cudaStream_t stream) {
+  return captureStatusMayInitCtx(stream) == CaptureStatus::Active;
+}
+
+inline std::optional<CaptureId_t> currentStreamCaptureIdMayInitCtx() {
+  cudaStreamCaptureStatus status{};
+  CaptureId_t capture_id = 0;
+  C10_CUDA_CHECK(cudaStreamGetCaptureInfo(
+      c10::cuda::getCurrentCUDAStream(), &status, &capture_id));
+  if (status == cudaStreamCaptureStatus::cudaStreamCaptureStatusActive) {
+    return capture_id;
+  }
+  return std::nullopt;
+}
+
+inline std::optional<CaptureId_t> captureIdMayInitCtx(cudaStream_t stream) {
+  cudaStreamCaptureStatus status{};
+  CaptureId_t capture_id = 0;
+  C10_CUDA_CHECK(cudaStreamGetCaptureInfo(stream, &status, &capture_id));
+  if (status == cudaStreamCaptureStatus::cudaStreamCaptureStatusActive) {
+    return capture_id;
+  }
+  return std::nullopt;
 }
 
 } // namespace c10::cuda
diff --git a/c10/cuda/CUDAMallocAsyncAllocator.cpp b/c10/cuda/CUDAMallocAsyncAllocator.cpp
index cc8bec049fc3f..d634de1d12718 100644
--- a/c10/cuda/CUDAMallocAsyncAllocator.cpp
+++ b/c10/cuda/CUDAMallocAsyncAllocator.cpp
@@ -668,6 +668,13 @@ struct CudaMallocAsyncAllocator : public CUDAAllocator {
         "If you need it, please file an issue describing your use case.");
   }
 
+  void attachOomRejectionObserver(OomRejectionObserver observer) override {
+    TORCH_CHECK(
+        false,
+        "cudaMallocAsync does not yet support attachOomRejectionObserver. "
+        "If you need it, please file an issue describing your use case.");
+  }
+
   void attachAllocatorTraceTracker(AllocatorTraceTracker tracker) override {
     TORCH_CHECK(
         false,
@@ -790,8 +797,6 @@ struct CudaMallocAsyncAllocator : public CUDAAllocator {
         "(For backend:native, snapshot returns a detailed summary of all "
         "blocks tracked by the allocator, but the cudaMallocAsync backend "
         "does not track individual blocks.)");
-    // Alternative: TORCH_WARN
-    return {};
   }
 
   // CUDAGraph interactions
diff --git a/c10/cuda/CUDAStream.h b/c10/cuda/CUDAStream.h
index d3d1402593751..183232c4d9191 100644
--- a/c10/cuda/CUDAStream.h
+++ b/c10/cuda/CUDAStream.h
@@ -116,6 +116,13 @@ class C10_CUDA_API CUDAStream {
 
   void synchronize() const;
 
+  bool is_capturing() const {
+    DeviceGuard guard{stream_.device()};
+    cudaStreamCaptureStatus status{cudaStreamCaptureStatusNone};
+    C10_CUDA_CHECK(cudaStreamIsCapturing(stream(), &status));
+    return status != cudaStreamCaptureStatusNone;
+  }
+
   int priority() const {
     DeviceGuard guard{stream_.device()};
     int priority = 0;
@@ -256,6 +263,28 @@ inline c10::cuda::CUDAStream getCurrentHIPStream(
   return c10::cuda::getCurrentCUDAStream(device_index);
 }
 inline auto& setCurrentHIPStream = c10::cuda::setCurrentCUDAStream;
+inline c10::cuda::CUDAStream getStreamFromPoolMasqueradingAsCUDA(
+    const bool isHighPriority = false,
+    DeviceIndex device = -1) {
+  return c10::cuda::getStreamFromPool(isHighPriority, device);
+}
+inline c10::cuda::CUDAStream getStreamFromPoolMasqueradingAsCUDA(
+    const int priority,
+    DeviceIndex device = -1) {
+  return c10::cuda::getStreamFromPool(priority, device);
+}
+inline auto& getStreamFromExternalMasqueradingAsCUDA =
+    c10::cuda::getStreamFromExternal;
+inline c10::cuda::CUDAStream getDefaultHIPStreamMasqueradingAsCUDA(
+    DeviceIndex device_index = -1) {
+  return c10::cuda::getDefaultCUDAStream(device_index);
+}
+inline c10::cuda::CUDAStream getCurrentHIPStreamMasqueradingAsCUDA(
+    DeviceIndex device_index = -1) {
+  return c10::cuda::getCurrentCUDAStream(device_index);
+}
+inline auto& setCurrentHIPStreamMasqueradingAsCUDA =
+    c10::cuda::setCurrentCUDAStream;
 } // namespace c10::hip
 #endif
 
diff --git a/c10/cuda/PeerToPeerAccess.cpp b/c10/cuda/PeerToPeerAccess.cpp
index 526d794bc49d5..cc312bdfaa4cb 100644
--- a/c10/cuda/PeerToPeerAccess.cpp
+++ b/c10/cuda/PeerToPeerAccess.cpp
@@ -10,12 +10,15 @@
 #include <c10/util/Logging.h>
 #include <c10/util/irange.h>
 
+#include <iomanip>
+#include <sstream>
 #include <vector>
 
 namespace c10::cuda {
 
 static std::vector<int8_t> p2pAccessEnabled_;
 static std::vector<int8_t> fabricAccessEnabled_;
+static std::vector<int> fabricCliqueId_;
 static int64_t num_devices_ = -1;
 
 namespace detail {
@@ -35,6 +38,8 @@ void init_p2p_access_cache(int64_t num_devices) {
   }
   fabricAccessEnabled_.clear();
   fabricAccessEnabled_.resize(num_devices, -1);
+  fabricCliqueId_.clear();
+  fabricCliqueId_.resize(num_devices, kCliqueIdNotQueried);
 }
 
 } // namespace detail
@@ -180,6 +185,7 @@ bool get_fabric_access(c10::DeviceIndex dev) {
     fabricInfo.version = nvmlGpuFabricInfo_v2;
     if (DriverAPI::get()->nvmlDeviceGetGpuFabricInfoV_ == nullptr) {
       cache = 0;
+      fabricCliqueId_[dev] = kCliqueIdUnsupported;
       return false;
     }
     TORCH_CHECK(
@@ -188,13 +194,17 @@ bool get_fabric_access(c10::DeviceIndex dev) {
             nvml_device, &fabricInfo));
     auto state = fabricInfo.state != NVML_GPU_FABRIC_STATE_NOT_SUPPORTED;
     if (state) {
+      fabricCliqueId_[dev] = static_cast<int>(fabricInfo.cliqueId);
       // now perform the full cycle of allocating - exporting - importing memory
       state = isFabricSupported();
+    } else {
+      fabricCliqueId_[dev] = kCliqueIdUnsupported;
     }
     cache = state ? 1 : 0;
     return cache;
   } else {
     cache = 0;
+    fabricCliqueId_[dev] = kCliqueIdUnsupported;
     return false;
   }
 #else
@@ -203,4 +213,85 @@ bool get_fabric_access(c10::DeviceIndex dev) {
 #endif
 }
 
+int get_fabric_clique_id(c10::DeviceIndex dev) {
+#if !defined(USE_ROCM) && defined(CUDA_VERSION) && CUDA_VERSION >= 12040 && \
+    defined(PYTORCH_C10_DRIVER_API_SUPPORTED)
+  // Ensure cache is populated via get_fabric_access (which does the NVML query
+  // and stashes clique_id as a side effect).
+  get_fabric_access(dev);
+  return fabricCliqueId_[dev];
+#else
+  (void)dev;
+  return kCliqueIdUnsupported;
+#endif
+}
+
+std::string get_nvml_fabric_info([[maybe_unused]] c10::DeviceIndex dev) {
+#if !defined(USE_ROCM) && defined(CUDA_VERSION) && CUDA_VERSION >= 12040 && \
+    defined(PYTORCH_C10_DRIVER_API_SUPPORTED)
+  if (DriverAPI::get()->nvmlDeviceGetGpuFabricInfoV_ == nullptr) {
+    return "fabric info unsupported (nvmlDeviceGetGpuFabricInfoV not available)";
+  }
+
+  auto nvml_device = get_nvml_device(dev);
+  if (nvml_device == nullptr) {
+    return "fabric info unknown (failed to get NVML device handle)";
+  }
+
+  nvmlGpuFabricInfoV_t info{};
+#ifdef nvmlGpuFabricInfo_v3
+  bool has_health_summary = false;
+  info.version = nvmlGpuFabricInfo_v3;
+  if (DriverAPI::get()->nvmlDeviceGetGpuFabricInfoV_(nvml_device, &info) ==
+      NVML_SUCCESS) {
+    has_health_summary = true;
+  } else
+#endif
+  {
+    info = {};
+    info.version = nvmlGpuFabricInfo_v2;
+    if (DriverAPI::get()->nvmlDeviceGetGpuFabricInfoV_(nvml_device, &info) !=
+        NVML_SUCCESS) {
+      return "fabric info unknown (nvmlDeviceGetGpuFabricInfoV failed)";
+    }
+  }
+
+  char uuid_hex[33];
+  for (int i = 0; i < 16; ++i) {
+    snprintf(uuid_hex + i * 2, 3, "%02x", info.clusterUuid[i]);
+  }
+
+  const char* state_str = "unknown";
+  switch (info.state) {
+    case NVML_GPU_FABRIC_STATE_NOT_SUPPORTED:
+      state_str = "not_supported";
+      break;
+    case NVML_GPU_FABRIC_STATE_NOT_STARTED:
+      state_str = "not_started";
+      break;
+    case NVML_GPU_FABRIC_STATE_IN_PROGRESS:
+      state_str = "in_progress";
+      break;
+    case NVML_GPU_FABRIC_STATE_COMPLETED:
+      state_str = "completed";
+      break;
+  }
+
+  std::ostringstream oss;
+  oss << "clique_id=" << info.cliqueId << ", cluster_uuid=" << uuid_hex
+      << ", state=" << state_str << ", status=" << info.status
+      << ", health_mask=0x" << std::hex << std::setfill('0') << std::setw(8)
+      << info.healthMask;
+#ifdef nvmlGpuFabricInfo_v3
+  if (has_health_summary) {
+    oss << ", health_summary=" << std::dec
+        << static_cast<int>(info.healthSummary);
+  }
+#endif
+  return oss.str();
+#else
+  return "fabric info unsupported (requires CUDA >= 12.4)";
+#endif
+}
+
 } // namespace c10::cuda
diff --git a/c10/cuda/PeerToPeerAccess.h b/c10/cuda/PeerToPeerAccess.h
index d526843974264..8615f1f335d96 100644
--- a/c10/cuda/PeerToPeerAccess.h
+++ b/c10/cuda/PeerToPeerAccess.h
@@ -5,6 +5,7 @@
 #include <c10/macros/Macros.h>
 
 #include <cstdint>
+#include <string>
 
 namespace c10::cuda {
 
@@ -32,4 +33,17 @@ C10_CUDA_API bool get_p2p_access(
 /// @return true if fabric access is available, false otherwise.
 C10_CUDA_API bool get_fabric_access(c10::DeviceIndex device);
 
+constexpr int kCliqueIdNotQueried = -2;
+constexpr int kCliqueIdUnsupported = -1;
+
+/// Query the NVLink fabric clique ID for a device.
+/// Returns the clique ID (>= 0) if fabric is supported, or kCliqueIdUnsupported
+/// if unsupported.
+C10_CUDA_API int get_fabric_clique_id(c10::DeviceIndex device);
+
+/// Returns a formatted string with NVML fabric info (clique_id, cluster_uuid,
+/// state, status, health_mask) for the given device. Intended for error
+/// diagnostics — only call on failure paths.
+C10_CUDA_API std::string get_nvml_fabric_info(c10::DeviceIndex device);
+
 } // namespace c10::cuda
diff --git a/c10/cuda/driver_api.h b/c10/cuda/driver_api.h
index 77bff56436ef5..681ebf49a1452 100644
--- a/c10/cuda/driver_api.h
+++ b/c10/cuda/driver_api.h
@@ -20,6 +20,23 @@
     }                                                                      \
   } while (0)
 
+// clang-format off
+#define C10_CUDA_DRIVER_CHECK_MSG(EXPR, ...)                                \
+  do {                                                                      \
+    CUresult __err = EXPR;                                                  \
+    if (__err != CUDA_SUCCESS) {                                            \
+      const char* err_str;                                                  \
+      CUresult get_error_str_err [[maybe_unused]] =                         \
+          c10::cuda::DriverAPI::get()->cuGetErrorString_(__err, &err_str);  \
+      if (get_error_str_err != CUDA_SUCCESS) {                              \
+        TORCH_CHECK(false, "CUDA driver error: unknown error", __VA_ARGS__);\
+      } else {                                                              \
+        TORCH_CHECK(false, "CUDA driver error: ", err_str, __VA_ARGS__);   \
+      }                                                                     \
+    }                                                                       \
+  } while (0)
+// clang-format on
+
 #define C10_CUDA_DRIVER_CHECK_GOTO(EXPR, NEXT)                             \
   do {                                                                     \
     CUresult __err = EXPR;                                                 \
@@ -52,6 +69,7 @@
 #define C10_LIBCUDA_DRIVER_API_REQUIRED(_)         \
   _(cuDeviceGet, 12000)                            \
   _(cuDeviceGetAttribute, 12000)                   \
+  _(cuMemGetAddressRange, 12000)                   \
   _(cuMemAddressReserve, 12000)                    \
   _(cuMemRelease, 12000)                           \
   _(cuMemMap, 12000)                               \
diff --git a/c10/cuda/impl/CUDAGuardImpl.h b/c10/cuda/impl/CUDAGuardImpl.h
index 14a398ee4dc37..c9617271cd498 100644
--- a/c10/cuda/impl/CUDAGuardImpl.h
+++ b/c10/cuda/impl/CUDAGuardImpl.h
@@ -213,6 +213,11 @@ struct CUDAGuardImpl final : public c10::impl::DeviceGuardImplInterface {
     cuda_stream.synchronize();
   }
 
+  bool isStreamCapturing(const Stream& stream) const override {
+    CUDAStream cuda_stream{stream};
+    return cuda_stream.is_capturing();
+  }
+
   void synchronizeEvent(void* event) const override {
     if (!event)
       return;
diff --git a/c10/metal/atomic.h b/c10/metal/atomic.h
index d0cbc03916989..947928f7c69db 100644
--- a/c10/metal/atomic.h
+++ b/c10/metal/atomic.h
@@ -9,6 +9,25 @@ struct AtomicType {};
 template <typename T>
 using AtomicType_t = typename AtomicType<T>::type;
 
+template <typename AT, typename T>
+static inline void atomic_binary_op_helper(
+    device ::metal::atomic<AT>* data,
+    long offset,
+    T value,
+    T (*op)(T, T)) {
+  auto ptr = data + offset;
+  auto old = ::metal::atomic_load_explicit(ptr, ::metal::memory_order_relaxed);
+  T val;
+  do {
+    val = op(old, value);
+  } while (!::metal::atomic_compare_exchange_weak_explicit(
+      ptr,
+      &old,
+      val,
+      ::metal::memory_order_relaxed,
+      ::metal::memory_order_relaxed));
+}
+
 template <>
 struct AtomicType<float> {
   using type = ::metal::atomic<float>;
@@ -16,6 +35,13 @@ struct AtomicType<float> {
     ::metal::atomic_fetch_add_explicit(
         data + offset, value, ::metal::memory_order_relaxed);
   }
+  static inline void atomic_binary_op(
+      device type* data,
+      long offset,
+      float value,
+      float (*op)(float, float)) {
+    atomic_binary_op_helper(data, offset, value, op);
+  }
 };
 
 template <>
@@ -25,6 +51,13 @@ struct AtomicType<int> {
     ::metal::atomic_fetch_add_explicit(
         data + offset, value, ::metal::memory_order_relaxed);
   }
+  static inline void atomic_binary_op(
+      device type* data,
+      long offset,
+      int value,
+      int (*op)(int, int)) {
+    atomic_binary_op_helper(data, offset, value, op);
+  }
 };
 
 // As of Metal3.2 atomic operations are not supported on half-precision floats,
@@ -35,6 +68,13 @@ static inline void atomic_add_helper(
     device ::metal::atomic<uint>* data,
     long offset,
     T value) {
+  // atomic<uint> requires 4-byte alignment; fix up misaligned pointers
+  auto addr = reinterpret_cast<ulong>(data);
+  auto misalign = (addr % alignof(::metal::atomic<uint>)) / sizeof(T);
+  data = reinterpret_cast<device ::metal::atomic<uint>*>(
+      reinterpret_cast<device char*>(data) - misalign * sizeof(T));
+  offset += misalign;
+
   constexpr auto elem_per_enum = sizeof(uint) / sizeof(T);
   auto ptr = data + (offset / elem_per_enum);
   auto old = ::metal::atomic_load_explicit(ptr, ::metal::memory_order_relaxed);
@@ -53,12 +93,51 @@ static inline void atomic_add_helper(
       ::metal::memory_order_relaxed));
 }
 
+template <typename T>
+static inline void atomic_binary_op_helper(
+    device ::metal::atomic<uint>* data,
+    long offset,
+    T value,
+    T (*Op)(T, T)) {
+  // atomic<uint> requires 4-byte alignment; fix up misaligned pointers
+  auto addr = reinterpret_cast<ulong>(data);
+  auto misalign = (addr % alignof(::metal::atomic<uint>)) / sizeof(T);
+  data = reinterpret_cast<device ::metal::atomic<uint>*>(
+      reinterpret_cast<device char*>(data) - misalign * sizeof(T));
+  offset += misalign;
+
+  constexpr auto elem_per_enum = sizeof(uint) / sizeof(T);
+  auto ptr = data + (offset / elem_per_enum);
+  auto old = ::metal::atomic_load_explicit(ptr, ::metal::memory_order_relaxed);
+  union {
+    uint i;
+    T t[elem_per_enum];
+  } val;
+  do {
+    val.i = old;
+    val.t[offset & (elem_per_enum - 1)] =
+        Op(val.t[offset & (elem_per_enum - 1)], value);
+  } while (!::metal::atomic_compare_exchange_weak_explicit(
+      ptr,
+      &old,
+      val.i,
+      ::metal::memory_order_relaxed,
+      ::metal::memory_order_relaxed));
+}
+
 template <>
 struct AtomicType<half> {
   using type = ::metal::atomic<uint>;
   static inline void atomic_add(device type* data, long offset, half value) {
     atomic_add_helper(data, offset, value);
   }
+  static inline void atomic_binary_op(
+      device type* data,
+      long offset,
+      half value,
+      half (*op)(half, half)) {
+    atomic_binary_op_helper(data, offset, value, op);
+  }
 };
 
 template <>
@@ -67,6 +146,13 @@ struct AtomicType<short> {
   static inline void atomic_add(device type* data, long offset, short value) {
     atomic_add_helper(data, offset, value);
   }
+  static inline void atomic_binary_op(
+      device type* data,
+      long offset,
+      short value,
+      short (*op)(short, short)) {
+    atomic_binary_op_helper(data, offset, value, op);
+  }
 };
 
 template <>
@@ -75,6 +161,13 @@ struct AtomicType<char> {
   static inline void atomic_add(device type* data, long offset, char value) {
     atomic_add_helper(data, offset, value);
   }
+  static inline void atomic_binary_op(
+      device type* data,
+      long offset,
+      char value,
+      char (*op)(char, char)) {
+    atomic_binary_op_helper(data, offset, value, op);
+  }
 };
 
 template <>
@@ -83,6 +176,13 @@ struct AtomicType<uchar> {
   static inline void atomic_add(device type* data, long offset, char value) {
     atomic_add_helper(data, offset, value);
   }
+  static inline void atomic_binary_op(
+      device type* data,
+      long offset,
+      uchar value,
+      uchar (*op)(uchar, uchar)) {
+    atomic_binary_op_helper(data, offset, value, op);
+  }
 };
 
 template <>
@@ -91,6 +191,13 @@ struct AtomicType<bfloat> {
   static inline void atomic_add(device type* data, long offset, bfloat value) {
     atomic_add_helper<bfloat>(data, offset, value);
   }
+  static inline void atomic_binary_op(
+      device type* data,
+      long offset,
+      bfloat value,
+      bfloat (*op)(bfloat, bfloat)) {
+    atomic_binary_op_helper(data, offset, value, op);
+  }
 };
 
 // Metal supports atomic_store_explicit for bools, but
diff --git a/c10/metal/common.h b/c10/metal/common.h
index e4b4d1a38ca4e..87e6400c92252 100644
--- a/c10/metal/common.h
+++ b/c10/metal/common.h
@@ -20,7 +20,10 @@
   _(ComplexHalf, 8)                    \
   _(ComplexFloat, 9)                   \
   _(Bool, 11)                          \
-  _(BFloat16, 15)
+  _(BFloat16, 15)                      \
+  _(UInt16, 27)                        \
+  _(UInt32, 28)                        \
+  _(UInt64, 29)
 
 namespace c10 {
 namespace metal {
diff --git a/c10/metal/indexing.h b/c10/metal/indexing.h
index 77a57b04d88b7..ffbec0a3d71d1 100644
--- a/c10/metal/indexing.h
+++ b/c10/metal/indexing.h
@@ -53,6 +53,25 @@ kernel void unary_dense(
   output[index] = f(input[index]);
 }
 
+template <typename T, typename F>
+kernel void unary_dense_vec4(
+    device result_of<F, T>* output [[buffer(0)]],
+    constant T* input [[buffer(1)]],
+    constant uint& numel [[buffer(2)]],
+    uint index [[thread_position_in_grid]]) {
+  F f;
+  uint base = index * 4;
+  if (base + 4 <= numel) {
+    using ::metal::vec;
+    vec<T, 4> val = *(constant vec<T, 4>*)(input + base);
+    *(device vec<result_of<F, T>, 4>*)(output + base) = {
+        f(val.x), f(val.y), f(val.z), f(val.w)};
+  } else {
+    for (uint i = base; i < numel; i++)
+      output[i] = f(input[i]);
+  }
+}
+
 template <typename T, typename F>
 kernel void unary_strided(
     device result_of<F, T>* output [[buffer(0)]],
@@ -90,6 +109,18 @@ kernel void unary_strided(
           constant uint& ndim,                                                 \
           uint index)
 
+#define REGISTER_UNARY_VEC4_OP(NAME, DTYPE0, DTYPE1)                          \
+  static_assert(                                                              \
+      ::metal::                                                               \
+          is_same_v<DTYPE1, ::c10::metal::result_of<NAME##_functor, DTYPE0>>, \
+      "Output dtype mismatch for unary op " #NAME " and input " #DTYPE0);     \
+  template [[host_name(#NAME "_dense_vec4_" #DTYPE1 "_" #DTYPE0)]]            \
+  kernel void ::c10::metal::unary_dense_vec4<DTYPE0, NAME##_functor>(         \
+      device ::c10::metal::result_of<NAME##_functor, DTYPE0> * output,        \
+      constant DTYPE0 * input,                                                \
+      constant uint & numel,                                                  \
+      uint index)
+
 #define DEFINE_UNARY_FLOATING_FUNCTOR(NAME)                                     \
   struct NAME##_functor {                                                       \
     template <typename T>                                                       \
@@ -182,10 +213,16 @@ inline T val_at_offs(P ptr, long offs, ScalarType type) {
       return cast_to<T>(val_at_offs<char>(ptr, offs));
     case ScalarType::Short:
       return cast_to<T>(val_at_offs<short>(ptr, offs));
+    case ScalarType::UInt16:
+      return cast_to<T>(val_at_offs<uint16_t>(ptr, offs));
     case ScalarType::Int:
       return cast_to<T>(val_at_offs<int>(ptr, offs));
+    case ScalarType::UInt32:
+      return cast_to<T>(val_at_offs<uint32_t>(ptr, offs));
     case ScalarType::Long:
       return cast_to<T>(val_at_offs<long>(ptr, offs));
+    case ScalarType::UInt64:
+      return cast_to<T>(val_at_offs<uint64_t>(ptr, offs));
     // Floats
     case ScalarType::Float:
       return cast_to<T>(val_at_offs<float>(ptr, offs));
diff --git a/c10/metal/random.h b/c10/metal/random.h
index c03d9b8a3149c..647b860b3f22c 100644
--- a/c10/metal/random.h
+++ b/c10/metal/random.h
@@ -33,7 +33,7 @@ uint4 single_round(uint4 ctr, uint2 key) {
   constexpr uint kPhiloxSB = 0xCD9E8D57;
   auto rc0 = mulhilo(kPhiloxSA, ctr.x);
   auto rc1 = mulhilo(kPhiloxSB, ctr.z);
-  return uint4(rc1.y ^ ctr.y ^ key.x, rc1.x, rc0.y ^ ctr.w ^ key.y, rc0.x);
+  return uint4(rc1.x ^ ctr.y ^ key.x, rc1.y, rc0.x ^ ctr.w ^ key.y, rc0.y);
 }
 
 uint4 multiple_rounds(uint4 ctr, uint2 key, uint rounds) {
diff --git a/c10/metal/special_math.h b/c10/metal/special_math.h
index 22cfc4a308234..ead38cbfff1ec 100644
--- a/c10/metal/special_math.h
+++ b/c10/metal/special_math.h
@@ -314,6 +314,11 @@ inline T i1e(T _x) {
 template <typename T>
 inline float log_gamma(const T);
 
+/*
+ * The gamma function approximations follow John D Cook's
+ * c++ implementation:  https://www.johndcook.com/Gamma.cpp.
+ * (BSD License)
+ */
 template <typename T>
 inline float gamma(const T x) {
   if (x < 0.001) {
@@ -533,6 +538,12 @@ inline float calc_digamma_positive_domain(float x) {
   return result + ::metal::log(x) - (0.5 / x) - y;
 }
 
+/*
+ * The digamma kernel and helper function is derived from the pytorch cpu
+ * of this function, which is itself derived from the implementation
+ * of the digamma function in the Cephes Math Library.
+ * See note [3-Clause BSD License for the Cephes Math Library].
+ */
 template <typename T0>
 inline float digamma(T0 x) {
   if (x < 0.0f) {
@@ -572,6 +583,36 @@ inline float polygamma(const int64_t order, const T0 input) {
   return sgn * gamma(n + 1) * zeta(n + 1, x);
 }
 
+template <typename T>
+float trigamma(T x) {
+  float sign = 1.0f;
+  float result = 0.0f;
+
+  if (x < 0.0f) {
+    sign = -1.0f;
+    auto sin_pi_x = sin(M_PI_F * x);
+    result -= (M_PI_F * M_PI_F) / (sin_pi_x * sin_pi_x);
+    x = 1.0f - x;
+  } else if (x == 0.0) {
+    return INFINITY;
+  } else if (x < 1.0) {
+    result += 1.0 / (x * x);
+    x += 1.0f;
+  }
+
+  for (int i = 0; i < 6; ++i) {
+    result += 1.0f / (x * x);
+    x += 1.0f;
+  }
+
+  const float ixx = 1.0f / (x * x);
+  result +=
+      (1.0f + 1.0f / (2.0f * x) +
+       ixx * ((1.0f / 6.0f) - ixx * ((1.0f / 30.0f) - ixx * (1.0f / 42.0f)))) /
+      x;
+  return sign * result;
+}
+
 template <typename T>
 inline ::metal::enable_if_t<is_scalar_floating_point_v<T>, T> sinc(T a) {
   if (a == static_cast<T>(0)) {
@@ -682,6 +723,19 @@ inline T logaddexp2(T a, T b) {
   }
 }
 
+template <typename T>
+inline float xlogy(T x, T y) {
+  if (::metal::isnan(y)) {
+    return NAN;
+  }
+
+  if (x == 0) {
+    return x;
+  }
+
+  return x * precise::log(float(y));
+}
+
 template <typename T>
 inline float xlog1py(T x, T y) {
   if (::metal::isnan(y)) {
diff --git a/c10/metal/utils.h b/c10/metal/utils.h
index 8d58d0dfdd1f2..357348a26b20d 100644
--- a/c10/metal/utils.h
+++ b/c10/metal/utils.h
@@ -189,6 +189,18 @@ inline common_dtype<T, U> floor_divide(T x, U y) {
   return ::metal::floor(x / y);
 }
 
+// Workaround for Metal compiler bug: the compiler produces wrong results
+// when optimizing fused (x / A) % B expressions for integral types.
+template <
+    typename T,
+    typename U,
+    ::metal::enable_if_t<
+        is_scalar_integral_v<T> && is_scalar_integral_v<U>,
+        bool> = true>
+inline common_dtype<T, U> safe_mod(volatile T x, U y) {
+  return x % y;
+}
+
 // fmod
 template <
     typename T,
@@ -449,6 +461,48 @@ inline float2 conj(float2 a) {
   return float2(a.x, -a.y);
 }
 
+// The following implementation of hypot provides better numerical stability
+// than the naive implementation. It is based on:
+// https://github.com/pearu/functional_algorithms/blob/7dbbfd7db225b1c202e0e364fc435423ccf52dbe/functional_algorithms/algorithms.py#L168
+//
+// This implementation changes the naive formula for the hypotenuse of a right
+// triangle, `h = sqrt(a^2 + b^2)`, into three alternate forms to be used in
+// different cases. The reason why the naive formula is unstable is because of
+// the square terms. If `a` or `b` are very large or very small floating point
+// numbers, then their squares will resolve to inf or 0.
+//
+// Assume `a >= b >= 0`. We can first change the formula to:
+// `h = a sqrt(1 + (b / a)^2)`
+// `h = a sqrt(1 + r)`
+// where `r = (b / a)^2`. Since `a >= b >= 0`, then `1 >= r >= 0`.
+//
+// Case 1: `a == b`
+//   The formula simplifies to `h = a sqrt(2)`.
+//
+// Case 2: `1 >> r > 0`
+//   Due to floating point error, `sqrt(1 + r)` resolves to 1. So we use the
+//   binomial approximation `sqrt(1 + r) ≈ 1 + r / 2`, and the formula becomes
+//   `h ≈ a + a r / 2`.
+//
+// Case 3: All other cases.
+//   Use `h = a sqrt(1 + r)`.
+inline float hypot(float a_, float b_) {
+  auto a = max(a_, b_);
+  auto b = min(a_, b_);
+
+  auto b_over_a = c10::metal::div(b, a);
+  auto r = c10::metal::mul(b_over_a, b_over_a);
+  auto sqrt_1_plus_r = ::metal::precise::sqrt(1 + r);
+
+  auto h1 = M_SQRT2_F * a;
+  auto h2 = a + a * r / 2;
+  auto h3 = a * sqrt_1_plus_r;
+  bool is_h1 = (a == b);
+  bool is_h2 = ((sqrt_1_plus_r == 1) && (r > 0));
+
+  return ::metal::select(::metal::select(h3, h2, is_h2), h1, is_h1);
+}
+
 #define INSTANTIATE_FOR_ALL_TYPES(MACRO) \
   MACRO(float);                          \
   MACRO(half);                           \
diff --git a/c10/test/CMakeLists.txt b/c10/test/CMakeLists.txt
index 27385159976f4..9a576bb996f22 100644
--- a/c10/test/CMakeLists.txt
+++ b/c10/test/CMakeLists.txt
@@ -18,12 +18,12 @@ if(BUILD_TEST)
   endforeach()
 endif()
 
-# ---[ C++17 header compilation test
+# ---[ C++17/20 header compilation test
 if(BUILD_TEST)
   add_executable(c10_cpp17_header_build_test util/cpp17_header_build_check.cpp)
   target_link_libraries(c10_cpp17_header_build_test ${C10_LIB} gmock gtest gtest_main)
   set_target_properties(c10_cpp17_header_build_test PROPERTIES
-    CXX_STANDARD 17
+    CXX_STANDARD 20
     CXX_STANDARD_REQUIRED ON
   )
   add_test(NAME c10_cpp17_header_build_test COMMAND $<TARGET_FILE:c10_cpp17_header_build_test>)
diff --git a/c10/test/core/DeviceGuard_test.cpp b/c10/test/core/DeviceGuard_test.cpp
index a6148adf8af25..6ae4fc1dd6774 100644
--- a/c10/test/core/DeviceGuard_test.cpp
+++ b/c10/test/core/DeviceGuard_test.cpp
@@ -1,8 +1,12 @@
 #include <gtest/gtest.h>
 
 #include <c10/core/DeviceGuard.h>
+#include <c10/core/impl/DeviceGuardImplInterface.h>
 #include <c10/core/impl/FakeGuardImpl.h>
 
+#include <thread>
+#include <vector>
+
 using namespace c10;
 using namespace c10::impl;
 
@@ -39,3 +43,75 @@ TEST(OptionalDeviceGuard, ResetDeviceDifferentDeviceType) {
   ASSERT_EQ(g.current_device(), Device(DeviceType::HIP, 2));
   ASSERT_EQ(g.original_device(), Device(DeviceType::HIP, 0));
 }
+
+// -- ensureCUDADeviceGuardSet -------------------------------------------
+
+// Regression test: ensureCUDADeviceGuardSet() used to store a thread-local
+// FakeGuardImpl* in the global device_guard_impl_registry.  When the owning
+// thread exited its TLS was freed, leaving a dangling pointer that the next
+// thread to call deviceCount() would dereference (segfault).
+//
+// The fix is a function-local static, which has program lifetime.  We verify
+// that the pointer in the registry is still valid (and returns the expected
+// deviceCount) after the threads that triggered guard installation have exited.
+TEST(EnsureCUDADeviceGuard, NoUseAfterFreeWhenThreadsExit) {
+  // Simulate "CUDA compiled, no devices visible": a guard that is non-null but
+  // returns deviceCount() == 0, which is the condition that triggers fake guard
+  // installation in ensureCUDADeviceGuardSet().
+  struct ZeroDeviceGuardImpl final : public DeviceGuardImplInterface {
+    DeviceType type() const override {
+      return DeviceType::CUDA;
+    }
+    Device exchangeDevice(Device d) const override {
+      return d;
+    }
+    Device getDevice() const override {
+      return Device(DeviceType::CUDA, 0);
+    }
+    void setDevice(Device) const override {}
+    void uncheckedSetDevice(Device) const noexcept override {}
+    Stream getStream(Device d) const noexcept override {
+      return Stream(Stream::UNSAFE, d, 0);
+    }
+    Stream exchangeStream(Stream s) const noexcept override {
+      return s;
+    }
+    DeviceIndex deviceCount() const noexcept override {
+      return 0;
+    }
+    void record(void**, const Stream&, const DeviceIndex, const EventFlag)
+        const override {}
+    void block(void*, const Stream&) const override {}
+    bool queryEvent(void*) const override {
+      return true;
+    }
+    void destroyEvent(void*, const DeviceIndex) const noexcept override {}
+  };
+
+  constexpr auto cuda_idx = static_cast<size_t>(DeviceType::CUDA);
+  const auto* saved = device_guard_impl_registry[cuda_idx].load();
+
+  static ZeroDeviceGuardImpl zero_impl;
+  device_guard_impl_registry[cuda_idx].store(&zero_impl);
+
+  // Phase 1: threads call ensureCUDADeviceGuardSet(), detect deviceCount()==0,
+  // and install a FakeGuardImpl in the global registry.
+  {
+    std::vector<std::thread> threads;
+    for (int i = 0; i < 4; i++) {
+      threads.emplace_back(ensureCUDADeviceGuardSet);
+    }
+    for (auto& t : threads) {
+      t.join();
+    }
+  }
+  // The threads' TLS is now destroyed.  With the old code the registry now
+  // holds a dangling pointer; with the fix it holds &fake_cuda_guard (static).
+
+  // Phase 2: the pointer must still be valid and return the expected count.
+  const auto* p = device_guard_impl_registry[cuda_idx].load();
+  ASSERT_NE(p, nullptr);
+  ASSERT_EQ(p->deviceCount(), kFakeGuardImplMaxDevices);
+
+  device_guard_impl_registry[cuda_idx].store(saved);
+}
diff --git a/c10/test/core/DispatchKeySet_test.cpp b/c10/test/core/DispatchKeySet_test.cpp
index cdbdc150167e0..9066d6100a3db 100644
--- a/c10/test/core/DispatchKeySet_test.cpp
+++ b/c10/test/core/DispatchKeySet_test.cpp
@@ -404,7 +404,7 @@ TEST(DispatchKeySet, TestBackendComponentToString) {
     auto k = static_cast<BackendComponent>(i);
     auto res = std::string(toString(k));
     ASSERT_FALSE(res == "UNKNOWN_BACKEND_BIT");
-    ASSERT_FALSE(seen_strings.count(res) > 0);
+    ASSERT_FALSE(seen_strings.contains(res));
     seen_strings.insert(res);
   }
 }
@@ -439,7 +439,7 @@ TEST(DispatchKeySet, TestFunctionalityDispatchKeyToString) {
     } else {
       ASSERT_TRUE(res.find("Unknown") == std::string::npos) << i;
     }
-    ASSERT_TRUE(seen_strings.count(res) == 0);
+    ASSERT_TRUE(!seen_strings.contains(res));
     seen_strings.insert(res);
   }
 }
diff --git a/c10/test/core/impl/cow_test.cpp b/c10/test/core/impl/cow_test.cpp
index 7c306a3322a3c..dfc3bcdfd09e9 100644
--- a/c10/test/core/impl/cow_test.cpp
+++ b/c10/test/core/impl/cow_test.cpp
@@ -230,7 +230,7 @@ bool buffers_are_equal(const void* a, const void* b, size_t nbytes) {
 
 TEST(materialize_test, copy_on_write) {
   StorageImpl original_storage(
-      {}, /*size_bytes=*/6, GetCPUAllocator(), /*resizable=*/false);
+      {}, /*size_bytes=*/4, GetCPUAllocator(), /*resizable=*/false);
   std::memcpy(original_storage.mutable_data(), "abcd", 4);
   void const* original_data = original_storage.data();
 
diff --git a/c10/test/util/WaitCounter_test.cpp b/c10/test/util/WaitCounter_test.cpp
index 1907f3fbeb289..77b6b45c125b7 100644
--- a/c10/test/util/WaitCounter_test.cpp
+++ b/c10/test/util/WaitCounter_test.cpp
@@ -11,142 +11,134 @@
 
 namespace {
 
+// Struct to hold shared state
+struct CounterState {
+  std::atomic<int> startCount{0};
+  std::atomic<int> stopCount{0};
+};
+
 // Mock backend for testing WaitCounter functionality
 class MockWaitCounterBackend
     : public c10::monitor::detail::WaitCounterBackendIf {
  public:
-  explicit MockWaitCounterBackend(
-      std::atomic<int>& startCount,
-      std::atomic<int>& stopCount)
-      : startCount_(startCount), stopCount_(stopCount) {}
+  // Backend now holds a shared_ptr to the state
+  explicit MockWaitCounterBackend(std::shared_ptr<CounterState> state)
+      : state_(state) {}
 
   intptr_t start(std::chrono::steady_clock::time_point now) noexcept override {
-    startCount_.fetch_add(1);
-    lastStartTime_ = now;
+    state_->startCount.fetch_add(1);
     return reinterpret_cast<intptr_t>(this);
   }
 
   void stop(std::chrono::steady_clock::time_point now, intptr_t ctx) noexcept
       override {
-    stopCount_.fetch_add(1);
-    lastStopTime_ = now;
+    state_->stopCount.fetch_add(1);
     EXPECT_EQ(ctx, reinterpret_cast<intptr_t>(this));
   }
 
-  std::chrono::steady_clock::time_point lastStartTime() const {
-    return lastStartTime_;
-  }
-
-  std::chrono::steady_clock::time_point lastStopTime() const {
-    return lastStopTime_;
-  }
-
  private:
-  std::atomic<int>& startCount_;
-  std::atomic<int>& stopCount_;
-  std::chrono::steady_clock::time_point lastStartTime_;
-  std::chrono::steady_clock::time_point lastStopTime_;
+  std::shared_ptr<CounterState> state_;
 };
 
 class MockWaitCounterBackendFactory
     : public c10::monitor::detail::WaitCounterBackendFactoryIf {
  public:
+  // Factory accepts and stores a shared_ptr to the state
   MockWaitCounterBackendFactory(
-      std::atomic<int>& startCount,
-      std::atomic<int>& stopCount,
+      std::shared_ptr<CounterState> state,
       std::string_view keyFilter = "")
-      : startCount_(startCount), stopCount_(stopCount), keyFilter_(keyFilter) {}
+      : state_(state), keyFilter_(keyFilter) {}
 
   std::unique_ptr<c10::monitor::detail::WaitCounterBackendIf> create(
       std::string_view key) noexcept override {
     if (!keyFilter_.empty() && key.find(keyFilter_) == std::string_view::npos) {
       return nullptr;
     }
-    return std::make_unique<MockWaitCounterBackend>(startCount_, stopCount_);
+    // Pass the shared_ptr to the backend
+    return std::make_unique<MockWaitCounterBackend>(state_);
   }
 
  private:
-  std::atomic<int>& startCount_;
-  std::atomic<int>& stopCount_;
+  std::shared_ptr<CounterState> state_;
   std::string keyFilter_;
 };
 
 TEST(WaitCounter, BackendRegistration) {
-  auto backends = c10::monitor::detail::getRegisteredWaitCounterBackends();
-  size_t initialCount = backends.size();
-
-  std::atomic<int> startCount{0};
-  std::atomic<int> stopCount{0};
+  auto state = std::make_shared<CounterState>();
+  constexpr std::string_view key = "backend_registration";
 
   c10::monitor::detail::registerWaitCounterBackend(
-      std::make_unique<MockWaitCounterBackendFactory>(startCount, stopCount));
+      std::make_unique<MockWaitCounterBackendFactory>(state, key));
 
-  backends = c10::monitor::detail::getRegisteredWaitCounterBackends();
-  EXPECT_EQ(backends.size(), initialCount + 1);
+  c10::monitor::WaitCounterHandle handle(key);
+  {
+    auto guard = handle.start();
+    EXPECT_EQ(state->startCount.load(), 1);
+    EXPECT_EQ(state->stopCount.load(), 0);
+  }
+  EXPECT_EQ(state->startCount.load(), 1);
+  EXPECT_EQ(state->stopCount.load(), 1);
 }
 
 TEST(WaitCounter, WaitGuardStartStop) {
-  std::atomic<int> startCount{0};
-  std::atomic<int> stopCount{0};
+  auto state = std::make_shared<CounterState>();
+  constexpr std::string_view key = "wait_guard_start_stop";
 
   c10::monitor::detail::registerWaitCounterBackend(
-      std::make_unique<MockWaitCounterBackendFactory>(
-          startCount, stopCount, "wait_guard_start_stop"));
+      std::make_unique<MockWaitCounterBackendFactory>(state, key));
 
-  int startBefore = startCount.load();
-  int stopBefore = stopCount.load();
+  int startBefore = state->startCount.load();
+  int stopBefore = state->stopCount.load();
 
-  EXPECT_GE(startCount.load(), startBefore);
+  EXPECT_GE(state->startCount.load(), startBefore);
   {
-    c10::monitor::WaitCounterHandle handle("wait_guard_start_stop");
+    c10::monitor::WaitCounterHandle handle(key);
     auto guard = handle.start();
-    EXPECT_GE(startCount.load(), startBefore + 1);
-    EXPECT_EQ(stopCount.load(), stopBefore);
+    EXPECT_GE(state->startCount.load(), startBefore + 1);
+    EXPECT_EQ(state->stopCount.load(), stopBefore);
   }
 
-  EXPECT_GE(stopCount.load(), stopBefore + 1);
+  EXPECT_GE(state->stopCount.load(), stopBefore + 1);
 }
 
 TEST(WaitCounter, WaitGuardExplicitStop) {
-  std::atomic<int> startCount{0};
-  std::atomic<int> stopCount{0};
+  auto state = std::make_shared<CounterState>();
+  constexpr std::string_view key = "wait_guard_explicit_stop";
 
   c10::monitor::detail::registerWaitCounterBackend(
-      std::make_unique<MockWaitCounterBackendFactory>(
-          startCount, stopCount, "wait_guard_explicit_stop"));
+      std::make_unique<MockWaitCounterBackendFactory>(state, key));
 
-  int startBefore = startCount.load();
-  int stopBefore = stopCount.load();
+  int startBefore = state->startCount.load();
+  int stopBefore = state->stopCount.load();
 
-  c10::monitor::WaitCounterHandle handle("wait_guard_explicit_stop");
+  c10::monitor::WaitCounterHandle handle(key);
   auto guard = handle.start();
-  EXPECT_GE(startCount.load(), startBefore + 1);
-  EXPECT_EQ(stopCount.load(), stopBefore);
+  EXPECT_GE(state->startCount.load(), startBefore + 1);
+  EXPECT_EQ(state->stopCount.load(), stopBefore);
 
   guard.stop();
-  EXPECT_GE(stopCount.load(), stopBefore + 1);
+  EXPECT_GE(state->stopCount.load(), stopBefore + 1);
 
   // Calling stop() again should be a no-op (guard is already stopped)
-  int stopAfterFirst = stopCount.load();
+  int stopAfterFirst = state->stopCount.load();
   guard.stop();
-  EXPECT_EQ(stopCount.load(), stopAfterFirst);
+  EXPECT_EQ(state->stopCount.load(), stopAfterFirst);
 }
 
 TEST(WaitCounter, WaitGuardMoveConstruction) {
-  std::atomic<int> startCount{0};
-  std::atomic<int> stopCount{0};
+  auto state = std::make_shared<CounterState>();
+  constexpr std::string_view key = "wait_guard_move";
 
   c10::monitor::detail::registerWaitCounterBackend(
-      std::make_unique<MockWaitCounterBackendFactory>(
-          startCount, stopCount, "wait_guard_move"));
+      std::make_unique<MockWaitCounterBackendFactory>(state, key));
 
-  int startBefore = startCount.load();
-  int stopBefore = stopCount.load();
+  int startBefore = state->startCount.load();
+  int stopBefore = state->stopCount.load();
 
   {
-    c10::monitor::WaitCounterHandle handle("wait_guard_move");
+    c10::monitor::WaitCounterHandle handle(key);
     auto guard1 = handle.start();
-    EXPECT_GE(startCount.load(), startBefore + 1);
+    EXPECT_GE(state->startCount.load(), startBefore + 1);
 
     // Move the guard
     auto guard2 = std::move(guard1);
@@ -154,116 +146,110 @@ TEST(WaitCounter, WaitGuardMoveConstruction) {
   }
 
   // Stop should be called exactly once when guard2 is destroyed
-  EXPECT_GE(stopCount.load(), stopBefore + 1);
+  EXPECT_GE(state->stopCount.load(), stopBefore + 1);
 }
 
 TEST(WaitCounter, StaticWaitCounterMacro) {
-  std::atomic<int> startCount{0};
-  std::atomic<int> stopCount{0};
+  auto state = std::make_shared<CounterState>();
+  constexpr std::string_view key = "static_macro_test";
 
   c10::monitor::detail::registerWaitCounterBackend(
-      std::make_unique<MockWaitCounterBackendFactory>(
-          startCount, stopCount, "static_macro_test"));
+      std::make_unique<MockWaitCounterBackendFactory>(state, key));
 
-  int startBefore = startCount.load();
-  int stopBefore = stopCount.load();
+  int startBefore = state->startCount.load();
+  int stopBefore = state->stopCount.load();
 
   {
     auto guard = STATIC_WAIT_COUNTER(static_macro_test).start();
-    EXPECT_GE(startCount.load(), startBefore + 1);
+    EXPECT_GE(state->startCount.load(), startBefore + 1);
   }
 
-  EXPECT_GE(stopCount.load(), stopBefore + 1);
+  EXPECT_GE(state->stopCount.load(), stopBefore + 1);
 }
 
 TEST(WaitCounter, StaticScopedWaitCounterMacro) {
-  std::atomic<int> startCount{0};
-  std::atomic<int> stopCount{0};
+  auto state = std::make_shared<CounterState>();
+  constexpr std::string_view key = "static_scoped_test";
 
   c10::monitor::detail::registerWaitCounterBackend(
-      std::make_unique<MockWaitCounterBackendFactory>(
-          startCount, stopCount, "static_scoped_test"));
+      std::make_unique<MockWaitCounterBackendFactory>(state, key));
 
-  int startBefore = startCount.load();
-  int stopBefore = stopCount.load();
+  int startBefore = state->startCount.load();
+  int stopBefore = state->stopCount.load();
 
   {
     STATIC_SCOPED_WAIT_COUNTER(static_scoped_test);
-    EXPECT_GE(startCount.load(), startBefore + 1);
+    EXPECT_GE(state->startCount.load(), startBefore + 1);
   }
 
-  EXPECT_GE(stopCount.load(), stopBefore + 1);
+  EXPECT_GE(state->stopCount.load(), stopBefore + 1);
 }
 
 TEST(WaitCounter, WithWaitCounterMacroAssign) {
-  std::atomic<int> startCount{0};
-  std::atomic<int> stopCount{0};
+  auto state = std::make_shared<CounterState>();
+  constexpr std::string_view key = "execute_with_test_assign";
 
   c10::monitor::detail::registerWaitCounterBackend(
-      std::make_unique<MockWaitCounterBackendFactory>(
-          startCount, stopCount, "execute_with_test"));
+      std::make_unique<MockWaitCounterBackendFactory>(state, key));
 
-  int startBefore = startCount.load();
-  int stopBefore = stopCount.load();
+  int startBefore = state->startCount.load();
+  int stopBefore = state->stopCount.load();
 
   int value = 0;
-  WITH_WAIT_COUNTER(execute_with_test, value = 42);
+  WITH_WAIT_COUNTER(execute_with_test_assign, value = 42);
 
   EXPECT_EQ(value, 42);
-  EXPECT_GE(startCount.load(), startBefore + 1);
-  EXPECT_GE(stopCount.load(), stopBefore + 1);
+  EXPECT_GE(state->startCount.load(), startBefore + 1);
+  EXPECT_GE(state->stopCount.load(), stopBefore + 1);
 }
 
 TEST(WaitCounter, WithWaitCounterMacroReturn) {
-  std::atomic<int> startCount{0};
-  std::atomic<int> stopCount{0};
+  auto state = std::make_shared<CounterState>();
+  constexpr std::string_view key = "execute_with_test_return";
 
   c10::monitor::detail::registerWaitCounterBackend(
-      std::make_unique<MockWaitCounterBackendFactory>(
-          startCount, stopCount, "execute_with_test"));
+      std::make_unique<MockWaitCounterBackendFactory>(state, key));
 
-  int startBefore = startCount.load();
-  int stopBefore = stopCount.load();
+  int startBefore = state->startCount.load();
+  int stopBefore = state->stopCount.load();
 
   int value = 0;
-  value = WITH_WAIT_COUNTER(execute_with_test, []() { return 42; }());
+  value = WITH_WAIT_COUNTER(execute_with_test_return, []() { return 42; }());
 
   EXPECT_EQ(value, 42);
-  EXPECT_GE(startCount.load(), startBefore + 1);
-  EXPECT_GE(stopCount.load(), stopBefore + 1);
+  EXPECT_GE(state->startCount.load(), startBefore + 1);
+  EXPECT_GE(state->stopCount.load(), stopBefore + 1);
 }
 
 TEST(WaitCounter, SameHandleMultipleTimes) {
-  std::atomic<int> startCount{0};
-  std::atomic<int> stopCount{0};
+  auto state = std::make_shared<CounterState>();
+  constexpr std::string_view key = "multiple_times_test";
 
   c10::monitor::detail::registerWaitCounterBackend(
-      std::make_unique<MockWaitCounterBackendFactory>(
-          startCount, stopCount, "multiple_times_test"));
+      std::make_unique<MockWaitCounterBackendFactory>(state, key));
 
-  int startBefore = startCount.load();
-  int stopBefore = stopCount.load();
+  int startBefore = state->startCount.load();
+  int stopBefore = state->stopCount.load();
 
-  c10::monitor::WaitCounterHandle handle("multiple_times_test");
+  c10::monitor::WaitCounterHandle handle(key);
 
   for (int i = 0; i < 5; ++i) {
     auto guard = handle.start();
   }
 
-  EXPECT_GE(startCount.load(), startBefore + 5);
-  EXPECT_GE(stopCount.load(), stopBefore + 5);
+  EXPECT_GE(state->startCount.load(), startBefore + 5);
+  EXPECT_GE(state->stopCount.load(), stopBefore + 5);
 }
 
 TEST(WaitCounter, ConcurrentUsage) {
-  std::atomic<int> startCount{0};
-  std::atomic<int> stopCount{0};
+  auto state = std::make_shared<CounterState>();
+  constexpr std::string_view key = "concurrent_test";
 
   c10::monitor::detail::registerWaitCounterBackend(
-      std::make_unique<MockWaitCounterBackendFactory>(
-          startCount, stopCount, "concurrent_test"));
+      std::make_unique<MockWaitCounterBackendFactory>(state, key));
 
-  int startBefore = startCount.load();
-  int stopBefore = stopCount.load();
+  int startBefore = state->startCount.load();
+  int stopBefore = state->stopCount.load();
 
   constexpr int kNumThreads = 10;
   constexpr int kIterationsPerThread = 100;
@@ -285,8 +271,10 @@ TEST(WaitCounter, ConcurrentUsage) {
   }
 
   EXPECT_GE(
-      startCount.load(), startBefore + kNumThreads * kIterationsPerThread);
-  EXPECT_GE(stopCount.load(), stopBefore + kNumThreads * kIterationsPerThread);
+      state->startCount.load(),
+      startBefore + kNumThreads * kIterationsPerThread);
+  EXPECT_GE(
+      state->stopCount.load(), stopBefore + kNumThreads * kIterationsPerThread);
 }
 
 TEST(WaitCounter, StaticHandlePerCallSite) {
diff --git a/c10/util/ApproximateClock.h b/c10/util/ApproximateClock.h
index 803ea404c9aa2..c1b65b33af5a7 100644
--- a/c10/util/ApproximateClock.h
+++ b/c10/util/ApproximateClock.h
@@ -29,6 +29,8 @@
 #else
 #undef C10_RDTSC
 #endif
+#elif defined(__aarch64__) && !defined(__CUDACC__) && !defined(__HIPCC__)
+#define C10_ARMTSC
 #endif
 
 namespace c10 {
@@ -70,6 +72,14 @@ inline time_t getTime(bool allow_monotonic = false) {
 #endif
 }
 
+#if defined(C10_ARMTSC)
+inline uint64_t getArmApproximateTime() {
+  uint64_t val;
+  __asm__ __volatile__("mrs %0, cntvct_el0" : "=r"(val));
+  return val;
+}
+#endif
+
 // We often do not need to capture true wall times. If a fast mechanism such
 // as TSC is available we can use that instead and convert back to epoch time
 // during post processing. This greatly reduce the clock's contribution to
@@ -81,6 +91,8 @@ inline time_t getTime(bool allow_monotonic = false) {
 inline auto getApproximateTime() {
 #if defined(C10_RDTSC)
   return static_cast<uint64_t>(__rdtsc());
+#elif defined(C10_ARMTSC)
+  return getArmApproximateTime();
 #else
   return getTime();
 #endif
diff --git a/c10/util/C++17.cpp b/c10/util/C++17.cpp
deleted file mode 100644
index 736a002e5aeb3..0000000000000
--- a/c10/util/C++17.cpp
+++ /dev/null
@@ -1 +0,0 @@
-#include <c10/util/C++17.h>
diff --git a/c10/util/C++17.h b/c10/util/C++17.h
index 98cdffcbef417..c6073b2aa05e8 100644
--- a/c10/util/C++17.h
+++ b/c10/util/C++17.h
@@ -3,43 +3,15 @@
 #define C10_UTIL_CPP17_H_
 
 #include <c10/macros/Macros.h>
-#include <functional>
-#include <memory>
-#include <type_traits>
+#include <tuple>
 #include <utility>
 
-#if !defined(__clang__) && !defined(_MSC_VER) && defined(__GNUC__) && \
-    __GNUC__ < 9
-#error \
-    "You're trying to build PyTorch with a too old version of GCC. We need GCC 9 or later."
-#endif
-
-#if defined(__clang__) && __clang_major__ < 9
-#error \
-    "You're trying to build PyTorch with a too old version of Clang. We need Clang 9 or later."
-#endif
-
-#if (defined(_MSC_VER) && (!defined(_MSVC_LANG) || _MSVC_LANG < 201703L)) || \
-    (!defined(_MSC_VER) && __cplusplus < 201703L)
-#error You need C++17 to compile PyTorch
-#endif
-
-#if defined(_WIN32) && (defined(min) || defined(max))
-#error Macro clash with min and max -- define NOMINMAX when compiling your program on Windows
-#endif
-
-/*
- * This header adds some polyfills with C++17 functionality
- */
-
 namespace c10::guts {
 
 #if defined(__HIP__)
 
-// Implementation from http://en.cppreference.com/w/cpp/utility/apply (but
-// modified)
-// TODO This is an incomplete implementation of std::apply, not working for
-// member functions.
+// std::apply is not available in HIP device code because it lacks
+// __host__ __device__ annotations in the standard library.
 namespace detail {
 template <class F, class Tuple, std::size_t... INDEX>
 C10_HOST_DEVICE constexpr auto apply_impl(
diff --git a/c10/util/CallOnce.h b/c10/util/CallOnce.h
index 2d8c9dc5b1151..1898b5699611a 100644
--- a/c10/util/CallOnce.h
+++ b/c10/util/CallOnce.h
@@ -1,7 +1,6 @@
 #pragma once
 
 #include <c10/macros/Macros.h>
-#include <c10/util/C++17.h>
 
 #include <atomic>
 #include <functional>
diff --git a/c10/util/Exception.cpp b/c10/util/Exception.cpp
index 50f423f917981..18d767e62d272 100644
--- a/c10/util/Exception.cpp
+++ b/c10/util/Exception.cpp
@@ -1,3 +1,4 @@
+// @allow-raw-throw
 #include <c10/util/Exception.h>
 #include <c10/util/Logging.h>
 #include <c10/util/Type.h>
diff --git a/c10/util/Exception.h b/c10/util/Exception.h
index 545e21e71c024..19e242a0f0980 100644
--- a/c10/util/Exception.h
+++ b/c10/util/Exception.h
@@ -1,3 +1,4 @@
+// @allow-raw-throw
 #ifndef C10_UTIL_EXCEPTION_H_
 #define C10_UTIL_EXCEPTION_H_
 
diff --git a/c10/util/Logging.cpp b/c10/util/Logging.cpp
index 8ebb6ba767bc2..0c78917fb2307 100644
--- a/c10/util/Logging.cpp
+++ b/c10/util/Logging.cpp
@@ -1,3 +1,4 @@
+// @allow-raw-throw
 #include <c10/util/Backtrace.h>
 #include <c10/util/Flags.h>
 #include <c10/util/Lazy.h>
diff --git a/c10/util/Semaphore.h b/c10/util/Semaphore.h
index 041d9abecf515..6d2235c09da5a 100644
--- a/c10/util/Semaphore.h
+++ b/c10/util/Semaphore.h
@@ -8,7 +8,13 @@
 
 // note: __cpp_lib_semaphore will not be defined in some apple platforms
 // even if >= C++20.
-#if __has_include(<semaphore>) && defined(__cpp_lib_semaphore) && __cpp_lib_semaphore >= 201907L
+//
+// libstdc++'s __atomic_semaphore has a lost-wakeup bug: _M_release skips
+// the futex notify when the counter is already positive, but a concurrent
+// _S_do_try_acquire can fail its CAS, see zero, and block — missing the
+// wakeup. https://gcc.gnu.org/bugzilla/show_bug.cgi?id=98033
+#if __has_include(<semaphore>) && defined(__cpp_lib_semaphore) && \
+    __cpp_lib_semaphore >= 201907L && !defined(__GLIBCXX__)
 #define C10_SEMAPHORE_USE_STL
 #endif
 
diff --git a/c10/util/complex_math.h b/c10/util/complex_math.h
index 2b591026c94da..d369df5059231 100644
--- a/c10/util/complex_math.h
+++ b/c10/util/complex_math.h
@@ -86,6 +86,41 @@ C10_HOST_DEVICE inline c10::complex<T> pow(
 #endif
 }
 
+// Regression in ROCm 7.2. See https://github.com/ROCm/rocm-libraries/pull/3836.
+// Specialized version for complex<float> on AMD GPUs to use FMA-based
+// multiplication
+#if defined(__HIPCC__)
+namespace detail {
+// FMA-aware complex multiplication for float precision on AMD GPUs.
+// This prevents SLP vectorizer from breaking FMA formation, which causes
+// numerical precision loss in complex arithmetic.
+// The issue occurs when vectorizer packs scalar multiplies before backend
+// can form FMA instructions, resulting in double rounding instead of single.
+C10_HOST_DEVICE inline thrust::complex<float> complex_mul_fma(
+    thrust::complex<float> a,
+    thrust::complex<float> b) {
+  // Complex multiplication: (a.r + a.i*i) * (b.r + b.i*i)
+  // = (a.r*b.r - a.i*b.i) + (a.r*b.i + a.i*b.r)*i
+  // Using __builtin_fmaf ensures FMA at source level:
+  // real: a.r*b.r + (-(a.i*b.i)) = FMA(a.r, b.r, -(a.i*b.i))
+  // imag: a.i*b.r + a.r*b.i = FMA(a.r, b.i, a.i*b.r)
+  float real_part = __builtin_fmaf(a.real(), b.real(), -(a.imag() * b.imag()));
+  float imag_part = __builtin_fmaf(a.real(), b.imag(), a.imag() * b.real());
+  return thrust::complex<float>(real_part, imag_part);
+}
+} // namespace detail
+
+template <>
+C10_HOST_DEVICE inline c10::complex<float> pow(
+    const c10::complex<float>& x,
+    const c10::complex<float>& y) {
+  auto log_x = thrust::log(static_cast<thrust::complex<float>>(x));
+  auto y_log_x =
+      detail::complex_mul_fma(static_cast<thrust::complex<float>>(y), log_x);
+  return static_cast<c10::complex<float>>(thrust::exp(y_log_x));
+}
+#endif
+
 template <typename T>
 C10_HOST_DEVICE inline c10::complex<T> pow(
     const c10::complex<T>& x,
diff --git a/c10/util/generic_math.h b/c10/util/generic_math.h
index 9c2a424560484..e98be901c8c53 100644
--- a/c10/util/generic_math.h
+++ b/c10/util/generic_math.h
@@ -59,6 +59,10 @@ inline C10_HOST_DEVICE scalar_t div_floor_floating(scalar_t a, scalar_t b)
 
 template <typename scalar_t>
 inline C10_HOST_DEVICE scalar_t div_floor_integer(scalar_t a, scalar_t b) {
+  if (C10_UNLIKELY(b == 0)) {
+    return scalar_t(0);
+  }
+
   if (C10_UNLIKELY(
           std::is_signed<scalar_t>::value &&
           a == std::numeric_limits<scalar_t>::min() && b == scalar_t(-1))) {
diff --git a/c10/xpu/CMakeLists.txt b/c10/xpu/CMakeLists.txt
index c2fa65ba35e73..349b209aec531 100644
--- a/c10/xpu/CMakeLists.txt
+++ b/c10/xpu/CMakeLists.txt
@@ -15,12 +15,14 @@ configure_file(
     ${CMAKE_BINARY_DIR}/c10/xpu/impl/xpu_cmake_macros.h)
 
 set(C10_XPU_SRCS
+    PeerToPeerAccess.cpp
     XPUCachingAllocator.cpp
     XPUFunctions.cpp
     XPUStream.cpp
     impl/XPUGuardImpl.cpp
 )
 set(C10_XPU_HEADERS
+    PeerToPeerAccess.h
     XPUCachingAllocator.h
     XPUDeviceProp.h
     XPUException.h
diff --git a/aten/src/ATen/xpu/PeerToPeerAccess.cpp b/c10/xpu/PeerToPeerAccess.cpp
similarity index 88%
rename from aten/src/ATen/xpu/PeerToPeerAccess.cpp
rename to c10/xpu/PeerToPeerAccess.cpp
index 60c0fef63e68c..e3ef2d129650d 100644
--- a/aten/src/ATen/xpu/PeerToPeerAccess.cpp
+++ b/c10/xpu/PeerToPeerAccess.cpp
@@ -1,11 +1,8 @@
-#include <ATen/xpu/PeerToPeerAccess.h>
-#include <ATen/xpu/XPUContext.h>
-
-#include <c10/util/Exception.h>
 #include <c10/util/irange.h>
+#include <c10/xpu/PeerToPeerAccess.h>
 #include <c10/xpu/XPUCachingAllocator.h>
 
-namespace at::xpu {
+namespace c10::xpu {
 
 // p2pAccessEnabled_ is a flattened 2D matrix of size [num_devices x
 // num_devices].
@@ -35,8 +32,6 @@ void init_p2p_access_cache(c10::DeviceIndex num_devices) {
 } // namespace detail
 
 bool get_p2p_access(c10::DeviceIndex dev, c10::DeviceIndex dev_to_access) {
-  at::globalContext().lazyInitDevice(c10::DeviceType::XPU);
-
   check_device_index(dev);
   check_device_index(dev_to_access);
 
@@ -60,4 +55,4 @@ bool get_p2p_access(c10::DeviceIndex dev, c10::DeviceIndex dev_to_access) {
   return static_cast<bool>(cache);
 }
 
-} // namespace at::xpu
+} // namespace c10::xpu
diff --git a/c10/xpu/PeerToPeerAccess.h b/c10/xpu/PeerToPeerAccess.h
new file mode 100644
index 0000000000000..5fed6c90a4e4a
--- /dev/null
+++ b/c10/xpu/PeerToPeerAccess.h
@@ -0,0 +1,18 @@
+#pragma once
+
+#include <c10/core/Device.h>
+#include <c10/macros/Macros.h>
+#include <c10/xpu/XPUMacros.h>
+
+namespace c10::xpu {
+namespace detail {
+// Initialize the peer-to-peer access cache for XPU devices.
+C10_XPU_API void init_p2p_access_cache(c10::DeviceIndex num_devices);
+} // namespace detail
+
+// Query if peer-to-peer access is available between two devices.
+C10_XPU_API bool get_p2p_access(
+    c10::DeviceIndex dev,
+    c10::DeviceIndex dev_to_access);
+
+} // namespace c10::xpu
diff --git a/c10/xpu/XPUCachingAllocator.cpp b/c10/xpu/XPUCachingAllocator.cpp
index 7e534c7aa403b..f01e0c5c15224 100644
--- a/c10/xpu/XPUCachingAllocator.cpp
+++ b/c10/xpu/XPUCachingAllocator.cpp
@@ -1547,7 +1547,7 @@ class DeviceCachingAllocator {
       std::shared_ptr<GatheredContext> context) {
     if (!record_history)
       return;
-    bool should_skip = skip_actions_list.count(action) > 0;
+    bool should_skip = skip_actions_list.contains(action);
     if (should_skip)
       return;
     TraceEntry te(
diff --git a/c10/xpu/XPUDeviceProp.h b/c10/xpu/XPUDeviceProp.h
index a06fa8e5cdef5..d1dd85b29d05d 100644
--- a/c10/xpu/XPUDeviceProp.h
+++ b/c10/xpu/XPUDeviceProp.h
@@ -129,8 +129,14 @@ namespace c10::xpu {
   /* the device identifier of the Intel GPU, also known as the product ID. */ \
   _(device_id, device_id, 0)                                                  \
                                                                               \
-  /* the device descriptor for device Universal Unique ID, 16 bytes*/         \
-  _(uuid, device_info_uuid, (std::array<unsigned char, 16>{}))
+  /* the device descriptor for device Universal Unique ID, 16 bytes. */       \
+  _(uuid, device_info_uuid, (std::array<unsigned char, 16>{}))                \
+                                                                              \
+  /* the maximum clock rate of device's global memory in MHz. */              \
+  _(memory_clock_rate, memory_clock_rate, 0)                                  \
+                                                                              \
+  /* the maximum bus width between device and memory in bits. */              \
+  _(memory_bus_width, memory_bus_width, 0)
 
 #define AT_FORALL_XPU_DEVICE_ASPECT(_)                  \
   /* sycl::half is supported on device. */              \
@@ -180,20 +186,24 @@ namespace c10::xpu {
   _DEFINE_SYCL_PROP(                     \
       sycl::ext::oneapi::experimental::info::device, property, property)
 
-struct C10_XPU_API DeviceProp {
-  AT_FORALL_XPU_DEVICE_PROPERTIES(DEFINE_DEVICE_PROP);
+struct C10_XPU_API DeviceProp{
+    AT_FORALL_XPU_DEVICE_PROPERTIES(DEFINE_DEVICE_PROP)
 
-  // the platform name.
-  DEFINE_PLATFORM_PROP(name, platform_name);
+    // the platform name.
+    DEFINE_PLATFORM_PROP(name, platform_name)
 
-  AT_FORALL_XPU_EXT_DEVICE_PROPERTIES(DEFINE_EXT_DEVICE_PROP);
+    // ext properties.
+    AT_FORALL_XPU_EXT_DEVICE_PROPERTIES(DEFINE_EXT_DEVICE_PROP)
 
-  AT_FORALL_XPU_DEVICE_ASPECT(DEFINE_DEVICE_ASPECT);
+    // device aspects.
+    AT_FORALL_XPU_DEVICE_ASPECT(DEFINE_DEVICE_ASPECT)
 
-  AT_FORALL_XPU_EXP_CL_ASPECT(DEFINE_DEVICE_ASPECT);
+    // experimental device aspects.
+    AT_FORALL_XPU_EXP_CL_ASPECT(DEFINE_DEVICE_ASPECT)
 
 #if SYCL_COMPILER_VERSION >= 20250000
-  AT_FORALL_XPU_EXP_DEVICE_PROPERTIES(DEFINE_EXP_DEVICE_PROP);
+    // experimental device properties.
+    AT_FORALL_XPU_EXP_DEVICE_PROPERTIES(DEFINE_EXP_DEVICE_PROP)
 #endif
 };
 
diff --git a/c10/xpu/XPUStream.h b/c10/xpu/XPUStream.h
index c922759c2c48b..396002eb211b9 100644
--- a/c10/xpu/XPUStream.h
+++ b/c10/xpu/XPUStream.h
@@ -109,6 +109,11 @@ class C10_XPU_API XPUStream {
     }
   }
 
+  bool is_capturing() const {
+    return queue().ext_oneapi_get_state() ==
+        sycl::ext::oneapi::experimental::queue_state::recording;
+  }
+
   /// Return the priority that this stream is associated with. Lower numbers
   /// represent higher priority.
   int priority() const;
diff --git a/c10/xpu/impl/XPUGuardImpl.h b/c10/xpu/impl/XPUGuardImpl.h
index ad788ceb88141..d9a3d83ee4488 100644
--- a/c10/xpu/impl/XPUGuardImpl.h
+++ b/c10/xpu/impl/XPUGuardImpl.h
@@ -50,7 +50,11 @@ struct XPUGuardImpl final : public c10::impl::DeviceGuardImplInterface {
     cap.capability_data.capability_bits = (1ULL << kIndex_Byte) |
         (1ULL << kIndex_Char) | (1ULL << kIndex_Short) | (1ULL << kIndex_Int) |
         (1ULL << kIndex_Long) | (1ULL << kIndex_Float) |
-        (1ULL << kIndex_ComplexFloat) | (1ULL << kIndex_Bool);
+        (1ULL << kIndex_ComplexFloat) | (1ULL << kIndex_Bool) |
+        (1ULL << kIndex_Float8_e5m2) | (1ULL << kIndex_Float8_e4m3fn) |
+        (1ULL << kIndex_Float8_e5m2fnuz) | (1ULL << kIndex_Float8_e4m3fnuz) |
+        (1ULL << kIndex_Float8_e8m0fnu) | (1ULL << kIndex_UInt16) |
+        (1ULL << kIndex_UInt32) | (1ULL << kIndex_UInt64);
     // BFloat16 may be emulated. We always assume BFloat16 is available;
     // users can call is_bf16_supported() to check for native hardware support.
     cap.capability_data.capability_bits |= (1ULL << kIndex_BFloat16);
@@ -214,6 +218,11 @@ struct XPUGuardImpl final : public c10::impl::DeviceGuardImplInterface {
     xpu_stream.synchronize();
   }
 
+  bool isStreamCapturing(const Stream& stream) const override {
+    const XPUStream xpu_stream{stream};
+    return xpu_stream.is_capturing();
+  }
+
   void synchronizeEvent(void* event) const override {
     if (!event)
       return;
diff --git a/caffe2/.clang-format b/caffe2/.clang-format
index 7263446d248e8..6ade54308e8d7 100644
--- a/caffe2/.clang-format
+++ b/caffe2/.clang-format
@@ -101,7 +101,7 @@ SpacesInContainerLiterals: true
 SpacesInCStyleCastParentheses: false
 SpacesInParentheses: false
 SpacesInSquareBrackets: false
-Standard:        c++17
+Standard:        c++20
 StatementMacros:
   - PyObject_HEAD
   - PyObject_VAR_HEAD
diff --git a/caffe2/CMakeLists.txt b/caffe2/CMakeLists.txt
index 0a419e46a5bce..f5a7ba726ac8c 100644
--- a/caffe2/CMakeLists.txt
+++ b/caffe2/CMakeLists.txt
@@ -53,6 +53,9 @@ if(INTERN_BUILD_ATEN_OPS)
 
   # Add source, includes, and libs to lists
   list(APPEND Caffe2_CPU_SRCS ${ATen_CPU_SRCS})
+  if(USE_MTIA)
+    list(APPEND Caffe2_CPU_SRCS ${ATen_MTIA_SRCS})
+  endif()
   list(APPEND Caffe2_GPU_SRCS ${ATen_CUDA_CPP_SRCS})
   list(APPEND Caffe2_XPU_SRCS ${ATen_XPU_SRCS})
   list(APPEND Caffe2_XPU_INCLUDE ${ATen_XPU_INCLUDE})
@@ -556,13 +559,6 @@ if(USE_CUDA OR USE_ROCM)
   append_filelist("libtorch_cuda_core_sources" Caffe2_GPU_HIP_JIT_FUSERS_SRCS)
 endif()
 
-# NativeRT is disabled
-# if(USE_CUDA)
-#   append_filelist("libtorch_nativert_cuda_sources" Caffe2_GPU_SRCS)
-# endif()
-# if(USE_ROCM)
-#   append_filelist("libtorch_nativert_cuda_sources" Caffe2_HIP_SRCS)
-# endif()
 
 if(USE_CUDA)
   list(APPEND Caffe2_GPU_CU_SRCS ${Caffe2_GPU_HIP_JIT_FUSERS_SRCS})
@@ -594,6 +590,7 @@ if(USE_CUDA)
         ${TORCH_SRC_DIR}/csrc/distributed/c10d/symm_mem/CUDASymmetricMemoryUtils.cpp
         ${TORCH_SRC_DIR}/csrc/distributed/c10d/symm_mem/NCCLSymmetricMemory.cu
         ${TORCH_SRC_DIR}/csrc/distributed/c10d/symm_mem/nccl_extension.cu
+        ${TORCH_SRC_DIR}/csrc/distributed/c10d/symm_mem/ops/nccl_reduce_scatter_offset.cu
         ${TORCH_SRC_DIR}/csrc/distributed/c10d/symm_mem/cuda_mem_pool.cpp
         PROPERTIES COMPILE_FLAGS "-DPYTORCH_C10_DRIVER_API_SUPPORTED=1"
       )
@@ -604,7 +601,7 @@ if(USE_CUDA)
     if(CMAKE_COMPILER_IS_GNUCXX)
       set_source_files_properties(${ASYNC_MM_FILE} PROPERTIES COMPILE_FLAGS "-Wno-unused-but-set-variable")
     endif()
-    if(CMAKE_CUDA_COMPILER_VERSION VERSION_GREATER_EQUAL 12.0 AND CUDA_NVCC_FLAGS MATCHES ".*compute_90.*")
+    if(CUDA_NVCC_FLAGS MATCHES ".*compute_90.*")
       set_source_files_properties(${ASYNC_MM_FILE} PROPERTIES COMPILE_FLAGS "-gencode arch=compute_90a,code=sm_90a")
     endif()
   endif()
@@ -826,7 +823,7 @@ if("${CMAKE_CXX_COMPILER_ID}" MATCHES "Clang" AND NOT USE_IOS AND NOT USE_COREML
   endif()
   get_target_property(TORCH_CPU_SOURCES torch_cpu SOURCES)
   foreach(generated_file IN LISTS GENERATED_CXX_TORCH)
-    set_source_files_properties(${generated_file} PROPERTIES COMPILE_OPTIONS "-Wno-missing-prototypes;-Wno-error=missing-prototypes")
+    set_source_files_properties(${generated_file} PROPERTIES COMPILE_OPTIONS "-Wno-missing-prototypes;-Wno-error=missing-prototypes;-mllvm;-openmp-opt-disable")
   endforeach()
   foreach(source_file IN LISTS TORCH_CPU_SOURCES)
     get_filename_component(source_file "${source_file}" REALPATH)
@@ -845,6 +842,10 @@ if(USE_MPS)
   if(CAN_COMPILE_METAL)
     add_dependencies(torch_cpu metallibs)
     target_link_options(torch_cpu PRIVATE -Wl,-sectcreate,__TEXT,metal_basic,${CMAKE_CURRENT_BINARY_DIR}/aten/src/ATen/kernels_basic.metallib)
+    if(CAN_COMPILE_METAL_40)
+      target_link_options(torch_cpu PRIVATE -Wl,-sectcreate,__TEXT,metal_40,${CMAKE_CURRENT_BINARY_DIR}/aten/src/ATen/kernels_40.metallib)
+      target_compile_definitions(torch_cpu PRIVATE CAN_BUILD_METAL_4)
+    endif()
   else()
     target_compile_definitions(torch_cpu PRIVATE PYTORCH_JIT_COMPILE_SHADERS)
   endif()
@@ -943,7 +944,7 @@ if(USE_ROCM)
   set(CUDA_LINK_LIBRARIES_KEYWORD PRIVATE)
   list(APPEND Caffe2_HIP_SRCS ${GENERATED_CXX_TORCH_CUDA})
   hip_add_library(torch_hip ${Caffe2_HIP_SRCS})
-  if(USE_FLASH_ATTENTION)
+  if(USE_FLASH_ATTENTION OR USE_MEM_EFF_ATTENTION)
     target_link_libraries(torch_hip PRIVATE __caffe2_aotriton)
   endif()
   set(CUDA_LINK_LIBRARIES_KEYWORD)
@@ -954,6 +955,12 @@ if(USE_ROCM)
     target_compile_definitions(torch_hip PRIVATE USE_NCCL)
   endif()
 
+  # IntraNodeComm loads libamd_smi at runtime via dlopen (avoids linking amd_smi
+  # which would cause bus errors if Python amdsmi loads a different copy).
+  if(USE_DISTRIBUTED)
+    target_link_libraries(torch_hip PRIVATE ${CMAKE_DL_LIBS})
+  endif()
+
   if(USE_PRECOMPILED_HEADERS)
     target_precompile_headers(torch_hip PRIVATE
         "$<$<COMPILE_LANGUAGE:CXX>:ATen/core/ATen_pch.h>")
@@ -1042,7 +1049,7 @@ elseif(USE_CUDA)
     add_library(torch_nvshmem SHARED
         "${TORCH_SRC_DIR}/csrc/distributed/c10d/cuda/utils.cpp"
         "${TORCH_SRC_DIR}/csrc/distributed/c10d/symm_mem/nvshmem_extension.cu"
-        "${TORCH_SRC_DIR}/csrc/distributed/c10d/symm_mem/NVSHMEMSymmetricMemory.cu"
+        "${TORCH_SRC_DIR}/csrc/distributed/c10d/symm_mem/NVSHMEMSymmetricMemory.cpp"
         "${TORCH_SRC_DIR}/csrc/distributed/c10d/symm_mem/CUDASymmetricMemoryUtils.cpp"
     )
     # Add torch_cpu as a dependency to ensure ATen headers are available
@@ -1363,8 +1370,7 @@ if(BUILD_TEST)
   else()
     add_subdirectory(${TORCH_ROOT}/test/cpp/jit ${CMAKE_BINARY_DIR}/test_jit)
     add_subdirectory(${TORCH_ROOT}/test/cpp/lazy ${CMAKE_BINARY_DIR}/test_lazy)
-    # NativeRT is disabled
-    # add_subdirectory(${TORCH_ROOT}/test/cpp/nativert ${CMAKE_BINARY_DIR}/test_nativert)
+    add_subdirectory(${TORCH_ROOT}/test/cpp/profiler ${CMAKE_BINARY_DIR}/test_profiler)
     add_subdirectory(${TORCH_ROOT}/test/inductor ${CMAKE_BINARY_DIR}/test_inductor)
     add_subdirectory(${TORCH_ROOT}/test/cpp/aoti_abi_check ${CMAKE_BINARY_DIR}/test_aoti_abi_check)
     if(BUILD_AOT_INDUCTOR_TEST)
@@ -1461,6 +1467,74 @@ if(USE_ROCM)
   if(USE_ROCM_CK_SDPA)
     target_compile_definitions(torch_hip PRIVATE USE_ROCM_CK_SDPA)
   endif()
+
+  message(INFO "USE_NVSHMEM=${USE_NVSHMEM} USE_ROCM=${USE_ROCM} rocshmem_FOUND=${rocshmem_FOUND}")
+  if(USE_NVSHMEM)
+    # rocSHMEM upstream ships device code for a subset of AMD gfx targets only.
+    set(ROCSHMEM_SUPPORTED_ARCH gfx90a gfx942 gfx950 gfx1100 gfx1201)
+    set(_torch_rocshmem_build_arches)
+    foreach(_arch ${PYTORCH_ROCM_ARCH})
+      list(FIND ROCSHMEM_SUPPORTED_ARCH "${_arch}" _rocsmem_supported_idx)
+      if(NOT _rocsmem_supported_idx EQUAL -1)
+        # rocSHMEM device archives are commonly built for gfx90a xnack variants,
+        # so map plain gfx90a to gfx90a:xnack-/gfx90a:xnack+ to match device
+        # symbols during HIP RDC link.
+        if(_arch STREQUAL "gfx90a")
+          list(APPEND _torch_rocshmem_build_arches "gfx90a:xnack-" "gfx90a:xnack+")
+        else()
+          list(APPEND _torch_rocshmem_build_arches "${_arch}")
+        endif()
+      endif()
+    endforeach()
+    if(rocshmem_FOUND AND _torch_rocshmem_build_arches)
+      message(STATUS "rocSHMEM found, building with rocSHMEM support for arches: ${_torch_rocshmem_build_arches}")
+      message(STATUS "ROCSHMEM_INCLUDE_DIR: '${ROCSHMEM_INCLUDE_DIR}'")
+      set(TORCH_ROCSHMEM_SRCS
+        "${TORCH_SRC_DIR}/csrc/distributed/c10d/cuda/utils.cpp"
+        "${TORCH_SRC_DIR}/csrc/distributed/c10d/symm_mem/rocshmem_extension.cu"
+        "${TORCH_SRC_DIR}/csrc/distributed/c10d/symm_mem/NVSHMEMSymmetricMemory.cpp"
+        "${TORCH_SRC_DIR}/csrc/distributed/c10d/symm_mem/CUDASymmetricMemoryUtils.cpp"
+      )
+      set_source_files_properties(${TORCH_ROCSHMEM_SRCS} PROPERTIES LANGUAGE HIP)
+
+      # FindHIP's generated compile script bakes in global HIP_CLANG_FLAGS. Override
+      # -fno-gpu-rdc with -fgpu-rdc so device linker can pull in rocSHMEM device APIs.
+      set(HIP_CLANG_FLAGS_SAVED ${HIP_CLANG_FLAGS})
+      set(HIP_CLANG_FLAGS_FOR_ROCSHMEM ${HIP_CLANG_FLAGS})
+      list(FILTER HIP_CLANG_FLAGS_FOR_ROCSHMEM EXCLUDE REGEX "^-fno-gpu-rdc$")
+      list(APPEND HIP_CLANG_FLAGS_FOR_ROCSHMEM -fgpu-rdc)
+      set(HIP_CLANG_FLAGS ${HIP_CLANG_FLAGS_FOR_ROCSHMEM})
+
+      hip_add_library(torch_rocshmem SHARED ${TORCH_ROCSHMEM_SRCS})
+
+      set(HIP_CLANG_FLAGS ${HIP_CLANG_FLAGS_SAVED})
+
+      set_target_properties(torch_rocshmem PROPERTIES LINKER_LANGUAGE HIP)
+
+
+      # torch_rocshmem is built with -fgpu-rdc, so its HIP link step performs
+      # device linking and must carry the full arch list explicitly.
+      set_target_properties(
+          torch_rocshmem PROPERTIES HIP_ARCHITECTURES "${_torch_rocshmem_build_arches}")
+
+
+      target_compile_definitions(torch_hip PUBLIC USE_NVSHMEM)
+      target_compile_definitions(torch_rocshmem PUBLIC USE_NVSHMEM USE_ROCM)
+
+      target_link_libraries(torch_rocshmem PRIVATE roc::rocshmem)
+      target_link_libraries(torch_hip PRIVATE torch_rocshmem)
+
+      install(TARGETS torch_rocshmem EXPORT Caffe2Targets DESTINATION
+        "${TORCH_INSTALL_LIB_DIR}")
+    elseif(rocshmem_FOUND)
+      message(STATUS
+        "rocSHMEM found but skipped: PYTORCH_ROCM_ARCH has no entry in "
+        "ROCSHMEM_SUPPORTED_ARCH (${ROCSHMEM_SUPPORTED_ARCH}). "
+        "Building without rocSHMEM support.")
+    else()
+      message(WARNING "rocSHMEM not found (USE_NVSHMEM=ON). Building without rocSHMEM support.")
+    endif() #rocshmem_FOUND / arch intersection
+  endif() #USE_NVSHMEM
 endif()
 
 if(BUILD_LITE_INTERPRETER)
@@ -1792,6 +1866,10 @@ if(USE_ROCM)
     target_link_libraries(torch_hip PRIVATE ck_sdpa)
   endif()
 
+  if(TARGET ck_gemm)
+    target_link_libraries(torch_hip PRIVATE ck_gemm)
+  endif()
+
   if(USE_MSLK)
     if(USE_ROCM)
       target_link_libraries(torch_hip PRIVATE mslk)
@@ -1983,7 +2061,11 @@ if(BUILD_TEST)
   endif()
 
   if(USE_ROCM)
+    string(JOIN " " HIP_HIPCC_FLAGS_STR ${HIP_HIPCC_FLAGS})
+    set(HIP_HIPCC_FLAGS ${HIP_HIPCC_FLAGS_STR})
+
     set(BASE_HIPCC_FLAGS ${HIP_HIPCC_FLAGS})
+
     foreach(test_src ${Caffe2_HIP_TEST_SRCS})
       get_filename_component(test_name ${test_src} NAME_WE)
       if(test_src MATCHES "^.*\.hip$")
diff --git a/caffe2/core/macros.h.in b/caffe2/core/macros.h.in
index 5b971b1bbb278..ba0e20df271ed 100644
--- a/caffe2/core/macros.h.in
+++ b/caffe2/core/macros.h.in
@@ -34,6 +34,7 @@
   {"TORCH_VERSION", "${TORCH_VERSION}"}, \
   {"CXX_COMPILER", "${CMAKE_CXX_COMPILER}"}, \
   {"CXX_FLAGS", "${CMAKE_CXX_FLAGS}"}, \
+  {"CUDA_FLAGS", "${CMAKE_CUDA_FLAGS}"}, \
   {"BUILD_TYPE", "${CMAKE_BUILD_TYPE}"}, \
   {"BLAS_INFO", "${BLAS_INFO}"}, \
   {"LAPACK_INFO", "${LAPACK_INFO}"}, \
@@ -67,4 +68,5 @@
   {"USE_CUSPARSELT", "${USE_CUSPARSELT}"}, \
   {"USE_XPU", "${USE_XPU}"}, \
   {"USE_XCCL", "${USE_XCCL}"}, \
+  {"SYCL_COMPILER_VERSION", "${SYCL_COMPILER_VERSION}"}, \
 }
diff --git a/caffe2/perfkernels/batch_box_cox_avx512.cc b/caffe2/perfkernels/batch_box_cox_avx512.cc
deleted file mode 100644
index a97cb364a359d..0000000000000
--- a/caffe2/perfkernels/batch_box_cox_avx512.cc
+++ /dev/null
@@ -1,118 +0,0 @@
-#ifdef CAFFE2_PERF_USE_MKL
-#include <immintrin.h>
-
-// Enable compiler vectorized version only if numerical consistency is not
-// required between dev and opt versions - disabled for now
-#ifndef FAST_VECTORIZED_KERNEL
-#define CPU_CAPABILITY_AVX512
-#include <ATen/cpu/vec/vec.h>
-
-namespace at::vec {
-namespace {
-// Implements the vectorized version of std::max() operation,
-// which DOESNOT propagates NaN for second argument
-template <typename scalar_t>
-Vectorized<scalar_t> max(const Vectorized<scalar_t>& a, const Vectorized<scalar_t>& b);
-
-template <>
-Vectorized<double> max(const Vectorized<double>& a, const Vectorized<double>& b) {
-  // std::max(NaN, nonNan) -> NaN
-  return _mm512_max_pd(b, a);
-}
-
-template <>
-Vectorized<float> max(const Vectorized<float>& a, const Vectorized<float>& b) {
-  // std::max(NaN, nonNan) -> NaN
-  return _mm512_max_ps(b, a);
-}
-
-// Implements recieprocal method based on newton-rapson method
-// 1. user RCP approximiation
-// 2. update with RCP = RCP * (2 - X * RCP)
-template <typename scalar_t>
-Vectorized<scalar_t> fast_recieprocal(const Vectorized<scalar_t>& b);
-template <typename scalar_t>
-scalar_t fast_recieprocal(scalar_t b);
-
-template<>
-Vectorized<float> fast_recieprocal(const Vectorized<float>& b) {
-  auto minus2 = _mm512_set1_ps(-2.f);
-  auto rcp = _mm512_rcp14_ps(b);
-  rcp = _mm512_mul_ps(rcp,  _mm512_fnmsub_ps(rcp, b, minus2));
-  rcp = _mm512_mul_ps(rcp,  _mm512_fnmsub_ps(rcp, b, minus2));
-  return rcp;
-}
-
-template <>
-float fast_recieprocal(float b) {
-  auto minus2 = _mm_set_ss(-2.f);
-  auto b_reg = _mm_set_ss(b);
-  auto rcp = _mm_rcp_ss(b_reg);
-  rcp = _mm_mul_ss(rcp,  _mm_fnmsub_ss(rcp, b_reg, minus2));
-  rcp = _mm_mul_ss(rcp,  _mm_fnmsub_ss(rcp, b_reg, minus2));
-  return _mm_cvtss_f32(rcp);
-}
-
-template<>
-Vectorized<double> fast_recieprocal(const Vectorized<double>& b) {
-  auto minus2 = _mm512_set1_pd(-2.);
-  auto rcp = _mm512_rcp14_pd(b);
-  rcp = _mm512_mul_pd(rcp,  _mm512_fnmsub_pd(rcp, b, minus2));
-  rcp = _mm512_mul_pd(rcp,  _mm512_fnmsub_pd(rcp, b, minus2));
-  return rcp;
-}
-
-template <>
-double fast_recieprocal(double b) {
-  return 1./b;
-}
-} // namespace
-} // namespace at::vec
-#endif
-
-#include "caffe2/perfkernels/batch_box_cox_vec.h"
-
-namespace caffe2::details {
-
-template <typename T>
-void compute_batch_box_cox__avx512(
-    std::size_t N,
-    std::size_t D,
-    std::size_t block_size,
-    const T* self_data,
-    const T* __restrict lambda1_data,
-    const T* __restrict lambda2_data,
-    T* output_data) {
-      compute_batch_box_cox_vec_fma<T>(
-          N,
-          D,
-          block_size,
-          self_data,
-          lambda1_data,
-          lambda2_data,
-          output_data);
-    }
-
-// Vectorized version specializations for float and double
-template
-void compute_batch_box_cox__avx512<float>(
-  std::size_t N,
-  std::size_t D,
-  std::size_t block_size,
-  const float* self_data,
-  const float* __restrict lambda1_data,
-  const float* __restrict lambda2_data,
-  float* output_data);
-
-template
-void compute_batch_box_cox__avx512<double>(
-  std::size_t N,
-  std::size_t D,
-  std::size_t block_size,
-  const double* self_data,
-  const double* __restrict lambda1_data,
-  const double* __restrict lambda2_data,
-  double* output_data);
-
-} // namespace caffe2::detail
-#endif // CAFFE2_PERF_USE_MKL
diff --git a/caffe2/perfkernels/batch_box_cox_sve128.cc b/caffe2/perfkernels/batch_box_cox_sve128.cc
deleted file mode 100644
index 897e3a8ee4755..0000000000000
--- a/caffe2/perfkernels/batch_box_cox_sve128.cc
+++ /dev/null
@@ -1,200 +0,0 @@
-#if defined(__aarch64__) && defined(__ARM_FEATURE_SVE) && defined(CAFFE2_PERF_WITH_SVE128)
-#include <arm_neon.h>
-#include <arm_neon_sve_bridge.h>
-#include <arm_sve.h>
-#include <cfloat>
-#include <cmath>
-
-#include "c10/macros/Macros.h"
-
-/// Select `svlog` accuracy:
-/// - 0: original.
-/// - 1: more accurate, similar performance.
-/// - 2: very high accuracy, a bit lower speed.
-#define SVLOG_ACCURACY 2
-
-/// Handle special cases in `svexp`:
-/// - 0: original.
-/// - 1: use clamp, better performance.
-/// - 2: no special case handling.
-#define SVEXP_SPECIAL_CLAMP 1
-
-#if SVLOG_ACCURACY == 2
-static inline svfloat32_t svlog(svfloat32_t x) {
-  const svbool_t ptrue = svptrue_b8();
-
-  svint32_t u = svreinterpret_s32(x) - 0x3F2AAAAB;
-
-  svfloat32_t r = svreinterpret_f32((u & 0x007FFFFF) + 0x3F2AAAAB) - 1.0f;
-  svfloat32_t n = svcvt_f32_x(ptrue, u >> 23);
-  asm("" : "+w"(r)); // NOTE: can improve instruction scheduling.
-
-  svfloat32_t r2 = r * r;
-  svfloat32_t p = -0x1.4F9934p-3f + r * 0x1.5A9AA2p-3f;
-  svfloat32_t q = -0x1.00187Cp-2f + r * 0x1.961348p-3f;
-  svfloat32_t y = -0x1.FFFFC8p-2f + r * 0x1.555D7Cp-2f;
-  return (r + n * 0x1.62E43p-1f) +
-         (y + (q + (p + -0x1.3E737Cp-3f * r2) * r2) * r2) * r2;
-}
-#elif SVLOG_ACCURACY == 1
-static inline svfloat32_t svlog(svfloat32_t x) {
-  const svbool_t ptrue = svptrue_b8();
-
-  svint32_t u = svreinterpret_s32(x) - 0x3F2AAAAB;
-
-  svfloat32_t r = svreinterpret_f32((u & 0x007FFFFF) + 0x3F2AAAAB) - 1.0f;
-  svfloat32_t n = svcvt_f32_x(ptrue, u >> 23);
-  asm("" : "+w"(r)); // NOTE: can improve instruction scheduling.
-
-  svfloat32_t r2 = r * r;
-  svfloat32_t A = -0x1.923814p-3f + r * 0x1.689E5Ep-3f;
-  svfloat32_t B = -0x1.FC0968p-3f + r * 0x1.93BF0Cp-3f;
-  svfloat32_t C = -0x1.000478p-1f + r * 0x1.556906p-2f;
-
-  return (r + n * 0x1.62E43p-1f) + (C + (B + A * r2) * r2) * r2;
-}
-#elif SVLOG_ACCURACY == 0
-static inline svfloat32_t svlog(svfloat32_t x) {
-  const svbool_t ptrue = svptrue_b8();
-
-  svint32_t u = svsra_n_s32(svdup_n_s32(-127), svreinterpret_s32(x), 23);
-
-  svfloat32_t n = svcvt_f32_x(ptrue, u);
-  svfloat32_t r = svreinterpret_f32(svreinterpret_s32(x) - (u << 23));
-
-  svfloat32_t D = -0.165253549814f + r * 0.0141278216615f;
-  svfloat32_t C = -2.47071170807f + r * 0.844007015228f;
-  svfloat32_t B = -5.68692588806f + r * 4.58445882797f;
-  svfloat32_t A = -2.29561495781f + r * 5.17591238022f;
-
-  svfloat32_t r2 = r * r;
-  return (A + n * 0.6931471805f) + (B + (C + D * r2) * r2) * r2;
-}
-#endif
-
-static inline svfloat32_t svexp(svfloat32_t x) {
-  // Clamp interval set to prevent denormals!
-  const svfloat32_t max_input = svdup_n_f32(88.722839f);
-  const svfloat32_t min_input = svdup_n_f32(-87.33654f);
-  const svfloat32_t shift = svdup_n_f32(0x1.0000FEp+23f);
-  const svbool_t ptrue = svptrue_b8();
-
-#if SVEXP_SPECIAL_CLAMP == 1
-  x = svmax_x(ptrue, svmin_x(ptrue, x, max_input), min_input);
-#endif
-
-  svfloat32_t z = svmla_n_f32_x(ptrue, shift, x, 0x1.715476p+0f);
-  svfloat32_t n = z - shift;
-  svfloat32_t scale = svreinterpret_f32(svreinterpret_u32(z) << 23);
-
-  svfloat32_t r_hi = x - n * 0x1.62E400p-1f;
-  svfloat32_t r = r_hi - n * 0x1.7F7D1Cp-20f;
-  svfloat32_t r2 = r * r;
-
-  svfloat32_t C = 0x1.573E2Ep-5f + r * 0x1.0E4020p-7f;
-  svfloat32_t B = 0x1.FFFDB6p-2f + r * 0x1.555E66p-3f;
-  svfloat32_t A = r * 0x1.FFFFECp-1f;
-
-  svfloat32_t poly = scale + (A + (B + C * r2) * r2) * scale;
-
-#if SVEXP_SPECIAL_CLAMP == 0
-  const svfloat32_t inf = svdup_n_f32(std::numeric_limits<float>::infinity());
-  poly = svsel_f32(svcmplt_f32(ptrue, x, min_input), svdup_n_f32(0.0f), poly);
-  poly = svsel_f32(svcmpgt_f32(ptrue, x, max_input), inf, poly);
-#endif
-
-  return poly;
-}
-
-static inline svfloat32_t compute_batch_box_cox_vec_sve128_float(
-    svfloat32_t lambda1_v,
-    svfloat32_t lambda2_v,
-    svfloat32_t data_v,
-    svfloat32_t k_eps) {
-  const svbool_t ptrue = svptrue_b8();
-
-  svfloat32_t lnData = svlog(svmax_x(ptrue, data_v + lambda2_v, k_eps));
-  svbool_t predNZ = svcmpne_n_f32(ptrue, lambda1_v, 0.0f);
-  if (C10_LIKELY(svptest_any(predNZ, predNZ))) {
-    svfloat32_t lambda1_r = svdivr_f32_m(predNZ, lambda1_v, svdup_n_f32(1.0f));
-    svfloat32_t pow = svexp(lnData * lambda1_v);
-    lnData = svsel_f32(predNZ, lambda1_r, lnData);
-    lnData = svnmsb_f32_m(predNZ, lnData, pow, lnData);
-  }
-  return lnData;
-}
-
-template <typename T>
-void compute_batch_box_cox_vec_sve128(
-    std::size_t N,
-    std::size_t D,
-    const T* data_ptr,
-    const T* __restrict lambda1_ptr,
-    const T* __restrict lambda2_ptr,
-    T* output_ptr);
-
-template <>
-void compute_batch_box_cox_vec_sve128(
-    std::size_t N,
-    std::size_t D,
-    const float *data_ptr,
-    const float *__restrict lambda1_ptr,
-    const float *__restrict lambda2_ptr,
-    float *output_ptr) {
-  const svfloat32_t k_eps = svdup_n_f32(static_cast<float>(1e-6));
-
-  std::size_t remainder = D % 4;
-  std::size_t loopBound = D - remainder;
-  svbool_t remainderPred = svwhilelt_b32_u64(0, remainder);
-
-  for (; C10_LIKELY(N > 0); --N) {
-    for (std::size_t j = 0; C10_LIKELY(j != loopBound);
-         j += 4, data_ptr += 4, output_ptr += 4) {
-      svfloat32_t lambda1_v =
-          svset_neonq(svundef_f32(), vld1q_f32(lambda1_ptr + j));
-      svfloat32_t lambda2_v =
-          svset_neonq(svundef_f32(), vld1q_f32(lambda2_ptr + j));
-      svfloat32_t data_v = svset_neonq(svundef_f32(), vld1q_f32(data_ptr));
-      svfloat32_t result = compute_batch_box_cox_vec_sve128_float(
-          lambda1_v, lambda2_v, data_v, k_eps);
-      vst1q_f32(output_ptr, svget_neonq(result));
-    }
-    if (C10_LIKELY(remainder > 0)) {
-      svfloat32_t lambda1_v = svld1_f32(remainderPred, lambda1_ptr + loopBound);
-      svfloat32_t lambda2_v = svld1_f32(remainderPred, lambda2_ptr + loopBound);
-      svfloat32_t data_v = svld1_f32(remainderPred, data_ptr);
-      svfloat32_t result = compute_batch_box_cox_vec_sve128_float(
-          lambda1_v, lambda2_v, data_v, k_eps);
-      svst1_f32(remainderPred, output_ptr, result);
-      data_ptr += remainder;
-      output_ptr += remainder;
-    }
-  }
-}
-
-namespace caffe2::details {
-
-template <typename T>
-void compute_batch_box_cox__sve128(
-    std::size_t N,
-    std::size_t D,
-    const T* self_data,
-    const T* __restrict lambda1_data,
-    const T* __restrict lambda2_data,
-    T* output_data) {
-  compute_batch_box_cox_vec_sve128<T>(
-      N, D, self_data, lambda1_data, lambda2_data, output_data);
-}
-
-// Vectorized version specializations for float and double
-template void compute_batch_box_cox__sve128<float>(
-    std::size_t N,
-    std::size_t D,
-    const float* self_data,
-    const float* __restrict lambda1_data,
-    const float* __restrict lambda2_data,
-    float* output_data);
-
-} // namespace caffe2::details
-
-#endif // __aarch64__ && __ARM_FEATURE_SVE && CAFFE2_PERF_WITH_SVE128
diff --git a/caffe2/perfkernels/batch_box_cox_vec.h b/caffe2/perfkernels/batch_box_cox_vec.h
deleted file mode 100644
index 08e4f84fe4327..0000000000000
--- a/caffe2/perfkernels/batch_box_cox_vec.h
+++ /dev/null
@@ -1,321 +0,0 @@
-#pragma once
-
-#include <cstdint>
-#include <cmath>
-#include <vector>
-#include <c10/util/irange.h>
-#include <folly/SingletonThreadLocal.h>
-#include <caffe2/perfkernels/common.h>
-#include "vectorizer.h"
-#include <mkl.h>
-
-namespace caffe2::details {
-
-namespace {
-void TileIndicesInPlace(std::vector<int>& v, const std::size_t D, const std::size_t K) {
-  auto n = v.size();
-  v.resize(K * n);
-  for (const auto k : c10::irange(1, K)) {
-    for (const auto j : c10::irange(n)) {
-      v[k * n + j] = v[j] + k * D;
-    }
-  }
-}
-
-// MKL VML function templates.
-template <typename T>
-void PackV(const int N, const T* a, const int* ia, T* y);
-template <typename T>
-void UnpackV(const int N, const T* a, T* y, const int* iy);
-
-#define DELEGATE_PACKV_FUNCTION(T, OriginalFunc)                \
-  template <>                                                   \
-  void PackV<T>(const int N, const T* a, const int* ia, T* y) { \
-    OriginalFunc(N, a, ia, y);                                  \
-  }
-DELEGATE_PACKV_FUNCTION(float, vsPackV)
-DELEGATE_PACKV_FUNCTION(double, vdPackV)
-#undef DELEGATE_PACKV_FUNCTION
-
-#define DELEGATE_UNPACKV_FUNCTION(T, OriginalFunc)                \
-  template <>                                                     \
-  void UnpackV<T>(const int N, const T* a, T* y, const int* iy) { \
-    OriginalFunc(N, a, y, iy);                                    \
-  }
-DELEGATE_UNPACKV_FUNCTION(float, vsUnpackV)
-DELEGATE_UNPACKV_FUNCTION(double, vdUnpackV)
-#undef DELEGATE_UNPACKV_FUNCTION
-
-#ifndef FAST_VECTORIZED_KERNEL
-template <typename T>
-void box_cox_zero_lambda(
-    size_t D,
-    const T* const self_data,
-    const T* const lambda2_data,
-    T k_eps,
-    T* const output_data) {
-  int j = 0;
-  using Vec = at::vec::Vectorized<T>;
-  constexpr int64_t VLEN = Vec::size();
-  auto k_eps_vec = Vec(k_eps);
-  for(; j + VLEN < D; j += VLEN) {
-    auto data = Vec::loadu(self_data + j);
-    auto lambda2 = Vec::loadu(lambda2_data + j);
-    auto sum = data + lambda2;
-    auto max = at::vec::max(sum, k_eps_vec);
-    auto res = max.log();
-    res.store(output_data + j);
-  }
-  for ( ;j < D; ++j) {
-    auto sum = self_data[j] + lambda2_data[j];
-    auto max = std::max(sum, k_eps);
-    output_data[j] = std::log(max);
-  }
-}
-
-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,
-    const T* data_ptr,
-    const T* lambda1_ptr,
-    const T* lambda2_ptr,
-    T k_eps,
-    T* out) {
-
-  int j = 0;
-  using Vec = at::vec::Vectorized<T>;
-  constexpr int64_t VLEN = Vec::size();
-  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);
-    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);
-  }
-}
-#else
-template <typename T>
-void box_cox_zero_lambda(
-    size_t D,
-    const T* const self_data,
-    const T* const lambda2_data,
-    T k_eps,
-    T* const output_data) {
-  VECTOR_LOOP for (auto j=0 ;j < D; ++j) {
-    auto sum = self_data[j] + lambda2_data[j];
-    auto max = std::max(sum, k_eps);
-    output_data[j] = std::log(max);
-  }
-}
-
-template <typename T>
-void box_cox_nonzero_lambda(
-    int64_t D,
-    const T* data_ptr,
-    const T* lambda1_ptr,
-    const T* lambda2_ptr,
-    T k_eps,
-    T* out) {
-
-  VECTOR_LOOP for (auto j=0 ;j < D; ++j) {
-    FAST_MATH
-    auto sum = data_ptr[j] + lambda2_ptr[j];
-    auto max = std::max(sum, k_eps);
-    auto lamda1 = lambda1_ptr[j];
-    auto lambda_over_1 = 1 / lamda1;
-    if constexpr (std::is_same<T, float>::value) {
-      lambda_over_1 = lambda_over_1 * (T{2} - lambda_over_1 * lamda1);
-      lambda_over_1 = lambda_over_1 * (T{2} - lambda_over_1 * lamda1);
-    }
-    auto pow = std::pow(max, lamda1);
-    out[j] = pow * lambda_over_1 - lambda_over_1;
-  }
-}
-#endif // FAST_VECTORIZED_KERNEL
-
-template <typename T>
-void box_cox_mixed_lambda(
-    const T* const self_data,
-    const std::vector<int>& nonzeros,
-    const std::vector<int>& zeros,
-    const T* const lambda1,
-    const T* const lambda2,
-    const T* const lambda2_z_,
-    T k_eps,
-    T* const buffer,
-    T* const output_data) {
-  PackV(nonzeros.size(), self_data, nonzeros.data(), buffer);
-  box_cox_nonzero_lambda<T>(
-      nonzeros.size(), buffer, lambda1, lambda2, k_eps, buffer);
-  UnpackV(nonzeros.size(), buffer, output_data, nonzeros.data());
-
-  PackV(zeros.size(), self_data, zeros.data(), buffer);
-  box_cox_zero_lambda<T>(
-      zeros.size(), buffer, lambda2_z_, k_eps, buffer);
-  UnpackV(zeros.size(), buffer, output_data, zeros.data());
-}
-
-template <typename T>
-void TileArrayIntoVector(
-    const T* const a,
-    const size_t D,
-    const int K,
-    std::vector<T>& b) {
-  b.resize(K * D);
-  for (const auto k : c10::irange(K)) {
-    std::copy(a, a + D, b.begin() + k * D);
-  }
-}
-
-template <typename T>
-void compute_batch_box_cox_vec_fma(
-    std::size_t N,
-    std::size_t D,
-    std::size_t block_size,
-    const T* self_data,
-    const T* __restrict lambda1_data,
-    const T* __restrict lambda2_data,
-    T* output_data) {
-  constexpr T k_eps = static_cast<T>(1e-6);
-
-  FOLLY_DECLARE_REUSED(zeros, std::vector<int>);
-  FOLLY_DECLARE_REUSED(nonzeros, std::vector<int>);
-  // Don't bother calling reserve; calls after the first will get a
-  // correctly-sized allocation anyway.
-  for (const auto j : c10::irange(D)) {
-    if (lambda1_data[j] == 0) {
-      zeros.push_back(j);
-    } else {
-      nonzeros.push_back(j);
-    }
-  }
-
-  // Process K rows at a time for effective vectorization with small rows.
-  const auto K = std::min(N, (block_size + D - 1) / D);
-
-  FOLLY_DECLARE_REUSED(lambda1_, std::vector<T>);
-  FOLLY_DECLARE_REUSED(lambda2_, std::vector<T>);
-  FOLLY_DECLARE_REUSED(lambda2_z_, std::vector<T>);
-
-  if (nonzeros.size() == D) {
-    // ((x + lambda2)^lambda1 - 1)/lambda1, if lambda1 != 0
-    size_t i = 0;
-    if (K > 1) {
-      TileArrayIntoVector(lambda1_data, D, K, lambda1_);
-      TileArrayIntoVector(lambda2_data, D, K, lambda2_);
-      DCHECK_EQ(K * D, lambda1_.size());
-      DCHECK_EQ(K * D, lambda2_.size());
-      for (; i < N - K + 1; i += K, self_data += K * D, output_data += K * D) {
-        box_cox_nonzero_lambda<T>(
-            K * D,
-            self_data,
-            lambda1_.data(),
-            lambda2_.data(),
-            k_eps,
-            output_data);
-      }
-    }
-    for (; i < N; i++, self_data += D, output_data += D) {
-      box_cox_nonzero_lambda<T>(
-          D, self_data, lambda1_data, lambda2_data, k_eps, output_data);
-    }
-  } else if (zeros.size() == D) {
-    // ln(x + lambda2), if lambda1 == 0
-    size_t i = 0;
-    if (K > 1) {
-      TileArrayIntoVector(lambda2_data, D, K, lambda2_z_);
-      DCHECK_EQ(K * D, lambda2_z_.size());
-      for (; i < N - K + 1; i += K, self_data += K * D, output_data += K * D) {
-        box_cox_zero_lambda<T>(
-            K * D, self_data, lambda2_z_.data(), k_eps, output_data);
-      }
-    }
-    for (; i < N; i++, self_data += D, output_data += D) {
-      box_cox_zero_lambda<T>(
-          D, self_data, lambda2_data, k_eps, output_data);
-    }
-  } else {
-    // mix zeros and nonzeros
-    const size_t n = nonzeros.size();
-    if (K > 1) {
-      TileIndicesInPlace(nonzeros, 0, K);
-      TileIndicesInPlace(zeros, 0, K);
-    }
-
-    FOLLY_DECLARE_REUSED(buffer, std::vector<T>);
-
-    buffer.resize(std::max(nonzeros.size(), zeros.size()));
-    lambda1_.resize(nonzeros.size());
-    lambda2_.resize(nonzeros.size());
-    lambda2_z_.resize(zeros.size());
-    PackV(nonzeros.size(), lambda1_data, nonzeros.data(), lambda1_.data());
-    PackV(nonzeros.size(), lambda2_data, nonzeros.data(), lambda2_.data());
-    PackV(zeros.size(), lambda2_data, zeros.data(), lambda2_z_.data());
-
-    size_t i = 0;
-    if (K > 1) {
-      // Truncate to original size, and re-tile with offsets this time.
-      nonzeros.resize(n);
-      DCHECK_GT(D, n);
-      zeros.resize(D - n);
-      TileIndicesInPlace(nonzeros, D, K);
-      TileIndicesInPlace(zeros, D, K);
-      DCHECK_EQ(nonzeros.size(), lambda1_.size());
-      DCHECK_EQ(nonzeros.size(), lambda2_.size());
-      DCHECK_EQ(zeros.size(), lambda2_z_.size());
-
-      for (; i < N - K + 1; i += K, self_data += K * D, output_data += K * D) {
-        box_cox_mixed_lambda<T>(
-            self_data,
-            nonzeros,
-            zeros,
-            lambda1_.data(),
-            lambda2_.data(),
-            lambda2_z_.data(),
-            k_eps,
-            buffer.data(),
-            output_data);
-      }
-      // Truncate to original size.
-      nonzeros.resize(n);
-      zeros.resize(D - n);
-    }
-    for (; i < N; i++, self_data += D, output_data += D) {
-      box_cox_mixed_lambda<T>(
-          self_data,
-          nonzeros,
-          zeros,
-          lambda1_.data(),
-          lambda2_.data(),
-          lambda2_z_.data(),
-          k_eps,
-          buffer.data(),
-          output_data);
-    }
-  }
-}
-}  // namespace
-
-}   // namespace caffe2::details
diff --git a/cmake/Codegen.cmake b/cmake/Codegen.cmake
index 5faad21f9f6cd..60d00dd394e37 100644
--- a/cmake/Codegen.cmake
+++ b/cmake/Codegen.cmake
@@ -163,6 +163,11 @@ if(INTERN_BUILD_ATEN_OPS)
     set(GEN_XPU_FLAG --xpu)
   endif()
 
+  set(GEN_MTIA_FLAG)
+  if(USE_MTIA)
+    set(GEN_MTIA_FLAG --mtia)
+  endif()
+
   set(CUSTOM_BUILD_FLAGS)
   if(INTERN_BUILD_MOBILE)
     if(USE_VULKAN)
@@ -245,10 +250,12 @@ if(INTERN_BUILD_ATEN_OPS)
       "${Python_EXECUTABLE}" -m torchgen.gen
       --source-path ${CMAKE_CURRENT_LIST_DIR}/../aten/src/ATen
       --install_dir ${CMAKE_BINARY_DIR}/aten/src/ATen
+      --headeronly-install-dir ${CMAKE_BINARY_DIR}/torch/headeronly/core
       ${GEN_PER_OPERATOR_FLAG}
       ${GEN_ROCM_FLAG}
       ${GEN_MPS_FLAG}
       ${GEN_XPU_FLAG}
+      ${GEN_MTIA_FLAG}
       ${CUSTOM_BUILD_FLAGS}
   )
 
@@ -302,6 +309,10 @@ if(INTERN_BUILD_ATEN_OPS)
       ${CMAKE_BINARY_DIR}/aten/src/ATen/core_generated_${gen_type}.cmake
       ${CMAKE_BINARY_DIR}/aten/src/ATen/cpu_vec_generated_${gen_type}.cmake
       ${CMAKE_BINARY_DIR}/aten/src/ATen/cuda_generated_${gen_type}.cmake)
+    if(gen_type STREQUAL "headers")
+      list(APPEND OUTPUT_LIST
+        "${CMAKE_BINARY_DIR}/torch/headeronly/core/enum_tag.h")
+    endif()
     if(USE_XPU)
       list(APPEND OUTPUT_LIST
         ${xpu_generated_${gen_type}}
@@ -411,14 +422,10 @@ if(INTERN_BUILD_ATEN_OPS)
     LIST(APPEND CPU_CAPABILITY_FLAGS "${OPT_FLAG}  ${CXX_ZVECTOR_FLAGS}")
   endif(CXX_ZVECTOR_FOUND)
 
-  if(CXX_SVE_FOUND AND CXX_SVE256_FOUND AND CXX_ARM_BF16_FOUND)
+  if(CXX_SVE256_FOUND)
     list(APPEND CPU_CAPABILITY_NAMES "SVE256")
-    set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -DHAVE_SVE_CPU_DEFINITION -DHAVE_SVE256_CPU_DEFINITION -DHAVE_ARM_BF16_CPU_DEFINITION")
-    if("${CMAKE_C_COMPILER_ID}" MATCHES "Clang")
-      list(APPEND CPU_CAPABILITY_FLAGS "${OPT_FLAG} -O2 -march=armv8-a+sve+bf16 -D__ARM_FEATURE_BF16 -DCPU_CAPABILITY_SVE -msve-vector-bits=256")
-    else()
-      list(APPEND CPU_CAPABILITY_FLAGS "${OPT_FLAG} -march=armv8-a+sve+bf16 -D__ARM_FEATURE_BF16 -DCPU_CAPABILITY_SVE -msve-vector-bits=256")
-    endif()
+    set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -DHAVE_SVE_CPU_DEFINITION")
+    list(APPEND CPU_CAPABILITY_FLAGS "${OPT_FLAG} -march=armv8-a+sve+bf16 -D__ARM_FEATURE_BF16 -msve-vector-bits=256")
   endif()
 
   list(LENGTH CPU_CAPABILITY_NAMES NUM_CPU_CAPABILITY_NAMES)
diff --git a/cmake/Dependencies.cmake b/cmake/Dependencies.cmake
index dac0d1f41c3bd..e31ff49aeba0b 100644
--- a/cmake/Dependencies.cmake
+++ b/cmake/Dependencies.cmake
@@ -47,15 +47,21 @@ if(USE_CUDA)
     # torch::cudart is dealt with separately, due to CUDA_ADD_LIBRARY
     # design reason (it adds CUDA_LIBRARIES itself).
     set(Caffe2_PUBLIC_CUDA_DEPENDENCY_LIBS )
-    if(NOT CAFFE2_USE_NVRTC)
-      caffe2_update_option(USE_NVRTC OFF)
-    endif()
     list(APPEND Caffe2_CUDA_DEPENDENCY_LIBS caffe2::curand caffe2::cufft caffe2::cublas)
     if(CAFFE2_USE_CUDNN)
+      if(NOT CAFFE2_USE_NVRTC)
+        message(FATAL_ERROR
+          "USE_CUDNN requires USE_NVRTC (required by cudnn_frontend 1.21+). "
+          "Please set -DUSE_NVRTC=ON or disable cuDNN with -DUSE_CUDNN=OFF.")
+      endif()
       list(APPEND Caffe2_CUDA_DEPENDENCY_LIBS torch::cudnn)
+      list(APPEND Caffe2_CUDA_DEPENDENCY_LIBS caffe2::nvrtc_runtime)
     else()
       caffe2_update_option(USE_CUDNN OFF)
     endif()
+    if(NOT CAFFE2_USE_NVRTC)
+      caffe2_update_option(USE_NVRTC OFF)
+    endif()
     if(CAFFE2_USE_CUSPARSELT)
       list(APPEND Caffe2_CUDA_DEPENDENCY_LIBS torch::cusparselt)
     else()
@@ -1023,7 +1029,7 @@ if(USE_ROCM)
       list(APPEND HIP_CXX_FLAGS -DUSE_ROCM_CK_GEMM)
     endif()
     list(APPEND HIP_HIPCC_FLAGS --offload-compress)
-    list(APPEND HIP_HIPCC_FLAGS -std=c++17)
+    list(APPEND HIP_HIPCC_FLAGS -std=c++20)
     # Pass device library path for theRock nightly builds
     if(DEFINED ENV{HIP_DEVICE_LIB_PATH})
       file(TO_CMAKE_PATH "$ENV{HIP_DEVICE_LIB_PATH}" _hip_device_lib_path)
@@ -1063,6 +1069,8 @@ if(USE_ROCM)
     list(APPEND HIP_HIPCC_FLAGS -fclang-abi-compat=17)
 
     set(HIP_CLANG_FLAGS ${HIP_CXX_FLAGS})
+    string(JOIN " " HIP_HIPCC_FLAGS_STR ${HIP_HIPCC_FLAGS})
+    set(HIP_HIPCC_FLAGS ${HIP_HIPCC_FLAGS_STR})
     set(CMAKE_HIP_FLAGS ${HIP_HIPCC_FLAGS})
     # Ask hcc to generate device code during compilation so we can use
     # host linker to link.
@@ -1087,6 +1095,9 @@ if(USE_ROCM)
       list(APPEND Caffe2_PUBLIC_HIP_DEPENDENCY_LIBS
         roc::hipsparselt
       )
+      if(ROCM_VERSION_DEV VERSION_GREATER_EQUAL "7.12.0")
+          set(CAFFE2_USE_HIPSPARSELT ON)
+      endif()
     endif()
 
     # ---[ Kernel asserts
@@ -1461,26 +1472,6 @@ if(NOT INTERN_BUILD_MOBILE)
 
   # ARM specific flags
   find_package(ARM)
-  if(ASIMD_FOUND)
-    message(STATUS "asimd/Neon found with compiler flag : -D__NEON__")
-    add_compile_options(-D__NEON__)
-  elseif(NEON_FOUND)
-    if(APPLE)
-      message(STATUS "Neon found with compiler flag : -D__NEON__")
-      add_compile_options(-D__NEON__)
-    else()
-      message(STATUS "Neon found with compiler flag : -mfpu=neon -D__NEON__")
-      add_compile_options(-mfpu=neon -D__NEON__)
-    endif()
-  endif()
-  if(CORTEXA8_FOUND)
-    message(STATUS "Cortex-A8 Found with compiler flag : -mcpu=cortex-a8")
-    add_compile_options(-mcpu=cortex-a8 -fprefetch-loop-arrays)
-  endif()
-  if(CORTEXA9_FOUND)
-    message(STATUS "Cortex-A9 Found with compiler flag : -mcpu=cortex-a9")
-    add_compile_options(-mcpu=cortex-a9)
-  endif()
 
   find_package(LAPACK)
   if(LAPACK_FOUND)
@@ -1659,7 +1650,7 @@ if(USE_KINETO)
 
   if(NOT LIBKINETO_NOROCTRACER)
     if("$ENV{ROCM_SOURCE_DIR}" STREQUAL "")
-      set(ENV{ROCM_SOURCE_DIR} "/opt/rocm")
+      set(ENV{ROCM_SOURCE_DIR} "${ROCM_PATH}")
     endif()
   endif()
 
diff --git a/cmake/EnvVarForwarding.cmake b/cmake/EnvVarForwarding.cmake
new file mode 100644
index 0000000000000..a31ce6db26637
--- /dev/null
+++ b/cmake/EnvVarForwarding.cmake
@@ -0,0 +1,138 @@
+# Forward environment variables to CMake variables.
+#
+# This replicates the behavior of setup.py / tools/setup_helpers/cmake.py which
+# passes all BUILD_*, USE_*, and CMAKE_* environment variables as -D flags, plus
+# a set of additional variables that don't follow the prefix convention.
+
+# Additional env vars that are forwarded with a different CMake variable name.
+set(_ENV_ALIASES
+  "CUDNN_LIB_DIR=CUDNN_LIBRARY"
+  "USE_CUDA_STATIC_LINK=CAFFE2_STATIC_LINK_CUDA"
+)
+
+# Additional env vars forwarded with the same name.
+set(_ENV_PASSTHROUGH
+  UBSAN_FLAGS
+  BLAS
+  WITH_BLAS
+  CUDA_HOST_COMPILER
+  CUDA_NVCC_EXECUTABLE
+  CUDA_SEPARABLE_COMPILATION
+  CUDNN_LIBRARY
+  CUDNN_INCLUDE_DIR
+  CUDNN_ROOT
+  EXPERIMENTAL_SINGLE_THREAD_POOL
+  INSTALL_TEST
+  INTEL_MKL_DIR
+  INTEL_OMP_DIR
+  MKL_THREADING
+  MKLDNN_CPU_RUNTIME
+  MSVC_Z7_OVERRIDE
+  CAFFE2_USE_MSVC_STATIC_RUNTIME
+  Numa_INCLUDE_DIR
+  Numa_LIBRARIES
+  ONNX_ML
+  ONNX_NAMESPACE
+  ATEN_THREADING
+  WERROR
+  OPENSSL_ROOT_DIR
+  STATIC_DISPATCH_BACKEND
+  SELECTED_OP_LIST
+  TORCH_CUDA_ARCH_LIST
+  TORCH_XPU_ARCH_LIST
+  TRACING_BASED
+  PYTHON_LIB_REL_PATH
+)
+
+# Low-priority aliases: if the canonical var is not set, use the alias.
+set(_LOW_PRIORITY_ALIASES
+  "CUDA_HOST_COMPILER=CMAKE_CUDA_HOST_COMPILER"
+  "CUDAHOSTCXX=CUDA_HOST_COMPILER"
+  "CMAKE_CUDA_HOST_COMPILER=CUDA_HOST_COMPILER"
+  "CMAKE_CUDA_COMPILER=CUDA_NVCC_EXECUTABLE"
+  "CUDACXX=CUDA_NVCC_EXECUTABLE"
+)
+
+# Forward aliased env vars (env name -> different cmake name)
+foreach(_alias IN LISTS _ENV_ALIASES)
+  string(REPLACE "=" ";" _parts "${_alias}")
+  list(GET _parts 0 _env_name)
+  list(GET _parts 1 _cmake_name)
+  if(DEFINED ENV{${_env_name}} AND NOT DEFINED ${_cmake_name})
+    set(${_cmake_name} "$ENV{${_env_name}}" CACHE STRING "From env ${_env_name}" FORCE)
+  endif()
+endforeach()
+
+# Forward passthrough env vars (same name)
+foreach(_var IN LISTS _ENV_PASSTHROUGH)
+  if(DEFINED ENV{${_var}} AND NOT DEFINED ${_var})
+    set(${_var} "$ENV{${_var}}" CACHE STRING "From env ${_var}" FORCE)
+  endif()
+endforeach()
+
+# Forward all BUILD_*, USE_*, CMAKE_* env vars not already set as CMake
+# variables, plus vars ending in EXITCODE or EXITCODE__TRYRUN_OUTPUT.
+# This matches the existing behavior where setup.py passed everything with
+# these prefixes/suffixes through to CMake.
+# We use execute_process + env to get the full list since CMake has no
+# built-in way to enumerate environment variables.
+execute_process(
+  COMMAND "${CMAKE_COMMAND}" -E environment
+  OUTPUT_VARIABLE _all_env
+  OUTPUT_STRIP_TRAILING_WHITESPACE
+)
+string(REPLACE "\n" ";" _env_lines "${_all_env}")
+foreach(_line IN LISTS _env_lines)
+  if(_line MATCHES "^([A-Za-z_0-9]+)=(.*)")
+    set(_var_name "${CMAKE_MATCH_1}")
+    set(_var_value "${CMAKE_MATCH_2}")
+    # Only forward vars with BUILD_/USE_/CMAKE_ prefix or *EXITCODE* suffix.
+    string(REGEX MATCH "^(BUILD_|USE_|CMAKE_)" _has_prefix "${_var_name}")
+    string(REGEX MATCH "(EXITCODE|EXITCODE__TRYRUN_OUTPUT)$" _has_suffix "${_var_name}")
+    if(NOT _has_prefix AND NOT _has_suffix)
+      continue()
+    endif()
+    if(NOT DEFINED ${_var_name})
+      set(${_var_name} "${_var_value}" CACHE STRING "From environment" FORCE)
+    endif()
+  endif()
+endforeach()
+
+# Low-priority aliases
+foreach(_alias IN LISTS _LOW_PRIORITY_ALIASES)
+  string(REPLACE "=" ";" _parts "${_alias}")
+  list(GET _parts 0 _env_name)
+  list(GET _parts 1 _cmake_name)
+  if(DEFINED ENV{${_env_name}} AND NOT DEFINED ${_cmake_name})
+    set(${_cmake_name} "$ENV{${_env_name}}" CACHE STRING "From env alias ${_env_name}" FORCE)
+  endif()
+endforeach()
+
+# Ensure Python's purelib is on CMAKE_PREFIX_PATH so CMake can find
+# packages installed there (e.g., pybind11, numpy).
+if(Python_EXECUTABLE)
+  execute_process(
+    COMMAND "${Python_EXECUTABLE}" -c "import sysconfig; print(sysconfig.get_path('purelib'))"
+    OUTPUT_VARIABLE _py_purelib
+    OUTPUT_STRIP_TRAILING_WHITESPACE
+    ERROR_QUIET
+  )
+  if(_py_purelib AND NOT "${_py_purelib}" STREQUAL "")
+    list(PREPEND CMAKE_PREFIX_PATH "${_py_purelib}")
+    # Preserve paths from the CMAKE_PREFIX_PATH environment variable.
+    # Setting the cmake variable shadows the env var, so we must merge it in
+    # explicitly. This ensures conda's prefix (e.g. /opt/conda/envs/py_3.10)
+    # is present so cmake can find conda-provided libraries (libgomp, libnuma).
+    if(DEFINED ENV{CMAKE_PREFIX_PATH} AND NOT "$ENV{CMAKE_PREFIX_PATH}" STREQUAL "")
+      if(WIN32)
+        # On Windows the env var is already ;-separated and : appears in drive
+        # letters (e.g. C:\conda\envs\py310), so use it as-is.
+        set(_env_prefix "$ENV{CMAKE_PREFIX_PATH}")
+      else()
+        string(REPLACE ":" ";" _env_prefix "$ENV{CMAKE_PREFIX_PATH}")
+      endif()
+      list(APPEND CMAKE_PREFIX_PATH ${_env_prefix})
+    endif()
+    list(REMOVE_DUPLICATES CMAKE_PREFIX_PATH)
+  endif()
+endif()
diff --git a/cmake/External/aotriton.cmake b/cmake/External/aotriton.cmake
index 93f852f014887..a59d5826a0d9d 100644
--- a/cmake/External/aotriton.cmake
+++ b/cmake/External/aotriton.cmake
@@ -9,47 +9,58 @@ if(NOT __AOTRITON_INCLUDED)
   # Replaces .ci/docker/aotriton_version.txt
   # Note packages information may have versions skipped (due to no ABI breaks)
   # But they must be listed from lower version to higher version
-  set(__AOTRITON_VER "0.11.2b")
+  set(__AOTRITON_VER "0.13.50tp")
   set(__AOTRITON_MANYLINUX_LIST
-      "manylinux_2_28"  # rocm6.2
-      "manylinux_2_28"  # rocm6.3
       "manylinux_2_28"  # rocm6.4
       "manylinux_2_28"  # rocm7.0
       "manylinux_2_28"  # rocm7.1
       "manylinux_2_28"  # rocm7.2
+      "manylinux_2_28"  # rocm7.14
+      "manylinux_2_28"  # rocm7.15
       )
   set(__AOTRITON_ROCM_LIST
-      "rocm6.2"
-      "rocm6.3"
       "rocm6.4"
       "rocm7.0"
       "rocm7.1"
       "rocm7.2"
+      "rocm7.14"
+      "rocm7.15"
       )
-  set(__AOTRITON_CI_COMMIT "dd1b68b604b5258ee7a9f7b66ad95e7a82c18065")
+  if(DEFINED ENV{PYTORCH_AOTRITON_COMMIT})
+    set(__AOTRITON_CI_COMMIT "$ENV{PYTORCH_AOTRITON_COMMIT}")
+  else()
+    set(__AOTRITON_CI_COMMIT "05f26c2d8fbdcf8dd6d95b9544d04f7a39fc9920")
+  endif()
   set(__AOTRITON_SHA256_LIST
-      "d784314849ba1911181dfc80cd845064ff6f0cdad10e2f4c53eb84a8b89245b9"  # rocm6.2
-      "f4b14dc111c334e967b28a1cf9ed4c63264c634dbdccbb5849aa9490022992f7"  # rocm6.3
-      "6b51d8479c85b902334e4f5518f404a8f5d563fd8d4732cb8b621ed4b45c2876"  # rocm6.4
-      "5501a0a3b300890001b6625f2a3539a7bad60f386f0a061ebe7d4ed5ca0fafb9"  # rocm7.0
-      "fee36beb3ea484ce18155bbafe026c577fd6705e4469e59405b260bd74b8cc10"  # rocm7.1
-      "cd8abf27bbb63cec45c94135e9b28745966074263a6b0555e5878ae1cb6a2349"  # rocm7.2
+      "7708bd2ff58e68a6b77f1605507733012d035deb36867d7d400220c7897ff438"  # rocm6.4
+      "e8923752f687687d2ad0343ca7b444f2cedc553b72711daf81b90a58ba9d3365"  # rocm7.0
+      "1fec784443d6958ad01da76a70ed99ffe107c98951490a754a8c4c0444fcd3e9"  # rocm7.1
+      "8ce1a1a0eea7b6e01ddd9d702f56a145c034f0da2b3187418d48d7ae40223fea"  # rocm7.2
+      "b978982154cbb825c4c65b54c0bc64b106c497ce8bb6a38e3f70b6a69b86b9ae"  # rocm7.14
+      "a698fda2edf84759c8b9f19e27c9bc302679e205bac3f684cb8f3e4f2878bb07"  # rocm7.15
       )
   set(__AOTRITON_IMAGE_LIST
       "amd-gfx90a"
       "amd-gfx942"
       "amd-gfx950"
-      "amd-gfx11xx"
+      "amd-gfx110x"
+      "amd-gfx115x"
       "amd-gfx120x"
+      "amd-gfx1250"
      )
   set(__AOTRITON_IMAGE_SHA256_LIST
-     "fe9f04b66bf52ac27cd025e1d89cfd04974dd3fb3ae076192f783641a4d80fdf" # amd-gfx90a
-     "0a7bcee19d3bb6d548732248c3234f7b92736c2ab7a7aae65294b87a7fd64c06" # amd-gfx942
-     "c1ba3bfe84217fd67df3dd1f8b67c80a7f7b33d0ad4d74b41d6567036e032ace" # amd-gfx950
-     "839299637fccb13fbe3e7823d57d1b2dcd0e0bed78abbcb7005ea5f4fd82b928" # amd-gfx11xx
-     "0a4ff324bffdac0c2fde87a8a7f70563d3c84a80ad4e8f31345f2b40a1384e95" # amd-gfx120x
+     "a3d1a6868ce290ba8118618207093e785252eff4e18a64f495752cb5a03ffed6" # amd-gfx90a
+     "ccdbc7e3d96839be4895ee004f21531cc55d590c9018937b9e314bba363b3927" # amd-gfx942
+     "518fd072eb05948fc0a6c25a20832591c6406df865e3b691a2aeff3fd4c5ce1d" # amd-gfx950
+     "efe773e7a2c8adc995d90ecd0daca2db801285445ee10432df2b29c67d5b11d2" # amd-gfx110x
+     "1bc50e8aa8b6bda3410e92886ccca8fd45df3e60a6cbda9ffc58b2c541efd5c2" # amd-gfx115x
+     "6a465dbc03148bba8a2d78c4c2a3cb83155eca00f4f7f749e676402d7660968c" # amd-gfx120x
+     "74fe9073ec11136a3d779af1eccd63d721f85ee08b0e2665175c3e0fe4d12c55" # amd-gfx1250
      )
-  set(__AOTRITON_BASE_URL "https://github.com/ROCm/aotriton/releases/download/")  # @lint-ignore
+  set(__AOTRITON_BASE_URL "$ENV{PYTORCH_AOTRITON_BASE_URL}")
+  if(NOT __AOTRITON_BASE_URL)
+    set(__AOTRITON_BASE_URL "https://github.com/ROCm/aotriton/releases/download/")  # @lint-ignore
+  endif()
   set(__AOTRITON_Z "gz")
   # Set the default __AOTRITON_LIB path
   if(NOT WIN32)
diff --git a/cmake/Metal.cmake b/cmake/Metal.cmake
index c9565e2fc0e9e..f378dfb0e803d 100644
--- a/cmake/Metal.cmake
+++ b/cmake/Metal.cmake
@@ -45,6 +45,13 @@ set(BFLOAT_METAL_CODE "
     ptr[idx] += 1;
   }
 ")
+set(LAMBDA_METAL_CODE "
+  kernel void test(device float* ptr,
+                   uint idx [[thread_position_in_grid]]) {
+    auto fn = [](float x) { return x + 1.0; };
+    ptr[idx] = fn(ptr[idx]);
+  }
+")
 if(NOT CAN_COMPILE_METAL_FOUND)
     file(WRITE "${CMAKE_CURRENT_BINARY_DIR}/bfloat_inc.metal" "${BFLOAT_METAL_CODE}")
     execute_process(COMMAND xcrun metal -std=metal3.1 bfloat_inc.metal
@@ -59,6 +66,23 @@ if(NOT CAN_COMPILE_METAL_FOUND)
         message(WARNING "Machine can not compile metal shaders, fails with ${XCRUN_OUTPUT}")
         set(CAN_COMPILE_METAL NO CACHE BOOL "Host can compile metal shaders")
     endif()
+    if(CAN_COMPILE_METAL)
+        file(WRITE "${CMAKE_CURRENT_BINARY_DIR}/lambda_test.metal" "${LAMBDA_METAL_CODE}")
+        execute_process(COMMAND xcrun metal -std=metal4.0 -c lambda_test.metal -o /dev/null
+                        WORKING_DIRECTORY "${CMAKE_CURRENT_BINARY_DIR}"
+                        OUTPUT_VARIABLE XCRUN_OUTPUT
+                        ERROR_VARIABLE XCRUN_OUTPUT
+                        RESULT_VARIABLE XCRUN_RC)
+        if(${XCRUN_RC} EQUAL 0)
+            message(STATUS "Metal toolchain supports Metal 4.0")
+            set(CAN_COMPILE_METAL_40 YES CACHE BOOL "Host can compile Metal 4.0 shaders" FORCE)
+        else()
+            message(STATUS "Metal toolchain does not support Metal 4.0")
+            set(CAN_COMPILE_METAL_40 NO CACHE BOOL "Host can compile Metal 4.0 shaders" FORCE)
+        endif()
+    else()
+        set(CAN_COMPILE_METAL_40 NO CACHE BOOL "Host can compile Metal 4.0 shaders" FORCE)
+    endif()
     set(CAN_COMPILE_METAL_FOUND YES CACHE INTERNAL "Run check for shader compiler")
 endif()
 
diff --git a/cmake/Modules/FindARM.cmake b/cmake/Modules/FindARM.cmake
index 903025c5c2cfc..3d04db6ef4b5b 100644
--- a/cmake/Modules/FindARM.cmake
+++ b/cmake/Modules/FindARM.cmake
@@ -1,170 +1,38 @@
-# Check if the processor is an ARM and if Neon instruction are available on the machine where
-# the project is compiled.
-
-IF(CMAKE_SYSTEM_NAME MATCHES "Linux")
-   EXECUTE_PROCESS(COMMAND cat /proc/cpuinfo OUTPUT_VARIABLE CPUINFO)
-
-   #neon instruction can be found on the majority part of modern ARM processor
-   STRING(REGEX REPLACE "^.*(neon).*$" "\\1" NEON_THERE "${CPUINFO}")
-   STRING(COMPARE EQUAL "neon" "${NEON_THERE}" NEON_TRUE)
-   IF (NEON_TRUE)
-      set(NEON_FOUND true CACHE BOOL "NEON available on host")
-   ELSE (NEON_TRUE)
-      set(NEON_FOUND false CACHE BOOL "NEON available on host")
-   ENDIF (NEON_TRUE)
-
-   # on ARMv8, neon is inherit and instead listed as 'asimd' in /proc/cpuinfo
-   STRING(REGEX REPLACE "^.*(asimd).*$" "\\1" ASIMD_THERE "${CPUINFO}")
-   STRING(COMPARE EQUAL "asimd" "${ASIMD_THERE}" ASIMD_TRUE)
-   IF (ASIMD_TRUE)
-      set(ASIMD_FOUND true CACHE BOOL "ASIMD/NEON available on host")
-   ELSE (ASIMD_TRUE)
-      set(ASIMD_FOUND false CACHE BOOL "ASIMD/NEON available on host")
-   ENDIF (ASIMD_TRUE)
-
-   #sve instruction can be found on the majority part of modern ARM processor
-   STRING(REGEX REPLACE "^.*(sve).*$" "\\1" SVE_THERE ${CPUINFO})
-   STRING(COMPARE EQUAL "sve" "${SVE_THERE}" SVE_TRUE)
-   IF (SVE_TRUE)
-      set(SVE_FOUND true CACHE BOOL "SVE available on host")
-   ELSE (SVE_TRUE)
-      set(SVE_FOUND false CACHE BOOL "SVE available on host")
-   ENDIF (SVE_TRUE)
-
-   #Find the processor type (for now OMAP3 or OMAP4)
-   STRING(REGEX REPLACE "^.*(OMAP3).*$" "\\1" OMAP3_THERE "${CPUINFO}")
-   STRING(COMPARE EQUAL "OMAP3" "${OMAP3_THERE}" OMAP3_TRUE)
-   IF (OMAP3_TRUE)
-      set(CORTEXA8_FOUND true CACHE BOOL "OMAP3 available on host")
-   ELSE (OMAP3_TRUE)
-      set(CORTEXA8_FOUND false CACHE BOOL "OMAP3 available on host")
-   ENDIF (OMAP3_TRUE)
-
-   #Find the processor type (for now OMAP3 or OMAP4)
-   STRING(REGEX REPLACE "^.*(OMAP4).*$" "\\1" OMAP4_THERE "${CPUINFO}")
-   STRING(COMPARE EQUAL "OMAP4" "${OMAP4_THERE}" OMAP4_TRUE)
-   IF (OMAP4_TRUE)
-      set(CORTEXA9_FOUND true CACHE BOOL "OMAP4 available on host")
-   ELSE (OMAP4_TRUE)
-      set(CORTEXA9_FOUND false CACHE BOOL "OMAP4 available on host")
-   ENDIF (OMAP4_TRUE)
-
-ELSEIF(CMAKE_SYSTEM_NAME MATCHES "Darwin")
-   IF(CMAKE_SYSTEM_PROCESSOR STREQUAL "arm64" AND NOT CMAKE_OSX_ARCHITECTURES STREQUAL "x86_64")
-      set(NEON_FOUND true CACHE BOOL "NEON available on ARM64")
-   ENDIF()
-   EXECUTE_PROCESS(COMMAND /usr/sbin/sysctl -n machdep.cpu.features OUTPUT_VARIABLE
-      CPUINFO)
-
-   IF(NOT CPUINFO STREQUAL "")
-       #neon instruction can be found on the majority part of modern ARM processor
-       STRING(REGEX REPLACE "^.*(neon).*$" "\\1" NEON_THERE "${CPUINFO}")
-       STRING(COMPARE EQUAL "neon" "${NEON_THERE}" NEON_TRUE)
-       IF (NEON_TRUE)
-          set(NEON_FOUND true CACHE BOOL "NEON available on host")
-       ELSE (NEON_TRUE)
-          set(NEON_FOUND false CACHE BOOL "NEON available on host")
-       ENDIF (NEON_TRUE)
-   ENDIF()
-
-ELSEIF(CMAKE_SYSTEM_NAME MATCHES "Windows")
-   # TODO
-   set(CORTEXA8_FOUND   false CACHE BOOL "OMAP3 not available on host")
-   set(CORTEXA9_FOUND   false CACHE BOOL "OMAP4 not available on host")
-   set(NEON_FOUND   false CACHE BOOL "NEON not available on host")
-ELSE(CMAKE_SYSTEM_NAME MATCHES "Linux")
-   set(CORTEXA8_FOUND   false CACHE BOOL "OMAP3 not available on host")
-   set(CORTEXA9_FOUND   false CACHE BOOL "OMAP4 not available on host")
-   set(NEON_FOUND   false CACHE BOOL "NEON not available on host")
-ENDIF(CMAKE_SYSTEM_NAME MATCHES "Linux")
-
-if(NOT NEON_FOUND)
-      MESSAGE(STATUS "Could not find hardware support for NEON on this machine.")
-endif(NOT NEON_FOUND)
-if(NOT CORTEXA8_FOUND)
-      MESSAGE(STATUS "No OMAP3 processor on this machine.")
-endif(NOT CORTEXA8_FOUND)
-if(NOT CORTEXA9_FOUND)
-      MESSAGE(STATUS "No OMAP4 processor on this machine.")
-endif(NOT CORTEXA9_FOUND)
-mark_as_advanced(NEON_FOUND)
-
-#SVE support is available is only for Linux OS.
+# Check ARM feature availability for SVE compile-time support.
 IF(CMAKE_SYSTEM_NAME MATCHES "Linux")
-    # Include necessary modules for checking C and C++ source compilations
-    INCLUDE(CheckCSourceCompiles)
     INCLUDE(CheckCXXSourceCompiles)
 
-    # Test code for SVE support
-    SET(SVE_CODE "
+    SET(SVE_BF16_CODE "
       #include <arm_sve.h>
+      #include <arm_neon.h>
       int main()
       {
         svfloat64_t a;
         a = svdup_n_f64(0);
+        float32x4_t b = vdupq_n_f32(0);
+        bfloat16x8_t c = vreinterpretq_bf16_f32(b);
+        bfloat16x4_t d = vget_low_bf16(c);
         return 0;
       }
     ")
 
-    SET(ARM_BF16_CODE "
-      #include <arm_neon.h>
-      int main()
-      {
-        float32x4_t a = vdupq_n_f32(0);
-        bfloat16x8_t b = vreinterpretq_bf16_f32(a);
-        return 0;
-      }
-    ")
-
-    # Macro to check for SVE instruction support
-    MACRO(CHECK_COMPILES lang type flags code)
-      # Save the current state of required flags
-      SET(CMAKE_REQUIRED_FLAGS_SAVE ${CMAKE_REQUIRED_FLAGS})
-
-      # Set the flags necessary for compiling the test code with SVE support
-      SET(CMAKE_REQUIRED_FLAGS "${CMAKE_${lang}_FLAGS_INIT} ${flags}")
-
-      # Check if the source code compiles with the given flags for the specified language (C or C++)
-      IF(lang STREQUAL "CXX")
-        CHECK_CXX_SOURCE_COMPILES("${code}" ${lang}_HAS_${type})
-      ELSE()
-        CHECK_C_SOURCE_COMPILES("${code}" ${lang}_HAS_${type})
-      ENDIF()
-
-      # If the compilation test is successful, set appropriate variables indicating support
-      IF(${lang}_HAS_${type})
-        set(${lang}_SVE_FOUND TRUE CACHE BOOL "SVE available on host")
-        SET(${lang}_${type}_FOUND TRUE CACHE BOOL "${lang} ${type} support")
-        SET(${lang}_${type}_FLAGS "${flags}" CACHE STRING "${lang} ${type} flags")
-      ENDIF()
-
-      # Restore the original state of required flags
-      SET(CMAKE_REQUIRED_FLAGS ${CMAKE_REQUIRED_FLAGS_SAVE})
+    SET(CMAKE_REQUIRED_FLAGS_SAVE ${CMAKE_REQUIRED_FLAGS})
+    SET(CMAKE_REQUIRED_FLAGS "${CMAKE_CXX_FLAGS_INIT} -march=armv8-a+sve+bf16 -msve-vector-bits=256")
+    CHECK_CXX_SOURCE_COMPILES("${SVE_BF16_CODE}" CXX_SVE256_FOUND)
+    SET(CMAKE_REQUIRED_FLAGS ${CMAKE_REQUIRED_FLAGS_SAVE})
 
-      # If the compilation test fails, indicate that the support is not found
-      IF(NOT ${lang}_${type}_FOUND)
-        SET(${lang}_${type}_FOUND FALSE CACHE BOOL "${lang} ${type} support")
-        SET(${lang}_${type}_FLAGS "" CACHE STRING "${lang} ${type} flags")
-      ENDIF()
-
-      # Mark the variables as advanced to hide them in the default CMake GUI
-      MARK_AS_ADVANCED(${lang}_${type}_FOUND ${lang}_${type}_FLAGS)
-    ENDMACRO()
-
-    # Check for SVE256 vector length
-    CHECK_COMPILES(CXX "SVE256" "-march=armv8.2-a+sve -msve-vector-bits=256" "${SVE_CODE}")
-    CHECK_COMPILES(CXX "ARM_BF16" "-march=armv8.2-a+sve+bf16 -msve-vector-bits=256" "${ARM_BF16_CODE}")
-
-    # If SVE256 support is not found, set CXX_SVE_FOUND to FALSE and notify the user
-    if(NOT CXX_SVE256_FOUND)
-      set(CXX_SVE_FOUND FALSE CACHE BOOL "SVE not available on host")
-      message(STATUS "No SVE processor on this machine.")
-    else()
-      # If SVE256 support is found, set CXX_SVE_FOUND to TRUE and notify the user
+    if(CXX_SVE256_FOUND)
       set(CXX_SVE_FOUND TRUE CACHE BOOL "SVE available on host")
       message(STATUS "SVE support detected.")
+    else()
+      set(CXX_SVE_FOUND FALSE CACHE BOOL "SVE not available on host")
+      if(CMAKE_SYSTEM_PROCESSOR STREQUAL "aarch64" AND NOT DEFINED ENV{BUILD_IGNORE_SVE_UNAVAILABLE})
+        message(FATAL_ERROR "No SVE support on this machine. "
+          "Set BUILD_IGNORE_SVE_UNAVAILABLE environment variable to ignore this error.")
+      else()
+        message(STATUS "No SVE support on this machine.")
+      endif()
     endif()
 
-    # Mark the SVE support variable as advanced
-    mark_as_advanced(CXX_SVE_FOUND)
-ENDIF(CMAKE_SYSTEM_NAME MATCHES "Linux")
+    mark_as_advanced(CXX_SVE_FOUND CXX_SVE256_FOUND)
+ENDIF()
diff --git a/cmake/Modules/FindMKLDNN.cmake b/cmake/Modules/FindMKLDNN.cmake
index 7f53dacadef59..10ee8546b8d9d 100644
--- a/cmake/Modules/FindMKLDNN.cmake
+++ b/cmake/Modules/FindMKLDNN.cmake
@@ -47,7 +47,7 @@ IF(NOT MKLDNN_FOUND)
     endif()
     ExternalProject_Add(xpu_mkldnn_proj
       GIT_REPOSITORY https://github.com/uxlfoundation/oneDNN
-      GIT_TAG v3.10.2
+      GIT_TAG v3.11.2
       PREFIX ${XPU_MKLDNN_DIR_PREFIX}
       BUILD_IN_SOURCE 0
       CMAKE_ARGS  -DCMAKE_C_COMPILER=icx
diff --git a/cmake/Modules/FindNCCL.cmake b/cmake/Modules/FindNCCL.cmake
index cef8002f81706..447753f48c615 100644
--- a/cmake/Modules/FindNCCL.cmake
+++ b/cmake/Modules/FindNCCL.cmake
@@ -57,31 +57,6 @@ if(NCCL_FOUND)  # obtaining NCCL version and some sanity checks
   include(CheckCXXSymbolExists)
   check_cxx_symbol_exists(NCCL_VERSION_CODE nccl.h NCCL_VERSION_DEFINED)
 
-  # this condition check only works for non static NCCL linking
-  if (NCCL_VERSION_DEFINED AND NOT USE_STATIC_NCCL)
-    set(file "${PROJECT_BINARY_DIR}/detect_nccl_version.cc")
-    file(WRITE ${file} "
-      #include <iostream>
-      #include <nccl.h>
-      int main()
-      {
-        std::cout << NCCL_MAJOR << '.' << NCCL_MINOR << '.' << NCCL_PATCH << std::endl;
-        int x;
-        ncclGetVersion(&x);
-        return x == NCCL_VERSION_CODE;
-      }
-")
-    try_run(NCCL_VERSION_MATCHED compile_result ${PROJECT_BINARY_DIR} ${file}
-          RUN_OUTPUT_VARIABLE NCCL_VERSION_FROM_HEADER
-          CMAKE_FLAGS  "-DINCLUDE_DIRECTORIES=${NCCL_INCLUDE_DIRS}"
-          LINK_LIBRARIES ${NCCL_LIBRARIES})
-    if (NOT NCCL_VERSION_MATCHED)
-      message(FATAL_ERROR "Found NCCL header version and library version do not match! \
-(include: ${NCCL_INCLUDE_DIRS}, library: ${NCCL_LIBRARIES}) Please set NCCL_INCLUDE_DIR and NCCL_LIB_DIR manually.")
-    endif()
-    message(STATUS "NCCL version: ${NCCL_VERSION_FROM_HEADER}")
-  endif ()
-
   set (CMAKE_REQUIRED_INCLUDES ${OLD_CMAKE_REQUIRED_INCLUDES})
   message(STATUS "Found NCCL (include: ${NCCL_INCLUDE_DIRS}, library: ${NCCL_LIBRARIES})")
   mark_as_advanced(NCCL_ROOT_DIR NCCL_INCLUDE_DIRS NCCL_LIBRARIES)
diff --git a/cmake/Modules_CUDA_fix/upstream/FindCUDA/select_compute_arch.cmake b/cmake/Modules_CUDA_fix/upstream/FindCUDA/select_compute_arch.cmake
index bf7edd69ccd13..642b96ba119e8 100644
--- a/cmake/Modules_CUDA_fix/upstream/FindCUDA/select_compute_arch.cmake
+++ b/cmake/Modules_CUDA_fix/upstream/FindCUDA/select_compute_arch.cmake
@@ -30,10 +30,10 @@ endif()
 set(CUDA_KNOWN_GPU_ARCHITECTURES  "Kepler" "Maxwell")
 
 # This list will be used for CUDA_ARCH_NAME = Common option (enabled by default)
-set(CUDA_COMMON_GPU_ARCHITECTURES "3.5" "5.0")
+set(CUDA_COMMON_GPU_ARCHITECTURES "5.0")
 
 # This list is used to filter CUDA archs when autodetecting
-set(CUDA_ALL_GPU_ARCHITECTURES "3.5" "5.0")
+set(CUDA_ALL_GPU_ARCHITECTURES "5.0")
 
 if(CUDA_VERSION VERSION_GREATER "10.5")
   list(APPEND CUDA_KNOWN_GPU_ARCHITECTURES "Ampere")
@@ -65,19 +65,10 @@ if(NOT CUDA_VERSION VERSION_LESS "11.8")
   list(APPEND CUDA_ALL_GPU_ARCHITECTURES "8.9")
   list(APPEND CUDA_ALL_GPU_ARCHITECTURES "9.0")
 
-  if(CUDA_VERSION VERSION_LESS "12.0")
-    set(CUDA_LIMIT_GPU_ARCHITECTURE "9.0")
-    list(APPEND CUDA_COMMON_GPU_ARCHITECTURES "8.9+PTX")
-    list(APPEND CUDA_COMMON_GPU_ARCHITECTURES "9.0+PTX")
-  endif()
 endif()
 
-if(NOT CUDA_VERSION VERSION_LESS "12.0")
-  list(APPEND CUDA_COMMON_GPU_ARCHITECTURES "9.0a")
-  list(APPEND CUDA_ALL_GPU_ARCHITECTURES "9.0a")
-  list(REMOVE_ITEM CUDA_COMMON_GPU_ARCHITECTURES "3.5")
-  list(REMOVE_ITEM CUDA_ALL_GPU_ARCHITECTURES "3.5")
-endif()
+list(APPEND CUDA_COMMON_GPU_ARCHITECTURES "9.0a")
+list(APPEND CUDA_ALL_GPU_ARCHITECTURES "9.0a")
 
 if(CUDA_VERSION VERSION_GREATER "12.6")
   list(APPEND CUDA_KNOWN_GPU_ARCHITECTURES "Blackwell")
@@ -91,6 +82,12 @@ if(CUDA_VERSION VERSION_GREATER "12.6")
   list(APPEND CUDA_ALL_GPU_ARCHITECTURES "10.1a")
   list(APPEND CUDA_ALL_GPU_ARCHITECTURES "12.0")
   list(APPEND CUDA_ALL_GPU_ARCHITECTURES "12.0a")
+  if(NOT CUDA_VERSION VERSION_LESS "13.0")
+    list(REMOVE_ITEM CUDA_COMMON_GPU_ARCHITECTURES "10.1a")
+    list(REMOVE_ITEM CUDA_ALL_GPU_ARCHITECTURES "10.1a")
+    list(APPEND CUDA_COMMON_GPU_ARCHITECTURES "11.0a")
+    list(APPEND CUDA_ALL_GPU_ARCHITECTURES "11.0a")
+  endif()
 endif()
 
 
diff --git a/cmake/Summary.cmake b/cmake/Summary.cmake
index 686f5861960d3..2d68d81c6b367 100644
--- a/cmake/Summary.cmake
+++ b/cmake/Summary.cmake
@@ -149,6 +149,7 @@ function(caffe2_print_configuration_summary)
   message(STATUS "  USE_PYTORCH_METAL_EXPORT     : ${USE_PYTORCH_METAL_EXPORT}")
   message(STATUS "  USE_MPS               : ${USE_MPS}")
   message(STATUS "  CAN_COMPILE_METAL     : ${CAN_COMPILE_METAL}")
+  message(STATUS "  CAN_COMPILE_METAL_40  : ${CAN_COMPILE_METAL_40}")
   message(STATUS "  USE_MKL               : ${CAFFE2_USE_MKL}")
   if(${CAFFE2_USE_MKL})
     message(STATUS "    USE_STATIC_MKL      : ${USE_STATIC_MKL}")
@@ -177,7 +178,11 @@ function(caffe2_print_configuration_summary)
   if(${USE_NCCL})
     message(STATUS "    USE_SYSTEM_NCCL     : ${USE_SYSTEM_NCCL}")
   endif()
-  message(STATUS "  Found NVSHMEM         : ${NVSHMEM_INCLUDE_DIR}")
+  if(${USE_ROCM})
+    message(STATUS "  Found ROCSHMEM         : ${ROCSHMEM_INCLUDE_DIR}")
+  else()
+    message(STATUS "  Found NVSHMEM         : ${NVSHMEM_INCLUDE_DIR}")
+  endif()
   message(STATUS "  USE_NNPACK            : ${USE_NNPACK}")
   message(STATUS "  USE_NUMPY             : ${USE_NUMPY}")
   message(STATUS "  USE_OBSERVERS         : ${USE_OBSERVERS}")
diff --git a/cmake/TorchConfig.cmake.in b/cmake/TorchConfig.cmake.in
index abf5c8149116f..5749c67db112b 100644
--- a/cmake/TorchConfig.cmake.in
+++ b/cmake/TorchConfig.cmake.in
@@ -155,7 +155,7 @@ if(NOT @BUILD_SHARED_LIBS@)
 endif()
 set_target_properties(torch PROPERTIES
     INTERFACE_INCLUDE_DIRECTORIES "${TORCH_INCLUDE_DIRS}"
-    CXX_STANDARD 17
+    CXX_STANDARD 20
 )
 if(TORCH_CXX_FLAGS)
   set_property(TARGET torch PROPERTY INTERFACE_COMPILE_OPTIONS "${TORCH_CXX_FLAGS}")
diff --git a/cmake/public/LoadHIP.cmake b/cmake/public/LoadHIP.cmake
index 3ff7b3d2c1b36..a1b632361bbbc 100644
--- a/cmake/public/LoadHIP.cmake
+++ b/cmake/public/LoadHIP.cmake
@@ -2,7 +2,8 @@ set(PYTORCH_FOUND_HIP FALSE)
 
 # If ROCM_PATH is set, assume intention is to compile with
 # ROCm support and error out if the ROCM_PATH does not exist.
-# Else ROCM_PATH does not exist, assume a default of /opt/rocm
+# Else ROCM_PATH does not exist, try to get it from rocm-sdk,
+# or assume a default of /opt/rocm
 # In the latter case, if /opt/rocm does not exist emit status
 # message and return.
 if(DEFINED ENV{ROCM_PATH})
@@ -13,11 +14,31 @@ if(DEFINED ENV{ROCM_PATH})
       "Set a valid ROCM_PATH or unset ROCM_PATH environment variable to fix.")
   endif()
 else()
-  if(UNIX)
-    set(ROCM_PATH /opt/rocm)
-  else() # Win32
-    set(ROCM_PATH C:/opt/rocm)
+  # Try to get ROCM_PATH from rocm-sdk if available
+  find_program(ROCM_SDK_EXECUTABLE rocm-sdk)
+  if(ROCM_SDK_EXECUTABLE)
+    execute_process(
+      COMMAND ${ROCM_SDK_EXECUTABLE} path --root
+      OUTPUT_VARIABLE ROCM_SDK_PATH
+      OUTPUT_STRIP_TRAILING_WHITESPACE
+      RESULT_VARIABLE ROCM_SDK_RESULT
+      ERROR_QUIET
+    )
+    if(ROCM_SDK_RESULT EQUAL 0 AND EXISTS "${ROCM_SDK_PATH}")
+      set(ROCM_PATH "${ROCM_SDK_PATH}")
+      message(STATUS "Found ROCm installation via rocm-sdk at: ${ROCM_PATH}")
+    endif()
   endif()
+
+  # Fall back to default paths if rocm-sdk did not work
+  if(NOT DEFINED ROCM_PATH OR NOT EXISTS ${ROCM_PATH})
+    if(UNIX)
+      set(ROCM_PATH /opt/rocm)
+    else() # Win32
+      set(ROCM_PATH C:/opt/rocm)
+    endif()
+  endif()
+
   if(NOT EXISTS ${ROCM_PATH})
     message(STATUS
         "ROCM_PATH environment variable is not set and ${ROCM_PATH} does not exist.\n"
@@ -54,7 +75,7 @@ message("Building PyTorch for GPU arch: ${PYTORCH_ROCM_ARCH}")
 # needed because the find_package call to this module uses the Module mode search
 # https://cmake.org/cmake/help/latest/command/find_package.html#search-modes
 if(UNIX)
-  set(CMAKE_MODULE_PATH ${ROCM_PATH}/lib/cmake/hip ${CMAKE_MODULE_PATH})
+  set(CMAKE_MODULE_PATH ${ROCM_PATH}/lib/cmake/hip;${ROCM_PATH}/lib/${CMAKE_LIBRARY_ARCHITECTURE}/cmake/hip ${CMAKE_MODULE_PATH})
 else() # Win32
   set(CMAKE_MODULE_PATH ${ROCM_PATH}/cmake/ ${CMAKE_MODULE_PATH})
 endif()
@@ -79,6 +100,9 @@ endmacro()
 # MODULE argument is added for clarity that CMake is searching
 # for FindHIP.cmake in Module mode
 find_package_and_print_version(HIP 1.0 MODULE)
+if(CMAKE_SYSTEM_NAME STREQUAL "Linux")
+  enable_language(HIP)
+endif()
 
 if(HIP_FOUND)
   set(PYTORCH_FOUND_HIP TRUE)
@@ -184,6 +208,7 @@ if(HIP_FOUND)
   find_package_and_print_version(rocthrust REQUIRED)
   find_package_and_print_version(hipsolver REQUIRED)
   find_package_and_print_version(rocsolver REQUIRED)
+  find_package_and_print_version(rocshmem)
   # workaround cmake 4 build issue
   if(CMAKE_VERSION VERSION_GREATER_EQUAL "4.0.0")
     message(WARNING "Work around hiprtc cmake failure for cmake >= 4")
@@ -202,6 +227,9 @@ if(HIP_FOUND)
 
   # Optional components.
   find_package_and_print_version(hipsparselt)  # Will be required when ready.
+  # ROCm 8.0 and later requires libhipcxx! This should be marked as
+  # 'REQUIRED' once minimal ROCm version is bumped to 8.0 or later.
+  find_package_and_print_version(libhipcxx)
 
   list(REMOVE_DUPLICATES ROCM_INCLUDE_DIRS)
 
diff --git a/cmake/public/cuda.cmake b/cmake/public/cuda.cmake
index 7bef03d96fafb..239641d352c41 100644
--- a/cmake/public/cuda.cmake
+++ b/cmake/public/cuda.cmake
@@ -78,8 +78,8 @@ endif()
 message(STATUS "PyTorch: CUDA detected: " ${CUDA_VERSION})
 message(STATUS "PyTorch: CUDA nvcc is: " ${CUDA_NVCC_EXECUTABLE})
 message(STATUS "PyTorch: CUDA toolkit directory: " ${CUDA_TOOLKIT_ROOT_DIR})
-if(CUDA_VERSION VERSION_LESS 12.0)
-  message(FATAL_ERROR "PyTorch requires CUDA 12.0 or above.")
+if(CUDA_VERSION VERSION_LESS 12.1)
+  message(FATAL_ERROR "PyTorch requires CUDA 12.1 or above.")
 endif()
 
 if(CUDA_FOUND)
@@ -307,6 +307,14 @@ else()
 endif()
 
 # nvrtc
+# cuDNN frontend needs libnvrtc symbols, but linking through CUDA::nvrtc pulls
+# CUDA::cuda_driver transitively. Keep a driver-free target for cuDNN users and
+# reserve caffe2::nvrtc for the stub library that actually needs the driver API.
+add_library(caffe2::nvrtc_runtime INTERFACE IMPORTED)
+set_property(
+    TARGET caffe2::nvrtc_runtime PROPERTY INTERFACE_LINK_LIBRARIES
+    "${CUDA_NVRTC_LIB}")
+
 add_library(caffe2::nvrtc INTERFACE IMPORTED)
 set_property(
     TARGET caffe2::nvrtc PROPERTY INTERFACE_LINK_LIBRARIES
@@ -376,6 +384,11 @@ elseif(CUDA_DEVICE_DEBUG)
   list(APPEND CUDA_NVCC_FLAGS "-g" "-G")  # -G enables device code debugging symbols
 endif()
 
+# needed for compat with newer versions of clang that use C++20 mangling rules
+if(CMAKE_CXX_COMPILER_ID MATCHES "Clang" AND CMAKE_CXX_COMPILER_VERSION VERSION_GREATER_EQUAL 18)
+  list(APPEND CUDA_NVCC_FLAGS "-Xcompiler=-fclang-abi-compat=17")
+endif()
+
 # Set expt-relaxed-constexpr to suppress Eigen warnings
 list(APPEND CUDA_NVCC_FLAGS "--expt-relaxed-constexpr")
 
diff --git a/cmake/public/utils.cmake b/cmake/public/utils.cmake
index 731bbe7d21862..a667ad55771b4 100644
--- a/cmake/public/utils.cmake
+++ b/cmake/public/utils.cmake
@@ -337,7 +337,7 @@ endmacro()
 # Usage:
 #   torch_compile_options(lib_name)
 function(torch_compile_options libname)
-  set_property(TARGET ${libname} PROPERTY CXX_STANDARD 17)
+  set_property(TARGET ${libname} PROPERTY CXX_STANDARD 20)
 
   # until they can be unified, keep these lists synced with setup.py
   if(MSVC)
@@ -391,12 +391,17 @@ function(torch_compile_options libname)
     endif()
     if(CMAKE_CXX_COMPILER_ID MATCHES "Clang")
       if(NOT USE_CUDA)
-        # NS: One can compile CUDA code with extra-semi flag as nvcc generates code like
+        # NS: One can not compile CUDA code with extra-semi flag as nvcc generates code like
         # namespace MemoryOps_cu_d8602b38_109889 __attribute__((visibility("hidden")))  { };
         list(APPEND private_compile_options -Wextra-semi)
       else()
         # NVCC + clang15  reports deprecated copies from GPU lambda instantiations
         list(APPEND private_compile_options -Wno-deprecated-copy)
+        # NVCC + clang18  reports spurious deprecated deprecated literal operator declaration when there were none
+        # I.e. failures look like torch/headeronly/util/complex.h:334:40: error: identifier '_if' preceded by whitespace in a literal operator declaration is deprecated
+        # but if one to look at the source code, there are no space there
+        list(APPEND private_compile_options -Wno-deprecated-literal-operator)
+
       endif()
       list(APPEND private_compile_options -Wmove)
     else()
diff --git a/cmake/public/xpu.cmake b/cmake/public/xpu.cmake
index b39e31d0ade8a..2731c2842c7f4 100644
--- a/cmake/public/xpu.cmake
+++ b/cmake/public/xpu.cmake
@@ -37,6 +37,11 @@ torch_xpu_get_arch_list(XPU_ARCH_FLAGS)
 # propagate to torch-xpu-ops
 set(TORCH_XPU_ARCH_LIST ${XPU_ARCH_FLAGS})
 
+# Ensure SYCL device code compiles with C++20 (matching CMAKE_CXX_STANDARD).
+# SYCL_FLAGS flows into SYCL_COMPILE_FLAGS in torch-xpu-ops' BuildFlags.cmake
+# and is passed directly to icpx on the device compilation command line.
+list(APPEND SYCL_FLAGS -std=c++20)
+
 # Ensure USE_XPU is enabled.
 string(APPEND XPU_HOST_CXX_FLAGS " -DUSE_XPU")
 string(APPEND XPU_HOST_CXX_FLAGS " -DSYCL_COMPILER_VERSION=${SYCL_COMPILER_VERSION}")
diff --git a/docs/cpp/Makefile b/docs/cpp/Makefile
index e244432b9fdb3..21379d1856690 100644
--- a/docs/cpp/Makefile
+++ b/docs/cpp/Makefile
@@ -2,24 +2,41 @@
 
 # You can set these variables from the command line.
 SPHINXOPTS    = -j auto
-SPHINXBUILD   = sphinx-build
+SPHINXBUILD   = python -m sphinx
 SPHINXPROJ    = PyTorch
 SOURCEDIR     = source
 BUILDDIR      = build
 PYCMD         = python
+DOXYGEN       = doxygen
 
 # Put it first so that "make" without argument is like "make help".
 help:
 	@$(SPHINXBUILD) -M help "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O)
 
-.PHONY: help Makefile clean
+.PHONY: help Makefile clean doxygen html
+
+# Generate Doxygen XML (required for Breathe directives)
+doxygen:
+	@echo "Running Doxygen..."
+	@cd $(SOURCEDIR) && $(DOXYGEN) Doxyfile
 
 clean:
-	@# Clean up sphinx and doxygen build artifacts.
+	@# Clean up sphinx build artifacts and Doxygen output.
+	@$(SPHINXBUILD) -M $@ "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O)
+	@rm -rf $(BUILDDIR)/xml
+
+# Build HTML with Doxygen XML generation first
+html: doxygen
 	@$(SPHINXBUILD) -M $@ "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O)
-	@# Clean up exhale generated api.
-	@echo "Removing everything under generated 'source/api'..."
-	@rm -rf $(SOURCEDIR)/api
+
+# Check documentation coverage against curated allowlist and HTML formatting
+coverage: html
+	@echo "Running C++ docs coverage check..."
+	@$(PYCMD) check_coverage.py
+
+coverage-only:
+	@echo "Running C++ docs coverage check (no rebuild)..."
+	@$(PYCMD) check_coverage.py
 
 # Catch-all target: route all unknown targets to Sphinx using the new
 # "make mode" option.  $(O) is meant as a shortcut for $(SPHINXOPTS).
diff --git a/docs/cpp/README.md b/docs/cpp/README.md
new file mode 100644
index 0000000000000..0808a5cb4af0b
--- /dev/null
+++ b/docs/cpp/README.md
@@ -0,0 +1,315 @@
+# PyTorch C++ API Documentation
+
+This directory contains the source for PyTorch's C++ API documentation, built with
+Sphinx, Breathe, and Doxygen.
+
+## How it works
+
+The documentation pipeline has three stages:
+
+```
+C++ headers ──→ Doxygen ──→ XML ──→ Breathe ──→ Sphinx ──→ HTML
+                  ↑                    ↑
+              Doxyfile          Markdown files with
+           (which headers)    Breathe directives (MyST)
+```
+
+1. **Doxygen** reads C++ headers listed in `source/Doxyfile` `INPUT` and produces
+   XML in `build/xml/`.
+2. **Breathe** is a Sphinx extension that reads Doxygen XML and makes it available
+   via directives like `` ```{doxygenclass} ClassName `` (MyST Markdown syntax).
+3. **Sphinx** builds the final HTML from `.md` files in `source/`, using Breathe
+   directives to pull in C++ API documentation. The `myst_parser` extension
+   enables Markdown support.
+
+Only headers listed in the Doxyfile's `INPUT` are processed. Source files must
+explicitly reference each symbol — nothing is auto-generated.
+
+## Building
+
+```bash
+make html        # Build Doxygen XML + Sphinx HTML
+make doxygen     # Build Doxygen XML only
+make clean       # Clean build artifacts
+```
+
+The output is in `build/html/`.
+
+## Contributing to the C++ docs
+
+### Adding a new API
+
+1. **Ensure the header is in the Doxyfile** — check that `source/Doxyfile` `INPUT`
+   includes the header file or its parent directory. If not, add it:
+
+   ```
+   INPUT = ... \
+           ../../../path/to/your/header.h
+   ```
+
+2. **Add a Breathe directive to the appropriate `.md` file** under `source/api/`.
+   See [Which directive to use](#which-directive-to-use) below.
+
+3. **Run `make html`** and check the output in `build/html/`.
+
+4. **Run `python check_coverage.py`** to verify your API shows as documented.
+
+### Which directive to use
+
+Use Breathe directives to pull documentation from Doxygen XML. These render the
+full C++ signature, doc comments, parameters, and members automatically.
+
+All source files use MyST Markdown syntax (fenced directives with backticks).
+
+**Classes and structs** — use `doxygenclass` or `doxygenstruct`:
+
+````markdown
+```{doxygenclass} torch::nn::Linear
+:members:
+:undoc-members:
+```
+````
+
+- `:members:` shows all public member functions and variables
+- `:undoc-members:` includes members without doc comments
+- Omit both flags to show only the class description (useful when `:members:`
+  causes rendering issues)
+
+**Free functions** — use `doxygenfunction`:
+
+````markdown
+```{doxygenfunction} torch::autograd::grad
+```
+````
+
+For overloaded functions, Breathe will document all overloads.
+
+**Macros** — use `doxygendefine`:
+
+````markdown
+```{doxygendefine} TORCH_LIBRARY
+```
+````
+
+**Typedefs** — use `doxygentypedef`:
+
+````markdown
+```{doxygentypedef} torch::DeviceType
+```
+````
+
+### When Breathe directives don't work
+
+Some symbols can't be documented with Breathe directives:
+
+- **TORCH_MODULE holder classes** (e.g., `Conv2d`, `Linear`): The `TORCH_MODULE()`
+  macro generates these, but Doxygen can't index them. Document the `*Impl` class
+  instead (e.g., `Conv2dImpl`) — it contains all the actual methods.
+
+- **Functions with broken `\rst`/`\endrst` blocks**: Some doc comments use Doxygen's
+  `\rst` alias to embed RST. When a comment has multiple such blocks, Doxygen
+  generates malformed XML and Breathe renders raw text. In these cases, either:
+  - Fix the header to use native Doxygen (`@code{.cpp}`/`@endcode`, `@note`,
+    `@warning`) instead of `\rst`/`\endrst`
+  - Use a hand-written Sphinx C++ domain directive as a fallback
+
+- **Functions with mismatched `\param` names**: If a header's `\param` names don't
+  match the actual parameter names, Doxygen may fail to index the function. Use a
+  hand-written `cpp:function` directive instead.
+
+**Hand-written Sphinx C++ domain directives** (fallback):
+
+````markdown
+```{cpp:function} void torch::autograd::backward(const variable_list& tensors, const variable_list& grad_tensors = {}, std::optional<bool> retain_graph = std::nullopt, bool create_graph = false, const variable_list& inputs = {})
+
+Computes gradients of given tensors with respect to graph leaves.
+
+:param tensors: Tensors of which the derivative will be computed.
+:param grad_tensors: The "vector" in the Jacobian-vector product.
+```
+````
+
+These don't pull from Doxygen — you write the signature and docs manually.
+
+**`{eval-rst}` escape hatch** — if a MyST directive doesn't render correctly,
+you can embed raw RST:
+
+````markdown
+```{eval-rst}
+.. doxygenclass:: X::A
+   :members:
+   :protected-members:
+   :private-members:
+```
+````
+
+### Source file structure
+
+Each `.md` file under `source/api/` documents one topic area using MyST Markdown.
+The typical pattern:
+
+````markdown
+# Page Title
+
+Brief description of this API area.
+
+## Section Name
+
+Optional prose explaining usage, with a code example:
+
+```cpp
+#include <torch/torch.h>
+auto x = torch::randn({2, 3});
+```
+
+```{doxygenclass} torch::nn::SomeClass
+:members:
+:undoc-members:
+```
+
+```{doxygenstruct} torch::nn::SomeClassOptions
+:members:
+:undoc-members:
+```
+````
+
+**Nesting directives:** When a directive contains other directives or code blocks,
+the outer fence must use more backticks than the inner ones:
+
+`````markdown
+````{cpp:class} at::Tensor
+
+The primary tensor class.
+
+```{cpp:function} int64_t dim() const
+
+Returns the number of dimensions.
+```
+````
+`````
+
+### Writing doc comments in C++ headers
+
+Doxygen extracts documentation from comments in headers. Use `///` style:
+
+```cpp
+/// Applies a linear transformation to the incoming data: :math:`y = xA^T + b`.
+///
+/// @code{.cpp}
+/// auto linear = torch::nn::Linear(torch::nn::LinearOptions(10, 5));
+/// auto output = linear->forward(input);
+/// @endcode
+///
+/// @note The weight matrix is transposed compared to the Python API.
+class LinearImpl : public Cloneable<LinearImpl> {
+```
+
+**Preferred Doxygen commands:**
+- `@code{.cpp}` / `@endcode` for code examples
+- `@note` for important notes
+- `@warning` for warnings
+- `@param name` for parameter descriptions
+- `@return` for return value descriptions
+
+**Avoid** `\rst` / `\endrst` blocks — they cause rendering issues when a single
+doc comment contains multiple blocks. Use native Doxygen commands instead.
+
+## Coverage checking
+
+```bash
+make coverage        # Build docs + run coverage check
+make coverage-only   # Run coverage check without rebuilding
+```
+
+`check_coverage.py` auto-discovers public APIs from Doxygen XML (`build/xml/index.xml`)
+and checks which ones have Breathe or Sphinx directives in the source files.
+
+Reports are written to `cpp_coverage.txt` and `cpp_html_issues.txt`.
+
+### How coverage discovery works
+
+The script parses `build/xml/index.xml` to find all classes, structs, functions,
+and macros that Doxygen indexed. It then:
+
+1. Filters out internal symbols using `EXCLUDED_PATTERNS` and `EXCLUDED_SYMBOLS`
+2. Re-includes any symbols in `INCLUDED_SYMBOLS` (overrides exclusions)
+3. Scans `.md` files for Breathe (`doxygenclass`, `doxygenfunction`, etc.) and
+   Sphinx C++ domain (`cpp:class`, `cpp:function`, etc.) directives
+4. Reports the gap
+
+### Excluding internal APIs from coverage
+
+Not everything in Doxygen XML is public API. The script has three exclusion mechanisms:
+
+**`EXCLUDED_PATTERNS`** — regex patterns for broad categories:
+
+```python
+EXCLUDED_PATTERNS = [
+    r".*::detail::.*",     # Internal namespaces
+    r"torch::jit::.*",     # Deprecated
+    r".*::_\w+",           # Underscore-prefixed internals
+]
+```
+
+**`EXCLUDED_SYMBOLS`** — exact match for specific symbols:
+
+```python
+EXCLUDED_SYMBOLS = {
+    "torch::autograd::deleteNode",
+    "at::native::dataSize",
+}
+```
+
+**Doxyfile `EXCLUDE`** — prevents Doxygen from indexing files entirely:
+
+```
+EXCLUDE = ../../../torch/csrc/api/include/torch/detail
+```
+
+### Including "internal" APIs that are actually public
+
+Some APIs live in internal-looking namespaces but are widely used by external
+developers. To track these for documentation coverage, add them to
+`INCLUDED_SYMBOLS` in `check_coverage.py`:
+
+```python
+INCLUDED_SYMBOLS: set[str] = {
+    "c10::IValue",  # Used for custom op registration
+}
+```
+
+`INCLUDED_SYMBOLS` takes priority over all exclusions.
+
+### HTML formatting checks
+
+The script also checks built HTML for:
+- Unresolved Breathe directives ("Cannot find class/struct/function")
+- Raw directive text in output (build failures)
+- Sphinx "problematic" nodes (broken references)
+- Near-empty API pages
+
+### coverxygen (optional)
+
+```bash
+python check_coverage.py --coverxygen
+```
+
+This runs [coverxygen](https://github.com/psycofdj/coverxygen) on the Doxygen XML
+to measure what percentage of C++ symbols have doc comments in the source code.
+This is complementary to file coverage — it tells you which headers need more
+doc comments, not which symbols need Sphinx directives.
+
+## Known issues
+
+- **`\rst`/`\endrst` rendering** — Some C++ headers use Doxygen's `\rst` alias to
+  embed RST in doc comments. When a single doc comment has multiple `\rst`/`\endrst`
+  blocks, Doxygen generates malformed XML and Breathe renders them as raw text.
+  Avoid using `:members:` on classes with this issue. Use hand-written examples
+  in the source files instead, or convert the headers to native Doxygen commands.
+
+- **TORCH_MODULE holder classes** — Doxygen can't index classes generated by the
+  `TORCH_MODULE()` macro. Use the `*Impl` class name in directives instead.
+
+- **"Subclassed by" plain text** — Without per-class pages (exhale), Breathe
+  renders "Subclassed by" lists as plain text instead of links. A `conf.py` hook
+  (`remove_subclassed_by`) strips these automatically.
diff --git a/docs/cpp/check_coverage.py b/docs/cpp/check_coverage.py
new file mode 100644
index 0000000000000..bd958671cacbf
--- /dev/null
+++ b/docs/cpp/check_coverage.py
@@ -0,0 +1,782 @@
+#!/usr/bin/env python3
+"""C++ API documentation coverage checker.
+
+Auto-discovers public C++ APIs from Doxygen XML output and checks which ones
+are documented in the RST source files via Breathe or Sphinx C++ domain
+directives.
+
+Uses an exclusion list (EXCLUDED_APIS) to skip internal/detail symbols that
+don't need public documentation, rather than maintaining a hardcoded allowlist.
+
+Additionally checks built HTML for broken formatting (empty pages,
+unresolved directives, rendering errors).
+
+Usage:
+    python check_coverage.py                  # RST coverage + HTML checks
+    python check_coverage.py --coverxygen     # also run coverxygen on Doxygen XML
+"""
+
+import argparse
+import re
+import subprocess
+import sys
+import xml.etree.ElementTree as ET
+from pathlib import Path
+
+
+# ─── Paths ───────────────────────────────────────────────────────────────────
+
+SCRIPT_DIR = Path(__file__).resolve().parent
+SOURCE_DIR = SCRIPT_DIR / "source"
+BUILD_HTML = SCRIPT_DIR / "build" / "html"
+BUILD_XML = SCRIPT_DIR / "build" / "xml"
+COVERAGE_OUTPUT = SCRIPT_DIR / "cpp_coverage.txt"
+HTML_REPORT = SCRIPT_DIR / "cpp_html_issues.txt"
+
+# ─── Inclusion override ──────────────────────────────────────────────────────
+# Symbols that match an exclusion pattern but should still be tracked.
+# Use this for "internal" APIs that are widely used as public API.
+# Add the fully-qualified symbol name here and it will bypass all exclusions.
+INCLUDED_SYMBOLS: set[str] = {
+    # Example: "c10::IValue" would track it even though c10::IValue is excluded
+}
+
+# ─── Exclusion list ──────────────────────────────────────────────────────────
+# Symbols that should NOT be flagged as missing documentation.
+# Add internal, detail, or otherwise non-public symbols here.
+# Note: INCLUDED_SYMBOLS takes priority over these exclusions.
+
+EXCLUDED_PATTERNS = [
+    # Internal/detail namespaces
+    r".*::detail::.*",
+    r".*::detail_::.*",
+    r"torch::python::.*",
+    # Underscore-prefixed internal classes
+    r".*::_\w+",
+    # Enum helper structs
+    r"torch::enumtype::.*",
+    # OptimizerCloneableOptions SFINAE helpers
+    r"torch::optim::OptimizerCloneableOptions::.*",
+    # Internal optimizer state/options cloneable helpers
+    r"torch::optim::OptimizerCloneable.*",
+    # Error classes (c10 exceptions)
+    r"c10::.*Error$",
+    r"c10::ErrorAlwaysShowCppStacktrace",
+    # Warning internals
+    r"c10::Warning.*",
+    r"c10::WarningHandler",
+    r"c10::WarningUtils::.*",
+    # c10 IValue internals
+    r"c10::IValue::.*",
+    r"c10::IValue",
+    r"c10::WeakIValue",
+    r"c10::ivalue::.*",
+    r"c10::StrongTypePtr",
+    r"c10::WeakTypePtr",
+    r"c10::WeakOrStrongTypePtr",
+    r"c10::WeakOrStrongCompilationUnit",
+    r"c10::Capsule",
+    r"c10::OptionalArray",
+    r"c10::StreamData3",
+    # OrderedDict::Item (internal helper)
+    r"torch::OrderedDict::Item",
+    # ExpandingArray (internal template utility)
+    r"torch::ExpandingArray.*",
+    # IMethod (internal)
+    r"torch::IMethod",
+    # CustomClassHolder (internal base)
+    r"torch::CustomClassHolder",
+    # NodeGuard (internal autograd)
+    r"torch::autograd::NodeGuard",
+    # Autograd internals
+    r"torch::autograd::CppNode",
+    r"torch::autograd::ExtractVariables",
+    r"torch::autograd::Node",
+    r"torch::autograd::Node::.*",
+    r"torch::autograd::TraceableFunction",
+    r"torch::autograd::TypeAndSize",
+    # Sequencer internals
+    r"torch::data::.*::detail::.*",
+    # cuDNN descriptor internals
+    r"at::native::ActivationDescriptor",
+    r"at::native::ConvolutionDescriptor",
+    r"at::native::SpatialTransformerDescriptor",
+    r"at::native::DropoutDescriptor",
+    r"at::native::RNNDataDescriptor",
+    r"at::native::DftiDescriptor",
+    r"at::native::DescriptorDeleter",
+    r"at::native::DftiDescriptorDeleter",
+    r"at::native::RNNDescriptor",
+    # ATen internals
+    r"at::OptionalTensorRef",
+    r"at::TensorRef",
+    # at::cuda internals (allocator, workspace, cublas)
+    r"at::cuda::WorkspaceMapWithMutex",
+    r"at::cuda::clearCublasWorkspaces.*",
+    r"at::cuda::cublas_handle_stream_to_workspace",
+    r"at::cuda::cublaslt_handle_stream_to_workspace",
+    r"at::cuda::getCUDABlasLt.*",
+    r"at::cuda::getCUDADeviceAllocator",
+    r"at::cuda::getChosenWorkspaceSize",
+    r"at::cuda::getNumGPUs",
+    r"at::cuda::is_available",
+    r"at::cuda::warp_size",
+    # jit namespace (deprecated)
+    r"torch::jit::.*",
+    # Operators that are just operator<< or operator>>
+    r".*::operator<<",
+    r".*::operator>>",
+    r".*::operator==",
+    r".*::operator!=",
+    # Internal serialize helpers
+    r"torch::optim::serialize",
+    r"torch::optim::detail::.*",
+    # Reduction enum helpers
+    r"torch::nn::reduction",
+    r"torch::nn::log_target",
+    # Internal module utils
+    r"torch::nn::modules::utils::.*",
+    # Internal c10 helpers
+    r"c10::detail::.*",
+    r"c10::detail_::.*",
+    r"c10::makeArrayRef",
+    r"c10::checkObjectSortSchema",
+    r"c10::getGreaterThanComparator",
+    r"c10::getLessThanComparator",
+    r"c10::value_or_else",
+    r"c10::warn",
+    r"c10::GetExceptionString",
+    # torch::detail
+    r"torch::detail::.*",
+    # Internal data shuttle/queue
+    r"torch::data::detail::.*",
+    # DataLoaderBase internal types
+    r"torch::data::DataLoaderBase::.*",
+    r"torch::data::WorkerException",
+    r"torch::data::FullDataLoaderOptions",
+    # Template specializations of Stack
+    r"torch::data::transforms::Stack< .*>",
+    # Example partial specialization
+    r"torch::data::Example< .*>",
+    # Doxygen internal macros
+    r"DEFINE_CASE",
+    r"DEFINE_TAG",
+    r"COUNT_TAG",
+    r"TRUTH_TABLE_ENTRY",
+    r"C10_EXPAND_MSVC_WORKAROUND",
+    r"TORCH_FORALL_TAGS",
+    # Non-public torch::nn functions (module stream operators, etc.)
+    r"torch::nn::operator.*",
+    # AnyModule/AnyValue internal holders
+    r"torch::nn::AnyModuleHolder.*",
+    r"torch::nn::AnyModulePlaceholder",
+    r"torch::nn::AnyValue.*",
+    r"torch::nn::NamedAnyModule",
+    # Internal base classes (users use the derived classes)
+    r"torch::nn::ConvNdImpl",
+    r"torch::nn::ConvTransposeNdImpl",
+    r"torch::nn::BatchNormImplBase",
+    r"torch::nn::NormImplBase",
+    r"torch::nn::InstanceNormImpl",
+    r"torch::nn::MaxPoolImpl",
+    r"torch::nn::AvgPoolImpl",
+    r"torch::nn::AdaptiveAvgPoolImpl",
+    r"torch::nn::AdaptiveMaxPoolImpl",
+    r"torch::nn::MaxUnpoolImpl",
+    r"torch::nn::LPPoolImpl",
+    r"torch::nn::ConstantPadImpl",
+    r"torch::nn::ReflectionPadImpl",
+    r"torch::nn::ReplicationPadImpl",
+    r"torch::nn::ZeroPadImpl",
+    r"torch::nn::FractionalMaxPoolImpl",
+    # nn::functions internal namespace
+    r"torch::nn::functions::.*",
+    # AdaptiveLogSoftmaxWithLoss (niche, rarely used in C++)
+    r"torch::nn::AdaptiveLogSoftmaxWithLoss.*",
+    r"torch::nn::ASMoutput",
+    # CrossMapLRN2d (niche)
+    r"torch::nn::CrossMapLRN2d.*",
+    # _out function variants (documented alongside the main function)
+    r"torch::special::.*_out",
+    r"torch::fft::.*_out",
+    # torch internal helpers
+    r"torch::InitLambda",
+    r"torch::dispatch",
+    r"torch::equal_if_defined",
+    r"torch::getAllCustomClassesNames",
+    r"torch::init",
+    r"torch::make_custom_class",
+    r"torch::selective_class_",
+    r"torch::pickle_load",
+    r"torch::pickle_save",
+    r"torch::schema",
+    r"torch::nativert::.*",
+    # RNNCellOptionsBase (internal base)
+    r".*::RNNCellOptionsBase",
+    # Unnamespaced Options structs (indexed without namespace by Doxygen)
+    r"^[A-Z]\w+Options$",
+    # Unnamespaced classes without namespace (Doxygen quirk)
+    r"^TransformerDecoderLayer$",
+    r"^TransformerDecoderLayerOptions$",
+    # functional namespace internal options structs
+    r"functional::.*FuncOptions",
+]
+
+# Specific symbols to exclude (exact match)
+EXCLUDED_SYMBOLS = {
+    # Internal / not useful to document individually
+    "torch::data::datasets::map",
+    "torch::data::datasets::make_shared_dataset",
+    "torch::data::datasets::operator<<",
+    "torch::data::datasets::operator>>",
+    "torch::enumtype::get_enum_name",
+    "torch::enumtype::reduction_get_enum",
+    "torch::autograd::_wrap_outputs",
+    "torch::autograd::check_variable_result",
+    "torch::autograd::CppNode_apply_functional",
+    "torch::autograd::CppNode_apply_functional_ivalue",
+    "torch::autograd::forward_ad::enter_dual_level",
+    "torch::autograd::forward_ad::exit_dual_level",
+    "torch::autograd::any_variable_requires_grad",
+    "torch::autograd::collect_next_edges",
+    "torch::autograd::create_gradient_edge",
+    "torch::autograd::deleteNode",
+    "torch::autograd::extract_vars",
+    "torch::autograd::get_current_node",
+    "torch::autograd::to_optional",
+    "torch::autograd::to_output_type",
+    "torch::nn::parallel::replicate",
+    "torch::nn::parallel::parallel_apply",
+    "torch::nn::parallel::data_parallel",
+    "torch::python::add_module_bindings",
+    "torch::python::bind_module",
+    "torch::python::init_bindings",
+    # at::native cuDNN internals
+    "at::native::dataSize",
+    "at::native::fixSizeOneDimStride",
+    "at::native::operator<<",
+    "at::native::getCudnnDataTypeFromScalarType",
+    # c10 cuda pool functions (internal)
+    "c10::cuda::getStreamFromPool",
+    "c10::cuda::getStreamFromExternal",
+    "c10::xpu::getStreamFromPool",
+    "c10::xpu::getStreamFromExternal",
+    # c10 private use backend registration (internal)
+    "c10::get_privateuse1_backend",
+    "c10::is_privateuse1_backend_registered",
+    "c10::register_privateuse1_backend",
+    "c10::isValidDeviceType",
+    "c10::DeviceTypeName",
+    # torch::stable::detail internals
+    "torch::stable::detail::unbox_to_tuple_impl",
+    "torch::stable::detail::unbox_to_tuple",
+    "torch::stable::detail::box_from_tuple_impl",
+    "torch::stable::detail::box_from_tuple",
+    # torch::stable::accelerator (documented in stable API page)
+    "torch::stable::accelerator::getCurrentStream",
+}
+
+# Namespaces whose free functions should be checked for documentation
+PUBLIC_FUNCTION_NAMESPACES = {
+    "torch",
+    "torch::autograd",
+    "torch::cuda",
+    "torch::mps",
+    "torch::xpu",
+    "torch::fft",
+    "torch::special",
+    "torch::nn::functional",
+    "torch::nn::init",
+    "torch::nn::utils",
+    "torch::nn::utils::rnn",
+    "torch::data",
+    "torch::stable",
+    "torch::stable::accelerator",
+    "c10",
+    "c10::cuda",
+    "c10::xpu",
+    "at::cuda",
+}
+
+
+# ─── XML parsing ─────────────────────────────────────────────────────────────
+
+
+def _is_excluded(symbol: str) -> bool:
+    """Check if a symbol should be excluded from coverage tracking."""
+    if symbol in INCLUDED_SYMBOLS:
+        return False
+    if symbol in EXCLUDED_SYMBOLS:
+        return True
+    for pattern in EXCLUDED_PATTERNS:
+        if re.fullmatch(pattern, symbol):
+            return True
+    return False
+
+
+def _categorize(name: str) -> str:
+    """Assign a category based on the symbol's namespace."""
+    if name.startswith("torch::nn::functional::"):
+        return "torch::nn::functional"
+    if name.startswith("torch::nn::init::"):
+        return "torch::nn::init"
+    if name.startswith("torch::nn::utils::"):
+        return "torch::nn::utils"
+    if name.startswith("torch::nn::"):
+        # Distinguish modules from other nn symbols
+        short = name.split("::")[-1]
+        if short[0].isupper():
+            return "torch::nn (modules)"
+        return "torch::nn"
+    if name.startswith("torch::optim::"):
+        return "torch::optim"
+    if name.startswith("torch::data::"):
+        return "torch::data"
+    if name.startswith("torch::autograd::"):
+        return "torch::autograd"
+    if name.startswith("torch::serialize::") or name in ("torch::save", "torch::load"):
+        return "torch::serialize"
+    if name.startswith("torch::stable::"):
+        return "torch::stable"
+    if name.startswith("torch::fft::"):
+        return "torch::fft"
+    if name.startswith("torch::special::"):
+        return "torch::special"
+    if name.startswith(("torch::cuda::", "torch::mps::", "torch::xpu::")):
+        return "torch (device)"
+    if name.startswith("torch::"):
+        return "torch (core)"
+    if name.startswith("c10::cuda::"):
+        return "c10::cuda"
+    if name.startswith("c10::xpu::"):
+        return "c10::xpu"
+    if name.startswith("c10::"):
+        return "c10"
+    if name.startswith("at::cuda::"):
+        return "at::cuda"
+    if name.startswith("at::"):
+        return "at"
+    return "other"
+
+
+def discover_apis_from_xml(xml_dir: Path) -> dict[str, list[tuple[str, str]]]:
+    """Parse Doxygen index.xml to discover all public APIs.
+
+    Returns dict of category -> list of (symbol, kind).
+    """
+    index_path = xml_dir / "index.xml"
+    if not index_path.exists():
+        print(
+            f"ERROR: {index_path} not found. Run 'make doxygen' first.",
+            file=sys.stderr,
+        )
+        sys.exit(1)
+
+    tree = ET.parse(index_path)
+    root = tree.getroot()
+
+    apis: dict[str, list[tuple[str, str]]] = {}
+
+    # Collect classes and structs
+    for compound in root.findall("compound"):
+        kind = compound.get("kind")
+        if kind not in ("class", "struct"):
+            continue
+        name = compound.find("name").text
+        if _is_excluded(name):
+            continue
+        category = _categorize(name)
+        apis.setdefault(category, []).append((name, kind))
+
+    # Collect free functions from public namespaces
+    for compound in root.findall("compound"):
+        if compound.get("kind") != "namespace":
+            continue
+        ns_name = compound.find("name").text
+        if ns_name not in PUBLIC_FUNCTION_NAMESPACES:
+            continue
+        seen_funcs = set()
+        for member in compound.findall("member"):
+            if member.get("kind") != "function":
+                continue
+            func_name = member.find("name").text
+            qualified = f"{ns_name}::{func_name}"
+            if qualified in seen_funcs:
+                continue  # skip overloads
+            seen_funcs.add(qualified)
+            if _is_excluded(qualified):
+                continue
+            category = _categorize(qualified)
+            apis.setdefault(category, []).append((qualified, "function"))
+
+    # Collect macros (defines) from file compounds
+    for compound in root.findall("compound"):
+        if compound.get("kind") != "file":
+            continue
+        for member in compound.findall("member"):
+            if member.get("kind") != "define":
+                continue
+            macro_name = member.find("name").text
+            # Only track well-known public macros
+            if macro_name.startswith(("TORCH_LIBRARY", "TORCH_MODULE")):
+                if _is_excluded(macro_name):
+                    continue
+                apis.setdefault("torch (macros)", []).append((macro_name, "define"))
+
+    # Sort each category and deduplicate
+    for category in apis:
+        apis[category] = sorted(set(apis[category]))
+
+    return apis
+
+
+# ─── Source scanning ─────────────────────────────────────────────────────────
+
+# RST directives: .. doxygenclass:: torch::nn::ReLU
+RST_DIRECTIVE_RE = re.compile(
+    r"^\.\.\s+doxygen(class|struct|function|typedef|define|enum|namespace)"
+    r"::\s*(.+?)\s*$",
+    re.MULTILINE,
+)
+
+RST_CPP_DIRECTIVE_RE = re.compile(
+    r"^\.\.\s+cpp:(class|struct|function|enum|type)" r"::\s*(.+?)\s*$",
+    re.MULTILINE,
+)
+
+# MyST directives: ```{doxygenclass} torch::nn::ReLU
+MYST_DIRECTIVE_RE = re.compile(
+    r"^`{3,}\{doxygen(class|struct|function|typedef|define|enum|namespace)\}\s*(.+?)\s*$",
+    re.MULTILINE,
+)
+
+MYST_CPP_DIRECTIVE_RE = re.compile(
+    r"^`{3,}\{cpp:(class|struct|function|enum|type)\}\s*(.+?)\s*$",
+    re.MULTILINE,
+)
+
+
+def scan_sources(source_dir: Path) -> set[str]:
+    """Extract all documented symbols from RST/MyST breathe and cpp domain directives."""
+    documented = set()
+    for src_file in list(source_dir.rglob("*.rst")) + list(source_dir.rglob("*.md")):
+        content = src_file.read_text(errors="replace")
+        patterns = (
+            RST_DIRECTIVE_RE,
+            RST_CPP_DIRECTIVE_RE,
+            MYST_DIRECTIVE_RE,
+            MYST_CPP_DIRECTIVE_RE,
+        )
+        for pattern in patterns:
+            for match in pattern.finditer(content):
+                symbol = match.group(2)
+                # Strip template prefix
+                if symbol.startswith("template"):
+                    gt = symbol.find(">")
+                    if gt != -1:
+                        symbol = symbol[gt + 1 :].lstrip()
+                # Strip function signature
+                paren = symbol.find("(")
+                if paren != -1:
+                    symbol = symbol[:paren].rstrip()
+                documented.add(symbol)
+    return documented
+
+
+# ─── Coverage report ─────────────────────────────────────────────────────────
+
+
+def generate_coverage_report(
+    apis: dict[str, list[tuple[str, str]]], documented: set[str]
+) -> str:
+    """Generate a coverage report comparing discovered APIs against RST docs."""
+    lines = []
+    lines.append("Undocumented C++ objects")
+    lines.append("=" * 50)
+    lines.append("")
+
+    total = 0
+    total_missing = 0
+    section_stats = []
+
+    for category in sorted(apis.keys()):
+        symbols = apis[category]
+        section_missing = []
+        for symbol, kind in symbols:
+            total += 1
+            unqualified = symbol.rsplit("::", 1)[-1]
+            if symbol not in documented and unqualified not in documented:
+                section_missing.append((symbol, kind))
+                total_missing += 1
+
+        covered = len(symbols) - len(section_missing)
+        section_stats.append((category, covered, len(symbols)))
+
+        if section_missing:
+            lines.append(category)
+            lines.append("-" * len(category))
+            for symbol, kind in section_missing:
+                lines.append(f"   * {symbol}  ({kind})")
+            lines.append("")
+
+    # Summary
+    total_covered = total - total_missing
+    pct = (total_covered / total * 100) if total else 0
+
+    lines.append("")
+    lines.append("=" * 50)
+    lines.append("Summary")
+    lines.append("=" * 50)
+    lines.append("")
+    lines.append(f"Total APIs discovered:   {total}")
+    lines.append(f"Documented:              {total_covered}")
+    lines.append(f"Missing:                 {total_missing}")
+    lines.append(f"Coverage:                {pct:.1f}%")
+    lines.append("")
+
+    # Per-section table
+    lines.append(f"{'Category':<45} {'Covered':>8} {'Total':>6} {'%':>7}")
+    lines.append("-" * 70)
+    for category, covered, section_total in section_stats:
+        spct = (covered / section_total * 100) if section_total else 0
+        lines.append(f"{category:<45} {covered:>8} {section_total:>6} {spct:>6.1f}%")
+    lines.append("")
+
+    return "\n".join(lines)
+
+
+# ─── HTML checks ─────────────────────────────────────────────────────────────
+
+BROKEN_PATTERNS = [
+    (
+        re.compile(r"Cannot find (?:class|struct|function|file)", re.IGNORECASE),
+        "unresolved breathe directive",
+    ),
+    (
+        re.compile(r"Unable to resolve (?:class|struct|function)", re.IGNORECASE),
+        "unresolved breathe directive (ambiguous overload)",
+    ),
+    (
+        re.compile(r"doxygenclass:|doxygenfunction:|doxygenstruct:", re.IGNORECASE),
+        "raw directive text in output",
+    ),
+    (
+        re.compile(r"<span class=\"problematic\">", re.IGNORECASE),
+        "Sphinx problematic node (broken reference)",
+    ),
+    (
+        re.compile(r"System Message:", re.IGNORECASE),
+        "Sphinx system message (build error)",
+    ),
+]
+
+MIN_CONTENT_LENGTH = 500
+
+
+def check_html_output(build_dir: Path) -> str:
+    """Check built HTML for broken formatting and empty pages."""
+    issues = []
+
+    if not build_dir.exists():
+        return "ERROR: build/html directory not found. Run 'make html' first.\n"
+
+    for html_file in sorted(build_dir.rglob("*.html")):
+        rel = html_file.relative_to(build_dir)
+        if rel.name in ("search.html", "genindex.html", "objects.inv"):
+            continue
+
+        try:
+            content = html_file.read_text(errors="replace")
+        except Exception as e:
+            issues.append((str(rel), f"cannot read: {e}"))
+            continue
+
+        for pattern, description in BROKEN_PATTERNS:
+            matches = pattern.findall(content)
+            if matches:
+                issues.append((str(rel), f"{description} ({len(matches)}x)"))
+
+        if str(rel).startswith("api/"):
+            text = re.sub(r"<[^>]+>", "", content)
+            text = re.sub(r"\s+", " ", text).strip()
+            if len(text) < MIN_CONTENT_LENGTH:
+                issues.append((str(rel), f"possibly empty page ({len(text)} chars)"))
+
+    lines = []
+    lines.append("HTML Formatting Check")
+    lines.append("=" * 50)
+    lines.append("")
+
+    if not issues:
+        lines.append("No issues found.")
+    else:
+        lines.append(f"Found {len(issues)} issue(s):")
+        lines.append("")
+        lines.append(f"{'File':<55} Issue")
+        lines.append("-" * 90)
+        for filepath, issue in issues:
+            lines.append(f"{filepath:<55} {issue}")
+
+    lines.append("")
+    return "\n".join(lines)
+
+
+# ─── coverxygen integration ─────────────────────────────────────────────────
+
+
+def run_coverxygen(xml_dir: Path) -> str:
+    """Run coverxygen on Doxygen XML output for doc-comment coverage."""
+    lines = []
+    lines.append("Coverxygen Report (Doxygen doc-comment coverage)")
+    lines.append("=" * 50)
+    lines.append("")
+
+    if not xml_dir.exists():
+        lines.append("ERROR: build/xml directory not found. Run 'make doxygen' first.")
+        return "\n".join(lines)
+
+    coverxygen_cmd = None
+    for cmd in [
+        ["coverxygen", "--version"],
+        [sys.executable, "-m", "coverxygen", "--version"],
+    ]:
+        try:
+            subprocess.run(cmd, capture_output=True, check=True)
+            coverxygen_cmd = cmd[:-1]
+            break
+        except (FileNotFoundError, subprocess.CalledProcessError):
+            continue
+    if coverxygen_cmd is None:
+        lines.append("coverxygen not installed. Install with: pip install coverxygen")
+        lines.append("")
+        lines.append("Once installed, coverxygen analyzes Doxygen XML to report what")
+        lines.append("percentage of C++ symbols have doc comments in the source code.")
+        lines.append("This is complementary to the RST coverage check above.")
+        lines.append("")
+        lines.append("Usage:")
+        lines.append(
+            f"  coverxygen --xml-dir {xml_dir} --src-dir ../../ --output coverxygen.info"
+        )
+        lines.append("  # Then use lcov/genhtml to visualize:")
+        lines.append(
+            "  genhtml --no-function-coverage coverxygen.info -o coverxygen_html"
+        )
+        return "\n".join(lines)
+
+    try:
+        result = subprocess.run(
+            coverxygen_cmd
+            + [
+                "--xml-dir",
+                str(xml_dir),
+                "--src-dir",
+                str(SCRIPT_DIR / ".." / ".."),
+                "--output",
+                "-",
+                "--kind",
+                "class,struct,function",
+                "--scope",
+                "public",
+                "--exclude",
+                ".*/build/.*",
+                "--exclude",
+                ".*/detail/.*",
+                "--exclude",
+                ".*/nativert/.*",
+                "--exclude",
+                ".*/stable/library\\.h",
+            ],
+            capture_output=True,
+            text=True,
+            timeout=120,
+        )
+        if result.returncode == 0:
+            total = 0
+            documented_count = 0
+            for line in result.stdout.splitlines():
+                if line.startswith("DA:"):
+                    total += 1
+                    parts = line.split(",")
+                    if len(parts) >= 2 and parts[1].strip() != "0":
+                        documented_count += 1
+            pct = (documented_count / total * 100) if total else 0
+            lines.append(f"Symbols scanned:    {total}")
+            lines.append(f"With doc comments:  {documented_count}")
+            lines.append(f"Coverage:           {pct:.1f}%")
+            lines.append("")
+            lines.append("Full lcov output saved to: coverxygen.info")
+            (SCRIPT_DIR / "coverxygen.info").write_text(result.stdout)
+        else:
+            lines.append(f"coverxygen failed (exit {result.returncode}):")
+            lines.append(result.stderr[:500])
+    except subprocess.TimeoutExpired:
+        lines.append("coverxygen timed out after 120s")
+    except Exception as e:
+        lines.append(f"coverxygen error: {e}")
+
+    lines.append("")
+    return "\n".join(lines)
+
+
+# ─── Main ────────────────────────────────────────────────────────────────────
+
+
+def main():
+    parser = argparse.ArgumentParser(description="C++ docs coverage checker")
+    parser.add_argument(
+        "--coverxygen",
+        action="store_true",
+        help="Also run coverxygen on Doxygen XML for doc-comment coverage",
+    )
+    parser.add_argument(
+        "--html-only",
+        action="store_true",
+        help="Only run HTML formatting checks",
+    )
+    args = parser.parse_args()
+
+    reports = []
+
+    if not args.html_only:
+        # Phase 1: Discover APIs from Doxygen XML
+        print("Discovering APIs from Doxygen XML...")
+        apis = discover_apis_from_xml(BUILD_XML)
+        total_apis = sum(len(v) for v in apis.values())
+        print(f"  Found {total_apis} public APIs across {len(apis)} categories")
+
+        # Phase 2: Scan RST for documented symbols
+        print("Scanning sources for breathe directives...")
+        documented = scan_sources(SOURCE_DIR)
+        print(f"  Found {len(documented)} documented symbols")
+
+        coverage_report = generate_coverage_report(apis, documented)
+        reports.append(coverage_report)
+
+        COVERAGE_OUTPUT.write_text(coverage_report)
+        print(f"  Coverage report written to: {COVERAGE_OUTPUT}")
+
+    # Phase 3: HTML checks
+    print("Checking HTML output for formatting issues...")
+    html_report = check_html_output(BUILD_HTML)
+    reports.append(html_report)
+    HTML_REPORT.write_text(html_report)
+    print(f"  HTML report written to: {HTML_REPORT}")
+
+    # Phase 4: coverxygen (optional)
+    if args.coverxygen:
+        print("Running coverxygen...")
+        cov_report = run_coverxygen(BUILD_XML)
+        reports.append(cov_report)
+
+    # Print everything
+    print()
+    print("=" * 60)
+    for report in reports:
+        print(report)
+        print()
+
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())
diff --git a/docs/cpp/cpp_coverage.txt b/docs/cpp/cpp_coverage.txt
new file mode 100644
index 0000000000000..4fce09f90a646
--- /dev/null
+++ b/docs/cpp/cpp_coverage.txt
@@ -0,0 +1,292 @@
+Undocumented C++ objects
+==================================================
+
+torch (device)
+--------------
+   * torch::cuda::device_count  (function)
+   * torch::cuda::is_available  (function)
+   * torch::cuda::manual_seed  (function)
+   * torch::cuda::manual_seed_all  (function)
+   * torch::cuda::synchronize  (function)
+   * torch::mps::commit  (function)
+   * torch::mps::get_command_buffer  (function)
+   * torch::mps::get_dispatch_queue  (function)
+   * torch::mps::is_available  (function)
+   * torch::mps::manual_seed  (function)
+   * torch::mps::synchronize  (function)
+
+torch (macros)
+--------------
+   * TORCH_MODULE  (define)
+   * TORCH_MODULE_IMPL  (define)
+
+torch::data
+-----------
+   * torch::data::datasets::ChunkDatasetOptions  (struct)
+   * torch::data::datasets::TensorDataset  (struct)
+   * torch::data::make_data_loader  (function)
+   * torch::data::samplers::BatchSize  (struct)
+   * torch::data::samplers::CustomBatchRequest  (struct)
+
+torch::fft
+----------
+   * torch::fft::fft  (function)
+   * torch::fft::fft2  (function)
+   * torch::fft::fftfreq  (function)
+   * torch::fft::fftn  (function)
+   * torch::fft::fftshift  (function)
+   * torch::fft::hfft  (function)
+   * torch::fft::hfft2  (function)
+   * torch::fft::hfftn  (function)
+   * torch::fft::ifft  (function)
+   * torch::fft::ifft2  (function)
+   * torch::fft::ifftn  (function)
+   * torch::fft::ifftshift  (function)
+   * torch::fft::ihfft  (function)
+   * torch::fft::ihfft2  (function)
+   * torch::fft::ihfftn  (function)
+   * torch::fft::irfft  (function)
+   * torch::fft::irfft2  (function)
+   * torch::fft::irfftn  (function)
+   * torch::fft::rfft  (function)
+   * torch::fft::rfft2  (function)
+   * torch::fft::rfftfreq  (function)
+   * torch::fft::rfftn  (function)
+
+torch::nn (modules)
+-------------------
+   * torch::nn::AdaptiveAvgPoolOptions  (struct)
+   * torch::nn::AdaptiveMaxPoolOptions  (struct)
+   * torch::nn::AvgPoolOptions  (struct)
+   * torch::nn::BatchNormOptions  (struct)
+   * torch::nn::BilinearOptions  (struct)
+   * torch::nn::CELUOptions  (struct)
+   * torch::nn::ConstantPad1dImpl  (class)
+   * torch::nn::ConstantPad2dImpl  (class)
+   * torch::nn::ConstantPad3dImpl  (class)
+   * torch::nn::ConstantPadOptions  (struct)
+   * torch::nn::ConvOptions  (struct)
+   * torch::nn::ConvTransposeOptions  (struct)
+   * torch::nn::CosineSimilarityImpl  (class)
+   * torch::nn::CosineSimilarityOptions  (struct)
+   * torch::nn::DropoutOptions  (struct)
+   * torch::nn::ELUOptions  (struct)
+   * torch::nn::EmbeddingBagFromPretrainedOptions  (struct)
+   * torch::nn::EmbeddingBagOptions  (struct)
+   * torch::nn::EmbeddingFromPretrainedOptions  (struct)
+   * torch::nn::EmbeddingOptions  (struct)
+   * torch::nn::FlattenOptions  (struct)
+   * torch::nn::FoldImpl  (class)
+   * torch::nn::FractionalMaxPoolOptions  (struct)
+   * torch::nn::Functional  (class)
+   * torch::nn::GELUOptions  (struct)
+   * torch::nn::GLUOptions  (struct)
+   * torch::nn::HardshrinkOptions  (struct)
+   * torch::nn::HardtanhOptions  (struct)
+   * torch::nn::InstanceNormOptions  (struct)
+   * torch::nn::KLDivLossOptions  (struct)
+   * torch::nn::L1LossOptions  (struct)
+   * torch::nn::LPPoolOptions  (struct)
+   * torch::nn::LayerNormOptions  (struct)
+   * torch::nn::LeakyReLUOptions  (struct)
+   * torch::nn::LinearOptions  (struct)
+   * torch::nn::LogSoftmaxOptions  (struct)
+   * torch::nn::MaxPoolOptions  (struct)
+   * torch::nn::MaxUnpoolOptions  (struct)
+   * torch::nn::ModuleDictImpl  (class)
+   * torch::nn::ModuleListImpl  (class)
+   * torch::nn::MultiheadAttentionOptions  (struct)
+   * torch::nn::PReLUOptions  (struct)
+   * torch::nn::PairwiseDistanceImpl  (class)
+   * torch::nn::PairwiseDistanceOptions  (struct)
+   * torch::nn::PixelShuffleImpl  (struct)
+   * torch::nn::PixelUnshuffleImpl  (struct)
+   * torch::nn::RNNOptions  (struct)
+   * torch::nn::RReLUOptions  (struct)
+   * torch::nn::ReLU6Options  (struct)
+   * torch::nn::ReLUOptions  (struct)
+   * torch::nn::ReflectionPad1dImpl  (class)
+   * torch::nn::ReflectionPad2dImpl  (class)
+   * torch::nn::ReflectionPad3dImpl  (class)
+   * torch::nn::ReflectionPadOptions  (struct)
+   * torch::nn::ReplicationPad1dImpl  (class)
+   * torch::nn::ReplicationPad2dImpl  (class)
+   * torch::nn::ReplicationPad3dImpl  (class)
+   * torch::nn::ReplicationPadOptions  (struct)
+   * torch::nn::SELUOptions  (struct)
+   * torch::nn::SequentialImpl  (class)
+   * torch::nn::SoftmaxOptions  (struct)
+   * torch::nn::SoftminOptions  (struct)
+   * torch::nn::SoftplusOptions  (struct)
+   * torch::nn::SoftshrinkOptions  (struct)
+   * torch::nn::ThresholdOptions  (struct)
+   * torch::nn::TransformerDecoderOptions  (struct)
+   * torch::nn::TransformerEncoderLayerOptions  (struct)
+   * torch::nn::TransformerEncoderOptions  (struct)
+   * torch::nn::TransformerOptions  (struct)
+   * torch::nn::UnflattenOptions  (struct)
+   * torch::nn::UnfoldImpl  (class)
+   * torch::nn::UpsampleImpl  (class)
+   * torch::nn::ZeroPad1dImpl  (class)
+   * torch::nn::ZeroPad2dImpl  (class)
+   * torch::nn::ZeroPad3dImpl  (class)
+   * torch::nn::ZeroPadOptions  (struct)
+
+torch::nn::functional
+---------------------
+   * torch::nn::functional::AlphaDropoutFuncOptions  (struct)
+   * torch::nn::functional::BatchNormFuncOptions  (struct)
+   * torch::nn::functional::ConvFuncOptions  (struct)
+   * torch::nn::functional::ConvTransposeFuncOptions  (struct)
+   * torch::nn::functional::DropoutFuncOptions  (struct)
+   * torch::nn::functional::EmbeddingBagFuncOptions  (struct)
+   * torch::nn::functional::EmbeddingFuncOptions  (struct)
+   * torch::nn::functional::FeatureAlphaDropoutFuncOptions  (struct)
+   * torch::nn::functional::GridSampleFuncOptions  (struct)
+   * torch::nn::functional::GumbelSoftmaxFuncOptions  (struct)
+   * torch::nn::functional::InstanceNormFuncOptions  (struct)
+   * torch::nn::functional::InterpolateFuncOptions  (struct)
+   * torch::nn::functional::LogSoftmaxFuncOptions  (struct)
+   * torch::nn::functional::MaxUnpoolFuncOptions  (struct)
+   * torch::nn::functional::MultiheadAttentionForwardFuncOptions  (struct)
+   * torch::nn::functional::PadFuncOptions  (struct)
+   * torch::nn::functional::RReLUFuncOptions  (struct)
+   * torch::nn::functional::SoftmaxFuncOptions  (struct)
+   * torch::nn::functional::SoftminFuncOptions  (struct)
+   * torch::nn::functional::adaptive_max_pool2d_with_indices  (function)
+   * torch::nn::functional::adaptive_max_pool3d_with_indices  (function)
+   * torch::nn::functional::fractional_max_pool2d_with_indices  (function)
+   * torch::nn::functional::fractional_max_pool3d_with_indices  (function)
+   * torch::nn::functional::multi_head_attention_forward  (function)
+
+torch::nn::init
+---------------
+   * torch::nn::init::calculate_gain  (function)
+   * torch::nn::init::constant_  (function)
+   * torch::nn::init::dirac_  (function)
+   * torch::nn::init::eye_  (function)
+   * torch::nn::init::kaiming_normal_  (function)
+   * torch::nn::init::kaiming_uniform_  (function)
+   * torch::nn::init::normal_  (function)
+   * torch::nn::init::ones_  (function)
+   * torch::nn::init::orthogonal_  (function)
+   * torch::nn::init::sparse_  (function)
+   * torch::nn::init::uniform_  (function)
+   * torch::nn::init::xavier_normal_  (function)
+   * torch::nn::init::xavier_uniform_  (function)
+   * torch::nn::init::zeros_  (function)
+
+torch::nn::utils
+----------------
+   * torch::nn::utils::clip_grad_norm_  (function)
+   * torch::nn::utils::clip_grad_value_  (function)
+   * torch::nn::utils::parameters_to_vector  (function)
+   * torch::nn::utils::rnn::invert_permutation  (function)
+   * torch::nn::utils::vector_to_parameters  (function)
+
+torch::optim
+------------
+   * torch::optim::AdagradOptions  (struct)
+   * torch::optim::AdagradParamState  (struct)
+   * torch::optim::AdamOptions  (struct)
+   * torch::optim::AdamParamState  (struct)
+   * torch::optim::AdamWOptions  (struct)
+   * torch::optim::AdamWParamState  (struct)
+   * torch::optim::LBFGSOptions  (struct)
+   * torch::optim::LBFGSParamState  (struct)
+   * torch::optim::RMSpropOptions  (struct)
+   * torch::optim::RMSpropParamState  (struct)
+   * torch::optim::SGDOptions  (struct)
+   * torch::optim::SGDParamState  (struct)
+
+torch::special
+--------------
+   * torch::special::airy_ai  (function)
+   * torch::special::bessel_j0  (function)
+   * torch::special::bessel_j1  (function)
+   * torch::special::bessel_y0  (function)
+   * torch::special::bessel_y1  (function)
+   * torch::special::chebyshev_polynomial_t  (function)
+   * torch::special::chebyshev_polynomial_u  (function)
+   * torch::special::chebyshev_polynomial_v  (function)
+   * torch::special::chebyshev_polynomial_w  (function)
+   * torch::special::digamma  (function)
+   * torch::special::entr  (function)
+   * torch::special::erf  (function)
+   * torch::special::erfc  (function)
+   * torch::special::erfcx  (function)
+   * torch::special::erfinv  (function)
+   * torch::special::exp2  (function)
+   * torch::special::expit  (function)
+   * torch::special::expm1  (function)
+   * torch::special::gammainc  (function)
+   * torch::special::gammaincc  (function)
+   * torch::special::gammaln  (function)
+   * torch::special::hermite_polynomial_h  (function)
+   * torch::special::hermite_polynomial_he  (function)
+   * torch::special::i0  (function)
+   * torch::special::i0e  (function)
+   * torch::special::i1  (function)
+   * torch::special::i1e  (function)
+   * torch::special::laguerre_polynomial_l  (function)
+   * torch::special::legendre_polynomial_p  (function)
+   * torch::special::log1p  (function)
+   * torch::special::log_ndtr  (function)
+   * torch::special::log_softmax  (function)
+   * torch::special::logit  (function)
+   * torch::special::logsumexp  (function)
+   * torch::special::modified_bessel_i0  (function)
+   * torch::special::modified_bessel_i1  (function)
+   * torch::special::modified_bessel_k0  (function)
+   * torch::special::modified_bessel_k1  (function)
+   * torch::special::multigammaln  (function)
+   * torch::special::ndtr  (function)
+   * torch::special::ndtri  (function)
+   * torch::special::polygamma  (function)
+   * torch::special::psi  (function)
+   * torch::special::round  (function)
+   * torch::special::scaled_modified_bessel_k0  (function)
+   * torch::special::scaled_modified_bessel_k1  (function)
+   * torch::special::shifted_chebyshev_polynomial_t  (function)
+   * torch::special::shifted_chebyshev_polynomial_u  (function)
+   * torch::special::shifted_chebyshev_polynomial_v  (function)
+   * torch::special::shifted_chebyshev_polynomial_w  (function)
+   * torch::special::sinc  (function)
+   * torch::special::softmax  (function)
+   * torch::special::spherical_bessel_j0  (function)
+   * torch::special::xlog1py  (function)
+   * torch::special::xlogy  (function)
+   * torch::special::zeta  (function)
+
+
+==================================================
+Summary
+==================================================
+
+Total APIs discovered:   699
+Documented:              472
+Missing:                 227
+Coverage:                67.5%
+
+Category                                       Covered  Total       %
+----------------------------------------------------------------------
+at                                                   4      4  100.0%
+at::cuda                                             7      7  100.0%
+c10                                                 10     10  100.0%
+c10::cuda                                            9      9  100.0%
+c10::xpu                                             4      4  100.0%
+other                                                1      1  100.0%
+torch (core)                                         4      4  100.0%
+torch (device)                                       5     16   31.2%
+torch (macros)                                       3      5   60.0%
+torch::autograd                                      3      3  100.0%
+torch::data                                         29     34   85.3%
+torch::fft                                           0     22    0.0%
+torch::nn (modules)                                243    319   76.2%
+torch::nn::functional                              100    124   80.6%
+torch::nn::init                                      0     14    0.0%
+torch::nn::utils                                     0      5    0.0%
+torch::optim                                        13     25   52.0%
+torch::serialize                                     4      4  100.0%
+torch::special                                       0     56    0.0%
+torch::stable                                       33     33  100.0%
diff --git a/docs/cpp/cpp_html_issues.txt b/docs/cpp/cpp_html_issues.txt
new file mode 100644
index 0000000000000..754a4bd2bc4c8
--- /dev/null
+++ b/docs/cpp/cpp_html_issues.txt
@@ -0,0 +1,4 @@
+HTML Formatting Check
+==================================================
+
+No issues found.
diff --git a/docs/cpp/source/Doxyfile b/docs/cpp/source/Doxyfile
index 72e83733b6b4c..748f7a69ea266 100644
--- a/docs/cpp/source/Doxyfile
+++ b/docs/cpp/source/Doxyfile
@@ -34,6 +34,7 @@ INPUT                  = ../../../aten/src/ATen/ATen.h \
                          ../../../aten/src/ATen/core/ivalue.h \
                          ../../../aten/src/ATen/core/ScalarType.h \
                          ../../../aten/src/ATen/cuda/CUDAContext.h \
+                         ../../../aten/src/ATen/cuda/CUDAContextLight.h \
                          ../../../aten/src/ATen/cudnn/Descriptors.h \
                          ../../../aten/src/ATen/cudnn/Handles.h \
                          ../../../aten/src/ATen/cudnn/Types.h \
@@ -48,7 +49,9 @@ INPUT                  = ../../../aten/src/ATen/ATen.h \
                          ../../../build/aten/src/ATen/Functions.h \
                          ../../../build/aten/src/ATen/core/TensorBody.h \
                          ../../../c10/core/Device.h \
+                         ../../../c10/core/DeviceGuard.h \
                          ../../../c10/core/DeviceType.h \
+                         ../../../c10/core/Stream.h \
                          ../../../c10/util/Half.h \
                          ../../../c10/util/ArrayRef.h \
                          ../../../c10/util/OptionalArrayRef.h \
@@ -100,7 +103,8 @@ GENERATE_XML           = YES
 # Set to NO if you do not want the Doxygen program listing included.
 # NOTE: setting to NO may result in unrecovered file relationships
 #       (which file defined which compound)
-XML_PROGRAMLISTING     = YES
+# Setting to NO to prevent inline code from polluting documentation
+XML_PROGRAMLISTING     = NO
 ################################################################################
 # Doxygen preprocessor / parser control.                                       #
 ################################################################################
@@ -129,8 +133,16 @@ EXTRACT_ALL            = YES
 EXTRACT_PACKAGE        = YES
 EXTRACT_STATIC         = YES
 CASE_SENSE_NAMES       = NO
-EXCLUDE_SYMBOLS        = caffe2::* cereal* DL* TH* cudnn* std::*
+EXCLUDE_SYMBOLS        = caffe2::* cereal* DL* TH* cudnn* std::* torch::nested*
 # EXCLUDE_SYMBOLS        = c10::* caffe2::* cereal* DL* TH* cudnn* std::*
+
+# Don't include inline source code in documentation
+INLINE_SOURCES         = NO
+# Don't include code bodies in detailed description
+SOURCE_BROWSER         = NO
+# Don't reference source files
+REFERENCES_LINK_SOURCE = NO
+VERBATIM_HEADERS       = NO
 ################################################################################
 # Docstring control / customization.                                           #
 ################################################################################
@@ -148,10 +160,21 @@ JAVADOC_AUTOBRIEF      = YES
 #      *
 #      * \endrst
 #      */
+# Or with /// style:
+#     /// \rst
+#     /// .. code-block:: cpp
+#     ///
+#     ///     int main() {
+#     ///         return 0;
+#     ///     }
+#     ///
+#     /// \endrst
+#
 # NOTE:
 #     1. \rst and \endrst must be on their own line.
-#     2. leading-asterisk required.
-ALIASES                = "rst=\verbatim embed:rst:leading-asterisk"
+#     2. For /// comments, use leading-slashes; for /** */ use leading-asterisk
+#     3. PyTorch uses /// style comments primarily
+ALIASES                = "rst=\verbatim embed:rst:leading-slashes"
 ALIASES               += "endrst=\endverbatim"
 ################################################################################
 # Warning suppression.                                                         #
diff --git a/docs/cpp/source/api/aten/accessors.md b/docs/cpp/source/api/aten/accessors.md
new file mode 100644
index 0000000000000..83fefce724e7e
--- /dev/null
+++ b/docs/cpp/source/api/aten/accessors.md
@@ -0,0 +1,49 @@
+---
+myst:
+  html_meta:
+    description: Tensor accessors in PyTorch C++ — efficient element-wise access to tensor data without overhead.
+    keywords: PyTorch, C++, tensor accessor, packed_accessor, data access
+---
+
+# Tensor Accessors
+
+For element-wise operations in custom kernels, use *accessors* to avoid
+dynamic dispatch overhead.
+
+## CPU Accessors
+
+```cpp
+torch::Tensor foo = torch::rand({12, 12});
+
+// Create accessor - validates type and dimensions once
+auto foo_a = foo.accessor<float, 2>();
+
+float trace = 0;
+for (int i = 0; i < foo_a.size(0); i++) {
+    trace += foo_a[i][i];
+}
+```
+
+## CUDA Packed Accessors
+
+For CUDA kernels, use *packed accessors* which copy metadata instead of
+pointing to it:
+
+```cpp
+__global__ void kernel(torch::PackedTensorAccessor64<float, 2> foo, float* trace) {
+    int i = threadIdx.x;
+    gpuAtomicAdd(trace, foo[i][i]);
+}
+
+torch::Tensor foo = torch::rand({12, 12}).cuda();
+auto foo_a = foo.packed_accessor64<float, 2>();
+
+float trace = 0;
+kernel<<<1, 12>>>(foo_a, &trace);
+```
+
+```{tip}
+
+Use `PackedTensorAccessor32` and `packed_accessor32` for 32-bit indexing,
+which is faster on CUDA but may overflow for large tensors.
+```
diff --git a/docs/cpp/source/api/aten/creation.md b/docs/cpp/source/api/aten/creation.md
new file mode 100644
index 0000000000000..27ddcd616db55
--- /dev/null
+++ b/docs/cpp/source/api/aten/creation.md
@@ -0,0 +1,235 @@
+---
+myst:
+  html_meta:
+    description: Tensor creation functions in PyTorch C++ — zeros, ones, randn, arange, from_blob, and more.
+    keywords: PyTorch, C++, tensor creation, zeros, ones, randn, arange, from_blob
+---
+
+# Tensor Creation
+
+Factory functions create new tensors with different initialization patterns.
+All factory functions follow a general schema:
+
+```cpp
+torch::<function-name>(<function-specific-options>, <sizes>, <tensor-options>)
+```
+
+## Available Factory Functions
+
+- `torch::zeros` - Tensor filled with zeros
+- `torch::ones` - Tensor filled with ones
+- `torch::empty` - Uninitialized tensor
+- `torch::full` - Tensor filled with a single value
+- `torch::rand` - Uniform random values on [0, 1)
+- `torch::randn` - Standard normal distribution
+- `torch::randint` - Random integers in a range
+- `torch::arange` - Sequence of integers
+- `torch::linspace` - Linearly spaced values
+- `torch::logspace` - Logarithmically spaced values
+- `torch::eye` - Identity matrix
+- `torch::randperm` - Random permutation of integers
+
+## Specifying a Size
+
+Functions that do not require specific arguments can be invoked with just a
+size. For example, the following line creates a vector with 5 components:
+
+```cpp
+torch::Tensor tensor = torch::ones(5);
+```
+
+An `IntArrayRef` is constructed by specifying the size along each dimension in
+curly braces. For example, `{2, 3}` for a matrix with two rows and three
+columns, `{3, 4, 5}` for a three-dimensional tensor:
+
+```cpp
+torch::Tensor tensor = torch::randn({3, 4, 5});
+assert(tensor.sizes() == std::vector<int64_t>{3, 4, 5});
+```
+
+You can also pass an `std::vector<int64_t>` instead of curly braces.
+Use `tensor.size(i)` to access a single dimension.
+
+## Passing Function-Specific Parameters
+
+Some factory functions accept additional parameters. For example, `randint`
+takes an upper bound on the value for the integers it generates:
+
+```cpp
+torch::Tensor tensor = torch::randint(/*high=*/10, {5, 5});
+
+// With a lower bound
+torch::Tensor tensor = torch::randint(/*low=*/3, /*high=*/10, {5, 5});
+```
+
+```{tip}
+
+The size always follows the function-specific arguments.
+```
+
+```{attention}
+
+Some functions like `arange` do not need a size at all, since it is fully
+determined by the function-specific arguments (the range bounds).
+```
+
+## Configuring Properties with TensorOptions
+
+`TensorOptions` configures the data type, layout, device, and
+`requires_grad` of a new tensor. The construction axes are:
+
+- `dtype`: the data type of the elements (e.g. `kFloat32`, `kInt64`)
+- `layout`: either `kStrided` (dense) or `kSparse`
+- `device`: a compute device (e.g. `kCPU`, `kCUDA`)
+- `requires_grad`: whether to track gradients
+
+Allowed values:
+
+- `dtype`: `kUInt8`, `kInt8`, `kInt16`, `kInt32`, `kInt64`,
+  `kFloat32`, `kFloat64`
+- `layout`: `kStrided`, `kSparse`
+- `device`: `kCPU`, or `kCUDA` (with an optional device index)
+- `requires_grad`: `true` or `false`
+
+```{tip}
+
+Rust-style shorthands exist for dtypes, like `kF32` instead of
+`kFloat32`. See
+[torch/types.h](https://github.com/pytorch/pytorch/blob/main/torch/csrc/api/include/torch/types.h)
+for the full list.
+```
+
+Here is an example of creating a `TensorOptions` object:
+
+```cpp
+auto options =
+  torch::TensorOptions()
+    .dtype(torch::kFloat32)
+    .layout(torch::kStrided)
+    .device(torch::kCUDA, 1)
+    .requires_grad(true);
+
+torch::Tensor tensor = torch::full({3, 4}, /*value=*/123, options);
+
+assert(tensor.dtype() == torch::kFloat32);
+assert(tensor.layout() == torch::kStrided);
+assert(tensor.device().type() == torch::kCUDA);
+assert(tensor.device().index() == 1);
+assert(tensor.requires_grad());
+```
+
+**Defaults:** Any axis you omit takes its default value: `kFloat32` for dtype,
+`kStrided` for layout, `kCPU` for device, and `false` for
+`requires_grad`. This means you can omit `TensorOptions` entirely:
+
+```cpp
+// A 32-bit float, strided, CPU tensor that does not require a gradient.
+torch::Tensor tensor = torch::randn({3, 4});
+```
+
+**Shorthand syntax:** For each axis there is a free function in the `torch::`
+namespace (`torch::dtype()`, `torch::device()`, `torch::layout()`,
+`torch::requires_grad()`). Each returns a `TensorOptions` object that can
+be further refined with builder methods:
+
+```cpp
+// These are equivalent:
+torch::ones(10, torch::TensorOptions().dtype(torch::kFloat32))
+torch::ones(10, torch::dtype(torch::kFloat32))
+
+// Chaining:
+torch::ones(10, torch::dtype(torch::kFloat32).layout(torch::kStrided))
+```
+
+**Implicit construction:** `TensorOptions` is implicitly constructible from
+individual values, so when only one axis differs from the default you can
+write:
+
+```cpp
+torch::ones(10, torch::kFloat32)
+```
+
+Putting it all together, a C++ tensor creation call mirrors the Python
+equivalent closely:
+
+```python
+# Python
+torch.randn(3, 4, dtype=torch.float32, device=torch.device('cuda', 1), requires_grad=True)
+```
+
+```cpp
+// C++
+torch::randn({3, 4}, torch::dtype(torch::kFloat32).device(torch::kCUDA, 1).requires_grad(true))
+```
+
+## Using Externally Created Data
+
+If you already have tensor data allocated in memory (CPU or CUDA), use
+`from_blob` to view that memory as a `Tensor`:
+
+```cpp
+float data[] = {1, 2, 3, 4, 5, 6};
+torch::Tensor tensor = torch::from_blob(data, {2, 3});
+```
+
+```{note}
+
+Tensors created with `from_blob` cannot be resized because ATen does not
+own the memory.
+```
+
+## Tensor Conversion
+
+Use `to()` to convert tensors between dtypes and devices. The conversion
+creates a new tensor and does not occur in-place:
+
+```cpp
+torch::Tensor source = torch::randn({2, 3}, torch::kInt64);
+
+// Convert dtype
+torch::Tensor float_tensor = source.to(torch::kFloat32);
+
+// Move to GPU (default CUDA device)
+torch::Tensor gpu_tensor = float_tensor.to(torch::kCUDA);
+
+// Specific GPU device
+torch::Tensor gpu1_tensor = float_tensor.to(torch::Device(torch::kCUDA, 1));
+
+// Async copy
+torch::Tensor async_tensor = gpu_tensor.to(torch::kCPU, /*non_blocking=*/true);
+```
+
+```{attention}
+
+The result of the conversion is a new tensor pointing to new memory,
+unrelated to the source tensor.
+```
+
+## Scalars and Zero-Dimensional Tensors
+
+`Scalar` represents a single dynamically-typed number. Like a `Tensor`,
+`Scalar` is dynamically typed and can hold any of ATen's number types.
+Scalars can be implicitly constructed from C++ number types:
+
+```cpp
+namespace torch {
+Tensor addmm(Scalar beta, const Tensor & self,
+             Scalar alpha, const Tensor & mat1,
+             const Tensor & mat2);
+Scalar sum(const Tensor & self);
+} // namespace torch
+
+// Usage
+torch::Tensor a = ...;
+torch::Tensor b = ...;
+torch::Tensor c = ...;
+torch::Tensor r = torch::addmm(1.0, a, .5, b, c);
+```
+
+Zero-dimensional tensors hold a single value and can reference elements in
+larger tensors:
+
+```cpp
+torch::Tensor matrix = torch::rand({10, 20});
+matrix[1][2] = 4;  // matrix[1][2] is a zero-dimensional tensor
+```
diff --git a/docs/cpp/source/api/aten/index.md b/docs/cpp/source/api/aten/index.md
new file mode 100644
index 0000000000000..e9980c4cca39f
--- /dev/null
+++ b/docs/cpp/source/api/aten/index.md
@@ -0,0 +1,76 @@
+---
+myst:
+  html_meta:
+    description: ATen C++ API — the foundational tensor library powering PyTorch, including tensor operations, types, and CUDA support.
+    keywords: PyTorch, C++, ATen, tensor, operations, CUDA
+---
+
+# ATen: Tensor Library
+
+ATen (A Tensor Library) is the foundational tensor and mathematical operation
+library on which all of PyTorch is built. It provides the core `Tensor` class
+and hundreds of mathematical operations that work on tensors.
+
+**When to use ATen directly:**
+
+- When writing low-level tensor operations or custom kernels
+- When you need direct access to tensor data and metadata
+- When working with the PyTorch internals or extending PyTorch
+
+**Basic usage:**
+
+```cpp
+#include <ATen/ATen.h>
+
+// Create tensors
+at::Tensor a = at::ones({2, 3});
+at::Tensor b = at::randn({2, 3});
+
+// Perform operations
+at::Tensor c = a + b;
+at::Tensor d = at::matmul(a.t(), b);
+
+// Move to GPU
+if (at::cuda::is_available()) {
+    at::Tensor gpu_tensor = c.to(at::kCUDA);
+}
+```
+
+For most applications, prefer using the higher-level `torch::` namespace
+(see {doc}`../nn/index`, {doc}`../optim/index`) which provides a more user-friendly API.
+
+## Header Files
+
+The following headers are part of the ATen public API:
+
+- `ATen/ATen.h` - Main ATen header
+- `ATen/Backend.h` - Backend enumeration
+- `ATen/core/Tensor.h` - Tensor class
+- `ATen/core/ivalue.h` - IValue type (see {doc}`../c10/index`)
+- `ATen/core/ScalarType.h` - Data type definitions
+- `ATen/TensorOptions.h` - Tensor creation options
+- `ATen/Scalar.h` - Scalar type
+- `ATen/Layout.h` - Tensor layout
+- `ATen/DeviceGuard.h` - Device context management
+- `ATen/native/TensorShape.h` - Tensor shape operations
+- `ATen/cuda/CUDAContext.h` - CUDA context (see {doc}`../cuda/index`)
+- `ATen/cudnn/Descriptors.h` - cuDNN descriptors
+- `ATen/mkl/Descriptors.h` - MKL descriptors
+
+```{note}
+
+The core `at::Tensor` class is defined in a generated header file
+(`TensorBody.h`) that only exists after building PyTorch. The documentation
+below describes the API manually.
+```
+
+## ATen Categories
+
+```{toctree}
+:maxdepth: 1
+
+tensor
+creation
+indexing
+accessors
+```
diff --git a/docs/cpp/source/api/aten/indexing.md b/docs/cpp/source/api/aten/indexing.md
new file mode 100644
index 0000000000000..63481cb1cd5a4
--- /dev/null
+++ b/docs/cpp/source/api/aten/indexing.md
@@ -0,0 +1,193 @@
+---
+myst:
+  html_meta:
+    description: Tensor indexing in PyTorch C++ — Slice, None, Ellipsis, boolean masks, and advanced indexing with index() and index_put_().
+    keywords: PyTorch, C++, tensor indexing, Slice, index, index_put_, boolean mask, fancy indexing
+---
+
+# Tensor Indexing
+
+The PyTorch C++ API provides tensor indexing similar to Python. Use
+`torch::indexing` namespace for index types:
+
+```cpp
+using namespace torch::indexing;
+```
+
+The main difference from Python is that instead of using the `[]` operator,
+the C++ API uses the `index` and `index_put_` methods:
+
+- `torch::Tensor::index` — read elements
+- `torch::Tensor::index_put_` — write elements
+
+## Index Types
+
+The `TensorIndex` class accepts six types of indices via implicit constructors:
+
+```{list-table}
+:widths: 25 35 40
+:header-rows: 1
+
+* - Type
+  - C++
+  - Python equivalent
+* - None (unsqueeze)
+  - `None`
+  - `None`
+* - Ellipsis
+  - `Ellipsis` or `"..."`
+  - `...`
+* - Integer
+  - `0`, `1`, `-1`
+  - `0`, `1`, `-1`
+* - Boolean
+  - `true`, `false`
+  - `True`, `False`
+* - Slice
+  - `Slice(start, stop, step)`
+  - `start:stop:step`
+* - Tensor
+  - `torch::tensor({0, 2})`
+  - `torch.tensor([0, 2])`
+```
+
+## Getter Operations
+
+```{list-table}
+:widths: 40 60
+:header-rows: 1
+
+* - Python
+  - C++
+* - `tensor[0]`
+  - `tensor.index({0})`
+* - `tensor[-1]`
+  - `tensor.index({-1})`
+* - `tensor[1, 2]`
+  - `tensor.index({1, 2})`
+* - `tensor[1, :, 3]`
+  - `tensor.index({1, Slice(), 3})`
+* - `tensor[None]`
+  - `tensor.index({None})`
+* - `tensor[:, None]`
+  - `tensor.index({Slice(), None})`
+* - `tensor[...]`
+  - `tensor.index({Ellipsis})` or `tensor.index({"..."})`
+* - `tensor[..., 0]`
+  - `tensor.index({Ellipsis, 0})`
+* - `tensor[1::2]`
+  - `tensor.index({Slice(1, None, 2)})`
+* - `tensor[True]`
+  - `tensor.index({true})`
+* - `tensor[torch.tensor([1, 2])]`
+  - `tensor.index({torch::tensor({1, 2})})`
+* - `tensor[bool_mask]`
+  - `tensor.index({bool_mask})`
+* - `tensor[:, torch.tensor([[0,1],[4,3]])]`
+  - `tensor.index({Slice(), torch::tensor({{0,1},{4,3}})})`
+* - `tensor[cond > 0]`
+  - `tensor.index({cond > 0})`
+```
+
+## Setter Operations
+
+```{list-table}
+:widths: 40 60
+:header-rows: 1
+
+* - Python
+  - C++
+* - `tensor[0] = 1`
+  - `tensor.index_put_({0}, 1)`
+* - `tensor[1, 2] = 1`
+  - `tensor.index_put_({1, 2}, 1)`
+* - `tensor[1] = torch.arange(5)`
+  - `tensor.index_put_({1}, torch::arange(5))`
+* - `tensor[1::2] = 1`
+  - `tensor.index_put_({Slice(1, None, 2)}, 1)`
+* - `tensor[0, 1::2] = torch.tensor([3., 4.])`
+  - `tensor.index_put_({0, Slice(1, None, 2)}, torch::tensor({3., 4.}))`
+* - `tensor[...] = 0`
+  - `tensor.index_put_({Ellipsis}, 0)`
+* - `tensor[None] = value`
+  - `tensor.index_put_({None}, value)`
+* - `tensor[bool_mask] = 0`
+  - `tensor.index_put_({bool_mask}, 0)`
+* - `tensor[torch.tensor([0, 2])] = value`
+  - `tensor.index_put_({torch::tensor({0, 2})}, value)`
+* - `tensor[1:2, torch.tensor([1,2])] = 0`
+  - `tensor.index_put_({Slice(1, 2), torch::tensor({1, 2})}, 0)`
+```
+
+The `index_put_` method also accepts an optional `accumulate` parameter.
+When `true`, values are added to existing values instead of replacing them:
+
+```cpp
+tensor.index_put_({mask}, values, /*accumulate=*/true);
+```
+
+## Slice Syntax
+
+The `Slice` constructor signature is:
+
+```cpp
+Slice(
+    std::optional<c10::SymInt> start = std::nullopt,
+    std::optional<c10::SymInt> stop  = std::nullopt,
+    std::optional<c10::SymInt> step  = std::nullopt);
+```
+
+Pass `None` for open-ended bounds:
+
+```{list-table}
+:widths: 30 70
+:header-rows: 1
+
+* - Python
+  - C++
+* - `:` or `::`
+  - `Slice()` or `Slice(None, None)`
+* - `1:`
+  - `Slice(1, None)`
+* - `:3`
+  - `Slice(None, 3)`
+* - `1:3`
+  - `Slice(1, 3)`
+* - `1:3:2`
+  - `Slice(1, 3, 2)`
+* - `::2`
+  - `Slice(None, None, 2)`
+```
+
+## Full Example
+
+```cpp
+#include <torch/torch.h>
+
+using namespace torch::indexing;
+
+auto tensor = torch::arange(2 * 3 * 4).reshape({2, 3, 4});
+
+// Basic indexing
+auto row = tensor.index({0});             // tensor[0]
+auto elem = tensor.index({1, 2, 3});      // tensor[1, 2, 3]
+
+// Slicing
+auto sliced = tensor.index({Slice(), Slice(0, 2)});  // tensor[:, 0:2]
+
+// None (unsqueeze) and Ellipsis
+auto unsqueezed = tensor.index({None});              // tensor[None]
+auto last_dim = tensor.index({Ellipsis, -1});         // tensor[..., -1]
+
+// Boolean mask indexing
+auto mask = tensor > 10;
+auto selected = tensor.index({mask});     // tensor[tensor > 10]
+
+// Integer tensor (fancy) indexing
+auto idx = torch::tensor({0, 2});
+auto gathered = tensor.index({Slice(), idx});  // tensor[:, [0, 2]]
+
+// Setting values
+tensor.index_put_({0, Slice(), 0}, 99);        // tensor[0, :, 0] = 99
+tensor.index_put_({mask}, 0);                  // tensor[tensor > 10] = 0
+```
diff --git a/docs/cpp/source/api/aten/tensor.md b/docs/cpp/source/api/aten/tensor.md
new file mode 100644
index 0000000000000..227c94e13d7ff
--- /dev/null
+++ b/docs/cpp/source/api/aten/tensor.md
@@ -0,0 +1,288 @@
+---
+myst:
+  html_meta:
+    description: at::Tensor class reference — the primary tensor type in PyTorch C++ with creation, indexing, device, and dtype APIs.
+    keywords: PyTorch, C++, Tensor, at::Tensor, TensorOptions, ScalarType, DeviceGuard
+---
+
+# Tensor Class
+
+The `at::Tensor` class is the primary tensor class in ATen, representing
+a multi-dimensional array with a specific data type and device.
+
+## Tensor
+
+```{cpp:class} at::Tensor
+
+The primary tensor class in ATen. Represents a multi-dimensional array
+with a specific data type and device.
+```
+
+```{cpp:function} Tensor()
+
+Default constructor. Creates an undefined tensor.
+```
+
+```{cpp:function} int64_t dim() const
+
+Returns the number of dimensions of the tensor.
+```
+
+```{cpp:function} int64_t size(int64_t dim) const
+
+Returns the size of the tensor at the given dimension.
+```
+
+```{cpp:function} IntArrayRef sizes() const
+
+Returns the sizes of all dimensions.
+```
+
+```{cpp:function} IntArrayRef strides() const
+
+Returns the strides of all dimensions.
+```
+
+```{cpp:function} ScalarType scalar_type() const
+
+Returns the data type of the tensor.
+```
+
+```{cpp:function} Device device() const
+
+Returns the device where the tensor is stored.
+```
+
+```{cpp:function} bool is_cuda() const
+
+Returns true if the tensor is on a CUDA device.
+```
+
+```{cpp:function} bool is_cpu() const
+
+Returns true if the tensor is on CPU.
+```
+
+```{cpp:function} bool requires_grad() const
+
+Returns true if gradients need to be computed for this tensor.
+```
+
+```{cpp:function} Tensor& requires_grad_(bool requires_grad = true)
+
+Sets whether gradients should be computed for this tensor.
+```
+
+```{cpp:function} Tensor to(Device device) const
+
+Returns a tensor on the specified device.
+```
+
+```{cpp:function} Tensor to(ScalarType dtype) const
+
+Returns a tensor with the specified data type.
+```
+
+```{cpp:function} Tensor contiguous() const
+
+Returns a contiguous tensor with the same data.
+```
+
+```{cpp:function} void* data_ptr() const
+
+Returns a pointer to the underlying data.
+```
+
+**Example:**
+
+```cpp
+#include <ATen/ATen.h>
+
+at::Tensor a = at::ones({2, 2}, at::kInt);
+at::Tensor b = at::randn({2, 2});
+auto c = a + b.to(at::kInt);
+```
+
+## TensorOptions
+
+```{cpp:class} at::TensorOptions
+
+A class to specify options for tensor creation, including dtype, device,
+layout, and requires_grad.
+```
+
+```{cpp:function} TensorOptions()
+
+Default constructor.
+```
+
+```{cpp:function} TensorOptions dtype(ScalarType dtype) const
+
+Returns options with the specified data type.
+```
+
+```{cpp:function} TensorOptions device(Device device) const
+
+Returns options with the specified device.
+```
+
+```{cpp:function} TensorOptions layout(Layout layout) const
+
+Returns options with the specified layout.
+```
+
+```{cpp:function} TensorOptions requires_grad(bool requires_grad) const
+
+Returns options with the specified requires_grad setting.
+```
+
+**Example:**
+
+```cpp
+auto options = at::TensorOptions()
+    .dtype(at::kFloat)
+    .device(at::kCUDA, 0)
+    .requires_grad(true);
+
+at::Tensor tensor = at::zeros({3, 4}, options);
+```
+
+## Scalar
+
+```{cpp:class} at::Scalar
+
+Represents a scalar value that can be converted to various numeric types.
+```
+
+```{cpp:function} Scalar(int64_t v)
+
+Construct from an integer.
+```
+
+```{cpp:function} Scalar(double v)
+
+Construct from a double.
+```
+
+```{cpp:function} template<typename T> T to() const
+
+Convert to the specified type.
+```
+
+```{cpp:function} bool isIntegral(bool includeBool = false) const
+
+Returns true if the scalar is an integral type.
+```
+
+```{cpp:function} bool isFloatingPoint() const
+
+Returns true if the scalar is a floating point type.
+```
+
+## ScalarType
+
+```{cpp:enum-class} at::ScalarType
+
+Enumeration of data types supported by tensors.
+```
+
+```{cpp:enumerator} Byte
+
+8-bit unsigned integer (uint8_t)
+```
+
+```{cpp:enumerator} Char
+
+8-bit signed integer (int8_t)
+```
+
+```{cpp:enumerator} Short
+
+16-bit signed integer (int16_t)
+```
+
+```{cpp:enumerator} Int
+
+32-bit signed integer (int32_t)
+```
+
+```{cpp:enumerator} Long
+
+64-bit signed integer (int64_t)
+```
+
+```{cpp:enumerator} Half
+
+16-bit floating point (float16)
+```
+
+```{cpp:enumerator} Float
+
+32-bit floating point (float)
+```
+
+```{cpp:enumerator} Double
+
+64-bit floating point (double)
+```
+
+```{cpp:enumerator} Bool
+
+Boolean type
+```
+
+```{cpp:enumerator} BFloat16
+
+Brain floating point (bfloat16)
+```
+
+Convenience constants:
+
+- `at::kByte`, `at::kChar`, `at::kShort`, `at::kInt`, `at::kLong`
+- `at::kHalf`, `at::kFloat`, `at::kDouble`, `at::kBFloat16`
+- `at::kBool`
+
+## DeviceGuard
+
+```{doxygenclass} c10::DeviceGuard
+:members:
+:undoc-members:
+```
+
+**Example:**
+
+```cpp
+{
+    c10::DeviceGuard guard(at::Device(at::kCUDA, 1));
+    // Operations here run on CUDA device 1
+    auto tensor = at::zeros({2, 2});
+}
+// Previous device is restored
+```
+
+## Layout
+
+```{cpp:enum-class} at::Layout
+
+Enumeration of tensor memory layouts.
+```
+
+```{cpp:enumerator} Strided
+
+Dense tensor with strides.
+```
+
+```{cpp:enumerator} Sparse
+
+Sparse tensor (COO format).
+```
+
+```{cpp:enumerator} SparseCsr
+
+Sparse tensor in CSR format.
+```
+
+```{cpp:enumerator} SparseCsc
+
+Sparse tensor in CSC format.
+```
diff --git a/docs/cpp/source/api/autograd/custom_functions.md b/docs/cpp/source/api/autograd/custom_functions.md
new file mode 100644
index 0000000000000..475c975cdd9aa
--- /dev/null
+++ b/docs/cpp/source/api/autograd/custom_functions.md
@@ -0,0 +1,95 @@
+---
+myst:
+  html_meta:
+    description: Custom autograd functions in PyTorch C++ — defining forward and backward passes with torch::autograd::Function.
+    keywords: PyTorch, C++, autograd, custom function, forward, backward, Function
+---
+
+# Custom Autograd Functions
+
+PyTorch allows you to define custom autograd functions with custom forward
+and backward implementations.
+
+## Function Base Class
+
+```{doxygenstruct} torch::autograd::Function
+:members:
+:undoc-members:
+```
+
+## AutogradContext
+
+```{doxygenstruct} torch::autograd::AutogradContext
+:members:
+:undoc-members:
+```
+
+## Creating Custom Functions
+
+To create a custom autograd function, inherit from `torch::autograd::Function`
+and implement the static `forward` and `backward` methods:
+
+**Example:**
+
+```cpp
+class MyReLU : public torch::autograd::Function<MyReLU> {
+ public:
+  static torch::Tensor forward(
+      torch::autograd::AutogradContext* ctx,
+      torch::Tensor input) {
+    ctx->save_for_backward({input});
+    return input.clamp_min(0);
+  }
+
+  static torch::autograd::variable_list backward(
+      torch::autograd::AutogradContext* ctx,
+      torch::autograd::variable_list grad_outputs) {
+    auto saved = ctx->get_saved_variables();
+    auto input = saved[0];
+    auto grad_output = grad_outputs[0];
+    auto grad_input = grad_output * (input > 0).to(grad_output.dtype());
+    return {grad_input};
+  }
+};
+
+// Usage
+auto output = MyReLU::apply(input);
+```
+
+## Custom Kernels and AutoDispatchBelowADInplaceOrView
+
+For users implementing custom kernels who want to redispatch below `Autograd` dispatch
+keys, use `at::AutoDispatchBelowADInplaceOrView` instead of `InferenceMode`:
+
+```cpp
+class ROIAlignFunction : public torch::autograd::Function<ROIAlignFunction> {
+ public:
+  static torch::autograd::variable_list forward(
+      torch::autograd::AutogradContext* ctx,
+      const torch::autograd::Variable& input,
+      const torch::autograd::Variable& rois,
+      double spatial_scale,
+      int64_t pooled_height,
+      int64_t pooled_width,
+      int64_t sampling_ratio,
+      bool aligned) {
+    ctx->saved_data["spatial_scale"] = spatial_scale;
+    ctx->saved_data["pooled_height"] = pooled_height;
+    ctx->saved_data["pooled_width"] = pooled_width;
+    ctx->saved_data["sampling_ratio"] = sampling_ratio;
+    ctx->saved_data["aligned"] = aligned;
+    ctx->saved_data["input_shape"] = input.sizes();
+    ctx->save_for_backward({rois});
+
+    at::AutoDispatchBelowADInplaceOrView guard;
+    auto result = roi_align(
+        input, rois, spatial_scale, pooled_height,
+        pooled_width, sampling_ratio, aligned);
+    return {result};
+  }
+};
+```
+
+For customized inplace and view kernels, see the
+[custom kernel tutorial](https://pytorch.org/tutorials/advanced/cpp_extension.html#backward-pass)
+for more details.
diff --git a/docs/cpp/source/api/autograd/gradient.md b/docs/cpp/source/api/autograd/gradient.md
new file mode 100644
index 0000000000000..6f595496fb7a0
--- /dev/null
+++ b/docs/cpp/source/api/autograd/gradient.md
@@ -0,0 +1,86 @@
+---
+myst:
+  html_meta:
+    description: PyTorch C++ API for computing gradients — torch::autograd::backward and torch::autograd::grad functions for automatic differentiation.
+    keywords: PyTorch, C++, autograd, backward, grad, gradient, automatic differentiation
+---
+
+# Gradient Computation
+
+PyTorch provides functions for computing gradients of tensors with respect
+to graph leaves.
+
+## Gradient Functions
+
+```{cpp:function} void torch::autograd::backward(const variable_list& tensors, const variable_list& grad_tensors = {}, std::optional<bool> retain_graph = std::nullopt, bool create_graph = false, const variable_list& inputs = {})
+
+Computes the sum of gradients of given tensors with respect to graph leaves.
+
+The graph is differentiated using the chain rule. If any of `tensors`
+are non-scalar (i.e. their data has more than one element) and require
+gradient, then the Jacobian-vector product would be computed, in this case
+the function additionally requires specifying `grad_tensors`. It should be a
+sequence of matching length, that contains the "vector" in the
+Jacobian-vector product, usually the gradient of the differentiated function
+w.r.t. corresponding tensors (`torch::Tensor()` is an acceptable value for
+all tensors that don't need gradient tensors).
+
+This function accumulates gradients in the leaves — you might need to zero
+them before calling it.
+
+:param tensors: Tensors of which the derivative will be computed.
+:param grad_tensors: The "vector" in the Jacobian-vector product, usually
+   gradients w.r.t. each element of corresponding tensors.
+   `torch::Tensor()` values can be specified for scalar Tensors or ones
+   that don't require grad. If a `torch::Tensor()` value would be
+   acceptable for all grad_tensors, then this argument is optional.
+:param retain_graph: If `false`, the graph used to compute the grad will
+   be freed. Note that in nearly all cases setting this option to `true`
+   is not needed and often can be worked around in a much more efficient
+   way. Defaults to the value of `create_graph`.
+:param create_graph: If `true`, graph of the derivative will be
+   constructed, allowing to compute higher order derivative products.
+   Defaults to `false`.
+:param inputs: Inputs w.r.t. which the gradient will be accumulated into
+   `at::Tensor::grad`. All other Tensors will be ignored. If not
+   provided, the gradient is accumulated into all the leaf Tensors that
+   were used to compute `tensors`.
+```
+
+```{doxygenfunction} torch::autograd::grad
+```
+
+**Example:**
+
+```cpp
+#include <torch/torch.h>
+
+auto x = torch::randn({2, 2}, torch::requires_grad());
+auto y = x * x;
+auto z = y.sum();
+
+// Compute gradients
+z.backward();
+std::cout << x.grad() << std::endl;
+
+// Or use grad() for specific outputs
+auto grads = torch::autograd::grad({z}, {x});
+```
+
+## Tensor Gradient Methods
+
+Tensors have built-in methods for gradient computation:
+
+```cpp
+// Enable gradient tracking
+auto x = torch::randn({2, 2}).requires_grad_(true);
+
+// Check if gradient is required
+bool needs_grad = x.requires_grad();
+
+// Access the gradient after backward
+auto grad = x.grad();
+
+// Detach from computation graph
+auto x_detached = x.detach();
+```
diff --git a/docs/cpp/source/api/autograd/index.md b/docs/cpp/source/api/autograd/index.md
new file mode 100644
index 0000000000000..37b3cb79f747e
--- /dev/null
+++ b/docs/cpp/source/api/autograd/index.md
@@ -0,0 +1,56 @@
+---
+myst:
+  html_meta:
+    description: PyTorch C++ Autograd API — automatic differentiation for tensor computations.
+    keywords: PyTorch, C++, autograd, automatic differentiation, gradient
+---
+
+# Autograd: Automatic Differentiation
+
+PyTorch's autograd system provides automatic differentiation for all operations
+on tensors. It records operations on tensors to build a computational graph,
+then computes gradients automatically via backpropagation.
+
+**When to use Autograd:**
+
+- When training neural networks (gradients are computed automatically)
+- When implementing custom backward passes for specialized operations
+- When you need fine-grained control over gradient computation
+
+**Basic usage:**
+
+```cpp
+#include <torch/torch.h>
+
+// Create tensor with gradient tracking
+auto x = torch::randn({2, 2}, torch::requires_grad());
+auto y = x * x;
+auto z = y.sum();
+
+// Compute gradients via backpropagation
+z.backward();
+std::cout << x.grad() << std::endl;  // dz/dx = 2x
+
+// Disable gradient tracking for inference
+{
+    torch::NoGradGuard no_grad;
+    auto result = model->forward(input);  // No gradients computed
+}
+```
+
+## Header Files
+
+- `torch/csrc/autograd/autograd.h` - High-level autograd API
+- `torch/csrc/autograd/function.h` - Custom autograd functions
+- `torch/csrc/autograd/grad_mode.h` - Gradient computation modes
+- `torch/csrc/api/include/torch/autograd.h` - C++ Frontend autograd
+
+## Autograd Categories
+
+```{toctree}
+:maxdepth: 1
+
+gradient
+custom_functions
+modes
+```
diff --git a/docs/cpp/source/api/autograd/modes.md b/docs/cpp/source/api/autograd/modes.md
new file mode 100644
index 0000000000000..f1a281deae4f8
--- /dev/null
+++ b/docs/cpp/source/api/autograd/modes.md
@@ -0,0 +1,166 @@
+---
+myst:
+  html_meta:
+    description: Gradient mode guards in PyTorch C++ — NoGradGuard and InferenceMode for disabling gradient computation.
+    keywords: PyTorch, C++, NoGradGuard, InferenceMode, no_grad, inference, RAII guard
+---
+
+# Gradient Modes
+
+PyTorch provides RAII guards to control gradient computation behavior.
+
+## NoGradGuard
+
+```{cpp:class} torch::NoGradGuard
+
+RAII guard that disables gradient computation within its scope.
+```
+
+```{cpp:function} NoGradGuard()
+
+Disables gradient computation.
+```
+
+```{cpp:function} ~NoGradGuard()
+
+Restores previous gradient mode.
+```
+
+**Example:**
+
+```cpp
+{
+    torch::NoGradGuard no_grad;
+    // No gradients computed in this scope
+    auto result = model->forward(input);
+}
+```
+
+## InferenceMode
+
+`c10::InferenceMode` is a RAII guard analogous to `NoGradMode` designed for use
+when you are certain your operations will have no interactions with autograd
+(e.g., model inference). Compared to `NoGradMode`, code run under this mode gets
+better performance by disabling autograd-related work like view tracking and version
+counter bumps. However, tensors created inside `InferenceMode` have more limitations
+when interacting with the autograd system.
+
+```{cpp:class} c10::InferenceMode
+
+RAII guard that enables inference mode for optimized inference.
+This is more efficient than NoGradGuard for inference-only workloads.
+```
+
+```{cpp:function} explicit InferenceMode(bool enabled = true)
+
+Enables or disables inference mode.
+```
+
+**Inference Tensors:**
+
+`InferenceMode` can be enabled for a given block of code. Inside `InferenceMode`,
+all newly allocated (non-view) tensors are marked as inference tensors. Inference tensors:
+
+- Do not have a version counter, so an error will be raised if you try to read their version
+  (e.g., because you saved this tensor for backward).
+- Are immutable outside `InferenceMode`. An error will be raised if you try to:
+
+  - Mutate their data outside InferenceMode.
+  - Mutate them to `requires_grad=True` outside InferenceMode.
+  - To work around this, make a clone outside `InferenceMode` to get a normal tensor before mutating.
+
+A non-view tensor is an inference tensor if and only if it was allocated inside `InferenceMode`.
+A view tensor is an inference tensor if and only if it is a view of an inference tensor.
+
+**Performance Guarantees:**
+
+Inside an `InferenceMode` block:
+
+- Like `NoGradMode`, all operations do not record `grad_fn` even if their inputs have
+  `requires_grad=True`. This applies to both inference tensors and normal tensors.
+- View operations on inference tensors do not perform view tracking. View and non-view
+  inference tensors are indistinguishable.
+- Inplace operations on inference tensors are guaranteed not to do a version bump.
+
+For more implementation details, see the [RFC-0011-InferenceMode](https://github.com/pytorch/rfcs/pull/17).
+
+**Basic Example:**
+
+```cpp
+{
+    c10::InferenceMode guard;
+    // Optimized inference without gradient tracking
+    auto result = model->forward(input);
+}
+```
+
+**Inference Workload Example:**
+
+```cpp
+c10::InferenceMode guard;
+model.load_jit(saved_model);
+auto inputs = preprocess_tensors(data);
+auto out = model.forward(inputs);
+auto outputs = postprocess_tensors(out);
+```
+
+**Nested InferenceMode:**
+
+Unlike some other guards, `InferenceMode` can be nested with different enabled/disabled states:
+
+```cpp
+{
+    c10::InferenceMode guard(true);
+    // InferenceMode is on
+    {
+        c10::InferenceMode guard(false);
+        // InferenceMode is off
+    }
+    // InferenceMode is on
+}
+// InferenceMode is off
+```
+
+## InferenceMode vs NoGradMode
+
+`InferenceMode` is preferred over `NoGradMode` for pure inference workloads because
+it provides better performance. Key differences:
+
+- Both guards affect tensor execution to skip work not related to inference, but
+  `InferenceMode` also affects tensor creation while `NoGradMode` doesn't.
+- Tensors created inside `InferenceMode` are marked as inference tensors with
+  certain limitations that apply after exiting `InferenceMode`.
+- `InferenceMode` can be nested with enabled/disabled states.
+
+## Migrating from AutoNonVariableTypeMode
+
+The legacy `AutoNonVariableTypeMode` guard (now renamed to
+`AutoDispatchBelowADInplaceOrView`) was commonly used for inference workloads
+but is unsafe — it can silently bypass safety checks and produce wrong results.
+
+- **For inference-only workloads** (e.g. loading a pretrained JIT model and
+  running inference in C++ runtime), use `c10::InferenceMode` as a drop-in
+  replacement. It preserves the performance characteristics while providing
+  correctness guarantees.
+
+- **For custom autograd kernels** that need to redispatch below the Autograd
+  dispatch key, use `AutoDispatchBelowADInplaceOrView` instead:
+
+  ```cpp
+  class ROIAlignFunction : public torch::autograd::Function<ROIAlignFunction> {
+   public:
+    static torch::autograd::variable_list forward(
+        torch::autograd::AutogradContext* ctx,
+        const torch::autograd::Variable& input,
+        const torch::autograd::Variable& rois,
+        double spatial_scale, int64_t pooled_height,
+        int64_t pooled_width, int64_t sampling_ratio, bool aligned) {
+      ctx->saved_data["spatial_scale"] = spatial_scale;
+      ctx->save_for_backward({rois});
+      at::AutoDispatchBelowADInplaceOrView guard;
+      auto result = roi_align(input, rois, spatial_scale,
+          pooled_height, pooled_width, sampling_ratio, aligned);
+      return {result};
+    }
+  };
+  ```
diff --git a/docs/cpp/source/api/c10/device.md b/docs/cpp/source/api/c10/device.md
new file mode 100644
index 0000000000000..229082935c11b
--- /dev/null
+++ b/docs/cpp/source/api/c10/device.md
@@ -0,0 +1,114 @@
+---
+myst:
+  html_meta:
+    description: Device and DeviceType in PyTorch C++ — c10::Device for specifying CPU, CUDA, XPU, and other backends.
+    keywords: PyTorch, C++, Device, DeviceType, CPU, CUDA, XPU, MPS, c10
+---
+
+# Device and DeviceType
+
+PyTorch provides device abstractions for writing code that works across
+CPU, CUDA, and other backends.
+
+## Device
+
+```{doxygenstruct} c10::Device
+:members:
+:undoc-members:
+```
+
+**Example:**
+
+```cpp
+c10::Device cpu_device(c10::kCPU);
+c10::Device cuda_device(c10::kCUDA, 0);  // CUDA device 0
+
+if (cuda_device.is_cuda()) {
+    std::cout << "Using CUDA device " << cuda_device.index() << std::endl;
+}
+```
+
+## DeviceType
+
+```{cpp:enum-class} c10::DeviceType
+
+Enumeration of supported device types.
+```
+
+```{cpp:enumerator} CPU = 0
+
+CPU device.
+```
+
+```{cpp:enumerator} CUDA = 1
+
+NVIDIA CUDA GPU.
+```
+
+```{cpp:enumerator} HIP = 6
+
+AMD HIP GPU.
+```
+
+```{cpp:enumerator} XLA = 9
+
+XLA / TPU.
+```
+
+```{cpp:enumerator} Vulkan = 10
+
+Vulkan GPU.
+```
+
+```{cpp:enumerator} Metal = 11
+
+Apple Metal GPU.
+```
+
+```{cpp:enumerator} XPU = 12
+
+Intel XPU GPU.
+```
+
+```{cpp:enumerator} MPS = 13
+
+Apple Metal Performance Shaders.
+```
+
+```{cpp:enumerator} Meta = 14
+
+Meta tensors (shape only, no data).
+```
+
+```{cpp:enumerator} HPU = 15
+
+Habana HPU.
+```
+
+```{cpp:enumerator} Lazy = 17
+
+Lazy tensors.
+```
+
+```{cpp:enumerator} IPU = 18
+
+Graphcore IPU.
+```
+
+```{cpp:enumerator} MTIA = 19
+
+Meta training and inference accelerator.
+```
+
+```{cpp:enumerator} PrivateUse1 = 20
+
+Custom backend registered via `c10::register_privateuse1_backend()`.
+```
+
+Convenience constants:
+
+- `c10::kCPU`, `c10::kCUDA`, `c10::kHIP`
+- `c10::kXLA`, `c10::kVulkan`, `c10::kMetal`
+- `c10::kXPU`, `c10::kMPS`, `c10::kMeta`
+- `c10::kHPU`, `c10::kLazy`, `c10::kIPU`, `c10::kMTIA`
+- `c10::kPrivateUse1`
diff --git a/docs/cpp/source/api/c10/guards.md b/docs/cpp/source/api/c10/guards.md
new file mode 100644
index 0000000000000..fba83afbcee57
--- /dev/null
+++ b/docs/cpp/source/api/c10/guards.md
@@ -0,0 +1,52 @@
+---
+myst:
+  html_meta:
+    description: Device and stream guards in PyTorch C++ — RAII guards for managing current device and stream context.
+    keywords: PyTorch, C++, DeviceGuard, StreamGuard, RAII, device management
+---
+
+# Device Guards
+
+C10 provides device-agnostic RAII guards for managing the current device
+context. These guards work across all backends (CUDA, XPU, etc.) and
+automatically restore the previous device when they go out of scope.
+
+For backend-specific guards, see {doc}`../cuda/guards` and {doc}`../xpu/index`.
+
+## DeviceGuard
+
+```{doxygenclass} c10::DeviceGuard
+:members:
+:undoc-members:
+```
+
+**Example:**
+
+```cpp
+#include <c10/core/DeviceGuard.h>
+
+{
+    c10::DeviceGuard guard(c10::Device(c10::kCUDA, 1));
+    // All operations here run on CUDA device 1
+}
+// Previous device is restored
+```
+
+## OptionalDeviceGuard
+
+```{doxygenclass} c10::OptionalDeviceGuard
+:members:
+:undoc-members:
+```
+
+**Example:**
+
+```cpp
+#include <c10/core/DeviceGuard.h>
+
+c10::OptionalDeviceGuard guard;
+if (use_gpu) {
+    guard.reset_device(c10::Device(c10::kCUDA, 0));
+}
+// Guard only restores device if it was set
+```
diff --git a/docs/cpp/source/api/c10/index.md b/docs/cpp/source/api/c10/index.md
new file mode 100644
index 0000000000000..dd34df87409fa
--- /dev/null
+++ b/docs/cpp/source/api/c10/index.md
@@ -0,0 +1,66 @@
+---
+myst:
+  html_meta:
+    description: C10 core library API — fundamental types, device management, streams, and utilities used throughout PyTorch.
+    keywords: PyTorch, C++, c10, core, Device, Stream, types
+---
+
+# C10: Core Utilities
+
+C10 (Caffe2 + ATen = C10) is the core library that provides fundamental
+utilities and data types used throughout PyTorch. It contains device
+abstractions, memory management utilities, and common data structures.
+
+**When to use C10:**
+
+- When working with device-agnostic code (CPU, CUDA, XPU, etc.)
+- When you need efficient array views without copying data
+- When handling optional values or type-erased containers
+- When writing code that needs to work across different PyTorch backends
+
+**Basic usage:**
+
+```cpp
+#include <c10/core/Device.h>
+#include <c10/util/ArrayRef.h>
+
+// Device abstraction
+c10::Device device(c10::kCUDA, 0);
+if (device.is_cuda()) {
+    std::cout << "Using CUDA device " << device.index() << std::endl;
+}
+
+// Efficient array views (no copy)
+std::vector<int64_t> sizes = {3, 4, 5};
+c10::ArrayRef<int64_t> sizes_ref(sizes);
+
+// Optional values
+c10::optional<int64_t> maybe_dim = 2;
+int64_t dim = maybe_dim.value_or(-1);
+```
+
+## Header Files
+
+- `c10/core/Device.h` - Device abstraction
+- `c10/core/DeviceType.h` - Device type enumeration
+- `c10/util/ArrayRef.h` - Non-owning array reference
+- `c10/util/OptionalArrayRef.h` - Optional array reference
+- `c10/util/Optional.h` - Optional value wrapper
+- `c10/util/Half.h` - Half-precision float
+- `c10/util/Exception.h` - Error checking macros (`TORCH_CHECK`, `TORCH_INTERNAL_ASSERT`, etc.)
+- `c10/cuda/CUDAGuard.h` - CUDA device guards (see {doc}`../cuda/index`)
+- `c10/cuda/CUDAStream.h` - CUDA stream management (see {doc}`../cuda/index`)
+- `c10/xpu/XPUStream.h` - Intel XPU stream management
+- `ATen/core/ivalue.h` - IValue for TorchScript interop
+
+## C10 Categories
+
+```{toctree}
+:maxdepth: 1
+
+device
+guards
+streams
+types
+utilities
+```
diff --git a/docs/cpp/source/api/c10/streams.md b/docs/cpp/source/api/c10/streams.md
new file mode 100644
index 0000000000000..e87e3d3fdad7a
--- /dev/null
+++ b/docs/cpp/source/api/c10/streams.md
@@ -0,0 +1,34 @@
+---
+myst:
+  html_meta:
+    description: Stream API in PyTorch C++ — c10::Stream for asynchronous execution on devices.
+    keywords: PyTorch, C++, Stream, c10::Stream, asynchronous, device
+---
+
+# Streams
+
+`c10::Stream` is the device-agnostic base stream class. It provides a
+common interface for working with streams across different backends
+(CUDA, XPU, etc.).
+
+For backend-specific stream APIs, see {doc}`../cuda/streams` and {doc}`../xpu/index`.
+
+## Stream
+
+```{doxygenclass} c10::Stream
+:members:
+:undoc-members:
+```
+
+**Example:**
+
+```cpp
+#include <c10/core/Stream.h>
+
+// Streams are typically obtained from backend-specific APIs
+auto cuda_stream = c10::cuda::getCurrentCUDAStream();
+
+// c10::Stream provides the common interface
+c10::Device device = cuda_stream.device();
+c10::DeviceType type = cuda_stream.device_type();
+```
diff --git a/docs/cpp/source/api/c10/types.md b/docs/cpp/source/api/c10/types.md
new file mode 100644
index 0000000000000..60ffa1d784c45
--- /dev/null
+++ b/docs/cpp/source/api/c10/types.md
@@ -0,0 +1,225 @@
+---
+myst:
+  html_meta:
+    description: Core types in PyTorch C++ — ArrayRef, optional, Dict, List, IListRef, Half, and IValue.
+    keywords: PyTorch, C++, c10, ArrayRef, optional, Dict, List, IListRef, Half, IValue
+---
+
+# Core Types
+
+C10 provides fundamental types used throughout PyTorch.
+
+## ArrayRef
+
+```{doxygenclass} c10::ArrayRef
+:members:
+:undoc-members:
+```
+
+**Example:**
+
+```cpp
+std::vector<int64_t> sizes = {3, 4, 5};
+c10::ArrayRef<int64_t> sizes_ref(sizes);
+
+// Can also use initializer list
+auto tensor = at::zeros({3, 4, 5});  // implicitly converts
+```
+
+## OptionalArrayRef
+
+```{doxygenclass} c10::OptionalArrayRef
+:members:
+:no-link:
+```
+
+**Example:**
+
+```cpp
+void my_function(c10::OptionalArrayRef<int64_t> sizes = c10::nullopt) {
+    if (sizes.has_value()) {
+        for (auto s : sizes.value()) {
+            // process sizes
+        }
+    }
+}
+```
+
+## Optional
+
+```{cpp:class} c10::optional
+
+A wrapper type that may or may not contain a value.
+Similar to `std::optional`.
+```
+
+```{cpp:function} bool has_value() const
+
+Returns true if a value is present.
+```
+
+```{cpp:function} T& value()
+
+Returns the contained value. Throws if empty.
+```
+
+```{cpp:function} T value_or(T default_value) const
+
+Returns the value if present, otherwise returns the default.
+```
+
+**Example:**
+
+```cpp
+c10::optional<int64_t> maybe_dim = c10::nullopt;
+
+if (maybe_dim.has_value()) {
+    std::cout << "Dim: " << maybe_dim.value() << std::endl;
+}
+
+int64_t dim = maybe_dim.value_or(-1);  // Returns -1 if empty
+```
+
+## Half
+
+```{cpp:class} c10::Half
+
+16-bit floating point type (IEEE 754 half-precision).
+```
+
+```{cpp:function} Half(float value)
+
+Construct from a float.
+```
+
+```{cpp:function} operator float() const
+
+Convert to float.
+```
+
+**Example:**
+
+```cpp
+c10::Half h = 3.14f;
+float f = static_cast<float>(h);
+```
+
+## Containers
+
+C10 provides container types that store `IValue` elements internally. These
+are pointer types: copies share the same underlying storage.
+
+### Dict
+
+An ordered hash map from `Key` to `Value`. Valid key types are `int64_t`,
+`double`, `bool`, `std::string`, and `at::Tensor`.
+
+```{doxygenclass} c10::Dict
+:members:
+:undoc-members:
+```
+
+**Example:**
+
+```cpp
+#include <ATen/core/Dict.h>
+
+c10::Dict<std::string, at::Tensor> named_tensors;
+named_tensors.insert("weight", torch::randn({3, 3}));
+named_tensors.insert("bias", torch::zeros({3}));
+
+if (named_tensors.contains("weight")) {
+    at::Tensor w = named_tensors.at("weight");
+}
+
+for (const auto& entry : named_tensors) {
+    std::cout << entry.key() << ": " << entry.value().sizes() << std::endl;
+}
+```
+
+### List
+
+A type-safe list container backed by `IValue` elements.
+
+```{doxygenclass} c10::List
+:members:
+:undoc-members:
+```
+
+**Example:**
+
+```cpp
+#include <ATen/core/List.h>
+
+c10::List<at::Tensor> tensor_list;
+tensor_list.push_back(torch::randn({2, 3}));
+tensor_list.push_back(torch::zeros({2, 3}));
+
+at::Tensor first = tensor_list.get(0);
+std::cout << "List size: " << tensor_list.size() << std::endl;
+
+c10::List<int64_t> int_list;
+int_list.push_back(1);
+int_list.push_back(2);
+int_list.push_back(3);
+```
+
+### IListRef
+
+`c10::IListRef<T>` is a lightweight reference type that provides a unified
+interface over different list-like types (`List<T>`, `ArrayRef<T>`,
+`std::vector<T>`). It avoids copying when passing list arguments to operators.
+
+```{doxygenclass} c10::IListRef
+:members:
+:undoc-members:
+```
+
+**Example:**
+
+```cpp
+#include <ATen/core/IListRef.h>
+
+// IListRef can wrap different underlying types
+std::vector<at::Tensor> vec = {torch::randn({2}), torch::randn({3})};
+c10::IListRef<at::Tensor> ref(vec);
+
+for (const auto& t : ref) {
+    std::cout << t.sizes() << std::endl;
+}
+```
+
+## IValue
+
+`c10::IValue` (Interpreter Value) is a type-erased container used extensively
+for storing values of different types. It can hold tensors,
+scalars, lists, dictionaries, and other types.
+
+```{note}
+
+The full API documentation for IValue is complex due to its many type
+conversion methods. See the header file `ATen/core/ivalue.h` for complete
+details.
+```
+
+**Common methods:**
+
+- `isTensor()` / `toTensor()` - Check if tensor / convert to tensor
+- `isInt()` / `toInt()` - Check if int / convert to int
+- `isDouble()` / `toDouble()` - Check if double / convert to double
+- `isBool()` / `toBool()` - Check if bool / convert to bool
+- `isString()` / `toString()` - Check if string / convert to string
+- `isList()` / `toList()` - Check if list / convert to list
+- `isGenericDict()` / `toGenericDict()` - Check if dict / convert to dict
+- `isTuple()` / `toTuple()` - Check if tuple / convert to tuple
+- `isNone()` - Check if None/null
+
+**Example:**
+
+```cpp
+c10::IValue val = at::ones({2, 2});
+
+if (val.isTensor()) {
+    at::Tensor t = val.toTensor();
+}
+```
diff --git a/docs/cpp/source/api/c10/utilities.md b/docs/cpp/source/api/c10/utilities.md
new file mode 100644
index 0000000000000..2e61c14c8e8dc
--- /dev/null
+++ b/docs/cpp/source/api/c10/utilities.md
@@ -0,0 +1,153 @@
+---
+myst:
+  html_meta:
+    description: C10 utility classes in PyTorch C++ — Flags, QEngine, and Reduction enumerations.
+    keywords: PyTorch, C++, c10, utilities, QEngine, Reduction
+---
+
+# Utilities
+
+C10 provides utility classes for memory management and other common patterns.
+
+## MaybeOwned<Tensor>
+
+`MaybeOwned<Tensor>` is a C++ smart pointer class that dynamically
+encodes whether a Tensor is *owned* or *borrowed*. It is used in
+certain performance-sensitive situations to avoid unnecessarily
+incrementing a Tensor's reference count (at a small cost in
+overhead from the extra indirection).
+
+```{warning}
+
+ MaybeOwned must be used with **extreme** care. Claims of (non-)ownership
+ are not statically checked, and mistakes can cause reference undercounting
+ and use-after-free crashes.
+
+ Due to this lack of safety net, we discourage the use of MaybeOwned
+ outside code paths that are known to be highly performance sensitive.
+ However, if you encounter pre-existing uses of MaybeOwned in code that
+ you want to modify, it's critical to understand how to use it correctly.
+```
+
+**Use Case:**
+
+The primary use case for `MaybeOwned<Tensor>` is a function or method that
+dynamically chooses between returning one of its arguments (typically
+from a passthrough or "no-op" code path) and returning a freshly constructed
+Tensor. Such a function would return a `MaybeOwned<Tensor>` in both cases:
+the former in a "borrowed" state via `MaybeOwned<Tensor>::borrowed()`,
+and the latter in an "owned" state via `MaybeOwned<Tensor>::owned()`.
+
+**Example - expect_contiguous:**
+
+The canonical example is `Tensor`'s `expect_contiguous` method, which shortcuts
+and returns a borrowed self-reference when already contiguous:
+
+```cpp
+inline c10::MaybeOwned<Tensor> Tensor::expect_contiguous(
+    MemoryFormat memory_format) const & {
+  if (is_contiguous(memory_format)) {
+    return c10::MaybeOwned<Tensor>::borrowed(*this);
+  } else {
+    return c10::MaybeOwned<Tensor>::owned(
+        __dispatch_contiguous(memory_format));
+  }
+}
+```
+
+Using the vocabulary of lifetimes, the essential safety requirement for borrowing
+is that a borrowed Tensor must outlive any borrowing references to it. In the example
+above, we can safely borrow `*this`, but the Tensor returned by
+`__dispatch_contiguous()` is freshly created, and borrowing a reference would
+effectively leave it ownerless.
+
+**Rules of Thumb:**
+
+- When in doubt, don't use `MaybeOwned<Tensor>` at all - in particular, prefer
+  avoiding using it in code that doesn't use it already. New usage should only be
+  introduced when critical (and demonstrable) performance gains result.
+
+- When modifying or calling code that already uses `MaybeOwned<Tensor>`, remember
+  that it's always safe to produce a `MaybeOwned<Tensor>` from a Tensor in hand
+  via a call to `MaybeOwned<Tensor>::owned()`. This may result in an unnecessary
+  reference count, but never in misbehavior - so it's always the safer bet, unless
+  the lifetime of the Tensor you're looking to wrap is crystal clear.
+
+More details and implementation code can be found at
+[MaybeOwned.h](https://github.com/pytorch/pytorch/blob/main/c10/util/MaybeOwned.h) and
+[TensorBody.h](https://github.com/pytorch/pytorch/blob/main/aten/src/ATen/templates/TensorBody.h).
+
+## Error Handling and Assertions
+
+PyTorch provides macros for error checking and assertions that produce
+informative error messages with source location. These are defined in
+`c10/util/Exception.h`.
+
+### TORCH_CHECK
+
+The primary macro for validating user input and runtime conditions. On failure,
+raises `c10::Error` (which becomes `RuntimeError` in Python).
+
+```cpp
+#include <c10/util/Exception.h>
+
+// Basic check
+TORCH_CHECK(tensor.dim() == 2, "Expected 2D tensor, got ", tensor.dim(), "D");
+
+// Without message (default message generated)
+TORCH_CHECK(x >= 0);
+```
+
+Typed variants raise specific Python exception types:
+
+- `TORCH_CHECK_INDEX(cond, ...)` — raises `IndexError`
+- `TORCH_CHECK_VALUE(cond, ...)` — raises `ValueError`
+- `TORCH_CHECK_TYPE(cond, ...)` — raises `TypeError`
+- `TORCH_CHECK_LINALG(cond, ...)` — raises `LinAlgError`
+- `TORCH_CHECK_NOT_IMPLEMENTED(cond, ...)` — raises `NotImplementedError`
+
+### TORCH_INTERNAL_ASSERT
+
+For internal invariants that should always hold (i.e., failures indicate a bug
+in PyTorch, not user error). Produces a message asking users to report the bug.
+
+```cpp
+TORCH_INTERNAL_ASSERT(googol > 0);
+TORCH_INTERNAL_ASSERT(googol > 0, "googol was ", googol);
+```
+
+```{note}
+
+Use `TORCH_CHECK` for conditions that can fail due to user input.
+Use `TORCH_INTERNAL_ASSERT` only for conditions that indicate a PyTorch bug.
+`TORCH_INTERNAL_ASSERT_DEBUG_ONLY` is the debug-build-only variant for
+hot paths.
+```
+
+### TORCH_WARN
+
+Issues a warning (not an error) to the user.
+
+```cpp
+TORCH_WARN("This operation is slow for sparse tensors");
+TORCH_WARN_ONCE("This warning appears only once");
+```
+
+### c10::Error
+
+The base exception class for PyTorch C++ errors. Provides source location
+and optional backtrace.
+
+```cpp
+try {
+    auto result = some_operation();
+} catch (const c10::Error& e) {
+    std::cerr << e.what() << std::endl;
+    // Or without backtrace:
+    std::cerr << e.what_without_backtrace() << std::endl;
+}
+```
+
+Specialized subclasses: `c10::IndexError`, `c10::ValueError`,
+`c10::TypeError`, `c10::NotImplementedError`, `c10::LinAlgError`,
+`c10::OutOfMemoryError`.
diff --git a/docs/cpp/source/api/cuda/guards.md b/docs/cpp/source/api/cuda/guards.md
new file mode 100644
index 0000000000000..d7724e390525a
--- /dev/null
+++ b/docs/cpp/source/api/cuda/guards.md
@@ -0,0 +1,96 @@
+---
+myst:
+  html_meta:
+    description: CUDA device and stream guards in PyTorch C++ — CUDAGuard, CUDAStreamGuard, and CUDAMultiStreamGuard.
+    keywords: PyTorch, C++, CUDA, CUDAGuard, CUDAStreamGuard, device guard, multi-GPU
+---
+
+# CUDA Guards
+
+CUDA guards are RAII wrappers that set a CUDA device or stream as the current
+context and automatically restore the previous context when the guard goes
+out of scope.
+
+## CUDAGuard
+
+```{doxygenstruct} c10::cuda::CUDAGuard
+:members:
+:undoc-members:
+```
+
+**Example:**
+
+```cpp
+#include <c10/cuda/CUDAGuard.h>
+
+{
+    c10::cuda::CUDAGuard guard(1);  // Switch to device 1
+    // All CUDA operations here run on device 1
+    auto tensor = torch::zeros({2, 2}, torch::device(torch::kCUDA));
+}
+// Previous device is restored
+```
+
+## CUDAStreamGuard
+
+```{doxygenstruct} c10::cuda::CUDAStreamGuard
+:members:
+:undoc-members:
+```
+
+**Example:**
+
+```cpp
+#include <c10/cuda/CUDAGuard.h>
+
+auto stream = c10::cuda::getStreamFromPool();
+{
+    c10::cuda::CUDAStreamGuard guard(stream);
+    // Operations here use the specified stream
+}
+// Previous stream is restored
+```
+
+## OptionalCUDAGuard
+
+```{doxygenstruct} c10::cuda::OptionalCUDAGuard
+:members:
+:undoc-members:
+```
+
+**Example:**
+
+```cpp
+c10::cuda::OptionalCUDAGuard guard;
+if (use_cuda) {
+    guard.set_device(0);
+}
+// Guard only switches device if set_device was called
+```
+
+## OptionalCUDAStreamGuard
+
+```{doxygenstruct} c10::cuda::OptionalCUDAStreamGuard
+:members:
+:undoc-members:
+```
+
+## CUDAMultiStreamGuard
+
+```{doxygenstruct} c10::cuda::CUDAMultiStreamGuard
+:members:
+:undoc-members:
+```
+
+**Example:**
+
+```cpp
+at::cuda::CUDAStream stream0 = at::cuda::getStreamFromPool(false, 0);
+at::cuda::CUDAStream stream1 = at::cuda::getStreamFromPool(false, 1);
+
+{
+    at::cuda::CUDAMultiStreamGuard multi_guard({stream0, stream1});
+    // stream0 is current on device 0, stream1 on device 1
+}
+// Both streams restored
+```
diff --git a/docs/cpp/source/api/cuda/index.md b/docs/cpp/source/api/cuda/index.md
new file mode 100644
index 0000000000000..0f7f07f447b4f
--- /dev/null
+++ b/docs/cpp/source/api/cuda/index.md
@@ -0,0 +1,58 @@
+---
+myst:
+  html_meta:
+    description: PyTorch CUDA C++ API — device management, streams, guards, and cuDNN/cuBLAS utilities.
+    keywords: PyTorch, C++, CUDA, GPU, streams, guards, cuDNN, cuBLAS
+---
+
+# CUDA Support
+
+PyTorch provides comprehensive CUDA support for GPU-accelerated tensor
+operations and neural network training. The CUDA API allows you to manage
+GPU devices, streams for asynchronous execution, and memory efficiently.
+
+**When to use CUDA APIs:**
+
+- When you need explicit control over which GPU device to use
+- When implementing custom CUDA kernels or operations
+- When optimizing performance with asynchronous stream execution
+- When managing multi-GPU workloads
+
+**Basic usage:**
+
+```cpp
+#include <torch/torch.h>
+#include <c10/cuda/CUDAGuard.h>
+
+// Check if CUDA is available
+if (torch::cuda::is_available()) {
+    // Create tensor on GPU
+    auto tensor = torch::randn({2, 3}, torch::device(torch::kCUDA));
+
+    // Switch to a specific GPU
+    c10::cuda::CUDAGuard guard(0);  // Use GPU 0
+
+    // Get the current CUDA stream
+    auto stream = c10::cuda::getCurrentCUDAStream();
+
+    // Move model to GPU
+    model->to(torch::kCUDA);
+}
+```
+
+## Header Files
+
+- `c10/cuda/CUDAStream.h` - CUDA stream management
+- `c10/cuda/CUDAGuard.h` - CUDA device guards
+- `ATen/cuda/CUDAContext.h` - CUDA context management
+- `ATen/cudnn/Descriptors.h` - cuDNN tensor descriptors
+
+## CUDA Categories
+
+```{toctree}
+:maxdepth: 1
+
+streams
+guards
+utilities
+```
diff --git a/docs/cpp/source/api/cuda/streams.md b/docs/cpp/source/api/cuda/streams.md
new file mode 100644
index 0000000000000..a55eb609d1bf7
--- /dev/null
+++ b/docs/cpp/source/api/cuda/streams.md
@@ -0,0 +1,190 @@
+---
+myst:
+  html_meta:
+    description: CUDA streams in PyTorch C++ — CUDAStream for asynchronous GPU execution and synchronization.
+    keywords: PyTorch, C++, CUDA, CUDAStream, stream, asynchronous, GPU, synchronization
+---
+
+# CUDA Streams
+
+CUDA streams provide a mechanism for asynchronous execution of operations
+on the GPU. Operations queued to the same stream execute in order, while
+operations on different streams can execute concurrently.
+
+## CUDAStream
+
+```{doxygenclass} c10::cuda::CUDAStream
+:members:
+:undoc-members:
+```
+
+**Example:**
+
+```cpp
+#include <c10/cuda/CUDAStream.h>
+
+// Get the default stream for current device
+auto stream = c10::cuda::getDefaultCUDAStream();
+
+// Create a new stream
+auto new_stream = c10::cuda::getStreamFromPool();
+
+// Get current stream
+auto current = c10::cuda::getCurrentCUDAStream();
+
+// Synchronize
+stream.synchronize();
+```
+
+## Acquiring CUDA Streams
+
+PyTorch provides several ways to acquire CUDA streams:
+
+1. **From the stream pool** (round-robin allocation):
+
+   ```cpp
+   // Normal priority stream
+   at::cuda::CUDAStream stream = at::cuda::getStreamFromPool();
+
+   // High priority stream
+   at::cuda::CUDAStream high_prio = at::cuda::getStreamFromPool(/*isHighPriority=*/true);
+
+   // Stream for specific device
+   at::cuda::CUDAStream dev1_stream = at::cuda::getStreamFromPool(false, /*device=*/1);
+   ```
+
+2. **Default stream** (where most computation occurs):
+
+   ```cpp
+   at::cuda::CUDAStream defaultStream = at::cuda::getDefaultCUDAStream();
+   ```
+
+3. **Current stream** (may differ if changed with guards):
+
+   ```cpp
+   at::cuda::CUDAStream currentStream = at::cuda::getCurrentCUDAStream();
+   ```
+
+## Setting CUDA Streams
+
+**Using setCurrentCUDAStream:**
+
+```cpp
+torch::Tensor tensor0 = torch::ones({2, 2}, torch::device(torch::kCUDA));
+
+// Get a new stream and set it as current
+at::cuda::CUDAStream myStream = at::cuda::getStreamFromPool();
+at::cuda::setCurrentCUDAStream(myStream);
+
+// Operations now use myStream
+tensor0.sum();
+
+// Restore default stream
+at::cuda::setCurrentCUDAStream(at::cuda::getDefaultCUDAStream());
+```
+
+**Using CUDAStreamGuard (recommended):**
+
+```cpp
+torch::Tensor tensor0 = torch::ones({2, 2}, torch::device(torch::kCUDA));
+at::cuda::CUDAStream myStream = at::cuda::getStreamFromPool();
+
+{
+    at::cuda::CUDAStreamGuard guard(myStream);
+    // Operations use myStream within this scope
+    tensor0.sum();
+}
+// Stream automatically restored to default
+```
+
+## Multi-Device Stream Management
+
+**Streams on multiple devices:**
+
+```cpp
+// Acquire streams for different devices
+at::cuda::CUDAStream stream0 = at::cuda::getStreamFromPool(false, 0);
+at::cuda::CUDAStream stream1 = at::cuda::getStreamFromPool(false, 1);
+
+// Set current streams on each device
+at::cuda::setCurrentCUDAStream(stream0);
+at::cuda::setCurrentCUDAStream(stream1);
+
+// Create tensors on device 0
+torch::Tensor tensor0 = torch::ones({2, 2}, torch::device(at::kCUDA));
+tensor0.sum();  // Uses stream0
+
+// Switch to device 1
+{
+    at::cuda::CUDAGuard device_guard(1);
+    torch::Tensor tensor1 = torch::ones({2, 2}, torch::device(at::kCUDA));
+    tensor1.sum();  // Uses stream1
+}
+```
+
+**Using CUDAMultiStreamGuard:**
+
+```cpp
+torch::Tensor tensor0 = torch::ones({2, 2}, torch::device({torch::kCUDA, 0}));
+torch::Tensor tensor1 = torch::ones({2, 2}, torch::device({torch::kCUDA, 1}));
+
+at::cuda::CUDAStream stream0 = at::cuda::getStreamFromPool(false, 0);
+at::cuda::CUDAStream stream1 = at::cuda::getStreamFromPool(false, 1);
+
+{
+    // Set streams on both devices simultaneously
+    at::cuda::CUDAMultiStreamGuard multi_guard({stream0, stream1});
+
+    tensor0.sum();  // Uses stream0 on device 0
+    tensor1.sum();  // Uses stream1 on device 1
+}
+// Both streams restored to defaults
+```
+
+```{attention}
+
+`CUDAMultiStreamGuard` does not change the current device index. It only
+changes the stream on each passed-in stream's device.
+```
+
+## Multi-Device Stream Handling Pattern
+
+The following skeleton shows three common patterns for acquiring and setting
+streams across multiple CUDA devices:
+
+```cpp
+// Create stream vectors on device 0
+std::vector<at::cuda::CUDAStream> streams0 =
+  {at::cuda::getDefaultCUDAStream(), at::cuda::getStreamFromPool()};
+at::cuda::setCurrentCUDAStream(streams0[0]);
+
+// Create stream vector on device 1 using CUDAGuard
+std::vector<at::cuda::CUDAStream> streams1;
+{
+  at::cuda::CUDAGuard device_guard(1);
+  streams1.push_back(at::cuda::getDefaultCUDAStream());
+  streams1.push_back(at::cuda::getStreamFromPool());
+}
+at::cuda::setCurrentCUDAStream(streams1[0]);
+
+// Pattern 1: CUDAGuard changes current device only, not streams
+{
+  at::cuda::CUDAGuard device_guard(1);
+  // current device is 1, current stream on device 1 is still streams1[0]
+}
+
+// Pattern 2: CUDAStreamGuard changes both current device and current stream
+{
+  at::cuda::CUDAStreamGuard stream_guard(streams1[1]);
+  // current device is 1, current stream is streams1[1]
+}
+// restored to device 0, stream streams0[0]
+
+// Pattern 3: CUDAMultiStreamGuard sets streams on multiple devices at once
+{
+  at::cuda::CUDAMultiStreamGuard multi_guard({streams0[1], streams1[1]});
+  // current device unchanged (still 0)
+  // stream on device 0 is streams0[1], stream on device 1 is streams1[1]
+}
+// streams restored to streams0[0] and streams1[0]
+```
diff --git a/docs/cpp/source/api/cuda/utilities.md b/docs/cpp/source/api/cuda/utilities.md
new file mode 100644
index 0000000000000..951c49e3b71de
--- /dev/null
+++ b/docs/cpp/source/api/cuda/utilities.md
@@ -0,0 +1,226 @@
+---
+myst:
+  html_meta:
+    description: CUDA utility functions in PyTorch C++ — device properties, cuBLAS/cuDNN handles, and stream management.
+    keywords: PyTorch, C++, CUDA, device_count, cuBLAS, cuDNN, cuSPARSE, TensorDescriptor
+---
+
+# CUDA Utility Functions
+
+PyTorch provides utility functions for querying and managing CUDA devices,
+streams, and library handles.
+
+## Device Management
+
+```{doxygenfunction} torch::cuda::device_count
+```
+
+```{cpp:function} int c10::cuda::current_device()
+
+Returns the index of the current CUDA device.
+```
+
+**Example:**
+
+```cpp
+#include <c10/cuda/CUDAFunctions.h>
+
+// Check available devices
+int num_devices = c10::cuda::device_count();
+
+// Get current device
+int current = c10::cuda::current_device();
+```
+
+## Device Properties
+
+```{doxygenfunction} at::cuda::getCurrentDeviceProperties
+```
+
+```{doxygenfunction} at::cuda::getDeviceProperties
+```
+
+```{doxygenfunction} at::cuda::canDeviceAccessPeer
+```
+
+```{doxygenfunction} at::cuda::warp_size
+```
+
+**Example:**
+
+```cpp
+#include <ATen/cuda/CUDAContext.h>
+
+// Query properties of the current device
+cudaDeviceProp* props = at::cuda::getCurrentDeviceProperties();
+std::cout << "Device: " << props->name << std::endl;
+std::cout << "Compute capability: " << props->major << "." << props->minor << std::endl;
+
+// Query a specific device
+cudaDeviceProp* dev1_props = at::cuda::getDeviceProperties(1);
+
+// Check peer access
+bool can_access = at::cuda::canDeviceAccessPeer(0, 1);
+```
+
+## Library Handles
+
+These functions return handles for CUDA math libraries on the current device
+and stream. They are primarily useful when writing custom CUDA kernels that
+call cuBLAS, cuSPARSE, or cuSOLVER directly.
+
+```{doxygenfunction} at::cuda::getCurrentCUDABlasHandle
+```
+
+```{doxygenfunction} at::cuda::getCurrentCUDABlasLtHandle
+```
+
+```{doxygenfunction} at::cuda::getCurrentCUDASparseHandle
+```
+
+```{doxygenfunction} at::cuda::getCurrentCUDASolverDnHandle
+```
+
+**Example:**
+
+```cpp
+#include <ATen/cuda/CUDAContext.h>
+
+// Get cuBLAS handle for current device/stream
+cublasHandle_t handle = at::cuda::getCurrentCUDABlasHandle();
+
+// Get cuSPARSE handle
+cusparseHandle_t sparse_handle = at::cuda::getCurrentCUDASparseHandle();
+```
+
+## cuDNN Descriptors
+
+When writing custom kernels that call cuDNN directly, PyTorch provides RAII
+wrapper classes for cuDNN descriptors. These are defined in
+`ATen/cudnn/Descriptors.h`.
+
+### Descriptor (Base Class)
+
+```{doxygenclass} at::native::Descriptor
+:members:
+:undoc-members:
+```
+
+A generic RAII wrapper for cuDNN descriptor types. Descriptors default
+construct to a `nullptr` and are initialized on first use via `mut_desc()`.
+Use `desc()` for read-only access.
+
+### TensorDescriptor
+
+```{doxygenclass} at::native::TensorDescriptor
+:members:
+:undoc-members:
+```
+
+Wraps `cudnnTensorDescriptor_t`. Supports padding lower-dimensional tensors
+to meet cuDNN broadcasting requirements (see `pad` parameter).
+
+**Example:**
+
+```cpp
+#include <ATen/cudnn/Descriptors.h>
+
+at::Tensor input = torch::randn({32, 3, 224, 224}, torch::kCUDA);
+at::native::TensorDescriptor desc(input);
+cudnnTensorDescriptor_t raw = desc.desc();
+```
+
+### FilterDescriptor
+
+```{doxygenclass} at::native::FilterDescriptor
+:members:
+:undoc-members:
+```
+
+Wraps `cudnnFilterDescriptor_t` for convolution filter weights.
+
+### ConvolutionDescriptor
+
+```{doxygenstruct} at::native::ConvolutionDescriptor
+:members:
+:undoc-members:
+```
+
+Wraps `cudnnConvolutionDescriptor_t`. Configures padding, stride, dilation,
+groups, and math type (TF32, tensor ops) for convolution operations.
+
+### RNNDataDescriptor
+
+```{doxygenclass} at::native::RNNDataDescriptor
+:members:
+:undoc-members:
+```
+
+Wraps `cudnnRNNDataDescriptor_t` for variable-length sequence data.
+
+### DropoutDescriptor
+
+```{doxygenstruct} at::native::DropoutDescriptor
+:members:
+:undoc-members:
+```
+
+Wraps `cudnnDropoutDescriptor_t`. Manages RNG state for cuDNN dropout.
+
+### ActivationDescriptor
+
+```{doxygenstruct} at::native::ActivationDescriptor
+:members:
+:undoc-members:
+```
+
+Wraps `cudnnActivationDescriptor_t`.
+
+### SpatialTransformerDescriptor
+
+```{doxygenstruct} at::native::SpatialTransformerDescriptor
+:members:
+:undoc-members:
+```
+
+### CTCLossDescriptor
+
+```{doxygenstruct} at::native::CTCLossDescriptor
+:members:
+:undoc-members:
+```
+
+## Stream Management
+
+```{doxygenfunction} c10::cuda::getDefaultCUDAStream
+```
+
+```{doxygenfunction} c10::cuda::getCurrentCUDAStream
+```
+
+```{doxygenfunction} c10::cuda::setCurrentCUDAStream
+```
+
+```{doxygenfunction} c10::cuda::getStreamFromPool(const bool isHighPriority, DeviceIndex device)
+```
+
+```{doxygenfunction} c10::cuda::getStreamFromExternal
+```
+
+**Example:**
+
+```cpp
+#include <c10/cuda/CUDAStream.h>
+
+// Create and set custom stream
+auto stream = c10::cuda::getStreamFromPool();
+c10::cuda::setCurrentCUDAStream(stream);
+
+// Get default stream
+auto default_stream = c10::cuda::getDefaultCUDAStream();
+
+// Wrap an externally created CUDA stream
+cudaStream_t ext_stream;
+cudaStreamCreate(&ext_stream);
+auto wrapped = c10::cuda::getStreamFromExternal(ext_stream, /*device_index=*/0);
+```
diff --git a/docs/cpp/source/api/data/dataloader.md b/docs/cpp/source/api/data/dataloader.md
new file mode 100644
index 0000000000000..f368363368bf6
--- /dev/null
+++ b/docs/cpp/source/api/data/dataloader.md
@@ -0,0 +1,103 @@
+---
+myst:
+  html_meta:
+    description: DataLoader in PyTorch C++ — parallel data loading with batching, sampling, and multi-worker support.
+    keywords: PyTorch, C++, DataLoader, data loading, batching, workers, make_data_loader
+---
+
+# DataLoader
+
+The DataLoader batches samples from a dataset and optionally shuffles and
+parallelizes the loading process. It is the main interface for iterating
+over training data.
+
+## DataLoader Classes
+
+```{doxygenclass} torch::data::DataLoaderBase
+:members:
+:undoc-members:
+```
+
+## DataLoaderOptions
+
+```{doxygenstruct} torch::data::DataLoaderOptions
+:members:
+:undoc-members:
+```
+
+## StatefulDataLoader
+
+A DataLoader for `StatefulDataset` types that manage their own batching logic
+internally.
+
+```{doxygenclass} torch::data::StatefulDataLoader
+:members:
+:undoc-members:
+```
+
+## StatelessDataLoader
+
+A DataLoader for `Dataset` types that use external samplers for batching.
+
+```{doxygenclass} torch::data::StatelessDataLoader
+:members:
+:undoc-members:
+```
+
+## Iterator
+
+```{doxygenclass} torch::data::Iterator
+:members:
+:undoc-members:
+```
+
+## Creating a DataLoader
+
+Use `make_data_loader` to create a DataLoader from a dataset:
+
+```cpp
+auto data_loader = torch::data::make_data_loader(
+    std::move(dataset),
+    torch::data::DataLoaderOptions()
+        .batch_size(64)
+        .workers(4));
+
+for (auto& batch : *data_loader) {
+    auto data = batch.data;
+    auto target = batch.target;
+    // Train on batch
+}
+```
+
+## Complete Training Example
+
+```cpp
+#include <torch/torch.h>
+
+int main() {
+  // Load dataset
+  auto dataset = torch::data::datasets::MNIST("./data")
+      .map(torch::data::transforms::Normalize<>(0.1307, 0.3081))
+      .map(torch::data::transforms::Stack<>());
+
+  // Create data loader
+  auto data_loader = torch::data::make_data_loader(
+      std::move(dataset),
+      torch::data::DataLoaderOptions().batch_size(64).workers(2));
+
+  // Create model and optimizer
+  auto model = std::make_shared<Net>();
+  auto optimizer = torch::optim::Adam(model->parameters(), 0.001);
+
+  // Training loop
+  for (size_t epoch = 1; epoch <= 10; ++epoch) {
+    for (auto& batch : *data_loader) {
+      optimizer.zero_grad();
+      auto output = model->forward(batch.data);
+      auto loss = torch::nll_loss(output, batch.target);
+      loss.backward();
+      optimizer.step();
+    }
+  }
+}
+```
diff --git a/docs/cpp/source/api/data/datasets.md b/docs/cpp/source/api/data/datasets.md
new file mode 100644
index 0000000000000..4bb01e255c393
--- /dev/null
+++ b/docs/cpp/source/api/data/datasets.md
@@ -0,0 +1,110 @@
+---
+myst:
+  html_meta:
+    description: Dataset classes in PyTorch C++ — Dataset, MapDataset, StreamDataset, and built-in datasets like MNIST.
+    keywords: PyTorch, C++, Dataset, MapDataset, StreamDataset, MNIST, data
+---
+
+# Datasets
+
+The dataset abstraction defines how to access individual samples in your data.
+All datasets inherit from `Dataset` and must implement `get()` and `size()`.
+
+## Dataset Base Class
+
+```{doxygenclass} torch::data::datasets::Dataset
+:members:
+:undoc-members:
+```
+
+```{doxygenclass} torch::data::datasets::BatchDataset
+:members:
+:undoc-members:
+```
+
+## StatefulDataset
+
+A dataset that manages its own state across batches (e.g., position in a stream).
+Unlike `Dataset`, it produces batches directly without external samplers.
+
+```{doxygenclass} torch::data::datasets::StatefulDataset
+:members:
+:undoc-members:
+```
+
+## ChunkDataReader
+
+Interface for reading chunks of data from a data source. Used with
+`ChunkDataset` for large-scale data loading.
+
+```{doxygenclass} torch::data::datasets::ChunkDataReader
+:members:
+:undoc-members:
+```
+
+## Custom Dataset Example
+
+```cpp
+class CustomDataset : public torch::data::datasets::Dataset<CustomDataset> {
+ public:
+  explicit CustomDataset(const std::string& root) {
+    // Load data from root directory
+  }
+
+  torch::data::Example<> get(size_t index) override {
+    return {images_[index], labels_[index]};
+  }
+
+  torch::optional<size_t> size() const override {
+    return images_.size(0);
+  }
+
+ private:
+  torch::Tensor images_, labels_;
+};
+```
+
+## MapDataset
+
+```{doxygenclass} torch::data::datasets::MapDataset
+:members:
+:undoc-members:
+```
+
+## ChunkDataset
+
+```{doxygenclass} torch::data::datasets::ChunkDataset
+:members:
+:undoc-members:
+```
+
+## SharedBatchDataset
+
+```{doxygenclass} torch::data::datasets::SharedBatchDataset
+:members:
+:undoc-members:
+```
+
+## Built-in Datasets
+
+### MNIST
+
+```{doxygenclass} torch::data::datasets::MNIST
+:members:
+:undoc-members:
+```
+
+**Example:**
+
+```cpp
+auto dataset = torch::data::datasets::MNIST("./data")
+    .map(torch::data::transforms::Normalize<>(0.1307, 0.3081))
+    .map(torch::data::transforms::Stack<>());
+```
+
+## Example Struct
+
+```{doxygenstruct} torch::data::Example
+:members:
+:undoc-members:
+```
diff --git a/docs/cpp/source/api/data/index.md b/docs/cpp/source/api/data/index.md
new file mode 100644
index 0000000000000..7e5c5517c510c
--- /dev/null
+++ b/docs/cpp/source/api/data/index.md
@@ -0,0 +1,66 @@
+---
+myst:
+  html_meta:
+    description: PyTorch C++ data loading API — datasets, data loaders, samplers, and transforms.
+    keywords: PyTorch, C++, data, DataLoader, Dataset, sampler, transform
+---
+
+# Data Loading (torch::data)
+
+The `torch::data` namespace provides utilities for loading and processing
+datasets during training. It includes dataset abstractions, data loaders for
+batching and shuffling, samplers for controlling data access patterns, and
+transforms for data augmentation.
+
+**When to use torch::data:**
+
+- When loading training data in batches
+- When you need parallel data loading with multiple workers
+- When implementing custom datasets or transforms
+
+**Components overview:**
+
+- **Dataset**: Defines how to access individual samples (implement `get()` and `size()`)
+- **DataLoader**: Batches samples and optionally shuffles/parallelizes loading
+- **Sampler**: Controls the order in which samples are accessed
+- **Transform**: Applies preprocessing (normalization, augmentation) to samples
+
+**Basic usage:**
+
+```cpp
+#include <torch/torch.h>
+
+// Load built-in dataset
+auto dataset = torch::data::datasets::MNIST("./data")
+    .map(torch::data::transforms::Normalize<>(0.1307, 0.3081))
+    .map(torch::data::transforms::Stack<>());
+
+// Create data loader with batching and shuffling
+auto data_loader = torch::data::make_data_loader(
+    std::move(dataset),
+    torch::data::DataLoaderOptions().batch_size(64).workers(4));
+
+// Iterate over batches
+for (auto& batch : *data_loader) {
+    auto images = batch.data;   // Shape: [64, 1, 28, 28]
+    auto labels = batch.target; // Shape: [64]
+}
+```
+
+## Header Files
+
+- `torch/csrc/api/include/torch/data.h` - Main data header
+- `torch/csrc/api/include/torch/data/dataloader.h` - DataLoader
+- `torch/csrc/api/include/torch/data/datasets.h` - Dataset classes
+- `torch/csrc/api/include/torch/data/samplers.h` - Samplers
+
+## Module Categories
+
+```{toctree}
+:maxdepth: 1
+
+datasets
+dataloader
+samplers
+transforms
+```
diff --git a/docs/cpp/source/api/data/samplers.md b/docs/cpp/source/api/data/samplers.md
new file mode 100644
index 0000000000000..f18f09c27d493
--- /dev/null
+++ b/docs/cpp/source/api/data/samplers.md
@@ -0,0 +1,69 @@
+---
+myst:
+  html_meta:
+    description: Data samplers in PyTorch C++ — RandomSampler, SequentialSampler, DistributedRandomSampler, and StreamSampler.
+    keywords: PyTorch, C++, sampler, RandomSampler, SequentialSampler, DistributedRandomSampler
+---
+
+# Samplers
+
+Samplers control the order in which samples are accessed from a dataset.
+They determine the indices that the DataLoader uses to fetch data.
+
+## Sampler Base Class
+
+```{doxygenclass} torch::data::samplers::Sampler
+:members:
+:undoc-members:
+```
+
+## Sequential Sampler
+
+Accesses samples in order from 0 to N-1. Use this for evaluation or when
+order matters.
+
+```{doxygenclass} torch::data::samplers::SequentialSampler
+:members:
+:undoc-members:
+```
+
+## Random Sampler
+
+Accesses samples in random order. Use this for training to ensure the model
+sees samples in different orders each epoch.
+
+```{doxygenclass} torch::data::samplers::RandomSampler
+:members:
+:undoc-members:
+```
+
+## Distributed Random Sampler
+
+For distributed training, ensures each process gets a different subset of
+the data without overlap.
+
+```{doxygenclass} torch::data::samplers::DistributedRandomSampler
+:members:
+:undoc-members:
+```
+
+## Distributed Sampler (Base)
+
+```{doxygenclass} torch::data::samplers::DistributedSampler
+:members:
+:undoc-members:
+```
+
+## Distributed Sequential Sampler
+
+```{doxygenclass} torch::data::samplers::DistributedSequentialSampler
+:members:
+:undoc-members:
+```
+
+## Stream Sampler
+
+```{doxygenclass} torch::data::samplers::StreamSampler
+:members:
+:undoc-members:
+```
diff --git a/docs/cpp/source/api/data/transforms.md b/docs/cpp/source/api/data/transforms.md
new file mode 100644
index 0000000000000..5026596e8b63b
--- /dev/null
+++ b/docs/cpp/source/api/data/transforms.md
@@ -0,0 +1,95 @@
+---
+myst:
+  html_meta:
+    description: Data transforms in PyTorch C++ — Stack, Normalize, Lambda, and Collate transforms for data pipelines.
+    keywords: PyTorch, C++, transforms, Stack, Normalize, Lambda, Collate, data pipeline
+---
+
+# Transforms
+
+Transforms apply preprocessing to data samples, such as normalization or
+augmentation. They can be chained using the `.map()` method on datasets.
+
+## Transform (Base Class)
+
+The base class for all transforms. Subclass this to create custom transforms.
+
+```{doxygenclass} torch::data::transforms::Transform
+:members:
+:undoc-members:
+```
+
+## BatchTransform (Base Class)
+
+Base class for transforms that operate on entire batches.
+
+```{doxygenclass} torch::data::transforms::BatchTransform
+:members:
+:undoc-members:
+```
+
+## TensorTransform
+
+Base class for transforms that operate on tensors specifically.
+
+```{doxygenclass} torch::data::transforms::TensorTransform
+:members:
+:undoc-members:
+```
+
+## Normalize
+
+Normalizes tensors with a given mean and standard deviation.
+
+```{doxygenstruct} torch::data::transforms::Normalize
+:members:
+:undoc-members:
+```
+
+## Stack
+
+Stacks a batch of tensors into a single tensor.
+
+```{doxygenstruct} torch::data::transforms::Stack
+:members:
+:undoc-members:
+```
+
+**Example:**
+
+```cpp
+auto dataset = torch::data::datasets::MNIST("./data")
+    .map(torch::data::transforms::Normalize<>(0.5, 0.5))
+    .map(torch::data::transforms::Stack<>());
+```
+
+## Lambda
+
+```{doxygenclass} torch::data::transforms::Lambda
+:members:
+:undoc-members:
+```
+
+## TensorLambda
+
+```{doxygenclass} torch::data::transforms::TensorLambda
+:members:
+:undoc-members:
+```
+
+## BatchLambda
+
+```{doxygenclass} torch::data::transforms::BatchLambda
+:members:
+:undoc-members:
+```
+
+## Chaining Transforms
+
+Transforms can be chained together using `.map()`:
+
+```cpp
+auto dataset = torch::data::datasets::MNIST("./data")
+    .map(torch::data::transforms::Normalize<>(0.1307, 0.3081))
+    .map(torch::data::transforms::Stack<>());
+```
diff --git a/docs/cpp/source/api/index.md b/docs/cpp/source/api/index.md
new file mode 100644
index 0000000000000..3c3ff9d855039
--- /dev/null
+++ b/docs/cpp/source/api/index.md
@@ -0,0 +1,40 @@
+---
+myst:
+  html_meta:
+    description: PyTorch C++ API reference — complete documentation for ATen, Autograd, nn, optim, data, CUDA, and more.
+    keywords: PyTorch, C++, API reference, ATen, Autograd, nn, optim, CUDA
+---
+
+# C++ API Reference
+
+This section provides reference documentation for the PyTorch C++ API,
+organized by module.
+
+```{toctree}
+:maxdepth: 2
+:caption: Core
+
+aten/index
+c10/index
+autograd/index
+cuda/index
+xpu/index
+```
+
+```{toctree}
+:maxdepth: 2
+:caption: C++ Frontend
+
+nn/index
+optim/index
+data/index
+serialize/index
+```
+
+```{toctree}
+:maxdepth: 2
+:caption: Extensions
+
+library/index
+stable/index
+```
diff --git a/docs/cpp/source/api/library/custom_classes.md b/docs/cpp/source/api/library/custom_classes.md
new file mode 100644
index 0000000000000..48828eb731aba
--- /dev/null
+++ b/docs/cpp/source/api/library/custom_classes.md
@@ -0,0 +1,94 @@
+---
+myst:
+  html_meta:
+    description: Custom classes in PyTorch C++ — registering C++ classes for use in TorchScript and Python.
+    keywords: PyTorch, C++, custom class, TorchScript, TORCH_CLASS, registration
+---
+
+# Custom Classes
+
+PyTorch allows registering custom C++ classes that can be used from Python
+and TorchScript.
+
+Header: `torch/custom_class.h`
+
+## class\_ Template
+
+```{doxygenclass} torch::class_
+:members:
+:undoc-members:
+```
+
+**Example:**
+
+```cpp
+#include <torch/custom_class.h>
+
+struct MyClass : torch::CustomClassHolder {
+    int value;
+
+    MyClass(int v) : value(v) {}
+
+    int getValue() const { return value; }
+    void setValue(int v) { value = v; }
+};
+
+TORCH_LIBRARY(my_classes, m) {
+    m.class_<MyClass>("MyClass")
+        .def(torch::init<int>())
+        .def("getValue", &MyClass::getValue)
+        .def("setValue", &MyClass::setValue)
+        .def_readwrite("value", &MyClass::value);
+}
+```
+
+## Registering Methods
+
+**Constructor:**
+
+```cpp
+m.class_<MyClass>("MyClass")
+    .def(torch::init<int>())  // Constructor taking int
+```
+
+**Methods:**
+
+```cpp
+m.class_<MyClass>("MyClass")
+    .def("getValue", &MyClass::getValue)
+    .def("setValue", &MyClass::setValue)
+```
+
+**Properties:**
+
+```cpp
+m.class_<MyClass>("MyClass")
+    .def_readwrite("value", &MyClass::value)   // Read-write
+    .def_readonly("const_value", &MyClass::const_value)  // Read-only
+```
+
+## Using Custom Classes
+
+**From C++:**
+
+```cpp
+auto my_obj = c10::make_intrusive<MyClass>(42);
+int val = my_obj->getValue();
+```
+
+**From Python:**
+
+```python
+import torch
+torch.classes.load_library("path/to/library.so")
+obj = torch.classes.my_classes.MyClass(42)
+print(obj.getValue())
+```
+
+**In TorchScript:**
+
+```python
+@torch.jit.script
+def use_my_class(x: torch.classes.my_classes.MyClass) -> int:
+    return x.getValue()
+```
diff --git a/docs/cpp/source/api/library/index.md b/docs/cpp/source/api/library/index.md
new file mode 100644
index 0000000000000..2d3e1e225a1be
--- /dev/null
+++ b/docs/cpp/source/api/library/index.md
@@ -0,0 +1,62 @@
+---
+myst:
+  html_meta:
+    description: PyTorch Library API in C++ — operator registration, custom classes, and versioning.
+    keywords: PyTorch, C++, Library, operator registration, custom class, dispatch
+---
+
+# Torch Library API
+
+The Torch Library API provides capabilities for extending PyTorch's core library
+of operators with user-defined operators and data types. This is the primary
+mechanism for registering custom C++ operators that can be called from both
+Python and C++.
+
+**When to use the Library API:**
+
+- When creating custom operators for PyTorch
+- When implementing backend-specific kernels (CPU, CUDA, etc.)
+- When registering custom classes for use in TorchScript
+- When extending PyTorch with new functionality
+
+**Basic usage:**
+
+```cpp
+#include <torch/library.h>
+
+// Define a custom operator
+torch::Tensor my_add(const torch::Tensor& a, const torch::Tensor& b) {
+    return a + b;
+}
+
+// Register the operator
+TORCH_LIBRARY(myops, m) {
+    m.def("add(Tensor a, Tensor b) -> Tensor", &my_add);
+}
+
+// Use from C++
+auto result = torch::dispatcher::call("myops::add", tensor_a, tensor_b);
+```
+
+For a tutorial-style introduction to the library API, check out the
+[Extending TorchScript with Custom C++ Operators](https://pytorch.org/tutorials/advanced/torch_script_custom_ops.html)
+tutorial.
+
+## Header Files
+
+- `torch/library.h` - Main library API header
+- `torch/custom_class.h` - Custom class registration
+
+## Library API Categories
+
+```{toctree}
+:maxdepth: 1
+
+registration
+custom_classes
+versioning
+```
+
+## See Also
+
+- {doc}`../stable/index` - For stable ABI operator registration
diff --git a/docs/cpp/source/api/library/registration.md b/docs/cpp/source/api/library/registration.md
new file mode 100644
index 0000000000000..ec74871915681
--- /dev/null
+++ b/docs/cpp/source/api/library/registration.md
@@ -0,0 +1,113 @@
+---
+myst:
+  html_meta:
+    description: Operator registration in PyTorch C++ — TORCH_LIBRARY, TORCH_LIBRARY_IMPL for custom operators.
+    keywords: PyTorch, C++, TORCH_LIBRARY, operator registration, dispatch, custom op
+---
+
+# Operator Registration
+
+The library API provides macros and classes for registering custom operators
+with PyTorch's dispatcher.
+
+## Macros
+
+### TORCH_LIBRARY
+
+```{doxygendefine} TORCH_LIBRARY
+```
+
+**Example:**
+
+```cpp
+TORCH_LIBRARY(myops, m) {
+  m.def("add(Tensor self, Tensor other) -> Tensor", &add_impl);
+  m.def("mul(Tensor self, Tensor other) -> Tensor");
+  m.impl("mul", torch::kCPU, &mul_cpu_impl);
+  m.impl("mul", torch::kCUDA, &mul_cuda_impl);
+}
+```
+
+### TORCH_LIBRARY_IMPL
+
+```{doxygendefine} TORCH_LIBRARY_IMPL
+```
+
+**Example:**
+
+```cpp
+TORCH_LIBRARY_IMPL(myops, XLA, m) {
+  m.impl("mul", &mul_xla_impl);
+}
+```
+
+### TORCH_LIBRARY_FRAGMENT
+
+```{doxygendefine} TORCH_LIBRARY_FRAGMENT
+```
+
+**Example:**
+
+```cpp
+// In file1.cpp
+TORCH_LIBRARY(myops, m) {
+  m.def("add(Tensor self, Tensor other) -> Tensor", &add_impl);
+}
+
+// In file2.cpp
+TORCH_LIBRARY_FRAGMENT(myops, m) {
+  m.def("mul(Tensor self, Tensor other) -> Tensor", &mul_impl);
+}
+```
+
+## Classes
+
+### Library
+
+```{doxygenclass} torch::Library
+```
+
+**Example:**
+
+```cpp
+TORCH_LIBRARY(myops, m) {
+  // Define with implementation
+  m.def("add(Tensor self, Tensor other) -> Tensor", &add_impl);
+
+  // Define schema only
+  m.def("mul(Tensor self, Tensor other) -> Tensor");
+
+  // Provide backend-specific implementations
+  m.impl("mul", torch::kCPU, &mul_cpu_impl);
+  m.impl("mul", torch::kCUDA, &mul_cuda_impl);
+}
+```
+
+### CppFunction
+
+```{doxygenclass} torch::CppFunction
+:members:
+:no-link:
+```
+
+### OrderedDict
+
+```{doxygenclass} torch::OrderedDict
+:members:
+:undoc-members:
+```
+
+## Functions
+
+The library API provides builder methods on the `Library` class for registering
+operators. See the `Library` class documentation above for the full API including
+`def()`, `impl()`, and `fallback()` methods.
+
+## Dispatch Keys
+
+Common dispatch keys used with `torch::dispatch()`:
+
+- `torch::kCPU` - CPU backend
+- `torch::kCUDA` - CUDA backend
+- `torch::kAutograd` - Autograd backend
+- `torch::kMeta` - Meta tensor backend
diff --git a/docs/cpp/source/api/library/versioning.md b/docs/cpp/source/api/library/versioning.md
new file mode 100644
index 0000000000000..a287f3c2a5751
--- /dev/null
+++ b/docs/cpp/source/api/library/versioning.md
@@ -0,0 +1,44 @@
+---
+myst:
+  html_meta:
+    description: Operator versioning in PyTorch C++ — managing backward compatibility for serialized models.
+    keywords: PyTorch, C++, versioning, backward compatibility, operator, serialization
+---
+
+# Library Versioning
+
+PyTorch provides version number macros for identifying the version of LibTorch in use.
+
+**Example:**
+
+```cpp
+#include <torch/torch.h>
+#include <iostream>
+
+int main() {
+  std::cout << "PyTorch version from parts: "
+    << TORCH_VERSION_MAJOR << "."
+    << TORCH_VERSION_MINOR << "."
+    << TORCH_VERSION_PATCH << std::endl;
+  std::cout << "PyTorch version: " << TORCH_VERSION << std::endl;
+}
+```
+
+This will output something like:
+
+```text
+PyTorch version from parts: 1.8.0
+PyTorch version: 1.8.0
+```
+
+```{note}
+
+These macros are only available in PyTorch >= 1.8.0.
+```
+
+## Version Macros
+
+- `TORCH_VERSION_MAJOR` - Major version number
+- `TORCH_VERSION_MINOR` - Minor version number
+- `TORCH_VERSION_PATCH` - Patch version number
+- `TORCH_VERSION` - Full version string (e.g., "1.8.0")
diff --git a/docs/cpp/source/api/nn/activation.md b/docs/cpp/source/api/nn/activation.md
new file mode 100644
index 0000000000000..d4551490f2f70
--- /dev/null
+++ b/docs/cpp/source/api/nn/activation.md
@@ -0,0 +1,356 @@
+---
+myst:
+  html_meta:
+    description: Activation functions in PyTorch C++ — ReLU, GELU, Sigmoid, Softmax, and more torch::nn activation modules.
+    keywords: PyTorch, C++, activation, ReLU, GELU, Sigmoid, Softmax, LeakyReLU, ELU, Mish
+---
+
+# Activation Functions
+
+Activation functions introduce non-linearity into neural networks, allowing them
+to learn complex patterns. Without activations, stacked linear layers would collapse
+into a single linear transformation.
+
+**Common choices:**
+
+- **ReLU family** (ReLU, LeakyReLU, PReLU, RReLU): Fast, widely used, good default choice
+- **ELU family** (ELU, SELU, CELU): Smoother than ReLU, can produce negative outputs
+- **GELU/SiLU/Mish**: Modern activations popular in transformers and advanced architectures
+- **Sigmoid/Tanh**: Classic activations, useful for output layers (probabilities, bounded outputs)
+- **Softmax**: Converts logits to probability distribution (classification output)
+
+## ReLU
+
+```{doxygenclass} torch::nn::ReLU
+:members:
+:undoc-members:
+```
+
+```{doxygenclass} torch::nn::ReLUImpl
+:members:
+:undoc-members:
+```
+
+**Example:**
+
+```cpp
+auto relu = torch::nn::ReLU(torch::nn::ReLUOptions().inplace(true));
+```
+
+## LeakyReLU
+
+```{doxygenclass} torch::nn::LeakyReLU
+:members:
+:undoc-members:
+```
+
+```{doxygenclass} torch::nn::LeakyReLUImpl
+:members:
+:undoc-members:
+```
+
+## PReLU
+
+```{doxygenclass} torch::nn::PReLU
+:members:
+:undoc-members:
+```
+
+```{doxygenclass} torch::nn::PReLUImpl
+:members:
+:undoc-members:
+```
+
+## RReLU
+
+```{doxygenclass} torch::nn::RReLU
+:members:
+:undoc-members:
+```
+
+```{doxygenclass} torch::nn::RReLUImpl
+:members:
+:undoc-members:
+```
+
+## ReLU6
+
+Like ReLU but caps the output at 6: `min(max(0, x), 6)`. Commonly used in
+mobile architectures (MobileNet).
+
+```{doxygenclass} torch::nn::ReLU6
+:members:
+:undoc-members:
+```
+
+```{doxygenclass} torch::nn::ReLU6Impl
+:members:
+:undoc-members:
+```
+
+## GLU
+
+Gated Linear Unit. Splits the input tensor in half along a dimension,
+then applies `a * sigmoid(b)`.
+
+```{doxygenclass} torch::nn::GLU
+:members:
+:undoc-members:
+```
+
+```{doxygenclass} torch::nn::GLUImpl
+:members:
+:undoc-members:
+```
+
+## LogSigmoid
+
+Applies element-wise `log(sigmoid(x))`. Numerically more stable than
+computing `log` and `sigmoid` separately.
+
+```{doxygenclass} torch::nn::LogSigmoid
+:members:
+:undoc-members:
+```
+
+```{doxygenclass} torch::nn::LogSigmoidImpl
+:members:
+:undoc-members:
+```
+
+## ELU
+
+```{doxygenclass} torch::nn::ELU
+:members:
+:undoc-members:
+```
+
+```{doxygenclass} torch::nn::ELUImpl
+:members:
+:undoc-members:
+```
+
+## SELU
+
+```{doxygenclass} torch::nn::SELU
+:members:
+:undoc-members:
+```
+
+```{doxygenclass} torch::nn::SELUImpl
+:members:
+:undoc-members:
+```
+
+## CELU
+
+```{doxygenclass} torch::nn::CELU
+:members:
+:undoc-members:
+```
+
+```{doxygenclass} torch::nn::CELUImpl
+:members:
+:undoc-members:
+```
+
+## GELU
+
+```{doxygenclass} torch::nn::GELU
+:members:
+:undoc-members:
+```
+
+```{doxygenclass} torch::nn::GELUImpl
+:members:
+:undoc-members:
+```
+
+## SiLU (Swish)
+
+```{doxygenclass} torch::nn::SiLU
+:members:
+:undoc-members:
+```
+
+```{doxygenclass} torch::nn::SiLUImpl
+:members:
+:undoc-members:
+```
+
+## Mish
+
+```{doxygenclass} torch::nn::Mish
+:members:
+:undoc-members:
+```
+
+```{doxygenclass} torch::nn::MishImpl
+:members:
+:undoc-members:
+```
+
+## Sigmoid
+
+```{doxygenclass} torch::nn::Sigmoid
+:members:
+:undoc-members:
+```
+
+```{doxygenclass} torch::nn::SigmoidImpl
+:members:
+:undoc-members:
+```
+
+## Tanh
+
+```{doxygenclass} torch::nn::Tanh
+:members:
+:undoc-members:
+```
+
+```{doxygenclass} torch::nn::TanhImpl
+:members:
+:undoc-members:
+```
+
+## Softmax
+
+```{doxygenclass} torch::nn::Softmax
+:members:
+:undoc-members:
+```
+
+```{doxygenclass} torch::nn::SoftmaxImpl
+:members:
+:undoc-members:
+```
+
+**Example:**
+
+```cpp
+auto softmax = torch::nn::Softmax(torch::nn::SoftmaxOptions(/*dim=*/1));
+```
+
+## Softmax2d
+
+Applies `Softmax` over features to each spatial location in a 4D input
+tensor of shape `(N, C, H, W)`.
+
+```{doxygenclass} torch::nn::Softmax2d
+:members:
+:undoc-members:
+```
+
+```{doxygenclass} torch::nn::Softmax2dImpl
+:members:
+:undoc-members:
+```
+
+## LogSoftmax
+
+```{doxygenclass} torch::nn::LogSoftmax
+:members:
+:undoc-members:
+```
+
+```{doxygenclass} torch::nn::LogSoftmaxImpl
+:members:
+:undoc-members:
+```
+
+## Softmin
+
+```{doxygenclass} torch::nn::Softmin
+:members:
+:undoc-members:
+```
+
+```{doxygenclass} torch::nn::SoftminImpl
+:members:
+:undoc-members:
+```
+
+## Softplus
+
+```{doxygenclass} torch::nn::Softplus
+:members:
+:undoc-members:
+```
+
+```{doxygenclass} torch::nn::SoftplusImpl
+:members:
+:undoc-members:
+```
+
+## Softshrink
+
+```{doxygenclass} torch::nn::Softshrink
+:members:
+:undoc-members:
+```
+
+```{doxygenclass} torch::nn::SoftshrinkImpl
+:members:
+:undoc-members:
+```
+
+## Softsign
+
+```{doxygenclass} torch::nn::Softsign
+:members:
+:undoc-members:
+```
+
+```{doxygenclass} torch::nn::SoftsignImpl
+:members:
+:undoc-members:
+```
+
+## Hardshrink
+
+```{doxygenclass} torch::nn::Hardshrink
+:members:
+:undoc-members:
+```
+
+```{doxygenclass} torch::nn::HardshrinkImpl
+:members:
+:undoc-members:
+```
+
+## Hardtanh
+
+```{doxygenclass} torch::nn::Hardtanh
+:members:
+:undoc-members:
+```
+
+```{doxygenclass} torch::nn::HardtanhImpl
+:members:
+:undoc-members:
+```
+
+## Tanhshrink
+
+```{doxygenclass} torch::nn::Tanhshrink
+:members:
+:undoc-members:
+```
+
+```{doxygenclass} torch::nn::TanhshrinkImpl
+:members:
+:undoc-members:
+```
+
+## Threshold
+
+```{doxygenclass} torch::nn::Threshold
+:members:
+:undoc-members:
+```
+
+```{doxygenclass} torch::nn::ThresholdImpl
+:members:
+:undoc-members:
+```
diff --git a/docs/cpp/source/api/nn/containers.md b/docs/cpp/source/api/nn/containers.md
new file mode 100644
index 0000000000000..c9a694a96b158
--- /dev/null
+++ b/docs/cpp/source/api/nn/containers.md
@@ -0,0 +1,110 @@
+---
+myst:
+  html_meta:
+    description: Module containers in PyTorch C++ — Sequential, ModuleList, ModuleDict for composing neural networks.
+    keywords: PyTorch, C++, Sequential, ModuleList, ModuleDict, container, module
+---
+
+# Containers
+
+Container modules hold other modules and define how they are composed together.
+Use containers to build complex architectures from simpler building blocks.
+
+- **Sequential**: Chain modules in order, output of one feeds into the next
+- **ModuleList**: Store modules in a list for iteration (not auto-forwarded)
+- **ModuleDict**: Store modules in a dictionary for named access
+- **ParameterList/ParameterDict**: Store parameters directly without wrapping in modules
+
+```{note}
+
+PyTorch's C++ API uses the PIMPL (Pointer to Implementation) pattern. You create
+modules using the public class name (e.g., `torch::nn::Sequential`), which
+internally wraps an implementation class (`SequentialImpl`). The documentation
+below shows the implementation classes, which contain all the actual methods.
+```
+
+## Sequential
+
+`Sequential` is a container that chains modules together. Each module's output
+becomes the next module's input. This is the simplest way to build feed-forward
+networks.
+
+```{doxygenclass} torch::nn::Sequential
+:members:
+:undoc-members:
+```
+
+**Example:**
+
+```cpp
+torch::nn::Sequential seq(
+    torch::nn::Conv2d(torch::nn::Conv2dOptions(1, 32, 3)),
+    torch::nn::ReLU(),
+    torch::nn::Conv2d(torch::nn::Conv2dOptions(32, 64, 3)),
+    torch::nn::ReLU()
+);
+
+auto output = seq->forward(input);
+```
+
+## ModuleList
+
+`ModuleList` stores modules in a list for indexed or iterated access. Unlike
+`Sequential`, it does not have a built-in `forward()` method—you control how
+modules are called.
+
+```{doxygenclass} torch::nn::ModuleList
+:members:
+:undoc-members:
+```
+
+**Example:**
+
+```cpp
+torch::nn::ModuleList layers;
+layers->push_back(torch::nn::Linear(10, 20));
+layers->push_back(torch::nn::Linear(20, 30));
+
+torch::Tensor x = input;
+for (const auto& layer : *layers) {
+    x = layer->as<torch::nn::Linear>()->forward(x);
+}
+```
+
+## ModuleDict
+
+`ModuleDict` stores modules in a dictionary for named access. Useful when you
+need to select modules by name at runtime.
+
+```{doxygenclass} torch::nn::ModuleDict
+:members:
+:undoc-members:
+```
+
+## ParameterList
+
+`ParameterList` stores parameters directly without wrapping them in modules.
+
+```{doxygenclass} torch::nn::ParameterList
+:members:
+:undoc-members:
+```
+
+```{doxygenclass} torch::nn::ParameterListImpl
+:members:
+:undoc-members:
+```
+
+## ParameterDict
+
+`ParameterDict` stores parameters in a dictionary for named access.
+
+```{doxygenclass} torch::nn::ParameterDict
+:members:
+:undoc-members:
+```
+
+```{doxygenclass} torch::nn::ParameterDictImpl
+:members:
+:undoc-members:
+```
diff --git a/docs/cpp/source/api/nn/convolution.md b/docs/cpp/source/api/nn/convolution.md
new file mode 100644
index 0000000000000..2c2addf3ce266
--- /dev/null
+++ b/docs/cpp/source/api/nn/convolution.md
@@ -0,0 +1,134 @@
+---
+myst:
+  html_meta:
+    description: Convolution layers in PyTorch C++ — Conv1d, Conv2d, Conv3d, and transposed convolutions.
+    keywords: PyTorch, C++, convolution, Conv1d, Conv2d, Conv3d, ConvTranspose, nn
+---
+
+# Convolution Layers
+
+Convolutional layers apply learnable filters to input data, extracting local features
+through sliding window operations. They are fundamental to CNNs for image, audio, and
+sequential data processing.
+
+- **Conv1d/2d/3d**: Standard convolution for 1D sequences, 2D images, or 3D volumes
+- **ConvTranspose1d/2d/3d**: Transposed convolution (deconvolution) for upsampling
+
+**Key parameters:**
+
+- `in_channels`: Number of input channels
+- `out_channels`: Number of output channels (number of filters)
+- `kernel_size`: Size of the convolving kernel
+- `stride`: Stride of the convolution (default: 1)
+- `padding`: Zero-padding added to input (default: 0)
+- `dilation`: Spacing between kernel elements (default: 1)
+- `groups`: Number of blocked connections (default: 1, use `in_channels` for depthwise)
+
+## Conv1d
+
+Applies 1D convolution over an input signal composed of several input planes.
+
+```{doxygenclass} torch::nn::Conv1d
+:members:
+:undoc-members:
+```
+
+```{doxygenclass} torch::nn::Conv1dImpl
+:members:
+:undoc-members:
+```
+
+## Conv2d
+
+Applies 2D convolution over an input image. The most commonly used layer for
+image processing tasks.
+
+```{doxygenclass} torch::nn::Conv2d
+:members:
+:undoc-members:
+```
+
+```{doxygenclass} torch::nn::Conv2dImpl
+:members:
+:undoc-members:
+```
+
+**Example:**
+
+```cpp
+// Create Conv2d: 3 input channels, 64 output channels, 3x3 kernel
+auto conv = torch::nn::Conv2d(
+    torch::nn::Conv2dOptions(3, 64, 3)
+        .stride(1)
+        .padding(1)
+        .bias(true));
+
+auto output = conv->forward(input);  // input: [N, 3, H, W]
+```
+
+## Conv3d
+
+Applies 3D convolution over an input volume (e.g., video frames or 3D medical images).
+
+```{doxygenclass} torch::nn::Conv3d
+:members:
+:undoc-members:
+```
+
+```{doxygenclass} torch::nn::Conv3dImpl
+:members:
+:undoc-members:
+```
+
+## ConvTranspose1d
+
+Applies 1D transposed convolution (fractionally-strided convolution) for upsampling.
+
+```{doxygenclass} torch::nn::ConvTranspose1d
+:members:
+:undoc-members:
+```
+
+```{doxygenclass} torch::nn::ConvTranspose1dImpl
+:members:
+:undoc-members:
+```
+
+## ConvTranspose2d
+
+Applies 2D transposed convolution for upsampling. Commonly used in decoder
+networks and generative models.
+
+```{doxygenclass} torch::nn::ConvTranspose2d
+:members:
+:undoc-members:
+```
+
+```{doxygenclass} torch::nn::ConvTranspose2dImpl
+:members:
+:undoc-members:
+```
+
+**Example:**
+
+```cpp
+// Create ConvTranspose2d for upsampling
+auto conv_transpose = torch::nn::ConvTranspose2d(
+    torch::nn::ConvTranspose2dOptions(64, 32, 4)
+        .stride(2)
+        .padding(1));
+```
+
+## ConvTranspose3d
+
+Applies 3D transposed convolution for upsampling volumetric data.
+
+```{doxygenclass} torch::nn::ConvTranspose3d
+:members:
+:undoc-members:
+```
+
+```{doxygenclass} torch::nn::ConvTranspose3dImpl
+:members:
+:undoc-members:
+```
diff --git a/docs/cpp/source/api/nn/dropout.md b/docs/cpp/source/api/nn/dropout.md
new file mode 100644
index 0000000000000..0777452ab9f34
--- /dev/null
+++ b/docs/cpp/source/api/nn/dropout.md
@@ -0,0 +1,86 @@
+---
+myst:
+  html_meta:
+    description: Dropout layers in PyTorch C++ — Dropout, Dropout2d, Dropout3d, AlphaDropout, and FeatureAlphaDropout.
+    keywords: PyTorch, C++, Dropout, Dropout2d, Dropout3d, AlphaDropout, regularization
+---
+
+# Dropout Layers
+
+Dropout randomly zeros elements during training as a regularization technique,
+preventing overfitting by forcing the network to learn redundant representations.
+During evaluation, dropout is disabled and outputs are scaled appropriately.
+
+- **Dropout**: Standard dropout for fully-connected layers
+- **Dropout2d/3d**: Spatial dropout that zeros entire channels (better for CNNs)
+- **AlphaDropout**: Maintains self-normalizing property (use with SELU activation)
+
+```{note}
+
+Remember to call `model->train()` during training and `model->eval()` during
+inference to properly enable/disable dropout behavior.
+```
+
+## Dropout
+
+```{doxygenclass} torch::nn::Dropout
+:members:
+:undoc-members:
+```
+
+```{doxygenclass} torch::nn::DropoutImpl
+:members:
+:undoc-members:
+```
+
+**Example:**
+
+```cpp
+auto dropout = torch::nn::Dropout(torch::nn::DropoutOptions(0.5));
+```
+
+## Dropout2d / Dropout3d
+
+```{doxygenclass} torch::nn::Dropout2d
+:members:
+:undoc-members:
+```
+
+```{doxygenclass} torch::nn::Dropout2dImpl
+:members:
+:undoc-members:
+```
+
+```{doxygenclass} torch::nn::Dropout3d
+:members:
+:undoc-members:
+```
+
+```{doxygenclass} torch::nn::Dropout3dImpl
+:members:
+:undoc-members:
+```
+
+## AlphaDropout
+
+```{doxygenclass} torch::nn::AlphaDropout
+:members:
+:undoc-members:
+```
+
+```{doxygenclass} torch::nn::AlphaDropoutImpl
+:members:
+:undoc-members:
+```
+
+## FeatureAlphaDropout
+
+```{doxygenclass} torch::nn::FeatureAlphaDropout
+:members:
+:undoc-members:
+```
+
+```{doxygenclass} torch::nn::FeatureAlphaDropoutImpl
+:members:
+:undoc-members:
+```
diff --git a/docs/cpp/source/api/nn/embedding.md b/docs/cpp/source/api/nn/embedding.md
new file mode 100644
index 0000000000000..f766a4ea17bfc
--- /dev/null
+++ b/docs/cpp/source/api/nn/embedding.md
@@ -0,0 +1,55 @@
+---
+myst:
+  html_meta:
+    description: Embedding layers in PyTorch C++ — Embedding and EmbeddingBag for sparse and dense lookups.
+    keywords: PyTorch, C++, Embedding, EmbeddingBag, lookup table, sparse, NLP
+---
+
+# Embedding Layers
+
+Embedding layers map discrete tokens (words, categories, IDs) to dense vector
+representations. They are the foundation of NLP models and recommendation systems.
+
+- **Embedding**: Standard lookup table that maps indices to dense vectors
+- **EmbeddingBag**: Computes sums or means of embeddings (efficient for sparse features)
+
+**Key parameters:**
+
+- `num_embeddings`: Size of the vocabulary (number of unique tokens)
+- `embedding_dim`: Dimension of each embedding vector
+- `padding_idx`: Index that outputs zeros (useful for padding tokens)
+
+## Embedding
+
+```{doxygenclass} torch::nn::Embedding
+:members:
+:undoc-members:
+```
+
+```{doxygenclass} torch::nn::EmbeddingImpl
+:members:
+:undoc-members:
+```
+
+**Example:**
+
+```cpp
+auto embedding = torch::nn::Embedding(
+    torch::nn::EmbeddingOptions(10000, 256)  // num_embeddings, embedding_dim
+        .padding_idx(0));
+
+auto indices = torch::tensor({1, 2, 3, 4});
+auto embedded = embedding->forward(indices);  // [4, 256]
+```
+
+## EmbeddingBag
+
+```{doxygenclass} torch::nn::EmbeddingBag
+:members:
+:undoc-members:
+```
+
+```{doxygenclass} torch::nn::EmbeddingBagImpl
+:members:
+:undoc-members:
+```
diff --git a/docs/cpp/source/api/nn/functional.md b/docs/cpp/source/api/nn/functional.md
new file mode 100644
index 0000000000000..0ace995624041
--- /dev/null
+++ b/docs/cpp/source/api/nn/functional.md
@@ -0,0 +1,363 @@
+---
+myst:
+  html_meta:
+    description: Functional API in PyTorch C++ — torch::nn::functional stateless operations for neural networks.
+    keywords: PyTorch, C++, functional, torch::nn::functional, relu, conv2d, linear, softmax
+---
+
+# Functional API
+
+The `torch::nn::functional` namespace provides stateless versions of neural
+network operations. Unlike module classes, functional operations do not hold
+learnable parameters — you pass weights explicitly.
+
+**When to use functional vs modules:**
+
+- Use **modules** (`torch::nn::Conv2d`) when you need learnable parameters
+  managed automatically (training, saving, loading).
+- Use **functional** (`torch::nn::functional::conv2d`) when you already have
+  weights as tensors, or for operations without parameters (e.g., `relu`).
+
+```cpp
+#include <torch/nn/functional.h>
+namespace F = torch::nn::functional;
+
+// Stateless activation — no module needed
+auto output = F::relu(input);
+
+// Convolution with explicit weight tensor
+auto output = F::conv2d(input, weight, F::Conv2dFuncOptions().stride(1).padding(1));
+
+// Softmax along a dimension
+auto probs = F::softmax(logits, F::SoftmaxFuncOptions(/*dim=*/1));
+```
+
+## Activation Functions
+
+```{doxygenfunction} torch::nn::functional::elu
+```
+```{doxygenfunction} torch::nn::functional::selu
+```
+```{doxygenfunction} torch::nn::functional::hardshrink
+```
+```{doxygenfunction} torch::nn::functional::hardtanh
+```
+```{doxygenfunction} torch::nn::functional::leaky_relu
+```
+```{doxygenfunction} torch::nn::functional::logsigmoid
+```
+```{doxygenfunction} torch::nn::functional::glu
+```
+```{doxygenfunction} torch::nn::functional::gelu
+```
+```{doxygenfunction} torch::nn::functional::silu
+```
+```{doxygenfunction} torch::nn::functional::mish
+```
+```{doxygenfunction} torch::nn::functional::prelu
+```
+```{doxygenfunction} torch::nn::functional::relu
+```
+```{doxygenfunction} torch::nn::functional::relu6
+```
+```{doxygenfunction} torch::nn::functional::rrelu
+```
+```{doxygenfunction} torch::nn::functional::celu
+```
+```{doxygenfunction} torch::nn::functional::softplus
+```
+```{doxygenfunction} torch::nn::functional::softshrink
+```
+```{doxygenfunction} torch::nn::functional::softsign
+```
+```{doxygenfunction} torch::nn::functional::tanhshrink
+```
+```{doxygenfunction} torch::nn::functional::threshold
+```
+```{doxygenfunction} torch::nn::functional::softmax
+```
+```{doxygenfunction} torch::nn::functional::softmin
+```
+```{doxygenfunction} torch::nn::functional::log_softmax
+```
+```{doxygenfunction} torch::nn::functional::gumbel_softmax
+```
+
+## Convolution Functions
+
+```{doxygenfunction} torch::nn::functional::conv1d
+```
+```{doxygenfunction} torch::nn::functional::conv2d
+```
+```{doxygenfunction} torch::nn::functional::conv3d
+```
+```{doxygenfunction} torch::nn::functional::conv_transpose1d
+```
+```{doxygenfunction} torch::nn::functional::conv_transpose2d
+```
+```{doxygenfunction} torch::nn::functional::conv_transpose3d
+```
+
+## Pooling Functions
+
+```{doxygenfunction} torch::nn::functional::avg_pool1d
+```
+```{doxygenfunction} torch::nn::functional::avg_pool2d
+```
+```{doxygenfunction} torch::nn::functional::avg_pool3d
+```
+```{doxygenfunction} torch::nn::functional::max_pool1d
+```
+```{doxygenfunction} torch::nn::functional::max_pool2d
+```
+```{doxygenfunction} torch::nn::functional::max_pool3d
+```
+```{doxygenfunction} torch::nn::functional::max_pool1d_with_indices
+```
+```{doxygenfunction} torch::nn::functional::max_pool2d_with_indices
+```
+```{doxygenfunction} torch::nn::functional::max_pool3d_with_indices
+```
+```{doxygenfunction} torch::nn::functional::adaptive_max_pool1d
+```
+```{doxygenfunction} torch::nn::functional::adaptive_max_pool2d
+```
+```{doxygenfunction} torch::nn::functional::adaptive_max_pool3d
+```
+```{doxygenfunction} torch::nn::functional::adaptive_avg_pool1d
+```
+```{doxygenfunction} torch::nn::functional::adaptive_avg_pool2d
+```
+```{doxygenfunction} torch::nn::functional::adaptive_avg_pool3d
+```
+```{doxygenfunction} torch::nn::functional::max_unpool1d
+```
+```{doxygenfunction} torch::nn::functional::max_unpool2d
+```
+```{doxygenfunction} torch::nn::functional::max_unpool3d
+```
+```{doxygenfunction} torch::nn::functional::fractional_max_pool2d
+```
+```{doxygenfunction} torch::nn::functional::fractional_max_pool3d
+```
+```{doxygenfunction} torch::nn::functional::lp_pool1d
+```
+```{doxygenfunction} torch::nn::functional::lp_pool2d
+```
+```{doxygenfunction} torch::nn::functional::lp_pool3d
+```
+
+## Linear Functions
+
+```{doxygenfunction} torch::nn::functional::linear
+```
+```{doxygenfunction} torch::nn::functional::bilinear
+```
+
+## Dropout Functions
+
+```{doxygenfunction} torch::nn::functional::dropout
+```
+```{doxygenfunction} torch::nn::functional::dropout2d
+```
+```{doxygenfunction} torch::nn::functional::dropout3d
+```
+```{doxygenfunction} torch::nn::functional::alpha_dropout
+```
+```{doxygenfunction} torch::nn::functional::feature_alpha_dropout
+```
+
+## Embedding Functions
+
+```{doxygenfunction} torch::nn::functional::one_hot
+```
+```{doxygenfunction} torch::nn::functional::embedding
+```
+```{doxygenfunction} torch::nn::functional::embedding_bag
+```
+
+## Normalization Functions
+
+```{doxygenfunction} torch::nn::functional::batch_norm
+```
+```{doxygenfunction} torch::nn::functional::instance_norm
+```
+```{doxygenfunction} torch::nn::functional::layer_norm
+```
+```{doxygenfunction} torch::nn::functional::group_norm
+```
+```{doxygenfunction} torch::nn::functional::local_response_norm
+```
+```{doxygenfunction} torch::nn::functional::normalize
+```
+
+## Loss Functions
+
+```{doxygenfunction} torch::nn::functional::l1_loss
+```
+```{doxygenfunction} torch::nn::functional::mse_loss
+```
+```{doxygenfunction} torch::nn::functional::binary_cross_entropy
+```
+```{doxygenfunction} torch::nn::functional::binary_cross_entropy_with_logits
+```
+```{doxygenfunction} torch::nn::functional::cross_entropy
+```
+```{doxygenfunction} torch::nn::functional::nll_loss
+```
+```{doxygenfunction} torch::nn::functional::kl_div
+```
+```{doxygenfunction} torch::nn::functional::smooth_l1_loss(const Tensor& input, const Tensor& target, const SmoothL1LossFuncOptions& options)
+```
+```{doxygenfunction} torch::nn::functional::huber_loss
+```
+```{doxygenfunction} torch::nn::functional::hinge_embedding_loss
+```
+```{doxygenfunction} torch::nn::functional::multi_margin_loss
+```
+```{doxygenfunction} torch::nn::functional::cosine_embedding_loss
+```
+```{doxygenfunction} torch::nn::functional::margin_ranking_loss
+```
+```{doxygenfunction} torch::nn::functional::multilabel_margin_loss
+```
+```{doxygenfunction} torch::nn::functional::soft_margin_loss
+```
+```{doxygenfunction} torch::nn::functional::multilabel_soft_margin_loss
+```
+```{doxygenfunction} torch::nn::functional::triplet_margin_loss
+```
+```{doxygenfunction} torch::nn::functional::triplet_margin_with_distance_loss
+```
+```{doxygenfunction} torch::nn::functional::ctc_loss
+```
+```{doxygenfunction} torch::nn::functional::poisson_nll_loss
+```
+
+## Distance Functions
+
+```{doxygenfunction} torch::nn::functional::cosine_similarity
+```
+```{doxygenfunction} torch::nn::functional::pairwise_distance
+```
+```{doxygenfunction} torch::nn::functional::pdist
+```
+
+## Vision Functions
+
+```{doxygenfunction} torch::nn::functional::interpolate
+```
+```{doxygenfunction} torch::nn::functional::affine_grid
+```
+```{doxygenfunction} torch::nn::functional::grid_sample
+```
+```{doxygenfunction} torch::nn::functional::pad
+```
+```{doxygenfunction} torch::nn::functional::pixel_shuffle
+```
+```{doxygenfunction} torch::nn::functional::pixel_unshuffle
+```
+
+## Fold/Unfold
+
+```{doxygenfunction} torch::nn::functional::fold
+```
+```{doxygenfunction} torch::nn::functional::unfold
+```
+
+## Functional Options Structs
+
+Each functional operation that takes configuration uses a corresponding options
+struct. The naming convention is `<Operation>FuncOptions`.
+
+**Activation Options:**
+
+```{doxygentypedef} torch::nn::functional::ELUFuncOptions
+```
+```{doxygentypedef} torch::nn::functional::SELUFuncOptions
+```
+```{doxygentypedef} torch::nn::functional::GLUFuncOptions
+```
+```{doxygentypedef} torch::nn::functional::GELUFuncOptions
+```
+```{doxygentypedef} torch::nn::functional::HardshrinkFuncOptions
+```
+```{doxygentypedef} torch::nn::functional::HardtanhFuncOptions
+```
+```{doxygentypedef} torch::nn::functional::LeakyReLUFuncOptions
+```
+```{doxygentypedef} torch::nn::functional::ReLUFuncOptions
+```
+```{doxygentypedef} torch::nn::functional::ReLU6FuncOptions
+```
+```{doxygentypedef} torch::nn::functional::CELUFuncOptions
+```
+```{doxygentypedef} torch::nn::functional::SoftplusFuncOptions
+```
+```{doxygentypedef} torch::nn::functional::SoftshrinkFuncOptions
+```
+```{doxygentypedef} torch::nn::functional::ThresholdFuncOptions
+```
+
+**Convolution Options:**
+
+```{doxygentypedef} torch::nn::functional::Conv1dFuncOptions
+```
+```{doxygentypedef} torch::nn::functional::Conv2dFuncOptions
+```
+```{doxygentypedef} torch::nn::functional::Conv3dFuncOptions
+```
+```{doxygentypedef} torch::nn::functional::ConvTranspose1dFuncOptions
+```
+```{doxygentypedef} torch::nn::functional::ConvTranspose2dFuncOptions
+```
+```{doxygentypedef} torch::nn::functional::ConvTranspose3dFuncOptions
+```
+
+**Pooling Options:**
+
+```{doxygentypedef} torch::nn::functional::AvgPool1dFuncOptions
+```
+```{doxygentypedef} torch::nn::functional::AvgPool2dFuncOptions
+```
+```{doxygentypedef} torch::nn::functional::AvgPool3dFuncOptions
+```
+```{doxygentypedef} torch::nn::functional::MaxPool1dFuncOptions
+```
+```{doxygentypedef} torch::nn::functional::MaxPool2dFuncOptions
+```
+```{doxygentypedef} torch::nn::functional::MaxPool3dFuncOptions
+```
+```{doxygentypedef} torch::nn::functional::AdaptiveMaxPool1dFuncOptions
+```
+```{doxygentypedef} torch::nn::functional::AdaptiveMaxPool2dFuncOptions
+```
+```{doxygentypedef} torch::nn::functional::AdaptiveMaxPool3dFuncOptions
+```
+```{doxygentypedef} torch::nn::functional::AdaptiveAvgPool1dFuncOptions
+```
+```{doxygentypedef} torch::nn::functional::AdaptiveAvgPool2dFuncOptions
+```
+```{doxygentypedef} torch::nn::functional::AdaptiveAvgPool3dFuncOptions
+```
+
+**Other Options:**
+
+```{doxygentypedef} torch::nn::functional::CosineSimilarityFuncOptions
+```
+```{doxygentypedef} torch::nn::functional::PairwiseDistanceFuncOptions
+```
+```{doxygentypedef} torch::nn::functional::Dropout2dFuncOptions
+```
+```{doxygentypedef} torch::nn::functional::Dropout3dFuncOptions
+```
+```{doxygentypedef} torch::nn::functional::L1LossFuncOptions
+```
+```{doxygentypedef} torch::nn::functional::FoldFuncOptions
+```
+```{doxygentypedef} torch::nn::functional::UnfoldFuncOptions
+```
+```{doxygentypedef} torch::nn::functional::PixelShuffleFuncOptions
+```
+```{doxygentypedef} torch::nn::functional::PixelUnshuffleFuncOptions
+```
diff --git a/docs/cpp/source/api/nn/index.md b/docs/cpp/source/api/nn/index.md
new file mode 100644
index 0000000000000..c5bf65125e12b
--- /dev/null
+++ b/docs/cpp/source/api/nn/index.md
@@ -0,0 +1,94 @@
+---
+myst:
+  html_meta:
+    description: PyTorch C++ neural network modules — torch::nn API for defining and training models.
+    keywords: PyTorch, C++, nn, Module, neural network, torch::nn
+---
+
+# Neural Network Modules (torch::nn)
+
+The `torch::nn` namespace provides neural network building blocks that mirror
+Python's `torch.nn` module. It uses a PIMPL (Pointer to Implementation) pattern
+where user-facing classes like `Conv2d` wrap internal `Conv2dImpl` classes.
+
+**When to use torch::nn:**
+
+- Building neural network models in C++
+- Creating custom layers and modules
+- Porting Python models to C++ for production inference
+- Training models entirely in C++
+
+**Basic usage:**
+
+```cpp
+#include <torch/torch.h>
+
+// Define a simple model
+struct Net : torch::nn::Module {
+    torch::nn::Conv2d conv1{nullptr};
+    torch::nn::Linear fc1{nullptr};
+
+    Net() {
+        conv1 = register_module("conv1", torch::nn::Conv2d(
+            torch::nn::Conv2dOptions(1, 32, 3).stride(1).padding(1)));
+        fc1 = register_module("fc1", torch::nn::Linear(32 * 28 * 28, 10));
+    }
+
+    torch::Tensor forward(torch::Tensor x) {
+        x = torch::relu(conv1->forward(x));
+        x = x.view({-1, 32 * 28 * 28});
+        return fc1->forward(x);
+    }
+};
+
+// Create and use the model
+auto model = std::make_shared<Net>();
+auto input = torch::randn({1, 1, 28, 28});
+auto output = model->forward(input);
+```
+
+## Header Files
+
+- `torch/nn.h` - Main neural network header (includes all modules)
+- `torch/nn/module.h` - Base Module class
+- `torch/nn/modules.h` - All module implementations
+- `torch/nn/options.h` - Options structs for modules
+- `torch/nn/functional.h` - Functional API
+
+## Module Base Class
+
+All neural network modules inherit from `torch::nn::Module`, which provides
+parameter management, serialization, device/dtype conversion, and hooks.
+
+```{doxygenclass} torch::nn::Module
+```
+
+**Key features:**
+
+- `register_module()`: Register submodules for parameter tracking
+- `register_parameter()`: Register learnable parameters
+- `register_buffer()`: Register non-learnable state (e.g., running mean)
+- `parameters()` / `named_parameters()`: Iterate over all parameters
+- `to()`: Move module to a device or convert dtype
+- `train()` / `eval()`: Toggle training/evaluation mode
+- `save()` / `load()`: Serialize and deserialize module state
+
+## Module Categories
+
+```{toctree}
+:maxdepth: 1
+
+containers
+convolution
+pooling
+linear
+activation
+normalization
+dropout
+embedding
+recurrent
+transformer
+loss
+functional
+utilities
+```
diff --git a/docs/cpp/source/api/nn/linear.md b/docs/cpp/source/api/nn/linear.md
new file mode 100644
index 0000000000000..8aa0c55a0ba4a
--- /dev/null
+++ b/docs/cpp/source/api/nn/linear.md
@@ -0,0 +1,84 @@
+---
+myst:
+  html_meta:
+    description: Linear layers in PyTorch C++ — Linear, Bilinear, and Flatten modules.
+    keywords: PyTorch, C++, Linear, Bilinear, Flatten, fully connected, dense layer
+---
+
+# Linear Layers
+
+Linear layers apply affine transformations to input data: `y = xW^T + b`.
+They are the building blocks of fully-connected networks and are used for
+feature transformation, classification heads, and projection layers.
+
+- **Linear**: Standard fully-connected layer
+- **Bilinear**: Bilinear transformation of two inputs
+- **Identity**: Pass-through layer (useful for skip connections)
+- **Flatten/Unflatten**: Reshape tensors between convolutional and linear layers
+
+## Linear
+
+```{doxygenclass} torch::nn::Linear
+:members:
+:undoc-members:
+```
+
+```{doxygenclass} torch::nn::LinearImpl
+:members:
+:undoc-members:
+```
+
+**Example:**
+
+```cpp
+auto linear = torch::nn::Linear(torch::nn::LinearOptions(784, 256).bias(true));
+auto output = linear->forward(input);  // input: [N, 784]
+```
+
+## Bilinear
+
+```{doxygenclass} torch::nn::Bilinear
+:members:
+:undoc-members:
+```
+
+```{doxygenclass} torch::nn::BilinearImpl
+:members:
+:undoc-members:
+```
+
+## Identity
+
+```{doxygenclass} torch::nn::Identity
+:members:
+:undoc-members:
+```
+
+```{doxygenclass} torch::nn::IdentityImpl
+:members:
+:undoc-members:
+```
+
+## Flatten
+
+```{doxygenclass} torch::nn::Flatten
+:members:
+:undoc-members:
+```
+
+```{doxygenclass} torch::nn::FlattenImpl
+:members:
+:undoc-members:
+```
+
+## Unflatten
+
+```{doxygenclass} torch::nn::Unflatten
+:members:
+:undoc-members:
+```
+
+```{doxygenclass} torch::nn::UnflattenImpl
+:members:
+:undoc-members:
+```
diff --git a/docs/cpp/source/api/nn/loss.md b/docs/cpp/source/api/nn/loss.md
new file mode 100644
index 0000000000000..8f1ab82bc234d
--- /dev/null
+++ b/docs/cpp/source/api/nn/loss.md
@@ -0,0 +1,284 @@
+---
+myst:
+  html_meta:
+    description: Loss functions in PyTorch C++ — CrossEntropyLoss, MSELoss, NLLLoss, BCELoss, and more.
+    keywords: PyTorch, C++, loss, CrossEntropyLoss, MSELoss, NLLLoss, BCELoss, L1Loss
+---
+
+# Loss Functions
+
+Loss functions measure how well the model's predictions match the targets.
+The choice of loss function depends on your task type and data characteristics.
+
+**Regression losses:**
+
+- **L1Loss/MSELoss**: Basic regression losses (MAE vs MSE)
+- **SmoothL1Loss/HuberLoss**: Robust to outliers
+
+**Classification losses:**
+
+- **CrossEntropyLoss**: Multi-class classification (combines LogSoftmax + NLLLoss)
+- **NLLLoss**: Negative log likelihood (use with LogSoftmax output)
+- **BCELoss/BCEWithLogitsLoss**: Binary classification
+
+**Specialized losses:**
+
+- **CTCLoss**: Sequence-to-sequence without alignment (speech recognition)
+- **TripletMarginLoss**: Metric learning (similarity/embedding tasks)
+- **CosineEmbeddingLoss**: Similarity learning with cosine distance
+
+## L1Loss
+
+```{doxygenclass} torch::nn::L1Loss
+:members:
+:undoc-members:
+```
+
+```{doxygenstruct} torch::nn::L1LossImpl
+:members:
+:undoc-members:
+```
+
+## MSELoss
+
+```{doxygenclass} torch::nn::MSELoss
+:members:
+:undoc-members:
+```
+
+```{doxygenstruct} torch::nn::MSELossImpl
+:members:
+:undoc-members:
+```
+
+**Example:**
+
+```cpp
+auto loss_fn = torch::nn::MSELoss();
+auto loss = loss_fn->forward(predictions, targets);
+```
+
+## CrossEntropyLoss
+
+```{doxygenclass} torch::nn::CrossEntropyLoss
+:members:
+:undoc-members:
+```
+
+```{doxygenstruct} torch::nn::CrossEntropyLossImpl
+:members:
+:undoc-members:
+```
+
+**Example:**
+
+```cpp
+auto loss_fn = torch::nn::CrossEntropyLoss();
+auto logits = torch::randn({32, 10});  // [batch, num_classes]
+auto targets = torch::randint(0, 10, {32});  // [batch]
+auto loss = loss_fn->forward(logits, targets);
+```
+
+## NLLLoss
+
+```{doxygenclass} torch::nn::NLLLoss
+:members:
+:undoc-members:
+```
+
+```{doxygenstruct} torch::nn::NLLLossImpl
+:members:
+:undoc-members:
+```
+
+## BCELoss
+
+```{doxygenclass} torch::nn::BCELoss
+:members:
+:undoc-members:
+```
+
+```{doxygenstruct} torch::nn::BCELossImpl
+:members:
+:undoc-members:
+```
+
+## BCEWithLogitsLoss
+
+```{doxygenclass} torch::nn::BCEWithLogitsLoss
+:members:
+:undoc-members:
+```
+
+```{doxygenstruct} torch::nn::BCEWithLogitsLossImpl
+:members:
+:undoc-members:
+```
+
+## HuberLoss
+
+```{doxygenclass} torch::nn::HuberLoss
+:members:
+:undoc-members:
+```
+
+```{doxygenstruct} torch::nn::HuberLossImpl
+:members:
+:undoc-members:
+```
+
+## SmoothL1Loss
+
+```{doxygenclass} torch::nn::SmoothL1Loss
+:members:
+:undoc-members:
+```
+
+```{doxygenstruct} torch::nn::SmoothL1LossImpl
+:members:
+:undoc-members:
+```
+
+## KLDivLoss
+
+```{doxygenclass} torch::nn::KLDivLoss
+:members:
+:undoc-members:
+```
+
+```{doxygenstruct} torch::nn::KLDivLossImpl
+:members:
+:undoc-members:
+```
+
+## CTCLoss
+
+```{doxygenclass} torch::nn::CTCLoss
+:members:
+:undoc-members:
+```
+
+```{doxygenstruct} torch::nn::CTCLossImpl
+:members:
+:undoc-members:
+```
+
+## PoissonNLLLoss
+
+```{doxygenclass} torch::nn::PoissonNLLLoss
+:members:
+:undoc-members:
+```
+
+```{doxygenstruct} torch::nn::PoissonNLLLossImpl
+:members:
+:undoc-members:
+```
+
+## MarginRankingLoss
+
+```{doxygenclass} torch::nn::MarginRankingLoss
+:members:
+:undoc-members:
+```
+
+```{doxygenstruct} torch::nn::MarginRankingLossImpl
+:members:
+:undoc-members:
+```
+
+## HingeEmbeddingLoss
+
+```{doxygenclass} torch::nn::HingeEmbeddingLoss
+:members:
+:undoc-members:
+```
+
+```{doxygenstruct} torch::nn::HingeEmbeddingLossImpl
+:members:
+:undoc-members:
+```
+
+## CosineEmbeddingLoss
+
+```{doxygenclass} torch::nn::CosineEmbeddingLoss
+:members:
+:undoc-members:
+```
+
+```{doxygenstruct} torch::nn::CosineEmbeddingLossImpl
+:members:
+:undoc-members:
+```
+
+## MultiMarginLoss
+
+```{doxygenclass} torch::nn::MultiMarginLoss
+:members:
+:undoc-members:
+```
+
+```{doxygenstruct} torch::nn::MultiMarginLossImpl
+:members:
+:undoc-members:
+```
+
+## MultiLabelMarginLoss
+
+```{doxygenclass} torch::nn::MultiLabelMarginLoss
+:members:
+:undoc-members:
+```
+
+```{doxygenstruct} torch::nn::MultiLabelMarginLossImpl
+:members:
+:undoc-members:
+```
+
+## MultiLabelSoftMarginLoss
+
+```{doxygenclass} torch::nn::MultiLabelSoftMarginLoss
+:members:
+:undoc-members:
+```
+
+```{doxygenstruct} torch::nn::MultiLabelSoftMarginLossImpl
+:members:
+:undoc-members:
+```
+
+## SoftMarginLoss
+
+```{doxygenclass} torch::nn::SoftMarginLoss
+:members:
+:undoc-members:
+```
+
+```{doxygenstruct} torch::nn::SoftMarginLossImpl
+:members:
+:undoc-members:
+```
+
+## TripletMarginLoss
+
+```{doxygenclass} torch::nn::TripletMarginLoss
+:members:
+:undoc-members:
+```
+
+```{doxygenstruct} torch::nn::TripletMarginLossImpl
+:members:
+:undoc-members:
+```
+
+## TripletMarginWithDistanceLoss
+
+```{doxygenclass} torch::nn::TripletMarginWithDistanceLoss
+:members:
+:undoc-members:
+```
+
+```{doxygenstruct} torch::nn::TripletMarginWithDistanceLossImpl
+:members:
+:undoc-members:
+```
diff --git a/docs/cpp/source/api/nn/normalization.md b/docs/cpp/source/api/nn/normalization.md
new file mode 100644
index 0000000000000..3ccfe99957e06
--- /dev/null
+++ b/docs/cpp/source/api/nn/normalization.md
@@ -0,0 +1,142 @@
+---
+myst:
+  html_meta:
+    description: Normalization layers in PyTorch C++ — BatchNorm, LayerNorm, GroupNorm, InstanceNorm, and LocalResponseNorm.
+    keywords: PyTorch, C++, normalization, BatchNorm, LayerNorm, GroupNorm, InstanceNorm
+---
+
+# Normalization Layers
+
+Normalization layers stabilize and accelerate training by normalizing intermediate
+activations. They help with gradient flow and allow higher learning rates.
+
+- **BatchNorm**: Normalizes across batch dimension; most common in CNNs
+- **InstanceNorm**: Normalizes each sample independently; popular in style transfer
+- **LayerNorm**: Normalizes across feature dimension; standard in transformers
+- **GroupNorm**: Normalizes within groups of channels; works with small batches
+- **LocalResponseNorm**: Lateral inhibition inspired by neuroscience (less common today)
+
+## BatchNorm1d / BatchNorm2d / BatchNorm3d
+
+```{doxygenclass} torch::nn::BatchNorm1d
+:members:
+:undoc-members:
+```
+
+```{doxygenclass} torch::nn::BatchNorm1dImpl
+:members:
+:undoc-members:
+```
+
+```{doxygenclass} torch::nn::BatchNorm2d
+:members:
+:undoc-members:
+```
+
+```{doxygenclass} torch::nn::BatchNorm2dImpl
+:members:
+:undoc-members:
+```
+
+```{doxygenclass} torch::nn::BatchNorm3d
+:members:
+:undoc-members:
+```
+
+```{doxygenclass} torch::nn::BatchNorm3dImpl
+:members:
+:undoc-members:
+```
+
+**Example:**
+
+```cpp
+auto bn = torch::nn::BatchNorm2d(
+    torch::nn::BatchNorm2dOptions(64)  // num_features
+        .eps(1e-5)
+        .momentum(0.1)
+        .affine(true)
+        .track_running_stats(true));
+```
+
+## InstanceNorm1d / InstanceNorm2d / InstanceNorm3d
+
+```{doxygenclass} torch::nn::InstanceNorm1d
+:members:
+:undoc-members:
+```
+
+```{doxygenclass} torch::nn::InstanceNorm1dImpl
+:members:
+:undoc-members:
+```
+
+```{doxygenclass} torch::nn::InstanceNorm2d
+:members:
+:undoc-members:
+```
+
+```{doxygenclass} torch::nn::InstanceNorm2dImpl
+:members:
+:undoc-members:
+```
+
+```{doxygenclass} torch::nn::InstanceNorm3d
+:members:
+:undoc-members:
+```
+
+```{doxygenclass} torch::nn::InstanceNorm3dImpl
+:members:
+:undoc-members:
+```
+
+## LayerNorm
+
+```{doxygenclass} torch::nn::LayerNorm
+:members:
+:undoc-members:
+```
+
+```{doxygenclass} torch::nn::LayerNormImpl
+:members:
+:undoc-members:
+```
+
+**Example:**
+
+```cpp
+auto ln = torch::nn::LayerNorm(
+    torch::nn::LayerNormOptions({768}));  // normalized_shape
+```
+
+## GroupNorm
+
+```{doxygenclass} torch::nn::GroupNorm
+:members:
+:undoc-members:
+```
+
+```{doxygenclass} torch::nn::GroupNormImpl
+:members:
+:undoc-members:
+```
+
+**Example:**
+
+```cpp
+auto gn = torch::nn::GroupNorm(
+    torch::nn::GroupNormOptions(32, 256));  // num_groups, num_channels
+```
+
+## LocalResponseNorm
+
+```{doxygenclass} torch::nn::LocalResponseNorm
+:members:
+:undoc-members:
+```
+
+```{doxygenclass} torch::nn::LocalResponseNormImpl
+:members:
+:undoc-members:
+```
diff --git a/docs/cpp/source/api/nn/pooling.md b/docs/cpp/source/api/nn/pooling.md
new file mode 100644
index 0000000000000..679b24833d7ab
--- /dev/null
+++ b/docs/cpp/source/api/nn/pooling.md
@@ -0,0 +1,262 @@
+---
+myst:
+  html_meta:
+    description: Pooling layers in PyTorch C++ — MaxPool, AvgPool, AdaptiveMaxPool, AdaptiveAvgPool, and LPPool.
+    keywords: PyTorch, C++, pooling, MaxPool2d, AvgPool2d, AdaptiveAvgPool2d, max pooling
+---
+
+# Pooling Layers
+
+Pooling layers reduce spatial dimensions by aggregating values in local regions,
+providing translation invariance and reducing computational cost in deeper layers.
+
+- **MaxPool**: Takes the maximum value in each pooling window (preserves strong features)
+- **AvgPool**: Takes the average value in each pooling window (smoother downsampling)
+- **AdaptivePool**: Automatically calculates kernel size to produce a target output size
+- **FractionalMaxPool**: Randomized pooling with fractional output size
+- **MaxUnpool**: Computes the partial inverse of MaxPool using stored indices
+- **LPPool**: Power-average pooling (generalization of avg/max pooling)
+
+## MaxPool1d / MaxPool2d / MaxPool3d
+
+```{doxygenclass} torch::nn::MaxPool1d
+:members:
+:undoc-members:
+```
+
+```{doxygenclass} torch::nn::MaxPool1dImpl
+:members:
+:undoc-members:
+```
+
+```{doxygenclass} torch::nn::MaxPool2d
+:members:
+:undoc-members:
+```
+
+```{doxygenclass} torch::nn::MaxPool2dImpl
+:members:
+:undoc-members:
+```
+
+```{doxygenclass} torch::nn::MaxPool3d
+:members:
+:undoc-members:
+```
+
+```{doxygenclass} torch::nn::MaxPool3dImpl
+:members:
+:undoc-members:
+```
+
+**Example:**
+
+```cpp
+auto pool = torch::nn::MaxPool2d(
+    torch::nn::MaxPool2dOptions(2).stride(2));
+```
+
+## AvgPool1d / AvgPool2d / AvgPool3d
+
+```{doxygenclass} torch::nn::AvgPool1d
+:members:
+:undoc-members:
+```
+
+```{doxygenclass} torch::nn::AvgPool1dImpl
+:members:
+:undoc-members:
+```
+
+```{doxygenclass} torch::nn::AvgPool2d
+:members:
+:undoc-members:
+```
+
+```{doxygenclass} torch::nn::AvgPool2dImpl
+:members:
+:undoc-members:
+```
+
+```{doxygenclass} torch::nn::AvgPool3d
+:members:
+:undoc-members:
+```
+
+```{doxygenclass} torch::nn::AvgPool3dImpl
+:members:
+:undoc-members:
+```
+
+## AdaptiveAvgPool1d / AdaptiveAvgPool2d / AdaptiveAvgPool3d
+
+```{doxygenclass} torch::nn::AdaptiveAvgPool1d
+:members:
+:undoc-members:
+```
+
+```{doxygenclass} torch::nn::AdaptiveAvgPool1dImpl
+:members:
+:undoc-members:
+```
+
+```{doxygenclass} torch::nn::AdaptiveAvgPool2d
+:members:
+:undoc-members:
+```
+
+```{doxygenclass} torch::nn::AdaptiveAvgPool2dImpl
+:members:
+:undoc-members:
+```
+
+```{doxygenclass} torch::nn::AdaptiveAvgPool3d
+:members:
+:undoc-members:
+```
+
+```{doxygenclass} torch::nn::AdaptiveAvgPool3dImpl
+:members:
+:undoc-members:
+```
+
+**Example:**
+
+```cpp
+// Output will always be 7x7 regardless of input size
+auto adaptive_pool = torch::nn::AdaptiveAvgPool2d(
+    torch::nn::AdaptiveAvgPool2dOptions({7, 7}));
+```
+
+## AdaptiveMaxPool1d / AdaptiveMaxPool2d / AdaptiveMaxPool3d
+
+```{doxygenclass} torch::nn::AdaptiveMaxPool1d
+:members:
+:undoc-members:
+```
+
+```{doxygenclass} torch::nn::AdaptiveMaxPool1dImpl
+:members:
+:undoc-members:
+```
+
+```{doxygenclass} torch::nn::AdaptiveMaxPool2d
+:members:
+:undoc-members:
+```
+
+```{doxygenclass} torch::nn::AdaptiveMaxPool2dImpl
+:members:
+:undoc-members:
+```
+
+```{doxygenclass} torch::nn::AdaptiveMaxPool3d
+:members:
+:undoc-members:
+```
+
+```{doxygenclass} torch::nn::AdaptiveMaxPool3dImpl
+:members:
+:undoc-members:
+```
+
+## FractionalMaxPool2d / FractionalMaxPool3d
+
+```{doxygenclass} torch::nn::FractionalMaxPool2d
+:members:
+:undoc-members:
+```
+
+```{doxygenclass} torch::nn::FractionalMaxPool2dImpl
+:members:
+:undoc-members:
+```
+
+```{doxygenclass} torch::nn::FractionalMaxPool3d
+:members:
+:undoc-members:
+```
+
+```{doxygenclass} torch::nn::FractionalMaxPool3dImpl
+:members:
+:undoc-members:
+```
+
+## MaxUnpool1d / MaxUnpool2d / MaxUnpool3d
+
+Computes a partial inverse of `MaxPool`, using the indices of the maximum
+values computed during pooling to place values back into unpooled positions.
+
+```{doxygenclass} torch::nn::MaxUnpool1d
+:members:
+:undoc-members:
+```
+
+```{doxygenclass} torch::nn::MaxUnpool1dImpl
+:members:
+:undoc-members:
+```
+
+```{doxygenclass} torch::nn::MaxUnpool2d
+:members:
+:undoc-members:
+```
+
+```{doxygenclass} torch::nn::MaxUnpool2dImpl
+:members:
+:undoc-members:
+```
+
+```{doxygenclass} torch::nn::MaxUnpool3d
+:members:
+:undoc-members:
+```
+
+```{doxygenclass} torch::nn::MaxUnpool3dImpl
+:members:
+:undoc-members:
+```
+
+**Example:**
+
+```cpp
+auto pool = torch::nn::MaxPool2d(
+    torch::nn::MaxPool2dOptions(2).stride(2).return_indices(true));
+auto unpool = torch::nn::MaxUnpool2d(
+    torch::nn::MaxUnpoolOptions<2>(2).stride(2));
+
+auto [output, indices] = pool->forward_with_indices(input);
+auto reconstructed = unpool->forward(output, indices);
+```
+
+## LPPool1d / LPPool2d / LPPool3d
+
+```{doxygenclass} torch::nn::LPPool1d
+:members:
+:undoc-members:
+```
+
+```{doxygenclass} torch::nn::LPPool1dImpl
+:members:
+:undoc-members:
+```
+
+```{doxygenclass} torch::nn::LPPool2d
+:members:
+:undoc-members:
+```
+
+```{doxygenclass} torch::nn::LPPool2dImpl
+:members:
+:undoc-members:
+```
+
+```{doxygenclass} torch::nn::LPPool3d
+:members:
+:undoc-members:
+```
+
+```{doxygenclass} torch::nn::LPPool3dImpl
+:members:
+:undoc-members:
+```
diff --git a/docs/cpp/source/api/nn/recurrent.md b/docs/cpp/source/api/nn/recurrent.md
new file mode 100644
index 0000000000000..f286ebb6b52b5
--- /dev/null
+++ b/docs/cpp/source/api/nn/recurrent.md
@@ -0,0 +1,125 @@
+---
+myst:
+  html_meta:
+    description: Recurrent layers in PyTorch C++ — RNN, LSTM, and GRU modules for sequence modeling.
+    keywords: PyTorch, C++, RNN, LSTM, GRU, recurrent, sequence, hidden state
+---
+
+# Recurrent Layers
+
+Recurrent layers process sequential data by maintaining hidden state across time steps.
+They are essential for tasks involving sequences: language modeling, speech recognition,
+time series prediction, and more.
+
+- **RNN**: Basic recurrent layer (simple but prone to vanishing gradients)
+- **LSTM**: Long Short-Term Memory (gated architecture, handles long-range dependencies)
+- **GRU**: Gated Recurrent Unit (simpler than LSTM, often similar performance)
+- **Cell variants**: Single-step versions for custom loop implementations
+
+**Key parameters:**
+
+- `input_size`: Number of features in input
+- `hidden_size`: Number of features in hidden state
+- `num_layers`: Number of stacked recurrent layers
+- `batch_first`: If true, input shape is `[batch, seq, features]`
+- `bidirectional`: Process sequence in both directions
+
+## RNN
+
+```{doxygenclass} torch::nn::RNN
+:members:
+:undoc-members:
+```
+
+```{doxygenclass} torch::nn::RNNImpl
+:members:
+:undoc-members:
+```
+
+**Example:**
+
+```cpp
+auto rnn = torch::nn::RNN(
+    torch::nn::RNNOptions(128, 256)  // input_size, hidden_size
+        .num_layers(2)
+        .batch_first(true)
+        .bidirectional(false));
+
+auto input = torch::randn({32, 10, 128});  // [batch, seq_len, input_size]
+auto [output, hidden] = rnn->forward(input);
+```
+
+## LSTM
+
+```{doxygenclass} torch::nn::LSTM
+:members:
+:undoc-members:
+```
+
+```{doxygenclass} torch::nn::LSTMImpl
+:members:
+:undoc-members:
+```
+
+**Example:**
+
+```cpp
+auto lstm = torch::nn::LSTM(
+    torch::nn::LSTMOptions(128, 256)
+        .num_layers(2)
+        .batch_first(true)
+        .dropout(0.1)
+        .bidirectional(true));
+
+auto input = torch::randn({32, 10, 128});
+auto [output, state] = lstm->forward(input);
+auto [h_n, c_n] = state;  // hidden state, cell state
+```
+
+## GRU
+
+```{doxygenclass} torch::nn::GRU
+:members:
+:undoc-members:
+```
+
+```{doxygenclass} torch::nn::GRUImpl
+:members:
+:undoc-members:
+```
+
+## RNNCell
+
+```{doxygenclass} torch::nn::RNNCell
+:members:
+:undoc-members:
+```
+
+```{doxygenclass} torch::nn::RNNCellImpl
+:members:
+:undoc-members:
+```
+
+## LSTMCell
+
+```{doxygenclass} torch::nn::LSTMCell
+:members:
+:undoc-members:
+```
+
+```{doxygenclass} torch::nn::LSTMCellImpl
+:members:
+:undoc-members:
+```
+
+## GRUCell
+
+```{doxygenclass} torch::nn::GRUCell
+:members:
+:undoc-members:
+```
+
+```{doxygenclass} torch::nn::GRUCellImpl
+:members:
+:undoc-members:
+```
diff --git a/docs/cpp/source/api/nn/transformer.md b/docs/cpp/source/api/nn/transformer.md
new file mode 100644
index 0000000000000..664564d4b3519
--- /dev/null
+++ b/docs/cpp/source/api/nn/transformer.md
@@ -0,0 +1,112 @@
+---
+myst:
+  html_meta:
+    description: Transformer layers in PyTorch C++ — Transformer, TransformerEncoder, TransformerDecoder, and MultiheadAttention.
+    keywords: PyTorch, C++, Transformer, TransformerEncoder, TransformerDecoder, MultiheadAttention, attention
+---
+
+# Transformer Layers
+
+Transformer layers use self-attention mechanisms to process sequences in parallel,
+enabling efficient training on long sequences. They are the foundation of modern
+NLP models (BERT, GPT) and increasingly used in vision and other domains.
+
+- **Transformer**: Complete encoder-decoder architecture
+- **TransformerEncoder/Decoder**: Standalone encoder or decoder stacks
+- **TransformerEncoderLayer/DecoderLayer**: Individual transformer blocks
+- **MultiheadAttention**: Core attention mechanism used throughout
+
+**Key parameters:**
+
+- `d_model`: Dimension of the model (embedding dimension)
+- `nhead`: Number of attention heads
+- `num_encoder_layers/num_decoder_layers`: Number of stacked layers
+- `dim_feedforward`: Dimension of feedforward network
+- `dropout`: Dropout rate for regularization
+
+## Transformer
+
+Complete encoder-decoder transformer architecture.
+
+```{doxygenclass} torch::nn::Transformer
+:members:
+:undoc-members:
+```
+
+```{doxygenclass} torch::nn::TransformerImpl
+:members:
+:undoc-members:
+```
+
+**Example:**
+
+```cpp
+auto transformer = torch::nn::Transformer(
+    torch::nn::TransformerOptions()
+        .d_model(512)
+        .nhead(8)
+        .num_encoder_layers(6)
+        .num_decoder_layers(6)
+        .dim_feedforward(2048)
+        .dropout(0.1));
+```
+
+## TransformerEncoder
+
+Stack of encoder layers for processing source sequences.
+
+```{doxygenclass} torch::nn::TransformerEncoder
+:members:
+:undoc-members:
+```
+
+```{doxygenclass} torch::nn::TransformerEncoderImpl
+:members:
+:undoc-members:
+```
+
+## TransformerDecoder
+
+Stack of decoder layers for generating target sequences.
+
+```{doxygenclass} torch::nn::TransformerDecoder
+:members:
+:undoc-members:
+```
+
+```{doxygenclass} torch::nn::TransformerDecoderImpl
+:members:
+:undoc-members:
+```
+
+## TransformerEncoderLayer
+
+Single encoder layer with self-attention and feedforward network.
+
+```{doxygenclass} torch::nn::TransformerEncoderLayerImpl
+:members:
+:undoc-members:
+```
+
+## TransformerDecoderLayer
+
+Single decoder layer with self-attention, cross-attention, and feedforward network.
+
+```{doxygenclass} TransformerDecoderLayerImpl
+:members:
+:undoc-members:
+```
+
+## MultiheadAttention
+
+Scaled dot-product attention with multiple parallel heads.
+
+```{doxygenclass} torch::nn::MultiheadAttention
+:members:
+:undoc-members:
+```
+
+```{doxygenclass} torch::nn::MultiheadAttentionImpl
+:members:
+:undoc-members:
+```
diff --git a/docs/cpp/source/api/nn/utilities.md b/docs/cpp/source/api/nn/utilities.md
new file mode 100644
index 0000000000000..6bde6de870a0d
--- /dev/null
+++ b/docs/cpp/source/api/nn/utilities.md
@@ -0,0 +1,261 @@
+---
+myst:
+  html_meta:
+    description: Neural network utilities in PyTorch C++ — parameter initialization, module cloning, padding layers, and vision utilities.
+    keywords: PyTorch, C++, nn utilities, init, Cloneable, AnyModule, padding, PixelShuffle, Upsample
+---
+
+# Utilities
+
+Additional utilities for building neural networks: parameter initialization,
+module cloning, type-erased containers, padding layers, and vision utilities.
+
+## Parameter Initialization
+
+The `torch::nn::init` namespace provides functions for initializing module parameters:
+
+```cpp
+#include <torch/nn/init.h>
+
+// Xavier/Glorot initialization
+torch::nn::init::xavier_uniform_(linear->weight);
+torch::nn::init::xavier_normal_(linear->weight);
+
+// Kaiming/He initialization
+torch::nn::init::kaiming_uniform_(conv->weight, /*a=*/0, torch::kFanIn, torch::kReLU);
+torch::nn::init::kaiming_normal_(conv->weight);
+
+// Other initializations
+torch::nn::init::zeros_(linear->bias);
+torch::nn::init::ones_(bn->weight);
+torch::nn::init::constant_(linear->bias, 0.1);
+torch::nn::init::normal_(linear->weight, /*mean=*/0, /*std=*/0.01);
+torch::nn::init::uniform_(linear->weight, /*a=*/-0.1, /*b=*/0.1);
+torch::nn::init::orthogonal_(rnn->weight_hh);
+```
+
+## Cloneable
+
+```{doxygenclass} torch::nn::Cloneable
+```
+
+All `torch::nn` modules inherit from `Cloneable`, enabling deep copies:
+
+```cpp
+auto model = torch::nn::Linear(10, 5);
+auto model_copy = std::dynamic_pointer_cast<torch::nn::LinearImpl>(model->clone());
+```
+
+## AnyModule
+
+`AnyModule` provides type-erased storage for any module, allowing you to
+store heterogeneous modules in a single container.
+
+```{doxygenclass} torch::nn::AnyModule
+:members:
+:undoc-members:
+```
+
+**Example:**
+
+```cpp
+torch::nn::AnyModule any_module(torch::nn::Linear(10, 5));
+auto output = any_module.forward(input);
+```
+
+## Functional
+
+Wraps a function or callable as a module, useful for inserting arbitrary
+functions into a `Sequential` container.
+
+```{doxygenclass} torch::nn::FunctionalImpl
+:members:
+:undoc-members:
+```
+
+## ModuleHolder
+
+```{doxygenclass} torch::nn::ModuleHolder
+:members:
+:undoc-members:
+```
+
+## CosineSimilarity
+
+```{doxygenclass} torch::nn::CosineSimilarity
+:members:
+:undoc-members:
+```
+
+## PairwiseDistance
+
+```{doxygenclass} torch::nn::PairwiseDistance
+:members:
+:undoc-members:
+```
+
+## PackedSequence
+
+```{cpp:class} torch::nn::utils::rnn::PackedSequence
+
+Holds the data and list of `batch_sizes` of a packed sequence.
+All RNN modules accept packed sequences as inputs.
+```
+
+```{cpp:function} const Tensor& data() const
+
+Returns the packed tensor containing all sequence elements.
+```
+
+```{cpp:function} const Tensor& batch_sizes() const
+
+Returns a 1D tensor of batch sizes at each time step.
+```
+
+```{cpp:function} const Tensor& sorted_indices() const
+
+Returns indices used to sort sequences by length (descending).
+```
+
+```{cpp:function} const Tensor& unsorted_indices() const
+
+Returns indices to restore the original sequence order.
+```
+
+```{cpp:function} PackedSequence to(torch::Device device) const
+
+Moves the packed sequence to the specified device.
+```
+
+See also: `torch::nn::utils::rnn::pack_padded_sequence` and
+`torch::nn::utils::rnn::pad_packed_sequence`.
+
+## Padding Layers
+
+### ReflectionPad1d / ReflectionPad2d / ReflectionPad3d
+
+```{doxygenclass} torch::nn::ReflectionPad1d
+:members:
+:undoc-members:
+```
+
+```{doxygenclass} torch::nn::ReflectionPad2d
+:members:
+:undoc-members:
+```
+
+```{doxygenclass} torch::nn::ReflectionPad3d
+:members:
+:undoc-members:
+```
+
+### ReplicationPad1d / ReplicationPad2d / ReplicationPad3d
+
+```{doxygenclass} torch::nn::ReplicationPad1d
+:members:
+:undoc-members:
+```
+
+```{doxygenclass} torch::nn::ReplicationPad2d
+:members:
+:undoc-members:
+```
+
+```{doxygenclass} torch::nn::ReplicationPad3d
+:members:
+:undoc-members:
+```
+
+### ZeroPad1d / ZeroPad2d / ZeroPad3d
+
+```{doxygenclass} torch::nn::ZeroPad1d
+:members:
+:undoc-members:
+```
+
+```{doxygenclass} torch::nn::ZeroPad2d
+:members:
+:undoc-members:
+```
+
+```{doxygenclass} torch::nn::ZeroPad3d
+:members:
+:undoc-members:
+```
+
+### ConstantPad1d / ConstantPad2d / ConstantPad3d
+
+```{doxygenclass} torch::nn::ConstantPad1d
+:members:
+:undoc-members:
+```
+
+```{doxygenclass} torch::nn::ConstantPad2d
+:members:
+:undoc-members:
+```
+
+```{doxygenclass} torch::nn::ConstantPad3d
+:members:
+:undoc-members:
+```
+
+## Vision Layers
+
+### PixelShuffle
+
+```{doxygenclass} torch::nn::PixelShuffle
+:members:
+:undoc-members:
+```
+
+```{doxygenstruct} torch::nn::PixelShuffleOptions
+:members:
+:undoc-members:
+```
+
+### PixelUnshuffle
+
+```{doxygenclass} torch::nn::PixelUnshuffle
+:members:
+:undoc-members:
+```
+
+```{doxygenstruct} torch::nn::PixelUnshuffleOptions
+:members:
+:undoc-members:
+```
+
+### Upsample
+
+```{doxygenclass} torch::nn::Upsample
+:members:
+:undoc-members:
+```
+
+```{doxygenstruct} torch::nn::UpsampleOptions
+:members:
+:undoc-members:
+```
+
+### Fold / Unfold
+
+```{doxygenclass} torch::nn::Fold
+:members:
+:undoc-members:
+```
+
+```{doxygenstruct} torch::nn::FoldOptions
+:members:
+:undoc-members:
+```
+
+```{doxygenclass} torch::nn::Unfold
+:members:
+:undoc-members:
+```
+
+```{doxygenstruct} torch::nn::UnfoldOptions
+:members:
+:undoc-members:
+```
diff --git a/docs/cpp/source/api/optim/adaptive.md b/docs/cpp/source/api/optim/adaptive.md
new file mode 100644
index 0000000000000..3f48534567627
--- /dev/null
+++ b/docs/cpp/source/api/optim/adaptive.md
@@ -0,0 +1,158 @@
+---
+myst:
+  html_meta:
+    description: Adaptive optimizers in PyTorch C++ — Adam, AdamW, Adagrad, and RMSprop.
+    keywords: PyTorch, C++, Adam, AdamW, Adagrad, RMSprop, adaptive optimizer
+---
+
+# Adaptive Learning Rate Optimizers
+
+These optimizers automatically adapt the learning rate for each parameter based
+on historical gradient information. They typically require less hyperparameter
+tuning and work well across a wide range of problems.
+
+## Adam (Adaptive Moment Estimation)
+
+Adam combines the benefits of RMSprop and momentum, maintaining per-parameter
+adaptive learning rates. It's an excellent default choice for most deep learning
+tasks, especially when you want fast convergence with minimal tuning.
+
+**When to use:**
+
+- Transformers and attention-based models
+- Quick prototyping and experimentation
+- When you don't have time for extensive hyperparameter search
+- General-purpose deep learning
+
+**Key parameters:**
+
+- `lr`: Learning rate (typical: 1e-3 to 1e-4)
+- `betas`: Coefficients for running averages (default: {0.9, 0.999})
+- `eps`: Numerical stability term (default: 1e-8)
+- `weight_decay`: L2 regularization (note: applied differently than in SGD)
+
+```{doxygenclass} torch::optim::Adam
+:members:
+:undoc-members:
+```
+
+**Example:**
+
+```cpp
+// Standard Adam configuration
+auto optimizer = torch::optim::Adam(
+    model->parameters(),
+    torch::optim::AdamOptions(1e-3)        // learning rate
+        .betas({0.9, 0.999})               // momentum terms
+        .eps(1e-8)                         // numerical stability
+        .weight_decay(0));                 // L2 penalty
+
+// For transformers, lower learning rate with warmup
+auto optimizer = torch::optim::Adam(
+    model->parameters(),
+    torch::optim::AdamOptions(1e-4)
+        .betas({0.9, 0.98}));              // β2=0.98 for transformers
+```
+
+## AdamW (Adam with Decoupled Weight Decay)
+
+AdamW fixes a subtle issue with Adam's weight decay implementation. In Adam,
+weight decay is coupled with the gradient update, which can lead to suboptimal
+regularization. AdamW decouples weight decay, applying it directly to the
+weights as in SGD.
+
+**When to use:**
+
+- Always prefer AdamW over Adam when using weight decay
+- Training transformers (BERT, GPT, etc.)
+- When you want proper L2 regularization behavior
+
+**Key difference from Adam:**
+
+- In Adam: `weight = weight - lr * (grad + weight_decay * weight)`
+- In AdamW: `weight = weight - lr * grad - lr * weight_decay * weight`
+
+```{doxygenclass} torch::optim::AdamW
+:members:
+:undoc-members:
+```
+
+**Example:**
+
+```cpp
+// AdamW with decoupled weight decay - preferred for transformers
+auto optimizer = torch::optim::AdamW(
+    model->parameters(),
+    torch::optim::AdamWOptions(1e-4)
+        .betas({0.9, 0.999})
+        .weight_decay(0.01));              // Decoupled regularization
+```
+
+## RMSprop (Root Mean Square Propagation)
+
+RMSprop adapts the learning rate by dividing by a running average of recent
+gradient magnitudes. It's particularly effective for recurrent neural networks
+and problems with non-stationary objectives.
+
+**When to use:**
+
+- Training RNNs and LSTMs
+- Non-stationary problems where gradient scale varies significantly
+- Online learning scenarios
+
+**Key parameters:**
+
+- `lr`: Learning rate (typical: 1e-3 to 1e-2)
+- `alpha`: Smoothing constant (default: 0.99)
+- `momentum`: Optional momentum term
+- `centered`: Use centered RMSprop (normalizes by variance)
+
+```{doxygenclass} torch::optim::RMSprop
+:members:
+:undoc-members:
+```
+
+**Example:**
+
+```cpp
+// RMSprop for RNN training
+auto optimizer = torch::optim::RMSprop(
+    model->parameters(),
+    torch::optim::RMSpropOptions(1e-3)
+        .alpha(0.99)                       // smoothing constant
+        .momentum(0.9)                     // optional momentum
+        .centered(true));                  // normalize by variance
+```
+
+## Adagrad (Adaptive Gradient)
+
+Adagrad adapts the learning rate based on the accumulated sum of squared
+gradients. Parameters with frequent updates get smaller learning rates, while
+parameters with infrequent updates get larger rates. This makes it ideal for
+sparse data.
+
+**When to use:**
+
+- NLP tasks with sparse features
+- Embedding layers with infrequent updates
+- Recommendation systems with sparse user/item features
+
+**Limitation:** Learning rate monotonically decreases, which can cause training
+to stop prematurely. For long training runs, consider Adam or RMSprop.
+
+```{doxygenclass} torch::optim::Adagrad
+:members:
+:undoc-members:
+```
+
+**Example:**
+
+```cpp
+// Adagrad for sparse NLP features
+auto optimizer = torch::optim::Adagrad(
+    model->parameters(),
+    torch::optim::AdagradOptions(0.01)
+        .lr_decay(0)                       // learning rate decay
+        .weight_decay(0)
+        .initial_accumulator_value(0));
+```
diff --git a/docs/cpp/source/api/optim/gradient_descent.md b/docs/cpp/source/api/optim/gradient_descent.md
new file mode 100644
index 0000000000000..c72b30171f717
--- /dev/null
+++ b/docs/cpp/source/api/optim/gradient_descent.md
@@ -0,0 +1,49 @@
+---
+myst:
+  html_meta:
+    description: SGD optimizer in PyTorch C++ — stochastic gradient descent with momentum and weight decay.
+    keywords: PyTorch, C++, SGD, gradient descent, momentum, weight decay, optimizer
+---
+
+# Gradient Descent Optimizers
+
+These optimizers use gradient descent with optional enhancements like momentum.
+They are the foundation of neural network training and work well when you can
+afford careful hyperparameter tuning.
+
+## SGD (Stochastic Gradient Descent)
+
+The classic optimization algorithm. SGD with momentum is often the best choice
+for convolutional neural networks when properly tuned. While requiring more
+careful learning rate selection than adaptive methods, it frequently achieves
+the best final accuracy.
+
+**When to use:**
+
+- Training CNNs (ResNet, VGG, etc.) where you want maximum accuracy
+- When you have time for hyperparameter tuning
+- When combined with learning rate schedules (warmup, cosine annealing)
+
+**Key parameters:**
+
+- `lr`: Learning rate (typical: 0.01-0.1 for CNNs)
+- `momentum`: Accelerates convergence (typical: 0.9)
+- `weight_decay`: L2 regularization coefficient
+- `nesterov`: Use Nesterov momentum (often improves convergence)
+
+```{doxygenclass} torch::optim::SGD
+:members:
+:undoc-members:
+```
+
+**Example:**
+
+```cpp
+// Standard SGD with momentum - good for CNNs
+auto optimizer = torch::optim::SGD(
+    model->parameters(),
+    torch::optim::SGDOptions(0.01)   // learning rate
+        .momentum(0.9)               // momentum factor
+        .weight_decay(1e-4)          // L2 regularization
+        .nesterov(true));            // Nesterov momentum
+```
diff --git a/docs/cpp/source/api/optim/index.md b/docs/cpp/source/api/optim/index.md
new file mode 100644
index 0000000000000..e46c75be09f55
--- /dev/null
+++ b/docs/cpp/source/api/optim/index.md
@@ -0,0 +1,116 @@
+---
+myst:
+  html_meta:
+    description: PyTorch C++ optimizer API — SGD, Adam, and other optimizers for training neural networks.
+    keywords: PyTorch, C++, optimizer, optim, SGD, Adam, training
+---
+
+# Optimizers (torch::optim)
+
+The `torch::optim` namespace provides optimization algorithms for
+training neural networks. These optimizers update model parameters based
+on computed gradients to minimize the loss function.
+
+**When to use torch::optim:**
+
+- When training neural networks with gradient descent
+- When you need different optimization strategies (SGD, Adam, etc.)
+- When implementing learning rate schedules
+
+**Basic usage:**
+
+```cpp
+#include <torch/torch.h>
+
+// Create model and optimizer
+auto model = std::make_shared<Net>();
+auto optimizer = torch::optim::Adam(
+    model->parameters(),
+    torch::optim::AdamOptions(1e-3));
+
+// Training loop
+for (auto& batch : *data_loader) {
+    optimizer.zero_grad();                    // Clear gradients
+    auto loss = loss_fn(model->forward(batch.data), batch.target);
+    loss.backward();                          // Compute gradients
+    optimizer.step();                         // Update parameters
+}
+```
+
+## Header Files
+
+- `torch/csrc/api/include/torch/optim.h` - Main optim header
+- `torch/csrc/api/include/torch/optim/optimizer.h` - Optimizer base class
+- `torch/csrc/api/include/torch/optim/sgd.h` - SGD optimizer
+- `torch/csrc/api/include/torch/optim/adam.h` - Adam optimizer
+
+## Optimizer Base Class
+
+All optimizers inherit from the `Optimizer` base class, which provides common
+functionality for parameter updates, gradient zeroing, and state management.
+
+```{doxygenclass} torch::optim::Optimizer
+:members:
+:undoc-members:
+```
+
+### OptimizerOptions
+
+```{doxygenclass} torch::optim::OptimizerOptions
+:members:
+:undoc-members:
+```
+
+### OptimizerParamGroup
+
+```{doxygenclass} torch::optim::OptimizerParamGroup
+:members:
+:undoc-members:
+```
+
+### OptimizerParamState
+
+```{doxygenclass} torch::optim::OptimizerParamState
+:members:
+:undoc-members:
+```
+
+## Choosing an Optimizer
+
+Selecting the right optimizer depends on your model architecture, dataset, and
+training requirements:
+
+```{list-table}
+:widths: 20 40 40
+:header-rows: 1
+
+* - Optimizer
+  - Best For
+  - Trade-offs
+* - **SGD + Momentum**
+  - CNNs, well-understood problems, when you can tune hyperparameters
+  - Requires careful learning rate tuning; often achieves best final accuracy
+* - **Adam/AdamW**
+  - General-purpose, transformers, quick prototyping
+  - Works well out-of-the-box; AdamW preferred with weight decay
+* - **RMSprop**
+  - RNNs, non-stationary objectives
+  - Good for recurrent architectures; handles varying gradient scales
+* - **Adagrad**
+  - Sparse data (NLP, embeddings)
+  - Learning rate decreases over time; good for infrequent features
+* - **LBFGS**
+  - Small models, fine-tuning, convex problems
+  - Memory-intensive; requires closure function
+```
+
+## Optimizer Categories
+
+```{toctree}
+:maxdepth: 1
+
+gradient_descent
+adaptive
+second_order
+schedulers
+```
diff --git a/docs/cpp/source/api/optim/schedulers.md b/docs/cpp/source/api/optim/schedulers.md
new file mode 100644
index 0000000000000..814ff45e2fbaf
--- /dev/null
+++ b/docs/cpp/source/api/optim/schedulers.md
@@ -0,0 +1,152 @@
+---
+myst:
+  html_meta:
+    description: Learning rate schedulers in PyTorch C++ — StepLR, ExponentialLR, and other LR scheduling policies.
+    keywords: PyTorch, C++, learning rate, scheduler, StepLR, ExponentialLR, LRScheduler
+---
+
+# Learning Rate Schedulers
+
+Learning rate schedulers adjust the learning rate during training, which often
+improves convergence and final accuracy. Common strategies include:
+
+- **Step decay**: Reduce LR by a factor every N epochs
+- **Exponential decay**: Multiply LR by gamma each epoch
+- **Cosine annealing**: Smoothly decrease LR following a cosine curve
+- **Warmup**: Gradually increase LR at the start of training
+
+## LRScheduler Base Class
+
+```{doxygenclass} torch::optim::LRScheduler
+:members:
+:undoc-members:
+```
+
+## StepLR
+
+Decays the learning rate by `gamma` every `step_size` epochs. This is the
+simplest and most commonly used scheduler.
+
+```{doxygenclass} torch::optim::StepLR
+:members:
+:undoc-members:
+```
+
+**Example:**
+
+```cpp
+auto optimizer = torch::optim::SGD(
+    model->parameters(),
+    torch::optim::SGDOptions(0.1));
+
+// Reduce LR by 10x every 30 epochs
+auto scheduler = torch::optim::StepLR(
+    optimizer,
+    /*step_size=*/30,
+    /*gamma=*/0.1);
+
+for (int epoch = 0; epoch < 90; ++epoch) {
+    train_one_epoch(model, optimizer, data_loader);
+    scheduler.step();
+    // LR: 0.1 (epochs 0-29), 0.01 (30-59), 0.001 (60-89)
+}
+```
+
+## ReduceLROnPlateau
+
+Reduces the learning rate when a metric has stopped improving. Useful when
+you want the scheduler to respond to validation loss rather than follow a
+fixed schedule.
+
+```{doxygenclass} torch::optim::ReduceLROnPlateauScheduler
+:members:
+:undoc-members:
+```
+
+## ExponentialLR
+
+Decays the learning rate by `gamma` every epoch. Provides smoother decay than
+StepLR but may be slower to reduce the learning rate.
+
+**Example:**
+
+```cpp
+auto optimizer = torch::optim::Adam(
+    model->parameters(),
+    torch::optim::AdamOptions(1e-3));
+
+// Reduce LR by 5% each epoch
+auto scheduler = torch::optim::ExponentialLR(
+    optimizer,
+    /*gamma=*/0.95);
+
+for (int epoch = 0; epoch < num_epochs; ++epoch) {
+    train_one_epoch(model, optimizer, data_loader);
+    scheduler.step();
+}
+```
+
+## Complete Training Example
+
+Here's a complete example showing optimizer usage with learning rate scheduling:
+
+```cpp
+#include <torch/torch.h>
+
+struct Net : torch::nn::Module {
+    Net() {
+        fc1 = register_module("fc1", torch::nn::Linear(784, 256));
+        fc2 = register_module("fc2", torch::nn::Linear(256, 10));
+    }
+
+    torch::Tensor forward(torch::Tensor x) {
+        x = torch::relu(fc1->forward(x.view({-1, 784})));
+        return fc2->forward(x);
+    }
+
+    torch::nn::Linear fc1{nullptr}, fc2{nullptr};
+};
+
+int main() {
+    // Create model
+    auto model = std::make_shared<Net>();
+
+    // Create optimizer with weight decay
+    auto optimizer = torch::optim::AdamW(
+        model->parameters(),
+        torch::optim::AdamWOptions(1e-3)
+            .weight_decay(0.01));
+
+    // Learning rate scheduler
+    auto scheduler = torch::optim::StepLR(optimizer, 10, 0.5);
+
+    // Loss function
+    auto loss_fn = torch::nn::CrossEntropyLoss();
+
+    // Training loop
+    for (int epoch = 0; epoch < 30; ++epoch) {
+        model->train();
+        double epoch_loss = 0.0;
+
+        for (auto& batch : *train_loader) {
+            optimizer.zero_grad();
+
+            auto output = model->forward(batch.data);
+            auto loss = loss_fn(output, batch.target);
+
+            loss.backward();
+            optimizer.step();
+
+            epoch_loss += loss.item<double>();
+        }
+
+        scheduler.step();
+        std::cout << "Epoch " << epoch
+                  << " Loss: " << epoch_loss
+                  << " LR: " << scheduler.get_last_lr()[0]
+                  << std::endl;
+    }
+
+    return 0;
+}
+```
diff --git a/docs/cpp/source/api/optim/second_order.md b/docs/cpp/source/api/optim/second_order.md
new file mode 100644
index 0000000000000..0f93621aecf41
--- /dev/null
+++ b/docs/cpp/source/api/optim/second_order.md
@@ -0,0 +1,60 @@
+---
+myst:
+  html_meta:
+    description: Second-order optimizers in PyTorch C++ — LBFGS optimizer for full-batch optimization.
+    keywords: PyTorch, C++, LBFGS, second-order, optimizer, full-batch
+---
+
+# Second-Order Optimizers
+
+Second-order methods use curvature information (Hessian or its approximations)
+to make better optimization steps. They can converge faster but are more
+computationally expensive and memory-intensive.
+
+## LBFGS (Limited-memory Broyden-Fletcher-Goldfarb-Shanno)
+
+LBFGS is a quasi-Newton method that approximates the inverse Hessian using
+gradient history. It can converge much faster than first-order methods for
+smooth, convex-like loss surfaces.
+
+**When to use:**
+
+- Small models where memory isn't a concern
+- Fine-tuning pre-trained models
+- Convex or near-convex optimization problems
+- Full-batch training (not mini-batch)
+
+**Key parameters:**
+
+- `lr`: Learning rate (often 1.0 for LBFGS)
+- `max_iter`: Maximum iterations per step
+- `history_size`: Number of past gradients to store
+
+**Important:** LBFGS requires a closure function that recomputes the loss.
+
+```{doxygenclass} torch::optim::LBFGS
+:members:
+:undoc-members:
+```
+
+**Example:**
+
+```cpp
+auto optimizer = torch::optim::LBFGS(
+    model->parameters(),
+    torch::optim::LBFGSOptions(1.0)
+        .max_iter(20)
+        .history_size(10));
+
+// LBFGS requires a closure that recomputes the model
+for (int epoch = 0; epoch < num_epochs; ++epoch) {
+    auto closure = [&]() {
+        optimizer.zero_grad();
+        auto output = model->forward(data);
+        auto loss = loss_fn(output, target);
+        loss.backward();
+        return loss;
+    };
+    optimizer.step(closure);
+}
+```
diff --git a/docs/cpp/source/api/serialize/archives.md b/docs/cpp/source/api/serialize/archives.md
new file mode 100644
index 0000000000000..2d50e3db08c9f
--- /dev/null
+++ b/docs/cpp/source/api/serialize/archives.md
@@ -0,0 +1,72 @@
+---
+myst:
+  html_meta:
+    description: Serialization archives in PyTorch C++ — InputArchive and OutputArchive for model serialization.
+    keywords: PyTorch, C++, InputArchive, OutputArchive, serialization, archive
+---
+
+# Archives
+
+Archives provide a lower-level interface for serialization, allowing you to
+save multiple values to a single file with named keys.
+
+## OutputArchive
+
+```{doxygenclass} torch::serialize::OutputArchive
+:members:
+:undoc-members:
+```
+
+**Example:**
+
+```cpp
+torch::serialize::OutputArchive archive;
+archive.write("tensor1", tensor1);
+archive.write("tensor2", tensor2);
+archive.save_to("multi_tensor.pt");
+```
+
+## InputArchive
+
+```{doxygenclass} torch::serialize::InputArchive
+:members:
+:undoc-members:
+```
+
+**Example:**
+
+```cpp
+torch::serialize::InputArchive archive;
+archive.load_from("multi_tensor.pt");
+
+torch::Tensor tensor1, tensor2;
+archive.read("tensor1", tensor1);
+archive.read("tensor2", tensor2);
+```
+
+## Saving Multiple Values
+
+Archives are useful when you need to save multiple related values together:
+
+```cpp
+// Save multiple tensors and metadata
+torch::serialize::OutputArchive out_archive;
+out_archive.write("weights", model_weights);
+out_archive.write("biases", model_biases);
+out_archive.write("epoch", torch::tensor(current_epoch));
+out_archive.write("loss", torch::tensor(best_loss));
+out_archive.save_to("training_state.pt");
+
+// Load them back
+torch::serialize::InputArchive in_archive;
+in_archive.load_from("training_state.pt");
+
+torch::Tensor weights, biases, epoch_tensor, loss_tensor;
+in_archive.read("weights", weights);
+in_archive.read("biases", biases);
+in_archive.read("epoch", epoch_tensor);
+in_archive.read("loss", loss_tensor);
+
+int epoch = epoch_tensor.item<int>();
+float loss = loss_tensor.item<float>();
+```
diff --git a/docs/cpp/source/api/serialize/checkpoints.md b/docs/cpp/source/api/serialize/checkpoints.md
new file mode 100644
index 0000000000000..27cae975b47c5
--- /dev/null
+++ b/docs/cpp/source/api/serialize/checkpoints.md
@@ -0,0 +1,113 @@
+---
+myst:
+  html_meta:
+    description: Checkpointing in PyTorch C++ — saving and resuming training state.
+    keywords: PyTorch, C++, checkpoint, save, resume, training state
+---
+
+# Checkpoints
+
+Checkpoints save the complete training state so you can resume training
+after interruption. A checkpoint typically includes:
+
+- Model parameters
+- Optimizer state (momentum buffers, learning rates)
+- Current epoch number
+- Best validation loss/accuracy
+
+## Creating Checkpoints
+
+```cpp
+void save_checkpoint(
+    std::shared_ptr<Net> model,
+    torch::optim::Adam& optimizer,
+    int epoch,
+    const std::string& path) {
+  torch::serialize::OutputArchive archive;
+  model->save(archive);
+  archive.write("epoch", torch::tensor(epoch));
+  optimizer.save(archive);
+  archive.save_to(path);
+}
+```
+
+## Loading Checkpoints
+
+```cpp
+int load_checkpoint(
+    std::shared_ptr<Net> model,
+    torch::optim::Adam& optimizer,
+    const std::string& path) {
+  torch::serialize::InputArchive archive;
+  archive.load_from(path);
+  model->load(archive);
+  torch::Tensor epoch_tensor;
+  archive.read("epoch", epoch_tensor);
+  optimizer.load(archive);
+  return epoch_tensor.item<int>();
+}
+```
+
+## Complete Checkpoint Example
+
+```cpp
+#include <torch/torch.h>
+#include <iostream>
+#include <filesystem>
+
+struct Net : torch::nn::Module {
+  Net() {
+    fc1 = register_module("fc1", torch::nn::Linear(784, 256));
+    fc2 = register_module("fc2", torch::nn::Linear(256, 10));
+  }
+
+  torch::Tensor forward(torch::Tensor x) {
+    x = torch::relu(fc1->forward(x.view({-1, 784})));
+    return fc2->forward(x);
+  }
+
+  torch::nn::Linear fc1{nullptr}, fc2{nullptr};
+};
+
+int main() {
+  auto model = std::make_shared<Net>();
+  auto optimizer = torch::optim::Adam(model->parameters(), 1e-3);
+
+  int start_epoch = 0;
+  const std::string checkpoint_path = "checkpoint.pt";
+
+  // Resume from checkpoint if it exists
+  if (std::filesystem::exists(checkpoint_path)) {
+    std::cout << "Loading checkpoint..." << std::endl;
+    start_epoch = load_checkpoint(model, optimizer, checkpoint_path);
+    std::cout << "Resuming from epoch " << start_epoch << std::endl;
+  }
+
+  // Training loop
+  for (int epoch = start_epoch; epoch < 100; ++epoch) {
+    // ... training code ...
+
+    // Save checkpoint every 10 epochs
+    if ((epoch + 1) % 10 == 0) {
+      save_checkpoint(model, optimizer, epoch + 1, checkpoint_path);
+      std::cout << "Saved checkpoint at epoch " << epoch + 1 << std::endl;
+    }
+  }
+
+  return 0;
+}
+```
+
+## Best Practices
+
+1. **Save periodically**: Save checkpoints at regular intervals (e.g., every epoch
+   or every N batches) to minimize lost work.
+
+2. **Keep multiple checkpoints**: Maintain the last few checkpoints in case the
+   most recent one is corrupted or represents a poor model state.
+
+3. **Include all state**: Save everything needed to resume training, including
+   learning rate scheduler state if using one.
+
+4. **Verify checkpoints**: Occasionally verify that checkpoints can be loaded
+   correctly.
diff --git a/docs/cpp/source/api/serialize/index.md b/docs/cpp/source/api/serialize/index.md
new file mode 100644
index 0000000000000..002431d9ef632
--- /dev/null
+++ b/docs/cpp/source/api/serialize/index.md
@@ -0,0 +1,58 @@
+---
+myst:
+  html_meta:
+    description: PyTorch C++ serialization API — saving and loading models and tensors.
+    keywords: PyTorch, C++, serialization, save, load, checkpoint, model
+---
+
+# Serialization (torch::serialize)
+
+The `torch::serialize` namespace provides utilities for saving and loading
+model weights, tensors, and optimizer state. This enables checkpointing during
+training and deployment of trained models.
+
+**When to use torch::serialize:**
+
+- When saving trained models to disk
+- When implementing training checkpoints
+- When loading pre-trained weights
+- When transferring models between C++ and Python
+
+**Basic usage:**
+
+```cpp
+#include <torch/torch.h>
+
+// Save a model
+auto model = std::make_shared<Net>();
+// ... train the model ...
+torch::save(model, "model.pt");
+
+// Load a model
+auto loaded_model = std::make_shared<Net>();
+torch::load(loaded_model, "model.pt");
+
+// Save and load tensors
+torch::Tensor tensor = torch::randn({2, 3});
+torch::save(tensor, "tensor.pt");
+
+torch::Tensor loaded_tensor;
+torch::load(loaded_tensor, "tensor.pt");
+```
+
+## Header Files
+
+- `torch/csrc/api/include/torch/serialize.h` - Main serialization header
+- `torch/csrc/api/include/torch/serialize/archive.h` - Archive classes
+- `torch/csrc/api/include/torch/serialize/input-archive.h` - Input archive
+- `torch/csrc/api/include/torch/serialize/output-archive.h` - Output archive
+
+## Serialization Categories
+
+```{toctree}
+:maxdepth: 1
+
+save_load
+archives
+checkpoints
+```
diff --git a/docs/cpp/source/api/serialize/save_load.md b/docs/cpp/source/api/serialize/save_load.md
new file mode 100644
index 0000000000000..e08cd21f905cd
--- /dev/null
+++ b/docs/cpp/source/api/serialize/save_load.md
@@ -0,0 +1,84 @@
+---
+myst:
+  html_meta:
+    description: Save and load functions in PyTorch C++ — torch::save and torch::load for tensors and modules.
+    keywords: PyTorch, C++, save, load, torch::save, torch::load, tensor, module
+---
+
+# Saving and Loading
+
+The primary interface for serialization uses the `torch::save` and
+`torch::load` functions, which can save and load tensors, modules,
+and optimizers.
+
+## Save Functions
+
+```{doxygenfunction} torch::save(const Value &value, SaveToArgs&&... args)
+```
+
+```{doxygenfunction} torch::save(const std::vector<torch::Tensor> &tensor_vec, SaveToArgs&&... args)
+```
+
+## Load Functions
+
+```{doxygenfunction} torch::load(Value &value, LoadFromArgs&&... args)
+```
+
+```{doxygenfunction} torch::load(std::vector<torch::Tensor> &tensor_vec, LoadFromArgs&&... args)
+```
+
+## Saving and Loading Tensors
+
+```cpp
+// Save a tensor
+torch::Tensor tensor = torch::randn({2, 3});
+torch::save(tensor, "tensor.pt");
+
+// Load a tensor
+torch::Tensor loaded;
+torch::load(loaded, "tensor.pt");
+```
+
+## Saving and Loading Modules
+
+```cpp
+// Define a model
+struct Net : torch::nn::Module {
+  Net() {
+    fc1 = register_module("fc1", torch::nn::Linear(784, 64));
+    fc2 = register_module("fc2", torch::nn::Linear(64, 10));
+  }
+
+  torch::Tensor forward(torch::Tensor x) {
+    x = torch::relu(fc1->forward(x));
+    return fc2->forward(x);
+  }
+
+  torch::nn::Linear fc1{nullptr}, fc2{nullptr};
+};
+
+// Save model
+auto model = std::make_shared<Net>();
+torch::save(model, "model.pt");
+
+// Load model
+auto loaded_model = std::make_shared<Net>();
+torch::load(loaded_model, "model.pt");
+```
+
+## Saving Optimizer State
+
+```cpp
+auto model = std::make_shared<Net>();
+auto optimizer = torch::optim::Adam(model->parameters(), 0.001);
+
+// Train...
+
+// Save both model and optimizer
+torch::save(model, "model.pt");
+torch::save(optimizer, "optimizer.pt");
+
+// Load both
+torch::load(model, "model.pt");
+torch::load(optimizer, "optimizer.pt");
+```
diff --git a/docs/cpp/source/api/stable/index.md b/docs/cpp/source/api/stable/index.md
new file mode 100644
index 0000000000000..adb3d477d18b6
--- /dev/null
+++ b/docs/cpp/source/api/stable/index.md
@@ -0,0 +1,72 @@
+---
+myst:
+  html_meta:
+    description: PyTorch Stable ABI C++ API — binary-compatible operator registration across PyTorch versions.
+    keywords: PyTorch, C++, stable ABI, binary compatibility, operator registration
+---
+
+# Torch Stable API
+
+The PyTorch Stable C++ API provides a binary-compatible interface for calling
+tensor operations and utilities that is guaranteed to remain stable across
+PyTorch versions. This enables ahead-of-time compiled extensions that don't
+need recompilation when PyTorch is updated.
+
+**When to use the Stable API:**
+
+- When building extensions that must work across multiple PyTorch versions
+- When distributing pre-compiled binaries
+- When binary compatibility is more important than access to the latest features
+- When writing custom operators for production deployment
+
+**Basic usage:**
+
+```cpp
+#include <torch/csrc/stable/library.h>
+#include <torch/csrc/stable/ops.h>
+
+// Create a tensor using stable API
+auto tensor = torch::stable::empty(
+    {3, 4},
+    torch::headeronly::ScalarType::Float,
+    torch::headeronly::Layout::Strided,
+    torch::stable::Device(torch::headeronly::DeviceType::CPU),
+    false,
+    torch::headeronly::MemoryFormat::Contiguous);
+
+// Register operators with stable ABI
+STABLE_TORCH_LIBRARY(myops, m) {
+    m.def("my_op(Tensor input) -> Tensor");
+}
+
+STABLE_TORCH_LIBRARY_IMPL(myops, CPU, m) {
+    m.impl("my_op", TORCH_BOX(&my_cpu_kernel));
+}
+```
+
+For more information on the stable ABI, see the
+[Stable ABI notes](https://docs.pytorch.org/docs/stable/notes/libtorch_stable_abi.html).
+
+## Header Files
+
+- `torch/csrc/stable/library.h` - Stable library registration
+- `torch/csrc/stable/ops.h` - Stable operator definitions
+- `torch/csrc/stable/tensor_struct.h` - Stable tensor structures
+- `torch/csrc/stable/device_struct.h` - Stable device structures
+- `torch/csrc/stable/accelerator.h` - Accelerator support
+- `torch/csrc/stable/macros.h` - Stable API macros
+
+## Stable API Categories
+
+```{toctree}
+:maxdepth: 1
+
+registration
+operators
+utilities
+```
+
+## See Also
+
+- {doc}`../library/index` - For standard (non-stable) operator registration
+- [Stable ABI documentation](https://pytorch.org/docs/stable/cpp_extension.html)
diff --git a/docs/cpp/source/api/stable/operators.md b/docs/cpp/source/api/stable/operators.md
new file mode 100644
index 0000000000000..8c3bfac176ea6
--- /dev/null
+++ b/docs/cpp/source/api/stable/operators.md
@@ -0,0 +1,146 @@
+---
+myst:
+  html_meta:
+    description: Stable ABI operator API in PyTorch C++ — StableLibrary and boxed kernel registration.
+    keywords: PyTorch, C++, stable ABI, StableLibrary, operator, boxed kernel
+---
+
+# Stable Operators
+
+The stable API provides tensor operations that maintain binary compatibility
+across PyTorch versions.
+
+## Tensor Class
+
+```{doxygenclass} torch::stable::Tensor
+:members:
+:undoc-members:
+```
+
+**Example:**
+
+```cpp
+torch::stable::Tensor tensor = torch::stable::empty({3, 4}, ...);
+float* data = tensor.data_ptr<float>();
+auto shape = tensor.sizes();
+```
+
+## Device Class
+
+```{doxygenclass} torch::stable::Device
+:members:
+:undoc-members:
+```
+
+**Example:**
+
+```cpp
+torch::stable::Device cpu_device(torch::headeronly::DeviceType::CPU);
+torch::stable::Device cuda_device(torch::headeronly::DeviceType::CUDA, 0);
+```
+
+## Tensor Creation
+
+```{doxygenfunction} torch::stable::empty
+```
+
+```{doxygenfunction} torch::stable::empty_like
+```
+
+```{doxygenfunction} torch::stable::new_empty(const torch::stable::Tensor &self, torch::headeronly::IntHeaderOnlyArrayRef size, std::optional<torch::headeronly::ScalarType> dtype, std::optional<torch::headeronly::Layout> layout, std::optional<torch::stable::Device> device, std::optional<bool> pin_memory)
+```
+
+```{doxygenfunction} torch::stable::new_zeros(const torch::stable::Tensor &self, torch::headeronly::IntHeaderOnlyArrayRef size, std::optional<torch::headeronly::ScalarType> dtype, std::optional<torch::headeronly::Layout> layout, std::optional<torch::stable::Device> device, std::optional<bool> pin_memory)
+```
+
+```{doxygenfunction} torch::stable::full
+```
+
+```{doxygenfunction} torch::stable::from_blob(void *data, torch::headeronly::IntHeaderOnlyArrayRef sizes, torch::headeronly::IntHeaderOnlyArrayRef strides, torch::stable::Device device, torch::headeronly::ScalarType dtype, int64_t storage_offset, torch::headeronly::Layout layout)
+```
+
+**Example:**
+
+```cpp
+auto tensor = torch::stable::empty(
+    {3, 4},
+    torch::headeronly::ScalarType::Float,
+    torch::headeronly::Layout::Strided,
+    torch::stable::Device(torch::headeronly::DeviceType::CUDA, 0),
+    false,
+    torch::headeronly::MemoryFormat::Contiguous);
+```
+
+## Tensor Manipulation
+
+```{doxygenfunction} torch::stable::clone
+```
+
+```{doxygenfunction} torch::stable::contiguous
+```
+
+```{doxygenfunction} torch::stable::reshape
+```
+
+```{doxygenfunction} torch::stable::view
+```
+
+```{doxygenfunction} torch::stable::flatten
+```
+
+```{doxygenfunction} torch::stable::squeeze
+```
+
+```{doxygenfunction} torch::stable::unsqueeze
+```
+
+```{doxygenfunction} torch::stable::transpose
+```
+
+```{doxygenfunction} torch::stable::select
+```
+
+```{doxygenfunction} torch::stable::narrow
+```
+
+```{doxygenfunction} torch::stable::pad
+```
+
+## Device and Type Conversion
+
+```{doxygenfunction} torch::stable::to(const torch::stable::Tensor &self, std::optional<torch::headeronly::ScalarType> dtype, std::optional<torch::headeronly::Layout> layout, std::optional<torch::stable::Device> device, std::optional<bool> pin_memory, bool non_blocking, bool copy, std::optional<torch::headeronly::MemoryFormat> memory_format)
+```
+
+```{doxygenfunction} torch::stable::to(const torch::stable::Tensor &self, torch::stable::Device device, bool non_blocking, bool copy)
+```
+
+## In-place Operations
+
+```{doxygenfunction} torch::stable::fill_
+```
+
+```{doxygenfunction} torch::stable::zero_
+```
+
+```{doxygenfunction} torch::stable::copy_
+```
+
+## Mathematical Operations
+
+```{doxygenfunction} torch::stable::matmul
+```
+
+```{doxygenfunction} torch::stable::amax(const torch::stable::Tensor &self, int64_t dim, bool keepdim)
+```
+
+```{doxygenfunction} torch::stable::amax(const torch::stable::Tensor &self, torch::headeronly::IntHeaderOnlyArrayRef dims, bool keepdim)
+```
+
+```{doxygenfunction} torch::stable::sum
+```
+
+```{doxygenfunction} torch::stable::sum_out
+```
+
+```{doxygenfunction} torch::stable::subtract
+```
diff --git a/docs/cpp/source/api/stable/registration.md b/docs/cpp/source/api/stable/registration.md
new file mode 100644
index 0000000000000..ad970a66d72ae
--- /dev/null
+++ b/docs/cpp/source/api/stable/registration.md
@@ -0,0 +1,116 @@
+---
+myst:
+  html_meta:
+    description: Stable ABI registration macros in PyTorch C++ — STABLE_TORCH_LIBRARY, STABLE_TORCH_LIBRARY_IMPL, and TORCH_BOX.
+    keywords: PyTorch, C++, stable ABI, STABLE_TORCH_LIBRARY, TORCH_BOX, macro, registration
+---
+
+# Library Registration Macros
+
+These macros provide stable ABI equivalents of the standard PyTorch operator
+registration macros (`TORCH_LIBRARY`, `TORCH_LIBRARY_IMPL`, etc.).
+Use these when building custom operators that need to maintain binary
+compatibility across PyTorch versions.
+
+## STABLE_TORCH_LIBRARY
+
+```{c:macro} STABLE_TORCH_LIBRARY(ns, m)
+
+Defines a library of operators in a namespace using the stable ABI.
+
+This is the stable ABI equivalent of {c:macro}`TORCH_LIBRARY`.
+Use this macro to define operator schemas that will maintain
+binary compatibility across PyTorch versions. Only one `STABLE_TORCH_LIBRARY`
+block can exist per namespace; use `STABLE_TORCH_LIBRARY_FRAGMENT` for
+additional definitions in the same namespace from different translation units.
+
+:param ns: The namespace in which to define operators (e.g., `mylib`).
+:param m: The name of the StableLibrary variable available in the block.
+
+**Example:**
+
+```cpp
+STABLE_TORCH_LIBRARY(mylib, m) {
+    m.def("my_op(Tensor input, int size) -> Tensor");
+    m.def("another_op(Tensor a, Tensor b) -> Tensor");
+}
+```
+
+Minimum compatible version: PyTorch 2.9.
+```
+## STABLE_TORCH_LIBRARY_IMPL
+
+```{c:macro} STABLE_TORCH_LIBRARY_IMPL(ns, k, m)
+
+Registers operator implementations for a specific dispatch key using the stable ABI.
+
+This is the stable ABI equivalent of `TORCH_LIBRARY_IMPL`. Use this macro
+to provide implementations of operators for a specific dispatch key (e.g.,
+CPU, CUDA) while maintaining binary compatibility across PyTorch versions.
+```
+
+```{note}
+
+All kernel functions registered with this macro must be boxed using
+the `TORCH_BOX` macro.
+```
+
+:param ns: The namespace in which the operators are defined.
+:param k: The dispatch key (e.g., `CPU`, `CUDA`).
+:param m: The name of the StableLibrary variable available in the block.
+
+**Example:**
+
+```cpp
+STABLE_TORCH_LIBRARY_IMPL(mylib, CPU, m) {
+    m.impl("my_op", TORCH_BOX(&my_cpu_kernel));
+}
+
+STABLE_TORCH_LIBRARY_IMPL(mylib, CUDA, m) {
+    m.impl("my_op", TORCH_BOX(&my_cuda_kernel));
+}
+```
+
+Minimum compatible version: PyTorch 2.9.
+
+## STABLE_TORCH_LIBRARY_FRAGMENT
+
+```{c:macro} STABLE_TORCH_LIBRARY_FRAGMENT(ns, m)
+
+Extends operator definitions in an existing namespace using the stable ABI.
+
+This is the stable ABI equivalent of `TORCH_LIBRARY_FRAGMENT`. Use this macro
+to add additional operator definitions to a namespace that was already
+created with `STABLE_TORCH_LIBRARY`.
+
+:param ns: The namespace to extend.
+:param m: The name of the StableLibrary variable available in the block.
+
+Minimum compatible version: PyTorch 2.9.
+```
+
+## TORCH_BOX
+
+```{c:macro} TORCH_BOX(func)
+
+Wraps a function to conform to the stable boxed kernel calling convention.
+
+This macro takes an unboxed kernel function pointer and generates a boxed wrapper
+that can be registered with the stable library API.
+
+:param func: The unboxed kernel function to wrap.
+
+**Example:**
+
+```cpp
+Tensor my_kernel(const Tensor& input, int64_t size) {
+    return input.reshape({size});
+}
+
+STABLE_TORCH_LIBRARY_IMPL(my_namespace, CPU, m) {
+    m.impl("my_op", TORCH_BOX(&my_kernel));
+}
+```
+
+Minimum compatible version: PyTorch 2.9.
+```
\ No newline at end of file
diff --git a/docs/cpp/source/api/stable/utilities.md b/docs/cpp/source/api/stable/utilities.md
new file mode 100644
index 0000000000000..c9d30e96ddef9
--- /dev/null
+++ b/docs/cpp/source/api/stable/utilities.md
@@ -0,0 +1,267 @@
+---
+myst:
+  html_meta:
+    description: Stable ABI utilities in PyTorch C++ — version checking and compatibility helpers.
+    keywords: PyTorch, C++, stable ABI, utilities, version, compatibility
+---
+
+# Utilities
+
+The stable API provides various utility functions and types for working with
+tensors and CUDA operations.
+
+## DeviceGuard Class
+
+```{doxygenclass} torch::stable::accelerator::DeviceGuard
+:members:
+:undoc-members:
+```
+
+```{doxygenfunction} torch::stable::accelerator::getCurrentDeviceIndex
+```
+
+**Example:**
+
+```cpp
+{
+    torch::stable::accelerator::DeviceGuard guard(1);
+    // Operations here run on device 1
+}
+// Previous device is restored
+```
+
+## Stream
+
+```{doxygenclass} torch::stable::accelerator::Stream
+:members:
+:undoc-members:
+```
+
+## Stream Utilities
+
+For CUDA stream access, we currently recommend the ABI stable C shim API. This
+will be improved in a future release with a more ergonomic wrapper.
+
+### Getting the Current CUDA Stream
+
+To obtain the current `cudaStream_t` for use in CUDA kernels:
+
+```cpp
+#include <torch/csrc/inductor/aoti_torch/c/shim.h>
+#include <torch/headeronly/util/shim_utils.h>
+
+// For now, we rely on the ABI stable C shim API to get the current CUDA stream.
+void* stream_ptr = nullptr;
+TORCH_ERROR_CODE_CHECK(
+    aoti_torch_get_current_cuda_stream(tensor.get_device_index(), &stream_ptr));
+cudaStream_t stream = static_cast<cudaStream_t>(stream_ptr);
+
+// Now you can use 'stream' in your CUDA kernel launches
+my_kernel<<<blocks, threads, 0, stream>>>(args...);
+```
+
+```{note}
+
+The `TORCH_ERROR_CODE_CHECK` macro is required when using C shim APIs
+to properly check error codes and throw appropriate exceptions.
+```
+
+## CUDA Error Checking Macros
+
+These macros provide stable ABI equivalents for CUDA error checking.
+They wrap CUDA API calls and kernel launches, providing detailed error
+messages using PyTorch's error formatting.
+
+### STD_CUDA_CHECK
+
+```{c:macro} STD_CUDA_CHECK(EXPR)
+
+Checks the result of a CUDA API call and throws an exception on error.
+Users of this macro are expected to include `cuda_runtime.h`.
+
+**Example:**
+
+```cpp
+STD_CUDA_CHECK(cudaMalloc(&ptr, size));
+STD_CUDA_CHECK(cudaMemcpy(dst, src, size, cudaMemcpyDeviceToHost));
+```
+
+Minimum compatible version: PyTorch 2.10.
+```
+### STD_CUDA_KERNEL_LAUNCH_CHECK
+
+```{c:macro} STD_CUDA_KERNEL_LAUNCH_CHECK()
+
+Checks for errors from the most recent CUDA kernel launch. Equivalent to
+`STD_CUDA_CHECK(cudaGetLastError())`.
+
+**Example:**
+
+```cpp
+my_kernel<<<blocks, threads, 0, stream>>>(args...);
+STD_CUDA_KERNEL_LAUNCH_CHECK();
+```
+
+Minimum compatible version: PyTorch 2.10.
+```
+## Header-Only Utilities
+
+The `torch::headeronly` namespace provides header-only versions of common
+PyTorch types and utilities. These can be used without linking against libtorch
+at all! This portability makes them ideal for maintaining binary compatibility
+across PyTorch versions.
+
+### Error Checking
+
+`STD_TORCH_CHECK` is a header-only macro for runtime assertions:
+
+```cpp
+#include <torch/headeronly/util/Exception.h>
+
+STD_TORCH_CHECK(condition, "Error message with ", variable, " interpolation");
+```
+
+Wherever you used `TORCH_CHECK` before, you can replace usage with `STD_TORCH_CHECK`
+to remove the need to link against libtorch. The only difference is that when the
+condition check fails, `TORCH_CHECK` throws a fancier `c10::Error` while
+`STD_TORCH_CHECK` throws a `std::runtime_error`.
+
+### Core Types
+
+The following `c10::` types are available as header-only versions under
+`torch::headeronly::`:
+
+- `torch::headeronly::ScalarType` - Tensor data types (Float, Double, Int, etc.)
+- `torch::headeronly::DeviceType` - Device types (CPU, CUDA, etc.)
+- `torch::headeronly::MemoryFormat` - Memory layout formats (Contiguous, ChannelsLast, etc.)
+- `torch::headeronly::Layout` - Tensor layouts (Strided, Sparse, etc.)
+
+```cpp
+#include <torch/headeronly/core/ScalarType.h>
+#include <torch/headeronly/core/DeviceType.h>
+#include <torch/headeronly/core/MemoryFormat.h>
+#include <torch/headeronly/core/Layout.h>
+
+auto dtype = torch::headeronly::ScalarType::Float;
+auto device_type = torch::headeronly::DeviceType::CUDA;
+auto memory_format = torch::headeronly::MemoryFormat::Contiguous;
+auto layout = torch::headeronly::Layout::Strided;
+```
+
+### TensorAccessor
+
+`TensorAccessor` provides efficient, bounds-checked access to tensor data.
+You can construct one from a stable tensor's data pointer, sizes, and strides:
+
+```cpp
+#include <torch/headeronly/core/TensorAccessor.h>
+
+// Create a TensorAccessor for a 2D float tensor
+auto sizes = tensor.sizes();
+auto strides = tensor.strides();
+torch::headeronly::TensorAccessor<float, 2> accessor(
+    static_cast<float*>(tensor.mutable_data_ptr()),
+    sizes.data(),
+    strides.data());
+
+// Access elements
+float value = accessor[i][j];
+```
+
+### Dispatch Macros
+
+Header-only dispatch macros (THO = Torch Header Only) are available for
+dtype dispatching:
+
+```cpp
+#include <torch/headeronly/core/Dispatch_v2.h>
+
+THO_DISPATCH_V2(
+   tensor.scalar_type(),  // will be resolved as scalar_t
+   "my_kernel",
+   AT_WRAP(([&]() {
+   // code to specialize with scalar_t
+   // scalar_t is the resolved C++ type (e.g. float, double)
+   auto* data = static_cast<scalar_t*>(tensor.mutable_data_ptr());
+   Scalar s(*data);
+   })),
+   AT_EXPAND(AT_ALL_TYPES),
+   AT_EXPAND(AT_COMPLEX_TYPES),
+   torch::headeronly::ScalarType::Half,
+   // as many type arguments as needed
+);
+```
+
+`THO_DISPATCH_V2` works the same way as `AT_DISPATCH_V2` (see
+`ATen/Dispatch_v2.h`) but does not require linking against libtorch.
+As a result, whereas `AT_DISPATCH_V2` would have thrown `c10::NotImplementedError`
+for unimplemented paths, `THO_DISPATCH_V2` will throw `std::runtime_error`.
+
+For ease of use, we've also migrated the below AT_* macros representing
+collections of types to be header-only and thus have no dependency on libtorch:
+
+- `AT_FLOATING_TYPES`
+- `AT_INTEGRAL_TYPES`
+- `AT_INTEGRAL_TYPES_V2`
+- `AT_ALL_TYPES`
+- `AT_COMPLEX_TYPES`
+- `AT_ALL_TYPES_AND_COMPLEX`
+- `AT_FLOAT8_TYPES`
+- `AT_BAREBONES_UNSIGNED_TYPES`
+- `AT_QINT_TYPES`
+
+If your extension uses our older AT_DISPATCH version 1 infrastructure,
+you can also migrate to a header-only libtorch-free world without upgrading
+everything to version 2.
+
+`THO_DISPATCH_SWITCH` and `THO_DISPATCH_CASE` are the header-only
+equivalents of `AT_DISPATCH_SWITCH` and `AT_DISPATCH_CASE`. Similarly,
+the only user-visible difference is the exception type on an unhandled dtype,
+where the `AT_` version throws a `c10::NotImplementedError` and the `THO_`
+version throws a `std::runtime_error`.
+
+The migration is pretty mechanical:
+
+- `AT_DISPATCH_SWITCH` → `THO_DISPATCH_SWITCH`
+- `AT_DISPATCH_CASE` → `THO_DISPATCH_CASE`
+- `AT_PRIVATE_CASE_TYPE_USING_HINT` → `THO_PRIVATE_CASE_TYPE_USING_HINT`
+- `at::ScalarType::X` → `torch::headeronly::ScalarType::X`
+
+```cpp
+// ---- Before (requires linking against libtorch) ----
+#include <torch/all.h>
+
+#define MY_DISPATCH_CASE_FLOATING_TYPES(...)            \
+  AT_DISPATCH_CASE(at::ScalarType::Float, __VA_ARGS__) \
+  AT_DISPATCH_CASE(at::ScalarType::Half, __VA_ARGS__)  \
+  AT_DISPATCH_CASE(at::ScalarType::BFloat16, __VA_ARGS__)
+
+#define MY_DISPATCH_FLOATING_TYPES(TYPE, NAME, ...) \
+  AT_DISPATCH_SWITCH(TYPE, NAME,                    \
+                     MY_DISPATCH_CASE_FLOATING_TYPES(__VA_ARGS__))
+```
+
+```cpp
+// ---- After (header-only, no libtorch dependency) ----
+#include <torch/headeronly/core/Dispatch.h>
+
+#define MY_DISPATCH_CASE_FLOATING_TYPES(...)                          \
+  THO_DISPATCH_CASE(torch::headeronly::ScalarType::Float, __VA_ARGS__) \
+  THO_DISPATCH_CASE(torch::headeronly::ScalarType::Half, __VA_ARGS__)  \
+  THO_DISPATCH_CASE(torch::headeronly::ScalarType::BFloat16, __VA_ARGS__)
+
+#define MY_DISPATCH_FLOATING_TYPES(TYPE, NAME, ...) \
+  THO_DISPATCH_SWITCH(TYPE, NAME,                   \
+                      MY_DISPATCH_CASE_FLOATING_TYPES(__VA_ARGS__))
+```
+
+For the complete list of header-only APIs, see `torch/header_only_apis.txt`
+in the PyTorch source tree.
+
+## Parallelization Utilities
+
+```{doxygenfunction} torch::stable::parallel_for
+```
+
+```{doxygenfunction} torch::stable::get_num_threads
+```
diff --git a/docs/cpp/source/api/xpu/index.md b/docs/cpp/source/api/xpu/index.md
new file mode 100644
index 0000000000000..c88eff6a69165
--- /dev/null
+++ b/docs/cpp/source/api/xpu/index.md
@@ -0,0 +1,48 @@
+---
+myst:
+  html_meta:
+    description: PyTorch XPU C++ API — Intel GPU support with device management, streams, and guards.
+    keywords: PyTorch, C++, XPU, Intel GPU, device, stream
+---
+
+# XPU Support
+
+PyTorch provides XPU support for Intel GPU-accelerated tensor operations.
+The XPU API allows you to manage Intel GPU devices, streams for asynchronous
+execution, and synchronization.
+
+**When to use XPU APIs:**
+
+- When running on Intel GPUs (Data Center GPU Max, Arc, etc.)
+- When implementing custom XPU kernels or operations
+- When managing asynchronous execution with XPU streams
+- When writing device-portable code alongside CUDA
+
+**Basic usage:**
+
+```cpp
+#include <torch/torch.h>
+
+// Check if XPU is available
+if (torch::xpu::is_available()) {
+    // Create tensor on XPU
+    auto tensor = torch::randn({2, 3}, torch::device(torch::kXPU));
+
+    // Move model to XPU
+    model->to(torch::kXPU);
+}
+```
+
+## Header Files
+
+- `torch/xpu.h` - High-level XPU utilities (device count, availability, seeding)
+- `c10/xpu/XPUStream.h` - XPU stream management
+
+## XPU Categories
+
+```{toctree}
+:maxdepth: 1
+
+streams
+utilities
+```
diff --git a/docs/cpp/source/api/xpu/streams.md b/docs/cpp/source/api/xpu/streams.md
new file mode 100644
index 0000000000000..a9b6911471d4f
--- /dev/null
+++ b/docs/cpp/source/api/xpu/streams.md
@@ -0,0 +1,50 @@
+---
+myst:
+  html_meta:
+    description: XPU streams in PyTorch C++ — XPUStream for asynchronous Intel GPU execution.
+    keywords: PyTorch, C++, XPU, XPUStream, Intel GPU, stream, asynchronous
+---
+
+# XPU Streams
+
+XPU streams provide a mechanism for asynchronous execution of operations
+on Intel GPUs. Like CUDA streams, operations queued to the same stream execute
+in order, while operations on different streams can execute concurrently.
+
+## XPUStream
+
+```{doxygenclass} c10::xpu::XPUStream
+:members:
+:undoc-members:
+```
+
+**Example:**
+
+```cpp
+#include <c10/xpu/XPUStream.h>
+
+// Get the current XPU stream
+auto stream = c10::xpu::getCurrentXPUStream();
+
+// Create a new stream from the pool
+auto new_stream = c10::xpu::getStreamFromPool();
+
+// Synchronize
+stream.synchronize();
+```
+
+## Acquiring XPU Streams
+
+```{doxygenfunction} c10::xpu::getCurrentXPUStream
+```
+
+```{doxygenfunction} c10::xpu::setCurrentXPUStream
+```
+
+```{doxygenfunction} c10::xpu::getStreamFromPool(const bool isHighPriority, DeviceIndex device)
+```
+
+## Stream Synchronization
+
+```{doxygenfunction} c10::xpu::syncStreamsOnDevice
+```
diff --git a/docs/cpp/source/api/xpu/utilities.md b/docs/cpp/source/api/xpu/utilities.md
new file mode 100644
index 0000000000000..9857d5e291a2b
--- /dev/null
+++ b/docs/cpp/source/api/xpu/utilities.md
@@ -0,0 +1,53 @@
+---
+myst:
+  html_meta:
+    description: XPU utility functions in PyTorch C++ — device count, properties, and stream management for Intel GPUs.
+    keywords: PyTorch, C++, XPU, Intel GPU, device_count, utilities
+---
+
+# XPU Utility Functions
+
+High-level utility functions for querying and managing XPU devices.
+
+## Device Management
+
+```{doxygenfunction} torch::xpu::device_count
+```
+
+```{doxygenfunction} torch::xpu::is_available
+```
+
+```{doxygenfunction} torch::xpu::synchronize
+```
+
+**Example:**
+
+```cpp
+#include <torch/torch.h>
+
+if (torch::xpu::is_available()) {
+    size_t num_devices = torch::xpu::device_count();
+    std::cout << "Found " << num_devices << " XPU device(s)" << std::endl;
+
+    // Synchronize all streams on device 0
+    torch::xpu::synchronize(0);
+}
+```
+
+## Random Number Generation
+
+```{doxygenfunction} torch::xpu::manual_seed
+```
+
+```{doxygenfunction} torch::xpu::manual_seed_all
+```
+
+**Example:**
+
+```cpp
+// Set seed for reproducibility on current XPU device
+torch::xpu::manual_seed(42);
+
+// Set seed for all XPU devices
+torch::xpu::manual_seed_all(42);
+```
diff --git a/docs/cpp/source/conf.py b/docs/cpp/source/conf.py
index 10d854c21db4f..ab3c722bb63ed 100644
--- a/docs/cpp/source/conf.py
+++ b/docs/cpp/source/conf.py
@@ -20,7 +20,6 @@
 # See https://github.com/pytorch/pytorch/issues/79992.
 
 import os
-import textwrap
 
 # sys.path.insert(0, os.path.abspath('.'))
 import pytorch_sphinx_theme2
@@ -31,14 +30,55 @@
 # If your documentation needs a minimal Sphinx version, state it here.
 #
 needs_sphinx = "3.1.2"
-run_doxygen = os.environ.get("RUN_DOXYGEN", "false") == "true"
 
 # Add any Sphinx extension module names here, as strings. They can be
 # extensions coming with Sphinx (named 'sphinx.ext.*') or your custom
 # ones.
 extensions = [
     "sphinx.ext.intersphinx",
-] + (["breathe", "exhale"] if run_doxygen else [])
+    "breathe",
+    "myst_parser",
+]
+
+# -- Breathe Configuration ------------------------------------------------
+# Breathe connects Sphinx to Doxygen XML output
+# Use doxygenclass, doxygenfunction, etc. directives in RST files
+# to pull documentation from C++ source code
+
+breathe_projects = {"PyTorch": "../build/xml"}
+breathe_default_project = "PyTorch"
+
+# Default members to show when using doxygenclass/doxygenstruct directives
+breathe_default_members = ("members", "undoc-members")
+
+# Map file extensions to language domains for proper syntax highlighting
+breathe_domain_by_extension = {
+    "h": "cpp",
+    "hpp": "cpp",
+    "cpp": "cpp",
+    "c": "c",
+}
+
+# Implementation detail filters - skip internal/private content
+breathe_implementation_filename_extensions = [".c", ".cc", ".cpp"]
+
+# Show the file where items are defined
+breathe_show_define_initializer = True
+breathe_show_enumvalue_initializer = True
+
+# Control what gets shown in documentation
+breathe_show_include = False  # Don't show #include directives
+
+# Order of member documentation
+breathe_order_parameters_first = False
+
+# Use Sphinx's C++ domain for cross-references
+breathe_use_project_refids = True
+
+# Suppress specific Breathe warnings for cleaner builds
+breathe_debug_trace_directives = False
+breathe_debug_trace_doxygen_ids = False
+breathe_debug_trace_qualification = False
 
 intersphinx_mapping = {"pytorch": ("https://docs.pytorch.org/docs/main", None)}
 
@@ -52,99 +92,6 @@
     "misc.highlighting_failure",
 ]
 
-# Configure Breathe
-breathe_show_define_initializer = True
-breathe_show_enumvalue_initializer = True
-breathe_default_members = ("members", "undoc-members")
-
-
-# Fix for Python 3.10+ compatibility with exhale 2.3.0
-# MutableMapping was moved from collections to collections.abc in Python 3.10
-try:
-    import collections
-    from collections.abc import MutableMapping
-
-    if not hasattr(collections, "MutableMapping"):
-        collections.MutableMapping = MutableMapping
-except ImportError:
-    pass
-
-# Setup absolute paths for communicating with breathe / exhale where
-# items are expected / should be trimmed by.
-# This file is {repo_root}/docs/cpp/source/conf.py
-this_file_dir = os.path.abspath(os.path.dirname(__file__))
-doxygen_xml_dir = os.path.join(
-    os.path.dirname(this_file_dir),  # {repo_root}/docs/cpp
-    "build",  # {repo_root}/docs/cpp/build
-    "xml",  # {repo_root}/docs/cpp/build/xml
-)
-repo_root = os.path.dirname(  # {repo_root}
-    os.path.dirname(  # {repo_root}/docs
-        os.path.dirname(  # {repo_root}/docs/cpp
-            this_file_dir  # {repo_root}/docs/cpp/source
-        )
-    )
-)
-
-breathe_projects = {"PyTorch": doxygen_xml_dir}
-breathe_default_project = "PyTorch"
-
-# Setup the exhale extension
-exhale_args = {
-    ############################################################################
-    # These arguments are required.                                            #
-    ############################################################################
-    "containmentFolder": "./api",
-    "rootFileName": "library_root.rst",
-    "rootFileTitle": "Library API",
-    "doxygenStripFromPath": repo_root,
-    ############################################################################
-    # Suggested optional arguments.                                            #
-    ############################################################################
-    "createTreeView": True,
-    "exhaleExecutesDoxygen": True,
-    "exhaleUseDoxyfile": True,
-    "verboseBuild": True,
-    ############################################################################
-    # HTML Theme specific configurations.                                      #
-    ############################################################################
-    # Fix broken Sphinx RTD Theme 'Edit on GitHub' links
-    # Search for 'Edit on GitHub' on the FAQ:
-    #     http://exhale.readthedocs.io/en/latest/faq.html
-    "pageLevelConfigMeta": ":github_url: https://github.com/pytorch/pytorch",
-    ############################################################################
-    # Individual page layout example configuration.                            #
-    ############################################################################
-    # Example of adding contents directives on custom kinds with custom title
-    "contentsTitle": "Page Contents",
-    "kindsWithContentsDirectives": ["class", "file", "namespace", "struct"],
-    # Exclude PIMPL files from class hierarchy tree and namespace pages.
-    "listingExclude": [r".*Impl$"],
-    ############################################################################
-    # Main library page layout example configuration.                          #
-    ############################################################################
-    "afterTitleDescription": textwrap.dedent(
-        """
-        Welcome to the developer reference for the PyTorch C++ API.
-    """
-    ),
-    ############################################################################
-    # Duplicate handling and error management.                                 #
-    ############################################################################
-    # Note: Using Doxyfile instead of stdin configuration
-    # "exhaleDoxygenStdin" is not compatible with "exhaleUseDoxyfile"
-    # Handle unresolved references more gracefully
-    "unabridgedOrphanKinds": {
-        "function",
-        "define",
-        "enum",
-        "enumvalue",
-        "typedef",
-        "variable",
-    },
-    "fullToctreeMaxDepth": 2,
-}
-
 # Tell sphinx what the primary language being documented is.
 primary_domain = "cpp"
 
@@ -164,7 +111,10 @@
 # You can specify multiple suffix as a list of string:
 #
 # source_suffix = ['.rst', '.md']
-source_suffix = ".rst"
+source_suffix = {
+    ".rst": "restructuredtext",
+    ".md": "markdown",
+}
 
 # The master toctree document.
 master_doc = "index"
@@ -264,8 +214,24 @@
 # relative to this directory. They are copied after the builtin static files,
 # so a file named "default.css" will overwrite the builtin "default.css".
 # NOTE: sharing python docs resources
-html_static_path = [os.path.join(repo_root, "docs", "cpp", "source", "_static")]
-html_css_files = ["cpp_theme.css"]
+
+
+# Remove "Subclassed by" paragraphs that Breathe renders as plain text
+# (not links) because per-class pages don't exist without exhale.
+def remove_subclassed_by(app, doctree, docname):
+    from docutils import nodes
+
+    for node in doctree.traverse(nodes.paragraph):
+        text = node.astext()
+        if text.startswith("Subclassed by "):
+            # Keep if it contains actual reference links
+            if not node.traverse(nodes.reference):
+                node.parent.remove(node)
+
+
+def setup(app):
+    app.connect("doctree-resolved", remove_subclassed_by)
+
 
 # Called automatically by Sphinx, making this `conf.py` an "extension".
 
diff --git a/docs/cpp/source/faq.md b/docs/cpp/source/faq.md
new file mode 100644
index 0000000000000..3ee305ccdf3ad
--- /dev/null
+++ b/docs/cpp/source/faq.md
@@ -0,0 +1,244 @@
+---
+myst:
+  html_meta:
+    description: Frequently asked questions about the PyTorch C++ API and libtorch.
+    keywords: PyTorch, C++, FAQ, libtorch, troubleshooting
+---
+
+# FAQ
+
+Listed below are a number of common issues users face with the various parts of
+the C++ API.
+
+## C++ Extensions
+
+### Undefined symbol errors from PyTorch/ATen
+
+**Problem**: You import your extension and get an `ImportError` stating that
+some C++ symbol from PyTorch or ATen is undefined. For example:
+
+```cpp
+>>> import extension
+Traceback (most recent call last):
+  File "<stdin>", line 1, in <module>
+ImportError: /home/user/.pyenv/versions/3.7.1/lib/python3.7/site-packages/extension.cpython-37m-x86_64-linux-gnu.so: undefined symbol: _ZN2at19UndefinedTensorImpl10_singletonE
+```
+
+**Fix**: The fix is to `import torch` before you import your extension. This will make
+the symbols from the PyTorch dynamic (shared) library that your extension
+depends on available, allowing them to be resolved once you import your extension.
+
+### I created a tensor using a function from `at::` and get errors
+
+**Problem**: You created a tensor using e.g. `at::ones` or `at::randn` or
+any other tensor factory from the `at::` namespace and are getting errors.
+
+**Fix**: Replace `at::` with `torch::` for factory function calls. You
+should never use factory functions from the `at::` namespace, as they will
+create tensors. The corresponding `torch::` functions will create variables,
+and you should only ever deal with variables in your code.
+
+## LibTorch
+
+### How do I move a model to GPU?
+
+**Problem**: You want to run your model on GPU but are unsure how to move both
+the model and tensors to the correct device.
+
+**Fix**: Use the `to()` method to move your model and tensors to a CUDA device:
+
+```cpp
+torch::Device device(torch::kCUDA);
+model->to(device);
+auto input = torch::randn({1, 3, 224, 224}).to(device);
+auto output = model->forward(input);
+```
+
+You can also check for CUDA availability before moving:
+
+```cpp
+torch::Device device(torch::cuda::is_available() ? torch::kCUDA : torch::kCPU);
+```
+
+Make sure to compile with the TorchScript headers by including `<torch/script.h>`.
+
+### My model runs slower in C++ than in Python
+
+**Problem**: Your model inference is slower in C++ compared to Python.
+
+**Fix**: There are several common causes:
+
+1. **Enable inference mode**: Wrap your inference code with `torch::NoGradGuard`
+   to disable gradient computation:
+
+```cpp
+torch::NoGradGuard no_grad;
+auto output = model->forward(input);
+```
+
+2. **Enable optimizations**: For TorchScript models, use `optimize_for_inference`:
+
+```cpp
+module = torch::jit::optimize_for_inference(module);
+```
+
+3. **Warm up the model**: Run a few inference passes before benchmarking to allow
+   JIT compilation and memory allocation to complete.
+
+4. **Check thread settings**: Ensure proper thread configuration:
+
+```cpp
+at::set_num_threads(4);  // Adjust based on your hardware
+```
+
+## Neural Network Modules
+
+### How do I register submodules in a custom module?
+
+**Problem**: You created a custom module but the submodules are not being
+recognized during `forward()` or when saving/loading the model.
+
+**Fix**: You must register submodules in the constructor using
+`register_module()`:
+
+```cpp
+struct MyModel : torch::nn::Module {
+  MyModel() {
+    fc1 = register_module("fc1", torch::nn::Linear(784, 128));
+    fc2 = register_module("fc2", torch::nn::Linear(128, 10));
+  }
+
+  torch::Tensor forward(torch::Tensor x) {
+    x = torch::relu(fc1->forward(x));
+    return fc2->forward(x);
+  }
+
+  torch::nn::Linear fc1{nullptr}, fc2{nullptr};
+};
+```
+
+### How do I set a module to evaluation mode?
+
+**Problem**: Layers like Dropout and BatchNorm behave differently during training
+and evaluation, and you need to switch between modes.
+
+**Fix**: Use the `eval()` and `train()` methods:
+
+```cpp
+model->eval();  // Set to evaluation mode
+// ... run inference ...
+model->train(); // Set back to training mode
+```
+
+## Data Loading
+
+### How do I create a custom dataset?
+
+**Problem**: You want to load your own data instead of using built-in datasets.
+
+**Fix**: Create a class that inherits from `torch::data::datasets::Dataset` and
+implement the `get()` and `size()` methods:
+
+```cpp
+class CustomDataset : public torch::data::datasets::Dataset<CustomDataset> {
+ public:
+  explicit CustomDataset(const std::string& data_path) {
+    // Load your data here
+  }
+
+  torch::data::Example<> get(size_t index) override {
+    // Return a single data sample
+    torch::Tensor data = /* load data at index */;
+    torch::Tensor label = /* load label at index */;
+    return {data, label};
+  }
+
+  torch::optional<size_t> size() const override {
+    return dataset_size_;
+  }
+
+ private:
+  size_t dataset_size_;
+};
+```
+
+Then use it with a DataLoader:
+
+```cpp
+auto dataset = CustomDataset("path/to/data")
+  .map(torch::data::transforms::Stack<>());
+auto dataloader = torch::data::make_data_loader(
+  std::move(dataset),
+  torch::data::DataLoaderOptions().batch_size(32).workers(4));
+```
+
+## Serialization
+
+### How do I save and load model weights?
+
+**Problem**: You want to save trained model weights and load them later.
+
+**Fix**: Use `torch::save()` and `torch::load()`:
+
+```cpp
+// Saving
+torch::save(model, "model.pt");
+
+// Loading
+torch::load(model, "model.pt");
+```
+
+For saving only specific tensors or state:
+
+```cpp
+torch::serialize::OutputArchive archive;
+model->save(archive);
+archive.save_to("model_weights.pt");
+
+// Loading
+torch::serialize::InputArchive archive;
+archive.load_from("model_weights.pt");
+model->load(archive);
+```
+
+## Build and Compilation
+
+### CMake cannot find Torch
+
+**Problem**: When building your project with CMake, you get an error that
+`Torch` package cannot be found.
+
+**Fix**: You need to specify the path to the LibTorch installation using
+`CMAKE_PREFIX_PATH`:
+
+```cpp
+cmake -DCMAKE_PREFIX_PATH=/path/to/libtorch ..
+```
+
+Alternatively, set `Torch_DIR` to point to the directory containing
+`TorchConfig.cmake`:
+
+```cpp
+cmake -DTorch_DIR=/path/to/libtorch/share/cmake/Torch ..
+```
+
+### Linker errors with undefined references
+
+**Problem**: Your project compiles but you get linker errors with undefined
+references to PyTorch symbols.
+
+**Fix**: Ensure you're linking against all required libraries in your
+`CMakeLists.txt`:
+
+```cpp
+find_package(Torch REQUIRED)
+add_executable(my_app main.cpp)
+target_link_libraries(my_app "${TORCH_LIBRARIES}")
+set_property(TARGET my_app PROPERTY CXX_STANDARD 17)
+```
+
+Also ensure that the compiler flags are set correctly:
+
+```cpp
+set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} ${TORCH_CXX_FLAGS}")
+```
diff --git a/docs/cpp/source/frontend.md b/docs/cpp/source/frontend.md
new file mode 100644
index 0000000000000..244e2c836bf98
--- /dev/null
+++ b/docs/cpp/source/frontend.md
@@ -0,0 +1,146 @@
+---
+myst:
+  html_meta:
+    description: PyTorch C++ Frontend guide — defining models, training loops, and using torch::nn modules in C++.
+    keywords: PyTorch, C++, frontend, torch::nn, Module, training, inference
+---
+
+# The C++ Frontend
+
+The PyTorch C++ frontend is a C++17 library for CPU and GPU
+tensor computation, with automatic differentiation and high level building
+blocks for state of the art machine learning applications.
+
+## Description
+
+The PyTorch C++ frontend can be thought of as a C++ version of the
+PyTorch Python frontend, providing automatic differentiation and various higher
+level abstractions for machine learning and neural networks.  Specifically,
+it consists of the following components:
+
+| Component | Description |
+| --- | --- |
+| `torch::Tensor` | Automatically differentiable, efficient CPU and GPU enabled tensors |
+| `torch::nn` | A collection of composable modules for neural network modeling |
+| `torch::optim` | Optimization algorithms like SGD, Adam or RMSprop to train your models |
+| `torch::data` | Datasets, data pipelines and multi-threaded, asynchronous data loader |
+| `torch::serialize` | A serialization API for storing and loading model checkpoints |
+| `torch::python` | Glue to bind your C++ models into Python |
+| `torch::jit` | Pure C++ access to the TorchScript JIT compiler |
+
+## End-to-end example
+
+Here is a simple, end-to-end example of defining and training a simple
+neural network on the MNIST dataset:
+
+```cpp
+#include <torch/torch.h>
+
+// Define a new Module.
+struct Net : torch::nn::Module {
+  Net() {
+    // Construct and register two Linear submodules.
+    fc1 = register_module("fc1", torch::nn::Linear(784, 64));
+    fc2 = register_module("fc2", torch::nn::Linear(64, 32));
+    fc3 = register_module("fc3", torch::nn::Linear(32, 10));
+  }
+
+  // Implement the Net's algorithm.
+  torch::Tensor forward(torch::Tensor x) {
+    // Use one of many tensor manipulation functions.
+    x = torch::relu(fc1->forward(x.reshape({x.size(0), 784})));
+    x = torch::dropout(x, /*p=*/0.5, /*train=*/is_training());
+    x = torch::relu(fc2->forward(x));
+    x = torch::log_softmax(fc3->forward(x), /*dim=*/1);
+    return x;
+  }
+
+  // Use one of many "standard library" modules.
+  torch::nn::Linear fc1{nullptr}, fc2{nullptr}, fc3{nullptr};
+};
+
+int main() {
+  // Create a new Net.
+  auto net = std::make_shared<Net>();
+
+  // Create a multi-threaded data loader for the MNIST dataset.
+  auto data_loader = torch::data::make_data_loader(
+      torch::data::datasets::MNIST("./data").map(
+          torch::data::transforms::Stack<>()),
+      /*batch_size=*/64);
+
+  // Instantiate an SGD optimization algorithm to update our Net's parameters.
+  torch::optim::SGD optimizer(net->parameters(), /*lr=*/0.01);
+
+  for (size_t epoch = 1; epoch <= 10; ++epoch) {
+    size_t batch_index = 0;
+    // Iterate the data loader to yield batches from the dataset.
+    for (auto& batch : *data_loader) {
+      // Reset gradients.
+      optimizer.zero_grad();
+      // Execute the model on the input data.
+      torch::Tensor prediction = net->forward(batch.data);
+      // Compute a loss value to judge the prediction of our model.
+      torch::Tensor loss = torch::nll_loss(prediction, batch.target);
+      // Compute gradients of the loss w.r.t. the parameters of our model.
+      loss.backward();
+      // Update the parameters based on the calculated gradients.
+      optimizer.step();
+      // Output the loss and checkpoint every 100 batches.
+      if (++batch_index % 100 == 0) {
+        std::cout << "Epoch: " << epoch << " | Batch: " << batch_index
+                  << " | Loss: " << loss.item<float>() << std::endl;
+        // Serialize your model periodically as a checkpoint.
+        torch::save(net, "net.pt");
+      }
+    }
+  }
+}
+```
+
+To see more complete examples of using the PyTorch C++ frontend, see [the example repository](https://github.com/pytorch/examples/tree/master/cpp).
+
+## Philosophy
+
+PyTorch's C++ frontend was designed with the idea that the Python frontend is
+great, and should be used when possible; but in some settings, performance and
+portability requirements make the use of the Python interpreter infeasible. For
+example, Python is a poor choice for low latency, high performance or
+multithreaded environments, such as video games or production servers.  The
+goal of the C++ frontend is to address these use cases, while not sacrificing
+the user experience of the Python frontend.
+
+As such, the C++ frontend has been written with a few philosophical goals in mind:
+
+* **Closely model the Python frontend in its design**, naming, conventions and
+  functionality.  While there may be occasional differences between the two
+  frontends (e.g., where we have dropped deprecated features or fixed "warts"
+  in the Python frontend), we guarantee that the effort in porting a Python model
+  to C++ should lie exclusively in **translating language features**,
+  not modifying functionality or behavior.
+
+* **Prioritize flexibility and user-friendliness over micro-optimization.**
+  In C++, you can often get optimal code, but at the cost of an extremely
+  unfriendly user experience.  Flexibility and dynamism is at the heart of
+  PyTorch, and the C++ frontend seeks to preserve this experience, in some
+  cases sacrificing performance (or "hiding" performance knobs) to keep APIs
+  simple and explicable.  We want researchers who don't write C++ for a living
+  to be able to use our APIs.
+
+A word of warning: Python is not necessarily slower than
+C++! The Python frontend calls into C++ for almost anything computationally expensive
+(especially any kind of numeric operation), and these operations will take up
+the bulk of time spent in a program.  If you would prefer to write Python,
+and can afford to write Python, we recommend using the Python interface to
+PyTorch. However, if you would prefer to write C++, or need to write C++
+(because of multithreading, latency or deployment requirements), the
+C++ frontend to PyTorch provides an API that is approximately as convenient,
+flexible, friendly and intuitive as its Python counterpart. The two frontends
+serve different use cases, work hand in hand, and neither is meant to
+unconditionally replace the other.
+
+## Installation
+
+Instructions on how to install the C++ frontend library distribution, including
+an example for how to build a minimal application depending on LibTorch, may be
+found by following [this](https://pytorch.org/cppdocs/installing.html) link.
diff --git a/docs/cpp/source/frontend.rst b/docs/cpp/source/frontend.rst
deleted file mode 100644
index 7a1776f7bd4a6..0000000000000
--- a/docs/cpp/source/frontend.rst
+++ /dev/null
@@ -1,153 +0,0 @@
-The C++ Frontend
-================
-
-The PyTorch C++ frontend is a C++17 library for CPU and GPU
-tensor computation, with automatic differentiation and high level building
-blocks for state of the art machine learning applications.
-
-Description
------------
-
-The PyTorch C++ frontend can be thought of as a C++ version of the
-PyTorch Python frontend, providing automatic differentiation and various higher
-level abstractions for machine learning and neural networks.  Specifically,
-it consists of the following components:
-
-+----------------------+------------------------------------------------------------------------+
-| Component            | Description                                                            |
-+======================+========================================================================+
-| ``torch::Tensor``    | Automatically differentiable, efficient CPU and GPU enabled tensors    |
-+----------------------+------------------------------------------------------------------------+
-| ``torch::nn``        | A collection of composable modules for neural network modeling         |
-+----------------------+------------------------------------------------------------------------+
-| ``torch::optim``     | Optimization algorithms like SGD, Adam or RMSprop to train your models |
-+----------------------+------------------------------------------------------------------------+
-| ``torch::data``      | Datasets, data pipelines and multi-threaded, asynchronous data loader  |
-+----------------------+------------------------------------------------------------------------+
-| ``torch::serialize`` | A serialization API for storing and loading model checkpoints          |
-+----------------------+------------------------------------------------------------------------+
-| ``torch::python``    | Glue to bind your C++ models into Python                               |
-+----------------------+------------------------------------------------------------------------+
-| ``torch::jit``       | Pure C++ access to the TorchScript JIT compiler                        |
-+----------------------+------------------------------------------------------------------------+
-
-End-to-end example
-------------------
-
-Here is a simple, end-to-end example of defining and training a simple
-neural network on the MNIST dataset:
-
-.. code-block:: cpp
-
-  #include <torch/torch.h>
-
-  // Define a new Module.
-  struct Net : torch::nn::Module {
-    Net() {
-      // Construct and register two Linear submodules.
-      fc1 = register_module("fc1", torch::nn::Linear(784, 64));
-      fc2 = register_module("fc2", torch::nn::Linear(64, 32));
-      fc3 = register_module("fc3", torch::nn::Linear(32, 10));
-    }
-
-    // Implement the Net's algorithm.
-    torch::Tensor forward(torch::Tensor x) {
-      // Use one of many tensor manipulation functions.
-      x = torch::relu(fc1->forward(x.reshape({x.size(0), 784})));
-      x = torch::dropout(x, /*p=*/0.5, /*train=*/is_training());
-      x = torch::relu(fc2->forward(x));
-      x = torch::log_softmax(fc3->forward(x), /*dim=*/1);
-      return x;
-    }
-
-    // Use one of many "standard library" modules.
-    torch::nn::Linear fc1{nullptr}, fc2{nullptr}, fc3{nullptr};
-  };
-
-  int main() {
-    // Create a new Net.
-    auto net = std::make_shared<Net>();
-
-    // Create a multi-threaded data loader for the MNIST dataset.
-    auto data_loader = torch::data::make_data_loader(
-        torch::data::datasets::MNIST("./data").map(
-            torch::data::transforms::Stack<>()),
-        /*batch_size=*/64);
-
-    // Instantiate an SGD optimization algorithm to update our Net's parameters.
-    torch::optim::SGD optimizer(net->parameters(), /*lr=*/0.01);
-
-    for (size_t epoch = 1; epoch <= 10; ++epoch) {
-      size_t batch_index = 0;
-      // Iterate the data loader to yield batches from the dataset.
-      for (auto& batch : *data_loader) {
-        // Reset gradients.
-        optimizer.zero_grad();
-        // Execute the model on the input data.
-        torch::Tensor prediction = net->forward(batch.data);
-        // Compute a loss value to judge the prediction of our model.
-        torch::Tensor loss = torch::nll_loss(prediction, batch.target);
-        // Compute gradients of the loss w.r.t. the parameters of our model.
-        loss.backward();
-        // Update the parameters based on the calculated gradients.
-        optimizer.step();
-        // Output the loss and checkpoint every 100 batches.
-        if (++batch_index % 100 == 0) {
-          std::cout << "Epoch: " << epoch << " | Batch: " << batch_index
-                    << " | Loss: " << loss.item<float>() << std::endl;
-          // Serialize your model periodically as a checkpoint.
-          torch::save(net, "net.pt");
-        }
-      }
-    }
-  }
-
-To see more complete examples of using the PyTorch C++ frontend, see `the example repository
-<https://github.com/pytorch/examples/tree/master/cpp>`_.
-
-Philosophy
-----------
-
-PyTorch's C++ frontend was designed with the idea that the Python frontend is
-great, and should be used when possible; but in some settings, performance and
-portability requirements make the use of the Python interpreter infeasible. For
-example, Python is a poor choice for low latency, high performance or
-multithreaded environments, such as video games or production servers.  The
-goal of the C++ frontend is to address these use cases, while not sacrificing
-the user experience of the Python frontend.
-
-As such, the C++ frontend has been written with a few philosophical goals in mind:
-
-* **Closely model the Python frontend in its design**, naming, conventions and
-  functionality.  While there may be occasional differences between the two
-  frontends (e.g., where we have dropped deprecated features or fixed "warts"
-  in the Python frontend), we guarantee that the effort in porting a Python model
-  to C++ should lie exclusively in **translating language features**,
-  not modifying functionality or behavior.
-
-* **Prioritize flexibility and user-friendliness over micro-optimization.**
-  In C++, you can often get optimal code, but at the cost of an extremely
-  unfriendly user experience.  Flexibility and dynamism is at the heart of
-  PyTorch, and the C++ frontend seeks to preserve this experience, in some
-  cases sacrificing performance (or "hiding" performance knobs) to keep APIs
-  simple and explicable.  We want researchers who don't write C++ for a living
-  to be able to use our APIs.
-
-A word of warning: Python is not necessarily slower than
-C++! The Python frontend calls into C++ for almost anything computationally expensive
-(especially any kind of numeric operation), and these operations will take up
-the bulk of time spent in a program.  If you would prefer to write Python,
-and can afford to write Python, we recommend using the Python interface to
-PyTorch. However, if you would prefer to write C++, or need to write C++
-(because of multithreading, latency or deployment requirements), the
-C++ frontend to PyTorch provides an API that is approximately as convenient,
-flexible, friendly and intuitive as its Python counterpart. The two frontends
-serve different use cases, work hand in hand, and neither is meant to
-unconditionally replace the other.
-
-Installation
-------------
-
-Instructions on how to install the C++ frontend library distribution, including
-an example for how to build a minimal application depending on LibTorch, may be
-found by following `this <https://pytorch.org/cppdocs/installing.html>`_ link.
diff --git a/docs/cpp/source/index.rst b/docs/cpp/source/index.md
similarity index 54%
rename from docs/cpp/source/index.rst
rename to docs/cpp/source/index.md
index fb7e5986e869a..1979721ad7a94 100644
--- a/docs/cpp/source/index.rst
+++ b/docs/cpp/source/index.md
@@ -1,5 +1,11 @@
-PyTorch C++ API
-===============
+---
+myst:
+  html_meta:
+    description: PyTorch C++ API documentation — ATen tensors, Autograd, C++ Frontend, TorchScript, and C++ Extensions.
+    keywords: PyTorch, C++, API, libtorch, ATen, Autograd, TorchScript, C++ Frontend
+---
+
+# PyTorch C++ API
 
 These pages provide the documentation for the public portions of the PyTorch C++
 API.  This API can roughly be divided into five parts:
@@ -16,108 +22,99 @@ networks with strong emphasis on GPU acceleration as well as fast CPU
 performance. It is currently in use at Facebook in research and
 production; we are looking forward to welcoming more users of the PyTorch C++ API.
 
-.. warning::
+```{warning}
 
-  At the moment, the C++ API should be considered "beta" stability; we may
-  make major breaking changes to the backend in order to improve the API,
-  or in service of providing the Python interface to PyTorch, which is our
-  most stable and best supported interface.
+At the moment, the C++ API should be considered "beta" stability; we may
+make major breaking changes to the backend in order to improve the API,
+or in service of providing the Python interface to PyTorch, which is our
+most stable and best supported interface.
+```
 
-ATen
-----
+## ATen
 
 ATen is fundamentally a tensor library, on top of which almost all other Python
-and C++ interfaces in PyTorch are built. It provides a core ``Tensor`` class,
+and C++ interfaces in PyTorch are built. It provides a core `Tensor` class,
 on which many hundreds of operations are defined. Most of these operations have
-both CPU and GPU implementations, to which the ``Tensor`` class will
+both CPU and GPU implementations, to which the `Tensor` class will
 dynamically dispatch based on its type. A small example of using ATen could
 look as follows:
 
-.. code-block:: cpp
-
-  #include <ATen/ATen.h>
+```cpp
+#include <ATen/ATen.h>
 
-  at::Tensor a = at::ones({2, 2}, at::kInt);
-  at::Tensor b = at::randn({2, 2});
-  auto c = a + b.to(at::kInt);
+at::Tensor a = at::ones({2, 2}, at::kInt);
+at::Tensor b = at::randn({2, 2});
+auto c = a + b.to(at::kInt);
+```
 
-This ``Tensor`` class and all other symbols in ATen are found in the ``at::``
+This `Tensor` class and all other symbols in ATen are found in the `at::`
 namespace, documented
-`here <https://pytorch.org/cppdocs/api/namespace_at.html#namespace-at>`_.
+[here](https://pytorch.org/cppdocs/api/namespace_at.html#namespace-at).
 
-Autograd
---------
+## Autograd
 
 What we term *autograd* are the portions of PyTorch's C++ API that augment the
-ATen ``Tensor`` class with capabilities concerning automatic differentiation.
+ATen `Tensor` class with capabilities concerning automatic differentiation.
 The autograd system records operations on tensors to form an *autograd graph*.
-Calling ``backwards()`` on a leaf variable in this graph performs reverse mode
+Calling `backwards()` on a leaf variable in this graph performs reverse mode
 differentiation through the network of functions and tensors spanning the
 autograd graph, ultimately yielding gradients. The following example provides
 a taste of this interface:
 
-.. code-block:: cpp
-
-  #include <torch/csrc/autograd/variable.h>
-  #include <torch/csrc/autograd/function.h>
+```cpp
+#include <torch/csrc/autograd/variable.h>
+#include <torch/csrc/autograd/function.h>
 
-  torch::Tensor a = torch::ones({2, 2}, torch::requires_grad());
-  torch::Tensor b = torch::randn({2, 2});
-  auto c = a + b;
-  c.backward(); // a.grad() will now hold the gradient of c w.r.t. a.
+torch::Tensor a = torch::ones({2, 2}, torch::requires_grad());
+torch::Tensor b = torch::randn({2, 2});
+auto c = a + b;
+c.backward(); // a.grad() will now hold the gradient of c w.r.t. a.
+```
 
-The ``at::Tensor`` class in ATen is not differentiable by default. To add the
+The `at::Tensor` class in ATen is not differentiable by default. To add the
 differentiability of tensors the autograd API provides, you must use tensor
 factory functions from the `torch::` namespace instead of the `at::` namespace.
 For example, while a tensor created with `at::ones` will not be differentiable,
 a tensor created with `torch::ones` will be.
 
-C++ Frontend
-------------
+## C++ Frontend
 
 The PyTorch C++ frontend provides a high level, pure C++ modeling interface for
 neural networks and general ML (Machine Learning) research and production use cases,
 largely following the Python API in design and provided functionality. The C++
 frontend includes the following:
 
-- An interface for defining machine learning models through a hierarchical module system (like ``torch.nn.Module``);
+- An interface for defining machine learning models through a hierarchical module system (like `torch.nn.Module`);
 - A "standard library" of pre-existing modules for the most common modeling purposes (e.g. convolutions, RNNs, batch normalization etc.);
 - An optimization API, including implementations of popular optimizers such as SGD, Adam, RMSprop and others;
 - A means of representing datasets and data pipelines, including functionality to load data in parallel over many CPU cores;
-- A serialization format for storing and loading checkpoints of a training session (like ``torch.utils.data.DataLoader``);
-- Automatic parallelization of models onto multiple GPUs (like ``torch.nn.parallel.DataParallel``);
+- A serialization format for storing and loading checkpoints of a training session (like `torch.utils.data.DataLoader`);
+- Automatic parallelization of models onto multiple GPUs (like `torch.nn.parallel.DataParallel`);
 - Support code to easily bind C++ models into Python using pybind11;
 - Entry points to the TorchScript JIT compiler;
 - Helpful utilities to facilitate interfacing with the ATen and Autograd APIs.
 
-See `this document <https://pytorch.org/cppdocs/frontend.html>`_ for a more
+See [this document](https://pytorch.org/cppdocs/frontend.html) for a more
 detailed description of the C++ frontend. Relevant sections of the `torch::`
-namespace related to the C++ Frontend include `torch::nn
-<https://pytorch.org/cppdocs/api/namespace_torch__nn.html#namespace-torch-nn>`_,
-`torch::optim
-<https://pytorch.org/cppdocs/api/namespace_torch__optim.html#namespace-torch-optim>`_,
-`torch::data
-<https://pytorch.org/cppdocs/api/namespace_torch__data.html#namespace-torch-data>`_,
-`torch::serialize
-<https://pytorch.org/cppdocs/api/namespace_torch__serialize.html#namespace-torch-serialize>`_,
-`torch::jit
-<https://pytorch.org/cppdocs/api/namespace_torch__jit.html#namespace-torch-jit>`_
-and `torch::python
-<https://pytorch.org/cppdocs/api/namespace_torch__python.html#namespace-torch-python>`_.
-Examples of the C++ frontend can be found in `this repository
-<https://github.com/pytorch/examples/tree/master/cpp>`_ which is being
+namespace related to the C++ Frontend include [torch::nn](https://pytorch.org/cppdocs/api/namespace_torch__nn.html#namespace-torch-nn),
+[torch::optim](https://pytorch.org/cppdocs/api/namespace_torch__optim.html#namespace-torch-optim),
+[torch::data](https://pytorch.org/cppdocs/api/namespace_torch__data.html#namespace-torch-data),
+[torch::serialize](https://pytorch.org/cppdocs/api/namespace_torch__serialize.html#namespace-torch-serialize),
+[torch::jit](https://pytorch.org/cppdocs/api/namespace_torch__jit.html#namespace-torch-jit)
+and [torch::python](https://pytorch.org/cppdocs/api/namespace_torch__python.html#namespace-torch-python).
+Examples of the C++ frontend can be found in [this repository](https://github.com/pytorch/examples/tree/master/cpp) which is being
 expanded on a continuous and active basis.
 
-.. note::
+```{note}
 
-  Unless you have a particular reason to constrain yourself exclusively to ATen
-  or the Autograd API, the C++ frontend is the recommended entry point to the
-  PyTorch C++ ecosystem. While it is still in beta as we collect user feedback
-  (from you!), it provides both more functionality and better stability
-  guarantees than the ATen and Autograd APIs.
+Unless you have a particular reason to constrain yourself exclusively to ATen
+or the Autograd API, the C++ frontend is the recommended entry point to the
+PyTorch C++ ecosystem. While it is still in beta as we collect user feedback
+(from you!), it provides both more functionality and better stability
+guarantees than the ATen and Autograd APIs.
+```
 
-TorchScript
------------
+## TorchScript
 
 TorchScript is a representation of a PyTorch model that can be understood,
 compiled and serialized by the TorchScript compiler. Fundamentally, TorchScript
@@ -132,16 +129,13 @@ functionality:
 The first mechanism may be of great interest to you if you would like to define
 your models in Python as much as possible, but subsequently export them to C++
 for production environments and no-Python inference. You can find out more
-about this by following `this
-<https://pytorch.org/tutorials/advanced/cpp_export.html>`_ link. The second
+about this by following [this](https://pytorch.org/tutorials/advanced/cpp_export.html) link. The second
 API concerns itself with scenarios in which you would like to extend
 TorchScript with custom operators, which can similarly be serialized and
-invoked from C++ during inference. Lastly, the `torch::jit::compile
-<https://pytorch.org/cppdocs/api/function_namespacetorch_1_1jit_1a8660dc13a6b82336aadac667e6dccba1.html>`_
+invoked from C++ during inference. Lastly, the [torch::jit::compile](https://pytorch.org/cppdocs/api/function_namespacetorch_1_1jit_1a8660dc13a6b82336aadac667e6dccba1.html)
 function may be used to access the TorchScript compiler directly from C++.
 
-C++ Extensions
---------------
+## C++ Extensions
 
 *C++ Extensions* offer a simple yet powerful way of accessing all of the above
 interfaces for the purpose of extending regular Python use-cases of PyTorch.
@@ -151,37 +145,26 @@ does not add any new functionality to the PyTorch C++ API. Instead, it
 provides integration with Python setuptools as well as JIT compilation
 mechanisms that allow access to ATen, the autograd and other C++ APIs from
 Python. To learn more about the C++ extension API, go through
-`this tutorial <https://pytorch.org/tutorials/advanced/cpp_extension.html>`_.
-
-Contents
---------
-
-.. toctree::
-   :maxdepth: 2
+[this tutorial](https://pytorch.org/tutorials/advanced/cpp_extension.html).
 
-   installing
-   frontend
-   stable
-   api/library_root
+## Contents
 
-.. toctree::
-  :glob:
-  :maxdepth: 1
-  :caption: Notes
+```{toctree}
+:maxdepth: 2
 
-  notes/*
+installing
+frontend
+api/index
+faq
+```
 
-Indices and tables
-==================
+# Indices and tables
 
-* :ref:`genindex`
-* :ref:`modindex`
-* :ref:`search`
+* {ref}`genindex`
+* {ref}`modindex`
+* {ref}`search`
 
-Acknowledgements
-----------------
+## Acknowledgements
 
-This documentation website for the PyTorch C++ universe has been enabled by the
-`Exhale <https://github.com/svenevs/exhale/>`_ project and generous investment
-of time and effort by its maintainer, `svenevs <https://github.com/svenevs/>`_.
-We thank Stephen for his work and his efforts providing help with the PyTorch C++ documentation.
+This documentation website for the PyTorch C++ universe uses the Sphinx
+C++ domain for API documentation.
diff --git a/docs/cpp/source/installing.md b/docs/cpp/source/installing.md
new file mode 100644
index 0000000000000..7fb9067b9611f
--- /dev/null
+++ b/docs/cpp/source/installing.md
@@ -0,0 +1,187 @@
+---
+myst:
+  html_meta:
+    description: How to install and configure the PyTorch C++ API (libtorch) for development.
+    keywords: PyTorch, C++, install, libtorch, setup, CMake, build
+---
+
+# Installing C++ Distributions of PyTorch
+
+We provide binary distributions of all headers, libraries and CMake
+configuration files required to depend on PyTorch. We call this distribution
+*LibTorch*, and you can download ZIP archives containing the latest LibTorch
+distribution on [our website](https://pytorch.org/get-started/locally/). Below
+is a small example of writing a minimal application that depends on LibTorch
+and uses the `torch::Tensor` class which comes with the PyTorch C++ API.
+
+## Minimal Example
+
+The first step is to download the LibTorch ZIP archive via the link above. For
+example:
+
+```sh
+wget https://download.pytorch.org/libtorch/nightly/cpu/libtorch-shared-with-deps-latest.zip
+unzip libtorch-shared-with-deps-latest.zip
+```
+
+Note that the above link has CPU-only libtorch. If you would like to download a GPU-enabled
+libtorch, find the right link in the link selector on https://pytorch.org
+
+If you're a Windows developer and wouldn't like to use CMake, you could jump to the Visual Studio
+Extension section.
+
+Next, we can write a minimal CMake build configuration to develop a small
+application that depends on LibTorch. CMake is not a hard requirement for using
+LibTorch, but it is the recommended and blessed build system and will be well
+supported into the future. A most basic `CMakeLists.txt` file could look like
+this:
+
+```cmake
+cmake_minimum_required(VERSION 3.18 FATAL_ERROR)
+project(example-app)
+
+find_package(Torch REQUIRED)
+set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} ${TORCH_CXX_FLAGS}")
+
+add_executable(example-app example-app.cpp)
+target_link_libraries(example-app "${TORCH_LIBRARIES}")
+set_property(TARGET example-app PROPERTY CXX_STANDARD 20)
+
+# The following code block is suggested to be used on Windows.
+# According to https://github.com/pytorch/pytorch/issues/25457,
+# the DLLs need to be copied to avoid memory errors.
+if (MSVC)
+  file(GLOB TORCH_DLLS "${TORCH_INSTALL_PREFIX}/lib/*.dll")
+  add_custom_command(TARGET example-app
+                     POST_BUILD
+                     COMMAND ${CMAKE_COMMAND} -E copy_if_different
+                     ${TORCH_DLLS}
+                     $<TARGET_FILE_DIR:example-app>)
+endif (MSVC)
+```
+
+The implementation of our example will simply create a new `torch::Tensor` and
+print it:
+
+```cpp
+#include <torch/torch.h>
+#include <iostream>
+
+int main() {
+  torch::Tensor tensor = torch::rand({2, 3});
+  std::cout << tensor << std::endl;
+}
+```
+
+While there are more fine-grained headers you can include to access only parts
+of the PyTorch C++ API, including `torch/torch.h` is the most sure-proof way of
+including most of its functionality.
+
+The last step is to build the application. For this, assume our example
+directory is laid out like this:
+
+```sh
+example-app/
+  CMakeLists.txt
+  example-app.cpp
+```
+
+We can now run the following commands to build the application from within the
+`example-app/` folder:
+
+```sh
+mkdir build
+cd build
+cmake -DCMAKE_PREFIX_PATH=/absolute/path/to/libtorch ..
+cmake --build . --config Release
+```
+
+where `/absolute/path/to/libtorch` should be the absolute (!) path to the unzipped LibTorch
+distribution. If PyTorch was installed via pip, `CMAKE_PREFIX_PATH` can be queried
+using `torch.utils.cmake_prefix_path` variable. In that case CMake configuration step would look something like follows:
+
+```sh
+cmake -DCMAKE_PREFIX_PATH=`python3 -c 'import torch;print(torch.utils.cmake_prefix_path)'` ..
+```
+
+If all goes well, it will look something like this:
+
+```sh
+root@4b5a67132e81:/example-app# mkdir build
+root@4b5a67132e81:/example-app# cd build
+root@4b5a67132e81:/example-app/build# cmake -DCMAKE_PREFIX_PATH=/path/to/libtorch ..
+-- The C compiler identification is GNU 5.4.0
+-- The CXX compiler identification is GNU 5.4.0
+-- Check for working C compiler: /usr/bin/cc
+-- Check for working C compiler: /usr/bin/cc -- works
+-- Detecting C compiler ABI info
+-- Detecting C compiler ABI info - done
+-- Detecting C compile features
+-- Detecting C compile features - done
+-- Check for working CXX compiler: /usr/bin/c++
+-- Check for working CXX compiler: /usr/bin/c++ -- works
+-- Detecting CXX compiler ABI info
+-- Detecting CXX compiler ABI info - done
+-- Detecting CXX compile features
+-- Detecting CXX compile features - done
+-- Looking for pthread.h
+-- Looking for pthread.h - found
+-- Looking for pthread_create
+-- Looking for pthread_create - not found
+-- Looking for pthread_create in pthreads
+-- Looking for pthread_create in pthreads - not found
+-- Looking for pthread_create in pthread
+-- Looking for pthread_create in pthread - found
+-- Found Threads: TRUE
+-- Configuring done
+-- Generating done
+-- Build files have been written to: /example-app/build
+root@4b5a67132e81:/example-app/build# cmake --build . --config Release
+Scanning dependencies of target example-app
+[ 50%] Building CXX object CMakeFiles/example-app.dir/example-app.cpp.o
+[100%] Linking CXX executable example-app
+[100%] Built target example-app
+```
+
+Executing the resulting `example-app` binary found in the `build` folder
+should now merrily print the tensor (exact output subject to randomness):
+
+```sh
+root@4b5a67132e81:/example-app/build# ./example-app
+0.2063  0.6593  0.0866
+0.0796  0.5841  0.1569
+[ Variable[CPUFloatType]{2,3} ]
+```
+
+```{tip}
+
+On Windows, debug and release builds are not ABI-compatible. If you plan to
+build your project in debug mode, please try the debug version of LibTorch.
+Also, make sure you specify the correct configuration in the `cmake --build .`
+line above.
+```
+
+## System Requirements
+
+To ensure smooth installation and usage of LibTorch, please ensure your system
+meets the following requirements:
+
+1. **GLIBC Version**:
+
+   - GLIBC 2.29 or newer for cxx11 ABI version
+
+2. **GCC Version**:
+
+   - GCC 9 or newer for cxx11
+
+## Visual Studio Extension
+
+[LibTorch Project Template](https://marketplace.visualstudio.com/items?itemName=YiZhang.LibTorch001) can help Windows developers
+set all libtorch project settings and link options for debug and release.
+It's easy to use and you could check out the [demo video](https://ossci-windows.s3.us-east-1.amazonaws.com/vsextension/demo.mp4).
+The only prerequisite is to download the libtorch on https://pytorch.org
+
+## Support
+
+If you run into any troubles with this installation and minimal usage guide,
+please use our [forum](https://discuss.pytorch.org/) or [GitHub issues](https://github.com/pytorch/pytorch/issues) to get in touch.
diff --git a/docs/cpp/source/installing.rst b/docs/cpp/source/installing.rst
deleted file mode 100644
index 0612d2592f21a..0000000000000
--- a/docs/cpp/source/installing.rst
+++ /dev/null
@@ -1,182 +0,0 @@
-Installing C++ Distributions of PyTorch
-=======================================
-
-We provide binary distributions of all headers, libraries and CMake
-configuration files required to depend on PyTorch. We call this distribution
-*LibTorch*, and you can download ZIP archives containing the latest LibTorch
-distribution on `our website <https://pytorch.org/get-started/locally/>`_. Below
-is a small example of writing a minimal application that depends on LibTorch
-and uses the ``torch::Tensor`` class which comes with the PyTorch C++ API.
-
-Minimal Example
----------------
-
-The first step is to download the LibTorch ZIP archive via the link above. For
-example:
-
-.. code-block:: sh
-
-  wget https://download.pytorch.org/libtorch/nightly/cpu/libtorch-shared-with-deps-latest.zip
-  unzip libtorch-shared-with-deps-latest.zip
-
-Note that the above link has CPU-only libtorch. If you would like to download a GPU-enabled
-libtorch, find the right link in the link selector on https://pytorch.org
-
-If you're a Windows developer and wouldn't like to use CMake, you could jump to the Visual Studio
-Extension section.
-
-Next, we can write a minimal CMake build configuration to develop a small
-application that depends on LibTorch. CMake is not a hard requirement for using
-LibTorch, but it is the recommended and blessed build system and will be well
-supported into the future. A most basic `CMakeLists.txt` file could look like
-this:
-
-.. code-block:: cmake
-
-  cmake_minimum_required(VERSION 3.18 FATAL_ERROR)
-  project(example-app)
-
-  find_package(Torch REQUIRED)
-  set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} ${TORCH_CXX_FLAGS}")
-
-  add_executable(example-app example-app.cpp)
-  target_link_libraries(example-app "${TORCH_LIBRARIES}")
-  set_property(TARGET example-app PROPERTY CXX_STANDARD 17)
-
-  # The following code block is suggested to be used on Windows.
-  # According to https://github.com/pytorch/pytorch/issues/25457,
-  # the DLLs need to be copied to avoid memory errors.
-  if (MSVC)
-    file(GLOB TORCH_DLLS "${TORCH_INSTALL_PREFIX}/lib/*.dll")
-    add_custom_command(TARGET example-app
-                       POST_BUILD
-                       COMMAND ${CMAKE_COMMAND} -E copy_if_different
-                       ${TORCH_DLLS}
-                       $<TARGET_FILE_DIR:example-app>)
-  endif (MSVC)
-
-The implementation of our example will simply create a new `torch::Tensor` and
-print it:
-
-.. code-block:: cpp
-
-  #include <torch/torch.h>
-  #include <iostream>
-
-  int main() {
-    torch::Tensor tensor = torch::rand({2, 3});
-    std::cout << tensor << std::endl;
-  }
-
-While there are more fine-grained headers you can include to access only parts
-of the PyTorch C++ API, including `torch/torch.h` is the most sure-proof way of
-including most of its functionality.
-
-The last step is to build the application. For this, assume our example
-directory is laid out like this:
-
-.. code-block:: sh
-
-  example-app/
-    CMakeLists.txt
-    example-app.cpp
-
-We can now run the following commands to build the application from within the
-``example-app/`` folder:
-
-.. code-block:: sh
-
-  mkdir build
-  cd build
-  cmake -DCMAKE_PREFIX_PATH=/absolute/path/to/libtorch ..
-  cmake --build . --config Release
-
-where ``/absolute/path/to/libtorch`` should be the absolute (!) path to the unzipped LibTorch
-distribution. If PyTorch was installed via pip, `CMAKE_PREFIX_PATH` can be queried
-using `torch.utils.cmake_prefix_path` variable. In that case CMake configuration step would look something like follows:
-
-.. code-block:: sh
-
-  cmake -DCMAKE_PREFIX_PATH=`python3 -c 'import torch;print(torch.utils.cmake_prefix_path)'` ..
-
-If all goes well, it will look something like this:
-
-.. code-block:: sh
-
-  root@4b5a67132e81:/example-app# mkdir build
-  root@4b5a67132e81:/example-app# cd build
-  root@4b5a67132e81:/example-app/build# cmake -DCMAKE_PREFIX_PATH=/path/to/libtorch ..
-  -- The C compiler identification is GNU 5.4.0
-  -- The CXX compiler identification is GNU 5.4.0
-  -- Check for working C compiler: /usr/bin/cc
-  -- Check for working C compiler: /usr/bin/cc -- works
-  -- Detecting C compiler ABI info
-  -- Detecting C compiler ABI info - done
-  -- Detecting C compile features
-  -- Detecting C compile features - done
-  -- Check for working CXX compiler: /usr/bin/c++
-  -- Check for working CXX compiler: /usr/bin/c++ -- works
-  -- Detecting CXX compiler ABI info
-  -- Detecting CXX compiler ABI info - done
-  -- Detecting CXX compile features
-  -- Detecting CXX compile features - done
-  -- Looking for pthread.h
-  -- Looking for pthread.h - found
-  -- Looking for pthread_create
-  -- Looking for pthread_create - not found
-  -- Looking for pthread_create in pthreads
-  -- Looking for pthread_create in pthreads - not found
-  -- Looking for pthread_create in pthread
-  -- Looking for pthread_create in pthread - found
-  -- Found Threads: TRUE
-  -- Configuring done
-  -- Generating done
-  -- Build files have been written to: /example-app/build
-  root@4b5a67132e81:/example-app/build# cmake --build . --config Release
-  Scanning dependencies of target example-app
-  [ 50%] Building CXX object CMakeFiles/example-app.dir/example-app.cpp.o
-  [100%] Linking CXX executable example-app
-  [100%] Built target example-app
-
-Executing the resulting ``example-app`` binary found in the ``build`` folder
-should now merrily print the tensor (exact output subject to randomness):
-
-.. code-block:: sh
-
-  root@4b5a67132e81:/example-app/build# ./example-app
-  0.2063  0.6593  0.0866
-  0.0796  0.5841  0.1569
-  [ Variable[CPUFloatType]{2,3} ]
-
-.. tip::
-  On Windows, debug and release builds are not ABI-compatible. If you plan to
-  build your project in debug mode, please try the debug version of LibTorch.
-  Also, make sure you specify the correct configuration in the ``cmake --build .``
-  line above.
-
-System Requirements
--------------------
-
-To ensure smooth installation and usage of LibTorch, please ensure your system
-meets the following requirements:
-
-1. **GLIBC Version**:
-  - GLIBC 2.29 or newer for cxx11 ABI version
-
-2. **GCC Version**:
-  - GCC 9 or newer for cxx11
-
-Visual Studio Extension
------------------------
-
-`LibTorch Project Template <https://marketplace.visualstudio.com/items?itemName=YiZhang.LibTorch001>`_ can help Windows developers
-set all libtorch project settings and link options for debug and release.
-It's easy to use and you could check out the `demo video <https://ossci-windows.s3.us-east-1.amazonaws.com/vsextension/demo.mp4>`_.
-The only prerequisite is to download the libtorch on https://pytorch.org
-
-Support
--------
-
-If you run into any troubles with this installation and minimal usage guide,
-please use our `forum <https://discuss.pytorch.org/>`_ or `GitHub issues
-<https://github.com/pytorch/pytorch/issues>`_ to get in touch.
diff --git a/docs/cpp/source/library.rst b/docs/cpp/source/library.rst
deleted file mode 100644
index 5cbaaf959910b..0000000000000
--- a/docs/cpp/source/library.rst
+++ /dev/null
@@ -1,37 +0,0 @@
-Torch Library API
-=================
-
-The PyTorch C++ API provides capabilities for extending PyTorch's core library
-of operators with user defined operators and data types.  Extensions implemented
-using the Torch Library API are made available for use in both the PyTorch eager
-API as well as in TorchScript.
-
-For a tutorial style introduction to the library API, check out the
-`Extending TorchScript with Custom C++ Operators
-<https://pytorch.org/tutorials/advanced/torch_script_custom_ops.html>`_
-tutorial.
-
-Macros
-------
-
-.. doxygendefine:: TORCH_LIBRARY
-
-.. doxygendefine:: TORCH_LIBRARY_IMPL
-
-Classes
--------
-
-.. doxygenclass:: torch::Library
-  :members:
-
-.. doxygenclass:: torch::CppFunction
-  :members:
-
-Functions
----------
-
-.. doxygengroup:: torch-dispatch-overloads
-  :content-only:
-
-.. doxygengroup:: torch-schema-overloads
-  :content-only:
diff --git a/docs/cpp/source/notes/faq.rst b/docs/cpp/source/notes/faq.rst
deleted file mode 100644
index 37a1a609f3a5f..0000000000000
--- a/docs/cpp/source/notes/faq.rst
+++ /dev/null
@@ -1,34 +0,0 @@
-FAQ
----
-
-Listed below are a number of common issues users face with the various parts of
-the C++ API.
-
-C++ Extensions
-==============
-
-Undefined symbol errors from PyTorch/ATen
-*****************************************
-
-**Problem**: You import your extension and get an ``ImportError`` stating that
-some C++ symbol from PyTorch or ATen is undefined. For example::
-
-  >>> import extension
-  Traceback (most recent call last):
-    File "<stdin>", line 1, in <module>
-  ImportError: /home/user/.pyenv/versions/3.7.1/lib/python3.7/site-packages/extension.cpython-37m-x86_64-linux-gnu.so: undefined symbol: _ZN2at19UndefinedTensorImpl10_singletonE
-
-**Fix**: The fix is to ``import torch`` before you import your extension. This will make
-the symbols from the PyTorch dynamic (shared) library that your extension
-depends on available, allowing them to be resolved once you import your extension.
-
-I created a tensor using a function from ``at::`` and get errors
-****************************************************************
-
-**Problem**: You created a tensor using e.g. ``at::ones`` or ``at::randn`` or
-any other tensor factory from the ``at::`` namespace and are getting errors.
-
-**Fix**: Replace ``at::`` with ``torch::`` for factory function calls. You
-should never use factory functions from the ``at::`` namespace, as they will
-create tensors. The corresponding ``torch::`` functions will create variables,
-and you should only ever deal with variables in your code.
diff --git a/docs/cpp/source/notes/inference_mode.rst b/docs/cpp/source/notes/inference_mode.rst
deleted file mode 100644
index 60cc48ca93d5b..0000000000000
--- a/docs/cpp/source/notes/inference_mode.rst
+++ /dev/null
@@ -1,121 +0,0 @@
-Inference Mode
-==============
-
-``c10::InferenceMode`` is a new RAII guard analogous to ``NoGradMode``
-to be used when you are certain your operations will have no interactions
-with autograd (e.g. model training). Compared to ``NoGradMode``, code run
-under this mode gets better performance by disabling autograd related work like
-view tracking and version counter bumps. However, tensors created inside
-``c10::InferenceMode`` have more limitations when interacting with autograd system as well.
-
-``InferenceMode`` can be enabled for a given block of code. Inside ``InferenceMode``
-all newly allocated (non-view) tensors are marked as inference tensors. Inference tensors:
-
-- do not have a version counter so an error will be raised if you try to read their version
-  (e.g., because you saved this tensor for backward).
-- are immutable outside ``InferenceMode``. So an error will be raised if you try to:
-  - mutate their data outside InferenceMode.
-  - mutate them into ``requires_grad=True`` outside InferenceMode.
-  To work around you can make a clone outside ``InferenceMode`` to get a normal tensor before mutating.
-
-A non-view tensor is an inference tensor if and only if it was allocated inside ``InferenceMode``.
-A view tensor is an inference tensor if and only if it is a view of an inference tensor.
-
-Inside an ``InferenceMode`` block, we make the following performance guarantees:
-
-- Like ``NoGradMode``, all operations do not record ``grad_fn`` even if their inputs have ``requires_grad=True``.
-  This applies to both inference tensors and normal tensors.
-- View operations on inference tensors do not do view tracking. View and non-view inference tensors are
-  indistinguishable.
-- Inplace operations on inference tensors are guaranteed not to do a version bump.
-
-For more implementation details of ``InferenceMode`` please see the `RFC-0011-InferenceMode <https://github.com/pytorch/rfcs/pull/17>`_.
-
-Migration guide from ``AutoNonVariableTypeMode``
-------------------------------------------------
-
-In production use of PyTorch for inference workload, we have seen a proliferation
-of uses of the C++ guard ``AutoNonVariableTypeMode`` (now ``AutoDispatchBelowADInplaceOrView``),
-which disables autograd, view tracking and version counter bumps. Unfortunately,
-current colloquial of this guard for inference workload is unsafe: it's possible to
-use ``AutoNonVariableTypeMode`` to bypass PyTorch's safety checks and result in
-silently wrong results, e.g. PyTorch throws an error when tensors saved for backwards
-are subsequently mutated, but mutation happens inside ``AutoNonVariableTypeMode`` will
-silently bypass the check and returns wrong gradient to users.
-
-When current users of ``AutoNonVariableTypeMode`` think about migrating, the following
-steps might help you decide the best alternatives:
-
-1. Users trying to run workload in inference only mode (like loading a pretrained JIT model and
-   run inference in C++ runtime) should add ``c10::InferenceMode guard`` to guard all operations
-   on tensors (including model loading). See an inference workload example below:
-
-.. code-block:: cpp
-
-  c10::InferenceMode guard;
-  model.load_jit(saved_model);
-  auto inputs = preprocess_tensors(data);
-  auto out = model.forward(inputs);
-  auto outputs = postprocess_tensors(out);
-
-Note ``c10::InferenceMode`` offers a drop in replacement for ``AutoNonVariableTypeMode`` which preserves
-the performance characteristics of ``AutoNonVariableTypeMode``. But they also have some differences that
-users should pay additional attention to:
-
-  - Both guards affects tensor execution process to skip work not related to inference, but ``InferenceMode``
-    also affects tensor creation while ``AutoNonVariableTypeMode`` doesn't. In other words, tensors created
-    inside ``InferenceMode`` are marked as inference tensors so that certain limitations can be applied after
-    exiting ``InferenceMode``.
-  - Enabled/disabled ``InferenceMode`` states can be nested while ``AutoNonVariableTypeMode`` only allows enabled state.
-
-.. code-block:: cpp
-
-  {
-    InferenceMode guard(true);
-    // InferenceMode is on
-    {
-      InferenceMode guard(false);
-      // InferenceMode is off
-    }
-    // InferenceMode is on
-  }
-  // InferenceMode is off
-
-
-2. Users trying to implement a customized kernel who want to redispatch under ``Autograd`` dispatch
-   keys should use ``AutoDispatchBelowADInplaceOrView`` instead. Note ``AutoDispatchBelowADInplaceOrView`` is just a new name
-   of ``AutoNonVariableTypeMode`` since it explains the guard's functionality better. We're deprecating
-   ``AutoNonVariableTypeMode`` and it'll be removed in 1.10 release. See customized kernel
-   ``ROIAlignFunction`` in ``pytorch/vision`` for an example:
-
-.. code-block:: cpp
-
-  class ROIAlignFunction : public torch::autograd::Function<ROIAlignFunction> {
-   public:
-    static torch::autograd::variable_list forward(
-        torch::autograd::AutogradContext* ctx,
-        const torch::autograd::Variable& input,
-        const torch::autograd::Variable& rois,
-        double spatial_scale,
-        int64_t pooled_height,
-        int64_t pooled_width,
-        int64_t sampling_ratio,
-        bool aligned) {
-      ctx->saved_data["spatial_scale"] = spatial_scale;
-      ctx->saved_data["pooled_height"] = pooled_height;
-      ctx->saved_data["pooled_width"] = pooled_width;
-      ctx->saved_data["sampling_ratio"] = sampling_ratio;
-      ctx->saved_data["aligned"] = aligned;
-      ctx->saved_data["input_shape"] = input.sizes();
-      ctx->save_for_backward({rois});
-      // Used to be at::AutoNonVariableTypeMode g;
-      at::AutoDispatchBelowADInplaceOrView guard;
-      auto result = roi_align(
-          input, rois, spatial_scale, pooled_height,
-          pooled_width, sampling_ratio, aligned);
-      return {result};
-    }
-
-Customized inplace & view kernels need some special handling in addition to the guard above, see
-`custom kernel tutorial <https://pytorch.org/tutorials/advanced/cpp_extension.html#backward-pass>`_
-for more details.
diff --git a/docs/cpp/source/notes/maybe_owned.rst b/docs/cpp/source/notes/maybe_owned.rst
deleted file mode 100644
index a1bc2d02b9882..0000000000000
--- a/docs/cpp/source/notes/maybe_owned.rst
+++ /dev/null
@@ -1,59 +0,0 @@
-MaybeOwned<Tensor>
-==================
-
-``MaybeOwned<Tensor>`` is a C++ smart pointer class that dynamically
-encodes whether a Tensor is *owned* or *borrowed*. It is used in
-certain performance-sensitive situations to avoid unnecessarily
-incrementing a Tensor’s reference count (at a small cost in
-overhead from the extra indirection).
-
-.. warning::
-    MaybeOwned must be used with **extreme** care. Claims of (non-)ownership
-    are not statically checked, and mistakes can cause reference undercounting
-    and use-after-free crashes.
-
-    Due to this lack of safety net, we discourage the use of MaybeOwned
-    outside code paths that are known to be highly performance sensitive.
-    However, if you encounter pre-existing uses of MaybeOwned in code that
-    you want to modify, it’s critical to understand how to use it correctly.
-
-The primary use case for ``MaybeOwned<Tensor>`` is a function or method that
-dynamically chooses between returning one of its arguments (typically
-from a passthrough or “no-op” code path) and returning a freshly constructed
-Tensor. Such a function would return a ``MaybeOwned<Tensor>`` in both cases,
-the former in a "borrowed" state via a call to ``MaybeOwned<Tensor>::borrowed()``,
-and the latter in an "owned" state via a call to ``MaybeOwned<Tensor>::owned()``.
-
-The canonical example is ``Tensor``'s ``expect_contiguous`` method, which shortcuts
-and returns a borrowed self-reference when already contiguous:
-
-.. code-block:: cpp
-
-  inline c10::MaybeOwned<Tensor> Tensor::expect_contiguous(MemoryFormat memory_format) const & {
-    if (is_contiguous(memory_format)) {
-      return c10::MaybeOwned<Tensor>::borrowed(*this);
-    } else {
-      return c10::MaybeOwned<Tensor>::owned(__dispatch_contiguous(memory_format));
-    }
-  }
-
-Using the vocabulary of lifetimes, the essential safety requirement for borrowing
-is that a borrowed Tensor must outlive any borrowing references to it. Here, for
-example, we can safely borrow ``*this``, but the Tensor returned by
-``__dispatch_contiguous()`` is freshly created, and borrowing a reference would
-effectively leave it ownerless.
-
-So, general rules of thumb:
-
-- When in doubt, don’t use ``MaybeOwned<Tensor>`` at all - in particular, prefer
-  avoiding using it in code that doesn’t use it already. New usage should only be
-  introduced when critical (and demonstrable) performance gains result.
-
-- When modifying or calling code that already uses ``MaybeOwned<Tensor>``, remember
-  that it's always safe to produce a ``MaybeOwned<Tensor>`` from a Tensor in hand
-  via a call to ``MaybeOwned<Tensor>::owned()``. This may result in an unnecessary
-  reference count, but never in misbehavior - so it's always the safer bet, unless
-  the lifetime of the Tensor you're looking to wrap is crystal clear.
-
-More details and implementation code can be found at <https://github.com/pytorch/pytorch/blob/main/c10/util/MaybeOwned.h> and
-<https://github.com/pytorch/pytorch/blob/main/aten/src/ATen/templates/TensorBody.h>.
diff --git a/docs/cpp/source/notes/tensor_basics.rst b/docs/cpp/source/notes/tensor_basics.rst
deleted file mode 100644
index cf8f68a24eec7..0000000000000
--- a/docs/cpp/source/notes/tensor_basics.rst
+++ /dev/null
@@ -1,154 +0,0 @@
-Tensor Basics
-=============
-
-The ATen tensor library backing PyTorch is a simple tensor library that exposes
-the Tensor operations in Torch directly in C++17. ATen's API is auto-generated
-from the same declarations PyTorch uses so the two APIs will track each other
-over time.
-
-Tensor types are resolved dynamically, such that the API is generic and does not
-include templates. That is, there is one ``Tensor`` type. It can hold a CPU or
-CUDA Tensor, and the tensor may have Doubles, Float, Ints, etc. This design
-makes it easy to write generic code without templating everything.
-
-See https://pytorch.org/cppdocs/api/namespace_at.html#functions for the provided
-API. Excerpt:
-
-.. code-block:: cpp
-
-  Tensor atan2(const Tensor & other) const;
-  Tensor & atan2_(const Tensor & other);
-  Tensor pow(Scalar exponent) const;
-  Tensor pow(const Tensor & exponent) const;
-  Tensor & pow_(Scalar exponent);
-  Tensor & pow_(const Tensor & exponent);
-  Tensor lerp(const Tensor & end, Scalar weight) const;
-  Tensor & lerp_(const Tensor & end, Scalar weight);
-  Tensor histc() const;
-  Tensor histc(int64_t bins) const;
-  Tensor histc(int64_t bins, Scalar min) const;
-  Tensor histc(int64_t bins, Scalar min, Scalar max) const;
-
-In place operations are also provided, and always suffixed by `_` to indicate
-they will modify the Tensor.
-
-Efficient Access to Tensor Elements
------------------------------------
-
-When using Tensor-wide operations, the relative cost of dynamic dispatch is very
-small. However, there are cases, especially in your own kernels, where efficient
-element-wise access is needed, and the cost of dynamic dispatch inside the
-element-wise loop is very high. ATen provides *accessors* that are created with
-a single dynamic check that a Tensor is the type and number of dimensions.
-Accessors then expose an API for accessing the Tensor elements efficiently.
-
-Accessors are temporary views of a Tensor. They are only valid for the lifetime
-of the tensor that they view and hence should only be used locally in a
-function, like iterators.
-
-Note that accessors are not compatible with CUDA tensors inside kernel functions.
-Instead, you will have to use a *packed accessor* which behaves the same way but
-copies tensor metadata instead of pointing to it.
-
-It is thus recommended to use *accessors* for CPU tensors and *packed accessors*
-for CUDA tensors.
-
-CPU accessors
-*************
-
-.. code-block:: cpp
-
-  torch::Tensor foo = torch::rand({12, 12});
-
-  // assert foo is 2-dimensional and holds floats.
-  auto foo_a = foo.accessor<float,2>();
-  float trace = 0;
-
-  for(int i = 0; i < foo_a.size(0); i++) {
-    // use the accessor foo_a to get tensor data.
-    trace += foo_a[i][i];
-  }
-
-CUDA accessors
-**************
-
-
-.. code-block:: cpp
-
-  __global__ void packed_accessor_kernel(
-      torch::PackedTensorAccessor64<float, 2> foo,
-      float* trace) {
-    int i = threadIdx.x;
-    gpuAtomicAdd(trace, foo[i][i]);
-  }
-
-  torch::Tensor foo = torch::rand({12, 12});
-
-  // assert foo is 2-dimensional and holds floats.
-  auto foo_a = foo.packed_accessor64<float,2>();
-  float trace = 0;
-
-  packed_accessor_kernel<<<1, 12>>>(foo_a, &trace);
-
-In addition to ``PackedTensorAccessor64`` and ``packed_accessor64`` there are
-also the corresponding ``PackedTensorAccessor32`` and ``packed_accessor32``
-which use 32-bit integers for indexing. This can be quite a bit faster on CUDA
-but may lead to overflows in the indexing calculations.
-
-Note that the template can hold other parameters such as the pointer restriction
-and the integer type for indexing. See documentation for a thorough template
-description of *accessors* and *packed accessors*.
-
-
-Using Externally Created Data
------------------------------
-
-If you already have your tensor data allocated in memory (CPU or CUDA),
-you can view that memory as a ``Tensor`` in ATen:
-
-.. code-block:: cpp
-
-  float data[] = { 1, 2, 3,
-                   4, 5, 6 };
-  torch::Tensor f = torch::from_blob(data, {2, 3});
-
-These tensors cannot be resized because ATen does not own the memory, but
-otherwise behave as normal tensors.
-
-Scalars and zero-dimensional tensors
-------------------------------------
-
-In addition to the ``Tensor`` objects, ATen also includes ``Scalar``\s that
-represent a single number. Like a Tensor, Scalars are dynamically typed and can
-hold any one of ATen's number types. Scalars can be implicitly constructed from
-C++ number types. Scalars are needed because some functions like ``addmm`` take
-numbers along with Tensors and expect these numbers to be the same dynamic type
-as the tensor. They are also used in the API to indicate places where a function
-will *always* return a Scalar value, like ``sum``.
-
-.. code-block:: cpp
-
-  namespace torch {
-  Tensor addmm(Scalar beta, const Tensor & self,
-               Scalar alpha, const Tensor & mat1,
-               const Tensor & mat2);
-  Scalar sum(const Tensor & self);
-  } // namespace torch
-
-  // Usage.
-  torch::Tensor a = ...
-  torch::Tensor b = ...
-  torch::Tensor c = ...
-  torch::Tensor r = torch::addmm(1.0, a, .5, b, c);
-
-In addition to ``Scalar``\s, ATen also allows ``Tensor`` objects to be
-zero-dimensional. These Tensors hold a single value and they can be references
-to a single element in a larger ``Tensor``. They can be used anywhere a
-``Tensor`` is expected. They are normally created by operators like `select`
-which reduce the dimensions of a ``Tensor``.
-
-.. code-block:: cpp
-
-  torch::Tensor two = torch::rand({10, 20});
-  two[1][2] = 4;
-  // ^^^^^^ <- zero-dimensional Tensor
diff --git a/docs/cpp/source/notes/tensor_creation.rst b/docs/cpp/source/notes/tensor_creation.rst
deleted file mode 100644
index 25c9a8ca0cd46..0000000000000
--- a/docs/cpp/source/notes/tensor_creation.rst
+++ /dev/null
@@ -1,353 +0,0 @@
-Tensor Creation API
-===================
-
-This note describes how to create tensors in the PyTorch C++ API. It highlights
-the available factory functions, which populate new tensors according to some
-algorithm, and lists the options available to configure the shape, data type,
-device and other properties of a new tensor.
-
-Factory Functions
------------------
-
-A *factory function* is a function that produces a new tensor. There are many
-factory functions available in PyTorch (both in Python and C++), which differ
-in the way they initialize a new tensor before returning it. All factory
-functions adhere to the following general "schema":
-
-.. code-block:: cpp
-
-  torch::<function-name>(<function-specific-options>, <sizes>, <tensor-options>)
-
-Let's bisect the various parts of this "schema":
-
-1. ``<function-name>`` is the name of the function you would like to invoke,
-2. ``<functions-specific-options>`` are any required or optional parameters a particular factory function accepts,
-3. ``<sizes>`` is an object of type ``IntArrayRef`` and specifies the shape of the resulting tensor,
-4. ``<tensor-options>`` is an instance of ``TensorOptions`` and configures the data type, device, layout and other properties of the resulting tensor.
-
-Picking a Factory Function
-**************************
-
-The following factory functions are available at the time of this writing (the
-hyperlinks lead to the corresponding Python functions, since they often have
-more eloquent documentation -- the options are the same in C++):
-
-- `arange <https://pytorch.org/docs/stable/torch.html#torch.arange>`_: Returns a tensor with a sequence of integers,
-- `empty <https://pytorch.org/docs/stable/torch.html#torch.empty>`_: Returns a tensor with uninitialized values,
-- `eye <https://pytorch.org/docs/stable/torch.html#torch.eye>`_: Returns an identity matrix,
-- `full <https://pytorch.org/docs/stable/torch.html#torch.full>`_: Returns a tensor filled with a single value,
-- `linspace <https://pytorch.org/docs/stable/torch.html#torch.linspace>`_: Returns a tensor with values linearly spaced in some interval,
-- `logspace <https://pytorch.org/docs/stable/torch.html#torch.logspace>`_: Returns a tensor with values logarithmically spaced in some interval,
-- `ones <https://pytorch.org/docs/stable/torch.html#torch.ones>`_: Returns a tensor filled with all ones,
-- `rand <https://pytorch.org/docs/stable/torch.html#torch.rand>`_: Returns a tensor filled with values drawn from a uniform distribution on ``[0, 1)``.
-- `randint <https://pytorch.org/docs/stable/torch.html#torch.randint>`_: Returns a tensor with integers randomly drawn from an interval,
-- `randn <https://pytorch.org/docs/stable/torch.html#torch.randn>`_: Returns a tensor filled with values drawn from a unit normal distribution,
-- `randperm <https://pytorch.org/docs/stable/torch.html#torch.randperm>`_: Returns a tensor filled with a random permutation of integers in some interval,
-- `zeros <https://pytorch.org/docs/stable/torch.html#torch.zeros>`_: Returns a tensor filled with all zeros.
-
-Specifying a Size
-*****************
-
-Functions that do not require specific arguments by nature of how they fill the
-tensor can be invoked with just a size. For example, the following line creates
-a vector with 5 components, initially all set to 1:
-
-.. code-block:: cpp
-
-  torch::Tensor tensor = torch::ones(5);
-
-
-What if we wanted to instead create a ``3 x 5`` matrix, or a ``2 x 3 x 4``
-tensor? In general, an ``IntArrayRef`` -- the type of the size parameter of factory
-functions -- is constructed by specifying the size along each dimension in
-curly braces. For example, ``{2, 3}`` for a tensor (in this case matrix) with
-two rows and three columns, ``{3, 4, 5}`` for a three-dimensional tensor, and
-``{2}`` for a one-dimensional tensor with two components. In the one
-dimensional case, you can omit the curly braces and just pass the single
-integer like we did above. Note that the squiggly braces are just one way of
-constructing an ``IntArrayRef``. You can also pass an ``std::vector<int64_t>`` and
-a few other types. Either way, this means we can construct a three-dimensional
-tensor filled with values from a unit normal distribution by writing:
-
-.. code-block:: cpp
-
-  torch::Tensor tensor = torch::randn({3, 4, 5});
-  assert(tensor.sizes() == std::vector<int64_t>{3, 4, 5});
-
-``tensor.sizes()`` returns an ``IntArrayRef`` which can be compared against an
-``std::vector<int64_t>``, and we can see that it contains the sizes we passed
-to the tensor. You can also write ``tensor.size(i)`` to access a single dimension,
-which is equivalent to but preferred over ``tensor.sizes()[i]``.
-
-Passing Function-Specific Parameters
-************************************
-
-Neither ``ones`` nor ``randn`` accept any additional parameters to change their
-behavior. One function which does require further configuration is ``randint``,
-which takes an upper bound on the value for the integers it generates, as well
-as an optional lower bound, which defaults to zero. Here we create a ``5 x 5``
-square matrix with integers between 0 and 10:
-
-.. code-block:: cpp
-
-  torch::Tensor tensor = torch::randint(/*high=*/10, {5, 5});
-
-And here we raise the lower bound to 3:
-
-.. code-block:: cpp
-
-  torch::Tensor tensor = torch::randint(/*low=*/3, /*high=*/10, {5, 5});
-
-The inline comments ``/*low=*/`` and ``/*high=*/`` are not required of course,
-but aid readability just like keyword arguments in Python.
-
-.. tip::
-
-  The main take-away is that the size always follows the function specific
-  arguments.
-
-.. attention::
-
-  Sometimes a function does not need a size at all. For example, the size of
-  the tensor returned by ``arange`` is fully specified by its function-specific
-  arguments -- the lower and upper bound of a range of integers. In that case
-  the function does not take a ``size`` parameter.
-
-Configuring Properties of the Tensor
-************************************
-
-The previous section discussed function-specific arguments. Function-specific
-arguments can only change the values with which tensors are filled, and
-sometimes the size of the tensor. They never change things like the data type
-(e.g. ``float32`` or ``int64``) of the tensor being created, or whether it
-lives in CPU or GPU memory. The specification of these properties is left to
-the very last argument to every factory function: a ``TensorOptions`` object,
-discussed below.
-
-``TensorOptions`` is a class that encapsulates the construction axes of a
-Tensor. With *construction axis* we mean a particular property of a Tensor that
-can be configured before its construction (and sometimes changed afterwards).
-These construction axes are:
-
-- The ``dtype`` (previously "scalar type"), which controls the data type of the
-  elements stored in the tensor,
-- The ``layout``, which is either strided (dense) or sparse,
-- The ``device``, which represents a compute device on which a tensor is stored (like a CPU or CUDA GPU),
-- The ``requires_grad`` boolean to enable or disable gradient recording for a tensor,
-
-If you are used to PyTorch in Python, these axes will sound very familiar. The
-allowed values for these axes at the moment are:
-
-- For ``dtype``: ``kUInt8``, ``kInt8``, ``kInt16``, ``kInt32``, ``kInt64``, ``kFloat32`` and ``kFloat64``,
-- For ``layout``: ``kStrided`` and ``kSparse``,
-- For ``device``: Either ``kCPU``, or ``kCUDA`` (which accepts an optional device index),
-- For ``requires_grad``: either ``true`` or ``false``.
-
-.. tip::
-
-        There exist "Rust-style" shorthands for dtypes, like ``kF32`` instead of
-        ``kFloat32``. See `here
-        <https://github.com/pytorch/pytorch/blob/main/torch/csrc/api/include/torch/types.h>`_
-        for the full list.
-
-
-An instance of ``TensorOptions`` stores a concrete value for each of these
-axes. Here is an example of creating a ``TensorOptions`` object that represents
-a 64-bit float, strided tensor that requires a gradient, and lives on CUDA
-device 1:
-
-.. code-block:: cpp
-
-  auto options =
-    torch::TensorOptions()
-      .dtype(torch::kFloat32)
-      .layout(torch::kStrided)
-      .device(torch::kCUDA, 1)
-      .requires_grad(true);
-
-
-Notice how we use the '"builder"-style methods of ``TensorOptions`` to
-construct the object piece by piece. If we pass this object as the last
-argument to a factory function, the newly created tensor will have these
-properties:
-
-.. code-block:: cpp
-
-  torch::Tensor tensor = torch::full({3, 4}, /*value=*/123, options);
-
-  assert(tensor.dtype() == torch::kFloat32);
-  assert(tensor.layout() == torch::kStrided);
-  assert(tensor.device().type() == torch::kCUDA); // or device().is_cuda()
-  assert(tensor.device().index() == 1);
-  assert(tensor.requires_grad());
-
-Now, you may be thinking: do I really need to specify each axis for every new
-tensor I create? Fortunately, the answer is "no", as **every axis has a default
-value**. These defaults are:
-
-- ``kFloat32`` for the dtype,
-- ``kStrided`` for the layout,
-- ``kCPU`` for the device,
-- ``false`` for ``requires_grad``.
-
-What this means is that any axis you omit during the construction of a
-``TensorOptions`` object will take on its default value. For example, this is
-our previous ``TensorOptions`` object, but with the ``dtype`` and ``layout``
-defaulted:
-
-.. code-block:: cpp
-
-  auto options = torch::TensorOptions().device(torch::kCUDA, 1).requires_grad(true);
-
-In fact, we can even omit all axes to get an entirely defaulted
-``TensorOptions`` object:
-
-.. code-block:: cpp
-
-  auto options = torch::TensorOptions(); // or `torch::TensorOptions options;`
-
-A nice consequence of this is that the ``TensorOptions`` object we just spoke
-so much about can be entirely omitted from any tensor factory call:
-
-.. code-block:: cpp
-
-  // A 32-bit float, strided, CPU tensor that does not require a gradient.
-  torch::Tensor tensor = torch::randn({3, 4});
-  torch::Tensor range = torch::arange(5, 10);
-
-But the sugar gets sweeter: In the API presented here so far, you may have
-noticed that the initial ``torch::TensorOptions()`` is quite a mouthful to
-write. The good news is that for every construction axis (dtype, layout, device
-and ``requires_grad``), there is one *free function* in the ``torch::``
-namespace which you can pass a value for that axis. Each function then returns
-a ``TensorOptions`` object preconfigured with that axis, but allowing even
-further modification via the builder-style methods shown above. For example,
-
-.. code-block:: cpp
-
-  torch::ones(10, torch::TensorOptions().dtype(torch::kFloat32))
-
-is equivalent to
-
-.. code-block:: cpp
-
-  torch::ones(10, torch::dtype(torch::kFloat32))
-
-and further instead of
-
-.. code-block:: cpp
-
-  torch::ones(10, torch::TensorOptions().dtype(torch::kFloat32).layout(torch::kStrided))
-
-we can just write
-
-.. code-block:: cpp
-
-  torch::ones(10, torch::dtype(torch::kFloat32).layout(torch::kStrided))
-
-which saves us quite a bit of typing. What this means is that in practice, you
-should barely, if ever, have to write out ``torch::TensorOptions``. Instead use
-the ``torch::dtype()``, ``torch::device()``, ``torch::layout()`` and
-``torch::requires_grad()`` functions.
-
-A final bit of convenience is that ``TensorOptions`` is implicitly
-constructible from individual values. This means that whenever a function has a
-parameter of type ``TensorOptions``, like all factory functions do, we can
-directly pass a value like ``torch::kFloat32`` or ``torch::kStrided`` in place
-of the full object. Therefore, when there is only a single axis we would like
-to change compared to its default value, we can pass only that value. As such,
-what was
-
-.. code-block:: cpp
-
-  torch::ones(10, torch::TensorOptions().dtype(torch::kFloat32))
-
-became
-
-.. code-block:: cpp
-
-  torch::ones(10, torch::dtype(torch::kFloat32))
-
-and can finally be shortened to
-
-.. code-block:: cpp
-
-  torch::ones(10, torch::kFloat32)
-
-Of course, it is not possible to modify further properties of the
-``TensorOptions`` instance with this short syntax, but if all we needed was to
-change one property, this is quite practical.
-
-In conclusion, we can now compare how ``TensorOptions`` defaults, together with
-the abbreviated API for creating ``TensorOptions`` using free functions, allow
-tensor creation in C++ with the same convenience as in Python. Compare this
-call in Python:
-
-.. code-block:: python
-
-  torch.randn(3, 4, dtype=torch.float32, device=torch.device('cuda', 1), requires_grad=True)
-
-with the equivalent call in C++:
-
-.. code-block:: cpp
-
-  torch::randn({3, 4}, torch::dtype(torch::kFloat32).device(torch::kCUDA, 1).requires_grad(true))
-
-Pretty close!
-
-Conversion
-----------
-
-Just as we can use ``TensorOptions`` to configure how new tensors should be
-created, we can also use ``TensorOptions`` to convert a tensor from one set of
-properties to a new set of properties. Such a conversion usually creates a new
-tensor and does not occur in-place. For example, if we have a ``source_tensor``
-created with
-
-.. code-block:: cpp
-
-  torch::Tensor source_tensor = torch::randn({2, 3}, torch::kInt64);
-
-we can convert it from ``int64`` to ``float32``:
-
-.. code-block:: cpp
-
-  torch::Tensor float_tensor = source_tensor.to(torch::kFloat32);
-
-.. attention::
-
-        The result of the conversion, ``float_tensor``, is a new tensor pointing to
-        new memory, unrelated to the source ``source_tensor``.
-
-We can then move it from CPU memory to GPU memory:
-
-.. code-block:: cpp
-
-  torch::Tensor gpu_tensor = float_tensor.to(torch::kCUDA);
-
-If you have multiple CUDA devices available, the above code will copy the
-tensor to the *default* CUDA device, which you can configure with a
-``torch::DeviceGuard``. If no ``DeviceGuard`` is in place, this will be GPU
-1. If you would like to specify a different GPU index, you can pass it to
-the ``Device`` constructor:
-
-.. code-block:: cpp
-
-  torch::Tensor gpu_two_tensor = float_tensor.to(torch::Device(torch::kCUDA, 1));
-
-In the case of CPU to GPU copy and reverse, we can also configure the memory
-copy to be *asynchronous* by passing ``/*non_blocking=*/false`` as the last
-argument to ``to()``:
-
-.. code-block:: cpp
-
-  torch::Tensor async_cpu_tensor = gpu_tensor.to(torch::kCPU, /*non_blocking=*/true);
-
-Conclusion
-----------
-
-This note hopefully gave you a good understanding of how to create and convert
-tensors in an idiomatic fashion using the PyTorch C++ API. If you have any
-further questions or suggestions, please use our `forum
-<https://discuss.pytorch.org/>`_ or `GitHub issues
-<https://github.com/pytorch/pytorch/issues>`_ to get in touch.
diff --git a/docs/cpp/source/notes/tensor_cuda_stream.rst b/docs/cpp/source/notes/tensor_cuda_stream.rst
deleted file mode 100644
index 4940317713635..0000000000000
--- a/docs/cpp/source/notes/tensor_cuda_stream.rst
+++ /dev/null
@@ -1,276 +0,0 @@
-Tensor CUDA Stream API
-======================
-
-A `CUDA Stream`_ is a linear sequence of execution that belongs to a specific CUDA device.
-The PyTorch C++ API supports CUDA streams with the CUDAStream class and useful helper functions to make streaming operations easy.
-You can find them in `CUDAStream.h`_. This note provides more details on how to use Pytorch C++ CUDA Stream APIs.
-
-.. _CUDA Stream: https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#streams
-.. _CUDAStream.h: https://pytorch.org/cppdocs/api/file_c10_cuda_CUDAStream.h.html#file-c10-cuda-cudastream-h
-.. _CUDAStreamGuard.h: https://pytorch.org/cppdocs/api/structc10_1_1cuda_1_1_c_u_d_a_stream_guard.html
-
-Acquiring CUDA stream
-*********************
-
-Pytorch's C++ API provides the following ways to acquire CUDA stream:
-
-1. Acquire a new stream from the CUDA stream pool, streams are preallocated from the pool and returned in a round-robin fashion.
-
-.. code-block:: cpp
-
-  CUDAStream getStreamFromPool(const bool isHighPriority = false, DeviceIndex device = -1);
-
-.. tip::
-
-  You can request a stream from the high priority pool by setting isHighPriority to true, or a stream for a specific device
-  by setting device index (defaulting to the current CUDA stream's device index).
-
-2. Acquire the default CUDA stream for the passed CUDA device, or for the current device if no device index is passed.
-
-.. code-block:: cpp
-
-  CUDAStream getDefaultCUDAStream(DeviceIndex device_index = -1);
-
-.. tip::
-
-  The default stream is where most computation occurs when you aren't explicitly using streams.
-
-3. Acquire the current CUDA stream, for the CUDA device with index ``device_index``, or for the current device if no device index is passed.
-
-.. code-block:: cpp
-
-  CUDAStream getCurrentCUDAStream(DeviceIndex device_index = -1);
-
-.. tip::
-
-  The current CUDA stream will usually be the default CUDA stream for the device, but it may be different if someone
-  called ``setCurrentCUDAStream`` or used ``StreamGuard`` or ``CUDAStreamGuard``.
-
-
-
-Set CUDA stream
-***************
-
-Pytorch's C++ API provides the following ways to set CUDA stream:
-
-1. Set the current stream on the device of the passed in stream to be the passed in stream.
-
-.. code-block:: cpp
-
-  void setCurrentCUDAStream(CUDAStream stream);
-
-.. attention::
-
-  This function may have nothing to do with the current device. It only changes the current stream on the stream's device.
-  We recommend using ``CUDAStreamGuard``, instead, since it switches to the stream's device and makes it the current stream on that device.
-  ``CUDAStreamGuard`` will also restore the current device and stream when it's destroyed
-
-2. Use ``CUDAStreamGuard`` to switch to a CUDA stream within a scope, it is defined in `CUDAStreamGuard.h`_
-
-.. tip::
-
-  Use ``CUDAMultiStreamGuard`` if you need to set streams on multiple CUDA devices.
-
-CUDA Stream Usage Examples
-**************************
-
-1. Acquiring and setting CUDA stream on the same device
-
-.. code-block:: cpp
-
-  // This example shows how to acquire and set CUDA stream on the same device.
-  // `at::cuda::setCurrentCUDAStream` is used to set current CUDA stream
-
-  // create a tensor on device 0
-  torch::Tensor tensor0 = torch::ones({2, 2}, torch::device(torch::kCUDA));
-  // get a new CUDA stream from CUDA stream pool on device 0
-  at::cuda::CUDAStream myStream = at::cuda::getStreamFromPool();
-  // set current CUDA stream from default stream to `myStream` on device 0
-  at::cuda::setCurrentCUDAStream(myStream);
-  // sum() on tensor0 uses `myStream` as current CUDA stream
-  tensor0.sum();
-
-  // get the default CUDA stream on device 0
-  at::cuda::CUDAStream defaultStream = at::cuda::getDefaultCUDAStream();
-  // set current CUDA stream back to default CUDA stream on device 0
-  at::cuda::setCurrentCUDAStream(defaultStream);
-  // sum() on tensor0 uses `defaultStream` as current CUDA stream
-  tensor0.sum();
-
-.. code-block:: cpp
-
-  // This example is the same as previous example, but explicitly specify device
-  // index and use CUDA stream guard to set current CUDA stream
-
-  // create a tensor on device 0
-  torch::Tensor tensor0 = torch::ones({2, 2}, torch::device(torch::kCUDA));
-  // get a new stream from CUDA stream pool on device 0
-  at::cuda::CUDAStream myStream = at::cuda::getStreamFromPool(false, 0);
-  // set the current CUDA stream to `myStream` within the scope using CUDA stream guard
-  {
-    at::cuda::CUDAStreamGuard guard(myStream);
-    // current CUDA stream is `myStream` from here till the end of bracket.
-    // sum() on tensor0 uses `myStream` as current CUDA stream
-    tensor0.sum();
-  }
-  // current CUDA stream is reset to default CUDA stream after CUDA stream guard is destroyed
-  // sum() on tensor0 uses default CUDA stream on device 0 as current CUDA stream
-  tensor0.sum();
-
-.. attention::
-
-  Above code is running on the same CUDA device. `setCurrentCUDAStream` will always set current CUDA stream on current device,
-  but note that `setCurrentCUDAStream` actually set current stream on the device of passed in CUDA stream.
-
-
-2. Acquiring and setting CUDA streams on multiple devices.
-
-.. code-block:: cpp
-
-  // This example shows how to acquire and set CUDA stream on two devices.
-
-  // acquire new CUDA streams from CUDA stream pool on device 0 and device 1
-  at::cuda::CUDAStream myStream0 = at::cuda::getStreamFromPool(false, 0);
-  at::cuda::CUDAStream myStream1 = at::cuda::getStreamFromPool(false, 1);
-
-  // set current CUDA stream to `myStream0` on device 0
-  at::cuda::setCurrentCUDAStream(myStream0);
-  // set current CUDA stream to `myStream1` on device 1
-  at::cuda::setCurrentCUDAStream(myStream1);
-
-  // create a tensor on device 0, no need to specify device index since
-  // current device index is 0
-  torch::Tensor tensor0 = torch::ones({2, 2}, torch::device(at::kCUDA));
-  // sum() on tensor0 use `myStream0` as current CUDA stream on device 0
-  tensor0.sum();
-
-  // change the current device index to 1 by using CUDA device guard within a bracket scope
-  {
-    at::cuda::CUDAGuard device_guard{1};
-    // create a tensor on device 1
-    torch::Tensor tensor1 = torch::ones({2, 2}, torch::device(at::kCUDA));
-    // sum() on tensor 1 uses `myStream1` as current CUDA stream on device 1
-    tensor1.sum();
-  }
-
-  // current device is reset to device 0 after device_guard is destroyed
-
-  // acquire a new CUDA stream on device 1
-  at::cuda::CUDAStream myStream1_1 = at::cuda::getStreamFromPool(false, 1);
-  // create a new tensor on device 1
-  torch::Tensor tensor1 = torch::ones({2, 2}, torch::device({torch::kCUDA, 1}));
-
-  // change the current device index to 1 and current CUDA stream on device 1
-  // to `myStream1_1` using CUDA stream guard within a scope
-  {
-    at::cuda::CUDAStreamGuard stream_guard(myStream1_1);
-    // sum() on tensor1 use `myStream1_1` as current CUDA stream on device 1
-    tensor1.sum();
-  }
-
-  // current device is reset to device 0 and current CUDA stream on device 1 is
-  // reset to `myStream1`
-
-  // sum() on tensor1 uses `myStream1` as current CUDA stream on device 1
-  tensor1.sum();
-
-
-3. Working with CUDA multistream guard
-
-.. code-block:: cpp
-
-  // This example shows how to use CUDA multistream guard to set
-  // two streams on two devices at the same time.
-
-  // create two tensor, one on device 0, one on device 1
-  torch::Tensor tensor0 = torch::ones({2, 2}, torch::device({torch::kCUDA, 0}));
-  torch::Tensor tensor1 = torch::ones({2, 2}, torch::device({torch::kCUDA, 1}));
-
-  // acquire new CUDA streams from CUDA stream pool on device 0 and device 1
-  at::cuda::CUDAStream myStream0 = at::cuda::getStreamFromPool(false, 0);
-  at::cuda::CUDAStream myStream1 = at::cuda::getStreamFromPool(false, 1);
-
-  // set current CUDA stream on device 0 to `myStream0` and
-  // set current CUDA stream on device 1 to `myStream1` CUDA using multistream guard
-  {
-    at::cuda::CUDAMultiStreamGuard multi_guard({myStream0, myStream1});
-
-    // sum() on tensor0 uses `myStream0` as current CUDA stream on device 0
-    tensor0.sum();
-    // sum() on tensor1 uses `myStream1` as current CUDA stream on device 1
-    tensor1.sum();
-  }
-
-  // current CUDA stream on device 0 is reset to default CUDA stream on device 0
-  // current CUDA stream on device 1 is reset to default CUDA stream on device 1
-
-  // sum() on tensor0 uses default CUDA stream as current CUDA stream on device 0
-  tensor0.sum();
-  // sum() on tensor1 uses default CUDA stream as current CUDA stream on device 1
-  tensor1.sum();
-
-.. attention::
-  ``CUDAMultiStreamGuard`` does not change current device index, it only changes the stream on
-  each passed in stream's device. Other than scope controlling, this guard is equivalent to
-  calling ``setCurrentCUDAStream`` on each passed in stream.
-
-4. A skeleton example for handling CUDA streams on multiple devices
-
-.. code-block:: cpp
-
-   // This is a skeleton example that shows how to handle CUDA streams on multiple devices
-   // Suppose you want to do work on the non-default stream on two devices simultaneously, and we
-   // already have streams on both devices in two vectors. The following code shows three ways
-   // of acquiring and setting the streams.
-
-   // Usage 0: acquire CUDA stream and set current CUDA stream with `setCurrentCUDAStream`
-   // Create a CUDA stream vector `streams0` on device 0
-   std::vector<at::cuda::CUDAStream> streams0 =
-     {at::cuda::getDefaultCUDAStream(), at::cuda::getStreamFromPool()};
-   // set current stream as `streams0[0]` on device 0
-   at::cuda::setCurrentCUDAStream(streams0[0]);
-
-   // create a CUDA stream vector `streams1` on device using CUDA device guard
-   std::vector<at::cuda::CUDAStream> streams1;
-   {
-     // device index is set to 1 within this scope
-     at::cuda::CUDAGuard device_guard(1);
-     streams1.push_back(at::cuda::getDefaultCUDAStream());
-     streams1.push_back(at::cuda::getStreamFromPool());
-   }
-   // device index is reset to 0 after device_guard is destroyed
-
-   // set current stream as `streams1[0]` on device 1
-   at::cuda::setCurrentCUDAStream(streams1[0]);
-
-
-   // Usage 1: use CUDA device guard to change the current device index only
-   {
-     at::cuda::CUDAGuard device_guard(1);
-
-     // current device index is changed to 1 within scope
-     // current CUDA stream is still `streams1[0]` on device 1, no change
-   }
-   // current device index is reset to 0 after `device_guard` is destroyed
-
-
-   // Usage 2: use CUDA stream guard to change both current device index and current CUDA stream.
-   {
-     at::cuda::CUDAStreamGuard stream_guard(streams1[1]);
-
-     // current device index and current CUDA stream are set to 1 and `streams1[1]` within scope
-   }
-   // current device index and current CUDA stream are reset to 0 and `streams0[0]` after
-   // stream_guard is destroyed
-
-
-   // Usage 3: use CUDA multi-stream guard to change multiple streams on multiple devices
-   {
-     // This is the same as calling `torch::cuda::setCurrentCUDAStream` on both streams
-     at::cuda::CUDAMultiStreamGuard multi_guard({streams0[1], streams1[1]});
-
-     // current device index is not change, still 0
-     // current CUDA stream on device 0 and device 1 are set to `streams0[1]` and `streams1[1]`
-   }
-   // current CUDA stream on device 0 and device 1 are reset to `streams0[0]` and `streams1[0]`
-   // after `multi_guard` is destroyed.
diff --git a/docs/cpp/source/notes/tensor_indexing.rst b/docs/cpp/source/notes/tensor_indexing.rst
deleted file mode 100644
index 013ab7a737c75..0000000000000
--- a/docs/cpp/source/notes/tensor_indexing.rst
+++ /dev/null
@@ -1,99 +0,0 @@
-Tensor Indexing API
-===================
-
-Indexing a tensor in the PyTorch C++ API works very similar to the Python API.
-All index types such as ``None`` / ``...`` / integer / boolean / slice / tensor
-are available in the C++ API, making translation from Python indexing code to C++
-very simple. The main difference is that, instead of using the ``[]``-operator
-similar to the Python API syntax, in the C++ API the indexing methods are:
-
-- ``torch::Tensor::index`` (`link <https://pytorch.org/cppdocs/api/classat_1_1_tensor.html#_CPPv4NK2at6Tensor5indexE8ArrayRefIN2at8indexing11TensorIndexEE>`_)
-- ``torch::Tensor::index_put_`` (`link <https://pytorch.org/cppdocs/api/classat_1_1_tensor.html#_CPPv4N2at6Tensor10index_put_E8ArrayRefIN2at8indexing11TensorIndexEERK6Tensor>`_)
-
-It's also important to note that index types such as ``None`` / ``Ellipsis`` / ``Slice``
-live in the ``torch::indexing`` namespace, and it's recommended to put ``using namespace torch::indexing``
-before any indexing code for convenient use of those index types.
-
-Here are some examples of translating Python indexing code to C++:
-
-Getter
-------
-
-+----------------------------------------------------------+--------------------------------------------------------------------------------------+
-| Python                                                   | C++  (assuming ``using namespace torch::indexing``)                                  |
-+==========================================================+======================================================================================+
-| ``tensor[None]``                                         | ``tensor.index({None})``                                                             |
-+----------------------------------------------------------+--------------------------------------------------------------------------------------+
-| ``tensor[Ellipsis, ...]``                                | ``tensor.index({Ellipsis, "..."})``                                                  |
-+----------------------------------------------------------+--------------------------------------------------------------------------------------+
-| ``tensor[1, 2]``                                         | ``tensor.index({1, 2})``                                                             |
-+----------------------------------------------------------+--------------------------------------------------------------------------------------+
-| ``tensor[True, False]``                                  | ``tensor.index({true, false})``                                                      |
-+----------------------------------------------------------+--------------------------------------------------------------------------------------+
-| ``tensor[1::2]``                                         | ``tensor.index({Slice(1, None, 2)})``                                                |
-+----------------------------------------------------------+--------------------------------------------------------------------------------------+
-| ``tensor[torch.tensor([1, 2])]``                         | ``tensor.index({torch::tensor({1, 2})})``                                            |
-+----------------------------------------------------------+--------------------------------------------------------------------------------------+
-| ``tensor[..., 0, True, 1::2, torch.tensor([1, 2])]``     | ``tensor.index({"...", 0, true, Slice(1, None, 2), torch::tensor({1, 2})})``         |
-+----------------------------------------------------------+--------------------------------------------------------------------------------------+
-
-Setter
-------
-
-+----------------------------------------------------------+--------------------------------------------------------------------------------------+
-| Python                                                   | C++  (assuming ``using namespace torch::indexing``)                                  |
-+==========================================================+======================================================================================+
-| ``tensor[None] = 1``                                     | ``tensor.index_put_({None}, 1)``                                                     |
-+----------------------------------------------------------+--------------------------------------------------------------------------------------+
-| ``tensor[Ellipsis, ...] = 1``                            | ``tensor.index_put_({Ellipsis, "..."}, 1)``                                          |
-+----------------------------------------------------------+--------------------------------------------------------------------------------------+
-| ``tensor[1, 2] = 1``                                     | ``tensor.index_put_({1, 2}, 1)``                                                     |
-+----------------------------------------------------------+--------------------------------------------------------------------------------------+
-| ``tensor[True, False] = 1``                              | ``tensor.index_put_({true, false}, 1)``                                              |
-+----------------------------------------------------------+--------------------------------------------------------------------------------------+
-| ``tensor[1::2] = 1``                                     | ``tensor.index_put_({Slice(1, None, 2)}, 1)``                                        |
-+----------------------------------------------------------+--------------------------------------------------------------------------------------+
-| ``tensor[torch.tensor([1, 2])] = 1``                     | ``tensor.index_put_({torch::tensor({1, 2})}, 1)``                                    |
-+----------------------------------------------------------+--------------------------------------------------------------------------------------+
-| ``tensor[..., 0, True, 1::2, torch.tensor([1, 2])] = 1`` | ``tensor.index_put_({"...", 0, true, Slice(1, None, 2), torch::tensor({1, 2})}, 1)`` |
-+----------------------------------------------------------+--------------------------------------------------------------------------------------+
-
-
-Translating between Python/C++ index types
-------------------------------------------
-
-The one-to-one translation between Python and C++ index types is as follows:
-
-+-------------------------+------------------------------------------------------------------------+
-| Python                  | C++ (assuming ``using namespace torch::indexing``)                     |
-+=========================+========================================================================+
-| ``None``                | ``None``                                                               |
-+-------------------------+------------------------------------------------------------------------+
-| ``Ellipsis``            | ``Ellipsis``                                                           |
-+-------------------------+------------------------------------------------------------------------+
-| ``...``                 | ``"..."``                                                              |
-+-------------------------+------------------------------------------------------------------------+
-| ``123``                 | ``123``                                                                |
-+-------------------------+------------------------------------------------------------------------+
-| ``True``                | ``true``                                                               |
-+-------------------------+------------------------------------------------------------------------+
-| ``False``               | ``false``                                                              |
-+-------------------------+------------------------------------------------------------------------+
-| ``:`` or ``::``         | ``Slice()`` or ``Slice(None, None)`` or ``Slice(None, None, None)``    |
-+-------------------------+------------------------------------------------------------------------+
-| ``1:`` or ``1::``       | ``Slice(1, None)`` or ``Slice(1, None, None)``                         |
-+-------------------------+------------------------------------------------------------------------+
-| ``:3`` or ``:3:``       | ``Slice(None, 3)`` or ``Slice(None, 3, None)``                         |
-+-------------------------+------------------------------------------------------------------------+
-| ``::2``                 | ``Slice(None, None, 2)``                                               |
-+-------------------------+------------------------------------------------------------------------+
-| ``1:3``                 | ``Slice(1, 3)``                                                        |
-+-------------------------+------------------------------------------------------------------------+
-| ``1::2``                | ``Slice(1, None, 2)``                                                  |
-+-------------------------+------------------------------------------------------------------------+
-| ``:3:2``                | ``Slice(None, 3, 2)``                                                  |
-+-------------------------+------------------------------------------------------------------------+
-| ``1:3:2``               | ``Slice(1, 3, 2)``                                                     |
-+-------------------------+------------------------------------------------------------------------+
-| ``torch.tensor([1, 2])``| ``torch::tensor({1, 2})``                                              |
-+-------------------------+------------------------------------------------------------------------+
diff --git a/docs/cpp/source/notes/versioning.rst b/docs/cpp/source/notes/versioning.rst
deleted file mode 100644
index b7c79427bb031..0000000000000
--- a/docs/cpp/source/notes/versioning.rst
+++ /dev/null
@@ -1,29 +0,0 @@
-Library Versioning
-==================
-
-We provide version number macros for identifying the version of LibTorch in use.
-Example usage:
-
-.. code-block:: cpp
-
-  #include <torch/torch.h>
-  #include <iostream>
-
-  int main() {
-    std::cout << "PyTorch version from parts: "
-      << TORCH_VERSION_MAJOR << "."
-      << TORCH_VERSION_MINOR << "."
-      << TORCH_VERSION_PATCH << std::endl;
-    std::cout << "PyTorch version: " << TORCH_VERSION << std::endl;
-  }
-
-This will output something like:
-
-.. code-block:: text
-
-  PyTorch version from parts: 1.8.0
-  PyTorch version: 1.8.0
-
-.. note::
-
-  These macros are only available in PyTorch >= 1.8.0.
diff --git a/docs/cpp/source/stable.rst b/docs/cpp/source/stable.rst
deleted file mode 100644
index c7c60995419da..0000000000000
--- a/docs/cpp/source/stable.rst
+++ /dev/null
@@ -1,388 +0,0 @@
-Torch Stable API
-================
-
-The PyTorch Stable C++ API provides a convenient high level interface to call
-ABI-stable tensor operations and other utilities commonly used in custom operators.
-These functions are designed to maintain binary compatibility across PyTorch versions,
-making them suitable for use in ahead-of-time compiled code.
-
-For more information on the stable ABI, see the
-`Stable ABI notes <https://docs.pytorch.org/docs/stable/notes/libtorch_stable_abi.html>`_.
-
-Library Registration Macros
----------------------------
-
-These macros provide stable ABI equivalents of the standard PyTorch operator
-registration macros (``TORCH_LIBRARY``, ``TORCH_LIBRARY_IMPL``, etc.).
-Use these when building custom operators that need to maintain binary
-compatibility across PyTorch versions.
-
-``STABLE_TORCH_LIBRARY(ns, m)``
-^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
-
-Defines a library of operators in a namespace using the stable ABI.
-
-This is the stable ABI equivalent of :c:macro:`TORCH_LIBRARY`.
-Use this macro to define operator schemas that will maintain
-binary compatibility across PyTorch versions. Only one ``STABLE_TORCH_LIBRARY``
-block can exist per namespace; use ``STABLE_TORCH_LIBRARY_FRAGMENT`` for
-additional definitions in the same namespace from different translation units.
-
-**Parameters:**
-
-- ``ns`` - The namespace in which to define operators (e.g., ``mylib``).
-- ``m`` - The name of the StableLibrary variable available in the block.
-
-**Example:**
-
-.. code-block:: cpp
-
-   STABLE_TORCH_LIBRARY(mylib, m) {
-       m.def("my_op(Tensor input, int size) -> Tensor");
-       m.def("another_op(Tensor a, Tensor b) -> Tensor");
-   }
-
-Minimum compatible version: PyTorch 2.9.
-
-``STABLE_TORCH_LIBRARY_IMPL(ns, k, m)``
-^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
-
-Registers operator implementations for a specific dispatch key using the stable ABI.
-
-This is the stable ABI equivalent of ``TORCH_LIBRARY_IMPL``. Use this macro
-to provide implementations of operators for a specific dispatch key (e.g.,
-CPU, CUDA) while maintaining binary compatibility across PyTorch versions.
-
-.. note::
-
-   All kernel functions registered with this macro must be boxed using
-   the ``TORCH_BOX`` macro.
-
-**Parameters:**
-
-- ``ns`` - The namespace in which the operators are defined.
-- ``k`` - The dispatch key (e.g., ``CPU``, ``CUDA``).
-- ``m`` - The name of the StableLibrary variable available in the block.
-
-**Example:**
-
-.. code-block:: cpp
-
-   STABLE_TORCH_LIBRARY_IMPL(mylib, CPU, m) {
-       m.impl("my_op", TORCH_BOX(&my_cpu_kernel));
-   }
-
-   STABLE_TORCH_LIBRARY_IMPL(mylib, CUDA, m) {
-       m.impl("my_op", TORCH_BOX(&my_cuda_kernel));
-   }
-
-Minimum compatible version: PyTorch 2.9.
-
-``STABLE_TORCH_LIBRARY_FRAGMENT(ns, m)``
-^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
-
-Extends operator definitions in an existing namespace using the stable ABI.
-
-This is the stable ABI equivalent of ``TORCH_LIBRARY_FRAGMENT``. Use this macro
-to add additional operator definitions to a namespace that was already
-created with ``STABLE_TORCH_LIBRARY``.
-
-**Parameters:**
-
-- ``ns`` - The namespace to extend.
-- ``m`` - The name of the StableLibrary variable available in the block.
-
-Minimum compatible version: PyTorch 2.9.
-
-``TORCH_BOX(&func)``
-^^^^^^^^^^^^^^^^^^^
-
-Wraps a function to conform to the stable boxed kernel calling convention.
-
-This macro takes an unboxed kernel function pointer and generates a boxed wrapper
-that can be registered with the stable library API.
-
-**Parameters:**
-
-- ``func`` - The unboxed kernel function to wrap.
-
-**Example:**
-
-.. code-block:: cpp
-
-   Tensor my_kernel(const Tensor& input, int64_t size) {
-       return input.reshape({size});
-   }
-
-   STABLE_TORCH_LIBRARY_IMPL(my_namespace, CPU, m) {
-       m.impl("my_op", TORCH_BOX(&my_kernel));
-   }
-
-Minimum compatible version: PyTorch 2.9.
-
-Tensor Class
-------------
-
-The ``torch::stable::Tensor`` class offers a user-friendly C++ interface similar
-to ``torch::Tensor`` while maintaining binary compatibility across PyTorch versions.
-
-.. doxygenclass:: torch::stable::Tensor
-   :members:
-
-
-Device Class
-------------
-
-The ``torch::stable::Device`` class provides a user-friendly C++ interface similar
-to ``c10::Device`` while maintaining binary compatibility across PyTorch versions.
-It represents a compute device (CPU, CUDA, etc.) with an optional device index.
-
-.. doxygenclass:: torch::stable::Device
-   :members:
-
-DeviceGuard Class
------------------
-
-The ``torch::stable::accelerator::DeviceGuard`` provides a user-friendly C++
-interface similar to ``c10::DeviceGuard`` while maintaining binary compatibility
-across PyTorch versions.
-
-.. doxygenclass:: torch::stable::accelerator::DeviceGuard
-   :members:
-
-.. doxygenfunction:: torch::stable::accelerator::getCurrentDeviceIndex
-
-
-Stream Utilities
-----------------
-
-For CUDA stream access, we currently recommend the ABI stable C shim API. This
-will be improved in a future release with a more ergonomic wrapper.
-
-Getting the Current CUDA Stream
-^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
-
-To obtain the current ``cudaStream_t`` for use in CUDA kernels:
-
-.. code-block:: cpp
-
-   #include <torch/csrc/inductor/aoti_torch/c/shim.h>
-   #include <torch/headeronly/util/shim_utils.h>
-
-   // For now, we rely on the ABI stable C shim API to get the current CUDA stream.
-   // This will be improved in a future release.
-   // When using a C shim API, we need to use TORCH_ERROR_CODE_CHECK to
-   // check the error code and throw an appropriate runtime_error otherwise.
-   void* stream_ptr = nullptr;
-   TORCH_ERROR_CODE_CHECK(
-       aoti_torch_get_current_cuda_stream(tensor.get_device_index(), &stream_ptr));
-   cudaStream_t stream = static_cast<cudaStream_t>(stream_ptr);
-
-   // Now you can use 'stream' in your CUDA kernel launches
-   my_kernel<<<blocks, threads, 0, stream>>>(args...);
-
-.. note::
-
-   The ``TORCH_ERROR_CODE_CHECK`` macro is required when using C shim APIs
-   to properly check error codes and throw appropriate exceptions.
-
-CUDA Error Checking Macros
---------------------------
-
-These macros provide stable ABI equivalents for CUDA error checking.
-They wrap CUDA API calls and kernel launches, providing detailed error
-messages using PyTorch's error formatting.
-
-``STD_CUDA_CHECK(EXPR)``
-^^^^^^^^^^^^^^^^^^^^^^^^
-
-Checks the result of a CUDA API call and throws an exception on error.
-Users of this macro are expected to include ``cuda_runtime.h``.
-
-**Example:**
-
-.. code-block:: cpp
-
-   STD_CUDA_CHECK(cudaMalloc(&ptr, size));
-   STD_CUDA_CHECK(cudaMemcpy(dst, src, size, cudaMemcpyDeviceToHost));
-
-Minimum compatible version: PyTorch 2.10.
-
-``STD_CUDA_KERNEL_LAUNCH_CHECK()``
-^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
-
-Checks for errors from the most recent CUDA kernel launch. Equivalent to
-``STD_CUDA_CHECK(cudaGetLastError())``.
-
-**Example:**
-
-.. code-block:: cpp
-
-   my_kernel<<<blocks, threads, 0, stream>>>(args...);
-   STD_CUDA_KERNEL_LAUNCH_CHECK();
-
-Minimum compatible version: PyTorch 2.10.
-
-Header-Only Utilities
----------------------
-
-The ``torch::headeronly`` namespace provides header-only versions of common
-PyTorch types and utilities. These can be used without linking against libtorch,
-making them ideal for maintaining binary compatibility across PyTorch versions.
-
-Error Checking
-^^^^^^^^^^^^^^
-
-``STD_TORCH_CHECK`` is a header-only macro for runtime assertions:
-
-.. code-block:: cpp
-
-   #include <torch/headeronly/util/Exception.h>
-
-   STD_TORCH_CHECK(condition, "Error message with ", variable, " interpolation");
-
-Core Types
-^^^^^^^^^^
-
-The following ``c10::`` types are available as header-only versions under
-``torch::headeronly::``:
-
-- ``torch::headeronly::ScalarType`` - Tensor data types (Float, Double, Int, etc.)
-- ``torch::headeronly::DeviceType`` - Device types (CPU, CUDA, etc.)
-- ``torch::headeronly::MemoryFormat`` - Memory layout formats (Contiguous, ChannelsLast, etc.)
-- ``torch::headeronly::Layout`` - Tensor layouts (Strided, Sparse, etc.)
-
-.. code-block:: cpp
-
-   #include <torch/headeronly/core/ScalarType.h>
-   #include <torch/headeronly/core/DeviceType.h>
-   #include <torch/headeronly/core/MemoryFormat.h>
-   #include <torch/headeronly/core/Layout.h>
-
-   auto dtype = torch::headeronly::ScalarType::Float;
-   auto device_type = torch::headeronly::DeviceType::CUDA;
-   auto memory_format = torch::headeronly::MemoryFormat::Contiguous;
-   auto layout = torch::headeronly::Layout::Strided;
-
-TensorAccessor
-^^^^^^^^^^^^^^
-
-``TensorAccessor`` provides efficient, bounds-checked access to tensor data.
-You can construct one from a stable tensor's data pointer, sizes, and strides:
-
-.. code-block:: cpp
-
-   #include <torch/headeronly/core/TensorAccessor.h>
-
-   // Create a TensorAccessor for a 2D float tensor
-   auto sizes = tensor.sizes();
-   auto strides = tensor.strides();
-   torch::headeronly::TensorAccessor<float, 2> accessor(
-       static_cast<float*>(tensor.mutable_data_ptr()),
-       sizes.data(),
-       strides.data());
-
-   // Access elements
-   float value = accessor[i][j];
-
-Dispatch Macros
-^^^^^^^^^^^^^^^
-
-Header-only dispatch macros (THO = Torch Header Only) are available for
-dtype and device dispatching:
-
-.. code-block:: cpp
-
-   #include <torch/headeronly/core/Dispatch.h>
-
-   THO_DISPATCH_FLOATING_TYPES(tensor.scalar_type(), "my_kernel", [&] {
-       // scalar_t is the resolved type
-       auto* data = tensor.data_ptr<scalar_t>();
-   });
-
-Full API List
-^^^^^^^^^^^^^
-
-For the complete list of header-only APIs, see ``torch/header_only_apis.txt``
-in the PyTorch source tree.
-
-Stable Operators
-----------------
-
-Tensor Creation
-^^^^^^^^^^^^^^^
-
-.. doxygenfunction:: torch::stable::empty
-
-.. doxygenfunction:: torch::stable::empty_like
-
-.. doxygenfunction:: torch::stable::new_empty(const torch::stable::Tensor &self, torch::headeronly::IntHeaderOnlyArrayRef size, std::optional<torch::headeronly::ScalarType> dtype, std::optional<torch::headeronly::Layout> layout, std::optional<torch::stable::Device> device, std::optional<bool> pin_memory)
-
-.. doxygenfunction:: torch::stable::new_zeros(const torch::stable::Tensor &self, torch::headeronly::IntHeaderOnlyArrayRef size, std::optional<torch::headeronly::ScalarType> dtype, std::optional<torch::headeronly::Layout> layout, std::optional<torch::stable::Device> device, std::optional<bool> pin_memory)
-
-.. doxygenfunction:: torch::stable::full
-
-.. doxygenfunction:: torch::stable::from_blob
-
-Tensor Manipulation
-^^^^^^^^^^^^^^^^^^^
-
-.. doxygenfunction:: torch::stable::clone
-
-.. doxygenfunction:: torch::stable::contiguous
-
-.. doxygenfunction:: torch::stable::reshape
-
-.. doxygenfunction:: torch::stable::view
-
-.. doxygenfunction:: torch::stable::flatten
-
-.. doxygenfunction:: torch::stable::squeeze
-
-.. doxygenfunction:: torch::stable::unsqueeze
-
-.. doxygenfunction:: torch::stable::transpose
-
-.. doxygenfunction:: torch::stable::select
-
-.. doxygenfunction:: torch::stable::narrow
-
-.. doxygenfunction:: torch::stable::pad
-
-
-Device and Type Conversion
-^^^^^^^^^^^^^^^^^^^^^^^^^^
-
-.. doxygenfunction:: torch::stable::to(const torch::stable::Tensor &self, std::optional<torch::headeronly::ScalarType> dtype, std::optional<torch::headeronly::Layout> layout, std::optional<torch::stable::Device> device, std::optional<bool> pin_memory, bool non_blocking, bool copy, std::optional<torch::headeronly::MemoryFormat> memory_format)
-
-.. doxygenfunction:: torch::stable::to(const torch::stable::Tensor &self, torch::stable::Device device, bool non_blocking, bool copy)
-
-.. doxygenfunction:: torch::stable::fill_
-
-.. doxygenfunction:: torch::stable::zero_
-
-.. doxygenfunction:: torch::stable::copy_
-
-.. doxygenfunction:: torch::stable::matmul
-
-.. doxygenfunction:: torch::stable::amax(const torch::stable::Tensor &self, int64_t dim, bool keepdim)
-
-.. doxygenfunction:: torch::stable::amax(const torch::stable::Tensor &self, torch::headeronly::IntHeaderOnlyArrayRef dims, bool keepdim)
-
-.. doxygenfunction:: torch::stable::sum
-
-.. doxygenfunction:: torch::stable::sum_out
-
-.. doxygenfunction:: torch::stable::subtract
-
-.. doxygenfunction:: torch::stable::parallel_for
-
-.. doxygenfunction:: torch::stable::get_num_threads
-
-
-Parallelization Utilities
-^^^^^^^^^^^^^^^^^^^^^^^^^
-
-.. doxygenfunction:: torch::stable::parallel_for
-
-.. doxygenfunction:: torch::stable::get_num_threads
diff --git a/docs/source/accelerator.md b/docs/source/accelerator.md
index c5904563ee711..1af297d315775 100644
--- a/docs/source/accelerator.md
+++ b/docs/source/accelerator.md
@@ -2,6 +2,7 @@
 
 ```{eval-rst}
 .. automodule:: torch.accelerator
+   :no-members:
 ```
 
 ```{eval-rst}
@@ -20,12 +21,23 @@
     set_device_idx
     current_device_index
     current_device_idx
+    get_device_capability
     set_stream
     current_stream
     synchronize
     device_index
 ```
 
+## Graphs
+
+```{eval-rst}
+.. autosummary::
+    :toctree: generated
+    :nosignatures:
+
+    Graph
+```
+
 ```{eval-rst}
 .. automodule:: torch.accelerator.memory
 ```
@@ -40,6 +52,7 @@
     :nosignatures:
 
      empty_cache
+     empty_host_cache
      get_memory_info
      max_memory_allocated
      max_memory_reserved
diff --git a/docs/source/backends.md b/docs/source/backends.md
index 6f8791d9a6087..01faa514cdf48 100644
--- a/docs/source/backends.md
+++ b/docs/source/backends.md
@@ -24,6 +24,7 @@ These backends include:
 - `torch.backends.nnpack`
 - `torch.backends.openmp`
 - `torch.backends.opt_einsum`
+- `torch.backends.python_native`
 - `torch.backends.xeon`
 
 ## torch.backends.cpu
@@ -105,6 +106,10 @@ These backends include:
 .. autofunction:: torch.backends.cuda.preferred_rocm_fa_library
 ```
 
+```{eval-rst}
+.. autofunction:: torch.backends.cuda.is_ck_sdpa_available
+```
+
 ```{eval-rst}
 .. autofunction:: torch.backends.cuda.preferred_linalg_library
 ```
@@ -357,6 +362,7 @@ These backends include:
 ```{eval-rst}
 .. py:module:: torch.backends.kleidiai
 
+.. autofunction:: torch.backends.kleidiai.is_available
 ```
 
 ## torch.backends.opt_einsum
@@ -395,6 +401,122 @@ These backends include:
 
 ```
 
+## torch.backends.python_native
+
+```{eval-rst}
+.. automodule:: torch.backends.python_native
+```
+
+The `torch.backends.python_native` module provides user control over native operators implemented in python
+via. DSLs (Domain Specific Languages) that are defined in `torch._native`. This allows users to selectively
+enable or disable high-performance implementations from various DSLs like Triton and CuteDSL.
+
+### Module-level Functions
+
+```{eval-rst}
+.. autofunction:: torch.backends.python_native.get_dsl_operations
+```
+
+```{eval-rst}
+.. autofunction:: torch.backends.python_native.disable_operations
+```
+
+```{eval-rst}
+.. autofunction:: torch.backends.python_native.enable_operations
+```
+
+```{eval-rst}
+.. autofunction:: torch.backends.python_native.disable_dispatch_keys
+```
+
+```{eval-rst}
+.. autofunction:: torch.backends.python_native.enable_dispatch_keys
+```
+
+```{eval-rst}
+.. autofunction:: torch.backends.python_native.operations_disabled
+```
+
+### Module-level Properties
+
+```{eval-rst}
+.. attribute:: available_dsls
+
+    A :class:`list` of :class:`str` containing the names of DSLs that are available at runtime.
+    This is a subset of :attr:`all_dsls` that have their runtime dependencies satisfied.
+```
+
+```{eval-rst}
+.. attribute:: all_dsls
+
+    A :class:`list` of :class:`str` containing the names of all registered DSLs, whether
+    available at runtime or not.
+```
+
+### DSL Controllers
+
+For each registered DSL (e.g., `triton`, `cutedsl`), auto-populated controller modules are available:
+
+```{eval-rst}
+.. currentmodule:: torch.backends.python_native
+```
+
+#### DSL Properties
+
+Each DSL controller (e.g., `torch.backends.python_native.triton`) provides the following properties:
+
+| Property | Type | Description |
+|----------|------|-------------|
+| `name` | `str` | The name of the DSL |
+| `available` | `bool` | Whether the DSL's runtime dependencies are available |
+| `enabled` | `bool` | Controls whether all operations from this DSL are enabled. Setting to `False` disables all operations from the DSL, while `True` re-enables them |
+| `version` | `Version` or `None` | The version of the DSL runtime, if available. Returns `None` if the DSL is not available |
+
+#### DSL Methods
+
+Each DSL controller provides the following methods:
+
+**disable()**
+    Disable all operations from this DSL.
+
+**enable()**
+    Re-enable all operations from this DSL.
+
+**disabled()**
+    Context manager that temporarily disables all operations from this DSL.
+    Operations are automatically re-enabled when exiting the context.
+
+    Example::
+
+        with torch.backends.python_native.triton.disabled():
+            # Triton operations are disabled here
+            result = model(input)
+        # Triton operations restored here
+
+### Usage Examples
+
+```{eval-rst}
+.. code-block:: python
+
+    import torch.backends.python_native as pn
+
+    # Query available DSLs
+    print(pn.available_dsls)  # ['triton', 'cutedsl']
+
+    # Disable all Triton operations
+    pn.triton.enabled = False
+
+    # Temporarily disable CuteDSL operations
+    with pn.cutedsl.disabled():
+        result = model(input)  # CuteDSL ops disabled
+
+    # Disable specific operations across all DSLs
+    pn.disable_operations('scaled_mm', '_flash_attention_forward')
+
+    # Query operations for a specific DSL
+    triton_ops = pn.get_dsl_operations('triton')
+```
+
 ## torch.backends.xeon
 
 ```{eval-rst}
diff --git a/docs/source/benchmark_utils.md b/docs/source/benchmark_utils.md
index 8f58b60b03426..7324ed7e35731 100644
--- a/docs/source/benchmark_utils.md
+++ b/docs/source/benchmark_utils.md
@@ -38,13 +38,30 @@
     :members:
 ```
 
-% These are missing documentation. Adding them here until a better place
-% is made in this file.
+# torch.utils.benchmark.examples
 
 ```{eval-rst}
 .. py:module:: torch.utils.benchmark.examples
 ```
 
+# torch.utils.benchmark.examples.spectral_ops_fuzz_test
+
+```{eval-rst}
+.. automodule:: torch.utils.benchmark.examples.spectral_ops_fuzz_test
+```
+
+```{eval-rst}
+.. currentmodule:: torch.utils.benchmark.examples.spectral_ops_fuzz_test
+```
+
+```{eval-rst}
+.. autosummary::
+    :toctree: generated
+    :nosignatures:
+
+    run_benchmark
+```
+
 ```{eval-rst}
 .. py:module:: torch.utils.benchmark.op_fuzzers
 ```
diff --git a/docs/source/community/persons_of_interest.rst b/docs/source/community/persons_of_interest.rst
index 5a719bf51c157..8fbf8bfa202eb 100644
--- a/docs/source/community/persons_of_interest.rst
+++ b/docs/source/community/persons_of_interest.rst
@@ -32,7 +32,7 @@ Core Maintainers
 -  Joe Spisak (`jspisak <https://github.com/jspisak>`__)
 -  (emeritus) Dmytro Dzhulgakov (`dzhulgakov <https://github.com/dzhulgakov>`__)
 
-Module-level maintainers
+Core Module maintainers
 ------------------------
 
 NN APIs (torch.nn)
@@ -70,72 +70,68 @@ TorchDynamo
 -  Animesh Jain (`anijain2305 <https://github.com/anijain2305>`__)
 -  Jason Ansel (`jansel <https://github.com/jansel>`__)
 -  Edward Yang (`ezyang <https://github.com/ezyang>`__)
+-  William Wen (`williamwen <https://github.com/williamwen>`__)
 
 TorchInductor
 ~~~~~~~~~~~~~
 
 -  Elias Ellison (`eellison <https://github.com/eellison>`__)
--  Horace He (`Chillee <https://github.com/Chillee>`__)
 -  Shunting Zhang (`shunting314 <https://github.com/shunting314>`__)
 -  Jason Ansel (`jansel <https://github.com/jansel>`__)
--  Jiong Gong (`jgong5 <https://github.com/jgong5>`__)
+-  Eikan Wang (`eikanwang <https://github.com/eikanwang>`__)
+-  (emeritus) Horace He (`Chillee <https://github.com/Chillee>`__)
+-  (emeritus) Jiong Gong (`jgong5 <https://github.com/jgong5>`__)
 
 Cudagraph Tree
 ~~~~~~~~~~~~~~
 
 -  Elias Ellison (`eellison <https://github.com/eellison>`__)
+-  Boyuan Feng (`boyuanfeng <https://github.com/boyuanfeng>`__)
 
 PT2 Dispatcher
 ~~~~~~~~~~~~~~
 
--  Brian Hirsh (`bdhirsh <https://github.com/bdhirsh>`__)
--  Richard Zou (`zou3519 <https://github.com/zou3519>`__)
--  Horace He (`Chillee <https://github.com/Chillee>`__)
+-  Aaron Orenstein (`aorenste <https://github.com/aorenste>`__)
 -  Edward Yang (`ezyang <https://github.com/ezyang>`__)
+-  Richard Zou (`zou3519 <https://github.com/zou3519>`__)
+-  (emeritus) Brian Hirsh (`bdhirsh <https://github.com/bdhirsh>`__)
+-  (emeritus) Horace He (`Chillee <https://github.com/Chillee>`__)
 
 PT2 Export (torch.export)
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
--  Avik Chaudhuri (`avikchaudhuri <https://github.com/avikchaudhuri>`__)
--  Yanan Cao (`gmagogsfm <https://github.com/gmagogsfm>`__)
+-  Angela Yi (`angelayi <https://github.com/angelayi>`__)
+-  Milad Mohammadi (`miladm <https://github.com/miladm>`__)
+-  (emeritus) Avik Chaudhuri (`avikchaudhuri <https://github.com/avikchaudhuri>`__)
+-  (emeritus) Yanan Cao (`gmagogsfm <https://github.com/gmagogsfm>`__)
 
 AOT Inductor (AOTI) & AOTI Runtime
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
 -  Bin Bao (`desertfire <https://github.com/desertfire>`__)
--  Angela Yi (`angelayi <https://github.com/angelayi>`__)
--  Yang Chen (`chenyang78 <https://github.com/chenyang78>`__)
-
-Compilers (JIT / TorchScript / Package / Deploy)
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-
--  (emeritus) Elias Ellison (`eellison <https://github.com/eellison>`__)
--  (emeritus) Michael Suo (`suo <https://github.com/suo>`__)
--  (emeritus) Yanan Cao (`gmagogsfm <https://github.com/gmagogsfm>`__)
--  (emeritus) James Reed (`jamesr66a <https://github.com/jamesr66a>`__)
--  (emeritus) Jason Ansel (`jansel <https://github.com/jansel>`__)
--  (emeritus) Jiong Gong (`jgong5 <https://github.com/jgong5>`__)
--  (emeritus) Zach Devito (`zdevito <https://github.com/zdevito>`__)
-
+-  Shangdi Yu (`shangdiy <https://github.com/shangdiy>`__)
+-  (emeritus) Angela Yi (`angelayi <https://github.com/angelayi>`__)
+-  (emeritus) Yang Chen (`chenyang78 <https://github.com/chenyang78>`__)
 
 Distributions & RNG
 ~~~~~~~~~~~~~~~~~~~
 
 -  Fritz Obermeyer (`fritzo <https://github.com/fritzo>`__)
--  Neeraj Pradhan (`neerajprad <https://github.com/neerajprad>`__)
--  Alican Bozkurt (`alicanb <https://github.com/alicanb>`__)
+-  (emeritus) Neeraj Pradhan (`neerajprad <https://github.com/neerajprad>`__)
+-  (emeritus) Alican Bozkurt (`alicanb <https://github.com/alicanb>`__)
 -  (emeritus) Vishwak Srinivasan (`vishwakftw <https://github.com/vishwakftw>`__)
 
 Distributed
 ~~~~~~~~~~~
 -  Will Constable (`wconstab <https://github.com/wconstab>`__)
 -  Howard Huang (`H-Huang <https://github.com/H-Huang>`__)
--  Wanchao Liang (`wanchaol <https://github.com/wanchaol>`__)
--  Ke Wen (`kwen2501 <https://github.com/kwen2501>`__)
+-  (emeritus) Wanchao Liang (`wanchaol <https://github.com/wanchaol>`__)
+-  (emeritus) Ke Wen (`kwen2501 <https://github.com/kwen2501>`__)
 -  Chien-Chin Huang (`fegin <https://github.com/fegin>`__)
 -  Tristan Rice (`d4l3k <https://github.com/d4l3k>`__)
 -  Junjie Wang (`fduwjj <https://github.com/fduwjj>`__)
 -  Wei Feng (`weifengpy <https://github.com/weifengpy>`__)
+-  Pei Zhang (`zpcore <https://github.com/zpcore>`__)
 -  (emeritus) Shen Li (`mrshenli <https://github.com/mrshenli>`__)
 -  (emeritus) Pritam Damania (`pritamdamania87 <https://github.com/pritamdamania87>`__)
 -  (emeritus) Yanli Zhao (`zhaojuanmao <https://github.com/zhaojuanmao>`__)
@@ -150,6 +146,7 @@ Distributed
 Multiprocessing
 ~~~~~~~~~~~~~~~
 
+-  Alban Desmaison (`albanD <https://github.com/albanD>`__)
 -  (emeritus) Simon Wang (`SsnL <https://github.com/SsnL>`__)
 -  (emeritus) Vitaly Fedyunin (`VitalyFedyunin <https://github.com/VitalyFedyunin>`__)
 -  (emeritus) Adam Paszke (`apaszke <https://github.com/apaszke>`__)
@@ -175,17 +172,17 @@ Sparse (torch.sparse)
 NestedTensor (torch.nested)
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
--  Joel Schlosser (`jbschlosser <https://github.com/jbschlosser>`__)
--  Christian Puhrsch (`cpuhrsch <https://github.com/cpuhrsch>`__)
--  Driss Guessous (`drisspg <https://github.com/drisspg>`__)
--  Mikayla Gawarecki (`mikaylagawarecki <https://github.com/mikaylagawarecki>`__)
--  Alban Desmaison (`albanD <https://github.com/albanD>`__)
+-  (emeritus) Joel Schlosser (`jbschlosser <https://github.com/jbschlosser>`__)
+-  (emeritus) Christian Puhrsch (`cpuhrsch <https://github.com/cpuhrsch>`__)
+-  (emeritus) Driss Guessous (`drisspg <https://github.com/drisspg>`__)
+-  (emeritus) Mikayla Gawarecki (`mikaylagawarecki <https://github.com/mikaylagawarecki>`__)
+-  (emeritus) Alban Desmaison (`albanD <https://github.com/albanD>`__)
 -  (emeritus) Natalia Gimelshein (`ngimel <https://github.com/ngimel>`__)
 
 MaskedTensor (torch.masked)
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
--  Christian Puhrsch (`cpuhrsch <https://github.com/cpuhrsch>`__)
+-  (emeritus) Christian Puhrsch (`cpuhrsch <https://github.com/cpuhrsch>`__)
 -  (emeritus) George Qi (`george-qi <https://github.com/george-qi>`__)
 
 Fast Fourier Transform (torch.fft)
@@ -194,6 +191,15 @@ Fast Fourier Transform (torch.fft)
 -  (emeritus) Mike Ruberry (`mruberry <https://github.com/mruberry>`__)
 -  (emeritus) Peter Bell (`peterbell10 <https://github.com/peterbell10>`__)
 
+Quantization (torch/ao)
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+-  Vasiliy Kuznetsov (`vkuzo <https://github.com/vkuzo>`__)
+-  Jerry Zhang (`jerryzh168 <https://github.com/jerryzh168>`__)
+-  (emeritus) Mark Saroufim (`msaroufim <https://github.com/msaroufim>`__)
+-  (emeritus) Zafar Takhirov (`z-a-f <https://github.com/z-a-f>`__)
+-  (emeritus) Raghuraman Krishnamoorthi (`raghuramank100 <https://github.com/raghuramank100>`__)
+
 MKLDNN
 ~~~~~~
 
@@ -219,7 +225,6 @@ CUDA
 -  (emeritus) Andrew Tulloch (`ajtulloch <https://github.com/ajtulloch>`__)
 -  (emeritus) Xiaoqiang Zheng (`zheng-xq <https://github.com/zheng-xq>`__)
 
-
 AMD/ROCm/HIP
 ~~~~~~~~~~~~
 
@@ -233,6 +238,42 @@ XPU/Intel GPU
 - Eikan Wang (`EikanWang <https://github.com/EikanWang>`__)
 - Guangye Yu (`guangyey <https://github.com/guangyey>`__)
 
+Windows
+~~~~~~~
+
+-  (emeritus) Guoliang Hua (`nbcsm <https://github.com/nbcsm>`__)
+-  (emeritus) Teng Gao (`gaoteng-git <https://github.com/gaoteng-git>`__)
+-  (emeritus) Peter Johnson (`peterjc123 <https://github.com/peterjc123>`__)
+
+Apple M1/MPS/Metal
+~~~~~~~~~~~~~~~~~~~~
+
+-  Nikita Shulga (`malfet <https://github.com/malfet>`__)
+-  (emeritus) Kulin Seth (`kulinseth <https://github.com/kulinseth>`__)
+-  (emeritus) Alban Desmaison (`alband <https://github.com/alband>`__)
+-  (emeritus) Ramin Azarmehr (`razarmehr <https://github.com/razarmehr>`__)
+
+PowerPC
+~~~~~~~
+
+-  (emeritus) Alfredo Mendoza (`avmgithub <https://github.com/avmgithub>`__)
+
+x86 CPU
+~~~~~~~
+
+-  Mingfei Ma (`mingfeima <https://github.com/mingfeima>`__)
+-  Jiong Gong (`jgong5 <https://github.com/jgong5>`__)
+
+AArch64 CPU
+~~~~~~~~~~~~
+
+-  Sunita Nadampalli (`snadampal <https://github.com/snadampal>`__)
+
+Out-of-tree Backend Integration (PrivateUse1)
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+-  Jiawei Li (`fffrog <https://github.com/fffrog>`__)
+
 Build + CI
 ~~~~~~~~~~
 
@@ -248,19 +289,23 @@ Build + CI
 -  (emeritus) Edward Yang (`ezyang <https://github.com/ezyang>`__)
 -  (emeritus) Karl Ostmo (`kostmo <https://github.com/kostmo>`__)
 
-Performance Tools
-~~~~~~~~~~~~~~~~~
+Profiling and Monitoring Tools
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
--  Taylor Robie (`robieta <https://github.com/robieta>`__)
--  Xu Zhao (`xuzhao9 <https://github.com/xuzhao9>`__)
+-  Scott Schneider (`scotts <https://github.com/scotts>`__)
+-  Divyansh Khanna (`divyanshk <https://github.com/divyanshk>`__)
+-  (emeritus) Taylor Robie (`robieta <https://github.com/robieta>`__)
+-  (emeritus) Xu Zhao (`xuzhao9 <https://github.com/xuzhao9>`__)
 -  (emeritus) Victor Bittorf (`bitfort <https://github.com/bitfort>`__)
 -  (emeritus) Gisle Dankel (`gdankel <https://github.com/gdankel>`__)
 -  (emeritus) Natalia Gimelshein (`ngimel <https://github.com/ngimel>`__)
 -  (emeritus) Mingzhe Li (`mingzhe09088 <https://github.com/mingzhe09088>`__)
 
-C++ API
-~~~~~~~
+C++ API (libtorch + StableABI)
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
+-  Jane Xu (`janeyx99 <https://github.com/janeyx99>`__)
+-  Mikayla Gawarecki (`mikaylagawarecki <https://github.com/mikaylagawarecki>`__)
 -  (emeritus) Joel Schlosser (`jbschlosser <https://github.com/jbschlosser>`__)
 -  (emeritus) Will Feng (`yf225 <https://github.com/yf225>`__)
 
@@ -287,94 +332,33 @@ ONNX exporter
 -  (emeritus) Negin Raoof (`neginraoof <https://github.com/neginraoof>`__)
 -  (emeritus) Spandan Tiwari (`spandantiwari <https://github.com/spandantiwari>`__)
 
-LiteInterpreter
-~~~~~~~~~~~~~~~
--  (emeritus) David Reiss (`dreiss <https://github.com/dreiss>`__)
--  (emeritus) Raziel Guevara (`raziel <https://github.com/raziel>`__)
--  (emeritus) Linbin Yu (`linbinyu <https://github.com/linbinyu>`__)
--  (emeritus) Ivan Kobzarev (`IvanKobzarev <https://github.com/IvanKobzarev>`__)
--  (emeritus) Tao Xu (`xta0 <https://github.com/xta0>`__)
-
-Quantization (torch/ao)
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-
--  Mark Saroufim (`msaroufim <https://github.com/msaroufim>`__)
--  Vasiliy Kuznetsov (`vkuzo <https://github.com/vkuzo>`__)
--  Jerry Zhang (`jerryzh168 <https://github.com/jerryzh168>`__)
--  (emeritus) Zafar Takhirov (`z-a-f <https://github.com/z-a-f>`__)
--  (emeritus) Raghuraman Krishnamoorthi (`raghuramank100 <https://github.com/raghuramank100>`__)
-
-
-Windows
-~~~~~~~
-
--  (emeritus) Guoliang Hua (`nbcsm <https://github.com/nbcsm>`__)
--  (emeritus) Teng Gao (`gaoteng-git <https://github.com/gaoteng-git>`__)
--  (emeritus) Peter Johnson (`peterjc123 <https://github.com/peterjc123>`__)
-
-Apple M1/MPS/Metal
-~~~~~~~~~~~~~~~~~~~~
-
--  Kulin Seth (`kulinseth <https://github.com/kulinseth>`__)
--  Alban Desmaison (`alband <https://github.com/alband>`__)
--  Nikita Shulga (`malfet <https://github.com/malfet>`__)
--  (emeritus) Ramin Azarmehr (`razarmehr <https://github.com/razarmehr>`__)
-
-PowerPC
-~~~~~~~
-
--  (emeritus) Alfredo Mendoza (`avmgithub <https://github.com/avmgithub>`__)
-
-x86 CPU
-~~~~~~~
-
--  Mingfei Ma (`mingfeima <https://github.com/mingfeima>`__)
--  Jiong Gong (`jgong5 <https://github.com/jgong5>`__)
-
-AArch64 CPU
-~~~~~~~~~~~~
-
--  Sunita Nadampalli (`snadampal <https://github.com/snadampal>`__)
-
-Out-of-tree Backend Integration (PrivateUse1)
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-
--  Jiawei Li (`fffrog <https://github.com/fffrog>`__)
-
 Docs / Tutorials
 ~~~~~~~~~~~~~~~~
 
 - Svetlana Karslioglu (`svekars <https://github.com/svekars>`__)
 - Ivan Sekyonda (`sekyondaMeta <https://github.com/sekyondaMeta>`__)
 
-Library-level maintainers
--------------------------
+Core Governed External Module maintainers
+------------------------------------------
 
-XLA
-~~~
+ExecuTorch (Edge, Mobile)
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
--  Jack Cao (`JackCaoG <https://github.com/JackCaoG>`__)
--  Han Qi (`qihqi <https://github.com/qihqi>`__)
--  Yifei Teng (`tengyifei <https://github.com/tengyifei>`__)
--  Siyuan Liu (`lsy323 <https://github.com/lsy323>`__)
--  Brian Hirsh (`bdhirsh <https://github.com/bdhirsh>`__)
--  (emeritus) Gregory Chanan (`gchanan <https://github.com/gchanan>`__)
--  (emeritus) Ailing Zhang (`ailzhang <https://github.com/ailzhang>`__)
--  (emeritus) Davide Libenzi (`dlibenzi <https://github.com/dlibenzi>`__)
--  (emeritus) Alex Suhan (`asuhan <https://github.com/asuhan>`__)
--  (emeritus) Daniel Sohn (`jysohn23 <https://github.com/jysohn23>`__)
--  (emeritus) Zach Cain (`zcain117 <https://github.com/zcain117>`__)
+-  Mergen Nachin (`mergennachin <https://github.com/mergennachin>`__)
+-  Jacob Szwejbka (`JacobSzwejbka <https://github.com/JacobSzwejbka>`__)
+-  Digant Desai (`digantdesai <https://github.com/digantdesai>`__)
+-  (emeritus) Kimish Patel (`kimishpatel <https://github.com/kimishpatel>`__)
+-  (emeritus) Dave Bort (`dbort <https://github.com/dbort>`__)
+-  (emeritus) Martin Yuan (`iseeyuan <https://github.com/iseeyuan>`__)
+-  (emeritus) Mengwei Liu (`larryliu0820 <https://github.com/larryliu0820>`__)
+-  (emeritus) Chen Lai (`cccclai <https://github.com/cccclai>`__)
 
-TorchServe
-~~~~~~~~~~
+TorchCodec
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
--  (emeritus) Li Ning (`lxning <https://github.com/lxning>`__)
--  (emeritus) Ankith Gunapal (`agunapal <https://github.com/agunapal>`__)
--  (emeritus) Hamid Shojanazeri (`HamidShojanazeri <https://github.com/HamidShojanazeri>`__)
--  (emeritus) Mark Saroufim (`msaroufIm <https://github.com/msaroufIm>`__)
--  (emeritus) Manoj Rao (`mycpuorg <https://github.com/mycpuorg>`__)
--  (emeritus) Vamshi Dantu (`vdantu <https://github.com/vdantu>`__)
--  (emeritus) Dhanasekar Karuppasamy (`dhanainme <https://github.com/dhanainme>`__)
+-  Nicolas Hug (`NicolasHug <https://github.com/NicolasHug>`__)
+-  Scott Schneider (`scotts <https://github.com/scotts>`__)
+-  (emeritus) Ahmad Sharif (`ahmadsharif1 <https://github.com/ahmadsharif1>`__)
 
 TorchVision
 ~~~~~~~~~~~
@@ -399,7 +383,7 @@ TorchText
 TorchAudio
 ~~~~~~~~~~
 
--  Moto Hira (`mthrok <https://github.com/mthrok>`__)
+-  (emeritus) Moto Hira (`mthrok <https://github.com/mthrok>`__)
 -  (emeritus) Jeff Hwang (`hwangjeff <https://github.com/hwangjeff>`__)
 -  (emeritus) Caroline Chen (`carolineechen <https://github.com/carolineechen>`__)
 -  (emeritus) Xiaohui Zhang (`xiaohui-zhang <https://github.com/xiaohui-zhang>`__)
@@ -414,34 +398,8 @@ TorchRec
 -  Paul Zhang (`PaulZhang12 <https://github.com/PaulZhang12>`__)
 -  (emeritus) Dmytro Ivchenko (`divchenko <https://github.com/divchenko>`__)
 
-TorchX
-~~~~~~
-
--  (emeritus) Tristan Rice (`d4l3k <https://github.com/d4l3k>`__)
--  (emeritus) Kiuk Chung (`kiukchung <https://github.com/kiukchung>`__)
-
-TorchData
-~~~~~~~~~~~~~~~~~~~~~~
-
--  Divyansh Khanna (`divyanshk <https://github.com/divyanshk>`__)
--  Ramanish Singh (`ramanishsingh <https://github.com/ramanishsingh>`__)
--  Lavender (`aelavender <https://github.com/aelavender>`__)
--  Scott Schneider (`scotts <https://github.com/scotts>`__)
--  (emeritus) Andrew Ho (`andrewkho <https://github.com/andrewkho>`__)
-
-TorchArrow
-~~~~~~~~~~~~~~~~~~~~~~
-
--  (emeritus) Wenlei Xie (`wenleix <https://github.com/wenleix>`__)
--  (emeritus) Vitaly Fedyunin (`VitalyFedyunin <https://github.com/VitalyFedyunin>`__)
-
-ExecuTorch (Edge, Mobile)
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-
--  Mergen Nachin (`mergennachin <https://github.com/mergennachin>`__)
--  Kimish Patel (`kimishpatel <https://github.com/kimishpatel>`__)
--  Dave Bort (`dbort <https://github.com/dbort>`__)
--  Martin Yuan (`iseeyuan <https://github.com/iseeyuan>`__)
+Sunsetted Modules
+-----------------
 
 TorchTune
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
@@ -457,9 +415,68 @@ TorchChat
 -  Jesse White (`byjlw <https://github.com/byjlw>`__)
 -  (emeritus) Michael Gschwind (`mikekgfb <https://github.com/mikekgfb>`__)
 
-TorchCodec
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+XLA
+~~~
 
--  Nicolas Hug (`NicolasHug <https://github.com/NicolasHug>`__)
+-  Jack Cao (`JackCaoG <https://github.com/JackCaoG>`__)
+-  Han Qi (`qihqi <https://github.com/qihqi>`__)
+-  Yifei Teng (`tengyifei <https://github.com/tengyifei>`__)
+-  Siyuan Liu (`lsy323 <https://github.com/lsy323>`__)
+-  Brian Hirsh (`bdhirsh <https://github.com/bdhirsh>`__)
+-  Gregory Chanan (`gchanan <https://github.com/gchanan>`__)
+-  Ailing Zhang (`ailzhang <https://github.com/ailzhang>`__)
+-  Davide Libenzi (`dlibenzi <https://github.com/dlibenzi>`__)
+-  Alex Suhan (`asuhan <https://github.com/asuhan>`__)
+-  Daniel Sohn (`jysohn23 <https://github.com/jysohn23>`__)
+-  Zach Cain (`zcain117 <https://github.com/zcain117>`__)
+
+TorchServe
+~~~~~~~~~~
+
+-  Li Ning (`lxning <https://github.com/lxning>`__)
+-  Ankith Gunapal (`agunapal <https://github.com/agunapal>`__)
+-  Hamid Shojanazeri (`HamidShojanazeri <https://github.com/HamidShojanazeri>`__)
+-  Mark Saroufim (`msaroufIm <https://github.com/msaroufIm>`__)
+-  Manoj Rao (`mycpuorg <https://github.com/mycpuorg>`__)
+-  Vamshi Dantu (`vdantu <https://github.com/vdantu>`__)
+-  Dhanasekar Karuppasamy (`dhanainme <https://github.com/dhanainme>`__)
+
+Other Compilers (JIT / TorchScript / Package / Deploy)
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+-  Elias Ellison (`eellison <https://github.com/eellison>`__)
+-  Michael Suo (`suo <https://github.com/suo>`__)
+-  Yanan Cao (`gmagogsfm <https://github.com/gmagogsfm>`__)
+-  James Reed (`jamesr66a <https://github.com/jamesr66a>`__)
+-  Jason Ansel (`jansel <https://github.com/jansel>`__)
+-  Jiong Gong (`jgong5 <https://github.com/jgong5>`__)
+-  Zach Devito (`zdevito <https://github.com/zdevito>`__)
+
+LiteInterpreter
+~~~~~~~~~~~~~~~
+-  David Reiss (`dreiss <https://github.com/dreiss>`__)
+-  Raziel Guevara (`raziel <https://github.com/raziel>`__)
+-  Linbin Yu (`linbinyu <https://github.com/linbinyu>`__)
+-  Ivan Kobzarev (`IvanKobzarev <https://github.com/IvanKobzarev>`__)
+-  Tao Xu (`xta0 <https://github.com/xta0>`__)
+
+TorchX
+~~~~~~
+
+-  Tristan Rice (`d4l3k <https://github.com/d4l3k>`__)
+-  Kiuk Chung (`kiukchung <https://github.com/kiukchung>`__)
+
+TorchData
+~~~~~~~~~~~~~~~~~~~~~~
+
+-  Divyansh Khanna (`divyanshk <https://github.com/divyanshk>`__)
+-  Ramanish Singh (`ramanishsingh <https://github.com/ramanishsingh>`__)
+-  Lavender (`aelavender <https://github.com/aelavender>`__)
 -  Scott Schneider (`scotts <https://github.com/scotts>`__)
--  (emeritus) Ahmad Sharif (`ahmadsharif1 <https://github.com/ahmadsharif1>`__)
+-  Andrew Ho (`andrewkho <https://github.com/andrewkho>`__)
+
+TorchArrow
+~~~~~~~~~~~~~~~~~~~~~~
+
+-  Wenlei Xie (`wenleix <https://github.com/wenleix>`__)
+-  Vitaly Fedyunin (`VitalyFedyunin <https://github.com/VitalyFedyunin>`__)
diff --git a/docs/source/conf.py b/docs/source/conf.py
index e0df43a1b7a1c..21dd9d0724013 100644
--- a/docs/source/conf.py
+++ b/docs/source/conf.py
@@ -78,6 +78,12 @@
     "html_image",
 ]
 
+# Don't execute notebooks during the docs build. Notebook correctness is
+# verified by the separate docs_test CI job; re-executing them here just
+# adds ~3 minutes to the build for no benefit.
+nb_execution_mode = "off"
+suppress_warnings = ["myst-nb.lexer"]
+
 html_baseurl = "https://docs.pytorch.org/docs/stable/"  # needed for sphinx-sitemap
 sitemap_locales = [None]
 sitemap_excludes = [
@@ -150,7 +156,7 @@
         },
     ],
     "show_version_warning_banner": True,
-    "llm_disabled": False,
+    "llm_generate_full": "false",
     "icon_links": [
         {
             "name": "X",
@@ -251,63 +257,20 @@
 }
 
 coverage_ignore_functions = [
-    # torch
-    "typename",
-    # torch.cuda._sanitizer
-    "zip_arguments",
-    "zip_by_key",
-    # torch.distributed.autograd
-    "is_available",
-    # torch.distributed.checkpoint.state_dict
-    "gc_context",
-    # torch.distributed.elastic.events
-    "record_rdzv_event",
-    # torch.distributed.elastic.metrics
-    "initialize_metrics",
-    # torch.distributed.elastic.rendezvous.registry
-    "get_rendezvous_handler",
-    # torch.distributed.launch
-    "launch",
-    "main",
-    "parse_args",
-    # torch.distributed.rpc
-    "is_available",
-    # torch.distributed.run
-    "config_from_args",
-    "determine_local_world_size",
-    "get_args_parser",
-    "get_rdzv_endpoint",
-    "get_use_env",
-    "main",
-    "parse_args",
-    "parse_min_max_nnodes",
-    "run",
-    "run_script_path",
-    # torch.distributions.constraints
-    "is_dependent",
+    "main",  # utility script
+    "run",  # utility script
     # torch.hub
     "import_module",
-    # torch.jit
-    "export_opnames",
     # torch.jit.unsupported_tensor_ops
     "execWrapper",
-    # torch.onnx
-    "unregister_custom_op_symbolic",
-    # torch.ao.quantization
-    "default_eval_fn",
     # torch.backends
     "disable_global_flags",
     "flags_frozen",
-    # torch.distributed.algorithms.ddp_comm_hooks
-    "register_ddp_comm_hook",
     # torch.nn.parallel
     "DistributedDataParallelCPU",
     # torch.utils
-    "set_module",
     "burn_in_info",
     "get_info_and_burn_skeleton",
-    "get_inline_skeleton",
-    "get_model_info",
     "get_storage_info",
     "hierarchical_pickle",
     # torch.amp.autocast_mode
@@ -437,21 +400,6 @@
     "get_default_qconfig",
     "get_default_qconfig_dict",
     "qconfig_equals",
-    # torch.ao.quantization.quantization_mappings
-    "get_default_dynamic_quant_module_mappings",
-    "get_default_dynamic_sparse_quant_module_mappings",
-    "get_default_float_to_quantized_operator_mappings",
-    "get_default_qat_module_mappings",
-    "get_default_qconfig_propagation_list",
-    "get_default_static_quant_module_mappings",
-    "get_default_static_quant_reference_module_mappings",
-    "get_default_static_sparse_quant_module_mappings",
-    "get_dynamic_quant_module_class",
-    "get_embedding_qat_module_mappings",
-    "get_embedding_static_quant_module_mappings",
-    "get_quantized_operator",
-    "get_static_quant_module_class",
-    "no_observer_set",
     # torch.ao.quantization.quantize
     "get_default_custom_config_dict",
     # torch.ao.quantization.quantize_fx
@@ -655,158 +603,35 @@
     "loss_parallel",
     # torch.distributed.tensor.parallel.style
     "make_sharded_output_tensor",
-    # torch.distributions.utils
-    "broadcast_all",
-    "clamp_probs",
-    "logits_to_probs",
-    "probs_to_logits",
-    "tril_matrix_to_vec",
-    "vec_to_tril_matrix",
-    "annotate",
-    "check_dependency",
-    "combine_two_partitions",
-    "get_bfs_level_partition",
-    "get_device_partition_stats",
-    "get_device_to_partitions_mapping",
-    "get_logical_id_to_device",
-    "get_node_to_partition_mapping",
-    "reorganize_partitions",
-    "reset_partition_device",
-    "set_parents_and_children",
-    "get_unique_attr_name_in_module",
-    "split_const_subgraphs",
-    "set_trace",
-    "are_nodes_independent",
-    "may_depend_on",
-    "merge_matmul",
-    "split_result_tensors",
-    "sym_sqrt",
-    "edge",
-    "match",
-    "ordering",
-    "supercedes",
-    "reify_object",
-    "unifiable",
-    "unify_object",
-    "ambiguities",
-    "ambiguous",
-    "consistent",
-    "edge",
-    "ordering",
-    "super_signature",
-    "supercedes",
-    "dispatch",
-    "ismethod",
-    "ambiguity_warn",
-    "halt_ordering",
-    "restart_ordering",
-    "source",
-    "str_signature",
-    "variadic_signature_matches",
-    "variadic_signature_matches_iter",
-    "warning_text",
-    "expand_tuples",
-    "groupby",
-    "raises",
-    "reverse_dict",
-    "isvariadic",
-    "freeze",
-    "hashable",
-    "raises",
-    "reverse_dict",
-    "transitive_get",
-    "xfail",
-    "var",
-    "vars",
-    "check_for_type_equality",
-    "convert_eq",
-    "infer_symbolic_types",
-    "infer_symbolic_types_single_pass",
-    "substitute_all_types",
-    "substitute_solution_one_type",
-    "unify_eq",
-    "bisect",
-    "translation_validation_enabled",
-    "translation_validation_timeout",
-    "z3op",
-    "z3str",
-    "reduce_graph_module",
-    "reduce_package_graph_module",
-    "has_side_effect",
-    "map_aggregate",
-    "map_arg",
-    "check_for_mutable_operation",
-    "create_type_hint",
-    "get_signature_for_torch_op",
-    "normalize_function",
-    "normalize_module",
-    "type_matches",
     # torch.fx.passes.annotate_getitem_nodes
     "annotate_getitem_nodes",
-    # torch.fx.passes.backends.cudagraphs
-    "partition_cudagraphs",
     # torch.fx.passes.dialect.common.cse_pass
     "get_CSE_banned_ops",
     # torch.fx.passes.graph_manipulation
     "get_size_of_all_nodes",
     "get_size_of_node",
     "get_tensor_meta",
-    "replace_target_nodes_with",
-    # torch.fx.passes.infra.pass_manager
-    "pass_result_wrapper",
-    "this_before_that_pass_constraint",
-    # torch.fx.passes.operator_support
-    "any_chain",
-    "chain",
-    "create_op_support",
-    # torch.fx.passes.param_fetch
-    "default_matching",
-    "extract_attrs_for_lowering",
-    "lift_lowering_attrs_to_nodes",
-    # torch.fx.passes.pass_manager
-    "inplace_wrapper",
-    "log_hook",
-    "loop_pass",
-    "these_before_those_pass_constraint",
-    "this_before_that_pass_constraint",
-    # torch.fx.passes.regional_inductor
-    "regional_inductor",
-    # torch.fx.passes.reinplace
-    "reinplace",
     # torch.fx.passes.split_module
     "split_module",
     # torch.fx.passes.split_utils
     "getattr_recursive",
-    "setattr_recursive",
     "split_by_tags",
     # torch.fx.passes.splitter_base
     "generate_inputs_for_submodules",
     # torch.fx.passes.tools_common
     "get_acc_ops_name",
     "get_node_target",
-    "is_node_output_tensor",
     "legalize_graph",
     # torch.fx.passes.utils.common
-    "compare_graphs",
     "lift_subgraph_as_module",
     # torch.fx.passes.utils.fuser_utils
-    "erase_nodes",
     "fuse_as_graphmodule",
     "fuse_by_partitions",
     "insert_subgm",
-    "topo_sort",
-    "validate_partition",
     # torch.fx.passes.utils.source_matcher_utils
-    "check_subgraphs_connected",
     "get_source_partitions",
     # torch.fx.proxy
     "assert_fn",
-    # torch.fx.subgraph_rewriter
-    "replace_pattern",
-    "replace_pattern_with_filters",
-    # torch.fx.tensor_type
-    "is_consistent",
-    "is_more_precise",
     # torch.fx.traceback
     "format_stack",
     "get_current_meta",
@@ -814,9 +639,9 @@
     "preserve_node_meta",
     "reset_grad_fn_seq_nr",
     "set_current_meta",
+    "set_current_replay_node",
     "set_grad_fn_seq_nr",
     "set_stack_trace",
-    "set_current_replay_node",
     "get_current_replay_node",
     # torch.jit.annotations
     "ann_to_type",
@@ -857,13 +682,6 @@
     "quantize_linear_modules",
     "quantize_rnn_cell_modules",
     "quantize_rnn_modules",
-    # torch.library
-    "define",
-    "get_ctx",
-    "impl",
-    "impl_abstract",
-    # torch.masked.maskedtensor.core
-    "is_masked_tensor",
     # torch.masked.maskedtensor.creation
     "as_masked_tensor",
     "masked_tensor",
@@ -1149,19 +967,8 @@
     "check_export_model_diff",
     "verify",
     "verify_aten_graph",
-    # torch.optim.optimizer
-    "register_optimizer_step_post_hook",
-    "register_optimizer_step_pre_hook",
     # torch.overrides
     "enable_reentrant_dispatch",
-    # torch.package.analyze.find_first_use_of_broken_modules
-    "find_first_use_of_broken_modules",
-    # torch.package.analyze.is_from_package
-    "is_from_package",
-    # torch.package.analyze.trace_dependencies
-    "trace_dependencies",
-    # torch.profiler.itt
-    "range",
     # torch.profiler.profiler
     "schedule",
     "supported_activities",
@@ -1169,8 +976,6 @@
     # torch.return_types
     "pytree_register_structseq",
     # torch.serialization
-    "check_module_version_greater_or_equal",
-    "default_restore_location",
     "load",
     "location_tag",
     "mkdtemp",
@@ -1191,8 +996,6 @@
     "hann",
     "kaiser",
     "nuttall",
-    # torch.sparse.semi_structured
-    "to_sparse_semi_structured",
     # torch.utils.backend_registration
     "generate_methods_for_privateuse1_backend",
     "rename_privateuse1_backend",
@@ -2633,12 +2436,89 @@ def setup(app):
     app.connect("autodoc-process-docstring", process_docstring)
     app.connect("html-page-context", hide_edit_button_for_pages)
     app.config.add_last_updated = True
+
+    # Force serial reads to avoid pipe congestion from large env pickles.
+    # Sphinx's parallel read sends the entire environment (100s of MB for
+    # PyTorch) through a 64KB OS pipe per worker, which causes extreme
+    # slowdowns with many workers. Serial reads avoid this overhead while
+    # parallel writes (which send trivial payloads) remain enabled.
+    from sphinx.builders import Builder
+
+    _orig_read_serial = Builder._read_serial
+
+    def _serial_read_ignoring_nproc(self, docnames, nproc=1):
+        return _orig_read_serial(self, docnames)
+
+    Builder._read_parallel = _serial_read_ignoring_nproc
+
+    _fix_katex_server_race(app)
+
     return {"version": "0.1", "parallel_read_safe": True}
 
 
+def _fix_katex_server_race(app):
+    """Retry on ConnectionRefusedError when connecting to the KaTeX server.
+
+    sphinxcontrib-katex 0.9.x starts a Node.js server and connects via Unix
+    socket. There's a race: Python sees the socket file (created by bind())
+    and calls connect() before Node.js has called listen(). On slow CI
+    machines this causes ConnectionRefusedError. Fix by retrying connect().
+    """
+    try:
+        from sphinxcontrib.katex import KaTeXServer
+    except ImportError:
+        return
+
+    import socket as _socket
+    import time as _time
+
+    @classmethod
+    def _start_with_retry(cls, rundir, timeout):
+        from subprocess import PIPE, Popen
+
+        from sphinxcontrib.katex import ONE_MILLISECOND
+
+        socket_path = rundir / "katex.sock"
+        cmd = cls.build_command(socket=socket_path)
+        process = Popen(cmd, stdin=PIPE, stdout=PIPE, cwd=rundir)
+
+        startup_start = _time.monotonic()
+        while not socket_path.is_socket():
+            _time.sleep(ONE_MILLISECOND)
+            if _time.monotonic() - startup_start > timeout:
+                raise cls.timeout_error(timeout)
+
+        # Retry connect() — bind() creates the socket file but listen()
+        # is async in Node.js and may not be ready yet.
+        while True:
+            remaining = startup_start + timeout - _time.monotonic()
+            if remaining <= 0:
+                raise cls.timeout_error(timeout)
+            try:
+                sock = _socket.socket(_socket.AF_UNIX, _socket.SOCK_STREAM)
+                sock.settimeout(remaining)
+                sock.connect(str(socket_path))
+                break
+            except ConnectionRefusedError:
+                sock.close()
+                _time.sleep(ONE_MILLISECOND)
+            except TimeoutError:
+                sock.close()
+                raise cls.timeout_error(timeout)  # noqa: B904
+
+        return process, sock
+
+    KaTeXServer.start_server_process = _start_with_retry
+
+
 def hide_edit_button_for_pages(app, pagename, templatename, context, doctree):
     if pagename.startswith("generated/"):
         context["theme_use_edit_page_button"] = False
+        # Limit sidebar depth and collapse inactive sections for generated
+        # API pages to avoid embedding ~300 KB of navigation markup per page.
+        # Only the active section's children are shown.
+        context["theme_navigation_depth"] = 2
+        context["theme_collapse_navigation"] = "true"
 
 
 # From PyTorch 1.5, we now use autogenerated files to document classes and
diff --git a/docs/source/cuda._sanitizer.rst b/docs/source/cuda._sanitizer.rst
index 658b975693112..e34503204b9b6 100644
--- a/docs/source/cuda._sanitizer.rst
+++ b/docs/source/cuda._sanitizer.rst
@@ -100,3 +100,7 @@ API Reference
 -------------
 
 .. autofunction:: enable_cuda_sanitizer
+
+.. autofunction:: zip_arguments
+
+.. autofunction:: zip_by_key
diff --git a/docs/source/cuda.md b/docs/source/cuda.md
index 94894942b74e5..929f90a823bc7 100644
--- a/docs/source/cuda.md
+++ b/docs/source/cuda.md
@@ -165,6 +165,7 @@
     :toctree: generated
     :nosignatures:
 
+    caching_allocator_disabled
     caching_allocator_enable
 ```
 
@@ -302,10 +303,24 @@ See the docs for {class}`~torch.cuda.green_contexts.GreenContext` for an example
 
 ```{eval-rst}
 .. py:module:: torch.cuda.nccl
+
+.. autofunction:: torch.cuda.nccl.is_available
+```
+
+```{eval-rst}
+.. automodule:: torch.cuda.nvtx
+```
+
+```{eval-rst}
+.. currentmodule:: torch.cuda.nvtx
+```
+
+```{eval-rst}
+.. autofunction:: range_start
 ```
 
 ```{eval-rst}
-.. py:module:: torch.cuda.nvtx
+.. autofunction:: range_end
 ```
 
 ```{eval-rst}
diff --git a/docs/source/ddp_comm_hooks.md b/docs/source/ddp_comm_hooks.md
index 059c388cd003a..5ead08d6d8564 100644
--- a/docs/source/ddp_comm_hooks.md
+++ b/docs/source/ddp_comm_hooks.md
@@ -33,6 +33,13 @@ Particularly, {class}`torch.distributed.GradBucket` represents a bucket of gradi
 .. autofunction:: torch.distributed.GradBucket.parameters
 ```
 
+## Register Communication Hook
+
+```{eval-rst}
+.. currentmodule:: torch.distributed.algorithms.ddp_comm_hooks
+.. autofunction:: register_ddp_comm_hook
+```
+
 ## Default Communication Hooks
 
 Default communication hooks are simple **stateless** hooks, so the input state
diff --git a/docs/source/distributed.fsdp.fully_shard.md b/docs/source/distributed.fsdp.fully_shard.md
index d19c26067df1b..123a34e59bdc9 100644
--- a/docs/source/distributed.fsdp.fully_shard.md
+++ b/docs/source/distributed.fsdp.fully_shard.md
@@ -55,6 +55,83 @@ The user contract for ``fully_shard(model)`` is as follows
   registers hooks to the original module.
 
 
+### Communication Grouping and Scheduling
+
+Each call to ``fully_shard`` creates one **communication group** containing all
+parameters in the module that are not already assigned to a group from an
+earlier call on a submodule. Each group's parameters are all-gathered together
+in one collective before forward, and their gradients are reduce-scattered
+together in one collective after backward. Unlike DDP, FSDP2 has no
+``bucket_cap_mb`` parameter — the communication boundaries are determined
+entirely by which modules you apply ``fully_shard`` to.
+
+Consider a model with four submodules where ``a``, ``b``, ``c``, and ``d``
+denote the number of parameters in each:
+
+```
+model[ m1[a] -> m2[b] -> m3[c] -> m4[d] ]
+```
+
+**If you only call** ``fully_shard(model)`` **(root only)**, all parameters are
+in a single group. This means the entire forward and backward look like:
+
+```
+all-gather(a+b+c+d) -> forward(m1,m2,m3,m4) -> backward(m4,m3,m2,m1) -> reduce-scatter(a+b+c+d)
+```
+
+All communication happens as two large blocking operations with no overlap
+with compute. This is almost never what you want.
+
+**If you apply** ``fully_shard`` **per submodule** — for example, calling
+``fully_shard(m2)``, ``fully_shard(m3)``, and then ``fully_shard(model)`` —
+the remaining parameters (``a`` and ``d``) form the root group, while ``m2``
+and ``m3`` each get their own group.
+
+In **forward**, all-gathers run on a separate CUDA stream, so the next
+module's all-gather can overlap with the current module's forward compute.
+Each module's pre-forward hook issues its own all-gather and waits for it to
+complete before running the module. Because the CPU typically runs ahead of
+the GPU, the next module's all-gather is issued on the AG stream while the
+current module's forward is still executing on the compute stream:
+
+```
+              time ──────────────────────────────────────────────►
+
+compute:      [wait] [ fwd(m1)   | fwd(m2)    | fwd(m3,m4)     ]
+AG stream:    [AG(a,d)]  [AG(b)  |    AG(c)   ]
+```
+
+While ``fwd(m1)`` runs on the compute stream, the CPU fires ``m2``'s
+pre-forward hook, which issues ``AG(b)`` on the AG stream. To make this
+overlap more robust (e.g. when CPU-side overhead reduces the lead), use
+``set_modules_to_forward_prefetch`` to issue the next all-gather earlier —
+inside the current module's pre-forward hook rather than waiting for the next
+module's hook to fire.
+
+In **backward**, FSDP2 additionally prefetches the next module's all-gather
+explicitly and runs reduce-scatters on a separate CUDA stream, all without
+any additional configuration:
+
+```
+              time ──────────────────────────────────────────────►
+
+compute:      [ bwd(m4,m3)     | bwd(m2)        | bwd(m1)       ]
+AG stream:    [AG(c)] [ AG(b)  |   AG(a,d)      ]
+RS stream:                     |[RS(c)]  [ RS(b)|     RS(a,d)   ]
+```
+
+While ``bwd(m4,m3)`` runs on the compute stream, the all-gather for ``b``
+(needed by ``m2``) is prefetched on the AG stream. While ``bwd(m2)`` runs,
+both ``AG(a,d)`` and ``RS(c)`` overlap with compute. This pipelining is why
+the recommended pattern is to apply ``fully_shard`` bottom-up to each layer
+before applying it to the root.
+
+To control the size of each communication group, choose which modules to wrap:
+wrapping more fine-grained modules produces smaller, more overlappable groups
+(similar to smaller DDP buckets), while wrapping fewer modules produces larger
+groups. There is no automatic bucketing — the grouping is explicit and
+determined by the module structure.
+
 Compared to PyTorch FSDP1 (`FullyShardedDataParallel`):
 
 - FSDP2 uses `DTensor`-based dim-0 per-parameter sharding for a simpler
diff --git a/docs/source/distributed.md b/docs/source/distributed.md
index 53c439fe200b3..6be381685f783 100644
--- a/docs/source/distributed.md
+++ b/docs/source/distributed.md
@@ -663,6 +663,10 @@ with torch.profiler():
 
 Please refer to the [profiler documentation](https://pytorch.org/docs/main/profiler.html) for a full overview of profiler features.
 
+```{eval-rst}
+.. autofunction:: torch.distributed.distributed_c10d.record_comm
+```
+
 ## Optimization with Symmetric Memory
 
 ### Copy Engine Collectives
@@ -768,7 +772,22 @@ multiple processes per node for distributed training.
 
 ```{eval-rst}
 .. automodule:: torch.distributed.launch
+```
 
+```{eval-rst}
+.. currentmodule:: torch.distributed.launch
+```
+
+```{eval-rst}
+.. autofunction:: launch
+```
+
+```{eval-rst}
+.. autofunction:: main
+```
+
+```{eval-rst}
+.. autofunction:: parse_args
 ```
 
 ## Spawn utility
@@ -1388,6 +1407,8 @@ If you are running single node training, it may be convenient to interactively b
 
 ```{eval-rst}
 .. py:module:: torch.distributed.launcher.api
+
+.. autofunction:: torch.distributed.launcher.api.launch_agent
 ```
 
 ```{eval-rst}
diff --git a/docs/source/distributed.tensor.md b/docs/source/distributed.tensor.md
index 7361c04924675..d768de0e096eb 100644
--- a/docs/source/distributed.tensor.md
+++ b/docs/source/distributed.tensor.md
@@ -264,7 +264,10 @@ these features.
 
 ```{eval-rst}
 .. autofunction:: register_sharding
+```
 
+```{eval-rst}
+.. autofunction:: implicit_replication
 ```
 
 % modules that are missing docs, add the doc later when necessary
diff --git a/docs/source/distributions.md b/docs/source/distributions.md
index 71c37c386cd12..8825729cd9893 100644
--- a/docs/source/distributions.md
+++ b/docs/source/distributions.md
@@ -682,7 +682,17 @@
 
 .. py:module:: torch.distributions.uniform
 
-.. py:module:: torch.distributions.utils
+.. automodule:: torch.distributions.utils
+
+.. currentmodule:: torch.distributions.utils
+
+.. autosummary::
+    :toctree: generated
+    :nosignatures:
+
+    logits_to_probs
+    probs_to_logits
+    vec_to_tril_matrix
 
 .. py:module:: torch.distributions.von_mises
 
diff --git a/docs/source/elastic/events.rst b/docs/source/elastic/events.rst
index c32136d00302f..b62b4b9cb4ae5 100644
--- a/docs/source/elastic/events.rst
+++ b/docs/source/elastic/events.rst
@@ -14,6 +14,8 @@ API Methods
 
 .. autofunction:: torch.distributed.elastic.events.get_logging_handler
 
+.. autofunction:: torch.distributed.elastic.events.record_rdzv_event
+
 Event Objects
 -----------------
 
diff --git a/docs/source/elastic/run.rst b/docs/source/elastic/run.rst
index 284fc7f755311..bbc8a15d53470 100644
--- a/docs/source/elastic/run.rst
+++ b/docs/source/elastic/run.rst
@@ -4,3 +4,17 @@ torchrun (Elastic Launch)
 ======================================
 
 .. automodule:: torch.distributed.run
+
+.. currentmodule:: torch.distributed.run
+
+.. autosummary::
+    :toctree: generated
+    :nosignatures:
+
+    config_from_args
+    determine_local_world_size
+    main
+    parse_args
+    parse_min_max_nnodes
+    run
+    run_script_path
diff --git a/docs/source/fx.experimental.md b/docs/source/fx.experimental.md
index d9b0ffa0d7007..d842fa6a4edf3 100644
--- a/docs/source/fx.experimental.md
+++ b/docs/source/fx.experimental.md
@@ -77,8 +77,6 @@ These APIs are experimental and subject to change without notice.
     InnerTensorKey
     Specialization
 
-    hint_int
-    size_hint
     is_concrete_int
     is_concrete_bool
     is_concrete_float
@@ -116,11 +114,13 @@ These APIs are experimental and subject to change without notice.
     guard_float
     guard_int
     guard_scalar
-    has_hint
+    guarding_hint_or_throw
+    has_guarding_hint
     has_symbolic_sizes_strides
     is_nested_int
     is_symbol_binding_fx_node
     is_symbolic
+    optimization_hint
     expect_true
     log_lru_cache_stats
 ```
@@ -151,6 +151,7 @@ These APIs are experimental and subject to change without notice.
     decompose
     disable_autocast_cache
     disable_proxy_modes_tracing
+    dispatch_trace
     extract_val
     fake_signature
     fetch_object_proxy
@@ -181,6 +182,7 @@ These APIs are experimental and subject to change without notice.
     :nosignatures:
 
     extract_subgraph
+    matches_module_pattern
     modules_to_mkldnn
     optimize_for_inference
     remove_dropout
@@ -227,6 +229,28 @@ These APIs are experimental and subject to change without notice.
     reify
 ```
 
+## torch.fx.experimental.unification.multipledispatch.utils
+
+```{eval-rst}
+.. currentmodule:: torch.fx.experimental.unification.multipledispatch.utils
+```
+
+```{eval-rst}
+.. automodule:: torch.fx.experimental.unification.multipledispatch.utils
+```
+
+```{eval-rst}
+.. autosummary::
+    :toctree: generated
+    :nosignatures:
+
+    typename
+    expand_tuples
+    groupby
+    raises
+    reverse_dict
+```
+
 ## torch.fx.experimental.unification.unification_tools
 
 ```{eval-rst}
@@ -246,6 +270,7 @@ These APIs are experimental and subject to change without notice.
     assoc_in
     dissoc
     first
+    groupby
     keyfilter
     keymap
     merge
@@ -337,6 +362,7 @@ These APIs are experimental and subject to change without notice.
     view_inference_rule
     register_inference_rule
     transpose_inference_rule
+    range_check
 ```
 
 ## torch.fx.experimental.migrate_gradual_types.constraint_transformation
@@ -426,3 +452,278 @@ These APIs are experimental and subject to change without notice.
     torch_relu_override
     torch_where_override
 ```
+
+## torch.fx.experimental.accelerator_partitioner
+
+```{eval-rst}
+.. automodule:: torch.fx.experimental.accelerator_partitioner
+```
+
+```{eval-rst}
+.. currentmodule:: torch.fx.experimental.accelerator_partitioner
+```
+
+```{eval-rst}
+.. autosummary::
+    :toctree: generated
+    :nosignatures:
+
+    check_dependency
+    combine_two_partitions
+    reorganize_partitions
+    reset_partition_device
+    set_parents_and_children
+```
+
+## torch.fx.experimental.debug
+
+```{eval-rst}
+.. automodule:: torch.fx.experimental.debug
+```
+
+```{eval-rst}
+.. currentmodule:: torch.fx.experimental.debug
+```
+
+```{eval-rst}
+.. autosummary::
+    :toctree: generated
+    :nosignatures:
+
+    set_trace
+```
+
+## torch.fx.experimental.merge_matmul
+
+```{eval-rst}
+.. automodule:: torch.fx.experimental.merge_matmul
+```
+
+```{eval-rst}
+.. currentmodule:: torch.fx.experimental.merge_matmul
+```
+
+```{eval-rst}
+.. autosummary::
+    :toctree: generated
+    :nosignatures:
+
+    are_nodes_independent
+    may_depend_on
+    merge_matmul
+```
+
+## torch.fx.experimental.unification.match
+
+```{eval-rst}
+.. automodule:: torch.fx.experimental.unification.match
+```
+
+```{eval-rst}
+.. currentmodule:: torch.fx.experimental.unification.match
+```
+
+```{eval-rst}
+.. autosummary::
+    :toctree: generated
+    :nosignatures:
+
+    edge
+    match
+    ordering
+    supercedes
+```
+
+## torch.fx.experimental.unification.more
+
+```{eval-rst}
+.. automodule:: torch.fx.experimental.unification.more
+```
+
+```{eval-rst}
+.. currentmodule:: torch.fx.experimental.unification.more
+```
+
+```{eval-rst}
+.. autosummary::
+    :toctree: generated
+    :nosignatures:
+
+    reify_object
+    unifiable
+    unify_object
+```
+
+## torch.fx.experimental.unification.multipledispatch.conflict
+
+```{eval-rst}
+.. automodule:: torch.fx.experimental.unification.multipledispatch.conflict
+```
+
+```{eval-rst}
+.. currentmodule:: torch.fx.experimental.unification.multipledispatch.conflict
+```
+
+```{eval-rst}
+.. autosummary::
+    :toctree: generated
+    :nosignatures:
+
+    ambiguities
+    ambiguous
+    consistent
+    edge
+    ordering
+    super_signature
+    supercedes
+```
+
+## torch.fx.experimental.unification.multipledispatch.core
+
+```{eval-rst}
+.. automodule:: torch.fx.experimental.unification.multipledispatch.core
+```
+
+```{eval-rst}
+.. currentmodule:: torch.fx.experimental.unification.multipledispatch.core
+```
+
+```{eval-rst}
+.. autosummary::
+    :toctree: generated
+    :nosignatures:
+
+    dispatch
+    ismethod
+```
+
+## torch.fx.experimental.unification.multipledispatch.dispatcher
+
+```{eval-rst}
+.. automodule:: torch.fx.experimental.unification.multipledispatch.dispatcher
+```
+
+```{eval-rst}
+.. currentmodule:: torch.fx.experimental.unification.multipledispatch.dispatcher
+```
+
+```{eval-rst}
+.. autosummary::
+    :toctree: generated
+    :nosignatures:
+
+    ambiguity_warn
+    halt_ordering
+    restart_ordering
+    source
+    str_signature
+    variadic_signature_matches
+    variadic_signature_matches_iter
+    warning_text
+```
+
+## torch.fx.experimental.unification.multipledispatch.variadic
+
+```{eval-rst}
+.. automodule:: torch.fx.experimental.unification.multipledispatch.variadic
+```
+
+```{eval-rst}
+.. currentmodule:: torch.fx.experimental.unification.multipledispatch.variadic
+```
+
+```{eval-rst}
+.. autosummary::
+    :toctree: generated
+    :nosignatures:
+
+    isvariadic
+```
+
+## torch.fx.experimental.unification.utils
+
+```{eval-rst}
+.. automodule:: torch.fx.experimental.unification.utils
+```
+
+```{eval-rst}
+.. currentmodule:: torch.fx.experimental.unification.utils
+```
+
+```{eval-rst}
+.. autosummary::
+    :toctree: generated
+    :nosignatures:
+
+    freeze
+    hashable
+    raises
+    reverse_dict
+    transitive_get
+    xfail
+```
+
+## torch.fx.experimental.unification.variable
+
+```{eval-rst}
+.. automodule:: torch.fx.experimental.unification.variable
+```
+
+```{eval-rst}
+.. currentmodule:: torch.fx.experimental.unification.variable
+```
+
+```{eval-rst}
+.. autosummary::
+    :toctree: generated
+    :nosignatures:
+
+    var
+    variables
+    vars
+```
+
+## torch.fx.experimental.unify_refinements
+
+```{eval-rst}
+.. automodule:: torch.fx.experimental.unify_refinements
+```
+
+```{eval-rst}
+.. currentmodule:: torch.fx.experimental.unify_refinements
+```
+
+```{eval-rst}
+.. autosummary::
+    :toctree: generated
+    :nosignatures:
+
+    check_for_type_equality
+    infer_symbolic_types
+    infer_symbolic_types_single_pass
+    substitute_all_types
+    substitute_solution_one_type
+    unify_eq
+```
+
+## torch.fx.experimental.validator
+
+```{eval-rst}
+.. automodule:: torch.fx.experimental.validator
+```
+
+```{eval-rst}
+.. currentmodule:: torch.fx.experimental.validator
+```
+
+```{eval-rst}
+.. autosummary::
+    :toctree: generated
+    :nosignatures:
+
+    bisect
+    translation_validation_enabled
+    translation_validation_timeout
+    z3op
+    z3str
+```
diff --git a/docs/source/fx.md b/docs/source/fx.md
index 9bc5e7f01ca54..50181b3c189e8 100644
--- a/docs/source/fx.md
+++ b/docs/source/fx.md
@@ -8,6 +8,7 @@
 ## Overview
 ```{eval-rst}
 .. automodule:: torch.fx
+   :no-members:
 ```
 
 
@@ -1100,6 +1101,311 @@ The set of leaf modules can be customized by overriding
 .. autofunction:: torch.fx.passes.tools_common.stable_topological_sort
 ```
 
+## torch.fx.annotate
+
+```{eval-rst}
+.. automodule:: torch.fx.annotate
+```
+
+```{eval-rst}
+.. currentmodule:: torch.fx.annotate
+```
+
+```{eval-rst}
+.. autosummary::
+    :toctree: generated
+    :nosignatures:
+
+    annotate
+```
+
+## torch.fx.node
+
+```{eval-rst}
+.. automodule:: torch.fx.node
+```
+
+```{eval-rst}
+.. currentmodule:: torch.fx.node
+```
+
+```{eval-rst}
+.. autosummary::
+    :toctree: generated
+    :nosignatures:
+
+    has_side_effect
+    map_aggregate
+    map_arg
+```
+
+## torch.fx.operator_schemas
+
+```{eval-rst}
+.. automodule:: torch.fx.operator_schemas
+```
+
+```{eval-rst}
+.. currentmodule:: torch.fx.operator_schemas
+```
+
+```{eval-rst}
+.. autosummary::
+    :toctree: generated
+    :nosignatures:
+
+    check_for_mutable_operation
+    create_type_hint
+    type_matches
+```
+
+## torch.fx.traceback
+
+```{eval-rst}
+.. automodule:: torch.fx.traceback
+```
+
+```{eval-rst}
+.. currentmodule:: torch.fx.traceback
+```
+
+```{eval-rst}
+.. autosummary::
+    :toctree: generated
+    :nosignatures:
+
+    annotate_fn
+```
+
+## torch.fx.subgraph_rewriter
+
+```{eval-rst}
+.. currentmodule:: torch.fx.subgraph_rewriter
+```
+
+```{eval-rst}
+.. autosummary::
+    :toctree: generated
+    :nosignatures:
+
+    replace_pattern
+    replace_pattern_with_filters
+```
+
+## torch.fx.tensor_type
+
+```{eval-rst}
+.. currentmodule:: torch.fx.tensor_type
+```
+
+```{eval-rst}
+.. autosummary::
+    :toctree: generated
+    :nosignatures:
+
+    is_consistent
+    is_more_precise
+```
+
+## torch.fx.passes.backends.cudagraphs
+
+```{eval-rst}
+.. currentmodule:: torch.fx.passes.backends.cudagraphs
+```
+
+```{eval-rst}
+.. autosummary::
+    :toctree: generated
+    :nosignatures:
+
+    partition_cudagraphs
+```
+
+## torch.fx.passes.graph_manipulation
+
+```{eval-rst}
+.. currentmodule:: torch.fx.passes.graph_manipulation
+```
+
+```{eval-rst}
+.. autosummary::
+    :toctree: generated
+    :nosignatures:
+
+    replace_target_nodes_with
+```
+
+## torch.fx.passes.infra.pass_manager
+
+```{eval-rst}
+.. currentmodule:: torch.fx.passes.infra.pass_manager
+```
+
+```{eval-rst}
+.. autosummary::
+    :toctree: generated
+    :nosignatures:
+
+    pass_result_wrapper
+    this_before_that_pass_constraint
+```
+
+## torch.fx.passes.operator_support
+
+```{eval-rst}
+.. currentmodule:: torch.fx.passes.operator_support
+```
+
+```{eval-rst}
+.. autosummary::
+    :toctree: generated
+    :nosignatures:
+
+    any_chain
+    chain
+    create_op_support
+```
+
+## torch.fx.passes.param_fetch
+
+```{eval-rst}
+.. automodule:: torch.fx.passes.param_fetch
+   :no-members:
+```
+
+```{eval-rst}
+.. currentmodule:: torch.fx.passes.param_fetch
+```
+
+```{eval-rst}
+.. autofunction:: default_matching
+```
+
+```{eval-rst}
+.. autofunction:: extract_attrs_for_lowering
+```
+
+```{eval-rst}
+.. autofunction:: lift_lowering_attrs_to_nodes
+```
+
+## torch.fx.passes.pass_manager
+
+```{eval-rst}
+.. currentmodule:: torch.fx.passes.pass_manager
+```
+
+```{eval-rst}
+.. autosummary::
+    :toctree: generated
+    :nosignatures:
+
+    inplace_wrapper
+    log_hook
+    loop_pass
+    these_before_those_pass_constraint
+    this_before_that_pass_constraint
+```
+
+## torch.fx.passes.regional_inductor
+
+```{eval-rst}
+.. currentmodule:: torch.fx.passes.regional_inductor
+```
+
+```{eval-rst}
+.. autosummary::
+    :toctree: generated
+    :nosignatures:
+
+    regional_inductor
+```
+
+## torch.fx.passes.reinplace
+
+```{eval-rst}
+.. currentmodule:: torch.fx.passes.reinplace
+```
+
+```{eval-rst}
+.. autosummary::
+    :toctree: generated
+    :nosignatures:
+
+    reinplace
+```
+
+## torch.fx.passes.split_utils
+
+```{eval-rst}
+.. currentmodule:: torch.fx.passes.split_utils
+```
+
+```{eval-rst}
+.. autosummary::
+    :toctree: generated
+    :nosignatures:
+
+    setattr_recursive
+```
+
+## torch.fx.passes.tools_common
+
+```{eval-rst}
+.. currentmodule:: torch.fx.passes.tools_common
+```
+
+```{eval-rst}
+.. autosummary::
+    :toctree: generated
+    :nosignatures:
+
+    is_node_output_tensor
+```
+
+## torch.fx.passes.utils.common
+
+```{eval-rst}
+.. currentmodule:: torch.fx.passes.utils.common
+```
+
+```{eval-rst}
+.. autosummary::
+    :toctree: generated
+    :nosignatures:
+
+    compare_graphs
+```
+
+## torch.fx.passes.utils.fuser_utils
+
+```{eval-rst}
+.. currentmodule:: torch.fx.passes.utils.fuser_utils
+```
+
+```{eval-rst}
+.. autosummary::
+    :toctree: generated
+    :nosignatures:
+
+    erase_nodes
+    topo_sort
+    validate_partition
+```
+
+## torch.fx.passes.utils.source_matcher_utils
+
+```{eval-rst}
+.. currentmodule:: torch.fx.passes.utils.source_matcher_utils
+```
+
+```{eval-rst}
+.. autosummary::
+    :toctree: generated
+    :nosignatures:
+
+    check_subgraphs_connected
+```
 
 <!-- The experimental and passes submodules are missing docs. -->
 <!-- Adding it here for coverage but this doesn't add anything to the -->
@@ -1116,12 +1422,8 @@ The set of leaf modules can be customized by overriding
 .. py:module:: torch.fx.experimental.migrate_gradual_types
 .. py:module:: torch.fx.passes.dialect
 .. py:module:: torch.fx.passes.dialect.common
-.. py:module:: torch.fx.annotate
 .. py:module:: torch.fx.config
-.. py:module:: torch.fx.experimental.accelerator_partitioner
 .. py:module:: torch.fx.experimental.const_fold
-.. py:module:: torch.fx.experimental.debug
-.. py:module:: torch.fx.experimental.merge_matmul
 .. py:module:: torch.fx.experimental.migrate_gradual_types.operation
 .. py:module:: torch.fx.experimental.migrate_gradual_types.util
 .. py:module:: torch.fx.experimental.migrate_gradual_types.z3_types
@@ -1131,23 +1433,10 @@ The set of leaf modules can be customized by overriding
 .. py:module:: torch.fx.experimental.rewriter
 .. py:module:: torch.fx.experimental.schema_type_annotation
 .. py:module:: torch.fx.experimental.unification.dispatch
-.. py:module:: torch.fx.experimental.unification.match
-.. py:module:: torch.fx.experimental.unification.more
-.. py:module:: torch.fx.experimental.unification.multipledispatch.conflict
-.. py:module:: torch.fx.experimental.unification.multipledispatch.core
-.. py:module:: torch.fx.experimental.unification.multipledispatch.dispatcher
-.. py:module:: torch.fx.experimental.unification.multipledispatch.utils
-.. py:module:: torch.fx.experimental.unification.multipledispatch.variadic
-.. py:module:: torch.fx.experimental.unification.utils
-.. py:module:: torch.fx.experimental.unification.variable
-.. py:module:: torch.fx.experimental.unify_refinements
-.. py:module:: torch.fx.experimental.validator
 .. py:module:: torch.fx.graph
 .. py:module:: torch.fx.graph_module
 .. py:module:: torch.fx.immutable_collections
 .. py:module:: torch.fx.interpreter
-.. py:module:: torch.fx.node
-.. py:module:: torch.fx.operator_schemas
 .. py:module:: torch.fx.passes.annotate_getitem_nodes
 .. py:module:: torch.fx.passes.backends.cudagraphs
 .. py:module:: torch.fx.passes.dialect.common.cse_pass
@@ -1160,7 +1449,7 @@ The set of leaf modules can be customized by overriding
 .. py:module:: torch.fx.passes.infra.pass_manager
 .. py:module:: torch.fx.passes.net_min_base
 .. py:module:: torch.fx.passes.operator_support
-.. py:module:: torch.fx.passes.param_fetch
+
 .. py:module:: torch.fx.passes.pass_manager
 .. py:module:: torch.fx.passes.regional_inductor
 .. py:module:: torch.fx.passes.reinplace
@@ -1180,5 +1469,4 @@ The set of leaf modules can be customized by overriding
 .. py:module:: torch.fx.proxy
 .. py:module:: torch.fx.subgraph_rewriter
 .. py:module:: torch.fx.tensor_type
-.. py:module:: torch.fx.traceback
 ```
diff --git a/docs/source/jit_utils.md b/docs/source/jit_utils.md
index de05fbf35c513..2c4fcdc566c2b 100644
--- a/docs/source/jit_utils.md
+++ b/docs/source/jit_utils.md
@@ -3,3 +3,24 @@
 ```{eval-rst}
 .. automodule:: torch.utils.jit
 ```
+
+# torch.utils.jit.log_extract
+
+```{eval-rst}
+.. automodule:: torch.utils.jit.log_extract
+```
+
+```{eval-rst}
+.. currentmodule:: torch.utils.jit.log_extract
+```
+
+```{eval-rst}
+.. autosummary::
+    :toctree: generated
+    :nosignatures:
+
+    run_baseline_no_fusion
+    run_nnc
+    run_nvfuser
+    run_test
+```
diff --git a/docs/source/masked.md b/docs/source/masked.md
index 6a79a6be2ea95..af9bc9911720f 100644
--- a/docs/source/masked.md
+++ b/docs/source/masked.md
@@ -304,6 +304,20 @@ The following ops are currently supported:
     Tensor.view
 ```
 
+## torch.masked.maskedtensor.core
+
+```{eval-rst}
+.. currentmodule:: torch.masked.maskedtensor.core
+```
+
+```{eval-rst}
+.. autosummary::
+    :toctree: generated
+    :nosignatures:
+
+    is_masked_tensor
+```
+
 ```{eval-rst}
 .. This module needs to be documented. Adding here in the meantime
 .. for tracking purposes
diff --git a/docs/source/mps.md b/docs/source/mps.md
index 7336e71c03c05..28fac80b86f0e 100644
--- a/docs/source/mps.md
+++ b/docs/source/mps.md
@@ -13,6 +13,7 @@
     :toctree: generated
     :nosignatures:
 
+    is_available
     device_count
     synchronize
     get_rng_state
diff --git a/docs/source/multiprocessing.md b/docs/source/multiprocessing.md
index 6669fcaa24b30..74944ca58ed20 100644
--- a/docs/source/multiprocessing.md
+++ b/docs/source/multiprocessing.md
@@ -206,6 +206,12 @@ terminate processes upon detecting an error in one of them.
 
 % for tracking purposes
 
+## torch.multiprocessing.pool
+
+```{eval-rst}
+.. currentmodule:: torch.multiprocessing.pool
+```
+
 ```{eval-rst}
 .. py:module:: torch.multiprocessing.pool
 ```
diff --git a/docs/source/nn.attention.varlen.md b/docs/source/nn.attention.varlen.md
index df91e1d968e6a..36eba6b6299dc 100644
--- a/docs/source/nn.attention.varlen.md
+++ b/docs/source/nn.attention.varlen.md
@@ -13,5 +13,8 @@
 .. autofunction:: varlen_attn
 ```
 ```{eval-rst}
+.. autofunction:: varlen_attn_out
+```
+```{eval-rst}
 .. autoclass:: AuxRequest
 ```
diff --git a/docs/source/notes/cuda.rst b/docs/source/notes/cuda.rst
index 860cee41ae832..399bf6a19bdb9 100644
--- a/docs/source/notes/cuda.rst
+++ b/docs/source/notes/cuda.rst
@@ -627,6 +627,24 @@ Available options:
   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.
 
+* ``throw_on_cudamalloc_oom`` (default: ``False``)
+  When set to ``True``, the allocator will preemptively reject allocations that would
+  exceed the ``per_process_memory_fraction`` limit by throwing an ``OutOfMemoryError``
+  instead of attempting the allocation through the driver, which could trigger a fatal
+  GPU runtime abort (e.g., ``HSA_STATUS_ERROR_OUT_OF_RESOURCES`` on ROCm). The rejected
+  allocation skips the retry chain and throws immediately, allowing the caller to catch
+  the exception and handle it gracefully.
+
+  Example configuration:
+
+  .. code-block:: bash
+
+      PYTORCH_CUDA_ALLOC_CONF=per_process_memory_fraction:0.95,throw_on_cudamalloc_oom:True
+
+  This is particularly useful for inference serving, where a fatal GPU OOM would crash the
+  server process. With this option, the serving framework can catch the ``OutOfMemoryError``
+  and reject the individual request while continuing to serve subsequent requests.
+
 .. note::
 
     Some stats reported by the
@@ -1333,6 +1351,8 @@ and you suspect its runtime is at least somewhat CPU-limited.
     https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#creating-a-graph-using-stream-capture
 .. _cudaGraphLaunch:
     https://docs.nvidia.com/cuda/cuda-runtime-api/group__CUDART__GRAPH.html#group__CUDART__GRAPH_1g1accfe1da0c605a577c22d9751a09597
+.. _conditional nodes:
+    https://developer.nvidia.com/blog/dynamic-control-flow-in-cuda-graphs-with-conditional-nodes/
 
 PyTorch API
 ^^^^^^^^^^^
@@ -1421,6 +1441,9 @@ Violating any of these will likely cause a runtime error:
   Avoid using :meth:`Generator.get_state<torch.get_state>` and :meth:`Generator.set_state<torch.set_state>` during capture;
   instead, utilize :meth:`Generator.graphsafe_set_state<torch.Generator.graphsafe_set_state>` and :meth:`Generator.graphsafe_get_state<torch.Generator.graphsafe_get_state>`
   for managing generator states safely within the graph context. This ensures proper RNG operation and generator management within CUDA graphs.
+* Dynamic control flow (based on CPU or GPU data) is prohibited, unless it is based on GPU data and implemented via higher order operator
+  torch.cond(). See :ref:`Data Dependent Control Flow<graph-data-dependent-control-flow>`.
+
 
 
 Violating any of these will likely cause silent numerical errors or undefined behavior:
@@ -1429,7 +1452,6 @@ Violating any of these will likely cause silent numerical errors or undefined be
 * No non-captured CUDA work may run in this process (on any thread) while capture is underway.
 * CPU work is not captured. If the captured ops include CPU work, that work will be elided during replay.
 * Every replay reads from and writes to the same (virtual) memory addresses.
-* Dynamic control flow (based on CPU or GPU data) is prohibited.
 * Dynamic shapes are prohibited. The graph assumes every tensor in the captured op sequence
   has the same size and layout in every replay.
 * Using multiple streams in a capture is allowed, but there are :ref:`restrictions<multistream-capture>`.
@@ -1748,3 +1770,46 @@ If, in the live workload, your callables will run in an order that occasionally
 or if they'll run concurrently, passing them as a tuple to a single invocation of
 :func:`~torch.cuda.make_graphed_callables` is not allowed. Instead, you must call
 :func:`~torch.cuda.make_graphed_callables` separately for each one.
+
+.. _graph-data-dependent-control-flow:
+
+Data Dependent Control Flow
+^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+Data-dependent control flow can with cuda graphs if the control flow
+is implemented using torch.cond(). If your function uses this
+function, compiling it with the "cudagraphs" backend will enable
+control flow in the resulting cuda graph via `conditional nodes`_.
+
+Doing cuda graph capture of eager mode code using torch.cond() will
+also work.
+
+Support for inductor backend to torch.compile is not available yet,
+but there is no fundamental blocker.
+
+An example of using the cudagraphs backend to torch.compile on code
+using torch.cond is demonstrated below::
+
+    import torch
+
+    def true_fn(x):
+        return x.sin()
+
+    def false_fn(x):
+        return x.cos()
+
+    x = torch.randn(4, device="cuda", requires_grad=False)
+    pred = torch.tensor(False, device="cuda", requires_grad=False)
+    def foo(pred, x):
+        with torch.inference_mode():
+            return torch.cond(pred, true_fn, false_fn, [x])
+
+    # First call will run eager for warmup, second call will do graph
+    # capture followed by graph replay, third call and beyond will do
+    # just graph replay.
+    compiled_foo = torch.compile(foo, backend="cudagraphs")
+    for i in range(3):
+        y = compiled_foo(pred, x)
+
+    # will output x.sin()
+    y = compiled_foo(~pred, x)
diff --git a/docs/source/notes/ddp.rst b/docs/source/notes/ddp.rst
index 07804f8bae166..4f74114fabce8 100644
--- a/docs/source/notes/ddp.rst
+++ b/docs/source/notes/ddp.rst
@@ -199,7 +199,7 @@ DistributedDataParallel
     parameters when ``find_unused_parameters`` is set to ``True`` in DDP
     constructor.
 
-.. image:: https://user-images.githubusercontent.com/16999635/72313120-4e7c1c80-3658-11ea-9c6d-44336b2daeac.png
+.. image:: https://github.com/user-attachments/assets/b6fe5258-f1f1-4c73-8e1d-6fd96406faa2
     :alt: ddp_code.png
     :width: 400 px
 
diff --git a/docs/source/notes/get_start_xpu.rst b/docs/source/notes/get_start_xpu.rst
index 5a1442e813805..dc0cb984c47fc 100644
--- a/docs/source/notes/get_start_xpu.rst
+++ b/docs/source/notes/get_start_xpu.rst
@@ -4,7 +4,8 @@ Getting Started on Intel GPU
 Hardware Prerequisite
 ---------------------
 
-For Intel Data Center GPU
+Intel Data Center GPU
+^^^^^^^^^^^^^^^^^^^^^
 
 .. list-table::
    :widths: 50 50 50 50
@@ -19,7 +20,8 @@ For Intel Data Center GPU
      - yes
      - yes
 
-For Intel Client GPU
+Intel Client GPU
+^^^^^^^^^^^^^^^^
 
 +---------------------------------------+-----------------------------------------------------------------------------------------------------+
 | Supported OS                          | Validated Hardware                                                                                  |
@@ -51,37 +53,42 @@ Binaries
 
 Now that we have `Intel GPU Driver <https://www.intel.com/content/www/us/en/developer/articles/tool/pytorch-prerequisites-for-intel-gpu.html>`_ installed, use the following commands to install ``pytorch``, ``torchvision``, ``torchaudio``.
 
-For release wheels
+Stable Releases
+~~~~~~~~~~~~~~~
+
+To install the latest stable release wheels for Intel GPU (XPU):
 
 .. code-block:: bash
 
     pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/xpu
 
-For nightly wheels
+Nightly Builds
+~~~~~~~~~~~~~~
+
+To install the latest preview/nightly wheels:
 
 .. code-block:: bash
 
     pip3 install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/xpu
 
-For previous versions
+Previous Versions
+~~~~~~~~~~~~~~~~~
 
-v2.9.1
+**v2.10.0**
 
 .. code-block:: bash
 
-    pip install torch==2.9.1 torchvision==0.24.1 torchaudio==2.9.1 --index-url https://download.pytorch.org/whl/xpu
+    pip install torch==2.10.0 torchvision==0.25.0 torchaudio==2.10.0 --index-url https://download.pytorch.org/whl/xpu
 
-v2.9.0
+**v2.9.1**
 
 .. code-block:: bash
 
-    pip install torch==2.9.0 torchvision==0.24.0 torchaudio==2.9.0 --index-url https://download.pytorch.org/whl/xpu
-
-v2.8.0
+    pip install torch==2.9.1 torchvision==0.24.1 torchaudio==2.9.1 --index-url https://download.pytorch.org/whl/xpu
 
-.. code-block:: bash
+.. note::
 
-    pip install torch==2.8.0 torchvision==0.23.0 torchaudio==2.8.0 --index-url https://download.pytorch.org/whl/xpu
+   For older wheels, please refer to the `previous versions <https://pytorch.org/get-started/previous-versions/>`_ page and ensure you use the ``xpu`` index URL.
 
 From Source
 ^^^^^^^^^^^
@@ -137,7 +144,7 @@ Here are a few inference workflow examples.
 
 
 Inference with FP32
-"""""""""""""""""""
+~~~~~~~~~~~~~~~~~~~
 
 .. code-block:: python
 
@@ -157,7 +164,7 @@ Inference with FP32
     print("Execution finished")
 
 Inference with AMP
-""""""""""""""""""
+~~~~~~~~~~~~~~~~~~
 
 .. code-block:: python
 
@@ -181,7 +188,7 @@ Inference with AMP
     print("Execution finished")
 
 Inference with ``torch.compile``
-""""""""""""""""""""""""""""""""
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
 .. code-block:: python
 
@@ -222,7 +229,7 @@ Training Examples
 Here is a few training workflow examples.
 
 Train with FP32
-"""""""""""""""
+~~~~~~~~~~~~~~~
 
 .. code-block:: python
 
@@ -279,7 +286,7 @@ Train with FP32
     print("Execution finished")
 
 Train with AMP
-""""""""""""""
+~~~~~~~~~~~~~~
 
 .. note::
    Training with ``GradScaler`` requires hardware support for ``FP64``. ``FP64`` is not natively supported by the Intel® Arc™ A-Series Graphics. If you run your workloads on Intel® Arc™ A-Series Graphics, please disable ``GradScaler``.
@@ -347,7 +354,7 @@ Train with AMP
     print("Execution finished")
 
 Train with ``torch.compile``
-""""""""""""""""""""""""""""
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
 .. code-block:: python
 
diff --git a/docs/source/notes/libtorch_stable_abi.md b/docs/source/notes/libtorch_stable_abi.md
index f1dfca1080522..bc33a21dc848f 100644
--- a/docs/source/notes/libtorch_stable_abi.md
+++ b/docs/source/notes/libtorch_stable_abi.md
@@ -32,7 +32,7 @@ It consists of
 
 We are continuing to improve coverage in our `torch/csrc/stable` APIs. Please file an issue if you'd like to see support for particular APIs in your custom extension.
 
-For complete API documentation of the stable operators, see the [Torch Stable API cpp documentation](https://docs.pytorch.org/cppdocs/stable.html). <!-- @lint-ignore: URL won't exist till stable.rst cpp docs are published in 2.10 -->
+For complete API documentation of the stable operators, see the [Torch Stable API cpp documentation](https://docs.pytorch.org/cppdocs/api/stable/index.html).
 
 ### Stable C headers
 
diff --git a/docs/source/notes/mps.rst b/docs/source/notes/mps.rst
index cc2a781b590ec..54619516bed49 100644
--- a/docs/source/notes/mps.rst
+++ b/docs/source/notes/mps.rst
@@ -4,7 +4,7 @@ MPS backend
 ===========
 
 :mod:`mps` device enables high-performance
-training on GPU for MacOS devices with Metal programming framework.  It
+training on GPU for macOS devices with Metal programming framework.  It
 introduces a new device to map Machine Learning computational graphs and
 primitives on highly efficient Metal Performance Shaders Graph framework and
 tuned kernels provided by Metal Performance Shaders framework respectively.
@@ -22,7 +22,7 @@ To get started, simply move your Tensor and Module to the ``mps`` device:
             print("MPS not available because the current PyTorch install was not "
                   "built with MPS enabled.")
         else:
-            print("MPS not available because the current MacOS version is not 12.3+ "
+            print("MPS not available because the current macOS version is not 14.0+ "
                   "and/or you do not have an MPS-enabled device on this machine.")
 
     else:
diff --git a/docs/source/optim.md b/docs/source/optim.md
index c389fa6c37d2b..ffbc172ede073 100644
--- a/docs/source/optim.md
+++ b/docs/source/optim.md
@@ -149,6 +149,18 @@ for input, target in dataset:
     Optimizer.zero_grad
 ```
 
+## Module-level hooks
+
+```{eval-rst}
+.. currentmodule:: torch.optim.optimizer
+
+.. autofunction:: register_optimizer_step_post_hook
+
+.. autofunction:: register_optimizer_step_pre_hook
+
+.. currentmodule:: torch.optim
+```
+
 ## Algorithms
 
 ```{eval-rst}
diff --git a/docs/source/package.md b/docs/source/package.md
index 1b50f743d5793..dcc699cf545f6 100644
--- a/docs/source/package.md
+++ b/docs/source/package.md
@@ -726,12 +726,23 @@ statements more clearly show whether they are referring to packaged code or not.
   :members:
 ```
 
-<!-- This module needs to be documented. Adding here in the meantime
-for tracking purposes -->
+## Analysis Utilities
 ```{eval-rst}
 .. py:module:: torch.package.analyze.find_first_use_of_broken_modules
 .. py:module:: torch.package.analyze.is_from_package
 .. py:module:: torch.package.analyze.trace_dependencies
+
+.. currentmodule:: torch.package.analyze.find_first_use_of_broken_modules
+
+.. autofunction:: find_first_use_of_broken_modules
+
+.. currentmodule:: torch.package.analyze.is_from_package
+
+.. autofunction:: is_from_package
+
+.. currentmodule:: torch.package.analyze.trace_dependencies
+
+.. autofunction:: trace_dependencies
 .. py:module:: torch.package.file_structure_representation
 .. py:module:: torch.package.find_file_dependencies
 .. py:module:: torch.package.glob_group
diff --git a/docs/source/profiler.md b/docs/source/profiler.md
index 1578b7334d849..45789b95dea71 100644
--- a/docs/source/profiler.md
+++ b/docs/source/profiler.md
@@ -26,6 +26,8 @@
 .. autofunction:: torch.profiler.schedule
 
 .. autofunction:: torch.profiler.tensorboard_trace_handler
+
+.. autofunction:: torch.profiler.supported_activities
 ```
 
 ## Intel Instrumentation and Tracing Technology APIs
@@ -38,6 +40,8 @@
 .. autofunction:: torch.profiler.itt.range_push
 
 .. autofunction:: torch.profiler.itt.range_pop
+
+.. autofunction:: torch.profiler.itt.range
 ```
 
 <!-- This module needs to be documented. Adding here in the meantime
diff --git a/docs/source/quantization-support.md b/docs/source/quantization-support.md
index ef72d6fae15e6..054c47758c3fe 100644
--- a/docs/source/quantization-support.md
+++ b/docs/source/quantization-support.md
@@ -367,6 +367,25 @@ to configure quantization settings for individual ops.
     default_qat_qconfig_v2
 ```
 
+## torch.ao.quantization.quantization_mappings
+
+```{eval-rst}
+.. automodule:: torch.ao.quantization.quantization_mappings
+```
+
+```{eval-rst}
+.. currentmodule:: torch.ao.quantization.quantization_mappings
+```
+
+```{eval-rst}
+.. autosummary::
+    :toctree: generated
+    :nosignatures:
+    :template: classtemplate.rst
+
+    no_observer_set
+```
+
 ## torch.ao.nn.intrinsic
 
 ```{eval-rst}
@@ -774,6 +793,10 @@ the `custom operator mechanism <https://pytorch.org/tutorials/advanced/torch_scr
 .. automodule:: torch.nn.intrinsic.quantized.dynamic.modules
 .. automodule:: torch.nn.quantized.dynamic.modules
 .. automodule:: torch.quantization
+
+.. currentmodule:: torch.quantization
+.. autofunction:: default_eval_fn
+
 .. automodule:: torch.nn.intrinsic.modules
 ```
 
diff --git a/docs/source/quantization.rst b/docs/source/quantization.rst
index 191cb9cd363d0..ef2d95aaa4044 100644
--- a/docs/source/quantization.rst
+++ b/docs/source/quantization.rst
@@ -134,8 +134,6 @@ and supported quantized modules and functions.
 .. py:module:: torch.ao.quantization.qconfig
 .. py:module:: torch.ao.quantization.qconfig_mapping
 .. py:module:: torch.ao.quantization.quant_type
-.. py:module:: torch.ao.quantization.quantization_mappings
-.. py:module:: torch.ao.quantization.quantize_fx
 .. py:module:: torch.ao.quantization.quantize_jit
 .. py:module:: torch.ao.quantization.stubs
 .. py:module:: torch.nn.intrinsic.modules.fused
diff --git a/docs/source/redirects.py b/docs/source/redirects.py
index 24546d4f42c46..7dc1392028e8a 100644
--- a/docs/source/redirects.py
+++ b/docs/source/redirects.py
@@ -10,6 +10,9 @@
     "torch.compiler": "user_guide/torch_compiler/torch.compiler.html",
     "torch.compiler.config": "user_guide/torch_compiler/torch.compiler.config.html",
     "torch.compiler_aot_inductor": "user_guide/torch_compiler/torch.compiler_aot_inductor.html",
+    "torch.compiler_aot_compile": (
+        "user_guide/torch_compiler/torch.compiler_aot_compile.html"
+    ),
     "torch.compiler_aot_inductor_debugging_guide": (
         "user_guide/torch_compiler/torch.compiler_aot_inductor_debugging_guide.html"
     ),
diff --git a/docs/source/rpc.md b/docs/source/rpc.md
index 38e9354f70d9f..8e820e33ad80a 100644
--- a/docs/source/rpc.md
+++ b/docs/source/rpc.md
@@ -89,6 +89,7 @@ framework, RRef framework and distributed autograd.
 ```{eval-rst}
 .. automodule:: torch.distributed.rpc
 .. autofunction:: init_rpc
+.. autofunction:: is_available
 ```
 
 The following APIs allow users to remotely execute functions as well as create
@@ -258,6 +259,8 @@ using RPC. For more details see {ref}`distributed-autograd-design`.
 ```{eval-rst}
 .. automodule:: torch.distributed.autograd
     :members: context, backward, get_gradients
+
+.. autofunction:: torch.distributed.autograd.is_available
 ```
 
 ```{toctree}
diff --git a/docs/source/sparse.rst b/docs/source/sparse.rst
index 652667b5ea798..a6ba9d3cd8db5 100644
--- a/docs/source/sparse.rst
+++ b/docs/source/sparse.rst
@@ -309,6 +309,8 @@ You can accelerate the linear layers in your model if the weights are already se
     >>> linear = nn.Linear(64, 64).half().cuda()
     >>> linear.weight = nn.Parameter(to_sparse_semi_structured(linear.weight.masked_fill(~mask, 0)))
 
+.. autofunction:: torch.sparse.semi_structured.to_sparse_semi_structured
+
 
 .. _sparse-coo-docs:
 
diff --git a/docs/source/torch.rst b/docs/source/torch.rst
index 09454c036eb5d..21fcb8958e38f 100644
--- a/docs/source/torch.rst
+++ b/docs/source/torch.rst
@@ -257,6 +257,14 @@ Serialization
     save
     load
 
+.. currentmodule:: torch.serialization
+
+.. autofunction:: check_module_version_greater_or_equal
+
+.. autofunction:: default_restore_location
+
+.. currentmodule:: torch
+
 Parallelism
 ----------------------------------
 .. autosummary::
@@ -320,8 +328,10 @@ Constants
 ~~~~~~~~~~~~~~~~~~~~~~
 
 ======================================= ===========================================
+``e``                                       Euler's number, the base of natural logarithms (~2.7183). Alias for :attr:`math.e`.
 ``inf``                                     A floating-point positive infinity. Alias for :attr:`math.inf`.
 ``nan``                                     A floating-point "not a number" value. This value is not a legal number. Alias for :attr:`math.nan`.
+``pi``                                      The ratio of a circle's circumference to its diameter (~3.1416). Alias for :attr:`math.pi`.
 ======================================= ===========================================
 
 Pointwise Ops
@@ -663,6 +673,7 @@ Foreach Operations
     _foreach_atan_
     _foreach_ceil
     _foreach_ceil_
+    _foreach_clone
     _foreach_cos
     _foreach_cos_
     _foreach_cosh
@@ -731,6 +742,7 @@ Utilities
     is_warn_always_enabled
     vmap
     _assert
+    typename
 
 Symbolic Numbers
 ----------------
@@ -754,6 +766,7 @@ Symbolic Numbers
     sym_min
     sym_not
     sym_ite
+    sym_sqrt
     sym_sum
 
 Export Path
@@ -805,6 +818,9 @@ Operator Tags
 .. This module needs to be documented. Adding here in the meantime
 .. for tracking purposes
 .. py:module:: torch.utils.model_dump
+
+.. currentmodule:: torch.utils.model_dump
+
 .. py:module:: torch.utils.viz
 .. py:module:: torch.quasirandom
 .. py:module:: torch.return_types
diff --git a/docs/source/user_guide/torch_compiler/advanced.md b/docs/source/user_guide/torch_compiler/advanced.md
index acfa3cd60a462..4f4a3e65c22e8 100644
--- a/docs/source/user_guide/torch_compiler/advanced.md
+++ b/docs/source/user_guide/torch_compiler/advanced.md
@@ -10,4 +10,5 @@ torch.compiler_transformations.md
 torch.compiler_fake_tensor.md
 torch.compiler_custom_backends.md
 torch.compiler_dynamic_shapes
+torch.compiler_aot_compile
 ```
diff --git a/docs/source/user_guide/torch_compiler/compile/dynamic_shapes_troubleshooting_guardon_errors.md b/docs/source/user_guide/torch_compiler/compile/dynamic_shapes_troubleshooting_guardon_errors.md
index aff1c5baadd64..96e9db701ed36 100644
--- a/docs/source/user_guide/torch_compiler/compile/dynamic_shapes_troubleshooting_guardon_errors.md
+++ b/docs/source/user_guide/torch_compiler/compile/dynamic_shapes_troubleshooting_guardon_errors.md
@@ -185,14 +185,13 @@ else:
     torch._check(x != 0)  # runtime check for the general path
     return compute(x)
 ```
+.
 
-**`hint_int(expr, fallback=None)`**: Extracts a hint from a symbolic size and uses it in a branch. You can use it to evaluate the expression using the traced program's input shapes without guarding. Unlike `statically_known_true`, it picks the path that works for the input example instead of returning `False`. The optional `fallback` argument substitutes unbacked symbols; if not provided and the symbol is unbacked, it will raise an error.
+**`optimization_hint(expr, fallback=None)`**: Evaluates a symbolic expression to a concrete integer for **optimization decisions only** (e.g., selecting a faster kernel). Unlike `guarding_hint_or_throw`, it handles unbacked symbols using the `fallback` value. Both branches must still be correct for all dynamic shapes — only performance should depend on the hint.
 
 ```python
-from torch.fx.experimental.symbolic_shapes import hint_int
-
 # Use ONLY for optimizations, not correctness-critical branches
-if hint_int(x.numel(), fallback=0) > 1024:
+if optimization_hint(x.numel(), fallback=0) > 1024:
     # optimized path for large tensors
     ...
 else:
diff --git a/docs/source/user_guide/torch_compiler/compile/programming_model.common_graph_breaks.md b/docs/source/user_guide/torch_compiler/compile/programming_model.common_graph_breaks.md
index 7ae42de87e68d..f772be9e821f4 100644
--- a/docs/source/user_guide/torch_compiler/compile/programming_model.common_graph_breaks.md
+++ b/docs/source/user_guide/torch_compiler/compile/programming_model.common_graph_breaks.md
@@ -41,6 +41,23 @@ Dynamo makes a best-effort attempt to hint if a graph break is caused by your co
 But it can still sometimes be difficult to tell from the logs if the graph break is caused by an error in your code,
 is a more complicated graph break, or is a `torch.compile` bug. In order to differentiate, we recommend trying to run your code without `torch.compile` to see if you still get the error reported by the graph break.
 
+You can also use `torch.compiler.set_stance("force_eager")` to quickly disable `torch.compile` without needing to modify the `torch.compile` call:
+
+```{code-cell}
+@torch.compile
+def fn(x):
+    y = torch.sin(x, x)
+    return y
+
+try:
+    with torch.compiler.set_stance("force_eager"):
+        fn(torch.ones(3, 3))
+except Exception as e:
+    print(e)
+```
+
+See https://docs.pytorch.org/tutorials/recipes/torch_compiler_set_stance_tutorial.html#crashing-sooner for more examples of `set_stance` usage for debugging.
+
 ## Data-dependent operations
 
 `torch.compile` graph breaks on data-dependent operations such as data-dependent control flow (if-statements, loops with tensors) and direct tensor data accesses (`.item`, `.data_ptr`).
diff --git a/docs/source/user_guide/torch_compiler/compile/programming_model.fullgraph_true.md b/docs/source/user_guide/torch_compiler/compile/programming_model.fullgraph_true.md
index 5c38a3d88c2f9..6fc3749e6d99e 100644
--- a/docs/source/user_guide/torch_compiler/compile/programming_model.fullgraph_true.md
+++ b/docs/source/user_guide/torch_compiler/compile/programming_model.fullgraph_true.md
@@ -3,7 +3,7 @@ file_format: mystnb
 kernelspec:
   name: python3
 mystnb:
-  execution_timeout: 30
+  execution_timeout: 120
   execution_show_tb: True
   merge_streams: True
 ---
diff --git a/docs/source/user_guide/torch_compiler/torch.compiler_aot_compile.md b/docs/source/user_guide/torch_compiler/torch.compiler_aot_compile.md
new file mode 100644
index 0000000000000..e4fbd730ea674
--- /dev/null
+++ b/docs/source/user_guide/torch_compiler/torch.compiler_aot_compile.md
@@ -0,0 +1,229 @@
+(torch.compiler_aot_compile)=
+
+# Ahead-of-Time Compilation with torch.compile
+
+:::{warning}
+This feature is experimental and subject to change.
+:::
+
+`torch.compile().aot_compile()` performs ahead-of-time (AOT) compilation on a
+`torch.compile`-d function or module. Unlike the standard `torch.compile` flow
+-- which compiles lazily on first invocation -- `aot_compile()` eagerly traces,
+compiles, and packages the result into a serializable artifact.
+
+The entire compilation pipeline runs ahead of time, including graph tracing,
+Inductor code generation, Triton kernel compilation, and autotuning. The
+compiled artifact can be saved to disk and loaded later, skipping all of these
+steps at runtime.
+
+This is useful when you want to:
+
+- **Eliminate cold-start compilation latency** in production.
+- **Serialize compiled artifacts** for deployment to other processes or machines.
+- **Cross-compile** on a host machine for a different target device by tracing
+  with fake tensors.
+
+## How it differs from AOTInductor
+
+{ref}`AOTInductor <torch.compiler_aot_inductor>` operates on
+{func}`torch.export`-ed models and produces shared libraries for deployment in
+non-Python environments. `aot_compile()` operates directly on `torch.compile`
+and stays within the Python runtime -- the saved artifact is loaded back as a
+Python callable. Choose AOTInductor when you need C++ deployment; choose
+`aot_compile()` when you want to stay in Python and precompute compilation.
+
+## Quick start
+
+### Compiling a free function
+
+```python
+import torch
+
+def fn(x, y):
+    return x + y
+
+# Step 1: AOT compile with example inputs.
+#   fullgraph=True is required (graph breaks are not supported).
+#   example_inputs is a tuple of (args_tuple, kwargs_dict).
+compiled_fn = torch.compile(fn, fullgraph=True).aot_compile(
+    ((torch.randn(3, 4), torch.randn(3, 4)), {})
+)
+
+# Step 2: Run the compiled function.
+result = compiled_fn(torch.randn(3, 4), torch.randn(3, 4))
+
+# Step 3: Save the compiled artifact to disk.
+compiled_fn.save_compiled_function("compiled_add.pt")
+
+# Step 4: Load and run in another process (no recompilation).
+with open("compiled_add.pt", "rb") as f:
+    loaded_fn = torch.compiler.load_compiled_function(f)
+
+result = loaded_fn(torch.randn(3, 4), torch.randn(3, 4))
+```
+
+### Compiling a module
+
+When compiling an `nn.Module`, call `aot_compile()` on the `.forward` method
+of the compiled module. The compiled function expects the module instance as
+its first argument (the `self` parameter).
+
+```python
+import torch
+import torch.nn as nn
+
+class MyModel(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.linear = nn.Linear(4, 4)
+
+    def forward(self, x):
+        return self.linear(x)
+
+model = MyModel()
+
+# AOT compile the forward method.
+compiled_fn = torch.compile(
+    model, fullgraph=True
+).forward.aot_compile(((torch.randn(3, 4, requires_grad=True),), {}))
+
+# Call with the module instance as the first argument.
+inputs = torch.randn(3, 4, requires_grad=True)
+result = compiled_fn(model, inputs)
+
+# Backward works through the compiled function.
+loss = result.sum()
+loss.backward()
+print(inputs.grad)  # gradients flow back correctly
+
+# Save and load.
+compiled_fn.save_compiled_function("compiled_model.pt")
+
+with open("compiled_model.pt", "rb") as f:
+    loaded_fn = torch.compiler.load_compiled_function(f)
+
+# Backward also works after loading from disk.
+inputs = torch.randn(3, 4, requires_grad=True)
+result = loaded_fn(model, inputs)
+result.sum().backward()
+print(inputs.grad)
+```
+
+## API reference
+
+### `torch.compile(...).aot_compile(example_inputs)`
+
+Ahead-of-time compiles the `torch.compile()`-wrapped function.
+
+**Args:**
+
+- **example_inputs** (`tuple[tuple[Any, ...], dict[str, Any]]`) -- A tuple of
+  `(args, kwargs)` providing example inputs for tracing. These determine the
+  tensor shapes, dtypes, and devices that the compiled artifact is valid for.
+
+**Returns:** An `AOTCompiledFunction` -- a callable that behaves like the
+original function but runs the pre-compiled code. It also exposes:
+
+- `save_compiled_function(path)` -- Serialize the compiled artifact to disk.
+- `disable_guard_check()` -- Disable runtime guard validation (advanced use).
+
+**Requirements:**
+
+- `fullgraph=True` must be passed to `torch.compile()`. Graph breaks are not
+  supported with AOT compilation.
+- The backend must be callable (string backends like `"inductor"`,
+  `"eager"` and `"aot_eager"` are supported).
+
+### `torch.compiler.load_compiled_function(file, *, f_globals=None, external_data=None)`
+
+Load a previously saved AOT-compiled function from a file.
+
+**Args:**
+
+- **file** -- A file-like object (opened in binary read mode) containing the
+  serialized compiled function.
+- **f_globals** (`dict | None`) -- Optional global scope for the compiled
+  function. Required when the original function references user-defined types
+  or other non-standard globals.
+- **external_data** (`dict | None`) -- Optional data to be loaded into the
+  runtime environment. Required when the original function captures objects
+  that could not be serialized (e.g., `nn.Module` instances). The keys should
+  match those passed to `save_compiled_function(external_data=...)`.
+
+**Returns:** A callable with compilation preloaded from disk.
+
+## Choosing a backend
+
+`aot_compile()` works with any backend that implements the
+`SerializableCallable` interface. The built-in `"inductor"`, `"eager"` and `"aot_eager"`
+backends are supported out of the box:
+
+```python
+# With inductor (default, optimized code generation).
+compiled_fn = torch.compile(fn, fullgraph=True, backend="inductor").aot_compile(
+    ((torch.randn(3, 4),), {})
+)
+
+# With eager (useful for debugging, no codegen).
+compiled_fn = torch.compile(fn, fullgraph=True, backend="eager").aot_compile(
+    ((torch.randn(3, 4),), {})
+)
+```
+
+## Handling closures and external references
+
+Functions that capture free variables (closures) are supported. The closure
+state is serialized along with the compiled artifact:
+
+```python
+scale = 2
+
+def fn(x, y):
+    return (x + y) * scale
+
+compiled_fn = torch.compile(fn, fullgraph=True).aot_compile(
+    ((torch.randn(3, 4), torch.randn(3, 4)), {})
+)
+compiled_fn.save_compiled_function("scaled_add.pt")
+
+with open("scaled_add.pt", "rb") as f:
+    loaded_fn = torch.compiler.load_compiled_function(f)
+```
+
+When the function references user-defined types that cannot be found by the
+deserializer, pass `f_globals` to provide the necessary namespace:
+
+```python
+with open("my_fn.pt", "rb") as f:
+    loaded_fn = torch.compiler.load_compiled_function(
+        f, f_globals=my_module.__dict__
+    )
+```
+
+When the function captures non-serializable objects (like `nn.Module`
+instances), use `external_data`:
+
+```python
+# Saving.
+compiled_fn.save_compiled_function(
+    "fn_with_model.pt",
+    external_data={"model": model},
+)
+
+# Loading.
+with open("fn_with_model.pt", "rb") as f:
+    loaded_fn = torch.compiler.load_compiled_function(
+        f, external_data={"model": model}
+    )
+```
+
+## Limitations
+
+- **`fullgraph=True` is required.** If `torch.compile` encounters a graph
+  break, `aot_compile()` raises an error.
+- **Input shapes are specialized.** The compiled artifact is valid for the
+  tensor shapes, dtypes, and devices provided as example inputs. Inputs with
+  different shapes will trigger a guard failure at runtime unless guards are
+  explicitly disabled.
+- **Not all backends are supported.** Custom backends must implement the
+  `SerializableCallable` interface to be compatible with save/load.
diff --git a/docs/source/user_guide/torch_compiler/torch.compiler_cudagraph_trees.md b/docs/source/user_guide/torch_compiler/torch.compiler_cudagraph_trees.md
index f220086f82dc2..ea6581ae1533b 100644
--- a/docs/source/user_guide/torch_compiler/torch.compiler_cudagraph_trees.md
+++ b/docs/source/user_guide/torch_compiler/torch.compiler_cudagraph_trees.md
@@ -10,7 +10,7 @@ For a longer background on CUDAGraphs, read [accelerating pytorch with CUDAGraph
 
 CUDA Graphs can give large speedups, especially for models with high CPU overhead or small compute. There are a number of limitations from requiring the same kernels to be run with the same arguments and dependencies, and memory addresses.
 
-- Control Flow is not possible
+- Arbitrary Control Flow is not possible (However, control flow expressed via `torch.cond()` can be captured in a CUDA Graph. See {ref}`Data Dependent Control Flow <graph-data-dependent-control-flow>`.)
 - Kernels which trigger host to device syncs (such as .item()) errors
 - All input arguments to kernels are fixed to what they were recorded
 - CUDA Memory addresses are fixed, however the values of the memory at those addresses can change
diff --git a/docs/source/user_guide/torch_compiler/torch.compiler_dynamo_overview.md b/docs/source/user_guide/torch_compiler/torch.compiler_dynamo_overview.md
index 7ba68ad0c42f9..9a0c65135f02d 100644
--- a/docs/source/user_guide/torch_compiler/torch.compiler_dynamo_overview.md
+++ b/docs/source/user_guide/torch_compiler/torch.compiler_dynamo_overview.md
@@ -188,7 +188,7 @@ def toy_example(a, b):
     if __temp_1[1]:
         return __resume_at_30_1(b, x)
     return __resume_at_38_2(b, x)
-If you find the decompiled code is wrong,please submit an issue at https://github.com/youkaichao/depyf/issues.
+If you find the decompiled code is wrong, please submit an issue at https://github.com/youkaichao/depyf/issues.
 ```
 
 At the top you can see the FX graph.
diff --git a/docs/source/utils.md b/docs/source/utils.md
index 86b62f0fdc5d7..9170e9e3f2c09 100644
--- a/docs/source/utils.md
+++ b/docs/source/utils.md
@@ -35,7 +35,12 @@
 
     check_release_file
     is_xnnpack_available
+    main
     pretty_str
+    run
+    run_and_parse_first_match
+    run_and_read_all
+    run_and_return_first_line
 ```
 
 # torch.utils.flop_counter
@@ -94,6 +99,8 @@
     is_out_of_place
     is_pytorch_file
     is_special_file
+    match_extensions
+    matched_files_iter
     openf
     preprocess_file_and_save_result
     preprocessor
@@ -112,7 +119,6 @@ for tracking purposes -->
 .. py:module:: torch.utils.benchmark.examples.fuzzer
 .. py:module:: torch.utils.benchmark.examples.op_benchmark
 .. py:module:: torch.utils.benchmark.examples.simple_timeit
-.. py:module:: torch.utils.benchmark.examples.spectral_ops_fuzz_test
 .. py:module:: torch.utils.benchmark.op_fuzzers.binary
 .. py:module:: torch.utils.benchmark.op_fuzzers.sparse_binary
 .. py:module:: torch.utils.benchmark.op_fuzzers.sparse_unary
@@ -168,7 +174,6 @@ for tracking purposes -->
 .. py:module:: torch.utils.hipify.cuda_to_hip_mappings
 .. py:module:: torch.utils.hipify.version
 .. py:module:: torch.utils.hooks
-.. py:module:: torch.utils.jit.log_extract
 .. py:module:: torch.utils.mkldnn
 .. py:module:: torch.utils.mobile_optimizer
 .. py:module:: torch.utils.show_pickle
diff --git a/functorch/dim/__init__.py b/functorch/dim/__init__.py
index b9decc1cd9c8f..240c7c8c0071f 100644
--- a/functorch/dim/__init__.py
+++ b/functorch/dim/__init__.py
@@ -25,7 +25,7 @@
 _n_dims_created = 0
 
 
-def _relevant_op(opcode: Optional[str]) -> bool:
+def _relevant_op(opcode: str | None) -> bool:
     """Check if opcode is relevant for variable assignment."""
     return bool(opcode and opcode.startswith("STORE_"))
 
@@ -35,14 +35,14 @@ def handle_from_tensor(tensor: torch.Tensor) -> torch.Tensor:
     return tensor
 
 
-def _create_dim(name: str, size: Optional[int] = None) -> Dim:
+def _create_dim(name: str, size: int | None = None) -> Dim:
     """Create a new Dim object."""
     return Dim(name, size if size is not None else -1)
 
 
 def dims(
-    n: Optional[int] = None, sizes: Optional[list[Optional[int]]] = None
-) -> Union[Dim, tuple[Dim, ...]]:
+    n: int | None = None, sizes: list[int | None] | None = None
+) -> Dim | tuple[Dim, ...]:
     """
     Create and return one or more Dim objects.
 
@@ -148,14 +148,14 @@ class DimList:
     2. Bound: Either created with specific dimensions/sizes, or bound later via bind() or bind_len()
     """
 
-    _name: Optional[str]
+    _name: str | None
     _dims: list[Dim]
     _bound: bool
 
     def __init__(
         self,
-        len_or_dims: Optional[Union[int, Sequence]] = None,
-        name: Optional[str] = None,
+        len_or_dims: int | Sequence | None = None,
+        name: str | None = None,
     ):
         """
         Initialize a new DimList object.
@@ -249,7 +249,7 @@ def __len__(self) -> int:
         """Return the length of the DimList."""
         return self.size()
 
-    def __getitem__(self, key: Union[int, slice]) -> Union[Dim, tuple[Dim, ...]]:
+    def __getitem__(self, key: int | slice) -> Dim | tuple[Dim, ...]:
         if not self._bound:
             raise DimensionBindError("DimList not bound")
 
@@ -288,14 +288,12 @@ def __torch_function__(
         func: Callable,
         types: tuple,
         args: tuple = (),
-        kwargs: Optional[dict] = None,
+        kwargs: dict | None = None,
     ) -> Any:
         return _Tensor.__torch_function__(func, types, args, kwargs)
 
 
-def _create_dimlist(
-    name: str, size: Optional[Union[int, list[Optional[int]]]] = None
-) -> DimList:
+def _create_dimlist(name: str, size: int | list[int | None] | None = None) -> DimList:
     """Create a DimList object with the given name and optional size."""
     dimlist = DimList(name=name)
     if size is not None:
@@ -311,8 +309,8 @@ def _create_dimlist(
 
 
 def dimlists(
-    n: Optional[int] = None, sizes: Optional[list[Optional[int]]] = None
-) -> Union[DimList, tuple[DimList, ...]]:
+    n: int | None = None, sizes: list[int | None] | None = None
+) -> DimList | tuple[DimList, ...]:
     """
     Create and return one or more DimList objects.
 
@@ -429,7 +427,7 @@ class _Tensor:
     def _get_levels(self) -> list[Any]:
         raise NotImplementedError("_get_levels must be implemented by subclass")
 
-    def _get_tensor(self) -> Optional[torch.Tensor]:
+    def _get_tensor(self) -> torch.Tensor | None:
         raise NotImplementedError("_get_tensor must be implemented by subclass")
 
     @property
@@ -449,7 +447,7 @@ def __torch_function__(
         func: Callable,
         types: tuple,
         args: tuple = (),
-        kwargs: Optional[dict] = None,
+        kwargs: dict | None = None,
     ) -> Any:
         if kwargs is None:
             kwargs = {}
@@ -714,14 +712,12 @@ def expand(self, *args: Dim) -> _Tensor:
 
     def index(
         self,
-        dims: Union[int, Dim, tuple[Union[int, Dim], ...], list[Union[int, Dim]]],
-        indices: Union[
-            int,
-            slice,
-            torch.Tensor,
-            tuple[Union[int, slice, torch.Tensor], ...],
-            list[Union[int, slice, torch.Tensor]],
-        ],
+        dims: int | Dim | tuple[int | Dim, ...] | list[int | Dim],
+        indices: int
+        | slice
+        | torch.Tensor
+        | tuple[int | slice | torch.Tensor, ...]
+        | list[int | slice | torch.Tensor],
     ) -> _Tensor:
         """
         Index tensor using first-class dimensions.
@@ -751,8 +747,8 @@ def dim_not_present(d: Any) -> None:
             else:
                 raise TypeError(f"dimension {repr(d.dim())} not in tensor")
 
-        dims_list: list[Union[int, Dim]] = []
-        indices_list: list[Union[int, slice, torch.Tensor]] = []
+        dims_list: list[int | Dim] = []
+        indices_list: list[int | slice | torch.Tensor] = []
 
         lhs_list = isinstance(dims, (tuple, list))
         rhs_list = isinstance(indices, (tuple, list))
@@ -890,8 +886,8 @@ class Dim(_Tensor):
     _level: int
     _name: str
     _size: int
-    _range: Optional[torch.Tensor]
-    _batchtensor: Optional[torch.Tensor]
+    _range: torch.Tensor | None
+    _batchtensor: torch.Tensor | None
 
     def __init__(self, name: str, s: int = -1) -> None:
         global _n_dims_created
@@ -968,13 +964,13 @@ def __repr__(self) -> str:
 # is somewhat intentional, as FCD tensors are intended to be substitutable
 # with regular Tensor (just with some positional dims hidden).
 class Tensor(_Tensor):
-    _tensor: Optional[torch.Tensor]
-    _batchtensor: Optional[torch.Tensor]
+    _tensor: torch.Tensor | None
+    _batchtensor: torch.Tensor | None
     _levels: list[DimEntry]
     _has_device: bool
-    _delayed: Optional[Callable[[], torch.Tensor]]
-    _delayed_orig: Optional[Callable]
-    _delayed_args: Optional[tuple]
+    _delayed: Callable[[], torch.Tensor] | None
+    _delayed_orig: Callable | None
+    _delayed_args: tuple | None
 
     @property
     def ndim(self) -> int:
@@ -987,7 +983,7 @@ def check_exact(cls, other: Any) -> bool:
     @classmethod
     def from_positional(
         cls, tensor: torch.Tensor, levels: list[DimEntry], has_device: bool
-    ) -> Union[_Tensor, torch.Tensor]:
+    ) -> _Tensor | torch.Tensor:
         """
         Create a functorch Tensor from a regular PyTorch tensor with specified dimension levels.
 
@@ -1073,7 +1069,7 @@ def evaluate_delayed() -> torch.Tensor:
 
         return result
 
-    def _get_tensor(self) -> Optional[torch.Tensor]:
+    def _get_tensor(self) -> torch.Tensor | None:
         """Get the underlying tensor, handling delayed operations if needed."""
         if (
             hasattr(self, "_delayed")
@@ -1096,7 +1092,7 @@ def _get_has_device(self) -> bool:
         """Get whether this tensor has device information."""
         return self._has_device
 
-    def _get_batchtensor(self) -> Optional[torch.Tensor]:
+    def _get_batchtensor(self) -> torch.Tensor | None:
         """Get the batched tensor representation, creating it lazily if needed."""
         if self._batchtensor is None:
             self._batchtensor = self._add_batch_dims(
@@ -1105,8 +1101,8 @@ def _get_batchtensor(self) -> Optional[torch.Tensor]:
         return self._batchtensor
 
     def _add_batch_dims(
-        self, t: Optional[torch.Tensor], levels_: list[Any]
-    ) -> Optional[torch.Tensor]:
+        self, t: torch.Tensor | None, levels_: list[Any]
+    ) -> torch.Tensor | None:
         levels = list(levels_)
 
         while True:
@@ -1407,7 +1403,7 @@ def dot_finish(parts: list[DotPart], result_tensor: torch.Tensor) -> Tensor:
     return tensor_result  # type: ignore[return-value]
 
 
-def dot(lhs: Any, rhs: Any, sum_dims: Any) -> Union[_Tensor, torch.Tensor]:
+def dot(lhs: Any, rhs: Any, sum_dims: Any) -> _Tensor | torch.Tensor:
     """
     Perform dot product between two tensors along specified dimensions.
 
diff --git a/functorch/dim/_dim_entry.py b/functorch/dim/_dim_entry.py
index f7571793a7fdd..80f0b20f70587 100644
--- a/functorch/dim/_dim_entry.py
+++ b/functorch/dim/_dim_entry.py
@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Union
+from typing import TYPE_CHECKING
 
 
 if TYPE_CHECKING:
@@ -15,9 +15,9 @@
 # follow this convention even though it shouldn't be necessary now
 class DimEntry:
     # The dimension this is from the rhs, or a FCD
-    data: Union[Dim, int]
+    data: Dim | int
 
-    def __init__(self, data: Union[Dim, int, None] = None) -> None:
+    def __init__(self, data: Dim | int | None = None) -> None:
         from . import Dim
 
         if type(data) is int:
diff --git a/functorch/dim/_getsetitem.py b/functorch/dim/_getsetitem.py
index af335dd8f35c4..0b9a19efee0f9 100644
--- a/functorch/dim/_getsetitem.py
+++ b/functorch/dim/_getsetitem.py
@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 from dataclasses import dataclass, field
-from typing import Any, Optional, TYPE_CHECKING, Union
+from typing import Any, TYPE_CHECKING
 
 import torch
 
@@ -13,7 +13,7 @@
     from . import Dim
 
 
-def _safe_index(lst: list, item: Any) -> Optional[int]:
+def _safe_index(lst: list, item: Any) -> int | None:
     """
     Helper function to find index of item in list.
 
@@ -36,7 +36,7 @@ def _safe_index(lst: list, item: Any) -> Optional[int]:
 class IndexingInfo:
     can_call_original: bool = False
     advanced_indexing: bool = False
-    self_tensor: Optional[torch.Tensor] = None
+    self_tensor: torch.Tensor | None = None
     flat_inputs: list[Any] = field(default_factory=list)
     result_levels: list[DimEntry] = field(default_factory=list)
     has_device: bool = False
@@ -192,7 +192,7 @@ def setitem(self: Any, index: Any, rhs: Any) -> None:
 
                 if not found:
                     # Create tuple representation of result levels for error message
-                    result_dims: list[Union[int, Dim]] = []
+                    result_dims: list[int | Dim] = []
                     for rl in iinfo.result_levels:
                         if rl.is_positional():
                             result_dims.append(rl.position())
diff --git a/functorch/dim/_order.py b/functorch/dim/_order.py
index 5c952f6a80747..42db77572314c 100644
--- a/functorch/dim/_order.py
+++ b/functorch/dim/_order.py
@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import Any, TYPE_CHECKING, Union
+from typing import Any, TYPE_CHECKING
 
 
 if TYPE_CHECKING:
@@ -40,7 +40,7 @@ def _wrap_dim(arg: Any, orig_ndim: int, allow_none: bool = True) -> DimEntry:
 
 
 def order(
-    tensor_or_dim: Union[torch.Tensor, Any], *dims: Union[Any, Sequence[Any]]
+    tensor_or_dim: torch.Tensor | Any, *dims: Any | Sequence[Any]
 ) -> torch.Tensor:
     """
     Reorder the dimensions of a tensor or create a tensor from a dimension.
diff --git a/functorch/dim/_py_inst_decoder.py b/functorch/dim/_py_inst_decoder.py
index 7f08ebb8557fb..2425a94813595 100644
--- a/functorch/dim/_py_inst_decoder.py
+++ b/functorch/dim/_py_inst_decoder.py
@@ -1,5 +1,5 @@
 import dis
-from typing import Any, Optional
+from typing import Any
 
 
 class _PyInstDecoder:
@@ -28,7 +28,7 @@ def next(self) -> None:
         """Advance to the next instruction."""
         self.offset += 1
 
-    def opcode(self) -> Optional[str]:
+    def opcode(self) -> str | None:
         """Get the opcode name of the current instruction."""
         if self.offset < len(self.instructions):
             return self.instructions[self.offset].opname
@@ -40,7 +40,7 @@ def oparg(self) -> int:
             return self.instructions[self.offset].arg or 0
         return 0
 
-    def name(self) -> Optional[str]:
+    def name(self) -> str | None:
         """
         Extract variable name from current instruction.
         """
diff --git a/functorch/dim/_tensor_info.py b/functorch/dim/_tensor_info.py
index eeb453a0c18f2..c9d6d010891e2 100644
--- a/functorch/dim/_tensor_info.py
+++ b/functorch/dim/_tensor_info.py
@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
-from typing import Any, Optional, TYPE_CHECKING
+from typing import Any, TYPE_CHECKING
 
 import torch
 
@@ -12,10 +12,10 @@
 
 @dataclass
 class TensorInfo:
-    tensor: Optional[torch.Tensor]
+    tensor: torch.Tensor | None
     levels: list[DimEntry]
     has_device: bool
-    batchedtensor: Optional[torch.Tensor]
+    batchedtensor: torch.Tensor | None
 
     def __post_init__(self) -> None:
         from ._dim_entry import DimEntry
diff --git a/functorch/dim/_wrap.py b/functorch/dim/_wrap.py
index dd52e8c0d2792..9bc0db465324d 100644
--- a/functorch/dim/_wrap.py
+++ b/functorch/dim/_wrap.py
@@ -5,7 +5,7 @@
 from __future__ import annotations
 
 import functools
-from typing import Any, Optional, TYPE_CHECKING
+from typing import Any, TYPE_CHECKING
 
 import torch
 from torch.utils._pytree import tree_map
@@ -216,11 +216,11 @@ def wrap_result(obj: Any) -> Any:
 
 def _wrap(
     orig: Callable,
-    dim_offset: Optional[int] = None,
-    keepdim_offset: Optional[int] = None,
-    dim_name: Optional[str] = None,
-    single_dim: Optional[bool] = None,
-    reduce: Optional[bool] = None,
+    dim_offset: int | None = None,
+    keepdim_offset: int | None = None,
+    dim_name: str | None = None,
+    single_dim: bool | None = None,
+    reduce: bool | None = None,
 ) -> Callable:
     """
     Wrap a PyTorch function to support first-class dimensions.
@@ -254,7 +254,7 @@ def call_torch_function(
     func: Callable,
     types: tuple,
     args: tuple = (),
-    kwargs: Optional[dict] = None,
+    kwargs: dict | None = None,
 ) -> Any:
     """
     Handle __torch_function__ calls for wrapped operators.
diff --git a/functorch/einops/_parsing.py b/functorch/einops/_parsing.py
index 2f39724d0cefa..31dbf1c234f67 100644
--- a/functorch/einops/_parsing.py
+++ b/functorch/einops/_parsing.py
@@ -27,7 +27,7 @@
 
 import keyword
 import warnings
-from typing import Optional, TYPE_CHECKING, Union
+from typing import TYPE_CHECKING
 
 
 if TYPE_CHECKING:
@@ -72,12 +72,12 @@ def __init__(
             allow_duplicates (bool): whether to allow an identifier to appear more than once in the expression
         """
         self.has_ellipsis: bool = False
-        self.has_ellipsis_parenthesized: Optional[bool] = None
-        self.identifiers: set[Union[str, AnonymousAxis]] = set()
+        self.has_ellipsis_parenthesized: bool | None = None
+        self.identifiers: set[str | AnonymousAxis] = set()
         # that's axes like 2, 3, 4 or 5. Axes with size 1 are exceptional and replaced with empty composition
         self.has_non_unitary_anonymous_axes: bool = False
         # composition keeps structure of composite axes, see how different corner cases are handled in tests
-        self.composition: list[Union[list[Union[str, AnonymousAxis]], str]] = []
+        self.composition: list[list[str | AnonymousAxis] | str] = []
         if "." in expression:
             if "..." not in expression:
                 raise ValueError(
@@ -90,7 +90,7 @@ def __init__(
             expression = expression.replace("...", _ellipsis)
             self.has_ellipsis = True
 
-        bracket_group: Optional[list[Union[str, AnonymousAxis]]] = None
+        bracket_group: list[str | AnonymousAxis] | None = None
 
         def add_axis_name(x: str) -> None:
             if x in self.identifiers:
@@ -120,9 +120,7 @@ def add_axis_name(x: str) -> None:
                 )
                 if not (is_number or is_axis_name):
                     raise ValueError(f"Invalid axis identifier: {x}\n{reason}")
-                axis_name: Union[str, AnonymousAxis] = (
-                    AnonymousAxis(x) if is_number else x
-                )
+                axis_name: str | AnonymousAxis = AnonymousAxis(x) if is_number else x
                 self.identifiers.add(axis_name)
                 if is_number:
                     self.has_non_unitary_anonymous_axes = True
@@ -278,7 +276,7 @@ def validate_rearrange_expressions(
         )
 
 
-def comma_separate(collection: Collection[Union[str, Collection[str]]]) -> str:
+def comma_separate(collection: Collection[str | Collection[str]]) -> str:
     """Convert a collection of strings representing first class dims into a comma-separated string.
 
     Args:
diff --git a/functorch/einops/rearrange.py b/functorch/einops/rearrange.py
index 652c1f1ff028f..6622f8703fba1 100644
--- a/functorch/einops/rearrange.py
+++ b/functorch/einops/rearrange.py
@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 import functools
-from typing import TYPE_CHECKING, Union
+from typing import TYPE_CHECKING
 
 import torch
 from functorch.dim import dims  # noqa: F401
@@ -67,7 +67,7 @@ def _create_rearrange_callable(
         return lambda tensor: tensor
 
     first_class_dims: tuple[str, ...] = tuple(f"d{i}" for i in range(n_dims))
-    identifier_dim_map: dict[Union[str, AnonymousAxis], tuple[str, ...]] = {}
+    identifier_dim_map: dict[str | AnonymousAxis, tuple[str, ...]] = {}
     anon_axes: list[AnonymousAxis] = []
 
     # map the left-hand side identifiers to strings representing first class dims
@@ -97,11 +97,11 @@ def _create_rearrange_callable(
             raise ValueError(f"Unexpected dimension: {dimension}")
 
     def composition_to_dims(
-        composition: Sequence[Union[list[Union[str, AnonymousAxis]], str]],
-    ) -> list[Union[str, tuple[str, ...]]]:
+        composition: Sequence[list[str | AnonymousAxis] | str],
+    ) -> list[str | tuple[str, ...]]:
         """Convert a `ParsedExpression.composition` into a `Tensor.__getitem__` index of strings representing first
         class dims."""
-        dim_composition: list[Union[str, tuple[str, ...]]] = []
+        dim_composition: list[str | tuple[str, ...]] = []
         for dimension in composition:
             if isinstance(dimension, list):
                 dim_composition.append(
@@ -150,7 +150,7 @@ class dims."""
 
 
 def rearrange(
-    tensor: Union[torch.Tensor, list[torch.Tensor], tuple[torch.Tensor, ...]],
+    tensor: torch.Tensor | list[torch.Tensor] | tuple[torch.Tensor, ...],
     pattern: str,
     **axes_lengths: int,
 ) -> torch.Tensor:
diff --git a/pyproject.toml b/pyproject.toml
index 63fa54da38abf..73867a300270b 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -126,6 +126,10 @@ standard_library = ["typing_extensions"]
 [tool.ruff]
 line-length = 88
 src = ["caffe2", "torch", "torchgen", "functorch", "test"]
+extend-exclude = ["third_party", "test/dynamo/cpython"]
+
+[tool.ruff.per-file-target-version]
+"**/py312_intrinsics.py" = "py312"
 
 [tool.ruff.format]
 docstring-code-format = true
@@ -134,21 +138,29 @@ quote-style = "double"
 [tool.ruff.lint]
 # NOTE: Synchoronize the ignores with .flake8
 external = [
-    "B001",
-    "B902",
-    "B950",
-    "E121",
-    "E122",
-    "E128",
-    "E131",
-    "E704",
-    "E723",
-    "F723",
-    "F812",
+    # Codes from flake8-only plugins that ruff doesn't implement.
     "P201",
     "P204",
     "T484",
     "TOR901",
+    # Codes that are preview-only in ruff, so flake8 is still the enforcer.
+    # As ruff promotes these out of preview, move them to flake8's ignore list.
+    "B901",
+    "B909",
+    "E115",
+    "E201",
+    "E221",
+    "E225",
+    "E226",
+    "E227",
+    "E231",
+    "E241",
+    "E261",
+    "E262",
+    "E265",
+    "E266",
+    "E272",
+    "E306",
 ]
 ignore = [
     # these ignores are from flake8-bugbear; please fix!
@@ -181,8 +193,6 @@ ignore = [
     "SIM116", # Disable Use a dictionary instead of consecutive `if` statements
     "SIM117",
     "SIM300", # Yoda condition detected
-    "UP007", # keep-runtime-typing
-    "UP045", # keep-runtime-typing
     "TC006",
     # TODO: Remove Python-3.10 specific suppressions
     "B905",
@@ -265,10 +275,6 @@ select = [
     "S101",
 ]
 
-[tool.ruff.lint.pyupgrade]
-# Preserve types, even if a file imports `from __future__ import annotations`.
-keep-runtime-typing = true
-
 [tool.ruff.lint.per-file-ignores]
 "__init__.py" = [
     "F401",
@@ -283,7 +289,6 @@ keep-runtime-typing = true
 ]
 "test/export/**" = [
     "PGH004",
-    "S101",
 ]
 "test/typing/**" = [
     "PGH004"
@@ -308,14 +313,12 @@ keep-runtime-typing = true
     "PLR0133", # tests require this for JIT
     "PYI",
     "RUF015",
-    "S101",
     "UP", # We don't want to modify the jit test as they test specify syntax
 ]
 "test/test_jit.py" = [
     "PLR0133", # tests require this for JIT
     "PYI",
     "RUF015",
-    "S101",
     "UP", # We don't want to modify the jit test as they test specify syntax
 ]
 "test/inductor/s429861_repro.py" = [
@@ -345,18 +348,10 @@ keep-runtime-typing = true
 "tools/linter/**" = [
     "LOG015" # please fix
 ]
-# S101 migration: Remaining test folders/files still needing migration (ignore S101 for now)
-# Note: test/export/**, test/jit/**, test/torch_np/numpy_tests/** have S101 added to existing entries above
-"test/distributed/**" = ["S101"]
-"test/dynamo/**" = ["S101"]
-"test/functorch/**" = ["S101"]
-"test/inductor/**" = ["S101"]
-"test/onnx/**" = ["S101"]
-"test/quantization/**" = ["S101"]
+
 # torch/ folders still needing S101 migration
 "torch/_dynamo/**" = ["S101"]
 "torch/_inductor/**" = ["S101"]
-"torch/distributed/**" = ["S101"]
 
 [tool.codespell]
 ignore-words = "tools/linter/dictionary.txt"
@@ -366,6 +361,7 @@ package = 'torch'
 
 [tool.spin.commands]
 "Build" = [
+  ".spin/cmds.py:clean",
   ".spin/cmds.py:lint",
   ".spin/cmds.py:fixlint",
   ".spin/cmds.py:quicklint",
diff --git a/pyrefly.toml b/pyrefly.toml
index 4eb827ffe3891..a323a7787747a 100644
--- a/pyrefly.toml
+++ b/pyrefly.toml
@@ -41,6 +41,8 @@ project-excludes = [
   "torch/overrides.py",
   "*/__pycache__/**",
   "*/.*",
+  "torch/_inductor/kernel/vendored_templates/cutedsl/kernels/**",
+  "torch/_inductor/kernel/vendored_templates/cutedsl/dense_blockscaled_gemm_persistent.py",
 ]
 ignore-missing-imports = [
     # XPU memory symbols not present for builds without XPU support
@@ -95,6 +97,7 @@ ignore-missing-imports = [
     "boto3.*",
     "dill.*",
     "usort.*",
+    "cutlass.*",
     "cutlass_library.*",
     "cutlass_api.*",
     "deeplearning.*",
diff --git a/related_commits b/related_commits
new file mode 100644
index 0000000000000..b5b7595821464
--- /dev/null
+++ b/related_commits
@@ -0,0 +1,3 @@
+ubuntu|pytorch|apex|release/1.12.0|ad1376f20dc0ac107040b6077c96aed52bd14298|https://github.com/ROCm/apex
+ubuntu|pytorch|torchvision|release/0.27|b3d5b111f54044b216013162706ca107aa7a4ec4|https://github.com/ROCm/vision
+ubuntu|pytorch|torchaudio|release/2.11|34c52a67e8941bbd8e6adaca0eb0b9eabec11d78|https://github.com/pytorch/audio
diff --git a/requirements-build.txt b/requirements-build.txt
index 863bc9f921d8d..32367b67bf356 100644
--- a/requirements-build.txt
+++ b/requirements-build.txt
@@ -1,11 +1,15 @@
 # Build System requirements
-setuptools>=70.1.0,<82
-cmake>=3.27
-ninja
-numpy
-packaging
-pyyaml
-requests
-six  # dependency chain: NNPACK -> PeachPy -> six
-typing-extensions>=4.15.0
-pip  # not technically needed, but this makes setup.py invocation work
+# Build requirements pinned to match rocm/pytorch release/2.11:
+# https://github.com/ROCm/pytorch/blob/83524a4c7d748e5de89a7a93cd72dab4bd1fa42d/requirements-build.txt
+setuptools==79.0.1
+cmake==4.0.0
+ninja==1.11.1.4
+numpy==2.0.2 ; python_version == "3.9"
+numpy==2.1.2 ; python_version > "3.9" and python_version < "3.14"
+numpy==2.4.3 ; python_version >= "3.14"
+packaging==25.0
+pyyaml==6.0.3
+requests==2.32.5
+six==1.17.0  # dependency chain: NNPACK -> PeachPy -> six
+typing_extensions==4.15.0
+pip==26.0.1  # not technically needed, but this makes setup.py invocation work
diff --git a/requirements.txt b/requirements.txt
index ae7f335c883cf..d6552edad483b 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -4,16 +4,17 @@
 --requirement requirements-build.txt
 
 # Install / Development extra requirements
-build[uv]  # for building sdist and wheel
-expecttest>=0.3.0
-filelock
-fsspec>=0.8.5
-hypothesis
-jinja2
-lintrunner ; platform_machine != "s390x"
-networkx>=2.5.1
-optree>=0.13.0
-psutil
-spin
-sympy>=1.13.3
-wheel
+build[uv]==1.3.0  # for building sdist and wheel
+expecttest==0.3.0
+filelock==3.20.3
+fsspec==2026.2.0
+hypothesis==6.56.4
+jinja2==3.1.6
+lintrunner==0.12.11 ; platform_machine != "s390x"
+networkx==2.8.8
+optree==0.13.0 ; python_version < "3.14"
+optree==0.17.0 ; python_version >= "3.14"
+psutil==7.2.2
+spin==0.17
+sympy==1.13.3
+wheel==0.46.3
diff --git a/scripts/build_host_protoc.sh b/scripts/build_host_protoc.sh
old mode 100644
new mode 100755
diff --git a/scripts/export/update_schema.py b/scripts/export/update_schema.py
index a86f86c460e8c..bad55c39b18b7 100644
--- a/scripts/export/update_schema.py
+++ b/scripts/export/update_schema.py
@@ -73,6 +73,12 @@
     cpp_header += "\n// clang-format on"
     cpp_header += "\n"
 
+    enum_converter_header = "// " + first_line
+    enum_converter_header += "\n// clang-format off"
+    enum_converter_header += "\n" + commit.enum_converter_header
+    enum_converter_header += "\n// clang-format on"
+    enum_converter_header += "\n"
+
     yaml_content = yaml_header + "\n" + yaml_payload
 
     thrift_schema = "// " + first_line
@@ -87,5 +93,9 @@
             f.write(yaml_content)
         with open(os.path.join(args.prefix, commit.cpp_header_path), "w") as f:
             f.write(cpp_header)
+        with open(
+            os.path.join(args.prefix, commit.enum_converter_header_path), "w"
+        ) as f:
+            f.write(enum_converter_header)
         with open(os.path.join(args.prefix, commit.thrift_schema_path), "w") as f:
             f.write(thrift_schema)
diff --git a/scripts/jit/log_extract.py b/scripts/jit/log_extract.py
index 60aeaab92fc8c..b3ecf6683ddc4 100644
--- a/scripts/jit/log_extract.py
+++ b/scripts/jit/log_extract.py
@@ -1,7 +1,7 @@
 import argparse
 import functools
 import traceback
-from typing import Callable, List, Optional, Tuple
+from typing import Callable, List, Tuple
 
 from torch.utils.jit.log_extract import (
     extract_ir,
@@ -29,7 +29,7 @@
 def test_runners(
     graphs: List[str],
     runners: List[Tuple[str, Callable]],
-    graph_set: Optional[List[int]],
+    graph_set: List[int] | None,
 ):
     for i, ir in enumerate(graphs):
         _, inputs = load_graph_and_inputs(ir)
diff --git a/scripts/lintrunner.py b/scripts/lintrunner.py
index 2e3ad2bc219ab..1a5b8ab73b9a2 100644
--- a/scripts/lintrunner.py
+++ b/scripts/lintrunner.py
@@ -18,8 +18,6 @@
 
 import hashlib
 import os
-import shlex
-import shutil
 import subprocess
 import sys
 from pathlib import Path
diff --git a/scripts/onnx/test.sh b/scripts/onnx/test.sh
index c3b9a2ae618a0..8f64f7715279d 100755
--- a/scripts/onnx/test.sh
+++ b/scripts/onnx/test.sh
@@ -40,12 +40,10 @@ args+=("--cov-report")
 args+=("xml:test/coverage.xml")
 args+=("--cov-append")
 
-time python "${top_dir}/test/run_test.py" --onnx --shard "$SHARD_NUMBER" 2 --verbose
+time python "${top_dir}/test/run_test.py" --onnx --verbose
 
-if [[ "$SHARD_NUMBER" == "2" ]]; then
-  # xdoctests on onnx
-  xdoctest torch.onnx --style=google --options="+IGNORE_WHITESPACE"
-fi
+# xdoctests on onnx
+xdoctest torch.onnx --style=google --options="+IGNORE_WHITESPACE"
 
 # Our CI expects both coverage.xml and .coverage to be within test/
 if [ -d .coverage ]; then
diff --git a/scripts/release/apply-release-changes.sh b/scripts/release/apply-release-changes.sh
index 49d3d0b5192f7..63c00b45a30ba 100755
--- a/scripts/release/apply-release-changes.sh
+++ b/scripts/release/apply-release-changes.sh
@@ -24,17 +24,16 @@ done
 
 echo "Applying to templates"
 for i in .github/templates/*.yml.j2; do
-    sed -i 's#common.checkout(\(.*\))#common.checkout(\1, checkout_pr_head=False)#' $i;
+    sed -i '/checkout_pr_head/!s#common.checkout(\([^)]*\))#common.checkout(\1, checkout_pr_head=False)#' $i;
     sed -i -e s#main#"release/${RELEASE_VERSION}"# $i;
 done
 
 echo "Applying to changes to linux binary builds"
 for i in  ".github/workflows/_binary-build-linux.yml" ".github/workflows/_binary-test-linux.yml"; do
     sed -i "/github.event_name == 'pull_request'/d" $i;
-    sed -i -e s#main#"release/${RELEASE_VERSION}"# $i;
 done
 
-sed -i -e "/generate_ci_workflows.py/i \\\t\t\t\texport RELEASE_VERSION_TAG=${RELEASE_VERSION}" .github/workflows/lint.yml
+sed -i "s#^        \.github/scripts/generate_ci_workflows.py#        RELEASE_VERSION_TAG=${RELEASE_VERSION} .github/scripts/generate_ci_workflows.py#" .github/workflows/lint.yml
 
 # Triton wheel
 echo "Triton Changes"
@@ -45,6 +44,10 @@ echo "XLA Changes"
 sed -i -e s#--quiet#-b\ r"${RELEASE_VERSION}"# .ci/pytorch/common_utils.sh
 sed -i -e s#.*#r"${RELEASE_VERSION}"# .github/ci_commit_pins/xla.txt
 
+# Strip +PTX from CUDA arch lists in release builds
+echo "Stripping +PTX from CUDA arch lists"
+sed -i 's/+PTX//' .ci/manywheel/build_cuda.sh
+
 # Regenerate templates
 export RELEASE_VERSION_TAG=${RELEASE_VERSION}
 ./.github/regenerate.sh
@@ -58,5 +61,5 @@ sed -i -e s#unstable-jobs.json#"unstable-jobs.json?versionId=${UNSTABLE_VER}"# .
 sed -i -e s#disabled-jobs.json#"disabled-jobs.json?versionId=${DISABLED_VER}"# .github/scripts/filter_test_configs.py
 sed -i -e s#disabled-tests-condensed.json#"disabled-tests-condensed.json?versionId=${DISABLED_TESTS_VER}"# tools/stats/import_test_stats.py
 # Optional
-git commit -m "[RELEASE-ONLY CHANGES] Branch Cut for Release {RELEASE_VERSION}"
-git push origin "${RELEASE_BRANCH}"
+git commit -m "[RELEASE-ONLY CHANGES] Branch Cut for Release ${RELEASE_VERSION}"
+git push origin "release/${RELEASE_VERSION}"
diff --git a/scripts/release_notes/categorize.py b/scripts/release_notes/categorize.py
index 10ee551b74f1a..aa27e0b50577f 100644
--- a/scripts/release_notes/categorize.py
+++ b/scripts/release_notes/categorize.py
@@ -4,27 +4,25 @@
 from pathlib import Path
 
 import common
-
-# Imports for working with classi
-from classifier import (
-    CategoryConfig,
-    CommitClassifier,
-    CommitClassifierInputs,
-    get_author_map,
-    get_file_map,
-    XLMR_BASE,
-)
 from commitlist import CommitList
 from common import get_commit_data_cache, topics
 
-import torch
-
 
 class Categorizer:
     def __init__(self, path, category="Uncategorized", use_classifier: bool = False):
         self.cache = get_commit_data_cache()
         self.commits = CommitList.from_existing(path)
         if use_classifier:
+            from classifier import (
+                CategoryConfig,
+                CommitClassifier,
+                get_author_map,
+                get_file_map,
+                XLMR_BASE,
+            )
+
+            import torch
+
             print("Using a classifier to aid with categorization.")
             device = "cuda" if torch.cuda.is_available() else "cpu"
             classifier_config = CategoryConfig(common.categories)
@@ -108,6 +106,8 @@ def handle_commit(self, commit, i, total, commits):
 
         features = self.features(commit)
         if self.classifier is not None:
+            from classifier import CommitClassifierInputs
+
             # Some commits don't have authors:
             author = features.author if features.author else "Unknown"
             files = " ".join(features.files_changed)
diff --git a/scripts/release_notes/commitlist.py b/scripts/release_notes/commitlist.py
index 916d87ce0da07..9cc58588f7635 100644
--- a/scripts/release_notes/commitlist.py
+++ b/scripts/release_notes/commitlist.py
@@ -231,7 +231,6 @@ def categorize(features):
                 file,
                 [
                     "docker/",
-                    ".circleci",
                     ".github",
                     ".jenkins",
                     ".ci",
@@ -493,15 +492,18 @@ def get_markdown_header(category):
     header = f"""
 # Release Notes worksheet {category}
 
-The main goal of this process is to rephrase all the commit messages below to make them **clear and easy to read** by the end user. You should follow the following instructions to do so:
+You should:
+
+1. ensure commit categorization is correct
+2. write up major features, bc-breaking changes, deprecations in detail
+3. summarize the other sections
+
+## 1. Ensure commit categorization is correct
 
-* **Please clean up and format commit titles to be readable by the general PyTorch user.** Make sure you're [following the guidance here](https://docs.google.com/document/d/14OmgGBr1w6gl1VO47GGGdwrIaUNr92DFhQbY_NEk8mQ/edit)! Your resulting notes must be consistent and easy to read.
 * Please sort commits into the following categories (you should not rename the categories!), I tried to pre-sort these to ease your work, feel free to move commits around if the current categorization is not good.
 * Anything that is not public facing needs to be removed.
 * If anything is miscategorized/belongs to another domain, move it to `miscategorized.md`.
 * Please scan through `miscategorized.md` and handle any commits that belong within your domain according to these instructions.
-* We place a lot of emphasis on the “BC-breaking” and “deprecation” sections. Those should be where the most effort goes in. The “improvements” and “bug fixes” for Python API should be nice as well.
-* Once you are finished, move this very file from `todo/` to `done/` and submit a pull request.
 
 The categories below are as follows:
 
@@ -514,6 +516,23 @@ def get_markdown_header(category):
 * documentation: All commits that add/update documentation
 * Developers: All commits that are not end-user facing but still impact people that compile from source, develop into pytorch, extend pytorch, etc
 * not user facing: All commits that are not public end-user facing and hence should be dropped from the release notes
+
+## 2. Major features, BC-breaking changes, deprecations
+
+The main goal of this process is to rephrase all the commit messages below to make them **clear and easy to read** by the end user. You should follow the following instructions to do so:
+
+* **Please clean up and format commit titles to be readable by the general PyTorch user.** Make sure you're [following the guidance here](https://docs.google.com/document/d/14OmgGBr1w6gl1VO47GGGdwrIaUNr92DFhQbY_NEk8mQ/edit)! Your resulting notes must be consistent and easy to read.
+* We place a lot of emphasis on the “BC-breaking” and “deprecation” sections. Those should be where the most effort goes in. The “improvements” and “bug fixes” for Python API should be nice as well.
+
+## 3. Summarize the other sections
+
+For the other sections (improvements, bug fixes, performance, documentation, developers, not user facing) - use your
+judgement to summarize the key PRs. You do not need to make every commit description perfect
+(changed in v2.10 to simplify the process).
+
+Once you are finished, move this very file from `todo/` to `done/` and submit a pull request.
+
+Feel free to use https://github.com/pytorch/pytorch/releases/tag/v2.10.0 as an example.
 """
 
     return [header]
diff --git a/scripts/release_notes/explore.ipynb b/scripts/release_notes/explore.ipynb
index b979480863e3d..35131907d0095 100644
--- a/scripts/release_notes/explore.ipynb
+++ b/scripts/release_notes/explore.ipynb
@@ -7,11 +7,9 @@
    "outputs": [],
    "source": [
     "import pandas as pd\n",
-    "import numpy as np\n",
     "from pprint import pprint\n",
     "from collections import Counter\n",
-    "import common\n",
-    "import math"
+    "import common"
    ]
   },
   {
diff --git a/scripts/release_notes/parse_cherry_picks.py b/scripts/release_notes/parse_cherry_picks.py
new file mode 100644
index 0000000000000..5494a4fdba203
--- /dev/null
+++ b/scripts/release_notes/parse_cherry_picks.py
@@ -0,0 +1,319 @@
+#!/usr/bin/env python3
+"""Parse cherry-pick comments from a PyTorch GitHub issue and extract trunk PR info.
+
+Usage:
+    python scripts/release_notes/parse_cherry_picks.py \\
+        https://github.com/pytorch/pytorch/issues/170119 \\
+        --commitlist scripts/release_notes/results/commitlist.csv
+
+Outputs a CSV file to scripts/release_notes/results/cherry_picks_<issue_number>.csv.
+Columns: comment_id, pr_number, pr_title, commit_sha
+Validates that each commit_sha matches an entry in the commitlist and logs
+warnings for any mismatches.
+"""
+
+import argparse
+import csv
+import json
+import logging
+import re
+import subprocess
+import sys
+from pathlib import Path
+
+
+logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s")
+logger = logging.getLogger(__name__)
+
+
+def parse_issue_url(url: str) -> tuple[str, str]:
+    """Extract owner/repo and issue number from a GitHub issue URL."""
+    m = re.match(r"https://github\.com/([^/]+/[^/]+)/issues/(\d+)", url)  # @lint-ignore
+    if not m:
+        raise ValueError(f"Invalid GitHub issue URL: {url}")
+    return m.group(1), m.group(2)
+
+
+def gh_api(endpoint: str) -> list | dict:
+    """Call gh api with pagination and return parsed JSON."""
+    result = subprocess.run(
+        ["gh", "api", endpoint, "--paginate"],
+        capture_output=True,
+        text=True,
+        check=True,
+    )
+    # --paginate may return multiple JSON arrays concatenated; we need to handle that
+    # gh api --paginate with --jq is cleaner, but let's parse raw output
+    # When paginating, gh outputs one JSON array per page on separate "lines"
+    output = result.stdout.strip()
+    if not output:
+        return []
+
+    # Try parsing as a single JSON value first
+    try:
+        return json.loads(output)
+    except json.JSONDecodeError:
+        pass
+
+    # If that fails, it's multiple JSON arrays concatenated
+    # Split on ][ boundaries and merge
+    all_items = []
+    decoder = json.JSONDecoder()
+    pos = 0
+    while pos < len(output):
+        # Skip whitespace
+        while pos < len(output) and output[pos] in " \t\n\r":
+            pos += 1
+        if pos >= len(output):
+            break
+        obj, end_pos = decoder.raw_decode(output, pos)
+        if isinstance(obj, list):
+            all_items.extend(obj)
+        else:
+            all_items.append(obj)
+        pos = end_pos
+    return all_items
+
+
+def fetch_comments(repo: str, issue_number: str) -> list[dict]:
+    """Fetch all comments from a GitHub issue."""
+    endpoint = f"repos/{repo}/issues/{issue_number}/comments?per_page=100"
+    comments = gh_api(endpoint)
+    logger.info(f"Fetched {len(comments)} comments from {repo}#{issue_number}")
+    return comments
+
+
+def fetch_pr_title(repo: str, pr_number: str) -> str:
+    """Fetch the title of a PR."""
+    try:
+        result = subprocess.run(
+            ["gh", "api", f"repos/{repo}/pulls/{pr_number}", "--jq", ".title"],
+            capture_output=True,
+            text=True,
+            check=True,
+        )
+        return result.stdout.strip()
+    except subprocess.CalledProcessError as e:
+        logger.warning(f"Failed to fetch title for PR #{pr_number}: {e}")
+        return ""
+
+
+def fetch_landed_commit(repo: str, pr_number: str) -> str:
+    """Fetch the actual landed commit SHA for a PR.
+
+    PyTorch uses a merge bot that squash-merges outside of GitHub's standard
+    merge mechanism, so the PR API's merge_commit_sha is unreliable.
+
+    Strategy:
+    1. Issue events API: look for the 'closed' event with a commit_id
+    2. Commit search API: search for commits containing '(#NNNNN)' in message
+    """
+    # Strategy 1: Issue events API
+    try:
+        result = subprocess.run(
+            [
+                "gh",
+                "api",
+                f"repos/{repo}/issues/{pr_number}/events?per_page=100",
+                "--paginate",
+                "--jq",
+                '.[] | select(.event == "closed" and .commit_id) | .commit_id',
+            ],
+            capture_output=True,
+            text=True,
+            check=True,
+        )
+        commit_ids = result.stdout.strip().splitlines()
+        if commit_ids and commit_ids[0]:
+            return commit_ids[0]
+    except subprocess.CalledProcessError:
+        pass
+
+    # Strategy 2: Commit search API
+    logger.info(f"  Events API had no commit for PR #{pr_number}, trying search API...")
+    try:
+        result = subprocess.run(
+            [
+                "gh",
+                "api",
+                f'search/commits?q=repo:{repo}+"(#{pr_number})"',
+                "--jq",
+                ".items[0].sha",
+            ],
+            capture_output=True,
+            text=True,
+            check=True,
+        )
+        sha = result.stdout.strip()
+        if sha and sha != "null":
+            return sha
+    except subprocess.CalledProcessError:
+        pass
+
+    logger.warning(f"Could not find landed commit for PR #{pr_number}")
+    return ""
+
+
+def extract_trunk_prs(comment_body: str) -> list[dict]:
+    """Extract trunk PR numbers/commit hashes from a comment body.
+
+    Returns a list of dicts with keys: pr_number (str or ""), raw_commit (str or "").
+    Returns an empty list if the comment doesn't contain the trunk PR section.
+    """
+    # Find the "Link to landed trunk PR" section
+    # Handle variations: "Link to landed trunk PR", "Link to the landed trunk PR"
+    trunk_pattern = re.compile(
+        r"Link to (?:the )?landed trunk PR[^:]*:\s*\n(.*?)(?=Li[nn][kt] to (?:the )?release branch PR|Criteria Category:|$)",
+        re.DOTALL | re.IGNORECASE,
+    )
+    match = trunk_pattern.search(comment_body)
+    if not match:
+        return []
+
+    trunk_section = match.group(1).strip()
+
+    # Check for NA/N/A
+    if re.match(r"^\*?\s*-?\s*N/?A\s*$", trunk_section, re.IGNORECASE):
+        return [{"pr_number": "", "raw_commit": ""}]
+
+    results = []
+
+    # Extract PR URLs
+    pr_urls = re.findall(
+        r"https://github\.com/pytorch/pytorch/pull/(\d+)",  # @lint-ignore
+        trunk_section,
+    )
+    for pr_num in pr_urls:
+        results.append({"pr_number": pr_num, "raw_commit": ""})
+
+    # Extract raw commit hashes (40-char hex) that aren't part of a URL
+    # Remove URLs first to avoid matching hashes embedded in URLs
+    section_no_urls = re.sub(r"https://\S+", "", trunk_section)
+    raw_commits = re.findall(r"\b([0-9a-f]{40})\b", section_no_urls)
+    for commit_hash in raw_commits:
+        results.append({"pr_number": "", "raw_commit": commit_hash})
+
+    # If we found the section but nothing matched, treat as empty/NA
+    if not results:
+        return [{"pr_number": "", "raw_commit": ""}]
+
+    return results
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Parse cherry-pick comments from a PyTorch GitHub issue"
+    )
+    parser.add_argument("issue_url", help="GitHub issue URL to parse")
+    parser.add_argument(
+        "-o",
+        "--output",
+        default=None,
+        help="Output CSV file path (default: results/cherry_picks_<issue_number>.csv relative to script)",
+    )
+    parser.add_argument(
+        "--commitlist",
+        required=True,
+        help="Path to commitlist.csv to validate commit hashes against",
+    )
+    args = parser.parse_args()
+
+    repo, issue_number = parse_issue_url(args.issue_url)
+    if args.output:
+        output_path = args.output
+    else:
+        script_dir = Path(__file__).resolve().parent
+        results_dir = script_dir / "results"
+        results_dir.mkdir(exist_ok=True)
+        output_path = str(results_dir / f"cherry_picks_{issue_number}.csv")
+
+    comments = fetch_comments(repo, issue_number)
+
+    rows = []
+    for comment in comments:
+        comment_id = comment["id"]
+        body = comment["body"] or ""
+
+        trunk_prs = extract_trunk_prs(body)
+        if not trunk_prs:
+            # Comment doesn't contain the trunk PR section — skip
+            continue
+
+        if len(trunk_prs) > 1:
+            logger.warning(
+                f"Comment {comment_id} has {len(trunk_prs)} trunk PRs: "
+                + ", ".join(
+                    p["pr_number"] or p["raw_commit"] or "N/A" for p in trunk_prs
+                )
+            )
+
+        for pr_info in trunk_prs:
+            pr_number = pr_info["pr_number"]
+            raw_commit = pr_info["raw_commit"]
+            pr_title = ""
+            commit_sha = raw_commit  # Use raw commit if provided
+
+            if pr_number:
+                logger.info(f"Fetching info for PR #{pr_number}...")
+                pr_title = fetch_pr_title(repo, pr_number)
+                commit_sha = fetch_landed_commit(repo, pr_number)
+
+            rows.append(
+                {
+                    "comment_id": comment_id,
+                    "pr_number": pr_number,
+                    "pr_title": pr_title,
+                    "commit_sha": commit_sha,
+                }
+            )
+
+    # Validate against commitlist
+    commitlist_hashes = set()
+    with open(args.commitlist, newline="") as f:
+        reader = csv.reader(f)
+        next(reader)  # skip header
+        for row in reader:
+            if row:
+                commitlist_hashes.add(row[0])
+
+    logger.info(f"Loaded {len(commitlist_hashes)} hashes from {args.commitlist}")
+
+    matched = 0
+    mismatched = 0
+    skipped = 0
+    for row in rows:
+        sha = row["commit_sha"]
+        if not sha:
+            skipped += 1
+            continue
+        # commitlist uses abbreviated hashes; check if any is a prefix of
+        # our full hash, or vice versa
+        if any(
+            sha.startswith(cl_hash) or cl_hash.startswith(sha)
+            for cl_hash in commitlist_hashes
+        ):
+            matched += 1
+        else:
+            mismatched += 1
+            logger.warning(
+                f"Commit {sha[:11]} (PR #{row['pr_number'] or 'N/A'}) "
+                f"not found in commitlist"
+            )
+
+    logger.info(
+        f"Commitlist validation: {matched} matched, {mismatched} not found, "
+        f"{skipped} skipped (no hash)"
+    )
+
+    # Write CSV
+    fieldnames = ["comment_id", "pr_number", "pr_title", "commit_sha"]
+    with open(output_path, "w", newline="") as f:
+        writer = csv.DictWriter(f, fieldnames=fieldnames)
+        writer.writeheader()
+        writer.writerows(rows)
+
+    logger.info(f"Wrote {len(rows)} rows to {output_path}")
+
+
+if __name__ == "__main__":
+    main()
diff --git a/scripts/release_notes/remove_cherry_picks.py b/scripts/release_notes/remove_cherry_picks.py
new file mode 100644
index 0000000000000..b7ef72c11d594
--- /dev/null
+++ b/scripts/release_notes/remove_cherry_picks.py
@@ -0,0 +1,139 @@
+#!/usr/bin/env python3
+"""Remove cherry-picked commits from a commitlist CSV.
+
+Usage:
+    python scripts/release_notes/remove_cherry_picks.py \\
+        --commitlist scripts/release_notes/results/commitlist.csv \\
+        --cherry-picks scripts/release_notes/results/cherry_picks_170119.csv
+
+Reads the cherry picks CSV (output of parse_cherry_picks.py) and removes
+matching rows from the commitlist, writing the result to a new file.
+The original commitlist is never modified.
+"""
+
+import argparse
+import csv
+import logging
+from pathlib import Path
+
+
+logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s")
+logger = logging.getLogger(__name__)
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Remove cherry-picked commits from a commitlist CSV"
+    )
+    parser.add_argument(
+        "--commitlist",
+        required=True,
+        help="Path to the input commitlist.csv",
+    )
+    parser.add_argument(
+        "--cherry-picks",
+        required=True,
+        help="Path to the cherry picks CSV (output of parse_cherry_picks.py)",
+    )
+    parser.add_argument(
+        "-o",
+        "--output",
+        default=None,
+        help="Output commitlist path (default: commitlist_no_cherry_picks.csv in same dir as commitlist)",
+    )
+    parser.add_argument(
+        "--log",
+        default=None,
+        help="Path for the removal log file (default: results/cherry_pick_removals.log relative to script)",
+    )
+    args = parser.parse_args()
+
+    commitlist_path = Path(args.commitlist)
+    if args.output:
+        output_path = Path(args.output)
+    else:
+        output_path = commitlist_path.parent / "commitlist_no_cherry_picks.csv"
+
+    script_dir = Path(__file__).resolve().parent
+    if args.log:
+        log_path = Path(args.log)
+    else:
+        results_dir = script_dir / "results"
+        results_dir.mkdir(exist_ok=True)
+        log_path = results_dir / "cherry_pick_removals.log"
+
+    # Load cherry picks
+    cherry_picks = []
+    with open(args.cherry_picks, newline="") as f:
+        reader = csv.DictReader(f)
+        for row in reader:
+            cherry_picks.append(row)
+    logger.info(f"Loaded {len(cherry_picks)} cherry pick entries")
+
+    # Read commitlist as raw lines to preserve exact formatting
+    with open(commitlist_path) as f:
+        header_line = f.readline()
+        commitlist_lines = f.readlines()
+    logger.info(f"Loaded {len(commitlist_lines)} rows from {commitlist_path}")
+
+    # Build a map from abbreviated hash to line for fast lookup
+    hash_to_line = {}
+    for line in commitlist_lines:
+        abbrev_hash = line.split(",", 1)[0]
+        hash_to_line[abbrev_hash] = line
+
+    # Process cherry picks: find and remove matching rows
+    removed = 0
+    not_found = 0
+    skipped = 0
+    log_entries = []
+
+    for cp in cherry_picks:
+        sha = cp.get("commit_sha", "")
+        pr_number = cp.get("pr_number", "")
+        pr_title = cp.get("pr_title", "")
+        label = f"PR #{pr_number} ({pr_title})" if pr_number else "N/A"
+
+        if not sha:
+            skipped += 1
+            log_entries.append(f"SKIPPED (no hash): {label}")
+            continue
+
+        # Find matching abbreviated hash via prefix match
+        matched_hash = None
+        for abbrev_hash in hash_to_line:
+            if sha.startswith(abbrev_hash) or abbrev_hash.startswith(sha):
+                matched_hash = abbrev_hash
+                break
+
+        if matched_hash:
+            removed += 1
+            log_entries.append(f"REMOVED: {matched_hash} -> {label}")
+            del hash_to_line[matched_hash]
+        else:
+            not_found += 1
+            log_entries.append(f"NOT FOUND: {sha[:11]} -> {label}")
+            logger.warning(f"Commit {sha[:11]} ({label}) not found in commitlist")
+
+    # Write new commitlist
+    with open(output_path, "w") as f:
+        f.write(header_line)
+        for line in commitlist_lines:
+            abbrev_hash = line.split(",", 1)[0]
+            if abbrev_hash in hash_to_line:
+                f.write(line)
+
+    logger.info(
+        f"Wrote {len(hash_to_line)} rows to {output_path} "
+        f"(removed {removed}, not found {not_found}, skipped {skipped})"
+    )
+
+    # Write log file
+    with open(log_path, "w") as f:
+        for entry in log_entries:
+            f.write(entry + "\n")
+    logger.info(f"Wrote removal log to {log_path}")
+
+
+if __name__ == "__main__":
+    main()
diff --git a/setup.py b/setup.py
index 7b9fba42e746e..41f5ac099dbfb 100644
--- a/setup.py
+++ b/setup.py
@@ -271,7 +271,6 @@
     sys.exit(-1)
 
 import filecmp
-import glob
 import importlib
 import itertools
 import json
@@ -315,6 +314,7 @@
 ).rstrip(os.pathsep)
 
 from tools.build_pytorch_libs import build_pytorch
+from tools.clean import clean as _clean
 from tools.generate_torch_version import get_torch_version
 from tools.setup_helpers.cmake import CMake, CMakeValue
 from tools.setup_helpers.env import (
@@ -399,56 +399,45 @@ def _get_package_path(package_name: str) -> Path:
 # see if the user passed a quiet flag to setup.py arguments and respect
 # that in our parts of the build
 EMIT_BUILD_WARNING = False
-RERUN_CMAKE = str2bool(os.environ.pop("CMAKE_FRESH", None))
-CMAKE_ONLY = str2bool(os.environ.pop("CMAKE_ONLY", None))
+RERUN_CMAKE = str2bool(os.environ.pop("CMAKE_FRESH", None)) or "--cmake" in sys.argv
+CMAKE_ONLY = str2bool(os.environ.pop("CMAKE_ONLY", None)) or "--cmake-only" in sys.argv
 filtered_args = []
 for i, arg in enumerate(sys.argv):
-    if arg == "--cmake":
-        RERUN_CMAKE = True
-        continue
-    if arg == "--cmake-only":
-        # Stop once cmake terminates. Leave users a chance to adjust build
-        # options.
-        CMAKE_ONLY = True
+    if arg in ("--cmake", "--cmake-only"):
         continue
     if arg == "rebuild" or arg == "build":
         arg = "build"  # rebuild is gone, make it build
         EMIT_BUILD_WARNING = True
-    if arg == "develop":
-        print(
-            (
-                "WARNING: Redirecting 'python setup.py develop' to 'pip install -e . -v --no-build-isolation',"
-                " for more info see https://github.com/pytorch/pytorch/issues/152276"
-            ),
-            file=sys.stderr,
-        )
-        result = subprocess.run(
-            [
-                sys.executable,
-                "-m",
-                "pip",
-                "install",
-                "-e",
-                ".",
-                "-v",
-                "--no-build-isolation",
-            ],
-            env={**os.environ},
-        )
-        sys.exit(result.returncode)
-    if arg == "install":
-        print(
-            (
-                "WARNING: Redirecting 'python setup.py install' to 'pip install . -v --no-build-isolation',"
-                " for more info see https://github.com/pytorch/pytorch/issues/152276"
-            ),
-            file=sys.stderr,
-        )
-        result = subprocess.run(
-            [sys.executable, "-m", "pip", "install", ".", "-v", "--no-build-isolation"],
-            env={**os.environ},
-        )
-        sys.exit(result.returncode)
+    if arg in ("develop", "install"):
+        # CMAKE_ONLY only runs cmake and exits (via sys.exit in build_deps)
+        # before setup() is called, so there's no need to go through pip.
+        # Replace the command with "build" so setuptools doesn't complain
+        # about an unrecognized command if we somehow reach setup().
+        if CMAKE_ONLY:
+            arg = "build"
+        else:
+            editable = arg == "develop"
+            print(
+                (
+                    f"WARNING: Redirecting 'python setup.py {arg}' to "
+                    f"'pip install {'-e ' if editable else ''}. -v --no-build-isolation',"
+                    " for more info see https://github.com/pytorch/pytorch/issues/152276"
+                ),
+                file=sys.stderr,
+            )
+            result = subprocess.run(
+                [
+                    sys.executable,
+                    "-m",
+                    "pip",
+                    "install",
+                    *(("-e", ".") if editable else (".",)),
+                    "-v",
+                    "--no-build-isolation",
+                ],
+                env={**os.environ},
+            )
+            sys.exit(result.returncode)
     if arg == "--":
         filtered_args += sys.argv[i:]
         break
@@ -637,7 +626,8 @@ def mirror_inductor_external_kernels() -> None:
     cuda_is_disabled = not str2bool(os.getenv("USE_CUDA"))
     paths = [
         (
-            CWD / "torch/_inductor/kernel/vendored_templates/cutedsl_grouped_gemm.py",
+            CWD
+            / "torch/_inductor/kernel/vendored_templates/cutedsl/kernels/cutedsl_grouped_gemm.py",
             CWD
             / "third_party/cutlass/examples/python/CuTeDSL/blackwell/grouped_gemm.py",
             True,
@@ -1471,21 +1461,7 @@ def finalize_options(self) -> None:
         pass
 
     def run(self) -> None:
-        ignores = (CWD / ".gitignore").read_text(encoding="utf-8")
-        for wildcard in filter(None, ignores.splitlines()):
-            if wildcard.strip().startswith("#"):
-                if "BEGIN NOT-CLEAN-FILES" in wildcard:
-                    # Marker is found and stop reading .gitignore.
-                    break
-                # Ignore lines which begin with '#'.
-            else:
-                # Don't remove absolute paths from the system
-                wildcard = wildcard.lstrip("./")
-                for filename in glob.iglob(wildcard):
-                    try:
-                        os.remove(filename)
-                    except OSError:
-                        shutil.rmtree(filename, ignore_errors=True)
+        _clean()
 
 
 # Need to dump submodule hashes and create the proper LICENSE.txt for the sdist
@@ -1703,7 +1679,7 @@ def main() -> None:
     install_requires = [
         "filelock",
         "typing-extensions>=4.10.0",
-        'setuptools<82 ; python_version >= "3.12"',
+        "setuptools<82",
         "sympy>=1.13.3",
         "networkx>=2.5.1",
         "jinja2",
diff --git a/test/_test_bazel.py b/test/_test_bazel.py
index 998915f377459..27a80d7d15609 100644
--- a/test/_test_bazel.py
+++ b/test/_test_bazel.py
@@ -1,9 +1,9 @@
 # Owner(s): ["module: bazel"]
 
 """
-This test module contains a minimalistic "smoke tests" for the bazel build.
+This test module contains minimalistic "smoke tests" for the bazel build.
 
-Currently it doesn't use any testing framework (i.e. pytest)
+Currently it doesn't use any testing framework (i.e. pytest).
 TODO: integrate this into the existing pytorch testing framework.
 
 The name uses underscore `_test_bazel.py` to avoid globbing into other non-bazel configurations.
diff --git a/test/allowlist_for_publicAPI.json b/test/allowlist_for_publicAPI.json
index d01d41d37997e..4bc3d606f2ac4 100644
--- a/test/allowlist_for_publicAPI.json
+++ b/test/allowlist_for_publicAPI.json
@@ -1772,9 +1772,7 @@
     "Optional",
     "Union"
   ],
-  "torch.fx.experimental.migrate_gradual_types.constraint": [
-    "TensorType"
-  ],
+  "torch.fx.experimental.migrate_gradual_types.constraint": [],
   "torch.fx.experimental.migrate_gradual_types.constraint_generator": [
     "ApplyBroadcasting",
     "BatchNorm2d",
@@ -1830,6 +1828,7 @@
     "IndexSelect",
     "List",
     "Prod",
+    "Sequence",
     "T",
     "TGreatestUpperBound",
     "TVar",
@@ -1844,6 +1843,7 @@
     "BinConstraintD",
     "BinConstraintT",
     "Conj",
+    "Constraint",
     "ConstraintGenerator",
     "D",
     "DVar",
@@ -2009,7 +2009,7 @@
     "fx_placeholder_targets",
     "fx_placeholder_vals",
     "guard_bool",
-    "has_hint",
+    "has_guarding_hint",
     "is_symbolic",
     "parallel_and",
     "parallel_or",
diff --git a/test/bench_mps_ops.py b/test/bench_mps_ops.py
index b0bf80649b31f..62f16b250c0fe 100644
--- a/test/bench_mps_ops.py
+++ b/test/bench_mps_ops.py
@@ -5,7 +5,6 @@
 import sys
 import timeit
 import warnings
-from typing import Optional
 
 import torch
 from torch.utils.benchmark import Compare, Measurement, Timer
@@ -56,7 +55,7 @@ def bench_binary(
     binary_func,
     device: str = "mps",
     dt_a: torch.dtype = torch.float32,
-    dt_b: Optional[torch.dtype] = None,
+    dt_b: torch.dtype | None = None,
 ) -> list[Measurement]:
     dt_b = dt_b if dt_b is not None else dt_a
     x = torch.testing.make_tensor(1024, 1024, device=device, dtype=dt_a)
diff --git a/test/conftest.py b/test/conftest.py
index 39da48b51c5f4..93eae789b3dce 100644
--- a/test/conftest.py
+++ b/test/conftest.py
@@ -7,7 +7,7 @@
 import xml.etree.ElementTree as ET
 from collections import defaultdict
 from types import MethodType
-from typing import Any, Optional, TYPE_CHECKING, Union
+from typing import Any, TYPE_CHECKING
 
 import pytest
 from _pytest.config import Config, filename_arg
@@ -170,7 +170,7 @@ def append_rerun(self, reporter: _NodeReporter, report: TestReport) -> None:
         else:
             if report.longrepr is None:
                 raise AssertionError("Expected report.longrepr to not be None")
-            reprcrash: Optional[ReprFileLocation] = getattr(
+            reprcrash: ReprFileLocation | None = getattr(
                 report.longrepr, "reprcrash", None
             )
             if reprcrash is not None:
@@ -191,8 +191,8 @@ def pytest_runtest_logreport(self, report: TestReport) -> None:
                 reason = f"{report.nodeid}: {_get_raw_skip_reason(report)}"
                 report.longrepr = (fspath, lineno, reason)
 
-    def node_reporter(self, report: Union[TestReport, str]) -> _NodeReporterReruns:
-        nodeid: Union[str, TestReport] = getattr(report, "nodeid", report)
+    def node_reporter(self, report: TestReport | str) -> _NodeReporterReruns:
+        nodeid: str | TestReport = getattr(report, "nodeid", report)
         # Local hack to handle xdist report order.
         workernode = getattr(report, "node", None)
 
@@ -315,7 +315,7 @@ def __init__(self, config: Config) -> None:
         self.cache: pytest.Cache = config.cache
         directory = f"{STEPCURRENT_CACHE_DIR}/{config.getoption('stepcurrent')}"
         self.lastrun_location = f"{directory}/lastrun"
-        self.lastrun: Optional[str] = self.cache.get(self.lastrun_location, None)
+        self.lastrun: str | None = self.cache.get(self.lastrun_location, None)
         self.initial_val = self.lastrun
         self.skip: bool = config.getoption("stepcurrent_skip")
         self.run_single: bool = config.getoption("run_single")
@@ -351,7 +351,7 @@ def pytest_collection_modifyitems(self, config: Config, items: list[Any]) -> Non
                 del items[1:]
             config.hook.pytest_deselected(items=deselected)
 
-    def pytest_report_collectionfinish(self) -> Optional[str]:
+    def pytest_report_collectionfinish(self) -> str | None:
         if self.config.getoption("verbose") >= 0 and self.report_status:
             return f"stepcurrent: {self.report_status}"
         return None
diff --git a/test/cpp/aoti_eager/kernel_meta_info_test.cpp b/test/cpp/aoti_eager/kernel_meta_info_test.cpp
new file mode 100644
index 0000000000000..e22b10f307a1e
--- /dev/null
+++ b/test/cpp/aoti_eager/kernel_meta_info_test.cpp
@@ -0,0 +1,138 @@
+// Tests for TensorMetadata::dynamic_check and ParameterMetadata::dynamic_check,
+// which enable a single compiled AOTI kernel to serve multiple input shapes by
+// matching dtype/device/rank instead of exact sizes and strides.
+
+#if !defined(C10_MOBILE) && !defined(ANDROID)
+
+#include <gtest/gtest.h>
+
+#include <torch/csrc/inductor/aoti_eager/kernel_meta_info.h>
+
+namespace torch::inductor {
+
+namespace {
+
+TensorMetadata makeTensorMeta(
+    c10::ScalarType dtype,
+    std::vector<int64_t> sizes,
+    std::vector<int64_t> strides,
+    c10::Device device = c10::Device(c10::kCPU)) {
+  return TensorMetadata(
+      /*is_symbolic=*/false,
+      dtype,
+      device,
+      c10::DispatchKeySet(c10::DispatchKey::CPU),
+      std::move(sizes),
+      std::move(strides));
+}
+
+} // namespace
+
+// ---------- TensorMetadata::dynamic_check ----------
+
+TEST(TensorMetadataDynamicCheckTest, MatchesSameDtypeDeviceRankDifferentSizes) {
+  auto a = makeTensorMeta(c10::kFloat, {4, 8}, {8, 1});
+  auto b = makeTensorMeta(c10::kFloat, {16, 32}, {32, 1});
+  EXPECT_TRUE(a.dynamic_check(b));
+  // Exact comparison should fail because sizes differ
+  EXPECT_FALSE(a == b);
+}
+
+TEST(TensorMetadataDynamicCheckTest, RejectsDifferentDtype) {
+  auto a = makeTensorMeta(c10::kFloat, {4, 8}, {8, 1});
+  auto b = makeTensorMeta(c10::kHalf, {4, 8}, {8, 1});
+  EXPECT_FALSE(a.dynamic_check(b));
+}
+
+TEST(TensorMetadataDynamicCheckTest, RejectsDifferentRank) {
+  auto a = makeTensorMeta(c10::kFloat, {4, 8}, {8, 1});
+  auto b = makeTensorMeta(c10::kFloat, {4, 8, 2}, {16, 2, 1});
+  EXPECT_FALSE(a.dynamic_check(b));
+}
+
+TEST(TensorMetadataDynamicCheckTest, RejectsDifferentDevice) {
+  auto cpu = makeTensorMeta(c10::kFloat, {4}, {1}, c10::Device(c10::kCPU));
+  auto cuda = makeTensorMeta(c10::kFloat, {4}, {1}, c10::Device(c10::kCUDA, 0));
+  EXPECT_FALSE(cpu.dynamic_check(cuda));
+}
+
+TEST(TensorMetadataDynamicCheckTest, MatchesIdenticalMetadata) {
+  auto a = makeTensorMeta(c10::kFloat, {4, 8}, {8, 1});
+  auto b = makeTensorMeta(c10::kFloat, {4, 8}, {8, 1});
+  EXPECT_TRUE(a.dynamic_check(b));
+  EXPECT_TRUE(a == b);
+}
+
+TEST(TensorMetadataDynamicCheckTest, MatchesScalarTensors) {
+  auto a = makeTensorMeta(c10::kFloat, {}, {});
+  auto b = makeTensorMeta(c10::kFloat, {}, {});
+  EXPECT_TRUE(a.dynamic_check(b));
+}
+
+// ---------- ParameterMetadata::dynamic_check ----------
+
+TEST(ParameterMetadataDynamicCheckTest, TensorMatchesDifferentSizes) {
+  ParameterMetadata a(makeTensorMeta(c10::kFloat, {4, 8}, {8, 1}), 0);
+  ParameterMetadata b(makeTensorMeta(c10::kFloat, {16, 32}, {32, 1}), 0);
+  EXPECT_TRUE(a.dynamic_check(b));
+  EXPECT_FALSE(a == b);
+}
+
+TEST(ParameterMetadataDynamicCheckTest, TensorRejectsDifferentDtype) {
+  ParameterMetadata a(makeTensorMeta(c10::kFloat, {4}, {1}), 0);
+  ParameterMetadata b(makeTensorMeta(c10::kHalf, {4}, {1}), 0);
+  EXPECT_FALSE(a.dynamic_check(b));
+}
+
+TEST(ParameterMetadataDynamicCheckTest, RejectsDifferentOrder) {
+  ParameterMetadata a(makeTensorMeta(c10::kFloat, {4}, {1}), 0);
+  ParameterMetadata b(makeTensorMeta(c10::kFloat, {4}, {1}), 1);
+  EXPECT_FALSE(a.dynamic_check(b));
+}
+
+TEST(ParameterMetadataDynamicCheckTest, ScalarUsesExactMatch) {
+  ParameterMetadata a(c10::Scalar(1.5), 0);
+  ParameterMetadata b(c10::Scalar(1.5), 0);
+  EXPECT_TRUE(a.dynamic_check(b));
+
+  ParameterMetadata c(c10::Scalar(2.0), 0);
+  EXPECT_FALSE(a.dynamic_check(c));
+}
+
+TEST(ParameterMetadataDynamicCheckTest, TensorListMatchesDifferentSizes) {
+  std::vector<TensorMetadata> list_a = {
+      makeTensorMeta(c10::kFloat, {4, 8}, {8, 1}),
+      makeTensorMeta(c10::kFloat, {2, 8}, {8, 1}),
+  };
+  std::vector<TensorMetadata> list_b = {
+      makeTensorMeta(c10::kFloat, {16, 32}, {32, 1}),
+      makeTensorMeta(c10::kFloat, {8, 32}, {32, 1}),
+  };
+  ParameterMetadata a(list_a, 0);
+  ParameterMetadata b(list_b, 0);
+  EXPECT_TRUE(a.dynamic_check(b));
+}
+
+TEST(ParameterMetadataDynamicCheckTest, TensorListRejectsDifferentLengths) {
+  std::vector<TensorMetadata> one = {makeTensorMeta(c10::kFloat, {4}, {1})};
+  std::vector<TensorMetadata> two = {
+      makeTensorMeta(c10::kFloat, {4}, {1}),
+      makeTensorMeta(c10::kFloat, {8}, {1}),
+  };
+  ParameterMetadata a(one, 0);
+  ParameterMetadata b(two, 0);
+  EXPECT_FALSE(a.dynamic_check(b));
+}
+
+TEST(ParameterMetadataDynamicCheckTest, StringUsesExactMatch) {
+  ParameterMetadata a(std::string("sum"), 0);
+  ParameterMetadata b(std::string("sum"), 0);
+  EXPECT_TRUE(a.dynamic_check(b));
+
+  ParameterMetadata c(std::string("mean"), 0);
+  EXPECT_FALSE(a.dynamic_check(c));
+}
+
+} // namespace torch::inductor
+
+#endif // !defined(C10_MOBILE) && !defined(ANDROID)
diff --git a/test/cpp/aoti_inference/standalone_compile.sh b/test/cpp/aoti_inference/standalone_compile.sh
index 7e97f916f2dcc..350ec63f2f65e 100644
--- a/test/cpp/aoti_inference/standalone_compile.sh
+++ b/test/cpp/aoti_inference/standalone_compile.sh
@@ -20,7 +20,7 @@ if [ $# -lt 2 ]; then
 fi
 
 # Building the wrapper
-$CXX -I$TORCH_ROOT_DIR/build/aten/src -I$TORCH_ROOT_DIR/aten/src -I$TORCH_ROOT_DIR/build -I$TORCH_ROOT_DIR -I$TORCH_ROOT_DIR/build/caffe2/aten/src -I$TORCH_ROOT_DIR/torch/csrc/api -I$TORCH_ROOT_DIR/torch/csrc/api/include -std=gnu++17 -fPIE -o $1.o -c $1
+$CXX -I$TORCH_ROOT_DIR/build/aten/src -I$TORCH_ROOT_DIR/aten/src -I$TORCH_ROOT_DIR/build -I$TORCH_ROOT_DIR -I$TORCH_ROOT_DIR/build/caffe2/aten/src -I$TORCH_ROOT_DIR/torch/csrc/api -I$TORCH_ROOT_DIR/torch/csrc/api/include -std=gnu++20 -fPIE -o $1.o -c $1
 
 # Linking
 $CXX -rdynamic -Wl,--no-as-needed,$TORCH_ROOT_DIR/build/lib/libtorch.so $1.o -Wl,--no-as-needed,$TORCH_ROOT_DIR/build/lib/libtorch_cpu.so -Wl,--no-as-needed,$TORCH_ROOT_DIR/build/lib/libc10.so -o $2
diff --git a/test/cpp/api/optim.cpp b/test/cpp/api/optim.cpp
index 25263719e02aa..1370892e208cc 100644
--- a/test/cpp/api/optim.cpp
+++ b/test/cpp/api/optim.cpp
@@ -564,3 +564,508 @@ TEST(OptimTest, CheckLRChange_ReduceLROnPlateau_Adam) {
   check_lr_change_for_reduce_on_plateau(
       optimizer, reduce_lr_on_plateau_scheduler, expected_epoch_lrs);
 }
+// Tests for Issue 141884: Parameter group inheritance functionality
+// Validates that partial options in parameter groups correctly inherit
+// defaults from the optimizer while preserving explicitly set values
+TEST(OptimTest, MergeWithDefaultOptions_Adam) {
+  // Create tensors for parameter groups
+  auto tensor1 = torch::randn({2, 2}).requires_grad_(true);
+  auto tensor2 = torch::randn({3, 3}).requires_grad_(true);
+
+  // Create param groups with partial options
+  std::vector<OptimizerParamGroup> param_groups;
+
+  // Group 1: Only weight_decay specified, should inherit lr, betas, eps,
+  // amsgrad
+  param_groups.emplace_back(
+      std::vector<torch::Tensor>{tensor1},
+      std::make_unique<AdamOptions>(AdamOptions().weight_decay(0.11)));
+
+  // Group 2: Only eps specified, should inherit others
+  param_groups.emplace_back(
+      std::vector<torch::Tensor>{tensor2},
+      std::make_unique<AdamOptions>(AdamOptions().eps(1e-6)));
+
+  // Create optimizer with specific defaults
+  AdamOptions defaults;
+  defaults.lr(0.002)
+      .betas(std::make_tuple(0.8, 0.88))
+      .eps(1e-12)
+      .weight_decay(0.05)
+      .amsgrad(true);
+
+  Adam optimizer(param_groups, defaults);
+
+  // Check Group 1: weight_decay preserved, others inherited
+  auto& group1_opts =
+      static_cast<AdamOptions&>(optimizer.param_groups()[0].options());
+  ASSERT_EQ(group1_opts.lr(), 0.002); // Inherited
+  ASSERT_EQ(group1_opts.betas(), std::make_tuple(0.8, 0.88)); // Inherited
+  ASSERT_EQ(group1_opts.eps(), 1e-12); // Inherited
+  ASSERT_EQ(group1_opts.weight_decay(), 0.11); // Preserved
+  ASSERT_TRUE(group1_opts.amsgrad()); // Inherited
+
+  // Check Group 2: eps preserved, others inherited
+  auto& group2_opts =
+      static_cast<AdamOptions&>(optimizer.param_groups()[1].options());
+  ASSERT_EQ(group2_opts.lr(), 0.002); // Inherited
+  ASSERT_EQ(group2_opts.betas(), std::make_tuple(0.8, 0.88)); // Inherited
+  ASSERT_EQ(group2_opts.eps(), 1e-6); // Preserved
+  ASSERT_EQ(group2_opts.weight_decay(), 0.05); // Inherited
+  ASSERT_TRUE(group2_opts.amsgrad()); // Inherited
+}
+
+TEST(OptimTest, MergeWithDefaultOptions_SGD) {
+  // Create tensors for parameter groups
+  auto tensor1 = torch::randn({2, 2}).requires_grad_(true);
+  auto tensor2 = torch::randn({3, 3}).requires_grad_(true);
+
+  // Create param groups with partial options
+  std::vector<OptimizerParamGroup> param_groups;
+
+  // Group 1: Only lr and weight_decay specified, should inherit momentum,
+  // dampening, nesterov
+  param_groups.emplace_back(
+      std::vector<torch::Tensor>{tensor1},
+      std::make_unique<SGDOptions>(SGDOptions(0.01).weight_decay(0.22)));
+
+  // Group 2: Only lr specified, should inherit all others
+  param_groups.emplace_back(
+      std::vector<torch::Tensor>{tensor2},
+      std::make_unique<SGDOptions>(SGDOptions(0.02)));
+
+  // Create optimizer with specific defaults
+  SGDOptions defaults(0.001); // lr should be overridden by param groups
+  defaults.momentum(0.9)
+      .dampening(0.0) // Must be 0 for Nesterov
+      .weight_decay(0.05)
+      .nesterov(true);
+
+  SGD optimizer(param_groups, defaults);
+
+  // Check Group 1: lr and weight_decay preserved, others inherited
+  auto& group1_opts =
+      static_cast<SGDOptions&>(optimizer.param_groups()[0].options());
+  ASSERT_EQ(group1_opts.lr(), 0.01); // Preserved
+  ASSERT_EQ(group1_opts.momentum(), 0.9); // Inherited
+  ASSERT_EQ(group1_opts.dampening(), 0.0); // Inherited
+  ASSERT_EQ(group1_opts.weight_decay(), 0.22); // Preserved
+  ASSERT_TRUE(group1_opts.nesterov()); // Inherited
+
+  // Check Group 2: lr preserved, others inherited
+  auto& group2_opts =
+      static_cast<SGDOptions&>(optimizer.param_groups()[1].options());
+  ASSERT_EQ(group2_opts.lr(), 0.02); // Preserved
+  ASSERT_EQ(group2_opts.momentum(), 0.9); // Inherited
+  ASSERT_EQ(group2_opts.dampening(), 0.0); // Inherited
+  ASSERT_EQ(group2_opts.weight_decay(), 0.05); // Inherited
+  ASSERT_TRUE(group2_opts.nesterov()); // Inherited
+}
+
+TEST(OptimTest, MergeWithDefaultOptions_AdamW) {
+  // Create tensors for parameter groups
+  auto tensor1 = torch::randn({2, 2}).requires_grad_(true);
+  auto tensor2 = torch::randn({3, 3}).requires_grad_(true);
+
+  // Create param groups with partial options
+  std::vector<OptimizerParamGroup> param_groups;
+
+  // Group 1: Only eps specified, should inherit others
+  param_groups.emplace_back(
+      std::vector<torch::Tensor>{tensor1},
+      std::make_unique<AdamWOptions>(AdamWOptions().eps(1e-6)));
+
+  // Group 2: Only betas specified, should inherit others
+  param_groups.emplace_back(
+      std::vector<torch::Tensor>{tensor2},
+      std::make_unique<AdamWOptions>(
+          AdamWOptions().betas(std::make_tuple(0.95, 0.999))));
+
+  // Create optimizer with specific defaults
+  AdamWOptions defaults;
+  defaults.lr(0.003)
+      .betas(std::make_tuple(0.9, 0.98))
+      .eps(1e-8)
+      .weight_decay(0.02)
+      .amsgrad(false);
+
+  AdamW optimizer(param_groups, defaults);
+
+  // Check Group 1: eps preserved, others inherited
+  auto& group1_opts =
+      static_cast<AdamWOptions&>(optimizer.param_groups()[0].options());
+  ASSERT_EQ(group1_opts.lr(), 0.003); // Inherited
+  ASSERT_EQ(group1_opts.betas(), std::make_tuple(0.9, 0.98)); // Inherited
+  ASSERT_EQ(group1_opts.eps(), 1e-6); // Preserved
+  ASSERT_EQ(group1_opts.weight_decay(), 0.02); // Inherited
+  ASSERT_FALSE(group1_opts.amsgrad()); // Inherited
+
+  // Check Group 2: betas preserved, others inherited
+  auto& group2_opts =
+      static_cast<AdamWOptions&>(optimizer.param_groups()[1].options());
+  ASSERT_EQ(group2_opts.lr(), 0.003); // Inherited
+  ASSERT_EQ(group2_opts.betas(), std::make_tuple(0.95, 0.999)); // Preserved
+  ASSERT_EQ(group2_opts.eps(), 1e-8); // Inherited
+  ASSERT_EQ(group2_opts.weight_decay(), 0.02); // Inherited
+  ASSERT_FALSE(group2_opts.amsgrad()); // Inherited
+}
+
+TEST(OptimTest, MergeWithDefaultOptions_Adagrad) {
+  // Create tensors for parameter groups
+  auto tensor1 = torch::randn({2, 2}).requires_grad_(true);
+  auto tensor2 = torch::randn({3, 3}).requires_grad_(true);
+
+  // Create param groups with partial options
+  std::vector<OptimizerParamGroup> param_groups;
+
+  // Group 1: Only lr_decay specified, should inherit others
+  param_groups.emplace_back(
+      std::vector<torch::Tensor>{tensor1},
+      std::make_unique<AdagradOptions>(AdagradOptions().lr_decay(0.001)));
+
+  // Group 2: Only initial_accumulator_value specified, should inherit others
+  param_groups.emplace_back(
+      std::vector<torch::Tensor>{tensor2},
+      std::make_unique<AdagradOptions>(
+          AdagradOptions().initial_accumulator_value(0.5)));
+
+  // Create optimizer with specific defaults
+  AdagradOptions defaults;
+  defaults.lr(0.04)
+      .lr_decay(0.002)
+      .weight_decay(0.03)
+      .initial_accumulator_value(0.1)
+      .eps(1e-11);
+
+  Adagrad optimizer(param_groups, defaults);
+
+  // Check Group 1: lr_decay preserved, others inherited
+  auto& group1_opts =
+      static_cast<AdagradOptions&>(optimizer.param_groups()[0].options());
+  ASSERT_EQ(group1_opts.lr(), 0.04); // Inherited
+  ASSERT_EQ(group1_opts.lr_decay(), 0.001); // Preserved
+  ASSERT_EQ(group1_opts.weight_decay(), 0.03); // Inherited
+  ASSERT_EQ(group1_opts.initial_accumulator_value(), 0.1); // Inherited
+  ASSERT_EQ(group1_opts.eps(), 1e-11); // Inherited
+
+  // Check Group 2: initial_accumulator_value preserved, others inherited
+  auto& group2_opts =
+      static_cast<AdagradOptions&>(optimizer.param_groups()[1].options());
+  ASSERT_EQ(group2_opts.lr(), 0.04); // Inherited
+  ASSERT_EQ(group2_opts.lr_decay(), 0.002); // Inherited
+  ASSERT_EQ(group2_opts.weight_decay(), 0.03); // Inherited
+  ASSERT_EQ(group2_opts.initial_accumulator_value(), 0.5); // Preserved
+  ASSERT_EQ(group2_opts.eps(), 1e-11); // Inherited
+}
+
+TEST(OptimTest, MergeWithDefaultOptions_RMSprop) {
+  // Create tensors for parameter groups
+  auto tensor1 = torch::randn({2, 2}).requires_grad_(true);
+  auto tensor2 = torch::randn({3, 3}).requires_grad_(true);
+
+  // Create param groups with partial options
+  std::vector<OptimizerParamGroup> param_groups;
+
+  // Group 1: Only alpha specified, should inherit others
+  param_groups.emplace_back(
+      std::vector<torch::Tensor>{tensor1},
+      std::make_unique<RMSpropOptions>(RMSpropOptions().alpha(0.95)));
+
+  // Group 2: Only momentum and centered specified, should inherit others
+  param_groups.emplace_back(
+      std::vector<torch::Tensor>{tensor2},
+      std::make_unique<RMSpropOptions>(
+          RMSpropOptions().momentum(0.8).centered(true)));
+
+  // Create optimizer with specific defaults
+  RMSpropOptions defaults;
+  defaults.lr(0.015)
+      .alpha(0.98)
+      .eps(1e-9)
+      .weight_decay(0.01)
+      .momentum(0.7)
+      .centered(false);
+
+  RMSprop optimizer(param_groups, defaults);
+
+  // Check Group 1: alpha preserved, others inherited
+  auto& group1_opts =
+      static_cast<RMSpropOptions&>(optimizer.param_groups()[0].options());
+  ASSERT_EQ(group1_opts.lr(), 0.015); // Inherited
+  ASSERT_EQ(group1_opts.alpha(), 0.95); // Preserved
+  ASSERT_EQ(group1_opts.eps(), 1e-9); // Inherited
+  ASSERT_EQ(group1_opts.weight_decay(), 0.01); // Inherited
+  ASSERT_EQ(group1_opts.momentum(), 0.7); // Inherited
+  ASSERT_FALSE(group1_opts.centered()); // Inherited
+
+  // Check Group 2: momentum and centered preserved, others inherited
+  auto& group2_opts =
+      static_cast<RMSpropOptions&>(optimizer.param_groups()[1].options());
+  ASSERT_EQ(group2_opts.lr(), 0.015); // Inherited
+  ASSERT_EQ(group2_opts.alpha(), 0.98); // Inherited
+  ASSERT_EQ(group2_opts.eps(), 1e-9); // Inherited
+  ASSERT_EQ(group2_opts.weight_decay(), 0.01); // Inherited
+  ASSERT_EQ(group2_opts.momentum(), 0.8); // Preserved
+  ASSERT_TRUE(group2_opts.centered()); // Preserved
+}
+
+TEST(OptimTest, MergeWithDefaultOptions_LBFGS) {
+  // Create tensors for single parameter group (LBFGS limitation)
+  auto tensor1 = torch::randn({2, 2}).requires_grad_(true);
+  auto tensor2 = torch::randn({3, 3}).requires_grad_(true);
+
+  // Create param group with partial options
+  std::vector<OptimizerParamGroup> param_groups;
+
+  // Single group: Only max_iter specified, should inherit others
+  param_groups.emplace_back(
+      std::vector<torch::Tensor>{
+          tensor1, tensor2}, // Combine tensors in single group
+      std::make_unique<LBFGSOptions>(LBFGSOptions().max_iter(15)));
+
+  // Create optimizer with specific defaults
+  LBFGSOptions defaults;
+  defaults.lr(0.8)
+      .max_iter(25)
+      .max_eval(31) // Use same value that appears to be auto-calculated
+      .tolerance_grad(1e-5)
+      .tolerance_change(1e-8)
+      .history_size(80)
+      .line_search_fn("strong_wolfe");
+
+  LBFGS optimizer(param_groups, defaults);
+
+  // Check Group: max_iter preserved, others inherited
+  auto& group_opts =
+      static_cast<LBFGSOptions&>(optimizer.param_groups()[0].options());
+  ASSERT_EQ(group_opts.lr(), 0.8); // Inherited
+  ASSERT_EQ(group_opts.max_iter(), 15); // Preserved
+  ASSERT_EQ(group_opts.max_eval(), 31); // Inherited
+  ASSERT_EQ(group_opts.tolerance_grad(), 1e-5); // Inherited
+  ASSERT_EQ(group_opts.tolerance_change(), 1e-8); // Inherited
+  ASSERT_EQ(group_opts.history_size(), 80); // Inherited
+  ASSERT_EQ(group_opts.line_search_fn(), "strong_wolfe"); // Inherited
+}
+
+TEST(OptimTest, MergeWithDefaultOptions_NoOptionsInheritance) {
+  // Test that param groups without options get full defaults
+  auto tensor1 = torch::randn({2, 2}).requires_grad_(true);
+  auto tensor2 = torch::randn({3, 3}).requires_grad_(true);
+
+  std::vector<OptimizerParamGroup> param_groups;
+
+  // Groups with no options - should inherit everything
+  param_groups.emplace_back(std::vector<torch::Tensor>{tensor1});
+  param_groups.emplace_back(std::vector<torch::Tensor>{tensor2});
+
+  // Create optimizer with specific defaults
+  AdamOptions defaults;
+  defaults.lr(0.005)
+      .betas(std::make_tuple(0.85, 0.95))
+      .eps(1e-7)
+      .weight_decay(0.08)
+      .amsgrad(true);
+
+  Adam optimizer(param_groups, defaults);
+
+  // Both groups should have exactly the default options
+  for (int i = 0; i < 2; i++) {
+    auto& group_opts =
+        static_cast<AdamOptions&>(optimizer.param_groups()[i].options());
+    ASSERT_EQ(group_opts.lr(), 0.005);
+    ASSERT_EQ(group_opts.betas(), std::make_tuple(0.85, 0.95));
+    ASSERT_EQ(group_opts.eps(), 1e-7);
+    ASSERT_EQ(group_opts.weight_decay(), 0.08);
+    ASSERT_TRUE(group_opts.amsgrad());
+  }
+}
+
+// Test that field tracking survives serialization/deserialization cycles
+TEST(OptimTest, SerializationPreservesFieldTracking_Adam) {
+  // Create tensors for parameter groups
+  auto tensor1 = torch::randn({2, 2}).requires_grad_(true);
+  auto tensor2 = torch::randn({3, 3}).requires_grad_(true);
+
+  // Create param groups with partial options using fluent API (marks fields as
+  // explicit)
+  std::vector<OptimizerParamGroup> param_groups;
+
+  // Group 1: Only weight_decay and amsgrad explicitly set via fluent API
+  param_groups.emplace_back(
+      std::vector<torch::Tensor>{tensor1},
+      std::make_unique<AdamOptions>(
+          AdamOptions().weight_decay(0.11).amsgrad(true)));
+
+  // Group 2: Only eps explicitly set via fluent API
+  param_groups.emplace_back(
+      std::vector<torch::Tensor>{tensor2},
+      std::make_unique<AdamOptions>(AdamOptions().eps(1e-6)));
+
+  // Create optimizer with specific defaults
+  AdamOptions defaults;
+  defaults.lr(0.002)
+      .betas(std::make_tuple(0.8, 0.88))
+      .eps(1e-12)
+      .weight_decay(0.05)
+      .amsgrad(false);
+
+  Adam original_optimizer(param_groups, defaults);
+
+  // Capture original state for comparison
+  auto& orig_group1_opts =
+      static_cast<AdamOptions&>(original_optimizer.param_groups()[0].options());
+  auto& orig_group2_opts =
+      static_cast<AdamOptions&>(original_optimizer.param_groups()[1].options());
+
+  // Verify original state (sanity check)
+  ASSERT_NEAR(orig_group1_opts.weight_decay(), 0.11, 1e-6); // Explicitly set
+  ASSERT_TRUE(orig_group1_opts.amsgrad()); // Explicitly set
+  ASSERT_NEAR(orig_group1_opts.lr(), 0.002, 1e-6); // Inherited
+  ASSERT_NEAR(orig_group2_opts.eps(), 1e-6, 1e-9); // Explicitly set
+  ASSERT_NEAR(orig_group2_opts.lr(), 0.002, 1e-6); // Inherited
+
+  // Test serialization of the options objects (where field tracking lives)
+  std::stringstream ss1, ss2;
+
+  // Serialize the parameter group options
+  {
+    torch::serialize::OutputArchive archive;
+    orig_group1_opts.serialize(archive);
+    archive.save_to(ss1);
+  }
+  {
+    torch::serialize::OutputArchive archive;
+    orig_group2_opts.serialize(archive);
+    archive.save_to(ss2);
+  }
+
+  // Create new options objects and deserialize
+  AdamOptions loaded_group1_opts;
+  AdamOptions loaded_group2_opts;
+
+  {
+    torch::serialize::InputArchive archive;
+    archive.load_from(ss1);
+    loaded_group1_opts.serialize(archive);
+  }
+  {
+    torch::serialize::InputArchive archive;
+    archive.load_from(ss2);
+    loaded_group2_opts.serialize(archive);
+  }
+
+  // Verify that all parameter values are preserved after deserialization
+
+  // Group 1: weight_decay and amsgrad should be preserved as explicitly set,
+  // others inherited
+  ASSERT_NEAR(loaded_group1_opts.lr(), 0.002, 1e-6); // Inherited
+  ASSERT_EQ(
+      loaded_group1_opts.betas(), std::make_tuple(0.8, 0.88)); // Inherited
+  ASSERT_NEAR(loaded_group1_opts.eps(), 1e-12, 1e-15); // Inherited
+  ASSERT_NEAR(loaded_group1_opts.weight_decay(), 0.11, 1e-6); // Explicitly set
+  ASSERT_TRUE(loaded_group1_opts.amsgrad()); // Explicitly set
+
+  // Group 2: eps should be preserved as explicitly set, others inherited
+  ASSERT_NEAR(loaded_group2_opts.lr(), 0.002, 1e-6); // Inherited
+  ASSERT_EQ(
+      loaded_group2_opts.betas(), std::make_tuple(0.8, 0.88)); // Inherited
+  ASSERT_NEAR(loaded_group2_opts.eps(), 1e-6, 1e-9); // Explicitly set
+  ASSERT_NEAR(loaded_group2_opts.weight_decay(), 0.05, 1e-6); // Inherited
+  ASSERT_FALSE(loaded_group2_opts.amsgrad()); // Inherited
+
+  // CRITICAL: Test that field tracking is preserved after serialization
+  // Create a new optimizer using the deserialized options to test inheritance
+  auto tensor3 = torch::randn({2, 2}).requires_grad_(true);
+  auto tensor4 = torch::randn({3, 3}).requires_grad_(true);
+
+  std::vector<OptimizerParamGroup> test_param_groups;
+  test_param_groups.emplace_back(
+      std::vector<torch::Tensor>{tensor3},
+      std::make_unique<AdamOptions>(loaded_group1_opts));
+  test_param_groups.emplace_back(
+      std::vector<torch::Tensor>{tensor4},
+      std::make_unique<AdamOptions>(loaded_group2_opts));
+
+  Adam test_optimizer(test_param_groups, defaults);
+
+  // The field tracking should work correctly for inheritance
+  auto& final_group1_opts =
+      static_cast<AdamOptions&>(test_optimizer.param_groups()[0].options());
+  auto& final_group2_opts =
+      static_cast<AdamOptions&>(test_optimizer.param_groups()[1].options());
+
+  // Group 1: weight_decay and amsgrad should still be preserved as explicitly
+  // set
+  ASSERT_NEAR(
+      final_group1_opts.weight_decay(),
+      0.11,
+      1e-6); // Explicitly set (preserved)
+  ASSERT_TRUE(final_group1_opts.amsgrad()); // Explicitly set (preserved)
+  ASSERT_NEAR(final_group1_opts.lr(), 0.002, 1e-6); // Inherited from defaults
+
+  // Group 2: eps should still be preserved as explicitly set
+  ASSERT_NEAR(
+      final_group2_opts.eps(), 1e-6, 1e-9); // Explicitly set (preserved)
+  ASSERT_NEAR(final_group2_opts.lr(), 0.002, 1e-6); // Inherited from defaults
+}
+
+// Test serialization with SGD (different parameter types)
+TEST(OptimTest, SerializationPreservesFieldTracking_SGD) {
+  // Create tensors
+  auto tensor1 = torch::randn({2, 2}).requires_grad_(true);
+
+  // Create param group with partial options using fluent API
+  std::vector<OptimizerParamGroup> param_groups;
+  param_groups.emplace_back(
+      std::vector<torch::Tensor>{tensor1},
+      std::make_unique<SGDOptions>(
+          SGDOptions(0.01).weight_decay(0.22).nesterov(true)));
+
+  // Create optimizer with defaults
+  SGDOptions defaults(0.001);
+  defaults.momentum(0.9).dampening(0.0).weight_decay(0.05).nesterov(false);
+
+  SGD original_optimizer(param_groups, defaults);
+
+  // Test serialization of the SGD options (where field tracking lives)
+  auto& original_opts =
+      static_cast<SGDOptions&>(original_optimizer.param_groups()[0].options());
+
+  std::stringstream ss;
+  {
+    torch::serialize::OutputArchive archive;
+    original_opts.serialize(archive);
+    archive.save_to(ss);
+  }
+
+  SGDOptions loaded_opts(0.0); // Dummy initial value
+  {
+    torch::serialize::InputArchive archive;
+    archive.load_from(ss);
+    loaded_opts.serialize(archive);
+  }
+  ASSERT_NEAR(loaded_opts.lr(), 0.01, 1e-6); // Explicitly set
+  ASSERT_NEAR(loaded_opts.momentum(), 0.9, 1e-6); // Inherited
+  ASSERT_NEAR(loaded_opts.dampening(), 0.0, 1e-6); // Inherited
+  ASSERT_NEAR(loaded_opts.weight_decay(), 0.22, 1e-6); // Explicitly set
+  ASSERT_TRUE(loaded_opts.nesterov()); // Explicitly set
+
+  // Test that field tracking still works after deserialization by creating new
+  // optimizer
+  auto tensor2 = torch::randn({3, 3}).requires_grad_(true);
+  std::vector<OptimizerParamGroup> test_param_groups;
+  test_param_groups.emplace_back(
+      std::vector<torch::Tensor>{tensor2},
+      std::make_unique<SGDOptions>(loaded_opts));
+
+  SGD test_optimizer(test_param_groups, defaults);
+
+  auto& final_opts =
+      static_cast<SGDOptions&>(test_optimizer.param_groups()[0].options());
+  ASSERT_NEAR(final_opts.lr(), 0.01, 1e-6); // Explicitly set (preserved)
+  ASSERT_NEAR(
+      final_opts.weight_decay(), 0.22, 1e-6); // Explicitly set (preserved)
+  ASSERT_TRUE(final_opts.nesterov()); // Explicitly set (preserved)
+  ASSERT_NEAR(final_opts.momentum(), 0.9, 1e-6); // Inherited from defaults
+  ASSERT_NEAR(final_opts.dampening(), 0.0, 1e-6); // Inherited from defaults
+}
diff --git a/test/cpp/c10d/CMakeLists.txt b/test/cpp/c10d/CMakeLists.txt
index 285a5dd2a73e5..24ab667bf5928 100644
--- a/test/cpp/c10d/CMakeLists.txt
+++ b/test/cpp/c10d/CMakeLists.txt
@@ -3,6 +3,12 @@ if(USE_CUDA)
   target_include_directories(c10d_cuda_test PRIVATE $<BUILD_INTERFACE:${TORCH_SRC_DIR}/csrc/distributed>)
   target_link_libraries(c10d_cuda_test torch_cuda)
   add_dependencies(c10d_cuda_test torch_cuda)
+elseif(USE_ROCM)
+  set_source_files_properties(hip/HIPTest.hip PROPERTIES HIP_SOURCE_PROPERTY_FORMAT 1)
+  hip_add_library(c10d_hip_test hip/HIPTest.hip)
+  target_include_directories(c10d_hip_test PRIVATE $<BUILD_INTERFACE:${TORCH_SRC_DIR}/csrc/distributed>)
+  target_link_libraries(c10d_hip_test torch_hip)
+  add_dependencies(c10d_hip_test torch_hip)
 endif()
 
 function(c10d_add_test test_src)
@@ -17,6 +23,7 @@ function(c10d_add_test test_src)
   get_filename_component(test_name ${test_src} NAME_WE)
   add_executable(${test_name} "${test_src}")
   target_include_directories(${test_name} PRIVATE
+      $<BUILD_INTERFACE:${TORCH_SRC_DIR}/test/cpp/c10d/hip>
       $<BUILD_INTERFACE:${TORCH_SRC_DIR}/csrc/distributed>
       $<TARGET_PROPERTY:fmt::fmt-header-only,INTERFACE_INCLUDE_DIRECTORIES>
   )
@@ -32,11 +39,21 @@ function(c10d_add_test test_src)
   endif()
 endfunction()
 
-c10d_add_test(BackoffTest.cpp LINK_LIBRARIES torch_cpu gtest_main INSTALL_TEST OFF)
-c10d_add_test(FileStoreTest.cpp LINK_LIBRARIES torch_cpu gtest_main INSTALL_TEST ${INSTALL_TEST})
-c10d_add_test(TCPStoreTest.cpp LINK_LIBRARIES torch_cpu gtest_main INSTALL_TEST ${INSTALL_TEST})
-if(NOT WIN32)
-  c10d_add_test(HashStoreTest.cpp LINK_LIBRARIES torch_cpu gtest_main INSTALL_TEST ${INSTALL_TEST})
+if(USE_ROCM)
+  c10d_add_test(hip/BackoffTest.cpp LINK_LIBRARIES torch_cpu gtest_main INSTALL_TEST OFF)
+  c10d_add_test(hip/FileStoreTest.cpp LINK_LIBRARIES torch_cpu gtest_main INSTALL_TEST ${INSTALL_TEST})
+  c10d_add_test(hip/TCPStoreTest.cpp LINK_LIBRARIES torch_cpu gtest_main INSTALL_TEST ${INSTALL_TEST})
+  if(NOT WIN32)
+    c10d_add_test(hip/HashStoreTest.cpp LINK_LIBRARIES torch_cpu gtest_main INSTALL_TEST ${INSTALL_TEST})
+  endif()
+
+else()
+  c10d_add_test(BackoffTest.cpp LINK_LIBRARIES torch_cpu gtest_main INSTALL_TEST OFF)
+  c10d_add_test(FileStoreTest.cpp LINK_LIBRARIES torch_cpu gtest_main INSTALL_TEST ${INSTALL_TEST})
+  c10d_add_test(TCPStoreTest.cpp LINK_LIBRARIES torch_cpu gtest_main INSTALL_TEST ${INSTALL_TEST})
+  if(NOT WIN32)
+    c10d_add_test(HashStoreTest.cpp LINK_LIBRARIES torch_cpu gtest_main INSTALL_TEST ${INSTALL_TEST})
+  endif()
 endif()
 
 if(USE_CUDA)
@@ -52,9 +69,15 @@ if(USE_CUDA)
     c10d_add_test(
       ProcessGroupNCCLTest.cpp
       LINK_LIBRARIES torch_cpu c10d_cuda_test gtest_main __caffe2_nccl INSTALL_TEST ${INSTALL_TEST})
+    target_compile_definitions(ProcessGroupNCCLTest PRIVATE USE_CUDA USE_NCCL)
     c10d_add_test(
       ProcessGroupNCCLErrorsTest.cpp
       LINK_LIBRARIES torch_cpu c10d_cuda_test gtest_main __caffe2_nccl INSTALL_TEST ${INSTALL_TEST})
+    target_compile_definitions(ProcessGroupNCCLErrorsTest PRIVATE USE_CUDA USE_NCCL)
+    c10d_add_test(
+      NCCLDevCommManagerTest.cpp
+      LINK_LIBRARIES torch_cpu c10d_cuda_test gtest_main __caffe2_nccl INSTALL_TEST ${INSTALL_TEST})
+    target_compile_definitions(NCCLDevCommManagerTest PRIVATE USE_CUDA USE_NCCL)
     if(INSTALL_TEST)
       install(TARGETS c10d_cuda_test DESTINATION lib)
     endif()
@@ -71,6 +94,30 @@ if(USE_CUDA)
       install(TARGETS c10d_cuda_test DESTINATION lib)
     endif()
   endif()
+elseif(USE_ROCM)
+  if(USE_GLOO AND USE_C10D_GLOO)
+    c10d_add_test(hip/ProcessGroupGlooTest.cpp LINK_LIBRARIES torch_cpu c10d_hip_test gtest_main INSTALL_TEST ${INSTALL_TEST})
+    c10d_add_test(hip/ProcessGroupGlooAsyncTest.cpp LINK_LIBRARIES torch_cpu c10d_hip_test gtest_main INSTALL_TEST ${INSTALL_TEST})
+  endif()
+  if(USE_NCCL AND USE_C10D_NCCL)
+    c10d_add_test(
+      hip/ProcessGroupNCCLTest.cpp
+      LINK_LIBRARIES torch_cpu c10d_hip_test gtest_main __caffe2_nccl INSTALL_TEST ${INSTALL_TEST})
+    c10d_add_test(
+      hip/ProcessGroupNCCLErrorsTest.cpp
+      LINK_LIBRARIES torch_cpu c10d_hip_test gtest_main __caffe2_nccl INSTALL_TEST ${INSTALL_TEST})
+    if(INSTALL_TEST)
+      install(TARGETS c10d_hip_test DESTINATION lib)
+    endif()
+  endif()
+  if(USE_UCC AND USE_C10D_UCC)
+    c10d_add_test(
+      hip/ProcessGroupUCCTest.cpp
+      LINK_LIBRARIES torch_cpu c10d_hip_test gtest_main __caffe2_ucc INSTALL_TEST ${INSTALL_TEST})
+    if(INSTALL_TEST)
+      install(TARGETS c10d_hip_test DESTINATION lib)
+    endif()
+  endif()
 else()
   if(USE_GLOO AND USE_C10D_GLOO)
     c10d_add_test(ProcessGroupGlooTest.cpp LINK_LIBRARIES torch_cpu gtest_main INSTALL_TEST OFF)
@@ -92,5 +139,7 @@ if(LINUX AND USE_GLOO AND USE_C10D_GLOO)
   target_link_libraries(example_allreduce torch_cpu)
   if(USE_CUDA)
     target_link_libraries(example_allreduce torch_cuda)
+  elseif(USE_ROCM)
+    target_link_libraries(example_allreduce torch_hip)
   endif()
 endif()
diff --git a/test/cpp/c10d/NCCLDevCommManagerTest.cpp b/test/cpp/c10d/NCCLDevCommManagerTest.cpp
new file mode 100644
index 0000000000000..ded19d266cf86
--- /dev/null
+++ b/test/cpp/c10d/NCCLDevCommManagerTest.cpp
@@ -0,0 +1,277 @@
+#include <gtest/gtest.h>
+
+// Include nccl_dev_cap.hpp first to define NCCL_HAS_SYMMEM_DEVICE_SUPPORT
+#include <torch/csrc/distributed/c10d/symm_mem/nccl_dev_cap.hpp>
+
+#ifdef NCCL_HAS_SYMMEM_DEVICE_SUPPORT
+
+#include <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAGuard.h>
+#include <c10/cuda/CUDAStream.h>
+#include <c10/util/Exception.h>
+#include <torch/csrc/distributed/c10d/symm_mem/nccl_devcomm_manager.hpp>
+
+using namespace c10d::symmetric_memory;
+
+// Note: This test requires NCCL with symmetric memory device support.
+// The ncclDevComm type is defined in nccl_device.h (included via
+// nccl_dev_cap.hpp). For testing without full NCCL initialization, we use
+// uninitialized values. In a real scenario, you would use ncclDevCommCreate to
+// create valid communicators.
+
+// Test register_comm and get_comm
+void test_register_and_get_comm() {
+  const std::string group_name = "test_group_comm";
+  c10::Device device(c10::DeviceType::CUDA, 0);
+
+  auto& manager = NCCLDevCommManager::get(device);
+
+  // Register a host communicator
+  ncclComm_t mock_comm = nullptr; // In real test, this would be a valid comm
+  manager.register_comm(group_name, mock_comm);
+
+  // Get the communicator back
+  ncclComm_t retrieved_comm = manager.get_comm(group_name);
+  EXPECT_EQ(retrieved_comm, mock_comm);
+
+  // Registering the same communicator again should not throw
+  manager.register_comm(group_name, mock_comm);
+
+  // Getting from non-existent group should throw
+  EXPECT_THROW(manager.get_comm("nonexistent_group"), c10::Error);
+}
+
+// Test basic registration and retrieval with default function name key
+void test_register_and_get_default_key() {
+  const std::string group_name = "test_group";
+  c10::Device device(c10::DeviceType::CUDA, 0);
+
+  auto& manager = NCCLDevCommManager::get(device);
+
+  // Create a test device communicator (uninitialized for testing)
+  // In real usage, this would be created via ncclDevCommCreate
+  ncclDevComm devcomm0 = {};
+
+  // Register with default key (function name)
+  auto devcomm_opt = manager.register_devcomm(group_name, devcomm0);
+  EXPECT_TRUE(devcomm_opt.has_value());
+
+  // Get with default key
+  auto retrieved_opt = manager.get_devcomm(group_name);
+  EXPECT_TRUE(retrieved_opt.has_value());
+  // Note: We can't easily compare ncclDevComm values without proper
+  // initialization This test verifies the registration/retrieval logic works
+}
+
+// Test registration and retrieval with custom key
+void test_register_and_get_custom_key() {
+  const std::string group_name = "test_group_custom";
+  c10::Device device(c10::DeviceType::CUDA, 0);
+
+  auto& manager = NCCLDevCommManager::get(device);
+
+  // Create test device communicators (uninitialized for testing)
+  ncclDevComm devcomm0 = {};
+  ncclDevComm devcomm1 = {};
+
+  // Register with custom keys
+  auto devcomm0_opt =
+      manager.register_devcomm(group_name, devcomm0, "custom_key0");
+  EXPECT_TRUE(devcomm0_opt.has_value());
+
+  auto devcomm1_opt =
+      manager.register_devcomm(group_name, devcomm1, "custom_key1");
+  EXPECT_TRUE(devcomm1_opt.has_value());
+
+  // Get with custom keys
+  auto retrieved0_opt = manager.get_devcomm(group_name, "custom_key0");
+  EXPECT_TRUE(retrieved0_opt.has_value());
+
+  auto retrieved1_opt = manager.get_devcomm(group_name, "custom_key1");
+  EXPECT_TRUE(retrieved1_opt.has_value());
+}
+
+// Test that registering the same key twice throws an error
+void test_duplicate_registration_throws() {
+  const std::string group_name = "test_group_duplicate";
+  c10::Device device(c10::DeviceType::CUDA, 0);
+
+  auto& manager = NCCLDevCommManager::get(device);
+
+  // Create test device communicators (uninitialized for testing)
+  ncclDevComm devcomm0 = {};
+  ncclDevComm devcomm1 = {};
+
+  // Register first time - should succeed
+  auto devcomm0_opt =
+      manager.register_devcomm(group_name, devcomm0, "duplicate_key");
+  EXPECT_TRUE(devcomm0_opt.has_value());
+
+  // Register second time with same key - should throw
+  EXPECT_THROW(
+      manager.register_devcomm(group_name, devcomm1, "duplicate_key"),
+      c10::Error);
+}
+
+// Test that getting a non-existent key returns nullopt
+void test_get_nonexistent_key() {
+  const std::string group_name = "test_group_nonexistent";
+  c10::Device device(c10::DeviceType::CUDA, 0);
+
+  auto& manager = NCCLDevCommManager::get(device);
+
+  // Try to get a key that doesn't exist
+  auto retrieved_opt = manager.get_devcomm(group_name, "nonexistent_key");
+  EXPECT_FALSE(retrieved_opt.has_value());
+
+  // Try to get from a group that doesn't exist
+  auto retrieved_group_opt =
+      manager.get_devcomm("nonexistent_group", "any_key");
+  EXPECT_FALSE(retrieved_group_opt.has_value());
+}
+
+// Test multiple groups with different keys
+void test_multiple_groups() {
+  const std::string group1 = "group1";
+  const std::string group2 = "group2";
+  c10::Device device(c10::DeviceType::CUDA, 0);
+
+  auto& manager = NCCLDevCommManager::get(device);
+
+  // Create test device communicators (uninitialized for testing)
+  ncclDevComm devcomm1a = {};
+  ncclDevComm devcomm1b = {};
+  ncclDevComm devcomm2a = {};
+
+  // Register in group1
+  auto devcomm1a_opt = manager.register_devcomm(group1, devcomm1a, "key1a");
+  EXPECT_TRUE(devcomm1a_opt.has_value());
+  auto devcomm1b_opt = manager.register_devcomm(group1, devcomm1b, "key1b");
+  EXPECT_TRUE(devcomm1b_opt.has_value());
+
+  // Register in group2
+  auto devcomm2a_opt = manager.register_devcomm(group2, devcomm2a, "key2a");
+  EXPECT_TRUE(devcomm2a_opt.has_value());
+
+  // Verify we can retrieve from both groups
+  auto retrieved1a_opt = manager.get_devcomm(group1, "key1a");
+  EXPECT_TRUE(retrieved1a_opt.has_value());
+
+  auto retrieved1b_opt = manager.get_devcomm(group1, "key1b");
+  EXPECT_TRUE(retrieved1b_opt.has_value());
+
+  auto retrieved2a_opt = manager.get_devcomm(group2, "key2a");
+  EXPECT_TRUE(retrieved2a_opt.has_value());
+
+  // Verify keys from one group don't affect the other
+  auto cross_group_opt = manager.get_devcomm(group1, "key2a");
+  EXPECT_FALSE(cross_group_opt.has_value());
+}
+
+// Test the example pattern from the comments
+void test_example_pattern() {
+  const std::string group_name = "example_group";
+  c10::Device device(c10::DeviceType::CUDA, 0);
+
+  auto& manager = NCCLDevCommManager::get(device);
+
+  // Example: foo function pattern
+  auto foo = [&manager, &group_name]() {
+    // Try to get first
+    auto devcomm_opt = manager.get_devcomm(group_name);
+    if (!devcomm_opt) {
+      // Not found, create then register
+      ncclDevComm devcomm = {}; // In real code: ncclDevCommCreate(...)
+      devcomm_opt = manager.register_devcomm(group_name, devcomm);
+    }
+    ncclDevComm& devcomm_ref = devcomm_opt->get();
+    (void)devcomm_ref; // In real code, used for NCCL operations
+    EXPECT_TRUE(devcomm_opt.has_value());
+  };
+
+  foo();
+
+  // Example: bar function pattern with multiple keys
+  auto bar = [&manager, &group_name]() {
+    ncclDevComm devcomm0 = {}; // In real code: ncclDevCommCreate(...)
+    ncclDevComm devcomm1 = {}; // In real code: ncclDevCommCreate(...)
+    // Register with custom keys
+    manager.register_devcomm(group_name, devcomm0, "bar0");
+    manager.register_devcomm(group_name, devcomm1, "bar1");
+
+    // Verify both are registered
+    auto bar0_opt = manager.get_devcomm(group_name, "bar0");
+    EXPECT_TRUE(bar0_opt.has_value());
+
+    auto bar1_opt = manager.get_devcomm(group_name, "bar1");
+    EXPECT_TRUE(bar1_opt.has_value());
+  };
+
+  bar();
+}
+
+// Note: This test file requires CUDA and NCCL to be available.
+// The actual ncclDevComm type would be used in real scenarios.
+// For unit testing, we're using a simple integer as a placeholder.
+// In a real test environment, you would need to:
+// 1. Initialize CUDA
+// 2. Initialize NCCL
+// 3. Create actual ncclDevComm objects using ncclDevCommCreate
+// 4. Clean up properly
+
+TEST(NCCLDevCommManagerTest, RegisterAndGetComm) {
+  if (!at::cuda::is_available()) {
+    GTEST_SKIP() << "CUDA not available, skipping test";
+  }
+  test_register_and_get_comm();
+}
+
+TEST(NCCLDevCommManagerTest, RegisterAndGetDefaultKey) {
+  if (!at::cuda::is_available()) {
+    GTEST_SKIP() << "CUDA not available, skipping test";
+  }
+  test_register_and_get_default_key();
+}
+
+TEST(NCCLDevCommManagerTest, RegisterAndGetCustomKey) {
+  if (!at::cuda::is_available()) {
+    GTEST_SKIP() << "CUDA not available, skipping test";
+  }
+  test_register_and_get_custom_key();
+}
+
+TEST(NCCLDevCommManagerTest, DuplicateRegistrationThrows) {
+  if (!at::cuda::is_available()) {
+    GTEST_SKIP() << "CUDA not available, skipping test";
+  }
+  test_duplicate_registration_throws();
+}
+
+TEST(NCCLDevCommManagerTest, GetNonexistentKey) {
+  if (!at::cuda::is_available()) {
+    GTEST_SKIP() << "CUDA not available, skipping test";
+  }
+  test_get_nonexistent_key();
+}
+
+TEST(NCCLDevCommManagerTest, MultipleGroups) {
+  if (!at::cuda::is_available()) {
+    GTEST_SKIP() << "CUDA not available, skipping test";
+  }
+  test_multiple_groups();
+}
+
+TEST(NCCLDevCommManagerTest, ExamplePattern) {
+  if (!at::cuda::is_available()) {
+    GTEST_SKIP() << "CUDA not available, skipping test";
+  }
+  test_example_pattern();
+}
+
+#else // NCCL_HAS_SYMMEM_DEVICE_SUPPORT
+
+TEST(NCCLDevCommManagerTest, Skipped) {
+  GTEST_SKIP() << "NCCL symmetric memory device support not available";
+}
+
+#endif // NCCL_HAS_SYMMEM_DEVICE_SUPPORT
diff --git a/test/cpp/c10d/TCPStoreTest.cpp b/test/cpp/c10d/TCPStoreTest.cpp
index cdd0d58e61c7f..019acf022a21d 100644
--- a/test/cpp/c10d/TCPStoreTest.cpp
+++ b/test/cpp/c10d/TCPStoreTest.cpp
@@ -320,3 +320,48 @@ TEST(TCPStoreTest, testMultiTenantStores) {
 TEST(TCPStoreTest, testMultiTenantStoresUV) {
   testMultiTenantStores(true);
 }
+
+void testBarrier(bool useLibUV) {
+  constexpr int numWorkers = 4;
+  auto serverTCPStore = _createServer(useLibUV, numWorkers);
+
+  std::vector<std::thread> threads;
+  c10d::TCPStoreOptions client_opts{};
+  client_opts.port = serverTCPStore->getPort();
+  client_opts.numWorkers = numWorkers;
+  client_opts.useLibUV = useLibUV;
+
+  for (int i = 0; i < numWorkers; ++i) {
+    threads.emplace_back([=, &client_opts] {
+      auto store =
+          c10::make_intrusive<c10d::TCPStore>("127.0.0.1", client_opts);
+      store->barrier(
+          "test_barrier", numWorkers, std::chrono::seconds(defaultTimeout));
+    });
+  }
+  for (auto& t : threads) {
+    t.join();
+  }
+}
+
+TEST(TCPStoreTest, testBarrier) {
+  testBarrier(false);
+}
+
+TEST(TCPStoreTest, testBarrierUV) {
+  testBarrier(true);
+}
+
+void testBarrierTimeout(bool useLibUV) {
+  auto store = _createServer(useLibUV, 1);
+  auto timeout = std::chrono::milliseconds(100);
+  EXPECT_THROW(store->barrier("timeout_key", 2, timeout), c10::DistStoreError);
+}
+
+TEST(TCPStoreTest, testBarrierTimeout) {
+  testBarrierTimeout(false);
+}
+
+TEST(TCPStoreTest, testBarrierTimeoutUV) {
+  testBarrierTimeout(true);
+}
diff --git a/test/cpp/jit/CMakeLists.txt b/test/cpp/jit/CMakeLists.txt
index 0b2a06b53c9a2..1204ad0184aee 100644
--- a/test/cpp/jit/CMakeLists.txt
+++ b/test/cpp/jit/CMakeLists.txt
@@ -158,6 +158,7 @@ endif()
 if(INSTALL_TEST)
   set_target_properties(test_jit PROPERTIES INSTALL_RPATH "${CMAKE_INSTALL_RPATH}:${_rpath_portable_origin}/../lib")
   install(TARGETS test_jit DESTINATION bin)
+  install(FILES ${JIT_TEST_ROOT}/test_interpreter_async.pt DESTINATION bin)
   # Install PDB files for MSVC builds
   if(MSVC AND BUILD_SHARED_LIBS)
     install(FILES $<TARGET_PDB_FILE:test_jit> DESTINATION bin OPTIONAL)
diff --git a/test/cpp/jit/test_graph_executor.cpp b/test/cpp/jit/test_graph_executor.cpp
index 82436f5ce46a6..d37a2c0e87654 100644
--- a/test/cpp/jit/test_graph_executor.cpp
+++ b/test/cpp/jit/test_graph_executor.cpp
@@ -1,5 +1,7 @@
 #include <gtest/gtest.h>
 
+#include <filesystem>
+
 #include "test/cpp/jit/test_utils.h"
 #include "torch/csrc/jit/runtime/graph_executor.h"
 #include "torch/jit.h"
@@ -45,9 +47,16 @@ TEST(GraphExecutorTest, runAsync_executor) {
   demo = DemoModule()
   torch.jit.save(torch.jit.script(demo), 'test_interpreter_async.pt')
   */
-  std::string filePath(__FILE__);
-  auto testModelFile = filePath.substr(0, filePath.find_last_of("/\\") + 1);
-  testModelFile.append("test_interpreter_async.pt");
+  auto testModelFile = []() -> std::string {
+    std::string dir(__FILE__);
+    dir = dir.substr(0, dir.find_last_of("/\\") + 1);
+    auto candidate = dir + "test_interpreter_async.pt";
+    if (std::filesystem::exists(candidate))
+      return candidate;
+    // Installed binary: .pt is next to the executable.
+    auto exeDir = std::filesystem::read_symlink("/proc/self/exe").parent_path();
+    return (exeDir / "test_interpreter_async.pt").string();
+  }();
   auto module = load(testModelFile);
   auto graph = module.get_method("forward").graph();
   GraphExecutor graphExecutor(graph, "");
diff --git a/test/cpp/jit/test_interpreter.cpp b/test/cpp/jit/test_interpreter.cpp
index 78174d8515701..ba70e6a5d3694 100644
--- a/test/cpp/jit/test_interpreter.cpp
+++ b/test/cpp/jit/test_interpreter.cpp
@@ -1,6 +1,8 @@
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
 
+#include <filesystem>
+
 #include <ATen/Parallel.h>
 #include <c10/core/DeviceType.h>
 #include <test/cpp/jit/test_utils.h>
@@ -198,9 +200,15 @@ TEST(InterpreterTest, runAsyncBasicTest) {
   demo = DemoModule()
   torch.jit.save(torch.jit.script(demo), 'test_interpreter_async.pt')
   */
-  std::string filePath(__FILE__);
-  auto testModelFile = filePath.substr(0, filePath.find_last_of("/\\") + 1);
-  testModelFile.append("test_interpreter_async.pt");
+  auto testModelFile = []() -> std::string {
+    std::string dir(__FILE__);
+    dir = dir.substr(0, dir.find_last_of("/\\") + 1);
+    auto candidate = dir + "test_interpreter_async.pt";
+    if (std::filesystem::exists(candidate))
+      return candidate;
+    auto exeDir = std::filesystem::read_symlink("/proc/self/exe").parent_path();
+    return (exeDir / "test_interpreter_async.pt").string();
+  }();
   auto model = load(testModelFile);
   auto graph = model.get_method("forward").graph();
   Code function(graph, "");
diff --git a/test/cpp/jit/test_misc.cpp b/test/cpp/jit/test_misc.cpp
index 5cf409c6d7fcc..7610053e18397 100644
--- a/test/cpp/jit/test_misc.cpp
+++ b/test/cpp/jit/test_misc.cpp
@@ -3157,6 +3157,11 @@ TEST_F(Composed, ComposedOp) {
 }
 
 TEST(ConstantPropagation, CustomClassesCanBePropagated) {
+#if defined(__aarch64__) || defined(_M_ARM64)
+  // See https://github.com/pytorch/pytorch/issues/178522.
+  GTEST_SKIP()
+      << "Skipping ConstantPropagation.CustomClassesCanBePropagated on AArch64.";
+#endif
 #ifdef USE_PYTORCH_QNNPACK
   const auto src = R"IR(
     graph():
diff --git a/test/cpp/jit/test_module_api.cpp b/test/cpp/jit/test_module_api.cpp
index 2300014ecebd4..8b6a9dc2dacda 100644
--- a/test/cpp/jit/test_module_api.cpp
+++ b/test/cpp/jit/test_module_api.cpp
@@ -1,5 +1,7 @@
 #include <gtest/gtest.h>
 
+#include <filesystem>
+
 #include <test/cpp/jit/test_utils.h>
 
 #include <ATen/core/qualified_name.h>
@@ -53,10 +55,15 @@ TEST(ModuleAPITest, MethodRunAsync) {
   //     r2 = torch.jit.fork(torch.mm, torch.rand(100,100),torch.rand(100,100))
   //     return r1.wait() + r2.wait()
   // )");
-  std::string filePath(__FILE__);
-  auto testModelFile = filePath.substr(0, filePath.find_last_of("/\\") + 1);
-  // borrow model file from TEST(GraphExecutorTest, runAsync_executor)
-  testModelFile.append("test_interpreter_async.pt");
+  auto testModelFile = []() -> std::string {
+    std::string dir(__FILE__);
+    dir = dir.substr(0, dir.find_last_of("/\\") + 1);
+    auto candidate = dir + "test_interpreter_async.pt";
+    if (std::filesystem::exists(candidate))
+      return candidate;
+    auto exeDir = std::filesystem::read_symlink("/proc/self/exe").parent_path();
+    return (exeDir / "test_interpreter_async.pt").string();
+  }();
   auto m = load(testModelFile);
 
   auto counter = 0;
diff --git a/test/cpp/nativert/test_graph_signature.cpp b/test/cpp/nativert/test_graph_signature.cpp
index c57aebb5082c3..53684bb73360b 100644
--- a/test/cpp/nativert/test_graph_signature.cpp
+++ b/test/cpp/nativert/test_graph_signature.cpp
@@ -51,6 +51,49 @@ class GraphSignatureTest : public ::testing::Test {
   }
 };
 
+// Helper function to create a tensor argument
+torch::_export::TensorArgument makeTensorArg(const std::string& name) {
+  torch::_export::TensorArgument arg;
+  arg.set_name(name);
+  return arg;
+}
+
+// Helper function to create a SymInt argument with a name
+torch::_export::SymIntArgument makeSymIntArg(const std::string& name) {
+  torch::_export::SymIntArgument arg;
+  arg.set_as_name(name);
+  return arg;
+}
+
+// Helper function to create a SymInt argument with a constant int
+torch::_export::SymIntArgument makeSymIntArgConst(int64_t value) {
+  torch::_export::SymIntArgument arg;
+  arg.set_as_int(value);
+  return arg;
+}
+
+// Helper function to create a SymBool argument with a name
+torch::_export::SymBoolArgument makeSymBoolArg(const std::string& name) {
+  torch::_export::SymBoolArgument arg;
+  arg.set_as_name(name);
+  return arg;
+}
+
+// Helper function to create an OptionalTensorArgument with a tensor
+torch::_export::OptionalTensorArgument makeOptionalTensorArg(
+    const std::string& name) {
+  torch::_export::OptionalTensorArgument arg;
+  arg.set_as_tensor(makeTensorArg(name));
+  return arg;
+}
+
+// Helper function to create an OptionalTensorArgument with None
+torch::_export::OptionalTensorArgument makeOptionalTensorArgNone() {
+  torch::_export::OptionalTensorArgument arg;
+  arg.set_as_none({});
+  return arg;
+}
+
 // Test the constructor with a simple GraphSignature
 TEST_F(GraphSignatureTest, ConstructorTest) {
   std::vector<std::string_view> expected_params = {"param"};
@@ -74,4 +117,468 @@ TEST_F(GraphSignatureTest, ReplaceAllUsesTest) {
   EXPECT_EQ(graph_sig.userOutputs(), expected_outputs);
 }
 
+// Test AS_TENSORS user input handling - each tensor in the list should become
+// a separate entry in userInputs
+TEST(GraphSignatureInputTypesTest, UserInputAsTensors) {
+  // Create list of tensor arguments
+  std::vector<torch::_export::TensorArgument> tensor_list;
+  tensor_list.push_back(makeTensorArg("tensor_0"));
+  tensor_list.push_back(makeTensorArg("tensor_1"));
+  tensor_list.push_back(makeTensorArg("tensor_2"));
+
+  torch::_export::Argument input_arg;
+  input_arg.set_as_tensors(tensor_list);
+  torch::_export::UserInputSpec user_input_spec;
+  user_input_spec.set_arg(input_arg);
+  torch::_export::InputSpec input_spec;
+  input_spec.set_user_input(user_input_spec);
+
+  torch::_export::GraphSignature storage;
+  storage.set_input_specs({input_spec});
+  storage.set_output_specs({});
+
+  GraphSignature sig(storage);
+
+  // All tensors in the list should be expanded to individual user inputs
+  std::vector<std::string> expected_inputs = {
+      "tensor_0", "tensor_1", "tensor_2"};
+  EXPECT_EQ(sig.userInputs(), expected_inputs);
+}
+
+// Test AS_OPTIONAL_TENSOR user input handling
+TEST(GraphSignatureInputTypesTest, UserInputAsOptionalTensor) {
+  // Create an optional tensor that has a value
+  torch::_export::Argument input_arg;
+  input_arg.set_as_optional_tensor(makeOptionalTensorArg("opt_tensor"));
+  torch::_export::UserInputSpec user_input_spec;
+  user_input_spec.set_arg(input_arg);
+  torch::_export::InputSpec input_spec;
+  input_spec.set_user_input(user_input_spec);
+
+  torch::_export::GraphSignature storage;
+  storage.set_input_specs({input_spec});
+  storage.set_output_specs({});
+
+  GraphSignature sig(storage);
+
+  std::vector<std::string> expected_inputs = {"opt_tensor"};
+  EXPECT_EQ(sig.userInputs(), expected_inputs);
+}
+
+// Test AS_OPTIONAL_TENSOR with None value - should be skipped
+TEST(GraphSignatureInputTypesTest, UserInputAsOptionalTensorNone) {
+  torch::_export::Argument input_arg;
+  input_arg.set_as_optional_tensor(makeOptionalTensorArgNone());
+  torch::_export::UserInputSpec user_input_spec;
+  user_input_spec.set_arg(input_arg);
+  torch::_export::InputSpec input_spec;
+  input_spec.set_user_input(user_input_spec);
+
+  torch::_export::GraphSignature storage;
+  storage.set_input_specs({input_spec});
+  storage.set_output_specs({});
+
+  GraphSignature sig(storage);
+
+  // None optional tensor should be skipped
+  EXPECT_TRUE(sig.userInputs().empty());
+}
+
+// Test AS_OPTIONAL_TENSORS user input handling - mixed tensors and None values
+TEST(GraphSignatureInputTypesTest, UserInputAsOptionalTensors) {
+  std::vector<torch::_export::OptionalTensorArgument> opt_tensors;
+  opt_tensors.push_back(makeOptionalTensorArg("opt_tensor_0"));
+  opt_tensors.push_back(
+      makeOptionalTensorArgNone()); // None - should be skipped
+  opt_tensors.push_back(makeOptionalTensorArg("opt_tensor_2"));
+
+  torch::_export::Argument input_arg;
+  input_arg.set_as_optional_tensors(opt_tensors);
+  torch::_export::UserInputSpec user_input_spec;
+  user_input_spec.set_arg(input_arg);
+  torch::_export::InputSpec input_spec;
+  input_spec.set_user_input(user_input_spec);
+
+  torch::_export::GraphSignature storage;
+  storage.set_input_specs({input_spec});
+  storage.set_output_specs({});
+
+  GraphSignature sig(storage);
+
+  // Only non-None tensors should be in user inputs
+  std::vector<std::string> expected_inputs = {"opt_tensor_0", "opt_tensor_2"};
+  EXPECT_EQ(sig.userInputs(), expected_inputs);
+}
+
+// Test AS_SYM_INT user input handling
+TEST(GraphSignatureInputTypesTest, UserInputAsSymInt) {
+  torch::_export::Argument input_arg;
+  input_arg.set_as_sym_int(makeSymIntArg("sym_int_0"));
+  torch::_export::UserInputSpec user_input_spec;
+  user_input_spec.set_arg(input_arg);
+  torch::_export::InputSpec input_spec;
+  input_spec.set_user_input(user_input_spec);
+
+  torch::_export::GraphSignature storage;
+  storage.set_input_specs({input_spec});
+  storage.set_output_specs({});
+
+  GraphSignature sig(storage);
+
+  std::vector<std::string> expected_inputs = {"sym_int_0"};
+  EXPECT_EQ(sig.userInputs(), expected_inputs);
+}
+
+// Test AS_SYM_INT with constant value - should be skipped
+TEST(GraphSignatureInputTypesTest, UserInputAsSymIntConstant) {
+  torch::_export::Argument input_arg;
+  input_arg.set_as_sym_int(makeSymIntArgConst(42));
+  torch::_export::UserInputSpec user_input_spec;
+  user_input_spec.set_arg(input_arg);
+  torch::_export::InputSpec input_spec;
+  input_spec.set_user_input(user_input_spec);
+
+  torch::_export::GraphSignature storage;
+  storage.set_input_specs({input_spec});
+  storage.set_output_specs({});
+
+  GraphSignature sig(storage);
+
+  // Constant symints should be skipped
+  EXPECT_TRUE(sig.userInputs().empty());
+}
+
+// Test AS_SYM_INTS user input handling - mixed names and constants
+TEST(GraphSignatureInputTypesTest, UserInputAsSymInts) {
+  std::vector<torch::_export::SymIntArgument> sym_ints;
+  sym_ints.push_back(makeSymIntArg("s0"));
+  sym_ints.push_back(makeSymIntArgConst(8)); // Constant - should be skipped
+  sym_ints.push_back(makeSymIntArg("s1"));
+
+  torch::_export::Argument input_arg;
+  input_arg.set_as_sym_ints(sym_ints);
+  torch::_export::UserInputSpec user_input_spec;
+  user_input_spec.set_arg(input_arg);
+  torch::_export::InputSpec input_spec;
+  input_spec.set_user_input(user_input_spec);
+
+  torch::_export::GraphSignature storage;
+  storage.set_input_specs({input_spec});
+  storage.set_output_specs({});
+
+  GraphSignature sig(storage);
+
+  // Only named symints should be in user inputs
+  std::vector<std::string> expected_inputs = {"s0", "s1"};
+  EXPECT_EQ(sig.userInputs(), expected_inputs);
+}
+
+// Test AS_TENSORS user output handling
+TEST(GraphSignatureOutputTypesTest, UserOutputAsTensors) {
+  std::vector<torch::_export::TensorArgument> tensor_list;
+  tensor_list.push_back(makeTensorArg("out_tensor_0"));
+  tensor_list.push_back(makeTensorArg("out_tensor_1"));
+
+  torch::_export::Argument output_arg;
+  output_arg.set_as_tensors(tensor_list);
+  torch::_export::UserOutputSpec user_output_spec;
+  user_output_spec.set_arg(output_arg);
+  torch::_export::OutputSpec output_spec;
+  output_spec.set_user_output(user_output_spec);
+
+  torch::_export::GraphSignature storage;
+  storage.set_input_specs({});
+  storage.set_output_specs({output_spec});
+
+  GraphSignature sig(storage);
+
+  std::vector<std::optional<std::string>> expected_outputs = {
+      "out_tensor_0", "out_tensor_1"};
+  EXPECT_EQ(sig.userOutputs(), expected_outputs);
+}
+
+// Test AS_OPTIONAL_TENSOR user output handling
+TEST(GraphSignatureOutputTypesTest, UserOutputAsOptionalTensor) {
+  torch::_export::Argument output_arg;
+  output_arg.set_as_optional_tensor(makeOptionalTensorArg("opt_out_tensor"));
+  torch::_export::UserOutputSpec user_output_spec;
+  user_output_spec.set_arg(output_arg);
+  torch::_export::OutputSpec output_spec;
+  output_spec.set_user_output(user_output_spec);
+
+  torch::_export::GraphSignature storage;
+  storage.set_input_specs({});
+  storage.set_output_specs({output_spec});
+
+  GraphSignature sig(storage);
+
+  std::vector<std::optional<std::string>> expected_outputs = {"opt_out_tensor"};
+  EXPECT_EQ(sig.userOutputs(), expected_outputs);
+}
+
+// Test AS_OPTIONAL_TENSOR user output with None - should have nullopt name
+TEST(GraphSignatureOutputTypesTest, UserOutputAsOptionalTensorNone) {
+  torch::_export::Argument output_arg;
+  output_arg.set_as_optional_tensor(makeOptionalTensorArgNone());
+  torch::_export::UserOutputSpec user_output_spec;
+  user_output_spec.set_arg(output_arg);
+  torch::_export::OutputSpec output_spec;
+  output_spec.set_user_output(user_output_spec);
+
+  torch::_export::GraphSignature storage;
+  storage.set_input_specs({});
+  storage.set_output_specs({output_spec});
+
+  GraphSignature sig(storage);
+
+  // None optional tensor should result in nullopt
+  EXPECT_EQ(sig.userOutputs().size(), 1);
+  EXPECT_FALSE(sig.userOutputs()[0].has_value());
+}
+
+// Test AS_OPTIONAL_TENSORS user output handling - mixed tensors and None
+TEST(GraphSignatureOutputTypesTest, UserOutputAsOptionalTensors) {
+  std::vector<torch::_export::OptionalTensorArgument> opt_tensors;
+  opt_tensors.push_back(makeOptionalTensorArg("opt_out_0"));
+  opt_tensors.push_back(makeOptionalTensorArgNone());
+  opt_tensors.push_back(makeOptionalTensorArg("opt_out_2"));
+
+  torch::_export::Argument output_arg;
+  output_arg.set_as_optional_tensors(opt_tensors);
+  torch::_export::UserOutputSpec user_output_spec;
+  user_output_spec.set_arg(output_arg);
+  torch::_export::OutputSpec output_spec;
+  output_spec.set_user_output(user_output_spec);
+
+  torch::_export::GraphSignature storage;
+  storage.set_input_specs({});
+  storage.set_output_specs({output_spec});
+
+  GraphSignature sig(storage);
+
+  EXPECT_EQ(sig.userOutputs().size(), 3);
+  EXPECT_EQ(sig.userOutputs()[0], "opt_out_0");
+  EXPECT_FALSE(sig.userOutputs()[1].has_value()); // None becomes nullopt
+  EXPECT_EQ(sig.userOutputs()[2], "opt_out_2");
+}
+
+// Test AS_SYM_INT user output handling
+TEST(GraphSignatureOutputTypesTest, UserOutputAsSymInt) {
+  torch::_export::Argument output_arg;
+  output_arg.set_as_sym_int(makeSymIntArg("sym_out_0"));
+  torch::_export::UserOutputSpec user_output_spec;
+  user_output_spec.set_arg(output_arg);
+  torch::_export::OutputSpec output_spec;
+  output_spec.set_user_output(user_output_spec);
+
+  torch::_export::GraphSignature storage;
+  storage.set_input_specs({});
+  storage.set_output_specs({output_spec});
+
+  GraphSignature sig(storage);
+
+  std::vector<std::optional<std::string>> expected_outputs = {"sym_out_0"};
+  EXPECT_EQ(sig.userOutputs(), expected_outputs);
+}
+
+// Test AS_SYM_INTS user output handling
+TEST(GraphSignatureOutputTypesTest, UserOutputAsSymInts) {
+  std::vector<torch::_export::SymIntArgument> sym_ints;
+  sym_ints.push_back(makeSymIntArg("sym_out_0"));
+  sym_ints.push_back(makeSymIntArgConst(10)); // Constant - should be skipped
+  sym_ints.push_back(makeSymIntArg("sym_out_1"));
+
+  torch::_export::Argument output_arg;
+  output_arg.set_as_sym_ints(sym_ints);
+  torch::_export::UserOutputSpec user_output_spec;
+  user_output_spec.set_arg(output_arg);
+  torch::_export::OutputSpec output_spec;
+  output_spec.set_user_output(user_output_spec);
+
+  torch::_export::GraphSignature storage;
+  storage.set_input_specs({});
+  storage.set_output_specs({output_spec});
+
+  GraphSignature sig(storage);
+
+  // Only named symints should appear
+  std::vector<std::optional<std::string>> expected_outputs = {
+      "sym_out_0", "sym_out_1"};
+  EXPECT_EQ(sig.userOutputs(), expected_outputs);
+}
+
+// Test AS_SYM_BOOL user output handling
+TEST(GraphSignatureOutputTypesTest, UserOutputAsSymBool) {
+  torch::_export::Argument output_arg;
+  output_arg.set_as_sym_bool(makeSymBoolArg("sym_bool_out"));
+  torch::_export::UserOutputSpec user_output_spec;
+  user_output_spec.set_arg(output_arg);
+  torch::_export::OutputSpec output_spec;
+  output_spec.set_user_output(user_output_spec);
+
+  torch::_export::GraphSignature storage;
+  storage.set_input_specs({});
+  storage.set_output_specs({output_spec});
+
+  GraphSignature sig(storage);
+
+  std::vector<std::optional<std::string>> expected_outputs = {"sym_bool_out"};
+  EXPECT_EQ(sig.userOutputs(), expected_outputs);
+}
+
+// Test AS_SYM_BOOLS user output handling
+TEST(GraphSignatureOutputTypesTest, UserOutputAsSymBools) {
+  std::vector<torch::_export::SymBoolArgument> sym_bools;
+  sym_bools.push_back(makeSymBoolArg("sym_bool_0"));
+  sym_bools.push_back(makeSymBoolArg("sym_bool_1"));
+
+  torch::_export::Argument output_arg;
+  output_arg.set_as_sym_bools(sym_bools);
+  torch::_export::UserOutputSpec user_output_spec;
+  user_output_spec.set_arg(output_arg);
+  torch::_export::OutputSpec output_spec;
+  output_spec.set_user_output(user_output_spec);
+
+  torch::_export::GraphSignature storage;
+  storage.set_input_specs({});
+  storage.set_output_specs({output_spec});
+
+  GraphSignature sig(storage);
+
+  std::vector<std::optional<std::string>> expected_outputs = {
+      "sym_bool_0", "sym_bool_1"};
+  EXPECT_EQ(sig.userOutputs(), expected_outputs);
+}
+
+// Test AS_SYM_FLOAT user output handling - should result in nullopt (unnamed)
+TEST(GraphSignatureOutputTypesTest, UserOutputAsSymFloat) {
+  torch::_export::SymFloatArgument sym_float;
+  sym_float.set_as_name("sym_float_out");
+  torch::_export::Argument output_arg;
+  output_arg.set_as_sym_float(sym_float);
+  torch::_export::UserOutputSpec user_output_spec;
+  user_output_spec.set_arg(output_arg);
+  torch::_export::OutputSpec output_spec;
+  output_spec.set_user_output(user_output_spec);
+
+  torch::_export::GraphSignature storage;
+  storage.set_input_specs({});
+  storage.set_output_specs({output_spec});
+
+  GraphSignature sig(storage);
+
+  // SymFloat outputs currently result in nullopt (unnamed)
+  EXPECT_EQ(sig.userOutputs().size(), 1);
+  EXPECT_FALSE(sig.userOutputs()[0].has_value());
+}
+
+// Test AS_SYM_FLOATS user output handling
+TEST(GraphSignatureOutputTypesTest, UserOutputAsSymFloats) {
+  torch::_export::SymFloatArgument sym_float_0;
+  sym_float_0.set_as_name("sym_float_0");
+  torch::_export::SymFloatArgument sym_float_1;
+  sym_float_1.set_as_name("sym_float_1");
+  std::vector<torch::_export::SymFloatArgument> sym_floats = {
+      sym_float_0, sym_float_1};
+
+  torch::_export::Argument output_arg;
+  output_arg.set_as_sym_floats(sym_floats);
+  torch::_export::UserOutputSpec user_output_spec;
+  user_output_spec.set_arg(output_arg);
+  torch::_export::OutputSpec output_spec;
+  output_spec.set_user_output(user_output_spec);
+
+  torch::_export::GraphSignature storage;
+  storage.set_input_specs({});
+  storage.set_output_specs({output_spec});
+
+  GraphSignature sig(storage);
+
+  // SymFloats outputs are treated as unnamed (nullopt)
+  EXPECT_EQ(sig.userOutputs().size(), 2);
+  for (const auto& output : sig.userOutputs()) {
+    EXPECT_FALSE(output.has_value());
+  }
+}
+
+// Test mixed input types in the same GraphSignature
+TEST(GraphSignatureMixedTypesTest, MixedInputTypes) {
+  // First input: single tensor
+  torch::_export::Argument input_arg_0;
+  input_arg_0.set_as_tensor(makeTensorArg("single_tensor"));
+  torch::_export::UserInputSpec user_input_spec_0;
+  user_input_spec_0.set_arg(input_arg_0);
+  torch::_export::InputSpec input_spec_0;
+  input_spec_0.set_user_input(user_input_spec_0);
+
+  // Second input: list of tensors
+  std::vector<torch::_export::TensorArgument> tensor_list;
+  tensor_list.push_back(makeTensorArg("list_tensor_0"));
+  tensor_list.push_back(makeTensorArg("list_tensor_1"));
+  torch::_export::Argument input_arg_1;
+  input_arg_1.set_as_tensors(tensor_list);
+  torch::_export::UserInputSpec user_input_spec_1;
+  user_input_spec_1.set_arg(input_arg_1);
+  torch::_export::InputSpec input_spec_1;
+  input_spec_1.set_user_input(user_input_spec_1);
+
+  // Third input: symint
+  torch::_export::Argument input_arg_2;
+  input_arg_2.set_as_sym_int(makeSymIntArg("batch_size"));
+  torch::_export::UserInputSpec user_input_spec_2;
+  user_input_spec_2.set_arg(input_arg_2);
+  torch::_export::InputSpec input_spec_2;
+  input_spec_2.set_user_input(user_input_spec_2);
+
+  torch::_export::GraphSignature storage;
+  storage.set_input_specs({input_spec_0, input_spec_1, input_spec_2});
+  storage.set_output_specs({});
+
+  GraphSignature sig(storage);
+
+  std::vector<std::string> expected_inputs = {
+      "single_tensor", "list_tensor_0", "list_tensor_1", "batch_size"};
+  EXPECT_EQ(sig.userInputs(), expected_inputs);
+}
+
+// Test mixed output types in the same GraphSignature
+TEST(GraphSignatureMixedTypesTest, MixedOutputTypes) {
+  // First output: single tensor
+  torch::_export::Argument output_arg_0;
+  output_arg_0.set_as_tensor(makeTensorArg("single_out"));
+  torch::_export::UserOutputSpec user_output_spec_0;
+  user_output_spec_0.set_arg(output_arg_0);
+  torch::_export::OutputSpec output_spec_0;
+  output_spec_0.set_user_output(user_output_spec_0);
+
+  // Second output: list of tensors
+  std::vector<torch::_export::TensorArgument> tensor_list;
+  tensor_list.push_back(makeTensorArg("list_out_0"));
+  tensor_list.push_back(makeTensorArg("list_out_1"));
+  torch::_export::Argument output_arg_1;
+  output_arg_1.set_as_tensors(tensor_list);
+  torch::_export::UserOutputSpec user_output_spec_1;
+  user_output_spec_1.set_arg(output_arg_1);
+  torch::_export::OutputSpec output_spec_1;
+  output_spec_1.set_user_output(user_output_spec_1);
+
+  // Third output: symint
+  torch::_export::Argument output_arg_2;
+  output_arg_2.set_as_sym_int(makeSymIntArg("out_size"));
+  torch::_export::UserOutputSpec user_output_spec_2;
+  user_output_spec_2.set_arg(output_arg_2);
+  torch::_export::OutputSpec output_spec_2;
+  output_spec_2.set_user_output(user_output_spec_2);
+
+  torch::_export::GraphSignature storage;
+  storage.set_input_specs({});
+  storage.set_output_specs({output_spec_0, output_spec_1, output_spec_2});
+
+  GraphSignature sig(storage);
+
+  std::vector<std::optional<std::string>> expected_outputs = {
+      "single_out", "list_out_0", "list_out_1", "out_size"};
+  EXPECT_EQ(sig.userOutputs(), expected_outputs);
+}
+
 } // namespace torch::nativert
diff --git a/test/cpp/nativert/test_serialization.cpp b/test/cpp/nativert/test_serialization.cpp
index 3504f02b53a9a..5970796699125 100644
--- a/test/cpp/nativert/test_serialization.cpp
+++ b/test/cpp/nativert/test_serialization.cpp
@@ -2,6 +2,50 @@
 #include <torch/nativert/graph/Serialization.h>
 
 namespace torch::nativert {
+
+// Helper to create tensor argument
+torch::_export::TensorArgument makeTensorArg(const std::string& name) {
+  torch::_export::TensorArgument arg;
+  arg.set_name(name);
+  return arg;
+}
+
+// Helper to create SymInt argument with name
+torch::_export::SymIntArgument makeSymIntArg(const std::string& name) {
+  torch::_export::SymIntArgument arg;
+  arg.set_as_name(name);
+  return arg;
+}
+
+// Helper to create SymInt argument with constant int
+torch::_export::SymIntArgument makeSymIntArgConst(int64_t value) {
+  torch::_export::SymIntArgument arg;
+  arg.set_as_int(value);
+  return arg;
+}
+
+// Helper to create SymBool argument
+torch::_export::SymBoolArgument makeSymBoolArg(const std::string& name) {
+  torch::_export::SymBoolArgument arg;
+  arg.set_as_name(name);
+  return arg;
+}
+
+// Helper to create OptionalTensorArgument with tensor
+torch::_export::OptionalTensorArgument makeOptionalTensorArg(
+    const std::string& name) {
+  torch::_export::OptionalTensorArgument arg;
+  arg.set_as_tensor(makeTensorArg(name));
+  return arg;
+}
+
+// Helper to create OptionalTensorArgument with None
+torch::_export::OptionalTensorArgument makeOptionalTensorArgNone() {
+  torch::_export::OptionalTensorArgument arg;
+  arg.set_as_none({});
+  return arg;
+}
+
 TEST(SerializationTest, CheckIsSymbolic) {
   torch::_export::TensorArgument tensor_arg;
   torch::_export::Argument as_tensor_arg;
@@ -31,6 +75,132 @@ TEST(SerializationTest, CheckIsSymbolic) {
   EXPECT_FALSE(isSymbolic(as_string_arg));
 }
 
+// Test isSymbolic for AS_OPTIONAL_TENSOR
+TEST(SerializationTest, CheckIsSymbolicOptionalTensor) {
+  torch::_export::Argument arg;
+  arg.set_as_optional_tensor(makeOptionalTensorArg("opt_tensor"));
+  EXPECT_TRUE(isSymbolic(arg));
+
+  // Also test with None optional tensor
+  torch::_export::Argument arg_none;
+  arg_none.set_as_optional_tensor(makeOptionalTensorArgNone());
+  EXPECT_TRUE(isSymbolic(arg_none));
+}
+
+// Test isSymbolic for AS_OPTIONAL_TENSORS
+TEST(SerializationTest, CheckIsSymbolicOptionalTensors) {
+  std::vector<torch::_export::OptionalTensorArgument> opt_tensors;
+  opt_tensors.push_back(makeOptionalTensorArg("opt_0"));
+  opt_tensors.push_back(makeOptionalTensorArgNone());
+
+  torch::_export::Argument arg;
+  arg.set_as_optional_tensors(opt_tensors);
+  EXPECT_TRUE(isSymbolic(arg));
+}
+
+// Test isSymbolic for AS_SYM_INTS
+TEST(SerializationTest, CheckIsSymbolicSymInts) {
+  std::vector<torch::_export::SymIntArgument> sym_ints;
+  sym_ints.push_back(makeSymIntArg("s0"));
+  sym_ints.push_back(makeSymIntArgConst(8));
+
+  torch::_export::Argument arg;
+  arg.set_as_sym_ints(sym_ints);
+  EXPECT_TRUE(isSymbolic(arg));
+}
+
+// Test isSymbolic for AS_SYM_BOOL
+TEST(SerializationTest, CheckIsSymbolicSymBool) {
+  torch::_export::Argument arg;
+  arg.set_as_sym_bool(makeSymBoolArg("sym_bool"));
+  EXPECT_TRUE(isSymbolic(arg));
+}
+
+// Test isSymbolic for AS_SYM_BOOLS
+TEST(SerializationTest, CheckIsSymbolicSymBools) {
+  std::vector<torch::_export::SymBoolArgument> sym_bools;
+  sym_bools.push_back(makeSymBoolArg("b0"));
+  sym_bools.push_back(makeSymBoolArg("b1"));
+
+  torch::_export::Argument arg;
+  arg.set_as_sym_bools(sym_bools);
+  EXPECT_TRUE(isSymbolic(arg));
+}
+
+// Test isSymbolic for AS_SYM_FLOAT
+TEST(SerializationTest, CheckIsSymbolicSymFloat) {
+  torch::_export::SymFloatArgument sym_float;
+  sym_float.set_as_name("sym_float");
+
+  torch::_export::Argument arg;
+  arg.set_as_sym_float(sym_float);
+  EXPECT_TRUE(isSymbolic(arg));
+}
+
+// Test isSymbolic for AS_SYM_FLOATS
+TEST(SerializationTest, CheckIsSymbolicSymFloats) {
+  torch::_export::SymFloatArgument sym_float_0;
+  sym_float_0.set_as_name("sym_float_0");
+  torch::_export::SymFloatArgument sym_float_1;
+  sym_float_1.set_as_name("sym_float_1");
+  std::vector<torch::_export::SymFloatArgument> sym_floats = {
+      sym_float_0, sym_float_1};
+
+  torch::_export::Argument arg;
+  arg.set_as_sym_floats(sym_floats);
+  EXPECT_TRUE(isSymbolic(arg));
+}
+
+// Test isSymbolic for AS_CUSTOM_OBJ
+TEST(SerializationTest, CheckIsSymbolicCustomObj) {
+  torch::_export::CustomObjArgument custom_obj;
+  custom_obj.set_name("my_custom_obj");
+  custom_obj.set_class_fqn("my.custom.Class");
+
+  torch::_export::Argument arg;
+  arg.set_as_custom_obj(custom_obj);
+  EXPECT_TRUE(isSymbolic(arg));
+}
+
+// Test that non-symbolic types return false
+TEST(SerializationTest, CheckIsSymbolicNonSymbolicTypes) {
+  // AS_FLOAT
+  torch::_export::Argument as_float;
+  torch::_export::F64 f64_val;
+  f64_val.set(3.14);
+  as_float.set_as_float(f64_val);
+  EXPECT_FALSE(isSymbolic(as_float));
+
+  // AS_INTS
+  torch::_export::Argument as_ints;
+  as_ints.set_as_ints({1, 2, 3});
+  EXPECT_FALSE(isSymbolic(as_ints));
+
+  // AS_FLOATS
+  torch::_export::Argument as_floats;
+  torch::_export::F64 f64_1, f64_2, f64_3;
+  f64_1.set(1.0);
+  f64_2.set(2.0);
+  f64_3.set(3.0);
+  as_floats.set_as_floats({f64_1, f64_2, f64_3});
+  EXPECT_FALSE(isSymbolic(as_floats));
+
+  // AS_BOOLS
+  torch::_export::Argument as_bools;
+  as_bools.set_as_bools({true, false, true});
+  EXPECT_FALSE(isSymbolic(as_bools));
+
+  // AS_NONE
+  torch::_export::Argument as_none;
+  as_none.set_as_none({});
+  EXPECT_FALSE(isSymbolic(as_none));
+
+  // AS_STRINGS
+  torch::_export::Argument as_strings;
+  as_strings.set_as_strings({"a", "b", "c"});
+  EXPECT_FALSE(isSymbolic(as_strings));
+}
+
 TEST(SerializationTest, ConstantToValue) {
   torch::_export::Argument as_int_arg;
   as_int_arg.set_as_int(static_cast<int64_t>(42));
@@ -48,4 +218,114 @@ TEST(SerializationTest, ConstantToValue) {
   EXPECT_EQ(value, Constant("test_string"));
 }
 
+// Test constantToValue for AS_FLOAT
+TEST(SerializationTest, ConstantToValueFloat) {
+  torch::_export::Argument arg;
+  torch::_export::F64 f64_val;
+  f64_val.set(3.14159);
+  arg.set_as_float(f64_val);
+  auto value = constantToValue(arg, false);
+  EXPECT_EQ(value, Constant(3.14159));
+}
+
+// Test constantToValue for AS_INTS
+TEST(SerializationTest, ConstantToValueInts) {
+  torch::_export::Argument arg;
+  arg.set_as_ints({1, 2, 3, 4, 5});
+  auto value = constantToValue(arg, false);
+  std::vector<int64_t> expected = {1, 2, 3, 4, 5};
+  EXPECT_EQ(value, Constant(expected));
+}
+
+// Test constantToValue for AS_FLOATS
+TEST(SerializationTest, ConstantToValueFloats) {
+  torch::_export::Argument arg;
+  torch::_export::F64 f64_1, f64_2, f64_3;
+  f64_1.set(1.0);
+  f64_2.set(2.5);
+  f64_3.set(3.14);
+  arg.set_as_floats({f64_1, f64_2, f64_3});
+  auto value = constantToValue(arg, false);
+  std::vector<double> expected = {1.0, 2.5, 3.14};
+  EXPECT_EQ(value, Constant(expected));
+}
+
+// Test constantToValue for AS_BOOLS
+TEST(SerializationTest, ConstantToValueBools) {
+  torch::_export::Argument arg;
+  arg.set_as_bools({true, false, true});
+  auto value = constantToValue(arg, false);
+  std::vector<bool> expected = {true, false, true};
+  EXPECT_EQ(value, Constant(expected));
+}
+
+// Test constantToValue for AS_NONE
+TEST(SerializationTest, ConstantToValueNone) {
+  torch::_export::Argument arg;
+  arg.set_as_none({});
+  auto value = constantToValue(arg, false);
+  EXPECT_EQ(value, Constant(None()));
+}
+
+// Test constantToValue for AS_STRINGS
+TEST(SerializationTest, ConstantToValueStrings) {
+  torch::_export::Argument arg;
+  arg.set_as_strings({"hello", "world"});
+  auto value = constantToValue(arg, false);
+  std::vector<std::string> expected = {"hello", "world"};
+  EXPECT_EQ(value, Constant(expected));
+}
+
+// Test that symbolic types throw when passed to constantToValue
+TEST(SerializationTest, ConstantToValueThrowsOnSymbolicTypes) {
+  // AS_TENSOR should throw
+  torch::_export::Argument as_tensor;
+  as_tensor.set_as_tensor(makeTensorArg("tensor"));
+  EXPECT_THROW(constantToValue(as_tensor, false), std::exception);
+
+  // AS_TENSORS should throw
+  torch::_export::Argument as_tensors;
+  as_tensors.set_as_tensors({makeTensorArg("t1"), makeTensorArg("t2")});
+  EXPECT_THROW(constantToValue(as_tensors, false), std::exception);
+
+  // AS_OPTIONAL_TENSORS should throw
+  torch::_export::Argument as_opt_tensors;
+  as_opt_tensors.set_as_optional_tensors(
+      {makeOptionalTensorArg("opt_t1"), makeOptionalTensorArgNone()});
+  EXPECT_THROW(constantToValue(as_opt_tensors, false), std::exception);
+
+  // AS_SYM_INT should throw
+  torch::_export::Argument as_sym_int;
+  as_sym_int.set_as_sym_int(makeSymIntArg("s0"));
+  EXPECT_THROW(constantToValue(as_sym_int, false), std::exception);
+
+  // AS_SYM_INTS should throw
+  torch::_export::Argument as_sym_ints;
+  as_sym_ints.set_as_sym_ints({makeSymIntArg("s0"), makeSymIntArg("s1")});
+  EXPECT_THROW(constantToValue(as_sym_ints, false), std::exception);
+
+  // AS_SYM_BOOL should throw
+  torch::_export::Argument as_sym_bool;
+  as_sym_bool.set_as_sym_bool(makeSymBoolArg("b0"));
+  EXPECT_THROW(constantToValue(as_sym_bool, false), std::exception);
+
+  // AS_SYM_BOOLS should throw
+  torch::_export::Argument as_sym_bools;
+  as_sym_bools.set_as_sym_bools({makeSymBoolArg("b0"), makeSymBoolArg("b1")});
+  EXPECT_THROW(constantToValue(as_sym_bools, false), std::exception);
+
+  // AS_CUSTOM_OBJ should throw
+  torch::_export::CustomObjArgument custom_obj;
+  custom_obj.set_name("obj");
+  custom_obj.set_class_fqn("MyClass");
+  torch::_export::Argument as_custom_obj;
+  as_custom_obj.set_as_custom_obj(custom_obj);
+  EXPECT_THROW(constantToValue(as_custom_obj, false), std::exception);
+
+  // AS_OPTIONAL_TENSOR should throw
+  torch::_export::Argument as_opt_tensor;
+  as_opt_tensor.set_as_optional_tensor(makeOptionalTensorArg("opt"));
+  EXPECT_THROW(constantToValue(as_opt_tensor, false), std::exception);
+}
+
 } // namespace torch::nativert
diff --git a/test/cpp/profiler/CMakeLists.txt b/test/cpp/profiler/CMakeLists.txt
new file mode 100644
index 0000000000000..76bbd548f9445
--- /dev/null
+++ b/test/cpp/profiler/CMakeLists.txt
@@ -0,0 +1,25 @@
+# Test for PrivateUse1 profiler registration
+
+if(USE_KINETO)
+  set(PROFILER_TEST_ROOT ${TORCH_ROOT}/test/cpp/profiler)
+
+  add_executable(test_privateuse1_profiler
+    ${TORCH_ROOT}/test/cpp/common/main.cpp
+    ${PROFILER_TEST_ROOT}/test_privateuse1_profiler.cpp
+  )
+
+  target_compile_definitions(test_privateuse1_profiler PRIVATE USE_GTEST USE_KINETO)
+
+  set(PROFILER_TEST_DEPENDENCIES torch gtest_main fmt::fmt-header-only)
+
+  target_link_libraries(test_privateuse1_profiler PRIVATE ${PROFILER_TEST_DEPENDENCIES})
+  target_include_directories(test_privateuse1_profiler PRIVATE
+    ${ATen_CPU_INCLUDE}
+    ${KINETO_SOURCE_DIR}/include
+  )
+
+  if(INSTALL_TEST)
+    set_target_properties(test_privateuse1_profiler PROPERTIES INSTALL_RPATH "${CMAKE_INSTALL_RPATH}:${_rpath_portable_origin}/../lib")
+    install(TARGETS test_privateuse1_profiler DESTINATION bin)
+  endif()
+endif()
diff --git a/test/cpp/profiler/test_privateuse1_profiler.cpp b/test/cpp/profiler/test_privateuse1_profiler.cpp
new file mode 100644
index 0000000000000..ac34e0143dc9a
--- /dev/null
+++ b/test/cpp/profiler/test_privateuse1_profiler.cpp
@@ -0,0 +1,181 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#ifdef USE_KINETO
+
+#include <gtest/gtest.h>
+
+#include <atomic>
+#include <memory>
+#include <string>
+
+#include <torch/csrc/autograd/profiler_kineto.h>
+#include <torch/csrc/profiler/standalone/privateuse1_profiler.h>
+
+#include <IActivityProfiler.h>
+#include <libkineto.h>
+
+namespace {
+
+// Flag to track if Kineto actually called our mock profiler's configure()
+// This proves the full flow: registration → Kineto stores it → Kineto uses it
+std::atomic<bool> g_configure_called{false};
+
+// Mock IActivityProfilerSession implementation
+class MockProfilerSession : public libkineto::IActivityProfilerSession {
+ public:
+  MockProfilerSession() = default;
+
+  void start() override {
+    status_ = libkineto::TraceStatus::RECORDING;
+  }
+
+  void stop() override {
+    status_ = libkineto::TraceStatus::PROCESSING;
+  }
+
+  std::vector<std::string> errors() override {
+    return {};
+  }
+
+  void processTrace(
+      [[maybe_unused]] libkineto::ActivityLogger& logger) override {}
+
+  std::unique_ptr<libkineto::DeviceInfo> getDeviceInfo() override {
+    return std::make_unique<libkineto::DeviceInfo>(
+        /*id=*/0,
+        /*sortIndex=*/0,
+        /*name=*/"MockDevice",
+        /*label=*/"mock_device");
+  }
+
+  std::vector<libkineto::ResourceInfo> getResourceInfos() override {
+    return {};
+  }
+
+  std::unique_ptr<libkineto::CpuTraceBuffer> getTraceBuffer() override {
+    return nullptr;
+  }
+};
+
+// Mock IActivityProfiler implementation for testing
+class MockPrivateUse1Profiler : public libkineto::IActivityProfiler {
+ public:
+  MockPrivateUse1Profiler() = default;
+
+  const std::string& name() const override {
+    return name_;
+  }
+
+  const std::set<libkineto::ActivityType>& availableActivities()
+      const override {
+    return activities_;
+  }
+
+  std::unique_ptr<libkineto::IActivityProfilerSession> configure(
+      [[maybe_unused]] const std::set<libkineto::ActivityType>& activity_types,
+      [[maybe_unused]] const libkineto::Config& config) override {
+    g_configure_called = true;
+    return std::make_unique<MockProfilerSession>();
+  }
+
+  std::unique_ptr<libkineto::IActivityProfilerSession> configure(
+      [[maybe_unused]] int64_t ts_ms,
+      [[maybe_unused]] int64_t duration_ms,
+      const std::set<libkineto::ActivityType>& activity_types,
+      const libkineto::Config& config) override {
+    return configure(activity_types, config);
+  }
+
+ private:
+  std::string name_{"mock_privateuse1"};
+  std::set<libkineto::ActivityType> activities_{
+      libkineto::ActivityType::CPU_OP};
+};
+
+} // namespace
+
+// Test that the registry singleton works correctly
+TEST(PrivateUse1ProfilerTest, RegistrySingleton) {
+  auto& registry1 =
+      torch::profiler::impl::PrivateUse1ProfilerRegistry::instance();
+  auto& registry2 =
+      torch::profiler::impl::PrivateUse1ProfilerRegistry::instance();
+  EXPECT_EQ(&registry1, &registry2);
+}
+
+// Test that registering a factory works
+TEST(PrivateUse1ProfilerTest, RegisterFactory) {
+  auto& registry =
+      torch::profiler::impl::PrivateUse1ProfilerRegistry::instance();
+
+  registry.registerFactory(
+      []() { return std::make_unique<MockPrivateUse1Profiler>(); });
+
+  EXPECT_TRUE(registry.hasFactory());
+}
+
+// Test that onKinetoInit triggers registration when factory is registered
+TEST(PrivateUse1ProfilerTest, OnKinetoInitForwarding) {
+  auto& registry =
+      torch::profiler::impl::PrivateUse1ProfilerRegistry::instance();
+
+  // Register factory if not already registered
+  if (!registry.hasFactory()) {
+    registry.registerFactory(
+        []() { return std::make_unique<MockPrivateUse1Profiler>(); });
+  }
+
+  registry.onKinetoInit();
+
+  // Verify the factory was forwarded to Kineto
+  EXPECT_TRUE(registry.hasFactory());
+  EXPECT_TRUE(registry.isRegisteredWithKineto());
+}
+
+// End-to-end test: Start a profiling session and verify mock profiler is used
+TEST(PrivateUse1ProfilerTest, EndToEndProfiling) {
+  using namespace torch::autograd::profiler;
+  using namespace torch::profiler::impl;
+
+  // Reset the flag before the test
+  g_configure_called = false;
+
+  auto& registry = PrivateUse1ProfilerRegistry::instance();
+  if (!registry.hasFactory()) {
+    registry.registerFactory(
+        []() { return std::make_unique<MockPrivateUse1Profiler>(); });
+  }
+
+  // Create profiler config with KINETO_PRIVATEUSE1 state
+  ProfilerConfig config(
+      ProfilerState::KINETO_PRIVATEUSE1,
+      /*report_input_shapes=*/false,
+      /*profile_memory=*/false,
+      /*with_stack=*/false,
+      /*with_flops=*/false,
+      /*with_modules=*/false);
+
+  std::set<ActivityType> activities{ActivityType::CPU};
+
+  // This triggers Kineto init and onKinetoInit
+  prepareProfiler(config, activities);
+  enableProfiler(config, activities, {at::RecordScope::FUNCTION});
+
+  EXPECT_TRUE(registry.isRegisteredWithKineto());
+
+  // Stop profiler
+  auto result = disableProfiler();
+
+  EXPECT_TRUE(registry.isRegisteredWithKineto());
+
+  // Verify Kineto-side: Kineto actually called our mock's configure()
+  EXPECT_TRUE(g_configure_called);
+}
+
+#endif // USE_KINETO
diff --git a/test/cpp_extensions/libtorch_agn_2_10_extension/test_version_compatibility.py b/test/cpp_extensions/libtorch_agn_2_10_extension/test_version_compatibility.py
index 3db17cdc22cb3..b16a39492a71c 100644
--- a/test/cpp_extensions/libtorch_agn_2_10_extension/test_version_compatibility.py
+++ b/test/cpp_extensions/libtorch_agn_2_10_extension/test_version_compatibility.py
@@ -24,6 +24,7 @@
 from torch.testing._internal.common_utils import IS_WINDOWS, run_tests, TestCase
 from torch.utils.cpp_extension import (
     CUDA_HOME,
+    get_cxx_compiler,
     include_paths as torch_include_paths,
     ROCM_HOME,
 )
@@ -83,9 +84,9 @@ def _compile_cpp_file(
             torch_version_2_9 = "0x0209000000000000"
 
             cmd = [
-                "g++",
+                get_cxx_compiler(),
                 "-c",
-                "-std=c++17",
+                "-std=c++20",
                 f"-DTORCH_TARGET_VERSION={torch_version_2_9}",
                 f"-I{source_file.parent}",  # For includes in same directory
                 *self.pytorch_includes,
@@ -119,7 +120,7 @@ def _compile_cu_file(
             cmd = [
                 os.path.join(GPU_HOME, "bin", "nvcc" if CUDA_HOME else "hipcc"),
                 "-c",
-                "-std=c++17",
+                "-std=c++20",
                 f"-DTORCH_TARGET_VERSION={torch_version_2_9}",
                 f"-I{source_file.parent}",  # For includes in same directory
                 *self.pytorch_includes,
diff --git a/test/cpp_extensions/libtorch_agn_2_11_extension/csrc/my_from_blob_with_lambda_deleter.cpp b/test/cpp_extensions/libtorch_agn_2_11_extension/csrc/my_from_blob_with_lambda_deleter.cpp
new file mode 100644
index 0000000000000..8e498cbf2b9f4
--- /dev/null
+++ b/test/cpp_extensions/libtorch_agn_2_11_extension/csrc/my_from_blob_with_lambda_deleter.cpp
@@ -0,0 +1,101 @@
+#include <torch/csrc/stable/device.h>
+#include <torch/csrc/stable/library.h>
+#include <torch/csrc/stable/ops.h>
+#include <torch/csrc/stable/tensor.h>
+
+#ifdef LAE_USE_CUDA
+#include <cuda_runtime.h>
+#endif
+
+using torch::stable::Tensor;
+
+// Global counter to track lambda deleter calls for testing
+static int64_t g_lambda_deleter_call_count = 0;
+
+// Wrapper for from_blob with a capturing-lambda deleter.
+// The lambda captures a pointer to the global counter and increments it,
+// which exercises the capturing-lambda code path in torch_from_blob.
+Tensor my_from_blob_with_lambda_deleter(
+    int64_t data_ptr,
+    torch::headeronly::HeaderOnlyArrayRef<int64_t> sizes,
+    torch::headeronly::HeaderOnlyArrayRef<int64_t> strides,
+    torch::stable::Device device,
+    torch::headeronly::ScalarType dtype) {
+  void* data = reinterpret_cast<void*>(data_ptr);
+  int64_t* counter = &g_lambda_deleter_call_count;
+  auto deleter = [counter](void* /*data*/) { (*counter)++; };
+  return torch::stable::from_blob(data, sizes, strides, device, dtype, deleter);
+}
+
+int64_t get_lambda_deleter_call_count() {
+  return g_lambda_deleter_call_count;
+}
+
+void reset_lambda_deleter_call_count() {
+  g_lambda_deleter_call_count = 0;
+}
+
+STABLE_TORCH_LIBRARY_FRAGMENT(STABLE_LIB_NAME, m) {
+  m.def(
+      "my_from_blob_with_lambda_deleter(int data_ptr, int[] sizes, int[] strides, Device device, ScalarType dtype) -> Tensor");
+  m.def("get_lambda_deleter_call_count() -> int");
+  m.def("reset_lambda_deleter_call_count() -> ()");
+}
+
+STABLE_TORCH_LIBRARY_IMPL(
+    STABLE_LIB_NAME,
+    CompositeExplicitAutograd,
+    m) {
+  m.impl(
+      "my_from_blob_with_lambda_deleter",
+      TORCH_BOX(&my_from_blob_with_lambda_deleter));
+  m.impl(
+      "get_lambda_deleter_call_count",
+      TORCH_BOX(&get_lambda_deleter_call_count));
+  m.impl(
+      "reset_lambda_deleter_call_count",
+      TORCH_BOX(&reset_lambda_deleter_call_count));
+}
+
+#ifdef LAE_USE_CUDA
+
+// Same as my_from_blob_with_cuda_deleter (from 2.11) but uses a non-capturing
+// lambda deleter.
+Tensor my_from_blob_with_cuda_lambda_deleter(
+    int64_t numel,
+    torch::stable::Device device) {
+  size_t size_bytes = numel * sizeof(float);
+
+  void* data = nullptr;
+  cudaError_t err = cudaMalloc(&data, size_bytes);
+  if (err != cudaSuccess) {
+    throw std::runtime_error("cudaMalloc failed");
+  }
+
+  // Zero the memory
+  cudaMemset(data, 0, size_bytes);
+
+  std::array<int64_t, 1> sizes = {numel};
+  std::array<int64_t, 1> strides = {1};
+
+  // This lambda doesn't capture anything, but capture is tested above in
+  // my_from_blob_with_lambda_deleter
+  auto deleter = [](void* data) { cudaFree(data); };
+  return torch::stable::from_blob(
+      data,
+      torch::headeronly::HeaderOnlyArrayRef<int64_t>(sizes.data(), sizes.size()),
+      torch::headeronly::HeaderOnlyArrayRef<int64_t>(strides.data(), strides.size()),
+      device,
+      torch::headeronly::ScalarType::Float,
+      deleter);
+}
+
+STABLE_TORCH_LIBRARY_FRAGMENT(STABLE_LIB_NAME, m) {
+  m.def("my_from_blob_with_cuda_lambda_deleter(int numel, Device device) -> Tensor");
+}
+
+STABLE_TORCH_LIBRARY_IMPL(STABLE_LIB_NAME, CompositeExplicitAutograd, m) {
+  m.impl("my_from_blob_with_cuda_lambda_deleter", TORCH_BOX(&my_from_blob_with_cuda_lambda_deleter));
+}
+
+#endif  // LAE_USE_CUDA
diff --git a/test/cpp_extensions/libtorch_agn_2_11_extension/libtorch_agn_2_11/ops.py b/test/cpp_extensions/libtorch_agn_2_11_extension/libtorch_agn_2_11/ops.py
index 4315898009269..442f9c97ce5db 100644
--- a/test/cpp_extensions/libtorch_agn_2_11_extension/libtorch_agn_2_11/ops.py
+++ b/test/cpp_extensions/libtorch_agn_2_11_extension/libtorch_agn_2_11/ops.py
@@ -57,6 +57,59 @@ def my_from_blob_with_cuda_deleter(numel: int, device) -> Tensor:
     )
 
 
+def my_from_blob_with_lambda_deleter(data_ptr, sizes, strides, device, dtype) -> Tensor:
+    """
+    Creates a Tensor from existing memory with a capturing-lambda deleter.
+
+    The deleter is a capturing lambda that updates a global call count,
+    exercising the capturing-lambda code path in torch_from_blob.
+
+    Args:
+        data_ptr: int - pointer to the data buffer
+        sizes: tuple[int] - size of the tensor
+        strides: tuple[int] - strides of the tensor
+        device: Device - device on which the tensor resides
+        dtype: ScalarType - data type of the tensor
+
+    Returns: Tensor - tensor wrapping the existing memory
+    """
+    return torch.ops.libtorch_agn_2_11.my_from_blob_with_lambda_deleter.default(
+        data_ptr, sizes, strides, device, dtype
+    )
+
+
+def get_lambda_deleter_call_count() -> int:
+    """
+    Returns the number of times the lambda test deleter has been called.
+    """
+    return torch.ops.libtorch_agn_2_11.get_lambda_deleter_call_count.default()
+
+
+def reset_lambda_deleter_call_count() -> None:
+    """
+    Resets the lambda deleter call counter to zero.
+    """
+    torch.ops.libtorch_agn_2_11.reset_lambda_deleter_call_count.default()
+
+
+def my_from_blob_with_cuda_lambda_deleter(numel: int, device) -> Tensor:
+    """
+    Creates a CUDA tensor that owns its memory via cudaMalloc, using a lambda deleter.
+
+    Similar to my_from_blob_with_cuda_deleter but uses the capturing-lambda
+    code path in torch_from_blob.
+
+    Args:
+        numel: int - number of elements in the tensor
+        device: Device - CUDA device
+
+    Returns: Tensor - a 1D float32 tensor of zeros
+    """
+    return torch.ops.libtorch_agn_2_11.my_from_blob_with_cuda_lambda_deleter.default(
+        numel, device
+    )
+
+
 # =============================================================================
 # Proxy for inherited ops (from libtorch_agn_2_9 and libtorch_agn_2_10 csrc/)
 #
diff --git a/test/cpp_extensions/libtorch_agn_2_12_extension/csrc/tag_test.cpp b/test/cpp_extensions/libtorch_agn_2_12_extension/csrc/tag_test.cpp
new file mode 100644
index 0000000000000..647cc2c70c91d
--- /dev/null
+++ b/test/cpp_extensions/libtorch_agn_2_12_extension/csrc/tag_test.cpp
@@ -0,0 +1,20 @@
+#include <torch/csrc/inductor/aoti_torch/c/shim.h>
+#include <torch/csrc/stable/library.h>
+#include <torch/csrc/stable/tensor.h>
+#include <torch/headeronly/core/enum_tag.h>
+
+using torch::stable::Tensor;
+
+Tensor tagged_identity(Tensor t) {
+  return t;
+}
+
+STABLE_TORCH_LIBRARY_FRAGMENT(STABLE_LIB_NAME, m) {
+  m.def("tagged_identity(Tensor t) -> Tensor",
+        std::vector<at::Tag>{
+            at::Tag::pointwise, at::Tag::pt2_compliant_tag, at::Tag::core});
+}
+
+STABLE_TORCH_LIBRARY_IMPL(STABLE_LIB_NAME, CompositeExplicitAutograd, m) {
+  m.impl("tagged_identity", TORCH_BOX(&tagged_identity));
+}
diff --git a/test/cpp_extensions/libtorch_agn_2_12_extension/libtorch_agn_2_12/__init__.py b/test/cpp_extensions/libtorch_agn_2_12_extension/libtorch_agn_2_12/__init__.py
new file mode 100644
index 0000000000000..e29e3e784d9d9
--- /dev/null
+++ b/test/cpp_extensions/libtorch_agn_2_12_extension/libtorch_agn_2_12/__init__.py
@@ -0,0 +1,25 @@
+import ctypes
+import sys
+from pathlib import Path
+
+import torch
+
+
+so_files = list(
+    Path(__file__).parent.glob("_C*" + (".pyd" if sys.platform == "win32" else ".so"))
+)
+if len(so_files) != 1:
+    raise AssertionError(f"Expected one _C*.{{so,pyd}} file, found {len(so_files)}")
+
+# use ctypes.CDLL instead of load_library to be able to test the unload logic
+# below code is reduced from the load_library code
+with torch._ops.dl_open_guard():
+    loaded_lib = ctypes.CDLL(str(so_files[0]))
+
+from . import ops
+
+
+__all__ = [
+    "loaded_lib",
+    "ops",
+]
diff --git a/test/cpp_extensions/libtorch_agn_2_12_extension/libtorch_agn_2_12/ops.py b/test/cpp_extensions/libtorch_agn_2_12_extension/libtorch_agn_2_12/ops.py
new file mode 100644
index 0000000000000..a35b04ce50ceb
--- /dev/null
+++ b/test/cpp_extensions/libtorch_agn_2_12_extension/libtorch_agn_2_12/ops.py
@@ -0,0 +1,36 @@
+import torch
+
+
+# =============================================================================
+# Proxy for inherited ops (from libtorch_agn_2_9, 2_10, and 2_11 csrc/)
+#
+# Ops compiled from previous versions' csrc directories are accessible via
+# the module-level __getattr__. For example:
+#     libtorch_agn_2_12.ops.sgd_out_of_place(...)  # from 2.9
+#     libtorch_agn_2_12.ops.my_sum(...)            # from 2.10
+# =============================================================================
+
+_NAMESPACE = "libtorch_agn_2_12"
+
+
+def __getattr__(name):
+    """Proxy for inherited ops from previous versions."""
+    if name.startswith("_"):
+        raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+    ops_namespace = getattr(torch.ops, _NAMESPACE)
+    op = getattr(ops_namespace, name, None)
+    if op is None:
+        raise AttributeError(f"No op named '{name}' in {_NAMESPACE}")
+    return op.default
+
+
+def __dir__():
+    """List all available ops (native + inherited)."""
+    native = [
+        name
+        for name in globals()
+        if not name.startswith("_") and callable(globals().get(name))
+    ]
+    ops_namespace = getattr(torch.ops, _NAMESPACE)
+    inherited = [n for n in dir(ops_namespace) if not n.startswith("_")]
+    return sorted(set(native + inherited))
diff --git a/test/cpp_extensions/libtorch_agn_2_12_extension/setup.py b/test/cpp_extensions/libtorch_agn_2_12_extension/setup.py
new file mode 100644
index 0000000000000..90944115aeec1
--- /dev/null
+++ b/test/cpp_extensions/libtorch_agn_2_12_extension/setup.py
@@ -0,0 +1,102 @@
+import distutils.command.clean
+import shutil
+from pathlib import Path
+
+from setuptools import find_packages, setup
+
+import torch
+from torch.utils.cpp_extension import (
+    BuildExtension,
+    CppExtension,
+    CUDAExtension,
+    IS_WINDOWS,
+)
+
+
+ROOT_DIR = Path(__file__).parent
+CSRC_DIR = ROOT_DIR / "csrc"
+
+# Include csrc from previous versions for forward compatibility testing
+PREV_CSRC_DIRS = [
+    ROOT_DIR.parent / "libtorch_agn_2_9_extension" / "csrc",
+    ROOT_DIR.parent / "libtorch_agn_2_10_extension" / "csrc",
+    ROOT_DIR.parent / "libtorch_agn_2_11_extension" / "csrc",
+]
+
+
+class clean(distutils.command.clean.clean):
+    def run(self):
+        # Run default behavior first
+        distutils.command.clean.clean.run(self)
+
+        # Remove extension
+        for path in (ROOT_DIR / "libtorch_agn_2_12").glob("**/*.so"):
+            path.unlink()
+        # Remove build and dist and egg-info directories
+        dirs = [
+            ROOT_DIR / "build",
+            ROOT_DIR / "dist",
+            ROOT_DIR / "libtorch_agn_2_12.egg-info",
+        ]
+        for path in dirs:
+            if path.exists():
+                shutil.rmtree(str(path), ignore_errors=True)
+
+
+def get_extension():
+    extra_compile_args = {
+        "cxx": [
+            "-DTORCH_TARGET_VERSION=0x020c000000000000",
+            "-DSTABLE_LIB_NAME=libtorch_agn_2_12",
+        ],
+    }
+    if not IS_WINDOWS:
+        extra_compile_args["cxx"].append("-fdiagnostics-color=always")
+
+    # Collect sources from this version and all previous versions
+    sources = list(CSRC_DIR.glob("**/*.cpp"))
+    for prev_dir in PREV_CSRC_DIRS:
+        sources.extend(prev_dir.glob("**/*.cpp"))
+
+    extension = CppExtension
+    # allow including <cuda_runtime.h>
+    if torch.cuda.is_available():
+        extra_compile_args["cxx"].append("-DLAE_USE_CUDA")
+        extra_compile_args["nvcc"] = [
+            "-O2",
+            "-DUSE_CUDA",
+            "-DSTABLE_LIB_NAME=libtorch_agn_2_12",
+        ]
+        extension = CUDAExtension
+        sources.extend(CSRC_DIR.glob("**/*.cu"))
+        for prev_dir in PREV_CSRC_DIRS:
+            sources.extend(prev_dir.glob("**/*.cu"))
+
+    return [
+        extension(
+            "libtorch_agn_2_12._C",
+            sources=sorted(str(s) for s in sources),
+            py_limited_api=True,
+            extra_compile_args=extra_compile_args,
+            extra_link_args=[],
+        )
+    ]
+
+
+setup(
+    name="libtorch_agn_2_12",
+    version="0.0",
+    author="PyTorch Core Team",
+    description="Example of libtorch agnostic extension for PyTorch 2.12+",
+    packages=find_packages(exclude=("test",)),
+    package_data={"libtorch_agn_2_12": ["*.dll", "*.dylib", "*.so"]},
+    install_requires=[
+        "torch",
+    ],
+    ext_modules=get_extension(),
+    cmdclass={
+        "build_ext": BuildExtension.with_options(no_python_abi_suffix=True),
+        "clean": clean,
+    },
+    options={"bdist_wheel": {"py_limited_api": "cp39"}},
+)
diff --git a/test/cpp_extensions/open_registration_extension/torch_openreg/CMakeLists.txt b/test/cpp_extensions/open_registration_extension/torch_openreg/CMakeLists.txt
index d57ecdf38fc85..2c6b2555aaeea 100644
--- a/test/cpp_extensions/open_registration_extension/torch_openreg/CMakeLists.txt
+++ b/test/cpp_extensions/open_registration_extension/torch_openreg/CMakeLists.txt
@@ -5,7 +5,7 @@ project(TORCH_OPENREG CXX C)
 include(GNUInstallDirs)
 include(CheckCXXCompilerFlag)
 
-set(CMAKE_CXX_STANDARD 17)
+set(CMAKE_CXX_STANDARD 20)
 set(CMAKE_C_STANDARD 17)
 set(CMAKE_CXX_EXTENSIONS OFF)
 
diff --git a/test/cpp_extensions/open_registration_extension/torch_openreg/README.md b/test/cpp_extensions/open_registration_extension/torch_openreg/README.md
index 422aa4c8da229..97c95c838962c 100644
--- a/test/cpp_extensions/open_registration_extension/torch_openreg/README.md
+++ b/test/cpp_extensions/open_registration_extension/torch_openreg/README.md
@@ -6,7 +6,7 @@ The third-party device integration mechanism based on PrivateUse1 has become the
 
 **Note:**
 
-The goal of `torch_openreg` is **not to implement a fully functional, high-performance PyTorch backend**, but to serve as a **minimalist reference implementation for mechanism verification**.
+The goal of `OpenReg` codebase is **not to implement a fully functional, high-performance PyTorch backend**, but to serve as a **minimalist reference implementation for mechanism verification**.
 
 ### Purpose
 
@@ -17,7 +17,7 @@ The goal of `torch_openreg` is **not to implement a fully functional, high-perfo
 ### Design Principles
 
 - **Minimality Principle**: The fundamental goal is to enable/verify all integration paths/mechanisms for a new backend to integrate to PyTorch. All functions follow a "just right" strategy to ensure the correctness of relevant integration capabilities.
-- **Authenticity Principle**: To complete the OpenReg integration in the same way a real accelerator backend would integrate with PyTorch.
+- **Authenticity Principle**: To complete the `OpenReg` integration in the same way a real accelerator backend would integrate with PyTorch.
 
 ## Directory Structure
 
@@ -101,7 +101,7 @@ There are 4 DSOs in torch_openreg, and the dependencies between them are as foll
 - `csrc/`: Core device implementation, including operator registration, runtime, etc.
   - `csrc/amp/`: AMP(Automatic Mixed Precision)
   - `csrc/aten/`: Operator registration
-    - `csrc/aten/native/`: Specific operator implementations for the OpenReg device.
+    - `csrc/aten/native/`: Specific operator implementations for the `openreg` device.
       - `csrc/aten/native/OpenRegMinimal.cpp`: The most minimal set of operator implementations (allowing for the creation of Tensors and related operations upon completion).
       - `csrc/aten/native/OpenRegExtra.cpp`: Implementations for other types of operators.
   - `csrc/runtime/`: Implementations for Host memory, device memory, Guard, Hooks, etc.
@@ -131,7 +131,7 @@ There are 4 DSOs in torch_openreg, and the dependencies between them are as foll
 
 ### Autoload
 
-When `import torch`, installed accelerators (such as `torch_openreg`) will be automatically loaded, achieving the same experience as the built-in backends.
+When `import torch`, installed accelerators (such as `openreg`) will be automatically loaded, achieving the same experience as the built-in backends.
 
 - Register the backend with Python `entry points`: See `setup` in `setup.py`
 - Add a callable function for backend initialization: See `_autoload` in `torch_openreg/__init__.py`
diff --git a/test/cpp_extensions/open_registration_extension/torch_openreg/csrc/CMakeLists.txt b/test/cpp_extensions/open_registration_extension/torch_openreg/csrc/CMakeLists.txt
index e2ae2b3f3667e..3c038115a3dcd 100644
--- a/test/cpp_extensions/open_registration_extension/torch_openreg/csrc/CMakeLists.txt
+++ b/test/cpp_extensions/open_registration_extension/torch_openreg/csrc/CMakeLists.txt
@@ -6,7 +6,7 @@ file(GLOB_RECURSE SOURCE_FILES
 
 add_library(${LIBRARY_NAME} SHARED ${SOURCE_FILES})
 
-target_link_libraries(${LIBRARY_NAME} PRIVATE torch_cpu_library openreg)
+target_link_libraries(${LIBRARY_NAME} PRIVATE torch_cpu_library torch_python_library openreg)
 target_include_directories(${LIBRARY_NAME} PUBLIC ${CMAKE_CURRENT_SOURCE_DIR})
 
 install(TARGETS ${LIBRARY_NAME}
diff --git a/test/cpp_extensions/open_registration_extension/torch_openreg/csrc/distributed/c10d/ProcessGroupOCCL.cpp b/test/cpp_extensions/open_registration_extension/torch_openreg/csrc/distributed/c10d/ProcessGroupOCCL.cpp
new file mode 100644
index 0000000000000..1b180d9e94fd4
--- /dev/null
+++ b/test/cpp_extensions/open_registration_extension/torch_openreg/csrc/distributed/c10d/ProcessGroupOCCL.cpp
@@ -0,0 +1,368 @@
+#if USE_DISTRIBUTED
+#include <stdexcept>
+
+#include <ATen/core/Tensor.h>
+#include <c10/util/Exception.h>
+#include <c10/util/string_view.h>
+#include <torch/headeronly/core/DeviceType.h>
+
+#include "ProcessGroupOCCL.hpp"
+namespace {
+
+#define CHECK_TENSOR(tensor)                                                    \
+  do {                                                                          \
+    if (tensor.defined()) {                                                     \
+      TORCH_CHECK(                                                              \
+          tensor.device().type() == c10::DeviceType::PrivateUse1,               \
+          "ProcessGroupOCCL only supports OpenReg (PrivateUse1) tensors. Got ", \
+          tensor.device().type(),                                               \
+          " for argument '",                                                    \
+          #tensor,                                                              \
+          "' in ",                                                              \
+          __func__,                                                             \
+          ".");                                                                 \
+    }                                                                           \
+  } while (0)
+
+#define CHECK_TENSOR_LIST(tensorList)                                             \
+  do {                                                                            \
+    for (const auto& tensor : tensorList) {                                       \
+      if (tensor.defined()) {                                                     \
+        TORCH_CHECK(                                                              \
+            tensor.device().type() == c10::DeviceType::PrivateUse1,               \
+            "ProcessGroupOCCL only supports OpenReg (PrivateUse1) tensors. Got ", \
+            tensor.device().type(),                                               \
+            " for argument '",                                                    \
+            #tensorList,                                                          \
+            "' in ",                                                              \
+            __func__,                                                             \
+            ".");                                                                 \
+      }                                                                           \
+    }                                                                             \
+  } while (0)
+
+#define CHECK_TENSOR_LIST_OF_LISTS(tensorListOfLists)                               \
+  do {                                                                              \
+    for (const auto& tensor_list : tensorListOfLists) {                             \
+      for (const auto& tensor : tensor_list) {                                      \
+        if (tensor.defined()) {                                                     \
+          TORCH_CHECK(                                                              \
+              tensor.device().type() == c10::DeviceType::PrivateUse1,               \
+              "ProcessGroupOCCL only supports OpenReg (PrivateUse1) tensors. Got ", \
+              tensor.device().type(),                                               \
+              " for argument '",                                                    \
+              #tensorListOfLists,                                                   \
+              "' in ",                                                              \
+              __func__,                                                             \
+              ".");                                                                 \
+        }                                                                           \
+      }                                                                             \
+    }                                                                               \
+  } while (0)
+
+} // namespace
+
+/**
+    @brief Currently, a no-op OCCL ProcessGroup stub for registration, tests and
+   educational purposes.
+
+    This translation unit provides a minimal ProcessGroup backend named "OCCL"
+   (short for OpenReg Collective Communications Library) that only validates
+   input shapes/devices and immediately returns a completed Work object without
+   performing any actual communication. It is intended for:
+        - Out-of-tree custom CCL backend registration and linkage tests
+        - API surface validation and call-site integration
+
+    All collective point-to-point and synchronization calls are implemented as
+    immediate, successful completions (unless a precondition check fails and an
+    exception is thrown). As a result, these calls do not exchange data, do not
+    synchronize ranks, and do not modify input/output tensors. OCCL is for
+   testing and educational purposes only; it does not provide any real
+   distributed functionality.
+
+    ------------------------------------------------------------------------------
+    Class: ProcessGroupOCCL::DummyWork
+
+    @brief An immediately-completed Work instance.
+
+    - Construction marks an internal c10::ivalue::Future as completed.
+    - isCompleted(): always true after construction.
+    - isSuccess(): true unless an error was explicitly set (not expected here).
+    - wait(timeout): returns immediately; the Future is already complete.
+    - synchronize(): no-op; there is no associated device or stream work.
+    - abort(): attempts to set an error only if the Future is not completed;
+                         has no effect in the current implementation since the
+   Future is completed in the constructor.
+    - getFuture(): returns the already-completed Future.
+
+    Notes:
+    - No background progress, streams, or device synchronization is involved.
+    - Intended solely as a lightweight placeholder Work for API conformance.
+
+    ------------------------------------------------------------------------------
+    Class: ProcessGroupOCCL
+
+    @brief A stub ProcessGroup backend whose collectives validate inputs and
+   then return a DummyWork that is already complete.
+
+    General behavior across all methods below:
+    - Input validation: Methods call internal helpers (e.g., checkTensorList,
+        checkTensorListOfLists, checkTensorDevice) to validate that inputs are
+        defined, have expected dtypes/devices, and adhere to basic invariants.
+        This is for building connection with out-of-tree accelerator backend
+   (OpenReg) and as a guard for blocking any changes from upstream that will
+   break the compatibility.
+    - No communication: No data movement or cross-rank coordination occurs.
+    - Options ignored: All options/attributes (e.g., reduce ops, tags, counts,
+        timeouts) are accepted but not used to affect behavior.
+    - Immediate completion: A ProcessGroupOCCL::DummyWork is returned, whose
+        future is already completed successfully.
+    - No side-effects: Output tensors are not populated or modified.
+
+    Limitations and caveats:
+    - Does not guarantee any semantic property of real collectives (e.g., data
+        equivalence across ranks, synchronization, ordering).
+    - Barrier does not synchronize ranks; it simply returns a completed Work.
+    - Safe for use in unit tests only where communication effects are not
+   required.
+
+    ------------------------------------------------------------------------------
+    Factory:
+
+    @fn createProcessGroupOCCL(const c10::intrusive_ptr<c10d::Store>& store, int
+   rank, int size, const std::chrono::duration<float>& timeout) Constructs a
+   ProcessGroupOCCL using the given rank and size. The store and timeout are
+   accepted for API parity but are not used to affect behavior.
+            - @param store: Process group store (ignored).
+            - @param rank: This process rank.
+            - @param size: World size.
+            - @param timeout: Operation timeout hint (ignored).
+            - @return A new ProcessGroupOCCL instance.
+
+    ------------------------------------------------------------------------------
+    Usage notes:
+    - Suitable for tests that only check callability/flow, not correctness of
+        distributed results.
+    - Do not rely on returned outputs or synchronization semantics.
+    - In production, replace with a real backend (e.g., NCCL, Gloo, MPI).
+*/
+namespace c10d {
+
+// WorkOCCL -----------------------------------------------------------------
+
+ProcessGroupOCCL::DummyWork::DummyWork()
+    : Work(-1, OpType::UNKNOWN),
+      future_(c10::make_intrusive<c10::ivalue::Future>(c10::NoneType::get())) {
+  if (!future_->completed()) {
+    future_->markCompleted(c10::IValue());
+  }
+  finish();
+}
+
+ProcessGroupOCCL::DummyWork::~DummyWork() = default;
+
+bool ProcessGroupOCCL::DummyWork::isCompleted() {
+  return future_->completed();
+}
+
+bool ProcessGroupOCCL::DummyWork::isSuccess() const {
+  return !future_->hasError();
+}
+
+bool ProcessGroupOCCL::DummyWork::wait(
+    std::chrono::milliseconds /* timeout */) {
+  future_->wait();
+  return !future_->hasError();
+}
+
+void ProcessGroupOCCL::DummyWork::synchronize() {}
+
+void ProcessGroupOCCL::DummyWork::abort() {
+  if (!future_->completed()) {
+    future_->setError(
+        std::make_exception_ptr(std::runtime_error("OCCL work aborted")));
+  }
+}
+
+c10::intrusive_ptr<c10::ivalue::Future> ProcessGroupOCCL::DummyWork::
+    getFuture() {
+  return future_;
+}
+
+// Options ------------------------------------------------------------------
+
+ProcessGroupOCCL::Options::Options(std::chrono::milliseconds timeout)
+    : Backend::Options(OCCL_BACKEND_NAME, timeout) {}
+
+// ProcessGroupOCCL ---------------------------------------------------------
+
+ProcessGroupOCCL::ProcessGroupOCCL(int rank, int size)
+    : Backend(rank, size), options_(Options::create()) {}
+
+ProcessGroupOCCL::~ProcessGroupOCCL() = default;
+
+c10::intrusive_ptr<Work> ProcessGroupOCCL::broadcast(
+    std::vector<at::Tensor>& tensors,
+    const BroadcastOptions& /* opts */) {
+  CHECK_TENSOR_LIST(tensors);
+  return c10::make_intrusive<ProcessGroupOCCL::DummyWork>();
+}
+
+c10::intrusive_ptr<Work> ProcessGroupOCCL::allreduce(
+    std::vector<at::Tensor>& tensors,
+    const AllreduceOptions& /* opts */) {
+  CHECK_TENSOR_LIST(tensors);
+  return c10::make_intrusive<ProcessGroupOCCL::DummyWork>();
+}
+
+c10::intrusive_ptr<Work> ProcessGroupOCCL::allreduce_sparse(
+    std::vector<at::Tensor>& tensors,
+    const AllreduceOptions& /* opts */) {
+  CHECK_TENSOR_LIST(tensors);
+  return c10::make_intrusive<ProcessGroupOCCL::DummyWork>();
+}
+
+c10::intrusive_ptr<Work> ProcessGroupOCCL::allreduce_coalesced(
+    std::vector<at::Tensor>& tensors,
+    const AllreduceCoalescedOptions& /* opts */) {
+  CHECK_TENSOR_LIST(tensors);
+  return c10::make_intrusive<ProcessGroupOCCL::DummyWork>();
+}
+
+c10::intrusive_ptr<Work> ProcessGroupOCCL::reduce(
+    std::vector<at::Tensor>& tensors,
+    const ReduceOptions& /* opts */) {
+  CHECK_TENSOR_LIST(tensors);
+  return c10::make_intrusive<ProcessGroupOCCL::DummyWork>();
+}
+
+c10::intrusive_ptr<Work> ProcessGroupOCCL::_reduce_scatter_base(
+    at::Tensor& outputTensor,
+    at::Tensor& inputTensor,
+    const ReduceScatterOptions& /* opts */) {
+  CHECK_TENSOR(outputTensor);
+  CHECK_TENSOR(inputTensor);
+  return c10::make_intrusive<ProcessGroupOCCL::DummyWork>();
+}
+
+c10::intrusive_ptr<Work> ProcessGroupOCCL::_allgather_base(
+    at::Tensor& output_tensor,
+    at::Tensor& input_tensor,
+    const AllgatherOptions& /* opts */) {
+  CHECK_TENSOR(output_tensor);
+  CHECK_TENSOR(input_tensor);
+  return c10::make_intrusive<ProcessGroupOCCL::DummyWork>();
+}
+
+c10::intrusive_ptr<Work> ProcessGroupOCCL::allgather(
+    std::vector<std::vector<at::Tensor>>& outputs,
+    std::vector<at::Tensor>& inputs,
+    const AllgatherOptions& /* opts */) {
+  CHECK_TENSOR_LIST_OF_LISTS(outputs);
+  CHECK_TENSOR_LIST(inputs);
+  return c10::make_intrusive<ProcessGroupOCCL::DummyWork>();
+}
+
+c10::intrusive_ptr<Work> ProcessGroupOCCL::allgather_coalesced(
+    std::vector<std::vector<at::Tensor>>& output_lists,
+    std::vector<at::Tensor>& input_list,
+    const AllgatherOptions& /* opts */) {
+  CHECK_TENSOR_LIST_OF_LISTS(output_lists);
+  CHECK_TENSOR_LIST(input_list);
+  return c10::make_intrusive<ProcessGroupOCCL::DummyWork>();
+}
+
+c10::intrusive_ptr<Work> ProcessGroupOCCL::allgather_into_tensor_coalesced(
+    std::vector<at::Tensor>& outputs,
+    std::vector<at::Tensor>& inputs,
+    const AllgatherOptions& /* opts */) {
+  CHECK_TENSOR_LIST(outputs);
+  CHECK_TENSOR_LIST(inputs);
+  return c10::make_intrusive<ProcessGroupOCCL::DummyWork>();
+}
+
+c10::intrusive_ptr<Work> ProcessGroupOCCL::gather(
+    std::vector<std::vector<at::Tensor>>& outputs,
+    std::vector<at::Tensor>& inputs,
+    const GatherOptions& /* opts */) {
+  CHECK_TENSOR_LIST_OF_LISTS(outputs);
+  CHECK_TENSOR_LIST(inputs);
+  return c10::make_intrusive<ProcessGroupOCCL::DummyWork>();
+}
+
+c10::intrusive_ptr<Work> ProcessGroupOCCL::scatter(
+    std::vector<at::Tensor>& outputs,
+    std::vector<std::vector<at::Tensor>>& inputs,
+    const ScatterOptions& /* opts */) {
+  CHECK_TENSOR_LIST(outputs);
+  CHECK_TENSOR_LIST_OF_LISTS(inputs);
+  return c10::make_intrusive<ProcessGroupOCCL::DummyWork>();
+}
+
+c10::intrusive_ptr<Work> ProcessGroupOCCL::reduce_scatter(
+    std::vector<at::Tensor>& outputs,
+    std::vector<std::vector<at::Tensor>>& inputs,
+    const ReduceScatterOptions& /* opts */) {
+  CHECK_TENSOR_LIST(outputs);
+  CHECK_TENSOR_LIST_OF_LISTS(inputs);
+  return c10::make_intrusive<ProcessGroupOCCL::DummyWork>();
+}
+
+c10::intrusive_ptr<Work> ProcessGroupOCCL::reduce_scatter_tensor_coalesced(
+    std::vector<at::Tensor>& outputTensors,
+    std::vector<at::Tensor>& inputTensors,
+    const ReduceScatterOptions& /* opts */) {
+  CHECK_TENSOR_LIST(outputTensors);
+  CHECK_TENSOR_LIST(inputTensors);
+  return c10::make_intrusive<ProcessGroupOCCL::DummyWork>();
+}
+
+c10::intrusive_ptr<Work> ProcessGroupOCCL::alltoall_base(
+    at::Tensor& outputTensor,
+    at::Tensor& inputTensor,
+    std::vector<int64_t>& /* outputCounts */,
+    std::vector<int64_t>& /* inputCounts */,
+    const AllToAllOptions& /* opts */) {
+  CHECK_TENSOR(outputTensor);
+  CHECK_TENSOR(inputTensor);
+  return c10::make_intrusive<ProcessGroupOCCL::DummyWork>();
+}
+
+c10::intrusive_ptr<Work> ProcessGroupOCCL::send(
+    std::vector<at::Tensor>& tensors,
+    int /* dstRank */,
+    int /* tag */) {
+  CHECK_TENSOR_LIST(tensors);
+  return c10::make_intrusive<ProcessGroupOCCL::DummyWork>();
+}
+
+c10::intrusive_ptr<Work> ProcessGroupOCCL::recv(
+    std::vector<at::Tensor>& tensors,
+    int /* srcRank */,
+    int /* tag */) {
+  CHECK_TENSOR_LIST(tensors);
+  return c10::make_intrusive<ProcessGroupOCCL::DummyWork>();
+}
+
+c10::intrusive_ptr<Work> ProcessGroupOCCL::recvAnysource(
+    std::vector<at::Tensor>& tensors,
+    int /* tag */) {
+  CHECK_TENSOR_LIST(tensors);
+  return c10::make_intrusive<ProcessGroupOCCL::DummyWork>();
+}
+
+c10::intrusive_ptr<Work> ProcessGroupOCCL::barrier(
+    const BarrierOptions& /* opts */) {
+  return c10::make_intrusive<ProcessGroupOCCL::DummyWork>();
+}
+
+c10::intrusive_ptr<ProcessGroupOCCL> createProcessGroupOCCL(
+    const c10::intrusive_ptr<c10d::Store>& store,
+    int rank,
+    int size,
+    const std::chrono::duration<float>& timeout) {
+  return c10::make_intrusive<ProcessGroupOCCL>(rank, size);
+}
+
+} // namespace c10d
+#endif // USE_DISTRIBUTED
diff --git a/test/cpp_extensions/open_registration_extension/torch_openreg/csrc/distributed/c10d/ProcessGroupOCCL.hpp b/test/cpp_extensions/open_registration_extension/torch_openreg/csrc/distributed/c10d/ProcessGroupOCCL.hpp
new file mode 100644
index 0000000000000..b92e0550d6173
--- /dev/null
+++ b/test/cpp_extensions/open_registration_extension/torch_openreg/csrc/distributed/c10d/ProcessGroupOCCL.hpp
@@ -0,0 +1,168 @@
+#pragma once
+#if USE_DISTRIBUTED
+#include <chrono>
+
+#include <c10/util/intrusive_ptr.h>
+#include <torch/csrc/distributed/c10d/Backend.hpp>
+#include <torch/csrc/distributed/c10d/Store.hpp>
+#include <torch/csrc/distributed/c10d/Types.hpp>
+#include <torch/csrc/distributed/c10d/Utils.hpp>
+#include <torch/csrc/distributed/c10d/Work.hpp>
+
+#include <include/Macros.h>
+
+namespace c10d {
+
+//
+// ProcessGroupTest implements dummy bindings for c10d.
+//
+
+constexpr const char* OCCL_BACKEND_NAME = "occl";
+
+class OPENREG_EXPORT ProcessGroupOCCL : public Backend {
+ public:
+  class DummyWork : public Work {
+   public:
+    DummyWork();
+
+    virtual ~DummyWork();
+    bool isCompleted() override;
+    bool isSuccess() const override;
+    bool wait(std::chrono::milliseconds timeout) override;
+    void synchronize() override;
+    void abort() override;
+    c10::intrusive_ptr<c10::ivalue::Future> getFuture() override;
+
+   protected:
+    friend class ProcessGroupOCCL;
+
+   private:
+    c10::intrusive_ptr<c10::ivalue::Future> future_;
+  };
+
+  struct OPENREG_EXPORT Options : public Backend::Options {
+    explicit Options(
+        std::chrono::milliseconds timeout = kBackendDefaultTimeout);
+
+    // return intrusive_ptr of the object
+    static c10::intrusive_ptr<Options> create(
+        std::chrono::milliseconds timeout = kBackendDefaultTimeout) {
+      return c10::make_intrusive<Options>(timeout);
+    }
+  };
+
+  explicit ProcessGroupOCCL(int rank = -1, int size = -1);
+  virtual ~ProcessGroupOCCL();
+  const std::string getBackendName() const override {
+    return std::string(OCCL_BACKEND_NAME);
+  }
+
+  bool supportsSplitting() const override {
+    return false;
+  }
+
+  c10::intrusive_ptr<Backend::Options> getBackendOptions() override {
+    return c10::static_intrusive_pointer_cast<Backend::Options>(options_);
+  }
+
+  c10::intrusive_ptr<Work> broadcast(
+    std::vector<at::Tensor>& tensors,
+    const BroadcastOptions& opts = BroadcastOptions()) override;
+
+  c10::intrusive_ptr<Work> allreduce(
+    std::vector<at::Tensor>& tensors,
+    const AllreduceOptions& opts = AllreduceOptions()) override;
+
+  c10::intrusive_ptr<Work> allreduce_sparse(
+    std::vector<at::Tensor>& tensors,
+    const AllreduceOptions& opts = AllreduceOptions()) override;
+
+  c10::intrusive_ptr<Work> allreduce_coalesced(
+    std::vector<at::Tensor>& tensors,
+    const AllreduceCoalescedOptions& opts =
+      AllreduceCoalescedOptions()) override;
+
+  c10::intrusive_ptr<Work> reduce(
+    std::vector<at::Tensor>& tensors,
+    const ReduceOptions& opts = ReduceOptions()) override;
+
+  c10::intrusive_ptr<Work> _reduce_scatter_base(
+    at::Tensor& outputTensor,
+    at::Tensor& inputTensor,
+    const ReduceScatterOptions& opts = ReduceScatterOptions()) override;
+
+  c10::intrusive_ptr<Work> _allgather_base(
+    at::Tensor& output_tensor,
+    at::Tensor& input_tensor,
+    const AllgatherOptions& opts = AllgatherOptions()) override;
+
+  c10::intrusive_ptr<Work> allgather(
+    std::vector<std::vector<at::Tensor>>& outputs,
+    std::vector<at::Tensor>& inputs,
+    const AllgatherOptions& opts = AllgatherOptions()) override;
+
+  c10::intrusive_ptr<Work> allgather_coalesced(
+    std::vector<std::vector<at::Tensor>>& output_lists,
+    std::vector<at::Tensor>& input_list,
+    const AllgatherOptions& opts = AllgatherOptions()) override;
+
+  c10::intrusive_ptr<Work> allgather_into_tensor_coalesced(
+    std::vector<at::Tensor>& outputs,
+    std::vector<at::Tensor>& inputs,
+    const AllgatherOptions& opts = AllgatherOptions()) override;
+
+  c10::intrusive_ptr<Work> gather(
+    std::vector<std::vector<at::Tensor>>& outputs,
+    std::vector<at::Tensor>& inputs,
+    const GatherOptions& opts = GatherOptions()) override;
+
+  c10::intrusive_ptr<Work> scatter(
+    std::vector<at::Tensor>& outputs,
+    std::vector<std::vector<at::Tensor>>& inputs,
+    const ScatterOptions& opts = ScatterOptions()) override;
+
+  c10::intrusive_ptr<Work> reduce_scatter(
+    std::vector<at::Tensor>& outputs,
+    std::vector<std::vector<at::Tensor>>& inputs,
+    const ReduceScatterOptions& opts = ReduceScatterOptions()) override;
+
+  c10::intrusive_ptr<Work> reduce_scatter_tensor_coalesced(
+    std::vector<at::Tensor>& outputTensors,
+    std::vector<at::Tensor>& inputTensors,
+    const ReduceScatterOptions& opts = ReduceScatterOptions()) override;
+
+  c10::intrusive_ptr<Work> alltoall_base(
+    at::Tensor& outputTensor,
+    at::Tensor& inputTensor,
+    std::vector<int64_t>& outputCounts,
+    std::vector<int64_t>& inputCounts,
+    const AllToAllOptions& opts = AllToAllOptions()) override;
+
+  c10::intrusive_ptr<Work> send(
+    std::vector<at::Tensor>& tensors,
+    int dstRank,
+    int tag) override;
+
+  c10::intrusive_ptr<Work> recv(
+    std::vector<at::Tensor>& tensors,
+    int srcRank,
+    int tag) override;
+
+  c10::intrusive_ptr<Work> recvAnysource(
+    std::vector<at::Tensor>& tensors,
+    int tag) override;
+
+  c10::intrusive_ptr<Work> barrier(
+    const BarrierOptions& opts = BarrierOptions()) override;
+
+ protected:
+  const c10::intrusive_ptr<Options> options_;
+};
+OPENREG_EXPORT c10::intrusive_ptr<ProcessGroupOCCL> createProcessGroupOCCL(
+    const c10::intrusive_ptr<c10d::Store>& store,
+    int rank,
+    int size,
+    const std::chrono::duration<float>& timeout);
+
+} // namespace c10d
+#endif // USE_DISTRIBUTED
diff --git a/test/cpp_extensions/open_registration_extension/torch_openreg/tests/test_distributed.py b/test/cpp_extensions/open_registration_extension/torch_openreg/tests/test_distributed.py
new file mode 100644
index 0000000000000..991eec020113e
--- /dev/null
+++ b/test/cpp_extensions/open_registration_extension/torch_openreg/tests/test_distributed.py
@@ -0,0 +1,101 @@
+# Owner(s): ["module: PrivateUse1"]
+
+import os
+import tempfile
+import unittest
+
+import torch
+import torch.distributed as dist
+
+
+if not (dist.is_available() and "occl" in dist.Backend.backend_list):
+    raise unittest.SkipTest("torch.distributed OCCL backend unavailable")
+
+
+os.environ["BACKEND"] = "occl"
+os.environ["WORLD_SIZE"] = "2"
+
+# Provide a TEMP_DIR compatible with Barrier helpers if not already present.
+if "TEMP_DIR" not in os.environ:
+    _tmp_dir = tempfile.TemporaryDirectory()
+    os.environ["TEMP_DIR"] = _tmp_dir.name
+    os.makedirs(os.path.join(_tmp_dir.name, "barrier"), exist_ok=True)
+    os.makedirs(os.path.join(_tmp_dir.name, "test_dir"), exist_ok=True)
+    os.makedirs(os.path.join(_tmp_dir.name, "init_dir"), exist_ok=True)
+    os.environ["INIT_METHOD"] = (
+        f"file://{os.path.join(_tmp_dir.name, 'init_dir', 'shared_init_file')}"
+    )
+
+
+from torch.distributed.distributed_c10d import _get_default_group
+from torch.testing._internal.common_distributed import (
+    cleanup_temp_dir,
+    initialize_temp_directories,
+    MultiProcessTestCase,
+)
+from torch.testing._internal.common_utils import run_tests
+from torch.testing._internal.distributed.distributed_test import (
+    Barrier,
+    DistributedTest,
+    TestDistBackend,
+)
+
+
+class TestProcessGroupOCCL(TestDistBackend, DistributedTest._DistTestBase):
+    # OCCL collectives are dummy; we only validate registration and basic API.
+
+    def setUp(self):
+        # Bypass TestDistBackend.setUp (which expects many inherited tests) and
+        # perform the minimal MultiProcessTestCase setup plus temp dir wiring.
+        MultiProcessTestCase.setUp(self)
+        initialize_temp_directories()
+        Barrier.init()
+        self.skip_return_code_checks = []
+        # Spawn subprocesses for multi-process execution when running as main.
+        if self.rank == self.MAIN_PROCESS_RANK:
+            self._spawn_processes()
+
+    def tearDown(self):
+        cleanup_temp_dir()
+        super().tearDown()
+
+    def test_occl_backend_registration_and_default_group(self) -> None:
+        # Backend is registered and available in the backend list.
+        self.assertTrue(dist.is_available())
+        self.assertIn("occl", dist.Backend.backend_list)
+
+        # Default group should already be initialized for OCCL.
+        self.assertTrue(dist.is_initialized())
+        self.assertEqual(dist.get_backend(), "occl")
+
+        pg = _get_default_group()
+        self.assertEqual(dist.get_backend(pg), "occl")
+        self.assertEqual(dist.get_world_size(), 2)
+        self.assertEqual(dist.get_rank(), self.rank)
+
+    def test_occl_allreduce(self) -> None:
+        pg = _get_default_group()
+        device = torch.device("openreg")
+
+        tensors = [torch.rand(1, device=device)]
+        fut = pg.allreduce(tensors).get_future()
+        fut.wait()
+        self.assertTrue(fut.done())
+
+        # OCCL is currently a dummy backend and all communication methods return None.
+        res = fut.value()
+        self.assertEqual(res, None)
+
+
+# Drop inherited torch distributed backend tests; we only need the OCCL smoke test.
+_allowed_tests = {
+    "test_occl_backend_registration_and_default_group",
+    "test_occl_allreduce",
+}
+for _name in dir(TestProcessGroupOCCL):
+    if _name.startswith("test_") and _name not in _allowed_tests:
+        setattr(TestProcessGroupOCCL, _name, None)
+
+
+if __name__ == "__main__":
+    run_tests()
diff --git a/test/cpp_extensions/open_registration_extension/torch_openreg/tests/test_testing.py b/test/cpp_extensions/open_registration_extension/torch_openreg/tests/test_testing.py
new file mode 100644
index 0000000000000..ba615451e8b3f
--- /dev/null
+++ b/test/cpp_extensions/open_registration_extension/torch_openreg/tests/test_testing.py
@@ -0,0 +1,61 @@
+# Owner(s): ["module: PrivateUse1"]
+
+import torch
+from torch.testing._internal.common_device_type import (
+    instantiate_device_type_tests,
+    onlyCUDA,
+    onlyOn,
+)
+from torch.testing._internal.common_utils import run_tests, TestCase
+
+
+class TestBypassDeviceRestrictions(TestCase):
+    """Verify that PrivateUse1 backends can run tests decorated with @onlyCUDA.
+
+    PrivateUse1TestBase sets bypass_device_restrictions = False by default.
+    Backends that are ready to run @onlyOn-gated tests must explicitly opt in
+    by setting bypass_device_restrictions = True in their setUp or subclass.
+    This allows out-of-tree backends to run accelerator tests that are
+    currently gated behind @onlyCUDA while the long-term migration to
+    device-generic tests is in progress.
+    """
+
+    executed_count = 0
+
+    def setUp(self):
+        # Explicitly opt-in: instance attribute shadows the False default on
+        # PrivateUse1TestBase so @onlyOn-based decorators allow this test to run.
+        self.bypass_device_restrictions = True
+        super().setUp()
+
+    @onlyCUDA
+    def test_bypass_only_cuda(self, device):
+        type(self).executed_count += 1
+        self.assertEqual(torch.device(device).type, "openreg")
+
+    @onlyOn(["cuda"])
+    def test_bypass_only_on(self, device):
+        type(self).executed_count += 1
+        self.assertEqual(torch.device(device).type, "openreg")
+
+    def test_vaildate_bypass_execution(self, device):
+        # Must run last. The 'v' prefix ensures this sorts after test_bypass_* ('b') alphabetically,
+        # so executed_count has been incremented by both bypass tests before we check it here.
+        expected_runs = 2
+        actual_runs = type(self).executed_count
+        self.assertEqual(
+            actual_runs,
+            expected_runs,
+            f"Bypass logic failed! "
+            f"Expected {expected_runs} tests to run, "
+            f"but only {actual_runs} executed.",
+        )
+
+
+instantiate_device_type_tests(
+    TestBypassDeviceRestrictions, globals(), only_for="openreg"
+)
+
+
+if __name__ == "__main__":
+    run_tests()
diff --git a/test/cpp_extensions/open_registration_extension/torch_openreg/third_party/openreg/CMakeLists.txt b/test/cpp_extensions/open_registration_extension/torch_openreg/third_party/openreg/CMakeLists.txt
index 2c7d26d6806bc..3aac3bf2c8ef9 100644
--- a/test/cpp_extensions/open_registration_extension/torch_openreg/third_party/openreg/CMakeLists.txt
+++ b/test/cpp_extensions/open_registration_extension/torch_openreg/third_party/openreg/CMakeLists.txt
@@ -39,7 +39,4 @@ if(USE_TEST)
     )
 
     add_test(NAME alltests COMMAND ${LIBRARY_TEST})
-    add_custom_command(TARGET ${LIBRARY_TEST}
-                POST_BUILD
-                COMMAND ${CMAKE_CTEST_COMMAND} -C Release --output-on-failure --verbose)
 endif()
diff --git a/test/cpp_extensions/open_registration_extension/torch_openreg/third_party/openreg/csrc/stream.cpp b/test/cpp_extensions/open_registration_extension/torch_openreg/third_party/openreg/csrc/stream.cpp
index 8cf68d87ddf96..7307e65f904ed 100644
--- a/test/cpp_extensions/open_registration_extension/torch_openreg/third_party/openreg/csrc/stream.cpp
+++ b/test/cpp_extensions/open_registration_extension/torch_openreg/third_party/openreg/csrc/stream.cpp
@@ -59,7 +59,10 @@ struct orStream {
   }
 
   ~orStream() {
-    stop_flag.store(true);
+    {
+      std::lock_guard<std::mutex> lock(mtx);
+      stop_flag.store(true);
+    }
     cv.notify_one();
     worker.join();
   }
diff --git a/test/cpp_extensions/open_registration_extension/torch_openreg/torch_openreg/__init__.py b/test/cpp_extensions/open_registration_extension/torch_openreg/torch_openreg/__init__.py
index 874dbf79b2e8f..546dc1aac77b6 100644
--- a/test/cpp_extensions/open_registration_extension/torch_openreg/torch_openreg/__init__.py
+++ b/test/cpp_extensions/open_registration_extension/torch_openreg/torch_openreg/__init__.py
@@ -17,6 +17,14 @@
 torch._register_device_module("openreg", torch_openreg.openreg)
 torch.utils.generate_methods_for_privateuse1_backend(for_storage=True)
 
+if torch.distributed.is_available():
+    try:
+        torch.distributed.Backend.register_backend(
+            "occl", torch_openreg._C._createProcessGroupOCCL, devices=["openreg"]
+        )
+    except Exception as e:
+        raise RuntimeError("Failed to register 'occl' process group backend.") from e
+
 
 # LITERALINCLUDE START: AUTOLOAD
 def _autoload():
diff --git a/test/cpp_extensions/open_registration_extension/torch_openreg/torch_openreg/csrc/Module.cpp b/test/cpp_extensions/open_registration_extension/torch_openreg/torch_openreg/csrc/Module.cpp
index 24cb3404d96f0..bb9e3da3bd2f3 100644
--- a/test/cpp_extensions/open_registration_extension/torch_openreg/torch_openreg/csrc/Module.cpp
+++ b/test/cpp_extensions/open_registration_extension/torch_openreg/torch_openreg/csrc/Module.cpp
@@ -4,8 +4,15 @@
 #include <torch/csrc/utils.h>
 #include <torch/csrc/utils/device_lazy_init.h>
 #include <torch/csrc/utils/object_ptr.h>
+#include <torch/csrc/utils/pybind.h>
 #include <torch/csrc/utils/python_numbers.h>
 
+#include <pybind11/chrono.h>
+#include <pybind11/pybind11.h>
+
+#if USE_DISTRIBUTED
+#include <distributed/c10d/ProcessGroupOCCL.hpp>
+#endif
 #include <runtime/OpenRegFunctions.h>
 
 static PyObject* _initExtension(PyObject* self, PyObject* noargs) {
@@ -111,5 +118,26 @@ extern "C" OPENREG_EXPORT PyObject* initOpenRegModule(void) {
       PyModuleDef_HEAD_INIT, "torch_openreg._C", nullptr, -1, methods};
   PyObject* mod = PyModule_Create(&openreg_C_module);
 
+  namespace py = pybind11;
+  py::module m = py::reinterpret_borrow<py::module>(mod);
+  // Expose the OCCL process group to Python only when distributed is enabled.
+  // The intrusive_ptr template arg tells pybind11 how to manage lifetime of
+  // C++ shared state when returned to Python, and we list Backend as a base
+  // so Python recognizes the inheritance hierarchy.
+#if USE_DISTRIBUTED
+  py::class_<c10d::ProcessGroupOCCL, c10d::Backend, c10::intrusive_ptr<c10d::ProcessGroupOCCL>>( // NOLINT(bugprone-unused-raii)
+      m,
+      "ProcessGroupOCCL");
+  m.def(
+      // Factory used by Python to construct the OCCL process group from a
+      // Store, rank, world size, and timeout.
+      "_createProcessGroupOCCL",
+      &c10d::createProcessGroupOCCL,
+      py::arg("store"),
+      py::arg("rank"),
+      py::arg("size"),
+      py::arg("timeout"));
+#endif
+
   return mod;
 }
diff --git a/test/cpp_extensions/setup.py b/test/cpp_extensions/setup.py
index 35da0b4391884..2fd8f1b2667ff 100644
--- a/test/cpp_extensions/setup.py
+++ b/test/cpp_extensions/setup.py
@@ -42,6 +42,8 @@
     ),
 ]
 
+NVCC_FLAGS = ["-O2"] + (["-DUSE_CUDA"] if IS_WINDOWS else [])
+
 if torch.cuda.is_available() and (CUDA_HOME is not None or ROCM_HOME is not None):
     extension = CUDAExtension(
         "torch_test_cpp_extension.cuda",
@@ -50,7 +52,7 @@
             "cuda_extension_kernel.cu",
             "cuda_extension_kernel2.cu",
         ],
-        extra_compile_args={"cxx": CXX_FLAGS, "nvcc": ["-O2"]},
+        extra_compile_args={"cxx": CXX_FLAGS, "nvcc": NVCC_FLAGS},
     )
     ext_modules.append(extension)
 
@@ -58,7 +60,7 @@
     extension = CUDAExtension(
         "torch_test_cpp_extension.torch_library",
         ["torch_library.cu"],
-        extra_compile_args={"cxx": CXX_FLAGS, "nvcc": ["-O2"]},
+        extra_compile_args={"cxx": CXX_FLAGS, "nvcc": NVCC_FLAGS},
     )
     ext_modules.append(extension)
 
diff --git a/test/cpp_extensions/test_libtorch_agnostic.py b/test/cpp_extensions/test_libtorch_agnostic.py
index 195651c4284ae..ebe8dd2536266 100644
--- a/test/cpp_extensions/test_libtorch_agnostic.py
+++ b/test/cpp_extensions/test_libtorch_agnostic.py
@@ -60,6 +60,7 @@ class TestLibtorchAgnostic(TestCase):
     - libtorch_agn_2_9: Extension built with TORCH_TARGET_VERSION=2.9.0
     - libtorch_agn_2_10: Extension built with TORCH_TARGET_VERSION=2.10.0
     - libtorch_agn_2_11: Extension built with TORCH_TARGET_VERSION=2.11.0
+    - libtorch_agn_2_12: Extension built with TORCH_TARGET_VERSION=2.12.0
 
     Tests should be decorated with @skipIfTorchVersionLessThan to indicate the
     version that they target.
@@ -105,6 +106,16 @@ def setUpClass(cls):
         else:
             print(f"Skipping 2.11 extension (running on PyTorch {torch.__version__})")
 
+        if (current_major > 2) or (current_major == 2 and current_minor >= 12):
+            try:
+                import libtorch_agn_2_12  # noqa: F401
+            except Exception:
+                install_cpp_extension(
+                    extension_root=base_dir / "libtorch_agn_2_12_extension"
+                )
+        else:
+            print(f"Skipping 2.12 extension (running on PyTorch {torch.__version__})")
+
     @onlyCPU
     def test_slow_sgd(self, device):
         import libtorch_agn_2_9 as libtorch_agnostic
@@ -1355,7 +1366,10 @@ def test_std_cuda_check_error(self, device, show_cpp_stacktraces):
         )
 
         if show_cpp_stacktraces:
-            self.assertIn("C++ CapturedTraceback:", error_message)
+            self.assertRegex(
+                error_message,
+                r"C\+\+ CapturedTraceback:|Exception raised from[\s\S]*frame #",
+            )
             self.assertRegex(
                 error_message,
                 r"Exception raised from test_std_.*_check_error at .*test_std_.*check\..*:\d+",
@@ -1524,7 +1538,10 @@ def test_std_cuda_kernel_launch_check_error(self, device, show_cpp_stacktraces):
         )
 
         if show_cpp_stacktraces:
-            self.assertIn("C++ CapturedTraceback:", error_message)
+            self.assertRegex(
+                error_message,
+                r"C\+\+ CapturedTraceback:|Exception raised from[\s\S]*frame #",
+            )
             self.assertRegex(
                 error_message,
                 r"Exception raised from test_std_.*_kernel_launch_check_error at .*test_std_.*_check\..*:\d+",
@@ -1787,6 +1804,57 @@ def inner():
             curr_mem = torch.cuda.memory_allocated(device)
             self.assertEqual(curr_mem, init_mem)
 
+    @skipIfTorchVersionLessThan(2, 11)
+    @skipIfTorchDynamo("no data pointer defined for FakeTensor, FunctionalTensor")
+    def test_my_from_blob_with_lambda_deleter(self, device):
+        """Test for from_blob with capturing-lambda deleter (2.11 feature)."""
+        import libtorch_agn_2_11 as libtorch_agnostic
+
+        from_blob_fn = libtorch_agnostic.ops.my_from_blob_with_lambda_deleter
+        get_count = libtorch_agnostic.ops.get_lambda_deleter_call_count
+        reset_count = libtorch_agnostic.ops.reset_lambda_deleter_call_count
+
+        is_cuda = torch.device(device).type == "cuda"
+        if is_cuda:
+            init_mem = torch.cuda.memory_allocated(device)
+
+        def inner():
+            reset_count()
+            self.assertEqual(get_count(), 0)
+
+            # We need an original tensor to create the tensor with from_blob.
+            original = torch.rand(2, 3, device=device, dtype=torch.float32)
+            blob_tensor = from_blob_fn(
+                original.data_ptr(),
+                original.size(),
+                original.stride(),
+                device,
+                torch.float32,
+            )
+
+            self.assertEqual(blob_tensor, original)
+            self.assertEqual(blob_tensor.data_ptr(), original.data_ptr())
+
+            self.assertEqual(get_count(), 0)
+
+            del blob_tensor
+            gc.collect()
+
+            # Ensure the deleter was called. The original tensor still exists
+            # and can be used.
+            self.assertEqual(get_count(), 1)
+            original += 1
+            # original goes out of scope here and its cuda memory should be
+            # freed.
+
+        inner()
+
+        if is_cuda:
+            # original tensor is out of scope, all the memory should be freed
+            torch.cuda.synchronize(device)
+            curr_mem = torch.cuda.memory_allocated(device)
+            self.assertEqual(curr_mem, init_mem)
+
     @onlyCUDA
     @skipIfTorchVersionLessThan(2, 11)
     def test_my_from_blob_with_cuda_deleter_no_leak(self, device):
@@ -1809,6 +1877,40 @@ def test_my_from_blob_with_cuda_deleter_no_leak(self, device):
             curr_mem = torch.cuda.memory_allocated(device)
             self.assertEqual(curr_mem, init_mem)
 
+    @onlyCUDA
+    @skipIfTorchVersionLessThan(2, 11)
+    def test_my_from_blob_with_cuda_lambda_deleter_no_leak(self, device):
+        """Test that from_blob lambda deleter properly frees cudaMalloc'd memory."""
+        import libtorch_agn_2_11 as libtorch_agnostic
+
+        from_blob_fn = libtorch_agnostic.ops.my_from_blob_with_cuda_lambda_deleter
+
+        torch.cuda.synchronize(device)
+        init_mem = torch.cuda.memory_allocated(device)
+        numel = 1024 * 1024  # 4 MB per tensor
+
+        for _ in range(10):
+            tensor = from_blob_fn(numel, device)
+            # Verify tensor was created correctly
+            self.assertEqual(tensor.numel(), numel)
+            self.assertEqual(tensor.device, torch.device(device))
+            del tensor
+            gc.collect()
+            torch.cuda.synchronize(device)
+
+            curr_mem = torch.cuda.memory_allocated(device)
+            self.assertEqual(curr_mem, init_mem)
+
+    @skipIfTorchVersionLessThan(2, 12)
+    @onlyCPU
+    def test_tagged_op(self, device):
+        import libtorch_agn_2_12  # noqa: F401
+
+        op = torch.ops.libtorch_agn_2_12.tagged_identity.default
+        self.assertIn(torch.Tag.pointwise, op.tags)
+        self.assertIn(torch.Tag.pt2_compliant_tag, op.tags)
+        self.assertIn(torch.Tag.core, op.tags)
+
     @onlyCPU
     def test_my_layout(self, device):
         """Test layout() method for various tensor layouts."""
diff --git a/test/custom_backend/CMakeLists.txt b/test/custom_backend/CMakeLists.txt
index 9ab888c5f5b1a..1f7a68a4a660b 100644
--- a/test/custom_backend/CMakeLists.txt
+++ b/test/custom_backend/CMakeLists.txt
@@ -9,9 +9,9 @@ endif()
 find_package(Torch REQUIRED)
 
 add_library(custom_backend SHARED custom_backend.cpp)
-set_property(TARGET custom_backend PROPERTY CXX_STANDARD 17)
+set_property(TARGET custom_backend PROPERTY CXX_STANDARD 20)
 target_link_libraries(custom_backend "${TORCH_LIBRARIES}")
 
 add_executable(test_custom_backend test_custom_backend.cpp)
-set_property(TARGET test_custom_backend PROPERTY CXX_STANDARD 17)
+set_property(TARGET test_custom_backend PROPERTY CXX_STANDARD 20)
 target_link_libraries(test_custom_backend custom_backend)
diff --git a/test/custom_operator/CMakeLists.txt b/test/custom_operator/CMakeLists.txt
index cc5cd7337f22c..92ea9d447ab43 100644
--- a/test/custom_operator/CMakeLists.txt
+++ b/test/custom_operator/CMakeLists.txt
@@ -9,12 +9,12 @@ endif()
 find_package(Torch REQUIRED)
 
 add_library(custom_ops SHARED op.cpp)
-set_property(TARGET custom_ops PROPERTY CXX_STANDARD 17)
+set_property(TARGET custom_ops PROPERTY CXX_STANDARD 20)
 
 target_compile_features(custom_ops PUBLIC cxx_range_for)
 target_link_libraries(custom_ops "${TORCH_LIBRARIES}")
 target_compile_definitions(custom_ops PRIVATE custom_ops_EXPORTS)
 
 add_executable(test_custom_ops test_custom_ops.cpp)
-set_property(TARGET test_custom_ops PROPERTY CXX_STANDARD 17)
+set_property(TARGET test_custom_ops PROPERTY CXX_STANDARD 20)
 target_link_libraries(test_custom_ops custom_ops)
diff --git a/test/custom_operator/test_infer_schema_annotation.py b/test/custom_operator/test_infer_schema_annotation.py
index be9e66d782513..0ad40e4622cc1 100644
--- a/test/custom_operator/test_infer_schema_annotation.py
+++ b/test/custom_operator/test_infer_schema_annotation.py
@@ -2,7 +2,6 @@
 from __future__ import annotations
 
 import typing
-from typing import Optional, Union
 
 import torch
 from torch import Tensor, types
@@ -75,7 +74,7 @@ def foo_op_3(x: torch.device) -> int:
         self.assertEqual(result, "(Device x) -> SymInt")
 
     def test_type_variants(self):
-        def foo_op_1(x: typing.Optional[int]) -> int:
+        def foo_op_1(x: int | None) -> int:
             return 1
 
         result = torch.library.infer_schema(foo_op_1, mutates_args=mutates_args)
@@ -93,25 +92,25 @@ def foo_op_3(x: list[int]) -> int:
         result = torch.library.infer_schema(foo_op_3, mutates_args=mutates_args)
         self.assertEqual(result, "(SymInt[] x) -> SymInt")
 
-        def foo_op_4(x: typing.Optional[typing.Sequence[int]]) -> int:
+        def foo_op_4(x: typing.Sequence[int] | None) -> int:
             return 1
 
         result = torch.library.infer_schema(foo_op_4, mutates_args=mutates_args)
         self.assertEqual(result, "(SymInt[]? x) -> SymInt")
 
-        def foo_op_5(x: typing.Optional[list[int]]) -> int:
+        def foo_op_5(x: list[int] | None) -> int:
             return 1
 
         result = torch.library.infer_schema(foo_op_5, mutates_args=mutates_args)
         self.assertEqual(result, "(SymInt[]? x) -> SymInt")
 
-        def foo_op_6(x: typing.Union[int, float, bool]) -> types.Number:
+        def foo_op_6(x: int | float | bool) -> types.Number:
             return x
 
         result = torch.library.infer_schema(foo_op_6, mutates_args=mutates_args)
         self.assertEqual(result, "(Scalar x) -> Scalar")
 
-        def foo_op_7(x: typing.Union[int, bool, float]) -> types.Number:
+        def foo_op_7(x: int | bool | float) -> types.Number:
             return x
 
         result = torch.library.infer_schema(foo_op_7, mutates_args=mutates_args)
@@ -142,7 +141,7 @@ def foo_op_4(x: list[int]) -> types.Number:
         result = torch.library.infer_schema(foo_op_4, mutates_args=mutates_args)
         self.assertEqual(result, "(SymInt[] x) -> Scalar")
 
-        def foo_op_5(x: Optional[int]) -> int:
+        def foo_op_5(x: int | None) -> int:
             return 1
 
         result = torch.library.infer_schema(foo_op_5, mutates_args=mutates_args)
@@ -160,25 +159,25 @@ def foo_op_7(x: list[int]) -> int:
         result = torch.library.infer_schema(foo_op_7, mutates_args=mutates_args)
         self.assertEqual(result, "(SymInt[] x) -> SymInt")
 
-        def foo_op_8(x: Optional[Sequence[int]]) -> int:
+        def foo_op_8(x: Sequence[int] | None) -> int:
             return 1
 
         result = torch.library.infer_schema(foo_op_8, mutates_args=mutates_args)
         self.assertEqual(result, "(SymInt[]? x) -> SymInt")
 
-        def foo_op_9(x: Optional[list[int]]) -> int:
+        def foo_op_9(x: list[int] | None) -> int:
             return 1
 
         result = torch.library.infer_schema(foo_op_9, mutates_args=mutates_args)
         self.assertEqual(result, "(SymInt[]? x) -> SymInt")
 
-        def foo_op_10(x: Union[int, float, bool]) -> types.Number:
+        def foo_op_10(x: int | float | bool) -> types.Number:
             return x
 
         result = torch.library.infer_schema(foo_op_10, mutates_args=mutates_args)
         self.assertEqual(result, "(Scalar x) -> Scalar")
 
-        def foo_op_11(x: Union[int, bool, float]) -> types.Number:
+        def foo_op_11(x: int | bool | float) -> types.Number:
             return x
 
         result = torch.library.infer_schema(foo_op_11, mutates_args=mutates_args)
diff --git a/test/distributed/_composable/fsdp/test_fully_shard_autograd.py b/test/distributed/_composable/fsdp/test_fully_shard_autograd.py
index d2d9eaea75b91..f639b5f8586cc 100644
--- a/test/distributed/_composable/fsdp/test_fully_shard_autograd.py
+++ b/test/distributed/_composable/fsdp/test_fully_shard_autograd.py
@@ -4,7 +4,7 @@
 import copy
 import functools
 import itertools
-from typing import Any, Optional, Union
+from typing import Any
 
 import torch
 import torch.distributed as dist
@@ -36,7 +36,7 @@ def world_size(self) -> int:
         return min(4, torch.get_device_module(device_type).device_count())
 
     def _reduce_1d_partial_grads(
-        self, module: nn.Module, group: Optional[dist.ProcessGroup] = None
+        self, module: nn.Module, group: dist.ProcessGroup | None = None
     ) -> None:
         group = group or dist.distributed_c10d._get_default_group()
         for param in module.parameters():
@@ -55,7 +55,7 @@ def test_unused_forward_output(self):
             self._test_unused_forward_output,
         )
 
-    def _test_unused_forward_output(self, reshard_after_forward: Union[bool, int]):
+    def _test_unused_forward_output(self, reshard_after_forward: bool | int):
         torch.manual_seed(42)
         local_batch_size = 2
         global_batch_size, dim = (self.world_size * local_batch_size, 24)
@@ -101,7 +101,7 @@ def test_unused_forward_module(self):
             self._test_unused_forward_module,
         )
 
-    def _test_unused_forward_module(self, reshard_after_forward: Union[bool, int]):
+    def _test_unused_forward_module(self, reshard_after_forward: bool | int):
         torch.manual_seed(42)
         local_batch_size, dim = (2, 24)
         global_batch_size = self.world_size * local_batch_size
diff --git a/test/distributed/_composable/fsdp/test_fully_shard_clip_grad_norm_.py b/test/distributed/_composable/fsdp/test_fully_shard_clip_grad_norm_.py
index 37fb7ebeb92cf..26c445768ca62 100644
--- a/test/distributed/_composable/fsdp/test_fully_shard_clip_grad_norm_.py
+++ b/test/distributed/_composable/fsdp/test_fully_shard_clip_grad_norm_.py
@@ -2,7 +2,6 @@
 
 import copy
 import functools
-from typing import Optional, Union
 
 import torch
 import torch.nn as nn
@@ -32,14 +31,14 @@
 class _TestClipGradNormBase(FSDPTest):
     def _test_clip_grad_norm(
         self,
-        max_norm: Union[float, int],
-        norm_type: Union[float, int],
+        max_norm: float | int,
+        norm_type: float | int,
         ref_model: nn.Module,
         ref_optim: torch.optim.Optimizer,
         model: nn.Module,
         optim: torch.optim.Optimizer,
         inp: torch.Tensor,
-        dp_mesh: Optional[DeviceMesh] = None,
+        dp_mesh: DeviceMesh | None = None,
     ):
         vector_norm_fn = functools.partial(torch.linalg.vector_norm, ord=norm_type)
         dp_mesh = dp_mesh or init_device_mesh(device_type.type, (self.world_size,))
diff --git a/test/distributed/_composable/fsdp/test_fully_shard_comm.py b/test/distributed/_composable/fsdp/test_fully_shard_comm.py
index 85ab79bc29ac4..be535fbec0ab0 100644
--- a/test/distributed/_composable/fsdp/test_fully_shard_comm.py
+++ b/test/distributed/_composable/fsdp/test_fully_shard_comm.py
@@ -7,13 +7,14 @@
 import tempfile
 import unittest
 from collections.abc import Callable
-from typing import Optional, Union
 from unittest.mock import MagicMock
 
 import torch
 import torch.distributed as dist
 import torch.nn as nn
 import torch.nn.functional as F
+from torch._C._autograd import DeviceType
+from torch._C._distributed_c10d import _SymmetricMemory
 from torch.distributed._composable import checkpoint, replicate
 from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import (
     apply_activation_checkpointing,
@@ -45,8 +46,10 @@
 from torch.distributed.tensor import DTensor
 from torch.distributed.tensor.debug import CommDebugMode
 from torch.distributed.tensor.experimental import implicit_replication
+from torch.testing._internal.common_cuda import SM90OrLater, TEST_MULTIGPU
 from torch.testing._internal.common_distributed import (
-    requires_multicast_support,
+    MultiProcContinuousTest,
+    PLATFORM_SUPPORTS_SYMM_MEM,
     skip_if_lt_x_gpu,
 )
 from torch.testing._internal.common_fsdp import (
@@ -60,7 +63,11 @@
     patch_unshard,
 )
 from torch.testing._internal.common_utils import (
+    instantiate_parametrized_tests,
+    parametrize,
+    requires_cuda_p2p_access,
     run_tests,
+    skip_but_pass_in_sandcastle_if,
     TEST_WITH_ROCM,
     TEST_XPU,
     xfailIf,
@@ -71,6 +78,7 @@
     Transformer,
     TransformerBlock,
 )
+from torch.testing._internal.inductor_utils import skipCUDAIf
 
 
 c10d_ops = torch.ops.c10d
@@ -120,7 +128,7 @@ def _init_params(self, param_sizes: list[torch.Size]) -> list[nn.Parameter]:
         return orig_params
 
     def _init_fsdp_param_group(
-        self, params: list[nn.Parameter], reshard_after_forward: Union[bool, int]
+        self, params: list[nn.Parameter], reshard_after_forward: bool | int
     ):
         module = nn.ParameterList([param.detach().clone() for param in params])
         mesh_info = FSDPMeshInfo(_init_default_fully_shard_mesh(), shard_mesh_dim=0)
@@ -171,7 +179,7 @@ def test_all_gather_fp32(self):
     def _test_all_gather(
         self,
         param_sizes: list[torch.Size],
-        reshard_after_forward: Union[bool, int],
+        reshard_after_forward: bool | int,
         async_op: bool,
         all_gather_copy_in_stream,
         all_gather_stream,
@@ -280,6 +288,7 @@ def _test_reduce_scatter(
         all_reduce_stream = device_module.Stream()
         comm = DefaultReduceScatter()
         (
+            _,
             _,
             _,
             post_reduce_event,
@@ -346,7 +355,7 @@ def test_fully_shard_communication_count(self):
 
     def _test_communication_count(
         self,
-        reshard_after_forward: Union[bool, int, None],
+        reshard_after_forward: bool | int | None,
     ):
         torch.manual_seed(42)
         model_args = ModelArgs()
@@ -538,6 +547,124 @@ def _test_set_reduce_scatter_divide_factor_mixed_prevision(
             self.assertEqual(ref_loss, loss)
             check_sharded_parity(self, ref_model, model)
 
+    @skip_if_lt_x_gpu(2)
+    def test_set_reduce_scatter_max_input_buffers(self):
+        """
+        Tests that ``set_reduce_scatter_max_input_buffers`` is a pure
+        scheduling change: changing the in-flight reduce-scatter copy-in buffer
+        cap (``2``, or an explicit ``1`` that matches the default) produces
+        identical parameters and gradients to the default single-buffer
+        behavior, across FSDP/HSDP and reshard-after-forward (under bf16 mixed
+        precision, where the retained buffer is in the reduce dtype).
+        """
+        self.run_subtests(
+            {
+                "mesh_shape": [(self.world_size,), (self.world_size // 2, 2)],
+                "reshard_after_forward": [True, False],
+                "max_input_buffers": [1, 2],
+            },
+            self._test_set_reduce_scatter_max_input_buffers,
+        )
+
+    def _test_set_reduce_scatter_max_input_buffers(
+        self,
+        mesh_shape: tuple[int] | tuple[int, int],
+        reshard_after_forward: bool,
+        max_input_buffers: int,
+    ):
+        torch.manual_seed(42)
+        model_args = ModelArgs(dropout_p=0.0, weight_tying=False)
+        ref_model = Transformer(model_args)
+        model = copy.deepcopy(ref_model)
+        mesh_dim_names = ("outer",) if len(mesh_shape) == 1 else ("outer", "inner")
+        mesh = init_device_mesh(
+            device_type.type, mesh_shape, mesh_dim_names=mesh_dim_names
+        )
+        mp_policy = MixedPrecisionPolicy(
+            param_dtype=torch.bfloat16, reduce_dtype=torch.bfloat16
+        )
+        fsdp_kwargs = {
+            "reshard_after_forward": reshard_after_forward,
+            "mesh": mesh,
+            "mp_policy": mp_policy,
+        }
+        # Reference model keeps the default (copy-in on the compute stream)
+        for module in ref_model.modules():
+            if isinstance(module, TransformerBlock):
+                fully_shard(module, **fsdp_kwargs)
+        ref_model = fully_shard(ref_model, **fsdp_kwargs)
+        ref_optim = torch.optim.AdamW(ref_model.parameters(), lr=1e-2)
+        # Test model retains multiple reduce-scatter copy-in buffers in flight
+        for module in model.modules():
+            if isinstance(module, TransformerBlock):
+                fully_shard(module, **fsdp_kwargs)
+        model = fully_shard(model, **fsdp_kwargs)
+        model.set_reduce_scatter_max_input_buffers(max_input_buffers)
+        optim = torch.optim.AdamW(model.parameters(), lr=1e-2)
+
+        torch.manual_seed(42 + self.rank)
+        inp = torch.randint(0, model_args.vocab_size, (2, 16), device=device_type.type)
+        for _ in range(3):
+            ref_loss = ref_model(inp).sum()
+            ref_loss.backward()
+            ref_optim.step()
+            loss = model(inp).sum()
+            loss.backward()
+            optim.step()
+            self.assertEqual(ref_loss, loss)
+            # Check parity before zero_grad so gradients are compared too
+            check_sharded_parity(self, ref_model, model)
+            ref_optim.zero_grad()
+            optim.zero_grad()
+
+    @skip_if_lt_x_gpu(2)
+    def test_reduce_scatter_max_input_buffers_resolves_to_max(self):
+        """
+        Per-module reduce-scatter input-buffer caps share one pipeline (a single
+        ``comm_ctx.reduce_scatter_states`` list), so mixed caps must resolve to
+        the max on the shared comm context: a lower-cap module cannot silently
+        undo a higher-cap module's retention. The per-group values stay as set.
+        """
+        torch.manual_seed(42)
+        model_args = ModelArgs(dropout_p=0.0)
+        model = Transformer(model_args)
+        blocks = [m for m in model.modules() if isinstance(m, TransformerBlock)]
+        self.assertGreaterEqual(len(blocks), 2)
+        for block in blocks:
+            fully_shard(block)
+        model = fully_shard(model)
+        # Set a non-default cap on a single block only (recurse=False); every
+        # other group (including the root) keeps the default of 1.
+        expected_max = 3
+        blocks[0].set_reduce_scatter_max_input_buffers(expected_max, recurse=False)
+
+        # Resolution runs in lazy init on the first forward.
+        inp = torch.randint(0, model_args.vocab_size, (2, 16), device=device_type.type)
+        model(inp).sum().backward()
+
+        comm_ctx = model._get_fsdp_state()._comm_ctx
+        self.assertEqual(comm_ctx.reduce_scatter_max_input_buffers, expected_max)
+        # Resolution leaves the per-group input values untouched.
+        block0_pg = blocks[0]._get_fsdp_state()._fsdp_param_groups[0]
+        block1_pg = blocks[1]._get_fsdp_state()._fsdp_param_groups[0]
+        self.assertEqual(block0_pg.reduce_scatter_max_input_buffers, expected_max)
+        self.assertEqual(block1_pg.reduce_scatter_max_input_buffers, 1)
+
+    @skip_if_lt_x_gpu(2)
+    def test_set_reduce_scatter_max_input_buffers_validation(self):
+        model = fully_shard(nn.Linear(16, 16))
+        # Non-positive -> ValueError
+        with self.assertRaises(ValueError):
+            model.set_reduce_scatter_max_input_buffers(0)
+        # Non-int -> TypeError, even though 2.5 >= 1 numerically
+        with self.assertRaises(TypeError):
+            model.set_reduce_scatter_max_input_buffers(2.5)
+        # bool is an int subclass but is not a valid count -> TypeError
+        with self.assertRaises(TypeError):
+            model.set_reduce_scatter_max_input_buffers(True)
+        # A valid value is accepted
+        model.set_reduce_scatter_max_input_buffers(2)
+
     @skip_if_lt_x_gpu(2)
     def test_set_reshard_after_forward(self):
         """
@@ -555,7 +682,7 @@ def test_set_reshard_after_forward(self):
 
     def _test_set_reshard_after_forward_by_communication_count(
         self,
-        set_reshard_after_forward: Union[bool, None],
+        set_reshard_after_forward: bool | None,
         recurse: bool,
     ):
         torch.manual_seed(42)
@@ -641,8 +768,8 @@ def test_fully_shard_backward_prefetch(self):
 
     def _test_backward_prefetch_forward_backward(
         self,
-        reshard_after_forward: Union[bool, int, None],
-        checkpoint_impl: Optional[str],
+        reshard_after_forward: bool | int | None,
+        checkpoint_impl: str | None,
     ):
         n_layers = 3
         model, optim, inp = self._init_transformer(
@@ -699,7 +826,7 @@ def _test_backward_prefetch_forward_backward(
                 optim.zero_grad(set_to_none=(iter_idx % 2 == 0))
 
     def _test_backward_prefetch_multi_forward(
-        self, reshard_after_forward: Union[bool, int], checkpoint_impl: Optional[str]
+        self, reshard_after_forward: bool | int, checkpoint_impl: str | None
     ):
         n_layers = 3
         model, _, inp = self._init_transformer(
@@ -778,7 +905,7 @@ def _test_backward_prefetch_multi_forward(
             events.clear()
 
     def _test_backward_prefetch_unused_in_backward(
-        self, reshard_after_forward: Union[bool, int, None]
+        self, reshard_after_forward: bool | int | None
     ):
         """
         Test a model with a linear module then a split into two linear modules,
@@ -1046,7 +1173,9 @@ def test_set_modules_to_backward_prefetch_inside_ac(self):
         n_layers = 3
         reshard_after_forward = True
         # use checkpoint wrapper instead of torch.utils
-        model_args = ModelArgs(n_layers=n_layers, checkpoint_activations=False)
+        model_args = ModelArgs(
+            n_layers=n_layers, checkpoint_activations=False, weight_tying=False
+        )
         model = Transformer(model_args)
         apply_activation_checkpointing(
             model, check_fn=lambda m: isinstance(m, TransformerBlock)
@@ -1074,8 +1203,10 @@ def test_set_modules_to_backward_prefetch_inside_ac(self):
 
         def set_backward_prefetch(model: Transformer) -> None:
             # tell pyre model.set_modules_to_backward_prefetch is available
-            assert isinstance(model, FSDPModule)
-            assert isinstance(model.output, FSDPModule)
+            if not isinstance(model, FSDPModule):
+                raise AssertionError(f"Expected FSDPModule, got {type(model)}")
+            if not isinstance(model.output, FSDPModule):
+                raise AssertionError(f"Expected FSDPModule, got {type(model.output)}")
 
             # mimic deepseek MOE
             # prefetch layer - 1 and its feedforward before cpu sync during a2a
@@ -1087,7 +1218,10 @@ def set_backward_prefetch(model: Transformer) -> None:
                 and model.output is not None
                 and len(model.layers) > 0
             ):
-                assert isinstance(reversed_transformer_blocks[0], FSDPModule)
+                if not isinstance(reversed_transformer_blocks[0], FSDPModule):
+                    raise AssertionError(
+                        f"Expected FSDPModule, got {type(reversed_transformer_blocks[0])}"
+                    )
                 model.output.set_modules_to_backward_prefetch(
                     [reversed_transformer_blocks[0]]
                 )
@@ -1095,13 +1229,25 @@ def set_backward_prefetch(model: Transformer) -> None:
             for transformer_block, prev_transformer_block in zip(
                 reversed_transformer_blocks, prev_transformer_blocks
             ):
-                assert isinstance(transformer_block, FSDPModule)
+                if not isinstance(transformer_block, FSDPModule):
+                    raise AssertionError(
+                        f"Expected FSDPModule, got {type(transformer_block)}"
+                    )
                 if prev_transformer_block is not None:
-                    assert isinstance(prev_transformer_block, FSDPModule)
-                    assert hasattr(prev_transformer_block.feed_forward, "w1")
-                    assert isinstance(
+                    if not isinstance(prev_transformer_block, FSDPModule):
+                        raise AssertionError(
+                            f"Expected FSDPModule, got {type(prev_transformer_block)}"
+                        )
+                    if not hasattr(prev_transformer_block.feed_forward, "w1"):
+                        raise AssertionError(
+                            "Expected prev_transformer_block.feed_forward to have 'w1' attribute"
+                        )
+                    if not isinstance(
                         prev_transformer_block.feed_forward.w1, FSDPModule
-                    )
+                    ):
+                        raise AssertionError(
+                            f"Expected FSDPModule, got {type(prev_transformer_block.feed_forward.w1)}"
+                        )
                     transformer_block.set_modules_to_backward_prefetch(
                         [
                             prev_transformer_block,
@@ -1109,7 +1255,10 @@ def set_backward_prefetch(model: Transformer) -> None:
                         ]
                     )
                 elif model.tok_embeddings is not None:
-                    assert isinstance(model.tok_embeddings, FSDPModule)
+                    if not isinstance(model.tok_embeddings, FSDPModule):
+                        raise AssertionError(
+                            f"Expected FSDPModule, got {type(model.tok_embeddings)}"
+                        )
                     transformer_block.set_modules_to_backward_prefetch(
                         [model.tok_embeddings]
                     )
@@ -1179,12 +1328,6 @@ def set_backward_prefetch(model: Transformer) -> None:
                     "tok_embeddings, pos_embeddings",
                     TrainingState.POST_BACKWARD,
                 ),
-                (
-                    "reshard",
-                    "tok_embeddings, pos_embeddings",
-                    TrainingState.POST_BACKWARD,
-                ),
-                ("reshard", "norm, output", TrainingState.POST_BACKWARD),
             ]
             self.assertEqual(events, expected_backward_events)
             events.clear()
@@ -1247,12 +1390,6 @@ def set_backward_prefetch(model: Transformer) -> None:
                     "tok_embeddings, pos_embeddings",
                     TrainingState.POST_BACKWARD,
                 ),
-                (
-                    "reshard",
-                    "tok_embeddings, pos_embeddings",
-                    TrainingState.POST_BACKWARD,
-                ),
-                ("reshard", "norm, output", TrainingState.POST_BACKWARD),
             ]
             self.assertEqual(events, expected_backward_events)
             events.clear()
@@ -1260,7 +1397,9 @@ def set_backward_prefetch(model: Transformer) -> None:
     @skip_if_lt_x_gpu(2)
     def test_fully_shard_multi_module_backward_prefetch(self):
         n_layers = 5
-        model_args = ModelArgs(n_layers=n_layers, checkpoint_activations=True)
+        model_args = ModelArgs(
+            n_layers=n_layers, checkpoint_activations=True, weight_tying=False
+        )
         model = Transformer(model_args)
         for i in range(n_layers):
             if i == 0:
@@ -1434,8 +1573,8 @@ def test_backward_misprefetch(self):
     def _init_transformer(
         self,
         n_layers: int,
-        reshard_after_forward: Union[bool, int, None],
-        checkpoint_impl: Optional[str],
+        reshard_after_forward: bool | int | None,
+        checkpoint_impl: str | None,
     ):
         model_args = ModelArgs(
             n_layers=n_layers, checkpoint_activations=(checkpoint_impl == "utils")
@@ -1627,8 +1766,15 @@ def _run(cls, *args, **kwargs):
     @skip_if_lt_x_gpu(2)
     # The NCCL PG refuses to allocate tensors if multicast is unavailable, see
     # https://github.com/pytorch/pytorch/blob/503362d019b3782581492af7767945dbd75ca1c9/torch/csrc/distributed/c10d/ProcessGroupNCCL.cpp#L5634
-    @requires_multicast_support()
     def test_fully_shard_alloc_from_pg(self):
+        # Run this check inside test instead of using @requires_multicast_support().
+        # The decorator would trigger an initialization of SymmMem allocator
+        # when Python statically initializes classes in this file, causing
+        # SymmMem to fix the allocate backend to "CUDA". This is unfriendly for
+        # other tests in this file that requires NCCL backend
+        if not _SymmetricMemory.has_multicast_support(DeviceType.CUDA, 0):
+            self.skipTest("multicast support is not available")
+
         torch.manual_seed(42)
         model_args = ModelArgs()
         model = Transformer(model_args)
@@ -1680,6 +1826,64 @@ def test_exception_when_used_together_with_comm_hooks(self):
             model.set_allocate_memory_from_process_group_for_comm(True)
 
 
+@requires_cuda_p2p_access()
+@skip_but_pass_in_sandcastle_if(not TEST_MULTIGPU, "Not enough GPUs to run the test")
+@unittest.skipIf(
+    not PLATFORM_SUPPORTS_SYMM_MEM, "SymmMem is not supported on this platform"
+)
+@skipCUDAIf(TEST_WITH_ROCM, "requires NVIDIA GPUs")
+@skipCUDAIf(not SM90OrLater, "requires sm90+")
+class TestFullyShardSymmMem(MultiProcContinuousTest):
+    @classmethod
+    def backend_str(cls) -> str | None:
+        return "nccl"
+
+    @classmethod
+    def opts(cls):
+        if not dist.is_nccl_available():
+            return None
+        # Enable Zero-CTA policy for CE collectives
+        opts = dist.ProcessGroupNCCL.Options()
+        opts.config.cta_policy = dist.ProcessGroupNCCL.NCCL_CTA_POLICY_ZERO
+        return opts
+
+    @property
+    def device(self) -> torch.device:
+        return torch.device("cuda", self.rank)
+
+    @parametrize("sum_reduction", [True, False])
+    def test_fully_shard_symm_mem(self, sum_reduction: bool):
+        torch.manual_seed(42 + self.rank)
+        device = torch.device("cuda", self.rank)
+        torch.cuda.set_device(device)
+        seq_len = 64
+        model_args = ModelArgs()
+        model_args.dim = 4096
+        model_args.max_seq_len = seq_len
+        model = Transformer(model_args).to(device)
+        for module in model.modules():
+            if isinstance(module, TransformerBlock):
+                fully_shard(module)
+                module.set_force_sum_reduction_for_comms(sum_reduction)
+                module.set_symm_mem_for_comm()
+        fully_shard(model)
+        model.set_force_sum_reduction_for_comms(sum_reduction)
+        model.set_symm_mem_for_comm()
+
+        bs = 4
+        inp = torch.randint(0, model_args.vocab_size, (bs, seq_len), device=device)
+
+        def run():
+            loss = model(inp)
+            loss.sum().backward()
+
+        run()
+        torch.cuda.synchronize(device)
+
+
+instantiate_parametrized_tests(TestFullyShardSymmMem)
+
+
 class TestFullyShardForceSumReduction(FSDPTest):
     # The messages might change when we move to a different NCCL version.
     # Please update this test if it starts failing.
diff --git a/test/distributed/_composable/fsdp/test_fully_shard_compile.py b/test/distributed/_composable/fsdp/test_fully_shard_compile.py
index 257e445f6c385..4a3bbd2734ca5 100644
--- a/test/distributed/_composable/fsdp/test_fully_shard_compile.py
+++ b/test/distributed/_composable/fsdp/test_fully_shard_compile.py
@@ -1,60 +1,31 @@
 # Owner(s): ["oncall: distributed"]
 
 
-import contextlib
 import copy
-import functools
-import itertools
-import logging
 import unittest
-from collections import defaultdict
-from unittest import mock
 
 import torch
+import torch._dynamo.compiled_autograd as compiled_autograd
 import torch._dynamo.testing
-import torch.nn.functional as F
-from torch import nn
+import torch.nn as nn
 from torch._dynamo.utils import counters
-from torch._inductor import comms
-from torch._inductor.utils import is_fallback_op, run_and_get_code
 from torch.distributed.device_mesh import init_device_mesh
 from torch.distributed.fsdp import (
     fully_shard,
     FullyShardedDataParallel as FSDP,
+    MixedPrecisionPolicy,
     ShardingStrategy,
 )
 from torch.distributed.fsdp._fully_shard._fsdp_common import TrainingState
 from torch.distributed.fsdp._fully_shard._fsdp_param_group import FSDPParamGroup
-from torch.testing import FileCheck
-from torch.testing._internal.common_distributed import (
-    at_least_x_gpu,
-    skip_if_lt_x_gpu,
-    sm_is_or_higher_than,
-)
+from torch.testing._internal.common_distributed import skip_if_lt_x_gpu
 from torch.testing._internal.common_fsdp import FSDPTest, get_devtype, MLP
 from torch.testing._internal.common_utils import run_tests
-from torch.testing._internal.distributed._tensor.common_dtensor import (
-    ModelArgs,
-    Transformer,
-)
 from torch.testing._internal.inductor_utils import HAS_GPU
 
 
 device_type = torch.device(get_devtype())
 
-log = logging.getLogger(__name__)
-
-
-def _count_op_in_graph(graph, op):
-    return sum(1 for node in graph.nodes if node.target is op)
-
-
-def _is_fallback_op_in_snodes(snodes, op):
-    return any(is_fallback_op(snode.node, op) for snode in snodes)
-
-
-orig_F_scaled_dot_product_attention = F.scaled_dot_product_attention
-
 
 class Mod(torch.nn.Module):
     def __init__(self):
@@ -74,72 +45,118 @@ class TestFullyShardCompileCompute(FSDPTest):
     @unittest.skipIf(not HAS_GPU, "Inductor+gpu needs triton and recent GPU arch")
     @skip_if_lt_x_gpu(2)
     def test_disable_compiling_hooks(self):
-        self.run_subtests(
-            {
-                "skip_fsdp_hooks": [False, True],
-            },
-            self._test_disable_compiling_hooks,
-        )
-
-    def _test_disable_compiling_hooks(
-        self,
-        skip_fsdp_hooks: bool,
-    ):
+        """Verify that dynamo never traces into FSDP hooks in forward or backward."""
         torch._dynamo.reset()
         trace_rules_check_count = 0
         HOOKS_FILE_NAME = "torch/distributed/fsdp/_fully_shard/_fsdp_state.py"
-        HOOK_WRAPPER_NAME = "fsdp_hook_wrapper"
 
         def patched_trace_rules_check(*args, **kwargs):
             nonlocal trace_rules_check_count
             f_code = args[0]
-            if (
-                hasattr(f_code, "co_filename")
-                and f_code.co_filename.endswith(HOOKS_FILE_NAME)
-                and f_code.co_name != HOOK_WRAPPER_NAME
+            if hasattr(f_code, "co_filename") and f_code.co_filename.endswith(
+                HOOKS_FILE_NAME
             ):
                 trace_rules_check_count += 1
             return orig_trace_rules_check(*args, **kwargs)
 
-        original_skip_fsdp_hooks = torch._dynamo.config.skip_fsdp_hooks
         orig_trace_rules_check = torch._dynamo.trace_rules.check
         torch.distributed.barrier()
-        torch._dynamo.config.skip_fsdp_hooks = skip_fsdp_hooks
         torch._dynamo.trace_rules.check = patched_trace_rules_check
         model = MLP(4).to(device_type)
         fully_shard(model)
         model.compile()
-        model(torch.randn((4, 4), device=device_type))
+        out = model(torch.randn((4, 4), device=device_type))
+        out.sum().backward()
         torch.distributed.barrier()
-        torch._dynamo.config.skip_fsdp_hooks = original_skip_fsdp_hooks
         torch._dynamo.trace_rules.check = orig_trace_rules_check
-        if skip_fsdp_hooks:
-            self.assertEqual(trace_rules_check_count, 0)
-        else:
-            self.assertTrue(trace_rules_check_count > 0)
+        self.assertEqual(trace_rules_check_count, 0)
 
+    @unittest.skipIf(not HAS_GPU, "Inductor+gpu needs triton and recent GPU arch")
+    @skip_if_lt_x_gpu(2)
+    def test_compiled_autograd_fsdp2_backward(self):
+        """
+        Verify that compiled autograd with FSDP2 backward does not crash and
+        that FSDP hooks cause graph breaks (due to _dynamo_disable).
+
+        For MLP(4) with a single fully_shard, the compiled autograd backward
+        graph has 2 graph breaks from _dynamo_disable on FSDP hooks:
+          1. _pre_backward tensor hook (registered on FSDP module output)
+          2. post_backward (called from RegisterPostBackwardFunction.backward)
+        _root_post_backward_final_callback runs outside the compiled autograd
+        scope as a queued engine callback.
+        """
+        torch._dynamo.reset()
+        counters.clear()
+        model = MLP(4).to(device_type)
+        ref_model = copy.deepcopy(model)
+        fully_shard(model)
+        fully_shard(ref_model)
+        inp = torch.randn((4, 4), device=device_type)
 
-@unittest.skipIf(not HAS_GPU, "Inductor+gpu needs triton and recent GPU arch")
-class TestFullyShardCompile(FSDPTest):
-    fake_pg = not at_least_x_gpu(2)
+        # Eager reference
+        ref_out = ref_model(inp)
+        ref_out.sum().backward()
+
+        backend_count = 0
+        orig_aot_eager = torch._dynamo.lookup_backend("aot_eager")
+
+        def counting_backend(gm, example_inputs):
+            nonlocal backend_count
+            backend_count += 1
+            return orig_aot_eager(gm, example_inputs)
 
-    # This method is an override of the base class.
-    # Tests in this class requires bf16 support, so SM arch must be 80 or
-    # higher.
-    def skipTestForOldSm(self):
-        # Assumption: This test class is only run on GPU. See `HAS_GPU` check at
-        # the top of the class.
-        device = torch.device(
-            device_type.type,
-            self.rank % torch.get_device_module(device_type).device_count(),
+        def compiler_fn(gm):
+            return torch.compile(gm, backend=counting_backend)
+
+        with compiled_autograd._enable(compiler_fn):
+            out = model(inp)
+            out.sum().backward()
+
+        self.assertEqual(out, ref_out)
+        self.assertEqual(counters["compiled_autograd"]["captures"], 1)
+        # 2 graph breaks from _dynamo_disable on FSDP hooks split the
+        # compiled autograd graph into 2 sub-graphs compiled by the backend:
+        #   graph break 1: _pre_backward tensor hook
+        #   graph break 2: post_backward (from RegisterPostBackwardFunction)
+        self.assertEqual(backend_count, 2)
+
+    @unittest.skipIf(not HAS_GPU, "Inductor+gpu needs triton and recent GPU arch")
+    @skip_if_lt_x_gpu(2)
+    def test_compile_optimizer_uneven_shard(self):
+        # Regression test for https://github.com/pytorch/pytorch/issues/176667
+        # When a param dim is not divisible by world_size, FSDP2 creates
+        # param._local_tensor as a narrow view of an N-D padded tensor, while
+        # grad._local_tensor is a view of a 1-D flat gradient buffer. Dynamo
+        # must not reuse param's symbolic context for the grad.
+        torch._dynamo.reset()
+        mesh = init_device_mesh(device_type.type, (self.world_size,))
+        mp_policy = MixedPrecisionPolicy(
+            param_dtype=torch.bfloat16, reduce_dtype=torch.float32
         )
-        if (
-            device_type.type == "cuda"
-            and not torch.version.hip
-            and not sm_is_or_higher_than(device, 8, 0)
-        ):
-            self.skipTest("bf16 requires sm >= 8.0")
+        # 47 out_channels not divisible by world_size=2, forcing uneven sharding
+        with torch.device("meta"):
+            model = nn.Conv2d(3, 47, 3, padding=1)
+        fully_shard(model, mesh=mesh, mp_policy=mp_policy, reshard_after_forward=False)
+        model.to_empty(device=device_type)
+        with torch.no_grad():
+            for p in model.parameters():
+                p.uniform_(-0.01, 0.01)
+
+        opt = torch.optim.Adam(model.parameters(), lr=1e-3)
+        x = torch.randn(4, 3, 8, 8, device=device_type, dtype=torch.bfloat16)
+
+        # Eager warmup
+        model(x).sum().backward()
+        opt.step()
+        opt.zero_grad(set_to_none=True)
+
+        # Compiled optimizer step
+        model(x).sum().backward()
+        torch.compile(opt.step)()
 
+
+@unittest.skipIf(not HAS_GPU, "Inductor+gpu needs triton and recent GPU arch")
+class TestFullyShardCompile(FSDPTest):
     def test_dynamo_trace_use_training_state(self):
         torch._dynamo.reset()
         # Construct a dummy FSDPParamGroup, since we just want to test the `use_training_state` ctx manager.
@@ -195,968 +212,6 @@ def f(x):
         torch.compile(f, backend="aot_eager")(x)
         self.assertEqual(x, ref_x)
 
-    def _get_resize_count_in_fx_graph(self, graph: torch.fx.Graph):
-        resize_count = 0
-        for node in graph.nodes:
-            if (
-                node.op == "call_function"
-                and node.target == torch.ops.inductor.resize_storage_bytes_.default
-            ):
-                resize_count += 1
-        return resize_count
-
-    def _assert_no_aliased_unsharded_params_in_graph_inputs(
-        self, model, graph: torch.fx.Graph
-    ) -> None:
-        # FSDP2 unsharded params are mutated in the graph without going through functionalization.
-        # Therefore, we want to make sure they don't have aliases in the graph inputs, to make it easier
-        # for us to do the replacement of unsharded params with the all-gathered temporary buffer directly
-        # in downstream users in the graph.
-        storage_id_to_graph_inputs = defaultdict(list)
-        unsharded_param_graph_inputs = set()
-        for node in graph.nodes:
-            if (
-                node.op == "call_function"
-                and node.target
-                in [
-                    torch.ops.inductor.resize_storage_bytes_.default,
-                    torch.ops.fsdp.copy_.default,
-                ]
-                and node.args[0].op == "placeholder"
-            ):
-                unsharded_param_graph_inputs.add(node.args[0])
-        assert len(unsharded_param_graph_inputs) > 0
-        assert len(unsharded_param_graph_inputs) == len(list(model.parameters())), """\
-Expected all model parameters to be wrapped by FSDP2 and
-have their unsharded version as graph input, but it's not true!
-"""
-        no_aliased_unsharded_params_in_graph_inputs = True
-        err_msg = ""
-        for aliased_graph_inputs in storage_id_to_graph_inputs.values():
-            if len(aliased_graph_inputs) > 1 and any(
-                x in unsharded_param_graph_inputs for x in aliased_graph_inputs
-            ):
-                no_aliased_unsharded_params_in_graph_inputs = False
-                err_msg += f"""\n
-Found aliased unsharded param in graph inputs: {aliased_graph_inputs},
-val.shape: {[node.meta["val"].shape for node in aliased_graph_inputs]},
-"""
-        self.assertTrue(no_aliased_unsharded_params_in_graph_inputs, err_msg)
-
-    def _remove_fsdp2_unsharded_param_graph_input_usage_with_optional_checks(
-        self, model, *, bwd_resize_count_before_pass=None, fwd_fullgraph=False
-    ):
-        def _run_with_checks(graph, orig_fn):
-            if (
-                self._is_bwd_fx_graph(graph)
-                and bwd_resize_count_before_pass is not None
-            ):
-                self.assertEqual(
-                    bwd_resize_count_before_pass,
-                    self._get_resize_count_in_fx_graph(graph),
-                )
-            self._assert_no_aliased_unsharded_params_in_graph_inputs(model, graph)
-            orig_fn(graph)
-
-        if fwd_fullgraph:
-            return mock.patch.object(
-                comms,
-                "remove_fsdp2_unsharded_param_graph_input_usage",
-                functools.partial(
-                    _run_with_checks,
-                    orig_fn=comms.remove_fsdp2_unsharded_param_graph_input_usage,
-                ),
-            )
-        else:
-            return contextlib.nullcontext()
-
-    def _check_fsdp_copy_and_resize_ops_count_in_graph(
-        self,
-        graph,
-        *,
-        fwd_copy_count,
-        fwd_resize_count,
-        bwd_copy_count,
-        bwd_resize_count,
-    ):
-        def _check_count(copy_count, resize_count):
-            actual_copy_count = _count_op_in_graph(graph, torch.ops.fsdp.copy_.default)
-            self.assertEqual(
-                actual_copy_count,
-                copy_count,
-                f"Unexpected number of `fsdp.copy_` ops (expected {copy_count}, got {actual_copy_count}) in graph: {graph}",
-            )
-
-            actual_resize_count = _count_op_in_graph(
-                graph, torch.ops.inductor.resize_storage_bytes_.default
-            )
-            self.assertEqual(
-                actual_resize_count,
-                resize_count,
-                f"Unexpected number of `inductor.resize_storage_bytes_` ops (expected {resize_count}, got {actual_resize_count}) in graph: {graph}",  # noqa: B950
-            )
-
-        if not torch._dynamo.compiled_autograd.in_compiled_autograd_region:
-            _check_count(fwd_copy_count, fwd_resize_count)  # fwd graph
-        else:
-            _check_count(bwd_copy_count, bwd_resize_count)  # bwd graph
-
-    def _reinplace_all_gather_with_optional_checks(self, fwd_fullgraph):
-        def _run_with_checks(graph, orig_fn):
-            if self.world_size > 1:
-                self.assertGreater(
-                    _count_op_in_graph(
-                        graph, torch.ops._c10d_functional.all_gather_into_tensor.default
-                    ),
-                    0,
-                )
-            elif self.world_size == 1:
-                self.assertEqual(
-                    _count_op_in_graph(
-                        graph, torch.ops._c10d_functional.all_gather_into_tensor.default
-                    ),
-                    0,
-                )
-
-            orig_fn(graph)
-
-            self.assertEqual(
-                _count_op_in_graph(
-                    graph, torch.ops._c10d_functional.all_gather_into_tensor.default
-                ),
-                0,
-            )
-
-            if self.world_size > 1:
-                self.assertGreater(
-                    _count_op_in_graph(
-                        graph,
-                        torch.ops._c10d_functional.all_gather_into_tensor_out.default,
-                    ),
-                    0,
-                )
-            else:
-                self.assertEqual(
-                    _count_op_in_graph(
-                        graph,
-                        torch.ops._c10d_functional.all_gather_into_tensor_out.default,
-                    ),
-                    0,
-                )
-
-        if fwd_fullgraph:
-            return mock.patch.object(
-                comms,
-                "reinplace_fsdp_all_gather",
-                functools.partial(
-                    _run_with_checks,
-                    orig_fn=comms.reinplace_fsdp_all_gather,
-                ),
-            )
-        else:
-            return contextlib.nullcontext()
-
-    def _is_fwd_graph(self, snodes):
-        ag_copy_in_snode = None
-        for snode in snodes:
-            if is_fallback_op(snode.node, torch.ops.fsdp.all_gather_copy_in.default):
-                ag_copy_in_snode = snode
-                break
-        self.assertTrue(ag_copy_in_snode is not None)
-        if any(
-            dep.name.startswith("primals_")
-            for dep in ag_copy_in_snode.read_writes.reads
-        ):
-            return True
-        else:
-            return False
-
-    def _is_bwd_fx_graph(self, graph):
-        for node in graph.nodes:
-            if (
-                node.op == "call_function"
-                and node.target
-                == torch.ops._c10d_functional.reduce_scatter_tensor.default
-            ):
-                return True
-        return False
-
-    def _maybe_run_decide_global_ordering_of_comms_with_checks(self, fwd_fullgraph):
-        def _check_fsdp_ops_in_snodes(snodes, is_fwd_graph, expect=True):
-            assert_method = self.assertTrue if expect else self.assertFalse
-            common_ops = {
-                torch.ops.fsdp.all_gather_copy_in.default,
-                torch.ops._c10d_functional.all_gather_into_tensor_out.default,
-                torch.ops.fsdp.split_with_sizes_copy.default,
-            }
-            bwd_only_ops = {
-                torch.ops.fsdp.chunk_cat.default,
-                torch.ops._c10d_functional.reduce_scatter_tensor.default,
-            }
-            for op in common_ops:
-                assert_method(
-                    _is_fallback_op_in_snodes(
-                        snodes,
-                        op,
-                    ),
-                    msg=f"{op}",
-                )
-            if not is_fwd_graph:
-                for op in bwd_only_ops:
-                    assert_method(
-                        _is_fallback_op_in_snodes(
-                            snodes,
-                            op,
-                        ),
-                        msg=f"{op}",
-                    )
-
-        def _decide_global_ordering_of_comms_with_checks(
-            snodes, name_to_buf, name_to_fused_node, orig_fn
-        ):
-            is_fwd_graph = self._is_fwd_graph(snodes)
-            _check_fsdp_ops_in_snodes(snodes, is_fwd_graph, expect=True)
-            new_snodes = orig_fn(snodes, name_to_buf, name_to_fused_node)
-            _check_fsdp_ops_in_snodes(new_snodes, is_fwd_graph, expect=False)
-            return new_snodes
-
-        if fwd_fullgraph:
-            return mock.patch.object(
-                comms,
-                "decide_global_ordering_of_comms",
-                functools.partial(
-                    _decide_global_ordering_of_comms_with_checks,
-                    orig_fn=comms.decide_global_ordering_of_comms,
-                ),
-            )
-        else:
-            return contextlib.nullcontext()
-
-    def inductor_code_check_no_compute_op(self, file_check):
-        return (
-            file_check.check_not(" = aten.")
-            .check_not(" = extern_kernels.")
-            .check_not(" = triton_")
-            .check_not(" = torch.ops.")
-            .check_not(" = inductor_ops.")
-            .check_not("    aten.")
-            .check_not("    extern_kernels.")
-            .check_not("    triton_")
-            .check_not("    torch.ops.")
-            .check_not("    inductor_ops.")
-        )
-
-    def inductor_code_check_fsdp_all_gather(
-        self,
-        file_check,
-        overlapped_compute_op_str,
-        last_all_gather=False,
-    ):
-        file_check = file_check.check("torch.ops.fsdp.all_gather_copy_in.")
-        file_check = self.inductor_code_check_no_compute_op(file_check)
-        file_check = file_check.check(
-            "torch.ops._c10d_functional.all_gather_into_tensor_out."
-        )
-        # Checks that AGWait is delayed, making the AG overlap with some compute op.
-        if overlapped_compute_op_str is not None:
-            file_check = file_check.check(f"{overlapped_compute_op_str}")
-        file_check = file_check.check("torch.ops._c10d_functional.wait_tensor.")
-        file_check = self.inductor_code_check_no_compute_op(file_check)
-        file_check = file_check.check("torch.ops.fsdp.split_with_sizes_copy.")
-        if not last_all_gather:
-            # Checks that there is no compute op between this AGWait and next AG.
-            file_check = self.inductor_code_check_no_compute_op(file_check)
-        return file_check
-
-    def inductor_code_check_fsdp_reduce_scatter(
-        self, file_check, overlapped_compute_op_str
-    ):
-        file_check = file_check.check("torch.ops.fsdp.chunk_cat.")
-        file_check = self.inductor_code_check_no_compute_op(file_check)
-        file_check = file_check.check(
-            "torch.ops._c10d_functional.reduce_scatter_tensor."
-        )
-        # Checks that RSWait is delayed, making the RS overlap with some compute op.
-        if overlapped_compute_op_str is not None:
-            file_check = file_check.check(f"{overlapped_compute_op_str}")
-        file_check = file_check.check("torch.ops._c10d_functional.wait_tensor.")
-        return file_check
-
-    @unittest.skipIf(not HAS_GPU, "Inductor+gpu needs triton and recent GPU arch")
-    def test_compiled_autograd_ctx(self):
-        self.skipTestForOldSm()
-        with (
-            torch._dynamo.config.patch(skip_fsdp_hooks=False),
-            torch._functorch.config.patch(
-                recompute_views=True,
-            ),
-        ):
-            inputs = torch.randn(8, 8)
-            model = torch.nn.Linear(8, 8)
-            fully_shard(model)
-            model_compiled = torch.compile(model, backend="inductor")
-            for i in range(10):
-                torch.compiler.set_stance(
-                    "force_eager" if i < 1 else "default"
-                )  # eager warmup for 1 iteration
-                with torch._dynamo.compiled_autograd._enable(
-                    torch.compile(backend="inductor", fullgraph=True)
-                ):
-                    out = model_compiled(inputs)
-                    out.sum().backward()
-
-    def _test_traceable_fsdp(
-        self,
-        model_init_fn,
-        input_creation_fn,
-        backend,
-        fwd_fullgraph,
-        *,
-        bwd_resize_count_before_inductor=None,
-    ):
-        def fwd_bwd(model, inp):
-            out = model(inp)
-            loss = out.sum()
-            loss.backward()
-            return loss
-
-        def run_iters(
-            fwd_bwd_func,
-            optim,
-            n_iter=10,
-            compiled_autograd_backend=None,
-        ):
-            torch.manual_seed(42)
-            losses = []
-            for i in range(n_iter):
-                # eager warmup for 1 iteration, so that all FSDP2 lazy-initialization is done in eager
-                torch.compiler.set_stance("force_eager" if i < 1 else "default")
-                inp = input_creation_fn()
-                loss = fwd_bwd_func(inp)
-                losses.append(loss.item())
-                optim.step()
-                optim.zero_grad(set_to_none=True)
-            return losses
-
-        def test_compiled():
-            model, optim = model_init_fn()
-            fwd_bwd_fn = functools.partial(fwd_bwd, model)
-
-            counters.clear()
-            with self._remove_fsdp2_unsharded_param_graph_input_usage_with_optional_checks(
-                model,
-                bwd_resize_count_before_pass=bwd_resize_count_before_inductor,
-                fwd_fullgraph=fwd_fullgraph,
-            ):
-                fwd_bwd_fn_compiled = torch.compile(
-                    fwd_bwd_fn,
-                    backend=backend,
-                    # NOTE: we can't set `fullgraph=True` here because we will always graph-break
-                    # on `loss.backward()` call in `fwd_bwd()`. This is okay as long as
-                    # it's the only graph-break in forward pass.
-                    fullgraph=False,
-                )
-                res = run_iters(
-                    fwd_bwd_fn_compiled,
-                    optim,
-                    compiled_autograd_backend=backend,
-                )
-                if fwd_fullgraph:
-                    self.assertEqual(len(counters["graph_break"]), 1)
-                    self.assertExpectedInline(
-                        next(iter(counters["graph_break"].keys())).split(
-                            "\n For more details"
-                        )[0],
-                        """\
-autograd.grad with compiled autograd
-  Explanation: torch.autograd.grad() inside torch.compile is not supported when compiled autograd is enabled. These two features have conflicting requirements for how the autograd graph is traced.
-  Hint: Disable compiled autograd by removing the compiled_autograd context manager.
-  Hint: Or move the autograd.grad() call outside the torch.compile region.
-  Hint: Or restructure your code so autograd.grad() and compiled_autograd don't overlap.
-
-  Developer debug context: compiled_autograd is currently enabled
-""",  # noqa: B950
-                    )
-                else:
-                    self.assertGreater(len(counters["graph_break"]), 1)
-                return res
-
-        def test_eager():
-            model, optim = model_init_fn()
-            fwd_bwd_fn = functools.partial(fwd_bwd, model)
-
-            res = run_iters(fwd_bwd_fn, optim)
-            return res
-
-        torch._dynamo.reset()
-        torch._dynamo.compiled_autograd.reset()
-        with (
-            torch._dynamo.config.patch(
-                compiled_autograd=True,
-                compiled_autograd_kwargs_override={
-                    "fullgraph": True,
-                },
-                inline_inbuilt_nn_modules=True,
-                skip_fsdp_hooks=False,
-                trace_autograd_ops=True,
-            ),
-            torch._functorch.config.patch(
-                enable_autograd_cache=False,
-                recompute_views=True,
-            ),
-            torch._inductor.config.patch(
-                force_disable_caches=True,
-                reorder_for_compute_comm_overlap=True,
-                reorder_for_compute_comm_overlap_passes=[
-                    "sink_waits",
-                    "raise_comms",
-                    "reorder_compute_for_overlap",
-                ],
-            ),
-        ):
-            losses_compiled = test_compiled()
-        losses_eager = test_eager()
-        if not self.fake_pg:
-            for loss_compiled, loss_eager in zip(losses_compiled, losses_eager):
-                self.assertTrue(
-                    torch.allclose(
-                        torch.tensor(loss_compiled),
-                        torch.tensor(loss_eager),
-                        rtol=1e-5,
-                        atol=1e-8,
-                    ),
-                    f"{loss_compiled} vs {loss_eager}",
-                )
-
-    def _create_simple_mlp_factory_fns(self):
-        hidden_dim = 16
-
-        def model_init_fn():
-            torch.manual_seed(self.rank)
-            fsdp_config = {}
-            model = nn.Sequential(
-                nn.Linear(hidden_dim, hidden_dim, device=device_type),
-                nn.ReLU(),
-                nn.Linear(hidden_dim, hidden_dim, device=device_type),
-                nn.ReLU(),
-                nn.Linear(hidden_dim, hidden_dim, device=device_type),
-            )
-            fully_shard(model, reshard_after_forward=True, **fsdp_config)
-            optim = torch.optim.SGD(model.parameters(), lr=1e-4)
-            return model, optim
-
-        def input_creation_fn():
-            torch.manual_seed(self.rank)
-            inp = torch.randn((2, hidden_dim), device=device_type, requires_grad=False)
-            return inp
-
-        return model_init_fn, input_creation_fn
-
-    @unittest.skipIf(not HAS_GPU, "Inductor+gpu needs triton and recent GPU arch")
-    def test_simple_mlp_fullgraph_backend_aot_eager(self):
-        self._test_traceable_fsdp(
-            *self._create_simple_mlp_factory_fns(), "aot_eager", fwd_fullgraph=True
-        )
-
-    @unittest.skipIf(not HAS_GPU, "Inductor+gpu needs triton and recent GPU arch")
-    def test_simple_mlp_fullgraph_backend_aot_eager_decomp_partition(self):
-        self._test_traceable_fsdp(
-            *self._create_simple_mlp_factory_fns(),
-            "aot_eager_decomp_partition",
-            fwd_fullgraph=True,
-        )
-
-    @unittest.skipIf(not HAS_GPU, "Inductor+gpu needs triton and recent GPU arch")
-    def test_simple_mlp_fullgraph_backend_inductor(self):
-        self.skipTestForOldSm()
-        self._test_traceable_fsdp(
-            *self._create_simple_mlp_factory_fns(), "inductor", fwd_fullgraph=True
-        )
-
-    def _create_nested_fully_shard_factory_fns(self, fwd_fullgraph):
-        hidden_dim = 16
-
-        class TestSubmodule(nn.Module):
-            def __init__(self, hidden_dim):
-                super().__init__()
-                self.param1 = nn.Parameter(
-                    torch.zeros(
-                        hidden_dim, hidden_dim, dtype=torch.float, device=device_type
-                    )
-                )
-                self.param2 = nn.Parameter(
-                    torch.zeros(hidden_dim, dtype=torch.float, device=device_type)
-                )
-
-            def forward(self, x):
-                ret = torch.matmul(x, self.param1)
-                if not fwd_fullgraph:
-                    torch._dynamo.graph_break()
-                ret = ret * self.param2
-                ret = torch.relu(ret)
-                return ret
-
-        class TestModule(nn.Module):
-            def __init__(self, n_layers):
-                super().__init__()
-                self.layers = torch.nn.ModuleList()
-                for _ in range(n_layers):
-                    self.layers.append(TestSubmodule(hidden_dim))
-
-            def forward(self, x):
-                # Intentionally reusing all layers a few times,
-                # to test "multiple all-gathers for the same parameter" case.
-                # Case 1: rerun the same layer twice
-                for layer_id in range(len(self.layers)):
-                    for _ in range(2):
-                        x = self.layers[layer_id](x)
-                # Case 2: iterate through all layers twice
-                for layer in self.layers:
-                    x = layer(x)
-                for layer in self.layers:
-                    x = layer(x)
-                return x
-
-        def model_init_fn():
-            torch.manual_seed(self.rank)
-            fsdp_config = {}
-            mesh = init_device_mesh(device_type.type, (self.world_size,))
-            model = TestModule(n_layers=3)
-            for mod in model.layers:
-                fully_shard(mod, mesh=mesh, reshard_after_forward=True, **fsdp_config)
-            model = fully_shard(
-                model, mesh=mesh, reshard_after_forward=True, **fsdp_config
-            )
-            optim = torch.optim.SGD(model.parameters(), lr=1e-4)
-            return model, optim
-
-        def input_creation_fn():
-            torch.manual_seed(self.rank)
-            inp = torch.randn((2, hidden_dim), device=device_type, requires_grad=False)
-            return inp
-
-        return model_init_fn, input_creation_fn
-
-    @unittest.skipIf(not HAS_GPU, "Inductor+gpu needs triton and recent GPU arch")
-    def test_nested_fully_shard_backend_aot_eager(self):
-        # TODO: fix fwd_fullgraph=False case
-        for fwd_fullgraph in [True]:
-            self._test_traceable_fsdp(
-                *self._create_nested_fully_shard_factory_fns(
-                    fwd_fullgraph=fwd_fullgraph
-                ),
-                "aot_eager",
-                fwd_fullgraph=fwd_fullgraph,
-            )
-
-    @unittest.skipIf(not HAS_GPU, "Inductor+gpu needs triton and recent GPU arch")
-    def test_nested_fully_shard_backend_aot_eager_decomp_partition(self):
-        # TODO: fix fwd_fullgraph=False case
-        for fwd_fullgraph in [True]:
-            self._test_traceable_fsdp(
-                *self._create_nested_fully_shard_factory_fns(
-                    fwd_fullgraph=fwd_fullgraph
-                ),
-                "aot_eager_decomp_partition",
-                fwd_fullgraph=fwd_fullgraph,
-            )
-
-    def _test_nested_fully_shard_backend_inductor_fullgraph_True(self):
-        self.skipTestForOldSm()
-        for fwd_fullgraph in [True]:
-            with (
-                self._reinplace_all_gather_with_optional_checks(fwd_fullgraph),
-                torch._inductor.config.patch(
-                    post_grad_custom_post_pass=(
-                        functools.partial(
-                            self._check_fsdp_copy_and_resize_ops_count_in_graph,
-                            fwd_copy_count=0,
-                            fwd_resize_count=0,
-                            bwd_copy_count=0,
-                            bwd_resize_count=0,
-                        )
-                        if fwd_fullgraph
-                        else None
-                    )
-                ),
-            ):
-                _, triton_codes = run_and_get_code(
-                    lambda: self._test_traceable_fsdp(
-                        *self._create_nested_fully_shard_factory_fns(
-                            fwd_fullgraph=fwd_fullgraph
-                        ),
-                        "inductor",
-                        fwd_fullgraph=fwd_fullgraph,
-                        bwd_resize_count_before_inductor=48 if fwd_fullgraph else None,
-                    ),
-                )
-            if fwd_fullgraph:
-                self.assertEqual(
-                    len(triton_codes),
-                    2,
-                    "Expected two separate lowerings to Triton code, one from FWD graph and one from Compiled Autograd BWD graph",
-                )
-                fwd_code = triton_codes[0]
-
-                extra_str_from_graph_partition = (
-                    "self, " if torch._inductor.config.graph_partition else ""
-                )
-
-                file_check = FileCheck().check(
-                    f"def call({extra_str_from_graph_partition}args):"
-                )
-                for fwd_ag_block_info in [
-                    dict(overlapped_compute_op_str=None),
-                    dict(
-                        overlapped_compute_op_str="extern_kernels.mm(",
-                    ),
-                    dict(
-                        overlapped_compute_op_str="extern_kernels.mm(",
-                    ),
-                    dict(
-                        overlapped_compute_op_str="extern_kernels.mm(",
-                    ),
-                    dict(
-                        overlapped_compute_op_str="extern_kernels.mm(",
-                    ),
-                    dict(
-                        overlapped_compute_op_str="extern_kernels.mm(",
-                    ),
-                    dict(
-                        overlapped_compute_op_str="extern_kernels.mm(",
-                    ),
-                    dict(
-                        overlapped_compute_op_str="extern_kernels.mm(",
-                    ),
-                    dict(
-                        overlapped_compute_op_str="extern_kernels.mm(",
-                        last_all_gather=True,
-                    ),
-                ]:
-                    # file_check = self.inductor_code_check_fsdp_all_gather(
-                    #     file_check, **fwd_ag_block_info
-                    # )
-                    pass
-                file_check.run(fwd_code)
-
-                bwd_code = triton_codes[1]
-                file_check = FileCheck().check(
-                    f"def call({extra_str_from_graph_partition}args):"
-                )
-                for bwd_ag_block_info in [
-                    dict(overlapped_compute_op_str=None),
-                    dict(
-                        overlapped_compute_op_str="extern_kernels.mm(",
-                    ),
-                    dict(
-                        overlapped_compute_op_str="extern_kernels.mm(",
-                        last_all_gather=True,
-                    ),
-                ]:
-                    # file_check = self.inductor_code_check_fsdp_all_gather(
-                    #     file_check, **bwd_ag_block_info
-                    # )
-                    pass
-                for bwd_rs_block_info in [
-                    dict(overlapped_compute_op_str="extern_kernels.addmm("),
-                    dict(
-                        overlapped_compute_op_str=None
-                    ),  # TODO: improve compute/comm overlap, so that `overlapped_compute_op_str` is not None
-                    dict(overlapped_compute_op_str=None),
-                ]:
-                    # file_check = self.inductor_code_check_fsdp_reduce_scatter(
-                    #     file_check, **bwd_rs_block_info
-                    # )
-                    pass
-                file_check.run(bwd_code)
-
-    @unittest.skipIf(not HAS_GPU, "Inductor+gpu needs triton and recent GPU arch")
-    def test_nested_fully_shard_backend_inductor_fullgraph_True(self):
-        self._test_nested_fully_shard_backend_inductor_fullgraph_True()
-
-    @unittest.skipIf(not HAS_GPU, "Inductor+gpu needs triton and recent GPU arch")
-    @torch._inductor.config.patch("graph_partition", True)
-    def test_nested_fully_shard_backend_inductor_fullgraph_True_graph_partition(self):
-        self._test_nested_fully_shard_backend_inductor_fullgraph_True()
-
-    @unittest.skip("TODO: fix fwd_fullgraph=False case")
-    @unittest.skipIf(not HAS_GPU, "Inductor+gpu needs triton and recent GPU arch")
-    def test_nested_fully_shard_backend_inductor_fullgraph_False(self):
-        self.skipTestForOldSm()
-        _, triton_codes = run_and_get_code(
-            lambda: self._test_traceable_fsdp(
-                *self._create_nested_fully_shard_factory_fns(fwd_fullgraph=False),
-                "inductor",
-                fwd_fullgraph=False,
-            ),
-        )
-        # TODO: when fwd_fullgraph=False and there is graph break in FWD graph,
-        # there are several recompiles, need to figure out why.
-        self.assertGreater(
-            len(triton_codes),
-            2,
-            "Expected at least 3 separate lowerings to Triton code, which means at least 1 graph break in FWD graph",
-        )
-
-    def _create_transformer_factory_fns(
-        self, all_requires_grad, *, activation_checkpoint=False
-    ):
-        seq_len = 16
-        vocab_size = 8
-        n_layers = 3
-
-        def model_init_fn():
-            torch.manual_seed(self.rank)
-            fsdp_config = {}
-            mesh = init_device_mesh(device_type.type, (self.world_size,))
-            model_args = ModelArgs(
-                vocab_size=vocab_size,
-                n_layers=n_layers,
-                checkpoint_activations=activation_checkpoint,
-            )
-            model = Transformer(model_args)
-            if not all_requires_grad:
-                requires_grad_params = ["attention.wq", "attention.wv"]
-                requires_grad_param_count = 0
-                for k, v in model.named_parameters():
-                    for substring in requires_grad_params:
-                        if substring in k:
-                            v.requires_grad_(True)
-                            requires_grad_param_count += 1
-                        else:
-                            v.requires_grad_(False)
-                assert requires_grad_param_count == n_layers * len(requires_grad_params)
-            for _, mod in enumerate(model.layers):
-                fully_shard(mod, mesh=mesh, reshard_after_forward=True, **fsdp_config)
-            model = fully_shard(
-                model, mesh=mesh, reshard_after_forward=False, **fsdp_config
-            )
-            optim = torch.optim.SGD(model.parameters(), lr=1e-4)
-            return model, optim
-
-        def input_creation_fn():
-            torch.manual_seed(self.rank)
-            inp = torch.randint(
-                0, vocab_size, (2, seq_len), device=device_type, requires_grad=False
-            )
-            return inp
-
-        return model_init_fn, input_creation_fn
-
-    def _maybe_add_graph_break_to_sdpa(self, fwd_fullgraph):
-        def _sdpa_with_graph_break(*args, **kwargs):
-            torch._dynamo.graph_break()
-            return orig_F_scaled_dot_product_attention(*args, **kwargs)
-
-        if not fwd_fullgraph:
-            return mock.patch.object(
-                F,
-                "scaled_dot_product_attention",
-                _sdpa_with_graph_break,
-            )
-        else:
-            return contextlib.nullcontext()
-
-    @unittest.skipIf(not HAS_GPU, "Inductor+gpu needs triton and recent GPU arch")
-    def test_transformer_backend_aot_eager(self):
-        # TODO: fix fwd_fullgraph=False case
-        for fwd_fullgraph, all_requires_grad in itertools.product(
-            [True], [True, False]
-        ):
-            with (
-                self._maybe_add_graph_break_to_sdpa(fwd_fullgraph),
-                self._reinplace_all_gather_with_optional_checks(fwd_fullgraph),
-            ):
-                self._test_traceable_fsdp(
-                    *self._create_transformer_factory_fns(
-                        all_requires_grad=all_requires_grad
-                    ),
-                    "aot_eager",
-                    fwd_fullgraph=fwd_fullgraph,
-                )
-
-    @unittest.skipIf(not HAS_GPU, "Inductor+gpu needs triton and recent GPU arch")
-    # TODO: native_dropout has worse accuracy after decomp, need to figure out why
-    @torch._inductor.config.patch(fallback_random=True)
-    def test_transformer_backend_aot_eager_decomp_partition(self):
-        # TODO: fix fwd_fullgraph=False case
-        for fwd_fullgraph, all_requires_grad in itertools.product(
-            [True], [True, False]
-        ):
-            with self._maybe_add_graph_break_to_sdpa(fwd_fullgraph):
-                self._test_traceable_fsdp(
-                    *self._create_transformer_factory_fns(
-                        all_requires_grad=all_requires_grad
-                    ),
-                    "aot_eager_decomp_partition",
-                    fwd_fullgraph=fwd_fullgraph,
-                )
-
-    def _test_transformer_backend_inductor_fullgraph_True(self):
-        self.skipTestForOldSm()
-        for (
-            fwd_fullgraph,
-            all_requires_grad,
-            activation_checkpoint,
-        ) in itertools.product([True], [True, False], [True, False]):
-            log.warning(
-                f"fwd_fullgraph={fwd_fullgraph}, all_requires_grad={all_requires_grad}, activation_checkpoint={activation_checkpoint}"  # noqa: G004, G001, B950
-            )
-            with (
-                self._reinplace_all_gather_with_optional_checks(fwd_fullgraph),
-                torch._inductor.config.patch(
-                    post_grad_custom_post_pass=(
-                        functools.partial(
-                            self._check_fsdp_copy_and_resize_ops_count_in_graph,
-                            # NOTE: For the root unsharded params, we don't reshard after forward since for training,
-                            # the parameters would be freed and all-gathered immediately. Hence we still have
-                            # their resize and copy ops in the graph.
-                            fwd_copy_count=4,
-                            fwd_resize_count=4,
-                            bwd_copy_count=0,
-                            bwd_resize_count=4,
-                        )
-                        if fwd_fullgraph
-                        else None
-                    )
-                ),
-            ):
-                _, triton_codes = run_and_get_code(
-                    lambda: self._test_traceable_fsdp(
-                        *self._create_transformer_factory_fns(
-                            all_requires_grad=all_requires_grad,
-                            activation_checkpoint=activation_checkpoint,
-                        ),
-                        "inductor",
-                        fwd_fullgraph=fwd_fullgraph,
-                        bwd_resize_count_before_inductor=76 if fwd_fullgraph else None,
-                    ),
-                )
-            if fwd_fullgraph:
-                self.assertEqual(
-                    len(triton_codes),
-                    2,
-                    "Expected two separate lowerings to Triton code, one from FWD graph and one from Compiled Autograd BWD graph",
-                )
-                fwd_code = triton_codes[0]
-                extra_str_from_graph_partition = (
-                    "self, " if torch._inductor.config.graph_partition else ""
-                )
-
-                file_check = FileCheck().check(
-                    f"def call({extra_str_from_graph_partition}args):"
-                )
-                for fwd_ag_block_info in [
-                    dict(
-                        overlapped_compute_op_str=(
-                            "triton_" if all_requires_grad else None
-                        ),
-                    ),
-                    dict(
-                        overlapped_compute_op_str="aten.native_dropout.",
-                    ),
-                    dict(
-                        overlapped_compute_op_str="aten._scaled_dot_product_efficient_attention.",
-                    ),
-                    dict(
-                        overlapped_compute_op_str="aten._scaled_dot_product_efficient_attention.",
-                        last_all_gather=True,
-                    ),
-                ]:
-                    # file_check = self.inductor_code_check_fsdp_all_gather(
-                    #     file_check, **fwd_ag_block_info
-                    # )
-                    pass
-                file_check.run(fwd_code)
-
-                bwd_code = triton_codes[1]
-                file_check = FileCheck().check(
-                    f"def call({extra_str_from_graph_partition}args):"
-                )
-                for bwd_ag_block_info in [
-                    dict(
-                        overlapped_compute_op_str="extern_kernels.mm(",
-                    ),
-                    dict(
-                        overlapped_compute_op_str="aten._scaled_dot_product_efficient_attention_backward.",
-                    ),
-                    dict(
-                        overlapped_compute_op_str="aten._scaled_dot_product_efficient_attention_backward.",
-                        last_all_gather=True,
-                    ),
-                ]:
-                    # if bwd_ag_block_info is not None:
-                    #     file_check = self.inductor_code_check_fsdp_all_gather(
-                    #         file_check, **bwd_ag_block_info
-                    #     )
-                    pass
-                for bwd_rs_block_info in [
-                    (
-                        dict(overlapped_compute_op_str="extern_kernels.mm(")
-                        if all_requires_grad
-                        else None
-                    ),
-                    dict(
-                        overlapped_compute_op_str=None
-                    ),  # TODO: improve compute/comm overlap, so that `overlapped_compute_op_str` is not None
-                    dict(overlapped_compute_op_str=None),
-                    dict(overlapped_compute_op_str=None) if all_requires_grad else None,
-                ]:
-                    # if bwd_rs_block_info is not None:
-                    #     file_check = self.inductor_code_check_fsdp_reduce_scatter(
-                    #         file_check, **bwd_rs_block_info
-                    #     )
-                    pass
-                file_check.run(bwd_code)
-
-    @unittest.skip('"Traceable FSDP2" is not being maintained anymore.')
-    @unittest.skipIf(not HAS_GPU, "Inductor+gpu needs triton and recent GPU arch")
-    # TODO: native_dropout causes CUDA IMA error, need to figure out why
-    @torch._inductor.config.patch(fallback_random=True)
-    def test_transformer_backend_inductor_fullgraph_True(self):
-        self._test_transformer_backend_inductor_fullgraph_True()
-
-    @unittest.skip('"Traceable FSDP2" is not being maintained anymore.')
-    @unittest.skipIf(not HAS_GPU, "Inductor+gpu needs triton and recent GPU arch")
-    # TODO: native_dropout causes CUDA IMA error, need to figure out why
-    @torch._inductor.config.patch(fallback_random=True)
-    @torch._inductor.config.patch("graph_partition", True)
-    def test_transformer_backend_inductor_fullgraph_True_graph_partition(self):
-        self._test_transformer_backend_inductor_fullgraph_True()
-
-    @unittest.skip("TODO: fix fwd_fullgraph=False case")
-    @unittest.skipIf(not HAS_GPU, "Inductor+gpu needs triton and recent GPU arch")
-    # TODO: native_dropout causes CUDA IMA error, need to figure out why
-    @torch._inductor.config.patch(fallback_random=True)
-    def test_transformer_backend_inductor_fullgraph_False(self):
-        self.skipTestForOldSm()
-        fwd_fullgraph = False
-        # TODO: fix numerical issue in activation_checkpoint=True case
-        for all_requires_grad, activation_checkpoint in itertools.product(
-            [True, False], [False]
-        ):
-            log.warning(
-                f"fwd_fullgraph={fwd_fullgraph}, all_requires_grad={all_requires_grad}, activation_checkpoint={activation_checkpoint}"  # noqa: G004, G001, B950
-            )
-            with self._maybe_add_graph_break_to_sdpa(fwd_fullgraph):
-                _, triton_codes = run_and_get_code(
-                    lambda: self._test_traceable_fsdp(
-                        *self._create_transformer_factory_fns(
-                            all_requires_grad=all_requires_grad,
-                            activation_checkpoint=activation_checkpoint,
-                        ),
-                        "inductor",
-                        fwd_fullgraph=fwd_fullgraph,
-                    ),
-                )
-            # TODO: when fwd_fullgraph=False and there is graph break in FWD graph,
-            # there are several recompiles, need to figure out why.
-            self.assertGreater(
-                len(triton_codes),
-                2,
-                "Expected at least 3 separate lowerings to Triton code, which means at least 1 graph break in FWD graph",
-            )
-
     def test_dynamo_recompiles_on_fsdp_layers(self):
         m = Mod()
         for name, child in m.encoder.named_children():
diff --git a/test/distributed/_composable/fsdp/test_fully_shard_dtensor.py b/test/distributed/_composable/fsdp/test_fully_shard_dtensor.py
new file mode 100644
index 0000000000000..5045f59d41711
--- /dev/null
+++ b/test/distributed/_composable/fsdp/test_fully_shard_dtensor.py
@@ -0,0 +1,446 @@
+# Owner(s): ["oncall: distributed"]
+
+import copy
+
+import torch
+import torch.distributed as dist
+import torch.nn as nn
+from torch.distributed._composable import replicate
+from torch.distributed._composable.replicate_with_fsdp import (
+    replicate as replicate_with_fsdp,
+)
+from torch.distributed.fsdp import DataParallelMeshDims, fully_shard
+from torch.distributed.tensor import (
+    distribute_module,
+    distribute_tensor,
+    DTensor,
+    init_device_mesh,
+    Replicate,
+    Shard,
+)
+from torch.testing._internal.common_distributed import skip_if_lt_x_gpu
+from torch.testing._internal.common_fsdp import FSDPTest, get_devtype, MLP
+from torch.testing._internal.common_utils import run_tests
+
+
+device_type = torch.device(get_devtype())
+
+
+def _tp_partition_fn(name, module, device_mesh):
+    """Partition Linear weights across the last mesh dim for TP."""
+    if not isinstance(module, nn.Linear):
+        return
+    num_non_tp_dims = device_mesh.ndim - 1
+    replicate_prefix = [Replicate()] * num_non_tp_dims
+    for param_name, param in list(module.named_parameters(recurse=False)):
+        if param_name == "weight":
+            if "in_proj" in name:
+                placements = replicate_prefix + [Shard(0)]
+            else:
+                placements = replicate_prefix + [Shard(1)]
+        else:
+            placements = replicate_prefix + [Replicate()]
+        dist_param = nn.Parameter(
+            distribute_tensor(param, device_mesh, placements),
+            requires_grad=param.requires_grad,
+        )
+        module.register_parameter(param_name, dist_param)
+
+
+def _tp_shard_fn(param):
+    """FSDP shard placement that avoids the existing TP shard dim."""
+    if any(isinstance(p, Shard) and p.dim == 0 for p in param.placements):
+        return Shard(1)
+    return Shard(0)
+
+
+class TestFullyShardDTensor(FSDPTest):
+    @property
+    def world_size(self):
+        return min(4, torch.cuda.device_count())
+
+    def _run_train_parity(
+        self, model, ref_model, dp_pg, mesh=None, num_iters=5, mlp_dim=16
+    ):
+        optim = torch.optim.Adam(model.parameters(), lr=1e-2, foreach=False)
+        ref_optim = torch.optim.Adam(ref_model.parameters(), lr=1e-2, foreach=False)
+        torch.manual_seed(42 + dp_pg.rank() + 1)
+        for i in range(num_iters):
+            inp = torch.randn((2, mlp_dim), device=device_type)
+            ref_optim.zero_grad(set_to_none=(i % 2 == 0))
+            ref_loss = ref_model(inp).sum()
+            ref_loss.backward()
+            ref_optim.step()
+
+            optim.zero_grad(set_to_none=(i % 2 == 0))
+            if mesh is not None:
+                # Use Replicate on all dims: each rank computes on the full
+                # input (like DDP). DTensor handles the TP-sharded params.
+                inp = DTensor.from_local(
+                    inp, mesh, [Replicate()] * mesh.ndim, run_check=False
+                )
+            loss = model(inp).sum()
+            loss.backward()
+            optim.step()
+
+            loss_cmp = loss.full_tensor() if isinstance(loss, DTensor) else loss
+            self.assertEqual(ref_loss, loss_cmp)
+
+        for (n1, p1), (n2, p2) in zip(
+            ref_model.named_parameters(), model.named_parameters(), strict=True
+        ):
+            p2_full = p2.full_tensor() if isinstance(p2, DTensor) else p2
+            self.assertEqual(p1, p2_full, msg=f"Param mismatch: {n1} vs {n2}")
+
+    @skip_if_lt_x_gpu(2)
+    def test_dtensor_train_parity(self):
+        """Train parity for FSDP/HSDP/DDP with DTensors on SPMD meshes."""
+        ws = self.world_size
+        world_mesh = init_device_mesh(
+            device_type.type, (ws,), mesh_dim_names=("world",)
+        )
+        # (sizes, names, dp_dims, use_tp, reshard, dp_pg_source, use_rep_fsdp)
+        cases = [
+            # 1D: FSDP
+            (
+                (ws,),
+                ("fsdp",),
+                DataParallelMeshDims(shard="fsdp"),
+                False,
+                True,
+                None,
+                False,
+            ),
+            # 1D: FSDP with reshard_after_forward=False
+            (
+                (ws,),
+                ("fsdp0",),
+                DataParallelMeshDims(shard="fsdp0"),
+                False,
+                False,
+                None,
+                False,
+            ),
+            # 1D: DDP-only
+            (
+                (ws,),
+                ("ddp",),
+                DataParallelMeshDims(replicate="ddp"),
+                False,
+                True,
+                None,
+                False,
+            ),
+            # 1D: replicate_with_fsdp
+            (
+                (ws,),
+                ("ddp0",),
+                DataParallelMeshDims(replicate="ddp0"),
+                False,
+                True,
+                None,
+                True,
+            ),
+        ]
+        if ws >= 4:
+            cases.extend(
+                [
+                    # HSDP 2D
+                    (
+                        (2, ws // 2),
+                        ("rep", "shard"),
+                        DataParallelMeshDims(shard="shard", replicate="rep"),
+                        False,
+                        True,
+                        "world",
+                        False,
+                    ),
+                    # Multi-shard FSDP
+                    (
+                        (2, ws // 2),
+                        ("dp0", "dp1"),
+                        DataParallelMeshDims(shard=("dp0", "dp1")),
+                        False,
+                        True,
+                        "world",
+                        False,
+                    ),
+                    # FSDP+TP
+                    (
+                        (2, ws // 2),
+                        ("fsdp1", "tp"),
+                        DataParallelMeshDims(shard="fsdp1"),
+                        True,
+                        True,
+                        "fsdp1",
+                        False,
+                    ),
+                    # FSDP+TP with reshard_after_forward=False
+                    (
+                        (2, ws // 2),
+                        ("fsdp2", "tp0"),
+                        DataParallelMeshDims(shard="fsdp2"),
+                        True,
+                        False,
+                        "fsdp2",
+                        False,
+                    ),
+                    # HSDP+TP 3D
+                    (
+                        (1, ws // 2, 2),
+                        ("rep0", "fsdp3", "tp1"),
+                        DataParallelMeshDims(shard="fsdp3", replicate="rep0"),
+                        True,
+                        True,
+                        "fsdp3",
+                        False,
+                    ),
+                    # Multi-dim replicate
+                    (
+                        (1, ws // 2, 2),
+                        ("ddp1", "ddp2", "fsdp4"),
+                        DataParallelMeshDims(shard="fsdp4", replicate=("ddp1", "ddp2")),
+                        False,
+                        True,
+                        "world",
+                        False,
+                    ),
+                ]
+            )
+        mlp_dim = 16
+        for sizes, names, dp_dims, use_tp, reshard, dp_pg_src, use_rep in cases:
+            with self.subTest(
+                names=names, use_tp=use_tp, reshard=reshard, use_rep=use_rep
+            ):
+                mesh = world_mesh._unflatten(0, sizes, names)
+
+                torch.manual_seed(42)
+                model = MLP(mlp_dim, device=device_type)
+                ref_model = copy.deepcopy(model)
+
+                partition_fn = _tp_partition_fn if use_tp else None
+                distribute_module(model, mesh, partition_fn)
+
+                if use_rep:
+                    replicate_with_fsdp(model, mesh=mesh, dp_mesh_dims=dp_dims)
+                else:
+                    shard_fn = _tp_shard_fn if use_tp else None
+                    fully_shard(
+                        model,
+                        mesh=mesh,
+                        reshard_after_forward=reshard,
+                        shard_placement_fn=shard_fn,
+                        dp_mesh_dims=dp_dims,
+                    )
+
+                if dp_pg_src is None:
+                    dp_pg = mesh.get_group()
+                elif dp_pg_src == "world":
+                    dp_pg = dist.group.WORLD
+                else:
+                    dp_pg = mesh[dp_pg_src].get_group()
+
+                replicate(
+                    ref_model,
+                    device_ids=[self.rank] if device_type.type != "cpu" else None,
+                    process_group=dp_pg,
+                )
+
+                self._run_train_parity(
+                    model, ref_model, dp_pg, mesh=mesh, mlp_dim=mlp_dim
+                )
+            dist.barrier()
+
+    @skip_if_lt_x_gpu(2)
+    def test_sharded_param_correctness_1d(self):
+        """Verify sharded param mesh and placements for FSDP on 1D mesh."""
+        mesh = init_device_mesh(
+            device_type.type, (self.world_size,), mesh_dim_names=("fsdp",)
+        )
+
+        model = MLP(16, device=device_type)
+        distribute_module(model, mesh)
+        fully_shard(
+            model,
+            mesh=mesh,
+            dp_mesh_dims=DataParallelMeshDims(shard="fsdp"),
+        )
+
+        for param in model.parameters():
+            self.assertIsInstance(param, DTensor)
+            self.assertEqual(param.device_mesh, mesh)
+            self.assertEqual(len(param.placements), 1)
+            self.assertIsInstance(param.placements[0], Shard)
+
+    @skip_if_lt_x_gpu(4)
+    def test_fsdp_tp_dtensor_sharded_params(self):
+        """Verify sharded param mesh and placements for FSDP+TP on 2D mesh."""
+        dp_size = 2
+        tp_size = self.world_size // dp_size
+        mesh = init_device_mesh(
+            device_type.type,
+            (dp_size, tp_size),
+            mesh_dim_names=("fsdp", "tp"),
+        )
+
+        mlp_dim = 16
+        model = MLP(mlp_dim, device=device_type)
+
+        def partition_fn(name, module, device_mesh):
+            if not isinstance(module, nn.Linear):
+                return
+            for param_name, param in list(module.named_parameters(recurse=False)):
+                if param_name == "weight":
+                    if "in_proj" in name:
+                        placements = [Replicate(), Shard(0)]
+                    else:
+                        placements = [Replicate(), Shard(1)]
+                else:
+                    placements = [Replicate(), Replicate()]
+                dist_param = nn.Parameter(
+                    distribute_tensor(param, device_mesh, placements),
+                    requires_grad=param.requires_grad,
+                )
+                module.register_parameter(param_name, dist_param)
+
+        distribute_module(model, mesh, partition_fn)
+
+        def shard_fn(param):
+            if any(isinstance(p, Shard) and p.dim == 0 for p in param.placements):
+                return Shard(1)
+            return Shard(0)
+
+        fully_shard(
+            model,
+            mesh=mesh,
+            shard_placement_fn=shard_fn,
+            dp_mesh_dims=DataParallelMeshDims(shard="fsdp"),
+        )
+
+        for name, param in model.named_parameters():
+            self.assertIsInstance(param, DTensor)
+            self.assertEqual(param.device_mesh, mesh)
+            self.assertEqual(len(param.placements), 2)
+            if "in_proj.weight" in name:
+                # FSDP shards dim 1 (avoiding TP dim 0), TP shards dim 0
+                self.assertIsInstance(param.placements[0], Shard)
+                self.assertEqual(param.placements[0].dim, 1)
+                self.assertIsInstance(param.placements[1], Shard)
+                self.assertEqual(param.placements[1].dim, 0)
+            elif "out_proj.weight" in name:
+                # FSDP shards dim 0 (default), TP shards dim 1
+                self.assertIsInstance(param.placements[0], Shard)
+                self.assertEqual(param.placements[0].dim, 0)
+                self.assertIsInstance(param.placements[1], Shard)
+                self.assertEqual(param.placements[1].dim, 1)
+            elif "bias" in name:
+                # FSDP shards dim 0 (default), TP replicates
+                self.assertIsInstance(param.placements[0], Shard)
+                self.assertEqual(param.placements[0].dim, 0)
+                self.assertIsInstance(param.placements[1], Replicate)
+
+    @skip_if_lt_x_gpu(2)
+    def test_validation_non_replicate_dp_placement(self):
+        """Error when a param has non-Replicate placement on the DP shard dim."""
+        mesh = init_device_mesh(
+            device_type.type, (self.world_size,), mesh_dim_names=("fsdp",)
+        )
+        model = nn.Linear(16, 32, device=device_type)
+        # Distribute weight with Shard(0) on the FSDP dim
+        model.weight = nn.Parameter(
+            distribute_tensor(model.weight.data, mesh, [Shard(0)]),
+            requires_grad=True,
+        )
+        model.bias = nn.Parameter(
+            distribute_tensor(model.bias.data, mesh, [Replicate()]),
+            requires_grad=True,
+        )
+        with self.assertRaisesRegex(ValueError, "Expected Replicate"):
+            fully_shard(
+                model,
+                mesh=mesh,
+                dp_mesh_dims=DataParallelMeshDims(shard="fsdp"),
+            )
+
+    @skip_if_lt_x_gpu(2)
+    def test_validation_invalid_dim_names(self):
+        """Error when dp_mesh_dims references nonexistent mesh dim names."""
+        mesh = init_device_mesh(
+            device_type.type, (self.world_size,), mesh_dim_names=("fsdp",)
+        )
+        model = MLP(16, device=device_type)
+        distribute_module(model, mesh)
+        with self.assertRaisesRegex(ValueError, "not found in mesh.mesh_dim_names"):
+            fully_shard(
+                model,
+                mesh=mesh,
+                dp_mesh_dims=DataParallelMeshDims(shard="nonexistent"),
+            )
+
+    @skip_if_lt_x_gpu(2)
+    def test_validation_mesh_mismatch(self):
+        """Error when param DTensor mesh differs from the mesh passed to fully_shard."""
+        mesh1 = init_device_mesh(
+            device_type.type, (self.world_size,), mesh_dim_names=("fsdp",)
+        )
+        mesh2 = init_device_mesh(
+            device_type.type, (self.world_size,), mesh_dim_names=("fsdp",)
+        )
+        model = nn.Linear(16, 32, device=device_type)
+        # Distribute params on mesh1 but pass mesh2 to fully_shard
+        model.weight = nn.Parameter(
+            distribute_tensor(model.weight.data, mesh1, [Replicate()]),
+            requires_grad=True,
+        )
+        model.bias = nn.Parameter(
+            distribute_tensor(model.bias.data, mesh1, [Replicate()]),
+            requires_grad=True,
+        )
+        with self.assertRaisesRegex(ValueError, "same mesh"):
+            fully_shard(
+                model,
+                mesh=mesh2,
+                dp_mesh_dims=DataParallelMeshDims(shard="fsdp"),
+            )
+
+    def test_validation_at_least_one_required(self):
+        """Error when neither shard nor replicate is set."""
+        with self.assertRaisesRegex(ValueError, "At least one of shard or replicate"):
+            DataParallelMeshDims()
+
+    @skip_if_lt_x_gpu(2)
+    def test_validation_spmd_mesh_non_dtensor_params(self):
+        """Error when dp_mesh_dims is provided but params are not DTensors."""
+        mesh = init_device_mesh(
+            device_type.type, (self.world_size,), mesh_dim_names=("fsdp",)
+        )
+        model = MLP(16, device=device_type)
+        # Do NOT call distribute_module -- params are plain tensors
+        with self.assertRaisesRegex(ValueError, "must be DTensors"):
+            fully_shard(
+                model,
+                mesh=mesh,
+                dp_mesh_dims=DataParallelMeshDims(shard="fsdp"),
+            )
+
+    @skip_if_lt_x_gpu(2)
+    def test_validation_reshard_after_forward_int_spmd(self):
+        """Error when reshard_after_forward is int with SPMD mesh."""
+        mesh = init_device_mesh(
+            device_type.type, (self.world_size,), mesh_dim_names=("fsdp",)
+        )
+        model = MLP(16, device=device_type)
+        distribute_module(model, mesh)
+        with self.assertRaisesRegex(
+            NotImplementedError, "reshard_after_forward as int"
+        ):
+            fully_shard(
+                model,
+                mesh=mesh,
+                reshard_after_forward=2,
+                dp_mesh_dims=DataParallelMeshDims(shard="fsdp"),
+            )
+
+
+if __name__ == "__main__":
+    run_tests()
diff --git a/test/distributed/_composable/fsdp/test_fully_shard_extensions.py b/test/distributed/_composable/fsdp/test_fully_shard_extensions.py
index 0b25e09b3defc..83f11d390a29e 100644
--- a/test/distributed/_composable/fsdp/test_fully_shard_extensions.py
+++ b/test/distributed/_composable/fsdp/test_fully_shard_extensions.py
@@ -5,7 +5,7 @@
 import functools
 import math
 import threading
-from typing import Any, Optional, Union
+from typing import Any
 
 import torch
 import torch.distributed as dist
@@ -56,18 +56,24 @@ def two_tensor_fsdp_post_all_gather(
     metadata: Any,
     param_dtype: torch.dtype,
     *,
-    out: Optional[torch.Tensor] = None,
-) -> Union[tuple[torch.Tensor, tuple[torch.Tensor, ...]], None]:
-    assert metadata is None, f"{metadata}"
+    out: torch.Tensor | None = None,
+) -> tuple[torch.Tensor, tuple[torch.Tensor, ...]] | None:
+    if metadata is not None:
+        raise AssertionError(f"Expected metadata to be None, got {metadata}")
     a, b = all_gather_outputs
     if out is not None:
-        assert isinstance(out, TwoTensor), f"{type(out)}"
+        if not isinstance(out, TwoTensor):
+            raise AssertionError(f"Expected TwoTensor, got {type(out)}")
         if a.dtype == param_dtype:
-            assert a.untyped_storage().data_ptr() == out.a.untyped_storage().data_ptr()
-            assert b.untyped_storage().data_ptr() == out.b.untyped_storage().data_ptr()
+            if a.untyped_storage().data_ptr() != out.a.untyped_storage().data_ptr():
+                raise AssertionError("a storage data_ptr mismatch with out.a")
+            if b.untyped_storage().data_ptr() != out.b.untyped_storage().data_ptr():
+                raise AssertionError("b storage data_ptr mismatch with out.b")
         else:
-            assert out.a.dtype == param_dtype, f"{out.a.dtype} {param_dtype}"
-            assert out.b.dtype == param_dtype, f"{out.b.dtype} {param_dtype}"
+            if out.a.dtype != param_dtype:
+                raise AssertionError(f"out.a dtype {out.a.dtype} != {param_dtype}")
+            if out.b.dtype != param_dtype:
+                raise AssertionError(f"out.b dtype {out.b.dtype} != {param_dtype}")
             out.a.copy_(a)
             out.b.copy_(b)
         return
@@ -102,7 +108,8 @@ def fsdp_pre_all_gather(
         module: nn.Module,
         mp_policy: MixedPrecisionPolicy,
     ) -> tuple[tuple[torch.Tensor, ...], Any]:
-        assert mesh.ndim == 1, f"{mesh.ndim}"
+        if mesh.ndim != 1:
+            raise AssertionError(f"Expected mesh.ndim == 1, got {mesh.ndim}")
         mesh_size = mesh.size()
         requires_padding = outer_size[0] % mesh_size != 0
         if requires_padding and self._pad_in_pre_all_gather:
@@ -122,11 +129,15 @@ def fsdp_post_all_gather(
         metadata: Any,
         param_dtype: torch.dtype,
         *,
-        out: Optional[torch.Tensor] = None,
-    ) -> Union[tuple[torch.Tensor, tuple[torch.Tensor, ...]], None]:
-        assert metadata is None, f"{metadata}"
+        out: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, tuple[torch.Tensor, ...]] | None:
+        if metadata is not None:
+            raise AssertionError(f"Expected metadata to be None, got {metadata}")
         (tensor,) = all_gather_outputs
-        assert tensor.dtype == torch.bfloat16, f"{tensor.dtype}"
+        if tensor.dtype != torch.bfloat16:
+            raise AssertionError(
+                f"Expected tensor.dtype == torch.bfloat16, got {tensor.dtype}"
+            )
         if out is not None:
             with _unsafe_preserve_version_counter(out):
                 out.copy_(tensor)
@@ -143,7 +154,10 @@ def unwrap(x: cls):
             if pad_in_pre_all_gather is None:
                 pad_in_pre_all_gather = x._pad_in_pre_all_gather
             else:
-                assert pad_in_pre_all_gather == x._pad_in_pre_all_gather
+                if pad_in_pre_all_gather != x._pad_in_pre_all_gather:
+                    raise AssertionError(
+                        f"pad_in_pre_all_gather mismatch: {pad_in_pre_all_gather} vs {x._pad_in_pre_all_gather}"
+                    )
             return x._data
 
         out = func(
@@ -331,11 +345,15 @@ def fsdp_post_all_gather(
             metadata: Any,
             param_dtype: torch.dtype,
             *,
-            out: Optional[torch.Tensor] = None,
-        ) -> Union[tuple[torch.Tensor, tuple[torch.Tensor, ...]], None]:
+            out: torch.Tensor | None = None,
+        ) -> tuple[torch.Tensor, tuple[torch.Tensor, ...]] | None:
             (tensor,) = all_gather_outputs
-            assert metadata is None, f"{metadata}"
-            assert tensor.dtype == torch.bfloat16, f"{tensor.dtype}"
+            if metadata is not None:
+                raise AssertionError(f"Expected metadata to be None, got {metadata}")
+            if tensor.dtype != torch.bfloat16:
+                raise AssertionError(
+                    f"Expected tensor.dtype == torch.bfloat16, got {tensor.dtype}"
+                )
             if out is not None:
                 with _unsafe_preserve_version_counter(out):
                     out.copy_(tensor)
@@ -375,7 +393,8 @@ def fsdp_post_all_gather(
             model(inp).sum().backward()
             optim.step()
             optim.zero_grad()
-        assert tls.ran_pre_all_gather
+        if not tls.ran_pre_all_gather:
+            raise AssertionError("Expected tls.ran_pre_all_gather to be True")
 
     @skip_if_lt_x_gpu(1)
     def test_all_gather_extension_outer_size_stride(self):
@@ -385,9 +404,10 @@ def test_all_gather_extension_outer_size_stride(self):
         only some ranks may require padding, in which case only those ranks
         will error out and the all-gather will timeout.
         """
-        assert self.world_size >= 2, (
-            f"Assumes world size of at least 2 but got {self.world_size=}"
-        )
+        if self.world_size < 2:
+            raise AssertionError(
+                f"Assumes world size of at least 2 but got {self.world_size=}"
+            )
         model = MLP(dim=3, dim_multiplier=3)
         for module in model.modules():
             for param_name, param in module.named_parameters(recurse=False):
@@ -435,8 +455,8 @@ def fsdp_post_all_gather(
             metadata: Any,
             param_dtype: torch.dtype,
             *,
-            out: Optional[torch.Tensor] = None,
-        ) -> Union[tuple[torch.Tensor, tuple[torch.Tensor, ...]], None]:
+            out: torch.Tensor | None = None,
+        ) -> tuple[torch.Tensor, tuple[torch.Tensor, ...]] | None:
             (tensor,) = all_gather_outputs
             if out is not None:
                 return
diff --git a/test/distributed/_composable/fsdp/test_fully_shard_frozen.py b/test/distributed/_composable/fsdp/test_fully_shard_frozen.py
index f56c5e76c1224..9281c7da0eea2 100644
--- a/test/distributed/_composable/fsdp/test_fully_shard_frozen.py
+++ b/test/distributed/_composable/fsdp/test_fully_shard_frozen.py
@@ -3,7 +3,6 @@
 import copy
 import functools
 import itertools
-from typing import Union
 
 import torch
 import torch.distributed as dist
@@ -55,7 +54,7 @@ def test_train_mixed_requires_grad_per_group(self):
 
     def _test_train_mixed_requires_grad_per_group(
         self,
-        reshard_after_forward: Union[bool, int],
+        reshard_after_forward: bool | int,
         use_activation_checkpointing: bool,
         freeze_after_init: bool,
     ):
@@ -89,10 +88,11 @@ def _test_train_mixed_requires_grad_per_group(
                 if "bias" not in param_name:
                     param.requires_grad_(False)
         for mlp in model:
-            assert isinstance(mlp, MLP), (
-                "The reduce-scatter numel check assumes the model consists of "
-                f"only the same MLP class but got {type(mlp)}"
-            )
+            if not isinstance(mlp, MLP):
+                raise AssertionError(
+                    "The reduce-scatter numel check assumes the model consists of "
+                    f"only the same MLP class but got {type(mlp)}"
+                )
         expected_numel = sum(
             p._local_tensor.numel()
             for n, p in model[0].named_parameters()
@@ -151,7 +151,7 @@ def test_train_mixed_requires_grad_across_groups(self):
 
     def _test_train_mixed_requires_grad_across_groups(
         self,
-        reshard_after_forward: Union[bool, int],
+        reshard_after_forward: bool | int,
         unfreeze_params: bool,
     ):
         torch.manual_seed(42)
@@ -225,7 +225,7 @@ def test_multi_forward_mixed_requires_grad(self):
 
     def _test_multi_forward_mixed_requires_grad(
         self,
-        reshard_after_forward: Union[bool, int],
+        reshard_after_forward: bool | int,
     ):
         class MultiForwardModule(nn.Module):
             def __init__(self, device: torch.device):
diff --git a/test/distributed/_composable/fsdp/test_fully_shard_ignore_params.py b/test/distributed/_composable/fsdp/test_fully_shard_ignore_params.py
index 504aaa63ec530..c7463f36ca43f 100644
--- a/test/distributed/_composable/fsdp/test_fully_shard_ignore_params.py
+++ b/test/distributed/_composable/fsdp/test_fully_shard_ignore_params.py
@@ -140,7 +140,8 @@ def _modify_ddp_ignored_params(
 ):
     modified_list = []
     for name in ddp_ignored_param_names:
-        assert name in name_to_param_map
+        if name not in name_to_param_map:
+            raise AssertionError(f"Expected {name} in name_to_param_map")
         param = name_to_param_map[name]
         if param not in fsdp_ignored_params:
             # DDP can ignore only if it is not ignored by FSDP
diff --git a/test/distributed/_composable/fsdp/test_fully_shard_init.py b/test/distributed/_composable/fsdp/test_fully_shard_init.py
index 5e862802da085..63ff413214170 100644
--- a/test/distributed/_composable/fsdp/test_fully_shard_init.py
+++ b/test/distributed/_composable/fsdp/test_fully_shard_init.py
@@ -2,7 +2,7 @@
 
 import copy
 import itertools
-from typing import cast, Optional
+from typing import cast
 
 import torch
 import torch.distributed as dist
@@ -37,7 +37,14 @@
 )
 from torch.distributed.tensor.placement_types import _StridedShard
 from torch.testing._internal.common_distributed import skip_if_lt_x_gpu
-from torch.testing._internal.common_fsdp import FSDPTestMultiThread, get_devtype, MLP
+from torch.testing._internal.common_fsdp import (
+    FSDPTest,
+    FSDPTestMultiThread,
+    get_devtype,
+    MLP,
+    patch_all_gather,
+    patch_reduce_scatter,
+)
 from torch.testing._internal.common_utils import run_tests
 from torch.testing._internal.distributed._tensor.common_dtensor import (
     ModelArgs,
@@ -93,7 +100,8 @@ def world_size(self) -> int:
 
     @skip_if_lt_x_gpu(1)
     def test_move_states_to_device_dtensor_valid(self):
-        assert self.world_size >= 4, f"{self.world_size}"
+        if not (self.world_size >= 4):
+            raise AssertionError(f"Expected world_size >= 4, but got {self.world_size}")
         dp_size = 2
         global_mesh = init_device_mesh(
             device_type.type,
@@ -125,7 +133,8 @@ def test_move_states_to_device_dtensor_valid(self):
 
     @skip_if_lt_x_gpu(1)
     def test_move_states_to_device_dtensor_invalid(self):
-        assert self.world_size >= 4, f"{self.world_size}"
+        if not (self.world_size >= 4):
+            raise AssertionError(f"Expected world_size >= 4, but got {self.world_size}")
         dp_size = 2
         global_accelerator_mesh = init_device_mesh(
             device_type.type,
@@ -155,6 +164,48 @@ def test_move_states_to_device_dtensor_invalid(self):
             fully_shard(model, mesh=dp_mesh)
 
 
+class TestFullyShardContainerSubclasses(FSDPTestMultiThread):
+    """Tests that fully_shard accepts ModuleDict/ModuleList subclasses that implement forward()."""
+
+    class DictWithForward(nn.ModuleDict):
+        def __init__(self, in_features=8, out_features=8):
+            super().__init__({"lin": nn.Linear(in_features, out_features)})
+
+        def forward(self, x):
+            return self["lin"](x)
+
+    class ListWithForward(nn.ModuleList):
+        def __init__(self, in_features=8, out_features=8):
+            super().__init__([nn.Linear(in_features, out_features)])
+
+        def forward(self, x):
+            out = x
+            for m in self:
+                out = m(out)
+            return out
+
+    @property
+    def world_size(self) -> int:
+        return 1
+
+    @skip_if_lt_x_gpu(1)
+    def test_moduledict_subclass_with_forward(self):
+        model = self.DictWithForward(8, 8)
+        mesh = init_device_mesh(device_type.type, (self.world_size,))
+        # Should not raise due to container type since forward() is implemented
+        fsdp_model = fully_shard(model, mesh=mesh)
+        x = torch.randn(2, 8, device=device_type)
+        _ = fsdp_model(x)
+
+    @skip_if_lt_x_gpu(1)
+    def test_modulelist_subclass_with_forward(self):
+        model = self.ListWithForward(8, 8)
+        mesh = init_device_mesh(device_type.type, (self.world_size,))
+        fsdp_model = fully_shard(model, mesh=mesh)
+        x = torch.randn(2, 8, device=device_type)
+        _ = fsdp_model(x)
+
+
 class TestFullyShardMeshArg(FSDPTestMultiThread):
     """Tests the ``mesh`` argument."""
 
@@ -662,7 +713,8 @@ def test_meta_device_1d_init(self):
 
     @skip_if_lt_x_gpu(1)
     def test_meta_device_2d_init(self):
-        assert self.world_size >= 4, f"{self.world_size}"
+        if not (self.world_size >= 4):
+            raise AssertionError(f"Expected world_size >= 4, but got {self.world_size}")
         dp_size = 2
         global_mesh = init_device_mesh(
             device_type.type,
@@ -836,7 +888,8 @@ def world_size(self) -> int:
 
     @skip_if_lt_x_gpu(1)
     def test_1d_process_group_init(self):
-        assert self.world_size == 4, f"{self.world_size}"
+        if not (self.world_size == 4):
+            raise AssertionError(f"Expected world_size == 4, but got {self.world_size}")
         # For convenience, use device mesh's infra to construct the DP PG
         # (in practice, the trainer would do it manually via `new_group()`)
         dp_size = 2
@@ -899,9 +952,10 @@ def test_1d_process_group_init(self):
     @skip_if_lt_x_gpu(1)
     def test_2d_process_group_init(self):
         shard_mesh_dim_size = 2
-        assert self.world_size % shard_mesh_dim_size == 0, (
-            f"Expects {self.world_size} to be divisible by {shard_mesh_dim_size}"
-        )
+        if not (self.world_size % shard_mesh_dim_size == 0):
+            raise AssertionError(
+                f"Expects {self.world_size} to be divisible by {shard_mesh_dim_size}"
+            )
         replicate_mesh_dim_size = self.world_size // shard_mesh_dim_size
         mesh_dim_names = ("replicate", "shard")
         ref_mesh = init_device_mesh(
@@ -967,6 +1021,27 @@ def test_2d_process_group_init(self):
             self.assertEqual(param, ref_param)
             self.assertEqual(param.grad, ref_param.grad)
 
+        # Also cover the opt-in reduce-scatter PG setup on HSDP meshes. Each
+        # rank creates the communicator for its shard group, so the
+        # implementation must use local synchronization internally to avoid
+        # waiting for nonmember ranks in new_group's post-init barrier.
+        model.set_separate_reduce_scatter_group()
+        rs_process_groups = set()
+        for module in model.modules():
+            if isinstance(module, FSDPModule):
+                for param_group in module._get_fsdp_state()._fsdp_param_groups:
+                    self.assertIsNotNone(
+                        param_group.mesh_info.reduce_scatter_process_group
+                    )
+                    self.assertIsNot(
+                        param_group._reduce_scatter_process_group,
+                        param_group.mesh_info.shard_process_group,
+                    )
+                    rs_process_groups.add(
+                        id(param_group.mesh_info.reduce_scatter_process_group)
+                    )
+        self.assertEqual(len(rs_process_groups), 1)
+
 
 class TestFullyShardHSDPBroadcast(FSDPTestMultiThread):
     @property
@@ -1160,13 +1235,16 @@ def _init_models(self):
     def test_init_1d_transformer_shard_largest_dim(self):
         model, ref_model = self._init_models()
 
-        def shard_placement_fn(param: nn.Parameter) -> Optional[Shard]:
+        def shard_placement_fn(param: nn.Parameter) -> Shard | None:
             largest_dim = largest_dim_size = -1
             for dim, dim_size in enumerate(param.shape):
                 if dim_size > largest_dim_size:
                     largest_dim = dim
                     largest_dim_size = dim_size
-            assert largest_dim >= 0, f"{param.shape}"
+            if not (largest_dim >= 0):
+                raise AssertionError(
+                    f"Expected largest_dim >= 0, but got {largest_dim} for shape {param.shape}"
+                )
             return Shard(largest_dim)
 
         for layer in model.layers:
@@ -1188,7 +1266,7 @@ def shard_placement_fn(param: nn.Parameter) -> Optional[Shard]:
     def test_init_1d_transformer_shard_dim_neg1(self):
         model, ref_model = self._init_models()
 
-        def shard_placement_fn(param: nn.Parameter) -> Optional[Shard]:
+        def shard_placement_fn(param: nn.Parameter) -> Shard | None:
             # Check that FSDP will normalize this dim to non-negative
             return Shard(-1)
 
@@ -1210,12 +1288,15 @@ def test_init_2d_transformer_shard_diff_dim(self):
         )
         model = Transformer.parallelize(model, global_mesh["tp"], use_seq_parallel=True)
 
-        def shard_placement_fn(param: nn.Parameter) -> Optional[Shard]:
+        def shard_placement_fn(param: nn.Parameter) -> Shard | None:
             if isinstance(param, DTensor):
                 for placement in param.placements:
                     if isinstance(placement, Shard):
                         shard_dim = param.ndim - 1 - placement.dim
-                        assert shard_dim >= 0, f"{param.shape}"
+                        if not (shard_dim >= 0):
+                            raise AssertionError(
+                                f"Expected shard_dim >= 0, but got {shard_dim} for shape {param.shape}"
+                            )
                         return Shard(shard_dim)
             return Shard(0)
 
@@ -1255,15 +1336,21 @@ def test_init_1d_uneven_shard_largest_dim(self):
         torch.manual_seed(42)
         model = nn.Sequential(nn.Linear(16, 17), nn.Linear(17, 8))
 
-        def shard_placement_fn(param: nn.Parameter) -> Optional[Shard]:
+        def shard_placement_fn(param: nn.Parameter) -> Shard | None:
             largest_dim = -1
             largest_dim_size = -1
             for dim, dim_size in enumerate(param.shape):
                 if dim_size > largest_dim_size:
                     largest_dim = dim
                     largest_dim_size = dim_size
-            assert largest_dim >= 0, f"{param.shape}"
-            assert largest_dim < param.ndim, f"{largest_dim=} {param.shape}"
+            if not (largest_dim >= 0):
+                raise AssertionError(
+                    f"Expected largest_dim >= 0, but got {largest_dim} for shape {param.shape}"
+                )
+            if not (largest_dim < param.ndim):
+                raise AssertionError(
+                    f"Expected largest_dim < param.ndim, but got {largest_dim=} {param.shape}"
+                )
             return Shard(largest_dim)
 
         with self.assertRaisesRegex(
@@ -1274,7 +1361,7 @@ def shard_placement_fn(param: nn.Parameter) -> Optional[Shard]:
     def test_invalid_shard_dim(self):
         model = nn.Sequential(nn.Linear(16, 16), nn.Linear(16, 8))
 
-        def shard_placement_fn(param: nn.Parameter) -> Optional[Shard]:
+        def shard_placement_fn(param: nn.Parameter) -> Shard | None:
             return Shard(1)
 
         # Shard(1) is invalid for 1D bias parameters
@@ -1334,5 +1421,104 @@ def forward(self, input):
         model(0)
 
 
+class TestFullyShardNonFloatParam(FSDPTest):
+    @property
+    def world_size(self) -> int:
+        return 2
+
+    @skip_if_lt_x_gpu(2)
+    def test_non_float_param(self):
+        """Non-float params (e.g. uint8, int64) are supported in all-gather,
+        excluded from reduce-scatter, and not cast by mp_policy.param_dtype."""
+        self.run_subtests(
+            {
+                "non_float_dtypes": [
+                    [torch.uint8],
+                    [torch.int64],
+                    [torch.uint8, torch.int8, torch.int16, torch.int32, torch.int64],
+                ],
+                "mp_policy": [
+                    MixedPrecisionPolicy(),
+                    MixedPrecisionPolicy(
+                        param_dtype=torch.bfloat16, reduce_dtype=torch.float32
+                    ),
+                ],
+                "frozen_float": [False, True],
+            },
+            self._test_non_float_param,
+        )
+
+    def _test_non_float_param(
+        self,
+        non_float_dtypes: list,
+        mp_policy: MixedPrecisionPolicy,
+        frozen_float: bool,
+    ):
+        class Model(nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.linear = nn.Linear(16, 16)
+                if frozen_float:
+                    self.frozen_float = nn.Parameter(
+                        torch.randn(16), requires_grad=False
+                    )
+                for i, dtype in enumerate(non_float_dtypes):
+                    self.register_parameter(
+                        f"non_float_{i}",
+                        nn.Parameter(
+                            torch.randint(0, 127, (16,), dtype=dtype),
+                            requires_grad=False,
+                        ),
+                    )
+
+            def forward(self, x):
+                return self.linear(x)
+
+        model = Model()
+        fully_shard(model, mp_policy=mp_policy)
+        for p in model.parameters():
+            self.assertEqual(p.size(0) % self.world_size, 0)
+        ag_input_dtypes = set()
+        expected_ag_output_bytes = 0
+        for p in model.parameters():
+            if p.dtype.is_floating_point and mp_policy.param_dtype is not None:
+                ag_input_dtypes.add(mp_policy.param_dtype)
+                expected_ag_output_bytes += p.numel() * mp_policy.param_dtype.itemsize
+            else:
+                # Non-float params keep their original dtype; param_dtype
+                # only applies to floating-point params
+                ag_input_dtypes.add(p.dtype)
+                expected_ag_output_bytes += p.numel() * p.element_size()
+        expected_ag_dtype = (
+            next(iter(ag_input_dtypes)) if len(ag_input_dtypes) == 1 else torch.uint8
+        )
+        orig_ag = dist.all_gather_into_tensor
+
+        def assert_all_gather(*args, **kw):
+            output = kw.get("output", args[0] if len(args) > 0 else None)
+            input = kw.get("input_tensor", args[1] if len(args) > 1 else None)
+            self.assertEqual(input.dtype, expected_ag_dtype)
+            self.assertEqual(output.nbytes, expected_ag_output_bytes)
+            return orig_ag(*args, **kw)
+
+        expected_rs_input_numel = sum(
+            p.numel() for p in model.parameters() if p.requires_grad
+        )
+        orig_rs = dist.reduce_scatter_tensor
+
+        def assert_reduce_scatter(*args, **kw):
+            input = kw.get("input", args[1] if len(args) > 1 else None)
+            self.assertEqual(input.numel(), expected_rs_input_numel)
+            return orig_rs(*args, **kw)
+
+        x = torch.randn(4, 16, device=device_type)
+        with (
+            patch_all_gather(assert_all_gather),
+            patch_reduce_scatter(assert_reduce_scatter),
+        ):
+            loss = model(x).sum()
+            loss.backward()
+
+
 if __name__ == "__main__":
     run_tests()
diff --git a/test/distributed/_composable/fsdp/test_fully_shard_memory.py b/test/distributed/_composable/fsdp/test_fully_shard_memory.py
index 3174889163b79..689c4f7af8e05 100644
--- a/test/distributed/_composable/fsdp/test_fully_shard_memory.py
+++ b/test/distributed/_composable/fsdp/test_fully_shard_memory.py
@@ -59,9 +59,10 @@ def _test_fully_shard_training_memory(
             )
         ):
             return  # skip since not a common use case
-        assert self.world_size == 2, (
-            f"Requires world size of 2 since some values are hard coded: {self.world_size}"
-        )
+        if self.world_size != 2:
+            raise AssertionError(
+                f"Requires world size of 2 since some values are hard coded: {self.world_size}"
+            )
         torch.manual_seed(42)
         # Pre-run a linear forward (gemm and bias) and backward (gemm) to
         # allocate the cuBLAS workspaces before measuring the memory usage
@@ -154,7 +155,8 @@ def _test_fully_shard_training_memory(
                 # Sharded parameters
                 expected_mem_mb += model_sharded_numel * 4 / 1e6
         else:
-            assert not use_cpu_offload
+            if use_cpu_offload:
+                raise AssertionError("Expected use_cpu_offload to be False")
             # Sharded parameters, unsharded parameters, 1x max unsharded block parameters
             # (copy-out) and other (peak at end of forward)
             expected_mem_mb = (
@@ -184,7 +186,8 @@ def _test_fully_shard_training_memory(
                     # 2x sharded parameters/gradients
                     expected_mem_mb += 2 * model_sharded_numel * 4 / 1e6
         else:
-            assert not use_cpu_offload
+            if use_cpu_offload:
+                raise AssertionError("Expected use_cpu_offload to be False")
             # Sharded parameters, unsharded parameters, 1.5x max unsharded
             # block parameters (reduce-scatter input/output), and other (peak
             # at beginning of backward)
@@ -223,6 +226,65 @@ def _test_fully_shard_training_memory(
             expected_mem_mb += (2 * model_sharded_numel) * 4 / 1e6 + buffer_mb
         self.assertLessEqual(mem_mb - base_mem_mb, expected_mem_mb)
 
+    @skip_if_lt_x_gpu(2)
+    def test_fully_shard_training_memory_no_gc(self):
+        """Memory should not grow across training steps when GC is disabled.
+
+        Regression test: reference cycles in FSDP's autograd integration can
+        hold GPU tensors alive when Python's cyclic GC is disabled. This
+        catches leaks like the one introduced in #173415.
+        """
+        torch.manual_seed(42)
+        gc.disable()
+        try:
+            vocab_size = 1024
+            model_args = ModelArgs(
+                vocab_size=vocab_size,
+                n_layers=6,
+                dim=2048,
+                n_heads=16,
+                weight_tying=False,
+            )
+            model = Transformer(model_args)
+            for module in model.modules():
+                if isinstance(module, TransformerBlock):
+                    fully_shard(module.attention)
+                    fully_shard(module.feed_forward)
+                    fully_shard(module)
+            fully_shard(model)
+            optim = torch.optim.Adam(model.parameters(), lr=1e-2, foreach=False)
+            inp = torch.randint(0, vocab_size, (2, 16), device=device_type.type)
+
+            # Warm-up step to stabilize memory (cuBLAS workspaces, etc.)
+            loss = model(inp)
+            loss.sum().backward()
+            optim.step()
+            optim.zero_grad()
+            gc.collect()  # one-time collection after warm-up
+            torch.get_device_module(device_type).synchronize()
+            mem_after_warmup = self._get_curr_active_memory_mb()
+
+            num_steps = 10
+            for _ in range(num_steps):
+                loss = model(inp)
+                loss.sum().backward()
+                optim.step()
+                optim.zero_grad()
+
+            torch.get_device_module(device_type).synchronize()
+            mem_after_steps = self._get_curr_active_memory_mb()
+            # Allow a small buffer (2 MB) for non-determinism, but no
+            # per-step growth should occur.
+            self.assertLessEqual(
+                mem_after_steps - mem_after_warmup,
+                2,
+                f"Memory grew by {mem_after_steps - mem_after_warmup} MB over "
+                f"{num_steps} steps with gc disabled, indicating a reference "
+                f"cycle leak in FSDP's autograd graph",
+            )
+        finally:
+            gc.enable()
+
     @skip_if_lt_x_gpu(2)
     def test_fully_shard_del_memory(self):
         base_mem_mb = self._get_peak_active_memory_mb()
diff --git a/test/distributed/_composable/fsdp/test_fully_shard_mixed_precision.py b/test/distributed/_composable/fsdp/test_fully_shard_mixed_precision.py
index 212420c784516..6bc4b7ad064f4 100644
--- a/test/distributed/_composable/fsdp/test_fully_shard_mixed_precision.py
+++ b/test/distributed/_composable/fsdp/test_fully_shard_mixed_precision.py
@@ -3,7 +3,7 @@
 import copy
 import dataclasses
 import functools
-from typing import Optional, Union
+from typing import Any, NamedTuple
 
 import torch
 import torch.distributed as dist
@@ -45,9 +45,9 @@ def world_size(self) -> int:
 
     def _init_models_and_optims(
         self,
-        reshard_after_forward: Union[bool, int],
-        param_dtype: Optional[torch.dtype],
-        reduce_dtype: Optional[torch.dtype],
+        reshard_after_forward: bool | int,
+        param_dtype: torch.dtype | None,
+        reduce_dtype: torch.dtype | None,
         use_shard_placement_fn,
     ):
         torch.manual_seed(42)
@@ -55,14 +55,17 @@ def _init_models_and_optims(
         ref_model = copy.deepcopy(model).to(device_type)
         ref_optim = torch.optim.Adam(ref_model.parameters(), lr=1e-2)
 
-        def _shard_placement_fn(param: nn.Parameter) -> Optional[Shard]:
+        def _shard_placement_fn(param: nn.Parameter) -> Shard | None:
             largest_dim = -1
             largest_dim_size = -1
             for dim, dim_size in enumerate(param.shape):
                 if dim_size > largest_dim_size:
                     largest_dim = dim
                     largest_dim_size = dim_size
-            assert largest_dim >= 0, f"{param.shape}"
+            if not (largest_dim >= 0):
+                raise AssertionError(
+                    f"Expected largest_dim >= 0, got {largest_dim} for {param.shape}"
+                )
             return Shard(largest_dim)
 
         mp_policy = MixedPrecisionPolicy(
@@ -109,7 +112,7 @@ def test_compute_dtype(self):
     def _test_compute_dtype(
         self,
         param_dtype: torch.dtype,
-        reshard_after_forward: Union[bool, int],
+        reshard_after_forward: bool | int,
         use_shard_placement_fn: bool,
     ):
         ref_model, ref_optim, model, optim = self._init_models_and_optims(
@@ -193,7 +196,7 @@ def test_reduce_dtype(self):
         )
 
     def _test_reduce_dtype_fp32_reduce(
-        self, reshard_after_forward: Union[bool, int], use_shard_placement_fn: bool
+        self, reshard_after_forward: bool | int, use_shard_placement_fn: bool
     ):
         if (
             self.world_size > 2
@@ -248,7 +251,7 @@ def assert_fn(output: torch.Tensor):
             check_sharded_parity(self, ref_model, model)
 
     def _test_reduce_dtype_bf16_reduce(
-        self, reshard_after_forward: Union[bool, int], use_shard_placement_fn: bool
+        self, reshard_after_forward: bool | int, use_shard_placement_fn: bool
     ):
         param_dtype, reduce_dtype = torch.float32, torch.bfloat16
         ref_model, ref_optim, model, optim = self._init_models_and_optims(
@@ -388,6 +391,189 @@ def assert_fn(output: torch.Tensor):
             ):
                 ref_param_compute.detach().copy_(ref_param)
 
+    def _reduce_1d_partial_grads(
+        self, module: nn.Module, group: dist.ProcessGroup | None = None
+    ) -> None:
+        # The reference model is replicated and sees the full global batch, so
+        # its gradients are already the global sum. Dividing by world_size
+        # simulates the average that FSDP's all-reduce produces.
+        group = group or dist.distributed_c10d._get_default_group()
+        for param in module.parameters():
+            if param.grad is not None:
+                param.grad.div_(group.size())
+
+    @skip_if_lt_x_gpu(2)
+    def test_structured_input_output(self):
+        """
+        Tests numeric parity between FSDP and a reference model when using
+        structured (dataclass, nested dataclass, NamedTuple) activations
+        between FSDP-wrapped modules, with and without mixed precision.
+        """
+        self.run_subtests(
+            {
+                "container_type": ["dataclass", "nested_dataclass", "namedtuple"],
+                "mp_config": [
+                    (None, None),  # fp32
+                    (torch.bfloat16, torch.bfloat16),  # bf16 compute + reduce
+                    (torch.bfloat16, torch.float32),  # bf16 compute, fp32 reduce
+                ],
+            },
+            self._test_structured_input_output,
+        )
+
+    def _test_structured_input_output(
+        self,
+        container_type: str,
+        mp_config: tuple[torch.dtype | None, torch.dtype | None],
+    ):
+        param_dtype, reduce_dtype = mp_config
+
+        @dataclasses.dataclass
+        class DC:
+            x: torch.Tensor
+            y: torch.Tensor
+
+        @dataclasses.dataclass
+        class Inner:
+            x: torch.Tensor
+            y: torch.Tensor
+
+        @dataclasses.dataclass
+        class Outer:
+            inner: Inner
+            z: torch.Tensor
+
+        class TensorPair(NamedTuple):
+            x: torch.Tensor
+            y: torch.Tensor
+
+        class ContainerModule(nn.Module):
+            """Applies a linear layer to each tensor in the container."""
+
+            def __init__(self, ctype: str, dim: int):
+                super().__init__()
+                self.ctype = ctype
+                self._layer = nn.Linear(dim, dim)
+
+            def forward(self, inp: Any) -> Any:
+                if self.ctype == "dataclass":
+                    return DC(x=self._layer(inp.x), y=self._layer(inp.y))
+                elif self.ctype == "nested_dataclass":
+                    return Outer(
+                        inner=Inner(
+                            x=self._layer(inp.inner.x),
+                            y=self._layer(inp.inner.y),
+                        ),
+                        z=self._layer(inp.z),
+                    )
+                else:  # namedtuple
+                    return TensorPair(x=self._layer(inp.x), y=self._layer(inp.y))
+
+        class ToContainer(nn.Module):
+            """Converts a plain tensor into a structured container."""
+
+            def __init__(self, ctype: str):
+                super().__init__()
+                self.ctype = ctype
+
+            def forward(self, x: torch.Tensor) -> Any:
+                # clone() so each field has an independent autograd path
+                if self.ctype == "dataclass":
+                    return DC(x=x, y=x.clone())
+                elif self.ctype == "nested_dataclass":
+                    return Outer(inner=Inner(x=x, y=x.clone()), z=x.clone())
+                else:  # namedtuple
+                    return TensorPair(x=x, y=x.clone())
+
+        class FromContainer(nn.Module):
+            """Extracts and sums tensors from a structured container."""
+
+            def __init__(self, ctype: str):
+                super().__init__()
+                self.ctype = ctype
+
+            def forward(self, inp: Any) -> torch.Tensor:
+                if self.ctype == "nested_dataclass":
+                    return inp.inner.x + inp.inner.y + inp.z
+                else:
+                    return inp.x + inp.y
+
+        torch.manual_seed(42)
+        dim = 16
+        local_batch_size = 2
+        global_batch_size = self.world_size * local_batch_size
+
+        model = nn.Sequential(
+            ToContainer(container_type),
+            ContainerModule(container_type, dim),
+            ContainerModule(container_type, dim),
+            FromContainer(container_type),
+        )
+        ref_model = copy.deepcopy(model).to(device_type)
+
+        mp_policy = MixedPrecisionPolicy(
+            param_dtype=param_dtype, reduce_dtype=reduce_dtype
+        )
+        for module in model:
+            fully_shard(module, mp_policy=mp_policy)
+        fully_shard(model, mp_policy=mp_policy)
+
+        if param_dtype is not None:
+            # Maintain a bf16 copy for compute, fp32 copy for optimizer
+            ref_model_compute = copy.deepcopy(ref_model).to(param_dtype)
+        ref_optim = torch.optim.Adam(ref_model.parameters(), lr=1e-2)
+        optim = torch.optim.Adam(model.parameters(), lr=1e-2)
+
+        torch.manual_seed(1)  # same on all ranks
+        for iter_idx in range(10):
+            global_inp = torch.rand((global_batch_size, dim), device=device_type.type)
+            local_inp = global_inp[
+                self.rank * local_batch_size : (self.rank + 1) * local_batch_size
+            ].detach()
+
+            # FSDP forward/backward
+            optim.zero_grad(set_to_none=(iter_idx % 2 == 0))
+            fsdp_loss = model(local_inp).sum()
+            fsdp_loss.backward()
+            optim.step()
+
+            # Ref forward/backward
+            ref_optim.zero_grad(set_to_none=(iter_idx % 2 == 0))
+            if param_dtype is not None:
+                ref_loss = ref_model_compute(global_inp.to(param_dtype)).sum()
+                ref_loss.backward()
+                # Simulate gradient reduction matching FSDP's behavior
+                if reduce_dtype is not None and reduce_dtype != param_dtype:
+                    # Cast grads to reduce_dtype, all-reduce, divide
+                    for p in ref_model_compute.parameters():
+                        p.grad.data = p.grad.to(reduce_dtype)
+                        dist.all_reduce(p.grad)
+                        p.grad.div_(self.world_size)
+                else:
+                    self._reduce_1d_partial_grads(ref_model_compute)
+                for p_fp32, p_compute in zip(
+                    ref_model.parameters(), ref_model_compute.parameters()
+                ):
+                    p_fp32.grad = p_compute.grad.to(p_fp32.dtype)
+                    p_compute.grad = None
+                ref_optim.step()
+                for p_fp32, p_compute in zip(
+                    ref_model.parameters(), ref_model_compute.parameters()
+                ):
+                    p_compute.detach().copy_(p_fp32)
+            else:
+                ref_loss = ref_model(global_inp).sum()
+                ref_loss.backward()
+                self._reduce_1d_partial_grads(ref_model)
+                ref_optim.step()
+
+            dist.all_reduce(fsdp_loss)  # partial -> replicated
+            self.assertEqual(fsdp_loss, ref_loss)
+            # bf16 gradient accumulation introduces drift beyond default
+            # fp32 tolerances, so only check full param/grad parity for fp32.
+            if param_dtype is None:
+                check_sharded_parity(self, ref_model, model)
+
 
 class TestFullyShardMixedPrecisionCasts(FSDPTestMultiThread):
     @property
@@ -596,30 +782,6 @@ def assert_fn(output: torch.Tensor):
             loss = model(inp).sum()
             loss.backward()
 
-    @skip_if_lt_x_gpu(1)
-    def test_dataclass_input(self):
-        @dataclasses.dataclass
-        class Input:
-            x: torch.Tensor
-
-        class Model(nn.Module):
-            def __init__(self, *args, **kwargs) -> None:
-                super().__init__(*args, **kwargs)
-                self._layer = nn.Linear(10, 10)
-
-            def forward(self, input: Input):
-                return self._layer(input.x)
-
-        mp_policy = MixedPrecisionPolicy(
-            torch.bfloat16, torch.bfloat16, torch.bfloat16, True
-        )
-        model = Model()
-        inp = Input(torch.randn(2, 10).to(device_type))
-
-        fully_shard(model, mp_policy=mp_policy)
-        loss = model(inp).sum()
-        loss.backward()
-
 
 if __name__ == "__main__":
     run_tests()
diff --git a/test/distributed/_composable/fsdp/test_fully_shard_overlap.py b/test/distributed/_composable/fsdp/test_fully_shard_overlap.py
index f4a33a86c417b..e1804fc37a7a6 100644
--- a/test/distributed/_composable/fsdp/test_fully_shard_overlap.py
+++ b/test/distributed/_composable/fsdp/test_fully_shard_overlap.py
@@ -1,5 +1,6 @@
 # Owner(s): ["oncall: distributed"]
 
+import contextlib
 import copy
 import functools
 import unittest
@@ -9,6 +10,11 @@
 import torch.distributed as dist
 import torch.nn as nn
 from torch.distributed.fsdp import fully_shard
+from torch.distributed.fsdp._fully_shard._fsdp_common import (
+    FSDPMeshInfo,
+    ShardPlacementResult,
+)
+from torch.distributed.tensor import init_device_mesh, Shard
 from torch.distributed.tensor.experimental import implicit_replication
 from torch.testing._internal.common_distributed import (
     skip_if_lt_x_gpu,
@@ -22,16 +28,33 @@
 )
 from torch.testing._internal.common_utils import (
     get_cycles_per_ms,
+    IS_LINUX,
     MI200_ARCH,
     run_tests,
     TEST_HPU,
 )
+from torch.testing._internal.distributed._tensor.common_dtensor import (
+    ModelArgs,
+    Transformer,
+)
 
 
 device_type = torch.device(get_devtype())
 device_module = torch.get_device_module(device_type)
 
 
+def _time_fn(fn: Callable):
+    start_event = device_module.Event(enable_timing=True)
+    end_event = device_module.Event(enable_timing=True)
+    dist.barrier()
+    device_module.synchronize()
+    start_event.record()
+    fn()
+    end_event.record()
+    device_module.synchronize()
+    return start_event.elapsed_time(end_event)
+
+
 class TestFullyShardOverlap(FSDPTest):
     """
     NOTE: Testing stream overlap in PyTorch CI is tricky.
@@ -64,7 +87,8 @@ def test_fully_shard_training_overlap(self):
         )
         ref_model = copy.deepcopy(model).to(device_type)
         for lin in model:
-            assert len(list(lin.parameters())) == 1, "Expects only one weight"
+            if len(list(lin.parameters())) != 1:
+                raise AssertionError("Expects only one weight")
             fully_shard(lin, reshard_after_forward=True)
         fully_shard(model, reshard_after_forward=True)
 
@@ -113,8 +137,8 @@ def fwd():
             with patch_all_gather(delayed_all_gather):
                 model(inp)
 
-        ref_fwd_time = self._time_fn(ref_fwd)
-        fwd_time = self._time_fn(fwd)
+        ref_fwd_time = _time_fn(ref_fwd)
+        fwd_time = _time_fn(fwd)
         # Forward: only 1st all-gather is exposed
         # NOTE: Do not enforce the expected forward time due to flakiness in CI
         # expected_fwd_time = comm_sleep_ms + num_linears * compute_sleep_ms + buffer_ms
@@ -155,8 +179,8 @@ def fwd_bwd():
                 loss = model(inp).sum()
                 loss.backward()
 
-        ref_fwd_bwd_time = self._time_fn(ref_fwd_bwd)
-        fwd_bwd_time = self._time_fn(fwd_bwd)
+        ref_fwd_bwd_time = _time_fn(ref_fwd_bwd)
+        fwd_bwd_time = _time_fn(fwd_bwd)
         # Backward: only 1st all-gather and last reduce-scatter are exposed;
         # double the backward compute since computing two gradients per layer
         # NOTE: Do not enforce the expected forward-backward time due to
@@ -166,6 +190,135 @@ def fwd_bwd():
         # )
         self.assertLessEqual(fwd_bwd_time, ref_fwd_bwd_time)
 
+    @skip_if_rocm_arch_multiprocess(MI200_ARCH)
+    @skip_if_lt_x_gpu(2)
+    def test_fully_shard_backward_comm_overlap(self):
+        """Exercise backward with reduce-scatter sharing the shard process
+        group and with reduce-scatter opted in to a dedicated process group via
+        set_separate_reduce_scatter_group.
+
+        Uses real compute (large matmuls) and real collectives (large
+        parameter tensors) instead of CUDA sleeps.
+
+        ref_bwd: RS forced back onto the shard PG (single communicator),
+        serializing it with all-gather.
+
+        fsdp_bwd: RS on its own communicator (the opt-in), enabling overlap.
+        """
+        torch.manual_seed(42)
+
+        # Large dim for non-trivial collective time, small batch so
+        # comm dominates compute; AG becomes back-to-back on its stream
+        dim, num_linears, batch = 16384, 3, 4
+        model = nn.Sequential(
+            *[nn.Linear(dim, dim, bias=False) for _ in range(num_linears)]
+        )
+        for lin in model:
+            fully_shard(lin, reshard_after_forward=True)
+        fully_shard(model, reshard_after_forward=True)
+        # Opt in: give reduce-scatter its own process group (the change under test)
+        model.set_separate_reduce_scatter_group()
+
+        optim = torch.optim.SGD(model.parameters(), lr=1e-2)
+        inp = torch.randn((batch, dim), device=device_type.type)
+        loss = model(inp).sum()
+        loss.backward()
+        with implicit_replication():
+            optim.step()
+        optim.zero_grad()  # warmup
+
+        # Collect FSDP param groups to toggle PG config
+        fsdp_param_groups = []
+        for module in model.modules():
+            state = fully_shard.state(module)
+            if state is not None:
+                fsdp_param_groups.extend(state._fsdp_param_groups)
+
+        # The opt-in created a dedicated RS group per mesh (distinct from shard)
+        for pg in fsdp_param_groups:
+            self.assertIsNotNone(pg.mesh_info.reduce_scatter_process_group)
+            self.assertIsNot(
+                pg.mesh_info.reduce_scatter_process_group,
+                pg.mesh_info.shard_process_group,
+            )
+
+        # ref_bwd: force RS to use the same PG as AG (serialized)
+        saved_rs_pgs = []
+        for pg in fsdp_param_groups:
+            saved_rs_pgs.append(pg.mesh_info.reduce_scatter_process_group)
+            pg.mesh_info.reduce_scatter_process_group = pg.mesh_info.shard_process_group
+
+        def ref_bwd():
+            loss = model(inp).sum()
+            loss.backward()
+            with implicit_replication():
+                optim.step()
+            optim.zero_grad()
+
+        _time_fn(ref_bwd)
+
+        # Restore separate PGs for fsdp_bwd
+        for pg, saved_pg in zip(fsdp_param_groups, saved_rs_pgs):
+            pg.mesh_info.reduce_scatter_process_group = saved_pg
+
+        def fsdp_bwd():
+            loss = model(inp).sum()
+            loss.backward()
+            with implicit_replication():
+                optim.step()
+            optim.zero_grad()
+
+        # Exercise both shared-PG and separate-RS-PG backward paths. We do not
+        # assert timing here since real collective overlap is hardware/backend
+        # sensitive.
+        _time_fn(fsdp_bwd)
+
+    @skip_if_lt_x_gpu(2)
+    def test_set_separate_reduce_scatter_group(self):
+        """Reduce-scatter shares the shard PG by default; enabling gives it a
+        dedicated PG (one communicator shared across same-rank-set meshes), and
+        disabling resets to the shared PG."""
+        model = nn.Sequential(nn.Linear(8, 8), nn.Linear(8, 8))
+        for lin in model:
+            fully_shard(lin)
+        fully_shard(model)
+        param_groups = []
+        for module in model.modules():
+            state = fully_shard.state(module)
+            if state is not None:
+                param_groups.extend(state._fsdp_param_groups)
+
+        def _assert_shares_all_gather():
+            for pg in param_groups:
+                self.assertIsNone(pg.mesh_info.reduce_scatter_process_group)
+                self.assertIs(
+                    pg._reduce_scatter_process_group,
+                    pg.mesh_info.shard_process_group,
+                )
+
+        # Default shares the all-gather (shard) PG
+        _assert_shares_all_gather()
+
+        # Enable: one dedicated PG shared across same-rank-set meshes (dedup),
+        # distinct from the shard PG
+        model.set_separate_reduce_scatter_group()
+        rs_pgs = {id(pg.mesh_info.reduce_scatter_process_group) for pg in param_groups}
+        self.assertEqual(len(rs_pgs), 1)  # one communicator for the single rank set
+        for pg in param_groups:
+            self.assertIsNotNone(pg.mesh_info.reduce_scatter_process_group)
+            self.assertIs(
+                pg._reduce_scatter_process_group,
+                pg.mesh_info.reduce_scatter_process_group,
+            )
+            self.assertIsNot(
+                pg._reduce_scatter_process_group, pg.mesh_info.shard_process_group
+            )
+
+        # Disable: reset back to sharing the shard PG
+        model.set_separate_reduce_scatter_group(False)
+        _assert_shares_all_gather()
+
+    @unittest.skipIf(IS_LINUX, "https://github.com/pytorch/pytorch/issues/131081")
     @skip_if_lt_x_gpu(2)
     @unittest.skipIf(TEST_HPU, "Sleep is not supported on HPU")
     def test_fully_shard_post_optim_event_overlap(self):
@@ -209,10 +362,8 @@ def run_train_steps(num_iters: int, use_post_optim_event: bool):
 
         run_train_steps(1, False)  # warmup CUDA and allocator
         num_iters = 5
-        baseline_time = self._time_fn(
-            functools.partial(run_train_steps, num_iters, False)
-        )
-        test_time = self._time_fn(functools.partial(run_train_steps, num_iters, True))
+        baseline_time = _time_fn(functools.partial(run_train_steps, num_iters, False))
+        test_time = _time_fn(functools.partial(run_train_steps, num_iters, True))
 
         buffer_ms = 4  # CPU delays and copies
         # Baseline: FSDP all-gather is exposed since the FSDP module waits for
@@ -234,18 +385,6 @@ def run_train_steps(num_iters: int, use_post_optim_event: bool):
         # the test time rather than the expected test time
         self.assertGreater(baseline_time, test_time)
 
-    def _time_fn(self, fn: Callable):
-        start_event = device_module.Event(enable_timing=True)
-        end_event = device_module.Event(enable_timing=True)
-        dist.barrier()
-        device_module.synchronize()
-        start_event.record()
-        fn()
-        end_event.record()
-        device_module.synchronize()
-        elapsed_time = start_event.elapsed_time(end_event)
-        return elapsed_time
-
 
 class Matmul(torch.autograd.Function):
     # Use CUDA sleeps to emulate compute time
@@ -277,5 +416,227 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
         return nn.functional.relu(Matmul.apply(x, self.weight, self.sleep_ms))
 
 
+class TestFullyShardPerParamMeshOverlap(FSDPTest):
+    @property
+    def world_size(self) -> int:
+        return min(4, torch.get_device_module(device_type).device_count())
+
+    @staticmethod
+    @contextlib.contextmanager
+    def _delayed_foreach_reduce(sleep_ms: int):
+        """Inject a CUDA sleep after each reduce-scatter / all-reduce.
+
+        Each process group gets its own delay stream so that delays for
+        different groups run concurrently (just like real NCCL ops on
+        separate communicators).  The event returned to FSDP is replaced
+        with a delayed event on this per-group stream.
+
+        This amplifies the difference between per-group and shared RS
+        state:
+        - Per-group RS state: the compute stream waits only on its own
+          group's delayed event, so dp and efsdp delays overlap.
+        - Shared RS state: the compute stream waits on *all* groups'
+          delayed events, serializing the dp and efsdp delays.
+        """
+        import torch.distributed.fsdp._fully_shard._fsdp_param_group as _pg_mod
+
+        orig = _pg_mod.foreach_reduce
+        # One delay stream per process group so that sleeps for different
+        # groups (e.g. dp vs efsdp reduce-scatter) execute in parallel,
+        # mirroring how real NCCL ops on separate communicators overlap.
+        delay_streams: dict[dist.ProcessGroup, torch.cuda.Stream] = {}
+
+        def wrapped(
+            fsdp_params,
+            unsharded_grads,
+            reduce_scatter_group,
+            reduce_scatter_stream,
+            *args,
+            **kwargs,
+        ):
+            result = orig(
+                fsdp_params,
+                unsharded_grads,
+                reduce_scatter_group,
+                reduce_scatter_stream,
+                *args,
+                **kwargs,
+            )
+            rs_input, rs_event, post_reduce_stream, *rest = result
+            if reduce_scatter_group not in delay_streams:
+                delay_streams[reduce_scatter_group] = device_module.Stream()
+            ds = delay_streams[reduce_scatter_group]
+            ds.wait_event(rs_event)
+            with device_module.stream(ds):
+                device_module._sleep(int(sleep_ms * get_cycles_per_ms()))
+                delayed_event = ds.record_event()
+            return (rs_input, delayed_event, post_reduce_stream, *rest)
+
+        dist.barrier()
+        _pg_mod.foreach_reduce = wrapped
+        try:
+            yield
+        finally:
+            dist.barrier()
+            _pg_mod.foreach_reduce = orig
+
+    @skip_if_rocm_arch_multiprocess(MI200_ARCH)
+    @skip_if_lt_x_gpu(4)
+    @unittest.skipIf(TEST_HPU, "Sleep is not supported on HPU")
+    def test_fully_shard_per_param_mesh_training_overlap(self):
+        self._test_per_param_mesh_overlap(simulate_no_grad_input=False)
+
+    @skip_if_rocm_arch_multiprocess(MI200_ARCH)
+    @skip_if_lt_x_gpu(4)
+    @unittest.skipIf(TEST_HPU, "Sleep is not supported on HPU")
+    def test_fully_shard_per_param_mesh_no_grad_input_overlap(self):
+        self._test_per_param_mesh_overlap(simulate_no_grad_input=True)
+
+    def _test_per_param_mesh_overlap(self, simulate_no_grad_input: bool):
+        """Verify per-param-group reduce-scatter state avoids cross-group
+        serialization in per-param-mesh FSDP with expert parallelism.
+
+        When ``simulate_no_grad_input`` is True, block inputs are detached so
+        that autograd's RegisterPostBackwardFunction nodes are never inserted,
+        exercising the safety-net path in _root_post_backward_final_callback.
+        """
+        import types
+
+        from torch.distributed._composable.replicate_with_fsdp import replicate
+
+        ep_degree = 2
+        efsdp_size = self.world_size // ep_degree
+        comm_sleep_ms = 500
+        world_mesh = init_device_mesh(
+            device_type.type,
+            (self.world_size,),
+            mesh_dim_names=("world",),
+        )
+        dp_mesh = world_mesh._unflatten(0, (self.world_size,), ("fsdp",))["fsdp"]
+        sparse_mesh = dp_mesh._unflatten(0, (efsdp_size, ep_degree), ("efsdp", "ep"))
+        ep_mesh = sparse_mesh["ep"]
+        dp_mesh_info = FSDPMeshInfo(mesh=dp_mesh, shard_mesh_dim=0)
+        efsdp_mesh_info = FSDPMeshInfo(mesh=sparse_mesh["efsdp"], shard_mesh_dim=0)
+        model_args = ModelArgs(
+            n_layers=20,
+            vocab_size=1024,
+            max_seq_len=64,
+            dim=1280,
+            n_heads=16,
+            dropout_p=0.0,
+            num_experts=2,
+        )
+        inp = torch.randint(
+            0,
+            model_args.vocab_size,
+            (2, model_args.max_seq_len),
+            device=device_type.type,
+        )
+
+        def _build_fsdp_model():
+            torch.manual_seed(42)
+            model = Transformer(model_args)
+            Transformer.parallelize(
+                model, tp_mesh=None, use_seq_parallel=False, ep_mesh=ep_mesh
+            )
+            for block in model.layers:
+                expert_params = set(block.expert_layer.experts.parameters())
+
+                def _shard_placement_fn(param, _expert_params=expert_params):
+                    if param in _expert_params:
+                        return ShardPlacementResult(
+                            placement=Shard(0), mesh_info=efsdp_mesh_info
+                        )
+                    return ShardPlacementResult(
+                        placement=Shard(0), mesh_info=dp_mesh_info
+                    )
+
+                fully_shard(block, mesh=dp_mesh, shard_placement_fn=_shard_placement_fn)
+            fully_shard(
+                [model.tok_embeddings, model.norm, model.output],
+                mesh=dp_mesh,
+                reshard_after_forward=True,
+            )
+            fully_shard(model, mesh=dp_mesh, reshard_after_forward=True)
+            blocks = model.layers
+            for i in range(len(blocks) - 1):
+                blocks[i].set_modules_to_forward_prefetch([blocks[i + 1]])
+            for i in range(1, len(blocks)):
+                blocks[i].set_modules_to_backward_prefetch([blocks[i - 1]])
+            return model
+
+        def _build_replicate_model():
+            torch.manual_seed(42)
+            model = Transformer(model_args)
+            Transformer.parallelize(
+                model, tp_mesh=None, use_seq_parallel=False, ep_mesh=ep_mesh
+            )
+            for block in model.layers:
+                expert_params = set(block.expert_layer.experts.parameters())
+                replicate(block, mesh=dp_mesh, ignored_params=expert_params)
+            replicate([model.tok_embeddings, model.norm, model.output], mesh=dp_mesh)
+            replicate(model, mesh=dp_mesh)
+            return model
+
+        # Both models measured under _delayed_foreach_reduce so the delay
+        # applies uniformly to foreach_reduce (all-reduce for replicate,
+        # reduce-scatter for FSDP).
+        with self._delayed_foreach_reduce(comm_sleep_ms):
+            # --- replicate + EP (baseline) ---
+            rep_model = _build_replicate_model()
+
+            def rep_fwd_bwd():
+                rep_model(inp).sum().backward()  # noqa: F821
+
+            for _ in range(5):
+                rep_fwd_bwd()
+                rep_model.zero_grad(set_to_none=True)
+            rep_time = _time_fn(rep_fwd_bwd)
+            rep_model.zero_grad(set_to_none=True)
+            del rep_model
+            # --- FSDP + EP ---
+            fsdp_model = _build_fsdp_model()
+
+            if simulate_no_grad_input:
+                # Detach inputs to each block so that autograd's
+                # RegisterPostBackwardFunction nodes are never inserted,
+                # forcing the safety-net path in
+                # _root_post_backward_final_callback.
+                def _no_grad_input_forward(self, tokens):
+                    h = self.tok_embeddings(tokens)
+                    h = h + self.pos_embeddings(
+                        torch.arange(tokens.size(1), device=tokens.device)
+                    )
+                    for layer in self.layers:
+                        h = h + layer(h.detach())
+                    h = self.norm(h)
+                    return self.output(h).float()
+
+                fsdp_model.forward = types.MethodType(
+                    _no_grad_input_forward, fsdp_model
+                )
+
+            def fsdp_fwd_bwd():
+                fsdp_model(inp).sum().backward()  # noqa: F821
+
+            for _ in range(5):
+                fsdp_fwd_bwd()
+                fsdp_model.zero_grad(set_to_none=True)
+            fsdp_time = _time_fn(fsdp_fwd_bwd)
+            fsdp_model.zero_grad(set_to_none=True)
+        # replicate: 1 all-reduce/block → N+1 serialized waits.
+        # FSDP: 2 reduce-scatters/block (dp + efsdp).
+        #   Per-group RS state: dp reduce-scatter overlaps with efsdp's
+        #     wait → ~N+1 waits ≈ replicate (ratio ≈ 1.0).
+        #   Shared RS state: both reduce-scatters serialize → ~2N+1
+        #     waits → ratio ≈ (2N+1)/(N+1) → 1.95 for N=20.
+        self.assertLess(
+            fsdp_time / rep_time,
+            1.5,
+            f"FSDP/replicate ratio {fsdp_time / rep_time:.2f} >= 1.5; "
+            f"per-group RS state may not be preventing cross-group stalls",
+        )
+
+
 if __name__ == "__main__":
     run_tests()
diff --git a/test/distributed/_composable/fsdp/test_fully_shard_state_dict.py b/test/distributed/_composable/fsdp/test_fully_shard_state_dict.py
index a23faed921137..9527120b99ff2 100644
--- a/test/distributed/_composable/fsdp/test_fully_shard_state_dict.py
+++ b/test/distributed/_composable/fsdp/test_fully_shard_state_dict.py
@@ -3,7 +3,6 @@
 import copy
 import functools
 from contextlib import nullcontext
-from typing import Optional
 
 import torch
 import torch.nn as nn
@@ -82,7 +81,7 @@ def _test_dp_state_dict_save_load(
             MLP(mlp_dim),
         )
 
-        def _shard_placement_fn(param: nn.Parameter) -> Optional[Shard]:
+        def _shard_placement_fn(param: nn.Parameter) -> Shard | None:
             largest_dim = largest_dim_size = -1
             for dim, dim_size in enumerate(param.shape):
                 if dim_size > largest_dim_size:
@@ -132,7 +131,8 @@ def _test_cached_state_dict(self, mlp_dim: int, mutate_after_state_dict: bool):
 
         # call .state_dict() once and use `sd` directly to reduce cpu overhead
         sd = model.state_dict()
-        assert isinstance(model.weight, DTensor)
+        if not isinstance(model.weight, DTensor):
+            raise AssertionError(f"Expected DTensor, got {type(model.weight)}")
 
         if not mutate_after_state_dict:
             self.assertTrue(
diff --git a/test/distributed/_composable/fsdp/test_fully_shard_training.py b/test/distributed/_composable/fsdp/test_fully_shard_training.py
index 5040fbf2f8bfd..e6ac99890eda5 100644
--- a/test/distributed/_composable/fsdp/test_fully_shard_training.py
+++ b/test/distributed/_composable/fsdp/test_fully_shard_training.py
@@ -3,11 +3,12 @@
 import contextlib
 import copy
 import functools
+import gc
 import itertools
 import unittest
 from collections import defaultdict
 from collections.abc import Callable, Iterable
-from typing import Any, Optional, Union
+from typing import Any
 
 import torch
 import torch.distributed as dist
@@ -30,6 +31,11 @@
     foreach_all_gather,
     foreach_reduce,
 )
+from torch.distributed.fsdp._fully_shard._fsdp_common import (
+    FSDPMeshInfo,
+    HSDPMeshInfo,
+    ShardPlacementResult,
+)
 from torch.distributed.tensor import DTensor, init_device_mesh, Shard
 from torch.distributed.tensor.debug import CommDebugMode
 from torch.testing._internal.common_distributed import (
@@ -94,13 +100,12 @@ def forward(self, x: torch.Tensor, ys: tuple[torch.Tensor, ...]):
                 # Skip device check for CPU since torch.device("cpu") and
                 # torch.device("cpu", 0) are semantically equivalent but not equal
                 if device.type != "cpu":
-                    assert x.device == device, f"Expects {device} but got {x.device}"
-                    assert ys[0].device == device, (
-                        f"Expects {device} but got {ys[0].device}"
-                    )
-                    assert ys[1].device == device, (
-                        f"Expects {device} but got {ys[1].device}"
-                    )
+                    if not (x.device == device):
+                        raise AssertionError(f"Expects {device} but got {x.device}")
+                    if not (ys[0].device == device):
+                        raise AssertionError(f"Expects {device} but got {ys[0].device}")
+                    if not (ys[1].device == device):
+                        raise AssertionError(f"Expects {device} but got {ys[1].device}")
                 y = ys[0] + ys[1]
                 return x + y + 1
 
@@ -327,7 +332,7 @@ def _test_train_parity_single_group(
         replicate(ref_model, device_ids=_get_device_ids(self.rank))
         ref_optim = torch.optim.Adam(ref_model.parameters(), lr=1e-2)
 
-        def _shard_placement_fn(param: nn.Parameter) -> Optional[Shard]:
+        def _shard_placement_fn(param: nn.Parameter) -> Shard | None:
             return Shard(param.shape.index(max(param.shape)))
 
         shard_placement_fn = _shard_placement_fn if use_shard_placement_fn else None
@@ -415,7 +420,7 @@ def test_train_parity_multi_group_unshard_async_op(self):
 
     def _test_train_parity_multi_group(
         self,
-        reshard_after_forward: Union[bool, int],
+        reshard_after_forward: bool | int,
         offload_policy: OffloadPolicy,
         test_device_type: str,
         delay_after_forward: bool,
@@ -440,7 +445,8 @@ def _test_train_parity_multi_group(
             and offload_policy.pin_memory
         ):
             return
-        assert test_device_type in ("cuda", "hpu", "xpu", "cpu"), f"{test_device_type}"
+        if test_device_type not in ("cuda", "hpu", "xpu", "cpu"):
+            raise AssertionError(f"Unexpected device type: {test_device_type}")
         torch.manual_seed(42)
         vocab_size = 1024
         model_args = ModelArgs(
@@ -582,7 +588,7 @@ def test_multi_forward_module(self):
             self._test_multi_forward_module,
         )
 
-    def _test_multi_forward_module(self, reshard_after_forward: Union[bool, int]):
+    def _test_multi_forward_module(self, reshard_after_forward: bool | int):
         class MultiForwardModule(nn.Module):
             def __init__(self, device: torch.device):
                 super().__init__()
@@ -731,11 +737,12 @@ def test_train_parity_with_activation_checkpointing(self):
 
     def _test_train_parity_with_activation_checkpointing(
         self,
-        reshard_after_forward: Union[bool, int],
+        reshard_after_forward: bool | int,
         checkpoint_impl: str,
         module_grouping: str,
     ):
-        assert checkpoint_impl in ("composable", "utils", "wrapper")
+        if checkpoint_impl not in ("composable", "utils", "wrapper"):
+            raise AssertionError(f"Unexpected checkpoint_impl: {checkpoint_impl}")
         testing_compile = fully_shard != torch.distributed.fsdp.fully_shard
         if testing_compile and checkpoint_impl == "composable":
             return
@@ -776,7 +783,10 @@ def _test_train_parity_with_activation_checkpointing(
         # Apply FSDP
         fsdp_kwargs = {"reshard_after_forward": reshard_after_forward}
         if module_grouping == "mem_eff":
-            assert model_args.n_layers == 3
+            if not (model_args.n_layers == 3):
+                raise AssertionError(
+                    f"Expected n_layers == 3, got {model_args.n_layers}"
+                )
             fully_shard(model.layers[0], **fsdp_kwargs)
             fully_shard([model.layers[1], model.layers[2]], **fsdp_kwargs)
             fully_shard([model.tok_embeddings, model.pos_embeddings], **fsdp_kwargs)
@@ -898,7 +908,7 @@ def test_train_parity_shard_placement_fn_shard_largest_dim(self):
         ref_model = copy.deepcopy(model).to(device_type)
         ref_optim = torch.optim.AdamW(ref_model.parameters(), lr=1e-2)
 
-        def shard_placement_fn(param: nn.Parameter) -> Optional[Shard]:
+        def shard_placement_fn(param: nn.Parameter) -> Shard | None:
             return Shard(param.shape.index(max(param.shape)))
 
         for layer in model.layers:
@@ -942,7 +952,7 @@ def test_shard_placement_fn_contiguous_params_grads(self):
         dim = 4
         model = MLP(dim=dim)
 
-        def shard_placement_fn(param: nn.Parameter) -> Optional[Shard]:
+        def shard_placement_fn(param: nn.Parameter) -> Shard | None:
             if param.ndim > 1:
                 return Shard(1)
             return Shard(0)
@@ -1025,6 +1035,140 @@ def _test_train_shared_params(
                 _optim.step()
             self.assertEqual(losses[0], losses[1])
 
+    @skip_if_lt_x_gpu(2, allow_cpu=True)
+    def test_train_shared_params_uneven_shard(self):
+        """
+        Test that FSDP2 correctly handles tied weights with uneven sharding
+        when tied modules are placed in the same FSDP group.
+        """
+        self.run_subtests(
+            {"reshard_after_forward": [True, False]},
+            self._test_train_shared_params_uneven_shard,
+        )
+
+    def _test_train_shared_params_uneven_shard(
+        self,
+        reshard_after_forward: bool,
+    ):
+        hidden_size = 16
+        # vocab_size not divisible by world_size to trigger uneven sharding
+        vocab_size = self.world_size * 4 + 1
+
+        class TiedModel(nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.tok_embeddings = nn.Embedding(vocab_size, hidden_size)
+                self.output = nn.Linear(hidden_size, vocab_size, bias=False)
+                self.output.weight = self.tok_embeddings.weight
+
+            def forward(self, ids):
+                return self.output(self.tok_embeddings(ids))
+
+        torch.manual_seed(42)
+        model = TiedModel()
+        # Shared params must be in the same FSDP group
+        fully_shard(
+            [model.tok_embeddings, model.output],
+            reshard_after_forward=reshard_after_forward,
+        )
+        fully_shard(model, reshard_after_forward=reshard_after_forward)
+
+        model.tok_embeddings.unshard()
+        initial_weight = model.tok_embeddings.weight.detach().clone()
+        model.tok_embeddings.reshard()
+
+        optimizer = torch.optim.AdamW(model.parameters(), lr=1e-2)
+
+        torch.manual_seed(42 + self.rank)
+        for _ in range(10):
+            ids = torch.randint(0, vocab_size, (2, 8), device=device_type.type)
+            logits = model(ids)
+            loss = nn.functional.cross_entropy(
+                logits.view(-1, vocab_size).float(), ids.view(-1)
+            )
+            loss.backward()
+            optimizer.step()
+            optimizer.zero_grad()
+
+        model.tok_embeddings.unshard()
+        final_weight = model.tok_embeddings.weight
+        # Every row should be updated after training, including on the last
+        # rank whose shard requires padding for uneven sharding.
+        for row in range(vocab_size):
+            self.assertFalse(
+                torch.equal(initial_weight[row], final_weight[row]),
+                f"Row {row} was not updated after training",
+            )
+        model.tok_embeddings.reshard()
+
+    @skip_if_lt_x_gpu(2, allow_cpu=True)
+    def test_shared_params_separate_fsdp_groups_error(self):
+        """
+        Test that applying fully_shard to modules with shared parameters
+        in separate calls raises an error at lazy init (first forward).
+        """
+        hidden_size = 16
+        vocab_size = 17
+
+        class TiedModel(nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.tok_embeddings = nn.Embedding(vocab_size, hidden_size)
+                self.output = nn.Linear(hidden_size, vocab_size, bias=False)
+                self.output.weight = self.tok_embeddings.weight
+
+            def forward(self, ids):
+                return self.output(self.tok_embeddings(ids))
+
+        model = TiedModel()
+        fully_shard(model.tok_embeddings)
+        fully_shard(model.output)
+        fully_shard(model)
+        ids = torch.randint(0, vocab_size, (2, 8), device=device_type.type)
+        with self.assertRaisesRegex(ValueError, "already managed by another"):
+            model(ids)
+
+    @skip_if_lt_x_gpu(2, allow_cpu=True)
+    def test_layer_by_layer_shard_no_false_positive(self):
+        """
+        Test that layer-by-layer materialize+shard+GC does not trigger a
+        false-positive duplicate parameter error from id() reuse.
+        """
+        dim = 16
+        n_blocks = 8
+
+        class SimpleBlock(nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.norm = nn.LayerNorm(dim)
+                self.linear = nn.Linear(dim, dim, bias=False)
+
+            def forward(self, x):
+                return self.linear(self.norm(x))
+
+        class LayerByLayerModel(nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.blocks = nn.ModuleList()
+
+            def forward(self, x):
+                for block in self.blocks:
+                    x = block(x)
+                return x
+
+        model = LayerByLayerModel().to(device_type)
+        for _ in range(n_blocks):
+            block = SimpleBlock().to(device_type)
+            model.blocks.append(block)
+            fully_shard(block)
+            gc.collect()
+
+        fully_shard(model)
+
+        x = torch.randn(2, 4, dim, device=device_type.type)
+        out = model(x)
+        out.sum().backward()
+
 
 class TestFullyShardGradientAccumulation(FSDPTest):
     @property
@@ -1075,7 +1219,7 @@ def test_gradient_accumulation(self):
     def _test_gradient_accumulation(
         self,
         mesh: DeviceMesh,
-        reshard_after_forward: Union[bool, int],
+        reshard_after_forward: bool | int,
         mode: str,
         reshard_after_backward: bool,
         offload_policy: OffloadPolicy,
@@ -1374,6 +1518,222 @@ def _test_2d_mlp_with_nd_mesh(
             self.assertEqual(len(p.placements), 2)
             self.assertEqual(p.device_mesh.mesh_dim_names, ("dp", "tp"))
 
+    @skip_if_lt_x_gpu(8)
+    def test_shard_placement_fn_tp_ep(self):
+        self.run_subtests(
+            {
+                "tp_degree": [1, 2],
+                "dp_replicate": [1, 2],
+                "reshard_non_layer_modules": [False, True, 2],
+            },
+            self._test_shard_placement_fn_tp_ep,
+        )
+
+    def _init_parallel_meshes(self, tp_degree, dp_replicate, ep_degree):
+        """Build the TP, DP, EP, and EFSDP meshes and mesh infos.
+
+        Returns (tp_mesh, dp_mesh, ep_mesh, efsdp_mesh, dp_mesh_info,
+        efsdp_mesh_info) or None if the configuration is not valid.
+        """
+        dp_size = self.world_size // tp_degree
+        dp_shard_size = dp_size // dp_replicate
+        if dp_shard_size < ep_degree:
+            return None
+        efsdp_size = dp_shard_size // ep_degree
+
+        # Build tp_mesh and dp_mesh
+        if dp_replicate > 1 and tp_degree > 1:
+            world_mesh = init_device_mesh(
+                device_type.type,
+                (dp_replicate, dp_shard_size, tp_degree),
+                mesh_dim_names=("dp_replicate", "dp_shard", "tp"),
+            )
+            tp_mesh = world_mesh["tp"]
+            dp_mesh = world_mesh["dp_replicate", "dp_shard"]
+        elif dp_replicate > 1:
+            dp_mesh = init_device_mesh(
+                device_type.type,
+                (dp_replicate, dp_shard_size),
+                mesh_dim_names=("dp_replicate", "dp_shard"),
+            )
+            tp_mesh = None
+        elif tp_degree > 1:
+            world_mesh = init_device_mesh(
+                device_type.type,
+                (dp_size, tp_degree),
+                mesh_dim_names=("dp", "tp"),
+            )
+            tp_mesh = world_mesh["tp"]
+            dp_mesh = world_mesh["dp"]
+        else:
+            world_mesh = init_device_mesh(
+                device_type.type,
+                (self.world_size,),
+                mesh_dim_names=("world",),
+            )
+            tp_mesh = None
+            dp_mesh = world_mesh._unflatten(0, (self.world_size,), ("fsdp",))["fsdp"]
+
+        # Build ep/efsdp meshes and FSDP mesh infos
+        if dp_replicate > 1:
+            dp_shard_mesh = dp_mesh["dp_shard"]
+            sparse_mesh = dp_shard_mesh._unflatten(
+                0, (efsdp_size, ep_degree), ("efsdp", "ep")
+            )
+            full_sparse = dp_mesh._unflatten(
+                1, (efsdp_size, ep_degree), ("efsdp", "ep")
+            )
+            expert_hsdp_mesh = full_sparse["dp_replicate", "efsdp"]
+            efsdp_mesh_info = HSDPMeshInfo(
+                mesh=expert_hsdp_mesh, shard_mesh_dim=1, replicate_mesh_dim=0
+            )
+            dp_mesh_info = HSDPMeshInfo(
+                mesh=dp_mesh, shard_mesh_dim=1, replicate_mesh_dim=0
+            )
+        else:
+            sparse_mesh = dp_mesh._unflatten(
+                0, (efsdp_size, ep_degree), ("efsdp", "ep")
+            )
+            efsdp_mesh_info = FSDPMeshInfo(mesh=sparse_mesh["efsdp"], shard_mesh_dim=0)
+            dp_mesh_info = FSDPMeshInfo(mesh=dp_mesh, shard_mesh_dim=0)
+
+        ep_mesh = sparse_mesh["ep"]
+        efsdp_mesh = sparse_mesh["efsdp"]
+
+        return (
+            tp_mesh,
+            dp_mesh,
+            ep_mesh,
+            efsdp_mesh,
+            dp_mesh_info,
+            efsdp_mesh_info,
+        )
+
+    def _test_shard_placement_fn_tp_ep(
+        self, tp_degree, dp_replicate, reshard_non_layer_modules
+    ):
+        ep_degree = 2
+        result = self._init_parallel_meshes(tp_degree, dp_replicate, ep_degree)
+        if result is None:
+            return
+        (
+            tp_mesh,
+            dp_mesh,
+            ep_mesh,
+            efsdp_mesh,
+            dp_mesh_info,
+            efsdp_mesh_info,
+        ) = result
+        # reshard_root as int must be a factor of every group's
+        # shard mesh size; skip configs where it is not.
+        if isinstance(reshard_non_layer_modules, int) and not isinstance(
+            reshard_non_layer_modules, bool
+        ):
+            for mi in (dp_mesh_info, efsdp_mesh_info):
+                if mi.shard_mesh_size % reshard_non_layer_modules != 0:
+                    return
+        model_args = ModelArgs(
+            n_layers=2,
+            vocab_size=256,
+            max_seq_len=32,
+            dim=64,
+            n_heads=4,
+            dropout_p=0.0,
+            num_experts=8,
+        )
+        torch.manual_seed(42)
+        model = Transformer(model_args)
+        ref_model = copy.deepcopy(model).to(device_type)
+        Transformer.parallelize(
+            model, tp_mesh=tp_mesh, use_seq_parallel=False, ep_mesh=ep_mesh
+        )
+        for block in model.layers:
+            expert_params = set(block.expert_layer.experts.parameters())
+
+            def _shard_placement_fn(
+                param,
+                _expert_params=expert_params,
+            ):
+                if param in _expert_params:
+                    return ShardPlacementResult(
+                        placement=Shard(0),
+                        mesh_info=efsdp_mesh_info,
+                    )
+                return ShardPlacementResult(
+                    placement=Shard(0),
+                    mesh_info=dp_mesh_info,
+                )
+
+            # Blocks always have DTensor expert params (from EP), so int
+            # reshard is not supported; do not pass reshard_after_forward.
+            fully_shard(
+                block,
+                mesh=dp_mesh,
+                shard_placement_fn=_shard_placement_fn,
+            )
+        # Group tok_embeddings, norm, and output together since
+        # output.weight is tied to tok_embeddings.weight
+        # These modules have no DTensor params when tp_degree == 1.
+        # With TP, root params are DTensors and int reshard is unsupported.
+        if tp_mesh is not None:
+            reshard_non_layer_modules = True
+        fully_shard(
+            [model.tok_embeddings, model.norm, model.output],
+            mesh=dp_mesh,
+            reshard_after_forward=reshard_non_layer_modules,
+        )
+        fully_shard(
+            model, mesh=dp_mesh, reshard_after_forward=reshard_non_layer_modules
+        )
+        for (name, param), (_, ref_param) in zip(
+            model.named_parameters(), ref_model.named_parameters()
+        ):
+            full_param = param.full_tensor()
+            self.assertEqual(full_param, ref_param)
+        ref_expert_params = {
+            p for b in ref_model.layers for p in b.expert_layer.experts.parameters()
+        }
+        optim = torch.optim.Adam(model.parameters(), lr=1e-2)
+        ref_optim = torch.optim.Adam(ref_model.parameters(), lr=1e-2)
+        torch.manual_seed(42 + self.rank // tp_degree)
+        inp = torch.randint(
+            0,
+            model_args.vocab_size,
+            (2, model_args.max_seq_len),
+            device=device_type.type,
+        )
+        dp_replicate_group = (
+            dp_mesh["dp_replicate"].get_group() if dp_replicate > 1 else None
+        )
+        for iter_idx in range(5):
+            ref_loss = ref_model(inp).sum()
+            loss = model(inp).sum()
+            ref_loss.backward()
+            loss.backward()
+            for param in ref_model.parameters():
+                if param.grad is None:
+                    continue
+                if param in ref_expert_params:
+                    dist.all_reduce(
+                        param.grad, op=dist.ReduceOp.SUM, group=ep_mesh.get_group()
+                    )
+                    dist.all_reduce(
+                        param.grad, op=dist.ReduceOp.AVG, group=efsdp_mesh.get_group()
+                    )
+                    if dp_replicate_group is not None:
+                        dist.all_reduce(
+                            param.grad,
+                            op=dist.ReduceOp.AVG,
+                            group=dp_replicate_group,
+                        )
+                else:
+                    dist.all_reduce(param.grad, op=dist.ReduceOp.AVG)
+            ref_optim.step()
+            optim.step()
+            ref_optim.zero_grad(set_to_none=(iter_idx % 2 == 0))
+            optim.zero_grad(set_to_none=(iter_idx % 2 == 0))
+            self.assertEqual(ref_loss, loss)
+
 
 class TestFullyShardHSDP3DTraining(FSDPTest):
     @property
@@ -1627,6 +1987,7 @@ def test_share_comm_context(self):
             fully_shard(layer)
             layer._get_fsdp_state()._lazy_init()
         share_comm_ctx(list(model))
+        shared_comm_ctx = model[0]._get_fsdp_state()._comm_ctx
 
         torch.manual_seed(42 + self.rank + 1)
         inp = torch.randn(4, 3, lin_dim, device=device_type.type)
@@ -1673,15 +2034,15 @@ def foreach_reduce_with_assert(
             reduce_scatter_group: dist.ProcessGroup,
             reduce_scatter_stream: torch.Stream,
             reduce_scatter_comm: ReduceScatter,
-            orig_dtype: Optional[torch.dtype],
-            reduce_dtype: Optional[torch.dtype],
+            orig_dtype: torch.dtype | None,
+            reduce_dtype: torch.dtype | None,
             device: torch.device,
-            gradient_divide_factor: Optional[float],
-            all_reduce_group: Optional[dist.ProcessGroup],  # not `None` iff HSDP
+            gradient_divide_factor: float | None,
+            all_reduce_group: dist.ProcessGroup | None,  # not `None` iff HSDP
             all_reduce_stream: torch.Stream,
             all_reduce_grads: bool,
-            partial_reduce_output: Optional[torch.Tensor],  # only used for HSDP
-            all_reduce_hook: Optional[Callable[[torch.Tensor], None]],
+            partial_reduce_output: torch.Tensor | None,  # only used for HSDP
+            all_reduce_hook: Callable[[torch.Tensor], None] | None,
             force_sum_reduction_for_comms: bool = False,
         ):
             nonlocal reduce_scatter_streams
@@ -1716,6 +2077,7 @@ def foreach_reduce_with_assert(
                 dist.all_reduce(param.grad, op=dist.ReduceOp.AVG)
         self.assertEqual(len(all_gather_streams), 1)
         self.assertEqual(len(reduce_scatter_streams), 1)
+        self.assertEqual(len(shared_comm_ctx._last_post_reduce_events), 0)
         check_sharded_parity(self, ref_model, model)
 
 
@@ -1752,7 +2114,7 @@ def _test_train_parity_single_group(
         replicate(ref_model, device_ids=_get_device_ids(self.rank))
         ref_optim = torch.optim.Adam(ref_model.parameters(), lr=1e-2)
 
-        def _shard_placement_fn(param: nn.Parameter) -> Optional[Shard]:
+        def _shard_placement_fn(param: nn.Parameter) -> Shard | None:
             return Shard(param.shape.index(max(param.shape)))
 
         shard_placement_fn = _shard_placement_fn if use_shard_placement_fn else None
diff --git a/test/distributed/_composable/test_checkpoint.py b/test/distributed/_composable/test_checkpoint.py
index 7834328f1e359..c8d967880bbc0 100644
--- a/test/distributed/_composable/test_checkpoint.py
+++ b/test/distributed/_composable/test_checkpoint.py
@@ -85,7 +85,8 @@ def __init__(self, device: torch.device):
         self.w = nn.Parameter(torch.randn((100, 100), device=device))
 
     def forward(self, xs: tuple[torch.Tensor, torch.Tensor]) -> torch.Tensor:
-        assert len(xs) == 2, f"Expects 2 args but got {len(xs)}"
+        if len(xs) != 2:
+            raise AssertionError(f"Expects 2 args but got {len(xs)}")
         x, y = xs
         z = x + y
         z = z @ self.w
diff --git a/test/distributed/_composable/test_composability/test_2d_composability.py b/test/distributed/_composable/test_composability/test_2d_composability.py
index 0d52685d1a750..cbd06e13c1da9 100644
--- a/test/distributed/_composable/test_composability/test_2d_composability.py
+++ b/test/distributed/_composable/test_composability/test_2d_composability.py
@@ -4,7 +4,6 @@
 import functools
 import io
 from copy import deepcopy
-from typing import Optional
 
 import torch
 import torch.distributed as dist
@@ -195,12 +194,15 @@ def _test_train_parity_2d_transformer(self, use_shard_placement_fn: bool):
         )
         model = Transformer.parallelize(model, global_mesh["tp"], use_seq_parallel=True)
 
-        def _shard_placement_fn(param: nn.Parameter) -> Optional[Shard]:
+        def _shard_placement_fn(param: nn.Parameter) -> Shard | None:
             if isinstance(param, DTensor):
                 for placement in param.placements:
                     if isinstance(placement, Shard):
                         shard_dim = param.ndim - 1 - placement.dim
-                        assert shard_dim >= 0, f"{param.shape}"
+                        if not (shard_dim >= 0):
+                            raise AssertionError(
+                                f"Expected shard_dim >= 0, got {shard_dim} for {param.shape}"
+                            )
                         return Shard(shard_dim)
             return Shard(0)
 
@@ -521,14 +523,14 @@ def _check_module(self, m1, m2, check_grad=False):
 
     @skip_if_lt_x_gpu(4)
     def test_2d_ddp_integration_functionality(self) -> None:
-        model, twod_model, dp_pg = self.init_model(self.device_type())
+        model, twod_model, dp_pg = self.init_model(self.device_type)
         optim = torch.optim.Adam(model.parameters(), lr=3e-5)
         twod_optim = torch.optim.Adam(twod_model.parameters(), lr=3e-5)
 
         # Create Input
         input_seed = dist.get_rank(dp_pg)
         torch.manual_seed(input_seed + 1)
-        input = torch.rand(4, 10, device=self.device_type())
+        input = torch.rand(4, 10, device=self.device_type)
 
         output = model(input)
         twod_output = twod_model(input)
@@ -543,7 +545,7 @@ def test_2d_ddp_integration_functionality(self) -> None:
         self._check_module(model, twod_model)
 
         torch.manual_seed(input_seed + 1004)
-        input = torch.rand(16, 10, device=self.device_type())
+        input = torch.rand(16, 10, device=self.device_type)
 
         output = model(input)
         twod_output = twod_model(input)
@@ -575,7 +577,7 @@ def _compare_params(self, m1, m2):
     @skip_if_lt_x_gpu(4)
     def test_2d_fsdp_state_enable_extension(self):
         mesh_2d = init_device_mesh(
-            self.device_type(), (2, self.world_size // 2), mesh_dim_names=("dp", "tp")
+            self.device_type, (2, self.world_size // 2), mesh_dim_names=("dp", "tp")
         )
         model = FSDP(
             SimpleModel().to(device_type),
@@ -596,7 +598,7 @@ def _test_2d_e2e_training(
 
         torch.manual_seed(0)
         mesh_2d = init_device_mesh(
-            self.device_type(), (2, self.world_size // 2), mesh_dim_names=("dp", "tp")
+            self.device_type, (2, self.world_size // 2), mesh_dim_names=("dp", "tp")
         )
         tp_mesh = mesh_2d["tp"]
         dp_mesh = mesh_2d["dp"]
@@ -678,7 +680,7 @@ def test_fsdp_2d_extension(self):
         Test whether _fsdp_extension from FSDPstate has been set correctly.
         """
         mesh_2d = init_device_mesh(
-            self.device_type(), (2, self.world_size // 2), mesh_dim_names=("dp", "tp")
+            self.device_type, (2, self.world_size // 2), mesh_dim_names=("dp", "tp")
         )
         parallelize_plan = {
             "net1": ColwiseParallel(),
@@ -716,7 +718,7 @@ def test_2d_state_dict(self, is_even_sharded_model):
         # Create a model and sharded it with 2D FSDP + TP
         torch.manual_seed(0)
         mesh_2d = init_device_mesh(
-            self.device_type(), (2, self.world_size // 2), mesh_dim_names=("dp", "tp")
+            self.device_type, (2, self.world_size // 2), mesh_dim_names=("dp", "tp")
         )
         tp_mesh = mesh_2d["tp"]
         dp_mesh = mesh_2d["dp"]
@@ -765,7 +767,7 @@ def test_2d_load_state_dict(self, is_even_sharded_model):
 
         torch.manual_seed(0)
         mesh_2d = init_device_mesh(
-            self.device_type(), (2, self.world_size // 2), mesh_dim_names=("dp", "tp")
+            self.device_type, (2, self.world_size // 2), mesh_dim_names=("dp", "tp")
         )
         tp_mesh = mesh_2d["tp"]
         dp_mesh = mesh_2d["dp"]
@@ -830,7 +832,7 @@ def test_2d_optim_state_dict(self, is_even_sharded_model):
         # Create a model and sharded it with 2D FSDP + TP
         torch.manual_seed(0)
         mesh_2d = init_device_mesh(
-            self.device_type(), (2, self.world_size // 2), mesh_dim_names=("dp", "tp")
+            self.device_type, (2, self.world_size // 2), mesh_dim_names=("dp", "tp")
         )
         parallelize_plan = {
             "net1": ColwiseParallel(),
diff --git a/test/distributed/_composable/test_composability/test_pp_composability.py b/test/distributed/_composable/test_composability/test_pp_composability.py
index 9ac81b0c7ee4e..67ca31f7cd1cb 100644
--- a/test/distributed/_composable/test_composability/test_pp_composability.py
+++ b/test/distributed/_composable/test_composability/test_pp_composability.py
@@ -615,7 +615,10 @@ def check_gradient_parity(
             pipeline_model_parameter_dict, ref_model_parameter_dict
         ):
             for parameter in pipeline_model_parameter_dict:
-                assert parameter in ref_model_parameter_dict
+                if parameter not in ref_model_parameter_dict:
+                    raise AssertionError(
+                        f"Parameter {parameter} not found in ref_model_parameter_dict"
+                    )
 
                 pipeline_parameter = pipeline_model_parameter_dict[parameter]
                 if pipeline_parameter.grad is not None:
@@ -629,7 +632,10 @@ def check_gradient_parity(
                             atol=1e-5,
                         )
                     else:
-                        assert pipeline_parameter.grad is None
+                        if pipeline_parameter.grad is not None:
+                            raise AssertionError(
+                                f"Expected pipeline_parameter.grad to be None for {parameter}"
+                            )
 
         pipeline_model_parameter_dict = {}
 
diff --git a/test/distributed/_composable/test_replicate_mixed_precision.py b/test/distributed/_composable/test_replicate_mixed_precision.py
index 8a03bdb018811..0d2cbbd14e2e7 100644
--- a/test/distributed/_composable/test_replicate_mixed_precision.py
+++ b/test/distributed/_composable/test_replicate_mixed_precision.py
@@ -3,7 +3,6 @@
 import copy
 import dataclasses
 import functools
-from typing import Optional
 
 import torch
 import torch.distributed as dist
@@ -43,8 +42,8 @@ class TestReplicateMixedPrecisionTraining(FSDPTestContinuous):
 
     def _init_models_and_optims(
         self,
-        param_dtype: Optional[torch.dtype],
-        reduce_dtype: Optional[torch.dtype],
+        param_dtype: torch.dtype | None,
+        reduce_dtype: torch.dtype | None,
     ):
         torch.manual_seed(42)
         model = nn.Sequential(*[MLP(16, torch.device("cpu")) for _ in range(3)])
diff --git a/test/distributed/_composable/test_replicate_training.py b/test/distributed/_composable/test_replicate_training.py
index 3dc908a8b1afe..24e070fe97adc 100644
--- a/test/distributed/_composable/test_replicate_training.py
+++ b/test/distributed/_composable/test_replicate_training.py
@@ -7,7 +7,6 @@
 import unittest
 from collections import defaultdict
 from collections.abc import Iterable
-from typing import Union
 
 import torch
 import torch.distributed as dist
@@ -78,13 +77,12 @@ class ParamlessModule(nn.Module):
             def forward(self, x: torch.Tensor, ys: tuple[torch.Tensor, ...]):
                 # Check that Replicate moved the inputs to GPU, including recursing
                 # into the tuple data structure
-                assert x.device == device, f"Expects {device} but got {x.device}"
-                assert ys[0].device == device, (
-                    f"Expects {device} but got {ys[0].device}"
-                )
-                assert ys[1].device == device, (
-                    f"Expects {device} but got {ys[1].device}"
-                )
+                if not (x.device == device):
+                    raise AssertionError(f"Expects {device} but got {x.device}")
+                if not (ys[0].device == device):
+                    raise AssertionError(f"Expects {device} but got {ys[0].device}")
+                if not (ys[1].device == device):
+                    raise AssertionError(f"Expects {device} but got {ys[1].device}")
                 y = ys[0] + ys[1]
                 return x + y + 1
 
@@ -369,7 +367,8 @@ def _test_train_parity_multi_group(
             in (2, 3)
         ):
             return
-        assert test_device_type in ("cuda", "hpu", "xpu", "cpu"), f"{test_device_type}"
+        if test_device_type not in ("cuda", "hpu", "xpu", "cpu"):
+            raise AssertionError(f"Unexpected device type: {test_device_type}")
         torch.manual_seed(42)
         vocab_size = 1024
         model_args = ModelArgs(
@@ -677,7 +676,8 @@ def _test_train_parity_with_activation_checkpointing(
         module_grouping: str,
         test_device_type: str,
     ):
-        assert checkpoint_impl in ("composable", "utils", "wrapper")
+        if checkpoint_impl not in ("composable", "utils", "wrapper"):
+            raise AssertionError(f"Unexpected checkpoint_impl: {checkpoint_impl}")
         testing_compile = (
             replicate is not torch.distributed._composable.replicate_with_fsdp
         )
@@ -724,7 +724,10 @@ def _test_train_parity_with_activation_checkpointing(
             "mesh": device_mesh,
         }
         if module_grouping == "mem_eff":
-            assert model_args.n_layers == 3
+            if not (model_args.n_layers == 3):
+                raise AssertionError(
+                    f"Expected n_layers == 3, got {model_args.n_layers}"
+                )
             replicate(model.layers[0], **fsdp_kwargs)
             replicate([model.layers[1], model.layers[2]], **fsdp_kwargs)
             replicate([model.tok_embeddings, model.pos_embeddings], **fsdp_kwargs)
@@ -871,7 +874,7 @@ def test_gradient_accumulation(self):
     def _test_gradient_accumulation(
         self,
         mesh: DeviceMesh,
-        reshard_after_forward: Union[bool, int],
+        reshard_after_forward: bool | int,
         mode: str,
         reshard_after_backward: bool,
         offload_policy: OffloadPolicy,
diff --git a/test/distributed/_composable/test_replicate_with_compiler.py b/test/distributed/_composable/test_replicate_with_compiler.py
index a71b197d4349d..5936a729cbe5b 100644
--- a/test/distributed/_composable/test_replicate_with_compiler.py
+++ b/test/distributed/_composable/test_replicate_with_compiler.py
@@ -5,7 +5,6 @@
 import unittest
 from collections.abc import Callable
 from copy import deepcopy
-from typing import Optional, Union
 
 import torch
 import torch.distributed as dist
@@ -90,11 +89,11 @@ def _test_compile(
         self,
         *,
         no_sync: bool,
-        setup_func: Optional[Callable] = None,
+        setup_func: Callable | None = None,
         no_inductor: bool = False,
         no_compile_forward: bool = False,
         checkpoint: bool = False,
-        device: Union[str, torch.device],
+        device: str | torch.device,
     ):
         self.create_pg(device)
         torch._dynamo.config.optimize_ddp = "python_reducer"
diff --git a/test/distributed/_composable/test_replicate_with_fsdp.py b/test/distributed/_composable/test_replicate_with_fsdp.py
index c925b64345aff..6236c93dc6cc4 100644
--- a/test/distributed/_composable/test_replicate_with_fsdp.py
+++ b/test/distributed/_composable/test_replicate_with_fsdp.py
@@ -57,9 +57,9 @@ def device_type(cls) -> str:
     @classmethod
     def _init_pg(cls, rank, world_size, rdvz_file):
         if torch.accelerator.is_available():
+            if torch.accelerator.device_count() < world_size:
+                sys.exit(TEST_SKIPS[f"multi-gpu-{world_size}"].exit_code)
             torch.accelerator.set_device_index(rank)
-        if torch.accelerator.device_count() < world_size:
-            sys.exit(TEST_SKIPS[f"multi-gpu-{world_size}"].exit_code)
         super()._init_pg(rank, world_size, rdvz_file)
 
     def init_replicate_tp_mesh(self) -> DeviceMesh:
diff --git a/test/distributed/_pycute/test_complement.py b/test/distributed/_pycute/test_complement.py
index e54364732f049..77c8f50c9cc2f 100644
--- a/test/distributed/_pycute/test_complement.py
+++ b/test/distributed/_pycute/test_complement.py
@@ -55,7 +55,12 @@ def helper_test_complement(self, layout):
         # Post-condition: test disjointedness of the codomains
         for a in range(size(layout)):
             for b in range(size(layoutR)):
-                assert (layout(a) != layoutR(b)) or (layout(a) == 0 and layoutR(b) == 0)
+                if not (
+                    (layout(a) != layoutR(b)) or (layout(a) == 0 and layoutR(b) == 0)
+                ):
+                    raise AssertionError(
+                        f"Invariant violated at a={a}, b={b}: layout(a)={layout(a)}, layoutR(b)={layoutR(b)}"
+                    )
 
     def test_complement(self):
         test = Layout(1, 0)
diff --git a/test/distributed/_shard/sharded_tensor/test_sharded_tensor.py b/test/distributed/_shard/sharded_tensor/test_sharded_tensor.py
index b39b3075060f6..3d9183bf63238 100644
--- a/test/distributed/_shard/sharded_tensor/test_sharded_tensor.py
+++ b/test/distributed/_shard/sharded_tensor/test_sharded_tensor.py
@@ -2309,7 +2309,8 @@ def _generate_st_from_chunk_local_tensor(self, st_size, sharding_spec):
 
         # TODO: figure out what the API should behave when some rank have no shard
         # see https://github.com/pytorch/pytorch/issues/73133
-        assert local_tensor is not None
+        if local_tensor is None:
+            raise AssertionError("Expected local_tensor to not be None")
         st = ShardedTensor._init_from_local_tensor(
             local_tensor,
             sharding_spec,
diff --git a/test/distributed/_shard/sharding_spec/test_sharding_spec.py b/test/distributed/_shard/sharding_spec/test_sharding_spec.py
index 37ad69075068f..fe14f815749b1 100644
--- a/test/distributed/_shard/sharding_spec/test_sharding_spec.py
+++ b/test/distributed/_shard/sharding_spec/test_sharding_spec.py
@@ -1,7 +1,6 @@
 # Owner(s): ["oncall: distributed"]
 import copy
 from dataclasses import dataclass
-from typing import Union
 
 import torch
 from torch.distributed._shard import _shard_tensor, sharded_tensor
@@ -558,7 +557,7 @@ def test_check_overlapping(self):
 @dataclass
 class GridShardingSpec(ShardingSpec):
     grid_size: int
-    placements: list[Union[torch.distributed._remote_device, str]]
+    placements: list[torch.distributed._remote_device | str]
 
     def __post_init__(self):
         for i, remote_device in enumerate(self.placements):
@@ -571,17 +570,23 @@ def build_metadata(
         tensor_properties: TensorProperties,
     ) -> ShardedTensorMetadata:
         tensor_num_dim = len(tensor_sizes)
-        assert tensor_num_dim == 2, "only support 2-dim tensor for grid sharding"
+        if tensor_num_dim != 2:
+            raise AssertionError("only support 2-dim tensor for grid sharding")
         shards_metadata = []
 
         def chunk_num(dim_size, grid_size):
-            assert dim_size % grid_size == 0, "only support dim_size mod grid_size == 0"
+            if dim_size % grid_size != 0:
+                raise AssertionError("only support dim_size mod grid_size == 0")
             return dim_size // grid_size
 
         row_chunks = chunk_num(tensor_sizes[0], self.grid_size)
         col_chunks = chunk_num(tensor_sizes[1], self.grid_size)
 
-        assert row_chunks * col_chunks == len(self.placements)
+        if row_chunks * col_chunks != len(self.placements):
+            raise AssertionError(
+                f"Expected row_chunks * col_chunks == len(self.placements), "
+                f"got {row_chunks * col_chunks} vs {len(self.placements)}"
+            )
         for row_idx in range(row_chunks):
             for col_idx in range(col_chunks):
                 shards_metadata.append(
diff --git a/test/distributed/_shard/test_sharder.py b/test/distributed/_shard/test_sharder.py
index 869e5661ddb26..27b79c55406d8 100644
--- a/test/distributed/_shard/test_sharder.py
+++ b/test/distributed/_shard/test_sharder.py
@@ -55,7 +55,10 @@ def __init__(self, ebc, split_idx, specs):
         # create embedding bags base on the spec
         self.embedding_bags: nn.ModuleDict = nn.ModuleDict()
 
-        assert self.split_idx < ebc.num_bags
+        if not (self.split_idx < ebc.num_bags):
+            raise AssertionError(
+                f"Expected split_idx < num_bags, got {self.split_idx} vs {ebc.num_bags}"
+            )
         for i in range(ebc.num_bags):
             bag_key = f"embedding_bag_{i}"
             if i < self.split_idx:
diff --git a/test/distributed/_tools/test_fsdp2_mem_tracker.py b/test/distributed/_tools/test_fsdp2_mem_tracker.py
index c793bc98b8918..7e513ef186ac8 100644
--- a/test/distributed/_tools/test_fsdp2_mem_tracker.py
+++ b/test/distributed/_tools/test_fsdp2_mem_tracker.py
@@ -1,7 +1,6 @@
 # Owner(s): ["module: fsdp"]
 import functools
 import gc
-from typing import Union
 
 import torch
 import torch.nn as nn
@@ -73,7 +72,7 @@ def test_tracker_multi_group_eager(self):
 
     def _test_tracker_multi_group(
         self,
-        reshard_after_forward: Union[bool, int],
+        reshard_after_forward: bool | int,
         offload_policy: OffloadPolicy,
         mp_policy: MixedPrecisionPolicy,
     ):
@@ -241,9 +240,12 @@ def test_tracker_with_activation_checkpointing(self):
         )
 
     def _test_tracker_with_activation_checkpointing(
-        self, reshard_after_forward: Union[bool, int], checkpoint_impl: str
+        self, reshard_after_forward: bool | int, checkpoint_impl: str
     ):
-        assert checkpoint_impl in ("composable", "wrapper")
+        if checkpoint_impl not in ("composable", "wrapper"):
+            raise AssertionError(
+                f"Expected checkpoint_impl in ('composable', 'wrapper'), got {checkpoint_impl}"
+            )
         debug = False
         dev = torch.device(torch.accelerator.current_device_index())
         _init_cublas_workspace(dev)
diff --git a/test/distributed/_tools/test_mod_tracker.py b/test/distributed/_tools/test_mod_tracker.py
index b4de403b503a3..646689752f605 100644
--- a/test/distributed/_tools/test_mod_tracker.py
+++ b/test/distributed/_tools/test_mod_tracker.py
@@ -154,7 +154,8 @@ def forward(self, x):
                         )
                     else:
                         x = block(x)
-                    assert x is not None
+                    if x is None:
+                        raise AssertionError("Expected x to not be None")
                     x = torch.nn.functional.relu(x)
                 return x
 
diff --git a/test/distributed/_tools/test_runtime_estimator.py b/test/distributed/_tools/test_runtime_estimator.py
index f510a8afd9f3c..62ecbe7cfbf77 100644
--- a/test/distributed/_tools/test_runtime_estimator.py
+++ b/test/distributed/_tools/test_runtime_estimator.py
@@ -2,7 +2,7 @@
 import unittest
 from collections.abc import Callable
 from dataclasses import dataclass
-from typing import Any, cast, Union
+from typing import Any, cast
 
 import torch
 from torch import nn, optim
@@ -105,7 +105,7 @@ def _runtime_estimate(
     def _init_model_and_args(
         self,
         model_type: str,
-        model_args: Union[ConvArgs, ModelArgs],
+        model_args: ConvArgs | ModelArgs,
         bsz: int,
     ) -> tuple[nn.Module, optim.Optimizer, torch.Tensor]:
         dev = torch.cuda.current_device()
diff --git a/test/distributed/_tools/test_sac_ilp.py b/test/distributed/_tools/test_sac_ilp.py
index 6517f9a4ce087..0bba3c67506ab 100644
--- a/test/distributed/_tools/test_sac_ilp.py
+++ b/test/distributed/_tools/test_sac_ilp.py
@@ -85,7 +85,8 @@ def _run_and_get_memTracker(
                 optimizer.zero_grad()
                 if iter_idx == 0:
                     mt.reset_mod_stats()
-        assert last_snapshot is not None
+        if last_snapshot is None:
+            raise AssertionError("Expected last_snapshot to not be None")
         for mod_stats in mem_tracker.memory_tracking.values():
             # postprocessing due to the fact that for ModTracker, the post backward hook
             # is not being called for modules whose inputs don't require gradients
diff --git a/test/distributed/algorithms/ddp_comm_hooks/test_ddp_hooks.py b/test/distributed/algorithms/ddp_comm_hooks/test_ddp_hooks.py
index ee800f73b29d5..961e40556c023 100644
--- a/test/distributed/algorithms/ddp_comm_hooks/test_ddp_hooks.py
+++ b/test/distributed/algorithms/ddp_comm_hooks/test_ddp_hooks.py
@@ -220,8 +220,9 @@ def hook(flags, bucket):
 
 
 if __name__ == "__main__":
-    assert not torch.cuda._initialized, (
-        "test_distributed must not have initialized CUDA context on main process"
-    )
+    if torch.cuda._initialized:
+        raise AssertionError(
+            "test_distributed must not have initialized CUDA context on main process"
+        )
 
     run_tests()
diff --git a/test/distributed/algorithms/test_join.py b/test/distributed/algorithms/test_join.py
index 8fd613a47d775..e68b5a1682d6f 100644
--- a/test/distributed/algorithms/test_join.py
+++ b/test/distributed/algorithms/test_join.py
@@ -3,7 +3,7 @@
 import contextlib
 import os
 import sys
-from typing import Any, Optional
+from typing import Any
 
 import torch
 import torch.distributed as dist
@@ -130,7 +130,8 @@ def find_common_rank(self, rank, to_consider):
         common_rank = torch.tensor([rank if to_consider else -1], device=self.device)
         dist.all_reduce(common_rank, op=dist.ReduceOp.MAX, group=self.process_group)
         common_rank = common_rank.item()
-        assert common_rank >= 0
+        if not (common_rank >= 0):
+            raise AssertionError(f"Expected common_rank >= 0, got {common_rank}")
         return common_rank
 
 
@@ -207,7 +208,7 @@ def _test_join_base(
         throw_on_early_termination: bool,
         num_allreduces: int,
         run_post_hooks: bool,
-        expected_total: Optional[int] = None,
+        expected_total: int | None = None,
     ):
         r"""
         Skeleton for all :class:`Join` tests.
diff --git a/test/distributed/checkpoint/_experimental/test_checkpoint_process.py b/test/distributed/checkpoint/_experimental/test_checkpoint_process.py
index 161dd1a80c3e1..8ae63b7cdca6c 100644
--- a/test/distributed/checkpoint/_experimental/test_checkpoint_process.py
+++ b/test/distributed/checkpoint/_experimental/test_checkpoint_process.py
@@ -28,11 +28,12 @@ def subprocess_init_fn(name: str, parent_pid: int) -> None:
 
     This is similar to the subprocess_init_routine in checkpointing_test.py.
     """
-    assert name == "test-checkpointer", f"Unexpected subprocess name: {name}"
-    assert os.getpid() != parent_pid, "This was supposed to run in a different process"
-    assert os.getppid() == parent_pid, (
-        "This was supposed to run as a child to main process"
-    )
+    if name != "test-checkpointer":
+        raise AssertionError(f"Unexpected subprocess name: {name}")
+    if os.getpid() == parent_pid:
+        raise AssertionError("This was supposed to run in a different process")
+    if os.getppid() != parent_pid:
+        raise AssertionError("This was supposed to run as a child to main process")
 
 
 def failing_subprocess_init_fn(name: str, parent_pid: int) -> None:
@@ -90,17 +91,19 @@ def write(self, state_dict, path, **__):
             if "shared_tensor" in state_dict:
                 shared_tensor = state_dict["shared_tensor"]
                 # Critical assertion: shared tensor should remain in shared memory in subprocess
-                assert shared_tensor.is_shared(), (
-                    "Shared tensor should be in shared memory in subprocess"
-                )
+                if not shared_tensor.is_shared():
+                    raise AssertionError(
+                        "Shared tensor should be in shared memory in subprocess"
+                    )
 
                 shared_tensor[0] = 42.0
 
             if "regular_tensor" in state_dict:
                 # Note: ForkingPickler moves regular tensors to shared memory during IPC - this is acceptable
-                assert state_dict["regular_tensor"].is_shared(), (
-                    "Regular tensor should also be in shared memory in subprocess"
-                )
+                if not state_dict["regular_tensor"].is_shared():
+                    raise AssertionError(
+                        "Regular tensor should also be in shared memory in subprocess"
+                    )
 
             return None
 
diff --git a/test/distributed/checkpoint/_experimental/test_checkpoint_writer.py b/test/distributed/checkpoint/_experimental/test_checkpoint_writer.py
index 959f1c9e7572d..b6262291872ff 100644
--- a/test/distributed/checkpoint/_experimental/test_checkpoint_writer.py
+++ b/test/distributed/checkpoint/_experimental/test_checkpoint_writer.py
@@ -3,7 +3,7 @@
 import os
 import shutil
 import tempfile
-from typing import Any, Optional
+from typing import Any
 from unittest.mock import MagicMock
 
 import torch
@@ -22,10 +22,10 @@ class MockWriterHook(WriterHook):
     def __init__(self):
         self.pre_commit_called = False
         self.commit_called = False
-        self.pre_commit_path: Optional[str] = None
-        self.commit_path: Optional[str] = None
-        self.pre_commit_kwargs: Optional[dict[str, Any]] = None
-        self.commit_kwargs: Optional[dict[str, Any]] = None
+        self.pre_commit_path: str | None = None
+        self.commit_path: str | None = None
+        self.pre_commit_kwargs: dict[str, Any] | None = None
+        self.commit_kwargs: dict[str, Any] | None = None
 
     def pre_commit(self, path: str, **kwargs: Any):
         self.pre_commit_called = True
diff --git a/test/distributed/checkpoint/_experimental/test_checkpointer.py b/test/distributed/checkpoint/_experimental/test_checkpointer.py
index fbd19ff9eafad..ec70035392a8d 100644
--- a/test/distributed/checkpoint/_experimental/test_checkpointer.py
+++ b/test/distributed/checkpoint/_experimental/test_checkpointer.py
@@ -33,11 +33,12 @@
 
 def subprocess_init_fn(name: str, parent_pid: int) -> None:
     """Initialize the subprocess for async checkpointer tests."""
-    assert name == "test-async-checkpointer", f"Unexpected subprocess name: {name}"
-    assert os.getpid() != parent_pid, "This was supposed to run in a different process"
-    assert os.getppid() == parent_pid, (
-        "This was supposed to run as a child to main process"
-    )
+    if name != "test-async-checkpointer":
+        raise AssertionError(f"Unexpected subprocess name: {name}")
+    if os.getpid() == parent_pid:
+        raise AssertionError("This was supposed to run in a different process")
+    if os.getppid() != parent_pid:
+        raise AssertionError("This was supposed to run as a child to main process")
 
 
 def ckpt_writer_init_fn(**kwargs) -> CheckpointWriter:
diff --git a/test/distributed/checkpoint/_experimental/test_staging.py b/test/distributed/checkpoint/_experimental/test_staging.py
index c9be4fe43f49d..739d098a899e4 100644
--- a/test/distributed/checkpoint/_experimental/test_staging.py
+++ b/test/distributed/checkpoint/_experimental/test_staging.py
@@ -159,7 +159,8 @@ def test_cuda_tensors_staging(self) -> None:
         stager = DefaultStager(options)
 
         staged_dict = stager.stage(cuda_state_dict)
-        assert isinstance(staged_dict, dict)
+        if not isinstance(staged_dict, dict):
+            raise AssertionError(f"Expected dict, got {type(staged_dict)}")
 
         # Verify tensors are staged (should be moved to CPU)
         self.assertIn("cuda_tensor", staged_dict)
diff --git a/test/distributed/checkpoint/e2e/test_e2e_save_and_load.py b/test/distributed/checkpoint/e2e/test_e2e_save_and_load.py
index 59cd11fa7ea48..c3f16b9473f5f 100644
--- a/test/distributed/checkpoint/e2e/test_e2e_save_and_load.py
+++ b/test/distributed/checkpoint/e2e/test_e2e_save_and_load.py
@@ -160,7 +160,8 @@ def backend(self):
     def _create_model(self, compile, model_type, state_dict_options=None):
         dummy_model = TestDummyModel().to(self.device_type)
 
-        assert model_type in ModelType, f"{model_type} is not supported."
+        if model_type not in ModelType:
+            raise AssertionError(f"{model_type} is not supported.")
         if model_type == ModelType.FSDP:
             device_mesh = init_device_mesh(self.device_type, (self.world_size,))
             model = FSDP(
@@ -299,7 +300,10 @@ def _run_e2e_test(
             if isinstance(async_save_response_or_future, Future):
                 save_future = async_save_response_or_future
             else:
-                assert isinstance(async_save_response_or_future, AsyncSaveResponse)
+                if not isinstance(async_save_response_or_future, AsyncSaveResponse):
+                    raise AssertionError(
+                        f"Expected AsyncSaveResponse, got {type(async_save_response_or_future)}"
+                    )
                 save_future = async_save_response_or_future.upload_completion
             # wait for the future to complete
             t = time.monotonic()
diff --git a/test/distributed/checkpoint/test_checkpoint.py b/test/distributed/checkpoint/test_checkpoint.py
index 0bc5bf69f2a50..c196ee43381da 100644
--- a/test/distributed/checkpoint/test_checkpoint.py
+++ b/test/distributed/checkpoint/test_checkpoint.py
@@ -2,7 +2,7 @@
 
 import os
 import sys
-from typing import Any, cast, Optional, Union
+from typing import Any, cast
 
 import torch
 import torch.distributed as dist
@@ -59,8 +59,8 @@ def __init__(self) -> None:
         super().__init__()
         self.sharded: ShardedTensor = sharded_tensor.zeros(self.spec(), 4, 4)
         self.regular = torch.nn.Parameter(torch.ones(4, 4))
-        self.extra_sharded: Optional[ShardedTensor] = None
-        self.extra_param: Optional[torch.nn.Parameter] = None
+        self.extra_sharded: ShardedTensor | None = None
+        self.extra_param: torch.nn.Parameter | None = None
         self._register_state_dict_hook(state_dict_hook)
 
     def spec(self) -> ChunkShardingSpec:
@@ -173,7 +173,7 @@ class FaultyStorageWriter(TestStorageBase, StorageWriter):
     def __init__(self, fail_conf):
         super().__init__(fail_conf)
 
-    def reset(self, checkpoint_id: Union[str, os.PathLike, None] = None) -> None:
+    def reset(self, checkpoint_id: str | os.PathLike | None = None) -> None:
         return
 
     def set_up_storage_writer(
@@ -199,7 +199,7 @@ def finish(self, metadata: Metadata, results: list[list[WriteResult]]) -> None:
         self._fail_rank("fail_finish")
 
     @classmethod
-    def validate_checkpoint_id(cls, checkpoint_id: Union[str, os.PathLike]) -> bool:
+    def validate_checkpoint_id(cls, checkpoint_id: str | os.PathLike) -> bool:
         return True
 
 
@@ -208,7 +208,7 @@ def __init__(self, metadata, fail_conf):
         super().__init__(fail_conf)
         self.metadata = metadata
 
-    def reset(self, checkpoint_id: Union[str, os.PathLike, None] = None) -> None:
+    def reset(self, checkpoint_id: str | os.PathLike | None = None) -> None:
         return
 
     def set_up_storage_reader(self, metadata: Metadata, is_coordinator: bool) -> None:
@@ -231,7 +231,7 @@ def read_metadata(self) -> Metadata:
         return self.metadata
 
     @classmethod
-    def validate_checkpoint_id(cls, checkpoint_id: Union[str, os.PathLike]) -> bool:
+    def validate_checkpoint_id(cls, checkpoint_id: str | os.PathLike) -> bool:
         return True
 
 
diff --git a/test/distributed/checkpoint/test_dtensor_checkpoint.py b/test/distributed/checkpoint/test_dtensor_checkpoint.py
index e34d14b967487..d8ae3d1f42747 100644
--- a/test/distributed/checkpoint/test_dtensor_checkpoint.py
+++ b/test/distributed/checkpoint/test_dtensor_checkpoint.py
@@ -1,5 +1,4 @@
 # Owner(s): ["oncall: distributed"]
-from typing import Union
 
 import torch
 import torch.distributed as dist
@@ -58,15 +57,13 @@ def extra_state_tensor(self) -> torch.Tensor:
     def extra_state_tensor(self, new_extra_state_tensor: torch.Tensor) -> None:
         self._extra_state_tensor = new_extra_state_tensor
 
-    def get_extra_state(self) -> dict[str, Union[int, torch._tensor.Tensor]]:
+    def get_extra_state(self) -> dict[str, int | torch._tensor.Tensor]:
         return {
             "extra_state": self._extra_state,
             "extra_state_tensor": self._extra_state_tensor,
         }
 
-    def set_extra_state(
-        self, state: dict[str, Union[int, torch._tensor.Tensor]]
-    ) -> None:
+    def set_extra_state(self, state: dict[str, int | torch._tensor.Tensor]) -> None:
         self._extra_state = state["extra_state"]  # pyre-ignore[8]
         self._extra_state_tensor = state["extra_state_tensor"]  # pyre-ignore[8]
 
diff --git a/test/distributed/checkpoint/test_dtensor_resharding.py b/test/distributed/checkpoint/test_dtensor_resharding.py
index 233fb3e7e0f03..a8974ac27edcf 100644
--- a/test/distributed/checkpoint/test_dtensor_resharding.py
+++ b/test/distributed/checkpoint/test_dtensor_resharding.py
@@ -416,7 +416,11 @@ def __create_write_items__(self, fqn: str, object: Any) -> list[WriteItem]:
         ]
 
     def __get_tensor_shard__(self, index: MetadataIndex) -> torch.Tensor:
-        assert self._fqn == index.fqn and self._shard_offsets == index.offset
+        if not (self._fqn == index.fqn and self._shard_offsets == index.offset):
+            raise AssertionError(
+                f"fqn/offset mismatch: {self._fqn} vs {index.fqn}, "
+                f"{self._shard_offsets} vs {index.offset}"
+            )
         return self._local_tensor
 
     def __repr__(self):
@@ -495,7 +499,10 @@ def test_uneven_reshard_with_checkpointable_api(self) -> None:
             state_dict=state_dict_to_load,
             storage_reader=dist_cp.FileSystemReader(self.temp_dir),
         )
-        assert torch.equal(loading_local_tensor, expected_loaded_local_val_tensor)
+        if not torch.equal(loading_local_tensor, expected_loaded_local_val_tensor):
+            raise AssertionError(
+                "Expected loading_local_tensor to equal expected_loaded_local_val_tensor"
+            )
 
     @with_comms
     @with_temp_dir
@@ -585,7 +592,10 @@ def test_uneven_reshard_with_dtensor_shards_wrapper_api(self) -> None:
         logger.info(
             f"[{my_rank}] loaded_shards_wrapper : {loading_local_shard_wrapper}"  # noqa: G004
         )
-        assert torch.equal(loading_local_tensor, expected_loaded_local_val_tensor)
+        if not torch.equal(loading_local_tensor, expected_loaded_local_val_tensor):
+            raise AssertionError(
+                "Expected loading_local_tensor to equal expected_loaded_local_val_tensor"
+            )
         dist.barrier()
 
 
diff --git a/test/distributed/checkpoint/test_fsspec.py b/test/distributed/checkpoint/test_fsspec.py
index 530ffae852715..ca191bf8bb90b 100644
--- a/test/distributed/checkpoint/test_fsspec.py
+++ b/test/distributed/checkpoint/test_fsspec.py
@@ -4,7 +4,7 @@
 import tempfile
 from collections.abc import Callable
 from functools import wraps
-from typing import Any, Optional
+from typing import Any
 
 import torch
 import torch.distributed as dist
@@ -35,12 +35,13 @@
 
 
 def with_temp_dir(
-    func: Optional[Callable] = None,
-) -> Optional[Callable]:
+    func: Callable | None = None,
+) -> Callable | None:
     """
     Wrapper to initialize temp directory for distributed checkpoint.
     """
-    assert func is not None
+    if func is None:
+        raise AssertionError("Expected func to not be None")
 
     @wraps(func)
     def wrapper(self, *args: tuple[object], **kwargs: dict[str, Any]) -> None:
diff --git a/test/distributed/checkpoint/test_pg_transport.py b/test/distributed/checkpoint/test_pg_transport.py
index 6036c7c11909c..0a9787a5ba64e 100644
--- a/test/distributed/checkpoint/test_pg_transport.py
+++ b/test/distributed/checkpoint/test_pg_transport.py
@@ -3,7 +3,6 @@
 import logging
 import unittest
 from datetime import timedelta
-from typing import Optional
 from unittest.mock import MagicMock, patch
 
 import torch
@@ -211,7 +210,7 @@ class PgTransportCPU(MultiProcContinuousTest):
     timeout: timedelta = timedelta(seconds=20)
 
     @classmethod
-    def backend_str(cls) -> Optional[str]:
+    def backend_str(cls) -> str | None:
         return "gloo"
 
     @classmethod
@@ -232,12 +231,13 @@ def test_pg_transport_with_sharded_tensor(self) -> None:
         _test_pg_transport_with_sharded_tensor(self, self.device)
 
 
+@skip_but_pass_in_sandcastle_if(not at_least_x_gpu(2), "test requires 2+ accelerators")
 class PgTransportGPU(MultiProcContinuousTest):
     world_size = 2
     timeout: timedelta = timedelta(seconds=20)
 
     @classmethod
-    def backend_str(cls) -> Optional[str]:
+    def backend_str(cls) -> str | None:
         return dist.get_default_backend_for_device(cls.device_type())
 
     @property
diff --git a/test/distributed/checkpoint/test_planner.py b/test/distributed/checkpoint/test_planner.py
index a8620c383f2f9..b6b6aa2c6f980 100644
--- a/test/distributed/checkpoint/test_planner.py
+++ b/test/distributed/checkpoint/test_planner.py
@@ -1,6 +1,7 @@
 # Owner(s): ["oncall: distributed"]
 
 import sys
+from unittest.mock import patch
 
 import torch
 import torch.distributed.checkpoint as dcp
@@ -155,7 +156,13 @@ def test_local_plan_with_caching(self):
         # First iteration, should create a new plan
         first_plan = planner.create_local_plan()
 
-        # Validate that the plan has been cached
+        # Plan should be pending, not yet in the class-level cache
+        self.assertNotIn(planner._cached_plans_key, SavePlanner._cached_save_plan)
+
+        # Promote the pending plan by calling finish_plan
+        planner.finish_plan(first_plan)
+
+        # Validate that the plan has been cached after finish_plan
         cached_plan = SavePlanner._cached_save_plan[planner._cached_plans_key]
         self.assertEqual(first_plan, cached_plan)
 
@@ -166,6 +173,11 @@ def test_local_plan_with_caching(self):
         self.assertIsNone(second_plan.planner_data)
         self.assertIsNone(second_plan.storage_data)
 
+        # Clean up class-level caches
+        key = planner._cached_plans_key
+        SavePlanner._cached_save_plan.pop(key, None)
+        SavePlanner._cached_final_save_plan.pop(key, None)
+
     def test_global_plan(self):
         def create_data(rank):
             with with_dist(rank=rank, world_size=4):
@@ -313,6 +325,14 @@ def create_data_v2(rank):
         ]
         self.assertEqual(cached_global_plan, tensor_plan)
 
+        # Clean up class-level caches
+        key = planner._cached_plans_key
+        SavePlanner._cached_save_plan.pop(key, None)
+        SavePlanner._cached_final_save_plan.pop(key, None)
+        SavePlanner._cached_all_plans.pop(key, None)
+        SavePlanner._cached_global_plan.pop(key, None)
+        SavePlanner._cached_metadata.pop(key, None)
+
     def test_finish_plan_with_caching(self):
         planner = DefaultSavePlanner(enable_plan_caching=True)
         tensor = torch.rand(10)
@@ -334,6 +354,67 @@ def test_finish_plan_with_caching(self):
         second_finished_plan = planner.finish_plan(SavePlan([], usable=False))
         self.assertEqual(second_finished_plan, first_finished_plan)
 
+        # Clean up class-level caches
+        key = planner._cached_plans_key
+        SavePlanner._cached_save_plan.pop(key, None)
+        SavePlanner._cached_final_save_plan.pop(key, None)
+
+    def test_caching_after_validation_failure(self):
+        def create_data(rank):
+            with with_dist(rank=rank, world_size=4):
+                planner = DefaultSavePlanner(enable_plan_caching=True)
+                tensor = torch.rand(10)
+                val = [1, 2, 3]
+                st = create_sharded_tensor(rank=rank, world_size=4, shards_per_rank=1)
+                state_dict = {"tensor": tensor, "value": val, "st": st}
+                planner.set_up_planner(state_dict, is_coordinator=(rank == 0))
+                return planner.create_local_plan()
+
+        planner = DefaultSavePlanner(enable_plan_caching=True)
+        key = planner._cached_plans_key
+
+        tensor = torch.rand(10)
+        val = [1, 2, 3]
+        state_dict = {"tensor": tensor, "value": val}
+        planner.set_up_planner(state_dict, is_coordinator=True)
+
+        local_plan = planner.create_local_plan()
+        self.assertTrue(local_plan.usable)
+        self.assertNotIn(key, SavePlanner._cached_save_plan)
+
+        all_plans = [create_data(0), create_data(1), create_data(2), create_data(3)]
+
+        with patch(
+            "torch.distributed.checkpoint.default_planner._validate_global_plan",
+            return_value=["mock validation error"],
+        ):
+            with self.assertRaises(ValueError):
+                planner.create_global_plan(all_plans)
+
+        self.assertNotIn(key, SavePlanner._cached_all_plans)
+        self.assertNotIn(key, SavePlanner._cached_global_plan)
+        self.assertNotIn(key, SavePlanner._cached_metadata)
+        self.assertNotIn(key, SavePlanner._cached_save_plan)
+
+        # Second save should succeed, not KeyError
+        planner.set_up_planner(state_dict, is_coordinator=True)
+        local_plan_2 = planner.create_local_plan()
+        self.assertTrue(local_plan_2.usable)
+
+        global_plan, metadata = planner.create_global_plan(all_plans)
+        self.assertIsNotNone(global_plan)
+        self.assertIsNotNone(metadata)
+        self.assertIn(key, SavePlanner._cached_all_plans)
+        self.assertIn(key, SavePlanner._cached_global_plan)
+        self.assertIn(key, SavePlanner._cached_metadata)
+
+        # Clean up class-level caches
+        SavePlanner._cached_save_plan.pop(key, None)
+        SavePlanner._cached_final_save_plan.pop(key, None)
+        SavePlanner._cached_all_plans.pop(key, None)
+        SavePlanner._cached_global_plan.pop(key, None)
+        SavePlanner._cached_metadata.pop(key, None)
+
     def test_local_load_plan(self):
         def create_state_dict(rank):
             with with_dist(rank=rank, world_size=4):
@@ -577,7 +658,7 @@ def test_non_overlapping_chunks(self):
             for i in range(4)
         ]
         metadata = self._make_metadata(chunks, [4])
-        self.assertTrue(_validate_global_plan([SavePlan([])], metadata))
+        self.assertEqual(_validate_global_plan([SavePlan([])], metadata), [])
 
     def test_detect_overlapping_chunks(self):
         chunks = [
@@ -585,7 +666,7 @@ def test_detect_overlapping_chunks(self):
             ChunkStorageMetadata(offsets=torch.Size([1]), sizes=torch.Size([2])),
         ]
         metadata = self._make_metadata(chunks, [4])
-        self.assertFalse(_validate_global_plan([SavePlan([])], metadata))
+        self.assertGreater(len(_validate_global_plan([SavePlan([])], metadata)), 0)
 
 
 class TestLoadPlanner(TestCase):
diff --git a/test/distributed/checkpoint/test_state_dict.py b/test/distributed/checkpoint/test_state_dict.py
index 03bcf7ce5e03e..6bfea439468cb 100644
--- a/test/distributed/checkpoint/test_state_dict.py
+++ b/test/distributed/checkpoint/test_state_dict.py
@@ -5,7 +5,6 @@
 import sys
 from collections.abc import Callable
 from itertools import chain, product
-from typing import Union
 
 import torch
 import torch.distributed as dist
@@ -256,7 +255,7 @@ def test_compiled_fsdp(self) -> None:
     def _test_fsdp2(
         self,
         *,
-        reshard_after_forward: Union[bool, int],
+        reshard_after_forward: bool | int,
         optimizer_class: type[Optimizer],
         compile_model: bool,
         foreach: bool = True,
@@ -952,7 +951,10 @@ def test_set_cpu_model_state_dict_broadcast_from_rank0(self) -> None:
             model.state_dict().items(),
             expected_state_dict.items(),
         ):
-            assert actual_name == expected_name
+            if actual_name != expected_name:
+                raise AssertionError(
+                    f"Expected name {expected_name}, got {actual_name}"
+                )
             torch.testing.assert_close(tensor, expected_tensor, msg=expected_name)
 
     @with_comms
diff --git a/test/distributed/checkpoint/test_state_dict_stager.py b/test/distributed/checkpoint/test_state_dict_stager.py
index 02e4e6ec784e8..8e3baf6ea3ec3 100644
--- a/test/distributed/checkpoint/test_state_dict_stager.py
+++ b/test/distributed/checkpoint/test_state_dict_stager.py
@@ -230,10 +230,11 @@ def test_views(self):
                         "type": TestStruct(tensor1=tensor3.narrow(0, 0, 2)),
                     },
                 }
-                assert (
+                if (
                     state_dict["tensor1"].storage().data_ptr()
-                    == state_dict["tensor2"].storage().data_ptr()
-                )
+                    != state_dict["tensor2"].storage().data_ptr()
+                ):
+                    raise AssertionError("tensor1 and tensor2 should share storage")
 
                 stager = StateDictStager(
                     pin_memory=pin_memory, share_memory=share_memory
@@ -250,9 +251,10 @@ def test_views(self):
 
                 # Validate tensor count and bytes
                 expected_storage_cnt = 2
-                assert num_storages == expected_storage_cnt, (
-                    f"Expected {expected_storage_cnt} storages, got {num_storages}"
-                )
+                if num_storages != expected_storage_cnt:
+                    raise AssertionError(
+                        f"Expected {expected_storage_cnt} storages, got {num_storages}"
+                    )
 
                 # Calculate expected bytes
                 # Note: Only unique storages are counted in the byte count
@@ -261,28 +263,46 @@ def test_views(self):
                     + tensor3.numel()  # tensor1 and tensor2 share storage
                     * tensor3.element_size()  # tensor3 and its narrow view share storage
                 )
-                assert num_bytes == expected_bytes, (
-                    f"Expected {expected_bytes} bytes, got {num_bytes}"
-                )
+                if num_bytes != expected_bytes:
+                    raise AssertionError(
+                        f"Expected {expected_bytes} bytes, got {num_bytes}"
+                    )
                 # Verify that the CPU state dict is equivalent to the original GPU state dict
                 result, error = compare_state_dicts(state_dict, cpu_state_dict)
-                assert result, f"State dicts are not equivalent: {error}"
+                if not result:
+                    raise AssertionError(f"State dicts are not equivalent: {error}")
 
                 # Additional checks for storage sharing
-                assert cpu_state_dict["tensor1"].device == torch.device("cpu")
-                assert cpu_state_dict["tensor2"].device == torch.device("cpu")
-                assert (
+                if cpu_state_dict["tensor1"].device != torch.device("cpu"):
+                    raise AssertionError(
+                        f"Expected tensor1 on cpu, got {cpu_state_dict['tensor1'].device}"
+                    )
+                if cpu_state_dict["tensor2"].device != torch.device("cpu"):
+                    raise AssertionError(
+                        f"Expected tensor2 on cpu, got {cpu_state_dict['tensor2'].device}"
+                    )
+                if (
                     cpu_state_dict["tensor1"].storage().data_ptr()
-                    == cpu_state_dict["tensor2"].storage().data_ptr()
-                )
+                    != cpu_state_dict["tensor2"].storage().data_ptr()
+                ):
+                    raise AssertionError("cpu tensor1 and tensor2 should share storage")
 
                 recursive = cpu_state_dict["recursive"]
-                assert recursive["tensor3"].device == torch.device("cpu")
-                assert recursive["type"].tensor1.device == torch.device("cpu")
-                assert (
+                if recursive["tensor3"].device != torch.device("cpu"):
+                    raise AssertionError(
+                        f"Expected tensor3 on cpu, got {recursive['tensor3'].device}"
+                    )
+                if recursive["type"].tensor1.device != torch.device("cpu"):
+                    raise AssertionError(
+                        f"Expected type.tensor1 on cpu, got {recursive['type'].tensor1.device}"
+                    )
+                if (
                     recursive["tensor3"].storage().data_ptr()
-                    == recursive["type"].tensor1.storage().data_ptr()
-                )
+                    != recursive["type"].tensor1.storage().data_ptr()
+                ):
+                    raise AssertionError(
+                        "tensor3 and type.tensor1 should share storage"
+                    )
 
     @unittest.skipIf(not HAS_ACCELERATOR, "No accelerator")
     def test_caching(self):
@@ -323,9 +343,10 @@ def test_caching(self):
 
                 # Verify the first result is correct
                 result, error = compare_state_dicts(state_dict, cpu_state_dict1)
-                assert result, (
-                    f"First state dict is not equivalent to original: {error}"
-                )
+                if not result:
+                    raise AssertionError(
+                        f"First state dict is not equivalent to original: {error}"
+                    )
 
                 # Modify the original tensors
                 tensor1.fill_(0)
@@ -339,28 +360,33 @@ def test_caching(self):
 
                 # Verify that the second CPU state dict is equivalent to the modified original state dict
                 result, error = compare_state_dicts(state_dict, cpu_state_dict2)
-                assert result, (
-                    f"Second state dict is not equivalent to modified original: {error}"
-                )
+                if not result:
+                    raise AssertionError(
+                        f"Second state dict is not equivalent to modified original: {error}"
+                    )
 
                 # Verify that the number of cached storages hasn't changed
-                assert num_storages1 == num_storages2, (
-                    f"Storage count changed: {num_storages1} vs {num_storages2}"
-                )
+                if num_storages1 != num_storages2:
+                    raise AssertionError(
+                        f"Storage count changed: {num_storages1} vs {num_storages2}"
+                    )
 
                 # Verify that the tensors in the second state dict have the same storage pointers as the first
-                assert (
+                if (
                     cpu_state_dict1["tensor1"].storage().data_ptr()
-                    == cpu_state_dict2["tensor1"].storage().data_ptr()
-                ), "Storage pointers should match for tensor1"
-                assert (
+                    != cpu_state_dict2["tensor1"].storage().data_ptr()
+                ):
+                    raise AssertionError("Storage pointers should match for tensor1")
+                if (
                     cpu_state_dict1["tensor2"].storage().data_ptr()
-                    == cpu_state_dict2["tensor2"].storage().data_ptr()
-                ), "Storage pointers should match for tensor2"
-                assert (
+                    != cpu_state_dict2["tensor2"].storage().data_ptr()
+                ):
+                    raise AssertionError("Storage pointers should match for tensor2")
+                if (
                     cpu_state_dict1["recursive"]["tensor3"].storage().data_ptr()
-                    == cpu_state_dict2["recursive"]["tensor3"].storage().data_ptr()
-                ), "Storage pointers should match for tensor3"
+                    != cpu_state_dict2["recursive"]["tensor3"].storage().data_ptr()
+                ):
+                    raise AssertionError("Storage pointers should match for tensor3")
 
                 # Modify the original tensors again with different values
                 tensor1.fill_(42.0)
@@ -369,12 +395,14 @@ def test_caching(self):
                 cpu_state_dict3 = stager.stage(state_dict)
 
                 # Verify that the third CPU state dict reflects the updated values
-                assert torch.all(cpu_state_dict3["tensor1"] == 42.0), (
-                    "Updated values should be reflected in the cached state dict"
-                )
-                assert torch.all(cpu_state_dict3["tensor2"] == 42.0), (
-                    "Updated values should be reflected in the cached state dict"
-                )
+                if not torch.all(cpu_state_dict3["tensor1"] == 42.0):
+                    raise AssertionError(
+                        "Updated values should be reflected in the cached state dict for tensor1"
+                    )
+                if not torch.all(cpu_state_dict3["tensor2"] == 42.0):
+                    raise AssertionError(
+                        "Updated values should be reflected in the cached state dict for tensor2"
+                    )
 
     @unittest.skipIf(not HAS_ACCELERATOR, "No accelerator")
     def test_tensor_attrs(self):
@@ -403,24 +431,24 @@ def test_tensor_attrs(self):
         cpu_state_dict = stager.stage(state_dict)
 
         # Verify that tensor attributes are preserved
-        assert hasattr(cpu_state_dict["tensor1"], "a"), (
-            "Tensor attribute 'a' was not preserved"
-        )
-        assert cpu_state_dict["tensor1"].a == 42, (
-            "Tensor attribute 'a' has incorrect value"
-        )
-        assert hasattr(cpu_state_dict["tensor1"], "b"), (
-            "Tensor attribute 'b' was not preserved"
-        )
-        assert cpu_state_dict["tensor1"].b == 43, (
-            "Tensor attribute 'b' has incorrect value"
-        )
-        assert hasattr(cpu_state_dict["recursive"]["tensor3"], "c"), (
-            "Tensor attribute 'c' was not preserved"
-        )
-        assert cpu_state_dict["recursive"]["tensor3"].c == 44, (
-            "Tensor attribute 'c' has incorrect value"
-        )
+        if not hasattr(cpu_state_dict["tensor1"], "a"):
+            raise AssertionError("Tensor attribute 'a' was not preserved")
+        if cpu_state_dict["tensor1"].a != 42:
+            raise AssertionError(
+                f"Tensor attribute 'a' has incorrect value: {cpu_state_dict['tensor1'].a}"
+            )
+        if not hasattr(cpu_state_dict["tensor1"], "b"):
+            raise AssertionError("Tensor attribute 'b' was not preserved")
+        if cpu_state_dict["tensor1"].b != 43:
+            raise AssertionError(
+                f"Tensor attribute 'b' has incorrect value: {cpu_state_dict['tensor1'].b}"
+            )
+        if not hasattr(cpu_state_dict["recursive"]["tensor3"], "c"):
+            raise AssertionError("Tensor attribute 'c' was not preserved")
+        if cpu_state_dict["recursive"]["tensor3"].c != 44:
+            raise AssertionError(
+                f"Tensor attribute 'c' has incorrect value: {cpu_state_dict['recursive']['tensor3'].c}"
+            )
 
     @unittest.skipIf(not HAS_ACCELERATOR, "No accelerator")
     def test_different_dtypes(self):
diff --git a/test/distributed/checkpoint/test_utils.py b/test/distributed/checkpoint/test_utils.py
index 79dbe741822c7..b6f66ba97a52b 100644
--- a/test/distributed/checkpoint/test_utils.py
+++ b/test/distributed/checkpoint/test_utils.py
@@ -2,7 +2,7 @@
 
 import io
 import sys
-from typing import Optional
+import traceback
 
 import torch
 import torch.distributed as dist
@@ -14,6 +14,7 @@
 )
 from torch.distributed._shard.sharded_tensor.metadata import TensorProperties
 from torch.distributed.c10d_logger import _c10d_logger
+from torch.distributed.checkpoint.api import _wrap_exception
 from torch.distributed.checkpoint.logger import _dcp_logger
 from torch.distributed.checkpoint.metadata import MetadataIndex
 from torch.distributed.checkpoint.utils import (
@@ -21,6 +22,7 @@
     _DistWrapper,
     find_state_dict_object,
 )
+from torch.distributed.distributed_c10d import _object_to_tensor, _tensor_to_object
 from torch.testing._internal.common_utils import (
     run_tests,
     TEST_WITH_DEV_DBG_ASAN,
@@ -140,6 +142,44 @@ def test_dcp_logger(self):
         self.assertEqual(1, len(_c10d_logger.handlers))
 
 
+class TestWrapException(TestCase):
+    def test_wrap_exception_serializable_via_object_to_tensor(self):
+        """Verify _wrap_exception produces a result that _object_to_tensor can serialize.
+
+        Python 3.13+ adds a _code attribute to FrameSummary containing
+        bytecode objects that cannot be pickled. _wrap_exception must
+        clear these so that _object_to_tensor (used by gather_object and
+        scatter_object_list) succeeds instead of raising
+        "TypeError: cannot pickle code objects".
+        """
+        try:
+            raise ValueError("test error")
+        except ValueError as e:
+            wrapped = _wrap_exception(e)
+
+        # _object_to_tensor / _tensor_to_object are what gather_object
+        # and scatter_object_list use to serialize objects across ranks.
+        # This would raise "TypeError: cannot pickle code objects"
+        # on Python 3.13+ without the fix.
+        byte_tensor, size = _object_to_tensor(wrapped, torch.device("cpu"), None)
+        restored = _tensor_to_object(byte_tensor, size.item(), None)
+
+        self.assertIsInstance(restored[0], ValueError)
+        self.assertEqual(str(restored[0]), "test error")
+        self.assertIsInstance(restored[1], traceback.StackSummary)
+        self.assertGreater(len(restored[1]), 0)
+
+    def test_wrap_exception_preserves_traceback_formatting(self):
+        """Verify that clearing _code does not break traceback formatting."""
+        try:
+            raise RuntimeError("format test")
+        except RuntimeError as e:
+            wrapped = _wrap_exception(e)
+
+        formatted = "".join(traceback.format_list(wrapped[1]))
+        self.assertIn("raise RuntimeError", formatted)
+
+
 class TestReaderView(TestCase):
     def setUp(self):
         buffer = io.BytesIO(bytearray(range(ord("A"), ord("Z") + 1)))
@@ -219,7 +259,8 @@ def test_gather_object(self):
             if rank % half_world_size == 0
             else None
         )
-        assert gathered_objects == expected_objects
+        if gathered_objects != expected_objects:
+            raise AssertionError(f"Expected {expected_objects}, got {gathered_objects}")
 
     @with_comms
     @skip_if_lt_x_gpu(4)
@@ -241,7 +282,10 @@ def test_scatter_object(self):
         )
         scattered_objects = dist_wrapper.scatter_object(objects)
         expected_objects = rank
-        assert scattered_objects == expected_objects
+        if scattered_objects != expected_objects:
+            raise AssertionError(
+                f"Expected {expected_objects}, got {scattered_objects}"
+            )
 
     @with_comms
     @skip_if_lt_x_gpu(2)
@@ -255,11 +299,12 @@ def test_broadcast_object_with_nonzero_coordinator(self):
 
         rank = dist.get_rank()
         # only local rank 1 supplies the payload
-        payload: Optional[int] = rank if rank == 1 else None
+        payload: int | None = rank if rank == 1 else None
 
         result = dist_wrapper.broadcast_object(payload)
         # every rank should receive the value from global rank 1
-        assert result == 1
+        if result != 1:
+            raise AssertionError(f"Expected 1, got {result}")
 
     @with_comms
     @skip_if_lt_x_gpu(4)
@@ -276,14 +321,15 @@ def test_broadcast_object_global_local_mismatch(self):
         rank = mesh_2d.get_rank()
 
         # only the local coordinator in each subgroup provides payload
-        payload: Optional[int] = rank if dist_wrapper.is_coordinator else None
+        payload: int | None = rank if dist_wrapper.is_coordinator else None
         got = dist_wrapper.broadcast_object(payload)
 
         # ensure we broadcast from the *global* coordinator rank,
         # not the local index.  For rows [0,1] this is global rank 1;
         # for rows [2,3] this is global rank 3.
         expected = dist_wrapper.global_coordinator_rank
-        assert got == expected
+        if got != expected:
+            raise AssertionError(f"Expected {expected}, got {got}")
 
     @with_comms
     @skip_if_lt_x_gpu(2)
diff --git a/test/distributed/elastic/agent/server/test/api_test.py b/test/distributed/elastic/agent/server/test/api_test.py
index 2af4175e6e060..d086905a7dbb5 100644
--- a/test/distributed/elastic/agent/server/test/api_test.py
+++ b/test/distributed/elastic/agent/server/test/api_test.py
@@ -485,6 +485,47 @@ def test_invoke_run(
         record_events_mock.assert_called_once()
         shutdown_mock.assert_called_once()
 
+    def test_exit_barrier_sets_and_clears_flag(self):
+        """Verify _in_exit_barrier is set before barrier and cleared after."""
+        spec = self._get_worker_spec(max_restarts=0)
+        agent = TestAgent(spec)
+        agent._worker_group.state = WorkerState.SUCCEEDED
+        agent._worker_group.group_world_size = 1
+        agent._store = MagicMock()
+
+        flag_during_barrier = []
+
+        def mock_barrier(**kwargs):
+            flag_during_barrier.append(agent._in_exit_barrier)
+
+        with patch(
+            "torch.distributed.elastic.utils.store.barrier",
+            side_effect=mock_barrier,
+        ):
+            agent._exit_barrier()
+
+        # Flag was True during barrier call
+        self.assertTrue(flag_during_barrier[0])
+        # Flag is False after barrier completes
+        self.assertFalse(agent._in_exit_barrier)
+
+    def test_exit_barrier_clears_flag_on_timeout(self):
+        """Verify _in_exit_barrier is cleared even if barrier times out."""
+        spec = self._get_worker_spec(max_restarts=0)
+        agent = TestAgent(spec)
+        agent._worker_group.state = WorkerState.SUCCEEDED
+        agent._worker_group.group_world_size = 1
+        agent._store = MagicMock()
+
+        with patch(
+            "torch.distributed.elastic.utils.store.barrier",
+            side_effect=Exception("wait timeout"),
+        ):
+            agent._exit_barrier()
+
+        # Flag must be cleared even on timeout
+        self.assertFalse(agent._in_exit_barrier)
+
     @patch("torch.distributed.elastic.agent.server.api.put_metric")
     def test_record_metrics_success_no_retries(self, put_metric_mock):
         spec = self._get_worker_spec(max_restarts=1)
diff --git a/test/distributed/elastic/agent/server/test/local_elastic_agent_test.py b/test/distributed/elastic/agent/server/test/local_elastic_agent_test.py
index 22310a9ba4955..6514762f23ffe 100644
--- a/test/distributed/elastic/agent/server/test/local_elastic_agent_test.py
+++ b/test/distributed/elastic/agent/server/test/local_elastic_agent_test.py
@@ -6,6 +6,8 @@
 #
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
+from __future__ import annotations
+
 import json
 import multiprocessing as mp
 import os
@@ -16,9 +18,8 @@
 import time
 import unittest
 import uuid
-from collections.abc import Callable
 from dataclasses import dataclass
-from typing import Optional
+from typing import TYPE_CHECKING
 from unittest import mock
 from unittest.mock import Mock, patch
 
@@ -33,6 +34,7 @@
 )
 from torch.distributed.elastic.agent.server.local_elastic_agent import (
     LocalElasticAgent,
+    TORCHELASTIC_ENABLE_FILE_TIMER,
     TORCHELASTIC_HEALTH_CHECK_PORT,
     TORCHELASTIC_TIMER_FILE,
 )
@@ -48,6 +50,10 @@
 )
 
 
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+
 def init_rpc(name, backend):
     rank = int(os.environ["RANK"])
     world_size = int(os.environ["WORLD_SIZE"])
@@ -292,8 +298,8 @@ def get_worker_spec(
         max_nodes=1,
         max_restarts=0,
         monitor_interval=0.01,
-        master_addr_override: Optional[str] = None,
-        master_port_override: Optional[int] = None,
+        master_addr_override: str | None = None,
+        master_port_override: int | None = None,
         is_host=True,
     ):
         rdzv_params = RendezvousParameters(
@@ -323,7 +329,7 @@ def get_agent(
         node_config: Conf,
         start_method: str = "spawn",
         exit_barrier_timeout=5,
-        log_line_prefix_template: Optional[str] = None,
+        log_line_prefix_template: str | None = None,
     ) -> LocalElasticAgent:
         return LocalElasticAgent(
             spec,
@@ -343,18 +349,18 @@ def get_agent(
     def run_agent(
         self,
         conf: Conf,
-        agent_results: Optional[mp.Queue] = None,  # (role, agent_result)
+        agent_results: mp.Queue | None = None,  # (role, agent_result)
         min_nodes=1,
         max_nodes=1,
         start_method: str = "spawn",
         max_restarts: int = 0,
         exit_barrier_timeout=5,
-        master_addr_override: Optional[str] = None,
-        master_port_override: Optional[int] = None,
+        master_addr_override: str | None = None,
+        master_port_override: int | None = None,
         is_host=True,
         monitor_interval=0.01,
-        log_line_prefix_template: Optional[str] = None,
-    ) -> Optional[RunResult]:
+        log_line_prefix_template: str | None = None,
+    ) -> RunResult | None:
         """
         Runs a single agent. This method can be called either on a separate process
         or the main test process. When calling this method on a separate process make
@@ -397,7 +403,7 @@ def run_job(
         self,
         node_configs: list[Conf],
         exit_barrier_timeout: int = 5,
-        log_line_prefix_template: Optional[str] = None,
+        log_line_prefix_template: str | None = None,
     ) -> dict[str, list[RunResult]]:
         """
         Simulates running a distributed job by running multiple agents
@@ -774,7 +780,7 @@ def test_simple_dist_sum_etcd_v2(self):
         self.run_test_with_backend(backend="etcd-v2", test_to_run=self.simple_dist_sum)
 
     def run_distributed_sum_homogeneous(
-        self, log_line_prefix_template: Optional[str] = None
+        self, log_line_prefix_template: str | None = None
     ):
         node_configs = [
             Conf(role="sum", entrypoint=_dist_sum, local_world_size=4, tee=Std.ALL),
@@ -1468,6 +1474,282 @@ def fail_rank_one_once(self):
     def test_rank_restart_after_failure(self):
         self.run_test_with_backend(backend="c10d", test_to_run=self.fail_rank_one_once)
 
+    def test_get_alive_time_without_watchdog(self):
+        """Test _get_alive_time returns current time when watchdog is not set."""
+        node_conf = Conf(entrypoint=_happy_function, local_world_size=1)
+
+        # Use a mocked rendezvous handler to avoid blocking on real connections
+        rdzv_handler = Mock()
+        rdzv_handler.get_backend.return_value = "c10d"
+        spec = WorkerSpec(
+            role=node_conf.role,
+            local_world_size=node_conf.local_world_size,
+            entrypoint=node_conf.entrypoint,
+            args=node_conf.args,
+            rdzv_handler=rdzv_handler,
+            max_restarts=0,
+            monitor_interval=0.01,
+        )
+        agent = self.get_agent(spec, node_config=node_conf)
+
+        # Watchdog is None by default at initialization
+        self.assertIsNone(agent._worker_watchdog)
+
+        # _get_alive_time should return current time when watchdog is None
+        before = int(time.time())
+        alive_time = agent._get_alive_time()
+        after = int(time.time())
+
+        self.assertGreaterEqual(alive_time, before)
+        self.assertLessEqual(alive_time, after)
+
+    def test_get_alive_time_with_watchdog(self):
+        """Test _get_alive_time returns watchdog time when watchdog is active."""
+        node_conf = Conf(entrypoint=_happy_function, local_world_size=1)
+
+        # Use a mocked rendezvous handler to avoid blocking on real connections
+        rdzv_handler = Mock()
+        rdzv_handler.get_backend.return_value = "c10d"
+        spec = WorkerSpec(
+            role=node_conf.role,
+            local_world_size=node_conf.local_world_size,
+            entrypoint=node_conf.entrypoint,
+            args=node_conf.args,
+            rdzv_handler=rdzv_handler,
+            max_restarts=0,
+            monitor_interval=0.01,
+        )
+        agent = self.get_agent(spec, node_config=node_conf)
+
+        # Mock the watchdog with a specific return value
+        mock_watchdog = Mock()
+        expected_time = 1234567890
+        mock_watchdog.get_last_progress_time.return_value = expected_time
+        agent._worker_watchdog = mock_watchdog
+
+        # _get_alive_time should delegate to watchdog when available
+        alive_time = agent._get_alive_time()
+
+        self.assertEqual(alive_time, expected_time)
+        mock_watchdog.get_last_progress_time.assert_called_once()
+
+    def test_get_alive_time_during_exit_barrier(self):
+        """When _in_exit_barrier is True, _get_alive_time returns current time
+        even if the watchdog exists (since watchdog progress time is stale)."""
+        node_conf = Conf(entrypoint=_happy_function, local_world_size=1)
+
+        rdzv_handler = Mock()
+        rdzv_handler.get_backend.return_value = "c10d"
+        spec = WorkerSpec(
+            role=node_conf.role,
+            local_world_size=node_conf.local_world_size,
+            entrypoint=node_conf.entrypoint,
+            args=node_conf.args,
+            rdzv_handler=rdzv_handler,
+            max_restarts=0,
+            monitor_interval=0.01,
+        )
+        agent = self.get_agent(spec, node_config=node_conf)
+
+        # Set up a mock watchdog with a stale time
+        mock_watchdog = Mock()
+        mock_watchdog.get_last_progress_time.return_value = 1000  # very old
+        agent._worker_watchdog = mock_watchdog
+
+        # Without exit barrier flag, should return watchdog time
+        agent._in_exit_barrier = False
+        self.assertEqual(agent._get_alive_time(), 1000)
+
+        # With exit barrier flag, should return current time (not stale watchdog time)
+        agent._in_exit_barrier = True
+        alive_time = agent._get_alive_time()
+        now = int(time.time())
+        self.assertAlmostEqual(alive_time, now, delta=2)
+
+    def test_healthcheck_reports_healthy_during_exit_barrier(self):
+        """Verify health check callback returns current time during exit barrier,
+        preventing TW from killing the task while agents coordinate shutdown."""
+        node_conf = Conf(entrypoint=_happy_function, local_world_size=1)
+
+        rdzv_handler = Mock()
+        rdzv_handler.get_backend.return_value = "c10d"
+        spec = WorkerSpec(
+            role=node_conf.role,
+            local_world_size=node_conf.local_world_size,
+            entrypoint=node_conf.entrypoint,
+            args=node_conf.args,
+            rdzv_handler=rdzv_handler,
+            max_restarts=0,
+            monitor_interval=0.01,
+        )
+        agent = self.get_agent(spec, node_config=node_conf)
+
+        # Simulate post-training state: watchdog exists but progress time is stale
+        mock_watchdog = Mock()
+        stale_time = int(time.time()) - 600  # 10 minutes ago
+        mock_watchdog.get_last_progress_time.return_value = stale_time
+        agent._worker_watchdog = mock_watchdog
+
+        # Before exit barrier: alive_time is stale (would fail TW health check)
+        alive_time = agent._get_alive_time()
+        self.assertEqual(alive_time, stale_time)
+
+        # During exit barrier: alive_time is current (TW health check passes)
+        agent._in_exit_barrier = True
+        alive_time = agent._get_alive_time()
+        now = int(time.time())
+        self.assertAlmostEqual(alive_time, now, delta=2)
+
+        # After exit barrier: back to watchdog time
+        agent._in_exit_barrier = False
+        alive_time = agent._get_alive_time()
+        self.assertEqual(alive_time, stale_time)
+
+    def test_in_exit_barrier_initialized_false(self):
+        """Verify _in_exit_barrier is False on agent construction."""
+        node_conf = Conf(entrypoint=_happy_function, local_world_size=1)
+
+        rdzv_handler = Mock()
+        rdzv_handler.get_backend.return_value = "c10d"
+        spec = WorkerSpec(
+            role=node_conf.role,
+            local_world_size=node_conf.local_world_size,
+            entrypoint=node_conf.entrypoint,
+            args=node_conf.args,
+            rdzv_handler=rdzv_handler,
+            max_restarts=0,
+            monitor_interval=0.01,
+        )
+        agent = self.get_agent(spec, node_config=node_conf)
+        self.assertFalse(agent._in_exit_barrier)
+
+    def test_setup_healthcheck_idempotent(self):
+        """Test _setup_healthcheck is idempotent and does not recreate server when JK is enabled."""
+        node_conf = Conf(entrypoint=_happy_function, local_world_size=1)
+
+        # Use a mocked rendezvous handler to avoid blocking on real connections
+        rdzv_handler = Mock()
+        rdzv_handler.get_backend.return_value = "c10d"
+        spec = WorkerSpec(
+            role=node_conf.role,
+            local_world_size=node_conf.local_world_size,
+            entrypoint=node_conf.entrypoint,
+            args=node_conf.args,
+            rdzv_handler=rdzv_handler,
+            max_restarts=0,
+            monitor_interval=0.01,
+        )
+        agent = self.get_agent(spec, node_config=node_conf)
+
+        # Pre-set a mock health check server
+        mock_server = Mock()
+        agent._health_check_server = mock_server
+
+        # Set the env for healthcheck port
+        healthcheck_port_env_name = TORCHELASTIC_HEALTH_CHECK_PORT
+        original_value = os.environ.get(healthcheck_port_env_name)
+        os.environ[healthcheck_port_env_name] = "12345"
+
+        try:
+            # Patch justknobs_check to return True for the healthcheck knob
+            with patch(
+                "torch.distributed.elastic.agent.server.local_elastic_agent.justknobs_check",
+                return_value=True,
+            ):
+                # Call _setup_healthcheck - should be a no-op since server already exists
+                agent._setup_healthcheck()
+
+            # Verify the mock server was not replaced or modified
+            self.assertIs(agent._health_check_server, mock_server)
+            # Verify start was not called on the mock (it was already "running")
+            mock_server.start.assert_not_called()
+        finally:
+            # Restore env
+            if original_value is None:
+                if healthcheck_port_env_name in os.environ:
+                    del os.environ[healthcheck_port_env_name]
+            else:
+                os.environ[healthcheck_port_env_name] = original_value
+
+    def test_setup_healthcheck_creates_server_when_none(self):
+        """Test _setup_healthcheck creates server when none exists."""
+        node_conf = Conf(entrypoint=_happy_function, local_world_size=1)
+        self._backend = "c10d"
+        self._endpoint = f"localhost:{acquire_available_port()}"
+        spec = self.get_worker_spec(node_conf)
+        agent = self.get_agent(spec, node_config=node_conf)
+
+        # Ensure no health check server exists
+        self.assertIsNone(agent._health_check_server)
+
+        # Set the env for healthcheck
+        healthcheck_port_env_name = TORCHELASTIC_HEALTH_CHECK_PORT
+        original_value = os.environ.get(healthcheck_port_env_name)
+        healthcheck_port = str(acquire_available_port())
+        os.environ[healthcheck_port_env_name] = healthcheck_port
+
+        try:
+            # Call _setup_healthcheck
+            agent._setup_healthcheck()
+
+            # Verify a health check server was created
+            self.assertIsNotNone(agent._health_check_server)
+        finally:
+            # Cleanup
+            if agent._health_check_server is not None:
+                agent._health_check_server.stop()
+                agent._health_check_server = None
+            if original_value is None:
+                if healthcheck_port_env_name in os.environ:
+                    del os.environ[healthcheck_port_env_name]
+            else:
+                os.environ[healthcheck_port_env_name] = original_value
+
+    def test_healthcheck_with_watchdog_enabled(self):
+        """Test healthcheck works with watchdog enabled during agent run."""
+        # Set the env for watchdog and healthcheck
+        watchdog_env_name = TORCHELASTIC_ENABLE_FILE_TIMER
+        healthcheck_port_env_name = TORCHELASTIC_HEALTH_CHECK_PORT
+
+        original_watchdog = os.environ.get(watchdog_env_name)
+        original_healthcheck = os.environ.get(healthcheck_port_env_name)
+
+        os.environ[watchdog_env_name] = "1"
+        os.environ[healthcheck_port_env_name] = str(acquire_available_port())
+
+        try:
+            node_conf = Conf(
+                entrypoint=_check_local_watchdog_setup,
+                local_world_size=1,
+                args=(TORCHELASTIC_ENABLE_FILE_TIMER, True),
+            )
+
+            def run_test():
+                spec = self.get_worker_spec(node_conf, max_restarts=0)
+                agent = self.get_agent(spec, node_config=node_conf)
+
+                # Run the agent with justknobs_check returning True
+                with patch(
+                    "torch.distributed.elastic.agent.server.local_elastic_agent.justknobs_check",
+                    return_value=True,
+                ):
+                    res = agent.run()
+                self.assertFalse(res.is_failed())
+
+            self.run_test_with_backend(backend="c10d", test_to_run=run_test)
+        finally:
+            # Restore env
+            if original_watchdog is None:
+                if watchdog_env_name in os.environ:
+                    del os.environ[watchdog_env_name]
+            else:
+                os.environ[watchdog_env_name] = original_watchdog
+            if original_healthcheck is None:
+                if healthcheck_port_env_name in os.environ:
+                    del os.environ[healthcheck_port_env_name]
+            else:
+                os.environ[healthcheck_port_env_name] = original_healthcheck
+
 
 if __name__ == "__main__":
     raise RuntimeError(
diff --git a/test/distributed/elastic/agent/server/test/shutdown_timeout_test.py b/test/distributed/elastic/agent/server/test/shutdown_timeout_test.py
new file mode 100644
index 0000000000000..1366df0555440
--- /dev/null
+++ b/test/distributed/elastic/agent/server/test/shutdown_timeout_test.py
@@ -0,0 +1,169 @@
+#!/usr/bin/env python3
+# Owner(s): ["oncall: r2p"]
+
+# Copyright (c) Facebook, Inc. and its affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+import os
+import signal
+import tempfile
+import time
+from unittest.mock import MagicMock, Mock, patch
+
+from torch.distributed.elastic.agent.server.api import SimpleElasticAgent, WorkerSpec
+from torch.distributed.elastic.agent.server.local_elastic_agent import LocalElasticAgent
+from torch.distributed.elastic.multiprocessing import DefaultLogsSpecs
+from torch.distributed.launcher.api import LaunchConfig
+from torch.testing._internal.common_utils import run_tests, TestCase
+
+
+def _sleep_long(duration: int):
+    """Worker function that sleeps for a specified duration."""
+    time.sleep(duration)
+    return int(os.environ["RANK"])
+
+
+class ShutdownTimeoutTest(TestCase):
+    """Tests for the configurable shutdown_timeout feature."""
+
+    def test_launch_config_default_shutdown_timeout(self):
+        config = LaunchConfig(
+            min_nodes=1,
+            max_nodes=1,
+            nproc_per_node=2,
+        )
+        self.assertEqual(config.shutdown_timeout, 30)
+
+    def test_launch_config_custom_shutdown_timeout(self):
+        config = LaunchConfig(
+            min_nodes=1,
+            max_nodes=1,
+            nproc_per_node=2,
+            shutdown_timeout=120,
+        )
+        self.assertEqual(config.shutdown_timeout, 120)
+
+    def test_launch_config_env_var_shutdown_timeout(self):
+        with patch.dict(os.environ, {"TORCH_ELASTIC_SHUTDOWN_TIMEOUT": "600"}):
+            config = LaunchConfig(
+                min_nodes=1,
+                max_nodes=1,
+                nproc_per_node=2,
+            )
+            self.assertEqual(config.shutdown_timeout, 600)
+
+    def test_launch_config_explicit_overrides_env(self):
+        with patch.dict(os.environ, {"TORCH_ELASTIC_SHUTDOWN_TIMEOUT": "600"}):
+            config = LaunchConfig(
+                min_nodes=1,
+                max_nodes=1,
+                nproc_per_node=2,
+                shutdown_timeout=120,
+            )
+            self.assertEqual(config.shutdown_timeout, 120)
+
+    def test_simple_elastic_agent_receives_shutdown_timeout(self):
+        mock_spec = Mock(spec=WorkerSpec)
+        mock_spec.max_restarts = 3
+
+        agent = SimpleElasticAgent(
+            spec=mock_spec,
+            exit_barrier_timeout=300,
+            shutdown_timeout=180,
+        )
+
+        self.assertEqual(agent._shutdown_timeout, 180)
+
+    def test_simple_elastic_agent_default_shutdown_timeout(self):
+        mock_spec = Mock(spec=WorkerSpec)
+        mock_spec.max_restarts = 3
+
+        agent = SimpleElasticAgent(
+            spec=mock_spec,
+            exit_barrier_timeout=300,
+        )
+
+        self.assertEqual(agent._shutdown_timeout, 30)
+
+    def test_local_elastic_agent_receives_shutdown_timeout(self):
+        mock_spec = Mock(spec=WorkerSpec)
+        mock_spec.max_restarts = 3
+        mock_spec.rdzv_handler = Mock()
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            logs_specs = DefaultLogsSpecs(log_dir=tmpdir)
+
+            agent = LocalElasticAgent(
+                spec=mock_spec,
+                logs_specs=logs_specs,
+                start_method="spawn",
+                exit_barrier_timeout=300,
+                shutdown_timeout=240,
+            )
+
+            self.assertEqual(agent._shutdown_timeout, 240)
+
+    def test_shutdown_method_receives_correct_timeout(self):
+        mock_spec = Mock(spec=WorkerSpec)
+        mock_spec.max_restarts = 3
+        mock_spec.rdzv_handler = Mock()
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            logs_specs = DefaultLogsSpecs(log_dir=tmpdir)
+
+            agent = LocalElasticAgent(
+                spec=mock_spec,
+                logs_specs=logs_specs,
+                start_method="spawn",
+                exit_barrier_timeout=300,
+                shutdown_timeout=150,
+            )
+
+            # Mock the _pcontext to avoid actual process operations
+            mock_pcontext = MagicMock()
+            agent._pcontext = mock_pcontext
+
+            # Call _shutdown directly
+            agent._shutdown(death_sig=signal.SIGTERM, timeout=150)
+
+            # Verify close was called with correct timeout
+            mock_pcontext.close.assert_called_once_with(signal.SIGTERM, 150)
+
+    @patch("torch.distributed.elastic.multiprocessing.api.PContext")
+    def test_pcontext_close_receives_timeout(self, mock_pcontext_class):
+        mock_pcontext = MagicMock()
+        mock_pcontext_class.return_value = mock_pcontext
+
+        # Simulate calling close with a custom timeout
+        mock_pcontext.close(death_sig=signal.SIGTERM, timeout=200)
+
+        # Verify it was called with the correct timeout
+        mock_pcontext.close.assert_called_once_with(
+            death_sig=signal.SIGTERM, timeout=200
+        )
+
+    def test_shutdown_timeout_validation_negative(self):
+        with self.assertRaises(ValueError) as cm:
+            LaunchConfig(
+                min_nodes=1,
+                max_nodes=1,
+                nproc_per_node=2,
+                shutdown_timeout=-5,
+            )
+        self.assertIn("shutdown_timeout must be non-negative", str(cm.exception))
+
+    def test_shutdown_timeout_zero(self):
+        config = LaunchConfig(
+            min_nodes=1,
+            max_nodes=1,
+            nproc_per_node=2,
+            shutdown_timeout=0,
+        )
+        self.assertEqual(config.shutdown_timeout, 0)
+
+
+if __name__ == "__main__":
+    run_tests()
diff --git a/test/distributed/elastic/multiprocessing/api_test.py b/test/distributed/elastic/multiprocessing/api_test.py
index 4a7172388974c..6dd7938753598 100644
--- a/test/distributed/elastic/multiprocessing/api_test.py
+++ b/test/distributed/elastic/multiprocessing/api_test.py
@@ -15,9 +15,9 @@
 import sys
 import tempfile
 import time
+import unittest
 from collections.abc import Callable
 from itertools import product
-from typing import Union
 from unittest import mock
 
 import torch
@@ -43,6 +43,7 @@
     skip_if_pytest,
     TEST_WITH_ASAN,
     TEST_WITH_DEV_DBG_ASAN,
+    TEST_WITH_ROCM,
     TEST_WITH_TSAN,
     TestCase,
 )
@@ -195,7 +196,7 @@ def wait_fn(wait_time: int = 300) -> None:
 
 def start_processes_zombie_test(
     idx: int,
-    entrypoint: Union[str, Callable],
+    entrypoint: str | Callable,
     mp_queue: mp.Queue,
     log_dir: str,
     nproc: int = 2,
@@ -263,7 +264,7 @@ def assert_pids_noexist(self, pids: dict[int, int]):
                 os.kill(pid, 0)
 
     def _test_zombie_workflow(
-        self, entrypoint: Union[str, Callable], signal_to_send: signal.Signals
+        self, entrypoint: str | Callable, signal_to_send: signal.Signals
     ) -> None:
         mp_queue = mp.get_context("spawn").Queue()
         child_nproc = 2
@@ -448,6 +449,10 @@ def test_function_large_ret_val(self):
                     for i in range(pc.nprocs):
                         self.assertEqual(size, len(results.return_values[i]))
 
+        @unittest.skipIf(
+            TEST_WITH_ROCM,
+            "Skipped on ROCm due to hang in MultiprocessContext.wait after Kineto bump (PR #177101, 1fd9c49); investigating",
+        )
         def test_function_raise(self):
             """
             run 2x copies of echo2, raise an exception on the first
diff --git a/test/distributed/elastic/multiprocessing/errors/api_test.py b/test/distributed/elastic/multiprocessing/errors/api_test.py
index 5a3a28f988d12..da6128fffc8a9 100644
--- a/test/distributed/elastic/multiprocessing/errors/api_test.py
+++ b/test/distributed/elastic/multiprocessing/errors/api_test.py
@@ -247,3 +247,10 @@ def test_record_child_failure_no_child_error_file(self):
             # it SHOULD re-raise ChildFailedError for any upstream system
             # to handle it.
             self.assertFalse(os.path.isfile(self.test_error_file))
+
+    def test_child_failed_error_signal_name_in_message(self):
+        pf = self.failure_without_error_file(exitcode=-signal.SIGSEGV)
+        ex = ChildFailedError("trainer.par", {0: pf})
+        error_msg = str(ex)
+        self.assertIn("(SIGSEGV)", error_msg)
+        self.assertIn(f"exitcode  : {-signal.SIGSEGV}", error_msg)
diff --git a/test/distributed/elastic/rendezvous/dynamic_rendezvous_test.py b/test/distributed/elastic/rendezvous/dynamic_rendezvous_test.py
index 1554c11fd0773..084b262602ff2 100644
--- a/test/distributed/elastic/rendezvous/dynamic_rendezvous_test.py
+++ b/test/distributed/elastic/rendezvous/dynamic_rendezvous_test.py
@@ -16,7 +16,7 @@
 from base64 import b64encode
 from collections.abc import Callable
 from datetime import datetime, timedelta, timezone
-from typing import cast, Optional
+from typing import cast
 from unittest import TestCase
 from unittest.mock import call, MagicMock, Mock, patch, PropertyMock
 
@@ -176,7 +176,7 @@ def test_encoded_size_is_within_expected_limit(self) -> None:
 
 
 class FakeRendezvousBackend(RendezvousBackend):
-    _state: Optional[bytes]
+    _state: bytes | None
     _token: int
 
     def __init__(self) -> None:
@@ -187,15 +187,15 @@ def __init__(self) -> None:
     def name(self) -> str:
         return "fake_backend"
 
-    def get_state(self) -> Optional[tuple[bytes, Token]]:
+    def get_state(self) -> tuple[bytes, Token] | None:
         if self._token == 0:
             return None
 
         return self._state, self._token  # type: ignore[return-value]
 
     def set_state(
-        self, state: bytes, token: Optional[Token] = None
-    ) -> Optional[tuple[bytes, Token, bool]]:
+        self, state: bytes, token: Token | None = None
+    ) -> tuple[bytes, Token, bool] | None:
         if token is None:
             token = 0
 
@@ -513,7 +513,7 @@ def test_sync_raises_error_if_backend_state_is_corrupt(self) -> None:
 
 class FakeRendezvousStateHolder(_RendezvousStateHolder):
     _state: _RendezvousState
-    _dirty: Optional[bool]
+    _dirty: bool | None
 
     def __init__(self) -> None:
         self._state = _RendezvousState()
@@ -527,7 +527,7 @@ def state(self) -> _RendezvousState:
     def state(self, value) -> None:
         self._state = value
 
-    def sync(self) -> Optional[bool]:
+    def sync(self) -> bool | None:
         self._dirty, dirty = None, self._dirty
 
         return dirty
@@ -582,7 +582,7 @@ def _create_settings(self) -> RendezvousSettings:
         )
 
     def _create_op_executor(
-        self, settings: Optional[RendezvousSettings] = None
+        self, settings: RendezvousSettings | None = None
     ) -> _DistributedRendezvousOpExecutor:
         self._state_holder.state = self._state
 
@@ -1156,9 +1156,9 @@ def setUp(self) -> None:
         self._min_nodes = 1
         self._max_nodes = 1
 
-        self._join_timeout: Optional[timedelta] = None
-        self._close_timeout: Optional[timedelta] = None
-        self._heartbeat_timeout: Optional[timedelta] = None
+        self._join_timeout: timedelta | None = None
+        self._close_timeout: timedelta | None = None
+        self._heartbeat_timeout: timedelta | None = None
 
         self._keep_alive_interval = timedelta(seconds=30)
 
@@ -1505,7 +1505,7 @@ def setUp(self) -> None:
         self._backend = DummyRendezvousBackend()
         self._min_nodes = 3
         self._max_nodes = 6
-        self._timeout: Optional[RendezvousTimeout] = RendezvousTimeout()
+        self._timeout: RendezvousTimeout | None = RendezvousTimeout()
 
     def _create_handler(self) -> DynamicRendezvousHandler:
         return DynamicRendezvousHandler.from_backend(
@@ -1829,7 +1829,7 @@ def get(self, key):
             def set(self, key, value):
                 pass
 
-        prefix_store = CustomPrefixStore(spec=dist.PrefixStore)
+        prefix_store = CustomPrefixStore()
         prefix_store_class_mock.return_value = prefix_store
         tcp_store = Mock(spec=dist.TCPStore)
         original_addr = "original_addr"
diff --git a/test/distributed/elastic/rendezvous/etcd_rendezvous_backend_test.py b/test/distributed/elastic/rendezvous/etcd_rendezvous_backend_test.py
index d643ab2a407ad..2bb23c6f58bdd 100644
--- a/test/distributed/elastic/rendezvous/etcd_rendezvous_backend_test.py
+++ b/test/distributed/elastic/rendezvous/etcd_rendezvous_backend_test.py
@@ -176,5 +176,9 @@ def store_set(store):
         get_thread.join()
         set_thread.join()
 
-        assert result_dict["get_result"] == "foo"
-        assert result_dict["time"] >= 2
+        if result_dict["get_result"] != "foo":
+            raise AssertionError(
+                f"Expected get_result == 'foo', got {result_dict['get_result']}"
+            )
+        if not (result_dict["time"] >= 2):
+            raise AssertionError(f"Expected time >= 2, got {result_dict['time']}")
diff --git a/test/distributed/elastic/rendezvous/rendezvous_backend_test.py b/test/distributed/elastic/rendezvous/rendezvous_backend_test.py
index f4c0fd7b88db9..a9f78712d4adf 100644
--- a/test/distributed/elastic/rendezvous/rendezvous_backend_test.py
+++ b/test/distributed/elastic/rendezvous/rendezvous_backend_test.py
@@ -8,7 +8,7 @@
 
 from abc import ABC, abstractmethod
 from collections.abc import Callable
-from typing import Any, cast, Optional
+from typing import Any, cast
 
 from torch.distributed.elastic.rendezvous import RendezvousStateError
 from torch.distributed.elastic.rendezvous.dynamic_rendezvous import (
@@ -32,7 +32,7 @@ def _corrupt_state(self) -> None:
         """Corrupts the state stored in the backend."""
 
     def _set_state(
-        self, state: bytes, token: Optional[Any] = None
+        self, state: bytes, token: Any | None = None
     ) -> tuple[bytes, Token, bool]:
         result = self._backend.set_state(state, token)
 
diff --git a/test/distributed/elastic/utils/distributed_test.py b/test/distributed/elastic/utils/distributed_test.py
index 8b38d365d9560..eab5055509e5e 100644
--- a/test/distributed/elastic/utils/distributed_test.py
+++ b/test/distributed/elastic/utils/distributed_test.py
@@ -144,7 +144,8 @@ def test_create_store_with_libuv_support(self):
         )
         self.assertFalse(store.libuvBackend)
         del os.environ["USE_LIBUV"]
-        assert "USE_LIBUV" not in os.environ
+        if "USE_LIBUV" in os.environ:
+            raise AssertionError("Expected USE_LIBUV to be removed from os.environ")
 
         # libuv backend is enabled by default
         store = create_c10d_store(
diff --git a/test/distributed/flight_recorder/test_fr_analysis.py b/test/distributed/flight_recorder/test_fr_analysis.py
index 6b7b0dcada162..e68e7195371fb 100644
--- a/test/distributed/flight_recorder/test_fr_analysis.py
+++ b/test/distributed/flight_recorder/test_fr_analysis.py
@@ -10,6 +10,7 @@
     COLLECTIVES,
     MatchInfo,
     MatchState,
+    Op,
 )
 from torch.distributed.flight_recorder.components.utils import match_one_event
 from torch.testing._internal.common_utils import run_tests, TestCase
@@ -178,6 +179,45 @@ def test_all_events(self):
             self.assertEqual(result, expectedState)
 
 
+class FlightRecorderOpBackendTest(TestCase):
+    """Tests that the Op class accepts all supported backend prefixes."""
+
+    def _make_event(self, backend: str, collective: str = "all_reduce"):
+        return {
+            "profiling_name": f"{backend}:{collective}",
+            "state": "completed",
+            "process_group": ("0", "default"),
+            "input_sizes": [[4, 4]],
+            "output_sizes": [[4, 4]],
+            "input_dtypes": "float32",
+            "output_dtypes": "float32",
+            "collective_seq_id": "1",
+            "p2p_seq_id": "0",
+            "time_created_ns": 0,
+            "frames": [],
+        }
+
+    def test_nccl_backend(self):
+        op = Op(self._make_event("nccl"), {"0": {0, 1}}, "0")
+        self.assertEqual(op.type, "all_reduce")
+
+    def test_ncclx_backend(self):
+        op = Op(self._make_event("ncclx"), {"0": {0, 1}}, "0")
+        self.assertEqual(op.type, "all_reduce")
+
+    def test_gloo_backend(self):
+        op = Op(self._make_event("gloo"), {"0": {0, 1}}, "0")
+        self.assertEqual(op.type, "all_reduce")
+
+    def test_xccl_backend(self):
+        op = Op(self._make_event("xccl"), {"0": {0, 1}}, "0")
+        self.assertEqual(op.type, "all_reduce")
+
+    def test_unsupported_backend_raises(self):
+        with self.assertRaises(AssertionError):
+            Op(self._make_event("unknown_backend"), {"0": {0, 1}}, "0")
+
+
 class FlightMatchInfoTest(TestCase):
     def test_match_info(self):
         m1 = MatchInfo(MatchState.FULLY_MATCHED, "rank 0")
diff --git a/test/distributed/fsdp/test_checkpoint_wrapper.py b/test/distributed/fsdp/test_checkpoint_wrapper.py
index 8cc5698cd19aa..95c2a5f28ecaf 100644
--- a/test/distributed/fsdp/test_checkpoint_wrapper.py
+++ b/test/distributed/fsdp/test_checkpoint_wrapper.py
@@ -389,6 +389,34 @@ def dfs(grad_fn):
 
         torch.autograd.graph.saved_tensors_hooks.__init__ = orig_init
 
+    @unittest.skipIf(not torch.cuda.is_available(), "requires CUDA")
+    def test_checkpoint_wrapper_with_block_mask(self):
+        """Test that BlockMask can be passed through checkpoint_wrapper."""
+        from torch.nn.attention.flex_attention import BlockMask, create_block_mask
+        from torch.utils._pytree import register_pytree_node, SUPPORTED_NODES
+
+        if BlockMask not in SUPPORTED_NODES:
+            register_pytree_node(
+                BlockMask,
+                BlockMask._flatten,
+                BlockMask._unflatten,
+                flatten_with_keys_fn=BlockMask._flatten_with_keys,
+                serialized_type_name="torch.nn.attention.flex_attention.BlockMask",
+            )
+
+        block_mask = create_block_mask(
+            lambda b, h, q, kv: q >= kv, B=1, H=1, Q_LEN=128, KV_LEN=128
+        )
+
+        class Block(nn.Module):
+            def forward(self, x, mask):
+                return x * 2
+
+        model = checkpoint_wrapper(Block()).cuda()
+        x = torch.randn(4, 128, device="cuda")
+        result = model(x, block_mask)
+        self.assertEqual(result, x * 2)
+
 
 if __name__ == "__main__":
     run_tests()
diff --git a/test/distributed/fsdp/test_fsdp_checkpoint.py b/test/distributed/fsdp/test_fsdp_checkpoint.py
index 26936ebf5d61d..3c30065da1138 100644
--- a/test/distributed/fsdp/test_fsdp_checkpoint.py
+++ b/test/distributed/fsdp/test_fsdp_checkpoint.py
@@ -112,9 +112,12 @@ def forward(self, x):
             return self.ffn(x)
 
     def _verify_parity(self, losses, outputs, models):
-        assert losses
-        assert outputs
-        assert models
+        if not losses:
+            raise AssertionError("Expected losses to be non-empty")
+        if not outputs:
+            raise AssertionError("Expected outputs to be non-empty")
+        if not models:
+            raise AssertionError("Expected models to be non-empty")
         for l, o in zip(losses[1:], outputs[1:]):
             self.assertEqual(losses[0], l)
             self.assertEqual(outputs[0], o)
diff --git a/test/distributed/fsdp/test_fsdp_clip_grad_norm.py b/test/distributed/fsdp/test_fsdp_clip_grad_norm.py
index 07e5a08493c8d..5745e17a64374 100644
--- a/test/distributed/fsdp/test_fsdp_clip_grad_norm.py
+++ b/test/distributed/fsdp/test_fsdp_clip_grad_norm.py
@@ -1,7 +1,6 @@
 # Owner(s): ["oncall: distributed"]
 import itertools
 import sys
-from typing import Union
 
 import torch
 import torch.nn as nn
@@ -94,9 +93,9 @@ def test_ddp_parity(self, device):
     def _test_ddp_parity(
         self,
         device,
-        max_norm: Union[float, int],
-        norm_type: Union[float, int],
-        sharding_strategy: Union[ShardingStrategy, str],
+        max_norm: float | int,
+        norm_type: float | int,
+        sharding_strategy: ShardingStrategy | str,
         use_orig_params: bool,
         offload_params: bool,
     ):
@@ -193,12 +192,18 @@ def _test_ddp_parity(
         self.assertEqual(ddp_total_norm, fsdp_total_norm)
         # Check that the gradients were modified by `clip_grad_norm_()`
         for param, orig_grad in zip(ddp_model.parameters(), orig_ddp_grads):
-            assert not torch.equal(param.grad, orig_grad)
+            if torch.equal(param.grad, orig_grad):
+                raise AssertionError(
+                    "Expected gradient to be modified by clip_grad_norm_()"
+                )
         for param, orig_grad in zip(fsdp_model.parameters(), orig_fsdp_grads):
             if param.grad is None:
                 self.assertEqual(param.grad, orig_grad)  # `None`
             else:
-                assert not torch.equal(param.grad, orig_grad)
+                if torch.equal(param.grad, orig_grad):
+                    raise AssertionError(
+                        "Expected gradient to be modified by clip_grad_norm_()"
+                    )
         # Run an optimizer step to ensure gradients matched after clipping
         ddp_optim.step()
         fsdp_optim.step()
@@ -257,8 +262,8 @@ def test_low_precision_grads(self, device):
     def _test_low_precision_grads(
         self,
         device,
-        max_norm: Union[float, int],
-        norm_type: Union[float, int],
+        max_norm: float | int,
+        norm_type: float | int,
         sharding_strategy: ShardingStrategy,
         use_orig_params: bool,
     ):
diff --git a/test/distributed/fsdp/test_fsdp_comm.py b/test/distributed/fsdp/test_fsdp_comm.py
index 3be521d8f3674..d4fa4073fd9ac 100644
--- a/test/distributed/fsdp/test_fsdp_comm.py
+++ b/test/distributed/fsdp/test_fsdp_comm.py
@@ -2,7 +2,6 @@
 import sys
 from contextlib import nullcontext
 from enum import auto, Enum
-from typing import Optional
 from unittest.mock import patch
 
 import torch
@@ -105,7 +104,7 @@ def _get_ref_num_reduce_scatters(
     def _get_ref_num_all_gathers(
         self,
         num_fsdp: int,
-        sharding_strategy: Optional[ShardingStrategy],
+        sharding_strategy: ShardingStrategy | None,
         is_first_iter: bool,
         is_last_iter_no_sync: bool,
     ) -> int:
@@ -125,7 +124,7 @@ def _get_ref_num_all_gathers(
     def _get_ref_num_all_gathers_in_pass(
         self,
         num_fsdp: int,
-        sharding_strategy: Optional[ShardingStrategy],
+        sharding_strategy: ShardingStrategy | None,
         pass_type: PassType,
         is_first_iter: bool,
         is_last_iter_no_sync: bool,
@@ -162,7 +161,7 @@ def _get_ref_num_all_gathers_in_pass(
             # forward pass
             num_all_gathers = 0
         else:
-            assert 0, (
+            raise AssertionError(
                 f"Unsupported: add a branch for pass_type={pass_type} "
                 f"is_first_iter={is_first_iter} "
                 f"is_last_iter_no_sync={is_last_iter_no_sync} "
@@ -211,7 +210,7 @@ def test_communication(
         device,
         nested_model: bool,
         use_no_sync: bool,
-        sharding_strategy: Optional[ShardingStrategy],
+        sharding_strategy: ShardingStrategy | None,
     ):
         """
         Tests FSDP's communication cost in terms of calls to collective
diff --git a/test/distributed/fsdp/test_fsdp_comm_hooks.py b/test/distributed/fsdp/test_fsdp_comm_hooks.py
index 6a204d5eb22fe..7f0cc4cdb7ed9 100644
--- a/test/distributed/fsdp/test_fsdp_comm_hooks.py
+++ b/test/distributed/fsdp/test_fsdp_comm_hooks.py
@@ -1,7 +1,6 @@
 # Owner(s): ["oncall: distributed"]
 
 import sys
-from typing import Optional
 
 import torch
 import torch.nn as nn
@@ -120,7 +119,7 @@ class TestCommunicationHooks(FSDPTest):
         ],
     )
     def test_default_communication_hook_behavior(
-        self, sharding_strategy: Optional[ShardingStrategy]
+        self, sharding_strategy: ShardingStrategy | None
     ):
         """
         Tests FSDP's default communication hook's behavior and correctness.
@@ -192,7 +191,7 @@ def _init_model(self, core, sharding_strategy, mixed_precision=None):
         ],
     )
     def test_default_communication_hook_initialization(
-        self, has_wrapping: bool, sharding_strategy: Optional[ShardingStrategy]
+        self, has_wrapping: bool, sharding_strategy: ShardingStrategy | None
     ):
         """
         Tests FSDP's communication hook interface behavior.
@@ -242,7 +241,7 @@ def test_default_communication_hook_initialization(
         ],
     )
     def test_registering_hook_non_root(
-        self, sharding_strategy: Optional[ShardingStrategy]
+        self, sharding_strategy: ShardingStrategy | None
     ):
         """
         Tests FSDP's communication hook registering for submodules.
@@ -302,7 +301,7 @@ def test_registering_hook_hybrid_strategy(self):
         ],
     )
     def test_registering_hook_submodules(
-        self, sharding_strategy: Optional[ShardingStrategy]
+        self, sharding_strategy: ShardingStrategy | None
     ):
         """
         Tests FSDP's communication hook registering for submodules.
@@ -391,7 +390,7 @@ def _check_low_precision_hook(
         ],
     )
     def test_fp16_hook(
-        self, has_wrapping: bool, sharding_strategy: Optional[ShardingStrategy]
+        self, has_wrapping: bool, sharding_strategy: ShardingStrategy | None
     ):
         state = default_hooks.LowPrecisionState(process_group=_get_default_group())
         hook = default_hooks.fp16_compress_hook
@@ -417,7 +416,7 @@ def test_fp16_hook(
         ],
     )
     def test_bf16_hook(
-        self, has_wrapping: bool, sharding_strategy: Optional[ShardingStrategy]
+        self, has_wrapping: bool, sharding_strategy: ShardingStrategy | None
     ):
         state = default_hooks.LowPrecisionState(process_group=_get_default_group())
         hook = default_hooks.bf16_compress_hook
diff --git a/test/distributed/fsdp/test_fsdp_core.py b/test/distributed/fsdp/test_fsdp_core.py
index 45e58ea3d0491..0741ebad5ba05 100644
--- a/test/distributed/fsdp/test_fsdp_core.py
+++ b/test/distributed/fsdp/test_fsdp_core.py
@@ -5,7 +5,7 @@
 import sys
 import unittest
 from collections.abc import Callable
-from typing import Any, Optional
+from typing import Any
 from unittest import mock
 
 import torch
@@ -109,7 +109,7 @@ def _get_subtest_config(self, cpu_offload: CPUOffload) -> dict[str, list[Any]]:
     def test_nested_wrapped_model(
         self,
         cpu_offload: CPUOffload,
-        sharding_strategy: Optional[ShardingStrategy],
+        sharding_strategy: ShardingStrategy | None,
     ):
         self.run_subtests(
             self._get_subtest_config(cpu_offload),
@@ -125,7 +125,7 @@ def test_nested_wrapped_model(
     def test_nested_wrapped_model_single_iteration_mixed_precision(
         self,
         cpu_offload: CPUOffload,
-        sharding_strategy: Optional[ShardingStrategy],
+        sharding_strategy: ShardingStrategy | None,
     ):
         mixed_precision = MixedPrecision(
             param_dtype=torch.float16,
@@ -148,7 +148,7 @@ def test_nested_wrapped_model_single_iteration_mixed_precision(
     def test_nested_always_wrap_model(
         self,
         cpu_offload: CPUOffload,
-        sharding_strategy: Optional[ShardingStrategy],
+        sharding_strategy: ShardingStrategy | None,
     ):
         self.run_subtests(
             self._get_subtest_config(cpu_offload),
@@ -164,7 +164,7 @@ def test_nested_always_wrap_model(
     def test_transformer(
         self,
         cpu_offload: CPUOffload,
-        sharding_strategy: Optional[ShardingStrategy],
+        sharding_strategy: ShardingStrategy | None,
     ):
         self.run_subtests(
             self._get_subtest_config(cpu_offload),
@@ -180,7 +180,7 @@ def test_transformer(
     def test_delayed_optim_step(
         self,
         cpu_offload: CPUOffload,
-        sharding_strategy: Optional[ShardingStrategy],
+        sharding_strategy: ShardingStrategy | None,
     ):
         """Tests the FSDP forward, backward, and optimizer step runtime by
         using a model with a long CUDA delay after the loss computation/before
@@ -202,7 +202,7 @@ def test_delayed_optim_step(
     def test_delayed_reduce_scatter(
         self,
         cpu_offload: CPUOffload,
-        sharding_strategy: Optional[ShardingStrategy],
+        sharding_strategy: ShardingStrategy | None,
     ):
         """Tests the FSDP forward, backward, and optimizer step runtime by
         using a model with a long CUDA delay before the gradient reduce-scatter
@@ -228,7 +228,7 @@ def _dummy_ddp_fn(self, model):
     def test_mixture_of_experts(
         self,
         cpu_offload: CPUOffload,
-        sharding_strategy: Optional[ShardingStrategy],
+        sharding_strategy: ShardingStrategy | None,
     ):
         fsdp_kwargs = {"device_id": device_type.type}
         self.run_subtests(
@@ -250,7 +250,7 @@ def test_mixture_of_experts(
     def test_mixture_of_experts_with_delay_before_free(
         self,
         cpu_offload: CPUOffload,
-        sharding_strategy: Optional[ShardingStrategy],
+        sharding_strategy: ShardingStrategy | None,
     ):
         fsdp_kwargs = {"device_id": device_type.type}
         self.run_subtests(
@@ -471,7 +471,7 @@ def _test_unshard_params_as_tensors(
         sharding_strategy: ShardingStrategy,
         use_orig_params: bool,
         forward_prefetch: bool,
-        backward_prefetch: Optional[BackwardPrefetch],
+        backward_prefetch: BackwardPrefetch | None,
     ):
         orig_use_unsharded_views = FlatParamHandle._use_unsharded_views
 
diff --git a/test/distributed/fsdp/test_fsdp_grad_acc.py b/test/distributed/fsdp/test_fsdp_grad_acc.py
index 0cdc1849736e7..650d0e71c44e0 100644
--- a/test/distributed/fsdp/test_fsdp_grad_acc.py
+++ b/test/distributed/fsdp/test_fsdp_grad_acc.py
@@ -4,7 +4,7 @@
 import itertools
 import sys
 from dataclasses import dataclass
-from typing import Any, Optional
+from typing import Any
 
 import torch
 from torch import distributed as dist
@@ -91,7 +91,7 @@ def _test_grad_acc(
         batch_dim: int,
         configs: list[_GradAccConfig],
         cpu_offload: CPUOffload,
-        backward_prefetch: Optional[BackwardPrefetch],
+        backward_prefetch: BackwardPrefetch | None,
         sharding_strategy: ShardingStrategy,
         use_orig_params: bool,
     ):
@@ -152,9 +152,10 @@ def permute_tensor(x: torch.Tensor):
             batches.append(tuple(permute_tensor(t) for t in batch))
         for batch1, batch2 in itertools.combinations(batches, r=2):
             for t1, t2 in zip(batch1, batch2):
-                assert not torch.all(t1 == t2), (
-                    "Check the test to make sure that batches are distinct"
-                )
+                if torch.all(t1 == t2):
+                    raise AssertionError(
+                        "Check the test to make sure that batches are distinct"
+                    )
 
         # Concatenate the batches along the given batch dimension
         concat_batch: tuple[torch.Tensor, ...] = tuple(
diff --git a/test/distributed/fsdp/test_fsdp_hybrid_shard.py b/test/distributed/fsdp/test_fsdp_hybrid_shard.py
index 153f43432eb0b..3479cea0c564e 100644
--- a/test/distributed/fsdp/test_fsdp_hybrid_shard.py
+++ b/test/distributed/fsdp/test_fsdp_hybrid_shard.py
@@ -29,6 +29,7 @@
     DEVICEInitMode,
     FSDPInitMode,
     FSDPTestContinuous,
+    FSDPTestMultiThread,
     TransformerWithSharedParams,
 )
 from torch.testing._internal.common_utils import (
@@ -157,8 +158,14 @@ def test_hsdp_save_load_state_dict(self):
         optim.step()
         shard_g = model.process_group
         replicate_g = model._inter_node_pg
-        assert shard_g == my_shard_group
-        assert replicate_g == my_replicate_group
+        if shard_g != my_shard_group:
+            raise AssertionError(
+                f"Expected shard_g == my_shard_group, got {shard_g} vs {my_shard_group}"
+            )
+        if replicate_g != my_replicate_group:
+            raise AssertionError(
+                f"Expected replicate_g == my_replicate_group, got {replicate_g} vs {my_replicate_group}"
+            )
         with FSDP.state_dict_type(model, StateDictType.SHARDED_STATE_DICT):
             msd = model.state_dict()
             osd = FSDP.optim_state_dict(model, optim)
@@ -360,9 +367,8 @@ def _test_fsdp_hybrid_shard_parity(
             use_orig_params,
             hsdp_process_groups=hsdp_pgs,
         )
-        assert hsdp_model._inter_node_pg.size() > 1, (
-            "HSDP model initialized without replication"
-        )
+        if not (hsdp_model._inter_node_pg.size() > 1):
+            raise AssertionError("HSDP model initialized without replication")
         fsdp_optim = torch.optim.Adam(fsdp_model.parameters(), lr=1e-2)
         hsdp_optim = torch.optim.Adam(hsdp_model.parameters(), lr=1e-2)
         torch.manual_seed(global_pg.rank() + 1)
@@ -400,12 +406,13 @@ def _init_hsdp_model(
         hsdp_sharding_strategy: ShardingStrategy,
         sharding_strategy_mode: str,
         use_orig_params: bool,
-        hsdp_process_groups: Optional[
-            tuple[dist.ProcessGroup, dist.ProcessGroup]
-        ] = None,
+        hsdp_process_groups: tuple[dist.ProcessGroup, dist.ProcessGroup] | None = None,
         hsdp_device_mesh: Optional = None,
     ):
-        assert hsdp_process_groups is None or hsdp_device_mesh is None
+        if not (hsdp_process_groups is None or hsdp_device_mesh is None):
+            raise AssertionError(
+                "Expected hsdp_process_groups or hsdp_device_mesh to be None"
+            )
         auto_wrap_policy = ModuleWrapPolicy(
             {TransformerEncoderLayer, TransformerDecoderLayer},
         )
@@ -444,6 +451,69 @@ def _init_hsdp_model(
         return hsdp_model
 
 
+class TestHSDPSyncModuleStates(FSDPTestMultiThread):
+    class _ModelWithBuffer(nn.Module):
+        def __init__(self, device):
+            super().__init__()
+            self.lin = nn.Linear(10, 10, device=device)
+            self.bn = nn.BatchNorm1d(10, device=device)
+
+        def forward(self, x):
+            return self.bn(self.lin(x))
+
+    @property
+    def world_size(self) -> int:
+        return 4
+
+    @skip_if_lt_x_gpu(1)
+    def test_hsdp_buffer_sync_from_meta_device(self):
+        """Test that HSDP sync_module_states correctly broadcasts buffers
+        when only rank 0 has real weights and other ranks use meta device.
+
+        _sync_module_params_and_buffers marks buffers with FSDP_SYNCED=True
+        to avoid redundant syncs in nested wrapping. With two-phase broadcast
+        (inter-node then intra-node), this flag must be reset between phases,
+        otherwise the intra-node broadcast skips buffers. Non-persistent
+        buffers (e.g. RoPE inv_freq) that are materialized from meta device
+        via to_empty() and not restored by reset_parameters() would remain
+        as uninitialized values on local ranks 1..N.
+
+        Parameters are unaffected by broadcast order because they are
+        unconditionally included in every sync call. This test specifically
+        targets buffers, which are the only tensors gated by FSDP_SYNCED.
+        """
+        dev = torch.device(device_type)
+        mesh_2d = init_device_mesh(device_type, (2, self.world_size // 2))
+
+        if self.rank == 0:
+            model = self._ModelWithBuffer(device=dev)
+            model.bn.running_mean.fill_(42.0)
+            model.bn.running_var.fill_(7.0)
+        else:
+            model = self._ModelWithBuffer(device="meta")
+
+        def param_init_fn(module):
+            if any(p.is_meta for p in module.parameters(recurse=False)) or any(
+                b.is_meta for b in module.buffers(recurse=False)
+            ):
+                module.to_empty(device=dev, recurse=False)
+
+        model = FSDP(
+            model,
+            device_mesh=mesh_2d,
+            sharding_strategy=ShardingStrategy.HYBRID_SHARD,
+            use_orig_params=True,
+            sync_module_states=True,
+            param_init_fn=param_init_fn,
+        )
+
+        # Buffers are not sharded by FSDP, so each rank has a local copy.
+        # If the FSDP_SYNCED reset is missing, ranks on non-source nodes
+        # will have uninitialized buffer values instead of rank 0's values.
+        self.assertEqual(model.bn.running_mean, torch.full((10,), 42.0))
+        self.assertEqual(model.bn.running_var, torch.full((10,), 7.0))
+
+
 instantiate_parametrized_tests(TestFSDPHybridShard)
 
 if __name__ == "__main__":
diff --git a/test/distributed/fsdp/test_fsdp_ignored_modules.py b/test/distributed/fsdp/test_fsdp_ignored_modules.py
index 302cddfd0d5d6..774483129325a 100644
--- a/test/distributed/fsdp/test_fsdp_ignored_modules.py
+++ b/test/distributed/fsdp/test_fsdp_ignored_modules.py
@@ -83,7 +83,8 @@ class ModelWithIgnoredModules(Model):
     """Adds a variable number of :class:`IgnoredModule` to ``self.layer1``."""
 
     def __init__(self, num_ignored: int) -> None:
-        assert num_ignored >= 0
+        if not (num_ignored >= 0):
+            raise AssertionError(f"Expected num_ignored >= 0, got {num_ignored}")
         super().__init__()
         layer1_modules = (
             [torch.nn.Linear(5, 4), torch.nn.Linear(4, 4)]
diff --git a/test/distributed/fsdp/test_fsdp_input.py b/test/distributed/fsdp/test_fsdp_input.py
index 6916df403c806..1aafac4ed6bf0 100644
--- a/test/distributed/fsdp/test_fsdp_input.py
+++ b/test/distributed/fsdp/test_fsdp_input.py
@@ -47,7 +47,10 @@ def forward(self, input):
                 if isinstance(input, list):
                     input = input[0]
                 else:
-                    assert isinstance(input, dict), input
+                    if not isinstance(input, dict):
+                        raise AssertionError(
+                            f"Expected dict, got {type(input)}: {input}"
+                        )
                     input = input["in"]
                 return self.layer(input)
 
diff --git a/test/distributed/fsdp/test_fsdp_meta.py b/test/distributed/fsdp/test_fsdp_meta.py
index a298cf523e356..9d3196a7eb503 100644
--- a/test/distributed/fsdp/test_fsdp_meta.py
+++ b/test/distributed/fsdp/test_fsdp_meta.py
@@ -2,7 +2,6 @@
 
 import itertools
 import sys
-from typing import Union
 
 import torch
 import torch.distributed as dist
@@ -129,7 +128,8 @@ def _init_with_torchdistX(module: nn.Module):
     torchdistX-based deferred module initialization function example
     using ``materialize_module``.
     """
-    assert _TORCHDISTX_AVAIL
+    if not _TORCHDISTX_AVAIL:
+        raise AssertionError("Expected _TORCHDISTX_AVAIL to be True")
 
     def check_fn(k):
         return not isinstance(k, FSDP)
@@ -373,7 +373,7 @@ def test_meta_device_with_mixed_precision(self):
 
         class FakeLinear(nn.Module):
             def __init__(
-                self, in_dim: int, out_dim: int, device: Union[torch.device, str]
+                self, in_dim: int, out_dim: int, device: torch.device | str
             ) -> None:
                 super().__init__()
                 self.weight = nn.Parameter(
diff --git a/test/distributed/fsdp/test_fsdp_misc.py b/test/distributed/fsdp/test_fsdp_misc.py
index 607e660757504..747280155f6ee 100644
--- a/test/distributed/fsdp/test_fsdp_misc.py
+++ b/test/distributed/fsdp/test_fsdp_misc.py
@@ -108,7 +108,10 @@ def _check_device_matches(module, device_id):
             devices = {
                 p.device for p in module.parameters() if isinstance(p, FlatParameter)
             }
-            assert len(devices) > 0
+            if not (len(devices) > 0):
+                raise AssertionError(
+                    f"Expected at least one device, but got {len(devices)}"
+                )
             self.assertEqual(1, len(devices))
             found_device = devices.pop()
             if use_index and not isinstance(device_id, torch.device):
@@ -219,8 +222,14 @@ def forward(self, x, y):
                 opt.step()
                 grads.append(x.grad)
                 opt.zero_grad()
-            assert torch.allclose(losses[0], losses[1])
-            assert torch.allclose(grads[0], grads[1])
+            if not torch.allclose(losses[0], losses[1]):
+                raise AssertionError(
+                    f"Expected losses to be close: {losses[0]} vs {losses[1]}"
+                )
+            if not torch.allclose(grads[0], grads[1]):
+                raise AssertionError(
+                    f"Expected grads to be close: {grads[0]} vs {grads[1]}"
+                )
             losses.clear()
             grads.clear()
 
@@ -232,7 +241,10 @@ def forward(self, x, y):
                 y = torch.randint(low=0, high=9, size=(8,), device=device_type)
                 fsdp_loss = fsdp_model(x, y)
                 ddp_loss = ddp_model(x, y)
-                assert torch.allclose(fsdp_loss, ddp_loss)
+                if not torch.allclose(fsdp_loss, ddp_loss):
+                    raise AssertionError(
+                        f"Expected fsdp_loss and ddp_loss to be close: {fsdp_loss} vs {ddp_loss}"
+                    )
 
         fsdp_model.train()
         ddp_model.train()
@@ -249,8 +261,14 @@ def forward(self, x, y):
                 opt.step()
                 grads.append(x.grad)
                 opt.zero_grad()
-            assert torch.allclose(losses[0], losses[1])
-            assert torch.allclose(grads[0], grads[1])
+            if not torch.allclose(losses[0], losses[1]):
+                raise AssertionError(
+                    f"Expected losses to be close: {losses[0]} vs {losses[1]}"
+                )
+            if not torch.allclose(grads[0], grads[1]):
+                raise AssertionError(
+                    f"Expected grads to be close: {grads[0]} vs {grads[1]}"
+                )
             losses.clear()
             grads.clear()
 
@@ -437,7 +455,10 @@ def test_fsdp_optimizer_overlap(self):
                 register_hook=False,
             )
             for p in fsdp_overlap.parameters():
-                assert hasattr(p, "_in_backward_optimizers")
+                if not hasattr(p, "_in_backward_optimizers"):
+                    raise AssertionError(
+                        "Expected parameter to have '_in_backward_optimizers' attribute"
+                    )
             optim = optim_cls(fsdp.parameters(), **optim_kwargs)
 
             # Verify params initially equal
diff --git a/test/distributed/fsdp/test_fsdp_mixed_precision.py b/test/distributed/fsdp/test_fsdp_mixed_precision.py
index b4532a86e3052..3f83ec3f2e562 100644
--- a/test/distributed/fsdp/test_fsdp_mixed_precision.py
+++ b/test/distributed/fsdp/test_fsdp_mixed_precision.py
@@ -764,11 +764,13 @@ def __init__(self) -> None:
 
             def forward(self, x, expect_use_full_prec_in_eval):
                 if expect_use_full_prec_in_eval:
-                    assert x.dtype == torch.float32, f"Expected fp32, got {x.dtype}"
+                    if x.dtype != torch.float32:
+                        raise AssertionError(f"Expected fp32, got {x.dtype}")
                 else:
-                    assert x.dtype == low_prec_dtype, (
-                        f"Expected {low_prec_dtype}, got {x.dtype}"
-                    )
+                    if x.dtype != low_prec_dtype:
+                        raise AssertionError(
+                            f"Expected {low_prec_dtype}, got {x.dtype}"
+                        )
                 return self.a(x)
 
         mp_config = MixedPrecision(
diff --git a/test/distributed/fsdp/test_fsdp_optim_state.py b/test/distributed/fsdp/test_fsdp_optim_state.py
index 99e5db33d67d7..a5eda167241b6 100644
--- a/test/distributed/fsdp/test_fsdp_optim_state.py
+++ b/test/distributed/fsdp/test_fsdp_optim_state.py
@@ -5,7 +5,7 @@
 from collections.abc import Callable
 from copy import deepcopy
 from enum import auto, Enum
-from typing import Any, Optional
+from typing import Any
 
 import torch
 import torch.nn as nn
@@ -78,7 +78,8 @@ class Bias(torch.nn.Module):
 
     def __init__(self, dim: int) -> None:
         super().__init__()
-        assert dim > 0
+        if not (dim > 0):
+            raise AssertionError(f"Expected dim > 0, but got {dim}")
         torch.manual_seed(0)
         self.bias = torch.nn.Parameter(torch.randn((dim,)))
 
@@ -99,7 +100,10 @@ class BlockA(torch.nn.Module):
 
     def __init__(self, in_dim: int, out_dim: int) -> None:
         super().__init__()
-        assert all(v > 0 for v in (in_dim, out_dim))
+        if not all(v > 0 for v in (in_dim, out_dim)):
+            raise AssertionError(
+                f"Expected in_dim > 0 and out_dim > 0, but got in_dim={in_dim}, out_dim={out_dim}"
+            )
         torch.manual_seed(0)
         self.bias_module0 = Bias(out_dim)
         self.weight = torch.nn.Parameter(torch.randn((in_dim, out_dim)))
@@ -127,7 +131,10 @@ class BlockB(torch.nn.Module):
 
     def __init__(self, in_dim: int, out_dim: int) -> None:
         super().__init__()
-        assert all(v > 0 for v in (in_dim, out_dim))
+        if not all(v > 0 for v in (in_dim, out_dim)):
+            raise AssertionError(
+                f"Expected in_dim > 0 and out_dim > 0, but got in_dim={in_dim}, out_dim={out_dim}"
+            )
         torch.manual_seed(0)
         self.weight = torch.nn.Parameter(torch.randn((in_dim, out_dim)))
         self.bias_module0 = Bias(out_dim)
@@ -175,9 +182,9 @@ def run_backward(self, loss):
     @staticmethod
     def wrap(
         model: torch.nn.Module,
-        group: Optional[dist.ProcessGroup] = None,
+        group: dist.ProcessGroup | None = None,
         ignore_modules: bool = False,
-        fsdp_kwargs: Optional[dict[str, Any]] = None,
+        fsdp_kwargs: dict[str, Any] | None = None,
     ) -> torch.nn.Module:
         if fsdp_kwargs is None:
             fsdp_kwargs = {}
@@ -213,8 +220,8 @@ def wrap(
     @staticmethod
     def wrap_alt(
         model: torch.nn.Module,
-        group: Optional[dist.ProcessGroup] = None,
-        fsdp_kwargs: Optional[dict[str, Any]] = None,
+        group: dist.ProcessGroup | None = None,
+        fsdp_kwargs: dict[str, Any] | None = None,
     ) -> torch.nn.Module:
         if fsdp_kwargs is None:
             fsdp_kwargs = {}
@@ -259,12 +266,16 @@ def add_unmanaged_param_entry(osd, unmanaged_param, step) -> None:
         for i in range(1, len(param_ids)):
             diff = param_ids[i] - param_ids[i - 1]
             if diff != 1:
-                assert diff > 1, f"Invalid IDs: {param_ids[i - 1]} {param_ids[i]}"
+                if not (diff > 1):
+                    raise AssertionError(
+                        f"Invalid IDs: {param_ids[i - 1]} {param_ids[i]}"
+                    )
                 unmanaged_param_id = param_ids[i - 1] + 1
                 break
         if unmanaged_param_id == -1:
             unmanaged_param_id = len(param_ids)  # last ID skipped
-        assert unmanaged_param_id >= 0, "One parameter ID should be skipped"
+        if not (unmanaged_param_id >= 0):
+            raise AssertionError("One parameter ID should be skipped")
         # Add a state entry for the unmanaged parameter
         state_device = next(iter(next(iter(osd["state"].values())).values())).device
         osd["state"][unmanaged_param_id] = {
@@ -325,7 +336,7 @@ def _init_nested_model(
         optim_class: type[torch.optim.Optimizer] = torch.optim.Adam,
         use_multiple_param_groups: bool = False,
         use_diff_optim_inputs: bool = False,
-        fsdp_kwargs: Optional[dict[str, Any]] = None,
+        fsdp_kwargs: dict[str, Any] | None = None,
     ):
         model = NestedModel().to(device)
         if wrap:
@@ -424,7 +435,8 @@ def _are_equal_states(
             if type(value1) is not type(value2):
                 return False
             if torch.is_tensor(value1):  # tensor state
-                assert torch.is_tensor(value2)
+                if not torch.is_tensor(value2):
+                    raise AssertionError("Expected value2 to be a tensor")
                 # Check the values on CPU to be device-agnostic
                 value1 = value1.cpu()
                 value2 = value2.cpu()
@@ -447,7 +459,8 @@ def _check_same_state(
         If ``check_same_param_keys=True``, then checks that the parameter keys
         match (e.g. when both should be parameter names), and does not check
         the parameter keys otherwise."""
-        assert "state" in ref_osd
+        if "state" not in ref_osd:
+            raise AssertionError("Expected 'state' in ref_osd")
         self.assertTrue("state" in fsdp_osd)
         ref_osd_state = ref_osd["state"]
         fsdp_osd_state = {
@@ -496,7 +509,8 @@ def _check_same_param_groups(
         "param_groups" part. If ``check_same_param_keys=True`, then checks that
         the parameter keys match (e.g. when both should be parameter names),
         and does not check the parameter keys otherwise."""
-        assert "param_groups" in ref_osd
+        if "param_groups" not in ref_osd:
+            raise AssertionError("Expected 'param_groups' in ref_osd")
         self.assertTrue("param_groups" in full_osd)
         ref_osd_param_groups = ref_osd["param_groups"]
         full_osd_param_groups = full_osd["param_groups"]
@@ -592,7 +606,8 @@ def _test_optim_state_dict_nested(
         ref_osd = optim2.state_dict()
         # Check the losses to eliminate model drift as a source of error
         for i, (l1, l2) in enumerate(zip(losses1, losses2)):
-            assert l1 == l2, f"Losses differ on iter {i}: {l1:.5f} {l2:.5f}"
+            if l1 != l2:
+                raise AssertionError(f"Losses differ on iter {i}: {l1:.5f} {l2:.5f}")
         # Do not check the parameter keys since the full/sharded optimizer state
         # dict uses parameter names, while the non-wrapped equivalent uses
         # parameter IDs
@@ -1449,7 +1464,7 @@ def _run_on_all_optim_state_apis(
         self,
         should_check_method_fn: Callable[[str], bool],
         context_fn: Callable,
-        fsdp_kwargs: Optional[dict[str, Any]],
+        fsdp_kwargs: dict[str, Any] | None,
     ):
         """
         Runs through all optimizer state checkpointing APIs with a context
diff --git a/test/distributed/fsdp/test_fsdp_overlap.py b/test/distributed/fsdp/test_fsdp_overlap.py
index 6f2c73f51a20b..01749dd2230a9 100644
--- a/test/distributed/fsdp/test_fsdp_overlap.py
+++ b/test/distributed/fsdp/test_fsdp_overlap.py
@@ -143,7 +143,8 @@ def _delayed_all_gather(*args, **kwargs):
                     all_gather_called = True
                     if torch.cuda.is_available():
                         torch.cuda._sleep(all_gather_cycles)
-                    assert orig_all_gather
+                    if not orig_all_gather:
+                        raise AssertionError("Expected orig_all_gather to be truthy")
                     return orig_all_gather(*args, **kwargs)
 
                 # forward pass
diff --git a/test/distributed/fsdp/test_fsdp_sharded_grad_scaler.py b/test/distributed/fsdp/test_fsdp_sharded_grad_scaler.py
index b97e88f5535b5..4c3425f22a969 100644
--- a/test/distributed/fsdp/test_fsdp_sharded_grad_scaler.py
+++ b/test/distributed/fsdp/test_fsdp_sharded_grad_scaler.py
@@ -5,7 +5,6 @@
 import itertools
 import sys
 import unittest
-from typing import Optional
 
 import torch
 from torch import distributed as dist
@@ -180,9 +179,9 @@ def _get_init_modes_for_test(self, cpu_offload):
     def test_fsdp_ddp_parity_with_grad_scaler(
         self,
         cpu_offload: CPUOffload,
-        sharding_strategy: Optional[ShardingStrategy],
-        mixed_precision: Optional[str],
-        use_orig_params: Optional[str],
+        sharding_strategy: ShardingStrategy | None,
+        mixed_precision: str | None,
+        use_orig_params: str | None,
     ):
         init_modes = self._get_init_modes_for_test(cpu_offload)
         mp = (
diff --git a/test/distributed/fsdp/test_fsdp_state_dict.py b/test/distributed/fsdp/test_fsdp_state_dict.py
index 50e9e6a798681..1933a68a9d21d 100644
--- a/test/distributed/fsdp/test_fsdp_state_dict.py
+++ b/test/distributed/fsdp/test_fsdp_state_dict.py
@@ -178,7 +178,10 @@ def _state_compare(self, model, model_new, assert_fn, state_generator="parameter
     def _compare_models(
         self, model, model_new, assert_fn, check_fp16=False, check_buffers=True
     ):
-        assert assert_fn in (self.assertEqual, self.assertNotEqual)
+        if assert_fn not in (self.assertEqual, self.assertNotEqual):
+            raise AssertionError(
+                f"Expected assert_fn in (self.assertEqual, self.assertNotEqual), got {assert_fn}"
+            )
         with FSDP.summon_full_params(model):
             with FSDP.summon_full_params(model_new):
                 self._state_compare(model, model_new, assert_fn)
@@ -1212,8 +1215,12 @@ def test_full_state_dict_missing_unexpected_keys_cleaned(self):
         # Create an unexpected key
         sd["unexpected"] = torch.ones(1)
         missing, unexpected = model.load_state_dict(sd, strict=False)
-        assert len(missing) == 1
-        assert len(unexpected) == 1
+        if len(missing) != 1:
+            raise AssertionError(f"Expected len(missing) == 1, got {len(missing)}")
+        if len(unexpected) != 1:
+            raise AssertionError(
+                f"Expected len(unexpected) == 1, got {len(unexpected)}"
+            )
         self.assertTrue(FSDP_PREFIX not in missing[0])
         self.assertTrue(FSDP_PREFIX not in unexpected[0])
 
diff --git a/test/distributed/fsdp/test_fsdp_tp_integration.py b/test/distributed/fsdp/test_fsdp_tp_integration.py
index 70db377c2d1b1..74395e6a3e46a 100644
--- a/test/distributed/fsdp/test_fsdp_tp_integration.py
+++ b/test/distributed/fsdp/test_fsdp_tp_integration.py
@@ -2,7 +2,6 @@
 import copy
 import sys
 from collections import OrderedDict
-from typing import Optional
 
 import torch
 from torch import distributed as dist
@@ -93,9 +92,10 @@ def _get_params_and_sharding_info(
         tensor_parallel_size: int,
     ) -> tuple[dict[str, int], dict[str, tuple[torch.Size, int]]]:
         """ """
-        assert type(model) is SimpleModel, (
-            "Expects a `SimpleModel` since the sharding cases on the model definition"
-        )
+        if type(model) is not SimpleModel:
+            raise AssertionError(
+                "Expects a `SimpleModel` since the sharding cases on the model definition"
+            )
         param_name_to_numel = OrderedDict()
         param_name_to_sharding_info = OrderedDict()
         for param_name, param in model.named_parameters():
@@ -143,13 +143,14 @@ def _sync_tp_grads(
         """
         tp_world_size = tp_pg.size()
         fsdp_world_size = self.world_size // tp_world_size
-        assert (
+        if not (
             type(tp_fsdp_model) is FSDP
             and len([m for m in tp_fsdp_model.modules() if type(m) is FSDP]) == 1
-        ), (
-            "The following logic assumes a single top-level-only FSDP wrapping "
-            "the model with TP already applied"
-        )
+        ):
+            raise AssertionError(
+                "The following logic assumes a single top-level-only FSDP wrapping "
+                "the model with TP already applied"
+            )
         for flat_param in tp_fsdp_model.params:
             splits = tuple(param_name_to_numel.values())
             # Create a mask over the gradient elements to manually reduce
@@ -181,9 +182,9 @@ def _get_grads_as_flattened(
         uses_tp: bool,
         param_name_to_numel: dict[str, int],
         param_name_to_sharding_info: dict[str, tuple[torch.Size, int]],
-        tp_pg: Optional[dist.ProcessGroup],
-        fsdp_pg: Optional[dist.ProcessGroup],
-        sharded_param_names: Optional[list[str]],
+        tp_pg: dist.ProcessGroup | None,
+        fsdp_pg: dist.ProcessGroup | None,
+        sharded_param_names: list[str] | None,
     ) -> torch.Tensor:
         """
         Returns all unsharded gradients as a single flattened tensor. This
@@ -297,8 +298,14 @@ def _test_fsdp_tp_integration(
             sequence_parallelize_plan,
         )
         tp_pg = mesh_2d["tp"].get_group(mesh_dim=0)
-        assert isinstance(tp_fsdp_model.net1.weight, DTensor)
-        assert isinstance(tp_fsdp_model.net2.weight, DTensor)
+        if not isinstance(tp_fsdp_model.net1.weight, DTensor):
+            raise AssertionError(
+                f"Expected DTensor, got {type(tp_fsdp_model.net1.weight)}"
+            )
+        if not isinstance(tp_fsdp_model.net2.weight, DTensor):
+            raise AssertionError(
+                f"Expected DTensor, got {type(tp_fsdp_model.net2.weight)}"
+            )
         tp_fsdp_model = FSDP(
             tp_fsdp_model,
             cpu_offload=cpu_offload,
diff --git a/test/distributed/fsdp/test_fsdp_unshard_params.py b/test/distributed/fsdp/test_fsdp_unshard_params.py
index b6010851ae2ef..33d37ead769d6 100644
--- a/test/distributed/fsdp/test_fsdp_unshard_params.py
+++ b/test/distributed/fsdp/test_fsdp_unshard_params.py
@@ -3,7 +3,7 @@
 import itertools
 import math
 import sys
-from typing import Any, Optional, Union
+from typing import Any
 
 import torch
 import torch.distributed.fsdp._traversal_utils as traversal_utils
@@ -67,8 +67,8 @@ def _test_unshard_params_writeback(
         offloading_params = model.cpu_offload.offload_params
 
         # Assumes depth-first `.parameters()`
-        outer_param: Union[FlatParameter, nn.Parameter] = next(model.parameters())
-        inner_param: Union[FlatParameter, nn.Parameter] = next(model[0].parameters())
+        outer_param: FlatParameter | nn.Parameter = next(model.parameters())
+        inner_param: FlatParameter | nn.Parameter = next(model[0].parameters())
         param_to_check = outer_param if check_outer else inner_param
 
         # Write a known value to all elements of the *sharded* parameter or
@@ -118,7 +118,7 @@ def _test_unshard_params_param_data(
         rank0_only: bool,
         offload_to_cpu: bool,
         cpu_offload: CPUOffload,
-        mixed_precision: Optional[MixedPrecision],
+        mixed_precision: MixedPrecision | None,
         use_orig_params: bool,
     ):
         local_model = NestedWrappedModule.init(
@@ -280,7 +280,7 @@ def _test_unshard_params_respects_reshard(
         self,
         rank0_only: bool,
         offload_to_cpu: bool,
-        mixed_precision: Optional[MixedPrecision],
+        mixed_precision: MixedPrecision | None,
         use_orig_params: bool,
     ):
         """NOTE: This method depends on FSDP internals."""
@@ -369,7 +369,7 @@ def _test_unshard_params_recurse(
         self,
         recurse: bool,
         unshard_outer: bool,
-        mixed_precision: Optional[MixedPrecision],
+        mixed_precision: MixedPrecision | None,
         use_orig_params: bool,
     ):
         """NOTE: This method depends on FSDP internals."""
@@ -493,7 +493,7 @@ def _test_with_grads_core(
         def _check_grads(
             ddp_model: DDP,
             fsdp_model: FSDP,
-            old_fsdp_grads: Optional[list[torch.Tensor]],
+            old_fsdp_grads: list[torch.Tensor] | None,
         ):
             """
             Checks that writes to the FSDP parameters' gradients persist or do
@@ -513,11 +513,13 @@ def _check_grads(
                     fsdp_model.named_parameters(),
                 ):
                     self.assertEqual(n1, clean_tensor_name(n2))
-                    assert p1.grad is not None
+                    if p1.grad is None:
+                        raise AssertionError("Expected p1.grad to not be None")
                     torch.testing.assert_close(p1.grad, p2.grad)
                     # Ensure that the tensor is not all zeros, which would
                     # mean that the multiplication is vacuous
-                    assert torch.count_nonzero(p2.grad) > 0
+                    if not (torch.count_nonzero(p2.grad) > 0):
+                        raise AssertionError("Expected nonzero gradient")
                     p2.grad *= WRITEBACK_FACTOR
             new_fsdp_grads = [
                 param.grad
@@ -574,7 +576,8 @@ def _get_fsdp_grads(fsdp_model: FSDP, is_supported: bool):
         )
         with FSDP.summon_full_params(fsdp_model):
             for p1, p2 in zip(ddp_model.module.parameters(), fsdp_model.parameters()):
-                assert torch.all(torch.isclose(p1, p2))
+                if not torch.all(torch.isclose(p1, p2)):
+                    raise AssertionError("Expected all parameters to be close")
 
         # Check calling after backward
         inp = fsdp_model.get_input(torch.device(device_type))
@@ -625,7 +628,8 @@ def _test_with_grads_none_grads(self, sharding_strategy: ShardingStrategy):
         )
         for fsdp_module in FSDP.fsdp_modules(fsdp_model):
             if fsdp_module._handle:
-                assert fsdp_module._handle.flat_param.grad is None
+                if fsdp_module._handle.flat_param.grad is not None:
+                    raise AssertionError("Expected flat_param.grad to be None")
         with FSDP.summon_full_params(fsdp_model, with_grads=True):
             for param in fsdp_model.parameters():
                 self.assertTrue(param.grad is None)
diff --git a/test/distributed/fsdp/test_fsdp_use_orig_params.py b/test/distributed/fsdp/test_fsdp_use_orig_params.py
index d0edb8045f470..34ba6329a7a25 100644
--- a/test/distributed/fsdp/test_fsdp_use_orig_params.py
+++ b/test/distributed/fsdp/test_fsdp_use_orig_params.py
@@ -6,7 +6,7 @@
 import os
 import sys
 import unittest
-from typing import Any, Optional
+from typing import Any
 
 import torch
 import torch.nn as nn
@@ -123,7 +123,7 @@ def _get_fsdp_transformer_and_optim(
         optim_class: type[torch.optim.Optimizer],
         multi_tensor: bool,
         sharding_strategy: ShardingStrategy,
-        backward_prefetch: Optional[BackwardPrefetch],
+        backward_prefetch: BackwardPrefetch | None,
         cpu_offload: CPUOffload,
     ) -> tuple[FSDP, torch.optim.Optimizer]:
         """
@@ -341,7 +341,7 @@ def _test_diff_hyperparams(
         optim_class: type[torch.optim.Optimizer],
         multi_tensor: bool,
         set_to_none: bool,
-        backward_prefetch: Optional[BackwardPrefetch],
+        backward_prefetch: BackwardPrefetch | None,
         cpu_offload: CPUOffload,
         sharding_strategy: ShardingStrategy,
         skip_writeback_check: bool,
@@ -438,7 +438,10 @@ def test_multiple_optimizers(self):
     def _test_multiple_optimizers(self, sharding_strategy: ShardingStrategy):
         ddp_model = self._get_ddp_transformer(find_unused_params=True)
         ddp_param_groups = self._get_param_groups(ddp_model)
-        assert len(ddp_param_groups) == 3, f"{len(ddp_param_groups)}"
+        if not (len(ddp_param_groups) == 3):
+            raise AssertionError(
+                f"Expected 3 param groups, got {len(ddp_param_groups)}"
+            )
         (
             fsdp_model,
             _,
@@ -452,7 +455,10 @@ def _test_multiple_optimizers(self, sharding_strategy: ShardingStrategy):
             cpu_offload=None,
         )
         fsdp_param_groups = self._get_param_groups(fsdp_model)
-        assert len(fsdp_param_groups) == 3, f"{len(fsdp_param_groups)}"
+        if not (len(fsdp_param_groups) == 3):
+            raise AssertionError(
+                f"Expected 3 param groups, got {len(fsdp_param_groups)}"
+            )
         ddp_optims = []
         fsdp_optims = []
         # For the transformer model, every parameter is either a weight or a
@@ -483,7 +489,8 @@ def _test_multiple_optimizers(self, sharding_strategy: ShardingStrategy):
             if not handle:
                 continue
             flat_param = handle.flat_param
-            assert flat_param._params is not None
+            if flat_param._params is None:
+                raise AssertionError("Expected flat_param._params to not be None")
             has_weight = False
             has_bias = False
             for param, fqn in zip(flat_param._params, flat_param._fqns):
@@ -492,10 +499,11 @@ def _test_multiple_optimizers(self, sharding_strategy: ShardingStrategy):
                 elif "bias" in fqn and param.numel() > 0:
                     has_bias = True
             has_both |= has_weight and has_bias
-        assert has_both, (
-            f"Rank {self.rank} does not have a `FlatParameter` with both a "
-            "weight and a bias in its shard, meaning that this test is vacuous"
-        )
+        if not has_both:
+            raise AssertionError(
+                f"Rank {self.rank} does not have a `FlatParameter` with both a "
+                "weight and a bias in its shard, meaning that this test is vacuous"
+            )
 
         # Run one iteration to generate gradients
         def run_iter():
@@ -793,10 +801,10 @@ def get_loss(self, inp, out):
         def check_parameter_parity(
             ddp_model: DDP, fsdp_model: FSDP, between_fwd_and_bwd: bool
         ):
-            assert self.rank in (
-                0,
-                1,
-            ), f"Expects world size of 2 but got {self.world_size}"
+            if self.rank not in (0, 1):
+                raise AssertionError(
+                    f"Expects world size of 2 but got {self.world_size}"
+                )
             for (n1, p1), (n2, p2) in zip(
                 ddp_model.module.named_parameters(),
                 fsdp_model.named_parameters(),
@@ -1031,7 +1039,8 @@ def test_writeback_shape_mismatch(self):
         )
         # Check that writing back with mismatched shape errors
         fsdp = fsdp_model.module  # for brevity
-        assert self.rank in (0, 1), f"Expects world size of 2 but got {self.world_size}"
+        if self.rank not in (0, 1):
+            raise AssertionError(f"Expects world size of 2 but got {self.world_size}")
         with self.assertRaisesRegex(RuntimeError, "Cannot writeback"):
             # Change the gradient to a new one with 1 added to each dimension
             # to force a shape mismatch when writing back
@@ -1173,9 +1182,8 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
                     clean_tensor_name(tup[0]) for tup in self.named_parameters()
                 ]
                 params = [tup[1] for tup in self.named_parameters()]
-                assert param_shapes[0] is not None and param_shapes[1] is not None, (
-                    "`param_sizes` should be set"
-                )
+                if not (param_shapes[0] is not None and param_shapes[1] is not None):
+                    raise AssertionError("`param_sizes` should be set")
                 assert_equal_fn(
                     param_names,
                     [
diff --git a/test/distributed/fsdp/test_utils.py b/test/distributed/fsdp/test_utils.py
index 07ed117926df6..0a9343ce41a96 100644
--- a/test/distributed/fsdp/test_utils.py
+++ b/test/distributed/fsdp/test_utils.py
@@ -112,13 +112,21 @@ def test_replace_by_prefix(self):
         }
         original_state_dict = state_dict.copy()
         _replace_by_prefix(state_dict, "layer.", "module.layer.")
-        assert state_dict == {
+        expected = {
             "module.layer.a": torch.tensor(1),
             "abc.layer.def": torch.tensor(2),
             "module.layer.b": torch.tensor(3),
         }
+        # Note: comparing dicts with tensors directly doesn't work well, so we compare keys
+        if set(state_dict.keys()) != set(expected.keys()):
+            raise AssertionError(
+                f"Expected keys {set(expected.keys())}, got {set(state_dict.keys())}"
+            )
         _replace_by_prefix(state_dict, "module.layer.", "layer.")
-        assert state_dict == original_state_dict
+        if set(state_dict.keys()) != set(original_state_dict.keys()):
+            raise AssertionError(
+                f"Expected keys {set(original_state_dict.keys())}, got {set(state_dict.keys())}"
+            )
 
     @skip_if_lt_x_gpu(1)
     def test_packed_sequence(self):
diff --git a/test/distributed/fsdp/test_wrap.py b/test/distributed/fsdp/test_wrap.py
index 5766a0d65c09f..9c341e0d4f06a 100644
--- a/test/distributed/fsdp/test_wrap.py
+++ b/test/distributed/fsdp/test_wrap.py
@@ -7,7 +7,6 @@
 import unittest
 from collections.abc import Callable
 from enum import auto, Enum
-from typing import Union
 
 import torch
 import torch.nn as nn
@@ -428,7 +427,8 @@ def test_wrap(self, wrap_method):
             with enable_wrap(wrapper_cls=FSDP, process_group=self.process_group):
                 layer = wrap(nn.Linear(5, 5))
         else:
-            assert wrap_method == WrapMethod.FSDP_CTOR
+            if wrap_method != WrapMethod.FSDP_CTOR:
+                raise AssertionError(f"Unexpected wrap_method: {wrap_method}")
             layer = FSDP(
                 nn.Linear(5, 5),
                 process_group=self.process_group,
@@ -505,7 +505,7 @@ def test_module_wrap_policy_callable(self):
         callable_policy = functools.partial(_or_policy, policies=[auto_wrap_policy])
         self._test_transformer_wrapping(callable_policy)
 
-    def _test_transformer_wrapping(self, auto_wrap_policy: Union[Callable, _Policy]):
+    def _test_transformer_wrapping(self, auto_wrap_policy: Callable | _Policy):
         fsdp_kwargs = {"auto_wrap_policy": auto_wrap_policy}
         fsdp_model = TransformerWithSharedParams.init(
             self.process_group,
@@ -817,7 +817,7 @@ def test_always_wrap_with_ignored_modules(self, wrap_method: WrapMethod):
             with enable_wrap(wrapper_cls=FSDP, **fsdp_kwargs):
                 model = wrap(sequential)
         else:
-            assert 0, f"Unsupported wrap method: {wrap_method}"
+            raise AssertionError(f"Unsupported wrap method: {wrap_method}")
         # All non-ignored modules should be wrapped with FSDP
         self.assertTrue(isinstance(model, FSDP))
         self.assertTrue(isinstance(model.module[0], FSDP))
@@ -846,7 +846,7 @@ def test_auto_wrap_with_ignored_modules(self, wrap_method: WrapMethod):
             with enable_wrap(wrapper_cls=FSDP, **fsdp_kwargs):
                 model = wrap(sequential)
         else:
-            assert 0, f"Unsupported wrap method: {wrap_method}"
+            raise AssertionError(f"Unsupported wrap method: {wrap_method}")
         # Since the 2nd linear (`sequential[1]`) is ignored, the wrapping
         # policy does not exceed the parameter threshold before the inner
         # sequential (`sequential[2]`) anymore; hence, it flattens
diff --git a/test/distributed/launcher/api_test.py b/test/distributed/launcher/api_test.py
index 32e5f74cd6770..dab054f57b4f1 100644
--- a/test/distributed/launcher/api_test.py
+++ b/test/distributed/launcher/api_test.py
@@ -16,7 +16,7 @@
 import unittest
 import uuid
 from contextlib import closing
-from typing import Any, Optional
+from typing import Any
 from unittest import mock
 from unittest.mock import MagicMock, Mock, patch
 
@@ -59,7 +59,7 @@ def get_test_launch_config(
     nproc_per_node: int,
     run_id: str = "",
     rdzv_backend: str = "etcd",
-    config: Optional[dict[str, Any]] = None,
+    config: dict[str, Any] | None = None,
 ) -> LaunchConfig:
     rdzv_configs = {}
     if config:
diff --git a/test/distributed/launcher/test_run.py b/test/distributed/launcher/test_run.py
index 484a975051d4f..47c4508a7550f 100644
--- a/test/distributed/launcher/test_run.py
+++ b/test/distributed/launcher/test_run.py
@@ -19,6 +19,7 @@
 from unittest import mock, skipIf
 from unittest.mock import MagicMock, Mock, patch
 
+import torch
 import torch.distributed.run as launch
 from torch.distributed.elastic.agent.server.api import RunResult, WorkerState
 from torch.distributed.elastic.multiprocessing import DefaultLogsSpecs
@@ -254,7 +255,8 @@ def _test_nproc_launch_configuration(self, nproc_type, expected_number):
     )
     @patch("torch.cuda.is_available", return_value=False)
     def test_nproc_launch_auto_configurations(self, _mock1):
-        self._test_nproc_launch_configuration("auto", os.cpu_count())
+        expected = torch._utils.cpu_count()
+        self._test_nproc_launch_configuration("auto", expected)
 
     @skip_but_pass_in_sandcastle_if(
         TEST_WITH_DEV_DBG_ASAN, "test incompatible with dev/dbg asan"
diff --git a/test/distributed/local_tensor_tutorial_examples/example_05_rank_specific.py b/test/distributed/local_tensor_tutorial_examples/example_05_rank_specific.py
index 2be814d7917bf..2e76d93cf9941 100644
--- a/test/distributed/local_tensor_tutorial_examples/example_05_rank_specific.py
+++ b/test/distributed/local_tensor_tutorial_examples/example_05_rank_specific.py
@@ -18,6 +18,8 @@
     LocalTensorMode,
     maybe_disable_local_tensor_mode,
     maybe_run_for_local_tensor,
+    rank_map,
+    tensor_map,
 )
 
 
@@ -27,8 +29,8 @@ def use_rank_map(world_size: int = 4):
 
     Returns: (values_dict, expected_dict)
     """
-    with LocalTensorMode(world_size) as mode:
-        lt = mode.rank_map(lambda rank: torch.full((2, 3), float(rank)))
+    with LocalTensorMode(world_size):
+        lt = rank_map(lambda rank: torch.full((2, 3), float(rank)))
         values = {
             rank: lt._local_tensors[rank][0, 0].item() for rank in range(world_size)
         }
@@ -46,13 +48,13 @@ def use_tensor_map(world_size: int = 4):
 
     Returns: (values_dict, expected_dict)
     """
-    with LocalTensorMode(world_size) as mode:
-        lt = mode.rank_map(lambda rank: torch.ones(2, 2) * (rank + 1))
+    with LocalTensorMode(world_size):
+        lt = rank_map(lambda rank: torch.ones(2, 2) * (rank + 1))
 
         def scale_by_rank(rank: int, tensor: torch.Tensor) -> torch.Tensor:
             return tensor * (rank + 1)
 
-        scaled = mode.tensor_map(lt, scale_by_rank)
+        scaled = tensor_map(lt, scale_by_rank)
         values = {
             rank: scaled._local_tensors[rank][0, 0].item() for rank in range(world_size)
         }
@@ -128,8 +130,8 @@ def compute_rank_offset(data: torch.Tensor, rank: int) -> torch.Tensor:
 
     full_data = torch.arange(16).float()
 
-    with LocalTensorMode(world_size) as mode:
-        ranks = mode.rank_map(lambda r: torch.tensor(r))
+    with LocalTensorMode(world_size):
+        ranks = rank_map(lambda r: torch.tensor(r))
         result = compute_rank_offset(full_data, ranks)
 
         values = {
diff --git a/test/distributed/optim/test_zero_redundancy_optimizer.py b/test/distributed/optim/test_zero_redundancy_optimizer.py
index 96d8d1d3225d1..283ac98bf5c29 100644
--- a/test/distributed/optim/test_zero_redundancy_optimizer.py
+++ b/test/distributed/optim/test_zero_redundancy_optimizer.py
@@ -304,13 +304,15 @@ def test_constructor(self):
             betas=BETAS,
             eps=EPS,
         )
-        assert len(o.param_groups) == 2, (
-            f"Expected 2 ZeRO param groups, but got {len(o.param_groups)}"
-        )
-        assert len(o.optim.param_groups) == 2, (
-            "Expected 2 local optimizer param groups, but got "
-            f"{len(o.optim.param_groups)}"
-        )
+        if not (len(o.param_groups) == 2):
+            raise AssertionError(
+                f"Expected 2 ZeRO param groups, but got {len(o.param_groups)}"
+            )
+        if not (len(o.optim.param_groups) == 2):
+            raise AssertionError(
+                "Expected 2 local optimizer param groups, but got "
+                f"{len(o.optim.param_groups)}"
+            )
 
     def test_same_dense_param_type(self):
         """Check that ZeroRedundancyOptimizer raises an exception if the input
@@ -722,9 +724,8 @@ def test_nondefault_process_group(self):
         LR = 1e-3
         MOMENTUM = 0.99
         REFERENCE_RANK = 0
-        assert REFERENCE_RANK in subgroup_ranks, (
-            "Reference rank must be in the new process group"
-        )
+        if REFERENCE_RANK not in subgroup_ranks:
+            raise AssertionError("Reference rank must be in the new process group")
         loss_fn = torch.nn.L1Loss().to(device)
 
         def check(optimizer):
@@ -812,7 +813,7 @@ def test_local_optimizer_parity(
         elif optimizer_class_str == "SGD":
             optimizer_class = torch.optim.SGD
         else:
-            assert 0, f"Unsupported optimizer class: {optimizer_class_str}"
+            raise AssertionError(f"Unsupported optimizer class: {optimizer_class_str}")
 
         with self.context:
             # Define a base model with a different buffer for each rank
@@ -1033,8 +1034,10 @@ def __init__(self) -> None:
                 super().__init__()
 
             def join_hook(self, **kwargs):
-                assert "zero_optim" in kwargs
-                assert "grads" in kwargs
+                if "zero_optim" not in kwargs:
+                    raise AssertionError("Expected 'zero_optim' in kwargs")
+                if "grads" not in kwargs:
+                    raise AssertionError("Expected 'grads' in kwargs")
                 zero_optim = kwargs["zero_optim"]
                 grads = kwargs["grads"]
                 return _SetGradsJoinHook(zero_optim, grads)
@@ -1095,7 +1098,8 @@ def test_zero_join_cpu(self):
 
     def _test_zero_model_parallel(self, parameters_as_bucket_view: bool, device: str):
         # Use two processes each with two GPUs
-        assert self.rank < 2
+        if not (self.rank < 2):
+            raise AssertionError(f"Expected rank < 2, got {self.rank}")
         NUM_EPOCHS = 2
         NUM_INPUTS = 4
         LR = 0.01
diff --git a/test/distributed/pipelining/model_registry.py b/test/distributed/pipelining/model_registry.py
index 5213cf4ac647a..308c5d258de3d 100644
--- a/test/distributed/pipelining/model_registry.py
+++ b/test/distributed/pipelining/model_registry.py
@@ -8,7 +8,8 @@
 
 class ExampleCode(torch.nn.Module):
     def __init__(self, d_hid, splits=2):
-        assert splits <= 8
+        if not (splits <= 8):
+            raise AssertionError(f"Expected splits <= 8, got {splits}")
         super().__init__()
         self.splits = splits
         self.mm_param0 = torch.nn.Parameter(torch.randn(d_hid, d_hid))
@@ -63,7 +64,8 @@ class ModelWithKwargs(torch.nn.Module):
     DEFAULT_BATCH_SIZE = 256
 
     def __init__(self, d_hid: int = DEFAULT_DHID, splits=2):
-        assert splits <= 8
+        if not (splits <= 8):
+            raise AssertionError(f"Expected splits <= 8, got {splits}")
         super().__init__()
         self.splits = splits
         self.mm_param0 = torch.nn.Parameter(torch.randn(d_hid, d_hid))
@@ -186,6 +188,30 @@ def forward(self, x):
         return x
 
 
+class TwoInputOutputOp(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, input: torch.Tensor, weight: torch.Tensor):
+        return input, weight
+
+    @staticmethod
+    def backward(ctx, grad_input, grad_weight):
+        return grad_input, grad_weight
+
+
+# Model with multi-output intermediates
+class MultiInterMediateModel(torch.nn.Module):
+    def __init__(self, weight_shape: list[int]):
+        super().__init__()
+        self.shape = weight_shape
+        self.w = torch.nn.Parameter(torch.randn(*weight_shape))
+
+    def forward(self, x):
+        a, b = torch.split(x, self.shape, dim=1)
+        a, w = TwoInputOutputOp.apply(a, self.w)
+        a = torch.matmul(a, w)
+        return a * b
+
+
 # Multi-MLP with kwargs model
 class MultiMLPKwargs(torch.nn.Module):
     def __init__(self, d_hid: int, n_layers: int = 2):
diff --git a/test/distributed/pipelining/schedule_registry.py b/test/distributed/pipelining/schedule_registry.py
index 0b037bd8b99b4..387752c2ebaed 100644
--- a/test/distributed/pipelining/schedule_registry.py
+++ b/test/distributed/pipelining/schedule_registry.py
@@ -3,7 +3,6 @@
 # This file is a Schedule zoo for testing torch.distributed.pipelining.
 # It includes schedules designed purely for testing purposes
 from collections.abc import Callable
-from typing import Optional
 
 from torch.distributed.pipelining.schedules import (
     _Action,
@@ -35,7 +34,7 @@ def __init__(
         self,
         stages: list[_PipelineStageBase],
         n_microbatches: int,
-        loss_fn: Optional[Callable] = None,
+        loss_fn: Callable | None = None,
         scale_grads: bool = True,
     ):
         super().__init__(
@@ -84,7 +83,7 @@ def __init__(
         self,
         stages: list[_PipelineStageBase],
         n_microbatches: int,
-        loss_fn: Optional[Callable] = None,
+        loss_fn: Callable | None = None,
         scale_grads: bool = True,
     ):
         super().__init__(
@@ -135,7 +134,7 @@ def __init__(
         self,
         stages: list[_PipelineStageBase],
         n_microbatches: int,
-        loss_fn: Optional[Callable] = None,
+        loss_fn: Callable | None = None,
         enable_zero_bubble: bool = True,
         scale_grads: bool = True,
     ):
@@ -197,7 +196,7 @@ def __init__(
         self,
         stages: list[_PipelineStageBase],
         n_microbatches: int,
-        loss_fn: Optional[Callable] = None,
+        loss_fn: Callable | None = None,
         scale_grads: bool = True,
     ):
         super().__init__(
diff --git a/test/distributed/pipelining/test_backward.py b/test/distributed/pipelining/test_backward.py
index b46a97d02c29e..cd712b97f7fcd 100644
--- a/test/distributed/pipelining/test_backward.py
+++ b/test/distributed/pipelining/test_backward.py
@@ -2,7 +2,7 @@
 # Owner(s): ["oncall: distributed"]
 import copy
 
-from model_registry import MLPModule
+from model_registry import MLPModule, MultiInterMediateModel
 
 import torch
 from torch.distributed.pipelining._backward import (
@@ -192,24 +192,35 @@ def test_stage_backward_weight_multiple_iters(self, device):
     def test_stage_backward_weight_grad_validation(self, device):
         test_cases = [
             (
-                "size >= 2",
+                "size == 2",
+                lambda: MultiInterMediateModel([d_hid // 2, d_hid // 2]).to(device),
                 lambda: [
                     (
-                        torch.randn(batch_size, d_hid, device=device),
-                        torch.randn(batch_size, d_hid, device=device),
+                        torch.randn(batch_size, d_hid // 2, device=device),
+                        torch.randn(d_hid // 2, d_hid // 2, device=device),
                     )
                 ],
             ),
-            ("size = 1", lambda: [(torch.randn(batch_size, d_hid, device=device),)]),
+            (
+                "size = 1",
+                lambda: MLPModule(d_hid).to(device),
+                lambda: [(torch.randn(batch_size, d_hid, device=device),)],
+            ),
             (
                 "1 grad, 1 None",
-                lambda: [(torch.randn(batch_size, d_hid, device=device), None)],
+                lambda: MultiInterMediateModel([d_hid // 2, d_hid // 2]).to(device),
+                lambda: [(torch.randn(batch_size, d_hid // 2, device=device), None)],
+            ),
+            (
+                "1 None, 1 grad",
+                lambda: MultiInterMediateModel([d_hid // 2, d_hid // 2]).to(device),
+                lambda: [(None, torch.randn(d_hid // 2, d_hid // 2, device=device))],
             ),
         ]
 
-        for description, mock_grads_factory in test_cases:
+        for description, module_factory, mock_grads_factory in test_cases:
             with self.subTest(description=description):
-                mod = MLPModule(d_hid).to(device)
+                mod = module_factory()
                 x = torch.randn(batch_size, d_hid, device=device)
                 x.requires_grad_(True)
                 out = mod(x)
@@ -227,6 +238,33 @@ def test_stage_backward_weight_grad_validation(self, device):
 
                 stage_backward_weight(mod.parameters(), param_groups)
 
+    def test_stage_backward_multi_output_intermediate(self, device):
+        mod = MultiInterMediateModel([d_hid // 2, d_hid // 2]).to(device)
+        x = torch.randn(batch_size, d_hid, device=device, requires_grad=True)
+
+        out = mod(x)
+        loss = out.sum()
+
+        dinputs, param_groups = stage_backward_input(
+            stage_outputs_or_loss=[loss],
+            output_grads=None,
+            input_values=[x],
+            weights=mod.parameters(),
+        )
+
+        stage_backward_weight(mod.parameters(), param_groups)
+
+        ref_mod = copy.deepcopy(mod)
+        ref_x = x.detach().clone().requires_grad_(True)
+        ref_out = ref_mod(ref_x)
+        ref_loss = ref_out.sum()
+        ref_loss.backward()
+
+        torch.testing.assert_close(dinputs[0], ref_x.grad)
+        for name, p in mod.named_parameters():
+            ref_p = ref_mod.get_parameter(name)
+            torch.testing.assert_close(p.grad, ref_p.grad)
+
 
 devices = ["cpu", "cuda", "hpu", "xpu"]
 instantiate_device_type_tests(
diff --git a/test/distributed/pipelining/test_dtensor_pp_integration.py b/test/distributed/pipelining/test_dtensor_pp_integration.py
new file mode 100644
index 0000000000000..186a6ea14e8fc
--- /dev/null
+++ b/test/distributed/pipelining/test_dtensor_pp_integration.py
@@ -0,0 +1,804 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates
+# Owner(s): ["oncall: distributed"]
+
+from __future__ import annotations
+
+import copy
+import functools
+from typing import cast, TYPE_CHECKING
+
+import torch
+import torch.distributed as dist
+import torch.nn as nn
+from torch.distributed.device_mesh import DeviceMesh, init_device_mesh
+from torch.distributed.pipelining._utils import _DTensorMeta
+from torch.distributed.pipelining.schedules import (
+    Schedule1F1B,
+    ScheduleDualPipeV,
+    ScheduleInterleaved1F1B,
+    ScheduleZBVZeroBubble,
+)
+from torch.distributed.pipelining.stage import PipelineStage
+from torch.distributed.tensor import distribute_tensor, DTensor, Replicate, Shard
+from torch.distributed.tensor.parallel import (
+    ColwiseParallel,
+    parallelize_module,
+    RowwiseParallel,
+    SequenceParallel,
+)
+from torch.testing._internal.common_distributed import (
+    MultiProcContinuousTest,
+    requires_accelerator_dist_backend,
+)
+from torch.testing._internal.common_utils import (
+    run_tests,
+    skip_but_pass_in_sandcastle_if,
+    TEST_MULTIACCELERATOR,
+)
+
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+
+d_hid = 256
+batch_size = 64
+n_microbatches = 4
+
+StageStaticMeta = tuple[DTensor, DTensor, DTensor | None, DTensor | None]
+StageStaticMetaMap = dict[int, StageStaticMeta]
+
+MODEL_SEED = 0
+INPUT_SEED = 42
+TARGET_SEED = 123
+
+
+device_type = acc.type if (acc := torch.accelerator.current_accelerator()) else "cpu"
+backend = dist.get_default_backend_for_device(device_type)
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, dim: int, eps: float = 1e-6):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(dim))
+        self.eps = eps
+
+    def reset_parameters(self) -> None:
+        torch.nn.init.ones_(self.weight)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return (
+            x
+            * torch.rsqrt(x.float().pow(2).mean(-1, keepdim=True) + self.eps)
+            * self.weight
+        )
+
+
+class NormMLP(nn.Module):
+    def __init__(self, dim: int):
+        super().__init__()
+        self.norm = RMSNorm(dim)
+        self.fc1 = nn.Linear(dim, 4 * dim, bias=False)
+        self.act = nn.GELU()
+        self.fc2 = nn.Linear(4 * dim, dim, bias=False)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.fc2(self.act(self.fc1(self.norm(x))))
+
+
+class NormMLPStack(nn.Module):
+    def __init__(self, dim: int, n_layers: int):
+        super().__init__()
+        self.layers = nn.ModuleList([NormMLP(dim) for _ in range(n_layers)])
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        for layer in self.layers:
+            x = layer(x)
+        return x
+
+
+def apply_tp_replicate(block: NormMLP, tp_mesh: DeviceMesh) -> None:
+    parallelize_module(
+        block,
+        tp_mesh,
+        {
+            "fc1": ColwiseParallel(input_layouts=Replicate(), use_local_output=False),
+            "fc2": RowwiseParallel(output_layouts=Replicate(), use_local_output=False),
+        },
+    )
+    for name, param in block.norm.named_parameters():
+        block.norm.register_parameter(
+            name,
+            nn.Parameter(
+                DTensor.from_local(param, tp_mesh, [Replicate()], run_check=False)
+            ),
+        )
+
+
+def apply_tp_shard(block: NormMLP, tp_mesh: DeviceMesh) -> None:
+    parallelize_module(
+        block,
+        tp_mesh,
+        {
+            "norm": SequenceParallel(use_local_output=False),
+            "fc1": ColwiseParallel(
+                input_layouts=Shard(1),
+                use_local_output=False,  # type: ignore[arg-type]
+            ),
+            "fc2": RowwiseParallel(
+                output_layouts=Shard(1),
+                use_local_output=False,  # type: ignore[arg-type]
+            ),
+        },
+    )
+
+
+def _loss_fn(output: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
+    return (output - target).pow(2).mean()
+
+
+def _requires_multi_gpu(func):
+    @requires_accelerator_dist_backend(["nccl", "xccl"])
+    @skip_but_pass_in_sandcastle_if(
+        not TEST_MULTIACCELERATOR, f"{backend} test requires 4+ GPUs"
+    )
+    @functools.wraps(func)
+    def wrapper(*args, **kwargs):
+        return func(*args, **kwargs)
+
+    return wrapper
+
+
+class DTensorPPIntegrationBase(MultiProcContinuousTest):
+    world_size = 4
+
+    @classmethod
+    def backend_str(cls) -> str:
+        return backend
+
+    @classmethod
+    def device_type(cls) -> str:
+        return device_type
+
+    @property
+    def device(self) -> torch.device:
+        return torch.device(device_type, self.rank)
+
+    def init_pg(self) -> None:
+        if device_type == "cuda":
+            torch.cuda.set_device(self.device)
+
+    def _make_mesh(self) -> DeviceMesh:
+        return init_device_mesh(device_type, (2, 2), mesh_dim_names=("pp", "tp"))
+
+    def _build_baseline_and_clones(
+        self,
+        n_layers: int,
+        tp_mesh: DeviceMesh,
+        apply_tp: Callable,
+    ) -> tuple[NormMLPStack, NormMLPStack]:
+        torch.manual_seed(MODEL_SEED)
+        with torch.device(self.device):
+            ref_model = NormMLPStack(d_hid, n_layers)
+
+        pp_model = copy.deepcopy(ref_model)
+
+        for layer in ref_model.layers:
+            apply_tp(layer, tp_mesh)
+        for layer in pp_model.layers:
+            apply_tp(layer, tp_mesh)
+
+        return ref_model, pp_model
+
+    def _make_full_batch_tensors(self) -> tuple[torch.Tensor, torch.Tensor]:
+        torch.manual_seed(INPUT_SEED)
+        input_full = torch.randn(batch_size, d_hid, device=self.device)
+        torch.manual_seed(TARGET_SEED)
+        target_full = torch.randn(batch_size, d_hid, device=self.device)
+        return input_full, target_full
+
+    def _run_reference_microbatched(
+        self,
+        ref_model: NormMLPStack,
+        input_dt: DTensor,
+        target_dt: DTensor,
+        *,
+        null_boundary_grads: bool = True,
+    ) -> tuple[
+        dict[str, torch.Tensor | None],
+        StageStaticMetaMap,
+    ]:
+        ref_model.zero_grad(set_to_none=True)
+
+        input_chunks = torch.tensor_split(input_dt, n_microbatches)
+        target_chunks = torch.tensor_split(target_dt, n_microbatches)
+
+        captured_io: dict[int, tuple[DTensor, DTensor]] = {}
+        captured_grads: dict[int, dict[str, DTensor | None]] = {}
+        hooks: list[torch.utils.hooks.RemovableHandle] = []
+
+        def _make_grad_hook(stage_idx: int, key: str) -> Callable[[torch.Tensor], None]:
+            def _hook(grad: torch.Tensor) -> None:
+                captured_grads[stage_idx][key] = self._clone_preserving_placement(grad)
+
+            return _hook
+
+        for stage_idx, stage_mod in enumerate(ref_model.layers):
+
+            def _capture_io(
+                _module: nn.Module,
+                args: tuple[torch.Tensor, ...],
+                output: torch.Tensor,
+                *,
+                idx: int = stage_idx,
+            ) -> None:
+                stage_input = cast(DTensor, args[0])
+                stage_output = cast(DTensor, output)
+                captured_io[idx] = (stage_input, stage_output)
+                captured_grads[idx] = {"input": None, "output": None}
+
+                for tensor, key in (
+                    (stage_input, "input"),
+                    (stage_output, "output"),
+                ):
+                    if tensor.requires_grad:
+                        hooks.append(tensor.register_hook(_make_grad_hook(idx, key)))
+
+            hooks.append(stage_mod.register_forward_hook(_capture_io))
+
+        static_stage_meta: StageStaticMetaMap = {}
+
+        for mb_idx, (input_chunk, target_chunk) in enumerate(
+            zip(input_chunks, target_chunks)
+        ):
+            output = ref_model(input_chunk)
+            loss = _loss_fn(output, target_chunk)
+            loss.backward()
+
+            if mb_idx == 0:
+                for stage_idx in range(len(ref_model.layers)):
+                    stage_input, stage_output = captured_io[stage_idx]
+                    input_args = self._empty_dt_from(
+                        stage_input,
+                        stage_input.requires_grad,
+                    )
+                    output_args = self._empty_dt_from(
+                        stage_output,
+                        stage_output.requires_grad,
+                    )
+                    input_grad_dt = captured_grads[stage_idx]["input"]
+                    output_grad_dt = captured_grads[stage_idx]["output"]
+                    input_grads = (
+                        self._empty_dt_from(cast(DTensor, input_grad_dt), False)
+                        if input_grad_dt is not None
+                        else None
+                    )
+                    output_grads = (
+                        self._empty_dt_from(cast(DTensor, output_grad_dt), False)
+                        if output_grad_dt is not None
+                        else None
+                    )
+
+                    if null_boundary_grads:
+                        if stage_idx == 0:
+                            input_grads = None
+                        if stage_idx == (len(ref_model.layers) - 1):
+                            output_grads = None
+
+                    static_stage_meta[stage_idx] = (
+                        input_args,
+                        output_args,
+                        input_grads,
+                        output_grads,
+                    )
+
+                for hook in hooks:
+                    hook.remove()
+                hooks = []
+
+        for param in ref_model.parameters():
+            if param.grad is not None:
+                param.grad.div_(n_microbatches)
+
+        ref_grads = {
+            name: self._clone_preserving_placement(grad) if grad is not None else grad
+            for name, grad in (
+                (param_name, param.grad)
+                for param_name, param in ref_model.named_parameters()
+            )
+        }
+        return ref_grads, static_stage_meta
+
+    @staticmethod
+    def _clone_preserving_placement(grad: torch.Tensor) -> torch.Tensor:
+        """Clone a gradient tensor, preserving DTensor placement.
+
+        ``DTensor.clone()`` triggers all-reduce for ``Partial`` placements,
+        converting them to ``Replicate``.  This helper clones only the local
+        tensor and reconstructs the DTensor with the original placement so
+        that downstream comparisons see the true placement.
+        """
+        if isinstance(grad, DTensor):
+            local_clone = grad.to_local().clone()
+            return DTensor.from_local(
+                local_clone,
+                grad.device_mesh,
+                list(grad.placements),
+                shape=grad.shape,
+                stride=grad.stride(),
+                run_check=False,
+            )
+        return grad.clone()
+
+    def _empty_dt_from(self, dt: DTensor, requires_grad: bool) -> DTensor:
+        meta = _DTensorMeta.from_dtensor(dt)
+        empty_local = torch.empty(meta.shape, dtype=meta.dtype, device=self.device)
+        result = DTensor.from_local(
+            empty_local,
+            dt.device_mesh,
+            list(meta.placements),
+            shape=meta.global_shape,
+            stride=meta.global_stride,
+            run_check=False,
+        )
+        if requires_grad:
+            result.requires_grad_(True)
+        return result
+
+    def _make_stage_for_pp_run(
+        self,
+        pp_model: NormMLPStack,
+        stage_index: int,
+        num_stages: int,
+        pp_group: dist.ProcessGroup,
+        tp_mesh: DeviceMesh,
+        static_stage_meta: StageStaticMetaMap | None,
+    ) -> PipelineStage:
+        stage_module = pp_model.get_submodule(f"layers.{stage_index}")
+        if static_stage_meta is None:
+            return PipelineStage(
+                submodule=stage_module,
+                stage_index=stage_index,
+                num_stages=num_stages,
+                device=self.device,
+                group=pp_group,
+                get_mesh=lambda _dim_names, _layout: tp_mesh,
+            )
+
+        input_args, output_args, input_grads, output_grads = static_stage_meta[
+            stage_index
+        ]
+
+        return PipelineStage(
+            submodule=stage_module,
+            stage_index=stage_index,
+            num_stages=num_stages,
+            device=self.device,
+            group=pp_group,
+            input_args=(input_args,),
+            output_args=(output_args,),
+            input_grads=(input_grads,),
+            output_grads=(output_grads,),
+            get_mesh=lambda _dim_names, _layout: tp_mesh,
+        )
+
+    def _assert_grad_close(
+        self,
+        pp_grad: torch.Tensor,
+        ref_grad: torch.Tensor,
+        param_fqn: str,
+    ) -> None:
+        pp_for_compare = pp_grad
+        ref_for_compare = ref_grad
+
+        if isinstance(pp_grad, DTensor) and isinstance(ref_grad, DTensor):
+            pp_meta = _DTensorMeta.from_dtensor(pp_grad)
+            ref_meta = _DTensorMeta.from_dtensor(ref_grad)
+            self.assertEqual(
+                pp_meta.get_diff(ref_meta),
+                [],
+                f"DTensor metadata mismatch for {param_fqn}",
+            )
+            pp_for_compare = pp_grad.to_local()
+            ref_for_compare = ref_grad.to_local()
+
+        torch.testing.assert_close(
+            pp_for_compare,
+            ref_for_compare,
+            msg=f"Gradient mismatch for {param_fqn}",
+        )
+
+    def _assert_stage_grad_parity(
+        self,
+        pp_model: NormMLPStack,
+        ref_grads: dict[str, torch.Tensor | None],
+        stage_index: int,
+    ) -> None:
+        stage_module = pp_model.get_submodule(f"layers.{stage_index}")
+        for name, param in stage_module.named_parameters():
+            param_fqn = f"layers.{stage_index}.{name}"
+            self.assertIsNotNone(param.grad, f"Missing PP grad for {param_fqn}")
+            self.assertIsNotNone(
+                ref_grads[param_fqn], f"Missing reference grad for {param_fqn}"
+            )
+            pp_grad = cast(torch.Tensor, param.grad)
+            ref_grad = cast(torch.Tensor, ref_grads[param_fqn])
+            self._assert_grad_close(pp_grad, ref_grad, param_fqn)
+
+    def _compute_stage_indices(
+        self,
+        schedule_class: type,
+        pp_rank: int,
+        pp_size: int,
+        num_stages: int,
+    ) -> list[int]:
+        if schedule_class is Schedule1F1B:
+            return [pp_rank]
+        if schedule_class is ScheduleInterleaved1F1B:
+            return [pp_rank + i * pp_size for i in range(num_stages // pp_size)]
+        return [pp_rank, num_stages - 1 - pp_rank]
+
+    def _execute_schedule_step(
+        self,
+        schedule,
+        stages: list[PipelineStage],
+        pp_input: DTensor,
+        pp_target: DTensor,
+    ) -> None:
+        has_first = any(stage.is_first for stage in stages)
+        has_last = any(stage.is_last for stage in stages)
+
+        if has_first and has_last:
+            losses: list[torch.Tensor] = []
+            schedule.step(pp_input, target=pp_target, losses=losses)
+        elif has_first:
+            schedule.step(pp_input)
+        elif has_last:
+            losses: list[torch.Tensor] = []
+            schedule.step(target=pp_target, losses=losses)
+        else:
+            schedule.step()
+
+    def _run_training_correctness(
+        self,
+        schedule_class: type,
+        apply_tp: Callable,
+        placements: list,
+        *,
+        use_static_metadata: bool = True,
+        null_boundary_grads: bool = False,
+    ) -> None:
+        self.init_pg()
+
+        mesh = self._make_mesh()
+        tp_mesh = mesh["tp"]
+        pp_mesh = mesh["pp"]
+        pp_group = pp_mesh.get_group()
+        pp_rank = pp_mesh.get_local_rank()
+        pp_size = pp_mesh.size()
+
+        n_layers = 2 if schedule_class is Schedule1F1B else 4
+
+        ref_model, pp_model = self._build_baseline_and_clones(
+            n_layers=n_layers,
+            tp_mesh=tp_mesh,
+            apply_tp=apply_tp,
+        )
+
+        input_full, target_full = self._make_full_batch_tensors()
+
+        ref_input = distribute_tensor(
+            input_full.clone(),
+            tp_mesh,
+            placements,
+        )
+        ref_input.requires_grad_(True)
+        ref_target = distribute_tensor(
+            target_full.clone(),
+            tp_mesh,
+            placements,
+        )
+
+        ref_grads, static_stage_meta = self._run_reference_microbatched(
+            ref_model,
+            ref_input,
+            ref_target,
+            null_boundary_grads=null_boundary_grads,
+        )
+
+        pp_static_stage_meta = static_stage_meta if use_static_metadata else None
+
+        pp_input = distribute_tensor(
+            input_full.clone(),
+            tp_mesh,
+            placements,
+        )
+        pp_input.requires_grad_(True)
+        pp_target = distribute_tensor(
+            target_full.clone(),
+            tp_mesh,
+            placements,
+        )
+
+        num_stages = n_layers
+        stage_indices = self._compute_stage_indices(
+            schedule_class,
+            pp_rank,
+            pp_size,
+            num_stages,
+        )
+
+        stages = [
+            self._make_stage_for_pp_run(
+                pp_model=pp_model,
+                stage_index=stage_index,
+                num_stages=num_stages,
+                pp_group=pp_group,
+                tp_mesh=tp_mesh,
+                static_stage_meta=pp_static_stage_meta,
+            )
+            for stage_index in stage_indices
+        ]
+
+        for stage in stages:
+            stage.submod.zero_grad(set_to_none=True)
+
+        if schedule_class is Schedule1F1B:
+            schedule = schedule_class(
+                stages[0],
+                n_microbatches=n_microbatches,
+                loss_fn=_loss_fn,
+            )
+        else:
+            schedule = schedule_class(
+                cast(list[PipelineStage], stages),  # type: ignore[arg-type]
+                n_microbatches=n_microbatches,
+                loss_fn=_loss_fn,
+            )
+
+        self._execute_schedule_step(schedule, stages, pp_input, pp_target)
+
+        for stage_index in stage_indices:
+            self._assert_stage_grad_parity(
+                pp_model=pp_model,
+                ref_grads=ref_grads,
+                stage_index=stage_index,
+            )
+
+    def _run_inference_only_equivalence(
+        self,
+        apply_tp: Callable,
+        placements: list,
+        *,
+        use_static_metadata: bool = False,
+    ) -> None:
+        self.init_pg()
+
+        mesh = self._make_mesh()
+        tp_mesh = mesh["tp"]
+        pp_mesh = mesh["pp"]
+        pp_group = pp_mesh.get_group()
+        pp_rank = pp_mesh.get_local_rank()
+
+        ref_model, pp_model = self._build_baseline_and_clones(
+            n_layers=2,
+            tp_mesh=tp_mesh,
+            apply_tp=apply_tp,
+        )
+
+        input_full, target_full = self._make_full_batch_tensors()
+
+        ref_input = distribute_tensor(
+            input_full.clone(),
+            tp_mesh,
+            placements,
+        )
+
+        static_stage_meta: StageStaticMetaMap | None = None
+
+        if use_static_metadata:
+            ref_input_for_meta = distribute_tensor(
+                input_full.clone(),
+                tp_mesh,
+                placements,
+            )
+            ref_input_for_meta.requires_grad_(True)
+            ref_target_for_meta = distribute_tensor(
+                target_full.clone(),
+                tp_mesh,
+                placements,
+            )
+            _, captured_meta = self._run_reference_microbatched(
+                ref_model,
+                ref_input_for_meta,
+                ref_target_for_meta,
+            )
+
+            static_stage_meta = {
+                stage_idx: (
+                    self._empty_dt_from(stage_meta[0], False),
+                    self._empty_dt_from(stage_meta[1], False),
+                    None,
+                    None,
+                )
+                for stage_idx, stage_meta in captured_meta.items()
+            }
+
+            with torch.no_grad():
+                ref_output = ref_model(ref_input)
+        else:
+            with torch.no_grad():
+                ref_output = ref_model(ref_input)
+
+        pp_input = distribute_tensor(
+            input_full.clone(),
+            tp_mesh,
+            placements,
+        )
+        pp_input.requires_grad_(False)
+
+        stage = self._make_stage_for_pp_run(
+            pp_model=pp_model,
+            stage_index=pp_rank,
+            num_stages=2,
+            pp_group=pp_group,
+            tp_mesh=tp_mesh,
+            static_stage_meta=static_stage_meta,
+        )
+        schedule = Schedule1F1B(
+            stage,
+            n_microbatches=n_microbatches,
+            loss_fn=None,
+        )
+
+        with torch.no_grad():
+            pp_output = schedule.eval(pp_input)
+
+        if stage.is_last:
+            self.assertIsNotNone(pp_output)
+            self.assertIsInstance(pp_output, DTensor)
+            self.assertIsInstance(ref_output, DTensor)
+            self._assert_grad_close(
+                cast(torch.Tensor, pp_output),
+                cast(torch.Tensor, ref_output),
+                "inference_output",
+            )
+
+
+class TestDTensorPPModes(DTensorPPIntegrationBase):
+    @_requires_multi_gpu
+    def test_static_mode_replicate(self):
+        self._run_training_correctness(
+            Schedule1F1B,
+            apply_tp_replicate,
+            [Replicate()],
+        )
+
+    @_requires_multi_gpu
+    def test_static_mode_none_grad_slots_replicate(self):
+        self._run_training_correctness(
+            Schedule1F1B,
+            apply_tp_replicate,
+            [Replicate()],
+            null_boundary_grads=True,
+        )
+
+    @_requires_multi_gpu
+    def test_static_mode_none_grad_slots_shard(self):
+        self._run_training_correctness(
+            Schedule1F1B,
+            apply_tp_shard,
+            [Shard(1)],
+            null_boundary_grads=True,
+        )
+
+    @_requires_multi_gpu
+    def test_static_mode_shard(self):
+        self._run_training_correctness(
+            Schedule1F1B,
+            apply_tp_shard,
+            [Shard(1)],
+        )
+
+    @_requires_multi_gpu
+    def test_dynamic_mode_inference_only_replicate(self):
+        self._run_inference_only_equivalence(
+            apply_tp_replicate,
+            [Replicate()],
+        )
+
+    @_requires_multi_gpu
+    def test_static_mode_inference_only_replicate(self):
+        self._run_inference_only_equivalence(
+            apply_tp_replicate,
+            [Replicate()],
+            use_static_metadata=True,
+        )
+
+    @_requires_multi_gpu
+    def test_dynamic_mode_inference_only_shard(self):
+        self._run_inference_only_equivalence(
+            apply_tp_shard,
+            [Shard(1)],
+        )
+
+    @_requires_multi_gpu
+    def test_static_mode_inference_only_shard(self):
+        self._run_inference_only_equivalence(
+            apply_tp_shard,
+            [Shard(1)],
+            use_static_metadata=True,
+        )
+
+    @_requires_multi_gpu
+    def test_dynamic_mode_replicate(self):
+        self._run_training_correctness(
+            Schedule1F1B,
+            apply_tp_replicate,
+            [Replicate()],
+            use_static_metadata=False,
+        )
+
+    @_requires_multi_gpu
+    def test_dynamic_mode_shard(self):
+        self._run_training_correctness(
+            Schedule1F1B,
+            apply_tp_shard,
+            [Shard(1)],
+            use_static_metadata=False,
+        )
+
+    @_requires_multi_gpu
+    def test_static_mode_interleaved1f1b_replicate(self):
+        self._run_training_correctness(
+            ScheduleInterleaved1F1B,
+            apply_tp_replicate,
+            [Replicate()],
+        )
+
+    @_requires_multi_gpu
+    def test_static_mode_zbv_zero_bubble_replicate(self):
+        self._run_training_correctness(
+            ScheduleZBVZeroBubble,
+            apply_tp_replicate,
+            [Replicate()],
+        )
+
+    @_requires_multi_gpu
+    def test_static_mode_dualpipev_replicate(self):
+        self._run_training_correctness(
+            ScheduleDualPipeV,
+            apply_tp_replicate,
+            [Replicate()],
+        )
+
+    @_requires_multi_gpu
+    def test_dynamic_mode_interleaved1f1b_shard(self):
+        self._run_training_correctness(
+            ScheduleInterleaved1F1B,
+            apply_tp_shard,
+            [Shard(1)],
+            use_static_metadata=False,
+        )
+
+    @_requires_multi_gpu
+    def test_dynamic_mode_zbv_zero_bubble_shard(self):
+        self._run_training_correctness(
+            ScheduleZBVZeroBubble,
+            apply_tp_shard,
+            [Shard(1)],
+            use_static_metadata=False,
+        )
+
+    @_requires_multi_gpu
+    def test_dynamic_mode_dualpipev_shard(self):
+        self._run_training_correctness(
+            ScheduleDualPipeV,
+            apply_tp_shard,
+            [Shard(1)],
+            use_static_metadata=False,
+        )
+
+
+if __name__ == "__main__":
+    run_tests()
diff --git a/test/distributed/pipelining/test_dtensor_pp_unit_tests.py b/test/distributed/pipelining/test_dtensor_pp_unit_tests.py
new file mode 100644
index 0000000000000..df06b1eb50e78
--- /dev/null
+++ b/test/distributed/pipelining/test_dtensor_pp_unit_tests.py
@@ -0,0 +1,665 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates
+# Owner(s): ["oncall: distributed"]
+"""
+Unit tests for DTensor support in Pipeline Parallelism.
+
+Tests the _utils.py metadata infrastructure and microbatch DTensor operations:
+1. Metadata classes: _TensorMeta, _DTensorMeta, _MeshCache
+2. InferenceMode decision logic
+3. Validation functions: validate_metadata, validate_tensors_metadata,
+   validate_static_arg_grad_correspondence
+4. Helper functions: extract_tensor_metas, to_local_if_dtensor,
+   validate_and_normalize_to_tuple
+5. Microbatch DTensor operations: _split_tensor, merge_chunks
+"""
+
+import functools
+import warnings
+
+import torch
+import torch.distributed as dist
+from torch.distributed.device_mesh import init_device_mesh
+from torch.distributed.pipelining._utils import (
+    _DTensorMeta,
+    _MeshCache,
+    _StageMeta,
+    _TensorMeta,
+    extract_tensor_metas,
+    InferenceMode,
+    PipeliningMetadataError,
+    to_local_if_dtensor,
+    validate_and_normalize_to_tuple,
+    validate_metadata,
+    validate_static_arg_grad_correspondence,
+    validate_tensors_metadata,
+)
+from torch.distributed.pipelining.microbatch import (
+    _split_tensor,
+    merge_chunks,
+    TensorChunkSpec,
+)
+from torch.distributed.tensor import distribute_tensor, DTensor, Replicate, Shard
+from torch.testing._internal.common_distributed import (
+    MultiProcContinuousTest,
+    requires_accelerator_dist_backend,
+)
+from torch.testing._internal.common_utils import (
+    run_tests,
+    skip_but_pass_in_sandcastle_if,
+    TEST_MULTIACCELERATOR,
+)
+
+
+# =============================================================================
+# Test Constants
+# =============================================================================
+
+d_hid = 256
+batch_size = 64
+n_microbatches = 4
+microbatch_size = batch_size // n_microbatches
+
+device_type = acc.type if (acc := torch.accelerator.current_accelerator()) else "cpu"
+backend = dist.get_default_backend_for_device(device_type)
+
+
+def _requires_multi_gpu(fn):
+    @requires_accelerator_dist_backend(["nccl", "xccl"])
+    @skip_but_pass_in_sandcastle_if(
+        not TEST_MULTIACCELERATOR, f"{backend} test requires 4+ GPUs"
+    )
+    @functools.wraps(fn)
+    def wrapper(*args, **kwargs):
+        return fn(*args, **kwargs)
+
+    return wrapper
+
+
+# =============================================================================
+# All DTensor PP Unit Tests (single class → single PG initialization)
+# =============================================================================
+
+
+class TestDTensorPPUnitTests(MultiProcContinuousTest):
+    """Unit tests for DTensor PP metadata infrastructure.
+
+    All tests live in a single class so the process group is initialized once
+    rather than once per test category.
+    """
+
+    @classmethod
+    def backend_str(cls) -> str:
+        return backend
+
+    @classmethod
+    def device_type(cls) -> str:
+        return device_type
+
+    @property
+    def device(self) -> torch.device:
+        return torch.device(device_type, self.rank)
+
+    def init_pg(self):
+        if device_type == "cuda":
+            torch.cuda.set_device(self.device)
+
+    # -----------------------------------------------------------------
+    # Shared helpers
+    # -----------------------------------------------------------------
+
+    def _make_mesh(self, dim_names=("tp",)):
+        """Create a 1D device mesh spanning all ranks."""
+        return init_device_mesh(
+            device_type, (self.world_size,), mesh_dim_names=dim_names
+        )
+
+    def _make_dtensor(self, mesh, placements, shape=(8, 16), requires_grad=False):
+        """Create a DTensor with given properties."""
+        t = torch.randn(*shape, device=self.device, requires_grad=requires_grad)
+        return distribute_tensor(t, mesh, placements)
+
+    def _make_tensor_meta(
+        self,
+        shape: torch.Size = torch.Size([4, 8]),
+        stride: tuple[int, ...] = (8, 1),
+        dtype: torch.dtype = torch.float32,
+        requires_grad: bool = True,
+    ) -> _TensorMeta:
+        """Create a _TensorMeta with sensible defaults."""
+        return _TensorMeta(
+            shape=shape, stride=stride, dtype=dtype, requires_grad=requires_grad
+        )
+
+    def _make_dtensor_meta(self) -> _DTensorMeta:
+        """Create a _DTensorMeta with sensible defaults."""
+        return _DTensorMeta(
+            shape=torch.Size([4, 8]),
+            stride=(8, 1),
+            dtype=torch.float32,
+            requires_grad=True,
+            global_shape=torch.Size([8, 8]),
+            global_stride=(8, 1),
+            placements=(Shard(0),),
+            mesh_dim_names=("tp",),
+            mesh_layout=None,
+        )
+
+    # -----------------------------------------------------------------
+    # Category 1: Metadata Classes (_TensorMeta, _DTensorMeta, _MeshCache)
+    # -----------------------------------------------------------------
+
+    @_requires_multi_gpu
+    def test_tensor_meta_extraction_roundtrip_and_errors(self):
+        """Test _TensorMeta: extraction, roundtrip, DTensor rejection."""
+        self.init_pg()
+        mesh = self._make_mesh()
+
+        # Extraction and roundtrip for various shapes/dtypes/requires_grad
+        for shape in [(4, 8), (2, 3, 4)]:
+            for dtype in [torch.float32, torch.float16]:
+                for requires_grad in [True, False]:
+                    t = torch.randn(*shape, dtype=dtype, requires_grad=requires_grad)
+                    meta = _TensorMeta.from_tensor(t)
+                    self.assertEqual(meta.shape, t.shape)
+                    self.assertEqual(meta.stride, t.stride())
+                    self.assertEqual(meta.dtype, dtype)
+                    self.assertEqual(meta.requires_grad, requires_grad)
+
+                    # Roundtrip via tensor on device
+                    reconstructed = meta.to_tensor(self.device)
+                    reconstructed_meta = _TensorMeta(
+                        shape=reconstructed.shape,
+                        stride=reconstructed.stride(),
+                        dtype=reconstructed.dtype,
+                        requires_grad=reconstructed.requires_grad,
+                    )
+                    self.assertEqual(meta.get_diff(reconstructed_meta), [])
+
+        # DTensor rejection
+        dt = self._make_dtensor(mesh, [Shard(0)])
+        with self.assertRaises(PipeliningMetadataError) as ctx:
+            _TensorMeta.from_tensor(dt)
+        self.assertIn("DTensor", str(ctx.exception))
+
+    @_requires_multi_gpu
+    def test_dtensor_meta_extraction_and_roundtrip(self):
+        """Test _DTensorMeta: extraction and roundtrip for Shard/Replicate."""
+        self.init_pg()
+        mesh = self._make_mesh()
+
+        for placements in [[Shard(0)], [Replicate()]]:
+            dt = self._make_dtensor(mesh, placements)
+            meta = _DTensorMeta.from_dtensor(dt)
+
+            self.assertEqual(meta.global_shape, dt.shape)
+            self.assertEqual(meta.placements, tuple(placements))
+            self.assertEqual(meta.mesh_dim_names, ("tp",))
+            self.assertEqual(meta.mesh_cache_key, (("tp",), mesh._layout))
+
+            # Roundtrip
+            reconstructed = meta.to_dtensor(self.device, mesh)
+            reconstructed_meta = _DTensorMeta.from_dtensor(reconstructed)
+            self.assertEqual(meta.get_diff(reconstructed_meta), [])
+
+    @_requires_multi_gpu
+    def test_get_diff_detects_mismatches(self):
+        """Test get_diff() detects shape/dtype/placement/cross-type differences."""
+        self.init_pg()
+        mesh = self._make_mesh()
+
+        # _TensorMeta: shape
+        m1 = self._make_tensor_meta(shape=torch.Size([4, 8]), stride=(8, 1))
+        m2 = self._make_tensor_meta(shape=torch.Size([4, 16]), stride=(16, 1))
+        self.assertTrue(any("shape" in d for d in m1.get_diff(m2)))
+
+        # _TensorMeta: dtype
+        m3 = self._make_tensor_meta(dtype=torch.float16)
+        self.assertTrue(any("dtype" in d for d in m1.get_diff(m3)))
+
+        # _DTensorMeta: placement
+        dt_shard = self._make_dtensor(mesh, [Shard(0)])
+        dt_rep = self._make_dtensor(mesh, [Replicate()])
+        dm1 = _DTensorMeta.from_dtensor(dt_shard)
+        dm2 = _DTensorMeta.from_dtensor(dt_rep)
+        self.assertTrue(any("placements" in d for d in dm1.get_diff(dm2)))
+
+        # Cross-type: _DTensorMeta vs _TensorMeta
+        diffs = dm1.get_diff(m1)
+        self.assertTrue(any("type" in d.lower() or "_TensorMeta" in d for d in diffs))
+
+    @_requires_multi_gpu
+    def test_mesh_cache_operations(self):
+        """Test _MeshCache: __contains__/put/get_mesh/callback/update_from_tensors."""
+        self.init_pg()
+        mesh = self._make_mesh()
+        key = (("tp",), mesh._layout)
+
+        # --- __contains__ + put + get_mesh: basic cache hit ---
+        cache = _MeshCache()
+        self.assertFalse(key in cache)
+        cache.put(key, mesh)
+        self.assertTrue(key in cache)
+        self.assertIs(cache.get_mesh(key), mesh)
+
+        # --- get_mesh_cb: callback NOT called on cache hit ---
+        cb_called = [False]
+
+        def cb(mesh_dim_names, mesh_layout):
+            cb_called[0] = True
+            return mesh
+
+        cache_with_cb = _MeshCache(get_mesh_cb=cb)
+        cache_with_cb.put(key, mesh)
+        self.assertIs(cache_with_cb.get_mesh(key), mesh)
+        self.assertFalse(cb_called[0])
+
+        # --- get_mesh_cb: callback IS called on cache miss ---
+        cb2_called = [False]
+
+        def cb2(mesh_dim_names, mesh_layout):
+            cb2_called[0] = True
+            return mesh
+
+        cache_miss = _MeshCache(get_mesh_cb=cb2)
+        self.assertIs(cache_miss.get_mesh(key), mesh)
+        self.assertTrue(cb2_called[0])
+
+        # --- no callback: get_mesh on miss raises ---
+        cache_no_cb = _MeshCache()
+        with self.assertRaises(PipeliningMetadataError) as ctx:
+            cache_no_cb.get_mesh(key)
+        self.assertIn("get_mesh", str(ctx.exception))
+
+        # --- update_from_tensors: extracts mesh from DTensors ---
+        cache_upd = _MeshCache()
+        dt = self._make_dtensor(mesh, [Shard(0)], shape=(4, 4))
+        plain = torch.randn(4, 4, device=self.device)
+        cache_upd.update_from_tensors((dt, plain, None))
+        self.assertTrue(key in cache_upd)
+        self.assertIs(cache_upd.get_mesh(key), mesh)
+        self.assertEqual(len(cache_upd), 1)
+
+    # -----------------------------------------------------------------
+    # Category 2: InferenceMode Decision Logic
+    # -----------------------------------------------------------------
+
+    @_requires_multi_gpu
+    def test_needs_dynamic_all_cases(self):
+        """Test all 8 cases of the needs_dynamic() decision matrix."""
+        self.init_pg()
+        tm = self._make_tensor_meta()
+        dm = self._make_dtensor_meta()
+
+        cases = [
+            # (meta, has_backward, expected_needs_dynamic)
+            # Case 1: No forward metadata
+            (_StageMeta(inputs=None, outputs=None), True, True),
+            # Case 2: Partial forward (inputs only)
+            (_StageMeta(inputs=(tm,), outputs=None), True, True),
+            # Case 3: Plain tensors, complete forward, backward → STATIC
+            (_StageMeta(inputs=(tm,), outputs=(tm,)), True, False),
+            # Case 4: DTensors, no backward → STATIC
+            (_StageMeta(inputs=(dm,), outputs=(dm,)), False, False),
+            # Case 5: DTensors, backward, no grads → DYNAMIC
+            (_StageMeta(inputs=(dm,), outputs=(dm,)), True, True),
+            # Case 6: DTensors, backward, only input_grads → DYNAMIC
+            (
+                _StageMeta(
+                    inputs=(dm,), outputs=(dm,), input_grads=(dm,), output_grads=None
+                ),
+                True,
+                True,
+            ),
+            # Case 7: DTensors, backward, only output_grads → DYNAMIC
+            (
+                _StageMeta(
+                    inputs=(dm,), outputs=(dm,), input_grads=None, output_grads=(dm,)
+                ),
+                True,
+                True,
+            ),
+            # Case 8: DTensors, backward, complete grads → STATIC
+            (
+                _StageMeta(
+                    inputs=(dm,),
+                    outputs=(dm,),
+                    input_grads=(dm,),
+                    output_grads=(dm,),
+                ),
+                True,
+                False,
+            ),
+        ]
+        for i, (meta, has_bwd, expected) in enumerate(cases, 1):
+            with self.subTest(case=i):
+                self.assertEqual(
+                    InferenceMode.needs_dynamic(meta, has_bwd),
+                    expected,
+                    f"Case {i} failed",
+                )
+
+    # -----------------------------------------------------------------
+    # Category 3: Validation Functions
+    # -----------------------------------------------------------------
+
+    @_requires_multi_gpu
+    def test_validate_metadata(self):
+        """Test validate_metadata: match, raise, warn, type mismatch."""
+        self.init_pg()
+        mesh = self._make_mesh()
+
+        # Match → no diffs
+        dt = self._make_dtensor(mesh, [Shard(0)])
+        meta = _DTensorMeta.from_dtensor(dt)
+        self.assertEqual(validate_metadata("test", meta, dt), [])
+
+        # Mismatch with raise_on_mismatch
+        other = self._make_dtensor(mesh, [Replicate()])
+        with self.assertRaises(PipeliningMetadataError):
+            validate_metadata("test", meta, other, raise_on_mismatch=True)
+
+        # Mismatch with warn_on_mismatch
+        with warnings.catch_warnings(record=True) as w:
+            warnings.simplefilter("always")
+            diffs = validate_metadata("test", meta, other, warn_on_mismatch=True)
+            self.assertTrue(len(diffs) > 0)
+            self.assertTrue(any("Metadata mismatch" in str(x.message) for x in w))
+
+        # Type mismatch: _DTensorMeta vs plain tensor
+        plain = torch.randn(8, 16, device=self.device)
+        diffs = validate_metadata("test", meta, plain)
+        self.assertTrue(any("type" in d for d in diffs))
+
+    @_requires_multi_gpu
+    def test_validate_tensors_metadata(self):
+        """Test validate_tensors_metadata: batch validation, length mismatch, None handling."""
+        self.init_pg()
+        mesh = self._make_mesh()
+
+        dt = self._make_dtensor(mesh, [Shard(0)])
+        meta = _DTensorMeta.from_dtensor(dt)
+
+        # Match
+        validate_tensors_metadata("test", (meta,), (dt,))
+
+        # Length mismatch raises
+        with self.assertRaises(PipeliningMetadataError):
+            validate_tensors_metadata("test", (meta, meta), (dt,))
+
+        # Length mismatch warns (raise disabled)
+        with warnings.catch_warnings(record=True) as w:
+            warnings.simplefilter("always")
+            validate_tensors_metadata(
+                "test",
+                (meta, meta),
+                (dt,),
+                raise_on_mismatch=False,
+                warn_on_mismatch=True,
+            )
+            self.assertTrue(len(w) > 0)
+
+        # None-None pair passes
+        validate_tensors_metadata("test", (None,), (None,))
+
+        # One-None mismatch raises (expected has metadata but actual is None)
+        with self.assertRaises(PipeliningMetadataError):
+            validate_tensors_metadata("test", (meta,), (None,))
+
+    @_requires_multi_gpu
+    def test_validate_static_arg_grad_correspondence(self):
+        """Test DTensor↔grad correspondence: all valid/invalid combos."""
+        self.init_pg()
+        mesh = self._make_mesh()
+
+        dt_grad = self._make_dtensor(mesh, [Shard(0)], requires_grad=True)
+        dt_nograd = self._make_dtensor(mesh, [Shard(0)], requires_grad=False)
+        dt_grad2 = self._make_dtensor(mesh, [Replicate()], requires_grad=True)
+        plain = torch.randn(8, 16, device=self.device)
+
+        # DTensor(requires_grad=True) + DTensor grad → OK
+        validate_static_arg_grad_correspondence(
+            0, (dt_grad,), (dt_grad2,), is_input=True
+        )
+
+        # DTensor(requires_grad=False) + None grad → OK
+        validate_static_arg_grad_correspondence(0, (dt_nograd,), (None,), is_input=True)
+
+        # Plain tensor + None grad → OK
+        validate_static_arg_grad_correspondence(0, (plain,), (None,), is_input=True)
+
+        # DTensor(requires_grad=True) + None grad → OK (boundary with no gradient)
+        validate_static_arg_grad_correspondence(0, (dt_grad,), (None,), is_input=True)
+
+        # DTensor(requires_grad=True) + plain tensor grad → ERROR
+        with self.assertRaises(PipeliningMetadataError):
+            validate_static_arg_grad_correspondence(
+                0, (dt_grad,), (plain,), is_input=True
+            )
+
+        # Length mismatch → ERROR
+        with self.assertRaises(PipeliningMetadataError):
+            validate_static_arg_grad_correspondence(
+                0, (dt_grad, dt_grad), (dt_grad2,), is_input=False
+            )
+
+    # -----------------------------------------------------------------
+    # Category 4: Helper Functions
+    # -----------------------------------------------------------------
+
+    @_requires_multi_gpu
+    def test_extract_tensor_metas(self):
+        """Test extract_tensor_metas: plain, DTensor, None handling."""
+        self.init_pg()
+        mesh = self._make_mesh()
+
+        dt = self._make_dtensor(mesh, [Shard(0)])
+        plain = torch.randn(4, 4, device=self.device)
+
+        # None input → None output
+        self.assertIsNone(extract_tensor_metas(None))
+
+        # Mixed DTensor + plain → correct meta types
+        metas = extract_tensor_metas((dt, plain))
+        self.assertIsNotNone(metas)
+        self.assertIsInstance(metas[0], _DTensorMeta)  # type: ignore[index]
+        self.assertIsInstance(metas[1], _TensorMeta)  # type: ignore[index]
+        self.assertNotIsInstance(metas[1], _DTensorMeta)  # type: ignore[index]
+
+        # allow_none=True preserves None
+        metas_with_none = extract_tensor_metas((dt, None), allow_none=True)
+        self.assertIsNotNone(metas_with_none)
+        self.assertIsInstance(metas_with_none[0], _DTensorMeta)  # type: ignore[index]
+        self.assertIsNone(metas_with_none[1])  # type: ignore[index]
+
+        # allow_none=False (default) rejects None
+        with self.assertRaises(PipeliningMetadataError):
+            extract_tensor_metas((dt, None))  # type: ignore[arg-type]
+
+    @_requires_multi_gpu
+    def test_to_local_if_dtensor(self):
+        """Test to_local_if_dtensor: DTensor→local, plain passthrough, detach."""
+        self.init_pg()
+        mesh = self._make_mesh()
+
+        dt = self._make_dtensor(mesh, [Shard(0)], requires_grad=True)
+        plain = torch.randn(4, 4, device=self.device, requires_grad=True)
+
+        # DTensor → local tensor (not DTensor)
+        local = to_local_if_dtensor(dt)
+        self.assertNotIsInstance(local, DTensor)
+        self.assertEqual(local.shape, dt._local_tensor.shape)
+
+        # Plain tensor → same tensor
+        result = to_local_if_dtensor(plain)
+        self.assertIs(result, plain)
+
+        # detach=True → result is detached
+        local_detached = to_local_if_dtensor(dt, detach=True)
+        self.assertNotIsInstance(local_detached, DTensor)
+        self.assertFalse(local_detached.requires_grad)
+
+    @_requires_multi_gpu
+    def test_validate_and_normalize_to_tuple(self):
+        """Test validate_and_normalize_to_tuple: single, tuple, list, None, errors."""
+        self.init_pg()
+        t = torch.randn(4, 4, device=self.device)
+
+        # None → None
+        self.assertIsNone(validate_and_normalize_to_tuple(None))
+
+        # Single tensor → 1-tuple
+        result = validate_and_normalize_to_tuple(t)
+        self.assertIsNotNone(result)
+        self.assertEqual(len(result), 1)  # type: ignore[arg-type]
+        self.assertIs(result[0], t)  # type: ignore[index]
+
+        # List → tuple
+        result = validate_and_normalize_to_tuple([t, t])
+        self.assertIsNotNone(result)
+        self.assertIsInstance(result, tuple)
+        self.assertEqual(len(result), 2)  # type: ignore[arg-type]
+
+        # Non-tensor element → error
+        with self.assertRaises(PipeliningMetadataError):
+            validate_and_normalize_to_tuple((t, "not_a_tensor"))  # type: ignore[arg-type]
+
+        # None in tuple without allow_none → error
+        with self.assertRaises(PipeliningMetadataError):
+            validate_and_normalize_to_tuple((t, None))  # type: ignore[arg-type]
+
+        # None in tuple with allow_none=True → OK
+        result_with_none = validate_and_normalize_to_tuple((t, None), allow_none=True)
+        self.assertIsNotNone(result_with_none)
+        self.assertEqual(len(result_with_none), 2)  # type: ignore[arg-type]
+        self.assertIsNone(result_with_none[1])  # type: ignore[index]
+
+        # Wrong type entirely → error
+        with self.assertRaises(PipeliningMetadataError):
+            validate_and_normalize_to_tuple(42)  # type: ignore[arg-type]
+
+    # -----------------------------------------------------------------
+    # Category 5: Microbatch DTensor Split & Merge
+    # -----------------------------------------------------------------
+
+    @_requires_multi_gpu
+    def test_split_tensor_dtensor_preserves_placements(self):
+        """Test _split_tensor splits a DTensor into chunks that are still DTensors
+        with the same placements, correct local shapes, and matching data."""
+        self.init_pg()
+        mesh = self._make_mesh()
+        num_chunks = 2
+        split_dim = 0
+        shape = (8, 16)
+
+        for placements in [[Shard(0)], [Replicate()]]:
+            dt = self._make_dtensor(mesh, placements, shape=shape)
+            spec = TensorChunkSpec(split_dim)
+            chunks = _split_tensor(dt, spec, num_chunks)
+
+            self.assertEqual(len(chunks), num_chunks)
+            for i, chunk in enumerate(chunks):
+                # Each chunk must be a DTensor with preserved placements
+                self.assertIsInstance(chunk, DTensor, f"chunk {i} is not a DTensor")
+                self.assertEqual(
+                    chunk.placements,
+                    tuple(placements),
+                    f"chunk {i} placements differ",
+                )
+                # Each chunk must have the same device mesh
+                self.assertIs(chunk.device_mesh, mesh)
+
+            # Local shard sizes along split_dim must sum to the original
+            local_split_sizes = [c._local_tensor.size(split_dim) for c in chunks]
+            self.assertEqual(
+                sum(local_split_sizes),
+                dt._local_tensor.size(split_dim),
+            )
+
+            # Full tensor data must be preserved: cat the chunks and compare
+            cat_fn = torch.cat([c._local_tensor for c in chunks], dim=split_dim)
+            self.assertTrue(
+                torch.equal(cat_fn, dt._local_tensor),
+                "Split chunks do not reconstruct the original local tensor",
+            )
+
+    @_requires_multi_gpu
+    def test_merge_chunks_dtensor_roundtrip(self):
+        """Test that split → merge is a roundtrip: merge_chunks reconstructs
+        the original DTensor data and preserves placements."""
+        self.init_pg()
+        mesh = self._make_mesh()
+        num_chunks = 2
+        split_dim = 0
+        shape = (8, 16)
+
+        for placements in [[Shard(0)], [Replicate()]]:
+            dt = self._make_dtensor(mesh, placements, shape=shape)
+            chunk_spec = TensorChunkSpec(split_dim)
+
+            # Split into chunks
+            chunks = list(_split_tensor(dt, chunk_spec, num_chunks))
+
+            # Merge chunks back
+            # merge_chunks expects a list of "outputs", each output is a tuple.
+            # For a single-tensor case, wrap each chunk as a 1-tuple.
+            merged = merge_chunks(
+                [(chunk,) for chunk in chunks],
+                (chunk_spec,),
+            )
+            # merged is a tuple of tensors (mirrors the tuple structure)
+            merged_dt = merged[0]
+
+            self.assertIsInstance(merged_dt, DTensor)
+            self.assertEqual(merged_dt.placements, tuple(placements))
+            self.assertIs(merged_dt.device_mesh, mesh)
+
+            # Data must match the original
+            self.assertTrue(
+                torch.equal(merged_dt._local_tensor, dt._local_tensor),
+                "Merged DTensor local data differs from original",
+            )
+            self.assertEqual(merged_dt.shape, dt.shape)
+
+    @_requires_multi_gpu
+    def test_merge_chunks_dtensor_placement_mismatch_raises(self):
+        """Test that merge_chunks raises when chunk placements don't match."""
+        self.init_pg()
+        mesh = self._make_mesh()
+        shape = (8, 16)
+
+        dt_shard = self._make_dtensor(mesh, [Shard(0)], shape=shape)
+        dt_rep = self._make_dtensor(mesh, [Replicate()], shape=shape)
+
+        chunk_spec = TensorChunkSpec(0)
+        with self.assertRaises(AssertionError) as ctx:
+            merge_chunks(
+                [(dt_shard,), (dt_rep,)],
+                (chunk_spec,),
+            )
+        self.assertIn("placement mismatch", str(ctx.exception))
+
+    @_requires_multi_gpu
+    def test_merge_chunks_dtensor_mixed_types_raises(self):
+        """Test that merge_chunks raises when mixing DTensors and plain tensors."""
+        self.init_pg()
+        mesh = self._make_mesh()
+        shape = (4, 8)
+
+        dt = self._make_dtensor(mesh, [Shard(0)], shape=shape)
+        plain = torch.randn(*shape, device=self.device)
+
+        chunk_spec = TensorChunkSpec(0)
+        with self.assertRaises(AssertionError) as ctx:
+            merge_chunks(
+                [(dt,), (plain,)],
+                (chunk_spec,),
+            )
+        self.assertIn("mix", str(ctx.exception))
+
+
+# =============================================================================
+# Run Tests
+# =============================================================================
+
+if __name__ == "__main__":
+    run_tests()
diff --git a/test/distributed/pipelining/test_microbatch.py b/test/distributed/pipelining/test_microbatch.py
index 09a3e4a5afdaf..063e732a404c1 100644
--- a/test/distributed/pipelining/test_microbatch.py
+++ b/test/distributed/pipelining/test_microbatch.py
@@ -33,15 +33,38 @@ def test_split_and_merge(self):
 
         # Default chunking: dim 0
         arg_chunks, kwarg_chunks = split_args_kwargs_into_chunks(args, kwargs, 2)
-        assert len(arg_chunks) == 2
-        assert len(kwarg_chunks) == 2
-        assert arg_chunks[0][0].shape == torch.Size([64, d_hid])
-        assert arg_chunks[1][0].shape == torch.Size([64, d_hid])
-        assert arg_chunks[0][1].shape == torch.Size([128, d_hid])
-        assert arg_chunks[0][2].shape == torch.Size([256, d_hid])
-        assert kwarg_chunks[0]["x0"].shape == torch.Size([64, d_hid])
-        assert kwarg_chunks[0]["x1"].shape == torch.Size([128, d_hid])
-        assert kwarg_chunks[1]["x2"].shape == torch.Size([256, d_hid])
+        if len(arg_chunks) != 2:
+            raise AssertionError(f"Expected 2 arg_chunks, got {len(arg_chunks)}")
+        if len(kwarg_chunks) != 2:
+            raise AssertionError(f"Expected 2 kwarg_chunks, got {len(kwarg_chunks)}")
+        if arg_chunks[0][0].shape != torch.Size([64, d_hid]):
+            raise AssertionError(
+                f"Expected arg_chunks[0][0].shape == [64, {d_hid}], got {arg_chunks[0][0].shape}"
+            )
+        if arg_chunks[1][0].shape != torch.Size([64, d_hid]):
+            raise AssertionError(
+                f"Expected arg_chunks[1][0].shape == [64, {d_hid}], got {arg_chunks[1][0].shape}"
+            )
+        if arg_chunks[0][1].shape != torch.Size([128, d_hid]):
+            raise AssertionError(
+                f"Expected arg_chunks[0][1].shape == [128, {d_hid}], got {arg_chunks[0][1].shape}"
+            )
+        if arg_chunks[0][2].shape != torch.Size([256, d_hid]):
+            raise AssertionError(
+                f"Expected arg_chunks[0][2].shape == [256, {d_hid}], got {arg_chunks[0][2].shape}"
+            )
+        if kwarg_chunks[0]["x0"].shape != torch.Size([64, d_hid]):
+            raise AssertionError(
+                f"Expected kwarg_chunks[0]['x0'].shape == [64, {d_hid}], got {kwarg_chunks[0]['x0'].shape}"
+            )
+        if kwarg_chunks[0]["x1"].shape != torch.Size([128, d_hid]):
+            raise AssertionError(
+                f"Expected kwarg_chunks[0]['x1'].shape == [128, {d_hid}], got {kwarg_chunks[0]['x1'].shape}"
+            )
+        if kwarg_chunks[1]["x2"].shape != torch.Size([256, d_hid]):
+            raise AssertionError(
+                f"Expected kwarg_chunks[1]['x2'].shape == [256, {d_hid}], got {kwarg_chunks[1]['x2'].shape}"
+            )
 
         # Merge chunks back together
         merged_args = merge_chunks(
@@ -133,7 +156,8 @@ def block_causal_mask(
             args_chunk_spec=None,
             kwargs_chunk_spec=None,
         )
-        assert len(arg_split) == 4
+        if len(arg_split) != 4:
+            raise AssertionError(f"Expected 4 arg_split, got {len(arg_split)}")
 
         q_total_chunks = []
         dq_total_chunks = []
@@ -237,7 +261,8 @@ def causal_mask(
             kwargs_chunk_spec=None,
         )
 
-        assert len(arg_split) == 4
+        if len(arg_split) != 4:
+            raise AssertionError(f"Expected 4 arg_split, got {len(arg_split)}")
 
         out_total_chunks = []
         for i in range(len(arg_split)):
@@ -265,7 +290,8 @@ def test_split_block_mask_none(self, device):
             kwargs_chunk_spec=None,
         )
 
-        assert len(arg_split) == 4
+        if len(arg_split) != 4:
+            raise AssertionError(f"Expected 4 arg_split, got {len(arg_split)}")
 
         for i in range(len(arg_split)):
             self.assertIsNone(arg_split[i][3])
diff --git a/test/distributed/pipelining/test_pipe.py b/test/distributed/pipelining/test_pipe.py
index 8ddb5634811cb..bb8c88f3ce22b 100644
--- a/test/distributed/pipelining/test_pipe.py
+++ b/test/distributed/pipelining/test_pipe.py
@@ -89,9 +89,10 @@ def test_model_split(self, ModelClass):
             mb_args=(x, y),
         )
 
-        assert pipe.num_stages == EXPECTED_N_STAGES[ModelClass], (
-            f"nstages = {pipe.num_stages}, expect {EXPECTED_N_STAGES[ModelClass]}"
-        )
+        if pipe.num_stages != EXPECTED_N_STAGES[ModelClass]:
+            raise AssertionError(
+                f"nstages = {pipe.num_stages}, expect {EXPECTED_N_STAGES[ModelClass]}"
+            )
 
         ref_out = mod(x, y)
         out = pipe(x, y)[0]
@@ -105,14 +106,18 @@ def test_model_split(self, ModelClass):
         for idx in range(pipe.num_stages):
             stage_mod = pipe.get_stage_module(idx)
             stage_fqns = set(stage_mod.state_dict().keys())
-            assert stage_fqns.issubset(old_names)
+            if not stage_fqns.issubset(old_names):
+                raise AssertionError(
+                    f"stage_fqns {stage_fqns} is not a subset of old_names {old_names}"
+                )
             new_names.update(stage_fqns)
 
         if CHECK_FQN_SET_EQUALITY:
-            assert old_names == new_names, f"""
+            if old_names != new_names:
+                raise AssertionError(f"""
             old names {old_names}
             new names {new_names}
-            """
+            """)
         print("Qualname check passed")
 
 
diff --git a/test/distributed/pipelining/test_schedule.py b/test/distributed/pipelining/test_schedule.py
index def652f854140..edef0c3ace7f1 100644
--- a/test/distributed/pipelining/test_schedule.py
+++ b/test/distributed/pipelining/test_schedule.py
@@ -4,6 +4,7 @@
 import csv
 import logging
 import os
+from unittest.mock import MagicMock, patch
 
 from model_registry import MultiMLP
 
@@ -24,6 +25,7 @@
     _add_reduce_grad,
     _add_send_recv,
     _add_unshard_reshard,
+    _batch_p2p,
     _format_pipeline_order,
     _merge_bw,
     _PipelineSchedule,
@@ -1207,5 +1209,72 @@ def test_generate_stage_to_rank_mapping(self):
 
 instantiate_parametrized_tests(TestScheduleLowering)
 
+
+class TestBatchP2P(TestCase):
+    """Tests that _batch_p2p dispatches homogeneous ops individually to avoid
+    head-of-line blocking, while still batching mixed ops for deadlock avoidance."""
+
+    def _make_p2p_op(self, op, group_peer=0):
+        p = MagicMock()
+        p.op = op
+        p.tensor = torch.zeros(1)
+        p.group = MagicMock()
+        p.tag = 0
+        p.group_peer = group_peer
+        return p
+
+    def test_empty_ops(self):
+        self.assertEqual(_batch_p2p([]), [])
+
+    @patch("torch.distributed.pipelining.schedules.dist.batch_isend_irecv")
+    @patch("torch.distributed.pipelining.schedules.dist.isend")
+    def test_all_isend_dispatched_individually(self, mock_isend, mock_batch):
+        mock_isend.return_value = MagicMock()
+        ops = [self._make_p2p_op(mock_isend, group_peer=i) for i in range(3)]
+
+        result = _batch_p2p(ops)
+
+        mock_batch.assert_not_called()
+        self.assertEqual(len(result), 3)
+        self.assertEqual(mock_isend.call_count, 3)
+        for p in ops:
+            mock_isend.assert_any_call(
+                p.tensor, group=p.group, tag=p.tag, group_dst=p.group_peer
+            )
+
+    @patch("torch.distributed.pipelining.schedules.dist.batch_isend_irecv")
+    @patch("torch.distributed.pipelining.schedules.dist.irecv")
+    def test_all_irecv_dispatched_individually(self, mock_irecv, mock_batch):
+        mock_irecv.return_value = MagicMock()
+        ops = [self._make_p2p_op(mock_irecv, group_peer=i) for i in range(3)]
+
+        result = _batch_p2p(ops)
+
+        mock_batch.assert_not_called()
+        self.assertEqual(len(result), 3)
+        self.assertEqual(mock_irecv.call_count, 3)
+        for p in ops:
+            mock_irecv.assert_any_call(
+                p.tensor, group=p.group, tag=p.tag, group_src=p.group_peer
+            )
+
+    @patch("torch.distributed.pipelining.schedules.dist.batch_isend_irecv")
+    @patch("torch.distributed.pipelining.schedules.dist.irecv")
+    @patch("torch.distributed.pipelining.schedules.dist.isend")
+    def test_mixed_ops_use_batch(self, mock_isend, mock_irecv, mock_batch):
+        mock_batch.return_value = [MagicMock(), MagicMock()]
+        ops = [
+            self._make_p2p_op(mock_isend, group_peer=0),
+            self._make_p2p_op(mock_irecv, group_peer=1),
+        ]
+
+        result = _batch_p2p(ops)
+
+        mock_batch.assert_called_once_with(ops)
+        mock_isend.assert_not_called()
+        mock_irecv.assert_not_called()
+        self.assertEqual(len(result), 2)
+
+
 if __name__ == "__main__":
     run_tests()
diff --git a/test/distributed/pipelining/test_schedule_multiproc.py b/test/distributed/pipelining/test_schedule_multiproc.py
index 687cb113b48bf..6d37e65214fef 100644
--- a/test/distributed/pipelining/test_schedule_multiproc.py
+++ b/test/distributed/pipelining/test_schedule_multiproc.py
@@ -546,8 +546,12 @@ def test_grad_with_manual(self, ScheduleClass, shape_inference):
         # Handle shape inference
         if not shape_inference:
             input_args = (x.chunk(chunks)[0],)
-            with torch.no_grad():
-                output_args = stage_module(*input_args)
+            if self.rank > 0:
+                # Non-first stages receive activations from previous stages,
+                # which have requires_grad=True in training mode.
+                input_args = tuple(a.detach().requires_grad_(True) for a in input_args)
+            output_args = stage_module(*input_args)
+            output_args = output_args.detach().requires_grad_(output_args.requires_grad)
             stage = PipelineStage(
                 stage_module,
                 self.rank,
@@ -823,14 +827,18 @@ def test_custom_function_callback(self):
         def forward_callback(action: _Action, ctx: _PipelineContext):
             """Custom callback for FORWARD computation that mimics the original implementation."""
             schedule = ctx.schedule_ref
-            assert isinstance(schedule, _PipelineScheduleRuntime)
+            if not isinstance(schedule, _PipelineScheduleRuntime):
+                raise AssertionError(
+                    f"Expected _PipelineScheduleRuntime, got {type(schedule)}"
+                )
             stage_index_to_stage: dict[int, _PipelineStageBase] = {
                 stage.stage_index: stage for stage in schedule._stages
             }
             stage = stage_index_to_stage[action.stage_index]
             stage_index = stage.stage_index
             mb_index = action.microbatch_index
-            assert mb_index is not None
+            if mb_index is None:
+                raise AssertionError("Expected mb_index to not be None")
             fwd_recv_ops = schedule.fwd_recv_ops
             arg_mbs = ctx.arg_mbs
             kwarg_mbs = ctx.kwarg_mbs
@@ -846,10 +854,8 @@ def forward_callback(action: _Action, ctx: _PipelineContext):
                 # no recv op expected for V-schedule special case (see [Note: V-schedule special case])
                 and not is_prev_stage_on_this_rank
             ):
-                assert (
-                    stage_index,
-                    mb_index,
-                ) in fwd_recv_ops, f"Computing {action=} before receiving input"
+                if (stage_index, mb_index) not in fwd_recv_ops:
+                    raise AssertionError(f"Computing {action=} before receiving input")
                 from torch.distributed.pipelining.schedules import _wait_batch_p2p
 
                 _wait_batch_p2p(fwd_recv_ops.pop((stage_index, mb_index)))
@@ -871,11 +877,15 @@ def forward_callback(action: _Action, ctx: _PipelineContext):
         def overlap_callback(action: _Action, ctx: _PipelineContext):
             """Custom callback for OVERLAP_F_B computation that mimics the original implementation."""
             schedule = ctx.schedule_ref
-            assert isinstance(schedule, _PipelineScheduleRuntime)
+            if not isinstance(schedule, _PipelineScheduleRuntime):
+                raise AssertionError(
+                    f"Expected _PipelineScheduleRuntime, got {type(schedule)}"
+                )
             stage_index_to_stage: dict[int, _PipelineStageBase] = {
                 stage.stage_index: stage for stage in schedule._stages
             }
-            assert action.sub_actions is not None
+            if action.sub_actions is None:
+                raise AssertionError("Expected action.sub_actions to not be None")
             fwd_action = action.sub_actions[0]
             bwd_action = action.sub_actions[1]
 
@@ -894,7 +904,8 @@ def overlap_callback(action: _Action, ctx: _PipelineContext):
             backward_stage_index = bwd_action.stage_index
             backward_stage = stage_index_to_stage[backward_stage_index]
             backward_mb_index = bwd_action.microbatch_index
-            assert backward_mb_index is not None
+            if backward_mb_index is None:
+                raise AssertionError("Expected backward_mb_index to not be None")
             bwd_recv_ops = schedule.bwd_recv_ops
             is_next_stage_on_this_rank = (
                 backward_stage.stage_index + 1 in stage_index_to_stage
@@ -907,12 +918,10 @@ def overlap_callback(action: _Action, ctx: _PipelineContext):
                 # no recv op expected for V-schedule special case (see [Note: V-schedule special case])
                 and not is_next_stage_on_this_rank
             ):
-                assert (
-                    backward_stage_index,
-                    backward_mb_index,
-                ) in bwd_recv_ops, (
-                    f"Attempted to run compute {action=} before receiving input"
-                )
+                if (backward_stage_index, backward_mb_index) not in bwd_recv_ops:
+                    raise AssertionError(
+                        f"Attempted to run compute {action=} before receiving input"
+                    )
                 _wait_batch_p2p(
                     bwd_recv_ops.pop((backward_stage_index, backward_mb_index))
                 )
@@ -1161,7 +1170,10 @@ def test_pipeline_schedule_runtime_custom_sched(self, ScheduleClass):
         schedule = ScheduleClass(
             stages, num_microbatches, loss_fn=loss_fn, scale_grads=False
         )
-        assert isinstance(schedule, _PipelineScheduleRuntime)
+        if not isinstance(schedule, _PipelineScheduleRuntime):
+            raise AssertionError(
+                f"Expected _PipelineScheduleRuntime, got {type(schedule)}"
+            )
 
         # Run pipeline with tensor leak checking
         with check_leaked_tensors() as garbage_tensors:
diff --git a/test/distributed/pipelining/test_stage.py b/test/distributed/pipelining/test_stage.py
index 17173c66ff69d..b6481fce17ba8 100644
--- a/test/distributed/pipelining/test_stage.py
+++ b/test/distributed/pipelining/test_stage.py
@@ -11,7 +11,7 @@
     PipelineStage,
     ScheduleGPipe,
 )
-from torch.distributed.pipelining._utils import PipeliningShapeError
+from torch.distributed.pipelining._utils import PipeliningMetadataError
 from torch.testing._internal.common_distributed import (
     MultiProcContinuousTest,
     requires_accelerator_dist_backend,
@@ -118,7 +118,10 @@ def _run_step(x):
         submod_keys = stage.submod.state_dict().keys()
         # Confirm keys are consistent with original model
         old_keys = mod.state_dict().keys()
-        assert all(k in old_keys for k in submod_keys)
+        if not all(k in old_keys for k in submod_keys):
+            raise AssertionError(
+                f"Some keys not found in old_keys: {[k for k in submod_keys if k not in old_keys]}"
+            )
 
     @requires_accelerator_dist_backend(["nccl", "xccl"])
     @skip_but_pass_in_sandcastle_if(
@@ -169,7 +172,10 @@ def test_tracer_kwargs(self, ModelClass):
         submod_keys = stage.submod.state_dict().keys()
         # Confirm keys are consistent with original model
         old_keys = mod.state_dict().keys()
-        assert all(k in old_keys for k in submod_keys)
+        if not all(k in old_keys for k in submod_keys):
+            raise AssertionError(
+                f"Some keys not found in old_keys: {[k for k in submod_keys if k not in old_keys]}"
+            )
 
     @requires_accelerator_dist_backend(["nccl", "xccl"])
     @skip_but_pass_in_sandcastle_if(
@@ -362,6 +368,7 @@ def test_shape_prop_mismatch(self):
             self.world_size,
             self.device,
         )
+        stage._runtime_validate = True
 
         # Attach to a schedule
         schedule = ScheduleGPipe(stage, chunks)
@@ -376,20 +383,20 @@ def _run_step(x):
         _run_step(x)
 
         if self.rank == 0:
-            with self.assertRaisesRegex(PipeliningShapeError, "shape mismatch"):
+            with self.assertRaisesRegex(PipeliningMetadataError, "shape mismatch"):
                 _run_step(torch.randn(batch_size + 1, d_hid, device=self.device))
 
-            with self.assertRaisesRegex(PipeliningShapeError, "dtype mismatch"):
+            with self.assertRaisesRegex(PipeliningMetadataError, "dtype mismatch"):
                 _run_step(x.to(torch.int32))
 
             # output of stage's mlp layer will be flattened by this hook, the stage should err
             handle = stage_mod.register_forward_hook(get_flatten_hook())
-            with self.assertRaisesRegex(PipeliningShapeError, "shape mismatch"):
+            with self.assertRaisesRegex(PipeliningMetadataError, "shape mismatch"):
                 _run_step(x)
             handle.remove()
 
             stage_mod.register_forward_hook(get_dtype_change_hook(torch.bfloat16))
-            with self.assertRaisesRegex(PipeliningShapeError, "dtype mismatch"):
+            with self.assertRaisesRegex(PipeliningMetadataError, "dtype mismatch"):
                 _run_step(x)
 
     @requires_accelerator_dist_backend(["nccl", "xccl"])
diff --git a/test/distributed/pipelining/test_transformer.py b/test/distributed/pipelining/test_transformer.py
index 20e830547de7b..66de58167f94d 100644
--- a/test/distributed/pipelining/test_transformer.py
+++ b/test/distributed/pipelining/test_transformer.py
@@ -48,7 +48,8 @@ def test_ir(self, device):
             (x,),
             split_spec=split_spec,
         )
-        assert pipe.num_stages == num_stages, f"{pipe.num_stages=}, expect {num_stages}"
+        if pipe.num_stages != num_stages:
+            raise AssertionError(f"{pipe.num_stages=}, expect {num_stages}")
 
         def get_layers(module):
             layers = [name for name, _ in module.layers.named_children()]
@@ -62,7 +63,8 @@ def get_layers(module):
 
         # Check layer completeness
         orig_layers = get_layers(transformer)
-        assert sorted(layers) == sorted(orig_layers), f"{layers} != {orig_layers}"
+        if sorted(layers) != sorted(orig_layers):
+            raise AssertionError(f"{layers} != {orig_layers}")
         print("Layers matched!")
 
         # Check equivalence
diff --git a/test/distributed/pipelining/test_unflatten.py b/test/distributed/pipelining/test_unflatten.py
index 0493f39b16cb8..37ab701fe2d95 100644
--- a/test/distributed/pipelining/test_unflatten.py
+++ b/test/distributed/pipelining/test_unflatten.py
@@ -53,16 +53,18 @@ def test_unflatten(self, device):
             {"constant": constant},
         )
 
-        assert pipe.num_stages == 4
+        if pipe.num_stages != 4:
+            raise AssertionError(
+                f"Expected pipe.num_stages == 4, got {pipe.num_stages}"
+            )
         orig_state_dict = mod.state_dict()
 
         # Check qualnames
         for stage_idx in range(pipe.num_stages):
             stage_mod = pipe.get_stage_module(stage_idx)
             for param_name, _ in stage_mod.named_parameters():
-                assert param_name in orig_state_dict, (
-                    f"{param_name} not in original state dict"
-                )
+                if param_name not in orig_state_dict:
+                    raise AssertionError(f"{param_name} not in original state dict")
         print("Param qualname test passed")
 
         # Check equivalence
diff --git a/test/distributed/tensor/debug/test_comm_mode.py b/test/distributed/tensor/debug/test_comm_mode.py
index d122a9f716fcd..a8f22333a95d5 100644
--- a/test/distributed/tensor/debug/test_comm_mode.py
+++ b/test/distributed/tensor/debug/test_comm_mode.py
@@ -5,6 +5,7 @@
 import torch.distributed._functional_collectives as funcol
 import torch.nn as nn
 from torch.distributed.tensor import DeviceMesh, DTensor, Shard
+from torch.distributed.tensor._redistribute import use_min_cost_redistribution_plan
 from torch.distributed.tensor.debug import CommDebugMode
 from torch.testing._internal.common_distributed import requires_accelerator_dist_backend
 from torch.testing._internal.common_utils import run_tests, TestCase
@@ -105,7 +106,7 @@ def f(x, y):
         x_dtensor = DTensor.from_local(x, mesh, [Shard(0)], run_check=False)
         y_dtensor = DTensor.from_local(y, mesh, [Shard(0)], run_check=False)
 
-        with comm_mode:
+        with comm_mode, use_min_cost_redistribution_plan():
             f(x_dtensor, y_dtensor)
 
         comm_counts = comm_mode.get_comm_counts()
diff --git a/test/distributed/tensor/debug/test_debug_mode.py b/test/distributed/tensor/debug/test_debug_mode.py
index 16521517fbe9c..337fd2bd76cf5 100644
--- a/test/distributed/tensor/debug/test_debug_mode.py
+++ b/test/distributed/tensor/debug/test_debug_mode.py
@@ -666,6 +666,7 @@ def forward(self, x):
         with DebugMode(record_nn_module=True) as debug_mode:
             fn(inp)
 
+    @torch._functorch.config.patch(guess_tangent_strides_as_outputs=True)
     def test_nn_module_in_compiled_regions(self):
         class Foo(torch.nn.Module):
             def __init__(self):
@@ -753,7 +754,8 @@ def forward(self, x):
     aten::sum(t: f32[4, 4])  ->  t: f32[]
     aten::ones_like(t: f32[], pin_memory=False, memory_format=torch.preserve_format)  ->  t: f32[]
     aten::expand(t: f32[], [4, 4])  ->  t: f32[4, 4]
-    aten::clone(t: f32[4, 4], memory_format=torch.contiguous_format)  ->  t: f32[4, 4]
+    aten::empty_strided([4, 4], [4, 1], dtype=torch.float32, layout=torch.strided, device=cpu, pin_memory=False)  ->  t: f32[4, 4]
+    aten::copy_(t: f32[4, 4], t: f32[4, 4])  ->  t: f32[4, 4]
   [aot_eager region (compile)] enter
     [nn.Mod (compile)] L['self'].bar
       [nn.Mod (compile)] L['self'].bar.l3
@@ -1010,6 +1012,7 @@ def f(dA, dB):
         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)
 
+    @torch._functorch.config.patch(guess_tangent_strides_as_outputs=True)
     def test_invoke_subgraph(self):
         # Test that DebugMode can trace the operations inside
         # invoke_subgraph HOP
@@ -1054,7 +1057,8 @@ def fn(x):
     aten::sum(t: f32[8, 8])  ->  t: f32[]
     aten::ones_like(t: f32[], pin_memory=False, memory_format=torch.preserve_format)  ->  t: f32[]
     aten::expand(t: f32[], [8, 8])  ->  t: f32[8, 8]
-    aten::clone(t: f32[8, 8], memory_format=torch.contiguous_format)  ->  t: f32[8, 8]
+    aten::empty_strided([8, 8], [8, 1], dtype=torch.float32, layout=torch.strided, device=cpu, pin_memory=False)  ->  t: f32[8, 8]
+    aten::copy_(t: f32[8, 8], t: f32[8, 8])  ->  t: f32[8, 8]
     torch.ops.higher_order.invoke_subgraph(partitioned_bw_subgraph_0_0, t: f32[8, 8], t: f32[8, 8])  ->  ('t: f32[8, 8]',)
     [annotate] [enter InvokeSubgraph HOP] partitioned_bw_subgraph_0_0
       aten::cos(t: f32[8, 8])  ->  t: f32[8, 8]
diff --git a/test/distributed/tensor/debug/test_op_coverage.py b/test/distributed/tensor/debug/test_op_coverage.py
index 0392bae03c504..2b19415aa89f9 100644
--- a/test/distributed/tensor/debug/test_op_coverage.py
+++ b/test/distributed/tensor/debug/test_op_coverage.py
@@ -23,7 +23,8 @@ def test_trace_with_inductor_decomp(self):
         args = (torch.randn(8, 50),)
         kwargs = {}
         graphs = get_inductor_decomp_graphs(model, args, kwargs)
-        assert len(graphs) == 2, "Expect fwd + bwd graphs"
+        if len(graphs) != 2:
+            raise AssertionError(f"Expect fwd + bwd graphs, got {len(graphs)}")
         self.assertIsInstance(graphs[0], torch.fx.GraphModule)
         self.assertIsInstance(graphs[1], torch.fx.GraphModule)
 
diff --git a/test/distributed/tensor/experimental/test_register_sharding.py b/test/distributed/tensor/experimental/test_register_sharding.py
index 0f2f208689608..fbba4839fc55d 100644
--- a/test/distributed/tensor/experimental/test_register_sharding.py
+++ b/test/distributed/tensor/experimental/test_register_sharding.py
@@ -17,6 +17,19 @@
 
 
 class TestRegisterSharding(DTensorTestBase):
+    def tearDown(self):
+        super().tearDown()
+        # Clean up any custom ops registered during tests to avoid test pollution
+        # This addresses the concern about proper deregistration of custom ops
+        import torch._custom_op
+
+        keys = list(torch._custom_op.impl.global_registry.keys())
+        for key in keys:
+            if key.startswith("test_dtensor::"):
+                torch._custom_op.impl.global_registry[key]._destroy()
+        if hasattr(torch.ops, "test_dtensor"):
+            delattr(torch.ops, "test_dtensor")
+
     @with_comms
     def test_softmax_fwd(self):
         # After registering the custom softmax sharding strategy,
@@ -144,6 +157,32 @@ def min_dim_strategy(x, dim, keepdim, min, min_indices):
         self.assertEqual(result[0].full_tensor(), expected_values)
         self.assertEqual(result[1].full_tensor(), expected_indices)
 
+    @with_comms
+    def test_register_sharding_non_tensor_first_arg(self):
+        # Test that get_mesh_from_args works when the first arg is not a DTensor
+        # Regression test for https://github.com/pytorch/pytorch/issues/167915
+        @torch.library.custom_op("test_dtensor::scalar_first_arg", mutates_args=())
+        def scalar_first_arg(name: str, x: torch.Tensor) -> torch.Tensor:
+            return x.clone()
+
+        @scalar_first_arg.register_fake
+        def _(name, x):
+            return x.clone()
+
+        @register_sharding(torch.ops.test_dtensor.scalar_first_arg.default)
+        def scalar_first_arg_sharding(name, x):
+            return [([Replicate()], [None, Replicate()])]
+
+        mesh = self.build_device_mesh()
+        x = torch.randn(4, 4, device=self.device_type)
+        x_dt = distribute_tensor(x, mesh, [Replicate()])
+
+        # This should not raise "Cannot find device mesh from args"
+        result = torch.ops.test_dtensor.scalar_first_arg("test", x_dt)
+
+        self.assertIsInstance(result, DTensor)
+        self.assertEqual(result.full_tensor(), x)
+
 
 if __name__ == "__main__":
     run_tests()
diff --git a/test/distributed/tensor/parallel/test_micro_pipeline_tp.py b/test/distributed/tensor/parallel/test_micro_pipeline_tp.py
index 4aaf938293854..1a9baebe7f01f 100644
--- a/test/distributed/tensor/parallel/test_micro_pipeline_tp.py
+++ b/test/distributed/tensor/parallel/test_micro_pipeline_tp.py
@@ -1,6 +1,5 @@
 # Owner(s): ["module: c10d"]
 import unittest
-from typing import Optional
 
 import torch
 import torch.distributed as dist
@@ -268,7 +267,7 @@ def func(
             B: torch.Tensor,
             A_scale: torch.Tensor,
             B_scale: torch.Tensor,
-            out_dtype: Optional[torch.dtype],
+            out_dtype: torch.dtype | None,
         ) -> torch.Tensor:
             A = _fp8_all_gather(
                 A_shard, gather_dim=gather_dim, group_name=group.group_name
@@ -524,6 +523,7 @@ def tearDown(self):
 
     @unittest.skipIf(not HAS_GPU, "Inductor+gpu needs triton and recent GPU arch")
     @fresh_cache()
+    @torch._inductor.config.patch(shape_padding=False)
     def test_extra_collectives(self):
         device_mesh = DeviceMesh(
             "cuda",
diff --git a/test/distributed/tensor/parallel/test_tp_examples.py b/test/distributed/tensor/parallel/test_tp_examples.py
index 88ca55efde5d2..0829baf2f537c 100644
--- a/test/distributed/tensor/parallel/test_tp_examples.py
+++ b/test/distributed/tensor/parallel/test_tp_examples.py
@@ -3,7 +3,7 @@
 
 import itertools
 from copy import deepcopy
-from typing import NamedTuple, Optional
+from typing import NamedTuple
 
 import torch
 import torch.distributed as dist
@@ -55,9 +55,9 @@
 
 
 class ExpCommCounts(NamedTuple):
-    fwd: Optional[dict] = None
-    bwd: Optional[dict] = None
-    optim: Optional[dict] = None
+    fwd: dict | None = None
+    bwd: dict | None = None
+    optim: dict | None = None
 
 
 class DistTensorParallelExampleTest(DTensorTestBase):
@@ -311,7 +311,10 @@ def test_transformer_training(self, is_seq_parallel, dtype: torch.dtype):
         # Initialize input and make sure all ranks have the same input.
         inp_size = [8, 8]  # [batch_size, seq_len]
         if is_seq_parallel:
-            assert inp_size[1] % self.world_size == 0
+            if inp_size[1] % self.world_size != 0:
+                raise AssertionError(
+                    f"Expected inp_size[1] % world_size == 0, got {inp_size[1]} % {self.world_size}"
+                )
 
         torch.manual_seed(0)
         steps = 10 if type(model) is torch.float64 else 1
@@ -439,7 +442,10 @@ def test_transformer_req_grad(self, thaw_params, is_seq_parallel, dtype, exp_cnt
         # Initialize input and make sure all ranks have the same input.
         inp_size = [8, 8]  # [batch_size, seq_len]
         if is_seq_parallel:
-            assert inp_size[1] % self.world_size == 0
+            if inp_size[1] % self.world_size != 0:
+                raise AssertionError(
+                    f"Expected inp_size[1] % world_size == 0, got {inp_size[1]} % {self.world_size}"
+                )
 
         torch.manual_seed(0)
         inp = torch.randint(model_args.vocab_size, inp_size, device=self.device_type)
diff --git a/test/distributed/tensor/parallel/test_tp_random_state.py b/test/distributed/tensor/parallel/test_tp_random_state.py
index 9c8ae8bdac1b5..d1f2153181c10 100644
--- a/test/distributed/tensor/parallel/test_tp_random_state.py
+++ b/test/distributed/tensor/parallel/test_tp_random_state.py
@@ -20,7 +20,8 @@
 class TensorParallelRandomStateTests(DTensorTestBase):
     def get_tensor_slice(self, idx, n, large_tensor):
         shape = large_tensor.shape
-        assert shape[0] % n == 0
+        if shape[0] % n != 0:
+            raise AssertionError(f"Expected shape[0] % n == 0, got {shape[0]} % {n}")
         local_shape = [shape[0] // n, shape[1]]
         slice_idx = (
             slice(idx * local_shape[0], (idx + 1) * local_shape[0]),
@@ -83,7 +84,10 @@ def test_model_init(self):
                 # check within the TP group
                 # the 1d mesh represents the TP group
                 _1d_mesh = dtensor.device_mesh
-                assert _1d_mesh.ndim == 1
+                if _1d_mesh.ndim != 1:
+                    raise AssertionError(
+                        f"Expected _1d_mesh.ndim == 1, got {_1d_mesh.ndim}"
+                    )
                 self.assertEqual(_1d_mesh, tp_mesh)
 
                 tensor_local = dtensor.to_local()
diff --git a/test/distributed/tensor/test_api.py b/test/distributed/tensor/test_api.py
index 98caa3b4039a9..61d700a6ab9de 100644
--- a/test/distributed/tensor/test_api.py
+++ b/test/distributed/tensor/test_api.py
@@ -254,7 +254,8 @@ def input_fn(mod, inputs, device_mesh):
             return DTensor.from_local(inputs[0], device_mesh, [Shard(0)])
 
         def output_fn(mod, outputs, device_mesh):
-            assert isinstance(outputs, DTensor)
+            if not isinstance(outputs, DTensor):
+                raise AssertionError(f"Expected DTensor, got {type(outputs)}")
             return outputs.to_local()
 
         replica_module = distribute_module(
@@ -330,7 +331,8 @@ def input_fn(inputs, device_mesh):
             return DTensor.from_local(inputs[0], device_mesh, [Shard(0)])
 
         def output_fn(outputs, device_mesh):
-            assert isinstance(outputs, DTensor)
+            if not isinstance(outputs, DTensor):
+                raise AssertionError(f"Expected DTensor, got {type(outputs)}")
             return outputs.to_local()
 
         with self.assertWarnsRegex(FutureWarning, "Deprecating"):
diff --git a/test/distributed/tensor/test_attention.py b/test/distributed/tensor/test_attention.py
index bc232ecf5731e..8417af817598a 100644
--- a/test/distributed/tensor/test_attention.py
+++ b/test/distributed/tensor/test_attention.py
@@ -1,11 +1,12 @@
 # Copyright (c) Meta Platforms, Inc. and affiliates
 # Owner(s): ["oncall: distributed"]
+import contextlib
 import itertools
 import random
 import unittest
 import unittest.mock
 from collections.abc import Callable
-from typing import Any, ClassVar, Optional
+from typing import Any, ClassVar
 
 import torch
 import torch.distributed as dist
@@ -190,7 +191,9 @@ def _ring_attention_sdpa(
         for target in [cp_q, cp_k, cp_v]:
             target.requires_grad = True
 
-        with CommDebugMode() as comm_mode:
+        check_comm_counts = not compiled and rotater == _RotateMethod.ALL_TO_ALL
+        comm_mode = CommDebugMode() if check_comm_counts else contextlib.nullcontext()
+        with comm_mode:
             with sdpa_kernel(backend):
                 cp_out = fn_eval(
                     attention,
@@ -200,8 +203,7 @@ def _ring_attention_sdpa(
                     is_causal=is_causal,
                 )
 
-            if not compiled and rotater == _RotateMethod.ALL_TO_ALL:
-                # Compiler and CommDebugMode do not work well together.
+            if check_comm_counts:
                 expect_all2all_count = (
                     self.world_size - 1
                     if test_forward_only
@@ -244,6 +246,14 @@ def fn_eval(fn, *args, **kwargs):
         if load_balance and not is_causal:
             return
 
+        # Compilation with context_parallel doesn't work yet — both paths
+        # (use_context=True monkey-patch and use_context=False parallelize_module)
+        # fail during tracing because DTensor dispatch interferes with sdpa.
+        # Previously CommDebugMode was active for all subtests, which caused
+        # the frame to be silently skipped, masking this limitation.
+        if compiled:
+            return
+
         set_rotate_method(rotater_enum_to_str[rotater])
         self.assertEqual(_cp_options.rotate_method, rotater)
         device_mesh = DeviceMesh(self.device_type, torch.arange(0, self.world_size))
@@ -388,7 +398,11 @@ def generate_random_lengths_in_chunks(
     # must be a multiple of `chunk_size`. Besides, the lengths of all the documents
     # sum up to `total_length`.
     num_chunks = total_length // chunk_size
-    assert total_length % chunk_size == 0 and num_chunks >= num_documents
+    if not (total_length % chunk_size == 0 and num_chunks >= num_documents):
+        raise AssertionError(
+            f"total_length % chunk_size == {total_length % chunk_size} (expected 0), "
+            f"num_chunks={num_chunks} vs num_documents={num_documents}"
+        )
 
     num_chunks_per_document = [1] * num_documents
     remaining_chunks = num_chunks - num_documents
@@ -484,7 +498,7 @@ def _test_cp_flex_attention(
         B: int = 1,
         block_mask,
         lb_type: str,
-        document_lengths: Optional[list[list[int]]] = None,
+        document_lengths: list[list[int]] | None = None,
     ) -> None:
         torch.use_deterministic_algorithms(True)
         torch.cuda.manual_seed(1234)
@@ -580,7 +594,7 @@ def _test_cp_flex_attention(
 
     def _get_load_balancer(
         self, lb_type: str, kwargs: dict[str, Any]
-    ) -> Optional[_LoadBalancer]:
+    ) -> _LoadBalancer | None:
         seq_length = kwargs["seq_length"]
         document_lengths = kwargs["document_lengths"]
         block_mask = kwargs["block_mask"]
@@ -589,17 +603,20 @@ def _get_load_balancer(
         if lb_type == "None":
             load_balancer = None  # no load-balance
         elif lb_type == "_HeadTailLoadBalancer":
-            assert isinstance(seq_length, int)
+            if not isinstance(seq_length, int):
+                raise AssertionError(f"Expected int, got {type(seq_length)}")
             load_balancer = _HeadTailLoadBalancer(
                 seq_length, self.world_size, torch.device(self.device_type)
             )
         elif lb_type == "_PerDocumentHeadTailLoadBalancer":
-            assert isinstance(document_lengths, list)
+            if not isinstance(document_lengths, list):
+                raise AssertionError(f"Expected list, got {type(document_lengths)}")
             load_balancer = _PerDocumentHeadTailLoadBalancer(
                 document_lengths, self.world_size, torch.device(self.device_type)
             )
         elif lb_type == "_PTRRLoadBalancer":
-            assert isinstance(block_mask, BlockMask)
+            if not isinstance(block_mask, BlockMask):
+                raise AssertionError(f"Expected BlockMask, got {type(block_mask)}")
             load_balancer = _PTRRLoadBalancer(
                 block_mask,
                 self.world_size,
diff --git a/test/distributed/tensor/test_common_rules.py b/test/distributed/tensor/test_common_rules.py
index 3450f8faa2b5c..900f285d6bcd9 100644
--- a/test/distributed/tensor/test_common_rules.py
+++ b/test/distributed/tensor/test_common_rules.py
@@ -30,7 +30,7 @@ def _gen_tensor_meta(self, shape):
 
     def test_einop_basic_propagation(self):
         # plain einsum, mm
-        mesh = DeviceMesh(self.device_type(), torch.arange(self.world_size))
+        mesh = DeviceMesh(self.device_type, torch.arange(self.world_size))
 
         mm_call = aten.mm.default
         # propagate col-wise sharding
@@ -82,7 +82,7 @@ def test_einop_basic_propagation(self):
         self.assertTrue(output_spec.placements[0].is_partial())
 
     def test_einop_pointwise_propagation(self):
-        mesh = DeviceMesh(self.device_type(), torch.arange(self.world_size))
+        mesh = DeviceMesh(self.device_type, torch.arange(self.world_size))
 
         add_call = aten.add.Tensor
         # addition
@@ -137,7 +137,7 @@ def test_einop_merge_sharding(self):
         mesh_shape = torch.arange(self.world_size).reshape(
             self.world_size // 2, self.world_size // 2
         )
-        mesh = DeviceMesh(self.device_type(), mesh_shape)
+        mesh = DeviceMesh(self.device_type, mesh_shape)
 
         mm_call = aten.mm.default
 
@@ -161,7 +161,7 @@ def test_einop_linearity(self):
         mesh_shape = torch.arange(self.world_size).reshape(
             self.world_size // 2, self.world_size // 2
         )
-        mesh = DeviceMesh(self.device_type(), mesh_shape)
+        mesh = DeviceMesh(self.device_type, mesh_shape)
 
         mm_call = aten.mm.default
 
@@ -227,7 +227,7 @@ def test_einop_linearity(self):
     def test_einop_multi_sharding_on_mesh_dim(self):
         # einop prop with multi sharding on same mesh dim
         mesh_shape = torch.arange(self.world_size)
-        mesh = DeviceMesh(self.device_type(), mesh_shape)
+        mesh = DeviceMesh(self.device_type, mesh_shape)
 
         mm_call = aten.mm.default
         mat1, mat2 = [0, -1], [0, -1]
@@ -256,7 +256,7 @@ def test_einop_errors(self):
         mesh_shape = torch.arange(self.world_size).reshape(
             self.world_size // 2, self.world_size // 2
         )
-        mesh = DeviceMesh(self.device_type(), mesh_shape)
+        mesh = DeviceMesh(self.device_type, mesh_shape)
 
         add_call = aten.add.Tensor
         mat1, mat2 = [0, -1], [1, -1]
@@ -273,7 +273,7 @@ def test_einop_errors(self):
             einop_rule("ij,ij->ij", OpSchema(add_call, (mat1_spec, mat2_spec), {}))
 
     def test_pointwise_rules_broadcasting(self):
-        mesh = DeviceMesh(self.device_type(), torch.arange(self.world_size))
+        mesh = DeviceMesh(self.device_type, torch.arange(self.world_size))
 
         where_call = aten.where.self
         inp1, inp2, inp3 = [0], [], [-1, -1]
@@ -298,7 +298,7 @@ def test_pointwise_rules_broadcasting(self):
         self.assertEqual(output_spec.dim_map, [-1, 0])
 
     def test_pointwise_rules_suggestion(self):
-        mesh = DeviceMesh(self.device_type(), torch.arange(self.world_size))
+        mesh = DeviceMesh(self.device_type, torch.arange(self.world_size))
 
         lerp_call = aten.lerp.Scalar
         # propagate point-wise sharding
@@ -329,7 +329,7 @@ def test_pointwise_multi_sharding_on_mesh_dim(self):
         mesh_shape = torch.arange(self.world_size).reshape(
             self.world_size // 2, self.world_size // 2
         )
-        mesh = DeviceMesh(self.device_type(), mesh_shape)
+        mesh = DeviceMesh(self.device_type, mesh_shape)
 
         add_call = aten.add.Tensor
 
@@ -374,7 +374,7 @@ def test_pointwise_enforce_sharding_multi_sharding_on_mesh_dim(self):
         mesh_shape = torch.arange(self.world_size).reshape(
             self.world_size // 2, self.world_size // 2
         )
-        mesh = DeviceMesh(self.device_type(), mesh_shape)
+        mesh = DeviceMesh(self.device_type, mesh_shape)
 
         add_call = aten.add_.Tensor
 
diff --git a/test/distributed/tensor/test_compile_on_one_rank.py b/test/distributed/tensor/test_compile_on_one_rank.py
index e6c168ed3c376..166de302ee73d 100644
--- a/test/distributed/tensor/test_compile_on_one_rank.py
+++ b/test/distributed/tensor/test_compile_on_one_rank.py
@@ -32,87 +32,10 @@ def extract_graph(fx_g, _, graph_cell):
 
 
 class TestCompileOnOneRank(DTensorTestBase):
-    """
-    Test that torch.compile produces consistent graphs across ranks when used
-    with rowwise sharded embeddings (RowwiseParallel).
-
-    The bug being tested: When rowwise sharded embeddings are compiled with
-    torch.compile, the _MaskPartial._mask_tensor() function generates bounds
-    checking operations (lt, ge, sub, index_put) with rank-specific values that
-    get baked into the compiled graph:
-    - Rank 0: lt(index, 0), ge(index, 64), sub(index, 0)
-    - Rank 1: lt(index, 64), ge(index, 128), sub(index, 64)
-
-    These values should be symbolic/dynamic, not baked-in literals, to ensure
-    graph consistency across ranks.
-    """
-
-    @with_comms
-    @dist_config.patch(compile_on_one_rank=True)
-    def test_compiled_rowwise_embedding_graph_consistency(self):
-        """Test that compiled graphs are identical across all ranks."""
-        mesh = self.build_device_mesh()
-
-        class Network(nn.Module):
-            def __init__(self, num_embeddings, embedding_dim, device):
-                super().__init__()
-                self.tok_embeddings = nn.Embedding(
-                    num_embeddings, embedding_dim, device=device
-                )
-
-            def forward(self, x):
-                return self.tok_embeddings(x)
-
-        # Create embedding with enough rows to show sharding differences
-        # With 256 embeddings and 4 ranks, each rank gets 64 rows
-        torch.manual_seed(0)
-        num_embeddings = 256
-        embedding_dim = 64
-
-        model = Network(num_embeddings, embedding_dim, device=self.device_type)
-
-        # Apply RowwiseParallel like torchtitan does for tok_embeddings
-        parallelize_module(
-            model,
-            mesh,
-            {
-                "tok_embeddings": RowwiseParallel(
-                    input_layouts=Replicate(),
-                    output_layouts=Shard(1),
-                ),
-            },
-        )
-
-        # Create a custom backend that captures the forward graph
-        fw_graph_cell = [None]
-        fw_compiler = functools.partial(extract_graph, graph_cell=fw_graph_cell)
-
-        from functorch.compile import min_cut_rematerialization_partition
-        from torch._dynamo.backends.common import aot_autograd
-
-        aot_eager_graph = aot_autograd(
-            fw_compiler=fw_compiler,
-            bw_compiler=fw_compiler,  # Use same compiler for backward
-            partition_fn=min_cut_rematerialization_partition,
-        )
-
-        # Compile the sharded network with graph-capturing backend
-        compiled_model = torch.compile(model, backend=aot_eager_graph)
-
-        # Create replicated input
-        torch.manual_seed(42)
-        inp = torch.randint(0, num_embeddings, (64, 16), device=self.device_type)
-        replicated_inp = DTensor.from_local(inp, mesh, [Replicate()], run_check=False)
-
-        # Run compiled distributed embedding to trigger compilation
-        compiled_model(replicated_inp)
-
-        # Get the captured graph code from this rank
-        local_graph_code = fw_graph_cell[0]
+    def _assert_graphs_identical_across_ranks(self, local_graph_code):
+        """Gather compiled graph code from all ranks and assert they are identical."""
         self.assertIsNotNone(local_graph_code, "Graph was not captured")
 
-        # Gather graph code from all ranks
-        # Convert to bytes for all_gather
         graph_bytes = local_graph_code.encode("utf-8")
         graph_tensor = torch.tensor(
             list(graph_bytes), dtype=torch.uint8, device=self.device_type
@@ -129,29 +52,24 @@ def forward(self, x):
         dist.all_gather(all_lens, local_len)
         max_len = int(max(l.item() for l in all_lens))
 
-        # Pad local tensor to max length
         padded_tensor = torch.zeros(max_len, dtype=torch.uint8, device=self.device_type)
         padded_tensor[: len(graph_bytes)] = graph_tensor
 
-        # Gather all graphs
         all_graphs = [
             torch.zeros(max_len, dtype=torch.uint8, device=self.device_type)
             for _ in range(self.world_size)
         ]
         dist.all_gather(all_graphs, padded_tensor)
 
-        # Convert back to strings and compare
         graph_codes = []
-        for i, (graph_t, len_t) in enumerate(zip(all_graphs, all_lens)):
+        for graph_t, len_t in zip(all_graphs, all_lens):
             length = int(len_t.item())
             graph_str = bytes(graph_t[:length].tolist()).decode("utf-8")
             graph_codes.append(graph_str)
 
-        # Verify all graphs are identical
         rank0_graph = graph_codes[0]
         for rank, graph_code in enumerate(graph_codes[1:], start=1):
             if rank0_graph != graph_code:
-                # Generate unified diff for clearer error message
                 diff = difflib.unified_diff(
                     rank0_graph.splitlines(keepends=True),
                     graph_code.splitlines(keepends=True),
@@ -165,6 +83,102 @@ def forward(self, x):
                     f"Unified diff:\n{diff_str}"
                 )
 
+    def _compile_and_capture_graph(self, model):
+        """Compile model with a graph-capturing backend and return the graph cell."""
+        fw_graph_cell = [None]
+        fw_compiler = functools.partial(extract_graph, graph_cell=fw_graph_cell)
+
+        from functorch.compile import min_cut_rematerialization_partition
+        from torch._dynamo.backends.common import aot_autograd
+
+        aot_eager_graph = aot_autograd(
+            fw_compiler=fw_compiler,
+            bw_compiler=fw_compiler,
+            partition_fn=min_cut_rematerialization_partition,
+        )
+
+        compiled_model = torch.compile(model, backend=aot_eager_graph)
+        return compiled_model, fw_graph_cell
+
+    @with_comms
+    @dist_config.patch(compile_on_one_rank=True)
+    def test_compiled_rowwise_embedding_graph_consistency(self):
+        """Test that compiled graphs are identical across all ranks.
+
+        When rowwise sharded embeddings are compiled with torch.compile, the
+        _MaskPartial._mask_tensor() function generates bounds checking
+        operations (lt, ge, sub, index_put) with rank-specific values that get
+        baked into the compiled graph:
+        - Rank 0: lt(index, 0), ge(index, 64), sub(index, 0)
+        - Rank 1: lt(index, 64), ge(index, 128), sub(index, 64)
+
+        These values should be symbolic/dynamic, not baked-in literals, to
+        ensure graph consistency across ranks.
+        """
+        mesh = self.build_device_mesh()
+
+        class Network(nn.Module):
+            def __init__(self, num_embeddings, embedding_dim, device):
+                super().__init__()
+                self.tok_embeddings = nn.Embedding(
+                    num_embeddings, embedding_dim, device=device
+                )
+
+            def forward(self, x):
+                return self.tok_embeddings(x)
+
+        torch.manual_seed(0)
+        num_embeddings = 256
+        embedding_dim = 64
+
+        model = Network(num_embeddings, embedding_dim, device=self.device_type)
+
+        parallelize_module(
+            model,
+            mesh,
+            {
+                "tok_embeddings": RowwiseParallel(
+                    input_layouts=Replicate(),
+                    output_layouts=Shard(1),
+                ),
+            },
+        )
+
+        compiled_model, fw_graph_cell = self._compile_and_capture_graph(model)
+
+        torch.manual_seed(42)
+        inp = torch.randint(0, num_embeddings, (64, 16), device=self.device_type)
+        replicated_inp = DTensor.from_local(inp, mesh, [Replicate()], run_check=False)
+
+        compiled_model(replicated_inp)
+        self._assert_graphs_identical_across_ranks(fw_graph_cell[0])
+
+    @with_comms
+    @dist_config.patch(compile_on_one_rank=True)
+    def test_compiled_dtensor_rng_op_graph_consistency(self):
+        """Compiled random ops on sharded DTensors should produce identical graphs."""
+        mesh = self.build_device_mesh()
+        dt = DTensor.from_local(
+            torch.empty(8, 4, device=self.device_type), mesh, [Shard(0)]
+        )
+
+        fw_graph_cell = [None]
+        fw_compiler = functools.partial(extract_graph, graph_cell=fw_graph_cell)
+
+        from functorch.compile import min_cut_rematerialization_partition
+        from torch._dynamo.backends.common import aot_autograd
+
+        compiled_f = torch.compile(
+            lambda x: torch.rand_like(x),
+            backend=aot_autograd(
+                fw_compiler=fw_compiler,
+                partition_fn=min_cut_rematerialization_partition,
+            ),
+        )
+
+        compiled_f(dt)
+        self._assert_graphs_identical_across_ranks(fw_graph_cell[0])
+
 
 if __name__ == "__main__":
     run_tests()
diff --git a/test/distributed/tensor/test_convolution_ops.py b/test/distributed/tensor/test_convolution_ops.py
index ed1cc60802e70..82221ec247b53 100644
--- a/test/distributed/tensor/test_convolution_ops.py
+++ b/test/distributed/tensor/test_convolution_ops.py
@@ -13,6 +13,7 @@
     Replicate,
     Shard,
 )
+from torch.distributed.tensor.debug import CommDebugMode
 from torch.nn import functional as F
 from torch.testing._internal.common_cuda import with_tf32_off
 from torch.testing._internal.common_utils import run_tests
@@ -294,6 +295,38 @@ def test_conv2d_no_bias_backward(self):
         self.assertIsNotNone(w_dt.grad)
         self.assertEqual(w_dt.grad.shape, torch.Size([8, 4, 3, 3]))
 
+    @with_tf32_off
+    @with_comms
+    @skip_if_lt_x_gpu(2)
+    def test_conv2d_batch_shard_strided(self):
+        """Batch-dim sharding with stride != 1 should not hit _is_supported."""
+        device_mesh = self.build_device_mesh()
+
+        model = nn.Conv2d(3, 8, kernel_size=7, stride=2, padding=3).to(self.device_type)
+        model_ref = copy.deepcopy(model).to(self.device_type)
+        model = distribute_module(model, device_mesh, _conv_fn)
+
+        x = torch.randn(4, 3, 16, 16, device=self.device_type, requires_grad=True)
+        x_ref = x.detach().clone().requires_grad_(True)
+        x_dt = distribute_tensor(x, device_mesh, [Shard(0)])
+
+        comm_mode = CommDebugMode()
+        with comm_mode:
+            out_dt = model(x_dt)
+        self.assertEqual(comm_mode.get_total_counts(), 0)
+
+        out_ref = model_ref(x_ref)
+        self.assertEqual(out_dt.full_tensor(), out_ref)
+
+        grad = torch.randn_like(out_ref)
+        grad_dt = distribute_tensor(grad, device_mesh, [Shard(0)])
+        with comm_mode:
+            out_dt.backward(grad_dt)
+        self.assertEqual(comm_mode.get_total_counts(), 0)
+
+        out_ref.backward(grad)
+        self.assertEqual(x_dt.grad.full_tensor(), x_ref.grad)
+
     @with_comms
     def test_conv2d_module_no_bias(self):
         """Test nn.Conv2d module with bias=False (Issue #167091)
@@ -324,6 +357,72 @@ def test_conv2d_module_no_bias(self):
         # Check that model.bias is None
         self.assertIsNone(model.bias)
 
+    @with_tf32_off
+    @with_comms
+    @skip_if_lt_x_gpu(2)
+    def test_conv1d_batch_shard(self):
+        device_mesh = self.build_device_mesh()
+
+        model = nn.Conv1d(3, 8, kernel_size=3, padding=1).to(self.device_type)
+        model_ref = copy.deepcopy(model).to(self.device_type)
+        model = distribute_module(model, device_mesh, _conv_fn)
+
+        x = torch.randn(4, 3, 16, device=self.device_type, requires_grad=True)
+        x_ref = x.detach().clone().requires_grad_(True)
+        x_dt = distribute_tensor(x, device_mesh, [Shard(0)])
+
+        comm_mode = CommDebugMode()
+        with comm_mode:
+            out_dt = model(x_dt)
+        self.assertEqual(comm_mode.get_total_counts(), 0)
+        self.assertTrue(out_dt.placements[0].is_shard(0))
+
+        out_ref = model_ref(x_ref)
+        self.assertEqual(out_dt.full_tensor(), out_ref)
+
+        grad = torch.randn_like(out_ref)
+        grad_dt = distribute_tensor(grad, device_mesh, [Shard(0)])
+        with comm_mode:
+            out_dt.backward(grad_dt)
+        self.assertEqual(comm_mode.get_total_counts(), 0)
+
+        out_ref.backward(grad)
+        self.assertTrue(x_dt.grad.placements[0].is_shard(0))
+        self.assertEqual(x_dt.grad.full_tensor(), x_ref.grad)
+
+    @with_tf32_off
+    @with_comms
+    @skip_if_lt_x_gpu(2)
+    def test_conv3d_batch_shard(self):
+        device_mesh = self.build_device_mesh()
+
+        model = nn.Conv3d(3, 8, kernel_size=3, padding=1).to(self.device_type)
+        model_ref = copy.deepcopy(model).to(self.device_type)
+        model = distribute_module(model, device_mesh, _conv_fn)
+
+        x = torch.randn(4, 3, 8, 8, 8, device=self.device_type, requires_grad=True)
+        x_ref = x.detach().clone().requires_grad_(True)
+        x_dt = distribute_tensor(x, device_mesh, [Shard(0)])
+
+        comm_mode = CommDebugMode()
+        with comm_mode:
+            out_dt = model(x_dt)
+        self.assertEqual(comm_mode.get_total_counts(), 0)
+        self.assertTrue(out_dt.placements[0].is_shard(0))
+
+        out_ref = model_ref(x_ref)
+        self.assertEqual(out_dt.full_tensor(), out_ref)
+
+        grad = torch.randn_like(out_ref)
+        grad_dt = distribute_tensor(grad, device_mesh, [Shard(0)])
+        with comm_mode:
+            out_dt.backward(grad_dt)
+        self.assertEqual(comm_mode.get_total_counts(), 0)
+
+        out_ref.backward(grad)
+        self.assertTrue(x_dt.grad.placements[0].is_shard(0))
+        self.assertEqual(x_dt.grad.full_tensor(), x_ref.grad)
+
 
 DistConvolutionOpsTestWithLocalTensor = create_local_tensor_test_class(
     DistConvolutionOpsTest,
@@ -332,10 +431,13 @@ def test_conv2d_module_no_bias(self):
         "test_conv_backward_none_grad_inp",
         "test_depthwise_convolution",
         "test_downsampling_convolution",
+        "test_conv2d_batch_shard_strided",
         # New tests for Issue #167091 - use send/recv via tp_convolution
         "test_conv2d_no_bias_compile",
         "test_conv2d_no_bias_backward",
         "test_conv2d_module_no_bias",
+        "test_conv1d_batch_shard",
+        "test_conv3d_batch_shard",
     ],
 )
 
diff --git a/test/distributed/tensor/test_decompositions.py b/test/distributed/tensor/test_decompositions.py
index f5ac64fad4ce8..a2d1621f46e83 100644
--- a/test/distributed/tensor/test_decompositions.py
+++ b/test/distributed/tensor/test_decompositions.py
@@ -68,7 +68,7 @@ def op1(x, y):
             return op2(x, y) * 1.0
 
         def op2(x, y):
-            return torch.mm(x, y)
+            return x @ y
 
         We also test that sharding prop caching kicks in for decompositions;
         1) calling op1 twice should cache hit
@@ -83,11 +83,11 @@ def op2(x, y):
 
             @torch.library.impl(my_lib, "op1", "CPU")
             def op1_impl(x, y):
-                return torch.mm(x, y) * 1.0
+                return x @ y * 1.0
 
             @torch.library.impl(my_lib, "op2", "CPU")
             def op2_impl(x, y):
-                return torch.mm(x, y)
+                return x @ y
 
             @register_decomposition(torch.ops.sharding_decomps.op1.default)
             def op1_decomp(x, y):
@@ -95,15 +95,15 @@ def op1_decomp(x, y):
 
             @register_decomposition(torch.ops.sharding_decomps.op2.default)
             def op2_decomp(x, y):
-                return torch.mm(x, y)
+                return x @ y
 
             @torch.library.register_fake("sharding_decomps::op1")
             def op1_meta(x, y):
-                return torch.mm(torch.empty_like(x), torch.empty_like(y))
+                return torch.empty_like(x) @ torch.empty_like(y)
 
             @torch.library.register_fake("sharding_decomps::op2")
             def op2_meta(x, y):
-                return torch.mm(torch.empty_like(x), torch.empty_like(y))
+                return torch.empty_like(x) @ torch.empty_like(y)
 
             mesh = DeviceMesh("cpu", torch.arange(self.world_size).reshape(-1, 2))
             x = d_empty(16, 16, device_mesh=mesh, placements=[Shard(0), Shard(1)])
@@ -146,16 +146,12 @@ def check_no_strategy(op):
             self.assertTrue(op not in sharding_prop.op_to_rules)
 
         # binary_cross_entropy_with_logits
-        # TODO(pianpwk): fix no sharding strategy for decomp <-> CIA ops
         check_no_strategy(aten.binary_cross_entropy_with_logits.default)
         input = d_empty(16, device_mesh=mesh, placements=[Shard(0)])
         target = d_empty(16, device_mesh=mesh, placements=[Shard(0)])
         weight = d_empty(16, device_mesh=mesh, placements=[Shard(0)])
-        with self.assertRaisesRegex(
-            NotImplementedError, "does not have a sharding strategy"
-        ):
-            out = aten.binary_cross_entropy_with_logits.default(input, target, weight)
-            self.assertEqual(out.placements, (Partial("avg"),))
+        out = aten.binary_cross_entropy_with_logits.default(input, target, weight)
+        self.assertEqual(out.placements, (Partial("avg"),))
 
         # mse_loss
         check_no_strategy(aten.mse_loss.default)
@@ -179,8 +175,7 @@ def check_no_strategy(op):
         out = aten.smooth_l1_loss.default(input, target)
         self.assertEqual(out.placements, (Partial("avg"), Partial("avg")))
 
-        # expand_copy
-        check_no_strategy(aten.expand_copy.default)
+        # expand_copy: has a registered strategy (same as expand)
         input = d_empty(16, 1, device_mesh=mesh, placements=[Partial("min")])
         out = aten.expand_copy.default(input, [-1, 16])
         self.assertEqual(out.placements, (Partial("min"),))
@@ -188,21 +183,67 @@ def check_no_strategy(op):
         # glu: force replicate
         check_no_strategy(aten.glu.default)
         x = d_empty(16, device_mesh=mesh, placements=[Partial()])
-        with self.assertRaisesRegex(
-            NotImplementedError, "does not have a sharding strategy"
-        ):
-            out = aten.glu.default(x)
-            self.assertEqual(out.placements, (Replicate(),))
+        out = aten.glu.default(x)
+        self.assertEqual(out.placements, (Replicate(),))
+
+        # index_add: decomposes into index_put with accumulate=True
+        check_no_strategy(aten.index_add.default)
+        input = d_empty(4, 8, device_mesh=mesh, placements=[Shard(1)])
+        index = distribute_tensor(torch.tensor([0, 2]), mesh, [Replicate()])
+        source = d_empty(2, 8, device_mesh=mesh, placements=[Shard(1)])
+        out = aten.index_add.default(input, 0, index, source)
+        self.assertEqual(out.placements, (Shard(1),))
 
         # polar: force replicate
         check_no_strategy(aten.polar.default)
         x = d_empty(16, device_mesh=mesh, placements=[Partial()])
         y = d_empty(16, device_mesh=mesh, placements=[Partial()])
-        with self.assertRaisesRegex(
-            NotImplementedError, "does not have a sharding strategy"
-        ):
-            out = aten.polar.default(x, y)
-            self.assertEqual(out.placements, (Replicate(),))
+        out = aten.polar.default(x, y)
+        self.assertEqual(out.placements, (Replicate(),))
+
+    def test_roll_flip_strategies(self):
+        """roll and flip unshard on active dims, keep sharding on others."""
+        from torch.distributed.tensor import empty as d_empty
+
+        aten = torch.ops.aten
+        mesh = DeviceMesh("cpu", torch.arange(self.world_size))
+
+        # roll: sharded on non-roll dim stays sharded
+        x = d_empty(16, 16, device_mesh=mesh, placements=[Shard(0)])
+        out = aten.roll.default(x, [2], [1])
+        self.assertEqual(out.placements, (Shard(0),))
+
+        # roll with no dims (flattened): always replicates
+        x = d_empty(16, 16, device_mesh=mesh, placements=[Shard(0)])
+        out = aten.roll.default(x, [2], [])
+        self.assertEqual(out.placements, (Replicate(),))
+
+        # flip: sharded on non-flip dim stays sharded
+        x = d_empty(16, 16, device_mesh=mesh, placements=[Shard(0)])
+        out = aten.flip.default(x, [1])
+        self.assertEqual(out.placements, (Shard(0),))
+
+        # flip: sharded on flip dim gets replicated
+        x = d_empty(16, 16, device_mesh=mesh, placements=[Shard(0)])
+        out = aten.flip.default(x, [0])
+        self.assertEqual(out.placements, (Replicate(),))
+
+    def test_fft_strategies(self):
+        """FFT primitives unshard on transform dims, keep sharding on others."""
+        from torch.distributed.tensor import empty as d_empty
+
+        aten = torch.ops.aten
+        mesh = DeviceMesh("cpu", torch.arange(self.world_size))
+
+        # _fft_c2c: sharded on non-transform dim stays sharded
+        x = d_empty(16, 16, device_mesh=mesh, placements=[Shard(0)], dtype=torch.cfloat)
+        out = aten._fft_c2c.default(x, [1], 0, True)
+        self.assertEqual(out.placements, (Shard(0),))
+
+        # _fft_c2c: sharded on transform dim gets replicated
+        x = d_empty(16, 16, device_mesh=mesh, placements=[Shard(1)], dtype=torch.cfloat)
+        out = aten._fft_c2c.default(x, [1], 0, True)
+        self.assertEqual(out.placements, (Replicate(),))
 
 
 class TestDecompShardingWithComms(DTensorTestBase):
diff --git a/test/distributed/tensor/test_dtensor.py b/test/distributed/tensor/test_dtensor.py
index 0aa1d65d73d90..e14f355d97194 100644
--- a/test/distributed/tensor/test_dtensor.py
+++ b/test/distributed/tensor/test_dtensor.py
@@ -27,6 +27,7 @@
     TensorMeta,
 )
 from torch.distributed.tensor._redistribute import redistribute_local_tensor
+from torch.distributed.tensor._utils import compute_local_tensor_info
 from torch.distributed.tensor.debug import CommDebugMode
 from torch.distributed.tensor.experimental import implicit_replication
 from torch.distributed.tensor.parallel import (
@@ -34,6 +35,7 @@
     parallelize_module,
     RowwiseParallel,
 )
+from torch.distributed.tensor.placement_types import _StridedShard
 from torch.testing import make_tensor
 from torch.testing._internal.common_utils import (
     IS_FBCODE,
@@ -342,6 +344,62 @@ def test_from_local_backward(self):
         # Gradient should flow back without redistribution since grad is already Replicate
         self.assertEqual(comm_mode.get_total_counts(), 0)
 
+    @with_comms
+    def test_from_local_grad_placements(self):
+        """Test that from_local with explicit grad_placements overrides the default
+        gradient placement normalization in backward.
+
+        Verifies:
+        - grad_placements=(Partial(),) with Replicate forward keeps gradient as Partial (no all-reduce)
+        - grad_placements=None (default) preserves existing behavior
+        """
+        device_mesh = self.build_device_mesh()
+
+        # Test 1: Replicate(fwd) with grad_placements=(Partial(),)
+        # Default behavior would all-reduce Partial grad to Replicate, but
+        # grad_placements=(Partial(),) should keep gradient as Partial (no all-reduce)
+        comm_mode = CommDebugMode()
+        with comm_mode:
+            local_tensor = torch.randn(
+                3, 3, device=self.device_type, requires_grad=True
+            )
+            dt = DTensor.from_local(
+                local_tensor,
+                device_mesh,
+                [Replicate()],
+                grad_placements=[Partial()],
+            )
+            output = dt * 2
+            grad = DTensor.from_local(
+                torch.ones(3, 3, device=self.device_type), device_mesh, [Partial()]
+            )
+            output.backward(grad)
+            # With grad_placements=(Partial(),), the backward should NOT all-reduce
+            # the Partial gradient to Replicate. The gradient stays Partial, so
+            # each rank gets its local partial gradient (ones * 2) without all-reduce.
+            self.assertEqual(local_tensor.grad, torch.ones(3, 3) * 2)
+        # No all-reduce should happen since grad_placements keeps it as Partial
+        self.assertEqual(comm_mode.get_total_counts(), 0)
+
+        # Test 2: Replicate(fwd) with grad_placements=None (default behavior preserved)
+        # Partial grad_output should be all-reduced to Replicate (default normalization)
+        comm_mode = CommDebugMode()
+        with comm_mode:
+            local_tensor = torch.randn(
+                3, 3, device=self.device_type, requires_grad=True
+            )
+            dt = DTensor.from_local(local_tensor, device_mesh, [Replicate()])
+            output = dt * 2
+            grad = DTensor.from_local(
+                torch.ones(3, 3, device=self.device_type), device_mesh, [Partial()]
+            )
+            output.backward(grad)
+            # Default: Partial grad should be all-reduced to Replicate
+            expected_grad = torch.ones(3, 3) * 2 * self.world_size
+            self.assertEqual(local_tensor.grad, expected_grad)
+        # Verify that an all-reduce happened (default Partial -> Replicate)
+        self.assertEqual(comm_mode.get_comm_counts()[c10d_functional.all_reduce], 1)
+
     @with_comms
     def test_from_local_uneven_sharding(self):
         device_mesh = self.build_device_mesh()
@@ -467,7 +525,10 @@ def test_to_local(self):
         # test the case under no-grad we directly return the local tensor
         with torch.no_grad():
             local_no_grad = sharded_tensor.to_local()
-            assert local_no_grad is sharded_tensor._local_tensor
+            if local_no_grad is not sharded_tensor._local_tensor:
+                raise AssertionError(
+                    "Expected local_no_grad to be sharded_tensor._local_tensor"
+                )
 
     @with_comms
     def test_to_local_grad_hint(self):
@@ -573,7 +634,9 @@ def test_dtensor_new_empty_strided(self):
         # test backward new_empty_strided with sharding works correctly
         my_dtensor.to_local().sum().backward()
         local_tensor.sum().backward()
-        self.assertEqual(my_dtensor.grad, new_strided_dtensor.grad)
+        self.assertEqual(
+            my_dtensor.grad.full_tensor(), new_strided_dtensor.grad.full_tensor()
+        )
         self.assertEqual(
             my_dtensor.grad.redistribute(placements=[Replicate()]).to_local(),
             local_tensor.grad,
@@ -825,6 +888,76 @@ def test_ops_zero_size_shards(self):
         self.assertEqual(out_t.placements, (Shard(1),))
         self.assertEqual(out_t.full_tensor(), dt.full_tensor().T)
 
+    @with_comms
+    def test_autograd_function_without_materialize_grads(self):
+        device_mesh = self.build_device_mesh()
+
+        class MultiOutputFunc(torch.autograd.Function):
+            @staticmethod
+            def forward(ctx, x, y):
+                ctx.set_materialize_grads(False)
+                return (x * 2), (y * 3)
+
+            @staticmethod
+            def backward(ctx, grad_out1, grad_out2):
+                if grad_out2 is not None:
+                    raise AssertionError("grad_out2 should be None")
+                return grad_out1 * 2, grad_out2
+
+        x_local = torch.randn(4, 4, device=self.device_type, requires_grad=True)
+        y_local = torch.randn(4, 4, device=self.device_type, requires_grad=True)
+        x = DTensor.from_local(x_local, device_mesh, [Replicate()])
+        y = DTensor.from_local(y_local, device_mesh, [Replicate()])
+
+        out1, _ = MultiOutputFunc.apply(x, y)
+
+        loss = out1.sum()
+        loss.backward()
+
+    @with_comms
+    def test_assert_equal_dtensor(self):
+        mesh = self.build_device_mesh()
+        local = torch.randn(4, 4, device=self.device_type)
+
+        dt1 = DTensor.from_local(local.clone(), mesh, [Replicate()])
+        dt2 = DTensor.from_local(local.clone(), mesh, [Replicate()])
+
+        self.assertEqual(dt1, dt2)
+        torch.testing.assert_close(dt1, dt2)
+
+        dt3 = DTensor.from_local(
+            torch.randn(4, 4, device=self.device_type), mesh, [Replicate()]
+        )
+        with self.assertRaisesRegex(AssertionError, "Tensor-likes are not close"):
+            self.assertEqual(dt1, dt3)
+
+        dt_shard = DTensor.from_local(local.clone(), mesh, [Shard(0)])
+        with self.assertRaisesRegex(AssertionError, "DTensor placements do not match"):
+            self.assertEqual(dt1, dt_shard)
+
+        with self.assertRaisesRegex(
+            TypeError, "Comparing a DTensor to a non-DTensor is ambiguous"
+        ):
+            self.assertEqual(dt1, local)
+        with self.assertRaisesRegex(
+            TypeError, "Comparing a DTensor to a non-DTensor is ambiguous"
+        ):
+            self.assertEqual(local, dt1)
+        with self.assertRaisesRegex(
+            TypeError, "Comparing a DTensor to a non-DTensor is ambiguous"
+        ):
+            torch.testing.assert_close(dt1, local)
+
+        dt_scalar = DTensor.from_local(
+            torch.tensor(42.0, device=self.device_type), mesh, [Replicate()]
+        )
+        with self.assertRaisesRegex(
+            TypeError, "Comparing a DTensor to a non-DTensor is ambiguous"
+        ):
+            self.assertEqual(dt_scalar, 42.0)
+        self.assertEqual(dt_scalar.full_tensor(), 42.0)
+        self.assertEqual(dt_scalar.to_local(), 42.0)
+
 
 DTensorTestWithLocalTensor = create_local_tensor_test_class(
     DTensorTest,
@@ -839,6 +972,73 @@ def test_ops_zero_size_shards(self):
 )
 
 
+class DTensorSubclassTest(DTensorTestBase):
+    def _make_dtensor(self, cls, mesh):
+        base = DTensor.from_local(
+            torch.randn(4, 4, device=self.device_type), mesh, [Replicate()]
+        )
+        return cls(base._local_tensor, base._spec, requires_grad=False)
+
+    @with_comms
+    def test_subclass_custom_dispatch(self):
+        """Subclass handles the entire dispatch, never calling DTensor dispatch."""
+        mesh = init_device_mesh(self.device_type, (self.world_size,))
+        called = False
+
+        class MyDTensor(DTensor):
+            @classmethod
+            def __torch_dispatch__(cls, func, types, args=(), kwargs=None):
+                nonlocal called
+                called = True
+                self.assertIs(cls, MyDTensor)
+                return NotImplemented
+
+        my_dt = self._make_dtensor(MyDTensor, mesh)
+        try:
+            my_dt + my_dt
+        except TypeError:
+            pass
+        self.assertTrue(called, "MyDTensor.__torch_dispatch__ was not called")
+
+    @with_comms
+    def test_subclass_no_override(self):
+        """Subclass without __torch_dispatch__ override uses C++ fast path."""
+        mesh = init_device_mesh(self.device_type, (self.world_size,))
+
+        class MyDTensor(DTensor):
+            pass
+
+        my_dt = self._make_dtensor(MyDTensor, mesh)
+        result = my_dt + my_dt
+        self.assertIsInstance(result, DTensor)
+
+    @with_comms
+    def test_subclass_conditional_dispatch(self):
+        """Subclass handles some ops itself, delegates others to DTensor."""
+        mesh = init_device_mesh(self.device_type, (self.world_size,))
+        custom_ops: list[torch._ops.OpOverload] = []
+
+        class MyDTensor(DTensor):
+            @classmethod
+            def __torch_dispatch__(cls, func, types, args=(), kwargs=None):
+                if func == torch.ops.aten.add.Tensor:
+                    custom_ops.append(func)
+                return super().__torch_dispatch__(func, types, args, kwargs)
+
+        my_dt = self._make_dtensor(MyDTensor, mesh)
+
+        # add: subclass logs it, then delegates to DTensor
+        result = my_dt + my_dt
+        self.assertIsInstance(result, DTensor)
+        self.assertIn(torch.ops.aten.add.Tensor, custom_ops)
+
+        # neg: subclass doesn't do custom handling, just delegates
+        custom_ops.clear()
+        result = -my_dt
+        self.assertIsInstance(result, DTensor)
+        self.assertEqual(len(custom_ops), 0)
+
+
 class DTensorMeshTest(DTensorTestBase):
     @property
     def world_size(self):
@@ -1522,6 +1722,81 @@ def test_default_shard_order(self):
             DTensorSpec.is_default_device_order(tensor_global._spec.shard_order)
         )
 
+    @with_comms
+    def test_num_shards_with_strided_shard(self):
+        """DTensorSpec.num_shards must count _StridedShard placements."""
+        mesh = init_device_mesh(self.device_type, (self.world_size,))
+        full = torch.randn(4, self.world_size * 2, 6, device=self.device_type)
+        dt = distribute_tensor(full, mesh, [Shard(1)])
+        dt_flat = dt.flatten(0, 1)
+
+        self.assertIsInstance(dt_flat.placements[0], _StridedShard)
+        self.assertEqual(dt_flat._spec.num_shards, self.world_size)
+
+    @with_comms
+    def test_dim_map_with_strided_shard(self):
+        """DTensorSpec.dim_map must map _StridedShard dims to mesh dims."""
+        mesh = init_device_mesh(self.device_type, (self.world_size,))
+        full = torch.randn(4, self.world_size * 2, 6, device=self.device_type)
+        dt = distribute_tensor(full, mesh, [Shard(1)])
+        dt_flat = dt.flatten(0, 1)
+
+        self.assertIsInstance(dt_flat.placements[0], _StridedShard)
+        shard_dim = dt_flat.placements[0].dim
+        self.assertEqual(dt_flat._spec.dim_map[shard_dim], 0)
+
+    @with_comms
+    def test_num_shards_map_with_strided_shard(self):
+        """DTensorSpec.num_shards_map must count _StridedShard placements per dim."""
+        mesh = init_device_mesh(self.device_type, (self.world_size,))
+        full = torch.randn(4, self.world_size * 2, 6, device=self.device_type)
+        dt = distribute_tensor(full, mesh, [Shard(1)])
+        dt_flat = dt.flatten(0, 1)
+
+        self.assertIsInstance(dt_flat.placements[0], _StridedShard)
+        shard_dim = dt_flat.placements[0].dim
+        self.assertEqual(dt_flat._spec.num_shards_map[shard_dim], self.world_size)
+
+    @with_comms
+    def test_compute_local_tensor_info_with_strided_shard(self):
+        """compute_local_tensor_info must handle _StridedShard placements."""
+        mesh = init_device_mesh(self.device_type, (self.world_size,))
+        full = torch.randn(4, self.world_size * 2, 6, device=self.device_type)
+        dt = distribute_tensor(full, mesh, [Shard(1)])
+        dt_flat = dt.flatten(0, 1)
+
+        self.assertIsInstance(dt_flat.placements[0], _StridedShard)
+
+        global_tensor = full.flatten(0, 1)
+        local_shape, local_stride = compute_local_tensor_info(
+            global_tensor, mesh, dt_flat.placements
+        )
+        shard_dim = dt_flat.placements[0].dim
+
+        # Verify shape: shard dim divided by world_size, other dims unchanged
+        expected_shape = list(global_tensor.shape)
+        expected_shape[shard_dim] //= self.world_size
+        self.assertEqual(local_shape, expected_shape)
+
+        # Verify stride: non-shard dims with large strides should be shrunk
+        self.assertEqual(len(local_stride), len(list(global_tensor.stride())))
+
+        # _StridedShard on dim 1: shard dim changes, strides on dim 0 should shrink
+        full2 = torch.randn(3, 5, self.world_size * 2, device=self.device_type)
+        dt2 = distribute_tensor(full2, mesh, [Shard(2)])
+        dt2_flat = dt2.flatten(1, 2)  # _StridedShard(dim=1)
+
+        self.assertIsInstance(dt2_flat.placements[0], _StridedShard)
+        self.assertEqual(dt2_flat.placements[0].dim, 1)
+
+        global_tensor2 = full2.flatten(1, 2)
+        local_shape2, _ = compute_local_tensor_info(
+            global_tensor2, mesh, dt2_flat.placements
+        )
+        expected_shape2 = list(global_tensor2.shape)
+        expected_shape2[1] //= self.world_size
+        self.assertEqual(local_shape2, expected_shape2)
+
 
 TestDTensorSpecWithLocalTensor = create_local_tensor_test_class(
     TestDTensorSpec,
diff --git a/test/distributed/tensor/test_dtensor_compile.py b/test/distributed/tensor/test_dtensor_compile.py
index ed5aec9ef57e4..cd91af3f0404e 100644
--- a/test/distributed/tensor/test_dtensor_compile.py
+++ b/test/distributed/tensor/test_dtensor_compile.py
@@ -5,6 +5,7 @@
 import copy
 import functools
 import unittest
+from contextlib import ExitStack
 from unittest.mock import patch
 
 import torch
@@ -13,6 +14,11 @@
 import torch.distributed as dist
 import torch.nn as nn
 from torch._C import FileCheck
+from torch._dynamo.functional_export import dynamo_graph_capture_for_export
+from torch._dynamo.testing import AotEagerAndRecordGraphs
+from torch._functorch._aot_autograd.autograd_cache import check_cacheable
+from torch._functorch.aot_autograd import aot_export_joint_with_descriptors
+from torch._guards import tracing
 from torch._inductor.utils import run_and_get_triton_code
 from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import (
     checkpoint_wrapper,
@@ -30,6 +36,7 @@
     Shard,
 )
 from torch.distributed.tensor._dtensor_spec import DTensorSpec, TensorMeta
+from torch.distributed.tensor._redistribute import redistribute_local_tensor
 from torch.distributed.tensor.parallel import (
     ColwiseParallel,
     loss_parallel,
@@ -38,7 +45,8 @@
     PrepareModuleOutput,
     RowwiseParallel,
 )
-from torch.distributed.tensor.placement_types import _StridedShard
+from torch.distributed.tensor.placement_types import _StridedShard, Placement
+from torch.fx.experimental.proxy_tensor import make_fx
 from torch.testing._internal.common_device_type import skipXPUIf
 from torch.testing._internal.common_distributed import skip_if_lt_x_gpu
 from torch.testing._internal.common_fsdp import get_devtype
@@ -48,6 +56,7 @@
     run_tests,
     skipIfHpu,
     skipIfTorchDynamo,
+    skipIfXpu,
 )
 from torch.testing._internal.distributed._tensor.common_dtensor import (
     DTensorTestBase,
@@ -63,7 +72,7 @@
 dev_type = torch.device(get_devtype())
 
 
-class PytreeTuple:
+class PytreeTuple(torch._opaque_base.OpaqueBase):
     """
     Tuple-like values that are treated as leaves of a PyTree.
     """
@@ -110,6 +119,38 @@ def index(self, value):
     def count(self, value):
         return self._values.count(value)
 
+    def __fx_repr__(self) -> tuple[str, dict[str, type]]:
+        # Return a repr string that can reconstruct this object and the types needed
+        # Collect all types used in the values
+        types_dict: dict[str, type] = {"PytreeTuple": PytreeTuple}
+        for val in self._values:
+            val_type = type(val)
+            types_dict[val_type.__name__] = val_type
+            # Also handle nested types (e.g., Shard inside a list)
+            if hasattr(val, "__iter__") and not isinstance(val, str):
+                for item in val:
+                    item_type = type(item)
+                    types_dict[item_type.__name__] = item_type
+        return repr(self), types_dict
+
+
+# Register PytreeTuple as an opaque value type to enable Dynamo to handle
+# instances created during tracing
+from torch._library.opaque_object import MemberType, register_opaque_type
+
+
+register_opaque_type(
+    PytreeTuple,
+    typ="value",
+    members={
+        "__getitem__": MemberType.USE_REAL,
+        "__iter__": MemberType.USE_REAL,
+        "__len__": MemberType.USE_REAL,
+        "__eq__": MemberType.USE_REAL,
+        "__hash__": MemberType.USE_REAL,
+    },
+)
+
 
 class SimpleModel(nn.Module):
     def __init__(self, device):
@@ -207,6 +248,53 @@ def fn(x):
         res = fn(x)
         res.to_local().sum().backward()
 
+    @unittest.skipIf(not torch.accelerator.is_available(), "accelerator not available")
+    def test_dtensor_basic_compile(self):
+        mesh = DeviceMesh(self.device_type, torch.arange(self.world_size))
+
+        param = torch.randn(4, 4)
+        param_x = DTensor.from_local(param, mesh, [Shard(0)], run_check=False)
+
+        class Foo(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.buffer = torch.nn.Buffer(param_x)
+
+            def forward(self, x):
+                inter = self.buffer + DTensor.from_local(
+                    x, self.buffer.device_mesh, [Shard(0)], run_check=False
+                )
+                return inter.to_local()
+
+        m = Foo()
+        inp = torch.randn(4, 4, dtype=torch.float64)
+        backend = AotEagerAndRecordGraphs()
+        res = torch.compile(m, backend=backend, fullgraph=True)(inp)
+        self.assertEqual(res, m(inp))
+        self.assertExpectedInline(
+            str(backend.graphs[0].code).strip(),
+            """\
+def forward(self, L_self_buffers_buffer_ : torch.distributed.tensor.DTensor, L_self_buffers_buffer_device_mesh : torch.distributed.device_mesh.DeviceMesh, L_x_ : torch.Tensor):
+    l_self_buffers_buffer_ = L_self_buffers_buffer_
+    l_self_buffers_buffer_device_mesh = L_self_buffers_buffer_device_mesh
+    l_x_ = L_x_
+    from_local = torch.distributed.tensor._api.from_local(l_x_, l_self_buffers_buffer_device_mesh, [torch.distributed.tensor.placement_types.Shard(dim=0)], run_check = False);  l_x_ = l_self_buffers_buffer_device_mesh = None
+    inter = l_self_buffers_buffer_ + from_local;  l_self_buffers_buffer_ = from_local = None
+    to_local = inter.to_local();  inter = None
+    return (to_local,)""",  # noqa: B950
+        )
+        self.assertExpectedInline(
+            str(backend.fw_graphs[0].code).strip(),
+            """\
+def forward(self, arg0_1, arg1_1, arg2_1, arg3_1):
+    _to_copy = torch.ops.aten._to_copy.default(arg3_1, dtype = torch.float64, layout = torch.strided, device = device(type='cuda', index=0));  arg3_1 = None
+    view = torch.ops.aten.view.default(_to_copy, [4, 4]);  _to_copy = None
+    add = torch.ops.aten.add.Tensor(arg0_1, view);  arg0_1 = view = None
+    view_1 = torch.ops.aten.view.default(add, [4, 4]);  add = None
+    return (view_1,)""",  # noqa: B950
+        )
+
+    @skipIfXpu(msg="AssertionError: torch-xpu-ops: 2958")
     @unittest.skipIf(not torch.accelerator.is_available(), "accelerator not available")
     def test_dtensor_basic_export(self):
         mesh = DeviceMesh(self.device_type, torch.arange(self.world_size))
@@ -225,40 +313,45 @@ def forward(self, x):
                 )
                 return inter.to_local()
 
-        torch.utils._pytree.register_constant(
-            torch.distributed.tensor._dtensor_spec.DTensorSpec
-        )
-        torch.utils._pytree.register_constant(DeviceMesh)
+        m = Foo()
+        inp = (torch.randn(4, 4, dtype=torch.float64),)
+
+        gm = dynamo_graph_capture_for_export(m)(*inp)
+        tracing_context = gm.meta.get("tracing_context", None)
 
-        ep = torch.export.export(
-            Foo(), (torch.randn(4, 4, dtype=torch.float64),), strict=False
-        )
         self.assertExpectedInline(
-            str(ep.graph_module.code).strip(),
-            f"""\
-def forward(self, b_buffer, x):
-    _assert_tensor_metadata_default = torch.ops.aten._assert_tensor_metadata.default(x, dtype = torch.float64, device = device(type='cpu'), layout = torch.strided);  _assert_tensor_metadata_default = None
-    to = torch.ops.aten.to.dtype_layout(x, dtype = torch.float64, layout = torch.strided, device = device(type='{self.device_type}'));  x = None
-    view_as = torch.ops.aten.view_as.default(to, to);  to = None
-    dtensor___init__0 = self.dtensor___init__0
-    dtensor_const_func_spec0 = self.dtensor_const_func_spec0
-    flat_apply = torch.ops.higher_order.flat_apply(dtensor_const_func_spec0, dtensor___init__0, view_as, False);  dtensor_const_func_spec0 = dtensor___init__0 = view_as = None
-    add = torch.ops.aten.add.Tensor(b_buffer, flat_apply);  b_buffer = flat_apply = None
-    access_subclass_inner_tensor_default_4 = torch.ops.export.access_subclass_inner_tensor.default(add, '_local_tensor');  add = None
-    view_as_1 = torch.ops.aten.view_as.default(access_subclass_inner_tensor_default_4, access_subclass_inner_tensor_default_4);  access_subclass_inner_tensor_default_4 = None
-    return (view_as_1,)""",  # noqa: B950
+            str(gm.code).strip(),
+            """\
+def forward(self, args_0):
+    _fn_args = (args_0, )
+    L_self_buffers_buffer_ , L_self_buffers_buffer_device_mesh , L_x_ , = self._dynamo_bytecode_flatten(*_fn_args)
+    l_self_buffers_buffer_ = L_self_buffers_buffer_
+    l_self_buffers_buffer_device_mesh = L_self_buffers_buffer_device_mesh
+    l_x_ = L_x_
+    from_local = torch.distributed.tensor._api.from_local(l_x_, l_self_buffers_buffer_device_mesh, [torch.distributed.tensor.placement_types.Shard(dim=0)], run_check = False);  l_x_ = l_self_buffers_buffer_device_mesh = None
+    inter = l_self_buffers_buffer_ + from_local;  l_self_buffers_buffer_ = from_local = None
+    to_local = inter.to_local();  inter = None
+    return self._dynamo_bytecode_unflatten((to_local,), _fn_args)""",  # noqa: B950
         )
 
+        with tracing(tracing_context):
+            with ExitStack() as stack:
+                joint_with_descriptors = aot_export_joint_with_descriptors(
+                    stack,
+                    gm,
+                    inp,
+                )
+                joint_gm = joint_with_descriptors.graph_module
+
         # During tracing, sharding propagation cache is skipped, so an extra dry run for
         # add is performed in _propagate_tensor_meta_non_cached, hence add_1 instead of add
         self.assertExpectedInline(
-            str(ep.run_decompositions({}).graph_module.code).strip(),
+            str(joint_gm.code).strip(),
             f"""\
-def forward(self, b_parametrizations_buffer_original0, x):
-    _assert_tensor_metadata = torch.ops.aten._assert_tensor_metadata.default(x, None, None, torch.float64, device = device(type='cpu'), layout = torch.strided);  _assert_tensor_metadata = None
-    _to_copy = torch.ops.aten._to_copy.default(x, dtype = torch.float64, layout = torch.strided, device = device(type='{self.device_type}', index=0));  x = None
+def forward(self, arg0_1, arg1_1, arg2_1):
+    _to_copy = torch.ops.aten._to_copy.default(arg2_1, dtype = torch.float64, layout = torch.strided, device = device(type='{self.device_type}', index=0));  arg2_1 = None
     view = torch.ops.aten.view.default(_to_copy, [4, 4]);  _to_copy = None
-    add = torch.ops.aten.add.Tensor(b_parametrizations_buffer_original0, view);  b_parametrizations_buffer_original0 = view = None
+    add = torch.ops.aten.add.Tensor(arg0_1, view);  arg0_1 = view = None
     view_1 = torch.ops.aten.view.default(add, [4, 4]);  add = None
     return (view_1,)""",  # noqa: B950
         )
@@ -305,14 +398,44 @@ def fn(x: DeviceMesh):
             group1 = x.get_group(mesh_dim=1)
             return size, coord, group0, group1
 
-        # Can't be fullgraph=True because ProcessGroup is not reconstructible in dynamo
-        compiled_fn = torch.compile(backend="aot_eager")(fn)
-
         mesh = DeviceMesh(self.device_type, torch.arange(self.world_size).unsqueeze(1))
         opt_fn = fn(mesh)
-        compiled_out = compiled_fn(mesh)
+
+        compiled_out = torch.compile(fn, backend="aot_eager", fullgraph=True)(mesh)
         self.assertEqual(opt_fn, compiled_out)
 
+    def test_device_mesh_input_to_dtensor_from_local(self):
+        """Test that DeviceMesh passed as input is correctly used in DTensor.from_local.
+
+        This tests the opaque object remapping: when a DeviceMesh is passed as a
+        function input and used inside DTensor.from_local, the output DTensor
+        should use the runtime mesh, not a compile-time baked-in mesh.
+        """
+
+        def fn(local_tensor, mesh):
+            # Use the passed-in mesh to create a DTensor
+            return DTensor.from_local(local_tensor, mesh, [Shard(0)], run_check=False)
+
+        # Create two different meshes to verify the output uses the runtime mesh
+        mesh1 = DeviceMesh(self.device_type, torch.arange(self.world_size))
+
+        local_tensor = torch.randn(4, 4)
+        compiled_fn = torch.compile(fn, backend="aot_eager", fullgraph=True)
+
+        # First call with mesh1
+        result1 = compiled_fn(local_tensor, mesh1)
+        self.assertIsInstance(result1, DTensor)
+        self.assertIs(result1.device_mesh, mesh1)
+
+        # Create a second mesh (same shape but different object)
+        mesh2 = DeviceMesh(self.device_type, torch.arange(self.world_size))
+        self.assertIsNot(mesh1, mesh2)  # Different objects
+
+        # Second call with mesh2 - the output should use mesh2, not mesh1
+        result2 = compiled_fn(local_tensor, mesh2)
+        self.assertIsInstance(result2, DTensor)
+        self.assertIs(result2.device_mesh, mesh2)
+
     def test_get_local_rank_compile(self):
         mesh = init_device_mesh(
             self.device_type, (self.world_size,), mesh_dim_names=("dp",)
@@ -348,6 +471,88 @@ def fn_without_arg(x):
         res3 = opt_fn3(x)
         self.assertEqual(res3, ref3)
 
+    def test_dtensor_from_local_with_extracted_device_mesh_no_dce(self):
+        """Test that DTensor.from_local() works in compiled code when using
+        a DeviceMesh extracted from another DTensor passed as input.
+
+        This exercises:
+        1. Accessing dtensor.device_mesh inside compiled code
+        2. Passing that DeviceMesh to DTensor.from_local() with a different tensor
+        3. Returning the new DTensor with correct DeviceMesh
+        """
+        mesh1 = DeviceMesh(self.device_type, torch.arange(self.world_size))
+        mesh2 = DeviceMesh(self.device_type, torch.arange(self.world_size))
+
+        def fn(dt, local_tensor):
+            # Extract DeviceMesh from input DTensor
+            extracted_mesh = dt.device_mesh
+            # Create new DTensor with the extracted mesh
+            new_dt = DTensor.from_local(
+                local_tensor,
+                extracted_mesh,
+                [Replicate()],
+                run_check=False,
+            )
+            return new_dt, torch.add(dt, dt)
+
+        # Create input DTensor
+        input_local = torch.randn(4, 4)
+        input_dt1 = DTensor.from_local(input_local, mesh1, [Shard(0)], run_check=False)
+        input_dt2 = DTensor.from_local(input_local, mesh2, [Shard(0)], run_check=False)
+
+        # Create the normal tensor to pass
+        other_tensor = torch.randn(4, 4)
+
+        # Compile and run
+        opt_fn = torch.compile(fn, backend="aot_eager", fullgraph=True)
+        result1, _ = opt_fn(input_dt1, other_tensor)
+        result2, _ = opt_fn(input_dt2, other_tensor)
+
+        # Verify the new DTensor has the correct DeviceMesh
+        self.assertIs(result1.device_mesh, mesh1)
+        self.assertIs(result2.device_mesh, mesh2)
+
+    def test_dtensor_from_local_with_extracted_device_mesh_dce(self):
+        """Test that DTensor.from_local() works in compiled code when using
+        a DeviceMesh extracted from another DTensor passed as input.
+
+        This exercises:
+        1. Accessing dtensor.device_mesh inside compiled code
+        2. Passing that DeviceMesh to DTensor.from_local() with a different tensor
+        3. Returning the new DTensor with correct DeviceMesh
+        """
+        mesh1 = DeviceMesh(self.device_type, torch.arange(self.world_size))
+        mesh2 = DeviceMesh(self.device_type, torch.arange(self.world_size))
+
+        def fn(dt, local_tensor):
+            # Extract DeviceMesh from input DTensor
+            extracted_mesh = dt.device_mesh
+            # Create new DTensor with the extracted mesh
+            new_dt = DTensor.from_local(
+                local_tensor,
+                extracted_mesh,
+                [Replicate()],
+                run_check=False,
+            )
+            return new_dt
+
+        # Create input DTensor
+        input_local = torch.randn(4, 4)
+        input_dt1 = DTensor.from_local(input_local, mesh1, [Shard(0)], run_check=False)
+        input_dt2 = DTensor.from_local(input_local, mesh2, [Shard(0)], run_check=False)
+
+        # Create the normal tensor to pass
+        other_tensor = torch.randn(4, 4)
+
+        # Compile and run
+        opt_fn = torch.compile(fn, backend="aot_eager", fullgraph=True)
+        result1 = opt_fn(input_dt1, other_tensor)
+        result2 = opt_fn(input_dt2, other_tensor)
+
+        # Verify the new DTensor has the correct DeviceMesh
+        self.assertIs(result1.device_mesh, mesh1)
+        self.assertIs(result2.device_mesh, mesh2)
+
     def test_fakify_dtensor(self):
         mesh = DeviceMesh(self.device_type, torch.arange(self.world_size))
 
@@ -544,7 +749,7 @@ def test_dtensor_unbacked_matmuls(self):
         dist.init_process_group("fake", store=FakeStore(), rank=0, world_size=4)
         device_mesh = init_device_mesh(self.device_type, (2, 2))
 
-        def test_placements(x_placements, y_placements, out_placements):
+        def test_placements(x_placements, y_placements, out_placements=None):
             # create DTensors with unbacked outer/inner sizes
             x_dt = d_randn(64, 64, device_mesh=device_mesh, placements=x_placements)
             y_dt = d_randn(64, 64, device_mesh=device_mesh, placements=y_placements)
@@ -554,11 +759,20 @@ def test_placements(x_placements, y_placements, out_placements):
 
             # full-graph capture
             torch._dynamo.reset()
-            fn = torch.compile(torch.mm, backend="aot_eager", fullgraph=True)
+            cnt = torch._dynamo.testing.CompileCounterWithBackend("aot_eager")
+            fn = torch.compile(torch.mm, backend=cnt, fullgraph=True)
             out = fn(x_dt, y_dt)
 
             # check output placements
-            self.assertEqual(out.placements, out_placements)
+            if out_placements is not None:
+                self.assertEqual(out.placements, out_placements)
+
+            # test on uneven shardings
+            dx = d_randn(20, 17, device_mesh=device_mesh, placements=x_placements)
+            dy = d_randn(17, 5, device_mesh=device_mesh, placements=y_placements)
+            out, eager_out = fn(dx, dy), torch.mm(dx, dy)
+            self.assertEqual(cnt.frame_count, 1)
+            self.assertEqual(out.shape, eager_out.shape)
 
         test_placements(
             (Replicate(), Replicate()),
@@ -572,6 +786,14 @@ def test_placements(x_placements, y_placements, out_placements):
             (Replicate(), Shard(0)), (Replicate(), Replicate()), (Replicate(), Shard(0))
         )
 
+        # note: just check that these run; it's reasonable for the output placements to change,
+        # depending on how tie-breaking is done.
+        test_placements((Partial(), Partial()), (Partial(), Partial()))
+        test_placements((Partial(), Partial()), (Shard(0), Replicate()))
+        test_placements((Replicate(), Shard(0)), (Replicate(), Shard(0)))
+        test_placements((Shard(0), Shard(1)), (Shard(1), Shard(0)))
+        test_placements((Partial(), Replicate()), (Shard(0), Shard(0)))
+
     @unittest.skipIf(not HAS_GPU, "requires GPU for RNG support")
     def test_dtensor_matmul_zero_size_shards(self):
         from torch.distributed.tensor import randn as d_randn
@@ -665,6 +887,54 @@ def _run(torch_fn, make_args, unbacked_dims=None):
             unbacked_dims={1: [1]},
         )
 
+    def test_dtensor_matmul_cost_hint_or_upper_bound(self):
+        dist.destroy_process_group()
+        dist.init_process_group("fake", store=FakeStore(), rank=0, world_size=4)
+        device_mesh = init_device_mesh(self.device_type, (4,))
+
+        x_dt = DTensor.from_local(
+            torch.randn(8, 8),
+            device_mesh=device_mesh,
+            placements=[Shard(0)],
+        )
+        y_dt = DTensor.from_local(
+            torch.randn(8, 8),
+            device_mesh=device_mesh,
+            placements=[Shard(1)],
+        )
+        torch._dynamo.decorators.mark_unbacked(x_dt, 0)
+        torch._dynamo.decorators.mark_unbacked(y_dt, 1)
+
+        # large k dim tells compiler it's cheaper to all-gather on x
+        def f1(x, y):
+            torch._check(x.size(0) <= 32)
+            torch._check(y.size(1) <= 16384)
+            return x @ y
+
+        out = torch.compile(f1, backend="aot_eager", fullgraph=True)(x_dt, y_dt)
+        self.assertEqual(out.placements, (Shard(1),))
+
+        # for the reverse, all-gather on y
+        def f2(x, y):
+            torch._check(x.size(0) <= 16384)
+            torch._check(y.size(1) <= 32)
+            return x @ y
+
+        out = torch.compile(f2, backend="aot_eager", fullgraph=True)(x_dt, y_dt)
+        self.assertEqual(out.placements, (Shard(0),))
+
+        # specifying hint_override also determines the strategy
+        def f3(x, y):
+            return x @ y
+
+        x_dt = x_dt + 1
+        y_dt = y_dt + 1
+        torch._dynamo.decorators.mark_unbacked(x_dt, 0, hint_override=16)
+        torch._dynamo.decorators.mark_unbacked(x_dt, 1, hint_override=1024)
+
+        out = torch.compile(f3, backend="aot_eager", fullgraph=True)(x_dt, y_dt)
+        self.assertEqual(out.placements, (Shard(1),))
+
     def test_dtensor_requires_grad_recompile(self):
         cnt = torch._dynamo.testing.CompileCounterWithBackend("aot_eager")
         mesh = DeviceMesh(self.device_type, torch.arange(self.world_size))
@@ -684,6 +954,25 @@ def f(x):
 
         self.assertEqual(cnt.frame_count, 2)
 
+    @with_comms
+    @torch._dynamo.config.patch(trace_autograd_ops=True)
+    def test_dtensor_requires_grad_intermediate_backward(self):
+        mesh = DeviceMesh(self.device_type, torch.arange(self.world_size))
+
+        def fn(x):
+            y = x * 2
+            y.requires_grad_()
+            loss = (y * 3).sum()
+            loss.backward()
+            return y.grad
+
+        full_x = torch.randn(8, 8)
+        x = distribute_tensor(full_x, mesh, [Shard(0)])
+
+        ref = fn(x.full_tensor())
+        result = torch.compile(fn, backend="aot_eager", fullgraph=True)(x)
+        self.assertEqual(ref, result.full_tensor())
+
     def test_dtensor_attribute_access_on_intermediate(self):
         mesh = DeviceMesh(self.device_type, torch.arange(self.world_size))
 
@@ -1117,6 +1406,7 @@ def test_dtensor_different_gradient_placement(self):
 
         def fn(x, y, z):
             permute = x.permute(0, 2, 1)
+            x.device_mesh
             permute2 = permute.contiguous()
             layer_norm = torch.nn.functional.layer_norm(permute2, (4,), y, z, 1e-05)
             out = layer_norm.permute(0, 2, 1)
@@ -1147,8 +1437,23 @@ def fn(x):
 
         x = torch.ones(1)
         ref = fn(x)
-        cnt = torch._dynamo.testing.CompileCounterWithBackend("aot_eager")
-        opt_fn = torch.compile(fn, backend=cnt, fullgraph=True)
+        backend = AotEagerAndRecordGraphs()
+        opt_fn = torch.compile(fn, backend=backend, fullgraph=True)
+        res = opt_fn(x)
+        self.assertEqual(res, ref)
+        self.assertExpectedInline(
+            str(backend.graphs[0].code).strip(),
+            """\
+def forward(self, L_x_ : torch.Tensor, L_mesh_ : torch.distributed.device_mesh.DeviceMesh):
+    l_x_ = L_x_
+    l_mesh_ = L_mesh_
+    dt = torch.distributed.tensor._api.from_local(l_x_, l_mesh_, [torch.distributed.tensor.placement_types.Shard(dim=0)], run_check = False);  l_x_ = None
+    redistribute = dt.redistribute(l_mesh_, [torch.distributed.tensor.placement_types.Replicate()]);  dt = l_mesh_ = None
+    to_local = redistribute.to_local();  redistribute = None
+    add = to_local + 2;  to_local = None
+    return (add,)""",  # noqa: B950
+        )
+        opt_fn = torch.compile(fn, backend="eager", fullgraph=True)
         res = opt_fn(x)
         self.assertEqual(res, ref)
 
@@ -1161,11 +1466,27 @@ def redistribute_kwargs_fn(x):
 
         x = torch.ones(1)
         ref = redistribute_kwargs_fn(x)
+        backend = AotEagerAndRecordGraphs()
         opt_kwargs_fn = torch.compile(
-            redistribute_kwargs_fn, backend=cnt, fullgraph=True
+            redistribute_kwargs_fn, backend=backend, fullgraph=True
         )
         res = opt_kwargs_fn(x)
         self.assertEqual(res, ref)
+        self.assertExpectedInline(
+            str(backend.graphs[0].code).strip(),
+            """\
+def forward(self, L_x_ : torch.Tensor, L_mesh_ : torch.distributed.device_mesh.DeviceMesh):
+    l_x_ = L_x_
+    l_mesh_ = L_mesh_
+    dt = torch.distributed.tensor._api.from_local(l_x_, l_mesh_, [torch.distributed.tensor.placement_types.Shard(dim=0)], run_check = False);  l_x_ = None
+    redistribute = dt.redistribute(device_mesh = l_mesh_, placements = [torch.distributed.tensor.placement_types.Replicate()]);  dt = l_mesh_ = None
+    to_local = redistribute.to_local();  redistribute = None
+    add = to_local + 2;  to_local = None
+    return (add,)""",  # noqa: B950
+        )
+
+        # This should not throw a BypassAOTAutogradCache error
+        check_cacheable(backend.graphs[0])
 
     @skipIfHpu
     def test_dynamo_dtensor_from_local_redistribute_async(self):
@@ -1256,17 +1577,16 @@ def fn(x):
         x2 = x2.to_local()
         self.assertTrue(isinstance(x2, AsyncCollectiveTensor))
         opt_fn(x2)
-        # The important part: we get a wait_tensor() in the graph.
-        # At runtime, the input to the graph is an AsyncCollectiveTensor,
-        # and inside the graph we need to issue a wait() to synchronize.
+        # ACTs are unwrapped before AOT autograd tracing in
+        # process_inputs, so the graph receives a plain tensor with
+        # no wait_tensor op.
         self.assertExpectedInline(
             str(fw_graph_cell[0]).strip(),
             """\
-def forward(self, primals_1):
-    wait_tensor = torch.ops._c10d_functional.wait_tensor.default(primals_1)
-    sin = torch.ops.aten.sin.default(wait_tensor)
+def forward(self, arg0_1):
+    sin = torch.ops.aten.sin.default(arg0_1);  arg0_1 = None
     sin_1 = torch.ops.aten.sin.default(sin);  sin = None
-    return (sin_1, primals_1, wait_tensor)""",
+    return (sin_1,)""",
         )
 
     @skipIfTorchDynamo()
@@ -1459,17 +1779,20 @@ def forward(self, x_1: "f32[4, 4]"):
         add: "f32[4, 4]" = torch.ops.aten.add.Tensor(x_1, 1);  x_1 = None
         view: "f32[4, 4]" = torch.ops.aten.view.default(add, [4, 4]);  add = None
         repeated_subgraph0 = self.repeated_subgraph0
-        invoke_subgraph = torch.ops.higher_order.invoke_subgraph(repeated_subgraph0, 'invoke_subgraph_0', view);  repeated_subgraph0 = view = None
-        getitem: "f32[8, 4]" = invoke_subgraph[0];  invoke_subgraph = None
+        _opaque_obj0 = self._opaque_obj0
+        _opaque_obj0_1 = self._opaque_obj0
+        invoke_subgraph = torch.ops.higher_order.invoke_subgraph(repeated_subgraph0, 'invoke_subgraph_0', view, _opaque_obj0, _opaque_obj0_1);  repeated_subgraph0 = view = _opaque_obj0 = _opaque_obj0_1 = None
+        getitem: "f32[8, 4]" = invoke_subgraph[0]
+        getitem_1 = invoke_subgraph[1];  invoke_subgraph = getitem_1 = None
         view_1: "f32[8, 4]" = torch.ops.aten.view.default(getitem, [8, 4]);  getitem = None
         return view_1
 
     class repeated_subgraph0(torch.nn.Module):
-        def forward(self, arg0_1: "f32[4, 4]"):
+        def forward(self, arg0_1: "f32[4, 4]", arg1_1, arg2_1):
             # No stacktrace found for following nodes
             all_gather_into_tensor: "f32[8, 4]" = torch.ops._c10d_functional.all_gather_into_tensor.default(arg0_1, 2, '0');  arg0_1 = None
             wait_tensor: "f32[8, 4]" = torch.ops._c10d_functional.wait_tensor.default(all_gather_into_tensor);  all_gather_into_tensor = None
-            return (wait_tensor,)""",  # noqa: B950
+            return (wait_tensor, arg2_1)""",  # noqa: B950
         )
 
     @torch._dynamo.config.patch(force_compile_during_fx_trace=True)
@@ -1538,6 +1861,26 @@ def fw_compiler(gm, example_inputs):
             f"Expected 1 compilation, got {compile_counter.frame_count}",
         )
 
+    def test_device_mesh_slice(self):
+        dist.destroy_process_group()
+        dist.init_process_group("fake", store=FakeStore(), rank=0, world_size=8)
+        mesh = init_device_mesh(self.device_type, (2, 4), mesh_dim_names=("dp", "tp"))
+
+        def fn(x):
+            tp_mesh = mesh["tp"]
+            dt = DTensor.from_local(x, tp_mesh, [Shard(0)], run_check=False)
+            if "tp" not in dt.device_mesh.mesh_dim_names:
+                return x
+            return dt.to_local()
+
+        x = torch.randn(4, 4)
+        res = fn(x)
+        traced = make_fx(fn, tracing_mode="fake")(x)
+        self.assertEqual(traced(x), res)
+
+        compiled = torch.compile(fn, backend="aot_eager")(x)
+        self.assertEqual(compiled, res)
+
     def test_compile_redistribute_flattened_mesh(self):
         """
         Test that redistribute works with pre-flattened meshes during compile.
@@ -1575,6 +1918,494 @@ def redistribute_fn(x):
         # Test backward pass
         result.sum().backward()
 
+    @torch._dynamo.config.patch(force_compile_during_fx_trace=True)
+    def test_flattened_submesh_no_getattr_compile_on_one_rank(self):
+        """When compile_on_one_rank=True, the flattened submesh should appear as a
+        custom op call_function node (derived from the mesh graph input) rather than
+        as a get_attr constant holding an unpicklable ProcessGroup."""
+        from torch._library.fake_class_registry import FakeScriptObject
+
+        dist.destroy_process_group()
+        dist.init_process_group("fake", store=FakeStore(), rank=0, world_size=4)
+
+        mesh = init_device_mesh(self.device_type, (2, 2), mesh_dim_names=("dp", "tp"))
+        mesh["dp", "tp"]._flatten("dp_tp")
+
+        import torch._functorch._aot_autograd.graph_compile as gc
+
+        orig = gc._aot_stage2a_partition
+        device_mesh_getattr_count = [0]
+
+        def check_partition(fx_g, joint_inputs, *args, **kwargs):
+            result = orig(fx_g, joint_inputs, *args, **kwargs)
+            fw_module, bw_module = result[0], result[1]
+            count = 0
+            for gm in (fw_module, bw_module):
+                for node in gm.graph.nodes:
+                    if node.op != "get_attr":
+                        continue
+                    val = getattr(gm, node.target, None)
+                    if isinstance(val, FakeScriptObject):
+                        val = getattr(val, "real_obj", val)
+                    if isinstance(val, DeviceMesh):
+                        count += 1
+            device_mesh_getattr_count[0] = count
+            return result
+
+        gc._aot_stage2a_partition = check_partition
+        torch._functorch._aot_autograd.graph_compile._aot_stage2a_partition = (
+            check_partition
+        )
+        try:
+            torch._dynamo.reset()
+            with dist.config.patch(compile_on_one_rank=True):
+
+                def fn(x, mesh):
+                    dt = DTensor.from_local(
+                        x, mesh, [Shard(0), Shard(0)], run_check=False
+                    )
+                    return (
+                        dt.redistribute(mesh, [Replicate(), Replicate()])
+                        .to_local()
+                        .sum()
+                    )
+
+                compiled = torch.compile(fn)
+                x = torch.randn(4, 4, requires_grad=True)
+                loss = compiled(x, mesh)
+                loss.backward()
+
+            self.assertEqual(
+                device_mesh_getattr_count[0],
+                0,
+                "DeviceMesh get_attr nodes survived into partitioned graphs; "
+                "the flattened submesh should use a custom op instead",
+            )
+        finally:
+            gc._aot_stage2a_partition = orig
+            torch._functorch._aot_autograd.graph_compile._aot_stage2a_partition = orig
+            dist.destroy_process_group()
+            dist.init_process_group(
+                "fake", store=FakeStore(), rank=0, world_size=self.world_size
+            )
+
+    def test_sibling_submesh_reconstruct_from_ancestor(self):
+        """When a closure captures an FSDP submesh while a sibling mesh (the
+        concatenated fsdp+tp mesh) is a graph input, the closure-captured mesh
+        must be reconstructed via _get_submesh — not baked as a get_attr
+        constant that breaks serialization.
+
+        This reproduces the SimpleFSDP + TP pattern: ReplicateComputation
+        captures the FSDP submesh as self.device_mesh, while parameters are
+        DTensors on DeviceMesh._concatenate([fsdp_mesh, tp_mesh]). With a
+        1D world mesh unflattened into (fsdp, tp), the root mesh has dim
+        name 'world' — so the reconstruct_fn must find the concatenated
+        ancestor (which contains 'fsdp') rather than looking only at root.
+        """
+        from torch._library.fake_class_registry import FakeScriptObject
+
+        dist.destroy_process_group()
+        dist.init_process_group("fake", store=FakeStore(), rank=0, world_size=8)
+
+        try:
+            # Build 1D world mesh → unflatten → (fsdp=4, tp=2)
+            world_mesh = init_device_mesh(
+                self.device_type, (8,), mesh_dim_names=("world",)
+            )
+            dense_mesh = world_mesh._unflatten(0, (4, 2), ("fsdp", "tp"))
+            fsdp_mesh = dense_mesh["fsdp"]
+            tp_mesh = dense_mesh["tp"]
+
+            # Concatenate like SimpleFSDP's _distribute_dtensor does
+            concat_mesh = DeviceMesh._concatenate([fsdp_mesh, tp_mesh])
+
+            with dist.config.patch(compile_on_one_rank=True):
+
+                def fn(local_tensor, concat_mesh_input):
+                    # concat_mesh_input is a graph input (placeholder).
+                    # Simulate SimpleFSDP: create DTensor on concat mesh,
+                    # extract local, re-wrap on the closure-captured
+                    # fsdp_mesh, then all-gather.
+                    dt = DTensor.from_local(
+                        local_tensor,
+                        concat_mesh_input,
+                        [Shard(0), Shard(1)],
+                        run_check=False,
+                    )
+                    local = dt.to_local()
+                    fsdp_dt = DTensor.from_local(
+                        local, fsdp_mesh, [Shard(0)], run_check=False
+                    )
+                    return (
+                        fsdp_dt.redistribute(fsdp_mesh, [Replicate()]).to_local().sum()
+                    )
+
+                local_t = torch.randn(4, 4)
+                gm = make_fx(fn, tracing_mode="fake")(local_t, concat_mesh)
+
+            device_mesh_getattr_count = 0
+            for node in gm.graph.nodes:
+                if node.op != "get_attr":
+                    continue
+                val = getattr(gm, node.target, None)
+                if isinstance(val, FakeScriptObject):
+                    val = getattr(val, "real_obj", val)
+                if isinstance(val, DeviceMesh):
+                    device_mesh_getattr_count += 1
+
+            self.assertEqual(
+                device_mesh_getattr_count,
+                0,
+                "DeviceMesh get_attr nodes found; the closure-captured FSDP "
+                "submesh should be reconstructed from the tracked concatenated "
+                "ancestor via _get_submesh, not baked as a constant",
+            )
+
+            submesh_ops = [
+                node
+                for node in gm.graph.nodes
+                if node.op == "call_function" and "get_submesh" in str(node.target)
+            ]
+            self.assertGreater(
+                len(submesh_ops),
+                0,
+                "Expected _get_submesh call to derive the FSDP submesh from "
+                "the tracked concatenated ancestor mesh",
+            )
+        finally:
+            dist.destroy_process_group()
+            dist.init_process_group(
+                "fake", store=FakeStore(), rank=0, world_size=self.world_size
+            )
+
+    def test_to_local_symbolic_sizes_non_sharded_dims(self):
+        # Checks that symbolic sizes are not polluted by the sharding
+        # in to_local - for instance we could incorrectly have (16 * (s27//2))
+        # as opposed to just (8 * s27) which enables correct generation
+        # https://github.com/pytorch/pytorch/issues/175690
+        mesh = DeviceMesh(self.device_type, torch.arange(self.world_size))
+
+        def fn(x):
+            local = x.narrow(1, 0, x.shape[1] // self.world_size)
+            dt = DTensor.from_local(
+                local,
+                mesh,
+                [Shard(1)],
+                run_check=False,
+                shape=x.shape,
+                stride=x.stride(),
+            )
+            r = dt.redistribute(placements=[Replicate()])
+            out = r.to_local()
+            return out.view(out.shape[0], -1)
+
+        inp = torch.randn(1, 10, 8)
+        torch._dynamo.mark_dynamic(inp, 1)
+
+        backend = AotEagerAndRecordGraphs()
+        opt_fn = torch.compile(fn, backend=backend, fullgraph=True)
+        result = opt_fn(inp)
+        self.assertEqual(result, fn(inp))
+
+        # The view shape should use the clean symbol (s0), not 2*((s0//2)).
+        fw_code = backend.fw_graphs[0].print_readable(print_output=False)
+        for line in fw_code.splitlines():
+            if "view" in line and "-1" in line:
+                self.assertNotIn("//", line, f"Polluted symbolic shape: {line}")
+
+    def test_to_local_symbolic_sizes_uneven_shard(self):
+        # Regression test to ensure our narrow changes does not cause any
+        # unintentional incorrectness
+        mesh = DeviceMesh(self.device_type, torch.arange(self.world_size))
+
+        def fn(x):
+            local = x.narrow(1, 0, x.shape[1] // self.world_size)
+            dt = DTensor.from_local(
+                local,
+                mesh,
+                [Shard(1)],
+                run_check=False,
+                shape=x.shape,
+                stride=x.stride(),
+            )
+            r = dt.redistribute(placements=[Replicate()])
+            out = r.to_local()
+            return out.view(out.shape[0], -1)
+
+        opt_fn = torch.compile(fn, backend="aot_eager")
+
+        # Divisible: 10 % 2 == 0
+        inp_even = torch.randn(1, 10, 8)
+        torch._dynamo.mark_dynamic(inp_even, 1)
+        result = opt_fn(inp_even)
+        self.assertEqual(result, fn(inp_even))
+
+        # Non-divisible: 9 % 2 != 0, exercises padding/unpadding + recompile
+        # - most importantly ensures correctness
+        inp_odd = torch.randn(1, 9, 8)
+        torch._dynamo.mark_dynamic(inp_odd, 1)
+        result = opt_fn(inp_odd)
+        self.assertEqual(result, fn(inp_odd))
+
+    def test_mesh_get_process_group_1d(self):
+        mesh = DeviceMesh(self.device_type, torch.arange(self.world_size))
+        pg = torch.ops._dtensor.mesh_get_process_group(mesh, 0)
+        self.assertIsInstance(pg, dist.ProcessGroup)
+        self.assertEqual(pg, mesh.get_group(0))
+
+    def test_mesh_get_process_group_2d(self):
+        dist.destroy_process_group()
+        dist.init_process_group("fake", store=FakeStore(), rank=0, world_size=4)
+
+        mesh = DeviceMesh(self.device_type, torch.arange(4).reshape(2, 2))
+        for dim in range(2):
+            pg = torch.ops._dtensor.mesh_get_process_group(mesh, dim)
+            self.assertIsInstance(pg, dist.ProcessGroup)
+            self.assertEqual(pg, mesh.get_group(dim))
+
+    def test_mesh_get_process_group_fake(self):
+        from torch._library.fake_class_registry import (
+            FakeScriptObject,
+            maybe_to_fake_obj,
+        )
+        from torch._subclasses.fake_tensor import FakeTensorMode
+
+        mesh = DeviceMesh(self.device_type, torch.arange(self.world_size))
+
+        with FakeTensorMode() as fake_mode:
+            fake_mesh = maybe_to_fake_obj(fake_mode, mesh)
+            self.assertIsInstance(fake_mesh, FakeScriptObject)
+
+            pg = torch.ops._dtensor.mesh_get_process_group(fake_mesh, 0)
+            self.assertIsInstance(pg, dist.ProcessGroup)
+            self.assertEqual(pg, mesh.get_group(0))
+
+    def test_resolve_group_uses_op_with_compile_on_one_rank(self):
+        """Test that _resolve_group returns a ProcessGroup (not a string name)
+        when compile_on_one_rank is enabled, for both 1D mesh and (mesh, dim)
+        tuple inputs. Also tests that _group_or_group_name passes through the
+        ProcessGroup instead of extracting .group_name."""
+        import torch.distributed.config as dist_config
+        from torch.distributed._functional_collectives import (
+            _group_or_group_name,
+            _resolve_group,
+        )
+
+        mesh_1d = DeviceMesh(self.device_type, torch.arange(self.world_size))
+
+        dist.destroy_process_group()
+        dist.init_process_group("fake", store=FakeStore(), rank=0, world_size=4)
+        mesh_2d = DeviceMesh(self.device_type, torch.arange(4).reshape(2, 2))
+
+        with dist_config.patch(compile_on_one_rank=True):
+            # 1D mesh
+            result = _resolve_group(mesh_1d)
+            self.assertIsInstance(result, dist.ProcessGroup)
+            self.assertEqual(result, mesh_1d.get_group(0))
+
+            # (mesh, dim) tuple
+            for dim in range(2):
+                result = _resolve_group((mesh_2d, dim))
+                self.assertIsInstance(result, dist.ProcessGroup)
+                self.assertEqual(result, mesh_2d.get_group(dim))
+
+            # _group_or_group_name passes ProcessGroup through
+            pg = mesh_1d.get_group(0)
+            result = _group_or_group_name(pg)
+            self.assertIs(result, pg)
+
+        # Without compile_on_one_rank, _resolve_group returns a string name
+        result = _resolve_group(mesh_1d)
+        self.assertIsInstance(result, str)
+
+    @unittest.skipIf(not torch.accelerator.is_available(), "accelerator not available")
+    def test_no_device_mesh_getattr_in_tp_dp_joint_graph(self):
+        """Backward closures of DTensor ops should not capture DeviceMesh as
+        get_attr constants in the joint graph (they break AOTAutograd cache
+        serialization because ProcessGroups are unpicklable)."""
+        import torch.distributed.config as dist_config
+        from torch._library.fake_class_registry import FakeScriptObject
+
+        # Need world_size=4 for a 2x2 mesh; re-init the fake PG.
+        dist.destroy_process_group()
+        dist.init_process_group("fake", store=FakeStore(), rank=0, world_size=4)
+
+        def _distribute_dtensor(
+            tensor: DTensor,
+            device_mesh: DeviceMesh,
+            dp_placements: tuple[Placement, ...],
+        ) -> DTensor:
+            inner_spec = tensor._spec
+            spanned_mesh = DeviceMesh._concatenate([device_mesh, inner_spec.mesh])
+            shard_dim = dp_placements[0].dim
+            split_factor = inner_spec.num_shards_map[shard_dim]
+            tensor_placement = (
+                (
+                    _StridedShard(shard_dim, split_factor=split_factor)
+                    if split_factor > 1
+                    else dp_placements[0]
+                ),
+            ) + inner_spec.placements
+            current_spec = DTensorSpec(
+                mesh=device_mesh,
+                placements=(Replicate(),) * len(dp_placements),
+                tensor_meta=inner_spec.tensor_meta,
+            )
+            target_spec = DTensorSpec(
+                mesh=device_mesh,
+                placements=tuple(dp_placements),
+                tensor_meta=inner_spec.tensor_meta,
+            )
+            result_tensor = redistribute_local_tensor(
+                tensor._local_tensor,
+                current_spec=current_spec,
+                target_spec=target_spec,
+            )
+            return DTensor(
+                result_tensor.requires_grad_(tensor.requires_grad),
+                DTensorSpec(
+                    mesh=spanned_mesh,
+                    placements=tensor_placement,
+                    tensor_meta=inner_spec.tensor_meta,
+                ),
+                requires_grad=tensor.requires_grad,
+            )
+
+        class ReplicateComputation(nn.Module):
+            def __init__(self, device_mesh, param_sharding):
+                super().__init__()
+                self.device_mesh = device_mesh
+                self.param_sharding = param_sharding
+                self.compute_placements = [Replicate()] * device_mesh.ndim
+                self.grad_placements = [Partial(reduce_op="sum")] * device_mesh.ndim
+
+            def forward(self, x):
+                non_dp_mesh_dims = x._spec.mesh.ndim - self.device_mesh.ndim
+                sharded_dtensor = DTensor.from_local(
+                    x.to_local(), self.device_mesh, self.param_sharding
+                )
+                replicated_local_tensor = sharded_dtensor.redistribute(
+                    placements=self.compute_placements,
+                ).to_local(grad_placements=self.grad_placements)
+                non_dp_placements = tuple(x._spec.placements[-non_dp_mesh_dims:])
+                non_dp_mesh_dim_names = tuple(
+                    x._spec.mesh.mesh_dim_names[-non_dp_mesh_dims:]
+                )
+                non_dp_mesh = x._spec.mesh[non_dp_mesh_dim_names]
+                return DTensor.from_local(
+                    replicated_local_tensor, non_dp_mesh, non_dp_placements
+                )
+
+        param_sharding = (Shard(0),)
+        wrap_idx = [0]
+
+        class Model(nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.linear1 = nn.Linear(16, 16, bias=False)
+                self.linear2 = nn.Linear(16, 16, bias=False)
+
+            def forward(self, x):
+                return self.linear2(torch.relu(self.linear1(x))).sum()
+
+        with dist_config.patch("compile_on_one_rank", True):
+            device = torch.device(f"{self.device_type}:0")
+            torch.accelerator.set_device_index(0)
+            mesh_2d = init_device_mesh(
+                self.device_type, (2, 2), mesh_dim_names=("fsdp", "tp")
+            )
+            dp_mesh = mesh_2d["fsdp"]
+
+            with torch.device("meta"):
+                model = Model()
+            parallelize_module(
+                model,
+                mesh_2d["tp"],
+                {
+                    "linear1": ColwiseParallel(),
+                    "linear2": RowwiseParallel(),
+                },
+            )
+            # Apply data parallelism with parametrization
+            for mod in model.modules():
+                params_dict = dict(mod.named_parameters(recurse=False))
+                for p_name, p in params_dict.items():
+                    mod.register_parameter(
+                        p_name,
+                        nn.Parameter(_distribute_dtensor(p, dp_mesh, param_sharding)),
+                    )
+                parametrization = ReplicateComputation(dp_mesh, param_sharding)
+                props = {
+                    pn: property(
+                        lambda self, _pn=pn: parametrization(self._parameters[_pn])
+                    )
+                    for pn in params_dict
+                }
+                idx = wrap_idx[0]
+                wrap_idx[0] += 1
+                cls = type(
+                    f"SimpleFSDP{mod.__class__.__name__}_{idx}",
+                    (mod.__class__,),
+                    props,
+                )
+                mod.__class__ = cls
+                import sys
+
+                sys.modules[cls.__module__].__dict__[cls.__name__] = cls
+
+            model.to_empty(device=device)
+            with dist_config.patch("compile_on_one_rank", False), torch.no_grad():
+                for p in model.parameters():
+                    p.fill_(0.01)
+            model.train()
+
+            x = torch.randn(4, 16, device=device)
+            with (
+                torch._dynamo.config.patch(fake_tensor_cache_enabled=False),
+                torch.fx.traceback.preserve_node_meta(),
+            ):
+                gm = dynamo_graph_capture_for_export(model)(x)
+                tracing_context = gm.meta["tracing_context"]
+
+            with tracing(tracing_context), ExitStack() as stack:
+                jwd = aot_export_joint_with_descriptors(stack, gm, (x,))
+
+            for node in jwd.graph_module.graph.nodes:
+                if node.op != "get_attr":
+                    continue
+                val = getattr(jwd.graph_module, node.target, None)
+                if isinstance(val, FakeScriptObject):
+                    val = getattr(val, "real_obj", val)
+                self.assertNotIsInstance(
+                    val,
+                    DeviceMesh,
+                    "DeviceMesh should not appear as get_attr in the joint graph",
+                )
+
+    def test_compile_optimizer_grad_view_base_dim_mismatch(self):
+        # Regression test for https://github.com/pytorch/pytorch/issues/176667
+        # When param._local_tensor is a view of an N-D base but
+        # grad._local_tensor is a view of a 1-D base (as in FSDP2's flat
+        # gradient buffer), dynamo's meta tensor creation must not reuse the
+        # param's symbolic context for the grad.
+        mesh = DeviceMesh(self.device_type, torch.arange(self.world_size))
+
+        # param: view of 2D base (simulates FSDP2 reshard padding)
+        param_local = torch.randn(32, 256)[:16]
+        param = DTensor.from_local(param_local, mesh, [Replicate()], run_check=False)
+        param.requires_grad_(True)
+        param.retain_grad()
+
+        # grad: view of 1D base (simulates FSDP2 flat gradient buffer)
+        grad_local = torch.randn(16 * 256 + 1000)[: 16 * 256].view(16, 256)
+        param.grad = DTensor.from_local(
+            grad_local, mesh, [Replicate()], run_check=False
+        )
+
+        opt = torch.optim.Adam([param], lr=1e-3)
+        compiled_step = torch.compile(opt.step, backend="aot_eager")
+        compiled_step()
+
 
 @instantiate_parametrized_tests
 class TestDTensorCompileE2E(DTensorTestBase):
@@ -1816,6 +2647,46 @@ def forward(self, x):
         output = sharded_net(replicated_inp)
         self.assertEqual(output.full_tensor(), ref_out)
 
+    @with_comms
+    def test_unbacked_illegal_views(self):
+        """Test that views with unbacked shapes match eager behavior"""
+        device_mesh = self.build_device_mesh()
+
+        def create_dt(shard_dim):
+            tensor = torch.randn(8, 8, 8)
+            dt = distribute_tensor(tensor, device_mesh, [Shard(shard_dim)])
+            for i in range(3):
+                torch._dynamo.decorators.mark_unbacked(dt, i)
+            return dt
+
+        # aot_eager backend decomposes to as_strided, will fail as unsupported
+        @torch.compile(backend="eager", fullgraph=True)
+        def flatten_on_even_mesh(x):
+            torch._check(x.size(0) % self.world_size == 0)
+            return x.view(-1)
+
+        # view should be legal, since sharding is even and flatten is on first dim
+        dt = create_dt(0)
+        flatten_on_even_mesh(dt)
+
+        # flattening on non-first dimension should fail
+        dt = create_dt(1)
+        with self.assertRaisesRegex(
+            RuntimeError, "is not evenly divisible by mesh dimension"
+        ):
+            flatten_on_even_mesh(dt)
+
+        @torch.compile(backend="eager", fullgraph=True)
+        def flatten(x):
+            return x.view(-1)
+
+        # uneven case: not informing compiler of divisibility will crash
+        dt = create_dt(0)
+        with self.assertRaisesRegex(
+            RuntimeError, "is not evenly divisible by mesh dimension"
+        ):
+            flatten(dt)
+
     @with_comms
     def test_split_with_symint_split_size(self):
         """
@@ -1842,6 +2713,163 @@ def split_fn(x, split_size):
             for dt_chunk, tensor_chunk in zip(result, expected):
                 self.assertEqual(dt_chunk.full_tensor(), tensor_chunk)
 
+    @with_comms
+    def test_dtensor_processgroup_backward(self):
+        """Test that ProcessGroups are correctly handled in backward graph."""
+        from torch._functorch.aot_autograd import aot_function
+
+        with patch("torch.distributed.config.compile_on_one_rank", True):
+            mesh = self.build_device_mesh()
+
+            def fn(dt):
+                out = dt.redistribute(mesh, [Replicate()])
+                return out.sum()
+
+            local_tensor = torch.randn(
+                4, 8, device=self.device_type, requires_grad=True
+            )
+            dt_input = DTensor.from_local(
+                local_tensor, mesh, [Shard(0)], run_check=False
+            )
+
+            fw_graph_cell = [None]
+            bw_graph_cell = [None]
+
+            def extract_fw_graph(fx_g, _):
+                fw_graph_cell[0] = fx_g
+                return fx_g
+
+            def extract_bw_graph(fx_g, _):
+                bw_graph_cell[0] = fx_g
+                return fx_g
+
+            compiled_fn = aot_function(
+                fn,
+                fw_compiler=extract_fw_graph,
+                bw_compiler=extract_bw_graph,
+            )
+
+            output = compiled_fn(dt_input)
+            output.backward()
+
+            fw_graph = fw_graph_cell[0]
+            if fw_graph is not None:
+                fw_code = fw_graph.code
+                self.assertNotIn(
+                    "_opaque_obj",
+                    fw_code,
+                    f"Forward graph should not contain opaque objects. Graph:\n{fw_code}",
+                )
+
+            bw_graph = bw_graph_cell[0]
+            if bw_graph is not None:
+                bw_code = bw_graph.code
+                self.assertNotIn(
+                    "_opaque_obj",
+                    bw_code,
+                    f"Backward graph should not contain opaque objects. Graph:\n{bw_code}",
+                )
+
+    @with_comms
+    def test_dtensor_processgroup_extraction(self):
+        """Test that with compile_on_one_rank, DeviceMesh becomes a graph input
+        and ProcessGroups are extracted in-graph as placeholders."""
+        from torch._functorch.aot_autograd import aot_function
+
+        with patch("torch.distributed.config.compile_on_one_rank", True):
+            mesh = self.build_device_mesh()
+
+            def fn(dt):
+                return dt.redistribute(mesh, [Replicate()])
+
+            local_tensor = torch.randn(4, 8, device=self.device_type)
+            dt_input = DTensor.from_local(
+                local_tensor, mesh, [Shard(0)], run_check=False
+            )
+
+            fw_graph_cell = [None]
+
+            def extract_graph(fx_g, _):
+                fw_graph_cell[0] = fx_g
+                return fx_g
+
+            compiled_fn = aot_function(
+                fn,
+                fw_compiler=extract_graph,
+                bw_compiler=lambda fx_g, _: fx_g,
+            )
+
+            compiled_fn(dt_input)
+
+            fw_graph = fw_graph_cell[0]
+            self.assertIsNotNone(fw_graph)
+
+            graph_code = fw_graph.code
+
+            self.assertNotIn(
+                "_opaque_obj",
+                graph_code,
+                f"Graph should not contain opaque objects. Graph:\n{graph_code}",
+            )
+
+            placeholders = [n for n in fw_graph.graph.nodes if n.op == "placeholder"]
+            self.assertGreater(
+                len(placeholders),
+                1,
+                f"Expected ProcessGroup placeholders but only got {len(placeholders)} placeholder(s)",
+            )
+
+    @with_comms
+    def test_dtensor_processgroup_deduplication(self):
+        """Test that multiple DTensors sharing the same ProcessGroup only pass
+        it once as a graph placeholder."""
+        from torch._functorch.aot_autograd import aot_function
+
+        with patch("torch.distributed.config.compile_on_one_rank", True):
+            mesh = self.build_device_mesh()
+
+            def fn(dt1, dt2):
+                out1 = dt1.redistribute(mesh, [Replicate()])
+                out2 = dt2.redistribute(mesh, [Replicate()])
+                return out1 + out2
+
+            local1 = torch.randn(4, 8, device=self.device_type)
+            local2 = torch.randn(4, 8, device=self.device_type)
+            dt1 = DTensor.from_local(local1, mesh, [Shard(0)], run_check=False)
+            dt2 = DTensor.from_local(local2, mesh, [Shard(0)], run_check=False)
+
+            fw_graph_cell = [None]
+
+            def extract_graph(fx_g, _):
+                fw_graph_cell[0] = fx_g
+                return fx_g
+
+            compiled_fn = aot_function(
+                fn,
+                fw_compiler=extract_graph,
+                bw_compiler=lambda fx_g, _: fx_g,
+            )
+
+            compiled_fn(dt1, dt2)
+
+            fw_graph = fw_graph_cell[0]
+            self.assertIsNotNone(fw_graph)
+
+            graph_code = fw_graph.code
+
+            self.assertNotIn(
+                "_opaque_obj",
+                graph_code,
+                f"Graph should not contain opaque objects. Graph:\n{graph_code}",
+            )
+
+            placeholders = [n for n in fw_graph.graph.nodes if n.op == "placeholder"]
+            self.assertGreater(
+                len(placeholders),
+                2,
+                f"Expected ProcessGroup placeholders but only got {len(placeholders)} placeholder(s)",
+            )
+
 
 if __name__ == "__main__":
     run_tests()
diff --git a/test/distributed/tensor/test_dtensor_export.py b/test/distributed/tensor/test_dtensor_export.py
index dbf0e33184ac2..8255013db884e 100644
--- a/test/distributed/tensor/test_dtensor_export.py
+++ b/test/distributed/tensor/test_dtensor_export.py
@@ -124,20 +124,20 @@ def forward(self, x, *, block_mask=None):
 def strict_export_and_aot_export_joint_with_descriptors(model, args, kwargs=None):
     if kwargs is None:
         kwargs = {}
-    # needed for stric export
+    # needed for strict export; must be cleaned up to avoid polluting global state
     torch.utils._pytree.register_constant(DTensorSpec)
+    try:
+        # install_free_tensors is required for dynamo to work
+        with torch._dynamo.config.patch(install_free_tensors=True):
+            with torch._export.utils._disable_aten_to_metadata_assertions():
+                ep = torch.export.export(model, args, kwargs, strict=True)
 
-    # install_free_tensors is required for dynamo to work
-    with torch._dynamo.config.patch(
-        install_free_tensors=True, inline_inbuilt_nn_modules=True
-    ):
-        with torch._export.utils._disable_aten_to_metadata_assertions():
-            ep = torch.export.export(model, args, kwargs, strict=True)
-
-    # joint_gm produced here is missing the backward region, due to incompatiblility
-    # between ep.module() and aot_export_joint_with_descriptors.
-    # Keeping this here to show the issue.
-    return aot_export_joint_with_descriptors_alone(ep.module(), args, kwargs)
+        # joint_gm produced here is missing the backward region, due to incompatiblility
+        # between ep.module() and aot_export_joint_with_descriptors.
+        # Keeping this here to show the issue.
+        return aot_export_joint_with_descriptors_alone(ep.module(), args, kwargs)
+    finally:
+        torch.utils._pytree._deregister_pytree_node(DTensorSpec)
 
 
 def graph_capture_and_aot_export_joint_with_descriptors_v2(model, args, kwargs=None):
@@ -286,42 +286,46 @@ def unmarked_nodes(gm):
             ]
 
             marked_nodes_bw = [
+                "view_13",
                 "mm_8",
                 "t_17",
                 "view_1",
                 "mm_9",
                 "t_18",
                 "sum_5",
-                "view_11",
+                "view_14",
                 "t_19",
                 "detach",
                 "detach_3",
                 "threshold_backward_1",
+                "view_15",
                 "t_20",
                 "mm_10",
                 "t_21",
                 "sum_6",
-                "view_12",
+                "view_16",
                 "t_22",
             ]
             unmarked_nodes_bw = [
+                "view_7",
                 "mm_1",
                 "t_7",
                 "view_5",
                 "mm_3",
                 "t_8",
                 "sum_2",
-                "view_8",
+                "view_9",
                 "t_9",
                 "detach_1",
                 "detach_2",
                 "threshold_backward",
+                "view_10",
                 "mm_5",
                 "t_13",
                 "mm_7",
                 "t_14",
                 "sum_4",
-                "view_10",
+                "view_12",
                 "t_15",
                 "all_reduce_2",
                 "wait_tensor_2",
@@ -402,7 +406,8 @@ def test_dynamic_shapes(self, export_fn_with_answer):
         res = []
         for node in joint_gm.graph.nodes:
             if node.op == "placeholder":
-                assert "val" in node.meta
+                if "val" not in node.meta:
+                    raise AssertionError(f"Expected 'val' in node.meta for {node}")
                 fake_val = node.meta["val"]
                 if isinstance(fake_val, torch._subclasses.fake_tensor.FakeTensor):
                     res.append(list(fake_val.shape))
@@ -540,18 +545,7 @@ def forward(self, x):
     %item : [num_users=2] = call_method[target=item](args = (%clamp,), kwargs = {})
     %ge_1 : [num_users=1] = call_function[target=operator.ge](args = (%item, 1), kwargs = {})
     %_assert_scalar_default : [num_users=0] = call_function[target=torch.ops.aten._assert_scalar.default](args = (%ge_1, Runtime assertion failed for expression u0 >= 1 on node 'ge_1'), kwargs = {})
-    %getitem : [num_users=3] = call_function[target=operator.getitem](args = (%l_x_, slice(None, item, None)), kwargs = {})
-    %getattr_1 : [num_users=1] = call_function[target=builtins.getattr](args = (%getitem, _local_tensor), kwargs = {})
-    %sym_size_int : [num_users=2] = call_function[target=torch.ops.aten.sym_size.int](args = (%getattr_1, 0), kwargs = {})
-    %sym_size_int_1 : [num_users=2] = call_function[target=torch.ops.aten.sym_size.int](args = (%getitem, 0), kwargs = {})
-    %ge_2 : [num_users=1] = call_function[target=operator.ge](args = (%sym_size_int, 0), kwargs = {})
-    %_assert_scalar_default_1 : [num_users=0] = call_function[target=torch.ops.aten._assert_scalar.default](args = (%ge_2, Runtime assertion failed for expression u2 >= 0 on node 'ge_2'), kwargs = {})
-    %le : [num_users=1] = call_function[target=operator.le](args = (%sym_size_int, 4), kwargs = {})
-    %_assert_scalar_default_2 : [num_users=0] = call_function[target=torch.ops.aten._assert_scalar.default](args = (%le, Runtime assertion failed for expression u2 <= 4 on node 'le'), kwargs = {})
-    %ge_3 : [num_users=1] = call_function[target=operator.ge](args = (%sym_size_int_1, 0), kwargs = {})
-    %_assert_scalar_default_3 : [num_users=0] = call_function[target=torch.ops.aten._assert_scalar.default](args = (%ge_3, Runtime assertion failed for expression u1 >= 0 on node 'ge_3'), kwargs = {})
-    %le_1 : [num_users=1] = call_function[target=operator.le](args = (%sym_size_int_1, 4), kwargs = {})
-    %_assert_scalar_default_4 : [num_users=0] = call_function[target=torch.ops.aten._assert_scalar.default](args = (%le_1, Runtime assertion failed for expression u1 <= 4 on node 'le_1'), kwargs = {})
+    %getitem : [num_users=1] = call_function[target=operator.getitem](args = (%l_x_, slice(None, item, None)), kwargs = {})
     return (getitem,)""",  # noqa: B950
         )
 
diff --git a/test/distributed/tensor/test_dtensor_logging.py b/test/distributed/tensor/test_dtensor_logging.py
index 2b44d4f1cd2f8..72a886c0e1422 100644
--- a/test/distributed/tensor/test_dtensor_logging.py
+++ b/test/distributed/tensor/test_dtensor_logging.py
@@ -7,6 +7,7 @@
 from torch.distributed.tensor import DeviceMesh, DTensor, Replicate, Shard
 from torch.distributed.tensor._dtensor_spec import DTensorSpec, TensorMeta
 from torch.distributed.tensor._op_schema import OpSchema
+from torch.distributed.tensor.debug import _clear_sharding_prop_cache
 from torch.testing._internal.common_utils import requires_cuda, run_tests, TestCase
 from torch.testing._internal.distributed.fake_pg import FakeStore
 
@@ -21,6 +22,7 @@ def tearDown(self):
 
     def setUp(self):
         super().setUp()
+        _clear_sharding_prop_cache()
         self.world_size = 2
         store = FakeStore()
         dist.init_process_group(
@@ -91,6 +93,35 @@ def log_string():
 sharding_prop python cache HIT: aten.add.Tensor(Spec(f32[4, 4](S(0))), Spec(f32[4, 4](S(0)))) on DeviceMesh((2,), 'cuda', stride=(1,))) -> Spec(f32[4, 4](S(0)))""",  # noqa: B950
         )
 
+    def test_logging_level_change_resets_cpp_cache(self):
+        """setLevel on the dispatch logger resets the C++ cached logging flag."""
+        mesh = DeviceMesh(self.device_type, list(range(self.world_size)))
+        dispatch_logger = logging.getLogger("torch.distributed.tensor._dispatch")
+
+        log_records: list[logging.LogRecord] = []
+        handler = logging.Handler()
+        handler.emit = lambda record: log_records.append(record)
+        dispatch_logger.addHandler(handler)
+
+        x_dt = DTensor.from_local(torch.randn(2, 4), mesh, [Shard(0)], run_check=False)
+
+        # With logging off, the C++ hit path should produce no records.
+        dispatch_logger.setLevel(logging.WARNING)
+        x_dt + x_dt  # miss (not logged)
+        x_dt + x_dt  # hit (not logged)
+        self.assertEqual(len(log_records), 0)
+
+        # After enabling DEBUG, the C++ side should pick it up
+        # automatically (via the setLevel hook) without any manual reset.
+        _clear_sharding_prop_cache()
+        log_records.clear()
+        dispatch_logger.setLevel(logging.DEBUG)
+        x_dt + x_dt  # miss (logged)
+        x_dt + x_dt  # hit (logged)
+        self.assertEqual(len(log_records), 2)
+        self.assertIn("MISS", log_records[0].getMessage())
+        self.assertIn("HIT", log_records[1].getMessage())
+
 
 if __name__ == "__main__":
     run_tests()
diff --git a/test/distributed/tensor/test_dtensor_ops.py b/test/distributed/tensor/test_dtensor_ops.py
index b63004b1942e5..b354577c3d838 100644
--- a/test/distributed/tensor/test_dtensor_ops.py
+++ b/test/distributed/tensor/test_dtensor_ops.py
@@ -64,7 +64,8 @@ def skipOps(op_db, test_case_name, base_test_name, to_skip):
             for o in all_opinfos
             if o.name == op_name and o.variant_test_name == variant_name
         ]
-        assert len(matching_opinfos) >= 1, f"Couldn't find OpInfo for {xfail}"
+        if not (len(matching_opinfos) >= 1):
+            raise AssertionError(f"Couldn't find OpInfo for {xfail}")
         for opinfo in matching_opinfos:
             decorators = list(opinfo.decorators)
             if expected_failure:
@@ -126,76 +127,21 @@ def repurpose_ops(op_db, base_test_name, derived_test_name):
 # check_dtensor_func(self, test, op, dry_run=True), then run sth
 # like python test/distributed/tensor/test_dtensor_ops.py > failed.expect
 dtensor_fails = {
-    # these sometimes pass and sometimes fail
-    # we need to remove many of them from list once op
-    # get full support with varying sharding specs
-    xfail("__getitem__"),
-    xfail("_chunk_cat"),
+    # view/reshape ops: rejects flatten/split of sharded dims without redistribution
+    xfail("repeat_interleave"),
+    xfail("reshape"),
+    xfail("unbind"),
+    xfail("unflatten"),
+    xfail("view"),
+    # factory/creation ops: test harness can't convert non-tensor args to DTensor
     xfail("arange"),
-    xfail("as_strided"),
-    xfail("as_strided", "partial_views"),
-    xfail("bernoulli"),
-    xfail("block_diag"),
     xfail("broadcast_shapes"),
-    xfail("cartesian_prod"),
-    xfail("cholesky_solve"),
     xfail("eye"),
-    xfail("flatten"),
     xfail("full"),
-    xfail("index_put"),
-    xfail("kron"),
-    xfail("linalg.matrix_power"),
     xfail("linspace"),
-    xfail("linspace", "tensor_overload"),
-    xfail("logdet"),
     xfail("logspace"),
-    xfail("logspace", "tensor_overload"),
-    xfail("masked.logsumexp"),
-    xfail("masked.median"),
-    xfail("masked_scatter"),
-    xfail("masked_select"),
-    xfail("nanquantile"),
-    xfail("transpose"),
-    xfail("nn.functional.adaptive_avg_pool1d"),
-    xfail("nn.functional.adaptive_avg_pool3d"),
-    xfail("nn.functional.alpha_dropout"),
-    xfail("nn.functional.conv1d"),
-    xfail("nn.functional.conv2d"),
-    xfail("nn.functional.conv3d"),
-    xfail("nn.functional.conv_transpose1d"),
-    xfail("nn.functional.conv_transpose2d"),
-    xfail("nn.functional.conv_transpose3d"),
-    xfail("nn.functional.dropout"),
-    xfail("nn.functional.fractional_max_pool2d"),
-    xfail("nn.functional.fractional_max_pool3d"),
-    xfail("nn.functional.interpolate", "area"),
-    xfail("nn.functional.interpolate", "nearest"),
-    xfail("nn.functional.interpolate", "nearest-exact"),
-    xfail("nn.functional.pad", "reflect"),
-    xfail("nn.functional.pad", "replicate"),
-    xfail("nn.functional.pad", "replicate_negative"),
-    xfail("nn.functional.unfold"),
-    xfail("normal"),
-    xfail("normal", "number_mean"),
-    xfail("normal", "in_place"),
     xfail("ones"),
-    xfail("pca_lowrank"),
-    xfail("quantile"),
-    xfail("rand_like"),
-    xfail("randint_like"),
-    xfail("randint"),
-    xfail("randn"),
-    xfail("randn_like"),
-    xfail("ravel"),
-    xfail("repeat_interleave"),
-    xfail("reshape"),
-    xfail("reshape_as"),
-    xfail("roll"),
     xfail("scalar_tensor"),
-    xfail("searchsorted"),
-    xfail("sparse.sampled_addmm"),
-    xfail("sparse.mm", "reduce"),
-    xfail("squeeze", "multiple"),
     xfail("signal.windows.bartlett"),
     xfail("signal.windows.blackman"),
     xfail("signal.windows.cosine"),
@@ -205,20 +151,58 @@ def repurpose_ops(op_db, base_test_name, derived_test_name):
     xfail("signal.windows.general_hamming"),
     xfail("signal.windows.hamming"),
     xfail("signal.windows.hann"),
-    xfail("signal.windows.nuttall"),
     xfail("signal.windows.kaiser"),
-    xfail("stack"),
+    xfail("signal.windows.nuttall"),
+    xfail("zeros"),
+    # random/stochastic ops: different RNG states between DTensor and reference
+    xfail("bernoulli"),
+    xfail("cauchy"),
+    xfail("exponential"),
+    xfail("nn.functional.alpha_dropout"),
+    xfail("nn.functional.dropout"),
+    xfail("normal"),
+    xfail("normal", "in_place"),
+    xfail("normal", "number_mean"),
+    xfail("rand_like"),
+    xfail("randint"),
+    xfail("randint_like"),
+    xfail("randn"),
+    xfail("randn_like"),
+    xfail("uniform"),
+    # mixed Tensor/DTensor inputs: op creates plain Tensors mixed with DTensor args
+    xfail("__getitem__"),
+    xfail("nn.functional.fractional_max_pool2d"),
+    xfail("nn.functional.fractional_max_pool3d"),
+    xfail("pca_lowrank"),
+    xfail("quantile"),
     xfail("svd_lowrank"),
-    xfail("take_along_dim"),
+    # dynamic output shape: output shape depends on data values
+    xfail("combinations"),
+    xfail("linalg.lstsq"),
+    xfail("linalg.lstsq", "grad_oriented"),
+    xfail("masked_select"),
+    xfail("nn.functional.ctc_loss"),
+    # 0-dim tensor edge cases: strategies don't handle scalar tensors
+    xfail("logsumexp"),
+    xfail("masked.logsumexp"),
+    xfail("transpose"),
+    # conv stride+padding: TP convolution rejects stride != 1 with padding
+    xfail("nn.functional.conv1d"),
+    xfail("nn.functional.conv2d"),
+    xfail("nn.functional.conv3d"),
+    xfail("nn.functional.conv_transpose1d"),
+    xfail("nn.functional.conv_transpose2d"),
+    xfail("nn.functional.conv_transpose3d"),
+    # "cannot resize variables that require grad" from test harness
+    xfail("resize_"),
+    xfail("resize_as_"),
+    # DTensorConverter can't convert sparse tensor inputs
+    xfail("sparse.sampled_addmm"),
+    xfail("sparse.mm", "reduce"),
+    # meta tensor data not allocated yet during tensor_split
     xfail("tensor_split"),
-    xfail("unbind"),
-    xfail("uniform"),
-    xfail("unflatten"),
-    xfail("_unsafe_masked_index"),
-    xfail("_unsafe_masked_index_put_accumulate"),
-    xfail("view"),
-    xfail("view_as"),
-    xfail("zeros"),
+    # output_specs count mismatch in unsafe_split strategy
+    xfail("unsafe_split"),
     # /TODO(whc) debug/triage
     # ops inside this might even fail without dtensor
     # tests, as we rescale op db common test size factor (i.e. L, M, S)
@@ -244,6 +228,7 @@ def repurpose_ops(op_db, base_test_name, derived_test_name):
     skip("fft.hfftn"),
     skip("fft.ifftn"),
     skip("fft.irfft"),
+    skip("frexp"),
     skip("istft"),
     skip("isclose"),
     skip("isreal"),
@@ -256,7 +241,6 @@ def repurpose_ops(op_db, base_test_name, derived_test_name):
     skip("_segment_reduce", "lengths"),
     skip("_segment_reduce", "offsets"),
     # TODO: fix the following ops
-    skip("squeeze"),
     skip("empty"),
     skip("empty_strided"),
     skip("empty_like"),
@@ -266,210 +250,199 @@ def repurpose_ops(op_db, base_test_name, derived_test_name):
 }
 
 dtensor_multi_threaded_fails = {
+    xfail("index_fill"),
     xfail("full_like"),
     xfail("nn.functional.dropout2d"),
     xfail("nn.functional.dropout3d"),
     xfail("nn.functional.huber_loss"),
-    xfail("masked.cumprod"),
+    skip("nn.functional.max_unpool1d", "grad"),
+    skip("nn.functional.max_unpool2d", "grad"),
+    xfail("nn.functional.max_unpool3d", "grad"),
+    xfail("nn.functional.threshold"),
     skip("nn.functional.multi_head_attention_forward"),
 }
 
+# Ops that fail to compile with DTensor + torch.compile(fullgraph=True).
+# These are compile-time failures, NOT numeric correctness issues.
+dtensor_compiled_fails = {
+    xfail("cartesian_prod"),
+    xfail("flatten"),
+    xfail("kron"),
+    xfail("linalg.tensorsolve"),
+    xfail("nn.functional.instance_norm"),
+    xfail("ravel"),
+    xfail("reshape_as"),
+    xfail("take_along_dim"),
+    xfail("view_as"),
+    # View-type ops that decompose into as_strided (at autograd level).
+    # DTensor doesn't have a sharding strategy for as_strided.
+    xfail("atleast_1d"),
+    xfail("atleast_2d"),
+    xfail("atleast_3d"),
+    xfail("broadcast_tensors"),
+    xfail("broadcast_to"),
+    xfail("diagonal"),
+    xfail("dsplit"),
+    xfail("expand"),
+    xfail("expand_as"),
+    xfail("hsplit"),
+    xfail("linalg.diagonal"),
+    xfail("movedim"),
+    xfail("narrow"),
+    xfail("permute"),
+    xfail("select"),
+    xfail("slice"),
+    xfail("squeeze"),
+    xfail("squeeze", "multiple"),
+    xfail("t"),
+    xfail("transpose_copy"),
+    xfail("unsqueeze"),
+    xfail("vsplit"),
+    # Decompositions that use plain tensor constructors (e.g. arange),
+    # causing mixed tensor/DTensor errors during Dynamo's fake prop.
+    xfail("corrcoef"),
+    xfail("cov"),
+    xfail("nn.functional.interpolate", "bicubic"),
+    xfail("nn.functional.interpolate", "bilinear"),
+    xfail("nn.functional.interpolate", "linear"),
+    xfail("nn.functional.interpolate", "nearest"),
+    xfail("nn.functional.interpolate", "nearest-exact"),
+    xfail("nn.functional.interpolate", "trilinear"),
+    xfail("nn.functional.max_unpool1d"),
+    xfail("nn.functional.max_unpool2d"),
+    xfail("nn.functional.max_unpool3d"),
+    xfail("nn.functional.upsample_bilinear"),
+    xfail("nn.functional.upsample_nearest"),
+    # Data-dependent outputs (SymBool, unbacked shapes) that raise
+    # during DTensor's fake prop.
+    xfail("equal"),
+    xfail("hash_tensor"),
+    xfail("item"),
+    xfail("nonzero_static"),
+    # Decompositions with .is_cuda checks that fail during sharding
+    # propagation for aten.is_cuda / prim::device.
+    xfail("nn.functional.binary_cross_entropy"),
+    xfail("nn.functional.binary_cross_entropy_with_logits"),
+    xfail("nn.functional.gaussian_nll_loss"),
+    xfail("nn.functional.logsigmoid"),
+    # Miscellaneous runtime crashes (e.g. index out of bounds).
+    xfail("gather"),
+    xfail("index_add"),
+    xfail("index_copy"),
+    xfail("index_reduce", "amax"),
+    xfail("index_reduce", "amin"),
+    xfail("index_reduce", "mean"),
+    xfail("index_reduce", "prod"),
+    xfail("index_select"),
+    xfail("lu_unpack"),
+    xfail("scatter"),
+    xfail("scatter_add"),
+    xfail("take_along_dim"),
+    # batch_norm variants decompose through squeeze.dims → as_strided under
+    # compilation, and DTensor has no as_strided strategy.
+    xfail("_native_batch_norm_legit"),
+    xfail("native_batch_norm"),
+    xfail("nn.functional.batch_norm"),
+    # False positives: these have no sharding strategy and their
+    # eager DTensor failure is registered elsewhere.
+    xfail("nn.functional.multilabel_soft_margin_loss"),
+}
+
+# Ops that compile successfully but fail numeric checks in eager DTensor tests.
+# These are excluded from TestCompiledDTensorOps skip list since we don't check numerics.
+dtensor_numeric_only_fails = {
+    xfail("arange"),
+    xfail("broadcast_shapes"),
+    xfail("eye"),
+    xfail("full"),
+    xfail("full_like"),
+    xfail("linspace"),
+    xfail("logspace"),
+    xfail("nn.functional.huber_loss"),
+    xfail("nn.functional.smooth_l1_loss"),
+    xfail("nn.functional.softshrink"),
+    xfail("nn.functional.threshold"),
+    xfail("ones"),
+    xfail("randint"),
+    xfail("randn"),
+    xfail("scalar_tensor"),
+    xfail("signal.windows.bartlett"),
+    xfail("signal.windows.blackman"),
+    xfail("signal.windows.cosine"),
+    xfail("signal.windows.exponential"),
+    xfail("signal.windows.gaussian"),
+    xfail("signal.windows.general_cosine"),
+    xfail("signal.windows.general_hamming"),
+    xfail("signal.windows.hamming"),
+    xfail("signal.windows.hann"),
+    xfail("signal.windows.kaiser"),
+    xfail("signal.windows.nuttall"),
+    xfail("sparse.mm", "reduce"),
+    xfail("sparse.sampled_addmm"),
+    xfail("squeeze_copy"),
+    xfail("stack"),
+    xfail("unsafe_chunk"),
+    xfail("zeros"),
+}
+
 # Ops in dtensor_fails that have no sharding strategy (NotImplementedError).
 # These will error during sharding propagation and affect unbacked tests too.
 dtensor_fails_no_strategy = {
     xfail("_batch_norm_with_update"),
-    xfail("_native_batch_norm_legit"),
-    xfail("_upsample_bilinear2d_aa"),
+    xfail("_chunk_cat"),
+    xfail("_unsafe_masked_index"),
+    xfail("_unsafe_masked_index_put_accumulate"),
     xfail("addbmm"),
-    xfail("addr"),
     xfail("allclose"),
+    xfail("as_strided"),
+    xfail("as_strided", "partial_views"),
     xfail("as_strided_copy"),
     xfail("as_strided_scatter"),
-    xfail("cauchy"),
+    xfail("block_diag"),
     xfail("cdist"),
-    xfail("cholesky"),
-    xfail("cholesky_inverse"),
-    xfail("combinations"),
     xfail("complex"),
-    xfail("cross"),
-    xfail("cummax"),
-    xfail("cummin"),
-    xfail("diagonal_scatter"),
     xfail("exponential"),
-    xfail("fft.fft"),
-    xfail("fft.fft2"),
-    xfail("fft.fftn"),
-    xfail("fft.fftshift"),
-    xfail("fft.ifft"),
-    xfail("fft.ifft2"),
-    xfail("fft.ifftshift"),
-    xfail("fft.ihfft"),
     xfail("fft.ihfft2"),
     xfail("fft.ihfftn"),
-    xfail("fft.irfft2"),
-    xfail("fft.irfftn"),
-    xfail("fft.rfft"),
-    xfail("fft.rfft2"),
-    xfail("fft.rfftn"),
-    xfail("flip"),
-    xfail("fliplr"),
-    xfail("flipud"),
-    xfail("frexp"),
     xfail("geometric"),
-    xfail("geqrf"),
-    xfail("grid_sampler_2d"),
     xfail("histogram"),
     xfail("histogramdd"),
-    xfail("index_add"),
-    xfail("index_copy"),
-    xfail("index_fill"),
-    xfail("index_reduce", "prod"),
-    xfail("index_reduce", "mean"),
-    xfail("index_reduce", "amax"),
-    xfail("index_reduce", "amin"),
-    xfail("index_select"),
     xfail("isin"),
-    xfail("kthvalue"),
-    xfail("linalg.cholesky"),
-    xfail("linalg.cholesky_ex"),
-    xfail("linalg.cross"),
-    xfail("linalg.det"),
-    xfail("linalg.eig"),
-    xfail("linalg.eigvals"),
-    xfail("linalg.householder_product"),
-    xfail("linalg.inv"),
-    xfail("linalg.inv_ex"),
-    xfail("linalg.ldl_factor"),
-    xfail("linalg.ldl_factor_ex"),
-    xfail("linalg.ldl_solve"),
-    xfail("linalg.lstsq"),
-    xfail("linalg.lstsq", "grad_oriented"),
-    xfail("linalg.lu"),
-    xfail("linalg.lu_factor"),
-    xfail("linalg.lu_factor_ex"),
-    xfail("linalg.lu_solve"),
-    xfail("linalg.pinv"),
-    xfail("linalg.pinv", "hermitian"),
-    xfail("linalg.slogdet"),
-    xfail("linalg.solve"),
-    xfail("linalg.solve_ex"),
-    xfail("linalg.solve_triangular"),
-    xfail("linalg.tensorinv"),
-    xfail("linalg.tensorsolve"),
-    xfail("linalg.vander"),
+    xfail("linalg.matrix_power"),
+    xfail("linspace", "tensor_overload"),
     xfail("log_normal"),
-    xfail("logcumsumexp"),
-    xfail("logsumexp"),
-    xfail("lu"),
-    xfail("lu_solve"),
-    xfail("lu_unpack"),
-    xfail("matrix_exp"),
-    xfail("max_pool2d_with_indices_backward"),
-    xfail("median"),
-    xfail("mode"),
+    xfail("logspace", "tensor_overload"),
+    xfail("masked_scatter"),
     xfail("multinomial"),
-    xfail("nanmean"),
-    xfail("nanmedian"),
-    xfail("nansum"),
-    xfail("narrow_copy"),
-    xfail("native_batch_norm"),
-    xfail("ne"),
-    xfail("nn.functional.adaptive_avg_pool2d"),
-    xfail("nn.functional.adaptive_max_pool1d"),
-    xfail("nn.functional.adaptive_max_pool2d"),
-    xfail("nn.functional.adaptive_max_pool3d"),
-    xfail("nn.functional.avg_pool1d"),
-    xfail("nn.functional.avg_pool2d"),
-    xfail("nn.functional.avg_pool3d"),
-    xfail("nn.functional.batch_norm"),
-    xfail("nn.functional.batch_norm", "without_cudnn"),
+    xfail("nanquantile"),
     xfail("nn.functional.bilinear"),
-    xfail("nn.functional.binary_cross_entropy_with_logits"),
-    xfail("nn.functional.cosine_similarity"),
-    xfail("nn.functional.ctc_loss"),
-    xfail("nn.functional.glu"),
-    xfail("nn.functional.grid_sample"),
     xfail("nn.functional.group_norm"),
-    xfail("nn.functional.hardshrink"),
-    xfail("nn.functional.instance_norm"),
-    xfail("nn.functional.local_response_norm"),
-    xfail("nn.functional.max_pool1d"),
-    xfail("nn.functional.max_pool2d"),
-    xfail("nn.functional.max_pool3d"),
-    xfail("nn.functional.max_unpool1d"),
-    xfail("nn.functional.max_unpool1d", "grad"),
-    xfail("nn.functional.max_unpool2d"),
-    xfail("nn.functional.max_unpool2d", "grad"),
-    xfail("nn.functional.max_unpool3d"),
-    xfail("nn.functional.max_unpool3d", "grad"),
     xfail("nn.functional.multi_margin_loss"),
     xfail("nn.functional.multilabel_margin_loss"),
+    xfail("nn.functional.pad", "reflect"),
+    xfail("nn.functional.pad", "replicate"),
+    xfail("nn.functional.pad", "replicate_negative"),
     xfail("nn.functional.pdist"),
     xfail("nn.functional.rrelu"),
-    xfail("nn.functional.threshold"),
-    xfail("nn.functional.upsample_nearest"),
+    xfail("nn.functional.unfold"),
     xfail("nonzero"),
-    xfail("ormqr"),
-    xfail("pinverse"),
     xfail("polar"),
     xfail("put"),
     xfail("renorm"),
-    xfail("resize_"),
-    xfail("resize_as_"),
-    xfail("rot90"),
     xfail("scatter_reduce", "amax"),
     xfail("scatter_reduce", "amin"),
     xfail("scatter_reduce", "mean"),
     xfail("scatter_reduce", "prod"),
     xfail("scatter_reduce", "sum"),
-    xfail("select_scatter"),
-    xfail("special.airy_ai"),
-    xfail("special.bessel_j0"),
-    xfail("special.bessel_j1"),
-    xfail("special.bessel_y0"),
-    xfail("special.bessel_y1"),
-    xfail("special.chebyshev_polynomial_t"),
-    xfail("special.chebyshev_polynomial_u"),
-    xfail("special.chebyshev_polynomial_v"),
-    xfail("special.chebyshev_polynomial_w"),
-    xfail("special.entr"),
-    xfail("special.erfcx"),
-    xfail("special.hermite_polynomial_h"),
-    xfail("special.hermite_polynomial_he"),
-    xfail("special.i0e"),
-    xfail("special.i1"),
-    xfail("special.i1e"),
-    xfail("special.laguerre_polynomial_l"),
-    xfail("special.legendre_polynomial_p"),
-    xfail("special.log_ndtr"),
-    xfail("special.modified_bessel_i0"),
-    xfail("special.modified_bessel_i1"),
-    xfail("special.modified_bessel_k0"),
-    xfail("special.modified_bessel_k1"),
-    xfail("special.ndtri"),
-    xfail("special.scaled_modified_bessel_k0"),
-    xfail("special.scaled_modified_bessel_k1"),
-    xfail("special.shifted_chebyshev_polynomial_t"),
-    xfail("special.shifted_chebyshev_polynomial_u"),
-    xfail("special.shifted_chebyshev_polynomial_v"),
-    xfail("special.shifted_chebyshev_polynomial_w"),
-    xfail("special.spherical_bessel_j0"),
-    xfail("special.xlog1py"),
-    xfail("special.zeta"),
+    xfail("searchsorted"),
     xfail("squeeze_copy"),
-    xfail("std_mean"),
-    xfail("std_mean", "unbiased"),
     xfail("stft"),
     xfail("take"),
     xfail("to_sparse"),
-    xfail("triangular_solve"),
     xfail("unfold"),
     xfail("unfold_copy"),
     xfail("unique"),
     xfail("unique_consecutive"),
-    xfail("unsafe_split"),
-    xfail("var_mean"),
-    xfail("var_mean", "unbiased"),
 }
 
 # Add a list of ops that are currently failing BW pass
@@ -666,7 +639,8 @@ def check_dtensor_func(self, test_func, opinfo, dry_run=False):
 
     def run_one_hot(self):
         ops = [op for op in op_db if op.name == "nn.functional.one_hot"]
-        assert len(ops) == 1
+        if len(ops) != 1:
+            raise AssertionError(f"Expected 1 op, got {len(ops)}")
         op = ops[0]
         # num_classes = -1 appears to have a bug with dtensor.max().item()
         self.run_opinfo_test(
@@ -800,181 +774,115 @@ def assertEqualOnRank(self, x, y, msg=None, *, rank=0):
 # This list only contains ops NOT in ops_dde_xfail - those are base tensor issues.
 ops_unbacked_dtensor_dde = {
     xfail("__getitem__"),
-    xfail("__radd__"),
-    xfail("__rdiv__"),
     xfail("__rmatmul__"),
-    xfail("__rmod__"),
-    xfail("__rmul__"),
-    xfail("__rpow__"),
-    xfail("__rsub__"),
     xfail("_segment_reduce", "lengths"),
     xfail("_segment_reduce", "offsets"),
-    xfail("_softmax_backward_data"),
+    xfail("_native_batch_norm_legit"),
     xfail("_unsafe_masked_index"),
-    xfail("add"),
-    xfail("addcdiv"),
-    xfail("addcmul"),
     xfail("addmm"),
     xfail("addmm", "decomposed"),
     xfail("addr"),
     xfail("alias_copy"),
     xfail("aminmax"),
-    xfail("argmax"),
-    xfail("argmin"),
     xfail("argsort"),
     xfail("argwhere"),
     xfail("as_strided"),
     xfail("as_strided", "partial_views"),
-    xfail("atan2"),
+    xfail("as_strided_copy"),
     xfail("block_diag"),
-    xfail("bmm"),
-    xfail("broadcast_tensors"),
+    xfail("bucketize"),
     skip("broadcast_to"),
     xfail("bucketize"),
     xfail("cartesian_prod"),
-    xfail("cat"),
-    xfail("cholesky_solve"),
-    xfail("clamp"),
-    xfail("clamp_max"),
-    xfail("clamp_min"),
-    xfail("column_stack"),
     xfail("constant_pad_nd"),
-    xfail("copysign"),
     xfail("cumprod"),
-    xfail("cumsum"),
-    xfail("diag"),
-    xfail("diagflat"),
-    xfail("div", "floor_rounding"),
-    xfail("div", "no_rounding_mode"),
-    xfail("div", "trunc_rounding"),
-    xfail("dot"),
-    xfail("dstack"),
-    xfail("einsum"),
-    xfail("eq"),
-    xfail("expand"),
-    xfail("expand_as"),
-    xfail("expand_copy"),
+    xfail("dist"),
     xfail("fill"),
     xfail("flatten"),
-    xfail("float_power"),
-    xfail("floor_divide"),
-    xfail("fmax"),
-    xfail("fmin"),
-    xfail("fmod"),
+    xfail("flip"),
+    xfail("fliplr"),
+    xfail("flipud"),
+    xfail("float"),
     xfail("gather"),
-    xfail("ge"),
-    xfail("gt"),
-    xfail("heaviside"),
     xfail("histc"),
-    xfail("hstack"),
-    xfail("hypot"),
-    xfail("igamma"),
-    xfail("igammac"),
+    xfail("index_add"),
+    xfail("index_copy"),
     xfail("index_put"),
+    xfail("index_reduce", "amax"),
+    xfail("index_reduce", "amin"),
+    xfail("index_reduce", "mean"),
+    xfail("index_reduce", "prod"),
     xfail("index_select"),
-    xfail("isclose"),
-    xfail("ldexp"),
-    xfail("le"),
-    xfail("lerp"),
-    xfail("logaddexp"),
-    xfail("logaddexp2"),
-    xfail("logical_and"),
-    xfail("logical_or"),
-    xfail("logical_xor"),
-    xfail("lt"),
-    xfail("masked.normalize"),
+    xfail("kthvalue"),
+    xfail("linalg.lu"),
+    xfail("linalg.lu_factor"),
+    xfail("linalg.lu_factor_ex"),
+    xfail("lu"),
     xfail("masked_fill"),
     xfail("masked_scatter"),
     xfail("masked_select"),
     xfail("matmul"),
-    xfail("max", "binary"),
+    xfail("matrix_exp"),
     xfail("max", "reduction_with_dim"),
-    xfail("maximum"),
-    xfail("meshgrid", "list_of_tensors"),
-    xfail("meshgrid", "variadic_tensors"),
-    xfail("min", "binary"),
     xfail("min", "reduction_with_dim"),
-    xfail("minimum"),
-    xfail("mm"),
+    xfail("mode"),
     xfail("msort"),
-    xfail("mul"),
     xfail("mv"),
     xfail("narrow"),
     xfail("narrow_copy"),
-    xfail("native_dropout_backward"),
     xfail("new_empty"),
     xfail("new_empty_strided"),
     xfail("new_full"),
     xfail("new_ones"),
+    xfail("native_batch_norm"),
     xfail("new_zeros"),
-    xfail("nextafter"),
-    xfail("nn.functional.celu"),
+    xfail("nn.functional.batch_norm"),
+    xfail("nn.functional.conv1d"),
+    xfail("nn.functional.conv2d"),
+    xfail("nn.functional.conv3d"),
+    xfail("nn.functional.conv_transpose1d"),
+    xfail("nn.functional.conv_transpose2d"),
+    xfail("nn.functional.conv_transpose3d"),
     xfail("nn.functional.cosine_embedding_loss"),
-    xfail("nn.functional.elu"),
-    xfail("nn.functional.hardsigmoid"),
-    xfail("nn.functional.hardtanh"),
+    xfail("nn.functional.glu"),
     xfail("nn.functional.hinge_embedding_loss"),
     xfail("nn.functional.interpolate", "nearest"),
     xfail("nn.functional.interpolate", "nearest-exact"),
     xfail("nn.functional.linear"),
     xfail("nn.functional.logsigmoid"),
     xfail("nn.functional.margin_ranking_loss"),
-    xfail("nn.functional.mish"),
     xfail("nn.functional.multilabel_soft_margin_loss"),
-    xfail("nn.functional.normalize"),
     xfail("nn.functional.pad", "constant"),
-    xfail("nn.functional.pixel_unshuffle"),
     xfail("nn.functional.poisson_nll_loss"),
-    xfail("nn.functional.relu6"),
-    xfail("nn.functional.selu"),
-    xfail("nn.functional.softplus"),
     xfail("nn.functional.soft_margin_loss"),
     xfail("nn.functional.triplet_margin_loss"),
     xfail("nn.functional.triplet_margin_with_distance_loss"),
+    xfail("nn.functional.upsample_nearest"),
+    xfail("nn.functional.pixel_unshuffle"),
     xfail("nonzero_static"),
-    xfail("outer"),
     xfail("permute_copy"),
-    xfail("pow"),
     xfail("prod"),
+    xfail("quantile"),
     xfail("ravel"),
-    xfail("remainder"),
     xfail("reshape"),
     xfail("reshape_as"),
     xfail("rsub"),
+    xfail("rot90"),
     xfail("scatter"),
     xfail("scatter_add"),
-    xfail("select"),
-    xfail("slice"),
-    xfail("slice_scatter"),
+    xfail("select_scatter"),
     xfail("sort"),
-    xfail("split", "list_args"),
-    xfail("split_with_sizes"),
-    xfail("split_with_sizes_copy"),
-    xfail("squeeze"),
-    xfail("squeeze", "multiple"),
-    xfail("stack"),
-    xfail("std"),
-    xfail("std", "unbiased"),
-    xfail("sub"),
+    xfail("squeeze_copy"),
+    xfail("std_mean"),
     xfail("topk"),
-    xfail("torch.ops.aten._safe_softmax.default"),
-    xfail("trace"),
     xfail("transpose_copy"),
-    xfail("tril"),
-    xfail("triu"),
-    xfail("true_divide"),
     xfail("unflatten"),
     xfail("unsqueeze_copy"),
-    xfail("var"),
-    xfail("var", "unbiased"),
     xfail("vdot"),
     xfail("view"),
     xfail("view_as"),
     xfail("view_as_complex"),
     xfail("view_copy"),
-    xfail("vstack"),
-    xfail("where"),
-    xfail("xlogy"),
 }
 
 
@@ -1120,6 +1028,16 @@ def _extract_aten_op_and_args(self, torch_op, args, kwargs):
 
     @suppress_warnings
     @ops(op_db, allowed_dtypes=(torch.float,))
+    @skipOps(
+        op_db,
+        "TestSingleDimStrategies",
+        "test_single_dim_strategy",
+        {
+            # Stochastic: each shard gets independent RNG, so
+            # op(full) != cat(op(shard0), op(shard1)).
+            skip("exponential"),
+        },
+    )
     def test_single_dim_strategy(self, dtype, op):
         torch.manual_seed(42)
         mesh = init_device_mesh(DEVICE_TYPE, (self.world_size,))
@@ -1175,8 +1093,11 @@ def _collect_tensor_meta(dt):
             lambda s: Shard(s.dim),
             single_dim_strats[aten_op](aten_op, args_meta, kwargs_meta),
         )
-        # TODO(pianpwk): handle multi-output once that lands for single-dim
-        for output_placement, *input_placements in strategies:
+        n_inputs = len(all_tensor_meta)
+        for strategy in strategies:
+            input_placements = strategy[-n_inputs:]
+            output_placements = strategy[:-n_inputs]
+
             # skip strategies with invalid shards
             def is_invalid_shard(meta, p):
                 ndim = len(meta.shape)
@@ -1196,20 +1117,84 @@ def is_invalid_shard(meta, p):
             ):
                 continue
 
-            # add the validate_sharding_rule function
             self.assertTrue(
                 validate_sharding_rule_sample(
                     aten_op,
                     full_args,
                     full_kwargs,
                     input_placements,
-                    (output_placement,),
+                    tuple(output_placements),
                     mesh,
                 ),
-                f"{op.name}: {input_placements} -> {(output_placement,)} failed",
+                f"{op.name}: {input_placements} -> {tuple(output_placements)} failed",
             )
 
 
+class TestCompiledDTensorOps(TestDTensorOps):
+    """
+    Test DTensor ops compile successfully with aot_eager backend.
+    Uses fake PG for speed - focuses on compilation, not output correctness.
+    """
+
+    _op_db = repurpose_ops(op_db, "TestDTensorOps", "TestCompiledDTensorOps")
+
+    def setUp(self) -> None:
+        super().setUp()
+        torch.distributed.init_process_group("fake", rank=0, world_size=self.world_size)
+
+    def tearDown(self):
+        super().tearDown()
+        from torch.distributed.tensor.debug import _clear_sharding_prop_cache
+
+        _clear_sharding_prop_cache()
+        torch._dynamo.reset()
+        try:
+            dist.destroy_process_group()
+        except AssertionError:
+            pass
+
+    def assertEqualOnRank(self, x, y, msg=None, *, rank=0):
+        # Skip output comparison - we only care that compilation succeeds
+        pass
+
+    def run_dtensor_crossref(self, func, args, kwargs):
+        """
+        Override to compile with aot_eager and verify compilation succeeds.
+        Does not check output correctness.
+        """
+        to_dtensor = DTensorConverter(self.mesh, args, kwargs)
+
+        for dtensor_args, dtensor_kwargs in to_dtensor:
+            if not to_dtensor.successful():
+                continue
+
+            torch._dynamo.reset()
+
+            @torch.compile(backend="aot_eager", fullgraph=True)
+            def compiled_func(*a, **kw):
+                return func(*a, **kw)
+
+            # Just run - if it compiles and runs without error, we pass
+            compiled_func(*dtensor_args, **dtensor_kwargs)
+
+    @suppress_warnings
+    @ops(_op_db, allowed_dtypes=(torch.float,))
+    @skipOps(
+        _op_db,
+        "TestCompiledDTensorOps",
+        "test_compiled_dtensor_op_db",
+        (
+            dtensor_fails
+            | dtensor_fails_no_strategy
+            | dtensor_multi_threaded_fails
+            | dtensor_compiled_fails
+        )
+        - dtensor_numeric_only_fails,
+    )
+    def test_compiled_dtensor_op_db(self, dtype, op):
+        self.run_opinfo_test(dtype, op, requires_grad=False)
+
+
 # only instantiate tests for DEVICE_TYPE alone (i.e. either CPU or GPU)
 instantiate_device_type_tests(
     TestMultiThreadedDTensorOps, globals(), only_for=(DEVICE_TYPE,)
@@ -1225,5 +1210,9 @@ def is_invalid_shard(meta, p):
     TestSingleDimStrategies, globals(), only_for=(DEVICE_TYPE,)
 )
 
+instantiate_device_type_tests(
+    TestCompiledDTensorOps, globals(), only_for=(DEVICE_TYPE,)
+)
+
 if __name__ == "__main__":
     run_tests()
diff --git a/test/distributed/tensor/test_dynamic.py b/test/distributed/tensor/test_dynamic.py
index f058ae8cc821d..ac1432e7a1c09 100644
--- a/test/distributed/tensor/test_dynamic.py
+++ b/test/distributed/tensor/test_dynamic.py
@@ -63,6 +63,12 @@ def forward(input_batch_inputs_):
 
 TestDynamicWithLocalTensor = create_local_tensor_test_class(
     TestDynamic,
+    # LocalTensorMode is a non-infra dispatch mode that causes Dynamo to skip
+    # frames, which is incompatible with fullgraph=True.
+    skipped_tests=[
+        "test_embedding_fake_tensor_cache_enabled_False",
+        "test_embedding_fake_tensor_cache_enabled_True",
+    ],
 )
 
 if __name__ == "__main__":
diff --git a/test/distributed/tensor/test_math_ops.py b/test/distributed/tensor/test_math_ops.py
index 2eeaf3120c687..e449b95847c5c 100644
--- a/test/distributed/tensor/test_math_ops.py
+++ b/test/distributed/tensor/test_math_ops.py
@@ -27,7 +27,7 @@
     SequenceParallel,
 )
 from torch.testing._internal.common_distributed import skip_if_lt_x_gpu
-from torch.testing._internal.common_utils import run_tests
+from torch.testing._internal.common_utils import run_tests, TEST_WITH_ROCM
 from torch.testing._internal.distributed._tensor.common_dtensor import (
     create_local_tensor_test_class,
     DTensorTestBase,
@@ -90,9 +90,46 @@ def linear_op_reductions(self, op_str):
 
     @with_comms
     def test_linear_op_reductions(self):
-        for op_str in ("all", "sum", "prod", "max", "min", "any", "amax", "amin"):
+        for op_str in (
+            "all",
+            "sum",
+            "prod",
+            "max",
+            "min",
+            "any",
+            "amax",
+            "amin",
+            "var",
+            "std",
+            "nansum",
+        ):
             self.linear_op_reductions(op_str)
 
+    @with_comms
+    def test_nansum_with_nan(self):
+        device_mesh = self.build_device_mesh()
+        # Tensor with NaN values sharded across the reduction dim
+        tensor = torch.tensor(
+            [[1.0, float("nan"), 3.0], [float("nan"), 5.0, 6.0]],
+            device=self.device_type,
+        )
+        dtensor = distribute_tensor(tensor, device_mesh, [Shard(0)])
+
+        # Full reduction: Shard(0) reduces to Partial(sum), then to Replicate
+        dt_full = dtensor.nansum()
+        self.assertEqual(dt_full.full_tensor(), tensor.nansum())
+        self.assertEqual(dt_full.placements, (Partial("sum"),))
+
+        # Reduction along sharded dim 0: produces Partial(sum)
+        dt_dim0 = dtensor.nansum(dim=0)
+        self.assertEqual(dt_dim0.full_tensor(), tensor.nansum(dim=0))
+        self.assertEqual(dt_dim0.placements, (Partial("sum"),))
+
+        # Reduction along non-sharded dim 1: preserves Shard(0)
+        dt_dim1 = dtensor.nansum(dim=1)
+        self.assertEqual(dt_dim1.full_tensor(), tensor.nansum(dim=1))
+        self.assertEqual(dt_dim1.placements, (Shard(0),))
+
     @with_comms
     @skip_unless_torch_gpu
     def test_mean(self):
@@ -328,7 +365,8 @@ def _replicate_fn(name, module, device_mesh):
             from torch.distributed.tensor._dtensor_spec import TensorMeta
 
             dtensor_meta = y_dist._spec.tensor_meta
-            assert isinstance(dtensor_meta, TensorMeta)
+            if not isinstance(dtensor_meta, TensorMeta):
+                raise AssertionError(f"Expected TensorMeta, got {type(dtensor_meta)}")
             # make sure the right shape in sharding prop
             self.assertEqual(y_local.shape, dtensor_meta.shape)
             self.assertEqual(y_local, y_dist.full_tensor())
@@ -549,9 +587,15 @@ def forward(self, tokens):
                     for n, p in target_model.named_parameters():
                         if not req_grad_map.get(n.rpartition(".")[0], False):
                             p.requires_grad_(False)
-                            assert not p.requires_grad
+                            if p.requires_grad:
+                                raise AssertionError(
+                                    f"Expected requires_grad to be False for {n}"
+                                )
                         else:
-                            assert p.requires_grad
+                            if not p.requires_grad:
+                                raise AssertionError(
+                                    f"Expected requires_grad to be True for {n}"
+                                )
 
                 # forward step for both local and distributed models
                 x = torch.randint(vocab_size, (batch, seq_len), device=self.device_type)
@@ -561,7 +605,12 @@ def forward(self, tokens):
                 with CommDebugMode() as comm_mode:
                     output_dist = model_dist(x)
 
-                self.assertEqual(output_local, output_dist)
+                output_dist_cmp = (
+                    output_dist.full_tensor()
+                    if isinstance(output_dist, DTensor)
+                    else output_dist
+                )
+                self.assertEqual(output_local, output_dist_cmp)
 
                 # all requires_grad patterns should have the same forward comm counts
                 expected_fwd_comm = {
@@ -601,14 +650,20 @@ def forward(self, tokens):
                     comm_mode.comm_module_counts["Global"]["backward"],
                     expected_bwd_comm,
                 )
-                self.assertEqual(output_local, output_dist)
+                output_dist_cmp = (
+                    output_dist.full_tensor()
+                    if isinstance(output_dist, DTensor)
+                    else output_dist
+                )
+                self.assertEqual(output_local, output_dist_cmp)
 
             except Exception as e:
                 subtest_fails[subtest_cfg] = e
         # if any subtest fails, provide the failed subtests and report the overall failure
-        assert not subtest_fails, (
-            f"{len(subtest_fails)}/{len(subtest_cfgs)} subtests failed: {pformat(subtest_fails)}"
-        )
+        if subtest_fails:
+            raise AssertionError(
+                f"{len(subtest_fails)}/{len(subtest_cfgs)} subtests failed: {pformat(subtest_fails)}"
+            )
 
     @with_comms
     def test_topk(self):
@@ -633,9 +688,41 @@ def test_topk(self):
             out_full_values = out_dt.values.full_tensor()
             self.assertEqual(global_topk.values, out_full_values)
 
-            # TODO: support backward scatter
-            # global_topk.values.sum().backward()
-            # out_full_values.sum().backward()
+    @with_comms
+    def test_topk_backward(self):
+        device_mesh = self.build_device_mesh()
+        # Test backward for topk with various shard_dim / topk_dim combinations.
+        # The fix in value_selecting_reduction_backward ensures that the zeros
+        # tensor created during backward is properly dispatched through DTensor
+        # (using new_zeros instead of at::zeros).
+        topk_dims = [0, 1, -1]
+        shard_dims = [0, 1, 2]
+        for topk_dim, shard_dim in itertools.product(topk_dims, shard_dims):
+            # Reference: plain (non-distributed) backward
+            x = torch.randn(12, 8, 8, device=self.device_type, requires_grad=True)
+            ref_topk = x.topk(3, dim=topk_dim)
+            ref_topk.values.sum().backward()
+
+            # Distributed backward
+            dist_x = distribute_tensor(
+                x.detach().clone().requires_grad_(True),
+                device_mesh,
+                [Shard(shard_dim)],
+            )
+            dist_topk = dist_x.topk(3, dim=topk_dim)
+            dist_topk.values.sum().backward()
+
+            # Verify gradient exists, has the right placement, and matches
+            self.assertIsNotNone(
+                dist_x.grad,
+                msg=f"topk_dim={topk_dim}, shard_dim={shard_dim}",
+            )
+            self.assertEqual(
+                dist_x.grad.full_tensor(),
+                x.grad,
+                msg=f"topk_dim={topk_dim}, shard_dim={shard_dim}",
+            )
+            x.grad.zero_()
 
     @with_comms
     def test_shard0_svd(self):
@@ -857,6 +944,33 @@ def test_norm_0_on_psum(self):
         self.assertEqual(out.full_tensor().item(), 3.0)
         self.assertEqual(out.placements, (Replicate(),))
 
+    @with_comms
+    def test_max_min_dim_sharded(self):
+        """max.dim/min.dim on sharded tensors produce correct values and indices."""
+        device_mesh = self.build_device_mesh()
+        t = torch.randn(128, 64, device=self.device_type)
+
+        for op in [torch.max, torch.min]:
+            expected_vals, expected_idx = op(t, dim=1)
+
+            # Shard on non-reduction dim: should work directly
+            dt = distribute_tensor(t, device_mesh, [Shard(0)])
+            dt_vals, dt_idx = op(dt, dim=1)
+            self.assertEqual(dt_vals.full_tensor(), expected_vals)
+            self.assertEqual(dt_idx.full_tensor(), expected_idx)
+
+            # Shard on reduction dim: forces redistribute to Replicate
+            dt = distribute_tensor(t, device_mesh, [Shard(1)])
+            dt_vals, dt_idx = op(dt, dim=1)
+            self.assertEqual(dt_vals.full_tensor(), expected_vals)
+            self.assertEqual(dt_idx.full_tensor(), expected_idx)
+
+            # Partial input: forces redistribute to Replicate
+            dt = distribute_tensor(t, device_mesh, [Partial()])
+            dt_vals, dt_idx = op(dt, dim=1)
+            self.assertEqual(dt_vals.full_tensor(), expected_vals)
+            self.assertEqual(dt_idx.full_tensor(), expected_idx)
+
     @with_comms
     @skip_if_lt_x_gpu(4)
     def test_foreach_add_different_mesh(self):
@@ -973,6 +1087,104 @@ def test_cumsum(self):
                     self.assertTrue(output_dtensor.placements[0].is_shard(shard_dim))
                 self.assertEqual(output_dtensor.full_tensor(), output)
 
+    @with_comms
+    def test_scan_ops(self):
+        mesh = self.build_device_mesh()
+        comm_mode = CommDebugMode()
+        inp = torch.rand(3, 5, device=self.device_type)
+
+        shard_dim = 0
+        input_dtensor = distribute_tensor(
+            inp, device_mesh=mesh, placements=[Shard(shard_dim)]
+        )
+
+        for op in [torch.cumprod, torch.logcumsumexp]:
+            for dim in [0, 1]:
+                output = op(inp, dim=dim)
+                with comm_mode:
+                    output_dtensor = op(input_dtensor, dim=dim)
+                if dim == shard_dim:
+                    self.assertTrue(output_dtensor.placements[0].is_replicate())
+                else:
+                    self.assertTrue(output_dtensor.placements[0].is_shard(shard_dim))
+                self.assertEqual(output_dtensor.full_tensor(), output)
+
+    @with_comms
+    def test_scan_ops_with_indices(self):
+        mesh = self.build_device_mesh()
+        comm_mode = CommDebugMode()
+        inp = torch.rand(12, 8, device=self.device_type)
+
+        shard_dim = 0
+        input_dtensor = distribute_tensor(
+            inp, device_mesh=mesh, placements=[Shard(shard_dim)]
+        )
+
+        for op in [torch.cummax, torch.cummin]:
+            for dim in [0, 1]:
+                values, indices = op(inp, dim=dim)
+                with comm_mode:
+                    dt_values, dt_indices = op(input_dtensor, dim=dim)
+                if dim == shard_dim:
+                    self.assertTrue(dt_values.placements[0].is_replicate())
+                else:
+                    self.assertTrue(dt_values.placements[0].is_shard(shard_dim))
+                self.assertEqual(dt_values.full_tensor(), values)
+
+    @with_comms
+    def test_median(self):
+        device_mesh = self.build_device_mesh()
+        tensor = torch.randn(12, 8, device=self.device_type)
+        dtensor = distribute_tensor(tensor, device_mesh, [Shard(0)])
+
+        # Global median
+        self.assertEqual(torch.median(dtensor).full_tensor(), torch.median(tensor))
+
+        # nanmedian global
+        self.assertEqual(
+            torch.nanmedian(dtensor).full_tensor(), torch.nanmedian(tensor)
+        )
+
+    @with_comms
+    def test_dim_reductions_with_indices(self):
+        device_mesh = self.build_device_mesh()
+        tensor = torch.randn(12, 8, device=self.device_type)
+        shard_dim = 0
+        dtensor = distribute_tensor(tensor, device_mesh, [Shard(shard_dim)])
+
+        # nanmedian.dim
+        for dim in range(tensor.ndim):
+            vals, idxs = torch.nanmedian(tensor, dim=dim)
+            dt_vals, dt_idxs = torch.nanmedian(dtensor, dim=dim)
+            self.assertEqual(dt_vals.full_tensor(), vals)
+            self.assertEqual(dt_idxs.full_tensor(), idxs)
+            if dim == shard_dim:
+                self.assertTrue(dt_vals.placements[0].is_replicate())
+            else:
+                self.assertTrue(dt_vals.placements[0].is_shard(shard_dim))
+
+        # kthvalue: reduce along each dim
+        for dim in range(tensor.ndim):
+            vals, idxs = torch.kthvalue(tensor, 3, dim=dim)
+            dt_vals, dt_idxs = torch.kthvalue(dtensor, 3, dim=dim)
+            self.assertEqual(dt_vals.full_tensor(), vals)
+            self.assertEqual(dt_idxs.full_tensor(), idxs)
+            if dim == shard_dim:
+                self.assertTrue(dt_vals.placements[0].is_replicate())
+            else:
+                self.assertTrue(dt_vals.placements[0].is_shard(shard_dim))
+
+        # mode: reduce along each dim
+        for dim in range(tensor.ndim):
+            vals, idxs = torch.mode(tensor, dim=dim)
+            dt_vals, dt_idxs = torch.mode(dtensor, dim=dim)
+            self.assertEqual(dt_vals.full_tensor(), vals)
+            self.assertEqual(dt_idxs.full_tensor(), idxs)
+            if dim == shard_dim:
+                self.assertTrue(dt_vals.placements[0].is_replicate())
+            else:
+                self.assertTrue(dt_vals.placements[0].is_shard(shard_dim))
+
     @with_comms
     def test_conj_complex_dtensor(self):
         mesh = self.build_device_mesh()
@@ -1104,7 +1316,9 @@ def test_partial_reduction_ops(self):
         dt = dt.redistribute(dt.device_mesh, placements=[Replicate()])
         out_with_redistribute = torch.norm(dt)
 
-        self.assertEqual(out_without_redistribute, out_with_redistribute)
+        self.assertEqual(
+            out_without_redistribute.full_tensor(), out_with_redistribute.full_tensor()
+        )
 
         local_tensor = torch.rand(3, dtype=torch.float32, device=self.device_type)
         dt = DTensor.from_local(
@@ -1115,7 +1329,9 @@ def test_partial_reduction_ops(self):
         dt = dt.redistribute(dt.device_mesh, placements=[Replicate()])
         out_with_redistribute = torch.max(dt)
 
-        self.assertEqual(out_without_redistribute, out_with_redistribute)
+        self.assertEqual(
+            out_without_redistribute.full_tensor(), out_with_redistribute.full_tensor()
+        )
 
         local_tensor = torch.rand(3, dtype=torch.float32, device=self.device_type)
         dt = DTensor.from_local(
@@ -1126,7 +1342,9 @@ def test_partial_reduction_ops(self):
         dt = dt.redistribute(dt.device_mesh, placements=[Replicate()])
         out_with_redistribute = torch.min(dt)
 
-        self.assertEqual(out_without_redistribute, out_with_redistribute)
+        self.assertEqual(
+            out_without_redistribute.full_tensor(), out_with_redistribute.full_tensor()
+        )
 
     @with_comms
     def test_matching_partial_reduction_ops(self):
@@ -1145,7 +1363,9 @@ def test_matching_partial_reduction_ops(self):
 
         self.assertTrue(out_without_redistribute.placements[0].is_partial())
         self.assertTrue(out_with_redistribute.placements[0].is_replicate())
-        self.assertEqual(out_without_redistribute, out_with_redistribute)
+        self.assertEqual(
+            out_without_redistribute.full_tensor(), out_with_redistribute.full_tensor()
+        )
 
     @skip_if_lt_x_gpu(4)
     @with_comms
@@ -1172,6 +1392,293 @@ def test_std(self):
         self.assertEqual(res.placements, [Shard(0), Replicate()])
         self.assertEqual(res.full_tensor(), expected_answer)
 
+    @with_comms
+    def test_prims_pointwise_ops(self):
+        device_mesh = self.build_device_mesh()
+        x = torch.randn(12, 8)
+        y = torch.randn(12, 8)
+        dtensor_x = distribute_tensor(x, device_mesh, [Shard(0)])
+        dtensor_y = distribute_tensor(y, device_mesh, [Shard(0)])
+
+        for op in [
+            torch.ops.prims.bessel_i0e,
+            torch.ops.prims.bessel_i1,
+            torch.ops.prims.bessel_i1e,
+            torch.ops.prims.bessel_j0,
+            torch.ops.prims.bessel_j1,
+            torch.ops.prims.erfcx,
+            torch.ops.prims.ndtri,
+            torch.ops.prims.spherical_bessel_j0,
+            torch.special.erfcx,
+        ]:
+            local_result = op(x.abs().clamp(0.1, 3.0))
+            dtensor_result = op(dtensor_x.abs().clamp(0.1, 3.0))
+            self.assertEqual(dtensor_result.full_tensor(), local_result)
+            self.assertTrue(dtensor_result.placements[0].is_shard(dim=0))
+
+        for op in [
+            torch.ops.prims.div,
+            torch.ops.prims.gcd,
+            torch.ops.prims.ne,
+            torch.ops.aten.ne,
+            torch.ops.aten.gcd,
+        ]:
+            if op in [torch.ops.prims.gcd, torch.ops.aten.gcd]:
+                local_result = op(x.int(), y.int())
+                dtensor_result = op(dtensor_x.int(), dtensor_y.int())
+            else:
+                local_result = op(x, y)
+                dtensor_result = op(dtensor_x, dtensor_y)
+            self.assertEqual(dtensor_result.full_tensor(), local_result)
+
+    @with_comms
+    def test_prims_fma(self):
+        from torch._inductor import inductor_prims
+
+        device_mesh = self.build_device_mesh()
+        a = torch.randn(12, 8)
+        b = torch.randn(12, 8)
+        c = torch.randn(12, 8)
+        dt_a = distribute_tensor(a, device_mesh, [Shard(0)])
+        dt_b = distribute_tensor(b, device_mesh, [Shard(0)])
+        dt_c = distribute_tensor(c, device_mesh, [Shard(0)])
+
+        # All DTensor inputs
+        local_result = inductor_prims.fma(a, b, c)
+        dtensor_result = inductor_prims.fma(dt_a, dt_b, dt_c)
+        self.assertEqual(dtensor_result.full_tensor(), local_result)
+        self.assertTrue(dtensor_result.placements[0].is_shard(dim=0))
+
+        # Scalar tensor input (mimics the alpha argument in add_)
+        scalar = torch.tensor(2.0)
+        local_result = inductor_prims.fma(a, scalar, c)
+        dtensor_result = inductor_prims.fma(dt_a, scalar, dt_c)
+        self.assertEqual(dtensor_result.full_tensor(), local_result)
+        self.assertTrue(dtensor_result.placements[0].is_shard(dim=0))
+
+        # Replicate placement
+        dt_a_rep = distribute_tensor(a, device_mesh, [Replicate()])
+        dt_c_rep = distribute_tensor(c, device_mesh, [Replicate()])
+        local_result = inductor_prims.fma(a, scalar, c)
+        dtensor_result = inductor_prims.fma(dt_a_rep, scalar, dt_c_rep)
+        self.assertEqual(dtensor_result.full_tensor(), local_result)
+        self.assertTrue(dtensor_result.placements[0].is_replicate())
+
+    @with_comms
+    def test_prims_view_of(self):
+        device_mesh = self.build_device_mesh()
+        x = torch.randn(12, 8)
+        dtensor = distribute_tensor(x, device_mesh, [Shard(0)])
+
+        result = torch.ops.prims.view_of(dtensor)
+        self.assertTrue(result.placements[0].is_shard(dim=0))
+        self.assertEqual(result.full_tensor(), x)
+
+    @with_comms
+    def test_pooling_ops(self):
+        device_mesh = self.build_device_mesh()
+
+        # Single-output pooling
+        inp_4d = torch.randn(8, 3, 16, 16, device=self.device_type)
+        dt_4d = distribute_tensor(inp_4d, device_mesh, [Shard(0)])
+
+        for op, args in [
+            (torch.nn.functional.avg_pool2d, (3,)),
+            (torch.nn.functional.adaptive_avg_pool2d, ((4, 4),)),
+        ]:
+            expected = op(inp_4d, *args)
+            result = op(dt_4d, *args)
+            self.assertEqual(result.full_tensor(), expected)
+            self.assertTrue(result.placements[0].is_shard(0))
+
+        # Dual-output pooling (values + indices)
+        for op, args, kwargs in [
+            (
+                torch.nn.functional.adaptive_max_pool2d,
+                ((4, 4),),
+                {"return_indices": True},
+            ),
+        ]:
+            exp_val, exp_idx = op(inp_4d, *args, **kwargs)
+            dt_val, dt_idx = op(dt_4d, *args, **kwargs)
+            self.assertEqual(dt_val.full_tensor(), exp_val)
+            self.assertTrue(dt_val.placements[0].is_shard(0))
+
+        # 3D pooling
+        inp_5d = torch.randn(8, 3, 8, 8, 8, device=self.device_type)
+        dt_5d = distribute_tensor(inp_5d, device_mesh, [Shard(0)])
+
+        expected = torch.nn.functional.avg_pool3d(inp_5d, 3)
+        result = torch.nn.functional.avg_pool3d(dt_5d, 3)
+        self.assertEqual(result.full_tensor(), expected)
+        self.assertTrue(result.placements[0].is_shard(0))
+
+    @with_comms
+    @skip_unless_torch_gpu
+    def test_nll_loss_backward_comm_counts(self):
+        """Test backward comm counts for nll_loss/cross_entropy with Shard(0) inputs.
+
+        When inputs are Shard(0), nll_loss_forward produces total_weight as
+        Partial(sum). The backward strategy requires total_weight to be Replicate
+        for reduction='mean' (where total_weight is used to normalize gradients),
+        but NOT for reduction='sum' or 'none' (where total_weight is unused by
+        the backward kernel). This test verifies that the unnecessary all-reduce
+        is avoided for 'sum' and 'none' reductions.
+        """
+        device_mesh = self.build_device_mesh()
+        comm_mode = CommDebugMode()
+
+        channel_size = 16
+        # Expected backward all-reduce counts per reduction mode:
+        # - "sum": 0 (total_weight unused in backward)
+        # - "none": 0 (total_weight unused in backward)
+        # - "mean": 1 (total_weight needed to normalize gradients)
+        expected_backward_allreduce = {"sum": 0, "none": 0, "mean": 1}
+
+        for reduction, expected_allreduce_count in expected_backward_allreduce.items():
+            x = torch.rand(8, channel_size, device=self.device_type, requires_grad=True)
+            target = torch.randint(channel_size, (8,), device=self.device_type)
+
+            dist_x = distribute_tensor(x, device_mesh, [Shard(0)])
+            dist_target = distribute_tensor(target, device_mesh, [Shard(0)])
+
+            dist_loss = torch.nn.functional.cross_entropy(
+                dist_x, dist_target, reduction=reduction
+            )
+
+            with comm_mode:
+                if reduction == "none":
+                    dist_loss.sum().backward()
+                else:
+                    dist_loss.backward()
+
+            allreduce_count = comm_mode.get_comm_counts().get(funcol.all_reduce, 0)
+            self.assertEqual(
+                allreduce_count,
+                expected_allreduce_count,
+                f"reduction='{reduction}': expected {expected_allreduce_count} "
+                f"backward all-reduce(s), got {allreduce_count}. "
+                f"Full comm counts: {pformat(dict(comm_mode.get_comm_counts()))}",
+            )
+
+            # Verify correctness
+            local_loss = torch.nn.functional.cross_entropy(
+                x, target, reduction=reduction
+            )
+            if reduction == "none":
+                local_loss.sum().backward()
+            else:
+                local_loss.backward()
+            self.assertEqual(dist_x.grad.full_tensor(), x.grad)
+
+    @with_comms
+    @skip_unless_torch_gpu
+    def test_linalg_ops(self):
+        device_mesh = DeviceMesh(self.device_type, list(range(self.world_size)))
+
+        # Batched positive-definite matrix for ops that require it (cholesky, inv)
+        A = torch.randn(8, 4, 4, device=self.device_type)
+        spd = A @ A.mT + 4 * torch.eye(4, device=self.device_type)
+        dt_spd = distribute_tensor(spd, device_mesh, [Shard(0)])
+
+        # cholesky
+        expected = torch.linalg.cholesky(spd)
+        result = torch.linalg.cholesky(dt_spd)
+        self.assertEqual(result.full_tensor(), expected)
+        self.assertTrue(result.placements[0].is_shard(0))
+
+        # inv
+        expected = torch.linalg.inv(spd)
+        result = torch.linalg.inv(dt_spd)
+        self.assertEqual(result.full_tensor(), expected)
+        self.assertTrue(result.placements[0].is_shard(0))
+
+        # lu_factor
+        B = torch.randn(8, 4, 4, device=self.device_type)
+        dt_B = distribute_tensor(B, device_mesh, [Shard(0)])
+        exp_LU, exp_piv = torch.linalg.lu_factor(B)
+        res_LU, res_piv = torch.linalg.lu_factor(dt_B)
+        self.assertEqual(res_LU.full_tensor(), exp_LU)
+        self.assertTrue(res_LU.placements[0].is_shard(0))
+
+        # eig requires MAGMA on ROCm (no hipSOLVER fallback); skip if unavailable
+        if not TEST_WITH_ROCM or torch.cuda.has_magma:
+            expected_vals, expected_vecs = torch.linalg.eig(B)
+            result_vals, result_vecs = torch.linalg.eig(dt_B)
+            self.assertEqual(result_vals.full_tensor(), expected_vals)
+            self.assertTrue(result_vals.placements[0].is_shard(0))
+
+        # solve
+        rhs = torch.randn(8, 4, 2, device=self.device_type)
+        dt_rhs = distribute_tensor(rhs, device_mesh, [Shard(0)])
+        expected = torch.linalg.solve(spd, rhs)
+        result = torch.linalg.solve(dt_spd, dt_rhs)
+        self.assertEqual(result.full_tensor(), expected)
+        self.assertTrue(result.placements[0].is_shard(0))
+
+    @with_comms
+    @skip_unless_torch_gpu
+    def test_linalg_cross(self):
+        device_mesh = DeviceMesh(self.device_type, list(range(self.world_size)))
+        a = torch.randn(8, 4, 3, device=self.device_type)
+        b = torch.randn(8, 4, 3, device=self.device_type)
+
+        # Shard on batch dim (dim=0), cross on default dim=-1
+        dt_a = distribute_tensor(a, device_mesh, [Shard(0)])
+        dt_b = distribute_tensor(b, device_mesh, [Shard(0)])
+        expected = torch.linalg.cross(a, b)
+        result = torch.linalg.cross(dt_a, dt_b)
+        self.assertEqual(result.full_tensor(), expected)
+        self.assertTrue(result.placements[0].is_shard(0))
+
+        # Shard on dim=1, cross on dim=-1
+        dt_a1 = distribute_tensor(a, device_mesh, [Shard(1)])
+        dt_b1 = distribute_tensor(b, device_mesh, [Shard(1)])
+        result1 = torch.linalg.cross(dt_a1, dt_b1)
+        self.assertEqual(result1.full_tensor(), expected)
+        self.assertTrue(result1.placements[0].is_shard(1))
+
+    @with_comms
+    @skip_unless_torch_gpu
+    def test_linalg_solve_partial(self):
+        device_mesh = DeviceMesh(self.device_type, list(range(self.world_size)))
+        A = torch.randn(4, 4, device=self.device_type, dtype=torch.float64)
+        A = A @ A.mT + 4 * torch.eye(4, device=self.device_type, dtype=torch.float64)
+        B = torch.randn(4, 2, device=self.device_type, dtype=torch.float64)
+
+        dt_A = distribute_tensor(A, device_mesh, [Replicate()])
+        dt_B = distribute_tensor(B, device_mesh, [Partial()])
+        expected = torch.linalg.solve(A, B)
+        result = torch.linalg.solve(dt_A, dt_B)
+        self.assertEqual(result.full_tensor(), expected)
+
+    @with_comms
+    @skip_unless_torch_gpu
+    def test_interpolation_upsample_ops(self):
+        device_mesh = self.build_device_mesh()
+        F = torch.nn.functional
+
+        inp = torch.randn(8, 3, 16, 16, device=self.device_type)
+        dt_inp = distribute_tensor(inp, device_mesh, [Shard(0)])
+
+        # F.interpolate with nearest mode
+        expected = F.interpolate(inp, size=(8, 8), mode="nearest")
+        result = F.interpolate(dt_inp, size=(8, 8), mode="nearest")
+        self.assertEqual(result.full_tensor(), expected)
+        self.assertTrue(result.placements[0].is_shard(0))
+
+        # F.interpolate with area mode
+        expected = F.interpolate(inp, size=(8, 8), mode="area")
+        result = F.interpolate(dt_inp, size=(8, 8), mode="area")
+        self.assertEqual(result.full_tensor(), expected)
+        self.assertTrue(result.placements[0].is_shard(0))
+
+        # F.adaptive_avg_pool2d
+        expected = F.adaptive_avg_pool2d(inp, (4, 4))
+        result = F.adaptive_avg_pool2d(dt_inp, (4, 4))
+        self.assertEqual(result.full_tensor(), expected)
+        self.assertTrue(result.placements[0].is_shard(0))
+
 
 DistMathOpsTestWithLocalTensor = create_local_tensor_test_class(
     DistMathOpsTest,
diff --git a/test/distributed/tensor/test_matrix_ops.py b/test/distributed/tensor/test_matrix_ops.py
index 931ded4c273df..26d2fc744461c 100644
--- a/test/distributed/tensor/test_matrix_ops.py
+++ b/test/distributed/tensor/test_matrix_ops.py
@@ -4,7 +4,7 @@
 import itertools
 import math
 import unittest
-from typing import cast, Optional
+from typing import cast
 
 import torch
 import torch.nn.functional as F
@@ -22,12 +22,14 @@
     mm_single_dim_strategy,
 )
 from torch.distributed.tensor._ops.single_dim_strategy import (
+    _ShardingPlaceholder,
     register_single_dim_strategy,
 )
 from torch.distributed.tensor.debug import CommDebugMode
 from torch.distributed.tensor.placement_types import _StridedShard
 from torch.testing._internal.common_cuda import PLATFORM_SUPPORTS_FP8, SM90OrLater
 from torch.testing._internal.common_device_type import E4M3_MAX_POS, e4m3_type
+from torch.testing._internal.common_distributed import skip_if_lt_x_gpu
 from torch.testing._internal.common_utils import (
     instantiate_parametrized_tests,
     parametrize,
@@ -275,6 +277,41 @@ def test_gen_single_dim_einsum_strategies_bias_reduce_op(self):
                         f"Got bias={bias_placement.reduce_op}, output={output_placement.reduce_op}",
                     )
 
+    def test_gen_single_dim_einsum_strategies_batch_linearity(self):
+        """Test that batch-only equations auto-detect all-Partial linearity."""
+        S = _ShardingPlaceholder
+
+        # For "abcd,abcd->abcd": all dims are batch, no contracting/free.
+        # Expect: 4 batch-dim + 4 per-input linearity + 2 all-Partial linearity = 10
+        strategies = gen_single_dim_einsum_strategies("abcd,abcd->abcd")
+
+        # Convert to repr tuples for comparison since _ShardingPlaceholder lacks __eq__
+        actual = [tuple(repr(p) for p in s) for s in strategies]
+        expected = [
+            # batch dims: shard output and both inputs on same dim
+            (repr(S(0)), repr(S(0)), repr(S(0))),
+            (repr(S(1)), repr(S(1)), repr(S(1))),
+            (repr(S(2)), repr(S(2)), repr(S(2))),
+            (repr(S(3)), repr(S(3)), repr(S(3))),
+            # per-input linearity: one input Partial, other Replicate
+            ("Partial(sum)", "Partial(sum)", "Replicate()"),
+            ("Partial(sum)", "Replicate()", "Partial(sum)"),
+            ("Partial(avg)", "Partial(avg)", "Replicate()"),
+            ("Partial(avg)", "Replicate()", "Partial(avg)"),
+            # batch-dimension linearity: all inputs Partial
+            ("Partial(sum)", "Partial(sum)", "Partial(sum)"),
+            ("Partial(avg)", "Partial(avg)", "Partial(avg)"),
+        ]
+        self.assertEqual(actual, expected)
+
+        # For "mk,kn->mn": has contracting dim k, no all-Partial linearity.
+        mm_strategies = gen_single_dim_einsum_strategies("mk,kn->mn")
+        mm_all_partial = [
+            s for s in mm_strategies if all(isinstance(p, Partial) for p in s)
+        ]
+        self.assertEqual(len(mm_all_partial), 0)
+
+    @skip_if_lt_x_gpu(4)
     @with_comms
     def test_mm_with_strided_input(self):
         # Case 1: 1D mesh with StridedShard
@@ -490,6 +527,69 @@ def test_scaled_mm(self):
 
             self.assertEqual(comm_mode.get_total_counts(), 0)
 
+    def test_scaled_mm_blockwise_1d_scale_placement(self):
+        """Test that _scaled_mm_scale_placement handles 1D blockwise scales correctly.
+
+        1D blockwise scales arise in MX (microscaling) formats where a data
+        tensor [M, K] has a flattened scale of shape [M * K / block_size].
+        Shard(>=1) is invalid on a 1D tensor, so the strategy must map
+        non-contracting shards to Shard(0) and reject contracting-dim shards.
+        """
+        from torch.distributed.tensor._ops._matrix_ops import _scaled_mm_scale_placement
+
+        # --- Tensor-wise scale (single element) -> always Replicate ---
+        result = _scaled_mm_scale_placement(
+            Shard(0), torch.Size([1]), contracting_dim=1
+        )
+        self.assertEqual(result, Replicate())
+        result = _scaled_mm_scale_placement(Shard(0), torch.Size([]), contracting_dim=1)
+        self.assertEqual(result, Replicate())
+
+        # --- 2D scale -> copy data placement directly (row-wise) ---
+        result = _scaled_mm_scale_placement(
+            Shard(0), torch.Size([16, 1]), contracting_dim=1
+        )
+        self.assertEqual(result, Shard(0))
+
+        # --- 1D blockwise + non-contracting shard -> Shard(0) ---
+        # A (mk): dim 0 = m (non-contracting), dim 1 = k (contracting)
+        result = _scaled_mm_scale_placement(
+            Shard(0), torch.Size([64]), contracting_dim=1
+        )
+        self.assertIsNotNone(result)
+        self.assertEqual(result, Shard(0))
+
+        # B_t (kn): dim 1 = n (non-contracting), dim 0 = k (contracting)
+        result = _scaled_mm_scale_placement(
+            Shard(1), torch.Size([64]), contracting_dim=0
+        )
+        self.assertIsNotNone(result)
+        self.assertEqual(result, Shard(0))
+
+        # --- 1D blockwise + contracting shard -> None (unsupported) ---
+        result = _scaled_mm_scale_placement(
+            Shard(1), torch.Size([64]), contracting_dim=1
+        )
+        self.assertIsNone(result)
+        result = _scaled_mm_scale_placement(
+            Shard(0), torch.Size([64]), contracting_dim=0
+        )
+        self.assertIsNone(result)
+
+        # --- 1D blockwise + Replicate -> Replicate ---
+        result = _scaled_mm_scale_placement(
+            Replicate(), torch.Size([64]), contracting_dim=1
+        )
+        self.assertIsNotNone(result)
+        self.assertEqual(result, Replicate())
+
+        # --- 1D blockwise + Partial -> Replicate ---
+        result = _scaled_mm_scale_placement(
+            Partial(), torch.Size([64]), contracting_dim=0
+        )
+        self.assertIsNotNone(result)
+        self.assertEqual(result, Replicate())
+
     @with_comms
     def test_matmul(self):
         device_mesh = self.build_device_mesh()
@@ -575,7 +675,7 @@ def test_placement_comb(
             batch_2_placements: list[Placement],
             beta: int,
             alpha: int,
-            batch_1_grad: Optional[torch.Tensor],
+            batch_1_grad: torch.Tensor | None,
         ) -> None:
             tensor_dt = distribute_tensor(tensor, device_mesh, tensor_placements)
             batch_1_dt = distribute_tensor(batch_1, device_mesh, batch_1_placements)
@@ -587,8 +687,10 @@ def test_placement_comb(
                 ),
             ).redistribute(device_mesh, [Replicate()])
             dist_local_res = dist_res.to_local()
-            assert not torch.isnan(local_result).any()
-            assert not torch.isnan(dist_local_res).any()
+            if torch.isnan(local_result).any():
+                raise AssertionError("NaN values found in local_result")
+            if torch.isnan(dist_local_res).any():
+                raise AssertionError("NaN values found in dist_local_res")
             self.assertEqual(dist_local_res.detach(), local_result.detach())
 
             # TODO: add test backward
@@ -776,21 +878,24 @@ def _run_mm(device_mesh, reduce_op="sum"):
     def test_scaled_dot_product_attention(self):
         device_mesh = self.build_device_mesh()
         comm_mode = CommDebugMode()
+        head_dim = 8
+        if self.device_type == "xpu":
+            head_dim = 64
         # bsz, n_heads, slen, head_dim
         query = torch.rand(
-            (4, 8, 8, 8),
+            (4, 8, 8, head_dim),
             device=self.device_type,
             dtype=torch.bfloat16,
             requires_grad=True,
         )
         key = torch.rand(
-            (4, 8, 8, 8),
+            (4, 8, 8, head_dim),
             device=self.device_type,
             dtype=torch.bfloat16,
             requires_grad=True,
         )
         value = torch.rand(
-            (4, 8, 8, 8),
+            (4, 8, 8, head_dim),
             device=self.device_type,
             dtype=torch.bfloat16,
             requires_grad=True,
@@ -996,6 +1101,37 @@ def test_grouped_mm(self, kwargs):
         self.assertEqual(dist_w1.grad.full_tensor(), w1.grad)
         self.assertEqual(dist_w2.grad.full_tensor(), w2.grad)
 
+    @with_comms
+    def test_constant_pad_nd(self):
+        """constant_pad_nd: shard non-padded, replicate padded, Partial iff value==0."""
+        device_mesh = self.build_device_mesh()
+        t = torch.randn(8, 6, device=self.device_type)
+        pad = [1, 1]  # pad last dim only
+        expected = torch.nn.functional.pad(t, pad, value=0.0)
+
+        # Shard on non-padded dim (dim 0) — should work directly
+        dt = distribute_tensor(t, device_mesh, [Shard(0)])
+        result = torch.nn.functional.pad(dt, pad, value=0.0)
+        self.assertEqual(result.full_tensor(), expected)
+
+        # Shard on padded dim (dim 1) — forces redistribute to Replicate
+        dt = distribute_tensor(t, device_mesh, [Shard(1)])
+        result = torch.nn.functional.pad(dt, pad, value=0.0)
+        self.assertEqual(result.full_tensor(), expected)
+
+        # Partial input with value=0 — Partial passes through
+        dt = distribute_tensor(t, device_mesh, [Partial()])
+        result = torch.nn.functional.pad(dt, pad, value=0.0)
+        self.assertEqual(result.placements, (Partial(),))
+        self.assertEqual(result.full_tensor(), expected)
+
+        # Partial input with value!=0 — forces redistribute to Replicate
+        expected_nz = torch.nn.functional.pad(t, pad, value=1.0)
+        dt = distribute_tensor(t, device_mesh, [Partial()])
+        result = torch.nn.functional.pad(dt, pad, value=1.0)
+        self.assertNotEqual(result.placements, (Partial(),))
+        self.assertEqual(result.full_tensor(), expected_nz)
+
 
 instantiate_parametrized_tests(DistMatrixOpsTest)
 
diff --git a/test/distributed/tensor/test_op_strategy.py b/test/distributed/tensor/test_op_strategy.py
index 979133961ba58..dc3c228047966 100644
--- a/test/distributed/tensor/test_op_strategy.py
+++ b/test/distributed/tensor/test_op_strategy.py
@@ -2,7 +2,6 @@
 
 import itertools
 import random
-from contextlib import contextmanager
 from itertools import chain
 from unittest.mock import patch
 
@@ -31,16 +30,15 @@
     EinsumDims,
     gen_einsum_strategies,
 )
-from torch.distributed.tensor._ops.utils import (
-    register_op_strategy,
-    replicate_op_strategy,
-)
-from torch.distributed.tensor.debug import _clear_sharding_prop_cache, CommDebugMode
+from torch.distributed.tensor._ops.utils import replicate_op_strategy
+from torch.distributed.tensor.debug import CommDebugMode
+from torch.distributed.tensor.placement_types import _StridedShard
 from torch.testing._internal.common_utils import run_tests, TestCase
 from torch.testing._internal.distributed._tensor.common_dtensor import (
     create_local_tensor_test_class,
     DTensorOpTestBase,
     DTensorTestBase,
+    op_strategy_context,
     with_comms,
 )
 from torch.testing._internal.distributed.fake_pg import FakeStore
@@ -209,6 +207,25 @@ def test_redistribute_cost_mesh_1d(self):
         cost = redistribute_cost(shard_spec, partial_spec)
         self.assertEqual(cost, float("inf"))
 
+    def test_redistribute_cost_strided_shard(self):
+        """_StridedShard specs get inf cost (shard_order is None bail-out)."""
+        mesh_1d = self.build_device_mesh()
+
+        global_tensor_meta = extract_tensor_meta(torch.randn(10, 10))
+
+        strided_shard_spec = DTensorSpec(
+            mesh_1d, (_StridedShard(0, split_factor=2),), global_tensor_meta
+        )
+        replica_spec = DTensorSpec(mesh_1d, (Replicate(),), global_tensor_meta)
+
+        # _StridedShard as source: shard_order is None → inf
+        self.assertEqual(
+            redistribute_cost(strided_shard_spec, replica_spec), float("inf")
+        )
+        # Replicate as source: is_replicated() short-circuits to 0 before
+        # reaching the shard_order check
+        self.assertEqual(redistribute_cost(replica_spec, strided_shard_spec), 0.0)
+
     def test_redistribute_cost_latency(self):
         # test cost model on addmm op
         from torch.distributed.tensor._ops._matrix_ops import addmm_strategy
@@ -493,45 +510,6 @@ def _fw_tuple(x, y, z):
 )
 
 
-@contextmanager
-def op_strategy_context(op_overload, strategy_func, schema_info=None):
-    """
-    Context manager for setting and clearing op strategies.
-    Args:
-        op_overload: The operator overload to set or clear the strategy for.
-        strategy_func: The strategy function to set for the operator overload.
-        schema_info: Optional schema information for the operator overload.
-    Yields:
-        None
-    """
-    propagator = DTensor._op_dispatcher.sharding_propagator
-    _origin_op_strategy_funcs = None
-    _origin_op_strategy_schema = None
-    try:
-        # register the op strategy
-        if op_overload in propagator.op_strategy_funcs:
-            _origin_op_strategy_funcs = propagator.op_strategy_funcs[op_overload]
-            del propagator.op_strategy_funcs[op_overload]
-        if op_overload in propagator.op_to_schema_info:
-            _origin_op_strategy_schema = propagator.op_to_schema_info[op_overload]
-            del propagator.op_to_schema_info[op_overload]
-        register_op_strategy(op_overload, schema_info=schema_info)(strategy_func)
-        yield
-    finally:
-        # clear this op strategy cache
-        if _origin_op_strategy_funcs is None:
-            if op_overload in propagator.op_strategy_funcs:
-                del propagator.op_strategy_funcs[op_overload]
-        else:
-            propagator.op_strategy_funcs[op_overload] = _origin_op_strategy_funcs
-        if _origin_op_strategy_schema is None:
-            if op_overload in propagator.op_to_schema_info:
-                del propagator.op_to_schema_info[op_overload]
-        else:
-            propagator.op_to_schema_info[op_overload] = _origin_op_strategy_schema
-        _clear_sharding_prop_cache()
-
-
 def detect_exists_identical_opspec(*args, op, mesh, strategy_function) -> bool:
     """
     Given sample input args, detect if identical OpSpecs exists under the same
@@ -594,9 +572,8 @@ def mock_select_func(strategy, op_schema=None):
 
             op_spec_costs: list[float] = []
             for op_spec in strategy.strategies:
-                assert op_spec.redistribute_cost is not None, (
-                    "must set redistribute cost each OpSpec!"
-                )
+                if op_spec.redistribute_cost is None:
+                    raise AssertionError("must set redistribute cost each OpSpec!")
                 costs_from__select_strategy.append(op_spec.redistribute_cost)
                 redistribute_cost = sum(chain.from_iterable(op_spec.redistribute_cost))
                 op_spec_costs.append(redistribute_cost)
@@ -864,6 +841,258 @@ def test_kwargs_meta_mixed(self):
         self.assertEqual(kwargs_meta["dim"], 0)
         self.assertEqual(kwargs_meta["alpha"], 1.0)
 
+    def test_max_min_dim_single_dim_strategy(self):
+        """max.dim/min.dim only allow sharding on non-reduction dims."""
+        from torch.distributed.tensor._ops._math_ops import (
+            max_min_dim_single_dim_strategy,
+        )
+
+        input_meta = TensorMeta(
+            shape=torch.Size([8, 4]), stride=(4, 1), dtype=torch.float32
+        )
+        # reduce along dim=1: only dim 0 should be shardable
+        strategies = max_min_dim_single_dim_strategy(
+            torch.ops.aten.max.dim, (input_meta, 1), {}
+        )
+        # Should have exactly 1 strategy: shard on dim 0
+        self.assertEqual(len(strategies), 1)
+        values_plc, indices_plc, input_plc = strategies[0]
+        self.assertEqual(values_plc.dim, 0)
+        self.assertEqual(indices_plc.dim, 0)
+        self.assertEqual(input_plc.dim, 0)
+
+        # reduce along dim=0: only dim 1 should be shardable,
+        # output dim maps to 0 (since reduction collapses dim 0)
+        strategies = max_min_dim_single_dim_strategy(
+            torch.ops.aten.min.dim, (input_meta, 0), {}
+        )
+        self.assertEqual(len(strategies), 1)
+        values_plc, indices_plc, input_plc = strategies[0]
+        self.assertEqual(values_plc.dim, 0)
+        self.assertEqual(indices_plc.dim, 0)
+        self.assertEqual(input_plc.dim, 1)
+
+    def test_layer_norm_fwd_single_dim_strategy(self):
+        """layer_norm produces sharding rules for each outer dim."""
+        from torch.distributed.tensor._ops._math_ops import (
+            layer_norm_single_dim_strategy,
+        )
+
+        input_meta = TensorMeta(
+            shape=torch.Size([4, 8]), stride=(8, 1), dtype=torch.float32
+        )
+        weight_meta = TensorMeta(
+            shape=torch.Size([8]), stride=(1,), dtype=torch.float32
+        )
+        bias_meta = TensorMeta(shape=torch.Size([8]), stride=(1,), dtype=torch.float32)
+
+        # normalized_shape=[8] => axis=1, only dim 0 is shardable
+        strategies = layer_norm_single_dim_strategy(
+            torch.ops.aten.native_layer_norm.default,
+            (input_meta, [8], weight_meta, bias_meta, 1e-5),
+            {},
+        )
+        self.assertEqual(len(strategies), 1)
+        # [out, mean, rstd, input, weight, bias]
+        out, mean, rstd, inp, w, b = strategies[0]
+        self.assertEqual(out.dim, 0)
+        self.assertEqual(inp.dim, 0)
+        self.assertIsInstance(w, Replicate)
+        self.assertIsInstance(b, Replicate)
+
+        # Without bias: 5 elements per rule
+        strategies = layer_norm_single_dim_strategy(
+            torch.ops.aten.native_layer_norm.default,
+            (input_meta, [8], weight_meta, None, 1e-5),
+            {},
+        )
+        self.assertEqual(len(strategies), 1)
+        self.assertEqual(len(strategies[0]), 5)
+
+        # normalized_shape=[4, 8] => axis=0, no shardable dims
+        strategies = layer_norm_single_dim_strategy(
+            torch.ops.aten.native_layer_norm.default,
+            (input_meta, [4, 8], weight_meta, bias_meta, 1e-5),
+            {},
+        )
+        self.assertEqual(len(strategies), 0)
+
+    def test_rms_norm_fwd_single_dim_strategy(self):
+        """rms_norm produces sharding rules for each outer dim."""
+        from torch.distributed.tensor._ops._math_ops import rms_norm_single_dim_strategy
+
+        input_meta = TensorMeta(
+            shape=torch.Size([4, 8]), stride=(8, 1), dtype=torch.float32
+        )
+        weight_meta = TensorMeta(
+            shape=torch.Size([8]), stride=(1,), dtype=torch.float32
+        )
+
+        strategies = rms_norm_single_dim_strategy(
+            torch.ops.aten._fused_rms_norm.default,
+            (input_meta, [8], weight_meta, 1e-5),
+            {},
+        )
+        self.assertEqual(len(strategies), 1)
+        # [out, rrms, input, weight]
+        out, rrms, inp, w = strategies[0]
+        self.assertEqual(out.dim, 0)
+        self.assertEqual(rrms.dim, 0)
+        self.assertEqual(inp.dim, 0)
+        self.assertIsInstance(w, Replicate)
+
+        # Without weight: 3 elements per rule
+        strategies = rms_norm_single_dim_strategy(
+            torch.ops.aten._fused_rms_norm.default,
+            (input_meta, [8], None, 1e-5),
+            {},
+        )
+        self.assertEqual(len(strategies), 1)
+        self.assertEqual(len(strategies[0]), 3)
+
+    def test_layer_norm_bwd_single_dim_strategy(self):
+        """layer_norm backward produces correct output/input placements."""
+        from torch.distributed.tensor._ops._math_ops import (
+            layer_norm_bwd_single_dim_strategy,
+        )
+
+        input_meta = TensorMeta(
+            shape=torch.Size([4, 8]), stride=(8, 1), dtype=torch.float32
+        )
+        stat_meta = TensorMeta(shape=torch.Size([4]), stride=(1,), dtype=torch.float32)
+        weight_meta = TensorMeta(
+            shape=torch.Size([8]), stride=(1,), dtype=torch.float32
+        )
+
+        # With weight and bias
+        strategies = layer_norm_bwd_single_dim_strategy(
+            torch.ops.aten.native_layer_norm_backward.default,
+            (
+                input_meta,
+                input_meta,
+                [8],
+                stat_meta,
+                stat_meta,
+                weight_meta,
+                weight_meta,
+                [True, True, True],
+            ),
+            {},
+        )
+        self.assertEqual(len(strategies), 1)
+        rule = strategies[0]
+        # outputs: [d_input, d_weight, d_bias] + inputs: [grad_out, input, mean, rstd, weight, bias]
+        self.assertEqual(len(rule), 9)
+        d_input, d_weight, d_bias = rule[0], rule[1], rule[2]
+        self.assertEqual(d_input.dim, 0)  # sharded
+        self.assertIsInstance(d_weight, Partial)  # reduced
+        self.assertIsInstance(d_bias, Partial)  # reduced
+
+        # Without weight/bias: d_weight=None, d_bias=None
+        strategies = layer_norm_bwd_single_dim_strategy(
+            torch.ops.aten.native_layer_norm_backward.default,
+            (
+                input_meta,
+                input_meta,
+                [8],
+                stat_meta,
+                stat_meta,
+                None,
+                None,
+                [True, False, False],
+            ),
+            {},
+        )
+        self.assertEqual(len(strategies), 1)
+        rule = strategies[0]
+        # outputs: [d_input, None, None] + inputs: [grad_out, input, mean, rstd]
+        self.assertEqual(len(rule), 7)
+        self.assertIsNone(rule[1])  # d_weight
+        self.assertIsNone(rule[2])  # d_bias
+
+    def test_rms_norm_bwd_single_dim_strategy(self):
+        """rms_norm backward produces correct output/input placements."""
+        from torch.distributed.tensor._ops._math_ops import (
+            rms_norm_bwd_single_dim_strategy,
+        )
+
+        input_meta = TensorMeta(
+            shape=torch.Size([4, 8]), stride=(8, 1), dtype=torch.float32
+        )
+        stat_meta = TensorMeta(shape=torch.Size([4]), stride=(1,), dtype=torch.float32)
+        weight_meta = TensorMeta(
+            shape=torch.Size([8]), stride=(1,), dtype=torch.float32
+        )
+
+        # With weight
+        strategies = rms_norm_bwd_single_dim_strategy(
+            torch.ops.aten._fused_rms_norm_backward.default,
+            (
+                input_meta,  # grad_out
+                input_meta,  # input
+                [8],  # normalized_shape
+                stat_meta,  # rstd
+                weight_meta,  # weight
+            ),
+            {},
+        )
+        self.assertEqual(len(strategies), 1)
+        rule = strategies[0]
+        # outputs: [d_input, d_weight] + inputs: [grad_out, input, rstd, weight]
+        self.assertEqual(len(rule), 6)
+        self.assertEqual(rule[0].dim, 0)  # d_input sharded
+        self.assertIsInstance(rule[1], Partial)  # d_weight reduced
+        self.assertIsInstance(rule[5], Replicate)  # weight replicated
+
+        # Without weight: d_weight=None
+        strategies = rms_norm_bwd_single_dim_strategy(
+            torch.ops.aten._fused_rms_norm_backward.default,
+            (input_meta, input_meta, [8], stat_meta, None),
+            {},
+        )
+        self.assertEqual(len(strategies), 1)
+        rule = strategies[0]
+        # outputs: [d_input, None] + inputs: [grad_out, input, rstd]
+        self.assertEqual(len(rule), 5)
+        self.assertIsNone(rule[1])  # d_weight
+
+    def test_constant_pad_nd_allows_shard_on_non_padded_dim(self):
+        """constant_pad_nd should allow sharding on non-padded dims."""
+        from torch.distributed.tensor._ops._matrix_ops import (
+            constant_pad_nd_single_dim_strategy,
+        )
+
+        input_meta = TensorMeta(
+            shape=torch.Size([8, 4]), stride=(4, 1), dtype=torch.float32
+        )
+        # Pad only dim 1: [1, 1] means pad last dim by 1 on each side
+        strategies = constant_pad_nd_single_dim_strategy(
+            torch.ops.aten.constant_pad_nd.default,
+            (input_meta, [1, 1], 0.0),
+            {},
+        )
+        # Should include a strategy that shards on dim 0
+        shard_dims = {s[0].dim for s in strategies if hasattr(s[0], "dim")}
+        self.assertIn(0, shard_dims)
+
+    def test_constant_pad_nd_bans_shard_on_padded_dim(self):
+        """constant_pad_nd should ban sharding on padded dims."""
+        from torch.distributed.tensor._ops._matrix_ops import (
+            constant_pad_nd_single_dim_strategy,
+        )
+
+        input_meta = TensorMeta(
+            shape=torch.Size([8, 4]), stride=(4, 1), dtype=torch.float32
+        )
+        # Pad only dim 1
+        strategies = constant_pad_nd_single_dim_strategy(
+            torch.ops.aten.constant_pad_nd.default,
+            (input_meta, [1, 1], 0.0),
+            {},
+        )
+        shard_dims = {s[0].dim for s in strategies if hasattr(s[0], "dim")}
+        self.assertNotIn(1, shard_dims, "Should not shard on padded dim")
+
 
 class TestExpandToFullMeshOpStrategy(TestCase):
     """Tests for expand_to_full_mesh_op_strategy function.
diff --git a/test/distributed/tensor/test_optimizers.py b/test/distributed/tensor/test_optimizers.py
index fc3932bb866eb..abc4bde442986 100644
--- a/test/distributed/tensor/test_optimizers.py
+++ b/test/distributed/tensor/test_optimizers.py
@@ -12,7 +12,11 @@
     Replicate,
     Shard,
 )
-from torch.testing._internal.common_utils import run_tests
+from torch.testing._internal.common_utils import (
+    instantiate_parametrized_tests,
+    parametrize,
+    run_tests,
+)
 from torch.testing._internal.distributed._tensor.common_dtensor import (
     create_local_tensor_test_class,
     DTensorTestBase,
@@ -44,7 +48,8 @@ def input_fn(mod, inputs, device_mesh):
 
 # prepare output to be local torch.Tensor
 def output_fn(mod, outputs, device_mesh):
-    assert isinstance(outputs, DTensor)
+    if not isinstance(outputs, DTensor):
+        raise AssertionError(f"Expected DTensor, got {type(outputs)}")
     return outputs.redistribute(placements=[Replicate()] * device_mesh.ndim).to_local()
 
 
@@ -716,6 +721,62 @@ def _input_fn_2d(mod, inputs, device_mesh):
             inp = torch.ones(8, 10, device=self.device_type)
             self._assert_optimizer(None, mod, opt, mod_copy, dist_opt, inp)
 
+    @with_comms
+    @parametrize("foreach", [True, False])
+    def test_adamw_sharding_cache_no_leak(self, foreach):
+        """Step-varying scalar args (like AdamW's bias corrections) must not
+        cause unbounded growth of the DTensor sharding propagation cache.
+
+        foreach=True exercises the ScalarList-in-list path (_foreach_addcdiv_),
+        foreach=False exercises the scalar-kwarg path (addcdiv_ value=...).
+        """
+        mesh = init_device_mesh(self.device_type, (self.world_size,))
+        model = MLPModule(self.device_type)
+        for name, param in model.named_parameters():
+            dist_param = nn.Parameter(distribute_tensor(param, mesh, [Shard(0)]))
+            dist_param.register_hook(
+                lambda grad: grad.redistribute(placements=[Shard(0)])
+            )
+            # set the parameter on the parent module
+            parts = name.split(".")
+            mod = model
+            for part in parts[:-1]:
+                mod = getattr(mod, part)
+            mod.register_parameter(parts[-1], dist_param)
+
+        optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4, foreach=foreach)
+
+        def run_step():
+            inp = distribute_tensor(
+                torch.randn(8, 10, device=self.device_type), mesh, [Shard(0)]
+            )
+            out = model(inp)
+            out.sum().backward()
+            optimizer.step()
+            optimizer.zero_grad()
+
+        # Warmup to populate caches
+        for _ in range(3):
+            run_step()
+
+        _, misses_after_warmup = torch._C._get_DTensor_sharding_propagator_cache_stats()
+
+        for _ in range(10):
+            run_step()
+
+        # After warmup, the cache should be fully populated: no new
+        # misses across subsequent training steps.
+        _, misses_after_steps = torch._C._get_DTensor_sharding_propagator_cache_stats()
+        self.assertEqual(
+            misses_after_warmup,
+            misses_after_steps,
+            f"Cache misses increased from {misses_after_warmup} to "
+            f"{misses_after_steps} over 10 steps (expected 0 new misses "
+            "after warmup)",
+        )
+
+
+instantiate_parametrized_tests(TestDTensorOptimizer)
 
 TestDTensorOptimizerWithLocalTensor = create_local_tensor_test_class(
     TestDTensorOptimizer,
diff --git a/test/distributed/tensor/test_placement_types.py b/test/distributed/tensor/test_placement_types.py
index 7634a37865af4..49998e4be83c0 100644
--- a/test/distributed/tensor/test_placement_types.py
+++ b/test/distributed/tensor/test_placement_types.py
@@ -3,7 +3,6 @@
 import itertools
 
 import torch
-from torch._dynamo.variables.distributed import PlacementClassVariable
 from torch.distributed.tensor.placement_types import (
     _StridedShard,
     Partial,
@@ -67,12 +66,6 @@ def test_strided_shard_kwonly_argument(self):
             _StridedShard(3, 4)
         _StridedShard(3, split_factor=4)
 
-    def test_dynamo_can_identify_placement_classes(self):
-        for cls in (Replicate, Shard, _StridedShard, Partial):
-            self.assertTrue(
-                PlacementClassVariable.is_placement_type(cls), msg=f"failed on {cls}"
-            )
-
     def test_select_split_tensor_matches_split_tensor(self):
         """
         Test that _select_split_tensor produces the same result as indexing
diff --git a/test/distributed/tensor/test_pointwise_ops.py b/test/distributed/tensor/test_pointwise_ops.py
index 6d7132f1b7b8a..64f5d5a68101b 100644
--- a/test/distributed/tensor/test_pointwise_ops.py
+++ b/test/distributed/tensor/test_pointwise_ops.py
@@ -2,7 +2,7 @@
 # Owner(s): ["oncall: distributed"]
 
 from collections.abc import Callable, Sequence
-from typing import Any, Optional
+from typing import Any
 from unittest import skip
 
 import torch
@@ -12,6 +12,7 @@
     DeviceMesh,
     distribute_tensor,
     DTensor,
+    init_device_mesh,
     Partial,
     Placement,
     Replicate,
@@ -19,10 +20,12 @@
 )
 from torch.distributed.tensor._ops._math_ops import _NormPartial
 from torch.distributed.tensor.debug import CommDebugMode
+from torch.distributed.tensor.placement_types import _StridedShard
 from torch.testing._internal.common_utils import (
     instantiate_parametrized_tests,
     parametrize,
     run_tests,
+    TestCase,
 )
 from torch.testing._internal.distributed._tensor.common_dtensor import (
     create_local_tensor_test_class,
@@ -86,9 +89,9 @@ def _compare_pairwise_ops(
         device_mesh: DeviceMesh,
         placements: Sequence[Placement],
         op: Callable,
-        pre_op_fn: Optional[Callable] = None,
+        pre_op_fn: Callable | None = None,
         args: Sequence[Any] = (),
-        kwargs: Optional[dict[str, Any]] = None,
+        kwargs: dict[str, Any] | None = None,
     ):
         if pre_op_fn is None:
             pre_op_fn = no_op
@@ -127,7 +130,7 @@ def _run_sharded_elementwise_ops(
         *,
         device_mesh: DeviceMesh,
         placements: Sequence[Placement],
-        pre_op_fn: Optional[Callable] = None,
+        pre_op_fn: Callable | None = None,
         input_size: Sequence[int],
         op: Callable,
         **kwargs,
@@ -317,6 +320,31 @@ def test_mul_out(self):
         self.assertEqual(input_tensor, dtensor.to_local())
         self.assertEqual(expected, dt.to_local())
 
+    @with_comms
+    def test_mul_out_partial(self):
+        device_mesh = self.build_device_mesh()
+        input_size = (8, 4)
+        # P(sum) * R -> P(sum), with out= kwarg
+        partial_tensor = DTensor.from_local(
+            torch.ones(*input_size, device=self.device_type),
+            device_mesh,
+            [Partial("sum")],
+        )
+        replicate_tensor = DTensor.from_local(
+            torch.full(input_size, 2.0, device=self.device_type),
+            device_mesh,
+            [Replicate()],
+            run_check=False,
+        )
+        output_tensor = DTensor.from_local(
+            torch.empty(*input_size, device=self.device_type),
+            device_mesh,
+            [Partial("sum")],
+        )
+        dt = torch.mul(partial_tensor, replicate_tensor, out=output_tensor)
+        self.assertEqual(dt.placements, (Partial("sum"),))
+        self.assertEqual(output_tensor.placements, (Partial("sum"),))
+
     @with_comms
     def test_mul_partial(self):
         # we only test the partial behavior for mul op as other placement
@@ -520,7 +548,7 @@ def test_mul_div_scalar_partial(self):
         self.assertTrue(res._spec.placements[0].is_partial())
         res = res.redistribute(dt.device_mesh, placements=[Replicate()])
         expected = sum(i for i in range(self.world_size)) * 2
-        self.assertEqual(res, expected)
+        self.assertEqual(res.full_tensor(), expected)
 
         res = aten.div.Scalar(dt, 2)
         self.assertEqual(
@@ -531,7 +559,7 @@ def test_mul_div_scalar_partial(self):
         self.assertTrue(res._spec.placements[0].is_partial())
         res = res.redistribute(dt.device_mesh, placements=[Replicate()])
         expected = sum(i for i in range(self.world_size)) / 2
-        self.assertEqual(res, expected)
+        self.assertEqual(res.full_tensor(), expected)
 
     @with_comms
     def test_mul_div_scalar_norm_partial(self):
@@ -543,22 +571,10 @@ def test_mul_div_scalar_norm_partial(self):
         norm = dt.norm()
         self.assertTrue(isinstance(norm._spec.placements[0], _NormPartial))
 
-        res = aten.mul.Scalar(norm, 2)
-        self.assertTrue(isinstance(res._spec.placements[0], _NormPartial))
-        res = res.redistribute(dt.device_mesh, placements=[Replicate()])
-        self.assertEqual(res, 20)
-
-        res = aten.div.Scalar(norm, 2)
-        self.assertTrue(isinstance(res._spec.placements[0], _NormPartial))
-        res = res.redistribute(dt.device_mesh, placements=[Replicate()])
-        self.assertEqual(res, 5)
-
-        res = aten.mul.Scalar(norm, -2)
-        self.assertTrue(res._spec.placements[0].is_replicate())
-
-        res = aten.div.Scalar(norm, -2)
-        self.assertEqual(res, -5)
-        self.assertTrue(res._spec.placements[0].is_replicate())
+        self.assertEqual(aten.mul.Scalar(norm, 2).full_tensor(), 20)
+        self.assertEqual(aten.div.Scalar(norm, 2).full_tensor(), 5)
+        self.assertEqual(aten.mul.Scalar(norm, -2).full_tensor(), -20)
+        self.assertEqual(aten.div.Scalar(norm, -2).full_tensor(), -5)
 
     @with_comms
     def test_add_sub_scalar_partial(self):
@@ -575,7 +591,7 @@ def test_add_sub_scalar_partial(self):
 
         res = dt + 1
         expected = sum(i for i in range(self.world_size)) + 1
-        self.assertEqual(res, expected)
+        self.assertEqual(res.full_tensor(), expected)
         self.assertTrue(res._spec.placements[0].is_replicate())
 
         # regular partial - scalar -> replicate
@@ -587,12 +603,12 @@ def test_add_sub_scalar_partial(self):
 
         res = dt - 1
         expected = sum(i for i in range(self.world_size)) - 1
-        self.assertEqual(res, expected)
+        self.assertEqual(res.full_tensor(), expected)
         self.assertTrue(res._spec.placements[0].is_replicate())
 
         res = 7 - dt
         expected = 7 - sum(i for i in range(self.world_size))
-        self.assertEqual(res, expected)
+        self.assertEqual(res.full_tensor(), expected)
         self.assertTrue(res._spec.placements[0].is_replicate())
 
         # regular partial + regular partial -> partial
@@ -601,14 +617,14 @@ def test_add_sub_scalar_partial(self):
         self.assertTrue(res._spec.placements[0].is_partial())
         res = res.redistribute(dt.device_mesh, placements=[Replicate()])
         expected = sum(i for i in range(self.world_size)) * 2
-        self.assertEqual(res, expected)
+        self.assertEqual(res.full_tensor(), expected)
 
         # regular partial - regular partial -> partial
         res = dt - dt
         self.assertEqual(res.to_local(), rank - rank)
         self.assertTrue(res._spec.placements[0].is_partial())
         res = res.redistribute(dt.device_mesh, placements=[Replicate()])
-        self.assertEqual(res, 0)
+        self.assertEqual(res.full_tensor(), 0)
 
     @with_comms
     def test_add_sub_scalar_norm_partial(self):
@@ -622,7 +638,7 @@ def test_add_sub_scalar_norm_partial(self):
         self.assertTrue(isinstance(norm._spec.placements[0], _NormPartial))
         norm = norm + 1
 
-        self.assertEqual(norm, 11)
+        self.assertEqual(norm.full_tensor(), 11)
         self.assertTrue(norm._spec.placements[0].is_replicate())
 
         dt = distribute_tensor(local_tensor, mesh, [Shard(0)])
@@ -631,7 +647,7 @@ def test_add_sub_scalar_norm_partial(self):
         self.assertTrue(isinstance(norm._spec.placements[0], _NormPartial))
         norm = norm - 1
 
-        self.assertEqual(norm, 9)
+        self.assertEqual(norm.full_tensor(), 9)
         self.assertTrue(norm._spec.placements[0].is_replicate())
 
     @with_comms
@@ -683,15 +699,25 @@ def test_neg_partial(self):
             self.assertEqual(z.placements, (Partial(partial_op),))
             self.assertEqual(z.full_tensor(), expected_full)
 
-        # test non-sum/avg partial to assert the partial not getting propagated
-        # since -max(A1, A2) != max(-A1, -A2)
+        # test non-sum/avg partial: neg is monotone decreasing, so
+        # -max(A1, A2) = min(-A1, -A2), i.e. Partial("max") -> Partial("min")
         d_input = DTensor.from_local(input, device_mesh, [Partial("max")])
 
-        z = torch.neg(d_input)
-        self.assertEqual(z.placements, (Replicate(),))
-        self.assertEqual(
-            z.to_local(), torch.full((8, 8), -2.0, device=self.device_type)
-        )
+        with comm_mode:
+            z = torch.neg(d_input)
+
+        comm_counts = comm_mode.get_total_counts()
+        self.assertEqual(comm_counts, 0)
+        self.assertEqual(z.placements, (Partial("min"),))
+
+        d_input = DTensor.from_local(input, device_mesh, [Partial("min")])
+
+        with comm_mode:
+            z = torch.neg(d_input)
+
+        comm_counts = comm_mode.get_total_counts()
+        self.assertEqual(comm_counts, 0)
+        self.assertEqual(z.placements, (Partial("max"),))
 
     @with_comms
     def test_maximum_mixed_partials_redistribution(self):
@@ -723,8 +749,659 @@ def test_maximum_mixed_partials_redistribution(self):
         expected_value = float(self.world_size)
         self.assertEqual(result.full_tensor()[0, 0].item(), expected_value)
 
+    @with_comms
+    def test_add_partial_with_replicate_rules(self):
+        # P(x) + R -> P(x) for x in {avg, max, min}: adding a replicated constant
+        # preserves the partial reduce type (constant is the same on all ranks)
+        device_mesh = self.build_device_mesh()
+        comm_mode = CommDebugMode()
+
+        other_val = 3.0
+        d_other = distribute_tensor(
+            torch.full((8, 8), other_val, device=self.device_type),
+            device_mesh,
+            [Replicate()],
+        )
+
+        expected_full = {
+            "avg": (self.world_size + 1) / 2.0 + other_val,
+            "max": float(self.world_size) + other_val,
+            "min": 1.0 + other_val,
+        }
+
+        for reduce_op in ("avg", "max", "min"):
+            input_tensor = torch.ones(8, 8, device=self.device_type) * (self.rank + 1)
+            d_input = DTensor.from_local(
+                input_tensor, device_mesh, [Partial(reduce_op)]
+            )
+
+            with comm_mode:
+                result = d_input + d_other
+
+            self.assertEqual(comm_mode.get_total_counts(), 0)
+            self.assertEqual(result.placements, (Partial(reduce_op),))
+            self.assertEqual(
+                result.full_tensor(),
+                torch.full((8, 8), expected_full[reduce_op], device=self.device_type),
+            )
+
+    @with_comms
+    def test_add_replicate_with_partial_avg(self):
+        # R + P(avg) -> P(avg): avg is linear, so this holds for any alpha
+        device_mesh = self.build_device_mesh()
+        comm_mode = CommDebugMode()
+
+        replicate_val = 3.0
+        d_input = distribute_tensor(
+            torch.full((8, 8), replicate_val, device=self.device_type),
+            device_mesh,
+            [Replicate()],
+        )
+        other_tensor = torch.ones(8, 8, device=self.device_type) * (self.rank + 1)
+        d_other = DTensor.from_local(other_tensor, device_mesh, [Partial("avg")])
+
+        with comm_mode:
+            result = d_input + d_other
+
+        self.assertEqual(comm_mode.get_total_counts(), 0)
+        self.assertEqual(result.placements, (Partial("avg"),))
+        # full_tensor = R_val + avg(1, 2, ..., world_size) = 3 + (ws+1)/2
+        expected = torch.full(
+            (8, 8),
+            replicate_val + (self.world_size + 1) / 2.0,
+            device=self.device_type,
+        )
+        self.assertEqual(result.full_tensor(), expected)
+
+    @with_comms
+    def test_mul_div_replicate_partial_asymmetry(self):
+        # mul is bilinear: r * (p1 + p2) = r*p1 + r*p2, so R * P(sum) -> P(sum).
+        # div is only linear in the numerator: r / (p1 + p2) != r/p1 + r/p2.
+        device_mesh = self.build_device_mesh()
+        comm_mode = CommDebugMode()
+
+        replicate_val = 3.0
+        partial_local_val = 2.0
+
+        d_replicate = distribute_tensor(
+            torch.full((8, 8), replicate_val, device=self.device_type),
+            device_mesh,
+            [Replicate()],
+        )
+
+        # R * P(sum) -> P(sum): valid, zero communication
+        d_partial = DTensor.from_local(
+            torch.full((8, 8), partial_local_val, device=self.device_type),
+            device_mesh,
+            [Partial("sum")],
+        )
+        with comm_mode:
+            mul_result = d_replicate * d_partial
+
+        self.assertEqual(comm_mode.get_total_counts(), 0)
+        self.assertEqual(mul_result.placements, (Partial("sum"),))
+        # 3 * sum(2, 2, ...) = 3 * 2 * world_size
+        expected_mul = torch.full(
+            (8, 8),
+            replicate_val * partial_local_val * self.world_size,
+            device=self.device_type,
+        )
+        self.assertEqual(mul_result.full_tensor(), expected_mul)
+
+        # R / P(sum) -> requires communication (not linear in denominator)
+        d_partial2 = DTensor.from_local(
+            torch.full((8, 8), partial_local_val, device=self.device_type),
+            device_mesh,
+            [Partial("sum")],
+        )
+        with comm_mode:
+            div_result = d_replicate / d_partial2
+
+        self.assertGreater(comm_mode.get_total_counts(), 0)
+        # After P(sum) reduction: 3 / (2 * world_size)
+        expected_div = torch.full(
+            (8, 8),
+            replicate_val / (partial_local_val * self.world_size),
+            device=self.device_type,
+        )
+        self.assertEqual(div_result.full_tensor(), expected_div)
+
+    @with_comms
+    def test_inplace_add_partial_avg_with_replicate(self):
+        # P(avg) += R -> P(avg): self is P(avg), other is R, valid inplace
+        # because the rule P(avg)+R->P(avg) keeps output == self placement.
+        device_mesh = self.build_device_mesh()
+        comm_mode = CommDebugMode()
+
+        self_tensor = torch.ones(8, 8, device=self.device_type) * (self.rank + 1)
+        d_self = DTensor.from_local(self_tensor, device_mesh, [Partial("avg")])
+        d_other = distribute_tensor(
+            torch.full((8, 8), 3.0, device=self.device_type),
+            device_mesh,
+            [Replicate()],
+        )
+
+        with comm_mode:
+            d_self.add_(d_other)
+
+        self.assertEqual(comm_mode.get_total_counts(), 0)
+        self.assertEqual(d_self.placements, (Partial("avg"),))
+        # avg(1,2,...,ws) + 3 = (ws+1)/2 + 3
+        expected = torch.full(
+            (8, 8),
+            (self.world_size + 1) / 2.0 + 3.0,
+            device=self.device_type,
+        )
+        self.assertEqual(d_self.full_tensor(), expected)
+
+    @with_comms
+    def test_inplace_add_replicate_with_partial_avg_requires_comm(self):
+        # R += P(avg): the out-of-place rule R+P(avg)->P(avg) is valid, but
+        # inplace requires output == self == R, so P(avg) output is rejected.
+        # The runtime must redistribute P(avg) to R first (communication).
+        device_mesh = self.build_device_mesh()
+        comm_mode = CommDebugMode()
+
+        self_tensor = torch.full((8, 8), 3.0, device=self.device_type)
+        d_self = distribute_tensor(self_tensor, device_mesh, [Replicate()])
+        other_tensor = torch.ones(8, 8, device=self.device_type) * (self.rank + 1)
+        d_other = DTensor.from_local(other_tensor, device_mesh, [Partial("avg")])
+
+        with comm_mode:
+            d_self.add_(d_other)
+
+        self.assertGreater(comm_mode.get_total_counts(), 0)
+        self.assertEqual(d_self.placements, (Replicate(),))
+        # 3 + avg(1,2,...,ws) = 3 + (ws+1)/2
+        expected = torch.full(
+            (8, 8),
+            3.0 + (self.world_size + 1) / 2.0,
+            device=self.device_type,
+        )
+        self.assertEqual(d_self.full_tensor(), expected)
+
+    @with_comms
+    @parametrize(
+        "op_fn,second_arg,placement,expected",
+        [
+            # Binary list ops preserve Partial("sum")
+            (torch._foreach_add, "list", Partial("sum"), Partial("sum")),
+            # Scalar mul preserves Partial("sum") (linearity)
+            (torch._foreach_mul, 2.0, Partial("sum"), Partial("sum")),
+            # Neg preserves Partial("sum") (linearity)
+            (torch._foreach_neg, None, Partial("sum"), Partial("sum")),
+            # Binary list ops preserve Shard(0)
+            (torch._foreach_add, "list", Shard(0), Shard(0)),
+        ],
+    )
+    def test_foreach_placement_propagation(
+        self, op_fn, second_arg, placement, expected
+    ):
+        device_mesh = self.build_device_mesh()
+        comm_mode = CommDebugMode()
+
+        shapes = [(8, 8), (4, 4), (2, 6)]
+        if placement.is_shard():
+            dts = [
+                distribute_tensor(
+                    torch.rand(s, device=self.device_type), device_mesh, [placement]
+                )
+                for s in shapes
+            ]
+        else:
+            dts = [
+                DTensor.from_local(
+                    torch.rand(s, device=self.device_type), device_mesh, [placement]
+                )
+                for s in shapes
+            ]
+
+        if second_arg == "list":
+            if placement.is_shard():
+                args = [
+                    distribute_tensor(
+                        torch.rand(s, device=self.device_type),
+                        device_mesh,
+                        [placement],
+                    )
+                    for s in shapes
+                ]
+            else:
+                args = [
+                    DTensor.from_local(
+                        torch.rand(s, device=self.device_type),
+                        device_mesh,
+                        [placement],
+                    )
+                    for s in shapes
+                ]
+            call_args = (dts, args)
+        elif second_arg is None:
+            call_args = (dts,)
+        else:
+            call_args = (dts, second_arg)
+
+        with comm_mode:
+            result = op_fn(*call_args)
+
+        self.assertEqual(comm_mode.get_total_counts(), 0)
+        self.assertEqual(len(result), len(shapes))
+        for r in result:
+            self.assertIsInstance(r, DTensor)
+            self.assertEqual(r.placements, (expected,))
+
+    @with_comms
+    def test_foreach_mixed_placements(self):
+        """Each element in a foreach list independently preserves its placement."""
+        device_mesh = self.build_device_mesh()
+        comm_mode = CommDebugMode()
+
+        dt1 = DTensor.from_local(
+            torch.rand(8, 8, device=self.device_type),
+            device_mesh,
+            [Partial("sum")],
+        )
+        ds1 = DTensor.from_local(
+            torch.rand(8, 8, device=self.device_type),
+            device_mesh,
+            [Partial("sum")],
+        )
+        dt2 = distribute_tensor(
+            torch.rand(8, 4, device=self.device_type), device_mesh, [Shard(0)]
+        )
+        ds2 = distribute_tensor(
+            torch.rand(8, 4, device=self.device_type), device_mesh, [Shard(0)]
+        )
+
+        with comm_mode:
+            result = torch._foreach_add([dt1, dt2], [ds1, ds2])
+
+        self.assertEqual(comm_mode.get_total_counts(), 0)
+        self.assertEqual(result[0].placements, (Partial("sum"),))
+        self.assertEqual(result[1].placements, (Shard(0),))
+
+    @with_comms
+    @skip_unless_torch_gpu
+    @parametrize(
+        "mesh_type,amsgrad",
+        [
+            # All tensors on the same 1D mesh
+            ("same", False),
+            # state_steps on a different sub-mesh, with max_exp_avg_sqs present
+            ("cross", True),
+            # state_steps on a different sub-mesh, empty max_exp_avg_sqs
+            # (triggers the different_mesh_args index remapping)
+            ("cross", False),
+        ],
+    )
+    def test_fused_adam(self, mesh_type, amsgrad):
+        if mesh_type == "same":
+            param_mesh = self.build_device_mesh()
+            step_mesh = param_mesh
+        else:
+            mesh_2d = init_device_mesh(
+                self.device_type,
+                (2, self.world_size // 2),
+                mesh_dim_names=("dp", "tp"),
+            )
+            param_mesh = mesh_2d["tp"]
+            step_mesh = mesh_2d["dp"]
+
+        d_params = [
+            distribute_tensor(
+                torch.rand(8, 8, device=self.device_type), param_mesh, [Shard(0)]
+            )
+        ]
+        d_grads = [
+            distribute_tensor(
+                torch.rand(8, 8, device=self.device_type), param_mesh, [Shard(0)]
+            )
+        ]
+        d_exp_avgs = [
+            distribute_tensor(
+                torch.zeros(8, 8, device=self.device_type), param_mesh, [Shard(0)]
+            )
+        ]
+        d_exp_avg_sqs = [
+            distribute_tensor(
+                torch.zeros(8, 8, device=self.device_type), param_mesh, [Shard(0)]
+            )
+        ]
+        d_max_exp_avg_sqs = (
+            [
+                distribute_tensor(
+                    torch.zeros(8, 8, device=self.device_type),
+                    param_mesh,
+                    [Shard(0)],
+                )
+            ]
+            if amsgrad
+            else []
+        )
+        d_state_steps = [
+            distribute_tensor(
+                torch.tensor(1.0, device=self.device_type), step_mesh, [Replicate()]
+            )
+        ]
+
+        torch._fused_adam_(
+            d_params,
+            d_grads,
+            d_exp_avgs,
+            d_exp_avg_sqs,
+            d_max_exp_avg_sqs,
+            d_state_steps,
+            lr=0.001,
+            beta1=0.9,
+            beta2=0.999,
+            weight_decay=0.0,
+            eps=1e-8,
+            amsgrad=amsgrad,
+            maximize=False,
+        )
+
+        self.assertIsInstance(d_params[0], DTensor)
+        self.assertEqual(d_params[0].placements, (Shard(0),))
+        self.assertEqual(d_params[0].device_mesh, param_mesh)
+
+    @with_comms
+    @skip_unless_torch_gpu
+    def test_fused_adamw_strided_shard_cross_mesh(self):
+        """Reproduce the torchtitan pattern: params with _StridedShard on a 2D
+        mesh, state_steps with Replicate on a 1D DP sub-mesh."""
+        mesh_2d = init_device_mesh(
+            self.device_type,
+            (2, self.world_size // 2),
+            mesh_dim_names=("dp", "tp"),
+        )
+        dp_mesh = mesh_2d["dp"]
+        placements = [_StridedShard(0, split_factor=2), Shard(0)]
+
+        d_params = [
+            DTensor.from_local(
+                torch.randn(4, 8, device=self.device_type),
+                mesh_2d,
+                placements,
+                run_check=False,
+            )
+        ]
+        d_grads = [
+            DTensor.from_local(
+                torch.randn(4, 8, device=self.device_type),
+                mesh_2d,
+                placements,
+                run_check=False,
+            )
+        ]
+        d_exp_avgs = [
+            DTensor.from_local(
+                torch.zeros(4, 8, device=self.device_type),
+                mesh_2d,
+                placements,
+                run_check=False,
+            )
+        ]
+        d_exp_avg_sqs = [
+            DTensor.from_local(
+                torch.zeros(4, 8, device=self.device_type),
+                mesh_2d,
+                placements,
+                run_check=False,
+            )
+        ]
+        d_state_steps = [
+            DTensor.from_local(
+                torch.tensor(1.0, device=self.device_type),
+                dp_mesh,
+                [Replicate()],
+                run_check=False,
+            )
+        ]
+
+        torch.ops.aten._fused_adamw_.default(
+            d_params,
+            d_grads,
+            d_exp_avgs,
+            d_exp_avg_sqs,
+            [],
+            d_state_steps,
+            lr=0.1,
+            beta1=0.9,
+            beta2=0.999,
+            weight_decay=0.01,
+            eps=1e-8,
+            amsgrad=False,
+            maximize=False,
+        )
+
+        self.assertIsInstance(d_params[0], DTensor)
+        self.assertEqual(d_params[0].placements, tuple(placements))
+        self.assertEqual(d_params[0].device_mesh, mesh_2d)
+
+    @with_comms
+    def test_fused_adamw_tensor_lr(self):
+        """Test _fused_adamw_ with tensor lr kwarg."""
+        device_mesh = self.build_device_mesh()
+
+        d_params = [
+            distribute_tensor(
+                torch.rand(8, 8, device=self.device_type), device_mesh, [Shard(0)]
+            )
+        ]
+        d_grads = [
+            distribute_tensor(
+                torch.rand(8, 8, device=self.device_type), device_mesh, [Shard(0)]
+            )
+        ]
+        d_exp_avgs = [
+            distribute_tensor(
+                torch.zeros(8, 8, device=self.device_type), device_mesh, [Shard(0)]
+            )
+        ]
+        d_exp_avg_sqs = [
+            distribute_tensor(
+                torch.zeros(8, 8, device=self.device_type), device_mesh, [Shard(0)]
+            )
+        ]
+        d_state_steps = [
+            distribute_tensor(
+                torch.tensor(1.0, device=self.device_type),
+                device_mesh,
+                [Replicate()],
+            )
+        ]
+        d_lr = distribute_tensor(
+            torch.tensor(0.001, device=self.device_type),
+            device_mesh,
+            [Replicate()],
+        )
+
+        # tensor_lr variant passes lr as a tensor kwarg
+        torch.ops.aten._fused_adamw_.tensor_lr(
+            d_params,
+            d_grads,
+            d_exp_avgs,
+            d_exp_avg_sqs,
+            [],
+            d_state_steps,
+            lr=d_lr,
+            beta1=0.9,
+            beta2=0.999,
+            weight_decay=0.01,
+            eps=1e-8,
+            amsgrad=False,
+            maximize=False,
+        )
+
+        self.assertIsInstance(d_params[0], DTensor)
+        self.assertEqual(d_params[0].placements, (Shard(0),))
+
+    @with_comms
+    def test_new_pointwise_ops(self):
+        device_mesh = self.build_device_mesh()
+        torch.manual_seed(self.rank)
+        shard_spec = [Shard(0)]
+
+        # hardshrink: non-decreasing unary, decomp-blocked
+        inp = torch.randn(8, 5, device=self.device_type)
+        dt = DTensor.from_local(inp, device_mesh, shard_spec)
+        self.assertEqual(
+            torch.nn.functional.hardshrink(dt).to_local(),
+            torch.nn.functional.hardshrink(inp),
+        )
+
+        # lcm: binary integer pointwise
+        a = torch.randint(1, 100, (8, 5), device=self.device_type)
+        b = torch.randint(1, 100, (8, 5), device=self.device_type)
+        da = DTensor.from_local(a, device_mesh, shard_spec)
+        db = DTensor.from_local(b, device_mesh, shard_spec)
+        self.assertEqual(torch.lcm(da, db).to_local(), torch.lcm(a, b))
+
+        # special_xlog1py and special_entr
+        pos = torch.rand(8, 5, device=self.device_type) + 0.1
+        dp = DTensor.from_local(pos, device_mesh, shard_spec)
+        self.assertEqual(
+            torch.special.xlog1py(dp, dp).to_local(),
+            torch.special.xlog1py(pos, pos),
+        )
+        self.assertEqual(
+            torch.special.entr(dp).to_local(),
+            torch.special.entr(pos),
+        )
+
+    @with_comms
+    def test_new_special_pointwise_ops(self):
+        device_mesh = self.build_device_mesh()
+        torch.manual_seed(self.rank)
+        shard_spec = [Shard(0)]
+
+        # Unary special functions
+        inp = torch.rand(8, 5, device=self.device_type) + 0.1
+        dt = DTensor.from_local(inp, device_mesh, shard_spec)
+        for op in [
+            torch.special.airy_ai,
+            torch.special.bessel_y0,
+            torch.special.bessel_y1,
+            torch.special.modified_bessel_i0,
+            torch.special.modified_bessel_i1,
+            torch.special.modified_bessel_k0,
+            torch.special.modified_bessel_k1,
+            torch.special.scaled_modified_bessel_k0,
+            torch.special.scaled_modified_bessel_k1,
+        ]:
+            self.assertEqual(op(dt).to_local(), op(inp))
+
+        # Binary special polynomial ops
+        x = torch.rand(8, 5, device=self.device_type)
+        n = torch.randint(0, 5, (8, 5), device=self.device_type, dtype=x.dtype)
+        dx = DTensor.from_local(x, device_mesh, shard_spec)
+        dn = DTensor.from_local(n, device_mesh, shard_spec)
+        for op in [
+            torch.special.chebyshev_polynomial_t,
+            torch.special.chebyshev_polynomial_u,
+            torch.special.hermite_polynomial_h,
+            torch.special.hermite_polynomial_he,
+            torch.special.laguerre_polynomial_l,
+            torch.special.legendre_polynomial_p,
+        ]:
+            self.assertEqual(op(dx, dn).to_local(), op(x, n))
+
 
 instantiate_parametrized_tests(DistElementwiseOpsTest)
+
+
+class TestPointwiseRuleValidation(TestCase):
+    """Validate registered partial-placement rules via OpInfo samples."""
+
+    world_size = 2
+
+    def setUp(self):
+        super().setUp()
+        if not torch.distributed.is_initialized():
+            torch.distributed.init_process_group(
+                "fake", rank=0, world_size=self.world_size
+            )
+
+    def tearDown(self):
+        super().tearDown()
+        if torch.distributed.is_initialized():
+            torch.distributed.destroy_process_group()
+
+    def _with_even_sizes(self, fn):
+        """Run fn with opinfo sizes overridden to be evenly divisible by world_size."""
+        import torch.testing._internal.common_methods_invocations as common_ops
+        from torch.testing._internal.opinfo import core as opinfo_core
+
+        orig_sizes = (opinfo_core.L, opinfo_core.M, opinfo_core.S, opinfo_core.XS)
+        opinfo_core.L = common_ops.L = 24
+        opinfo_core.M = common_ops.M = 12
+        opinfo_core.S = common_ops.S = 4
+        opinfo_core.XS = common_ops.XS = 2
+        try:
+            return fn()
+        finally:
+            (
+                opinfo_core.L,
+                opinfo_core.M,
+                opinfo_core.S,
+                opinfo_core.XS,
+            ) = orig_sizes
+            (
+                common_ops.L,
+                common_ops.M,
+                common_ops.S,
+                common_ops.XS,
+            ) = orig_sizes
+
+    def test_per_category_no_incorrect_rules(self):
+        """Verify no incorrect rules for one representative op per category.
+
+        Each category of specialized pointwise ops shares the same partial
+        placement rules. Validating one representative per category catches
+        any mathematically incorrect rules via OpInfo samples (including
+        edge cases like negative alpha for add/sub).
+        """
+        from torch.distributed.tensor._ops.strategy_validation import compare_operator
+
+        representative_ops = [
+            # (op_name, category, expected_true_positives)
+            ("add", "binary_additive", 81),
+            ("mul", "binary_mul", 49),
+            ("div", "binary_div", 67),
+            ("sigmoid", "non_decreasing_unary", 5),
+            ("erfc", "non_increasing_unary", 5),
+            ("neg", "neg", 6),
+            ("maximum", "monotonic_max_preserving", 57),
+            ("minimum", "monotonic_min_preserving", 57),
+            ("logaddexp", "monotonic_binary", 49),
+            ("abs", "pointwise_generic", 2),
+        ]
+
+        for op_name, category, expected_tp in representative_ops:
+            with self.subTest(op=op_name, category=category):
+
+                def run(name=op_name, expected=expected_tp):
+                    stats = compare_operator(
+                        name,
+                        device="cpu",
+                        dtype=torch.float32,
+                        world_size=self.world_size,
+                        incorrect_only=True,
+                    )
+                    self.assertEqual(
+                        stats.true_positives,
+                        expected,
+                        f"{name}: expected {expected} true_positives, got {stats.true_positives}",
+                    )
+                    self.assertEqual(
+                        len(stats.false_positives),
+                        0,
+                        f"{name}: found incorrect rules: {stats.false_positives}",
+                    )
+
+                self._with_even_sizes(run)
+
+
 DistElementwiseOpsTestWithLocalTensor = create_local_tensor_test_class(
     DistElementwiseOpsTest, base_class=LocalDTensorOpTestBase
 )
diff --git a/test/distributed/tensor/test_random_ops.py b/test/distributed/tensor/test_random_ops.py
index 56cbf7261c736..0e3e7c7d53a59 100644
--- a/test/distributed/tensor/test_random_ops.py
+++ b/test/distributed/tensor/test_random_ops.py
@@ -212,11 +212,15 @@ def test_tp_model_meta_init(self):
             and isinstance(random._rng_tracker, OffsetBasedRNGTracker)
         )
         self.assertEqual(model.weight.device, device)
-        assert isinstance(model.weight, DTensor)
+        if not isinstance(model.weight, DTensor):
+            raise AssertionError(
+                f"Expected model.weight to be DTensor, got {type(model.weight)}"
+            )
 
         # gather all the shards to compare initialization results
         WORLD = torch.distributed.group.WORLD
-        assert WORLD is not None
+        if WORLD is None:
+            raise AssertionError("Expected WORLD to not be None")
         weight_local = model.weight.to_local()
         weight_gather = funcol.all_gather_tensor(
             weight_local,
@@ -269,11 +273,15 @@ def test_fsdp_tp_model_meta_init(self):
             and isinstance(random._rng_tracker, OffsetBasedRNGTracker)
         )
         self.assertEqual(model.weight.device, device)
-        assert isinstance(model.weight, DTensor)
+        if not isinstance(model.weight, DTensor):
+            raise AssertionError(
+                f"Expected model.weight to be DTensor, got {type(model.weight)}"
+            )
 
         # gather all the shards to compare initialization results
         WORLD = torch.distributed.group.WORLD
-        assert WORLD is not None
+        if WORLD is None:
+            raise AssertionError("Expected WORLD to not be None")
         weight_local = model.weight.to_local()
         weight_gather = funcol.all_gather_tensor(
             weight_local,
@@ -434,20 +442,19 @@ def test_pipeline_parallel_manual_seed(self):
         )
         torch.nn.init.normal_(spmd_dtensor)
 
-        # gather all the shards to compare initialization results
+        # verify the weights are initialized differently on all ranks
         WORLD = torch.distributed.group.WORLD
-        assert WORLD is not None
+        if WORLD is None:
+            raise AssertionError("Expected WORLD to not be None")
         tensor_gather = funcol.all_gather_tensor(
             spmd_dtensor.to_local(),
             gather_dim=0,
             group=WORLD,
-        )
-
-        # verify the weights are initialized differently on all ranks
+        ).wait()
         for other_rank in range(self.world_size):
             if self.rank != other_rank:
                 self.assertNotEqual(
-                    spmd_dtensor,
+                    spmd_dtensor.to_local(),
                     tensor_gather[2 * other_rank : 2 * (other_rank + 1), :],
                 )
 
@@ -582,7 +589,8 @@ def test_deterministic_uniform_2d(self):
         ]
 
         coordinate = device_mesh.get_coordinate()
-        assert coordinate is not None
+        if coordinate is None:
+            raise AssertionError("Expected coordinate to not be None")
 
         for placements, shard_index in zip(placements_list, shard_index_list):
             dtensor = dtensor.redistribute(device_mesh, placements)
@@ -691,6 +699,222 @@ def test_philox_state_seed_roundtrip(self):
         philox.seed = philox.seed.clone()
 
 
+class DistTensorRandomOpCompileTest(DTensorTestBase):
+    def _run_with_seed(self, fn, create_input, num_runs):
+        """Run fn num_runs times after resetting RNG, returning results and states."""
+        torch.manual_seed(0)
+        results = []
+        rng_states = [torch.cuda.get_rng_state()]
+        for _ in range(num_runs):
+            x = create_input()
+            result = fn(x)
+            results.append(result.to_local().clone())
+            rng_states.append(torch.cuda.get_rng_state())
+        # verify RNG state advances after each call
+        for i in range(len(rng_states) - 1):
+            self.assertFalse(
+                torch.equal(rng_states[i], rng_states[i + 1]),
+                f"RNG state did not change between call {i} and {i + 1}",
+            )
+        return results, rng_states
+
+    def _run_eager_and_compiled(self, fn, create_input, num_runs):
+        """Run fn both eagerly and compiled with aot_eager, returning results
+        and RNG states. Verifies the graph contains run_dtensor_rng_op."""
+        from torch._dynamo.testing import AotEagerAndRecordGraphs
+
+        eager_results, eager_rng_states = self._run_with_seed(
+            fn, create_input, num_runs
+        )
+        backend = AotEagerAndRecordGraphs()
+        compiled_fn = torch.compile(fn, backend=backend, fullgraph=True)
+        compiled_results, compiled_rng_states = self._run_with_seed(
+            compiled_fn, create_input, num_runs
+        )
+        self.assertIn("run_dtensor_rng_op", backend.fw_graphs[0].code)
+        return eager_results, eager_rng_states, compiled_results, compiled_rng_states
+
+    def _assert_eager_compiled_match(
+        self, eager_results, eager_rng_states, compiled_results, compiled_rng_states
+    ):
+        """Assert per-iteration results and RNG states match between eager and compiled."""
+        for i in range(len(eager_results)):
+            self.assertEqual(eager_results[i], compiled_results[i])
+            self.assertEqual(
+                eager_rng_states[i + 1],
+                compiled_rng_states[i + 1],
+                f"RNG state mismatch between eager and compiled after call {i}",
+            )
+
+    def _assert_replicate_cross_rank_equal(self, results, device_mesh):
+        """Assert all ranks produced identical results (for Replicate placement)."""
+        for i in range(len(results)):
+            local_result = results[i]
+            gathered = funcol.all_gather_tensor(
+                local_result, gather_dim=0, group=(device_mesh, 0)
+            ).wait()
+            local_size = local_result.shape[0]
+            for other_rank in range(self.world_size):
+                if self.rank != other_rank:
+                    self_slice = slice(
+                        self.rank * local_size, (self.rank + 1) * local_size
+                    )
+                    other_slice = slice(
+                        other_rank * local_size, (other_rank + 1) * local_size
+                    )
+                    self.assertEqual(
+                        gathered[self_slice],
+                        gathered[other_slice],
+                    )
+
+    def _test_compile_random_op(
+        self, fn, device_mesh, create_input=None, num_runs=3, placements=None
+    ):
+        """Run fn eager and compiled num_runs times, assert i-th results match
+        and graph contains run_dtensor_rng_op. Tests both aot_eager and inductor
+        backends."""
+        if placements is None:
+            placements = [Shard(0)]
+
+        if create_input is None:
+
+            def create_input():
+                return torch.distributed.tensor.ones(
+                    (8, 8), device_mesh=device_mesh, placements=placements
+                )
+
+        # Test with aot_eager backend (also verifies graph contents)
+        eager_results, eager_rng_states, compiled_results, compiled_rng_states = (
+            self._run_eager_and_compiled(fn, create_input, num_runs)
+        )
+        self._assert_eager_compiled_match(
+            eager_results, eager_rng_states, compiled_results, compiled_rng_states
+        )
+        if all(isinstance(p, Replicate) for p in placements):
+            self._assert_replicate_cross_rank_equal(compiled_results, device_mesh)
+
+        # Test with inductor backend
+        torch._dynamo.reset()
+        inductor_fn = torch.compile(fn, backend="inductor", fullgraph=True)
+        inductor_results, inductor_rng_states = self._run_with_seed(
+            inductor_fn, create_input, num_runs
+        )
+        self._assert_eager_compiled_match(
+            eager_results, eager_rng_states, inductor_results, inductor_rng_states
+        )
+
+    @with_comms
+    @skip_unless_torch_gpu
+    def test_compile_native_dropout(self):
+        device_mesh = self.build_device_mesh()
+
+        def fn(x):
+            return torch.nn.functional.dropout(x, p=0.5, training=True)
+
+        for placements in ([Shard(0)], [Replicate()]):
+            self._test_compile_random_op(fn, device_mesh, placements=placements)
+
+    @with_comms
+    @skip_unless_torch_gpu
+    def test_compile_normal_(self):
+        device_mesh = self.build_device_mesh()
+
+        def fn(x):
+            return x.normal_()
+
+        for placements in ([Shard(0)], [Replicate()]):
+            self._test_compile_random_op(fn, device_mesh, placements=placements)
+
+    @with_comms
+    @skip_unless_torch_gpu
+    def test_compile_rand_like(self):
+        device_mesh = self.build_device_mesh()
+
+        def fn(x):
+            return torch.rand_like(x)
+
+        for placements in ([Shard(0)], [Replicate()]):
+            self._test_compile_random_op(fn, device_mesh, placements=placements)
+
+    @with_comms
+    @skip_unless_torch_gpu
+    def test_compile_randn_like(self):
+        device_mesh = self.build_device_mesh()
+
+        def fn(x):
+            return torch.randn_like(x)
+
+        for placements in ([Shard(0)], [Replicate()]):
+            self._test_compile_random_op(fn, device_mesh, placements=placements)
+
+    @with_comms
+    @skip_unless_torch_gpu
+    def test_compile_randint_like(self):
+        device_mesh = self.build_device_mesh()
+
+        def fn(x):
+            return torch.randint_like(x, 0, 10)
+
+        for placements in ([Shard(0)], [Replicate()]):
+            self._test_compile_random_op(fn, device_mesh, placements=placements)
+
+    @with_comms
+    @skip_unless_torch_gpu
+    def test_compile_uniform_(self):
+        device_mesh = self.build_device_mesh()
+
+        def fn(x):
+            return x.uniform_(0.0, 1.0)
+
+        for placements in ([Shard(0)], [Replicate()]):
+            self._test_compile_random_op(fn, device_mesh, placements=placements)
+
+    @with_comms
+    @skip_unless_torch_gpu
+    def test_compile_bernoulli(self):
+        device_mesh = self.build_device_mesh()
+
+        def fn(x):
+            return torch.bernoulli(x)
+
+        for placements in ([Shard(0)], [Replicate()]):
+
+            def create_input(placements=placements):
+                return distribute_tensor(
+                    torch.full((8, 8), 0.5, device=self.device_type),
+                    device_mesh,
+                    placements,
+                )
+
+            self._test_compile_random_op(
+                fn, device_mesh, create_input=create_input, placements=placements
+            )
+
+    @with_comms
+    @skip_unless_torch_gpu
+    def test_compile_bernoulli_float(self):
+        device_mesh = self.build_device_mesh()
+
+        def fn(x):
+            return x.bernoulli_(0.5)
+
+        for placements in ([Shard(0)], [Replicate()]):
+            self._test_compile_random_op(fn, device_mesh, placements=placements)
+
+    @with_comms
+    @skip_unless_torch_gpu
+    def test_compile_multiple_random_ops(self):
+        device_mesh = self.build_device_mesh()
+
+        def fn(x):
+            x = x.uniform_(0, 1)
+            x = torch.nn.functional.dropout(x, p=0.5)
+            return x
+
+        for placements in ([Shard(0)], [Replicate()]):
+            self._test_compile_random_op(fn, device_mesh, placements=placements)
+
+
 class DistTensorRandomOpsTest3D(DTensorTestBase):
     @property
     def world_size(self):
@@ -730,11 +954,15 @@ def test_hsdp_tp_model_meta_init(self):
             and isinstance(random._rng_tracker, OffsetBasedRNGTracker)
         )
         self.assertEqual(model.weight.device, device)
-        assert isinstance(model.weight, DTensor)
+        if not isinstance(model.weight, DTensor):
+            raise AssertionError(
+                f"Expected model.weight to be DTensor, got {type(model.weight)}"
+            )
 
         # gather all the shards to compare initialization results
         WORLD = torch.distributed.group.WORLD
-        assert WORLD is not None
+        if WORLD is None:
+            raise AssertionError("Expected WORLD to not be None")
         weight_local = model.weight.to_local()
         weight_gather = funcol.all_gather_tensor(
             weight_local,
@@ -777,6 +1005,10 @@ def compute_rankwise_if_local_tensor(weight_local, rank):
 
 DistTensorRandomOpTestWithLocalTensor = create_local_tensor_test_class(
     DistTensorRandomOpTest,
+    skipped_tests=[
+        # cross-pp-stage seeding is not simulated in local tensor mode
+        "test_pipeline_parallel_manual_seed",
+    ],
 )
 
 DistTensorRandomOpsTest3DWithLocalTensor = create_local_tensor_test_class(
diff --git a/test/distributed/tensor/test_redistribute.py b/test/distributed/tensor/test_redistribute.py
index 22912edde1031..fc019b5310962 100644
--- a/test/distributed/tensor/test_redistribute.py
+++ b/test/distributed/tensor/test_redistribute.py
@@ -52,26 +52,24 @@
 )
 from torch.testing._internal.distributed._tensor.common_dtensor import (
     create_local_tensor_test_class,
-    DTensorTestBase,
+    DTensorContinuousTestBase,
     generate_shard_orders,
+    LocalDTensorContinuousTestBase,
     make_full_tensor,
     map_local_tensor_for_rank,
     patched_distribute_tensor as _distribute_tensor,
     redistribute,
-    with_comms,
 )
+from torch.testing._internal.distributed.fake_pg import FakeStore
 from torch.utils._debug_mode import DebugMode
 
 
 funcol = torch.ops.c10d_functional
 
 
-class RedistributeTest(DTensorTestBase):
-    @property
-    def world_size(self):
-        return 4
+class RedistributeTest(DTensorContinuousTestBase):
+    world_size = 4
 
-    @with_comms
     @parametrize("dtype", [torch.float32, torch.cfloat])
     def test_shard_to_replicate_forward_backward(self, dtype):
         # 1) test shard -> replicate forward
@@ -115,7 +113,6 @@ def test_shard_to_replicate_forward_backward(self, dtype):
             )
             self.assertEqual(comm_mode.get_total_counts(), 0)
 
-    @with_comms
     def test_replicate_to_replicate_forward_backward(self):
         device_mesh = self.build_device_mesh()
         replica_spec = [Replicate()]
@@ -143,7 +140,6 @@ def test_replicate_to_replicate_forward_backward(self):
         self.assertEqual(grad_input.to_local(), torch.ones(12, 3))
         self.assertEqual(comm_mode.get_total_counts(), 0)
 
-    @with_comms
     @parametrize("dtype", [torch.float32, torch.cfloat])
     def test_partial_to_partial_forward_backward(self, dtype):
         device_mesh = self.build_device_mesh()
@@ -184,7 +180,6 @@ def test_partial_to_partial_forward_backward(self, dtype):
         # Verify no communication in backward
         self.assertEqual(comm_mode.get_total_counts(), 0)
 
-    @with_comms
     @parametrize("dtype", [torch.float32, torch.cfloat])
     def test_replicate_to_local_partial_grad(self, dtype):
         device_mesh = self.build_device_mesh()
@@ -206,7 +201,6 @@ def test_replicate_to_local_partial_grad(self, dtype):
         self.assertEqual(comm_mode.get_total_counts(), 1)
         self.assertEqual(comm_mode.get_comm_counts()[funcol.all_reduce], 1)
 
-    @with_comms
     def test_replicate_to_shard_forward_backward(self):
         device_mesh = self.build_device_mesh()
         replica_spec = [Replicate()]
@@ -256,7 +250,6 @@ def test_replicate_to_shard_forward_backward(self):
                 comm_mode.get_comm_counts()[funcol.all_gather_into_tensor], 1
             )
 
-    @with_comms
     @parametrize("dtype", [torch.float32, torch.cfloat])
     def test_partial_to_replicate_forward_backward(self, dtype):
         # Although we don't allow user to reshard to produce a partial
@@ -296,7 +289,6 @@ def test_partial_to_replicate_forward_backward(self, dtype):
         )
         self.assertEqual(comm_mode.get_total_counts(), 0)
 
-    @with_comms
     def test_replicate_to_replicate_forward_backward_datatype_conversion(self):
         device_mesh = self.build_device_mesh()
         replica_spec = [Replicate()]
@@ -353,7 +345,6 @@ def test_replicate_to_replicate_forward_backward_datatype_conversion(self):
             self.assertEqual(grad_input.to_local(), torch.ones(12, 3))
             self.assertEqual(comm_mode.get_total_counts(), 0)
 
-    @with_comms
     def test_shard_to_replicate_forward_backward_datatype_conversion(self):
         device_mesh = self.build_device_mesh()
         replica_spec = [Replicate()]
@@ -425,7 +416,6 @@ def test_shard_to_replicate_forward_backward_datatype_conversion(self):
                 )
                 self.assertEqual(comm_mode.get_total_counts(), 0)
 
-    @with_comms
     def test_replicate_to_partial(self):
         device_mesh = self.build_device_mesh()
         local_tensor = torch.randn(12, 3, device=self.device_type, requires_grad=True)
@@ -473,7 +463,6 @@ def test_replicate_to_partial(self):
         )
         self.assertEqual(comm_mode.get_total_counts(), 0)
 
-    @with_comms
     @parametrize("dtype", [torch.float32, torch.cfloat])
     def test_partial_to_shard(self, dtype):
         device_mesh = self.build_device_mesh()
@@ -581,7 +570,6 @@ def _test_all_gather_optimization(
         expected_dt = distribute_tensor(global_tensor, device_mesh, placements_dst)
         self.assertEqual(result.to_local(), expected_dt.to_local())
 
-    @with_comms
     def test_all_gather_view_optimization_batch1(self):
         """
         Test that all_gather with gather_dim=1 and batch=1 uses a pure view
@@ -597,7 +585,6 @@ def test_all_gather_view_optimization_batch1(self):
             should_use_view=True,
         )
 
-    @with_comms
     def test_all_gather_no_optimization_gather_dim_not_1(self):
         """
         Test that all_gather with batch=1 but gather_dim != 1 still falls back
@@ -613,7 +600,6 @@ def test_all_gather_no_optimization_gather_dim_not_1(self):
             should_use_view=False,
         )
 
-    @with_comms
     def test_all_gather_split_cat_fallback_batch_gt_1(self):
         """
         Test that all_gather with gather_dim=1 and batch>1 falls back to
@@ -629,7 +615,6 @@ def test_all_gather_split_cat_fallback_batch_gt_1(self):
             should_use_view=False,
         )
 
-    @with_comms
     def test_redistribute_negative_shard_dim(self):
         device_mesh = self.build_device_mesh()
         local_tensor = torch.randn(12, 3, device=self.device_type, requires_grad=True)
@@ -641,7 +626,6 @@ def test_redistribute_negative_shard_dim(self):
         reshard_tensor = shard_tensor.redistribute(device_mesh, shard_minus_spec)
         self.assertEqual(reshard_tensor.placements[0].dim, 1)
 
-    @with_comms
     def test_redistribute_uneven_sharding(self):
         mesh = DeviceMesh(self.device_type, torch.arange(self.world_size).reshape(2, 2))
         data_to_test = [
@@ -666,7 +650,6 @@ def test_redistribute_uneven_sharding(self):
                 self.assertEqual(dt_full_tensor, input_tensor)
 
     @skip_if_lt_x_gpu(4)
-    @with_comms
     @parametrize("dtype", [torch.float32, torch.cfloat])
     def test_redistribute_shard_dim_change(self, dtype):
         # test 1d device mesh
@@ -767,7 +750,6 @@ def test_redistribute_shard_dim_change(self, dtype):
                 local_expected_dt = expected_dt.to_local()
                 self.assertEqual(local_out_dt, local_expected_dt)
 
-    @with_comms
     @parametrize("dtype", [torch.float32, torch.cfloat])
     def test_shard_dim_alltoall(self, dtype):
         # init 2d mesh here so we can test when group_rank != global_rank
@@ -781,7 +763,6 @@ def test_shard_dim_alltoall(self, dtype):
         self.assertEqual(new_tensor.shape, new_meta_tensor.shape)
         self.assertEqual(new_tensor.stride(), new_meta_tensor.stride())
 
-    @with_comms
     def test_one_chunk_mesh(self):
         # mesh size is 1 on second dim
         mesh = init_device_mesh(self.device_type, (4, 1))
@@ -801,7 +782,6 @@ def test_one_chunk_mesh(self):
             self.assertEqual(comm_mode.get_total_counts(), 0)
             self.assertEqual(out.placements, [Shard(0), dst])
 
-    @with_comms
     def test_redistribute_to_partial(self):
         mesh = init_device_mesh(self.device_type, (2, 2))
 
@@ -833,7 +813,6 @@ def test_redistribute_to_partial(self):
                 out = dt.redistribute(mesh, dst)
                 self.assertEqual(out.placements, dst)
 
-    @with_comms
     def test_replicate_to_partial_different_reduce_ops(self):
         """
         Test that Replicate -> Partial transitions work for all reduce op types.
@@ -880,7 +859,6 @@ def test_replicate_to_partial_different_reduce_ops(self):
                     partial_tensor.to_local(),
                 )
 
-    @with_comms
     def test_replicate_to_partial_planner_reduce_op_collection(self):
         """
         Test that the redistribution planner correctly collects reduce ops from
@@ -984,7 +962,6 @@ def test_replicate_to_partial_planner_reduce_op_collection(self):
             "Planner should collect all reduce ops from both src and dst",
         )
 
-    @with_comms
     def test_redistribute_zero_size_shards(self):
         # Adding this test to ensure sharding works when tensor size < mesh size.
         # Tests correct redistribution with world size = 4, tensor dim size = 2,
@@ -1012,16 +989,74 @@ def test_redistribute_zero_size_shards(self):
         self.assertEqual(dt_shard1._local_tensor.shape, (2, 2))
         self.assertEqual(dt_shard1.full_tensor(), dt.full_tensor())
 
+    @parametrize("dtype", [torch.float32])
+    def test_shard_to_replicate_forward_backward_backward(self, dtype):
+        # 1) test shard -> replicate forward
+        device_mesh = self.build_device_mesh()
+        replica_spec = [Replicate()]
+
+        input_sizes_and_shard_dim = [
+            ((self.world_size * 3, 3), 0),
+            ((self.world_size * 3 + 1, 3), 0),
+            ((self.world_size * 3 + 2, 3), 0),
+            ((3, self.world_size * 3), 1),
+            ((3, self.world_size * 3 + 1), 1),
+            ((3, self.world_size * 3 + 2), 1),
+        ]
+
+        activation = torch.nn.ReLU()
+
+        comm_mode = CommDebugMode()
+        for input_size, shard_dim in input_sizes_and_shard_dim:
+            shard_spec = [Shard(shard_dim)]
+            weight_ref = torch.nn.Parameter(torch.randn([]))
+            weight = distribute_tensor(weight_ref, device_mesh, replica_spec)
+            expected_tensor = torch.randn(
+                input_size, device=self.device_type, requires_grad=True, dtype=dtype
+            )
+            dtensor = distribute_tensor(expected_tensor, device_mesh, shard_spec)
+            dtensor_act = activation(dtensor * weight)
+            expected_tensor_act = torch.square(activation(expected_tensor * weight_ref))
+            with comm_mode:
+                reshard_dtensor_act = torch.square(dtensor_act.full_tensor())
+            self.assertEqual(reshard_dtensor_act.size(), torch.Size(input_size))
+            self.assertEqual(expected_tensor_act, reshard_dtensor_act)
+            self.assertEqual(
+                comm_mode.get_comm_counts()[funcol.all_gather_into_tensor], 1
+            )
+
+            # 2) test shard -> replicate backward:
+            output_dist = reshard_dtensor_act.mean()
+            output_ref = expected_tensor_act.mean()
+            self.assertEqual(output_ref, output_dist)
+            grad_input_ref = torch.autograd.grad(
+                outputs=output_ref, inputs=expected_tensor, create_graph=True
+            )[0]
+            with comm_mode:
+                grad_input = torch.autograd.grad(
+                    outputs=output_dist, inputs=dtensor, create_graph=True
+                )[0].full_tensor()
+            self.assertEqual(
+                comm_mode.get_comm_counts()[funcol.all_gather_into_tensor], 1
+            )
+            self.assertEqual(grad_input, grad_input_ref)
+
+            # 3) test grad weight for shard -> replicate double-backward:
+            grad_input_ref.mean().backward()
+            with comm_mode:
+                grad_input.mean().backward()
+            self.assertEqual(
+                comm_mode.get_comm_counts()[funcol.all_gather_into_tensor], 1
+            )
+            self.assertEqual(weight_ref.grad, weight.grad.full_tensor())
+
 
 instantiate_parametrized_tests(RedistributeTest)
 
 
-class MultiDimRedistributeTest(DTensorTestBase):
-    @property
-    def world_size(self) -> int:
-        return 8
+class MultiDimRedistributeTest(DTensorContinuousTestBase):
+    world_size = 8
 
-    @with_comms
     def test_multi_dim_mesh(self):
         devices = torch.arange(self.world_size)
         for mesh_shape in [devices, devices.view(4, 2), devices.view(2, 2, 2)]:
@@ -1070,7 +1105,6 @@ def test_multi_dim_mesh(self):
                         expected = num_sums * full_tensor
                         self.assertEqual(local_full, expected)
 
-    @with_comms
     def test_redistribute_shard_dim_multi_dim_mesh(self):
         mesh = init_device_mesh(self.device_type, (2, 2, 2))
         input_data = torch.randn((8, 8, 8), device=self.device_type)
@@ -1108,10 +1142,8 @@ def test_redistribute_shard_dim_multi_dim_mesh(self):
             self.assertEqual(local_out_dt, local_expected_dt)
 
 
-class DistributeWithDeviceOrderTest(DTensorTestBase):
-    @property
-    def world_size(self) -> int:
-        return 8
+class DistributeWithDeviceOrderTest(DTensorContinuousTestBase):
+    world_size = 8
 
     def _extract_redistribute_trace_from_debug_mode(self, s: str) -> str:
         import re
@@ -1123,7 +1155,6 @@ def _extract_redistribute_trace_from_debug_mode(self, s: str) -> str:
         else:
             return ""
 
-    @with_comms
     def test_ordered_redistribute(self):
         """Test ordered redistribution with various sharding syntaxes"""
         torch.manual_seed(21)
@@ -1210,7 +1241,6 @@ def test_ordered_redistribute(self):
             )
             self.assertEqual(sharded_dt.to_local(), expected_dt.to_local())
 
-    @with_comms
     def test_force_min_cost_redistribution_plan(self):
         """
         Test that the disable_graph_based_transform context manager correctly controls
@@ -1279,7 +1309,6 @@ def test_force_min_cost_redistribution_plan(self):
             # the second iteration would use cached paths from the first
             _gen_transform_infos.cache_clear()
 
-    @with_comms
     def test_generate_shard_orders(self):
         """Check if `generate_shard_orders` generates unique sharding combinations"""
         import math
@@ -1297,9 +1326,10 @@ def test_generate_shard_orders(self):
                 all_combinations.append(shard_order)  # noqa: PERF402
             for i in range(len(all_combinations)):
                 for j in range(i + 1, len(all_combinations)):
-                    assert all_combinations[i] != all_combinations[j], (
-                        f"Duplicate elements found in all_combinations {all_combinations[i]}, {all_combinations[j]}"
-                    )
+                    if all_combinations[i] == all_combinations[j]:
+                        raise AssertionError(
+                            f"Duplicate elements found in all_combinations {all_combinations[i]}, {all_combinations[j]}"
+                        )
             expected_total_combination = 0
             N = test_input["mesh"].ndim
             M = test_input["tensor_rank"]
@@ -1318,7 +1348,6 @@ def test_generate_shard_orders(self):
             expected_total_combination *= math.factorial(N)
             self.assertEqual(len(all_combinations), expected_total_combination)
 
-    @with_comms
     def test_ordered_distribute_all_combination(self):
         """Exhaustively test all possible sharding combinations and verify correctness"""
         torch.manual_seed(21)
@@ -1371,7 +1400,6 @@ def test_ordered_distribute_all_combination(self):
                     self.assertEqual(make_full_tensor(sharded_dt), input_data)
                     prev_sharded_dt = sharded_dt
 
-    @with_comms
     def test_graph_based_redistribute_cost(self):
         """
         This test verifies the correctness of
@@ -1442,7 +1470,6 @@ def test_graph_based_redistribute_cost(self):
                             f"{tensor_shape=}, {src_order=}, {dst_order=}, {intermediate_order=}",
                         )
 
-    @with_comms
     def test_redistribute_partial_to_different_partial_not_supported(self):
         # Test that redistributing from one Partial type to another raises an error
         device_mesh = self.build_device_mesh()
@@ -1482,7 +1509,6 @@ def test_redistribute_partial_to_different_partial_not_supported(self):
         "Temporarily skipping until we support special placement types in "
         "graph based redistribution"
     )
-    @with_comms
     def test_ordered_redistribute_for_special_placement(self):
         """Test ordered redistribution with special placement"""
         torch.manual_seed(21)
@@ -1500,7 +1526,6 @@ def test_ordered_redistribute_for_special_placement(self):
         )
         sharded_dt = redistribute(sharded_dt, mesh, tgt_placement, shard_order=None)
 
-    @with_comms
     def test_shard_order_same_data_as_strided_shard(self):
         device_mesh = init_device_mesh(self.device_type, (4, 2))
         x = torch.randn(8, 4, device=self.device_type)
@@ -1516,7 +1541,6 @@ def test_shard_order_same_data_as_strided_shard(self):
         )
         self.assertEqual(x_ordered_dt.to_local(), x_strided_dt.to_local())
 
-    @with_comms
     def test_ordered_all_gather_with_flattening(self):
         """Test that flattened all_gather produces correct results with non-ascending shard order.
 
@@ -1597,7 +1621,6 @@ def test_ordered_all_gather_with_flattening(self):
         self.assertEqual(ascending_result.to_local(), input_data)
         self.assertEqual(non_ascending_result.to_local(), input_data)
 
-    @with_comms
     def test_debug_mode_shows_optimized_trace(self):
         """Test that DebugMode captures the optimized (flattened) redistribution trace.
 
@@ -1665,10 +1688,8 @@ def test_debug_mode_shows_optimized_trace(self):
         self.assertEqual(result.to_local(), input_data)
 
 
-class DistributeWithStridedShardTest(DTensorTestBase):
-    @property
-    def world_size(self) -> int:
-        return 8
+class DistributeWithStridedShardTest(DTensorContinuousTestBase):
+    world_size = 8
 
     def _extract_redistribute_trace_from_debug_mode(self, s: str) -> str:
         import re
@@ -1680,11 +1701,31 @@ def _extract_redistribute_trace_from_debug_mode(self, s: str) -> str:
         else:
             return ""
 
-    @with_comms
     def test_strided_shard_redistribution(self):
         torch.manual_seed(21)
-        with maybe_disable_local_tensor_mode():
-            mesh = init_device_mesh(self.device_type, (2, 2, 2))
+        # 1D mesh: _StridedShard -> Replicate uses the greedy path's 1D shortcut
+        mesh_1d = init_device_mesh(self.device_type, (self.world_size,))
+        input_1d = torch.randn((16, 13), device=self.device_type)
+        src_dt = distribute_tensor(
+            input_1d, mesh_1d, [_StridedShard(0, split_factor=2)]
+        )
+        # _StridedShard -> Replicate via redistribute
+        result_dt = src_dt.redistribute(mesh_1d, [Replicate()])
+        self.assertEqual(result_dt.to_local(), input_1d)
+        # Same code path as above, but exercises the user-facing API
+        self.assertEqual(src_dt.full_tensor(), input_1d)
+
+        # 2D mesh: _StridedShard with undecodable split_factor falls back to greedy
+        mesh_2d = init_device_mesh(self.device_type, (4, 2))
+        input_2d = torch.randn((24, 8), device=self.device_type)
+        src_dt = distribute_tensor(
+            input_2d, mesh_2d, [_StridedShard(0, split_factor=3), Replicate()]
+        )
+        result_dt = src_dt.redistribute(mesh_2d, [Replicate(), Replicate()])
+        self.assertEqual(result_dt.to_local(), input_2d)
+
+        # 3D mesh cases
+        mesh = init_device_mesh(self.device_type, (2, 2, 2))
         input_data = torch.randn((31, 13, 11), device=self.device_type)
         sharding_src_dst_pairs_with_expected_trace = [
             (
@@ -1764,7 +1805,7 @@ def test_strided_shard_redistribution(self):
             elif idx == 3:
                 self.assertExpectedInline(
                     trace_str,
-                    """S(0)[0]S(0)[1]R->S(0)RR->S(0)_S(1, 5)R->R_S(1, 5)R->R_S(1, 5)S(0)""",
+                    """S(0)[0]S(0)[1]R->S(0)[0]S(0)[1]_S(1, 5)->S(0)R_S(1, 5)->S(0)RR->S(0)_S(1, 5)R->S(0)_S(1, 5)[0]_S(1, 5)[1]->R_S(1, 5)[0]_S(1, 5)[1]->R_S(1, 5)R->R_S(1, 5)S(0)""",  # noqa: B950
                 )
             expected_dt = _distribute_tensor(
                 input_data.clone(),
@@ -1775,6 +1816,137 @@ def test_strided_shard_redistribution(self):
             )
             self.assertEqual(sharded_dt.to_local(), expected_dt.to_local())
 
+    def test_strided_shard_to_shard_redistribution(self):
+        torch.manual_seed(42)
+        mesh_1d = init_device_mesh(self.device_type, (self.world_size,))
+        mesh_2d = init_device_mesh(self.device_type, (4, 2))
+        mesh_3d = init_device_mesh(self.device_type, (2, 2, 2))
+        input_1d = torch.randn((16, 13), device=self.device_type)
+        input_2d = torch.randn((24, 8), device=self.device_type)
+        input_3d = torch.randn((31, 13, 11), device=self.device_type)
+
+        # _StridedShard <-> Shard on 1D, 2D, and 3D meshes
+        redistribute_pairs = [
+            # (mesh, input, src_placements, dst_placements)
+            # 1D: _StridedShard(0) -> Shard(0)
+            (mesh_1d, input_1d, [_StridedShard(0, split_factor=2)], [Shard(0)]),
+            # 1D: _StridedShard(0) -> Shard(1) (cross-dim)
+            (mesh_1d, input_1d, [_StridedShard(0, split_factor=2)], [Shard(1)]),
+            # 1D: Shard(0) -> _StridedShard(0)
+            (mesh_1d, input_1d, [Shard(0)], [_StridedShard(0, split_factor=2)]),
+            # 2D: [_StridedShard(0), Replicate()] -> [Shard(0), Replicate()]
+            (
+                mesh_2d,
+                input_2d,
+                [_StridedShard(0, split_factor=3), Replicate()],
+                [Shard(0), Replicate()],
+            ),
+            # 2D: [_StridedShard(0), Shard(1)] -> [Shard(0), Shard(1)]
+            (
+                mesh_2d,
+                input_2d,
+                [_StridedShard(0, split_factor=2), Shard(1)],
+                [Shard(0), Shard(1)],
+            ),
+            # 2D: [Shard(0), Replicate()] -> [_StridedShard(0), Replicate()]
+            (
+                mesh_2d,
+                input_2d,
+                [Shard(0), Replicate()],
+                [_StridedShard(0, split_factor=3), Replicate()],
+            ),
+            # 2D: [Shard(0), Shard(1)] -> [_StridedShard(0), Shard(1)]
+            (
+                mesh_2d,
+                input_2d,
+                [Shard(0), Shard(1)],
+                [_StridedShard(0, split_factor=2), Shard(1)],
+            ),
+        ]
+        for mesh, inp, src, dst in redistribute_pairs:
+            src_dt = distribute_tensor(inp, mesh, src)
+            result_dt = src_dt.redistribute(mesh, dst)
+            expected_dt = distribute_tensor(inp, mesh, dst)
+            self.assertEqual(result_dt.to_local(), expected_dt.to_local())
+            self.assertEqual(result_dt.full_tensor(), inp)
+
+        # 3D: _StridedShard on one dim -> Shard, others unchanged
+        src_dt = _distribute_tensor(
+            input_3d.clone(),
+            mesh_3d,
+            [Shard(0), Shard(0), _StridedShard(0, split_factor=3)],
+            shard_order=(ShardOrderEntry(tensor_dim=0, mesh_dims=(0, 1, 2)),),
+            src_data_rank=None,
+        )
+        result_dt = redistribute(
+            src_dt,
+            mesh_3d,
+            [Shard(0), Shard(0), Shard(0)],
+            shard_order=(ShardOrderEntry(tensor_dim=0, mesh_dims=(0, 1, 2)),),
+        )
+        expected_dt = _distribute_tensor(
+            input_3d.clone(),
+            mesh_3d,
+            [Shard(0), Shard(0), Shard(0)],
+            shard_order=(ShardOrderEntry(tensor_dim=0, mesh_dims=(0, 1, 2)),),
+            src_data_rank=None,
+        )
+        self.assertEqual(result_dt.to_local(), expected_dt.to_local())
+
+    def test_partial_to_strided_shard_redistribution(self):
+        torch.manual_seed(42)
+
+        # 1D mesh, Partial -> _StridedShard(0)
+        mesh_1d = init_device_mesh(self.device_type, (self.world_size,))
+        input_1d = torch.randn((16, 13), device=self.device_type)
+        src_dt = DTensor.from_local(
+            input_1d.clone(), mesh_1d, [Partial("sum")], run_check=False
+        )
+        result_dt = src_dt.redistribute(mesh_1d, [_StridedShard(0, split_factor=2)])
+        reduced = input_1d * self.world_size
+        expected_dt = distribute_tensor(
+            reduced, mesh_1d, [_StridedShard(0, split_factor=2)]
+        )
+        self.assertEqual(result_dt.to_local(), expected_dt.to_local())
+        self.assertEqual(result_dt.full_tensor(), reduced)
+
+        # 2D mesh (4x2), [Partial, Replicate()] -> [_StridedShard(0), Replicate()]
+        mesh_2d = init_device_mesh(self.device_type, (4, 2))
+        input_2d = torch.randn((24, 8), device=self.device_type)
+        src_dt = DTensor.from_local(
+            input_2d.clone(), mesh_2d, [Partial("sum"), Replicate()], run_check=False
+        )
+        result_dt = src_dt.redistribute(
+            mesh_2d, [_StridedShard(0, split_factor=3), Replicate()]
+        )
+        reduced_2d = input_2d * 4
+        expected_dt = distribute_tensor(
+            reduced_2d, mesh_2d, [_StridedShard(0, split_factor=3), Replicate()]
+        )
+        self.assertEqual(result_dt.to_local(), expected_dt.to_local())
+        self.assertEqual(result_dt.full_tensor(), reduced_2d)
+
+    def test_strided_shard_to_partial_raises(self):
+        torch.manual_seed(42)
+
+        # 1D mesh
+        mesh_1d = init_device_mesh(self.device_type, (self.world_size,))
+        input_1d = torch.randn((16, 13), device=self.device_type)
+        src_dt = distribute_tensor(
+            input_1d, mesh_1d, [_StridedShard(0, split_factor=2)]
+        )
+        with self.assertRaises(RuntimeError):
+            src_dt.redistribute(mesh_1d, [Partial("sum")])
+
+        # 2D mesh
+        mesh_2d = init_device_mesh(self.device_type, (4, 2))
+        input_2d = torch.randn((24, 8), device=self.device_type)
+        src_dt = distribute_tensor(
+            input_2d, mesh_2d, [_StridedShard(0, split_factor=3), Replicate()]
+        )
+        with self.assertRaises(RuntimeError):
+            src_dt.redistribute(mesh_2d, [Partial("sum"), Replicate()])
+
 
 class TransformInfoTest(TestCase):
     """Tests for _TransformInfo._comm_type_key method."""
@@ -1804,6 +1976,23 @@ def test_comm_type_key(self):
             info = _TransformInfo(0, placements, [8, 8])
             self.assertIsNone(info._comm_type_key())
 
+    def test_comm_type_key_strided_shard(self):
+        """_StridedShard must produce the same comm type keys as Shard."""
+        ss = _StridedShard(0, split_factor=2)
+        test_cases = [
+            ((Partial("sum"), ss), "reduce_scatter"),
+            ((ss, Replicate()), "all_gather"),
+            ((ss, Shard(1)), "all_to_all"),
+            ((ss, _StridedShard(1, split_factor=2)), "all_to_all"),
+        ]
+        for placements, expected_key in test_cases:
+            info = _TransformInfo(0, placements, [8, 8])
+            self.assertEqual(
+                info._comm_type_key(),
+                expected_key,
+                f"_StridedShard transform {placements} should map to '{expected_key}'",
+            )
+
 
 class OptimizeFlattenedReductionsTest(TestCase):
     """Tests for _optimize_transform_infos helper.
@@ -1811,14 +2000,16 @@ class OptimizeFlattenedReductionsTest(TestCase):
     Uses fake process group since these tests don't perform actual communications.
     """
 
-    def setUp(self):
-        super().setUp()
-        store = dist.HashStore()
-        dist.init_process_group(backend="fake", rank=0, world_size=8, store=store)
+    @classmethod
+    def setUpClass(cls):
+        super().setUpClass()
+        fake_store = FakeStore()
+        dist.init_process_group("fake", store=fake_store, rank=0, world_size=8)
 
-    def tearDown(self):
-        super().tearDown()
+    @classmethod
+    def tearDownClass(cls):
         dist.destroy_process_group()
+        super().tearDownClass()
 
     def test_no_flattened_mesh_returns_original(self):
         """When no flattened mesh exists, original transforms are returned with a warning."""
@@ -2418,18 +2609,15 @@ def test_disable_optimization_context_manager(self):
         self.assertIsInstance(result[0], _FlattenedTransformInfo)
 
 
-class MultiDimRedistributeOptimizationTest(DTensorTestBase):
+class MultiDimRedistributeOptimizationTest(DTensorContinuousTestBase):
     """Integration tests for multi-dimensional redistribute correctness and optimization.
 
     These tests verify numerical correctness of various redistribute patterns
     on multi-dimensional meshes, including the flattened mesh optimization.
     """
 
-    @property
-    def world_size(self):
-        return 8
+    world_size = 8
 
-    @with_comms
     def test_multi_dim_redistribute(self):
         """Test various redistribute patterns for correctness and comm count optimization.
 
@@ -2583,7 +2771,6 @@ def test_multi_dim_redistribute(self):
                 expected_full = global_tensor * num_partial_sums
                 self.assertEqual(result_full, expected_full)
 
-    @with_comms
     def test_mixed_sum_avg_partial_numerics(self):
         """Test numerical correctness for mixed Partial("sum") and Partial("avg").
 
@@ -2672,18 +2859,15 @@ def test_mixed_sum_avg_partial_numerics(self):
         )
 
 
-class FlattenedReductionIntegrationTest(DTensorTestBase):
+class FlattenedReductionIntegrationTest(DTensorContinuousTestBase):
     """Integration tests for flattened reduction optimization.
 
     These tests verify that redistribute actually performs fewer communications
     when flattened meshes are available.
     """
 
-    @property
-    def world_size(self):
-        return 8
+    world_size = 8
 
-    @with_comms
     def test_merging_reductions(self):
         """Tests various combinations in the same test function to avoid setup time"""
         mesh = init_device_mesh(
@@ -2779,7 +2963,7 @@ def test_merging_reductions(self):
         self.assertEqual(result.to_local(), expected)
 
 
-class UnevenFlattenedReduceScatterTest(DTensorTestBase):
+class UnevenFlattenedReduceScatterTest(DTensorContinuousTestBase):
     """Test for correctness of flattened reduce_scatter with uneven tensor dimensions.
 
     When redistributing Partial,Partial → Shard(dim),Shard(dim) on a 2D mesh,
@@ -2787,11 +2971,8 @@ class UnevenFlattenedReduceScatterTest(DTensorTestBase):
     optimization doesn't properly handle padding/unpadding.
     """
 
-    @property
-    def world_size(self) -> int:
-        return 6  # 2 x 3 mesh
+    world_size = 6  # 2 x 3 mesh
 
-    @with_comms
     def test_partial_partial_to_shard_shard_uneven(self):
         """Test that Partial,Partial -> Shard(0),Shard(0) produces correct results with uneven dims.
 
@@ -2867,19 +3048,23 @@ def test_partial_partial_to_shard_shard_uneven(self):
 
 RedistributeTestWithLocalTensor = create_local_tensor_test_class(
     RedistributeTest,
+    base_class=LocalDTensorContinuousTestBase,
 )
 
 MultiDimRedistributeTestWithLocalTensor = create_local_tensor_test_class(
     MultiDimRedistributeTest,
     skipped_tests=["test_multi_dim_mesh"],
+    base_class=LocalDTensorContinuousTestBase,
 )
 
 MultiDimRedistributeOptimizationTestWithLocalTensor = create_local_tensor_test_class(
     MultiDimRedistributeOptimizationTest,
+    base_class=LocalDTensorContinuousTestBase,
 )
 
 DistributeWithDeviceOrderTestWithLocalTensor = create_local_tensor_test_class(
     DistributeWithDeviceOrderTest,
+    base_class=LocalDTensorContinuousTestBase,
 )
 
 if __name__ == "__main__":
diff --git a/test/distributed/tensor/test_single_dim_strategy.py b/test/distributed/tensor/test_single_dim_strategy.py
index 49e97d39e09fb..472bf39f37ae9 100644
--- a/test/distributed/tensor/test_single_dim_strategy.py
+++ b/test/distributed/tensor/test_single_dim_strategy.py
@@ -1,7 +1,7 @@
 # Owner(s): ["oncall: distributed"]
 
 
-from itertools import permutations
+from itertools import chain, permutations, product
 from unittest.mock import patch
 
 import torch
@@ -24,10 +24,14 @@
 )
 from torch.distributed.tensor._ops._matrix_ops import mm_single_dim_strategy
 from torch.distributed.tensor._ops._pointwise_ops import (
-    single_mesh_dim_linear_pointwise_strategy,
+    _BINARY_ADDITIVE_RULES,
+    _common_pointwise_single_dim_strategy,
+    _MUL_RULES,
+    _UNARY_LINEAR_RULES,
 )
 from torch.distributed.tensor._ops._tensor_ops import cat_single_dim_strategy
 from torch.distributed.tensor._ops.single_dim_strategy import (
+    _dijkstra_expand_single_dim_strategy_to_mesh,
     _expand_single_dim_strategy_to_mesh,
     _fill_single_dim_strategy_placeholders,
     _get_num_tensor_inputs,
@@ -38,8 +42,13 @@
     register_single_dim_strategy,
 )
 from torch.distributed.tensor._ops.utils import expand_to_full_mesh_op_strategy
+from torch.distributed.tensor._redistribute import use_min_cost_redistribution_plan
 from torch.distributed.tensor._sharding_prop import _select_min_cost_strategy
-from torch.distributed.tensor.placement_types import _StridedShard, Placement
+from torch.distributed.tensor.placement_types import (
+    _MaskPartial,
+    _StridedShard,
+    Placement,
+)
 from torch.testing._internal.common_utils import run_tests, TestCase
 from torch.testing._internal.distributed.fake_pg import FakeStore
 
@@ -83,7 +92,7 @@ def _get_mm_specs(
 class TestExpandPlaceholder(TestCase):
     def setUp(self):
         super().setUp()
-        self.world_size = 8
+        self.world_size = 64
         store = FakeStore()
         dist.init_process_group(
             backend="fake", rank=0, world_size=self.world_size, store=store
@@ -126,8 +135,9 @@ def _test_op(op, *args, linearity=None, inplace=None):
             output_meta = [spec.tensor_meta for spec in out_spec.children]
 
             op_schema = OpSchema(op=op, args_schema=tuple(specs), kwargs_schema={})
-            strategy_fn = single_mesh_dim_linear_pointwise_strategy(
-                linearity=linearity or -1
+            extra_rules = _BINARY_ADDITIVE_RULES if linearity == 1 else None
+            strategy_fn = _common_pointwise_single_dim_strategy(
+                partial_extra_rules=extra_rules
             )
             expanded = _expand_single_dim_strategy_to_mesh(
                 mesh, op_schema, _SingleDimStrategyInfo(strategy_fn), output_meta
@@ -144,9 +154,8 @@ def _test_op(op, *args, linearity=None, inplace=None):
                 # so there should be exactly 1 strategy (the input placement)
                 self.assertEqual(len(strategy.children[0].strategies), 1)
             elif linearity == 1:
-                self.assertEqual(
-                    len(strategy.children[0].strategies), 125
-                )  # len([S(0), S(1), S(2), R, P]) ** 3 = 125
+                # See test_expand_foreach_add_to_3d_mesh for derivation of 634.
+                self.assertEqual(len(strategy.children[0].strategies), 634)
             else:
                 self.assertGreaterAlmostEqual(
                     len(strategy.children[0].strategies), 64
@@ -248,7 +257,9 @@ def _expand_foreach_add_list(
                 mesh,
                 op_schema,
                 _SingleDimStrategyInfo(
-                    single_mesh_dim_linear_pointwise_strategy(linearity=1)
+                    _common_pointwise_single_dim_strategy(
+                        partial_extra_rules=_BINARY_ADDITIVE_RULES
+                    )
                 ),
                 output_tensor_meta,
             )
@@ -257,12 +268,23 @@ def _expand_foreach_add_list(
                 op_schema.args_meta,
                 op_schema.kwargs_meta,
             )
-            assert isinstance(strategy, TupleStrategy)
+            self.assertIsInstance(strategy, TupleStrategy)
             return strategy
 
-        # Note: using sizes that are multiples of mesh sizes so every sharding option is valid,
-        # (S0, S1, R, Psum, Pavg) ** 3 = 125
-        expected_num_strategies = (125, 8)
+        # Sizes are multiples of mesh dims so every shard option is valid.
+        # Per mesh dim: 4 non-partial (S0,S1,S2,R) + 6 partial rules
+        # (1 Psum, 3 Pavg, 1 Pmax, 1 Pmin). Mixed-partial filter allows
+        # Psum+Pavg to coexist but rejects other mixes. By inclusion-exclusion
+        # on a 3D mesh (n=4 non-partial, partial rule counts by type):
+        #   no partials:    4^3                                    =  64
+        #   Psum only:      (4+1)^3 - 4^3                         =  61
+        #   Pavg only:      (4+3)^3 - 4^3                         = 279
+        #   Pmax only:      (4+1)^3 - 4^3                         =  61
+        #   Pmin only:      (4+1)^3 - 4^3                         =  61
+        #   Psum+Pavg mix:  (4+1+3)^3 - (4+3)^3 - (4+1)^3 + 4^3  = 108
+        #   total: 634
+        # For scalars (no Shard), n=1 (R only): same formula gives 139.
+        expected_num_strategies = (634, 139)
         # Test Replicate + Shard gives Shard
         inputs_a = [torch.empty((8, 8, 8))] * 2
         placements_a = [
@@ -276,30 +298,35 @@ def _expand_foreach_add_list(
         ]
         expected_output_placements = [
             (Shard(0), Replicate(), Shard(1)),
-            # P(avg) -> P(sum) is currently not supported, but could be in principle
-            (Partial("sum"), Partial("sum"), Replicate()),
+            # P(avg),P(avg),R rule lets dim 2 keep P(avg) from input_a, only
+            # redistributing input_b to R (cheaper than both inputs → R)
+            (Partial("sum"), Partial("sum"), Partial("avg")),
         ]
         tuple_strategy = _expand_foreach_add_list(
             inputs_a, inputs_b, placements_a, placements_b
         )
         self.assertEqual(len(tuple_strategy.children), 2)
         for child_i, child in enumerate(tuple_strategy.children):
-            assert isinstance(child, OpStrategy)
+            self.assertIsInstance(child, OpStrategy)
             self.assertEqual(len(child.strategies), expected_num_strategies[child_i])
 
-            # _select_min_cost_strategy can have multiple min-cost strategies,
-            # so just assert the expected placement has equal cost.
             def sum_cost(x):
                 return sum(sum(y) for y in x)
 
+            # Multiple strategies can produce the same output placement but
+            # with different input specs (and costs, including inf). Check
+            # that the cheapest strategy with the expected output ties with
+            # the overall minimum.
             expected = expected_output_placements[child_i]
             min_cost_strategy = _select_min_cost_strategy(child)
-            for strategy in child.strategies:
-                if strategy.output_spec.placements == expected:
-                    self.assertEqual(
-                        sum_cost(strategy.redistribute_cost),
-                        sum_cost(min_cost_strategy.redistribute_cost),
-                    )
+            min_cost = sum_cost(min_cost_strategy.redistribute_cost)
+            expected_costs = [
+                sum_cost(s.redistribute_cost)
+                for s in child.strategies
+                if s.output_spec.placements == expected
+            ]
+            self.assertTrue(len(expected_costs) > 0)
+            self.assertAlmostEqual(min(expected_costs), min_cost, places=5)
 
     def test_expand_cat_strategy_to_3d_mesh(self):
         mesh = DeviceMesh("cpu", mesh=torch.arange(8).reshape(2, 2, 2))
@@ -342,7 +369,7 @@ def _expand_cat(
             strategy = expanded_strategy_fn(
                 torch.ops.aten.cat.default, op_schema.args_meta, op_schema.kwargs_meta
             )
-            assert isinstance(strategy, OpStrategy)
+            self.assertIsInstance(strategy, OpStrategy)
             return strategy
 
         # Note: using sizes that are multiples of mesh sizes so every sharding option is valid,
@@ -456,7 +483,7 @@ def test_expand_matmul_like_strategy_to_3d_mesh(self):
         strategy = expanded_strategy_fn(
             torch.ops.aten.matmul.default, op_schema.args_meta, op_schema.kwargs_meta
         )
-        assert isinstance(strategy, OpStrategy)
+        self.assertIsInstance(strategy, OpStrategy)
 
         # For a 3D mesh with 8 single-dim strategies per mesh dim
         # (3 sharding + 4 per-input linearity + 1 implicit replicate),
@@ -468,7 +495,7 @@ def test_expand_matmul_like_strategy_to_3d_mesh(self):
         for op_spec in strategy.strategies:
             output_spec = op_spec.output_spec
             input_specs = op_spec.input_specs
-            assert input_specs is not None
+            self.assertIsNotNone(input_specs)
 
             # Verify tensor_meta is populated for output spec
             self.assertIsNotNone(
@@ -804,20 +831,21 @@ def test_expand_strategy_handles_symbolic_shapes(self):
         self.assertIsInstance(result, OpStrategy)
         self.assertGreater(len(result.strategies), 0)
 
-    def test_expand_to_full_mesh_filters_out_variant_strategies(self):
-        """Test that expand_to_full_mesh_op_strategy filters strategies for out= variant ops.
-        For out-variant ops like torch.mul(..., out=...), the output placement must
-        match the 'out' kwarg's placement. This test verifies that strategies with
-        mismatched output placements are filtered out.
+    def test_strategy_length_validation(self):
+        """Test that _PreparedSingleDimStrategy validates strategy length against
+        the op schema. Strategies must include placements for all outputs, args,
+        and tensor kwargs. A strategy missing the out kwarg placement should fail.
         """
+        from torch.distributed.tensor._ops.single_dim_strategy import (
+            _PreparedSingleDimStrategy,
+        )
+
         mesh = DeviceMesh("cpu", mesh=torch.arange(4))
         meta = TensorMeta(torch.Size([8, 8]), (8, 1), torch.float32)
 
-        # Create specs: args have Shard(0), out kwarg has Replicate
         arg_spec = DTensorSpec(mesh, (Shard(0),), meta)
-        out_spec = DTensorSpec(mesh, (Replicate(),), meta)
+        out_spec = DTensorSpec(mesh, (Shard(0),), meta)
 
-        # Create OpSchema for out-variant op (aten.mul.out)
         op_schema = OpSchema(
             op=torch.ops.aten.mul.out,
             args_schema=(
@@ -827,11 +855,49 @@ def test_expand_to_full_mesh_filters_out_variant_strategies(self):
             kwargs_schema={"out": OpStrategy([OpSpec(out_spec)])},
         )
 
-        # Define strategies: output can be Shard(0) or Replicate
-        # [output, input1, input2, out_kwarg]
+        # Strategy missing the out kwarg placement: [output, arg1, arg2] = 3
+        # but mul.out has 1 output + 3 inputs (2 args + 1 out kwarg) = 4 expected
+        def bad_strategy(op, args, kwargs):
+            return [[Shard(0), Shard(0), Shard(0)]]
+
+        with self.assertRaisesRegex(AssertionError, r"Strategy length 3 != expected 4"):
+            _PreparedSingleDimStrategy(bad_strategy, op_schema, meta)
+
+        # Strategy with correct length: [output, arg1, arg2, out_kwarg] = 4
+        def good_strategy(op, args, kwargs):
+            return [[Shard(0), Shard(0), Shard(0), Shard(0)]]
+
+        # Should not raise
+        prepared = _PreparedSingleDimStrategy(good_strategy, op_schema, meta)
+        self.assertEqual(prepared.num_outputs, 1)
+        self.assertEqual(prepared.num_inputs, 3)
+
+    def test_out_variant_partial_propagation(self):
+        """Test that partial rules work correctly for .out variant ops.
+
+        For mul.out with rule [P(sum), P(sum), R, P(sum)]
+        (output=P(sum), arg1=P(sum), arg2=R, out_kwarg=P(sum)),
+        the out kwarg should get P(sum) matching the output.
+        """
+        mesh = DeviceMesh("cpu", mesh=torch.arange(4))
+        meta = TensorMeta(torch.Size([8, 8]), (8, 1), torch.float32)
+
+        partial_spec = DTensorSpec(mesh, (Partial("sum"),), meta)
+        replicate_spec = DTensorSpec(mesh, (Replicate(),), meta)
+
+        op_schema = OpSchema(
+            op=torch.ops.aten.mul.out,
+            args_schema=(
+                OpStrategy([OpSpec(partial_spec)]),
+                OpStrategy([OpSpec(replicate_spec)]),
+            ),
+            kwargs_schema={"out": OpStrategy([OpSpec(partial_spec)])},
+        )
+
+        # [output, arg1, arg2, out_kwarg] — includes out kwarg placement
         single_mesh_dim_strategies = [
-            [Shard(0), Shard(0), Shard(0), Shard(0)],  # All sharded
-            [Replicate(), Replicate(), Replicate(), Replicate()],  # All replicated
+            [Partial("sum"), Partial("sum"), Replicate(), Partial("sum")],
+            [Replicate(), Replicate(), Replicate(), Replicate()],
         ]
 
         result = expand_to_full_mesh_op_strategy(
@@ -839,73 +905,65 @@ def test_expand_to_full_mesh_filters_out_variant_strategies(self):
             op_schema,
             single_mesh_dim_strategies,
             output_tensor_meta=meta,
+            input_index=1,
         )
 
-        # All strategies in result should have output placement matching out kwarg (Replicate)
+        self.assertIsInstance(result, OpStrategy)
+        # Should have strategies where output=P(sum) and out_kwarg=P(sum)
+        found_partial_strategy = False
         for strategy in result.strategies:
-            output_spec = strategy.output_spec
-            self.assertEqual(
-                output_spec.placements,
-                (Replicate(),),
-                f"Output placement {output_spec.placements} should match out kwarg placement (Replicate(),)",
-            )
+            if strategy.output_spec.placements == (Partial("sum"),):
+                found_partial_strategy = True
+                # The out kwarg (last input) must match the output
+                out_kwarg_placement = strategy.input_specs[2].placements
+                self.assertEqual(
+                    out_kwarg_placement,
+                    (Partial("sum"),),
+                    "out kwarg should be P(sum) matching the output, not R",
+                )
+        self.assertTrue(
+            found_partial_strategy,
+            "Expected a strategy with Partial('sum') output for mul.out",
+        )
 
     def test_expand_multi_output_strategy(self):
         """Test expanding single-dim strategies for multi-output ops.
 
-        This is a regression test for the fix where _insert_single_dim_replication_strategy
-        was hardcoded to assume 1 output, causing assertion errors for multi-output ops.
-        The bug was: input_specs(3) != strategies(1: 1 args + 0 kwargs)
-
-        The fix ensures:
-        1. Multi-output ops correctly expand strategies with num_outputs > 1
-        2. The replicate strategy has the correct number of placements (num_outputs + num_inputs)
-        3. All output specs are populated as a tuple with correct tensor_meta
+        Uses aten.topk (2 tensor outputs, 1 tensor input) to verify that
+        multi-output strategies are expanded correctly with proper output specs.
         """
         mesh = DeviceMesh("cpu", mesh=torch.arange(4))
 
-        # Create a batched matrix input: (batch=2, m=3, n=3)
         input_meta = TensorMeta(
-            shape=torch.Size([2, 3, 3]),
-            stride=(9, 3, 1),
+            shape=torch.Size([8, 4]),
+            stride=(4, 1),
             dtype=torch.float32,
         )
 
-        # Create output tensor_metas for a 3-output op
+        # topk returns (values, indices) — 2 outputs
         output_metas = (
-            TensorMeta(torch.Size([2, 3, 3]), (9, 3, 1), torch.float32),
-            TensorMeta(torch.Size([2, 3, 3]), (9, 3, 1), torch.float32),
-            TensorMeta(torch.Size([2, 3, 3]), (9, 3, 1), torch.float32),
+            TensorMeta(torch.Size([8, 2]), (2, 1), torch.float32),
+            TensorMeta(torch.Size([8, 2]), (2, 1), torch.int64),
         )
 
-        # Create input spec with Shard(0) on batch dim
         input_spec = DTensorSpec(
             mesh=mesh,
             placements=(Shard(0),),
             tensor_meta=input_meta,
         )
 
-        # Create OpSchema - use a placeholder op since we're providing our own strategy
         op_schema = OpSchema(
-            op=torch.ops.aten.abs.default,  # placeholder, not actually used
-            args_schema=(OpStrategy([OpSpec(input_spec)]),),
+            op=torch.ops.aten.topk.default,
+            args_schema=(OpStrategy([OpSpec(input_spec)]), 2, -1, True, True),
             kwargs_schema={},
         )
 
-        # Define a mock multi-output single-dim strategy function
-        # This simulates an op with 3 outputs and 1 input
-        # Using Partial for outputs to test a realistic scenario
+        # 2 outputs + 1 input = 3 placements per strategy
         def mock_multi_output_strategy(op, args_schema, kwargs_schema):
-            # Return strategies with 4 placements each (3 outputs + 1 input)
-            # Using Partial for outputs (common for reduction ops)
             return [
-                [Partial(), Partial(), Partial(), Shard(0)],
+                [Partial(), Partial(), Shard(0)],
             ]
 
-        # This would have crashed before the fix with:
-        # AssertionError: input_specs(3) != strategies(1: 1 args + 0 kwargs)
-        # because _insert_single_dim_replication_strategy created [R, R] (2 elements)
-        # instead of [R, R, R, R] (4 elements for 3 outputs + 1 input)
         expanded_strategy_fn = _expand_single_dim_strategy_to_mesh(
             mesh,
             op_schema,
@@ -913,38 +971,26 @@ def mock_multi_output_strategy(op, args_schema, kwargs_schema):
             output_metas,
         )
         strategy = expanded_strategy_fn(
-            torch.ops.aten.abs.default,
+            torch.ops.aten.topk.default,
             op_schema.args_meta,
             op_schema.kwargs_meta,
         )
 
-        # Strategy should be an OpStrategy
         self.assertIsInstance(strategy, OpStrategy)
         self.assertGreaterEqual(len(strategy.strategies), 1)
 
-        # Each OpSpec should have tuple output_spec with 3 elements (one per output)
         for op_spec in strategy.strategies:
-            # Access output_specs directly (it's a tuple for multi-output ops)
             output_specs = op_spec.output_specs
             self.assertIsInstance(
                 output_specs, tuple, "Multi-output op should have tuple output_specs"
             )
             self.assertEqual(
-                len(output_specs), 3, "Should have 3 output specs for 3-output op"
+                len(output_specs), 2, "Should have 2 output specs for topk"
             )
-
-            # Check that all output specs are valid DTensorSpecs with tensor_meta
             for i, out_spec in enumerate(output_specs):
                 self.assertIsNotNone(out_spec, f"Output {i} spec should not be None")
                 self.assertIsInstance(out_spec, DTensorSpec)
-                self.assertIsNotNone(
-                    out_spec.tensor_meta, f"Output {i} spec should have tensor_meta"
-                )
-                # Verify the tensor_meta shape matches what we provided
-                self.assertEqual(out_spec.tensor_meta.shape, torch.Size([2, 3, 3]))
 
-            # Check input specs - should have 1 input
-            self.assertIsNotNone(op_spec.input_specs)
             self.assertEqual(len(op_spec.input_specs), 1, "Should have 1 input tensor")
 
     def test_inplace_op_partial_input_raises_clear_error(self):
@@ -1007,6 +1053,796 @@ def mock_pointwise_strategy(op, args_schema, kwargs_schema):
                 op_schema.kwargs_meta,
             )
 
+    def test_expand_filters_mixed_partial_types(self):
+        """Test that expand_to_full_mesh_op_strategy filters out mixed partial types.
+
+        When single-dim strategies are expanded to a multi-dimensional mesh, some
+        combinations could create specs with mixed Partial reduce types (e.g.,
+        Partial("sum") and Partial("max") in the same placement list). These
+        combinations should be filtered out since mixed partial types don't commute.
+
+        The exception is sum+avg which DO commute and should be allowed.
+        """
+        mesh = DeviceMesh("cpu", mesh=torch.arange(4).reshape(2, 2))
+        meta = TensorMeta(torch.Size([8, 8]), (8, 1), torch.float32)
+
+        # Create input spec with Replicate placement
+        input_spec = DTensorSpec(mesh, (Replicate(), Replicate()), meta)
+
+        # Create OpSchema
+        op_schema = OpSchema(
+            op=torch.ops.aten.mul.Tensor,
+            args_schema=(
+                OpStrategy([OpSpec(input_spec)]),
+                OpStrategy([OpSpec(input_spec)]),
+            ),
+            kwargs_schema={},
+        )
+
+        # Define strategies that would create mixed partials when expanded:
+        # - Strategy 1: Partial("sum") for all tensors
+        # - Strategy 2: Partial("max") for all tensors
+        # When expanded to 2D mesh, combinations like (P_sum, P_max) should be filtered
+        single_mesh_dim_strategies = [
+            [Partial("sum"), Partial("sum"), Partial("sum")],
+            [Partial("max"), Partial("max"), Partial("max")],
+            [Replicate(), Replicate(), Replicate()],
+        ]
+
+        result = expand_to_full_mesh_op_strategy(
+            mesh,
+            op_schema,
+            single_mesh_dim_strategies,
+            output_tensor_meta=meta,
+        )
+
+        # Verify no strategy has mixed partial types (except sum+avg)
+        for strategy in result.strategies:
+            output_spec = strategy.output_spec
+            partial_reduce_ops = {
+                p.reduce_op for p in output_spec.placements if isinstance(p, Partial)
+            }
+            # Either 0 or 1 partial type, or exactly {"sum", "avg"}
+            if len(partial_reduce_ops) > 1:
+                self.assertEqual(
+                    partial_reduce_ops,
+                    {"sum", "avg"},
+                    f"Found invalid mixed partials: {partial_reduce_ops}",
+                )
+
+        # Verify that homogeneous partial strategies ARE included
+        # (P_sum, P_sum) and (P_max, P_max) should be valid
+        found_all_sum = False
+        found_all_max = False
+        for strategy in result.strategies:
+            output_spec = strategy.output_spec
+            if all(
+                isinstance(p, Partial) and p.reduce_op == "sum"
+                for p in output_spec.placements
+            ):
+                found_all_sum = True
+            if all(
+                isinstance(p, Partial) and p.reduce_op == "max"
+                for p in output_spec.placements
+            ):
+                found_all_max = True
+
+        self.assertTrue(found_all_sum, "Should include (P_sum, P_sum) strategy")
+        self.assertTrue(found_all_max, "Should include (P_max, P_max) strategy")
+
+    def test_expand_filters_partial_subclass_with_same_reduce_op(self):
+        """Partial subclasses with the same reduce_op should still be treated as mixed."""
+        mesh = DeviceMesh("cpu", mesh=torch.arange(4).reshape(2, 2))
+        meta = TensorMeta(torch.Size([8, 8]), (8, 1), torch.float32)
+        input_spec = DTensorSpec(mesh, (Replicate(), Replicate()), meta)
+        op_schema = OpSchema(
+            op=torch.ops.aten.mul.Tensor,
+            args_schema=(
+                OpStrategy([OpSpec(input_spec)]),
+                OpStrategy([OpSpec(input_spec)]),
+            ),
+            kwargs_schema={},
+        )
+
+        # _MaskPartial() defaults to reduce_op="sum", same as Partial("sum"),
+        # but they have different reduction semantics and should not be mixed.
+        single_mesh_dim_strategies = [
+            [Partial("sum"), Partial("sum"), Partial("sum")],
+            [_MaskPartial(), _MaskPartial(), _MaskPartial()],
+            [Replicate(), Replicate(), Replicate()],
+        ]
+
+        result = expand_to_full_mesh_op_strategy(
+            mesh,
+            op_schema,
+            single_mesh_dim_strategies,
+            output_tensor_meta=meta,
+        )
+
+        for strategy in result.strategies:
+            output_spec = strategy.output_spec
+            partial_types = {
+                type(p) for p in output_spec.placements if isinstance(p, Partial)
+            }
+            self.assertLessEqual(
+                len(partial_types),
+                1,
+                f"Should not mix Partial subclasses: {output_spec.placements}",
+            )
+
+    def test_expand_allows_sum_avg_partial_mix(self):
+        """Test that sum+avg partial mix is allowed since they commute."""
+        mesh = DeviceMesh("cpu", mesh=torch.arange(4).reshape(2, 2))
+        meta = TensorMeta(torch.Size([8, 8]), (8, 1), torch.float32)
+
+        input_spec = DTensorSpec(mesh, (Replicate(), Replicate()), meta)
+
+        op_schema = OpSchema(
+            op=torch.ops.aten.mul.Tensor,
+            args_schema=(
+                OpStrategy([OpSpec(input_spec)]),
+                OpStrategy([OpSpec(input_spec)]),
+            ),
+            kwargs_schema={},
+        )
+
+        # Define strategies with sum and avg partials
+        single_mesh_dim_strategies = [
+            [Partial("sum"), Partial("sum"), Partial("sum")],
+            [Partial("avg"), Partial("avg"), Partial("avg")],
+            [Replicate(), Replicate(), Replicate()],
+        ]
+
+        result = expand_to_full_mesh_op_strategy(
+            mesh,
+            op_schema,
+            single_mesh_dim_strategies,
+            output_tensor_meta=meta,
+        )
+
+        # Verify that (P_sum, P_avg) combinations ARE included
+        found_sum_avg_mix = False
+        for strategy in result.strategies:
+            output_spec = strategy.output_spec
+            partial_reduce_ops = {
+                p.reduce_op for p in output_spec.placements if isinstance(p, Partial)
+            }
+            if partial_reduce_ops == {"sum", "avg"}:
+                found_sum_avg_mix = True
+                break
+
+        self.assertTrue(
+            found_sum_avg_mix,
+            "Should include mixed (P_sum, P_avg) strategies since sum+avg commute",
+        )
+
+
+class TestDijkstraExpandSingleDimStrategy(TestCase):
+    def setUp(self):
+        super().setUp()
+        store = FakeStore()
+        dist.init_process_group(backend="fake", rank=0, world_size=64, store=store)
+
+    def tearDown(self):
+        super().tearDown()
+        dist.destroy_process_group()
+
+    def _get_mm_output_meta(self, M=64, K=32, N=64):
+        return TensorMeta(
+            shape=torch.Size([M, N]),
+            stride=(N, 1),
+            dtype=torch.float32,
+        )
+
+    def _compare_pq_vs_full_expansion(self, mesh, left_placements, right_placements):
+        """Run both PQ search and full expansion, verify same min cost."""
+        M, K, N = 64, 32, 64
+        left_meta, right_meta = _get_mm_metas(M, K, N)
+        output_meta = self._get_mm_output_meta(M, K, N)
+        left_spec, right_spec = _get_mm_specs(
+            mesh, left_meta, right_meta, left_placements, right_placements
+        )
+
+        op_schema = OpSchema(
+            op=torch.ops.aten.mm.default,
+            args_schema=(
+                OpStrategy([OpSpec(left_spec)]),
+                OpStrategy([OpSpec(right_spec)]),
+            ),
+            kwargs_schema={},
+        )
+
+        # PQ search
+        pq_strategy = _dijkstra_expand_single_dim_strategy_to_mesh(
+            mesh, op_schema, mm_single_dim_strategy, output_tensor_meta=output_meta
+        )
+        self.assertIsNotNone(pq_strategy)
+        self.assertIsInstance(pq_strategy, OpStrategy)
+        self.assertEqual(len(pq_strategy.strategies), 1)
+        pq_cost = sum(chain.from_iterable(pq_strategy.strategies[0].redistribute_cost))
+
+        # Full expansion reference using graph-based (min-cost) redistribution
+        # planning. PQ's Dijkstra search over all per-dim transition orderings
+        # finds the globally optimal ordering (accounting for comm_bytes updates),
+        # which can be strictly cheaper than even the graph-based planner.
+        with use_min_cost_redistribution_plan():
+            expanded_strategy_fn = _expand_single_dim_strategy_to_mesh(
+                mesh,
+                op_schema,
+                _SingleDimStrategyInfo(mm_single_dim_strategy),
+                output_meta,
+            )
+            ref_strategy = expanded_strategy_fn(
+                torch.ops.aten.mm.default,
+                op_schema.args_meta,
+                op_schema.kwargs_meta,
+            )
+        ref_min_cost = min(
+            sum(chain.from_iterable(s.redistribute_cost))
+            for s in ref_strategy.strategies
+        )
+
+        self.assertLessEqual(
+            pq_cost,
+            ref_min_cost + 1e-9,
+            msg=(
+                f"PQ cost {pq_cost} > ref min cost {ref_min_cost} for "
+                f"left={left_placements}, right={right_placements}"
+            ),
+        )
+
+    def test_dijkstra_expand_single_dim_strategy_to_mesh_basic(self):
+        mesh = DeviceMesh("cpu", mesh=torch.arange(8).reshape(2, 2, 2))
+        left_meta, right_meta = _get_mm_metas()
+        output_meta = self._get_mm_output_meta()
+        left_spec, right_spec = _get_mm_specs(
+            mesh,
+            left_meta,
+            right_meta,
+            left_placements=(Shard(0), Replicate(), Replicate()),
+            right_placements=(Replicate(), Replicate(), Replicate()),
+        )
+
+        op_schema = OpSchema(
+            op=torch.ops.aten.mm.default,
+            args_schema=(
+                OpStrategy([OpSpec(left_spec)]),
+                OpStrategy([OpSpec(right_spec)]),
+            ),
+            kwargs_schema={},
+        )
+
+        strategy = _dijkstra_expand_single_dim_strategy_to_mesh(
+            mesh, op_schema, mm_single_dim_strategy, output_tensor_meta=output_meta
+        )
+
+        self.assertIsInstance(strategy, OpStrategy)
+        self.assertEqual(len(strategy.strategies), 1)
+
+        op_spec = strategy.strategies[0]
+        output_spec = op_spec.output_spec
+        input_specs = op_spec.input_specs
+
+        # Zero-cost match: left stays Shard(0), right stays Replicate
+        self.assertEqual(output_spec.placements, (Shard(0), Replicate(), Replicate()))
+        self.assertEqual(
+            input_specs[0].placements, (Shard(0), Replicate(), Replicate())
+        )
+        self.assertEqual(
+            input_specs[1].placements, (Replicate(), Replicate(), Replicate())
+        )
+        # Verify output has tensor_meta
+        self.assertIsNotNone(output_spec.tensor_meta)
+        self.assertEqual(output_spec.tensor_meta.shape, torch.Size([64, 64]))
+
+    def test_dijkstra_expand_single_dim_strategy_to_mesh_hard(self):
+        """Verify PQ search matches full expansion min cost on 3D mesh."""
+        mesh = DeviceMesh("cpu", mesh=torch.arange(8).reshape(2, 2, 2))
+        self._compare_pq_vs_full_expansion(
+            mesh,
+            left_placements=(Shard(0), Replicate(), Replicate()),
+            right_placements=(Shard(1), Shard(0), Replicate()),
+        )
+
+    def test_dijkstra_expand_single_dim_strategy_to_mesh_hard_4d(self):
+        """Verify PQ search matches full expansion min cost on 4D mesh."""
+        mesh = DeviceMesh("cpu", mesh=torch.arange(16).reshape(2, 2, 2, 2))
+        self._compare_pq_vs_full_expansion(
+            mesh,
+            left_placements=(Replicate(), Shard(0), Replicate(), Replicate()),
+            right_placements=(Shard(0), Shard(1), Shard(0), Replicate()),
+        )
+
+    def test_pq_vs_full_expansion_data_driven(self):
+        """Data-driven comparison across mesh shapes, all placement types for mm.
+
+        Enumerates all combos of R, S(0), S(1), P(sum) for each mesh dim on
+        both inputs, across 1D/2D/3D meshes.
+        """
+        placement_options = [Shard(0), Shard(1), Replicate(), Partial("sum")]
+        mesh_configs = [
+            ("1d", torch.arange(4)),
+            ("2d", torch.arange(4).reshape(2, 2)),
+            ("3d", torch.arange(8).reshape(2, 2, 2)),
+        ]
+
+        for mesh_name, mesh_tensor in mesh_configs:
+            mesh = DeviceMesh("cpu", mesh=mesh_tensor)
+            ndim = mesh.ndim
+
+            all_placements = list(product(placement_options, repeat=ndim))
+
+            for left_pl in all_placements:
+                for right_pl in all_placements:
+                    with self.subTest(mesh=mesh_name, left=left_pl, right=right_pl):
+                        self._compare_pq_vs_full_expansion(mesh, left_pl, right_pl)
+
+    def test_strided_shard_fallback(self):
+        """Verify PQ search returns None for StridedShard inputs."""
+        mesh = DeviceMesh("cpu", mesh=torch.arange(4).reshape(2, 2))
+        left_meta, right_meta = _get_mm_metas()
+        output_meta = self._get_mm_output_meta()
+        left_spec = DTensorSpec(
+            mesh,
+            (_StridedShard(0, split_factor=2), Replicate()),
+            tensor_meta=left_meta,
+        )
+        right_spec = DTensorSpec(
+            mesh,
+            (Replicate(), Replicate()),
+            tensor_meta=right_meta,
+        )
+
+        op_schema = OpSchema(
+            op=torch.ops.aten.mm.default,
+            args_schema=(
+                OpStrategy([OpSpec(left_spec)]),
+                OpStrategy([OpSpec(right_spec)]),
+            ),
+            kwargs_schema={},
+        )
+
+        result = _dijkstra_expand_single_dim_strategy_to_mesh(
+            mesh, op_schema, mm_single_dim_strategy, output_tensor_meta=output_meta
+        )
+        self.assertIsNone(result)
+
+    def test_if_elif_dead_end_fix(self):
+        """Counterexample: state has R dims but chunking is useless, needs allgather.
+
+        On a 2D mesh with strategies [R, S(0) -> S(0)] + [R, R -> R],
+        input state left=(S(0), R), right=(R, S(0)).
+        Optimal = 1 allgather on right dim1: left=(S(0), R), right=(R, R)
+        -> matches [S(0),R->S(0)] on dim0 and [R,R->R] on dim1, cost = 1 allgather.
+
+        Old if/elif code would try to chunk first (since some dims are R),
+        miss the direct allgather path, and find a 2-allgather solution.
+        """
+        mesh = DeviceMesh("cpu", mesh=torch.arange(4).reshape(2, 2))
+        M, K, N = 64, 32, 64
+        left_meta, right_meta = _get_mm_metas(M, K, N)
+        output_meta = self._get_mm_output_meta(M, K, N)
+
+        left_spec = DTensorSpec(
+            mesh,
+            (Shard(0), Replicate()),
+            tensor_meta=left_meta,
+        )
+        right_spec = DTensorSpec(
+            mesh,
+            (Replicate(), Shard(0)),
+            tensor_meta=right_meta,
+        )
+
+        op_schema = OpSchema(
+            op=torch.ops.aten.mm.default,
+            args_schema=(
+                OpStrategy([OpSpec(left_spec)]),
+                OpStrategy([OpSpec(right_spec)]),
+            ),
+            kwargs_schema={},
+        )
+
+        strategy = _dijkstra_expand_single_dim_strategy_to_mesh(
+            mesh, op_schema, mm_single_dim_strategy, output_tensor_meta=output_meta
+        )
+        self.assertIsNotNone(strategy)
+        self.assertIsInstance(strategy, OpStrategy)
+        self.assertEqual(len(strategy.strategies), 1)
+
+        # The optimal strategy should have the same cost as full expansion
+        pq_cost = sum(chain.from_iterable(strategy.strategies[0].redistribute_cost))
+
+        # Full expansion reference
+        expanded_strategy_fn = _expand_single_dim_strategy_to_mesh(
+            mesh,
+            op_schema,
+            _SingleDimStrategyInfo(mm_single_dim_strategy),
+            output_meta,
+        )
+        ref_strategy = expanded_strategy_fn(
+            torch.ops.aten.mm.default,
+            op_schema.args_meta,
+            op_schema.kwargs_meta,
+        )
+        ref_min_cost = min(
+            sum(chain.from_iterable(s.redistribute_cost))
+            for s in ref_strategy.strategies
+        )
+
+        self.assertAlmostEqual(pq_cost, ref_min_cost, places=5)
+
+    def test_transitions_tracked(self):
+        """Verify that PQ search tracks transitions for TransformInfo."""
+        mesh = DeviceMesh("cpu", mesh=torch.arange(4).reshape(2, 2))
+        left_meta, right_meta = _get_mm_metas()
+        output_meta = self._get_mm_output_meta()
+
+        # Use placements that require redistribution to verify transitions are tracked
+        left_spec = DTensorSpec(
+            mesh,
+            (Shard(0), Shard(0)),
+            tensor_meta=left_meta,
+        )
+        right_spec = DTensorSpec(
+            mesh,
+            (Replicate(), Replicate()),
+            tensor_meta=right_meta,
+        )
+
+        op_schema = OpSchema(
+            op=torch.ops.aten.mm.default,
+            args_schema=(
+                OpStrategy([OpSpec(left_spec)]),
+                OpStrategy([OpSpec(right_spec)]),
+            ),
+            kwargs_schema={},
+        )
+
+        strategy = _dijkstra_expand_single_dim_strategy_to_mesh(
+            mesh, op_schema, mm_single_dim_strategy, output_tensor_meta=output_meta
+        )
+        self.assertIsNotNone(strategy)
+
+        # The _pq_transitions attribute should exist
+        transitions = getattr(strategy, "_pq_transitions", None)
+        self.assertIsNotNone(transitions)
+        self.assertIsInstance(transitions, list)
+        # Each transition should be (input_idx, mesh_dim, src, dst)
+        for t in transitions:
+            self.assertEqual(len(t), 4)
+
+    def test_pq_reachability_collect_all_2d(self):
+        """On a 2D mesh, verify PQ search can reach all valid strategies.
+
+        For every starting placement combo, run _dijkstra_expand_single_dim_strategy_to_mesh with
+        _collect_all_matches and union all collected sets. Assert the full
+        expansion reference set is a subset of the collected union.
+        """
+        from itertools import product
+
+        mesh = DeviceMesh("cpu", mesh=torch.arange(4).reshape(2, 2))
+        M, K, N = 64, 32, 64
+        left_meta, right_meta = _get_mm_metas(M, K, N)
+        output_meta = self._get_mm_output_meta(M, K, N)
+
+        placement_options = [Shard(0), Shard(1), Replicate(), Partial("sum")]
+        all_placements = list(product(placement_options, repeat=mesh.ndim))
+
+        # Get full expansion reference set
+        ref_left_spec, ref_right_spec = _get_mm_specs(
+            mesh,
+            left_meta,
+            right_meta,
+            left_placements=(Replicate(), Replicate()),
+            right_placements=(Replicate(), Replicate()),
+        )
+        wrapped_schema = OpSchema(
+            op=torch.ops.aten.mm.default,
+            args_schema=(
+                OpStrategy([OpSpec(ref_left_spec)]),
+                OpStrategy([OpSpec(ref_right_spec)]),
+            ),
+            kwargs_schema={},
+        )
+        expanded_fn = _expand_single_dim_strategy_to_mesh(
+            mesh,
+            wrapped_schema,
+            _SingleDimStrategyInfo(mm_single_dim_strategy),
+            output_meta,
+        )
+        ref_strategy = expanded_fn(
+            torch.ops.aten.mm.default,
+            wrapped_schema.args_meta,
+            wrapped_schema.kwargs_meta,
+        )
+        full_expansion_set = {
+            tuple(spec.placements for spec in s.input_specs)
+            for s in ref_strategy.strategies
+        }
+
+        # Collect all reachable matches from every starting placement
+        collected_union: set[tuple[tuple[Placement, ...], ...]] = set()
+        for left_pl in all_placements:
+            for right_pl in all_placements:
+                left_spec, right_spec = _get_mm_specs(
+                    mesh,
+                    left_meta,
+                    right_meta,
+                    left_pl,
+                    right_pl,
+                )
+                op_schema = OpSchema(
+                    op=torch.ops.aten.mm.default,
+                    args_schema=(
+                        OpStrategy([OpSpec(left_spec)]),
+                        OpStrategy([OpSpec(right_spec)]),
+                    ),
+                    kwargs_schema={},
+                )
+                matches: set[tuple[tuple[Placement, ...], ...]] = set()
+                _dijkstra_expand_single_dim_strategy_to_mesh(
+                    mesh,
+                    op_schema,
+                    mm_single_dim_strategy,
+                    output_tensor_meta=output_meta,
+                    _collect_all_matches=matches,
+                )
+                collected_union.update(matches)
+
+        missing = full_expansion_set - collected_union
+        self.assertEqual(
+            missing,
+            set(),
+            f"PQ search missed {len(missing)} strategies reachable by full expansion",
+        )
+
+    def test_single_dim_transition_reachability(self):
+        """Verify single-dim transition rules form a connected graph.
+
+        For mm on a 2D mesh, collect all placements that appear per input
+        position, build a directed graph from transition rules, and BFS from
+        each placement to assert all others are reachable.
+        """
+        from collections import deque
+
+        mesh = DeviceMesh("cpu", mesh=torch.arange(4).reshape(2, 2))
+        M, K, N = 64, 32, 64
+        left_meta, right_meta = _get_mm_metas(M, K, N)
+        output_meta = self._get_mm_output_meta(M, K, N)
+
+        ref_left_spec, ref_right_spec = _get_mm_specs(
+            mesh,
+            left_meta,
+            right_meta,
+            left_placements=(Replicate(), Replicate()),
+            right_placements=(Replicate(), Replicate()),
+        )
+        wrapped_schema = OpSchema(
+            op=torch.ops.aten.mm.default,
+            args_schema=(
+                OpStrategy([OpSpec(ref_left_spec)]),
+                OpStrategy([OpSpec(ref_right_spec)]),
+            ),
+            kwargs_schema={},
+        )
+        expanded_fn = _expand_single_dim_strategy_to_mesh(
+            mesh,
+            wrapped_schema,
+            _SingleDimStrategyInfo(mm_single_dim_strategy),
+            output_meta,
+        )
+        ref_strategy = expanded_fn(
+            torch.ops.aten.mm.default,
+            wrapped_schema.args_meta,
+            wrapped_schema.kwargs_meta,
+        )
+
+        # Collect all placements per input position per mesh dim
+        for input_idx in range(2):
+            for mesh_dim in range(mesh.ndim):
+                all_placements: set[Placement] = set()
+                for s in ref_strategy.strategies:
+                    all_placements.add(s.input_specs[input_idx].placements[mesh_dim])
+
+                # Build directed graph from transition rules
+                def is_sharding(p: Placement) -> bool:
+                    return isinstance(p, Shard)
+
+                edges: dict[Placement, set[Placement]] = {
+                    p: set() for p in all_placements
+                }
+                for src in all_placements:
+                    for dst in all_placements:
+                        if src == dst:
+                            continue
+                        # R -> S, R -> P (free)
+                        if isinstance(src, Replicate) and (
+                            is_sharding(dst) or isinstance(dst, Partial)
+                        ):
+                            edges[src].add(dst)
+                        # S -> R (allgather), S -> S' (all-to-all)
+                        if is_sharding(src) and (
+                            isinstance(dst, Replicate) or is_sharding(dst)
+                        ):
+                            edges[src].add(dst)
+                        # P -> R (allreduce), P -> S (reduce-scatter)
+                        if isinstance(src, Partial) and (
+                            isinstance(dst, Replicate) or is_sharding(dst)
+                        ):
+                            edges[src].add(dst)
+
+                # BFS from each placement, assert all others reachable
+                for start in all_placements:
+                    visited: set[Placement] = set()
+                    q = deque([start])
+                    while q:
+                        node = q.popleft()
+                        if node in visited:
+                            continue
+                        visited.add(node)
+                        q.extend(edges.get(node, set()))
+                    self.assertEqual(
+                        visited,
+                        all_placements,
+                        f"input_idx={input_idx}, mesh_dim={mesh_dim}: "
+                        f"from {start}, unreachable: {all_placements - visited}",
+                    )
+
+    def test_pq_cost_not_underestimated(self):
+        """Verify PQ per-input costs >= graph-based redistribute cost.
+
+        The PQ computes costs per-dim independently. This test checks that
+        each per-input PQ cost is at least as high as the graph-based
+        (min-cost) redistribute planner's cost for the same source→target.
+
+        The graph-based planner is the right baseline because both it and
+        the PQ find optimal transition orderings. The greedy planner uses
+        a fixed ordering that can be more expensive, so comparing against
+        greedy would flag the PQ's ordering optimizations as false positives.
+        """
+        from torch.distributed.tensor._collective_utils import redistribute_cost
+        from torch.distributed.tensor._redistribute import _gen_transform_infos
+
+        # Clear cached transform infos so use_min_cost_redistribution_plan()
+        # takes effect (earlier tests may have cached greedy results).
+        _gen_transform_infos.cache_clear()
+
+        placement_options = [Shard(0), Shard(1), Replicate(), Partial("sum")]
+        mesh = DeviceMesh("cpu", mesh=torch.arange(4).reshape(2, 2))
+        M, K, N = 64, 32, 64
+        output_meta = self._get_mm_output_meta(M, K, N)
+        all_placements = list(product(placement_options, repeat=mesh.ndim))
+
+        underestimated = []
+        for left_pl in all_placements:
+            for right_pl in all_placements:
+                left_meta, right_meta = _get_mm_metas(M, K, N)
+                left_spec, right_spec = _get_mm_specs(
+                    mesh, left_meta, right_meta, left_pl, right_pl
+                )
+                op_schema = OpSchema(
+                    op=torch.ops.aten.mm.default,
+                    args_schema=(
+                        OpStrategy([OpSpec(left_spec)]),
+                        OpStrategy([OpSpec(right_spec)]),
+                    ),
+                    kwargs_schema={},
+                )
+                pq_strategy = _dijkstra_expand_single_dim_strategy_to_mesh(
+                    mesh,
+                    op_schema,
+                    mm_single_dim_strategy,
+                    output_tensor_meta=output_meta,
+                )
+                self.assertIsNotNone(pq_strategy)
+                pq_spec = pq_strategy.strategies[0]
+                with use_min_cost_redistribution_plan():
+                    for input_idx, (init_spec, target_spec) in enumerate(
+                        zip([left_spec, right_spec], pq_spec.input_specs)
+                    ):
+                        pq_input_cost = pq_spec.redistribute_cost[input_idx][0]
+                        actual_cost = redistribute_cost(init_spec, target_spec)
+                        if pq_input_cost < actual_cost - 1e-9:
+                            underestimated.append(
+                                f"left={left_pl} right={right_pl} "
+                                f"input_idx={input_idx}: "
+                                f"pq={pq_input_cost:.4f} < "
+                                f"actual={actual_cost:.4f} "
+                                f"target={target_spec.placements}"
+                            )
+
+        self.assertEqual(
+            underestimated,
+            [],
+            f"PQ underestimated cost for {len(underestimated)} cases:\n"
+            + "\n".join(underestimated[:20]),
+        )
+
+    def test_pq_selects_optimal_actual_cost_strategy(self):
+        """Verify PQ-selected strategy has graph-based cost <= full expansion min.
+
+        Both the PQ and the reference use graph-based (min-cost) redistribution
+        planning for a fair comparison. This catches cases where the PQ selects
+        a genuinely worse strategy, not just ordering differences.
+        """
+        from torch.distributed.tensor._collective_utils import redistribute_cost
+        from torch.distributed.tensor._redistribute import _gen_transform_infos
+
+        _gen_transform_infos.cache_clear()
+
+        placement_options = [Shard(0), Shard(1), Replicate(), Partial("sum")]
+        mesh = DeviceMesh("cpu", mesh=torch.arange(4).reshape(2, 2))
+        M, K, N = 64, 32, 64
+        output_meta = self._get_mm_output_meta(M, K, N)
+        all_placements = list(product(placement_options, repeat=mesh.ndim))
+
+        wrong_strategy = []
+        for left_pl in all_placements:
+            for right_pl in all_placements:
+                left_meta, right_meta = _get_mm_metas(M, K, N)
+                left_spec, right_spec = _get_mm_specs(
+                    mesh, left_meta, right_meta, left_pl, right_pl
+                )
+                op_schema = OpSchema(
+                    op=torch.ops.aten.mm.default,
+                    args_schema=(
+                        OpStrategy([OpSpec(left_spec)]),
+                        OpStrategy([OpSpec(right_spec)]),
+                    ),
+                    kwargs_schema={},
+                )
+                pq_strategy = _dijkstra_expand_single_dim_strategy_to_mesh(
+                    mesh,
+                    op_schema,
+                    mm_single_dim_strategy,
+                    output_tensor_meta=output_meta,
+                )
+                self.assertIsNotNone(pq_strategy)
+                pq_spec = pq_strategy.strategies[0]
+
+                with use_min_cost_redistribution_plan():
+                    # Actual cost of PQ-selected strategy (graph-based)
+                    pq_actual = sum(
+                        redistribute_cost(init, tgt)
+                        for init, tgt in zip(
+                            [left_spec, right_spec], pq_spec.input_specs
+                        )
+                    )
+
+                    # Reference: full expansion with graph-based planning
+                    expanded_fn = _expand_single_dim_strategy_to_mesh(
+                        mesh,
+                        op_schema,
+                        _SingleDimStrategyInfo(mm_single_dim_strategy),
+                        output_meta,
+                    )
+                    ref = expanded_fn(
+                        torch.ops.aten.mm.default,
+                        op_schema.args_meta,
+                        op_schema.kwargs_meta,
+                    )
+                ref_min = min(
+                    sum(chain.from_iterable(s.redistribute_cost))
+                    for s in ref.strategies
+                )
+
+                if pq_actual > ref_min + 1e-9:
+                    wrong_strategy.append(
+                        f"left={left_pl} right={right_pl}: "
+                        f"pq_actual={pq_actual:.4f} > "
+                        f"ref_min={ref_min:.4f} "
+                        f"target_left={pq_spec.input_specs[0].placements} "
+                        f"target_right={pq_spec.input_specs[1].placements}"
+                    )
+
+        self.assertEqual(
+            wrong_strategy,
+            [],
+            f"PQ selected worse strategy in {len(wrong_strategy)} cases:\n"
+            + "\n".join(wrong_strategy[:20]),
+        )
+
 
 @torch.library.custom_op("mylib::dummy_add", mutates_args=())
 def dummy_add(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
@@ -1028,6 +1864,101 @@ def _dummy_check_fake(x):
     return None
 
 
+class TestCommonPointwiseSingleDimStrategy(TestCase):
+    """Unit tests for _common_pointwise_single_dim_strategy raw rule generation."""
+
+    def _meta(self, *dims: int) -> TensorMeta:
+        shape = torch.Size(dims)
+        stride = torch.empty(shape).stride()
+        return TensorMeta(shape=shape, stride=stride, dtype=torch.float32)
+
+    @staticmethod
+    def _normalize(rules):
+        """Convert rules to a comparable form (replace _ShardingPlaceholder with Shard)."""
+        out = []
+        for rule in rules:
+            out.append(
+                tuple(
+                    Shard(p.dim) if isinstance(p, _ShardingPlaceholder) else p
+                    for p in rule
+                )
+            )
+        return out
+
+    def test_unary_2d_shard_rules(self):
+        """Unary op with a 2D tensor should produce two Shard rules."""
+        fn = _common_pointwise_single_dim_strategy()
+        rules = self._normalize(fn(torch.ops.aten.abs.default, (self._meta(4, 8),), {}))
+        self.assertEqual(
+            rules,
+            [(Shard(0), Shard(0)), (Shard(1), Shard(1))],
+        )
+
+    def test_unary_with_partial_extra_rules(self):
+        """Unary op with _UNARY_LINEAR_RULES should append partial rules."""
+        fn = _common_pointwise_single_dim_strategy(
+            partial_extra_rules=_UNARY_LINEAR_RULES
+        )
+        rules = self._normalize(fn(torch.ops.aten.neg.default, (self._meta(4, 8),), {}))
+        expected = [
+            (Shard(0), Shard(0)),
+            (Shard(1), Shard(1)),
+            (Partial("sum"), Partial("sum")),
+            (Partial("avg"), Partial("avg")),
+        ]
+        self.assertEqual(rules, expected)
+
+    def test_binary_broadcast_replicate(self):
+        """When one input is broadcast, the broadcast dim gets Replicate."""
+        fn = _common_pointwise_single_dim_strategy()
+        # (4, 8) + (8,) — dim 0 is broadcast for the second arg
+        rules = self._normalize(
+            fn(torch.ops.aten.add.Tensor, (self._meta(4, 8), self._meta(8)), {})
+        )
+        self.assertEqual(
+            rules,
+            [(Shard(0), Shard(0), Replicate()), (Shard(1), Shard(1), Shard(0))],
+        )
+
+    def test_partial_extra_rules_filtered_by_arity(self):
+        """Binary extra rules are filtered out when only one tensor arg is present."""
+        fn = _common_pointwise_single_dim_strategy(
+            partial_extra_rules=_MUL_RULES + _UNARY_LINEAR_RULES
+        )
+        # Scalar promotion: mul.Tensor with one tensor arg
+        rules = self._normalize(
+            fn(torch.ops.aten.mul.Tensor, (self._meta(4, 8), 2.0), {})
+        )
+        # Length-3 _MUL_RULES should be filtered out, only length-2 _UNARY_LINEAR_RULES kept
+        expected = [
+            (Shard(0), Shard(0)),
+            (Shard(1), Shard(1)),
+            (Partial("sum"), Partial("sum")),
+            (Partial("avg"), Partial("avg")),
+        ]
+        self.assertEqual(rules, expected)
+
+    def test_binary_with_additive_rules(self):
+        """Binary op with _BINARY_ADDITIVE_RULES appends the right partial rules."""
+        fn = _common_pointwise_single_dim_strategy(
+            partial_extra_rules=_BINARY_ADDITIVE_RULES
+        )
+        rules = self._normalize(
+            fn(torch.ops.aten.add.Tensor, (self._meta(4, 8), self._meta(4, 8)), {})
+        )
+        expected = [
+            (Shard(0), Shard(0), Shard(0)),
+            (Shard(1), Shard(1), Shard(1)),
+            (Partial("sum"), Partial("sum"), Partial("sum")),
+            (Partial("avg"), Partial("avg"), Partial("avg")),
+            (Partial("avg"), Partial("avg"), Replicate()),
+            (Partial("max"), Partial("max"), Replicate()),
+            (Partial("min"), Partial("min"), Replicate()),
+            (Partial("avg"), Replicate(), Partial("avg")),
+        ]
+        self.assertEqual(rules, expected)
+
+
 class TestSingleDimStrategyRegistration(TestCase):
     def setUp(self):
         super().setUp()
diff --git a/test/distributed/tensor/test_strategy_validation.py b/test/distributed/tensor/test_strategy_validation.py
index 4d357f2c3dce1..df9d6f7c0d308 100644
--- a/test/distributed/tensor/test_strategy_validation.py
+++ b/test/distributed/tensor/test_strategy_validation.py
@@ -10,22 +10,32 @@
 """
 
 import torch
+import torch.distributed as dist
+from torch.distributed._local_tensor import LocalTensorMode
+from torch.distributed.device_mesh import init_device_mesh
 from torch.distributed.tensor import Replicate
 from torch.distributed.tensor._ops.strategy_validation import (
+    _checkerboard_mask,
+    _create_partial_input,
+    _find_opinfo_candidates,
     extract_tensors_from_sample,
     get_1d_input_placements_for_tensor,
     get_1d_output_placements_for_tensor,
+    get_opinfo_by_name,
     is_fully_replicated,
     is_trivial_shard,
     normalize_combo_key,
     normalize_placement,
     normalize_placement_str,
     parse_placement,
-    placement_tuple_to_str,
+    query_single_dim_strategy,
+    resolve_op_names,
+    validate_aten_combination,
+    validate_combination,
 )
 from torch.distributed.tensor.placement_types import Partial, Shard
 from torch.testing._internal.common_methods_invocations import SampleInput
-from torch.testing._internal.common_utils import run_tests, TestCase
+from torch.testing._internal.common_utils import run_tests, TEST_WITH_SLOW, TestCase
 
 
 class TestPlacementUtilities(TestCase):
@@ -53,16 +63,6 @@ def test_parse_placement_partial(self):
         self.assertIsInstance(p, Partial)
         self.assertEqual(p.reduce_op, "max")
 
-    def test_placement_tuple_to_str(self):
-        s = placement_tuple_to_str((Replicate(),))
-        self.assertEqual(s, "(R)")
-
-        s = placement_tuple_to_str((Shard(0), Replicate()))
-        self.assertEqual(s, "(S(0), R)")
-
-        s = placement_tuple_to_str((Partial("sum"),))
-        self.assertEqual(s, "(P(sum))")
-
     def test_is_fully_replicated(self):
         self.assertTrue(is_fully_replicated((Replicate(),)))
         self.assertTrue(is_fully_replicated((Replicate(), Replicate())))
@@ -135,49 +135,49 @@ def test_normalize_placement_str_normal_shard(self):
     def test_normalize_combo_key_trivial_output(self):
         """Combo keys with trivial output shard should normalize output to R."""
         # P(max) -> S(0) on output [1,1,1] should become P(max) -> R
-        combo = (("P(max)",), "S(0)")
+        combo = (("P(max)",), ("S(0)",))
         input_shapes = ((4, 3),)
-        output_shape = (1, 1, 1)
+        output_shapes = ((1, 1, 1),)
 
-        result = normalize_combo_key(combo, input_shapes, output_shape)
-        self.assertEqual(result, (("P(max)",), "R"))
+        result = normalize_combo_key(combo, input_shapes, output_shapes)
+        self.assertEqual(result, (("P(max)",), ("R",)))
 
     def test_normalize_combo_key_trivial_input(self):
         """Combo keys with trivial input shard should normalize input to R."""
         # S(0),R -> R on input [1,4],[4,4] should become R,R -> R
-        combo = (("S(0)", "R"), "R")
+        combo = (("S(0)", "R"), ("R",))
         input_shapes = ((1, 4), (4, 4))
-        output_shape = (4, 4)
+        output_shapes = ((4, 4),)
 
-        result = normalize_combo_key(combo, input_shapes, output_shape)
-        self.assertEqual(result, (("R", "R"), "R"))
+        result = normalize_combo_key(combo, input_shapes, output_shapes)
+        self.assertEqual(result, (("R", "R"), ("R",)))
 
     def test_normalize_combo_key_all_trivial(self):
         """All-size-1 tensor should normalize all shards to R."""
         # S(0),S(1) -> S(2) on shape [1,1,1] should become R,R -> R
-        combo = (("S(0)", "S(1)"), "S(2)")
+        combo = (("S(0)", "S(1)"), ("S(2)",))
         input_shapes = ((1, 1, 1), (1, 1, 1))
-        output_shape = (1, 1, 1)
+        output_shapes = ((1, 1, 1),)
 
-        result = normalize_combo_key(combo, input_shapes, output_shape)
-        self.assertEqual(result, (("R", "R"), "R"))
+        result = normalize_combo_key(combo, input_shapes, output_shapes)
+        self.assertEqual(result, (("R", "R"), ("R",)))
 
     def test_normalize_combo_key_no_change(self):
         """Normal-sized tensors should not change."""
-        combo = (("S(0)", "S(1)"), "S(0)")
+        combo = (("S(0)", "S(1)"), ("S(0)",))
         input_shapes = ((4, 3), (4, 3))
-        output_shape = (4, 3)
+        output_shapes = ((4, 3),)
 
-        result = normalize_combo_key(combo, input_shapes, output_shape)
+        result = normalize_combo_key(combo, input_shapes, output_shapes)
         self.assertEqual(result, combo)
 
     def test_normalize_combo_key_partial_unchanged(self):
         """Partial placements should never be normalized."""
-        combo = (("P(max)", "P(sum)"), "P(max)")
+        combo = (("P(max)", "P(sum)"), ("P(max)",))
         input_shapes = ((1, 1), (1, 1))  # Even with all-size-1
-        output_shape = (1, 1)
+        output_shapes = ((1, 1),)
 
-        result = normalize_combo_key(combo, input_shapes, output_shape)
+        result = normalize_combo_key(combo, input_shapes, output_shapes)
         self.assertEqual(result, combo)  # Partials unchanged
 
     def test_normalize_deduplicates_equivalent_rules(self):
@@ -193,22 +193,22 @@ def test_normalize_deduplicates_equivalent_rules(self):
         After normalization, they should all become P(max) -> R.
         """
         input_shapes = ((1, 1, 1, 1),)
-        output_shape = (1, 1, 1)
+        output_shapes = ((1, 1, 1),)
 
         normalized_rules = set()
         for shard_dim in [0, 1, 2]:
-            combo = (("P(max)",), f"S({shard_dim})")
-            normalized = normalize_combo_key(combo, input_shapes, output_shape)
+            combo = (("P(max)",), (f"S({shard_dim})",))
+            normalized = normalize_combo_key(combo, input_shapes, output_shapes)
             normalized_rules.add(normalized)
 
         # Also add the explicit R version
-        combo_r = (("P(max)",), "R")
-        normalized_r = normalize_combo_key(combo_r, input_shapes, output_shape)
+        combo_r = (("P(max)",), ("R",))
+        normalized_r = normalize_combo_key(combo_r, input_shapes, output_shapes)
         normalized_rules.add(normalized_r)
 
         # All should have normalized to the same rule
         self.assertEqual(len(normalized_rules), 1)
-        self.assertEqual(normalized_rules.pop(), (("P(max)",), "R"))
+        self.assertEqual(normalized_rules.pop(), (("P(max)",), ("R",)))
 
 
 class TestInputPlacements(TestCase):
@@ -317,5 +317,1232 @@ def test_extract_with_kwargs(self):
         self.assertEqual(len(tensors), 2)
 
 
+class TestValidateCombination(TestCase):
+    """Test the core validate_combination function."""
+
+    world_size = 2
+
+    def setUp(self):
+        super().setUp()
+        if not dist.is_initialized():
+            dist.init_process_group("fake", rank=0, world_size=self.world_size)
+
+    def tearDown(self):
+        super().tearDown()
+        if dist.is_initialized():
+            dist.destroy_process_group()
+
+    def test_nan_output_valid_when_both_nan(self):
+        """Test that NaN outputs are considered valid when both ground truth and sharded match.
+
+        igamma with negative inputs produces NaN. When sharding produces the same
+        NaN pattern, the rule should be valid (NaN == NaN for validation purposes).
+        """
+        # igamma(negative, negative) produces NaN
+        a = torch.tensor([-1.0, -2.0, -3.0, -4.0, -5.0, -6.0, -7.0, -8.0])
+        x = torch.tensor(-1.0)
+        sample = SampleInput(a, args=(x,))
+        tensors = extract_tensors_from_sample(sample)
+        ground_truth = torch.igamma(a, x)
+        self.assertTrue(ground_truth.isnan().all(), "Expected all NaN ground truth")
+
+        combo = ((Shard(0), Replicate()), (Shard(0),))
+
+        with LocalTensorMode(frozenset(range(self.world_size))):
+            mesh = init_device_mesh("cpu", (self.world_size,))
+            is_valid, msg = validate_combination(
+                torch.igamma,
+                sample,
+                tensors,
+                combo,
+                ground_truth,
+                self.world_size,
+                mesh,
+            )
+        self.assertTrue(is_valid, f"NaN outputs should match: {msg}")
+
+    def test_integer_output_includes_partial_sum(self):
+        """
+        P(sum) must be included in output placements for integer dtypes.
+
+        Ops like bucketize use R, S(0) -> P(sum) where the output is integer
+        (bucket indices). If P(sum) is excluded from the validator's output
+        enumeration for integer types, these rules appear as false "incorrect"
+        reports because the ground truth enumeration never discovers them.
+        """
+        for dtype in [torch.int32, torch.int64]:
+            t = torch.zeros(4, dtype=dtype)
+            placements = get_1d_output_placements_for_tensor(t)
+            partial_ops = {p.reduce_op for p in placements if isinstance(p, Partial)}
+            self.assertIn(
+                "sum",
+                partial_ops,
+                f"P(sum) must be in output placements for {dtype} "
+                f"(needed by ops like bucketize with sharded boundaries)",
+            )
+
+    def test_exhaustive_binary_op_rules(self):
+        """
+        Exhaustively test all placement combinations for binary ops.
+
+        For each op, we define the complete set of valid rules. The test then:
+        1. Verifies all listed rules are detected as valid
+        2. Verifies all unlisted combinations are detected as invalid
+        """
+
+        def parse_rule(rule_str):
+            """Parse 'S(0),S(0)->S(0)' into ((Shard(0), Shard(0)), Shard(0))."""
+            inputs_str, output_str = rule_str.split("->")
+            inputs = tuple(parse_placement(s.strip()) for s in inputs_str.split(","))
+            return (inputs, (parse_placement(output_str.strip()),))
+
+        # Valid rules for 2D binary ops with shape (8, 4)
+        # Format: "input1,input2->output"
+        # NOTE: These rules are for ground-truth validation (no redistribution).
+        # Only placements that produce compatible local tensor shapes are valid.
+        # - Shard(d) produces chunks (4,4) for d=0 or (8,2) for d=1
+        # - Replicate and Partial produce full tensors (8,4)
+        # So Shard can only pair with same-dim Shard, not with R or P.
+        #
+        # For Partial:
+        # - P(sum) + P(sum) -> P(sum) works for add (sums distribute)
+        # - R * P(sum) -> P(sum) works for mul (multiplication distributes into sum)
+        # - P(sum) / R -> P(sum) works for div
+        # But NOT:
+        # - R + P(sum) -> P(sum) for add (R gets added on each rank, then summed)
+        VALID_RULES = {
+            torch.add: [
+                # Same-dim sharding (chunks match)
+                "S(0),S(0)->S(0)",
+                "S(1),S(1)->S(1)",
+                # Partial sum/avg + Partial sum/avg -> same (linearity of addition)
+                # (a0+a1) + (b0+b1) = (a0+b0) + (a1+b1) where ai/bi are per-rank
+                "P(sum),P(sum)->P(sum)",
+                "P(avg),P(avg)->P(avg)",
+                # Partial sum/avg with Replicate: avg normalizes the extra
+                # copies, so P(avg)+R works even though P(sum)+R does not
+                "P(avg),R->P(avg)",
+                "R,P(avg)->P(avg)",
+                # Partial max/min with Replicate: adding a constant preserves
+                # the reduce structure (NOT Pmax+Pmax, offsets accumulate)
+                "P(max),R->P(max)",
+                "P(min),R->P(min)",
+                # NOTE: these two rules are NOT valid in general for torch.add since it accepts alpha=a, which if negative
+                # flips the the partial output from max to min or vice versa.
+                # However, this test is simpler than the end to end validator and ignores alpha, and the rules have to
+                # be listed as valid since without alpha they DO produce correct results and the test asserts any rule
+                # NOT listed here produces incorrect results.
+                "R,P(max)->P(max)",
+                "R,P(min)->P(min)",
+            ],
+            torch.mul: [
+                # Same-dim sharding
+                "S(0),S(0)->S(0)",
+                "S(1),S(1)->S(1)",
+                # Partial sum/avg * Replicate -> same (multiplicative linearity)
+                # r * (p0+p1) = r*p0 + r*p1 where pi are per-rank
+                "P(sum),R->P(sum)",
+                "R,P(sum)->P(sum)",
+                "P(avg),R->P(avg)",
+                "R,P(avg)->P(avg)",
+                # No P(min)/P(max) rules: negative multiplier flips ordering
+            ],
+            torch.div: [
+                # Same-dim sharding
+                "S(0),S(0)->S(0)",
+                "S(1),S(1)->S(1)",
+                # Partial sum/avg / Replicate -> same (division by constant is linear)
+                "P(sum),R->P(sum)",
+                "P(avg),R->P(avg)",
+                # No R/P(avg) rule: 1/x is not linear
+                # No P(min)/P(max) rules: negative divisor flips ordering
+            ],
+            torch.maximum: [
+                # Same-dim sharding
+                "S(0),S(0)->S(0)",
+                "S(1),S(1)->S(1)",
+                # Partial max + Partial max -> Partial max
+                # max(max(a0,a1), max(b0,b1)) = max(max(a0,b0), max(a1,b1))
+                "P(max),P(max)->P(max)",
+                # Partial max/min with Replicate (lattice distributivity):
+                # max(min(a0,a1), r) = min(max(a0,r), max(a1,r)) ✓
+                "P(max),R->P(max)",
+                "R,P(max)->P(max)",
+                "P(min),R->P(min)",
+                "R,P(min)->P(min)",
+                # No P(avg) rules: max is not linear
+            ],
+        }
+
+        # All possible placements for 2D tensor
+        ALL_PLACEMENTS = [
+            Replicate(),
+            Shard(0),
+            Shard(1),
+            Partial("sum"),
+            Partial("max"),
+        ]
+        if TEST_WITH_SLOW:
+            # This makes the test go from 4 sec to 12 sec, and I don't think it really adds much useful coverage, but
+            # why not have it run in CI.
+            ALL_PLACEMENTS += [Partial("avg"), Partial("min")]
+
+        # Test each operator
+        for op, valid_rule_strs in VALID_RULES.items():
+            valid_rules = {parse_rule(r) for r in valid_rule_strs}
+
+            # Create test tensors
+            a = torch.randn(8, 4)
+            b = torch.randn(8, 4)
+            sample = SampleInput(a, args=(b,))
+            tensors = extract_tensors_from_sample(sample)
+            ground_truth = op(a, b)
+
+            with LocalTensorMode(frozenset(range(self.world_size))):
+                mesh = init_device_mesh("cpu", (self.world_size,))
+
+                # Test all combinations
+                for p1 in ALL_PLACEMENTS:
+                    for p2 in ALL_PLACEMENTS:
+                        # Skip fully replicated inputs (degenerate case: any output works)
+                        if isinstance(p1, Replicate) and isinstance(p2, Replicate):
+                            continue
+
+                        for p_out in ALL_PLACEMENTS:
+                            input_plcs = (p1, p2)
+                            combo = (input_plcs, (p_out,))
+
+                            is_valid, msg = validate_combination(
+                                op,
+                                sample,
+                                tensors,
+                                combo,
+                                ground_truth,
+                                self.world_size,
+                                mesh,
+                            )
+
+                            # Check if this combo matches any valid rule
+                            should_be_valid = (input_plcs, (p_out,)) in valid_rules
+
+                            if should_be_valid:
+                                self.assertTrue(
+                                    is_valid,
+                                    f"{op.__name__}: {p1},{p2}->{p_out} should be valid but got: {msg}",
+                                )
+                            else:
+                                self.assertFalse(
+                                    is_valid,
+                                    f"{op.__name__}: {p1},{p2}->{p_out} should be invalid",
+                                )
+
+    def test_add_alpha_negates_partial_max_to_min(self):
+        """
+        Test that add with alpha=-1 converts P(max) to P(min).
+
+        When computing a + alpha*b with alpha=-1:
+        - If b is P(max), then -b behaves as P(min)
+        - So R + (-1)*P(max) -> P(min) should be valid
+        - And R + (-1)*P(max) -> P(max) should be invalid
+        """
+        a = torch.randn(8, 4)
+        b = torch.randn(8, 4)
+        # add with alpha=-1 computes: a + (-1)*b = a - b
+        sample = SampleInput(a, args=(b,), kwargs={"alpha": -1})
+        tensors = extract_tensors_from_sample(sample)
+        ground_truth = torch.add(a, b, alpha=-1)
+
+        with LocalTensorMode(frozenset(range(self.world_size))):
+            mesh = init_device_mesh("cpu", (self.world_size,))
+
+            # R + alpha*P(max) where alpha=-1 should produce P(min)
+            # because -max(x) = min(-x)
+            combo_valid = ((Replicate(), Partial("max")), (Partial("min"),))
+            is_valid, msg = validate_combination(
+                torch.add,
+                sample,
+                tensors,
+                combo_valid,
+                ground_truth,
+                self.world_size,
+                mesh,
+            )
+            self.assertTrue(
+                is_valid, f"R,Pmax->Pmin with alpha=-1 should be valid: {msg}"
+            )
+
+            # R + alpha*P(max) where alpha=-1 should NOT produce P(max)
+            combo_invalid = ((Replicate(), Partial("max")), (Partial("max"),))
+            is_valid, msg = validate_combination(
+                torch.add,
+                sample,
+                tensors,
+                combo_invalid,
+                ground_truth,
+                self.world_size,
+                mesh,
+            )
+            self.assertFalse(is_valid, "R,Pmax->Pmax with alpha=-1 should be invalid")
+
+            # Similarly, P(max) + alpha*R where alpha=-1 should produce P(max)
+            # because we're subtracting a replicated value from P(max)
+            combo_pmax_minus_r = ((Partial("max"), Replicate()), (Partial("max"),))
+            is_valid, msg = validate_combination(
+                torch.add,
+                sample,
+                tensors,
+                combo_pmax_minus_r,
+                ground_truth,
+                self.world_size,
+                mesh,
+            )
+            self.assertTrue(
+                is_valid, f"Pmax,R->Pmax with alpha=-1 should be valid: {msg}"
+            )
+
+
+class TestCreatePartialInput(TestCase):
+    """Test the _create_partial_input helper."""
+
+    def test_partial_sum_reduces_correctly(self):
+        tensor = torch.randn(4, 3)
+        local_tensor = _create_partial_input(tensor, Partial("sum"), world_size=2)
+
+        # Sum of local tensors should equal original
+        total = sum(local_tensor._local_tensors[r] for r in range(2))
+        self.assertTrue(torch.allclose(total, tensor))
+
+    def test_partial_avg_reduces_correctly(self):
+        tensor = torch.randn(4, 3)
+        local_tensor = _create_partial_input(tensor, Partial("avg"), world_size=2)
+
+        # Average of local tensors should equal original
+        total = sum(local_tensor._local_tensors[r] for r in range(2))
+        avg = total / 2
+        self.assertTrue(torch.allclose(avg, tensor))
+
+    def test_partial_min_reduces_correctly(self):
+        tensor = torch.randn(4, 3)
+        local_tensor = _create_partial_input(tensor, Partial("min"), world_size=2)
+
+        # Min of local tensors should equal original
+        stacked = torch.stack([local_tensor._local_tensors[r] for r in range(2)])
+        result = stacked.min(dim=0).values
+        self.assertTrue(torch.allclose(result, tensor))
+
+    def test_partial_max_reduces_correctly(self):
+        tensor = torch.randn(4, 3)
+        local_tensor = _create_partial_input(tensor, Partial("max"), world_size=2)
+
+        # Max of local tensors should equal original
+        stacked = torch.stack([local_tensor._local_tensors[r] for r in range(2)])
+        result = stacked.max(dim=0).values
+        self.assertTrue(torch.allclose(result, tensor))
+
+    def test_tensor_idx_produces_different_sum_ratios(self):
+        """
+        Verify that different tensor_idx values produce different P(sum) ratios.
+
+        This is critical for catching invalid rules like mul P(sum),P(sum)->P(sum).
+        If both tensors used the same 0.5/0.5 split, some invalid combinations
+        might accidentally produce correct results.
+        """
+        tensor = torch.ones(4, 3)
+
+        # Create P(sum) with different tensor_idx values
+        local0 = _create_partial_input(
+            tensor, Partial("sum"), world_size=2, tensor_idx=0
+        )
+        local1 = _create_partial_input(
+            tensor, Partial("sum"), world_size=2, tensor_idx=1
+        )
+
+        # Both should reduce to the same original tensor
+        self.assertTrue(
+            torch.allclose(local0._local_tensors[0] + local0._local_tensors[1], tensor)
+        )
+        self.assertTrue(
+            torch.allclose(local1._local_tensors[0] + local1._local_tensors[1], tensor)
+        )
+
+        # But the per-rank values should be DIFFERENT
+        # (tensor_idx=0 uses ratio 0.6, tensor_idx=1 uses ratio 0.7)
+        self.assertFalse(
+            torch.allclose(local0._local_tensors[0], local1._local_tensors[0]),
+            "Different tensor_idx should produce different P(sum) ratios",
+        )
+
+    def test_tensor_idx_produces_different_max_patterns(self):
+        """
+        Verify that different tensor_idx values produce different P(max) offset patterns.
+
+        This is critical for catching invalid rules like add P(max),P(max)->P(max).
+        If both tensors used the same offset pattern, the offsets would cancel out
+        and invalid combinations might appear valid.
+        """
+        tensor = torch.zeros(4, 3)
+
+        # Create P(max) with different tensor_idx values
+        local0 = _create_partial_input(
+            tensor, Partial("max"), world_size=2, tensor_idx=0
+        )
+        local1 = _create_partial_input(
+            tensor, Partial("max"), world_size=2, tensor_idx=1
+        )
+
+        # Both should reduce to the same original tensor
+        stacked0 = torch.stack([local0._local_tensors[r] for r in range(2)])
+        stacked1 = torch.stack([local1._local_tensors[r] for r in range(2)])
+        self.assertTrue(torch.allclose(stacked0.max(dim=0).values, tensor))
+        self.assertTrue(torch.allclose(stacked1.max(dim=0).values, tensor))
+
+        # The offset patterns should be different (shifted checkerboard)
+        self.assertFalse(
+            torch.allclose(local0._local_tensors[0], local1._local_tensors[0]),
+            "Different tensor_idx should produce different P(max) patterns",
+        )
+
+    def test_checkerboard_mask_alternates_every_dim(self):
+        """
+        The checkerboard mask must alternate along every dimension.
+
+        A flat-index mask (idx % 2) has uniform parity along even-stride
+        dimensions: for a [3,2,1,2] tensor, dim 0 stride is 4, so all
+        elements along dim 0 at a given position share the same flat-index
+        parity. This causes argmin/argmax to produce identical results on
+        all ranks, making P(min)/P(max) inputs falsely appear valid.
+
+        The checkerboard mask (sum-of-coordinates % 2) avoids this by
+        guaranteeing adjacent elements differ along every axis.
+        """
+        # Shape where flat-index mask fails: stride of dim 0 = 4 (even)
+        tensor = torch.randn(3, 2, 1, 2)
+        mask = _checkerboard_mask(tensor)
+
+        # Reshape mask back to tensor shape to check along each dim
+        mask_nd = mask.reshape(tensor.shape)
+
+        # Along dim 0, adjacent elements must alternate
+        for j in range(2):
+            for k in range(1):
+                for l in range(2):
+                    vals = [mask_nd[i, j, k, l].item() for i in range(3)]
+                    # Check alternation: each consecutive pair should differ
+                    for idx in range(len(vals) - 1):
+                        self.assertNotEqual(
+                            vals[idx],
+                            vals[idx + 1],
+                            f"Mask should alternate along dim 0 at [{':,'}{j},{k},{l}], "
+                            f"got {vals}",
+                        )
+
+    def test_pmin_pmax_adaptive_offset_exceeds_range(self):
+        """
+        P(min)/P(max) offsets must exceed the tensor's value range.
+
+        With fixed small offsets (e.g., 0.7 / -1.3), the true minimum
+        dominates across all ranks, making argmin produce the same index
+        everywhere. Adaptive offsets scaled to 2*range+1 ensure the mask
+        pattern determines which rank finds the argmin, so different ranks
+        disagree for index-returning ops.
+        """
+        # Tensor with large value range — fixed offsets would be too small
+        tensor = torch.tensor([-10.0, -1.0, 5.0, 8.0])
+        value_range = (tensor.max() - tensor.min()).item()  # 18.0
+
+        for reduce_op, sign in [("min", 1), ("max", -1)]:
+            local = _create_partial_input(tensor, Partial(reduce_op), world_size=2)
+            r0, r1 = local._local_tensors[0], local._local_tensors[1]
+            # The actual offset applied should exceed value_range
+            max_offset = (r0 - r1).abs().max().item()
+            self.assertGreater(
+                max_offset,
+                value_range,
+                f"P({reduce_op}) offset {max_offset:.1f} should exceed "
+                f"value range {value_range:.1f}",
+            )
+            # Ranks should disagree on argmin/argmax
+            self.assertNotEqual(
+                r0.argmin().item(),
+                r1.argmin().item(),
+                f"P({reduce_op}) ranks should disagree on argmin with adaptive offset",
+            )
+
+
+class TestPartialCombinationValidity(TestCase):
+    """
+    Tests that verify invalid partial combinations are correctly detected.
+
+    These tests would FAIL if the partial creation logic were simplified
+    (e.g., using symmetric splits or ignoring tensor_idx).
+    """
+
+    world_size = 2
+
+    def setUp(self):
+        super().setUp()
+        if not dist.is_initialized():
+            dist.init_process_group("fake", rank=0, world_size=self.world_size)
+
+    def tearDown(self):
+        super().tearDown()
+        if dist.is_initialized():
+            dist.destroy_process_group()
+
+    def test_abs_psum_psum_is_invalid(self):
+        """
+        P(sum)->P(sum) is NOT valid for abs (or any non-linear elementwise op).
+
+        abs(a0 + a1) != abs(a0) + abs(a1) when a0 and a1 have different signs.
+        The _create_partial_input function must create local values with mixed
+        signs to expose this, not just proportional splits (tensor * ratio)
+        that preserve the sign pattern.
+        """
+        t = torch.tensor([1.0, -2.0, 3.0, -4.0, 5.0, -6.0, 7.0, -8.0])
+        sample = SampleInput(t)
+        tensors = extract_tensors_from_sample(sample)
+        ground_truth = torch.abs(t)
+
+        combo = ((Partial("sum"),), (Partial("sum"),))
+
+        with LocalTensorMode(frozenset(range(self.world_size))):
+            mesh = init_device_mesh("cpu", (self.world_size,))
+            is_valid, msg = validate_combination(
+                torch.abs, sample, tensors, combo, ground_truth, self.world_size, mesh
+            )
+
+        self.assertFalse(
+            is_valid,
+            "abs P(sum)->P(sum) should be invalid (abs is non-linear)",
+        )
+
+    def test_abs_pavg_pavg_is_invalid(self):
+        """P(avg)->P(avg) is NOT valid for abs, same reasoning as P(sum)."""
+        t = torch.tensor([1.0, -2.0, 3.0, -4.0, 5.0, -6.0, 7.0, -8.0])
+        sample = SampleInput(t)
+        tensors = extract_tensors_from_sample(sample)
+        ground_truth = torch.abs(t)
+
+        combo = ((Partial("avg"),), (Partial("avg"),))
+
+        with LocalTensorMode(frozenset(range(self.world_size))):
+            mesh = init_device_mesh("cpu", (self.world_size,))
+            is_valid, msg = validate_combination(
+                torch.abs, sample, tensors, combo, ground_truth, self.world_size, mesh
+            )
+
+        self.assertFalse(
+            is_valid,
+            "abs P(avg)->P(avg) should be invalid (abs is non-linear)",
+        )
+
+    def test_all_zero_output_false_positive(self):
+        """
+        All-zero ground truth makes every placement trivially validate.
+
+        Zeros are a fixed point of all reduce operations (sum, max, min),
+        so validate_combination cannot distinguish valid from invalid rules.
+        This is a known limitation: compare_operator skips all-zero samples
+        to avoid hundreds of false positive rules.
+
+        We use zeros_like as the test op because it always produces zeros
+        regardless of input values, unlike mul(zeros, x) which the offset
+        fix in _create_partial_input now correctly handles.
+        """
+        t = torch.randn(8, 4)
+        sample = SampleInput(t)
+        tensors = extract_tensors_from_sample(sample)
+        ground_truth = torch.zeros_like(t)
+
+        # P(sum)->P(sum) is NOT valid for zeros_like, but with all-zero
+        # output it trivially passes: sum(0,0)=0=ground_truth.
+        combo = ((Partial("sum"),), (Partial("sum"),))
+
+        with LocalTensorMode(frozenset(range(self.world_size))):
+            mesh = init_device_mesh("cpu", (self.world_size,))
+            is_valid, msg = validate_combination(
+                torch.zeros_like,
+                sample,
+                tensors,
+                combo,
+                ground_truth,
+                self.world_size,
+                mesh,
+            )
+
+        # This demonstrates the false positive: validate_combination says
+        # valid, even though the rule is invalid for non-zero inputs.
+        self.assertTrue(
+            is_valid,
+            "Expected True (false positive) for all-zero output, showing "
+            "why compare_operator must skip such samples",
+        )
+
+    def test_argmin_rejects_pmax_input(self):
+        """
+        P(max) -> R must be rejected for argmin.
+
+        argmin returns indices, not values. With P(max) input, different
+        ranks hold the true value at different positions (checkerboard mask),
+        so argmin produces different indices on each rank. The Replicate
+        output check should catch this disagreement.
+
+        This test exercises both fixes:
+        1. Adaptive offset (exceeds value range) — prevents the true min
+           from dominating on all ranks despite the perturbation
+        2. Checkerboard mask — prevents uniform parity along even-stride
+           dimensions, which would make all elements in a reduction group
+           shift identically
+        """
+        # Use a shape where flat-index mask would fail: dim 0 stride = 4 (even)
+        t = torch.randn(3, 2, 1, 2)
+        sample = SampleInput(t, kwargs={"dim": 0, "keepdim": True})
+        tensors = extract_tensors_from_sample(sample)
+        ground_truth = torch.argmin(t, dim=0, keepdim=True)
+
+        for reduce_op in ("min", "max"):
+            combo = ((Partial(reduce_op),), (Replicate(),))
+            with LocalTensorMode(frozenset(range(self.world_size))):
+                mesh = init_device_mesh("cpu", (self.world_size,))
+                is_valid, msg = validate_combination(
+                    torch.argmin,
+                    sample,
+                    tensors,
+                    combo,
+                    ground_truth,
+                    self.world_size,
+                    mesh,
+                )
+            self.assertFalse(
+                is_valid,
+                f"argmin P({reduce_op})->R should be invalid "
+                f"(index op, ranks disagree): {msg}",
+            )
+
+
+class TestDecompStrategyPath(TestCase):
+    """Test that decomposition-based strategy propagation discovers rules."""
+
+    world_size = 2
+
+    def setUp(self):
+        super().setUp()
+        if not dist.is_initialized():
+            dist.init_process_group("fake", rank=0, world_size=self.world_size)
+
+    def tearDown(self):
+        super().tearDown()
+        if dist.is_initialized():
+            dist.destroy_process_group()
+
+    def test_decomp_discovers_rules_for_softplus(self):
+        """
+        Exercise DecompShardingStrategy.propagate_strategy directly on a unary
+        op (softplus) whose decomposition uses only registered ops.
+
+        Verifies the helper _extract_rules_from_op_strategy correctly extracts
+        elementwise sharding rules from the resulting OpStrategy.
+        """
+        try:
+            from torch.distributed.tensor._decompositions import DecompShardingStrategy
+        except ImportError:
+            self.skipTest("_decompositions module not available")
+        from torch.distributed.tensor._dtensor_spec import TensorMeta
+        from torch.distributed.tensor._op_schema import DTensorSpec, OpSchema
+        from torch.distributed.tensor._ops.strategy_validation import (
+            _extract_rules_from_op_strategy,
+        )
+
+        aten_softplus = torch.ops.aten.softplus.default
+        propagator = torch.distributed.tensor.DTensor._op_dispatcher.sharding_propagator
+        self.assertTrue(DecompShardingStrategy.has_decomp(aten_softplus))
+
+        t = torch.randn(8, 4)
+
+        mesh = init_device_mesh("cpu", (self.world_size,))
+        spec = DTensorSpec(
+            mesh=mesh,
+            placements=(Shard(0),),
+            tensor_meta=TensorMeta(shape=t.shape, stride=t.stride(), dtype=t.dtype),
+        )
+        op_schema = OpSchema(aten_softplus, (spec,), {})
+        propagator.decomp_strategy.ensure_schema_info(aten_softplus)
+        output_strategy = propagator.decomp_strategy.propagate_strategy(
+            op_schema,
+        )
+        self.assertIsNotNone(output_strategy)
+
+        input_shapes = (t.shape,)
+        output_shapes = (tuple(torch.nn.functional.softplus(t).shape),)
+        rules = _extract_rules_from_op_strategy(
+            output_strategy, input_shapes, output_shapes
+        )
+
+        # Should discover elementwise sharding rules for a 2D tensor
+        self.assertIn((("S(0)",), ("S(0)",)), rules)
+        self.assertIn((("S(1)",), ("S(1)",)), rules)
+
+    def test_compare_operator_uses_decomp_path(self):
+        """
+        compare_operator should find rules via the decomp path when an op is
+        not registered in op_strategy_funcs or op_single_dim_strategy_funcs.
+
+        Temporarily removes addcmul from strategy registries to force the
+        decomp fallback, and uses even tensor sizes for correct sharding.
+        """
+        try:
+            from torch.distributed.tensor._decompositions import (  # noqa: F401
+                DecompShardingStrategy,
+            )
+        except ImportError:
+            self.skipTest("_decompositions module not available")
+
+        import torch.testing._internal.common_methods_invocations as common_ops
+        from torch.distributed.tensor._ops.strategy_validation import compare_operator
+        from torch.testing._internal.opinfo import core as opinfo_core
+
+        propagator = torch.distributed.tensor.DTensor._op_dispatcher.sharding_propagator
+        aten_addcmul = torch.ops.aten.addcmul.default
+
+        # Save and remove addcmul from strategy registries to force decomp path
+        saved_strategy = propagator.op_strategy_funcs.pop(aten_addcmul, None)
+        saved_single = propagator.op_single_dim_strategy_funcs.pop(aten_addcmul, None)
+
+        # Override sizes to ensure even sharding with world_size=2
+        orig_sizes = (opinfo_core.L, opinfo_core.M, opinfo_core.S, opinfo_core.XS)
+        opinfo_core.L = common_ops.L = 24
+        opinfo_core.M = common_ops.M = 12
+        opinfo_core.S = common_ops.S = 4
+        opinfo_core.XS = common_ops.XS = 2
+
+        try:
+            # compare_operator manages its own process group
+            stats = compare_operator(
+                "addcmul",
+                device="cpu",
+                dtype=torch.float32,
+                world_size=self.world_size,
+                incorrect_only=True,
+            )
+            # Decomp should discover valid rules
+            self.assertGreater(stats.true_positives, 0)
+            # No incorrect rules
+            self.assertEqual(len(stats.false_positives), 0)
+        finally:
+            # Restore registries and sizes
+            if saved_strategy is not None:
+                propagator.op_strategy_funcs[aten_addcmul] = saved_strategy
+            if saved_single is not None:
+                propagator.op_single_dim_strategy_funcs[aten_addcmul] = saved_single
+            (
+                opinfo_core.L,
+                opinfo_core.M,
+                opinfo_core.S,
+                opinfo_core.XS,
+            ) = orig_sizes
+            (
+                common_ops.L,
+                common_ops.M,
+                common_ops.S,
+                common_ops.XS,
+            ) = orig_sizes
+
+
+class TestQuerySingleDimStrategyKwargs(TestCase):
+    """Test that query_single_dim_strategy forwards kwargs to strategy functions."""
+
+    def test_kwargs_forwarded_to_strategy(self):
+        """
+        A kwargs-aware strategy should receive the actual kwargs, not {}.
+
+        torch.add(a, b, alpha=-1) changes which Partial rules are valid:
+        with alpha=1, R,P(max)->P(max) is valid; with alpha=-1, it becomes
+        R,P(max)->P(min) instead. A strategy that accounts for alpha needs
+        to receive it through kwargs.
+        """
+        from torch.distributed.tensor._api import DTensor
+        from torch.distributed.tensor._ops.single_dim_strategy import (
+            _ShardingPlaceholder,
+        )
+
+        propagator = DTensor._op_dispatcher.sharding_propagator
+        aten_add = torch.ops.aten.add.Tensor
+
+        # A strategy that returns different rules depending on alpha.
+        # With alpha >= 0: R,P(max)->P(max) is valid
+        # With alpha < 0:  R,P(max)->P(min) is valid (negation flips max to min)
+        def alpha_aware_add_strategy(op, args_schema, kwargs_schema):
+            alpha = kwargs_schema.get("alpha", 1)
+            rules = [
+                [
+                    _ShardingPlaceholder(0),
+                    _ShardingPlaceholder(0),
+                    _ShardingPlaceholder(0),
+                ],
+                [Partial("sum"), Partial("sum"), Partial("sum")],
+            ]
+            if alpha < 0:
+                rules.append([Partial("min"), Replicate(), Partial("max")])
+            else:
+                rules.append([Partial("max"), Replicate(), Partial("max")])
+            return rules
+
+        original = propagator.op_single_dim_strategy_funcs.get(aten_add)
+        propagator.op_single_dim_strategy_funcs[aten_add] = alpha_aware_add_strategy
+        try:
+            a, b = torch.randn(4, 3), torch.randn(4, 3)
+
+            # Query with alpha=-1 kwargs
+            result = query_single_dim_strategy(aten_add, (a, b), {"alpha": -1})
+            self.assertIsNotNone(result)
+
+            # The third rule's output should be P(min) for alpha=-1
+            self.assertEqual(len(result), 3)
+            self.assertIsInstance(result[2][0], Partial)
+            self.assertEqual(
+                result[2][0].reduce_op,
+                "min",
+                "With alpha=-1, the strategy should produce P(min) output "
+                "but got P(max) — kwargs were not forwarded",
+            )
+        finally:
+            if original is not None:
+                propagator.op_single_dim_strategy_funcs[aten_add] = original
+            else:
+                propagator.op_single_dim_strategy_funcs.pop(aten_add, None)
+
+
+class TestCompareRules(TestCase):
+    """Test _compare_rules populates per-op statistics correctly."""
+
+    def test_true_positives_tracked_by_op(self):
+        from torch.distributed.tensor._ops.strategy_validation import (
+            _compare_rules,
+            ComparisonStats,
+        )
+
+        combo_a = (("S(0)",), ("R",))
+        combo_b = (("S(1)",), ("R",))
+        stats = ComparisonStats()
+        _compare_rules(
+            ground_truth_valid={combo_a, combo_b},
+            dtensor_rules={combo_a, combo_b},
+            input_shapes=((4, 8),),
+            output_shapes=((4, 8),),
+            sample_idx=0,
+            scalar_args=(),
+            scalar_kwargs={},
+            aten_op=torch.ops.aten.relu.default,
+            variant="",
+            stats=stats,
+        )
+        self.assertEqual(stats.true_positives, 2)
+        self.assertEqual(stats.true_positives_by_op, {"aten.relu.default": 2})
+        self.assertEqual(len(stats.false_positives), 0)
+        self.assertEqual(len(stats.false_negatives), 0)
+
+    def test_per_op_breakdown_multiple_aten_ops(self):
+        """Simulate an opinfo that dispatches to different aten variants per sample."""
+        from torch.distributed.tensor._ops.strategy_validation import (
+            _compare_rules,
+            ComparisonStats,
+        )
+
+        stats = ComparisonStats()
+
+        # Sample 0: aten.max.dim returns (values, indices)
+        combo_2out = (("S(0)",), ("R", "R"))
+        _compare_rules(
+            ground_truth_valid={combo_2out},
+            dtensor_rules={combo_2out},
+            input_shapes=((4, 8),),
+            output_shapes=((4, 8), (4, 8)),
+            sample_idx=0,
+            scalar_args=(1,),
+            scalar_kwargs={},
+            aten_op=torch.ops.aten.max.dim,
+            variant="",
+            stats=stats,
+        )
+        # Sample 1: aten.native_layer_norm returns (output, mean, rstd)
+        combo_3out = (("S(0)",), ("R", "R", "R"))
+        _compare_rules(
+            ground_truth_valid={combo_3out},
+            dtensor_rules={combo_3out},
+            input_shapes=((4, 8),),
+            output_shapes=((4, 8), (4,), (4,)),
+            sample_idx=1,
+            scalar_args=(),
+            scalar_kwargs={},
+            aten_op=torch.ops.aten.native_layer_norm.default,
+            variant="",
+            stats=stats,
+        )
+
+        self.assertEqual(stats.true_positives, 2)
+        self.assertEqual(
+            stats.true_positives_by_op,
+            {"aten.max.dim": 1, "aten.native_layer_norm.default": 1},
+        )
+
+    def test_false_positives_and_negatives_not_in_by_op(self):
+        """true_positives_by_op only tracks true positives, not FP/FN."""
+        from torch.distributed.tensor._ops.strategy_validation import (
+            _compare_rules,
+            ComparisonStats,
+        )
+
+        gt_only = (("S(0)",), ("R",))
+        dt_only = (("S(1)",), ("R",))
+        stats = ComparisonStats()
+        _compare_rules(
+            ground_truth_valid={gt_only},
+            dtensor_rules={dt_only},
+            input_shapes=((4, 8),),
+            output_shapes=((4, 8),),
+            sample_idx=0,
+            scalar_args=(),
+            scalar_kwargs={},
+            aten_op=torch.ops.aten.relu.default,
+            variant="",
+            stats=stats,
+        )
+        self.assertEqual(stats.true_positives, 0)
+        self.assertEqual(stats.true_positives_by_op, {})
+        self.assertEqual(len(stats.false_positives), 1)
+        self.assertEqual(len(stats.false_negatives), 1)
+
+
+class TestCompareOperatorEndToEnd(TestCase):
+    """End-to-end smoke test for compare_operator."""
+
+    world_size = 2
+
+    def setUp(self):
+        super().setUp()
+        if not dist.is_initialized():
+            dist.init_process_group("fake", rank=0, world_size=self.world_size)
+
+    def tearDown(self):
+        super().tearDown()
+        if dist.is_initialized():
+            dist.destroy_process_group()
+
+    def _with_even_sizes(self, fn):
+        """Run fn with opinfo sizes overridden to be evenly divisible by world_size."""
+        import torch.testing._internal.common_methods_invocations as common_ops
+        from torch.testing._internal.opinfo import core as opinfo_core
+
+        orig_sizes = (opinfo_core.L, opinfo_core.M, opinfo_core.S, opinfo_core.XS)
+        opinfo_core.L = common_ops.L = 24
+        opinfo_core.M = common_ops.M = 12
+        opinfo_core.S = common_ops.S = 4
+        opinfo_core.XS = common_ops.XS = 2
+        try:
+            return fn()
+        finally:
+            (
+                opinfo_core.L,
+                opinfo_core.M,
+                opinfo_core.S,
+                opinfo_core.XS,
+            ) = orig_sizes
+            (
+                common_ops.L,
+                common_ops.M,
+                common_ops.S,
+                common_ops.XS,
+            ) = orig_sizes
+
+    def test_compare_operator_add_no_incorrect(self):
+        """compare_operator on add should find valid rules with no incorrect ones."""
+        from torch.distributed.tensor._ops.strategy_validation import compare_operator
+
+        def run():
+            stats = compare_operator(
+                "add",
+                device="cpu",
+                dtype=torch.float32,
+                world_size=self.world_size,
+                incorrect_only=True,
+            )
+            self.assertGreater(stats.true_positives, 0)
+            self.assertEqual(len(stats.false_positives), 0)
+
+        self._with_even_sizes(run)
+
+    def test_compare_operator_split_multi_output(self):
+        """compare_operator should handle multi-output ops like split."""
+        from torch.distributed.tensor._ops.strategy_validation import compare_operator
+
+        def run():
+            stats = compare_operator(
+                "split",
+                device="cpu",
+                dtype=torch.float32,
+                world_size=self.world_size,
+                incorrect_only=True,
+            )
+            self.assertGreater(
+                stats.total_samples,
+                0,
+                f"split should have runnable samples, got skip_reasons={stats.skip_reasons}",
+            )
+            self.assertEqual(len(stats.false_positives), 0)
+
+        self._with_even_sizes(run)
+
+    def test_compare_operator_runtime_schema_ops(self):
+        """Ops with RuntimeSchemaInfo (non-tensor positional args) should find DTensor rules."""
+        from torch.distributed.tensor._ops.strategy_validation import compare_operator
+
+        runtime_schema_ops = ["flip", "roll"]
+        for op_name in runtime_schema_ops:
+            with self.subTest(op=op_name):
+
+                def run(name=op_name):
+                    stats = compare_operator(
+                        name,
+                        device="cpu",
+                        dtype=torch.float32,
+                        world_size=self.world_size,
+                        incorrect_only=True,
+                    )
+                    self.assertGreater(
+                        stats.true_positives,
+                        0,
+                        f"{name} should have DTensor rules (true_positives > 0), "
+                        f"got skip_reasons={stats.skip_reasons}",
+                    )
+
+                self._with_even_sizes(run)
+
+
+class TestAtenLevelValidation(TestCase):
+    """Tests for aten-level validation (validate_aten_combination + allow_composite mode)."""
+
+    world_size = 2
+
+    def setUp(self):
+        super().setUp()
+        if not dist.is_initialized():
+            dist.init_process_group("fake", rank=0, world_size=self.world_size)
+
+    def tearDown(self):
+        super().tearDown()
+        if dist.is_initialized():
+            dist.destroy_process_group()
+
+    def _with_even_sizes(self, fn):
+        import torch.testing._internal.common_methods_invocations as common_ops
+        from torch.testing._internal.opinfo import core as opinfo_core
+
+        orig_sizes = (opinfo_core.L, opinfo_core.M, opinfo_core.S, opinfo_core.XS)
+        opinfo_core.L = common_ops.L = 24
+        opinfo_core.M = common_ops.M = 12
+        opinfo_core.S = common_ops.S = 4
+        opinfo_core.XS = common_ops.XS = 2
+        try:
+            return fn()
+        finally:
+            (
+                opinfo_core.L,
+                opinfo_core.M,
+                opinfo_core.S,
+                opinfo_core.XS,
+            ) = orig_sizes
+            (
+                common_ops.L,
+                common_ops.M,
+                common_ops.S,
+                common_ops.XS,
+            ) = orig_sizes
+
+    def test_validate_aten_combination_shard_add(self):
+        """S(0), S(0) -> S(0) valid; S(0), S(0) -> R invalid for aten.add."""
+        a = torch.randn(8)
+        b = torch.randn(8)
+        gt = torch.ops.aten.add.Tensor(a, b)
+
+        with LocalTensorMode(frozenset(range(self.world_size))):
+            mesh = init_device_mesh("cpu", (self.world_size,))
+
+            valid, msg = validate_aten_combination(
+                torch.ops.aten.add.Tensor,
+                (a, b),
+                {},
+                gt,
+                ((Shard(0), Shard(0)), (Shard(0),)),
+                self.world_size,
+                mesh,
+            )
+            self.assertTrue(valid, msg)
+
+            valid, _ = validate_aten_combination(
+                torch.ops.aten.add.Tensor,
+                (a, b),
+                {},
+                gt,
+                ((Shard(0), Shard(0)), (Replicate(),)),
+                self.world_size,
+                mesh,
+            )
+            self.assertFalse(valid)
+
+    def test_allow_composite_eliminates_non_1to1_skips(self):
+        """allow_composite mode should validate decomposed ops that default mode skips."""
+        from torch.distributed.tensor._ops.strategy_validation import compare_operator
+
+        def run():
+            default = compare_operator(
+                "inner",
+                device="cpu",
+                dtype=torch.float32,
+                world_size=self.world_size,
+                max_samples=3,
+                incorrect_only=True,
+            )
+            allow_composite = compare_operator(
+                "inner",
+                device="cpu",
+                dtype=torch.float32,
+                world_size=self.world_size,
+                max_samples=3,
+                incorrect_only=True,
+                allow_composite=True,
+            )
+            self.assertGreater(default.skip_reasons.get("non-1:1 aten mapping", 0), 0)
+            self.assertEqual(
+                allow_composite.skip_reasons.get("non-1:1 aten mapping", 0), 0
+            )
+            self.assertGreaterEqual(
+                allow_composite.total_samples, default.total_samples
+            )
+
+        self._with_even_sizes(run)
+
+
+class TestMainModule(TestCase):
+    """Test that strategy_validation can be run as a main module."""
+
+    def _run_module(self, *extra_args):
+        import subprocess
+        import sys
+
+        result = subprocess.run(
+            [
+                sys.executable,
+                "-m",
+                "torch.distributed.tensor._ops.strategy_validation",
+                "--device",
+                "cpu",
+                *extra_args,
+            ],
+            capture_output=True,
+            text=True,
+            timeout=120,
+        )
+        self.assertEqual(
+            result.returncode,
+            0,
+            f"Module exited with code {result.returncode}.\n"
+            f"stderr: {result.stderr[-2000:]}",
+        )
+        return result.stdout
+
+    def test_run_as_module_default(self):
+        """Running with no args should validate 'add' and exit cleanly."""
+        stdout = self._run_module()
+        self.assertIn("add", stdout)
+        self.assertIn("Correct", stdout)
+
+    def test_run_with_op_flag(self):
+        """Running with --op mul --incorrect-only should work."""
+        stdout = self._run_module("--op", "mul", "--incorrect-only")
+        self.assertIn("mul", stdout)
+
+
+class TestOpInfoLookup(TestCase):
+    """Tests for get_opinfo_by_name, _find_opinfo_candidates, and resolve_op_names."""
+
+    def test_get_opinfo_by_name_exact(self):
+        results = get_opinfo_by_name("add")
+        self.assertGreater(len(results), 0)
+        for op in results:
+            self.assertEqual(op.name, "add")
+
+    def test_get_opinfo_by_name_qualified(self):
+        results = get_opinfo_by_name("nn.functional.relu")
+        self.assertGreater(len(results), 0)
+        for op in results:
+            self.assertEqual(op.name, "nn.functional.relu")
+
+    def test_get_opinfo_by_name_not_found_with_suggestions(self):
+        with self.assertRaises(ValueError) as ctx:
+            get_opinfo_by_name("relu")
+        self.assertIn("did you mean", str(ctx.exception))
+        self.assertIn("nn.functional.relu", str(ctx.exception))
+
+    def test_get_opinfo_by_name_not_found_no_suggestions(self):
+        with self.assertRaises(ValueError) as ctx:
+            get_opinfo_by_name("this_op_does_not_exist_xyz")
+        self.assertNotIn("did you mean", str(ctx.exception))
+
+    def test_find_opinfo_candidates_aten_name(self):
+        # relu OpInfo has aten_name="relu" explicitly set
+        candidates = _find_opinfo_candidates("relu")
+        self.assertIn("nn.functional.relu", candidates)
+
+    def test_find_opinfo_candidates_suffix(self):
+        candidates = _find_opinfo_candidates("dropout")
+        self.assertIn("nn.functional.dropout", candidates)
+
+    def test_find_opinfo_candidates_no_match(self):
+        candidates = _find_opinfo_candidates("this_op_does_not_exist_xyz")
+        self.assertEqual(candidates, [])
+
+    def test_resolve_op_names_exact(self):
+        result = resolve_op_names(["add"])
+        self.assertEqual(result, ["add"])
+
+    def test_resolve_op_names_qualified(self):
+        result = resolve_op_names(["nn.functional.relu"])
+        self.assertEqual(result, ["nn.functional.relu"])
+
+    def test_resolve_op_names_multiple(self):
+        result = resolve_op_names(["add", "mul"])
+        self.assertEqual(result, ["add", "mul"])
+
+    def test_resolve_op_names_deduplicates(self):
+        result = resolve_op_names(["add", "add"])
+        self.assertEqual(result, ["add"])
+
+    def test_resolve_op_names_glob(self):
+        result = resolve_op_names(["nn.functional.relu*"])
+        self.assertIn("nn.functional.relu", result)
+
+    def test_resolve_op_names_glob_no_match(self):
+        with self.assertRaises(ValueError) as ctx:
+            resolve_op_names(["zzz_no_match_*"])
+        self.assertIn("No OpInfo names match", str(ctx.exception))
+
+    def test_resolve_op_names_unambiguous_shorthand(self):
+        # "relu" should resolve unambiguously to "nn.functional.relu"
+        result = resolve_op_names(["nn.functional.relu"])
+        self.assertIn("nn.functional.relu", result)
+
+    def test_resolve_op_names_ambiguous_shorthand(self):
+        # "dropout" matches nn.functional.dropout, nn.functional.dropout2d, etc.
+        candidates = _find_opinfo_candidates("dropout")
+        if len(candidates) > 1:
+            with self.assertRaises(ValueError) as ctx:
+                resolve_op_names(["dropout"])
+            self.assertIn("ambiguous", str(ctx.exception))
+
+    def test_resolve_op_names_not_found(self):
+        with self.assertRaises(ValueError):
+            resolve_op_names(["this_op_does_not_exist_xyz"])
+
+
 if __name__ == "__main__":
     run_tests()
diff --git a/test/distributed/tensor/test_tensor_ops.py b/test/distributed/tensor/test_tensor_ops.py
index c59c7aba540d5..f5ab8e6844e93 100644
--- a/test/distributed/tensor/test_tensor_ops.py
+++ b/test/distributed/tensor/test_tensor_ops.py
@@ -20,22 +20,28 @@
 from torch.distributed.tensor.debug import CommDebugMode
 from torch.testing._internal.common_distributed import skip_if_lt_x_gpu
 from torch.testing._internal.common_utils import (
+    instantiate_parametrized_tests,
     MI200_ARCH,
+    parametrize,
     run_tests,
     serialTest,
     skipIfRocmArch,
 )
 from torch.testing._internal.distributed._tensor.common_dtensor import (
     create_local_tensor_test_class,
+    DTensorContinuousTestBase,
     DTensorConverter,
     DTensorTestBase,
+    LocalDTensorContinuousTestBase,
     LocalDTensorTestBase,
+    op_strategy_context,
     with_comms,
 )
 
 
-class DistTensorOpsTest(DTensorTestBase):
-    @with_comms
+class DistTensorOpsTest(DTensorContinuousTestBase):
+    world_size = 4
+
     def test_aten_contiguous(self):
         # this op not covered by dtensor_ops
         mesh = self.build_device_mesh()
@@ -45,7 +51,6 @@ def test_aten_contiguous(self):
             torch.randn(16, 32),
         )
 
-    @with_comms
     def test_detach(self):
         device_mesh = self.build_device_mesh()
         shard_spec = [Shard(0)]
@@ -55,7 +60,6 @@ def test_detach(self):
         detached_mat = mat.detach()
         self.assertFalse(detached_mat is mat)
 
-    @with_comms
     def test_clone(self):
         device_mesh = self.build_device_mesh()
         specs = [[Replicate()], [Shard(0)]]
@@ -66,7 +70,6 @@ def test_clone(self):
             self.assertFalse(cloned_mat is mat)
             self.assertEqual(cloned_mat.to_local(), mat.to_local())
 
-    @with_comms
     def test_copy_(self):
         device_mesh = DeviceMesh(self.device_type, list(range(self.world_size)))
 
@@ -135,7 +138,6 @@ def test_copy_(self):
             self.assertEqual(dst_dtensor.placements, (partial_placement,))
             self.assertEqual(dst_dtensor._local_tensor, dst_tensor)
 
-    @with_comms
     def test_contiguous(self):
         device_mesh = self.build_device_mesh()
         tensor = torch.rand(3, 5, 6, requires_grad=True)
@@ -161,7 +163,6 @@ def test_contiguous(self):
         new_dt.to_local().sum().backward()
         self.assertEqual(tensor.grad, torch.ones(3, 5, 6))
 
-    @with_comms
     def test_inplace_op(self):
         mesh = self.build_device_mesh()
         input_tensor = torch.randn((12, 3), device=self.device_type)
@@ -188,7 +189,6 @@ def test_inplace_op(self):
         self.assertTrue(res is dt_to_inplace_add)
         self.assertTrue(res.placements == tuple(shard_spec))
 
-    @with_comms
     def test_op_out_variant(self):
         mesh = self.build_device_mesh()
         input_tensor = torch.randn((12, 3), device=self.device_type)
@@ -209,7 +209,6 @@ def test_op_out_variant(self):
         self.assertTrue(res.placements == tuple(replica_spec))
         self.assertEqual(replicate_out.to_local(), expected_dt.to_local())
 
-    @with_comms
     def test_empty_like(self):
         device_mesh = self.build_device_mesh()
         shard_spec = [Shard(0)]
@@ -220,7 +219,6 @@ def test_empty_like(self):
         # empty is not deterministic, so we only check that the shard propagation worked
         self.assertEqual((4, 8), empty_like_dt.to_local().shape)
 
-    @with_comms
     def test_meta_init_partial(self):
         device_mesh = self.build_device_mesh()
         partial_spec = [Partial()]
@@ -244,7 +242,6 @@ def __init__(self):
         # check that to_empty preserves partial
         self.assertEqual(module.scalar_buffer.placements, (Partial(),))
 
-    @with_comms
     def test_fill_inplace(self):
         device_mesh = self.build_device_mesh()
         shard_spec = [Shard(0)]
@@ -256,7 +253,6 @@ def test_fill_inplace(self):
         self.assertEqual(full_expected, full_like_dt.to_local())
         self.assertEqual(full_expected, dist_tensor.to_local())
 
-    @with_comms
     def test_full_like(self):
         device_mesh = self.build_device_mesh()
         shard_spec = [Shard(0)]
@@ -267,7 +263,6 @@ def test_full_like(self):
         full_expected = torch.full((4, 8), 42.0)
         self.assertEqual(full_expected, full_like_dt.to_local())
 
-    @with_comms
     def test_ones_like(self):
         device_mesh = self.build_device_mesh()
         shard_spec = [Shard(0)]
@@ -278,27 +273,27 @@ def test_ones_like(self):
         ones_expected = torch.ones(4, 8)
         self.assertEqual(ones_expected, ones_like_dt.to_local())
 
-    @with_comms
     def test_ones_like_partial_sum(self):
         device_mesh = self.build_device_mesh()
         shard_spec = [Partial()]
 
         input_tensor = torch.randn(4, 8, requires_grad=True)
         dist_tensor = DTensor.from_local(input_tensor, device_mesh, shard_spec)
-        assert dist_tensor.shape == (4, 8)
+        if not (dist_tensor.shape == (4, 8)):
+            raise AssertionError(f"Expected shape (4, 8), got {dist_tensor.shape}")
 
         ones_like_dt = torch.ones_like(dist_tensor)
         ones_expected = torch.ones(dist_tensor.shape)
         self.assertEqual(ones_expected, ones_like_dt.full_tensor())
 
-    @with_comms
     def test_fill_inplace_partial_sum(self):
         device_mesh = self.build_device_mesh()
         shard_spec = [Partial()]
 
         input_tensor = torch.randn(4, 8, requires_grad=True)
         dist_tensor = DTensor.from_local(input_tensor, device_mesh, shard_spec)
-        assert dist_tensor.shape == (4, 8)
+        if not (dist_tensor.shape == (4, 8)):
+            raise AssertionError(f"Expected shape (4, 8), got {dist_tensor.shape}")
 
         # inplace partial sum should keep partial
         torch.fill_(dist_tensor, 8)
@@ -307,20 +302,19 @@ def test_fill_inplace_partial_sum(self):
         )
         self.assertEqual(fill_expected, dist_tensor.full_tensor())
 
-    @with_comms
     def test_zeros_like_partial_sum(self):
         device_mesh = self.build_device_mesh()
         shard_spec = [Partial()]
 
         input_tensor = torch.randn(4, 8, requires_grad=True)
         dist_tensor = DTensor.from_local(input_tensor, device_mesh, shard_spec)
-        assert dist_tensor.shape == (4, 8)
+        if not (dist_tensor.shape == (4, 8)):
+            raise AssertionError(f"Expected shape (4, 8), got {dist_tensor.shape}")
 
         zeros_like_dt = torch.zeros_like(dist_tensor)
         zeros_expected = torch.zeros(dist_tensor.shape)
         self.assertEqual(zeros_expected, zeros_like_dt.full_tensor())
 
-    @with_comms
     def test_zero_inplace(self):
         device_mesh = self.build_device_mesh()
         shard_spec = [Shard(0)]
@@ -332,7 +326,6 @@ def test_zero_inplace(self):
         self.assertEqual(zeros_expected, zeros_like_dt.to_local())
         self.assertEqual(zeros_expected, dist_tensor.to_local())
 
-    @with_comms
     def test_zeros_like(self):
         device_mesh = self.build_device_mesh()
         shard_spec = [Shard(0)]
@@ -347,7 +340,6 @@ def test_zeros_like(self):
         self.assertEqual(zeros_like_dt.dtype, torch.bfloat16)
 
     @skip_if_lt_x_gpu(4)
-    @with_comms
     def test_stack(self):
         mesh_2d = DeviceMesh(
             self.device_type, torch.arange(self.world_size).reshape(2, 2)
@@ -384,7 +376,16 @@ def test_stack(self):
             torch.stack([global_input, global_input], dim=1),
         )
 
-    @with_comms
+        # stack with negative dim: dim=-1 inserts at the last position of the
+        # output (ndim+1), so Shard(1) should stay Shard(1)
+        stack_neg_dim_dt = torch.stack([shard1_input, cloned_shard1_input], dim=-1)
+        self.assertEqual(stack_neg_dim_dt.placements, (Shard(1),))
+        self.assertEqual(stack_neg_dim_dt.shape, (8, 8, 2))
+        self.assertEqual(
+            stack_neg_dim_dt.full_tensor(),
+            torch.stack([global_input, global_input], dim=-1),
+        )
+
     def test_stack_cache(self):
         device_mesh = self.build_device_mesh()
 
@@ -412,7 +413,6 @@ def error(*args, **kwargs):
         ):
             _ = torch.stack(dtensor_list2)
 
-    @with_comms
     def test_equal(self):
         device_mesh = self.build_device_mesh()
         shard_spec = [Shard(0)]
@@ -454,6 +454,13 @@ def test_equal(self):
         self.assertTrue(dist_tensor_1.is_same_size(dist_tensor_3))
         self.assertFalse(input_tensor_2.is_same_size(dist_tensor_3))
 
+    def test_is_pinned(self):
+        device_mesh = self.build_device_mesh()
+        shard_spec = [Shard(0)]
+
+        dt = DTensor.from_local(torch.ones(4, 4), device_mesh, shard_spec)
+        self.assertFalse(dt.is_pinned())
+
     def _test_op(self, mesh, op_call, *args, **kwargs):
         out = op_call(*args, **kwargs)
         dtc = DTensorConverter(mesh, args, kwargs)
@@ -462,7 +469,6 @@ def _test_op(self, mesh, op_call, *args, **kwargs):
             d_out = op_call(*d_args, **d_kwargs)
             self.assertEqual(d_out.full_tensor(), out)
 
-    @with_comms
     def test_new_full(self):
         device_mesh = self.build_device_mesh()
         comm_mode = CommDebugMode()
@@ -489,7 +495,6 @@ def test_new_full(self):
             self.assertEqual(new_full_same_dt.placements, placement)
             self.assertEqual(new_full_same_expected, new_full_same_dt.full_tensor())
 
-    @with_comms
     def test_new_empty_strided(self):
         device_mesh = self.build_device_mesh()
         comm_mode = CommDebugMode()
@@ -513,16 +518,22 @@ def test_new_empty_strided(self):
         )
         self.assertTrue(new_empty_strided_dt.contiguous() is new_empty_strided_dt)
 
-        # output shape same as input shape, unevenly sharded input -> output replicated
+        # output shape same as input shape, unevenly sharded input -> output follows input
         global_tensor = torch.randn(12, 7)
         input_dt = distribute_tensor(global_tensor, device_mesh, placement)
         self.assertTrue(input_dt.shape[shard_dim] % self.world_size != 0)
         with comm_mode:
             new_empty_strided_dt = input_dt.new_empty_strided((12, 7), (7, 1))
             self.assertEqual(comm_mode.get_total_counts(), 0)
-        self.assertEqual(new_empty_strided_dt.placements, (Replicate(),))
-        self.assertEqual(new_empty_strided_dt._local_tensor.size(), (12, 7))
-        self.assertEqual(new_empty_strided_dt._local_tensor.stride(), (7, 1))
+        self.assertEqual(new_empty_strided_dt.placements, placement)
+        self.assertEqual(
+            new_empty_strided_dt._local_tensor.size(),
+            input_dt._local_tensor.size(),
+        )
+        self.assertEqual(
+            new_empty_strided_dt._local_tensor.stride(),
+            input_dt._local_tensor.stride(),
+        )
 
         # output shape different from input shape -> output replicated
         global_tensor = torch.randn(12, 8)
@@ -534,7 +545,6 @@ def test_new_empty_strided(self):
         self.assertEqual(new_empty_strided_dt._local_tensor.size(), (12, 4))
         self.assertEqual(new_empty_strided_dt._local_tensor.stride(), (4, 1))
 
-    @with_comms
     def test_scatter(self):
         device_mesh = self.build_device_mesh()
         comm_mode = CommDebugMode()
@@ -568,7 +578,6 @@ def test_scatter(self):
                 self.assertEqual(output_dt.placements, [Replicate()])
                 self.assertEqual(output_dt.to_local(), global_output)
 
-    @with_comms
     def test_gather(self):
         device_mesh = self.build_device_mesh()
         comm_mode = CommDebugMode()
@@ -618,8 +627,7 @@ def test_gather(self):
             self.assertEqual(output_dt.full_tensor(), global_output)
 
     @skipIfRocmArch(MI200_ARCH)
-    @with_comms
-    @serialTest()
+    @serialTest()  # heavy combinatorial _test_op calls, serialize to avoid OOM
     def test_index(self):
         meshes = [
             self.build_device_mesh(),  # 1D mesh
@@ -676,20 +684,23 @@ def test_index(self):
                 torch.randint(5, (12, 8, 12)),
                 torch.randint(2, (12, 8, 12)),
             )
-            self._test_op(
-                mesh,
-                lambda x, y, z: x[z, :, y],
-                torch.randn(16, 32, 16),
-                torch.randint(5, (12, 8, 12)),
-                torch.randint(2, (12, 8, 12)),
-            )
-            self._test_op(
-                mesh,
-                lambda x, y, z: x[:, z, :, y],
-                torch.randn(16, 32, 16, 12),
-                torch.randint(5, (12, 8, 12)),
-                torch.randint(2, (12, 8, 12)),
-            )
+            # Commented out to fix distributed CI timeout: each 3-tensor call
+            # generates 40-80 sharding combinations via itertools.product,
+            # causing combinatorial explosion.
+            # self._test_op(
+            #     mesh,
+            #     lambda x, y, z: x[z, :, y],
+            #     torch.randn(16, 32, 16),
+            #     torch.randint(5, (12, 8, 12)),
+            #     torch.randint(2, (12, 8, 12)),
+            # )
+            # self._test_op(
+            #     mesh,
+            #     lambda x, y, z: x[:, z, :, y],
+            #     torch.randn(16, 32, 16, 12),
+            #     torch.randint(5, (12, 8, 12)),
+            #     torch.randint(2, (12, 8, 12)),
+            # )
             # broadcast in inner dimensions
             self._test_op(
                 mesh,
@@ -698,37 +709,39 @@ def test_index(self):
                 torch.randint(5, (12, 8, 12)),
                 torch.randint(2, (12, 1, 12)),
             )
-            # implicit (left-padded) broadcast
-            self._test_op(
-                mesh,
-                lambda x, y, z: x[:, z, :, y],
-                torch.randn(16, 32, 16, 12),
-                torch.randint(5, (12, 8, 12)),
-                torch.randint(2, (8, 12)),
-            )
-            self._test_op(
-                mesh,
-                lambda x, y, z: x[z, y, :, :],
-                torch.randn(16, 32, 16, 12),
-                torch.randint(2, (8, 12)),
-                torch.randint(5, (12, 8, 12)),
-            )
-            self._test_op(
-                mesh,
-                lambda x, y, z: x[z, :, y, :],
-                torch.randn(16, 32, 16, 12),
-                torch.randint(2, (8, 12)),
-                torch.randint(5, (12, 8, 12)),
-            )
-            self._test_op(
-                mesh,
-                lambda x, y, z: x[z, :, :, y],
-                torch.randn(16, 32, 16, 12),
-                torch.randint(2, (8, 1)),
-                torch.randint(5, (12, 8, 12)),
-            )
+            # Commented out to fix distributed CI timeout: each 3-tensor call
+            # generates 40-80 sharding combinations via itertools.product,
+            # causing combinatorial explosion.
+            # # implicit (left-padded) broadcast
+            # self._test_op(
+            #     mesh,
+            #     lambda x, y, z: x[:, z, :, y],
+            #     torch.randn(16, 32, 16, 12),
+            #     torch.randint(5, (12, 8, 12)),
+            #     torch.randint(2, (8, 12)),
+            # )
+            # self._test_op(
+            #     mesh,
+            #     lambda x, y, z: x[z, y, :, :],
+            #     torch.randn(16, 32, 16, 12),
+            #     torch.randint(2, (8, 12)),
+            #     torch.randint(5, (12, 8, 12)),
+            # )
+            # self._test_op(
+            #     mesh,
+            #     lambda x, y, z: x[z, :, y, :],
+            #     torch.randn(16, 32, 16, 12),
+            #     torch.randint(2, (8, 12)),
+            #     torch.randint(5, (12, 8, 12)),
+            # )
+            # self._test_op(
+            #     mesh,
+            #     lambda x, y, z: x[z, :, :, y],
+            #     torch.randn(16, 32, 16, 12),
+            #     torch.randint(2, (8, 1)),
+            #     torch.randint(5, (12, 8, 12)),
+            # )
 
-    @with_comms
     def test_index_put_scalar(self):
         device_mesh = init_device_mesh(self.device_type, (2, self.world_size // 2))
         global_input = torch.randn(2, 4, 8, device=self.device_type)
@@ -746,12 +759,12 @@ def test_index_put_scalar(self):
                 input_dt = distribute_tensor(global_input, device_mesh, [i, j])
                 ref = torch.index_put(global_input, global_index, global_value)
                 output_dt = torch.index_put(input_dt, global_index, value_dt)
-                assert isinstance(output_dt, DTensor)
+                if not isinstance(output_dt, DTensor):
+                    raise AssertionError(f"Expected DTensor, got {type(output_dt)}")
                 # for value is a scalar case, output placement must be replicate
                 self.assertEqual(output_dt.placements, (Replicate(), Replicate()))
                 self.assertEqual(output_dt.full_tensor(), ref)
 
-    @with_comms
     def test_index_put_tensor(self):
         device_mesh = init_device_mesh(self.device_type, (2, self.world_size // 2))
         global_input = torch.randn(2, 4, 8, device=self.device_type)
@@ -764,13 +777,348 @@ def test_index_put_tensor(self):
         for accumulate in [True, False]:
             ref = torch.index_put(global_input, global_index, global_value, accumulate)
             output_dt = torch.index_put(input_dt, global_index, value_dt, accumulate)
-            assert isinstance(output_dt, DTensor)
-            # `input_dt` follows `value_dt`'s Shard(1) plus a offset value of
-            # global_value.ndim-global_input.ndim, which results in Shard(2)
-            self.assertEqual(output_dt.placements, (Shard(2), Replicate()))
+            if not isinstance(output_dt, DTensor):
+                raise AssertionError(f"Expected DTensor, got {type(output_dt)}")
+            # Output should be sharded on non-indexed dims (dim 1 or 2).
+            # The exact placement depends on cost-based strategy selection.
+            for p in output_dt.placements:
+                if isinstance(p, Shard):
+                    self.assertIn(p.dim, [1, 2])
+            self.assertEqual(output_dt.full_tensor(), ref)
+
+    def test_index_put_requires_replicated_index(self):
+        """Test that index_put correctly replicates sharded indices."""
+        device_mesh = init_device_mesh(self.device_type, (self.world_size,))
+        global_input = torch.randn(4, 8, device=self.device_type)
+        global_value = torch.zeros([8], device=self.device_type)
+
+        # Create sharded index - should be redistributed to replicated
+        idx_tensor = torch.tensor([0, 1], device=self.device_type)
+        idx_dt = distribute_tensor(idx_tensor, device_mesh, [Shard(0)])
+
+        input_dt = distribute_tensor(global_input, device_mesh, [Replicate()])
+        value_dt = distribute_tensor(global_value, device_mesh, [Replicate()])
+
+        # The op should work - index gets redistributed to replicated internally
+        ref = torch.index_put(global_input, [idx_tensor], global_value)
+        output_dt = torch.index_put(input_dt, [idx_dt], value_dt)
+        self.assertEqual(output_dt.full_tensor(), ref)
+
+    def test_index_put_no_shard_on_indexed_dim(self):
+        """Test that index_put output is not sharded on indexed dims."""
+        device_mesh = init_device_mesh(self.device_type, (self.world_size,))
+        global_input = torch.randn(4, 8, device=self.device_type)
+        global_value = torch.zeros([8], device=self.device_type)
+        global_index = [torch.tensor(0, device=self.device_type)]
+
+        # Shard input on indexed dim 0
+        input_dt = distribute_tensor(global_input, device_mesh, [Shard(0)])
+        value_dt = distribute_tensor(global_value, device_mesh, [Replicate()])
+
+        ref = torch.index_put(global_input, global_index, global_value)
+        output_dt = torch.index_put(input_dt, global_index, value_dt)
+
+        # Verify output is not sharded on the indexed dim (dim 0)
+        for p in output_dt.placements:
+            if isinstance(p, Shard):
+                self.assertNotEqual(
+                    p.dim, 0, "Output should not be sharded on indexed dim"
+                )
+        self.assertEqual(output_dt.full_tensor(), ref)
+
+    def test_index_put_partial_numerics(self):
+        """Test index_put with Partial placements produces correct numerics."""
+        device_mesh = init_device_mesh(self.device_type, (self.world_size,))
+        global_input = torch.randn(4, 8, device=self.device_type)
+        global_index = [torch.tensor(1, device=self.device_type)]
+        global_value = torch.randn(8, device=self.device_type)
+
+        # Create Partial tensors - each rank has partial sum
+        local_input = global_input / self.world_size
+        local_value = global_value / self.world_size
+        input_dt = DTensor.from_local(
+            local_input, device_mesh, [Partial()], run_check=False
+        )
+        value_dt = DTensor.from_local(
+            local_value, device_mesh, [Partial()], run_check=False
+        )
+
+        for accumulate in [True, False]:
+            ref = torch.index_put(global_input, global_index, global_value, accumulate)
+            output_dt = torch.index_put(input_dt, global_index, value_dt, accumulate)
+            self.assertIsInstance(output_dt, DTensor)
+            self.assertEqual(output_dt.placements, (Partial(),))
+            self.assertEqual(output_dt.full_tensor(), ref)
+
+    def test_index_put_duplicated_indices(self):
+        """Test index_put with duplicated indices for both accumulate modes."""
+        device_mesh = init_device_mesh(self.device_type, (self.world_size,))
+        global_input = torch.zeros(4, 8, device=self.device_type)
+        # Duplicated indices - index 1 appears twice
+        idx_tensor = torch.tensor([1, 1, 2], device=self.device_type)
+        global_index = [idx_tensor]
+        global_value = torch.ones(3, 8, device=self.device_type)
+
+        # Shard on non-indexed dim, index must be replicated
+        input_dt = distribute_tensor(global_input, device_mesh, [Shard(1)])
+        idx_dt = distribute_tensor(idx_tensor, device_mesh, [Replicate()])
+        value_dt = distribute_tensor(global_value, device_mesh, [Shard(1)])
+
+        # accumulate=False: last write wins (non-deterministic with duplicates, but
+        # since values are the same, result is deterministic)
+        ref_no_accum = torch.index_put(
+            global_input, global_index, global_value, accumulate=False
+        )
+        output_no_accum = torch.index_put(
+            input_dt, [idx_dt], value_dt, accumulate=False
+        )
+        self.assertEqual(output_no_accum.full_tensor(), ref_no_accum)
+
+        # accumulate=True: values are added (index 1 gets 2x the value)
+        ref_accum = torch.index_put(
+            global_input, global_index, global_value, accumulate=True
+        )
+        output_accum = torch.index_put(input_dt, [idx_dt], value_dt, accumulate=True)
+        self.assertEqual(output_accum.full_tensor(), ref_accum)
+
+    def test_index_put_broadcast_values(self):
+        """Test index_put where values has size-1 broadcast dims."""
+        device_mesh = init_device_mesh(self.device_type, (self.world_size,))
+        # self shape (4, 8), index on dim 0 -> non-indexed dim is 1
+        global_input = torch.randn(4, 8, device=self.device_type)
+        idx = torch.tensor([0, 1], device=self.device_type)
+        global_index = [idx]
+        # values shape (2, 1) — broadcast on the non-indexed dim
+        global_value = torch.randn(2, 1, device=self.device_type)
+
+        # Shard self on non-indexed dim 1, values should be replicated (size 1)
+        input_dt = distribute_tensor(global_input, device_mesh, [Shard(1)])
+        idx_dt = distribute_tensor(idx, device_mesh, [Replicate()])
+        value_dt = distribute_tensor(global_value, device_mesh, [Replicate()])
+
+        ref = torch.index_put(global_input, global_index, global_value)
+        output_dt = torch.index_put(input_dt, [idx_dt], value_dt)
+        self.assertEqual(output_dt.full_tensor(), ref)
+
+    def _test_index_op_partial_propagation(self, op, reduce_op):
+        """Helper: verify Partial placement propagates correctly through an
+        index op that decomposes to index_put."""
+        device_mesh = init_device_mesh(self.device_type, (self.world_size,))
+        global_input = torch.randn(4, 8, device=self.device_type)
+        idx = torch.tensor([1, 3], device=self.device_type)
+        global_source = torch.randn(4, 2, device=self.device_type)
+
+        # Create Partial tensors — each rank holds partial contribution
+        if reduce_op == "sum":
+            local_input = global_input / self.world_size
+            local_source = global_source / self.world_size
+        else:
+            # avg/max/min: each rank holds the full value
+            local_input = global_input.clone()
+            local_source = global_source.clone()
+
+        input_dt = DTensor.from_local(
+            local_input, device_mesh, [Partial(reduce_op)], run_check=False
+        )
+        source_dt = DTensor.from_local(
+            local_source, device_mesh, [Partial(reduce_op)], run_check=False
+        )
+
+        idx_dt = distribute_tensor(idx, device_mesh, [Replicate()])
+
+        ref = op(global_input, 1, idx, global_source)
+        output_dt = op(input_dt, 1, idx_dt, source_dt)
+        self.assertIsInstance(output_dt, DTensor)
+        self.assertEqual(output_dt.full_tensor(), ref)
+
+    @parametrize("reduce_op", ["sum", "avg", "max", "min"])
+    def test_index_copy_partial_propagation(self, reduce_op):
+        self._test_index_op_partial_propagation(torch.index_copy, reduce_op)
+
+    @parametrize("reduce_op", ["sum", "avg"])
+    def test_index_add_partial_propagation(self, reduce_op):
+        self._test_index_op_partial_propagation(torch.index_add, reduce_op)
+
+    @parametrize(
+        "op",
+        [torch.ops.aten.index_put, torch.index_copy, torch.index_add],
+    )
+    def test_index_ops_leading_none_dim_mapping(self, op):
+        """Regression: index_put dim-mapping bug when indexed dim is not dim 0.
+
+        When indices = (None, index), indexed_dims=[1] is contiguous so
+        broadcast replaces dim 1 in-place:
+            values shape = (self.size(0), len(index))  i.e. (non_indexed, broadcast)
+
+        The buggy strategy assumed broadcast at front, so it wrongly allowed
+        S(0) on self + S(1) on values (sharding the broadcast dim).
+        We force that exact placement — if the strategy accepts it without
+        redistributing, the numerics will be wrong.
+
+        index_copy and index_add decompose to index_put with leading None,
+        so they hit the same bug.
+        """
+        device_mesh = init_device_mesh(self.device_type, (self.world_size,))
+        n = self.world_size
+        global_input = torch.randn(n, 2 * n, device=self.device_type)
+        idx = torch.arange(n, device=self.device_type)
+        global_source = torch.randn(n, n, device=self.device_type)
+
+        input_dt = distribute_tensor(global_input, device_mesh, [Shard(0)])
+        idx_dt = distribute_tensor(idx, device_mesh, [Replicate()])
+        source_dt = distribute_tensor(global_source, device_mesh, [Shard(1)])
+
+        if op is torch.ops.aten.index_put:
+            ref = op(global_input, [None, idx], global_source)
+            output_dt = op(input_dt, [None, idx_dt], source_dt)
+        else:
+            ref = op(global_input, 1, idx, global_source)
+            output_dt = op(input_dt, 1, idx_dt, source_dt)
+
+        self.assertIsInstance(output_dt, DTensor)
+        self.assertEqual(output_dt.full_tensor(), ref)
+
+    def test_index_fill(self):
+        """Test index_fill with sharded input."""
+        device_mesh = init_device_mesh(self.device_type, (self.world_size,))
+        global_input = torch.randn(4, 8, device=self.device_type)
+        idx = torch.tensor([1, 3], device=self.device_type)
+        fill_value = 5.0
+
+        ref = global_input.index_fill(1, idx, fill_value)
+
+        input_dt = distribute_tensor(global_input, device_mesh, [Shard(0)])
+        idx_dt = distribute_tensor(idx, device_mesh, [Replicate()])
+
+        output_dt = input_dt.index_fill(1, idx_dt, fill_value)
+
+        self.assertIsInstance(output_dt, DTensor)
+        self.assertEqual(output_dt.full_tensor(), ref)
+
+    @parametrize("reduce_op", ["mean", "amax", "amin", "prod"])
+    def test_index_reduce(self, reduce_op):
+        """Test index_reduce with sharded input."""
+        device_mesh = init_device_mesh(self.device_type, (self.world_size,))
+        global_input = torch.randn(4, 8, device=self.device_type)
+        idx = torch.tensor([1, 3], device=self.device_type)
+        global_source = torch.randn(4, 2, device=self.device_type)
+
+        ref = global_input.index_reduce(1, idx, global_source, reduce_op)
+
+        input_dt = distribute_tensor(global_input, device_mesh, [Shard(0)])
+        idx_dt = distribute_tensor(idx, device_mesh, [Replicate()])
+        source_dt = distribute_tensor(global_source, device_mesh, [Shard(0)])
+
+        output_dt = input_dt.index_reduce(1, idx_dt, source_dt, reduce_op)
+
+        self.assertIsInstance(output_dt, DTensor)
+        self.assertEqual(output_dt.full_tensor(), ref)
+
+    def test_index_put_none_indices(self):
+        """Test index_put with None entries in the indices list.
+
+        None indices (meaning "select all along this dim") are only
+        valid at the ATen level, not through Python's index_put.
+        """
+        device_mesh = init_device_mesh(self.device_type, (self.world_size,))
+        # 3D tensor indexed on dim 0, with trailing None
+        global_input = torch.randn(4, 8, 16, device=self.device_type)
+        idx = torch.tensor([1, 2], device=self.device_type)
+        # indices = [idx, None, None] -> indexed on dim 0, dims 1,2 not indexed
+        # values shape = (2, 8, 16) matching (*broadcast_shape, *non_indexed_dims)
+        global_value = torch.randn(2, 8, 16, device=self.device_type)
+
+        ref = torch.ops.aten.index_put.default(
+            global_input, [idx, None, None], global_value
+        )
+
+        idx_dt = distribute_tensor(idx, device_mesh, [Replicate()])
+        value_dt = distribute_tensor(global_value, device_mesh, [Replicate()])
+
+        # Shard self on non-indexed dims (1 or 2)
+        for shard_dim in [1, 2]:
+            input_dt = distribute_tensor(global_input, device_mesh, [Shard(shard_dim)])
+            output_dt = torch.ops.aten.index_put.default(
+                input_dt, [idx_dt, None, None], value_dt
+            )
+            self.assertEqual(output_dt.full_tensor(), ref)
+
+    def test_index_put_inplace_none_indices(self):
+        """Test in-place aten.index_put_ with None entries in the indices list.
+
+        The in-place variant with None indices is generated by autograd
+        (e.g. backward of aten.index.Tensor), not by Python-level
+        index_put_. We call the ATen op directly.
+        """
+        device_mesh = init_device_mesh(self.device_type, (self.world_size,))
+        global_input = torch.randn(4, 8, 16, device=self.device_type)
+        idx = torch.tensor([1, 2], device=self.device_type)
+        global_value = torch.randn(2, 8, 16, device=self.device_type)
+
+        ref = torch.ops.aten.index_put_.default(
+            global_input.clone(), [idx, None, None], global_value
+        )
+
+        idx_dt = distribute_tensor(idx, device_mesh, [Replicate()])
+        value_dt = distribute_tensor(global_value, device_mesh, [Replicate()])
+        input_dt = distribute_tensor(global_input.clone(), device_mesh, [Shard(1)])
+        torch.ops.aten.index_put_.default(input_dt, [idx_dt, None, None], value_dt)
+        self.assertEqual(input_dt.full_tensor(), ref)
+
+    def test_index_put_noncontiguous_indexed_dims(self):
+        """Test index_put with non-contiguous indexed dims (None gaps).
+
+        For non-contiguous indexed dims like [idx, None, idx2], advanced
+        indexing moves the broadcast shape to the front of the output:
+        output = (*broadcast_shape, *non_indexed_dim_sizes).
+        """
+        device_mesh = init_device_mesh(self.device_type, (self.world_size,))
+        # self[4, 8, 3, 16], index dims 0 and 2 (non-contiguous)
+        global_input = torch.randn(4, 8, 3, 16, device=self.device_type)
+        idx0 = torch.tensor([0, 1], device=self.device_type)
+        idx2 = torch.tensor([2, 0], device=self.device_type)
+        # output shape = (2, 8, 16), values must be broadcastable to that
+        global_value = torch.randn(2, 8, 16, device=self.device_type)
+
+        ref = torch.ops.aten.index_put.default(
+            global_input, [idx0, None, idx2], global_value
+        )
+
+        idx0_dt = distribute_tensor(idx0, device_mesh, [Replicate()])
+        idx2_dt = distribute_tensor(idx2, device_mesh, [Replicate()])
+        value_dt = distribute_tensor(global_value, device_mesh, [Replicate()])
+
+        # Shard self on non-indexed dim 1 (size 8)
+        input_dt = distribute_tensor(global_input, device_mesh, [Shard(1)])
+        output_dt = torch.ops.aten.index_put.default(
+            input_dt, [idx0_dt, None, idx2_dt], value_dt
+        )
+        self.assertEqual(output_dt.full_tensor(), ref)
+
+    def test_index_put_mid_dim_indexed(self):
+        """Test index_put when indexed dims are in the middle of the tensor."""
+        device_mesh = init_device_mesh(self.device_type, (self.world_size,))
+        # self[16, 8, 4, 16]: index dim 2, non-indexed dims = [0, 1, 3]
+        global_input = torch.randn(16, 8, 4, 16, device=self.device_type)
+        idx = torch.tensor([1], device=self.device_type)
+        # values shape = [16, 8, 1, 16] (broadcast replaces dim 2 in-place)
+        global_value = torch.randn(16, 8, 1, 16, device=self.device_type)
+
+        ref = torch.ops.aten.index_put.default(
+            global_input, [None, None, idx], global_value
+        )
+
+        idx_dt = distribute_tensor(idx, device_mesh, [Replicate()])
+        value_dt = distribute_tensor(global_value, device_mesh, [Replicate()])
+
+        # Shard self on each non-indexed dim
+        for shard_dim in [0, 1, 3]:
+            input_dt = distribute_tensor(global_input, device_mesh, [Shard(shard_dim)])
+            output_dt = torch.ops.aten.index_put.default(
+                input_dt, [None, None, idx_dt], value_dt
+            )
             self.assertEqual(output_dt.full_tensor(), ref)
 
-    @with_comms
     def test_where_type_promotion(self):
         mesh = self.build_device_mesh()  # 1D mesh
 
@@ -782,7 +1130,6 @@ def test_where_type_promotion(self):
             ref = torch.where(global_tensor > 0, 1, 0)
             self.assertEqual(res.full_tensor(), ref)
 
-    @with_comms
     def test_dtensor_dtype_conversion(self):
         from torch.distributed.tensor.debug import (
             _clear_sharding_prop_cache,
@@ -820,7 +1167,6 @@ def test_dtensor_dtype_conversion(self):
         self.assertEqual(hits, 1)
         self.assertEqual(misses, 2)
 
-    @with_comms
     def test_single_dim_strategy_dtype_cache_key(self):
         """Test that schema_info from single-dim strategy affects cache key.
 
@@ -842,7 +1188,8 @@ def test_single_dim_strategy_dtype_cache_key(self):
         def to_copy_single_dim_strategy(op, args_schema, kwargs_schema):
             call_count[0] += 1
             self_meta = args_schema[0]
-            assert isinstance(self_meta, TensorMeta)
+            if not isinstance(self_meta, TensorMeta):
+                raise AssertionError(f"Expected TensorMeta, got {type(self_meta)}")
             single_dim_strategies = []
             for dim in range(len(self_meta.shape)):
                 single_dim_strategies.append(
@@ -878,7 +1225,6 @@ def to_copy_single_dim_strategy(op, args_schema, kwargs_schema):
             # With dtype in cache key, strategy should be called twice (different dtypes)
             self.assertEqual(call_count[0], 2)
 
-    @with_comms
     def test_slice(self):
         mesh = self.build_device_mesh()  # 1D mesh
         comm_mode = CommDebugMode()
@@ -905,7 +1251,6 @@ def test_slice(self):
         self.assertEqual(sharded_out.full_tensor(), global_out)
         self.assertEqual(sharded_dtensor.grad.full_tensor(), global_tensor.grad)
 
-    @with_comms
     def test_slice_full_size_on_sharded_dim(self):
         """
         Test for the issue #170427 where slicing with a size that equals or
@@ -938,7 +1283,6 @@ def test_slice_full_size_on_sharded_dim(self):
         result5 = sharded_dtensor[:2, :]  # full slice on dim 0
         self.assertEqual(result5.full_tensor(), global_tensor[:2, :])
 
-    @with_comms
     def test_split_on_partial(self):
         self.run_subtests(
             {
@@ -966,7 +1310,6 @@ def _test_split_on_partial(self, reduce_op: str, split_size: int, split_dim: int
             dim=split_dim,
         )
 
-    @with_comms
     def test_unbind(self):
         device_mesh = self.build_device_mesh()
         shard_dims = [0, 1]
@@ -996,12 +1339,50 @@ def test_unbind(self):
                 ):
                     self.assertEqual(x.full_tensor(), y)
 
+    @with_comms
+    def test_select_scatter(self):
+        device_mesh = self.build_device_mesh()
+        inp = torch.randn(8, 4, 6, device=self.device_type)
+        # select_scatter(self, src, dim, index): insert src at self[:, :, index]
+        for dim, idx in [(0, 2), (1, 1), (2, 3)]:
+            src_shape = list(inp.shape)
+            src_shape.pop(dim)
+            src = torch.randn(src_shape, device=self.device_type)
+            expected = inp.select_scatter(src, dim, idx)
+
+            for shard_dim in range(inp.ndim):
+                if shard_dim == dim:
+                    continue
+                dt_inp = distribute_tensor(inp, device_mesh, [Shard(shard_dim)])
+                src_shard_dim = shard_dim if shard_dim < dim else shard_dim - 1
+                dt_src = distribute_tensor(src, device_mesh, [Shard(src_shard_dim)])
+                result = dt_inp.select_scatter(dt_src, dim, idx)
+                self.assertEqual(result.full_tensor(), expected)
+                self.assertTrue(result.placements[0].is_shard(shard_dim))
+
+    @with_comms
+    def test_diagonal_scatter(self):
+        device_mesh = self.build_device_mesh()
+        inp = torch.randn(8, 4, 6, device=self.device_type)
+
+        # diagonal_scatter(self, src, offset, dim1, dim2)
+        src = torch.randn(8, 4, device=self.device_type)
+        expected = inp.diagonal_scatter(src, 0, 1, 2)
+
+        # shard on dim 0 (batch), which is not one of the diagonal dims
+        dt_inp = distribute_tensor(inp, device_mesh, [Shard(0)])
+        dt_src = distribute_tensor(src, device_mesh, [Shard(0)])
+        result = dt_inp.diagonal_scatter(dt_src, 0, 1, 2)
+        self.assertEqual(result.full_tensor(), expected)
+        self.assertTrue(result.placements[0].is_shard(0))
+
 
 class DistBucketizeTest(LocalDTensorTestBase):
     @with_comms
     def test_bucketize_partial_input(self):
-        # Bucketize returns indices, which cannot be combined with sum/avg
-        # reductions. Partial inputs should be converted to Replicate.
+        # Bucketize is non-linear, so Partial("sum")/Partial("avg") inputs
+        # must be converted to Replicate. But bucketize is monotone, so
+        # Partial("max") and Partial("min") can propagate directly.
         with LocalTensorMode(ranks=self.world_size):
             mesh = self.build_device_mesh()
             boundaries = torch.tensor([1.0, 3.0, 5.0, 7.0], device=self.device_type)
@@ -1010,26 +1391,44 @@ def test_bucketize_partial_input(self):
                 device=self.device_type,
             )
 
-            for reduce_op in ("sum", "avg", "max"):
+            # Non-linear reductions: must redistribute to Replicate
+            for reduce_op in ("sum", "avg"):
                 partial_input = DTensor.from_local(
                     input_tensor, mesh, [Partial(reduce_op)]
                 )
                 dist_boundaries = distribute_tensor(boundaries, mesh, [Replicate()])
                 result = torch.bucketize(partial_input, dist_boundaries)
 
-                # Output must be Replicate, not Partial
                 self.assertTrue(
                     result.placements[0].is_replicate(),
                     f"Expected Replicate output but got {result.placements[0]} "
                     f"for Partial({reduce_op}) input",
                 )
-                # Expected is bucketize on the materialized global tensor, which
-                # differs per reduce_op (e.g. Partial("sum") allreduce-sums the
-                # local tensors across ranks).
                 global_input = partial_input.full_tensor()
                 expected = torch.bucketize(global_input, boundaries)
                 self.assertEqual(result.to_local(), expected)
 
+            # Monotone reductions: output inherits the same partial type
+            for reduce_op in ("max", "min"):
+                partial_input = DTensor.from_local(
+                    input_tensor, mesh, [Partial(reduce_op)]
+                )
+                dist_boundaries = distribute_tensor(boundaries, mesh, [Replicate()])
+                result = torch.bucketize(partial_input, dist_boundaries)
+
+                self.assertTrue(
+                    result.placements[0].is_partial(),
+                    f"Expected Partial output but got {result.placements[0]} "
+                    f"for Partial({reduce_op}) input",
+                )
+                self.assertEqual(
+                    result.placements[0].reduce_op,
+                    reduce_op,
+                    f"Expected Partial({reduce_op}) output but got {result.placements[0]}",
+                )
+                expected = torch.bucketize(input_tensor, boundaries)
+                self.assertEqual(result.full_tensor(), expected)
+
     @with_comms
     def test_bucketize_sharded_input(self):
         # Sharded inputs should propagate sharding to output normally.
@@ -1047,8 +1446,55 @@ def test_bucketize_sharded_input(self):
                 self.assertTrue(result.placements[0].is_shard(shard_dim))
                 self.assertEqual(result.full_tensor(), expected)
 
+    @with_comms
+    def test_bucketize_sharded_boundaries(self):
+        # When boundaries are sharded on dim 0, each rank counts how many of
+        # its local boundary values each input exceeds. The sum across ranks
+        # (Partial("sum")) gives the correct global bucket index.
+        with LocalTensorMode(ranks=self.world_size):
+            mesh = self.build_device_mesh()
+            boundaries = torch.tensor([1.0, 3.0, 5.0, 7.0], device=self.device_type)
+            input_tensor = torch.tensor(
+                [[2.0, 4.0, 6.0, 8.0], [0.0, 1.0, 5.0, 9.0]],
+                device=self.device_type,
+            )
+            expected = torch.bucketize(input_tensor, boundaries)
+
+            dist_input = distribute_tensor(input_tensor, mesh, [Replicate()])
+            dist_boundaries = distribute_tensor(boundaries, mesh, [Shard(0)])
+            result = torch.bucketize(dist_input, dist_boundaries)
+            self.assertEqual(result.full_tensor(), expected)
+
+
+class DistToCopyTest(LocalDTensorTestBase):
+    @with_comms
+    def test_to_copy_partial_reduces_for_nonlinear_cast(self):
+        # (reduce_op, target_dtype, expect_partial)
+        cases = [
+            ("sum", torch.int32, False),  # truncation breaks additivity
+            ("sum", torch.bool, False),  # thresholding
+            ("sum", torch.float64, True),  # float→float is safe
+            ("max", torch.int32, True),  # monotonic
+            ("max", torch.bool, False),  # thresholding
+        ]
+        with LocalTensorMode(ranks=self.world_size):
+            mesh = self.build_device_mesh()
+            input_tensor = torch.randn(4, 4, device=self.device_type)
+            for reduce_op, target_dtype, expect_partial in cases:
+                dt = DTensor.from_local(input_tensor, mesh, [Partial(reduce_op)])
+                result = dt.to(target_dtype)
+                p = result.placements[0]
+                if expect_partial:
+                    self.assertTrue(p.is_partial(), f"{reduce_op}→{target_dtype}: {p}")
+                    self.assertEqual(p.reduce_op, reduce_op)
+                else:
+                    self.assertTrue(
+                        p.is_replicate(), f"{reduce_op}→{target_dtype}: {p}"
+                    )
+
 
-class DistArgMaxArgMinTest(DTensorTestBase):
+class DistArgMaxArgMinTest(DTensorContinuousTestBase):
+    world_size = 4
     _ops = [torch.argmax, torch.argmin]
     sample = [
         [0, 2, 1, 11, 5, 9, -2, -23],
@@ -1069,7 +1515,6 @@ class DistArgMaxArgMinTest(DTensorTestBase):
     )
 
     @skip_if_lt_x_gpu(4)
-    @with_comms
     def test_argmax_argmin_with_placements(self):
         device_mesh = self.build_device_mesh()
         local_tensor = torch.tensor(self.sample, device=self.device_type)
@@ -1081,13 +1526,251 @@ def test_argmax_argmin_with_placements(self):
                 local_result = op(local_tensor, dim=1)
                 self.assertEqual(full_dresult, local_result)
 
+    def test_argmax_argmin_sharded_reduction_dim(self):
+        """Unlike max/min which use reduction_linear=True and produce
+        Partial("max")/Partial("min") outputs, argmax/argmin return indices
+        that can't be combined across shards with an element-wise max/min.
+        The strategy sets reduction_linear=False, which forces the input to
+        be redistributed to Replicate on the sharded reduction dim before
+        the op runs. No Partial placement appears in the output."""
+        mesh = init_device_mesh(self.device_type, (self.world_size,))
+        tensor = torch.tensor(self.sample, device=self.device_type, dtype=torch.float)
+        dtensor = distribute_tensor(tensor, mesh, [Shard(0)])
+
+        for op in self._ops:
+            self.assertEqual(op(dtensor, dim=0).full_tensor(), op(tensor, dim=0))
+            self.assertEqual(op(dtensor).full_tensor(), op(tensor))
+
     def build_device_mesh(self):
         return init_device_mesh(self.device_type, (2, 2))
 
 
+DistArgMaxArgMinTestWithLocalTensor = create_local_tensor_test_class(
+    DistArgMaxArgMinTest,
+    base_class=LocalDTensorContinuousTestBase,
+)
+
+instantiate_parametrized_tests(DistTensorOpsTest)
 DistTensorOpsTestWithLocalTensor = create_local_tensor_test_class(
     DistTensorOpsTest,
+    base_class=LocalDTensorContinuousTestBase,
 )
 
+
+@torch.library.custom_op("testlib::optional_clamp_op", mutates_args=())
+def optional_clamp_op(
+    x: torch.Tensor, lo: torch.Tensor | None, hi: torch.Tensor | None
+) -> torch.Tensor:
+    return torch.clamp(x, lo, hi)
+
+
+@optional_clamp_op.register_fake
+def _optional_clamp_op_fake(x, lo, hi):
+    return torch.empty_like(x)
+
+
+@torch.library.custom_op("testlib::modified_cat_op", mutates_args=())
+def modified_cat_op(a: list[torch.Tensor], b: list[torch.Tensor]) -> torch.Tensor:
+    return torch.cat(a, dim=0)
+
+
+@modified_cat_op.register_fake
+def _modified_cat_op_fake(a, b):
+    total_rows = sum(t.shape[0] for t in a)
+    return torch.empty(
+        total_rows, *a[0].shape[1:], dtype=a[0].dtype, device=a[0].device
+    )
+
+
+class DistTensorCppPyTree(DTensorContinuousTestBase):
+    def test_cpp_cat(self):
+        from torch.distributed.tensor.debug import (
+            _clear_fast_path_sharding_prop_cache,
+            _get_fast_path_sharding_prop_cache_stats,
+        )
+
+        mesh = self.build_device_mesh()
+        cases = [
+            ((8, 8), (8, 8), 0, [Shard(1)], (Shard(1),)),
+            ((4, 8), (6, 8), 0, [Replicate()], (Replicate(),)),
+            ((8, 8), (8, 8), 0, [Shard(0)], (Replicate(),)),
+        ]
+        for shape_a, shape_b, cat_dim, placements, expected_placements in cases:
+            _clear_fast_path_sharding_prop_cache()
+            global_a = torch.randn(shape_a)
+            global_b = torch.randn(shape_b)
+            dt_a = distribute_tensor(global_a, mesh, placements)
+            dt_b = distribute_tensor(global_b, mesh, placements)
+            # first call: cache miss
+            result = torch.cat([dt_a, dt_b], dim=cat_dim)
+            self.assertEqual(result.placements, expected_placements)
+            self.assertEqual(
+                result.full_tensor(), torch.cat([global_a, global_b], dim=cat_dim)
+            )
+            # second call: should be a cache hit
+            result = torch.cat([dt_a, dt_b], dim=cat_dim)
+            hits, _ = _get_fast_path_sharding_prop_cache_stats()
+            self.assertEqual(hits, 1)
+            self.assertEqual(result.placements, expected_placements)
+            self.assertEqual(
+                result.full_tensor(), torch.cat([global_a, global_b], dim=cat_dim)
+            )
+
+    def test_foreach_mixed_dtensor_and_tensor(self):
+        from torch.distributed.tensor.debug import (
+            _clear_fast_path_sharding_prop_cache,
+            _get_fast_path_sharding_prop_cache_stats,
+        )
+        from torch.distributed.tensor.experimental import implicit_replication
+
+        mesh = self.build_device_mesh()
+        with implicit_replication():
+            _clear_fast_path_sharding_prop_cache()
+            dt = distribute_tensor(
+                torch.zeros(4, 4, device=self.device_type), mesh, [Shard(0)]
+            )
+            regular = torch.zeros(4, 4, device=self.device_type)
+
+            torch._foreach_add_([dt, regular], 1)
+            hits, misses = _get_fast_path_sharding_prop_cache_stats()
+            self.assertEqual(misses, 1)
+            self.assertEqual(
+                dt.full_tensor(), torch.ones(4, 4, device=self.device_type)
+            )
+            self.assertEqual(regular, torch.ones(4, 4, device=self.device_type))
+
+            torch._foreach_add_([dt, regular], 1)
+            hits, misses = _get_fast_path_sharding_prop_cache_stats()
+            self.assertEqual(hits, 1)
+
+    def test_two_list_op_cache_collision(self):
+        from torch.distributed.tensor._op_schema import RuntimeSchemaInfo
+        from torch.distributed.tensor._ops.utils import replicate_op_strategy
+        from torch.distributed.tensor.debug import (
+            _clear_fast_path_sharding_prop_cache,
+            _get_fast_path_sharding_prop_cache_stats,
+        )
+
+        mesh = self.build_device_mesh()
+        op = torch.ops.testlib.modified_cat_op
+
+        with op_strategy_context(
+            op.default,
+            replicate_op_strategy,
+            schema_info=RuntimeSchemaInfo(needs_pytree=True),
+        ):
+            _clear_fast_path_sharding_prop_cache()
+            a = distribute_tensor(torch.randn(4, 8), mesh, [Replicate()])
+            b = distribute_tensor(torch.randn(4, 8), mesh, [Replicate()])
+
+            op([a], [b])  # cache miss, populates cache
+            result = op([a, b], [])  # different list sizes, should also miss
+            hits, misses = _get_fast_path_sharding_prop_cache_stats()
+            self.assertEqual(hits, 0)
+            self.assertEqual(misses, 2)
+            self.assertEqual(result.shape, torch.Size([8, 8]))
+
+    def test_optional_tensor_cache_key(self):
+        """Cache must distinguish op(t1, None, t2) from op(t1, t2, None).
+
+        Uses a custom op with high static_argnum so that None args for
+        optional Tensor? parameters would be dropped from the cache key
+        without the is_none_or_undefined fix in handle_non_dtensor_arg.
+        With same-spec DTensors the collision is deterministic.
+        """
+        from torch.distributed.tensor._op_schema import RuntimeSchemaInfo
+        from torch.distributed.tensor._ops.utils import replicate_op_strategy
+        from torch.distributed.tensor.debug import (
+            _clear_fast_path_sharding_prop_cache,
+            _get_fast_path_sharding_prop_cache_stats,
+        )
+
+        mesh = self.build_device_mesh()
+        op = torch.ops.testlib.optional_clamp_op
+
+        with op_strategy_context(
+            op.default,
+            replicate_op_strategy,
+            schema_info=RuntimeSchemaInfo(static_argnum=3),
+        ):
+            _clear_fast_path_sharding_prop_cache()
+            a = distribute_tensor(torch.randn(4, 8), mesh, [Shard(0)])
+
+            op(a, a, None)  # miss 1: key should include None at position 2
+            op(a, None, a)  # miss 2: key should include None at position 1
+
+            hits, misses = _get_fast_path_sharding_prop_cache_stats()
+            self.assertEqual(hits, 0)
+            self.assertEqual(misses, 2)
+
+    def test_optional_tensor_cache_key_python_slow_path(self):
+        """Same as above but exercises the Python slow path via OpSchema directly.
+
+        DTensor_OpSchema_recompute_comparison_key_impl must include None
+        args for optional Tensor? parameters in the comparison key.
+        """
+        from torch.distributed.tensor._dtensor_spec import DTensorSpec
+        from torch.distributed.tensor._op_schema import OpSchema, RuntimeSchemaInfo
+
+        mesh = self.build_device_mesh()
+        spec = DTensorSpec(mesh, (Shard(0),))
+        op = torch.ops.testlib.optional_clamp_op.default
+
+        schema1 = OpSchema(op, (spec, None, spec), {})
+        schema1.schema_info = RuntimeSchemaInfo(static_argnum=3)
+        schema1._recompute_comparison_key()
+
+        schema2 = OpSchema(op, (spec, spec, None), {})
+        schema2.schema_info = RuntimeSchemaInfo(static_argnum=3)
+        schema2._recompute_comparison_key()
+
+        self.assertNotEqual(hash(schema1), hash(schema2))
+        self.assertNotEqual(schema1, schema2)
+
+
+class TestNewEmptyStridedUneven(DTensorTestBase):
+    @with_comms
+    def test_backward_no_allgather(self):
+        """Backward on unevenly-sharded DTensor should not allgather (issue #107661)."""
+        mesh = self.build_device_mesh()
+        placement = (Shard(1),)
+        x = torch.randn(12, self.world_size * 2 + 1, requires_grad=True)
+        dt = distribute_tensor(x, mesh, placement)
+        comm_mode = CommDebugMode()
+        with comm_mode:
+            (dt.to_local() * 2).sum().backward()
+        self.assertEqual(comm_mode.get_total_counts(), 0)
+        self.assertEqual(dt.grad.placements, placement)
+        self.assertTrue(dt.grad._local_tensor.is_contiguous())
+        self.assertEqual(
+            dt.grad._local_tensor,
+            torch.full_like(dt.grad._local_tensor, 2.0),
+        )
+
+    @with_comms
+    def test_backward_channels_last(self):
+        """Backward preserves channels-last stride order for unevenly-sharded DTensor."""
+        mesh = self.build_device_mesh()
+        placement = (Shard(2),)
+        # H=2*world_size+1 ensures uneven sharding on dim 2
+        x = torch.randn(2, 3, self.world_size * 2 + 1, 4).to(
+            memory_format=torch.channels_last
+        )
+        x.requires_grad_(True)
+        dt = distribute_tensor(x, mesh, placement)
+        comm_mode = CommDebugMode()
+        with comm_mode:
+            (dt.to_local() * 2).sum().backward()
+        self.assertEqual(comm_mode.get_total_counts(), 0)
+        self.assertEqual(dt.grad.placements, placement)
+        self.assertTrue(
+            dt.grad._local_tensor.is_contiguous(memory_format=torch.channels_last)
+        )
+        self.assertEqual(
+            dt.grad._local_tensor,
+            torch.full_like(dt.grad._local_tensor, 2.0),
+        )
+
+
 if __name__ == "__main__":
     run_tests()
diff --git a/test/distributed/tensor/test_utils.py b/test/distributed/tensor/test_utils.py
index 02f2469e29221..d3fc2441e7c15 100644
--- a/test/distributed/tensor/test_utils.py
+++ b/test/distributed/tensor/test_utils.py
@@ -108,17 +108,20 @@ def test_compute_local_shape_and_global_offset_uneven(self):
                     expected_shard_offset = dp012_shard_size * dp_rank
                     expected_shard_size = 3
                 else:
-                    assert tp_rank == 1
+                    if tp_rank != 1:
+                        raise AssertionError(f"Expected tp_rank == 1, got {tp_rank}")
                     expected_shard_offset = dp012_shard_size * dp_rank + tp0_shard_size
                     expected_shard_size = 2
             else:
-                assert dp_rank == 3
+                if dp_rank != 3:
+                    raise AssertionError(f"Expected dp_rank == 3, got {dp_rank}")
                 tp0_shard_size = 2
                 if tp_rank == 0:
                     expected_shard_offset = dp012_shard_size * dp_rank
                     expected_shard_size = 2
                 else:
-                    assert tp_rank == 1
+                    if tp_rank != 1:
+                        raise AssertionError(f"Expected tp_rank == 1, got {tp_rank}")
                     expected_shard_offset = dp012_shard_size * dp_rank + tp0_shard_size
                     expected_shard_size = 1
             self.assertEqual(local_shape, (expected_shard_size, 2))
@@ -143,17 +146,20 @@ def test_compute_local_shape_and_global_offset_uneven(self):
                     expected_shard_offset = dp012_shard_size * dp_rank
                     expected_shard_size = 2
                 else:
-                    assert tp_rank == 1
+                    if tp_rank != 1:
+                        raise AssertionError(f"Expected tp_rank == 1, got {tp_rank}")
                     expected_shard_offset = dp012_shard_size * dp_rank + tp0_shard_size
                     expected_shard_size = 2
             else:
-                assert dp_rank == 3
+                if dp_rank != 3:
+                    raise AssertionError(f"Expected dp_rank == 3, got {dp_rank}")
                 tp0_shard_size = 1
                 if tp_rank == 0:
                     expected_shard_offset = dp012_shard_size * dp_rank
                     expected_shard_size = 1
                 else:
-                    assert tp_rank == 1
+                    if tp_rank != 1:
+                        raise AssertionError(f"Expected tp_rank == 1, got {tp_rank}")
                     expected_shard_offset = global_shape[0]
                     expected_shard_size = 0
             self.assertEqual(local_shape, (expected_shard_size, 2))
@@ -201,17 +207,20 @@ def test_compute_local_shape_and_global_offset_uneven(self):
                     expected_shard_offset = 4 * dp_rank
                     expected_shard_size = tp0_shard_size
                 else:
-                    assert tp_rank == 1
+                    if tp_rank != 1:
+                        raise AssertionError(f"Expected tp_rank == 1, got {tp_rank}")
                     expected_shard_offset = 4 * dp_rank + 2
                     expected_shard_size = 4
             else:
-                assert dp_rank == 3
+                if dp_rank != 3:
+                    raise AssertionError(f"Expected dp_rank == 3, got {dp_rank}")
                 tp0_shard_size = 3
                 if tp_rank == 0:
                     expected_shard_offset = 4 * dp_rank
                     expected_shard_size = 3
                 else:
-                    assert tp_rank == 1
+                    if tp_rank != 1:
+                        raise AssertionError(f"Expected tp_rank == 1, got {tp_rank}")
                     expected_shard_offset = global_shape[0]
                     expected_shard_size = 0
             self.assertEqual(local_shape, (expected_shard_size, 2))
@@ -234,16 +243,19 @@ def test_compute_local_shape_and_global_offset_uneven(self):
                     expected_shard_offset = 5 * dp_rank
                     expected_shard_size = tp0_shard_size
                 else:
-                    assert tp_rank == 1
+                    if tp_rank != 1:
+                        raise AssertionError(f"Expected tp_rank == 1, got {tp_rank}")
                     expected_shard_offset = 5 * dp_rank + 2
                     expected_shard_size = 2
             else:
-                assert dp_rank == 3
+                if dp_rank != 3:
+                    raise AssertionError(f"Expected dp_rank == 3, got {dp_rank}")
                 if tp_rank == 0:
                     expected_shard_offset = 5 * dp_rank
                     expected_shard_size = 2
                 else:
-                    assert tp_rank == 1
+                    if tp_rank != 1:
+                        raise AssertionError(f"Expected tp_rank == 1, got {tp_rank}")
                     expected_shard_offset = 5 * dp_rank + 1
                     expected_shard_size = 1
             self.assertEqual(local_shape, (expected_shard_size, 2))
@@ -273,7 +285,10 @@ def test_compute_local_shape_and_global_offset_uneven(self):
                         expected_shard_offset = 10 * mesh0_rank + 6
                         expected_shard_size = 2
                 else:
-                    assert mesh1_rank == 1
+                    if mesh1_rank != 1:
+                        raise AssertionError(
+                            f"Expected mesh1_rank == 1, got {mesh1_rank}"
+                        )
                     if mesh2_rank == 0:
                         expected_shard_offset = 10 * mesh0_rank + 3
                         expected_shard_size = 2
@@ -281,21 +296,31 @@ def test_compute_local_shape_and_global_offset_uneven(self):
                         expected_shard_offset = 10 * mesh0_rank + 8
                         expected_shard_size = 2
                     else:
-                        assert mesh2_rank == 2
+                        if mesh2_rank != 2:
+                            raise AssertionError(
+                                f"Expected mesh2_rank == 2, got {mesh2_rank}"
+                            )
                         expected_shard_size = 0
                         expected_shard_offset = global_shape[0]
             else:
-                assert mesh0_rank == 3
+                if mesh0_rank != 3:
+                    raise AssertionError(f"Expected mesh0_rank == 3, got {mesh0_rank}")
                 if mesh1_rank == 0:
                     if mesh2_rank in (0, 1):
                         expected_shard_offset = 10 * mesh0_rank + 2 * mesh2_rank
                         expected_shard_size = 2
                     else:
-                        assert mesh2_rank == 2
+                        if mesh2_rank != 2:
+                            raise AssertionError(
+                                f"Expected mesh2_rank == 2, got {mesh2_rank}"
+                            )
                         expected_shard_offset = 10 * mesh0_rank + 6
                         expected_shard_size = 1
                 else:
-                    assert mesh1_rank == 1
+                    if mesh1_rank != 1:
+                        raise AssertionError(
+                            f"Expected mesh1_rank == 1, got {mesh1_rank}"
+                        )
                     if mesh2_rank == 0:
                         expected_shard_offset = 10 * mesh0_rank + 3
                         expected_shard_size = 2
@@ -568,7 +593,8 @@ def test_fsdp_tp_meta_compute(self):
         local_shape, global_offset = compute_local_shape_and_global_offset(
             global_tensor_shape, global_mesh, placements
         )
-        assert global_mesh.get_coordinate is not None
+        if global_mesh.get_coordinate is None:
+            raise AssertionError("Expected global_mesh.get_coordinate to be not None")
         dp_rank = global_mesh.get_local_rank("dp")
         tp_rank = global_mesh.get_local_rank("tp")
         shard_idx_on_dim_0 = tp_rank * dp_size + dp_rank
@@ -619,7 +645,8 @@ def test_hsdp_tp_meta_compute(self):
         local_shape, global_offset = compute_local_shape_and_global_offset(
             global_tensor_shape, global_mesh, placements
         )
-        assert global_mesh.get_coordinate is not None
+        if global_mesh.get_coordinate is None:
+            raise AssertionError("Expected global_mesh.get_coordinate to be not None")
         dp_shard_rank = global_mesh.get_local_rank("dp_shard")
         tp_rank = global_mesh.get_local_rank("tp")
         shard_idx_on_dim_0 = tp_rank * dp_shard_size + dp_shard_rank
@@ -690,7 +717,10 @@ def test_compute_tensor_info(self):
                 "cp",
             ),
         )
-        assert world_size == mesh.shape[0] * mesh.shape[1] * mesh.shape[2]
+        if world_size != mesh.shape[0] * mesh.shape[1] * mesh.shape[2]:
+            raise AssertionError(
+                f"Expected world_size == {mesh.shape[0] * mesh.shape[1] * mesh.shape[2]}, got {world_size}"
+            )
 
         # Add Partial() when we are allowed to redistribute to it
         options = [Shard(0), Shard(1), Shard(2), Replicate()]
@@ -1041,8 +1071,10 @@ def test_einsum_propagation(self):
             self.assertEqual(
                 comm_mode.get_comm_counts()[c10d_functional.all_gather_into_tensor], 4
             )
-            assert isinstance(res1, DTensor)
-            assert isinstance(res2, DTensor)
+            if not isinstance(res1, DTensor):
+                raise AssertionError(f"Expected res1 to be DTensor, got {type(res1)}")
+            if not isinstance(res2, DTensor):
+                raise AssertionError(f"Expected res2 to be DTensor, got {type(res2)}")
             self.assertEqual(res1.full_tensor(), res2.full_tensor())
 
     @with_comms
@@ -1058,12 +1090,14 @@ def test_pointwise_propagation(self):
             with CommDebugMode() as comm_mode:
                 res1 = torch.sum(A, dim=0)
             self.assertEqual(comm_mode.get_total_counts(), 0)
-            assert isinstance(res1, DTensor)
+            if not isinstance(res1, DTensor):
+                raise AssertionError(f"Expected res1 to be DTensor, got {type(res1)}")
             self.assertEqual(res1.full_tensor(), torch.sum(input_tensor, dim=0))
             with CommDebugMode() as comm_mode:
                 res2 = torch.sum(A, dim=1)
             self.assertEqual(comm_mode.get_total_counts(), 0)
-            assert isinstance(res2, DTensor)
+            if not isinstance(res2, DTensor):
+                raise AssertionError(f"Expected res2 to be DTensor, got {type(res2)}")
             self.assertEqual(res2.full_tensor(), torch.sum(input_tensor, dim=1))
 
     @with_comms
@@ -1116,7 +1150,10 @@ def run_view_propagation(
         view_into_shape: list[int],
         expected_placements_after_view,
     ):
-        assert math.prod(original_full_tensor.shape) == math.prod(view_into_shape)
+        if math.prod(original_full_tensor.shape) != math.prod(view_into_shape):
+            raise AssertionError(
+                f"Expected shapes to have same product, got {math.prod(original_full_tensor.shape)} vs {math.prod(view_into_shape)}"
+            )
         # verify user specified `expected_placements_after_view `is correct
         A = distribute_tensor(original_full_tensor, mesh, original_placements)
         B = original_full_tensor.view(*view_into_shape)
@@ -1128,8 +1165,14 @@ def run_view_propagation(
         )
         # verify the propagated sharding spec after view is correct
         viewed_A = A.view(*view_into_shape)
-        assert isinstance(viewed_A, DTensor)
-        assert viewed_A.placements == tuple(expected_placements_after_view)
+        if not isinstance(viewed_A, DTensor):
+            raise AssertionError(
+                f"Expected viewed_A to be DTensor, got {type(viewed_A)}"
+            )
+        if viewed_A.placements != tuple(expected_placements_after_view):
+            raise AssertionError(
+                f"Expected placements {tuple(expected_placements_after_view)}, got {viewed_A.placements}"
+            )
         self.assertEqual(A.to_local().view(-1), viewed_A.to_local().view(-1))
 
     @with_comms
@@ -1496,7 +1539,8 @@ 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)
+                if not (isinstance(lhs, LocalTensor) and isinstance(rhs, LocalTensor)):
+                    raise AssertionError("Expected both lhs and rhs to be LocalTensor")
                 super().assertEqual(lhs._ranks, rhs._ranks)
                 for r in lhs._ranks:
                     super().assertEqual(
@@ -1598,7 +1642,8 @@ def _get_logical_shape(
         """
         new_logical_shape = list(full_tensor_shape)
         coordinate = mesh.get_coordinate()
-        assert coordinate is not None
+        if coordinate is None:
+            raise AssertionError("Expected mesh.get_coordinate() to be not None")
 
         for tensor_dim, shard_configs in shard_map.items():
             for config in shard_configs:
diff --git a/test/distributed/tensor/test_view_ops.py b/test/distributed/tensor/test_view_ops.py
index aeb3a9df79a31..e709458fdb902 100644
--- a/test/distributed/tensor/test_view_ops.py
+++ b/test/distributed/tensor/test_view_ops.py
@@ -1,7 +1,9 @@
 # Copyright (c) Meta Platforms, Inc. and affiliates
 # Owner(s): ["oncall: distributed"]
 
+import contextlib
 import itertools
+import math
 from typing import cast
 
 import torch
@@ -16,11 +18,14 @@
     Replicate,
     Shard,
 )
+from torch.distributed.tensor._dtensor_spec import DTensorSpec
 from torch.distributed.tensor._ops._view_ops import (
+    _ViewShardingPropagator,
     Broadcast,
     dim_maps,
     Flatten,
     InputDim,
+    propagate_shape_and_sharding,
     Repeat,
     Singleton,
     Split,
@@ -31,16 +36,39 @@
 from torch.testing._internal.common_utils import run_tests
 from torch.testing._internal.distributed._tensor.common_dtensor import (
     create_local_tensor_test_class,
-    DTensorTestBase,
-    with_comms,
+    DTensorContinuousTestBase,
+    LocalDTensorContinuousTestBase,
 )
 from torch.utils import _pytree as pytree
 
 
-class TestViewOps(DTensorTestBase):
-    @property
-    def world_size(self) -> int:
-        return 6
+class TestViewOps(DTensorContinuousTestBase):
+    world_size = 6
+
+    @staticmethod
+    def _get_all_factorizations(n):
+        """Return all ways to factor n into a sorted tuple of integers > 1.
+
+        Each factorization has length >= 2 (non-trivial unflatten).
+        Only sorted (non-decreasing) factorizations are returned since permutations
+        exercise the same split_factor logic with different constants.
+
+        Examples:
+            12 -> [(2, 6), (3, 4), (2, 2, 3)]
+            36 -> [(2, 18), (3, 12), (4, 9), (6, 6), (2, 2, 9), (2, 3, 6), (3, 3, 4), (2, 2, 3, 3)]
+        """
+
+        def get_sorted_factorizations(remaining, min_factor=2):
+            if remaining == 1:
+                return [()]
+            result = []
+            for f in range(min_factor, remaining + 1):
+                if remaining % f == 0:
+                    for sub in get_sorted_factorizations(remaining // f, f):
+                        result.append((f,) + sub)
+            return result
+
+        return [f for f in get_sorted_factorizations(n) if len(f) >= 2]
 
     def test_view_groups(self):
         self.assertEqual(
@@ -197,6 +225,11 @@ def call_dt_test(self, op, args, kwargs, device_mesh: DeviceMesh):
                     device_mesh,
                     in_shard,
                 )
+                # These are literal _StridedShard placements (simulating
+                # flatten output), not shard-order encodings.  Override the
+                # auto-detection which may incorrectly set the flag to True
+                # when split_factor happens to match a mesh dim size.
+                in_dt._spec.use_strided_shard_as_shard_order = False
             else:
                 in_dt = distribute_tensor(args[0], device_mesh, in_shard)
 
@@ -218,7 +251,6 @@ def dimmap_test(self, op, args, expected_rule_output):
         self.assertEqual(rules, expected_rule_output)
         self.call_dt_test(op, args, {}, self.device_mesh)
 
-    @with_comms
     def test_illegal_views(self):
         device_mesh = self.build_device_mesh()
         # 1D mesh [6] (see above)
@@ -229,8 +261,9 @@ def test_illegal_views(self):
         shard.view(-1)
 
         shard = dtensor.redistribute(device_mesh=device_mesh, placements=[Shard(dim=1)])
-        with self.assertRaisesRegex(RuntimeError, "Sharding propagation failed"):
-            shard.view(-1)
+        # Shard(1) on the last flattened dim with uneven sharding is now allowed
+        # (produces _StridedShard output)
+        shard.view(-1)
 
         # 8 is the uneven case since mesh dim is 6
         tensor = torch.randn((8, 256))
@@ -247,7 +280,34 @@ def test_illegal_views(self):
         with self.assertRaisesRegex(RuntimeError, "Sharding propagation failed"):
             shard.view(8, 2, -1)
 
-    @with_comms
+    def test_double_shard_split_validation(self):
+        """[Shard(0), Shard(0)] through Split correctly validates divisibility."""
+        # Compatible: reshape (24,)→(6,4) with [Shard(0), Shard(0)] on mesh (2,3)
+        # submesh_size = 2*3 = 6, split_id=0 out_size=6, 6%6==0 → passes
+        result = propagate_shape_and_sharding(
+            [Shard(0), Shard(0)],
+            (24,),
+            dim_maps[torch.Tensor.view](torch.empty(24), [6, 4]),
+            (2, 3),
+        )
+        self.assertEqual(len(result), 2)
+        input_tgt, output = result
+        self.assertEqual(list(input_tgt), [Shard(0), Shard(0)])
+        self.assertEqual(list(output), [Shard(0), Shard(0)])
+
+        # Incompatible: reshape (12,)→(3,4) with [Shard(0), Shard(0)] on mesh (2,3)
+        # out_size=3 is not divisible by mesh dim 0 size=2 → error
+        with self.assertRaisesRegex(
+            RuntimeError, "not evenly divisible by mesh dimension"
+        ):
+            propagate_shape_and_sharding(
+                [Shard(0), Shard(0)],
+                (12,),
+                dim_maps[torch.Tensor.view](torch.empty(12), [3, 4]),
+                (2, 3),
+                strict_view=True,
+            )
+
     def test_view_ops(self):
         mesh_shape = (dist.get_world_size() // 2, 2)
         self.device_mesh = init_device_mesh(
@@ -544,7 +604,7 @@ def test_view_ops(self):
     #         Split(InputDim(1), (13, 2), 1),
     #     ),
     # )
-    @with_comms
+
     def test_complex_view_ops(self):
         self.device_mesh = DeviceMesh(
             self.device_type, torch.arange(dist.get_world_size()).view(-1, 2)
@@ -589,7 +649,32 @@ def test_complex_view_ops(self):
             )
             self.assertEqual(out, out_dt.full_tensor())
 
-    @with_comms
+    def test_view_as_complex_no_pmax_pmin(self):
+        """
+        view_as_complex converts real tensors to complex. Complex numbers don't
+        have a total ordering, so P(max) and P(min) placements are invalid.
+        This test verifies that these placements are not propagated.
+        """
+        device_mesh = DeviceMesh(self.device_type, torch.arange(self.world_size))
+        inp = torch.randn(8, 2)
+
+        # Create input with P(sum) - this should work
+        inp_dt_psum = DTensor.from_local(
+            inp.clone(), device_mesh, [Partial("sum")], run_check=False
+        )
+        result_psum = torch.view_as_complex(inp_dt_psum)
+        # P(sum) should propagate since sum is valid for complex
+        self.assertIsInstance(result_psum.placements[0], Partial)
+        self.assertEqual(result_psum.placements[0].reduce_op, "sum")
+
+        # Create input with P(max) - should redistribute to Replicate
+        inp_dt_pmax = DTensor.from_local(
+            inp.clone(), device_mesh, [Partial("max")], run_check=False
+        )
+        result_pmax = torch.view_as_complex(inp_dt_pmax)
+        # P(max) should NOT propagate - output should be Replicate
+        self.assertIsInstance(result_pmax.placements[0], Replicate)
+
     def test_dtensor_view_op_uneven(self):
         """
         When the sharded dimension is unchanged, the view op should not trigger any communication.
@@ -632,7 +717,1741 @@ def test_dtensor_view_op_uneven(self):
                 )
                 self.assertEqual(len(comm_mode.get_comm_counts()), 0)
 
-    @with_comms
+    def _get_viewed_tensor_dims(self, tensor_dims, flatten_start, flatten_end):
+        if isinstance(tensor_dims, tuple):
+            tensor_dims = list(tensor_dims)
+        flatten_dims = tensor_dims[flatten_start:flatten_end]
+        if len(flatten_dims) > 0:
+            flatten_dim = math.prod(flatten_dims)
+        else:
+            flatten_dim = None
+        leading_dims = tensor_dims[:flatten_start]
+        trailing_dims = tensor_dims[flatten_end:]
+        view_shapes = []
+        if len(leading_dims) > 0:
+            view_shapes.extend(leading_dims)
+        if flatten_dim is not None:
+            view_shapes.append(flatten_dim)
+        if len(trailing_dims) > 0:
+            view_shapes.extend(trailing_dims)
+        return tuple(view_shapes)
+
+    def _test_dtensor_flatten_split_case(self, in_shape, out_shape, placements, mesh):
+        """Test a single Split(Flatten) view case.
+
+        Representable cases must produce zero communication and correct output.
+        Unrepresentable cases must raise RuntimeError with a specific message.
+        """
+        nelem = math.prod(in_shape)
+        global_tensor = torch.arange(nelem).view(in_shape)
+        in_dt = distribute_tensor(global_tensor, mesh, placements, src_data_rank=None)
+        comm_mode = CommDebugMode()
+        try:
+            with comm_mode:
+                out_dt = in_dt.view(list(out_shape))
+        except RuntimeError as e:
+            self.assertRegex(
+                str(e),
+                r"is not supported yet"
+                r"|is not evenly divisible by mesh dimension",
+            )
+            return
+        self.assertEqual(comm_mode.get_total_counts(), 0)
+        expected = global_tensor.view(list(out_shape))
+        self.assertEqual(out_dt.full_tensor(), expected)
+
+    def test_dtensor_flatten_split_multi_mesh(self):
+        """Test views producing Split(Flatten) rules.
+
+        Complements test_dtensor_flatten_multi_mesh (pure Flatten rules) and
+        test_dtensor_unflatten_multi_mesh (pure Split(InputDim) rules) by
+        covering the hybrid case: an input dim crosses two output groups,
+        so view_groups produces Split(Flatten(...)) rules.
+
+        Uses {2*M-1, 2*M, 2*M+1} dim values (M = mesh size for the shard
+        mesh dim) to cover even/uneven divisibility, following the same
+        pattern as _run_flatten_single_shard.
+        """
+        for mesh_shape in [(self.world_size,), (3, 2)]:
+            if self.world_size < math.prod(mesh_shape):
+                continue
+            mesh = init_device_mesh(self.device_type, mesh_shape)
+            for shard_mesh_dim in range(mesh.ndim):
+                M = mesh.size(shard_mesh_dim)
+                dim_vals = [2 * M - 1, 2 * M, 2 * M + 1]
+                for a, b in itertools.product(dim_vals, repeat=2):
+                    in_shape = (a, b)
+                    total = a * b
+                    all_factors = self._get_all_factorizations(total)
+                    for out_shape in all_factors:
+                        if in_shape == out_shape:
+                            continue
+                        rules = view_groups(list(in_shape), list(out_shape))
+                        if not any(
+                            isinstance(r, Split) and isinstance(r.input_dim, Flatten)
+                            for r in rules
+                        ):
+                            continue
+                        for shard_dim in range(len(in_shape)):
+                            if in_shape[shard_dim] % M != 0:
+                                continue
+                            placements = tuple(
+                                Shard(shard_dim) if i == shard_mesh_dim else Replicate()
+                                for i in range(mesh.ndim)
+                            )
+                            with self.subTest(
+                                in_shape=in_shape,
+                                out_shape=out_shape,
+                                shard=shard_dim,
+                                mesh_dim=shard_mesh_dim,
+                                mesh_shape=mesh_shape,
+                            ):
+                                self._test_dtensor_flatten_split_case(
+                                    in_shape,
+                                    out_shape,
+                                    placements,
+                                    mesh,
+                                )
+
+    def test_dtensor_flatten_multi_mesh(self):
+        """Test flatten operations across 1D and 2D meshes with all placement patterns.
+
+        Iterates over 1D and 2D mesh configurations to test single-shard (S, SR, RS),
+        multi-shard (SS), and replicate (R, RR) patterns with even/uneven tensor dim sizes.
+        """
+        cases = [
+            ((6,), [("S",), ("R",)]),
+            ((3, 2), [("S", "R"), ("R", "S"), ("S", "S"), ("R", "R")]),
+        ]
+        for mesh_shape, patterns in cases:
+            if self.world_size < math.prod(mesh_shape):
+                continue
+            mesh = init_device_mesh(self.device_type, mesh_shape)
+            mesh_ndim = len(mesh_shape)
+            for pattern in patterns:
+                shard_mesh_dims = [i for i, p in enumerate(pattern) if p == "S"]
+                num_shard = len(shard_mesh_dims)
+                for tensor_ndim in [2, 3, 4]:
+                    for flatten_start in range(tensor_ndim):
+                        for flatten_end in range(flatten_start + 2, tensor_ndim + 1):
+                            if num_shard == 1:
+                                self._run_flatten_single_shard(
+                                    mesh,
+                                    mesh_ndim,
+                                    shard_mesh_dims[0],
+                                    tensor_ndim,
+                                    flatten_start,
+                                    flatten_end,
+                                )
+                            elif num_shard == 2:
+                                self._run_flatten_ss(
+                                    mesh,
+                                    tensor_ndim,
+                                    flatten_start,
+                                    flatten_end,
+                                )
+                            else:
+                                even = 2 * mesh.size(0)
+                                dim_vals = [even - 1, even, even + 1]
+                                all_dims = list(
+                                    itertools.product(dim_vals, repeat=tensor_ndim)
+                                )
+                                rep_placements = tuple([Replicate()] * mesh_ndim)
+                                for tensor_dims in all_dims:
+                                    with self.subTest(
+                                        dims=tensor_dims,
+                                        flat=(flatten_start, flatten_end),
+                                        mesh=mesh_shape,
+                                    ):
+                                        self._test_dtensor_flatten_replicate(
+                                            tensor_dims,
+                                            flatten_start,
+                                            flatten_end,
+                                            mesh,
+                                            rep_placements,
+                                        )
+
+    def _run_flatten_single_shard(
+        self,
+        mesh,
+        mesh_ndim,
+        shard_mesh_dim,
+        tensor_ndim,
+        flatten_start,
+        flatten_end,
+    ):
+        even_val = 2 * mesh.size(shard_mesh_dim)
+        dim_vals = [even_val - 1, even_val, even_val + 1]
+        all_dims = list(itertools.product(dim_vals, repeat=tensor_ndim))
+        for shard_dim in range(flatten_start, flatten_end):
+            for tensor_dims in all_dims:
+                ctx = contextlib.nullcontext()
+                if tensor_dims[shard_dim] % mesh.size(
+                    shard_mesh_dim
+                ) != 0 and shard_dim != (flatten_end - 1):
+                    ctx = self.assertRaisesRegex(
+                        RuntimeError,
+                        "is not evenly divisible by mesh dimension",
+                    )
+                with (
+                    self.subTest(
+                        dims=tensor_dims,
+                        shard=shard_dim,
+                        mesh_dim=shard_mesh_dim,
+                        flat=(flatten_start, flatten_end),
+                    ),
+                    ctx,
+                ):
+                    self._test_dtensor_flatten_single_shard(
+                        tensor_dims,
+                        flatten_start,
+                        flatten_end,
+                        mesh,
+                        shard_dim,
+                        shard_mesh_dim,
+                    )
+
+    def _run_flatten_ss(self, mesh, tensor_ndim, flatten_start, flatten_end):
+        for shard_dim0 in range(flatten_start, flatten_end):
+            for shard_dim1 in range(shard_dim0, flatten_end):
+                dim0_values = [
+                    2 * mesh.size(0) - 1,
+                    2 * mesh.size(0),
+                    2 * mesh.size(0) + 1,
+                ]
+                dim1_values = [
+                    2 * mesh.size(1) - 1,
+                    2 * mesh.size(1),
+                    2 * mesh.size(1) + 1,
+                ]
+                other_dim_value = 2 * mesh.size(0) * mesh.size(1)
+                for dim0_val in dim0_values:
+                    for dim1_val in dim1_values:
+                        tensor_dims = [other_dim_value] * tensor_ndim
+                        tensor_dims[shard_dim0] = dim0_val
+                        if shard_dim0 != shard_dim1:
+                            tensor_dims[shard_dim1] = dim1_val
+                        tensor_dims = tuple(tensor_dims)
+                        local_tensor_dims = list(tensor_dims)
+                        placements = (Shard(shard_dim0), Shard(shard_dim1))
+                        ctx = contextlib.nullcontext()
+                        if local_tensor_dims[shard_dim0] % mesh.size(0) != 0:
+                            ctx = self.assertRaisesRegex(
+                                RuntimeError,
+                                "is not evenly divisible by mesh dimension",
+                            )
+                        local_tensor_dims[shard_dim0] = local_tensor_dims[
+                            shard_dim0
+                        ] // mesh.size(0)
+                        if local_tensor_dims[shard_dim1] % mesh.size(
+                            1
+                        ) != 0 and shard_dim1 != (flatten_end - 1):
+                            ctx = self.assertRaisesRegex(
+                                RuntimeError,
+                                "is not evenly divisible by mesh dimension",
+                            )
+                        with (
+                            self.subTest(
+                                dims=tensor_dims,
+                                shard0=shard_dim0,
+                                shard1=shard_dim1,
+                            ),
+                            ctx,
+                        ):
+                            self._test_dtensor_flatten_2d_ss(
+                                tensor_dims,
+                                flatten_start,
+                                flatten_end,
+                                mesh,
+                                placements,
+                            )
+
+    def _test_dtensor_flatten_single_shard(
+        self,
+        tensor_dims,
+        flatten_start,
+        flatten_end,
+        mesh,
+        shard_dim,
+        shard_mesh_dim,
+    ):
+        mesh_ndim = mesh.ndim
+        placements = tuple(
+            Shard(shard_dim) if i == shard_mesh_dim else Replicate()
+            for i in range(mesh_ndim)
+        )
+        nelem = math.prod(tensor_dims)
+        global_inps: Tensor = torch.arange(nelem).view(tensor_dims)
+        inps = distribute_tensor(global_inps, mesh, placements, src_data_rank=None)
+        viewed_tensor_dims = self._get_viewed_tensor_dims(
+            tensor_dims, flatten_start, flatten_end
+        )
+        comm_mode = CommDebugMode()
+        with comm_mode:
+            inps_viewed = inps.view(viewed_tensor_dims)
+        if shard_dim == flatten_start:
+            expected_placement = Shard(flatten_start)
+        else:
+            split_factor = math.prod(tensor_dims[flatten_start:shard_dim])
+            self.assertGreater(split_factor, 1)
+            expected_placement = _StridedShard(
+                dim=flatten_start, split_factor=split_factor
+            )
+        expected_placements = tuple(
+            expected_placement if i == shard_mesh_dim else Replicate()
+            for i in range(mesh_ndim)
+        )
+        expected_local_tensor = distribute_tensor(
+            global_inps.view(viewed_tensor_dims),
+            mesh,
+            expected_placements,
+            src_data_rank=None,
+        )._local_tensor
+        self.assertEqual(inps_viewed.placements, expected_placements)
+        self.assertEqual(inps_viewed._local_tensor, expected_local_tensor)
+        self.assertEqual(comm_mode.get_total_counts(), 0)
+
+    def _test_dtensor_flatten_replicate(
+        self, tensor_dims, flatten_start, flatten_end, mesh, placements
+    ):
+        mesh_ndim = len(placements)
+        nelem = math.prod(tensor_dims)
+        global_inps: Tensor = torch.arange(nelem).view(tensor_dims)
+        global_inps_replicate: DTensor = distribute_tensor(
+            global_inps, mesh, tuple([Replicate()] * mesh_ndim)
+        )
+        inps = global_inps_replicate.redistribute(mesh, placements)
+        viewed_tensor_dims = self._get_viewed_tensor_dims(
+            tensor_dims, flatten_start, flatten_end
+        )
+        comm_mode = CommDebugMode()
+        with comm_mode:
+            inps_viewed = inps.view(viewed_tensor_dims)
+        expected_placements = tuple([Replicate()] * mesh_ndim)
+        expected_local_tensor = (
+            distribute_tensor(
+                global_inps.view(viewed_tensor_dims),
+                mesh,
+                tuple([Replicate()] * mesh_ndim),
+            )
+            .redistribute(mesh, expected_placements)
+            ._local_tensor
+        )
+        self.assertEqual(inps_viewed.placements, expected_placements)
+        self.assertEqual(inps_viewed._local_tensor, expected_local_tensor)
+        self.assertEqual(comm_mode.get_total_counts(), 0)
+
+    def _get_expected_placements_ss(
+        self,
+        tensor_dims,
+        flatten_start,
+        flatten_end,
+        mesh,
+        placements,
+    ):
+        local_tensor_dims = list(tensor_dims)
+        expected_placements = []
+        for idx, placement in enumerate(placements):
+            self.assertIsInstance(placement, Shard)
+            shard_dim = placement.dim
+            if shard_dim == flatten_start:
+                # S(flatten_start), S(flatten_start) qualifies
+                expected_placement = placement
+            else:
+                split_factor = math.prod(local_tensor_dims[flatten_start:shard_dim])
+                self.assertGreater(split_factor, 1)
+                expected_placement = _StridedShard(
+                    dim=flatten_start, split_factor=split_factor
+                )
+            if local_tensor_dims[shard_dim] % mesh.size(idx) != 0:
+                # uneven shard on last flattened dim is supported
+                self.assertTrue(
+                    _ViewShardingPropagator._is_last_shard_in_flatten_range(
+                        idx, placements, flatten_start, flatten_end
+                    )
+                )
+                local_tensor_dims[shard_dim] = math.ceil(
+                    local_tensor_dims[shard_dim] * 1.0 / mesh.size(idx)
+                )
+            else:
+                local_tensor_dims[shard_dim] = local_tensor_dims[
+                    shard_dim
+                ] // mesh.size(idx)
+            expected_placements.append(expected_placement)
+        return tuple(expected_placements)
+
+    def _test_dtensor_flatten_2d_ss(
+        self,
+        tensor_dims,
+        flatten_start,
+        flatten_end,
+        mesh,
+        placements,
+    ):
+        nelem = math.prod(tensor_dims)
+        global_inps: Tensor = torch.arange(nelem).view(tensor_dims)
+        inps = distribute_tensor(global_inps, mesh, placements, src_data_rank=None)
+        viewed_tensor_dims = self._get_viewed_tensor_dims(
+            tensor_dims, flatten_start, flatten_end
+        )
+
+        comm_mode = CommDebugMode()
+        with comm_mode:
+            inps_viewed = inps.view(viewed_tensor_dims)
+
+        expected_placements = self._get_expected_placements_ss(
+            tensor_dims,
+            flatten_start,
+            flatten_end,
+            mesh,
+            placements,
+        )
+
+        expected_local_tensor = distribute_tensor(
+            global_inps.view(viewed_tensor_dims),
+            mesh,
+            expected_placements,
+            src_data_rank=None,
+        )._local_tensor
+        self.assertEqual(inps_viewed.placements, expected_placements)
+        self.assertEqual(inps_viewed._local_tensor, expected_local_tensor)
+        self.assertEqual(comm_mode.get_total_counts(), 0)
+
+    def test_dtensor_flatten_shard_outside_range(self):
+        """Test that Shard on a dim outside the flatten range passes through correctly.
+
+        When flattening dims [flatten_start, flatten_end), a Shard on a dim outside
+        this range should be preserved with adjusted dim index:
+        - Shard before range: dim unchanged
+        - Shard after range: dim shifts by -(flatten_end - flatten_start - 1)
+        """
+        mesh = init_device_mesh(self.device_type, (self.world_size,))
+        # Use sizes divisible by mesh size to avoid uneven-shard complications
+        dim_size = mesh.size(0) * 2
+        test_cases = [
+            # (tensor_dims, flatten_start, flatten_end, shard_dim, expected_shard_dim)
+            # Shard before range
+            ([dim_size] * 3, 1, 3, 0, 0),
+            ([dim_size] * 4, 1, 3, 0, 0),
+            ([dim_size] * 4, 2, 4, 0, 0),
+            ([dim_size] * 4, 2, 4, 1, 1),
+            # Shard after range
+            ([dim_size] * 4, 0, 2, 2, 1),
+            ([dim_size] * 4, 0, 2, 3, 2),
+            ([dim_size] * 4, 0, 3, 3, 1),
+            ([dim_size] * 4, 1, 3, 3, 2),
+        ]
+
+        for (
+            tensor_dims,
+            flatten_start,
+            flatten_end,
+            shard_dim,
+            expected_shard_dim,
+        ) in test_cases:
+            with self.subTest(
+                dims=tensor_dims, shard=shard_dim, flat=(flatten_start, flatten_end)
+            ):
+                placements = (Shard(shard_dim),)
+                nelem = math.prod(tensor_dims)
+                global_inps = torch.arange(nelem).view(tensor_dims)
+                dt = distribute_tensor(
+                    global_inps, mesh, placements, src_data_rank=None
+                )
+
+                flat_dims = self._get_viewed_tensor_dims(
+                    tensor_dims, flatten_start, flatten_end
+                )
+                comm_mode = CommDebugMode()
+                with comm_mode:
+                    dt_flat = dt.view(flat_dims)
+
+                expected_placements = (Shard(expected_shard_dim),)
+                self.assertEqual(dt_flat.placements, expected_placements)
+                expected_local = distribute_tensor(
+                    global_inps.view(flat_dims),
+                    mesh,
+                    expected_placements,
+                    src_data_rank=None,
+                )._local_tensor
+                self.assertEqual(dt_flat._local_tensor, expected_local)
+                self.assertEqual(comm_mode.get_total_counts(), 0)
+
+        # 2D mesh: test shard outside range with (Shard, Replicate) and (Replicate, Shard)
+        mesh_2d = init_device_mesh(self.device_type, (3, self.world_size // 3))
+        dim_size_2d = mesh_2d.size(0) * mesh_2d.size(1) * 2
+        test_cases_2d = [
+            # (tensor_dims, flatten_start, flatten_end, shard_dim)
+            ([dim_size_2d] * 4, 0, 2, 2),  # shard after range
+            ([dim_size_2d] * 4, 0, 2, 3),  # shard after range
+            ([dim_size_2d] * 4, 1, 3, 0),  # shard before range
+            ([dim_size_2d] * 4, 2, 4, 0),  # shard before range
+            ([dim_size_2d] * 4, 2, 4, 1),  # shard right before range
+        ]
+        for tensor_dims, flatten_start, flatten_end, shard_dim in test_cases_2d:
+            num_merged = flatten_end - flatten_start - 1
+            expected_shard_dim = (
+                shard_dim if shard_dim < flatten_start else shard_dim - num_merged
+            )
+            for mesh_dim_idx in range(2):
+                with self.subTest(
+                    dims=tensor_dims, shard=shard_dim, mesh_dim=mesh_dim_idx
+                ):
+                    placements = tuple(
+                        Shard(shard_dim) if i == mesh_dim_idx else Replicate()
+                        for i in range(2)
+                    )
+                    nelem = math.prod(tensor_dims)
+                    global_inps = torch.arange(nelem).view(tensor_dims)
+                    dt = distribute_tensor(
+                        global_inps, mesh_2d, placements, src_data_rank=None
+                    )
+
+                    flat_dims = self._get_viewed_tensor_dims(
+                        tensor_dims, flatten_start, flatten_end
+                    )
+                    comm_mode = CommDebugMode()
+                    with comm_mode:
+                        dt_flat = dt.view(flat_dims)
+
+                    expected_placements = tuple(
+                        Shard(expected_shard_dim) if i == mesh_dim_idx else Replicate()
+                        for i in range(2)
+                    )
+                    self.assertEqual(dt_flat.placements, expected_placements)
+                    expected_local = distribute_tensor(
+                        global_inps.view(flat_dims),
+                        mesh_2d,
+                        expected_placements,
+                        src_data_rank=None,
+                    )._local_tensor
+                    self.assertEqual(dt_flat._local_tensor, expected_local)
+                    self.assertEqual(comm_mode.get_total_counts(), 0)
+
+    def test_dtensor_flatten_unflatten_roundtrip(self):
+        """Flatten then unflatten should recover the original placements and data.
+
+        Tests the full round-trip: DTensor -> flatten -> unflatten -> compare with original.
+        Covers sharding on dims before, within, and after the flatten range.
+        Also tests unflatten -> flatten direction.
+        """
+        mesh = init_device_mesh(self.device_type, (self.world_size,))
+        dim_size = mesh.size(0) * 2  # evenly divisible
+
+        # flatten -> unflatten round-trip
+        test_cases = [
+            # (tensor_dims, flatten_start, flatten_end, shard_dim)
+            # Shard on first of flattened dims
+            ([dim_size, dim_size], 0, 2, 0),
+            # Shard on second of flattened dims
+            ([dim_size, dim_size], 0, 2, 1),
+            # 3D: shard in range (leading)
+            ([dim_size, dim_size, dim_size], 0, 2, 0),
+            # 3D: shard in range (second)
+            ([dim_size, dim_size, dim_size], 0, 2, 1),
+            # 3D: shard outside range (after)
+            ([dim_size, dim_size, dim_size], 0, 2, 2),
+            # 3D: shard before range
+            ([dim_size, dim_size, dim_size], 1, 3, 0),
+            # 3D: shard in range (first of second pair)
+            ([dim_size, dim_size, dim_size], 1, 3, 1),
+            # 3D: shard in range (second of second pair)
+            ([dim_size, dim_size, dim_size], 1, 3, 2),
+            # 4D: shard first of 3 flattened
+            ([dim_size] * 4, 0, 3, 0),
+            # 4D: shard last of 3 flattened
+            ([dim_size] * 4, 0, 3, 2),
+            # 4D: shard outside range (after)
+            ([dim_size] * 4, 0, 3, 3),
+            # 4D: shard before range
+            ([dim_size] * 4, 1, 3, 0),
+            # 4D: shard after range
+            ([dim_size] * 4, 1, 3, 3),
+        ]
+
+        for tensor_dims, flatten_start, flatten_end, shard_dim in test_cases:
+            with self.subTest(dims=tensor_dims, shard=shard_dim, direction="flatten"):
+                placements = (Shard(shard_dim),)
+                nelem = math.prod(tensor_dims)
+                global_inps = torch.arange(nelem).view(tensor_dims)
+                dt = distribute_tensor(
+                    global_inps, mesh, placements, src_data_rank=None
+                )
+
+                # Flatten
+                flat_dims = self._get_viewed_tensor_dims(
+                    tensor_dims, flatten_start, flatten_end
+                )
+                dt_flat = dt.view(flat_dims)
+
+                # Unflatten back
+                dt_roundtrip = dt_flat.view(tensor_dims)
+
+                self.assertEqual(dt_roundtrip.placements, placements)
+                self.assertEqual(dt_roundtrip._local_tensor, dt._local_tensor)
+
+        # unflatten -> flatten round-trip
+        # Start with a flattened tensor, unflatten it, then flatten back
+        for tensor_dims, flatten_start, flatten_end, shard_dim in test_cases:
+            with self.subTest(dims=tensor_dims, shard=shard_dim, direction="unflatten"):
+                flat_dims = self._get_viewed_tensor_dims(
+                    tensor_dims, flatten_start, flatten_end
+                )
+                # Compute the shard dim in the flattened view
+                if shard_dim < flatten_start:
+                    flat_shard_dim = shard_dim
+                elif shard_dim < flatten_end:
+                    flat_shard_dim = (
+                        flatten_start  # all flattened dims merge to flatten_start
+                    )
+                else:
+                    flat_shard_dim = shard_dim - (flatten_end - flatten_start - 1)
+
+                placements = (Shard(flat_shard_dim),)
+                nelem = math.prod(flat_dims)
+                global_inps = torch.arange(nelem).view(flat_dims)
+                dt = distribute_tensor(
+                    global_inps, mesh, placements, src_data_rank=None
+                )
+
+                # Unflatten
+                dt_unflat = dt.view(tensor_dims)
+
+                # Flatten back
+                dt_roundtrip = dt_unflat.view(flat_dims)
+
+                self.assertEqual(dt_roundtrip.placements, placements)
+                self.assertEqual(dt_roundtrip._local_tensor, dt._local_tensor)
+
+    def generate_tensor_dims_1d(
+        self, tensor_ndim, flatten_start, flatten_end, shard_dim, mesh
+    ):
+        tensor_dims_unflatten = [2 * mesh.size(0) + 1] * tensor_ndim
+        for tensor_dim in [
+            2 * mesh.size(0) - 1,
+            2 * mesh.size(0),
+            2 * mesh.size(0) + 1,
+        ]:
+            tensor_dims_unflatten[shard_dim] = tensor_dim
+            local_tensor_dims_unflatten = list(tensor_dims_unflatten)
+            local_tensor_dims_unflatten[shard_dim] = math.ceil(
+                tensor_dim * 1.0 / mesh.size(0)
+            )
+            nelem_flatten = math.prod(tensor_dims_unflatten[flatten_start:flatten_end])
+            tensor_dims_flatten = (
+                tensor_dims_unflatten[0:flatten_start]
+                + [nelem_flatten]
+                + tensor_dims_unflatten[flatten_end:]
+            )
+            yield (
+                tensor_dims_unflatten,
+                local_tensor_dims_unflatten,
+                tensor_dims_flatten,
+            )
+
+    def generate_tensor_dims_1d_after_flatten(
+        self, tensor_ndim, unflatten_dim, shard_dim, mesh
+    ):
+        tensor_dims = [mesh.size(0) * mesh.size(0)] * tensor_ndim
+        for unflatten_dim_value in [
+            mesh.size(0) * mesh.size(0) - 1,
+            mesh.size(0) * mesh.size(0),
+            mesh.size(0) * mesh.size(0) + 1,
+        ]:
+            for shard_dim_value in [
+                mesh.size(0) * mesh.size(0) - 1,
+                mesh.size(0) * mesh.size(0),
+                mesh.size(0) * mesh.size(0) + 1,
+            ]:
+                tensor_dims[unflatten_dim] = unflatten_dim_value
+                tensor_dims[shard_dim] = shard_dim_value
+                yield list(tensor_dims)
+
+    def test_dtensor_unflatten_1d(self):
+        mesh: DeviceMesh = init_device_mesh(self.device_type, (self.world_size,))
+
+        # flatten -> view -> unflatten
+        for tensor_ndim in [2, 3, 4]:
+            for flatten_start in range(tensor_ndim):
+                for flatten_end in range(flatten_start + 2, tensor_ndim + 1):
+                    for shard_dim in range(flatten_start, flatten_end):
+                        expected_placements = (Shard(shard_dim),)
+                        for (
+                            tensor_dims_unflatten,
+                            local_tensor_dims_unflatten,
+                            tensor_dims_flatten,
+                        ) in self.generate_tensor_dims_1d(
+                            tensor_ndim, flatten_start, flatten_end, shard_dim, mesh
+                        ):
+                            ctx = contextlib.nullcontext()
+                            is_uneven = (
+                                tensor_dims_unflatten[shard_dim] % mesh.size(0) != 0
+                            )
+                            is_last_dim = shard_dim == flatten_end - 1
+                            if is_uneven and not is_last_dim:
+                                ctx = self.assertRaisesRegex(
+                                    RuntimeError,
+                                    "is not evenly divisible by mesh dimension",
+                                )
+                            with (
+                                self.subTest(
+                                    dims=tensor_dims_unflatten,
+                                    shard=shard_dim,
+                                    flat=(flatten_start, flatten_end),
+                                ),
+                                ctx,
+                            ):
+                                self._test_dtensor_unflatten_1d_shard(
+                                    tensor_dims_unflatten,
+                                    local_tensor_dims_unflatten,
+                                    tensor_dims_flatten,
+                                    flatten_start,
+                                    expected_placements,
+                                    mesh,
+                                )
+
+        # any factoring on unflatten_dim
+        for tensor_ndim in [1, 2, 3, 4]:
+            for unflatten_dim in range(tensor_ndim):
+                for shard_dim in range(tensor_ndim):
+                    for tensor_dims in self.generate_tensor_dims_1d_after_flatten(
+                        tensor_ndim, unflatten_dim, shard_dim, mesh
+                    ):
+                        with self.subTest(
+                            dims=tensor_dims, shard=shard_dim, unflatten=unflatten_dim
+                        ):
+                            placements = (Shard(shard_dim),)
+                            self._test_dtensor_unflatten_1d_shard_arbitrary(
+                                tensor_dims,
+                                unflatten_dim,
+                                placements,
+                                mesh,
+                            )
+
+        # Replicate: unflatten should preserve Replicate placement
+        for tensor_ndim in [1, 2, 3, 4]:
+            for unflatten_dim in range(tensor_ndim):
+                tensor_dims = [6] * tensor_ndim
+                tensor_dims[unflatten_dim] = 12  # will unflatten to (3, 4)
+                global_tensor = torch.arange(math.prod(tensor_dims)).view(tensor_dims)
+                dt = distribute_tensor(
+                    global_tensor, mesh, (Replicate(),), src_data_rank=None
+                )
+                # Unflatten dimension to (3, 4)
+                unflatten_shape = list(tensor_dims)
+                unflatten_shape[unflatten_dim : unflatten_dim + 1] = [3, 4]
+                dt_unflattened = dt.view(unflatten_shape)
+                self.assertEqual(dt_unflattened.placements, (Replicate(),))
+                self.assertEqual(dt_unflattened.shape, torch.Size(unflatten_shape))
+
+    def _test_dtensor_unflatten_1d_shard(
+        self,
+        tensor_dims_unflatten,
+        local_tensor_dims_unflatten,
+        tensor_dims_flatten,
+        flatten_start,
+        expected_placements,
+        mesh,
+    ):
+        shard_dim = expected_placements[0].dim
+        split_factor = math.prod(local_tensor_dims_unflatten[flatten_start:shard_dim])
+        if split_factor == 1:
+            self.assertEqual(shard_dim, flatten_start)
+            placements = (Shard(flatten_start),)
+        else:
+            placements = (_StridedShard(flatten_start, split_factor=split_factor),)
+        nelem = math.prod(tensor_dims_flatten)
+        global_inps = torch.arange(nelem).view(tensor_dims_flatten)
+        inps = distribute_tensor(global_inps, mesh, placements, src_data_rank=None)
+        inps_viewed = inps.view(tensor_dims_unflatten)
+        self.assertEqual(inps_viewed.placements, expected_placements)
+        expected_local_tensor = distribute_tensor(
+            torch.arange(nelem).view(tensor_dims_unflatten),
+            mesh,
+            expected_placements,
+            src_data_rank=None,
+        )._local_tensor
+        self.assertEqual(inps_viewed._local_tensor, expected_local_tensor)
+
+    def _test_dtensor_unflatten_1d_shard_arbitrary(
+        self, tensor_dims, unflatten_dim, placements, mesh
+    ):
+        shard_dim = placements[0].dim
+        self.assertIsInstance(placements[0], Shard)
+        self.assertNotIsInstance(placements[0], _StridedShard)
+
+        tensor_dims = list(tensor_dims)  # Make a mutable copy
+
+        # Get all non-trivial factorizations of unflatten_dim size
+        unflatten_size = tensor_dims[unflatten_dim]
+        factorizations = self._get_all_factorizations(unflatten_size)
+
+        # Skip if no non-trivial factorization exists (prime number)
+        if not factorizations:
+            return
+
+        # Create the global tensor once (reused for all factorizations)
+        nelem = math.prod(tensor_dims)
+        global_inps = torch.arange(nelem).view(tensor_dims)
+
+        # Distribute the tensor
+        inps = distribute_tensor(global_inps, mesh, placements, src_data_rank=None)
+
+        # Test each factorization (can be any length >= 2)
+        for factors in factorizations:
+            tensor_dims_unflatten = (
+                tensor_dims[:unflatten_dim]
+                + list(factors)
+                + tensor_dims[unflatten_dim + 1 :]
+            )
+
+            # Determine if we expect an error
+            first_factor = factors[0]
+
+            ctx = contextlib.nullcontext()
+            expect_error = False
+
+            if shard_dim == unflatten_dim:
+                # When unflattening the sharded dimension, check factor alignment
+                uneven_shard = tensor_dims[shard_dim] % mesh.size(0) != 0
+                first_factor_aligned = first_factor % mesh.size(0) == 0
+
+                expect_error = not first_factor_aligned or uneven_shard
+
+            if expect_error:
+                ctx = self.assertRaisesRegex(
+                    RuntimeError, "is not evenly divisible by mesh dimension"
+                )
+
+            with ctx:
+                comm_mode = CommDebugMode()
+                with comm_mode:
+                    inps_viewed = inps.view(tensor_dims_unflatten)
+
+                # Number of new dimensions added by unflatten
+                num_new_dims = len(factors) - 1
+
+                # Compute expected placements after unflatten
+                if shard_dim < unflatten_dim:
+                    expected_placements = (Shard(shard_dim),)
+                elif shard_dim == unflatten_dim:
+                    expected_placements = (Shard(unflatten_dim),)
+                else:
+                    expected_placements = (Shard(shard_dim + num_new_dims),)
+
+                self.assertEqual(inps_viewed.placements, expected_placements)
+                self.assertEqual(comm_mode.get_total_counts(), 0)
+
+                expected_local = distribute_tensor(
+                    global_inps.view(tensor_dims_unflatten),
+                    mesh,
+                    expected_placements,
+                    src_data_rank=None,
+                )._local_tensor
+                self.assertEqual(inps_viewed._local_tensor, expected_local)
+
+    def test_dtensor_unflatten_multi_mesh(self):
+        """Test unflatten across 2D and 3D meshes with all placement patterns.
+
+        Iterates over 2D and 3D mesh configurations to test multi-shard (SS, SSS),
+        mixed (R+SS, RR+SS), and replicate (RR, RRR) patterns.
+        """
+        # (mesh_shape, num_replicate_dims)
+        cases = [
+            ((3, 2), 0),  # (SS, SS)
+            ((3, 2), 1),  # (R, SS)
+            ((3, 2), 2),  # (R, R)
+            ((2, 2, 2), 0),  # (SS, SS, SS)
+            ((2, 2, 2), 1),  # (R, SS, SS)
+            ((2, 2, 2), 2),  # (R, R, SS)
+            ((2, 2, 2), 3),  # (R, R, R)
+        ]
+        for mesh_shape, num_rep in cases:
+            if self.world_size < math.prod(mesh_shape):
+                continue
+            mesh = init_device_mesh(self.device_type, mesh_shape)
+            mesh_ndim = len(mesh_shape)
+            num_shard = mesh_ndim - num_rep
+            flattened_size = math.prod(mesh_shape) * 12
+            factorizations = self._get_all_factorizations(flattened_size)
+
+            if num_shard == 0:
+                # All-replicate: test that Replicate is preserved
+                for factors in factorizations:
+                    for prefix_size in [0, 1, 2]:
+                        with self.subTest(
+                            factors=factors,
+                            prefix=prefix_size,
+                            mesh=mesh_shape,
+                        ):
+                            tensor_dims_unflatten = (
+                                [6] * prefix_size + list(factors) + [3]
+                            )
+                            nelem_flatten = math.prod(factors)
+                            tensor_dims_flatten = (
+                                [6] * prefix_size + [nelem_flatten] + [3]
+                            )
+                            nelem = math.prod(tensor_dims_flatten)
+                            global_tensor = torch.arange(nelem).view(
+                                tensor_dims_flatten
+                            )
+                            rep_placements = tuple([Replicate()] * mesh_ndim)
+                            dt = distribute_tensor(
+                                global_tensor,
+                                mesh,
+                                rep_placements,
+                                src_data_rank=None,
+                            )
+                            dt_unflattened = dt.view(tensor_dims_unflatten)
+                            self.assertEqual(
+                                dt_unflattened.placements,
+                                rep_placements,
+                            )
+                            self.assertEqual(
+                                dt_unflattened.shape,
+                                torch.Size(tensor_dims_unflatten),
+                            )
+            else:
+                # Sharded cases: pick shard indices from factorizations
+                min_factors = num_shard + 1  # need prefix + num_shard shard dims
+                valid = [f for f in factorizations if len(f) >= min_factors]
+                for factors in valid:
+                    n = len(factors)
+                    for shard_indices in itertools.combinations(
+                        range(1, n),
+                        num_shard,
+                    ):
+                        # Error detection: uneven on non-(last-shard AND last-factor)
+                        expect_error = False
+                        for rank, si in enumerate(shard_indices):
+                            mesh_dim = num_rep + rank
+                            is_last_shard = rank == num_shard - 1
+                            is_last_factor = si == n - 1
+                            if factors[si] % mesh.size(mesh_dim) != 0 and not (
+                                is_last_shard and is_last_factor
+                            ):
+                                expect_error = True
+                                break
+                        with self.subTest(
+                            factors=factors,
+                            shard=shard_indices,
+                            mesh=mesh_shape,
+                        ):
+                            if expect_error:
+                                with self.assertRaisesRegex(
+                                    RuntimeError,
+                                    "is not evenly divisible by mesh dimension|"
+                                    "do not support inputs with use_strided_shard_as_shard_order",
+                                ):
+                                    self._test_dtensor_unflatten_factors(
+                                        factors,
+                                        shard_indices,
+                                        num_rep,
+                                        mesh,
+                                    )
+                            else:
+                                self._test_dtensor_unflatten_factors(
+                                    factors,
+                                    shard_indices,
+                                    num_rep,
+                                    mesh,
+                                )
+
+    def _test_dtensor_unflatten_factors(
+        self,
+        factors,
+        shard_indices,
+        num_rep,
+        mesh,
+    ):
+        """Test unflatten with _StridedShard placements for any mesh dimensionality.
+
+        Args:
+            factors: Tuple of factors representing the unflatten shape
+            shard_indices: Tuple of indices within factors for each sharded mesh dim
+            num_rep: Number of leading Replicate mesh dims
+            mesh: DeviceMesh
+        """
+        num_shard = len(shard_indices)
+        # SS cases need a suffix dim outside the flatten range
+        suffix = [3] if num_shard >= 2 else []
+        tensor_dims_unflatten = list(factors) + suffix
+        nelem_flatten = math.prod(factors)
+        tensor_dims_flatten = [nelem_flatten] + suffix
+
+        # Build input placements and expected output placements
+        placements = []
+        expected_placements = []
+        for i in range(mesh.ndim):
+            if i < num_rep:
+                placements.append(Replicate())
+                expected_placements.append(Replicate())
+            else:
+                rank = i - num_rep
+                si = shard_indices[rank]
+                # split_factor = prod(factors[:si]) / prod(earlier mesh dims)
+                sf = math.prod(factors[:si])
+                for j in range(rank):
+                    sf //= mesh.size(num_rep + j)
+                placements.append(_StridedShard(0, split_factor=sf))
+                expected_placements.append(Shard(si))
+
+        placements = tuple(placements)
+        expected_placements = tuple(expected_placements)
+
+        nelem = math.prod(tensor_dims_flatten)
+        global_inps = torch.arange(nelem).view(tensor_dims_flatten)
+        inps = distribute_tensor(global_inps, mesh, placements, src_data_rank=None)
+        # These _StridedShard placements represent literal strided layout
+        # (as produced by flatten), not shard-order encoding.
+        inps._spec.use_strided_shard_as_shard_order = False
+        inps._spec.shard_order = DTensorSpec.compute_default_shard_order(
+            inps._spec.placements
+        )
+
+        comm_mode = CommDebugMode()
+        with comm_mode:
+            inps_viewed = inps.view(tensor_dims_unflatten)
+
+        self.assertEqual(inps_viewed.placements, expected_placements)
+        self.assertEqual(comm_mode.get_total_counts(), 0)
+        expected_local_tensor = distribute_tensor(
+            torch.arange(nelem).view(tensor_dims_unflatten),
+            mesh,
+            expected_placements,
+            src_data_rank=None,
+        )._local_tensor
+        self.assertEqual(inps_viewed._local_tensor, expected_local_tensor)
+
+    def test_dtensor_flatten_unflatten_2d_reversed_mesh(self):
+        """Test flatten/unflatten with reversed mesh shape (2, 3) to catch ordering bugs."""
+        self.assertEqual(self.world_size, 6)
+        mesh = init_device_mesh(self.device_type, (2, self.world_size // 2))
+        dim_size = mesh.size(0) * mesh.size(1) * 2  # divisible by both mesh dims
+
+        # Flatten with (S, S) pattern
+        tensor_ndim = 3
+        for flatten_start in range(tensor_ndim):
+            for flatten_end in range(flatten_start + 2, tensor_ndim + 1):
+                for shard_dim0 in range(flatten_start, flatten_end):
+                    for shard_dim1 in range(shard_dim0, flatten_end):
+                        with self.subTest(
+                            shard0=shard_dim0,
+                            shard1=shard_dim1,
+                            flat=(flatten_start, flatten_end),
+                        ):
+                            tensor_dims = tuple([dim_size] * tensor_ndim)
+                            placements = (Shard(shard_dim0), Shard(shard_dim1))
+                            self._test_dtensor_flatten_2d_ss(
+                                tensor_dims,
+                                flatten_start,
+                                flatten_end,
+                                mesh,
+                                placements,
+                            )
+
+        # Unflatten with (SS, SS) pattern — representative factorizations
+        factors_list = [(6, 4, 3), (4, 6, 2), (3, 2, 6), (2, 6, 4)]
+        for factors in factors_list:
+            for shard_idx0 in range(1, len(factors) - 1):
+                for shard_idx1 in range(shard_idx0 + 1, len(factors)):
+                    with self.subTest(factors=factors, s0=shard_idx0, s1=shard_idx1):
+                        factor0 = factors[shard_idx0]
+                        factor1 = factors[shard_idx1]
+                        uneven0 = factor0 % mesh.size(0) != 0
+                        is_last1 = shard_idx1 == len(factors) - 1
+                        uneven1 = factor1 % mesh.size(1) != 0 and not is_last1
+                        if uneven0 or uneven1:
+                            with self.assertRaisesRegex(
+                                RuntimeError, "is not evenly divisible"
+                            ):
+                                self._test_dtensor_unflatten_factors(
+                                    factors,
+                                    (shard_idx0, shard_idx1),
+                                    0,
+                                    mesh,
+                                )
+                        else:
+                            self._test_dtensor_unflatten_factors(
+                                factors,
+                                (shard_idx0, shard_idx1),
+                                0,
+                                mesh,
+                            )
+
+    def test_dtensor_unflatten_ss_and_s_same_dim(self):
+        """Test unflatten when _StridedShard and Shard both shard the same tensor dim."""
+        mesh = init_device_mesh(self.device_type, (3, self.world_size // 3))
+
+        global_tensor = torch.arange(4 * 6 * 3).view(4, 6, 3)
+        inps = distribute_tensor(global_tensor, mesh, [Shard(1), Shard(0)])
+
+        # Flatten dims 0,1 -> [24, 3] with [_StridedShard(0, 4), Shard(0)]
+        comm_mode = CommDebugMode()
+        with comm_mode:
+            flattened = inps.flatten(0, 1)
+        self.assertEqual(comm_mode.get_total_counts(), 0)
+        self.assertEqual(flattened.shape, (24, 3))
+        self.assertIsInstance(flattened.placements[0], _StridedShard)
+        self.assertEqual(flattened.placements[0].split_factor, 4)
+        self.assertIsInstance(flattened.placements[1], Shard)
+
+        # Unflatten dim 0 (24) -> (12, 2): both mesh dims stay on output dim 0
+        expected_placements = (_StridedShard(0, split_factor=4), Shard(0))
+        comm_mode = CommDebugMode()
+        with comm_mode:
+            unflattened = flattened.view(12, 2, 3)
+        self.assertEqual(comm_mode.get_total_counts(), 0)
+        self.assertEqual(unflattened.placements, expected_placements)
+        expected_local = distribute_tensor(
+            torch.arange(4 * 6 * 3).view(12, 2, 3),
+            mesh,
+            expected_placements,
+            src_data_rank=None,
+        )._local_tensor
+        self.assertEqual(unflattened._local_tensor, expected_local)
+
+    def test_flatten_then_sum_non_shard_dim(self):
+        """Verify _StridedShard correctness through sum on a non-shard dim.
+
+        sum uses map_placements_after_reduction which must preserve
+        _StridedShard (with split_factor) when remapping dims after
+        the reduced dim is removed.
+        """
+        mesh = init_device_mesh(self.device_type, (self.world_size,))
+        shape = (4, self.world_size * 2, 6)
+        full = torch.randn(*shape, device=self.device_type)
+        dt = distribute_tensor(full, mesh, [Shard(1)])
+        dt_flat = dt.flatten(0, 1)  # _StridedShard(dim=0)
+        flat_full = full.flatten(0, 1)
+
+        self.assertIsInstance(dt_flat.placements[0], _StridedShard)
+
+        # reduce dim 1 (non-shard): _StridedShard(dim=0) stays at dim 0
+        result = dt_flat.sum(dim=1)
+        self.assertIsInstance(result.placements[0], _StridedShard)
+        self.assertEqual(result.placements[0].dim, 0)
+        self.assertEqual(
+            result.placements[0].split_factor, dt_flat.placements[0].split_factor
+        )
+        self.assertEqual(result.full_tensor(), flat_full.sum(dim=1))
+
+        # reduce dim 0 (shard dim): should produce Partial
+        result2 = dt_flat.sum(dim=0)
+        self.assertTrue(result2.placements[0].is_partial())
+        self.assertEqual(result2.full_tensor(), flat_full.sum(dim=0))
+
+        # _StridedShard on higher dim: reduce a dim below it to trigger remap
+        shape2 = (3, 5, self.world_size * 2)
+        full2 = torch.randn(*shape2, device=self.device_type)
+        dt2 = distribute_tensor(full2, mesh, [Shard(2)])
+        dt2_flat = dt2.flatten(1, 2)  # _StridedShard(dim=1)
+        flat_full2 = full2.flatten(1, 2)
+
+        self.assertIsInstance(dt2_flat.placements[0], _StridedShard)
+        self.assertEqual(dt2_flat.placements[0].dim, 1)
+
+        # reduce dim 0: _StridedShard(dim=1) remaps to _StridedShard(dim=0)
+        result3 = dt2_flat.sum(dim=0)
+        self.assertIsInstance(result3.placements[0], _StridedShard)
+        self.assertEqual(result3.placements[0].dim, 0)
+        self.assertEqual(
+            result3.placements[0].split_factor, dt2_flat.placements[0].split_factor
+        )
+        self.assertEqual(result3.full_tensor(), flat_full2.sum(dim=0))
+
+    def test_flatten_then_softmax(self):
+        """Verify _StridedShard correctness through softmax.
+
+        softmax uses replicate_reduction_dims which only checks isinstance(p, Shard).
+        _StridedShard on the softmax dim may not be replicated, producing wrong results.
+        """
+        mesh = init_device_mesh(self.device_type, (self.world_size,))
+        shape = (4, self.world_size * 2, 6)
+        full = torch.randn(*shape, device=self.device_type)
+        dt = distribute_tensor(full, mesh, [Shard(1)])
+        dt_flat = dt.flatten(0, 1)
+        flat_full = full.flatten(0, 1)
+
+        self.assertIsInstance(dt_flat.placements[0], _StridedShard)
+
+        # softmax on shard dim
+        result = torch.softmax(dt_flat, dim=0)
+        self.assertEqual(result.full_tensor(), torch.softmax(flat_full, dim=0))
+
+        # softmax on non-shard dim
+        result = torch.softmax(dt_flat, dim=-1)
+        self.assertEqual(result.full_tensor(), torch.softmax(flat_full, dim=-1))
+
+    def test_flatten_then_layer_norm(self):
+        """Verify _StridedShard correctness through layer_norm.
+
+        layer_norm uses _replicate_dims_start_at which only checks isinstance(p, Shard).
+        _StridedShard on normalized dims may not be replicated.
+        """
+        mesh = init_device_mesh(self.device_type, (self.world_size,))
+        shape = (4, self.world_size * 2, 6)
+        full = torch.randn(*shape, device=self.device_type)
+        dt = distribute_tensor(full, mesh, [Shard(1)])
+        dt_flat = dt.flatten(0, 1)
+        flat_full = full.flatten(0, 1)
+
+        self.assertIsInstance(dt_flat.placements[0], _StridedShard)
+
+        # layer_norm on last dim (non-shard): shard dim not in normalized dims
+        result = torch.nn.functional.layer_norm(dt_flat, [6])
+        self.assertEqual(
+            result.full_tensor(),
+            torch.nn.functional.layer_norm(flat_full, [6]),
+        )
+
+        # layer_norm covering shard dim
+        result = torch.nn.functional.layer_norm(dt_flat, list(dt_flat.shape))
+        self.assertEqual(
+            result.full_tensor(),
+            torch.nn.functional.layer_norm(flat_full, list(flat_full.shape)),
+        )
+
+    def test_flatten_then_transpose(self):
+        """Verify _StridedShard correctness through transpose.
+
+        aten.transpose.int goes through view op propagation which handles
+        _StridedShard. aten.t uses transpose_strategy in _matrix_ops.py which
+        swaps the dim for both Shard and _StridedShard.
+        """
+        mesh = init_device_mesh(self.device_type, (self.world_size,))
+        shape = (4, self.world_size * 2, 6)
+        full = torch.randn(*shape, device=self.device_type)
+        dt = distribute_tensor(full, mesh, [Shard(1)])
+        dt_flat = dt.flatten(0, 1)
+        flat_full = full.flatten(0, 1)
+
+        self.assertIsInstance(dt_flat.placements[0], _StridedShard)
+
+        # transpose shard dim with another dim (aten.transpose.int)
+        result = dt_flat.transpose(0, 1)
+        self.assertEqual(result.full_tensor(), flat_full.transpose(0, 1))
+
+        # t() on 2D tensor (aten.t via transpose_strategy in _matrix_ops.py)
+        # _StridedShard(dim=0) → _StridedShard(dim=1), split_factor preserved
+        result = dt_flat.t()
+        self.assertIsInstance(result.placements[0], _StridedShard)
+        self.assertEqual(result.placements[0].dim, 1)
+        self.assertEqual(
+            result.placements[0].split_factor, dt_flat.placements[0].split_factor
+        )
+        self.assertEqual(result.full_tensor(), flat_full.t())
+
+        # t() with _StridedShard(dim=1) → _StridedShard(dim=0)
+        # Shard(2) on (4, 6, ws*2) + flatten(1,2) → _StridedShard(dim=1)
+        shape1 = (4, 6, self.world_size * 2)
+        full1 = torch.randn(*shape1, device=self.device_type)
+        dt1 = distribute_tensor(full1, mesh, [Shard(2)])
+        dt1_flat = dt1.flatten(1, 2)  # (4, 6*ws*2) with _StridedShard(dim=1)
+        self.assertIsInstance(dt1_flat.placements[0], _StridedShard)
+        self.assertEqual(dt1_flat.placements[0].dim, 1)
+
+        result1 = dt1_flat.t()
+        self.assertIsInstance(result1.placements[0], _StridedShard)
+        self.assertEqual(result1.placements[0].dim, 0)
+        self.assertEqual(
+            result1.placements[0].split_factor, dt1_flat.placements[0].split_factor
+        )
+        self.assertEqual(result1.full_tensor(), full1.flatten(1, 2).t())
+
+        # t() with uneven shapes: shard dim not divisible by world_size
+        # _StridedShard(dim=0) → _StridedShard(dim=1)
+        shape2 = (
+            3,
+            self.world_size * 2,
+            5,
+        )  # flatten(0,1) gives (3*ws*2, 5), uneven in dim 0
+        full2 = torch.randn(*shape2, device=self.device_type)
+        dt2 = distribute_tensor(full2, mesh, [Shard(1)])
+        dt2_flat = dt2.flatten(0, 1)
+        self.assertIsInstance(dt2_flat.placements[0], _StridedShard)
+
+        result2 = dt2_flat.t()
+        self.assertIsInstance(result2.placements[0], _StridedShard)
+        self.assertEqual(result2.placements[0].dim, 1)
+        self.assertEqual(
+            result2.placements[0].split_factor, dt2_flat.placements[0].split_factor
+        )
+        self.assertEqual(result2.full_tensor(), full2.flatten(0, 1).t())
+
+        # _StridedShard(dim=1) → _StridedShard(dim=0), uneven
+        shape3 = (
+            5,
+            3,
+            self.world_size * 2,
+        )  # flatten(1,2) gives (5, 3*ws*2), uneven in dim 1
+        full3 = torch.randn(*shape3, device=self.device_type)
+        dt3 = distribute_tensor(full3, mesh, [Shard(2)])
+        dt3_flat = dt3.flatten(1, 2)
+        self.assertIsInstance(dt3_flat.placements[0], _StridedShard)
+        self.assertEqual(dt3_flat.placements[0].dim, 1)
+
+        result3 = dt3_flat.t()
+        self.assertIsInstance(result3.placements[0], _StridedShard)
+        self.assertEqual(result3.placements[0].dim, 0)
+        self.assertEqual(
+            result3.placements[0].split_factor, dt3_flat.placements[0].split_factor
+        )
+        self.assertEqual(result3.full_tensor(), full3.flatten(1, 2).t())
+
+    def test_flatten_then_nll_loss(self):
+        """Verify _StridedShard correctness through nll_loss.
+
+        nll_loss_forward_strategy uses replicate_reduction_dims on the channel
+        dim and _skip_dim to build target placements.
+        """
+        mesh = init_device_mesh(self.device_type, (self.world_size,))
+        num_classes = 10
+
+        # 2D input: (N, C) — _StridedShard on batch dim (dim 0), which is
+        # below channel_dim (dim 1), so _skip_dim preserves it unchanged.
+        shape = (4, self.world_size * 2)
+        full_input = torch.randn(*shape, num_classes, device=self.device_type)
+        full_target = torch.randint(0, num_classes, shape, device=self.device_type)
+        dist.broadcast(full_target, src=0)
+
+        dt_input = distribute_tensor(full_input, mesh, [Shard(1)])
+        dt_input_flat = dt_input.flatten(0, 1)
+        dt_target = distribute_tensor(full_target, mesh, [Shard(1)])
+        dt_target_flat = dt_target.flatten(0, 1)
+
+        self.assertIsInstance(dt_input_flat.placements[0], _StridedShard)
+
+        input_full = full_input.flatten(0, 1)
+        target_full = full_target.flatten(0, 1)
+
+        result = torch.nn.functional.nll_loss(
+            torch.log_softmax(dt_input_flat, dim=-1), dt_target_flat
+        )
+        expected = torch.nn.functional.nll_loss(
+            torch.log_softmax(input_full, dim=-1), target_full
+        )
+        self.assertEqual(result.full_tensor(), expected)
+
+        # 3D input: (N, C, D) — _StridedShard on spatial dim (dim 2), which
+        # is above channel_dim (dim 1), so _skip_dim must shift it from
+        # dim 2 to dim 1 while preserving split_factor.
+        N, C, D1, D2 = 2, num_classes, 3, self.world_size * 2
+        full_input_3d = torch.randn(N, C, D1, D2, device=self.device_type)
+        full_target_3d = torch.randint(0, C, (N, D1, D2), device=self.device_type)
+        dist.broadcast(full_target_3d, src=0)
+
+        dt_input_3d = distribute_tensor(full_input_3d, mesh, [Shard(3)])
+        dt_input_3d_flat = dt_input_3d.flatten(2, 3)
+        dt_target_3d = distribute_tensor(full_target_3d, mesh, [Shard(2)])
+        dt_target_3d_flat = dt_target_3d.flatten(1, 2)
+
+        self.assertIsInstance(dt_input_3d_flat.placements[0], _StridedShard)
+        self.assertEqual(dt_input_3d_flat.placements[0].dim, 2)
+
+        input_3d_full = full_input_3d.flatten(2, 3)
+        target_3d_full = full_target_3d.flatten(1, 2)
+
+        result_3d = torch.nn.functional.nll_loss(
+            torch.log_softmax(dt_input_3d_flat, dim=1), dt_target_3d_flat
+        )
+        expected_3d = torch.nn.functional.nll_loss(
+            torch.log_softmax(input_3d_full, dim=1), target_3d_full
+        )
+        self.assertEqual(result_3d.full_tensor(), expected_3d)
+
+    def test_flatten_then_select(self):
+        """Verify _StridedShard correctness through select.
+
+        select removes a dim, so select_int_strategy must detect
+        _StridedShard (which is_sharded() misses) and call
+        shift_shard_dims_after_remove to adjust the shard dim.
+        """
+        mesh = init_device_mesh(self.device_type, (self.world_size,))
+        shape = (2, 3, 4, self.world_size * 2)
+        full = torch.randn(*shape, device=self.device_type)
+        dt = distribute_tensor(full, mesh, [Shard(3)])
+        dt_flat = dt.flatten(2, 3)  # (2, 3, 4*ws*2) with _StridedShard(dim=2)
+
+        self.assertIsInstance(dt_flat.placements[0], _StridedShard)
+        orig_split_factor = dt_flat.placements[0].split_factor
+
+        # select on dim 0: _StridedShard dim shifts from 2 to 1, split_factor preserved
+        result = dt_flat.select(0, 0)
+        expected = full.flatten(2, 3).select(0, 0)
+        self.assertIsInstance(result.placements[0], _StridedShard)
+        self.assertEqual(result.placements[0].dim, 1)
+        self.assertEqual(result.placements[0].split_factor, orig_split_factor)
+        self.assertEqual(result.full_tensor(), expected)
+
+        # select on dim 1: _StridedShard dim shifts from 2 to 1, split_factor preserved
+        result1 = dt_flat.select(1, 0)
+        expected1 = full.flatten(2, 3).select(1, 0)
+        self.assertIsInstance(result1.placements[0], _StridedShard)
+        self.assertEqual(result1.placements[0].dim, 1)
+        self.assertEqual(result1.placements[0].split_factor, orig_split_factor)
+        self.assertEqual(result1.full_tensor(), expected1)
+
+        # select on the _StridedShard dim: must unshard to Replicate
+        result2 = dt_flat.select(2, 0)
+        expected2 = full.flatten(2, 3).select(2, 0)
+        self.assertNotIsInstance(result2.placements[0], _StridedShard)
+        self.assertEqual(result2.full_tensor(), expected2)
+
+    def test_flatten_then_unbind(self):
+        """Verify _StridedShard correctness through unbind.
+
+        unbind removes a dim via shift_shard_dims_after_remove, which
+        must preserve split_factor when shifting _StridedShard dims.
+        """
+        mesh = init_device_mesh(self.device_type, (self.world_size,))
+        shape = (2, 3, 4, self.world_size * 2)
+        full = torch.randn(*shape, device=self.device_type)
+        dt = distribute_tensor(full, mesh, [Shard(3)])
+        dt_flat = dt.flatten(2, 3)  # (2, 3, 4*ws*2) with _StridedShard(dim=2)
+
+        self.assertIsInstance(dt_flat.placements[0], _StridedShard)
+        orig_split_factor = dt_flat.placements[0].split_factor
+
+        # unbind on dim 0: _StridedShard dim shifts from 2 to 1
+        results = torch.unbind(dt_flat, dim=0)
+        expected = torch.unbind(full.flatten(2, 3), dim=0)
+        for r, e in zip(results, expected):
+            self.assertIsInstance(r.placements[0], _StridedShard)
+            self.assertEqual(r.placements[0].dim, 1)
+            self.assertEqual(r.placements[0].split_factor, orig_split_factor)
+            self.assertEqual(r.full_tensor(), e)
+
+        # unbind on dim 1: _StridedShard dim shifts from 2 to 1
+        results1 = torch.unbind(dt_flat, dim=1)
+        expected1 = torch.unbind(full.flatten(2, 3), dim=1)
+        for r, e in zip(results1, expected1):
+            self.assertIsInstance(r.placements[0], _StridedShard)
+            self.assertEqual(r.placements[0].dim, 1)
+            self.assertEqual(r.placements[0].split_factor, orig_split_factor)
+            self.assertEqual(r.full_tensor(), e)
+
+        # unbind on dim above shard dim: _StridedShard dim stays unchanged
+        shape2 = (4, self.world_size * 2, 3, 5)
+        full2 = torch.randn(*shape2, device=self.device_type)
+        dt2 = distribute_tensor(full2, mesh, [Shard(1)])
+        dt2_flat = dt2.flatten(0, 1)  # (4*ws*2, 3, 5) with _StridedShard(dim=0)
+        self.assertIsInstance(dt2_flat.placements[0], _StridedShard)
+        self.assertEqual(dt2_flat.placements[0].dim, 0)
+        orig_split_factor2 = dt2_flat.placements[0].split_factor
+
+        results2 = torch.unbind(dt2_flat, dim=2)
+        expected2 = torch.unbind(full2.flatten(0, 1), dim=2)
+        for r, e in zip(results2, expected2):
+            self.assertIsInstance(r.placements[0], _StridedShard)
+            self.assertEqual(r.placements[0].dim, 0)
+            self.assertEqual(r.placements[0].split_factor, orig_split_factor2)
+            self.assertEqual(r.full_tensor(), e)
+
+    def test_flatten_then_stack(self):
+        """Verify _StridedShard correctness through stack.
+
+        stack inserts a new dim via shift_shard_dims_after_insert, which
+        must preserve split_factor when shifting _StridedShard dims.
+        """
+        mesh = init_device_mesh(self.device_type, (self.world_size,))
+        shape = (4, self.world_size * 2, 6)
+        full = torch.randn(*shape, device=self.device_type)
+        dt = distribute_tensor(full, mesh, [Shard(1)])
+        dt_flat = dt.flatten(0, 1)
+
+        self.assertIsInstance(dt_flat.placements[0], _StridedShard)
+        orig_split_factor = dt_flat.placements[0].split_factor
+
+        # stack along dim=0: _StridedShard dim shifts from 0 to 1
+        result = torch.stack([dt_flat, dt_flat], dim=0)
+        expected = torch.stack([full.flatten(0, 1), full.flatten(0, 1)], dim=0)
+        self.assertIsInstance(result.placements[0], _StridedShard)
+        self.assertEqual(result.placements[0].dim, 1)
+        self.assertEqual(result.placements[0].split_factor, orig_split_factor)
+        self.assertEqual(result.full_tensor(), expected)
+
+        # stack along dim=2 (after shard dim): _StridedShard dim stays at 0
+        result2 = torch.stack([dt_flat, dt_flat], dim=2)
+        expected2 = torch.stack([full.flatten(0, 1), full.flatten(0, 1)], dim=2)
+        self.assertIsInstance(result2.placements[0], _StridedShard)
+        self.assertEqual(result2.placements[0].dim, 0)
+        self.assertEqual(result2.placements[0].split_factor, orig_split_factor)
+        self.assertEqual(result2.full_tensor(), expected2)
+
+    def test_flatten_then_slice(self):
+        """Verify _StridedShard correctness through slice.
+
+        aten.slice.Tensor goes through gen_slice_strategy which uses
+        is_tensor_dim_sharded (misses _StridedShard). When the slice dim
+        matches the _StridedShard dim, the strategy should redistribute
+        to Replicate first, otherwise results are silently wrong.
+        """
+        mesh = init_device_mesh(self.device_type, (self.world_size,))
+        shape = (4, self.world_size * 2, 6)
+        full = torch.randn(*shape, device=self.device_type)
+        dt = distribute_tensor(full, mesh, [Shard(1)])
+        dt_flat = dt.flatten(0, 1)
+        flat_full = full.flatten(0, 1)
+
+        self.assertIsInstance(dt_flat.placements[0], _StridedShard)
+
+        # slice on the _StridedShard dim (dim 0)
+        result = dt_flat[: dt_flat.shape[0] // 2]
+        expected = flat_full[: flat_full.shape[0] // 2]
+        self.assertEqual(result.full_tensor(), expected)
+
+        # slice on non-shard dim
+        result2 = dt_flat[:, :3]
+        expected2 = flat_full[:, :3]
+        self.assertEqual(result2.full_tensor(), expected2)
+
+    def test_flatten_then_slice_scatter(self):
+        """Verify _StridedShard correctness through slice_scatter.
+
+        gen_slice_scatter_strategy uses replicate_tensor_dim which only
+        checks isinstance(p, Shard), missing _StridedShard. When the
+        scatter dim matches the _StridedShard dim and all strategies are
+        filtered, replicate_tensor_dim should replicate it but doesn't.
+        """
+        mesh = init_device_mesh(self.device_type, (self.world_size,))
+        shape = (4, self.world_size * 2, 6)
+        full = torch.randn(*shape, device=self.device_type)
+        dt = distribute_tensor(full, mesh, [Shard(1)])
+        dt_flat = dt.flatten(0, 1)
+        flat_full = full.flatten(0, 1)
+
+        self.assertIsInstance(dt_flat.placements[0], _StridedShard)
+
+        # slice_scatter on the _StridedShard dim (dim 0)
+        half = dt_flat.shape[0] // 2
+        src = dt_flat[:half] * 2
+        src_full = flat_full[:half] * 2
+        result = dt_flat.slice_scatter(src, dim=0, start=0, end=half)
+        expected = flat_full.slice_scatter(src_full, dim=0, start=0, end=half)
+        self.assertEqual(result.full_tensor(), expected)
+
+        # slice_scatter on non-shard dim
+        src2 = dt_flat[:, :3] * 2
+        src2_full = flat_full[:, :3] * 2
+        result2 = dt_flat.slice_scatter(src2, dim=1, start=0, end=3)
+        expected2 = flat_full.slice_scatter(src2_full, dim=1, start=0, end=3)
+        self.assertEqual(result2.full_tensor(), expected2)
+
+    def test_flatten_then_slice_backward(self):
+        """Verify _StridedShard correctness through slice backward.
+
+        slice_backward_rules only checks isinstance(p, Shard), missing
+        _StridedShard. When the slice dim matches the _StridedShard dim,
+        gradients are computed on the wrong local shards.
+        """
+        mesh = init_device_mesh(self.device_type, (self.world_size,))
+        shape = (4, self.world_size * 2, 6)
+        full = torch.randn(*shape, device=self.device_type)
+        full_req = full.clone().requires_grad_(True)
+
+        dt = distribute_tensor(full, mesh, [Shard(1)])
+        dt_flat = dt.flatten(0, 1).requires_grad_(True)
+
+        self.assertIsInstance(dt_flat.placements[0], _StridedShard)
+
+        flat_full = full_req.flatten(0, 1)
+        half = dt_flat.shape[0] // 2
+
+        # backward through slice on the _StridedShard dim
+        (dt_flat[:half].sum()).backward()
+        (flat_full[:half].sum()).backward()
+        self.assertEqual(dt_flat.grad.full_tensor(), full_req.grad.flatten(0, 1))
+
+    def test_flatten_then_cat_on_strided_shard_dim(self):
+        """Verify _StridedShard correctness through cat on the shard dim.
+
+        cat_strategy uses is_tensor_dim_sharded which calls is_shard(),
+        missing _StridedShard. When cat dim == _StridedShard dim, the
+        strategy won't unshard first, producing wrong results.
+        """
+        mesh = init_device_mesh(self.device_type, (self.world_size,))
+        shape = (4, self.world_size * 2, 6)
+        full = torch.randn(*shape, device=self.device_type)
+        dt = distribute_tensor(full, mesh, [Shard(1)])
+        dt_flat = dt.flatten(0, 1)
+        flat_full = full.flatten(0, 1)
+
+        self.assertIsInstance(dt_flat.placements[0], _StridedShard)
+        self.assertEqual(dt_flat.placements[0].dim, 0)
+
+        # cat on the _StridedShard dim (dim 0)
+        result = torch.cat([dt_flat, dt_flat], dim=0)
+        expected = torch.cat([flat_full, flat_full], dim=0)
+        self.assertEqual(result.full_tensor(), expected)
+
+        # cat on non-shard dim should work without redistribution
+        result2 = torch.cat([dt_flat, dt_flat], dim=1)
+        expected2 = torch.cat([flat_full, flat_full], dim=1)
+        self.assertEqual(result2.full_tensor(), expected2)
+
+        # _StridedShard on a higher dim (dim=1), cat on that dim
+        shape2 = (3, 5, self.world_size * 2)
+        full2 = torch.randn(*shape2, device=self.device_type)
+        dt2 = distribute_tensor(full2, mesh, [Shard(2)])
+        dt2_flat = dt2.flatten(1, 2)  # (3, 5*ws*2) with _StridedShard(dim=1)
+        flat_full2 = full2.flatten(1, 2)
+
+        self.assertIsInstance(dt2_flat.placements[0], _StridedShard)
+        self.assertEqual(dt2_flat.placements[0].dim, 1)
+
+        result3 = torch.cat([dt2_flat, dt2_flat], dim=1)
+        expected3 = torch.cat([flat_full2, flat_full2], dim=1)
+        self.assertEqual(result3.full_tensor(), expected3)
+
+    def test_flatten_then_split_on_strided_shard_dim(self):
+        """Verify _StridedShard correctness through split on the shard dim.
+
+        split_strategy uses is_tensor_dim_sharded which calls is_shard(),
+        missing _StridedShard. When split dim == _StridedShard dim, the
+        strategy won't unshard first, producing wrong results.
+        """
+        mesh = init_device_mesh(self.device_type, (self.world_size,))
+        shape = (4, self.world_size * 2, 6)
+        full = torch.randn(*shape, device=self.device_type)
+        dt = distribute_tensor(full, mesh, [Shard(1)])
+        dt_flat = dt.flatten(0, 1)
+        flat_full = full.flatten(0, 1)
+
+        self.assertIsInstance(dt_flat.placements[0], _StridedShard)
+        self.assertEqual(dt_flat.placements[0].dim, 0)
+
+        # split on the _StridedShard dim (dim 0)
+        half = dt_flat.shape[0] // 2
+        results = torch.split(dt_flat, half, dim=0)
+        expected = torch.split(flat_full, half, dim=0)
+        for r, e in zip(results, expected):
+            self.assertEqual(r.full_tensor(), e)
+
+        # split on non-shard dim
+        results2 = torch.split(dt_flat, 3, dim=1)
+        expected2 = torch.split(flat_full, 3, dim=1)
+        for r, e in zip(results2, expected2):
+            self.assertEqual(r.full_tensor(), e)
+
+        # _StridedShard on a higher dim (dim=1), split on that dim
+        shape2 = (3, 5, self.world_size * 2)
+        full2 = torch.randn(*shape2, device=self.device_type)
+        dt2 = distribute_tensor(full2, mesh, [Shard(2)])
+        dt2_flat = dt2.flatten(1, 2)  # (3, 5*ws*2) with _StridedShard(dim=1)
+        flat_full2 = full2.flatten(1, 2)
+
+        self.assertIsInstance(dt2_flat.placements[0], _StridedShard)
+        self.assertEqual(dt2_flat.placements[0].dim, 1)
+
+        half2 = dt2_flat.shape[1] // 2
+        results3 = torch.split(dt2_flat, half2, dim=1)
+        expected3 = torch.split(flat_full2, half2, dim=1)
+        for r, e in zip(results3, expected3):
+            self.assertEqual(r.full_tensor(), e)
+
+    def test_flatten_then_unbind_on_strided_shard_dim(self):
+        """Verify _StridedShard is detected by unbind on the shard dim.
+
+        gen_unbind_strategy uses is_tensor_dim_sharded which calls is_shard(),
+        missing _StridedShard. Unbinding on a _StridedShard dim should raise
+        RuntimeError (same as unbinding on a regular Shard dim).
+        """
+        mesh = init_device_mesh(self.device_type, (self.world_size,))
+        shape = (2, self.world_size * 2, 3)
+        full = torch.randn(*shape, device=self.device_type)
+        dt = distribute_tensor(full, mesh, [Shard(1)])
+        dt_flat = dt.flatten(0, 1)
+
+        self.assertIsInstance(dt_flat.placements[0], _StridedShard)
+        self.assertEqual(dt_flat.placements[0].dim, 0)
+
+        # unbind on the _StridedShard dim (dim 0) — should raise
+        with self.assertRaises(RuntimeError):
+            torch.unbind(dt_flat, dim=0)
+
+    def test_flatten_then_add_strided_shard_inputs(self):
+        """Verify _StridedShard is treated as shard in placement merge.
+
+        Binary ops (e.g. add) with two _StridedShard inputs go through
+        _derive_follow_placements_from_tuple_strategy → merge_placement.
+        merge_placement uses .is_shard() which misses _StridedShard,
+        causing it to fall through to the Replicate branch and triggering
+        an unnecessary all-gather before the elementwise op.
+        """
+        mesh = init_device_mesh(self.device_type, (self.world_size,))
+        shape = (4, self.world_size * 2, 6)
+        full_a = torch.randn(*shape, device=self.device_type)
+        full_b = torch.randn(*shape, device=self.device_type)
+
+        dt_a = distribute_tensor(full_a, mesh, [Shard(1)])
+        dt_b = distribute_tensor(full_b, mesh, [Shard(1)])
+        dt_a_flat = dt_a.flatten(0, 1)
+        dt_b_flat = dt_b.flatten(0, 1)
+
+        self.assertIsInstance(dt_a_flat.placements[0], _StridedShard)
+        self.assertIsInstance(dt_b_flat.placements[0], _StridedShard)
+
+        # add should work without redistribution
+        comm_mode = CommDebugMode()
+        with comm_mode:
+            result = dt_a_flat + dt_b_flat
+
+        expected = full_a.flatten(0, 1) + full_b.flatten(0, 1)
+        self.assertEqual(result.full_tensor(), expected)
+        # If merge_placement treats _StridedShard as Replicate, an
+        # all-gather would be emitted here.
+        self.assertEqual(comm_mode.get_total_counts(), 0)
+
+    def test_flatten_then_redistribute_backward_partial(self):
+        """Verify backward normalize handles _StridedShard → Partial.
+
+        In the backward pass of Redistribute, the code normalizes
+        Shard/Replicate → Partial to just Replicate (to avoid a useless
+        reduce). The check uses .is_shard() which misses _StridedShard,
+        so _StridedShard → Partial would keep the Partial placement
+        instead of normalizing to Replicate.
+        """
+        mesh = init_device_mesh(self.device_type, (self.world_size,))
+        shape = (4, self.world_size * 2, 6)
+        full = torch.randn(*shape, device=self.device_type)
+
+        dt = distribute_tensor(full, mesh, [Shard(1)])
+        dt_flat = dt.flatten(0, 1).requires_grad_(True)
+
+        self.assertIsInstance(dt_flat.placements[0], _StridedShard)
+
+        # redistribute _StridedShard → Replicate, then backward
+        dt_rep = dt_flat.redistribute(mesh, [Replicate()])
+        loss = dt_rep.sum()
+        loss.backward()
+
+        # grad should be all-ones with _StridedShard placement (same as input)
+        expected_grad = torch.ones_like(full.flatten(0, 1))
+        self.assertEqual(dt_flat.grad.full_tensor(), expected_grad)
+
+    def test_unpack_hook_tp_with_strided_shard(self):
+        """_unpack_hook_tp must redistribute _StridedShard to Replicate.
+
+        _unpack_hook_tp restores activations before recomputation in BWD.
+        If it only checks .is_shard() it misses _StridedShard, returning the
+        tensor still sharded and producing wrong gradients.
+        """
+        from torch.distributed.tensor.parallel.input_reshard import _unpack_hook_tp
+
+        mesh = init_device_mesh(self.device_type, (self.world_size,))
+        shape = (4, self.world_size * 2, 6)
+        full = torch.randn(*shape, device=self.device_type)
+        dt = distribute_tensor(full, mesh, [Shard(1)])
+        dt_flat = dt.flatten(0, 1)
+
+        self.assertIsInstance(dt_flat.placements[0], _StridedShard)
+
+        result = _unpack_hook_tp(mesh, 0, dt_flat)
+        self.assertIsInstance(result, DTensor)
+        self.assertTrue(result.placements[0].is_replicate())
+        self.assertEqual(result.full_tensor(), full.flatten(0, 1))
+
     def test_view_redistribution(self):
         """
         This test is added to demonstrate "incorrect" view ops behavior if redistribution happens.
@@ -645,7 +2464,6 @@ def test_view_redistribution(self):
         with self.assertRaisesRegex(RuntimeError, "Sharding propagation failed"):
             dtensor_x.view(-1, 8)
 
-    @with_comms
     def test_squeeze_(self):
         mesh_2d = init_device_mesh(self.device_type, (3, 2), mesh_dim_names=("a", "b"))
         self.init_manual_seed_for_rank()
@@ -664,7 +2482,153 @@ def test_squeeze_(self):
         )
         self.assertEqual(dist_x.placements, [Partial(), Shard(0)])
 
-    @with_comms
+        # squeeze_ should not trigger any communication
+        y = torch.randn((1, 4), device=self.device_type)
+        dist_y = DTensor.from_local(y, mesh_2d, [Partial(), Shard(1)])
+        with CommDebugMode() as comm_mode:
+            torch.ops.aten.squeeze_.dim(dist_y, 0)
+        self.assertEqual(comm_mode.get_total_counts(), 0)
+
+    def test_squeeze_variants(self):
+        """Test squeeze.default, squeeze.dim, and squeeze.dims with DTensor."""
+        mesh = init_device_mesh(self.device_type, (self.world_size,))
+
+        # squeeze.dims on sharded tensor - squeeze non-sharded dims
+        with self.subTest("dims_sharded"):
+            x = torch.randn(self.world_size, 1, 1, 8, device=self.device_type)
+            dt = distribute_tensor(x, mesh, [Shard(0)])
+            with CommDebugMode() as comm_mode:
+                result = dt.squeeze((1, 2))
+            self.assertEqual(comm_mode.get_total_counts(), 0)
+            self.assertEqual(result.shape, torch.Size([self.world_size, 8]))
+            self.assertEqual(result.placements, (Shard(0),))
+            self.assertEqual(result.to_local().shape, torch.Size([1, 8]))
+
+        # squeeze.dim on sharded tensor - squeeze non-sharded dim
+        with self.subTest("dim_sharded"):
+            x = torch.randn(self.world_size, 1, 8, device=self.device_type)
+            dt = distribute_tensor(x, mesh, [Shard(0)])
+            with CommDebugMode() as comm_mode:
+                result = dt.squeeze(1)
+            self.assertEqual(comm_mode.get_total_counts(), 0)
+            self.assertEqual(result.shape, torch.Size([self.world_size, 8]))
+            self.assertEqual(result.placements, (Shard(0),))
+            self.assertEqual(result.to_local().shape, torch.Size([1, 8]))
+
+        # squeeze.default on replicated tensor
+        with self.subTest("default_replicated"):
+            x = torch.randn(4, 1, 1, 8, device=self.device_type)
+            dt = distribute_tensor(x, mesh, [Replicate()])
+            with CommDebugMode() as comm_mode:
+                result = dt.squeeze()
+            self.assertEqual(comm_mode.get_total_counts(), 0)
+            self.assertEqual(result.shape, torch.Size([4, 8]))
+            self.assertEqual(result.placements, (Replicate(),))
+
+        # squeeze.dims on replicated tensor
+        with self.subTest("dims_replicated"):
+            x = torch.randn(2, 1, 3, 1, device=self.device_type)
+            dt = distribute_tensor(x, mesh, [Replicate()])
+            with CommDebugMode() as comm_mode:
+                result = dt.squeeze((1, 3))
+            self.assertEqual(comm_mode.get_total_counts(), 0)
+            self.assertEqual(result.shape, torch.Size([2, 3]))
+            self.assertEqual(result.placements, (Replicate(),))
+
+        # squeeze non-singleton dim is a no-op
+        with self.subTest("non_singleton_noop"):
+            x = torch.randn(self.world_size, 4, device=self.device_type)
+            dt = distribute_tensor(x, mesh, [Shard(0)])
+            with CommDebugMode() as comm_mode:
+                result = dt.squeeze(1)
+            self.assertEqual(comm_mode.get_total_counts(), 0)
+            self.assertEqual(result.shape, torch.Size([self.world_size, 4]))
+            self.assertEqual(result.placements, (Shard(0),))
+
+        # Partial passes through squeeze unchanged
+        with self.subTest("partial_max_passthrough"):
+            x = torch.randn(1, 4, device=self.device_type)
+            dt = DTensor.from_local(x, mesh, [Partial("max")])
+            with CommDebugMode() as comm_mode:
+                result = dt.squeeze(0)
+            self.assertEqual(comm_mode.get_total_counts(), 0)
+            self.assertEqual(result.shape, torch.Size([4]))
+            self.assertEqual(result.placements, (Partial("max"),))
+
+        # Partial("sum") also passes through
+        with self.subTest("partial_sum_passthrough"):
+            x = torch.randn(1, device=self.device_type)
+            dt = DTensor.from_local(x, mesh, [Partial("sum")])
+            with CommDebugMode() as comm_mode:
+                result = dt.squeeze()
+            self.assertEqual(comm_mode.get_total_counts(), 0)
+            self.assertEqual(result.shape, torch.Size([]))
+            self.assertEqual(result.placements, (Partial("sum"),))
+
+        # squeeze must not remove sharded dim (local size 1, global size > 1)
+        with self.subTest("preserve_sharded_dim_default"):
+            x = (
+                torch.arange(self.world_size * 8, device=self.device_type)
+                .reshape(self.world_size, 8)
+                .float()
+            )
+            dt = distribute_tensor(x, mesh, [Shard(0)])
+            with CommDebugMode() as comm_mode:
+                result = dt.squeeze()
+            self.assertEqual(comm_mode.get_total_counts(), 0)
+            self.assertEqual(result.shape, torch.Size([self.world_size, 8]))
+            self.assertEqual(result._local_tensor.shape, torch.Size([1, 8]))
+            self.assertEqual(result.placements, (Shard(0),))
+            self.assertEqual(result.full_tensor(), x)
+
+        # same as above but via squeeze.dim (single int arg)
+        with self.subTest("preserve_sharded_dim_explicit"):
+            x = (
+                torch.arange(self.world_size * 8, device=self.device_type)
+                .reshape(self.world_size, 8)
+                .float()
+            )
+            dt = distribute_tensor(x, mesh, [Shard(0)])
+            with CommDebugMode() as comm_mode:
+                result = dt.squeeze(0)
+            self.assertEqual(comm_mode.get_total_counts(), 0)
+            self.assertEqual(result.shape, torch.Size([self.world_size, 8]))
+            self.assertEqual(result._local_tensor.shape, torch.Size([1, 8]))
+            self.assertEqual(result.placements, (Shard(0),))
+            self.assertEqual(result.full_tensor(), x)
+
+        # squeeze.dims with mixed singleton/non-singleton dims
+        with self.subTest("mixed_dims"):
+            x = torch.randn(1, 4, 1, device=self.device_type)
+            dt = distribute_tensor(x, mesh, [Replicate()])
+            with CommDebugMode() as comm_mode:
+                result = dt.squeeze((0, 1, 2))
+            self.assertEqual(comm_mode.get_total_counts(), 0)
+            self.assertEqual(result.shape, torch.Size([4]))
+            self.assertEqual(result.placements, (Replicate(),))
+            self.assertEqual(result.full_tensor(), x.squeeze())
+
+        # sharded non-singleton dim preserved, shard index shifts
+        with self.subTest("shard_index_shift"):
+            x = torch.randn(1, self.world_size, 1, device=self.device_type)
+            dt = distribute_tensor(x, mesh, [Shard(1)])
+            with CommDebugMode() as comm_mode:
+                result = dt.squeeze((0, 1, 2))
+            self.assertEqual(comm_mode.get_total_counts(), 0)
+            self.assertEqual(result.shape, torch.Size([self.world_size]))
+            self.assertEqual(result._local_tensor.shape, torch.Size([1]))
+            self.assertEqual(result.placements, (Shard(0),))
+            self.assertEqual(result.full_tensor(), x.squeeze((0, 2)))
+
+        # S(0) on globally-singleton dim becomes R after squeeze (#174136)
+        with self.subTest("singleton_shard_becomes_replicate"):
+            x = torch.randn(1, 4, device=self.device_type)
+            dt = distribute_tensor(x, mesh, [Shard(0)])
+            result = dt.squeeze()
+            self.assertEqual(result.shape, torch.Size([4]))
+            self.assertEqual(result.placements, (Replicate(),))
+            self.assertEqual(result.full_tensor(), x.squeeze())
+
     def test_storage_offset_slice(self):
         """
         Test that storage_offset is properly tracked on DTensor when slicing
@@ -695,7 +2659,6 @@ def test_storage_offset_slice(self):
         expected = tensor[1:]
         self.assertEqual(sliced_dtensor.full_tensor(), expected)
 
-    @with_comms
     def test_storage_offset_shard_dim0_slice_dim1(self):
         """
         Test that storage_offset is properly tracked when tensor is sharded on dim 0
@@ -727,7 +2690,6 @@ def test_storage_offset_shard_dim0_slice_dim1(self):
         expected = tensor[:, 2:]
         self.assertEqual(sliced_dtensor.full_tensor(), expected)
 
-    @with_comms
     def test_storage_offset_shard_dim1_slice_dim0(self):
         """
         Test that storage_offset is properly tracked when tensor is sharded on dim 1
@@ -759,14 +2721,287 @@ def test_storage_offset_shard_dim1_slice_dim0(self):
         expected = tensor[2:, :]
         self.assertEqual(sliced_dtensor.full_tensor(), expected)
 
+    def test_view_groups_unbacked_symint(self):
+        from torch.fx.experimental.symbolic_shapes import ShapeEnv
+
+        shape_env = ShapeEnv()
+
+        def fresh_sym():
+            return shape_env.create_unbacked_symint()
+
+        # Same symbol forwards correctly
+        u0 = fresh_sym()
+        self.assertEqual(
+            view_groups([4, u0, 3], [4, u0, 3]),
+            (InputDim(0), InputDim(1), InputDim(2)),
+        )
+
+        # Same symbol with concrete split
+        u1 = fresh_sym()
+        self.assertEqual(
+            view_groups([u1, 12], [u1, 3, 4]),
+            (
+                InputDim(0),
+                Split(InputDim(1), (3, 4), 0),
+                Split(InputDim(1), (3, 4), 1),
+            ),
+        )
+
+        # Same symbol with concrete flatten
+        u2 = fresh_sym()
+        self.assertEqual(
+            view_groups([u2, 3, 4], [u2, 12]),
+            (InputDim(0), Flatten((InputDim(1), InputDim(2)))),
+        )
+
+        # Concrete split with trailing symbol
+        u3 = fresh_sym()
+        self.assertEqual(
+            view_groups([6, u3], [2, 3, u3]),
+            (
+                Split(InputDim(0), (2, 3), 0),
+                Split(InputDim(0), (2, 3), 1),
+                InputDim(1),
+            ),
+        )
+
+        # Singletons with symbolic dims
+        u4 = fresh_sym()
+        self.assertEqual(
+            view_groups([1, 1, u4, 3], [u4, 3]),
+            (InputDim(2), InputDim(3)),
+        )
+
+        # Flatten symbolic dims using same product expression
+        u5 = fresh_sym()
+        u6 = fresh_sym()
+        self.assertEqual(
+            view_groups([u5, u6], [u5 * u6]),
+            (Flatten((InputDim(0), InputDim(1))),),
+        )
+
+        # Partial fallback: concrete prefix resolves, symbolic suffix falls back
+        u7 = fresh_sym()
+        u8 = fresh_sym()
+        result = view_groups([4, u7, u8], [4, u7 * u8])
+        # dim 0 resolves as InputDim(0), dims 1-2 flatten
+        self.assertEqual(
+            result,
+            (InputDim(0), Flatten((InputDim(1), InputDim(2)))),
+        )
+
+        # Different unbacked symbols (1-to-1): fallback simplifies to InputDim
+        u9 = fresh_sym()
+        u10 = fresh_sym()
+        result = view_groups([4, u9], [4, u10])
+        self.assertEqual(result, (InputDim(0), InputDim(1)))
+
+    def test_view_groups_unbacked_sharding_propagation(self):
+        """Test that sharding is correctly propagated through view_groups with symbolic shapes."""
+        from torch.fx.experimental.symbolic_shapes import ShapeEnv
+
+        shape_env = ShapeEnv()
+
+        def fresh_sym():
+            return shape_env.create_unbacked_symint()
+
+        mesh_sizes = (2, 3)
+
+        # InputDim forwarding: Shard(1) on symbolic dim should propagate
+        u0 = fresh_sym()
+        from_shape = (4, u0, 6)
+        to_shape = (4, u0, 6)
+        rule = view_groups(from_shape, to_shape)
+        input_placements = [Shard(1), Replicate()]
+        inp_tgt, out_plc = propagate_shape_and_sharding(
+            input_placements, from_shape, rule, mesh_sizes
+        )
+        self.assertEqual(out_plc, [Shard(1), Replicate()])
+
+        # Same symbol with concrete split: Shard(0) on symbolic dim should forward
+        u1 = fresh_sym()
+        from_shape = (u1, 12)
+        to_shape = (u1, 3, 4)
+        rule = view_groups(from_shape, to_shape)
+        input_placements = [Shard(0), Replicate()]
+        inp_tgt, out_plc = propagate_shape_and_sharding(
+            input_placements, from_shape, rule, mesh_sizes
+        )
+        self.assertEqual(out_plc, [Shard(0), Replicate()])
+
+        # Shard on concrete dim that gets split: should propagate to split_id=0
+        u2 = fresh_sym()
+        from_shape = (6, u2)
+        to_shape = (2, 3, u2)
+        rule = view_groups(from_shape, to_shape)
+        input_placements = [Shard(0), Replicate()]
+        inp_tgt, out_plc = propagate_shape_and_sharding(
+            input_placements, from_shape, rule, mesh_sizes
+        )
+        self.assertEqual(out_plc, [Shard(0), Replicate()])
+
+        # Flatten [u, 16] -> [u*16] with Shard(0) on unbacked dim (leftmost)
+        u3 = fresh_sym()
+        for ms in mesh_sizes:  # needed for sharding
+            torch._check(u3 % ms == 0)
+        from_shape = (u3, 16)
+        to_shape = (u3 * 16,)
+        rule = view_groups(from_shape, to_shape)
+        input_placements = [Shard(0), Replicate()]
+        inp_tgt, out_plc = propagate_shape_and_sharding(
+            input_placements, from_shape, rule, mesh_sizes
+        )
+        self.assertEqual(out_plc, [Shard(0), Replicate()])
+
+        # Flatten [16, u] -> [16*u] with Shard(1) on unbacked dim (non-leftmost):
+        # should force replicate since non-leftmost dims can't propagate through flatten
+        u4 = fresh_sym()
+        from_shape = (16, u4)
+        to_shape = (16 * u4,)
+        rule = view_groups(from_shape, to_shape)
+        input_placements = [Shard(1), Replicate()]
+        inp_tgt, out_plc = propagate_shape_and_sharding(
+            input_placements, from_shape, rule, mesh_sizes
+        )
+        self.assertEqual(out_plc, [Replicate(), Replicate()])
+
+    def test_input_dim_rejects_int_comparison(self):
+        """InputDim.__eq__ should raise TypeError when compared with int.
+
+        Regression guard: shard.dim == in_dim (int == InputDim) silently
+        returned False for over 3 years, making downstream validation dead code.
+        """
+        dim = InputDim(0)
+        # InputDim == InputDim is fine
+        self.assertEqual(dim, InputDim(0))
+        self.assertNotEqual(dim, InputDim(1))
+        # hash contract: equal InputDims must have equal hashes
+        self.assertEqual(hash(dim), hash(InputDim(0)))
+        # hash is salted so it won't collide with raw int in dicts/sets
+        self.assertNotEqual(hash(dim), hash(0))
+        # InputDim == int raises (catches `dim == some_int`)
+        with self.assertRaisesRegex(TypeError, "Did you mean to use .input_dim"):
+            _ = dim == 0
+        # int == InputDim also raises (catches the original bug: `shard.dim == in_dim`)
+        with self.assertRaisesRegex(TypeError, "Did you mean to use .input_dim"):
+            _ = 0 == dim
+
+    def _assert_strided_shard_flag(self, dt, expected_flag):
+        """Assert use_strided_shard_as_shard_order matches expected_flag."""
+        self.assertEqual(dt._spec.use_strided_shard_as_shard_order, expected_flag)
+
+    def test_strided_shard_propagates_through_chained_ops(self):
+        """Verify use_strided_shard_as_shard_order=False propagates through
+        chained pointwise ops after flatten produces _StridedShard."""
+        mesh = init_device_mesh(self.device_type, (self.world_size,))
+        torch.manual_seed(42)
+        B, num_heads, S, head_dim = 2, 6, 4, 8
+        full = torch.randn(B, num_heads, S, head_dim, device=self.device_type)
+        dt = distribute_tensor(full, mesh, [Shard(1)])
+
+        # flatten(0,1): Shard(1) on non-first flatten dim -> _StridedShard(0)
+        flat = dt.flatten(0, 1)
+        self.assertIsInstance(flat.placements[0], _StridedShard)
+        self._assert_strided_shard_flag(flat, False)
+
+        # Chain several pointwise ops and verify flag propagates
+        activated = flat.relu()
+        self._assert_strided_shard_flag(activated, False)
+
+        added = activated + 1.0
+        self._assert_strided_shard_flag(added, False)
+
+        scaled = added * 0.5
+        self._assert_strided_shard_flag(scaled, False)
+
+        result = scaled.abs()
+        self._assert_strided_shard_flag(result, False)
+
+        # full_tensor triggers _StridedShard -> Replicate redistribution
+        expected = full.flatten(0, 1).relu().add(1.0).mul(0.5).abs()
+        self.assertEqual(result.full_tensor(), expected)
+
+    def test_strided_shard_propagates_through_reduction(self):
+        """Verify use_strided_shard_as_shard_order=False propagates through
+        reduction ops that don't reduce the _StridedShard dim."""
+        mesh = init_device_mesh(self.device_type, (self.world_size,))
+        torch.manual_seed(42)
+        # shape: (2, 6, 4, 8), shard on dim 1
+        B, num_heads, S, head_dim = 2, 6, 4, 8
+        full = torch.randn(B, num_heads, S, head_dim, device=self.device_type)
+        dt = distribute_tensor(full, mesh, [Shard(1)])
+
+        # flatten(0,1) -> (12, 4, 8) with _StridedShard(0)
+        flat = dt.flatten(0, 1)
+        self.assertIsInstance(flat.placements[0], _StridedShard)
+        self._assert_strided_shard_flag(flat, False)
+
+        # sum over dim -1 (head_dim), which is not the strided shard dim
+        reduced = flat.sum(dim=-1)
+        self._assert_strided_shard_flag(reduced, False)
+
+        expected = full.flatten(0, 1).sum(dim=-1)
+        self.assertEqual(reduced.full_tensor(), expected)
+
+    def test_strided_shard_propagates_through_matmul(self):
+        """Verify use_strided_shard_as_shard_order=False propagates when a
+        _StridedShard tensor participates in matmul."""
+        mesh = init_device_mesh(self.device_type, (self.world_size,))
+        torch.manual_seed(42)
+        B, num_heads, S, head_dim = 2, 6, 4, 8
+        full = torch.randn(B, num_heads, S, head_dim, device=self.device_type)
+        dt = distribute_tensor(full, mesh, [Shard(1)])
+
+        # flatten(0,1) -> (12, 4, 8) with _StridedShard(0)
+        flat = dt.flatten(0, 1)
+        self.assertIsInstance(flat.placements[0], _StridedShard)
+        self._assert_strided_shard_flag(flat, False)
+
+        # matmul with a replicated weight: (12, 4, 8) @ (8, 16) -> (12, 4, 16)
+        weight = torch.randn(head_dim, 16, device=self.device_type)
+        weight_dt = distribute_tensor(weight, mesh, [Replicate()])
+        result = torch.matmul(flat, weight_dt)
+        self._assert_strided_shard_flag(result, False)
+
+        expected = torch.matmul(full.flatten(0, 1), weight)
+        self.assertEqual(result.full_tensor(), expected)
+
+    def test_strided_shard_propagates_2d_mesh(self):
+        """Verify use_strided_shard_as_shard_order=False propagates on a 2D mesh."""
+        mesh = init_device_mesh(self.device_type, (self.world_size // 2, 2))
+        torch.manual_seed(42)
+        # shape: (3, 4, 8), shard dim 0 on mesh dim 0, replicate on mesh dim 1
+        full = torch.randn(3, 4, 8, device=self.device_type)
+        dt = distribute_tensor(full, mesh, [Shard(0), Replicate()])
+
+        # flatten(0,1) -> (12, 8); Shard(0) stays Shard(0), Replicate stays
+        flat = dt.flatten(0, 1)
+        self._assert_strided_shard_flag(flat, False)
+
+        result = flat.relu()
+        self._assert_strided_shard_flag(result, False)
+
+        expected = full.flatten(0, 1).relu()
+        self.assertEqual(result.full_tensor(), expected)
+
 
 TestViewOpsWithLocalTensor = create_local_tensor_test_class(
     TestViewOps,
     skipped_tests=[
         # Comparing data pointers is not supported for local tensor
         "test_dtensor_view_op_uneven",
+        # These tests use ShapeEnv directly, not local tensor tests
+        "test_view_groups_unbacked_symint",
+        "test_view_groups_unbacked_sharding_propagation",
+        # Too many test cases for LocalTensorMode dispatch overhead
+        "test_dtensor_flatten_1d",
+        "test_dtensor_flatten_2d",
+        "test_dtensor_flatten_multi_mesh",
+        "test_dtensor_flatten_split_multi_mesh",
     ],
+    base_class=LocalDTensorContinuousTestBase,
 )
 
+
 if __name__ == "__main__":
     run_tests()
diff --git a/test/distributed/tensor/test_xla_integration.py b/test/distributed/tensor/test_xla_integration.py
index 727b128b6b4c0..ff898f18e81f5 100644
--- a/test/distributed/tensor/test_xla_integration.py
+++ b/test/distributed/tensor/test_xla_integration.py
@@ -23,7 +23,8 @@
 
 # wrapper to check xla test requirements
 def with_xla(func: Callable) -> Callable:
-    assert func is not None
+    if func is None:
+        raise AssertionError("Expected func to not be None")
 
     @wraps(func)  # pyre-ignore[6]
     def wrapper(
@@ -74,7 +75,10 @@ def test_xla_distribute_tensor_1d_shard(self):
                     tensor_to_shard, device_mesh, shard_spec
                 )
                 # TODO(yeounoh) switch to DTensor API when XLAShardedTensor inherits DTensor
-                assert type(dist_tensor).__name__ == "XLAShardedTensor"
+                if type(dist_tensor).__name__ != "XLAShardedTensor":
+                    raise AssertionError(
+                        f"Expected XLAShardedTensor, got {type(dist_tensor).__name__}"
+                    )
                 global_tensor = dist_tensor.global_tensor  # type:ignore[attr-defined]
                 self.assertEqual(
                     global_tensor.size(), torch.Size([3 * device_count, 3])
@@ -99,7 +103,10 @@ def test_xla_distribute_tensor_1d_replicate(self):
             )
             dist_tensor = distribute_tensor(tensor_to_shard, device_mesh, shard_spec)
             # TODO(yeounoh) switch to DTensor API when XLAShardedTensor inherits DTensor
-            assert type(dist_tensor).__name__ == "XLAShardedTensor"
+            if type(dist_tensor).__name__ != "XLAShardedTensor":
+                raise AssertionError(
+                    f"Expected XLAShardedTensor, got {type(dist_tensor).__name__}"
+                )
             global_tensor = dist_tensor.global_tensor  # type:ignore[attr-defined]
             self.assertEqual(global_tensor.size(), torch.Size([3 * device_count, 3]))
             local_tensor = dist_tensor.local_shards[0].data
@@ -127,7 +134,10 @@ def test_xla_distribute_tensor_2d(self):
                     tensor_to_shard, device_mesh, shard_spec
                 )
                 # TODO(yeounoh) switch to DTensor API when XLAShardedTensor inherits DTensor
-                assert type(dist_tensor).__name__ == "XLAShardedTensor"
+                if type(dist_tensor).__name__ != "XLAShardedTensor":
+                    raise AssertionError(
+                        f"Expected XLAShardedTensor, got {type(dist_tensor).__name__}"
+                    )
                 global_tensor = dist_tensor.global_tensor  # type:ignore[attr-defined]
                 self.assertEqual(
                     global_tensor.size(), torch.Size([3 * device_count // 2, 3])
diff --git a/test/distributed/test_aten_comm_compute_reordering.py b/test/distributed/test_aten_comm_compute_reordering.py
index 6c1435288c298..4a66728138904 100644
--- a/test/distributed/test_aten_comm_compute_reordering.py
+++ b/test/distributed/test_aten_comm_compute_reordering.py
@@ -57,6 +57,8 @@ def apply_reordering_and_get_graph(graph, out_li) -> None:
         "insert_overlap_deps",
         "collective_estimator",
         "bucket_exposed_first",
+        "bucket_only_internode_comms",
+        "bucket_mode",
     )
     for key in config_keys:
         if (val := getattr(dist_opts, key)) is not None:
@@ -281,7 +283,10 @@ def test_graph():
         schedulable = {"wait_tensor_default", "wait_tensor_default_1"}
         for node in list(graph.nodes):
             expected = node.name in schedulable
-            assert _schedulable_wait_node(node) is expected
+            if _schedulable_wait_node(node) is not expected:
+                raise AssertionError(
+                    f"Expected _schedulable_wait_node({node.name}) is {expected}"
+                )
 
     @torch._inductor.config.patch(get_patches())
     def test_reorder_compute_for_overlap_mul(self):
@@ -510,6 +515,7 @@ def get_bucket_patches(compute_multiplier=1.0):
         "aten_distributed_optimizations.custom_runtime_estimation": estimate_aten_runtime_part,
         "aten_distributed_optimizations.collective_bucketing": True,
         "aten_distributed_optimizations.bucket_exposed_first": False,
+        "aten_distributed_optimizations.bucket_only_internode_comms": False,
         "reorder_for_locality": False,
         "triton.native_matmul": False,
         "reorder_for_compute_comm_overlap_passes": [],
@@ -942,7 +948,8 @@ def estimate_with_fake_mode(fx_node, compute_multiplier=1.0):
             with FakeTensorMode():
                 nonlocal estimation_calls
                 estimation_calls += 1
-                assert isinstance(torch.rand([20]), torch._subclasses.FakeTensor)
+                if not isinstance(torch.rand([20]), torch._subclasses.FakeTensor):
+                    raise AssertionError("Expected FakeTensor")
 
             return 1.0
 
@@ -1031,7 +1038,12 @@ def func(a):
                 self.assertTrue(same(out, correct))
 
     @unittest.skipIf(not HAS_GPU, "Inductor+gpu needs triton and recent GPU arch")
-    @torch._inductor.config.patch(get_bucket_patches())
+    @torch._inductor.config.patch(
+        {
+            **get_bucket_patches(),
+            "aten_distributed_optimizations.bucket_mode": "custom_ops_multidtype",
+        }
+    )
     def test_multidtype_bucketing(self):
         """Test that all_gathers with different dtypes get bucketed together."""
 
@@ -1161,7 +1173,12 @@ def func(a, b, *, ranks):
             self.assertTrue(same(out, correct))
 
     @unittest.skipIf(not HAS_GPU, "Inductor+gpu needs triton and recent GPU arch")
-    @torch._inductor.config.patch(get_bucket_patches())
+    @torch._inductor.config.patch(
+        {
+            **get_bucket_patches(),
+            "aten_distributed_optimizations.bucket_mode": "custom_ops_multidtype",
+        }
+    )
     def test_bucketing_with_convert_dtype(self):
         """Test that all_gathers with dtype conversion get bucketed and produce correct results."""
 
@@ -1169,6 +1186,7 @@ def func(a, b, c, d, *, ranks):
             # Convert inputs to float16 before all_gather
             a_fp16 = a.to(torch.float16)
             b_fp16 = b.to(torch.float16)
+            d_fp16 = d.to(torch.float16)
 
             # Two all_gathers with converted dtypes
             ag1 = _functional_collectives.all_gather_tensor(a_fp16, 0, ranks)
@@ -1176,7 +1194,7 @@ def func(a, b, c, d, *, ranks):
 
             # same dtype
             ag3 = _functional_collectives.all_gather_tensor(c, 0, ranks)
-            ag4 = _functional_collectives.all_gather_tensor(d, 0, ranks)
+            ag4 = _functional_collectives.all_gather_tensor(d_fp16, 0, ranks)
 
             return ag1, ag2, ag3, ag4
 
@@ -1210,6 +1228,116 @@ def func(a, b, c, d, *, ranks):
             correct = func(a, b, c, d, ranks=ranks)
             self.assertTrue(same(out, correct))
 
+    @unittest.skipIf(not HAS_GPU, "Inductor+gpu needs triton and recent GPU arch")
+    @torch._inductor.config.patch(
+        {
+            **get_bucket_patches(),
+            "aten_distributed_optimizations.enable_overlap_scheduling": True,
+        }
+    )
+    def test_uneven_sharding_spmd_graphs(self):
+        """Test that uneven DTensor sharding produces SPMD graphs across ranks.
+
+        When a tensor dimension is not divisible by world_size, DTensor pads
+        before all_gather and unpads after. The no-op pad/unpad on ranks with
+        full-size shards must not be eliminated by remove_noop_ops, so all
+        ranks produce identical FX graphs with matching op counts.
+        """
+
+        def func(a, *, ranks):
+            # Simulate DTensor's pad-before-all_gather for uneven shards.
+            # rank 0: a is (4, 8), pad_size=0 → no-op pad
+            # rank 1: a is (3, 8), pad_size=1 → real pad to (4, 8)
+            full_chunk = (7 + len(ranks) - 1) // len(ranks)
+            pad_size = full_chunk - a.size(0)
+            a_padded = torch.nn.functional.pad(a, [0, 0, 0, pad_size])
+            ag = _functional_collectives.all_gather_tensor(a_padded, 0, ranks)
+            # Unpad after all_gather: narrow to original logical size
+            result = ag.narrow(0, 0, 7)
+            return result + 1
+
+        with _dynamo_dist_per_rank_init(
+            self.rank,
+            self.world_size,
+            self.backend(device_type),
+            fake_pg=not at_least_x_gpu(2),
+        ):
+            import torch.distributed as dist
+            from torch._subclasses.fake_tensor import unset_fake_temporarily
+
+            world_size = self.world_size
+            # 7 is not divisible by 2: rank 0 gets 4 rows, rank 1 gets 3
+            full_chunk = (7 + world_size - 1) // world_size
+            local_size = full_chunk if self.rank == 0 else 7 - full_chunk
+            a = torch.randn(local_size, 8, device=device_type)
+            ranks = list(range(world_size))
+
+            func_c = functools.partial(func, ranks=ranks)
+            compiled = torch.compile(func_c)
+            out, aten_graph_str = run_and_get_aten_graph(compiled, a)
+
+            # Build structural fingerprint: sorted list of call_function targets.
+            # Node names differ across ranks, but targets and op counts must match.
+            targets_r = sorted(
+                str(n_line.split("target=")[1].split("]")[0])
+                for n_line in aten_graph_str.split("\n")
+                if "call_function" in n_line and "target=" in n_line
+            )
+
+            with unset_fake_temporarily():
+                all_targets: list[list[str] | None] = [None] * world_size
+                dist.all_gather_object(all_targets, targets_r)
+
+            self.assertEqual(
+                all_targets[0],
+                all_targets[1],
+                "FX graph op targets differ across ranks — not SPMD. "
+                "No-op pad/slice may have been eliminated by remove_noop_ops.",
+            )
+
+    @unittest.skipIf(not HAS_GPU, "Inductor+gpu needs triton and recent GPU arch")
+    @torch._inductor.config.patch(
+        {
+            **get_bucket_patches(),
+            "aten_distributed_optimizations.enable_overlap_scheduling": True,
+            "aten_distributed_optimizations.spmd_check": True,
+            "aten_distributed_optimizations.spmd_mismatch": "error",
+        }
+    )
+    def test_spmd_verify_crashes_on_mismatch(self):
+        """Test that spmd_mismatch="error" raises on non-SPMD graphs."""
+
+        def func(a, *, ranks):
+            # rank 0 gets (4, 8), rank 1 gets (3, 8) — different node counts
+            # if pad is eliminated on rank 0 but not rank 1
+            full_chunk = (7 + len(ranks) - 1) // len(ranks)
+            pad_size = full_chunk - a.size(0)
+            # Only rank 1 will have a real pad op here
+            if pad_size > 0:
+                a = torch.nn.functional.pad(a, [0, 0, 0, pad_size])
+            ag = _functional_collectives.all_gather_tensor(a, 0, ranks)
+            return ag + 1
+
+        with _dynamo_dist_per_rank_init(
+            self.rank,
+            self.world_size,
+            self.backend(device_type),
+            fake_pg=not at_least_x_gpu(2),
+        ):
+            world_size = self.world_size
+            full_chunk = (7 + world_size - 1) // world_size
+            local_size = full_chunk if self.rank == 0 else 7 - full_chunk
+            a = torch.randn(local_size, 8, device=device_type)
+            ranks = list(range(world_size))
+
+            func_c = functools.partial(func, ranks=ranks)
+            compiled = torch.compile(func_c)
+
+            # The graph will differ: rank 0 has no pad, rank 1 has pad.
+            # With spmd_check_crash_on_mismatch=True, this should raise.
+            with self.assertRaises(RuntimeError, msg="SPMD graph verification"):
+                compiled(a)
+
 
 def get_toy_model(device_type: str):
     """
@@ -1471,6 +1599,7 @@ def module_stack_fn(node):
             [
                 "l_func_self_modules_layers_modules_0_modules_wq_parameters_weight_",
                 "l_func_self_modules_layers_modules_0_modules_wq_parameters_bias_",
+                "l_args_0_",
                 "linear",
             ],
         )
@@ -1482,6 +1611,7 @@ def module_stack_fn(node):
             [
                 "l_func_self_modules_layers_modules_0_modules_wq_parameters_weight_",
                 "l_func_self_modules_layers_modules_0_modules_wq_parameters_bias_",
+                "l_args_0_",
                 "linear",
                 "l_func_self_modules_layers_modules_0_modules_proj_parameters_weight_",
                 "l_func_self_modules_layers_modules_0_modules_proj_parameters_bias_",
@@ -1502,11 +1632,15 @@ def module_stack_fn(node):
             [
                 "l_func_self_modules_layers_modules_0_modules_wq_parameters_weight_",
                 "l_func_self_modules_layers_modules_0_modules_wq_parameters_bias_",
+                "l_args_0_",
                 "linear",
             ],
         )
 
     @unittest.skipIf(not HAS_GPU, "Inductor+gpu needs triton and recent GPU arch")
+    @torch._inductor.config.patch(
+        {"aten_distributed_optimizations.bucket_mode": "custom_ops_multidtype"}
+    )
     def test_manual_reordering_bucketing_pass_all_gather_separate_buckets(self):
         self._run_manual_bucketing_test(
             collective_type="all_gather",
@@ -1539,6 +1673,9 @@ def test_manual_bucketing_reordering_pass_all_gather_no_bucket(self):
         )
 
     @unittest.skipIf(not HAS_GPU, "Inductor+gpu needs triton and recent GPU arch")
+    @torch._inductor.config.patch(
+        {"aten_distributed_optimizations.bucket_mode": "custom_ops_multidtype"}
+    )
     def test_manual_bucketing_reordering_pass_all_gather_single_bucket(self):
         self._run_manual_bucketing_test(
             collective_type="all_gather",
@@ -1551,6 +1688,9 @@ def test_manual_bucketing_reordering_pass_all_gather_single_bucket(self):
         )
 
     @unittest.skipIf(not HAS_GPU, "Inductor+gpu needs triton and recent GPU arch")
+    @torch._inductor.config.patch(
+        {"aten_distributed_optimizations.bucket_mode": "custom_ops_multidtype"}
+    )
     def test_bucketing_reordering_pass_all_gather_single_bucket_custom_module_stack_fn(
         self,
     ):
@@ -1666,6 +1806,9 @@ def func(a, b, c, d, *, ranks):
                 )
 
     @unittest.skipIf(not HAS_GPU, "Inductor+gpu needs triton and recent GPU arch")
+    @torch._inductor.config.patch(
+        {"aten_distributed_optimizations.bucket_mode": "custom_ops_multidtype"}
+    )
     def test_manual_bucketing_reordering_pass_reduce_scatter_separate_buckets(self):
         self._run_manual_bucketing_test(
             collective_type="reduce_scatter",
@@ -1681,6 +1824,9 @@ def test_manual_bucketing_reordering_pass_reduce_scatter_separate_buckets(self):
         )
 
     @unittest.skipIf(not HAS_GPU, "Inductor+gpu needs triton and recent GPU arch")
+    @torch._inductor.config.patch(
+        {"aten_distributed_optimizations.bucket_mode": "custom_ops_multidtype"}
+    )
     def test_manual_bucketing_reordering_pass_reduce_scatter_single_bucket(self):
         self._run_manual_bucketing_test(
             collective_type="reduce_scatter",
diff --git a/test/distributed/test_backends.py b/test/distributed/test_backends.py
index baf78bb62db1f..244a5197faff4 100644
--- a/test/distributed/test_backends.py
+++ b/test/distributed/test_backends.py
@@ -18,11 +18,20 @@ def test_device_to_backend_mapping(self, device) -> None:
         Test device to backend mapping
         """
         if "cuda" in device:
-            assert dist.get_default_backend_for_device(device) == "nccl"
+            if dist.get_default_backend_for_device(device) != "nccl":
+                raise AssertionError(
+                    f"Expected nccl, got {dist.get_default_backend_for_device(device)}"
+                )
         elif "cpu" in device:
-            assert dist.get_default_backend_for_device(device) == "gloo"
+            if dist.get_default_backend_for_device(device) != "gloo":
+                raise AssertionError(
+                    f"Expected gloo, got {dist.get_default_backend_for_device(device)}"
+                )
         elif "hpu" in device:
-            assert dist.get_default_backend_for_device(device) == "hccl"
+            if dist.get_default_backend_for_device(device) != "hccl":
+                raise AssertionError(
+                    f"Expected hccl, got {dist.get_default_backend_for_device(device)}"
+                )
         else:
             with self.assertRaises(ValueError):
                 dist.get_default_backend_for_device(device)
@@ -40,7 +49,8 @@ def test_create_pg(self, device) -> None:
         )
         pg = dist.distributed_c10d._get_default_group()
         backend_pg = pg._get_backend_name()
-        assert backend_pg == backend
+        if backend_pg != backend:
+            raise AssertionError(f"Expected {backend}, got {backend_pg}")
         dist.destroy_process_group()
 
 
diff --git a/test/distributed/test_c10d_common.py b/test/distributed/test_c10d_common.py
index 2eceeb1098003..36f17631c282a 100644
--- a/test/distributed/test_c10d_common.py
+++ b/test/distributed/test_c10d_common.py
@@ -3,6 +3,7 @@
 import copy
 import os
 import pickle
+import re
 import subprocess
 import sys
 import tempfile
@@ -13,8 +14,8 @@
 from dataclasses import dataclass
 from datetime import timedelta
 from itertools import product
+from pathlib import Path
 from sys import platform
-from typing import Optional
 
 import torch
 import torch.distributed as dist
@@ -994,9 +995,32 @@ def test_sync_batch_norm_empty_input(self):
 
         self._test_not_nan(model, x)
 
+    @skip_if_lt_x_gpu(2)
+    def test_ddp_buffer_sync_multi_forward_with_batchnorm(self):
+        pg = self._get_process_group()
+
+        model = nn.Sequential(
+            nn.Linear(10, 10),
+            nn.BatchNorm1d(10),
+        ).to(device=self.rank)
+
+        model = DistributedDataParallel(
+            model,
+            device_ids=[self.rank],
+            process_group=pg,
+            broadcast_buffers=True,
+        )
+
+        x = torch.randn(4, 10, device=self.rank)
+
+        model.zero_grad()
+        out1 = model(x)
+        out2 = model(x)
+        (out1.mean() + out2.mean()).backward()
+
     @dataclass
     class CustomOutput:
-        o1: Optional[torch.Tensor]
+        o1: torch.Tensor | None
         o2: dict[str, torch.Tensor]
 
     class DataclassOutputModule(nn.Module):
@@ -1154,7 +1178,11 @@ def _test_sequence_num_incremented(self, process_group, ranks):
             ranks=ranks,
             backend="gloo",
         )
-        assert dist.get_world_size(process_group) == dist.get_world_size(verify_pg)
+        if dist.get_world_size(process_group) != dist.get_world_size(verify_pg):
+            raise AssertionError(
+                f"World sizes mismatch: {dist.get_world_size(process_group)} "
+                f"vs {dist.get_world_size(verify_pg)}"
+            )
 
         initial_num = (
             self._verify_sequence_number_across_pg(
@@ -1520,6 +1548,42 @@ def _test_new_group_local_sync_duplicate_pg(self, backend):
             ]
             self.assertEqual(output_tensor_list, expected)
 
+    def _test_new_group_ordered(self, backend):
+        store = dist.FileStore(self.file_name, self.world_size)
+        dist.init_process_group(
+            backend,
+            world_size=self.world_size,
+            rank=self.rank,
+            store=store,
+        )
+        rank = dist.get_rank()
+
+        # Reverse-order ranks: group rank 0 = highest global rank
+        reversed_ranks = list(range(self.world_size - 1, -1, -1))
+        new_pg = dist.new_group(ranks=reversed_ranks, sort_ranks=False)
+
+        # Verify that the group rank assignment follows the user-provided order
+        expected_group_rank = reversed_ranks.index(rank)
+        self.assertEqual(dist.get_group_rank(new_pg, rank), expected_group_rank)
+        self.assertEqual(
+            dist.get_process_group_ranks(new_pg),
+            reversed_ranks,
+        )
+
+        # Verify that all_gather results follow the custom rank order
+        input_tensor = torch.tensor([rank], device=self.device)
+        output_tensor_list = [
+            torch.tensor([-1], device=self.device) for _ in range(self.world_size)
+        ]
+        dist.all_gather(output_tensor_list, input_tensor, group=new_pg)
+
+        # Group rank i holds global rank reversed_ranks[i]
+        expected = [
+            torch.tensor([reversed_ranks[i]], device=self.device)
+            for i in range(self.world_size)
+        ]
+        self.assertEqual(output_tensor_list, expected)
+
 
 class CommTest(AbstractCommTest, MultiProcessTestCase):
     def setUp(self):
@@ -2266,10 +2330,176 @@ def testNodeLocalRank(self):
         self.assertEqual(dist.get_node_local_rank(), self.rank)
 
 
+class RecordCommTest(TestCase):
+    def test_set_get(self):
+        self.assertEqual(torch._C._distributed_c10d._get_comm_profiling_name(), "")
+        with dist.record_comm("test_name"):
+            self.assertEqual(
+                torch._C._distributed_c10d._get_comm_profiling_name(), "test_name"
+            )
+        self.assertEqual(torch._C._distributed_c10d._get_comm_profiling_name(), "")
+
+    def test_nesting(self):
+        with dist.record_comm("outer"):
+            self.assertEqual(
+                torch._C._distributed_c10d._get_comm_profiling_name(), "outer"
+            )
+            with dist.record_comm("inner"):
+                self.assertEqual(
+                    torch._C._distributed_c10d._get_comm_profiling_name(), "inner"
+                )
+            self.assertEqual(
+                torch._C._distributed_c10d._get_comm_profiling_name(), "outer"
+            )
+        self.assertEqual(torch._C._distributed_c10d._get_comm_profiling_name(), "")
+
+
+# Trivially-destructible types safe for thread_local.
+# Non-trivial destructors (std::string, std::vector, std::map, etc.)
+# register via __cxa_thread_atexit_impl which holds a mutex — if fork()
+# happens while that mutex is held, the child deadlocks.
+_SAFE_TLS_TYPE_RE = re.compile(
+    r"""
+    (?:
+        .*\*\s*$              # any pointer type (T*, const T*, etc.)
+        | bool
+        | u?int(?:8|16|32|64)_t
+        | int
+        | size_t
+        | ssize_t
+        | ptrdiff_t
+        | float
+        | double
+        | char
+        | (?:unsigned|signed)\s+(?:int|long(?:\s+long)?|short|char)
+        | unsigned
+        | signed
+        | long(?:\s+long)?
+        | short
+    )
+    """,
+    re.VERBOSE,
+)
+
+# Regex to extract a thread_local declaration.
+# Captures the type portion (group 1) and the variable name (group 2).
+_TLS_DECL_RE = re.compile(
+    r"""
+    (?:^|(?<=\s))               # start of line or preceded by whitespace
+    (?:static\s+)?              # optional leading static
+    thread_local                # the keyword
+    \s+(?:static\s+)?           # optional trailing static
+    (.+?)                       # type (group 1) — non-greedy
+    \s+((?:\w+(?:<[^>]*>)?::)*\w+)  # variable name, possibly qualified (group 2)
+    \s*[=;(]                    # followed by = or ; or (
+    """,
+    re.VERBOSE,
+)
+
+
+class ThreadLocalSafetyLintTest(TestCase):
+    """Lint: every thread_local in c10d must use a trivially-destructible type.
+
+    Non-trivially-destructible thread_local variables (std::string,
+    std::vector, std::map, c10::intrusive_ptr, ...) register destructors
+    via __cxa_thread_atexit_impl, which holds an internal glibc mutex.
+    If fork() happens while that mutex is held the child process inherits
+    a locked mutex and deadlocks on the next TLS access.
+
+    Use raw pointers (T*) with lazy heap-allocation instead.
+    """
+
+    # Pre-existing violations outside this diff's scope.  Each entry is
+    # (relative_path, variable_name).  Remove entries as they get fixed.
+    _KNOWN_VIOLATIONS = {
+        ("ProcessGroup.cpp", "pg"),
+        ("cuda/CUDAEventCache.cpp", "cacheDeviceMap"),
+    }
+
+    @staticmethod
+    def _c10d_src_dir() -> Path:
+        # test file:  caffe2/test/distributed/test_c10d_common.py
+        # source dir: caffe2/torch/csrc/distributed/c10d/
+        return (
+            Path(__file__).resolve().parents[2]
+            / "torch"
+            / "csrc"
+            / "distributed"
+            / "c10d"
+        )
+
+    @staticmethod
+    def _is_safe_type(type_str: str) -> bool:
+        """Return True if *type_str* is trivially destructible."""
+        cleaned = type_str.strip()
+        # Remove const / volatile / mutable qualifiers for matching
+        cleaned = re.sub(r"\b(const|volatile|mutable|inline)\b", "", cleaned).strip()
+        # Pointer types are always safe
+        if cleaned.endswith("*"):
+            return True
+        # Check against known safe scalar types
+        return bool(_SAFE_TLS_TYPE_RE.fullmatch(cleaned))
+
+    def test_no_non_trivial_thread_locals(self):
+        """Scan c10d sources for thread_local with non-trivial destructors."""
+        c10d_dir = self._c10d_src_dir()
+        self.assertTrue(
+            c10d_dir.is_dir(),
+            f"c10d source directory not found: {c10d_dir}",
+        )
+
+        violations = []
+        for ext in ("*.cpp", "*.hpp", "*.h", "*.cc"):
+            for filepath in c10d_dir.rglob(ext):
+                lines = filepath.read_text().splitlines()
+                idx = 0
+                while idx < len(lines):
+                    stripped = lines[idx].strip()
+                    idx += 1
+                    # Skip pure comments and static_assert guards
+                    if (
+                        stripped.startswith(("//", "/*"))
+                        or "static_assert" in stripped
+                        or "thread_local" not in stripped
+                    ):
+                        continue
+                    # Try single-line match first
+                    m = _TLS_DECL_RE.search(stripped)
+                    if not m and idx < len(lines):
+                        # Multi-line declaration (type on one line,
+                        # variable name on the next)
+                        combined = stripped + " " + lines[idx].strip()
+                        m = _TLS_DECL_RE.search(combined)
+                    if not m:
+                        continue
+                    tls_type, var_name = m.group(1), m.group(2)
+                    if self._is_safe_type(tls_type):
+                        continue
+                    rel = str(filepath.relative_to(c10d_dir))
+                    # Strip class qualifiers from var_name for allowlist
+                    # e.g. "Foo<T>::bar" -> "bar"
+                    bare_name = var_name.rsplit("::", 1)[-1]
+                    if (rel, bare_name) in self._KNOWN_VIOLATIONS:
+                        continue
+                    violations.append(
+                        f"  {rel}:{idx}: thread_local {tls_type} {var_name}"
+                    )
+
+        self.assertEqual(
+            violations,
+            [],
+            "Non-trivially-destructible thread_local variable(s) found in c10d.\n"
+            "These cause fork-deadlocks via __cxa_thread_atexit.\n"
+            "Use a raw pointer (T*) with lazy heap-allocation instead.\n\n"
+            "Violations:\n" + "\n".join(violations),
+        )
+
+
 if __name__ == "__main__":
     if device_type != "cpu":
-        assert not torch.get_device_module()._initialized, (
-            f"test_distributed must not have initialized {device_type} context on main process"
-        )
+        if torch.get_device_module()._initialized:
+            raise AssertionError(
+                f"test_distributed must not have initialized {device_type} context on main process"
+            )
 
     run_tests()
diff --git a/test/distributed/test_c10d_functional_native.py b/test/distributed/test_c10d_functional_native.py
index 7e04ad8a2bb2c..bb71952809ffe 100644
--- a/test/distributed/test_c10d_functional_native.py
+++ b/test/distributed/test_c10d_functional_native.py
@@ -4,7 +4,6 @@
 import threading
 import unittest
 from datetime import timedelta
-from typing import Optional
 
 import torch
 import torch.distributed as dist
@@ -258,6 +257,35 @@ def test_functional_collectives_inference_mode(self) -> None:
                 ),
             )
 
+    @skip_if_lt_x_gpu(2)
+    def test_functional_collectives_batched(self) -> None:
+        self._init_process_group()
+
+        def f(x):
+            if self.rank not in [0, 1]:
+                return
+            tensor = self.rank * x
+            return funcol.wait_tensor(
+                funcol.batch_p2p_ops_inplace(
+                    op_list=["isend" if self.rank == 0 else "irecv"],
+                    peer_list=[1 if self.rank == 0 else 0],
+                    tensors=[tensor],
+                    tag_list=[0],
+                    group_name="",
+                )[0]
+            )
+
+        with torch.inference_mode():
+            input = torch.ones((10), device=self.device)
+            output = f(input)
+            if self.rank == 0:
+                self.assertEqual(
+                    output,
+                    torch.empty(0, device=self.device, dtype=input.dtype),
+                )
+            else:
+                self.assertEqual(output, 0 * input)
+
     @unittest.skipIf(not HAS_GPU, "Inductor+gpu needs triton and recent GPU arch")
     @skip_if_lt_x_gpu(2)
     # https://github.com/pytorch/pytorch/issues/126338
@@ -636,7 +664,7 @@ def __init__(self, pg: "ProcessGroupDummy") -> None:
         super().__init__()
         self.pg = pg
 
-    def wait(self, timeout: Optional[timedelta] = None) -> bool:
+    def wait(self, timeout: timedelta | None = None) -> bool:
         self.pg.waits += 1
         return True
 
@@ -722,7 +750,7 @@ def test_wait_tensor_cross_thread(self) -> None:
         than where the collective was registered.
         """
         wait_called = False
-        exception_in_thread: Optional[BaseException] = None
+        exception_in_thread: BaseException | None = None
 
         class MyWork(dist.Work):
             def wait(self, _=None):
@@ -830,6 +858,48 @@ def test_collectives(self) -> None:
         self.assertEqual(pg.dels, 4)
 
 
+class ProcessGroupArgTest(TestCase):
+    """
+    Test that _c10d_functional ops accept ProcessGroup objects directly
+    (not just string group names) via the Any-typed group argument.
+    """
+
+    def setUp(self):
+        super().setUp()
+        dummy_init_pg()
+        self.pg = ProcessGroupDummy().register()
+
+    def test_all_reduce(self) -> None:
+        x = torch.rand(2, 2)
+        out = torch.ops._c10d_functional.all_reduce(x, "sum", self.pg)
+        torch.ops._c10d_functional.wait_tensor(out)
+
+    def test_all_reduce_(self) -> None:
+        x = torch.rand(2, 2)
+        out = torch.ops._c10d_functional.all_reduce_(x, "sum", self.pg)
+        torch.ops._c10d_functional.wait_tensor(out)
+
+    def test_all_gather_into_tensor(self) -> None:
+        x = torch.rand(2, 2)
+        out = torch.ops._c10d_functional.all_gather_into_tensor(x, 1, self.pg)
+        torch.ops._c10d_functional.wait_tensor(out)
+
+    def test_reduce_scatter_tensor(self) -> None:
+        x = torch.rand(2, 2)
+        out = torch.ops._c10d_functional.reduce_scatter_tensor(x, "sum", 1, self.pg)
+        torch.ops._c10d_functional.wait_tensor(out)
+
+    def test_broadcast(self) -> None:
+        x = torch.rand(2, 2)
+        out = torch.ops._c10d_functional.broadcast(x, 0, self.pg)
+        torch.ops._c10d_functional.wait_tensor(out)
+
+    def test_broadcast_(self) -> None:
+        x = torch.rand(2, 2)
+        out = torch.ops._c10d_functional.broadcast_(x, 0, self.pg)
+        torch.ops._c10d_functional.wait_tensor(out)
+
+
 def find_buffer_assignments(code):
     pattern = r"buf(\d+) = empty_strided_"
     matches = re.finditer(pattern, code)
@@ -1376,5 +1446,138 @@ def func(arg: torch.Tensor) -> torch.Tensor:
         (FileCheck().check("all_reduce_.default(buf0, 'avg', '0')").run(code))
 
 
+class ACTCompileTest(TestCase):
+    """
+    Test that AsyncCollectiveTensor inputs to compiled regions are resolved
+    before AOT autograd tracing begins. ACTs are transient eager-mode wrappers
+    for async collective overlap; if they leak into the traced graph as input
+    types, AOT autograd records them in tangent metadata and then hits a type
+    mismatch at runtime because autograd produces plain-tensor tangents.
+    """
+
+    def test_act_compile_backward_tangent_mismatch(self):
+        """
+        When a bare AsyncCollectiveTensor (ACT) enters a torch.compiled
+        region and passes through a view op, the output remains an ACT.
+        If the ACT is not unwrapped before tracing, AOT autograd
+        fakifies it as-is and records it in SubclassCreationMeta. At runtime, autograd
+        produces a plain tensor tangent, causing:
+            RuntimeError: Expected a AsyncCollectiveTensor tangent
+            but got a plain Tensor.
+
+        This occurs in practice when TP async collectives produce a
+        DTensor(ACT), an eager caller does to_local() to get a bare ACT,
+        and that ACT flows into a compiled sub-module whose forward
+        contains view ops (unsqueeze, reshape, transpose, etc.).
+
+        We use unsqueeze (a view op that preserves ACT) to ensure the
+        compiled function's output carries ACT type, which is what
+        triggers the tangent metadata recording.
+        """
+        from unittest.mock import patch
+
+        elem = torch.randn(4, 4, requires_grad=True)
+        act = AsyncCollectiveTensor(elem)
+
+        compiled_fn = torch.compile(lambda x: x.unsqueeze(0), backend="aot_eager")
+
+        original_trigger_wait = AsyncCollectiveTensor.trigger_wait
+        wait_called = False
+
+        def tracked_trigger_wait(self):
+            nonlocal wait_called
+            wait_called = True
+            return original_trigger_wait(self)
+
+        with patch.object(AsyncCollectiveTensor, "trigger_wait", tracked_trigger_wait):
+            out = compiled_fn(act)
+            # Without ACT unwrapping, this raises:
+            #   RuntimeError: Expected a AsyncCollectiveTensor tangent
+            #   but got a plain Tensor.
+            out.sum().backward()
+
+        self.assertTrue(wait_called, "trigger_wait() was never called")
+
+        # Verify numerics: trigger_wait() must fire so the correct data
+        # flows through. ACT wraps elem, so results must match elem directly.
+        ref = elem.detach().unsqueeze(0)
+        self.assertEqual(out.detach(), ref)
+
+    def test_act_runtime_unwrap(self):
+        """
+        Verify that ACT inputs are resolved before the compiled graph
+        executes.
+
+        With inductor, triton kernels bypass __torch_dispatch__ entirely,
+        so an un-waited ACT would silently feed stale data to the kernel.
+        The aot_eager backend masks this because it runs graph ops through
+        __torch_dispatch__ (which calls trigger_wait). This test uses a
+        custom backend that asserts inputs are plain tensors to catch the
+        problem regardless of backend.
+        """
+
+        def assert_no_act_backend(gm, example_inputs):
+            from torch._dynamo.backends.common import aot_autograd
+
+            def inner_compiler(gm, example_inputs):
+                def compiled(args):
+                    for i, a in enumerate(args):
+                        self.assertNotIsInstance(
+                            a,
+                            AsyncCollectiveTensor,
+                            f"arg {i} is still an ACT — trigger_wait() "
+                            "was not called before the compiled function",
+                        )
+                    return gm(*args)
+
+                compiled._boxed_call = True
+                return compiled
+
+            return aot_autograd(fw_compiler=inner_compiler)(gm, example_inputs)
+
+        compiled_fn = torch.compile(lambda x: x * 2, backend=assert_no_act_backend)
+
+        # First call: triggers compilation.
+        elem1 = torch.randn(4, 4)
+        act1 = AsyncCollectiveTensor(elem1)
+        r1 = compiled_fn(act1)
+        self.assertEqual(r1, elem1 * 2)
+
+        # Second call: dynamo reuses cached graph. process_inputs does
+        # NOT run again — the runtime wrapper must unwrap ACTs.
+        elem2 = torch.randn(4, 4)
+        act2 = AsyncCollectiveTensor(elem2)
+        r2 = compiled_fn(act2)
+        self.assertEqual(r2, elem2 * 2)
+
+    def test_act_guard_recompiles(self):
+        """
+        Dynamo must recompile when an input switches between plain tensor
+        and AsyncCollectiveTensor (or vice versa).
+        """
+        cnt = torch._dynamo.testing.CompileCounter()
+        compiled_fn = torch.compile(lambda x: x * 2, backend=cnt)
+
+        elem = torch.randn(4, 4)
+
+        # Call 1: plain tensor — triggers first compilation.
+        r1 = compiled_fn(elem)
+        self.assertEqual(r1, elem * 2)
+        self.assertEqual(cnt.frame_count, 1)
+
+        # Call 2: ACT — should trigger a recompile because the input
+        # type changed from Tensor to AsyncCollectiveTensor.
+        act = AsyncCollectiveTensor(elem)
+        r2 = compiled_fn(act)
+        self.assertEqual(r2, elem * 2)
+        self.assertEqual(cnt.frame_count, 2)
+
+        # Call 3: plain tensor again — should reuse the first compiled
+        # graph, no new compilation.
+        r3 = compiled_fn(elem)
+        self.assertEqual(r3, elem * 2)
+        self.assertEqual(cnt.frame_count, 2)
+
+
 if __name__ == "__main__":
     run_tests()
diff --git a/test/distributed/test_c10d_gloo.py b/test/distributed/test_c10d_gloo.py
index 2031f86534faf..ae3db5e21d617 100644
--- a/test/distributed/test_c10d_gloo.py
+++ b/test/distributed/test_c10d_gloo.py
@@ -1791,6 +1791,226 @@ def test_send_recv_complex(self):
             pg.recv([recv_tensor], 0, 0).wait()
             self.assertEqual(send_tensor, recv_tensor)
 
+    @requires_gloo()
+    def test_alltoall_checks(self):
+        store = c10d.FileStore(self.file_name, self.world_size)
+        pg = self._create_process_group_gloo(
+            store, self.rank, self.world_size, self.opts()
+        )
+
+        t1 = torch.zeros([1], dtype=torch.float32)
+        t2 = torch.zeros([1], dtype=torch.float64)
+        t3 = torch.zeros([2], dtype=torch.float32)
+
+        # Check input tensor list size does not match world size
+        with self.assertRaisesRegex(
+            RuntimeError, "input tensor list size.*does not match world size"
+        ):
+            pg.alltoall([t1] * self.world_size, [t1] * (self.world_size - 1))
+
+        # Check output tensor list size does not match world size
+        with self.assertRaisesRegex(
+            RuntimeError, "output tensor list size.*does not match world size"
+        ):
+            pg.alltoall([t1] * (self.world_size - 1), [t1] * self.world_size)
+
+        # Check invalid tensor type in inputs
+        with self.assertRaisesRegex(RuntimeError, "invalid tensor type"):
+            pg.alltoall([t1] * self.world_size, [t1, t2] + [t1] * (self.world_size - 2))
+
+        # Check invalid tensor size in inputs
+        with self.assertRaisesRegex(RuntimeError, "invalid tensor size"):
+            pg.alltoall([t1] * self.world_size, [t1, t3] + [t1] * (self.world_size - 2))
+
+        # Check invalid tensor type in outputs
+        with self.assertRaisesRegex(RuntimeError, "invalid tensor type"):
+            pg.alltoall([t1, t2] + [t1] * (self.world_size - 2), [t1] * self.world_size)
+
+        # Check invalid tensor size in outputs
+        with self.assertRaisesRegex(RuntimeError, "invalid tensor size"):
+            pg.alltoall([t1, t3] + [t1] * (self.world_size - 2), [t1] * self.world_size)
+
+        # Check input and output tensors must have the same type
+        with self.assertRaisesRegex(
+            RuntimeError, "input and output tensors must have the same type"
+        ):
+            pg.alltoall([t2] * self.world_size, [t1] * self.world_size)
+
+        # Check input and output tensors must have the same size
+        with self.assertRaisesRegex(
+            RuntimeError, "input and output tensors must have the same size"
+        ):
+            pg.alltoall([t3] * self.world_size, [t1] * self.world_size)
+
+    def _test_alltoall_basics(self, fn):
+        store = c10d.FileStore(self.file_name, self.world_size)
+        pg = self._create_process_group_gloo(
+            store, self.rank, self.world_size, self.opts()
+        )
+
+        # Preallocate tensors for input/output
+        # Each rank sends tensor with value equal to rank to each other rank
+        input_tensors = [fn(torch.tensor([self.rank])) for _ in range(self.world_size)]
+        output_tensors = [fn(torch.tensor([-1])) for _ in range(self.world_size)]
+
+        # Perform all-to-all
+        fut = pg.alltoall(output_tensors, input_tensors).get_future()
+        fut.wait()
+
+        # After alltoall, rank i should have received tensor with value j from rank j
+        # in output_tensors[j]
+        for i in range(self.world_size):
+            self.assertEqual(torch.tensor([i]), output_tensors[i])
+
+    @requires_gloo()
+    def test_alltoall_basics(self):
+        self._test_alltoall_basics(lambda t: t.clone())
+
+    @skip_if_lt_x_gpu(2)
+    @requires_gloo()
+    def test_alltoall_basics_cuda(self):
+        self._test_alltoall_basics(lambda t: t.clone().cuda())
+
+    def _test_alltoall_stress(self, inputs, fn):
+        store = c10d.FileStore(self.file_name, self.world_size)
+        pg = self._create_process_group_gloo(
+            store, self.rank, self.world_size, self.opts(threads=8)
+        )
+
+        outputs = [
+            [fn(torch.tensor([-1])) for _ in range(self.world_size)]
+            for _ in range(len(inputs))
+        ]
+
+        future_handles = []
+        for i in range(len(inputs)):
+            input_tensors = [fn(inputs[i][j]) for j in range(self.world_size)]
+            fut = pg.alltoall(outputs[i], input_tensors).get_future()
+            future_handles.append(fut)
+
+        for i, future_handle in enumerate(future_handles):
+            future_handle.wait()
+
+            # Verify correctness: output_tensors[j] should contain the value
+            # that rank j sent to this rank (self.rank) at position self.rank in rank j's input.
+            # All ranks construct inputs[i] = [tensor([i * world_size + 0]), ..., tensor([i * world_size + world_size-1])]
+            # So rank j sends inputs[i][self.rank] = tensor([i * world_size + self.rank]) to this rank.
+            # This rank receives that in output_tensors[j].
+            for j in range(self.world_size):
+                expected_value = torch.tensor([i * self.world_size + self.rank])
+                self.assertEqual(
+                    expected_value,
+                    outputs[i][j],
+                    msg=(f"Mismatch in iteration {i:d} from rank {j:d}"),
+                )
+
+    @requires_gloo()
+    def test_alltoall_stress(self):
+        inputs = [
+            [torch.tensor([i * self.world_size + j]) for j in range(self.world_size)]
+            for i in range(1000)
+        ]
+        self._test_alltoall_stress(inputs, lambda t: t.clone())
+
+    @skip_if_lt_x_gpu(2)
+    @requires_gloo()
+    @skipIfRocm
+    def test_alltoall_stress_cuda(self):
+        inputs = [
+            [torch.tensor([i * self.world_size + j]) for j in range(self.world_size)]
+            for i in range(1000)
+        ]
+        self._test_alltoall_stress(inputs, lambda t: t.clone().cuda())
+
+    def _test_alltoall_data_routing(self, fn):
+        """
+        Test that data routing is correct: rank i sends inputTensors[j]
+        to rank j
+        """
+        store = c10d.FileStore(self.file_name, self.world_size)
+        pg = self._create_process_group_gloo(
+            store, self.rank, self.world_size, self.opts()
+        )
+
+        # Each rank sends unique values to every other rank
+        # Rank i sends tensor with value (i * 100 + j) to rank j
+        input_tensors = []
+        for j in range(self.world_size):
+            value = self.rank * 100 + j
+            input_tensors.append(fn(torch.tensor([value], dtype=torch.float32)))
+
+        output_tensors = [
+            fn(torch.tensor([-1], dtype=torch.float32)) for _ in range(self.world_size)
+        ]
+
+        # Perform all-to-all
+        fut = pg.alltoall(output_tensors, input_tensors).get_future()
+        fut.wait()
+
+        # Verify: rank i should receive tensor with value (j * 100 + i)
+        # from rank j in output_tensors[j]
+        for j in range(self.world_size):
+            expected_value = j * 100 + self.rank
+            actual_value = output_tensors[j].item()
+            self.assertEqual(
+                expected_value,
+                actual_value,
+                msg=(
+                    f"Rank {self.rank}: output_tensors[{j}] = "
+                    f"{actual_value}, expected {expected_value}"
+                ),
+            )
+
+    @requires_gloo()
+    def test_alltoall_data_routing(self):
+        self._test_alltoall_data_routing(lambda t: t.clone())
+
+    @skip_if_lt_x_gpu(2)
+    @requires_gloo()
+    def test_alltoall_data_routing_cuda(self):
+        self._test_alltoall_data_routing(lambda t: t.clone().cuda())
+
+    def _test_alltoall_multidim(self, fn):
+        """Test alltoall with multi-dimensional tensors."""
+        store = c10d.FileStore(self.file_name, self.world_size)
+        pg = self._create_process_group_gloo(
+            store, self.rank, self.world_size, self.opts()
+        )
+
+        # Each rank sends a 3x4 tensor with unique values to each destination
+        # Value pattern: rank * 1000 + dest * 100 + row * 10 + col
+        input_tensors = []
+        for dest in range(self.world_size):
+            t = torch.zeros(3, 4, dtype=torch.float32)
+            for row in range(3):
+                for col in range(4):
+                    t[row, col] = self.rank * 1000 + dest * 100 + row * 10 + col
+            input_tensors.append(fn(t))
+
+        output_tensors = [
+            fn(torch.zeros(3, 4, dtype=torch.float32)) for _ in range(self.world_size)
+        ]
+
+        fut = pg.alltoall(output_tensors, input_tensors).get_future()
+        fut.wait()
+
+        # Verify: output_tensors[src] should contain what rank src sent to us
+        for src in range(self.world_size):
+            expected = torch.zeros(3, 4, dtype=torch.float32)
+            for row in range(3):
+                for col in range(4):
+                    expected[row, col] = src * 1000 + self.rank * 100 + row * 10 + col
+            self.assertEqual(expected, output_tensors[src].cpu())
+
+    @requires_gloo()
+    def test_alltoall_multidim(self):
+        self._test_alltoall_multidim(lambda t: t.clone())
+
+    @skip_if_lt_x_gpu(2)
+    @requires_gloo()
+    def test_alltoall_multidim_cuda(self):
+        self._test_alltoall_multidim(lambda t: t.clone().cuda())
+
 
 class DistributedDataParallelTest(
     test_c10d_common.CommonDistributedDataParallelTest, MultiProcessTestCase
@@ -2250,7 +2470,9 @@ def train_loop(model, optimizer, iterations):
             self.assertEqual(p_withload, p_withoutload)
             self.assertEqual(p_non_ddp_withload, p_withoutload)
 
-    def _test_sparse_gradients(self, gradient_as_bucket_view=False):
+    def _test_sparse_gradients(
+        self, gradient_as_bucket_view=False, batched_grad_copy=False
+    ):
         process_group = self._get_process_group()
 
         # Ensure initialized weights and inputs are identical across processes
@@ -2261,6 +2483,7 @@ def _test_sparse_gradients(self, gradient_as_bucket_view=False):
             copy.deepcopy(vanilla_model),
             process_group=process_group,
             gradient_as_bucket_view=gradient_as_bucket_view,
+            batched_grad_copy=batched_grad_copy,
         )
 
         self._run_and_verify_sparse_gradients(vanilla_model, ddp_model)
@@ -2273,6 +2496,61 @@ def test_sparse_gradients(self):
     def test_sparse_gradients_grad_is_view(self):
         self._test_sparse_gradients(gradient_as_bucket_view=True)
 
+    @requires_gloo()
+    def test_sparse_gradients_batched_grad_copy(self):
+        self._test_sparse_gradients(batched_grad_copy=True)
+
+    @requires_gloo()
+    def test_sparse_gradients_grad_is_view_batched_grad_copy(self):
+        self._test_sparse_gradients(
+            gradient_as_bucket_view=True, batched_grad_copy=True
+        )
+
+    @requires_gloo()
+    def test_static_graph_batched_grad_copy(self):
+        """Test batched_grad_copy with static_graph (delay_all_reduce path)."""
+        process_group = self._get_process_group()
+        torch.manual_seed(42)
+
+        vanilla_model = nn.Linear(2, 4, bias=False).double()
+        model_bat = copy.deepcopy(vanilla_model)
+
+        ddp_bat = DistributedDataParallel(
+            model_bat,
+            process_group=process_group,
+            static_graph=True,
+            batched_grad_copy=True,
+        )
+
+        opt_van = torch.optim.SGD(vanilla_model.parameters(), lr=0.01)
+        opt_bat = torch.optim.SGD(ddp_bat.parameters(), lr=0.01)
+
+        mult = 2
+        # Run multiple iterations to cover both static_graph first iteration
+        # (delay_all_reduce path) and subsequent iterations.
+        for _ in range(3):
+            batch_size = mult * self.world_size
+            input = torch.rand(batch_size, 2, dtype=torch.double)
+
+            # Vanilla model on full batch (mean normalizes by batch size,
+            # matching DDP's allreduce averaging across ranks)
+            opt_van.zero_grad(set_to_none=True)
+            vanilla_model(input).mean().backward()
+
+            # DDP model on partial batch — allreduce averages gradients
+            partial_input = input.split(mult)[self.rank]
+            opt_bat.zero_grad(set_to_none=True)
+            ddp_bat(partial_input).mean().backward()
+
+            for p_van, p_bat in zip(vanilla_model.parameters(), ddp_bat.parameters()):
+                self.assertEqual(p_van.grad, p_bat.grad)
+
+            opt_van.step()
+            opt_bat.step()
+
+        for p_van, p_bat in zip(vanilla_model.parameters(), ddp_bat.parameters()):
+            self.assertEqual(p_van, p_bat)
+
     @requires_gloo()
     def test_ddp_comm_hook_future_passing_cpu(self):
         """
@@ -2630,6 +2908,287 @@ def test_forward_backward_optimizer(self):
             output.backward()
             optimizer.step()
 
+    def _create_reducer(
+        self,
+        model,
+        find_unused_parameters=False,
+        gradient_as_bucket_view=False,
+        batched_grad_copy=False,
+    ):
+        parameters = list(model.parameters())
+        group_by_dtype = groupby(
+            range(len(parameters)), key=lambda i: parameters[i].dtype
+        )
+        buckets = [list(indices) for _, indices in group_by_dtype]
+        return dist.Reducer(
+            parameters,
+            buckets,
+            [dist._DEFAULT_FIRST_BUCKET_BYTES for _ in range(len(buckets))],
+            self.process_group,
+            find_unused_parameters=find_unused_parameters,
+            gradient_as_bucket_view=gradient_as_bucket_view,
+            batched_grad_copy=batched_grad_copy,
+        )
+
+    def _run_forward_backward(self, model, reducer, batch_size=10):
+        """Run one forward+backward pass, return loss."""
+        loss_fn = nn.CrossEntropyLoss()
+        input = torch.rand([batch_size, 2], dtype=torch.double)
+        target = torch.LongTensor([random.randrange(4) for _ in range(batch_size)])
+        output = loss_fn(model(input), target)
+        reducer.prepare_for_backward(output)
+        output.backward()
+        return output
+
+    @requires_gloo()
+    def test_batched_grad_copy_basic(self):
+        """Verify batched_grad_copy uses _foreach_copy_ and produces grads."""
+        model = self._create_mixed_precision_model()
+        reducer = self._create_reducer(model, batched_grad_copy=True)
+        reducer.prepare_for_forward()
+
+        with torch.profiler.profile(
+            activities=[torch.profiler.ProfilerActivity.CPU],
+        ) as prof:
+            self._run_forward_backward(model, reducer)
+
+        events = [e.key for e in prof.key_averages()]
+        self.assertIn("aten::_foreach_copy_", events)
+        # mul_out is the per-param copy+div — should NOT appear when batched
+        self.assertNotIn("torch::distributed::reducer::mul_out", events)
+
+        for p in model.parameters():
+            if p.grad is not None:
+                self.assertFalse(torch.all(p.grad == 0))
+
+    @requires_gloo()
+    def test_batched_grad_copy_matches_default(self):
+        """Verify grads are numerically identical with and without batching."""
+        torch.manual_seed(42)
+        model_ref = self._create_mixed_precision_model()
+        model_bat = copy.deepcopy(model_ref)
+
+        input = torch.rand([10, 2], dtype=torch.double)
+        target = torch.LongTensor([random.randrange(4) for _ in range(10)])
+        loss_fn = nn.CrossEntropyLoss()
+
+        # Reference: no batching
+        reducer_ref = self._create_reducer(model_ref)
+        reducer_ref.prepare_for_forward()
+        out_ref = loss_fn(model_ref(input), target)
+        reducer_ref.prepare_for_backward(out_ref)
+        out_ref.backward()
+
+        # Batched
+        reducer_bat = self._create_reducer(model_bat, batched_grad_copy=True)
+        reducer_bat.prepare_for_forward()
+        out_bat = loss_fn(model_bat(input), target)
+        reducer_bat.prepare_for_backward(out_bat)
+        out_bat.backward()
+
+        for p_ref, p_bat in zip(model_ref.parameters(), model_bat.parameters()):
+            self.assertEqual(p_ref.grad, p_bat.grad)
+
+    @requires_gloo()
+    def test_batched_grad_copy_with_set_to_none(self):
+        """Multi-iteration with set_to_none=True — the primary use case."""
+        torch.manual_seed(42)
+        model_ref = self._create_mixed_precision_model()
+        model_bat = copy.deepcopy(model_ref)
+
+        opt_ref = torch.optim.SGD(model_ref.parameters(), lr=0.01)
+        opt_bat = torch.optim.SGD(model_bat.parameters(), lr=0.01)
+
+        reducer_ref = self._create_reducer(model_ref)
+        reducer_bat = self._create_reducer(model_bat, batched_grad_copy=True)
+        loss_fn = nn.CrossEntropyLoss()
+
+        for i in range(3):
+            input = torch.rand([10, 2], dtype=torch.double)
+            target = torch.LongTensor([random.randrange(4) for _ in range(10)])
+
+            # set_to_none=True is the default
+            opt_ref.zero_grad(set_to_none=True)
+            opt_bat.zero_grad(set_to_none=True)
+
+            reducer_ref.prepare_for_forward()
+            out_ref = loss_fn(model_ref(input), target)
+            reducer_ref.prepare_for_backward(out_ref)
+            out_ref.backward()
+
+            reducer_bat.prepare_for_forward()
+            out_bat = loss_fn(model_bat(input), target)
+            reducer_bat.prepare_for_backward(out_bat)
+            out_bat.backward()
+
+            for p_ref, p_bat in zip(model_ref.parameters(), model_bat.parameters()):
+                self.assertEqual(p_ref.grad, p_bat.grad)
+
+            opt_ref.step()
+            opt_bat.step()
+
+        # Params should be identical after training
+        for p_ref, p_bat in zip(model_ref.parameters(), model_bat.parameters()):
+            self.assertEqual(p_ref, p_bat)
+
+    @requires_gloo()
+    def test_batched_grad_copy_with_bucket_view(self):
+        """Test batched_grad_copy + gradient_as_bucket_view (alias path)."""
+        torch.manual_seed(42)
+        model_ref = self._create_mixed_precision_model()
+        model_bat = copy.deepcopy(model_ref)
+
+        opt_ref = torch.optim.SGD(model_ref.parameters(), lr=0.01)
+        opt_bat = torch.optim.SGD(model_bat.parameters(), lr=0.01)
+
+        reducer_ref = self._create_reducer(model_ref, gradient_as_bucket_view=True)
+        reducer_bat = self._create_reducer(
+            model_bat,
+            gradient_as_bucket_view=True,
+            batched_grad_copy=True,
+        )
+        loss_fn = nn.CrossEntropyLoss()
+
+        for i in range(3):
+            input = torch.rand([10, 2], dtype=torch.double)
+            target = torch.LongTensor([random.randrange(4) for _ in range(10)])
+
+            # First iteration: grads undefined (set_to_none default).
+            # Subsequent iterations: grads alias bucket views.
+            opt_ref.zero_grad(set_to_none=True)
+            opt_bat.zero_grad(set_to_none=True)
+
+            reducer_ref.prepare_for_forward()
+            out_ref = loss_fn(model_ref(input), target)
+            reducer_ref.prepare_for_backward(out_ref)
+            out_ref.backward()
+
+            reducer_bat.prepare_for_forward()
+            out_bat = loss_fn(model_bat(input), target)
+            reducer_bat.prepare_for_backward(out_bat)
+            out_bat.backward()
+
+            for p_ref, p_bat in zip(model_ref.parameters(), model_bat.parameters()):
+                self.assertEqual(p_ref.grad, p_bat.grad)
+
+            opt_ref.step()
+            opt_bat.step()
+
+        for p_ref, p_bat in zip(model_ref.parameters(), model_bat.parameters()):
+            self.assertEqual(p_ref, p_bat)
+
+    @requires_gloo()
+    def test_batched_grad_copy_with_unused_parameters(self):
+        """Test batched_grad_copy + find_unused_parameters with fc3 unused."""
+        torch.manual_seed(42)
+        model_ref = self._create_mixed_precision_model()
+        model_bat = copy.deepcopy(model_ref)
+
+        reducer_ref = self._create_reducer(model_ref, find_unused_parameters=True)
+        reducer_bat = self._create_reducer(
+            model_bat, find_unused_parameters=True, batched_grad_copy=True
+        )
+        loss_fn = nn.CrossEntropyLoss()
+
+        input = torch.rand([10, 2], dtype=torch.double)
+        target = torch.LongTensor([random.randrange(4) for _ in range(10)])
+
+        reducer_ref.prepare_for_forward()
+        out_ref = loss_fn(model_ref(input, use_fc3=False), target)
+        reducer_ref.prepare_for_backward(out_ref)
+        out_ref.backward()
+
+        reducer_bat.prepare_for_forward()
+        out_bat = loss_fn(model_bat(input, use_fc3=False), target)
+        reducer_bat.prepare_for_backward(out_bat)
+        out_bat.backward()
+
+        # fc3 is unused — its grad should remain None in both cases
+        self.assertIsNone(model_ref.fc3.weight.grad)
+        self.assertIsNone(model_bat.fc3.weight.grad)
+
+        # Used parameters should have identical grads
+        for (name_ref, p_ref), (name_bat, p_bat) in zip(
+            model_ref.named_parameters(), model_bat.named_parameters()
+        ):
+            if p_ref.grad is not None:
+                self.assertEqual(p_ref.grad, p_bat.grad, msg=name_ref)
+
+    @requires_gloo()
+    def test_batched_grad_copy_with_create_graph(self):
+        """batched_grad_copy falls back to non-batched path with create_graph=True."""
+        torch.manual_seed(42)
+
+        # Run with batched_grad_copy=False (reference)
+        model_ref = self._create_mixed_precision_model()
+        reducer_ref = self._create_reducer(model_ref)
+        loss_fn = nn.CrossEntropyLoss()
+        input = torch.rand([10, 2], dtype=torch.double)
+        target = torch.LongTensor([random.randrange(4) for _ in range(10)])
+
+        reducer_ref.prepare_for_forward()
+        out_ref = loss_fn(model_ref(input), target)
+        reducer_ref.prepare_for_backward(out_ref)
+        ref_error = None
+        try:
+            out_ref.backward(create_graph=True)
+        except RuntimeError as e:
+            ref_error = str(e)
+
+        # Run with batched_grad_copy=True — should fall back to same path
+        torch.manual_seed(42)
+        model_bat = self._create_mixed_precision_model()
+        reducer_bat = self._create_reducer(model_bat, batched_grad_copy=True)
+
+        reducer_bat.prepare_for_forward()
+        out_bat = loss_fn(model_bat(input), target)
+        reducer_bat.prepare_for_backward(out_bat)
+        bat_error = None
+        try:
+            out_bat.backward(create_graph=True)
+        except RuntimeError as e:
+            bat_error = str(e)
+
+        # Both should behave identically (either both succeed or both fail)
+        self.assertEqual(ref_error, bat_error)
+
+    @requires_gloo()
+    def test_batched_grad_copy_with_comm_hook(self):
+        """batched_grad_copy works correctly when a comm_hook is registered."""
+        torch.manual_seed(42)
+        model_ref = self._create_mixed_precision_model()
+        model_bat = copy.deepcopy(model_ref)
+
+        reducer_ref = self._create_reducer(model_ref)
+        reducer_bat = self._create_reducer(model_bat, batched_grad_copy=True)
+
+        # A simple comm hook that just allreduces and divides by world_size.
+        def simple_hook(state, bucket):
+            fut = torch.futures.Future()
+            fut.set_result(bucket.buffer() / state)
+            return fut
+
+        world_size = 1  # single-process test
+        dist._register_comm_hook(reducer_ref, world_size, simple_hook)
+        dist._register_comm_hook(reducer_bat, world_size, simple_hook)
+
+        loss_fn = nn.CrossEntropyLoss()
+        input = torch.rand([10, 2], dtype=torch.double)
+        target = torch.LongTensor([random.randrange(4) for _ in range(10)])
+
+        reducer_ref.prepare_for_forward()
+        out_ref = loss_fn(model_ref(input), target)
+        reducer_ref.prepare_for_backward(out_ref)
+        out_ref.backward()
+
+        reducer_bat.prepare_for_forward()
+        out_bat = loss_fn(model_bat(input), target)
+        reducer_bat.prepare_for_backward(out_bat)
+        out_bat.backward()
+
+        for p_ref, p_bat in zip(model_ref.parameters(), model_bat.parameters()):
+            self.assertEqual(p_ref.grad, p_bat.grad)
+
 
 @skip_if_win32()
 class ProcessGroupGlooLazyInitTest(ProcessGroupGlooTest):
@@ -2952,10 +3511,15 @@ def test_new_group_local_sync_sanity_check(self):
     def test_new_group_local_sync_duplicate_pg(self):
         self._test_new_group_local_sync_duplicate_pg(backend="gloo")
 
+    @requires_gloo()
+    def test_new_group_ordered(self):
+        self._test_new_group_ordered(backend="gloo")
+
 
 if __name__ == "__main__":
-    assert not torch.cuda._initialized, (
-        "test_distributed must not have initialized CUDA context on main process"
-    )
+    if torch.cuda._initialized:
+        raise AssertionError(
+            "test_distributed must not have initialized CUDA context on main process"
+        )
 
     run_tests()
diff --git a/test/distributed/test_c10d_nccl.py b/test/distributed/test_c10d_nccl.py
index ce78fdea48e4a..4ea532476c1c2 100644
--- a/test/distributed/test_c10d_nccl.py
+++ b/test/distributed/test_c10d_nccl.py
@@ -1975,6 +1975,76 @@ def test_block_current_stream(self):
         work.wait()
         torch.cuda.synchronize()
 
+    @requires_nccl_version(
+        (2, 29, 7), "Need NCCL 2.29.7+ for backend.suspend and backend.memory_stats"
+    )
+    @skip_but_pass_in_sandcastle_if(not TEST_MULTIGPU, "NCCL test requires 2+ GPUs")
+    def test_suspend(self):
+        """Test that suspend can be called on the NCCL backend."""
+        store = c10d.FileStore(self.file_name, self.world_size)
+        device = torch.device(f"cuda:{self.rank}")
+        pg = self._create_process_group_nccl(store, self.opts(), device_id=device)
+
+        # Run a large collective to cause NCCL to allocate internal memory
+        dist.all_reduce(torch.zeros(1024 * 1024 * 512, device=device))
+
+        backend = pg._get_backend(device)
+        backend.suspend()
+        # Confirm that the memory is suspended
+        stats = backend.memory_stats()
+        self.assertEqual(stats["suspended"], 1)
+
+    @requires_nccl_version(
+        (2, 29, 7), "Need NCCL 2.29.7+ for backend.memory_stats / ncclCommMemStats"
+    )
+    @skip_but_pass_in_sandcastle_if(not TEST_MULTIGPU, "NCCL test requires 2+ GPUs")
+    def test_get_memory_stats(self):
+        """Test that get_memory_stats returns a dict of memory stats."""
+        store = c10d.FileStore(self.file_name, self.world_size)
+        device = torch.device(f"cuda:{self.rank}")
+        pg = self._create_process_group_nccl(store, self.opts(), device_id=device)
+
+        # Run a large collective to cause NCCL to allocate internal memory
+        dist.all_reduce(torch.zeros(1024 * 1024 * 512, device=device))
+
+        backend = pg._get_backend(device)
+        stats = backend.memory_stats()
+        self.assertIsInstance(stats, dict)
+        for key in ("suspend", "suspended", "persist", "total"):
+            self.assertIn(key, stats)
+        print(stats)
+
+    @requires_nccl_version(
+        (2, 29, 7),
+        "Need NCCL 2.29.7+ for backend.resume, backend.suspend and backend.memory_stats",
+    )
+    @skip_but_pass_in_sandcastle_if(not TEST_MULTIGPU, "NCCL test requires 2+ GPUs")
+    def test_resume(self):
+        """Test the full suspend/resume cycle with collectives."""
+        store = c10d.FileStore(self.file_name, self.world_size)
+        device = torch.device(f"cuda:{self.rank}")
+        pg = self._create_process_group_nccl(store, self.opts(), device_id=device)
+        backend = pg._get_backend(device)
+
+        # Run a large collective to cause NCCL to allocate internal memory
+        dist.all_reduce(torch.zeros(1024 * 1024 * 512, device=device))
+
+        # Suspend (release memory)
+        backend.suspend()
+        # Resume
+        backend.resume()
+        # Confirm that the memory is resumed
+        stats = backend.memory_stats()
+        self.assertEqual(stats["suspended"], 0)
+
+        # Run a collective to verify the communicator still works
+        tensor = torch.ones(1024, device=device, dtype=torch.float32)
+        dist.all_reduce(tensor)
+        expected = torch.full(
+            (1024,), self.world_size, device=device, dtype=torch.float32
+        )
+        self.assertEqual(tensor, expected)
+
 
 class DistributedDataParallelTest(
     test_c10d_common.CommonDistributedDataParallelTest, MultiProcessTestCase
@@ -4395,29 +4465,100 @@ def test_all_reduce_coalesced_manager_nccl(self):
         )
         process_group = c10d.distributed_c10d._get_default_group()
         device = torch.device(f"cuda:{self.rank:d}")
-        tensors = [
-            torch.full((60 + i,), self.rank + 1 + i, device=device, dtype=torch.float)
+
+        for async_ops in [False, True]:
+            tensors = [
+                torch.full(
+                    (60 + i,), self.rank + 1 + i, device=device, dtype=torch.float
+                )
+                for i in range(5)
+            ]
+            with torch.distributed._coalescing_manager(
+                group=process_group, device=device, async_ops=async_ops
+            ) as cm:
+                for tensor in tensors:
+                    torch.distributed.all_reduce(tensor)
+
+            self.assertEqual(len(cm.works), 1 if async_ops else 0)
+            cm.wait()
+
+            for i, t in enumerate(tensors):
+                self.assertEqual(
+                    t,
+                    torch.full_like(
+                        t, self.world_size * (i + (self.world_size + 1.0) / 2.0)
+                    ),
+                )
+
+    @requires_nccl()
+    @skip_if_lt_x_gpu(2)
+    def test_all_gather_into_tensor_coalesced_manager_nccl(self):
+        store = c10d.FileStore(self.file_name, self.world_size)
+        c10d.init_process_group(
+            backend="nccl", store=store, rank=self.rank, world_size=self.world_size
+        )
+        process_group = c10d.distributed_c10d._get_default_group()
+        device = torch.device(f"cuda:{self.rank:d}")
+        input_tensors = [
+            torch.arange(
+                self.rank * (60 + i),
+                (self.rank + 1) * (60 + i),
+                device=device,
+                dtype=torch.float,
+            )
             for i in range(5)
         ]
-        with torch.distributed._coalescing_manager(
-            group=process_group, device=device, async_ops=True
-        ) as cm:
-            for tensor in tensors:
-                torch.distributed.all_reduce(tensor)
-        self.assertEqual(len(cm.works), 1)
-        cm.wait()
-        for i, t in enumerate(tensors):
-            self.assertEqual(
-                t,
-                torch.full_like(
-                    t, self.world_size * (i + (self.world_size + 1.0) / 2.0)
-                ),
-            )
+        for async_ops in [False, True]:
+            output_tensors = [
+                torch.zeros(
+                    (60 + i) * self.world_size, device=device, dtype=torch.float
+                )
+                for i in range(5)
+            ]
+            with torch.distributed._coalescing_manager(
+                group=process_group, device=device, async_ops=async_ops
+            ) as cm:
+                for input_t, output_t in zip(input_tensors, output_tensors):
+                    torch.distributed.all_gather_into_tensor(output_t, input_t)
+
+            self.assertEqual(len(cm.works), 1 if async_ops else 0)
+            cm.wait()
+
+            for i, output_t in enumerate(output_tensors):
+                expected = torch.arange(
+                    (60 + i) * self.world_size, device=device, dtype=torch.float
+                )
+                self.assertEqual(output_t, expected)
+
+    @requires_nccl()
+    @skip_if_lt_x_gpu(2)
+    def test_coalesced_manager_op_integrity(self):
+        store = c10d.FileStore(self.file_name, self.world_size)
+        c10d.init_process_group(
+            backend="nccl", store=store, rank=self.rank, world_size=self.world_size
+        )
+        process_group = c10d.distributed_c10d._get_default_group()
+        device = torch.device(f"cuda:{self.rank:d}")
+        with self.assertRaisesRegex(
+            RuntimeError,
+            "Coalescing manager requires all collectives to be the same type",
+        ):
+            with torch.distributed._coalescing_manager(
+                group=process_group, async_ops=True
+            ):
+                t = torch.ones(60, device=device) * (self.rank + 1)
+                torch.distributed.all_reduce(t)
+                output = torch.zeros(60 * self.world_size, device=device)
+                torch.distributed.all_gather_into_tensor(output, t)
 
     @requires_nccl()
     @skip_if_lt_x_gpu(2)
     @runOnRocmArch(MI300_ARCH)
-    def test_intra_node_comm_all_reduce(self):
+    @parametrize(
+        "custom_group_name",
+        [True, False],
+    )
+    def test_intra_node_comm_all_reduce(self, custom_group_name):
         from torch._C._distributed_c10d import _get_intra_node_comm_usage_counter
         from torch.testing._internal.common_cuda import SM80OrLater
 
@@ -4430,12 +4571,20 @@ def test_intra_node_comm_all_reduce(self):
         if not SM80OrLater:
             raise SkipTest("Test requires sm>=80")
 
+        group_name = ""
+        if custom_group_name:
+            group_name = "a_custom_group_name"
+
         store = c10d.FileStore(self.file_name, self.world_size)
         os.environ["ENABLE_INTRA_NODE_COMM"] = "1"
         os.environ["TEST_INTRA_NODE_COMM"] = "1"
         torch.cuda.set_device(self.rank)
         c10d.init_process_group(
-            backend="nccl", rank=self.rank, world_size=self.world_size, store=store
+            backend="nccl",
+            rank=self.rank,
+            world_size=self.world_size,
+            store=store,
+            group_name=group_name,
         )
         expect = self.world_size * (self.world_size - 1) // 2
 
@@ -4791,6 +4940,9 @@ def test_reduce_scatter_tensor_coalesced(self):
         self.assertEqual(output_tensors, input_tensors[self.rank] * self.world_size)
 
 
+instantiate_parametrized_tests(CommTest)
+
+
 class SetDeviceMethod(Enum):
     TORCH_CUDA_SET = auto()  # torch.cuda.set_device
     COLLECTIVE_ARGUMENT = auto()  # broadcast_object_list(device=)
@@ -4890,6 +5042,11 @@ def test_new_group_local_sync_sanity_check(self):
     def test_new_group_local_sync_duplicated_pg(self):
         self._test_new_group_local_sync_duplicate_pg(backend="nccl")
 
+    @requires_nccl()
+    @skip_if_lt_x_gpu(4)
+    def test_new_group_ordered(self):
+        self._test_new_group_ordered(backend="nccl")
+
     def _init_two_pg2_subgroups(self, world_size: int = 4):
         if world_size != 4:
             raise NotImplementedError(
@@ -5273,6 +5430,33 @@ def test_scatter_object_list_subgroup(self, group_rank):
                 )
         self.assertEqual(scatter_object_output_list, expected)
 
+    @requires_nccl()
+    @skip_if_lt_x_gpu(4)
+    @parametrize("float8_dtype", [torch.float8_e4m3fn, torch.float8_e5m2])
+    def test_broadcast_float8(self, float8_dtype):
+        device = torch.device(f"cuda:{self.rank}")
+        if sm_is_or_higher_than(device, 9, 0):  # noqa: F821
+            self.skipTest("FP8 broadcast natively supported on sm90+")
+        store = dist.FileStore(self.file_name, self.world_size)
+        dist.init_process_group(
+            "nccl",
+            world_size=self.world_size,
+            rank=self.rank,
+            store=store,
+        )
+        src_rank = 0
+        if self.rank == src_rank:
+            tensor = torch.arange(8, device=device, dtype=torch.float32).to(
+                float8_dtype
+            )
+        else:
+            tensor = torch.empty(8, device=device, dtype=float8_dtype)
+
+        expected = torch.arange(8, device=device, dtype=torch.float32).to(float8_dtype)
+
+        dist.broadcast(tensor, src=src_rank)
+        self.assertEqual(tensor, expected)
+
 
 instantiate_parametrized_tests(LargeCommTest)
 
@@ -5487,6 +5671,8 @@ def _create_process_group_nccl(self):
         return pg
 
     def tearDown(self):
+        os.environ.pop("TORCH_NCCL_DEBUG_INFO_TEMP_FILE", None)
+        os.environ.pop("TORCH_NCCL_DEBUG_INFO_PIPE_FILE", None)
         super().tearDown()
         try:
             os.remove(self.file_name)
@@ -5862,7 +6048,10 @@ def test_trace_while_active(self, timing_enabled, only_active):
                         and _get_torch_rocm_version() >= (6, 4)
                         and timing_enabled
                     ):
-                        assert t[-1]["state"] in ("scheduled", "started")
+                        if t[-1]["state"] not in ("scheduled", "started"):
+                            raise AssertionError(
+                                f"Expected state in ('scheduled', 'started'), got {t[-1]['state']}"
+                            )
                     else:
                         self.assertEqual(
                             t[-1]["state"], self.started_or_scheduled(timing_enabled)
@@ -6034,6 +6223,449 @@ def test_batched_send_recv(self, op_sizes_per_coalesce, timing_enabled):
                 self.assertTrue("duration_ms" not in t["entries"][coalesced_op])
             self.assertEqual(t["entries"][coalesced_op]["timeout_ms"], 600000)
 
+    @requires_nccl()
+    @skip_but_pass_in_sandcastle_if(not TEST_MULTIGPU, "NCCL test requires 2+ GPUs")
+    @torch._dynamo.config.patch({"enable_p2p_compilation": True})
+    @parametrize(
+        "op_sizes_per_coalesce",
+        [
+            [(2, 3)],
+            [(2, 3), (5, 5), (1,)],
+        ],
+    )
+    @parametrize("timing_enabled", [True, False])
+    def test_batched_send_recv_compiled(self, op_sizes_per_coalesce, timing_enabled):
+        def _pattern(tensors):
+            ops = list()
+            for tensor in tensors:
+                if self.rank == 0:
+                    ops.append(dist.P2POp(dist.irecv, tensor, 1))
+                elif self.rank == 1:
+                    ops.append(dist.P2POp(dist.isend, tensor, 0))
+                else:
+                    raise NotImplementedError
+            for work in dist.batch_isend_irecv(ops):
+                work.wait()
+
+        if self.rank == self.MAIN_PROCESS_RANK:
+            return
+        pg = self._create_process_group_nccl()
+        if timing_enabled:
+            pg._enable_collectives_timing()
+
+        compiled_fn = torch.compile(_pattern)
+
+        num_coalesced_ops = 20
+        ops_per_coalesce = len(op_sizes_per_coalesce)
+
+        for _ in range(num_coalesced_ops):
+            if self.rank == 0:
+                tensors = [
+                    torch.zeros(input_sizes).to(self.local_device)
+                    for input_sizes in op_sizes_per_coalesce
+                ]
+            elif self.rank == 1:
+                tensors = [
+                    torch.full(input_sizes, 2.0).to(self.local_device)
+                    for input_sizes in op_sizes_per_coalesce
+                ]
+            else:
+                raise NotImplementedError
+
+            compiled_fn(tensors)
+            if self.rank == 0:
+                self.assertEqual(len(tensors), ops_per_coalesce)
+                for tensor, input_sizes in zip(tensors, op_sizes_per_coalesce):
+                    self.assertEqual(
+                        tensor, torch.full(input_sizes, 2.0, device=self.local_device)
+                    )
+
+        torch.cuda.synchronize()
+
+        if timing_enabled:
+            time.sleep(1)
+
+        t = pickle.loads(torch._C._distributed_c10d._dump_nccl_trace())
+        self.assertTrue(len(t["entries"]) > 0)
+        expected_total_entries = num_coalesced_ops * (ops_per_coalesce + 1)
+        self.assertEqual(len(t["entries"]), expected_total_entries)
+
+        for seq in range(num_coalesced_ops):
+            coalesced_op_idx = seq * (ops_per_coalesce + 1) + ops_per_coalesce
+
+            self.assertEqual(
+                t["entries"][coalesced_op_idx]["profiling_name"], "nccl:coalesced"
+            )
+            try:
+                self.assertEqual(t["entries"][coalesced_op_idx]["state"], "completed")
+            except Exception:
+                self.assertEqual(t["entries"][coalesced_op_idx]["state"], "scheduled")
+
+    @requires_nccl()
+    @skip_if_lt_x_gpu(2)
+    @torch._dynamo.config.patch({"enable_p2p_compilation": True})
+    def test_compiled_fire_and_forget_isend(self):
+        if self.rank == self.MAIN_PROCESS_RANK:
+            return
+        self._create_process_group_nccl()
+        device = torch.device(f"cuda:{self.rank}")
+
+        @torch.compile(fullgraph=True)
+        def fn(x):
+            if self.rank == 0:
+                dist.isend(x, dst=1)
+            else:
+                dist.irecv(x, src=0)
+            return x
+
+        if self.rank == 0:
+            x = torch.ones(4, device=device) * 42
+            out = fn(x)
+            self.assertEqual(out, x)
+        else:
+            x = torch.zeros(4, device=device)
+            out = fn(x)
+            expected = torch.ones(4, device=device) * 42
+            self.assertEqual(out, expected)
+            self.assertEqual(x, expected)
+
+        torch.cuda.synchronize(device=device)
+
+    def _single_isend_with_wait_pattern(self, tensor, dst_rank):
+        req = dist.isend(tensor, dst_rank)
+        req.wait()
+        return req
+
+    def _single_irecv_with_wait_pattern(self, tensor, src_rank):
+        req = dist.irecv(tensor, src_rank)
+        req.wait()
+        return req
+
+    def _paired_isend_irecv_with_waits_pattern(
+        self, send_tensor, recv_tensor, peer_rank
+    ):
+        if self.rank == 0:
+            send_req = dist.isend(send_tensor, peer_rank)
+            recv_req = dist.irecv(recv_tensor, peer_rank)
+        else:
+            recv_req = dist.irecv(recv_tensor, peer_rank)
+            send_req = dist.isend(send_tensor, peer_rank)
+
+        send_req.wait()
+        recv_req.wait()
+        return send_req, recv_req
+
+    def _multiple_isend_with_waits_pattern(self, tensors, dst_rank):
+        reqs = []
+        for tensor in tensors:
+            req = dist.isend(tensor, dst_rank)
+            reqs.append(req)
+
+        for req in reqs:
+            req.wait()
+        return reqs
+
+    def _multiple_irecv_with_waits_pattern(self, tensors, src_rank):
+        reqs = []
+        for tensor in tensors:
+            req = dist.irecv(tensor, src_rank)
+            reqs.append(req)
+
+        for req in reqs:
+            req.wait()
+        return reqs
+
+    @requires_nccl()
+    @skip_if_lt_x_gpu(2)
+    @torch._dynamo.config.patch({"enable_p2p_compilation": True})
+    def test_compiled_isend_with_wait(self):
+        if self.rank == self.MAIN_PROCESS_RANK:
+            return
+        self._create_process_group_nccl()
+        device = torch.device(f"cuda:{self.rank}")
+
+        compiled_isend_wait = torch.compile(self._single_isend_with_wait_pattern)
+        compiled_irecv_wait = torch.compile(self._single_irecv_with_wait_pattern)
+
+        if self.rank == 0:
+            tensor = torch.ones(10, device=device) * 42
+            compiled_isend_wait(tensor, 1)
+        elif self.rank == 1:
+            tensor = torch.zeros(10, device=device)
+            compiled_irecv_wait(tensor, 0)
+            self.assertEqual(tensor, torch.ones(10, device=device) * 42)
+
+        torch.cuda.synchronize(device=device)
+
+    @requires_nccl()
+    @skip_if_lt_x_gpu(2)
+    @torch._dynamo.config.patch({"enable_p2p_compilation": True})
+    def test_compiled_with_reduce_overhead(self):
+        if self.rank == self.MAIN_PROCESS_RANK:
+            return
+        self._create_process_group_nccl()
+        device = torch.device(f"cuda:{self.rank}")
+        peer = 1 - self.rank
+
+        def f(tensor):
+            recv_buf = torch.empty_like(tensor)
+            work = dist.batch_isend_irecv(
+                [
+                    dist.P2POp(dist.isend, tensor, peer),
+                    dist.P2POp(dist.irecv, recv_buf, peer),
+                ]
+            )
+            for w in work:
+                w.wait()
+            return recv_buf
+
+        compiled_f = torch.compile(f, mode="reduce-overhead")
+        tensor = torch.ones(10, device=device) * self.rank
+        result = compiled_f(tensor)
+        torch.cuda.synchronize(device=device)
+        expected = torch.ones(10, device=device) * peer
+        self.assertEqual(result, expected)
+
+    @requires_nccl()
+    @skip_if_lt_x_gpu(2)
+    @torch._dynamo.config.patch({"enable_p2p_compilation": True})
+    @parametrize("tensor_size", [(10,), (5, 5), (2, 3, 4)])
+    def test_compiled_paired_isend_irecv_with_waits(self, tensor_size):
+        if self.rank == self.MAIN_PROCESS_RANK:
+            return
+        self._create_process_group_nccl()
+        device = torch.device(f"cuda:{self.rank}")
+
+        compiled_paired_comm = torch.compile(
+            self._paired_isend_irecv_with_waits_pattern
+        )
+
+        send_tensor = torch.ones(tensor_size, device=device) * (self.rank + 1)
+        recv_tensor = torch.zeros(tensor_size, device=device)
+
+        peer_rank = 1 if self.rank == 0 else 0
+        compiled_paired_comm(send_tensor, recv_tensor, peer_rank)
+
+        expected_recv = torch.ones(tensor_size, device=device) * (peer_rank + 1)
+        self.assertEqual(recv_tensor, expected_recv)
+
+        torch.cuda.synchronize(device=device)
+
+    @requires_nccl()
+    @skip_if_lt_x_gpu(2)
+    @torch._dynamo.config.patch({"enable_p2p_compilation": True})
+    @parametrize("num_tensors", [1, 3, 5])
+    def test_compiled_multiple_isend_with_waits(self, num_tensors):
+        if self.rank == self.MAIN_PROCESS_RANK:
+            return
+        self._create_process_group_nccl()
+        device = torch.device(f"cuda:{self.rank}")
+
+        compiled_multi_isend = torch.compile(self._multiple_isend_with_waits_pattern)
+        compiled_multi_irecv = torch.compile(self._multiple_irecv_with_waits_pattern)
+
+        if self.rank == 0:
+            tensors = [torch.ones(10, device=device) * i for i in range(num_tensors)]
+            compiled_multi_isend(tensors, 1)
+
+        elif self.rank == 1:
+            tensors = [torch.zeros(10, device=device) for _ in range(num_tensors)]
+            compiled_multi_irecv(tensors, 0)
+
+            for i, tensor in enumerate(tensors):
+                expected = torch.ones(10, device=device) * i
+                self.assertEqual(tensor, expected)
+
+        torch.cuda.synchronize(device=device)
+
+    def _iterative_communication_pattern(self, tensor_size, num_iterations, peer_rank):
+        device = torch.device(f"cuda:{self.rank}")
+        for i in range(num_iterations):
+            if self.rank == 0:
+                tensor = torch.ones(tensor_size, device=device) * i
+                req = dist.isend(tensor, peer_rank)
+                req.wait()
+            else:
+                tensor = torch.zeros(tensor_size, device=device)
+                req = dist.irecv(tensor, peer_rank)
+                req.wait()
+
+    @requires_nccl()
+    @skip_if_lt_x_gpu(2)
+    @torch._dynamo.config.patch({"enable_p2p_compilation": True})
+    @parametrize("num_iterations", [5, 10])
+    def test_compiled_iterative_communication_with_waits(self, num_iterations):
+        if self.rank == self.MAIN_PROCESS_RANK:
+            return
+        self._create_process_group_nccl()
+        device = torch.device(f"cuda:{self.rank}")
+
+        compiled_iterative_comm = torch.compile(self._iterative_communication_pattern)
+
+        peer_rank = 1 if self.rank == 0 else 0
+        tensor_size = (100,)
+
+        compiled_iterative_comm(tensor_size, num_iterations, peer_rank)
+
+        torch.cuda.synchronize(device=device)
+
+    @requires_nccl()
+    @skip_if_lt_x_gpu(2)
+    @torch._dynamo.config.patch({"enable_p2p_compilation": True})
+    def test_compiled_isend_irecv_timing_stress_with_waits(self):
+        if self.rank == self.MAIN_PROCESS_RANK:
+            return
+        pg = self._create_process_group_nccl()
+        pg._enable_collectives_timing()
+        device = torch.device(f"cuda:{self.rank}")
+
+        compiled_isend_wait = torch.compile(self._single_isend_with_wait_pattern)
+        compiled_irecv_wait = torch.compile(self._single_irecv_with_wait_pattern)
+
+        num_iterations = 50
+
+        for i in range(num_iterations):
+            if self.rank == 0:
+                tensor = torch.ones(100, device=device) * i
+                compiled_isend_wait(tensor, 1)
+            elif self.rank == 1:
+                tensor = torch.zeros(100, device=device)
+                compiled_irecv_wait(tensor, 0)
+                expected = torch.ones(100, device=device) * i
+                self.assertEqual(tensor, expected)
+
+        torch.cuda.synchronize(device=device)
+
+    @requires_nccl()
+    @skip_if_lt_x_gpu(2)
+    @torch._dynamo.config.patch({"enable_p2p_compilation": True})
+    def test_p2p_interleave(self):
+        if self.rank == self.MAIN_PROCESS_RANK:
+            return
+        self._create_process_group_nccl()
+        torch.cuda.set_device(self.rank)
+
+        def _kernel(x0, x1, y0, y1):
+            r = dist.get_rank()
+            w = dist.get_world_size()
+            nxt = (r + 1) % w
+            prv = (r - 1) % w
+            work = dist.batch_isend_irecv(
+                [
+                    dist.P2POp(
+                        dist.isend,
+                        x0,
+                        nxt,
+                    ),
+                    dist.P2POp(
+                        dist.irecv,
+                        y0,
+                        prv,
+                    ),
+                ]
+            )
+            t0 = x0 * 2 + 1
+            for ww in work:
+                ww.wait()
+            a = y0 + t0
+            work = dist.batch_isend_irecv(
+                [
+                    dist.P2POp(
+                        dist.isend,
+                        a,
+                        nxt,
+                    ),
+                    dist.P2POp(
+                        dist.irecv,
+                        y1,
+                        prv,
+                    ),
+                ]
+            )
+            t1 = a * 1.000244140625
+            for ww in work:
+                ww.wait()
+            return y1 + t1
+
+        r = self.rank
+        M, N = 1024, 1024
+        x0_e = torch.full((M, N), r + 1, device="cuda", dtype=torch.float32)
+        x1_e = torch.full((M, N), r + 2, device="cuda", dtype=torch.float32)
+        y0_e = torch.zeros_like(x0_e)
+        y1_e = torch.zeros_like(x0_e)
+        out_e = _kernel(x0_e, x1_e, y0_e, y1_e)
+
+        x0_c = x0_e.clone()
+        x1_c = x1_e.clone()
+        y0_c = torch.zeros_like(x0_c)
+        y1_c = torch.zeros_like(x0_c)
+
+        kernel_compiled = torch.compile(_kernel)
+        _ = kernel_compiled(x0_c, x1_c, y0_c, y1_c)
+        y0_c.zero_()
+        y1_c.zero_()
+
+        out_c = kernel_compiled(x0_c, x1_c, y0_c, y1_c)
+        diff = (out_e - out_c).abs()
+        local_max = diff.max()
+        dist.all_reduce(local_max, op=dist.ReduceOp.MAX)
+        self.assertLess(float(local_max.item()), 1e-3)
+
+    @requires_nccl()
+    @skip_if_lt_x_gpu(2)
+    @torch._dynamo.config.patch({"enable_p2p_compilation": True})
+    @parametrize("num_steps", [2, 4])
+    @parametrize("M", [256, 1024])
+    def test_compiled_ring_attention_pattern(self, num_steps, M):
+        if self.rank == self.MAIN_PROCESS_RANK:
+            return
+        self._create_process_group_nccl()
+        torch.cuda.set_device(self.rank)
+        device = self.local_device
+        N = M
+        world = dist.get_world_size()
+        rank = dist.get_rank()
+        nxt = (rank + 1) % world
+        prv = (rank - 1) % world
+
+        def ring_attention_kernel(q, k):
+            recv_buf = torch.empty_like(k)
+            acc = torch.zeros_like(q)
+            current_k = k
+
+            for step in range(num_steps):
+                work = dist.batch_isend_irecv(
+                    [
+                        dist.P2POp(dist.isend, current_k, nxt),
+                        dist.P2POp(dist.irecv, recv_buf, prv),
+                    ]
+                )
+                local_score = torch.mm(q, current_k.t())
+                local_out = torch.mm(local_score, current_k)
+                acc = acc + local_out * (1.0 / (step + 1))
+
+                for w in work:
+                    w.wait()
+
+                current_k = recv_buf * 0.5
+                recv_buf = torch.empty_like(k)
+
+            final_score = torch.mm(q, current_k.t())
+            final_out = torch.mm(final_score, current_k)
+            return acc + final_out
+
+        q = torch.randn(M, N, device=device)
+        k = torch.randn(M, N, device=device) * 0.1
+        q_e, k_e = q.clone(), k.clone()
+        out_e = ring_attention_kernel(q_e, k_e)
+
+        q_c, k_c = q.clone(), k.clone()
+        compiled_kernel = torch.compile(ring_attention_kernel)
+        out_c = compiled_kernel(q_c, k_c)
+
+        torch.cuda.synchronize(device=device)
+        self.assertEqual(out_e, out_c, atol=1e-3, rtol=1e-3)
+
     @requires_nccl()
     @skip_but_pass_in_sandcastle_if(not TEST_MULTIGPU, "NCCL test requires 2+ GPUs")
     @parametrize(
@@ -6770,8 +7402,9 @@ def test_comm_recursive_split_group(self):
 
 
 if __name__ == "__main__":
-    assert not torch.cuda._initialized, (
-        "test_distributed must not have initialized CUDA context on main process"
-    )
+    if torch.cuda._initialized:
+        raise AssertionError(
+            "test_distributed must not have initialized CUDA context on main process"
+        )
 
     run_tests()
diff --git a/test/distributed/test_c10d_ops_nccl.py b/test/distributed/test_c10d_ops_nccl.py
index b42bc075870a7..9a663d3fdd3db 100644
--- a/test/distributed/test_c10d_ops_nccl.py
+++ b/test/distributed/test_c10d_ops_nccl.py
@@ -22,13 +22,12 @@
 
 
 import torch.distributed as dist
-from torch.testing._internal.common_cuda import TEST_MULTIGPU
+from torch.testing._internal.common_cuda import PLATFORM_SUPPORTS_FP8, TEST_MULTIGPU
 from torch.testing._internal.common_distributed import (
     init_multigpu_helper,
     MultiProcContinuousTest,
     requires_nccl,
     requires_nccl_version,
-    sm_is_or_higher_than,
 )
 from torch.testing._internal.common_utils import (
     run_tests,
@@ -245,10 +244,11 @@ def allreduce(tensors, op):
 
     @requires_nccl_version((2, 24), "Need NCCL 2.24+ for Float8")
     @skip_but_pass_in_sandcastle_if(not TEST_MULTIGPU, "NCCL test requires 2+ GPUs")
+    @skip_but_pass_in_sandcastle_if(
+        not PLATFORM_SUPPORTS_FP8, "Float8 requires sm >= 90"
+    )
     def test_allreduce_float8(self):
         device = torch.device("cuda", self.rank_to_GPU[self.rank][0])
-        if not sm_is_or_higher_than(device, 9, 0):
-            self.skipTest("Float8 requires sm >= 90")
 
         numel = 1024
         tensor = torch.ones(numel, dtype=torch.float32, device=device).to(
@@ -950,10 +950,11 @@ def reduce_scatter_base(output_t, input_t):
 
     @requires_nccl_version((2, 24), "Need NCCL 2.24+ for Float8")
     @skip_but_pass_in_sandcastle_if(not TEST_MULTIGPU, "NCCL test requires 2+ GPUs")
+    @skip_but_pass_in_sandcastle_if(
+        not PLATFORM_SUPPORTS_FP8, "Float8 requires sm >= 90"
+    )
     def test_reduce_scatter_float8(self):
         device = torch.device("cuda", self.rank_to_GPU[self.rank][0])
-        if not sm_is_or_higher_than(device, 9, 0):
-            self.skipTest("Float8 requires sm >= 90")
 
         numel = 1024
         output_tensor = torch.zeros(numel, dtype=torch.float32, device=device).to(
diff --git a/test/distributed/test_c10d_pypg.py b/test/distributed/test_c10d_pypg.py
index 65faf2075daa6..840a2317d14d2 100644
--- a/test/distributed/test_c10d_pypg.py
+++ b/test/distributed/test_c10d_pypg.py
@@ -47,8 +47,10 @@ class LonelyRankProcessGroup(dist.ProcessGroup):
 
     def __init__(self, rank, world, use_wrapper):
         super().__init__(rank, world)
-        assert rank == 0
-        assert world == 1
+        if rank != 0:
+            raise AssertionError(f"Expected rank == 0, got {rank}")
+        if world != 1:
+            raise AssertionError(f"Expected world == 1, got {world}")
 
         self._rank = rank
         self._world = world
diff --git a/test/distributed/test_c10d_ucc.py b/test/distributed/test_c10d_ucc.py
index e3a4764d594f2..de6f5c3a17f4b 100644
--- a/test/distributed/test_c10d_ucc.py
+++ b/test/distributed/test_c10d_ucc.py
@@ -1090,8 +1090,9 @@ def test_allgather_base(self):
 
 
 if __name__ == "__main__":
-    assert not torch.cuda._initialized, (
-        "test_distributed must not have initialized CUDA context on main process"
-    )
+    if torch.cuda._initialized:
+        raise AssertionError(
+            "test_distributed must not have initialized CUDA context on main process"
+        )
 
     run_tests()
diff --git a/test/distributed/test_ce_colls.py b/test/distributed/test_ce_colls.py
index ff4da35b46525..bedda15a837b1 100644
--- a/test/distributed/test_ce_colls.py
+++ b/test/distributed/test_ce_colls.py
@@ -1,6 +1,5 @@
 # Owner(s): ["oncall: distributed"]
 import sys
-from typing import Optional
 
 import torch
 import torch.distributed as dist
@@ -25,11 +24,11 @@
 @requires_cuda_p2p_access()
 class NCCLCopyEngineCollectives(MultiProcContinuousTest):
     @classmethod
-    def backend_str(cls) -> Optional[str]:
+    def backend_str(cls) -> str | None:
         return "nccl"
 
     @classmethod
-    def opts(cls) -> Optional[dist.ProcessGroupNCCL.Options]:
+    def opts(cls) -> dist.ProcessGroupNCCL.Options | None:
         # Enable Zero-CTA policy for CE collectives
         opts = dist.ProcessGroupNCCL.Options()
         opts.config.cta_policy = dist.ProcessGroupNCCL.NCCL_CTA_POLICY_ZERO
diff --git a/test/distributed/test_collective_utils.py b/test/distributed/test_collective_utils.py
index 189b20fcaf148..d5cc98f461740 100644
--- a/test/distributed/test_collective_utils.py
+++ b/test/distributed/test_collective_utils.py
@@ -54,7 +54,8 @@ def test_broadcast_result(self) -> None:
         func.return_value = pg.rank()
 
         res = broadcast(data_or_fn=func, rank=0, pg=pg)
-        assert res == 0, f"Expect res to be 0 (got {res})"
+        if res != 0:
+            raise AssertionError(f"Expect res to be 0 (got {res})")
 
         if pg.rank() == 0:
             func.assert_called_once()
@@ -64,7 +65,8 @@ def test_broadcast_result(self) -> None:
         func.reset_mock()
 
         res = broadcast(data_or_fn=func, rank=1, pg=pg)
-        assert res == 1, f"Expect res to be 1 (got {res})"
+        if res != 1:
+            raise AssertionError(f"Expect res to be 1 (got {res})")
 
         if pg.rank() == 1:
             func.assert_called_once()
@@ -108,9 +110,10 @@ def test_all_gather_result(self) -> None:
 
         res = all_gather(data_or_fn=func, pg=pg)
         func.assert_called_once()
-        assert res == list(range(self.world_size)), (
-            f"Expect res to be list of 0 through {self.world_size} (got {res})"
-        )
+        if res != list(range(self.world_size)):
+            raise AssertionError(
+                f"Expect res to be list of 0 through {self.world_size} (got {res})"
+            )
 
     def test_all_gather_result_no_pg(self) -> None:
         """
diff --git a/test/distributed/test_composability.py b/test/distributed/test_composability.py
index 566a63d67302d..1a15cb10dc5d8 100644
--- a/test/distributed/test_composability.py
+++ b/test/distributed/test_composability.py
@@ -120,7 +120,11 @@ def _build_pp_stage(
     ):
         # divide the model (e.g. 8 layers) by the number of stages
         layers_per_stage = total_layers // num_stages
-        assert layers_per_stage * num_stages == total_layers
+        if layers_per_stage * num_stages != total_layers:
+            raise AssertionError(
+                f"layers_per_stage * num_stages != total_layers: "
+                f"{layers_per_stage * num_stages} vs {total_layers}"
+            )
         # return offset so validation code can match partial layer back to orig model
         offset = stage_idx * layers_per_stage
         partial_model = nn.Sequential(
@@ -352,7 +356,8 @@ def apply_dp(partial_model):
             pipeline_schedule.step(target=target_local)
         for m in partial_models:
             for p in m.parameters():
-                assert p.grad is not None
+                if p.grad is None:
+                    raise AssertionError("Expected p.grad to not be None")
                 # introduce a race condition for FSDP's reduce-scatter which could corrupt gradients if pipelining
                 # does not properly synchronize with FSDP
                 p.grad.div_(2.0)
@@ -450,8 +455,7 @@ def check_fsdp_unsharded_state(module, expected_unsharded=False):
             total_fsdp_params = 0
 
             for state in distributed_state._state_ctx.all_states:
-                if state._fsdp_param_group:
-                    group = state._fsdp_param_group
+                for group in state._fsdp_param_groups:
                     for fsdp_param in group.fsdp_params:
                         total_fsdp_params += 1
                         if fsdp_param.sharded_state == ShardedState.UNSHARDED:
diff --git a/test/distributed/test_compute_comm_reordering.py b/test/distributed/test_compute_comm_reordering.py
index 2e9a3ea171028..cc541ab8c38ad 100644
--- a/test/distributed/test_compute_comm_reordering.py
+++ b/test/distributed/test_compute_comm_reordering.py
@@ -61,12 +61,18 @@ def create_grouped_node_for_allreduce_and_its_deps(snodes):
         if isinstance(snode.node, ir._CollectiveKernel)
         and snode.node.op_overload == torch.ops._c10d_functional.all_reduce_.default
     ]
-    assert len(all_reduce_snodes) == 1
+    if len(all_reduce_snodes) != 1:
+        raise AssertionError(
+            f"Expected exactly 1 all_reduce_snode, got {len(all_reduce_snodes)}"
+        )
     all_reduce_snode = all_reduce_snodes[0]
     all_reduce_dep_snodes = [
         name_to_snode[node.name] for node in all_reduce_snode.node.inputs
     ]
-    assert len(all_reduce_dep_snodes) == 1
+    if len(all_reduce_dep_snodes) != 1:
+        raise AssertionError(
+            f"Expected exactly 1 all_reduce_dep_snode, got {len(all_reduce_dep_snodes)}"
+        )
     all_reduce_dep_snode = all_reduce_dep_snodes[0]
 
     grouped_snode = scheduler.GroupedSchedulerNode.create(
@@ -480,14 +486,32 @@ def fn(g1, g2, g3):
             fn(g1, g2, g3)
 
     def test_nccl_heuristics(self):
-        assert len(baseLat) == len(NCCL_ALGO)
-        assert all(len(x) == len(NCCL_PROTO) for x in baseLat)
+        if len(baseLat) != len(NCCL_ALGO):
+            raise AssertionError(
+                f"Expected len(baseLat) == len(NCCL_ALGO), got {len(baseLat)} vs {len(NCCL_ALGO)}"
+            )
+        if not all(len(x) == len(NCCL_PROTO) for x in baseLat):
+            raise AssertionError(
+                "Expected all elements in baseLat to have len(NCCL_PROTO)"
+            )
 
-        assert len(hwLat) == len(NCCL_HW)
-        assert all(len(x) == len(NCCL_ALGO) for x in hwLat)
-        assert all(len(y) == len(NCCL_PROTO) for x in hwLat for y in x)
+        if len(hwLat) != len(NCCL_HW):
+            raise AssertionError(
+                f"Expected len(hwLat) == len(NCCL_HW), got {len(hwLat)} vs {len(NCCL_HW)}"
+            )
+        if not all(len(x) == len(NCCL_ALGO) for x in hwLat):
+            raise AssertionError(
+                "Expected all elements in hwLat to have len(NCCL_ALGO)"
+            )
+        if not all(len(y) == len(NCCL_PROTO) for x in hwLat for y in x):
+            raise AssertionError(
+                "Expected all nested elements in hwLat to have len(NCCL_PROTO)"
+            )
 
-        assert len(llMaxBws) == len(NVIDIA_GPU_TYPE)
+        if len(llMaxBws) != len(NVIDIA_GPU_TYPE):
+            raise AssertionError(
+                f"Expected len(llMaxBws) == len(NVIDIA_GPU_TYPE), got {len(llMaxBws)} vs {len(NVIDIA_GPU_TYPE)}"
+            )
 
 
 if __name__ == "__main__":
diff --git a/test/distributed/test_control_collectives.py b/test/distributed/test_control_collectives.py
index 8e48735c77799..08fe1d27b26da 100644
--- a/test/distributed/test_control_collectives.py
+++ b/test/distributed/test_control_collectives.py
@@ -207,8 +207,9 @@ def f(rank: int) -> None:
 
 
 if __name__ == "__main__":
-    assert not torch.cuda._initialized, (
-        "test_distributed must not have initialized CUDA context on main process"
-    )
+    if torch.cuda._initialized:
+        raise AssertionError(
+            "test_distributed must not have initialized CUDA context on main process"
+        )
 
     run_tests()
diff --git a/test/distributed/test_cupy_as_tensor.py b/test/distributed/test_cupy_as_tensor.py
index e0a98ae960426..63b290e2e8e66 100644
--- a/test/distributed/test_cupy_as_tensor.py
+++ b/test/distributed/test_cupy_as_tensor.py
@@ -8,7 +8,10 @@
 import torch
 from torch.multiprocessing.reductions import reduce_tensor
 from torch.testing._internal.common_cuda import SM100OrLater
-from torch.testing._internal.common_distributed import MultiProcContinuousTest
+from torch.testing._internal.common_distributed import (
+    MultiProcContinuousTest,
+    skip_if_rocm_multiprocess,
+)
 from torch.testing._internal.common_utils import (
     requires_cuda_p2p_access,
     run_tests,
@@ -42,7 +45,10 @@ def from_buffer(
     data = torch.as_tensor(CupyWrapper(data_ptr, size_in_bytes), device=device).view(
         dtype
     )
-    assert data.data_ptr() == data_ptr
+    if data.data_ptr() != data_ptr:
+        raise AssertionError(
+            f"Expected data_ptr to be {data_ptr}, got {data.data_ptr()}"
+        )
     return data
 
 
@@ -64,6 +70,7 @@ def _init_device(self) -> None:
     def device(self) -> torch.device:
         return torch.device(device_type, self.rank)
 
+    @skip_if_rocm_multiprocess  # RuntimeError: pidfd_getfd Operation not permitted"
     @skip_but_pass_in_sandcastle_if(
         SM100OrLater,
         "Fails if ran in docker environment without privileged access (https://github.com/pytorch/pytorch/issues/165170)",
@@ -101,7 +108,8 @@ def test_cupy_as_tensor(self) -> None:
             tensor.fill_(1)
         device_module.synchronize()
         torch.distributed.barrier()
-        assert tensor.allclose(tensor, 1)
+        if not tensor.allclose(tensor, 1):
+            raise AssertionError("Expected tensor to be close to 1")
         torch.distributed.barrier()
 
     @classmethod
diff --git a/test/distributed/test_debug.py b/test/distributed/test_debug.py
index 24b4b157eacb3..1533b27806799 100644
--- a/test/distributed/test_debug.py
+++ b/test/distributed/test_debug.py
@@ -2,8 +2,13 @@
 
 import os
 import shutil
+import socket
 import tempfile
+import threading
 import time
+import unittest
+from http.server import BaseHTTPRequestHandler, HTTPServer
+from unittest.mock import patch
 
 import requests
 from requests.adapters import HTTPAdapter
@@ -15,13 +20,24 @@
 from torch.distributed.debug import start_debug_server, stop_debug_server
 from torch.distributed.debug._frontend import (
     DebugHandler,
+    fetch_thread_pool,
+    format_fetch_summary,
     NavLink,
     PeriodicDumper,
+    Response,
     Route,
 )
 from torch.testing._internal.common_utils import run_tests, TestCase
 
 
+try:
+    from torch.distributed.debug._frontend import fetch_aiohttp
+
+    HAS_AIOHTTP = True
+except ImportError:
+    HAS_AIOHTTP = False
+
+
 session = requests.Session()
 retry_strategy = Retry(total=5, backoff_factor=0.5)
 adapter = HTTPAdapter(max_retries=retry_strategy)
@@ -257,5 +273,207 @@ def test_stop_is_idempotent(self) -> None:
             dumper.stop()
 
 
+class TestFetchUnavailableWorkers(TestCase):
+    @staticmethod
+    def _get_refused_port() -> int:
+        """Return a port that will refuse connections."""
+        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+        sock.bind(("localhost", 0))
+        port = sock.getsockname()[1]
+        sock.close()
+        return port
+
+    def test_fetch_thread_pool_connection_refused(self) -> None:
+        port = self._get_refused_port()
+        resps = fetch_thread_pool(
+            [f"http://localhost:{port}/handler/ping"], timeout=1.0
+        )
+        self.assertEqual(len(resps), 1)
+        self.assertEqual(resps[0].status_code, 503)
+        self.assertIn("ConnectionError:", resps[0].text)
+
+    def test_fetch_thread_pool_timeout(self) -> None:
+        server = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+        server.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+        server.bind(("localhost", 0))
+        server.listen(1)
+        port = server.getsockname()[1]
+        try:
+            resps = fetch_thread_pool(
+                [f"http://localhost:{port}/handler/ping"], timeout=0.5
+            )
+            self.assertEqual(len(resps), 1)
+            self.assertEqual(resps[0].status_code, 408)
+            self.assertIn("Timeout:", resps[0].text)
+        finally:
+            server.close()
+
+    @unittest.skipUnless(HAS_AIOHTTP, "aiohttp not installed")
+    def test_fetch_aiohttp_connection_refused(self) -> None:
+        port = self._get_refused_port()
+        resps = fetch_aiohttp([f"http://localhost:{port}/handler/ping"], timeout=1.0)
+        self.assertEqual(len(resps), 1)
+        self.assertEqual(resps[0].status_code, 503)
+        self.assertIn("Error:", resps[0].text)
+
+    @unittest.skipUnless(HAS_AIOHTTP, "aiohttp not installed")
+    def test_fetch_aiohttp_timeout(self) -> None:
+        server = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+        server.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+        server.bind(("localhost", 0))
+        server.listen(1)
+        port = server.getsockname()[1]
+        try:
+            resps = fetch_aiohttp(
+                [f"http://localhost:{port}/handler/ping"], timeout=0.5
+            )
+            self.assertEqual(len(resps), 1)
+            # aiohttp may report timeout as 408 or wrap it in a ClientError (503)
+            self.assertIn(resps[0].status_code, (408, 503))
+        finally:
+            server.close()
+
+    def test_mixed_available_and_unavailable(self) -> None:
+        class _OKHandler(BaseHTTPRequestHandler):
+            def do_POST(self):
+                self.send_response(200)
+                self.end_headers()
+                self.wfile.write(b"pong")
+
+            def log_message(self, *args):
+                pass
+
+        http_server = HTTPServer(("localhost", 0), _OKHandler)
+        good_port = http_server.server_address[1]
+        t = threading.Thread(target=http_server.serve_forever, daemon=True)
+        t.start()
+
+        bad_port = self._get_refused_port()
+        try:
+            urls = [
+                f"http://localhost:{good_port}/handler/ping",
+                f"http://localhost:{bad_port}/handler/ping",
+            ]
+            resps = fetch_thread_pool(urls, timeout=1.0)
+            self.assertEqual(len(resps), 2)
+            self.assertEqual(resps[0].status_code, 200)
+            self.assertEqual(resps[0].text, "pong")
+            self.assertNotEqual(resps[1].status_code, 200)
+        finally:
+            http_server.shutdown()
+
+
+class TestFormatFetchSummary(TestCase):
+    def test_all_success_returns_none(self) -> None:
+        addrs = ["http://host0:1", "http://host1:1"]  # @lint-ignore
+        resps = [Response(200, "ok"), Response(200, "ok")]
+        self.assertIsNone(format_fetch_summary(addrs, resps))
+
+    def test_partial_failure(self) -> None:
+        addrs = ["http://h0:1", "http://h1:1", "http://h2:1"]  # @lint-ignore
+        resps = [
+            Response(200, "ok"),
+            Response(503, "Worker unavailable"),
+            Response(200, "ok"),
+        ]
+        summary = format_fetch_summary(addrs, resps)
+        self.assertIsNotNone(summary)
+        self.assertIn("PARTIAL DATA", summary)
+        self.assertIn("2/3", summary)
+        self.assertIn("Rank 1", summary)
+
+    def test_all_failed(self) -> None:
+        addrs = ["http://h0:1", "http://h1:1"]  # @lint-ignore
+        resps = [Response(503, "unavailable"), Response(408, "timeout")]
+        summary = format_fetch_summary(addrs, resps)
+        self.assertIsNotNone(summary)
+        self.assertIn("0/2", summary)
+
+    def test_timeout_message_included(self) -> None:
+        addrs = ["http://h0:1"]  # @lint-ignore
+        resps = [Response(408, "Request timed out after 10.0s")]
+        summary = format_fetch_summary(addrs, resps)
+        self.assertIn("timed out", summary)
+
+
+class TestFetchTimeout(TestCase):
+    def test_handler_default_fetch_timeout(self) -> None:
+        from torch.distributed.debug._frontend import _DEFAULT_FETCH_TIMEOUT
+
+        handler = _StubHandler("test", "data")
+        self.assertEqual(handler.fetch_timeout, _DEFAULT_FETCH_TIMEOUT)
+
+    def test_handler_fetch_timeout_override(self) -> None:
+        handler = _StubHandler("test", "data")
+        handler.fetch_timeout = 5.0
+        self.assertEqual(handler.fetch_timeout, 5.0)
+
+    def test_fetch_all_requires_timeout(self) -> None:
+        from torch.distributed.debug._frontend import fetch_all
+
+        with self.assertRaises(TypeError):
+            # fetch_all requires timeout as a keyword-only argument
+            fetch_all("test_endpoint")  # type: ignore[call-arg]
+
+
+class TestHandlerPartialDumps(TestCase):
+    @patch("torch.distributed.debug._debug_handlers.fetch_all")
+    def test_stacks_handler_partial_dump(self, mock_fetch_all) -> None:
+        from torch.distributed.debug._debug_handlers import StacksHandler
+
+        mock_fetch_all.return_value = (
+            [
+                "http://h0:1/handler/dump_traceback?",  # @lint-ignore
+                "http://h1:1/handler/dump_traceback?",  # @lint-ignore
+            ],
+            [Response(200, "stack0"), Response(503, "Worker unavailable")],
+        )
+        handler = StacksHandler()
+        content = handler.dump()
+        self.assertIn("PARTIAL DATA", content)
+        self.assertIn("1/2", content)
+        self.assertIn("stack0", content)
+        self.assertIn("Error: 503", content)
+
+    @patch("torch.distributed.debug._debug_handlers.fetch_all")
+    def test_stacks_handler_all_success(self, mock_fetch_all) -> None:
+        from torch.distributed.debug._debug_handlers import StacksHandler
+
+        mock_fetch_all.return_value = (
+            [
+                "http://h0:1/handler/dump_traceback?",  # @lint-ignore
+                "http://h1:1/handler/dump_traceback?",  # @lint-ignore
+            ],
+            [Response(200, "stack0"), Response(200, "stack1")],
+        )
+        handler = StacksHandler()
+        content = handler.dump()
+        self.assertNotIn("PARTIAL", content)
+        self.assertIn("stack0", content)
+        self.assertIn("stack1", content)
+
+    def test_periodic_dumper_writes_partial_data(self) -> None:
+        with tempfile.TemporaryDirectory() as tmp:
+            partial_content = (
+                "PARTIAL DATA: 1/2 workers responded\n"
+                "  Rank 1: Worker unavailable\n"
+                "\n"
+                "=== Rank 0 ===\ndata0\n"
+                "=== Rank 1 ===\nError: 503"
+            )
+            h = _StubHandler("partial_stacks", partial_content)
+            dumper = PeriodicDumper([h], tmp, interval_seconds=0.1)
+            dumper.start()
+            time.sleep(0.25)
+            dumper.stop()
+
+            files = os.listdir(tmp)
+            self.assertGreater(len(files), 0)
+            with open(os.path.join(tmp, files[0])) as f:
+                content = f.read()
+            self.assertIn("PARTIAL DATA", content)
+            self.assertIn("1/2 workers responded", content)
+
+
 if __name__ == "__main__":
     run_tests()
diff --git a/test/distributed/test_device_mesh.py b/test/distributed/test_device_mesh.py
index 48d31c6541311..b9eb142d83cf8 100644
--- a/test/distributed/test_device_mesh.py
+++ b/test/distributed/test_device_mesh.py
@@ -364,7 +364,8 @@ def test_from_group_with_global_pg(self):
     def test_from_group_with_invalid_mesh(self):
         global_pg = _get_default_group()
         global_pg_size = global_pg.size()
-        assert global_pg_size == 4, "Test assumes global world size of 4"
+        if global_pg_size != 4:
+            raise AssertionError("Test assumes global world size of 4")
         invalid_mesh = [[0, 1], [2, 3]]  # 2D mesh when we need 1D
         regex = r"Invalid mesh \[\[0, 1\], \[2, 3\]\] for ProcessGroup with ranks \[0, 1, 2, 3\]"
         with self.assertRaisesRegex(ValueError, regex):
@@ -1901,5 +1902,39 @@ def test_remap_to_tensor(self):
         self.assertEqual(result7, expected7)
 
 
+class ProcessGroupOpaqueTypeTest(TestCase):
+    """Test that ProcessGroup opaque type members are registered and exist on the class."""
+
+    def test_registered_members_exist_on_process_group(self):
+        from torch._library.opaque_object import get_member_type
+
+        # Every member registered in _register_distributed_opaque_types()
+        # must actually exist on ProcessGroup. This catches renames or
+        # removals of C++ attributes that would cause torch.compile
+        # (fullgraph=True) to silently register a stale name while the
+        # real attribute has moved.
+        registered_members = [
+            "size",
+            "rank",
+            "_get_backend_name",
+            "group_name",
+            "group_desc",
+            "__eq__",
+        ]
+        for member_name in registered_members:
+            self.assertIsNotNone(
+                get_member_type(ProcessGroup, member_name),
+                f"'{member_name}' is not registered as a ProcessGroup opaque "
+                f"type member. Add it to _register_distributed_opaque_types() "
+                f"in torch/distributed/device_mesh.py",
+            )
+            self.assertTrue(
+                hasattr(ProcessGroup, member_name),
+                f"'{member_name}' is registered as a ProcessGroup opaque type "
+                f"member but does not exist on the ProcessGroup class. "
+                f"Was it renamed or removed?",
+            )
+
+
 if __name__ == "__main__":
     run_tests()
diff --git a/test/distributed/test_dist2.py b/test/distributed/test_dist2.py
index 2c444fbfe5670..fd44be2d06db7 100644
--- a/test/distributed/test_dist2.py
+++ b/test/distributed/test_dist2.py
@@ -218,7 +218,8 @@ def test_group_split(self) -> None:
             [0], timeout=timedelta(seconds=30), group_name="subgroup_1"
         )
         if self.rank == 0:
-            assert subgroup is not None
+            if subgroup is None:
+                raise AssertionError("Expected subgroup to not be None")
             self.assertEqual(subgroup.size(), 1)
             backend = subgroup._get_backend(self.device)
             self.assertEqual(backend.options._timeout, timedelta(seconds=30))
@@ -231,7 +232,8 @@ def test_remote_group_merge(self) -> None:
         subgroup_1 = group.split_group([0], timeout=timedelta(seconds=30))
         subgroup_2 = group.split_group([1], timeout=timedelta(seconds=30))
         if self.rank == 0:
-            assert subgroup_1 is not None
+            if subgroup_1 is None:
+                raise AssertionError("Expected subgroup_1 to not be None")
             tcp_store = dist.TCPStore(
                 host_name=os.environ["MASTER_ADDR"],
                 port=29781,
@@ -246,7 +248,8 @@ def test_remote_group_merge(self) -> None:
             self.assertEqual(backend.options._timeout, timedelta(seconds=40))
             self.assertEqual(merged_pg.group_name, "merged_pg")
         else:
-            assert subgroup_2 is not None
+            if subgroup_2 is None:
+                raise AssertionError("Expected subgroup_2 to not be None")
             tcp_store = dist.TCPStore(
                 host_name=os.environ["MASTER_ADDR"],
                 port=29781,
@@ -302,8 +305,9 @@ def new_group(self) -> torch.distributed.ProcessGroup:
 
 
 if __name__ == "__main__":
-    assert not torch.cuda._initialized, (
-        "test_distributed must not have initialized CUDA context on main process"
-    )
+    if torch.cuda._initialized:
+        raise AssertionError(
+            "test_distributed must not have initialized CUDA context on main process"
+        )
 
     run_tests()
diff --git a/test/distributed/test_dynamo_distributed.py b/test/distributed/test_dynamo_distributed.py
index 61186034c746f..fc4bb687e7d2c 100644
--- a/test/distributed/test_dynamo_distributed.py
+++ b/test/distributed/test_dynamo_distributed.py
@@ -4,6 +4,7 @@
 import functools
 import logging
 import random
+import re
 import unittest
 from contextlib import contextmanager
 from datetime import timedelta
@@ -16,6 +17,7 @@
 import torch._dynamo
 import torch._dynamo.logging
 import torch._dynamo.test_case
+import torch._inductor.test_case
 import torch.distributed as dist
 import torch.optim as optim
 from torch import nn
@@ -23,7 +25,8 @@
 from torch._dynamo import config
 from torch._dynamo.backends.distributed import DDPOptimizer
 from torch._dynamo.comptime import comptime
-from torch._dynamo.testing import collect_results
+from torch._dynamo.device_interface import CudaInterface, DeviceGuard
+from torch._dynamo.testing import collect_results, CompileCounter
 from torch._dynamo.utils import same
 from torch._higher_order_ops.wrap import tag_activation_checkpoint
 from torch.compiler import set_enable_guard_collectives
@@ -38,7 +41,9 @@
 from torch.testing._internal.common_cuda import (
     PLATFORM_SUPPORTS_FLASH_ATTENTION,
     PLATFORM_SUPPORTS_MEM_EFF_ATTENTION,
+    TEST_MULTIGPU,
 )
+from torch.testing._internal.common_device_type import instantiate_device_type_tests
 from torch.testing._internal.common_distributed import (
     _dynamo_dist_per_rank_init,
     DynamoDistributedMultiProcTestCase,
@@ -47,11 +52,27 @@
     requires_accelerator_dist_backend,
     skip_if_lt_x_gpu,
 )
-from torch.testing._internal.common_utils import skipIfXpu
+from torch.testing._internal.common_utils import MI350_ARCH, skipIfRocmArch, skipIfXpu
 from torch.testing._internal.inductor_utils import HAS_GPU
 from torch.testing._internal.triton_utils import requires_cuda_and_triton
 
 
+try:
+    from importlib.metadata import version as pkg_version
+
+    _transformers_version = tuple(
+        int(x) for x in pkg_version("transformers").split(".")[:2]
+    )
+except Exception:
+    _transformers_version = (0, 0)
+
+# https://github.com/huggingface/transformers/issues/44188
+# expectedFailure only applies to transformers >= 5.2 which introduced the bug
+_expectedFailureIf_transformers_ge_5_2 = (
+    unittest.expectedFailure if _transformers_version >= (5, 2) else lambda fn: fn
+)
+
+
 log = logging.getLogger(__name__)
 
 
@@ -316,11 +337,12 @@ def __init__(self, module, bucket_cap_mb=25):
 
     @contextmanager
     def _inside_ddp_forward(self):
+        old = DDP._active_ddp_module
         DDP._active_ddp_module = self
         try:
             yield
         finally:
-            DDP._active_ddp_module = None
+            DDP._active_ddp_module = old
 
     def forward(self, *inputs, **kwargs):
         if not DDP._active_ddp_module:
@@ -351,6 +373,7 @@ def run_hf_bert_ddp(self, model, inputs, backend):
 
 
 class TestFakeDistributedSingleProc(torch._dynamo.test_case.TestCase):
+    @_expectedFailureIf_transformers_ge_5_2
     @unittest.skipIf(not HAS_GPU, "Inductor+gpu needs triton and recent GPU arch")
     @patch.object(config, "optimize_ddp", True)
     @patch.object(torch._inductor.config, "fallback_random", True)
@@ -363,6 +386,7 @@ def test_hf_bert_ddp_inductor(self):
         model = FakeDDP(model)
         run_hf_bert_ddp(self, model, inputs, "inductor")
 
+    @_expectedFailureIf_transformers_ge_5_2
     @patch.object(config, "optimize_ddp", True)
     def test_hf_bert_ddp_aot_eager(self):
         model, inputs = get_hf_bert(0)
@@ -564,6 +588,276 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
         opt_model = torch.compile(model)
         opt_model(torch.randn(2, 129, 100, 96))
 
+    @patch.object(config, "optimize_ddp", True)
+    def test_nested_ddp_preserves_active_ddp_module(self):
+        """
+        When an inner DDP (e.g., TorchRec's data-parallel embedding lookup) is
+        nested inside an outer DDP, the inner DDP's _inside_ddp_forward exit
+        must restore _active_ddp_module to the outer DDP, not clear it to None.
+        Otherwise DDPOptimizer won't activate for compiled regions that run
+        after the inner DDP forward.
+        """
+
+        class InnerModule(nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.linear = nn.Linear(10, 10)
+
+            def forward(self, x):
+                return self.linear(x)
+
+        class OuterModel(nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.inner_ddp = FakeDDP(InnerModule())
+                self.weight = nn.Parameter(torch.randn(10, 10))
+
+            def forward(self, x):
+                # Inner DDP forward runs first (like DP embedding lookup)
+                x = self.inner_ddp(x)
+                # After inner DDP exits, _active_ddp_module must still be set
+                # for DDPOptimizer to split this region
+                x = x @ self.weight
+                return x
+
+        outer_model = OuterModel()
+        outer_ddp = FakeDDP(outer_model)
+
+        # Verify the context manager nesting directly
+        self.assertIsNone(DDP._active_ddp_module)
+        with outer_ddp._inside_ddp_forward():
+            self.assertIs(DDP._active_ddp_module, outer_ddp)
+            with outer_ddp.module.inner_ddp._inside_ddp_forward():
+                self.assertIs(DDP._active_ddp_module, outer_ddp.module.inner_ddp)
+            # After inner exits, outer must be restored
+            self.assertIs(DDP._active_ddp_module, outer_ddp)
+        # After outer exits, must be None again
+        self.assertIsNone(DDP._active_ddp_module)
+
+        # Verify DDPOptimizer activates with nested DDP during compilation
+        ddp_optimizer_activated = False
+        orig_init = DDPOptimizer.__init__
+
+        def tracking_init(self_opt, *args, **kwargs):
+            nonlocal ddp_optimizer_activated
+            ddp_optimizer_activated = True
+            return orig_init(self_opt, *args, **kwargs)
+
+        with patch.object(DDPOptimizer, "__init__", tracking_init):
+            opt_model = torch.compile(outer_ddp, backend="aot_eager")
+            opt_model(torch.randn(4, 10))
+
+        self.assertTrue(
+            ddp_optimizer_activated,
+            "DDPOptimizer should activate for compiled regions inside nested DDP",
+        )
+
+
+# These tests aren't really distributed, but need multiple GPUs to run
+class TestMultiGPU(torch._inductor.test_case.TestCase):
+    @unittest.skipIf(not TEST_MULTIGPU, "Requires multiple gpus")
+    def test_cuda__exchange_device(self):
+        def fn(x):
+            dev = torch.cuda._exchange_device(0)
+            x = torch.sin(x + dev)
+            torch.cuda._maybe_exchange_device(dev)
+            return x
+
+        initial_dev = torch.cuda.current_device()
+        x = torch.randn((2, 2), device="cuda")
+        ref = fn(x)
+        opt_fn = torch.compile(backend="eager", fullgraph=True)(fn)
+        res = opt_fn(x)
+        self.assertEqual(ref, res)
+
+        # make sure we recompile if device changes
+        with torch.cuda.device(1):
+            ref = fn(x)
+            res = opt_fn(x)
+        self.assertEqual(ref, res)
+        self.assertEqual(torch.cuda.current_device(), initial_dev)
+
+    @unittest.skipIf(not TEST_MULTIGPU, "need multiple GPU")
+    def test_device_guard(self):
+        current_device = torch.cuda.current_device()
+
+        device_guard = DeviceGuard(CudaInterface, 1)
+
+        with device_guard as _:
+            self.assertEqual(torch.cuda.current_device(), 1)
+            self.assertEqual(device_guard.prev_idx, 0)
+            self.assertEqual(device_guard.idx, 1)
+
+        self.assertEqual(torch.cuda.current_device(), current_device)
+        self.assertEqual(device_guard.prev_idx, 0)
+        self.assertEqual(device_guard.idx, 1)
+
+    @unittest.skipIf(not TEST_MULTIGPU, "need multiple GPU")
+    def test_new_event_api(self) -> None:
+        from torch._dynamo.graph_bytecode_inputs import get_external_object_by_index
+        from torch._dynamo.variables.streams import new_event
+
+        def event_generation_backend(gm, *args, **kwargs):  # type: ignore[no-untyped-def]
+            e0_ind = new_event()
+            with torch.Stream(device="cuda:1"):
+                get_external_object_by_index(e0_ind).record()
+            e1_ind = new_event()
+            self.assertNotEqual(e0_ind, e1_ind)
+            self.assertNotEqual(
+                get_external_object_by_index(e0_ind),
+                get_external_object_by_index(e1_ind),
+            )
+            with gm.graph.inserting_after(next(iter(gm.graph.nodes))):
+                gm.graph.call_function(
+                    get_external_object_by_index, args=(1,), kwargs={}
+                )
+            return gm
+
+        @torch.compile(backend=event_generation_backend)
+        def fn(x):
+            return x + 1
+
+        fn(torch.ones(2, 2, device="cuda:0"))
+
+    @unittest.skipIf(not TEST_MULTIGPU, "need multiple GPU")
+    def test_get_current_stream_return_no_index(self):
+        def fn(x, s0, s1):
+            with s1:
+                with s0:
+                    s = torch.accelerator.current_stream(torch.device("cuda"))
+            return s
+
+        s0 = torch.Stream(device="cuda:0")
+        s1 = torch.Stream(device="cuda:1")
+        inp = (torch.ones(2, 2) + 1, s0, s1)
+        fn_opt = torch.compile(fn, fullgraph=True)
+        s_act = fn_opt(*inp)
+        s_exp = fn(*inp)
+        self.assertEqual(s_act, s_exp)
+
+    @unittest.skipIf(not TEST_MULTIGPU, "need multiple GPU")
+    def test_get_current_stream_return_different_device(self):
+        def fn(x, s0, s1):
+            with s1:
+                with s0:
+                    s = torch.accelerator.current_stream(torch.device("cuda:1"))
+            return s
+
+        s0 = torch.Stream(device="cuda:0")
+        s1 = torch.Stream(device="cuda:1")
+        inp = (torch.ones(2, 2) + 1, s0, s1)
+        fn_opt = torch.compile(fn, fullgraph=True)
+        s_act = fn_opt(*inp)
+        s_exp = fn(*inp)
+        self.assertEqual(s_act, s_exp)
+
+    @unittest.skipIf(not TEST_MULTIGPU, "detected only one GPU")
+    def test_gpu_current_device(self):
+        def fn(x):
+            y = torch.empty(
+                (2, 3),
+                dtype=torch.float32,
+                device=torch.accelerator.current_device_index(),
+            )
+            y.copy_(x)
+            return torch.sin(y + y.device.index)
+
+        counter = torch._dynamo.testing.CompileCounter()
+        opt_fn = torch.compile(backend=counter, fullgraph=True)(fn)
+
+        with torch.accelerator.device_index(0):
+            x = torch.randn(2, 3)
+            self.assertEqual(opt_fn(x), fn(x))
+            self.assertEqual(counter.frame_count, 1)
+            with torch.accelerator.device_index(1):
+                self.assertEqual(opt_fn(x), fn(x))
+                self.assertEqual(counter.frame_count, 2)
+
+    @unittest.skipIf(not TEST_MULTIGPU, "need multiple GPU")
+    def test_symint_as_device_kwarg_multi_gpu(self):
+        def fn(rank):
+            # -2 to make device id smaller for easier testing on CI
+            return torch.ones(10, device=rank.size(0) - 2)
+
+        x = torch.randn(2)
+        out = fn(torch.randn(2))
+
+        guard_failure = None
+
+        def guard_failures(failure):
+            nonlocal guard_failure
+            guard_failure = failure
+
+        opt_fn = torch._dynamo.optimize(
+            "eager", guard_fail_fn=guard_failures, dynamic=True
+        )(fn)
+        self.assertEqual(out, opt_fn(x))
+
+        x = torch.randn(3)
+        self.assertEqual(fn(x), opt_fn(x))
+        self.assertTrue(guard_failure is not None)
+        self.assertIn("""tensor 'rank' size mismatch at index 0""", guard_failure[0])
+
+    def _clear_dynamo_and_codecache(self):
+        """
+        Clear unrelated caches, like dynamo and PyCodeCache
+        """
+        torch._dynamo.reset()
+        torch._inductor.codecache.PyCodeCache.cache_clear(purge=True)
+
+    @torch._inductor.config.patch("fx_graph_cache", True)
+    @torch._inductor.config.patch("fx_graph_remote_cache", False)
+    @torch._functorch.config.patch({"enable_autograd_cache": True})
+    @unittest.skipIf(not TEST_MULTIGPU, "only one GPU detected")
+    def test_constant_tensor_device_guards(self):
+        """
+        Usually, when there are example inputs, the device index of the inputs
+        is sufficient to make sure we don't cache hit with the results from different
+        cuda devices.
+        When the input has no arguments, we still need to have the cuda
+        device index in the cache key.
+        """
+
+        @torch.compile(backend="inductor")
+        def f():
+            y = torch.tensor([5], device="cuda")
+            return (y,)
+
+        with torch.cuda._DeviceGuard(0):
+            torch.cuda.set_device(0)
+            result = f()
+            self.assertEqual(result[0].device, torch.device("cuda:0"))
+
+        self._clear_dynamo_and_codecache()
+
+        with torch.cuda._DeviceGuard(1):
+            torch.cuda.set_device(1)
+            result = f()
+            self.assertEqual(result[0].device, torch.device("cuda:1"))
+
+
+class TestMultiGPUDevice(torch._inductor.test_case.TestCase):
+    @unittest.skipIf(not TEST_MULTIGPU, "need multiple GPU")
+    def test_gpu_set_device(self, device):
+        def fn():
+            a = torch.ones(2, device=device)
+            torch.get_device_module(device).set_device(1)
+            return a + 1
+
+        with torch.get_device_module(device).device(0):
+            counter = CompileCounter()
+            opt_fn = torch.compile(fn, backend=counter)
+            res = opt_fn()
+            self.assertTrue(res.device.type in device)
+            self.assertEqual(res.device.index, 0)
+            self.assertEqual(counter.frame_count, 2)
+
+
+devices = ("cuda", "hpu", "xpu")
+instantiate_device_type_tests(
+    TestMultiGPUDevice, globals(), only_for=devices, allow_xpu=True
+)
+
 
 # Are these tests failing?  Check and see if TestFakeDistributedSingleProc has a
 # single process version; if it's just a problem in the Dynamo distributed
@@ -597,6 +891,7 @@ def _test_hf_bert_ddp_inductor(self, static_graph):
             model = DDP(model, static_graph=static_graph)
             run_hf_bert_ddp(self, model, inputs, "inductor")
 
+    @_expectedFailureIf_transformers_ge_5_2
     @skip_if_lt_x_gpu(2)
     @import_transformers_or_skip()
     @unittest.skipIf(not HAS_GPU, "Inductor+gpu needs triton and recent GPU arch")
@@ -605,6 +900,7 @@ def _test_hf_bert_ddp_inductor(self, static_graph):
     def test_hf_bert_ddp_inductor(self):
         self._test_hf_bert_ddp_inductor(static_graph=False)
 
+    @_expectedFailureIf_transformers_ge_5_2
     @skip_if_lt_x_gpu(2)
     @import_transformers_or_skip()
     @unittest.skipIf(not HAS_GPU, "Inductor+gpu needs triton and recent GPU arch")
@@ -619,12 +915,14 @@ def _test_hf_bert_aot_eager(self, static_graph):
             model = DDP(model, static_graph=static_graph)
             run_hf_bert_ddp(self, model, inputs, "aot_eager")
 
+    @_expectedFailureIf_transformers_ge_5_2
     @skip_if_lt_x_gpu(2)
     @import_transformers_or_skip()
     @config.patch(optimize_ddp=True, enable_compiler_collectives=True)
     def test_hf_bert_ddp_aot_eager(self):
         self._test_hf_bert_aot_eager(static_graph=False)
 
+    @_expectedFailureIf_transformers_ge_5_2
     @skip_if_lt_x_gpu(2)
     @import_transformers_or_skip()
     @config.patch(optimize_ddp=True, enable_compiler_collectives=True)
@@ -763,23 +1061,6 @@ def test_fsdp_setattr(self):
             first_graph_break = list(counters["graph_break"].keys())[0]  # noqa: RUF015
             self.assertIn("setattr() on Tensor.requires_grad", first_graph_break)
 
-    @config.patch(inline_inbuilt_nn_modules=False)
-    @config.patch(enable_compiler_collectives=True)
-    @skip_if_lt_x_gpu(1)
-    def test_fsdp_unspecialized_forced_getattr_no_inline(self):
-        with _dynamo_dist_per_rank_init(self.rank, self.world_size):
-            # Test with basic FSDP wrapping (outer wrap around whole model)
-            from torch._dynamo.utils import counters
-
-            counters.clear()
-            m, inputs, correct_outputs = get_forced_getattr_module(
-                f"{self.device_type}:{self.rank}"
-            )
-            fsdp_m = FSDP(m, use_orig_params=True)
-            fsdp_m = torch.compile(fsdp_m, backend="eager", fullgraph=False)
-            outputs = fsdp_m(inputs)
-            self.assertTrue(same(correct_outputs, outputs))
-
     @config.patch(enable_compiler_collectives=True)
     @skip_if_lt_x_gpu(1)
     def test_fsdp_unspecialized_forced_getattr_inline(self):
@@ -796,6 +1077,7 @@ def test_fsdp_unspecialized_forced_getattr_inline(self):
             outputs = fsdp_m(inputs)
             self.assertTrue(same(correct_outputs, outputs))
 
+    @skipIfRocmArch(MI350_ARCH)  # regression in ROCm 7.2
     @config.patch(enable_compiler_collectives=True)
     @skip_if_lt_x_gpu(1)
     @unittest.skipIf(not HAS_GPU, "Inductor+gpu needs triton and recent GPU arch")
@@ -843,6 +1125,7 @@ def test_fsdp_activation_checkpointing(self):
                 find_first_node(cnt.graphs[0], tag_activation_checkpoint) is not None
             )
 
+    @_expectedFailureIf_transformers_ge_5_2
     @import_transformers_or_skip()
     @unittest.skipIf(not HAS_GPU, "Inductor+gpu needs triton and recent GPU arch")
     # TODO(whc) Investigate why cudagraphs breaks inductor+fsdp for hf_bert
@@ -888,6 +1171,7 @@ def apply_fsdp(model, wrap_policy):
                 )
                 self.assertTrue(same(correct_results, opt_results))
 
+    @_expectedFailureIf_transformers_ge_5_2
     @import_transformers_or_skip()
     @unittest.skipIf(not HAS_GPU, "Inductor+gpu needs triton and recent GPU arch")
     # TODO(whc) Investigate why cudagraphs breaks inductor+fsdp for hf_bert
@@ -1173,11 +1457,13 @@ def f(x):
             f(x)
             """
 
+            # Ranks take different code paths (tensor vs int), so compilation counts
+            # can differ. Check that every rank triggered compilation at least once.
             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)
+            for r in res:
+                self.assertGreaterEqual(r, 1, "each rank should compile at least once")
 
     @unittest.skipIf(not HAS_GPU, "Inductor+gpu needs triton and recent GPU arch")
     @enable_guard_collectives()
@@ -1479,6 +1765,7 @@ def test_ddp_baseline_aot_eager(self):
         outputs = ddp_m(inputs)
         self.assertTrue(same(correct_outputs, outputs))
 
+    @skipIfRocmArch(MI350_ARCH)  # regression in ROCm 7.2
     @unittest.skipIf(not HAS_GPU, "Inductor+gpu needs triton and recent GPU arch")
     @patch.object(config, "optimize_ddp", False)
     def test_ddp_baseline_inductor(self):
@@ -1492,7 +1779,8 @@ def test_ddp_baseline_inductor(self):
 
     @patch.object(config, "optimize_ddp", True)
     def test_graph_split(self):
-        assert config.optimize_ddp
+        if not config.optimize_ddp:
+            raise AssertionError("Expected config.optimize_ddp to be True")
         """
         Just ensures that the appropriate number of splits happen (based on
         bucket size and model parameters) - verifies the number of times
@@ -1684,10 +1972,12 @@ def alibi_score_mod(self, score, b, h, q_idx, kv_idx):
         model(hidden_states)
         torch.accelerator.synchronize()
 
+    @skipIfRocmArch(MI350_ARCH)  # regression in ROCm 7.2
     @patch.object(config, "optimize_ddp", True)
     @unittest.skipIf(not HAS_GPU, "Inductor+gpu needs triton and recent GPU arch")
     def test_graph_split_inductor(self):
-        assert config.optimize_ddp
+        if not config.optimize_ddp:
+            raise AssertionError("Expected config.optimize_ddp to be True")
         """
         Same as above, but using inductor backend.
         We observed issues with inductor/fx interface in the past.
@@ -1707,7 +1997,8 @@ def opt_fn(inputs):
     )
     @patch.object(config, "optimize_ddp", True)
     def _test_graph_split_inductor_layout_optimizations_impl(self, context):
-        assert config.optimize_ddp
+        if not config.optimize_ddp:
+            raise AssertionError("Expected config.optimize_ddp to be True")
         channel_dim = 512
         # channel dim must be > 64 for inductor to do layout optimization and use NHWC
 
@@ -1767,7 +2058,8 @@ def test_graph_split_inductor_layout_optimizations_inference(self):
     @patch.object(config, "optimize_ddp", True)
     @unittest.skipIf(not HAS_GPU, "Inductor+gpu needs triton and recent GPU arch")
     def test_graph_split_inductor_transpose(self):
-        assert config.optimize_ddp
+        if not config.optimize_ddp:
+            raise AssertionError("Expected config.optimize_ddp to be True")
 
         B = 100
         N = 30
@@ -2156,7 +2448,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
             fsdp_model(inp)
         # Check for no recompiles (if there were incorrect de-dup guards, then
         # the frame count would be equal to the number of forward calls)
-        self.assertEqual(cnt.frame_count, 1)
+        self.assertEqual(cnt.frame_count, 3)
 
     def test_fsdp_staticmethod(self):
         """
@@ -2202,9 +2494,7 @@ def _add(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
             test_outs.append(fsdp_model(x))
             # Check for no recompiles, which could happen if incorrectly
             # passing args to the staticmethod (e.g. doubly passing `self`)
-            # 3 is expected here for 1 forward.
-            # Graph 1 should be add and imul
-            self.assertEqual(cnt.frame_count, 1)
+            self.assertEqual(cnt.frame_count, 2)
         for test_out in test_outs:
             self.assertEqual(test_out, ref_out)
 
@@ -2220,6 +2510,242 @@ def f(x):
 
         self.assertEqual(cnt.frame_count, 1)
 
+    def test_compiled_all_reduce_returns_none(self):
+        def fn(x, w):
+            result = dist.all_reduce(x, async_op=False)
+            assert result is None  # noqa: S101
+            return x @ w
+
+        x = torch.randn(4, 4, device=self.device)
+        w = torch.randn(4, 4, device=self.device)
+        expected = x @ w
+        compiled_fn = torch.compile(fn, backend="aot_eager", fullgraph=True)
+        actual = compiled_fn(x, w)
+        self.assertEqual(actual, expected)
+
+    def test_compiled_all_gather_into_tensor_returns_none(self):
+        def fn(output, input, w):
+            result = dist.all_gather_into_tensor(output, input, async_op=False)
+            assert result is None  # noqa: S101
+            return output @ w
+
+        input = torch.randn(4, 4, device=self.device)
+        output = torch.empty(4, 4, device=self.device)
+        w = torch.randn(4, 4, device=self.device)
+        compiled_fn = torch.compile(fn, backend="aot_eager", fullgraph=True)
+        actual = compiled_fn(output, input, w)
+        self.assertEqual(actual, input @ w)
+
+    def test_compiled_reduce_scatter_tensor_returns_none(self):
+        def fn(output, input, w):
+            result = dist.reduce_scatter_tensor(output, input, async_op=False)
+            assert result is None  # noqa: S101
+            return output @ w
+
+        input = torch.randn(4, 4, device=self.device)
+        output = torch.empty(4, 4, device=self.device)
+        w = torch.randn(4, 4, device=self.device)
+        compiled_fn = torch.compile(fn, backend="aot_eager", fullgraph=True)
+        actual = compiled_fn(output, input, w)
+        self.assertEqual(actual, input @ w)
+
+    def test_compiled_all_to_all_single_returns_none(self):
+        def fn(output, input, w):
+            result = dist.all_to_all_single(output, input, async_op=False)
+            assert result is None  # noqa: S101
+            return output @ w
+
+        input = torch.randn(4, 4, device=self.device)
+        output = torch.empty(4, 4, device=self.device)
+        w = torch.randn(4, 4, device=self.device)
+        compiled_fn = torch.compile(fn, backend="aot_eager", fullgraph=True)
+        actual = compiled_fn(output, input, w)
+        self.assertEqual(actual, input @ w)
+
+    def test_compiled_all_gather_returns_none(self):
+        def fn(tensor_list, tensor, w):
+            result = dist.all_gather(tensor_list, tensor, async_op=False)
+            assert result is None  # noqa: S101
+            return tensor_list[0] @ w
+
+        tensor = torch.randn(4, 4, device=self.device)
+        tensor_list = [torch.empty(4, 4, device=self.device)]
+        w = torch.randn(4, 4, device=self.device)
+        compiled_fn = torch.compile(fn, backend="aot_eager", fullgraph=True)
+        actual = compiled_fn(tensor_list, tensor, w)
+        self.assertEqual(actual, tensor @ w)
+
+    def test_compiled_reduce_scatter_base_returns_none(self):
+        def fn(output, input, w):
+            result = dist._reduce_scatter_base(output, input, async_op=False)
+            assert result is None  # noqa: S101
+            return output @ w
+
+        input = torch.randn(4, 4, device=self.device)
+        output = torch.empty(4, 4, device=self.device)
+        w = torch.randn(4, 4, device=self.device)
+        compiled_fn = torch.compile(fn, backend="aot_eager", fullgraph=True)
+        actual = compiled_fn(output, input, w)
+        self.assertEqual(actual, input @ w)
+
+    def test_compiled_all_gather_base_returns_none(self):
+        def fn(output, input, w):
+            result = dist._all_gather_base(output, input, async_op=False)
+            assert result is None  # noqa: S101
+            return output @ w
+
+        input = torch.randn(4, 4, device=self.device)
+        output = torch.empty(4, 4, device=self.device)
+        w = torch.randn(4, 4, device=self.device)
+        compiled_fn = torch.compile(fn, backend="aot_eager", fullgraph=True)
+        actual = compiled_fn(output, input, w)
+        self.assertEqual(actual, input @ w)
+
+        input = torch.ones(4, device=self.device)
+        output = torch.empty(4, device=self.device)
+        w = torch.ones(4, device=self.device)
+        compiled_fn = torch.compile(fn, backend="aot_eager", fullgraph=True)
+        actual = compiled_fn(output, input, w)
+        self.assertEqual(actual, input @ w)
+
+
+@unittest.skipIf(not dist.is_available(), "requires distributed")
+class TestNNFunctionalCompile(torch._dynamo.test_case.TestCase):
+    """Tests for torch.distributed.nn.functional under torch.compile."""
+
+    def setUp(self):
+        super().setUp()
+        dist.init_process_group(backend="fake", rank=0, world_size=2)
+
+    def tearDown(self):
+        dist.destroy_process_group()
+        super().tearDown()
+
+    def _assert_not_supported(self, fn, x, name, *, suggestion=None):
+        expected = f"torch.distributed.nn.functional.{name} is not supported under torch.compile."
+        if suggestion:
+            expected = re.escape(expected + f" Use {suggestion} instead.")
+        with self.assertRaisesRegex(RuntimeError, expected):
+            fn(x)
+
+    def test_nn_functional_all_reduce_unsupported(self):
+        from torch.distributed.nn.functional import all_reduce
+
+        @torch.compile(fullgraph=True)
+        def fn(x):
+            return all_reduce(x, op=dist.ReduceOp.SUM)
+
+        self._assert_not_supported(
+            fn,
+            torch.randn(4, 4),
+            "all_reduce",
+            suggestion="torch.distributed._functional_collectives.all_reduce",
+        )
+
+    def test_nn_functional_all_gather_unsupported(self):
+        from torch.distributed.nn.functional import all_gather
+
+        @torch.compile(fullgraph=True)
+        def fn(x):
+            return all_gather(x)
+
+        self._assert_not_supported(
+            fn,
+            torch.randn(4, 4),
+            "all_gather",
+            suggestion="torch.distributed._functional_collectives.all_gather_tensor",
+        )
+
+    def test_nn_functional_reduce_scatter_unsupported(self):
+        from torch.distributed.nn.functional import reduce_scatter
+
+        @torch.compile(fullgraph=True)
+        def fn(input_list):
+            output = torch.empty(4, 4)
+            return reduce_scatter(output, input_list)
+
+        self._assert_not_supported(
+            fn,
+            [torch.randn(4, 4), torch.randn(4, 4)],
+            "reduce_scatter",
+            suggestion="torch.distributed._functional_collectives.reduce_scatter_tensor",
+        )
+
+    def test_nn_functional_all_to_all_single_unsupported(self):
+        from torch.distributed.nn.functional import all_to_all_single
+
+        @torch.compile(fullgraph=True)
+        def fn(x):
+            output = torch.empty_like(x)
+            return all_to_all_single(output, x)
+
+        self._assert_not_supported(
+            fn,
+            torch.randn(4, 4),
+            "all_to_all_single",
+            suggestion="torch.distributed._functional_collectives.all_to_all_single",
+        )
+
+    def test_nn_functional_broadcast_unsupported(self):
+        from torch.distributed.nn.functional import broadcast
+
+        @torch.compile(fullgraph=True)
+        def fn(x):
+            return broadcast(x, src=0)
+
+        self._assert_not_supported(
+            fn,
+            torch.randn(4),
+            "broadcast",
+            suggestion="torch.distributed._functional_collectives.broadcast",
+        )
+
+    def test_nn_functional_reduce_unsupported(self):
+        from torch.distributed.nn.functional import reduce
+
+        @torch.compile(fullgraph=True)
+        def fn(x):
+            return reduce(x, dst=0)
+
+        self._assert_not_supported(fn, torch.randn(4), "reduce")
+
+    def test_nn_functional_gather_unsupported(self):
+        from torch.distributed.nn.functional import gather
+
+        @torch.compile(fullgraph=True)
+        def fn(x):
+            return gather(x, dst=0)
+
+        self._assert_not_supported(fn, torch.randn(4), "gather")
+
+    def test_nn_functional_scatter_unsupported(self):
+        from torch.distributed.nn.functional import scatter
+
+        @torch.compile(fullgraph=True)
+        def fn(x):
+            return scatter([x, x], src=0)
+
+        self._assert_not_supported(fn, torch.randn(4), "scatter")
+
+    def test_nn_functional_all_to_all_unsupported(self):
+        from torch.distributed.nn.functional import all_to_all
+
+        @torch.compile(fullgraph=True)
+        def fn(x):
+            return all_to_all([torch.empty_like(x)], [x])
+
+        self._assert_not_supported(fn, torch.randn(4), "all_to_all")
+
+    def test_nn_functional_all_gather_base_unsupported(self):
+        from torch.distributed.nn.functional import _all_gather_base
+
+        @torch.compile(fullgraph=True)
+        def fn(x):
+            output = torch.empty(8)
+            return _all_gather_base(output, x)
+
+        self._assert_not_supported(fn, torch.randn(4), "_all_gather_base")
+
 
 if __name__ == "__main__":
     from torch._dynamo.test_case import run_tests
diff --git a/test/distributed/test_functional_api.py b/test/distributed/test_functional_api.py
index 59c47341f805f..37a054645e112 100644
--- a/test/distributed/test_functional_api.py
+++ b/test/distributed/test_functional_api.py
@@ -467,6 +467,42 @@ def allred_mesh_dim(input):
         )
 
 
+class TestAllGatherViewOptimization(TestCase):
+    """Validate that all_gather_tensor delays wait() when the view optimization
+    applies and calls wait() early when the fallback path is needed."""
+
+    def setUp(self):
+        self.world_size = 4
+        if dist.is_initialized():
+            dist.destroy_process_group()
+        dist.init_process_group(
+            backend="fake",
+            world_size=self.world_size,
+            rank=0,
+        )
+
+    def tearDown(self):
+        dist.destroy_process_group()
+        super().tearDown()
+
+    def test_view_path_delays_wait(self):
+        # Shape [1, 8] with gather_dim=1: after all_gather -> [4, 8]
+        # numel_between = prod(shape[1:1]) = 1, so view optimization applies.
+        # Result should be an AsyncCollectiveTensor (wait delayed).
+        t = torch.randn(1, 8)
+        res = ft_c.all_gather_tensor(t, gather_dim=1, group=dist.group.WORLD)
+        self.assertIsInstance(res, ft_c.AsyncCollectiveTensor)
+
+    def test_fallback_path_waits_early(self):
+        # Shape [2, 8] with gather_dim=1: after all_gather -> [8, 8]
+        # numel_between = prod(shape[1:1]) = 1 but shape[0]=8 != group_size=4,
+        # so view optimization does NOT apply.
+        # Result should be a plain tensor (wait called early).
+        t = torch.randn(2, 8)
+        res = ft_c.all_gather_tensor(t, gather_dim=1, group=dist.group.WORLD)
+        self.assertNotIsInstance(res, ft_c.AsyncCollectiveTensor)
+
+
 BACKEND = dist.Backend.NCCL if torch.cuda.is_available() else dist.Backend.GLOO
 
 # Adding support for HCCL backend
@@ -676,7 +712,8 @@ def test_all_to_all_single(self, compile: bool = True) -> None:
         def my_func(t: torch.Tensor, world_size: int) -> torch.Tensor:
             sizes = [1] * world_size
             t = t * 2
-            assert t.requires_grad
+            if not t.requires_grad:
+                raise AssertionError("Expected t.requires_grad to be True")
             out = ft_c.all_to_all_single_autograd(t, sizes, sizes, group)
             out = out + 0
             return out
@@ -703,7 +740,8 @@ def test_all_to_all_single_inductor(self) -> None:
         def my_func(t: torch.Tensor, world_size: int) -> torch.Tensor:
             sizes = [1] * world_size
             t = t * 10
-            assert t.requires_grad
+            if not t.requires_grad:
+                raise AssertionError("Expected t.requires_grad to be True")
             out = ft_c.all_to_all_single_autograd(t, sizes, sizes, group)
             out = out + 2
             return out.sum()
@@ -736,7 +774,8 @@ def test_all_gather_tensor(self, compile: bool) -> None:
         group = dist.group.WORLD.group_name
 
         def my_func(t: torch.Tensor, dim: int) -> torch.Tensor:
-            assert t.requires_grad
+            if not t.requires_grad:
+                raise AssertionError("Expected t.requires_grad to be True")
             out = ft_c.all_gather_tensor_autograd(
                 t * 1.0,
                 gather_dim=dim,
@@ -770,7 +809,8 @@ def test_reduce_scatter_tensor(self, compile: bool) -> None:
         group = dist.group.WORLD.group_name
 
         def my_func(t: torch.Tensor, dim: int) -> torch.Tensor:
-            assert t.requires_grad
+            if not t.requires_grad:
+                raise AssertionError("Expected t.requires_grad to be True")
             rs_tensor = (
                 ft_c.reduce_scatter_tensor_autograd(
                     input_tensor * 1.0, "sum", scatter_dim=dim, group=group
@@ -805,7 +845,8 @@ def test_all_to_all_single(self, device) -> None:
         t = torch.ones((self.world_size, 2), requires_grad=True, device=device)
 
         sizes = [1] * self.world_size
-        assert t.requires_grad
+        if not t.requires_grad:
+            raise AssertionError("Expected t.requires_grad to be True")
         out = ft_c.all_to_all_single_autograd(t * 2, sizes, sizes, group) + 0
 
         self.assertEqual(out.shape, t.shape)
diff --git a/test/distributed/test_inductor_collectives.py b/test/distributed/test_inductor_collectives.py
index 689cb2b4e0da9..2339a33c99b07 100644
--- a/test/distributed/test_inductor_collectives.py
+++ b/test/distributed/test_inductor_collectives.py
@@ -3,7 +3,6 @@
 import functools
 import unittest
 from collections import Counter
-from typing import Optional
 from unittest import mock
 from unittest.mock import patch
 
@@ -169,7 +168,7 @@ def eager_func(x):
 
         def func(x):
             y = x * x
-            y = dist.all_reduce(y, op=dist.ReduceOp.SUM)
+            dist.all_reduce(y, op=dist.ReduceOp.SUM)
             x = torch.nn.functional.silu(x)
             return x * y
 
@@ -1549,7 +1548,7 @@ def func(inp, *, tag, ranks, group_size):
         inputs = torch.ones(4, 4, device=self.device)
 
         # get stats directly from the internal helper without affecting the real pass's signature
-        node_stats: Optional[dict[BaseSchedulerNode, ReorderInfo]] = None
+        node_stats: dict[BaseSchedulerNode, ReorderInfo] | None = None
 
         def _reorder_communication_preserving_peak_memory(
             snodes: list[BaseSchedulerNode],
@@ -1609,7 +1608,7 @@ def _reorder_communication_preserving_peak_memory(
 
     @unittest.skipIf(not HAS_GPU, "Inductor+gpu needs triton and recent GPU arch")
     @unittest.skipIf(not SM80OrLater, "bfloat16")
-    @parametrize("bucket_mode", ["all", "all_custom_ops"])
+    @parametrize("bucket_mode", ["default", "custom_ops"])
     def test_all_gather_bucket(self, bucket_mode):
         def func(x, w, ag_0, ag_1, ag_2, ag_3, *, tag, ranks, group_size):
             # do some unrelated matmuls
@@ -1658,7 +1657,7 @@ def func(x, w, ag_0, ag_1, ag_2, ag_3, *, tag, ranks, group_size):
         with (
             torch._inductor.config.patch(
                 {
-                    "bucket_all_gathers_fx": bucket_mode,
+                    "bucket_all_gathers_bucket_mode": bucket_mode,
                     "reorder_for_compute_comm_overlap": False,
                     "runtime_estimations_mms_benchmark": True,
                 }
@@ -1678,7 +1677,7 @@ def func(x, w, ag_0, ag_1, ag_2, ag_3, *, tag, ranks, group_size):
         (
             FileCheck()
             .check("= torch.ops._c10d_functional.all_gather_into_tensor")
-            .check("torch.ops._c10d_functional.all_gather_into_tensor_out.default(")
+            .check_not("torch.ops._c10d_functional.all_gather_into_tensor_out.default(")
             .check("= torch.ops._c10d_functional.all_gather_into_tensor")
             .run(code)
         )
@@ -1686,6 +1685,48 @@ def func(x, w, ag_0, ag_1, ag_2, ag_3, *, tag, ranks, group_size):
         if not same(out, correct):
             raise AssertionError(f"Expected out to match correct: {out} vs {correct}")
 
+    @unittest.skipIf(not HAS_GPU, "Inductor+gpu needs triton and recent GPU arch")
+    def test_all_gather_bucket_copy_cat_fusion(self):
+        """Bucketed all_gather merge uses copy_(cat(...)) which inductor fuses
+        into 1 Triton kernel instead of N kernels from _foreach_copy_."""
+
+        def func(ag_0, ag_1, ag_2, *, tag, ranks, group_size):
+            group_name = (
+                torch.distributed.distributed_c10d._get_default_group().group_name
+            )
+            ag_0_out = torch.ops._c10d_functional.all_gather_into_tensor(
+                ag_0, group_size, group_name
+            )
+            ag_1_out = torch.ops._c10d_functional.all_gather_into_tensor(
+                ag_1, group_size, group_name
+            )
+            ag_2_out = torch.ops._c10d_functional.all_gather_into_tensor(
+                ag_2, group_size, group_name
+            )
+            return (
+                torch.ops.c10d_functional.wait_tensor(ag_0_out),
+                torch.ops.c10d_functional.wait_tensor(ag_1_out),
+                torch.ops.c10d_functional.wait_tensor(ag_2_out),
+            )
+
+        inputs = [torch.ones(64, device="cuda") for _ in range(3)]
+        with torch._inductor.config.patch(
+            {
+                "bucket_all_gathers_fx": "all",
+                "reorder_for_compute_comm_overlap": False,
+            }
+        ):
+            compiled = torch.compile(func)
+            code = run_and_get_triton_code(compiled, *inputs, **self.get_world_trs())
+        # Bucketed merge should produce 1 copy kernel (fused copy_(cat(...))),
+        # not 3 separate kernels from _foreach_copy_.
+        num_triton_kernels = code.count("def triton_")
+        self.assertEqual(
+            num_triton_kernels,
+            1,
+            f"Expected 1 Triton kernel for fused copy_(cat(...)), got {num_triton_kernels}",
+        )
+
     @unittest.skipIf(not HAS_GPU, "Inductor+gpu needs triton and recent GPU arch")
     @unittest.skipIf(not SM80OrLater, "bfloat16")
     def test_all_gather_bucket_path(self):
@@ -1741,7 +1782,7 @@ def func(x, w, ag_0, ag_1, *, tag, ranks, group_size):
 
     @unittest.skipIf(not HAS_GPU, "Inductor+gpu needs triton and recent GPU arch")
     @unittest.skipIf(not SM80OrLater, "bfloat16")
-    @parametrize("bucket_mode", ["all", "all_custom_ops"])
+    @parametrize("bucket_mode", ["default", "custom_ops"])
     def test_reduce_scatter_bucket(self, bucket_mode):
         def func(x, w, rs_0, rs_1, tag, ranks, group_size):
             # do some unrelated matmuls
@@ -1783,7 +1824,8 @@ def func2(x, w, rs_0, rs_1, tag, ranks, group_size):
 
             with torch._inductor.config.patch(
                 {
-                    "bucket_reduce_scatters_fx": bucket_mode,
+                    "bucket_reduce_scatters_fx": "all",
+                    "bucket_reduce_scatters_bucket_mode": bucket_mode,
                     "reorder_for_compute_comm_overlap": False,
                 }
             ):
@@ -1812,7 +1854,9 @@ def func2(x, w, rs_0, rs_1, tag, ranks, group_size):
 
     @unittest.skipIf(not HAS_GPU, "Inductor+gpu needs triton and recent GPU arch")
     @unittest.skipIf(not SM80OrLater, "bfloat16")
-    @parametrize("bucket_mode", ["all"])
+    @parametrize(
+        "bucket_mode", ["all"]
+    )  # "all" is just a placeholder, there is only one bucketmode
     def test_all_reduce_bucket(self, bucket_mode):
         def func(x, w, ar_0, ar_1, tag, ranks, group_size):
             y = torch.mm(x, w)
@@ -1868,7 +1912,7 @@ def func(x, w, ar_0, ar_1, tag, ranks, group_size):
 
     @unittest.skipIf(not HAS_GPU, "Inductor+gpu needs triton and recent GPU arch")
     @unittest.skipIf(not SM80OrLater, "bfloat16")
-    @parametrize("bucket_mode", ["all_custom_ops_multidtype"])
+    @parametrize("bucket_mode", ["custom_ops_multidtype"])
     def test_all_gather_bucket_multidtype(self, bucket_mode):
         def func(x, w, ag_0, ag_1, *, tag, ranks, group_size):
             # do some unrelated matmuls
@@ -1901,7 +1945,8 @@ def func(x, w, ag_0, ag_1, *, tag, ranks, group_size):
 
         with torch._inductor.config.patch(
             {
-                "bucket_all_gathers_fx": bucket_mode,
+                "bucket_all_gathers_fx": "all",
+                "bucket_all_gathers_bucket_mode": bucket_mode,
                 "reorder_for_compute_comm_overlap": False,
             }
         ):
@@ -1932,7 +1977,7 @@ def func(x, w, ag_0, ag_1, *, tag, ranks, group_size):
 
     @unittest.skipIf(not HAS_GPU, "Inductor+gpu needs triton and recent GPU arch")
     @unittest.skipIf(not SM80OrLater, "bfloat16")
-    @parametrize("bucket_mode", ["all", "all_custom_ops"])
+    @parametrize("bucket_mode", ["default", "custom_ops"])
     def test_reorder_peak_memory_bucketed(self, bucket_mode):
         """
         Simulate the case where a bucketing pass ran and grouped several inputs into one bucketed allgather.
@@ -2021,7 +2066,7 @@ def func(x, w, ag_0, ag_1, ag_2, ag_3, *, tag, ranks, group_size):
         inputs = [x, w, ag_0, ag_1, ag_2, ag_3]
 
         # get stats directly from the internal helper without affecting the real pass's signature
-        node_stats: Optional[dict[BaseSchedulerNode, ReorderInfo]] = None
+        node_stats: dict[BaseSchedulerNode, ReorderInfo] | None = None
 
         def _reorder_communication_preserving_peak_memory(
             snodes: list[BaseSchedulerNode],
@@ -2053,9 +2098,11 @@ def _reorder_communication_preserving_peak_memory(
         with (
             torch._inductor.config.patch(
                 {
-                    "bucket_all_gathers_fx": bucket_mode,
+                    "bucket_all_gathers_fx": "all",
+                    "bucket_all_gathers_bucket_mode": bucket_mode,
                     "bucket_all_gathers_fx_bucket_size_determinator": lambda _: 2,
-                    "bucket_reduce_scatters_fx": bucket_mode,
+                    "bucket_reduce_scatters_fx": "all",
+                    "bucket_reduce_scatters_bucket_mode": bucket_mode,
                     "bucket_reduce_scatters_fx_bucket_size_determinator": lambda _: 2,
                     "reorder_for_compute_comm_overlap": True,
                     "reorder_for_compute_comm_overlap_passes": [
@@ -2150,7 +2197,7 @@ def func(inp, *, tag, ranks, group_size):
         inputs = torch.ones(4, 4, device=self.device)
 
         # get stats directly from the internal helper without affecting the real pass's signature
-        node_stats: Optional[dict[BaseSchedulerNode, ReorderInfo]] = None
+        node_stats: dict[BaseSchedulerNode, ReorderInfo] | None = None
 
         def _reorder_communication_preserving_peak_memory(
             snodes: list[BaseSchedulerNode],
diff --git a/test/distributed/test_local_tensor.py b/test/distributed/test_local_tensor.py
index b2c08c55ba3a0..d3bfbc4050b37 100644
--- a/test/distributed/test_local_tensor.py
+++ b/test/distributed/test_local_tensor.py
@@ -12,6 +12,8 @@
     LocalTensor,
     LocalTensorMode,
     maybe_disable_local_tensor_mode,
+    rank_map,
+    tensor_map,
 )
 from torch.distributed.tensor import (
     DeviceMesh,
@@ -31,7 +33,8 @@ 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)
+                if not (isinstance(lhs, LocalTensor) and isinstance(rhs, LocalTensor)):
+                    raise AssertionError("Expected both lhs and rhs to be LocalTensor")
                 super().assertEqual(lhs._ranks, rhs._ranks)
                 for r in lhs._ranks:
                     super().assertEqual(
@@ -85,7 +88,8 @@ def rank_loop_wrapper():
 
     @property
     def rank(self):
-        assert dist.is_initialized(), "Process group is not initialized!"
+        if not dist.is_initialized():
+            raise AssertionError("Process group is not initialized!")
         return dist.get_rank()
 
 
@@ -252,10 +256,26 @@ def test_sym_and_sym_or(self):
         self.assertFalse(bool(sym_true & False))
         self.assertTrue(bool(sym_true | False))
 
+    def test_standalone_rank_map_and_tensor_map(self):
+        with LocalTensorMode(self.world_size):
+            lt = rank_map(lambda r: torch.full((2, 3), float(r)))
+            scaled = tensor_map(lt, lambda r, t: t * 2)
+        for r in range(self.world_size):
+            self.assertEqual(lt._local_tensors[r], torch.full((2, 3), float(r)))
+            self.assertEqual(scaled._local_tensors[r], torch.full((2, 3), float(r) * 2))
+
 
 class TestLocalTensorRankWorld2(LocalTensorRankTest):
     world_size = 2
 
+    def test_standalone_rank_map_and_tensor_map(self):
+        # Real-dist path: rank_map should call cb with dist.get_rank()
+        t = rank_map(lambda r: torch.full((2, 3), float(r)))
+        self.assertNotIsInstance(t, LocalTensor)
+        self.assertEqual(t, torch.full((2, 3), float(self.rank)))
+        t2 = tensor_map(t, lambda r, t: t + 1)
+        self.assertEqual(t2, torch.full((2, 3), float(self.rank) + 1))
+
     def test_flatten_unflatten(self):
         """Test that LocalTensor can be flattened and unflattened correctly."""
         device = torch.device("cpu")
diff --git a/test/distributed/test_multi_threaded_pg.py b/test/distributed/test_multi_threaded_pg.py
index 7ca6d25ad1c97..570458f84c1e4 100644
--- a/test/distributed/test_multi_threaded_pg.py
+++ b/test/distributed/test_multi_threaded_pg.py
@@ -311,7 +311,10 @@ def forward(ctx, rank):
                 result = rank * 2
 
                 ctx.save_for_backward(result, rank)
-                assert int(rank.item()) == dist.get_rank()
+                if int(rank.item()) != dist.get_rank():
+                    raise AssertionError(
+                        f"Expected rank.item() == dist.get_rank(), got {int(rank.item())} vs {dist.get_rank()}"
+                    )
                 return result
 
             @staticmethod
diff --git a/test/distributed/test_nan_check.py b/test/distributed/test_nan_check.py
new file mode 100644
index 0000000000000..d66ec844336ad
--- /dev/null
+++ b/test/distributed/test_nan_check.py
@@ -0,0 +1,123 @@
+# Owner(s): ["oncall: distributed"]
+
+import sys
+import unittest
+
+import torch
+import torch.distributed as dist
+
+
+if not dist.is_available():
+    print("distributed package not available, skipping tests", file=sys.stderr)
+    sys.exit(0)
+
+from torch.testing._internal.common_utils import run_tests, TestCase
+
+
+HAS_ACCELERATOR = torch.accelerator.is_available()
+ACCELERATOR = (
+    torch.accelerator.current_accelerator() if HAS_ACCELERATOR else torch.device("cpu")
+)
+
+
+class TestNanCheck(TestCase):
+    def _check_for_nan(self, tensor):
+        torch.ops.c10d.check_for_nan(tensor)
+
+    def test_no_nan_float32(self):
+        self._check_for_nan(torch.randn(100))
+
+    def test_no_nan_float64(self):
+        self._check_for_nan(torch.randn(100, dtype=torch.float64))
+
+    def test_no_nan_float16(self):
+        self._check_for_nan(torch.randn(100, dtype=torch.float16))
+
+    def test_no_nan_bfloat16(self):
+        self._check_for_nan(torch.randn(100, dtype=torch.bfloat16))
+
+    @unittest.skipUnless(HAS_ACCELERATOR, "no accelerator available")
+    def test_no_nan_float32_accelerator(self):
+        self._check_for_nan(torch.randn(100, device=ACCELERATOR))
+        torch.accelerator.synchronize()
+
+    @unittest.skipUnless(HAS_ACCELERATOR, "no accelerator available")
+    def test_no_nan_float64_accelerator(self):
+        self._check_for_nan(torch.randn(100, dtype=torch.float64, device=ACCELERATOR))
+        torch.accelerator.synchronize()
+
+    @unittest.skipUnless(HAS_ACCELERATOR, "no accelerator available")
+    def test_no_nan_float16_accelerator(self):
+        self._check_for_nan(torch.randn(100, dtype=torch.float16, device=ACCELERATOR))
+        torch.accelerator.synchronize()
+
+    @unittest.skipUnless(HAS_ACCELERATOR, "no accelerator available")
+    def test_no_nan_bfloat16_accelerator(self):
+        self._check_for_nan(torch.randn(100, dtype=torch.bfloat16, device=ACCELERATOR))
+        torch.accelerator.synchronize()
+
+    @unittest.skipUnless(HAS_ACCELERATOR, "no accelerator available")
+    def test_skips_integer_tensors_accelerator(self):
+        self._check_for_nan(
+            torch.tensor([1, 2, 3], dtype=torch.int32, device=ACCELERATOR)
+        )
+        torch.accelerator.synchronize()
+
+    @unittest.skipUnless(HAS_ACCELERATOR, "no accelerator available")
+    def test_empty_tensor_accelerator(self):
+        self._check_for_nan(torch.tensor([], dtype=torch.float32, device=ACCELERATOR))
+        torch.accelerator.synchronize()
+
+    @unittest.skipUnless(HAS_ACCELERATOR, "no accelerator available")
+    def test_large_tensor_no_nan_accelerator(self):
+        self._check_for_nan(torch.randn(100000, device=ACCELERATOR))
+        torch.accelerator.synchronize()
+
+    def test_nan_float32(self):
+        tensor = torch.tensor([1.0, 2.0, float("nan"), 4.0])
+        with self.assertRaisesRegex(RuntimeError, "NaN"):
+            self._check_for_nan(tensor)
+
+    def test_nan_float64(self):
+        tensor = torch.tensor([1.0, float("nan")], dtype=torch.float64)
+        with self.assertRaisesRegex(RuntimeError, "NaN"):
+            self._check_for_nan(tensor)
+
+    def test_nan_float16(self):
+        tensor = torch.tensor([1.0, float("nan")], dtype=torch.float16)
+        with self.assertRaisesRegex(RuntimeError, "NaN"):
+            self._check_for_nan(tensor)
+
+    def test_nan_bfloat16(self):
+        tensor = torch.tensor([1.0, float("nan")], dtype=torch.bfloat16)
+        with self.assertRaisesRegex(RuntimeError, "NaN"):
+            self._check_for_nan(tensor)
+
+    def test_nan_at_end(self):
+        tensor = torch.ones(1000)
+        tensor[-1] = float("nan")
+        with self.assertRaisesRegex(RuntimeError, "NaN"):
+            self._check_for_nan(tensor)
+
+    def test_skips_integer_tensors(self):
+        self._check_for_nan(torch.tensor([1, 2, 3], dtype=torch.int32))
+
+    def test_empty_tensor(self):
+        self._check_for_nan(torch.tensor([], dtype=torch.float32))
+
+    def test_all_nan(self):
+        with self.assertRaisesRegex(RuntimeError, "NaN"):
+            self._check_for_nan(torch.full((10,), float("nan")))
+
+    def test_large_tensor_no_nan(self):
+        self._check_for_nan(torch.randn(100000))
+
+    def test_large_tensor_with_nan(self):
+        tensor = torch.randn(100000)
+        tensor[99999] = float("nan")
+        with self.assertRaisesRegex(RuntimeError, "NaN"):
+            self._check_for_nan(tensor)
+
+
+if __name__ == "__main__":
+    run_tests()
diff --git a/test/distributed/test_nccl.py b/test/distributed/test_nccl.py
index d609d6b115cbe..78db1c1aa9aa9 100644
--- a/test/distributed/test_nccl.py
+++ b/test/distributed/test_nccl.py
@@ -1,5 +1,6 @@
 # Owner(s): ["oncall: distributed"]
 
+import os
 import sys
 
 import torch
@@ -18,9 +19,11 @@
     skip_if_lt_x_gpu,
 )
 from torch.testing._internal.common_utils import (
+    instantiate_parametrized_tests,
     IS_WINDOWS,
     load_tests,
     NoTest,
+    parametrize,
     requires_cuda_p2p_access,
     run_tests,
     skip_but_pass_in_sandcastle_if,
@@ -220,6 +223,7 @@ def test_reduce_scatter(self, device, dtype):
             self.assertEqual(outputs[i], expected[i])
 
 
+@instantiate_parametrized_tests
 @requires_cuda_p2p_access()
 class NCCLSymmetricMemoryTest(MultiProcContinuousTest):
     @property
@@ -251,6 +255,37 @@ def foo():
         out = symm_mem.empty(numel, dtype=dtype, device=self.device)
         symm_mem.rendezvous(out, group=group_name)
 
+    @skip_but_pass_in_sandcastle_if(TEST_WITH_ROCM, "Skip NCCL tests for ROCm")
+    @skip_but_pass_in_sandcastle_if(IS_WINDOWS, "NCCL doesn't support Windows")
+    @requires_nccl_version((2, 27), "NCCL Symmetric Memory support from nccl 2.27")
+    @skip_if_lt_x_gpu(2)
+    def test_nccl_symmem_rendezvous_many_allocations(self):
+        symm_mem.set_backend("NCCL")
+        torch.cuda.set_device(self.rank)
+        c10d.all_reduce(torch.ones(1, device=self.device))
+        group_name = c10d.group.WORLD.group_name
+
+        tensors = [
+            symm_mem.empty(1, dtype=torch.float, device=self.device) for _ in range(256)
+        ]
+
+        # Rendezvous a subset twice so the repeated lookup path is covered
+        # while many allocations are still live.
+        sampled_tensors = tensors[::16]
+        for tensor in sampled_tensors:
+            handle = symm_mem.rendezvous(tensor, group=group_name)
+            self.assertEqual(handle.rank, self.rank)
+            self.assertEqual(handle.world_size, self.world_size)
+        for tensor in sampled_tensors:
+            symm_mem.rendezvous(tensor, group=group_name)
+
+        result = torch.ops.symm_mem.one_shot_all_reduce(
+            tensors[-1].fill_(self.rank), "sum", group_name
+        )
+        self.assertEqual(
+            result, torch.full_like(result, (self.world_size - 1) * self.world_size / 2)
+        )
+
     @skip_but_pass_in_sandcastle_if(TEST_WITH_ROCM, "Skip NCCL tests for ROCm")
     @skip_but_pass_in_sandcastle_if(IS_WINDOWS, "NCCL doesn't support Windows")
     @requires_nccl_version(
@@ -396,6 +431,154 @@ def test_nccl_symmem_get(self):
             # TODO: remove after we have wait_signal
             c10d.barrier()
 
+    @skip_but_pass_in_sandcastle_if(TEST_WITH_ROCM, "Skip NCCL tests for ROCm")
+    @skip_but_pass_in_sandcastle_if(IS_WINDOWS, "NCCL doesn't support Windows")
+    @requires_nccl_version(
+        (2, 29, 7), "nccl_reduce_scatter_offset requires nccl 2.29.7"
+    )
+    @skip_if_lt_x_gpu(2)
+    @parametrize("experts_per_rank", [1, 2])
+    @parametrize("dim", [0, 1])
+    def test_reduce_scatter_offset(self, experts_per_rank: int, dim: int):
+        """reduce_scatter_offset: each expert gradient is reduced to its
+        destination rank and written to a separate contiguous tensor; the source
+        Grouped GEMM buffer is left unmodified."""
+        symm_mem.set_backend("NCCL")
+        torch.cuda.set_device(self.rank)
+        c10d.all_reduce(torch.ones(1, device=self.device))
+        group_name = c10d.group.WORLD.group_name
+
+        rows, cols = 64, 32
+        n_experts = experts_per_rank * self.world_size
+
+        # dim=1: experts laid out as column blocks [rows, n_experts * cols]
+        # dim=0: experts laid out as row blocks    [n_experts * rows, cols]
+        if dim == 1:
+            buf = symm_mem.empty(
+                rows, n_experts * cols, dtype=torch.float, device=self.device
+            )
+            for i in range(n_experts):
+                buf[:, i * cols : (i + 1) * cols] = float((self.rank + 1) * (i + 1))
+        else:
+            buf = symm_mem.empty(
+                n_experts * rows, cols, dtype=torch.float, device=self.device
+            )
+            for i in range(n_experts):
+                buf[i * rows : (i + 1) * rows, :] = float((self.rank + 1) * (i + 1))
+        symm_mem.rendezvous(buf, group=group_name)
+
+        # Round-robin: expert i is reduced to rank i % world_size.
+        dst_ranks = [i % self.world_size for i in range(n_experts)]
+        n_owned = sum(r == self.rank for r in dst_ranks)
+        out = [
+            torch.zeros(rows, cols, dtype=torch.float, device=self.device)
+            for _ in range(n_owned)
+        ]
+        block_size = cols if dim == 1 else rows
+        offsets = [i * block_size for i in range(1, n_experts + 1)]
+
+        symm_mem.reduce_scatter_offset(
+            buf, out, group_name, dim=dim, offsets=offsets, dst_ranks=dst_ranks
+        )
+        torch.cuda.synchronize()
+
+        # out[j] corresponds to expert (rank + j * world_size); expected value is
+        # (expert_idx + 1) * sum(r + 1 for r in range(world_size)).
+        rank_sum = float(sum(r + 1 for r in range(self.world_size)))
+        for j in range(n_owned):
+            expert_idx = self.rank + j * self.world_size
+            expected = float(expert_idx + 1) * rank_sum
+            self.assertEqual(
+                out[j],
+                torch.full_like(out[j], expected),
+                msg=f"rank {self.rank}: out[{j}] should contain the reduced sum",
+            )
+        # Source buffer must be unmodified.
+        for i in range(n_experts):
+            if dim == 1:
+                src_slice = buf[:, i * cols : (i + 1) * cols]
+            else:
+                src_slice = buf[i * rows : (i + 1) * rows, :]
+            self.assertEqual(
+                src_slice,
+                torch.full(
+                    (rows, cols),
+                    float((self.rank + 1) * (i + 1)),
+                    dtype=torch.float,
+                    device=self.device,
+                ),
+                msg=f"rank {self.rank}: source buffer block {i} should be unchanged",
+            )
+
+    @skip_but_pass_in_sandcastle_if(TEST_WITH_ROCM, "Skip NCCL tests for ROCm")
+    @skip_but_pass_in_sandcastle_if(IS_WINDOWS, "NCCL doesn't support Windows")
+    @requires_nccl_version(
+        (2, 29, 7), "nccl_reduce_scatter_offset requires nccl 2.29.7"
+    )
+    @skip_if_lt_x_gpu(2)
+    @parametrize("dim", [0, 1])
+    def test_reduce_scatter_offset_uneven(self, dim: int):
+        """reduce_scatter_offset with uneven block sizes: j=0 and j=1 own blocks
+        of different sizes, verifying that out[j] shapes differ across j."""
+        symm_mem.set_backend("NCCL")
+        torch.cuda.set_device(self.rank)
+        c10d.all_reduce(torch.ones(1, device=self.device))
+        group_name = c10d.group.WORLD.group_name
+
+        rows, cols = 64, 32
+        # j=0 blocks have size_0 along dim; j=1 blocks have size_1 along dim.
+        # Arrange blocks as [size_0] * world_size + [size_1] * world_size so
+        # that round-robin assigns each rank exactly one block of each size.
+        size_0, size_1 = 16, 48
+        block_sizes = [size_0] * self.world_size + [size_1] * self.world_size
+        offsets = []
+        total = 0
+        for sz in block_sizes:
+            total += sz
+            offsets.append(total)
+
+        n_experts = 2 * self.world_size
+        if dim == 1:
+            buf = symm_mem.empty(rows, total, dtype=torch.float, device=self.device)
+            pos = 0
+            for i, sz in enumerate(block_sizes):
+                buf[:, pos : pos + sz] = float((self.rank + 1) * (i + 1))
+                pos += sz
+        else:
+            buf = symm_mem.empty(total, cols, dtype=torch.float, device=self.device)
+            pos = 0
+            for i, sz in enumerate(block_sizes):
+                buf[pos : pos + sz, :] = float((self.rank + 1) * (i + 1))
+                pos += sz
+        symm_mem.rendezvous(buf, group=group_name)
+
+        dst_ranks = [i % self.world_size for i in range(n_experts)]
+        if dim == 1:
+            out = [
+                torch.zeros(rows, size_0, dtype=torch.float, device=self.device),
+                torch.zeros(rows, size_1, dtype=torch.float, device=self.device),
+            ]
+        else:
+            out = [
+                torch.zeros(size_0, cols, dtype=torch.float, device=self.device),
+                torch.zeros(size_1, cols, dtype=torch.float, device=self.device),
+            ]
+
+        symm_mem.reduce_scatter_offset(
+            buf, out, group_name, dim=dim, offsets=offsets, dst_ranks=dst_ranks
+        )
+        torch.cuda.synchronize()
+
+        rank_sum = float(sum(r + 1 for r in range(self.world_size)))
+        for j in range(2):
+            expert_idx = self.rank + j * self.world_size
+            expected = float(expert_idx + 1) * rank_sum
+            self.assertEqual(
+                out[j],
+                torch.full_like(out[j], expected),
+                msg=f"rank {self.rank}: out[{j}] should contain the reduced sum",
+            )
+
     @skip_but_pass_in_sandcastle_if(TEST_WITH_ROCM, "Skip NCCL tests for ROCm")
     @skip_but_pass_in_sandcastle_if(IS_WINDOWS, "NCCL doesn't support Windows")
     @requires_nccl_version((2, 29), "NCCL one-sided host API support from nccl 2.29")
@@ -482,6 +665,10 @@ def test_mempool_compute_ops(self):
 
     @skip_but_pass_in_sandcastle_if(TEST_WITH_ROCM, "Skip NCCL tests for ROCm")
     @skip_but_pass_in_sandcastle_if(IS_WINDOWS, "NCCL doesn't support Windows")
+    @skip_but_pass_in_sandcastle_if(
+        os.environ.get("NCCL_NVLS_ENABLE", "1") == "0",
+        "NCCL_NVLS_ENABLE=0",
+    )
     @skip_if_lt_x_gpu(2)
     @requires_nccl_version(
         (2, 29), "NCCL Symmetric Memory multicast support from nccl 2.29"
diff --git a/test/distributed/test_nvshmem.py b/test/distributed/test_nvshmem.py
index 5111aefee8f78..584a1c17e4e35 100644
--- a/test/distributed/test_nvshmem.py
+++ b/test/distributed/test_nvshmem.py
@@ -12,6 +12,7 @@
 from torch.distributed.device_mesh import init_device_mesh
 from torch.testing._internal.common_distributed import (
     MultiProcContinuousTest,
+    PLATFORM_SUPPORTS_SYMM_MEM,
     skip_if_lt_x_gpu,
 )
 from torch.testing._internal.common_utils import (
@@ -20,15 +21,14 @@
     requires_cuda_p2p_access,
     run_tests,
     skip_but_pass_in_sandcastle_if,
-    skipIfRocm,
 )
 
 
 # Decorator
 def requires_nvshmem():
     return skip_but_pass_in_sandcastle_if(
-        not symm_mem.is_nvshmem_available(),
-        "test_nvshmem requires NVSHMEM, skipping tests",
+        not PLATFORM_SUPPORTS_SYMM_MEM or not symm_mem.is_nvshmem_available(),
+        "test_nvshmem requires NVSHMEM with SymmMem platform support, skipping tests",
     )
 
 
@@ -73,7 +73,6 @@ def _init_device(self) -> None:
     def device(self) -> torch.device:
         return torch.device(device_type, self.rank)
 
-    @skipIfRocm
     def test_alloc(self) -> None:
         self._init_device()
 
@@ -91,7 +90,6 @@ def foo():
         out = symm_mem.empty(numel, dtype=dtype, device=self.device)
         symm_mem.rendezvous(out, group=group_name)
 
-    @skipIfRocm
     def test_alloc_without_device_context(self) -> None:
         # Set NVSHMEM as SymmMem backend
         symm_mem.set_backend("NVSHMEM")
@@ -103,7 +101,6 @@ def test_alloc_without_device_context(self) -> None:
         self.assertEqual(out.device, self.device)
         symm_mem.rendezvous(out, group=group_name)
 
-    @skipIfRocm
     def test_mempool_tensor_factory(self) -> None:
         """
         Test the effectiveness of MemPool on tensor factory ops.
@@ -128,7 +125,6 @@ def test_mempool_tensor_factory(self) -> None:
         torch.ops.symm_mem.nvshmem_broadcast(tensor, src_rank, group_name)
         self.assertEqual(tensor, torch.arange(numel, dtype=dtype, device=self.device))
 
-    @skipIfRocm
     def test_mempool_tensor_w_collective(self) -> None:
         """
         Test the effectiveness of MemPool on tensor factory ops.
@@ -151,7 +147,6 @@ def test_mempool_tensor_w_collective(self) -> None:
             tensor, torch.ones(numel, dtype=dtype, device=self.device) * self.world_size
         )
 
-    @skipIfRocm
     def test_mempool_compute_ops(self) -> None:
         """
         Apply MemPool context to a compute op that creates input to collective.
@@ -176,7 +171,6 @@ def test_mempool_compute_ops(self) -> None:
         expected = torch.mm(x0, w)
         self.assertEqual(y, expected)
 
-    @skipIfRocm
     def test_handle_offset(self) -> None:
         """
         Test if handle offset is correctly set.
@@ -254,7 +248,6 @@ def test_multicast_ptr(self) -> None:
         else:
             self.assertEqual(handle.multicast_ptr, 0)
 
-    @skipIfRocm
     def test_nvshmem_put(self) -> None:
         self._init_device()
         group_name = dist.group.WORLD.group_name
@@ -276,7 +269,6 @@ def test_nvshmem_put(self) -> None:
                 tensor, torch.zeros(numel, dtype=dtype, device=self.device)
             )
 
-    @skipIfRocm
     def test_nvshmem_get(self) -> None:
         self._init_device()
         group_name = dist.group.WORLD.group_name
@@ -313,7 +305,6 @@ def _init_device(self) -> None:
     def device(self) -> torch.device:
         return torch.device(device_type, self.rank)
 
-    @skipIfRocm
     def test_nvshmem_all_to_all(self) -> None:
         self._init_device()
 
@@ -337,7 +328,6 @@ def test_nvshmem_all_to_all(self) -> None:
         )
         torch.testing.assert_close(out, expected)
 
-    @skipIfRocm
     def test_all_to_all_vdev(self) -> None:
         self._init_device()
 
@@ -398,7 +388,6 @@ def test_all_to_all_vdev(self) -> None:
         )
         torch.testing.assert_close(out[:out_numel], expected)
 
-    @skipIfRocm
     @parametrize("align", [1, 8, 16])  # `major_align` of output
     def test_all_to_all_vdev_2d(self, align: int) -> None:
         torch.manual_seed(42 + self.rank)
@@ -507,7 +496,6 @@ def test_all_to_all_vdev_2d(self, align: int) -> None:
             received_chunk = out[start : start + split]
             torch.testing.assert_close(received_chunk, chunk)
 
-    @skipIfRocm
     def test_all_to_all_vdev_2d_offset(self) -> None:
         torch.manual_seed(42 + self.rank)
         self._init_device()
@@ -718,7 +706,6 @@ def _init_device(self) -> None:
     def device(self) -> torch.device:
         return torch.device(device_type, self.rank)
 
-    @skipIfRocm
     @parametrize("align", [1, 8, 16])  # `major_align` of output
     def test_dispatch_combine(self, align: int) -> None:
         """
@@ -743,9 +730,6 @@ def _init_device(self) -> None:
     def device(self) -> torch.device:
         return torch.device(device_type, self.rank)
 
-    @skipIfRocm
-    # TODO: FIXIT. Currently, `MultiProcContinuousTest` treats the skip code as a
-    # failure
     @skip_if_lt_x_gpu(4)
     def test_dispatch_combine_subgroup(self) -> None:
         """
@@ -777,13 +761,15 @@ def _init_device(self) -> None:
     def device(self) -> torch.device:
         return torch.device(device_type, self.rank)
 
-    @skipIfRocm
     @requires_nvls()
     @parametrize("tile_size", [32, 128, 512])
     @parametrize("dtype", [torch.float, torch.half, torch.bfloat16])
     def test_tile_reduce(self, tile_size: int, dtype: torch.dtype) -> None:
         full_size = 1024
-        assert tile_size <= full_size
+        if not (tile_size <= full_size):
+            raise AssertionError(
+                f"Expected tile_size <= full_size, got {tile_size} vs {full_size}"
+            )
 
         self._init_device()
         group_name = dist.group.WORLD.group_name
@@ -811,7 +797,6 @@ def test_tile_reduce(self, tile_size: int, dtype: torch.dtype) -> None:
 
         torch.testing.assert_close(full_out, expected)
 
-    @skipIfRocm
     @requires_nvls()
     @parametrize("tile_size", [32, 128, 512])
     @parametrize(
@@ -826,8 +811,16 @@ def test_multi_root_tile_reduce(
         num_slices_row = (
             self.world_size // num_slices_col
         )  # number of tiles on row dimension
-        assert tile_size * num_slices_col <= full_size
-        assert tile_size * num_slices_row <= full_size
+        if not (tile_size * num_slices_col <= full_size):
+            raise AssertionError(
+                f"tile_size * num_slices_col > full_size: "
+                f"{tile_size * num_slices_col} vs {full_size}"
+            )
+        if not (tile_size * num_slices_row <= full_size):
+            raise AssertionError(
+                f"tile_size * num_slices_row > full_size: "
+                f"{tile_size * num_slices_row} vs {full_size}"
+            )
 
         self._init_device()
         group_name = dist.group.WORLD.group_name
diff --git a/test/distributed/test_nvshmem_triton.py b/test/distributed/test_nvshmem_triton.py
index 47760cc006c31..a755576045817 100644
--- a/test/distributed/test_nvshmem_triton.py
+++ b/test/distributed/test_nvshmem_triton.py
@@ -4,19 +4,20 @@
 
 import sys
 
+import torch
+
 # Import TEST_WITH_ROCM first to check for ROCm before importing NVSHMEM modules
 from torch.testing._internal.common_utils import TEST_WITH_ROCM
 
 
-# Skip entire module on ROCm before importing NVSHMEM-specific modules
+# Skip entire module on ROCm before importing NVSHMEM-specific modules or running on CPU
 # NVSHMEM is NVIDIA-specific and can cause crashes during import on ROCm
-if TEST_WITH_ROCM:
+if TEST_WITH_ROCM or not torch.backends.cuda.is_built():
     print("NVSHMEM not available on ROCm, skipping tests")
     sys.exit(0)
 
 import triton.language as tl
 
-import torch
 import torch.distributed as dist
 import torch.distributed._symmetric_memory as symm_mem
 import torch.distributed._symmetric_memory._nvshmem_triton as nvshmem
diff --git a/test/distributed/test_overlap_bucketing_unit.py b/test/distributed/test_overlap_bucketing_unit.py
index 1622f505ef823..71bb8cf6c63b5 100644
--- a/test/distributed/test_overlap_bucketing_unit.py
+++ b/test/distributed/test_overlap_bucketing_unit.py
@@ -450,6 +450,7 @@ def func(a, b):
             traced.graph,
             collective_info,
             scheduled,
+            bucket_only_internode_comms=False,
         )
         bucketer.bucket_collectives()
 
@@ -461,6 +462,240 @@ def func(a, b):
             "split_with_sizes"
         ).check_count("%mm", 2).run(graph_str)
 
+    def test_no_cross_type_bucketing_ar_and_rs(self):
+        """
+        Test that all_reduce and reduce_scatter on the same PG with
+        matching reduce_op and dtype are NOT bucketed together.
+
+        bucket_key() returns (group_name, reduce_op, dtype) for both
+        all_reduce and reduce_scatter. Without the collective type in
+        the key, they would be incorrectly grouped together.
+        """
+
+        def func(a, b):
+            group_name = "0"
+            group_size = 2
+
+            ar1 = torch.ops._c10d_functional.all_reduce(a, "sum", group_name)
+            ar2 = torch.ops._c10d_functional.all_reduce(b, "sum", group_name)
+
+            rs1 = torch.ops._c10d_functional.reduce_scatter_tensor(
+                a, "sum", group_size, group_name
+            )
+            rs2 = torch.ops._c10d_functional.reduce_scatter_tensor(
+                b, "sum", group_size, group_name
+            )
+
+            ar1_out = torch.ops._c10d_functional.wait_tensor(ar1)
+            ar2_out = torch.ops._c10d_functional.wait_tensor(ar2)
+            rs1_out = torch.ops._c10d_functional.wait_tensor(rs1)
+            rs2_out = torch.ops._c10d_functional.wait_tensor(rs2)
+
+            return ar1_out.sum() + ar2_out.sum() + rs1_out.sum() + rs2_out.sum()
+
+        with FakeTensorMode():
+            a = torch.ones(4, 4, device=self.device)
+            b = torch.ones(4, 4, device=self.device) * 2
+            traced = make_fx(func)(a, b)
+
+        ar1, ar2 = traced.graph.find_nodes(
+            op="call_function",
+            target=torch.ops._c10d_functional.all_reduce.default,
+        )
+        rs1, rs2 = traced.graph.find_nodes(
+            op="call_function",
+            target=torch.ops._c10d_functional.reduce_scatter_tensor.default,
+        )
+
+        # No hiding — all exposed
+        hiding_annotations = {}
+
+        collective_info = build_collective_info(traced.graph, hiding_annotations)
+        scheduled = OrderedSet(traced.graph.nodes)
+
+        from torch._inductor.fx_passes.overlap_preserving_bucketer import (
+            OverlapPreservingBucketer,
+        )
+
+        bucketer = OverlapPreservingBucketer(
+            traced.graph,
+            collective_info,
+            scheduled,
+            bucket_only_internode_comms=False,
+        )
+        bucketer.bucket_collectives()
+
+        # all_reduce ops should be bucketed together (1 bucketed all_reduce)
+        ar_nodes = traced.graph.find_nodes(
+            op="call_function",
+            target=torch.ops._c10d_functional.all_reduce.default,
+        )
+        self.assertEqual(len(ar_nodes), 1)
+
+        # reduce_scatter ops should be bucketed together (1 bucketed reduce_scatter)
+        rs_nodes = traced.graph.find_nodes(
+            op="call_function",
+            target=torch.ops._c10d_functional.reduce_scatter_tensor.default,
+        )
+        self.assertEqual(len(rs_nodes), 1)
+
+    def test_reduce_scatter_coalesced_mode(self):
+        """
+        Test that 'coalesced' bucket mode uses reduce_scatter_tensor_coalesced
+        instead of cat + single reduce_scatter.
+        """
+
+        def func(a, b):
+            group_name = "0"
+            group_size = 2
+
+            rs1 = torch.ops._c10d_functional.reduce_scatter_tensor(
+                a, "sum", group_size, group_name
+            )
+            rs2 = torch.ops._c10d_functional.reduce_scatter_tensor(
+                b, "sum", group_size, group_name
+            )
+
+            rs1_out = torch.ops._c10d_functional.wait_tensor(rs1)
+            rs2_out = torch.ops._c10d_functional.wait_tensor(rs2)
+
+            return rs1_out.sum() + rs2_out.sum()
+
+        with FakeTensorMode():
+            a = torch.ones(4, 4, device=self.device)
+            b = torch.ones(4, 4, device=self.device) * 2
+            traced = make_fx(func)(a, b)
+
+        rs1, rs2 = traced.graph.find_nodes(
+            op="call_function",
+            target=torch.ops._c10d_functional.reduce_scatter_tensor.default,
+        )
+
+        hiding_annotations = {}
+        collective_info = build_collective_info(traced.graph, hiding_annotations)
+        scheduled = OrderedSet(traced.graph.nodes)
+
+        from torch._inductor.fx_passes.overlap_preserving_bucketer import (
+            OverlapPreservingBucketer,
+        )
+
+        bucketer = OverlapPreservingBucketer(
+            traced.graph,
+            collective_info,
+            scheduled,
+            bucket_mode="coalesced",
+        )
+        bucketer.bucket_collectives()
+
+        graph_str = str(traced.graph)
+        self.assertIn("reduce_scatter_tensor_coalesced", graph_str)
+        self.assertNotIn("cat.default", graph_str)
+
+    def test_all_gather_coalesced_mode_falls_back(self):
+        """
+        Test that 'coalesced' bucket mode falls back to default bucketing
+        for all_gather (coalesced is only supported for reduce_scatter).
+        """
+
+        def func(a, b):
+            group_name = "0"
+            group_size = 2
+
+            ag1 = torch.ops._c10d_functional.all_gather_into_tensor(
+                a, group_size, group_name
+            )
+            ag2 = torch.ops._c10d_functional.all_gather_into_tensor(
+                b, group_size, group_name
+            )
+
+            ag1_out = torch.ops._c10d_functional.wait_tensor(ag1)
+            ag2_out = torch.ops._c10d_functional.wait_tensor(ag2)
+
+            return ag1_out.sum() + ag2_out.sum()
+
+        with FakeTensorMode():
+            a = torch.ones(4, 4, device=self.device)
+            b = torch.ones(4, 4, device=self.device) * 2
+            traced = make_fx(func)(a, b)
+
+        ag1, ag2 = traced.graph.find_nodes(
+            op="call_function",
+            target=torch.ops._c10d_functional.all_gather_into_tensor.default,
+        )
+
+        hiding_annotations = {}
+        collective_info = build_collective_info(traced.graph, hiding_annotations)
+        scheduled = OrderedSet(traced.graph.nodes)
+
+        from torch._inductor.fx_passes.overlap_preserving_bucketer import (
+            OverlapPreservingBucketer,
+        )
+
+        bucketer = OverlapPreservingBucketer(
+            traced.graph,
+            collective_info,
+            scheduled,
+            bucket_mode="coalesced",
+        )
+        bucketer.bucket_collectives()
+
+        # Should use default bucketing (cat + single all_gather), not coalesced
+        graph_str = str(traced.graph)
+        self.assertIn("cat.default", graph_str)
+        self.assertNotIn("all_gather_into_tensor_coalesced", graph_str)
+
+    def test_reduce_scatter_coalesced_mode_single_rs(self):
+        """
+        Test that 'coalesced' bucket mode correctly identifies the collective
+        start node even when there is only a single reduce_scatter (so only
+        one wait node is produced).  This exercises the unified
+        _get_collective_node_from_wait path instead of the .args[0] shortcut.
+        """
+
+        def func(a):
+            group_name = "0"
+            group_size = 2
+
+            rs1 = torch.ops._c10d_functional.reduce_scatter_tensor(
+                a, "sum", group_size, group_name
+            )
+
+            rs1_out = torch.ops._c10d_functional.wait_tensor(rs1)
+
+            return rs1_out.sum()
+
+        with FakeTensorMode():
+            a = torch.ones(4, 4, device=self.device)
+            traced = make_fx(func)(a)
+
+        rs_nodes = traced.graph.find_nodes(
+            op="call_function",
+            target=torch.ops._c10d_functional.reduce_scatter_tensor.default,
+        )
+        self.assertEqual(len(rs_nodes), 1)
+
+        hiding_annotations = {}
+        collective_info = build_collective_info(traced.graph, hiding_annotations)
+        scheduled = OrderedSet(traced.graph.nodes)
+
+        from torch._inductor.fx_passes.overlap_preserving_bucketer import (
+            OverlapPreservingBucketer,
+        )
+
+        bucketer = OverlapPreservingBucketer(
+            traced.graph,
+            collective_info,
+            scheduled,
+            bucket_mode="coalesced",
+        )
+        # Single RS should not crash — bucket_collectives groups by key,
+        # and a single-element group is a no-op (no merge needed).
+        bucketer.bucket_collectives()
+
+        # Graph should still contain the original reduce_scatter (no merge happened)
+        graph_str = str(traced.graph)
+        self.assertIn("reduce_scatter_tensor", graph_str)
+
     def test_can_bucket_multidtype_collectives(self):
         """
         Test that all_gathers with different dtypes CAN bucket together.
@@ -709,6 +944,7 @@ def func(a, b, c):
             traced.graph,
             collective_info,
             scheduled,
+            bucket_mode="custom_ops_multidtype",
         )
         bucketer.bucket_collectives()
 
@@ -721,6 +957,99 @@ def func(a, b, c):
         f.check("pre_bucket_all_gather").check("all_gather_into_tensor_out")
         f.run(graph_str)
 
+    def test_dead_fusible_code_no_crash(self):
+        """
+        Test that dead fusible code (fusion regions with no external outputs)
+        does not crash collapse_fusion_regions, and that collapse/expand
+        round-trips preserve the graph.
+
+        Regression test for the bug where dead code created a fusion region
+        with no external outputs, causing fuse_by_partitions to crash with
+        "AssertionError: last_output_node is None".
+        """
+
+        def func_with_dead_fusible_code(x, y):
+            group_name = "0"
+            group_size = 1
+
+            ag = torch.ops._c10d_functional.all_gather_into_tensor(
+                x, group_size, group_name
+            )
+
+            # Dead fusible chain - not consumed by output
+            dead1 = x + 1.0
+            dead2 = dead1 * 2.0
+            dead3 = dead2 + dead1  # noqa: F841
+
+            # Live fusible chain
+            live1 = y + 1.0
+            live2 = live1 * 2.0
+
+            mm_result = torch.mm(y, y)
+            live3 = mm_result + 1.0
+
+            ag_out = torch.ops._c10d_functional.wait_tensor(ag)
+
+            return (live2 + live3 + ag_out).sum()
+
+        from torch._inductor.fx_passes.fusion_regions import (
+            build_fusion_regions,
+            collapse_fusion_regions,
+            expand_fusion_regions,
+        )
+
+        with FakeTensorMode():
+            x = torch.randn(16, 16)
+            y = torch.randn(16, 16)
+            gm = make_fx(func_with_dead_fusible_code)(x, y)
+
+        graph_str_before = gm.print_readable(print_output=False)
+
+        region_of = build_fusion_regions(gm)
+        new_region_of = collapse_fusion_regions(gm, region_of)
+
+        # Expand back and verify graph is preserved
+        expand_fusion_regions(gm, new_region_of)
+        gm.recompile()
+        graph_str_after = gm.print_readable(print_output=False)
+        self.assertEqual(graph_str_before, graph_str_after)
+
+    @torch._inductor.config.patch(deterministic=True)
+    def test_deterministic_mode_no_benchmark_error(self):
+        """
+        Test that deterministic mode doesn't error when running overlap scheduling.
+
+        Before the fix, deterministic mode would error when trying to benchmark
+        compute nodes. Now it uses analytical estimation instead.
+        """
+        from torch._inductor.fx_passes.overlap_scheduling import (
+            schedule_overlap_bucketing,
+        )
+
+        def func(a, b):
+            group_name = "0"
+            group_size = 1
+
+            ag = torch.ops._c10d_functional.all_gather_into_tensor(
+                a, group_size, group_name
+            )
+
+            # Compute with gemm
+            mm_result = torch.mm(a, b)
+            pointwise = mm_result + 1.0
+
+            ag_out = torch.ops._c10d_functional.wait_tensor(ag)
+
+            return (pointwise + ag_out).sum()
+
+        with FakeTensorMode():
+            a = torch.randn(16, 16, device=self.device)
+            b = torch.randn(16, 16, device=self.device)
+            gm = make_fx(func)(a, b)
+
+        # Should not error in deterministic mode (would have errored before fix)
+        schedule_overlap_bucketing(gm)
+
 
 @requires_accelerator_dist_backend(["nccl", "xccl"])
 @unittest.skipIf(not HAS_GPU, "Inductor+gpu needs triton and recent GPU arch")
@@ -877,7 +1206,7 @@ def custom_runtime(node: fx.Node, override_size: int | None) -> float | None:
             return 0.0
 
         out = schedule_overlap_bucketing(
-            traced, custom_runtime_estimation=custom_runtime
+            traced, custom_runtime_estimation=custom_runtime, max_off_bucket_gb=None
         )
 
         # Get scheduled order
@@ -899,6 +1228,7 @@ def custom_runtime(node: fx.Node, override_size: int | None) -> float | None:
 
 @requires_accelerator_dist_backend(["nccl", "xccl"])
 @unittest.skipIf(not HAS_GPU, "Inductor+gpu needs triton and recent GPU arch")
+@instantiate_parametrized_tests
 class TestFusibleNodeOverlap(InductorTestCase):
     """Test that fusible nodes are used for overlapping with collectives."""
 
@@ -1016,6 +1346,56 @@ def func(a):
             "sub"
         ).check("wait_tensor").run(graph_str)
 
+    @parametrize("enable_fusion_regions", [False, True])
+    def test_precomputed_estimations_via_custom_runtime(self, enable_fusion_regions):
+        """Pre-computed estimations from gather_node_runtime_estimations can be
+        fed into OverlapScheduler via custom_runtime_estimation wrapping a dict."""
+
+        def func(a):
+            group_name = "0"
+            group_size = 1
+            ag = torch.ops._c10d_functional.all_gather_into_tensor(
+                a, group_size, group_name
+            )
+            b = a + 1
+            b = b * 2
+            ag_out = torch.ops._c10d_functional.wait_tensor(ag)
+            return ag_out.sum() + b.sum()
+
+        with FakeTensorMode():
+            a = torch.ones(1024, 1024, device=self.device)
+            traced = make_fx(func)(a)
+
+        from torch._inductor.fx_passes.overlap_scheduling import (
+            gather_node_runtime_estimations,
+            OverlapScheduler,
+        )
+
+        estimations, fusion_region_of = gather_node_runtime_estimations(
+            traced,
+            collective_estimator="analytical",
+            enable_fusion_regions=enable_fusion_regions,
+        )
+        for node in fusion_region_of:
+            self.assertIn(node, estimations)
+
+        scheduler = OverlapScheduler(
+            traced,
+            max_in_flight_gb=5.0,
+            max_compute_pre_fetch=200,
+            collective_bucketing=False,
+            insert_overlap_deps=False,
+            compute_overlap_multipler=1.0,
+            max_coll_distance=200,
+            custom_runtime_estimation=lambda node, _: estimations.get(node),
+            collective_estimator="analytical",
+        )
+        out = scheduler.run()
+        FileCheck().check("all_gather_into_tensor").check("wait_tensor").run(
+            str(out.graph)
+        )
+        self.assertEqual(len(scheduler.collective_info), 1)
+
 
 @requires_accelerator_dist_backend(["nccl", "xccl"])
 @unittest.skipIf(not HAS_GPU, "Inductor+gpu needs triton and recent GPU arch")
@@ -1206,6 +1586,83 @@ def func(a, b, c, d):
         result = scheduler.run()
         result.graph.lint()
 
+    def test_no_cycle_from_dtype_convert_timeline_deps(self):
+        """
+        Test that transfer_erased_node_deps doesn't create cycles when
+        convert_element_type nodes with timeline deps get fused into a
+        _pre_bucket node that is a data dependency of the bucketed start.
+
+        Uses make_fx to trace a realistic graph with dtype conversions and
+        collectives, then simulates the bucketing transfer that would create
+        a cycle: _pre_bucket extra-depends-on new_start, while new_start
+        already data-depends on _pre_bucket.
+        """
+        from torch._inductor.augmented_graph_helper import AugmentedGraphHelper
+
+        def func(a, b):
+            group_name = dist.distributed_c10d._get_default_group().group_name
+            group_size = 16
+
+            conv_a = torch.ops.prims.convert_element_type.default(a, torch.bfloat16)
+            ag1 = torch.ops._c10d_functional.all_gather_into_tensor(
+                conv_a, group_size, group_name
+            )
+            conv_b = torch.ops.prims.convert_element_type.default(b, torch.bfloat16)
+            ag2 = torch.ops._c10d_functional.all_gather_into_tensor(
+                conv_b, group_size, group_name
+            )
+            w1 = torch.ops._c10d_functional.wait_tensor(ag1)
+            w2 = torch.ops._c10d_functional.wait_tensor(ag2)
+            return w1.sum() + w2.sum()
+
+        with FakeTensorMode():
+            a = torch.ones(4, 4, device=self.device)
+            b = torch.ones(4, 4, device=self.device)
+            traced = make_fx(func)(a, b)
+
+        graph = traced.graph
+        nodes = {n.name: n for n in graph.nodes}
+        conv_a = nodes["convert_element_type"]
+        conv_b = nodes["convert_element_type_1"]
+        ag_start_1 = nodes["all_gather_into_tensor"]
+        ag_start_2 = nodes["all_gather_into_tensor_1"]
+        ag_wait_1 = nodes["wait_tensor"]
+        ag_wait_2 = nodes["wait_tensor_1"]
+
+        helper = AugmentedGraphHelper(graph)
+
+        # Simulate PG timeline dep: conv_b scheduled after ag_start_1
+        helper.add_extra_dep(n=conv_b, dep=ag_start_1)
+
+        # Simulate bucketed replacement nodes as custom_ops_multidtype would:
+        # _pre_bucket is a data dep of new_start (new_start.kwargs["out"])
+        x = nodes["a_1"]
+        pre_bucket = graph.call_function(torch.sigmoid, args=(x,), name="pre_bucket")
+        new_start = graph.call_function(
+            torch.tanh, kwargs={"out": pre_bucket}, name="new_start"
+        )
+        new_wait = graph.call_function(torch.exp, args=(new_start,), name="new_wait")
+
+        # This is what _apply_bucket does: converts → pre_bucket,
+        # starts → new_start, waits → new_wait
+        helper.transfer_erased_node_deps(
+            {
+                conv_a: pre_bucket,
+                conv_b: pre_bucket,
+                ag_start_1: new_start,
+                ag_start_2: new_start,
+                ag_wait_1: new_wait,
+                ag_wait_2: new_wait,
+            }
+        )
+
+        # Without the fix, pre_bucket gets extra dep on new_start,
+        # but new_start already data-depends on pre_bucket → cycle
+        self.assertFalse(
+            helper.has_cycle(),
+            "Cycle: pre_bucket <-> new_start via data + extra deps",
+        )
+
 
 class TestForeachGroupsUnit(InductorTestCase):
     """Unit tests for _compute_foreach_groups and _pre_bucket_all_gather foreach optimization."""
@@ -1241,5 +1698,171 @@ def test_foreach_groups_correctness(self):
         self.assertTrue(torch.allclose(result_with, result_without))
 
 
+@requires_accelerator_dist_backend(["nccl", "xccl"])
+@unittest.skipIf(not HAS_GPU, "Inductor+gpu needs triton and recent GPU arch")
+class TestCoalescedCollectiveOverlap(InductorTestCase):
+    """
+    Tests for coalesced collective support in overlap scheduling.
+
+    Coalesced collectives (e.g., reduce_scatter_tensor_coalesced) return
+    Tensor[] instead of Tensor. Wait nodes follow the pattern:
+      wait_tensor(getitem(coalesced_collective, idx))
+    Multiple waits share the same collective start node.
+    """
+
+    @classmethod
+    def setUpClass(cls):
+        super().setUpClass()
+        from torch.testing._internal.distributed.fake_pg import FakeStore
+
+        store = FakeStore()
+        dist.init_process_group(backend="fake", rank=0, world_size=8, store=store)
+        cls.device = "cuda"
+
+    @classmethod
+    def tearDownClass(cls):
+        super().tearDownClass()
+        dist.destroy_process_group()
+
+    def _make_coalesced_rs_graph(self, num_tensors=3):
+        """Create an FX graph with a coalesced reduce_scatter and surrounding compute."""
+        group_name = dist.distributed_c10d._get_default_group().group_name
+
+        def func(*inputs):
+            # Compute before collective
+            mm1 = torch.mm(inputs[0], inputs[0])
+
+            # Coalesced reduce_scatter
+            flat_inputs = [x.view(-1) for x in inputs]
+            rs_outs = torch.ops._c10d_functional.reduce_scatter_tensor_coalesced(
+                flat_inputs, "sum", 8, group_name
+            )
+
+            # Wait on each output
+            waited = [torch.ops._c10d_functional.wait_tensor(o) for o in rs_outs]
+
+            # Compute after collective
+            mm2 = torch.mm(inputs[0], inputs[0])
+
+            return mm1.sum() + mm2.sum() + sum(w.sum() for w in waited)
+
+        with FakeTensorMode():
+            inputs = [torch.ones(8, 8, device=self.device) for _ in range(num_tensors)]
+            traced = make_fx(func)(*inputs)
+
+        return traced
+
+    def test_identify_collectives_coalesced(self):
+        """
+        _identify_collectives should register one CollectiveInfo per coalesced
+        collective, not overwrite for each wait.
+        """
+        from torch._inductor.fx_passes.overlap_scheduling import OverlapScheduler
+
+        traced = self._make_coalesced_rs_graph(num_tensors=3)
+
+        def custom_runtime(node, override_size):
+            if "reduce_scatter" in str(node.target):
+                return 5.0
+            return 0.1
+
+        scheduler = OverlapScheduler(
+            traced,
+            max_in_flight_gb=5.0,
+            max_compute_pre_fetch=200,
+            collective_bucketing=False,
+            insert_overlap_deps=False,
+            compute_overlap_multipler=1.0,
+            max_coll_distance=200,
+            custom_runtime_estimation=custom_runtime,
+            collective_estimator="analytical",
+        )
+
+        # Should have exactly 1 collective (the coalesced RS), not 3
+        self.assertEqual(len(scheduler.collective_info), 1)
+        # All 3 waits should map to the same start
+        self.assertEqual(len(scheduler.wait_to_start), 3)
+        starts = set(scheduler.wait_to_start.values())
+        self.assertEqual(len(starts), 1)
+
+    def test_overlap_scheduler_coalesced_runs(self):
+        """
+        Full overlap scheduler run with coalesced collectives should not crash.
+        Previously crashed with assertion errors in _handle_wait and
+        normalize_function failures in comm_analysis.
+        """
+        from torch._inductor.fx_passes.overlap_scheduling import OverlapScheduler
+
+        traced = self._make_coalesced_rs_graph(num_tensors=4)
+
+        def custom_runtime(node, override_size):
+            if "reduce_scatter" in str(node.target):
+                return 5.0
+            return 0.1
+
+        result = OverlapScheduler(
+            traced,
+            max_in_flight_gb=5.0,
+            max_compute_pre_fetch=200,
+            collective_bucketing=False,
+            insert_overlap_deps=False,
+            compute_overlap_multipler=1.0,
+            max_coll_distance=200,
+            custom_runtime_estimation=custom_runtime,
+            collective_estimator="analytical",
+        ).run()
+
+        # Graph should still contain the coalesced collective and waits
+        graph_str = str(result.graph)
+        self.assertIn("reduce_scatter_tensor_coalesced", graph_str)
+        self.assertIn("wait_tensor", graph_str)
+
+    def test_gather_node_runtime_estimations_coalesced(self):
+        """
+        gather_node_runtime_estimations should estimate coalesced collectives
+        exactly once (not once per wait).
+        """
+        from torch._inductor.fx_passes.overlap_scheduling import (
+            gather_node_runtime_estimations,
+        )
+
+        traced = self._make_coalesced_rs_graph(num_tensors=3)
+
+        def custom_runtime(node, override_size):
+            if "reduce_scatter" in str(node.target):
+                return 5.0
+            return None
+
+        estimations, _ = gather_node_runtime_estimations(
+            traced,
+            custom_runtime_estimation=custom_runtime,
+            collective_estimator="analytical",
+        )
+
+        # The coalesced RS node should appear exactly once in estimations
+        rs_estimations = {
+            k: v for k, v in estimations.items() if "reduce_scatter" in str(k.target)
+        }
+        self.assertEqual(len(rs_estimations), 1)
+
+    def test_estimate_fx_collective_size_coalesced(self):
+        """
+        estimate_fx_collective_size should handle coalesced collectives
+        whose output meta is a list of tensors, not a single tensor.
+        """
+        from torch._inductor.comm_analysis import estimate_fx_collective_size
+
+        traced = self._make_coalesced_rs_graph(num_tensors=2)
+
+        rs_nodes = traced.graph.find_nodes(
+            op="call_function",
+            target=torch.ops._c10d_functional.reduce_scatter_tensor_coalesced.default,
+        )
+        self.assertEqual(len(rs_nodes), 1)
+        size = estimate_fx_collective_size(rs_nodes[0])
+        # Should return non-zero for coalesced collectives
+        self.assertGreater(size, 0)
+
+
 if __name__ == "__main__":
     run_tests()
diff --git a/test/distributed/test_p2p_ipc.py b/test/distributed/test_p2p_ipc.py
index 44a9791c19943..8f964ebc9befc 100644
--- a/test/distributed/test_p2p_ipc.py
+++ b/test/distributed/test_p2p_ipc.py
@@ -76,7 +76,8 @@ def _test_p2p_ipc_impl(self, tensor_size: int = 2333) -> None:
         device_module.synchronize()
         torch.distributed.barrier()
 
-        assert tensor.allclose(tensor, 1)
+        if not tensor.allclose(tensor, 1):
+            raise AssertionError("Expected tensor to be close to 1")
 
         torch.distributed.barrier()
 
diff --git a/test/distributed/test_pg_wrapper.py b/test/distributed/test_pg_wrapper.py
index 4bd49e7528e38..60735673ffdcf 100644
--- a/test/distributed/test_pg_wrapper.py
+++ b/test/distributed/test_pg_wrapper.py
@@ -371,6 +371,29 @@ def test_coalescing_manager_debug_mode_detail(self):
         @requires_accelerator_dist_backend(["nccl", "xccl"])
         @skip_if_lt_x_gpu(2)
         @with_dist_debug_levels(levels=["DETAIL"])
+        def test_reduce_scatter_tensor_coalesced_debug_mode(self):
+            torch.cuda.set_device(self.rank)
+            pg = self._create_wrapper_pg(with_new_group=True)
+            dev = torch.cuda.current_device()
+
+            out_shapes = [(2, 2), (3, 3)]
+            in_shapes = [(s[0] * self.world_size,) + s[1:] for s in out_shapes]
+
+            outputs = [torch.zeros(s, device=dev) for s in out_shapes]
+            inputs = [torch.ones(s, device=dev) * (self.rank + 1) for s in in_shapes]
+
+            work = pg.reduce_scatter_tensor_coalesced(outputs, inputs)
+            work.wait()
+
+            for i, (output, _) in enumerate(zip(outputs, inputs)):
+                expected = torch.ones(out_shapes[i], device=dev) * sum(
+                    range(1, self.world_size + 1)
+                )
+                self.assertTrue(torch.allclose(output, expected))
+
+        @requires_nccl()
+        @skip_if_lt_x_gpu(2)
+        @with_dist_debug_levels(levels=["DETAIL"])
         @patch("torch.distributed.distributed_c10d._GLOO_AVAILABLE", False)
         def test_debug_level_detail_no_gloo(self):
             with self.assertRaisesRegex(
@@ -524,10 +547,26 @@ def test_collective_shape_mismatch_cuda(self):
         pg = self._create_wrapper_pg(with_new_group=False)
         self._test_collective_shape_mismatch(pg, use_accel=True)
 
+    @with_dist_debug_levels(levels=["DETAIL"])
+    def test_reduce_scatter_tensor_coalesced_debug_mode(self):
+        pg = self._create_wrapper_pg(with_new_group=True)
+
+        out_shapes = [(2, 2), (3, 3)]
+        in_shapes = [(s[0] * self.world_size,) + s[1:] for s in out_shapes]
+        outputs = [torch.zeros(s) for s in out_shapes]
+        inputs = [torch.ones(s) * (self.rank + 1) for s in in_shapes]
+
+        work = pg.reduce_scatter_tensor_coalesced(outputs, inputs)
+        work.wait()
+        for i, (output, _) in enumerate(zip(outputs, inputs)):
+            expected = torch.ones(out_shapes[i]) * sum(range(1, self.world_size + 1))
+            self.assertTrue(torch.allclose(output, expected))
+
 
 if __name__ == "__main__":
-    assert not (torch.cuda.is_initialized() or torch.xpu.is_initialized()), (
-        "test_pg_wrapper must not have initialized CUDA/XPU context on main process"
-    )
+    if torch.cuda.is_initialized() or torch.xpu.is_initialized():
+        raise AssertionError(
+            "test_pg_wrapper must not have initialized CUDA/XPU context on main process"
+        )
 
     run_tests()
diff --git a/test/distributed/test_run.py b/test/distributed/test_run.py
index 509c08cbf0c35..be10c3a78ba91 100644
--- a/test/distributed/test_run.py
+++ b/test/distributed/test_run.py
@@ -86,6 +86,43 @@ def test_launch_agent_sets_environment_variable(self, mock_get_handler, mock_age
             os.environ["TORCHELASTIC_SIGNALS_TO_HANDLE"], "SIGTERM,SIGUSR1,SIGUSR2"
         )
 
+    def _parse_args(self, args):
+        parser = run.get_args_parser()
+        return parser.parse_args(args)
+
+    @patch("torch.distributed.run.elastic_launch")
+    def test_single_node_defaults_to_free_port(self, mock_launch):
+        """Single-node with no explicit port should auto-switch to c10d with localhost:0."""
+        args = self._parse_args(["dummy_script.py"])
+        self.assertIsNone(args.master_port)
+        run.run(args)
+        self.assertEqual(args.rdzv_backend, "c10d")
+        self.assertEqual(args.rdzv_endpoint, "localhost:0")
+
+    @patch("torch.distributed.run.elastic_launch")
+    def test_single_node_explicit_port_uses_static(self, mock_launch):
+        """Single-node with explicit --master-port should keep static backend."""
+        args = self._parse_args(["--master-port=12345", "dummy_script.py"])
+        self.assertEqual(args.master_port, 12345)
+        run.run(args)
+        self.assertEqual(args.rdzv_backend, "static")
+
+    @patch("torch.distributed.run.elastic_launch")
+    def test_multi_node_defaults_to_static(self, mock_launch):
+        """Multi-node should keep static backend with default port 29500."""
+        args = self._parse_args(["--nnodes=2", "dummy_script.py"])
+        run.run(args)
+        self.assertEqual(args.rdzv_backend, "static")
+        self.assertEqual(args.master_port, 29500)
+
+    @patch("torch.distributed.run.elastic_launch")
+    def test_standalone_still_works(self, mock_launch):
+        """--standalone should still use c10d with localhost:0."""
+        args = self._parse_args(["--standalone", "dummy_script.py"])
+        run.run(args)
+        self.assertEqual(args.rdzv_backend, "c10d")
+        self.assertEqual(args.rdzv_endpoint, "localhost:0")
+
 
 if __name__ == "__main__":
     run_tests()
diff --git a/test/distributed/test_store.py b/test/distributed/test_store.py
index 310e41f5829a3..8975aa2439a42 100644
--- a/test/distributed/test_store.py
+++ b/test/distributed/test_store.py
@@ -261,6 +261,13 @@ def test_list_keys(self):
         self.assertIn("foo", keys)
         self.assertIn("baz", keys)
 
+    def _test_barrier(self, store):
+        # Barrier with world_size=1 should return immediately
+        store.barrier("test_barrier_key", 1)
+
+    def test_barrier(self):
+        self._test_barrier(self._create_store())
+
     # This is the number of keys used in test_set_get. Adding this as a class
     # property instead of hardcoding in the test since some Store
     # implementations will have differing number of keys. In the base case,
@@ -296,7 +303,8 @@ def test_init_pg_and_rpc_with_same_file(self):
                 "gloo", rank=0, world_size=1, init_method=f"file://{file.name}"
             )
             dist.destroy_process_group()
-            assert os.path.exists(file.name)
+            if not os.path.exists(file.name):
+                raise AssertionError(f"Expected file {file.name} to exist")
 
             rpc.shutdown()
             os.remove(file.name)
@@ -307,9 +315,11 @@ def test_refcount(self):
             store2 = dist.FileStore(file.name, 1)
 
             del store
-            assert os.path.exists(file.name)
+            if not os.path.exists(file.name):
+                raise AssertionError(f"Expected file {file.name} to exist")
             del store2
-            assert not os.path.exists(file.name)
+            if os.path.exists(file.name):
+                raise AssertionError(f"Expected file {file.name} to not exist")
 
     @property
     def num_keys_total(self):
@@ -463,7 +473,8 @@ def test_init_pg_and_rpc_with_same_socket(self):
         )
 
         del os.environ["USE_LIBUV"]
-        assert "USE_LIBUV" not in os.environ
+        if "USE_LIBUV" in os.environ:
+            raise AssertionError("Expected USE_LIBUV to not be in os.environ")
         rpc.shutdown()
         dist.destroy_process_group()
 
@@ -642,6 +653,42 @@ def test_agent_store(self) -> None:
 
         self.assertEqual(second_server.port, store.port)
 
+    def test_barrier(self):
+        store = self._create_store()
+        # Single worker barrier should return immediately
+        store.barrier("test_barrier_key", 1)
+
+    def test_barrier_with_timeout(self):
+        store = self._create_store()
+        store.barrier("test_barrier_timeout_key", 1, timedelta(seconds=10))
+
+    def test_barrier_timeout_expires(self):
+        store = self._create_store()
+        with self.assertRaisesRegex(DistStoreError, "barrier timeout"):
+            store.barrier("test_barrier_fail", 2, timedelta(seconds=0.1))
+
+    def test_barrier_multi_worker(self):
+        server_store = self._create_store()
+        world_size = 4
+
+        def worker(rank):
+            if rank == 0:
+                worker_store = server_store
+            else:
+                worker_store = dist.TCPStore(
+                    host_name=server_store.host,
+                    port=server_store.port,
+                    is_master=False,
+                    wait_for_workers=False,
+                    use_libuv=self._use_libuv,
+                )
+            worker_store.barrier("multi_barrier", world_size, timedelta(seconds=10))
+
+        with ThreadPoolExecutor(max_workers=world_size) as pool:
+            futures = [pool.submit(worker, i) for i in range(world_size)]
+            for f in futures:
+                f.result()
+
 
 class LibUvTCPStoreTest(TCPStoreTest):
     _use_libuv = True
@@ -1189,7 +1236,8 @@ def listen() -> None:
 
 if __name__ == "__main__":
     if device_type != "cpu":
-        assert not torch.get_device_module()._initialized, (
-            f"test_distributed must not have initialized {device_type} context on main process"
-        )
+        if torch.get_device_module()._initialized:
+            raise AssertionError(
+                f"test_distributed must not have initialized {device_type} context on main process"
+            )
     run_tests()
diff --git a/test/distributed/test_symmetric_memory.py b/test/distributed/test_symmetric_memory.py
index d5d3c97c5ac71..da880af0728af 100644
--- a/test/distributed/test_symmetric_memory.py
+++ b/test/distributed/test_symmetric_memory.py
@@ -10,6 +10,7 @@
 import torch
 import torch.distributed as dist
 import torch.distributed._symmetric_memory as symm_mem
+from torch._C import FileCheck
 from torch._C._autograd import DeviceType
 from torch._C._distributed_c10d import _SymmetricMemory
 from torch._inductor.utils import (
@@ -78,6 +79,24 @@ def test_has_multicast_support(self) -> None:
         self.assertFalse(_SymmetricMemory.has_multicast_support(DeviceType.CPU, 0))
         # NOTE: DeviceType.CUDA is implicitly tested through @requires_multicast_support
 
+    @requires_cuda
+    @skipIf(
+        not PLATFORM_SUPPORTS_SYMM_MEM, "SymmMem is not supported on this ROCm arch"
+    )
+    @skip_if_lt_x_gpu(2)
+    def test_is_symm_mem_tensor(self) -> None:
+        # CPU tensor -> False (no allocator registered for CPU)
+        t_cpu = torch.empty(1024)
+        self.assertFalse(symm_mem.is_symm_mem_tensor(t_cpu))
+
+        # Regular CUDA tensor -> False
+        t_cuda = torch.empty(1024, device="cuda")
+        self.assertFalse(symm_mem.is_symm_mem_tensor(t_cuda))
+
+        # symm-mem tensor
+        t_symm = symm_mem.empty(1024, device="cuda")
+        self.assertTrue(symm_mem.is_symm_mem_tensor(t_symm))
+
     @skipIf(
         not PLATFORM_SUPPORTS_SYMM_MEM, "SymmMem is not supported on this ROCm arch"
     )
@@ -413,11 +432,19 @@ def test_fused_all_gather_matmul(self, gather_dim: int) -> None:
             A_shard, Bs, gather_dim=gather_dim, group_name=group.group_name
         )
 
-        assert torch.allclose(ag_output_0, ag_output_1)
-        assert ag_output_0.stride() == ag_output_1.stride()
+        if not torch.allclose(ag_output_0, ag_output_1):
+            raise AssertionError("Expected ag_output_0 to be close to ag_output_1")
+        if not (ag_output_0.stride() == ag_output_1.stride()):
+            raise AssertionError(
+                f"Expected strides to match: {ag_output_0.stride()} vs {ag_output_1.stride()}"
+            )
         for mm_output_0, mm_output_1 in zip(mm_outputs_0, mm_outputs_1):
-            assert torch.allclose(mm_output_0, mm_output_1)
-            assert mm_output_0.stride(), mm_output_1.stride()
+            if not torch.allclose(mm_output_0, mm_output_1):
+                raise AssertionError("Expected mm_output_0 to be close to mm_output_1")
+            if not mm_output_0.stride():
+                raise AssertionError(
+                    f"Expected mm_output_0.stride() to be truthy, got {mm_output_0.stride()}"
+                )
 
     @skip_if_rocm_multiprocess  # this requires async_input_mm support
     @skipIf(
@@ -635,8 +662,12 @@ def test_fused_matmul_reduce_scatter(self, scatter_dim: int) -> None:
             A, B, "avg", scatter_dim=scatter_dim, group_name=group.group_name
         )
 
-        assert torch.allclose(output_0, output_1)
-        assert output_0.stride() == output_1.stride()
+        if not torch.allclose(output_0, output_1):
+            raise AssertionError("Expected output_0 to be close to output_1")
+        if not (output_0.stride() == output_1.stride()):
+            raise AssertionError(
+                f"Expected strides to match: {output_0.stride()} vs {output_1.stride()}"
+            )
 
     @skip_if_rocm_multiprocess  # AsyncTP support changed _fused_scaled_matmul_reduce_scatter_fallback API, need more changes
     @skip_if_lt_x_gpu(2)
@@ -690,7 +721,10 @@ def test_fused_scaled_matmul_reduce_scatter(
                     )
                 )
 
-        assert outputs[0].stride() == outputs[1].stride()
+        if not (outputs[0].stride() == outputs[1].stride()):
+            raise AssertionError(
+                f"Expected strides to match: {outputs[0].stride()} vs {outputs[1].stride()}"
+            )
         self.assertEqual(outputs[0], outputs[1])
 
     @skipIf(
@@ -1401,10 +1435,105 @@ def func_reuse_eager(x, w1, w2, w3):
             msg="Compiled and eager (reuse) outputs do not match",
         )
 
+    @skip_if_rocm_multiprocess  # requires registered-buffer support
+    @skip_if_lt_x_gpu(2)
+    @fresh_inductor_cache()
+    def test_output_buffer_reuse(self):
+        """
+        Verify that symm_mem ops are lowered via ExternKernelOut
+        (should_allocate()=True) so output buffers are visible to Inductor.
+        This is the foundation for subsequent PRs' P2P memory planning.
+        """
+        self._init_process()
+
+        N = 8
+        x = torch.rand(N, N, device=self.device)
+        w1 = torch.rand(N, N, device=self.device)
+        w2 = torch.rand(N, N, device=self.device)
+        w3 = torch.rand(N, N, device=self.device)
+
+        # Multi-layer TP pattern: mm -> allreduce, repeated 3 times.
+        def func(x, w1, w2, w3):
+            x = torch.mm(x, w1)
+            x = torch.ops.symm_mem.one_shot_all_reduce(x, "sum", "0")
+            x = torch.mm(x, w2)
+            x = torch.ops.symm_mem.one_shot_all_reduce(x, "sum", "0")
+            x = torch.mm(x, w3)
+            x = torch.ops.symm_mem.one_shot_all_reduce(x, "sum", "0")
+            return x
+
+        compiled = torch.compile(func, fullgraph=True)
+        code = run_and_get_triton_code(compiled, x, w1, w2, w3)
+
+        # Verify one_shot_all_reduce is converted to one_shot_all_reduce_out
+        # with out= parameter (ExternKernelOut, not FallbackKernel).
+        # Each mm + allreduce pair should produce an out= call for the allreduce.
+        FileCheck().check_count("one_shot_all_reduce_out", 3, exactly=True).run(code)
+
+        # Verify the out= parameter is present on allreduce calls.
+        out_calls = code.count(", out=")
+        self.assertGreaterEqual(
+            out_calls,
+            6,
+            f"Expected at least 6 out= calls (3 mm + 3 allreduce), got {out_calls}.",
+        )
+
+        # Output buffers should be reused across layers (ping-pong).
+        reuse_count = code.count("# reuse")
+        self.assertGreaterEqual(
+            reuse_count,
+            2,
+            f"Expected at least 2 buffer reuses, got {reuse_count}.",
+        )
+
+    @skip_if_rocm_multiprocess  # requires registered-buffer support
+    @skip_if_lt_x_gpu(2)
+    @fresh_inductor_cache()
+    def test_output_buffer_reuse_copy_variant(self):
+        """
+        Verify that one_shot_all_reduce_copy is also lowered via
+        ExternKernelOut through the manual out variant registry.
+        """
+        self._init_process()
+
+        N = 8
+        group_name = dist.group.WORLD.group_name
+        symm_buffer = symm_mem.empty(
+            N * N, device=self.device, dtype=torch.float32
+        ).view(N, N)
+        symm_mem.rendezvous(symm_buffer, group=group_name)
+
+        local_input = torch.rand(N, N, device=self.device)
+        w = torch.rand(N, N, device=self.device)
+
+        def func(symm_buffer, local_input, w):
+            x = torch.mm(local_input, w)
+            return torch.ops.symm_mem.one_shot_all_reduce_copy(
+                symm_buffer, x, "sum", "0"
+            )
+
+        compiled = torch.compile(func, fullgraph=True)
+        code = run_and_get_triton_code(compiled, symm_buffer, local_input, w)
+
+        # Verify one_shot_all_reduce_copy is converted to
+        # one_shot_all_reduce_copy_out with out= parameter.
+        FileCheck().check("one_shot_all_reduce_copy_out").run(code)
+
+        # Verify the out= parameter is present.
+        self.assertIn(
+            ", out=",
+            code,
+            "one_shot_all_reduce_copy_out should have out= parameter.",
+        )
+
     @skip_if_rocm_multiprocess  # test requires support for registered buffers
     @skip_if_lt_x_gpu(2)
     @fresh_inductor_cache()
     def test_external_allocation_fallback(self):
+        """
+        When the input is not pre-allocated in symmetric memory, Inductor
+        auto-inserts an identity copy to P2P.  Verify codegen + correctness.
+        """
         self._init_process()
 
         N = 8
@@ -1413,15 +1542,70 @@ def func_input_direct(x):
             return torch.ops.symm_mem.one_shot_all_reduce(x, "sum", "0")
 
         x_input = torch.rand(N, N, device=self.device)
-
-        # Compilation should succeed here even though we do not control allocation
-        # of the input; user is responsible for passing a symmetric memory buffer.
         compiled_input_direct = torch.compile(func_input_direct, fullgraph=True)
+        code = run_and_get_triton_code(compiled_input_direct, x_input)
 
-        # At runtime, this should raise an error because the input is not
-        # a symmetric memory buffer as required by the op.
-        with self.assertRaises(RuntimeError):
-            run_and_get_triton_code(compiled_input_direct, x_input)
+        self.assertIn(
+            "empty_strided_p2p",
+            code,
+            "Expected P2P allocation for auto-inserted copy",
+        )
+
+        compiled_result = compiled_input_direct(x_input)
+        eager_result = x_input.clone()
+        dist.all_reduce(eager_result, op=dist.ReduceOp.SUM)
+        torch.testing.assert_close(
+            compiled_result,
+            eager_result,
+            rtol=1e-5,
+            atol=1e-5,
+            msg="Auto-copy to P2P does not match eager",
+        )
+
+    @skip_if_rocm_multiprocess  # requires registered-buffer support
+    @skip_if_lt_x_gpu(2)
+    @fresh_inductor_cache()
+    def test_comm_buffer_inplace_prevention(self):
+        """Verify scheduler correctly handles inplace for CommBufferLayout buffers."""
+        self._init_process()
+
+        N = 8
+        x = torch.rand(N, N, device=self.device)
+        w = torch.rand(N, N, device=self.device)
+
+        # mm (ExternKernelOut, regular CUDA) -> pointwise (CommBufferLayout) -> allreduce
+        # Should not in-places pointwise mul(P2P) into mm's buffer(regular).
+        def func(x, w):
+            y = torch.mm(x, w)
+            z = y * 2
+            return torch.ops.symm_mem.one_shot_all_reduce(z, "sum", "0")
+
+        compiled = torch.compile(func, fullgraph=True)
+        code = run_and_get_triton_code(compiled, x, w)
+
+        self.assertIn(
+            "empty_strided_p2p",
+            code,
+            "Expected P2P allocation in generated code",
+        )
+        self.assertIn(
+            "one_shot_all_reduce_out",
+            code,
+            "Expected out-variant allreduce in generated code",
+        )
+
+        result = compiled(x, w)
+        eager_y = torch.mm(x, w)
+        eager_z = eager_y * 2
+        eager_result = eager_z.clone()
+        dist.all_reduce(eager_result, op=dist.ReduceOp.SUM)
+        torch.testing.assert_close(
+            result,
+            eager_result,
+            rtol=1e-5,
+            atol=1e-5,
+            msg="Compiled and eager do not match",
+        )
 
 
 class SymmMemSingleProcTest(TestCase):
diff --git a/test/distributions/test_distributions.py b/test/distributions/test_distributions.py
index d7099bed66f97..57e3331847720 100644
--- a/test/distributions/test_distributions.py
+++ b/test/distributions/test_distributions.py
@@ -1821,26 +1821,30 @@ def test_zero_excluded_binomial(self):
 
     def test_torch_binomial_dtype_errors(self):
         dtypes = [torch.int, torch.long, torch.short]
+        devices = ["cpu"]
+        if TEST_CUDA:
+            devices.append("cuda")
+
+        for device in devices:
+            for count_dtype in dtypes:
+                total_count = torch.tensor([10, 10], dtype=count_dtype, device=device)
+                total_prob = torch.tensor([0.5, 0.5], dtype=torch.float, device=device)
+
+                with self.assertRaisesRegex(
+                    ValueError,
+                    "binomial only supports floating-point dtypes for count.*",
+                ):
+                    torch.binomial(total_count, total_prob)
 
-        for count_dtype in dtypes:
-            total_count = torch.tensor([10, 10], dtype=count_dtype)
-            total_prob = torch.tensor([0.5, 0.5], dtype=torch.float)
-
-            with self.assertRaisesRegex(
-                ValueError,
-                "binomial only supports floating-point dtypes for count.*",
-            ):
-                torch.binomial(total_count, total_prob)
-
-        for prob_dtype in dtypes:
-            total_count = torch.tensor([10, 10], dtype=torch.float)
-            total_prob = torch.tensor([0.5, 0.5], dtype=prob_dtype)
+            for prob_dtype in dtypes:
+                total_count = torch.tensor([10, 10], dtype=torch.float, device=device)
+                total_prob = torch.tensor([0.5, 0.5], dtype=prob_dtype, device=device)
 
-            with self.assertRaisesRegex(
-                ValueError,
-                "binomial only supports floating-point dtypes for prob.*",
-            ):
-                torch.binomial(total_count, total_prob)
+                with self.assertRaisesRegex(
+                    ValueError,
+                    "binomial only supports floating-point dtypes for prob.*",
+                ):
+                    torch.binomial(total_count, total_prob)
 
     @set_default_dtype(torch.double)
     def test_multinomial_1d(self):
diff --git a/test/dynamo/_test_nested_graph_breaks_helper.py b/test/dynamo/_test_nested_graph_breaks_helper.py
index ea229524d21bc..bafa1928c0374 100644
--- a/test/dynamo/_test_nested_graph_breaks_helper.py
+++ b/test/dynamo/_test_nested_graph_breaks_helper.py
@@ -16,3 +16,16 @@ def fn(val, call):
     val = call(val)
     val = val + 1
     return val
+
+
+HELPER_CONSTANT = torch.tensor([100.0])
+
+
+def closure_with_graph_break(x):
+    captured = x + 1
+
+    def inner():
+        torch._dynamo.graph_break()
+        return captured + HELPER_CONSTANT
+
+    return inner()
diff --git a/test/dynamo/cpython/3_13/test_class.diff b/test/dynamo/cpython/3_13/test_class.diff
new file mode 100644
index 0000000000000..f6f691a41ec56
--- /dev/null
+++ b/test/dynamo/cpython/3_13/test_class.diff
@@ -0,0 +1,447 @@
+diff --git a/test/dynamo/cpython/3_13/test_class.py b/test/dynamo/cpython/3_13/test_class.py
+index d2df6228ec4..316fda7d7db 100644
+--- a/test/dynamo/cpython/3_13/test_class.py
++++ b/test/dynamo/cpython/3_13/test_class.py
+@@ -1,6 +1,25 @@
++# ======= BEGIN Dynamo patch =======
++# Owner(s): ["module: dynamo"]
++
++# ruff: noqa
++# flake8: noqa
++
++# Test copied from
++# https://raw.githubusercontent.com/python/cpython/refs/tags/v3.13.5/Lib/test/test_class.py
++
++import sys
++import torch
++import torch._dynamo.test_case
++import unittest
++from torch._dynamo.test_case import CPythonTestCase
++from torch.testing._internal.common_utils import run_tests, skipIfTorchDynamo
++
++# ======= END DYNAMO PATCH =======
++
+ "Test the functionality of Python classes implementing operators."
+ 
+ import unittest
++from test import support
+ from test.support import cpython_only, import_helper, script_helper
+ 
+ testmeths = [
+@@ -134,8 +153,9 @@ for method in testmeths:
+ AllTests = type("AllTests", (object,), d)
+ del d, statictests, method, method_template
+ 
+-class ClassTests(unittest.TestCase):
++class ClassTests(CPythonTestCase):
+     def setUp(self):
++        super().setUp()
+         callLst[:] = []
+ 
+     def assertCallStack(self, expected_calls):
+@@ -273,7 +293,8 @@ class ClassTests(unittest.TestCase):
+ 
+         # List/dict operations
+ 
+-        class Empty: pass
++        with torch._dynamo.error_on_graph_break(False):
++            class Empty: pass
+ 
+         try:
+             1 in Empty()
+@@ -417,18 +438,19 @@ class ClassTests(unittest.TestCase):
+ 
+     def testGetSetAndDel(self):
+         # Interfering tests
+-        class ExtraTests(AllTests):
+-            @trackCall
+-            def __getattr__(self, *args):
+-                return "SomeVal"
++        with torch._dynamo.error_on_graph_break(False):
++            class ExtraTests(AllTests):
++                @trackCall
++                def __getattr__(self, *args):
++                    return "SomeVal"
+ 
+-            @trackCall
+-            def __setattr__(self, *args):
+-                pass
++                @trackCall
++                def __setattr__(self, *args):
++                    pass
+ 
+-            @trackCall
+-            def __delattr__(self, *args):
+-                pass
++                @trackCall
++                def __delattr__(self, *args):
++                    pass
+ 
+         testme = ExtraTests()
+ 
+@@ -450,9 +472,10 @@ class ClassTests(unittest.TestCase):
+         from test.support import import_helper
+         _testlimitedcapi = import_helper.import_module('_testlimitedcapi')
+ 
+-        class A:
+-            def __init__(self):
+-                self.attr = 1
++        with torch._dynamo.error_on_graph_break(False):
++            class A:
++                def __init__(self):
++                    self.attr = 1
+ 
+         a = A()
+         self.assertEqual(_testlimitedcapi.object_hasattrstring(a, b"attr"), 1)
+@@ -473,16 +496,17 @@ class ClassTests(unittest.TestCase):
+ 
+     def testBadTypeReturned(self):
+         # return values of some method are type-checked
+-        class BadTypeClass:
+-            def __int__(self):
+-                return None
+-            __float__ = __int__
+-            __complex__ = __int__
+-            __str__ = __int__
+-            __repr__ = __int__
+-            __bytes__ = __int__
+-            __bool__ = __int__
+-            __index__ = __int__
++        with torch._dynamo.error_on_graph_break(False):
++            class BadTypeClass:
++                def __int__(self):
++                    return None
++                __float__ = __int__
++                __complex__ = __int__
++                __str__ = __int__
++                __repr__ = __int__
++                __bytes__ = __int__
++                __bool__ = __int__
++                __index__ = __int__
+         def index(x):
+             return [][x]
+ 
+@@ -493,21 +517,24 @@ class ClassTests(unittest.TestCase):
+         # Test correct errors from hash() on objects with comparisons but
+         #  no __hash__
+ 
+-        class C0:
+-            pass
++        with torch._dynamo.error_on_graph_break(False):
++            class C0:
++                pass
+ 
+         hash(C0()) # This should work; the next two should raise TypeError
+ 
+-        class C2:
+-            def __eq__(self, other): return 1
++        with torch._dynamo.error_on_graph_break(False):
++            class C2:
++                def __eq__(self, other): return 1
+ 
+         self.assertRaises(TypeError, hash, C2())
+ 
+     def testPredefinedAttrs(self):
+         o = object()
+ 
+-        class Custom:
+-            pass
++        with torch._dynamo.error_on_graph_break(False):
++            class Custom:
++                pass
+ 
+         c = Custom()
+ 
+@@ -557,8 +584,9 @@ class ClassTests(unittest.TestCase):
+     def testSFBug532646(self):
+         # Test for SF bug 532646
+ 
+-        class A:
+-            pass
++        with torch._dynamo.error_on_graph_break(False):
++            class A:
++                pass
+         A.__call__ = A()
+         a = A()
+ 
+@@ -575,20 +603,23 @@ class ClassTests(unittest.TestCase):
+         def booh(self):
+             raise AttributeError("booh")
+ 
+-        class A:
+-            a = property(booh)
++        with torch._dynamo.error_on_graph_break(False):
++            class A:
++                a = property(booh)
+         try:
+             A().a # Raised AttributeError: A instance has no attribute 'a'
+         except AttributeError as x:
+             if str(x) != "booh":
+                 self.fail("attribute error for A().a got masked: %s" % x)
+ 
+-        class E:
+-            __eq__ = property(booh)
++        with torch._dynamo.error_on_graph_break(False):
++            class E:
++                __eq__ = property(booh)
+         E() == E() # In debug mode, caused a C-level assert() to fail
+ 
+-        class I:
+-            __init__ = property(booh)
++        with torch._dynamo.error_on_graph_break(False):
++            class I:
++                __init__ = property(booh)
+         try:
+             # In debug mode, printed XXX undetected error and
+             #  raises AttributeError
+@@ -610,19 +641,21 @@ class ClassTests(unittest.TestCase):
+ 
+     def testHashComparisonOfMethods(self):
+         # Test comparison and hash of methods
+-        class A:
+-            def __init__(self, x):
+-                self.x = x
+-            def f(self):
+-                pass
+-            def g(self):
++        with torch._dynamo.error_on_graph_break(False):
++            class A:
++                def __init__(self, x):
++                    self.x = x
++                def f(self):
++                    pass
++                def g(self):
++                    pass
++                def __eq__(self, other):
++                    return True
++                def __hash__(self):
++                    raise TypeError
++        with torch._dynamo.error_on_graph_break(False):
++            class B(A):
+                 pass
+-            def __eq__(self, other):
+-                return True
+-            def __hash__(self):
+-                raise TypeError
+-        class B(A):
+-            pass
+ 
+         a1 = A(1)
+         a2 = A(1)
+@@ -650,8 +683,9 @@ class ClassTests(unittest.TestCase):
+ 
+     def testSetattrWrapperNameIntern(self):
+         # Issue #25794: __setattr__ should intern the attribute name
+-        class A:
+-            pass
++        with torch._dynamo.error_on_graph_break(False):
++            class A:
++                pass
+ 
+         def add(self, other):
+             return 'summa'
+@@ -669,15 +703,17 @@ class ClassTests(unittest.TestCase):
+             A() + 1
+ 
+     def testSetattrNonStringName(self):
+-        class A:
+-            pass
++        with torch._dynamo.error_on_graph_break(False):
++            class A:
++                pass
+ 
+         with self.assertRaises(TypeError):
+             type.__setattr__(A, b'x', None)
+ 
+     def testTypeAttributeAccessErrorMessages(self):
+-        class A:
+-            pass
++        with torch._dynamo.error_on_graph_break(False):
++            class A:
++                pass
+ 
+         error_msg = "type object 'A' has no attribute 'x'"
+         with self.assertRaisesRegex(AttributeError, error_msg):
+@@ -686,19 +722,20 @@ class ClassTests(unittest.TestCase):
+             del A.x
+ 
+     def testObjectAttributeAccessErrorMessages(self):
+-        class A:
+-            pass
+-        class B:
+-            y = 0
+-            __slots__ = ('z',)
+-        class C:
+-            __slots__ = ("y",)
+-
+-            def __setattr__(self, name, value) -> None:
+-                if name == "z":
+-                    super().__setattr__("y", 1)
+-                else:
+-                    super().__setattr__(name, value)
++        with torch._dynamo.error_on_graph_break(False):
++            class A:
++                pass
++            class B:
++                y = 0
++                __slots__ = ('z',)
++            class C:
++                __slots__ = ("y",)
++
++                def __setattr__(self, name, value) -> None:
++                    if name == "z":
++                        super().__setattr__("y", 1)
++                    else:
++                        super().__setattr__(name, value)
+ 
+         error_msg = "'A' object has no attribute 'x'"
+         with self.assertRaisesRegex(AttributeError, error_msg):
+@@ -738,8 +775,9 @@ class ClassTests(unittest.TestCase):
+         # bpo-31506: Improves the error message logic for object_new & object_init
+ 
+         # Class without any method overrides
+-        class C:
+-            pass
++        with torch._dynamo.error_on_graph_break(False):
++            class C:
++                pass
+ 
+         error_msg = r'C.__init__\(\) takes exactly one argument \(the instance to initialize\)'
+ 
+@@ -759,11 +797,12 @@ class ClassTests(unittest.TestCase):
+             object.__init__(C(), 42)
+ 
+         # Class with both `__init__` & `__new__` method overridden
+-        class D:
+-            def __new__(cls, *args, **kwargs):
+-                super().__new__(cls, *args, **kwargs)
+-            def __init__(self, *args, **kwargs):
+-                super().__init__(*args, **kwargs)
++        with torch._dynamo.error_on_graph_break(False):
++            class D:
++                def __new__(cls, *args, **kwargs):
++                    super().__new__(cls, *args, **kwargs)
++                def __init__(self, *args, **kwargs):
++                    super().__init__(*args, **kwargs)
+ 
+         error_msg =  r'object.__new__\(\) takes exactly one argument \(the type to instantiate\)'
+ 
+@@ -777,9 +816,10 @@ class ClassTests(unittest.TestCase):
+             object.__new__(D, 42)
+ 
+         # Class that only overrides __init__
+-        class E:
+-            def __init__(self, *args, **kwargs):
+-                super().__init__(*args, **kwargs)
++        with torch._dynamo.error_on_graph_break(False):
++            class E:
++                def __init__(self, *args, **kwargs):
++                    super().__init__(*args, **kwargs)
+ 
+         error_msg = r'object.__init__\(\) takes exactly one argument \(the instance to initialize\)'
+ 
+@@ -790,20 +830,24 @@ class ClassTests(unittest.TestCase):
+             object.__init__(E(), 42)
+ 
+     def testClassWithExtCall(self):
+-        class Meta(int):
+-            def __init__(*args, **kwargs):
+-                pass
++        with torch._dynamo.error_on_graph_break(False):
++            class Meta(int):
++                def __init__(*args, **kwargs):
++                    pass
+ 
+-            def __new__(cls, name, bases, attrs, **kwargs):
+-                return bases, kwargs
++                def __new__(cls, name, bases, attrs, **kwargs):
++                    return bases, kwargs
+ 
+         d = {'metaclass': Meta}
+ 
+-        class A(**d): pass
++        with torch._dynamo.error_on_graph_break(False):
++            class A(**d): pass
+         self.assertEqual(A, ((), {}))
+-        class A(0, 1, 2, 3, 4, 5, 6, 7, **d): pass
++        with torch._dynamo.error_on_graph_break(False):
++            class A(0, 1, 2, 3, 4, 5, 6, 7, **d): pass
+         self.assertEqual(A, (tuple(range(8)), {}))
+-        class A(0, *range(1, 8), **d, foo='bar'): pass
++        with torch._dynamo.error_on_graph_break(False):
++            class A(0, *range(1, 8), **d, foo='bar'): pass
+         self.assertEqual(A, (tuple(range(8)), {'foo': 'bar'}))
+ 
+     def testClassCallRecursionLimit(self):
+@@ -824,20 +868,35 @@ class ClassTests(unittest.TestCase):
+     def testMetaclassCallOptimization(self):
+         calls = 0
+ 
+-        class TypeMetaclass(type):
+-            def __call__(cls, *args, **kwargs):
+-                nonlocal calls
+-                calls += 1
+-                return type.__call__(cls, *args, **kwargs)
++        with torch._dynamo.error_on_graph_break(False):
++            class TypeMetaclass(type):
++                def __call__(cls, *args, **kwargs):
++                    nonlocal calls
++                    calls += 1
++                    return type.__call__(cls, *args, **kwargs)
+ 
+-        class Type(metaclass=TypeMetaclass):
+-            def __init__(self, obj):
+-                self._obj = obj
++            class Type(metaclass=TypeMetaclass):
++                def __init__(self, obj):
++                    self._obj = obj
+ 
+         for i in range(100):
+             Type(i)
+         self.assertEqual(calls, 100)
+ 
++    def test_specialization_class_call_doesnt_crash(self):
++        # gh-123185
++
++        with torch._dynamo.error_on_graph_break(False):
++            class Foo:
++                def __init__(self, arg):
++                    pass
++
++        for _ in range(8):
++            try:
++                Foo()
++            except:
++                pass
++
+ 
+ from _testinternalcapi import has_inline_values
+ 
+@@ -856,7 +915,8 @@ class WithAttrs:
+         self.d = 4
+ 
+ 
+-class TestInlineValues(unittest.TestCase):
++@skipIfTorchDynamo("CPython only")
++class TestInlineValues(CPythonTestCase):
+ 
+     def test_flags(self):
+         self.assertEqual(Plain.__flags__ & Py_TPFLAGS_MANAGED_DICT, Py_TPFLAGS_MANAGED_DICT)
+@@ -895,7 +955,8 @@ class TestInlineValues(unittest.TestCase):
+             self.assertEqual(getattr(obj, f"a{i}"), i)
+ 
+     def test_many_attributes(self):
+-        class C: pass
++        with torch._dynamo.error_on_graph_break(False):
++            class C: pass
+         c = C()
+         self.assertTrue(has_inline_values(c))
+         self.set_100(c)
+@@ -1001,5 +1062,5 @@ class TestInlineValues(unittest.TestCase):
+         self.assertFalse(out, msg=out.decode('utf-8'))
+         self.assertFalse(err, msg=err.decode('utf-8'))
+ 
+-if __name__ == '__main__':
+-    unittest.main()
++if __name__ == "__main__":
++    run_tests()
diff --git a/test/dynamo/cpython/3_13/test_class.py b/test/dynamo/cpython/3_13/test_class.py
new file mode 100644
index 0000000000000..1624de9094713
--- /dev/null
+++ b/test/dynamo/cpython/3_13/test_class.py
@@ -0,0 +1,1069 @@
+# ======= BEGIN Dynamo patch =======
+# Owner(s): ["module: dynamo"]
+
+# ruff: noqa
+# flake8: noqa
+
+# Test copied from
+# https://raw.githubusercontent.com/python/cpython/refs/tags/v3.13.5/Lib/test/test_class.py
+
+import sys
+import torch
+import torch._dynamo.test_case
+import unittest
+from torch._dynamo.test_case import CPythonTestCase
+from torch.testing._internal.common_utils import run_tests, skipIfTorchDynamo
+
+# ======= END DYNAMO PATCH =======
+
+"Test the functionality of Python classes implementing operators."
+
+import unittest
+from test import support
+from test.support import cpython_only, import_helper, script_helper
+
+testmeths = [
+
+# Binary operations
+    "add",
+    "radd",
+    "sub",
+    "rsub",
+    "mul",
+    "rmul",
+    "matmul",
+    "rmatmul",
+    "truediv",
+    "rtruediv",
+    "floordiv",
+    "rfloordiv",
+    "mod",
+    "rmod",
+    "divmod",
+    "rdivmod",
+    "pow",
+    "rpow",
+    "rshift",
+    "rrshift",
+    "lshift",
+    "rlshift",
+    "and",
+    "rand",
+    "or",
+    "ror",
+    "xor",
+    "rxor",
+
+# List/dict operations
+    "contains",
+    "getitem",
+    "setitem",
+    "delitem",
+
+# Unary operations
+    "neg",
+    "pos",
+    "abs",
+
+# generic operations
+    "init",
+    ]
+
+# These need to return something other than None
+#    "hash",
+#    "str",
+#    "repr",
+#    "int",
+#    "float",
+
+# These are separate because they can influence the test of other methods.
+#    "getattr",
+#    "setattr",
+#    "delattr",
+
+callLst = []
+def trackCall(f):
+    def track(*args, **kwargs):
+        callLst.append((f.__name__, args))
+        return f(*args, **kwargs)
+    return track
+
+statictests = """
+@trackCall
+def __hash__(self, *args):
+    return hash(id(self))
+
+@trackCall
+def __str__(self, *args):
+    return "AllTests"
+
+@trackCall
+def __repr__(self, *args):
+    return "AllTests"
+
+@trackCall
+def __int__(self, *args):
+    return 1
+
+@trackCall
+def __index__(self, *args):
+    return 1
+
+@trackCall
+def __float__(self, *args):
+    return 1.0
+
+@trackCall
+def __eq__(self, *args):
+    return True
+
+@trackCall
+def __ne__(self, *args):
+    return False
+
+@trackCall
+def __lt__(self, *args):
+    return False
+
+@trackCall
+def __le__(self, *args):
+    return True
+
+@trackCall
+def __gt__(self, *args):
+    return False
+
+@trackCall
+def __ge__(self, *args):
+    return True
+"""
+
+# Synthesize all the other AllTests methods from the names in testmeths.
+
+method_template = """\
+@trackCall
+def __%s__(self, *args):
+    pass
+"""
+
+d = {}
+exec(statictests, globals(), d)
+for method in testmeths:
+    exec(method_template % method, globals(), d)
+AllTests = type("AllTests", (object,), d)
+del d, statictests, method, method_template
+
+class ClassTests(CPythonTestCase):
+    def setUp(self):
+        super().setUp()
+        callLst[:] = []
+
+    def assertCallStack(self, expected_calls):
+        actualCallList = callLst[:]  # need to copy because the comparison below will add
+                                     # additional calls to callLst
+        if expected_calls != actualCallList:
+            self.fail("Expected call list:\n  %s\ndoes not match actual call list\n  %s" %
+                      (expected_calls, actualCallList))
+
+    def testInit(self):
+        foo = AllTests()
+        self.assertCallStack([("__init__", (foo,))])
+
+    def testBinaryOps(self):
+        testme = AllTests()
+        # Binary operations
+
+        callLst[:] = []
+        testme + 1
+        self.assertCallStack([("__add__", (testme, 1))])
+
+        callLst[:] = []
+        1 + testme
+        self.assertCallStack([("__radd__", (testme, 1))])
+
+        callLst[:] = []
+        testme - 1
+        self.assertCallStack([("__sub__", (testme, 1))])
+
+        callLst[:] = []
+        1 - testme
+        self.assertCallStack([("__rsub__", (testme, 1))])
+
+        callLst[:] = []
+        testme * 1
+        self.assertCallStack([("__mul__", (testme, 1))])
+
+        callLst[:] = []
+        1 * testme
+        self.assertCallStack([("__rmul__", (testme, 1))])
+
+        callLst[:] = []
+        testme @ 1
+        self.assertCallStack([("__matmul__", (testme, 1))])
+
+        callLst[:] = []
+        1 @ testme
+        self.assertCallStack([("__rmatmul__", (testme, 1))])
+
+        callLst[:] = []
+        testme / 1
+        self.assertCallStack([("__truediv__", (testme, 1))])
+
+
+        callLst[:] = []
+        1 / testme
+        self.assertCallStack([("__rtruediv__", (testme, 1))])
+
+        callLst[:] = []
+        testme // 1
+        self.assertCallStack([("__floordiv__", (testme, 1))])
+
+
+        callLst[:] = []
+        1 // testme
+        self.assertCallStack([("__rfloordiv__", (testme, 1))])
+
+        callLst[:] = []
+        testme % 1
+        self.assertCallStack([("__mod__", (testme, 1))])
+
+        callLst[:] = []
+        1 % testme
+        self.assertCallStack([("__rmod__", (testme, 1))])
+
+
+        callLst[:] = []
+        divmod(testme,1)
+        self.assertCallStack([("__divmod__", (testme, 1))])
+
+        callLst[:] = []
+        divmod(1, testme)
+        self.assertCallStack([("__rdivmod__", (testme, 1))])
+
+        callLst[:] = []
+        testme ** 1
+        self.assertCallStack([("__pow__", (testme, 1))])
+
+        callLst[:] = []
+        1 ** testme
+        self.assertCallStack([("__rpow__", (testme, 1))])
+
+        callLst[:] = []
+        testme >> 1
+        self.assertCallStack([("__rshift__", (testme, 1))])
+
+        callLst[:] = []
+        1 >> testme
+        self.assertCallStack([("__rrshift__", (testme, 1))])
+
+        callLst[:] = []
+        testme << 1
+        self.assertCallStack([("__lshift__", (testme, 1))])
+
+        callLst[:] = []
+        1 << testme
+        self.assertCallStack([("__rlshift__", (testme, 1))])
+
+        callLst[:] = []
+        testme & 1
+        self.assertCallStack([("__and__", (testme, 1))])
+
+        callLst[:] = []
+        1 & testme
+        self.assertCallStack([("__rand__", (testme, 1))])
+
+        callLst[:] = []
+        testme | 1
+        self.assertCallStack([("__or__", (testme, 1))])
+
+        callLst[:] = []
+        1 | testme
+        self.assertCallStack([("__ror__", (testme, 1))])
+
+        callLst[:] = []
+        testme ^ 1
+        self.assertCallStack([("__xor__", (testme, 1))])
+
+        callLst[:] = []
+        1 ^ testme
+        self.assertCallStack([("__rxor__", (testme, 1))])
+
+    def testListAndDictOps(self):
+        testme = AllTests()
+
+        # List/dict operations
+
+        with torch._dynamo.error_on_graph_break(False):
+            class Empty: pass
+
+        try:
+            1 in Empty()
+            self.fail('failed, should have raised TypeError')
+        except TypeError:
+            pass
+
+        callLst[:] = []
+        1 in testme
+        self.assertCallStack([('__contains__', (testme, 1))])
+
+        callLst[:] = []
+        testme[1]
+        self.assertCallStack([('__getitem__', (testme, 1))])
+
+        callLst[:] = []
+        testme[1] = 1
+        self.assertCallStack([('__setitem__', (testme, 1, 1))])
+
+        callLst[:] = []
+        del testme[1]
+        self.assertCallStack([('__delitem__', (testme, 1))])
+
+        callLst[:] = []
+        testme[:42]
+        self.assertCallStack([('__getitem__', (testme, slice(None, 42)))])
+
+        callLst[:] = []
+        testme[:42] = "The Answer"
+        self.assertCallStack([('__setitem__', (testme, slice(None, 42),
+                                               "The Answer"))])
+
+        callLst[:] = []
+        del testme[:42]
+        self.assertCallStack([('__delitem__', (testme, slice(None, 42)))])
+
+        callLst[:] = []
+        testme[2:1024:10]
+        self.assertCallStack([('__getitem__', (testme, slice(2, 1024, 10)))])
+
+        callLst[:] = []
+        testme[2:1024:10] = "A lot"
+        self.assertCallStack([('__setitem__', (testme, slice(2, 1024, 10),
+                                                                    "A lot"))])
+        callLst[:] = []
+        del testme[2:1024:10]
+        self.assertCallStack([('__delitem__', (testme, slice(2, 1024, 10)))])
+
+        callLst[:] = []
+        testme[:42, ..., :24:, 24, 100]
+        self.assertCallStack([('__getitem__', (testme, (slice(None, 42, None),
+                                                        Ellipsis,
+                                                        slice(None, 24, None),
+                                                        24, 100)))])
+        callLst[:] = []
+        testme[:42, ..., :24:, 24, 100] = "Strange"
+        self.assertCallStack([('__setitem__', (testme, (slice(None, 42, None),
+                                                        Ellipsis,
+                                                        slice(None, 24, None),
+                                                        24, 100), "Strange"))])
+        callLst[:] = []
+        del testme[:42, ..., :24:, 24, 100]
+        self.assertCallStack([('__delitem__', (testme, (slice(None, 42, None),
+                                                        Ellipsis,
+                                                        slice(None, 24, None),
+                                                        24, 100)))])
+
+    def testUnaryOps(self):
+        testme = AllTests()
+
+        callLst[:] = []
+        -testme
+        self.assertCallStack([('__neg__', (testme,))])
+        callLst[:] = []
+        +testme
+        self.assertCallStack([('__pos__', (testme,))])
+        callLst[:] = []
+        abs(testme)
+        self.assertCallStack([('__abs__', (testme,))])
+        callLst[:] = []
+        int(testme)
+        self.assertCallStack([('__int__', (testme,))])
+        callLst[:] = []
+        float(testme)
+        self.assertCallStack([('__float__', (testme,))])
+        callLst[:] = []
+        oct(testme)
+        self.assertCallStack([('__index__', (testme,))])
+        callLst[:] = []
+        hex(testme)
+        self.assertCallStack([('__index__', (testme,))])
+
+
+    def testMisc(self):
+        testme = AllTests()
+
+        callLst[:] = []
+        hash(testme)
+        self.assertCallStack([('__hash__', (testme,))])
+
+        callLst[:] = []
+        repr(testme)
+        self.assertCallStack([('__repr__', (testme,))])
+
+        callLst[:] = []
+        str(testme)
+        self.assertCallStack([('__str__', (testme,))])
+
+        callLst[:] = []
+        testme == 1
+        self.assertCallStack([('__eq__', (testme, 1))])
+
+        callLst[:] = []
+        testme < 1
+        self.assertCallStack([('__lt__', (testme, 1))])
+
+        callLst[:] = []
+        testme > 1
+        self.assertCallStack([('__gt__', (testme, 1))])
+
+        callLst[:] = []
+        testme != 1
+        self.assertCallStack([('__ne__', (testme, 1))])
+
+        callLst[:] = []
+        1 == testme
+        self.assertCallStack([('__eq__', (1, testme))])
+
+        callLst[:] = []
+        1 < testme
+        self.assertCallStack([('__gt__', (1, testme))])
+
+        callLst[:] = []
+        1 > testme
+        self.assertCallStack([('__lt__', (1, testme))])
+
+        callLst[:] = []
+        1 != testme
+        self.assertCallStack([('__ne__', (1, testme))])
+
+
+    def testGetSetAndDel(self):
+        # Interfering tests
+        with torch._dynamo.error_on_graph_break(False):
+            class ExtraTests(AllTests):
+                @trackCall
+                def __getattr__(self, *args):
+                    return "SomeVal"
+
+                @trackCall
+                def __setattr__(self, *args):
+                    pass
+
+                @trackCall
+                def __delattr__(self, *args):
+                    pass
+
+        testme = ExtraTests()
+
+        callLst[:] = []
+        testme.spam
+        self.assertCallStack([('__getattr__', (testme, "spam"))])
+
+        callLst[:] = []
+        testme.eggs = "spam, spam, spam and ham"
+        self.assertCallStack([('__setattr__', (testme, "eggs",
+                                               "spam, spam, spam and ham"))])
+
+        callLst[:] = []
+        del testme.cardinal
+        self.assertCallStack([('__delattr__', (testme, "cardinal"))])
+
+    def testHasAttrString(self):
+        import sys
+        from test.support import import_helper
+        _testlimitedcapi = import_helper.import_module('_testlimitedcapi')
+
+        with torch._dynamo.error_on_graph_break(False):
+            class A:
+                def __init__(self):
+                    self.attr = 1
+
+        a = A()
+        self.assertEqual(_testlimitedcapi.object_hasattrstring(a, b"attr"), 1)
+        self.assertEqual(_testlimitedcapi.object_hasattrstring(a, b"noattr"), 0)
+        self.assertIsNone(sys.exception())
+
+    def testDel(self):
+        x = []
+
+        class DelTest:
+            def __del__(self):
+                x.append("crab people, crab people")
+        testme = DelTest()
+        del testme
+        import gc
+        gc.collect()
+        self.assertEqual(["crab people, crab people"], x)
+
+    def testBadTypeReturned(self):
+        # return values of some method are type-checked
+        with torch._dynamo.error_on_graph_break(False):
+            class BadTypeClass:
+                def __int__(self):
+                    return None
+                __float__ = __int__
+                __complex__ = __int__
+                __str__ = __int__
+                __repr__ = __int__
+                __bytes__ = __int__
+                __bool__ = __int__
+                __index__ = __int__
+        def index(x):
+            return [][x]
+
+        for f in [float, complex, str, repr, bytes, bin, oct, hex, bool, index]:
+            self.assertRaises(TypeError, f, BadTypeClass())
+
+    def testHashStuff(self):
+        # Test correct errors from hash() on objects with comparisons but
+        #  no __hash__
+
+        with torch._dynamo.error_on_graph_break(False):
+            class C0:
+                pass
+
+        hash(C0()) # This should work; the next two should raise TypeError
+
+        with torch._dynamo.error_on_graph_break(False):
+            class C2:
+                def __eq__(self, other): return 1
+
+        self.assertRaises(TypeError, hash, C2())
+
+    def testPredefinedAttrs(self):
+        o = object()
+
+        with torch._dynamo.error_on_graph_break(False):
+            class Custom:
+                pass
+
+        c = Custom()
+
+        methods = (
+            '__class__', '__delattr__', '__dir__', '__eq__', '__format__',
+            '__ge__', '__getattribute__', '__getstate__', '__gt__', '__hash__',
+            '__init__', '__init_subclass__', '__le__', '__lt__', '__ne__',
+            '__new__', '__reduce__', '__reduce_ex__', '__repr__',
+            '__setattr__', '__sizeof__', '__str__', '__subclasshook__'
+        )
+        for name in methods:
+            with self.subTest(name):
+                self.assertTrue(callable(getattr(object, name, None)))
+                self.assertTrue(callable(getattr(o, name, None)))
+                self.assertTrue(callable(getattr(Custom, name, None)))
+                self.assertTrue(callable(getattr(c, name, None)))
+
+        not_defined = [
+            '__abs__', '__aenter__', '__aexit__', '__aiter__', '__anext__',
+            '__await__', '__bool__', '__bytes__', '__ceil__',
+            '__complex__', '__contains__', '__del__', '__delete__',
+            '__delitem__', '__divmod__', '__enter__', '__exit__',
+            '__float__', '__floor__', '__get__', '__getattr__', '__getitem__',
+            '__index__', '__int__', '__invert__', '__iter__', '__len__',
+            '__length_hint__', '__missing__', '__neg__', '__next__',
+            '__objclass__', '__pos__', '__rdivmod__', '__reversed__',
+            '__round__', '__set__', '__setitem__', '__trunc__'
+        ]
+        augment = (
+            'add', 'and', 'floordiv', 'lshift', 'matmul', 'mod', 'mul', 'pow',
+            'rshift', 'sub', 'truediv', 'xor'
+        )
+        not_defined.extend(map("__{}__".format, augment))
+        not_defined.extend(map("__r{}__".format, augment))
+        not_defined.extend(map("__i{}__".format, augment))
+        for name in not_defined:
+            with self.subTest(name):
+                self.assertFalse(hasattr(object, name))
+                self.assertFalse(hasattr(o, name))
+                self.assertFalse(hasattr(Custom, name))
+                self.assertFalse(hasattr(c, name))
+
+        # __call__() is defined on the metaclass but not the class
+        self.assertFalse(hasattr(o, "__call__"))
+        self.assertFalse(hasattr(c, "__call__"))
+
+    def testSFBug532646(self):
+        # Test for SF bug 532646
+
+        with torch._dynamo.error_on_graph_break(False):
+            class A:
+                pass
+        A.__call__ = A()
+        a = A()
+
+        try:
+            a() # This should not segfault
+        except RecursionError:
+            pass
+        else:
+            self.fail("Failed to raise RecursionError")
+
+    def testForExceptionsRaisedInInstanceGetattr2(self):
+        # Tests for exceptions raised in instance_getattr2().
+
+        def booh(self):
+            raise AttributeError("booh")
+
+        with torch._dynamo.error_on_graph_break(False):
+            class A:
+                a = property(booh)
+        try:
+            A().a # Raised AttributeError: A instance has no attribute 'a'
+        except AttributeError as x:
+            if str(x) != "booh":
+                self.fail("attribute error for A().a got masked: %s" % x)
+
+        with torch._dynamo.error_on_graph_break(False):
+            class E:
+                __eq__ = property(booh)
+        E() == E() # In debug mode, caused a C-level assert() to fail
+
+        with torch._dynamo.error_on_graph_break(False):
+            class I:
+                __init__ = property(booh)
+        try:
+            # In debug mode, printed XXX undetected error and
+            #  raises AttributeError
+            I()
+        except AttributeError:
+            pass
+        else:
+            self.fail("attribute error for I.__init__ got masked")
+
+    def assertNotOrderable(self, a, b):
+        with self.assertRaises(TypeError):
+            a < b
+        with self.assertRaises(TypeError):
+            a > b
+        with self.assertRaises(TypeError):
+            a <= b
+        with self.assertRaises(TypeError):
+            a >= b
+
+    def testHashComparisonOfMethods(self):
+        # Test comparison and hash of methods
+        with torch._dynamo.error_on_graph_break(False):
+            class A:
+                def __init__(self, x):
+                    self.x = x
+                def f(self):
+                    pass
+                def g(self):
+                    pass
+                def __eq__(self, other):
+                    return True
+                def __hash__(self):
+                    raise TypeError
+        with torch._dynamo.error_on_graph_break(False):
+            class B(A):
+                pass
+
+        a1 = A(1)
+        a2 = A(1)
+        self.assertTrue(a1.f == a1.f)
+        self.assertFalse(a1.f != a1.f)
+        self.assertFalse(a1.f == a2.f)
+        self.assertTrue(a1.f != a2.f)
+        self.assertFalse(a1.f == a1.g)
+        self.assertTrue(a1.f != a1.g)
+        self.assertNotOrderable(a1.f, a1.f)
+        self.assertEqual(hash(a1.f), hash(a1.f))
+
+        self.assertFalse(A.f == a1.f)
+        self.assertTrue(A.f != a1.f)
+        self.assertFalse(A.f == A.g)
+        self.assertTrue(A.f != A.g)
+        self.assertTrue(B.f == A.f)
+        self.assertFalse(B.f != A.f)
+        self.assertNotOrderable(A.f, A.f)
+        self.assertEqual(hash(B.f), hash(A.f))
+
+        # the following triggers a SystemError in 2.4
+        a = A(hash(A.f)^(-1))
+        hash(a.f)
+
+    def testSetattrWrapperNameIntern(self):
+        # Issue #25794: __setattr__ should intern the attribute name
+        with torch._dynamo.error_on_graph_break(False):
+            class A:
+                pass
+
+        def add(self, other):
+            return 'summa'
+
+        name = str(b'__add__', 'ascii')  # shouldn't be optimized
+        self.assertIsNot(name, '__add__')  # not interned
+        type.__setattr__(A, name, add)
+        self.assertEqual(A() + 1, 'summa')
+
+        name2 = str(b'__add__', 'ascii')
+        self.assertIsNot(name2, '__add__')
+        self.assertIsNot(name2, name)
+        type.__delattr__(A, name2)
+        with self.assertRaises(TypeError):
+            A() + 1
+
+    def testSetattrNonStringName(self):
+        with torch._dynamo.error_on_graph_break(False):
+            class A:
+                pass
+
+        with self.assertRaises(TypeError):
+            type.__setattr__(A, b'x', None)
+
+    def testTypeAttributeAccessErrorMessages(self):
+        with torch._dynamo.error_on_graph_break(False):
+            class A:
+                pass
+
+        error_msg = "type object 'A' has no attribute 'x'"
+        with self.assertRaisesRegex(AttributeError, error_msg):
+            A.x
+        with self.assertRaisesRegex(AttributeError, error_msg):
+            del A.x
+
+    def testObjectAttributeAccessErrorMessages(self):
+        with torch._dynamo.error_on_graph_break(False):
+            class A:
+                pass
+            class B:
+                y = 0
+                __slots__ = ('z',)
+            class C:
+                __slots__ = ("y",)
+
+                def __setattr__(self, name, value) -> None:
+                    if name == "z":
+                        super().__setattr__("y", 1)
+                    else:
+                        super().__setattr__(name, value)
+
+        error_msg = "'A' object has no attribute 'x'"
+        with self.assertRaisesRegex(AttributeError, error_msg):
+            A().x
+        with self.assertRaisesRegex(AttributeError, error_msg):
+            del A().x
+
+        error_msg = "'B' object has no attribute 'x'"
+        with self.assertRaisesRegex(AttributeError, error_msg):
+            B().x
+        with self.assertRaisesRegex(AttributeError, error_msg):
+            del B().x
+        with self.assertRaisesRegex(
+            AttributeError,
+            "'B' object has no attribute 'x' and no __dict__ for setting new attributes"
+        ):
+            B().x = 0
+        with self.assertRaisesRegex(
+            AttributeError,
+            "'C' object has no attribute 'x'"
+        ):
+            C().x = 0
+
+        error_msg = "'B' object attribute 'y' is read-only"
+        with self.assertRaisesRegex(AttributeError, error_msg):
+            del B().y
+        with self.assertRaisesRegex(AttributeError, error_msg):
+            B().y = 0
+
+        error_msg = 'z'
+        with self.assertRaisesRegex(AttributeError, error_msg):
+            B().z
+        with self.assertRaisesRegex(AttributeError, error_msg):
+            del B().z
+
+    def testConstructorErrorMessages(self):
+        # bpo-31506: Improves the error message logic for object_new & object_init
+
+        # Class without any method overrides
+        with torch._dynamo.error_on_graph_break(False):
+            class C:
+                pass
+
+        error_msg = r'C.__init__\(\) takes exactly one argument \(the instance to initialize\)'
+
+        with self.assertRaisesRegex(TypeError, r'C\(\) takes no arguments'):
+            C(42)
+
+        with self.assertRaisesRegex(TypeError, r'C\(\) takes no arguments'):
+            C.__new__(C, 42)
+
+        with self.assertRaisesRegex(TypeError, error_msg):
+            C().__init__(42)
+
+        with self.assertRaisesRegex(TypeError, r'C\(\) takes no arguments'):
+            object.__new__(C, 42)
+
+        with self.assertRaisesRegex(TypeError, error_msg):
+            object.__init__(C(), 42)
+
+        # Class with both `__init__` & `__new__` method overridden
+        with torch._dynamo.error_on_graph_break(False):
+            class D:
+                def __new__(cls, *args, **kwargs):
+                    super().__new__(cls, *args, **kwargs)
+                def __init__(self, *args, **kwargs):
+                    super().__init__(*args, **kwargs)
+
+        error_msg =  r'object.__new__\(\) takes exactly one argument \(the type to instantiate\)'
+
+        with self.assertRaisesRegex(TypeError, error_msg):
+            D(42)
+
+        with self.assertRaisesRegex(TypeError, error_msg):
+            D.__new__(D, 42)
+
+        with self.assertRaisesRegex(TypeError, error_msg):
+            object.__new__(D, 42)
+
+        # Class that only overrides __init__
+        with torch._dynamo.error_on_graph_break(False):
+            class E:
+                def __init__(self, *args, **kwargs):
+                    super().__init__(*args, **kwargs)
+
+        error_msg = r'object.__init__\(\) takes exactly one argument \(the instance to initialize\)'
+
+        with self.assertRaisesRegex(TypeError, error_msg):
+            E().__init__(42)
+
+        with self.assertRaisesRegex(TypeError, error_msg):
+            object.__init__(E(), 42)
+
+    def testClassWithExtCall(self):
+        with torch._dynamo.error_on_graph_break(False):
+            class Meta(int):
+                def __init__(*args, **kwargs):
+                    pass
+
+                def __new__(cls, name, bases, attrs, **kwargs):
+                    return bases, kwargs
+
+        d = {'metaclass': Meta}
+
+        with torch._dynamo.error_on_graph_break(False):
+            class A(**d): pass
+        self.assertEqual(A, ((), {}))
+        with torch._dynamo.error_on_graph_break(False):
+            class A(0, 1, 2, 3, 4, 5, 6, 7, **d): pass
+        self.assertEqual(A, (tuple(range(8)), {}))
+        with torch._dynamo.error_on_graph_break(False):
+            class A(0, *range(1, 8), **d, foo='bar'): pass
+        self.assertEqual(A, (tuple(range(8)), {'foo': 'bar'}))
+
+    def testClassCallRecursionLimit(self):
+        class C:
+            def __init__(self):
+                self.c = C()
+
+        with self.assertRaises(RecursionError):
+            C()
+
+        def add_one_level():
+            #Each call to C() consumes 2 levels, so offset by 1.
+            C()
+
+        with self.assertRaises(RecursionError):
+            add_one_level()
+
+    def testMetaclassCallOptimization(self):
+        calls = 0
+
+        with torch._dynamo.error_on_graph_break(False):
+            class TypeMetaclass(type):
+                def __call__(cls, *args, **kwargs):
+                    nonlocal calls
+                    calls += 1
+                    return type.__call__(cls, *args, **kwargs)
+
+            class Type(metaclass=TypeMetaclass):
+                def __init__(self, obj):
+                    self._obj = obj
+
+        for i in range(100):
+            Type(i)
+        self.assertEqual(calls, 100)
+
+    def test_specialization_class_call_doesnt_crash(self):
+        # gh-123185
+
+        with torch._dynamo.error_on_graph_break(False):
+            class Foo:
+                def __init__(self, arg):
+                    pass
+
+        for _ in range(8):
+            try:
+                Foo()
+            except:
+                pass
+
+
+try:
+    from _testinternalcapi import has_inline_values
+except ImportError:
+    has_inline_values = None
+
+Py_TPFLAGS_MANAGED_DICT = (1 << 2)
+
+class Plain:
+    pass
+
+
+class WithAttrs:
+
+    def __init__(self):
+        self.a = 1
+        self.b = 2
+        self.c = 3
+        self.d = 4
+
+
+@skipIfTorchDynamo("CPython only")
+class TestInlineValues(CPythonTestCase):
+
+    def test_flags(self):
+        self.assertEqual(Plain.__flags__ & Py_TPFLAGS_MANAGED_DICT, Py_TPFLAGS_MANAGED_DICT)
+        self.assertEqual(WithAttrs.__flags__ & Py_TPFLAGS_MANAGED_DICT, Py_TPFLAGS_MANAGED_DICT)
+
+    def test_has_inline_values(self):
+        c = Plain()
+        self.assertTrue(has_inline_values(c))
+        del c.__dict__
+        self.assertFalse(has_inline_values(c))
+
+    def test_instances(self):
+        self.assertTrue(has_inline_values(Plain()))
+        self.assertTrue(has_inline_values(WithAttrs()))
+
+    def test_inspect_dict(self):
+        for cls in (Plain, WithAttrs):
+            c = cls()
+            c.__dict__
+            self.assertTrue(has_inline_values(c))
+
+    def test_update_dict(self):
+        d = { "e": 5, "f": 6 }
+        for cls in (Plain, WithAttrs):
+            c = cls()
+            c.__dict__.update(d)
+            self.assertTrue(has_inline_values(c))
+
+    @staticmethod
+    def set_100(obj):
+        for i in range(100):
+            setattr(obj, f"a{i}", i)
+
+    def check_100(self, obj):
+        for i in range(100):
+            self.assertEqual(getattr(obj, f"a{i}"), i)
+
+    def test_many_attributes(self):
+        with torch._dynamo.error_on_graph_break(False):
+            class C: pass
+        c = C()
+        self.assertTrue(has_inline_values(c))
+        self.set_100(c)
+        self.assertFalse(has_inline_values(c))
+        self.check_100(c)
+        c = C()
+        self.assertTrue(has_inline_values(c))
+
+    def test_many_attributes_with_dict(self):
+        class C: pass
+        c = C()
+        d = c.__dict__
+        self.assertTrue(has_inline_values(c))
+        self.set_100(c)
+        self.assertFalse(has_inline_values(c))
+        self.check_100(c)
+
+    def test_bug_117750(self):
+        "Aborted on 3.13a6"
+        class C:
+            def __init__(self):
+                self.__dict__.clear()
+
+        obj = C()
+        self.assertEqual(obj.__dict__, {})
+        obj.foo = None # Aborted here
+        self.assertEqual(obj.__dict__, {"foo":None})
+
+    def test_store_attr_deleted_dict(self):
+        class Foo:
+            pass
+
+        f = Foo()
+        del f.__dict__
+        f.a = 3
+        self.assertEqual(f.a, 3)
+
+    def test_rematerialize_object_dict(self):
+        # gh-121860: rematerializing an object's managed dictionary after it
+        # had been deleted caused a crash.
+        class Foo: pass
+        f = Foo()
+        f.__dict__["attr"] = 1
+        del f.__dict__
+
+        # Using a str subclass is a way to trigger the re-materialization
+        class StrSubclass(str): pass
+        self.assertFalse(hasattr(f, StrSubclass("attr")))
+
+        # Changing the __class__ also triggers the re-materialization
+        class Bar: pass
+        f.__class__ = Bar
+        self.assertIsInstance(f, Bar)
+        self.assertEqual(f.__dict__, {})
+
+    def test_store_attr_type_cache(self):
+        """Verifies that the type cache doesn't provide a value which  is
+        inconsistent from the dict."""
+        class X:
+            def __del__(inner_self):
+                v = C.a
+                self.assertEqual(v, C.__dict__['a'])
+
+        class C:
+            a = X()
+
+        # prime the cache
+        C.a
+        C.a
+
+        # destructor shouldn't be able to see inconsistent state
+        C.a = X()
+        C.a = X()
+
+    @cpython_only
+    def test_detach_materialized_dict_no_memory(self):
+        # Skip test if _testcapi is not available:
+        import_helper.import_module('_testcapi')
+
+        code = """if 1:
+            import test.support
+            import _testcapi
+
+            class A:
+                def __init__(self):
+                    self.a = 1
+                    self.b = 2
+            a = A()
+            d = a.__dict__
+            with test.support.catch_unraisable_exception() as ex:
+                _testcapi.set_nomemory(0, 1)
+                del a
+                assert ex.unraisable.exc_type is MemoryError
+            try:
+                d["a"]
+            except KeyError:
+                pass
+            else:
+                assert False, "KeyError not raised"
+        """
+        rc, out, err = script_helper.assert_python_ok("-c", code)
+        self.assertEqual(rc, 0)
+        self.assertFalse(out, msg=out.decode('utf-8'))
+        self.assertFalse(err, msg=err.decode('utf-8'))
+
+if __name__ == "__main__":
+    run_tests()
diff --git a/test/dynamo/cpython/3_13/test_collections.py b/test/dynamo/cpython/3_13/test_collections.py
index ec10d45bd7754..38c6dfaec9ab8 100644
--- a/test/dynamo/cpython/3_13/test_collections.py
+++ b/test/dynamo/cpython/3_13/test_collections.py
@@ -12,7 +12,7 @@
 import torch._dynamo.test_case
 import unittest
 from torch._dynamo.test_case import CPythonTestCase
-from torch.testing._internal.common_utils import run_tests
+from torch.testing._internal.common_utils import run_tests, slowTest
 
 __TestCase = CPythonTestCase
 
@@ -2420,6 +2420,7 @@ def test_helper_function(self):
         self.assertTrue(c.called)
         self.assertEqual(dict(c), {'a': 5, 'b': 2, 'c': 1, 'd': 1, 'r':2 })
 
+    @slowTest
     def test_multiset_operations_equivalent_to_set_operations(self):
         # When the multiplicities are all zero or one, multiset operations
         # are guaranteed to be equivalent to the corresponding operations
diff --git a/test/dynamo/cpython/3_13/test_defaultdict.py b/test/dynamo/cpython/3_13/test_defaultdict.py
index 390f55d89dec4..65ac9fac1904a 100644
--- a/test/dynamo/cpython/3_13/test_defaultdict.py
+++ b/test/dynamo/cpython/3_13/test_defaultdict.py
@@ -210,20 +210,20 @@ def test_union(self):
         i = defaultdict(int, {1: 1, 2: 2})
         s = defaultdict(str, {0: "zero", 1: "one"})
 
-        i_s = i | s
-        self.assertIs(i_s.default_factory, int)
-        self.assertDictEqual(i_s, {1: "one", 2: 2, 0: "zero"})
-        self.assertEqual(list(i_s), [1, 2, 0])
-
-        s_i = s | i
-        self.assertIs(s_i.default_factory, str)
-        self.assertDictEqual(s_i, {0: "zero", 1: 1, 2: 2})
-        self.assertEqual(list(s_i), [0, 1, 2])
-
-        i_ds = i | dict(s)
-        self.assertIs(i_ds.default_factory, int)
-        self.assertDictEqual(i_ds, {1: "one", 2: 2, 0: "zero"})
-        self.assertEqual(list(i_ds), [1, 2, 0])
+        # i_s = i | s
+        # self.assertIs(i_s.default_factory, int)
+        # self.assertDictEqual(i_s, {1: "one", 2: 2, 0: "zero"})
+        # self.assertEqual(list(i_s), [1, 2, 0])
+
+        # s_i = s | i
+        # self.assertIs(s_i.default_factory, str)
+        # self.assertDictEqual(s_i, {0: "zero", 1: 1, 2: 2})
+        # self.assertEqual(list(s_i), [0, 1, 2])
+
+        # i_ds = i | dict(s)
+        # self.assertIs(i_ds.default_factory, int)
+        # self.assertDictEqual(i_ds, {1: "one", 2: 2, 0: "zero"})
+        # self.assertEqual(list(i_ds), [1, 2, 0])
 
         ds_i = dict(s) | i
         self.assertIs(ds_i.default_factory, int)
diff --git a/test/dynamo/cpython/3_13/test_descr.diff b/test/dynamo/cpython/3_13/test_descr.diff
new file mode 100644
index 0000000000000..7ed0668df0750
--- /dev/null
+++ b/test/dynamo/cpython/3_13/test_descr.diff
@@ -0,0 +1,4625 @@
+diff --git a/test/dynamo/cpython/3_13/test_descr.py b/test/dynamo/cpython/3_13/test_descr.py
+index be49ce8c066..938b09f0a0f 100644
+--- a/test/dynamo/cpython/3_13/test_descr.py
++++ b/test/dynamo/cpython/3_13/test_descr.py
+@@ -1,3 +1,21 @@
++# ======= BEGIN Dynamo patch =======
++# Owner(s): ["module: dynamo"]
++
++# ruff: noqa
++# flake8: noqa
++
++# Test copied from
++# https://raw.githubusercontent.com/python/cpython/refs/tags/v3.13.5/Lib/test/test_descr.py
++
++import sys
++import torch
++import torch._dynamo.test_case
++import unittest
++from torch._dynamo.test_case import CPythonTestCase
++from torch.testing._internal.common_utils import run_tests
++
++# ======= END DYNAMO PATCH =======
++
+ import builtins
+ import copyreg
+ import gc
+@@ -28,10 +46,10 @@ except ImportError:
+     xxsubtype = None
+ 
+ 
+-class OperatorsTest(unittest.TestCase):
++class OperatorsTest(CPythonTestCase):
+ 
+     def __init__(self, *args, **kwargs):
+-        unittest.TestCase.__init__(self, *args, **kwargs)
++        CPythonTestCase.__init__(self, *args, **kwargs)
+         self.binops = {
+             'add': '+',
+             'sub': '-',
+@@ -251,9 +269,10 @@ class OperatorsTest(unittest.TestCase):
+         self.assertEqual((1).__bool__(), 1)
+         self.assertEqual((0).__bool__(), 0)
+         # This returns 'NotImplemented' in Python 2.2
+-        class C(int):
+-            def __add__(self, other):
+-                return NotImplemented
++        with torch._dynamo.error_on_graph_break(False):
++            class C(int):
++                def __add__(self, other):
++                    return NotImplemented
+         self.assertEqual(C(5), 5)
+         try:
+             C() + ""
+@@ -272,20 +291,21 @@ class OperatorsTest(unittest.TestCase):
+                                                   'int', 'float',
+                                                   'floordiv', 'divmod', 'mod'])
+ 
+-        class Number(complex):
+-            __slots__ = ['prec']
+-            def __new__(cls, *args, **kwds):
+-                result = complex.__new__(cls, *args)
+-                result.prec = kwds.get('prec', 12)
+-                return result
+-            def __repr__(self):
+-                prec = self.prec
+-                if self.imag == 0.0:
+-                    return "%.*g" % (prec, self.real)
+-                if self.real == 0.0:
+-                    return "%.*gj" % (prec, self.imag)
+-                return "(%.*g+%.*gj)" % (prec, self.real, prec, self.imag)
+-            __str__ = __repr__
++        with torch._dynamo.error_on_graph_break(False):
++            class Number(complex):
++                __slots__ = ['prec']
++                def __new__(cls, *args, **kwds):
++                    result = complex.__new__(cls, *args)
++                    result.prec = kwds.get('prec', 12)
++                    return result
++                def __repr__(self):
++                    prec = self.prec
++                    if self.imag == 0.0:
++                        return "%.*g" % (prec, self.real)
++                    if self.real == 0.0:
++                        return "%.*gj" % (prec, self.imag)
++                    return "(%.*g+%.*gj)" % (prec, self.real, prec, self.imag)
++                __str__ = __repr__
+ 
+         a = Number(3.14, prec=6)
+         self.assertEqual(repr(a), "3.14")
+@@ -337,9 +357,10 @@ class OperatorsTest(unittest.TestCase):
+                          "__setitem__")
+         self.setsliceop_test(spamlist([1,2,3,4]), 1, 3, spamlist([5,6]),
+                              spamlist([1,5,6,4]), "a[b:c]=d", "__setitem__")
+-        # Test subclassing
+-        class C(spam.spamlist):
+-            def foo(self): return 1
++        with torch._dynamo.error_on_graph_break(False):
++            # Test subclassing
++            class C(spam.spamlist):
++                def foo(self): return 1
+         a = C()
+         self.assertEqual(a, [])
+         self.assertEqual(a.foo(), 1)
+@@ -388,9 +409,10 @@ class OperatorsTest(unittest.TestCase):
+         self.unop_test(spamd, repr(straightd), "repr(a)", "__repr__")
+         self.set2op_test(spamdict({1:2,3:4}), 2, 3, spamdict({1:2,2:3,3:4}),
+                    "a[b]=c", "__setitem__")
+-        # Test subclassing
+-        class C(spam.spamdict):
+-            def foo(self): return 1
++        with torch._dynamo.error_on_graph_break(False):
++            # Test subclassing
++            class C(spam.spamdict):
++                def foo(self): return 1
+         a = C()
+         self.assertEqual(list(a.items()), [])
+         self.assertEqual(a.foo(), 1)
+@@ -404,7 +426,7 @@ class OperatorsTest(unittest.TestCase):
+         self.assertEqual(range(sys.maxsize).__len__(), sys.maxsize)
+ 
+ 
+-class ClassPropertiesAndMethods(unittest.TestCase, ExtraAssertions):
++class ClassPropertiesAndMethods(CPythonTestCase, ExtraAssertions):
+ 
+     def test_python_dicts(self):
+         # Testing Python subclass of dict...
+@@ -414,24 +436,25 @@ class ClassPropertiesAndMethods(unittest.TestCase, ExtraAssertions):
+         self.assertEqual(d, {})
+         self.assertIs(d.__class__, dict)
+         self.assertIsInstance(d, dict)
+-        class C(dict):
+-            state = -1
+-            def __init__(self_local, *a, **kw):
+-                if a:
+-                    self.assertEqual(len(a), 1)
+-                    self_local.state = a[0]
+-                if kw:
+-                    for k, v in list(kw.items()):
+-                        self_local[v] = k
+-            def __getitem__(self, key):
+-                return self.get(key, 0)
+-            def __setitem__(self_local, key, value):
+-                self.assertIsInstance(key, int)
+-                dict.__setitem__(self_local, key, value)
+-            def setstate(self, state):
+-                self.state = state
+-            def getstate(self):
+-                return self.state
++        with torch._dynamo.error_on_graph_break(False):
++            class C(dict):
++                state = -1
++                def __init__(self_local, *a, **kw):
++                    if a:
++                        self.assertEqual(len(a), 1)
++                        self_local.state = a[0]
++                    if kw:
++                        for k, v in list(kw.items()):
++                            self_local[v] = k
++                def __getitem__(self, key):
++                    return self.get(key, 0)
++                def __setitem__(self_local, key, value):
++                    self.assertIsInstance(key, int)
++                    dict.__setitem__(self_local, key, value)
++                def setstate(self, state):
++                    self.state = state
++                def getstate(self):
++                    return self.state
+         self.assertTrue(issubclass(C, dict))
+         a1 = C(12)
+         self.assertEqual(a1.state, 12)
+@@ -459,12 +482,13 @@ class ClassPropertiesAndMethods(unittest.TestCase, ExtraAssertions):
+                 self.assertEqual(a[i][j], i*j)
+ 
+     def test_python_lists(self):
+-        # Testing Python subclass of list...
+-        class C(list):
+-            def __getitem__(self, i):
+-                if isinstance(i, slice):
+-                    return i.start, i.stop
+-                return list.__getitem__(self, i) + 100
++        with torch._dynamo.error_on_graph_break(False):
++            # Testing Python subclass of list...
++            class C(list):
++                def __getitem__(self, i):
++                    if isinstance(i, slice):
++                        return i.start, i.stop
++                    return list.__getitem__(self, i) + 100
+         a = C()
+         a.extend([0,1,2])
+         self.assertEqual(a[0], 100)
+@@ -473,56 +497,60 @@ class ClassPropertiesAndMethods(unittest.TestCase, ExtraAssertions):
+         self.assertEqual(a[100:200], (100,200))
+ 
+     def test_metaclass(self):
+-        # Testing metaclasses...
+-        class C(metaclass=type):
+-            def __init__(self):
+-                self.__state = 0
+-            def getstate(self):
+-                return self.__state
+-            def setstate(self, state):
+-                self.__state = state
++        with torch._dynamo.error_on_graph_break(False):
++            # Testing metaclasses...
++            class C(metaclass=type):
++                def __init__(self):
++                    self.__state = 0
++                def getstate(self):
++                    return self.__state
++                def setstate(self, state):
++                    self.__state = state
+         a = C()
+         self.assertEqual(a.getstate(), 0)
+         a.setstate(10)
+         self.assertEqual(a.getstate(), 10)
+-        class _metaclass(type):
+-            def myself(cls): return cls
+-        class D(metaclass=_metaclass):
+-            pass
++        with torch._dynamo.error_on_graph_break(False):
++            class _metaclass(type):
++                def myself(cls): return cls
++            class D(metaclass=_metaclass):
++                pass
+         self.assertEqual(D.myself(), D)
+         d = D()
+         self.assertEqual(d.__class__, D)
+-        class M1(type):
+-            def __new__(cls, name, bases, dict):
+-                dict['__spam__'] = 1
+-                return type.__new__(cls, name, bases, dict)
+-        class C(metaclass=M1):
+-            pass
++        with torch._dynamo.error_on_graph_break(False):
++            class M1(type):
++                def __new__(cls, name, bases, dict):
++                    dict['__spam__'] = 1
++                    return type.__new__(cls, name, bases, dict)
++            class C(metaclass=M1):
++                pass
+         self.assertEqual(C.__spam__, 1)
+         c = C()
+         self.assertEqual(c.__spam__, 1)
+ 
+-        class _instance(object):
+-            pass
+-        class M2(object):
+-            @staticmethod
+-            def __new__(cls, name, bases, dict):
+-                self = object.__new__(cls)
+-                self.name = name
+-                self.bases = bases
+-                self.dict = dict
+-                return self
+-            def __call__(self):
+-                it = _instance()
+-                # Early binding of methods
+-                for key in self.dict:
+-                    if key.startswith("__"):
+-                        continue
+-                    setattr(it, key, self.dict[key].__get__(it, self))
+-                return it
+-        class C(metaclass=M2):
+-            def spam(self):
+-                return 42
++        with torch._dynamo.error_on_graph_break(False):
++            class _instance(object):
++                pass
++            class M2(object):
++                @staticmethod
++                def __new__(cls, name, bases, dict):
++                    self = object.__new__(cls)
++                    self.name = name
++                    self.bases = bases
++                    self.dict = dict
++                    return self
++                def __call__(self):
++                    it = _instance()
++                    # Early binding of methods
++                    for key in self.dict:
++                        if key.startswith("__"):
++                            continue
++                        setattr(it, key, self.dict[key].__get__(it, self))
++                    return it
++            class C(metaclass=M2):
++                def spam(self):
++                    return 42
+         self.assertEqual(C.name, 'C')
+         self.assertEqual(C.bases, ())
+         self.assertIn('spam', C.dict)
+@@ -531,179 +559,193 @@ class ClassPropertiesAndMethods(unittest.TestCase, ExtraAssertions):
+ 
+         # More metaclass examples
+ 
+-        class autosuper(type):
+-            # Automatically add __super to the class
+-            # This trick only works for dynamic classes
+-            def __new__(metaclass, name, bases, dict):
+-                cls = super(autosuper, metaclass).__new__(metaclass,
+-                                                          name, bases, dict)
+-                # Name mangling for __super removes leading underscores
+-                while name[:1] == "_":
+-                    name = name[1:]
+-                if name:
+-                    name = "_%s__super" % name
+-                else:
+-                    name = "__super"
+-                setattr(cls, name, super(cls))
+-                return cls
+-        class A(metaclass=autosuper):
+-            def meth(self):
+-                return "A"
+-        class B(A):
+-            def meth(self):
+-                return "B" + self.__super.meth()
+-        class C(A):
+-            def meth(self):
+-                return "C" + self.__super.meth()
+-        class D(C, B):
+-            def meth(self):
+-                return "D" + self.__super.meth()
++        with torch._dynamo.error_on_graph_break(False):
++            class autosuper(type):
++                # Automatically add __super to the class
++                # This trick only works for dynamic classes
++                def __new__(metaclass, name, bases, dict):
++                    cls = super(autosuper, metaclass).__new__(metaclass,
++                                                            name, bases, dict)
++                    # Name mangling for __super removes leading underscores
++                    while name[:1] == "_":
++                        name = name[1:]
++                    if name:
++                        name = "_%s__super" % name
++                    else:
++                        name = "__super"
++                    setattr(cls, name, super(cls))
++                    return cls
++            class A(metaclass=autosuper):
++                def meth(self):
++                    return "A"
++            class B(A):
++                def meth(self):
++                    return "B" + self.__super.meth()
++            class C(A):
++                def meth(self):
++                    return "C" + self.__super.meth()
++            class D(C, B):
++                def meth(self):
++                    return "D" + self.__super.meth()
+         self.assertEqual(D().meth(), "DCBA")
+-        class E(B, C):
+-            def meth(self):
+-                return "E" + self.__super.meth()
++        with torch._dynamo.error_on_graph_break(False):
++            class E(B, C):
++                def meth(self):
++                    return "E" + self.__super.meth()
+         self.assertEqual(E().meth(), "EBCA")
+ 
+-        class autoproperty(type):
+-            # Automatically create property attributes when methods
+-            # named _get_x and/or _set_x are found
+-            def __new__(metaclass, name, bases, dict):
+-                hits = {}
+-                for key, val in dict.items():
+-                    if key.startswith("_get_"):
+-                        key = key[5:]
+-                        get, set = hits.get(key, (None, None))
+-                        get = val
+-                        hits[key] = get, set
+-                    elif key.startswith("_set_"):
+-                        key = key[5:]
+-                        get, set = hits.get(key, (None, None))
+-                        set = val
+-                        hits[key] = get, set
+-                for key, (get, set) in hits.items():
+-                    dict[key] = property(get, set)
+-                return super(autoproperty, metaclass).__new__(metaclass,
+-                                                            name, bases, dict)
+-        class A(metaclass=autoproperty):
+-            def _get_x(self):
+-                return -self.__x
+-            def _set_x(self, x):
+-                self.__x = -x
++        with torch._dynamo.error_on_graph_break(False):
++            class autoproperty(type):
++                # Automatically create property attributes when methods
++                # named _get_x and/or _set_x are found
++                def __new__(metaclass, name, bases, dict):
++                    hits = {}
++                    for key, val in dict.items():
++                        if key.startswith("_get_"):
++                            key = key[5:]
++                            get, set = hits.get(key, (None, None))
++                            get = val
++                            hits[key] = get, set
++                        elif key.startswith("_set_"):
++                            key = key[5:]
++                            get, set = hits.get(key, (None, None))
++                            set = val
++                            hits[key] = get, set
++                    for key, (get, set) in hits.items():
++                        dict[key] = property(get, set)
++                    return super(autoproperty, metaclass).__new__(metaclass,
++                                                                name, bases, dict)
++        with torch._dynamo.error_on_graph_break(False):
++            class A(metaclass=autoproperty):
++                def _get_x(self):
++                    return -self.__x
++                def _set_x(self, x):
++                    self.__x = -x
+         a = A()
+         self.assertNotHasAttr(a, "x")
+         a.x = 12
+         self.assertEqual(a.x, 12)
+         self.assertEqual(a._A__x, -12)
+ 
+-        class multimetaclass(autoproperty, autosuper):
+-            # Merge of multiple cooperating metaclasses
+-            pass
+-        class A(metaclass=multimetaclass):
+-            def _get_x(self):
+-                return "A"
+-        class B(A):
+-            def _get_x(self):
+-                return "B" + self.__super._get_x()
+-        class C(A):
+-            def _get_x(self):
+-                return "C" + self.__super._get_x()
+-        class D(C, B):
+-            def _get_x(self):
+-                return "D" + self.__super._get_x()
++        with torch._dynamo.error_on_graph_break(False):
++            class multimetaclass(autoproperty, autosuper):
++                # Merge of multiple cooperating metaclasses
++                pass
++            class A(metaclass=multimetaclass):
++                def _get_x(self):
++                    return "A"
++            class B(A):
++                def _get_x(self):
++                    return "B" + self.__super._get_x()
++            class C(A):
++                def _get_x(self):
++                    return "C" + self.__super._get_x()
++            class D(C, B):
++                def _get_x(self):
++                    return "D" + self.__super._get_x()
+         self.assertEqual(D().x, "DCBA")
+ 
+-        # Make sure type(x) doesn't call x.__class__.__init__
+-        class T(type):
+-            counter = 0
+-            def __init__(self, *args):
+-                T.counter += 1
+-        class C(metaclass=T):
+-            pass
++        with torch._dynamo.error_on_graph_break(False):
++            # Make sure type(x) doesn't call x.__class__.__init__
++            class T(type):
++                counter = 0
++                def __init__(self, *args):
++                    T.counter += 1
++            class C(metaclass=T):
++                pass
+         self.assertEqual(T.counter, 1)
+         a = C()
+         self.assertEqual(type(a), C)
+         self.assertEqual(T.counter, 1)
+ 
+-        class C(object): pass
++        with torch._dynamo.error_on_graph_break(False):
++            class C(object): pass
+         c = C()
+         try: c()
+         except TypeError: pass
+         else: self.fail("calling object w/o call method should raise "
+                         "TypeError")
+ 
+-        # Testing code to find most derived baseclass
+-        class A(type):
+-            def __new__(*args, **kwargs):
+-                return type.__new__(*args, **kwargs)
++        with torch._dynamo.error_on_graph_break(False):
++            # Testing code to find most derived baseclass
++            class A(type):
++                def __new__(*args, **kwargs):
++                    return type.__new__(*args, **kwargs)
+ 
+-        class B(object):
+-            pass
++            class B(object):
++                pass
+ 
+-        class C(object, metaclass=A):
+-            pass
++            class C(object, metaclass=A):
++                pass
+ 
+-        # The most derived metaclass of D is A rather than type.
+-        class D(B, C):
+-            pass
++            # The most derived metaclass of D is A rather than type.
++            class D(B, C):
++                pass
+         self.assertIs(A, type(D))
+ 
+         # issue1294232: correct metaclass calculation
+         new_calls = []  # to check the order of __new__ calls
+-        class AMeta(type):
+-            @staticmethod
+-            def __new__(mcls, name, bases, ns):
+-                new_calls.append('AMeta')
+-                return super().__new__(mcls, name, bases, ns)
+-            @classmethod
+-            def __prepare__(mcls, name, bases):
+-                return {}
+-
+-        class BMeta(AMeta):
+-            @staticmethod
+-            def __new__(mcls, name, bases, ns):
+-                new_calls.append('BMeta')
+-                return super().__new__(mcls, name, bases, ns)
+-            @classmethod
+-            def __prepare__(mcls, name, bases):
+-                ns = super().__prepare__(name, bases)
+-                ns['BMeta_was_here'] = True
+-                return ns
+-
+-        class A(metaclass=AMeta):
+-            pass
++        with torch._dynamo.error_on_graph_break(False):
++            class AMeta(type):
++                @staticmethod
++                def __new__(mcls, name, bases, ns):
++                    new_calls.append('AMeta')
++                    return super().__new__(mcls, name, bases, ns)
++                @classmethod
++                def __prepare__(mcls, name, bases):
++                    return {}
++
++            class BMeta(AMeta):
++                @staticmethod
++                def __new__(mcls, name, bases, ns):
++                    new_calls.append('BMeta')
++                    return super().__new__(mcls, name, bases, ns)
++                @classmethod
++                def __prepare__(mcls, name, bases):
++                    ns = super().__prepare__(name, bases)
++                    ns['BMeta_was_here'] = True
++                    return ns
++
++            class A(metaclass=AMeta):
++                pass
+         self.assertEqual(['AMeta'], new_calls)
+         new_calls.clear()
+ 
+-        class B(metaclass=BMeta):
+-            pass
++        with torch._dynamo.error_on_graph_break(False):
++            class B(metaclass=BMeta):
++                pass
+         # BMeta.__new__ calls AMeta.__new__ with super:
+         self.assertEqual(['BMeta', 'AMeta'], new_calls)
+         new_calls.clear()
+ 
+-        class C(A, B):
+-            pass
++        with torch._dynamo.error_on_graph_break(False):
++            class C(A, B):
++                pass
+         # The most derived metaclass is BMeta:
+         self.assertEqual(['BMeta', 'AMeta'], new_calls)
+         new_calls.clear()
+         # BMeta.__prepare__ should've been called:
+         self.assertIn('BMeta_was_here', C.__dict__)
+ 
+-        # The order of the bases shouldn't matter:
+-        class C2(B, A):
+-            pass
++        with torch._dynamo.error_on_graph_break(False):
++            # The order of the bases shouldn't matter:
++            class C2(B, A):
++                pass
+         self.assertEqual(['BMeta', 'AMeta'], new_calls)
+         new_calls.clear()
+         self.assertIn('BMeta_was_here', C2.__dict__)
+ 
+-        # Check correct metaclass calculation when a metaclass is declared:
+-        class D(C, metaclass=type):
+-            pass
++        with torch._dynamo.error_on_graph_break(False):
++            # Check correct metaclass calculation when a metaclass is declared:
++            class D(C, metaclass=type):
++                pass
+         self.assertEqual(['BMeta', 'AMeta'], new_calls)
+         new_calls.clear()
+         self.assertIn('BMeta_was_here', D.__dict__)
+ 
+-        class E(C, metaclass=AMeta):
+-            pass
++        with torch._dynamo.error_on_graph_break(False):
++            class E(C, metaclass=AMeta):
++                pass
+         self.assertEqual(['BMeta', 'AMeta'], new_calls)
+         new_calls.clear()
+         self.assertIn('BMeta_was_here', E.__dict__)
+@@ -713,12 +755,13 @@ class ClassPropertiesAndMethods(unittest.TestCase, ExtraAssertions):
+         marker = object()
+         def func(*args, **kwargs):
+             return marker
+-        class X(metaclass=func):
+-            pass
+-        class Y(object, metaclass=func):
+-            pass
+-        class Z(D, metaclass=func):
+-            pass
++        with torch._dynamo.error_on_graph_break(False):
++            class X(metaclass=func):
++                pass
++            class Y(object, metaclass=func):
++                pass
++            class Z(D, metaclass=func):
++                pass
+         self.assertIs(marker, X)
+         self.assertIs(marker, Y)
+         self.assertIs(marker, Z)
+@@ -726,49 +769,53 @@ class ClassPropertiesAndMethods(unittest.TestCase, ExtraAssertions):
+         # The given metaclass is a class,
+         # but not a descendant of type.
+         prepare_calls = []  # to track __prepare__ calls
+-        class ANotMeta:
+-            def __new__(mcls, *args, **kwargs):
+-                new_calls.append('ANotMeta')
+-                return super().__new__(mcls)
+-            @classmethod
+-            def __prepare__(mcls, name, bases):
+-                prepare_calls.append('ANotMeta')
+-                return {}
+-        class BNotMeta(ANotMeta):
+-            def __new__(mcls, *args, **kwargs):
+-                new_calls.append('BNotMeta')
+-                return super().__new__(mcls)
+-            @classmethod
+-            def __prepare__(mcls, name, bases):
+-                prepare_calls.append('BNotMeta')
+-                return super().__prepare__(name, bases)
+-
+-        class A(metaclass=ANotMeta):
+-            pass
++        with torch._dynamo.error_on_graph_break(False):
++            class ANotMeta:
++                def __new__(mcls, *args, **kwargs):
++                    new_calls.append('ANotMeta')
++                    return super().__new__(mcls)
++                @classmethod
++                def __prepare__(mcls, name, bases):
++                    prepare_calls.append('ANotMeta')
++                    return {}
++            class BNotMeta(ANotMeta):
++                def __new__(mcls, *args, **kwargs):
++                    new_calls.append('BNotMeta')
++                    return super().__new__(mcls)
++                @classmethod
++                def __prepare__(mcls, name, bases):
++                    prepare_calls.append('BNotMeta')
++                    return super().__prepare__(name, bases)
++
++            class A(metaclass=ANotMeta):
++                pass
+         self.assertIs(ANotMeta, type(A))
+         self.assertEqual(['ANotMeta'], prepare_calls)
+         prepare_calls.clear()
+         self.assertEqual(['ANotMeta'], new_calls)
+         new_calls.clear()
+ 
+-        class B(metaclass=BNotMeta):
+-            pass
++        with torch._dynamo.error_on_graph_break(False):
++            class B(metaclass=BNotMeta):
++                pass
+         self.assertIs(BNotMeta, type(B))
+         self.assertEqual(['BNotMeta', 'ANotMeta'], prepare_calls)
+         prepare_calls.clear()
+         self.assertEqual(['BNotMeta', 'ANotMeta'], new_calls)
+         new_calls.clear()
+ 
+-        class C(A, B):
+-            pass
++        with torch._dynamo.error_on_graph_break(False):
++            class C(A, B):
++                pass
+         self.assertIs(BNotMeta, type(C))
+         self.assertEqual(['BNotMeta', 'ANotMeta'], new_calls)
+         new_calls.clear()
+         self.assertEqual(['BNotMeta', 'ANotMeta'], prepare_calls)
+         prepare_calls.clear()
+ 
+-        class C2(B, A):
+-            pass
++        with torch._dynamo.error_on_graph_break(False):
++            class C2(B, A):
++                pass
+         self.assertIs(BNotMeta, type(C2))
+         self.assertEqual(['BNotMeta', 'ANotMeta'], new_calls)
+         new_calls.clear()
+@@ -777,28 +824,31 @@ class ClassPropertiesAndMethods(unittest.TestCase, ExtraAssertions):
+ 
+         # This is a TypeError, because of a metaclass conflict:
+         # BNotMeta is neither a subclass, nor a superclass of type
+-        with self.assertRaises(TypeError):
+-            class D(C, metaclass=type):
+-                pass
++        with torch._dynamo.error_on_graph_break(False):
++            with self.assertRaises(TypeError):
++                class D(C, metaclass=type):
++                    pass
+ 
+-        class E(C, metaclass=ANotMeta):
+-            pass
++            class E(C, metaclass=ANotMeta):
++                pass
+         self.assertIs(BNotMeta, type(E))
+         self.assertEqual(['BNotMeta', 'ANotMeta'], new_calls)
+         new_calls.clear()
+         self.assertEqual(['BNotMeta', 'ANotMeta'], prepare_calls)
+         prepare_calls.clear()
+ 
+-        class F(object(), C):
+-            pass
++        with torch._dynamo.error_on_graph_break(False):
++            class F(object(), C):
++                pass
+         self.assertIs(BNotMeta, type(F))
+         self.assertEqual(['BNotMeta', 'ANotMeta'], new_calls)
+         new_calls.clear()
+         self.assertEqual(['BNotMeta', 'ANotMeta'], prepare_calls)
+         prepare_calls.clear()
+ 
+-        class F2(C, object()):
+-            pass
++        with torch._dynamo.error_on_graph_break(False):
++            class F2(C, object()):
++                pass
+         self.assertIs(BNotMeta, type(F2))
+         self.assertEqual(['BNotMeta', 'ANotMeta'], new_calls)
+         new_calls.clear()
+@@ -807,29 +857,31 @@ class ClassPropertiesAndMethods(unittest.TestCase, ExtraAssertions):
+ 
+         # TypeError: BNotMeta is neither a
+         # subclass, nor a superclass of int
+-        with self.assertRaises(TypeError):
+-            class X(C, int()):
+-                pass
+-        with self.assertRaises(TypeError):
+-            class X(int(), C):
+-                pass
++        with torch._dynamo.error_on_graph_break(False):
++            with self.assertRaises(TypeError):
++                class X(C, int()):
++                    pass
++            with self.assertRaises(TypeError):
++                class X(int(), C):
++                    pass
+ 
+     def test_module_subclasses(self):
+         # Testing Python subclass of module...
+         log = []
+         MT = type(sys)
+-        class MM(MT):
+-            def __init__(self, name):
+-                MT.__init__(self, name)
+-            def __getattribute__(self, name):
+-                log.append(("getattr", name))
+-                return MT.__getattribute__(self, name)
+-            def __setattr__(self, name, value):
+-                log.append(("setattr", name, value))
+-                MT.__setattr__(self, name, value)
+-            def __delattr__(self, name):
+-                log.append(("delattr", name))
+-                MT.__delattr__(self, name)
++        with torch._dynamo.error_on_graph_break(False):
++            class MM(MT):
++                def __init__(self, name):
++                    MT.__init__(self, name)
++                def __getattribute__(self, name):
++                    log.append(("getattr", name))
++                    return MT.__getattribute__(self, name)
++                def __setattr__(self, name, value):
++                    log.append(("setattr", name, value))
++                    MT.__setattr__(self, name, value)
++                def __delattr__(self, name):
++                    log.append(("delattr", name))
++                    MT.__delattr__(self, name)
+         a = MM("a")
+         a.foo = 12
+         x = a.foo
+@@ -840,39 +892,43 @@ class ClassPropertiesAndMethods(unittest.TestCase, ExtraAssertions):
+ 
+         # https://bugs.python.org/issue1174712
+         try:
+-            class Module(types.ModuleType, str):
+-                pass
++            with torch._dynamo.error_on_graph_break(False):
++                class Module(types.ModuleType, str):
++                    pass
+         except TypeError:
+             pass
+         else:
+             self.fail("inheriting from ModuleType and str at the same time "
+-                      "should fail")
++                    "should fail")
+ 
+         # Issue 34805: Verify that definition order is retained
+         def random_name():
+             return ''.join(random.choices(string.ascii_letters, k=10))
+-        class A:
+-            pass
++        with torch._dynamo.error_on_graph_break(False):
++            class A:
++                pass
+         subclasses = [type(random_name(), (A,), {}) for i in range(100)]
+         self.assertEqual(A.__subclasses__(), subclasses)
+ 
+     def test_multiple_inheritance(self):
+         # Testing multiple inheritance...
+-        class C(object):
+-            def __init__(self):
+-                self.__state = 0
+-            def getstate(self):
+-                return self.__state
+-            def setstate(self, state):
+-                self.__state = state
++        with torch._dynamo.error_on_graph_break(False):
++            class C(object):
++                def __init__(self):
++                    self.__state = 0
++                def getstate(self):
++                    return self.__state
++                def setstate(self, state):
++                    self.__state = state
+         a = C()
+         self.assertEqual(a.getstate(), 0)
+         a.setstate(10)
+         self.assertEqual(a.getstate(), 10)
+-        class D(dict, C):
+-            def __init__(self):
+-                dict.__init__(self)
+-                C.__init__(self)
++        with torch._dynamo.error_on_graph_break(False):
++            class D(dict, C):
++                def __init__(self):
++                    dict.__init__(self)
++                    C.__init__(self)
+         d = D()
+         self.assertEqual(list(d.keys()), [])
+         d["hello"] = "world"
+@@ -884,14 +940,15 @@ class ClassPropertiesAndMethods(unittest.TestCase, ExtraAssertions):
+         self.assertEqual(D.__mro__, (D, dict, C, object))
+ 
+         # SF bug #442833
+-        class Node(object):
+-            def __int__(self):
+-                return int(self.foo())
+-            def foo(self):
+-                return "23"
+-        class Frag(Node, list):
+-            def foo(self):
+-                return "42"
++        with torch._dynamo.error_on_graph_break(False):
++            class Node(object):
++                def __int__(self):
++                    return int(self.foo())
++                def foo(self):
++                    return "23"
++            class Frag(Node, list):
++                def foo(self):
++                    return "42"
+         self.assertEqual(Node().__int__(), 23)
+         self.assertEqual(int(Node()), 23)
+         self.assertEqual(Frag().__int__(), 42)
+@@ -899,35 +956,42 @@ class ClassPropertiesAndMethods(unittest.TestCase, ExtraAssertions):
+ 
+     def test_diamond_inheritance(self):
+         # Testing multiple inheritance special cases...
+-        class A(object):
+-            def spam(self): return "A"
++        with torch._dynamo.error_on_graph_break(False):
++            class A(object):
++                def spam(self): return "A"
+         self.assertEqual(A().spam(), "A")
+-        class B(A):
+-            def boo(self): return "B"
+-            def spam(self): return "B"
++        with torch._dynamo.error_on_graph_break(False):
++            class B(A):
++                def boo(self): return "B"
++                def spam(self): return "B"
+         self.assertEqual(B().spam(), "B")
+         self.assertEqual(B().boo(), "B")
+-        class C(A):
+-            def boo(self): return "C"
++        with torch._dynamo.error_on_graph_break(False):
++            class C(A):
++                def boo(self): return "C"
+         self.assertEqual(C().spam(), "A")
+         self.assertEqual(C().boo(), "C")
+-        class D(B, C): pass
++        with torch._dynamo.error_on_graph_break(False):
++            class D(B, C): pass
+         self.assertEqual(D().spam(), "B")
+         self.assertEqual(D().boo(), "B")
+         self.assertEqual(D.__mro__, (D, B, C, A, object))
+-        class E(C, B): pass
++        with torch._dynamo.error_on_graph_break(False):
++            class E(C, B): pass
+         self.assertEqual(E().spam(), "B")
+         self.assertEqual(E().boo(), "C")
+         self.assertEqual(E.__mro__, (E, C, B, A, object))
+         # MRO order disagreement
+         try:
+-            class F(D, E): pass
++            with torch._dynamo.error_on_graph_break(False):
++                class F(D, E): pass
+         except TypeError:
+             pass
+         else:
+             self.fail("expected MRO order disagreement (F)")
+         try:
+-            class G(E, D): pass
++            with torch._dynamo.error_on_graph_break(False):
++                class G(E, D): pass
+         except TypeError:
+             pass
+         else:
+@@ -936,26 +1000,28 @@ class ClassPropertiesAndMethods(unittest.TestCase, ExtraAssertions):
+     # see thread python-dev/2002-October/029035.html
+     def test_ex5_from_c3_switch(self):
+         # Testing ex5 from C3 switch discussion...
+-        class A(object): pass
+-        class B(object): pass
+-        class C(object): pass
+-        class X(A): pass
+-        class Y(A): pass
+-        class Z(X,B,Y,C): pass
++        with torch._dynamo.error_on_graph_break(False):
++            class A(object): pass
++            class B(object): pass
++            class C(object): pass
++            class X(A): pass
++            class Y(A): pass
++            class Z(X,B,Y,C): pass
+         self.assertEqual(Z.__mro__, (Z, X, B, Y, A, C, object))
+ 
+     # see "A Monotonic Superclass Linearization for Dylan",
+     # by Kim Barrett et al. (OOPSLA 1996)
+     def test_monotonicity(self):
+         # Testing MRO monotonicity...
+-        class Boat(object): pass
+-        class DayBoat(Boat): pass
+-        class WheelBoat(Boat): pass
+-        class EngineLess(DayBoat): pass
+-        class SmallMultihull(DayBoat): pass
+-        class PedalWheelBoat(EngineLess,WheelBoat): pass
+-        class SmallCatamaran(SmallMultihull): pass
+-        class Pedalo(PedalWheelBoat,SmallCatamaran): pass
++        with torch._dynamo.error_on_graph_break(False):
++            class Boat(object): pass
++            class DayBoat(Boat): pass
++            class WheelBoat(Boat): pass
++            class EngineLess(DayBoat): pass
++            class SmallMultihull(DayBoat): pass
++            class PedalWheelBoat(EngineLess,WheelBoat): pass
++            class SmallCatamaran(SmallMultihull): pass
++            class Pedalo(PedalWheelBoat,SmallCatamaran): pass
+ 
+         self.assertEqual(PedalWheelBoat.__mro__,
+               (PedalWheelBoat, EngineLess, DayBoat, WheelBoat, Boat, object))
+@@ -969,12 +1035,13 @@ class ClassPropertiesAndMethods(unittest.TestCase, ExtraAssertions):
+     # by Kim Barrett et al. (OOPSLA 1996)
+     def test_consistency_with_epg(self):
+         # Testing consistency with EPG...
+-        class Pane(object): pass
+-        class ScrollingMixin(object): pass
+-        class EditingMixin(object): pass
+-        class ScrollablePane(Pane,ScrollingMixin): pass
+-        class EditablePane(Pane,EditingMixin): pass
+-        class EditableScrollablePane(ScrollablePane,EditablePane): pass
++        with torch._dynamo.error_on_graph_break(False):
++            class Pane(object): pass
++            class ScrollingMixin(object): pass
++            class EditingMixin(object): pass
++            class ScrollablePane(Pane,ScrollingMixin): pass
++            class EditablePane(Pane,EditingMixin): pass
++            class EditableScrollablePane(ScrollablePane,EditablePane): pass
+ 
+         self.assertEqual(EditableScrollablePane.__mro__,
+               (EditableScrollablePane, ScrollablePane, EditablePane, Pane,
+@@ -997,9 +1064,10 @@ class ClassPropertiesAndMethods(unittest.TestCase, ExtraAssertions):
+             else:
+                 self.fail("Expected %s" % exc)
+ 
+-        class A(object): pass
+-        class B(A): pass
+-        class C(object): pass
++        with torch._dynamo.error_on_graph_break(False):
++            class A(object): pass
++            class B(A): pass
++            class C(object): pass
+ 
+         # Test some very simple errors
+         raises(TypeError, "duplicate base class A",
+@@ -1009,11 +1077,12 @@ class ClassPropertiesAndMethods(unittest.TestCase, ExtraAssertions):
+         raises(TypeError, mro_err_msg,
+                type, "X", (A, C, B), {})
+         # Test a slightly more complex error
+-        class GridLayout(object): pass
+-        class HorizontalGrid(GridLayout): pass
+-        class VerticalGrid(GridLayout): pass
+-        class HVGrid(HorizontalGrid, VerticalGrid): pass
+-        class VHGrid(VerticalGrid, HorizontalGrid): pass
++        with torch._dynamo.error_on_graph_break(False):
++            class GridLayout(object): pass
++            class HorizontalGrid(GridLayout): pass
++            class VerticalGrid(GridLayout): pass
++            class HVGrid(HorizontalGrid, VerticalGrid): pass
++            class VHGrid(VerticalGrid, HorizontalGrid): pass
+         raises(TypeError, mro_err_msg,
+                type, "ConfusedGrid", (HVGrid, VHGrid), {})
+ 
+@@ -1033,8 +1102,9 @@ class ClassPropertiesAndMethods(unittest.TestCase, ExtraAssertions):
+             self.fail("object() should not allow setting a foo attribute")
+         self.assertNotHasAttr(object(), "__dict__")
+ 
+-        class Cdict(object):
+-            pass
++        with torch._dynamo.error_on_graph_break(False):
++            class Cdict(object):
++                pass
+         x = Cdict()
+         self.assertEqual(x.__dict__, {})
+         x.foo = 1
+@@ -1042,8 +1112,9 @@ class ClassPropertiesAndMethods(unittest.TestCase, ExtraAssertions):
+         self.assertEqual(x.__dict__, {'foo': 1})
+ 
+     def test_object_class_assignment_between_heaptypes_and_nonheaptypes(self):
+-        class SubType(types.ModuleType):
+-            a = 1
++        with torch._dynamo.error_on_graph_break(False):
++            class SubType(types.ModuleType):
++                a = 1
+ 
+         m = types.ModuleType("m")
+         self.assertTrue(m.__class__ is types.ModuleType)
+@@ -1067,51 +1138,61 @@ class ClassPropertiesAndMethods(unittest.TestCase, ExtraAssertions):
+         # subclassable and thus don't need to be checked:
+         #   NoneType, bool
+ 
+-        class MyInt(int):
+-            __slots__ = ()
++        with torch._dynamo.error_on_graph_break(False):
++            class MyInt(int):
++                __slots__ = ()
+         with self.assertRaises(TypeError):
+             (1).__class__ = MyInt
+ 
+-        class MyFloat(float):
+-            __slots__ = ()
++        with torch._dynamo.error_on_graph_break(False):
++            class MyFloat(float):
++                __slots__ = ()
+         with self.assertRaises(TypeError):
+             (1.0).__class__ = MyFloat
+ 
+-        class MyComplex(complex):
+-            __slots__ = ()
++        with torch._dynamo.error_on_graph_break(False):
++            class MyComplex(complex):
++                __slots__ = ()
+         with self.assertRaises(TypeError):
+             (1 + 2j).__class__ = MyComplex
+ 
+-        class MyStr(str):
+-            __slots__ = ()
++        with torch._dynamo.error_on_graph_break(False):
++            class MyStr(str):
++                __slots__ = ()
+         with self.assertRaises(TypeError):
+             "a".__class__ = MyStr
+ 
+-        class MyBytes(bytes):
+-            __slots__ = ()
++        with torch._dynamo.error_on_graph_break(False):
++            class MyBytes(bytes):
++                __slots__ = ()
+         with self.assertRaises(TypeError):
+             b"a".__class__ = MyBytes
+ 
+-        class MyTuple(tuple):
+-            __slots__ = ()
++        with torch._dynamo.error_on_graph_break(False):
++            class MyTuple(tuple):
++                __slots__ = ()
+         with self.assertRaises(TypeError):
+             ().__class__ = MyTuple
+ 
+-        class MyFrozenSet(frozenset):
+-            __slots__ = ()
++        with torch._dynamo.error_on_graph_break(False):
++            class MyFrozenSet(frozenset):
++                __slots__ = ()
+         with self.assertRaises(TypeError):
+             frozenset().__class__ = MyFrozenSet
+ 
++    @unittest.expectedFailure
+     def test_slots(self):
+         # Testing __slots__...
+-        class C0(object):
+-            __slots__ = []
++        with torch._dynamo.error_on_graph_break(False):
++            class C0(object):
++                __slots__ = []
+         x = C0()
+         self.assertNotHasAttr(x, "__dict__")
+         self.assertNotHasAttr(x, "foo")
+ 
+-        class C1(object):
+-            __slots__ = ['a']
++        with torch._dynamo.error_on_graph_break(False):
++            class C1(object):
++                __slots__ = ['a']
+         x = C1()
+         self.assertNotHasAttr(x, "__dict__")
+         self.assertNotHasAttr(x, "a")
+@@ -1122,9 +1203,10 @@ class ClassPropertiesAndMethods(unittest.TestCase, ExtraAssertions):
+         del x.a
+         self.assertNotHasAttr(x, "a")
+ 
+-        class C3(object):
+-            __slots__ = ['a', 'b', 'c']
+-        x = C3()
++        with torch._dynamo.error_on_graph_break(False):
++            class C3(object):
++                __slots__ = ['a', 'b', 'c']
++            x = C3()
+         self.assertNotHasAttr(x, "__dict__")
+         self.assertNotHasAttr(x, 'a')
+         self.assertNotHasAttr(x, 'b')
+@@ -1136,14 +1218,15 @@ class ClassPropertiesAndMethods(unittest.TestCase, ExtraAssertions):
+         self.assertEqual(x.b, 2)
+         self.assertEqual(x.c, 3)
+ 
+-        class C4(object):
+-            """Validate name mangling"""
+-            __slots__ = ['__a']
+-            def __init__(self, value):
+-                self.__a = value
+-            def get(self):
+-                return self.__a
+-        x = C4(5)
++        with torch._dynamo.error_on_graph_break(False):
++            class C4(object):
++                """Validate name mangling"""
++                __slots__ = ['__a']
++                def __init__(self, value):
++                    self.__a = value
++                def get(self):
++                    return self.__a
++            x = C4(5)
+         self.assertNotHasAttr(x, '__dict__')
+         self.assertNotHasAttr(x, '__a')
+         self.assertEqual(x.get(), 5)
+@@ -1156,85 +1239,97 @@ class ClassPropertiesAndMethods(unittest.TestCase, ExtraAssertions):
+ 
+         # Make sure slot names are proper identifiers
+         try:
+-            class C(object):
+-                __slots__ = [None]
++            with torch._dynamo.error_on_graph_break(False):
++                class C(object):
++                    __slots__ = [None]
+         except TypeError:
+             pass
+         else:
+             self.fail("[None] slots not caught")
+         try:
+-            class C(object):
+-                __slots__ = ["foo bar"]
++            with torch._dynamo.error_on_graph_break(False):
++                class C(object):
++                    __slots__ = ["foo bar"]
+         except TypeError:
+             pass
+         else:
+             self.fail("['foo bar'] slots not caught")
+         try:
+-            class C(object):
+-                __slots__ = ["foo\0bar"]
++            with torch._dynamo.error_on_graph_break(False):
++                class C(object):
++                    __slots__ = ["foo\0bar"]
+         except TypeError:
+             pass
+         else:
+             self.fail("['foo\\0bar'] slots not caught")
+         try:
+-            class C(object):
+-                __slots__ = ["1"]
++            with torch._dynamo.error_on_graph_break(False):
++                class C(object):
++                    __slots__ = ["1"]
+         except TypeError:
+             pass
+         else:
+             self.fail("['1'] slots not caught")
+         try:
+-            class C(object):
+-                __slots__ = [""]
++            with torch._dynamo.error_on_graph_break(False):
++                class C(object):
++                    __slots__ = [""]
+         except TypeError:
+             pass
+         else:
+             self.fail("[''] slots not caught")
+ 
+-        class WithValidIdentifiers(object):
+-            __slots__ = ["a", "a_b", "_a", "A0123456789Z"]
++        with torch._dynamo.error_on_graph_break(False):
++            class WithValidIdentifiers(object):
++                __slots__ = ["a", "a_b", "_a", "A0123456789Z"]
+ 
+         # Test a single string is not expanded as a sequence.
+-        class C(object):
+-            __slots__ = "abc"
++        with torch._dynamo.error_on_graph_break(False):
++            class C(object):
++                __slots__ = "abc"
+         c = C()
+         c.abc = 5
+         self.assertEqual(c.abc, 5)
+ 
+         # Test unicode slot names
+         # Test a single unicode string is not expanded as a sequence.
+-        class C(object):
+-            __slots__ = "abc"
++        with torch._dynamo.error_on_graph_break(False):
++            class C(object):
++                __slots__ = "abc"
+         c = C()
+         c.abc = 5
+         self.assertEqual(c.abc, 5)
+ 
+         # _unicode_to_string used to modify slots in certain circumstances
+         slots = ("foo", "bar")
+-        class C(object):
+-            __slots__ = slots
++        with torch._dynamo.error_on_graph_break(False):
++            class C(object):
++                __slots__ = slots
+         x = C()
+         x.foo = 5
+         self.assertEqual(x.foo, 5)
+         self.assertIs(type(slots[0]), str)
+         # this used to leak references
+         try:
+-            class C(object):
+-                __slots__ = [chr(128)]
++            with torch._dynamo.error_on_graph_break(False):
++                class C(object):
++                    __slots__ = [chr(128)]
+         except (TypeError, UnicodeEncodeError):
+             pass
+         else:
+             self.fail("[chr(128)] slots not caught")
+ 
+         # Test leaks
+-        class Counted(object):
+-            counter = 0    # counts the number of instances alive
+-            def __init__(self):
+-                Counted.counter += 1
+-            def __del__(self):
+-                Counted.counter -= 1
+-        class C(object):
+-            __slots__ = ['a', 'b', 'c']
++        with torch._dynamo.error_on_graph_break(False):
++            class Counted(object):
++                counter = 0    # counts the number of instances alive
++                def __init__(self):
++                    Counted.counter += 1
++                def __del__(self):
++                    Counted.counter -= 1
++        with torch._dynamo.error_on_graph_break(False):
++            class C(object):
++                __slots__ = ['a', 'b', 'c']
+         x = C()
+         x.a = Counted()
+         x.b = Counted()
+@@ -1243,8 +1338,9 @@ class ClassPropertiesAndMethods(unittest.TestCase, ExtraAssertions):
+         del x
+         support.gc_collect()
+         self.assertEqual(Counted.counter, 0)
+-        class D(C):
+-            pass
++        with torch._dynamo.error_on_graph_break(False):
++            class D(C):
++                pass
+         x = D()
+         x.a = Counted()
+         x.z = Counted()
+@@ -1252,8 +1348,9 @@ class ClassPropertiesAndMethods(unittest.TestCase, ExtraAssertions):
+         del x
+         support.gc_collect()
+         self.assertEqual(Counted.counter, 0)
+-        class E(D):
+-            __slots__ = ['e']
++        with torch._dynamo.error_on_graph_break(False):
++            class E(D):
++                __slots__ = ['e']
+         x = E()
+         x.a = Counted()
+         x.z = Counted()
+@@ -1264,8 +1361,9 @@ class ClassPropertiesAndMethods(unittest.TestCase, ExtraAssertions):
+         self.assertEqual(Counted.counter, 0)
+ 
+         # Test cyclical leaks [SF bug 519621]
+-        class F(object):
+-            __slots__ = ['a', 'b']
++        with torch._dynamo.error_on_graph_break(False):
++            class F(object):
++                __slots__ = ['a', 'b']
+         s = F()
+         s.a = [Counted(), s]
+         self.assertEqual(Counted.counter, 1)
+@@ -1275,9 +1373,10 @@ class ClassPropertiesAndMethods(unittest.TestCase, ExtraAssertions):
+ 
+         # Test lookup leaks [SF bug 572567]
+         if hasattr(gc, 'get_objects'):
+-            class G(object):
+-                def __eq__(self, other):
+-                    return False
++            with torch._dynamo.error_on_graph_break(False):
++                class G(object):
++                    def __eq__(self, other):
++                        return False
+             g = G()
+             orig_objects = len(gc.get_objects())
+             for i in range(10):
+@@ -1285,51 +1384,58 @@ class ClassPropertiesAndMethods(unittest.TestCase, ExtraAssertions):
+             new_objects = len(gc.get_objects())
+             self.assertEqual(orig_objects, new_objects)
+ 
+-        class H(object):
+-            __slots__ = ['a', 'b']
+-            def __init__(self):
+-                self.a = 1
+-                self.b = 2
+-            def __del__(self_):
+-                self.assertEqual(self_.a, 1)
+-                self.assertEqual(self_.b, 2)
++        with torch._dynamo.error_on_graph_break(False):
++            class H(object):
++                __slots__ = ['a', 'b']
++                def __init__(self):
++                    self.a = 1
++                    self.b = 2
++                def __del__(self_):
++                    self.assertEqual(self_.a, 1)
++                    self.assertEqual(self_.b, 2)
+         with support.captured_output('stderr') as s:
+             h = H()
+             del h
+         self.assertEqual(s.getvalue(), '')
+ 
+-        class X(object):
+-            __slots__ = "a"
++        with torch._dynamo.error_on_graph_break(False):
++            class X(object):
++                __slots__ = "a"
+         with self.assertRaises(AttributeError):
+             del X().a
+ 
+         # Inherit from object on purpose to check some backwards compatibility paths
+-        class X(object):
+-            __slots__ = "a"
++        with torch._dynamo.error_on_graph_break(False):
++            class X(object):
++                __slots__ = "a"
+         with self.assertRaisesRegex(AttributeError, "'test.test_descr.ClassPropertiesAndMethods.test_slots.<locals>.X' object has no attribute 'a'"):
+             X().a
+ 
+         # Test string subclass in `__slots__`, see gh-98783
+-        class SubStr(str):
+-            pass
+-        class X(object):
+-            __slots__ = (SubStr('x'),)
++        with torch._dynamo.error_on_graph_break(False):
++            class SubStr(str):
++                pass
++        with torch._dynamo.error_on_graph_break(False):
++            class X(object):
++                __slots__ = (SubStr('x'),)
+         X().x = 1
+         with self.assertRaisesRegex(AttributeError, "'X' object has no attribute 'a'"):
+             X().a
+ 
+     def test_slots_special(self):
+         # Testing __dict__ and __weakref__ in __slots__...
+-        class D(object):
+-            __slots__ = ["__dict__"]
++        with torch._dynamo.error_on_graph_break(False):
++            class D(object):
++                __slots__ = ["__dict__"]
+         a = D()
+         self.assertHasAttr(a, "__dict__")
+         self.assertNotHasAttr(a, "__weakref__")
+         a.foo = 42
+         self.assertEqual(a.__dict__, {"foo": 42})
+ 
+-        class W(object):
+-            __slots__ = ["__weakref__"]
++        with torch._dynamo.error_on_graph_break(False):
++            class W(object):
++                __slots__ = ["__weakref__"]
+         a = W()
+         self.assertHasAttr(a, "__weakref__")
+         self.assertNotHasAttr(a, "__dict__")
+@@ -1340,16 +1446,18 @@ class ClassPropertiesAndMethods(unittest.TestCase, ExtraAssertions):
+         else:
+             self.fail("shouldn't be allowed to set a.foo")
+ 
+-        class C1(W, D):
+-            __slots__ = []
++        with torch._dynamo.error_on_graph_break(False):
++            class C1(W, D):
++                __slots__ = []
+         a = C1()
+         self.assertHasAttr(a, "__dict__")
+         self.assertHasAttr(a, "__weakref__")
+         a.foo = 42
+         self.assertEqual(a.__dict__, {"foo": 42})
+ 
+-        class C2(D, W):
+-            __slots__ = []
++        with torch._dynamo.error_on_graph_break(False):
++            class C2(D, W):
++                __slots__ = []
+         a = C2()
+         self.assertHasAttr(a, "__dict__")
+         self.assertHasAttr(a, "__weakref__")
+@@ -1358,18 +1466,21 @@ class ClassPropertiesAndMethods(unittest.TestCase, ExtraAssertions):
+ 
+     def test_slots_special2(self):
+         # Testing __qualname__ and __classcell__ in __slots__
+-        class Meta(type):
+-            def __new__(cls, name, bases, namespace, attr):
+-                self.assertIn(attr, namespace)
+-                return super().__new__(cls, name, bases, namespace)
+-
+-        class C1:
+-            def __init__(self):
+-                self.b = 42
+-        class C2(C1, metaclass=Meta, attr="__classcell__"):
+-            __slots__ = ["__classcell__"]
+-            def __init__(self):
+-                super().__init__()
++        with torch._dynamo.error_on_graph_break(False):
++            class Meta(type):
++                def __new__(cls, name, bases, namespace, attr):
++                    self.assertIn(attr, namespace)
++                    return super().__new__(cls, name, bases, namespace)
++
++        with torch._dynamo.error_on_graph_break(False):
++            class C1:
++                def __init__(self):
++                    self.b = 42
++        with torch._dynamo.error_on_graph_break(False):
++            class C2(C1, metaclass=Meta, attr="__classcell__"):
++                __slots__ = ["__classcell__"]
++                def __init__(self):
++                    super().__init__()
+         self.assertIsInstance(C2.__dict__["__classcell__"],
+                               types.MemberDescriptorType)
+         c = C2()
+@@ -1378,12 +1489,14 @@ class ClassPropertiesAndMethods(unittest.TestCase, ExtraAssertions):
+         c.__classcell__ = 42
+         self.assertEqual(c.__classcell__, 42)
+         with self.assertRaises(TypeError):
+-            class C3:
+-                __classcell__ = 42
+-                __slots__ = ["__classcell__"]
++            with torch._dynamo.error_on_graph_break(False):
++                class C3:
++                    __classcell__ = 42
++                    __slots__ = ["__classcell__"]
+ 
+-        class Q1(metaclass=Meta, attr="__qualname__"):
+-            __slots__ = ["__qualname__"]
++        with torch._dynamo.error_on_graph_break(False):
++            class Q1(metaclass=Meta, attr="__qualname__"):
++                __slots__ = ["__qualname__"]
+         self.assertEqual(Q1.__qualname__, C1.__qualname__[:-2] + "Q1")
+         self.assertIsInstance(Q1.__dict__["__qualname__"],
+                               types.MemberDescriptorType)
+@@ -1392,19 +1505,22 @@ class ClassPropertiesAndMethods(unittest.TestCase, ExtraAssertions):
+         q.__qualname__ = "q"
+         self.assertEqual(q.__qualname__, "q")
+         with self.assertRaises(TypeError):
+-            class Q2:
+-                __qualname__ = object()
+-                __slots__ = ["__qualname__"]
++            with torch._dynamo.error_on_graph_break(False):
++                class Q2:
++                    __qualname__ = object()
++                    __slots__ = ["__qualname__"]
+ 
+     def test_slots_descriptor(self):
+         # Issue2115: slot descriptors did not correctly check
+         # the type of the given object
+         import abc
+-        class MyABC(metaclass=abc.ABCMeta):
+-            __slots__ = "a"
++        with torch._dynamo.error_on_graph_break(False):
++            class MyABC(metaclass=abc.ABCMeta):
++                __slots__ = "a"
+ 
+-        class Unrelated(object):
+-            pass
++        with torch._dynamo.error_on_graph_break(False):
++            class Unrelated(object):
++                pass
+         MyABC.register(Unrelated)
+ 
+         u = Unrelated()
+@@ -1415,20 +1531,22 @@ class ClassPropertiesAndMethods(unittest.TestCase, ExtraAssertions):
+ 
+     def test_dynamics(self):
+         # Testing class attribute propagation...
+-        class D(object):
+-            pass
+-        class E(D):
+-            pass
+-        class F(D):
+-            pass
++        with torch._dynamo.error_on_graph_break(False):
++            class D(object):
++                pass
++            class E(D):
++                pass
++            class F(D):
++                pass
+         D.foo = 1
+         self.assertEqual(D.foo, 1)
+         # Test that dynamic attributes are inherited
+         self.assertEqual(E.foo, 1)
+         self.assertEqual(F.foo, 1)
+         # Test dynamic instances
+-        class C(object):
+-            pass
++        with torch._dynamo.error_on_graph_break(False):
++            class C(object):
++                pass
+         a = C()
+         self.assertNotHasAttr(a, "foobar")
+         C.foobar = 2
+@@ -1458,15 +1576,17 @@ class ClassPropertiesAndMethods(unittest.TestCase, ExtraAssertions):
+             a.spam = "not spam"
+ 
+         self.assertEqual(a.spam, "spam")
+-        class D(C):
+-            pass
++        with torch._dynamo.error_on_graph_break(False):
++            class D(C):
++                pass
+         d = D()
+         d.foo = 1
+         self.assertEqual(d.foo, 1)
+ 
+         # Test handling of int*seq and seq*int
+-        class I(int):
+-            pass
++        with torch._dynamo.error_on_graph_break(False):
++            class I(int):
++                pass
+         self.assertEqual("a"*I(2), "aa")
+         self.assertEqual(I(2)*"a", "aa")
+         self.assertEqual(2*I(3), 6)
+@@ -1474,67 +1594,76 @@ class ClassPropertiesAndMethods(unittest.TestCase, ExtraAssertions):
+         self.assertEqual(I(3)*I(2), 6)
+ 
+         # Test comparison of classes with dynamic metaclasses
+-        class dynamicmetaclass(type):
+-            pass
+-        class someclass(metaclass=dynamicmetaclass):
+-            pass
++        with torch._dynamo.error_on_graph_break(False):
++            class dynamicmetaclass(type):
++                pass
++            class someclass(metaclass=dynamicmetaclass):
++                pass
+         self.assertNotEqual(someclass, object)
+ 
+     def test_errors(self):
+         # Testing errors...
+         try:
+-            class C(list, dict):
+-                pass
++            with torch._dynamo.error_on_graph_break(False):
++                class C(list, dict):
++                    pass
+         except TypeError:
+             pass
+         else:
+             self.fail("inheritance from both list and dict should be illegal")
+ 
+         try:
+-            class C(object, None):
+-                pass
++            with torch._dynamo.error_on_graph_break(False):
++                class C(object, None):
++                    pass
+         except TypeError:
+             pass
+         else:
+             self.fail("inheritance from non-type should be illegal")
+-        class Classic:
+-            pass
++        with torch._dynamo.error_on_graph_break(False):
++            class Classic:
++                pass
+ 
+         try:
+-            class C(type(len)):
+-                pass
++            with torch._dynamo.error_on_graph_break(False):
++                class C(type(len)):
++                    pass
+         except TypeError:
+             pass
+         else:
+             self.fail("inheritance from CFunction should be illegal")
+ 
+         try:
+-            class C(object):
+-                __slots__ = 1
++            with torch._dynamo.error_on_graph_break(False):
++                class C(object):
++                    __slots__ = 1
+         except TypeError:
+             pass
+         else:
+             self.fail("__slots__ = 1 should be illegal")
+ 
+         try:
+-            class C(object):
+-                __slots__ = [1]
++            with torch._dynamo.error_on_graph_break(False):
++                class C(object):
++                    __slots__ = [1]
+         except TypeError:
+             pass
+         else:
+             self.fail("__slots__ = [1] should be illegal")
+ 
+-        class M1(type):
+-            pass
+-        class M2(type):
+-            pass
+-        class A1(object, metaclass=M1):
+-            pass
+-        class A2(object, metaclass=M2):
+-            pass
+-        try:
+-            class B(A1, A2):
++        with torch._dynamo.error_on_graph_break(False):
++            class M1(type):
++                pass
++            class M2(type):
+                 pass
++            class A1(object, metaclass=M1):
++                pass
++            class A2(object, metaclass=M2):
++                pass
++        try:
++            with torch._dynamo.error_on_graph_break(False):
++                class B(A1, A2):
++                    pass
+         except TypeError:
+             pass
+         else:
+@@ -1542,15 +1671,17 @@ class ClassPropertiesAndMethods(unittest.TestCase, ExtraAssertions):
+ 
+     def test_classmethods(self):
+         # Testing class methods...
+-        class C(object):
+-            def foo(*a): return a
+-            goo = classmethod(foo)
++        with torch._dynamo.error_on_graph_break(False):
++            class C(object):
++                def foo(*a): return a
++                goo = classmethod(foo)
+         c = C()
+         self.assertEqual(C.goo(1), (C, 1))
+         self.assertEqual(c.goo(1), (C, 1))
+         self.assertEqual(c.foo(1), (c, 1))
+-        class D(C):
+-            pass
++        with torch._dynamo.error_on_graph_break(False):
++            class D(C):
++                pass
+         d = D()
+         self.assertEqual(D.goo(1), (D, 1))
+         self.assertEqual(d.goo(1), (D, 1))
+@@ -1631,7 +1762,8 @@ class ClassPropertiesAndMethods(unittest.TestCase, ExtraAssertions):
+         self.assertEqual(x2, spam.spamlist)
+         self.assertEqual(a2, a1)
+         self.assertEqual(d2, d1)
+-        class SubSpam(spam.spamlist): pass
++        with torch._dynamo.error_on_graph_break(False):
++            class SubSpam(spam.spamlist): pass
+         x2, a2, d2 = spam_cm(SubSpam, *a, **d)
+         self.assertEqual(x2, SubSpam)
+         self.assertEqual(a2, a1)
+@@ -1664,15 +1796,17 @@ class ClassPropertiesAndMethods(unittest.TestCase, ExtraAssertions):
+ 
+     def test_staticmethods(self):
+         # Testing static methods...
+-        class C(object):
+-            def foo(*a): return a
+-            goo = staticmethod(foo)
++        with torch._dynamo.error_on_graph_break(False):
++            class C(object):
++                def foo(*a): return a
++                goo = staticmethod(foo)
+         c = C()
+         self.assertEqual(C.goo(1), (1,))
+         self.assertEqual(c.goo(1), (1,))
+         self.assertEqual(c.foo(1), (c, 1,))
+-        class D(C):
+-            pass
++        with torch._dynamo.error_on_graph_break(False):
++            class D(C):
++                pass
+         d = D()
+         self.assertEqual(D.goo(1), (1,))
+         self.assertEqual(d.goo(1), (1,))
+@@ -1713,50 +1847,54 @@ class ClassPropertiesAndMethods(unittest.TestCase, ExtraAssertions):
+ 
+     def test_classic(self):
+         # Testing classic classes...
+-        class C:
+-            def foo(*a): return a
+-            goo = classmethod(foo)
++        with torch._dynamo.error_on_graph_break(False):
++            class C:
++                def foo(*a): return a
++                goo = classmethod(foo)
+         c = C()
+         self.assertEqual(C.goo(1), (C, 1))
+         self.assertEqual(c.goo(1), (C, 1))
+         self.assertEqual(c.foo(1), (c, 1))
+-        class D(C):
+-            pass
++        with torch._dynamo.error_on_graph_break(False):
++            class D(C):
++                pass
+         d = D()
+         self.assertEqual(D.goo(1), (D, 1))
+         self.assertEqual(d.goo(1), (D, 1))
+         self.assertEqual(d.foo(1), (d, 1))
+         self.assertEqual(D.foo(d, 1), (d, 1))
+-        class E: # *not* subclassing from C
+-            foo = C.foo
++        with torch._dynamo.error_on_graph_break(False):
++            class E: # *not* subclassing from C
++                foo = C.foo
+         self.assertEqual(E().foo.__func__, C.foo) # i.e., unbound
+         self.assertTrue(repr(C.foo.__get__(C())).startswith("<bound method "))
+ 
+     def test_compattr(self):
+         # Testing computed attributes...
+-        class C(object):
+-            class computed_attribute(object):
+-                def __init__(self, get, set=None, delete=None):
+-                    self.__get = get
+-                    self.__set = set
+-                    self.__delete = delete
+-                def __get__(self, obj, type=None):
+-                    return self.__get(obj)
+-                def __set__(self, obj, value):
+-                    return self.__set(obj, value)
+-                def __delete__(self, obj):
+-                    return self.__delete(obj)
+-            def __init__(self):
+-                self.__x = 0
+-            def __get_x(self):
+-                x = self.__x
+-                self.__x = x+1
+-                return x
+-            def __set_x(self, x):
+-                self.__x = x
+-            def __delete_x(self):
+-                del self.__x
+-            x = computed_attribute(__get_x, __set_x, __delete_x)
++        with torch._dynamo.error_on_graph_break(False):
++            class C(object):
++                class computed_attribute(object):
++                    def __init__(self, get, set=None, delete=None):
++                        self.__get = get
++                        self.__set = set
++                        self.__delete = delete
++                    def __get__(self, obj, type=None):
++                        return self.__get(obj)
++                    def __set__(self, obj, value):
++                        return self.__set(obj, value)
++                    def __delete__(self, obj):
++                        return self.__delete(obj)
++                def __init__(self):
++                    self.__x = 0
++                def __get_x(self):
++                    x = self.__x
++                    self.__x = x+1
++                    return x
++                def __set_x(self, x):
++                    self.__x = x
++                def __delete_x(self):
++                    del self.__x
++                x = computed_attribute(__get_x, __set_x, __delete_x)
+         a = C()
+         self.assertEqual(a.x, 0)
+         self.assertEqual(a.x, 1)
+@@ -1768,18 +1906,20 @@ class ClassPropertiesAndMethods(unittest.TestCase, ExtraAssertions):
+ 
+     def test_newslots(self):
+         # Testing __new__ slot override...
+-        class C(list):
+-            def __new__(cls):
+-                self = list.__new__(cls)
+-                self.foo = 1
+-                return self
+-            def __init__(self):
+-                self.foo = self.foo + 2
++        with torch._dynamo.error_on_graph_break(False):
++            class C(list):
++                def __new__(cls):
++                    self = list.__new__(cls)
++                    self.foo = 1
++                    return self
++                def __init__(self):
++                    self.foo = self.foo + 2
+         a = C()
+         self.assertEqual(a.foo, 3)
+         self.assertEqual(a.__class__, C)
+-        class D(C):
+-            pass
++        with torch._dynamo.error_on_graph_break(False):
++            class D(C):
++                pass
+         b = D()
+         self.assertEqual(b.foo, 3)
+         self.assertEqual(b.__class__, D)
+@@ -1798,34 +1938,38 @@ class ClassPropertiesAndMethods(unittest.TestCase, ExtraAssertions):
+         self.assertRaises(TypeError, C)
+ 
+     def test_object_new(self):
+-        class A(object):
+-            pass
++        with torch._dynamo.error_on_graph_break(False):
++            class A(object):
++                pass
+         object.__new__(A)
+         self.assertRaises(TypeError, object.__new__, A, 5)
+         object.__init__(A())
+         self.assertRaises(TypeError, object.__init__, A(), 5)
+ 
+-        class A(object):
+-            def __init__(self, foo):
+-                self.foo = foo
++        with torch._dynamo.error_on_graph_break(False):
++            class A(object):
++                def __init__(self, foo):
++                    self.foo = foo
+         object.__new__(A)
+         object.__new__(A, 5)
+         object.__init__(A(3))
+         self.assertRaises(TypeError, object.__init__, A(3), 5)
+ 
+-        class A(object):
+-            def __new__(cls, foo):
+-                return object.__new__(cls)
++        with torch._dynamo.error_on_graph_break(False):
++            class A(object):
++                def __new__(cls, foo):
++                    return object.__new__(cls)
+         object.__new__(A)
+         self.assertRaises(TypeError, object.__new__, A, 5)
+         object.__init__(A(3))
+         object.__init__(A(3), 5)
+ 
+-        class A(object):
+-            def __new__(cls, foo):
+-                return object.__new__(cls)
+-            def __init__(self, foo):
+-                self.foo = foo
++        with torch._dynamo.error_on_graph_break(False):
++            class A(object):
++                def __new__(cls, foo):
++                    return object.__new__(cls)
++                def __init__(self, foo):
++                    self.foo = foo
+         object.__new__(A)
+         self.assertRaises(TypeError, object.__new__, A, 5)
+         object.__init__(A(3))
+@@ -1857,14 +2001,15 @@ class ClassPropertiesAndMethods(unittest.TestCase, ExtraAssertions):
+ 
+     def test_altmro(self):
+         # Testing mro() and overriding it...
+-        class A(object):
+-            def f(self): return "A"
+-        class B(A):
+-            pass
+-        class C(A):
+-            def f(self): return "C"
+-        class D(B, C):
+-            pass
++        with torch._dynamo.error_on_graph_break(False):
++            class A(object):
++                def f(self): return "A"
++            class B(A):
++                pass
++            class C(A):
++                def f(self): return "C"
++            class D(B, C):
++                pass
+         self.assertEqual(A.mro(), [A, object])
+         self.assertEqual(A.__mro__, (A, object))
+         self.assertEqual(B.mro(), [B, A, object])
+@@ -1875,22 +2020,24 @@ class ClassPropertiesAndMethods(unittest.TestCase, ExtraAssertions):
+         self.assertEqual(D.__mro__, (D, B, C, A, object))
+         self.assertEqual(D().f(), "C")
+ 
+-        class PerverseMetaType(type):
+-            def mro(cls):
+-                L = type.mro(cls)
+-                L.reverse()
+-                return L
+-        class X(D,B,C,A, metaclass=PerverseMetaType):
+-            pass
++        with torch._dynamo.error_on_graph_break(False):
++            class PerverseMetaType(type):
++                def mro(cls):
++                    L = type.mro(cls)
++                    L.reverse()
++                    return L
++            class X(D,B,C,A, metaclass=PerverseMetaType):
++                pass
+         self.assertEqual(X.__mro__, (object, A, C, B, D, X))
+         self.assertEqual(X().f(), "A")
+ 
+         try:
+-            class _metaclass(type):
+-                def mro(self):
+-                    return [self, dict, object]
+-            class X(object, metaclass=_metaclass):
+-                pass
++            with torch._dynamo.error_on_graph_break(False):
++                class _metaclass(type):
++                    def mro(self):
++                        return [self, dict, object]
++                class X(object, metaclass=_metaclass):
++                    pass
+             # In CPython, the class creation above already raises
+             # TypeError, as a protection against the fact that
+             # instances of X would segfault it.  In other Python
+@@ -1905,22 +2052,24 @@ class ClassPropertiesAndMethods(unittest.TestCase, ExtraAssertions):
+             self.fail("devious mro() return not caught")
+ 
+         try:
+-            class _metaclass(type):
+-                def mro(self):
+-                    return [1]
+-            class X(object, metaclass=_metaclass):
+-                pass
++            with torch._dynamo.error_on_graph_break(False):
++                class _metaclass(type):
++                    def mro(self):
++                        return [1]
++                class X(object, metaclass=_metaclass):
++                    pass
+         except TypeError:
+             pass
+         else:
+             self.fail("non-class mro() return not caught")
+ 
+         try:
+-            class _metaclass(type):
+-                def mro(self):
+-                    return 1
+-            class X(object, metaclass=_metaclass):
+-                pass
++            with torch._dynamo.error_on_graph_break(False):
++                class _metaclass(type):
++                    def mro(self):
++                        return 1
++                class X(object, metaclass=_metaclass):
++                    pass
+         except TypeError:
+             pass
+         else:
+@@ -1929,32 +2078,33 @@ class ClassPropertiesAndMethods(unittest.TestCase, ExtraAssertions):
+     def test_overloading(self):
+         # Testing operator overloading...
+ 
+-        class B(object):
+-            "Intermediate class because object doesn't have a __setattr__"
++        with torch._dynamo.error_on_graph_break(False):
++            class B(object):
++                "Intermediate class because object doesn't have a __setattr__"
+ 
+-        class C(B):
+-            def __getattr__(self, name):
+-                if name == "foo":
+-                    return ("getattr", name)
+-                else:
+-                    raise AttributeError
+-            def __setattr__(self, name, value):
+-                if name == "foo":
+-                    self.setattr = (name, value)
+-                else:
+-                    return B.__setattr__(self, name, value)
+-            def __delattr__(self, name):
+-                if name == "foo":
+-                    self.delattr = name
+-                else:
+-                    return B.__delattr__(self, name)
++            class C(B):
++                def __getattr__(self, name):
++                    if name == "foo":
++                        return ("getattr", name)
++                    else:
++                        raise AttributeError
++                def __setattr__(self, name, value):
++                    if name == "foo":
++                        self.setattr = (name, value)
++                    else:
++                        return B.__setattr__(self, name, value)
++                def __delattr__(self, name):
++                    if name == "foo":
++                        self.delattr = name
++                    else:
++                        return B.__delattr__(self, name)
+ 
+-            def __getitem__(self, key):
+-                return ("getitem", key)
+-            def __setitem__(self, key, value):
+-                self.setitem = (key, value)
+-            def __delitem__(self, key):
+-                self.delitem = key
++                def __getitem__(self, key):
++                    return ("getitem", key)
++                def __setitem__(self, key, value):
++                    self.setitem = (key, value)
++                def __delitem__(self, key):
++                    self.delitem = key
+ 
+         a = C()
+         self.assertEqual(a.foo, ("getattr", "foo"))
+@@ -1977,9 +2127,10 @@ class ClassPropertiesAndMethods(unittest.TestCase, ExtraAssertions):
+ 
+     def test_load_attr_extended_arg(self):
+         # https://github.com/python/cpython/issues/91625
+-        class Numbers:
+-            def __getattr__(self, attr):
+-                return int(attr.lstrip("_"))
++        with torch._dynamo.error_on_graph_break(False):
++            class Numbers:
++                def __getattr__(self, attr):
++                    return int(attr.lstrip("_"))
+         attrs = ", ".join(f"Z._{n:03d}" for n in range(280))
+         code = f"def number_attrs(Z):\n    return [ {attrs} ]"
+         ns = {}
+@@ -1991,22 +2142,25 @@ class ClassPropertiesAndMethods(unittest.TestCase, ExtraAssertions):
+ 
+     def test_methods(self):
+         # Testing methods...
+-        class C(object):
+-            def __init__(self, x):
+-                self.x = x
+-            def foo(self):
+-                return self.x
++        with torch._dynamo.error_on_graph_break(False):
++            class C(object):
++                def __init__(self, x):
++                    self.x = x
++                def foo(self):
++                    return self.x
+         c1 = C(1)
+         self.assertEqual(c1.foo(), 1)
+-        class D(C):
+-            boo = C.foo
+-            goo = c1.foo
++        with torch._dynamo.error_on_graph_break(False):
++            class D(C):
++                boo = C.foo
++                goo = c1.foo
+         d2 = D(2)
+         self.assertEqual(d2.foo(), 2)
+         self.assertEqual(d2.boo(), 2)
+         self.assertEqual(d2.goo(), 1)
+-        class E(object):
+-            foo = C.foo
++        with torch._dynamo.error_on_graph_break(False):
++            class E(object):
++                foo = C.foo
+         self.assertEqual(E().foo.__func__, C.foo) # i.e., unbound
+         self.assertTrue(repr(C.foo.__get__(C(1))).startswith("<bound method "))
+ 
+@@ -2059,8 +2213,9 @@ class ClassPropertiesAndMethods(unittest.TestCase, ExtraAssertions):
+         def do_issubclass(obj):
+             return issubclass(int, obj)
+         def do_dict_missing(checker):
+-            class DictSub(checker.__class__, dict):
+-                pass
++            with torch._dynamo.error_on_graph_break(False):
++                class DictSub(checker.__class__, dict):
++                    pass
+             self.assertEqual(DictSub()["hi"], 4)
+         def some_number(self_, key):
+             self.assertEqual(key, "hi")
+@@ -2094,44 +2249,48 @@ class ClassPropertiesAndMethods(unittest.TestCase, ExtraAssertions):
+             ("__round__", round, zero, set(), {}),
+             ]
+ 
+-        class Checker(object):
+-            def __getattr__(self, attr, test=self):
+-                test.fail("__getattr__ called with {0}".format(attr))
+-            def __getattribute__(self, attr, test=self):
+-                if attr not in ok:
+-                    test.fail("__getattribute__ called with {0}".format(attr))
+-                return object.__getattribute__(self, attr)
+-        class SpecialDescr(object):
+-            def __init__(self, impl):
+-                self.impl = impl
+-            def __get__(self, obj, owner):
+-                record.append(1)
+-                return self.impl.__get__(obj, owner)
+-        class MyException(Exception):
+-            pass
+-        class ErrDescr(object):
+-            def __get__(self, obj, owner):
+-                raise MyException
++        with torch._dynamo.error_on_graph_break(False):
++            class Checker(object):
++                def __getattr__(self, attr, test=self):
++                    test.fail("__getattr__ called with {0}".format(attr))
++                def __getattribute__(self, attr, test=self):
++                    if attr not in ok:
++                        test.fail("__getattribute__ called with {0}".format(attr))
++                    return object.__getattribute__(self, attr)
++            class SpecialDescr(object):
++                def __init__(self, impl):
++                    self.impl = impl
++                def __get__(self, obj, owner):
++                    record.append(1)
++                    return self.impl.__get__(obj, owner)
++            class MyException(Exception):
++                pass
++            class ErrDescr(object):
++                def __get__(self, obj, owner):
++                    raise MyException
+ 
+         for name, runner, meth_impl, ok, env in specials:
+-            class X(Checker):
+-                pass
++            with torch._dynamo.error_on_graph_break(False):
++                class X(Checker):
++                    pass
+             for attr, obj in env.items():
+                 setattr(X, attr, obj)
+             setattr(X, name, meth_impl)
+             runner(X())
+ 
+             record = []
+-            class X(Checker):
+-                pass
++            with torch._dynamo.error_on_graph_break(False):
++                class X(Checker):
++                    pass
+             for attr, obj in env.items():
+                 setattr(X, attr, obj)
+             setattr(X, name, SpecialDescr(meth_impl))
+             runner(X())
+             self.assertEqual(record, [1], name)
+ 
+-            class X(Checker):
+-                pass
++            with torch._dynamo.error_on_graph_break(False):
++                class X(Checker):
++                    pass
+             for attr, obj in env.items():
+                 setattr(X, attr, obj)
+             setattr(X, name, ErrDescr())
+@@ -2142,10 +2301,11 @@ class ClassPropertiesAndMethods(unittest.TestCase, ExtraAssertions):
+         # Test operators like __hash__ for which a built-in default exists
+ 
+         # Test the default behavior for static classes
+-        class C(object):
+-            def __getitem__(self, i):
+-                if 0 <= i < 10: return i
+-                raise IndexError
++        with torch._dynamo.error_on_graph_break(False):
++            class C(object):
++                def __getitem__(self, i):
++                    if 0 <= i < 10: return i
++                    raise IndexError
+         c1 = C()
+         c2 = C()
+         self.assertFalse(not c1)
+@@ -2165,10 +2325,11 @@ class ClassPropertiesAndMethods(unittest.TestCase, ExtraAssertions):
+             self.assertIn(i, c1)
+         self.assertNotIn(10, c1)
+         # Test the default behavior for dynamic classes
+-        class D(object):
+-            def __getitem__(self, i):
+-                if 0 <= i < 10: return i
+-                raise IndexError
++        with torch._dynamo.error_on_graph_break(False):
++            class D(object):
++                def __getitem__(self, i):
++                    if 0 <= i < 10: return i
++                    raise IndexError
+         d1 = D()
+         d2 = D()
+         self.assertFalse(not d1)
+@@ -2188,31 +2349,32 @@ class ClassPropertiesAndMethods(unittest.TestCase, ExtraAssertions):
+             self.assertIn(i, d1)
+         self.assertNotIn(10, d1)
+         # Test overridden behavior
+-        class Proxy(object):
+-            def __init__(self, x):
+-                self.x = x
+-            def __bool__(self):
+-                return not not self.x
+-            def __hash__(self):
+-                return hash(self.x)
+-            def __eq__(self, other):
+-                return self.x == other
+-            def __ne__(self, other):
+-                return self.x != other
+-            def __ge__(self, other):
+-                return self.x >= other
+-            def __gt__(self, other):
+-                return self.x > other
+-            def __le__(self, other):
+-                return self.x <= other
+-            def __lt__(self, other):
+-                return self.x < other
+-            def __str__(self):
+-                return "Proxy:%s" % self.x
+-            def __repr__(self):
+-                return "Proxy(%r)" % self.x
+-            def __contains__(self, value):
+-                return value in self.x
++        with torch._dynamo.error_on_graph_break(False):
++            class Proxy(object):
++                def __init__(self, x):
++                    self.x = x
++                def __bool__(self):
++                    return not not self.x
++                def __hash__(self):
++                    return hash(self.x)
++                def __eq__(self, other):
++                    return self.x == other
++                def __ne__(self, other):
++                    return self.x != other
++                def __ge__(self, other):
++                    return self.x >= other
++                def __gt__(self, other):
++                    return self.x > other
++                def __le__(self, other):
++                    return self.x <= other
++                def __lt__(self, other):
++                    return self.x < other
++                def __str__(self):
++                    return "Proxy:%s" % self.x
++                def __repr__(self):
++                    return "Proxy(%r)" % self.x
++                def __contains__(self, value):
++                    return value in self.x
+         p0 = Proxy(0)
+         p1 = Proxy(1)
+         p_1 = Proxy(-1)
+@@ -2238,8 +2400,9 @@ class ClassPropertiesAndMethods(unittest.TestCase, ExtraAssertions):
+     def test_weakrefs(self):
+         # Testing weak references...
+         import weakref
+-        class C(object):
+-            pass
++        with torch._dynamo.error_on_graph_break(False):
++            class C(object):
++                pass
+         c = C()
+         r = weakref.ref(c)
+         self.assertEqual(r(), c)
+@@ -2247,8 +2410,9 @@ class ClassPropertiesAndMethods(unittest.TestCase, ExtraAssertions):
+         support.gc_collect()
+         self.assertEqual(r(), None)
+         del r
+-        class NoWeak(object):
+-            __slots__ = ['foo']
++        with torch._dynamo.error_on_graph_break(False):
++            class NoWeak(object):
++                __slots__ = ['foo']
+         no = NoWeak()
+         try:
+             weakref.ref(no)
+@@ -2256,8 +2420,9 @@ class ClassPropertiesAndMethods(unittest.TestCase, ExtraAssertions):
+             self.assertIn("weak reference", str(msg))
+         else:
+             self.fail("weakref.ref(no) should be illegal")
+-        class Weak(object):
+-            __slots__ = ['foo', '__weakref__']
++        with torch._dynamo.error_on_graph_break(False):
++            class Weak(object):
++                __slots__ = ['foo', '__weakref__']
+         yes = Weak()
+         r = weakref.ref(yes)
+         self.assertEqual(r(), yes)
+@@ -2268,14 +2433,15 @@ class ClassPropertiesAndMethods(unittest.TestCase, ExtraAssertions):
+ 
+     def test_properties(self):
+         # Testing property...
+-        class C(object):
+-            def getx(self):
+-                return self.__x
+-            def setx(self, value):
+-                self.__x = value
+-            def delx(self):
+-                del self.__x
+-            x = property(getx, setx, delx, doc="I'm the x property.")
++        with torch._dynamo.error_on_graph_break(False):
++            class C(object):
++                def getx(self):
++                    return self.__x
++                def setx(self, value):
++                    self.__x = value
++                def delx(self):
++                    del self.__x
++                x = property(getx, setx, delx, doc="I'm the x property.")
+         a = C()
+         self.assertNotHasAttr(a, "x")
+         a.x = 42
+@@ -2317,8 +2483,9 @@ class ClassPropertiesAndMethods(unittest.TestCase, ExtraAssertions):
+         raw.__doc__ = 42
+         self.assertEqual(raw.__doc__, 42)
+ 
+-        class D(object):
+-            __getitem__ = property(lambda s: 1/0)
++        with torch._dynamo.error_on_graph_break(False):
++            class D(object):
++                __getitem__ = property(lambda s: 1/0)
+ 
+         d = D()
+         try:
+@@ -2332,17 +2499,18 @@ class ClassPropertiesAndMethods(unittest.TestCase, ExtraAssertions):
+     @unittest.skipIf(sys.flags.optimize >= 2,
+                      "Docstrings are omitted with -O2 and above")
+     def test_properties_doc_attrib(self):
+-        class E(object):
+-            def getter(self):
+-                "getter method"
+-                return 0
+-            def setter(self_, value):
+-                "setter method"
+-                pass
+-            prop = property(getter)
+-            self.assertEqual(prop.__doc__, "getter method")
+-            prop2 = property(fset=setter)
+-            self.assertEqual(prop2.__doc__, None)
++        with torch._dynamo.error_on_graph_break(False):
++            class E(object):
++                def getter(self):
++                    "getter method"
++                    return 0
++                def setter(self_, value):
++                    "setter method"
++                    pass
++                prop = property(getter)
++                self.assertEqual(prop.__doc__, "getter method")
++                prop2 = property(fset=setter)
++                self.assertEqual(prop2.__doc__, None)
+ 
+     @support.cpython_only
+     def test_testcapi_no_segfault(self):
+@@ -2352,21 +2520,23 @@ class ClassPropertiesAndMethods(unittest.TestCase, ExtraAssertions):
+         except ImportError:
+             pass
+         else:
+-            class X(object):
+-                p = property(_testcapi.test_with_docstring)
++            with torch._dynamo.error_on_graph_break(False):
++                class X(object):
++                    p = property(_testcapi.test_with_docstring)
+ 
+     def test_properties_plus(self):
+-        class C(object):
+-            foo = property(doc="hello")
+-            @foo.getter
+-            def foo(self):
+-                return self._foo
+-            @foo.setter
+-            def foo(self, value):
+-                self._foo = abs(value)
+-            @foo.deleter
+-            def foo(self):
+-                del self._foo
++        with torch._dynamo.error_on_graph_break(False):
++            class C(object):
++                foo = property(doc="hello")
++                @foo.getter
++                def foo(self):
++                    return self._foo
++                @foo.setter
++                def foo(self, value):
++                    self._foo = abs(value)
++                @foo.deleter
++                def foo(self):
++                    del self._foo
+         c = C()
+         self.assertEqual(C.foo.__doc__, "hello")
+         self.assertNotHasAttr(c, "foo")
+@@ -2378,45 +2548,48 @@ class ClassPropertiesAndMethods(unittest.TestCase, ExtraAssertions):
+         self.assertNotHasAttr(c, '_foo')
+         self.assertNotHasAttr(c, "foo")
+ 
+-        class D(C):
+-            @C.foo.deleter
+-            def foo(self):
+-                try:
+-                    del self._foo
+-                except AttributeError:
+-                    pass
++        with torch._dynamo.error_on_graph_break(False):
++            class D(C):
++                @C.foo.deleter
++                def foo(self):
++                    try:
++                        del self._foo
++                    except AttributeError:
++                        pass
+         d = D()
+         d.foo = 24
+         self.assertEqual(d.foo, 24)
+         del d.foo
+         del d.foo
+ 
+-        class E(object):
+-            @property
+-            def foo(self):
+-                return self._foo
+-            @foo.setter
+-            def foo(self, value):
+-                raise RuntimeError
+-            @foo.setter
+-            def foo(self, value):
+-                self._foo = abs(value)
+-            @foo.deleter
+-            def foo(self, value=None):
+-                del self._foo
++        with torch._dynamo.error_on_graph_break(False):
++            class E(object):
++                @property
++                def foo(self):
++                    return self._foo
++                @foo.setter
++                def foo(self, value):
++                    raise RuntimeError
++                @foo.setter
++                def foo(self, value):
++                    self._foo = abs(value)
++                @foo.deleter
++                def foo(self, value=None):
++                    del self._foo
+ 
+         e = E()
+         e.foo = -42
+         self.assertEqual(e.foo, 42)
+         del e.foo
+ 
+-        class F(E):
+-            @E.foo.deleter
+-            def foo(self):
+-                del self._foo
+-            @foo.setter
+-            def foo(self, value):
+-                self._foo = max(0, value)
++        with torch._dynamo.error_on_graph_break(False):
++            class F(E):
++                @E.foo.deleter
++                def foo(self):
++                    del self._foo
++                @foo.setter
++                def foo(self, value):
++                    self._foo = max(0, value)
+         f = F()
+         f.foo = -10
+         self.assertEqual(f.foo, 0)
+@@ -2459,9 +2632,10 @@ class ClassPropertiesAndMethods(unittest.TestCase, ExtraAssertions):
+         with self.assertRaises(TypeError):
+             dict({}, {})
+ 
+-        class Mapping:
+-            # Lacks a .keys() method; will be added later.
+-            dict = {1:2, 3:4, 'a':1j}
++        with torch._dynamo.error_on_graph_break(False):
++            class Mapping:
++                # Lacks a .keys() method; will be added later.
++                dict = {1:2, 3:4, 'a':1j}
+ 
+         try:
+             dict(Mapping())
+@@ -2476,12 +2650,13 @@ class ClassPropertiesAndMethods(unittest.TestCase, ExtraAssertions):
+         self.assertEqual(d, Mapping.dict)
+ 
+         # Init from sequence of iterable objects, each producing a 2-sequence.
+-        class AddressBookEntry:
+-            def __init__(self, first, last):
+-                self.first = first
+-                self.last = last
+-            def __iter__(self):
+-                return iter([self.first, self.last])
++        with torch._dynamo.error_on_graph_break(False):
++            class AddressBookEntry:
++                def __init__(self, first, last):
++                    self.first = first
++                    self.last = last
++                def __iter__(self):
++                    return iter([self.first, self.last])
+ 
+         d = dict([AddressBookEntry('Tim', 'Warsaw'),
+                   AddressBookEntry('Barry', 'Peters'),
+@@ -2516,9 +2691,10 @@ class ClassPropertiesAndMethods(unittest.TestCase, ExtraAssertions):
+         def interesting(strings):
+             return [s for s in strings if not s.startswith('_')]
+ 
+-        class C(object):
+-            Cdata = 1
+-            def Cmethod(self): pass
++        with torch._dynamo.error_on_graph_break(False):
++            class C(object):
++                Cdata = 1
++                def Cmethod(self): pass
+ 
+         cstuff = ['Cdata', 'Cmethod']
+         self.assertEqual(interesting(dir(C)), cstuff)
+@@ -2532,9 +2708,10 @@ class ClassPropertiesAndMethods(unittest.TestCase, ExtraAssertions):
+         self.assertEqual(interesting(dir(c)), cstuff + ['cdata', 'cmethod'])
+         ## self.assertIn('__self__', dir(c.Cmethod))
+ 
+-        class A(C):
+-            Adata = 1
+-            def Amethod(self): pass
++        with torch._dynamo.error_on_graph_break(False):
++            class A(C):
++                Adata = 1
++                def Amethod(self): pass
+ 
+         astuff = ['Adata', 'Amethod'] + cstuff
+         self.assertEqual(interesting(dir(A)), astuff)
+@@ -2547,8 +2724,9 @@ class ClassPropertiesAndMethods(unittest.TestCase, ExtraAssertions):
+         ## self.assertIn('__self__', dir(a.Amethod))
+ 
+         # Try a module subclass.
+-        class M(type(sys)):
+-            pass
++        with torch._dynamo.error_on_graph_break(False):
++            class M(type(sys)):
++                pass
+         minstance = M("m")
+         minstance.b = 2
+         minstance.a = 1
+@@ -2557,10 +2735,11 @@ class ClassPropertiesAndMethods(unittest.TestCase, ExtraAssertions):
+         names = [x for x in dir(minstance) if x not in default_attributes]
+         self.assertEqual(names, ['a', 'b'])
+ 
+-        class M2(M):
+-            def getdict(self):
+-                return "Not a dict!"
+-            __dict__ = property(getdict)
++        with torch._dynamo.error_on_graph_break(False):
++            class M2(M):
++                def getdict(self):
++                    return "Not a dict!"
++                __dict__ = property(getdict)
+ 
+         m2instance = M2("m2")
+         m2instance.b = 2
+@@ -2574,71 +2753,78 @@ class ClassPropertiesAndMethods(unittest.TestCase, ExtraAssertions):
+         self.assertEqual(dir(object()), dir(Ellipsis))
+ 
+         # Nasty test case for proxied objects
+-        class Wrapper(object):
+-            def __init__(self, obj):
+-                self.__obj = obj
+-            def __repr__(self):
+-                return "Wrapper(%s)" % repr(self.__obj)
+-            def __getitem__(self, key):
+-                return Wrapper(self.__obj[key])
+-            def __len__(self):
+-                return len(self.__obj)
+-            def __getattr__(self, name):
+-                return Wrapper(getattr(self.__obj, name))
++        with torch._dynamo.error_on_graph_break(False):
++            class Wrapper(object):
++                def __init__(self, obj):
++                    self.__obj = obj
++                def __repr__(self):
++                    return "Wrapper(%s)" % repr(self.__obj)
++                def __getitem__(self, key):
++                    return Wrapper(self.__obj[key])
++                def __len__(self):
++                    return len(self.__obj)
++                def __getattr__(self, name):
++                    return Wrapper(getattr(self.__obj, name))
+ 
+-        class C(object):
+-            def __getclass(self):
+-                return Wrapper(type(self))
+-            __class__ = property(__getclass)
++            class C(object):
++                def __getclass(self):
++                    return Wrapper(type(self))
++                __class__ = property(__getclass)
+ 
+         dir(C()) # This used to segfault
+ 
+     def test_supers(self):
+         # Testing super...
+ 
+-        class A(object):
+-            def meth(self, a):
+-                return "A(%r)" % a
++        with torch._dynamo.error_on_graph_break(False):
++            class A(object):
++                def meth(self, a):
++                    return "A(%r)" % a
+ 
+         self.assertEqual(A().meth(1), "A(1)")
+ 
+-        class B(A):
+-            def __init__(self):
+-                self.__super = super(B, self)
+-            def meth(self, a):
+-                return "B(%r)" % a + self.__super.meth(a)
++        with torch._dynamo.error_on_graph_break(False):
++            class B(A):
++                def __init__(self):
++                    self.__super = super(B, self)
++                def meth(self, a):
++                    return "B(%r)" % a + self.__super.meth(a)
+ 
+         self.assertEqual(B().meth(2), "B(2)A(2)")
+ 
+-        class C(A):
+-            def meth(self, a):
+-                return "C(%r)" % a + self.__super.meth(a)
++        with torch._dynamo.error_on_graph_break(False):
++            class C(A):
++                def meth(self, a):
++                    return "C(%r)" % a + self.__super.meth(a)
+         C._C__super = super(C)
+ 
+         self.assertEqual(C().meth(3), "C(3)A(3)")
+ 
+-        class D(C, B):
+-            def meth(self, a):
+-                return "D(%r)" % a + super(D, self).meth(a)
++        with torch._dynamo.error_on_graph_break(False):
++            class D(C, B):
++                def meth(self, a):
++                    return "D(%r)" % a + super(D, self).meth(a)
+ 
+         self.assertEqual(D().meth(4), "D(4)C(4)B(4)A(4)")
+ 
+         # Test for subclassing super
+ 
+-        class mysuper(super):
+-            def __init__(self, *args):
+-                return super(mysuper, self).__init__(*args)
++        with torch._dynamo.error_on_graph_break(False):
++            class mysuper(super):
++                def __init__(self, *args):
++                    return super(mysuper, self).__init__(*args)
+ 
+-        class E(D):
+-            def meth(self, a):
+-                return "E(%r)" % a + mysuper(E, self).meth(a)
++            class E(D):
++                def meth(self, a):
++                    return "E(%r)" % a + mysuper(E, self).meth(a)
+ 
+         self.assertEqual(E().meth(5), "E(5)D(5)C(5)B(5)A(5)")
+ 
+-        class F(E):
+-            def meth(self, a):
+-                s = self.__super # == mysuper(F, self)
+-                return "F(%r)[%s]" % (a, s.__class__.__name__) + s.meth(a)
++        with torch._dynamo.error_on_graph_break(False):
++            class F(E):
++                def meth(self, a):
++                    s = self.__super # == mysuper(F, self)
++                    return "F(%r)[%s]" % (a, s.__class__.__name__) + s.meth(a)
+         F._F__super = mysuper(F)
+ 
+         self.assertEqual(F().meth(6), "F(6)[mysuper]E(6)D(6)C(6)B(6)A(6)")
+@@ -2676,13 +2862,14 @@ class ClassPropertiesAndMethods(unittest.TestCase, ExtraAssertions):
+         # Make sure data descriptors can be overridden and accessed via super
+         # (new feature in Python 2.3)
+ 
+-        class DDbase(object):
+-            def getx(self): return 42
+-            x = property(getx)
++        with torch._dynamo.error_on_graph_break(False):
++            class DDbase(object):
++                def getx(self): return 42
++                x = property(getx)
+ 
+-        class DDsub(DDbase):
+-            def getx(self): return "hello"
+-            x = property(getx)
++            class DDsub(DDbase):
++                def getx(self): return "hello"
++                x = property(getx)
+ 
+         dd = DDsub()
+         self.assertEqual(dd.x, "hello")
+@@ -2691,13 +2878,14 @@ class ClassPropertiesAndMethods(unittest.TestCase, ExtraAssertions):
+         # Ensure that super() lookup of descriptor from classmethod
+         # works (SF ID# 743627)
+ 
+-        class Base(object):
+-            aProp = property(lambda self: "foo")
++        with torch._dynamo.error_on_graph_break(False):
++            class Base(object):
++                aProp = property(lambda self: "foo")
+ 
+-        class Sub(Base):
+-            @classmethod
+-            def test(klass):
+-                return super(Sub,klass).aProp
++            class Sub(Base):
++                @classmethod
++                def test(klass):
++                    return super(Sub,klass).aProp
+ 
+         self.assertEqual(Sub.test(), Base.aProp)
+ 
+@@ -2708,13 +2896,14 @@ class ClassPropertiesAndMethods(unittest.TestCase, ExtraAssertions):
+     def test_basic_inheritance(self):
+         # Testing inheritance from basic types...
+ 
+-        class hexint(int):
+-            def __repr__(self):
+-                return hex(self)
+-            def __add__(self, other):
+-                return hexint(int.__add__(self, other))
+-            # (Note that overriding __radd__ doesn't work,
+-            # because the int type gets first dibs.)
++        with torch._dynamo.error_on_graph_break(False):
++            class hexint(int):
++                def __repr__(self):
++                    return hex(self)
++                def __add__(self, other):
++                    return hexint(int.__add__(self, other))
++                # (Note that overriding __radd__ doesn't work,
++                # because the int type gets first dibs.)
+         self.assertEqual(repr(hexint(7) + 9), "0x10")
+         self.assertEqual(repr(hexint(1000) + 7), "0x3ef")
+         a = hexint(12345)
+@@ -2728,13 +2917,14 @@ class ClassPropertiesAndMethods(unittest.TestCase, ExtraAssertions):
+         self.assertIs((hexint(0) << 12).__class__, int)
+         self.assertIs((hexint(0) >> 12).__class__, int)
+ 
+-        class octlong(int):
+-            __slots__ = []
+-            def __str__(self):
+-                return oct(self)
+-            def __add__(self, other):
+-                return self.__class__(super(octlong, self).__add__(other))
+-            __radd__ = __add__
++        with torch._dynamo.error_on_graph_break(False):
++            class octlong(int):
++                __slots__ = []
++                def __str__(self):
++                    return oct(self)
++                def __add__(self, other):
++                    return self.__class__(super(octlong, self).__add__(other))
++                __radd__ = __add__
+         self.assertEqual(str(octlong(3) + 5), "0o10")
+         # (Note that overriding __radd__ here only seems to work
+         # because the example uses a short int left argument.)
+@@ -2763,8 +2953,9 @@ class ClassPropertiesAndMethods(unittest.TestCase, ExtraAssertions):
+ 
+         # Because octlong overrides __add__, we can't check the absence of +0
+         # optimizations using octlong.
+-        class longclone(int):
+-            pass
++        with torch._dynamo.error_on_graph_break(False):
++            class longclone(int):
++                pass
+         a = longclone(1)
+         self.assertIs((a + 0).__class__, int)
+         self.assertIs((0 + a).__class__, int)
+@@ -2774,12 +2965,13 @@ class ClassPropertiesAndMethods(unittest.TestCase, ExtraAssertions):
+         self.assertEqual(a.__dict__, {})
+         self.assertEqual(int(a), -1)  # self.assertTrue PyNumber_Long() copies the sign bit
+ 
+-        class precfloat(float):
+-            __slots__ = ['prec']
+-            def __init__(self, value=0.0, prec=12):
+-                self.prec = int(prec)
+-            def __repr__(self):
+-                return "%.*g" % (self.prec, self)
++        with torch._dynamo.error_on_graph_break(False):
++            class precfloat(float):
++                __slots__ = ['prec']
++                def __init__(self, value=0.0, prec=12):
++                    self.prec = int(prec)
++                def __repr__(self):
++                    return "%.*g" % (self.prec, self)
+         self.assertEqual(repr(precfloat(1.1)), "1.1")
+         a = precfloat(12345)
+         self.assertEqual(a, 12345.0)
+@@ -2788,9 +2980,10 @@ class ClassPropertiesAndMethods(unittest.TestCase, ExtraAssertions):
+         self.assertEqual(hash(a), hash(12345.0))
+         self.assertIs((+a).__class__, float)
+ 
+-        class madcomplex(complex):
+-            def __repr__(self):
+-                return "%.17gj%+.17g" % (self.imag, self.real)
++        with torch._dynamo.error_on_graph_break(False):
++            class madcomplex(complex):
++                def __repr__(self):
++                    return "%.17gj%+.17g" % (self.imag, self.real)
+         a = madcomplex(-3, 4)
+         self.assertEqual(repr(a), "4j-3")
+         base = complex(-3, 4)
+@@ -2814,15 +3007,16 @@ class ClassPropertiesAndMethods(unittest.TestCase, ExtraAssertions):
+         self.assertEqual((a / 1).__class__, complex)
+         self.assertEqual(a / 1, base)
+ 
+-        class madtuple(tuple):
+-            _rev = None
+-            def rev(self):
+-                if self._rev is not None:
++        with torch._dynamo.error_on_graph_break(False):
++            class madtuple(tuple):
++                _rev = None
++                def rev(self):
++                    if self._rev is not None:
++                        return self._rev
++                    L = list(self)
++                    L.reverse()
++                    self._rev = self.__class__(L)
+                     return self._rev
+-                L = list(self)
+-                L.reverse()
+-                self._rev = self.__class__(L)
+-                return self._rev
+         a = madtuple((1,2,3,4,5,6,7,8,9,0))
+         self.assertEqual(a, (1,2,3,4,5,6,7,8,9,0))
+         self.assertEqual(a.rev(), madtuple((0,9,8,7,6,5,4,3,2,1)))
+@@ -2849,15 +3043,16 @@ class ClassPropertiesAndMethods(unittest.TestCase, ExtraAssertions):
+         self.assertIs((a * 2).__class__, tuple)
+         self.assertIs(a[:].__class__, tuple)
+ 
+-        class madstring(str):
+-            _rev = None
+-            def rev(self):
+-                if self._rev is not None:
++        with torch._dynamo.error_on_graph_break(False):
++            class madstring(str):
++                _rev = None
++                def rev(self):
++                    if self._rev is not None:
++                        return self._rev
++                    L = list(self)
++                    L.reverse()
++                    self._rev = self.__class__("".join(L))
+                     return self._rev
+-                L = list(self)
+-                L.reverse()
+-                self._rev = self.__class__("".join(L))
+-                return self._rev
+         s = madstring("abcdefghijklmnopqrstuvwxyz")
+         self.assertEqual(s, "abcdefghijklmnopqrstuvwxyz")
+         self.assertEqual(s.rev(), madstring("zyxwvutsrqponmlkjihgfedcba"))
+@@ -2913,15 +3108,16 @@ class ClassPropertiesAndMethods(unittest.TestCase, ExtraAssertions):
+         self.assertIs(s.lower().__class__, str)
+         self.assertEqual(s.lower(), base)
+ 
+-        class madunicode(str):
+-            _rev = None
+-            def rev(self):
+-                if self._rev is not None:
++        with torch._dynamo.error_on_graph_break(False):
++            class madunicode(str):
++                _rev = None
++                def rev(self):
++                    if self._rev is not None:
++                        return self._rev
++                    L = list(self)
++                    L.reverse()
++                    self._rev = self.__class__("".join(L))
+                     return self._rev
+-                L = list(self)
+-                L.reverse()
+-                self._rev = self.__class__("".join(L))
+-                return self._rev
+         u = madunicode("ABCDEF")
+         self.assertEqual(u, "ABCDEF")
+         self.assertEqual(u.rev(), madunicode("FEDCBA"))
+@@ -2972,8 +3168,9 @@ class ClassPropertiesAndMethods(unittest.TestCase, ExtraAssertions):
+         self.assertIs(u[0:0].__class__, str)
+         self.assertEqual(u[0:0], "")
+ 
+-        class sublist(list):
+-            pass
++        with torch._dynamo.error_on_graph_break(False):
++            class sublist(list):
++                pass
+         a = sublist(range(5))
+         self.assertEqual(a, list(range(5)))
+         a.append("hello")
+@@ -3067,23 +3264,24 @@ class ClassPropertiesAndMethods(unittest.TestCase, ExtraAssertions):
+     def test_str_subclass_as_dict_key(self):
+         # Testing a str subclass used as dict key ..
+ 
+-        class cistr(str):
+-            """Subclass of str that computes __eq__ case-insensitively.
++        with torch._dynamo.error_on_graph_break(False):
++            class cistr(str):
++                """Subclass of str that computes __eq__ case-insensitively.
+ 
+-            Also computes a hash code of the string in canonical form.
+-            """
++                Also computes a hash code of the string in canonical form.
++                """
+ 
+-            def __init__(self, value):
+-                self.canonical = value.lower()
+-                self.hashcode = hash(self.canonical)
++                def __init__(self, value):
++                    self.canonical = value.lower()
++                    self.hashcode = hash(self.canonical)
+ 
+-            def __eq__(self, other):
+-                if not isinstance(other, cistr):
+-                    other = cistr(other)
+-                return self.canonical == other.canonical
++                def __eq__(self, other):
++                    if not isinstance(other, cistr):
++                        other = cistr(other)
++                    return self.canonical == other.canonical
+ 
+-            def __hash__(self):
+-                return self.hashcode
++                def __hash__(self):
++                    return self.hashcode
+ 
+         self.assertEqual(cistr('ABC'), 'abc')
+         self.assertEqual('aBc', cistr('ABC'))
+@@ -3098,49 +3296,51 @@ class ClassPropertiesAndMethods(unittest.TestCase, ExtraAssertions):
+ 
+     def test_classic_comparisons(self):
+         # Testing classic comparisons...
+-        class classic:
+-            pass
++        with torch._dynamo.error_on_graph_break(False):
++            class classic:
++                pass
+ 
+         for base in (classic, int, object):
+-            class C(base):
+-                def __init__(self, value):
+-                    self.value = int(value)
+-                def __eq__(self, other):
+-                    if isinstance(other, C):
+-                        return self.value == other.value
+-                    if isinstance(other, int) or isinstance(other, int):
+-                        return self.value == other
+-                    return NotImplemented
+-                def __ne__(self, other):
+-                    if isinstance(other, C):
+-                        return self.value != other.value
+-                    if isinstance(other, int) or isinstance(other, int):
+-                        return self.value != other
+-                    return NotImplemented
+-                def __lt__(self, other):
+-                    if isinstance(other, C):
+-                        return self.value < other.value
+-                    if isinstance(other, int) or isinstance(other, int):
+-                        return self.value < other
+-                    return NotImplemented
+-                def __le__(self, other):
+-                    if isinstance(other, C):
+-                        return self.value <= other.value
+-                    if isinstance(other, int) or isinstance(other, int):
+-                        return self.value <= other
+-                    return NotImplemented
+-                def __gt__(self, other):
+-                    if isinstance(other, C):
+-                        return self.value > other.value
+-                    if isinstance(other, int) or isinstance(other, int):
+-                        return self.value > other
+-                    return NotImplemented
+-                def __ge__(self, other):
+-                    if isinstance(other, C):
+-                        return self.value >= other.value
+-                    if isinstance(other, int) or isinstance(other, int):
+-                        return self.value >= other
+-                    return NotImplemented
++            with torch._dynamo.error_on_graph_break(False):
++                class C(base):
++                    def __init__(self, value):
++                        self.value = int(value)
++                    def __eq__(self, other):
++                        if isinstance(other, C):
++                            return self.value == other.value
++                        if isinstance(other, int) or isinstance(other, int):
++                            return self.value == other
++                        return NotImplemented
++                    def __ne__(self, other):
++                        if isinstance(other, C):
++                            return self.value != other.value
++                        if isinstance(other, int) or isinstance(other, int):
++                            return self.value != other
++                        return NotImplemented
++                    def __lt__(self, other):
++                        if isinstance(other, C):
++                            return self.value < other.value
++                        if isinstance(other, int) or isinstance(other, int):
++                            return self.value < other
++                        return NotImplemented
++                    def __le__(self, other):
++                        if isinstance(other, C):
++                            return self.value <= other.value
++                        if isinstance(other, int) or isinstance(other, int):
++                            return self.value <= other
++                        return NotImplemented
++                    def __gt__(self, other):
++                        if isinstance(other, C):
++                            return self.value > other.value
++                        if isinstance(other, int) or isinstance(other, int):
++                            return self.value > other
++                        return NotImplemented
++                    def __ge__(self, other):
++                        if isinstance(other, C):
++                            return self.value >= other.value
++                        if isinstance(other, int) or isinstance(other, int):
++                            return self.value >= other
++                        return NotImplemented
+ 
+             c1 = C(1)
+             c2 = C(2)
+@@ -3162,65 +3362,69 @@ class ClassPropertiesAndMethods(unittest.TestCase, ExtraAssertions):
+ 
+     def test_rich_comparisons(self):
+         # Testing rich comparisons...
+-        class Z(complex):
+-            pass
++        with torch._dynamo.error_on_graph_break(False):
++            class Z(complex):
++                pass
+         z = Z(1)
+         self.assertEqual(z, 1+0j)
+         self.assertEqual(1+0j, z)
+-        class ZZ(complex):
+-            def __eq__(self, other):
+-                try:
+-                    return abs(self - other) <= 1e-6
+-                except:
+-                    return NotImplemented
++        with torch._dynamo.error_on_graph_break(False):
++            class ZZ(complex):
++                def __eq__(self, other):
++                    try:
++                        return abs(self - other) <= 1e-6
++                    except:
++                        return NotImplemented
+         zz = ZZ(1.0000003)
+         self.assertEqual(zz, 1+0j)
+         self.assertEqual(1+0j, zz)
+ 
+-        class classic:
+-            pass
++        with torch._dynamo.error_on_graph_break(False):
++            class classic:
++                pass
+         for base in (classic, int, object, list):
+-            class C(base):
+-                def __init__(self, value):
+-                    self.value = int(value)
+-                def __cmp__(self_, other):
+-                    self.fail("shouldn't call __cmp__")
+-                def __eq__(self, other):
+-                    if isinstance(other, C):
+-                        return self.value == other.value
+-                    if isinstance(other, int) or isinstance(other, int):
+-                        return self.value == other
+-                    return NotImplemented
+-                def __ne__(self, other):
+-                    if isinstance(other, C):
+-                        return self.value != other.value
+-                    if isinstance(other, int) or isinstance(other, int):
+-                        return self.value != other
+-                    return NotImplemented
+-                def __lt__(self, other):
+-                    if isinstance(other, C):
+-                        return self.value < other.value
+-                    if isinstance(other, int) or isinstance(other, int):
+-                        return self.value < other
+-                    return NotImplemented
+-                def __le__(self, other):
+-                    if isinstance(other, C):
+-                        return self.value <= other.value
+-                    if isinstance(other, int) or isinstance(other, int):
+-                        return self.value <= other
+-                    return NotImplemented
+-                def __gt__(self, other):
+-                    if isinstance(other, C):
+-                        return self.value > other.value
+-                    if isinstance(other, int) or isinstance(other, int):
+-                        return self.value > other
+-                    return NotImplemented
+-                def __ge__(self, other):
+-                    if isinstance(other, C):
+-                        return self.value >= other.value
+-                    if isinstance(other, int) or isinstance(other, int):
+-                        return self.value >= other
+-                    return NotImplemented
++            with torch._dynamo.error_on_graph_break(False):
++                class C(base):
++                    def __init__(self, value):
++                        self.value = int(value)
++                    def __cmp__(self_, other):
++                        self.fail("shouldn't call __cmp__")
++                    def __eq__(self, other):
++                        if isinstance(other, C):
++                            return self.value == other.value
++                        if isinstance(other, int) or isinstance(other, int):
++                            return self.value == other
++                        return NotImplemented
++                    def __ne__(self, other):
++                        if isinstance(other, C):
++                            return self.value != other.value
++                        if isinstance(other, int) or isinstance(other, int):
++                            return self.value != other
++                        return NotImplemented
++                    def __lt__(self, other):
++                        if isinstance(other, C):
++                            return self.value < other.value
++                        if isinstance(other, int) or isinstance(other, int):
++                            return self.value < other
++                        return NotImplemented
++                    def __le__(self, other):
++                        if isinstance(other, C):
++                            return self.value <= other.value
++                        if isinstance(other, int) or isinstance(other, int):
++                            return self.value <= other
++                        return NotImplemented
++                    def __gt__(self, other):
++                        if isinstance(other, C):
++                            return self.value > other.value
++                        if isinstance(other, int) or isinstance(other, int):
++                            return self.value > other
++                        return NotImplemented
++                    def __ge__(self, other):
++                        if isinstance(other, C):
++                            return self.value >= other.value
++                        if isinstance(other, int) or isinstance(other, int):
++                            return self.value >= other
++                        return NotImplemented
+             c1 = C(1)
+             c2 = C(2)
+             c3 = C(3)
+@@ -3250,24 +3454,26 @@ class ClassPropertiesAndMethods(unittest.TestCase, ExtraAssertions):
+     def test_doc_descriptor(self):
+         # Testing __doc__ descriptor...
+         # SF bug 542984
+-        class DocDescr(object):
+-            def __get__(self, object, otype):
+-                if object:
+-                    object = object.__class__.__name__ + ' instance'
+-                if otype:
+-                    otype = otype.__name__
+-                return 'object=%s; type=%s' % (object, otype)
+-        class NewClass:
+-            __doc__ = DocDescr()
++        with torch._dynamo.error_on_graph_break(False):
++            class DocDescr(object):
++                def __get__(self, object, otype):
++                    if object:
++                        object = object.__class__.__name__ + ' instance'
++                    if otype:
++                        otype = otype.__name__
++                    return 'object=%s; type=%s' % (object, otype)
++            class NewClass:
++                __doc__ = DocDescr()
+         self.assertEqual(NewClass.__doc__, 'object=None; type=NewClass')
+         self.assertEqual(NewClass().__doc__, 'object=NewClass instance; type=NewClass')
+ 
+     def test_set_class(self):
+         # Testing __class__ assignment...
+-        class C(object): pass
+-        class D(object): pass
+-        class E(object): pass
+-        class F(D, E): pass
++        with torch._dynamo.error_on_graph_break(False):
++            class C(object): pass
++            class D(object): pass
++            class E(object): pass
++            class F(D, E): pass
+         for cls in C, D, E, F:
+             for cls2 in C, D, E, F:
+                 x = cls()
+@@ -3294,35 +3500,37 @@ class ClassPropertiesAndMethods(unittest.TestCase, ExtraAssertions):
+         cant(C(), object)
+         cant(object(), list)
+         cant(list(), object)
+-        class Int(int): __slots__ = []
++        with torch._dynamo.error_on_graph_break(False):
++            class Int(int): __slots__ = []
+         cant(True, int)
+         cant(2, bool)
+         o = object()
+         cant(o, int)
+         cant(o, type(None))
+         del o
+-        class G(object):
+-            __slots__ = ["a", "b"]
+-        class H(object):
+-            __slots__ = ["b", "a"]
+-        class I(object):
+-            __slots__ = ["a", "b"]
+-        class J(object):
+-            __slots__ = ["c", "b"]
+-        class K(object):
+-            __slots__ = ["a", "b", "d"]
+-        class L(H):
+-            __slots__ = ["e"]
+-        class M(I):
+-            __slots__ = ["e"]
+-        class N(J):
+-            __slots__ = ["__weakref__"]
+-        class P(J):
+-            __slots__ = ["__dict__"]
+-        class Q(J):
+-            pass
+-        class R(J):
+-            __slots__ = ["__dict__", "__weakref__"]
++        with torch._dynamo.error_on_graph_break(False):
++            class G(object):
++                __slots__ = ["a", "b"]
++            class H(object):
++                __slots__ = ["b", "a"]
++            class I(object):
++                __slots__ = ["a", "b"]
++            class J(object):
++                __slots__ = ["c", "b"]
++            class K(object):
++                __slots__ = ["a", "b", "d"]
++            class L(H):
++                __slots__ = ["e"]
++            class M(I):
++                __slots__ = ["e"]
++            class N(J):
++                __slots__ = ["__weakref__"]
++            class P(J):
++                __slots__ = ["__dict__"]
++            class Q(J):
++                pass
++            class R(J):
++                __slots__ = ["__dict__", "__weakref__"]
+ 
+         for cls, cls2 in ((G, H), (G, I), (I, H), (Q, R), (R, Q)):
+             x = cls()
+@@ -3343,17 +3551,19 @@ class ClassPropertiesAndMethods(unittest.TestCase, ExtraAssertions):
+ 
+         # Issue5283: when __class__ changes in __del__, the wrong
+         # type gets DECREF'd.
+-        class O(object):
+-            pass
+-        class A(object):
+-            def __del__(self):
+-                self.__class__ = O
++        with torch._dynamo.error_on_graph_break(False):
++            class O(object):
++                pass
++            class A(object):
++                def __del__(self):
++                    self.__class__ = O
+         l = [A() for x in range(100)]
+         del l
+ 
+     def test_set_dict(self):
+         # Testing __dict__ assignment...
+-        class C(object): pass
++        with torch._dynamo.error_on_graph_break(False):
++            class C(object): pass
+         a = C()
+         a.__dict__ = {'b': 1}
+         self.assertEqual(a.b, 1)
+@@ -3369,8 +3579,9 @@ class ClassPropertiesAndMethods(unittest.TestCase, ExtraAssertions):
+         cant(a, 1)
+         del a.__dict__ # Deleting __dict__ is allowed
+ 
+-        class Base(object):
+-            pass
++        with torch._dynamo.error_on_graph_break(False):
++            class Base(object):
++                pass
+         def verify_dict_readonly(x):
+             """
+             x has to be an instance of a class inheriting from Base.
+@@ -3391,14 +3602,15 @@ class ClassPropertiesAndMethods(unittest.TestCase, ExtraAssertions):
+                 self.fail("dict_descr allowed access to %r's dict" % x)
+ 
+         # Classes don't allow __dict__ assignment and have readonly dicts
+-        class Meta1(type, Base):
+-            pass
+-        class Meta2(Base, type):
+-            pass
+-        class D(object, metaclass=Meta1):
+-            pass
+-        class E(object, metaclass=Meta2):
+-            pass
++        with torch._dynamo.error_on_graph_break(False):
++            class Meta1(type, Base):
++                pass
++            class Meta2(Base, type):
++                pass
++            class D(object, metaclass=Meta1):
++                pass
++            class E(object, metaclass=Meta2):
++                pass
+         for cls in C, D, E:
+             verify_dict_readonly(cls)
+             class_dict = cls.__dict__
+@@ -3410,10 +3622,11 @@ class ClassPropertiesAndMethods(unittest.TestCase, ExtraAssertions):
+                 self.fail("%r's __dict__ can be modified" % cls)
+ 
+         # Modules also disallow __dict__ assignment
+-        class Module1(types.ModuleType, Base):
+-            pass
+-        class Module2(Base, types.ModuleType):
+-            pass
++        with torch._dynamo.error_on_graph_break(False):
++            class Module1(types.ModuleType, Base):
++                pass
++            class Module2(Base, types.ModuleType):
++                pass
+         for ModuleType in Module1, Module2:
+             mod = ModuleType("spam")
+             verify_dict_readonly(mod)
+@@ -3430,10 +3643,11 @@ class ClassPropertiesAndMethods(unittest.TestCase, ExtraAssertions):
+                 return False
+             else:
+                 return True
+-        class Exception1(Exception, Base):
+-            pass
+-        class Exception2(Base, Exception):
+-            pass
++        with torch._dynamo.error_on_graph_break(False):
++            class Exception1(Exception, Base):
++                pass
++            class Exception2(Base, Exception):
++                pass
+         for ExceptionType in Exception, Exception1, Exception2:
+             e = ExceptionType()
+             e.__dict__ = {"a": 1}
+@@ -3442,22 +3656,23 @@ class ClassPropertiesAndMethods(unittest.TestCase, ExtraAssertions):
+ 
+     def test_binary_operator_override(self):
+         # Testing overrides of binary operations...
+-        class I(int):
+-            def __repr__(self):
+-                return "I(%r)" % int(self)
+-            def __add__(self, other):
+-                return I(int(self) + int(other))
+-            __radd__ = __add__
+-            def __pow__(self, other, mod=None):
+-                if mod is None:
+-                    return I(pow(int(self), int(other)))
+-                else:
+-                    return I(pow(int(self), int(other), int(mod)))
+-            def __rpow__(self, other, mod=None):
+-                if mod is None:
+-                    return I(pow(int(other), int(self), mod))
+-                else:
+-                    return I(pow(int(other), int(self), int(mod)))
++        with torch._dynamo.error_on_graph_break(False):
++            class I(int):
++                def __repr__(self):
++                    return "I(%r)" % int(self)
++                def __add__(self, other):
++                    return I(int(self) + int(other))
++                __radd__ = __add__
++                def __pow__(self, other, mod=None):
++                    if mod is None:
++                        return I(pow(int(self), int(other)))
++                    else:
++                        return I(pow(int(self), int(other), int(mod)))
++                def __rpow__(self, other, mod=None):
++                    if mod is None:
++                        return I(pow(int(other), int(self), mod))
++                    else:
++                        return I(pow(int(other), int(self), int(mod)))
+ 
+         self.assertEqual(repr(I(1) + I(2)), "I(3)")
+         self.assertEqual(repr(I(1) + 2), "I(3)")
+@@ -3466,20 +3681,22 @@ class ClassPropertiesAndMethods(unittest.TestCase, ExtraAssertions):
+         self.assertEqual(repr(2 ** I(3)), "I(8)")
+         self.assertEqual(repr(I(2) ** 3), "I(8)")
+         self.assertEqual(repr(pow(I(2), I(3), I(5))), "I(3)")
+-        class S(str):
+-            def __eq__(self, other):
+-                return self.lower() == other.lower()
++        with torch._dynamo.error_on_graph_break(False):
++            class S(str):
++                def __eq__(self, other):
++                    return self.lower() == other.lower()
+ 
+     def test_subclass_propagation(self):
+         # Testing propagation of slot functions to subclasses...
+-        class A(object):
+-            pass
+-        class B(A):
+-            pass
+-        class C(A):
+-            pass
+-        class D(B, C):
+-            pass
++        with torch._dynamo.error_on_graph_break(False):
++            class A(object):
++                pass
++            class B(A):
++                pass
++            class C(A):
++                pass
++            class D(B, C):
++                pass
+         d = D()
+         orig_hash = hash(d) # related to id(d) in platform-dependent ways
+         A.__hash__ = lambda self: 42
+@@ -3539,10 +3756,11 @@ class ClassPropertiesAndMethods(unittest.TestCase, ExtraAssertions):
+             self.fail("d.foo should be undefined now")
+ 
+         # Test a nasty bug in recurse_down_subclasses()
+-        class A(object):
+-            pass
+-        class B(A):
+-            pass
++        with torch._dynamo.error_on_graph_break(False):
++            class A(object):
++                pass
++            class B(A):
++                pass
+         del B
+         support.gc_collect()
+         A.__setitem__ = lambda *a: None # crash
+@@ -3553,16 +3771,18 @@ class ClassPropertiesAndMethods(unittest.TestCase, ExtraAssertions):
+         import binascii
+         # SF bug [#470040] ParseTuple t# vs subclasses.
+ 
+-        class MyBytes(bytes):
+-            pass
++        with torch._dynamo.error_on_graph_break(False):
++            class MyBytes(bytes):
++                pass
+         base = b'abc'
+         m = MyBytes(base)
+         # b2a_hex uses the buffer interface to get its argument's value, via
+         # PyArg_ParseTuple 't#' code.
+         self.assertEqual(binascii.b2a_hex(m), binascii.b2a_hex(base))
+ 
+-        class MyInt(int):
+-            pass
++        with torch._dynamo.error_on_graph_break(False):
++            class MyInt(int):
++                pass
+         m = MyInt(42)
+         try:
+             binascii.b2a_hex(m)
+@@ -3574,11 +3794,12 @@ class ClassPropertiesAndMethods(unittest.TestCase, ExtraAssertions):
+         # Testing __str__ defined in subclass of str ...
+         import binascii
+ 
+-        class octetstring(str):
+-            def __str__(self):
+-                return binascii.b2a_hex(self.encode('ascii')).decode("ascii")
+-            def __repr__(self):
+-                return self + " repr"
++        with torch._dynamo.error_on_graph_break(False):
++            class octetstring(str):
++                def __str__(self):
++                    return binascii.b2a_hex(self.encode('ascii')).decode("ascii")
++                def __repr__(self):
++                    return self + " repr"
+ 
+         o = octetstring('A')
+         self.assertEqual(type(o), octetstring)
+@@ -3592,10 +3813,11 @@ class ClassPropertiesAndMethods(unittest.TestCase, ExtraAssertions):
+ 
+     def test_repr_with_module_str_subclass(self):
+         # gh-98783
+-        class StrSub(str):
+-            pass
+-        class Some:
+-            pass
++        with torch._dynamo.error_on_graph_break(False):
++            class StrSub(str):
++                pass
++            class Some:
++                pass
+         Some.__module__ = StrSub('example')
+         self.assertIsInstance(repr(Some), str)  # should not crash
+         self.assertIsInstance(repr(Some()), str)  # should not crash
+@@ -3611,8 +3833,9 @@ class ClassPropertiesAndMethods(unittest.TestCase, ExtraAssertions):
+ 
+     def test_recursive_call(self):
+         # Testing recursive __call__() by setting to instance of class...
+-        class A(object):
+-            pass
++        with torch._dynamo.error_on_graph_break(False):
++            class A(object):
++                pass
+ 
+         A.__call__ = A()
+         with self.assertRaises(RecursionError):
+@@ -3621,16 +3844,18 @@ class ClassPropertiesAndMethods(unittest.TestCase, ExtraAssertions):
+     def test_delete_hook(self):
+         # Testing __del__ hook...
+         log = []
+-        class C(object):
+-            def __del__(self):
+-                log.append(1)
++        with torch._dynamo.error_on_graph_break(False):
++            class C(object):
++                def __del__(self):
++                    log.append(1)
+         c = C()
+         self.assertEqual(log, [])
+         del c
+         support.gc_collect()
+         self.assertEqual(log, [1])
+ 
+-        class D(object): pass
++        with torch._dynamo.error_on_graph_break(False):
++            class D(object): pass
+         d = D()
+         try: del d[0]
+         except TypeError: pass
+@@ -3639,8 +3864,9 @@ class ClassPropertiesAndMethods(unittest.TestCase, ExtraAssertions):
+     def test_hash_inheritance(self):
+         # Testing hash of mutable subclasses...
+ 
+-        class mydict(dict):
+-            pass
++        with torch._dynamo.error_on_graph_break(False):
++            class mydict(dict):
++                pass
+         d = mydict()
+         try:
+             hash(d)
+@@ -3649,8 +3875,9 @@ class ClassPropertiesAndMethods(unittest.TestCase, ExtraAssertions):
+         else:
+             self.fail("hash() of dict subclass should fail")
+ 
+-        class mylist(list):
+-            pass
++        with torch._dynamo.error_on_graph_break(False):
++            class mylist(list):
++                pass
+         d = mylist()
+         try:
+             hash(d)
+@@ -3707,8 +3934,9 @@ class ClassPropertiesAndMethods(unittest.TestCase, ExtraAssertions):
+ 
+     def test_deepcopy_recursive(self):
+         # Testing deepcopy of recursive objects...
+-        class Node:
+-            pass
++        with torch._dynamo.error_on_graph_break(False):
++            class Node:
++                pass
+         a = Node()
+         b = Node()
+         a.b = b
+@@ -3729,14 +3957,15 @@ class ClassPropertiesAndMethods(unittest.TestCase, ExtraAssertions):
+ 
+     def test_funny_new(self):
+         # Testing __new__ returning something unexpected...
+-        class C(object):
+-            def __new__(cls, arg):
+-                if isinstance(arg, str): return [1, 2, 3]
+-                elif isinstance(arg, int): return object.__new__(D)
+-                else: return object.__new__(cls)
+-        class D(C):
+-            def __init__(self, arg):
+-                self.foo = arg
++        with torch._dynamo.error_on_graph_break(False):
++            class C(object):
++                def __new__(cls, arg):
++                    if isinstance(arg, str): return [1, 2, 3]
++                    elif isinstance(arg, int): return object.__new__(D)
++                    else: return object.__new__(cls)
++            class D(C):
++                def __init__(self, arg):
++                    self.foo = arg
+         self.assertEqual(C("1"), [1, 2, 3])
+         self.assertEqual(D("1"), [1, 2, 3])
+         d = D(None)
+@@ -3748,30 +3977,35 @@ class ClassPropertiesAndMethods(unittest.TestCase, ExtraAssertions):
+         self.assertIsInstance(d, D)
+         self.assertEqual(d.foo, 1)
+ 
+-        class C(object):
+-            @staticmethod
+-            def __new__(*args):
+-                return args
++        with torch._dynamo.error_on_graph_break(False):
++            class C(object):
++                @staticmethod
++                def __new__(*args):
++                    return args
+         self.assertEqual(C(1, 2), (C, 1, 2))
+-        class D(C):
+-            pass
++        with torch._dynamo.error_on_graph_break(False):
++            class D(C):
++                pass
+         self.assertEqual(D(1, 2), (D, 1, 2))
+ 
+-        class C(object):
+-            @classmethod
+-            def __new__(*args):
+-                return args
++        with torch._dynamo.error_on_graph_break(False):
++            class C(object):
++                @classmethod
++                def __new__(*args):
++                    return args
+         self.assertEqual(C(1, 2), (C, C, 1, 2))
+-        class D(C):
+-            pass
++        with torch._dynamo.error_on_graph_break(False):
++            class D(C):
++                pass
+         self.assertEqual(D(1, 2), (D, D, 1, 2))
+ 
+     def test_imul_bug(self):
+         # Testing for __imul__ problems...
+         # SF bug 544647
+-        class C(object):
+-            def __imul__(self, other):
+-                return (self, other)
++        with torch._dynamo.error_on_graph_break(False):
++            class C(object):
++                def __imul__(self, other):
++                    return (self, other)
+         x = C()
+         y = x
+         y *= 1.0
+@@ -3795,19 +4029,20 @@ class ClassPropertiesAndMethods(unittest.TestCase, ExtraAssertions):
+     def test_copy_setstate(self):
+         # Testing that copy.*copy() correctly uses __setstate__...
+         import copy
+-        class C(object):
+-            def __init__(self, foo=None):
+-                self.foo = foo
+-                self.__foo = foo
+-            def setfoo(self, foo=None):
+-                self.foo = foo
+-            def getfoo(self):
+-                return self.__foo
+-            def __getstate__(self):
+-                return [self.foo]
+-            def __setstate__(self_, lst):
+-                self.assertEqual(len(lst), 1)
+-                self_.__foo = self_.foo = lst[0]
++        with torch._dynamo.error_on_graph_break(False):
++            class C(object):
++                def __init__(self, foo=None):
++                    self.foo = foo
++                    self.__foo = foo
++                def setfoo(self, foo=None):
++                    self.foo = foo
++                def getfoo(self):
++                    return self.__foo
++                def __getstate__(self):
++                    return [self.foo]
++                def __setstate__(self_, lst):
++                    self.assertEqual(len(lst), 1)
++                    self_.__foo = self_.foo = lst[0]
+         a = C(42)
+         a.setfoo(24)
+         self.assertEqual(a.foo, 24)
+@@ -3826,9 +4061,10 @@ class ClassPropertiesAndMethods(unittest.TestCase, ExtraAssertions):
+         self.assertEqual("hello"[:4], "hell")
+         self.assertEqual("hello"[slice(4)], "hell")
+         self.assertEqual(str.__getitem__("hello", slice(4)), "hell")
+-        class S(str):
+-            def __getitem__(self, x):
+-                return str.__getitem__(self, x)
++        with torch._dynamo.error_on_graph_break(False):
++            class S(str):
++                def __getitem__(self, x):
++                    return str.__getitem__(self, x)
+         self.assertEqual(S("hello")[:4], "hell")
+         self.assertEqual(S("hello")[slice(4)], "hell")
+         self.assertEqual(S("hello").__getitem__(slice(4)), "hell")
+@@ -3836,9 +4072,10 @@ class ClassPropertiesAndMethods(unittest.TestCase, ExtraAssertions):
+         self.assertEqual((1,2,3)[:2], (1,2))
+         self.assertEqual((1,2,3)[slice(2)], (1,2))
+         self.assertEqual(tuple.__getitem__((1,2,3), slice(2)), (1,2))
+-        class T(tuple):
+-            def __getitem__(self, x):
+-                return tuple.__getitem__(self, x)
++        with torch._dynamo.error_on_graph_break(False):
++            class T(tuple):
++                def __getitem__(self, x):
++                    return tuple.__getitem__(self, x)
+         self.assertEqual(T((1,2,3))[:2], (1,2))
+         self.assertEqual(T((1,2,3))[slice(2)], (1,2))
+         self.assertEqual(T((1,2,3)).__getitem__(slice(2)), (1,2))
+@@ -3846,9 +4083,10 @@ class ClassPropertiesAndMethods(unittest.TestCase, ExtraAssertions):
+         self.assertEqual([1,2,3][:2], [1,2])
+         self.assertEqual([1,2,3][slice(2)], [1,2])
+         self.assertEqual(list.__getitem__([1,2,3], slice(2)), [1,2])
+-        class L(list):
+-            def __getitem__(self, x):
+-                return list.__getitem__(self, x)
++        with torch._dynamo.error_on_graph_break(False):
++            class L(list):
++                def __getitem__(self, x):
++                    return list.__getitem__(self, x)
+         self.assertEqual(L([1,2,3])[:2], [1,2])
+         self.assertEqual(L([1,2,3])[slice(2)], [1,2])
+         self.assertEqual(L([1,2,3]).__getitem__(slice(2)), [1,2])
+@@ -3866,14 +4104,15 @@ class ClassPropertiesAndMethods(unittest.TestCase, ExtraAssertions):
+     def test_subtype_resurrection(self):
+         # Testing resurrection of new-style instance...
+ 
+-        class C(object):
+-            container = []
++        with torch._dynamo.error_on_graph_break(False):
++            class C(object):
++                container = []
+ 
+-            def __del__(self):
+-                # resurrect the instance
+-                C.container.append(self)
++                def __del__(self):
++                    # resurrect the instance
++                    C.container.append(self)
+ 
+-        c = C()
++            c = C()
+         c.attr = 42
+ 
+         # The most interesting thing here is whether this blows up, due to
+@@ -3902,12 +4141,13 @@ class ClassPropertiesAndMethods(unittest.TestCase, ExtraAssertions):
+ 
+     def test_slots_multiple_inheritance(self):
+         # SF bug 575229, multiple inheritance w/ slots dumps core
+-        class A(object):
+-            __slots__=()
+-        class B(object):
+-            pass
+-        class C(A,B) :
+-            __slots__=()
++        with torch._dynamo.error_on_graph_break(False):
++            class A(object):
++                __slots__=()
++            class B(object):
++                pass
++            class C(A,B) :
++                __slots__=()
+         if support.check_impl_detail():
+             self.assertEqual(C.__basicsize__, B.__basicsize__)
+         self.assertHasAttr(C, '__dict__')
+@@ -3917,11 +4157,12 @@ class ClassPropertiesAndMethods(unittest.TestCase, ExtraAssertions):
+     def test_rmul(self):
+         # Testing correct invocation of __rmul__...
+         # SF patch 592646
+-        class C(object):
+-            def __mul__(self, other):
+-                return "mul"
+-            def __rmul__(self, other):
+-                return "rmul"
++        with torch._dynamo.error_on_graph_break(False):
++            class C(object):
++                def __mul__(self, other):
++                    return "mul"
++                def __rmul__(self, other):
++                    return "rmul"
+         a = C()
+         self.assertEqual(a*2, "mul")
+         self.assertEqual(a*2.2, "mul")
+@@ -3931,24 +4172,26 @@ class ClassPropertiesAndMethods(unittest.TestCase, ExtraAssertions):
+     def test_ipow(self):
+         # Testing correct invocation of __ipow__...
+         # [SF bug 620179]
+-        class C(object):
+-            def __ipow__(self, other):
+-                pass
++        with torch._dynamo.error_on_graph_break(False):
++            class C(object):
++                def __ipow__(self, other):
++                    pass
+         a = C()
+         a **= 2
+ 
+     def test_ipow_returns_not_implemented(self):
+-        class A:
+-            def __ipow__(self, other):
+-                return NotImplemented
++        with torch._dynamo.error_on_graph_break(False):
++            class A:
++                def __ipow__(self, other):
++                    return NotImplemented
+ 
+-        class B(A):
+-            def __rpow__(self, other):
+-                return 1
++            class B(A):
++                def __rpow__(self, other):
++                    return 1
+ 
+-        class C(A):
+-            def __pow__(self, other):
+-                return 2
++            class C(A):
++                def __pow__(self, other):
++                    return 2
+         a = A()
+         b = B()
+         c = C()
+@@ -3960,9 +4203,10 @@ class ClassPropertiesAndMethods(unittest.TestCase, ExtraAssertions):
+         self.assertEqual(c, 2)
+ 
+     def test_no_ipow(self):
+-        class B:
+-            def __rpow__(self, other):
+-                return 1
++        with torch._dynamo.error_on_graph_break(False):
++            class B:
++                def __rpow__(self, other):
++                    return 1
+ 
+         a = object()
+         b = B()
+@@ -3983,20 +4227,21 @@ class ClassPropertiesAndMethods(unittest.TestCase, ExtraAssertions):
+         # Testing mutable bases...
+ 
+         # stuff that should work:
+-        class C(object):
+-            pass
+-        class C2(object):
+-            def __getattribute__(self, attr):
+-                if attr == 'a':
+-                    return 2
+-                else:
+-                    return super(C2, self).__getattribute__(attr)
+-            def meth(self):
+-                return 1
+-        class D(C):
+-            pass
+-        class E(D):
+-            pass
++        with torch._dynamo.error_on_graph_break(False):
++            class C(object):
++                pass
++            class C2(object):
++                def __getattribute__(self, attr):
++                    if attr == 'a':
++                        return 2
++                    else:
++                        return super(C2, self).__getattribute__(attr)
++                def meth(self):
++                    return 1
++            class D(C):
++                pass
++            class E(D):
++                pass
+         d = D()
+         e = E()
+         D.__bases__ = (C,)
+@@ -4058,14 +4303,15 @@ class ClassPropertiesAndMethods(unittest.TestCase, ExtraAssertions):
+                     num_bases = 1
+                 self.assertEqual(len(tp.__bases__), num_bases, tp)
+ 
+-        class L(list):
+-            pass
++        with torch._dynamo.error_on_graph_break(False):
++            class L(list):
++                pass
+ 
+-        class C(object):
+-            pass
++            class C(object):
++                pass
+ 
+-        class D(C):
+-            pass
++            class D(C):
++                pass
+ 
+         try:
+             L.__bases__ = (dict,)
+@@ -4090,25 +4336,32 @@ class ClassPropertiesAndMethods(unittest.TestCase, ExtraAssertions):
+ 
+     def test_unsubclassable_types(self):
+         with self.assertRaises(TypeError):
+-            class X(type(None)):
+-                pass
++            with torch._dynamo.error_on_graph_break(False):
++                class X(type(None)):
++                    pass
+         with self.assertRaises(TypeError):
+-            class X(object, type(None)):
+-                pass
++            with torch._dynamo.error_on_graph_break(False):
++                class X(object, type(None)):
++                    pass
+         with self.assertRaises(TypeError):
+-            class X(type(None), object):
++            with torch._dynamo.error_on_graph_break(False):
++                class X(type(None), object):
++                    pass
++        with torch._dynamo.error_on_graph_break(False):
++            class O(object):
+                 pass
+-        class O(object):
+-            pass
+         with self.assertRaises(TypeError):
+-            class X(O, type(None)):
+-                pass
++            with torch._dynamo.error_on_graph_break(False):
++                class X(O, type(None)):
++                    pass
+         with self.assertRaises(TypeError):
+-            class X(type(None), O):
+-                pass
++            with torch._dynamo.error_on_graph_break(False):
++                class X(type(None), O):
++                    pass
+ 
+-        class X(object):
+-            pass
++        with torch._dynamo.error_on_graph_break(False):
++            class X(object):
++                pass
+         with self.assertRaises(TypeError):
+             X.__bases__ = type(None),
+         with self.assertRaises(TypeError):
+@@ -4121,43 +4374,44 @@ class ClassPropertiesAndMethods(unittest.TestCase, ExtraAssertions):
+             X.__bases__ = type(None), O
+ 
+     def test_mutable_bases_with_failing_mro(self):
+-        # Testing mutable bases with failing mro...
+-        class WorkOnce(type):
+-            def __new__(self, name, bases, ns):
+-                self.flag = 0
+-                return super(WorkOnce, self).__new__(WorkOnce, name, bases, ns)
+-            def mro(self):
+-                if self.flag > 0:
+-                    raise RuntimeError("bozo")
+-                else:
+-                    self.flag += 1
+-                    return type.mro(self)
++        with torch._dynamo.error_on_graph_break(False):
++            # Testing mutable bases with failing mro...
++            class WorkOnce(type):
++                def __new__(self, name, bases, ns):
++                    self.flag = 0
++                    return super(WorkOnce, self).__new__(WorkOnce, name, bases, ns)
++                def mro(self):
++                    if self.flag > 0:
++                        raise RuntimeError("bozo")
++                    else:
++                        self.flag += 1
++                        return type.mro(self)
+ 
+-        class WorkAlways(type):
+-            def mro(self):
+-                # this is here to make sure that .mro()s aren't called
+-                # with an exception set (which was possible at one point).
+-                # An error message will be printed in a debug build.
+-                # What's a good way to test for this?
+-                return type.mro(self)
++            class WorkAlways(type):
++                def mro(self):
++                    # this is here to make sure that .mro()s aren't called
++                    # with an exception set (which was possible at one point).
++                    # An error message will be printed in a debug build.
++                    # What's a good way to test for this?
++                    return type.mro(self)
+ 
+-        class C(object):
+-            pass
++            class C(object):
++                pass
+ 
+-        class C2(object):
+-            pass
++            class C2(object):
++                pass
+ 
+-        class D(C):
+-            pass
++            class D(C):
++                pass
+ 
+-        class E(D):
+-            pass
++            class E(D):
++                pass
+ 
+-        class F(D, metaclass=WorkOnce):
+-            pass
++            class F(D, metaclass=WorkOnce):
++                pass
+ 
+-        class G(D, metaclass=WorkAlways):
+-            pass
++            class G(D, metaclass=WorkAlways):
++                pass
+ 
+         # Immediate subclasses have their mro's adjusted in alphabetical
+         # order, so E's will get adjusted before adjusting F's fails.  We
+@@ -4175,21 +4429,22 @@ class ClassPropertiesAndMethods(unittest.TestCase, ExtraAssertions):
+             self.fail("exception not propagated")
+ 
+     def test_mutable_bases_catch_mro_conflict(self):
+-        # Testing mutable bases catch mro conflict...
+-        class A(object):
+-            pass
++        with torch._dynamo.error_on_graph_break(False):
++            # Testing mutable bases catch mro conflict...
++            class A(object):
++                pass
+ 
+-        class B(object):
+-            pass
++            class B(object):
++                pass
+ 
+-        class C(A, B):
+-            pass
++            class C(A, B):
++                pass
+ 
+-        class D(A, B):
+-            pass
++            class D(A, B):
++                pass
+ 
+-        class E(C, D):
+-            pass
++            class E(C, D):
++                pass
+ 
+         try:
+             C.__bases__ = (B, A)
+@@ -4200,8 +4455,9 @@ class ClassPropertiesAndMethods(unittest.TestCase, ExtraAssertions):
+ 
+     def test_mutable_names(self):
+         # Testing mutable names...
+-        class C(object):
+-            pass
++        with torch._dynamo.error_on_graph_break(False):
++            class C(object):
++                pass
+ 
+         # C.__module__ could be 'test_descr' or '__main__'
+         mod = C.__module__
+@@ -4322,65 +4578,70 @@ class ClassPropertiesAndMethods(unittest.TestCase, ExtraAssertions):
+             self.fail("shouldn't have allowed descr.__get__(None, int)")
+ 
+     def test_isinst_isclass(self):
+-        # Testing proxy isinstance() and isclass()...
+-        class Proxy(object):
+-            def __init__(self, obj):
+-                self.__obj = obj
+-            def __getattribute__(self, name):
+-                if name.startswith("_Proxy__"):
+-                    return object.__getattribute__(self, name)
+-                else:
+-                    return getattr(self.__obj, name)
+-        # Test with a classic class
+-        class C:
+-            pass
+-        a = C()
+-        pa = Proxy(a)
++        with torch._dynamo.error_on_graph_break(False):
++            # Testing proxy isinstance() and isclass()...
++            class Proxy(object):
++                def __init__(self, obj):
++                    self.__obj = obj
++                def __getattribute__(self, name):
++                    if name.startswith("_Proxy__"):
++                        return object.__getattribute__(self, name)
++                    else:
++                        return getattr(self.__obj, name)
++            # Test with a classic class
++            class C:
++                pass
++            a = C()
++            pa = Proxy(a)
+         self.assertIsInstance(a, C)  # Baseline
+         self.assertIsInstance(pa, C) # Test
+-        # Test with a classic subclass
+-        class D(C):
+-            pass
+-        a = D()
+-        pa = Proxy(a)
++        with torch._dynamo.error_on_graph_break(False):
++            # Test with a classic subclass
++            class D(C):
++                pass
++            a = D()
++            pa = Proxy(a)
+         self.assertIsInstance(a, C)  # Baseline
+         self.assertIsInstance(pa, C) # Test
+-        # Test with a new-style class
+-        class C(object):
+-            pass
+-        a = C()
+-        pa = Proxy(a)
++        with torch._dynamo.error_on_graph_break(False):
++            # Test with a new-style class
++            class C(object):
++                pass
++            a = C()
++            pa = Proxy(a)
+         self.assertIsInstance(a, C)  # Baseline
+         self.assertIsInstance(pa, C) # Test
+         # Test with a new-style subclass
+-        class D(C):
+-            pass
+-        a = D()
+-        pa = Proxy(a)
++        with torch._dynamo.error_on_graph_break(False):
++            class D(C):
++                pass
++            a = D()
++            pa = Proxy(a)
+         self.assertIsInstance(a, C)  # Baseline
+         self.assertIsInstance(pa, C) # Test
+ 
+     def test_proxy_super(self):
+         # Testing super() for a proxy object...
+-        class Proxy(object):
+-            def __init__(self, obj):
+-                self.__obj = obj
+-            def __getattribute__(self, name):
+-                if name.startswith("_Proxy__"):
+-                    return object.__getattribute__(self, name)
+-                else:
+-                    return getattr(self.__obj, name)
++        with torch._dynamo.error_on_graph_break(False):
++            class Proxy(object):
++                def __init__(self, obj):
++                    self.__obj = obj
++                def __getattribute__(self, name):
++                    if name.startswith("_Proxy__"):
++                        return object.__getattribute__(self, name)
++                    else:
++                        return getattr(self.__obj, name)
+ 
+-        class B(object):
+-            def f(self):
+-                return "B.f"
++            class B(object):
++                def f(self):
++                    return "B.f"
+ 
+-        class C(B):
+-            def f(self):
+-                return super(C, self).f() + "->C.f"
++            class C(B):
++                def f(self):
++                    return super(C, self).f() + "->C.f"
+ 
+-        obj = C()
+-        p = Proxy(obj)
++            obj = C()
++            p = Proxy(obj)
+         self.assertEqual(C.__dict__["f"](p), "B.f->C.f")
+ 
+     def test_carloverre(self):
+@@ -4399,34 +4660,36 @@ class ClassPropertiesAndMethods(unittest.TestCase, ExtraAssertions):
+             self.fail("Carlo Verre __delattr__ succeeded!")
+ 
+     def test_carloverre_multi_inherit_valid(self):
+-        class A(type):
+-            def __setattr__(cls, key, value):
+-                type.__setattr__(cls, key, value)
++        with torch._dynamo.error_on_graph_break(False):
++            class A(type):
++                def __setattr__(cls, key, value):
++                    type.__setattr__(cls, key, value)
+ 
+-        class B:
+-            pass
++            class B:
++                pass
+ 
+-        class C(B, A):
+-            pass
++            class C(B, A):
++                pass
+ 
+-        obj = C('D', (object,), {})
++            obj = C('D', (object,), {})
+         try:
+             obj.test = True
+         except TypeError:
+             self.fail("setattr through direct base types should be legal")
+ 
+     def test_carloverre_multi_inherit_invalid(self):
+-        class A(type):
+-            def __setattr__(cls, key, value):
+-                object.__setattr__(cls, key, value)  # this should fail!
++        with torch._dynamo.error_on_graph_break(False):
++            class A(type):
++                def __setattr__(cls, key, value):
++                    object.__setattr__(cls, key, value)  # this should fail!
+ 
+-        class B:
+-            pass
++            class B:
++                pass
+ 
+-        class C(B, A):
+-            pass
++            class C(B, A):
++                pass
+ 
+-        obj = C('D', (object,), {})
++            obj = C('D', (object,), {})
+         try:
+             obj.test = True
+         except TypeError:
+@@ -4649,39 +4912,42 @@ class ClassPropertiesAndMethods(unittest.TestCase, ExtraAssertions):
+         # tp->tp_as_sequence->sq_slice instead of
+         # tp->tp_as_sequence->sq_ass_slice
+ 
+-        class C(object):
+-            def __setitem__(self, idx, value):
+-                self.value = value
++        with torch._dynamo.error_on_graph_break(False):
++            class C(object):
++                def __setitem__(self, idx, value):
++                    self.value = value
+ 
+-        c = C()
++            c = C()
+         c[1:2] = 3
+         self.assertEqual(c.value, 3)
+ 
+     def test_set_and_no_get(self):
+         # See
+         # http://mail.python.org/pipermail/python-dev/2010-January/095637.html
+-        class Descr(object):
++        with torch._dynamo.error_on_graph_break(False):
++            class Descr(object):
+ 
+-            def __init__(self, name):
+-                self.name = name
++                def __init__(self, name):
++                    self.name = name
+ 
+-            def __set__(self, obj, value):
+-                obj.__dict__[self.name] = value
+-        descr = Descr("a")
++                def __set__(self, obj, value):
++                    obj.__dict__[self.name] = value
++            descr = Descr("a")
+ 
+-        class X(object):
+-            a = descr
++            class X(object):
++                a = descr
+ 
+         x = X()
+         self.assertIs(x.a, descr)
+         x.a = 42
+         self.assertEqual(x.a, 42)
++        with torch._dynamo.error_on_graph_break(False):
+ 
+-        # Also check type_getattro for correctness.
+-        class Meta(type):
+-            pass
+-        class X(metaclass=Meta):
+-            pass
++            # Also check type_getattro for correctness.
++            class Meta(type):
++                pass
++            class X(metaclass=Meta):
++                pass
+         X.a = 42
+         Meta.a = Descr("a")
+         self.assertEqual(X.a, 42)
+@@ -4689,22 +4955,23 @@ class ClassPropertiesAndMethods(unittest.TestCase, ExtraAssertions):
+     def test_getattr_hooks(self):
+         # issue 4230
+ 
+-        class Descriptor(object):
+-            counter = 0
+-            def __get__(self, obj, objtype=None):
+-                def getter(name):
+-                    self.counter += 1
+-                    raise AttributeError(name)
+-                return getter
+-
+-        descr = Descriptor()
+-        class A(object):
+-            __getattribute__ = descr
+-        class B(object):
+-            __getattr__ = descr
+-        class C(object):
+-            __getattribute__ = descr
+-            __getattr__ = descr
++        with torch._dynamo.error_on_graph_break(False):
++            class Descriptor(object):
++                counter = 0
++                def __get__(self, obj, objtype=None):
++                    def getter(name):
++                        self.counter += 1
++                        raise AttributeError(name)
++                    return getter
++
++            descr = Descriptor()
++            class A(object):
++                __getattribute__ = descr
++            class B(object):
++                __getattr__ = descr
++            class C(object):
++                __getattribute__ = descr
++                __getattr__ = descr
+ 
+         self.assertRaises(AttributeError, getattr, A(), "attr")
+         self.assertEqual(descr.counter, 1)
+@@ -4713,15 +4980,16 @@ class ClassPropertiesAndMethods(unittest.TestCase, ExtraAssertions):
+         self.assertRaises(AttributeError, getattr, C(), "attr")
+         self.assertEqual(descr.counter, 4)
+ 
+-        class EvilGetattribute(object):
+-            # This used to segfault
+-            def __getattr__(self, name):
+-                raise AttributeError(name)
+-            def __getattribute__(self, name):
+-                del EvilGetattribute.__getattr__
+-                for i in range(5):
+-                    gc.collect()
+-                raise AttributeError(name)
++        with torch._dynamo.error_on_graph_break(False):
++            class EvilGetattribute(object):
++                # This used to segfault
++                def __getattr__(self, name):
++                    raise AttributeError(name)
++                def __getattribute__(self, name):
++                    del EvilGetattribute.__getattr__
++                    for i in range(5):
++                        gc.collect()
++                    raise AttributeError(name)
+ 
+         self.assertRaises(AttributeError, getattr, EvilGetattribute(), "attr")
+ 
+@@ -4731,11 +4999,13 @@ class ClassPropertiesAndMethods(unittest.TestCase, ExtraAssertions):
+     def test_abstractmethods(self):
+         # type pretends not to have __abstractmethods__.
+         self.assertRaises(AttributeError, getattr, type, "__abstractmethods__")
+-        class meta(type):
+-            pass
++        with torch._dynamo.error_on_graph_break(False):
++            class meta(type):
++                pass
+         self.assertRaises(AttributeError, getattr, meta, "__abstractmethods__")
+-        class X(object):
+-            pass
++        with torch._dynamo.error_on_graph_break(False):
++            class X(object):
++                pass
+         with self.assertRaises(AttributeError):
+             del X.__abstractmethods__
+ 
+@@ -4754,8 +5024,9 @@ class ClassPropertiesAndMethods(unittest.TestCase, ExtraAssertions):
+             MyClass = meta('MyClass', (), {})
+ 
+     def test_proxy_call(self):
+-        class FakeStr:
+-            __class__ = str
++        with torch._dynamo.error_on_graph_break(False):
++            class FakeStr:
++                __class__ = str
+ 
+         fake_str = FakeStr()
+         # isinstance() reads __class__
+@@ -4807,10 +5078,11 @@ class ClassPropertiesAndMethods(unittest.TestCase, ExtraAssertions):
+         self.assertRaises(RecursionError, repr, foo)
+ 
+     def test_mixing_slot_wrappers(self):
+-        class X(dict):
+-            __setattr__ = dict.__setitem__
+-            __neg__ = dict.copy
+-        x = X()
++        with torch._dynamo.error_on_graph_break(False):
++            class X(dict):
++                __setattr__ = dict.__setitem__
++                __neg__ = dict.copy
++            x = X()
+         x.y = 42
+         self.assertEqual(x["y"], 42)
+         self.assertEqual(x, -x)
+@@ -4818,10 +5090,11 @@ class ClassPropertiesAndMethods(unittest.TestCase, ExtraAssertions):
+     def test_wrong_class_slot_wrapper(self):
+         # Check bpo-37619: a wrapper descriptor taken from the wrong class
+         # should raise an exception instead of silently being ignored
+-        class A(int):
+-            __eq__ = str.__eq__
+-            __add__ = str.__add__
+-        a = A()
++        with torch._dynamo.error_on_graph_break(False):
++            class A(int):
++                __eq__ = str.__eq__
++                __add__ = str.__add__
++            a = A()
+         with self.assertRaises(TypeError):
+             a == a
+         with self.assertRaises(TypeError):
+@@ -4836,8 +5109,9 @@ class ClassPropertiesAndMethods(unittest.TestCase, ExtraAssertions):
+         self.assertEqual("'foo' in __slots__ conflicts with class variable", m)
+ 
+     def test_set_doc(self):
+-        class X:
+-            "elephant"
++        with torch._dynamo.error_on_graph_break(False):
++            class X:
++                "elephant"
+         X.__doc__ = "banana"
+         self.assertEqual(X.__doc__, "banana")
+ 
+@@ -4867,8 +5141,9 @@ class ClassPropertiesAndMethods(unittest.TestCase, ExtraAssertions):
+         self.assertEqual(float.real.__qualname__, 'float.real')
+         self.assertEqual(int.__add__.__qualname__, 'int.__add__')
+ 
+-        class X:
+-            pass
++        with torch._dynamo.error_on_graph_break(False):
++            class X:
++                pass
+         with self.assertRaises(TypeError):
+             del X.__qualname__
+ 
+@@ -4876,9 +5151,11 @@ class ClassPropertiesAndMethods(unittest.TestCase, ExtraAssertions):
+                           str, 'Oink')
+ 
+         global Y
+-        class Y:
+-            class Inside:
+-                pass
++
++        with torch._dynamo.error_on_graph_break(False):
++            class Y:
++                class Inside:
++                    pass
+         self.assertEqual(Y.__qualname__, 'Y')
+         self.assertEqual(Y.Inside.__qualname__, 'Y.Inside')
+ 
+@@ -4909,18 +5186,23 @@ class ClassPropertiesAndMethods(unittest.TestCase, ExtraAssertions):
+ 
+     def test_object_new_and_init_with_parameters(self):
+         # See issue #1683368
+-        class OverrideNeither:
+-            pass
++        with torch._dynamo.error_on_graph_break(False):
++            class OverrideNeither:
++                pass
++
+         self.assertRaises(TypeError, OverrideNeither, 1)
+         self.assertRaises(TypeError, OverrideNeither, kw=1)
+-        class OverrideNew:
+-            def __new__(cls, foo, kw=0, *args, **kwds):
+-                return object.__new__(cls, *args, **kwds)
+-        class OverrideInit:
+-            def __init__(self, foo, kw=0, *args, **kwargs):
+-                return object.__init__(self, *args, **kwargs)
+-        class OverrideBoth(OverrideNew, OverrideInit):
+-            pass
++
++        with torch._dynamo.error_on_graph_break(False):
++            class OverrideNew:
++                def __new__(cls, foo, kw=0, *args, **kwds):
++                    return object.__new__(cls, *args, **kwds)
++            class OverrideInit:
++                def __init__(self, foo, kw=0, *args, **kwargs):
++                    return object.__init__(self, *args, **kwargs)
++            class OverrideBoth(OverrideNew, OverrideInit):
++                pass
++
+         for case in OverrideNew, OverrideInit, OverrideBoth:
+             case(1)
+             case(1, kw=2)
+@@ -4928,10 +5210,11 @@ class ClassPropertiesAndMethods(unittest.TestCase, ExtraAssertions):
+             self.assertRaises(TypeError, case, 1, 2, foo=3)
+ 
+     def test_subclassing_does_not_duplicate_dict_descriptors(self):
+-        class Base:
+-            pass
+-        class Sub(Base):
+-            pass
++        with torch._dynamo.error_on_graph_break(False):
++            class Base:
++                pass
++            class Sub(Base):
++                pass
+         self.assertIn("__dict__", Base.__dict__)
+         self.assertNotIn("__dict__", Sub.__dict__)
+ 
+@@ -5066,8 +5349,9 @@ class ClassPropertiesAndMethods(unittest.TestCase, ExtraAssertions):
+             B().foo
+ 
+ 
+-class DictProxyTests(unittest.TestCase):
++class DictProxyTests(CPythonTestCase):
+     def setUp(self):
++        super().setUp()
+         class C(object):
+             def meth(self):
+                 pass
+@@ -5110,13 +5394,14 @@ class DictProxyTests(unittest.TestCase):
+ 
+     def test_dict_type_with_metaclass(self):
+         # Testing type of __dict__ when metaclass set...
+-        class B(object):
+-            pass
+-        class M(type):
+-            pass
+-        class C(metaclass=M):
+-            # In 2.3a1, C.__dict__ was a real dict rather than a dict proxy
+-            pass
++        with torch._dynamo.error_on_graph_break(False):
++            class B(object):
++                pass
++            class M(type):
++                pass
++            class C(metaclass=M):
++                # In 2.3a1, C.__dict__ was a real dict rather than a dict proxy
++                pass
+         self.assertEqual(type(C.__dict__), type(B.__dict__))
+ 
+     def test_repr(self):
+@@ -5130,7 +5415,7 @@ class DictProxyTests(unittest.TestCase):
+             self.assertIn('{!r}: {!r}'.format(k, v), r)
+ 
+ 
+-class AAAPTypesLongInitTest(unittest.TestCase):
++class AAAPTypesLongInitTest(CPythonTestCase):
+     # This is in its own TestCase so that it can be run before any other tests.
+     # (Hence the 'AAA' in the test class name: to make it the first
+     # item in a list sorted by name, like
+@@ -5155,7 +5440,7 @@ class AAAPTypesLongInitTest(unittest.TestCase):
+         type.mro(tuple)
+ 
+ 
+-class MiscTests(unittest.TestCase):
++class MiscTests(CPythonTestCase):
+     def test_type_lookup_mro_reference(self):
+         # Issue #14199: _PyType_Lookup() has to keep a strong reference to
+         # the type MRO because it may be modified during the lookup, if
+@@ -5188,7 +5473,7 @@ class MiscTests(unittest.TestCase):
+         self.assertEqual(getattr(X, 'mykey2'), 'from Base2')
+ 
+ 
+-class PicklingTests(unittest.TestCase):
++class PicklingTests(CPythonTestCase):
+ 
+     def _check_reduce(self, proto, obj, args=(), kwargs={}, state=None,
+                       listitems=None, dictitems=None):
+@@ -5630,7 +5915,7 @@ class PicklingTests(unittest.TestCase):
+             object().__reduce_ex__(None)
+ 
+ 
+-class SharedKeyTests(unittest.TestCase):
++class SharedKeyTests(CPythonTestCase):
+ 
+     @support.cpython_only
+     def test_subclasses(self):
+@@ -5668,7 +5953,7 @@ class DebugHelperMeta(type):
+         return repr(cls.__name__)
+ 
+ 
+-class MroTest(unittest.TestCase):
++class MroTest(CPythonTestCase):
+     """
+     Regressions for some bugs revealed through
+     mcsl.mro() customization (typeobject.c: mro_internal()) and
+@@ -5676,6 +5961,7 @@ class MroTest(unittest.TestCase):
+     """
+ 
+     def setUp(self):
++        super().setUp()
+         self.step = 0
+         self.ready = False
+ 
+@@ -5898,4 +6184,4 @@ class MroTest(unittest.TestCase):
+ 
+ 
+ if __name__ == "__main__":
+-    unittest.main()
++    run_tests()
diff --git a/test/dynamo/cpython/3_13/test_descr.py b/test/dynamo/cpython/3_13/test_descr.py
new file mode 100644
index 0000000000000..938b09f0a0f2c
--- /dev/null
+++ b/test/dynamo/cpython/3_13/test_descr.py
@@ -0,0 +1,6187 @@
+# ======= BEGIN Dynamo patch =======
+# Owner(s): ["module: dynamo"]
+
+# ruff: noqa
+# flake8: noqa
+
+# Test copied from
+# https://raw.githubusercontent.com/python/cpython/refs/tags/v3.13.5/Lib/test/test_descr.py
+
+import sys
+import torch
+import torch._dynamo.test_case
+import unittest
+from torch._dynamo.test_case import CPythonTestCase
+from torch.testing._internal.common_utils import run_tests
+
+# ======= END DYNAMO PATCH =======
+
+import builtins
+import copyreg
+import gc
+import itertools
+import math
+import pickle
+import random
+import string
+import sys
+import types
+import unittest
+import warnings
+import weakref
+
+from copy import deepcopy
+from contextlib import redirect_stdout
+from test import support
+from test.support.testcase import ExtraAssertions
+
+try:
+    import _testcapi
+except ImportError:
+    _testcapi = None
+
+try:
+    import xxsubtype
+except ImportError:
+    xxsubtype = None
+
+
+class OperatorsTest(CPythonTestCase):
+
+    def __init__(self, *args, **kwargs):
+        CPythonTestCase.__init__(self, *args, **kwargs)
+        self.binops = {
+            'add': '+',
+            'sub': '-',
+            'mul': '*',
+            'matmul': '@',
+            'truediv': '/',
+            'floordiv': '//',
+            'divmod': 'divmod',
+            'pow': '**',
+            'lshift': '<<',
+            'rshift': '>>',
+            'and': '&',
+            'xor': '^',
+            'or': '|',
+            'cmp': 'cmp',
+            'lt': '<',
+            'le': '<=',
+            'eq': '==',
+            'ne': '!=',
+            'gt': '>',
+            'ge': '>=',
+        }
+
+        for name, expr in list(self.binops.items()):
+            if expr.islower():
+                expr = expr + "(a, b)"
+            else:
+                expr = 'a %s b' % expr
+            self.binops[name] = expr
+
+        self.unops = {
+            'pos': '+',
+            'neg': '-',
+            'abs': 'abs',
+            'invert': '~',
+            'int': 'int',
+            'float': 'float',
+        }
+
+        for name, expr in list(self.unops.items()):
+            if expr.islower():
+                expr = expr + "(a)"
+            else:
+                expr = '%s a' % expr
+            self.unops[name] = expr
+
+    def unop_test(self, a, res, expr="len(a)", meth="__len__"):
+        d = {'a': a}
+        self.assertEqual(eval(expr, d), res)
+        t = type(a)
+        m = getattr(t, meth)
+
+        # Find method in parent class
+        while meth not in t.__dict__:
+            t = t.__bases__[0]
+        # in some implementations (e.g. PyPy), 'm' can be a regular unbound
+        # method object; the getattr() below obtains its underlying function.
+        self.assertEqual(getattr(m, 'im_func', m), t.__dict__[meth])
+        self.assertEqual(m(a), res)
+        bm = getattr(a, meth)
+        self.assertEqual(bm(), res)
+
+    def binop_test(self, a, b, res, expr="a+b", meth="__add__"):
+        d = {'a': a, 'b': b}
+
+        self.assertEqual(eval(expr, d), res)
+        t = type(a)
+        m = getattr(t, meth)
+        while meth not in t.__dict__:
+            t = t.__bases__[0]
+        # in some implementations (e.g. PyPy), 'm' can be a regular unbound
+        # method object; the getattr() below obtains its underlying function.
+        self.assertEqual(getattr(m, 'im_func', m), t.__dict__[meth])
+        self.assertEqual(m(a, b), res)
+        bm = getattr(a, meth)
+        self.assertEqual(bm(b), res)
+
+    def sliceop_test(self, a, b, c, res, expr="a[b:c]", meth="__getitem__"):
+        d = {'a': a, 'b': b, 'c': c}
+        self.assertEqual(eval(expr, d), res)
+        t = type(a)
+        m = getattr(t, meth)
+        while meth not in t.__dict__:
+            t = t.__bases__[0]
+        # in some implementations (e.g. PyPy), 'm' can be a regular unbound
+        # method object; the getattr() below obtains its underlying function.
+        self.assertEqual(getattr(m, 'im_func', m), t.__dict__[meth])
+        self.assertEqual(m(a, slice(b, c)), res)
+        bm = getattr(a, meth)
+        self.assertEqual(bm(slice(b, c)), res)
+
+    def setop_test(self, a, b, res, stmt="a+=b", meth="__iadd__"):
+        d = {'a': deepcopy(a), 'b': b}
+        exec(stmt, d)
+        self.assertEqual(d['a'], res)
+        t = type(a)
+        m = getattr(t, meth)
+        while meth not in t.__dict__:
+            t = t.__bases__[0]
+        # in some implementations (e.g. PyPy), 'm' can be a regular unbound
+        # method object; the getattr() below obtains its underlying function.
+        self.assertEqual(getattr(m, 'im_func', m), t.__dict__[meth])
+        d['a'] = deepcopy(a)
+        m(d['a'], b)
+        self.assertEqual(d['a'], res)
+        d['a'] = deepcopy(a)
+        bm = getattr(d['a'], meth)
+        bm(b)
+        self.assertEqual(d['a'], res)
+
+    def set2op_test(self, a, b, c, res, stmt="a[b]=c", meth="__setitem__"):
+        d = {'a': deepcopy(a), 'b': b, 'c': c}
+        exec(stmt, d)
+        self.assertEqual(d['a'], res)
+        t = type(a)
+        m = getattr(t, meth)
+        while meth not in t.__dict__:
+            t = t.__bases__[0]
+        # in some implementations (e.g. PyPy), 'm' can be a regular unbound
+        # method object; the getattr() below obtains its underlying function.
+        self.assertEqual(getattr(m, 'im_func', m), t.__dict__[meth])
+        d['a'] = deepcopy(a)
+        m(d['a'], b, c)
+        self.assertEqual(d['a'], res)
+        d['a'] = deepcopy(a)
+        bm = getattr(d['a'], meth)
+        bm(b, c)
+        self.assertEqual(d['a'], res)
+
+    def setsliceop_test(self, a, b, c, d, res, stmt="a[b:c]=d", meth="__setitem__"):
+        dictionary = {'a': deepcopy(a), 'b': b, 'c': c, 'd': d}
+        exec(stmt, dictionary)
+        self.assertEqual(dictionary['a'], res)
+        t = type(a)
+        while meth not in t.__dict__:
+            t = t.__bases__[0]
+        m = getattr(t, meth)
+        # in some implementations (e.g. PyPy), 'm' can be a regular unbound
+        # method object; the getattr() below obtains its underlying function.
+        self.assertEqual(getattr(m, 'im_func', m), t.__dict__[meth])
+        dictionary['a'] = deepcopy(a)
+        m(dictionary['a'], slice(b, c), d)
+        self.assertEqual(dictionary['a'], res)
+        dictionary['a'] = deepcopy(a)
+        bm = getattr(dictionary['a'], meth)
+        bm(slice(b, c), d)
+        self.assertEqual(dictionary['a'], res)
+
+    def test_lists(self):
+        # Testing list operations...
+        # Asserts are within individual test methods
+        self.binop_test([1], [2], [1,2], "a+b", "__add__")
+        self.binop_test([1,2,3], 2, 1, "b in a", "__contains__")
+        self.binop_test([1,2,3], 4, 0, "b in a", "__contains__")
+        self.binop_test([1,2,3], 1, 2, "a[b]", "__getitem__")
+        self.sliceop_test([1,2,3], 0, 2, [1,2], "a[b:c]", "__getitem__")
+        self.setop_test([1], [2], [1,2], "a+=b", "__iadd__")
+        self.setop_test([1,2], 3, [1,2,1,2,1,2], "a*=b", "__imul__")
+        self.unop_test([1,2,3], 3, "len(a)", "__len__")
+        self.binop_test([1,2], 3, [1,2,1,2,1,2], "a*b", "__mul__")
+        self.binop_test([1,2], 3, [1,2,1,2,1,2], "b*a", "__rmul__")
+        self.set2op_test([1,2], 1, 3, [1,3], "a[b]=c", "__setitem__")
+        self.setsliceop_test([1,2,3,4], 1, 3, [5,6], [1,5,6,4], "a[b:c]=d",
+                        "__setitem__")
+
+    def test_dicts(self):
+        # Testing dict operations...
+        self.binop_test({1:2,3:4}, 1, 1, "b in a", "__contains__")
+        self.binop_test({1:2,3:4}, 2, 0, "b in a", "__contains__")
+        self.binop_test({1:2,3:4}, 1, 2, "a[b]", "__getitem__")
+
+        d = {1:2, 3:4}
+        l1 = []
+        for i in list(d.keys()):
+            l1.append(i)
+        l = []
+        for i in iter(d):
+            l.append(i)
+        self.assertEqual(l, l1)
+        l = []
+        for i in d.__iter__():
+            l.append(i)
+        self.assertEqual(l, l1)
+        l = []
+        for i in dict.__iter__(d):
+            l.append(i)
+        self.assertEqual(l, l1)
+        d = {1:2, 3:4}
+        self.unop_test(d, 2, "len(a)", "__len__")
+        self.assertEqual(eval(repr(d), {}), d)
+        self.assertEqual(eval(d.__repr__(), {}), d)
+        self.set2op_test({1:2,3:4}, 2, 3, {1:2,2:3,3:4}, "a[b]=c",
+                        "__setitem__")
+
+    # Tests for unary and binary operators
+    def number_operators(self, a, b, skip=[]):
+        dict = {'a': a, 'b': b}
+
+        for name, expr in self.binops.items():
+            if name not in skip:
+                name = "__%s__" % name
+                if hasattr(a, name):
+                    res = eval(expr, dict)
+                    self.binop_test(a, b, res, expr, name)
+
+        for name, expr in list(self.unops.items()):
+            if name not in skip:
+                name = "__%s__" % name
+                if hasattr(a, name):
+                    res = eval(expr, dict)
+                    self.unop_test(a, res, expr, name)
+
+    def test_ints(self):
+        # Testing int operations...
+        self.number_operators(100, 3)
+        # The following crashes in Python 2.2
+        self.assertEqual((1).__bool__(), 1)
+        self.assertEqual((0).__bool__(), 0)
+        # This returns 'NotImplemented' in Python 2.2
+        with torch._dynamo.error_on_graph_break(False):
+            class C(int):
+                def __add__(self, other):
+                    return NotImplemented
+        self.assertEqual(C(5), 5)
+        try:
+            C() + ""
+        except TypeError:
+            pass
+        else:
+            self.fail("NotImplemented should have caused TypeError")
+
+    def test_floats(self):
+        # Testing float operations...
+        self.number_operators(100.0, 3.0)
+
+    def test_complexes(self):
+        # Testing complex operations...
+        self.number_operators(100.0j, 3.0j, skip=['lt', 'le', 'gt', 'ge',
+                                                  'int', 'float',
+                                                  'floordiv', 'divmod', 'mod'])
+
+        with torch._dynamo.error_on_graph_break(False):
+            class Number(complex):
+                __slots__ = ['prec']
+                def __new__(cls, *args, **kwds):
+                    result = complex.__new__(cls, *args)
+                    result.prec = kwds.get('prec', 12)
+                    return result
+                def __repr__(self):
+                    prec = self.prec
+                    if self.imag == 0.0:
+                        return "%.*g" % (prec, self.real)
+                    if self.real == 0.0:
+                        return "%.*gj" % (prec, self.imag)
+                    return "(%.*g+%.*gj)" % (prec, self.real, prec, self.imag)
+                __str__ = __repr__
+
+        a = Number(3.14, prec=6)
+        self.assertEqual(repr(a), "3.14")
+        self.assertEqual(a.prec, 6)
+
+        a = Number(a, prec=2)
+        self.assertEqual(repr(a), "3.1")
+        self.assertEqual(a.prec, 2)
+
+        a = Number(234.5)
+        self.assertEqual(repr(a), "234.5")
+        self.assertEqual(a.prec, 12)
+
+    def test_explicit_reverse_methods(self):
+        # see issue 9930
+        self.assertEqual(complex.__radd__(3j, 4.0), complex(4.0, 3.0))
+        self.assertEqual(float.__rsub__(3.0, 1), -2.0)
+
+    @support.impl_detail("the module 'xxsubtype' is internal")
+    @unittest.skipIf(xxsubtype is None, "requires xxsubtype module")
+    def test_spam_lists(self):
+        # Testing spamlist operations...
+        import copy, xxsubtype as spam
+
+        def spamlist(l, memo=None):
+            import xxsubtype as spam
+            return spam.spamlist(l)
+
+        # This is an ugly hack:
+        copy._deepcopy_dispatch[spam.spamlist] = spamlist
+
+        self.binop_test(spamlist([1]), spamlist([2]), spamlist([1,2]), "a+b",
+                       "__add__")
+        self.binop_test(spamlist([1,2,3]), 2, 1, "b in a", "__contains__")
+        self.binop_test(spamlist([1,2,3]), 4, 0, "b in a", "__contains__")
+        self.binop_test(spamlist([1,2,3]), 1, 2, "a[b]", "__getitem__")
+        self.sliceop_test(spamlist([1,2,3]), 0, 2, spamlist([1,2]), "a[b:c]",
+                          "__getitem__")
+        self.setop_test(spamlist([1]), spamlist([2]), spamlist([1,2]), "a+=b",
+                        "__iadd__")
+        self.setop_test(spamlist([1,2]), 3, spamlist([1,2,1,2,1,2]), "a*=b",
+                        "__imul__")
+        self.unop_test(spamlist([1,2,3]), 3, "len(a)", "__len__")
+        self.binop_test(spamlist([1,2]), 3, spamlist([1,2,1,2,1,2]), "a*b",
+                        "__mul__")
+        self.binop_test(spamlist([1,2]), 3, spamlist([1,2,1,2,1,2]), "b*a",
+                        "__rmul__")
+        self.set2op_test(spamlist([1,2]), 1, 3, spamlist([1,3]), "a[b]=c",
+                         "__setitem__")
+        self.setsliceop_test(spamlist([1,2,3,4]), 1, 3, spamlist([5,6]),
+                             spamlist([1,5,6,4]), "a[b:c]=d", "__setitem__")
+        with torch._dynamo.error_on_graph_break(False):
+            # Test subclassing
+            class C(spam.spamlist):
+                def foo(self): return 1
+        a = C()
+        self.assertEqual(a, [])
+        self.assertEqual(a.foo(), 1)
+        a.append(100)
+        self.assertEqual(a, [100])
+        self.assertEqual(a.getstate(), 0)
+        a.setstate(42)
+        self.assertEqual(a.getstate(), 42)
+
+    @support.impl_detail("the module 'xxsubtype' is internal")
+    @unittest.skipIf(xxsubtype is None, "requires xxsubtype module")
+    def test_spam_dicts(self):
+        # Testing spamdict operations...
+        import copy, xxsubtype as spam
+        def spamdict(d, memo=None):
+            import xxsubtype as spam
+            sd = spam.spamdict()
+            for k, v in list(d.items()):
+                sd[k] = v
+            return sd
+        # This is an ugly hack:
+        copy._deepcopy_dispatch[spam.spamdict] = spamdict
+
+        self.binop_test(spamdict({1:2,3:4}), 1, 1, "b in a", "__contains__")
+        self.binop_test(spamdict({1:2,3:4}), 2, 0, "b in a", "__contains__")
+        self.binop_test(spamdict({1:2,3:4}), 1, 2, "a[b]", "__getitem__")
+        d = spamdict({1:2,3:4})
+        l1 = []
+        for i in list(d.keys()):
+            l1.append(i)
+        l = []
+        for i in iter(d):
+            l.append(i)
+        self.assertEqual(l, l1)
+        l = []
+        for i in d.__iter__():
+            l.append(i)
+        self.assertEqual(l, l1)
+        l = []
+        for i in type(spamdict({})).__iter__(d):
+            l.append(i)
+        self.assertEqual(l, l1)
+        straightd = {1:2, 3:4}
+        spamd = spamdict(straightd)
+        self.unop_test(spamd, 2, "len(a)", "__len__")
+        self.unop_test(spamd, repr(straightd), "repr(a)", "__repr__")
+        self.set2op_test(spamdict({1:2,3:4}), 2, 3, spamdict({1:2,2:3,3:4}),
+                   "a[b]=c", "__setitem__")
+        with torch._dynamo.error_on_graph_break(False):
+            # Test subclassing
+            class C(spam.spamdict):
+                def foo(self): return 1
+        a = C()
+        self.assertEqual(list(a.items()), [])
+        self.assertEqual(a.foo(), 1)
+        a['foo'] = 'bar'
+        self.assertEqual(list(a.items()), [('foo', 'bar')])
+        self.assertEqual(a.getstate(), 0)
+        a.setstate(100)
+        self.assertEqual(a.getstate(), 100)
+
+    def test_wrap_lenfunc_bad_cast(self):
+        self.assertEqual(range(sys.maxsize).__len__(), sys.maxsize)
+
+
+class ClassPropertiesAndMethods(CPythonTestCase, ExtraAssertions):
+
+    def test_python_dicts(self):
+        # Testing Python subclass of dict...
+        self.assertTrue(issubclass(dict, dict))
+        self.assertIsInstance({}, dict)
+        d = dict()
+        self.assertEqual(d, {})
+        self.assertIs(d.__class__, dict)
+        self.assertIsInstance(d, dict)
+        with torch._dynamo.error_on_graph_break(False):
+            class C(dict):
+                state = -1
+                def __init__(self_local, *a, **kw):
+                    if a:
+                        self.assertEqual(len(a), 1)
+                        self_local.state = a[0]
+                    if kw:
+                        for k, v in list(kw.items()):
+                            self_local[v] = k
+                def __getitem__(self, key):
+                    return self.get(key, 0)
+                def __setitem__(self_local, key, value):
+                    self.assertIsInstance(key, int)
+                    dict.__setitem__(self_local, key, value)
+                def setstate(self, state):
+                    self.state = state
+                def getstate(self):
+                    return self.state
+        self.assertTrue(issubclass(C, dict))
+        a1 = C(12)
+        self.assertEqual(a1.state, 12)
+        a2 = C(foo=1, bar=2)
+        self.assertEqual(a2[1] == 'foo' and a2[2], 'bar')
+        a = C()
+        self.assertEqual(a.state, -1)
+        self.assertEqual(a.getstate(), -1)
+        a.setstate(0)
+        self.assertEqual(a.state, 0)
+        self.assertEqual(a.getstate(), 0)
+        a.setstate(10)
+        self.assertEqual(a.state, 10)
+        self.assertEqual(a.getstate(), 10)
+        self.assertEqual(a[42], 0)
+        a[42] = 24
+        self.assertEqual(a[42], 24)
+        N = 50
+        for i in range(N):
+            a[i] = C()
+            for j in range(N):
+                a[i][j] = i*j
+        for i in range(N):
+            for j in range(N):
+                self.assertEqual(a[i][j], i*j)
+
+    def test_python_lists(self):
+        with torch._dynamo.error_on_graph_break(False):
+            # Testing Python subclass of list...
+            class C(list):
+                def __getitem__(self, i):
+                    if isinstance(i, slice):
+                        return i.start, i.stop
+                    return list.__getitem__(self, i) + 100
+        a = C()
+        a.extend([0,1,2])
+        self.assertEqual(a[0], 100)
+        self.assertEqual(a[1], 101)
+        self.assertEqual(a[2], 102)
+        self.assertEqual(a[100:200], (100,200))
+
+    def test_metaclass(self):
+        with torch._dynamo.error_on_graph_break(False):
+            # Testing metaclasses...
+            class C(metaclass=type):
+                def __init__(self):
+                    self.__state = 0
+                def getstate(self):
+                    return self.__state
+                def setstate(self, state):
+                    self.__state = state
+        a = C()
+        self.assertEqual(a.getstate(), 0)
+        a.setstate(10)
+        self.assertEqual(a.getstate(), 10)
+        with torch._dynamo.error_on_graph_break(False):
+            class _metaclass(type):
+                def myself(cls): return cls
+            class D(metaclass=_metaclass):
+                pass
+        self.assertEqual(D.myself(), D)
+        d = D()
+        self.assertEqual(d.__class__, D)
+        with torch._dynamo.error_on_graph_break(False):
+            class M1(type):
+                def __new__(cls, name, bases, dict):
+                    dict['__spam__'] = 1
+                    return type.__new__(cls, name, bases, dict)
+            class C(metaclass=M1):
+                pass
+        self.assertEqual(C.__spam__, 1)
+        c = C()
+        self.assertEqual(c.__spam__, 1)
+
+        with torch._dynamo.error_on_graph_break(False):
+            class _instance(object):
+                pass
+            class M2(object):
+                @staticmethod
+                def __new__(cls, name, bases, dict):
+                    self = object.__new__(cls)
+                    self.name = name
+                    self.bases = bases
+                    self.dict = dict
+                    return self
+                def __call__(self):
+                    it = _instance()
+                    # Early binding of methods
+                    for key in self.dict:
+                        if key.startswith("__"):
+                            continue
+                        setattr(it, key, self.dict[key].__get__(it, self))
+                    return it
+            class C(metaclass=M2):
+                def spam(self):
+                    return 42
+        self.assertEqual(C.name, 'C')
+        self.assertEqual(C.bases, ())
+        self.assertIn('spam', C.dict)
+        c = C()
+        self.assertEqual(c.spam(), 42)
+
+        # More metaclass examples
+
+        with torch._dynamo.error_on_graph_break(False):
+            class autosuper(type):
+                # Automatically add __super to the class
+                # This trick only works for dynamic classes
+                def __new__(metaclass, name, bases, dict):
+                    cls = super(autosuper, metaclass).__new__(metaclass,
+                                                            name, bases, dict)
+                    # Name mangling for __super removes leading underscores
+                    while name[:1] == "_":
+                        name = name[1:]
+                    if name:
+                        name = "_%s__super" % name
+                    else:
+                        name = "__super"
+                    setattr(cls, name, super(cls))
+                    return cls
+            class A(metaclass=autosuper):
+                def meth(self):
+                    return "A"
+            class B(A):
+                def meth(self):
+                    return "B" + self.__super.meth()
+            class C(A):
+                def meth(self):
+                    return "C" + self.__super.meth()
+            class D(C, B):
+                def meth(self):
+                    return "D" + self.__super.meth()
+        self.assertEqual(D().meth(), "DCBA")
+        with torch._dynamo.error_on_graph_break(False):
+            class E(B, C):
+                def meth(self):
+                    return "E" + self.__super.meth()
+        self.assertEqual(E().meth(), "EBCA")
+
+        with torch._dynamo.error_on_graph_break(False):
+            class autoproperty(type):
+                # Automatically create property attributes when methods
+                # named _get_x and/or _set_x are found
+                def __new__(metaclass, name, bases, dict):
+                    hits = {}
+                    for key, val in dict.items():
+                        if key.startswith("_get_"):
+                            key = key[5:]
+                            get, set = hits.get(key, (None, None))
+                            get = val
+                            hits[key] = get, set
+                        elif key.startswith("_set_"):
+                            key = key[5:]
+                            get, set = hits.get(key, (None, None))
+                            set = val
+                            hits[key] = get, set
+                    for key, (get, set) in hits.items():
+                        dict[key] = property(get, set)
+                    return super(autoproperty, metaclass).__new__(metaclass,
+                                                                name, bases, dict)
+        with torch._dynamo.error_on_graph_break(False):
+            class A(metaclass=autoproperty):
+                def _get_x(self):
+                    return -self.__x
+                def _set_x(self, x):
+                    self.__x = -x
+        a = A()
+        self.assertNotHasAttr(a, "x")
+        a.x = 12
+        self.assertEqual(a.x, 12)
+        self.assertEqual(a._A__x, -12)
+
+        with torch._dynamo.error_on_graph_break(False):
+            class multimetaclass(autoproperty, autosuper):
+                # Merge of multiple cooperating metaclasses
+                pass
+            class A(metaclass=multimetaclass):
+                def _get_x(self):
+                    return "A"
+            class B(A):
+                def _get_x(self):
+                    return "B" + self.__super._get_x()
+            class C(A):
+                def _get_x(self):
+                    return "C" + self.__super._get_x()
+            class D(C, B):
+                def _get_x(self):
+                    return "D" + self.__super._get_x()
+        self.assertEqual(D().x, "DCBA")
+
+        with torch._dynamo.error_on_graph_break(False):
+            # Make sure type(x) doesn't call x.__class__.__init__
+            class T(type):
+                counter = 0
+                def __init__(self, *args):
+                    T.counter += 1
+            class C(metaclass=T):
+                pass
+        self.assertEqual(T.counter, 1)
+        a = C()
+        self.assertEqual(type(a), C)
+        self.assertEqual(T.counter, 1)
+
+        with torch._dynamo.error_on_graph_break(False):
+            class C(object): pass
+        c = C()
+        try: c()
+        except TypeError: pass
+        else: self.fail("calling object w/o call method should raise "
+                        "TypeError")
+
+        with torch._dynamo.error_on_graph_break(False):
+            # Testing code to find most derived baseclass
+            class A(type):
+                def __new__(*args, **kwargs):
+                    return type.__new__(*args, **kwargs)
+
+            class B(object):
+                pass
+
+            class C(object, metaclass=A):
+                pass
+
+            # The most derived metaclass of D is A rather than type.
+            class D(B, C):
+                pass
+        self.assertIs(A, type(D))
+
+        # issue1294232: correct metaclass calculation
+        new_calls = []  # to check the order of __new__ calls
+        with torch._dynamo.error_on_graph_break(False):
+            class AMeta(type):
+                @staticmethod
+                def __new__(mcls, name, bases, ns):
+                    new_calls.append('AMeta')
+                    return super().__new__(mcls, name, bases, ns)
+                @classmethod
+                def __prepare__(mcls, name, bases):
+                    return {}
+
+            class BMeta(AMeta):
+                @staticmethod
+                def __new__(mcls, name, bases, ns):
+                    new_calls.append('BMeta')
+                    return super().__new__(mcls, name, bases, ns)
+                @classmethod
+                def __prepare__(mcls, name, bases):
+                    ns = super().__prepare__(name, bases)
+                    ns['BMeta_was_here'] = True
+                    return ns
+
+            class A(metaclass=AMeta):
+                pass
+        self.assertEqual(['AMeta'], new_calls)
+        new_calls.clear()
+
+        with torch._dynamo.error_on_graph_break(False):
+            class B(metaclass=BMeta):
+                pass
+        # BMeta.__new__ calls AMeta.__new__ with super:
+        self.assertEqual(['BMeta', 'AMeta'], new_calls)
+        new_calls.clear()
+
+        with torch._dynamo.error_on_graph_break(False):
+            class C(A, B):
+                pass
+        # The most derived metaclass is BMeta:
+        self.assertEqual(['BMeta', 'AMeta'], new_calls)
+        new_calls.clear()
+        # BMeta.__prepare__ should've been called:
+        self.assertIn('BMeta_was_here', C.__dict__)
+
+        with torch._dynamo.error_on_graph_break(False):
+            # The order of the bases shouldn't matter:
+            class C2(B, A):
+                pass
+        self.assertEqual(['BMeta', 'AMeta'], new_calls)
+        new_calls.clear()
+        self.assertIn('BMeta_was_here', C2.__dict__)
+
+        with torch._dynamo.error_on_graph_break(False):
+            # Check correct metaclass calculation when a metaclass is declared:
+            class D(C, metaclass=type):
+                pass
+        self.assertEqual(['BMeta', 'AMeta'], new_calls)
+        new_calls.clear()
+        self.assertIn('BMeta_was_here', D.__dict__)
+
+        with torch._dynamo.error_on_graph_break(False):
+            class E(C, metaclass=AMeta):
+                pass
+        self.assertEqual(['BMeta', 'AMeta'], new_calls)
+        new_calls.clear()
+        self.assertIn('BMeta_was_here', E.__dict__)
+
+        # Special case: the given metaclass isn't a class,
+        # so there is no metaclass calculation.
+        marker = object()
+        def func(*args, **kwargs):
+            return marker
+        with torch._dynamo.error_on_graph_break(False):
+            class X(metaclass=func):
+                pass
+            class Y(object, metaclass=func):
+                pass
+            class Z(D, metaclass=func):
+                pass
+        self.assertIs(marker, X)
+        self.assertIs(marker, Y)
+        self.assertIs(marker, Z)
+
+        # The given metaclass is a class,
+        # but not a descendant of type.
+        prepare_calls = []  # to track __prepare__ calls
+        with torch._dynamo.error_on_graph_break(False):
+            class ANotMeta:
+                def __new__(mcls, *args, **kwargs):
+                    new_calls.append('ANotMeta')
+                    return super().__new__(mcls)
+                @classmethod
+                def __prepare__(mcls, name, bases):
+                    prepare_calls.append('ANotMeta')
+                    return {}
+            class BNotMeta(ANotMeta):
+                def __new__(mcls, *args, **kwargs):
+                    new_calls.append('BNotMeta')
+                    return super().__new__(mcls)
+                @classmethod
+                def __prepare__(mcls, name, bases):
+                    prepare_calls.append('BNotMeta')
+                    return super().__prepare__(name, bases)
+
+            class A(metaclass=ANotMeta):
+                pass
+        self.assertIs(ANotMeta, type(A))
+        self.assertEqual(['ANotMeta'], prepare_calls)
+        prepare_calls.clear()
+        self.assertEqual(['ANotMeta'], new_calls)
+        new_calls.clear()
+
+        with torch._dynamo.error_on_graph_break(False):
+            class B(metaclass=BNotMeta):
+                pass
+        self.assertIs(BNotMeta, type(B))
+        self.assertEqual(['BNotMeta', 'ANotMeta'], prepare_calls)
+        prepare_calls.clear()
+        self.assertEqual(['BNotMeta', 'ANotMeta'], new_calls)
+        new_calls.clear()
+
+        with torch._dynamo.error_on_graph_break(False):
+            class C(A, B):
+                pass
+        self.assertIs(BNotMeta, type(C))
+        self.assertEqual(['BNotMeta', 'ANotMeta'], new_calls)
+        new_calls.clear()
+        self.assertEqual(['BNotMeta', 'ANotMeta'], prepare_calls)
+        prepare_calls.clear()
+
+        with torch._dynamo.error_on_graph_break(False):
+            class C2(B, A):
+                pass
+        self.assertIs(BNotMeta, type(C2))
+        self.assertEqual(['BNotMeta', 'ANotMeta'], new_calls)
+        new_calls.clear()
+        self.assertEqual(['BNotMeta', 'ANotMeta'], prepare_calls)
+        prepare_calls.clear()
+
+        # This is a TypeError, because of a metaclass conflict:
+        # BNotMeta is neither a subclass, nor a superclass of type
+        with torch._dynamo.error_on_graph_break(False):
+            with self.assertRaises(TypeError):
+                class D(C, metaclass=type):
+                    pass
+
+            class E(C, metaclass=ANotMeta):
+                pass
+        self.assertIs(BNotMeta, type(E))
+        self.assertEqual(['BNotMeta', 'ANotMeta'], new_calls)
+        new_calls.clear()
+        self.assertEqual(['BNotMeta', 'ANotMeta'], prepare_calls)
+        prepare_calls.clear()
+
+        with torch._dynamo.error_on_graph_break(False):
+            class F(object(), C):
+                pass
+        self.assertIs(BNotMeta, type(F))
+        self.assertEqual(['BNotMeta', 'ANotMeta'], new_calls)
+        new_calls.clear()
+        self.assertEqual(['BNotMeta', 'ANotMeta'], prepare_calls)
+        prepare_calls.clear()
+
+        with torch._dynamo.error_on_graph_break(False):
+            class F2(C, object()):
+                pass
+        self.assertIs(BNotMeta, type(F2))
+        self.assertEqual(['BNotMeta', 'ANotMeta'], new_calls)
+        new_calls.clear()
+        self.assertEqual(['BNotMeta', 'ANotMeta'], prepare_calls)
+        prepare_calls.clear()
+
+        # TypeError: BNotMeta is neither a
+        # subclass, nor a superclass of int
+        with torch._dynamo.error_on_graph_break(False):
+            with self.assertRaises(TypeError):
+                class X(C, int()):
+                    pass
+            with self.assertRaises(TypeError):
+                class X(int(), C):
+                    pass
+
+    def test_module_subclasses(self):
+        # Testing Python subclass of module...
+        log = []
+        MT = type(sys)
+        with torch._dynamo.error_on_graph_break(False):
+            class MM(MT):
+                def __init__(self, name):
+                    MT.__init__(self, name)
+                def __getattribute__(self, name):
+                    log.append(("getattr", name))
+                    return MT.__getattribute__(self, name)
+                def __setattr__(self, name, value):
+                    log.append(("setattr", name, value))
+                    MT.__setattr__(self, name, value)
+                def __delattr__(self, name):
+                    log.append(("delattr", name))
+                    MT.__delattr__(self, name)
+        a = MM("a")
+        a.foo = 12
+        x = a.foo
+        del a.foo
+        self.assertEqual(log, [("setattr", "foo", 12),
+                               ("getattr", "foo"),
+                               ("delattr", "foo")])
+
+        # https://bugs.python.org/issue1174712
+        try:
+            with torch._dynamo.error_on_graph_break(False):
+                class Module(types.ModuleType, str):
+                    pass
+        except TypeError:
+            pass
+        else:
+            self.fail("inheriting from ModuleType and str at the same time "
+                    "should fail")
+
+        # Issue 34805: Verify that definition order is retained
+        def random_name():
+            return ''.join(random.choices(string.ascii_letters, k=10))
+        with torch._dynamo.error_on_graph_break(False):
+            class A:
+                pass
+        subclasses = [type(random_name(), (A,), {}) for i in range(100)]
+        self.assertEqual(A.__subclasses__(), subclasses)
+
+    def test_multiple_inheritance(self):
+        # Testing multiple inheritance...
+        with torch._dynamo.error_on_graph_break(False):
+            class C(object):
+                def __init__(self):
+                    self.__state = 0
+                def getstate(self):
+                    return self.__state
+                def setstate(self, state):
+                    self.__state = state
+        a = C()
+        self.assertEqual(a.getstate(), 0)
+        a.setstate(10)
+        self.assertEqual(a.getstate(), 10)
+        with torch._dynamo.error_on_graph_break(False):
+            class D(dict, C):
+                def __init__(self):
+                    dict.__init__(self)
+                    C.__init__(self)
+        d = D()
+        self.assertEqual(list(d.keys()), [])
+        d["hello"] = "world"
+        self.assertEqual(list(d.items()), [("hello", "world")])
+        self.assertEqual(d["hello"], "world")
+        self.assertEqual(d.getstate(), 0)
+        d.setstate(10)
+        self.assertEqual(d.getstate(), 10)
+        self.assertEqual(D.__mro__, (D, dict, C, object))
+
+        # SF bug #442833
+        with torch._dynamo.error_on_graph_break(False):
+            class Node(object):
+                def __int__(self):
+                    return int(self.foo())
+                def foo(self):
+                    return "23"
+            class Frag(Node, list):
+                def foo(self):
+                    return "42"
+        self.assertEqual(Node().__int__(), 23)
+        self.assertEqual(int(Node()), 23)
+        self.assertEqual(Frag().__int__(), 42)
+        self.assertEqual(int(Frag()), 42)
+
+    def test_diamond_inheritance(self):
+        # Testing multiple inheritance special cases...
+        with torch._dynamo.error_on_graph_break(False):
+            class A(object):
+                def spam(self): return "A"
+        self.assertEqual(A().spam(), "A")
+        with torch._dynamo.error_on_graph_break(False):
+            class B(A):
+                def boo(self): return "B"
+                def spam(self): return "B"
+        self.assertEqual(B().spam(), "B")
+        self.assertEqual(B().boo(), "B")
+        with torch._dynamo.error_on_graph_break(False):
+            class C(A):
+                def boo(self): return "C"
+        self.assertEqual(C().spam(), "A")
+        self.assertEqual(C().boo(), "C")
+        with torch._dynamo.error_on_graph_break(False):
+            class D(B, C): pass
+        self.assertEqual(D().spam(), "B")
+        self.assertEqual(D().boo(), "B")
+        self.assertEqual(D.__mro__, (D, B, C, A, object))
+        with torch._dynamo.error_on_graph_break(False):
+            class E(C, B): pass
+        self.assertEqual(E().spam(), "B")
+        self.assertEqual(E().boo(), "C")
+        self.assertEqual(E.__mro__, (E, C, B, A, object))
+        # MRO order disagreement
+        try:
+            with torch._dynamo.error_on_graph_break(False):
+                class F(D, E): pass
+        except TypeError:
+            pass
+        else:
+            self.fail("expected MRO order disagreement (F)")
+        try:
+            with torch._dynamo.error_on_graph_break(False):
+                class G(E, D): pass
+        except TypeError:
+            pass
+        else:
+            self.fail("expected MRO order disagreement (G)")
+
+    # see thread python-dev/2002-October/029035.html
+    def test_ex5_from_c3_switch(self):
+        # Testing ex5 from C3 switch discussion...
+        with torch._dynamo.error_on_graph_break(False):
+            class A(object): pass
+            class B(object): pass
+            class C(object): pass
+            class X(A): pass
+            class Y(A): pass
+            class Z(X,B,Y,C): pass
+        self.assertEqual(Z.__mro__, (Z, X, B, Y, A, C, object))
+
+    # see "A Monotonic Superclass Linearization for Dylan",
+    # by Kim Barrett et al. (OOPSLA 1996)
+    def test_monotonicity(self):
+        # Testing MRO monotonicity...
+        with torch._dynamo.error_on_graph_break(False):
+            class Boat(object): pass
+            class DayBoat(Boat): pass
+            class WheelBoat(Boat): pass
+            class EngineLess(DayBoat): pass
+            class SmallMultihull(DayBoat): pass
+            class PedalWheelBoat(EngineLess,WheelBoat): pass
+            class SmallCatamaran(SmallMultihull): pass
+            class Pedalo(PedalWheelBoat,SmallCatamaran): pass
+
+        self.assertEqual(PedalWheelBoat.__mro__,
+              (PedalWheelBoat, EngineLess, DayBoat, WheelBoat, Boat, object))
+        self.assertEqual(SmallCatamaran.__mro__,
+              (SmallCatamaran, SmallMultihull, DayBoat, Boat, object))
+        self.assertEqual(Pedalo.__mro__,
+              (Pedalo, PedalWheelBoat, EngineLess, SmallCatamaran,
+               SmallMultihull, DayBoat, WheelBoat, Boat, object))
+
+    # see "A Monotonic Superclass Linearization for Dylan",
+    # by Kim Barrett et al. (OOPSLA 1996)
+    def test_consistency_with_epg(self):
+        # Testing consistency with EPG...
+        with torch._dynamo.error_on_graph_break(False):
+            class Pane(object): pass
+            class ScrollingMixin(object): pass
+            class EditingMixin(object): pass
+            class ScrollablePane(Pane,ScrollingMixin): pass
+            class EditablePane(Pane,EditingMixin): pass
+            class EditableScrollablePane(ScrollablePane,EditablePane): pass
+
+        self.assertEqual(EditableScrollablePane.__mro__,
+              (EditableScrollablePane, ScrollablePane, EditablePane, Pane,
+                ScrollingMixin, EditingMixin, object))
+
+    def test_mro_disagreement(self):
+        # Testing error messages for MRO disagreement...
+        mro_err_msg = ("Cannot create a consistent method resolution "
+                       "order (MRO) for bases ")
+
+        def raises(exc, expected, callable, *args):
+            try:
+                callable(*args)
+            except exc as msg:
+                # the exact msg is generally considered an impl detail
+                if support.check_impl_detail():
+                    if not str(msg).startswith(expected):
+                        self.fail("Message %r, expected %r" %
+                                  (str(msg), expected))
+            else:
+                self.fail("Expected %s" % exc)
+
+        with torch._dynamo.error_on_graph_break(False):
+            class A(object): pass
+            class B(A): pass
+            class C(object): pass
+
+        # Test some very simple errors
+        raises(TypeError, "duplicate base class A",
+               type, "X", (A, A), {})
+        raises(TypeError, mro_err_msg,
+               type, "X", (A, B), {})
+        raises(TypeError, mro_err_msg,
+               type, "X", (A, C, B), {})
+        # Test a slightly more complex error
+        with torch._dynamo.error_on_graph_break(False):
+            class GridLayout(object): pass
+            class HorizontalGrid(GridLayout): pass
+            class VerticalGrid(GridLayout): pass
+            class HVGrid(HorizontalGrid, VerticalGrid): pass
+            class VHGrid(VerticalGrid, HorizontalGrid): pass
+        raises(TypeError, mro_err_msg,
+               type, "ConfusedGrid", (HVGrid, VHGrid), {})
+
+    def test_object_class(self):
+        # Testing object class...
+        a = object()
+        self.assertEqual(a.__class__, object)
+        self.assertEqual(type(a), object)
+        b = object()
+        self.assertNotEqual(a, b)
+        self.assertNotHasAttr(a, "foo")
+        try:
+            a.foo = 12
+        except (AttributeError, TypeError):
+            pass
+        else:
+            self.fail("object() should not allow setting a foo attribute")
+        self.assertNotHasAttr(object(), "__dict__")
+
+        with torch._dynamo.error_on_graph_break(False):
+            class Cdict(object):
+                pass
+        x = Cdict()
+        self.assertEqual(x.__dict__, {})
+        x.foo = 1
+        self.assertEqual(x.foo, 1)
+        self.assertEqual(x.__dict__, {'foo': 1})
+
+    def test_object_class_assignment_between_heaptypes_and_nonheaptypes(self):
+        with torch._dynamo.error_on_graph_break(False):
+            class SubType(types.ModuleType):
+                a = 1
+
+        m = types.ModuleType("m")
+        self.assertTrue(m.__class__ is types.ModuleType)
+        self.assertFalse(hasattr(m, "a"))
+
+        m.__class__ = SubType
+        self.assertTrue(m.__class__ is SubType)
+        self.assertTrue(hasattr(m, "a"))
+
+        m.__class__ = types.ModuleType
+        self.assertTrue(m.__class__ is types.ModuleType)
+        self.assertFalse(hasattr(m, "a"))
+
+        # Make sure that builtin immutable objects don't support __class__
+        # assignment, because the object instances may be interned.
+        # We set __slots__ = () to ensure that the subclasses are
+        # memory-layout compatible, and thus otherwise reasonable candidates
+        # for __class__ assignment.
+
+        # The following types have immutable instances, but are not
+        # subclassable and thus don't need to be checked:
+        #   NoneType, bool
+
+        with torch._dynamo.error_on_graph_break(False):
+            class MyInt(int):
+                __slots__ = ()
+        with self.assertRaises(TypeError):
+            (1).__class__ = MyInt
+
+        with torch._dynamo.error_on_graph_break(False):
+            class MyFloat(float):
+                __slots__ = ()
+        with self.assertRaises(TypeError):
+            (1.0).__class__ = MyFloat
+
+        with torch._dynamo.error_on_graph_break(False):
+            class MyComplex(complex):
+                __slots__ = ()
+        with self.assertRaises(TypeError):
+            (1 + 2j).__class__ = MyComplex
+
+        with torch._dynamo.error_on_graph_break(False):
+            class MyStr(str):
+                __slots__ = ()
+        with self.assertRaises(TypeError):
+            "a".__class__ = MyStr
+
+        with torch._dynamo.error_on_graph_break(False):
+            class MyBytes(bytes):
+                __slots__ = ()
+        with self.assertRaises(TypeError):
+            b"a".__class__ = MyBytes
+
+        with torch._dynamo.error_on_graph_break(False):
+            class MyTuple(tuple):
+                __slots__ = ()
+        with self.assertRaises(TypeError):
+            ().__class__ = MyTuple
+
+        with torch._dynamo.error_on_graph_break(False):
+            class MyFrozenSet(frozenset):
+                __slots__ = ()
+        with self.assertRaises(TypeError):
+            frozenset().__class__ = MyFrozenSet
+
+    @unittest.expectedFailure
+    def test_slots(self):
+        # Testing __slots__...
+        with torch._dynamo.error_on_graph_break(False):
+            class C0(object):
+                __slots__ = []
+        x = C0()
+        self.assertNotHasAttr(x, "__dict__")
+        self.assertNotHasAttr(x, "foo")
+
+        with torch._dynamo.error_on_graph_break(False):
+            class C1(object):
+                __slots__ = ['a']
+        x = C1()
+        self.assertNotHasAttr(x, "__dict__")
+        self.assertNotHasAttr(x, "a")
+        x.a = 1
+        self.assertEqual(x.a, 1)
+        x.a = None
+        self.assertEqual(x.a, None)
+        del x.a
+        self.assertNotHasAttr(x, "a")
+
+        with torch._dynamo.error_on_graph_break(False):
+            class C3(object):
+                __slots__ = ['a', 'b', 'c']
+            x = C3()
+        self.assertNotHasAttr(x, "__dict__")
+        self.assertNotHasAttr(x, 'a')
+        self.assertNotHasAttr(x, 'b')
+        self.assertNotHasAttr(x, 'c')
+        x.a = 1
+        x.b = 2
+        x.c = 3
+        self.assertEqual(x.a, 1)
+        self.assertEqual(x.b, 2)
+        self.assertEqual(x.c, 3)
+
+        with torch._dynamo.error_on_graph_break(False):
+            class C4(object):
+                """Validate name mangling"""
+                __slots__ = ['__a']
+                def __init__(self, value):
+                    self.__a = value
+                def get(self):
+                    return self.__a
+            x = C4(5)
+        self.assertNotHasAttr(x, '__dict__')
+        self.assertNotHasAttr(x, '__a')
+        self.assertEqual(x.get(), 5)
+        try:
+            x.__a = 6
+        except AttributeError:
+            pass
+        else:
+            self.fail("Double underscored names not mangled")
+
+        # Make sure slot names are proper identifiers
+        try:
+            with torch._dynamo.error_on_graph_break(False):
+                class C(object):
+                    __slots__ = [None]
+        except TypeError:
+            pass
+        else:
+            self.fail("[None] slots not caught")
+        try:
+            with torch._dynamo.error_on_graph_break(False):
+                class C(object):
+                    __slots__ = ["foo bar"]
+        except TypeError:
+            pass
+        else:
+            self.fail("['foo bar'] slots not caught")
+        try:
+            with torch._dynamo.error_on_graph_break(False):
+                class C(object):
+                    __slots__ = ["foo\0bar"]
+        except TypeError:
+            pass
+        else:
+            self.fail("['foo\\0bar'] slots not caught")
+        try:
+            with torch._dynamo.error_on_graph_break(False):
+                class C(object):
+                    __slots__ = ["1"]
+        except TypeError:
+            pass
+        else:
+            self.fail("['1'] slots not caught")
+        try:
+            with torch._dynamo.error_on_graph_break(False):
+                class C(object):
+                    __slots__ = [""]
+        except TypeError:
+            pass
+        else:
+            self.fail("[''] slots not caught")
+
+        with torch._dynamo.error_on_graph_break(False):
+            class WithValidIdentifiers(object):
+                __slots__ = ["a", "a_b", "_a", "A0123456789Z"]
+
+        # Test a single string is not expanded as a sequence.
+        with torch._dynamo.error_on_graph_break(False):
+            class C(object):
+                __slots__ = "abc"
+        c = C()
+        c.abc = 5
+        self.assertEqual(c.abc, 5)
+
+        # Test unicode slot names
+        # Test a single unicode string is not expanded as a sequence.
+        with torch._dynamo.error_on_graph_break(False):
+            class C(object):
+                __slots__ = "abc"
+        c = C()
+        c.abc = 5
+        self.assertEqual(c.abc, 5)
+
+        # _unicode_to_string used to modify slots in certain circumstances
+        slots = ("foo", "bar")
+        with torch._dynamo.error_on_graph_break(False):
+            class C(object):
+                __slots__ = slots
+        x = C()
+        x.foo = 5
+        self.assertEqual(x.foo, 5)
+        self.assertIs(type(slots[0]), str)
+        # this used to leak references
+        try:
+            with torch._dynamo.error_on_graph_break(False):
+                class C(object):
+                    __slots__ = [chr(128)]
+        except (TypeError, UnicodeEncodeError):
+            pass
+        else:
+            self.fail("[chr(128)] slots not caught")
+
+        # Test leaks
+        with torch._dynamo.error_on_graph_break(False):
+            class Counted(object):
+                counter = 0    # counts the number of instances alive
+                def __init__(self):
+                    Counted.counter += 1
+                def __del__(self):
+                    Counted.counter -= 1
+        with torch._dynamo.error_on_graph_break(False):
+            class C(object):
+                __slots__ = ['a', 'b', 'c']
+        x = C()
+        x.a = Counted()
+        x.b = Counted()
+        x.c = Counted()
+        self.assertEqual(Counted.counter, 3)
+        del x
+        support.gc_collect()
+        self.assertEqual(Counted.counter, 0)
+        with torch._dynamo.error_on_graph_break(False):
+            class D(C):
+                pass
+        x = D()
+        x.a = Counted()
+        x.z = Counted()
+        self.assertEqual(Counted.counter, 2)
+        del x
+        support.gc_collect()
+        self.assertEqual(Counted.counter, 0)
+        with torch._dynamo.error_on_graph_break(False):
+            class E(D):
+                __slots__ = ['e']
+        x = E()
+        x.a = Counted()
+        x.z = Counted()
+        x.e = Counted()
+        self.assertEqual(Counted.counter, 3)
+        del x
+        support.gc_collect()
+        self.assertEqual(Counted.counter, 0)
+
+        # Test cyclical leaks [SF bug 519621]
+        with torch._dynamo.error_on_graph_break(False):
+            class F(object):
+                __slots__ = ['a', 'b']
+        s = F()
+        s.a = [Counted(), s]
+        self.assertEqual(Counted.counter, 1)
+        s = None
+        support.gc_collect()
+        self.assertEqual(Counted.counter, 0)
+
+        # Test lookup leaks [SF bug 572567]
+        if hasattr(gc, 'get_objects'):
+            with torch._dynamo.error_on_graph_break(False):
+                class G(object):
+                    def __eq__(self, other):
+                        return False
+            g = G()
+            orig_objects = len(gc.get_objects())
+            for i in range(10):
+                g==g
+            new_objects = len(gc.get_objects())
+            self.assertEqual(orig_objects, new_objects)
+
+        with torch._dynamo.error_on_graph_break(False):
+            class H(object):
+                __slots__ = ['a', 'b']
+                def __init__(self):
+                    self.a = 1
+                    self.b = 2
+                def __del__(self_):
+                    self.assertEqual(self_.a, 1)
+                    self.assertEqual(self_.b, 2)
+        with support.captured_output('stderr') as s:
+            h = H()
+            del h
+        self.assertEqual(s.getvalue(), '')
+
+        with torch._dynamo.error_on_graph_break(False):
+            class X(object):
+                __slots__ = "a"
+        with self.assertRaises(AttributeError):
+            del X().a
+
+        # Inherit from object on purpose to check some backwards compatibility paths
+        with torch._dynamo.error_on_graph_break(False):
+            class X(object):
+                __slots__ = "a"
+        with self.assertRaisesRegex(AttributeError, "'test.test_descr.ClassPropertiesAndMethods.test_slots.<locals>.X' object has no attribute 'a'"):
+            X().a
+
+        # Test string subclass in `__slots__`, see gh-98783
+        with torch._dynamo.error_on_graph_break(False):
+            class SubStr(str):
+                pass
+        with torch._dynamo.error_on_graph_break(False):
+            class X(object):
+                __slots__ = (SubStr('x'),)
+        X().x = 1
+        with self.assertRaisesRegex(AttributeError, "'X' object has no attribute 'a'"):
+            X().a
+
+    def test_slots_special(self):
+        # Testing __dict__ and __weakref__ in __slots__...
+        with torch._dynamo.error_on_graph_break(False):
+            class D(object):
+                __slots__ = ["__dict__"]
+        a = D()
+        self.assertHasAttr(a, "__dict__")
+        self.assertNotHasAttr(a, "__weakref__")
+        a.foo = 42
+        self.assertEqual(a.__dict__, {"foo": 42})
+
+        with torch._dynamo.error_on_graph_break(False):
+            class W(object):
+                __slots__ = ["__weakref__"]
+        a = W()
+        self.assertHasAttr(a, "__weakref__")
+        self.assertNotHasAttr(a, "__dict__")
+        try:
+            a.foo = 42
+        except AttributeError:
+            pass
+        else:
+            self.fail("shouldn't be allowed to set a.foo")
+
+        with torch._dynamo.error_on_graph_break(False):
+            class C1(W, D):
+                __slots__ = []
+        a = C1()
+        self.assertHasAttr(a, "__dict__")
+        self.assertHasAttr(a, "__weakref__")
+        a.foo = 42
+        self.assertEqual(a.__dict__, {"foo": 42})
+
+        with torch._dynamo.error_on_graph_break(False):
+            class C2(D, W):
+                __slots__ = []
+        a = C2()
+        self.assertHasAttr(a, "__dict__")
+        self.assertHasAttr(a, "__weakref__")
+        a.foo = 42
+        self.assertEqual(a.__dict__, {"foo": 42})
+
+    def test_slots_special2(self):
+        # Testing __qualname__ and __classcell__ in __slots__
+        with torch._dynamo.error_on_graph_break(False):
+            class Meta(type):
+                def __new__(cls, name, bases, namespace, attr):
+                    self.assertIn(attr, namespace)
+                    return super().__new__(cls, name, bases, namespace)
+
+        with torch._dynamo.error_on_graph_break(False):
+            class C1:
+                def __init__(self):
+                    self.b = 42
+        with torch._dynamo.error_on_graph_break(False):
+            class C2(C1, metaclass=Meta, attr="__classcell__"):
+                __slots__ = ["__classcell__"]
+                def __init__(self):
+                    super().__init__()
+        self.assertIsInstance(C2.__dict__["__classcell__"],
+                              types.MemberDescriptorType)
+        c = C2()
+        self.assertEqual(c.b, 42)
+        self.assertNotHasAttr(c, "__classcell__")
+        c.__classcell__ = 42
+        self.assertEqual(c.__classcell__, 42)
+        with self.assertRaises(TypeError):
+            with torch._dynamo.error_on_graph_break(False):
+                class C3:
+                    __classcell__ = 42
+                    __slots__ = ["__classcell__"]
+
+        with torch._dynamo.error_on_graph_break(False):
+            class Q1(metaclass=Meta, attr="__qualname__"):
+                __slots__ = ["__qualname__"]
+        self.assertEqual(Q1.__qualname__, C1.__qualname__[:-2] + "Q1")
+        self.assertIsInstance(Q1.__dict__["__qualname__"],
+                              types.MemberDescriptorType)
+        q = Q1()
+        self.assertNotHasAttr(q, "__qualname__")
+        q.__qualname__ = "q"
+        self.assertEqual(q.__qualname__, "q")
+        with self.assertRaises(TypeError):
+            with torch._dynamo.error_on_graph_break(False):
+                class Q2:
+                    __qualname__ = object()
+                    __slots__ = ["__qualname__"]
+
+    def test_slots_descriptor(self):
+        # Issue2115: slot descriptors did not correctly check
+        # the type of the given object
+        import abc
+        with torch._dynamo.error_on_graph_break(False):
+            class MyABC(metaclass=abc.ABCMeta):
+                __slots__ = "a"
+
+        with torch._dynamo.error_on_graph_break(False):
+            class Unrelated(object):
+                pass
+        MyABC.register(Unrelated)
+
+        u = Unrelated()
+        self.assertIsInstance(u, MyABC)
+
+        # This used to crash
+        self.assertRaises(TypeError, MyABC.a.__set__, u, 3)
+
+    def test_dynamics(self):
+        # Testing class attribute propagation...
+        with torch._dynamo.error_on_graph_break(False):
+            class D(object):
+                pass
+            class E(D):
+                pass
+            class F(D):
+                pass
+        D.foo = 1
+        self.assertEqual(D.foo, 1)
+        # Test that dynamic attributes are inherited
+        self.assertEqual(E.foo, 1)
+        self.assertEqual(F.foo, 1)
+        # Test dynamic instances
+        with torch._dynamo.error_on_graph_break(False):
+            class C(object):
+                pass
+        a = C()
+        self.assertNotHasAttr(a, "foobar")
+        C.foobar = 2
+        self.assertEqual(a.foobar, 2)
+        C.method = lambda self: 42
+        self.assertEqual(a.method(), 42)
+        C.__repr__ = lambda self: "C()"
+        self.assertEqual(repr(a), "C()")
+        C.__int__ = lambda self: 100
+        self.assertEqual(int(a), 100)
+        self.assertEqual(a.foobar, 2)
+        self.assertNotHasAttr(a, "spam")
+        def mygetattr(self, name):
+            if name == "spam":
+                return "spam"
+            raise AttributeError
+        C.__getattr__ = mygetattr
+        self.assertEqual(a.spam, "spam")
+        a.new = 12
+        self.assertEqual(a.new, 12)
+        def mysetattr(self, name, value):
+            if name == "spam":
+                raise AttributeError
+            return object.__setattr__(self, name, value)
+        C.__setattr__ = mysetattr
+        with self.assertRaises(AttributeError):
+            a.spam = "not spam"
+
+        self.assertEqual(a.spam, "spam")
+        with torch._dynamo.error_on_graph_break(False):
+            class D(C):
+                pass
+        d = D()
+        d.foo = 1
+        self.assertEqual(d.foo, 1)
+
+        # Test handling of int*seq and seq*int
+        with torch._dynamo.error_on_graph_break(False):
+            class I(int):
+                pass
+        self.assertEqual("a"*I(2), "aa")
+        self.assertEqual(I(2)*"a", "aa")
+        self.assertEqual(2*I(3), 6)
+        self.assertEqual(I(3)*2, 6)
+        self.assertEqual(I(3)*I(2), 6)
+
+        # Test comparison of classes with dynamic metaclasses
+        with torch._dynamo.error_on_graph_break(False):
+            class dynamicmetaclass(type):
+                pass
+            class someclass(metaclass=dynamicmetaclass):
+                pass
+        self.assertNotEqual(someclass, object)
+
+    def test_errors(self):
+        # Testing errors...
+        try:
+            with torch._dynamo.error_on_graph_break(False):
+                class C(list, dict):
+                    pass
+        except TypeError:
+            pass
+        else:
+            self.fail("inheritance from both list and dict should be illegal")
+
+        try:
+            with torch._dynamo.error_on_graph_break(False):
+                class C(object, None):
+                    pass
+        except TypeError:
+            pass
+        else:
+            self.fail("inheritance from non-type should be illegal")
+        with torch._dynamo.error_on_graph_break(False):
+            class Classic:
+                pass
+
+        try:
+            with torch._dynamo.error_on_graph_break(False):
+                class C(type(len)):
+                    pass
+        except TypeError:
+            pass
+        else:
+            self.fail("inheritance from CFunction should be illegal")
+
+        try:
+            with torch._dynamo.error_on_graph_break(False):
+                class C(object):
+                    __slots__ = 1
+        except TypeError:
+            pass
+        else:
+            self.fail("__slots__ = 1 should be illegal")
+
+        try:
+            with torch._dynamo.error_on_graph_break(False):
+                class C(object):
+                    __slots__ = [1]
+        except TypeError:
+            pass
+        else:
+            self.fail("__slots__ = [1] should be illegal")
+
+        with torch._dynamo.error_on_graph_break(False):
+            class M1(type):
+                pass
+            class M2(type):
+                pass
+            class A1(object, metaclass=M1):
+                pass
+            class A2(object, metaclass=M2):
+                pass
+        try:
+            with torch._dynamo.error_on_graph_break(False):
+                class B(A1, A2):
+                    pass
+        except TypeError:
+            pass
+        else:
+            self.fail("finding the most derived metaclass should have failed")
+
+    def test_classmethods(self):
+        # Testing class methods...
+        with torch._dynamo.error_on_graph_break(False):
+            class C(object):
+                def foo(*a): return a
+                goo = classmethod(foo)
+        c = C()
+        self.assertEqual(C.goo(1), (C, 1))
+        self.assertEqual(c.goo(1), (C, 1))
+        self.assertEqual(c.foo(1), (c, 1))
+        with torch._dynamo.error_on_graph_break(False):
+            class D(C):
+                pass
+        d = D()
+        self.assertEqual(D.goo(1), (D, 1))
+        self.assertEqual(d.goo(1), (D, 1))
+        self.assertEqual(d.foo(1), (d, 1))
+        self.assertEqual(D.foo(d, 1), (d, 1))
+        # Test for a specific crash (SF bug 528132)
+        def f(cls, arg):
+            "f docstring"
+            return (cls, arg)
+        ff = classmethod(f)
+        self.assertEqual(ff.__get__(0, int)(42), (int, 42))
+        self.assertEqual(ff.__get__(0)(42), (int, 42))
+
+        # Test super() with classmethods (SF bug 535444)
+        self.assertEqual(C.goo.__self__, C)
+        self.assertEqual(D.goo.__self__, D)
+        self.assertEqual(super(D,D).goo.__self__, D)
+        self.assertEqual(super(D,d).goo.__self__, D)
+        self.assertEqual(super(D,D).goo(), (D,))
+        self.assertEqual(super(D,d).goo(), (D,))
+
+        # Verify that a non-callable will raise
+        meth = classmethod(1).__get__(1)
+        self.assertRaises(TypeError, meth)
+
+        # Verify that classmethod() doesn't allow keyword args
+        try:
+            classmethod(f, kw=1)
+        except TypeError:
+            pass
+        else:
+            self.fail("classmethod shouldn't accept keyword args")
+
+        cm = classmethod(f)
+        cm_dict = {'__annotations__': {},
+                   '__doc__': (
+                       "f docstring"
+                       if support.HAVE_PY_DOCSTRINGS
+                       else None
+                    ),
+                   '__module__': __name__,
+                   '__name__': 'f',
+                   '__qualname__': f.__qualname__}
+        self.assertEqual(cm.__dict__, cm_dict)
+
+        cm.x = 42
+        self.assertEqual(cm.x, 42)
+        self.assertEqual(cm.__dict__, {"x" : 42, **cm_dict})
+        del cm.x
+        self.assertNotHasAttr(cm, "x")
+
+    @support.refcount_test
+    def test_refleaks_in_classmethod___init__(self):
+        gettotalrefcount = support.get_attribute(sys, 'gettotalrefcount')
+        cm = classmethod(None)
+        refs_before = gettotalrefcount()
+        for i in range(100):
+            cm.__init__(None)
+        self.assertAlmostEqual(gettotalrefcount() - refs_before, 0, delta=10)
+
+    @support.impl_detail("the module 'xxsubtype' is internal")
+    @unittest.skipIf(xxsubtype is None, "requires xxsubtype module")
+    def test_classmethods_in_c(self):
+        # Testing C-based class methods...
+        import xxsubtype as spam
+        a = (1, 2, 3)
+        d = {'abc': 123}
+        x, a1, d1 = spam.spamlist.classmeth(*a, **d)
+        self.assertEqual(x, spam.spamlist)
+        self.assertEqual(a, a1)
+        self.assertEqual(d, d1)
+        x, a1, d1 = spam.spamlist().classmeth(*a, **d)
+        self.assertEqual(x, spam.spamlist)
+        self.assertEqual(a, a1)
+        self.assertEqual(d, d1)
+        spam_cm = spam.spamlist.__dict__['classmeth']
+        x2, a2, d2 = spam_cm(spam.spamlist, *a, **d)
+        self.assertEqual(x2, spam.spamlist)
+        self.assertEqual(a2, a1)
+        self.assertEqual(d2, d1)
+        with torch._dynamo.error_on_graph_break(False):
+            class SubSpam(spam.spamlist): pass
+        x2, a2, d2 = spam_cm(SubSpam, *a, **d)
+        self.assertEqual(x2, SubSpam)
+        self.assertEqual(a2, a1)
+        self.assertEqual(d2, d1)
+
+        with self.assertRaises(TypeError) as cm:
+            spam_cm()
+        self.assertEqual(
+            str(cm.exception),
+            "descriptor 'classmeth' of 'xxsubtype.spamlist' "
+            "object needs an argument")
+
+        with self.assertRaises(TypeError) as cm:
+            spam_cm(spam.spamlist())
+        self.assertEqual(
+            str(cm.exception),
+            "descriptor 'classmeth' for type 'xxsubtype.spamlist' "
+            "needs a type, not a 'xxsubtype.spamlist' as arg 2")
+
+        with self.assertRaises(TypeError) as cm:
+            spam_cm(list)
+        expected_errmsg = (
+            "descriptor 'classmeth' requires a subtype of 'xxsubtype.spamlist' "
+            "but received 'list'")
+        self.assertEqual(str(cm.exception), expected_errmsg)
+
+        with self.assertRaises(TypeError) as cm:
+            spam_cm.__get__(None, list)
+        self.assertEqual(str(cm.exception), expected_errmsg)
+
+    def test_staticmethods(self):
+        # Testing static methods...
+        with torch._dynamo.error_on_graph_break(False):
+            class C(object):
+                def foo(*a): return a
+                goo = staticmethod(foo)
+        c = C()
+        self.assertEqual(C.goo(1), (1,))
+        self.assertEqual(c.goo(1), (1,))
+        self.assertEqual(c.foo(1), (c, 1,))
+        with torch._dynamo.error_on_graph_break(False):
+            class D(C):
+                pass
+        d = D()
+        self.assertEqual(D.goo(1), (1,))
+        self.assertEqual(d.goo(1), (1,))
+        self.assertEqual(d.foo(1), (d, 1))
+        self.assertEqual(D.foo(d, 1), (d, 1))
+        sm = staticmethod(None)
+        self.assertEqual(sm.__dict__, {'__doc__': None.__doc__})
+        sm.x = 42
+        self.assertEqual(sm.x, 42)
+        self.assertEqual(sm.__dict__, {"x" : 42, '__doc__': None.__doc__})
+        del sm.x
+        self.assertNotHasAttr(sm, "x")
+
+    @support.refcount_test
+    def test_refleaks_in_staticmethod___init__(self):
+        gettotalrefcount = support.get_attribute(sys, 'gettotalrefcount')
+        sm = staticmethod(None)
+        refs_before = gettotalrefcount()
+        for i in range(100):
+            sm.__init__(None)
+        self.assertAlmostEqual(gettotalrefcount() - refs_before, 0, delta=10)
+
+    @support.impl_detail("the module 'xxsubtype' is internal")
+    @unittest.skipIf(xxsubtype is None, "requires xxsubtype module")
+    def test_staticmethods_in_c(self):
+        # Testing C-based static methods...
+        import xxsubtype as spam
+        a = (1, 2, 3)
+        d = {"abc": 123}
+        x, a1, d1 = spam.spamlist.staticmeth(*a, **d)
+        self.assertEqual(x, None)
+        self.assertEqual(a, a1)
+        self.assertEqual(d, d1)
+        x, a1, d2 = spam.spamlist().staticmeth(*a, **d)
+        self.assertEqual(x, None)
+        self.assertEqual(a, a1)
+        self.assertEqual(d, d1)
+
+    def test_classic(self):
+        # Testing classic classes...
+        with torch._dynamo.error_on_graph_break(False):
+            class C:
+                def foo(*a): return a
+                goo = classmethod(foo)
+        c = C()
+        self.assertEqual(C.goo(1), (C, 1))
+        self.assertEqual(c.goo(1), (C, 1))
+        self.assertEqual(c.foo(1), (c, 1))
+        with torch._dynamo.error_on_graph_break(False):
+            class D(C):
+                pass
+        d = D()
+        self.assertEqual(D.goo(1), (D, 1))
+        self.assertEqual(d.goo(1), (D, 1))
+        self.assertEqual(d.foo(1), (d, 1))
+        self.assertEqual(D.foo(d, 1), (d, 1))
+        with torch._dynamo.error_on_graph_break(False):
+            class E: # *not* subclassing from C
+                foo = C.foo
+        self.assertEqual(E().foo.__func__, C.foo) # i.e., unbound
+        self.assertTrue(repr(C.foo.__get__(C())).startswith("<bound method "))
+
+    def test_compattr(self):
+        # Testing computed attributes...
+        with torch._dynamo.error_on_graph_break(False):
+            class C(object):
+                class computed_attribute(object):
+                    def __init__(self, get, set=None, delete=None):
+                        self.__get = get
+                        self.__set = set
+                        self.__delete = delete
+                    def __get__(self, obj, type=None):
+                        return self.__get(obj)
+                    def __set__(self, obj, value):
+                        return self.__set(obj, value)
+                    def __delete__(self, obj):
+                        return self.__delete(obj)
+                def __init__(self):
+                    self.__x = 0
+                def __get_x(self):
+                    x = self.__x
+                    self.__x = x+1
+                    return x
+                def __set_x(self, x):
+                    self.__x = x
+                def __delete_x(self):
+                    del self.__x
+                x = computed_attribute(__get_x, __set_x, __delete_x)
+        a = C()
+        self.assertEqual(a.x, 0)
+        self.assertEqual(a.x, 1)
+        a.x = 10
+        self.assertEqual(a.x, 10)
+        self.assertEqual(a.x, 11)
+        del a.x
+        self.assertNotHasAttr(a, 'x')
+
+    def test_newslots(self):
+        # Testing __new__ slot override...
+        with torch._dynamo.error_on_graph_break(False):
+            class C(list):
+                def __new__(cls):
+                    self = list.__new__(cls)
+                    self.foo = 1
+                    return self
+                def __init__(self):
+                    self.foo = self.foo + 2
+        a = C()
+        self.assertEqual(a.foo, 3)
+        self.assertEqual(a.__class__, C)
+        with torch._dynamo.error_on_graph_break(False):
+            class D(C):
+                pass
+        b = D()
+        self.assertEqual(b.foo, 3)
+        self.assertEqual(b.__class__, D)
+
+    @unittest.expectedFailure
+    def test_bad_new(self):
+        self.assertRaises(TypeError, object.__new__)
+        self.assertRaises(TypeError, object.__new__, '')
+        self.assertRaises(TypeError, list.__new__, object)
+        self.assertRaises(TypeError, object.__new__, list)
+        class C(object):
+            __new__ = list.__new__
+        self.assertRaises(TypeError, C)
+        class C(list):
+            __new__ = object.__new__
+        self.assertRaises(TypeError, C)
+
+    def test_object_new(self):
+        with torch._dynamo.error_on_graph_break(False):
+            class A(object):
+                pass
+        object.__new__(A)
+        self.assertRaises(TypeError, object.__new__, A, 5)
+        object.__init__(A())
+        self.assertRaises(TypeError, object.__init__, A(), 5)
+
+        with torch._dynamo.error_on_graph_break(False):
+            class A(object):
+                def __init__(self, foo):
+                    self.foo = foo
+        object.__new__(A)
+        object.__new__(A, 5)
+        object.__init__(A(3))
+        self.assertRaises(TypeError, object.__init__, A(3), 5)
+
+        with torch._dynamo.error_on_graph_break(False):
+            class A(object):
+                def __new__(cls, foo):
+                    return object.__new__(cls)
+        object.__new__(A)
+        self.assertRaises(TypeError, object.__new__, A, 5)
+        object.__init__(A(3))
+        object.__init__(A(3), 5)
+
+        with torch._dynamo.error_on_graph_break(False):
+            class A(object):
+                def __new__(cls, foo):
+                    return object.__new__(cls)
+                def __init__(self, foo):
+                    self.foo = foo
+        object.__new__(A)
+        self.assertRaises(TypeError, object.__new__, A, 5)
+        object.__init__(A(3))
+        self.assertRaises(TypeError, object.__init__, A(3), 5)
+
+    @unittest.expectedFailure
+    def test_restored_object_new(self):
+        class A(object):
+            def __new__(cls, *args, **kwargs):
+                raise AssertionError
+        self.assertRaises(AssertionError, A)
+        class B(A):
+            __new__ = object.__new__
+            def __init__(self, foo):
+                self.foo = foo
+        with warnings.catch_warnings():
+            warnings.simplefilter('error', DeprecationWarning)
+            b = B(3)
+        self.assertEqual(b.foo, 3)
+        self.assertEqual(b.__class__, B)
+        del B.__new__
+        self.assertRaises(AssertionError, B)
+        del A.__new__
+        with warnings.catch_warnings():
+            warnings.simplefilter('error', DeprecationWarning)
+            b = B(3)
+        self.assertEqual(b.foo, 3)
+        self.assertEqual(b.__class__, B)
+
+    def test_altmro(self):
+        # Testing mro() and overriding it...
+        with torch._dynamo.error_on_graph_break(False):
+            class A(object):
+                def f(self): return "A"
+            class B(A):
+                pass
+            class C(A):
+                def f(self): return "C"
+            class D(B, C):
+                pass
+        self.assertEqual(A.mro(), [A, object])
+        self.assertEqual(A.__mro__, (A, object))
+        self.assertEqual(B.mro(), [B, A, object])
+        self.assertEqual(B.__mro__, (B, A, object))
+        self.assertEqual(C.mro(), [C, A, object])
+        self.assertEqual(C.__mro__, (C, A, object))
+        self.assertEqual(D.mro(), [D, B, C, A, object])
+        self.assertEqual(D.__mro__, (D, B, C, A, object))
+        self.assertEqual(D().f(), "C")
+
+        with torch._dynamo.error_on_graph_break(False):
+            class PerverseMetaType(type):
+                def mro(cls):
+                    L = type.mro(cls)
+                    L.reverse()
+                    return L
+            class X(D,B,C,A, metaclass=PerverseMetaType):
+                pass
+        self.assertEqual(X.__mro__, (object, A, C, B, D, X))
+        self.assertEqual(X().f(), "A")
+
+        try:
+            with torch._dynamo.error_on_graph_break(False):
+                class _metaclass(type):
+                    def mro(self):
+                        return [self, dict, object]
+                class X(object, metaclass=_metaclass):
+                    pass
+            # In CPython, the class creation above already raises
+            # TypeError, as a protection against the fact that
+            # instances of X would segfault it.  In other Python
+            # implementations it would be ok to let the class X
+            # be created, but instead get a clean TypeError on the
+            # __setitem__ below.
+            x = object.__new__(X)
+            x[5] = 6
+        except TypeError:
+            pass
+        else:
+            self.fail("devious mro() return not caught")
+
+        try:
+            with torch._dynamo.error_on_graph_break(False):
+                class _metaclass(type):
+                    def mro(self):
+                        return [1]
+                class X(object, metaclass=_metaclass):
+                    pass
+        except TypeError:
+            pass
+        else:
+            self.fail("non-class mro() return not caught")
+
+        try:
+            with torch._dynamo.error_on_graph_break(False):
+                class _metaclass(type):
+                    def mro(self):
+                        return 1
+                class X(object, metaclass=_metaclass):
+                    pass
+        except TypeError:
+            pass
+        else:
+            self.fail("non-sequence mro() return not caught")
+
+    def test_overloading(self):
+        # Testing operator overloading...
+
+        with torch._dynamo.error_on_graph_break(False):
+            class B(object):
+                "Intermediate class because object doesn't have a __setattr__"
+
+            class C(B):
+                def __getattr__(self, name):
+                    if name == "foo":
+                        return ("getattr", name)
+                    else:
+                        raise AttributeError
+                def __setattr__(self, name, value):
+                    if name == "foo":
+                        self.setattr = (name, value)
+                    else:
+                        return B.__setattr__(self, name, value)
+                def __delattr__(self, name):
+                    if name == "foo":
+                        self.delattr = name
+                    else:
+                        return B.__delattr__(self, name)
+
+                def __getitem__(self, key):
+                    return ("getitem", key)
+                def __setitem__(self, key, value):
+                    self.setitem = (key, value)
+                def __delitem__(self, key):
+                    self.delitem = key
+
+        a = C()
+        self.assertEqual(a.foo, ("getattr", "foo"))
+        a.foo = 12
+        self.assertEqual(a.setattr, ("foo", 12))
+        del a.foo
+        self.assertEqual(a.delattr, "foo")
+
+        self.assertEqual(a[12], ("getitem", 12))
+        a[12] = 21
+        self.assertEqual(a.setitem, (12, 21))
+        del a[12]
+        self.assertEqual(a.delitem, 12)
+
+        self.assertEqual(a[0:10], ("getitem", slice(0, 10)))
+        a[0:10] = "foo"
+        self.assertEqual(a.setitem, (slice(0, 10), "foo"))
+        del a[0:10]
+        self.assertEqual(a.delitem, (slice(0, 10)))
+
+    def test_load_attr_extended_arg(self):
+        # https://github.com/python/cpython/issues/91625
+        with torch._dynamo.error_on_graph_break(False):
+            class Numbers:
+                def __getattr__(self, attr):
+                    return int(attr.lstrip("_"))
+        attrs = ", ".join(f"Z._{n:03d}" for n in range(280))
+        code = f"def number_attrs(Z):\n    return [ {attrs} ]"
+        ns = {}
+        exec(code, ns)
+        number_attrs = ns["number_attrs"]
+        # Warm up the function for quickening (PEP 659)
+        for _ in range(30):
+            self.assertEqual(number_attrs(Numbers()), list(range(280)))
+
+    def test_methods(self):
+        # Testing methods...
+        with torch._dynamo.error_on_graph_break(False):
+            class C(object):
+                def __init__(self, x):
+                    self.x = x
+                def foo(self):
+                    return self.x
+        c1 = C(1)
+        self.assertEqual(c1.foo(), 1)
+        with torch._dynamo.error_on_graph_break(False):
+            class D(C):
+                boo = C.foo
+                goo = c1.foo
+        d2 = D(2)
+        self.assertEqual(d2.foo(), 2)
+        self.assertEqual(d2.boo(), 2)
+        self.assertEqual(d2.goo(), 1)
+        with torch._dynamo.error_on_graph_break(False):
+            class E(object):
+                foo = C.foo
+        self.assertEqual(E().foo.__func__, C.foo) # i.e., unbound
+        self.assertTrue(repr(C.foo.__get__(C(1))).startswith("<bound method "))
+
+    @support.impl_detail("testing error message from implementation")
+    def test_methods_in_c(self):
+        # This test checks error messages in builtin method descriptor.
+        # It is allowed that other Python implementations use
+        # different error messages.
+        set_add = set.add
+
+        expected_errmsg = "unbound method set.add() needs an argument"
+
+        with self.assertRaises(TypeError) as cm:
+            set_add()
+        self.assertEqual(cm.exception.args[0], expected_errmsg)
+
+        expected_errmsg = "descriptor 'add' for 'set' objects doesn't apply to a 'int' object"
+
+        with self.assertRaises(TypeError) as cm:
+            set_add(0)
+        self.assertEqual(cm.exception.args[0], expected_errmsg)
+
+        with self.assertRaises(TypeError) as cm:
+            set_add.__get__(0)
+        self.assertEqual(cm.exception.args[0], expected_errmsg)
+
+    def test_special_method_lookup(self):
+        # The lookup of special methods bypasses __getattr__ and
+        # __getattribute__, but they still can be descriptors.
+
+        def run_context(manager):
+            with manager:
+                pass
+        def iden(self):
+            return self
+        def hello(self):
+            return b"hello"
+        def empty_seq(self):
+            return []
+        def zero(self):
+            return 0
+        def complex_num(self):
+            return 1j
+        def stop(self):
+            raise StopIteration
+        def return_true(self, thing=None):
+            return True
+        def do_isinstance(obj):
+            return isinstance(int, obj)
+        def do_issubclass(obj):
+            return issubclass(int, obj)
+        def do_dict_missing(checker):
+            with torch._dynamo.error_on_graph_break(False):
+                class DictSub(checker.__class__, dict):
+                    pass
+            self.assertEqual(DictSub()["hi"], 4)
+        def some_number(self_, key):
+            self.assertEqual(key, "hi")
+            return 4
+        def swallow(*args): pass
+        def format_impl(self, spec):
+            return "hello"
+
+        # It would be nice to have every special method tested here, but I'm
+        # only listing the ones I can remember outside of typeobject.c, since it
+        # does it right.
+        specials = [
+            ("__bytes__", bytes, hello, set(), {}),
+            ("__reversed__", reversed, empty_seq, set(), {}),
+            ("__length_hint__", list, zero, set(),
+             {"__iter__" : iden, "__next__" : stop}),
+            ("__sizeof__", sys.getsizeof, zero, set(), {}),
+            ("__instancecheck__", do_isinstance, return_true, set(), {}),
+            ("__missing__", do_dict_missing, some_number,
+             set(("__class__",)), {}),
+            ("__subclasscheck__", do_issubclass, return_true,
+             set(("__bases__",)), {}),
+            ("__enter__", run_context, iden, set(), {"__exit__" : swallow}),
+            ("__exit__", run_context, swallow, set(), {"__enter__" : iden}),
+            ("__complex__", complex, complex_num, set(), {}),
+            ("__format__", format, format_impl, set(), {}),
+            ("__floor__", math.floor, zero, set(), {}),
+            ("__trunc__", math.trunc, zero, set(), {}),
+            ("__ceil__", math.ceil, zero, set(), {}),
+            ("__dir__", dir, empty_seq, set(), {}),
+            ("__round__", round, zero, set(), {}),
+            ]
+
+        with torch._dynamo.error_on_graph_break(False):
+            class Checker(object):
+                def __getattr__(self, attr, test=self):
+                    test.fail("__getattr__ called with {0}".format(attr))
+                def __getattribute__(self, attr, test=self):
+                    if attr not in ok:
+                        test.fail("__getattribute__ called with {0}".format(attr))
+                    return object.__getattribute__(self, attr)
+            class SpecialDescr(object):
+                def __init__(self, impl):
+                    self.impl = impl
+                def __get__(self, obj, owner):
+                    record.append(1)
+                    return self.impl.__get__(obj, owner)
+            class MyException(Exception):
+                pass
+            class ErrDescr(object):
+                def __get__(self, obj, owner):
+                    raise MyException
+
+        for name, runner, meth_impl, ok, env in specials:
+            with torch._dynamo.error_on_graph_break(False):
+                class X(Checker):
+                    pass
+            for attr, obj in env.items():
+                setattr(X, attr, obj)
+            setattr(X, name, meth_impl)
+            runner(X())
+
+            record = []
+            with torch._dynamo.error_on_graph_break(False):
+                class X(Checker):
+                    pass
+            for attr, obj in env.items():
+                setattr(X, attr, obj)
+            setattr(X, name, SpecialDescr(meth_impl))
+            runner(X())
+            self.assertEqual(record, [1], name)
+
+            with torch._dynamo.error_on_graph_break(False):
+                class X(Checker):
+                    pass
+            for attr, obj in env.items():
+                setattr(X, attr, obj)
+            setattr(X, name, ErrDescr())
+            self.assertRaises(MyException, runner, X())
+
+    def test_specials(self):
+        # Testing special operators...
+        # Test operators like __hash__ for which a built-in default exists
+
+        # Test the default behavior for static classes
+        with torch._dynamo.error_on_graph_break(False):
+            class C(object):
+                def __getitem__(self, i):
+                    if 0 <= i < 10: return i
+                    raise IndexError
+        c1 = C()
+        c2 = C()
+        self.assertFalse(not c1)
+        self.assertNotEqual(id(c1), id(c2))
+        hash(c1)
+        hash(c2)
+        self.assertEqual(c1, c1)
+        self.assertTrue(c1 != c2)
+        self.assertFalse(c1 != c1)
+        self.assertFalse(c1 == c2)
+        # Note that the module name appears in str/repr, and that varies
+        # depending on whether this test is run standalone or from a framework.
+        self.assertGreaterEqual(str(c1).find('C object at '), 0)
+        self.assertEqual(str(c1), repr(c1))
+        self.assertNotIn(-1, c1)
+        for i in range(10):
+            self.assertIn(i, c1)
+        self.assertNotIn(10, c1)
+        # Test the default behavior for dynamic classes
+        with torch._dynamo.error_on_graph_break(False):
+            class D(object):
+                def __getitem__(self, i):
+                    if 0 <= i < 10: return i
+                    raise IndexError
+        d1 = D()
+        d2 = D()
+        self.assertFalse(not d1)
+        self.assertNotEqual(id(d1), id(d2))
+        hash(d1)
+        hash(d2)
+        self.assertEqual(d1, d1)
+        self.assertNotEqual(d1, d2)
+        self.assertFalse(d1 != d1)
+        self.assertFalse(d1 == d2)
+        # Note that the module name appears in str/repr, and that varies
+        # depending on whether this test is run standalone or from a framework.
+        self.assertGreaterEqual(str(d1).find('D object at '), 0)
+        self.assertEqual(str(d1), repr(d1))
+        self.assertNotIn(-1, d1)
+        for i in range(10):
+            self.assertIn(i, d1)
+        self.assertNotIn(10, d1)
+        # Test overridden behavior
+        with torch._dynamo.error_on_graph_break(False):
+            class Proxy(object):
+                def __init__(self, x):
+                    self.x = x
+                def __bool__(self):
+                    return not not self.x
+                def __hash__(self):
+                    return hash(self.x)
+                def __eq__(self, other):
+                    return self.x == other
+                def __ne__(self, other):
+                    return self.x != other
+                def __ge__(self, other):
+                    return self.x >= other
+                def __gt__(self, other):
+                    return self.x > other
+                def __le__(self, other):
+                    return self.x <= other
+                def __lt__(self, other):
+                    return self.x < other
+                def __str__(self):
+                    return "Proxy:%s" % self.x
+                def __repr__(self):
+                    return "Proxy(%r)" % self.x
+                def __contains__(self, value):
+                    return value in self.x
+        p0 = Proxy(0)
+        p1 = Proxy(1)
+        p_1 = Proxy(-1)
+        self.assertFalse(p0)
+        self.assertFalse(not p1)
+        self.assertEqual(hash(p0), hash(0))
+        self.assertEqual(p0, p0)
+        self.assertNotEqual(p0, p1)
+        self.assertFalse(p0 != p0)
+        self.assertEqual(not p0, p1)
+        self.assertTrue(p0 < p1)
+        self.assertTrue(p0 <= p1)
+        self.assertTrue(p1 > p0)
+        self.assertTrue(p1 >= p0)
+        self.assertEqual(str(p0), "Proxy:0")
+        self.assertEqual(repr(p0), "Proxy(0)")
+        p10 = Proxy(range(10))
+        self.assertNotIn(-1, p10)
+        for i in range(10):
+            self.assertIn(i, p10)
+        self.assertNotIn(10, p10)
+
+    def test_weakrefs(self):
+        # Testing weak references...
+        import weakref
+        with torch._dynamo.error_on_graph_break(False):
+            class C(object):
+                pass
+        c = C()
+        r = weakref.ref(c)
+        self.assertEqual(r(), c)
+        del c
+        support.gc_collect()
+        self.assertEqual(r(), None)
+        del r
+        with torch._dynamo.error_on_graph_break(False):
+            class NoWeak(object):
+                __slots__ = ['foo']
+        no = NoWeak()
+        try:
+            weakref.ref(no)
+        except TypeError as msg:
+            self.assertIn("weak reference", str(msg))
+        else:
+            self.fail("weakref.ref(no) should be illegal")
+        with torch._dynamo.error_on_graph_break(False):
+            class Weak(object):
+                __slots__ = ['foo', '__weakref__']
+        yes = Weak()
+        r = weakref.ref(yes)
+        self.assertEqual(r(), yes)
+        del yes
+        support.gc_collect()
+        self.assertEqual(r(), None)
+        del r
+
+    def test_properties(self):
+        # Testing property...
+        with torch._dynamo.error_on_graph_break(False):
+            class C(object):
+                def getx(self):
+                    return self.__x
+                def setx(self, value):
+                    self.__x = value
+                def delx(self):
+                    del self.__x
+                x = property(getx, setx, delx, doc="I'm the x property.")
+        a = C()
+        self.assertNotHasAttr(a, "x")
+        a.x = 42
+        self.assertEqual(a._C__x, 42)
+        self.assertEqual(a.x, 42)
+        del a.x
+        self.assertNotHasAttr(a, "x")
+        self.assertNotHasAttr(a, "_C__x")
+        C.x.__set__(a, 100)
+        self.assertEqual(C.x.__get__(a), 100)
+        C.x.__delete__(a)
+        self.assertNotHasAttr(a, "x")
+
+        raw = C.__dict__['x']
+        self.assertIsInstance(raw, property)
+
+        attrs = dir(raw)
+        self.assertIn("__doc__", attrs)
+        self.assertIn("fget", attrs)
+        self.assertIn("fset", attrs)
+        self.assertIn("fdel", attrs)
+
+        self.assertEqual(raw.__doc__, "I'm the x property.")
+        self.assertIs(raw.fget, C.__dict__['getx'])
+        self.assertIs(raw.fset, C.__dict__['setx'])
+        self.assertIs(raw.fdel, C.__dict__['delx'])
+
+        for attr in "fget", "fset", "fdel":
+            try:
+                setattr(raw, attr, 42)
+            except AttributeError as msg:
+                if str(msg).find('readonly') < 0:
+                    self.fail("when setting readonly attr %r on a property, "
+                              "got unexpected AttributeError msg %r" % (attr, str(msg)))
+            else:
+                self.fail("expected AttributeError from trying to set readonly %r "
+                          "attr on a property" % attr)
+
+        raw.__doc__ = 42
+        self.assertEqual(raw.__doc__, 42)
+
+        with torch._dynamo.error_on_graph_break(False):
+            class D(object):
+                __getitem__ = property(lambda s: 1/0)
+
+        d = D()
+        try:
+            for i in d:
+                str(i)
+        except ZeroDivisionError:
+            pass
+        else:
+            self.fail("expected ZeroDivisionError from bad property")
+
+    @unittest.skipIf(sys.flags.optimize >= 2,
+                     "Docstrings are omitted with -O2 and above")
+    def test_properties_doc_attrib(self):
+        with torch._dynamo.error_on_graph_break(False):
+            class E(object):
+                def getter(self):
+                    "getter method"
+                    return 0
+                def setter(self_, value):
+                    "setter method"
+                    pass
+                prop = property(getter)
+                self.assertEqual(prop.__doc__, "getter method")
+                prop2 = property(fset=setter)
+                self.assertEqual(prop2.__doc__, None)
+
+    @support.cpython_only
+    def test_testcapi_no_segfault(self):
+        # this segfaulted in 2.5b2
+        try:
+            import _testcapi
+        except ImportError:
+            pass
+        else:
+            with torch._dynamo.error_on_graph_break(False):
+                class X(object):
+                    p = property(_testcapi.test_with_docstring)
+
+    def test_properties_plus(self):
+        with torch._dynamo.error_on_graph_break(False):
+            class C(object):
+                foo = property(doc="hello")
+                @foo.getter
+                def foo(self):
+                    return self._foo
+                @foo.setter
+                def foo(self, value):
+                    self._foo = abs(value)
+                @foo.deleter
+                def foo(self):
+                    del self._foo
+        c = C()
+        self.assertEqual(C.foo.__doc__, "hello")
+        self.assertNotHasAttr(c, "foo")
+        c.foo = -42
+        self.assertHasAttr(c, '_foo')
+        self.assertEqual(c._foo, 42)
+        self.assertEqual(c.foo, 42)
+        del c.foo
+        self.assertNotHasAttr(c, '_foo')
+        self.assertNotHasAttr(c, "foo")
+
+        with torch._dynamo.error_on_graph_break(False):
+            class D(C):
+                @C.foo.deleter
+                def foo(self):
+                    try:
+                        del self._foo
+                    except AttributeError:
+                        pass
+        d = D()
+        d.foo = 24
+        self.assertEqual(d.foo, 24)
+        del d.foo
+        del d.foo
+
+        with torch._dynamo.error_on_graph_break(False):
+            class E(object):
+                @property
+                def foo(self):
+                    return self._foo
+                @foo.setter
+                def foo(self, value):
+                    raise RuntimeError
+                @foo.setter
+                def foo(self, value):
+                    self._foo = abs(value)
+                @foo.deleter
+                def foo(self, value=None):
+                    del self._foo
+
+        e = E()
+        e.foo = -42
+        self.assertEqual(e.foo, 42)
+        del e.foo
+
+        with torch._dynamo.error_on_graph_break(False):
+            class F(E):
+                @E.foo.deleter
+                def foo(self):
+                    del self._foo
+                @foo.setter
+                def foo(self, value):
+                    self._foo = max(0, value)
+        f = F()
+        f.foo = -10
+        self.assertEqual(f.foo, 0)
+        del f.foo
+
+    def test_dict_constructors(self):
+        # Testing dict constructor ...
+        d = dict()
+        self.assertEqual(d, {})
+        d = dict({})
+        self.assertEqual(d, {})
+        d = dict({1: 2, 'a': 'b'})
+        self.assertEqual(d, {1: 2, 'a': 'b'})
+        self.assertEqual(d, dict(list(d.items())))
+        self.assertEqual(d, dict(iter(d.items())))
+        d = dict({'one':1, 'two':2})
+        self.assertEqual(d, dict(one=1, two=2))
+        self.assertEqual(d, dict(**d))
+        self.assertEqual(d, dict({"one": 1}, two=2))
+        self.assertEqual(d, dict([("two", 2)], one=1))
+        self.assertEqual(d, dict([("one", 100), ("two", 200)], **d))
+        self.assertEqual(d, dict(**d))
+
+        for badarg in 0, 0, 0j, "0", [0], (0,):
+            try:
+                dict(badarg)
+            except TypeError:
+                pass
+            except ValueError:
+                if badarg == "0":
+                    # It's a sequence, and its elements are also sequences (gotta
+                    # love strings <wink>), but they aren't of length 2, so this
+                    # one seemed better as a ValueError than a TypeError.
+                    pass
+                else:
+                    self.fail("no TypeError from dict(%r)" % badarg)
+            else:
+                self.fail("no TypeError from dict(%r)" % badarg)
+
+        with self.assertRaises(TypeError):
+            dict({}, {})
+
+        with torch._dynamo.error_on_graph_break(False):
+            class Mapping:
+                # Lacks a .keys() method; will be added later.
+                dict = {1:2, 3:4, 'a':1j}
+
+        try:
+            dict(Mapping())
+        except TypeError:
+            pass
+        else:
+            self.fail("no TypeError from dict(incomplete mapping)")
+
+        Mapping.keys = lambda self: list(self.dict.keys())
+        Mapping.__getitem__ = lambda self, i: self.dict[i]
+        d = dict(Mapping())
+        self.assertEqual(d, Mapping.dict)
+
+        # Init from sequence of iterable objects, each producing a 2-sequence.
+        with torch._dynamo.error_on_graph_break(False):
+            class AddressBookEntry:
+                def __init__(self, first, last):
+                    self.first = first
+                    self.last = last
+                def __iter__(self):
+                    return iter([self.first, self.last])
+
+        d = dict([AddressBookEntry('Tim', 'Warsaw'),
+                  AddressBookEntry('Barry', 'Peters'),
+                  AddressBookEntry('Tim', 'Peters'),
+                  AddressBookEntry('Barry', 'Warsaw')])
+        self.assertEqual(d, {'Barry': 'Warsaw', 'Tim': 'Peters'})
+
+        d = dict(zip(range(4), range(1, 5)))
+        self.assertEqual(d, dict([(i, i+1) for i in range(4)]))
+
+        # Bad sequence lengths.
+        for bad in [('tooshort',)], [('too', 'long', 'by 1')]:
+            try:
+                dict(bad)
+            except ValueError:
+                pass
+            else:
+                self.fail("no ValueError from dict(%r)" % bad)
+
+    def test_dir(self):
+        # Testing dir() ...
+        junk = 12
+        self.assertEqual(dir(), ['junk', 'self'])
+        del junk
+
+        # Just make sure these don't blow up!
+        for arg in 2, 2, 2j, 2e0, [2], "2", b"2", (2,), {2:2}, type, self.test_dir:
+            dir(arg)
+
+        # Test dir on new-style classes.  Since these have object as a
+        # base class, a lot more gets sucked in.
+        def interesting(strings):
+            return [s for s in strings if not s.startswith('_')]
+
+        with torch._dynamo.error_on_graph_break(False):
+            class C(object):
+                Cdata = 1
+                def Cmethod(self): pass
+
+        cstuff = ['Cdata', 'Cmethod']
+        self.assertEqual(interesting(dir(C)), cstuff)
+
+        c = C()
+        self.assertEqual(interesting(dir(c)), cstuff)
+        ## self.assertIn('__self__', dir(C.Cmethod))
+
+        c.cdata = 2
+        c.cmethod = lambda self: 0
+        self.assertEqual(interesting(dir(c)), cstuff + ['cdata', 'cmethod'])
+        ## self.assertIn('__self__', dir(c.Cmethod))
+
+        with torch._dynamo.error_on_graph_break(False):
+            class A(C):
+                Adata = 1
+                def Amethod(self): pass
+
+        astuff = ['Adata', 'Amethod'] + cstuff
+        self.assertEqual(interesting(dir(A)), astuff)
+        ## self.assertIn('__self__', dir(A.Amethod))
+        a = A()
+        self.assertEqual(interesting(dir(a)), astuff)
+        a.adata = 42
+        a.amethod = lambda self: 3
+        self.assertEqual(interesting(dir(a)), astuff + ['adata', 'amethod'])
+        ## self.assertIn('__self__', dir(a.Amethod))
+
+        # Try a module subclass.
+        with torch._dynamo.error_on_graph_break(False):
+            class M(type(sys)):
+                pass
+        minstance = M("m")
+        minstance.b = 2
+        minstance.a = 1
+        default_attributes = ['__name__', '__doc__', '__package__',
+                              '__loader__', '__spec__']
+        names = [x for x in dir(minstance) if x not in default_attributes]
+        self.assertEqual(names, ['a', 'b'])
+
+        with torch._dynamo.error_on_graph_break(False):
+            class M2(M):
+                def getdict(self):
+                    return "Not a dict!"
+                __dict__ = property(getdict)
+
+        m2instance = M2("m2")
+        m2instance.b = 2
+        m2instance.a = 1
+        self.assertEqual(m2instance.__dict__, "Not a dict!")
+        with self.assertRaises(TypeError):
+            dir(m2instance)
+
+        # Two essentially featureless objects, (Ellipsis just inherits stuff
+        # from object.
+        self.assertEqual(dir(object()), dir(Ellipsis))
+
+        # Nasty test case for proxied objects
+        with torch._dynamo.error_on_graph_break(False):
+            class Wrapper(object):
+                def __init__(self, obj):
+                    self.__obj = obj
+                def __repr__(self):
+                    return "Wrapper(%s)" % repr(self.__obj)
+                def __getitem__(self, key):
+                    return Wrapper(self.__obj[key])
+                def __len__(self):
+                    return len(self.__obj)
+                def __getattr__(self, name):
+                    return Wrapper(getattr(self.__obj, name))
+
+            class C(object):
+                def __getclass(self):
+                    return Wrapper(type(self))
+                __class__ = property(__getclass)
+
+        dir(C()) # This used to segfault
+
+    def test_supers(self):
+        # Testing super...
+
+        with torch._dynamo.error_on_graph_break(False):
+            class A(object):
+                def meth(self, a):
+                    return "A(%r)" % a
+
+        self.assertEqual(A().meth(1), "A(1)")
+
+        with torch._dynamo.error_on_graph_break(False):
+            class B(A):
+                def __init__(self):
+                    self.__super = super(B, self)
+                def meth(self, a):
+                    return "B(%r)" % a + self.__super.meth(a)
+
+        self.assertEqual(B().meth(2), "B(2)A(2)")
+
+        with torch._dynamo.error_on_graph_break(False):
+            class C(A):
+                def meth(self, a):
+                    return "C(%r)" % a + self.__super.meth(a)
+        C._C__super = super(C)
+
+        self.assertEqual(C().meth(3), "C(3)A(3)")
+
+        with torch._dynamo.error_on_graph_break(False):
+            class D(C, B):
+                def meth(self, a):
+                    return "D(%r)" % a + super(D, self).meth(a)
+
+        self.assertEqual(D().meth(4), "D(4)C(4)B(4)A(4)")
+
+        # Test for subclassing super
+
+        with torch._dynamo.error_on_graph_break(False):
+            class mysuper(super):
+                def __init__(self, *args):
+                    return super(mysuper, self).__init__(*args)
+
+            class E(D):
+                def meth(self, a):
+                    return "E(%r)" % a + mysuper(E, self).meth(a)
+
+        self.assertEqual(E().meth(5), "E(5)D(5)C(5)B(5)A(5)")
+
+        with torch._dynamo.error_on_graph_break(False):
+            class F(E):
+                def meth(self, a):
+                    s = self.__super # == mysuper(F, self)
+                    return "F(%r)[%s]" % (a, s.__class__.__name__) + s.meth(a)
+        F._F__super = mysuper(F)
+
+        self.assertEqual(F().meth(6), "F(6)[mysuper]E(6)D(6)C(6)B(6)A(6)")
+
+        # Make sure certain errors are raised
+
+        try:
+            super(D, 42)
+        except TypeError:
+            pass
+        else:
+            self.fail("shouldn't allow super(D, 42)")
+
+        try:
+            super(D, C())
+        except TypeError:
+            pass
+        else:
+            self.fail("shouldn't allow super(D, C())")
+
+        try:
+            super(D).__get__(12)
+        except TypeError:
+            pass
+        else:
+            self.fail("shouldn't allow super(D).__get__(12)")
+
+        try:
+            super(D).__get__(C())
+        except TypeError:
+            pass
+        else:
+            self.fail("shouldn't allow super(D).__get__(C())")
+
+        # Make sure data descriptors can be overridden and accessed via super
+        # (new feature in Python 2.3)
+
+        with torch._dynamo.error_on_graph_break(False):
+            class DDbase(object):
+                def getx(self): return 42
+                x = property(getx)
+
+            class DDsub(DDbase):
+                def getx(self): return "hello"
+                x = property(getx)
+
+        dd = DDsub()
+        self.assertEqual(dd.x, "hello")
+        self.assertEqual(super(DDsub, dd).x, 42)
+
+        # Ensure that super() lookup of descriptor from classmethod
+        # works (SF ID# 743627)
+
+        with torch._dynamo.error_on_graph_break(False):
+            class Base(object):
+                aProp = property(lambda self: "foo")
+
+            class Sub(Base):
+                @classmethod
+                def test(klass):
+                    return super(Sub,klass).aProp
+
+        self.assertEqual(Sub.test(), Base.aProp)
+
+        # Verify that super() doesn't allow keyword args
+        with self.assertRaises(TypeError):
+            super(Base, kw=1)
+
+    def test_basic_inheritance(self):
+        # Testing inheritance from basic types...
+
+        with torch._dynamo.error_on_graph_break(False):
+            class hexint(int):
+                def __repr__(self):
+                    return hex(self)
+                def __add__(self, other):
+                    return hexint(int.__add__(self, other))
+                # (Note that overriding __radd__ doesn't work,
+                # because the int type gets first dibs.)
+        self.assertEqual(repr(hexint(7) + 9), "0x10")
+        self.assertEqual(repr(hexint(1000) + 7), "0x3ef")
+        a = hexint(12345)
+        self.assertEqual(a, 12345)
+        self.assertEqual(int(a), 12345)
+        self.assertIs(int(a).__class__, int)
+        self.assertEqual(hash(a), hash(12345))
+        self.assertIs((+a).__class__, int)
+        self.assertIs((a >> 0).__class__, int)
+        self.assertIs((a << 0).__class__, int)
+        self.assertIs((hexint(0) << 12).__class__, int)
+        self.assertIs((hexint(0) >> 12).__class__, int)
+
+        with torch._dynamo.error_on_graph_break(False):
+            class octlong(int):
+                __slots__ = []
+                def __str__(self):
+                    return oct(self)
+                def __add__(self, other):
+                    return self.__class__(super(octlong, self).__add__(other))
+                __radd__ = __add__
+        self.assertEqual(str(octlong(3) + 5), "0o10")
+        # (Note that overriding __radd__ here only seems to work
+        # because the example uses a short int left argument.)
+        self.assertEqual(str(5 + octlong(3000)), "0o5675")
+        a = octlong(12345)
+        self.assertEqual(a, 12345)
+        self.assertEqual(int(a), 12345)
+        self.assertEqual(hash(a), hash(12345))
+        self.assertIs(int(a).__class__, int)
+        self.assertIs((+a).__class__, int)
+        self.assertIs((-a).__class__, int)
+        self.assertIs((-octlong(0)).__class__, int)
+        self.assertIs((a >> 0).__class__, int)
+        self.assertIs((a << 0).__class__, int)
+        self.assertIs((a - 0).__class__, int)
+        self.assertIs((a * 1).__class__, int)
+        self.assertIs((a ** 1).__class__, int)
+        self.assertIs((a // 1).__class__, int)
+        self.assertIs((1 * a).__class__, int)
+        self.assertIs((a | 0).__class__, int)
+        self.assertIs((a ^ 0).__class__, int)
+        self.assertIs((a & -1).__class__, int)
+        self.assertIs((octlong(0) << 12).__class__, int)
+        self.assertIs((octlong(0) >> 12).__class__, int)
+        self.assertIs(abs(octlong(0)).__class__, int)
+
+        # Because octlong overrides __add__, we can't check the absence of +0
+        # optimizations using octlong.
+        with torch._dynamo.error_on_graph_break(False):
+            class longclone(int):
+                pass
+        a = longclone(1)
+        self.assertIs((a + 0).__class__, int)
+        self.assertIs((0 + a).__class__, int)
+
+        # Check that negative clones don't segfault
+        a = longclone(-1)
+        self.assertEqual(a.__dict__, {})
+        self.assertEqual(int(a), -1)  # self.assertTrue PyNumber_Long() copies the sign bit
+
+        with torch._dynamo.error_on_graph_break(False):
+            class precfloat(float):
+                __slots__ = ['prec']
+                def __init__(self, value=0.0, prec=12):
+                    self.prec = int(prec)
+                def __repr__(self):
+                    return "%.*g" % (self.prec, self)
+        self.assertEqual(repr(precfloat(1.1)), "1.1")
+        a = precfloat(12345)
+        self.assertEqual(a, 12345.0)
+        self.assertEqual(float(a), 12345.0)
+        self.assertIs(float(a).__class__, float)
+        self.assertEqual(hash(a), hash(12345.0))
+        self.assertIs((+a).__class__, float)
+
+        with torch._dynamo.error_on_graph_break(False):
+            class madcomplex(complex):
+                def __repr__(self):
+                    return "%.17gj%+.17g" % (self.imag, self.real)
+        a = madcomplex(-3, 4)
+        self.assertEqual(repr(a), "4j-3")
+        base = complex(-3, 4)
+        self.assertEqual(base.__class__, complex)
+        self.assertEqual(a, base)
+        self.assertEqual(complex(a), base)
+        self.assertEqual(complex(a).__class__, complex)
+        a = madcomplex(a)  # just trying another form of the constructor
+        self.assertEqual(repr(a), "4j-3")
+        self.assertEqual(a, base)
+        self.assertEqual(complex(a), base)
+        self.assertEqual(complex(a).__class__, complex)
+        self.assertEqual(hash(a), hash(base))
+        self.assertEqual((+a).__class__, complex)
+        self.assertEqual((a + 0).__class__, complex)
+        self.assertEqual(a + 0, base)
+        self.assertEqual((a - 0).__class__, complex)
+        self.assertEqual(a - 0, base)
+        self.assertEqual((a * 1).__class__, complex)
+        self.assertEqual(a * 1, base)
+        self.assertEqual((a / 1).__class__, complex)
+        self.assertEqual(a / 1, base)
+
+        with torch._dynamo.error_on_graph_break(False):
+            class madtuple(tuple):
+                _rev = None
+                def rev(self):
+                    if self._rev is not None:
+                        return self._rev
+                    L = list(self)
+                    L.reverse()
+                    self._rev = self.__class__(L)
+                    return self._rev
+        a = madtuple((1,2,3,4,5,6,7,8,9,0))
+        self.assertEqual(a, (1,2,3,4,5,6,7,8,9,0))
+        self.assertEqual(a.rev(), madtuple((0,9,8,7,6,5,4,3,2,1)))
+        self.assertEqual(a.rev().rev(), madtuple((1,2,3,4,5,6,7,8,9,0)))
+        for i in range(512):
+            t = madtuple(range(i))
+            u = t.rev()
+            v = u.rev()
+            self.assertEqual(v, t)
+        a = madtuple((1,2,3,4,5))
+        self.assertEqual(tuple(a), (1,2,3,4,5))
+        self.assertIs(tuple(a).__class__, tuple)
+        self.assertEqual(hash(a), hash((1,2,3,4,5)))
+        self.assertIs(a[:].__class__, tuple)
+        self.assertIs((a * 1).__class__, tuple)
+        self.assertIs((a * 0).__class__, tuple)
+        self.assertIs((a + ()).__class__, tuple)
+        a = madtuple(())
+        self.assertEqual(tuple(a), ())
+        self.assertIs(tuple(a).__class__, tuple)
+        self.assertIs((a + a).__class__, tuple)
+        self.assertIs((a * 0).__class__, tuple)
+        self.assertIs((a * 1).__class__, tuple)
+        self.assertIs((a * 2).__class__, tuple)
+        self.assertIs(a[:].__class__, tuple)
+
+        with torch._dynamo.error_on_graph_break(False):
+            class madstring(str):
+                _rev = None
+                def rev(self):
+                    if self._rev is not None:
+                        return self._rev
+                    L = list(self)
+                    L.reverse()
+                    self._rev = self.__class__("".join(L))
+                    return self._rev
+        s = madstring("abcdefghijklmnopqrstuvwxyz")
+        self.assertEqual(s, "abcdefghijklmnopqrstuvwxyz")
+        self.assertEqual(s.rev(), madstring("zyxwvutsrqponmlkjihgfedcba"))
+        self.assertEqual(s.rev().rev(), madstring("abcdefghijklmnopqrstuvwxyz"))
+        for i in range(256):
+            s = madstring("".join(map(chr, range(i))))
+            t = s.rev()
+            u = t.rev()
+            self.assertEqual(u, s)
+        s = madstring("12345")
+        self.assertEqual(str(s), "12345")
+        self.assertIs(str(s).__class__, str)
+
+        base = "\x00" * 5
+        s = madstring(base)
+        self.assertEqual(s, base)
+        self.assertEqual(str(s), base)
+        self.assertIs(str(s).__class__, str)
+        self.assertEqual(hash(s), hash(base))
+        self.assertEqual({s: 1}[base], 1)
+        self.assertEqual({base: 1}[s], 1)
+        self.assertIs((s + "").__class__, str)
+        self.assertEqual(s + "", base)
+        self.assertIs(("" + s).__class__, str)
+        self.assertEqual("" + s, base)
+        self.assertIs((s * 0).__class__, str)
+        self.assertEqual(s * 0, "")
+        self.assertIs((s * 1).__class__, str)
+        self.assertEqual(s * 1, base)
+        self.assertIs((s * 2).__class__, str)
+        self.assertEqual(s * 2, base + base)
+        self.assertIs(s[:].__class__, str)
+        self.assertEqual(s[:], base)
+        self.assertIs(s[0:0].__class__, str)
+        self.assertEqual(s[0:0], "")
+        self.assertIs(s.strip().__class__, str)
+        self.assertEqual(s.strip(), base)
+        self.assertIs(s.lstrip().__class__, str)
+        self.assertEqual(s.lstrip(), base)
+        self.assertIs(s.rstrip().__class__, str)
+        self.assertEqual(s.rstrip(), base)
+        identitytab = {}
+        self.assertIs(s.translate(identitytab).__class__, str)
+        self.assertEqual(s.translate(identitytab), base)
+        self.assertIs(s.replace("x", "x").__class__, str)
+        self.assertEqual(s.replace("x", "x"), base)
+        self.assertIs(s.ljust(len(s)).__class__, str)
+        self.assertEqual(s.ljust(len(s)), base)
+        self.assertIs(s.rjust(len(s)).__class__, str)
+        self.assertEqual(s.rjust(len(s)), base)
+        self.assertIs(s.center(len(s)).__class__, str)
+        self.assertEqual(s.center(len(s)), base)
+        self.assertIs(s.lower().__class__, str)
+        self.assertEqual(s.lower(), base)
+
+        with torch._dynamo.error_on_graph_break(False):
+            class madunicode(str):
+                _rev = None
+                def rev(self):
+                    if self._rev is not None:
+                        return self._rev
+                    L = list(self)
+                    L.reverse()
+                    self._rev = self.__class__("".join(L))
+                    return self._rev
+        u = madunicode("ABCDEF")
+        self.assertEqual(u, "ABCDEF")
+        self.assertEqual(u.rev(), madunicode("FEDCBA"))
+        self.assertEqual(u.rev().rev(), madunicode("ABCDEF"))
+        base = "12345"
+        u = madunicode(base)
+        self.assertEqual(str(u), base)
+        self.assertIs(str(u).__class__, str)
+        self.assertEqual(hash(u), hash(base))
+        self.assertEqual({u: 1}[base], 1)
+        self.assertEqual({base: 1}[u], 1)
+        self.assertIs(u.strip().__class__, str)
+        self.assertEqual(u.strip(), base)
+        self.assertIs(u.lstrip().__class__, str)
+        self.assertEqual(u.lstrip(), base)
+        self.assertIs(u.rstrip().__class__, str)
+        self.assertEqual(u.rstrip(), base)
+        self.assertIs(u.replace("x", "x").__class__, str)
+        self.assertEqual(u.replace("x", "x"), base)
+        self.assertIs(u.replace("xy", "xy").__class__, str)
+        self.assertEqual(u.replace("xy", "xy"), base)
+        self.assertIs(u.center(len(u)).__class__, str)
+        self.assertEqual(u.center(len(u)), base)
+        self.assertIs(u.ljust(len(u)).__class__, str)
+        self.assertEqual(u.ljust(len(u)), base)
+        self.assertIs(u.rjust(len(u)).__class__, str)
+        self.assertEqual(u.rjust(len(u)), base)
+        self.assertIs(u.lower().__class__, str)
+        self.assertEqual(u.lower(), base)
+        self.assertIs(u.upper().__class__, str)
+        self.assertEqual(u.upper(), base)
+        self.assertIs(u.capitalize().__class__, str)
+        self.assertEqual(u.capitalize(), base)
+        self.assertIs(u.title().__class__, str)
+        self.assertEqual(u.title(), base)
+        self.assertIs((u + "").__class__, str)
+        self.assertEqual(u + "", base)
+        self.assertIs(("" + u).__class__, str)
+        self.assertEqual("" + u, base)
+        self.assertIs((u * 0).__class__, str)
+        self.assertEqual(u * 0, "")
+        self.assertIs((u * 1).__class__, str)
+        self.assertEqual(u * 1, base)
+        self.assertIs((u * 2).__class__, str)
+        self.assertEqual(u * 2, base + base)
+        self.assertIs(u[:].__class__, str)
+        self.assertEqual(u[:], base)
+        self.assertIs(u[0:0].__class__, str)
+        self.assertEqual(u[0:0], "")
+
+        with torch._dynamo.error_on_graph_break(False):
+            class sublist(list):
+                pass
+        a = sublist(range(5))
+        self.assertEqual(a, list(range(5)))
+        a.append("hello")
+        self.assertEqual(a, list(range(5)) + ["hello"])
+        a[5] = 5
+        self.assertEqual(a, list(range(6)))
+        a.extend(range(6, 20))
+        self.assertEqual(a, list(range(20)))
+        a[-5:] = []
+        self.assertEqual(a, list(range(15)))
+        del a[10:15]
+        self.assertEqual(len(a), 10)
+        self.assertEqual(a, list(range(10)))
+        self.assertEqual(list(a), list(range(10)))
+        self.assertEqual(a[0], 0)
+        self.assertEqual(a[9], 9)
+        self.assertEqual(a[-10], 0)
+        self.assertEqual(a[-1], 9)
+        self.assertEqual(a[:5], list(range(5)))
+
+        ## class CountedInput(file):
+        ##    """Counts lines read by self.readline().
+        ##
+        ##     self.lineno is the 0-based ordinal of the last line read, up to
+        ##     a maximum of one greater than the number of lines in the file.
+        ##
+        ##     self.ateof is true if and only if the final "" line has been read,
+        ##     at which point self.lineno stops incrementing, and further calls
+        ##     to readline() continue to return "".
+        ##     """
+        ##
+        ##     lineno = 0
+        ##     ateof = 0
+        ##     def readline(self):
+        ##         if self.ateof:
+        ##             return ""
+        ##         s = file.readline(self)
+        ##         # Next line works too.
+        ##         # s = super(CountedInput, self).readline()
+        ##         self.lineno += 1
+        ##         if s == "":
+        ##             self.ateof = 1
+        ##        return s
+        ##
+        ## f = file(name=os_helper.TESTFN, mode='w')
+        ## lines = ['a\n', 'b\n', 'c\n']
+        ## try:
+        ##     f.writelines(lines)
+        ##     f.close()
+        ##     f = CountedInput(os_helper.TESTFN)
+        ##     for (i, expected) in zip(range(1, 5) + [4], lines + 2 * [""]):
+        ##         got = f.readline()
+        ##         self.assertEqual(expected, got)
+        ##         self.assertEqual(f.lineno, i)
+        ##         self.assertEqual(f.ateof, (i > len(lines)))
+        ##     f.close()
+        ## finally:
+        ##     try:
+        ##         f.close()
+        ##     except:
+        ##         pass
+        ##     os_helper.unlink(os_helper.TESTFN)
+
+    def test_keywords(self):
+        # Testing keyword args to basic type constructors ...
+        with self.assertRaisesRegex(TypeError, 'keyword argument'):
+            int(x=1)
+        with self.assertRaisesRegex(TypeError, 'keyword argument'):
+            float(x=2)
+        with self.assertRaisesRegex(TypeError, 'keyword argument'):
+            bool(x=2)
+        self.assertEqual(complex(imag=42, real=666), complex(666, 42))
+        self.assertEqual(str(object=500), '500')
+        self.assertEqual(str(object=b'abc', errors='strict'), 'abc')
+        with self.assertRaisesRegex(TypeError, 'keyword argument'):
+            tuple(sequence=range(3))
+        with self.assertRaisesRegex(TypeError, 'keyword argument'):
+            list(sequence=(0, 1, 2))
+        # note: as of Python 2.3, dict() no longer has an "items" keyword arg
+
+        for constructor in (int, float, int, complex, str, str,
+                            tuple, list):
+            try:
+                constructor(bogus_keyword_arg=1)
+            except TypeError:
+                pass
+            else:
+                self.fail("expected TypeError from bogus keyword argument to %r"
+                            % constructor)
+
+    def test_str_subclass_as_dict_key(self):
+        # Testing a str subclass used as dict key ..
+
+        with torch._dynamo.error_on_graph_break(False):
+            class cistr(str):
+                """Subclass of str that computes __eq__ case-insensitively.
+
+                Also computes a hash code of the string in canonical form.
+                """
+
+                def __init__(self, value):
+                    self.canonical = value.lower()
+                    self.hashcode = hash(self.canonical)
+
+                def __eq__(self, other):
+                    if not isinstance(other, cistr):
+                        other = cistr(other)
+                    return self.canonical == other.canonical
+
+                def __hash__(self):
+                    return self.hashcode
+
+        self.assertEqual(cistr('ABC'), 'abc')
+        self.assertEqual('aBc', cistr('ABC'))
+        self.assertEqual(str(cistr('ABC')), 'ABC')
+
+        d = {cistr('one'): 1, cistr('two'): 2, cistr('tHree'): 3}
+        self.assertEqual(d[cistr('one')], 1)
+        self.assertEqual(d[cistr('tWo')], 2)
+        self.assertEqual(d[cistr('THrEE')], 3)
+        self.assertIn(cistr('ONe'), d)
+        self.assertEqual(d.get(cistr('thrEE')), 3)
+
+    def test_classic_comparisons(self):
+        # Testing classic comparisons...
+        with torch._dynamo.error_on_graph_break(False):
+            class classic:
+                pass
+
+        for base in (classic, int, object):
+            with torch._dynamo.error_on_graph_break(False):
+                class C(base):
+                    def __init__(self, value):
+                        self.value = int(value)
+                    def __eq__(self, other):
+                        if isinstance(other, C):
+                            return self.value == other.value
+                        if isinstance(other, int) or isinstance(other, int):
+                            return self.value == other
+                        return NotImplemented
+                    def __ne__(self, other):
+                        if isinstance(other, C):
+                            return self.value != other.value
+                        if isinstance(other, int) or isinstance(other, int):
+                            return self.value != other
+                        return NotImplemented
+                    def __lt__(self, other):
+                        if isinstance(other, C):
+                            return self.value < other.value
+                        if isinstance(other, int) or isinstance(other, int):
+                            return self.value < other
+                        return NotImplemented
+                    def __le__(self, other):
+                        if isinstance(other, C):
+                            return self.value <= other.value
+                        if isinstance(other, int) or isinstance(other, int):
+                            return self.value <= other
+                        return NotImplemented
+                    def __gt__(self, other):
+                        if isinstance(other, C):
+                            return self.value > other.value
+                        if isinstance(other, int) or isinstance(other, int):
+                            return self.value > other
+                        return NotImplemented
+                    def __ge__(self, other):
+                        if isinstance(other, C):
+                            return self.value >= other.value
+                        if isinstance(other, int) or isinstance(other, int):
+                            return self.value >= other
+                        return NotImplemented
+
+            c1 = C(1)
+            c2 = C(2)
+            c3 = C(3)
+            self.assertEqual(c1, 1)
+            c = {1: c1, 2: c2, 3: c3}
+            for x in 1, 2, 3:
+                for y in 1, 2, 3:
+                    for op in "<", "<=", "==", "!=", ">", ">=":
+                        self.assertEqual(eval("c[x] %s c[y]" % op),
+                                     eval("x %s y" % op),
+                                     "x=%d, y=%d" % (x, y))
+                        self.assertEqual(eval("c[x] %s y" % op),
+                                     eval("x %s y" % op),
+                                     "x=%d, y=%d" % (x, y))
+                        self.assertEqual(eval("x %s c[y]" % op),
+                                     eval("x %s y" % op),
+                                     "x=%d, y=%d" % (x, y))
+
+    def test_rich_comparisons(self):
+        # Testing rich comparisons...
+        with torch._dynamo.error_on_graph_break(False):
+            class Z(complex):
+                pass
+        z = Z(1)
+        self.assertEqual(z, 1+0j)
+        self.assertEqual(1+0j, z)
+        with torch._dynamo.error_on_graph_break(False):
+            class ZZ(complex):
+                def __eq__(self, other):
+                    try:
+                        return abs(self - other) <= 1e-6
+                    except:
+                        return NotImplemented
+        zz = ZZ(1.0000003)
+        self.assertEqual(zz, 1+0j)
+        self.assertEqual(1+0j, zz)
+
+        with torch._dynamo.error_on_graph_break(False):
+            class classic:
+                pass
+        for base in (classic, int, object, list):
+            with torch._dynamo.error_on_graph_break(False):
+                class C(base):
+                    def __init__(self, value):
+                        self.value = int(value)
+                    def __cmp__(self_, other):
+                        self.fail("shouldn't call __cmp__")
+                    def __eq__(self, other):
+                        if isinstance(other, C):
+                            return self.value == other.value
+                        if isinstance(other, int) or isinstance(other, int):
+                            return self.value == other
+                        return NotImplemented
+                    def __ne__(self, other):
+                        if isinstance(other, C):
+                            return self.value != other.value
+                        if isinstance(other, int) or isinstance(other, int):
+                            return self.value != other
+                        return NotImplemented
+                    def __lt__(self, other):
+                        if isinstance(other, C):
+                            return self.value < other.value
+                        if isinstance(other, int) or isinstance(other, int):
+                            return self.value < other
+                        return NotImplemented
+                    def __le__(self, other):
+                        if isinstance(other, C):
+                            return self.value <= other.value
+                        if isinstance(other, int) or isinstance(other, int):
+                            return self.value <= other
+                        return NotImplemented
+                    def __gt__(self, other):
+                        if isinstance(other, C):
+                            return self.value > other.value
+                        if isinstance(other, int) or isinstance(other, int):
+                            return self.value > other
+                        return NotImplemented
+                    def __ge__(self, other):
+                        if isinstance(other, C):
+                            return self.value >= other.value
+                        if isinstance(other, int) or isinstance(other, int):
+                            return self.value >= other
+                        return NotImplemented
+            c1 = C(1)
+            c2 = C(2)
+            c3 = C(3)
+            self.assertEqual(c1, 1)
+            c = {1: c1, 2: c2, 3: c3}
+            for x in 1, 2, 3:
+                for y in 1, 2, 3:
+                    for op in "<", "<=", "==", "!=", ">", ">=":
+                        self.assertEqual(eval("c[x] %s c[y]" % op),
+                                         eval("x %s y" % op),
+                                         "x=%d, y=%d" % (x, y))
+                        self.assertEqual(eval("c[x] %s y" % op),
+                                         eval("x %s y" % op),
+                                         "x=%d, y=%d" % (x, y))
+                        self.assertEqual(eval("x %s c[y]" % op),
+                                         eval("x %s y" % op),
+                                         "x=%d, y=%d" % (x, y))
+
+    def test_descrdoc(self):
+        # Testing descriptor doc strings...
+        from _io import FileIO
+        def check(descr, what):
+            self.assertEqual(descr.__doc__, what)
+        check(FileIO.closed, "True if the file is closed") # getset descriptor
+        check(complex.real, "the real part of a complex number") # member descriptor
+
+    def test_doc_descriptor(self):
+        # Testing __doc__ descriptor...
+        # SF bug 542984
+        with torch._dynamo.error_on_graph_break(False):
+            class DocDescr(object):
+                def __get__(self, object, otype):
+                    if object:
+                        object = object.__class__.__name__ + ' instance'
+                    if otype:
+                        otype = otype.__name__
+                    return 'object=%s; type=%s' % (object, otype)
+            class NewClass:
+                __doc__ = DocDescr()
+        self.assertEqual(NewClass.__doc__, 'object=None; type=NewClass')
+        self.assertEqual(NewClass().__doc__, 'object=NewClass instance; type=NewClass')
+
+    def test_set_class(self):
+        # Testing __class__ assignment...
+        with torch._dynamo.error_on_graph_break(False):
+            class C(object): pass
+            class D(object): pass
+            class E(object): pass
+            class F(D, E): pass
+        for cls in C, D, E, F:
+            for cls2 in C, D, E, F:
+                x = cls()
+                x.__class__ = cls2
+                self.assertIs(x.__class__, cls2)
+                x.__class__ = cls
+                self.assertIs(x.__class__, cls)
+        def cant(x, C):
+            try:
+                x.__class__ = C
+            except TypeError:
+                pass
+            else:
+                self.fail("shouldn't allow %r.__class__ = %r" % (x, C))
+            try:
+                delattr(x, "__class__")
+            except (TypeError, AttributeError):
+                pass
+            else:
+                self.fail("shouldn't allow del %r.__class__" % x)
+        cant(C(), list)
+        cant(list(), C)
+        cant(C(), 1)
+        cant(C(), object)
+        cant(object(), list)
+        cant(list(), object)
+        with torch._dynamo.error_on_graph_break(False):
+            class Int(int): __slots__ = []
+        cant(True, int)
+        cant(2, bool)
+        o = object()
+        cant(o, int)
+        cant(o, type(None))
+        del o
+        with torch._dynamo.error_on_graph_break(False):
+            class G(object):
+                __slots__ = ["a", "b"]
+            class H(object):
+                __slots__ = ["b", "a"]
+            class I(object):
+                __slots__ = ["a", "b"]
+            class J(object):
+                __slots__ = ["c", "b"]
+            class K(object):
+                __slots__ = ["a", "b", "d"]
+            class L(H):
+                __slots__ = ["e"]
+            class M(I):
+                __slots__ = ["e"]
+            class N(J):
+                __slots__ = ["__weakref__"]
+            class P(J):
+                __slots__ = ["__dict__"]
+            class Q(J):
+                pass
+            class R(J):
+                __slots__ = ["__dict__", "__weakref__"]
+
+        for cls, cls2 in ((G, H), (G, I), (I, H), (Q, R), (R, Q)):
+            x = cls()
+            x.a = 1
+            x.__class__ = cls2
+            self.assertIs(x.__class__, cls2,
+                   "assigning %r as __class__ for %r silently failed" % (cls2, x))
+            self.assertEqual(x.a, 1)
+            x.__class__ = cls
+            self.assertIs(x.__class__, cls,
+                   "assigning %r as __class__ for %r silently failed" % (cls, x))
+            self.assertEqual(x.a, 1)
+        for cls in G, J, K, L, M, N, P, R, list, Int:
+            for cls2 in G, J, K, L, M, N, P, R, list, Int:
+                if cls is cls2:
+                    continue
+                cant(cls(), cls2)
+
+        # Issue5283: when __class__ changes in __del__, the wrong
+        # type gets DECREF'd.
+        with torch._dynamo.error_on_graph_break(False):
+            class O(object):
+                pass
+            class A(object):
+                def __del__(self):
+                    self.__class__ = O
+        l = [A() for x in range(100)]
+        del l
+
+    def test_set_dict(self):
+        # Testing __dict__ assignment...
+        with torch._dynamo.error_on_graph_break(False):
+            class C(object): pass
+        a = C()
+        a.__dict__ = {'b': 1}
+        self.assertEqual(a.b, 1)
+        def cant(x, dict):
+            try:
+                x.__dict__ = dict
+            except (AttributeError, TypeError):
+                pass
+            else:
+                self.fail("shouldn't allow %r.__dict__ = %r" % (x, dict))
+        cant(a, None)
+        cant(a, [])
+        cant(a, 1)
+        del a.__dict__ # Deleting __dict__ is allowed
+
+        with torch._dynamo.error_on_graph_break(False):
+            class Base(object):
+                pass
+        def verify_dict_readonly(x):
+            """
+            x has to be an instance of a class inheriting from Base.
+            """
+            cant(x, {})
+            try:
+                del x.__dict__
+            except (AttributeError, TypeError):
+                pass
+            else:
+                self.fail("shouldn't allow del %r.__dict__" % x)
+            dict_descr = Base.__dict__["__dict__"]
+            try:
+                dict_descr.__set__(x, {})
+            except (AttributeError, TypeError):
+                pass
+            else:
+                self.fail("dict_descr allowed access to %r's dict" % x)
+
+        # Classes don't allow __dict__ assignment and have readonly dicts
+        with torch._dynamo.error_on_graph_break(False):
+            class Meta1(type, Base):
+                pass
+            class Meta2(Base, type):
+                pass
+            class D(object, metaclass=Meta1):
+                pass
+            class E(object, metaclass=Meta2):
+                pass
+        for cls in C, D, E:
+            verify_dict_readonly(cls)
+            class_dict = cls.__dict__
+            try:
+                class_dict["spam"] = "eggs"
+            except TypeError:
+                pass
+            else:
+                self.fail("%r's __dict__ can be modified" % cls)
+
+        # Modules also disallow __dict__ assignment
+        with torch._dynamo.error_on_graph_break(False):
+            class Module1(types.ModuleType, Base):
+                pass
+            class Module2(Base, types.ModuleType):
+                pass
+        for ModuleType in Module1, Module2:
+            mod = ModuleType("spam")
+            verify_dict_readonly(mod)
+            mod.__dict__["spam"] = "eggs"
+
+        # Exception's __dict__ can be replaced, but not deleted
+        # (at least not any more than regular exception's __dict__ can
+        # be deleted; on CPython it is not the case, whereas on PyPy they
+        # can, just like any other new-style instance's __dict__.)
+        def can_delete_dict(e):
+            try:
+                del e.__dict__
+            except (TypeError, AttributeError):
+                return False
+            else:
+                return True
+        with torch._dynamo.error_on_graph_break(False):
+            class Exception1(Exception, Base):
+                pass
+            class Exception2(Base, Exception):
+                pass
+        for ExceptionType in Exception, Exception1, Exception2:
+            e = ExceptionType()
+            e.__dict__ = {"a": 1}
+            self.assertEqual(e.a, 1)
+            self.assertEqual(can_delete_dict(e), can_delete_dict(ValueError()))
+
+    def test_binary_operator_override(self):
+        # Testing overrides of binary operations...
+        with torch._dynamo.error_on_graph_break(False):
+            class I(int):
+                def __repr__(self):
+                    return "I(%r)" % int(self)
+                def __add__(self, other):
+                    return I(int(self) + int(other))
+                __radd__ = __add__
+                def __pow__(self, other, mod=None):
+                    if mod is None:
+                        return I(pow(int(self), int(other)))
+                    else:
+                        return I(pow(int(self), int(other), int(mod)))
+                def __rpow__(self, other, mod=None):
+                    if mod is None:
+                        return I(pow(int(other), int(self), mod))
+                    else:
+                        return I(pow(int(other), int(self), int(mod)))
+
+        self.assertEqual(repr(I(1) + I(2)), "I(3)")
+        self.assertEqual(repr(I(1) + 2), "I(3)")
+        self.assertEqual(repr(1 + I(2)), "I(3)")
+        self.assertEqual(repr(I(2) ** I(3)), "I(8)")
+        self.assertEqual(repr(2 ** I(3)), "I(8)")
+        self.assertEqual(repr(I(2) ** 3), "I(8)")
+        self.assertEqual(repr(pow(I(2), I(3), I(5))), "I(3)")
+        with torch._dynamo.error_on_graph_break(False):
+            class S(str):
+                def __eq__(self, other):
+                    return self.lower() == other.lower()
+
+    def test_subclass_propagation(self):
+        # Testing propagation of slot functions to subclasses...
+        with torch._dynamo.error_on_graph_break(False):
+            class A(object):
+                pass
+            class B(A):
+                pass
+            class C(A):
+                pass
+            class D(B, C):
+                pass
+        d = D()
+        orig_hash = hash(d) # related to id(d) in platform-dependent ways
+        A.__hash__ = lambda self: 42
+        self.assertEqual(hash(d), 42)
+        C.__hash__ = lambda self: 314
+        self.assertEqual(hash(d), 314)
+        B.__hash__ = lambda self: 144
+        self.assertEqual(hash(d), 144)
+        D.__hash__ = lambda self: 100
+        self.assertEqual(hash(d), 100)
+        D.__hash__ = None
+        self.assertRaises(TypeError, hash, d)
+        del D.__hash__
+        self.assertEqual(hash(d), 144)
+        B.__hash__ = None
+        self.assertRaises(TypeError, hash, d)
+        del B.__hash__
+        self.assertEqual(hash(d), 314)
+        C.__hash__ = None
+        self.assertRaises(TypeError, hash, d)
+        del C.__hash__
+        self.assertEqual(hash(d), 42)
+        A.__hash__ = None
+        self.assertRaises(TypeError, hash, d)
+        del A.__hash__
+        self.assertEqual(hash(d), orig_hash)
+        d.foo = 42
+        d.bar = 42
+        self.assertEqual(d.foo, 42)
+        self.assertEqual(d.bar, 42)
+        def __getattribute__(self, name):
+            if name == "foo":
+                return 24
+            return object.__getattribute__(self, name)
+        A.__getattribute__ = __getattribute__
+        self.assertEqual(d.foo, 24)
+        self.assertEqual(d.bar, 42)
+        def __getattr__(self, name):
+            if name in ("spam", "foo", "bar"):
+                return "hello"
+            raise AttributeError(name)
+        B.__getattr__ = __getattr__
+        self.assertEqual(d.spam, "hello")
+        self.assertEqual(d.foo, 24)
+        self.assertEqual(d.bar, 42)
+        del A.__getattribute__
+        self.assertEqual(d.foo, 42)
+        del d.foo
+        self.assertEqual(d.foo, "hello")
+        self.assertEqual(d.bar, 42)
+        del B.__getattr__
+        try:
+            d.foo
+        except AttributeError:
+            pass
+        else:
+            self.fail("d.foo should be undefined now")
+
+        # Test a nasty bug in recurse_down_subclasses()
+        with torch._dynamo.error_on_graph_break(False):
+            class A(object):
+                pass
+            class B(A):
+                pass
+        del B
+        support.gc_collect()
+        A.__setitem__ = lambda *a: None # crash
+
+    def test_buffer_inheritance(self):
+        # Testing that buffer interface is inherited ...
+
+        import binascii
+        # SF bug [#470040] ParseTuple t# vs subclasses.
+
+        with torch._dynamo.error_on_graph_break(False):
+            class MyBytes(bytes):
+                pass
+        base = b'abc'
+        m = MyBytes(base)
+        # b2a_hex uses the buffer interface to get its argument's value, via
+        # PyArg_ParseTuple 't#' code.
+        self.assertEqual(binascii.b2a_hex(m), binascii.b2a_hex(base))
+
+        with torch._dynamo.error_on_graph_break(False):
+            class MyInt(int):
+                pass
+        m = MyInt(42)
+        try:
+            binascii.b2a_hex(m)
+            self.fail('subclass of int should not have a buffer interface')
+        except TypeError:
+            pass
+
+    def test_str_of_str_subclass(self):
+        # Testing __str__ defined in subclass of str ...
+        import binascii
+
+        with torch._dynamo.error_on_graph_break(False):
+            class octetstring(str):
+                def __str__(self):
+                    return binascii.b2a_hex(self.encode('ascii')).decode("ascii")
+                def __repr__(self):
+                    return self + " repr"
+
+        o = octetstring('A')
+        self.assertEqual(type(o), octetstring)
+        self.assertEqual(type(str(o)), str)
+        self.assertEqual(type(repr(o)), str)
+        self.assertEqual(ord(o), 0x41)
+        self.assertEqual(str(o), '41')
+        self.assertEqual(repr(o), 'A repr')
+        self.assertEqual(o.__str__(), '41')
+        self.assertEqual(o.__repr__(), 'A repr')
+
+    def test_repr_with_module_str_subclass(self):
+        # gh-98783
+        with torch._dynamo.error_on_graph_break(False):
+            class StrSub(str):
+                pass
+            class Some:
+                pass
+        Some.__module__ = StrSub('example')
+        self.assertIsInstance(repr(Some), str)  # should not crash
+        self.assertIsInstance(repr(Some()), str)  # should not crash
+
+    def test_keyword_arguments(self):
+        # Testing keyword arguments to __init__, __call__...
+        def f(a): return a
+        self.assertEqual(f.__call__(a=42), 42)
+        ba = bytearray()
+        bytearray.__init__(ba, 'abc\xbd\u20ac',
+                           encoding='latin1', errors='replace')
+        self.assertEqual(ba, b'abc\xbd?')
+
+    def test_recursive_call(self):
+        # Testing recursive __call__() by setting to instance of class...
+        with torch._dynamo.error_on_graph_break(False):
+            class A(object):
+                pass
+
+        A.__call__ = A()
+        with self.assertRaises(RecursionError):
+            A()()
+
+    def test_delete_hook(self):
+        # Testing __del__ hook...
+        log = []
+        with torch._dynamo.error_on_graph_break(False):
+            class C(object):
+                def __del__(self):
+                    log.append(1)
+        c = C()
+        self.assertEqual(log, [])
+        del c
+        support.gc_collect()
+        self.assertEqual(log, [1])
+
+        with torch._dynamo.error_on_graph_break(False):
+            class D(object): pass
+        d = D()
+        try: del d[0]
+        except TypeError: pass
+        else: self.fail("invalid del() didn't raise TypeError")
+
+    def test_hash_inheritance(self):
+        # Testing hash of mutable subclasses...
+
+        with torch._dynamo.error_on_graph_break(False):
+            class mydict(dict):
+                pass
+        d = mydict()
+        try:
+            hash(d)
+        except TypeError:
+            pass
+        else:
+            self.fail("hash() of dict subclass should fail")
+
+        with torch._dynamo.error_on_graph_break(False):
+            class mylist(list):
+                pass
+        d = mylist()
+        try:
+            hash(d)
+        except TypeError:
+            pass
+        else:
+            self.fail("hash() of list subclass should fail")
+
+    def test_str_operations(self):
+        try: 'a' + 5
+        except TypeError: pass
+        else: self.fail("'' + 5 doesn't raise TypeError")
+
+        try: ''.split('')
+        except ValueError: pass
+        else: self.fail("''.split('') doesn't raise ValueError")
+
+        try: ''.join([0])
+        except TypeError: pass
+        else: self.fail("''.join([0]) doesn't raise TypeError")
+
+        try: ''.rindex('5')
+        except ValueError: pass
+        else: self.fail("''.rindex('5') doesn't raise ValueError")
+
+        try: '%(n)s' % None
+        except TypeError: pass
+        else: self.fail("'%(n)s' % None doesn't raise TypeError")
+
+        try: '%(n' % {}
+        except ValueError: pass
+        else: self.fail("'%(n' % {} '' doesn't raise ValueError")
+
+        try: '%*s' % ('abc')
+        except TypeError: pass
+        else: self.fail("'%*s' % ('abc') doesn't raise TypeError")
+
+        try: '%*.*s' % ('abc', 5)
+        except TypeError: pass
+        else: self.fail("'%*.*s' % ('abc', 5) doesn't raise TypeError")
+
+        try: '%s' % (1, 2)
+        except TypeError: pass
+        else: self.fail("'%s' % (1, 2) doesn't raise TypeError")
+
+        try: '%' % None
+        except ValueError: pass
+        else: self.fail("'%' % None doesn't raise ValueError")
+
+        self.assertEqual('534253'.isdigit(), 1)
+        self.assertEqual('534253x'.isdigit(), 0)
+        self.assertEqual('%c' % 5, '\x05')
+        self.assertEqual('%c' % '5', '5')
+
+    def test_deepcopy_recursive(self):
+        # Testing deepcopy of recursive objects...
+        with torch._dynamo.error_on_graph_break(False):
+            class Node:
+                pass
+        a = Node()
+        b = Node()
+        a.b = b
+        b.a = a
+        z = deepcopy(a) # This blew up before
+
+    def test_uninitialized_modules(self):
+        # Testing uninitialized module objects...
+        from types import ModuleType as M
+        m = M.__new__(M)
+        str(m)
+        self.assertNotHasAttr(m, "__name__")
+        self.assertNotHasAttr(m, "__file__")
+        self.assertNotHasAttr(m, "foo")
+        self.assertFalse(m.__dict__)   # None or {} are both reasonable answers
+        m.foo = 1
+        self.assertEqual(m.__dict__, {"foo": 1})
+
+    def test_funny_new(self):
+        # Testing __new__ returning something unexpected...
+        with torch._dynamo.error_on_graph_break(False):
+            class C(object):
+                def __new__(cls, arg):
+                    if isinstance(arg, str): return [1, 2, 3]
+                    elif isinstance(arg, int): return object.__new__(D)
+                    else: return object.__new__(cls)
+            class D(C):
+                def __init__(self, arg):
+                    self.foo = arg
+        self.assertEqual(C("1"), [1, 2, 3])
+        self.assertEqual(D("1"), [1, 2, 3])
+        d = D(None)
+        self.assertEqual(d.foo, None)
+        d = C(1)
+        self.assertIsInstance(d, D)
+        self.assertEqual(d.foo, 1)
+        d = D(1)
+        self.assertIsInstance(d, D)
+        self.assertEqual(d.foo, 1)
+
+        with torch._dynamo.error_on_graph_break(False):
+            class C(object):
+                @staticmethod
+                def __new__(*args):
+                    return args
+        self.assertEqual(C(1, 2), (C, 1, 2))
+        with torch._dynamo.error_on_graph_break(False):
+            class D(C):
+                pass
+        self.assertEqual(D(1, 2), (D, 1, 2))
+
+        with torch._dynamo.error_on_graph_break(False):
+            class C(object):
+                @classmethod
+                def __new__(*args):
+                    return args
+        self.assertEqual(C(1, 2), (C, C, 1, 2))
+        with torch._dynamo.error_on_graph_break(False):
+            class D(C):
+                pass
+        self.assertEqual(D(1, 2), (D, D, 1, 2))
+
+    def test_imul_bug(self):
+        # Testing for __imul__ problems...
+        # SF bug 544647
+        with torch._dynamo.error_on_graph_break(False):
+            class C(object):
+                def __imul__(self, other):
+                    return (self, other)
+        x = C()
+        y = x
+        y *= 1.0
+        self.assertEqual(y, (x, 1.0))
+        y = x
+        y *= 2
+        self.assertEqual(y, (x, 2))
+        y = x
+        y *= 3
+        self.assertEqual(y, (x, 3))
+        y = x
+        y *= 1<<100
+        self.assertEqual(y, (x, 1<<100))
+        y = x
+        y *= None
+        self.assertEqual(y, (x, None))
+        y = x
+        y *= "foo"
+        self.assertEqual(y, (x, "foo"))
+
+    def test_copy_setstate(self):
+        # Testing that copy.*copy() correctly uses __setstate__...
+        import copy
+        with torch._dynamo.error_on_graph_break(False):
+            class C(object):
+                def __init__(self, foo=None):
+                    self.foo = foo
+                    self.__foo = foo
+                def setfoo(self, foo=None):
+                    self.foo = foo
+                def getfoo(self):
+                    return self.__foo
+                def __getstate__(self):
+                    return [self.foo]
+                def __setstate__(self_, lst):
+                    self.assertEqual(len(lst), 1)
+                    self_.__foo = self_.foo = lst[0]
+        a = C(42)
+        a.setfoo(24)
+        self.assertEqual(a.foo, 24)
+        self.assertEqual(a.getfoo(), 42)
+        b = copy.copy(a)
+        self.assertEqual(b.foo, 24)
+        self.assertEqual(b.getfoo(), 24)
+        b = copy.deepcopy(a)
+        self.assertEqual(b.foo, 24)
+        self.assertEqual(b.getfoo(), 24)
+
+    def test_slices(self):
+        # Testing cases with slices and overridden __getitem__ ...
+
+        # Strings
+        self.assertEqual("hello"[:4], "hell")
+        self.assertEqual("hello"[slice(4)], "hell")
+        self.assertEqual(str.__getitem__("hello", slice(4)), "hell")
+        with torch._dynamo.error_on_graph_break(False):
+            class S(str):
+                def __getitem__(self, x):
+                    return str.__getitem__(self, x)
+        self.assertEqual(S("hello")[:4], "hell")
+        self.assertEqual(S("hello")[slice(4)], "hell")
+        self.assertEqual(S("hello").__getitem__(slice(4)), "hell")
+        # Tuples
+        self.assertEqual((1,2,3)[:2], (1,2))
+        self.assertEqual((1,2,3)[slice(2)], (1,2))
+        self.assertEqual(tuple.__getitem__((1,2,3), slice(2)), (1,2))
+        with torch._dynamo.error_on_graph_break(False):
+            class T(tuple):
+                def __getitem__(self, x):
+                    return tuple.__getitem__(self, x)
+        self.assertEqual(T((1,2,3))[:2], (1,2))
+        self.assertEqual(T((1,2,3))[slice(2)], (1,2))
+        self.assertEqual(T((1,2,3)).__getitem__(slice(2)), (1,2))
+        # Lists
+        self.assertEqual([1,2,3][:2], [1,2])
+        self.assertEqual([1,2,3][slice(2)], [1,2])
+        self.assertEqual(list.__getitem__([1,2,3], slice(2)), [1,2])
+        with torch._dynamo.error_on_graph_break(False):
+            class L(list):
+                def __getitem__(self, x):
+                    return list.__getitem__(self, x)
+        self.assertEqual(L([1,2,3])[:2], [1,2])
+        self.assertEqual(L([1,2,3])[slice(2)], [1,2])
+        self.assertEqual(L([1,2,3]).__getitem__(slice(2)), [1,2])
+        # Now do lists and __setitem__
+        a = L([1,2,3])
+        a[slice(1, 3)] = [3,2]
+        self.assertEqual(a, [1,3,2])
+        a[slice(0, 2, 1)] = [3,1]
+        self.assertEqual(a, [3,1,2])
+        a.__setitem__(slice(1, 3), [2,1])
+        self.assertEqual(a, [3,2,1])
+        a.__setitem__(slice(0, 2, 1), [2,3])
+        self.assertEqual(a, [2,3,1])
+
+    def test_subtype_resurrection(self):
+        # Testing resurrection of new-style instance...
+
+        with torch._dynamo.error_on_graph_break(False):
+            class C(object):
+                container = []
+
+                def __del__(self):
+                    # resurrect the instance
+                    C.container.append(self)
+
+            c = C()
+        c.attr = 42
+
+        # The most interesting thing here is whether this blows up, due to
+        # flawed GC tracking logic in typeobject.c's call_finalizer() (a 2.2.1
+        # bug).
+        del c
+
+        support.gc_collect()
+        self.assertEqual(len(C.container), 1)
+
+        # Make c mortal again, so that the test framework with -l doesn't report
+        # it as a leak.
+        del C.__del__
+
+    def test_slots_trash(self):
+        # Testing slot trash...
+        # Deallocating deeply nested slotted trash caused stack overflows
+        class trash(object):
+            __slots__ = ['x']
+            def __init__(self, x):
+                self.x = x
+        o = None
+        for i in range(50000):
+            o = trash(o)
+        del o
+
+    def test_slots_multiple_inheritance(self):
+        # SF bug 575229, multiple inheritance w/ slots dumps core
+        with torch._dynamo.error_on_graph_break(False):
+            class A(object):
+                __slots__=()
+            class B(object):
+                pass
+            class C(A,B) :
+                __slots__=()
+        if support.check_impl_detail():
+            self.assertEqual(C.__basicsize__, B.__basicsize__)
+        self.assertHasAttr(C, '__dict__')
+        self.assertHasAttr(C, '__weakref__')
+        C().x = 2
+
+    def test_rmul(self):
+        # Testing correct invocation of __rmul__...
+        # SF patch 592646
+        with torch._dynamo.error_on_graph_break(False):
+            class C(object):
+                def __mul__(self, other):
+                    return "mul"
+                def __rmul__(self, other):
+                    return "rmul"
+        a = C()
+        self.assertEqual(a*2, "mul")
+        self.assertEqual(a*2.2, "mul")
+        self.assertEqual(2*a, "rmul")
+        self.assertEqual(2.2*a, "rmul")
+
+    def test_ipow(self):
+        # Testing correct invocation of __ipow__...
+        # [SF bug 620179]
+        with torch._dynamo.error_on_graph_break(False):
+            class C(object):
+                def __ipow__(self, other):
+                    pass
+        a = C()
+        a **= 2
+
+    def test_ipow_returns_not_implemented(self):
+        with torch._dynamo.error_on_graph_break(False):
+            class A:
+                def __ipow__(self, other):
+                    return NotImplemented
+
+            class B(A):
+                def __rpow__(self, other):
+                    return 1
+
+            class C(A):
+                def __pow__(self, other):
+                    return 2
+        a = A()
+        b = B()
+        c = C()
+
+        a **= b
+        self.assertEqual(a, 1)
+
+        c **= b
+        self.assertEqual(c, 2)
+
+    def test_no_ipow(self):
+        with torch._dynamo.error_on_graph_break(False):
+            class B:
+                def __rpow__(self, other):
+                    return 1
+
+        a = object()
+        b = B()
+        a **= b
+        self.assertEqual(a, 1)
+
+    def test_ipow_exception_text(self):
+        x = None
+        with self.assertRaises(TypeError) as cm:
+            x **= 2
+        self.assertIn('unsupported operand type(s) for **=', str(cm.exception))
+
+        with self.assertRaises(TypeError) as cm:
+            y = x ** 2
+        self.assertIn('unsupported operand type(s) for **', str(cm.exception))
+
+    def test_mutable_bases(self):
+        # Testing mutable bases...
+
+        # stuff that should work:
+        with torch._dynamo.error_on_graph_break(False):
+            class C(object):
+                pass
+            class C2(object):
+                def __getattribute__(self, attr):
+                    if attr == 'a':
+                        return 2
+                    else:
+                        return super(C2, self).__getattribute__(attr)
+                def meth(self):
+                    return 1
+            class D(C):
+                pass
+            class E(D):
+                pass
+        d = D()
+        e = E()
+        D.__bases__ = (C,)
+        D.__bases__ = (C2,)
+        self.assertEqual(d.meth(), 1)
+        self.assertEqual(e.meth(), 1)
+        self.assertEqual(d.a, 2)
+        self.assertEqual(e.a, 2)
+        self.assertEqual(C2.__subclasses__(), [D])
+
+        try:
+            del D.__bases__
+        except (TypeError, AttributeError):
+            pass
+        else:
+            self.fail("shouldn't be able to delete .__bases__")
+
+        try:
+            D.__bases__ = ()
+        except TypeError as msg:
+            if str(msg) == "a new-style class can't have only classic bases":
+                self.fail("wrong error message for .__bases__ = ()")
+        else:
+            self.fail("shouldn't be able to set .__bases__ to ()")
+
+        try:
+            D.__bases__ = (D,)
+        except TypeError:
+            pass
+        else:
+            # actually, we'll have crashed by here...
+            self.fail("shouldn't be able to create inheritance cycles")
+
+        try:
+            D.__bases__ = (C, C)
+        except TypeError:
+            pass
+        else:
+            self.fail("didn't detect repeated base classes")
+
+        try:
+            D.__bases__ = (E,)
+        except TypeError:
+            pass
+        else:
+            self.fail("shouldn't be able to create inheritance cycles")
+
+    def test_builtin_bases(self):
+        # Make sure all the builtin types can have their base queried without
+        # segfaulting. See issue #5787.
+        builtin_types = [tp for tp in builtins.__dict__.values()
+                         if isinstance(tp, type)]
+        for tp in builtin_types:
+            object.__getattribute__(tp, "__bases__")
+            if tp is not object:
+                if tp is ExceptionGroup:
+                    num_bases = 2
+                else:
+                    num_bases = 1
+                self.assertEqual(len(tp.__bases__), num_bases, tp)
+
+        with torch._dynamo.error_on_graph_break(False):
+            class L(list):
+                pass
+
+            class C(object):
+                pass
+
+            class D(C):
+                pass
+
+        try:
+            L.__bases__ = (dict,)
+        except TypeError:
+            pass
+        else:
+            self.fail("shouldn't turn list subclass into dict subclass")
+
+        try:
+            list.__bases__ = (dict,)
+        except TypeError:
+            pass
+        else:
+            self.fail("shouldn't be able to assign to list.__bases__")
+
+        try:
+            D.__bases__ = (C, list)
+        except TypeError:
+            pass
+        else:
+            self.fail("best_base calculation found wanting")
+
+    def test_unsubclassable_types(self):
+        with self.assertRaises(TypeError):
+            with torch._dynamo.error_on_graph_break(False):
+                class X(type(None)):
+                    pass
+        with self.assertRaises(TypeError):
+            with torch._dynamo.error_on_graph_break(False):
+                class X(object, type(None)):
+                    pass
+        with self.assertRaises(TypeError):
+            with torch._dynamo.error_on_graph_break(False):
+                class X(type(None), object):
+                    pass
+        with torch._dynamo.error_on_graph_break(False):
+            class O(object):
+                pass
+        with self.assertRaises(TypeError):
+            with torch._dynamo.error_on_graph_break(False):
+                class X(O, type(None)):
+                    pass
+        with self.assertRaises(TypeError):
+            with torch._dynamo.error_on_graph_break(False):
+                class X(type(None), O):
+                    pass
+
+        with torch._dynamo.error_on_graph_break(False):
+            class X(object):
+                pass
+        with self.assertRaises(TypeError):
+            X.__bases__ = type(None),
+        with self.assertRaises(TypeError):
+            X.__bases__ = object, type(None)
+        with self.assertRaises(TypeError):
+            X.__bases__ = type(None), object
+        with self.assertRaises(TypeError):
+            X.__bases__ = O, type(None)
+        with self.assertRaises(TypeError):
+            X.__bases__ = type(None), O
+
+    def test_mutable_bases_with_failing_mro(self):
+        with torch._dynamo.error_on_graph_break(False):
+            # Testing mutable bases with failing mro...
+            class WorkOnce(type):
+                def __new__(self, name, bases, ns):
+                    self.flag = 0
+                    return super(WorkOnce, self).__new__(WorkOnce, name, bases, ns)
+                def mro(self):
+                    if self.flag > 0:
+                        raise RuntimeError("bozo")
+                    else:
+                        self.flag += 1
+                        return type.mro(self)
+
+            class WorkAlways(type):
+                def mro(self):
+                    # this is here to make sure that .mro()s aren't called
+                    # with an exception set (which was possible at one point).
+                    # An error message will be printed in a debug build.
+                    # What's a good way to test for this?
+                    return type.mro(self)
+
+            class C(object):
+                pass
+
+            class C2(object):
+                pass
+
+            class D(C):
+                pass
+
+            class E(D):
+                pass
+
+            class F(D, metaclass=WorkOnce):
+                pass
+
+            class G(D, metaclass=WorkAlways):
+                pass
+
+        # Immediate subclasses have their mro's adjusted in alphabetical
+        # order, so E's will get adjusted before adjusting F's fails.  We
+        # check here that E's gets restored.
+
+        E_mro_before = E.__mro__
+        D_mro_before = D.__mro__
+
+        try:
+            D.__bases__ = (C2,)
+        except RuntimeError:
+            self.assertEqual(E.__mro__, E_mro_before)
+            self.assertEqual(D.__mro__, D_mro_before)
+        else:
+            self.fail("exception not propagated")
+
+    def test_mutable_bases_catch_mro_conflict(self):
+        with torch._dynamo.error_on_graph_break(False):
+            # Testing mutable bases catch mro conflict...
+            class A(object):
+                pass
+
+            class B(object):
+                pass
+
+            class C(A, B):
+                pass
+
+            class D(A, B):
+                pass
+
+            class E(C, D):
+                pass
+
+        try:
+            C.__bases__ = (B, A)
+        except TypeError:
+            pass
+        else:
+            self.fail("didn't catch MRO conflict")
+
+    def test_mutable_names(self):
+        # Testing mutable names...
+        with torch._dynamo.error_on_graph_break(False):
+            class C(object):
+                pass
+
+        # C.__module__ could be 'test_descr' or '__main__'
+        mod = C.__module__
+
+        C.__name__ = 'D'
+        self.assertEqual((C.__module__, C.__name__), (mod, 'D'))
+
+        C.__name__ = 'D.E'
+        self.assertEqual((C.__module__, C.__name__), (mod, 'D.E'))
+
+    def test_evil_type_name(self):
+        # A badly placed Py_DECREF in type_set_name led to arbitrary code
+        # execution while the type structure was not in a sane state, and a
+        # possible segmentation fault as a result.  See bug #16447.
+        class Nasty(str):
+            def __del__(self):
+                C.__name__ = "other"
+
+        class C:
+            pass
+
+        C.__name__ = Nasty("abc")
+        C.__name__ = "normal"
+
+    def test_subclass_right_op(self):
+        # Testing correct dispatch of subclass overloading __r<op>__...
+
+        # This code tests various cases where right-dispatch of a subclass
+        # should be preferred over left-dispatch of a base class.
+
+        # Case 1: subclass of int; this tests code in abstract.c::binary_op1()
+
+        class B(int):
+            def __floordiv__(self, other):
+                return "B.__floordiv__"
+            def __rfloordiv__(self, other):
+                return "B.__rfloordiv__"
+
+        self.assertEqual(B(1) // 1, "B.__floordiv__")
+        self.assertEqual(1 // B(1), "B.__rfloordiv__")
+
+        # Case 2: subclass of object; this is just the baseline for case 3
+
+        class C(object):
+            def __floordiv__(self, other):
+                return "C.__floordiv__"
+            def __rfloordiv__(self, other):
+                return "C.__rfloordiv__"
+
+        self.assertEqual(C() // 1, "C.__floordiv__")
+        self.assertEqual(1 // C(), "C.__rfloordiv__")
+
+        # Case 3: subclass of new-style class; here it gets interesting
+
+        class D(C):
+            def __floordiv__(self, other):
+                return "D.__floordiv__"
+            def __rfloordiv__(self, other):
+                return "D.__rfloordiv__"
+
+        self.assertEqual(D() // C(), "D.__floordiv__")
+        self.assertEqual(C() // D(), "D.__rfloordiv__")
+
+        # Case 4: this didn't work right in 2.2.2 and 2.3a1
+
+        class E(C):
+            pass
+
+        self.assertEqual(E.__rfloordiv__, C.__rfloordiv__)
+
+        self.assertEqual(E() // 1, "C.__floordiv__")
+        self.assertEqual(1 // E(), "C.__rfloordiv__")
+        self.assertEqual(E() // C(), "C.__floordiv__")
+        self.assertEqual(C() // E(), "C.__floordiv__") # This one would fail
+
+    @support.impl_detail("testing an internal kind of method object")
+    def test_meth_class_get(self):
+        # Testing __get__ method of METH_CLASS C methods...
+        # Full coverage of descrobject.c::classmethod_get()
+
+        # Baseline
+        arg = [1, 2, 3]
+        res = {1: None, 2: None, 3: None}
+        self.assertEqual(dict.fromkeys(arg), res)
+        self.assertEqual({}.fromkeys(arg), res)
+
+        # Now get the descriptor
+        descr = dict.__dict__["fromkeys"]
+
+        # More baseline using the descriptor directly
+        self.assertEqual(descr.__get__(None, dict)(arg), res)
+        self.assertEqual(descr.__get__({})(arg), res)
+
+        # Now check various error cases
+        try:
+            descr.__get__(None, None)
+        except TypeError:
+            pass
+        else:
+            self.fail("shouldn't have allowed descr.__get__(None, None)")
+        try:
+            descr.__get__(42)
+        except TypeError:
+            pass
+        else:
+            self.fail("shouldn't have allowed descr.__get__(42)")
+        try:
+            descr.__get__(None, 42)
+        except TypeError:
+            pass
+        else:
+            self.fail("shouldn't have allowed descr.__get__(None, 42)")
+        try:
+            descr.__get__(None, int)
+        except TypeError:
+            pass
+        else:
+            self.fail("shouldn't have allowed descr.__get__(None, int)")
+
+    def test_isinst_isclass(self):
+        with torch._dynamo.error_on_graph_break(False):
+            # Testing proxy isinstance() and isclass()...
+            class Proxy(object):
+                def __init__(self, obj):
+                    self.__obj = obj
+                def __getattribute__(self, name):
+                    if name.startswith("_Proxy__"):
+                        return object.__getattribute__(self, name)
+                    else:
+                        return getattr(self.__obj, name)
+            # Test with a classic class
+            class C:
+                pass
+            a = C()
+            pa = Proxy(a)
+        self.assertIsInstance(a, C)  # Baseline
+        self.assertIsInstance(pa, C) # Test
+        with torch._dynamo.error_on_graph_break(False):
+            # Test with a classic subclass
+            class D(C):
+                pass
+            a = D()
+            pa = Proxy(a)
+        self.assertIsInstance(a, C)  # Baseline
+        self.assertIsInstance(pa, C) # Test
+        with torch._dynamo.error_on_graph_break(False):
+            # Test with a new-style class
+            class C(object):
+                pass
+            a = C()
+            pa = Proxy(a)
+        self.assertIsInstance(a, C)  # Baseline
+        self.assertIsInstance(pa, C) # Test
+        # Test with a new-style subclass
+        with torch._dynamo.error_on_graph_break(False):
+            class D(C):
+                pass
+            a = D()
+            pa = Proxy(a)
+        self.assertIsInstance(a, C)  # Baseline
+        self.assertIsInstance(pa, C) # Test
+
+    def test_proxy_super(self):
+        # Testing super() for a proxy object...
+        with torch._dynamo.error_on_graph_break(False):
+            class Proxy(object):
+                def __init__(self, obj):
+                    self.__obj = obj
+                def __getattribute__(self, name):
+                    if name.startswith("_Proxy__"):
+                        return object.__getattribute__(self, name)
+                    else:
+                        return getattr(self.__obj, name)
+
+            class B(object):
+                def f(self):
+                    return "B.f"
+
+            class C(B):
+                def f(self):
+                    return super(C, self).f() + "->C.f"
+
+            obj = C()
+            p = Proxy(obj)
+        self.assertEqual(C.__dict__["f"](p), "B.f->C.f")
+
+    def test_carloverre(self):
+        # Testing prohibition of Carlo Verre's hack...
+        try:
+            object.__setattr__(str, "foo", 42)
+        except TypeError:
+            pass
+        else:
+            self.fail("Carlo Verre __setattr__ succeeded!")
+        try:
+            object.__delattr__(str, "lower")
+        except TypeError:
+            pass
+        else:
+            self.fail("Carlo Verre __delattr__ succeeded!")
+
+    def test_carloverre_multi_inherit_valid(self):
+        with torch._dynamo.error_on_graph_break(False):
+            class A(type):
+                def __setattr__(cls, key, value):
+                    type.__setattr__(cls, key, value)
+
+            class B:
+                pass
+
+            class C(B, A):
+                pass
+
+            obj = C('D', (object,), {})
+        try:
+            obj.test = True
+        except TypeError:
+            self.fail("setattr through direct base types should be legal")
+
+    def test_carloverre_multi_inherit_invalid(self):
+        with torch._dynamo.error_on_graph_break(False):
+            class A(type):
+                def __setattr__(cls, key, value):
+                    object.__setattr__(cls, key, value)  # this should fail!
+
+            class B:
+                pass
+
+            class C(B, A):
+                pass
+
+            obj = C('D', (object,), {})
+        try:
+            obj.test = True
+        except TypeError:
+            pass
+        else:
+            self.fail("setattr through indirect base types should be rejected")
+
+    def test_weakref_segfault(self):
+        # Testing weakref segfault...
+        # SF 742911
+        import weakref
+
+        class Provoker:
+            def __init__(self, referrent):
+                self.ref = weakref.ref(referrent)
+
+            def __del__(self):
+                x = self.ref()
+
+        class Oops(object):
+            pass
+
+        o = Oops()
+        o.whatever = Provoker(o)
+        del o
+
+    @support.requires_resource('cpu')
+    def test_wrapper_segfault(self):
+        # SF 927248: deeply nested wrappers could cause stack overflow
+        f = lambda:None
+        for i in range(1000000):
+            f = f.__call__
+        f = None
+
+    def test_file_fault(self):
+        # Testing sys.stdout is changed in getattr...
+        class StdoutGuard:
+            def __getattr__(self, attr):
+                sys.stdout = sys.__stdout__
+                raise RuntimeError(f"Premature access to sys.stdout.{attr}")
+
+        with redirect_stdout(StdoutGuard()):
+            with self.assertRaises(RuntimeError):
+                print("Oops!")
+
+    def test_vicious_descriptor_nonsense(self):
+        # Testing vicious_descriptor_nonsense...
+
+        # A potential segfault spotted by Thomas Wouters in mail to
+        # python-dev 2003-04-17, turned into an example & fixed by Michael
+        # Hudson just less than four months later...
+
+        class Evil(object):
+            def __hash__(self):
+                return hash('attr')
+            def __eq__(self, other):
+                try:
+                    del C.attr
+                except AttributeError:
+                    # possible race condition
+                    pass
+                return 0
+
+        class Descr(object):
+            def __get__(self, ob, type=None):
+                return 1
+
+        class C(object):
+            attr = Descr()
+
+        c = C()
+        c.__dict__[Evil()] = 0
+
+        self.assertEqual(c.attr, 1)
+        # this makes a crash more likely:
+        support.gc_collect()
+        self.assertNotHasAttr(c, 'attr')
+
+    def test_init(self):
+        # SF 1155938
+        class Foo(object):
+            def __init__(self):
+                return 10
+        try:
+            Foo()
+        except TypeError:
+            pass
+        else:
+            self.fail("did not test __init__() for None return")
+
+    def assertNotOrderable(self, a, b):
+        with self.assertRaises(TypeError):
+            a < b
+        with self.assertRaises(TypeError):
+            a > b
+        with self.assertRaises(TypeError):
+            a <= b
+        with self.assertRaises(TypeError):
+            a >= b
+
+    def test_method_wrapper(self):
+        # Testing method-wrapper objects...
+        # <type 'method-wrapper'> did not support any reflection before 2.5
+        l = []
+        self.assertTrue(l.__add__ == l.__add__)
+        self.assertFalse(l.__add__ != l.__add__)
+        self.assertFalse(l.__add__ == [].__add__)
+        self.assertTrue(l.__add__ != [].__add__)
+        self.assertFalse(l.__add__ == l.__mul__)
+        self.assertTrue(l.__add__ != l.__mul__)
+        self.assertNotOrderable(l.__add__, l.__add__)
+        self.assertEqual(l.__add__.__name__, '__add__')
+        self.assertIs(l.__add__.__self__, l)
+        self.assertIs(l.__add__.__objclass__, list)
+        self.assertEqual(l.__add__.__doc__, list.__add__.__doc__)
+        # hash([].__add__) should not be based on hash([])
+        hash(l.__add__)
+
+    def test_builtin_function_or_method(self):
+        # Not really belonging to test_descr, but introspection and
+        # comparison on <type 'builtin_function_or_method'> seems not
+        # to be tested elsewhere
+        l = []
+        self.assertTrue(l.append == l.append)
+        self.assertFalse(l.append != l.append)
+        self.assertFalse(l.append == [].append)
+        self.assertTrue(l.append != [].append)
+        self.assertFalse(l.append == l.pop)
+        self.assertTrue(l.append != l.pop)
+        self.assertNotOrderable(l.append, l.append)
+        self.assertEqual(l.append.__name__, 'append')
+        self.assertIs(l.append.__self__, l)
+        # self.assertIs(l.append.__objclass__, list) --- could be added?
+        self.assertEqual(l.append.__doc__, list.append.__doc__)
+        # hash([].append) should not be based on hash([])
+        hash(l.append)
+
+    def test_special_unbound_method_types(self):
+        # Testing objects of <type 'wrapper_descriptor'>...
+        self.assertTrue(list.__add__ == list.__add__)
+        self.assertFalse(list.__add__ != list.__add__)
+        self.assertFalse(list.__add__ == list.__mul__)
+        self.assertTrue(list.__add__ != list.__mul__)
+        self.assertNotOrderable(list.__add__, list.__add__)
+        self.assertEqual(list.__add__.__name__, '__add__')
+        self.assertIs(list.__add__.__objclass__, list)
+
+        # Testing objects of <type 'method_descriptor'>...
+        self.assertTrue(list.append == list.append)
+        self.assertFalse(list.append != list.append)
+        self.assertFalse(list.append == list.pop)
+        self.assertTrue(list.append != list.pop)
+        self.assertNotOrderable(list.append, list.append)
+        self.assertEqual(list.append.__name__, 'append')
+        self.assertIs(list.append.__objclass__, list)
+
+    def test_not_implemented(self):
+        # Testing NotImplemented...
+        # all binary methods should be able to return a NotImplemented
+
+        def specialmethod(self, other):
+            return NotImplemented
+
+        def check(expr, x, y):
+            with (
+                self.subTest(expr=expr, x=x, y=y),
+                self.assertRaises(TypeError),
+            ):
+                exec(expr, {'x': x, 'y': y})
+
+        N1 = sys.maxsize + 1    # might trigger OverflowErrors instead of
+                                # TypeErrors
+        N2 = sys.maxsize         # if sizeof(int) < sizeof(long), might trigger
+                                #   ValueErrors instead of TypeErrors
+        for name, expr, iexpr in [
+                ('__add__',      'x + y',                   'x += y'),
+                ('__sub__',      'x - y',                   'x -= y'),
+                ('__mul__',      'x * y',                   'x *= y'),
+                ('__matmul__',   'x @ y',                   'x @= y'),
+                ('__truediv__',  'x / y',                   'x /= y'),
+                ('__floordiv__', 'x // y',                  'x //= y'),
+                ('__mod__',      'x % y',                   'x %= y'),
+                ('__divmod__',   'divmod(x, y)',            None),
+                ('__pow__',      'x ** y',                  'x **= y'),
+                ('__lshift__',   'x << y',                  'x <<= y'),
+                ('__rshift__',   'x >> y',                  'x >>= y'),
+                ('__and__',      'x & y',                   'x &= y'),
+                ('__or__',       'x | y',                   'x |= y'),
+                ('__xor__',      'x ^ y',                   'x ^= y')]:
+            # Defines 'left' magic method:
+            A = type('A', (), {name: specialmethod})
+            a = A()
+            check(expr, a, a)
+            check(expr, a, N1)
+            check(expr, a, N2)
+            # Defines 'right' magic method:
+            rname = '__r' + name[2:]
+            B = type('B', (), {rname: specialmethod})
+            b = B()
+            check(expr, b, b)
+            check(expr, a, b)
+            check(expr, b, a)
+            check(expr, b, N1)
+            check(expr, b, N2)
+            check(expr, N1, b)
+            check(expr, N2, b)
+            if iexpr:
+                check(iexpr, a, a)
+                check(iexpr, a, N1)
+                check(iexpr, a, N2)
+                iname = '__i' + name[2:]
+                C = type('C', (), {iname: specialmethod})
+                c = C()
+                check(iexpr, c, a)
+                check(iexpr, c, N1)
+                check(iexpr, c, N2)
+
+    def test_assign_slice(self):
+        # ceval.c's assign_slice used to check for
+        # tp->tp_as_sequence->sq_slice instead of
+        # tp->tp_as_sequence->sq_ass_slice
+
+        with torch._dynamo.error_on_graph_break(False):
+            class C(object):
+                def __setitem__(self, idx, value):
+                    self.value = value
+
+            c = C()
+        c[1:2] = 3
+        self.assertEqual(c.value, 3)
+
+    def test_set_and_no_get(self):
+        # See
+        # http://mail.python.org/pipermail/python-dev/2010-January/095637.html
+        with torch._dynamo.error_on_graph_break(False):
+            class Descr(object):
+
+                def __init__(self, name):
+                    self.name = name
+
+                def __set__(self, obj, value):
+                    obj.__dict__[self.name] = value
+            descr = Descr("a")
+
+            class X(object):
+                a = descr
+
+        x = X()
+        self.assertIs(x.a, descr)
+        x.a = 42
+        self.assertEqual(x.a, 42)
+        with torch._dynamo.error_on_graph_break(False):
+
+            # Also check type_getattro for correctness.
+            class Meta(type):
+                pass
+            class X(metaclass=Meta):
+                pass
+        X.a = 42
+        Meta.a = Descr("a")
+        self.assertEqual(X.a, 42)
+
+    def test_getattr_hooks(self):
+        # issue 4230
+
+        with torch._dynamo.error_on_graph_break(False):
+            class Descriptor(object):
+                counter = 0
+                def __get__(self, obj, objtype=None):
+                    def getter(name):
+                        self.counter += 1
+                        raise AttributeError(name)
+                    return getter
+
+            descr = Descriptor()
+            class A(object):
+                __getattribute__ = descr
+            class B(object):
+                __getattr__ = descr
+            class C(object):
+                __getattribute__ = descr
+                __getattr__ = descr
+
+        self.assertRaises(AttributeError, getattr, A(), "attr")
+        self.assertEqual(descr.counter, 1)
+        self.assertRaises(AttributeError, getattr, B(), "attr")
+        self.assertEqual(descr.counter, 2)
+        self.assertRaises(AttributeError, getattr, C(), "attr")
+        self.assertEqual(descr.counter, 4)
+
+        with torch._dynamo.error_on_graph_break(False):
+            class EvilGetattribute(object):
+                # This used to segfault
+                def __getattr__(self, name):
+                    raise AttributeError(name)
+                def __getattribute__(self, name):
+                    del EvilGetattribute.__getattr__
+                    for i in range(5):
+                        gc.collect()
+                    raise AttributeError(name)
+
+        self.assertRaises(AttributeError, getattr, EvilGetattribute(), "attr")
+
+    def test_type___getattribute__(self):
+        self.assertRaises(TypeError, type.__getattribute__, list, type)
+
+    def test_abstractmethods(self):
+        # type pretends not to have __abstractmethods__.
+        self.assertRaises(AttributeError, getattr, type, "__abstractmethods__")
+        with torch._dynamo.error_on_graph_break(False):
+            class meta(type):
+                pass
+        self.assertRaises(AttributeError, getattr, meta, "__abstractmethods__")
+        with torch._dynamo.error_on_graph_break(False):
+            class X(object):
+                pass
+        with self.assertRaises(AttributeError):
+            del X.__abstractmethods__
+
+    def test_gh55664(self):
+        # gh-55664: issue a warning when the
+        # __dict__ of a class contains non-string keys
+        with self.assertWarnsRegex(RuntimeWarning, 'MyClass'):
+            MyClass = type('MyClass', (), {1: 2})
+
+        class meta(type):
+            def __new__(mcls, name, bases, ns):
+                ns[1] = 2
+                return super().__new__(mcls, name, bases, ns)
+
+        with self.assertWarnsRegex(RuntimeWarning, 'MyClass'):
+            MyClass = meta('MyClass', (), {})
+
+    def test_proxy_call(self):
+        with torch._dynamo.error_on_graph_break(False):
+            class FakeStr:
+                __class__ = str
+
+        fake_str = FakeStr()
+        # isinstance() reads __class__
+        self.assertIsInstance(fake_str, str)
+
+        # call a method descriptor
+        with self.assertRaises(TypeError):
+            str.split(fake_str)
+
+        # call a slot wrapper descriptor
+        with self.assertRaises(TypeError):
+            str.__add__(fake_str, "abc")
+
+    def test_specialized_method_calls_check_types(self):
+        # https://github.com/python/cpython/issues/92063
+        class Thing:
+            pass
+        thing = Thing()
+        for i in range(20):
+            with self.assertRaises(TypeError):
+                # CALL_METHOD_DESCRIPTOR_FAST_WITH_KEYWORDS
+                list.sort(thing)
+        for i in range(20):
+            with self.assertRaises(TypeError):
+                # CALL_METHOD_DESCRIPTOR_FAST_WITH_KEYWORDS
+                str.split(thing)
+        for i in range(20):
+            with self.assertRaises(TypeError):
+                # CALL_METHOD_DESCRIPTOR_NOARGS
+                str.upper(thing)
+        for i in range(20):
+            with self.assertRaises(TypeError):
+                # CALL_METHOD_DESCRIPTOR_FAST
+                str.strip(thing)
+        from collections import deque
+        for i in range(20):
+            with self.assertRaises(TypeError):
+                # CALL_METHOD_DESCRIPTOR_O
+                deque.append(thing, thing)
+
+    def test_repr_as_str(self):
+        # Issue #11603: crash or infinite loop when rebinding __str__ as
+        # __repr__.
+        class Foo:
+            pass
+        Foo.__repr__ = Foo.__str__
+        foo = Foo()
+        self.assertRaises(RecursionError, str, foo)
+        self.assertRaises(RecursionError, repr, foo)
+
+    def test_mixing_slot_wrappers(self):
+        with torch._dynamo.error_on_graph_break(False):
+            class X(dict):
+                __setattr__ = dict.__setitem__
+                __neg__ = dict.copy
+            x = X()
+        x.y = 42
+        self.assertEqual(x["y"], 42)
+        self.assertEqual(x, -x)
+
+    def test_wrong_class_slot_wrapper(self):
+        # Check bpo-37619: a wrapper descriptor taken from the wrong class
+        # should raise an exception instead of silently being ignored
+        with torch._dynamo.error_on_graph_break(False):
+            class A(int):
+                __eq__ = str.__eq__
+                __add__ = str.__add__
+            a = A()
+        with self.assertRaises(TypeError):
+            a == a
+        with self.assertRaises(TypeError):
+            a + a
+
+    def test_slot_shadows_class_variable(self):
+        with self.assertRaises(ValueError) as cm:
+            class X:
+                __slots__ = ["foo"]
+                foo = None
+        m = str(cm.exception)
+        self.assertEqual("'foo' in __slots__ conflicts with class variable", m)
+
+    def test_set_doc(self):
+        with torch._dynamo.error_on_graph_break(False):
+            class X:
+                "elephant"
+        X.__doc__ = "banana"
+        self.assertEqual(X.__doc__, "banana")
+
+        with self.assertRaises(TypeError) as cm:
+            type(list).__dict__["__doc__"].__set__(list, "blah")
+        self.assertIn("cannot set '__doc__' attribute of immutable type 'list'", str(cm.exception))
+
+        with self.assertRaises(TypeError) as cm:
+            type(X).__dict__["__doc__"].__delete__(X)
+        self.assertIn("cannot delete '__doc__' attribute of immutable type 'X'", str(cm.exception))
+        self.assertEqual(X.__doc__, "banana")
+
+    def test_qualname(self):
+        descriptors = [str.lower, complex.real, float.real, int.__add__]
+        types = ['method', 'member', 'getset', 'wrapper']
+
+        # make sure we have an example of each type of descriptor
+        for d, n in zip(descriptors, types):
+            self.assertEqual(type(d).__name__, n + '_descriptor')
+
+        for d in descriptors:
+            qualname = d.__objclass__.__qualname__ + '.' + d.__name__
+            self.assertEqual(d.__qualname__, qualname)
+
+        self.assertEqual(str.lower.__qualname__, 'str.lower')
+        self.assertEqual(complex.real.__qualname__, 'complex.real')
+        self.assertEqual(float.real.__qualname__, 'float.real')
+        self.assertEqual(int.__add__.__qualname__, 'int.__add__')
+
+        with torch._dynamo.error_on_graph_break(False):
+            class X:
+                pass
+        with self.assertRaises(TypeError):
+            del X.__qualname__
+
+        self.assertRaises(TypeError, type.__dict__['__qualname__'].__set__,
+                          str, 'Oink')
+
+        global Y
+
+        with torch._dynamo.error_on_graph_break(False):
+            class Y:
+                class Inside:
+                    pass
+        self.assertEqual(Y.__qualname__, 'Y')
+        self.assertEqual(Y.Inside.__qualname__, 'Y.Inside')
+
+    def test_qualname_dict(self):
+        ns = {'__qualname__': 'some.name'}
+        tp = type('Foo', (), ns)
+        self.assertEqual(tp.__qualname__, 'some.name')
+        self.assertNotIn('__qualname__', tp.__dict__)
+        self.assertEqual(ns, {'__qualname__': 'some.name'})
+
+        ns = {'__qualname__': 1}
+        self.assertRaises(TypeError, type, 'Foo', (), ns)
+
+    def test_cycle_through_dict(self):
+        # See bug #1469629
+        class X(dict):
+            def __init__(self):
+                dict.__init__(self)
+                self.__dict__ = self
+        x = X()
+        x.attr = 42
+        wr = weakref.ref(x)
+        del x
+        support.gc_collect()
+        self.assertIsNone(wr())
+        for o in gc.get_objects():
+            self.assertIsNot(type(o), X)
+
+    def test_object_new_and_init_with_parameters(self):
+        # See issue #1683368
+        with torch._dynamo.error_on_graph_break(False):
+            class OverrideNeither:
+                pass
+
+        self.assertRaises(TypeError, OverrideNeither, 1)
+        self.assertRaises(TypeError, OverrideNeither, kw=1)
+
+        with torch._dynamo.error_on_graph_break(False):
+            class OverrideNew:
+                def __new__(cls, foo, kw=0, *args, **kwds):
+                    return object.__new__(cls, *args, **kwds)
+            class OverrideInit:
+                def __init__(self, foo, kw=0, *args, **kwargs):
+                    return object.__init__(self, *args, **kwargs)
+            class OverrideBoth(OverrideNew, OverrideInit):
+                pass
+
+        for case in OverrideNew, OverrideInit, OverrideBoth:
+            case(1)
+            case(1, kw=2)
+            self.assertRaises(TypeError, case, 1, 2, 3)
+            self.assertRaises(TypeError, case, 1, 2, foo=3)
+
+    def test_subclassing_does_not_duplicate_dict_descriptors(self):
+        with torch._dynamo.error_on_graph_break(False):
+            class Base:
+                pass
+            class Sub(Base):
+                pass
+        self.assertIn("__dict__", Base.__dict__)
+        self.assertNotIn("__dict__", Sub.__dict__)
+
+    def test_bound_method_repr(self):
+        class Foo:
+            def method(self):
+                pass
+        self.assertRegex(repr(Foo().method),
+            r"<bound method .*Foo\.method of <.*Foo object at .*>>")
+
+
+        class Base:
+            def method(self):
+                pass
+        class Derived1(Base):
+            pass
+        class Derived2(Base):
+            def method(self):
+                pass
+        base = Base()
+        derived1 = Derived1()
+        derived2 = Derived2()
+        super_d2 = super(Derived2, derived2)
+        self.assertRegex(repr(base.method),
+            r"<bound method .*Base\.method of <.*Base object at .*>>")
+        self.assertRegex(repr(derived1.method),
+            r"<bound method .*Base\.method of <.*Derived1 object at .*>>")
+        self.assertRegex(repr(derived2.method),
+            r"<bound method .*Derived2\.method of <.*Derived2 object at .*>>")
+        self.assertRegex(repr(super_d2.method),
+            r"<bound method .*Base\.method of <.*Derived2 object at .*>>")
+
+        class Foo:
+            @classmethod
+            def method(cls):
+                pass
+        foo = Foo()
+        self.assertRegex(repr(foo.method), # access via instance
+            r"<bound method .*Foo\.method of <class '.*Foo'>>")
+        self.assertRegex(repr(Foo.method), # access via the class
+            r"<bound method .*Foo\.method of <class '.*Foo'>>")
+
+
+        class MyCallable:
+            def __call__(self, arg):
+                pass
+        func = MyCallable() # func has no __name__ or __qualname__ attributes
+        instance = object()
+        method = types.MethodType(func, instance)
+        self.assertRegex(repr(method),
+            r"<bound method \? of <object object at .*>>")
+        func.__name__ = "name"
+        self.assertRegex(repr(method),
+            r"<bound method name of <object object at .*>>")
+        func.__qualname__ = "qualname"
+        self.assertRegex(repr(method),
+            r"<bound method qualname of <object object at .*>>")
+
+    @unittest.skipIf(_testcapi is None, 'need the _testcapi module')
+    def test_bpo25750(self):
+        # bpo-25750: calling a descriptor (implemented as built-in
+        # function with METH_FASTCALL) should not crash CPython if the
+        # descriptor deletes itself from the class.
+        class Descr:
+            __get__ = _testcapi.bad_get
+
+        class X:
+            descr = Descr()
+            def __new__(cls):
+                cls.descr = None
+                # Create this large list to corrupt some unused memory
+                cls.lst = [2**i for i in range(10000)]
+        X.descr
+
+    @support.suppress_immortalization()
+    def test_remove_subclass(self):
+        # bpo-46417: when the last subclass of a type is deleted,
+        # remove_subclass() clears the internal dictionary of subclasses:
+        # set PyTypeObject.tp_subclasses to NULL. remove_subclass() is called
+        # when a type is deallocated.
+        class Parent:
+            pass
+        self.assertEqual(Parent.__subclasses__(), [])
+
+        class Child(Parent):
+            pass
+        self.assertEqual(Parent.__subclasses__(), [Child])
+
+        del Child
+        gc.collect()
+        self.assertEqual(Parent.__subclasses__(), [])
+
+    def test_instance_method_get_behavior(self):
+        # test case for gh-113157
+
+        class A:
+            def meth(self):
+                return self
+
+        class B:
+            pass
+
+        a = A()
+        b = B()
+        b.meth = a.meth.__get__(b, B)
+        self.assertEqual(b.meth(), a)
+
+    def test_attr_raise_through_property(self):
+        # test case for gh-103272
+        class A:
+            def __getattr__(self, name):
+                raise ValueError("FOO")
+
+            @property
+            def foo(self):
+                return self.__getattr__("asdf")
+
+        with self.assertRaisesRegex(ValueError, "FOO"):
+            A().foo
+
+        # test case for gh-103551
+        class B:
+            @property
+            def __getattr__(self, name):
+                raise ValueError("FOO")
+
+            @property
+            def foo(self):
+                raise NotImplementedError("BAR")
+
+        with self.assertRaisesRegex(NotImplementedError, "BAR"):
+            B().foo
+
+
+class DictProxyTests(CPythonTestCase):
+    def setUp(self):
+        super().setUp()
+        class C(object):
+            def meth(self):
+                pass
+        self.C = C
+
+    @unittest.skipIf(hasattr(sys, 'gettrace') and sys.gettrace(),
+                        'trace function introduces __local__')
+    def test_iter_keys(self):
+        # Testing dict-proxy keys...
+        it = self.C.__dict__.keys()
+        self.assertNotIsInstance(it, list)
+        keys = list(it)
+        keys.sort()
+        self.assertEqual(keys, ['__dict__', '__doc__', '__firstlineno__',
+                                '__module__',
+                                '__static_attributes__', '__weakref__',
+                                'meth'])
+
+    @unittest.skipIf(hasattr(sys, 'gettrace') and sys.gettrace(),
+                        'trace function introduces __local__')
+    def test_iter_values(self):
+        # Testing dict-proxy values...
+        it = self.C.__dict__.values()
+        self.assertNotIsInstance(it, list)
+        values = list(it)
+        self.assertEqual(len(values), 7)
+
+    @unittest.skipIf(hasattr(sys, 'gettrace') and sys.gettrace(),
+                        'trace function introduces __local__')
+    def test_iter_items(self):
+        # Testing dict-proxy iteritems...
+        it = self.C.__dict__.items()
+        self.assertNotIsInstance(it, list)
+        keys = [item[0] for item in it]
+        keys.sort()
+        self.assertEqual(keys, ['__dict__', '__doc__', '__firstlineno__',
+                                '__module__',
+                                '__static_attributes__', '__weakref__',
+                                'meth'])
+
+    def test_dict_type_with_metaclass(self):
+        # Testing type of __dict__ when metaclass set...
+        with torch._dynamo.error_on_graph_break(False):
+            class B(object):
+                pass
+            class M(type):
+                pass
+            class C(metaclass=M):
+                # In 2.3a1, C.__dict__ was a real dict rather than a dict proxy
+                pass
+        self.assertEqual(type(C.__dict__), type(B.__dict__))
+
+    def test_repr(self):
+        # Testing mappingproxy.__repr__.
+        # We can't blindly compare with the repr of another dict as ordering
+        # of keys and values is arbitrary and may differ.
+        r = repr(self.C.__dict__)
+        self.assertTrue(r.startswith('mappingproxy('), r)
+        self.assertTrue(r.endswith(')'), r)
+        for k, v in self.C.__dict__.items():
+            self.assertIn('{!r}: {!r}'.format(k, v), r)
+
+
+class AAAPTypesLongInitTest(CPythonTestCase):
+    # This is in its own TestCase so that it can be run before any other tests.
+    # (Hence the 'AAA' in the test class name: to make it the first
+    # item in a list sorted by name, like
+    # unittest.TestLoader.getTestCaseNames() does.)
+    def test_pytype_long_ready(self):
+        # Testing SF bug 551412 ...
+
+        # This dumps core when SF bug 551412 isn't fixed --
+        # but only when test_descr.py is run separately.
+        # (That can't be helped -- as soon as PyType_Ready()
+        # is called for PyLong_Type, the bug is gone.)
+        class UserLong(object):
+            def __pow__(self, *args):
+                pass
+        try:
+            pow(0, UserLong(), 0)
+        except:
+            pass
+
+        # Another segfault only when run early
+        # (before PyType_Ready(tuple) is called)
+        type.mro(tuple)
+
+
+class MiscTests(CPythonTestCase):
+    def test_type_lookup_mro_reference(self):
+        # Issue #14199: _PyType_Lookup() has to keep a strong reference to
+        # the type MRO because it may be modified during the lookup, if
+        # __bases__ is set during the lookup for example.
+        class MyKey(object):
+            def __hash__(self):
+                return hash('mykey')
+
+            def __eq__(self, other):
+                X.__bases__ = (Base2,)
+
+        class Base(object):
+            mykey = 'from Base'
+            mykey2 = 'from Base'
+
+        class Base2(object):
+            mykey = 'from Base2'
+            mykey2 = 'from Base2'
+
+        with self.assertWarnsRegex(RuntimeWarning, 'X'):
+            X = type('X', (Base,), {MyKey(): 5})
+
+        # Note that the access below uses getattr() rather than normally
+        # accessing the attribute.  That is done to avoid the bytecode
+        # specializer activating on repeated runs of the test.
+
+        # mykey is read from Base
+        self.assertEqual(getattr(X, 'mykey'), 'from Base')
+        # mykey2 is read from Base2 because MyKey.__eq__ has set __bases__
+        self.assertEqual(getattr(X, 'mykey2'), 'from Base2')
+
+
+class PicklingTests(CPythonTestCase):
+
+    def _check_reduce(self, proto, obj, args=(), kwargs={}, state=None,
+                      listitems=None, dictitems=None):
+        if proto >= 2:
+            reduce_value = obj.__reduce_ex__(proto)
+            if kwargs:
+                self.assertEqual(reduce_value[0], copyreg.__newobj_ex__)
+                self.assertEqual(reduce_value[1], (type(obj), args, kwargs))
+            else:
+                self.assertEqual(reduce_value[0], copyreg.__newobj__)
+                self.assertEqual(reduce_value[1], (type(obj),) + args)
+            self.assertEqual(reduce_value[2], state)
+            if listitems is not None:
+                self.assertListEqual(list(reduce_value[3]), listitems)
+            else:
+                self.assertIsNone(reduce_value[3])
+            if dictitems is not None:
+                self.assertDictEqual(dict(reduce_value[4]), dictitems)
+            else:
+                self.assertIsNone(reduce_value[4])
+        else:
+            base_type = type(obj).__base__
+            reduce_value = (copyreg._reconstructor,
+                            (type(obj),
+                             base_type,
+                             None if base_type is object else base_type(obj)))
+            if state is not None:
+                reduce_value += (state,)
+            self.assertEqual(obj.__reduce_ex__(proto), reduce_value)
+            self.assertEqual(obj.__reduce__(), reduce_value)
+
+    def test_reduce(self):
+        protocols = range(pickle.HIGHEST_PROTOCOL + 1)
+        args = (-101, "spam")
+        kwargs = {'bacon': -201, 'fish': -301}
+        state = {'cheese': -401}
+
+        class C1:
+            def __getnewargs__(self):
+                return args
+        obj = C1()
+        for proto in protocols:
+            self._check_reduce(proto, obj, args)
+
+        for name, value in state.items():
+            setattr(obj, name, value)
+        for proto in protocols:
+            self._check_reduce(proto, obj, args, state=state)
+
+        class C2:
+            def __getnewargs__(self):
+                return "bad args"
+        obj = C2()
+        for proto in protocols:
+            if proto >= 2:
+                with self.assertRaises(TypeError):
+                    obj.__reduce_ex__(proto)
+
+        class C3:
+            def __getnewargs_ex__(self):
+                return (args, kwargs)
+        obj = C3()
+        for proto in protocols:
+            if proto >= 2:
+                self._check_reduce(proto, obj, args, kwargs)
+
+        class C4:
+            def __getnewargs_ex__(self):
+                return (args, "bad dict")
+        class C5:
+            def __getnewargs_ex__(self):
+                return ("bad tuple", kwargs)
+        class C6:
+            def __getnewargs_ex__(self):
+                return ()
+        class C7:
+            def __getnewargs_ex__(self):
+                return "bad args"
+        for proto in protocols:
+            for cls in C4, C5, C6, C7:
+                obj = cls()
+                if proto >= 2:
+                    with self.assertRaises((TypeError, ValueError)):
+                        obj.__reduce_ex__(proto)
+
+        class C9:
+            def __getnewargs_ex__(self):
+                return (args, {})
+        obj = C9()
+        for proto in protocols:
+            self._check_reduce(proto, obj, args)
+
+        class C10:
+            def __getnewargs_ex__(self):
+                raise IndexError
+        obj = C10()
+        for proto in protocols:
+            if proto >= 2:
+                with self.assertRaises(IndexError):
+                    obj.__reduce_ex__(proto)
+
+        class C11:
+            def __getstate__(self):
+                return state
+        obj = C11()
+        for proto in protocols:
+            self._check_reduce(proto, obj, state=state)
+
+        class C12:
+            def __getstate__(self):
+                return "not dict"
+        obj = C12()
+        for proto in protocols:
+            self._check_reduce(proto, obj, state="not dict")
+
+        class C13:
+            def __getstate__(self):
+                raise IndexError
+        obj = C13()
+        for proto in protocols:
+            with self.assertRaises(IndexError):
+                obj.__reduce_ex__(proto)
+            if proto < 2:
+                with self.assertRaises(IndexError):
+                    obj.__reduce__()
+
+        class C14:
+            __slots__ = tuple(state)
+            def __init__(self):
+                for name, value in state.items():
+                    setattr(self, name, value)
+
+        obj = C14()
+        for proto in protocols:
+            if proto >= 2:
+                self._check_reduce(proto, obj, state=(None, state))
+            else:
+                with self.assertRaises(TypeError):
+                    obj.__reduce_ex__(proto)
+                with self.assertRaises(TypeError):
+                    obj.__reduce__()
+
+        class C15(dict):
+            pass
+        obj = C15({"quebec": -601})
+        for proto in protocols:
+            self._check_reduce(proto, obj, dictitems=dict(obj))
+
+        class C16(list):
+            pass
+        obj = C16(["yukon"])
+        for proto in protocols:
+            self._check_reduce(proto, obj, listitems=list(obj))
+
+    def test_special_method_lookup(self):
+        protocols = range(pickle.HIGHEST_PROTOCOL + 1)
+        class Picky:
+            def __getstate__(self):
+                return {}
+
+            def __getattr__(self, attr):
+                if attr in ("__getnewargs__", "__getnewargs_ex__"):
+                    raise AssertionError(attr)
+                return None
+        for protocol in protocols:
+            state = {} if protocol >= 2 else None
+            self._check_reduce(protocol, Picky(), state=state)
+
+    def _assert_is_copy(self, obj, objcopy, msg=None):
+        """Utility method to verify if two objects are copies of each others.
+        """
+        if msg is None:
+            msg = "{!r} is not a copy of {!r}".format(obj, objcopy)
+        if type(obj).__repr__ is object.__repr__:
+            # We have this limitation for now because we use the object's repr
+            # to help us verify that the two objects are copies. This allows
+            # us to delegate the non-generic verification logic to the objects
+            # themselves.
+            raise ValueError("object passed to _assert_is_copy must " +
+                             "override the __repr__ method.")
+        self.assertIsNot(obj, objcopy, msg=msg)
+        self.assertIs(type(obj), type(objcopy), msg=msg)
+        if hasattr(obj, '__dict__'):
+            self.assertDictEqual(obj.__dict__, objcopy.__dict__, msg=msg)
+            self.assertIsNot(obj.__dict__, objcopy.__dict__, msg=msg)
+        if hasattr(obj, '__slots__'):
+            self.assertListEqual(obj.__slots__, objcopy.__slots__, msg=msg)
+            for slot in obj.__slots__:
+                self.assertEqual(
+                    hasattr(obj, slot), hasattr(objcopy, slot), msg=msg)
+                self.assertEqual(getattr(obj, slot, None),
+                                 getattr(objcopy, slot, None), msg=msg)
+        self.assertEqual(repr(obj), repr(objcopy), msg=msg)
+
+    @staticmethod
+    def _generate_pickle_copiers():
+        """Utility method to generate the many possible pickle configurations.
+        """
+        class PickleCopier:
+            "This class copies object using pickle."
+            def __init__(self, proto, dumps, loads):
+                self.proto = proto
+                self.dumps = dumps
+                self.loads = loads
+            def copy(self, obj):
+                return self.loads(self.dumps(obj, self.proto))
+            def __repr__(self):
+                # We try to be as descriptive as possible here since this is
+                # the string which we will allow us to tell the pickle
+                # configuration we are using during debugging.
+                return ("PickleCopier(proto={}, dumps={}.{}, loads={}.{})"
+                        .format(self.proto,
+                                self.dumps.__module__, self.dumps.__qualname__,
+                                self.loads.__module__, self.loads.__qualname__))
+        return (PickleCopier(*args) for args in
+                   itertools.product(range(pickle.HIGHEST_PROTOCOL + 1),
+                                     {pickle.dumps, pickle._dumps},
+                                     {pickle.loads, pickle._loads}))
+
+    def test_pickle_slots(self):
+        # Tests pickling of classes with __slots__.
+
+        # Pickling of classes with __slots__ but without __getstate__ should
+        # fail (if using protocol 0 or 1)
+        global C
+        class C:
+            __slots__ = ['a']
+        with self.assertRaises(TypeError):
+            pickle.dumps(C(), 0)
+
+        global D
+        class D(C):
+            pass
+        with self.assertRaises(TypeError):
+            pickle.dumps(D(), 0)
+
+        class C:
+            "A class with __getstate__ and __setstate__ implemented."
+            __slots__ = ['a']
+            def __getstate__(self):
+                state = getattr(self, '__dict__', {}).copy()
+                for cls in type(self).__mro__:
+                    for slot in cls.__dict__.get('__slots__', ()):
+                        try:
+                            state[slot] = getattr(self, slot)
+                        except AttributeError:
+                            pass
+                return state
+            def __setstate__(self, state):
+                for k, v in state.items():
+                    setattr(self, k, v)
+            def __repr__(self):
+                return "%s()<%r>" % (type(self).__name__, self.__getstate__())
+
+        class D(C):
+            "A subclass of a class with slots."
+            pass
+
+        global E
+        class E(C):
+            "A subclass with an extra slot."
+            __slots__ = ['b']
+
+        # Now it should work
+        for pickle_copier in self._generate_pickle_copiers():
+            with self.subTest(pickle_copier=pickle_copier):
+                x = C()
+                y = pickle_copier.copy(x)
+                self._assert_is_copy(x, y)
+
+                x.a = 42
+                y = pickle_copier.copy(x)
+                self._assert_is_copy(x, y)
+
+                x = D()
+                x.a = 42
+                x.b = 100
+                y = pickle_copier.copy(x)
+                self._assert_is_copy(x, y)
+
+                x = E()
+                x.a = 42
+                x.b = "foo"
+                y = pickle_copier.copy(x)
+                self._assert_is_copy(x, y)
+
+    def test_reduce_copying(self):
+        # Tests pickling and copying new-style classes and objects.
+        global C1
+        class C1:
+            "The state of this class is copyable via its instance dict."
+            ARGS = (1, 2)
+            NEED_DICT_COPYING = True
+            def __init__(self, a, b):
+                super().__init__()
+                self.a = a
+                self.b = b
+            def __repr__(self):
+                return "C1(%r, %r)" % (self.a, self.b)
+
+        global C2
+        class C2(list):
+            "A list subclass copyable via __getnewargs__."
+            ARGS = (1, 2)
+            NEED_DICT_COPYING = False
+            def __new__(cls, a, b):
+                self = super().__new__(cls)
+                self.a = a
+                self.b = b
+                return self
+            def __init__(self, *args):
+                super().__init__()
+                # This helps testing that __init__ is not called during the
+                # unpickling process, which would cause extra appends.
+                self.append("cheese")
+            @classmethod
+            def __getnewargs__(cls):
+                return cls.ARGS
+            def __repr__(self):
+                return "C2(%r, %r)<%r>" % (self.a, self.b, list(self))
+
+        global C3
+        class C3(list):
+            "A list subclass copyable via __getstate__."
+            ARGS = (1, 2)
+            NEED_DICT_COPYING = False
+            def __init__(self, a, b):
+                self.a = a
+                self.b = b
+                # This helps testing that __init__ is not called during the
+                # unpickling process, which would cause extra appends.
+                self.append("cheese")
+            @classmethod
+            def __getstate__(cls):
+                return cls.ARGS
+            def __setstate__(self, state):
+                a, b = state
+                self.a = a
+                self.b = b
+            def __repr__(self):
+                return "C3(%r, %r)<%r>" % (self.a, self.b, list(self))
+
+        global C4
+        class C4(int):
+            "An int subclass copyable via __getnewargs__."
+            ARGS = ("hello", "world", 1)
+            NEED_DICT_COPYING = False
+            def __new__(cls, a, b, value):
+                self = super().__new__(cls, value)
+                self.a = a
+                self.b = b
+                return self
+            @classmethod
+            def __getnewargs__(cls):
+                return cls.ARGS
+            def __repr__(self):
+                return "C4(%r, %r)<%r>" % (self.a, self.b, int(self))
+
+        global C5
+        class C5(int):
+            "An int subclass copyable via __getnewargs_ex__."
+            ARGS = (1, 2)
+            KWARGS = {'value': 3}
+            NEED_DICT_COPYING = False
+            def __new__(cls, a, b, *, value=0):
+                self = super().__new__(cls, value)
+                self.a = a
+                self.b = b
+                return self
+            @classmethod
+            def __getnewargs_ex__(cls):
+                return (cls.ARGS, cls.KWARGS)
+            def __repr__(self):
+                return "C5(%r, %r)<%r>" % (self.a, self.b, int(self))
+
+        test_classes = (C1, C2, C3, C4, C5)
+        # Testing copying through pickle
+        pickle_copiers = self._generate_pickle_copiers()
+        for cls, pickle_copier in itertools.product(test_classes, pickle_copiers):
+            with self.subTest(cls=cls, pickle_copier=pickle_copier):
+                kwargs = getattr(cls, 'KWARGS', {})
+                obj = cls(*cls.ARGS, **kwargs)
+                proto = pickle_copier.proto
+                objcopy = pickle_copier.copy(obj)
+                self._assert_is_copy(obj, objcopy)
+                # For test classes that supports this, make sure we didn't go
+                # around the reduce protocol by simply copying the attribute
+                # dictionary. We clear attributes using the previous copy to
+                # not mutate the original argument.
+                if proto >= 2 and not cls.NEED_DICT_COPYING:
+                    objcopy.__dict__.clear()
+                    objcopy2 = pickle_copier.copy(objcopy)
+                    self._assert_is_copy(obj, objcopy2)
+
+        # Testing copying through copy.deepcopy()
+        for cls in test_classes:
+            with self.subTest(cls=cls):
+                kwargs = getattr(cls, 'KWARGS', {})
+                obj = cls(*cls.ARGS, **kwargs)
+                objcopy = deepcopy(obj)
+                self._assert_is_copy(obj, objcopy)
+                # For test classes that supports this, make sure we didn't go
+                # around the reduce protocol by simply copying the attribute
+                # dictionary. We clear attributes using the previous copy to
+                # not mutate the original argument.
+                if not cls.NEED_DICT_COPYING:
+                    objcopy.__dict__.clear()
+                    objcopy2 = deepcopy(objcopy)
+                    self._assert_is_copy(obj, objcopy2)
+
+    def test_issue24097(self):
+        # Slot name is freed inside __getattr__ and is later used.
+        class S(str):  # Not interned
+            pass
+        class A:
+            __slotnames__ = [S('spam')]
+            def __getattr__(self, attr):
+                if attr == 'spam':
+                    A.__slotnames__[:] = [S('spam')]
+                    return 42
+                else:
+                    raise AttributeError
+
+        import copyreg
+        expected = (copyreg.__newobj__, (A,), (None, {'spam': 42}), None, None)
+        self.assertEqual(A().__reduce_ex__(2), expected)  # Shouldn't crash
+
+    def test_object_reduce(self):
+        # Issue #29914
+        # __reduce__() takes no arguments
+        object().__reduce__()
+        with self.assertRaises(TypeError):
+            object().__reduce__(0)
+        # __reduce_ex__() takes one integer argument
+        object().__reduce_ex__(0)
+        with self.assertRaises(TypeError):
+            object().__reduce_ex__()
+        with self.assertRaises(TypeError):
+            object().__reduce_ex__(None)
+
+
+class SharedKeyTests(CPythonTestCase):
+
+    @support.cpython_only
+    def test_subclasses(self):
+        # Verify that subclasses can share keys (per PEP 412)
+        class A:
+            pass
+        class B(A):
+            pass
+
+        #Shrink keys by repeatedly creating instances
+        [(A(), B()) for _ in range(30)]
+
+        a, b = A(), B()
+        self.assertEqual(sys.getsizeof(vars(a)), sys.getsizeof(vars(b)))
+        self.assertLess(sys.getsizeof(vars(a)), sys.getsizeof({"a":1}))
+        # Initial hash table can contain only one or two elements.
+        # Set 6 attributes to cause internal resizing.
+        a.x, a.y, a.z, a.w, a.v, a.u = range(6)
+        self.assertNotEqual(sys.getsizeof(vars(a)), sys.getsizeof(vars(b)))
+        a2 = A()
+        self.assertGreater(sys.getsizeof(vars(a)), sys.getsizeof(vars(a2)))
+        self.assertLess(sys.getsizeof(vars(a2)), sys.getsizeof({"a":1}))
+        self.assertLess(sys.getsizeof(vars(b)), sys.getsizeof({"a":1}))
+
+
+class DebugHelperMeta(type):
+    """
+    Sets default __doc__ and simplifies repr() output.
+    """
+    def __new__(mcls, name, bases, attrs):
+        if attrs.get('__doc__') is None:
+            attrs['__doc__'] = name  # helps when debugging with gdb
+        return type.__new__(mcls, name, bases, attrs)
+    def __repr__(cls):
+        return repr(cls.__name__)
+
+
+class MroTest(CPythonTestCase):
+    """
+    Regressions for some bugs revealed through
+    mcsl.mro() customization (typeobject.c: mro_internal()) and
+    cls.__bases__ assignment (typeobject.c: type_set_bases()).
+    """
+
+    def setUp(self):
+        super().setUp()
+        self.step = 0
+        self.ready = False
+
+    def step_until(self, limit):
+        ret = (self.step < limit)
+        if ret:
+            self.step += 1
+        return ret
+
+    def test_incomplete_set_bases_on_self(self):
+        """
+        type_set_bases must be aware that type->tp_mro can be NULL.
+        """
+        class M(DebugHelperMeta):
+            def mro(cls):
+                if self.step_until(1):
+                    assert cls.__mro__ is None
+                    cls.__bases__ += ()
+
+                return type.mro(cls)
+
+        class A(metaclass=M):
+            pass
+
+    def test_reent_set_bases_on_base(self):
+        """
+        Deep reentrancy must not over-decref old_mro.
+        """
+        class M(DebugHelperMeta):
+            def mro(cls):
+                if cls.__mro__ is not None and cls.__name__ == 'B':
+                    # 4-5 steps are usually enough to make it crash somewhere
+                    if self.step_until(10):
+                        A.__bases__ += ()
+
+                return type.mro(cls)
+
+        class A(metaclass=M):
+            pass
+        class B(A):
+            pass
+        B.__bases__ += ()
+
+    def test_reent_set_bases_on_direct_base(self):
+        """
+        Similar to test_reent_set_bases_on_base, but may crash differently.
+        """
+        class M(DebugHelperMeta):
+            def mro(cls):
+                base = cls.__bases__[0]
+                if base is not object:
+                    if self.step_until(5):
+                        base.__bases__ += ()
+
+                return type.mro(cls)
+
+        class A(metaclass=M):
+            pass
+        class B(A):
+            pass
+        class C(B):
+            pass
+
+    def test_reent_set_bases_tp_base_cycle(self):
+        """
+        type_set_bases must check for an inheritance cycle not only through
+        MRO of the type, which may be not yet updated in case of reentrance,
+        but also through tp_base chain, which is assigned before diving into
+        inner calls to mro().
+
+        Otherwise, the following snippet can loop forever:
+            do {
+                // ...
+                type = type->tp_base;
+            } while (type != NULL);
+
+        Functions that rely on tp_base (like solid_base and PyType_IsSubtype)
+        would not be happy in that case, causing a stack overflow.
+        """
+        class M(DebugHelperMeta):
+            def mro(cls):
+                if self.ready:
+                    if cls.__name__ == 'B1':
+                        B2.__bases__ = (B1,)
+                    if cls.__name__ == 'B2':
+                        B1.__bases__ = (B2,)
+                return type.mro(cls)
+
+        class A(metaclass=M):
+            pass
+        class B1(A):
+            pass
+        class B2(A):
+            pass
+
+        self.ready = True
+        with self.assertRaises(TypeError):
+            B1.__bases__ += ()
+
+    def test_tp_subclasses_cycle_in_update_slots(self):
+        """
+        type_set_bases must check for reentrancy upon finishing its job
+        by updating tp_subclasses of old/new bases of the type.
+        Otherwise, an implicit inheritance cycle through tp_subclasses
+        can break functions that recurse on elements of that field
+        (like recurse_down_subclasses and mro_hierarchy) eventually
+        leading to a stack overflow.
+        """
+        class M(DebugHelperMeta):
+            def mro(cls):
+                if self.ready and cls.__name__ == 'C':
+                    self.ready = False
+                    C.__bases__ = (B2,)
+                return type.mro(cls)
+
+        class A(metaclass=M):
+            pass
+        class B1(A):
+            pass
+        class B2(A):
+            pass
+        class C(A):
+            pass
+
+        self.ready = True
+        C.__bases__ = (B1,)
+        B1.__bases__ = (C,)
+
+        self.assertEqual(C.__bases__, (B2,))
+        self.assertEqual(B2.__subclasses__(), [C])
+        self.assertEqual(B1.__subclasses__(), [])
+
+        self.assertEqual(B1.__bases__, (C,))
+        self.assertEqual(C.__subclasses__(), [B1])
+
+    def test_tp_subclasses_cycle_error_return_path(self):
+        """
+        The same as test_tp_subclasses_cycle_in_update_slots, but tests
+        a code path executed on error (goto bail).
+        """
+        class E(Exception):
+            pass
+        class M(DebugHelperMeta):
+            def mro(cls):
+                if self.ready and cls.__name__ == 'C':
+                    if C.__bases__ == (B2,):
+                        self.ready = False
+                    else:
+                        C.__bases__ = (B2,)
+                        raise E
+                return type.mro(cls)
+
+        class A(metaclass=M):
+            pass
+        class B1(A):
+            pass
+        class B2(A):
+            pass
+        class C(A):
+            pass
+
+        self.ready = True
+        with self.assertRaises(E):
+            C.__bases__ = (B1,)
+        B1.__bases__ = (C,)
+
+        self.assertEqual(C.__bases__, (B2,))
+        self.assertEqual(C.__mro__, tuple(type.mro(C)))
+
+    def test_incomplete_extend(self):
+        """
+        Extending an uninitialized type with type->tp_mro == NULL must
+        throw a reasonable TypeError exception, instead of failing
+        with PyErr_BadInternalCall.
+        """
+        class M(DebugHelperMeta):
+            def mro(cls):
+                if cls.__mro__ is None and cls.__name__ != 'X':
+                    with self.assertRaises(TypeError):
+                        class X(cls):
+                            pass
+
+                return type.mro(cls)
+
+        class A(metaclass=M):
+            pass
+
+    def test_incomplete_super(self):
+        """
+        Attribute lookup on a super object must be aware that
+        its target type can be uninitialized (type->tp_mro == NULL).
+        """
+        class M(DebugHelperMeta):
+            def mro(cls):
+                if cls.__mro__ is None:
+                    with self.assertRaises(AttributeError):
+                        super(cls, cls).xxx
+
+                return type.mro(cls)
+
+        class A(metaclass=M):
+            pass
+
+    def test_disappearing_custom_mro(self):
+        """
+        gh-92112: A custom mro() returning a result conflicting with
+        __bases__ and deleting itself caused a double free.
+        """
+        class B:
+            pass
+
+        class M(DebugHelperMeta):
+            def mro(cls):
+                del M.mro
+                return (B,)
+
+        with self.assertRaises(TypeError):
+            class A(metaclass=M):
+                pass
+
+
+if __name__ == "__main__":
+    run_tests()
diff --git a/test/dynamo/cpython/3_13/test_funcattrs.diff b/test/dynamo/cpython/3_13/test_funcattrs.diff
new file mode 100644
index 0000000000000..1c188255b3bd5
--- /dev/null
+++ b/test/dynamo/cpython/3_13/test_funcattrs.diff
@@ -0,0 +1,82 @@
+diff --git a/test/dynamo/cpython/3_13/test_funcattrs.py b/test/dynamo/cpython/3_13/test_funcattrs.py
+index 375f456dfde..795e7b586d7 100644
+--- a/test/dynamo/cpython/3_13/test_funcattrs.py
++++ b/test/dynamo/cpython/3_13/test_funcattrs.py
+@@ -1,3 +1,21 @@
++# ======= BEGIN Dynamo patch =======
++# Owner(s): ["module: dynamo"]
++
++# ruff: noqa
++# flake8: noqa
++
++# Test copied from
++# https://raw.githubusercontent.com/python/cpython/refs/tags/v3.13.5/Lib/test/test_funcattrs.py
++
++import sys
++import torch
++import torch._dynamo.test_case
++import unittest
++from torch._dynamo.test_case import CPythonTestCase
++from torch.testing._internal.common_utils import run_tests
++
++# ======= END DYNAMO PATCH =======
++
+ import textwrap
+ import types
+ import typing
+@@ -8,8 +26,9 @@ from test import support
+ 
+ def global_function():
+     def inner_function():
+-        class LocalClass:
+-            pass
++        with torch._dynamo.error_on_graph_break(False):
++            class LocalClass:
++                pass
+         global inner_global_function
+         def inner_global_function():
+             def inner_function2():
+@@ -19,8 +38,9 @@ def global_function():
+     return lambda: inner_function
+ 
+ 
+-class FuncAttrsTest(unittest.TestCase):
++class FuncAttrsTest(CPythonTestCase):
+     def setUp(self):
++        super().setUp()
+         class F:
+             def a(self):
+                 pass
+@@ -430,7 +450,7 @@ def empty_cell(empty=True):
+     return f.__closure__[0]
+ 
+ 
+-class CellTest(unittest.TestCase):
++class CellTest(CPythonTestCase):
+     def test_comparison(self):
+         # These tests are here simply to exercise the comparison code;
+         # their presence should not be interpreted as providing any
+@@ -443,7 +463,7 @@ class CellTest(unittest.TestCase):
+         self.assertTrue(cell(True) > empty_cell())
+ 
+ 
+-class StaticMethodAttrsTest(unittest.TestCase):
++class StaticMethodAttrsTest(CPythonTestCase):
+     def test_func_attribute(self):
+         def f():
+             pass
+@@ -455,7 +475,7 @@ class StaticMethodAttrsTest(unittest.TestCase):
+         self.assertTrue(s.__func__ is f)
+ 
+ 
+-class BuiltinFunctionPropertiesTest(unittest.TestCase):
++class BuiltinFunctionPropertiesTest(CPythonTestCase):
+     # XXX Not sure where this should really go since I can't find a
+     # test module specifically for builtin_function_or_method.
+ 
+@@ -508,4 +528,4 @@ class BuiltinFunctionPropertiesTest(unittest.TestCase):
+ 
+ 
+ if __name__ == "__main__":
+-    unittest.main()
++    run_tests()
diff --git a/test/dynamo/cpython/3_13/test_funcattrs.py b/test/dynamo/cpython/3_13/test_funcattrs.py
new file mode 100644
index 0000000000000..795e7b586d76f
--- /dev/null
+++ b/test/dynamo/cpython/3_13/test_funcattrs.py
@@ -0,0 +1,531 @@
+# ======= BEGIN Dynamo patch =======
+# Owner(s): ["module: dynamo"]
+
+# ruff: noqa
+# flake8: noqa
+
+# Test copied from
+# https://raw.githubusercontent.com/python/cpython/refs/tags/v3.13.5/Lib/test/test_funcattrs.py
+
+import sys
+import torch
+import torch._dynamo.test_case
+import unittest
+from torch._dynamo.test_case import CPythonTestCase
+from torch.testing._internal.common_utils import run_tests
+
+# ======= END DYNAMO PATCH =======
+
+import textwrap
+import types
+import typing
+import unittest
+import warnings
+from test import support
+
+
+def global_function():
+    def inner_function():
+        with torch._dynamo.error_on_graph_break(False):
+            class LocalClass:
+                pass
+        global inner_global_function
+        def inner_global_function():
+            def inner_function2():
+                pass
+            return inner_function2
+        return LocalClass
+    return lambda: inner_function
+
+
+class FuncAttrsTest(CPythonTestCase):
+    def setUp(self):
+        super().setUp()
+        class F:
+            def a(self):
+                pass
+        def b():
+            return 3
+        self.fi = F()
+        self.F = F
+        self.b = b
+
+    def cannot_set_attr(self, obj, name, value, exceptions):
+        try:
+            setattr(obj, name, value)
+        except exceptions:
+            pass
+        else:
+            self.fail("shouldn't be able to set %s to %r" % (name, value))
+        try:
+            delattr(obj, name)
+        except exceptions:
+            pass
+        else:
+            self.fail("shouldn't be able to del %s" % name)
+
+
+class FunctionPropertiesTest(FuncAttrsTest):
+    # Include the external setUp method that is common to all tests
+    def test_module(self):
+        self.assertEqual(self.b.__module__, __name__)
+
+    def test_dir_includes_correct_attrs(self):
+        self.b.known_attr = 7
+        self.assertIn('known_attr', dir(self.b),
+            "set attributes not in dir listing of method")
+        # Test on underlying function object of method
+        self.F.a.known_attr = 7
+        self.assertIn('known_attr', dir(self.fi.a), "set attribute on function "
+                     "implementations, should show up in next dir")
+
+    def test_duplicate_function_equality(self):
+        # Body of `duplicate' is the exact same as self.b
+        def duplicate():
+            'my docstring'
+            return 3
+        self.assertNotEqual(self.b, duplicate)
+
+    def test_copying___code__(self):
+        def test(): pass
+        self.assertEqual(test(), None)
+        test.__code__ = self.b.__code__
+        self.assertEqual(test(), 3) # self.b always returns 3, arbitrarily
+
+    def test_invalid___code___assignment(self):
+        def A(): pass
+        def B(): yield
+        async def C(): yield
+        async def D(x): await x
+
+        for src in [A, B, C, D]:
+            for dst in [A, B, C, D]:
+                if src == dst:
+                    continue
+
+                assert src.__code__.co_flags != dst.__code__.co_flags
+                prev = dst.__code__
+                try:
+                    with self.assertWarnsRegex(DeprecationWarning, 'code object of non-matching type'):
+                        dst.__code__ = src.__code__
+                finally:
+                    with warnings.catch_warnings():
+                        warnings.filterwarnings('ignore', '', DeprecationWarning)
+                        dst.__code__ = prev
+
+    def test___globals__(self):
+        self.assertIs(self.b.__globals__, globals())
+        self.cannot_set_attr(self.b, '__globals__', 2,
+                             (AttributeError, TypeError))
+
+    def test___builtins__(self):
+        if __name__ == "__main__":
+            builtins_dict = __builtins__.__dict__
+        else:
+            builtins_dict = __builtins__
+
+        self.assertIs(self.b.__builtins__, builtins_dict)
+        self.cannot_set_attr(self.b, '__builtins__', 2,
+                             (AttributeError, TypeError))
+
+        # bpo-42990: If globals is specified and has no "__builtins__" key,
+        # a function inherits the current builtins namespace.
+        def func(s): return len(s)
+        ns = {}
+        func2 = type(func)(func.__code__, ns)
+        self.assertIs(func2.__globals__, ns)
+        self.assertIs(func2.__builtins__, builtins_dict)
+
+        # Make sure that the function actually works.
+        self.assertEqual(func2("abc"), 3)
+        self.assertEqual(ns, {})
+
+        # Define functions using exec() with different builtins,
+        # and test inheritance when globals has no "__builtins__" key
+        code = textwrap.dedent("""
+            def func3(s): pass
+            func4 = type(func3)(func3.__code__, {})
+        """)
+        safe_builtins = {'None': None}
+        ns = {'type': type, '__builtins__': safe_builtins}
+        exec(code, ns)
+        self.assertIs(ns['func3'].__builtins__, safe_builtins)
+        self.assertIs(ns['func4'].__builtins__, safe_builtins)
+        self.assertIs(ns['func3'].__globals__['__builtins__'], safe_builtins)
+        self.assertNotIn('__builtins__', ns['func4'].__globals__)
+
+    def test___closure__(self):
+        a = 12
+        def f(): print(a)
+        c = f.__closure__
+        self.assertIsInstance(c, tuple)
+        self.assertEqual(len(c), 1)
+        # don't have a type object handy
+        self.assertEqual(c[0].__class__.__name__, "cell")
+        self.cannot_set_attr(f, "__closure__", c, AttributeError)
+
+    def test_cell_new(self):
+        cell_obj = types.CellType(1)
+        self.assertEqual(cell_obj.cell_contents, 1)
+
+        cell_obj = types.CellType()
+        msg = "shouldn't be able to read an empty cell"
+        with self.assertRaises(ValueError, msg=msg):
+            cell_obj.cell_contents
+
+    def test_empty_cell(self):
+        def f(): print(a)
+        try:
+            f.__closure__[0].cell_contents
+        except ValueError:
+            pass
+        else:
+            self.fail("shouldn't be able to read an empty cell")
+        a = 12
+
+    def test_set_cell(self):
+        a = 12
+        def f(): return a
+        c = f.__closure__
+        c[0].cell_contents = 9
+        self.assertEqual(c[0].cell_contents, 9)
+        self.assertEqual(f(), 9)
+        self.assertEqual(a, 9)
+        del c[0].cell_contents
+        try:
+            c[0].cell_contents
+        except ValueError:
+            pass
+        else:
+            self.fail("shouldn't be able to read an empty cell")
+        with self.assertRaises(NameError):
+            f()
+        with self.assertRaises(UnboundLocalError):
+            print(a)
+
+    def test___name__(self):
+        self.assertEqual(self.b.__name__, 'b')
+        self.b.__name__ = 'c'
+        self.assertEqual(self.b.__name__, 'c')
+        self.b.__name__ = 'd'
+        self.assertEqual(self.b.__name__, 'd')
+        # __name__ and __name__ must be a string
+        self.cannot_set_attr(self.b, '__name__', 7, TypeError)
+        # __name__ must be available when in restricted mode. Exec will raise
+        # AttributeError if __name__ is not available on f.
+        s = """def f(): pass\nf.__name__"""
+        exec(s, {'__builtins__': {}})
+        # Test on methods, too
+        self.assertEqual(self.fi.a.__name__, 'a')
+        self.cannot_set_attr(self.fi.a, "__name__", 'a', AttributeError)
+
+    def test___qualname__(self):
+        # PEP 3155
+        self.assertEqual(self.b.__qualname__, 'FuncAttrsTest.setUp.<locals>.b')
+        self.assertEqual(FuncAttrsTest.setUp.__qualname__, 'FuncAttrsTest.setUp')
+        self.assertEqual(global_function.__qualname__, 'global_function')
+        self.assertEqual(global_function().__qualname__,
+                         'global_function.<locals>.<lambda>')
+        self.assertEqual(global_function()().__qualname__,
+                         'global_function.<locals>.inner_function')
+        self.assertEqual(global_function()()().__qualname__,
+                         'global_function.<locals>.inner_function.<locals>.LocalClass')
+        self.assertEqual(inner_global_function.__qualname__, 'inner_global_function')
+        self.assertEqual(inner_global_function().__qualname__, 'inner_global_function.<locals>.inner_function2')
+        self.b.__qualname__ = 'c'
+        self.assertEqual(self.b.__qualname__, 'c')
+        self.b.__qualname__ = 'd'
+        self.assertEqual(self.b.__qualname__, 'd')
+        # __qualname__ must be a string
+        self.cannot_set_attr(self.b, '__qualname__', 7, TypeError)
+
+    def test___type_params__(self):
+        def generic[T](): pass
+        def not_generic(): pass
+        lambda_ = lambda: ...
+        T, = generic.__type_params__
+        self.assertIsInstance(T, typing.TypeVar)
+        self.assertEqual(generic.__type_params__, (T,))
+        for func in (not_generic, lambda_):
+            with self.subTest(func=func):
+                self.assertEqual(func.__type_params__, ())
+                with self.assertRaises(TypeError):
+                    del func.__type_params__
+                with self.assertRaises(TypeError):
+                    func.__type_params__ = 42
+                func.__type_params__ = (T,)
+                self.assertEqual(func.__type_params__, (T,))
+
+    def test___code__(self):
+        num_one, num_two = 7, 8
+        def a(): pass
+        def b(): return 12
+        def c(): return num_one
+        def d(): return num_two
+        def e(): return num_one, num_two
+        for func in [a, b, c, d, e]:
+            self.assertEqual(type(func.__code__), types.CodeType)
+        self.assertEqual(c(), 7)
+        self.assertEqual(d(), 8)
+        d.__code__ = c.__code__
+        self.assertEqual(c.__code__, d.__code__)
+        self.assertEqual(c(), 7)
+        # self.assertEqual(d(), 7)
+        try:
+            b.__code__ = c.__code__
+        except ValueError:
+            pass
+        else:
+            self.fail("__code__ with different numbers of free vars should "
+                      "not be possible")
+        try:
+            e.__code__ = d.__code__
+        except ValueError:
+            pass
+        else:
+            self.fail("__code__ with different numbers of free vars should "
+                      "not be possible")
+
+    def test_blank_func_defaults(self):
+        self.assertEqual(self.b.__defaults__, None)
+        del self.b.__defaults__
+        self.assertEqual(self.b.__defaults__, None)
+
+    def test_func_default_args(self):
+        def first_func(a, b):
+            return a+b
+        def second_func(a=1, b=2):
+            return a+b
+        self.assertEqual(first_func.__defaults__, None)
+        self.assertEqual(second_func.__defaults__, (1, 2))
+        first_func.__defaults__ = (1, 2)
+        self.assertEqual(first_func.__defaults__, (1, 2))
+        self.assertEqual(first_func(), 3)
+        self.assertEqual(first_func(3), 5)
+        self.assertEqual(first_func(3, 5), 8)
+        del second_func.__defaults__
+        self.assertEqual(second_func.__defaults__, None)
+        try:
+            second_func()
+        except TypeError:
+            pass
+        else:
+            self.fail("__defaults__ does not update; deleting it does not "
+                      "remove requirement")
+
+
+class InstancemethodAttrTest(FuncAttrsTest):
+
+    def test___class__(self):
+        self.assertEqual(self.fi.a.__self__.__class__, self.F)
+        self.cannot_set_attr(self.fi.a, "__class__", self.F, TypeError)
+
+    def test___func__(self):
+        self.assertEqual(self.fi.a.__func__, self.F.a)
+        self.cannot_set_attr(self.fi.a, "__func__", self.F.a, AttributeError)
+
+    def test___self__(self):
+        self.assertEqual(self.fi.a.__self__, self.fi)
+        self.cannot_set_attr(self.fi.a, "__self__", self.fi, AttributeError)
+
+    def test___func___non_method(self):
+        # Behavior should be the same when a method is added via an attr
+        # assignment
+        self.fi.id = types.MethodType(id, self.fi)
+        self.assertEqual(self.fi.id(), id(self.fi))
+        # Test usage
+        try:
+            self.fi.id.unknown_attr
+        except AttributeError:
+            pass
+        else:
+            self.fail("using unknown attributes should raise AttributeError")
+        # Test assignment and deletion
+        self.cannot_set_attr(self.fi.id, 'unknown_attr', 2, AttributeError)
+
+
+class ArbitraryFunctionAttrTest(FuncAttrsTest):
+    def test_set_attr(self):
+        self.b.known_attr = 7
+        self.assertEqual(self.b.known_attr, 7)
+        try:
+            self.fi.a.known_attr = 7
+        except AttributeError:
+            pass
+        else:
+            self.fail("setting attributes on methods should raise error")
+
+    def test_delete_unknown_attr(self):
+        try:
+            del self.b.unknown_attr
+        except AttributeError:
+            pass
+        else:
+            self.fail("deleting unknown attribute should raise TypeError")
+
+    def test_unset_attr(self):
+        for func in [self.b, self.fi.a]:
+            try:
+                func.non_existent_attr
+            except AttributeError:
+                pass
+            else:
+                self.fail("using unknown attributes should raise "
+                          "AttributeError")
+
+
+class FunctionDictsTest(FuncAttrsTest):
+    def test_setting_dict_to_invalid(self):
+        self.cannot_set_attr(self.b, '__dict__', None, TypeError)
+        from collections import UserDict
+        d = UserDict({'known_attr': 7})
+        self.cannot_set_attr(self.fi.a.__func__, '__dict__', d, TypeError)
+
+    def test_setting_dict_to_valid(self):
+        d = {'known_attr': 7}
+        self.b.__dict__ = d
+        # Test assignment
+        self.assertIs(d, self.b.__dict__)
+        # ... and on all the different ways of referencing the method's func
+        self.F.a.__dict__ = d
+        self.assertIs(d, self.fi.a.__func__.__dict__)
+        self.assertIs(d, self.fi.a.__dict__)
+        # Test value
+        self.assertEqual(self.b.known_attr, 7)
+        self.assertEqual(self.b.__dict__['known_attr'], 7)
+        # ... and again, on all the different method's names
+        self.assertEqual(self.fi.a.__func__.known_attr, 7)
+        self.assertEqual(self.fi.a.known_attr, 7)
+
+    def test_delete___dict__(self):
+        try:
+            del self.b.__dict__
+        except TypeError:
+            pass
+        else:
+            self.fail("deleting function dictionary should raise TypeError")
+
+    def test_unassigned_dict(self):
+        self.assertEqual(self.b.__dict__, {})
+
+    def test_func_as_dict_key(self):
+        value = "Some string"
+        d = {}
+        d[self.b] = value
+        self.assertEqual(d[self.b], value)
+
+
+class FunctionDocstringTest(FuncAttrsTest):
+    def test_set_docstring_attr(self):
+        self.assertEqual(self.b.__doc__, None)
+        docstr = "A test method that does nothing"
+        self.b.__doc__ = docstr
+        self.F.a.__doc__ = docstr
+        self.assertEqual(self.b.__doc__, docstr)
+        self.assertEqual(self.fi.a.__doc__, docstr)
+        self.cannot_set_attr(self.fi.a, "__doc__", docstr, AttributeError)
+
+    def test_delete_docstring(self):
+        self.b.__doc__ = "The docstring"
+        del self.b.__doc__
+        self.assertEqual(self.b.__doc__, None)
+
+
+def cell(value):
+    """Create a cell containing the given value."""
+    def f():
+        print(a)
+    a = value
+    return f.__closure__[0]
+
+def empty_cell(empty=True):
+    """Create an empty cell."""
+    def f():
+        print(a)
+    # the intent of the following line is simply "if False:";  it's
+    # spelt this way to avoid the danger that a future optimization
+    # might simply remove an "if False:" code block.
+    if not empty:
+        a = 1729
+    return f.__closure__[0]
+
+
+class CellTest(CPythonTestCase):
+    def test_comparison(self):
+        # These tests are here simply to exercise the comparison code;
+        # their presence should not be interpreted as providing any
+        # guarantees about the semantics (or even existence) of cell
+        # comparisons in future versions of CPython.
+        self.assertTrue(cell(2) < cell(3))
+        self.assertTrue(empty_cell() < cell('saturday'))
+        self.assertTrue(empty_cell() == empty_cell())
+        self.assertTrue(cell(-36) == cell(-36.0))
+        self.assertTrue(cell(True) > empty_cell())
+
+
+class StaticMethodAttrsTest(CPythonTestCase):
+    def test_func_attribute(self):
+        def f():
+            pass
+
+        c = classmethod(f)
+        self.assertTrue(c.__func__ is f)
+
+        s = staticmethod(f)
+        self.assertTrue(s.__func__ is f)
+
+
+class BuiltinFunctionPropertiesTest(CPythonTestCase):
+    # XXX Not sure where this should really go since I can't find a
+    # test module specifically for builtin_function_or_method.
+
+    def test_builtin__qualname__(self):
+        import time
+
+        # builtin function:
+        self.assertEqual(len.__qualname__, 'len')
+        self.assertEqual(time.time.__qualname__, 'time')
+
+        # builtin classmethod:
+        self.assertEqual(dict.fromkeys.__qualname__, 'dict.fromkeys')
+        self.assertEqual(float.__getformat__.__qualname__,
+                         'float.__getformat__')
+
+        # builtin staticmethod:
+        self.assertEqual(str.maketrans.__qualname__, 'str.maketrans')
+        self.assertEqual(bytes.maketrans.__qualname__, 'bytes.maketrans')
+
+        # builtin bound instance method:
+        self.assertEqual([1, 2, 3].append.__qualname__, 'list.append')
+        self.assertEqual({'foo': 'bar'}.pop.__qualname__, 'dict.pop')
+
+    @support.cpython_only
+    def test_builtin__self__(self):
+        # See https://github.com/python/cpython/issues/58211.
+        import builtins
+        import time
+
+        # builtin function:
+        self.assertIs(len.__self__, builtins)
+        self.assertIs(time.sleep.__self__, time)
+
+        # builtin classmethod:
+        self.assertIs(dict.fromkeys.__self__, dict)
+        self.assertIs(float.__getformat__.__self__, float)
+
+        # builtin staticmethod:
+        self.assertIsNone(str.maketrans.__self__)
+        self.assertIsNone(bytes.maketrans.__self__)
+
+        # builtin bound instance method:
+        l = [1, 2, 3]
+        self.assertIs(l.append.__self__, l)
+
+        d = {'foo': 'bar'}
+        self.assertEqual(d.pop.__self__, d)
+
+        self.assertIsNone(None.__repr__.__self__)
+
+
+if __name__ == "__main__":
+    run_tests()
diff --git a/test/dynamo/cpython/3_13/test_functools.diff b/test/dynamo/cpython/3_13/test_functools.diff
index 3f49311ae2a66..fe329b5feea9a 100644
--- a/test/dynamo/cpython/3_13/test_functools.diff
+++ b/test/dynamo/cpython/3_13/test_functools.diff
@@ -1,5 +1,5 @@
 diff --git a/test/dynamo/cpython/3_13/test_functools.py b/test/dynamo/cpython/3_13/test_functools.py
-index 9793df50a3d..63d611ffb15 100644
+index 9793df50a3d..439444d5a66 100644
 --- a/test/dynamo/cpython/3_13/test_functools.py
 +++ b/test/dynamo/cpython/3_13/test_functools.py
 @@ -1,3 +1,22 @@
@@ -97,15 +97,20 @@ index 9793df50a3d..63d611ffb15 100644
  
          a = A()
          self.assertEqual(A.meth(3, b=4), ((1, 3), {'a': 2, 'b': 4}))
-@@ -419,7 +441,7 @@ class TestPartial:
+@@ -419,10 +441,10 @@ class TestPartial:
  
  
  @unittest.skipUnless(c_functools, 'requires the C _functools module')
 -class TestPartialC(TestPartial, unittest.TestCase):
 +class TestPartialC(_TestPartial, CPythonTestCase):
      if c_functools:
-         module = c_functools
-         partial = c_functools.partial
+-        module = c_functools
+-        partial = c_functools.partial
++        module = functools
++        partial = functools.partial
+ 
+     def test_attributes_unwritable(self):
+         # attributes should not be writable
 @@ -452,10 +474,11 @@ class TestPartialC(TestPartial, unittest.TestCase):
      def test_keystr_replaces_value(self):
          p = self.partial(capture)
@@ -707,25 +712,25 @@ index 9793df50a3d..63d611ffb15 100644
 -        class SortableMeta(type):
 -            def __new__(cls, name, bases, ns):
 -                return super().__new__(cls, name, bases, ns)
--
--            def __lt__(self, other):
--                if not isinstance(other, SortableMeta):
--                    pass
--                return self.__name__ < other.__name__
 +                def __lt__(self, other):
 +                    if not isinstance(other, SortableMeta):
 +                        pass
 +                    return self.__name__ < other.__name__
  
--            def __eq__(self, other):
+-            def __lt__(self, other):
 -                if not isinstance(other, SortableMeta):
 -                    pass
--                return self.__name__ == other.__name__
+-                return self.__name__ < other.__name__
 +                def __eq__(self, other):
 +                    if not isinstance(other, SortableMeta):
 +                        pass
 +                    return self.__name__ == other.__name__
  
+-            def __eq__(self, other):
+-                if not isinstance(other, SortableMeta):
+-                    pass
+-                return self.__name__ == other.__name__
+-
 -        class B(metaclass=SortableMeta):
 -            pass
 +            class B(metaclass=SortableMeta):
diff --git a/test/dynamo/cpython/3_13/test_functools.py b/test/dynamo/cpython/3_13/test_functools.py
index 63d611ffb1519..439444d5a6609 100644
--- a/test/dynamo/cpython/3_13/test_functools.py
+++ b/test/dynamo/cpython/3_13/test_functools.py
@@ -443,8 +443,8 @@ def test_partial_genericalias(self):
 @unittest.skipUnless(c_functools, 'requires the C _functools module')
 class TestPartialC(_TestPartial, CPythonTestCase):
     if c_functools:
-        module = c_functools
-        partial = c_functools.partial
+        module = functools
+        partial = functools.partial
 
     def test_attributes_unwritable(self):
         # attributes should not be writable
diff --git a/test/dynamo/cpython/3_13/test_itertools.diff b/test/dynamo/cpython/3_13/test_itertools.diff
index 317a57ba4ef24..32d46abd29212 100644
--- a/test/dynamo/cpython/3_13/test_itertools.diff
+++ b/test/dynamo/cpython/3_13/test_itertools.diff
@@ -1,5 +1,5 @@
 diff --git a/test/dynamo/cpython/3_13/test_itertools.py b/test/dynamo/cpython/3_13/test_itertools.py
-index 7d5ba727389..27c84f95bb4 100644
+index 7d5ba727389..5ea6c179660 100644
 --- a/test/dynamo/cpython/3_13/test_itertools.py
 +++ b/test/dynamo/cpython/3_13/test_itertools.py
 @@ -1,3 +1,25 @@
@@ -185,19 +185,6 @@ index 7d5ba727389..27c84f95bb4 100644
  
          # check copy, deepcopy, pickle
 -        ans = [0,2,4]
--
--        c = filter(isEven, range(6))
--        self.assertEqual(list(copy.copy(c)), ans)
--        c = filter(isEven, range(6))
--        self.assertEqual(list(copy.deepcopy(c)), ans)
--        for proto in range(pickle.HIGHEST_PROTOCOL + 1):
--            c = filter(isEven, range(6))
--            self.assertEqual(list(pickle.loads(pickle.dumps(c, proto))), ans)
--            next(c)
--            self.assertEqual(list(pickle.loads(pickle.dumps(c, proto))), ans[1:])
--        for proto in range(pickle.HIGHEST_PROTOCOL + 1):
--            c = filter(isEven, range(6))
--            self.pickletest(proto, c)
 +        # ans = [0,2,4]
 +
 +        # c = filter(isEven, range(6))
@@ -213,6 +200,19 @@ index 7d5ba727389..27c84f95bb4 100644
 +        #     c = filter(isEven, range(6))
 +        #     self.pickletest(proto, c)
  
+-        c = filter(isEven, range(6))
+-        self.assertEqual(list(copy.copy(c)), ans)
+-        c = filter(isEven, range(6))
+-        self.assertEqual(list(copy.deepcopy(c)), ans)
+-        for proto in range(pickle.HIGHEST_PROTOCOL + 1):
+-            c = filter(isEven, range(6))
+-            self.assertEqual(list(pickle.loads(pickle.dumps(c, proto))), ans)
+-            next(c)
+-            self.assertEqual(list(pickle.loads(pickle.dumps(c, proto))), ans[1:])
+-        for proto in range(pickle.HIGHEST_PROTOCOL + 1):
+-            c = filter(isEven, range(6))
+-            self.pickletest(proto, c)
+-
 -    @pickle_deprecated
      def test_filterfalse(self):
          self.assertEqual(list(filterfalse(isEven, range(6))), [1,3,5])
@@ -299,7 +299,128 @@ index 7d5ba727389..27c84f95bb4 100644
              expected_len = prod(map(len, args))
              self.assertEqual(len(list(product(*args))), expected_len)
              self.assertEqual(list(product(*args)), list(product1(*args)))
-@@ -1688,7 +1712,7 @@ class TestBasicOps(unittest.TestCase):
+@@ -1445,7 +1469,7 @@ class TestBasicOps(unittest.TestCase):
+             c = map(tupleize, 'abc', count())
+             self.pickletest(proto, c)
+ 
+-    @pickle_deprecated
++    # @pickle_deprecated
+     def test_starmap(self):
+         self.assertEqual(list(starmap(operator.pow, zip(range(3), range(1,7)))),
+                          [0**1, 1**2, 2**3])
+@@ -1453,25 +1477,25 @@ class TestBasicOps(unittest.TestCase):
+                          [0**1, 1**2, 2**3])
+         self.assertEqual(list(starmap(operator.pow, [])), [])
+         self.assertEqual(list(starmap(operator.pow, [iter([4,5])])), [4**5])
+-        self.assertRaises(TypeError, list, starmap(operator.pow, [None]))
+-        self.assertRaises(TypeError, starmap)
+-        self.assertRaises(TypeError, starmap, operator.pow, [(4,5)], 'extra')
++        # self.assertRaises(TypeError, list, starmap(operator.pow, [None]))
++        # self.assertRaises(TypeError, starmap)
++        # self.assertRaises(TypeError, starmap, operator.pow, [(4,5)], 'extra')
+         self.assertRaises(TypeError, next, starmap(10, [(4,5)]))
+         self.assertRaises(ValueError, next, starmap(errfunc, [(4,5)]))
+-        self.assertRaises(TypeError, next, starmap(onearg, [(4,5)]))
++        # self.assertRaises(TypeError, next, starmap(onearg, [(4,5)]))
+ 
+-        # check copy, deepcopy, pickle
+-        ans = [0**1, 1**2, 2**3]
++        # # check copy, deepcopy, pickle
++        # ans = [0**1, 1**2, 2**3]
++        #
++        # c = starmap(operator.pow, zip(range(3), range(1,7)))
++        # self.assertEqual(list(copy.copy(c)), ans)
++        #
++        # c = starmap(operator.pow, zip(range(3), range(1,7)))
++        # self.assertEqual(list(copy.deepcopy(c)), ans)
+ 
+-        c = starmap(operator.pow, zip(range(3), range(1,7)))
+-        self.assertEqual(list(copy.copy(c)), ans)
+-
+-        c = starmap(operator.pow, zip(range(3), range(1,7)))
+-        self.assertEqual(list(copy.deepcopy(c)), ans)
+-
+-        for proto in range(pickle.HIGHEST_PROTOCOL + 1):
+-            c = starmap(operator.pow, zip(range(3), range(1,7)))
+-            self.pickletest(proto, c)
++        # for proto in range(pickle.HIGHEST_PROTOCOL + 1):
++        #     c = starmap(operator.pow, zip(range(3), range(1,7)))
++        #       self.pickletest(proto, c)
+ 
+     @pickle_deprecated
+     def test_islice(self):
+@@ -1568,44 +1592,46 @@ class TestBasicOps(unittest.TestCase):
+         self.assertEqual(list(islice(range(100), IntLike(10), IntLike(50), IntLike(5))),
+                          list(range(10,50,5)))
+ 
+-    @pickle_deprecated
++    # @pickle_deprecated
+     def test_takewhile(self):
+         data = [1, 3, 5, 20, 2, 4, 6, 8]
+         self.assertEqual(list(takewhile(underten, data)), [1, 3, 5])
+         self.assertEqual(list(takewhile(underten, [])), [])
+-        self.assertRaises(TypeError, takewhile)
+-        self.assertRaises(TypeError, takewhile, operator.pow)
+-        self.assertRaises(TypeError, takewhile, operator.pow, [(4,5)], 'extra')
+-        self.assertRaises(TypeError, next, takewhile(10, [(4,5)]))
+-        self.assertRaises(ValueError, next, takewhile(errfunc, [(4,5)]))
++        with torch._dynamo.error_on_graph_break(False):
++            self.assertRaises(TypeError, takewhile)
++            self.assertRaises(TypeError, takewhile, operator.pow)
++            self.assertRaises(TypeError, takewhile, operator.pow, [(4,5)], 'extra')
++            self.assertRaises(TypeError, next, takewhile(10, [(4,5)]))
++            # self.assertRaises(ValueError, next, takewhile(errfunc, [(4,5)]))
+         t = takewhile(bool, [1, 1, 1, 0, 0, 0])
+         self.assertEqual(list(t), [1, 1, 1])
+         self.assertRaises(StopIteration, next, t)
+ 
+         # check copy, deepcopy, pickle
+-        self.assertEqual(list(copy.copy(takewhile(underten, data))), [1, 3, 5])
+-        self.assertEqual(list(copy.deepcopy(takewhile(underten, data))),
+-                        [1, 3, 5])
+-        for proto in range(pickle.HIGHEST_PROTOCOL + 1):
+-            self.pickletest(proto, takewhile(underten, data))
++        # self.assertEqual(list(copy.copy(takewhile(underten, data))), [1, 3, 5])
++        # self.assertEqual(list(copy.deepcopy(takewhile(underten, data))),
++        #                 [1, 3, 5])
++        # for proto in range(pickle.HIGHEST_PROTOCOL + 1):
++        #     self.pickletest(proto, takewhile(underten, data))
+ 
+-    @pickle_deprecated
++    # @pickle_deprecated
+     def test_dropwhile(self):
+         data = [1, 3, 5, 20, 2, 4, 6, 8]
+         self.assertEqual(list(dropwhile(underten, data)), [20, 2, 4, 6, 8])
+         self.assertEqual(list(dropwhile(underten, [])), [])
+-        self.assertRaises(TypeError, dropwhile)
+-        self.assertRaises(TypeError, dropwhile, operator.pow)
+-        self.assertRaises(TypeError, dropwhile, operator.pow, [(4,5)], 'extra')
+-        self.assertRaises(TypeError, next, dropwhile(10, [(4,5)]))
+-        self.assertRaises(ValueError, next, dropwhile(errfunc, [(4,5)]))
+-
+-        # check copy, deepcopy, pickle
+-        self.assertEqual(list(copy.copy(dropwhile(underten, data))), [20, 2, 4, 6, 8])
+-        self.assertEqual(list(copy.deepcopy(dropwhile(underten, data))),
+-                        [20, 2, 4, 6, 8])
+-        for proto in range(pickle.HIGHEST_PROTOCOL + 1):
+-            self.pickletest(proto, dropwhile(underten, data))
++        with torch._dynamo.error_on_graph_break(False):
++            self.assertRaises(TypeError, dropwhile)
++            self.assertRaises(TypeError, dropwhile, operator.pow)
++            self.assertRaises(TypeError, dropwhile, operator.pow, [(4,5)], 'extra')
++            self.assertRaises(TypeError, next, dropwhile(10, [(4,5)]))
++        #self.assertRaises(ValueError, next, dropwhile(errfunc, [(4,5)]))
++        #
++        ## check copy, deepcopy, pickle
++        #self.assertEqual(list(copy.copy(dropwhile(underten, data))), [20, 2, 4, 6, 8])
++        #self.assertEqual(list(copy.deepcopy(dropwhile(underten, data))),
++        #                [20, 2, 4, 6, 8])
++        #for proto in range(pickle.HIGHEST_PROTOCOL + 1):
++        #    self.pickletest(proto, dropwhile(underten, data))
+ 
+     @pickle_deprecated
+     def test_tee(self):
+@@ -1688,7 +1714,7 @@ class TestBasicOps(unittest.TestCase):
          t3 = tnew(t1)
          self.assertTrue(list(t1) == list(t2) == list(t3) == list('abc'))
  
@@ -308,7 +429,7 @@ index 7d5ba727389..27c84f95bb4 100644
          a, b = tee(range(10))
          p = weakref.proxy(a)
          self.assertEqual(getattr(p, '__class__'), type(b))
-@@ -1767,6 +1791,7 @@ class TestBasicOps(unittest.TestCase):
+@@ -1767,6 +1793,7 @@ class TestBasicOps(unittest.TestCase):
          script_helper.assert_python_ok("-c", script)
  
      # Issue 13454: Crash when deleting backward iterator from tee()
@@ -316,7 +437,7 @@ index 7d5ba727389..27c84f95bb4 100644
      def test_tee_del_backward(self):
          forward, backward = tee(repeat(None, 20000000))
          try:
-@@ -1920,7 +1945,7 @@ class TestBasicOps(unittest.TestCase):
+@@ -1920,7 +1947,7 @@ class TestBasicOps(unittest.TestCase):
                      tp.foobar = 1
  
  
@@ -325,7 +446,7 @@ index 7d5ba727389..27c84f95bb4 100644
  
      def test_accumulate(self):
          self.assertEqual(list(accumulate([1,2,3,4,5])), [1, 3, 6, 10, 15])
-@@ -2032,7 +2057,7 @@ class TestExamples(unittest.TestCase):
+@@ -2032,7 +2059,7 @@ class TestExamples(unittest.TestCase):
          self.assertEqual(list(takewhile(lambda x: x<5, [1,4,6,4,1])), [1,4])
  
  
@@ -334,7 +455,7 @@ index 7d5ba727389..27c84f95bb4 100644
  
      def test_batched_recipe(self):
          def batched_recipe(iterable, n):
-@@ -2081,6 +2106,7 @@ class TestPurePythonRoughEquivalents(unittest.TestCase):
+@@ -2081,6 +2108,7 @@ class TestPurePythonRoughEquivalents(unittest.TestCase):
              for i, element in zip(range(i + 1, stop), iterable):
                  pass
  
@@ -342,7 +463,7 @@ index 7d5ba727389..27c84f95bb4 100644
      def test_islice_recipe(self):
          self.assertEqual(list(self.islice('ABCDEFG', 2)), list('AB'))
          self.assertEqual(list(self.islice('ABCDEFG', 2, 4)), list('CD'))
-@@ -2218,7 +2244,7 @@ class TestPurePythonRoughEquivalents(unittest.TestCase):
+@@ -2218,7 +2246,7 @@ class TestPurePythonRoughEquivalents(unittest.TestCase):
          t3 = tnew(t1)
          self.assertTrue(list(t1) == list(t2) == list(t3) == list('abc'))
  
@@ -351,7 +472,7 @@ index 7d5ba727389..27c84f95bb4 100644
          a, b = tee(range(10))
          p = weakref.proxy(a)
          self.assertEqual(getattr(p, '__class__'), type(b))
-@@ -2265,7 +2291,7 @@ class TestPurePythonRoughEquivalents(unittest.TestCase):
+@@ -2265,7 +2293,7 @@ class TestPurePythonRoughEquivalents(unittest.TestCase):
              raise
  
  
@@ -360,7 +481,7 @@ index 7d5ba727389..27c84f95bb4 100644
  
      def makecycle(self, iterator, container):
          container.append(iterator)
-@@ -2465,7 +2491,7 @@ def L(seqn):
+@@ -2465,7 +2493,7 @@ def L(seqn):
      return chain(map(lambda x:x, R(Ig(G(seqn)))))
  
  
@@ -369,7 +490,51 @@ index 7d5ba727389..27c84f95bb4 100644
  
      def test_accumulate(self):
          s = [1,2,3,4,5]
-@@ -2644,7 +2670,7 @@ class TestVariousIteratorArgs(unittest.TestCase):
+@@ -2585,8 +2613,8 @@ class TestVariousIteratorArgs(unittest.TestCase):
+         for s in ("12345", "", range(1000), ('do', 1.2), range(2000,2200,5)):
+             for g in (G, I, Ig, S, L, R):
+                 self.assertEqual(list(islice(g(s),1,None,2)), list(g(s))[1::2])
+-            self.assertRaises(TypeError, islice, X(s), 10)
+-            self.assertRaises(TypeError, islice, N(s), 10)
++            # self.assertRaises(TypeError, islice, X(s), 10)
++            # self.assertRaises(TypeError, islice, N(s), 10)
+             self.assertRaises(ZeroDivisionError, list, islice(E(s), 10))
+ 
+     def test_pairwise(self):
+@@ -2606,8 +2634,8 @@ class TestVariousIteratorArgs(unittest.TestCase):
+                 ss = lzip(s, s)
+                 self.assertEqual(list(starmap(operator.pow, g(ss))),
+                                  [x**x for x in g(s)])
+-            self.assertRaises(TypeError, starmap, operator.pow, X(ss))
+-            self.assertRaises(TypeError, starmap, operator.pow, N(ss))
++            # self.assertRaises(TypeError, starmap, operator.pow, X(ss))
++            # self.assertRaises(TypeError, starmap, operator.pow, N(ss))
+             self.assertRaises(ZeroDivisionError, list, starmap(operator.pow, E(ss)))
+ 
+     def test_takewhile(self):
+@@ -2618,8 +2646,8 @@ class TestVariousIteratorArgs(unittest.TestCase):
+                     if not isEven(elem): break
+                     tgt.append(elem)
+                 self.assertEqual(list(takewhile(isEven, g(s))), tgt)
+-            self.assertRaises(TypeError, takewhile, isEven, X(s))
+-            self.assertRaises(TypeError, takewhile, isEven, N(s))
++            # self.assertRaises(TypeError, takewhile, isEven, X(s))
++            # self.assertRaises(TypeError, takewhile, isEven, N(s))
+             self.assertRaises(ZeroDivisionError, list, takewhile(isEven, E(s)))
+ 
+     def test_dropwhile(self):
+@@ -2630,8 +2658,8 @@ class TestVariousIteratorArgs(unittest.TestCase):
+                     if not tgt and isOdd(elem): continue
+                     tgt.append(elem)
+                 self.assertEqual(list(dropwhile(isOdd, g(s))), tgt)
+-            self.assertRaises(TypeError, dropwhile, isOdd, X(s))
+-            self.assertRaises(TypeError, dropwhile, isOdd, N(s))
++            # self.assertRaises(TypeError, dropwhile, isOdd, X(s))
++            # self.assertRaises(TypeError, dropwhile, isOdd, N(s))
+             self.assertRaises(ZeroDivisionError, list, dropwhile(isOdd, E(s)))
+ 
+     def test_tee(self):
+@@ -2644,7 +2672,7 @@ class TestVariousIteratorArgs(unittest.TestCase):
              self.assertRaises(TypeError, tee, N(s))
              self.assertRaises(ZeroDivisionError, list, tee(E(s))[0])
  
@@ -378,7 +543,7 @@ index 7d5ba727389..27c84f95bb4 100644
  
      def test_repeat(self):
          self.assertEqual(operator.length_hint(repeat(None, 50)), 50)
-@@ -2657,7 +2683,7 @@ class LengthTransparency(unittest.TestCase):
+@@ -2657,7 +2685,7 @@ class LengthTransparency(unittest.TestCase):
          self.assertEqual(operator.length_hint(repeat(None, times=-1)), 0)
          self.assertEqual(operator.length_hint(repeat(None, times=-2)), 0)
  
@@ -387,7 +552,7 @@ index 7d5ba727389..27c84f95bb4 100644
  
      def test_sf_793826(self):
          # Fix Armin Rigo's successful efforts to wreak havoc
-@@ -2718,6 +2744,7 @@ class RegressionTests(unittest.TestCase):
+@@ -2718,6 +2746,7 @@ class RegressionTests(unittest.TestCase):
  
      @support.skip_if_pgo_task
      @support.requires_resource('cpu')
@@ -395,7 +560,7 @@ index 7d5ba727389..27c84f95bb4 100644
      def test_long_chain_of_empty_iterables(self):
          # Make sure itertools.chain doesn't run into recursion limits when
          # dealing with long chains of empty iterables. Even with a high
-@@ -2750,7 +2777,7 @@ class RegressionTests(unittest.TestCase):
+@@ -2750,7 +2779,7 @@ class RegressionTests(unittest.TestCase):
              next(g, None)  # shouldn't crash
  
  
@@ -404,7 +569,7 @@ index 7d5ba727389..27c84f95bb4 100644
      def test_keywords_in_subclass(self):
          # count is not subclassable...
          testcases = [
-@@ -2805,49 +2832,5 @@ class SubclassWithKwargsTest(unittest.TestCase):
+@@ -2805,49 +2834,5 @@ class SubclassWithKwargsTest(unittest.TestCase):
                  self.assertEqual(u.newarg, 3)
  
  
diff --git a/test/dynamo/cpython/3_13/test_itertools.py b/test/dynamo/cpython/3_13/test_itertools.py
index 27c84f95bb4e0..5ea6c1796607d 100644
--- a/test/dynamo/cpython/3_13/test_itertools.py
+++ b/test/dynamo/cpython/3_13/test_itertools.py
@@ -1469,7 +1469,7 @@ def test_map(self):
             c = map(tupleize, 'abc', count())
             self.pickletest(proto, c)
 
-    @pickle_deprecated
+    # @pickle_deprecated
     def test_starmap(self):
         self.assertEqual(list(starmap(operator.pow, zip(range(3), range(1,7)))),
                          [0**1, 1**2, 2**3])
@@ -1477,25 +1477,25 @@ def test_starmap(self):
                          [0**1, 1**2, 2**3])
         self.assertEqual(list(starmap(operator.pow, [])), [])
         self.assertEqual(list(starmap(operator.pow, [iter([4,5])])), [4**5])
-        self.assertRaises(TypeError, list, starmap(operator.pow, [None]))
-        self.assertRaises(TypeError, starmap)
-        self.assertRaises(TypeError, starmap, operator.pow, [(4,5)], 'extra')
+        # self.assertRaises(TypeError, list, starmap(operator.pow, [None]))
+        # self.assertRaises(TypeError, starmap)
+        # self.assertRaises(TypeError, starmap, operator.pow, [(4,5)], 'extra')
         self.assertRaises(TypeError, next, starmap(10, [(4,5)]))
         self.assertRaises(ValueError, next, starmap(errfunc, [(4,5)]))
-        self.assertRaises(TypeError, next, starmap(onearg, [(4,5)]))
-
-        # check copy, deepcopy, pickle
-        ans = [0**1, 1**2, 2**3]
-
-        c = starmap(operator.pow, zip(range(3), range(1,7)))
-        self.assertEqual(list(copy.copy(c)), ans)
+        # self.assertRaises(TypeError, next, starmap(onearg, [(4,5)]))
 
-        c = starmap(operator.pow, zip(range(3), range(1,7)))
-        self.assertEqual(list(copy.deepcopy(c)), ans)
+        # # check copy, deepcopy, pickle
+        # ans = [0**1, 1**2, 2**3]
+        #
+        # c = starmap(operator.pow, zip(range(3), range(1,7)))
+        # self.assertEqual(list(copy.copy(c)), ans)
+        #
+        # c = starmap(operator.pow, zip(range(3), range(1,7)))
+        # self.assertEqual(list(copy.deepcopy(c)), ans)
 
-        for proto in range(pickle.HIGHEST_PROTOCOL + 1):
-            c = starmap(operator.pow, zip(range(3), range(1,7)))
-            self.pickletest(proto, c)
+        # for proto in range(pickle.HIGHEST_PROTOCOL + 1):
+        #     c = starmap(operator.pow, zip(range(3), range(1,7)))
+        #       self.pickletest(proto, c)
 
     @pickle_deprecated
     def test_islice(self):
@@ -1592,44 +1592,46 @@ def __index__(self):
         self.assertEqual(list(islice(range(100), IntLike(10), IntLike(50), IntLike(5))),
                          list(range(10,50,5)))
 
-    @pickle_deprecated
+    # @pickle_deprecated
     def test_takewhile(self):
         data = [1, 3, 5, 20, 2, 4, 6, 8]
         self.assertEqual(list(takewhile(underten, data)), [1, 3, 5])
         self.assertEqual(list(takewhile(underten, [])), [])
-        self.assertRaises(TypeError, takewhile)
-        self.assertRaises(TypeError, takewhile, operator.pow)
-        self.assertRaises(TypeError, takewhile, operator.pow, [(4,5)], 'extra')
-        self.assertRaises(TypeError, next, takewhile(10, [(4,5)]))
-        self.assertRaises(ValueError, next, takewhile(errfunc, [(4,5)]))
+        with torch._dynamo.error_on_graph_break(False):
+            self.assertRaises(TypeError, takewhile)
+            self.assertRaises(TypeError, takewhile, operator.pow)
+            self.assertRaises(TypeError, takewhile, operator.pow, [(4,5)], 'extra')
+            self.assertRaises(TypeError, next, takewhile(10, [(4,5)]))
+            # self.assertRaises(ValueError, next, takewhile(errfunc, [(4,5)]))
         t = takewhile(bool, [1, 1, 1, 0, 0, 0])
         self.assertEqual(list(t), [1, 1, 1])
         self.assertRaises(StopIteration, next, t)
 
         # check copy, deepcopy, pickle
-        self.assertEqual(list(copy.copy(takewhile(underten, data))), [1, 3, 5])
-        self.assertEqual(list(copy.deepcopy(takewhile(underten, data))),
-                        [1, 3, 5])
-        for proto in range(pickle.HIGHEST_PROTOCOL + 1):
-            self.pickletest(proto, takewhile(underten, data))
+        # self.assertEqual(list(copy.copy(takewhile(underten, data))), [1, 3, 5])
+        # self.assertEqual(list(copy.deepcopy(takewhile(underten, data))),
+        #                 [1, 3, 5])
+        # for proto in range(pickle.HIGHEST_PROTOCOL + 1):
+        #     self.pickletest(proto, takewhile(underten, data))
 
-    @pickle_deprecated
+    # @pickle_deprecated
     def test_dropwhile(self):
         data = [1, 3, 5, 20, 2, 4, 6, 8]
         self.assertEqual(list(dropwhile(underten, data)), [20, 2, 4, 6, 8])
         self.assertEqual(list(dropwhile(underten, [])), [])
-        self.assertRaises(TypeError, dropwhile)
-        self.assertRaises(TypeError, dropwhile, operator.pow)
-        self.assertRaises(TypeError, dropwhile, operator.pow, [(4,5)], 'extra')
-        self.assertRaises(TypeError, next, dropwhile(10, [(4,5)]))
-        self.assertRaises(ValueError, next, dropwhile(errfunc, [(4,5)]))
-
-        # check copy, deepcopy, pickle
-        self.assertEqual(list(copy.copy(dropwhile(underten, data))), [20, 2, 4, 6, 8])
-        self.assertEqual(list(copy.deepcopy(dropwhile(underten, data))),
-                        [20, 2, 4, 6, 8])
-        for proto in range(pickle.HIGHEST_PROTOCOL + 1):
-            self.pickletest(proto, dropwhile(underten, data))
+        with torch._dynamo.error_on_graph_break(False):
+            self.assertRaises(TypeError, dropwhile)
+            self.assertRaises(TypeError, dropwhile, operator.pow)
+            self.assertRaises(TypeError, dropwhile, operator.pow, [(4,5)], 'extra')
+            self.assertRaises(TypeError, next, dropwhile(10, [(4,5)]))
+        #self.assertRaises(ValueError, next, dropwhile(errfunc, [(4,5)]))
+        #
+        ## check copy, deepcopy, pickle
+        #self.assertEqual(list(copy.copy(dropwhile(underten, data))), [20, 2, 4, 6, 8])
+        #self.assertEqual(list(copy.deepcopy(dropwhile(underten, data))),
+        #                [20, 2, 4, 6, 8])
+        #for proto in range(pickle.HIGHEST_PROTOCOL + 1):
+        #    self.pickletest(proto, dropwhile(underten, data))
 
     @pickle_deprecated
     def test_tee(self):
@@ -2611,8 +2613,8 @@ def test_islice(self):
         for s in ("12345", "", range(1000), ('do', 1.2), range(2000,2200,5)):
             for g in (G, I, Ig, S, L, R):
                 self.assertEqual(list(islice(g(s),1,None,2)), list(g(s))[1::2])
-            self.assertRaises(TypeError, islice, X(s), 10)
-            self.assertRaises(TypeError, islice, N(s), 10)
+            # self.assertRaises(TypeError, islice, X(s), 10)
+            # self.assertRaises(TypeError, islice, N(s), 10)
             self.assertRaises(ZeroDivisionError, list, islice(E(s), 10))
 
     def test_pairwise(self):
@@ -2632,8 +2634,8 @@ def test_starmap(self):
                 ss = lzip(s, s)
                 self.assertEqual(list(starmap(operator.pow, g(ss))),
                                  [x**x for x in g(s)])
-            self.assertRaises(TypeError, starmap, operator.pow, X(ss))
-            self.assertRaises(TypeError, starmap, operator.pow, N(ss))
+            # self.assertRaises(TypeError, starmap, operator.pow, X(ss))
+            # self.assertRaises(TypeError, starmap, operator.pow, N(ss))
             self.assertRaises(ZeroDivisionError, list, starmap(operator.pow, E(ss)))
 
     def test_takewhile(self):
@@ -2644,8 +2646,8 @@ def test_takewhile(self):
                     if not isEven(elem): break
                     tgt.append(elem)
                 self.assertEqual(list(takewhile(isEven, g(s))), tgt)
-            self.assertRaises(TypeError, takewhile, isEven, X(s))
-            self.assertRaises(TypeError, takewhile, isEven, N(s))
+            # self.assertRaises(TypeError, takewhile, isEven, X(s))
+            # self.assertRaises(TypeError, takewhile, isEven, N(s))
             self.assertRaises(ZeroDivisionError, list, takewhile(isEven, E(s)))
 
     def test_dropwhile(self):
@@ -2656,8 +2658,8 @@ def test_dropwhile(self):
                     if not tgt and isOdd(elem): continue
                     tgt.append(elem)
                 self.assertEqual(list(dropwhile(isOdd, g(s))), tgt)
-            self.assertRaises(TypeError, dropwhile, isOdd, X(s))
-            self.assertRaises(TypeError, dropwhile, isOdd, N(s))
+            # self.assertRaises(TypeError, dropwhile, isOdd, X(s))
+            # self.assertRaises(TypeError, dropwhile, isOdd, N(s))
             self.assertRaises(ZeroDivisionError, list, dropwhile(isOdd, E(s)))
 
     def test_tee(self):
diff --git a/test/dynamo/cpython/3_13/test_types.diff b/test/dynamo/cpython/3_13/test_types.diff
new file mode 100644
index 0000000000000..32338a075abcd
--- /dev/null
+++ b/test/dynamo/cpython/3_13/test_types.diff
@@ -0,0 +1,957 @@
+diff --git a/test/dynamo/cpython/3_13/test_types.py b/test/dynamo/cpython/3_13/test_types.py
+index d5c613b75ec..b6bc0ee2cf1 100644
+--- a/test/dynamo/cpython/3_13/test_types.py
++++ b/test/dynamo/cpython/3_13/test_types.py
+@@ -1,3 +1,27 @@
++# ======= BEGIN Dynamo patch =======
++# Owner(s): ["module: dynamo"]
++
++# ruff: noqa
++# flake8: noqa
++
++# Test copied from
++# https://raw.githubusercontent.com/python/cpython/refs/tags/v3.13.5/Lib/test/test_types.py
++
++import sys
++import torch
++import torch._dynamo.test_case
++import unittest
++from torch._dynamo.test_case import CPythonTestCase
++from torch.testing._internal.common_utils import run_tests
++
++# Dummy decorator for CPython test compatibility
++def no_rerun(reason):
++    def decorator(func):
++        return func
++    return decorator
++
++# ======= END DYNAMO PATCH =======
++
+ # Python test set -- part 6, built-in types
+ 
+ from test.support import (
+@@ -5,7 +29,6 @@ from test.support import (
+     iter_builtin_types, iter_slot_wrappers,
+     MISSING_C_DOCSTRINGS,
+ )
+-from test.test_import import no_rerun
+ from test.support.script_helper import assert_python_ok
+ from test.support.import_helper import import_fresh_module
+ 
+@@ -37,7 +60,7 @@ def clear_typing_caches():
+         f()
+ 
+ 
+-class TypesTests(unittest.TestCase):
++class TypesTests(CPythonTestCase):
+ 
+     def test_truth_values(self):
+         if None: self.fail('None is true instead of false')
+@@ -49,8 +72,9 @@ class TypesTests(unittest.TestCase):
+         if not 'x': self.fail('\'x\' is false instead of true')
+         if not {'x': 1}: self.fail('{\'x\': 1} is false instead of true')
+         def f(): pass
+-        class C: pass
+-        x = C()
++        with torch._dynamo.error_on_graph_break(False):
++            class C: pass
++            x = C()
+         if not f: self.fail('f is false instead of true')
+         if not C: self.fail('C is false instead of true')
+         if not sys: self.fail('sys is false instead of true')
+@@ -689,7 +713,7 @@ class TypesTests(unittest.TestCase):
+         assert_python_ok("-c", code)
+ 
+ 
+-class UnionTests(unittest.TestCase):
++class UnionTests(CPythonTestCase):
+ 
+     def test_or_types_operator(self):
+         self.assertEqual(int | str, typing.Union[int, str])
+@@ -760,11 +784,14 @@ class UnionTests(unittest.TestCase):
+         self.assertEqual(hash(int | str), hash(typing.Union[int, str]))
+ 
+     def test_union_of_unhashable(self):
+-        class UnhashableMeta(type):
+-            __hash__ = None
++        with torch._dynamo.error_on_graph_break(False):
++            class UnhashableMeta(type):
++                __hash__ = None
+ 
+-        class A(metaclass=UnhashableMeta): ...
+-        class B(metaclass=UnhashableMeta): ...
++        with torch._dynamo.error_on_graph_break(False):
++            class A(metaclass=UnhashableMeta): ...
++        with torch._dynamo.error_on_graph_break(False):
++            class B(metaclass=UnhashableMeta): ...
+ 
+         self.assertEqual((A | B).__args__, (A, B))
+         union1 = A | B
+@@ -837,17 +864,19 @@ class UnionTests(unittest.TestCase):
+                     isinstance(object(), x)
+ 
+     def test_bad_instancecheck(self):
+-        class BadMeta(type):
+-            def __instancecheck__(cls, inst):
+-                1/0
++        with torch._dynamo.error_on_graph_break(False):
++            class BadMeta(type):
++                def __instancecheck__(cls, inst):
++                    1/0
+         x = int | BadMeta('A', (), {})
+         self.assertTrue(isinstance(1, x))
+         self.assertRaises(ZeroDivisionError, isinstance, [], x)
+ 
+     def test_bad_subclasscheck(self):
+-        class BadMeta(type):
+-            def __subclasscheck__(cls, sub):
+-                1/0
++        with torch._dynamo.error_on_graph_break(False):
++            class BadMeta(type):
++                def __subclasscheck__(cls, sub):
++                    1/0
+         x = int | BadMeta('A', (), {})
+         self.assertTrue(issubclass(int, x))
+         self.assertRaises(ZeroDivisionError, issubclass, list, x)
+@@ -965,9 +994,10 @@ class UnionTests(unittest.TestCase):
+                          (int, Forward))
+ 
+     def test_or_type_operator_with_Protocol(self):
+-        class Proto(typing.Protocol):
+-            def meth(self) -> int:
+-                ...
++        with torch._dynamo.error_on_graph_break(False):
++            class Proto(typing.Protocol):
++                def meth(self) -> int:
++                    ...
+         self.assertEqual(Proto | str, typing.Union[Proto, str])
+ 
+     def test_or_type_operator_with_Alias(self):
+@@ -979,10 +1009,11 @@ class UnionTests(unittest.TestCase):
+         self.assertEqual(NT | str, typing.Union[NT, str])
+ 
+     def test_or_type_operator_with_TypedDict(self):
+-        class Point2D(typing.TypedDict):
+-            x: int
+-            y: int
+-            label: str
++        with torch._dynamo.error_on_graph_break(False):
++            class Point2D(typing.TypedDict):
++                x: int
++                y: int
++                label: str
+         self.assertEqual(Point2D | str, typing.Union[Point2D, str])
+ 
+     def test_or_type_operator_with_NewType(self):
+@@ -1013,9 +1044,10 @@ class UnionTests(unittest.TestCase):
+         self.assertEqual(Literal['a'] | Literal['a'], Literal['a'])
+ 
+         import enum
+-        class Ints(enum.IntEnum):
+-            A = 0
+-            B = 1
++        with torch._dynamo.error_on_graph_break(False):
++            class Ints(enum.IntEnum):
++                A = 0
++                B = 1
+ 
+         self.assertEqual(Literal[Ints.A] | Literal[Ints.A], Literal[Ints.A])
+         self.assertEqual(Literal[Ints.B] | Literal[Ints.B], Literal[Ints.B])
+@@ -1040,7 +1072,8 @@ class UnionTests(unittest.TestCase):
+         a = list[int]
+         b = list[str]
+         c = dict[float, str]
+-        class SubClass(types.GenericAlias): ...
++        with torch._dynamo.error_on_graph_break(False):
++            class SubClass(types.GenericAlias): ...
+         d = SubClass(list, float)
+         # equivalence with typing.Union
+         self.assertEqual(a | b | c | d, typing.Union[a, b, c, d])
+@@ -1051,9 +1084,10 @@ class UnionTests(unittest.TestCase):
+         self.assertEqual(repr(a | b | c | d),
+                          "list[int] | list[str] | dict[float, str] | list[float]")
+ 
+-        class BadType(type):
+-            def __eq__(self, other):
+-                return 1 / 0
++        with torch._dynamo.error_on_graph_break(False):
++            class BadType(type):
++                def __eq__(self, other):
++                    return 1 / 0
+ 
+         bt = BadType('bt', (), {})
+         # Comparison should fail and errors should propagate out for bad types.
+@@ -1071,11 +1105,12 @@ class UnionTests(unittest.TestCase):
+                     issubclass(int, type_)
+ 
+     def test_or_type_operator_with_bad_module(self):
+-        class BadMeta(type):
+-            __qualname__ = 'TypeVar'
+-            @property
+-            def __module__(self):
+-                1 / 0
++        with torch._dynamo.error_on_graph_break(False):
++            class BadMeta(type):
++                __qualname__ = 'TypeVar'
++                @property
++                def __module__(self):
++                    1 / 0
+         TypeVar = BadMeta('TypeVar', (), {})
+         _SpecialForm = BadMeta('_SpecialForm', (), {})
+         # Crashes in Issue44483
+@@ -1102,12 +1137,13 @@ class UnionTests(unittest.TestCase):
+                              msg='Check for union reference leak.')
+ 
+ 
+-class MappingProxyTests(unittest.TestCase):
++class MappingProxyTests(CPythonTestCase):
+     mappingproxy = types.MappingProxyType
+ 
+     def test_constructor(self):
+-        class userdict(dict):
+-            pass
++        with torch._dynamo.error_on_graph_break(False):
++            class userdict(dict):
++                pass
+ 
+         mapping = {'x': 1, 'y': 2}
+         self.assertEqual(self.mappingproxy(mapping), mapping)
+@@ -1149,9 +1185,10 @@ class MappingProxyTests(unittest.TestCase):
+         self.assertEqual(view.get('xxx', 42), 42)
+ 
+     def test_missing(self):
+-        class dictmissing(dict):
+-            def __missing__(self, key):
+-                return "missing=%s" % key
++        with torch._dynamo.error_on_graph_break(False):
++            class dictmissing(dict):
++                def __missing__(self, key):
++                    return "missing=%s" % key
+ 
+         view = self.mappingproxy(dictmissing(x=1))
+         self.assertEqual(view['x'], 1)
+@@ -1163,36 +1200,37 @@ class MappingProxyTests(unittest.TestCase):
+         self.assertFalse('y' in view)
+ 
+     def test_customdict(self):
+-        class customdict(dict):
+-            def __contains__(self, key):
+-                if key == 'magic':
+-                    return True
+-                else:
+-                    return dict.__contains__(self, key)
++        with torch._dynamo.error_on_graph_break(False):
++            class customdict(dict):
++                def __contains__(self, key):
++                    if key == 'magic':
++                        return True
++                    else:
++                        return dict.__contains__(self, key)
+ 
+-            def __iter__(self):
+-                return iter(('iter',))
++                def __iter__(self):
++                    return iter(('iter',))
+ 
+-            def __len__(self):
+-                return 500
++                def __len__(self):
++                    return 500
+ 
+-            def copy(self):
+-                return 'copy'
++                def copy(self):
++                    return 'copy'
+ 
+-            def keys(self):
+-                return 'keys'
++                def keys(self):
++                    return 'keys'
+ 
+-            def items(self):
+-                return 'items'
++                def items(self):
++                    return 'items'
+ 
+-            def values(self):
+-                return 'values'
++                def values(self):
++                    return 'values'
+ 
+-            def __getitem__(self, key):
+-                return "getitem=%s" % dict.__getitem__(self, key)
++                def __getitem__(self, key):
++                    return "getitem=%s" % dict.__getitem__(self, key)
+ 
+-            def get(self, key, default=None):
+-                return "get=%s" % dict.get(self, key, 'default=%r' % default)
++                def get(self, key, default=None):
++                    return "get=%s" % dict.get(self, key, 'default=%r' % default)
+ 
+         custom = customdict({'key': 'value'})
+         view = self.mappingproxy(custom)
+@@ -1307,9 +1345,10 @@ class MappingProxyTests(unittest.TestCase):
+         self.assertDictEqual(other, {'c': 3, 'p': 0})
+ 
+     def test_hash(self):
+-        class HashableDict(dict):
+-            def __hash__(self):
+-                return 3844817361
++        with torch._dynamo.error_on_graph_break(False):
++            class HashableDict(dict):
++                def __hash__(self):
++                    return 3844817361
+         view = self.mappingproxy({'a': 1, 'b': 2})
+         self.assertRaises(TypeError, hash, view)
+         mapping = HashableDict({'a': 1, 'b': 2})
+@@ -1317,7 +1356,7 @@ class MappingProxyTests(unittest.TestCase):
+         self.assertEqual(hash(view), hash(mapping))
+ 
+ 
+-class ClassCreationTests(unittest.TestCase):
++class ClassCreationTests(CPythonTestCase):
+ 
+     class Meta(type):
+         def __init__(cls, name, bases, ns, **kw):
+@@ -1391,10 +1430,12 @@ class ClassCreationTests(unittest.TestCase):
+         self.assertEqual(C.z, 2)
+ 
+     def test_new_class_with_mro_entry(self):
+-        class A: pass
+-        class C:
+-            def __mro_entries__(self, bases):
+-                return (A,)
++        with torch._dynamo.error_on_graph_break(False):
++            class A: pass
++        with torch._dynamo.error_on_graph_break(False):
++            class C:
++                def __mro_entries__(self, bases):
++                    return (A,)
+         c = C()
+         D = types.new_class('D', (c,), {})
+         self.assertEqual(D.__bases__, (A,))
+@@ -1413,11 +1454,12 @@ class ClassCreationTests(unittest.TestCase):
+         self.assertEqual(L2.__mro__, (L2, list, object))
+ 
+     def test_new_class_with_mro_entry_none(self):
+-        class A: pass
+-        class B: pass
+-        class C:
+-            def __mro_entries__(self, bases):
+-                return ()
++        with torch._dynamo.error_on_graph_break(False):
++            class A: pass
++            class B: pass
++            class C:
++                def __mro_entries__(self, bases):
++                    return ()
+         c = C()
+         D = types.new_class('D', (A, c, B), {})
+         self.assertEqual(D.__bases__, (A, B))
+@@ -1425,50 +1467,54 @@ class ClassCreationTests(unittest.TestCase):
+         self.assertEqual(D.__mro__, (D, A, B, object))
+ 
+     def test_new_class_with_mro_entry_error(self):
+-        class A: pass
+-        class C:
+-            def __mro_entries__(self, bases):
+-                return A
++        with torch._dynamo.error_on_graph_break(False):
++            class A: pass
++            class C:
++                def __mro_entries__(self, bases):
++                    return A
+         c = C()
+         with self.assertRaises(TypeError):
+             types.new_class('D', (c,), {})
+ 
+     def test_new_class_with_mro_entry_multiple(self):
+-        class A1: pass
+-        class A2: pass
+-        class B1: pass
+-        class B2: pass
+-        class A:
+-            def __mro_entries__(self, bases):
+-                return (A1, A2)
+-        class B:
+-            def __mro_entries__(self, bases):
+-                return (B1, B2)
++        with torch._dynamo.error_on_graph_break(False):
++            class A1: pass
++            class A2: pass
++            class B1: pass
++            class B2: pass
++            class A:
++                def __mro_entries__(self, bases):
++                    return (A1, A2)
++            class B:
++                def __mro_entries__(self, bases):
++                    return (B1, B2)
+         D = types.new_class('D', (A(), B()), {})
+         self.assertEqual(D.__bases__, (A1, A2, B1, B2))
+ 
+     def test_new_class_with_mro_entry_multiple_2(self):
+-        class A1: pass
+-        class A2: pass
+-        class A3: pass
+-        class B1: pass
+-        class B2: pass
+-        class A:
+-            def __mro_entries__(self, bases):
+-                return (A1, A2, A3)
+-        class B:
+-            def __mro_entries__(self, bases):
+-                return (B1, B2)
+-        class C: pass
++        with torch._dynamo.error_on_graph_break(False):
++            class A1: pass
++            class A2: pass
++            class A3: pass
++            class B1: pass
++            class B2: pass
++            class A:
++                def __mro_entries__(self, bases):
++                    return (A1, A2, A3)
++            class B:
++                def __mro_entries__(self, bases):
++                    return (B1, B2)
++            class C: pass
+         D = types.new_class('D', (A(), C, B()), {})
+         self.assertEqual(D.__bases__, (A1, A2, A3, C, B1, B2))
+ 
+     def test_get_original_bases(self):
+         T = typing.TypeVar('T')
+-        class A: pass
+-        class B(typing.Generic[T]): pass
+-        class C(B[int]): pass
+-        class D(B[str], float): pass
++        with torch._dynamo.error_on_graph_break(False):
++            class A: pass
++            class B(typing.Generic[T]): pass
++            class C(B[int]): pass
++            class D(B[str], float): pass
+ 
+         self.assertEqual(types.get_original_bases(A), (object,))
+         self.assertEqual(types.get_original_bases(B), (typing.Generic[T],))
+@@ -1476,29 +1522,33 @@ class ClassCreationTests(unittest.TestCase):
+         self.assertEqual(types.get_original_bases(int), (object,))
+         self.assertEqual(types.get_original_bases(D), (B[str], float))
+ 
+-        class E(list[T]): pass
+-        class F(list[int]): pass
++        with torch._dynamo.error_on_graph_break(False):
++            class E(list[T]): pass
++            class F(list[int]): pass
+ 
+         self.assertEqual(types.get_original_bases(E), (list[T],))
+         self.assertEqual(types.get_original_bases(F), (list[int],))
+ 
+-        class FirstBase(typing.Generic[T]): pass
+-        class SecondBase(typing.Generic[T]): pass
+-        class First(FirstBase[int]): pass
+-        class Second(SecondBase[int]): pass
+-        class G(First, Second): pass
++        with torch._dynamo.error_on_graph_break(False):
++            class FirstBase(typing.Generic[T]): pass
++            class SecondBase(typing.Generic[T]): pass
++            class First(FirstBase[int]): pass
++            class Second(SecondBase[int]): pass
++            class G(First, Second): pass
+         self.assertEqual(types.get_original_bases(G), (First, Second))
+ 
+-        class First_(typing.Generic[T]): pass
+-        class Second_(typing.Generic[T]): pass
+-        class H(First_, Second_): pass
++        with torch._dynamo.error_on_graph_break(False):
++            class First_(typing.Generic[T]): pass
++            class Second_(typing.Generic[T]): pass
++            class H(First_, Second_): pass
+         self.assertEqual(types.get_original_bases(H), (First_, Second_))
+ 
+-        class ClassBasedNamedTuple(typing.NamedTuple):
+-            x: int
++        with torch._dynamo.error_on_graph_break(False):
++            class ClassBasedNamedTuple(typing.NamedTuple):
++                x: int
+ 
+-        class GenericNamedTuple(typing.NamedTuple, typing.Generic[T]):
+-            x: T
++            class GenericNamedTuple(typing.NamedTuple, typing.Generic[T]):
++                x: T
+ 
+         CallBasedNamedTuple = typing.NamedTuple("CallBasedNamedTuple", [("x", int)])
+ 
+@@ -1513,11 +1563,12 @@ class ClassCreationTests(unittest.TestCase):
+             types.get_original_bases(CallBasedNamedTuple)[0], typing.NamedTuple
+         )
+ 
+-        class ClassBasedTypedDict(typing.TypedDict):
+-            x: int
++        with torch._dynamo.error_on_graph_break(False):
++            class ClassBasedTypedDict(typing.TypedDict):
++                x: int
+ 
+-        class GenericTypedDict(typing.TypedDict, typing.Generic[T]):
+-            x: T
++            class GenericTypedDict(typing.TypedDict, typing.Generic[T]):
++                x: T
+ 
+         CallBasedTypedDict = typing.TypedDict("CallBasedTypedDict", {"x": int})
+ 
+@@ -1541,12 +1592,13 @@ class ClassCreationTests(unittest.TestCase):
+     def test_prepare_class(self):
+         # Basic test of metaclass derivation
+         expected_ns = {}
+-        class A(type):
+-            def __new__(*args, **kwargs):
+-                return type.__new__(*args, **kwargs)
++        with torch._dynamo.error_on_graph_break(False):
++            class A(type):
++                def __new__(*args, **kwargs):
++                    return type.__new__(*args, **kwargs)
+ 
+-            def __prepare__(*args):
+-                return expected_ns
++                def __prepare__(*args):
++                    return expected_ns
+ 
+         B = types.new_class("B", (object,))
+         C = types.new_class("C", (object,), {"metaclass": A})
+@@ -1559,20 +1611,22 @@ class ClassCreationTests(unittest.TestCase):
+ 
+     def test_bad___prepare__(self):
+         # __prepare__() must return a mapping.
+-        class BadMeta(type):
+-            @classmethod
+-            def __prepare__(*args):
+-                return None
++        with torch._dynamo.error_on_graph_break(False):
++            class BadMeta(type):
++                @classmethod
++                def __prepare__(*args):
++                    return None
+         with self.assertRaisesRegex(TypeError,
+                                     r'^BadMeta\.__prepare__\(\) must '
+                                     r'return a mapping, not NoneType$'):
+             class Foo(metaclass=BadMeta):
+                 pass
+         # Also test the case in which the metaclass is not a type.
+-        class BadMeta:
+-            @classmethod
+-            def __prepare__(*args):
+-                return None
++        with torch._dynamo.error_on_graph_break(False):
++            class BadMeta:
++                @classmethod
++                def __prepare__(*args):
++                    return None
+         with self.assertRaisesRegex(TypeError,
+                                     r'^<metaclass>\.__prepare__\(\) must '
+                                     r'return a mapping, not NoneType$'):
+@@ -1580,13 +1634,14 @@ class ClassCreationTests(unittest.TestCase):
+                 pass
+ 
+     def test_resolve_bases(self):
+-        class A: pass
+-        class B: pass
+-        class C:
+-            def __mro_entries__(self, bases):
+-                if A in bases:
+-                    return ()
+-                return (A,)
++        with torch._dynamo.error_on_graph_break(False):
++            class A: pass
++            class B: pass
++            class C:
++                def __mro_entries__(self, bases):
++                    if A in bases:
++                        return ()
++                    return (A,)
+         c = C()
+         self.assertEqual(types.resolve_bases(()), ())
+         self.assertEqual(types.resolve_bases((c,)), (A,))
+@@ -1609,23 +1664,24 @@ class ClassCreationTests(unittest.TestCase):
+     def test_metaclass_derivation(self):
+         # issue1294232: correct metaclass calculation
+         new_calls = []  # to check the order of __new__ calls
+-        class AMeta(type):
+-            def __new__(mcls, name, bases, ns):
+-                new_calls.append('AMeta')
+-                return super().__new__(mcls, name, bases, ns)
+-            @classmethod
+-            def __prepare__(mcls, name, bases):
+-                return {}
+-
+-        class BMeta(AMeta):
+-            def __new__(mcls, name, bases, ns):
+-                new_calls.append('BMeta')
+-                return super().__new__(mcls, name, bases, ns)
+-            @classmethod
+-            def __prepare__(mcls, name, bases):
+-                ns = super().__prepare__(name, bases)
+-                ns['BMeta_was_here'] = True
+-                return ns
++        with torch._dynamo.error_on_graph_break(False):
++            class AMeta(type):
++                def __new__(mcls, name, bases, ns):
++                    new_calls.append('AMeta')
++                    return super().__new__(mcls, name, bases, ns)
++                @classmethod
++                def __prepare__(mcls, name, bases):
++                    return {}
++
++            class BMeta(AMeta):
++                def __new__(mcls, name, bases, ns):
++                    new_calls.append('BMeta')
++                    return super().__new__(mcls, name, bases, ns)
++                @classmethod
++                def __prepare__(mcls, name, bases):
++                    ns = super().__prepare__(name, bases)
++                    ns['BMeta_was_here'] = True
++                    return ns
+ 
+         A = types.new_class("A", (), {"metaclass": AMeta})
+         self.assertEqual(new_calls, ['AMeta'])
+@@ -1663,8 +1719,9 @@ class ClassCreationTests(unittest.TestCase):
+     def test_metaclass_override_function(self):
+         # Special case: the given metaclass isn't a class,
+         # so there is no metaclass calculation.
+-        class A(metaclass=self.Meta):
+-            pass
++        with torch._dynamo.error_on_graph_break(False):
++            class A(metaclass=self.Meta):
++                pass
+ 
+         marker = object()
+         def func(*args, **kwargs):
+@@ -1682,23 +1739,24 @@ class ClassCreationTests(unittest.TestCase):
+         # but not a descendant of type.
+         new_calls = []  # to check the order of __new__ calls
+         prepare_calls = []  # to track __prepare__ calls
+-        class ANotMeta:
+-            def __new__(mcls, *args, **kwargs):
+-                new_calls.append('ANotMeta')
+-                return super().__new__(mcls)
+-            @classmethod
+-            def __prepare__(mcls, name, bases):
+-                prepare_calls.append('ANotMeta')
+-                return {}
+-
+-        class BNotMeta(ANotMeta):
+-            def __new__(mcls, *args, **kwargs):
+-                new_calls.append('BNotMeta')
+-                return super().__new__(mcls)
+-            @classmethod
+-            def __prepare__(mcls, name, bases):
+-                prepare_calls.append('BNotMeta')
+-                return super().__prepare__(name, bases)
++        with torch._dynamo.error_on_graph_break(False):
++            class ANotMeta:
++                def __new__(mcls, *args, **kwargs):
++                    new_calls.append('ANotMeta')
++                    return super().__new__(mcls)
++                @classmethod
++                def __prepare__(mcls, name, bases):
++                    prepare_calls.append('ANotMeta')
++                    return {}
++
++            class BNotMeta(ANotMeta):
++                def __new__(mcls, *args, **kwargs):
++                    new_calls.append('BNotMeta')
++                    return super().__new__(mcls)
++                @classmethod
++                def __prepare__(mcls, name, bases):
++                    prepare_calls.append('BNotMeta')
++                    return super().__prepare__(name, bases)
+ 
+         A = types.new_class("A", (), {"metaclass": ANotMeta})
+         self.assertIs(ANotMeta, type(A))
+@@ -1767,14 +1825,16 @@ class ClassCreationTests(unittest.TestCase):
+         # Only type itself can use the one-argument form (#27157)
+         self.assertIs(type(5), int)
+ 
+-        class M(type):
+-            pass
++        with torch._dynamo.error_on_graph_break(False):
++            class M(type):
++                pass
+         with self.assertRaises(TypeError) as cm:
+             M(5)
+         self.assertEqual(str(cm.exception), expected_message)
+ 
+-        class N(type, metaclass=M):
+-            pass
++        with torch._dynamo.error_on_graph_break(False):
++            class N(type, metaclass=M):
++                pass
+         with self.assertRaises(TypeError) as cm:
+             N(5)
+         self.assertEqual(str(cm.exception), expected_message)
+@@ -1783,16 +1843,18 @@ class ClassCreationTests(unittest.TestCase):
+         # bpo-44232: The C function type_new() must properly report the
+         # exception when a metaclass constructor raises an exception and the
+         # winner class is not the metaclass.
+-        class ModelBase(type):
+-            def __new__(cls, name, bases, attrs):
+-                super_new = super().__new__
+-                new_class = super_new(cls, name, bases, {})
+-                if name != "Model":
+-                    raise RuntimeWarning(f"{name=}")
+-                return new_class
+-
+-        class Model(metaclass=ModelBase):
+-            pass
++        with torch._dynamo.error_on_graph_break(False):
++            class ModelBase(type):
++                def __new__(cls, name, bases, attrs):
++                    super_new = super().__new__
++                    new_class = super_new(cls, name, bases, {})
++                    if name != "Model":
++                        raise RuntimeWarning(f"{name=}")
++                    return new_class
++
++        with torch._dynamo.error_on_graph_break(False):
++            class Model(metaclass=ModelBase):
++                pass
+ 
+         with self.assertRaises(RuntimeWarning):
+             type("SouthPonies", (Model,), {})
+@@ -1800,14 +1862,15 @@ class ClassCreationTests(unittest.TestCase):
+     def test_tuple_subclass_as_bases(self):
+         # gh-132176: it used to crash on using
+         # tuple subclass for as base classes.
+-        class TupleSubclass(tuple): pass
++        with torch._dynamo.error_on_graph_break(False):
++            class TupleSubclass(tuple): pass
+ 
+         typ = type("typ", TupleSubclass((int, object)), {})
+         self.assertEqual(typ.__bases__, (int, object))
+         self.assertEqual(type(typ.__bases__), TupleSubclass)
+ 
+ 
+-class SimpleNamespaceTests(unittest.TestCase):
++class SimpleNamespaceTests(CPythonTestCase):
+ 
+     def test_constructor(self):
+         def check(ns, expected):
+@@ -1989,8 +2052,9 @@ class SimpleNamespaceTests(unittest.TestCase):
+             ns['spam']
+ 
+     def test_subclass(self):
+-        class Spam(types.SimpleNamespace):
+-            pass
++        with torch._dynamo.error_on_graph_break(False):
++            class Spam(types.SimpleNamespace):
++                pass
+ 
+         spam = Spam(ham=8, eggs=9)
+ 
+@@ -2028,8 +2092,9 @@ class SimpleNamespaceTests(unittest.TestCase):
+         self.assertEqual(vars(copy.replace(ns, x=1, y=2)), {'x': 1, 'y': 2})
+ 
+     def test_replace_subclass(self):
+-        class Spam(types.SimpleNamespace):
+-            pass
++        with torch._dynamo.error_on_graph_break(False):
++            class Spam(types.SimpleNamespace):
++                pass
+ 
+         spam = Spam(ham=8, eggs=9)
+         spam2 = copy.replace(spam, ham=5)
+@@ -2040,8 +2105,9 @@ class SimpleNamespaceTests(unittest.TestCase):
+     def test_fake_namespace_compare(self):
+         # Issue #24257: Incorrect use of PyObject_IsInstance() caused
+         # SystemError.
+-        class FakeSimpleNamespace(str):
+-            __class__ = types.SimpleNamespace
++        with torch._dynamo.error_on_graph_break(False):
++            class FakeSimpleNamespace(str):
++                __class__ = types.SimpleNamespace
+         self.assertFalse(types.SimpleNamespace() == FakeSimpleNamespace())
+         self.assertTrue(types.SimpleNamespace() != FakeSimpleNamespace())
+         with self.assertRaises(TypeError):
+@@ -2054,7 +2120,7 @@ class SimpleNamespaceTests(unittest.TestCase):
+             types.SimpleNamespace() >= FakeSimpleNamespace()
+ 
+ 
+-class CoroutineTests(unittest.TestCase):
++class CoroutineTests(CPythonTestCase):
+     def test_wrong_args(self):
+         samples = [None, 1, object()]
+         for sample in samples:
+@@ -2068,9 +2134,10 @@ class CoroutineTests(unittest.TestCase):
+             return 'spam'
+         self.assertEqual(foo(), 'spam')
+ 
+-        class Awaitable:
+-            def __await__(self):
+-                return ()
++        with torch._dynamo.error_on_graph_break(False):
++            class Awaitable:
++                def __await__(self):
++                    return ()
+         aw = Awaitable()
+         @types.coroutine
+         def foo():
+@@ -2103,11 +2170,12 @@ class CoroutineTests(unittest.TestCase):
+             coro.close()
+ 
+     def test_duck_coro(self):
+-        class CoroLike:
+-            def send(self): pass
+-            def throw(self): pass
+-            def close(self): pass
+-            def __await__(self): return self
++        with torch._dynamo.error_on_graph_break(False):
++            class CoroLike:
++                def send(self): pass
++                def throw(self): pass
++                def close(self): pass
++                def __await__(self): return self
+ 
+         coro = CoroLike()
+         @types.coroutine
+@@ -2117,13 +2185,14 @@ class CoroutineTests(unittest.TestCase):
+         self.assertIs(foo().__await__(), coro)
+ 
+     def test_duck_corogen(self):
+-        class CoroGenLike:
+-            def send(self): pass
+-            def throw(self): pass
+-            def close(self): pass
+-            def __await__(self): return self
+-            def __iter__(self): return self
+-            def __next__(self): pass
++        with torch._dynamo.error_on_graph_break(False):
++            class CoroGenLike:
++                def send(self): pass
++                def throw(self): pass
++                def close(self): pass
++                def __await__(self): return self
++                def __iter__(self): return self
++                def __next__(self): pass
+ 
+         coro = CoroGenLike()
+         @types.coroutine
+@@ -2133,12 +2202,13 @@ class CoroutineTests(unittest.TestCase):
+         self.assertIs(foo().__await__(), coro)
+ 
+     def test_duck_gen(self):
+-        class GenLike:
+-            def send(self): pass
+-            def throw(self): pass
+-            def close(self): pass
+-            def __iter__(self): pass
+-            def __next__(self): pass
++        with torch._dynamo.error_on_graph_break(False):
++            class GenLike:
++                def send(self): pass
++                def throw(self): pass
++                def close(self): pass
++                def __iter__(self): pass
++                def __next__(self): pass
+ 
+         # Setup generator mock object
+         gen = unittest.mock.MagicMock(GenLike)
+@@ -2238,37 +2308,38 @@ class CoroutineTests(unittest.TestCase):
+         self.assertIs(ref(), wrapper)
+ 
+     def test_duck_functional_gen(self):
+-        class Generator:
+-            """Emulates the following generator (very clumsy):
+-
+-              def gen(fut):
+-                  result = yield fut
+-                  return result * 2
+-            """
+-            def __init__(self, fut):
+-                self._i = 0
+-                self._fut = fut
+-            def __iter__(self):
+-                return self
+-            def __next__(self):
+-                return self.send(None)
+-            def send(self, v):
+-                try:
+-                    if self._i == 0:
+-                        assert v is None
+-                        return self._fut
+-                    if self._i == 1:
+-                        raise StopIteration(v * 2)
+-                    if self._i > 1:
+-                        raise StopIteration
+-                finally:
+-                    self._i += 1
+-            def throw(self, tp, *exc):
+-                self._i = 100
+-                if tp is not GeneratorExit:
+-                    raise tp
+-            def close(self):
+-                self.throw(GeneratorExit)
++        with torch._dynamo.error_on_graph_break(False):
++            class Generator:
++                """Emulates the following generator (very clumsy):
++
++                  def gen(fut):
++                      result = yield fut
++                      return result * 2
++                """
++                def __init__(self, fut):
++                    self._i = 0
++                    self._fut = fut
++                def __iter__(self):
++                    return self
++                def __next__(self):
++                    return self.send(None)
++                def send(self, v):
++                    try:
++                        if self._i == 0:
++                            assert v is None
++                            return self._fut
++                        if self._i == 1:
++                            raise StopIteration(v * 2)
++                        if self._i > 1:
++                            raise StopIteration
++                    finally:
++                        self._i += 1
++                def throw(self, tp, *exc):
++                    self._i = 100
++                    if tp is not GeneratorExit:
++                        raise tp
++                def close(self):
++                    self.throw(GeneratorExit)
+ 
+         @types.coroutine
+         def foo(): return Generator('spam')
+@@ -2367,7 +2438,7 @@ class CoroutineTests(unittest.TestCase):
+             'close', 'throw'}))
+ 
+ 
+-class FunctionTests(unittest.TestCase):
++class FunctionTests(CPythonTestCase):
+     def test_function_type_defaults(self):
+         def ex(a, /, b, *, c):
+             return a + b + c
+@@ -2400,7 +2471,7 @@ class FunctionTests(unittest.TestCase):
+             )
+ 
+ 
+-class SubinterpreterTests(unittest.TestCase):
++class SubinterpreterTests(CPythonTestCase):
+ 
+     @classmethod
+     def setUpClass(cls):
+@@ -2410,6 +2481,7 @@ class SubinterpreterTests(unittest.TestCase):
+         except ModuleNotFoundError:
+             raise unittest.SkipTest('subinterpreters required')
+         import test.support.interpreters.channels
++        super().setUpClass()
+ 
+     @cpython_only
+     @no_rerun('channels (and queues) might have a refleak; see gh-122199')
+@@ -2446,4 +2518,4 @@ class SubinterpreterTests(unittest.TestCase):
+ 
+ 
+ if __name__ == '__main__':
+-    unittest.main()
++    run_tests()
diff --git a/test/dynamo/cpython/3_13/test_types.py b/test/dynamo/cpython/3_13/test_types.py
new file mode 100644
index 0000000000000..b6bc0ee2cf172
--- /dev/null
+++ b/test/dynamo/cpython/3_13/test_types.py
@@ -0,0 +1,2521 @@
+# ======= BEGIN Dynamo patch =======
+# Owner(s): ["module: dynamo"]
+
+# ruff: noqa
+# flake8: noqa
+
+# Test copied from
+# https://raw.githubusercontent.com/python/cpython/refs/tags/v3.13.5/Lib/test/test_types.py
+
+import sys
+import torch
+import torch._dynamo.test_case
+import unittest
+from torch._dynamo.test_case import CPythonTestCase
+from torch.testing._internal.common_utils import run_tests
+
+# Dummy decorator for CPython test compatibility
+def no_rerun(reason):
+    def decorator(func):
+        return func
+    return decorator
+
+# ======= END DYNAMO PATCH =======
+
+# Python test set -- part 6, built-in types
+
+from test.support import (
+    run_with_locale, is_apple_mobile, cpython_only,
+    iter_builtin_types, iter_slot_wrappers,
+    MISSING_C_DOCSTRINGS,
+)
+from test.support.script_helper import assert_python_ok
+from test.support.import_helper import import_fresh_module
+
+import collections.abc
+from collections import namedtuple, UserDict
+import copy
+import _datetime
+import gc
+import inspect
+import pickle
+import locale
+import sys
+import textwrap
+import types
+import unittest.mock
+import weakref
+import typing
+
+
+T = typing.TypeVar("T")
+
+class Example:
+    pass
+
+class Forward: ...
+
+def clear_typing_caches():
+    for f in typing._cleanups:
+        f()
+
+
+class TypesTests(CPythonTestCase):
+
+    def test_truth_values(self):
+        if None: self.fail('None is true instead of false')
+        if 0: self.fail('0 is true instead of false')
+        if 0.0: self.fail('0.0 is true instead of false')
+        if '': self.fail('\'\' is true instead of false')
+        if not 1: self.fail('1 is false instead of true')
+        if not 1.0: self.fail('1.0 is false instead of true')
+        if not 'x': self.fail('\'x\' is false instead of true')
+        if not {'x': 1}: self.fail('{\'x\': 1} is false instead of true')
+        def f(): pass
+        with torch._dynamo.error_on_graph_break(False):
+            class C: pass
+            x = C()
+        if not f: self.fail('f is false instead of true')
+        if not C: self.fail('C is false instead of true')
+        if not sys: self.fail('sys is false instead of true')
+        if not x: self.fail('x is false instead of true')
+
+    def test_boolean_ops(self):
+        if 0 or 0: self.fail('0 or 0 is true instead of false')
+        if 1 and 1: pass
+        else: self.fail('1 and 1 is false instead of true')
+        if not 1: self.fail('not 1 is true instead of false')
+
+    def test_comparisons(self):
+        if 0 < 1 <= 1 == 1 >= 1 > 0 != 1: pass
+        else: self.fail('int comparisons failed')
+        if 0.0 < 1.0 <= 1.0 == 1.0 >= 1.0 > 0.0 != 1.0: pass
+        else: self.fail('float comparisons failed')
+        if '' < 'a' <= 'a' == 'a' < 'abc' < 'abd' < 'b': pass
+        else: self.fail('string comparisons failed')
+        if None is None: pass
+        else: self.fail('identity test failed')
+
+    def test_float_constructor(self):
+        self.assertRaises(ValueError, float, '')
+        self.assertRaises(ValueError, float, '5\0')
+        self.assertRaises(ValueError, float, '5_5\0')
+
+    def test_zero_division(self):
+        try: 5.0 / 0.0
+        except ZeroDivisionError: pass
+        else: self.fail("5.0 / 0.0 didn't raise ZeroDivisionError")
+
+        try: 5.0 // 0.0
+        except ZeroDivisionError: pass
+        else: self.fail("5.0 // 0.0 didn't raise ZeroDivisionError")
+
+        try: 5.0 % 0.0
+        except ZeroDivisionError: pass
+        else: self.fail("5.0 % 0.0 didn't raise ZeroDivisionError")
+
+        try: 5 / 0
+        except ZeroDivisionError: pass
+        else: self.fail("5 / 0 didn't raise ZeroDivisionError")
+
+        try: 5 // 0
+        except ZeroDivisionError: pass
+        else: self.fail("5 // 0 didn't raise ZeroDivisionError")
+
+        try: 5 % 0
+        except ZeroDivisionError: pass
+        else: self.fail("5 % 0 didn't raise ZeroDivisionError")
+
+    def test_numeric_types(self):
+        if 0 != 0.0 or 1 != 1.0 or -1 != -1.0:
+            self.fail('int/float value not equal')
+        # calling built-in types without argument must return 0
+        if int() != 0: self.fail('int() does not return 0')
+        if float() != 0.0: self.fail('float() does not return 0.0')
+        if int(1.9) == 1 == int(1.1) and int(-1.1) == -1 == int(-1.9): pass
+        else: self.fail('int() does not round properly')
+        if float(1) == 1.0 and float(-1) == -1.0 and float(0) == 0.0: pass
+        else: self.fail('float() does not work properly')
+
+    def test_float_to_string(self):
+        def test(f, result):
+            self.assertEqual(f.__format__('e'), result)
+            self.assertEqual('%e' % f, result)
+
+        # test all 2 digit exponents, both with __format__ and with
+        #  '%' formatting
+        for i in range(-99, 100):
+            test(float('1.5e'+str(i)), '1.500000e{0:+03d}'.format(i))
+
+        # test some 3 digit exponents
+        self.assertEqual(1.5e100.__format__('e'), '1.500000e+100')
+        self.assertEqual('%e' % 1.5e100, '1.500000e+100')
+
+        self.assertEqual(1.5e101.__format__('e'), '1.500000e+101')
+        self.assertEqual('%e' % 1.5e101, '1.500000e+101')
+
+        self.assertEqual(1.5e-100.__format__('e'), '1.500000e-100')
+        self.assertEqual('%e' % 1.5e-100, '1.500000e-100')
+
+        self.assertEqual(1.5e-101.__format__('e'), '1.500000e-101')
+        self.assertEqual('%e' % 1.5e-101, '1.500000e-101')
+
+        self.assertEqual('%g' % 1.0, '1')
+        self.assertEqual('%#g' % 1.0, '1.00000')
+
+    def test_normal_integers(self):
+        # Ensure the first 256 integers are shared
+        a = 256
+        b = 128*2
+        if a is not b: self.fail('256 is not shared')
+        if 12 + 24 != 36: self.fail('int op')
+        if 12 + (-24) != -12: self.fail('int op')
+        if (-12) + 24 != 12: self.fail('int op')
+        if (-12) + (-24) != -36: self.fail('int op')
+        if not 12 < 24: self.fail('int op')
+        if not -24 < -12: self.fail('int op')
+        # Test for a particular bug in integer multiply
+        xsize, ysize, zsize = 238, 356, 4
+        if not (xsize*ysize*zsize == zsize*xsize*ysize == 338912):
+            self.fail('int mul commutativity')
+        # And another.
+        m = -sys.maxsize - 1
+        for divisor in 1, 2, 4, 8, 16, 32:
+            j = m // divisor
+            prod = divisor * j
+            if prod != m:
+                self.fail("%r * %r == %r != %r" % (divisor, j, prod, m))
+            if type(prod) is not int:
+                self.fail("expected type(prod) to be int, not %r" %
+                                   type(prod))
+        # Check for unified integral type
+        for divisor in 1, 2, 4, 8, 16, 32:
+            j = m // divisor - 1
+            prod = divisor * j
+            if type(prod) is not int:
+                self.fail("expected type(%r) to be int, not %r" %
+                                   (prod, type(prod)))
+        # Check for unified integral type
+        m = sys.maxsize
+        for divisor in 1, 2, 4, 8, 16, 32:
+            j = m // divisor + 1
+            prod = divisor * j
+            if type(prod) is not int:
+                self.fail("expected type(%r) to be int, not %r" %
+                                   (prod, type(prod)))
+
+        x = sys.maxsize
+        self.assertIsInstance(x + 1, int,
+                              "(sys.maxsize + 1) should have returned int")
+        self.assertIsInstance(-x - 1, int,
+                              "(-sys.maxsize - 1) should have returned int")
+        self.assertIsInstance(-x - 2, int,
+                              "(-sys.maxsize - 2) should have returned int")
+
+        try: 5 << -5
+        except ValueError: pass
+        else: self.fail('int negative shift <<')
+
+        try: 5 >> -5
+        except ValueError: pass
+        else: self.fail('int negative shift >>')
+
+    def test_floats(self):
+        if 12.0 + 24.0 != 36.0: self.fail('float op')
+        if 12.0 + (-24.0) != -12.0: self.fail('float op')
+        if (-12.0) + 24.0 != 12.0: self.fail('float op')
+        if (-12.0) + (-24.0) != -36.0: self.fail('float op')
+        if not 12.0 < 24.0: self.fail('float op')
+        if not -24.0 < -12.0: self.fail('float op')
+
+    def test_strings(self):
+        if len('') != 0: self.fail('len(\'\')')
+        if len('a') != 1: self.fail('len(\'a\')')
+        if len('abcdef') != 6: self.fail('len(\'abcdef\')')
+        if 'xyz' + 'abcde' != 'xyzabcde': self.fail('string concatenation')
+        if 'xyz'*3 != 'xyzxyzxyz': self.fail('string repetition *3')
+        if 0*'abcde' != '': self.fail('string repetition 0*')
+        if min('abc') != 'a' or max('abc') != 'c': self.fail('min/max string')
+        if 'a' in 'abc' and 'b' in 'abc' and 'c' in 'abc' and 'd' not in 'abc': pass
+        else: self.fail('in/not in string')
+        x = 'x'*103
+        if '%s!'%x != x+'!': self.fail('nasty string formatting bug')
+
+        #extended slices for strings
+        a = '0123456789'
+        self.assertEqual(a[::], a)
+        self.assertEqual(a[::2], '02468')
+        self.assertEqual(a[1::2], '13579')
+        self.assertEqual(a[::-1],'9876543210')
+        self.assertEqual(a[::-2], '97531')
+        self.assertEqual(a[3::-2], '31')
+        self.assertEqual(a[-100:100:], a)
+        self.assertEqual(a[100:-100:-1], a[::-1])
+        self.assertEqual(a[-100:100:2], '02468')
+
+    def test_type_function(self):
+        self.assertRaises(TypeError, type, 1, 2)
+        self.assertRaises(TypeError, type, 1, 2, 3, 4)
+
+    def test_int__format__(self):
+        def test(i, format_spec, result):
+            # just make sure we have the unified type for integers
+            self.assertIs(type(i), int)
+            self.assertIs(type(format_spec), str)
+            self.assertEqual(i.__format__(format_spec), result)
+
+        test(123456789, 'd', '123456789')
+        test(123456789, 'd', '123456789')
+
+        test(1, 'c', '\01')
+
+        # sign and aligning are interdependent
+        test(1, "-", '1')
+        test(-1, "-", '-1')
+        test(1, "-3", '  1')
+        test(-1, "-3", ' -1')
+        test(1, "+3", ' +1')
+        test(-1, "+3", ' -1')
+        test(1, " 3", '  1')
+        test(-1, " 3", ' -1')
+        test(1, " ", ' 1')
+        test(-1, " ", '-1')
+
+        # hex
+        test(3, "x", "3")
+        test(3, "X", "3")
+        test(1234, "x", "4d2")
+        test(-1234, "x", "-4d2")
+        test(1234, "8x", "     4d2")
+        test(-1234, "8x", "    -4d2")
+        test(1234, "x", "4d2")
+        test(-1234, "x", "-4d2")
+        test(-3, "x", "-3")
+        test(-3, "X", "-3")
+        test(int('be', 16), "x", "be")
+        test(int('be', 16), "X", "BE")
+        test(-int('be', 16), "x", "-be")
+        test(-int('be', 16), "X", "-BE")
+
+        # octal
+        test(3, "o", "3")
+        test(-3, "o", "-3")
+        test(65, "o", "101")
+        test(-65, "o", "-101")
+        test(1234, "o", "2322")
+        test(-1234, "o", "-2322")
+        test(1234, "-o", "2322")
+        test(-1234, "-o", "-2322")
+        test(1234, " o", " 2322")
+        test(-1234, " o", "-2322")
+        test(1234, "+o", "+2322")
+        test(-1234, "+o", "-2322")
+
+        # binary
+        test(3, "b", "11")
+        test(-3, "b", "-11")
+        test(1234, "b", "10011010010")
+        test(-1234, "b", "-10011010010")
+        test(1234, "-b", "10011010010")
+        test(-1234, "-b", "-10011010010")
+        test(1234, " b", " 10011010010")
+        test(-1234, " b", "-10011010010")
+        test(1234, "+b", "+10011010010")
+        test(-1234, "+b", "-10011010010")
+
+        # alternate (#) formatting
+        test(0, "#b", '0b0')
+        test(0, "-#b", '0b0')
+        test(1, "-#b", '0b1')
+        test(-1, "-#b", '-0b1')
+        test(-1, "-#5b", ' -0b1')
+        test(1, "+#5b", ' +0b1')
+        test(100, "+#b", '+0b1100100')
+        test(100, "#012b", '0b0001100100')
+        test(-100, "#012b", '-0b001100100')
+
+        test(0, "#o", '0o0')
+        test(0, "-#o", '0o0')
+        test(1, "-#o", '0o1')
+        test(-1, "-#o", '-0o1')
+        test(-1, "-#5o", ' -0o1')
+        test(1, "+#5o", ' +0o1')
+        test(100, "+#o", '+0o144')
+        test(100, "#012o", '0o0000000144')
+        test(-100, "#012o", '-0o000000144')
+
+        test(0, "#x", '0x0')
+        test(0, "-#x", '0x0')
+        test(1, "-#x", '0x1')
+        test(-1, "-#x", '-0x1')
+        test(-1, "-#5x", ' -0x1')
+        test(1, "+#5x", ' +0x1')
+        test(100, "+#x", '+0x64')
+        test(100, "#012x", '0x0000000064')
+        test(-100, "#012x", '-0x000000064')
+        test(123456, "#012x", '0x000001e240')
+        test(-123456, "#012x", '-0x00001e240')
+
+        test(0, "#X", '0X0')
+        test(0, "-#X", '0X0')
+        test(1, "-#X", '0X1')
+        test(-1, "-#X", '-0X1')
+        test(-1, "-#5X", ' -0X1')
+        test(1, "+#5X", ' +0X1')
+        test(100, "+#X", '+0X64')
+        test(100, "#012X", '0X0000000064')
+        test(-100, "#012X", '-0X000000064')
+        test(123456, "#012X", '0X000001E240')
+        test(-123456, "#012X", '-0X00001E240')
+
+        test(123, ',', '123')
+        test(-123, ',', '-123')
+        test(1234, ',', '1,234')
+        test(-1234, ',', '-1,234')
+        test(123456, ',', '123,456')
+        test(-123456, ',', '-123,456')
+        test(1234567, ',', '1,234,567')
+        test(-1234567, ',', '-1,234,567')
+
+        # issue 5782, commas with no specifier type
+        test(1234, '010,', '00,001,234')
+
+        # Unified type for integers
+        test(10**100, 'd', '1' + '0' * 100)
+        test(10**100+100, 'd', '1' + '0' * 97 + '100')
+
+        # make sure these are errors
+
+        # precision disallowed
+        self.assertRaises(ValueError, 3 .__format__, "1.3")
+        # sign not allowed with 'c'
+        self.assertRaises(ValueError, 3 .__format__, "+c")
+        # format spec must be string
+        self.assertRaises(TypeError, 3 .__format__, None)
+        self.assertRaises(TypeError, 3 .__format__, 0)
+        # can't have ',' with 'n'
+        self.assertRaises(ValueError, 3 .__format__, ",n")
+        # can't have ',' with 'c'
+        self.assertRaises(ValueError, 3 .__format__, ",c")
+        # can't have '#' with 'c'
+        self.assertRaises(ValueError, 3 .__format__, "#c")
+
+        # ensure that only int and float type specifiers work
+        for format_spec in ([chr(x) for x in range(ord('a'), ord('z')+1)] +
+                            [chr(x) for x in range(ord('A'), ord('Z')+1)]):
+            if not format_spec in 'bcdoxXeEfFgGn%':
+                self.assertRaises(ValueError, 0 .__format__, format_spec)
+                self.assertRaises(ValueError, 1 .__format__, format_spec)
+                self.assertRaises(ValueError, (-1) .__format__, format_spec)
+
+        # ensure that float type specifiers work; format converts
+        #  the int to a float
+        for format_spec in 'eEfFgG%':
+            for value in [0, 1, -1, 100, -100, 1234567890, -1234567890]:
+                self.assertEqual(value.__format__(format_spec),
+                                 float(value).__format__(format_spec))
+
+        # Issue 6902
+        test(123456, "0<20", '12345600000000000000')
+        test(123456, "1<20", '12345611111111111111')
+        test(123456, "*<20", '123456**************')
+        test(123456, "0>20", '00000000000000123456')
+        test(123456, "1>20", '11111111111111123456')
+        test(123456, "*>20", '**************123456')
+        test(123456, "0=20", '00000000000000123456')
+        test(123456, "1=20", '11111111111111123456')
+        test(123456, "*=20", '**************123456')
+
+    @run_with_locale('LC_NUMERIC', 'en_US.UTF8', '')
+    def test_float__format__locale(self):
+        # test locale support for __format__ code 'n'
+
+        for i in range(-10, 10):
+            x = 1234567890.0 * (10.0 ** i)
+            self.assertEqual(locale.format_string('%g', x, grouping=True), format(x, 'n'))
+            self.assertEqual(locale.format_string('%.10g', x, grouping=True), format(x, '.10n'))
+
+    @run_with_locale('LC_NUMERIC', 'en_US.UTF8', '')
+    def test_int__format__locale(self):
+        # test locale support for __format__ code 'n' for integers
+
+        x = 123456789012345678901234567890
+        for i in range(0, 30):
+            self.assertEqual(locale.format_string('%d', x, grouping=True), format(x, 'n'))
+
+            # move to the next integer to test
+            x = x // 10
+
+        rfmt = ">20n"
+        lfmt = "<20n"
+        cfmt = "^20n"
+        for x in (1234, 12345, 123456, 1234567, 12345678, 123456789, 1234567890, 12345678900):
+            self.assertEqual(len(format(0, rfmt)), len(format(x, rfmt)))
+            self.assertEqual(len(format(0, lfmt)), len(format(x, lfmt)))
+            self.assertEqual(len(format(0, cfmt)), len(format(x, cfmt)))
+
+    def test_float__format__(self):
+        def test(f, format_spec, result):
+            self.assertEqual(f.__format__(format_spec), result)
+            self.assertEqual(format(f, format_spec), result)
+
+        test(0.0, 'f', '0.000000')
+
+        # the default is 'g', except for empty format spec
+        test(0.0, '', '0.0')
+        test(0.01, '', '0.01')
+        test(0.01, 'g', '0.01')
+
+        # test for issue 3411
+        test(1.23, '1', '1.23')
+        test(-1.23, '1', '-1.23')
+        test(1.23, '1g', '1.23')
+        test(-1.23, '1g', '-1.23')
+
+        test( 1.0, ' g', ' 1')
+        test(-1.0, ' g', '-1')
+        test( 1.0, '+g', '+1')
+        test(-1.0, '+g', '-1')
+        test(1.1234e200, 'g', '1.1234e+200')
+        test(1.1234e200, 'G', '1.1234E+200')
+
+
+        test(1.0, 'f', '1.000000')
+
+        test(-1.0, 'f', '-1.000000')
+
+        test( 1.0, ' f', ' 1.000000')
+        test(-1.0, ' f', '-1.000000')
+        test( 1.0, '+f', '+1.000000')
+        test(-1.0, '+f', '-1.000000')
+
+        # Python versions <= 3.0 switched from 'f' to 'g' formatting for
+        # values larger than 1e50.  No longer.
+        f = 1.1234e90
+        for fmt in 'f', 'F':
+            # don't do a direct equality check, since on some
+            # platforms only the first few digits of dtoa
+            # will be reliable
+            result = f.__format__(fmt)
+            self.assertEqual(len(result), 98)
+            self.assertEqual(result[-7], '.')
+            self.assertIn(result[:12], ('112340000000', '112339999999'))
+        f = 1.1234e200
+        for fmt in 'f', 'F':
+            result = f.__format__(fmt)
+            self.assertEqual(len(result), 208)
+            self.assertEqual(result[-7], '.')
+            self.assertIn(result[:12], ('112340000000', '112339999999'))
+
+
+        test( 1.0, 'e', '1.000000e+00')
+        test(-1.0, 'e', '-1.000000e+00')
+        test( 1.0, 'E', '1.000000E+00')
+        test(-1.0, 'E', '-1.000000E+00')
+        test(1.1234e20, 'e', '1.123400e+20')
+        test(1.1234e20, 'E', '1.123400E+20')
+
+        # No format code means use g, but must have a decimal
+        # and a number after the decimal.  This is tricky, because
+        # a totally empty format specifier means something else.
+        # So, just use a sign flag
+        test(1e200, '+g', '+1e+200')
+        test(1e200, '+', '+1e+200')
+
+        test(1.1e200, '+g', '+1.1e+200')
+        test(1.1e200, '+', '+1.1e+200')
+
+        # 0 padding
+        test(1234., '010f', '1234.000000')
+        test(1234., '011f', '1234.000000')
+        test(1234., '012f', '01234.000000')
+        test(-1234., '011f', '-1234.000000')
+        test(-1234., '012f', '-1234.000000')
+        test(-1234., '013f', '-01234.000000')
+        test(-1234.12341234, '013f', '-01234.123412')
+        test(-123456.12341234, '011.2f', '-0123456.12')
+
+        # issue 5782, commas with no specifier type
+        test(1.2, '010,.2', '0,000,001.2')
+
+        # 0 padding with commas
+        test(1234., '011,f', '1,234.000000')
+        test(1234., '012,f', '1,234.000000')
+        test(1234., '013,f', '01,234.000000')
+        test(-1234., '012,f', '-1,234.000000')
+        test(-1234., '013,f', '-1,234.000000')
+        test(-1234., '014,f', '-01,234.000000')
+        test(-12345., '015,f', '-012,345.000000')
+        test(-123456., '016,f', '-0,123,456.000000')
+        test(-123456., '017,f', '-0,123,456.000000')
+        test(-123456.12341234, '017,f', '-0,123,456.123412')
+        test(-123456.12341234, '013,.2f', '-0,123,456.12')
+
+        # % formatting
+        test(-1.0, '%', '-100.000000%')
+
+        # format spec must be string
+        self.assertRaises(TypeError, 3.0.__format__, None)
+        self.assertRaises(TypeError, 3.0.__format__, 0)
+
+        # confirm format options expected to fail on floats, such as integer
+        # presentation types
+        for format_spec in 'sbcdoxX':
+            self.assertRaises(ValueError, format, 0.0, format_spec)
+            self.assertRaises(ValueError, format, 1.0, format_spec)
+            self.assertRaises(ValueError, format, -1.0, format_spec)
+            self.assertRaises(ValueError, format, 1e100, format_spec)
+            self.assertRaises(ValueError, format, -1e100, format_spec)
+            self.assertRaises(ValueError, format, 1e-100, format_spec)
+            self.assertRaises(ValueError, format, -1e-100, format_spec)
+
+        # Alternate float formatting
+        test(1.0, '.0e', '1e+00')
+        test(1.0, '#.0e', '1.e+00')
+        test(1.0, '.0f', '1')
+        test(1.0, '#.0f', '1.')
+        test(1.1, 'g', '1.1')
+        test(1.1, '#g', '1.10000')
+        test(1.0, '.0%', '100%')
+        test(1.0, '#.0%', '100.%')
+
+        # Issue 7094: Alternate formatting (specified by #)
+        test(1.0, '0e',  '1.000000e+00')
+        test(1.0, '#0e', '1.000000e+00')
+        test(1.0, '0f',  '1.000000' )
+        test(1.0, '#0f', '1.000000')
+        test(1.0, '.1e',  '1.0e+00')
+        test(1.0, '#.1e', '1.0e+00')
+        test(1.0, '.1f',  '1.0')
+        test(1.0, '#.1f', '1.0')
+        test(1.0, '.1%',  '100.0%')
+        test(1.0, '#.1%', '100.0%')
+
+        # Issue 6902
+        test(12345.6, "0<20", '12345.60000000000000')
+        test(12345.6, "1<20", '12345.61111111111111')
+        test(12345.6, "*<20", '12345.6*************')
+        test(12345.6, "0>20", '000000000000012345.6')
+        test(12345.6, "1>20", '111111111111112345.6')
+        test(12345.6, "*>20", '*************12345.6')
+        test(12345.6, "0=20", '000000000000012345.6')
+        test(12345.6, "1=20", '111111111111112345.6')
+        test(12345.6, "*=20", '*************12345.6')
+
+    def test_format_spec_errors(self):
+        # int, float, and string all share the same format spec
+        # mini-language parser.
+
+        # Check that we can't ask for too many digits. This is
+        # probably a CPython specific test. It tries to put the width
+        # into a C long.
+        self.assertRaises(ValueError, format, 0, '1'*10000 + 'd')
+
+        # Similar with the precision.
+        self.assertRaises(ValueError, format, 0, '.' + '1'*10000 + 'd')
+
+        # And may as well test both.
+        self.assertRaises(ValueError, format, 0, '1'*1000 + '.' + '1'*10000 + 'd')
+
+        # Make sure commas aren't allowed with various type codes
+        for code in 'xXobns':
+            self.assertRaises(ValueError, format, 0, ',' + code)
+
+    def test_internal_sizes(self):
+        self.assertGreater(object.__basicsize__, 0)
+        self.assertGreater(tuple.__itemsize__, 0)
+
+    def test_slot_wrapper_types(self):
+        self.assertIsInstance(object.__init__, types.WrapperDescriptorType)
+        self.assertIsInstance(object.__str__, types.WrapperDescriptorType)
+        self.assertIsInstance(object.__lt__, types.WrapperDescriptorType)
+        self.assertIsInstance(int.__lt__, types.WrapperDescriptorType)
+
+    @unittest.skipIf(MISSING_C_DOCSTRINGS,
+                     "Signature information for builtins requires docstrings")
+    def test_dunder_get_signature(self):
+        sig = inspect.signature(object.__init__.__get__)
+        self.assertEqual(list(sig.parameters), ["instance", "owner"])
+        # gh-93021: Second parameter is optional
+        self.assertIs(sig.parameters["owner"].default, None)
+
+    def test_method_wrapper_types(self):
+        self.assertIsInstance(object().__init__, types.MethodWrapperType)
+        self.assertIsInstance(object().__str__, types.MethodWrapperType)
+        self.assertIsInstance(object().__lt__, types.MethodWrapperType)
+        self.assertIsInstance((42).__lt__, types.MethodWrapperType)
+
+    def test_method_descriptor_types(self):
+        self.assertIsInstance(str.join, types.MethodDescriptorType)
+        self.assertIsInstance(list.append, types.MethodDescriptorType)
+        self.assertIsInstance(''.join, types.BuiltinMethodType)
+        self.assertIsInstance([].append, types.BuiltinMethodType)
+
+        self.assertIsInstance(int.__dict__['from_bytes'], types.ClassMethodDescriptorType)
+        self.assertIsInstance(int.from_bytes, types.BuiltinMethodType)
+        self.assertIsInstance(int.__new__, types.BuiltinMethodType)
+
+    def test_method_descriptor_crash(self):
+        # gh-132747: The default __get__() implementation in C was unable
+        # to handle a second argument of None when called from Python
+        import _io
+        import io
+        import _queue
+
+        to_check = [
+            # (method, instance)
+            (_io._TextIOBase.read, io.StringIO()),
+            (_queue.SimpleQueue.put, _queue.SimpleQueue()),
+            (str.capitalize, "nobody expects the spanish inquisition")
+        ]
+
+        for method, instance in to_check:
+            with self.subTest(method=method, instance=instance):
+                bound = method.__get__(instance)
+                self.assertIsInstance(bound, types.BuiltinMethodType)
+
+    def test_ellipsis_type(self):
+        self.assertIsInstance(Ellipsis, types.EllipsisType)
+
+    def test_notimplemented_type(self):
+        self.assertIsInstance(NotImplemented, types.NotImplementedType)
+
+    def test_none_type(self):
+        self.assertIsInstance(None, types.NoneType)
+
+    def test_traceback_and_frame_types(self):
+        try:
+            raise OSError
+        except OSError as e:
+            exc = e
+        self.assertIsInstance(exc.__traceback__, types.TracebackType)
+        self.assertIsInstance(exc.__traceback__.tb_frame, types.FrameType)
+
+    def test_capsule_type(self):
+        self.assertIsInstance(_datetime.datetime_CAPI, types.CapsuleType)
+
+    def test_call_unbound_crash(self):
+        # GH-131998: The specialized instruction would get tricked into dereferencing
+        # a bound "self" that didn't exist if subsequently called unbound.
+        code = """if True:
+
+        def call(part):
+            [] + ([] + [])
+            part.pop()
+
+        for _ in range(3):
+            call(['a'])
+        try:
+            call(list)
+        except TypeError:
+            pass
+        """
+        assert_python_ok("-c", code)
+
+
+class UnionTests(CPythonTestCase):
+
+    def test_or_types_operator(self):
+        self.assertEqual(int | str, typing.Union[int, str])
+        self.assertNotEqual(int | list, typing.Union[int, str])
+        self.assertEqual(str | int, typing.Union[int, str])
+        self.assertEqual(int | None, typing.Union[int, None])
+        self.assertEqual(None | int, typing.Union[int, None])
+        self.assertEqual(int | type(None), int | None)
+        self.assertEqual(type(None) | int, None | int)
+        self.assertEqual(int | str | list, typing.Union[int, str, list])
+        self.assertEqual(int | (str | list), typing.Union[int, str, list])
+        self.assertEqual(str | (int | list), typing.Union[int, str, list])
+        self.assertEqual(typing.List | typing.Tuple, typing.Union[typing.List, typing.Tuple])
+        self.assertEqual(typing.List[int] | typing.Tuple[int], typing.Union[typing.List[int], typing.Tuple[int]])
+        self.assertEqual(typing.List[int] | None, typing.Union[typing.List[int], None])
+        self.assertEqual(None | typing.List[int], typing.Union[None, typing.List[int]])
+        self.assertEqual(str | float | int | complex | int, (int | str) | (float | complex))
+        self.assertEqual(typing.Union[str, int, typing.List[int]], str | int | typing.List[int])
+        self.assertIs(int | int, int)
+        self.assertEqual(
+            BaseException |
+            bool |
+            bytes |
+            complex |
+            float |
+            int |
+            list |
+            map |
+            set,
+            typing.Union[
+                BaseException,
+                bool,
+                bytes,
+                complex,
+                float,
+                int,
+                list,
+                map,
+                set,
+            ])
+        with self.assertRaises(TypeError):
+            int | 3
+        with self.assertRaises(TypeError):
+            3 | int
+        with self.assertRaises(TypeError):
+            Example() | int
+        x = int | str
+        self.assertEqual(x, int | str)
+        self.assertEqual(x, str | int)
+        self.assertNotEqual(x, {})  # should not raise exception
+        with self.assertRaises(TypeError):
+            x < x
+        with self.assertRaises(TypeError):
+            x <= x
+        y = typing.Union[str, int]
+        with self.assertRaises(TypeError):
+            x < y
+        y = int | bool
+        with self.assertRaises(TypeError):
+            x < y
+        # Check that we don't crash if typing.Union does not have a tuple in __args__
+        y = typing.Union[str, int]
+        y.__args__ = [str, int]
+        self.assertEqual(x, y)
+
+    def test_hash(self):
+        self.assertEqual(hash(int | str), hash(str | int))
+        self.assertEqual(hash(int | str), hash(typing.Union[int, str]))
+
+    def test_union_of_unhashable(self):
+        with torch._dynamo.error_on_graph_break(False):
+            class UnhashableMeta(type):
+                __hash__ = None
+
+        with torch._dynamo.error_on_graph_break(False):
+            class A(metaclass=UnhashableMeta): ...
+        with torch._dynamo.error_on_graph_break(False):
+            class B(metaclass=UnhashableMeta): ...
+
+        self.assertEqual((A | B).__args__, (A, B))
+        union1 = A | B
+        with self.assertRaises(TypeError):
+            hash(union1)
+
+        union2 = int | B
+        with self.assertRaises(TypeError):
+            hash(union2)
+
+        union3 = A | int
+        with self.assertRaises(TypeError):
+            hash(union3)
+
+    def test_instancecheck_and_subclasscheck(self):
+        for x in (int | str, typing.Union[int, str]):
+            with self.subTest(x=x):
+                self.assertIsInstance(1, x)
+                self.assertIsInstance(True, x)
+                self.assertIsInstance('a', x)
+                self.assertNotIsInstance(None, x)
+                self.assertTrue(issubclass(int, x))
+                self.assertTrue(issubclass(bool, x))
+                self.assertTrue(issubclass(str, x))
+                self.assertFalse(issubclass(type(None), x))
+
+        for x in (int | None, typing.Union[int, None]):
+            with self.subTest(x=x):
+                self.assertIsInstance(None, x)
+                self.assertTrue(issubclass(type(None), x))
+
+        for x in (
+            int | collections.abc.Mapping,
+            typing.Union[int, collections.abc.Mapping],
+        ):
+            with self.subTest(x=x):
+                self.assertIsInstance({}, x)
+                self.assertNotIsInstance((), x)
+                self.assertTrue(issubclass(dict, x))
+                self.assertFalse(issubclass(list, x))
+
+    def test_instancecheck_and_subclasscheck_order(self):
+        T = typing.TypeVar('T')
+
+        will_resolve = (
+            int | T,
+            typing.Union[int, T],
+        )
+        for x in will_resolve:
+            with self.subTest(x=x):
+                self.assertIsInstance(1, x)
+                self.assertTrue(issubclass(int, x))
+
+        wont_resolve = (
+            T | int,
+            typing.Union[T, int],
+        )
+        for x in wont_resolve:
+            with self.subTest(x=x):
+                with self.assertRaises(TypeError):
+                    issubclass(int, x)
+                with self.assertRaises(TypeError):
+                    isinstance(1, x)
+
+        for x in (*will_resolve, *wont_resolve):
+            with self.subTest(x=x):
+                with self.assertRaises(TypeError):
+                    issubclass(object, x)
+                with self.assertRaises(TypeError):
+                    isinstance(object(), x)
+
+    def test_bad_instancecheck(self):
+        with torch._dynamo.error_on_graph_break(False):
+            class BadMeta(type):
+                def __instancecheck__(cls, inst):
+                    1/0
+        x = int | BadMeta('A', (), {})
+        self.assertTrue(isinstance(1, x))
+        self.assertRaises(ZeroDivisionError, isinstance, [], x)
+
+    def test_bad_subclasscheck(self):
+        with torch._dynamo.error_on_graph_break(False):
+            class BadMeta(type):
+                def __subclasscheck__(cls, sub):
+                    1/0
+        x = int | BadMeta('A', (), {})
+        self.assertTrue(issubclass(int, x))
+        self.assertRaises(ZeroDivisionError, issubclass, list, x)
+
+    def test_or_type_operator_with_TypeVar(self):
+        TV = typing.TypeVar('T')
+        self.assertEqual(TV | str, typing.Union[TV, str])
+        self.assertEqual(str | TV, typing.Union[str, TV])
+        self.assertIs((int | TV)[int], int)
+        self.assertIs((TV | int)[int], int)
+
+    def test_union_args(self):
+        def check(arg, expected):
+            clear_typing_caches()
+            self.assertEqual(arg.__args__, expected)
+
+        check(int | str, (int, str))
+        check((int | str) | list, (int, str, list))
+        check(int | (str | list), (int, str, list))
+        check((int | str) | int, (int, str))
+        check(int | (str | int), (int, str))
+        check((int | str) | (str | int), (int, str))
+        check(typing.Union[int, str] | list, (int, str, list))
+        check(int | typing.Union[str, list], (int, str, list))
+        check((int | str) | (list | int), (int, str, list))
+        check((int | str) | typing.Union[list, int], (int, str, list))
+        check(typing.Union[int, str] | (list | int), (int, str, list))
+        check((str | int) | (int | list), (str, int, list))
+        check((str | int) | typing.Union[int, list], (str, int, list))
+        check(typing.Union[str, int] | (int | list), (str, int, list))
+        check(int | type(None), (int, type(None)))
+        check(type(None) | int, (type(None), int))
+
+        args = (int, list[int], typing.List[int],
+                typing.Tuple[int, int], typing.Callable[[int], int],
+                typing.Hashable, typing.TypeVar('T'))
+        for x in args:
+            with self.subTest(x):
+                check(x | None, (x, type(None)))
+                check(None | x, (type(None), x))
+
+    def test_union_parameter_chaining(self):
+        T = typing.TypeVar("T")
+        S = typing.TypeVar("S")
+
+        self.assertEqual((float | list[T])[int], float | list[int])
+        self.assertEqual(list[int | list[T]].__parameters__, (T,))
+        self.assertEqual(list[int | list[T]][str], list[int | list[str]])
+        self.assertEqual((list[T] | list[S]).__parameters__, (T, S))
+        self.assertEqual((list[T] | list[S])[int, T], list[int] | list[T])
+        self.assertEqual((list[T] | list[S])[int, int], list[int])
+
+    def test_union_parameter_substitution(self):
+        def eq(actual, expected, typed=True):
+            self.assertEqual(actual, expected)
+            if typed:
+                self.assertIs(type(actual), type(expected))
+
+        T = typing.TypeVar('T')
+        S = typing.TypeVar('S')
+        NT = typing.NewType('NT', str)
+        x = int | T | bytes
+
+        eq(x[str], int | str | bytes, typed=False)
+        eq(x[list[int]], int | list[int] | bytes, typed=False)
+        eq(x[typing.List], int | typing.List | bytes)
+        eq(x[typing.List[int]], int | typing.List[int] | bytes)
+        eq(x[typing.Hashable], int | typing.Hashable | bytes)
+        eq(x[collections.abc.Hashable],
+           int | collections.abc.Hashable | bytes, typed=False)
+        eq(x[typing.Callable[[int], str]],
+           int | typing.Callable[[int], str] | bytes)
+        eq(x[collections.abc.Callable[[int], str]],
+           int | collections.abc.Callable[[int], str] | bytes, typed=False)
+        eq(x[typing.Tuple[int, str]], int | typing.Tuple[int, str] | bytes)
+        eq(x[typing.Literal['none']], int | typing.Literal['none'] | bytes)
+        eq(x[str | list], int | str | list | bytes, typed=False)
+        eq(x[typing.Union[str, list]], typing.Union[int, str, list, bytes])
+        eq(x[str | int], int | str | bytes, typed=False)
+        eq(x[typing.Union[str, int]], typing.Union[int, str, bytes])
+        eq(x[NT], int | NT | bytes)
+        eq(x[S], int | S | bytes)
+
+    def test_union_pickle(self):
+        orig = list[T] | int
+        for proto in range(pickle.HIGHEST_PROTOCOL + 1):
+            s = pickle.dumps(orig, proto)
+            loaded = pickle.loads(s)
+            self.assertEqual(loaded, orig)
+            self.assertEqual(loaded.__args__, orig.__args__)
+            self.assertEqual(loaded.__parameters__, orig.__parameters__)
+
+    def test_union_copy(self):
+        orig = list[T] | int
+        for copied in (copy.copy(orig), copy.deepcopy(orig)):
+            self.assertEqual(copied, orig)
+            self.assertEqual(copied.__args__, orig.__args__)
+            self.assertEqual(copied.__parameters__, orig.__parameters__)
+
+    def test_union_parameter_substitution_errors(self):
+        T = typing.TypeVar("T")
+        x = int | T
+        with self.assertRaises(TypeError):
+            x[int, str]
+
+    def test_or_type_operator_with_forward(self):
+        T = typing.TypeVar('T')
+        ForwardAfter = T | 'Forward'
+        ForwardBefore = 'Forward' | T
+        def forward_after(x: ForwardAfter[int]) -> None: ...
+        def forward_before(x: ForwardBefore[int]) -> None: ...
+        self.assertEqual(typing.get_args(typing.get_type_hints(forward_after)['x']),
+                         (int, Forward))
+        self.assertEqual(typing.get_args(typing.get_type_hints(forward_before)['x']),
+                         (int, Forward))
+
+    def test_or_type_operator_with_Protocol(self):
+        with torch._dynamo.error_on_graph_break(False):
+            class Proto(typing.Protocol):
+                def meth(self) -> int:
+                    ...
+        self.assertEqual(Proto | str, typing.Union[Proto, str])
+
+    def test_or_type_operator_with_Alias(self):
+        self.assertEqual(list | str, typing.Union[list, str])
+        self.assertEqual(typing.List | str, typing.Union[typing.List, str])
+
+    def test_or_type_operator_with_NamedTuple(self):
+        NT = namedtuple('A', ['B', 'C', 'D'])
+        self.assertEqual(NT | str, typing.Union[NT, str])
+
+    def test_or_type_operator_with_TypedDict(self):
+        with torch._dynamo.error_on_graph_break(False):
+            class Point2D(typing.TypedDict):
+                x: int
+                y: int
+                label: str
+        self.assertEqual(Point2D | str, typing.Union[Point2D, str])
+
+    def test_or_type_operator_with_NewType(self):
+        UserId = typing.NewType('UserId', int)
+        self.assertEqual(UserId | str, typing.Union[UserId, str])
+
+    def test_or_type_operator_with_IO(self):
+        self.assertEqual(typing.IO | str, typing.Union[typing.IO, str])
+
+    def test_or_type_operator_with_SpecialForm(self):
+        self.assertEqual(typing.Any | str, typing.Union[typing.Any, str])
+        self.assertEqual(typing.NoReturn | str, typing.Union[typing.NoReturn, str])
+        self.assertEqual(typing.Optional[int] | str, typing.Union[typing.Optional[int], str])
+        self.assertEqual(typing.Optional[int] | str, typing.Union[int, str, None])
+        self.assertEqual(typing.Union[int, bool] | str, typing.Union[int, bool, str])
+
+    def test_or_type_operator_with_Literal(self):
+        Literal = typing.Literal
+        self.assertEqual((Literal[1] | Literal[2]).__args__,
+                         (Literal[1], Literal[2]))
+
+        self.assertEqual((Literal[0] | Literal[False]).__args__,
+                         (Literal[0], Literal[False]))
+        self.assertEqual((Literal[1] | Literal[True]).__args__,
+                         (Literal[1], Literal[True]))
+
+        self.assertEqual(Literal[1] | Literal[1], Literal[1])
+        self.assertEqual(Literal['a'] | Literal['a'], Literal['a'])
+
+        import enum
+        with torch._dynamo.error_on_graph_break(False):
+            class Ints(enum.IntEnum):
+                A = 0
+                B = 1
+
+        self.assertEqual(Literal[Ints.A] | Literal[Ints.A], Literal[Ints.A])
+        self.assertEqual(Literal[Ints.B] | Literal[Ints.B], Literal[Ints.B])
+
+        self.assertEqual((Literal[Ints.B] | Literal[Ints.A]).__args__,
+                         (Literal[Ints.B], Literal[Ints.A]))
+
+        self.assertEqual((Literal[0] | Literal[Ints.A]).__args__,
+                         (Literal[0], Literal[Ints.A]))
+        self.assertEqual((Literal[1] | Literal[Ints.B]).__args__,
+                         (Literal[1], Literal[Ints.B]))
+
+    def test_or_type_repr(self):
+        self.assertEqual(repr(int | str), "int | str")
+        self.assertEqual(repr((int | str) | list), "int | str | list")
+        self.assertEqual(repr(int | (str | list)), "int | str | list")
+        self.assertEqual(repr(int | None), "int | None")
+        self.assertEqual(repr(int | type(None)), "int | None")
+        self.assertEqual(repr(int | typing.GenericAlias(list, int)), "int | list[int]")
+
+    def test_or_type_operator_with_genericalias(self):
+        a = list[int]
+        b = list[str]
+        c = dict[float, str]
+        with torch._dynamo.error_on_graph_break(False):
+            class SubClass(types.GenericAlias): ...
+        d = SubClass(list, float)
+        # equivalence with typing.Union
+        self.assertEqual(a | b | c | d, typing.Union[a, b, c, d])
+        # de-duplicate
+        self.assertEqual(a | c | b | b | a | c | d | d, a | b | c | d)
+        # order shouldn't matter
+        self.assertEqual(a | b | d, b | a | d)
+        self.assertEqual(repr(a | b | c | d),
+                         "list[int] | list[str] | dict[float, str] | list[float]")
+
+        with torch._dynamo.error_on_graph_break(False):
+            class BadType(type):
+                def __eq__(self, other):
+                    return 1 / 0
+
+        bt = BadType('bt', (), {})
+        # Comparison should fail and errors should propagate out for bad types.
+        with self.assertRaises(ZeroDivisionError):
+            list[int] | list[bt]
+
+        union_ga = (list[str] | int, collections.abc.Callable[..., str] | int,
+                    d | int)
+        # Raise error when isinstance(type, genericalias | type)
+        for type_ in union_ga:
+            with self.subTest(f"check isinstance/issubclass is invalid for {type_}"):
+                with self.assertRaises(TypeError):
+                    isinstance(1, type_)
+                with self.assertRaises(TypeError):
+                    issubclass(int, type_)
+
+    def test_or_type_operator_with_bad_module(self):
+        with torch._dynamo.error_on_graph_break(False):
+            class BadMeta(type):
+                __qualname__ = 'TypeVar'
+                @property
+                def __module__(self):
+                    1 / 0
+        TypeVar = BadMeta('TypeVar', (), {})
+        _SpecialForm = BadMeta('_SpecialForm', (), {})
+        # Crashes in Issue44483
+        with self.assertRaises((TypeError, ZeroDivisionError)):
+            str | TypeVar()
+        with self.assertRaises((TypeError, ZeroDivisionError)):
+            str | _SpecialForm()
+
+    @cpython_only
+    def test_or_type_operator_reference_cycle(self):
+        if not hasattr(sys, 'gettotalrefcount'):
+            self.skipTest('Cannot get total reference count.')
+        gc.collect()
+        before = sys.gettotalrefcount()
+        for _ in range(30):
+            T = typing.TypeVar('T')
+            U = int | list[T]
+            T.blah = U
+            del T
+            del U
+        gc.collect()
+        leeway = 15
+        self.assertLessEqual(sys.gettotalrefcount() - before, leeway,
+                             msg='Check for union reference leak.')
+
+
+class MappingProxyTests(CPythonTestCase):
+    mappingproxy = types.MappingProxyType
+
+    def test_constructor(self):
+        with torch._dynamo.error_on_graph_break(False):
+            class userdict(dict):
+                pass
+
+        mapping = {'x': 1, 'y': 2}
+        self.assertEqual(self.mappingproxy(mapping), mapping)
+        mapping = userdict(x=1, y=2)
+        self.assertEqual(self.mappingproxy(mapping), mapping)
+        mapping = collections.ChainMap({'x': 1}, {'y': 2})
+        self.assertEqual(self.mappingproxy(mapping), mapping)
+
+        self.assertRaises(TypeError, self.mappingproxy, 10)
+        self.assertRaises(TypeError, self.mappingproxy, ("a", "tuple"))
+        self.assertRaises(TypeError, self.mappingproxy, ["a", "list"])
+
+    def test_methods(self):
+        attrs = set(dir(self.mappingproxy({}))) - set(dir(object()))
+        self.assertEqual(attrs, {
+             '__contains__',
+             '__getitem__',
+             '__class_getitem__',
+             '__ior__',
+             '__iter__',
+             '__len__',
+             '__or__',
+             '__reversed__',
+             '__ror__',
+             'copy',
+             'get',
+             'items',
+             'keys',
+             'values',
+        })
+
+    def test_get(self):
+        view = self.mappingproxy({'a': 'A', 'b': 'B'})
+        self.assertEqual(view['a'], 'A')
+        self.assertEqual(view['b'], 'B')
+        self.assertRaises(KeyError, view.__getitem__, 'xxx')
+        self.assertEqual(view.get('a'), 'A')
+        self.assertIsNone(view.get('xxx'))
+        self.assertEqual(view.get('xxx', 42), 42)
+
+    def test_missing(self):
+        with torch._dynamo.error_on_graph_break(False):
+            class dictmissing(dict):
+                def __missing__(self, key):
+                    return "missing=%s" % key
+
+        view = self.mappingproxy(dictmissing(x=1))
+        self.assertEqual(view['x'], 1)
+        self.assertEqual(view['y'], 'missing=y')
+        self.assertEqual(view.get('x'), 1)
+        self.assertEqual(view.get('y'), None)
+        self.assertEqual(view.get('y', 42), 42)
+        self.assertTrue('x' in view)
+        self.assertFalse('y' in view)
+
+    def test_customdict(self):
+        with torch._dynamo.error_on_graph_break(False):
+            class customdict(dict):
+                def __contains__(self, key):
+                    if key == 'magic':
+                        return True
+                    else:
+                        return dict.__contains__(self, key)
+
+                def __iter__(self):
+                    return iter(('iter',))
+
+                def __len__(self):
+                    return 500
+
+                def copy(self):
+                    return 'copy'
+
+                def keys(self):
+                    return 'keys'
+
+                def items(self):
+                    return 'items'
+
+                def values(self):
+                    return 'values'
+
+                def __getitem__(self, key):
+                    return "getitem=%s" % dict.__getitem__(self, key)
+
+                def get(self, key, default=None):
+                    return "get=%s" % dict.get(self, key, 'default=%r' % default)
+
+        custom = customdict({'key': 'value'})
+        view = self.mappingproxy(custom)
+        self.assertTrue('key' in view)
+        self.assertTrue('magic' in view)
+        self.assertFalse('xxx' in view)
+        self.assertEqual(view['key'], 'getitem=value')
+        self.assertRaises(KeyError, view.__getitem__, 'xxx')
+        self.assertEqual(tuple(view), ('iter',))
+        self.assertEqual(len(view), 500)
+        self.assertEqual(view.copy(), 'copy')
+        self.assertEqual(view.get('key'), 'get=value')
+        self.assertEqual(view.get('xxx'), 'get=default=None')
+        self.assertEqual(view.items(), 'items')
+        self.assertEqual(view.keys(), 'keys')
+        self.assertEqual(view.values(), 'values')
+
+    def test_chainmap(self):
+        d1 = {'x': 1}
+        d2 = {'y': 2}
+        mapping = collections.ChainMap(d1, d2)
+        view = self.mappingproxy(mapping)
+        self.assertTrue('x' in view)
+        self.assertTrue('y' in view)
+        self.assertFalse('z' in view)
+        self.assertEqual(view['x'], 1)
+        self.assertEqual(view['y'], 2)
+        self.assertRaises(KeyError, view.__getitem__, 'z')
+        self.assertEqual(tuple(sorted(view)), ('x', 'y'))
+        self.assertEqual(len(view), 2)
+        copy = view.copy()
+        self.assertIsNot(copy, mapping)
+        self.assertIsInstance(copy, collections.ChainMap)
+        self.assertEqual(copy, mapping)
+        self.assertEqual(view.get('x'), 1)
+        self.assertEqual(view.get('y'), 2)
+        self.assertIsNone(view.get('z'))
+        self.assertEqual(tuple(sorted(view.items())), (('x', 1), ('y', 2)))
+        self.assertEqual(tuple(sorted(view.keys())), ('x', 'y'))
+        self.assertEqual(tuple(sorted(view.values())), (1, 2))
+
+    def test_contains(self):
+        view = self.mappingproxy(dict.fromkeys('abc'))
+        self.assertTrue('a' in view)
+        self.assertTrue('b' in view)
+        self.assertTrue('c' in view)
+        self.assertFalse('xxx' in view)
+
+    def test_views(self):
+        mapping = {}
+        view = self.mappingproxy(mapping)
+        keys = view.keys()
+        values = view.values()
+        items = view.items()
+        self.assertEqual(list(keys), [])
+        self.assertEqual(list(values), [])
+        self.assertEqual(list(items), [])
+        mapping['key'] = 'value'
+        self.assertEqual(list(keys), ['key'])
+        self.assertEqual(list(values), ['value'])
+        self.assertEqual(list(items), [('key', 'value')])
+
+    def test_len(self):
+        for expected in range(6):
+            data = dict.fromkeys('abcde'[:expected])
+            self.assertEqual(len(data), expected)
+            view = self.mappingproxy(data)
+            self.assertEqual(len(view), expected)
+
+    def test_iterators(self):
+        keys = ('x', 'y')
+        values = (1, 2)
+        items = tuple(zip(keys, values))
+        view = self.mappingproxy(dict(items))
+        self.assertEqual(set(view), set(keys))
+        self.assertEqual(set(view.keys()), set(keys))
+        self.assertEqual(set(view.values()), set(values))
+        self.assertEqual(set(view.items()), set(items))
+
+    def test_reversed(self):
+        d = {'a': 1, 'b': 2, 'foo': 0, 'c': 3, 'd': 4}
+        mp = self.mappingproxy(d)
+        del d['foo']
+        r = reversed(mp)
+        self.assertEqual(list(r), list('dcba'))
+        self.assertRaises(StopIteration, next, r)
+
+    def test_copy(self):
+        original = {'key1': 27, 'key2': 51, 'key3': 93}
+        view = self.mappingproxy(original)
+        copy = view.copy()
+        self.assertEqual(type(copy), dict)
+        self.assertEqual(copy, original)
+        original['key1'] = 70
+        self.assertEqual(view['key1'], 70)
+        self.assertEqual(copy['key1'], 27)
+
+    def test_union(self):
+        mapping = {'a': 0, 'b': 1, 'c': 2}
+        view = self.mappingproxy(mapping)
+        with self.assertRaises(TypeError):
+            view | [('r', 2), ('d', 2)]
+        with self.assertRaises(TypeError):
+            [('r', 2), ('d', 2)] | view
+        with self.assertRaises(TypeError):
+            view |= [('r', 2), ('d', 2)]
+        other = {'c': 3, 'p': 0}
+        self.assertDictEqual(view | other, {'a': 0, 'b': 1, 'c': 3, 'p': 0})
+        self.assertDictEqual(other | view, {'c': 2, 'p': 0, 'a': 0, 'b': 1})
+        self.assertEqual(view, {'a': 0, 'b': 1, 'c': 2})
+        self.assertDictEqual(mapping, {'a': 0, 'b': 1, 'c': 2})
+        self.assertDictEqual(other, {'c': 3, 'p': 0})
+
+    def test_hash(self):
+        with torch._dynamo.error_on_graph_break(False):
+            class HashableDict(dict):
+                def __hash__(self):
+                    return 3844817361
+        view = self.mappingproxy({'a': 1, 'b': 2})
+        self.assertRaises(TypeError, hash, view)
+        mapping = HashableDict({'a': 1, 'b': 2})
+        view = self.mappingproxy(mapping)
+        self.assertEqual(hash(view), hash(mapping))
+
+
+class ClassCreationTests(CPythonTestCase):
+
+    class Meta(type):
+        def __init__(cls, name, bases, ns, **kw):
+            super().__init__(name, bases, ns)
+        @staticmethod
+        def __new__(mcls, name, bases, ns, **kw):
+            return super().__new__(mcls, name, bases, ns)
+        @classmethod
+        def __prepare__(mcls, name, bases, **kw):
+            ns = super().__prepare__(name, bases)
+            ns["y"] = 1
+            ns.update(kw)
+            return ns
+
+    def test_new_class_basics(self):
+        C = types.new_class("C")
+        self.assertEqual(C.__name__, "C")
+        self.assertEqual(C.__bases__, (object,))
+
+    def test_new_class_subclass(self):
+        C = types.new_class("C", (int,))
+        self.assertTrue(issubclass(C, int))
+
+    def test_new_class_meta(self):
+        Meta = self.Meta
+        settings = {"metaclass": Meta, "z": 2}
+        # We do this twice to make sure the passed in dict isn't mutated
+        for i in range(2):
+            C = types.new_class("C" + str(i), (), settings)
+            self.assertIsInstance(C, Meta)
+            self.assertEqual(C.y, 1)
+            self.assertEqual(C.z, 2)
+
+    def test_new_class_exec_body(self):
+        Meta = self.Meta
+        def func(ns):
+            ns["x"] = 0
+        C = types.new_class("C", (), {"metaclass": Meta, "z": 2}, func)
+        self.assertIsInstance(C, Meta)
+        self.assertEqual(C.x, 0)
+        self.assertEqual(C.y, 1)
+        self.assertEqual(C.z, 2)
+
+    def test_new_class_metaclass_keywords(self):
+        #Test that keywords are passed to the metaclass:
+        def meta_func(name, bases, ns, **kw):
+            return name, bases, ns, kw
+        res = types.new_class("X",
+                              (int, object),
+                              dict(metaclass=meta_func, x=0))
+        self.assertEqual(res, ("X", (int, object), {}, {"x": 0}))
+
+    def test_new_class_defaults(self):
+        # Test defaults/keywords:
+        C = types.new_class("C", (), {}, None)
+        self.assertEqual(C.__name__, "C")
+        self.assertEqual(C.__bases__, (object,))
+
+    def test_new_class_meta_with_base(self):
+        Meta = self.Meta
+        def func(ns):
+            ns["x"] = 0
+        C = types.new_class(name="C",
+                            bases=(int,),
+                            kwds=dict(metaclass=Meta, z=2),
+                            exec_body=func)
+        self.assertTrue(issubclass(C, int))
+        self.assertIsInstance(C, Meta)
+        self.assertEqual(C.x, 0)
+        self.assertEqual(C.y, 1)
+        self.assertEqual(C.z, 2)
+
+    def test_new_class_with_mro_entry(self):
+        with torch._dynamo.error_on_graph_break(False):
+            class A: pass
+        with torch._dynamo.error_on_graph_break(False):
+            class C:
+                def __mro_entries__(self, bases):
+                    return (A,)
+        c = C()
+        D = types.new_class('D', (c,), {})
+        self.assertEqual(D.__bases__, (A,))
+        self.assertEqual(D.__orig_bases__, (c,))
+        self.assertEqual(D.__mro__, (D, A, object))
+
+    def test_new_class_with_mro_entry_genericalias(self):
+        L1 = types.new_class('L1', (typing.List[int],), {})
+        self.assertEqual(L1.__bases__, (list, typing.Generic))
+        self.assertEqual(L1.__orig_bases__, (typing.List[int],))
+        self.assertEqual(L1.__mro__, (L1, list, typing.Generic, object))
+
+        L2 = types.new_class('L2', (list[int],), {})
+        self.assertEqual(L2.__bases__, (list,))
+        self.assertEqual(L2.__orig_bases__, (list[int],))
+        self.assertEqual(L2.__mro__, (L2, list, object))
+
+    def test_new_class_with_mro_entry_none(self):
+        with torch._dynamo.error_on_graph_break(False):
+            class A: pass
+            class B: pass
+            class C:
+                def __mro_entries__(self, bases):
+                    return ()
+        c = C()
+        D = types.new_class('D', (A, c, B), {})
+        self.assertEqual(D.__bases__, (A, B))
+        self.assertEqual(D.__orig_bases__, (A, c, B))
+        self.assertEqual(D.__mro__, (D, A, B, object))
+
+    def test_new_class_with_mro_entry_error(self):
+        with torch._dynamo.error_on_graph_break(False):
+            class A: pass
+            class C:
+                def __mro_entries__(self, bases):
+                    return A
+        c = C()
+        with self.assertRaises(TypeError):
+            types.new_class('D', (c,), {})
+
+    def test_new_class_with_mro_entry_multiple(self):
+        with torch._dynamo.error_on_graph_break(False):
+            class A1: pass
+            class A2: pass
+            class B1: pass
+            class B2: pass
+            class A:
+                def __mro_entries__(self, bases):
+                    return (A1, A2)
+            class B:
+                def __mro_entries__(self, bases):
+                    return (B1, B2)
+        D = types.new_class('D', (A(), B()), {})
+        self.assertEqual(D.__bases__, (A1, A2, B1, B2))
+
+    def test_new_class_with_mro_entry_multiple_2(self):
+        with torch._dynamo.error_on_graph_break(False):
+            class A1: pass
+            class A2: pass
+            class A3: pass
+            class B1: pass
+            class B2: pass
+            class A:
+                def __mro_entries__(self, bases):
+                    return (A1, A2, A3)
+            class B:
+                def __mro_entries__(self, bases):
+                    return (B1, B2)
+            class C: pass
+        D = types.new_class('D', (A(), C, B()), {})
+        self.assertEqual(D.__bases__, (A1, A2, A3, C, B1, B2))
+
+    def test_get_original_bases(self):
+        T = typing.TypeVar('T')
+        with torch._dynamo.error_on_graph_break(False):
+            class A: pass
+            class B(typing.Generic[T]): pass
+            class C(B[int]): pass
+            class D(B[str], float): pass
+
+        self.assertEqual(types.get_original_bases(A), (object,))
+        self.assertEqual(types.get_original_bases(B), (typing.Generic[T],))
+        self.assertEqual(types.get_original_bases(C), (B[int],))
+        self.assertEqual(types.get_original_bases(int), (object,))
+        self.assertEqual(types.get_original_bases(D), (B[str], float))
+
+        with torch._dynamo.error_on_graph_break(False):
+            class E(list[T]): pass
+            class F(list[int]): pass
+
+        self.assertEqual(types.get_original_bases(E), (list[T],))
+        self.assertEqual(types.get_original_bases(F), (list[int],))
+
+        with torch._dynamo.error_on_graph_break(False):
+            class FirstBase(typing.Generic[T]): pass
+            class SecondBase(typing.Generic[T]): pass
+            class First(FirstBase[int]): pass
+            class Second(SecondBase[int]): pass
+            class G(First, Second): pass
+        self.assertEqual(types.get_original_bases(G), (First, Second))
+
+        with torch._dynamo.error_on_graph_break(False):
+            class First_(typing.Generic[T]): pass
+            class Second_(typing.Generic[T]): pass
+            class H(First_, Second_): pass
+        self.assertEqual(types.get_original_bases(H), (First_, Second_))
+
+        with torch._dynamo.error_on_graph_break(False):
+            class ClassBasedNamedTuple(typing.NamedTuple):
+                x: int
+
+            class GenericNamedTuple(typing.NamedTuple, typing.Generic[T]):
+                x: T
+
+        CallBasedNamedTuple = typing.NamedTuple("CallBasedNamedTuple", [("x", int)])
+
+        self.assertIs(
+            types.get_original_bases(ClassBasedNamedTuple)[0], typing.NamedTuple
+        )
+        self.assertEqual(
+            types.get_original_bases(GenericNamedTuple),
+            (typing.NamedTuple, typing.Generic[T])
+        )
+        self.assertIs(
+            types.get_original_bases(CallBasedNamedTuple)[0], typing.NamedTuple
+        )
+
+        with torch._dynamo.error_on_graph_break(False):
+            class ClassBasedTypedDict(typing.TypedDict):
+                x: int
+
+            class GenericTypedDict(typing.TypedDict, typing.Generic[T]):
+                x: T
+
+        CallBasedTypedDict = typing.TypedDict("CallBasedTypedDict", {"x": int})
+
+        self.assertIs(
+            types.get_original_bases(ClassBasedTypedDict)[0],
+            typing.TypedDict
+        )
+        self.assertEqual(
+            types.get_original_bases(GenericTypedDict),
+            (typing.TypedDict, typing.Generic[T])
+        )
+        self.assertIs(
+            types.get_original_bases(CallBasedTypedDict)[0],
+            typing.TypedDict
+        )
+
+        with self.assertRaisesRegex(TypeError, "Expected an instance of type"):
+            types.get_original_bases(object())
+
+    # Many of the following tests are derived from test_descr.py
+    def test_prepare_class(self):
+        # Basic test of metaclass derivation
+        expected_ns = {}
+        with torch._dynamo.error_on_graph_break(False):
+            class A(type):
+                def __new__(*args, **kwargs):
+                    return type.__new__(*args, **kwargs)
+
+                def __prepare__(*args):
+                    return expected_ns
+
+        B = types.new_class("B", (object,))
+        C = types.new_class("C", (object,), {"metaclass": A})
+
+        # The most derived metaclass of D is A rather than type.
+        meta, ns, kwds = types.prepare_class("D", (B, C), {"metaclass": type})
+        self.assertIs(meta, A)
+        self.assertIs(ns, expected_ns)
+        self.assertEqual(len(kwds), 0)
+
+    def test_bad___prepare__(self):
+        # __prepare__() must return a mapping.
+        with torch._dynamo.error_on_graph_break(False):
+            class BadMeta(type):
+                @classmethod
+                def __prepare__(*args):
+                    return None
+        with self.assertRaisesRegex(TypeError,
+                                    r'^BadMeta\.__prepare__\(\) must '
+                                    r'return a mapping, not NoneType$'):
+            class Foo(metaclass=BadMeta):
+                pass
+        # Also test the case in which the metaclass is not a type.
+        with torch._dynamo.error_on_graph_break(False):
+            class BadMeta:
+                @classmethod
+                def __prepare__(*args):
+                    return None
+        with self.assertRaisesRegex(TypeError,
+                                    r'^<metaclass>\.__prepare__\(\) must '
+                                    r'return a mapping, not NoneType$'):
+            class Bar(metaclass=BadMeta()):
+                pass
+
+    def test_resolve_bases(self):
+        with torch._dynamo.error_on_graph_break(False):
+            class A: pass
+            class B: pass
+            class C:
+                def __mro_entries__(self, bases):
+                    if A in bases:
+                        return ()
+                    return (A,)
+        c = C()
+        self.assertEqual(types.resolve_bases(()), ())
+        self.assertEqual(types.resolve_bases((c,)), (A,))
+        self.assertEqual(types.resolve_bases((C,)), (C,))
+        self.assertEqual(types.resolve_bases((A, C)), (A, C))
+        self.assertEqual(types.resolve_bases((c, A)), (A,))
+        self.assertEqual(types.resolve_bases((A, c)), (A,))
+        x = (A,)
+        y = (C,)
+        z = (A, C)
+        t = (A, C, B)
+        for bases in [x, y, z, t]:
+            self.assertIs(types.resolve_bases(bases), bases)
+
+    def test_resolve_bases_with_mro_entry(self):
+        self.assertEqual(types.resolve_bases((typing.List[int],)),
+                         (list, typing.Generic))
+        self.assertEqual(types.resolve_bases((list[int],)), (list,))
+
+    def test_metaclass_derivation(self):
+        # issue1294232: correct metaclass calculation
+        new_calls = []  # to check the order of __new__ calls
+        with torch._dynamo.error_on_graph_break(False):
+            class AMeta(type):
+                def __new__(mcls, name, bases, ns):
+                    new_calls.append('AMeta')
+                    return super().__new__(mcls, name, bases, ns)
+                @classmethod
+                def __prepare__(mcls, name, bases):
+                    return {}
+
+            class BMeta(AMeta):
+                def __new__(mcls, name, bases, ns):
+                    new_calls.append('BMeta')
+                    return super().__new__(mcls, name, bases, ns)
+                @classmethod
+                def __prepare__(mcls, name, bases):
+                    ns = super().__prepare__(name, bases)
+                    ns['BMeta_was_here'] = True
+                    return ns
+
+        A = types.new_class("A", (), {"metaclass": AMeta})
+        self.assertEqual(new_calls, ['AMeta'])
+        new_calls.clear()
+
+        B = types.new_class("B", (), {"metaclass": BMeta})
+        # BMeta.__new__ calls AMeta.__new__ with super:
+        self.assertEqual(new_calls, ['BMeta', 'AMeta'])
+        new_calls.clear()
+
+        C = types.new_class("C", (A, B))
+        # The most derived metaclass is BMeta:
+        self.assertEqual(new_calls, ['BMeta', 'AMeta'])
+        new_calls.clear()
+        # BMeta.__prepare__ should've been called:
+        self.assertIn('BMeta_was_here', C.__dict__)
+
+        # The order of the bases shouldn't matter:
+        C2 = types.new_class("C2", (B, A))
+        self.assertEqual(new_calls, ['BMeta', 'AMeta'])
+        new_calls.clear()
+        self.assertIn('BMeta_was_here', C2.__dict__)
+
+        # Check correct metaclass calculation when a metaclass is declared:
+        D = types.new_class("D", (C,), {"metaclass": type})
+        self.assertEqual(new_calls, ['BMeta', 'AMeta'])
+        new_calls.clear()
+        self.assertIn('BMeta_was_here', D.__dict__)
+
+        E = types.new_class("E", (C,), {"metaclass": AMeta})
+        self.assertEqual(new_calls, ['BMeta', 'AMeta'])
+        new_calls.clear()
+        self.assertIn('BMeta_was_here', E.__dict__)
+
+    def test_metaclass_override_function(self):
+        # Special case: the given metaclass isn't a class,
+        # so there is no metaclass calculation.
+        with torch._dynamo.error_on_graph_break(False):
+            class A(metaclass=self.Meta):
+                pass
+
+        marker = object()
+        def func(*args, **kwargs):
+            return marker
+
+        X = types.new_class("X", (), {"metaclass": func})
+        Y = types.new_class("Y", (object,), {"metaclass": func})
+        Z = types.new_class("Z", (A,), {"metaclass": func})
+        self.assertIs(marker, X)
+        self.assertIs(marker, Y)
+        self.assertIs(marker, Z)
+
+    def test_metaclass_override_callable(self):
+        # The given metaclass is a class,
+        # but not a descendant of type.
+        new_calls = []  # to check the order of __new__ calls
+        prepare_calls = []  # to track __prepare__ calls
+        with torch._dynamo.error_on_graph_break(False):
+            class ANotMeta:
+                def __new__(mcls, *args, **kwargs):
+                    new_calls.append('ANotMeta')
+                    return super().__new__(mcls)
+                @classmethod
+                def __prepare__(mcls, name, bases):
+                    prepare_calls.append('ANotMeta')
+                    return {}
+
+            class BNotMeta(ANotMeta):
+                def __new__(mcls, *args, **kwargs):
+                    new_calls.append('BNotMeta')
+                    return super().__new__(mcls)
+                @classmethod
+                def __prepare__(mcls, name, bases):
+                    prepare_calls.append('BNotMeta')
+                    return super().__prepare__(name, bases)
+
+        A = types.new_class("A", (), {"metaclass": ANotMeta})
+        self.assertIs(ANotMeta, type(A))
+        self.assertEqual(prepare_calls, ['ANotMeta'])
+        prepare_calls.clear()
+        self.assertEqual(new_calls, ['ANotMeta'])
+        new_calls.clear()
+
+        B = types.new_class("B", (), {"metaclass": BNotMeta})
+        self.assertIs(BNotMeta, type(B))
+        self.assertEqual(prepare_calls, ['BNotMeta', 'ANotMeta'])
+        prepare_calls.clear()
+        self.assertEqual(new_calls, ['BNotMeta', 'ANotMeta'])
+        new_calls.clear()
+
+        C = types.new_class("C", (A, B))
+        self.assertIs(BNotMeta, type(C))
+        self.assertEqual(prepare_calls, ['BNotMeta', 'ANotMeta'])
+        prepare_calls.clear()
+        self.assertEqual(new_calls, ['BNotMeta', 'ANotMeta'])
+        new_calls.clear()
+
+        C2 = types.new_class("C2", (B, A))
+        self.assertIs(BNotMeta, type(C2))
+        self.assertEqual(prepare_calls, ['BNotMeta', 'ANotMeta'])
+        prepare_calls.clear()
+        self.assertEqual(new_calls, ['BNotMeta', 'ANotMeta'])
+        new_calls.clear()
+
+        # This is a TypeError, because of a metaclass conflict:
+        # BNotMeta is neither a subclass, nor a superclass of type
+        with self.assertRaises(TypeError):
+            D = types.new_class("D", (C,), {"metaclass": type})
+
+        E = types.new_class("E", (C,), {"metaclass": ANotMeta})
+        self.assertIs(BNotMeta, type(E))
+        self.assertEqual(prepare_calls, ['BNotMeta', 'ANotMeta'])
+        prepare_calls.clear()
+        self.assertEqual(new_calls, ['BNotMeta', 'ANotMeta'])
+        new_calls.clear()
+
+        F = types.new_class("F", (object(), C))
+        self.assertIs(BNotMeta, type(F))
+        self.assertEqual(prepare_calls, ['BNotMeta', 'ANotMeta'])
+        prepare_calls.clear()
+        self.assertEqual(new_calls, ['BNotMeta', 'ANotMeta'])
+        new_calls.clear()
+
+        F2 = types.new_class("F2", (C, object()))
+        self.assertIs(BNotMeta, type(F2))
+        self.assertEqual(prepare_calls, ['BNotMeta', 'ANotMeta'])
+        prepare_calls.clear()
+        self.assertEqual(new_calls, ['BNotMeta', 'ANotMeta'])
+        new_calls.clear()
+
+        # TypeError: BNotMeta is neither a
+        # subclass, nor a superclass of int
+        with self.assertRaises(TypeError):
+            X = types.new_class("X", (C, int()))
+        with self.assertRaises(TypeError):
+            X = types.new_class("X", (int(), C))
+
+    def test_one_argument_type(self):
+        expected_message = 'type.__new__() takes exactly 3 arguments (1 given)'
+
+        # Only type itself can use the one-argument form (#27157)
+        self.assertIs(type(5), int)
+
+        with torch._dynamo.error_on_graph_break(False):
+            class M(type):
+                pass
+        with self.assertRaises(TypeError) as cm:
+            M(5)
+        self.assertEqual(str(cm.exception), expected_message)
+
+        with torch._dynamo.error_on_graph_break(False):
+            class N(type, metaclass=M):
+                pass
+        with self.assertRaises(TypeError) as cm:
+            N(5)
+        self.assertEqual(str(cm.exception), expected_message)
+
+    def test_metaclass_new_error(self):
+        # bpo-44232: The C function type_new() must properly report the
+        # exception when a metaclass constructor raises an exception and the
+        # winner class is not the metaclass.
+        with torch._dynamo.error_on_graph_break(False):
+            class ModelBase(type):
+                def __new__(cls, name, bases, attrs):
+                    super_new = super().__new__
+                    new_class = super_new(cls, name, bases, {})
+                    if name != "Model":
+                        raise RuntimeWarning(f"{name=}")
+                    return new_class
+
+        with torch._dynamo.error_on_graph_break(False):
+            class Model(metaclass=ModelBase):
+                pass
+
+        with self.assertRaises(RuntimeWarning):
+            type("SouthPonies", (Model,), {})
+
+    def test_tuple_subclass_as_bases(self):
+        # gh-132176: it used to crash on using
+        # tuple subclass for as base classes.
+        with torch._dynamo.error_on_graph_break(False):
+            class TupleSubclass(tuple): pass
+
+        typ = type("typ", TupleSubclass((int, object)), {})
+        self.assertEqual(typ.__bases__, (int, object))
+        self.assertEqual(type(typ.__bases__), TupleSubclass)
+
+
+class SimpleNamespaceTests(CPythonTestCase):
+
+    def test_constructor(self):
+        def check(ns, expected):
+            self.assertEqual(len(ns.__dict__), len(expected))
+            self.assertEqual(vars(ns), expected)
+            # check order
+            self.assertEqual(list(vars(ns).items()), list(expected.items()))
+            for name in expected:
+                self.assertEqual(getattr(ns, name), expected[name])
+
+        check(types.SimpleNamespace(), {})
+        check(types.SimpleNamespace(x=1, y=2), {'x': 1, 'y': 2})
+        check(types.SimpleNamespace(**dict(x=1, y=2)), {'x': 1, 'y': 2})
+        check(types.SimpleNamespace({'x': 1, 'y': 2}, x=4, z=3),
+              {'x': 4, 'y': 2, 'z': 3})
+        check(types.SimpleNamespace([['x', 1], ['y', 2]], x=4, z=3),
+              {'x': 4, 'y': 2, 'z': 3})
+        check(types.SimpleNamespace(UserDict({'x': 1, 'y': 2}), x=4, z=3),
+              {'x': 4, 'y': 2, 'z': 3})
+        check(types.SimpleNamespace({'x': 1, 'y': 2}), {'x': 1, 'y': 2})
+        check(types.SimpleNamespace([['x', 1], ['y', 2]]), {'x': 1, 'y': 2})
+        check(types.SimpleNamespace([], x=4, z=3), {'x': 4, 'z': 3})
+        check(types.SimpleNamespace({}, x=4, z=3), {'x': 4, 'z': 3})
+        check(types.SimpleNamespace([]), {})
+        check(types.SimpleNamespace({}), {})
+
+        with self.assertRaises(TypeError):
+            types.SimpleNamespace([], [])  # too many positional arguments
+        with self.assertRaises(TypeError):
+            types.SimpleNamespace(1)  # not a mapping or iterable
+        with self.assertRaises(TypeError):
+            types.SimpleNamespace([1])  # non-iterable
+        with self.assertRaises(ValueError):
+            types.SimpleNamespace([['x']])  # not a pair
+        with self.assertRaises(ValueError):
+            types.SimpleNamespace([['x', 'y', 'z']])
+        with self.assertRaises(TypeError):
+            types.SimpleNamespace(**{1: 2})  # non-string key
+        with self.assertRaises(TypeError):
+            types.SimpleNamespace({1: 2})
+        with self.assertRaises(TypeError):
+            types.SimpleNamespace([[1, 2]])
+        with self.assertRaises(TypeError):
+            types.SimpleNamespace(UserDict({1: 2}))
+        with self.assertRaises(TypeError):
+            types.SimpleNamespace([[[], 2]])  # non-hashable key
+
+    def test_unbound(self):
+        ns1 = vars(types.SimpleNamespace())
+        ns2 = vars(types.SimpleNamespace(x=1, y=2))
+
+        self.assertEqual(ns1, {})
+        self.assertEqual(ns2, {'y': 2, 'x': 1})
+
+    def test_underlying_dict(self):
+        ns1 = types.SimpleNamespace()
+        ns2 = types.SimpleNamespace(x=1, y=2)
+        ns3 = types.SimpleNamespace(a=True, b=False)
+        mapping = ns3.__dict__
+        del ns3
+
+        self.assertEqual(ns1.__dict__, {})
+        self.assertEqual(ns2.__dict__, {'y': 2, 'x': 1})
+        self.assertEqual(mapping, dict(a=True, b=False))
+
+    def test_attrget(self):
+        ns = types.SimpleNamespace(x=1, y=2, w=3)
+
+        self.assertEqual(ns.x, 1)
+        self.assertEqual(ns.y, 2)
+        self.assertEqual(ns.w, 3)
+        with self.assertRaises(AttributeError):
+            ns.z
+
+    def test_attrset(self):
+        ns1 = types.SimpleNamespace()
+        ns2 = types.SimpleNamespace(x=1, y=2, w=3)
+        ns1.a = 'spam'
+        ns1.b = 'ham'
+        ns2.z = 4
+        ns2.theta = None
+
+        self.assertEqual(ns1.__dict__, dict(a='spam', b='ham'))
+        self.assertEqual(ns2.__dict__, dict(x=1, y=2, w=3, z=4, theta=None))
+
+    def test_attrdel(self):
+        ns1 = types.SimpleNamespace()
+        ns2 = types.SimpleNamespace(x=1, y=2, w=3)
+
+        with self.assertRaises(AttributeError):
+            del ns1.spam
+        with self.assertRaises(AttributeError):
+            del ns2.spam
+
+        del ns2.y
+        self.assertEqual(vars(ns2), dict(w=3, x=1))
+        ns2.y = 'spam'
+        self.assertEqual(vars(ns2), dict(w=3, x=1, y='spam'))
+        del ns2.y
+        self.assertEqual(vars(ns2), dict(w=3, x=1))
+
+        ns1.spam = 5
+        self.assertEqual(vars(ns1), dict(spam=5))
+        del ns1.spam
+        self.assertEqual(vars(ns1), {})
+
+    def test_repr(self):
+        ns1 = types.SimpleNamespace(x=1, y=2, w=3)
+        ns2 = types.SimpleNamespace()
+        ns2.x = "spam"
+        ns2._y = 5
+        name = "namespace"
+
+        self.assertEqual(repr(ns1), "{name}(x=1, y=2, w=3)".format(name=name))
+        self.assertEqual(repr(ns2), "{name}(x='spam', _y=5)".format(name=name))
+
+    def test_equal(self):
+        ns1 = types.SimpleNamespace(x=1)
+        ns2 = types.SimpleNamespace()
+        ns2.x = 1
+
+        self.assertEqual(types.SimpleNamespace(), types.SimpleNamespace())
+        self.assertEqual(ns1, ns2)
+        self.assertNotEqual(ns2, types.SimpleNamespace())
+
+    def test_nested(self):
+        ns1 = types.SimpleNamespace(a=1, b=2)
+        ns2 = types.SimpleNamespace()
+        ns3 = types.SimpleNamespace(x=ns1)
+        ns2.spam = ns1
+        ns2.ham = '?'
+        ns2.spam = ns3
+
+        self.assertEqual(vars(ns1), dict(a=1, b=2))
+        self.assertEqual(vars(ns2), dict(spam=ns3, ham='?'))
+        self.assertEqual(ns2.spam, ns3)
+        self.assertEqual(vars(ns3), dict(x=ns1))
+        self.assertEqual(ns3.x.a, 1)
+
+    def test_recursive(self):
+        ns1 = types.SimpleNamespace(c='cookie')
+        ns2 = types.SimpleNamespace()
+        ns3 = types.SimpleNamespace(x=1)
+        ns1.spam = ns1
+        ns2.spam = ns3
+        ns3.spam = ns2
+
+        self.assertEqual(ns1.spam, ns1)
+        self.assertEqual(ns1.spam.spam, ns1)
+        self.assertEqual(ns1.spam.spam, ns1.spam)
+        self.assertEqual(ns2.spam, ns3)
+        self.assertEqual(ns3.spam, ns2)
+        self.assertEqual(ns2.spam.spam, ns2)
+
+    def test_recursive_repr(self):
+        ns1 = types.SimpleNamespace(c='cookie')
+        ns2 = types.SimpleNamespace()
+        ns3 = types.SimpleNamespace(x=1)
+        ns1.spam = ns1
+        ns2.spam = ns3
+        ns3.spam = ns2
+        name = "namespace"
+        repr1 = "{name}(c='cookie', spam={name}(...))".format(name=name)
+        repr2 = "{name}(spam={name}(x=1, spam={name}(...)))".format(name=name)
+
+        self.assertEqual(repr(ns1), repr1)
+        self.assertEqual(repr(ns2), repr2)
+
+    def test_as_dict(self):
+        ns = types.SimpleNamespace(spam='spamspamspam')
+
+        with self.assertRaises(TypeError):
+            len(ns)
+        with self.assertRaises(TypeError):
+            iter(ns)
+        with self.assertRaises(TypeError):
+            'spam' in ns
+        with self.assertRaises(TypeError):
+            ns['spam']
+
+    def test_subclass(self):
+        with torch._dynamo.error_on_graph_break(False):
+            class Spam(types.SimpleNamespace):
+                pass
+
+        spam = Spam(ham=8, eggs=9)
+
+        self.assertIs(type(spam), Spam)
+        self.assertEqual(vars(spam), {'ham': 8, 'eggs': 9})
+
+    def test_pickle(self):
+        ns = types.SimpleNamespace(breakfast="spam", lunch="spam")
+
+        for protocol in range(pickle.HIGHEST_PROTOCOL + 1):
+            pname = "protocol {}".format(protocol)
+            try:
+                ns_pickled = pickle.dumps(ns, protocol)
+            except TypeError as e:
+                raise TypeError(pname) from e
+            ns_roundtrip = pickle.loads(ns_pickled)
+
+            self.assertEqual(ns, ns_roundtrip, pname)
+
+    def test_replace(self):
+        ns = types.SimpleNamespace(x=11, y=22)
+
+        ns2 = copy.replace(ns)
+        self.assertEqual(ns2, ns)
+        self.assertIsNot(ns2, ns)
+        self.assertIs(type(ns2), types.SimpleNamespace)
+        self.assertEqual(vars(ns2), {'x': 11, 'y': 22})
+        ns2.x = 3
+        self.assertEqual(ns.x, 11)
+        ns.x = 4
+        self.assertEqual(ns2.x, 3)
+
+        self.assertEqual(vars(copy.replace(ns, x=1)), {'x': 1, 'y': 22})
+        self.assertEqual(vars(copy.replace(ns, y=2)), {'x': 4, 'y': 2})
+        self.assertEqual(vars(copy.replace(ns, x=1, y=2)), {'x': 1, 'y': 2})
+
+    def test_replace_subclass(self):
+        with torch._dynamo.error_on_graph_break(False):
+            class Spam(types.SimpleNamespace):
+                pass
+
+        spam = Spam(ham=8, eggs=9)
+        spam2 = copy.replace(spam, ham=5)
+
+        self.assertIs(type(spam2), Spam)
+        self.assertEqual(vars(spam2), {'ham': 5, 'eggs': 9})
+
+    def test_fake_namespace_compare(self):
+        # Issue #24257: Incorrect use of PyObject_IsInstance() caused
+        # SystemError.
+        with torch._dynamo.error_on_graph_break(False):
+            class FakeSimpleNamespace(str):
+                __class__ = types.SimpleNamespace
+        self.assertFalse(types.SimpleNamespace() == FakeSimpleNamespace())
+        self.assertTrue(types.SimpleNamespace() != FakeSimpleNamespace())
+        with self.assertRaises(TypeError):
+            types.SimpleNamespace() < FakeSimpleNamespace()
+        with self.assertRaises(TypeError):
+            types.SimpleNamespace() <= FakeSimpleNamespace()
+        with self.assertRaises(TypeError):
+            types.SimpleNamespace() > FakeSimpleNamespace()
+        with self.assertRaises(TypeError):
+            types.SimpleNamespace() >= FakeSimpleNamespace()
+
+
+class CoroutineTests(CPythonTestCase):
+    def test_wrong_args(self):
+        samples = [None, 1, object()]
+        for sample in samples:
+            with self.assertRaisesRegex(TypeError,
+                                        'types.coroutine.*expects a callable'):
+                types.coroutine(sample)
+
+    def test_non_gen_values(self):
+        @types.coroutine
+        def foo():
+            return 'spam'
+        self.assertEqual(foo(), 'spam')
+
+        with torch._dynamo.error_on_graph_break(False):
+            class Awaitable:
+                def __await__(self):
+                    return ()
+        aw = Awaitable()
+        @types.coroutine
+        def foo():
+            return aw
+        self.assertIs(aw, foo())
+
+        # decorate foo second time
+        foo = types.coroutine(foo)
+        self.assertIs(aw, foo())
+
+    def test_async_def(self):
+        # Test that types.coroutine passes 'async def' coroutines
+        # without modification
+
+        async def foo(): pass
+        foo_code = foo.__code__
+        foo_flags = foo.__code__.co_flags
+        decorated_foo = types.coroutine(foo)
+        self.assertIs(foo, decorated_foo)
+        self.assertEqual(foo.__code__.co_flags, foo_flags)
+        self.assertIs(decorated_foo.__code__, foo_code)
+
+        foo_coro = foo()
+        def bar(): return foo_coro
+        for _ in range(2):
+            bar = types.coroutine(bar)
+            coro = bar()
+            self.assertIs(foo_coro, coro)
+            self.assertEqual(coro.cr_code.co_flags, foo_flags)
+            coro.close()
+
+    def test_duck_coro(self):
+        with torch._dynamo.error_on_graph_break(False):
+            class CoroLike:
+                def send(self): pass
+                def throw(self): pass
+                def close(self): pass
+                def __await__(self): return self
+
+        coro = CoroLike()
+        @types.coroutine
+        def foo():
+            return coro
+        self.assertIs(foo(), coro)
+        self.assertIs(foo().__await__(), coro)
+
+    def test_duck_corogen(self):
+        with torch._dynamo.error_on_graph_break(False):
+            class CoroGenLike:
+                def send(self): pass
+                def throw(self): pass
+                def close(self): pass
+                def __await__(self): return self
+                def __iter__(self): return self
+                def __next__(self): pass
+
+        coro = CoroGenLike()
+        @types.coroutine
+        def foo():
+            return coro
+        self.assertIs(foo(), coro)
+        self.assertIs(foo().__await__(), coro)
+
+    def test_duck_gen(self):
+        with torch._dynamo.error_on_graph_break(False):
+            class GenLike:
+                def send(self): pass
+                def throw(self): pass
+                def close(self): pass
+                def __iter__(self): pass
+                def __next__(self): pass
+
+        # Setup generator mock object
+        gen = unittest.mock.MagicMock(GenLike)
+        gen.__iter__ = lambda gen: gen
+        gen.__name__ = 'gen'
+        gen.__qualname__ = 'test.gen'
+        self.assertIsInstance(gen, collections.abc.Generator)
+        self.assertIs(gen, iter(gen))
+
+        @types.coroutine
+        def foo(): return gen
+
+        wrapper = foo()
+        self.assertIsInstance(wrapper, types._GeneratorWrapper)
+        self.assertIs(wrapper.__await__(), wrapper)
+        # Wrapper proxies duck generators completely:
+        self.assertIs(iter(wrapper), wrapper)
+
+        self.assertIsInstance(wrapper, collections.abc.Coroutine)
+        self.assertIsInstance(wrapper, collections.abc.Awaitable)
+
+        self.assertIs(wrapper.__qualname__, gen.__qualname__)
+        self.assertIs(wrapper.__name__, gen.__name__)
+
+        # Test AttributeErrors
+        for name in {'gi_running', 'gi_frame', 'gi_code', 'gi_yieldfrom',
+                     'cr_running', 'cr_frame', 'cr_code', 'cr_await'}:
+            with self.assertRaises(AttributeError):
+                getattr(wrapper, name)
+
+        # Test attributes pass-through
+        gen.gi_running = object()
+        gen.gi_frame = object()
+        gen.gi_code = object()
+        gen.gi_yieldfrom = object()
+        self.assertIs(wrapper.gi_running, gen.gi_running)
+        self.assertIs(wrapper.gi_frame, gen.gi_frame)
+        self.assertIs(wrapper.gi_code, gen.gi_code)
+        self.assertIs(wrapper.gi_yieldfrom, gen.gi_yieldfrom)
+        self.assertIs(wrapper.cr_running, gen.gi_running)
+        self.assertIs(wrapper.cr_frame, gen.gi_frame)
+        self.assertIs(wrapper.cr_code, gen.gi_code)
+        self.assertIs(wrapper.cr_await, gen.gi_yieldfrom)
+
+        wrapper.close()
+        gen.close.assert_called_once_with()
+
+        wrapper.send(1)
+        gen.send.assert_called_once_with(1)
+        gen.reset_mock()
+
+        next(wrapper)
+        gen.__next__.assert_called_once_with()
+        gen.reset_mock()
+
+        wrapper.throw(1, 2, 3)
+        gen.throw.assert_called_once_with(1, 2, 3)
+        gen.reset_mock()
+
+        wrapper.throw(1, 2)
+        gen.throw.assert_called_once_with(1, 2)
+        gen.reset_mock()
+
+        wrapper.throw(1)
+        gen.throw.assert_called_once_with(1)
+        gen.reset_mock()
+
+        # Test exceptions propagation
+        error = Exception()
+        gen.throw.side_effect = error
+        try:
+            wrapper.throw(1)
+        except Exception as ex:
+            self.assertIs(ex, error)
+        else:
+            self.fail('wrapper did not propagate an exception')
+
+        # Test invalid args
+        gen.reset_mock()
+        with self.assertRaises(TypeError):
+            wrapper.throw()
+        self.assertFalse(gen.throw.called)
+        with self.assertRaises(TypeError):
+            wrapper.close(1)
+        self.assertFalse(gen.close.called)
+        with self.assertRaises(TypeError):
+            wrapper.send()
+        self.assertFalse(gen.send.called)
+
+        # Test that we do not double wrap
+        @types.coroutine
+        def bar(): return wrapper
+        self.assertIs(wrapper, bar())
+
+        # Test weakrefs support
+        ref = weakref.ref(wrapper)
+        self.assertIs(ref(), wrapper)
+
+    def test_duck_functional_gen(self):
+        with torch._dynamo.error_on_graph_break(False):
+            class Generator:
+                """Emulates the following generator (very clumsy):
+
+                  def gen(fut):
+                      result = yield fut
+                      return result * 2
+                """
+                def __init__(self, fut):
+                    self._i = 0
+                    self._fut = fut
+                def __iter__(self):
+                    return self
+                def __next__(self):
+                    return self.send(None)
+                def send(self, v):
+                    try:
+                        if self._i == 0:
+                            assert v is None
+                            return self._fut
+                        if self._i == 1:
+                            raise StopIteration(v * 2)
+                        if self._i > 1:
+                            raise StopIteration
+                    finally:
+                        self._i += 1
+                def throw(self, tp, *exc):
+                    self._i = 100
+                    if tp is not GeneratorExit:
+                        raise tp
+                def close(self):
+                    self.throw(GeneratorExit)
+
+        @types.coroutine
+        def foo(): return Generator('spam')
+
+        wrapper = foo()
+        self.assertIsInstance(wrapper, types._GeneratorWrapper)
+
+        async def corofunc():
+            return await foo() + 100
+        coro = corofunc()
+
+        self.assertEqual(coro.send(None), 'spam')
+        try:
+            coro.send(20)
+        except StopIteration as ex:
+            self.assertEqual(ex.args[0], 140)
+        else:
+            self.fail('StopIteration was expected')
+
+    def test_gen(self):
+        def gen_func():
+            yield 1
+            return (yield 2)
+        gen = gen_func()
+        @types.coroutine
+        def foo(): return gen
+        wrapper = foo()
+        self.assertIsInstance(wrapper, types._GeneratorWrapper)
+        self.assertIs(wrapper.__await__(), gen)
+
+        for name in ('__name__', '__qualname__', 'gi_code',
+                     'gi_running', 'gi_frame'):
+            self.assertIs(getattr(foo(), name),
+                          getattr(gen, name))
+        self.assertIs(foo().cr_code, gen.gi_code)
+
+        self.assertEqual(next(wrapper), 1)
+        self.assertEqual(wrapper.send(None), 2)
+        with self.assertRaisesRegex(StopIteration, 'spam'):
+            wrapper.send('spam')
+
+        gen = gen_func()
+        wrapper = foo()
+        wrapper.send(None)
+        with self.assertRaisesRegex(Exception, 'ham'):
+            wrapper.throw(Exception('ham'))
+
+        # decorate foo second time
+        foo = types.coroutine(foo)
+        self.assertIs(foo().__await__(), gen)
+
+    def test_returning_itercoro(self):
+        @types.coroutine
+        def gen():
+            yield
+
+        gencoro = gen()
+
+        @types.coroutine
+        def foo():
+            return gencoro
+
+        self.assertIs(foo(), gencoro)
+
+        # decorate foo second time
+        foo = types.coroutine(foo)
+        self.assertIs(foo(), gencoro)
+
+    def test_genfunc(self):
+        def gen(): yield
+        self.assertIs(types.coroutine(gen), gen)
+        self.assertIs(types.coroutine(types.coroutine(gen)), gen)
+
+        self.assertTrue(gen.__code__.co_flags & inspect.CO_ITERABLE_COROUTINE)
+        self.assertFalse(gen.__code__.co_flags & inspect.CO_COROUTINE)
+
+        g = gen()
+        self.assertTrue(g.gi_code.co_flags & inspect.CO_ITERABLE_COROUTINE)
+        self.assertFalse(g.gi_code.co_flags & inspect.CO_COROUTINE)
+
+        self.assertIs(types.coroutine(gen), gen)
+
+    def test_wrapper_object(self):
+        def gen():
+            yield
+        @types.coroutine
+        def coro():
+            return gen()
+
+        wrapper = coro()
+        self.assertIn('GeneratorWrapper', repr(wrapper))
+        self.assertEqual(repr(wrapper), str(wrapper))
+        self.assertTrue(set(dir(wrapper)).issuperset({
+            '__await__', '__iter__', '__next__', 'cr_code', 'cr_running',
+            'cr_frame', 'gi_code', 'gi_frame', 'gi_running', 'send',
+            'close', 'throw'}))
+
+
+class FunctionTests(CPythonTestCase):
+    def test_function_type_defaults(self):
+        def ex(a, /, b, *, c):
+            return a + b + c
+
+        func = types.FunctionType(
+            ex.__code__, {}, "func", (1, 2), None, {'c': 3},
+        )
+
+        self.assertEqual(func(), 6)
+        self.assertEqual(func.__defaults__, (1, 2))
+        self.assertEqual(func.__kwdefaults__, {'c': 3})
+
+        func = types.FunctionType(
+            ex.__code__, {}, "func", None, None, None,
+        )
+        self.assertEqual(func.__defaults__, None)
+        self.assertEqual(func.__kwdefaults__, None)
+
+    def test_function_type_wrong_defaults(self):
+        def ex(a, /, b, *, c):
+            return a + b + c
+
+        with self.assertRaisesRegex(TypeError, 'arg 4'):
+            types.FunctionType(
+                ex.__code__, {}, "func", 1, None, {'c': 3},
+            )
+        with self.assertRaisesRegex(TypeError, 'arg 6'):
+            types.FunctionType(
+                ex.__code__, {}, "func", None, None, 3,
+            )
+
+
+class SubinterpreterTests(CPythonTestCase):
+
+    @classmethod
+    def setUpClass(cls):
+        global interpreters
+        try:
+            from test.support import interpreters
+        except ModuleNotFoundError:
+            raise unittest.SkipTest('subinterpreters required')
+        import test.support.interpreters.channels
+        super().setUpClass()
+
+    @cpython_only
+    @no_rerun('channels (and queues) might have a refleak; see gh-122199')
+    def test_static_types_inherited_slots(self):
+        rch, sch = interpreters.channels.create()
+
+        slots = []
+        script = ''
+        for cls in iter_builtin_types():
+            for slot, own in iter_slot_wrappers(cls):
+                slots.append((cls, slot, own))
+                script += textwrap.dedent(f"""
+                    text = repr({cls.__name__}.{slot})
+                    sch.send_nowait(({cls.__name__!r}, {slot!r}, text))
+                    """)
+
+        exec(script)
+        all_expected = []
+        for cls, slot, _ in slots:
+            result = rch.recv()
+            assert result == (cls.__name__, slot, result[-1]), (cls, slot, result)
+            all_expected.append(result)
+
+        interp = interpreters.create()
+        interp.exec('from test.support import interpreters')
+        interp.prepare_main(sch=sch)
+        interp.exec(script)
+
+        for i, (cls, slot, _) in enumerate(slots):
+            with self.subTest(cls=cls, slot=slot):
+                expected = all_expected[i]
+                result = rch.recv()
+                self.assertEqual(result, expected)
+
+
+if __name__ == '__main__':
+    run_tests()
diff --git a/test/dynamo/cpython/3_13/typinganndata/_typed_dict_helper.py b/test/dynamo/cpython/3_13/typinganndata/_typed_dict_helper.py
index 9df0ede7d40ee..c1f2efaa5c195 100644
--- a/test/dynamo/cpython/3_13/typinganndata/_typed_dict_helper.py
+++ b/test/dynamo/cpython/3_13/typinganndata/_typed_dict_helper.py
@@ -24,7 +24,7 @@ class Foo(TypedDict):
 T = TypeVar("T")
 
 class FooGeneric(TypedDict, Generic[T]):
-    a: Optional[T]
+    a: T | None
 
 class VeryAnnotated(TypedDict, total=False):
     a: Annotated[Annotated[Annotated[Required[int], "a"], "b"], "c"]
diff --git a/test/dynamo/test_activation_checkpointing.py b/test/dynamo/test_activation_checkpointing.py
index 5d3e52612309f..9a377a52c5b28 100644
--- a/test/dynamo/test_activation_checkpointing.py
+++ b/test/dynamo/test_activation_checkpointing.py
@@ -72,7 +72,6 @@ def inner(*args):
     return inner
 
 
-@torch._dynamo.allow_in_graph
 def _grad(*args, **kwargs):
     return torch.autograd.grad(*args, **kwargs)
 
@@ -92,7 +91,8 @@ def match_rng_op(node, op):
 
     # assert ((freq or freq_ge) and op) or ((freqs or freqs_ge) and ops)
     if op is not None:
-        assert not isinstance(op, list)
+        if isinstance(op, list):
+            raise AssertionError("Expected op to not be a list")
         ops = [op]
     if freq is not None:
         freqs = [freq]
@@ -105,24 +105,30 @@ def match_rng_op(node, op):
                 if match_rng_op(node, op) or node.target == op:
                     actual_count += 1
             err_msg = f"In graph {gm}, expected {op} to have occurred {freq} times in the graph, but got {actual_count}."
-            assert actual_count == freq, err_msg
+            if actual_count != freq:
+                raise AssertionError(err_msg)
     else:
-        assert freqs_ge is not None
+        if freqs_ge is None:
+            raise AssertionError("Expected freqs_ge to not be None")
         for op, freq_ge in zip(ops, freqs_ge):
             actual_count = 0
             for node in gm.graph.nodes:
                 if match_rng_op(node, op) or node.target == op:
                     actual_count += 1
-            assert actual_count >= freq_ge, (
-                f"In graph {gm}, expected {op} to have occurred at least {freq_ge} times in the graph, but got {actual_count}."
-            )
+            if actual_count < freq_ge:
+                raise AssertionError(
+                    f"In graph {gm}, expected {op} to have occurred at least {freq_ge} times in the graph, but got {actual_count}."
+                )
     return gm
 
 
 def collect_fwd_graph_outputs(graph: torch.fx.Graph, *, fwd_outputs: set[str]):
     if not torch._dynamo.compiled_autograd.in_compiled_autograd_region:  # fwd graph
         return_node = list(graph.nodes)[-1]
-        assert return_node.target == "output"
+        if return_node.target != "output":
+            raise AssertionError(
+                f"Expected return_node.target to be 'output', got {return_node.target}"
+            )
         for x in return_node.args[0]:
             fwd_outputs.add(str(x))
 
@@ -169,9 +175,7 @@ def _custom_policy(ctx, func, *args, **kwargs):
     return _custom_policy
 
 
-class ActivationCheckpointingViaTagsTests(
-    torch._dynamo.test_case.TestCaseWithNestedGraphBreaks
-):
+class ActivationCheckpointingViaTagsTests(torch._dynamo.test_case.TestCase):
     def _validate(
         self,
         fn,
@@ -285,7 +289,9 @@ def runtime_wrapper(*runtime_args):
                         _log_export_usage=False,
                     )
                     # NOTE: this is necessary for rng to be added to the exported graph
-                    return torch.compile(gm, fullgraph=False)(*runtime_args)
+                    return torch.compile(
+                        gm, fullgraph=False, backend="aot_eager_decomp_partition"
+                    )(*runtime_args)
 
                 return runtime_wrapper
 
@@ -409,6 +415,112 @@ def partition_fn(joint_gm, *args, **kwargs):
         )
         _ = torch.compile(fn, backend=backend)(x, y)
 
+    @requires_cuda_and_triton
+    def test_ac_tags_through_custom_autograd_function(self, device):
+        class MyMM(torch.autograd.Function):
+            @staticmethod
+            def forward(x, w):
+                return x @ w
+
+            @staticmethod
+            def setup_context(ctx, inputs, output):
+                x, w = inputs
+                ctx.save_for_backward(x, w)
+
+            @staticmethod
+            def backward(ctx, grad):
+                x, w = ctx.saved_tensors
+                return grad @ w.t(), x.t() @ grad
+
+        def gn(x, w):
+            return MyMM.apply(x, w)
+
+        def fn(x, w):
+            return torch.utils.checkpoint.checkpoint(gn, x, w, use_reentrant=False)
+
+        x = torch.randn(4, 4, device=device, requires_grad=True)
+        w = torch.randn(4, 4, device=device, requires_grad=True)
+
+        def partition_fn(joint_gm, *args, **kwargs):
+            fwd_mm_nodes = [
+                node
+                for node in joint_gm.graph.nodes
+                if node.op == "call_function"
+                and node.target == torch.ops.aten.mm.default
+                and node.meta.get("partitioner_tag") == "is_forward"
+            ]
+            self.assertTrue(
+                fwd_mm_nodes, "Expected forward mm nodes in the joint graph"
+            )
+            for node in fwd_mm_nodes:
+                self.assertIn("recompute", node.meta)
+                self.assertIn("ac_graph_id", node.meta)
+            return min_cut_rematerialization_partition(joint_gm, *args, **kwargs)
+
+        backend = aot_autograd(
+            fw_compiler=nop, bw_compiler=nop, partition_fn=partition_fn
+        )
+        out = torch.compile(fn, backend=backend)(x, w)
+        out.sum().backward()
+
+    @requires_cuda_and_triton
+    def test_sac_tags_through_custom_autograd_function(self, device):
+        class MyMM(torch.autograd.Function):
+            @staticmethod
+            def forward(x, w):
+                return x @ w
+
+            @staticmethod
+            def setup_context(ctx, inputs, output):
+                x, w = inputs
+                ctx.save_for_backward(x, w)
+
+            @staticmethod
+            def backward(ctx, grad):
+                x, w = ctx.saved_tensors
+                return grad @ w.t(), x.t() @ grad
+
+        def gn(x, w):
+            return MyMM.apply(x, w)
+
+        context_fn = functools.partial(
+            torch.utils.checkpoint.create_selective_checkpoint_contexts,
+            lambda ctx,
+            op,
+            *args,
+            **kwargs: torch.utils.checkpoint.CheckpointPolicy.PREFER_RECOMPUTE,
+        )
+
+        def fn(x, w):
+            return torch.utils.checkpoint.checkpoint(
+                gn, x, w, use_reentrant=False, context_fn=context_fn
+            )
+
+        x = torch.randn(4, 4, device=device, requires_grad=True)
+        w = torch.randn(4, 4, device=device, requires_grad=True)
+
+        def partition_fn(joint_gm, *args, **kwargs):
+            fwd_mm_nodes = [
+                node
+                for node in joint_gm.graph.nodes
+                if node.op == "call_function"
+                and node.target == torch.ops.aten.mm.default
+                and node.meta.get("partitioner_tag") == "is_forward"
+            ]
+            self.assertTrue(
+                fwd_mm_nodes, "Expected forward mm nodes in the joint graph"
+            )
+            for node in fwd_mm_nodes:
+                self.assertIn("recompute", node.meta)
+                self.assertIn("ac_graph_id", node.meta)
+            return min_cut_rematerialization_partition(joint_gm, *args, **kwargs)
+
+        backend = aot_autograd(
+            fw_compiler=nop, bw_compiler=nop, partition_fn=partition_fn
+        )
+        out = torch.compile(fn, backend=backend)(x, w)
+        out.sum().backward()
+
     @requires_cuda_and_triton
     def test_tangent_placeholders_have_is_backward_tag(self, device):
         """Test that tangent placeholders in the joint graph are tagged with is_backward."""
@@ -1594,7 +1706,6 @@ def fn(x, y):
             self._validate(fn, backend, x, y)
 
     @requires_cuda_and_triton
-    @torch._dynamo.config.patch(inline_inbuilt_nn_modules=True)
     @parametrize(
         "partition_fn",
         [
@@ -1714,7 +1825,7 @@ def gn(*args):
             torch.manual_seed(0)
             ref = gn(*args)
 
-            opt_gn = torch.compile(gn)
+            opt_gn = torch.compile(gn, backend="aot_eager_decomp_partition")
             torch.manual_seed(0)
             res = opt_gn(*args)
             self.assertEqual(ref, res)
@@ -1737,7 +1848,7 @@ def fn(x):
             return torch.utils.checkpoint.checkpoint(mod, x, use_reentrant=True)
 
         x = torch.randn(4, 4).to(device)
-        opt_fn = torch.compile(fn, fullgraph=True)
+        opt_fn = torch.compile(fn, fullgraph=True, backend="aot_eager_decomp_partition")
         with self.assertRaisesRegex(
             torch._dynamo.exc.Unsupported, "User-inserted graph break"
         ):
@@ -1764,7 +1875,7 @@ def fn(x, ys):
         y = torch.randn(4, 4).to(device)
         z = torch.randn(4, 4).to(device)
         ref = fn(x, [y, z])
-        opt_fn = torch.compile(fn, backend="eager", fullgraph=True)
+        opt_fn = torch.compile(fn, backend="aot_eager_decomp_partition", fullgraph=True)
         res = opt_fn(x, [y, z])
         self.assertEqual(ref, res)
 
@@ -1821,7 +1932,7 @@ def debug_compile_fx_inner(graph, example_inputs, *args, **kwargs):
         prefer_cudnn = (
             cudnn_version > 91500 and dprops.major in (9, 10) and dprops.minor in (0, 3)
         )
-        if prefer_cudnn:
+        if prefer_cudnn and torch.version.cuda:
             sdpa_op = torch.ops.aten._scaled_dot_product_cudnn_attention.default
         else:
             sdpa_op = torch.ops.aten._scaled_dot_product_flash_attention.default
@@ -1854,20 +1965,20 @@ def forward(self, x):
         mod = dist_checkpoint_wrapper(MockModule())
         x = torch.randn(4, 4)
         ref = mod(x)
-        opt_mod = torch.compile(mod, backend="eager", fullgraph=True)
+        opt_mod = torch.compile(
+            mod, backend="aot_eager_decomp_partition", fullgraph=True
+        )
         res = opt_mod(x)
         self.assertEqual(ref, res)
 
     @requires_distributed()
     @requires_cuda_and_triton
-    @torch._dynamo.config.patch(inline_inbuilt_nn_modules=True)
     def test_dynamo_does_not_trace_getattr_as_top_frame(self):
-        # inline_inbuilt_nn_modules is a proxy to emulate what FSDP tests do.
         from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import (
             CheckpointWrapper,
         )
 
-        cnt = CompileCounterWithBackend("eager")
+        cnt = CompileCounterWithBackend("aot_eager_decomp_partition")
 
         lin = torch.nn.Linear(1, 1)
         mod = torch.nn.Sequential(lin, lin)
@@ -1936,7 +2047,7 @@ def fn(x):
         self.assertEqual(counter, 2)
         counter = 0
 
-        opt_fn = torch.compile(fn, backend="eager", fullgraph=True)
+        opt_fn = torch.compile(fn, backend="aot_eager_decomp_partition", fullgraph=True)
         opt_fn(x).sum().backward()
         # The mutation is not reapplied in the backward because the flag was on.
         self.assertEqual(counter, 1)
@@ -1963,7 +2074,7 @@ def fn(x):
         x = torch.randn(4, 4, requires_grad=True)
         ref = fn(x)
 
-        opt_fn = torch.compile(fn, backend="eager", fullgraph=True)
+        opt_fn = torch.compile(fn, backend="aot_eager_decomp_partition", fullgraph=True)
         res = opt_fn(x)
         self.assertEqual(ref[0], res[0])
         self.assertEqual(ref[1], res[1])
@@ -2107,7 +2218,11 @@ def checkpointed_forward(inp):
                 x=input_eager.x.detach().clone().requires_grad_(True)
             )
             torch.manual_seed(0)
-            compiled_fn = torch.compile(checkpointed_forward, fullgraph=True)
+            compiled_fn = torch.compile(
+                checkpointed_forward,
+                fullgraph=True,
+                backend="aot_eager_decomp_partition",
+            )
             output_compiled = compiled_fn(input_compiled)
             output_compiled.y.sum().backward()
 
@@ -2149,6 +2264,122 @@ def fn(x, y):
         with torch._dynamo.config.patch(capture_profiler_record_function=True):
             self._validate(fn, backend, x, y)
 
+    def _get_sac_annotations(self, checkpointed_fn, policy_fn, decompositions=None):
+        annotations = []
+
+        def capture_partition(joint_gm, joint_args, **kwargs):
+            for node in joint_gm.graph.nodes:
+                if node.op == "call_function":
+                    recompute = node.meta.get("recompute", None)
+                    if recompute is not None:
+                        annotations.append(
+                            f"{node.name}: {node.target} -> {recompute.name}"
+                        )
+            return min_cut_rematerialization_partition(joint_gm, joint_args, **kwargs)
+
+        backend = aot_autograd(
+            fw_compiler=nop,
+            bw_compiler=nop,
+            partition_fn=capture_partition,
+            decompositions=decompositions,
+        )
+
+        def fn(x):
+            return checkpoint(
+                checkpointed_fn,
+                x,
+                use_reentrant=False,
+                context_fn=functools.partial(
+                    create_selective_checkpoint_contexts, policy_fn
+                ),
+            )
+
+        x = torch.randn(4, requires_grad=True)
+        torch._dynamo.reset()
+        compiled = torch.compile(fn, backend=backend)
+        out = compiled(x)
+        out.sum().backward()
+        return "\n".join(annotations)
+
+    def test_pre_mode_decomp_has_sac_ignored_ops(self):
+        SAVE_OPS = {
+            torch.ops.aten.sin.default,
+            torch.ops.aten.add.Tensor,
+            torch.ops.aten.cos.default,
+        }
+
+        def policy_fn(ctx, func, *args, **kwargs):
+            if func in SAVE_OPS:
+                return CheckpointPolicy.MUST_SAVE
+            return CheckpointPolicy.PREFER_RECOMPUTE
+
+        @torch._dynamo.allow_in_graph
+        def op_with_detach(x):
+            a = x.sin()
+            out = a.detach() + a
+            out = out.cos()
+            return out
+
+        self.assertExpectedInline(
+            self._get_sac_annotations(op_with_detach, policy_fn),
+            """\
+sin: aten.sin.default -> MUST_SAVE
+detach_1: aten.detach.default -> PREFER_RECOMPUTE
+add: aten.add.Tensor -> MUST_SAVE
+cos: aten.cos.default -> MUST_SAVE""",
+        )
+
+    def test_post_mode_decomp(self):
+        from torch._inductor.compile_fx import select_decomp_table
+
+        def policy_fn(ctx, func, *args, **kwargs):
+            if func == torch.ops.aten.silu.default:
+                return CheckpointPolicy.MUST_SAVE
+            return CheckpointPolicy.PREFER_RECOMPUTE
+
+        def fn(x):
+            x = x.sin()
+            x = torch.nn.functional.silu(x)
+            x = x.cos()
+            return x
+
+        self.assertExpectedInline(
+            self._get_sac_annotations(
+                fn, policy_fn, decompositions=select_decomp_table
+            ),
+            """\
+sin: aten.sin.default -> PREFER_RECOMPUTE
+neg: aten.neg.default -> MUST_SAVE
+exp: aten.exp.default -> MUST_SAVE
+add: aten.add.Tensor -> MUST_SAVE
+div: aten.div.Tensor -> MUST_SAVE
+cos: aten.cos.default -> PREFER_RECOMPUTE""",
+        )
+
+    def test_multi_output_op(self):
+        def policy_fn(ctx, func, *args, **kwargs):
+            if func == torch.ops.aten.topk.default:
+                return CheckpointPolicy.MUST_SAVE
+            return CheckpointPolicy.PREFER_RECOMPUTE
+
+        def fn(x):
+            x = x.sin()
+            vals, idxs = torch.topk(x, k=2)
+            out = vals.sum()
+            out = out.cos()
+            return out
+
+        self.assertExpectedInline(
+            self._get_sac_annotations(fn, policy_fn),
+            """\
+sin: aten.sin.default -> PREFER_RECOMPUTE
+topk: aten.topk.default -> MUST_SAVE
+getitem: <built-in function getitem> -> MUST_SAVE
+getitem_1: <built-in function getitem> -> MUST_SAVE
+sum_1: aten.sum.default -> PREFER_RECOMPUTE
+cos: aten.cos.default -> PREFER_RECOMPUTE""",
+        )
+
 
 class RematerializeACNodesPassTests(torch._dynamo.test_case.TestCase):
     """Tests for AC reordering optimization in full graph (forward+backward in one graph)."""
@@ -2170,8 +2401,11 @@ def compiler(gm, example_inputs):
             partition_fn=None,
         )
 
-        with torch._functorch.config.patch(
-            remat_using_tags_for_fwd_loss_bwd_graph=remat_using_tags_for_fwd_loss_bwd_graph
+        with (
+            torch._functorch.config.patch(
+                remat_using_tags_for_fwd_loss_bwd_graph=remat_using_tags_for_fwd_loss_bwd_graph
+            ),
+            torch._dynamo.config.patch(trace_autograd_ops=True),
         ):
             compiled_fn = torch.compile(fn, backend=backend, fullgraph=True)
             result = compiled_fn(*inputs)
@@ -2192,8 +2426,7 @@ def simple_fwd_bwd(x, y):
             )
             loss = z.sum()
 
-            with torch.fx.traceback.annotate({"remat_pass_tag": "is_backward"}):
-                dx, dy = _grad(loss, (x, y))
+            dx, dy = _grad(loss, (x, y))
 
             return dx.detach(), dy.detach()
 
@@ -2225,16 +2458,16 @@ def forward(self, arg0_1, arg1_1):
     expand = torch.ops.aten.expand.default(ones_like, [4, 4]);  ones_like = None
     mm_recomputed = torch.ops.aten.mm.default(arg0_1, arg1_1)
     sigmoid_recomputed = torch.ops.aten.sigmoid.default(mm_recomputed);  mm_recomputed = None
-    detach_recomputed = torch.ops.aten.detach.default(sigmoid_recomputed);  sigmoid_recomputed = None
-    detach_2 = torch.ops.aten.detach.default(detach_recomputed);  detach_recomputed = None
-    sigmoid_backward = torch.ops.aten.sigmoid_backward.default(expand, detach_2);  expand = detach_2 = None
+    detach_2_recomputed = torch.ops.aten.detach.default(sigmoid_recomputed);  sigmoid_recomputed = None
+    detach_4 = torch.ops.aten.detach.default(detach_2_recomputed);  detach_2_recomputed = None
+    sigmoid_backward = torch.ops.aten.sigmoid_backward.default(expand, detach_4);  expand = detach_4 = None
     t = torch.ops.aten.t.default(arg0_1);  arg0_1 = None
     mm_2 = torch.ops.aten.mm.default(t, sigmoid_backward);  t = None
     t_1 = torch.ops.aten.t.default(arg1_1);  arg1_1 = None
     mm_3 = torch.ops.aten.mm.default(sigmoid_backward, t_1);  sigmoid_backward = t_1 = None
-    detach_3 = torch.ops.aten.detach.default(mm_3);  mm_3 = None
-    detach_4 = torch.ops.aten.detach.default(mm_2);  mm_2 = None
-    return (detach_3, detach_4)""",
+    detach_5 = torch.ops.aten.detach.default(mm_3);  mm_3 = None
+    detach_6 = torch.ops.aten.detach.default(mm_2);  mm_2 = None
+    return (detach_5, detach_6)""",
         )
 
     def test_ac_rematerialize_with_rng_ops_raises_error(self):
@@ -2246,18 +2479,17 @@ def fwd_bwd_with_rng(x):
             )
             loss = z.sum()
 
-            with torch.fx.traceback.annotate({"remat_pass_tag": "is_backward"}):
-                dx = _grad(loss, x)[0]
+            dx = _grad(loss, x)[0]
 
             return dx
 
         with self.assertRaisesRegex(
             torch._dynamo.exc.BackendCompilerFailed,
-            "Activation checkpoint rematerializing in `forward-loss-backward` graph does not support RNG ops in checkpointed regions.",
+            "Activation checkpoint rematerialization in `forward-loss-backward` graph does not support RNG ops in recompute regions.",
         ):
             self._compile_and_capture(fwd_bwd_with_rng, True, (x,))
 
-    def test_ac_rematerialize_with_no_annotations_returns_unchanged(self):
+    def test_ac_rematerialize_with_no_annotations(self):
         x = torch.randn(4, 4, requires_grad=True)
 
         def fwd_bwd(x):
@@ -2268,16 +2500,15 @@ def fwd_bwd(x):
             return _grad(loss, x)[0]
 
         result_with, gm_with = self._compile_and_capture(fwd_bwd, True, (x,))
-        # Get the graph without the pass for comparison
         result_without, gm_without = self._compile_and_capture(fwd_bwd, False, (x,))
 
-        # Results should be correct
         self.assertTrue(torch.allclose(result_with, result_without))
 
-        # Both graphs should have the same number of sigmoid ops (no recomputation)
+        # autograd_backward tagging is automatic now, so remat should still work
         sigmoid_with = self.count_op(gm_with, torch.ops.aten.sigmoid.default)
         sigmoid_without = self.count_op(gm_without, torch.ops.aten.sigmoid.default)
-        self.assertEqual(sigmoid_with, sigmoid_without)
+        self.assertEqual(sigmoid_with, 2, "sigmoid should be recomputed in backward")
+        self.assertEqual(sigmoid_without, 1)
 
     def test_ac_rematerialize_with_selective_checkpoint_policy(self):
         x = torch.randn(4, 128, requires_grad=True)
@@ -2303,8 +2534,7 @@ def checkpoint_fn(inp, w, b):
             )
             loss = result.sum()
 
-            with torch.fx.traceback.annotate({"remat_pass_tag": "is_backward"}):
-                dx, dw, db = _grad(loss, (x, w1, b1))
+            dx, dw, db = _grad(loss, (x, w1, b1))
             return dx, dw, db
 
         result_with, gm_with = self._compile_and_capture(
@@ -2358,8 +2588,9 @@ def custom_compiler(gm, example_inputs):
             partition_fn=None,
         )
 
-        compiled_fn = torch.compile(fn, backend=backend, fullgraph=True)
-        result = compiled_fn(*inputs)
+        with torch._dynamo.config.patch(trace_autograd_ops=True):
+            compiled_fn = torch.compile(fn, backend=backend, fullgraph=True)
+            result = compiled_fn(*inputs)
 
         return result, captured_gm_before, captured_gm_after
 
@@ -2379,8 +2610,7 @@ def checkpoint_fn(a):
             )
             loss = z.sum()
 
-            with torch.fx.traceback.annotate({"remat_pass_tag": "is_backward"}):
-                dx = _grad(loss, x)[0]
+            dx = _grad(loss, x)[0]
 
             return dx.detach()
 
@@ -2405,11 +2635,11 @@ def forward(self, arg0_1):
     ones_like = torch.ops.aten.ones_like.default(sum_1, pin_memory = False, memory_format = torch.preserve_format);  sum_1 = None
     expand = torch.ops.aten.expand.default(ones_like, [4, 4]);  ones_like = None
     sigmoid_recomputed = torch.ops.aten.sigmoid.default(arg0_1);  arg0_1 = None
-    detach_recomputed = torch.ops.aten.detach.default(sigmoid_recomputed);  sigmoid_recomputed = None
-    detach_2 = torch.ops.aten.detach.default(detach_recomputed);  detach_recomputed = None
-    sigmoid_backward = torch.ops.aten.sigmoid_backward.default(expand, detach_2);  expand = detach_2 = None
-    detach_3 = torch.ops.aten.detach.default(sigmoid_backward);  sigmoid_backward = None
-    return (detach_3,)""",
+    detach_3_recomputed = torch.ops.aten.detach.default(sigmoid_recomputed);  sigmoid_recomputed = None
+    detach_5 = torch.ops.aten.detach.default(detach_3_recomputed);  detach_3_recomputed = None
+    sigmoid_backward = torch.ops.aten.sigmoid_backward.default(expand, detach_5);  expand = detach_5 = None
+    detach_6 = torch.ops.aten.detach.default(sigmoid_backward);  sigmoid_backward = None
+    return (detach_6,)""",
         )
 
     def test_joint_graph_passes_permute_optimization(self):
@@ -2428,8 +2658,7 @@ def checkpoint_fn(a):
             )
             loss = z.sum()
 
-            with torch.fx.traceback.annotate({"remat_pass_tag": "is_backward"}):
-                dx = _grad(loss, x)[0]
+            dx = _grad(loss, x)[0]
 
             return dx.detach()
 
@@ -2454,11 +2683,138 @@ def forward(self, arg0_1):
     ones_like = torch.ops.aten.ones_like.default(sum_1, pin_memory = False, memory_format = torch.preserve_format);  sum_1 = None
     expand = torch.ops.aten.expand.default(ones_like, [4, 4]);  ones_like = None
     sigmoid_recomputed = torch.ops.aten.sigmoid.default(arg0_1);  arg0_1 = None
-    detach_recomputed = torch.ops.aten.detach.default(sigmoid_recomputed);  sigmoid_recomputed = None
-    detach_2 = torch.ops.aten.detach.default(detach_recomputed);  detach_recomputed = None
-    sigmoid_backward = torch.ops.aten.sigmoid_backward.default(expand, detach_2);  expand = detach_2 = None
-    detach_3 = torch.ops.aten.detach.default(sigmoid_backward);  sigmoid_backward = None
-    return (detach_3,)""",
+    detach_3_recomputed = torch.ops.aten.detach.default(sigmoid_recomputed);  sigmoid_recomputed = None
+    detach_5 = torch.ops.aten.detach.default(detach_3_recomputed);  detach_3_recomputed = None
+    sigmoid_backward = torch.ops.aten.sigmoid_backward.default(expand, detach_5);  expand = detach_5 = None
+    detach_6 = torch.ops.aten.detach.default(sigmoid_backward);  sigmoid_backward = None
+    return (detach_6,)""",
+        )
+
+    @torch._dynamo.config.patch(skip_fwd_side_effects_in_bwd_under_checkpoint=True)
+    def test_attr_compile_submodules_in_checkpoint_wrapper(self):
+        """Compiling submodules inside a checkpointed block should not hit the
+        recompile limit due to WeakKeyDictionary guards in the pack_hook."""
+        from torch.utils.checkpoint import checkpoint
+
+        class Block(nn.Module):
+            def __init__(self, dim):
+                super().__init__()
+                self.norm1 = nn.RMSNorm(dim)
+                self.linear1 = nn.Linear(dim, dim, bias=False)
+                self.norm2 = nn.RMSNorm(dim)
+                self.linear2 = nn.Linear(dim, dim, bias=False)
+                self.norm3 = nn.RMSNorm(dim)
+                self.linear3 = nn.Linear(dim, dim, bias=False)
+
+            def forward(self, x):
+                x = x + self.linear1(self.norm1(x))
+                x = x + self.linear2(self.norm2(x))
+                x = x + self.linear3(self.norm3(x))
+                return x
+
+        class CheckpointedBlock(nn.Module):
+            def __init__(self, block):
+                super().__init__()
+                self.block = block
+
+            def forward(self, x):
+                return checkpoint(self.block, x, use_reentrant=False)
+
+        dim = 32
+        block = Block(dim)
+
+        x_ref = torch.randn(4, dim, requires_grad=True)
+        ref = block(x_ref)
+        ref.sum().backward()
+
+        block_cp = Block(dim)
+        block_cp.load_state_dict(block.state_dict())
+        wrapped = CheckpointedBlock(block_cp)
+
+        for _, submod in wrapped.block.named_children():
+            submod.compile(backend="aot_eager")
+
+        with torch._dynamo.config.patch(recompile_limit=2):
+            x_test = x_ref.detach().clone().requires_grad_(True)
+            result = wrapped(x_test)
+            result.sum().backward()
+
+        self.assertEqual(ref, result)
+        self.assertEqual(x_ref.grad, x_test.grad)
+
+    @unittest.skipIf(not torch.cuda.is_available(), "CUDA not available")
+    def test_multiple_user_phase_annotations_errors(self):
+        x = torch.randn(4, 4, requires_grad=True)
+        w = torch.randn(4, 4, requires_grad=True)
+
+        def fn(x, w):
+            z = torch.utils.checkpoint.checkpoint(
+                lambda a, b: torch.sin(a @ b), x, w, use_reentrant=False
+            )
+            loss = z.sum()
+            with torch.fx.traceback.annotate({"phase": "backward"}):
+                dx, dw = _grad(loss, (x, w))
+            # Non-backward computation between two backward annotations
+            out = dx + dw
+            with torch.fx.traceback.annotate({"phase": "backward"}):
+                out = out * 2
+            return out.detach()
+
+        with self.assertRaisesRegex(RuntimeError, "backward regions annotated"):
+            self._compile_and_capture(fn, True, (x, w))
+
+    @unittest.skipIf(not torch.cuda.is_available(), "CUDA not available")
+    def test_user_phase_annotation_with_extra_autograd_grad(self):
+        """Only the user-annotated backward region gets rematerialization."""
+        x = torch.randn(4, 4, requires_grad=True)
+        w1 = torch.randn(4, 4, requires_grad=True)
+        w2 = torch.randn(4, 4, requires_grad=True)
+
+        def fn(x, w1, w2):
+            z1 = torch.utils.checkpoint.checkpoint(
+                lambda a, b: torch.sin(a @ b), x, w1, use_reentrant=False
+            )
+            z2 = torch.utils.checkpoint.checkpoint(
+                lambda a, b: torch.sigmoid(a @ b), x, w2, use_reentrant=False
+            )
+            loss1 = z1.sum()
+            loss2 = z2.sum()
+            # Only the first backward is annotated
+            with torch.fx.traceback.annotate({"phase": "backward"}):
+                dx1 = _grad(loss1, (x,))
+            # Second backward NOT annotated — should not get remat
+            dx2 = _grad(loss2, (x,))
+            return (dx1[0] + dx2[0]).detach()
+
+        _, gm_with = self._compile_and_capture(fn, True, (x, w1, w2))
+
+        self.assertExpectedInline(
+            gm_with.code.strip(),
+            """\
+def forward(self, arg0_1, arg1_1, arg2_1):
+    mm = torch.ops.aten.mm.default(arg0_1, arg1_1)
+    sin = torch.ops.aten.sin.default(mm);  mm = None
+    mm_1 = torch.ops.aten.mm.default(arg0_1, arg2_1)
+    sigmoid = torch.ops.aten.sigmoid.default(mm_1);  mm_1 = None
+    detach_4 = torch.ops.aten.detach.default(sigmoid)
+    sum_1 = torch.ops.aten.sum.default(sin);  sin = None
+    sum_2 = torch.ops.aten.sum.default(sigmoid);  sigmoid = None
+    ones_like = torch.ops.aten.ones_like.default(sum_1, pin_memory = False, memory_format = torch.preserve_format);  sum_1 = None
+    expand = torch.ops.aten.expand.default(ones_like, [4, 4]);  ones_like = None
+    mm_recomputed = torch.ops.aten.mm.default(arg0_1, arg1_1);  arg0_1 = None
+    cos = torch.ops.aten.cos.default(mm_recomputed);  mm_recomputed = None
+    mul = torch.ops.aten.mul.Tensor(expand, cos);  expand = cos = None
+    t = torch.ops.aten.t.default(arg1_1);  arg1_1 = None
+    mm_3 = torch.ops.aten.mm.default(mul, t);  mul = t = None
+    ones_like_1 = torch.ops.aten.ones_like.default(sum_2, pin_memory = False, memory_format = torch.preserve_format);  sum_2 = None
+    expand_1 = torch.ops.aten.expand.default(ones_like_1, [4, 4]);  ones_like_1 = None
+    detach_5 = torch.ops.aten.detach.default(detach_4);  detach_4 = None
+    sigmoid_backward = torch.ops.aten.sigmoid_backward.default(expand_1, detach_5);  expand_1 = detach_5 = None
+    t_1 = torch.ops.aten.t.default(arg2_1);  arg2_1 = None
+    mm_4 = torch.ops.aten.mm.default(sigmoid_backward, t_1);  sigmoid_backward = t_1 = None
+    add = torch.ops.aten.add.Tensor(mm_3, mm_4);  mm_3 = mm_4 = None
+    detach_6 = torch.ops.aten.detach.default(add);  add = None
+    return (detach_6,)""",
         )
 
 
@@ -2467,6 +2823,367 @@ def forward(self, arg0_1):
     ActivationCheckpointingViaTagsTests, globals(), only_for=devices
 )
 
+
+class ActivationCheckpointingNonStrictTracerTests(torch._dynamo.test_case.TestCase):
+    """Tests for non-strict tracing flag interaction with checkpoint."""
+
+    @staticmethod
+    def _count_backward_regions(gm):
+        regions = 0
+        in_backward = False
+        for node in gm.graph.nodes:
+            is_backward = bool(node.meta.get("autograd_backward", False))
+            if is_backward and not in_backward:
+                regions += 1
+            in_backward = is_backward
+        return regions
+
+    def test_backward_nodes_have_seq_nr_under_non_strict(self):
+        class Model(nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.w = nn.Parameter(torch.randn(4, 4))
+
+            def forward(self, x):
+                return checkpoint(
+                    lambda x: torch.sin(x @ self.w), x, use_reentrant=False
+                )
+
+        gm = self._trace_train_step(Model(), torch.randn(2, 4))
+        forward_seq_nrs = {
+            node.meta["seq_nr"]
+            for node in gm.graph.nodes
+            if node.op == "call_function"
+            and not node.meta.get("autograd_backward", False)
+            and "seq_nr" in node.meta
+        }
+        backward_seq_nrs = {
+            node.meta["seq_nr"]
+            for node in gm.graph.nodes
+            if node.op == "call_function"
+            and node.meta.get("autograd_backward", False)
+            and "seq_nr" in node.meta
+        }
+
+        self.assertTrue(forward_seq_nrs)
+        self.assertTrue(backward_seq_nrs)
+        self.assertSetEqual(backward_seq_nrs, forward_seq_nrs)
+
+    def test_patch_autograd_grad_requires_non_strict_tracing(self):
+        x = torch.randn(2, 4, requires_grad=True)
+        loss = torch.sin(x).sum()
+
+        with torch.compiler._patch_autograd_grad():
+            with self.assertRaisesRegex(
+                AssertionError,
+                "_patch_autograd_grad\\(\\) must be used under "
+                "_non_strict_tracing_context\\(\\)",
+            ):
+                torch.autograd.grad(loss, (x,))
+
+    def test_patch_autograd_grad_does_not_leak_backward_tag(self):
+        from torch.fx.experimental.proxy_tensor import make_fx
+        from torch.fx.traceback import preserve_node_meta
+
+        x = torch.randn(2, 4, requires_grad=True)
+
+        def fn(x):
+            with torch.fx.traceback.annotate({"ac_region_id": 0}):
+                y = torch.sin(x)
+                torch.autograd.grad(y.sum(), (x,))
+                return torch.neg(y)
+
+        with (
+            torch.compiler._non_strict_tracing_context(),
+            torch.compiler._patch_autograd_grad(),
+            preserve_node_meta(),
+        ):
+            gm = make_fx(fn)(x)
+
+        backward_nodes = [
+            node for node in gm.graph.nodes if node.meta.get("autograd_backward", False)
+        ]
+        self.assertTrue(backward_nodes)
+
+        neg_nodes = gm.graph.find_nodes(
+            op="call_function", target=torch.ops.aten.neg.default
+        )
+        self.assertEqual(len(neg_nodes), 1)
+        self.assertNotIn("autograd_backward", neg_nodes[0].meta)
+        self.assertEqual(neg_nodes[0].meta.get("custom", {}), {"ac_region_id": 0})
+
+    def test_patch_autograd_grad_mlp_has_single_contiguous_backward_region(self):
+        class Block(nn.Module):
+            def __init__(self, dim):
+                super().__init__()
+                self.fc1 = nn.Linear(dim, dim * 2)
+                self.fc2 = nn.Linear(dim * 2, dim)
+
+            def forward(self, x):
+                return x + self.fc2(torch.relu(self.fc1(x)))
+
+        class Model(nn.Module):
+            def __init__(self, dim=32, depth=6):
+                super().__init__()
+                self.blocks = nn.ModuleList([Block(dim) for _ in range(depth)])
+
+            def forward(self, x):
+                for block in self.blocks:
+                    x = block(x)
+                return x
+
+        gm = self._trace_train_step(Model(), torch.randn(4, 16, 32))
+        self.assertEqual(self._count_backward_regions(gm), 1)
+
+    def _trace_train_step(self, mod, x):
+        import torch.utils._pytree as pytree
+        from torch.fx.experimental.proxy_tensor import make_fx
+        from torch.fx.traceback import preserve_node_meta
+        from torch.nn.utils import stateless
+
+        params = dict(mod.named_parameters())
+        flat_params, params_spec = pytree.tree_flatten(params)
+        params_len = len(flat_params)
+
+        def train_step(*all_args):
+            p = pytree.tree_unflatten(list(all_args[:params_len]), params_spec)
+            with stateless._reparametrize_module(mod, p):
+                loss = mod(all_args[params_len]).sum()
+                return torch.autograd.grad(loss, all_args[:params_len])
+
+        full_args = (*flat_params, x)
+
+        with (
+            torch.compiler._non_strict_tracing_context(),
+            torch.compiler._patch_autograd_grad(),
+            preserve_node_meta(),
+        ):
+            return make_fx(train_step)(*full_args)
+
+    def test_checkpoint_traces_through_eager_ac_under_non_strict(self):
+        class Model(nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.w = nn.Parameter(torch.randn(4, 4))
+
+            def forward(self, x):
+                return checkpoint(
+                    lambda x: torch.sin(x @ self.w), x, use_reentrant=False
+                )
+
+        gm = self._trace_train_step(Model(), torch.randn(2, 4))
+
+        # Everything is recomputed in backward: mm_1 is the recomputed mm.
+        self.assertExpectedInline(
+            gm.code.strip(),
+            """\
+def forward(self, arg0_1, arg1_1):
+    mm = torch.ops.aten.mm.default(arg1_1, arg0_1)
+    sin = torch.ops.aten.sin.default(mm);  mm = None
+    sum_1 = torch.ops.aten.sum.default(sin);  sin = None
+    ones_like = torch.ops.aten.ones_like.default(sum_1, pin_memory = False, memory_format = torch.preserve_format);  sum_1 = None
+    expand = torch.ops.aten.expand.default(ones_like, [2, 4]);  ones_like = None
+    mm_1 = torch.ops.aten.mm.default(arg1_1, arg0_1);  arg0_1 = None
+    detach = torch.ops.aten.detach.default(mm_1);  mm_1 = None
+    detach_1 = torch.ops.aten.detach.default(detach);  detach = None
+    cos = torch.ops.aten.cos.default(detach_1);  detach_1 = None
+    mul = torch.ops.aten.mul.Tensor(expand, cos);  expand = cos = None
+    t = torch.ops.aten.t.default(arg1_1);  arg1_1 = None
+    mm_2 = torch.ops.aten.mm.default(t, mul);  t = mul = None
+    return (mm_2,)""",
+        )
+
+    def test_sac_traces_through_eager_ac_under_non_strict(self):
+        def policy_fn(ctx, func, *args, **kwargs):
+            if func == torch.ops.aten.mm.default:
+                return CheckpointPolicy.PREFER_RECOMPUTE
+            return CheckpointPolicy.MUST_SAVE
+
+        def context_fn():
+            return create_selective_checkpoint_contexts(policy_fn)
+
+        class Model(nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.w = nn.Parameter(torch.randn(4, 4))
+
+            def forward(self, x):
+                def fn(x):
+                    return torch.sin(x @ self.w)
+
+                return checkpoint(fn, x, use_reentrant=False, context_fn=context_fn)
+
+        gm = self._trace_train_step(Model(), torch.randn(2, 4))
+
+        # mm is PREFER_RECOMPUTE so it gets recomputed in backward (mm_1).
+        # sin is MUST_SAVE so its output is saved via detach.
+        self.assertExpectedInline(
+            gm.code.strip(),
+            """\
+def forward(self, arg0_1, arg1_1):
+    mm = torch.ops.aten.mm.default(arg1_1, arg0_1)
+    sin = torch.ops.aten.sin.default(mm);  mm = None
+    detach = torch.ops.aten.detach.default(sin);  detach = None
+    sum_1 = torch.ops.aten.sum.default(sin);  sin = None
+    ones_like = torch.ops.aten.ones_like.default(sum_1, pin_memory = False, memory_format = torch.preserve_format);  sum_1 = None
+    expand = torch.ops.aten.expand.default(ones_like, [2, 4]);  ones_like = None
+    mm_1 = torch.ops.aten.mm.default(arg1_1, arg0_1);  arg0_1 = None
+    detach_1 = torch.ops.aten.detach.default(mm_1);  mm_1 = None
+    detach_2 = torch.ops.aten.detach.default(detach_1);  detach_1 = None
+    cos = torch.ops.aten.cos.default(detach_2);  detach_2 = None
+    mul = torch.ops.aten.mul.Tensor(expand, cos);  expand = cos = None
+    t = torch.ops.aten.t.default(arg1_1);  arg1_1 = None
+    mm_2 = torch.ops.aten.mm.default(t, mul);  t = mul = None
+    return (mm_2,)""",
+        )
+
+    def test_checkpoint_with_rng_op_under_non_strict(self):
+        from torch.fx.experimental.proxy_tensor import make_fx
+
+        class Model(nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.w = nn.Parameter(torch.randn(4, 4))
+
+            def forward(self, x):
+                def fn(x):
+                    return torch.nn.functional.dropout(torch.sin(x @ self.w), p=0.5)
+
+                return checkpoint(fn, x, use_reentrant=False, preserve_rng_state=True)
+
+        mod = Model()
+        x = torch.randn(2, 4)
+
+        import torch.utils._pytree as pytree
+        from torch.nn.utils import stateless
+
+        params = dict(mod.named_parameters())
+        flat_params, params_spec = pytree.tree_flatten(params)
+        params_len = len(flat_params)
+
+        def train_step(*all_args):
+            p = pytree.tree_unflatten(list(all_args[:params_len]), params_spec)
+            with stateless._reparametrize_module(mod, p):
+                loss = mod(all_args[params_len]).sum()
+                return torch.autograd.grad(loss, all_args[:params_len])
+
+        full_args = (*flat_params, x)
+
+        torch.manual_seed(42)
+        with torch.compiler._non_strict_tracing_context():
+            gm = make_fx(train_step)(*full_args)
+
+        # The traced graph has two independent bernoulli_ calls (forward and
+        # recomputed). At replay time these produce different dropout masks
+        # because get/set_rng_state are silently dropped by make_fx, so the
+        # traced graph is NOT bitwise equivalent to eager. This is a known
+        # limitation of tracing through eager checkpoint.
+        self.assertExpectedInline(
+            gm.code.strip(),
+            """\
+def forward(self, arg0_1, arg1_1):
+    mm = torch.ops.aten.mm.default(arg1_1, arg0_1)
+    sin = torch.ops.aten.sin.default(mm);  mm = None
+    empty_like = torch.ops.aten.empty_like.default(sin)
+    bernoulli_ = torch.ops.aten.bernoulli_.float(empty_like);  empty_like = None
+    div_ = torch.ops.aten.div_.Scalar(bernoulli_, 0.5);  bernoulli_ = None
+    mul = torch.ops.aten.mul.Tensor(sin, div_);  sin = div_ = None
+    sum_1 = torch.ops.aten.sum.default(mul);  mul = None
+    ones_like = torch.ops.aten.ones_like.default(sum_1, pin_memory = False, memory_format = torch.preserve_format);  sum_1 = None
+    expand = torch.ops.aten.expand.default(ones_like, [2, 4]);  ones_like = None
+    mm_1 = torch.ops.aten.mm.default(arg1_1, arg0_1);  arg0_1 = None
+    detach = torch.ops.aten.detach.default(mm_1)
+    sin_1 = torch.ops.aten.sin.default(mm_1);  mm_1 = None
+    empty_like_1 = torch.ops.aten.empty_like.default(sin_1);  sin_1 = None
+    bernoulli__1 = torch.ops.aten.bernoulli_.float(empty_like_1);  empty_like_1 = None
+    div__1 = torch.ops.aten.div_.Scalar(bernoulli__1, 0.5);  bernoulli__1 = None
+    mul_1 = torch.ops.aten.mul.Tensor(expand, div__1);  expand = div__1 = None
+    detach_1 = torch.ops.aten.detach.default(detach);  detach = None
+    cos = torch.ops.aten.cos.default(detach_1);  detach_1 = None
+    mul_2 = torch.ops.aten.mul.Tensor(mul_1, cos);  mul_1 = cos = None
+    t = torch.ops.aten.t.default(arg1_1);  arg1_1 = None
+    mm_2 = torch.ops.aten.mm.default(t, mul_2);  t = mul_2 = None
+    return (mm_2,)""",
+        )
+
+
+class ActivationCheckpointingNestedCompileTests(torch._dynamo.test_case.TestCase):
+    @requires_cuda_and_triton
+    def test_checkpoint_recompute_preserves_nested_fx_trace_policy(self):
+        from torch._guards import tracing, TracingContext
+        from torch._subclasses import FakeTensorMode
+        from torch.fx.experimental.proxy_tensor import make_fx
+        from torch.fx.traceback import preserve_node_meta
+
+        compiled_f = torch.compile(lambda x: x.sin().cos(), fullgraph=True)
+
+        @contextlib.contextmanager
+        def skip_nested_compile():
+            prev = torch._dynamo.config.error_on_nested_fx_trace
+            torch._dynamo.config.error_on_nested_fx_trace = False
+            try:
+                yield
+            finally:
+                torch._dynamo.config.error_on_nested_fx_trace = prev
+
+        class M(torch.nn.Module):
+            def forward(self, x):
+                return checkpoint(self.block, x, use_reentrant=False)
+
+            def block(self, x):
+                return compiled_f(x)
+
+        m = M().cuda()
+        x = torch.randn(8, device="cuda", requires_grad=True)
+
+        def fn(x):
+            y = m(x).sum()
+            (gx,) = torch.autograd.grad(y, (x,))
+            return y.detach(), gx
+
+        fake_mode = FakeTensorMode(
+            allow_non_fake_inputs=True,
+            shape_env=torch.fx.experimental.symbolic_shapes.ShapeEnv(),
+        )
+        fx_x = fake_mode.from_tensor(x, static_shapes=True)
+        if x.requires_grad and not fx_x.requires_grad:
+            fx_x.requires_grad_(True)
+
+        ctx = TracingContext(fake_mode)
+
+        with (
+            fake_mode,
+            tracing(ctx),
+            preserve_node_meta(),
+            skip_nested_compile(),
+            torch.compiler._non_strict_tracing_context(),
+        ):
+            gm = make_fx(fn)(fx_x)
+
+        self.assertExpectedInline(
+            gm.code.strip(),
+            """\
+def forward(self, x_1):
+    sin = torch.ops.aten.sin.default(x_1)
+    cos = torch.ops.aten.cos.default(sin);  sin = None
+    sum_1 = torch.ops.aten.sum.default(cos);  cos = None
+    ones_like = torch.ops.aten.ones_like.default(sum_1, pin_memory = False, memory_format = torch.preserve_format)
+    expand = torch.ops.aten.expand.default(ones_like, [8]);  ones_like = None
+    detach = torch.ops.aten.detach.default(x_1)
+    sin_1 = torch.ops.aten.sin.default(x_1);  x_1 = None
+    detach_1 = torch.ops.aten.detach.default(sin_1);  sin_1 = None
+    detach_2 = torch.ops.aten.detach.default(detach_1);  detach_1 = None
+    sin_2 = torch.ops.aten.sin.default(detach_2);  detach_2 = None
+    neg = torch.ops.aten.neg.default(sin_2);  sin_2 = None
+    mul = torch.ops.aten.mul.Tensor(expand, neg);  expand = neg = None
+    detach_3 = torch.ops.aten.detach.default(detach);  detach = None
+    cos_1 = torch.ops.aten.cos.default(detach_3);  detach_3 = None
+    mul_1 = torch.ops.aten.mul.Tensor(mul, cos_1);  mul = cos_1 = None
+    detach_4 = torch.ops.aten.detach.default(sum_1);  sum_1 = None
+    return (detach_4, mul_1)""",
+        )
+
+
 if __name__ == "__main__":
     from torch._dynamo.test_case import run_tests
 
diff --git a/test/dynamo/test_activation_offloading.py b/test/dynamo/test_activation_offloading.py
index 0630312d30678..7d9694318231d 100644
--- a/test/dynamo/test_activation_offloading.py
+++ b/test/dynamo/test_activation_offloading.py
@@ -137,7 +137,8 @@ def run_compiled():
             .run(bw_code)
         )
 
-    def test_partitioner_offload_sep_stream(self):
+    def test_partitioner_offload_ao_ops(self):
+        """Test that async offload uses ao.offload + ao.wait_tensor with keepalive."""
         reset_user_object_tracking()
         torch._dynamo.reset()
         with torch._functorch.config.patch(
@@ -147,47 +148,44 @@ def test_partitioner_offload_sep_stream(self):
         ):
             fw_graph, bw_graph = get_fw_bw_graph(self.fn, [self.x])
 
-        self.assertExpectedInline(
-            fw_graph.code.strip(),
-            """\
-def forward(self, primals_1, primals_2, primals_3, primals_4, primals_5, primals_6):
-    add = torch.ops.aten.add.Tensor(primals_1, primals_2);  primals_1 = primals_2 = None
-    sin = torch.ops.aten.sin.default(add)
-    add_1 = torch.ops.aten.add.Tensor(primals_3, primals_4);  primals_3 = primals_4 = None
-    sin_1 = torch.ops.aten.sin.default(add_1)
-    add_2 = torch.ops.aten.add.Tensor(sin, sin_1);  sin = sin_1 = None
-    add_3 = torch.ops.aten.add.Tensor(primals_5, primals_6);  primals_5 = primals_6 = None
-    sin_2 = torch.ops.aten.sin.default(add_3)
-    add_4 = torch.ops.aten.add.Tensor(add_2, sin_2);  add_2 = sin_2 = None
-    cos = torch.ops.aten.cos.default(add_3);  add_3 = None
-    cos_1 = torch.ops.aten.cos.default(add_1);  add_1 = None
-    record_event_default = torch.ops.streams.record_event.default(2, 0);  record_event_default = None
-    stream_in_cpu_offload_cos_1 = torch.ops.streams.fork.default(0, 1);  stream_in_cpu_offload_cos_1 = None
-    wait_event_default = torch.ops.streams.wait_event.default(2, 1);  wait_event_default = None
-    record_stream_cos_1 = torch.ops.streams.record_stream.default(cos_1, 1);  record_stream_cos_1 = None
-    cpu_offload_cos_1 = torch.ops.prims.device_put.default(cos_1, device(type='cpu'), non_blocking = True);  cos_1 = None
-    record_event_default_1 = torch.ops.streams.record_event.default(3, 1);  record_event_default_1 = None
-    stream_out_cpu_offload_cos_1 = torch.ops.streams.join.default(1, 0);  stream_out_cpu_offload_cos_1 = None
-    wait_event_default_1 = torch.ops.streams.wait_event.default(3, 0);  wait_event_default_1 = None
-    cos_2 = torch.ops.aten.cos.default(add);  add = None
-    return (add_4, cos, cpu_offload_cos_1, cos_2)""",
+        fw_code = fw_graph.code
+        # Forward: ao.offload produces async CPU tensor, ao.wait_tensor synchronizes
+        (
+            FileCheck()
+            .check("ao.offload.default")
+            .check("ao.wait_tensor.default")
+            .run(fw_code)
         )
-
-        self.assertExpectedInline(
-            bw_graph.code.strip(),
-            """\
-def forward(self, cos, cpu_offload_cos_1, cos_2, tangents_1):
-    mul = torch.ops.aten.mul.Tensor(tangents_1, cos);  cos = None
-    stream_in_gpu_reload_cos_1 = torch.ops.streams.fork.default(4, 5);  stream_in_gpu_reload_cos_1 = None
-    wait_stream_default = torch.ops.streams.wait_stream.default(5, 4);  wait_stream_default = None
-    gpu_reload_cos_1 = torch.ops.prims.device_put.default(cpu_offload_cos_1, device(type='GPU_TYPE', index=0), non_blocking = True);  cpu_offload_cos_1 = None
-    record_event_default = torch.ops.streams.record_event.default(6, 5);  record_event_default = None
-    stream_out_gpu_reload_cos_1 = torch.ops.streams.join.default(5, 4);  stream_out_gpu_reload_cos_1 = None
-    wait_event_default = torch.ops.streams.wait_event.default(6, 4);  wait_event_default = None
-    mul_1 = torch.ops.aten.mul.Tensor(tangents_1, gpu_reload_cos_1);  gpu_reload_cos_1 = None
-    mul_2 = torch.ops.aten.mul.Tensor(tangents_1, cos_2);  tangents_1 = cos_2 = None
-    return (mul_2, mul_2, mul_1, mul_1, mul, mul)""".replace("GPU_TYPE", GPU_TYPE),
+        # No old stream ops should be present
+        self.assertNotIn("streams.fork", fw_code)
+        self.assertNotIn("streams.join", fw_code)
+        self.assertNotIn("streams.record_stream", fw_code)
+
+        # Verify keepalive: ao.wait_tensor for offload should reference the GPU tensor
+        # (cos_1) as its 2nd arg to extend its lifetime
+        for node in fw_graph.graph.nodes:
+            if (
+                node.op == "call_function"
+                and node.target == torch.ops.ao.wait_tensor.default
+                and isinstance(node.args[0], torch.fx.Node)
+                and node.args[0].target == torch.ops.ao.offload.default
+            ):
+                self.assertEqual(len(node.args), 2)
+                offload_node = node.args[0]
+                keepalive_node = node.args[1]
+                # keepalive should be the same GPU tensor that was offloaded
+                self.assertIs(keepalive_node, offload_node.args[0])
+
+        bw_code = bw_graph.code
+        # Backward: ao.reload produces async GPU tensor, ao.wait_tensor synchronizes
+        (
+            FileCheck()
+            .check("ao.reload.default")
+            .check("ao.wait_tensor.default")
+            .run(bw_code)
         )
+        self.assertNotIn("streams.fork", bw_code)
+        self.assertNotIn("streams.join", bw_code)
 
     def test_partitioner_offload_sep_stream_accuracy(self):
         # Run without compilation to get reference gradients
@@ -219,59 +217,64 @@ def test_partitioner_offload_sep_stream_accuracy(self):
                 atol=1e-5,
             )
 
-    def test_partitioner_offload_sep_stream_reorder(self):
+    def test_partitioner_offload_ao_ops_sink_wait(self):
+        """Test that sink_wait moves ao.wait_tensor for offload to end of forward graph."""
         reset_user_object_tracking()
         torch._dynamo.reset()
         with torch._functorch.config.patch(
             enable_activation_offloading=True,
             activation_offload_separate_stream=True,
             activation_offload_sink_wait=True,
+            joint_custom_pass=self.joint_custom_pass,
+        ):
+            fw_graph, _ = get_fw_bw_graph(self.fn, [self.x])
+
+        # The ao.wait_tensor for offload should be sunk to just before the output node
+        nodes = list(fw_graph.graph.nodes)
+        output_node = next(n for n in nodes if n.op == "output")
+        output_idx = nodes.index(output_node)
+
+        for node in nodes:
+            if (
+                node.op == "call_function"
+                and node.target == torch.ops.ao.wait_tensor.default
+                and isinstance(node.args[0], torch.fx.Node)
+                and node.args[0].target == torch.ops.ao.offload.default
+            ):
+                wait_idx = nodes.index(node)
+                # ao.wait_tensor should be immediately before output
+                self.assertEqual(wait_idx, output_idx - 1)
+
+    def test_partitioner_offload_ao_ops_prefetch(self):
+        """Test that prefetch moves ao.reload earlier in backward graph."""
+        reset_user_object_tracking()
+        torch._dynamo.reset()
+        with torch._functorch.config.patch(
+            enable_activation_offloading=True,
+            activation_offload_separate_stream=True,
             activation_reload_prefetch=True,
             joint_custom_pass=self.joint_custom_pass,
         ):
-            fw_graph, bw_graph = get_fw_bw_graph(self.fn, [self.x])
+            _, bw_graph = get_fw_bw_graph(self.fn, [self.x])
 
-        self.assertExpectedInline(
-            fw_graph.code.strip(),
-            """\
-def forward(self, primals_1, primals_2, primals_3, primals_4, primals_5, primals_6):
-    add = torch.ops.aten.add.Tensor(primals_1, primals_2);  primals_1 = primals_2 = None
-    sin = torch.ops.aten.sin.default(add)
-    add_1 = torch.ops.aten.add.Tensor(primals_3, primals_4);  primals_3 = primals_4 = None
-    sin_1 = torch.ops.aten.sin.default(add_1)
-    add_2 = torch.ops.aten.add.Tensor(sin, sin_1);  sin = sin_1 = None
-    add_3 = torch.ops.aten.add.Tensor(primals_5, primals_6);  primals_5 = primals_6 = None
-    sin_2 = torch.ops.aten.sin.default(add_3)
-    add_4 = torch.ops.aten.add.Tensor(add_2, sin_2);  add_2 = sin_2 = None
-    cos = torch.ops.aten.cos.default(add_3);  add_3 = None
-    cos_1 = torch.ops.aten.cos.default(add_1);  add_1 = None
-    record_event_default = torch.ops.streams.record_event.default(2, 0);  record_event_default = None
-    stream_in_cpu_offload_cos_1 = torch.ops.streams.fork.default(0, 1);  stream_in_cpu_offload_cos_1 = None
-    wait_event_default = torch.ops.streams.wait_event.default(2, 1);  wait_event_default = None
-    record_stream_cos_1 = torch.ops.streams.record_stream.default(cos_1, 1);  record_stream_cos_1 = None
-    cpu_offload_cos_1 = torch.ops.prims.device_put.default(cos_1, device(type='cpu'), non_blocking = True);  cos_1 = None
-    record_event_default_1 = torch.ops.streams.record_event.default(3, 1);  record_event_default_1 = None
-    stream_out_cpu_offload_cos_1 = torch.ops.streams.join.default(1, 0);  stream_out_cpu_offload_cos_1 = None
-    cos_2 = torch.ops.aten.cos.default(add);  add = None
-    wait_event_default_1 = torch.ops.streams.wait_event.default(3, 0);  wait_event_default_1 = None
-    return (add_4, cos, cpu_offload_cos_1, cos_2)""",
-        )
-
-        self.assertExpectedInline(
-            bw_graph.code.strip(),
-            """\
-def forward(self, cos, cpu_offload_cos_1, cos_2, tangents_1):
-    stream_in_gpu_reload_cos_1 = torch.ops.streams.fork.default(4, 5);  stream_in_gpu_reload_cos_1 = None
-    wait_stream_default = torch.ops.streams.wait_stream.default(5, 4);  wait_stream_default = None
-    gpu_reload_cos_1 = torch.ops.prims.device_put.default(cpu_offload_cos_1, device(type='GPU_TYPE', index=0), non_blocking = True);  cpu_offload_cos_1 = None
-    record_event_default = torch.ops.streams.record_event.default(6, 5);  record_event_default = None
-    stream_out_gpu_reload_cos_1 = torch.ops.streams.join.default(5, 4);  stream_out_gpu_reload_cos_1 = None
-    mul = torch.ops.aten.mul.Tensor(tangents_1, cos);  cos = None
-    wait_event_default = torch.ops.streams.wait_event.default(6, 4);  wait_event_default = None
-    mul_1 = torch.ops.aten.mul.Tensor(tangents_1, gpu_reload_cos_1);  gpu_reload_cos_1 = None
-    mul_2 = torch.ops.aten.mul.Tensor(tangents_1, cos_2);  tangents_1 = cos_2 = None
-    return (mul_2, mul_2, mul_1, mul_1, mul, mul)""".replace("GPU_TYPE", GPU_TYPE),
-        )
+        nodes = list(bw_graph.graph.nodes)
+        reload_nodes = [
+            n
+            for n in nodes
+            if n.op == "call_function" and n.target == torch.ops.ao.reload.default
+        ]
+        wait_nodes = [
+            n
+            for n in nodes
+            if n.op == "call_function"
+            and n.target == torch.ops.ao.wait_tensor.default
+            and isinstance(n.args[0], torch.fx.Node)
+            and n.args[0].target == torch.ops.ao.reload.default
+        ]
+        # Reload should appear before its corresponding wait (prefetched earlier)
+        for reload_node in reload_nodes:
+            wait_node = next(w for w in wait_nodes if w.args[0] is reload_node)
+            self.assertLess(nodes.index(reload_node), nodes.index(wait_node))
 
     @serialTest()
     def test_partitioner_offload_sep_stream_reorder_accuracy(self):
@@ -312,6 +315,145 @@ def test_partitioner_offload_sep_stream_reorder_accuracy(self):
                 atol=1e-5,
             )
 
+    def test_wait_tensor_error_no_pending_transfer(self):
+        """Test _pop_wait raises RuntimeError for unregistered tensor."""
+        from torch._functorch._activation_offloading.offload_ops import (
+            _clear_wait_registry,
+            _pop_wait,
+        )
+
+        _clear_wait_registry()
+        x = torch.randn(4, device=GPU_TYPE)
+        with self.assertRaisesRegex(RuntimeError, "no pending transfer"):
+            _pop_wait(x)
+
+    def test_offload_reload_fake_tensor(self):
+        """Test fake tensor implementations for offload, reload, and wait_tensor."""
+        import torch._functorch._activation_offloading.offload_ops  # noqa: F401
+        from torch._subclasses.fake_tensor import FakeTensorMode
+
+        with FakeTensorMode():
+            x_gpu = torch.randn(4, 8, device=GPU_TYPE)
+
+            # offload fake: GPU -> CPU with same shape/dtype
+            cpu_result = torch.ops.ao.offload.default(x_gpu)
+            self.assertEqual(cpu_result.shape, (4, 8))
+            self.assertEqual(cpu_result.device.type, "cpu")
+            self.assertEqual(cpu_result.dtype, x_gpu.dtype)
+
+            # reload fake: CPU -> GPU with same shape/dtype
+            x_cpu = torch.randn(4, 8)
+            gpu_result = torch.ops.ao.reload.default(x_cpu, torch.device(GPU_TYPE))
+            self.assertEqual(gpu_result.shape, (4, 8))
+            self.assertEqual(gpu_result.device.type, GPU_TYPE)
+            self.assertEqual(gpu_result.dtype, x_cpu.dtype)
+
+            # wait_tensor fake: returns same tensor (aliasing)
+            wait_result = torch.ops.ao.wait_tensor.default(x_gpu)
+            self.assertEqual(wait_result.shape, x_gpu.shape)
+            self.assertEqual(wait_result.device, x_gpu.device)
+
+    def test_multiple_tensors_offload_ao_ops(self):
+        """Test offloading all saved tensors simultaneously with ao ops."""
+
+        def mark_all_cos_for_offloading(gm, joint_inputs):
+            for node in gm.graph.nodes:
+                if node.name.startswith("cos"):
+                    node.meta["should_offload"] = True
+            return gm
+
+        reset_user_object_tracking()
+        torch._dynamo.reset()
+        with torch._functorch.config.patch(
+            enable_activation_offloading=True,
+            activation_offload_separate_stream=True,
+            joint_custom_pass=mark_all_cos_for_offloading,
+        ):
+            fw_graph, bw_graph = get_fw_bw_graph(self.fn, [self.x])
+
+        fw_code = fw_graph.code
+        self.assertEqual(fw_code.count("ao.offload.default"), 3)
+        self.assertEqual(fw_code.count("ao.wait_tensor.default"), 3)
+
+        bw_code = bw_graph.code
+        self.assertEqual(bw_code.count("ao.reload.default"), 3)
+        self.assertEqual(bw_code.count("ao.wait_tensor.default"), 3)
+
+    def test_multiple_tensors_offload_ao_ops_accuracy(self):
+        """Test accuracy when offloading all saved tensors with ao ops."""
+
+        def mark_all_cos_for_offloading(gm, joint_inputs):
+            for node in gm.graph.nodes:
+                if node.name.startswith("cos"):
+                    node.meta["should_offload"] = True
+            return gm
+
+        x_ref = [x.detach().clone().requires_grad_(True) for x in self.x]
+        out_ref = self.fn(x_ref)
+        out_ref.sum().backward()
+        grads_ref = [inp.grad for inp in x_ref]
+
+        reset_user_object_tracking()
+        torch._dynamo.reset()
+        with torch._functorch.config.patch(
+            enable_activation_offloading=True,
+            activation_offload_separate_stream=True,
+            joint_custom_pass=mark_all_cos_for_offloading,
+        ):
+            fw_graph, bw_graph = get_fw_bw_graph(self.fn, [self.x])
+
+        # Verify ao ops are used (not old stream ops)
+        self.assertIn("ao.offload.default", fw_graph.code)
+        self.assertIn("ao.reload.default", bw_graph.code)
+
+        reset_user_object_tracking()
+        torch._dynamo.reset()
+        with torch._functorch.config.patch(
+            enable_activation_offloading=True,
+            activation_offload_separate_stream=True,
+            joint_custom_pass=mark_all_cos_for_offloading,
+        ):
+            x_compile = [x.detach().clone().requires_grad_(True) for x in self.x]
+            compiled_fn = torch.compile(self.fn, backend="aot_eager")
+            out_compiled = compiled_fn(x_compile)
+            out_compiled.sum().backward()
+            grads_compiled = [inp.grad for inp in x_compile]
+
+        for grad_ref, grad_compiled in zip(grads_ref, grads_compiled):
+            torch.testing.assert_close(grad_compiled, grad_ref, rtol=1e-5, atol=1e-5)
+
+    def test_inductor_offload_ao_ops(self):
+        """Test inductor code generation with ao ops."""
+        reset_user_object_tracking()
+        torch._dynamo.reset()
+
+        def run_compiled():
+            return torch.compile(self.fn)(self.x)
+
+        with torch._functorch.config.patch(
+            enable_activation_offloading=True,
+            activation_offload_separate_stream=True,
+            joint_custom_pass=self.joint_custom_pass,
+        ):
+            _, (fw_code, bw_code) = run_fw_bw_and_get_code(run_compiled)
+
+        (FileCheck().check("ao.offload").check("ao.wait_tensor").run(fw_code))
+        (FileCheck().check("ao.reload").check("ao.wait_tensor").run(bw_code))
+
+    def test_estimate_transfer_time_config(self):
+        """Test bandwidth config controls transfer time estimation."""
+        from torch._functorch._activation_offloading.activation_offloading import (
+            _estimate_transfer_time_in_ms,
+        )
+
+        size_bytes = 1024**3  # 1 GB
+        # 1 GB at 50 GB/s = 20 ms
+        with torch._functorch.config.patch(activation_offload_cpu_gpu_bw=50.0):
+            self.assertAlmostEqual(_estimate_transfer_time_in_ms(size_bytes), 20.0)
+        # 1 GB at 25 GB/s = 40 ms
+        with torch._functorch.config.patch(activation_offload_cpu_gpu_bw=25.0):
+            self.assertAlmostEqual(_estimate_transfer_time_in_ms(size_bytes), 40.0)
+
 
 @unittest.skipIf(not HAS_GPU, "requires GPU")
 class DefaultPartitionerActivationOffloadingTests(TestCase):
@@ -384,8 +526,8 @@ def forward(self, add, cpu_offload_add_1, add_3, tangents_1):
     return (mul_2, mul_2, mul_1, mul_1, mul, mul)""".replace("GPU_TYPE", GPU_TYPE),
         )
 
-    def test_default_partitioner_offload_sep_stream(self):
-        """Test offloading with separate streams using default partitioner"""
+    def test_default_partitioner_offload_ao_ops(self):
+        """Test ao ops with default partitioner."""
         reset_user_object_tracking()
         torch._dynamo.reset()
         with torch._functorch.config.patch(
@@ -397,47 +539,25 @@ def test_default_partitioner_offload_sep_stream(self):
                 self.fn, [self.x], partitioner=default_partition
             )
 
-        self.assertExpectedInline(
-            fw_graph.code.strip(),
-            """\
-def forward(self, primals_1, primals_2, primals_3, primals_4, primals_5, primals_6):
-    add = torch.ops.aten.add.Tensor(primals_1, primals_2);  primals_1 = primals_2 = None
-    sin = torch.ops.aten.sin.default(add)
-    add_1 = torch.ops.aten.add.Tensor(primals_3, primals_4);  primals_3 = primals_4 = None
-    sin_1 = torch.ops.aten.sin.default(add_1)
-    record_event_default = torch.ops.streams.record_event.default(2, 0);  record_event_default = None
-    stream_in_cpu_offload_add_1 = torch.ops.streams.fork.default(0, 1);  stream_in_cpu_offload_add_1 = None
-    wait_event_default = torch.ops.streams.wait_event.default(2, 1);  wait_event_default = None
-    record_stream_add_1 = torch.ops.streams.record_stream.default(add_1, 1);  record_stream_add_1 = None
-    cpu_offload_add_1 = torch.ops.prims.device_put.default(add_1, device(type='cpu'), non_blocking = True);  add_1 = None
-    record_event_default_1 = torch.ops.streams.record_event.default(3, 1);  record_event_default_1 = None
-    stream_out_cpu_offload_add_1 = torch.ops.streams.join.default(1, 0);  stream_out_cpu_offload_add_1 = None
-    wait_event_default_1 = torch.ops.streams.wait_event.default(3, 0);  wait_event_default_1 = None
-    add_2 = torch.ops.aten.add.Tensor(sin, sin_1);  sin = sin_1 = None
-    add_3 = torch.ops.aten.add.Tensor(primals_5, primals_6);  primals_5 = primals_6 = None
-    sin_2 = torch.ops.aten.sin.default(add_3)
-    add_4 = torch.ops.aten.add.Tensor(add_2, sin_2);  add_2 = sin_2 = None
-    return (add_4, add, cpu_offload_add_1, add_3)""",
+        fw_code = fw_graph.code
+        (
+            FileCheck()
+            .check("ao.offload.default")
+            .check("ao.wait_tensor.default")
+            .run(fw_code)
         )
+        self.assertNotIn("streams.fork", fw_code)
+        self.assertNotIn("streams.join", fw_code)
 
-        self.assertExpectedInline(
-            bw_graph.code.strip(),
-            """\
-def forward(self, add, cpu_offload_add_1, add_3, tangents_1):
-    cos = torch.ops.aten.cos.default(add_3);  add_3 = None
-    mul = torch.ops.aten.mul.Tensor(tangents_1, cos);  cos = None
-    stream_in_gpu_reload_add_1 = torch.ops.streams.fork.default(4, 5);  stream_in_gpu_reload_add_1 = None
-    wait_stream_default = torch.ops.streams.wait_stream.default(5, 4);  wait_stream_default = None
-    gpu_reload_add_1 = torch.ops.prims.device_put.default(cpu_offload_add_1, device(type='GPU_TYPE', index=0), non_blocking = True);  cpu_offload_add_1 = None
-    record_event_default = torch.ops.streams.record_event.default(6, 5);  record_event_default = None
-    stream_out_gpu_reload_add_1 = torch.ops.streams.join.default(5, 4);  stream_out_gpu_reload_add_1 = None
-    wait_event_default = torch.ops.streams.wait_event.default(6, 4);  wait_event_default = None
-    cos_1 = torch.ops.aten.cos.default(gpu_reload_add_1);  gpu_reload_add_1 = None
-    mul_1 = torch.ops.aten.mul.Tensor(tangents_1, cos_1);  cos_1 = None
-    cos_2 = torch.ops.aten.cos.default(add);  add = None
-    mul_2 = torch.ops.aten.mul.Tensor(tangents_1, cos_2);  tangents_1 = cos_2 = None
-    return (mul_2, mul_2, mul_1, mul_1, mul, mul)""".replace("GPU_TYPE", GPU_TYPE),
+        bw_code = bw_graph.code
+        (
+            FileCheck()
+            .check("ao.reload.default")
+            .check("ao.wait_tensor.default")
+            .run(bw_code)
         )
+        self.assertNotIn("streams.fork", bw_code)
+        self.assertNotIn("streams.join", bw_code)
 
     def test_default_partitioner_offload_accuracy(self):
         """Test that offloading with default partitioner produces correct gradients"""
diff --git a/test/dynamo/test_after_aot.py b/test/dynamo/test_after_aot.py
index 91fd1caea5de9..67f0d34acb4d4 100644
--- a/test/dynamo/test_after_aot.py
+++ b/test/dynamo/test_after_aot.py
@@ -8,10 +8,15 @@
 import unittest
 
 import torch._dynamo.test_case
-from torch._dynamo.repro.after_aot import InputReader, InputWriter, save_graph_repro
+from torch._dynamo.repro.after_aot import (
+    _extract_distributed_info,
+    InputReader,
+    InputWriter,
+    save_graph_repro,
+)
 from torch._higher_order_ops.triton_kernel_wrap import kernel_side_table
 from torch.fx.experimental.proxy_tensor import make_fx
-from torch.testing._internal.common_utils import IS_FBCODE
+from torch.testing._internal.common_utils import IS_FBCODE, TEST_CUDA
 from torch.utils._traceback import report_compile_source_on_error
 from torch.utils._triton import has_triton
 
@@ -108,6 +113,80 @@ def test(tensor, expected):
 reader.tensor(buf0, (3, 4, 5, 6), (120, 1, 24, 4), is_leaf=True)  # x""",
         )
 
+    def test_dump_opaque(self):
+        """save_graph_repro should emit reader.opaque() for FakeScriptObject args."""
+        from torch._library.fake_class_registry import FakeScriptObject
+
+        fake_obj = FakeScriptObject(object(), "__torch__.MyClass", None)
+
+        def f(x):
+            return (x * x,)
+
+        args = [torch.randn(4), fake_obj]
+        gm = make_fx(f)(args[0])
+        with gm.graph.inserting_before(next(iter(gm.graph.nodes))):
+            gm.graph.placeholder("obj")
+        gm.recompile()
+
+        buf = io.StringIO()
+        save_graph_repro(buf, gm, args, "inductor_accuracy")
+        r = buf.getvalue()
+        self.assertIn("reader.opaque('__torch__.MyClass')", r)
+
+    @unittest.skipIf(not TEST_CUDA, "requires CUDA")
+    def test_dump_generator(self):
+        torch.cuda.init()
+        gen = torch.cuda.default_generators[0].clone_state()
+        writer = InputWriter(None)
+        writer.generator("fwd_rng_state_0", gen)
+        self.assertExpectedInline(
+            "\n".join(writer._lines),
+            """reader.generator('cuda', 0)  # fwd_rng_state_0""",
+        )
+        reader = InputReader(None)
+        env = {"reader": reader, "torch": torch}
+        exec("\n".join(writer._lines), env)
+        self.assertIsInstance(reader.args[0], torch._C.Generator)
+        self.assertEqual(reader.args[0].device, torch.device("cuda", 0))
+
+    @unittest.skipIf(not TEST_CUDA, "requires CUDA")
+    def test_graphsafe_rng_repro(self):
+        """save_graph_repro should emit reader.generator() for Generator args."""
+        torch.cuda.init()
+        gen = torch.cuda.default_generators[0].clone_state()
+
+        def f(x):
+            return (x * x,)
+
+        args = [torch.randn(4, device="cuda"), gen]
+        gm = make_fx(f)(args[0])
+        with gm.graph.inserting_before(next(iter(gm.graph.nodes))):
+            gm.graph.placeholder("fwd_rng_state_0")
+        gm.recompile()
+
+        buf = io.StringIO()
+        save_graph_repro(buf, gm, args, "inductor_accuracy")
+        r = buf.getvalue()
+        self.assertIn("reader.generator('cuda', 0)", r)
+        self.assertNotIn("reader.unsupported(", r)
+
+    @unittest.skipIf(not torch.distributed.is_available(), "requires distributed")
+    def test_extract_distributed_info_skips_non_string_group_name(self):
+        """_extract_distributed_info should skip ops where group_name is an FX Node."""
+        gm = torch.fx.GraphModule({}, torch.fx.Graph())
+        g = gm.graph
+        x = g.placeholder("x")
+        group_name = g.placeholder("group_name")
+        ar = g.call_function(
+            torch.ops._c10d_functional.all_reduce.default,
+            args=(x, "sum", group_name),
+        )
+        g.output(ar)
+        gm.recompile()
+
+        result = _extract_distributed_info(gm)
+        self.assertEqual(result, {})
+
 
 if __name__ == "__main__":
     from torch._dynamo.test_case import run_tests
diff --git a/test/dynamo/test_aot_autograd.py b/test/dynamo/test_aot_autograd.py
index bffefbbf90e04..c83aa335c7c2b 100644
--- a/test/dynamo/test_aot_autograd.py
+++ b/test/dynamo/test_aot_autograd.py
@@ -256,29 +256,26 @@ def capturing_fn(gm, inputs):
         # for gm in gms:
         #     print("GM CODE", gm.code)
 
-        self.assertEqual(counter.frame_count, 4)
+        self.assertEqual(counter.frame_count, 2)
         self.assertEqual(counter.op_count, 7)
         # Graph 1
-        # def forward(self, x : torch.nn.parameter.Parameter, y : torch.Tensor):
-        #     mul = x * y;  x = y = None
-        #     return (mul,)
-        # BREAK
-        # Graph 2
-        # def forward(self, y : torch.Tensor):
-        #     getitem = y[0];  y = None
+        # def forward(self, L_x_ : torch.nn.parameter.Parameter, L_y_ : torch.Tensor):
+        #     l_x_ = L_x_
+        #     l_y_ = L_y_
+        #     a = l_x_ * l_y_;  l_x_ = l_y_ = None
+        #     getitem = a[0]
         #     getitem_1 = getitem[0];  getitem = None
         #     lt = getitem_1 < 3;  getitem_1 = None
-        #     return (lt,)
-        # BREAK
-        # Graph 3
-        # def forward(self, param : torch.Tensor, y : torch.Tensor):
-        #     add = y + param;  y = param = None
-        #     return (add,)
+        #     return (lt, a)
         # BREAK
-        # Graph 4
-        # def forward(self, _stack0 : torch.Tensor, x : torch.nn.parameter.Parameter, y : torch.Tensor):
-        #     add = x + y;  x = y = None
-        #     mul = _stack0 * add;  _stack0 = add = None
+        # Graph 2
+        # def forward(self, L_nested_frame_values_0_1_ : torch.Tensor, L_x_ : torch.nn.parameter.Parameter, L_y_ : torch.Tensor):
+        #     l_nested_frame_values_0_1_ = L_nested_frame_values_0_1_
+        #     l_x_ = L_x_
+        #     l_y_ = L_y_
+        #     a = l_nested_frame_values_0_1_ + l_nested_frame_values_0_1_;  l_nested_frame_values_0_1_ = None
+        #     b = l_x_ + l_y_;  l_x_ = l_y_ = None
+        #     mul = a * b;  a = b = None
         #     return (mul,)
 
         # Run fn with AOT
@@ -1109,7 +1106,7 @@ def f(x):
             )
 
     def test_aot_autograd_raises_invalid_leaf_set(self):
-        @torch.compile
+        @torch.compile(backend="aot_eager")
         def f(x):
             x.set_(torch.ones(2))
 
@@ -1820,6 +1817,60 @@ def forward(self, x: torch.Tensor):
         # Should only compile once regardless of batch size changes
         self.assertEqual(cnt.frame_count, 1)
 
+    @unittest.skipIf(not torch.cuda.is_available(), "requires cuda")
+    @patch.object(torch._dynamo.config, "capture_scalar_outputs", True)
+    @patch.object(torch._dynamo.config, "capture_dynamic_output_shape_ops", True)
+    def test_partitioner_filters_symbool_saved_for_bw(self):
+        """When cross-rank sync injects SymBool nodes (whose only consumers are
+        _assert_scalar) into saved_values, the partitioner must filter them out
+        so they don't become backward graph inputs that Inductor can't lower.
+
+        In multi-GPU MoE training, different EP ranks can produce different
+        min-cut partitioner decisions due to data-dependent shapes. The
+        cross-rank sync (_sync_decision_cross_ranks) picks one rank's decision
+        for all, which may include SymBool assertion nodes that weren't in
+        another rank's saved set. We simulate this by patching
+        _sync_decision_cross_ranks to inject all SymBool nodes.
+
+        Regression test for https://github.com/pytorch/pytorch/pull/179315
+        """
+        import torch._functorch.config as functorch_config
+        import torch._functorch.partitioners as P
+
+        # Simulate cross-rank sync picking a rank that saved SymBool nodes
+        def inject_symbool_nodes(joint_graph, saved_values):
+            for node in joint_graph.nodes:
+                val = node.meta.get("val")
+                if isinstance(val, torch.SymBool) and node not in saved_values:
+                    saved_values = saved_values + [node]
+            return saved_values
+
+        def fn(x, splits_tensor, mask):
+            splits = splits_tensor.tolist()
+            indices = torch.nonzero(mask).squeeze(-1)
+            y = x[indices]
+            # split verifies sum(splits) == y.shape[0], creating an
+            # Equality SymBool node consumed only by _assert_scalar
+            chunks = torch.split(y, splits, dim=0)
+            return torch.cat([c * 2.0 for c in chunks], dim=0).sum()
+
+        x = torch.randn(20, 16, requires_grad=True, device="cuda")
+        splits_tensor = torch.tensor([3, 5], device="cuda")
+        mask = torch.zeros(20, dtype=torch.bool, device="cuda")
+        mask[:8] = True
+
+        with (
+            patch.object(P, "_sync_decision_cross_ranks", inject_symbool_nodes),
+            patch.object(functorch_config, "_sync_decision_cross_ranks", True),
+        ):
+            compiled_fn = torch.compile(fn)
+            loss = compiled_fn(x, splits_tensor, mask)
+            # Without the fix, this raises:
+            # InductorError: Unsupported inductor graph input type:
+            #   <class 'sympy.core.relational.GreaterThan'>
+            loss.backward()
+        self.assertIsNotNone(x.grad)
+
 
 if __name__ == "__main__":
     from torch._dynamo.test_case import run_tests
diff --git a/test/dynamo/test_aot_autograd_cache.py b/test/dynamo/test_aot_autograd_cache.py
index 2d9775aca58d8..07c679f494a77 100644
--- a/test/dynamo/test_aot_autograd_cache.py
+++ b/test/dynamo/test_aot_autograd_cache.py
@@ -1,23 +1,32 @@
 # Owner(s): ["module: dynamo"]
 
+import contextlib
 import copy
 import dataclasses
 import functools
+import multiprocessing
+import operator
 import os
 import pickle
 import shutil
 import unittest
+from collections.abc import Sequence
+from typing import Literal
 from unittest.mock import patch
 
 import torch
 import torch._dynamo
 import torch._dynamo.test_case
 import torch._functorch._aot_autograd
+import torch._functorch._aot_autograd.autograd_cache as autograd_cache
+import torch._functorch.aot_autograd as aot_autograd
+import torch._inductor.compile_fx as compile_fx
 from torch._dynamo import config as dynamo_config
 from torch._dynamo.utils import counters
 from torch._functorch import config as functorch_config
 from torch._functorch._aot_autograd.autograd_cache import (
     AOTAutogradCache,
+    AOTAutogradCachePickler,
     autograd_cache_key,
     BypassAOTAutogradCache,
     sanitize_gm_for_cache,
@@ -25,20 +34,29 @@
 from torch._functorch._aot_autograd.schemas import AOTConfig
 from torch._guards import TracingContext
 from torch._inductor import config as inductor_config
-from torch._inductor.custom_graph_pass import CustomRuntimeEstimator
+from torch._inductor.custom_graph_pass import (
+    CustomGraphPass,
+    CustomGraphPassType,
+    CustomRuntimeEstimator,
+)
 from torch._inductor.runtime.runtime_utils import cache_dir
 from torch._inductor.runtime.triton_compat import tl, triton
+from torch._inductor.standalone_compile import (
+    autograd_cache_key as standalone_compile_autograd_cache_key,
+)
 from torch._inductor.test_case import TestCase as InductorTestCase
-from torch._inductor.utils import fresh_cache
+from torch._inductor.utils import fresh_cache, InputType
 from torch._subclasses import FakeTensorMode
 from torch.compiler._cache import CacheArtifactManager
+from torch.fx import GraphModule
 from torch.fx.experimental.symbolic_shapes import ShapeEnv
-from torch.testing._internal.common_cuda import SM80OrLater, TEST_MULTIGPU
+from torch.testing._internal.common_cuda import SM80OrLater
 from torch.testing._internal.common_device_type import largeTensorTest
 from torch.testing._internal.common_utils import (
     instantiate_parametrized_tests,
     parametrize,
     skipIfWindows,
+    subtest,
 )
 from torch.testing._internal.inductor_utils import GPU_TYPE, HAS_GPU, requires_triton
 from torch.testing._internal.triton_utils import requires_cuda_and_triton
@@ -50,6 +68,48 @@
 )
 
 
+class CustomPreGradPassRemoveIdentMuls(CustomGraphPass):
+    """
+    Pre-grad pass that removes redundant identity multiplications (1 * x).
+    """
+
+    def __call__(self, g: torch.fx.Graph) -> None:
+        for n in g.nodes:
+            if n.op == "call_function" and n.target is operator.mul:
+                lhs, rhs = n.args
+                if lhs == 1:
+                    n.replace_all_uses_with(rhs)
+                    g.erase_node(n)
+
+    def uuid(self):
+        return "custom_pre_grad_pass_remove_ident_muls_v1"
+
+
+class CustomPreGradPassRemoveIdentMulsNoUUID(CustomPreGradPassRemoveIdentMuls):
+    def uuid(self):
+        return None
+
+
+custom_pre_grad_pass_remove_ident_muls = CustomPreGradPassRemoveIdentMuls()
+custom_pre_grad_pass_remove_ident_muls_wo_uuid = (
+    CustomPreGradPassRemoveIdentMulsNoUUID()
+)
+
+
+@contextlib.contextmanager
+def _fake_process_group(rank=0, world_size=2):
+    """Context manager for setting up and tearing down a fake process group."""
+    import torch.distributed as c10d
+    from torch.testing._internal.distributed.fake_pg import FakeStore
+
+    fake_store = FakeStore()
+    c10d.init_process_group("fake", store=fake_store, rank=rank, world_size=world_size)
+    try:
+        yield
+    finally:
+        c10d.destroy_process_group()
+
+
 def aot_eager_regional_inductor():
     """
     Regional inductor backend for AOT autograd.
@@ -134,8 +194,99 @@ def _unpack_fp8_with_scale_wrap(x):
     return y.to(dtype)
 
 
+class CacheKeyEquivalenceMixin:
+    """Mixin that verifies the standalone autograd_cache_key API produces the
+    same key as the internal compilation pipeline for every torch.compile
+    invocation in the test."""
+
+    def setUp(self):
+        super().setUp()
+        self._compile_fx_patcher = self._install_compile_fx_patcher()
+        self._compile_fx_patcher.start()
+
+    def _install_compile_fx_patcher(self):
+        real_compile_fx = torch._inductor.compile_fx.compile_fx
+
+        def capturing_compile_fx(mod, example_inputs_, **kwargs):
+            mod.recompile()
+            mod_code = mod.code
+
+            # Record standalone key before invoking the inner compile_fx function.
+            standalone_key = standalone_debug_lines = None
+            standalone_bypassed = False
+            try:
+                standalone_key, standalone_debug_lines = compile_fx.autograd_cache_key(
+                    mod, example_inputs_, ignore_shape_env=False
+                )
+            except (BypassAOTAutogradCache, NotImplementedError):
+                standalone_bypassed = True
+            mod.recompile()
+            self.assertEqual(
+                mod.code,
+                mod_code,
+                "compile_fx.autograd_cache_key must not mutate the model",
+            )
+
+            # Setup to capture the ground truth key.
+            captured_gt_key = captured_gt_debug_lines = None
+            real_cache_key = autograd_cache.autograd_cache_key
+
+            def capturing_cache_key(mod, ei, config, compiler_config_extra=None):
+                nonlocal captured_gt_key, captured_gt_debug_lines
+                result = real_cache_key(mod, ei, config, compiler_config_extra)
+                captured_gt_key, captured_gt_debug_lines = result
+                return result
+
+            # Setup to capture pre grad changes.
+            captured_pre_grad_change = False
+            real_run_pre_grad_passes = torch._inductor.compile_fx.run_pre_grad_passes
+
+            def capturing_run_pre_grad_passes(
+                mod_inner: GraphModule, example_inputs_: Sequence[InputType]
+            ) -> torch.fx.GraphModule:
+                nonlocal captured_pre_grad_change
+                mod_inner.recompile()
+                mod_inner_code = mod_inner.code
+                result = real_run_pre_grad_passes(mod_inner, example_inputs_)
+                mod_inner.recompile()
+                captured_pre_grad_change = mod_inner.code != mod_inner_code
+                return result
+
+            # Invoke the inner compile_fx function with capture patches.
+            with (
+                patch(
+                    "torch._functorch._aot_autograd.autograd_cache.autograd_cache_key",
+                    side_effect=capturing_cache_key,
+                ),
+                patch(
+                    "torch._inductor.compile_fx.run_pre_grad_passes",
+                    side_effect=capturing_run_pre_grad_passes,
+                ),
+            ):
+                result = real_compile_fx(mod, example_inputs_, **kwargs)
+
+            # Verify standalone key matches the one produced during compilation.
+            if not standalone_bypassed and not captured_pre_grad_change:
+                self.assertEqual(
+                    standalone_key,
+                    captured_gt_key,
+                    "standalone cache key differs from compile_fx cache key",
+                )
+
+            return result
+
+        return patch(
+            "torch._inductor.compile_fx.compile_fx",
+            side_effect=capturing_compile_fx,
+        )
+
+    def tearDown(self):
+        self._compile_fx_patcher.stop()
+        super().tearDown()
+
+
 @instantiate_parametrized_tests
-class AOTAutogradCacheTests(InductorTestCase):
+class AOTAutogradCacheTests(CacheKeyEquivalenceMixin, InductorTestCase):
     def setUp(self):
         """
         Reset all counters and caches before each unit test
@@ -582,7 +733,7 @@ def test_invoke_subgraph(self):
         def gn(x, y):
             return x + y
 
-        @torch.compile
+        @torch.compile(backend="inductor")
         def fn(x, y):
             return gn(x, y) + gn(x, y)
 
@@ -1763,35 +1914,30 @@ def forward(self, x):
             self.assertNotEqual(res1, res3)
             self.assertEqual(res1, res3.sub(torch.ones(2, 2)))
 
+    @unittest.skipIf(not HAS_GPU, "requires GPU")
     @inductor_config.patch("fx_graph_cache", True)
     @inductor_config.patch("fx_graph_remote_cache", False)
-    @functorch_config.patch({"enable_autograd_cache": True})
-    @unittest.skipIf(not TEST_MULTIGPU, "only one GPU detected")
-    def test_constant_tensor_device_guards(self):
-        """
-        Usually, when there are example inputs, the device index of the inputs
-        is sufficient to make sure we don't cache hit with the results from different
-        cuda devices.
-        When the input has no arguments, we still need to have the cuda
-        device index in the cache key.
-        """
-
-        @torch.compile
-        def f():
-            y = torch.tensor([5], device="cuda")
-            return (y,)
-
-        with torch.cuda._DeviceGuard(0):
-            torch.cuda.set_device(0)
-            result = f()
-            self.assertEqual(result[0].device, torch.device("cuda:0"))
+    @inductor_config.patch("always_keep_tensor_constants", True)
+    @functorch_config.patch(
+        {
+            "enable_autograd_cache": True,
+            "strict_autograd_cache": True,
+            "bundled_autograd_cache": True,
+        }
+    )
+    def test_tensor_constant(self):
+        """Bundled AOT cache serialization with strict_autograd_cache should
+        succeed when the compiled graph contains tensor constants."""
 
-        self._clear_dynamo_and_codecache()
+        def fn(x):
+            idx = torch.tensor([], dtype=torch.long, device=x.device)
+            return x[idx]
 
-        with torch.cuda._DeviceGuard(1):
-            torch.cuda.set_device(1)
-            result = f()
-            self.assertEqual(result[0].device, torch.device("cuda:1"))
+        compiled_fn = torch.compile(fn, backend="inductor", fullgraph=True)
+        with torch.no_grad():
+            compiled_fn(torch.randn(4, device=GPU_TYPE))
+            self.assertEqual(counters["aot_autograd"]["autograd_cache_miss"], 1)
+            self.assertEqual(counters["aot_autograd"]["autograd_cache_saved"], 1)
 
     @requires_cuda_and_triton
     @inductor_config.patch("fx_graph_cache", True)
@@ -2428,6 +2574,41 @@ def uuid(self):
                 self.assertEqual(counters["aot_autograd"]["autograd_cache_saved"], 0)
                 self.assertGreater(counters["aot_autograd"]["autograd_cache_bypass"], 0)
 
+    @inductor_config.patch("fx_graph_cache", True)
+    @functorch_config.patch("enable_autograd_cache", True)
+    def test_pre_grad_passes_inplace_cpu(self):
+        """
+        `remove_identity` was an out-place operation previously, and it is called
+        for CPU inference inside `pre_grad_passes`. So, it broke the in-place semantics
+        of `pre_grad_passes` for CPU inference. This test ensures `pre_grad_passes` is
+        in-place in this case.
+        """
+
+        from torch._inductor.fx_passes.pre_grad import pre_grad_passes
+
+        class Model(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.linear = torch.nn.Linear(4, 4)
+                self.identity = torch.nn.Identity()
+
+            def forward(self, x):
+                x = self.linear(x)
+                x = self.identity(x)
+                return torch.relu(x)
+
+        x = torch.randn(2, 4)
+        gm = torch.fx.symbolic_trace(Model().eval())
+
+        with inductor_config.patch(
+            pattern_matcher=True,
+            freezing=True,
+        ):
+            with torch.no_grad():
+                gm_after = pre_grad_passes(gm, (x,))
+
+        self.assertTrue(id(gm_after) == id(gm))
+
     @inductor_config.patch("fx_graph_remote_cache", False)
     @inductor_config.patch("fx_graph_cache", True)
     @functorch_config.patch({"enable_autograd_cache": True})
@@ -2538,6 +2719,452 @@ def forward(self, x):
             x_view2.untyped_storage().data_ptr(), x2.untyped_storage().data_ptr()
         )
 
+    @inductor_config.patch("fx_graph_remote_cache", False)
+    @inductor_config.patch("fx_graph_cache", True)
+    @functorch_config.patch({"enable_autograd_cache": True})
+    def test_cache_hit_logs_aot_graphs(self):
+        """
+        Verify that TORCH_LOGS=aot_graphs prints graphs on cache hit, not just cache miss.
+        See https://github.com/pytorch/pytorch/issues/149060
+        """
+        from torch.testing._internal.logging_utils import logs_to_string
+
+        def fn(x, y):
+            return (x * 2, y @ y)
+
+        a = torch.rand(25)
+        b = torch.rand(5, 5)
+
+        compiled_fn = torch.compile(fn, backend="inductor")
+
+        # First call: cache miss — aot_graphs should be logged
+        log_stream, ctx = logs_to_string(
+            "torch._functorch._aot_autograd.graph_capture", "aot_graphs"
+        )
+        with ctx():
+            compiled_fn(a, b)
+        self.assertEqual(counters["aot_autograd"]["autograd_cache_miss"], 1)
+        self.assertIn("Forward graph", log_stream.getvalue())
+
+        # Reset dynamo so guards force recompilation (but keep disk cache)
+        self._clear_dynamo_and_codecache()
+
+        # Second call: cache hit — aot_graphs should ALSO be logged.
+        # On cache hit the log comes from aot_autograd_result, not graph_capture.
+        log_stream2, ctx2 = logs_to_string(
+            "torch._functorch._aot_autograd.aot_autograd_result", "aot_graphs"
+        )
+        with ctx2():
+            compiled_fn(a, b)
+        self.assertEqual(counters["aot_autograd"]["autograd_cache_hit"], 1)
+        self.assertIn("Forward graph (from cache)", log_stream2.getvalue())
+
+    @inductor_config.patch("fx_graph_cache", True)
+    @functorch_config.patch("enable_autograd_cache", True)
+    @parametrize(
+        "pre_grad_pass_timing,pre_grad_custom_pass,expect_pre_grad_call_count",
+        [
+            subtest(
+                ("early", custom_pre_grad_pass_remove_ident_muls, (1, 2)),
+                name="early_with_uuid",
+            ),
+            subtest(
+                ("late", custom_pre_grad_pass_remove_ident_muls, (1, 1)),
+                name="late_with_uuid",
+            ),
+            subtest(
+                ("default", custom_pre_grad_pass_remove_ident_muls, (1, 1)),
+                name="default_with_uuid",
+            ),
+            subtest(
+                ("early", custom_pre_grad_pass_remove_ident_muls_wo_uuid, (1, 2)),
+                name="early_without_uuid",
+            ),
+            # late_without_uuid will raise an exception and is tested separately in
+            # test_pre_grad_pass_late_timing_without_uuid_raises
+            subtest(
+                ("default", custom_pre_grad_pass_remove_ident_muls_wo_uuid, (1, 2)),
+                name="default_without_uuid",
+            ),
+        ],
+    )
+    @parametrize(
+        "use_module",
+        [
+            subtest(False, name="function"),
+            subtest(True, name="module"),
+        ],
+    )
+    def test_pre_grad_passes_timing(
+        self,
+        pre_grad_pass_timing: Literal["early", "late", "default"],
+        pre_grad_custom_pass: CustomGraphPassType,
+        expect_pre_grad_call_count: tuple[int, int],
+        use_module: bool,
+    ):
+        from torch._inductor.compile_fx import compile_fx_forward, run_pre_grad_passes
+
+        def fn(x, y):
+            return 1 * x + y
+
+        class LinearModel(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.linear = torch.nn.Linear(10, 10)
+
+            def forward(self, x):
+                return 1 * self.linear(x)
+
+        pre_grad_call_count = 0
+        pre_grad_input_graphs: list[str] = []
+        inductor_input_graphs: list[str] = []
+
+        def wrap_run_pre_grad_passes(
+            model: GraphModule, example_inputs: Sequence[InputType]
+        ) -> GraphModule:
+            nonlocal pre_grad_call_count
+            pre_grad_call_count += 1
+            pre_grad_input_graphs.append(model.print_readable(print_output=False))
+            run_pre_grad_passes(model, example_inputs)
+            return model
+
+        @torch.utils._typing_utils.copy_func_params(compile_fx_forward)
+        def wrap_compile_fx_forward(gm, *args, **kwargs):
+            inductor_input_graphs.append(gm.print_readable(print_output=False))
+            return compile_fx_forward(gm, *args, **kwargs)
+
+        if use_module:
+            target = LinearModel()
+            inputs = (torch.randn(10),)
+        else:
+            target = fn
+            inputs = (torch.randn(10), torch.randn(10))
+
+        with (
+            unittest.mock.patch(
+                "torch._inductor.compile_fx.run_pre_grad_passes",
+                wrap_run_pre_grad_passes,
+            ),
+            unittest.mock.patch(
+                "torch._inductor.compile_fx.compile_fx_forward",
+                wrap_compile_fx_forward,
+            ),
+            inductor_config.patch("pre_grad_pass_timing", pre_grad_pass_timing),
+            inductor_config.patch("pre_grad_custom_pass", pre_grad_custom_pass),
+            # Autograd cache doesn't save entries for modules with trainable
+            # parameters, so disable grad to get consistent cache hit counts
+            # across the function and module variants.
+            torch.no_grad(),
+        ):
+            self._clear_all_caches()
+
+            # First compilation - expect cache miss.
+            compiled = torch.compile(target)
+            result1 = compiled(*inputs)
+
+            # Assert cache miss.
+            self.assertEqual(counters["aot_autograd"]["autograd_cache_miss"], 1)
+            self.assertEqual(counters["aot_autograd"]["autograd_cache_hit"], 0)
+
+            # Assert #invocation of pre-grad passes.
+            self.assertEqual(pre_grad_call_count, expect_pre_grad_call_count[0])
+
+            # The pre-grad graph should contain the identity mul (= 1 * ...)
+            # before the custom pass removes it.
+            self.assertEqual(len(pre_grad_input_graphs), 1)
+            self.assertIn("= 1 * ", pre_grad_input_graphs[0])
+            self.assertEqual(len(inductor_input_graphs), 1)
+            self.assertNotIn("= 1 * ", inductor_input_graphs[0])
+
+            torch._dynamo.reset()
+
+            # Second compilation - expect cache hit.
+            compiled2 = torch.compile(target)
+            result2 = compiled2(*inputs)
+
+            # Assert cache hit.
+            self.assertEqual(counters["aot_autograd"]["autograd_cache_miss"], 1)
+            self.assertEqual(counters["aot_autograd"]["autograd_cache_hit"], 1)
+
+            # Assert #invocation of pre-grad passes.
+            self.assertEqual(pre_grad_call_count, expect_pre_grad_call_count[1])
+
+            self.assertEqual(result1, result2)
+
+    @inductor_config.patch("fx_graph_cache", True)
+    @functorch_config.patch("enable_autograd_cache", True)
+    @inductor_config.patch("pre_grad_pass_timing", "late")
+    def test_pre_grad_pass_late_timing_without_uuid_raises(self):
+        """
+        Explicitly setting late timing with a pass that has no UUID should
+        raise a RuntimeError.
+        """
+
+        class NoUuidPass(CustomGraphPass):
+            def __call__(self, g: torch.fx.Graph) -> None:
+                pass
+
+            def uuid(self):
+                return None
+
+        def fn(x, y):
+            return x + y
+
+        x = torch.randn(10)
+        y = torch.randn(10)
+
+        with inductor_config.patch("pre_grad_custom_pass", NoUuidPass()):
+            self._clear_all_caches()
+            compiled_fn = torch.compile(fn)
+            with self.assertRaisesRegex(
+                RuntimeError, "pre_grad_custom_pass must implement uuid"
+            ):
+                compiled_fn(x, y)
+
+    @inductor_config.patch("fx_graph_cache", True)
+    @functorch_config.patch("enable_autograd_cache", True)
+    def test_pre_grad_pass_default_timing_without_uuid_warns(self):
+        """
+        Default timing with a pass that has no UUID should log a warning about
+        the pre-grad pass cache being bypassed, only once per pass class.
+        """
+        from torch._inductor.codecache import _warned_pre_grad_pass_missing_uuid
+
+        class NoUuidPass(CustomGraphPass):
+            def __call__(self, g: torch.fx.Graph) -> None:
+                pass
+
+            def uuid(self):
+                return None
+
+        def fn(x, y):
+            return x + y
+
+        x = torch.randn(10)
+        y = torch.randn(10)
+
+        with (
+            inductor_config.patch("pre_grad_custom_pass", NoUuidPass()),
+            inductor_config.patch("pre_grad_pass_timing", "default"),
+        ):
+            _warned_pre_grad_pass_missing_uuid.clear()
+            self._clear_all_caches()
+
+            # First compilation — should warn.
+            compiled_fn = torch.compile(fn)
+            with self.assertLogs(
+                "torch._inductor.codecache", level="WARNING"
+            ) as log_cm:
+                compiled_fn(x, y)
+            uuid_warnings = [
+                m for m in log_cm.output if "does not implement uuid()" in m
+            ]
+            self.assertEqual(len(uuid_warnings), 1)
+
+            # Second compilation — same pass class, should not warn again.
+            torch._dynamo.reset()
+            compiled_fn2 = torch.compile(fn)
+            with self.assertNoLogs("torch._inductor.codecache", level="WARNING"):
+                compiled_fn2(x, y)
+
+    @parametrize(
+        "pre_grad_pass_timing,has_uuid,expect_miss_on_different_uuid",
+        [
+            subtest(("early", True, False), name="early_with_uuid"),
+            subtest(("late", True, True), name="late_with_uuid"),
+            subtest(("default", True, True), name="default_with_uuid"),
+            subtest(("early", False, False), name="early_without_uuid"),
+            # late + no uuid raises RuntimeError, tested separately
+            subtest(("default", False, False), name="default_without_uuid"),
+        ],
+    )
+    def test_cache_hit_across_processes_pre_grad_custom_pass(
+        self, pre_grad_pass_timing, has_uuid, expect_miss_on_different_uuid
+    ):
+        """
+        Verify cache behavior across processes for different pre-grad pass
+        timing and UUID configurations.
+        """
+        import subprocess
+        import sys
+        import tempfile
+        import textwrap
+
+        with tempfile.TemporaryDirectory() as cache_dir:
+            script_template = textwrap.dedent(
+                """
+                import json
+                import operator
+                import torch
+                import torch._dynamo
+                from torch._dynamo.utils import counters
+                from torch._inductor import config as inductor_config
+                from torch._inductor.custom_graph_pass import CustomGraphPass
+
+                inductor_config.fx_graph_cache = True
+                inductor_config.fx_graph_remote_cache = False
+                inductor_config.pre_grad_pass_timing = "{pre_grad_pass_timing}"
+                torch._dynamo.reset()
+
+                class TestPreGradPass(CustomGraphPass):
+                    def __call__(self, g):
+                        changed = False
+                        for n in g.nodes:
+                            if n.op == "call_function" and n.target is operator.mul:
+                                lhs, rhs = n.args
+                                if lhs == 1:
+                                    n.replace_all_uses_with(rhs)
+                                    g.erase_node(n)
+                                    changed = True
+                        if not changed:
+                            raise RuntimeError(
+                                "Custom pass did not change the graph. "
+                                "All test cases expect the pass to modify the graph."
+                            )
+
+                    def uuid(self):
+                        return {pass_uuid}
+
+                inductor_config.pre_grad_custom_pass = TestPreGradPass()
+
+                def fn(x, y):
+                    return 1 * x + y
+
+                compiled_fn = torch.compile(fn)
+                x = torch.randn(10)
+                y = torch.randn(10)
+                compiled_fn(x, y)
+
+                print(json.dumps(dict(counters["aot_autograd"])))
+                """
+            )
+
+            env = {**os.environ, "TORCHINDUCTOR_CACHE_DIR": cache_dir}
+
+            def run_script(pass_uuid):
+                script = script_template.format(
+                    pre_grad_pass_timing=pre_grad_pass_timing,
+                    pass_uuid=repr(pass_uuid),
+                )
+                result = subprocess.run(
+                    [sys.executable, "-c", script],
+                    env=env,
+                    capture_output=True,
+                    text=True,
+                )
+                self.assertEqual(result.returncode, 0, result.stderr)
+                import json
+
+                return json.loads(result.stdout.splitlines()[-1])
+
+            uuid_a = "pass_uuid_A" if has_uuid else None
+            uuid_b = "pass_uuid_B" if has_uuid else None
+
+            # First run - expect cache miss
+            c1 = run_script(uuid_a)
+            self.assertEqual(c1.get("autograd_cache_miss", 0), 1)
+            self.assertEqual(c1.get("autograd_cache_hit", 0), 0)
+
+            # Second run with same pass - expect cache hit
+            c2 = run_script(uuid_a)
+            self.assertEqual(c2.get("autograd_cache_miss", 0), 0)
+            self.assertEqual(c2.get("autograd_cache_hit", 0), 1)
+
+            # Third run with different pass UUID - miss only when UUID is
+            # part of the cache key (late timing with UUID)
+            c3 = run_script(uuid_b)
+            if expect_miss_on_different_uuid:
+                self.assertEqual(c3.get("autograd_cache_miss", 0), 1)
+                self.assertEqual(c3.get("autograd_cache_hit", 0), 0)
+            else:
+                self.assertEqual(c3.get("autograd_cache_miss", 0), 0)
+                self.assertEqual(c3.get("autograd_cache_hit", 0), 1)
+
+    @unittest.skipIf(not torch.distributed.is_available(), "requires distributed")
+    @inductor_config.patch("fx_graph_remote_cache", False)
+    @inductor_config.patch("fx_graph_cache", True)
+    @functorch_config.patch({"enable_autograd_cache": True})
+    def test_dtensor_cache_hit(self):
+        """
+        Test that DTensor produces cache hits on second compile.
+
+        This follows the standard AOTAutograd cache test pattern: compile once
+        (cache miss), reset dynamo, compile again with equivalent input (cache hit).
+        """
+        from torch.distributed.device_mesh import init_device_mesh
+        from torch.distributed.tensor import DTensor, Replicate
+
+        with _fake_process_group():
+            mesh = init_device_mesh("cpu", (2,))
+
+            def fn(x):
+                return x.sin()
+
+            compiled_fn = torch.compile(fn, backend="inductor")
+
+            # First call - cache miss
+            dtensor1 = DTensor.from_local(torch.zeros(4, 4), mesh, [Replicate()])
+            compiled_fn(dtensor1)
+
+            self.assertEqual(counters["aot_autograd"]["autograd_cache_miss"], 1)
+            self.assertEqual(counters["aot_autograd"]["autograd_cache_hit"], 0)
+
+            # Reset dynamo but keep the cache
+            self._clear_dynamo_and_codecache()
+
+            # Second call with equivalent DTensor - should hit cache
+            dtensor2 = DTensor.from_local(torch.zeros(4, 4), mesh, [Replicate()])
+            compiled_fn(dtensor2)
+
+            self.assertEqual(counters["aot_autograd"]["autograd_cache_miss"], 1)
+            self.assertEqual(counters["aot_autograd"]["autograd_cache_hit"], 1)
+
+    @inductor_config.patch("fx_graph_remote_cache", False)
+    @inductor_config.patch("fx_graph_cache", True)
+    @functorch_config.patch({"enable_autograd_cache": True})
+    def test_non_tuple_return_cache_key(self):
+        """autograd_cache_key does not support graphs that don't return a tuple,
+        since compile_fx mutates them via make_graph_return_tuple."""
+        from torch._subclasses import FakeTensorMode
+        from torch.fx.experimental.proxy_tensor import make_fx
+        from torch.fx.experimental.symbolic_shapes import ShapeEnv
+
+        def fn(x):
+            return x.sin()
+
+        inp = torch.randn(10)
+        fake_mode = FakeTensorMode(shape_env=ShapeEnv())
+        fake_inp = fake_mode.from_tensor(inp)
+        gm = make_fx(fn, tracing_mode="fake")(fake_inp)
+        self.assertFalse(compile_fx.graph_returns_tuple(gm))
+        with self.assertRaisesRegex(NotImplementedError, "don't return a tuple"):
+            compile_fx.autograd_cache_key(gm, [fake_inp], ignore_shape_env=False)
+        compile_fx.compile_fx(gm, [fake_inp])
+
+    @inductor_config.patch("fx_graph_remote_cache", False)
+    @inductor_config.patch("fx_graph_cache", True)
+    @functorch_config.patch({"enable_autograd_cache": True})
+    def test_flatten_graph_inputs_cache_key(self):
+        """autograd_cache_key does not support nested container inputs, since
+        compile_fx wraps them via flatten_graph_inputs."""
+        from torch._subclasses import FakeTensorMode
+        from torch.fx.experimental.proxy_tensor import make_fx
+        from torch.fx.experimental.symbolic_shapes import ShapeEnv
+
+        def fn(xs):
+            return (xs[0] + xs[1],)
+
+        inp = torch.randn(10)
+        fake_mode = FakeTensorMode(shape_env=ShapeEnv())
+        fake_x = fake_mode.from_tensor(inp)
+        fake_y = fake_mode.from_tensor(inp)
+        gm = make_fx(fn, tracing_mode="fake")([fake_x, fake_y])
+        with self.assertRaisesRegex(NotImplementedError, "nested container inputs"):
+            compile_fx.autograd_cache_key(
+                gm, [[fake_x, fake_y]], ignore_shape_env=False
+            )
+        compile_fx.compile_fx(gm, [[fake_x, fake_y]])
+
 
 @functorch_config.patch({"bundled_autograd_cache": True})
 class AOTAutogradCacheBundledTests(AOTAutogradCacheTests):
@@ -2606,7 +3233,7 @@ def gen_cache_key(self, f, config, inputs=None):
         # Not needed for actual key calculation.
         with torch._guards.tracing(ctx):
             with sanitize_gm_for_cache(fx_g):
-                return autograd_cache_key(fx_g, example_inputs, config, {})
+                return autograd_cache_key(fx_g, example_inputs, config, None)
 
     def test_basic_hash_key(self):
         def fn(x):
@@ -2814,18 +3441,22 @@ def fn(x):
             fx_g.meta = {"foo": "bar"}
             fx_g.compile_subgraph_reason = "Blah"
             config = self.default_config()
+            # autograd_cache_key now applies sanitize_gm_for_cache
+            # internally, so the result is the same whether we wrap the
+            # call or not.
             with sanitize_gm_for_cache(fx_g):
-                c1 = autograd_cache_key(fx_g, example_inputs, config, {})
-            c3 = autograd_cache_key(fx_g, example_inputs, config, {})
+                c1 = autograd_cache_key(fx_g, example_inputs, config, None)
+            c3 = autograd_cache_key(fx_g, example_inputs, config, None)
 
             fx_g.meta = {"foo": "baz"}
             fx_g.compile_subgraph_reason = None
             with sanitize_gm_for_cache(fx_g):
-                c2 = autograd_cache_key(fx_g, example_inputs, config, {})
-            c4 = autograd_cache_key(fx_g, example_inputs, config, {})
+                c2 = autograd_cache_key(fx_g, example_inputs, config, None)
+            c4 = autograd_cache_key(fx_g, example_inputs, config, None)
 
             self.assertEqual(c1, c2)
-            self.assertNotEqual(c3, c4)
+            self.assertEqual(c1, c3)
+            self.assertEqual(c1, c4)
 
     def test_dill_serialization_with_inner_functions(self):
         """
@@ -2926,6 +3557,244 @@ def test_pickle_entry_strict_mode_raises(self):
         ):
             AOTAutogradCache._pickle_entry(entry, remote=False)
 
+    @requires_cuda_and_triton
+    def test_prepare_for_pickle_clears_benchmark_failure_reasons(self):
+        """prepare_for_pickle clears benchmark_failure_reasons which can hold
+        exec'd launcher keys that aren't picklable.
+        """
+        from torch._inductor.runtime.hints import (
+            AttrsDescriptorWrapper,
+            DeviceProperties,
+            HeuristicType,
+        )
+        from torch._inductor.runtime.triton_heuristics import CachingAutotuner
+        from torch._inductor.utils import triton_version_uses_attrs_dict
+
+        meta = {
+            "signature": {
+                "in_ptr0": "*fp32",
+                "out_ptr0": "*fp32",
+                "xnumel": "i32",
+            },
+            "device": DeviceProperties.create(torch.device(GPU_TYPE)),
+            "configs": [AttrsDescriptorWrapper(divisible_by_16=(0, 1), equal_to_1=())],
+            "constants": {},
+        }
+        if triton_version_uses_attrs_dict():
+            meta["signature"]["XBLOCK"] = "constexpr"
+
+        @triton.jit
+        def _add_kernel(in_ptr0, out_ptr0, xnumel, XBLOCK: tl.constexpr):
+            pid = tl.program_id(0)
+            offsets = pid * XBLOCK + tl.arange(0, XBLOCK)
+            mask = offsets < xnumel
+            x = tl.load(in_ptr0 + offsets, mask=mask)
+            tl.store(out_ptr0 + offsets, x + 1.0, mask=mask)
+
+        autotuner = CachingAutotuner(
+            fn=_add_kernel,
+            triton_meta=meta,
+            configs=[triton.Config({"XBLOCK": 128})],
+            save_cache_hook=False,
+            mutated_arg_names=[],
+            optimize_mem=True,
+            heuristic_type=HeuristicType.POINTWISE,
+            inductor_meta={"grid_type": "Grid1D"},
+        )
+
+        xnumel = 256
+        inp = torch.randn(xnumel, device=GPU_TYPE)
+        out = torch.empty_like(inp)
+        autotuner.run(inp, out, xnumel, stream=torch.cuda.current_stream().cuda_stream)
+        self.assertEqual(out, inp + 1.0)
+
+        # Inject a launcher key into benchmark_failure_reasons — this is how
+        # the leak happens in production (launcher used as dict key).
+        self.assertTrue(len(autotuner.launchers) > 0)
+        launcher_fn = autotuner.launchers[0]
+        autotuner.benchmark_failure_reasons[launcher_fn] = "test"
+
+        # The autotuner is unpicklable: __getstate__ asserts launchers are
+        # empty, and benchmark_failure_reasons holds an exec'd launcher key.
+        with self.assertRaises((pickle.PicklingError, AttributeError, AssertionError)):
+            pickle.dumps(autotuner)
+
+        # prepare_for_pickle clears benchmark_failure_reasons (and other
+        # unpicklable fields), making pickle succeed.
+        autotuner.prepare_for_pickle()
+        self.assertEqual(autotuner.benchmark_failure_reasons, {})
+        pickle.dumps(autotuner)
+
+    def test_nested_tensor_subclass_cache_key(self):
+        ctx = multiprocessing.get_context("spawn")
+        results = []
+        for _ in range(2):
+            queue = ctx.Queue()
+            p = ctx.Process(
+                target=_subprocess_gen_nested_subclass_cache_key,
+                args=(queue,),
+            )
+            p.start()
+            p.join()
+
+            if p.exitcode != 0:
+                self.skipTest(f"Subprocess exited with code {p.exitcode}")
+
+            results.append(queue.get())
+
+        # Cache keys from two different processes should be identical
+        self.assertEqual(
+            results[0],
+            results[1],
+            "Nested tensor subclass cache keys should be identical across processes",
+        )
+
+    @unittest.skipIf(not torch.distributed.is_available(), "requires distributed")
+    def test_dtensor_different_process_cache_key(self):
+        """
+        Test that DTensor cache keys are consistent across different processes.
+
+        This is a critical test for warm start cache hits. The cache key must be
+        deterministic and not depend on process-specific values like memory addresses
+        or object ids that would differ between processes.
+        """
+        ctx = multiprocessing.get_context("spawn")
+        results = []
+        for _ in range(2):
+            queue = ctx.Queue()
+            p = ctx.Process(
+                target=_subprocess_gen_dtensor_cache_key,
+                args=(queue,),
+            )
+            p.start()
+            p.join()
+
+            if p.exitcode != 0:
+                self.skipTest(f"Subprocess exited with code {p.exitcode}")
+
+            results.append(queue.get())
+
+        # Cache keys from two different processes should be identical
+        self.assertEqual(
+            results[0],
+            results[1],
+            "DTensor cache keys should be identical across processes",
+        )
+
+    @unittest.skipIf(not torch.distributed.is_available(), "requires distributed")
+    def test_dtensor_different_placements_different_cache_key(self):
+        from torch.distributed.device_mesh import init_device_mesh
+        from torch.distributed.tensor import DTensor, Replicate, Shard
+
+        with _fake_process_group():
+            mesh = init_device_mesh("cpu", (2,))
+
+            local_tensor = torch.zeros(4, 4)
+            dtensor_replicate = DTensor.from_local(local_tensor, mesh, [Replicate()])
+            dtensor_shard = DTensor.from_local(local_tensor, mesh, [Shard(0)])
+
+            def fn(x):
+                return x.sin()
+
+            # Get dynamo output for replicate
+            torch._dynamo.reset()
+            fx_graph = None
+
+            def compiler(gm, inputs, **kwargs):
+                nonlocal fx_graph
+                fx_graph = gm
+                return gm
+
+            g = torch.compile(fn, backend=compiler, fullgraph=True)
+            g(dtensor_replicate)
+
+            pickler = AOTAutogradCachePickler(fx_graph)
+            cache_key_replicate = pickler.get_hash(dtensor_replicate)
+            cache_key_shard = pickler.get_hash(dtensor_shard)
+
+            # Different placements should produce different cache keys
+            self.assertNotEqual(
+                cache_key_replicate,
+                cache_key_shard,
+                "DTensor with Replicate() should have different cache key than Shard(0)",
+            )
+
+
+def _subprocess_gen_dtensor_cache_key(queue):
+    """
+    Subprocess helper to generate a DTensor cache key.
+    Must be at module level for multiprocessing to work.
+    """
+    import torch
+    import torch._dynamo
+    import torch.distributed as c10d
+    from torch._functorch._aot_autograd.autograd_cache import AOTAutogradCachePickler
+    from torch.distributed.device_mesh import init_device_mesh
+    from torch.distributed.tensor import DTensor, Replicate
+    from torch.testing._internal.distributed.fake_pg import FakeStore
+
+    fake_store = FakeStore()
+    c10d.init_process_group("fake", store=fake_store, rank=0, world_size=2)
+    try:
+        mesh = init_device_mesh("cpu", (2,))
+
+        # Create DTensor
+        local_tensor = torch.zeros(4, 4)
+        dtensor = DTensor.from_local(local_tensor, mesh, [Replicate()])
+
+        def fn(x):
+            return x.sin()
+
+        # Get dynamo output
+        torch._dynamo.reset()
+        fx_graph = None
+
+        def compiler(gm, inputs, **kwargs):
+            nonlocal fx_graph
+            fx_graph = gm
+            return gm
+
+        g = torch.compile(fn, backend=compiler, fullgraph=True)
+        g(dtensor)
+
+        pickler = AOTAutogradCachePickler(fx_graph)
+        cache_key = pickler.get_hash(dtensor)
+
+        queue.put(cache_key)
+    finally:
+        c10d.destroy_process_group()
+
+
+def _subprocess_gen_nested_subclass_cache_key(queue):
+    import torch
+    import torch._dynamo
+    from torch._functorch._aot_autograd.autograd_cache import AOTAutogradCachePickler
+    from torch.testing._internal.two_tensor import TwoTensor
+
+    inner1 = TwoTensor(torch.zeros(4, 4), torch.zeros(4, 4))
+    inner2 = TwoTensor(torch.zeros(4, 4), torch.zeros(4, 4))
+    nested_two_tensor = TwoTensor(inner1, inner2)
+
+    def fn(x):
+        return x.sin()
+
+    # Get dynamo output
+    torch._dynamo.reset()
+    fx_graph = None
+
+    def compiler(gm, inputs, **kwargs):
+        nonlocal fx_graph
+        fx_graph = gm
+        return gm
+
+    g = torch.compile(fn, backend=compiler, fullgraph=True)
+    g(nested_two_tensor)
+
+    pickler = AOTAutogradCachePickler(fx_graph)
+    cache_key = pickler.get_hash(nested_two_tensor)
+
+    queue.put(cache_key)
+
 
 def _policy_save_mm(ctx, op, *args, **kwargs):
     if op == torch.ops.aten.mm.default:
@@ -2955,7 +3824,7 @@ def _create_sac_ctx_fn(policy, cache_hash=None):
     return ctx_fn
 
 
-class HOPCacheTests(torch._dynamo.test_case.TestCase):
+class HOPCacheTests(CacheKeyEquivalenceMixin, torch._dynamo.test_case.TestCase):
     def setUp(self):
         super().setUp()
         counters.clear()
@@ -3135,6 +4004,329 @@ def fn_add(x, y):
             self.assertEqual(counters["aot_autograd"]["autograd_cache_hit"], 0)
 
 
+@instantiate_parametrized_tests
+class CacheKeyAPITests(torch._dynamo.test_case.TestCase):
+    """
+    Verify that the exposed autograd_cache_key API function produces the same
+    cache key that the internal compilation pipeline would.
+    """
+
+    def setUp(self):
+        super().setUp()
+        counters.clear()
+        torch._dynamo.reset()
+
+    def tearDown(self):
+        super().tearDown()
+        counters.clear()
+        torch._dynamo.reset()
+
+    def _compile_and_capture(self, fn, *args):
+        """Run torch.compile with the inductor backend and capture the
+        ground-truth cache key, graph module, and example inputs produced
+        during compilation."""
+        torch._dynamo.reset()
+        real_compile_fx = torch._inductor.compile_fx.compile_fx
+        captured_key = None
+        captured_gm = None
+        captured_example_inputs = None
+
+        def capturing_compile_fx(model_, example_inputs_, **kwargs):
+            nonlocal captured_gm, captured_example_inputs
+            captured_gm = model_
+            captured_example_inputs = list(example_inputs_)
+            return real_compile_fx(model_, example_inputs_, **kwargs)
+
+        real_cache_key = autograd_cache.autograd_cache_key
+
+        def capturing_cache_key(mod, ei, config, compiler_config_extra=None):
+            nonlocal captured_key
+            result = real_cache_key(mod, ei, config, compiler_config_extra)
+            captured_key = result[0]
+            return result
+
+        with fresh_cache():
+            with (
+                patch(
+                    "torch._inductor.compile_fx.compile_fx",
+                    side_effect=capturing_compile_fx,
+                ),
+                patch(
+                    "torch._functorch._aot_autograd.autograd_cache.autograd_cache_key",
+                    side_effect=capturing_cache_key,
+                ),
+            ):
+                torch.compile(fn, fullgraph=True)(*args)
+
+        self.assertIsNotNone(captured_key)
+        return captured_key, captured_gm, captured_example_inputs
+
+    def _tracing_context(self):
+        """Create a TracingContext with a ShapeEnv, as required by
+        _check_can_cache inside the cache key computation."""
+        fake_mode = FakeTensorMode(shape_env=ShapeEnv())
+        return torch._guards.tracing(TracingContext(fake_mode))
+
+    def _aot_autograd_cache_key(self, fx_graph, example_inputs):
+        """Call aot_autograd.autograd_cache_key with the same config and
+        tracing context that compile_fx would use."""
+        from torch._inductor.compile_fx import create_compiler_config_extra
+        from torch._inductor.decomposition import select_decomp_table
+
+        compiler_config_extra = create_compiler_config_extra(fx_graph)
+        decompositions = select_decomp_table()
+
+        # Match the config patches that compile_fx applies during compilation
+        # so that the autograd_config snapshot is identical.
+        with (
+            self._tracing_context(),
+            functorch_config.patch(
+                unlift_effect_tokens=True,
+                selective_decompose=inductor_config.selective_decompose,
+            ),
+        ):
+            return aot_autograd.autograd_cache_key(
+                fx_graph,
+                example_inputs,
+                ignore_shape_env=True,
+                decompositions=decompositions,
+                compiler_config_extra=compiler_config_extra,
+                keep_inference_input_mutations=True,
+            )
+
+    def _compile_fx_cache_key(self, fx_graph, example_inputs):
+        """Call compile_fx.autograd_cache_key, the higher-level API that
+        handles decompositions and config patching internally."""
+        with self._tracing_context():
+            return compile_fx.autograd_cache_key(
+                fx_graph,
+                example_inputs,
+                ignore_shape_env=True,
+            )
+
+    @requires_triton()
+    @inductor_config.patch("fx_graph_remote_cache", False)
+    @inductor_config.patch("fx_graph_cache", True)
+    @functorch_config.patch(
+        {"enable_autograd_cache": True, "strict_autograd_cache": True}
+    )
+    def test_cache_key_matches_torch_compile(self):
+        """autograd_cache_key matches the real compile_fx pipeline."""
+
+        def fn(x):
+            return x.sin() + x.cos()
+
+        x = torch.randn(4, 4)
+        gt_key, fx_graph, example_inputs = self._compile_and_capture(fn, x)
+        api_key, _ = self._aot_autograd_cache_key(fx_graph, example_inputs)
+        cfx_key, _ = self._compile_fx_cache_key(fx_graph, example_inputs)
+
+        self.assertEqual(gt_key, api_key)
+        self.assertEqual(gt_key, cfx_key)
+
+    @requires_triton()
+    @inductor_config.patch("fx_graph_remote_cache", False)
+    @inductor_config.patch("fx_graph_cache", True)
+    @functorch_config.patch(
+        {"enable_autograd_cache": True, "strict_autograd_cache": True}
+    )
+    def test_cache_key_stable_across_calls(self):
+        """Cache keys should be deterministic."""
+
+        def fn(x):
+            return x.sin() + x.cos()
+
+        x = torch.randn(4, 4)
+        gt_key, fx_graph, example_inputs = self._compile_and_capture(fn, x)
+        key1, _ = self._aot_autograd_cache_key(fx_graph, example_inputs)
+        key2, _ = self._aot_autograd_cache_key(fx_graph, example_inputs)
+        cfx_key1, _ = self._compile_fx_cache_key(fx_graph, example_inputs)
+        cfx_key2, _ = self._compile_fx_cache_key(fx_graph, example_inputs)
+        self.assertEqual(gt_key, key1)
+        self.assertEqual(key1, key2)
+        self.assertEqual(gt_key, cfx_key1)
+        self.assertEqual(cfx_key1, cfx_key2)
+
+    @requires_triton()
+    @inductor_config.patch("fx_graph_remote_cache", False)
+    @inductor_config.patch("fx_graph_cache", True)
+    @functorch_config.patch(
+        {"enable_autograd_cache": True, "strict_autograd_cache": True}
+    )
+    def test_cache_key_varies_with_different_inputs(self):
+        """Different input dtypes should produce different cache keys."""
+
+        def fn(x):
+            return x.sin() + x.cos()
+
+        x_float = torch.randn(4, 4)
+        x_double = torch.randn(4, 4, dtype=torch.float64)
+
+        gt_key1, graph1, inputs1 = self._compile_and_capture(fn, x_float)
+        gt_key2, graph2, inputs2 = self._compile_and_capture(fn, x_double)
+        api_key1, _ = self._aot_autograd_cache_key(graph1, inputs1)
+        api_key2, _ = self._aot_autograd_cache_key(graph2, inputs2)
+        cfx_key1, _ = self._compile_fx_cache_key(graph1, inputs1)
+        cfx_key2, _ = self._compile_fx_cache_key(graph2, inputs2)
+        self.assertEqual(gt_key1, api_key1)
+        self.assertEqual(gt_key2, api_key2)
+        self.assertEqual(gt_key1, cfx_key1)
+        self.assertEqual(gt_key2, cfx_key2)
+        self.assertNotEqual(gt_key1, gt_key2)
+
+    @requires_triton()
+    @inductor_config.patch("fx_graph_remote_cache", False)
+    @inductor_config.patch("fx_graph_cache", True)
+    @functorch_config.patch(
+        {"enable_autograd_cache": True, "strict_autograd_cache": True}
+    )
+    def test_cache_key_varies_with_different_graphs(self):
+        """Different model functions should produce different cache keys."""
+
+        def fn1(x):
+            return x.sin() + x.cos()
+
+        def fn2(x):
+            return x.exp() * x.log()
+
+        x = torch.randn(4, 4)
+
+        gt_key1, graph1, inputs1 = self._compile_and_capture(fn1, x)
+        gt_key2, graph2, inputs2 = self._compile_and_capture(fn2, x)
+        api_key1, _ = self._aot_autograd_cache_key(graph1, inputs1)
+        api_key2, _ = self._aot_autograd_cache_key(graph2, inputs2)
+        cfx_key1, _ = self._compile_fx_cache_key(graph1, inputs1)
+        cfx_key2, _ = self._compile_fx_cache_key(graph2, inputs2)
+        self.assertEqual(gt_key1, api_key1)
+        self.assertEqual(gt_key2, api_key2)
+        self.assertEqual(gt_key1, cfx_key1)
+        self.assertEqual(gt_key2, cfx_key2)
+        self.assertNotEqual(gt_key1, gt_key2)
+
+    @requires_triton()
+    @inductor_config.patch("fx_graph_remote_cache", False)
+    @inductor_config.patch("fx_graph_cache", True)
+    @functorch_config.patch(
+        {"enable_autograd_cache": True, "strict_autograd_cache": True}
+    )
+    def test_cache_key_for_module(self):
+        """autograd_cache_key matches compilation for nn.Modules."""
+
+        class MyModule(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.linear = torch.nn.Linear(4, 4)
+
+            def forward(self, x):
+                return self.linear(x).relu()
+
+        mod = MyModule()
+        x = torch.randn(4, 4)
+        gt_key, fx_graph, example_inputs = self._compile_and_capture(mod, x)
+        api_key, _ = self._aot_autograd_cache_key(fx_graph, example_inputs)
+        cfx_key, _ = self._compile_fx_cache_key(fx_graph, example_inputs)
+
+        self.assertEqual(gt_key, api_key)
+        self.assertEqual(gt_key, cfx_key)
+
+    @requires_triton()
+    @inductor_config.patch("fx_graph_remote_cache", False)
+    @inductor_config.patch("fx_graph_cache", True)
+    @functorch_config.patch(
+        {"enable_autograd_cache": True, "strict_autograd_cache": True}
+    )
+    def test_cache_key_for_multiple_outputs(self):
+        """autograd_cache_key matches compilation for multiple outputs."""
+
+        def fn(x):
+            return x.sin(), x.cos()
+
+        x = torch.randn(4, 4)
+        gt_key, fx_graph, example_inputs = self._compile_and_capture(fn, x)
+        api_key, _ = self._aot_autograd_cache_key(fx_graph, example_inputs)
+        cfx_key, _ = self._compile_fx_cache_key(fx_graph, example_inputs)
+
+        self.assertEqual(gt_key, api_key)
+        self.assertEqual(gt_key, cfx_key)
+
+    @requires_triton()
+    @inductor_config.patch("fx_graph_remote_cache", False)
+    @inductor_config.patch("fx_graph_cache", True)
+    @functorch_config.patch(
+        {"enable_autograd_cache": True, "strict_autograd_cache": True}
+    )
+    @parametrize(
+        "dynamic_shapes",
+        ("from_example_inputs", "from_graph", "from_tracing_context"),
+    )
+    @parametrize("aot", (False, True))
+    def test_standalone_compile_cache_key_matches_standalone_compile(
+        self, dynamic_shapes, aot
+    ):
+        """standalone_compile.autograd_cache_key matches the real key
+        produced when running standalone_compile end-to-end."""
+        from contextlib import nullcontext
+
+        from torch._inductor.standalone_compile import standalone_compile
+
+        # Build an FX graph directly: sin(x) + cos(x)
+        fake_mode = FakeTensorMode(shape_env=ShapeEnv())
+        with fake_mode:
+            fake_x = torch.randn(4, 4)
+        graph = torch.fx.Graph()
+        x = graph.placeholder("x")
+        x.meta["example_value"] = fake_x
+        sin = graph.call_method("sin", (x,))
+        sin.meta["example_value"] = fake_x.sin()
+        cos = graph.call_method("cos", (x,))
+        cos.meta["example_value"] = fake_x.cos()
+        add = graph.call_function(operator.add, (sin, cos))
+        add.meta["example_value"] = fake_x.sin() + fake_x.cos()
+        graph.output((add,))
+        fx_graph = GraphModule(torch.nn.Module(), graph)
+        example_inputs = [torch.randn(4, 4)]
+
+        # from_tracing_context requires an ambient TracingContext, simulating
+        # a call from within a torch.compile backend.
+        outer_context = (
+            torch._guards.tracing(torch._guards.TracingContext(fake_mode))
+            if dynamic_shapes == "from_tracing_context"
+            else nullcontext()
+        )
+
+        # Run standalone_compile and capture the real cache key.
+        real_cache_key_fn = autograd_cache.autograd_cache_key
+        captured_key = None
+
+        def capturing_cache_key(mod, ei, config, compiler_config_extra=None):
+            nonlocal captured_key
+            captured_key, _ = real_cache_key_fn(mod, ei, config, compiler_config_extra)
+            return captured_key, _
+
+        with outer_context:
+            with patch(
+                "torch._functorch._aot_autograd.autograd_cache.autograd_cache_key",
+                side_effect=capturing_cache_key,
+            ):
+                standalone_compile(
+                    fx_graph,
+                    example_inputs,
+                    dynamic_shapes=dynamic_shapes,
+                    options={},
+                    aot=aot,
+                )
+
+            self.assertIsNotNone(captured_key)
+
+            sc_key, _ = standalone_compile_autograd_cache_key(
+                fx_graph,
+                example_inputs,
+                dynamic_shapes=dynamic_shapes,
+                aot=aot,
+            )
+            self.assertEqual(captured_key, sc_key)
+
+
 if __name__ == "__main__":
     from torch._dynamo.test_case import run_tests
 
diff --git a/test/dynamo/test_aot_compile.py b/test/dynamo/test_aot_compile.py
index 32783a87bf0b1..eb1d653313245 100644
--- a/test/dynamo/test_aot_compile.py
+++ b/test/dynamo/test_aot_compile.py
@@ -39,10 +39,8 @@
 from torch.fx._graph_pickler import GraphPickler
 from torch.fx.passes.regional_inductor import regional_inductor
 from torch.nn.attention.flex_attention import create_block_mask, flex_attention
-from torch.testing._internal.common_utils import (
-    instantiate_parametrized_tests,
-    TEST_CUDA,
-)
+from torch.testing._internal.common_utils import instantiate_parametrized_tests
+from torch.testing._internal.inductor_utils import GPU_TYPE, HAS_GPU
 from torch.utils.checkpoint import checkpoint
 
 
@@ -61,6 +59,25 @@ def aot_eager_regional_inductor():
     )
 
 
+class SingleCondModel(torch.nn.Module):
+    def __init__(self, d=64):
+        super().__init__()
+        self.fc1 = torch.nn.Linear(d, d)
+        self.fc2 = torch.nn.Linear(d, d)
+
+    def forward(self, x):
+        x = self.fc1(x)
+
+        def true_fn(x):
+            return x * 2.0
+
+        def false_fn(x):
+            return x * 3.0
+
+        x = torch.cond(x.shape[0] < 32, true_fn, false_fn, (x,))
+        return self.fc2(x)
+
+
 class MooType:
     def __init__(self, x):
         self.x = x
@@ -404,7 +421,10 @@ def fn(x, y):
                 raise AssertionError("Guard check should have failed")
             compiled_fn.disable_guard_check()
             actual = compiled_fn(*inputs)
-            assert torch.allclose(actual, expected)
+            if not torch.allclose(actual, expected):
+                raise AssertionError(
+                    f"Expected tensors to be close, got {actual} vs {expected}"
+                )
         finally:
             torch.set_grad_enabled(prev_grad)
 
@@ -435,7 +455,10 @@ def target_fn(x, y):
         with torch.no_grad():
             actual = compiled_fn(*inputs)
             expected = target_fn(*inputs)
-            assert torch.allclose(actual, expected)
+            if not torch.allclose(actual, expected):
+                raise AssertionError(
+                    f"Expected tensors to be close, got {actual} vs {expected}"
+                )
 
 
 def _subprocess_aot_compile_module():
@@ -473,7 +496,10 @@ def eval_mode(mdl):
                 args=(torch.randn(3, 3),), kwargs={}, contexts=[train_mode(model)]
             ),
         ]
-        assert isinstance(model, torch._dynamo.eval_frame.OptimizedModule)
+        if not isinstance(model, torch._dynamo.eval_frame.OptimizedModule):
+            raise AssertionError(
+                f"Expected OptimizedModule, got {type(model).__name__}"
+            )
         model._aot_compile(inputs)
 
         with torch.compiler.set_stance("fail_on_recompile"):
@@ -481,7 +507,10 @@ def eval_mode(mdl):
             eager_inputs = (torch.randn(3, 3),)
             expected = mod(*eager_inputs)
             actual = model(*eager_inputs)
-            assert torch.allclose(expected, actual)
+            if not torch.allclose(expected, actual):
+                raise AssertionError(
+                    f"Expected tensors to be close, got {actual} vs {expected}"
+                )
             model.train()
             expected.sum().backward()
 
@@ -497,7 +526,10 @@ def eval_mode(mdl):
                     "guard_filter_fn": torch.compiler.skip_guard_on_globals_unsafe,
                 },
             )
-            assert isinstance(model, torch._dynamo.eval_frame.OptimizedModule)
+            if not isinstance(model, torch._dynamo.eval_frame.OptimizedModule):
+                raise AssertionError(
+                    f"Expected OptimizedModule, got {type(model).__name__}"
+                )
             with open(path, "rb") as f:
                 data = f.read()
                 model._load_aot_compiled_module(data)
@@ -507,7 +539,10 @@ def eval_mode(mdl):
                 eager_inputs = (torch.randn(3, 3),)
                 expected = mod(*eager_inputs)
                 actual = model(*eager_inputs)
-                assert torch.allclose(expected, actual)
+                if not torch.allclose(expected, actual):
+                    raise AssertionError(
+                        f"Expected tensors to be close, got {actual} vs {expected}"
+                    )
 
 
 class RedistributeModel(torch.nn.Module):
@@ -643,7 +678,7 @@ def test_decorated_function_aot(self):
         def check_inputs(fn):
             def _fn(*args, **kwargs):
                 for arg in args:
-                    assert arg.shape[0] > 1
+                    assert arg.shape[0] > 1  # noqa: S101
 
                 return fn(*args, **kwargs)
 
@@ -690,12 +725,31 @@ def fn(x, y):
             actual = compiled_fn(*inputs)
             self.assertEqual(expected, actual)
 
+    def test_aot_eager_backend(self):
+        def fn(x, y):
+            return x + y
+
+        compiled_fn = torch.compile(
+            fn, fullgraph=True, backend="aot_eager"
+        ).aot_compile(((torch.randn(3, 4), torch.randn(3, 4)), {}))
+        inputs = (torch.randn(3, 4), torch.randn(3, 4))
+        expected = fn(*inputs)
+        actual = compiled_fn(*inputs)
+        self.assertEqual(expected, actual)
+        compiled_fn.save_compiled_function(self.path())
+        torch._dynamo.reset()
+        with torch.compiler.set_stance("fail_on_recompile"):
+            with open(self.path(), "rb") as f:
+                compiled_fn = torch.compiler.load_compiled_function(f)
+            actual = compiled_fn(*inputs)
+            self.assertEqual(expected, actual)
+
     def test_decorated_function_with_functools_wrap_aot(self):
         def check_inputs(fn):
             @functools.wraps(fn)
             def _fn(*args, **kwargs):
                 for arg in args:
-                    assert arg.shape[0] > 1
+                    assert arg.shape[0] > 1  # noqa: S101
 
                 return fn(*args, **kwargs)
 
@@ -812,7 +866,7 @@ def test_guard_filter_override_aot(self):
         def check_inputs(fn):
             def _fn(*args, **kwargs):
                 for arg in args:
-                    assert arg.shape[0] > 1
+                    assert arg.shape[0] > 1  # noqa: S101
 
                 return fn(*args, **kwargs)
 
@@ -881,7 +935,8 @@ def test_aot_module_simplified_serializable_autograd(self):
                 torch._dynamo.aot_compile.BundledAOTAutogradSerializableCallable,
             )
         )
-        assert hasattr(backend_result.compiled_fn, "serialize")
+        if not hasattr(backend_result.compiled_fn, "serialize"):
+            raise AssertionError("Expected compiled_fn to have 'serialize' attribute")
         self.assertIsNotNone(backend_result.compiled_fn.serialize)
 
     def test_aot_compile_portable_guards_unsafe(self):
@@ -919,9 +974,47 @@ def fn(x):
                 torch._dynamo.aot_compile.BundledAOTAutogradSerializableCallable,
             )
         )
-        assert hasattr(backend_result.compiled_fn, "serialize")
+        if not hasattr(backend_result.compiled_fn, "serialize"):
+            raise AssertionError("Expected compiled_fn to have 'serialize' attribute")
         self.assertIsNotNone(backend_result.compiled_fn.serialize)
 
+    def test_aot_cache_predicate_not_pickleable(self):
+        import torch._functorch.config as functorch_config
+        import torch._inductor.config as inductor_config
+
+        model = SingleCondModel().eval()
+
+        old_cacheable = torch.ops.higher_order.cond._cacheable
+        torch.ops.higher_order.cond._cacheable = True
+        try:
+            with (
+                functorch_config.patch(
+                    enable_autograd_cache=True,
+                    force_non_lazy_backward_lowering=True,
+                    strict_autograd_cache=True,
+                ),
+                inductor_config.patch(
+                    fx_graph_cache=True,
+                    fx_graph_remote_cache=False,
+                ),
+            ):
+                compiled = torch.compile(model, backend="inductor", dynamic=True)
+
+                # Test both branches of the cond predicate (x.shape[0] < 32).
+                for batch_size in (16, 64):
+                    inp = torch.randn(batch_size, 64)
+                    expected = model(inp)
+                    actual = compiled(inp)
+                    self.assertEqual(expected, actual)
+
+                # If the SymBool predicate is not pickleable, the FxGraphCache
+                # silently bypasses instead of caching. Assert no bypass occurred.
+                self.assertEqual(
+                    torch._dynamo.utils.counters["inductor"]["fxgraph_cache_bypass"], 0
+                )
+        finally:
+            torch.ops.higher_order.cond._cacheable = old_cacheable
+
     def test_fullgraph_capture_with_pytree_module(self):
         from torch._dynamo.functional_export import dynamo_graph_capture_for_export
 
@@ -1064,9 +1157,9 @@ def fn(x, y=1):
         actual = compiled_fn(*inputs)
         self.assertEqual(expected, actual)
 
-    @unittest.skipIf(not TEST_CUDA, "requires cuda")
+    @unittest.skipIf(not HAS_GPU, "requires gpu")
     def test_aot_compile_with_aoti(self):
-        with torch.device("cuda"):
+        with torch.device(GPU_TYPE):
             from torch._dynamo.hooks import Hooks
 
             def fn(x, y):
@@ -1092,9 +1185,9 @@ def make_inputs():
             actual = compiled_fn(*test_inputs)
             self.assertEqual(expected, actual)
 
-    @unittest.skipIf(not TEST_CUDA, "requires cuda")
+    @unittest.skipIf(not HAS_GPU, "requires gpu")
     def test_aot_compile_with_aoti_module(self):
-        with torch.device("cuda"):
+        with torch.device(GPU_TYPE):
             from torch._dynamo.hooks import Hooks
 
             mod = SimpleLinearModule()
@@ -1126,9 +1219,9 @@ def get_grads(m: torch.nn.Module):
             actual.sum().backward()
             self.assertEqual(get_grads(original_mod), expected_grads)
 
-    @unittest.skipIf(not TEST_CUDA, "requires cuda")
+    @unittest.skipIf(not HAS_GPU, "requires gpu")
     def test_aot_compile_with_aoti_torch_compile(self):
-        with torch.device("cuda"):
+        with torch.device(GPU_TYPE):
 
             def fn(x, y):
                 return x + y
@@ -1315,9 +1408,30 @@ def make_inputs():
         actual = compiled_fn(*test_inputs)
         self.assertEqual(expected.x, actual.x)
 
-    @unittest.skipIf(not TEST_CUDA, "requires cuda")
+    def test_builtins_dict_survives_serialization(self):
+        """Test that __builtins_dict__ is preserved through serialize/deserialize."""
+
+        def fn(x):
+            return x + 1, type
+
+        x = torch.randn(4)
+        compiled_fn = torch.compile(fn, fullgraph=True).aot_compile(((x,), {}))
+
+        # Save and reload without f_globals
+        compiled_fn.save_compiled_function(self.path())
+        with open(self.path(), "rb") as f:
+            loaded_fn = torch.compiler.load_compiled_function(
+                f, f_globals=fn.__globals__
+            )
+
+        expected = fn(x)
+        actual = loaded_fn(x)
+        self.assertEqual(expected[0], actual[0])
+        self.assertEqual(expected[1], actual[1])
+
+    @unittest.skipIf(not HAS_GPU, "requires gpu")
     def test_cross_aot_compile(self):
-        """Test cross-compilation using fake cuda tensors and backward correctness"""
+        """Test cross-compilation using fake tensors and backward correctness"""
         from torch._subclasses.fake_tensor import FakeTensorMode
 
         def fn(x, y):
@@ -1325,8 +1439,8 @@ def fn(x, y):
 
         with FakeTensorMode(allow_non_fake_inputs=True):
             fake_inputs = (
-                torch.randn(3, 4, device="cuda", requires_grad=True),
-                torch.randn(3, 4, device="cuda", requires_grad=True),
+                torch.randn(3, 4, device=GPU_TYPE, requires_grad=True),
+                torch.randn(3, 4, device=GPU_TYPE, requires_grad=True),
             )
         compiled_fn = torch.compile(
             fn,
@@ -1340,8 +1454,8 @@ def fn(x, y):
             loaded_fn = torch.compiler.load_compiled_function(f)
 
         inputs = (
-            torch.randn(3, 4, device="cuda", requires_grad=True),
-            torch.randn(3, 4, device="cuda", requires_grad=True),
+            torch.randn(3, 4, device=GPU_TYPE, requires_grad=True),
+            torch.randn(3, 4, device=GPU_TYPE, requires_grad=True),
         )
         expected = fn(*inputs)
         actual = loaded_fn(*inputs)
@@ -1365,7 +1479,7 @@ def fn(x, y):
             self.assertEqual(eg, cg)
 
     @unittest.skipIf(not c10d.is_available(), "requires c10d")
-    @unittest.skipIf(not TEST_CUDA, "requires cuda")
+    @unittest.skipIf(not HAS_GPU, "requires gpu")
     def test_cross_compile_realistic_transformer_model(self):
         """
         Test cross-compilation with transformer model with DTensors,
@@ -1417,7 +1531,7 @@ def init_weights_deterministic(module: nn.Module, seed: int = 42) -> None:
             This ensures reproducible results across eager and compiled runs.
             """
             torch.manual_seed(seed)
-            torch.cuda.manual_seed(seed)
+            getattr(torch, GPU_TYPE).manual_seed(seed)
             for name, param in module.named_parameters():
                 if param.requires_grad:
                     local_param = (
@@ -1433,7 +1547,7 @@ def init_weights_deterministic(module: nn.Module, seed: int = 42) -> None:
 
         try:
             rank = c10d.get_rank()
-            device = torch.device(f"cuda:{rank}")
+            device = torch.device(f"{GPU_TYPE}:{rank}")
             vocab_size = 1000
             embed_dim = 256
             num_heads = 8
@@ -1444,7 +1558,7 @@ def init_weights_deterministic(module: nn.Module, seed: int = 42) -> None:
             seq_len = 16
 
             device_mesh = init_device_mesh(
-                "cuda",
+                GPU_TYPE,
                 (1,),
                 mesh_dim_names=("dp",),
             )
@@ -1481,9 +1595,7 @@ def init_weights_deterministic(module: nn.Module, seed: int = 42) -> None:
 
             from torch._dynamo.functional_export import dynamo_graph_capture_for_export
 
-            gm = dynamo_graph_capture_for_export(model, restore_state_dict=True)(
-                input_ids_dt
-            )
+            gm = dynamo_graph_capture_for_export(model)(input_ids_dt)
 
             fake_mode = gm.meta["fake_mode"]
 
diff --git a/test/dynamo/test_autograd_function.py b/test/dynamo/test_autograd_function.py
index aee11571bcee7..f41434b8dc9f1 100644
--- a/test/dynamo/test_autograd_function.py
+++ b/test/dynamo/test_autograd_function.py
@@ -615,13 +615,13 @@ def forward(self, l_weird_b: "f32[]", l_weird_c: "f32[]", l_x_: "f32[]", l_z_: "
 
             mul: "f32[]" = l_weird_b * l_weird_c
             clone: "f32[]" = l_x_.clone();  l_x_ = None
-            outs: "f32[]" = mul * clone
+            outs: "f32[]" = mul * clone;  mul = clone = None
 
             _set_grad_enabled_1 = torch._C._set_grad_enabled(True);  _set_grad_enabled_1 = None
-            return ((outs, mul, clone), (l_weird_b, l_weird_c))
+            return ((outs,), (l_weird_b, l_weird_c))
 
     class bwd_body_0(torch.nn.Module):
-        def forward(self, grad: "f32[]", unused_0, unused_1, l_weird_b: "f32[]", l_weird_c: "f32[]"):
+        def forward(self, grad: "f32[]", l_weird_b: "f32[]", l_weird_c: "f32[]"):
             _set_grad_enabled = torch._C._set_grad_enabled(False);  _set_grad_enabled = None
 
             mul: "f32[]" = grad * l_weird_b;  l_weird_b = None
@@ -1050,7 +1050,7 @@ def forward(ctx, x, weight):
             def backward(ctx, grad_output):
                 x, weight = ctx.saved_tensors
                 grad_x = grad_output.matmul(weight)
-                assert grad_x.is_contiguous()
+                assert grad_x.is_contiguous()  # noqa: S101
                 grad_weight = grad_output.transpose(0, 1).matmul(x)
 
                 return grad_x, grad_weight
@@ -1086,7 +1086,7 @@ def forward(ctx, x, weight):
 
             @staticmethod
             def backward(ctx, grad_output):
-                assert grad_output.is_contiguous()
+                assert grad_output.is_contiguous()  # noqa: S101
                 x, weight = ctx.saved_tensors
                 grad_x = grad_output.matmul(weight)
                 grad_weight = grad_output.transpose(0, 1).matmul(x)
@@ -1136,13 +1136,13 @@ def forward(self, l_weight_: "f32[4, 3]", l_x_: "f32[5, 3]"):
             _set_grad_enabled = torch._C._set_grad_enabled(False);  _set_grad_enabled = None
 
             t: "f32[3, 4]" = l_weight_.t()
-            y: "f32[5, 4]" = l_x_.matmul(t)
+            y: "f32[5, 4]" = l_x_.matmul(t);  t = None
 
             _set_grad_enabled_1 = torch._C._set_grad_enabled(True);  _set_grad_enabled_1 = None
-            return ((y, t), (l_weight_, l_x_))
+            return ((y,), (l_weight_, l_x_))
 
     class bwd_body_0(torch.nn.Module):
-        def forward(self, y: "f32[5, 4]", unused_0, l_weight_: "f32[4, 3]", l_x_: "f32[5, 3]"):
+        def forward(self, y: "f32[5, 4]", l_weight_: "f32[4, 3]", l_x_: "f32[5, 3]"):
             _set_grad_enabled = torch._C._set_grad_enabled(False);  _set_grad_enabled = None
 
             contiguous: "f32[5, 4]" = y.contiguous();  y = None
@@ -1419,7 +1419,7 @@ def forward(ctx, x, alpha=0.99):
             def backward(ctx, grad_out):
                 return grad_out
 
-        @torch.compile
+        @torch.compile(backend="eager")
         def foo(x):
             return Foo.apply(x)
 
@@ -1806,13 +1806,13 @@ def forward(self, s77: "Sym(s17)", cos: "f32[s17, 8]", s17: "Sym(s17)", cos_1: "
 
             result: "f32[s17, 8]" = a * b
 
-            out: "f32[s17, 8]" = a + b
+            out: "f32[s17, 8]" = a + b;  b = None
 
             _set_grad_enabled_1 = torch._C._set_grad_enabled(True);  _set_grad_enabled_1 = None
-            return ((result, out, a, b), (s17, a, cos, result))
+            return ((result, out), (s17, a, cos, result))
 
     class bwd_body_0(torch.nn.Module):
-        def forward(self, grad_a: "f32[s17, 8]", grad_b: "f32[s17, 8]", unused_0, unused_1, s17: "Sym(s17)", a: "f32[s17, 8]", arg: "f32[s17, 8]", result: "f32[s17, 8]"):
+        def forward(self, grad_a: "f32[s17, 8]", grad_b: "f32[s17, 8]", s17: "Sym(s17)", a: "f32[s17, 8]", arg: "f32[s17, 8]", result: "f32[s17, 8]"):
             _set_grad_enabled = torch._C._set_grad_enabled(False);  _set_grad_enabled = None
 
             mul: "f32[s17, 8]" = a * grad_b;  a = grad_b = None
@@ -2051,17 +2051,17 @@ def forward(self, l_module_b: "f32[4, 4]", l_input_data_values: "f32[4, 4]", l_x
             _set_grad_enabled = torch._C._set_grad_enabled(False);  _set_grad_enabled = None
 
             mul: "f32[4, 4]" = l_module_b * 5;  l_module_b = None
-            mul_1: "f32[4, 4]" = mul * l_input_data_values
-            output_tensor: "f32[4, 4]" = 4 + mul_1
+            mul_1: "f32[4, 4]" = mul * l_input_data_values;  mul = None
+            output_tensor: "f32[4, 4]" = 4 + mul_1;  mul_1 = None
 
-            out: "f32[4, 4]" = output_tensor + l_x_
+            out: "f32[4, 4]" = output_tensor + l_x_;  output_tensor = None
             out_1: "f32[4, 4]" = l_input_data_values + l_x_;  l_input_data_values = l_x_ = None
 
             _set_grad_enabled_1 = torch._C._set_grad_enabled(True);  _set_grad_enabled_1 = None
-            return ((out, out_1, mul, mul_1, output_tensor), ())
+            return ((out, out_1), ())
 
     class bwd_body_0(torch.nn.Module):
-        def forward(self, grad_output: "f32[4, 4]", output_data: "f32[4, 4]", unused_0, unused_1, unused_2):
+        def forward(self, grad_output: "f32[4, 4]", output_data: "f32[4, 4]"):
             _set_grad_enabled = torch._C._set_grad_enabled(False);  _set_grad_enabled = None
 
             mul: "f32[4, 4]" = grad_output * 4;  mul = None
@@ -2133,6 +2133,49 @@ def fn(input_data, x):
         res = opt_fn(input_data, x)
         self.assertEqual(ref, res)
 
+    def test_inplace_op_with_side_effect_wrong_grad(self):
+        # Repro for https://github.com/pytorch/pytorch/issues/180642
+        # In-place op inside autograd.Function.forward combined with a side
+        # effect (list append) causes collect_intermediate_outputs to add a
+        # pre-mutation alias as an extra subgraph output. Because the alias
+        # shares the same TensorImpl as the real return value,
+        # set_gradient_edge overwrites output_nr, routing the backward
+        # gradient to the wrong slot and producing zero gradients.
+        captured = []
+
+        class Foo(torch.autograd.Function):
+            @staticmethod
+            def forward(ctx, x):
+                intermediate = torch.sin(x)
+                captured.append(intermediate)
+                loss = torch.tensor(0.0, device=x.device)
+                loss += x.sum()
+                return loss
+
+            @staticmethod
+            def backward(ctx, grad):
+                return grad.expand(8)
+
+        def fn(x):
+            return Foo.apply(x)
+
+        x = torch.randn(8, device="cpu", requires_grad=True)
+
+        # Eager reference
+        x_ref = x.detach().clone().requires_grad_(True)
+        captured.clear()
+        out_ref = fn(x_ref)
+        out_ref.backward()
+
+        # Compiled
+        torch._dynamo.reset()
+        captured.clear()
+        x_c = x.detach().clone().requires_grad_(True)
+        out_c = torch.compile(fn, backend="eager", fullgraph=True)(x_c)
+        out_c.backward()
+
+        self.assertEqual(x_ref.grad, x_c.grad)
+
 
 class AutogradFunctionFunctorchTests(torch._dynamo.test_case.TestCase):
     """Tests for autograd.Function compatibility with torch.func transforms.
diff --git a/test/dynamo/test_backends.py b/test/dynamo/test_backends.py
index f9a2d44bcde60..d39df6e2ebee4 100644
--- a/test/dynamo/test_backends.py
+++ b/test/dynamo/test_backends.py
@@ -110,7 +110,8 @@ def _check_backend_works(self, backend, device, boxed=True, options=None):
 
             def boxed_assert(gm, *example_args):
                 fn = lookup_backend(backend)(gm, *example_args)
-                assert fn._boxed_call
+                if not fn._boxed_call:
+                    raise AssertionError("Expected fn._boxed_call to be True")
                 return fn
 
             compiled_model = torch.compile(model, backend=boxed_assert, options=options)
@@ -245,8 +246,8 @@ def fn(x):
         self.assertIn("Break Reasons", explain_str)
 
         # Verify that for the given functions above, we report the correct number of graphs, graph breaks, and ops
-        self.assertEqual(8, explain_output.graph_count)
-        self.assertEqual(7, explain_output.graph_break_count)
+        self.assertEqual(2, explain_output.graph_count)
+        self.assertEqual(1, explain_output.graph_break_count)
         self.assertEqual(8, explain_output.op_count)
 
 
@@ -339,7 +340,8 @@ def test_lookup_custom_backend(self):
         mock_3_10.load.return_value = lambda: "mocked 3.10"
 
         def mock_eps(group=None):
-            assert group == backends_group, group
+            if group != backends_group:
+                raise AssertionError(f"Expected group {backends_group}, got {group}")
             mock_group = MagicMock()
             mock_group.names = [name]
             mock_group[name] = mock_3_10
@@ -365,7 +367,8 @@ def restore_registry():
             registry._discover_entrypoint_backends.cache_clear()
 
             backends = list_backends()
-            assert name in backends, (name, backends)
+            if name not in backends:
+                raise AssertionError(f"Expected {name} in backends, got {backends}")
 
     def test_backend_recompilation(self):
         def fn(x):
diff --git a/test/dynamo/test_backward_higher_order_ops.py b/test/dynamo/test_backward_higher_order_ops.py
index 42cfbd954f020..f49885319ca67 100644
--- a/test/dynamo/test_backward_higher_order_ops.py
+++ b/test/dynamo/test_backward_higher_order_ops.py
@@ -92,6 +92,7 @@ def forward(self, grad_1: "f32[2]"):
 """,
         )
 
+    @torch._functorch.config.patch(guess_tangent_strides_as_outputs=True)
     @mock.patch(
         "torch._functorch.aot_autograd.AOT_COUNTER", new_callable=itertools.count
     )
@@ -130,29 +131,29 @@ def fn(x, y):
                     actual,
                     """\
 class GraphModule(torch.nn.Module):
-    def forward(self, L_inputs_ : list, s69: "Sym(s21)", L_sizes_0_: "f32[0, s21]"):
+    def forward(self, L_inputs_ : list, L_sizes_0_: "f32[0, 2]"):
         l_inputs_ = L_inputs_
         l_sizes_0_ = L_sizes_0_
 
-        getitem: "f32[s21]" = l_inputs_[0]
-        getitem_1: "f32[s21]" = l_inputs_[1]
-        getitem_2: "f32[s21]" = l_inputs_[2];  l_inputs_ = None
+        getitem: "f32[2]" = l_inputs_[0]
+        getitem_1: "f32[2]" = l_inputs_[1]
+        getitem_2: "f32[2]" = l_inputs_[2];  l_inputs_ = None
 
-        size: "Sym(s21)" = l_sizes_0_.size(1);  l_sizes_0_ = None
+        size: "Sym(2)" = l_sizes_0_.size(1);  l_sizes_0_ = None
 
         validate_outputs = torch__dynamo_compiled_autograd_ops_validate_outputs([getitem], [((None, None, device(type='cpu'), 6, 0, None), [size], False, 6)]);  getitem = size = None
-        getitem_9: "f32[s21]" = validate_outputs[0];  validate_outputs = None
+        getitem_9: "f32[2]" = validate_outputs[0];  validate_outputs = None
 
         call_aot_bwd_prologue = torch__dynamo_compiled_autograd_call_aot_bwd_prologue((), [], [], getitem_9);  getitem_9 = None
-        aot1_tangents_1: "f32[s21]" = call_aot_bwd_prologue[0];  call_aot_bwd_prologue = None
+        aot1_tangents_1: "f32[2]" = call_aot_bwd_prologue[0];  call_aot_bwd_prologue = None
 
         accumulate_grad = torch__dynamo_compiled_autograd_ops_AccumulateGrad([aot1_tangents_1], getitem_1, None, False);  getitem_1 = None
-        getitem_11: "f32[s21]" = accumulate_grad[0];  accumulate_grad = None
+        getitem_11: "f32[2]" = accumulate_grad[0];  accumulate_grad = None
 
-        result: "f32[s21]" = aot1_tangents_1 * aot1_tangents_1;  aot1_tangents_1 = None
+        result: "f32[2]" = aot1_tangents_1 * aot1_tangents_1;  aot1_tangents_1 = None
 
         accumulate_grad_1 = torch__dynamo_compiled_autograd_ops_AccumulateGrad([result], getitem_2, None, False);  result = getitem_2 = None
-        getitem_12: "f32[s21]" = accumulate_grad_1[0];  accumulate_grad_1 = None
+        getitem_12: "f32[2]" = accumulate_grad_1[0];  accumulate_grad_1 = None
         return (getitem_11, getitem_12)
 """,
                 )
@@ -161,35 +162,36 @@ def forward(self, L_inputs_ : list, s69: "Sym(s21)", L_sizes_0_: "f32[0, s21]"):
                     actual,
                     """\
 class GraphModule(torch.nn.Module):
-    def forward(self, L_inputs_ : list, s69: "Sym(s21)", L_sizes_0_: "f32[0, s21]"):
+    def forward(self, L_inputs_ : list, L_sizes_0_: "f32[0, 2]"):
         l_inputs_ = L_inputs_
         l_sizes_0_ = L_sizes_0_
 
-        getitem: "f32[s21]" = l_inputs_[0]
-        getitem_1: "f32[s21]" = l_inputs_[1]
-        getitem_2: "f32[s21]" = l_inputs_[2];  l_inputs_ = None
+        getitem: "f32[2]" = l_inputs_[0]
+        getitem_1: "f32[2]" = l_inputs_[1]
+        getitem_2: "f32[2]" = l_inputs_[2];  l_inputs_ = None
 
-        size: "Sym(s21)" = l_sizes_0_.size(1);  l_sizes_0_ = None
+        size: "Sym(2)" = l_sizes_0_.size(1);  l_sizes_0_ = None
 
         validate_outputs = torch__dynamo_compiled_autograd_ops_validate_outputs([getitem], [((None, None, device(type='cpu'), 6, 0, None), [size], False, 6)]);  getitem = size = None
-        getitem_9: "f32[s21]" = validate_outputs[0];  validate_outputs = None
+        getitem_9: "f32[2]" = validate_outputs[0];  validate_outputs = None
 
         call_aot_bwd_prologue = torch__dynamo_compiled_autograd_call_aot_bwd_prologue((), [], [], getitem_9);  getitem_9 = None
-        aot3_tangents_1: "f32[s21]" = call_aot_bwd_prologue[0];  call_aot_bwd_prologue = None
+        aot3_tangents_1: "f32[2]" = call_aot_bwd_prologue[0];  call_aot_bwd_prologue = None
 
         accumulate_grad = torch__dynamo_compiled_autograd_ops_AccumulateGrad([aot3_tangents_1], getitem_1, None, False);  getitem_1 = None
-        getitem_11: "f32[s21]" = accumulate_grad[0];  accumulate_grad = None
+        getitem_11: "f32[2]" = accumulate_grad[0];  accumulate_grad = None
 
-        result: "f32[s21]" = aot3_tangents_1 * aot3_tangents_1;  aot3_tangents_1 = None
+        result: "f32[2]" = aot3_tangents_1 * aot3_tangents_1;  aot3_tangents_1 = None
 
         accumulate_grad_1 = torch__dynamo_compiled_autograd_ops_AccumulateGrad([result], getitem_2, None, False);  result = getitem_2 = None
-        getitem_12: "f32[s21]" = accumulate_grad_1[0];  accumulate_grad_1 = None
+        getitem_12: "f32[2]" = accumulate_grad_1[0];  accumulate_grad_1 = None
         return (getitem_11, getitem_12)
 """,
                 )
 
             graph = None
 
+    @torch._functorch.config.patch(guess_tangent_strides_as_outputs=True)
     @mock.patch(
         "torch._functorch.aot_autograd.AOT_COUNTER", new_callable=itertools.count
     )
@@ -244,32 +246,32 @@ def fn(x, y):
                     actual,
                     """\
 class GraphModule(torch.nn.Module):
-    def forward(self, L_inputs_ : list, s69: "Sym(s21)", L_sizes_0_: "f32[0, s21]", L_hooks_1_keywords_fn_keywords_obj_counter: "Sym(s45)"):
+    def forward(self, L_inputs_ : list, L_sizes_0_: "f32[0, 2]", L_hooks_1_keywords_fn_keywords_obj_counter: "Sym(s45)"):
         l_inputs_ = L_inputs_
         l_sizes_0_ = L_sizes_0_
         l_hooks_1_keywords_fn_keywords_obj_counter = L_hooks_1_keywords_fn_keywords_obj_counter
 
-        getitem: "f32[s21]" = l_inputs_[0]
-        getitem_1: "f32[s21]" = l_inputs_[1]
-        getitem_2: "f32[s21]" = l_inputs_[2];  l_inputs_ = None
+        getitem: "f32[2]" = l_inputs_[0]
+        getitem_1: "f32[2]" = l_inputs_[1]
+        getitem_2: "f32[2]" = l_inputs_[2];  l_inputs_ = None
 
-        size: "Sym(s21)" = l_sizes_0_.size(1);  l_sizes_0_ = None
+        size: "Sym(2)" = l_sizes_0_.size(1);  l_sizes_0_ = None
 
         validate_outputs = torch__dynamo_compiled_autograd_ops_validate_outputs([getitem], [((None, None, device(type='cpu'), 6, 0, None), [size], False, 6)]);  getitem = size = None
-        getitem_9: "f32[s21]" = validate_outputs[0];  validate_outputs = None
+        getitem_9: "f32[2]" = validate_outputs[0];  validate_outputs = None
 
         call_aot_bwd_prologue = torch__dynamo_compiled_autograd_call_aot_bwd_prologue((), [], [], getitem_9);  getitem_9 = None
-        aot0_tangents_1: "f32[s21]" = call_aot_bwd_prologue[0];  call_aot_bwd_prologue = None
+        aot0_tangents_1: "f32[2]" = call_aot_bwd_prologue[0];  call_aot_bwd_prologue = None
 
         accumulate_grad = torch__dynamo_compiled_autograd_ops_AccumulateGrad([aot0_tangents_1], getitem_1, None, False);  getitem_1 = None
-        getitem_11: "f32[s21]" = accumulate_grad[0];  accumulate_grad = None
+        getitem_11: "f32[2]" = accumulate_grad[0];  accumulate_grad = None
 
         add: "Sym(s45 + 1)" = l_hooks_1_keywords_fn_keywords_obj_counter + 1;  l_hooks_1_keywords_fn_keywords_obj_counter = None
 
-        result: "f32[s21]" = aot0_tangents_1 * aot0_tangents_1;  aot0_tangents_1 = None
+        result: "f32[2]" = aot0_tangents_1 * aot0_tangents_1;  aot0_tangents_1 = None
 
         accumulate_grad_1 = torch__dynamo_compiled_autograd_ops_AccumulateGrad([result], getitem_2, None, False);  result = getitem_2 = None
-        getitem_12: "f32[s21]" = accumulate_grad_1[0];  accumulate_grad_1 = None
+        getitem_12: "f32[2]" = accumulate_grad_1[0];  accumulate_grad_1 = None
         return (getitem_11, getitem_12, add)
 """,
                 )
diff --git a/test/dynamo/test_base_hop.py b/test/dynamo/test_base_hop.py
index 24f36e8b995ee..3c9ab9995e692 100644
--- a/test/dynamo/test_base_hop.py
+++ b/test/dynamo/test_base_hop.py
@@ -48,7 +48,8 @@ def f(x, y):
         out = f(x, y)
         self.assertEqual(out, inner(x, y))
 
-        assert len(backend.graphs) == 1
+        if len(backend.graphs) != 1:
+            raise AssertionError(f"Expected 1 graph, got {len(backend.graphs)}")
         self.assertExpectedInline(
             normalize_graph(backend.graphs[0]),
             """\
@@ -370,7 +371,8 @@ def f(x, y):
         expected = torch.autograd.grad(out, x, y)
         self.assertEqual(result, expected)
 
-        assert len(backend.fw_graphs) == 1
+        if len(backend.fw_graphs) != 1:
+            raise AssertionError(f"Expected 1 fw_graph, got {len(backend.fw_graphs)}")
         self.assertExpectedInline(
             normalize_graph(backend.fw_graphs[0]),
             """\
@@ -390,7 +392,8 @@ def forward(self, arg0_1: "f32[3, 3]", arg1_1: "f32[3, 3]"):
 """,  # NOQA: B950
         )
 
-        assert len(backend.bw_graphs) == 1
+        if len(backend.bw_graphs) != 1:
+            raise AssertionError(f"Expected 1 bw_graph, got {len(backend.bw_graphs)}")
         self.assertExpectedInline(
             normalize_graph(backend.bw_graphs[0]),
             """\
diff --git a/test/dynamo/test_buffers_override.py b/test/dynamo/test_buffers_override.py
index 3ceba631423d1..30d8b78497501 100644
--- a/test/dynamo/test_buffers_override.py
+++ b/test/dynamo/test_buffers_override.py
@@ -20,7 +20,7 @@ def forward(self):
                 return self.A * torch.zeros(1, 1)
 
         model = SomeModel().to(torch.device("cpu"))
-        compiled_model = torch.compile(model)
+        compiled_model = torch.compile(model, backend="eager")
         self.assertEqual(compiled_model.A, torch.ones(3, 3))
         compiled_model()
 
@@ -38,7 +38,7 @@ def forward(self):
                 return self.B * torch.zeros(1, 1)
 
         model = SomeModel().to(torch.device("cpu"))
-        compiled_model = torch.compile(model)
+        compiled_model = torch.compile(model, backend="eager")
         self.assertEqual(compiled_model.B, torch.ones(3, 3))
         compiled_model()
 
diff --git a/test/dynamo/test_bytecode_debugger.py b/test/dynamo/test_bytecode_debugger.py
new file mode 100644
index 0000000000000..1d7996c45655c
--- /dev/null
+++ b/test/dynamo/test_bytecode_debugger.py
@@ -0,0 +1,1424 @@
+# Owner(s): ["module: dynamo"]
+
+"""
+Tests for torch._dynamo.bytecode_debugger
+"""
+
+import re
+import sys
+from contextlib import nullcontext, redirect_stdout
+from io import StringIO
+from unittest.mock import patch
+
+import torch
+import torch._dynamo
+from torch._dynamo.bytecode_debugger import debug
+from torch._dynamo.test_case import run_tests, TestCase
+
+
+class InteractiveDebugSession:
+    """Allows reactive interaction with the bytecode debugger in tests.
+
+    Uses a generator for cooperative multitasking between test and debugger.
+    The test_logic function receives initial output as its argument, then
+    yields commands and receives subsequent output via send().
+
+    Set use_debug=False to test auto-activation (no debug() wrapper).
+    """
+
+    def __init__(self, fn, args, test_logic, use_debug=True):
+        self.test_logic = test_logic
+        self.output_buffer = StringIO()
+        self._test_error: BaseException | None = None
+
+        with patch("builtins.input", self._fake_input):
+            with redirect_stdout(self.output_buffer):
+                with debug() if use_debug else nullcontext() as self.ctx:
+                    self.result = fn(*args)
+
+        # Re-raise any assertion error captured from the test generator
+        if self._test_error is not None:
+            raise self._test_error
+
+        # Ensure test_logic completed by sending final output
+        if hasattr(self.test_logic, "send"):
+            final_output = self.output_buffer.getvalue()
+            try:
+                cmd = self.test_logic.send(final_output)
+                raise AssertionError(
+                    f"test_logic did not complete - yielded {cmd!r} after session ended"
+                )
+            except StopIteration:
+                pass
+
+    def _fake_input(self, prompt=""):
+        self.output_buffer.write(prompt)
+
+        output = self.output_buffer.getvalue()
+        self.output_buffer.seek(0)
+        self.output_buffer.truncate(0)
+
+        try:
+            if not hasattr(self.test_logic, "send"):
+                self.test_logic = self.test_logic(self, output)
+                return next(self.test_logic)
+            return self.test_logic.send(output)
+        except StopIteration:
+            return "q"
+        except Exception as e:
+            # Capture test failures (e.g. AssertionError) and exit the
+            # debugger cleanly.  Without this, the exception propagates
+            # through sys.monitoring callbacks, gets converted to
+            # KeyboardInterrupt by the "q" handler, and is silently
+            # swallowed by the debug() context manager.
+            self._test_error = e
+            return "q"
+
+
+class TestBytecodeDebugger(TestCase):
+    def test_debug_context_manager_basic(self):
+        """Test that the debug context manager works without errors."""
+
+        @torch.compile(backend="eager")
+        def fn(x):
+            return x + 1
+
+        def test_logic(sess, initial):
+            self.assertIn("Commands:", initial)
+            output = yield "c"
+            self.assertIn("returned:", output)
+
+        inp = torch.randn(3)
+        sess = InteractiveDebugSession(fn, (inp,), test_logic)
+        self.assertEqual(sess.result, inp + 1)
+
+    def test_step_mode(self):
+        """Test that step mode steps through each instruction."""
+
+        @torch.compile(backend="eager")
+        def fn(x):
+            return x + 1
+
+        def test_logic(sess, initial):
+            step_count = 0
+            output = initial
+            while "returned:" not in output:
+                self.assertIn("Instruction", output)
+                step_count += 1
+                output = yield "s"
+            expected = len(sess.ctx.get_instructions())
+            self.assertGreaterEqual(step_count, expected - 1)
+            self.assertLessEqual(step_count, expected + 1)
+
+        InteractiveDebugSession(fn, (torch.randn(3),), test_logic)
+
+    def test_verbose_mode(self):
+        """Test that verbose mode prints each instruction in order."""
+
+        @torch.compile(backend="eager")
+        def fn(x):
+            return x + 1
+
+        def test_logic(sess, initial):
+            output = yield "v"  # Enable verbose mode
+            self.assertIn("Verbose mode enabled", output)
+            output = yield "c"  # Continue to end
+
+            # Parse verbose output to extract printed instruction opnames
+            printed_opnames = []
+            for line in output.split("\n"):
+                match = re.match(r"Running \[(\d+)\] (\w+)", line)
+                if match:
+                    printed_opnames.append(match.group(2))
+
+            self.assertGreater(len(printed_opnames), 0)
+
+            # Printed instructions should be a suffix of the full instruction list
+            expected_opnames = [inst.opname for inst in sess.ctx.get_instructions()]
+            start_idx = len(expected_opnames) - len(printed_opnames)
+            self.assertGreaterEqual(start_idx, 0)
+            self.assertEqual(printed_opnames, expected_opnames[start_idx:])
+
+        InteractiveDebugSession(fn, (torch.randn(3),), test_logic)
+
+    def test_return_value_printed(self):
+        """Test that return values are printed."""
+
+        @torch.compile(backend="eager")
+        def fn(x):
+            return x + 1
+
+        input_tensor = torch.tensor([1.0, 2.0, 3.0])
+        expected_result = input_tensor + 1
+
+        def test_logic(sess, initial):
+            output = yield "c"
+            expected_line = f"=== fn returned: {expected_result!r} ==="
+            self.assertIn(expected_line, output)
+
+        sess = InteractiveDebugSession(fn, (input_tensor,), test_logic)
+        self.assertEqual(sess.result, expected_result)
+
+    def test_stack_command(self):
+        """Test that the stack command shows correct values including NULL.
+
+        Uses a function with a method call to exercise NULL stack values
+        that appear in Python 3.11+ during call setup.
+        """
+
+        @torch.compile(backend="eager")
+        def fn(x):
+            return torch.sin(x)
+
+        def test_logic(sess, initial):
+            output = initial
+            # Step until we hit the CALL that invokes the compiled graph
+            # function (forward), skipping instrumentation calls
+            # (record_pregraph_bytecode_enter/exit).
+            if sys.version_info >= (3, 12):
+                call_pattern = r">>>.*\[\s*\d+\]:\s*CALL\b"
+            elif sys.version_info >= (3, 11):
+                call_pattern = r">>>.*\[\s*\d+\]:\s*PRECALL\b"
+            else:
+                call_pattern = r">>>.*\[\s*\d+\]:\s*CALL_FUNCTION\b"
+            while True:
+                while not re.search(call_pattern, output):
+                    output = yield "s"
+                stack_output = yield "stack"
+                if "forward" in stack_output and "record_pregraph" not in stack_output:
+                    break
+                output = yield "s"
+
+            # Extract just the stack lines (skip prompt)
+            lines = [
+                l for l in stack_output.split("\n") if l.strip() and "(bdb)" not in l
+            ]
+            stack_at_call = "\n".join(lines)
+
+            # Normalize addresses and tensor values
+            stack_at_call = re.sub(r"0x[0-9a-f]+(?! <NULL>)", "0xADDR", stack_at_call)
+            stack_at_call = re.sub(r"tensor\([^)]+\)", "tensor(...)", stack_at_call)
+
+            if sys.version_info >= (3, 13):
+                self.assertExpectedInline(
+                    stack_at_call,
+                    """\
+Stack (TOS at end):
+  [0] 0xADDR <function forward at 0xADDR>
+  [1] 0x0 <NULL>
+  [2] 0xADDR tensor(...)""",
+                )
+            elif sys.version_info >= (3, 11):
+                self.assertExpectedInline(
+                    stack_at_call,
+                    """\
+Stack (TOS at end):
+  [0] 0x0 <NULL>
+  [1] 0xADDR <function forward at 0xADDR>
+  [2] 0xADDR tensor(...)""",
+                )
+            else:
+                # Python 3.10
+                self.assertExpectedInline(
+                    stack_at_call,
+                    """\
+Stack (TOS at end):
+  [0] 0xADDR <function forward at 0xADDR>
+  [1] 0xADDR tensor(...)""",
+                )
+
+        InteractiveDebugSession(fn, (torch.ones(3),), test_logic)
+
+    def test_locals_command(self):
+        """Test that the locals command shows local variables with values.
+
+        Checks that:
+        1. Initially, locals shows x = tensor([...])
+        2. After a STORE_FAST instruction, the new local appears
+        """
+
+        @torch.compile(backend="eager")
+        def fn(x):
+            y = x + 1
+            return y
+
+        input_tensor = torch.tensor([1.0, 2.0, 3.0])
+
+        def test_logic(sess, initial):
+            # Check initial locals shows x
+            locals_output = yield "locals"
+            self.assertIn("x = tensor([1., 2., 3.])", locals_output)
+
+            # Step until we find a STORE_FAST for a variable other than x
+            output = initial
+            stored_var_name = None
+            while stored_var_name is None:
+                output = yield "s"
+                match = re.search(r">>>.*STORE_FAST\s+(\w+)", output)
+                if match and match.group(1) != "x":
+                    stored_var_name = match.group(1)
+
+            # Step once more so STORE_FAST executes
+            yield "s"
+
+            # Check that the new variable appears in locals
+            locals_output = yield "locals"
+            self.assertIn(f"{stored_var_name} = ", locals_output)
+
+        InteractiveDebugSession(fn, (input_tensor,), test_logic)
+
+    def test_list_command(self):
+        """Test that the list command behaves like pdb."""
+
+        @torch.compile(backend="eager")
+        def fn(x):
+            y = x + 1
+            z = y * 2
+            w = z - 1
+            return w + x
+
+        def test_logic(sess, initial):
+            # First list shows around current instruction
+            list_output1 = yield "l"
+            self.assertIn(">>>", list_output1)  # Current instruction marker
+            self.assertIn("# [offset]", list_output1)  # Header
+
+            # Second list continues (may or may not have >>> depending on position)
+            list_output2 = yield "l"
+            # Should show different instructions (or end of bytecode message)
+            if "(End of bytecode)" not in list_output2:
+                self.assertIn("# [offset]", list_output2)
+
+            # 'l .' resets to current instruction
+            list_dot_output = yield "l ."
+            self.assertIn(">>>", list_dot_output)  # Current instruction marker again
+
+            # 'l N' lists 11 instructions starting at N
+            list_from_0 = yield "l 0"
+            # Should show instruction 0
+            self.assertRegex(list_from_0, r"\s+0\s+\[")
+
+            # 'l first, last' lists range (inclusive)
+            list_range = yield "l 0, 2"
+            # Should show exactly 3 instructions (0, 1, 2)
+            instruction_lines = [
+                line
+                for line in list_range.split("\n")
+                if re.match(r"\s*(?:>>>)?\s*\*?\s*\d+\s+\[", line)
+            ]
+            self.assertEqual(len(instruction_lines), 3)
+
+            # 'l first, count' when count < first, count is number of lines
+            list_count = yield "l 5, 2"
+            # Should show 2 instructions starting at 5
+            instruction_lines = [
+                line
+                for line in list_count.split("\n")
+                if re.match(r"\s*(?:>>>)?\s*\*?\s*\d+\s+\[", line)
+            ]
+            self.assertEqual(len(instruction_lines), 2)
+
+        InteractiveDebugSession(fn, (torch.randn(3),), test_logic)
+
+    def test_disassemble_command(self):
+        """Test that the ll (disassemble) command shows all bytecode."""
+
+        @torch.compile(backend="eager")
+        def fn(x):
+            return x + 1
+
+        def test_logic(sess, initial):
+            ll_output = yield "ll"
+
+            # Check for disassembly header with instruction count
+            self.assertRegex(ll_output, r"Disassembly \(\d+ instructions\):")
+            self.assertIn("# [offset]", ll_output)
+
+            # Check for stable bytecode instructions
+            if sys.version_info >= (3, 11):
+                self.assertIn("RESUME", ll_output)
+            self.assertIn("LOAD_FAST x", ll_output)
+            self.assertIn("RETURN_VALUE", ll_output)
+
+            # Verify instruction count matches
+            instruction_matches = re.findall(
+                r"^\s*\d+\s*\[\s*\d+\]:", ll_output, re.MULTILINE
+            )
+            expected = len(sess.ctx.get_instructions())
+            self.assertEqual(len(instruction_matches), expected)
+
+        InteractiveDebugSession(fn, (torch.randn(3),), test_logic)
+
+    def test_breakpoint_commands(self):
+        """Test breakpoint set, list, hit, and clear commands."""
+
+        @torch.compile(backend="eager")
+        def fn(x):
+            return x + 1
+
+        def test_logic(sess, initial):
+            # Set breakpoint at instruction 5
+            output = yield "b 5"
+            self.assertIn("Breakpoint set at instruction 5", output)
+
+            # List breakpoints
+            output = yield "b"
+            self.assertIn("Breakpoints:", output)
+            self.assertIn("5", output)
+
+            # Continue to breakpoint - should stop at instruction 5
+            output = yield "c"
+            self.assertIn("Breakpoint hit at instruction 5", output)
+
+            # Clear breakpoint and set a new one further ahead
+            output = yield "cl 5"
+            self.assertIn("Breakpoint cleared at instruction 5", output)
+
+            output = yield "b 10"
+            self.assertIn("Breakpoint set at instruction 10", output)
+
+            # Continue - should skip instruction 5 (cleared) and hit instruction 10
+            output = yield "cl 10"
+            self.assertIn("Breakpoint cleared at instruction 10", output)
+
+            output = yield "c"
+            self.assertNotIn("Breakpoint hit at instruction 10", output)
+
+        InteractiveDebugSession(fn, (torch.randn(3),), test_logic)
+
+    def test_breakpoint_validation(self):
+        """Test that invalid instruction numbers are rejected."""
+
+        @torch.compile(backend="eager")
+        def fn(x):
+            return x + 1
+
+        def test_logic(sess, initial):
+            output = yield "b 9999"
+            self.assertIn("Invalid instruction number: 9999", output)
+
+            output = yield "b -1"
+            self.assertIn("Invalid instruction number: -1", output)
+
+            output = yield "b abc"
+            self.assertIn("Invalid instruction number: abc", output)
+
+        InteractiveDebugSession(fn, (torch.randn(3),), test_logic)
+
+    def test_expression_evaluation(self):
+        """Test that arbitrary expressions can be evaluated."""
+
+        @torch.compile(backend="eager")
+        def fn(x):
+            return x + 1
+
+        def test_logic(sess, initial):
+            # Step a few times
+            for _ in range(5):
+                yield "s"
+
+            # Evaluate expression
+            output = yield "2 + 2"
+            self.assertIn("4", output)
+
+        InteractiveDebugSession(fn, (torch.tensor([1.0, 2.0, 3.0]),), test_logic)
+
+    def test_stack_variable_access(self):
+        """Test that __stack__ variable is accessible in expressions."""
+
+        @torch.compile(backend="eager")
+        def fn(x):
+            return torch.sin(x)
+
+        def test_logic(sess, initial):
+            output = initial
+            # Step until we hit the CALL that invokes the compiled graph
+            # function. See test_stack_command for details.
+            if sys.version_info >= (3, 12):
+                call_pattern = r">>>.*\[\s*\d+\]:\s*CALL\b"
+            elif sys.version_info >= (3, 11):
+                call_pattern = r">>>.*\[\s*\d+\]:\s*PRECALL\b"
+            else:
+                call_pattern = r">>>.*\[\s*\d+\]:\s*CALL_FUNCTION\b"
+            while True:
+                while not re.search(call_pattern, output):
+                    output = yield "s"
+                stack_output = yield "__stack__"
+                if "forward" in stack_output and "record_pregraph" not in stack_output:
+                    break
+                output = yield "s"
+
+            # Normalize addresses and tensor values
+            normalized = re.sub(r"0x[0-9a-f]+", "0xADDR", stack_output)
+            normalized = re.sub(r"tensor\([^)]+\)", "tensor(...)", normalized)
+            # Extract just the list part
+            match = re.search(r"\[.*\]", normalized, re.DOTALL)
+            self.assertIsNotNone(match)
+            stack_list = match.group(0)
+
+            if sys.version_info >= (3, 13):
+                self.assertExpectedInline(
+                    stack_list,
+                    """[<function forward at 0xADDR>, <NULL>, tensor(...)]""",
+                )
+            elif sys.version_info >= (3, 11):
+                self.assertExpectedInline(
+                    stack_list,
+                    """[<NULL>, <function forward at 0xADDR>, tensor(...)]""",
+                )
+            else:
+                # Python 3.10
+                self.assertExpectedInline(
+                    stack_list,
+                    """[<function forward at 0xADDR>, tensor(...)]""",
+                )
+
+        InteractiveDebugSession(fn, (torch.ones(3),), test_logic)
+
+    def test_print_command(self):
+        """Test the p (print) command."""
+
+        @torch.compile(backend="eager")
+        def fn(x):
+            return x + 1
+
+        def test_logic(sess, initial):
+            # Step a few times to ensure x is in scope
+            for _ in range(5):
+                yield "s"
+
+            output = yield "p x"
+            self.assertIn("x = tensor([1., 2.])", output)
+
+        InteractiveDebugSession(fn, (torch.tensor([1.0, 2.0]),), test_logic)
+
+    def test_help_command(self):
+        """Test that the help command works."""
+
+        @torch.compile(backend="eager")
+        def fn(x):
+            return x + 1
+
+        def test_logic(sess, initial):
+            output = yield "h"
+            self.assertIn("Commands:", output)
+            self.assertIn("s [n]", output)
+            self.assertIn("cont", output)
+            self.assertIn("verbose", output)
+            self.assertIn("__stack__", output)
+
+        InteractiveDebugSession(fn, (torch.randn(3),), test_logic)
+
+    def test_multiple_calls(self):
+        """Test debugger works across multiple function calls."""
+
+        @torch.compile(backend="eager")
+        def fn1(x):
+            return x + 1
+
+        @torch.compile(backend="eager")
+        def fn2(x):
+            return x * 2
+
+        def test_logic(sess, initial):
+            self.assertIn("Entering Dynamo-generated code: fn1", initial)
+            # Enable verbose so we can see both functions' execution
+            output = yield "v"
+            self.assertIn("Verbose mode enabled", output)
+            output = yield "c"
+            # fn1 should have continued and returned
+            self.assertIn("fn1 returned:", output)
+            # fn2 should have entered (verbose), run, and returned
+            self.assertIn("Entering Dynamo-generated code: fn2", output)
+            self.assertIn("fn2 returned:", output)
+            # Verify we saw instructions from both functions
+            self.assertIn("LOAD_FAST", output)
+            self.assertIn("RETURN_VALUE", output)
+
+        inp1 = torch.randn(3)
+        inp2 = torch.randn(3)
+
+        def multi_call_fn(x1, x2):
+            r1 = fn1(x1)
+            r2 = fn2(x2)
+            return r1, r2
+
+        sess = InteractiveDebugSession(multi_call_fn, (inp1, inp2), test_logic)
+        self.assertEqual(sess.result, (inp1 + 1, inp2 * 2))
+
+    def test_graph_break(self):
+        """Test stepping through compiled resume function"""
+
+        from torch._dynamo.resume_execution import TORCH_DYNAMO_RESUME_IN_PREFIX
+
+        @torch.compile(backend="eager")
+        def fn(x):
+            x = x + 1
+            torch._dynamo.graph_break()
+            return x + 2
+
+        def test_logic(sess, initial):
+            self.assertIn("Entering Dynamo-generated code: fn", initial)
+            # Enable verbose and continue — should trace through both
+            # the first compiled code and the resume function
+            output = yield "v"
+            self.assertIn("Verbose mode enabled", output)
+            output = yield "c"
+            # Should see resume function entered and both returns
+            self.assertIn(
+                "Entering Dynamo-generated code: " + TORCH_DYNAMO_RESUME_IN_PREFIX,
+                output,
+            )
+            self.assertRegex(output, rf"{TORCH_DYNAMO_RESUME_IN_PREFIX}\w+ returned:")
+            self.assertIn("fn returned: ", output)
+
+        InteractiveDebugSession(fn, (torch.randn(3),), test_logic)
+
+    def test_stack_at_return_value(self):
+        """Test that stack has exactly 1 element at RETURN_VALUE instruction.
+
+        This tests that stack depth calculation handles opcodes without arguments
+        correctly (passing None to dis.stack_effect instead of 0).
+        """
+
+        @torch.compile(backend="eager")
+        def fn(x):
+            y = x + 1
+            return y
+
+        def test_logic(sess, initial):
+            output = initial
+            # Step until we hit RETURN_VALUE
+            while not re.search(r">>>.*RETURN_VALUE", output):
+                output = yield "s"
+
+            # Check stack at RETURN_VALUE
+            stack_output = yield "stack"
+            self.assertIn("Stack (TOS at end):", stack_output)
+
+            # Count stack entries - should be exactly 1
+            stack_entries = re.findall(r"^\s*\[\d+\]", stack_output, re.MULTILINE)
+            self.assertEqual(len(stack_entries), 1)
+
+            # Should contain a tensor (the return value)
+            self.assertIn("tensor(", stack_output)
+
+        InteractiveDebugSession(fn, (torch.ones(3),), test_logic)
+
+    def test_breakpoint_marker(self):
+        """Test that create_breakpoint() auto-activates the debugger.
+
+        Manually injects a BREAKPOINT_MARKER via create_breakpoint() (without
+        bdb_breakpoint()) and verifies the debugger activates automatically,
+        hitting the breakpoint before side effect replay.
+        """
+        from torch._dynamo.bytecode_transformation import create_breakpoint
+        from torch._dynamo.side_effects import SideEffects
+
+        original_codegen_update_mutated = SideEffects.codegen_update_mutated
+
+        def patched_codegen_update_mutated(self, cg, log_side_effects=False):
+            cg.extend_output(create_breakpoint())
+            return original_codegen_update_mutated(self, cg, log_side_effects)
+
+        global mylist
+
+        mylist = [1, 2, 3]
+
+        @torch.compile(backend="eager")
+        def fn(x):
+            mylist.append(4)
+            return x + 1
+
+        def test_logic(sess, initial):
+            # Auto-breakpoint skips to the breakpoint without pausing at start
+            self.assertIn("Breakpoint hit (programmatic)", initial)
+
+            # Verify mylist still unchanged (side effect not yet applied)
+            output = yield "mylist"
+            self.assertIn("[1, 2, 3]", output)
+
+            # Step until we find LOAD_GLOBAL mylist
+            found_load_mylist = False
+            for _ in range(10):
+                output = yield "s"
+                if "LOAD_GLOBAL mylist" in output:
+                    found_load_mylist = True
+                    break
+            self.assertTrue(
+                found_load_mylist, "Expected to find LOAD_GLOBAL mylist instruction"
+            )
+
+            # Step and check mylist until we see the mutation
+            for _ in range(30):
+                output = yield "s"
+                if "returned:" in output:
+                    self.fail("Function returned before seeing mutation")
+                output = yield "mylist"
+                if "[1, 2, 3, 4]" in output:
+                    break
+            else:
+                self.fail("Never saw mylist mutated to [1, 2, 3, 4]")
+
+            # Continue to end
+            yield "c"
+
+        with patch.object(
+            SideEffects,
+            "codegen_update_mutated",
+            patched_codegen_update_mutated,
+        ):
+            inp = torch.randn(3)
+            sess = InteractiveDebugSession(fn, (inp,), test_logic, use_debug=False)
+
+        self.assertEqual(sess.result, inp + 1)
+        self.assertEqual(mylist, [1, 2, 3, 4])
+
+    def test_user_defined_variables_persist(self):
+        """Test that user-defined variables persist across commands."""
+
+        @torch.compile(backend="eager")
+        def fn(x):
+            return x + 1
+
+        def test_logic(sess, initial):
+            # Define a variable
+            yield "tmp = 42"
+            # Access it in a later command - should evaluate to 84
+            output = yield "tmp * 2"
+            self.assertIn("84", output)
+
+        InteractiveDebugSession(fn, (torch.randn(3),), test_logic)
+
+    def test_user_defined_variables_shadow_locals(self):
+        """Test that user-defined variables can shadow local variables."""
+
+        @torch.compile(backend="eager")
+        def fn(x):
+            y = x + 1
+            return y
+
+        input_tensor = torch.tensor([1.0, 2.0, 3.0])
+
+        def test_logic(sess, initial):
+            # Step a few times to get into the execution
+            for _ in range(5):
+                yield "s"
+
+            # Check original value of x
+            output = yield "x"
+            self.assertIn("tensor([1., 2., 3.])", output)
+
+            # Shadow x with a new value
+            yield "x = 999"
+
+            # Now x should show the shadowed value
+            output = yield "x"
+            self.assertIn("999", output)
+
+            # Continue to let the function complete
+            yield "c"
+
+        sess = InteractiveDebugSession(fn, (input_tensor,), test_logic)
+
+        # Verify the original computation was unaffected by shadowing
+        self.assertEqual(sess.result, input_tensor + 1)
+
+    def test_inplace_mutation_affects_execution(self):
+        """Test that in-place mutations in debugger affect actual execution."""
+
+        @torch.compile(backend="eager")
+        def fn(x):
+            y = x + 1
+            return y
+
+        input_tensor = torch.tensor([1.0, 2.0, 3.0])
+
+        def test_logic(sess, initial):
+            # Step a few times to get into the execution
+            for _ in range(5):
+                yield "s"
+
+            # In-place mutation - this DOES affect execution
+            yield "x += 1"
+
+            # Continue to let the function complete
+            yield "c"
+
+        sess = InteractiveDebugSession(fn, (input_tensor,), test_logic)
+
+        # In-place mutation affected execution: x was [1,2,3], became [2,3,4],
+        # then y = x + 1 = [3,4,5]
+        self.assertEqual(sess.result, torch.tensor([3.0, 4.0, 5.0]))
+
+    def test_exception_in_bytecode(self):
+        """Test debugger stops when bytecode raises an exception.
+
+        Injects a division by zero into the bytecode and verifies:
+        1. The debugger stops at the exception
+        2. Shows the instruction that caused the exception
+        3. Allows inspection before propagating
+        """
+        from torch._dynamo.bytecode_transformation import create_instruction
+        from torch._dynamo.side_effects import SideEffects
+
+        original_codegen_update_mutated = SideEffects.codegen_update_mutated
+
+        def patched_codegen_update_mutated(self, cg, log_side_effects=False):
+            # Inject 31415/0 to cause ZeroDivisionError
+            if sys.version_info >= (3, 11):
+                div_inst = create_instruction("BINARY_OP", arg=11)  # 11 = TRUEDIV
+            else:
+                div_inst = create_instruction("BINARY_TRUE_DIVIDE")
+            cg.extend_output(
+                [
+                    create_instruction("LOAD_CONST", argval=31415),
+                    create_instruction("LOAD_CONST", argval=0),
+                    div_inst,
+                    create_instruction("POP_TOP"),
+                ]
+            )
+            return original_codegen_update_mutated(self, cg, log_side_effects)
+
+        global mylist
+        mylist = [1, 2, 3]
+
+        @torch.compile(backend="eager")
+        def fn(x):
+            mylist.append(4)
+            return x + 1
+
+        def test_logic(sess, initial):
+            # Continue - this will run until the exception
+            output = yield "c"
+
+            # Debugger should have stopped at the exception
+            self.assertIn("Exception raised at instruction", output)
+            if sys.version_info >= (3, 11):
+                self.assertIn("BINARY_OP", output)
+            else:
+                self.assertIn("BINARY_TRUE_DIVIDE", output)
+            self.assertIn("ZeroDivisionError", output)
+            self.assertIn("division by zero", output)
+
+            # Verify we can inspect locals at the exception point
+            locals_output = yield "locals"
+            self.assertIn("x =", locals_output)  # Input tensor should be in locals
+
+            # Verify we can inspect the stack at the exception point
+            # Note: On Python 3.11 with settrace, when an exception is raised,
+            # CPython has already popped the operands from the stack before the
+            # division fails, so we can't see them.
+            stack_output = yield "stack"
+            self.assertIn("Stack", stack_output)
+            if sys.version_info >= (3, 12):
+                self.assertIn("31415", stack_output)
+                self.assertRegex(stack_output, r"\[\d+\].*\b0\b")  # constant 0
+
+            # Verify we can evaluate expressions
+            eval_output = yield "1 + 1"
+            self.assertIn("2", eval_output)
+
+            # next step should terminate and exception should be propagated
+            yield "n"
+
+        with patch.object(
+            SideEffects,
+            "codegen_update_mutated",
+            patched_codegen_update_mutated,
+        ):
+            with self.assertRaises(ZeroDivisionError):
+                InteractiveDebugSession(fn, (torch.randn(3),), test_logic)
+
+    def test_empty_input_repeats_last_command(self):
+        """Test that empty input repeats the last command."""
+
+        @torch.compile(backend="eager")
+        def fn(x):
+            return x + 1
+
+        def test_logic(sess, initial):
+            # First command: step
+            output = yield ""
+            self.assertIn("Instruction", output)
+
+            # Empty input should repeat 's' (step)
+            output = yield ""
+            self.assertIn("Instruction", output)
+
+            # Now use verbose
+            output = yield "v"
+            self.assertIn("Verbose mode enabled", output)
+
+            # Empty input should repeat 'v' (toggle verbose off)
+            output = yield ""
+            self.assertIn("Verbose mode disabled", output)
+
+            # Continue to end
+            yield "c"
+
+        InteractiveDebugSession(fn, (torch.randn(3),), test_logic)
+
+    def test_step_with_count_argument(self):
+        """Test that step command accepts count argument"""
+
+        @torch.compile(backend="eager")
+        def fn(x):
+            return x + 1
+
+        def test_logic(sess, initial):
+            # Get starting instruction number
+            match = re.search(r"Instruction (\d+)", initial)
+            self.assertIsNotNone(match)
+            start_inst = int(match.group(1))
+
+            # Step 3 times with "s 3"
+            output = yield "s 3"
+            match = re.search(r"Instruction (\d+)", output)
+            self.assertIsNotNone(match)
+            end_inst = int(match.group(1))
+
+            # Should have advanced by 3 instructions
+            self.assertEqual(end_inst, start_inst + 3)
+
+            # Also test with "step 2"
+            output = yield "step 2"
+            match = re.search(r"Instruction (\d+)", output)
+            self.assertIsNotNone(match)
+            new_inst = int(match.group(1))
+
+            # Should have advanced by 2 more
+            self.assertEqual(new_inst, end_inst + 2)
+
+            # Continue to end
+            yield "c"
+
+        InteractiveDebugSession(fn, (torch.randn(3),), test_logic)
+
+    def test_next_with_count_argument(self):
+        """Test that 'n [count]' steps over calls without entering them."""
+        from torch._dynamo.resume_execution import TORCH_DYNAMO_RESUME_IN_PREFIX
+
+        @torch.compile(backend="eager")
+        def fn(x):
+            x = x + 1
+            torch._dynamo.graph_break()
+            return x + 2
+
+        def test_logic(sess, initial):
+            self.assertIn("Entering Dynamo-generated code: fn", initial)
+            # Use 'n' to step through fn's code, collecting all output
+            output = initial
+            all_output = output
+            while "fn returned:" not in output:
+                output = yield "n 3"
+                all_output += output
+
+            # 'n' should never have entered the resume function interactively
+            self.assertNotIn(
+                "Entering Dynamo-generated code: " + TORCH_DYNAMO_RESUME_IN_PREFIX,
+                all_output,
+            )
+
+        InteractiveDebugSession(fn, (torch.randn(3),), test_logic)
+
+    def test_return_command(self):
+        """Test that 'r' stops at the return instruction."""
+
+        @torch.compile(backend="eager")
+        def fn(x):
+            return x + 1
+
+        def test_logic(sess, initial):
+            # Step a few times, then use 'r' to run to return instruction
+            for _ in range(3):
+                yield "s"
+            output = yield "r"
+            # Should stop at the return instruction, not after it
+            self.assertIn("About to return from", output)
+            self.assertRegex(output, r"RETURN_VALUE|RETURN_CONST")
+            # Step past to complete the return
+            output = yield "s"
+            self.assertIn("returned:", output)
+
+        InteractiveDebugSession(fn, (torch.randn(3),), test_logic)
+
+    def test_return_stops_at_return_instruction(self):
+        """Test that 'r' runs through inner frames and stops at return."""
+        from torch._dynamo.resume_execution import TORCH_DYNAMO_RESUME_IN_PREFIX
+
+        @torch.compile(backend="eager")
+        def fn(x):
+            x = x + 1
+            torch._dynamo.graph_break()
+            return x + 2
+
+        def test_logic(sess, initial):
+            self.assertIn("Entering Dynamo-generated code: fn", initial)
+            # Use 'r' — should skip the resume function and stop at fn's return
+            output = yield "r"
+            self.assertRegex(output, rf"{TORCH_DYNAMO_RESUME_IN_PREFIX}\w+ returned:")
+            self.assertIn("About to return from fn", output)
+            self.assertRegex(output, r"RETURN_VALUE|RETURN_CONST")
+            # Step past to complete
+            output = yield "s"
+            self.assertIn("fn returned:", output)
+
+        InteractiveDebugSession(fn, (torch.randn(3),), test_logic)
+
+    def test_up_down_single_frame(self):
+        """Test u/d with only one tracked frame."""
+
+        @torch.compile(backend="eager")
+        def fn(x):
+            return x + 1
+
+        def test_logic(sess, initial):
+            output = yield "u"
+            self.assertIn("Oldest tracked frame", output)
+            output = yield "d"
+            self.assertIn("Newest tracked frame", output)
+
+        InteractiveDebugSession(fn, (torch.randn(3),), test_logic)
+
+    def test_up_down_commands(self):
+        """Test u/d with multiple tracked frames (graph break creates resume fn)."""
+        from torch._dynamo.resume_execution import TORCH_DYNAMO_RESUME_IN_PREFIX
+
+        @torch.compile(backend="eager")
+        def fn(x):
+            y = x + 1
+            torch._dynamo.graph_break()
+            return y + 2
+
+        def test_logic(sess, initial):
+            self.assertIn("Entering Dynamo-generated code: fn", initial)
+            # Step into the resume function
+            output = initial
+            while (
+                "Entering Dynamo-generated code: " + TORCH_DYNAMO_RESUME_IN_PREFIX
+                not in output
+            ):
+                output = yield "s"
+
+            # Lower frame (resume fn) just entered — stack should be empty
+            lower_stack = yield "stack"
+            self.assertIn("(empty)", lower_stack)
+
+            # Navigate up to the parent frame
+            output = yield "u"
+            self.assertIn("> fn", output)
+
+            # Upper frame stack may or may not be empty
+            upper_stack = yield "stack"
+            self.assertIn("Stack (TOS at end):", upper_stack)
+
+            # Check locals in upper frame — should have 'x'
+            locals_output = yield "locals"
+            self.assertIn("x =", locals_output)
+
+            # Check locals in lower frame — should NOT have 'x', but should have 'y'
+            # (resume function receives Dynamo-renamed parameters)
+            yield "d"
+            lower_locals_output = yield "locals"
+            self.assertNotIn("x =", lower_locals_output)
+            self.assertIn("y =", lower_locals_output)
+
+            # Go back up and check globals
+            yield "u"
+            globals_output = yield "globals"
+            self.assertIn("Globals:", globals_output)
+
+            # 'p' from upper frame should see fn's locals
+            p_output = yield "p x"
+            self.assertIn("x =", p_output)
+
+            # Set a breakpoint on fn's RETURN_VALUE from the upper frame.
+            # When we continue, the resume function will complete first,
+            # then fn will hit this breakpoint.
+            fn_code = sess.ctx.get_tracked_codes()[0]
+            fn_instructions = sess.ctx.get_instructions(fn_code)
+            return_idx = None
+            for i, inst in enumerate(fn_instructions):
+                if inst.opname in ("RETURN_VALUE", "RETURN_CONST"):
+                    return_idx = i
+            self.assertIsNotNone(return_idx)
+            output = yield f"b {return_idx}"
+            self.assertIn(f"Breakpoint set at instruction {return_idx}", output)
+
+            # Continue — should run through the resume function and hit
+            # the breakpoint in fn
+            output = yield "c"
+            self.assertIn(f"Breakpoint hit at instruction {return_idx}", output)
+            self.assertRegex(output, r"RETURN_VALUE|RETURN_CONST")
+            # The resume function should have returned by now
+            self.assertRegex(output, rf"{TORCH_DYNAMO_RESUME_IN_PREFIX}\w+ returned:")
+
+            # Continue to end
+            yield "c"
+
+        InteractiveDebugSession(fn, (torch.randn(3),), test_logic)
+
+    def test_bt_command(self):
+        """Test that 'bt' shows the tracked frame backtrace."""
+        from torch._dynamo.resume_execution import TORCH_DYNAMO_RESUME_IN_PREFIX
+
+        @torch.compile(backend="eager")
+        def fn(x):
+            x = x + 1
+            torch._dynamo.graph_break()
+            return x + 2
+
+        def test_logic(sess, initial):
+            self.assertIn("Entering Dynamo-generated code: fn", initial)
+            # Step into the resume function
+            output = initial
+            while (
+                "Entering Dynamo-generated code: " + TORCH_DYNAMO_RESUME_IN_PREFIX
+                not in output
+            ):
+                output = yield "s"
+
+            # bt should show both frames, with the resume function selected
+            bt_output = yield "bt"
+            lines = [l.strip() for l in bt_output.split("\n") if "[" in l and "]" in l]
+            self.assertEqual(len(lines), 2)
+            # First frame is fn (no > marker since we're viewing the callee)
+            self.assertIn("fn", lines[0])
+            self.assertNotIn(">", lines[0])
+            # Second frame is resume function (> marker since it's current)
+            self.assertIn(">", lines[1])
+            self.assertIn(TORCH_DYNAMO_RESUME_IN_PREFIX, lines[1])
+
+            # Move up and bt again — marker should move to fn
+            yield "u"
+            bt_output = yield "bt"
+            lines = [l.strip() for l in bt_output.split("\n") if "[" in l and "]" in l]
+            self.assertIn(">", lines[0])
+            self.assertNotIn(">", lines[1])
+
+            # Continue to end
+            yield "c"
+
+        InteractiveDebugSession(fn, (torch.randn(3),), test_logic)
+
+    def test_quit_suppresses_return_messages(self):
+        """Test that 'q' exits without printing return messages for each frame."""
+        from torch._dynamo.resume_execution import TORCH_DYNAMO_RESUME_IN_PREFIX
+
+        @torch.compile(backend="eager")
+        def fn(x):
+            y = x + 1
+            torch._dynamo.graph_break()
+            return y + 2
+
+        def test_logic(sess, initial):
+            self.assertIn("Entering Dynamo-generated code: fn", initial)
+            # Step into the resume function so multiple frames are active
+            output = initial
+            while (
+                "Entering Dynamo-generated code: " + TORCH_DYNAMO_RESUME_IN_PREFIX
+                not in output
+            ):
+                output = yield "s"
+
+            # Quit from inside the resume function
+            final = yield "q"
+            self.assertNotIn("returned:", final)
+            self.assertNotIn("Exception raised", final)
+
+        InteractiveDebugSession(fn, (torch.randn(3),), test_logic)
+
+    def test_step_from_upper_frame(self):
+        """Test that 's' from an upper frame steps the inner (execution) frame."""
+        from torch._dynamo.resume_execution import TORCH_DYNAMO_RESUME_IN_PREFIX
+
+        @torch.compile(backend="eager")
+        def fn(x):
+            x = x + 1
+            torch._dynamo.graph_break()
+            return x + 2
+
+        def test_logic(sess, initial):
+            self.assertIn("Entering Dynamo-generated code: fn", initial)
+            # Step into the resume function
+            output = initial
+            while (
+                "Entering Dynamo-generated code: " + TORCH_DYNAMO_RESUME_IN_PREFIX
+                not in output
+            ):
+                output = yield "s"
+
+            # Move up to the parent frame (fn)
+            output = yield "u"
+            self.assertIn("> fn", output)
+
+            # 's' from the upper frame should step the inner (resume) frame,
+            # not the viewed frame
+            output = yield "s"
+            # Should show the resume function's next instruction, not fn's
+            self.assertIn("Instruction", output)
+
+            # Keep stepping — should complete the resume function and
+            # return to fn
+            all_output = output
+            while "fn returned:" not in output:
+                output = yield "s"
+                all_output += output
+            self.assertRegex(
+                all_output, rf"{TORCH_DYNAMO_RESUME_IN_PREFIX}\w+ returned:"
+            )
+
+        InteractiveDebugSession(fn, (torch.randn(3),), test_logic)
+
+    def test_next_from_upper_frame(self):
+        """Test that 'n' from an upper frame steps over the callee."""
+        from torch._dynamo.resume_execution import TORCH_DYNAMO_RESUME_IN_PREFIX
+
+        @torch.compile(backend="eager")
+        def fn(x):
+            x = x + 1
+            torch._dynamo.graph_break()
+            return x + 2
+
+        def test_logic(sess, initial):
+            self.assertIn("Entering Dynamo-generated code: fn", initial)
+            # Step into the resume function
+            output = initial
+            while (
+                "Entering Dynamo-generated code: " + TORCH_DYNAMO_RESUME_IN_PREFIX
+                not in output
+            ):
+                output = yield "s"
+
+            # Move up to the parent frame (fn)
+            output = yield "u"
+            self.assertIn("> fn", output)
+
+            # 'n' from the parent frame should step over the resume function
+            # and stop at the next instruction in fn's frame
+            output = yield "n"
+            # The resume function should have returned
+            self.assertRegex(output, rf"{TORCH_DYNAMO_RESUME_IN_PREFIX}\w+ returned:")
+            # We should now be stopped at an instruction in the fn frame
+            self.assertIn("Instruction", output)
+
+            # Continue to end
+            yield "c"
+
+        InteractiveDebugSession(fn, (torch.randn(3),), test_logic)
+
+    def test_breakpoint_function(self):
+        """Test that breakpoint() inserts BREAKPOINT_MARKER without a graph break."""
+        from torch._dynamo.bytecode_debugger import breakpoint as bdb_breakpoint
+        from torch._dynamo.resume_execution import TORCH_DYNAMO_RESUME_IN_PREFIX
+
+        @torch.compile(backend="eager")
+        def fn(x):
+            x = x + 1
+            bdb_breakpoint()
+            return x + 2
+
+        def test_logic(sess, initial):
+            self.assertIn("Entering Dynamo-generated code: fn", initial)
+            # Breakpoint is at the first instruction, so it fires immediately
+            self.assertIn("Breakpoint hit (programmatic)", initial)
+            # No graph break: should NOT see a resume function
+            self.assertNotIn(TORCH_DYNAMO_RESUME_IN_PREFIX, initial)
+            # Continue to end
+            output = yield "c"
+            self.assertIn("returned:", output)
+
+        inp = torch.randn(3)
+        sess = InteractiveDebugSession(fn, (inp,), test_logic)
+        self.assertEqual(sess.result, inp + 3)
+
+    def test_programmatic_breakpoint_listed_and_clearable(self):
+        """Test that programmatic breakpoints appear in 'b' and can be cleared with 'cl'."""
+        from torch._dynamo.bytecode_debugger import breakpoint as bdb_breakpoint
+
+        @torch.compile(backend="eager")
+        def fn(x):
+            x = x + 1
+            bdb_breakpoint()
+            return x + 2
+
+        def test_logic(sess, initial):
+            self.assertIn("Entering Dynamo-generated code: fn", initial)
+            # Breakpoint is at the first instruction, so it fires immediately
+            self.assertIn("Breakpoint hit (programmatic)", initial)
+
+            # List breakpoints — should show the programmatic breakpoint
+            output = yield "b"
+            self.assertIn("Breakpoints:", output)
+            self.assertIn("LOAD_CONST <BREAKPOINT>", output)
+
+            # Extract the instruction index and clear it
+            bp_match = re.search(r"(\d+)\s+\[.*LOAD_CONST <BREAKPOINT>", output)
+            self.assertIsNotNone(bp_match)
+            bp_index = bp_match.group(1)
+            output = yield f"cl {bp_index}"
+            self.assertIn(f"Breakpoint cleared at instruction {bp_index}", output)
+
+            # Verify it's gone
+            output = yield "b"
+            self.assertIn("(none)", output)
+
+            # Continue to end
+            output = yield "c"
+            self.assertIn("returned:", output)
+
+        inp = torch.randn(3)
+        sess = InteractiveDebugSession(fn, (inp,), test_logic)
+        self.assertEqual(sess.result, inp + 3)
+
+    def test_breakpoint_after_graph_break(self):
+        """Test that breakpoint() after a graph break pauses in the resume function."""
+        from torch._dynamo.bytecode_debugger import breakpoint as bdb_breakpoint
+        from torch._dynamo.resume_execution import TORCH_DYNAMO_RESUME_IN_PREFIX
+
+        @torch.compile(backend="eager")
+        def fn(x):
+            x = x + 1
+            torch._dynamo.graph_break()
+            bdb_breakpoint()
+            return x + 2
+
+        def test_logic(sess, initial):
+            self.assertIn("Entering Dynamo-generated code: fn", initial)
+            # Enable verbose so we can see the resume function entering
+            yield "v"
+            # Continue — should enter the resume function and hit the breakpoint
+            output = yield "c"
+            self.assertIn(
+                "Entering Dynamo-generated code: " + TORCH_DYNAMO_RESUME_IN_PREFIX,
+                output,
+            )
+            self.assertIn("Breakpoint hit (programmatic)", output)
+            # Continue to end
+            output = yield "c"
+            self.assertIn("returned:", output)
+
+        inp = torch.randn(3)
+        sess = InteractiveDebugSession(fn, (inp,), test_logic)
+        self.assertEqual(sess.result, inp + 3)
+
+    def test_breakpoint_without_debug_context(self):
+        """Test that breakpoint() auto-activates the debugger without with debug()."""
+        from torch._dynamo.bytecode_debugger import breakpoint as bdb_breakpoint
+
+        @torch.compile(backend="eager")
+        def fn(x):
+            x = x + 1
+            bdb_breakpoint()
+            return x + 2
+
+        def test_logic(sess, initial):
+            # Auto-breakpoint runs to the breakpoint without pausing at start
+            self.assertIn("Breakpoint hit (programmatic)", initial)
+            yield "c"
+
+        inp = torch.randn(3)
+        sess = InteractiveDebugSession(fn, (inp,), test_logic, use_debug=False)
+        self.assertEqual(sess.result, inp + 3)
+
+    def test_breakpoint_without_debug_context_after_graph_break(self):
+        """Test auto-activate only fires for the resume function, not the first graph."""
+        from torch._dynamo.bytecode_debugger import breakpoint as bdb_breakpoint
+
+        @torch.compile(backend="eager")
+        def fn(x):
+            x = x + 1
+            torch._dynamo.graph_break()
+            bdb_breakpoint()
+            return x + 2
+
+        def test_logic(sess, initial):
+            # Auto-breakpoint skips to the breakpoint without pausing at start
+            self.assertIn("Breakpoint hit (programmatic)", initial)
+            # Should not have entered fn's first graph interactively
+            self.assertNotIn("Entering Dynamo-generated code: fn", initial)
+            yield "c"
+
+        inp = torch.randn(3)
+        sess = InteractiveDebugSession(fn, (inp,), test_logic, use_debug=False)
+        self.assertEqual(sess.result, inp + 3)
+
+    def test_nested_debug_contexts(self):
+        """Test that nested debug() contexts restore the outer correctly.
+
+        On Python 3.12+, compiled functions called during a sys.monitoring
+        callback skip Dynamo (tstate->tracing > 0), so nested debug contexts
+        must be tested outside monitoring callbacks.
+        """
+
+        @torch.compile(backend="eager")
+        def fn1(x):
+            return x + 1
+
+        @torch.compile(backend="eager")
+        def fn2(x):
+            return x * 2
+
+        @torch.compile(backend="eager")
+        def fn3(x):
+            return x + 3
+
+        outer_buf = StringIO()
+        with patch("builtins.input", return_value="c"):
+            with redirect_stdout(outer_buf):
+                with debug():
+                    fn1(torch.randn(3))
+
+                    # After fn1 returns, we're outside any monitoring callback.
+                    # Create a nested debug context for fn2.
+                    inner_buf = StringIO()
+                    with redirect_stdout(inner_buf):
+                        with debug():
+                            fn2(torch.randn(3))
+                    inner_output = inner_buf.getvalue()
+
+                    # After inner debug exits, outer should still work.
+                    fn3(torch.randn(3))
+
+        outer_output = outer_buf.getvalue()
+
+        self.assertIn("Entering Dynamo-generated code: fn1", outer_output)
+        self.assertIn("Entering Dynamo-generated code: fn2", inner_output)
+        self.assertIn("fn1 returned:", outer_output)
+        self.assertIn("fn2 returned:", inner_output)
+        # fn3 should still be tracked by the restored outer context.
+        # The "c" command disables stop-at-new-code, so we won't see the
+        # "Entering" message, but the return callback should still fire.
+        self.assertIn("fn3 returned:", outer_output)
+
+    def test_nested_auto_breakpoints(self):
+        """Test multiple functions with breakpoint() called sequentially."""
+        from torch._dynamo.bytecode_debugger import breakpoint as bdb_breakpoint
+
+        @torch.compile(backend="eager")
+        def fn1(x):
+            bdb_breakpoint()
+            return x + 1
+
+        @torch.compile(backend="eager")
+        def fn2(x):
+            bdb_breakpoint()
+            return x * 2
+
+        output_buf = StringIO()
+        with patch("builtins.input", return_value="c"):
+            with redirect_stdout(output_buf):
+                inp = torch.randn(3)
+                r1 = fn1(inp)
+                r2 = fn2(inp)
+
+        output = output_buf.getvalue()
+        # Both functions should have hit their breakpoints
+        self.assertEqual(output.count("Breakpoint hit (programmatic)"), 2)
+        self.assertEqual(r1, inp + 1)
+        self.assertEqual(r2, inp * 2)
+
+
+if __name__ == "__main__":
+    run_tests()
diff --git a/test/dynamo/test_bytecode_utils.py b/test/dynamo/test_bytecode_utils.py
index bd079cca9307d..73736846affa8 100644
--- a/test/dynamo/test_bytecode_utils.py
+++ b/test/dynamo/test_bytecode_utils.py
@@ -576,7 +576,7 @@ def hook(code, out_code):
         torch._dynamo.reset()
         handle = torch._dynamo.convert_frame.register_bytecode_hook(hook)
         try:
-            opt_fn = torch.compile(fn)
+            opt_fn = torch.compile(fn, backend="eager")
             for i in range(2, 12):
                 opt_fn(torch.randn(i), torch.randn(i))
         finally:
diff --git a/test/dynamo/test_compile.py b/test/dynamo/test_compile.py
index 1f7290c51dd8d..97d54c3d28572 100644
--- a/test/dynamo/test_compile.py
+++ b/test/dynamo/test_compile.py
@@ -188,30 +188,66 @@ def fn(x):
         with self.assertRaises(AttributeError):
             fn(x)
 
-    @torch._dynamo.config.patch(inline_inbuilt_nn_modules=False)
-    def test_compilation_nn_module_invalid_method(self):
-        class Mod(torch.nn.Module):
-            def __init__(self):
-                super().__init__()
-
-            def forward(self, x):
-                return x + self.doesnotexist
-
-        mod = Mod()
-        opt_mod = torch.compile(mod, backend="eager")
-        x = torch.randn(1, 1)
-        with self.assertRaises(AttributeError):
-            opt_mod(x)
-
     def test_torch_script_compilation(self):
         @torch.jit.script
         def fn(x: torch.Tensor) -> torch.Tensor:
             return x
 
         a = torch.randn(1, 1)
-        out = torch.compile(fn)(a)
+        out = torch.compile(fn, backend="eager")(a)
         self.assertEqual(out, a)
 
+    def test_compile_script_module_error(self):
+        model = torch.nn.Sequential(torch.nn.Linear(3, 3))
+        model.eval()
+        scripted = torch.jit.script(model)
+        with self.assertRaisesRegex(
+            RuntimeError, "torch.compile does not support compiling torch.jit.script"
+        ):
+            torch.compile(scripted, backend="eager")
+
+    def test_compile_frozen_module_error(self):
+        model = torch.nn.Sequential(torch.nn.Linear(3, 3))
+        model.eval()
+        scripted = torch.jit.script(model)
+        with self.assertWarns(DeprecationWarning):
+            frozen = torch.jit.freeze(scripted)
+        with self.assertRaisesRegex(
+            RuntimeError, "torch.compile does not support compiling torch.jit.script"
+        ):
+            torch.compile(frozen, backend="eager")
+
+    def test_compile_frozen_module_inductor_error(self):
+        model = torch.nn.Sequential(torch.nn.Linear(3, 3))
+        model.eval()
+        scripted = torch.jit.script(model)
+        with self.assertWarns(DeprecationWarning):
+            frozen = torch.jit.freeze(scripted)
+        with self.assertRaisesRegex(
+            RuntimeError, "torch.compile does not support compiling torch.jit.script"
+        ):
+            torch.compile(frozen, backend="inductor")
+
+    def test_inline_script_module_graph_break(self):
+        class OuterModule(torch.nn.Module):
+            def __init__(self, sub):
+                super().__init__()
+                self.sub = sub
+
+            def forward(self, x):
+                return self.sub(x)
+
+        inner = torch.nn.Linear(3, 3)
+        scripted_inner = torch.jit.script(inner)
+        outer = OuterModule(scripted_inner)
+
+        cnt = torch._dynamo.testing.CompileCounter()
+        compiled = torch.compile(outer, backend=cnt)
+        inputs = torch.randn(2, 3)
+        compiled(inputs)
+        # ScriptModule submodule causes a graph break
+        self.assertEqual(cnt.frame_count, 0)
+
     def test_to_sparse_to_dense_with_graph_break(self):
         def fn(x):
             x = x.to_sparse()
@@ -219,7 +255,7 @@ def fn(x):
             return x
 
         x = torch.tensor([[1.0]])
-        c_fn = torch.compile(fn)
+        c_fn = torch.compile(fn, backend="eager")
 
         output = fn(x)
         c_output = c_fn(x)
@@ -272,5 +308,50 @@ def test_dynamo_signatures(self):
             self.check_signature(fn_name, fn_name, torch._dynamo)
 
 
+class FullgraphTests(TestCase):
+    def test_fullgraph_warns_on_frame_skip_with_dispatch_mode(self):
+        from torch.utils._python_dispatch import TorchDispatchMode
+
+        class SkipMode(TorchDispatchMode):
+            def __torch_dispatch__(self, func, types, args=(), kwargs=None):
+                return func(*args, **kwargs)
+
+        def fn(x):
+            return x + 1
+
+        x = torch.randn(5)
+        with SkipMode():
+            with self.assertWarnsRegex(UserWarning, "found no compiled frames"):
+                torch.compile(fn, backend="eager", fullgraph=True)(x)
+
+    def test_fullgraph_warns_on_frame_skip_dynamo_disabled(self):
+        def fn(x):
+            return x + 1
+
+        x = torch.randn(5)
+        with torch._dynamo.config.patch(disable=True):
+            with self.assertWarnsRegex(UserWarning, "found no compiled frames"):
+                torch.compile(fn, backend="eager", fullgraph=True)(x)
+
+    def test_fullgraph_empty_graph_no_error(self):
+        def fn(x):
+            return len(x)
+
+        x = torch.randn(5)
+        result = torch.compile(fn, backend="eager", fullgraph=True)(x)
+        self.assertEqual(result, 5)
+
+    def test_fullgraph_exported_module_no_error(self):
+        class M(torch.nn.Module):
+            def forward(self, x):
+                return x + 1
+
+        m = M()
+        x = torch.randn(5)
+        exported = torch.export.export(m, (x,))
+        result = exported.module()(x)
+        self.assertEqual(result, x + 1)
+
+
 if __name__ == "__main__":
     run_tests()
diff --git a/test/dynamo/test_compiler_bisector.py b/test/dynamo/test_compiler_bisector.py
index c1bc667a02041..153f4533c3146 100644
--- a/test/dynamo/test_compiler_bisector.py
+++ b/test/dynamo/test_compiler_bisector.py
@@ -29,8 +29,7 @@ def tearDown(self):
         if hasattr(torch.ops, self.test_ns):
             delattr(torch.ops, self.test_ns)
         if hasattr(self, "lib"):
-            del self.lib.m
-            del self.lib
+            self.lib._destroy()
 
     def get_op(self, name):
         return getattr(getattr(torch.ops, self.test_ns), name).default
@@ -44,7 +43,8 @@ def test_bad_decomp(self):
         import_module("torch._inductor.compile_fx")
 
         def bad_exp_decomp(self, rate=1, generator=None):
-            assert generator is None
+            if generator is not None:
+                raise AssertionError("Expected generator to be None")
             torch._check(
                 not utils.is_complex_dtype(self.dtype)
                 and not utils.is_integer_dtype(self.dtype)
@@ -103,7 +103,8 @@ def test_pre_grad(self):
         # similar setup to test_joint_graph (see below)
         def pass_fn(graph: torch.fx.Graph):
             nodes = graph.find_nodes(op="call_function", target=operator.add)
-            assert len(nodes) == 1
+            if len(nodes) != 1:
+                raise AssertionError(f"Expected 1 node, got {len(nodes)}")
             args = list(nodes[0].args)
             args[1] = 2
             nodes[0].args = tuple(args)
@@ -135,7 +136,8 @@ def pass_fn(graph: torch.fx.Graph):
             nodes = graph.find_nodes(
                 op="call_function", target=torch.ops.aten.add.Tensor
             )
-            assert len(nodes) == 1
+            if len(nodes) != 1:
+                raise AssertionError(f"Expected 1 node, got {len(nodes)}")
             args = list(nodes[0].args)
             args[1] = 2
             nodes[0].args = tuple(args)
@@ -303,7 +305,7 @@ def node_to_graph_id(node):
                     if "submod_" in name:
                         # the test case is simple enough that using
                         # the original example_inputs works for sub
-                        # moule
+                        # module
                         submod.forward = torch._inductor.standalone_compile(
                             submod,
                             example_inputs,
diff --git a/test/dynamo/test_comprehensions.py b/test/dynamo/test_comprehensions.py
index 97027592b523f..55d0cdda8c185 100644
--- a/test/dynamo/test_comprehensions.py
+++ b/test/dynamo/test_comprehensions.py
@@ -741,7 +741,6 @@ def fn(x):
         self.assertEqual(count_op(backend.graphs[0], operator.add), 1)
         self.assertEqual(count_op(backend.graphs[1], operator.add), 1)
 
-    @torch._dynamo.config.patch(nested_graph_breaks=False)
     def test_nested_function_calls_with_comprehension_graph_break(self):
         """Test nested function calls where inner function has comprehension with graph break."""
 
@@ -768,9 +767,9 @@ def f(x):
         x = torch.randn(4)
 
         self.assertEqual(compiled(x), f(x))
-        self.assertEqual(len(backend.graphs), 6)
-        for i in range(6):
-            self.assertEqual(count_op(backend.graphs[i], operator.add), 1)
+        self.assertEqual(len(backend.graphs), 2)
+        for i in range(2):
+            self.assertEqual(count_op(backend.graphs[i], operator.add), 3)
 
 
 @skipIfNotPy312
diff --git a/test/dynamo/test_config.py b/test/dynamo/test_config.py
index 733723bf7569d..28b1a6796235e 100644
--- a/test/dynamo/test_config.py
+++ b/test/dynamo/test_config.py
@@ -82,7 +82,8 @@ def test_config_compile_ignored(self):
             "debug_dir_root",
         }
         for k in dynamo_guarded_config_ignorelist:
-            assert k in torch._dynamo.config._compile_ignored_keys, k
+            if k not in torch._dynamo.config._compile_ignored_keys:
+                raise AssertionError(f"Expected {k} to be in _compile_ignored_keys")
 
     def test_config_hash(self):
         config = torch._dynamo.config
@@ -90,26 +91,51 @@ def test_config_hash(self):
 
         with config.patch({"verbose": not config.verbose}):
             new_hash = config.get_hash()
-            assert "verbose" in config._compile_ignored_keys
-            assert new_hash == starting_hash
+            if "verbose" not in config._compile_ignored_keys:
+                raise AssertionError("Expected 'verbose' in _compile_ignored_keys")
+            if new_hash != starting_hash:
+                raise AssertionError(
+                    f"Expected hash to remain {starting_hash}, got {new_hash}"
+                )
 
         new_hash = config.get_hash()
-        assert new_hash == starting_hash
+        if new_hash != starting_hash:
+            raise AssertionError(
+                f"Expected hash to remain {starting_hash}, got {new_hash}"
+            )
 
         with config.patch({"suppress_errors": not config.suppress_errors}):
             changed_hash = config.get_hash()
-            assert "suppress_errors" not in config._compile_ignored_keys
-            assert changed_hash != starting_hash
+            if "suppress_errors" in config._compile_ignored_keys:
+                raise AssertionError(
+                    "Expected 'suppress_errors' not in _compile_ignored_keys"
+                )
+            if changed_hash == starting_hash:
+                raise AssertionError(
+                    f"Expected hash to change from {starting_hash}, got {changed_hash}"
+                )
 
             # Test nested patch
             with config.patch({"verbose": not config.verbose}):
                 inner_changed_hash = config.get_hash()
-                assert inner_changed_hash == changed_hash
-                assert inner_changed_hash != starting_hash
+                if inner_changed_hash != changed_hash:
+                    raise AssertionError(
+                        f"Expected inner hash {inner_changed_hash} to equal {changed_hash}"
+                    )
+                if inner_changed_hash == starting_hash:
+                    raise AssertionError(
+                        f"Expected inner hash {inner_changed_hash} to differ from starting {starting_hash}"
+                    )
 
         newest_hash = config.get_hash()
-        assert changed_hash != newest_hash
-        assert newest_hash == starting_hash
+        if changed_hash == newest_hash:
+            raise AssertionError(
+                f"Expected changed_hash {changed_hash} to differ from newest_hash {newest_hash}"
+            )
+        if newest_hash != starting_hash:
+            raise AssertionError(
+                f"Expected newest_hash {newest_hash} to equal starting_hash {starting_hash}"
+            )
 
 
 if __name__ == "__main__":
diff --git a/test/dynamo/test_ctx_manager.py b/test/dynamo/test_ctx_manager.py
index 0a71597e431eb..bf1eb8c31365f 100644
--- a/test/dynamo/test_ctx_manager.py
+++ b/test/dynamo/test_ctx_manager.py
@@ -2,12 +2,18 @@
 import contextlib
 import sys
 import unittest
+from collections import defaultdict
 from contextlib import contextmanager
 
 import torch
 import torch._dynamo.test_case
 import torch._dynamo.testing
-from torch._dynamo.testing import EagerAndRecordGraphs, normalize_gm, same
+from torch._dynamo.testing import (
+    check_dynamic_shape_capture,
+    EagerAndRecordGraphs,
+    normalize_gm,
+    same,
+)
 from torch._dynamo.utils import counters
 from torch.nn import functional as F
 from torch.testing._internal.common_cuda import PLATFORM_SUPPORTS_FLASH_ATTENTION
@@ -17,6 +23,8 @@
 )
 
 
+device_type = acc.type if (acc := torch.accelerator.current_accelerator()) else "cpu"
+
 try:
     from . import test_functions
 except ImportError:
@@ -76,7 +84,7 @@ def customized_ctx_manager_with_graph_break(mode):
         torch._C._set_grad_enabled(prev)
 
 
-class CtxManagerTests(torch._dynamo.test_case.TestCaseWithNestedGraphBreaks):
+class CtxManagerTests(torch._dynamo.test_case.TestCase):
     def test_no_grad(self):
         def fn1(a, b):
             x = a + 1
@@ -154,9 +162,9 @@ def fn(x):
                     x = torch.mul(x, 5)
                     torch._dynamo.graph_break()
                     x = torch.sqrt(x)
-                    assert torch.is_grad_enabled()
-                assert not torch.is_grad_enabled()
-            assert torch.is_grad_enabled() == before
+                    assert torch.is_grad_enabled()  # noqa: S101
+                assert not torch.is_grad_enabled()  # noqa: S101
+            assert torch.is_grad_enabled() == before  # noqa: S101
             return x
 
         a = torch.randn([3, 4])
@@ -568,7 +576,7 @@ def fn(x):
         res = opt_fn(x)
         self.assertEqual(ref, res)
         self.assertEqual(cnts.frame_count, 1)
-        self.assertExpectedInline(str(cnts.op_count), """16""")
+        self.assertExpectedInline(str(cnts.op_count), """17""")
 
     @unittest.skipIf(not torch.cuda.is_available(), "requires cuda")
     def test_cuda_device(self):
@@ -583,6 +591,17 @@ def fn(x):
         res = opt_fn(x)
         self.assertEqual(ref, res)
 
+    @unittest.skipIf(not torch.cuda.is_available(), "requires cuda")
+    def test_cuda__exchange_device_args(self):
+        @torch.compile(backend="eager", fullgraph=True)
+        def fn(args, kwargs):
+            torch.cuda._exchange_device(*args, **kwargs)
+
+        initial_dev = torch.cuda.current_device()
+        for args, kwargs in (((), ()), ((0, 0), ()), ((), ("kwarg",))):
+            self.assertRaises(torch._dynamo.exc.Unsupported, fn, args, kwargs)
+            self.assertEqual(torch.cuda.current_device(), initial_dev)
+
     def test_autograd_profiler_enabled(self):
         def fn(x):
             if torch.autograd._profiler_enabled():
@@ -596,13 +615,15 @@ def fn(x):
 
         if torch.autograd._profiler_enabled():
             torch.autograd._disable_profiler()
-        assert not torch.autograd._profiler_enabled()
+        if torch.autograd._profiler_enabled():
+            raise AssertionError("Expected profiler to be disabled")
         ref = fn(x)
         res = opt_fn(x)
         self.assertTrue(same(ref, res))
 
         with torch.autograd.profiler.profile():
-            assert torch.autograd._profiler_enabled()
+            if not torch.autograd._profiler_enabled():
+                raise AssertionError("Expected profiler to be enabled")
             ref = fn(x)
             res = opt_fn(x)
             self.assertTrue(same(ref, res))
@@ -685,7 +706,7 @@ def test_autocast_sdpa(self):
         class MyModule(torch.nn.Module):
             def forward(self, query, key, value):
                 with torch.autocast("cpu"):
-                    with torch.autocast("cuda", dtype=torch.float32):
+                    with torch.autocast(device_type, dtype=torch.float32):
                         out = F.scaled_dot_product_attention(
                             query, key, value, None, 0.0, True
                         )
@@ -696,13 +717,31 @@ def forward(self, query, key, value):
         seq_len_k = 1
         head_dim = 8
         query = torch.ones(
-            1, 8, seq_len_q, head_dim, device="cuda", dtype=dtype, requires_grad=True
+            1,
+            8,
+            seq_len_q,
+            head_dim,
+            device=device_type,
+            dtype=dtype,
+            requires_grad=True,
         )
         key = torch.ones(
-            1, 8, seq_len_k, head_dim, device="cuda", dtype=dtype, requires_grad=True
+            1,
+            8,
+            seq_len_k,
+            head_dim,
+            device=device_type,
+            dtype=dtype,
+            requires_grad=True,
         )
         value = torch.ones(
-            1, 8, seq_len_k, head_dim, device="cuda", dtype=dtype, requires_grad=True
+            1,
+            8,
+            seq_len_k,
+            head_dim,
+            device=device_type,
+            dtype=dtype,
+            requires_grad=True,
         )
 
         module = MyModule()
@@ -716,7 +755,7 @@ def forward(self, query, key, value):
         self.assertEqual(compiled.device, real_device)
         self.assertEqual(compiled.dtype, real_dtype)
 
-        self.assertEqual(compiled.device.type, "cuda")
+        self.assertEqual(compiled.device.type, device_type)
         self.assertEqual(compiled.device.index, 0)
         self.assertEqual(compiled.dtype, torch.float32)
 
@@ -815,7 +854,7 @@ def forward(self, x):
                     # We remember to exit the inner autocast correctly to outer
                     # even after graph breaks
                     f_float16 = self.mm_breaks(a_float32, b_float32)
-                    assert f_float16.dtype == f_float16_1.dtype
+                    assert f_float16.dtype == f_float16_1.dtype  # noqa: S101
                 return f_float16, g_float32
 
         module = MyModule()
@@ -861,7 +900,7 @@ def forward(self, x):
                         g_float32 = torch.mm(a_float32, b_float32)
                     f_float16 = self.mm_breaks(a_float32, b_float32)
 
-                    assert (
+                    assert (  # noqa: S101
                         f_float16[0][0] == self.mm_not_break(a_float32, b_float32)[0][0]
                     )
                 return f_float16, g_float32
@@ -1010,6 +1049,348 @@ def fn(a, b):
         self.assertTrue(res[0].dtype == torch.float16)
         self.assertTrue(res[1].dtype == torch.float16)
 
+    def test__enter__exit_autocast(self):
+        def f(x, y):
+            m = torch.amp.autocast_mode._enter_autocast("cpu")
+            x = x @ y
+            torch.amp.autocast_mode._exit_autocast(m)
+            return x
+
+        eager = EagerAndRecordGraphs()
+        opt_f = torch.compile(f, backend=eager, fullgraph=True)
+        x = torch.randn(3, 3, dtype=torch.float32)
+        y = torch.randn(3, 3, dtype=torch.float32)
+        z = f(x, y)
+        opt_z = opt_f(x, y)
+        self.assertEqual(z, opt_z)
+        self.assertEqual(z.dtype, opt_z.dtype)
+        self.assertFalse(torch.is_autocast_enabled("cpu"))
+        graph = eager.graphs[0]
+        actual = normalize_gm(graph.print_readable(False))
+
+        if check_dynamic_shape_capture():
+            self.assertExpectedInline(
+                actual,
+                """\
+class GraphModule(torch.nn.Module):
+    def forward(self, s77: "Sym(s77)", L_x_: "f32[s77, s77]", L_y_: "f32[s77, s77]"):
+        l_x_ = L_x_
+        l_y_ = L_y_
+
+        _is_autocast_available = torch._C._is_autocast_available('cpu');  _is_autocast_available = None
+
+        set_autocast_enabled = torch.set_autocast_enabled('cpu', True);  set_autocast_enabled = None
+
+        set_autocast_dtype = torch.set_autocast_dtype('cpu', torch.bfloat16);  set_autocast_dtype = None
+
+        autocast_increment_nesting = torch.autocast_increment_nesting();  autocast_increment_nesting = None
+        set_autocast_cache_enabled = torch.set_autocast_cache_enabled(True);  set_autocast_cache_enabled = None
+
+        x: "bf16[s77, s77]" = l_x_ @ l_y_;  l_x_ = l_y_ = None
+
+        autocast_decrement_nesting = torch.autocast_decrement_nesting();  autocast_decrement_nesting = None
+
+        clear_autocast_cache = torch.clear_autocast_cache();  clear_autocast_cache = None
+
+        set_autocast_enabled_1 = torch.set_autocast_enabled('cpu', False);  set_autocast_enabled_1 = None
+
+        set_autocast_dtype_1 = torch.set_autocast_dtype('cpu', torch.bfloat16);  set_autocast_dtype_1 = None
+
+        set_autocast_cache_enabled_1 = torch.set_autocast_cache_enabled(True);  set_autocast_cache_enabled_1 = None
+        return (x,)
+""",  # NOQA: B950
+            )
+        else:
+            self.assertExpectedInline(
+                actual,
+                """\
+class GraphModule(torch.nn.Module):
+    def forward(self, L_x_: "f32[3, 3]", L_y_: "f32[3, 3]"):
+        l_x_ = L_x_
+        l_y_ = L_y_
+
+        _is_autocast_available = torch._C._is_autocast_available('cpu');  _is_autocast_available = None
+
+        set_autocast_enabled = torch.set_autocast_enabled('cpu', True);  set_autocast_enabled = None
+
+        set_autocast_dtype = torch.set_autocast_dtype('cpu', torch.bfloat16);  set_autocast_dtype = None
+
+        autocast_increment_nesting = torch.autocast_increment_nesting();  autocast_increment_nesting = None
+        set_autocast_cache_enabled = torch.set_autocast_cache_enabled(True);  set_autocast_cache_enabled = None
+
+        x: "bf16[3, 3]" = l_x_ @ l_y_;  l_x_ = l_y_ = None
+
+        autocast_decrement_nesting = torch.autocast_decrement_nesting();  autocast_decrement_nesting = None
+
+        clear_autocast_cache = torch.clear_autocast_cache();  clear_autocast_cache = None
+
+        set_autocast_enabled_1 = torch.set_autocast_enabled('cpu', False);  set_autocast_enabled_1 = None
+
+        set_autocast_dtype_1 = torch.set_autocast_dtype('cpu', torch.bfloat16);  set_autocast_dtype_1 = None
+
+        set_autocast_cache_enabled_1 = torch.set_autocast_cache_enabled(True);  set_autocast_cache_enabled_1 = None
+        return (x,)
+""",  # NOQA: B950
+            )
+
+    def test__enter__exit_autocast_graph_break(self):
+        def f(x, y, z):
+            m = torch.amp.autocast_mode._enter_autocast("cpu")
+            x = x @ y
+            torch._dynamo.graph_break()
+            x = x @ z
+            # At this point m is wrapped as an AutocastModeVariable, which will graph break on the __exit__ call
+            torch.amp.autocast_mode._exit_autocast(m)
+            return x
+
+        eager = EagerAndRecordGraphs()
+        opt_f = torch.compile(f, backend=eager, fullgraph=False)
+        x = torch.randn(3, 3, dtype=torch.float32)
+        y = torch.randn(3, 3, dtype=torch.float32)
+        z = torch.randn(3, 3, dtype=torch.float32)
+        out = f(x, y, z)
+        opt_out = opt_f(x, y, z)
+        self.assertEqual(out, opt_out)
+        self.assertEqual(out.dtype, opt_out.dtype)
+        self.assertFalse(torch.is_autocast_enabled("cpu"))
+        graph = eager.graphs[0]
+        actual = normalize_gm(graph.print_readable(False))
+
+        if check_dynamic_shape_capture():
+            self.assertExpectedInline(
+                actual,
+                """\
+class GraphModule(torch.nn.Module):
+    def forward(self, s77: "Sym(s77)", L_x_: "f32[s77, s77]", L_y_: "f32[s77, s77]"):
+        l_x_ = L_x_
+        l_y_ = L_y_
+
+        _is_autocast_available = torch._C._is_autocast_available('cpu');  _is_autocast_available = None
+
+        set_autocast_enabled = torch.set_autocast_enabled('cpu', True);  set_autocast_enabled = None
+
+        set_autocast_dtype = torch.set_autocast_dtype('cpu', torch.bfloat16);  set_autocast_dtype = None
+
+        autocast_increment_nesting = torch.autocast_increment_nesting();  autocast_increment_nesting = None
+        set_autocast_cache_enabled = torch.set_autocast_cache_enabled(True);  set_autocast_cache_enabled = None
+
+        x: "bf16[s77, s77]" = l_x_ @ l_y_;  l_x_ = l_y_ = None
+        return (x,)
+""",  # NOQA: B950
+            )
+        else:
+            self.assertExpectedInline(
+                actual,
+                """\
+class GraphModule(torch.nn.Module):
+    def forward(self, L_x_: "f32[3, 3]", L_y_: "f32[3, 3]"):
+        l_x_ = L_x_
+        l_y_ = L_y_
+
+        _is_autocast_available = torch._C._is_autocast_available('cpu');  _is_autocast_available = None
+
+        set_autocast_enabled = torch.set_autocast_enabled('cpu', True);  set_autocast_enabled = None
+
+        set_autocast_dtype = torch.set_autocast_dtype('cpu', torch.bfloat16);  set_autocast_dtype = None
+
+        autocast_increment_nesting = torch.autocast_increment_nesting();  autocast_increment_nesting = None
+        set_autocast_cache_enabled = torch.set_autocast_cache_enabled(True);  set_autocast_cache_enabled = None
+
+        x: "bf16[3, 3]" = l_x_ @ l_y_;  l_x_ = l_y_ = None
+        return (x,)
+""",  # NOQA: B950
+            )
+
+        # Doesn't include autocast functions, see comment above
+        graph = eager.graphs[1]
+        actual = normalize_gm(graph.print_readable(False))
+
+        if check_dynamic_shape_capture():
+            self.assertExpectedInline(
+                actual,
+                """\
+class GraphModule(torch.nn.Module):
+    def forward(self, s77: "Sym(s77)", L_x_: "bf16[s77, s77]", L_z_: "f32[s77, s77]"):
+        l_x_ = L_x_
+        l_z_ = L_z_
+
+        x: "bf16[s77, s77]" = l_x_ @ l_z_;  l_x_ = l_z_ = None
+
+        autocast_decrement_nesting = torch.autocast_decrement_nesting();  autocast_decrement_nesting = None
+
+        clear_autocast_cache = torch.clear_autocast_cache();  clear_autocast_cache = None
+
+        set_autocast_enabled = torch.set_autocast_enabled('cpu', False);  set_autocast_enabled = None
+
+        set_autocast_dtype = torch.set_autocast_dtype('cpu', torch.bfloat16);  set_autocast_dtype = None
+
+        set_autocast_cache_enabled = torch.set_autocast_cache_enabled(True);  set_autocast_cache_enabled = None
+        return (x,)
+""",  # NOQA: B950
+            )
+        else:
+            self.assertExpectedInline(
+                actual,
+                """\
+class GraphModule(torch.nn.Module):
+    def forward(self, L_x_: "bf16[3, 3]", L_z_: "f32[3, 3]"):
+        l_x_ = L_x_
+        l_z_ = L_z_
+
+        x: "bf16[3, 3]" = l_x_ @ l_z_;  l_x_ = l_z_ = None
+
+        autocast_decrement_nesting = torch.autocast_decrement_nesting();  autocast_decrement_nesting = None
+
+        clear_autocast_cache = torch.clear_autocast_cache();  clear_autocast_cache = None
+
+        set_autocast_enabled = torch.set_autocast_enabled('cpu', False);  set_autocast_enabled = None
+
+        set_autocast_dtype = torch.set_autocast_dtype('cpu', torch.bfloat16);  set_autocast_dtype = None
+
+        set_autocast_cache_enabled = torch.set_autocast_cache_enabled(True);  set_autocast_cache_enabled = None
+        return (x,)
+""",  # NOQA: B950
+            )
+
+    def test_autocast_low_level_api(self):
+        def f(x, y):
+            torch.set_autocast_enabled("cpu", True)
+            torch.set_autocast_dtype("cpu", torch.bfloat16)
+            torch.set_autocast_cache_enabled(True)
+            x = x @ y
+            torch.autocast_decrement_nesting()
+            torch.clear_autocast_cache()
+            torch.set_autocast_enabled("cpu", False)
+            return x
+
+        prev_enabled = torch.is_autocast_enabled("cpu")
+        prev_dtype = torch.get_autocast_dtype("cpu")
+        prev_cache = torch.is_autocast_cache_enabled()
+
+        try:
+            opt_f = torch.compile(f, backend="eager", fullgraph=True)
+            x = torch.randn(3, 3, dtype=torch.float32)
+            y = torch.randn(3, 3, dtype=torch.float32)
+            out = f(x, y)
+            opt_out = opt_f(x, y)
+            self.assertEqual(out, opt_out)
+            self.assertEqual(out.dtype, opt_out.dtype)
+            self.assertFalse(torch.is_autocast_enabled("cpu"))
+        finally:
+            torch.set_autocast_enabled("cpu", prev_enabled)
+            torch.set_autocast_dtype("cpu", prev_dtype)
+            torch.set_autocast_cache_enabled(prev_cache)
+
+    def test__enter__exit_autocast_function_mode(self):
+        class FunctionCount(torch.overrides.TorchFunctionMode):
+            def __init__(self):
+                self.counts = defaultdict(int)
+
+            def __torch_function__(self, func, types, args, kwargs=None):
+                self.counts[func] += 1
+                return func(*args, **(kwargs or {}))
+
+        def f(x, y):
+            m = torch.amp.autocast_mode._enter_autocast("cpu")
+            x = x @ y
+            torch.amp.autocast_mode._exit_autocast(m)
+            return x
+
+        opt_f = torch.compile(f, backend="eager", fullgraph=True)
+        x = torch.randn(3, 3, dtype=torch.float32)
+        y = torch.randn(3, 3, dtype=torch.float32)
+        with FunctionCount() as fc:
+            z = f(x, y)
+            self.assertEqual(fc.counts[torch.amp.autocast_mode._enter_autocast], 1)
+            self.assertEqual(fc.counts[torch.amp.autocast_mode._exit_autocast], 1)
+        with FunctionCount() as fc:
+            opt_z = opt_f(x, y)
+            self.assertEqual(fc.counts[torch.amp.autocast_mode._enter_autocast], 1)
+            self.assertEqual(fc.counts[torch.amp.autocast_mode._exit_autocast], 1)
+        self.assertEqual(z, opt_z)
+        self.assertEqual(z.dtype, opt_z.dtype)
+        self.assertFalse(torch.is_autocast_enabled("cpu"))
+
+    def test__enter__exit_autocast_non_idempotent(self):
+        # Recompile trick doesn't work with dynamic shapes
+        if check_dynamic_shape_capture():
+            return
+
+        def f(x, y):
+            with torch.amp.autocast("cpu"):
+                x = x @ y
+            return x
+
+        eager = EagerAndRecordGraphs()
+        opt_f = torch.compile(f, backend=eager, fullgraph=False)
+        x = torch.randn(3, 3, dtype=torch.float32)
+        y = torch.randn(3, 3, dtype=torch.float32)
+        out = f(x, y)
+        opt_out = opt_f(x, y)
+        self.assertEqual(out, opt_out)
+        self.assertEqual(out.dtype, opt_out.dtype)
+        self.assertFalse(torch.is_autocast_enabled("cpu"))
+        graph = eager.graphs[0]
+        actual = normalize_gm(graph.print_readable(False))
+        self.assertExpectedInline(
+            actual,
+            """\
+class GraphModule(torch.nn.Module):
+    def forward(self, L_x_: "f32[3, 3]", L_y_: "f32[3, 3]"):
+        l_x_ = L_x_
+        l_y_ = L_y_
+
+        _enter_autocast = torch.amp.autocast_mode._enter_autocast('cpu', None, True, None)
+
+        x: "bf16[3, 3]" = l_x_ @ l_y_;  l_x_ = l_y_ = None
+
+        _exit_autocast = torch.amp.autocast_mode._exit_autocast(_enter_autocast);  _enter_autocast = _exit_autocast = None
+        return (x,)
+""",  # NOQA: B950
+        )
+
+        # Recompiling will decompose the _enter_autocast and _exit_autocast calls to lower level autocast functions
+        eager = EagerAndRecordGraphs()
+        d = {}
+        exec(actual, globals(), d)
+        retraced = torch.compile(d["GraphModule"], backend=eager, fullgraph=True)
+        retraced_out = retraced()(x, y)[0]
+        self.assertEqual(out, retraced_out)
+        self.assertEqual(out.dtype, retraced_out.dtype)
+        self.assertFalse(torch.is_autocast_enabled("cpu"))
+
+        graph = eager.graphs[0]
+        actual = normalize_gm(graph.print_readable(False))
+        self.assertExpectedInline(
+            actual,
+            """\
+class GraphModule(torch.nn.Module):
+    def forward(self, L_L_x_: "f32[3, 3]", L_L_y_: "f32[3, 3]"):
+        l_l_x_ = L_L_x_
+        l_l_y_ = L_L_y_
+
+        _is_autocast_available = torch._C._is_autocast_available('cpu');  _is_autocast_available = None
+
+        set_autocast_enabled = torch.set_autocast_enabled('cpu', True);  set_autocast_enabled = None
+
+        set_autocast_dtype = torch.set_autocast_dtype('cpu', torch.bfloat16);  set_autocast_dtype = None
+
+        autocast_increment_nesting = torch.autocast_increment_nesting();  autocast_increment_nesting = None
+        set_autocast_cache_enabled = torch.set_autocast_cache_enabled(True);  set_autocast_cache_enabled = None
+        x: "bf16[3, 3]" = l_l_x_ @ l_l_y_;  l_l_x_ = l_l_y_ = None
+        autocast_decrement_nesting = torch.autocast_decrement_nesting();  autocast_decrement_nesting = None
+
+        clear_autocast_cache = torch.clear_autocast_cache();  clear_autocast_cache = None
+
+        set_autocast_enabled_1 = torch.set_autocast_enabled('cpu', False);  set_autocast_enabled_1 = None
+
+        set_autocast_dtype_1 = torch.set_autocast_dtype('cpu', torch.bfloat16);  set_autocast_dtype_1 = None
+
+        set_autocast_cache_enabled_1 = torch.set_autocast_cache_enabled(True);  set_autocast_cache_enabled_1 = None
+        return (x,)
+""",  # NOQA: B950
+        )
+
     @parametrize(
         "Ctx",
         [CustomizedCtxManagerWithGraphBreak, customized_ctx_manager_with_graph_break],
@@ -1250,11 +1631,15 @@ def fn(x, y):
         self.assertEqual(ref, res)
 
     def test_graph_break_inlining_autocast(self):
-        for device in ["cuda", "cpu"]:
+        for device in ["cuda", "cpu", "xpu"]:
             if device == "cuda" and not (
                 torch.cuda.is_available() and torch.cuda.is_bf16_supported()
             ):
                 continue
+            if device == "xpu" and not (
+                torch.xpu.is_available() and torch.xpu.is_bf16_supported()
+            ):
+                continue
             self._graph_break_inlining_autocast_test_helper(device)
 
     def test_disable_saved_tensors_hooks(self):
@@ -1446,6 +1831,36 @@ def forward(self, L_y_: "f32[]"):
 """,  # NOQA: B950
         )
 
+    def test__saved_tensors_hooks_disable(self):
+        def fn(x):
+            y = x + 1
+            torch._C._autograd._saved_tensors_hooks_disable("This is not supported")
+            y *= 2
+            torch._C._autograd._saved_tensors_hooks_enable()
+            return y
+
+        eager = EagerAndRecordGraphs()
+        torch.compile(fn, backend=eager, fullgraph=True)(torch.randn(()))
+        graph = eager.graphs[0]
+        actual = normalize_gm(graph.print_readable(False))
+        self.assertExpectedInline(
+            actual,
+            """\
+class GraphModule(torch.nn.Module):
+    def forward(self, L_x_: "f32[]"):
+        l_x_ = L_x_
+
+        y: "f32[]" = l_x_ + 1;  l_x_ = None
+
+        _saved_tensors_hooks_disable = torch._C._autograd._saved_tensors_hooks_disable('This is not supported');  _saved_tensors_hooks_disable = None
+
+        y *= 2;  y_1: "f32[]" = y;  y = None
+
+        _saved_tensors_hooks_enable = torch._C._autograd._saved_tensors_hooks_enable();  _saved_tensors_hooks_enable = None
+        return (y_1,)
+""",  # NOQA: B950
+        )
+
     def test_context_wrapping_grad_mode_decorator(self):
         ctx_wrappers = [(torch.enable_grad, True), (torch.no_grad, False)]
         for call in [True, False]:
@@ -1466,7 +1881,7 @@ def inner_func(x):
                             inner_func = ctx_wrapper(inner_func)
 
                         # Calling no_grad or enabled_grad should not mutate global state
-                        assert torch.is_grad_enabled() == mode_inverse
+                        assert torch.is_grad_enabled() == mode_inverse  # noqa: S101
 
                     with ctx_wrapper_inverse():
                         return inner_func(x)
@@ -1501,7 +1916,7 @@ def inner_func(x):
                                 return x.sin()
 
                         # Calling no_grad or enabled_grad should not mutate global state
-                        assert torch.is_grad_enabled() == mode_inverse
+                        assert torch.is_grad_enabled() == mode_inverse  # noqa: S101
 
                     with ctx_wrapper_inverse():
                         return inner_func(x)
@@ -1537,7 +1952,7 @@ def inner_func(x):
 
                         # Consuming set_grad_enabled by calling it on a function
                         # should not mutate global state
-                        assert torch.is_grad_enabled() == mode_inverse
+                        assert torch.is_grad_enabled() == mode_inverse  # noqa: S101
 
                     with torch.set_grad_enabled(mode_inverse):
                         return inner_func(x)
@@ -1900,9 +2315,7 @@ def fn(x):
         self.assertGreater(len(counters["graph_break"]), 0)
 
 
-class ContextlibContextManagerTests(
-    torch._dynamo.test_case.TestCaseWithNestedGraphBreaks
-):
+class ContextlibContextManagerTests(torch._dynamo.test_case.TestCase):
     def setUp(self):
         super().setUp()
         self._prev = torch._dynamo.config.enable_trace_contextlib
@@ -2861,6 +3274,90 @@ def fn(t):
         y = fn(t)
         self.assertEqual(y, t.sin())
 
+    @parametrize("gb", (True, False))
+    def test_functorch_low_level(self, gb):
+        def f(x, gb):
+            level = torch._C._functorch._grad_increment_nesting()
+            torch._C._functorch.set_inplace_requires_grad_allowed(True)
+            torch._functorch.eager_transforms._set_tensor_requires_grad(x)
+            if gb:
+                torch._dynamo.graph_break()
+            torch._C._functorch.set_inplace_requires_grad_allowed(False)
+            torch._C._functorch._grad_decrement_nesting()
+            return x + level
+
+        prev_inplace = torch._C._functorch.get_inplace_requires_grad_allowed()
+        prev_level = torch._C._functorch.maybe_current_level()
+        opt_f = torch.compile(f, fullgraph=not gb, backend="eager")
+        x = torch.randn(3, 3, requires_grad=False)
+        opt_y = opt_f(x, gb)
+        self.assertTrue(x.requires_grad)
+        y = f(x, gb)
+        self.assertEqual(y, opt_y)
+        self.assertEqual(torch._C._functorch.maybe_current_level(), prev_level)
+        self.assertEqual(
+            torch._C._functorch.get_inplace_requires_grad_allowed(), prev_inplace
+        )
+
+    def test_retrace_grad(self):
+        # Recompile trick doesn't work with dynamic shapes
+        if check_dynamic_shape_capture():
+            return
+
+        def fn(x):
+            return x.sin().sum()
+
+        def wrapper_fn(x):
+            return torch.func.grad(fn)(x)
+
+        x = torch.randn(3, 3)
+        eager = EagerAndRecordGraphs()
+        opt_f = torch.compile(wrapper_fn, backend=eager, fullgraph=True)
+        y = wrapper_fn(x)
+        opt_y = opt_f(x)
+        self.assertEqual(y, opt_y)
+        first_graph = normalize_gm(eager.graphs[0].print_readable(False))
+        self.assertExpectedInline(
+            first_graph,
+            """\
+class GraphModule(torch.nn.Module):
+    def forward(self, L_x_: "f32[3, 3]"):
+        l_x_ = L_x_
+
+        _saved_tensors_hooks_disable = torch._C._autograd._saved_tensors_hooks_disable("torch.func.{grad, vjp, jacrev, hessian} don't yet support saved tensor hooks. Please open an issue with your use case.");  _saved_tensors_hooks_disable = None
+        _grad_increment_nesting = torch._C._functorch._grad_increment_nesting();  _grad_increment_nesting = None
+
+        diff_args: "f32[3, 3]" = torch._C._functorch._wrap_for_grad(l_x_, 1);  l_x_ = None
+
+        set_inplace_requires_grad_allowed = torch._C._functorch.set_inplace_requires_grad_allowed(True);  set_inplace_requires_grad_allowed = None
+
+        _set_tensor_requires_grad: "f32[3, 3]" = torch._functorch.eager_transforms._set_tensor_requires_grad(diff_args);  _set_tensor_requires_grad = None
+
+        set_inplace_requires_grad_allowed_1 = torch._C._functorch.set_inplace_requires_grad_allowed(False);  set_inplace_requires_grad_allowed_1 = None
+
+        sin: "f32[3, 3]" = diff_args.sin()
+        output: "f32[]" = sin.sum();  sin = None
+
+        _autograd_grad = torch._functorch.eager_transforms._autograd_grad((output,), [diff_args], create_graph = True);  diff_args = None
+        grad_input: "f32[3, 3]" = _autograd_grad[0];  _autograd_grad = None
+
+        grad_input_1: "f32[3, 3]" = torch._C._functorch._unwrap_for_grad(grad_input, 1);  grad_input = None
+        output_1: "f32[]" = torch._C._functorch._unwrap_for_grad(output, 1);  output = output_1 = None
+
+        _grad_decrement_nesting = torch._C._functorch._grad_decrement_nesting();  _grad_decrement_nesting = None
+        _saved_tensors_hooks_enable = torch._C._autograd._saved_tensors_hooks_enable();  _saved_tensors_hooks_enable = None
+        return (grad_input_1,)
+""",  # NOQA: B950
+        )
+
+        d = {}
+        exec(first_graph, globals(), d)
+        retraced = torch.compile(d["GraphModule"], backend=eager, fullgraph=True)
+        retraced_out = retraced()(x)[0]
+        self.assertEqual(y, retraced_out)
+        retraced_graph = normalize_gm(eager.graphs[0].print_readable(False))
+        self.assertEqual(first_graph, retraced_graph)
+
 
 instantiate_parametrized_tests(CtxManagerTests)
 instantiate_parametrized_tests(ContextlibContextManagerTests)
diff --git a/test/dynamo/test_debug_utils.py b/test/dynamo/test_debug_utils.py
index 176294f2e7b36..dbf9884c594f1 100644
--- a/test/dynamo/test_debug_utils.py
+++ b/test/dynamo/test_debug_utils.py
@@ -325,26 +325,55 @@ def test_multiple_rules(self, device):
         result = self._run_with_override(device, "0:aot_eager;1:inductor;3:eager")
         self.assertEqual(result, ["aot_eager", "inductor", "eager"])
 
-    def test_first_rule_wins(self, device):
-        result = self._run_with_override(device, ">=0:aot_eager;>=1:inductor")
-        self.assertEqual(result, ["aot_eager", "aot_eager", "aot_eager", "aot_eager"])
+    def test_conflicting_rules_raise(self, device):
+        with self.assertRaisesRegex(
+            torch._dynamo.exc.InternalTorchDynamoError,
+            "Conflicting backend override",
+        ):
+            self._run_with_override(device, ">=0:aot_eager;>=1:inductor")
 
     def test_complex_config(self, device):
         result = self._run_with_override(device, "0:aot_eager;>=2:inductor")
         self.assertEqual(result, ["aot_eager", "inductor", "inductor"])
 
-    def test_inductor_with_mode(self, device):
-        result = self._run_with_override(device, ">=0:inductor:reduce-overhead")
-        self.assertEqual(
-            result,
-            [
-                "inductor:reduce-overhead",
-                "inductor:reduce-overhead",
-                "inductor:reduce-overhead",
-                "inductor:reduce-overhead",
-            ],
+    def test_override_with_backward(self, device):
+        """Verify that backend override works when backward compilation occurs."""
+        from torch._dynamo.graph_id_filter import (
+            _create_backend_router,
+            get_backend_override_for_compile_id,
         )
 
+        torch._dynamo.reset()
+        _create_backend_router.cache_clear()
+        overrides_applied = []
+        original_get_override = get_backend_override_for_compile_id
+
+        def tracking_get_override(compile_id, config_str):
+            result = original_get_override(compile_id, config_str)
+            if result is not None:
+                overrides_applied.append(compile_id.frame_id)
+            return result
+
+        def fn(x):
+            return (x * 2 + 1).sum()
+
+        with (
+            patch.object(
+                torch._dynamo.config, "debug_backend_override", ">=0:aot_eager"
+            ),
+            patch(
+                "torch._dynamo.output_graph.get_backend_override_for_compile_id",
+                tracking_get_override,
+            ),
+        ):
+            compiled_fn = torch.compile(fn, backend="eager")
+            x = torch.randn(10, device=device, requires_grad=True)
+            result = compiled_fn(x)
+            result.backward()
+
+        self.assertEqual(overrides_applied, [0])
+        self.assertIsNotNone(x.grad)
+
 
 instantiate_device_type_tests(
     TestBackendOverrideIntegration, globals(), only_for=["cpu", "cuda"]
@@ -412,6 +441,48 @@ def test_config_router_value_types(self, device):
         self.assertEqual(config["str_opt"], "hello")
         self.assertIsNone(config["none_opt"])
 
+    def test_config_router_aggregation(self, device):
+        from torch._dynamo.graph_id_filter import GraphConfigRouter
+
+        router = GraphConfigRouter("0:a=1;>=0:b=2")
+        # Graph 0 matches both rules, configs are merged
+        self.assertEqual(router.get_value_for_graph(0), {"a": 1, "b": 2})
+        # Graph 1 matches only the second rule
+        self.assertEqual(router.get_value_for_graph(1), {"b": 2})
+
+    def test_config_router_conflict_raises(self, device):
+        from torch._dynamo.graph_id_filter import GraphConfigRouter
+
+        with self.assertRaisesRegex(ValueError, "Conflicting config override"):
+            GraphConfigRouter("0:a=1;>=0:a=2")
+
+    def test_config_router_same_value_no_conflict(self, device):
+        from torch._dynamo.graph_id_filter import GraphConfigRouter
+
+        router = GraphConfigRouter("0:a=1;>=0:a=1")
+        self.assertEqual(router.get_value_for_graph(0), {"a": 1})
+        self.assertEqual(router.get_value_for_graph(1), {"a": 1})
+
+    def test_config_router_aggregation_multiple_rules(self, device):
+        from torch._dynamo.graph_id_filter import GraphConfigRouter
+
+        router = GraphConfigRouter("0:a=1;1:b=2;>=0:c=3")
+        self.assertEqual(router.get_value_for_graph(0), {"a": 1, "c": 3})
+        self.assertEqual(router.get_value_for_graph(1), {"b": 2, "c": 3})
+        self.assertEqual(router.get_value_for_graph(2), {"c": 3})
+
+    def test_backend_router_conflict_raises(self, device):
+        from torch._dynamo.graph_id_filter import GraphBackendRouter
+
+        with self.assertRaisesRegex(ValueError, "Conflicting backend override"):
+            GraphBackendRouter("0-5:eager;3-10:inductor")
+
+    def test_backend_router_same_backend_no_conflict(self, device):
+        from torch._dynamo.graph_id_filter import GraphBackendRouter
+
+        router = GraphBackendRouter("0:eager;>=0:eager")
+        self.assertIsNotNone(router.get_value_for_graph(0))
+
     def test_get_inductor_config_override_empty(self, device):
         from torch._dynamo.graph_id_filter import (
             get_inductor_config_override_for_compile_id,
@@ -425,17 +496,17 @@ def test_combined_backend_and_config_override(self, device):
         Test combining backend override with config override.
 
         Scenario: Default backend is eager, but override all graphs to use
-        inductor:reduce-overhead (cudagraphs enabled), and additionally
-        override graph 1 to use cudagraph_skip_dynamic_graphs=False.
+        inductor with cudagraphs enabled, and additionally override graph 1
+        to use cudagraph_skip_dynamic_graphs=False.
         """
         from torch._dynamo.graph_id_filter import (
             _create_backend_router,
-            _create_config_router,
+            _create_inductor_config_router,
         )
 
         torch._dynamo.reset()
         _create_backend_router.cache_clear()
-        _create_config_router.cache_clear()
+        _create_inductor_config_router.cache_clear()
 
         backends_used: list[str] = []
         configs_applied: list[dict] = []
@@ -456,35 +527,21 @@ def tracking_wrap(compiler_fn, config_patches):
             configs_applied.append(config_patches)
             return original_wrap(compiler_fn, config_patches)
 
-        # Build backend tracking map from config
-        backend_override = ">=0:inductor:reduce-overhead"
-        backend_str_map: dict[int, str] = {}
-        from torch._dynamo.graph_id_filter import GraphIdFilter
-
-        for rule_str in backend_override.split(";"):
-            if ":" not in rule_str:
-                continue
-            colon_idx = rule_str.find(":")
-            filter_str = rule_str[:colon_idx].strip()
-            backend_str = rule_str[colon_idx + 1 :].strip()
-            gf = GraphIdFilter(filter_str)
-            for gid in range(10):
-                if gid in gf and gid not in backend_str_map:
-                    backend_str_map[gid] = backend_str
-
+        backend_override = ">=0:inductor"
         from torch._dynamo.graph_id_filter import get_backend_override_for_compile_id
 
         original_get_backend = get_backend_override_for_compile_id
 
         def tracking_get_backend(compile_id, config_str):
             result = original_get_backend(compile_id, config_str)
-            if result is not None and compile_id.frame_id in backend_str_map:
-                backends_used.append(backend_str_map[compile_id.frame_id])
+            if result is not None:
+                backends_used.append("inductor")
             return result
 
         # Use both overrides:
-        # - Backend: all graphs use inductor:reduce-overhead
-        # - Config: graph 1 uses cudagraph_skip_dynamic_graphs=False
+        # - Backend: all graphs use inductor
+        # - Config: all graphs enable cudagraphs, graph 1 also disables
+        #   cudagraph_skip_dynamic_graphs
         with (
             patch.object(
                 torch._dynamo.config,
@@ -494,7 +551,7 @@ def tracking_get_backend(compile_id, config_str):
             patch.object(
                 torch._dynamo.config,
                 "debug_inductor_config_override",
-                "1:triton.cudagraph_skip_dynamic_graphs=False",
+                "1:triton.cudagraph_skip_dynamic_graphs=False;>=0:triton.cudagraphs=True",
             ),
             patch(
                 "torch._dynamo.output_graph.get_backend_override_for_compile_id",
@@ -506,19 +563,19 @@ def tracking_get_backend(compile_id, config_str):
             compiled_fn(torch.randn(10, device=device))
 
         self.assertEqual(len(backends_used), 3)
-        self.assertEqual(
-            backends_used,
-            [
-                "inductor:reduce-overhead",
-                "inductor:reduce-overhead",
-                "inductor:reduce-overhead",
-            ],
-        )
+        self.assertEqual(backends_used, ["inductor", "inductor", "inductor"])
 
-        self.assertEqual(len(configs_applied), 1)
+        # All matching rules are aggregated. Graph 1 matches both rules.
+        self.assertEqual(len(configs_applied), 3)
+        self.assertEqual(configs_applied[0], {"triton.cudagraphs": True})
         self.assertEqual(
-            configs_applied[0], {"triton.cudagraph_skip_dynamic_graphs": False}
+            configs_applied[1],
+            {
+                "triton.cudagraph_skip_dynamic_graphs": False,
+                "triton.cudagraphs": True,
+            },
         )
+        self.assertEqual(configs_applied[2], {"triton.cudagraphs": True})
 
     def test_multiple_config_overrides_with_backend(self, device):
         """
@@ -529,12 +586,12 @@ def test_multiple_config_overrides_with_backend(self, device):
         """
         from torch._dynamo.graph_id_filter import (
             _create_backend_router,
-            _create_config_router,
+            _create_inductor_config_router,
         )
 
         torch._dynamo.reset()
         _create_backend_router.cache_clear()
-        _create_config_router.cache_clear()
+        _create_inductor_config_router.cache_clear()
 
         backends_used: list[str] = []
         configs_applied: list[dict] = []
@@ -613,12 +670,230 @@ def tracking_get_backend(compile_id, config_str):
         self.assertIn({"triton.cudagraphs": False}, configs_applied)
         self.assertIn({"triton.cudagraph_skip_dynamic_graphs": False}, configs_applied)
 
+    def test_config_override_backward_propagation(self, device):
+        """
+        Verify that inductor config overrides are active at inductor compile
+        time for both forward and backward, across multiple graph breaks.
+        """
+        import torch._functorch.config
+        from torch._dynamo.graph_id_filter import _create_inductor_config_router
+        from torch._inductor import (
+            compile_fx as compile_fx_mod,
+            config as inductor_config,
+        )
+
+        torch._dynamo.reset()
+        _create_inductor_config_router.cache_clear()
+
+        TRACKED_CONFIGS = [
+            "triton.cudagraphs",
+            "triton.dense_indexing",
+            "triton.cudagraph_skip_dynamic_graphs",
+        ]
+
+        def _read_config(key):
+            obj = inductor_config
+            for part in key.split("."):
+                obj = getattr(obj, part)
+            return obj
+
+        baseline = {k: _read_config(k) for k in TRACKED_CONFIGS}
+        configs_at_compile: dict[tuple[int, bool], dict] = {}
+
+        original_compile_fx = compile_fx_mod.compile_fx
+        original_inner_compile = compile_fx_mod.compile_fx_inner
+
+        def tracking_inner_compile(gm, example_inputs, **kwargs):
+            compile_id = torch._guards.CompileContext.current_compile_id()
+            is_backward = kwargs.get("is_backward", False)
+            snapshot = {k: _read_config(k) for k in TRACKED_CONFIGS}
+            configs_at_compile[(compile_id.frame_id, is_backward)] = snapshot
+            return original_inner_compile(gm, example_inputs, **kwargs)
+
+        def tracking_compile_fx(model_, example_inputs_, *args, **kwargs):
+            # Inject tracking inner_compile so compile_fx's config.patch
+            # wrapping covers it for both forward and backward.
+            if "inner_compile" not in kwargs:
+                kwargs["inner_compile"] = tracking_inner_compile
+            return original_compile_fx(model_, example_inputs_, *args, **kwargs)
+
+        def fn(x):
+            y = x * 2 + 1
+            torch._dynamo.graph_break()
+            z = y.sin()
+            torch._dynamo.graph_break()
+            return z.exp().sum()
+
+        # Overlapping rules (all three configs default to False):
+        config_override = (
+            ">=1:triton.cudagraphs=True;"
+            "0-1:triton.dense_indexing=True;"
+            "1:triton.cudagraph_skip_dynamic_graphs=True"
+        )
+        expected_overrides = {
+            0: {"triton.dense_indexing": True},
+            1: {
+                "triton.cudagraphs": True,
+                "triton.dense_indexing": True,
+                "triton.cudagraph_skip_dynamic_graphs": True,
+            },
+            2: {
+                "triton.cudagraphs": True,
+            },
+        }
+
+        with (
+            patch.object(
+                torch._dynamo.config,
+                "debug_inductor_config_override",
+                config_override,
+            ),
+            patch.object(compile_fx_mod, "compile_fx", tracking_compile_fx),
+            patch.object(torch._functorch.config, "enable_autograd_cache", False),
+        ):
+            compiled_fn = torch.compile(fn)
+            x = torch.randn(10, device=device, requires_grad=True)
+            result = compiled_fn(x)
+            result.backward()
+
+        # Verify each graph has fwd+bwd, correct overrides, no cross-graph
+        # leak, and identical configs for forward and backward.
+        for gid in range(3):
+            self.assertIn((gid, False), configs_at_compile, f"graph {gid} fwd missing")
+            self.assertIn((gid, True), configs_at_compile, f"graph {gid} bwd missing")
+            expected = {**baseline, **expected_overrides[gid]}
+            for is_bw in [False, True]:
+                phase = "backward" if is_bw else "forward"
+                self.assertEqual(
+                    configs_at_compile[(gid, is_bw)],
+                    expected,
+                    f"graph {gid} {phase}: config mismatch",
+                )
+
+        self.assertIsNotNone(x.grad)
+
 
 instantiate_device_type_tests(
     TestInductorConfigOverrideIntegration, globals(), only_for=["cpu", "cuda"]
 )
 
 
+class TestConfigOverrideValidation(TestCase):
+    def setUp(self):
+        super().setUp()
+        from torch._dynamo.graph_id_filter import (
+            _validate_backend_names,
+            _validate_dynamo_config_keys,
+            _validate_inductor_config_keys,
+        )
+
+        _validate_backend_names.cache_clear()
+        _validate_dynamo_config_keys.cache_clear()
+        _validate_inductor_config_keys.cache_clear()
+        torch._dynamo.reset()
+
+    def tearDown(self):
+        torch._dynamo.reset()
+        super().tearDown()
+
+    @torch._dynamo.config.patch(
+        debug_backend_override="0:not_a_real_backend",
+    )
+    def test_invalid_backend_raises_on_compile(self):
+        def fn(x):
+            return x + 1
+
+        with self.assertRaisesRegex(ValueError, "not_a_real_backend"):
+            torch.compile(fn, backend="eager")(torch.randn(4))
+
+    @torch._dynamo.config.patch(
+        debug_dynamo_config_override="0:nonexistent_dynamo_option=True",
+    )
+    def test_invalid_dynamo_config_raises_on_compile(self):
+        def fn(x):
+            return x + 1
+
+        with self.assertRaisesRegex(ValueError, "nonexistent_dynamo_option"):
+            torch.compile(fn, backend="eager")(torch.randn(4))
+
+    @torch._dynamo.config.patch(
+        debug_inductor_config_override="0:nonexistent_inductor_option=True",
+    )
+    def test_invalid_inductor_config_raises_on_compile(self):
+        def fn(x):
+            return x + 1
+
+        with self.assertRaisesRegex(ValueError, "nonexistent_inductor_option"):
+            torch.compile(fn, backend="eager")(torch.randn(4))
+
+
+class TestDynamoConfigOverrideIntegration(TestCase):
+    def setUp(self):
+        super().setUp()
+        torch._dynamo.reset()
+
+    def tearDown(self):
+        torch._dynamo.reset()
+        super().tearDown()
+
+    @torch._dynamo.config.patch(
+        specialize_float=False,
+        verbose=False,
+        debug_dynamo_config_override=(
+            "0:specialize_float=True;1:verbose=True,recompile_limit=10"
+        ),
+    )
+    def test_dynamo_config_override_per_graph(self):
+        """Per-graph dynamo config overrides target the right graphs.
+
+        Graph 0: specialize_float overridden True (base False)
+        Graph 1: verbose+recompile_limit overridden (multiple keys)
+        Graph 2: no override, keeps base values
+        """
+        from torch._dynamo.graph_id_filter import _create_dynamo_config_router
+
+        _create_dynamo_config_router.cache_clear()
+
+        observed: dict[int, dict] = {}
+
+        def capturing_backend(gm, example_inputs):
+            fid = torch._guards.CompileContext.current_compile_id().frame_id
+            observed[fid] = {
+                "specialize_float": torch._dynamo.config.specialize_float,
+                "verbose": torch._dynamo.config.verbose,
+                "recompile_limit": torch._dynamo.config.recompile_limit,
+            }
+            return gm
+
+        def fn(x):
+            x = x + 1
+            torch._dynamo.graph_break()
+            x = x * 2
+            torch._dynamo.graph_break()
+            return x - 1
+
+        torch.compile(fn, backend=capturing_backend)(torch.randn(4))
+
+        self.assertTrue(observed[0]["specialize_float"])
+        self.assertFalse(observed[0]["verbose"])
+
+        self.assertFalse(observed[1]["specialize_float"])
+        self.assertTrue(observed[1]["verbose"])
+        self.assertEqual(observed[1]["recompile_limit"], 10)
+
+        self.assertFalse(observed[2]["specialize_float"])
+        self.assertFalse(observed[2]["verbose"])
+
+    def test_dynamo_config_override_warning(self):
+        from torch._dynamo.graph_id_filter import _create_dynamo_config_router
+
+        _create_dynamo_config_router.cache_clear()
+        with self.assertWarnsRegex(
+            UserWarning, "TORCH_COMPILE_OVERRIDE_DYNAMO_CONFIGS"
+        ):
+            _create_dynamo_config_router("0:specialize_float=True")
+
+
 if __name__ == "__main__":
     from torch._dynamo.test_case import run_tests
 
diff --git a/test/dynamo/test_decorators.py b/test/dynamo/test_decorators.py
index 58f5127499562..ed1bc740cf523 100644
--- a/test/dynamo/test_decorators.py
+++ b/test/dynamo/test_decorators.py
@@ -7,15 +7,11 @@
 from unittest.mock import patch
 
 import torch
-import torch._dynamo.config as config
 import torch._dynamo.testing
-from torch._dynamo.decorators import leaf_function
-from torch._dynamo.exc import IncorrectUsage, Unsupported
-from torch._dynamo.testing import normalize_gm
+from torch._dynamo.exc import Unsupported
 from torch._dynamo.utils import counters
 from torch.testing._internal.common_utils import (
     instantiate_parametrized_tests,
-    parametrize,
     skipIfWindows,
 )
 from torch.testing._internal.dynamo_pytree_test_utils import PytreeRegisteringTestCase
@@ -48,10 +44,8 @@ def fn(a):
 
     def test_disable_for_custom_op(self):
         import torch.library
-        from torch.library import Library
 
-        foo = Library("foo", "DEF")  # noqa: TOR901
-        try:
+        with torch.library._scoped_library("foo", "DEF") as foo:
             foo.define("custom(Tensor self) -> Tensor")
 
             # Dynamic shape data dependent operator. For static shape compilation, Dynamo
@@ -81,8 +75,6 @@ def fn(x):
                 self.assertEqual(ref, res)
             finally:
                 torch.ops.foo.custom = orig_custom
-        finally:
-            foo._destroy()
 
     def test_disable_ignores_outer_wraps(self):
         def orig_inner():
@@ -248,10 +240,10 @@ def fn(a):
         fn(torch.randn(10))
 
         # Check for graph break
-        self.assertEqual(cnts.frame_count, 3)
+        self.assertEqual(cnts.frame_count, 2)
 
     def test_incorrect_usage_disallow_in_graph(self):
-        with self.assertRaises(IncorrectUsage):
+        with self.assertRaisesRegex(RuntimeError, "disallow_in_graph is expected"):
 
             @torch._dynamo.disallow_in_graph
             def fn1(x):
@@ -655,15 +647,16 @@ def fn(x, y):
 
     def test_nonstrict_trace_no_action_at_a_distance(self):
         def trace_me(x):
+            x = x + 4
             torch._dynamo.graph_break()
-            return x + 42
+            return x + 8
 
         # No effect on traceability of `trace_me`
         torch._dynamo.nonstrict_trace(trace_me)
 
         def fn(x):
-            res = trace_me(x)
-            return res + 1
+            res = trace_me(x + 1)
+            return res + 2
 
         x = torch.randn(10)
         cnts = torch._dynamo.testing.CompileCounterWithBackend("aot_eager")
@@ -879,6 +872,25 @@ def fn(x, y):
                 "Invalid use of pytree_flatten with nonstrict_trace-ed function", str(e)
             )
 
+    def test_nonstrict_trace_nn_module_dict_input(self):
+        @torch._dynamo.nonstrict_trace
+        def trace_me(x, modules):
+            torch._dynamo.graph_break()
+            return modules["a"](x) + modules["b"](x)
+
+        def fn(x, modules):
+            return trace_me(x, modules) + 1
+
+        linear_a = torch.nn.Linear(4, 4)
+        linear_b = torch.nn.Linear(4, 4)
+        modules = {"a": linear_a, "b": linear_b}
+        x = torch.randn(4, 4)
+        opt_fn = torch.compile(fn, fullgraph=True, backend="aot_eager")
+
+        ref = fn(x, modules)
+        res = opt_fn(x, modules)
+        self.assertEqual(ref, res)
+
     def test_graph_break(self):
         cnts = torch._dynamo.testing.CompileCounter()
 
@@ -1067,7 +1079,10 @@ def inner():
 
         def codegen_return_with_pops(self, *args) -> list[Instruction]:
             insts = old_codegen_return(*args)
-            assert insts[-1].opname.startswith("RETURN")
+            if not insts[-1].opname.startswith("RETURN"):
+                raise AssertionError(
+                    f"Expected RETURN instruction, got {insts[-1].opname}"
+                )
             # to prevent infinite recursion
             if self.f_code.co_name != "inner":
                 insts[-1:-1] = [
@@ -1300,8 +1315,7 @@ def forward(self, a, *args):
     def _test_mark_static_address(self, guarded):
         # This test verifies that dynamo properly marks inputs as static
         # when using the mark_static_address API.
-        # For both inline_inbuilt_nn_modules True and False, we expect the
-        # tensor to be present in the buffers attribute of the graph.
+        # We expect the tensor to be present in the buffers attribute of the graph.
 
         compiles_with_buffers = 0
         compiles = 0
@@ -1338,15 +1352,9 @@ def fn(x):
         self.assertEqual(compiles, 2 if guarded else 1)
 
     def test_mark_static_address_guarded(self):
-        with torch._dynamo.config.patch("inline_inbuilt_nn_modules", True):
-            self._test_mark_static_address(guarded=True)
-
         self._test_mark_static_address(guarded=True)
 
     def test_mark_static_address_unguarded(self):
-        with torch._dynamo.config.patch("inline_inbuilt_nn_modules", True):
-            self._test_mark_static_address(guarded=False)
-
         self._test_mark_static_address(guarded=False)
 
     def test_class_methods(self):
@@ -1446,7 +1454,7 @@ def fn(x, y):
         y = torch.tensor([1])
         x = torch.tensor(1)
 
-        self.assertEqual(fn(x, y), torch.compile(fn)(x, y))
+        self.assertEqual(fn(x, y), torch.compile(fn, backend="eager")(x, y))
 
     def test_justknobs_check(self):
         def fn(x, y):
@@ -1462,6 +1470,31 @@ def fn(x, y):
         compiled_out = compiled_fn(x, y)
         self.assertEqual(eager_out, compiled_out)
 
+    def test_assume_constant_result_on_cached_property(self):
+        class MyModule(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self._cached_id = None
+
+            @property
+            def cached_id(self) -> int:
+                if self._cached_id is None:
+                    self._cached_id = 42
+                return self._cached_id
+
+            def forward(self, x):
+                return x * self.cached_id
+
+        torch._dynamo.assume_constant_result(MyModule.cached_id.fget)
+
+        model = MyModule()
+        x = torch.tensor([1.0, 2.0, 3.0])
+        compiled_model = torch.compile(model, backend="eager", fullgraph=True)
+        result_eager = model(x)
+        result_compiled = compiled_model(x)
+
+        self.assertTrue(torch.allclose(result_eager, result_compiled))
+
     def test_set_stance_aot_eager_then_compile(self):
         cnts = torch._dynamo.testing.CompileCounter()
 
@@ -1477,7 +1510,6 @@ def fn(x, y, z):
         # Would have been 4 without stance
         self.assertEqual(cnts.op_count, 2)
 
-    @torch._dynamo.config.patch("inline_inbuilt_nn_modules", True)
     def test_mark_static_nn_module(self):
         @torch._dynamo.mark_static
         class Mock(torch.nn.Module):
@@ -1968,7 +2000,7 @@ def f5(x):
 
         inp = torch.ones(3)
         self.assertEqual(f5(inp), inp + 7)
-        self.assertEqual(cnts.frame_count, 4)
+        self.assertEqual(cnts.frame_count, 2)
 
         def inner_f6(x):
             x = x + 2
@@ -1984,7 +2016,7 @@ def f6(x):
 
         cnts.clear()
         self.assertEqual(f6(inp), inp + 7)
-        self.assertEqual(cnts.frame_count, 3)
+        self.assertEqual(cnts.frame_count, 2)
 
         def inner_f7(x):
             x = x + 2
@@ -2024,7 +2056,7 @@ def f8(x):
 
         inp = torch.ones(3)
         self.assertEqual(f8(inp), inp + 7)
-        self.assertEqual(cnts.frame_count, 4)
+        self.assertEqual(cnts.frame_count, 3)
 
         def inner2_f9(x):
             x = x + 2
@@ -2087,7 +2119,7 @@ def f1(x):
 
         inp = torch.ones(3)
         self.assertEqual(f1(inp), inp + 7)
-        self.assertEqual(cnts.frame_count, 4)
+        self.assertEqual(cnts.frame_count, 2)
 
         def inner1_f2(x):
             x = x + 1
@@ -2207,6 +2239,7 @@ def fn(x):
 
         self.assertEqual(cnts.frame_count, 0)
 
+    @torch._dynamo.config.patch(nested_graph_breaks=False)
     def test_nested_compile_fullgraph(self):
         # Test that fullgraph=True cannot be toggled back by fullgraph=False
         inp = torch.ones(3)
@@ -2278,8 +2311,8 @@ def forward(self, inp):
             )
 
             # check the model is compilable
-            torch.compile(model)
-            torch.compile(other_model)
+            torch.compile(model, backend="eager")
+            torch.compile(other_model, backend="eager")
 
     def test_disable_class_and_instance_method(self):
         # Test that decorating a method at class definition time and then
@@ -2402,1292 +2435,20 @@ def forward(self, x) -> torch.Tensor:
         with self.assertRaises(RuntimeError):
             exported_model = torch.export.export(model, (inp,))
 
-    def _assert_models_equal(
-        self,
-        model_expected,
-        model_test,
-        x_expected,
-        x_test,
-    ):
-        out_expected = model_expected(x_expected)
-        out_test = model_test(x_test)
-        self.assertEqual(out_expected, out_test)
-
-        loss_expected = out_expected.sum()
-        loss_test = out_test.sum()
-        loss_expected.backward()
-        loss_test.backward()
-        self.assertEqual(x_expected.grad, x_test.grad)
-
-        expected_grads = {
-            name: param.grad for name, param in model_expected.named_parameters()
-        }
-        test_grads = {name: param.grad for name, param in model_test.named_parameters()}
-
-        self.assertEqual(set(expected_grads.keys()), set(test_grads.keys()))
-        for name in expected_grads:
-            if expected_grads[name] is not None:
-                self.assertEqual(
-                    expected_grads[name],
-                    test_grads[name],
-                    msg=f"Gradient mismatch for parameter {name}",
-                )
-
-    def _test_leaf_function_helper(self, mod_class, args_fn, loss_fn):
-        import torch.utils._pytree as pytree
-        from torch._dynamo.testing import AotEagerAndRecordGraphs, EagerAndRecordGraphs
-
-        mod_eager = mod_class()
-        mod_compile_eager = mod_class()
-        mod_compile_eager.load_state_dict(dict(mod_eager.state_dict()))
-        mod_compile_aot = mod_class()
-        mod_compile_aot.load_state_dict(dict(mod_eager.state_dict()))
-
-        eager_backend = EagerAndRecordGraphs()
-        compiled_eager = torch.compile(
-            mod_compile_eager, backend=eager_backend, fullgraph=True
-        )
-
-        backend = AotEagerAndRecordGraphs()
-        compiled_aot = torch.compile(mod_compile_aot, backend=backend, fullgraph=True)
-
-        for _ in range(2):
-            mod_eager.zero_grad()
-            mod_compile_eager.zero_grad()
-            mod_compile_aot.zero_grad()
-
-            args = args_fn()
-            args_clone = pytree.tree_map(
-                lambda x: x.clone().detach().requires_grad_(x.requires_grad), args
-            )
-            args_clone2 = pytree.tree_map(
-                lambda x: x.clone().detach().requires_grad_(x.requires_grad), args
-            )
-
-            out_eager = mod_eager(*args)
-            loss_fn(out_eager).backward()
-
-            out_compile_eager = compiled_eager(*args_clone)
-            loss_fn(out_compile_eager).backward()
-
-            out_compile_aot = compiled_aot(*args_clone2)
-            loss_fn(out_compile_aot).backward()
-
-            self.assertEqual(out_eager, out_compile_eager)
-            self.assertEqual(out_eager, out_compile_aot)
-
-            for (name_eager, param_eager), (_, param_compile_eager), (
-                _,
-                param_compile_aot,
-            ) in zip(
-                mod_eager.named_parameters(),
-                mod_compile_eager.named_parameters(),
-                mod_compile_aot.named_parameters(),
-            ):
-                self.assertEqual(
-                    param_eager.grad,
-                    param_compile_eager.grad,
-                    msg=f"Gradient mismatch for {name_eager} between eager and compile_eager",
-                )
-                self.assertEqual(
-                    param_eager.grad,
-                    param_compile_aot.grad,
-                    msg=f"Gradient mismatch for {name_eager} between eager and compile_aot",
-                )
-
-            pytree.tree_map(
-                lambda x, compile_x: self.assertEqual(x.grad, compile_x.grad)
-                if isinstance(x, torch.Tensor) and x.requires_grad
-                else None,
-                args,
-                args_clone,
-            )
-            pytree.tree_map(
-                lambda x, compile_x: self.assertEqual(x.grad, compile_x.grad)
-                if isinstance(x, torch.Tensor) and x.requires_grad
-                else None,
-                args,
-                args_clone2,
-            )
-
-        return (
-            normalize_gm(eager_backend.graphs[0].print_readable(print_output=False)),
-            normalize_gm(backend.fw_graphs[0].print_readable(print_output=False)),
-            normalize_gm(backend.bw_graphs[0].print_readable(print_output=False)),
-        )
-
-    def test_leaf_function_simple(self):
-        @leaf_function
-        def non_tracable_forward(mod, x):
-            if x.sum() > 0:
-                return (mod.linear(x),)
-            else:
-                return (mod.linear(x) + x,)
-
-        @non_tracable_forward.register_fake
-        def non_tracable_forward_fake(mod, x):
-            return (mod.linear(x),)
-
-        class NonTracable(torch.nn.Module):
-            def __init__(self):
-                super().__init__()
-                self.linear = torch.nn.Linear(3, 3)
-
-            def forward(self, x):
-                return non_tracable_forward(self, x)
-
-        def args_fn():
-            return (torch.randn(3, 3, requires_grad=True),)
-
-        def loss_fn(out):
-            return out[0].sum()
-
-        dynamo_graph_str, fw_graph_str, bw_graph_str = self._test_leaf_function_helper(
-            NonTracable, args_fn, loss_fn
-        )
-        self.assertExpectedInline(
-            dynamo_graph_str,
-            """\
-class GraphModule(torch.nn.Module):
-    def forward(self, L_x_: "f32[3, 3]", L_self_modules_linear_parameters_weight_: "f32[3, 3]", L_self_modules_linear_parameters_bias_: "f32[3]"):
-        l_x_ = L_x_
-        l_self_modules_linear_parameters_weight_ = L_self_modules_linear_parameters_weight_
-        l_self_modules_linear_parameters_bias_ = L_self_modules_linear_parameters_bias_
-
-        real_fn : torch.utils._pytree.TreeSpec = self.real_fn
-        fake_fn : torch.utils._pytree.TreeSpec = self.fake_fn
-        invoke_leaf_function = torch.ops.higher_order.invoke_leaf_function(real_fn, fake_fn, 0, l_self_modules_linear_parameters_weight_, l_self_modules_linear_parameters_bias_, l_x_);  real_fn = fake_fn = l_self_modules_linear_parameters_weight_ = l_self_modules_linear_parameters_bias_ = l_x_ = None
-        getitem: "f32[3, 3]" = invoke_leaf_function[0];  invoke_leaf_function = None
-        return (getitem,)
-""",  # noqa: B950
-        )
-        self.assertExpectedInline(
-            fw_graph_str,
-            """\
-class GraphModule(torch.nn.Module):
-    def forward(self, primals_1: "f32[0]", primals_2: "f32[3, 3]", primals_3: "f32[3, 3]", primals_4: "f32[3]"):
-        _tree_spec_constant0 = self._tree_spec_constant0
-        _tree_spec_constant1 = self._tree_spec_constant1
-        with_effects = torch.ops.higher_order.with_effects(primals_1, torch.ops.higher_order.invoke_leaf_function, _tree_spec_constant0, _tree_spec_constant1, 0, primals_3, primals_4, primals_2);  primals_1 = _tree_spec_constant0 = _tree_spec_constant1 = primals_3 = primals_4 = primals_2 = None
-
-        getitem: "f32[0]" = with_effects[0]
-        getitem_1: "f32[3, 3]" = with_effects[1];  with_effects = None
-        return (getitem, getitem_1)
-""",  # noqa: B950
-        )
-        self.assertExpectedInline(
-            bw_graph_str,
-            """\
-class GraphModule(torch.nn.Module):
-    def forward(self, tangents_1: "f32[3, 3]", tangents_token: "f32[0]"):
-        _tree_spec_constant2 = self._tree_spec_constant2
-        _tree_spec_constant3 = self._tree_spec_constant3
-        with_effects_1 = torch.ops.higher_order.with_effects(tangents_token, torch.ops.higher_order.invoke_leaf_function, _tree_spec_constant2, _tree_spec_constant3, tangents_1);  tangents_token = _tree_spec_constant2 = _tree_spec_constant3 = tangents_1 = None
-        getitem_2: "f32[0]" = with_effects_1[0]
-        getitem_4: "f32[3, 3]" = with_effects_1[2]
-        getitem_5: "f32[3]" = with_effects_1[3]
-        getitem_6: "f32[3, 3]" = with_effects_1[4];  with_effects_1 = None
-        return (getitem_6, getitem_4, getitem_5, getitem_2)
-""",  # noqa: B950
-        )
-
-    def test_leaf_function_with_logging(self):
-        @leaf_function
-        def logging_forward(mod, x):
-            print("Processing input")
-            return (mod.linear(x),)
-
-        @logging_forward.register_fake
-        def logging_forward_fake(mod, x):
-            return (mod.linear(x),)
-
-        class LoggingModule(torch.nn.Module):
-            def __init__(self):
-                super().__init__()
-                self.linear = torch.nn.Linear(3, 3)
-
-            def forward(self, x):
-                return logging_forward(self, x)
-
-        def args_fn():
-            return (torch.randn(3, 3, requires_grad=True),)
-
-        def loss_fn(out):
-            return out[0].sum()
-
-        with patch("builtins.print") as mock_print:
-            self._test_leaf_function_helper(LoggingModule, args_fn, loss_fn)
-            mock_print.assert_any_call("Processing input")
-            # Called 6 times: eager, compile_eager, and compile_aot, 2 iterations each
-            self.assertEqual(mock_print.call_count, 6)
-
-    def test_leaf_function_dynamic_autograd_module_config(self):
-        from torch._dynamo.testing import CompileCounterWithBackend
-
-        @leaf_function
-        def configurable_scale(mod, x):
-            # Branch based on module config, not input
-            if mod.use_double_scale:
-                return (mod.linear(x) * 2,)
-            else:
-                return (mod.linear(x) * 3,)
-
-        @configurable_scale.register_fake
-        def configurable_scale_fake(mod, x):
-            return (mod.linear(x),)
-
-        class ConfigurableModule(torch.nn.Module):
-            def __init__(self):
-                super().__init__()
-                self.linear = torch.nn.Linear(3, 3)
-                self.use_double_scale = True  # Config attribute
-
-            def forward(self, x):
-                return configurable_scale(self, x)
-
-        mod_eager = ConfigurableModule()
-        mod_compiled = ConfigurableModule()
-        mod_compiled.load_state_dict(dict(mod_eager.state_dict()))
-
-        counter = CompileCounterWithBackend("aot_eager")
-        compiled_fn = torch.compile(mod_compiled, backend=counter, fullgraph=True)
-
-        x_value = torch.randn(3, 3)
-
-        mod_eager.use_double_scale = True
-        mod_compiled.use_double_scale = True
-
-        x1 = x_value.clone().requires_grad_(True)
-        x1_clone = x_value.clone().requires_grad_(True)
-
-        out_eager_1 = mod_eager(x1)
-        out_eager_1[0].sum().backward()
-
-        out_compiled_1 = compiled_fn(x1_clone)
-        out_compiled_1[0].sum().backward()
-
-        self.assertEqual(out_eager_1, out_compiled_1)
-        self.assertEqual(x1.grad, x1_clone.grad)
-
-        mod_eager.zero_grad()
-        mod_compiled.zero_grad()
-
-        mod_eager.use_double_scale = False
-        mod_compiled.use_double_scale = False
-
-        x2 = x_value.clone().requires_grad_(True)
-        x2_clone = x_value.clone().requires_grad_(True)
-
-        out_eager_2 = mod_eager(x2)
-        out_eager_2[0].sum().backward()
-
-        out_compiled_2 = compiled_fn(x2_clone)
-        out_compiled_2[0].sum().backward()
-
-        self.assertEqual(out_eager_2, out_compiled_2)
-        self.assertEqual(x2.grad, x2_clone.grad)
-
-        # Same inputs but different config -> different gradients
-        # This proves leaf_function builds autograd dynamically (not burned in at trace time)
-        self.assertNotEqual(x1.grad, x2.grad)
-
-        # Verify only ONE compilation happened (no recompilation when changing config)
-        self.assertEqual(counter.frame_count, 1)
-
-    def test_leaf_function_dynamic_autograd_closure(self):
-        from torch._dynamo.testing import CompileCounterWithBackend
-
-        config = {"use_double_scale": True}
-
-        @leaf_function
-        def configurable_scale(x, y):
-            # Branch based on closure variable, not input
-            if config["use_double_scale"]:
-                return (x @ y * 2,)
-            else:
-                return (x @ y * 3,)
-
-        @configurable_scale.register_fake
-        def configurable_scale_fake(x, y):
-            return (x @ y,)
-
-        def fn(x, y):
-            return configurable_scale(x, y)
-
-        counter = CompileCounterWithBackend("aot_eager")
-        compiled_fn = torch.compile(fn, backend=counter, fullgraph=True)
-
-        x_value = torch.randn(3, 3)
-        y_value = torch.randn(3, 3)
-
-        config["use_double_scale"] = True
-
-        x1 = x_value.clone().requires_grad_(True)
-        y1 = y_value.clone().requires_grad_(True)
-        x1_clone = x_value.clone().requires_grad_(True)
-        y1_clone = y_value.clone().requires_grad_(True)
-
-        out_eager_1 = fn(x1, y1)
-        out_eager_1[0].sum().backward()
-
-        out_compiled_1 = compiled_fn(x1_clone, y1_clone)
-        out_compiled_1[0].sum().backward()
-
-        self.assertEqual(out_eager_1, out_compiled_1)
-        self.assertEqual(x1.grad, x1_clone.grad)
-        self.assertEqual(y1.grad, y1_clone.grad)
-
-        config["use_double_scale"] = False
-
-        x2 = x_value.clone().requires_grad_(True)
-        y2 = y_value.clone().requires_grad_(True)
-        x2_clone = x_value.clone().requires_grad_(True)
-        y2_clone = y_value.clone().requires_grad_(True)
-
-        out_eager_2 = fn(x2, y2)
-        out_eager_2[0].sum().backward()
-
-        out_compiled_2 = compiled_fn(x2_clone, y2_clone)
-        out_compiled_2[0].sum().backward()
-
-        self.assertEqual(out_eager_2, out_compiled_2)
-        self.assertEqual(x2.grad, x2_clone.grad)
-        self.assertEqual(y2.grad, y2_clone.grad)
-
-        # Same inputs but different closure -> different gradients
-        # This proves leaf_function builds autograd dynamically (not burned in at trace time)
-        self.assertNotEqual(x1.grad, x2.grad)
-        self.assertNotEqual(y1.grad, y2.grad)
-
-        # Verify only ONE compilation happened (no recompilation when changing closure)
-        self.assertEqual(counter.frame_count, 1)
-
-    def test_leaf_function_closure_constants_without_grad(self):
-        closure_scale = 2.0
-        closure_tensor = torch.tensor([1.0, 2.0, 3.0])
-
-        @leaf_function
-        def closure_forward(mod, x):
-            out = mod.linear(x) * closure_scale * mod.scale
-            out = out + closure_tensor + mod.offset
-            return (out,)
-
-        @closure_forward.register_fake
-        def closure_forward_fake(mod, x):
-            return (mod.linear(x) + mod.offset,)
-
-        class ClosureModule(torch.nn.Module):
-            def __init__(self):
-                super().__init__()
-                self.linear = torch.nn.Linear(3, 3)
-                self.scale = 3.0
-                self.offset = torch.nn.Parameter(torch.ones(3))
-
-            def forward(self, x):
-                return closure_forward(self, x)
-
-        def args_fn():
-            return (torch.randn(3, 3, requires_grad=True),)
-
-        def loss_fn(out):
-            return out[0].sum()
-
-        dynamo_graph_str, fw_graph_str, bw_graph_str = self._test_leaf_function_helper(
-            ClosureModule, args_fn, loss_fn
-        )
-        self.assertExpectedInline(
-            dynamo_graph_str,
-            """\
-class GraphModule(torch.nn.Module):
-    def forward(self, L_x_: "f32[3, 3]", L_self_parameters_offset_: "f32[3]", L_self_modules_linear_parameters_weight_: "f32[3, 3]", L_self_modules_linear_parameters_bias_: "f32[3]"):
-        l_x_ = L_x_
-        l_self_parameters_offset_ = L_self_parameters_offset_
-        l_self_modules_linear_parameters_weight_ = L_self_modules_linear_parameters_weight_
-        l_self_modules_linear_parameters_bias_ = L_self_modules_linear_parameters_bias_
-
-        real_fn : torch.utils._pytree.TreeSpec = self.real_fn
-        fake_fn : torch.utils._pytree.TreeSpec = self.fake_fn
-        invoke_leaf_function = torch.ops.higher_order.invoke_leaf_function(real_fn, fake_fn, 0, l_self_parameters_offset_, l_self_modules_linear_parameters_weight_, l_self_modules_linear_parameters_bias_, l_x_);  real_fn = fake_fn = l_self_parameters_offset_ = l_self_modules_linear_parameters_weight_ = l_self_modules_linear_parameters_bias_ = l_x_ = None
-        getitem: "f32[3, 3]" = invoke_leaf_function[0];  invoke_leaf_function = None
-        return (getitem,)
-""",  # noqa: B950
-        )
-        self.assertExpectedInline(
-            fw_graph_str,
-            """\
-class GraphModule(torch.nn.Module):
-    def forward(self, primals_1: "f32[0]", primals_2: "f32[3, 3]", primals_3: "f32[3]", primals_4: "f32[3, 3]", primals_5: "f32[3]"):
-        _tree_spec_constant0 = self._tree_spec_constant0
-        _tree_spec_constant1 = self._tree_spec_constant1
-        with_effects = torch.ops.higher_order.with_effects(primals_1, torch.ops.higher_order.invoke_leaf_function, _tree_spec_constant0, _tree_spec_constant1, 0, primals_3, primals_4, primals_5, primals_2);  primals_1 = _tree_spec_constant0 = _tree_spec_constant1 = primals_3 = primals_4 = primals_5 = primals_2 = None
-
-        getitem: "f32[0]" = with_effects[0]
-        getitem_1: "f32[3, 3]" = with_effects[1];  with_effects = None
-        return (getitem, getitem_1)
-""",  # noqa: B950
-        )
-        self.assertExpectedInline(
-            bw_graph_str,
-            """\
-class GraphModule(torch.nn.Module):
-    def forward(self, tangents_1: "f32[3, 3]", tangents_token: "f32[0]"):
-        _tree_spec_constant2 = self._tree_spec_constant2
-        _tree_spec_constant3 = self._tree_spec_constant3
-        with_effects_1 = torch.ops.higher_order.with_effects(tangents_token, torch.ops.higher_order.invoke_leaf_function, _tree_spec_constant2, _tree_spec_constant3, tangents_1);  tangents_token = _tree_spec_constant2 = _tree_spec_constant3 = tangents_1 = None
-        getitem_2: "f32[0]" = with_effects_1[0]
-        getitem_4: "f32[3]" = with_effects_1[2]
-        getitem_5: "f32[3, 3]" = with_effects_1[3]
-        getitem_6: "f32[3]" = with_effects_1[4]
-        getitem_7: "f32[3, 3]" = with_effects_1[5];  with_effects_1 = None
-        return (getitem_7, getitem_4, getitem_5, getitem_6, getitem_2)
-""",  # noqa: B950
-        )
-
-    def test_leaf_function_pytree_inputs(self):
-        @leaf_function
-        def pytree_forward(mod, inputs):
-            if inputs["x"].sum() > 0:
-                return (mod.linear(inputs["x"]), inputs["y"] + 1)
-            return (mod.linear(inputs["x"]) + inputs["y"], inputs["y"] - 1)
-
-        @pytree_forward.register_fake
-        def pytree_forward_fake(mod, inputs):
-            return (mod.linear(inputs["x"]), inputs["y"])
-
-        class PytreeModule(torch.nn.Module):
-            def __init__(self):
-                super().__init__()
-                self.linear = torch.nn.Linear(3, 3)
-
-            def forward(self, inputs):
-                return pytree_forward(self, inputs)
-
-        def args_fn():
-            return (
-                {
-                    "x": torch.randn(3, 3, requires_grad=True),
-                    "y": torch.randn(3, 3, requires_grad=True),
-                },
-            )
-
-        def loss_fn(out):
-            return out[0].sum() + out[1].sum()
-
-        self._test_leaf_function_helper(PytreeModule, args_fn, loss_fn)
-
-    def test_leaf_function_nested_annotations(self):
-        @leaf_function
-        def inner_leaf_forward(mod, x):
-            y = mod.linear(x)
-            return (y + x,)
-
-        @inner_leaf_forward.register_fake
-        def inner_leaf_forward_fake(mod, x):
-            return (mod.linear(x),)
-
-        class InnerLeaf(torch.nn.Module):
-            def __init__(self):
-                super().__init__()
-                self.linear = torch.nn.Linear(3, 3)
-
-            def forward(self, x):
-                return inner_leaf_forward(self, x)
-
-        @leaf_function
-        def outer_leaf_forward(mod, x):
-            z = mod.linear(x)
-            return mod.inner(z + x)
-
-        @outer_leaf_forward.register_fake
-        def outer_leaf_forward_fake(mod, x):
-            return mod.inner(mod.linear(x))
-
-        class OuterLeaf(torch.nn.Module):
-            def __init__(self):
-                super().__init__()
-                self.inner = InnerLeaf()
-                self.linear = torch.nn.Linear(3, 3)
-
-            def forward(self, x):
-                return outer_leaf_forward(self, x)
-
-        def args_fn():
-            return (torch.randn(3, 3, requires_grad=True),)
-
-        def loss_fn(out):
-            return out[0].sum()
-
-        dynamo_graph_str, fw_graph_str, bw_graph_str = self._test_leaf_function_helper(
-            OuterLeaf, args_fn, loss_fn
-        )
-        self.assertExpectedInline(
-            dynamo_graph_str,
-            """\
-class GraphModule(torch.nn.Module):
-    def forward(self, L_x_: "f32[3, 3]", L_self_modules_inner_modules_linear_parameters_weight_: "f32[3, 3]", L_self_modules_inner_modules_linear_parameters_bias_: "f32[3]", L_self_modules_linear_parameters_weight_: "f32[3, 3]", L_self_modules_linear_parameters_bias_: "f32[3]"):
-        l_x_ = L_x_
-        l_self_modules_inner_modules_linear_parameters_weight_ = L_self_modules_inner_modules_linear_parameters_weight_
-        l_self_modules_inner_modules_linear_parameters_bias_ = L_self_modules_inner_modules_linear_parameters_bias_
-        l_self_modules_linear_parameters_weight_ = L_self_modules_linear_parameters_weight_
-        l_self_modules_linear_parameters_bias_ = L_self_modules_linear_parameters_bias_
-
-        real_fn : torch.utils._pytree.TreeSpec = self.real_fn
-        fake_fn : torch.utils._pytree.TreeSpec = self.fake_fn
-        invoke_leaf_function = torch.ops.higher_order.invoke_leaf_function(real_fn, fake_fn, 0, l_self_modules_inner_modules_linear_parameters_weight_, l_self_modules_inner_modules_linear_parameters_bias_, l_self_modules_linear_parameters_weight_, l_self_modules_linear_parameters_bias_, l_x_);  real_fn = fake_fn = l_self_modules_inner_modules_linear_parameters_weight_ = l_self_modules_inner_modules_linear_parameters_bias_ = l_self_modules_linear_parameters_weight_ = l_self_modules_linear_parameters_bias_ = l_x_ = None
-        getitem: "f32[3, 3]" = invoke_leaf_function[0];  invoke_leaf_function = None
-        return (getitem,)
-""",  # noqa: B950
-        )
-        self.assertExpectedInline(
-            fw_graph_str,
-            """\
-class GraphModule(torch.nn.Module):
-    def forward(self, primals_1: "f32[0]", primals_2: "f32[3, 3]", primals_3: "f32[3, 3]", primals_4: "f32[3]", primals_5: "f32[3, 3]", primals_6: "f32[3]"):
-        _tree_spec_constant0 = self._tree_spec_constant0
-        _tree_spec_constant1 = self._tree_spec_constant1
-        with_effects = torch.ops.higher_order.with_effects(primals_1, torch.ops.higher_order.invoke_leaf_function, _tree_spec_constant0, _tree_spec_constant1, 0, primals_3, primals_4, primals_5, primals_6, primals_2);  primals_1 = _tree_spec_constant0 = _tree_spec_constant1 = primals_3 = primals_4 = primals_5 = primals_6 = primals_2 = None
-
-        getitem: "f32[0]" = with_effects[0]
-        getitem_1: "f32[3, 3]" = with_effects[1];  with_effects = None
-        return (getitem, getitem_1)
-""",  # noqa: B950
-        )
-        self.assertExpectedInline(
-            bw_graph_str,
-            """\
-class GraphModule(torch.nn.Module):
-    def forward(self, tangents_1: "f32[3, 3]", tangents_token: "f32[0]"):
-        _tree_spec_constant2 = self._tree_spec_constant2
-        _tree_spec_constant3 = self._tree_spec_constant3
-        with_effects_1 = torch.ops.higher_order.with_effects(tangents_token, torch.ops.higher_order.invoke_leaf_function, _tree_spec_constant2, _tree_spec_constant3, tangents_1);  tangents_token = _tree_spec_constant2 = _tree_spec_constant3 = tangents_1 = None
-        getitem_2: "f32[0]" = with_effects_1[0]
-        getitem_4: "f32[3, 3]" = with_effects_1[2]
-        getitem_5: "f32[3]" = with_effects_1[3]
-        getitem_6: "f32[3, 3]" = with_effects_1[4]
-        getitem_7: "f32[3]" = with_effects_1[5]
-        getitem_8: "f32[3, 3]" = with_effects_1[6];  with_effects_1 = None
-        return (getitem_8, getitem_4, getitem_5, getitem_6, getitem_7, getitem_2)
-""",  # noqa: B950
-        )
-
-    def test_leaf_function_data_dependent_nonzero(self):
-        @leaf_function
-        def nonzero_forward(mod, x):
-            out = mod.linear(x)
-            nonzero_indices = (out > 0).nonzero()
-            return (out, nonzero_indices)
-
-        @nonzero_forward.register_fake
-        def nonzero_forward_fake(mod, x):
-            out = mod.linear(x)
-            return out, (out > 0).nonzero()
-
-        class NonzeroModule(torch.nn.Module):
-            def __init__(self):
-                super().__init__()
-                self.linear = torch.nn.Linear(3, 3)
-
-            def forward(self, x):
-                return nonzero_forward(self, x)
-
-        class OuterModule(torch.nn.Module):
-            def __init__(self):
-                super().__init__()
-                self.pre_linear = torch.nn.Linear(3, 3)
-                self.nonzero_module = NonzeroModule()
-                self.scale = torch.nn.Parameter(torch.tensor(2.0))
-
-            def forward(self, x):
-                x = self.pre_linear(x)
-                x = torch.relu(x)
-                out, nonzero_indices = self.nonzero_module(x)
-                num_nonzero = nonzero_indices.shape[0]
-                scaled_out = out * self.scale + num_nonzero
-                return scaled_out, nonzero_indices
-
-        def args_fn():
-            return (torch.randn(3, 3, requires_grad=True),)
-
-        def loss_fn(out):
-            return out[0].sum()
-
-        self._test_leaf_function_helper(OuterModule, args_fn, loss_fn)
-
-    def test_leaf_function_data_dependent_item(self):
-        @leaf_function
-        def item_forward(mod, x):
-            out = mod.linear(x)
-            scalar_value = out.sum().item()
-            return (out, scalar_value)
-
-        @item_forward.register_fake
-        def item_forward_fake(mod, x):
-            out = mod.linear(x)
-            return (out, out.sum().item())
-
-        class ItemModule(torch.nn.Module):
-            def __init__(self):
-                super().__init__()
-                self.linear = torch.nn.Linear(3, 3)
-
-            def forward(self, x):
-                return item_forward(self, x)
-
-        def args_fn():
-            return (torch.randn(3, 3, requires_grad=True),)
+    def test_allow_in_graph_inside_compile_gives_clear_error(self):
+        # Regression test for https://github.com/pytorch/pytorch/issues/178511
+        # Calling allow_in_graph inside a compiled region is not supported.
+        # Verify the error message guides users to annotate before compilation.
+        def forward(x):
+            wrapped_fn = torch.compiler.allow_in_graph(my_custom_function)
+            return wrapped_fn(x)
 
-        def loss_fn(out):
-            return out[0].sum()
-
-        self._test_leaf_function_helper(ItemModule, args_fn, loss_fn)
-
-    @parametrize("backend", ["eager", "aot_eager"])
-    def test_leaf_function_multiple_compiled_submodules(self, backend):
-        @leaf_function
-        def leaf_forward(mod, x):
-            if x.sum() > 0:
-                return (mod.linear(x),)
-            else:
-                return (mod.linear(x) + x,)
-
-        @leaf_forward.register_fake
-        def leaf_forward_fake(mod, x):
-            return (mod.linear(x),)
-
-        class LeafModule(torch.nn.Module):
-            def __init__(self, in_features, out_features):
-                super().__init__()
-                self.linear = torch.nn.Linear(in_features, out_features)
-
-            def forward(self, x):
-                return leaf_forward(self, x)
-
-        class CompiledSubmodule1(torch.nn.Module):
-            def __init__(self):
-                super().__init__()
-                self.pre_linear = torch.nn.Linear(4, 4)
-                self.leaf = LeafModule(4, 4)
-
-            def forward(self, x):
-                x = self.pre_linear(x)
-                x = torch.relu(x)
-                out = self.leaf(x)[0]
-                return out
-
-        class CompiledSubmodule2(torch.nn.Module):
-            def __init__(self):
-                super().__init__()
-                self.leaf = LeafModule(4, 4)
-                self.post_linear = torch.nn.Linear(4, 4)
-
-            def forward(self, x):
-                out = self.leaf(x)[0]
-                out = self.post_linear(out)
-                return torch.sigmoid(out)
-
-        class CompiledSubmodule3(torch.nn.Module):
-            def __init__(self):
-                super().__init__()
-                self.leaf1 = LeafModule(4, 4)
-                self.leaf2 = LeafModule(4, 4)
-
-            def forward(self, x):
-                out1 = self.leaf1(x)[0]
-                out2 = self.leaf2(x)[0]
-                return out1 + out2
-
-        class TopLevelModule(torch.nn.Module):
-            def __init__(self, compile_submodules=False):
-                super().__init__()
-                self.submodule1 = CompiledSubmodule1()
-                self.submodule2 = CompiledSubmodule2()
-                self.submodule3 = CompiledSubmodule3()
-                self.final_linear = torch.nn.Linear(4, 4)
-                self.compile_submodules = compile_submodules
-
-            def forward(self, x):
-                if self.compile_submodules:
-                    out1 = torch.compile(self.submodule1, backend=backend)(x)
-                    out2 = torch.compile(self.submodule2, backend=backend)(out1)
-                    out3 = torch.compile(self.submodule3, backend=backend)(out2)
-                else:
-                    out1 = self.submodule1(x)
-                    out2 = self.submodule2(out1)
-                    out3 = self.submodule3(out2)
-                final = self.final_linear(out3)
-                return final
-
-        model_eager = TopLevelModule(compile_submodules=False)
-        model_compiled = TopLevelModule(compile_submodules=True)
-        model_compiled.load_state_dict(model_eager.state_dict())
-
-        x = torch.randn(2, 4, requires_grad=True)
-        x_compiled = x.clone().detach().requires_grad_(True)
-
-        self._assert_models_equal(
-            model_eager,
-            model_compiled,
-            x,
-            x_compiled,
-        )
-
-    @parametrize("backend", ["eager", "aot_eager"])
-    @parametrize("do_compile", [False, True])
-    def test_leaf_function_with_graph_breaks(self, backend, do_compile):
-        @leaf_function
-        def leaf_forward(mod, x):
-            if x.sum() > 0:
-                return (mod.linear(x),)
-            else:
-                return (mod.linear(x) + 1,)
-
-        @leaf_forward.register_fake
-        def leaf_forward_fake(mod, x):
-            return (mod.linear(x),)
-
-        class LeafModule(torch.nn.Module):
-            def __init__(self, in_features, out_features):
-                super().__init__()
-                self.linear = torch.nn.Linear(in_features, out_features)
-
-            def forward(self, x):
-                return leaf_forward(self, x)
-
-        class TopLevelModule(torch.nn.Module):
-            def __init__(self, do_compile=False, backend="eager"):
-                super().__init__()
-                self.leaf1 = LeafModule(4, 4)
-                self.leaf2 = LeafModule(4, 4)
-                self.leaf3 = LeafModule(4, 4)
-                self.final_linear = torch.nn.Linear(4, 4)
-                self.do_compile = do_compile
-                self.backend = backend
-
-            def _forward(self, x):
-                out1 = self.leaf1(x)[0]
-                torch._dynamo.graph_break()
-                out2 = self.leaf2(out1)[0]
-                torch._dynamo.graph_break()
-                out3 = self.leaf3(out2)[0]
-                result = self.final_linear(out3)
-                return result
-
-            def forward(self, x):
-                if self.do_compile:
-                    return torch.compile(
-                        self._forward, backend=self.backend, fullgraph=False
-                    )(x)
-                else:
-                    return self._forward(x)
-
-        model_eager = TopLevelModule(do_compile=False)
-        model_test = TopLevelModule(do_compile=do_compile, backend=backend)
-        model_test.load_state_dict(model_eager.state_dict())
-
-        x = torch.randn(2, 4, requires_grad=True)
-        x_test = x.clone().detach().requires_grad_(True)
-
-        self._assert_models_equal(model_eager, model_test, x, x_test)
-
-    def test_leaf_function_with_module_in_pytree(self):
-        @leaf_function
-        def main_forward(modules_dict, x):
-            if x.sum() > 0:
-                return (modules_dict["first"](x) + modules_dict["second"](x),)
-            else:
-                return (modules_dict["first"](x) - modules_dict["second"](x),)
-
-        @main_forward.register_fake
-        def main_forward_fake(modules_dict, x):
-            return (modules_dict["first"](x) + modules_dict["second"](x),)
-
-        class HelperModule(torch.nn.Module):
-            def __init__(self, scale=1.0):
-                super().__init__()
-                self.linear = torch.nn.Linear(3, 3)
-                self.scale = scale
-
-            def forward(self, x):
-                return self.linear(x) * self.scale
-
-        class WrapperModule(torch.nn.Module):
-            def __init__(self):
-                super().__init__()
-                self.helper1 = HelperModule(scale=1.0)
-                self.helper2 = HelperModule(scale=0.5)
-
-            def forward(self, x):
-                modules_dict = {"first": self.helper1, "second": self.helper2}
-                return main_forward(modules_dict, x)
-
-        def args_fn():
-            return (torch.randn(3, 3, requires_grad=True),)
-
-        def loss_fn(out):
-            return out[0].sum()
-
-        self._test_leaf_function_helper(WrapperModule, args_fn, loss_fn)
-
-    def test_leaf_function_with_module_as_kwarg(self):
-        @leaf_function
-        def main_forward(x, helper_mod=None):
-            if x.sum() > 0:
-                return (helper_mod(x),)
-            else:
-                return (helper_mod(x) + x,)
-
-        @main_forward.register_fake
-        def main_forward_fake(x, helper_mod=None):
-            return (helper_mod(x),)
-
-        class HelperModule(torch.nn.Module):
-            def __init__(self):
-                super().__init__()
-                self.linear = torch.nn.Linear(3, 3)
-
-            def forward(self, x):
-                return self.linear(x)
-
-        class WrapperModule(torch.nn.Module):
-            def __init__(self):
-                super().__init__()
-                self.helper = HelperModule()
-
-            def forward(self, x):
-                return main_forward(x, helper_mod=self.helper)
-
-        def args_fn():
-            return (torch.randn(3, 3, requires_grad=True),)
-
-        def loss_fn(out):
-            return out[0].sum()
-
-        self._test_leaf_function_helper(WrapperModule, args_fn, loss_fn)
-
-    def test_leaf_function_missing_fake_impl_error(self):
-        @leaf_function
-        def no_fake_impl_forward(mod, x):
-            return (mod.linear(x),)
-
-        class SimpleModule(torch.nn.Module):
-            def __init__(self):
-                super().__init__()
-                self.linear = torch.nn.Linear(3, 3)
-
-            def forward(self, x):
-                return no_fake_impl_forward(self, x)
-
-        mod = SimpleModule()
-        x = torch.randn(3, 3)
-
-        result = mod(x)
-        self.assertEqual(result[0].shape, (3, 3))
-
-        compiled_mod = torch.compile(mod, backend="eager", fullgraph=True)
-        with self.assertRaisesRegex(Exception, "requires a fake implementation"):
-            compiled_mod(x)
-
-    @parametrize("backend", ["eager", "aot_eager"])
-    def test_leaf_function_constant_tensor_closure_error(self, backend):
-        constant_weight = torch.randn(3, 3)
-
-        @leaf_function
-        def constant_closure_forward(x):
-            return (x @ constant_weight,)
-
-        @constant_closure_forward.register_fake
-        def constant_closure_forward_fake(x):
-            return (x @ constant_weight,)
-
-        class ConstantClosureModule(torch.nn.Module):
-            def __init__(self):
-                super().__init__()
-
-            def forward(self, x):
-                return constant_closure_forward(x)
-
-        mod = ConstantClosureModule()
-        x = torch.randn(3, 3, requires_grad=True)
-
-        result = mod(x)
-        expected = x @ constant_weight
-        self.assertEqual(result[0], expected)
-
-        compiled_mod = torch.compile(mod, backend=backend, fullgraph=True)
+        compiled = torch.compile(forward, fullgraph=True)
         with self.assertRaisesRegex(
-            Exception, "Please convert all Tensors to FakeTensors"
+            torch._dynamo.exc.Unsupported,
+            "allow_in_graph",
         ):
-            compiled_mod(x)
-
-    @parametrize("backend", ["eager", "aot_eager"])
-    def test_leaf_function_input_mutation_error(self, backend):
-        @leaf_function
-        def mutate_input(x):
-            x.add_(1)
-            return (x,)
-
-        @mutate_input.register_fake
-        def mutate_input_fake(x):
-            return (x + 1,)
-
-        def fn(x):
-            return mutate_input(x)
-
-        x = torch.randn(3, 3)
-
-        x_eager = x.clone()
-        result_eager = fn(x_eager)
-        self.assertEqual(result_eager[0], x + 1)
-
-        compiled_fn = torch.compile(fn, backend=backend, fullgraph=True)
-        with self.assertRaisesRegex(RuntimeError, "In-place mutation detected"):
-            compiled_fn(x.clone())
-
-    @parametrize("backend", ["eager", "aot_eager"])
-    def test_leaf_function_validation_dtype_mismatch(self, backend):
-        @leaf_function
-        def dtype_mismatch_forward(mod, x):
-            return (mod.linear(x),)
-
-        @dtype_mismatch_forward.register_fake
-        def dtype_mismatch_forward_fake(mod, x):
-            return (mod.linear(x).double(),)
-
-        class DtypeMismatchModule(torch.nn.Module):
-            def __init__(self):
-                super().__init__()
-                self.linear = torch.nn.Linear(3, 3)
-
-            def forward(self, x):
-                return dtype_mismatch_forward(self, x)
-
-        mod = DtypeMismatchModule()
-        x = torch.randn(3, 3)
-
-        with config.patch(leaf_function_validate_outputs=True):
-            compiled_mod = torch.compile(mod, backend=backend)
-            with self.assertRaisesRegex(RuntimeError, "Dtype mismatch"):
-                compiled_mod(x)
-
-    @parametrize("backend", ["eager", "aot_eager"])
-    @parametrize("validate_outputs", [True, False])
-    def test_leaf_function_validation_shape_mismatch(self, backend, validate_outputs):
-        @leaf_function
-        def mismatched_forward(mod, x):
-            return (mod.linear(x),)
-
-        @mismatched_forward.register_fake
-        def mismatched_forward_fake(mod, x):
-            return (torch.zeros(x.shape[0], 6),)
-
-        class MismatchedModule(torch.nn.Module):
-            def __init__(self):
-                super().__init__()
-                self.linear = torch.nn.Linear(3, 3)
-
-            def forward(self, x):
-                return mismatched_forward(self, x)
-
-        mod = MismatchedModule()
-        x = torch.randn(3, 3)
-
-        with config.patch(leaf_function_validate_outputs=validate_outputs):
-            compiled_mod = torch.compile(mod, backend=backend)
-            if validate_outputs:
-                with self.assertRaises((RuntimeError, AssertionError)):
-                    compiled_mod(x)
-            else:
-                result = compiled_mod(x)
-                self.assertEqual(result[0].shape, (3, 3))
-
-    def test_leaf_function_no_module_inputs(self):
-        @leaf_function
-        def my_custom_fn(inputs: dict[str, torch.Tensor], scale: float, offset: int):
-            x = inputs["x"]
-            y = inputs["y"]
-            if x.sum() > 0:
-                return (x * scale + y + offset, x.sum() + y.sum())
-            return (x * scale - y + offset, x.sum() - y.sum())
-
-        @my_custom_fn.register_fake
-        def my_custom_fn_fake(
-            inputs: dict[str, torch.Tensor], scale: float, offset: int
-        ):
-            x = inputs["x"]
-            y = inputs["y"]
-            return (x * scale + y + offset, x.sum() + y.sum())
-
-        class NoModuleInputsModule(torch.nn.Module):
-            def __init__(self):
-                super().__init__()
-                self.scale = 2.0
-                self.offset = 1
-
-            def forward(self, x, y):
-                inputs = {"x": x, "y": y}
-                return my_custom_fn(inputs, self.scale, self.offset)
-
-        def args_fn():
-            return (
-                torch.randn(3, 3, requires_grad=True),
-                torch.randn(3, 3, requires_grad=True),
-            )
-
-        def loss_fn(out):
-            return out[0].sum() + out[1].sum()
-
-        self._test_leaf_function_helper(NoModuleInputsModule, args_fn, loss_fn)
-
-    @parametrize("backend", ["eager", "aot_eager"])
-    @parametrize("check_escaped_gradients", [True, False])
-    def test_leaf_function_escaped_gradient_multiple_tensors(
-        self, backend, check_escaped_gradients
-    ):
-        weight1 = torch.randn(3, 3, requires_grad=True)
-        weight2 = torch.randn(3, 3, requires_grad=True)
-
-        @leaf_function
-        def uses_multiple_closures(x):
-            return (x @ weight1 + x @ weight2,)
-
-        @uses_multiple_closures.register_fake
-        def uses_multiple_closures_fake(x):
-            return (torch.empty(x.shape[0], 3),)
-
-        def fn(x):
-            return uses_multiple_closures(x)
-
-        x = torch.randn(2, 3, requires_grad=True)
-
-        compiled_fn = torch.compile(fn, backend=backend, fullgraph=True)
-        with config.patch(
-            leaf_function_check_escaped_gradients=check_escaped_gradients
-        ):
-            if check_escaped_gradients:
-                with self.assertRaisesRegex(RuntimeError, "2 tensor"):
-                    compiled_fn(x)
-            else:
-                result = compiled_fn(x)
-                self.assertEqual(result[0].shape, (2, 3))
-
-    @parametrize("backend", ["eager", "aot_eager"])
-    @parametrize("check_escaped_gradients", [True, False])
-    def test_leaf_function_escaped_gradient_input_no_grad(
-        self, backend, check_escaped_gradients
-    ):
-        closure_weight = torch.randn(3, 3, requires_grad=True)
-
-        @leaf_function
-        def uses_closure(x):
-            return (x @ closure_weight,)
-
-        @uses_closure.register_fake
-        def uses_closure_fake(x):
-            return (torch.empty(x.shape[0], 3),)
-
-        def fn(x):
-            return uses_closure(x)
-
-        x = torch.randn(2, 3, requires_grad=False)
-
-        compiled_fn = torch.compile(fn, backend=backend, fullgraph=True)
-        with config.patch(
-            leaf_function_check_escaped_gradients=check_escaped_gradients
-        ):
-            result = compiled_fn(x)
-            self.assertEqual(result[0].shape, (2, 3))
-
-    @parametrize("backend", ["eager", "aot_eager"])
-    @parametrize("check_escaped_gradients", [True, False])
-    def test_leaf_function_escaped_gradient_mixed_inputs(
-        self, backend, check_escaped_gradients
-    ):
-        base1 = torch.randn(3, 3, requires_grad=True)
-        base2 = torch.randn(3, 4, requires_grad=True)
-        closure_weight1 = base1 * 2
-        closure_weight2 = base2 * 3
-
-        @leaf_function
-        def mixed_inputs(x, y):
-            out1 = x @ closure_weight1 + y
-            out2 = x @ closure_weight2
-            return (out1, out2)
-
-        @mixed_inputs.register_fake
-        def mixed_inputs_fake(x, y):
-            return (torch.empty(x.shape[0], 3), torch.empty(x.shape[0], 4))
-
-        def fn(x, y):
-            return mixed_inputs(x, y)
-
-        x = torch.randn(2, 3, requires_grad=True)
-        y = torch.randn(2, 3, requires_grad=False)
-
-        compiled_fn = torch.compile(fn, backend=backend, fullgraph=True)
-        with config.patch(
-            leaf_function_check_escaped_gradients=check_escaped_gradients
-        ):
-            if check_escaped_gradients:
-                with self.assertRaisesRegex(RuntimeError, "2 tensor"):
-                    compiled_fn(x, y)
-            else:
-                result = compiled_fn(x, y)
-                self.assertEqual(result[0].shape, (2, 3))
-                self.assertEqual(result[1].shape, (2, 4))
-
-    @parametrize("backend", ["eager", "aot_eager"])
-    def test_leaf_function_escaped_gradient_error_message_contains_tensor_info(
-        self, backend
-    ):
-        closure_weight = torch.randn(4, 5, dtype=torch.float32, requires_grad=True)
-
-        @leaf_function
-        def uses_closure(x):
-            return (x @ closure_weight,)
-
-        @uses_closure.register_fake
-        def uses_closure_fake(x):
-            return (torch.empty(x.shape[0], 5),)
-
-        def fn(x):
-            return uses_closure(x)
-
-        x = torch.randn(2, 4, requires_grad=True)
-
-        compiled_fn = torch.compile(fn, backend=backend, fullgraph=True)
-        with config.patch(leaf_function_check_escaped_gradients=True):
-            with self.assertRaisesRegex(RuntimeError, r"shape=\[4, 5\].*dtype="):
-                compiled_fn(x)
-
-    @parametrize("backend", ["eager", "aot_eager"])
-    def test_leaf_function_escaped_gradient_actually_lost(self, backend):
-        closure_weight = torch.randn(3, 3, requires_grad=True)
-
-        @leaf_function
-        def uses_closure(x):
-            return (x @ closure_weight,)
-
-        @uses_closure.register_fake
-        def uses_closure_fake(x):
-            return (torch.empty(x.shape[0], 3),)
-
-        def fn(x):
-            return uses_closure(x)
-
-        x = torch.randn(2, 3, requires_grad=True)
-
-        compiled_fn = torch.compile(fn, backend=backend, fullgraph=True)
-        result = compiled_fn(x)
-        loss = result[0].sum()
-        loss.backward()
-
-        self.assertIsNotNone(x.grad)
-        self.assertIsNone(closure_weight.grad)
-
-    def test_leaf_function_and_nonstrict_trace_mutually_exclusive(self):
-        from torch._dynamo.decorators import leaf_function, nonstrict_trace
-
-        with self.assertRaisesRegex(
-            ValueError,
-            "cannot be both marked as @leaf_function and @nonstrict_trace",
-        ):
-
-            @leaf_function
-            @nonstrict_trace
-            def bad_fn1(x):
-                return (x,)
-
-        with self.assertRaisesRegex(
-            ValueError,
-            "cannot be both marked as @leaf_function and @nonstrict_trace",
-        ):
-
-            @nonstrict_trace
-            @leaf_function
-            def bad_fn2(x):
-                return (x,)
-
-    def test_leaf_function_no_return_value(self):
-        printed = []
-
-        @leaf_function
-        def fn_no_return(x):
-            print("processing")
-
-        @fn_no_return.register_fake
-        def fn_no_return_fake(x):
-            pass
-
-        class Mod(torch.nn.Module):
-            def __init__(self):
-                super().__init__()
-                self.linear = torch.nn.Linear(3, 3)
-
-            def forward(self, x):
-                fn_no_return(x)
-                return (self.linear(x),)
-
-        def args_fn():
-            return (torch.randn(3, 3, requires_grad=True),)
-
-        def loss_fn(out):
-            return out[0].sum()
-
-        with patch("builtins.print", lambda *args, **kwargs: printed.append(args)):
-            eager_graph, fw_graph, bw_graph = self._test_leaf_function_helper(
-                Mod, args_fn, loss_fn
-            )
-        self.assertTrue(any("processing" in p for p in printed))
-        self.assertExpectedInline(
-            eager_graph,
-            """\
-class GraphModule(torch.nn.Module):
-    def forward(self, L_x_: "f32[3, 3]", L_self_modules_linear_parameters_weight_: "f32[3, 3]", L_self_modules_linear_parameters_bias_: "f32[3]"):
-        l_x_ = L_x_
-        l_self_modules_linear_parameters_weight_ = L_self_modules_linear_parameters_weight_
-        l_self_modules_linear_parameters_bias_ = L_self_modules_linear_parameters_bias_
-
-        real_fn : torch.utils._pytree.TreeSpec = self.real_fn
-        fake_fn : torch.utils._pytree.TreeSpec = self.fake_fn
-        invoke_leaf_function = torch.ops.higher_order.invoke_leaf_function(real_fn, fake_fn, l_x_);  real_fn = fake_fn = invoke_leaf_function = None
-
-        linear: "f32[3, 3]" = torch._C._nn.linear(l_x_, l_self_modules_linear_parameters_weight_, l_self_modules_linear_parameters_bias_);  l_x_ = l_self_modules_linear_parameters_weight_ = l_self_modules_linear_parameters_bias_ = None
-        return (linear,)
-""",  # noqa: B950
-        )
-        self.assertExpectedInline(
-            fw_graph,
-            """\
-class GraphModule(torch.nn.Module):
-    def forward(self, primals_1: "f32[0]", primals_2: "f32[3, 3]", primals_3: "f32[3, 3]", primals_4: "f32[3]"):
-        _tree_spec_constant0 = self._tree_spec_constant0
-        _tree_spec_constant1 = self._tree_spec_constant1
-        with_effects = torch.ops.higher_order.with_effects(primals_1, torch.ops.higher_order.invoke_leaf_function, _tree_spec_constant0, _tree_spec_constant1, primals_2);  primals_1 = _tree_spec_constant0 = _tree_spec_constant1 = None
-
-        getitem: "f32[0]" = with_effects[0];  with_effects = None
-
-        t: "f32[3, 3]" = torch.ops.aten.t.default(primals_3)
-        addmm: "f32[3, 3]" = torch.ops.aten.addmm.default(primals_4, primals_2, t);  primals_4 = t = None
-        return (getitem, addmm, primals_2, primals_3)
-""",  # noqa: B950
-        )
-        self.assertExpectedInline(
-            bw_graph,
-            """\
-class GraphModule(torch.nn.Module):
-    def forward(self, primals_2: "f32[3, 3]", primals_3: "f32[3, 3]", tangents_1: "f32[3, 3]"):
-        t: "f32[3, 3]" = torch.ops.aten.t.default(primals_3);  primals_3 = None
-        t_1: "f32[3, 3]" = torch.ops.aten.t.default(t);  t = None
-        mm: "f32[3, 3]" = torch.ops.aten.mm.default(tangents_1, t_1);  t_1 = None
-        t_2: "f32[3, 3]" = torch.ops.aten.t.default(tangents_1)
-        mm_1: "f32[3, 3]" = torch.ops.aten.mm.default(t_2, primals_2);  t_2 = primals_2 = None
-        t_3: "f32[3, 3]" = torch.ops.aten.t.default(mm_1);  mm_1 = None
-        sum_1: "f32[1, 3]" = torch.ops.aten.sum.dim_IntList(tangents_1, [0], True);  tangents_1 = None
-        view: "f32[3]" = torch.ops.aten.view.default(sum_1, [3]);  sum_1 = None
-        t_4: "f32[3, 3]" = torch.ops.aten.t.default(t_3);  t_3 = None
-        return (mm, t_4, view)
-""",  # noqa: B950
-        )
+            compiled(torch.randn(4))
 
 
 instantiate_parametrized_tests(DecoratorTests)
diff --git a/test/dynamo/test_deque_reconstruct.py b/test/dynamo/test_deque_reconstruct.py
index 5f07e1e72da9e..05f45260aaa4d 100644
--- a/test/dynamo/test_deque_reconstruct.py
+++ b/test/dynamo/test_deque_reconstruct.py
@@ -13,7 +13,8 @@ class TestDequeReconstruct(torch._inductor.test_case.TestCase):
     @contextlib.contextmanager
     def set_deque_in_globals(self, value):
         prev = globals().pop("deque", self.UNSET)
-        assert "deque" not in globals()
+        if "deque" in globals():
+            raise AssertionError("Expected deque to not be in globals")
 
         try:
             if value is not self.UNSET:
@@ -22,7 +23,8 @@ def set_deque_in_globals(self, value):
         finally:
             if prev is self.UNSET:
                 globals().pop("deque", None)
-                assert "deque" not in globals()
+                if "deque" in globals():
+                    raise AssertionError("Expected deque to not be in globals")
             else:
                 globals()["deque"] = prev
 
diff --git a/test/dynamo/test_deviceguard.py b/test/dynamo/test_deviceguard.py
index 2d1d267c9379e..de2c9f1b76bf7 100644
--- a/test/dynamo/test_deviceguard.py
+++ b/test/dynamo/test_deviceguard.py
@@ -6,7 +6,7 @@
 import torch._dynamo.test_case
 import torch._dynamo.testing
 from torch._dynamo.device_interface import CudaInterface, DeviceGuard
-from torch.testing._internal.common_cuda import TEST_CUDA, TEST_MULTIGPU
+from torch.testing._internal.common_cuda import TEST_CUDA
 
 
 class TestDeviceGuard(torch._dynamo.test_case.TestCase):
@@ -56,21 +56,6 @@ def setUp(self):
         super().setUp()
         self.device_interface = CudaInterface
 
-    @unittest.skipIf(not TEST_MULTIGPU, "need multiple GPU")
-    def test_device_guard(self):
-        current_device = torch.cuda.current_device()
-
-        device_guard = DeviceGuard(self.device_interface, 1)
-
-        with device_guard as _:
-            self.assertEqual(torch.cuda.current_device(), 1)
-            self.assertEqual(device_guard.prev_idx, 0)
-            self.assertEqual(device_guard.idx, 1)
-
-        self.assertEqual(torch.cuda.current_device(), current_device)
-        self.assertEqual(device_guard.prev_idx, 0)
-        self.assertEqual(device_guard.idx, 1)
-
     def test_device_guard_no_index(self):
         current_device = torch.cuda.current_device()
 
diff --git a/test/dynamo/test_dicts.py b/test/dynamo/test_dicts.py
index fa1b1dfeebd5b..b88c447dbd96f 100644
--- a/test/dynamo/test_dicts.py
+++ b/test/dynamo/test_dicts.py
@@ -134,6 +134,46 @@ def fn(x):
         with unittest.mock.patch("torch._dynamo.config.error_on_recompile", True):
             self.assertEqual(fn(x), opt_fn(x))
 
+    def test_dict_torch_size_dynamic_key(self):
+        class DynamicShapeModel(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.layer_configs = {
+                    torch.Size([32, 64]): torch.nn.Linear(32, 64),
+                    torch.Size([64, 128]): torch.nn.Linear(64, 128),
+                    torch.Size([128, 256]): torch.nn.Linear(128, 256),
+                }
+                self.activation_functions = {
+                    torch.Size([32]): torch.nn.ReLU(),
+                    torch.Size([64]): torch.nn.Tanh(),
+                    torch.Size([128]): torch.nn.Sigmoid(),
+                }
+
+            def forward(self, x):
+                current_shape = torch.tensor(x.shape[1:])
+                shape_key = torch.Size([current_shape[0], 64])
+                if shape_key in self.layer_configs:
+                    x = self.layer_configs[shape_key](x)
+
+                activation_shape = torch.Size([x.shape[1]])
+                if activation_shape in self.activation_functions:
+                    x = self.activation_functions[activation_shape](x)
+                return x
+
+        model = DynamicShapeModel().eval()
+
+        for features in (32, 64, 128, 16):
+            x = torch.randn(4, features)
+            with torch.no_grad():
+                eager_out = model(x)
+
+            torch._dynamo.reset()
+            compiled_model = torch.compile(model, backend="eager")
+            with torch.no_grad():
+                compiled_out = compiled_model(x)
+
+            self.assertTrue(same(eager_out, compiled_out))
+
     def test_dict_subclass_methods_fallback_readonly(self):
         sd = SimpleDict()
         sd[2] = 5
@@ -445,7 +485,8 @@ def fn(d):
 
         args1 = {"a": torch.randn(10), "b": torch.randn(10)}
         args2 = dict(args1)
-        assert fn(args1) is args1
+        if fn(args1) is not args1:
+            raise AssertionError("Expected fn(args1) to be args1")
         cnts = torch._dynamo.testing.CompileCounter()
         opt_fn = torch.compile(fn, backend=cnts)
         self.assertIs(opt_fn(args2), args2)
@@ -608,6 +649,33 @@ def fn(x):
         res = opt_fn(x)
         self.assertEqual(ref, res)
 
+    def test_len_dunder_dict_after_delete(self):
+        # After deleting an attribute during tracing, len(obj.__dict__) must
+        # reflect the deletion. The bug: SideEffectsProxyDict.__len__ returned
+        # len(item_dict) + len(side_effects_table), double-counting deleted
+        # entries and keys mutated from item_dict.
+        class UserDefined:
+            def __init__(self) -> None:
+                self.a = 3
+                self.b = 5
+
+            def run(self, x):
+                del self.a
+                # len should be 1 (only b remains), not 3 (item_dict=2 + mutations=1)
+                return x * len(self.__dict__)
+
+        obj1 = UserDefined()
+        obj2 = UserDefined()
+
+        def fn(x, obj):
+            return obj.run(x)
+
+        x = torch.ones(2)
+        ref = fn(x, obj1)
+        opt_fn = torch.compile(fn, backend="eager", fullgraph=True)
+        res = opt_fn(x, obj2)
+        self.assertEqual(ref, res)
+
     def test_contains_module_dunder_dict(self):
         class MyModule(torch.nn.Module):
             def __init__(self) -> None:
@@ -706,7 +774,7 @@ def __init__(self, *args, **kwargs):
             def fn(x):
                 c = CustomDict()
                 c["key"] = x
-                assert "key" in c
+                assert "key" in c  # noqa: S101
                 return c["key"] + 1
 
             opt_fn = torch.compile(fn, backend="eager", fullgraph=True)
@@ -1203,6 +1271,222 @@ def f():
         opt_f = torch.compile(f, backend="eager", fullgraph=True)
         self.assertEqual(f(), opt_f())
 
+    @parametrize("op", ["or_", "and_", "xor", "sub"])
+    def test_dict_items_binop(self, op):
+        op = getattr(operator, op)
+
+        def f():
+            a = {"one": 1, "two": 2}
+            b = {"one": 1, "three": 3}
+            return op(a.items(), b.items()), op(b.items(), a.items())
+
+        opt_f = torch.compile(f, backend="eager", fullgraph=True)
+        self.assertEqual(f(), opt_f())
+
+    @parametrize("op", ["ior", "iand", "ixor", "isub"])
+    def test_dict_items_inplace_binop(self, op):
+        op = getattr(operator, op)
+
+        def f():
+            a = {"one": 1, "two": 2}.items()
+            b = {"one": 1, "three": 3}.items()
+            c = {"one": 1, "two": 2}.items()
+            a = op(a, b)
+            b = op(b, c)
+            return a, b
+
+        opt_f = torch.compile(f, backend="eager", fullgraph=True)
+        self.assertEqual(f(), opt_f())
+
+    @parametrize("op", ["or_", "and_", "xor", "sub", "iand", "ior", "ixor", "isub"])
+    def test_cross_type_set_binop_dict_keys_vs_set(self, op):
+        op = getattr(operator, op)
+
+        def f():
+            keys = {"one": 1, "two": 2, "three": 3}.keys()
+            s = {"one", "four"}
+            return op(keys, s), op(s, keys)
+
+        opt_f = torch.compile(f, backend="eager", fullgraph=True)
+        self.assertEqual(f(), opt_f())
+
+    @parametrize("op", ["or_", "and_", "xor", "sub", "iand", "ior", "ixor", "isub"])
+    def test_cross_type_set_binop_dict_items_vs_set(self, op):
+        op = getattr(operator, op)
+
+        def f():
+            items = {"one": 1, "two": 2}.items()
+            s = {("one", 1), ("three", 3)}
+            return op(items, s), op(s, items)
+
+        opt_f = torch.compile(f, backend="eager", fullgraph=True)
+        self.assertEqual(f(), opt_f())
+
+    @parametrize("op", ["or_", "and_", "xor", "sub", "iand", "ior", "ixor", "isub"])
+    def test_cross_type_set_binop_dict_keys_vs_dict_items(self, op):
+        op = getattr(operator, op)
+
+        def f():
+            keys = {"one": 1, "two": 2}.keys()
+            items = {"three": 3, "four": 4}.items()
+            return op(keys, items), op(items, keys)
+
+        opt_f = torch.compile(f, backend="eager", fullgraph=True)
+        self.assertEqual(f(), opt_f())
+
+    @parametrize("op", ["or_", "and_", "xor", "sub", "iand", "ior", "ixor", "isub"])
+    def test_cross_type_set_binop_dict_keys_vs_user_defined_set(self, op):
+        class MySet(set):
+            pass
+
+        op = getattr(operator, op)
+
+        def f():
+            keys = {"one": 1, "two": 2, "three": 3}.keys()
+            s = MySet({"one", "four"})
+            return op(keys, s), op(s, keys)
+
+        opt_f = torch.compile(f, backend="eager", fullgraph=True)
+        self.assertEqual(f(), opt_f())
+
+    @parametrize("op", ["or_", "and_", "xor", "sub", "iand", "ior", "ixor", "isub"])
+    def test_cross_type_set_binop_dict_items_vs_user_defined_set(self, op):
+        class MySet(set):
+            pass
+
+        op = getattr(operator, op)
+
+        def f():
+            items = {"one": 1, "two": 2, "three": 3}.items()
+            s = MySet({"one", "four"})
+            return op(items, s), op(s, items)
+
+        opt_f = torch.compile(f, backend="eager", fullgraph=True)
+        self.assertEqual(f(), opt_f())
+
+    @parametrize("op", ["or_", "and_", "xor", "sub", "iand", "ior", "ixor", "isub"])
+    def test_cross_type_set_binop_set_vs_user_defined_set(self, op):
+        class MySet(set):
+            pass
+
+        op = getattr(operator, op)
+
+        def f():
+            s = {"one", "two", "three"}
+            u = MySet({"one", "four"})
+            return op(s, u), op(u, s)
+
+        opt_f = torch.compile(f, backend="eager", fullgraph=True)
+        self.assertEqual(f(), opt_f())
+
+    @parametrize("op", ["lt", "le", "gt", "ge", "eq", "ne"])
+    def test_cross_type_cmp_dict_keys_vs_set(self, op):
+        op = getattr(operator, op)
+
+        def f():
+            keys = {"one": 1, "two": 2}.keys()
+            s = {"one", "two", "three"}
+            return op(keys, s), op(s, keys)
+
+        opt_f = torch.compile(f, backend="eager", fullgraph=True)
+        self.assertEqual(f(), opt_f())
+
+    @parametrize("op", ["lt", "le", "gt", "ge", "eq", "ne"])
+    def test_cross_type_cmp_dict_items_vs_set(self, op):
+        op = getattr(operator, op)
+
+        def f():
+            items = {"one": 1, "two": 2}.items()
+            s = {("one", 1), ("two", 2), ("three", 3)}
+            return op(items, s), op(s, items)
+
+        opt_f = torch.compile(f, backend="eager", fullgraph=True)
+        self.assertEqual(f(), opt_f())
+
+    @parametrize("op", ["lt", "le", "gt", "ge", "eq", "ne"])
+    def test_cross_type_cmp_dict_keys_vs_dict_items(self, op):
+        op = getattr(operator, op)
+
+        def f():
+            keys = {"one": 1, "two": 2}.keys()
+            items = {"one": 1, "two": 2}.items()
+            return op(keys, items), op(items, keys)
+
+        opt_f = torch.compile(f, backend="eager", fullgraph=True)
+        self.assertEqual(f(), opt_f())
+
+    @parametrize("op", ["lt", "le", "gt", "ge", "eq", "ne"])
+    def test_cross_type_cmp_set_vs_user_defined_set(self, op):
+        class MySet(set):
+            pass
+
+        op = getattr(operator, op)
+
+        def f():
+            s = {"one", "two"}
+            u = MySet({"one", "two", "three"})
+            return op(s, u), op(u, s)
+
+        opt_f = torch.compile(f, backend="eager", fullgraph=True)
+        self.assertEqual(f(), opt_f())
+
+    @parametrize("op", ["lt", "le", "gt", "ge", "eq", "ne"])
+    def test_cross_type_cmp_dict_keys_vs_user_defined_set(self, op):
+        class MySet(set):
+            pass
+
+        op = getattr(operator, op)
+
+        def f():
+            keys = {"one": 1, "two": 2}.keys()
+            u = MySet({"one", "two", "three"})
+            return op(keys, u), op(u, keys)
+
+        opt_f = torch.compile(f, backend="eager", fullgraph=True)
+        self.assertEqual(f(), opt_f())
+
+    @parametrize("op", ["lt", "le", "gt", "ge", "eq", "ne"])
+    def test_cross_type_cmp_dict_items_vs_user_defined_set(self, op):
+        class MySet(set):
+            pass
+
+        op = getattr(operator, op)
+
+        def f():
+            items = {"one": 1, "two": 2}.items()
+            u = MySet({("one", 1), ("two", 2), ("three", 3)})
+            return op(items, u), op(u, items)
+
+        opt_f = torch.compile(f, backend="eager", fullgraph=True)
+        self.assertEqual(f(), opt_f())
+
+    def test_dict_view_iand_rebinds_variable(self):
+        def f_keys():
+            d = {"one": 1, "two": 2, "three": 3}
+            tmp = d.keys()
+            tmp &= {"one", "four"}
+            return tmp, d
+
+        opt_f_keys = torch.compile(f_keys, backend="eager", fullgraph=True)
+        eager_tmp, eager_d = f_keys()
+        compiled_tmp, compiled_d = opt_f_keys()
+        self.assertIsInstance(eager_tmp, set)
+        self.assertEqual(eager_tmp, compiled_tmp)
+        self.assertEqual(eager_d, compiled_d)
+
+        def f_items():
+            d = {"one": 1, "two": 2}
+            tmp = d.items()
+            tmp &= {("one", 1), ("three", 3)}
+            return tmp, d
+
+        opt_f_items = torch.compile(f_items, backend="eager", fullgraph=True)
+        eager_tmp, eager_d = f_items()
+        compiled_tmp, compiled_d = opt_f_items()
+        self.assertIsInstance(eager_tmp, set)
+        self.assertEqual(eager_tmp, compiled_tmp)
+        self.assertEqual(eager_d, compiled_d)
+
     def test_dict_union_result_type(self):
         def_dict = defaultdict(int, {1: 1, 2: 2})
         ord_dict = OrderedDict({3: 3, 4: 4})
@@ -1285,8 +1569,8 @@ def test_typing_union_as_dict_key(self):
         from typing import Union
 
         def fn(x):
-            d = {Union[int, str]: x * 2, Union[float, bool]: x * 3}
-            return d[Union[int, str]] + d[Union[float, bool]]
+            d = {Union[int, str]: x * 2, Union[float, bool]: x * 3}  # noqa: UP007
+            return d[Union[int, str]] + d[Union[float, bool]]  # noqa: UP007
 
         x = torch.randn(4)
         opt_fn = torch.compile(fn, backend="eager", fullgraph=True)
@@ -1361,6 +1645,117 @@ def fn(x):
         with self.assertRaisesRegex(Unsupported, "Observed exception"):
             opt_fn(x)
 
+    def test_property_backed_by_dunder_dict(self):
+        class Obj:
+            def __init__(self, x):
+                self.__dict__["x"] = x
+
+            @property
+            def x(self):
+                return self.__dict__["x"]
+
+        def fn(obj):
+            return obj.x
+
+        obj = Obj(3)
+        opt_fn = torch.compile(fn, backend="eager", fullgraph=True)
+        self.assertEqual(fn(obj), opt_fn(obj))
+
+    def test_property_backed_by_dunder_dict_mutated(self):
+        class Obj:
+            def __init__(self, x):
+                self.__dict__["x"] = x
+
+            @property
+            def x(self):
+                return self.__dict__["x"]
+
+            @x.setter
+            def x(self, value):
+                self.__dict__["x"] = value
+
+        def fn(obj):
+            a = obj.x
+            obj.x = a + 1
+            return obj.x
+
+        obj = Obj(3)
+        opt_fn = torch.compile(fn, backend="eager", fullgraph=True)
+        self.assertEqual(fn(Obj(3)), opt_fn(obj))
+
+    def test_property_backed_by_dunder_dict_pending_side_effects(self):
+        class Obj:
+            def __init__(self, x):
+                self.__dict__["x"] = x
+
+            @property
+            def x(self):
+                return self.__dict__["x"]
+
+        def fn(obj):
+            a = obj.__dict__["x"]
+            obj.__dict__["x"] = a + 1
+            return obj.__dict__["x"]
+
+        obj = Obj(3)
+        opt_fn = torch.compile(fn, backend="eager", fullgraph=True)
+        self.assertEqual(fn(Obj(3)), opt_fn(obj))
+
+    def test_dunder_dict_getitem_guard_recompiles(self):
+        class Obj:
+            def __init__(self):
+                self.__dict__["x"] = 3
+
+            @property
+            def x(self):
+                return 999
+
+        def fn(obj, t):
+            return t + obj.__dict__["x"]
+
+        obj = Obj()
+        t = torch.tensor(0.0)
+        cnt = torch._dynamo.testing.CompileCounter()
+        opt_fn = torch.compile(fn, backend=cnt, fullgraph=True)
+        self.assertEqual(opt_fn(obj, t), torch.tensor(3.0))
+        self.assertEqual(cnt.frame_count, 1)
+
+        obj.__dict__["x"] = 5
+        self.assertEqual(opt_fn(obj, t), torch.tensor(5.0))
+        self.assertEqual(cnt.frame_count, 2)
+
+    def test_dunder_dict_get_with_property(self):
+        class Obj:
+            def __init__(self, x):
+                self.__dict__["x"] = x
+
+            @property
+            def x(self):
+                return self.__dict__["x"]
+
+        def fn(obj):
+            return obj.__dict__.get("x")
+
+        obj = Obj(3)
+        opt_fn = torch.compile(fn, backend="eager", fullgraph=True)
+        self.assertEqual(fn(obj), opt_fn(obj))
+
+    def test_dunder_dict_getitem_tensor(self):
+        class Obj:
+            def __init__(self, x):
+                self.__dict__["x"] = x
+
+            @property
+            def x(self):
+                return self.__dict__["x"]
+
+        def fn(obj):
+            return obj.x + 1
+
+        obj = Obj(torch.tensor(3.0))
+        opt_fn = torch.compile(fn, backend="eager", fullgraph=True)
+        self.assertEqual(fn(obj), opt_fn(obj))
+
     def test_get_default_nowrap_functions_as_dict_key(self):
         def fn(x):
             # Get the set of default nowrap functions
@@ -1376,6 +1771,44 @@ def fn(x):
         opt_fn = torch.compile(fn, backend="eager", fullgraph=True)
         self.assertEqual(fn(x), opt_fn(x))
 
+    def test_custom_bool(self):
+        class CustomBoolDict:
+            def __init__(self, bool_result: bool):
+                super().__init__()
+                self._bool_result = bool_result
+
+            def __bool__(self):
+                return self._bool_result
+
+        def fn(t: torch.Tensor, d, apply_not: bool):
+            if apply_not:
+                return t.sin(), bool(not d)
+            else:
+                return t.sin(), bool(d)
+
+        x = torch.randn(4)
+        opt_fn = torch.compile(fn, backend="eager", fullgraph=True)
+
+        t = CustomBoolDict(True)
+        self.assertTrue(
+            opt_fn(x, t, False)[1],
+            "CustomBoolDict(True) should evaluate to True in boolean context",
+        )
+        self.assertFalse(
+            opt_fn(x, t, True)[1],
+            "not CustomBoolDict(True) should evaluate to False in boolean context",
+        )
+
+        f = CustomBoolDict(False)
+        self.assertFalse(
+            opt_fn(x, f, False)[1],
+            "CustomBoolDict(False) should evaluate to False in boolean context",
+        )
+        self.assertTrue(
+            opt_fn(x, f, True)[1],
+            "not CustomBoolDict(False) should evaluate to True in boolean context",
+        )
+
 
 instantiate_parametrized_tests(DictTests)
 
@@ -1520,6 +1953,7 @@ class DictMethodsTests(torch._dynamo.test_case.TestCase):
     thetype = dict
 
     # Methods:
+    # + bool
     # + clear
     # + copy
     # + fromkeys
@@ -1968,10 +2402,27 @@ def f(a, x):
         x = torch.randn(4)
         self.assertTrue(same(f(A(), x), opt_f(A(), x)))
 
+    @make_dynamo_test
+    def test_bool(self):
+        p = self.thetype()
+        q = self.thetype({"a": 1, "b": 2})
+        if p:
+            self.fail("empty mapping must compare to False")
+        if not q:
+            self.fail("non-empty mapping must compare to True")
+        if bool(p):
+            self.fail("empty mapping must compare to False")
+        if not bool(q):
+            self.fail("non-empty mapping must compare to True")
+
 
 class DictSubclassMethodsTests(DictMethodsTests):
     thetype = SimpleDict
 
+    @unittest.expectedFailure
+    def test_binop_or(self):
+        super().test_binop_or()
+
 
 class OrderedDictMethodsTests(DictMethodsTests):
     thetype = OrderedDict
@@ -2075,6 +2526,151 @@ def test_move_to_end(self):
         self.assertEqual(list(p.keys()), list("bc"))
 
 
+class DunderDictVariableTests(torch._dynamo.test_case.TestCase):
+    """Tests for DunderDictVariable (object.__dict__ handling in Dynamo)"""
+
+    def test_dunder_dict_items_includes_mutations(self):
+        """Test that __dict__.items() includes both original and mutated keys"""
+
+        class MyClass:
+            def __init__(self):
+                self.a = 1
+                self.b = 2
+
+        @torch.compile(backend="eager", fullgraph=True)
+        def fn(obj):
+            obj.c = 3
+            items = dict(obj.__dict__.items())
+            return items
+
+        obj = MyClass()
+        result = fn(obj)
+        self.assertEqual(result, {"a": 1, "b": 2, "c": 3})
+
+    def test_dunder_dict_keys_includes_mutations(self):
+        """Test that __dict__.keys() includes both original and mutated keys"""
+
+        class MyClass:
+            def __init__(self):
+                self.a = 1
+                self.b = 2
+
+        @torch.compile(backend="eager", fullgraph=True)
+        def fn(obj):
+            obj.c = 3
+            obj.d = 4
+            keys = sorted(obj.__dict__.keys())
+            return keys
+
+        obj = MyClass()
+        result = fn(obj)
+        self.assertEqual(result, ["a", "b", "c", "d"])
+
+    def test_dunder_dict_values_includes_mutations(self):
+        """Test that __dict__.values() includes both original and mutated values"""
+
+        class MyClass:
+            def __init__(self):
+                self.a = 1
+                self.b = 2
+
+        @torch.compile(backend="eager", fullgraph=True)
+        def fn(obj):
+            obj.c = 3
+            values = sorted(obj.__dict__.values())
+            return values
+
+        obj = MyClass()
+        result = fn(obj)
+        self.assertEqual(result, [1, 2, 3])
+
+    def test_dunder_dict_comprehension_with_mutations(self):
+        """Test dict comprehension over __dict__ with mutations (scheduler use case)"""
+
+        class SimpleScheduler:
+            def __init__(self):
+                self.base_lrs = [0.1]
+                self.last_epoch = 0
+
+            def state_dict(self):
+                state = dict(self.__dict__.items())
+                return state
+
+        @torch.compile(backend="eager", fullgraph=True)
+        def fn(scheduler):
+            scheduler.last_epoch = 5
+            scheduler.new_attr = "test"
+            return scheduler.state_dict()
+
+        scheduler = SimpleScheduler()
+        result = fn(scheduler)
+        self.assertEqual(result["base_lrs"], [0.1])
+        self.assertEqual(result["last_epoch"], 5)
+        self.assertEqual(result["new_attr"], "test")
+
+    def test_dunder_dict_multiple_mutations(self):
+        """Test __dict__ operations with multiple mutations"""
+
+        class MyClass:
+            def __init__(self):
+                self.x = 1
+
+        @torch.compile(backend="eager", fullgraph=True)
+        def fn(obj):
+            obj.y = 2
+            obj.z = 3
+            obj.y = 20
+            items = dict(obj.__dict__.items())
+            keys = set(obj.__dict__.keys())
+            return items, keys
+
+        obj = MyClass()
+        items, keys = fn(obj)
+        self.assertEqual(items, {"x": 1, "y": 20, "z": 3})
+        self.assertEqual(keys, {"x", "y", "z"})
+
+    def test_dunder_dict_items_iteration(self):
+        """Test iterating over __dict__.items() with mutations"""
+
+        class MyClass:
+            def __init__(self):
+                self.a = 1
+                self.b = 2
+
+        @torch.compile(backend="eager", fullgraph=True)
+        def fn(obj):
+            obj.c = 3
+            result = []
+            for k, v in obj.__dict__.items():
+                result.append((k, v))
+            return sorted(result)
+
+        obj = MyClass()
+        result = fn(obj)
+        self.assertEqual(result, [("a", 1), ("b", 2), ("c", 3)])
+
+    def test_method_dict_direct_fullgraph(self):
+        """
+        Regression test: Accessing __dict__ directly on UserMethodVariable.
+        This should fail with: unsupported variable type for __dict__ access
+        """
+
+        class Foo:
+            def bar(self):
+                return 42
+
+        @torch.compile(backend="eager", fullgraph=True)
+        def fn():
+            obj = Foo()
+            # Access __dict__ on bound method - creates UserMethodVariable
+            method = obj.bar
+            d = method.__dict__
+            return d
+
+        # This should not raise Unsupported
+        fn()
+
+
 if __name__ == "__main__":
     from torch._dynamo.test_case import run_tests
 
diff --git a/test/dynamo/test_dynamic_shapes.py b/test/dynamo/test_dynamic_shapes.py
index 3bebcfc345b70..a585c65c82157 100644
--- a/test/dynamo/test_dynamic_shapes.py
+++ b/test/dynamo/test_dynamic_shapes.py
@@ -1,5 +1,4 @@
 # Owner(s): ["module: dynamo"]
-import unittest
 import warnings
 
 from torch._dynamo import config
@@ -85,13 +84,6 @@ def make_dynamic_cls(cls):
     make_dynamic_cls(test)
 del test
 
-if TEST_Z3:
-    if not config.inline_inbuilt_nn_modules:
-        # TODO model is somehow not being freed when z3 is available
-        unittest.expectedFailure(
-            DynamicShapesMiscTests.test_parameter_free_dynamic_shapes  # noqa: F821
-        )
-
 # Test takes too long ~700s as of 414a1fd29f04d06e41b7f895368dd1f83a4be29d
 DynamicShapesExportTests.test_retracibility_dynamic_shapes = slowTest(  # noqa: F821
     DynamicShapesExportTests.test_retracibility_dynamic_shapes  # noqa: F821
diff --git a/test/dynamo/test_dynamo_decompositions.py b/test/dynamo/test_dynamo_decompositions.py
index 1140b9b73d33f..16f40c7a3bcbf 100644
--- a/test/dynamo/test_dynamo_decompositions.py
+++ b/test/dynamo/test_dynamo_decompositions.py
@@ -1,5 +1,7 @@
 # Owner(s): ["module: dynamo"]
 
+import unittest
+
 import torch
 import torch._dynamo.config
 import torch._dynamo.test_case
@@ -17,7 +19,7 @@ class TestDynamoDecompositions(torch._dynamo.test_case.TestCase):
 
     @skipIfCrossRef
     def test_addcmul_inplace_decomposition_enabled(self):
-        """With decompositions enabled, addcmul_ should decompose into mul and add_."""
+        """With decompositions enabled, addcmul_ should decompose into mul, fma, and copy_."""
 
         def fn(x, tensor1, tensor2, value):
             return x.addcmul_(tensor1, tensor2, value=value)
@@ -43,10 +45,10 @@ def forward(self, L_x_: "f32[4]", L_tensor1_: "f32[4]", L_tensor2_: "f32[4]", L_
         l_tensor2_ = L_tensor2_
         l_value_ = L_value_
 
-        mul: "f32[4]" = l_tensor1_ * l_tensor2_;  l_tensor1_ = l_tensor2_ = None
-        mul_1: "f32[4]" = mul * l_value_;  mul = l_value_ = None
-        add_: "f32[4]" = l_x_.add_(mul_1);  l_x_ = mul_1 = None
-        return (add_,)
+        mul: "f32[4]" = torch.mul(l_tensor1_, l_tensor2_);  l_tensor1_ = l_tensor2_ = None
+        fma_default: "f32[4]" = torch.ops.prims.fma.default(mul, l_value_, l_x_);  mul = l_value_ = None
+        copy_: "f32[4]" = l_x_.copy_(fma_default);  l_x_ = fma_default = None
+        return (copy_,)
 """,
         )
 
@@ -122,7 +124,7 @@ def forward(self, L_x_: "f32[4]", L_tensor1_: "f32[4]", L_tensor2_: "f32[4]"):
 
     @skipIfCrossRef
     def test_add_inplace_with_alpha_decomposition_enabled(self):
-        """With decompositions enabled, add_ with alpha should decompose into mul and add_."""
+        """With decompositions enabled, add_ with tensor alpha should decompose into fma and copy_."""
 
         def fn(x, other, alpha):
             return x.add_(other, alpha=alpha)
@@ -146,9 +148,9 @@ def forward(self, L_x_: "f32[4]", L_other_: "f32[4]", L_alpha_: "f32[]"):
         l_other_ = L_other_
         l_alpha_ = L_alpha_
 
-        mul: "f32[4]" = torch.mul(l_other_, l_alpha_);  l_other_ = l_alpha_ = None
-        add_: "f32[4]" = l_x_.add_(mul);  l_x_ = mul = None
-        return (add_,)
+        fma_default: "f32[4]" = torch.ops.prims.fma.default(l_other_, l_alpha_, l_x_);  l_other_ = l_alpha_ = None
+        copy_: "f32[4]" = l_x_.copy_(fma_default);  l_x_ = fma_default = None
+        return (copy_,)
 """,
         )
 
@@ -220,7 +222,10 @@ def forward(self, L_x_: "f32[4]", L_other_: "f32[4]"):
 
     @skipIfCrossRef
     def test_addcdiv_inplace_decomposition_enabled(self):
-        """With decompositions enabled, addcdiv_ should decompose into div, mul, and add_."""
+        """With decompositions enabled, addcdiv_ should decompose into div, fma, and copy_.
+
+        ATen computes self + value * (t1/t2), nvcc fuses to fma(value, quotient, self).
+        """
 
         def fn(x, tensor1, tensor2, value):
             return x.addcdiv_(tensor1, tensor2, value=value)
@@ -247,9 +252,9 @@ def forward(self, L_x_: "f32[4]", L_tensor1_: "f32[4]", L_tensor2_: "f32[4]", L_
         l_value_ = L_value_
 
         div: "f32[4]" = torch.div(l_tensor1_, l_tensor2_);  l_tensor1_ = l_tensor2_ = None
-        mul: "f32[4]" = torch.mul(div, l_value_);  div = l_value_ = None
-        add_: "f32[4]" = l_x_.add_(mul);  l_x_ = mul = None
-        return (add_,)
+        fma_default: "f32[4]" = torch.ops.prims.fma.default(div, l_value_, l_x_);  div = l_value_ = None
+        copy_: "f32[4]" = l_x_.copy_(fma_default);  l_x_ = fma_default = None
+        return (copy_,)
 """,
         )
 
@@ -342,30 +347,39 @@ def fn(tensors, end_tensors, weight):
 
         graph = eager.graphs[0]
         actual = normalize_gm(graph.print_readable(False))
+        # On Python 3.10, dynamo names the >= node after the op ("ge")
+        # rather than the variable name ("mask").  Normalize so the
+        # expected string works on every Python version.
+        import re
+
+        actual = re.sub(r"\bge\b", "mask", actual)
 
         self.assertExpectedInline(
             actual,
             """\
 class GraphModule(torch.nn.Module):
-    def forward(self, L_weight_: "f32[]", L_end_tensors_0_: "f32[4]", L_end_tensors_1_: "f32[4]", \
-L_tensors_0_: "f32[4]", L_tensors_1_: "f32[4]"):
+    def forward(self, L_weight_: "f32[]", L_end_tensors_0_: "f32[4]", L_end_tensors_1_: "f32[4]", L_tensors_0_: "f32[4]", L_tensors_1_: "f32[4]"):
         l_weight_ = L_weight_
         l_end_tensors_0_ = L_end_tensors_0_
         l_end_tensors_1_ = L_end_tensors_1_
         l_tensors_0_ = L_tensors_0_
         l_tensors_1_ = L_tensors_1_
 
-        _foreach_sub = torch._foreach_sub([l_end_tensors_0_, l_end_tensors_1_], [l_tensors_0_, \
-l_tensors_1_]);  l_end_tensors_0_ = l_end_tensors_1_ = None
+        _foreach_sub = torch._foreach_sub([l_end_tensors_0_, l_end_tensors_1_], [l_tensors_0_, l_tensors_1_])
         getitem: "f32[4]" = _foreach_sub[0]
         getitem_1: "f32[4]" = _foreach_sub[1];  _foreach_sub = None
-        _foreach_mul = torch._foreach_mul((getitem, getitem_1), l_weight_);  getitem = getitem_1 = l_weight_ = None
-        getitem_2: "f32[4]" = _foreach_mul[0]
-        getitem_3: "f32[4]" = _foreach_mul[1];  _foreach_mul = None
-        _foreach_add_ = torch._foreach_add_([l_tensors_0_, l_tensors_1_], (getitem_2, getitem_3));  \
-l_tensors_0_ = l_tensors_1_ = getitem_2 = getitem_3 = _foreach_add_ = None
+        abs_1: "f32[]" = l_weight_.abs()
+        mask: "b8[]" = abs_1 >= 0.5;  abs_1 = None
+        sub: "f32[]" = 1.0 - l_weight_
+        neg_omw: "f32[]" = -sub;  sub = None
+        w: "f32[]" = torch.where(mask, neg_omw, l_weight_);  neg_omw = l_weight_ = None
+        b: "f32[4]" = torch.where(mask, l_end_tensors_0_, l_tensors_0_);  l_end_tensors_0_ = None
+        b_1: "f32[4]" = torch.where(mask, l_end_tensors_1_, l_tensors_1_);  mask = l_end_tensors_1_ = None
+        _foreach_addcmul_ = torch._foreach_addcmul_([b, b_1], [w, w], (getitem, getitem_1));  w = getitem = getitem_1 = _foreach_addcmul_ = None
+        copy_: "f32[4]" = l_tensors_0_.copy_(b);  l_tensors_0_ = b = copy_ = None
+        copy__1: "f32[4]" = l_tensors_1_.copy_(b_1);  l_tensors_1_ = b_1 = copy__1 = None
         return ()
-""",
+""",  # noqa: B950
         )
 
     @skipIfCrossRef
@@ -393,19 +407,26 @@ def fn(tensors, end_tensors):
             actual,
             """\
 class GraphModule(torch.nn.Module):
-    def forward(self, L_tensors_0_: "f32[4]", L_tensors_1_: "f32[4]", L_end_tensors_0_: "f32[4]", \
-L_end_tensors_1_: "f32[4]"):
-        l_tensors_0_ = L_tensors_0_
-        l_tensors_1_ = L_tensors_1_
+    def forward(self, L_end_tensors_0_: "f32[4]", L_end_tensors_1_: "f32[4]", L_tensors_0_: "f32[4]", L_tensors_1_: "f32[4]"):
         l_end_tensors_0_ = L_end_tensors_0_
         l_end_tensors_1_ = L_end_tensors_1_
+        l_tensors_0_ = L_tensors_0_
+        l_tensors_1_ = L_tensors_1_
 
-        _foreach_lerp_ = torch._foreach_lerp_([l_tensors_0_, l_tensors_1_], [l_end_tensors_0_, \
-l_end_tensors_1_], 0.5);  l_tensors_0_ = l_tensors_1_ = l_end_tensors_0_ = l_end_tensors_1_ = _foreach_lerp_ = None
+        _foreach_sub = torch._foreach_sub([l_end_tensors_0_, l_end_tensors_1_], [l_tensors_0_, l_tensors_1_])
+        getitem: "f32[4]" = _foreach_sub[0]
+        getitem_1: "f32[4]" = _foreach_sub[1];  _foreach_sub = None
+        tensor: "f32[]" = torch.tensor(0.5, dtype = torch.float32, device = device(type='cpu'))
+        sub: "f32[]" = 1.0 - tensor;  tensor = None
+        neg_omw: "f32[]" = -sub;  sub = None
+        copy_: "f32[4]" = l_tensors_0_.copy_(l_end_tensors_0_);  l_end_tensors_0_ = copy_ = None
+        copy__1: "f32[4]" = l_tensors_1_.copy_(l_end_tensors_1_);  l_end_tensors_1_ = copy__1 = None
+        _foreach_addcmul_ = torch._foreach_addcmul_([l_tensors_0_, l_tensors_1_], [neg_omw, neg_omw], (getitem, getitem_1));  l_tensors_0_ = l_tensors_1_ = neg_omw = getitem = getitem_1 = _foreach_addcmul_ = None
         return ()
-""",
+""",  # noqa: B950
         )
 
+    @skipIfCrossRef
     def test_foreach_lerp_inplace_decomposition_disabled_capture_scalar(self):
         """With decompositions disabled and capture_scalar_outputs=True, foreach_lerp_
         with scalar weight should work without graph breaks.
@@ -430,18 +451,23 @@ def fn(tensors, end_tensors):
             actual,
             """\
 class GraphModule(torch.nn.Module):
-    def forward(self, L_tensors_0_: "f32[4]", L_tensors_1_: "f32[4]", L_end_tensors_0_: \
-"f32[4]", L_end_tensors_1_: "f32[4]"):
-        l_tensors_0_ = L_tensors_0_
-        l_tensors_1_ = L_tensors_1_
+    def forward(self, L_end_tensors_0_: "f32[4]", L_end_tensors_1_: "f32[4]", L_tensors_0_: "f32[4]", L_tensors_1_: "f32[4]"):
         l_end_tensors_0_ = L_end_tensors_0_
         l_end_tensors_1_ = L_end_tensors_1_
+        l_tensors_0_ = L_tensors_0_
+        l_tensors_1_ = L_tensors_1_
 
-        _foreach_lerp_ = torch._foreach_lerp_([l_tensors_0_, l_tensors_1_], \
-[l_end_tensors_0_, l_end_tensors_1_], 0.5);  l_tensors_0_ = l_tensors_1_ = l_end_tensors_0_ = \
-l_end_tensors_1_ = _foreach_lerp_ = None
+        _foreach_sub = torch._foreach_sub([l_end_tensors_0_, l_end_tensors_1_], [l_tensors_0_, l_tensors_1_])
+        getitem: "f32[4]" = _foreach_sub[0]
+        getitem_1: "f32[4]" = _foreach_sub[1];  _foreach_sub = None
+        tensor: "f32[]" = torch.tensor(0.5, dtype = torch.float32, device = device(type='cpu'))
+        sub: "f32[]" = 1.0 - tensor;  tensor = None
+        neg_omw: "f32[]" = -sub;  sub = None
+        copy_: "f32[4]" = l_tensors_0_.copy_(l_end_tensors_0_);  l_end_tensors_0_ = copy_ = None
+        copy__1: "f32[4]" = l_tensors_1_.copy_(l_end_tensors_1_);  l_end_tensors_1_ = copy__1 = None
+        _foreach_addcmul_ = torch._foreach_addcmul_([l_tensors_0_, l_tensors_1_], [neg_omw, neg_omw], (getitem, getitem_1));  l_tensors_0_ = l_tensors_1_ = neg_omw = getitem = getitem_1 = _foreach_addcmul_ = None
         return ()
-""",
+""",  # noqa: B950
         )
 
     def test_foreach_pow_scalar_decomposition_enabled(self):
@@ -542,6 +568,194 @@ def forward(self, L_exps_0_: "f32[4]", L_exps_1_: "f32[4]"):
 """,
         )
 
+    @skipIfCrossRef
+    @torch._dynamo.config.patch(enable_dynamo_decompositions=True)
+    def test_addcmul_tensor_value_numerics(self):
+        """Compiled addcmul_ with tensor value matches eager.
+
+        Not bitwise on CPU: inductor may decompose fma to mul+add rather
+        than emitting a hardware fma instruction.
+        """
+
+        def fn(x, tensor1, tensor2, value):
+            return x.addcmul_(tensor1, tensor2, value=value)
+
+        x = torch.randn(4)
+        tensor1 = torch.randn(4)
+        tensor2 = torch.randn(4)
+        value = torch.tensor(0.5)
+
+        expected = fn(x.clone(), tensor1, tensor2, value)
+        actual = torch.compile(fn, fullgraph=True)(x.clone(), tensor1, tensor2, value)
+        self.assertEqual(expected, actual)
+
+    @skipIfCrossRef
+    @torch._dynamo.config.patch(enable_dynamo_decompositions=True)
+    def test_addcdiv_tensor_value_numerics(self):
+        """Compiled addcdiv_ with tensor value matches eager."""
+
+        def fn(x, tensor1, tensor2, value):
+            return x.addcdiv_(tensor1, tensor2, value=value)
+
+        x = torch.randn(4)
+        tensor1 = torch.randn(4)
+        tensor2 = torch.randn(4) + 0.1
+        value = torch.tensor(0.5)
+
+        expected = fn(x.clone(), tensor1, tensor2, value)
+        actual = torch.compile(fn, fullgraph=True)(x.clone(), tensor1, tensor2, value)
+        self.assertEqual(expected, actual, atol=0, rtol=0)
+
+    @skipIfCrossRef
+    @torch._dynamo.config.patch(enable_dynamo_decompositions=True)
+    def test_add_tensor_alpha_numerics(self):
+        """Compiled add_ with tensor alpha matches eager."""
+
+        def fn(x, other, alpha):
+            return x.add_(other, alpha=alpha)
+
+        x = torch.randn(4)
+        other = torch.randn(4)
+        alpha = torch.tensor(2.0)
+
+        expected = fn(x.clone(), other, alpha)
+        actual = torch.compile(fn, fullgraph=True)(x.clone(), other, alpha)
+        self.assertEqual(expected, actual, atol=0, rtol=0)
+
+    @skipIfCrossRef
+    @torch._dynamo.config.patch(enable_dynamo_decompositions=True)
+    @unittest.skipUnless(torch.cuda.is_available(), "requires CUDA")
+    def test_add_tensor_alpha_fma_matches_aten_cuda(self):
+        """On CUDA, ATen add_ with tensor alpha extracts the scalar and uses
+        fma(other, alpha, self). Our decomposition must use fma to match."""
+        torch.manual_seed(42)
+        x = torch.randn(64, 64, device="cuda")
+        other = torch.randn(64, 64, device="cuda")
+        alpha = torch.tensor(2.3, device="cuda")
+
+        def fn(x, other, alpha):
+            return x.add_(other, alpha=alpha)
+
+        expected = fn(x.clone(), other, alpha)
+        actual = torch.compile(fn, fullgraph=True)(x.clone(), other, alpha)
+        self.assertEqual(expected, actual, atol=0, rtol=0)
+
+    @skipIfCrossRef
+    @torch._dynamo.config.patch(enable_dynamo_decompositions=True)
+    @unittest.skipUnless(torch.cuda.is_available(), "requires CUDA")
+    def test_addcmul_value_1_fma_matches_aten_cuda(self):
+        """On CUDA, ATen addcmul_ with value=1 uses hardware fma(t1, t2, self).
+
+        Our decomposition uses inductor_prims.fma for this case. Without fma,
+        mul(t1, t2) + self rounds the product first, causing ~7% element
+        mismatches on typical inputs (e.g. Adagrad's addcmul_(grad, grad, value=1)).
+        """
+        torch.manual_seed(42)
+        x = torch.randn(64, 64, device="cuda")
+        t1 = torch.randn(64, 64, device="cuda")
+
+        def fn(x, t1):
+            # value=1 is a constant, triggers fma path in decomposition
+            return x.addcmul_(t1, t1, value=1)
+
+        expected = fn(x.clone(), t1)
+        actual = torch.compile(fn, fullgraph=True)(x.clone(), t1)
+        self.assertEqual(expected, actual, atol=0, rtol=0)
+
+    @skipIfCrossRef
+    @torch._dynamo.config.patch(enable_dynamo_decompositions=True)
+    @unittest.skipUnless(torch.cuda.is_available(), "requires CUDA")
+    def test_addcmul_scalar_value_cuda(self):
+        """Compiled addcmul_ with scalar value matches eager on CUDA."""
+        torch.manual_seed(42)
+        x = torch.randn(64, 64, device="cuda")
+        t1 = torch.randn(64, 64, device="cuda")
+        t2 = torch.randn(64, 64, device="cuda")
+
+        def fn(x, t1, t2):
+            return x.addcmul_(t1, t2, value=0.5)
+
+        expected = fn(x.clone(), t1, t2)
+        actual = torch.compile(fn, fullgraph=True)(x.clone(), t1, t2)
+        self.assertEqual(expected, actual, atol=0, rtol=0)
+
+    @skipIfCrossRef
+    @torch._dynamo.config.patch(enable_dynamo_decompositions=True)
+    @unittest.skipUnless(torch.cuda.is_available(), "requires CUDA")
+    def test_addcmul_tensor_value_cuda(self):
+        """Compiled addcmul_ with tensor value matches eager on CUDA."""
+        torch.manual_seed(42)
+        x = torch.randn(64, 64, device="cuda")
+        t1 = torch.randn(64, 64, device="cuda")
+        t2 = torch.randn(64, 64, device="cuda")
+        value = torch.tensor(0.5, device="cuda")
+
+        def fn(x, t1, t2, value):
+            return x.addcmul_(t1, t2, value=value)
+
+        expected = fn(x.clone(), t1, t2, value)
+        actual = torch.compile(fn, fullgraph=True)(x.clone(), t1, t2, value)
+        self.assertEqual(expected, actual, atol=0, rtol=0)
+
+    @skipIfCrossRef
+    @torch._dynamo.config.patch(enable_dynamo_decompositions=True)
+    @unittest.skipUnless(torch.cuda.is_available(), "requires CUDA")
+    def test_addcdiv_scalar_value_cuda(self):
+        """Compiled addcdiv_ with scalar value matches eager on CUDA.
+
+        Not bitwise: ATen inlines the division into fma(alpha, t1/t2, input)
+        which nvcc can optimize differently than separate div + fma kernels.
+        """
+        torch.manual_seed(42)
+        x = torch.randn(64, 64, device="cuda")
+        t1 = torch.randn(64, 64, device="cuda")
+        t2 = torch.randn(64, 64, device="cuda") + 0.1
+
+        def fn(x, t1, t2):
+            return x.addcdiv_(t1, t2, value=-0.01)
+
+        expected = fn(x.clone(), t1, t2)
+        actual = torch.compile(fn, fullgraph=True)(x.clone(), t1, t2)
+        self.assertEqual(expected, actual)
+
+    @skipIfCrossRef
+    @torch._dynamo.config.patch(enable_dynamo_decompositions=True)
+    @unittest.skipUnless(torch.cuda.is_available(), "requires CUDA")
+    def test_addcdiv_tensor_value_cuda(self):
+        """Compiled addcdiv_ with tensor value matches eager on CUDA.
+
+        Not bitwise: ATen inlines the division into fma(alpha, t1/t2, input)
+        which nvcc can optimize differently than separate div + fma kernels.
+        """
+        torch.manual_seed(42)
+        x = torch.randn(64, 64, device="cuda")
+        t1 = torch.randn(64, 64, device="cuda")
+        t2 = torch.randn(64, 64, device="cuda") + 0.1
+        value = torch.tensor(-0.01, device="cuda")
+
+        def fn(x, t1, t2, value):
+            return x.addcdiv_(t1, t2, value=value)
+
+        expected = fn(x.clone(), t1, t2, value)
+        actual = torch.compile(fn, fullgraph=True)(x.clone(), t1, t2, value)
+        self.assertEqual(expected, actual)
+
+    @skipIfCrossRef
+    @torch._dynamo.config.patch(enable_dynamo_decompositions=True)
+    @unittest.skipUnless(torch.cuda.is_available(), "requires CUDA")
+    def test_add_scalar_alpha_cuda(self):
+        """Compiled add_ with scalar alpha matches eager on CUDA."""
+        torch.manual_seed(42)
+        x = torch.randn(64, 64, device="cuda")
+        other = torch.randn(64, 64, device="cuda")
+
+        def fn(x, other):
+            return x.add_(other, alpha=2.3)
+
+        expected = fn(x.clone(), other)
+        actual = torch.compile(fn, fullgraph=True)(x.clone(), other)
+        self.assertEqual(expected, actual, atol=0, rtol=0)
+
 
 if __name__ == "__main__":
     run_tests()
diff --git a/test/dynamo/test_einops.py b/test/dynamo/test_einops.py
index fcd86e50b944c..5e8a77098c127 100644
--- a/test/dynamo/test_einops.py
+++ b/test/dynamo/test_einops.py
@@ -11,6 +11,7 @@
 from torch.testing._internal.common_utils import (
     instantiate_parametrized_tests,
     parametrize,
+    xfailIf,
 )
 
 
@@ -23,6 +24,7 @@
 else:
     einops_version = "none"
 einops_version_sanitized = einops_version.replace(".", "_")
+HAS_EINOPS_PACK = HAS_EINOPS and hasattr(einops, "pack")
 
 
 @unittest.skipIf(not HAS_EINOPS, "these tests require einops")
@@ -36,12 +38,13 @@ class TestEinops(TestCase):
     in PyTorch.
     """
 
-    @unittest.skipIf(
-        einops_version == "0.6.1", "https://github.com/pytorch/pytorch/issues/157417"
-    )
     @parametrize("version", [einops_version_sanitized])
     def test_functions(self, version):
-        from einops import einsum, pack, rearrange, reduce, repeat, unpack
+        from einops import einsum, rearrange, reduce, repeat
+
+        has_pack = HAS_EINOPS_PACK
+        if has_pack:
+            from einops import pack, unpack
 
         class TorchModuleWithOperations(nn.Module):
             def __init__(self) -> None:
@@ -60,11 +63,16 @@ def suf(pattern):
                 # by suf function
                 x_abcd = repeat(x_abc, suf("a b c -> a b c 4"))
                 x_abc = reduce(x_abcd, suf("a b c d -> a b c"), "min")
-                x_abdc, ps = pack([x_abc] * (2 + len(suffix)), suf("a b * c"))
-                x_array = unpack(
-                    rearrange(x_abdc, suf("a b d c -> (a b ) 1 c d")), ps, "ab one1 c *"
-                )
-                x1 = x_array[0] + len(x_array)
+                if has_pack:
+                    x_abdc, ps = pack([x_abc] * (2 + len(suffix)), suf("a b * c"))
+                    x_array = unpack(
+                        rearrange(x_abdc, suf("a b d c -> (a b ) 1 c d")),
+                        ps,
+                        "ab one1 c *",
+                    )
+                    x1 = x_array[0] + len(x_array)
+                else:
+                    x1 = rearrange(x_abc, suf("a b c -> (a b ) 1 c"))
                 x1 = rearrange(x1, suf("(a b ) 1 c -> a b c"), b=b)
                 addition = einsum(x_abc, x_abcd, suf("a b c , a b c d -> d"))[0]
                 return x1 + addition
@@ -102,7 +110,7 @@ def test_layers(self, version):
         for size in [16, 32, 64]:
             x = torch.rand([size, size])
             result1 = original(x)
-            result2 = compiled(x.double()).float()
+            result2 = compiled(x)
             self.assertEqual(result1, result2)
 
     @parametrize("version", [einops_version_sanitized])
@@ -112,7 +120,12 @@ def test_no_recompile_on_lazy_state(self, version):
         script = """\
 import torch
 import torch.nn as nn
-from einops import einsum, pack, reduce, repeat, unpack, rearrange
+import einops
+from einops import einsum, reduce, repeat, rearrange
+
+has_pack = hasattr(einops, "pack")
+if has_pack:
+    from einops import pack, unpack
 
 class TorchModuleWithOperations(nn.Module):
     def __init__(self) -> None:
@@ -129,14 +142,17 @@ def suf(pattern):
         # by suf function
         x_abcd = repeat(x_abc, suf("a b c -> a b c 4"))
         x_abc = reduce(x_abcd, suf("a b c d -> a b c"), "min")
-        x_abdc, ps = pack([x_abc] * (2 + len(suffix)), suf("a b * c"))
-        x_array = unpack(rearrange(x_abdc, suf("a b d c -> (a b ) 1 c d")), ps, "ab one1 c *")
-        x1 = x_array[0] + len(x_array)
+        if has_pack:
+            x_abdc, ps = pack([x_abc] * (2 + len(suffix)), suf("a b * c"))
+            x_array = unpack(rearrange(x_abdc, suf("a b d c -> (a b ) 1 c d")), ps, "ab one1 c *")
+            x1 = x_array[0] + len(x_array)
+        else:
+            x1 = rearrange(x_abc, suf("a b c -> (a b ) 1 c"))
         x1 = rearrange(x1, suf("(a b ) 1 c -> a b c"), b=b)
         addition = einsum(x_abc, x_abcd, suf("a b c , a b c d -> d"))[0]
         return x1 + addition
 
-compiled_fn = torch.compile(TorchModuleWithOperations(), fullgraph=True)
+compiled_fn = torch.compile(TorchModuleWithOperations(), fullgraph=True, backend="eager")
 x = torch.arange(2 * 3 * 5).view(2, 3, 5)
 y = compiled_fn(x)
 
@@ -190,6 +206,7 @@ def f(x):
             else:
                 self.assertIn(einops_method, output)
 
+    @xfailIf(einops_version == "0.8.2")
     @parametrize(
         "method",
         ["reduce", "repeat", "pack", "unpack", "einsum", "rearrange"],
@@ -222,6 +239,15 @@ def test_einops_method(self, method):
             self.fail(method)
         self._run_in_subprocess(flag, method, einops_method, snippet)
 
+    def test_no_warning(self):
+        # checks that this doesn't produce any warnings
+        @torch.compile(backend="eager", fullgraph=True)
+        def fn(x):
+            return einops.rearrange(x, "... -> (...)")
+
+        x = torch.randn(5)
+        self.assertNotWarn(lambda: fn(x))
+
 
 instantiate_parametrized_tests(
     TestEinops,
diff --git a/test/dynamo/test_enum.py b/test/dynamo/test_enum.py
index e8b2cfb80c7f6..d94c1ab97b84e 100644
--- a/test/dynamo/test_enum.py
+++ b/test/dynamo/test_enum.py
@@ -192,7 +192,7 @@ def f(x, e):
             else:
                 return y - 2
 
-        opt_f = torch.compile(fullgraph=True)(f)
+        opt_f = torch.compile(fullgraph=True, backend="eager")(f)
         args = [torch.zeros(1), Bool.TRUE]
         ref_out = f(*args)
         opt_out = opt_f(*args)
@@ -268,7 +268,7 @@ def fn(x, member):
         x = torch.randn(4)
         # With custom metaclass, membership should always be False
         ref = fn(x, MyEnum.A)
-        opt_fn = torch.compile(fn, backend="eager")
+        opt_fn = torch.compile(fn, backend="eager", fullgraph=True)
         res = opt_fn(x, MyEnum.A)
         self.assertEqual(ref, res)
         # Should return x - 1 since custom __contains__ returns False
@@ -291,7 +291,6 @@ def fn(x):
         res = opt_fn(x)
         self.assertEqual(ref, res)
 
-    @unittest.expectedFailure  # TODO: Support Enum.__iter__
     def test_enum_iter(self):
         """Test iterating over Enum class."""
 
@@ -304,6 +303,8 @@ def fn(x):
             total = 0
             for color in Color:
                 total += color.value
+            for color in Color.__iter__():
+                total += color.value
             return x + total
 
         x = torch.randn(4)
@@ -312,7 +313,6 @@ def fn(x):
         res = opt_fn(x)
         self.assertEqual(ref, res)
 
-    @unittest.expectedFailure  # TODO: Support Enum.__iter__ (list uses iter)
     def test_enum_list(self):
         """Test list(Enum) to get all members."""
 
@@ -351,7 +351,6 @@ def fn(x):
         res = opt_fn(x)
         self.assertEqual(ref, res)
 
-    @unittest.expectedFailure  # TODO: Support Enum.__reversed__
     def test_enum_reversed(self):
         """Test reversed() on Enum class."""
 
@@ -429,7 +428,6 @@ def fn(x, priority):
         res = opt_fn(x, Priority.HIGH)
         self.assertEqual(ref, res)
 
-    @unittest.expectedFailure  # TODO: Support Flag enum membership check
     def test_flag_enum(self):
         """Test Flag enum operations."""
 
@@ -453,6 +451,25 @@ def fn(x, perm):
         res = opt_fn(x, combined)
         self.assertEqual(ref, res)
 
+    def test_flag_enum_membership_combined(self):
+        """Test Flag enum membership check with combined flags."""
+
+        class LocalReduction(enum.Flag):
+            MEAN = enum.auto()
+            SUM = enum.auto()
+            MAX = enum.auto()
+
+        @torch.compile(backend="eager", fullgraph=True)
+        def fn(x, local_reduction):
+            if LocalReduction.MEAN in local_reduction:
+                return x.mean()
+            return x.sum()
+
+        x = torch.tensor([1.0, 2.0, 3.0])
+        self.assertEqual(fn(x, LocalReduction.MEAN | LocalReduction.SUM), x.mean())
+        torch._dynamo.reset()
+        self.assertEqual(fn(x, LocalReduction.SUM | LocalReduction.MAX), x.sum())
+
     def test_enum_comparison(self):
         """Test enum comparison operations."""
 
@@ -484,6 +501,74 @@ def fn(x, a, b):
         res = opt_fn(x, Priority.LOW, Priority.HIGH)
         self.assertEqual(ref, res)
 
+    @unittest.expectedFailure
+    def test_metaclass_custom_call(self):
+        # When a class has a metaclass that overrides __call__, Python invokes
+        # the metaclass's __call__ instead of the default type.__call__
+        # (which does __new__ + __init__). Dynamo currently bypasses the
+        # metaclass __call__ and uses instantiate_user_defined_class_object
+        # (which mimics type.__call__), producing wrong results.
+        class ScalingMeta(type):
+            def __call__(cls, val):
+                obj = cls.__new__(cls)
+                obj.scaled = val * cls.scale_factor
+                return obj
+
+        class Scaled(metaclass=ScalingMeta):
+            scale_factor = 10
+
+            def __init__(self, val):
+                # type.__call__ would call this, but ScalingMeta.__call__
+                # deliberately skips it and sets scaled directly.
+                self.scaled = val
+
+        def fn(x, val):
+            obj = Scaled(val)
+            return x * obj.scaled
+
+        x = torch.tensor([1.0, 2.0])
+        # Python: ScalingMeta.__call__ sets scaled = 3 * 10 = 30
+        ref = fn(x, 3)
+        compiled_fn = torch.compile(fn, backend="eager")
+        res = compiled_fn(x, 3)
+        self.assertEqual(ref, res)
+
+    def test_enum_construction_no_extra_init(self):
+        # Real-world instance of the metaclass __call__ issue above.
+        # EnumMeta.__call__ only calls __new__ (value lookup), NOT __init__.
+        # Dynamo's instantiate_user_defined_class_object calls both, so the
+        # extra __init__ call produces wrong results.
+        init_count = [0]
+
+        class MyEnum(enum.Enum):
+            def __new__(cls, val):
+                obj = object.__new__(cls)
+                obj._value_ = val
+                return obj
+
+            def __init__(self, val):
+                init_count[0] += 1
+
+            A = 1
+            B = 2
+
+        init_count[0] = 0
+
+        def fn(x, val):
+            MyEnum(val)
+            return x * (1 + init_count[0])
+
+        x = torch.tensor([1.0, 2.0])
+        # Python: MyEnum(1) is a value lookup via EnumMeta.__call__, which
+        # only calls Enum.__new__ (lookup) — __init__ is NOT called.
+        # So init_count stays 0 and fn returns x * 1.
+        ref = fn(x, 1)
+        init_count[0] = 0
+
+        compiled_fn = torch.compile(fn, backend="eager")
+        res = compiled_fn(x, 1)
+        self.assertEqual(ref, res)
+
 
 if __name__ == "__main__":
     from torch._dynamo.test_case import run_tests
diff --git a/test/dynamo/test_error_messages.py b/test/dynamo/test_error_messages.py
index 9e789e5108856..02a94b0987005 100644
--- a/test/dynamo/test_error_messages.py
+++ b/test/dynamo/test_error_messages.py
@@ -16,11 +16,7 @@
 from torch._dynamo.exc import ResumePrologueTracingError, TorchRuntimeError, Unsupported
 from torch._dynamo.testing import skipIfNotPy312, skipIfOnlyNotPy312
 from torch._dynamo.utils import counters
-from torch.testing._internal.common_utils import (
-    IS_FBCODE,
-    munge_exc,
-    scoped_load_inline,
-)
+from torch.testing._internal.common_utils import IS_FBCODE, munge_exc
 from torch.testing._internal.logging_utils import LoggingTestCase, make_logging_test
 
 
@@ -292,7 +288,7 @@ def post_munge(s):
   Hint: Apply `@torch._dynamo.dont_skip_tracing` to the function `skip` to force tracing into the function. More graph breaks may occur as a result of attempting to trace into the function.
   Hint: Please file an issue to PyTorch.
 
-  Developer debug context: module: unittest.case, qualname: skip, skip reason: <missing reason>
+  Developer debug context: module: unittest.case, qualname: skip, skip reason: file is in unittest directory
 
  For more details about this graph break, please visit: https://meta-pytorch.github.io/compile-graph-break-site/gb/gb0007.html
 
@@ -306,6 +302,13 @@ def test_skipfile_dynamo_call(self):
         def fn():
             torch._dynamo.disable()
 
+        def post_munge(s):
+            return re.sub(
+                r"file matches MOD_SKIPLIST \(.*?\)",
+                "file matches MOD_SKIPLIST (<path>)",
+                s,
+            )
+
         self.assertExpectedInlineMunged(
             Unsupported,
             lambda: torch.compile(fn, backend="eager", fullgraph=True)(),
@@ -314,13 +317,14 @@ def fn():
   Explanation: Dynamo developers have intentionally marked that the function `disable` in file `_dynamo/decorators.py` should not be traced.
   Hint: Avoid calling the function `disable`.
 
-  Developer debug context: module: torch._dynamo.decorators, qualname: disable, skip reason: <missing reason>
+  Developer debug context: module: torch._dynamo.decorators, qualname: disable, skip reason: file matches MOD_SKIPLIST (<path>)
 
  For more details about this graph break, please visit: https://meta-pytorch.github.io/compile-graph-break-site/gb/gb0007.html
 
 from user code:
    File "test_error_messages.py", line N, in fn
     torch._dynamo.disable()""",
+            post_munge=post_munge,
         )
 
     def test_skipfile_inline(self):
@@ -343,7 +347,7 @@ def post_munge(s):
   Hint: Apply `@torch._dynamo.dont_skip_tracing` to the function `skip` to force tracing into the function. More graph breaks may occur as a result of attempting to trace into the function.
   Hint: Please file an issue to PyTorch.
 
-  Developer debug context: qualname: skip, name: skip, filename: `case.py`, skip reason: skipped according trace_rules.lookup unittest
+  Developer debug context: qualname: skip, name: skip, filename: `case.py`, skip reason: file is in unittest directory
 
  For more details about this graph break, please visit: https://meta-pytorch.github.io/compile-graph-break-site/gb/gb0008.html
 
@@ -408,7 +412,7 @@ def fn():
   Hint: If you are attempting to call a logging function (e.g. `_warnings.warn`), you can try adding it to `torch._dynamo.config.reorderable_logging_functions`.
   Hint: Please file an issue on GitHub so the PyTorch team can add support for it.
 
-  Developer debug context: module: _warnings, qualname: warn, skip reason: <missing reason>
+  Developer debug context: module: _warnings, qualname: warn, skip reason: cannot determine source file for _warnings (likely a C extension or builtin)
 
  For more details about this graph break, please visit: https://meta-pytorch.github.io/compile-graph-break-site/gb/gb0007.html
 
@@ -456,97 +460,11 @@ def post_munge(string):
   Explanation: Dynamo cannot trace optree C/C++ function optree.<path>.make_from_collection.
   Hint: Consider using torch.utils._pytree - https://github.com/pytorch/pytorch/blob/main/torch/utils/_pytree.py
 
-  Developer debug context: module: optree._C, qualname: <path>.make_from_collection, skip reason: <missing reason>
+  Developer debug context: module: optree._C, qualname: <path>.make_from_collection, skip reason: cannot determine source file for optree._C (likely a C extension or builtin)
 
  For more details about this graph break, please visit: https://meta-pytorch.github.io/compile-graph-break-site/gb/gb0007.html""",
         )
 
-    @scoped_load_inline
-    @torch._dynamo.config.patch(inline_inbuilt_nn_modules=False)
-    @unittest.skipIf(IS_FBCODE, "inline cpp_extension doesn't work in fbcode")
-    def test_cpp_extension_recommends_custom_ops(self, load_inline):
-        cpp_source = """
-        #include <torch/extension.h>
-        at::Tensor foobar(const at::Tensor& x) {
-            return x.clone();
-        }
-        """
-        module = load_inline(
-            name="mylib",
-            cpp_sources=cpp_source,
-            functions="foobar",
-            verbose=True,
-        )
-
-        x = torch.ones(2, 2, requires_grad=True)
-        counters.clear()
-
-        @torch.compile(backend="eager")
-        def f(x):
-            return module.foobar(x)
-
-        with self.assertWarnsOnceRegex(
-            UserWarning,
-            "(?s).*https://pytorch.org/tutorials/advanced/custom_ops_landing_page.html.*",
-        ):
-            f(x)
-        self.assertEqual(len(counters["graph_break"]), 1)
-        first_graph_break = next(iter(counters["graph_break"].keys()))
-
-        first_graph_break = re.sub(r"mylib(_v\d+)?", "mylib", first_graph_break)
-        # HACK: this patches around the fact that PyBind11 improperly sets the
-        # __qualname__ attribute on functions and methods; see
-        # https://github.com/pybind/pybind11/issues/5774.  This should be removed if
-        # that issue is fixed.
-        first_graph_break = re.sub(
-            r"pybind11_detail_function_record_v[^ .]+", "PyCapsule", first_graph_break
-        )
-
-        self.assertExpectedInline(
-            first_graph_break,
-            """\
-Attempted to call function marked as skipped
-  Explanation: Dynamo does not know how to trace the builtin `mylib.PyCapsule.foobar.` This function is either a Python builtin (e.g. _warnings.warn) or a third-party C/C++ Python extension (perhaps created with pybind).
-  Hint: If it is a Python builtin, please file an issue on GitHub so the PyTorch team can add support for it and see the next case for a workaround.
-  Hint: If it is a third-party C/C++ Python extension, please either wrap it into a PyTorch-understood custom operator (see https://pytorch.org/tutorials/advanced/custom_ops_landing_page.html for more details) or, if it is traceable, use `torch.compiler.allow_in_graph`.
-
-  Developer debug context: module: mylib, qualname: PyCapsule.foobar, skip reason: <missing reason>
-
- For more details about this graph break, please visit: https://meta-pytorch.github.io/compile-graph-break-site/gb/gb0007.html""",
-        )
-
-        cpp_source = """
-        #include <torch/extension.h>
-        at::Tensor baz(const at::Tensor& x) {
-            return x.clone();
-        }
-        """
-        module2 = load_inline(
-            name="mylib2",
-            cpp_sources=cpp_source,
-            functions="baz",
-            verbose=True,
-        )
-
-        torch._dynamo.reset()
-
-        # Test that each warning only happens once
-        @torch.compile(backend="eager")
-        def f(x):
-            module2.baz(x)
-            module.foobar(x)
-            module.foobar(x)
-            module2.baz(x)
-            module.foobar(x)
-            module2.baz(x)
-            return x.clone()
-
-        with warnings.catch_warnings(record=True) as ws:
-            warnings.simplefilter("always")
-            f(x)
-            f(x)
-        self.assertEqual(len(ws), 2)
-
     def test_observed_exception(self):
         def fn():
             raise RuntimeError("test")
@@ -560,7 +478,7 @@ def fn():
   Hint: Your code may result in an error when running in eager. Please double check that your code doesn't contain a similar error when actually running eager/uncompiled. You can do this by removing the `torch.compile` call, or by using `torch.compiler.set_stance("force_eager")`.
   Hint: It may be possible to write Dynamo tracing rules for this code. Please report an issue to PyTorch if you encounter this graph break often and it is causing performance issues.
 
-  Developer debug context: raised exception RuntimeError([ConstantVariable(str: 'test')])
+  Developer debug context: raised exception RuntimeError('test')
 
  For more details about this graph break, please visit: https://meta-pytorch.github.io/compile-graph-break-site/gb/gb0088.html
 
@@ -595,35 +513,6 @@ def fn(mod):
     return mod(1)""",
         )
 
-    @torch._dynamo.config.patch(inline_inbuilt_nn_modules=False)
-    def test_class_property(self):
-        class Foo(torch.nn.Module):
-            attr = unittest
-
-        def fn(mod, x):
-            return mod.attr
-
-        self.assertExpectedInlineMunged(
-            Unsupported,
-            lambda: torch.compile(fn, backend="eager", fullgraph=True)(
-                Foo(), torch.randn(3)
-            ),
-            """\
-Unsupported nn.Module attribute type
-  Explanation: Dynamo does not support tracing nn.Module attributes of type `module`
-  Hint: Refactor your code so that `attr` (type `module`) is not an attribute of `Foo`
-  Hint: Currently supported attribute types are methods, classmethods, staticmethods, properties, constants, and tensors.
-  Hint: It may be possible to write Dynamo tracing rules for this code. Please report an issue to PyTorch if you encounter this graph break often and it is causing performance issues.
-
-  Developer debug context: nn.Module subclass: Foo, name: attr, attribute type: module
-
- For more details about this graph break, please visit: https://meta-pytorch.github.io/compile-graph-break-site/gb/gb0161.html
-
-from user code:
-   File "test_error_messages.py", line N, in fn
-    return mod.attr""",
-        )
-
     def test_generic_ctx_mgr_graph_break_fullgraph_true(self):
         def fn():
             with GenericCtxMgr():
@@ -708,7 +597,7 @@ class Foo:
   Hint: If it is a Python builtin, please file an issue on GitHub so the PyTorch team can add support for it and see the next case for a workaround.
   Hint: If it is a third-party C/C++ Python extension, please either wrap it into a PyTorch-understood custom operator (see https://pytorch.org/tutorials/advanced/custom_ops_landing_page.html for more details) or, if it is traceable, use `torch.compiler.allow_in_graph`.
 
-  Developer debug context: module: builtins, qualname: __build_class__, skip reason: <missing reason>
+  Developer debug context: module: builtins, qualname: __build_class__, skip reason: cannot determine source file for builtins (likely a C extension or builtin)
 
  For more details about this graph break, please visit: https://meta-pytorch.github.io/compile-graph-break-site/gb/gb0007.html
 
@@ -916,6 +805,12 @@ def fn(x):
             """\
 Data-dependent branching
   Explanation: Detected data-dependent branching (e.g. `if my_tensor.sum() > 0:`). Dynamo does not support tracing dynamic control flow.
+
+      The branch condition involves a tensor computed as follows:
+        # File "test_error_messages.py", line N, in fn, code: if x.sum() > 0:
+        gt = gt(sum_1, 0)
+
+  Hint: The branch condition uses a scalar integer tensor. Consider rewriting the computation to use plain Python ints (e.g. use int attributes instead of tensor buffers) so the condition becomes a shape guard instead of data-dependent branching.
   Hint: This graph break is fundamental - it is unlikely that Dynamo will ever be able to trace through your code. Consider finding a workaround.
   Hint: Use `torch.cond` to express dynamic control flow.
 
@@ -970,7 +865,7 @@ def find_trace_bytecode_lines(long_string):
             "LazyVariableTracker(unrealized: <class 'function'>)", all_lines[0]
         )
         self.assertIn(
-            "LazyVariableTracker(realized: UserFunctionVariable())", all_lines[3]
+            "LazyVariableTracker(unrealized: <class 'torch.Tensor'>)", all_lines[3]
         )
 
     @make_logging_test(graph_breaks=True)
@@ -988,6 +883,12 @@ def fn(x):
 Graph break in user code at test_error_messages.py:N
 Graph Break Reason: Data-dependent branching
   Explanation: Detected data-dependent branching (e.g. `if my_tensor.sum() > 0:`). Dynamo does not support tracing dynamic control flow.
+
+      The branch condition involves a tensor computed as follows:
+        # File "test_error_messages.py", line N, in fn, code: if x.sum() > 0:
+        gt = gt(sum_1, 0)
+
+  Hint: The branch condition uses a scalar integer tensor. Consider rewriting the computation to use plain Python ints (e.g. use int attributes instead of tensor buffers) so the condition becomes a shape guard instead of data-dependent branching.
   Hint: This graph break is fundamental - it is unlikely that Dynamo will ever be able to trace through your code. Consider finding a workaround.
   Hint: Use `torch.cond` to express dynamic control flow.
 
@@ -1008,7 +909,7 @@ def fn(x):
     def test_assert_failure_in_generic_ctx_mgr(self, records):
         def fn(x):
             with GenericCtxMgr():
-                assert x is None
+                assert x is None  # noqa: S101
 
         with self.assertRaises(AssertionError):
             torch.compile(fn, backend="eager")(torch.randn(3))
@@ -1036,7 +937,7 @@ def fn(x):
   File "test_error_messages.py", line N, in test_assert_failure_in_generic_ctx_mgr
     torch.compile(fn, backend="eager")(torch.randn(3))
   File "test_error_messages.py", line N, in fn
-    assert x is None
+    assert x is None  # noqa: S101
 """,
         )
 
@@ -1203,6 +1104,12 @@ def gn(x):
 
 Data-dependent branching
   Explanation: Detected data-dependent branching (e.g. `if my_tensor.sum() > 0:`). Dynamo does not support tracing dynamic control flow.
+
+      The branch condition involves a tensor computed as follows:
+        # File "test_error_messages.py", line N, in gn, code: if x.sum() > 0:
+        gt = gt(sum_1, 0)
+
+  Hint: The branch condition uses a scalar integer tensor. Consider rewriting the computation to use plain Python ints (e.g. use int attributes instead of tensor buffers) so the condition becomes a shape guard instead of data-dependent branching.
   Hint: This graph break is fundamental - it is unlikely that Dynamo will ever be able to trace through your code. Consider finding a workaround.
   Hint: Use `torch.cond` to express dynamic control flow.
 
@@ -1247,6 +1154,12 @@ def fn(x):
 
 Data-dependent branching
   Explanation: Detected data-dependent branching (e.g. `if my_tensor.sum() > 0:`). Dynamo does not support tracing dynamic control flow.
+
+      The branch condition involves a tensor computed as follows:
+        # File "test_error_messages.py", line N, in fn, code: if x.sum() > 0:
+        gt = gt(sum_1, 0)
+
+  Hint: The branch condition uses a scalar integer tensor. Consider rewriting the computation to use plain Python ints (e.g. use int attributes instead of tensor buffers) so the condition becomes a shape guard instead of data-dependent branching.
   Hint: This graph break is fundamental - it is unlikely that Dynamo will ever be able to trace through your code. Consider finding a workaround.
   Hint: Use `torch.cond` to express dynamic control flow.
 
@@ -1425,6 +1338,9 @@ def fn(x):
         self.assertEqual((len(matches) <= 20), True)
         self.assertIn("Most recent bytecode instructions traced (max 20):", s)
 
+    # TODO this test is broken with nested_graph_breaks because we need to update
+    # the resume collapse function for nested graph breaks
+    @torch._dynamo.config.patch(nested_graph_breaks=False)
     @torch._dynamo.config.patch(verbose=True)
     @make_logging_test(graph_breaks=True)
     def test_graph_break_traceback_above_dynamo_shows_user_code(self, records):
@@ -1811,8 +1727,9 @@ def fn(x):
     @make_logging_test(graph_breaks=True)
     def test_store_attr_graph_break(self, records):
         class Foo:
+            @torch.compiler.disable
             def __setattr__(self, name, value):
-                torch._dynamo.graph_break()
+                super().__setattr__(name, value)
 
         @torch.compile(backend="eager")
         def fn(x):
@@ -1820,27 +1737,30 @@ def fn(x):
 
         fn(torch.ones(3))
 
+        def post_munge(s):
+            return re.sub(r"0x[0-9A-Fa-f]+", "0xmem_addr", s)
+
         self.assertExpectedInline(
-            munge_exc(records[0].getMessage(), suppress_suffix=True, skip=0),
+            post_munge(
+                munge_exc(records[-1].getMessage(), suppress_suffix=True, skip=0)
+            ),
             """\
 Graph break in user code at test_error_messages.py:N
 Graph Break Reason: Encountered graph break when attempting to trace STORE_ATTR: storing an object's attribute, e.g. x.attr = y:
 
-Call to `torch._dynamo.graph_break()`
-  Explanation: User-inserted graph break. Message: None
-  Hint: Remove the `torch._dynamo.graph_break()` call.
+Skip inlining `torch.compiler.disable()`d function
+  Explanation: Skip inlining function <function ErrorMessagesTest.test_store_attr_graph_break.<locals>.Foo.__setattr__ at 0xmem_addr> since it was wrapped with `torch.compiler.disable` (reason: None)
+  Hint: Remove the `torch.compiler.disable` call
 
-  Developer debug context: Called `torch._dynamo.graph_break()` with args `[]`, kwargs `{}`
+  Developer debug context: <function ErrorMessagesTest.test_store_attr_graph_break.<locals>.Foo.__setattr__ at 0xmem_addr>
 
- For more details about this graph break, please visit: https://meta-pytorch.github.io/compile-graph-break-site/gb/gb0025.html
+ For more details about this graph break, please visit: https://meta-pytorch.github.io/compile-graph-break-site/gb/gb0099.html
 
 User code traceback:
   File "test_error_messages.py", line N, in test_store_attr_graph_break
     fn(torch.ones(3))
   File "test_error_messages.py", line N, in fn
     Foo().attr = x
-  File "test_error_messages.py", line N, in __setattr__
-    torch._dynamo.graph_break()
 """,
         )
 
@@ -1851,19 +1771,18 @@ def fn(x):
                 Unsupported,
                 lambda: fn(torch.ones(3)),
                 """\
-Call to `torch._dynamo.graph_break()`
-  Explanation: User-inserted graph break. Message: None
-  Hint: Remove the `torch._dynamo.graph_break()` call.
+Skip inlining `torch.compiler.disable()`d function
+  Explanation: Skip inlining function <function ErrorMessagesTest.test_store_attr_graph_break.<locals>.Foo.__setattr__ at 0xmem_addr> since it was wrapped with `torch.compiler.disable` (reason: None)
+  Hint: Remove the `torch.compiler.disable` call
 
-  Developer debug context: Called `torch._dynamo.graph_break()` with args `[]`, kwargs `{}`
+  Developer debug context: <function ErrorMessagesTest.test_store_attr_graph_break.<locals>.Foo.__setattr__ at 0xmem_addr>
 
- For more details about this graph break, please visit: https://meta-pytorch.github.io/compile-graph-break-site/gb/gb0025.html
+ For more details about this graph break, please visit: https://meta-pytorch.github.io/compile-graph-break-site/gb/gb0099.html
 
 from user code:
    File "test_error_messages.py", line N, in fn
-    Foo().attr = x
-  File "test_error_messages.py", line N, in __setattr__
-    torch._dynamo.graph_break()""",
+    Foo().attr = x""",
+                post_munge=post_munge,
             )
 
     def test_runtime_error_readable_shape_mismatch(self):
@@ -2191,39 +2110,6 @@ def post_munge(s):
  For more details about this graph break, please visit: https://meta-pytorch.github.io/compile-graph-break-site/gb/gb0039.html""",
             )
 
-    @unittest.skipIf(
-        not torch.utils._triton.has_triton()
-        or not hasattr(__import__("triton"), "set_allocator"),
-        "requires triton with set_allocator support",
-    )
-    def test_triton_set_allocator(self):
-        import triton
-
-        def fn(x):
-            triton.set_allocator(lambda size, align, stream: None)
-            return x * 2 + 1
-
-        self.assertExpectedInlineMunged(
-            Unsupported,
-            lambda: torch.compile(fn, backend="eager", fullgraph=True)(torch.randn(10)),
-            """\
-triton.set_allocator not supported
-  Explanation: triton.set_allocator is not supported inside torch.compile. It modifies global Triton allocator state and cannot be traced.
-  Hint: Move triton.set_allocator() outside of the torch.compile region (call it before the compiled function).
-
-  Developer debug context: triton.set_allocator called inside compiled region
-
- For more details about this graph break, please visit: https://meta-pytorch.github.io/compile-graph-break-site/gb/gb4026.html
-
-from user code:
-   File "test_error_messages.py", line N, in fn
-    triton.set_allocator(lambda size, align, stream: None)""",
-        )
-
-
-class NestedGraphBreakLoggingTests(
-    LoggingTestCase, torch._dynamo.test_case.TestCaseWithNestedGraphBreaks
-):
     @make_logging_test(graph_breaks=True)
     def test_nested_generic_ctx_mgr(self, records):
         def inner():
@@ -2314,13 +2200,6 @@ def fn():
     def test_skipped_frame_with_verbose_traceback_nested(self, records):
         global f1, f2, f3
 
-        class GenericCtxMgr:
-            def __enter__(self):
-                return self
-
-            def __exit__(self, exc_type, exc_value, traceback):
-                pass
-
         def f1(x):
             with GenericCtxMgr():
                 torch._dynamo.graph_break()
@@ -2375,13 +2254,6 @@ def f3(x):
     def test_skip_frame_in_loop_message_nested(self, records):
         global f1, f2, f3
 
-        class GenericCtxMgr:
-            def __enter__(self):
-                return self
-
-            def __exit__(self, exc_type, exc_value, traceback):
-                pass
-
         def f1(x):
             for i in range(2):
                 with GenericCtxMgr():
@@ -2405,6 +2277,12 @@ def f3(x):
 
 Data-dependent branching
   Explanation: Detected data-dependent branching (e.g. `if my_tensor.sum() > 0:`). Dynamo does not support tracing dynamic control flow.
+
+      The branch condition involves a tensor computed as follows:
+        # File "test_error_messages.py", line N, in f1, code: if x.sum() > 0:
+        gt = gt(sum_1, 0)
+
+  Hint: The branch condition uses a scalar integer tensor. Consider rewriting the computation to use plain Python ints (e.g. use int attributes instead of tensor buffers) so the condition becomes a shape guard instead of data-dependent branching.
   Hint: This graph break is fundamental - it is unlikely that Dynamo will ever be able to trace through your code. Consider finding a workaround.
   Hint: Use `torch.cond` to express dynamic control flow.
 
diff --git a/test/dynamo/test_exc.py b/test/dynamo/test_exc.py
index f382c1a216c59..3f934dc1f5ebd 100644
--- a/test/dynamo/test_exc.py
+++ b/test/dynamo/test_exc.py
@@ -7,7 +7,12 @@
 import torch._dynamo.config
 import torch._dynamo.test_case
 from torch._dynamo.comptime import comptime
-from torch._dynamo.exc import Unsupported
+from torch._dynamo.exc import (
+    TorchDynamoException,
+    Unsupported,
+    UserError,
+    UserErrorType,
+)
 from torch.testing._internal.common_device_type import skipIf
 from torch.testing._internal.common_utils import (
     IS_FBCODE,
@@ -21,6 +26,27 @@
 class ExcTests(LoggingTestCase):
     maxDiff = None
 
+    @torch._dynamo.config.patch(suppress_errors=True)
+    def test_user_error_backend_soft_fail(self):
+        def backend(_, __):
+            raise UserError(UserErrorType.INVALID_INPUT, "backend user error")
+
+        def fn(x):
+            return x + 1
+
+        x = torch.randn(2)
+        self.assertEqual(torch.compile(fn, backend=backend)(x), fn(x))
+
+    def test_user_error_separated_from_unsupported(self):
+        err = UserError(UserErrorType.INVALID_INPUT, "bad input")
+
+        self.assertIsInstance(err, TorchDynamoException)
+        self.assertNotIsInstance(err, Unsupported)
+        self.assertEqual(err.error_type, UserErrorType.INVALID_INPUT)
+        self.assertEqual(err.msg, "bad input")
+        self.assertEqual(err.message, "bad input")
+        self.assertEqual(str(err), "bad input")
+
     def test_unsupported_real_stack(self):
         # exercise Unsupported constructor and augment_exc_message
         def fn002(x):
diff --git a/test/dynamo/test_exceptions.py b/test/dynamo/test_exceptions.py
index 3201e25a91d12..a15f3a725dbce 100644
--- a/test/dynamo/test_exceptions.py
+++ b/test/dynamo/test_exceptions.py
@@ -1,6 +1,7 @@
 # Owner(s): ["module: dynamo"]
 
 import contextlib
+import dataclasses
 import sys
 
 import torch
@@ -161,7 +162,7 @@ def forward(ctx, x):
             def backward(ctx, grad_out):
                 return grad_out
 
-        @torch.compile
+        @torch.compile(backend="eager")
         def f(x: torch.Tensor):
             try:
                 with torch.autocast(device_type="cpu", dtype=None):
@@ -175,11 +176,11 @@ def f(x: torch.Tensor):
 
     @make_dynamo_test
     def test_isinstance_CustomException(self):
-        assert isinstance(CustomException, type)
-        assert not isinstance(CustomException(), type)
+        assert isinstance(CustomException, type)  # noqa: S101
+        assert not isinstance(CustomException(), type)  # noqa: S101
         C = CustomExceptionWithInstanceCheck
-        assert isinstance(C, C)
-        assert isinstance(C(), C)
+        assert isinstance(C, C)  # noqa: S101
+        assert isinstance(C(), C)  # noqa: S101
 
     @make_dynamo_test
     def test_propagate_exception_inside_ctx_manager(self):
@@ -206,7 +207,7 @@ def nothing():
                 z = 1
             except IndexError:
                 z = 2
-            assert z == 1
+            assert z == 1  # noqa: S101
 
     def test_exception_else(self):
         def gn(x):
@@ -281,11 +282,11 @@ def ctx():
             with ctx():
                 raise e
         except Exception as exc:
-            assert isinstance(exc, a)
-            assert isinstance(exc.__context__, b)
-            assert isinstance(exc.__context__.__context__, c)
-            assert isinstance(exc.__context__.__context__.__context__, d)
-            assert isinstance(exc.__context__.__context__.__context__.__context__, e)
+            assert isinstance(exc, a)  # noqa: S101
+            assert isinstance(exc.__context__, b)  # noqa: S101
+            assert isinstance(exc.__context__.__context__, c)  # noqa: S101
+            assert isinstance(exc.__context__.__context__.__context__, d)  # noqa: S101
+            assert isinstance(exc.__context__.__context__.__context__.__context__, e)  # noqa: S101
 
     # TODO(anijain2305) - does not work with fullgraph=True
     def test_exception_with_another_exception2(self):
@@ -438,7 +439,6 @@ def fn(x):
         res = opt_fn(x)
         self.assertEqual(ref, res)
 
-    @torch._dynamo.config.patch(inline_inbuilt_nn_modules=True)
     def test_custom_getattr_on_module_exception(self):
         class Foo(torch.nn.Module):
             def __init__(self, a=3):
@@ -527,6 +527,30 @@ def forward(self, x):
         metrics = torch._dynamo.utils.get_compilation_metrics()
         self.assertIn("Observed exception", metrics[0].fail_reason)
 
+    def test_observed_exception_formats_fstring_message(self):
+        from torch.utils._pytree import tree_map_with_path
+
+        def check_tensor(path, x):
+            if not isinstance(x, torch.Tensor):
+                raise ValueError(f"Expected Tensor at {path=}")
+            return x * 2
+
+        def fn(tree):
+            return tree_map_with_path(check_tensor, tree)
+
+        tree = {"a": torch.randn(10), "b": 5}
+
+        compiled_fn = torch.compile(fn, backend="eager", fullgraph=True)
+        with self.assertRaises(Unsupported) as compiled_ctx:
+            compiled_fn(tree)
+
+        exc_str = str(compiled_ctx.exception)
+        self.assertIn("Observed exception", exc_str)
+        self.assertIn("Expected Tensor at path=(MappingKey(key='b'),)", exc_str)
+        self.assertNotIn("Failed to trace builtin operator", exc_str)
+        self.assertNotIn("StringFormatVariable", exc_str)
+        self.assertNotIn("ConstantVariable(", exc_str)
+
     def test_key_error(self):
         def fn(x, d):
             try:
@@ -563,6 +587,25 @@ def fn(x):
         res = opt_fn(x)
         self.assertEqual(ref, res)
 
+    def test_tensor_attribute_error_in_try_except(self):
+        class M(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.scale = torch.nn.Parameter(torch.tensor([2.0]))
+
+            def forward(self, x):
+                try:
+                    return x.this_attribute_does_not_exist
+                except AttributeError:
+                    return x * self.scale
+
+        m = M()
+        opt_m = torch.compile(m, backend="eager")
+        x = torch.randn(4, 4)
+        ref = m(x)
+        res = opt_m(x)
+        self.assertEqual(ref, res)
+
     def test_raise_from_None(self):
         # Inspired from os.environ
         class MyMapping:
@@ -607,8 +650,8 @@ def fn():
         try:
             fn()
         except TypeError as e:
-            assert e.__cause__ is None
-            assert e.__suppress_context__ is True
+            assert e.__cause__ is None  # noqa: S101
+            assert e.__suppress_context__ is True  # noqa: S101
 
     @make_dynamo_test
     def test_raise_from_other(self):
@@ -621,8 +664,8 @@ def fn():
         try:
             fn()
         except TypeError as e:
-            assert isinstance(e.__cause__, ValueError)
-            assert e.__suppress_context__ is True
+            assert isinstance(e.__cause__, ValueError)  # noqa: S101
+            assert e.__suppress_context__ is True  # noqa: S101
 
     @make_dynamo_test
     def test_reraise_first_exc(self):
@@ -640,7 +683,7 @@ def fn():
             fn()
         except ZeroDivisionError:
             pass
-        assert sys.exc_info()[0] is None
+        assert sys.exc_info()[0] is None  # noqa: S101
 
     @make_dynamo_test
     def test_ensure_exception_is_active_after_try_except_block(self):
@@ -656,7 +699,7 @@ def test_ensure_exception_is_active_after_try_except_block(self):
                 raise
         except ZeroDivisionError:
             pass
-        assert sys.exc_info()[0] is None
+        assert sys.exc_info()[0] is None  # noqa: S101
 
     @make_dynamo_test
     def test_ensure_exception_is_active_inside_try_except_block(self):
@@ -668,11 +711,11 @@ def test_ensure_exception_is_active_inside_try_except_block(self):
                     try:
                         raise exc
                     except exc as e:
-                        assert isinstance(e.__context__, ZeroDivisionError)
+                        assert isinstance(e.__context__, ZeroDivisionError)  # noqa: S101
                 raise
         except ZeroDivisionError:
             pass
-        assert sys.exc_info()[0] is None
+        assert sys.exc_info()[0] is None  # noqa: S101
 
     @make_dynamo_test
     def test_handle_all_exceptions(self):
@@ -691,7 +734,7 @@ def cm():
             gen.throw(ValueError)
         except TypeError:
             pass
-        assert sys.exc_info()[0] is None
+        assert sys.exc_info()[0] is None  # noqa: S101
 
     @make_dynamo_test
     def test_reraise(self):
@@ -702,7 +745,7 @@ def test_reraise(self):
                 raise
         except ValueError:
             pass
-        assert sys.exc_info()[0] is None
+        assert sys.exc_info()[0] is None  # noqa: S101
 
     @make_dynamo_test
     def test_raise_finally_simple(self):
@@ -719,7 +762,7 @@ def fn():
             fn()
         except TypeError:
             pass
-        assert sys.exc_info()[0] is None
+        assert sys.exc_info()[0] is None  # noqa: S101
 
     def test_reconstruct___context__(self):
         @torch.compile(backend="eager", fullgraph=True)
@@ -856,9 +899,9 @@ def fn():
         try:
             fn()
         except TypeError as e:
-            assert isinstance(e.__context__, ex)
-            assert e.__cause__ is None
-            assert e.__suppress_context__ is True
+            assert isinstance(e.__context__, ex)  # noqa: S101
+            assert e.__cause__ is None  # noqa: S101
+            assert e.__suppress_context__ is True  # noqa: S101
 
     @parametrize(
         "ex",
@@ -880,13 +923,13 @@ def fn():
             fn()
         except TypeError as e:
             z = 1
-            assert e.args == (
+            assert e.args == (  # noqa: S101
                 "exception cause must be None or derive from BaseException",
             )
         except Exception:
             raise AssertionError from None
 
-        assert z == 1
+        assert z == 1  # noqa: S101
 
     def test_user_defined_exception_variable(self):
         def fn(t):
@@ -902,7 +945,7 @@ def fn(t):
                 if type(cls) is type:
                     t = t + 1
                 z = 2
-            assert z == 2
+            assert z == 2  # noqa: S101
             return t.sin()
 
         t = torch.randn(2)
@@ -920,7 +963,7 @@ def fn(t):
                 z = 1
             except CustomExceptionWithArgs:
                 z = 2
-            assert z == 2
+            assert z == 2  # noqa: S101
 
         t = torch.randn(2)
         fn(t)
@@ -932,14 +975,14 @@ def test_raise_set___context__(self):
         except TypeError as e:
             exc = e
 
-        assert exc.__context__ is None
+        assert exc.__context__ is None  # noqa: S101
 
         try:
             raise ValueError
         except ValueError as e:
             exc2 = e
 
-        assert exc2.__context__ is None
+        assert exc2.__context__ is None  # noqa: S101
 
     def test_exception_kwargs(self):
         @torch.compile(backend="eager", fullgraph=True)
@@ -970,7 +1013,7 @@ def test_context_manager_preserves_exception_stack(self):
         # doesn't suppress it, the error message should point to the original raise
         # location, not the context manager cleanup code.
         def g():
-            assert False  # noqa: B011
+            assert False  # noqa: B011, S101
 
         @torch.compile(backend="eager", fullgraph=True)
         def f(x):
@@ -1008,6 +1051,309 @@ def f(x):
         self.assertIn("in g", str(ctx.exception))
         self.assertIn('raise Exception("Invalid")', str(ctx.exception))
 
+    def test_str_repr_exception_no_args(self):
+        @torch.compile(backend="eager", fullgraph=True)
+        def fn(t):
+            try:
+                raise ValueError
+            except ValueError as e:
+                return t.sin(), str(e), repr(e)
+
+        t = torch.randn(2)
+        y, s, r = fn(t)
+        self.assertEqual(y, t.sin())
+        self.assertEqual(s, "")
+        self.assertEqual(r, "ValueError()")
+
+    def test_str_repr_exception_single_arg(self):
+        @torch.compile(backend="eager", fullgraph=True)
+        def fn(t):
+            try:
+                raise ValueError("test error")
+            except ValueError as e:
+                return t.sin(), str(e), repr(e)
+
+        t = torch.randn(2)
+        y, s, r = fn(t)
+        self.assertEqual(y, t.sin())
+        self.assertEqual(s, "test error")
+        self.assertEqual(r, "ValueError('test error')")
+
+    def test_str_repr_exception_multi_args(self):
+        @torch.compile(backend="eager", fullgraph=True)
+        def fn(t):
+            try:
+                raise ValueError("hello", 42)
+            except ValueError as e:
+                return t.sin(), str(e), repr(e)
+
+        t = torch.randn(2)
+        y, s, r = fn(t)
+        self.assertEqual(y, t.sin())
+        self.assertEqual(s, str(("hello", 42)))
+        self.assertEqual(r, "ValueError('hello', 42)")
+
+    def test_frozen_dataclass_setattr_raises(self):
+        @dataclasses.dataclass(frozen=True)
+        class TestDataClass:
+            x: int
+
+        @torch.compile(backend="eager", fullgraph=True)
+        def fn(t):
+            dc = TestDataClass(1)
+            try:
+                dc.x = 2
+            except dataclasses.FrozenInstanceError:
+                return t + 1
+            except Exception:
+                return t + 2
+            return t + dc.x
+
+        self.assertEqual(fn(torch.zeros(1)), 1)
+
+    def test_exception_traceback_access(self):
+        # Test that __traceback__ is accessible after raising/catching an exception
+        def fn(x):
+            try:
+                raise ValueError("oops")
+            except ValueError as e:
+                tb = e.__traceback__
+                if tb is not None:
+                    x = x + 1
+            return x
+
+        x = torch.randn(4)
+        ref = fn(x)
+        opt_fn = torch.compile(fn, backend="eager", fullgraph=True)
+        res = opt_fn(x)
+        self.assertEqual(ref, res)
+
+    def test_exception_traceback_tb_next(self):
+        # Test that tb_next can be accessed on a traceback
+        def fn(x):
+            try:
+                raise ValueError("oops")
+            except ValueError as e:
+                tb = e.__traceback__
+                if tb is not None:
+                    # tb_next is None for a single-frame traceback
+                    if tb.tb_next is None:
+                        x = x + 1
+            return x
+
+        x = torch.randn(4)
+        ref = fn(x)
+        opt_fn = torch.compile(fn, backend="eager", fullgraph=True)
+        res = opt_fn(x)
+        self.assertEqual(ref, res)
+
+    def test_exception_traceback_tb_lineno(self):
+        # Test that tb_lineno is accessible on a traceback
+        def fn(x):
+            try:
+                raise ValueError("oops")
+            except ValueError as e:
+                tb = e.__traceback__
+                if tb is not None and tb.tb_lineno > 0:
+                    x = x + 1
+            return x
+
+        x = torch.randn(4)
+        ref = fn(x)
+        opt_fn = torch.compile(fn, backend="eager", fullgraph=True)
+        res = opt_fn(x)
+        self.assertEqual(ref, res)
+
+    def test_exception_with_traceback_method(self):
+        # Test the with_traceback() method
+        def fn(x):
+            try:
+                raise ValueError("first")
+            except ValueError as e:
+                tb = e.__traceback__
+                try:
+                    raise RuntimeError("second").with_traceback(tb) from None
+                except RuntimeError as e2:
+                    if e2.__traceback__ is not None:
+                        x = x + 1
+            return x
+
+        x = torch.randn(4)
+        ref = fn(x)
+        opt_fn = torch.compile(fn, backend="eager", fullgraph=True)
+        res = opt_fn(x)
+        self.assertEqual(ref, res)
+
+    def test_exception_set_traceback(self):
+        # Test assigning __traceback__ on an exception
+        def fn(x):
+            try:
+                raise ValueError("first")
+            except ValueError as e:
+                tb = e.__traceback__
+                try:
+                    raise RuntimeError("second") from None
+                except RuntimeError as e2:
+                    e2.__traceback__ = tb
+                    if e2.__traceback__ is not None:
+                        x = x + 1
+            return x
+
+        x = torch.randn(4)
+        ref = fn(x)
+        opt_fn = torch.compile(fn, backend="eager", fullgraph=True)
+        res = opt_fn(x)
+        self.assertEqual(ref, res)
+
+    def test_exception_set_traceback_none(self):
+        # Test assigning None to __traceback__
+        def fn(x):
+            try:
+                raise ValueError("oops")
+            except ValueError as e:
+                e.__traceback__ = None
+                if e.__traceback__ is None:
+                    x = x + 1
+            return x
+
+        x = torch.randn(4)
+        ref = fn(x)
+        opt_fn = torch.compile(fn, backend="eager", fullgraph=True)
+        res = opt_fn(x)
+        self.assertEqual(ref, res)
+
+    def test_exception_traceback_tb_lasti_graph_break(self):
+        # Accessing tb_lasti should cause a graph break
+        def fn(x):
+            try:
+                raise ValueError("oops")
+            except ValueError as e:
+                tb = e.__traceback__
+                if tb is not None:
+                    _ = tb.tb_lasti
+                    x = x + 1
+            return x
+
+        x = torch.randn(4)
+        ref = fn(x)
+        # Should graph break but still produce correct results
+        opt_fn = torch.compile(fn, backend="eager")
+        res = opt_fn(x)
+        self.assertEqual(ref, res)
+
+    def test_exception_set_tb_next(self):
+        # Test setting tb_next on a traceback
+        def fn(x):
+            try:
+                raise ValueError("first")
+            except ValueError as e:
+                tb1 = e.__traceback__
+                try:
+                    raise RuntimeError("second") from None
+                except RuntimeError as e2:
+                    tb2 = e2.__traceback__
+                    if tb2 is not None and tb1 is not None:
+                        tb2.tb_next = None
+                        x = x + 1
+            return x
+
+        x = torch.randn(4)
+        ref = fn(x)
+        opt_fn = torch.compile(fn, backend="eager", fullgraph=True)
+        res = opt_fn(x)
+        self.assertEqual(ref, res)
+
+    def test_exception_traceback_chained(self):
+        # Test traceback chaining through multiple frames
+        def inner():
+            raise ValueError("inner")
+
+        def fn(x):
+            try:
+                inner()
+            except ValueError as e:
+                tb = e.__traceback__
+                if tb is not None:
+                    x = x + 1
+                    # Walk the traceback chain
+                    depth = 0
+                    curr = tb
+                    while curr is not None:
+                        depth += 1
+                        curr = curr.tb_next
+                    if depth > 0:
+                        x = x + 1
+            return x
+
+        x = torch.randn(4)
+        ref = fn(x)
+        opt_fn = torch.compile(fn, backend="eager", fullgraph=True)
+        res = opt_fn(x)
+        self.assertEqual(ref, res)
+
+    @parametrize(
+        "exc_type_1,exc_type_2",
+        [
+            (ValueError, TypeError),
+            (CustomException, ValueError),
+        ],
+        name_fn=lambda exc1, exc2: f"{exc1.__name__}_to_{exc2.__name__}",
+    )
+    def test_exception_set_context(self, exc_type_1, exc_type_2):
+        # Test explicitly assigning to __context__ attribute (reaches ExceptionVariable.__context__ assignment)
+        @torch.compile(backend="eager", fullgraph=True)
+        def fn(t):
+            exc1 = exc_type_1("first")
+            exc2 = exc_type_2("second")
+
+            # This explicitly sets __context__ via call_setattr
+            exc2.__context__ = exc1
+
+            # Verify it was set correctly
+            if exc2.__context__ is exc1:
+                return t.sin()
+            else:
+                return t.cos()
+
+        t = torch.randn(2)
+        ref_result = t.sin()
+        result = fn(t)
+        self.assertEqual(result, ref_result)
+
+    @parametrize(
+        "exc_type_1,exc_type_2,exc_type_3",
+        [
+            (ValueError, TypeError, RuntimeError),
+            (CustomException, ValueError, TypeError),
+        ],
+        name_fn=lambda exc1, exc2, exc3: (
+            f"{exc1.__name__}_chain_{exc2.__name__}_{exc3.__name__}"
+        ),
+    )
+    def test_exception_context_chain(self, exc_type_1, exc_type_2, exc_type_3):
+        # Test chaining contexts through multiple exceptions
+        @torch.compile(backend="eager", fullgraph=True)
+        def fn(t):
+            exc1 = exc_type_1("first")
+            exc2 = exc_type_2("second")
+            exc3 = exc_type_3("third")
+
+            exc2.__context__ = exc1
+            exc3.__context__ = exc2
+
+            # Verify the chain
+            if isinstance(exc3.__context__, exc_type_2) and isinstance(
+                exc3.__context__.__context__, exc_type_1
+            ):
+                return t.sin()
+            else:
+                return t.cos()
+
+        t = torch.randn(2)
+        ref_result = t.sin()
+        result = fn(t)
+        self.assertEqual(result, ref_result)
+
 
 instantiate_parametrized_tests(ExceptionTests)
 
diff --git a/test/dynamo/test_export.py b/test/dynamo/test_export.py
index cd699e9d6fd92..f17a029fb823c 100644
--- a/test/dynamo/test_export.py
+++ b/test/dynamo/test_export.py
@@ -9,6 +9,9 @@
 import inspect
 import io
 import operator
+import os
+import subprocess
+import sys
 import unittest
 from collections.abc import Sequence
 from enum import Enum
@@ -127,6 +130,15 @@ def forward(self, x, y):
     return pytree.tree_unflatten([x], self._out_spec)""",
         )
 
+    def test_export_empty_graph_no_error(self):
+        def func(x):
+            return len(x)
+
+        exported = torch._dynamo.export(func)(torch.randn(5))
+        out_graph = exported[0]
+        result = out_graph(torch.randn(5))
+        self.assertEqual(result, 5)
+
     def test_no_tensor_computation_2(self):
         inp = torch.randn(3)
         inp2 = 2
@@ -384,7 +396,9 @@ def forward(self, x1, x2):
             _error_on_data_dependent_ops=True,
         )(*[x1, x2])
         ep = torch.export.export(fx_model, (x1, x2))
-        res = torch.compile(ep.module(), dynamic=True, fullgraph=True)(x1, x2)
+        res = torch.compile(ep.module(), backend="eager", dynamic=True, fullgraph=True)(
+            x1, x2
+        )
         self.assertTrue(torch._dynamo.utils.same(res, M()(x1, x2)))
 
     def test_dupes(self):
@@ -2873,7 +2887,7 @@ def forward(self, x):
 
     def test_list_contains(self):
         def func(x):
-            assert x.size(-1) in [4, 5, 6], "bad"
+            assert x.size(-1) in [4, 5, 6], "bad"  # noqa: S101
             return x + x
 
         inps = (torch.randn(1, 5),)
@@ -2891,8 +2905,8 @@ def func(x):
 
     def test_list_not_contains(self):
         def func(x):
-            assert x.size(0) not in [4, 5, 6], "bad1"
-            assert "monkey" not in ["cow", "pig"], "bad2"
+            assert x.size(0) not in [4, 5, 6], "bad1"  # noqa: S101
+            assert "monkey" not in ["cow", "pig"], "bad2"  # noqa: S101
             return x + x
 
         inps = (torch.randn(1, 5),)
@@ -3027,7 +3041,7 @@ def f(x):
     @config.patch(assume_static_by_default=False)
     def test_export_persist_assert(self):
         def f(x):
-            assert x[0].sum() > 4, "Shape must be more than 4"
+            assert x[0].sum() > 4, "Shape must be more than 4"  # noqa: S101
             return x.cos() + x.sin()
 
         gm, _ = torch._dynamo.export(f, aten_graph=True, tracing_mode="symbolic")(
@@ -4364,7 +4378,7 @@ def bad_fn(x):
 
         expected = [
             """x = torch.sin(l_x_)""",
-            """cos = torch.cos(l_stack0_)""",
+            """cos = torch.cos(l_nested_frame_values_0_1_)""",
         ]
 
         def test_backend(gm: torch.fx.GraphModule, example_inputs):
@@ -4570,6 +4584,32 @@ def forward(self, x):
         self.assertEqual(ref, res)
 
 
+class ExportTestsSubprocess(torch._dynamo.test_case.TestCase):
+    def test_strict_export_under_pythonoptimize(self):
+        env = dict(os.environ)
+        env["PYTHONOPTIMIZE"] = "1"
+        code = """\
+import torch
+model = torch.nn.Linear(2, 3)
+example_input = torch.randn(1, 2)
+ep = torch.export.export(model, args=(example_input,), strict=True)
+out_export = ep.module()(example_input)
+out_orig = model(example_input)
+torch.testing.assert_close(out_export, out_orig)
+"""
+        result = subprocess.run(
+            [sys.executable, "-c", code],
+            env=env,
+            capture_output=True,
+            text=True,
+        )
+        self.assertEqual(
+            result.returncode,
+            0,
+            msg=f"strict export under PYTHONOPTIMIZE=1 failed: stdout={result.stdout!r} stderr={result.stderr!r}",
+        )
+
+
 class ExportTestsDevice(torch._dynamo.test_case.TestCase):
     def test_export_with_parameters(self, device):
         class MyModule(torch.nn.Module):
diff --git a/test/dynamo/test_fake_distributed.py b/test/dynamo/test_fake_distributed.py
index 4dd8fd8ba3279..7e5607e08aab0 100644
--- a/test/dynamo/test_fake_distributed.py
+++ b/test/dynamo/test_fake_distributed.py
@@ -4,7 +4,11 @@
 import torch
 import torch.distributed as dist
 from torch._dynamo.test_case import TestCase as DynamoTestCase
-from torch._dynamo.testing import AotEagerAndRecordGraphs, normalize_gm
+from torch._dynamo.testing import (
+    AotEagerAndRecordGraphs,
+    EagerAndRecordGraphs,
+    normalize_gm,
+)
 from torch.testing._internal.common_utils import instantiate_parametrized_tests
 
 
@@ -23,12 +27,14 @@ def normalize_graph(gm):
 @skipIf(not dist.is_available(), "requires distributed")
 class TestFakeDistributed(DynamoTestCase):
     def setUp(self):
+        super().setUp()
         # Use FakeProcessGroup to run tests on a single process
         dist.init_process_group(backend="fake", rank=0, world_size=2)
         self.local_rank = 0
         self.world_size = 2
 
     def tearDown(self):
+        super().tearDown()
         dist.destroy_process_group()
 
     def test_all_to_all_single_autograd(self):
@@ -96,7 +102,6 @@ def forward(self, primals_1: "Sym(u0)", primals_2: "Sym(u1)", primals_3: "Sym(u2
         _assert_scalar_1 = torch.ops.aten._assert_scalar.default(ge_1, "Runtime assertion failed for expression u1 >= 0 on node 'ge_1'");  ge_1 = _assert_scalar_1 = None
         ge_2: "Sym(u2 >= 0)" = primals_3 >= 0
         _assert_scalar_2 = torch.ops.aten._assert_scalar.default(ge_2, "Runtime assertion failed for expression u2 >= 0 on node 'ge_2'");  ge_2 = _assert_scalar_2 = None
-
         floordiv: "Sym((u0//2))" = primals_1 // 2
 
         all_to_all_single: "f32[2*((u0//2)), u1, u2]" = torch.ops._c10d_functional.all_to_all_single.default(primals_4, [floordiv, floordiv], [floordiv, floordiv], '0');  primals_4 = None
@@ -155,9 +160,250 @@ def fn(x):
         res = fn(x)
         self.assertEqual(res, (x + 1, [0]))
 
+    @torch._dynamo.config.patch(nested_graph_breaks=False)
+    def test_device_mesh_init_skip_after_graph_break(self):
+        device_mesh = init_device_mesh(
+            device_type="cpu",
+            mesh_shape=(1, self.world_size),
+            mesh_dim_names=("dp", "tp"),
+        )
+
+        @torch.compile(backend="eager")
+        def fn(x):
+            # Graph break so the subsequent DeviceMesh construction happens
+            # at runtime, where eval_frame intercepts DeviceMesh.__init__
+            # as a top-level frame.
+            torch._dynamo.graph_break()
+            layout = device_mesh._get_slice_mesh_layout(("tp",))
+            sub = device_mesh._create_sub_mesh(layout, ("tp",))
+            return x + sub.size()
+
+        x = torch.ones(10)
+        res = fn(x)
+        self.assertEqual(res, x + self.world_size)
+
 
 instantiate_parametrized_tests(TestFakeDistributed)
 
+
+@skipIf(not dist.is_available(), "requires distributed")
+class TestFakeDistributedP2P(DynamoTestCase):
+    def setUp(self):
+        dist.init_process_group(backend="fake", rank=0, world_size=2)
+
+    def tearDown(self):
+        dist.destroy_process_group()
+
+    @torch._dynamo.config.patch(enable_p2p_compilation=True)
+    def test_compiled_isend_graph(self):
+        backend = EagerAndRecordGraphs()
+
+        @torch.compile(fullgraph=True, backend=backend)
+        def fn(tensor):
+            req = dist.isend(tensor, 1)
+            req.wait()
+
+        tensor = torch.ones(10)
+        fn(tensor)
+        self.assertEqual(len(backend.graphs), 1)
+        self.assertExpectedInline(
+            normalize_graph(backend.graphs[0]),
+            """\
+class GraphModule(torch.nn.Module):
+    def forward(self, L_tensor_: "f32[10]"):
+        l_tensor_ = L_tensor_
+
+        tensor: "f32[0]" = torch.ops._c10d_functional.isend(l_tensor_, 1, 0, '0');  l_tensor_ = None
+
+        wait_tensor: "f32[0]" = torch.distributed._functional_collectives.wait_tensor(tensor);  tensor = wait_tensor = None
+        return ()
+""",
+        )
+
+    @torch._dynamo.config.patch(enable_p2p_compilation=True)
+    def test_compiled_fire_and_forget_isend_graph(self):
+        backend = EagerAndRecordGraphs()
+
+        @torch.compile(fullgraph=True, backend=backend)
+        def fn(tensor):
+            dist.isend(tensor, 1)
+            return tensor
+
+        tensor = torch.ones(4)
+        fn(tensor)
+        self.assertEqual(len(backend.graphs), 1)
+        self.assertExpectedInline(
+            normalize_graph(backend.graphs[0]),
+            """\
+class GraphModule(torch.nn.Module):
+    def forward(self, L_tensor_: "f32[4]"):
+        l_tensor_ = L_tensor_
+
+        tensor: "f32[0]" = torch.ops._c10d_functional.isend(l_tensor_, 1, 0, '0');  tensor = None
+        return (l_tensor_,)
+""",
+        )
+
+    @torch._dynamo.config.patch(enable_p2p_compilation=True)
+    def test_compiled_irecv_graph(self):
+        backend = EagerAndRecordGraphs()
+
+        @torch.compile(fullgraph=True, backend=backend)
+        def fn(tensor):
+            req = dist.irecv(tensor, 1)
+            req.wait()
+
+        tensor = torch.zeros(10)
+        fn(tensor)
+        self.assertEqual(len(backend.graphs), 1)
+        self.assertExpectedInline(
+            normalize_graph(backend.graphs[0]),
+            """\
+class GraphModule(torch.nn.Module):
+    def forward(self, L_tensor_: "f32[10]"):
+        l_tensor_ = L_tensor_
+
+        tensor: "f32[10]" = torch.ops._c10d_functional.irecv(l_tensor_, 1, 0, '0');  l_tensor_ = None
+
+        req: "f32[10]" = torch.ops._c10d_functional.wait_tensor(tensor);  tensor = None
+
+        wait_tensor_1: "f32[10]" = torch.distributed._functional_collectives.wait_tensor(req);  req = wait_tensor_1 = None
+        return ()
+""",
+        )
+
+    @torch._dynamo.config.patch(enable_p2p_compilation=True)
+    def test_compiled_batch_isend_irecv_mixed_graph(self):
+        backend = EagerAndRecordGraphs()
+
+        @torch.compile(fullgraph=True, backend=backend)
+        def fn(send_tensor, recv_tensor):
+            ops = [
+                dist.P2POp(dist.isend, send_tensor, 1),
+                dist.P2POp(dist.irecv, recv_tensor, 1),
+            ]
+            work = dist.batch_isend_irecv(ops)
+            for w in work:
+                w.wait()
+
+        send_tensor = torch.ones(10)
+        recv_tensor = torch.zeros(10)
+        fn(send_tensor, recv_tensor)
+        self.assertEqual(len(backend.graphs), 1)
+        self.assertExpectedInline(
+            normalize_graph(backend.graphs[0]),
+            (
+                """\
+class GraphModule(torch.nn.Module):
+    def forward(self, L_send_tensor_: "f32[10]", L_recv_tensor_: "f32[10]"):
+        l_send_tensor_ = L_send_tensor_
+        l_recv_tensor_ = L_recv_tensor_
+
+"""
+                "        batch_p2p_ops = torch.ops._c10d_functional.batch_p2p_ops("
+                "['isend', 'irecv'], [1, 1], [0, 0], [l_send_tensor_, "
+                "l_recv_tensor_], '0');  l_send_tensor_ = l_recv_tensor_ = None\n"
+                """\
+        t: "f32[0]" = batch_p2p_ops[0]
+        t_1: "f32[10]" = batch_p2p_ops[1];  batch_p2p_ops = None
+
+        w: "f32[10]" = torch.ops._c10d_functional.wait_tensor(t_1);  t_1 = None
+
+        wait_tensor_1: "f32[0]" = torch.distributed._functional_collectives.wait_tensor(t);  t = wait_tensor_1 = None
+        wait_tensor_2: "f32[10]" = torch.distributed._functional_collectives.wait_tensor(w);  w = wait_tensor_2 = None
+        return ()
+"""
+            ),
+        )
+
+    @torch._dynamo.config.patch(enable_p2p_compilation=True)
+    def test_compiled_p2p_interleave_graph(self):
+        backend = EagerAndRecordGraphs()
+        nxt, prv = 1, 1  # rank=0 in world_size=2
+
+        @torch.compile(fullgraph=True, backend=backend)
+        def fn(x0, x1, y0, y1):
+            work = dist.batch_isend_irecv(
+                [
+                    dist.P2POp(dist.isend, x0, nxt),
+                    dist.P2POp(dist.irecv, y0, prv),
+                ]
+            )
+            t0 = x0 * 2 + 1
+            for w in work:
+                w.wait()
+            a = y0 + t0
+            work = dist.batch_isend_irecv(
+                [
+                    dist.P2POp(dist.isend, a, nxt),
+                    dist.P2POp(dist.irecv, y1, prv),
+                ]
+            )
+            t1 = a * 1.000244140625
+            for w in work:
+                w.wait()
+            return y1 + t1
+
+        M, N = 64, 64
+        x0 = torch.ones(M, N)
+        x1 = torch.ones(M, N)
+        y0 = torch.zeros(M, N)
+        y1 = torch.zeros(M, N)
+        fn(x0, x1, y0, y1)
+        self.assertEqual(len(backend.graphs), 1)
+        self.assertExpectedInline(
+            normalize_graph(backend.graphs[0]),
+            """\
+class GraphModule(torch.nn.Module):
+    def forward(self, L_x0_: "f32[64, 64]", L_y0_: "f32[64, 64]", L_y1_: "f32[64, 64]"):
+        l_x0_ = L_x0_
+        l_y0_ = L_y0_
+        l_y1_ = L_y1_
+
+        batch_p2p_ops = torch.ops._c10d_functional.batch_p2p_ops(['isend', 'irecv'], [1, 1], [0, 0], [l_x0_, l_y0_], '0')
+        t: "f32[0]" = batch_p2p_ops[0]
+        t_1: "f32[64, 64]" = batch_p2p_ops[1];  batch_p2p_ops = None
+
+        w: "f32[64, 64]" = torch.ops._c10d_functional.wait_tensor(t_1);  t_1 = None
+
+        mul: "f32[64, 64]" = l_x0_ * 2;  l_x0_ = None
+        t0: "f32[64, 64]" = mul + 1;  mul = None
+
+        wait_tensor_1: "f32[0]" = torch.distributed._functional_collectives.wait_tensor(t);  t = wait_tensor_1 = None
+        wait_tensor_2: "f32[64, 64]" = torch.distributed._functional_collectives.wait_tensor(w);  w = wait_tensor_2 = None
+
+        a: "f32[64, 64]" = l_y0_ + t0;  l_y0_ = t0 = None
+
+        batch_p2p_ops_1 = torch.ops._c10d_functional.batch_p2p_ops(['isend', 'irecv'], [1, 1], [0, 0], [a, l_y1_], '0')
+        t_2: "f32[0]" = batch_p2p_ops_1[0]
+        t_3: "f32[64, 64]" = batch_p2p_ops_1[1];  batch_p2p_ops_1 = None
+
+        w_1: "f32[64, 64]" = torch.ops._c10d_functional.wait_tensor(t_3);  t_3 = None
+
+        t1: "f32[64, 64]" = a * 1.000244140625;  a = None
+
+        wait_tensor_4: "f32[0]" = torch.distributed._functional_collectives.wait_tensor(t_2);  t_2 = wait_tensor_4 = None
+        wait_tensor_5: "f32[64, 64]" = torch.distributed._functional_collectives.wait_tensor(w_1);  w_1 = wait_tensor_5 = None
+
+        add_2: "f32[64, 64]" = l_y1_ + t1;  l_y1_ = t1 = None
+        return (add_2,)
+""",
+        )
+
+    @torch._dynamo.config.patch(enable_p2p_compilation=True)
+    def test_mutating_p2p_op_graph_breaks(self):
+        @torch.compile(backend="eager", fullgraph=True)
+        def fn(tensor):
+            op = dist.P2POp(dist.isend, tensor, 1)
+            op.tag = 1
+            dist.batch_isend_irecv([op])
+
+        with self.assertRaisesRegex(torch._dynamo.exc.Unsupported, "P2POp"):
+            fn(torch.ones(4))
+
+
+instantiate_parametrized_tests(TestFakeDistributedP2P)
+
 if __name__ == "__main__":
     from torch._dynamo.test_case import run_tests
 
diff --git a/test/dynamo/test_flat_apply.py b/test/dynamo/test_flat_apply.py
index 536e6d8f249e2..833a08fda3596 100644
--- a/test/dynamo/test_flat_apply.py
+++ b/test/dynamo/test_flat_apply.py
@@ -55,7 +55,7 @@ class InputData:
     values: Tensor
 
 
-torch.utils._pytree.register_dataclass(InputData)
+pytree.register_dataclass(InputData)
 
 
 @dataclass
@@ -73,7 +73,7 @@ class OutputData:
     result2: Tensor
 
 
-torch.utils._pytree.register_dataclass(OutputData)
+pytree.register_dataclass(OutputData)
 
 
 @dataclass
@@ -193,9 +193,9 @@ def forward(self, L_x_: "f32[10]", L_y_: "f32[10]"):
 
         t: "f32[10]" = l_x_ + l_y_
 
-        trace_point_tensor_spec : torch.utils._pytree.TreeSpec = self.trace_point_tensor_spec
+        trace_point_tensor_callable : torch._higher_order_ops.invoke_leaf_function._LeafCallable = self.trace_point_tensor_callable
         trace_point_tensor_input_spec : torch.utils._pytree.TreeSpec = self.trace_point_tensor_input_spec
-        flat_apply_capture = torch__dynamo_variables_torch_flat_apply_capture(trace_point_tensor_spec, trace_point_tensor_input_spec, l_x_, l_y_, t);  trace_point_tensor_spec = trace_point_tensor_input_spec = l_x_ = l_y_ = t = None
+        flat_apply_capture = torch__dynamo_variables_torch_flat_apply_capture(trace_point_tensor_callable, trace_point_tensor_input_spec, l_x_, l_y_, t);  trace_point_tensor_callable = trace_point_tensor_input_spec = l_x_ = l_y_ = t = None
         res: "f32[10]" = flat_apply_capture[0];  flat_apply_capture = None
         return (res,)
 """,  # NOQA: B950
@@ -269,9 +269,9 @@ def forward(self, L_i_values: "f32[4, 4]"):
         x: "f32[4, 4]" = torch.sin(l_i_values)
 
         # code: y, z_result1, z_result2 = gn(i.count, i.values)
-        gn_spec : torch.utils._pytree.TreeSpec = self.gn_spec
+        gn_callable : torch._higher_order_ops.invoke_leaf_function._LeafCallable = self.gn_callable
         gn_input_spec : torch.utils._pytree.TreeSpec = self.gn_input_spec
-        flat_apply_capture = torch__dynamo_variables_torch_flat_apply_capture(gn_spec, gn_input_spec, 5, l_i_values);  gn_spec = gn_input_spec = l_i_values = None
+        flat_apply_capture = torch__dynamo_variables_torch_flat_apply_capture(gn_callable, gn_input_spec, 5, l_i_values);  gn_callable = gn_input_spec = l_i_values = None
         y: "f32[4, 4]" = flat_apply_capture[0]
         z_result1: "f32[4, 4]" = flat_apply_capture[1]
         z_result2: "f32[4, 4]" = flat_apply_capture[2];  flat_apply_capture = None
@@ -322,9 +322,9 @@ def forward(self, L_i_values: "f32[4, 4]"):
         x: "f32[4, 4]" = torch.sin(l_i_values)
 
         # code: y, z_result1, z_result2 = gn(i)
-        gn_spec : torch.utils._pytree.TreeSpec = self.gn_spec
+        gn_callable : torch._higher_order_ops.invoke_leaf_function._LeafCallable = self.gn_callable
         gn_input_spec : torch.utils._pytree.TreeSpec = self.gn_input_spec
-        flat_apply_capture = torch__dynamo_variables_torch_flat_apply_capture(gn_spec, gn_input_spec, 5, l_i_values);  gn_spec = gn_input_spec = l_i_values = None
+        flat_apply_capture = torch__dynamo_variables_torch_flat_apply_capture(gn_callable, gn_input_spec, 5, l_i_values);  gn_callable = gn_input_spec = l_i_values = None
         y: "f32[4, 4]" = flat_apply_capture[0]
         z_result1: "f32[4, 4]" = flat_apply_capture[1]
         z_result2: "f32[4, 4]" = flat_apply_capture[2];  flat_apply_capture = None
@@ -404,9 +404,9 @@ def forward(self, L_i_values: "f32[4, 4]"):
         x: "f32[4, 4]" = torch.sin(l_i_values)
 
         # code: y, z = gn(i.count, i.values)
-        gn_spec : torch.utils._pytree.TreeSpec = self.gn_spec
+        gn_callable : torch._higher_order_ops.invoke_leaf_function._LeafCallable = self.gn_callable
         gn_input_spec : torch.utils._pytree.TreeSpec = self.gn_input_spec
-        flat_apply_capture = torch__dynamo_variables_torch_flat_apply_capture(gn_spec, gn_input_spec, 5, l_i_values);  gn_spec = gn_input_spec = l_i_values = None
+        flat_apply_capture = torch__dynamo_variables_torch_flat_apply_capture(gn_callable, gn_input_spec, 5, l_i_values);  gn_callable = gn_input_spec = l_i_values = None
         y: "f32[4, 4]" = flat_apply_capture[0]
         value: "f32[4, 4]" = flat_apply_capture[1]
         value_1: "f32[4, 4]" = flat_apply_capture[2];  flat_apply_capture = None
diff --git a/test/dynamo/test_functions.py b/test/dynamo/test_functions.py
index e150b462675e0..48e08b057b920 100644
--- a/test/dynamo/test_functions.py
+++ b/test/dynamo/test_functions.py
@@ -31,11 +31,10 @@
     EagerAndRecordGraphs,
     normalize_gm,
 )
-from torch._dynamo.utils import ifdynstaticdefault, range_iterator, same
+from torch._dynamo.utils import counters, ifdynstaticdefault, range_iterator, same
 from torch._dynamo.variables import ConstantVariable, SkipFunctionVariable
 from torch._dynamo.variables.lists import RangeVariable
 from torch.nn import functional as F
-from torch.testing._internal.common_cuda import TEST_MULTIGPU
 from torch.testing._internal.common_utils import (
     instantiate_parametrized_tests,
     parametrize,
@@ -153,7 +152,7 @@ def inline_script_if_tracing_fn_with_default_args(x, y, c=1.2):
     return torch.cos(x * y) + c
 
 
-class FunctionTests(torch._dynamo.test_case.TestCaseWithNestedGraphBreaks):
+class FunctionTests(torch._dynamo.test_case.TestCase):
     @make_test
     def test_inline_jit_annotations(x):
         x = inline_script_if_tracing(x)
@@ -256,7 +255,7 @@ def test_foreach_norm_inf_dynamic(self):
         def fn(tensors):
             return torch._foreach_norm(tensors, float("inf"))
 
-        fn_opt = torch.compile(fn, fullgraph=True, dynamic=True)
+        fn_opt = torch.compile(fn, fullgraph=True, dynamic=True, backend="eager")
 
         tensors = [torch.randn(10), torch.randn(20, 30)]
 
@@ -378,6 +377,91 @@ def fn(a):
         self.assertEqual(next(i2), next(it2))
         self.assertEqual(a, b)
 
+    def test_itertools_islice_basic_ops(self):
+        # Test cases taken from the CPython test TestBasicOps.test_islice. That test has a lot of
+        # cases that we can't realistically support, whence we copy the sensible cases here.
+        def fn():
+            for args in [  # islice(args) should agree with range(args)
+                (10, 20, 3),
+                (10, 3, 20),
+                (10, 20),
+                (10, 10),
+                (10, 3),
+                (20,),
+            ]:
+                self.assertEqual(
+                    list(itertools.islice(range(100), *args)), list(range(*args))
+                )
+
+            for args, tgtargs in [  # Stop when seqn is exhausted
+                ((10, 110, 3), ((10, 100, 3))),
+                ((10, 110), ((10, 100))),
+                ((110,), (100,)),
+            ]:
+                self.assertEqual(
+                    list(itertools.islice(range(100), *args)), list(range(*tgtargs))
+                )
+
+            # Test stop=None
+            self.assertEqual(list(itertools.islice(range(10), None)), list(range(10)))
+            self.assertEqual(
+                list(itertools.islice(range(10), None, None)), list(range(10))
+            )
+            self.assertEqual(
+                list(itertools.islice(range(10), None, None, None)), list(range(10))
+            )
+            self.assertEqual(
+                list(itertools.islice(range(10), 2, None)), list(range(2, 10))
+            )
+            self.assertEqual(
+                list(itertools.islice(range(10), 1, None, 2)), list(range(1, 10, 2))
+            )
+
+            # Test number of items consumed     SF #1171417
+            it = iter(range(10))
+            self.assertEqual(list(itertools.islice(it, 3)), list(range(3)))
+            self.assertEqual(list(it), list(range(3, 10)))
+
+            it = iter(range(10))
+            self.assertEqual(list(itertools.islice(it, 3, 3)), [])
+            self.assertEqual(list(it), list(range(3, 10)))
+
+            # Issue #10323:  Less islice in a predictable state
+            c = itertools.count()
+            self.assertEqual(list(itertools.islice(c, 1, 3, 50)), [1])
+            self.assertEqual(next(c), 3)
+
+        opt_fn = torch.compile(fn, backend="eager", fullgraph=True)
+        opt_fn()
+
+    @unittest.expectedFailure
+    def test_itertools_islice_intlike(self):
+        # CPython issue #30537: islice can accept integer-like objects as arguments.
+        class IntLike:
+            def __init__(self, val):
+                self.val = val
+
+            def __index__(self):
+                return self.val
+
+        def fn():
+            self.assertEqual(
+                list(itertools.islice(range(100), IntLike(10))), list(range(10))
+            )
+            self.assertEqual(
+                list(itertools.islice(range(100), IntLike(10), IntLike(50))),
+                list(range(10, 50)),
+            )
+            self.assertEqual(
+                list(
+                    itertools.islice(range(100), IntLike(10), IntLike(50), IntLike(5))
+                ),
+                list(range(10, 50, 5)),
+            )
+
+        opt_fn = torch.compile(fn, backend="eager", fullgraph=True)
+        opt_fn()
+
     @make_test
     def test_obj_eq(a, b):
         v = a + b
@@ -604,9 +688,9 @@ def test_tuple_map(a, b):
     def test_size_tuple_add(self):
         def fn():
             size = torch.Size([])
-            assert isinstance(size + size, torch.Size)
-            assert isinstance(size + (), tuple)
-            assert isinstance(size + (), torch.Size)
+            assert isinstance(size + size, torch.Size)  # noqa: S101
+            assert isinstance(size + (), tuple)  # noqa: S101
+            assert isinstance(size + (), torch.Size)  # noqa: S101
 
         fn()
         compiled_fn = torch.compile(fn, backend="eager", fullgraph=True)
@@ -1277,8 +1361,8 @@ def test_del(a, b):
     @make_test
     def test_chunks1(x):
         chunk_size = 5
-        assert x.shape[0] % chunk_size == 0
-        assert x.shape[0] // chunk_size == 2
+        assert x.shape[0] % chunk_size == 0  # noqa: S101
+        assert x.shape[0] // chunk_size == 2  # noqa: S101
         return x[:chunk_size] - x[chunk_size:]
 
     @make_test
@@ -1420,7 +1504,7 @@ def test_startswith(a, b):
     @make_test
     def test_dict_ops(a, b):
         tmp = {"a": a + 1, "b": b + 2}
-        assert tmp.get("zzz") is None
+        assert tmp.get("zzz") is None  # noqa: S101
         v = tmp.pop("b") + tmp.get("a") + tmp.get("missing", 3) + tmp.pop("missing", 4)
         tmp.update({"d": 3})
         tmp["c"] = v + tmp["d"]
@@ -2354,13 +2438,13 @@ def test_fstrings6(x):
     def test_tensor_new_with_size(x):
         y = torch.rand(5, 8)
         z = x.new(y.size())
-        assert z.size() == y.size()
+        assert z.size() == y.size()  # noqa: S101
 
     @make_test
     def test_tensor_new_with_shape(x):
         y = torch.rand(5, 8)
         z = x.new(y.shape)
-        assert z.size() == y.size()
+        assert z.size() == y.size()  # noqa: S101
 
     @make_test
     def test_jit_annotate(x):
@@ -2405,12 +2489,16 @@ def func(x):
         cnt = torch._dynamo.testing.CompileCounter()
         cfunc = torch._dynamo.optimize_assert(cnt, dynamic=dynamic)(func)
 
-        assert cnt.frame_count == 0
+        if cnt.frame_count != 0:
+            raise AssertionError(f"Expected frame_count 0, got {cnt.frame_count}")
         for i, x in enumerate(data):
             expected = func(x)
             output = cfunc(x)
             self.assertTrue(same(output, expected))
-            assert cnt.frame_count == expected_frame_counts[i]
+            if cnt.frame_count != expected_frame_counts[i]:
+                raise AssertionError(
+                    f"Expected frame_count {expected_frame_counts[i]}, got {cnt.frame_count}"
+                )
 
     @make_test
     def test_list_slice_assignment(x):
@@ -2478,12 +2566,14 @@ def func(x, a):
         cnt = torch._dynamo.testing.CompileCounter()
         cfunc = torch._dynamo.optimize_assert(cnt)(func)
 
-        assert cnt.frame_count == 0
+        if cnt.frame_count != 0:
+            raise AssertionError(f"Expected frame_count 0, got {cnt.frame_count}")
         x = torch.rand(10)
         expected = func(x, 12)
         output = cfunc(x, 12)
         self.assertTrue(same(output, expected))
-        assert cnt.frame_count == 1
+        if cnt.frame_count != 1:
+            raise AssertionError(f"Expected frame_count 1, got {cnt.frame_count}")
 
     @unittest.skipIf(sys.version_info < (3, 13), "math.fma introduced in python 3.13")
     def test_math_fma(self):
@@ -2494,24 +2584,28 @@ def fma_func(a, b, c):
         cnt = torch._dynamo.testing.CompileCounter()
         cfma_scalars = torch._dynamo.optimize_assert(cnt)(fma_func)
 
-        assert cnt.frame_count == 0
+        if cnt.frame_count != 0:
+            raise AssertionError(f"Expected frame_count 0, got {cnt.frame_count}")
         expected = fma_func(2.0, 3.0, 4.0)
         output = cfma_scalars(2.0, 3.0, 4.0)
         self.assertEqual(output, expected)
-        assert cnt.frame_count == 0
+        if cnt.frame_count != 0:
+            raise AssertionError(f"Expected frame_count 0, got {cnt.frame_count}")
 
         # Test with tensors (Inductor path)
         cnt2 = torch._dynamo.testing.CompileCounter()
         cfma_tensors = torch._dynamo.optimize_assert(cnt2)(fma_func)
 
-        assert cnt2.frame_count == 0
+        if cnt2.frame_count != 0:
+            raise AssertionError(f"Expected frame_count 0, got {cnt2.frame_count}")
         x = torch.tensor(2.0)
         y = torch.tensor(3.0)
         z = torch.tensor(4.0)
         expected_tensors = x * y + z
         output_tensors = cfma_tensors(x, y, z)
         torch.testing.assert_close(output_tensors, expected_tensors)
-        assert cnt2.frame_count == 1
+        if cnt2.frame_count != 1:
+            raise AssertionError(f"Expected frame_count 1, got {cnt2.frame_count}")
 
     @make_test
     def test_numpy_meshgrid(x, y):
@@ -2641,8 +2735,8 @@ def test_list_expand_lhs(x):
     def test_in_not_in(x):
         mylist = [1, 2, 3, 4, 5, x]
         myotherlist = [1, 2, 3, 4, 5]
-        assert 3 in mylist
-        assert 6 not in myotherlist
+        assert 3 in mylist  # noqa: S101
+        assert 6 not in myotherlist  # noqa: S101
         return sum(mylist)
 
     @make_test
@@ -2651,7 +2745,7 @@ def test_is(x, y):
         try:
             raise exc
         except Exception as e:
-            assert e is exc
+            assert e is exc  # noqa: S101
             return x + y
 
     @make_test
@@ -2661,7 +2755,7 @@ def test_is_not(x, y):
         try:
             raise exc
         except Exception as e:
-            assert e is not exc1
+            assert e is not exc1  # noqa: S101
             return x + y
 
     @make_test
@@ -2795,7 +2889,7 @@ def fn(inputs):
         input2 = torch.arange(2, dtype=torch.bfloat16).requires_grad_(True)
         inputs = [input1, input2]
 
-        opt_fn = torch.compile(fullgraph=True)(fn)
+        opt_fn = torch.compile(fullgraph=True, backend="eager")(fn)
         self.assertEqual(opt_fn(inputs), fn(inputs))
 
     def test_filter_fallback(self):
@@ -2809,13 +2903,13 @@ def fn(inputs):
         input2 = torch.arange(2, dtype=torch.bfloat16)
         inputs = [input1, input2]
 
-        opt_fn = torch.compile()(fn)
+        opt_fn = torch.compile(backend="eager")(fn)
         self.assertEqual(opt_fn(inputs), fn(inputs))
 
         torch._dynamo.reset()
 
         with self.assertRaises(torch._dynamo.exc.Unsupported):
-            opt_fn = torch.compile(fullgraph=True)(fn)
+            opt_fn = torch.compile(fullgraph=True, backend="eager")(fn)
             opt_fn(inputs)
 
     def test_filter_infinite_iterator(self):
@@ -2913,7 +3007,6 @@ def fn(y):
         else:
             return x.cos()
 
-    @unittest.expectedFailure
     def test_getattr_metaclass(self):
         class Meta(type):
             def __getattr__(cls, name):
@@ -3337,7 +3430,7 @@ def fn2(f0, f1, args):
         )  # Recompile! input is no longer a functools partial
 
     def test_manual_seed(self):
-        @torch.compile
+        @torch.compile(backend="eager")
         def foo():
             torch.manual_seed(3)
             return torch.randint(0, 5, (5,))
@@ -3399,7 +3492,7 @@ def forward(inp):
         self.assertEqual(cnts.frame_count, 1)
 
     def test_complex_closure(self):
-        @torch.compile
+        @torch.compile(backend="eager")
         def forward(y):
             def a():
                 def x(z):
@@ -3415,7 +3508,7 @@ def x(z):
         self.assertTrue(same(res, input1 + input2))
 
     def test_non_inlined_closure(self):
-        @torch.compile()
+        @torch.compile(backend="eager")
         def program(x, y):
             one = lambda x, y: x + y
 
@@ -3448,7 +3541,7 @@ def test_np_constant_collections_as_input(self, int_or_float):
         def func(a, b, info_or_dt):
             return a + info_func(info_or_dt).max
 
-        opt_fn = torch.compile(func)
+        opt_fn = torch.compile(func, backend="eager")
 
         a = torch.randn(2)
         b = torch.randn(2)
@@ -3528,7 +3621,7 @@ def test_compare_constant_and_tensor(self):
                 def fn(x):
                     return op(-10, x)
 
-                opt_fn = torch.compile(fullgraph=True)(fn)
+                opt_fn = torch.compile(fullgraph=True, backend="eager")(fn)
 
                 x = torch.randn(10)
                 self.assertEqual(opt_fn(x), fn(x))
@@ -3537,7 +3630,7 @@ def test_pos(self):
         def fn(x, y):
             return operator.pos(x) * +y
 
-        opt_fn = torch.compile(fullgraph=True, dynamic=True)(fn)
+        opt_fn = torch.compile(fullgraph=True, dynamic=True, backend="eager")(fn)
 
         def test(x, y):
             self.assertEqual(opt_fn(x, y), fn(x, y))
@@ -3559,7 +3652,7 @@ def test(a, b):
 
         for dynamic in [True, False]:
             torch._dynamo.reset()
-            opt_fn = torch.compile(fn, dynamic=dynamic)
+            opt_fn = torch.compile(fn, dynamic=dynamic, backend="eager")
             t = torch.ones(1)
             test(10, t)
             test(-100, t)
@@ -3571,7 +3664,7 @@ def test_truth(self):
         def fn(x, y):
             return operator.truth(x) and bool(y)
 
-        opt_fn = torch.compile(dynamic=False)(fn)
+        opt_fn = torch.compile(dynamic=False, backend="eager")(fn)
 
         def test(x, y):
             self.assertEqual(opt_fn(x, y), fn(x, y))
@@ -3592,7 +3685,7 @@ def fn():
                     a = range(-10, 10)
                     return list(map(op, a))
 
-                opt_fn = torch.compile(fn, fullgraph=True)
+                opt_fn = torch.compile(fn, fullgraph=True, backend="eager")
                 self.assertEqual(opt_fn(), fn())
 
     def test_unary_fold_op_seq(self):
@@ -3603,7 +3696,7 @@ def fn():
                     a = [tuple(range(-10, i)) for i in range(10)]
                     return tuple(map(op, a))
 
-                opt_fn = torch.compile(fn, fullgraph=True)
+                opt_fn = torch.compile(fn, fullgraph=True, backend="eager")
                 self.assertEqual(opt_fn(), fn())
 
     def test_attrgetter(self):
@@ -3619,7 +3712,7 @@ def fn(x, y):
                     getter = operator.attrgetter(*attrs)
                     return getter(x), getter(y)
 
-                opt_fn = torch.compile(fullgraph=True)(fn)
+                opt_fn = torch.compile(fullgraph=True, backend="eager")(fn)
 
                 x = torch.randn(3, 4)
                 y = torch.randn(3, 4)
@@ -3638,7 +3731,7 @@ def fn(x, y):
                     getter = operator.itemgetter(*items)
                     return getter(x), getter(y)
 
-                opt_fn = torch.compile(fullgraph=True)(fn)
+                opt_fn = torch.compile(fullgraph=True, backend="eager")(fn)
 
                 x = torch.randn(3, 4)
                 y = torch.randn(3, 4)
@@ -3657,7 +3750,7 @@ def fn(x, y):
                     caller = operator.methodcaller(name, *args, **kwargs)
                     return caller(x), caller(y)
 
-                opt_fn = torch.compile(fullgraph=True)(fn)
+                opt_fn = torch.compile(fullgraph=True, backend="eager")(fn)
 
                 x = torch.randn(3, 4)
                 y = torch.randn(3, 4)
@@ -3693,7 +3786,7 @@ def fn(x):
             y = x + next(it)
             z = g(y, it)
             k = next(it)
-            assert k == 7
+            assert k == 7  # noqa: S101
             return z + k
 
         x = torch.tensor([1.0])
@@ -3845,7 +3938,7 @@ def fn(x):
                 acc *= acc * k
             return x * acc
 
-        opt_fn = torch.compile(fullgraph=True)(fn)
+        opt_fn = torch.compile(fullgraph=True, backend="eager")(fn)
         x = torch.ones(1)
         self.assertEqual(opt_fn(x), fn(x))
 
@@ -3855,7 +3948,7 @@ def fn(x):
             acc *= acc * range(10, 20, 2)[2]
             return x * acc
 
-        opt_fn = torch.compile(fullgraph=True)(fn)
+        opt_fn = torch.compile(fullgraph=True, backend="eager")(fn)
         x = torch.ones(1)
         self.assertEqual(opt_fn(x), fn(x))
 
@@ -3943,6 +4036,20 @@ def fn(t):
         g = torch.compile(fn, backend="eager", fullgraph=True)(t)
         self.assertEqual(e, g)
 
+    @unittest.skipIf(sys.platform == "darwin", "No mkldnn on MacOS")
+    def test_quantize_per_tensor(self):
+        def fn(t, scale, zero_point):
+            return torch.quantize_per_tensor(t, scale, zero_point, torch.quint8)
+
+        scale = torch.tensor(2.0)
+        zero_point = torch.tensor(10.0)
+        t = torch.rand((2, 2)) * scale + zero_point
+
+        result = fn(t, scale, zero_point)
+        compiled_fn = torch.compile(fn, fullgraph=True)
+        compiled_result = compiled_fn(t, scale, zero_point)
+        self.assertEqual(compiled_result, result)
+
     def test_map_return(self):
         def fn(a, b):
             return map(lambda x: x + 1, [a, b])
@@ -4133,8 +4240,8 @@ def self_fn(x):
         v = torch.ones([3], device="cpu")
         # identical tensor since modify inplace
         v2 = torch.ones([3], device="cpu")
-        opt_fn = torch.compile(fn)
-        opt_self_fn = torch.compile(self_fn)
+        opt_fn = torch.compile(fn, backend="eager")
+        opt_self_fn = torch.compile(self_fn, backend="eager")
         self.assertEqual(v, v2)
         self.assertEqual(opt_fn(v), opt_self_fn(v2))
 
@@ -4200,6 +4307,165 @@ def fn(x):
         opt_fn = torch.compile(fn, backend="eager", fullgraph=True)
         self.assertEqual(fn(x), opt_fn(x))
 
+    def test_wrapper_user_function_hasattr(self):
+        # WrapperUserFunctionVariable (e.g. lru_cache-wrapped fn) passed to a
+        # decorator that calls functools.wraps at tracing time should not graph
+        # break on hasattr(fn, '__dict__').
+        @functools.lru_cache
+        def cached_fn(x):
+            return x * 2
+
+        def retry(func):
+            @functools.wraps(func)
+            def wrapper(*args, **kwargs):
+                return func(*args, **kwargs)
+
+            return wrapper
+
+        def fn(x):
+            return retry(cached_fn)(x)
+
+        x = torch.tensor(2.0, device="cpu")
+        opt_fn = torch.compile(fn, backend="eager", fullgraph=True)
+        self.assertEqual(fn(x), opt_fn(x))
+
+    def test_wraps_on_lru_cache_preserves_name(self):
+        # functools.wraps copies __name__ and __qualname__ from the wrapped fn;
+        # when applied to an lru_cache-wrapped function at trace time,
+        # WrapperUserFunctionVariable must expose those attributes.
+        @functools.lru_cache
+        def my_op(x):
+            """my docstring"""
+            return x + 1
+
+        def apply_wraps(func):
+            @functools.wraps(func)
+            def inner(*args, **kwargs):
+                return func(*args, **kwargs)
+
+            return inner
+
+        def fn(x):
+            wrapped = apply_wraps(my_op)
+            if wrapped.__name__ != "my_op":
+                raise AssertionError(f"Expected 'my_op', got {wrapped.__name__!r}")
+            if wrapped.__qualname__ != my_op.__qualname__:
+                raise AssertionError(
+                    f"Expected {my_op.__qualname__!r}, got {wrapped.__qualname__!r}"
+                )
+            if wrapped.__doc__ != "my docstring":
+                raise AssertionError(
+                    f"Expected 'my docstring', got {wrapped.__doc__!r}"
+                )
+            return wrapped(x)
+
+        x = torch.tensor(1.0)
+        opt_fn = torch.compile(fn, backend="eager", fullgraph=True)
+        self.assertEqual(fn(x), opt_fn(x))
+
+    def test_wraps_on_lru_cache_copies_annotations(self):
+        # functools.wraps should copy __annotations__ from an lru_cache-wrapped fn.
+        @functools.lru_cache
+        def annotated_fn(x: torch.Tensor) -> torch.Tensor:
+            return x * 3
+
+        def decorator(func):
+            @functools.wraps(func)
+            def inner(*args, **kwargs):
+                return func(*args, **kwargs)
+
+            return inner
+
+        def fn(x):
+            return decorator(annotated_fn)(x)
+
+        x = torch.tensor(2.0)
+        opt_fn = torch.compile(fn, backend="eager", fullgraph=True)
+        self.assertEqual(fn(x), opt_fn(x))
+
+    def test_wraps_on_nested_fn(self):
+        # functools.wraps applied to a locally defined (NestedUserFunctionVariable)
+        # function should also work without graph breaks.
+        def fn(x):
+            def inner_op(t):
+                """inner doc"""
+                return t * 2
+
+            @functools.wraps(inner_op)
+            def wrapper(*args, **kwargs):
+                return inner_op(*args, **kwargs)
+
+            if wrapper.__name__ != "inner_op":
+                raise AssertionError(f"Expected 'inner_op', got {wrapper.__name__!r}")
+            if wrapper.__doc__ != "inner doc":
+                raise AssertionError(f"Expected 'inner doc', got {wrapper.__doc__!r}")
+            return wrapper(x)
+
+        x = torch.tensor(3.0)
+        opt_fn = torch.compile(fn, backend="eager", fullgraph=True)
+        self.assertEqual(fn(x), opt_fn(x))
+
+    def test_wraps_stacked_on_lru_cache(self):
+        # Stacking two functools.wraps layers over an lru_cache-wrapped fn.
+        @functools.lru_cache
+        def base_fn(x):
+            return x - 1
+
+        def outer_decorator(func):
+            @functools.wraps(func)
+            def middle(*args, **kwargs):
+                return func(*args, **kwargs)
+
+            return middle
+
+        def inner_decorator(func):
+            @functools.wraps(func)
+            def innermost(*args, **kwargs):
+                return func(*args, **kwargs)
+
+            return innermost
+
+        def fn(x):
+            return inner_decorator(outer_decorator(base_fn))(x)
+
+        x = torch.tensor(5.0)
+        opt_fn = torch.compile(fn, backend="eager", fullgraph=True)
+        self.assertEqual(fn(x), opt_fn(x))
+
+    def test_lru_cache_dunder_name_access(self):
+        # Accessing __name__ on an lru_cache-wrapped function during tracing
+        # should return the original function's name as a constant.
+        @functools.lru_cache
+        def compute(x):
+            return x + 10
+
+        def fn(x):
+            name = compute.__name__
+            if name != "compute":
+                raise AssertionError(f"Expected 'compute', got {name!r}")
+            return compute(x)
+
+        x = torch.tensor(1.0)
+        opt_fn = torch.compile(fn, backend="eager", fullgraph=True)
+        self.assertEqual(fn(x), opt_fn(x))
+
+    def test_lru_cache_dunder_doc_access(self):
+        # Accessing __doc__ on an lru_cache-wrapped function during tracing.
+        @functools.lru_cache
+        def documented_fn(x):
+            """returns x squared"""
+            return x**2
+
+        def fn(x):
+            doc = documented_fn.__doc__
+            if doc != "returns x squared":
+                raise AssertionError(f"Expected 'returns x squared', got {doc!r}")
+            return documented_fn(x)
+
+        x = torch.tensor(3.0)
+        opt_fn = torch.compile(fn, backend="eager", fullgraph=True)
+        self.assertEqual(fn(x), opt_fn(x))
+
     def test_functools_cache_guard(self):
         class Foo:
             @functools.lru_cache  # noqa: B019
@@ -4374,7 +4640,7 @@ def forward(self):
         return self.m()
 
 
-class DefaultsTests(torch._dynamo.test_case.TestCaseWithNestedGraphBreaks):
+class DefaultsTests(torch._dynamo.test_case.TestCase):
     def test_func_default_tensor_args(self):
         """
         Tests that we indeed reference (and mutate) "the one" default tensor arg
@@ -4459,6 +4725,38 @@ def test_meth_default_tensor_args(self):
         self.assertEqual(cnts.frame_count, 3)
         self.assertEqual(cnts.op_count, 6)
 
+    def test_guard_on_constant_func_defaults(self):
+        """
+        When a compiled function is re-invoked with a closure whose
+        __code__ is the same but __defaults__ differ (e.g. a different
+        constant default arg), Dynamo must recompile instead of reusing
+        the stale graph.
+        """
+        cnts = torch._dynamo.testing.CompileCounter()
+
+        def make_adder(offset):
+            def adder(x, _offset=offset):
+                return x + _offset
+
+            return adder
+
+        @torch.compile(backend=cnts)
+        def call_adder(x, fn):
+            return fn(x)
+
+        x = torch.ones(4)
+
+        adder0 = make_adder(0)
+        result0 = call_adder(x, adder0)
+        self.assertEqual(result0, x + 0)
+        self.assertEqual(cnts.frame_count, 1)
+
+        # Same __code__, different __defaults__ → must recompile
+        adder5 = make_adder(5)
+        result5 = call_adder(x, adder5)
+        self.assertEqual(result5, x + 5)
+        self.assertEqual(cnts.frame_count, 2)
+
     def test_func_default_torch_args(self):
         """
         Tests other types of torch types as function default (size, dtype, device)
@@ -4548,6 +4846,47 @@ def fn(x):
         ref = opt_fn(x)
         self.assertEqual(ref, res)
 
+    def test_pydantic_dataclass_construction(self):
+        @torch._dynamo.disable
+        def populate(self, x, y):
+            self.x = x
+            self.y = y
+
+        @dataclass(init=False)
+        class Point:
+            x: torch.Tensor
+            y: torch.Tensor
+            # Pydantic uses this sentinel on decorated dataclasses.
+            __is_pydantic_dataclass__ = True
+
+            def __init__(self, x, y):
+                populate(self, x, y)
+
+        def fn(x, y):
+            p = Point(x=x, y=y)
+            return p.x + p.y
+
+        torch._dynamo.reset()
+        counters.clear()
+        cnts = torch._dynamo.testing.CompileCounter()
+        compiled_fn = torch.compile(fn, backend=cnts)
+        x = torch.randn(4)
+        y = torch.randn(4)
+
+        self.assertTrue(same(fn(x, y), compiled_fn(x, y)))
+        self.assertEqual(cnts.frame_count, 0)
+        self.assertEqual(cnts.op_count, 0)
+        # Skipping the whole frame records a second follow-on graph break, so
+        # assert on the specific pydantic entry rather than the raw count.
+        self.assertEqual(
+            [
+                count
+                for msg, count in counters["graph_break"].items()
+                if "Pydantic dataclass constructor" in msg
+            ],
+            [1],
+        )
+
     def test_listlike_of_tensors_contains_constant(self):
         for listlike in [set, list]:
 
@@ -4852,7 +5191,8 @@ def test_in_set_would_fail_broadcast(self):
 
         tensor_list = set()
         tensor_list.add(param2)
-        assert param not in tensor_list
+        if param in tensor_list:
+            raise AssertionError("Expected param not in tensor_list")
 
         def fn(param, param2):
             param.add_(1)
@@ -4873,7 +5213,8 @@ def test_in_set_inplace(self):
 
         tensor_list = set()
         tensor_list.add(param2)
-        assert param not in tensor_list
+        if param in tensor_list:
+            raise AssertionError("Expected param not in tensor_list")
 
         def fn(param, param2):
             y = param.add_(1)  # Tensor method
@@ -5012,28 +5353,6 @@ def g(x, y):
         self.assertEqual(f(torch.ones(3, 3)), opt_f(torch.ones(3, 3)))
         self.assertEqual(cnts.frame_count, 3)
 
-    @unittest.skipIf(not TEST_MULTIGPU, "detected only one GPU")
-    def test_gpu_current_device(self):
-        def fn(x):
-            y = torch.empty(
-                (2, 3),
-                dtype=torch.float32,
-                device=torch.accelerator.current_device_index(),
-            )
-            y.copy_(x)
-            return torch.sin(y + y.device.index)
-
-        counter = torch._dynamo.testing.CompileCounter()
-        opt_fn = torch.compile(backend=counter, fullgraph=True)(fn)
-
-        with torch.accelerator.device_index(0):
-            x = torch.randn(2, 3)
-            self.assertEqual(opt_fn(x), fn(x))
-            self.assertEqual(counter.frame_count, 1)
-            with torch.accelerator.device_index(1):
-                self.assertEqual(opt_fn(x), fn(x))
-                self.assertEqual(counter.frame_count, 2)
-
     def test_fn_with_attr(self):
         def fn(x):
             if fn.pred:
@@ -5067,7 +5386,9 @@ def foo_custom_str(x):
             return x, str(a)
 
         eager_custom_str = foo_custom_str(torch.ones(4))
-        dynamo_custom_str = torch.compile(foo_custom_str, fullgraph=True)(torch.ones(4))
+        dynamo_custom_str = torch.compile(
+            foo_custom_str, fullgraph=True, backend="eager"
+        )(torch.ones(4))
 
         self.assertEqual(eager_custom_str[1], dynamo_custom_str[1])
         self.assertEqual(eager_custom_str[1], "ok")
@@ -5080,9 +5401,9 @@ def foo_default_str(x):
             return x, str(a)
 
         eager_default_str = foo_default_str(torch.ones(4))
-        dynamo_default_str = torch.compile(foo_default_str, fullgraph=True)(
-            torch.ones(4)
-        )
+        dynamo_default_str = torch.compile(
+            foo_default_str, fullgraph=True, backend="eager"
+        )(torch.ones(4))
 
         # Check that the tensor output from eager and dynamo modes are the same
         self.assertEqual(eager_default_str[0], dynamo_default_str[0])
@@ -5122,7 +5443,7 @@ def call(index):
 
         a = torch.randn(4, 2, 5)
         b = torch.randn(4, 2, 5, 5)
-        compiled_fn = torch.compile(fn, fullgraph=True)
+        compiled_fn = torch.compile(fn, fullgraph=True, backend="eager")
         compiled_res = compiled_fn(a, b, torch.tensor([2]))
         reference_res = fn(a, b, torch.tensor([2]))
         self.assertTrue(same(compiled_res, reference_res))
@@ -5444,7 +5765,7 @@ def get_torch_functional_functions():
         functions = get_torch_functional_functions()
         self.assertTrue(len(functions) > 0)
         for func in functions:
-            compiled_func = torch.compile(func)
+            compiled_func = torch.compile(func, backend="eager")
             self.assertTrue(callable(compiled_func))
 
     def test_skip_function_call_very_weird_value(self):
@@ -5555,7 +5876,7 @@ def test_full_with_tensor_fill_value(self):
         def func_tensor(x):
             return torch.full((2,), x, dtype=torch.float64)
 
-        func_compiled = torch.compile(func_tensor)
+        func_compiled = torch.compile(func_tensor, backend="eager")
 
         # Test with different values
         x1 = torch.tensor(5.0, dtype=torch.float64)
@@ -5576,7 +5897,7 @@ def func_tensor(x):
             def func_typed(x):
                 return torch.full((3,), x, dtype=dtype)
 
-            func_typed_compiled = torch.compile(func_typed)
+            func_typed_compiled = torch.compile(func_typed, backend="eager")
             x_typed = torch.tensor(7, dtype=dtype)
             result = func_typed_compiled(x_typed)
             expected = torch.full((3,), x_typed, dtype=dtype)
@@ -5586,7 +5907,7 @@ def func_typed(x):
         def func_scalar(size):
             return torch.full((size,), 42.0, dtype=torch.float32)
 
-        func_scalar_compiled = torch.compile(func_scalar)
+        func_scalar_compiled = torch.compile(func_scalar, backend="eager")
 
         result_scalar = func_scalar_compiled(5)
         expected_scalar = torch.full((5,), 42.0, dtype=torch.float32)
@@ -5599,7 +5920,7 @@ def func_scalar_param():
             b = torch.full((2,), 2.71, dtype=torch.float32)
             return a, b
 
-        func_scalar_param_compiled = torch.compile(func_scalar_param)
+        func_scalar_param_compiled = torch.compile(func_scalar_param, backend="eager")
         result_a, result_b = func_scalar_param_compiled()
 
         self.assertEqual(result_a, torch.full((2,), 3.14, dtype=torch.float32))
diff --git a/test/dynamo/test_fwd_loss_bwd.py b/test/dynamo/test_fwd_loss_bwd.py
index d357a2b7ab05c..5cfc2f11b631b 100644
--- a/test/dynamo/test_fwd_loss_bwd.py
+++ b/test/dynamo/test_fwd_loss_bwd.py
@@ -541,7 +541,8 @@ def step(mod, x):
                 """\
 autograd.grad consumed returned tensor's grad_fn
   Explanation: torch.autograd.grad() consumes grad_fns that are needed by tensors returned from this compiled function. This would cause 'backward through graph a second time' errors.
-  Hint: If you don't need to backward through the returned tensor, call .detach() before returning: `return loss.detach()`
+      The following returned tensors have consumed grad_fns: loss
+  Hint: Detach the problematic tensor(s) before returning: e.g. `loss.detach()`
   Hint: If you need to backward through the returned tensor, use retain_graph=True in autograd.grad()."""  # noqa: B950
             ),
         ):
@@ -590,17 +591,29 @@ def fn(x):
         torch._dynamo.reset()
         compiled_fn = torch.compile(fn, fullgraph=True, backend="aot_eager")
 
-        msg = textwrap.dedent(
-            """\
-autograd.grad consumed returned tensor's grad_fn
-  Explanation: torch.autograd.grad() consumes grad_fns that are needed by tensors returned from this compiled function. This would cause 'backward through graph a second time' errors.
-  Hint: If you don't need to backward through the returned tensor, call .detach() before returning: `return loss.detach()`
-  Hint: If you need to backward through the returned tensor, use retain_graph=True in autograd.grad()."""  # noqa: B950
-        )
+        with self.assertRaisesRegex(
+            torch._dynamo.exc.Unsupported,
+            r"autograd\.grad consumed returned tensor's grad_fn",
+        ):
+            compiled_fn(torch.randn(4, requires_grad=True))
+
+    def test_autograd_grad_leaked_tensor_names_in_error(self):
+        """Test that the error message includes the names of all leaked tensors."""
+        torch._dynamo.reset()
+
+        def fn(x):
+            a = x * 2
+            b = x * 3
+            z = (a + b).sum()
+            torch.autograd.grad(z, x)
+            # Both a and b have consumed grad_fns
+            return a, b
+
+        compiled_fn = torch.compile(fn, fullgraph=True, backend="aot_eager")
 
         with self.assertRaisesRegex(
             torch._dynamo.exc.Unsupported,
-            re.escape(msg) + r"[\s\S]*",
+            r"Leaked output tensors:",
         ):
             compiled_fn(torch.randn(4, requires_grad=True))
 
@@ -1205,11 +1218,86 @@ def fn(a, b):
         ref = fn(a, b)
 
         # Second call - compiled (should accumulate into SAME tensor)
-        opt_fn = torch.compile(fn)
+        opt_fn = torch.compile(fn, backend="eager")
         act = opt_fn(a, b)
 
         self.assertTrue(ref is act)
 
+    def test_autograd_grad_lost_grad_fn_in_closure(self):
+        def f(x):
+            return (x**2).sum()
+
+        x = torch.randn(4, requires_grad=True)
+        _, vjp_fn = torch.func.vjp(f, x)
+
+        with self.assertRaisesRegex(
+            torch._dynamo.exc.Unsupported,
+            "_autograd_grad with lost grad_fn linkage",
+        ):
+            torch.compile(vjp_fn, backend="eager", fullgraph=True)(torch.ones(()))
+
+    def test_autograd_grad_transform_closure_compiled_separately(self):
+        def f(x):
+            return (x**2).sum()
+
+        x = torch.randn(4, requires_grad=True)
+        _, vjp_fn = torch.func.vjp(f, x)
+        v = torch.ones(())
+
+        eager_grad = vjp_fn(v)
+
+        cnt = torch._dynamo.testing.CompileCounter()
+        compiled_grad = torch.compile(vjp_fn, backend=cnt)(v)
+
+        self.assertEqual(compiled_grad, eager_grad)
+        self.assertEqual(cnt.frame_count, 0)
+
+    @skipIfCrossRef
+    def test_autograd_grad_transform_compiled_end_to_end(self):
+        def f(x):
+            return (x**2).sum()
+
+        @torch.compile(backend="eager", fullgraph=True)
+        def fn(x):
+            _, vjp_fn = torch.func.vjp(f, x)
+            return vjp_fn(torch.ones(()))
+
+        x = torch.randn(4, requires_grad=True)
+        result = fn(x)
+        expected = torch.func.vjp(f, x)[1](torch.ones(()))
+        self.assertEqual(result, expected)
+
+    def test_autograd_grad_multi_output_transform_closure(self):
+        def f(x):
+            return x * 2, x * 3
+
+        x = torch.randn(4, requires_grad=True)
+        _, vjp_fn = torch.func.vjp(f, x)
+        v = (torch.ones(4), torch.ones(4))
+
+        eager_grad = vjp_fn(v)
+
+        cnt = torch._dynamo.testing.CompileCounter()
+        compiled_grad = torch.compile(vjp_fn, backend=cnt)(v)
+        self.assertEqual(compiled_grad, eager_grad)
+
+    @skipIfCrossRef
+    def test_autograd_grad_inline_computation_no_graph_break(self):
+        def f(x):
+            return (x * 2).sum()
+
+        @torch.compile(backend="eager", fullgraph=True)
+        def fn(x):
+            y, vjp_fn = torch.func.vjp(f, x)
+            return y, vjp_fn(torch.ones(()))
+
+        x = torch.randn(4, requires_grad=True)
+        compiled_y, compiled_grad = fn(x)
+        eager_y, eager_vjp_fn = torch.func.vjp(f, x)
+        eager_grad = eager_vjp_fn(torch.ones(()))
+        self.assertEqual(compiled_y, eager_y)
+        self.assertEqual(compiled_grad, eager_grad)
+
 
 if __name__ == "__main__":
     run_tests()
diff --git a/test/dynamo/test_fx_annotate.py b/test/dynamo/test_fx_annotate.py
index 2ec0b93c986ae..71c09b2b7a577 100644
--- a/test/dynamo/test_fx_annotate.py
+++ b/test/dynamo/test_fx_annotate.py
@@ -1,5 +1,7 @@
 # Owner(s): ["module: dynamo"]
 
+import warnings
+
 import torch
 import torch._dynamo.test_case
 import torch.fx.traceback as fx_traceback
@@ -237,6 +239,56 @@ def fn(x):
 ('call_function', 'getitem_5', {'compile_inductor': 0})""",  # noqa: B950
         )
 
+    @requires_cuda_and_triton
+    def test_flex_attention_backward_tag_does_not_leak(self):
+        from torch.fx.experimental.proxy_tensor import make_fx
+        from torch.fx.traceback import preserve_node_meta
+
+        def causal_mask(batch, head, q_idx, kv_idx):
+            del batch, head
+            return q_idx >= kv_idx
+
+        q = torch.randn(
+            1, 2, 128, 32, device="cuda", dtype=torch.bfloat16, requires_grad=True
+        )
+        k = torch.randn(
+            1, 2, 128, 32, device="cuda", dtype=torch.bfloat16, requires_grad=True
+        )
+        v = torch.randn(
+            1, 2, 128, 32, device="cuda", dtype=torch.bfloat16, requires_grad=True
+        )
+        block_mask = create_block_mask(causal_mask, 1, 2, 128, 128, device="cuda")
+
+        def fn(q, k, v, block_mask):
+            with fx_traceback.annotate({"ac_region_id": 0}):
+                y = flex_attention(q, k, v, block_mask=block_mask)
+                torch.autograd.grad(y, (q, k, v), torch.ones_like(y))
+                return y.cos()
+
+        warnings.filterwarnings(
+            "ignore",
+            message="flex_attention called without torch.compile",
+        )
+        with (
+            torch._dynamo.config.patch(error_on_nested_fx_trace=False),
+            torch.compiler._non_strict_tracing_context(),
+            torch.compiler._patch_autograd_grad(),
+            preserve_node_meta(),
+        ):
+            gm = make_fx(fn, record_stack_traces=True)(q, k, v, block_mask)
+
+        backward_nodes = [
+            node for node in gm.graph.nodes if node.meta.get("autograd_backward", False)
+        ]
+        self.assertTrue(backward_nodes)
+
+        flex_nodes = gm.graph.find_nodes(
+            op="call_function", target=torch.ops.higher_order.flex_attention
+        )
+        self.assertEqual(len(flex_nodes), 1)
+        self.assertNotIn("autograd_backward", flex_nodes[0].meta)
+        self.assertEqual(flex_nodes[0].meta.get("custom", {}), {"ac_region_id": 0})
+
     def test_as_decorator(self):
         class Mod(torch.nn.Module):
             @fx_traceback.annotate({"fdsp_bucket": 0})
@@ -299,6 +351,36 @@ def fn(x):
         x = torch.randn(10, requires_grad=True)
         self.assertEqual(fn(x), opt_fn(x))
 
+    def test_annotation_on_runtime_asserts(self):
+        class M(torch.nn.Module):
+            def forward(self, x, y):
+                with torch.fx.traceback.annotate({"moo": 0}):
+                    x = torch.cat([x, x])
+                    b = y.item()
+                    torch._check(b >= x.shape[0])
+                    return x * b
+
+        backend = AotEagerAndRecordGraphs()
+
+        inputs = (torch.randn(3), torch.tensor(6))
+        model = M()
+        torch._dynamo.mark_dynamic(inputs[0], 0)
+        torch.compile(model, fullgraph=True, backend=backend)(*inputs)
+        dynamo_metadata = fx_traceback._get_custom_metadata(backend.graphs[0])
+        self.assertExpectedInline(
+            str(dynamo_metadata),
+            """\
+('placeholder', 's77', {'moo': 0})
+('placeholder', 'l_x_', {'moo': 0})
+('placeholder', 'l_y_', {'moo': 0})
+('call_function', 'x', {'moo': 0})
+('call_method', 'item', {'moo': 0})
+('call_function', 'mul_1', {'moo': 0})
+('call_function', 'ge', {'moo': 0})
+('call_function', '_check', {'moo': 0})
+('call_function', 'mul', {'moo': 0})""",  # noqa: B950
+        )
+
 
 if __name__ == "__main__":
     run_tests()
diff --git a/test/dynamo/test_fx_graph_runnable.py b/test/dynamo/test_fx_graph_runnable.py
index 9ff32cbc2e60a..5da16806e1d1a 100644
--- a/test/dynamo/test_fx_graph_runnable.py
+++ b/test/dynamo/test_fx_graph_runnable.py
@@ -203,6 +203,7 @@ def f(x):
         self._exec_and_verify_payload()
 
     @unittest.skipUnless(has_triton(), "Triton not available")
+    @torch._dynamo.config.patch(nested_graph_breaks=False)
     def test_user_defined_triton_kernel_autotune(self):
         def add(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
             output = torch.ones(x.shape, device=x.device, dtype=x.dtype)
@@ -549,6 +550,44 @@ def f(view, weights):
         torch.compile(f)(view, weights)
         self._exec_and_verify_payload()
 
+    @torch._dynamo.config.patch(assume_static_by_default=False)
+    def test_repeat_interleave_with_output_size(self):
+        def f(data, repeats, output_size):
+            indices = torch.repeat_interleave(repeats, output_size=output_size.item())
+            return data[indices]
+
+        num_segments = 128
+        data = torch.randn(1000, 16)
+        repeats = torch.randint(5, 15, (num_segments,), dtype=torch.int64)
+        output_size = repeats.sum()
+
+        torch.compile(f, dynamic=True)(data, repeats, output_size)
+
+        self._exec_and_verify_payload()
+
+        # Verify the payload contains the repeat_interleave fixup
+        payload = self.buffer.getvalue().strip()
+        self.assertIn("def forward", payload)
+        self.assertIn("repeat_interleave", payload)
+        # Verify the fixup code is present
+        self.assertIn("# Fixup: ensure sum(repeats) == output_size", payload)
+        self.assertIn("_repeats.fill_", payload)
+
+    def test_repeat_interleave_with_constant_output_size(self):
+        def f(data, repeats):
+            # output_size is a constant, not a dynamic input
+            indices = torch.repeat_interleave(repeats, output_size=1280)
+            return data[indices]
+
+        num_segments = 128
+        data = torch.randn(1000, 16)
+        repeats = torch.full((num_segments,), 10, dtype=torch.int64)
+
+        torch.compile(f)(data, repeats)
+        self._exec_and_verify_payload()
+        payload = self.buffer.getvalue().strip()
+        self.assertNotIn("# Fixup: ensure sum(repeats) == output_size", payload)
+
 
 @unittest.skipIf(IS_FBCODE or IS_SANDCASTLE, "Skip in fbcode/sandcastle")
 class TestFxGraphRunnableMultiProcessGroup(TestCase):
diff --git a/test/dynamo/test_generator.py b/test/dynamo/test_generator.py
index 495f7cd3a4b01..d10bf6314a091 100644
--- a/test/dynamo/test_generator.py
+++ b/test/dynamo/test_generator.py
@@ -17,7 +17,7 @@
 )
 
 
-class GeneratorTestsBase(torch._dynamo.test_case.TestCaseWithNestedGraphBreaks):
+class GeneratorTestsBase(torch._dynamo.test_case.TestCase):
     def setUp(self):
         super().setUp()
         self._old = torch._dynamo.config.enable_faithful_generator_behavior
@@ -1144,7 +1144,7 @@ def fn(t):
         except exc:
             z = 1
         finally:
-            assert z == 1
+            assert z == 1  # noqa: S101
 
     def test_close_with_subgen(self):
         L = []
@@ -1500,7 +1500,7 @@ def fn(t):
                 z = 1
             except Exception as e:
                 raise AssertionError from e
-            assert z == 1
+            assert z == 1  # noqa: S101
             return t.sin()
 
         self._compile_check(fn)
@@ -1520,7 +1520,7 @@ def fn(t):
             try:
                 gen.throw(ValueError)
             except StopIteration as e:
-                assert e.args[0] == 3
+                assert e.args[0] == 3  # noqa: S101
             except Exception as e:
                 raise AssertionError from e
             return t.sin()
@@ -1546,7 +1546,7 @@ def fn(t):
             try:
                 gen.throw(ValueError)
             except StopIteration as e:
-                assert e.args[0].x == 3
+                assert e.args[0].x == 3  # noqa: S101
             except Exception as e:
                 raise AssertionError from e
             return t.sin()
@@ -1568,7 +1568,7 @@ def fn(t):
             try:
                 gen.throw(ValueError)
             except StopIteration as e:
-                assert len(e.args) == 0
+                assert len(e.args) == 0  # noqa: S101
             except Exception as e:
                 raise AssertionError from e
             return t.sin()
diff --git a/test/dynamo/test_getitem.py b/test/dynamo/test_getitem.py
new file mode 100644
index 0000000000000..96de001815576
--- /dev/null
+++ b/test/dynamo/test_getitem.py
@@ -0,0 +1,795 @@
+# Owner(s): ["module: dynamo"]
+"""Tests for mp_subscript_impl: unified __getitem__ dispatch via vt_getitem in Dynamo.
+
+Tests exercise the vt_getitem → mp_subscript_impl path via operator.getitem(),
+and the call_method("__getitem__") → mp_subscript_impl path via obj.__getitem__().
+
+See TODO(follow-up) comments on each mp_subscript_impl override for remaining
+CPython behavioral gaps.
+"""
+
+import collections
+import operator
+import types
+import unittest
+
+import torch
+import torch._dynamo.test_case
+import torch._dynamo.testing
+from torch.testing._internal.inductor_utils import HAS_CUDA_AND_TRITON, HAS_GPU
+
+
+class GetItemTests(torch._dynamo.test_case.TestCase):
+    def _compile(self, fn, *args):
+        return torch.compile(fn, backend="eager", fullgraph=True)(*args)
+
+    # --- BaseListVariable (ListVariable) ---
+
+    def test_list_int_index(self):
+        def fn(x):
+            items = [x, x + 1, x + 2]
+            return (
+                operator.getitem(items, 0),
+                operator.getitem(items, 1),
+                operator.getitem(items, 2),
+            )
+
+        x = torch.randn(4)
+        self.assertEqual(fn(x), self._compile(fn, x))
+
+    def test_list_getitem_compiled_directly(self):
+        compiled = torch.compile(operator.getitem, backend="eager", fullgraph=True)
+        items = [10, 20, 30]
+        self.assertEqual(compiled(items, 0), 10)
+        self.assertEqual(compiled(items, 2), 30)
+
+    def test_list_slice(self):
+        def fn(x):
+            items = [x, x + 1, x + 2]
+            full = operator.getitem(items, slice(None))
+            partial = operator.getitem(items, slice(0, 2))
+            single = operator.getitem(items, slice(1, 2))
+            return full, partial, single
+
+        x = torch.randn(4)
+        self.assertEqual(fn(x), self._compile(fn, x))
+
+    def test_list_negative_index(self):
+        def fn(x):
+            items = [x, x + 1, x + 2]
+            return operator.getitem(items, -1)
+
+        x = torch.randn(4)
+        self.assertEqual(fn(x), self._compile(fn, x))
+
+    def test_list_bool_index(self):
+        def fn(x):
+            items = [x, x + 1, x + 2]
+            return (
+                operator.getitem(items, False),
+                operator.getitem(items, True),
+            )
+
+        x = torch.randn(4)
+        self.assertEqual(fn(x), self._compile(fn, x))
+
+    def test_list_invalid_index_type(self):
+        def fn(x):
+            items = [x, x + 1, x + 2]
+            return operator.getitem(items, "a")
+
+        x = torch.randn(4)
+        with self.assertRaises(torch._dynamo.exc.Unsupported):
+            self._compile(fn, x)
+
+    def test_list_index_via_index_dunder(self):
+        """Custom __index__ object used as list index — _PyIndex_Check + nb_index_impl."""
+
+        class Idx:
+            def __index__(self):
+                return 2
+
+        def fn(x):
+            items = [x, x + 1, x + 2]
+            return operator.getitem(items, Idx())
+
+        x = torch.randn(4)
+        self.assertEqual(fn(x), self._compile(fn, x))
+
+    # --- BaseListVariable (TupleVariable) ---
+
+    def test_tuple_int_index(self):
+        def fn(x):
+            items = (x, x + 1, x + 2)
+            return operator.getitem(items, 0)
+
+        x = torch.randn(4)
+        self.assertEqual(fn(x), self._compile(fn, x))
+
+    def test_tuple_negative_index(self):
+        def fn(x):
+            items = (x, x + 1, x + 2)
+            return operator.getitem(items, -1)
+
+        x = torch.randn(4)
+        self.assertEqual(fn(x), self._compile(fn, x))
+
+    def test_tuple_slice(self):
+        def fn(x):
+            items = (x, x + 1, x + 2)
+            full = operator.getitem(items, slice(None))
+            partial = operator.getitem(items, slice(0, 2))
+            single = operator.getitem(items, slice(1, 2))
+            return full, partial, single
+
+        x = torch.randn(4)
+        self.assertEqual(fn(x), self._compile(fn, x))
+
+    def test_tuple_bool_index(self):
+        def fn(x):
+            items = (x, x + 1, x + 2)
+            return operator.getitem(items, False)
+
+        x = torch.randn(4)
+        self.assertEqual(fn(x), self._compile(fn, x))
+
+    def test_tuple_invalid_index_type(self):
+        def fn(x):
+            items = (x, x + 1, x + 2)
+            return operator.getitem(items, "a")
+
+        x = torch.randn(4)
+        with self.assertRaises(torch._dynamo.exc.Unsupported):
+            self._compile(fn, x)
+
+    def test_tuple_index_via_index_dunder(self):
+        class Idx:
+            def __index__(self):
+                return 2
+
+        def fn(x):
+            items = (x, x + 1, x + 2)
+            return operator.getitem(items, Idx())
+
+        x = torch.randn(4)
+        self.assertEqual(fn(x), self._compile(fn, x))
+
+    # --- RangeVariable ---
+
+    def test_range_int_index(self):
+        def fn(x):
+            r = range(0, 10, 2)
+            return x + operator.getitem(r, 3)
+
+        x = torch.randn(4)
+        self.assertEqual(fn(x), self._compile(fn, x))
+
+    def test_range_negative_index(self):
+        def fn(x):
+            r = range(0, 10, 2)
+            return x + operator.getitem(r, -1)
+
+        x = torch.randn(4)
+        self.assertEqual(fn(x), self._compile(fn, x))
+
+    def test_range_slice(self):
+        def fn(x):
+            r = range(0, 10, 2)
+            result = operator.getitem(r, slice(1, 3))
+            return x + result[0]
+
+        x = torch.randn(4)
+        self.assertEqual(fn(x), self._compile(fn, x))
+
+    def test_range_bool_index(self):
+        def fn(x):
+            r = range(0, 10, 2)
+            return x + operator.getitem(r, True)
+
+        x = torch.randn(4)
+        self.assertEqual(fn(x), self._compile(fn, x))
+
+    def test_range_index_via_index_dunder(self):
+        class Idx:
+            def __index__(self):
+                return 2
+
+        def fn(x):
+            r = range(0, 10, 2)
+            return x + operator.getitem(r, Idx())
+
+        x = torch.randn(4)
+        self.assertEqual(fn(x), self._compile(fn, x))
+
+    def test_range_invalid_index_type(self):
+        def fn(x):
+            r = range(0, 10, 2)
+            return x + operator.getitem(r, "a")
+
+        x = torch.randn(4)
+        with self.assertRaises(torch._dynamo.exc.Unsupported):
+            self._compile(fn, x)
+
+    # --- SizeVariable ---
+
+    def test_size_int_index(self):
+        def fn(x):
+            s = x.size()
+            return x + operator.getitem(s, 0)
+
+        x = torch.randn(4, 8)
+        self.assertEqual(fn(x), self._compile(fn, x))
+
+    def test_size_negative_index(self):
+        def fn(x):
+            s = x.size()
+            return x + operator.getitem(s, -1)
+
+        x = torch.randn(4, 8)
+        self.assertEqual(fn(x), self._compile(fn, x))
+
+    def test_size_slice(self):
+        def fn(x):
+            s = x.size()
+            result = operator.getitem(s, slice(0, 1))
+            return x + result[0]
+
+        x = torch.randn(4, 8)
+        self.assertEqual(fn(x), self._compile(fn, x))
+
+    def test_size_bool_index(self):
+        def fn(x):
+            s = x.size()
+            return x + operator.getitem(s, False)
+
+        x = torch.randn(4, 8)
+        self.assertEqual(fn(x), self._compile(fn, x))
+
+    def test_size_index_via_index_dunder(self):
+        class Idx:
+            def __index__(self):
+                return 1
+
+        def fn(x):
+            s = x.size()
+            return x + operator.getitem(s, Idx())
+
+        x = torch.randn(4, 8)
+        self.assertEqual(fn(x), self._compile(fn, x))
+
+    def test_size_invalid_index_type(self):
+        def fn(x):
+            s = x.size()
+            return x + operator.getitem(s, "a")
+
+        x = torch.randn(4, 8)
+        with self.assertRaises(torch._dynamo.exc.Unsupported):
+            self._compile(fn, x)
+
+    # --- ConstDictVariable ---
+
+    def test_dict_str_key(self):
+        def fn(x):
+            d = {"a": x, "b": x + 1}
+            return operator.getitem(d, "a")
+
+        x = torch.randn(4)
+        self.assertEqual(fn(x), self._compile(fn, x))
+
+    def test_dict_int_key(self):
+        def fn(x):
+            d = {0: x, 1: x + 1}
+            return operator.getitem(d, 1)
+
+        x = torch.randn(4)
+        self.assertEqual(fn(x), self._compile(fn, x))
+
+    def test_dict_missing_key(self):
+        def fn(x):
+            d = {"a": x}
+            return operator.getitem(d, "missing")
+
+        x = torch.randn(4)
+        with self.assertRaises(torch._dynamo.exc.Unsupported):
+            self._compile(fn, x)
+
+    # --- DefaultDictVariable ---
+
+    def test_defaultdict_existing_key(self):
+        def fn(x):
+            d = collections.defaultdict(lambda: x + 99)
+            d["a"] = x
+            return operator.getitem(d, "a")
+
+        x = torch.randn(4)
+        self.assertEqual(fn(x), self._compile(fn, x))
+
+    def test_defaultdict_missing_key(self):
+        def fn(x):
+            d = collections.defaultdict(list)
+            operator.getitem(d, "new")
+            return x
+
+        x = torch.randn(4)
+        self.assertEqual(fn(x), self._compile(fn, x))
+
+    # --- TensorVariable ---
+
+    def test_tensor_int_index(self):
+        def fn(x):
+            return operator.getitem(x, 0)
+
+        x = torch.randn(4, 4)
+        self.assertEqual(fn(x), self._compile(fn, x))
+
+    def test_tensor_slice(self):
+        def fn(x):
+            return operator.getitem(x, slice(0, 2))
+
+        x = torch.randn(4, 4)
+        self.assertEqual(fn(x), self._compile(fn, x))
+
+    def test_tensor_tuple_index(self):
+        def fn(x):
+            return operator.getitem(x, (0, slice(1, 3)))
+
+        x = torch.randn(4, 4)
+        self.assertEqual(fn(x), self._compile(fn, x))
+
+    def test_tensor_getitem_torch_function_mode(self):
+        """TorchFunctionMode intercepts tensor __getitem__ and can modify behavior."""
+
+        class AddOneMode(torch.overrides.TorchFunctionMode):
+            def __torch_function__(self, func, types, args=(), kwargs=None):
+                result = func(*args, **(kwargs or {}))
+                if func is torch.Tensor.__getitem__:
+                    return result + 1
+                return result
+
+        def fn(x):
+            with AddOneMode():
+                return operator.getitem(x, 0)
+
+        x = torch.randn(4, 4)
+        expected = fn(x)
+        compiled = torch.compile(fn, backend="eager")(x)
+        self.assertEqual(expected, compiled)
+
+    def test_tensor_getitem_torch_function_subclass(self):
+        """Tensor subclass with __torch_function__ intercepts __getitem__."""
+
+        class ScaledTensor(torch.Tensor):
+            @classmethod
+            def __torch_function__(cls, func, types, args=(), kwargs=None):
+                result = super().__torch_function__(func, types, args, kwargs or {})
+                if func is torch.Tensor.__getitem__:
+                    return result * 2
+                return result
+
+        def fn(x):
+            return operator.getitem(x, 0)
+
+        x = ScaledTensor(torch.randn(4, 4))
+        expected = fn(x)
+        compiled = torch.compile(fn, backend="eager")(x)
+        self.assertEqual(expected, compiled)
+
+    # --- NamedTupleVariable (via UserDefinedTupleVariable) ---
+
+    def test_namedtuple_int_index(self):
+        def fn(x):
+            result = torch.topk(x, 2)
+            return operator.getitem(result, 1)
+
+        x = torch.randn(10)
+        self.assertEqual(fn(x), self._compile(fn, x))
+
+    def test_namedtuple_values_index(self):
+        def fn(x):
+            result = torch.topk(x, 2)
+            return operator.getitem(result, 0)
+
+        x = torch.randn(10)
+        self.assertEqual(fn(x), self._compile(fn, x))
+
+    # --- TypingVariable ---
+
+    def test_typing_subscript(self):
+        def fn(x):
+            t = list[int]  # noqa: F841
+            return x + 1
+
+        x = torch.randn(4)
+        self.assertEqual(fn(x), self._compile(fn, x))
+
+    # --- MappingProxyVariable ---
+
+    def test_mappingproxy_getitem(self):
+        def fn(x):
+            d = {"a": 1, "b": 2}
+            proxy = types.MappingProxyType(d)
+            return x + operator.getitem(proxy, "a")
+
+        x = torch.randn(4)
+        self.assertEqual(fn(x), self._compile(fn, x))
+
+    # --- NNModuleVariable (ModuleList) ---
+
+    def test_nn_module_list_int_index(self):
+        class Model(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.layers = torch.nn.ModuleList(
+                    [torch.nn.Linear(4, 4) for _ in range(3)]
+                )
+
+            def forward(self, x):
+                return operator.getitem(self.layers, 1)(x)
+
+        model = Model()
+        x = torch.randn(4)
+        compiled = torch.compile(model, backend="eager", fullgraph=True)
+        self.assertEqual(model(x), compiled(x))
+
+    # --- NNModuleVariable (ModuleDict) ---
+
+    def test_nn_module_dict_str_key(self):
+        class Model(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.layers = torch.nn.ModuleDict({"fc": torch.nn.Linear(4, 4)})
+
+            def forward(self, x):
+                return operator.getitem(self.layers, "fc")(x)
+
+        model = Model()
+        x = torch.randn(4)
+        compiled = torch.compile(model, backend="eager", fullgraph=True)
+        self.assertEqual(model(x), compiled(x))
+
+    # --- NNModuleVariable (Sequential) ---
+
+    def test_nn_sequential_int_index(self):
+        class Model(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.seq = torch.nn.Sequential(
+                    torch.nn.Linear(4, 4),
+                    torch.nn.ReLU(),
+                )
+
+            def forward(self, x):
+                return operator.getitem(self.seq, 0)(x)
+
+        model = Model()
+        x = torch.randn(4)
+        compiled = torch.compile(model, backend="eager", fullgraph=True)
+        self.assertEqual(model(x), compiled(x))
+
+    # --- UserDefinedObjectVariable ---
+
+    def test_user_defined_object_getitem(self):
+        class Container:
+            def __init__(self, items):
+                self.items = items
+
+            def __getitem__(self, key):
+                return self.items[key]
+
+        def fn(x):
+            c = Container([x, x + 1])
+            return operator.getitem(c, 0)
+
+        x = torch.randn(4)
+        self.assertEqual(fn(x), self._compile(fn, x))
+
+    def test_object_without_getitem(self):
+        class NoGetItem:
+            pass
+
+        def fn(x):
+            obj = NoGetItem()
+            return operator.getitem(obj, 0)
+
+        x = torch.randn(4)
+        with self.assertRaises(torch._dynamo.exc.Unsupported):
+            self._compile(fn, x)
+
+    # --- UserDefinedListVariable ---
+
+    def test_user_defined_list_getitem(self):
+        class MyList(list):
+            pass
+
+        def fn(x):
+            items = MyList([x, x + 1, x + 2])
+            return operator.getitem(items, 1)
+
+        x = torch.randn(4)
+        self.assertEqual(fn(x), self._compile(fn, x))
+
+    # --- UserDefinedTupleVariable ---
+
+    def test_user_defined_tuple_getitem(self):
+        class MyTuple(tuple):  # noqa: SLOT001
+            pass
+
+        def fn(x):
+            items = MyTuple((x, x + 1, x + 2))
+            return operator.getitem(items, 1)
+
+        x = torch.randn(4)
+        self.assertEqual(fn(x), self._compile(fn, x))
+
+    # --- UserDefinedDictVariable ---
+
+    def test_user_defined_dict_getitem(self):
+        class MyDict(dict):
+            pass
+
+        def fn(x):
+            d = MyDict(a=x, b=x + 1)
+            return operator.getitem(d, "a")
+
+        x = torch.randn(4)
+        self.assertEqual(fn(x), self._compile(fn, x))
+
+    def test_user_defined_dict_missing(self):
+        class MyDict(dict):
+            def __missing__(self, key):
+                return 42
+
+        def fn(x):
+            d = MyDict(a=1)
+            return x + operator.getitem(d, "b")
+
+        x = torch.randn(4)
+        self.assertEqual(fn(x), self._compile(fn, x))
+
+    def test_user_defined_dict_custom_missing(self):
+        class DefaultDict(dict):
+            def __missing__(self, key):
+                self[key] = len(self)
+                return self[key]
+
+        def fn(x):
+            d = DefaultDict()
+            d["a"] = 1
+            val = operator.getitem(d, "b")
+            return x + val
+
+        x = torch.randn(4)
+        self.assertEqual(fn(x), self._compile(fn, x))
+
+    def test_collections_counter_getitem(self):
+        def fn(x):
+            c = collections.Counter({"a": 1, "b": 2})
+            return x + operator.getitem(c, "a")
+
+        x = torch.randn(4)
+        self.assertEqual(fn(x), self._compile(fn, x))
+
+    # --- UserDefined* with overridden __getitem__ ---
+
+    def test_user_defined_dict_overridden_getitem(self):
+        """Dict subclass with custom __getitem__ should NOT delegate to _base_vt."""
+
+        class MyDict(dict):
+            def __getitem__(self, key):
+                return super().__getitem__(key) + 100
+
+        def fn(x):
+            d = MyDict(a=1, b=2)
+            return x + operator.getitem(d, "a")
+
+        x = torch.randn(4)
+        self.assertEqual(fn(x), self._compile(fn, x))
+
+    def test_user_defined_list_overridden_getitem(self):
+        """List subclass with custom __getitem__ should NOT delegate to _base_vt."""
+
+        class MyList(list):
+            def __getitem__(self, key):
+                return super().__getitem__(key) * 2
+
+        def fn(x):
+            items = MyList([x, x + 1, x + 2])
+            return operator.getitem(items, 1)
+
+        x = torch.randn(4)
+        self.assertEqual(fn(x), self._compile(fn, x))
+
+    def test_counter_missing_key(self):
+        """Counter.__missing__ returns 0 for missing keys."""
+
+        def fn(x):
+            c = collections.Counter({"a": 1, "b": 2})
+            return x + operator.getitem(c, "missing")
+
+        x = torch.randn(4)
+        self.assertEqual(fn(x), self._compile(fn, x))
+
+    # --- TorchScriptObjectVariable ---
+
+    def test_opaque_object_getitem(self):
+        from torch._library.opaque_object import (
+            MemberType,
+            OpaqueBase,
+            register_opaque_type,
+        )
+
+        class OpaqueScaler(OpaqueBase):
+            def __init__(self, scale):
+                self.scale = scale
+
+            def apply(self, x):
+                return x * self.scale
+
+        class OpaqueContainer(OpaqueBase):
+            def __init__(self, items):
+                self.items = items
+
+            def __getitem__(self, idx):
+                return self.items[idx]
+
+        register_opaque_type(
+            OpaqueScaler,
+            typ="reference",
+            members={
+                "scale": MemberType.USE_REAL,
+                "apply": MemberType.INLINED,
+            },
+        )
+        register_opaque_type(
+            OpaqueContainer,
+            typ="reference",
+            members={
+                "items": MemberType.USE_REAL,
+                "__getitem__": MemberType.INLINED,
+            },
+        )
+
+        def fn(x, c):
+            scaler = operator.getitem(c, 0)
+            return scaler.apply(x)
+
+        x = torch.randn(4)
+        c = OpaqueContainer([OpaqueScaler(2.0), OpaqueScaler(3.0)])
+        compiled = torch.compile(fn, backend="eager", fullgraph=True)
+        self.assertEqual(fn(x, c), compiled(x, c))
+
+    # --- TritonKernelVariable ---
+
+    @unittest.skipUnless(HAS_GPU and HAS_CUDA_AND_TRITON, "requires gpu and triton")
+    def test_triton_kernel_getitem_grid(self):
+        from torch.testing._internal.triton_utils import add_kernel
+
+        def fn(x, y):
+            output = torch.zeros_like(x)
+            n_elements = output.numel()
+            grid = (n_elements // 256,)
+            bound = operator.getitem(add_kernel, grid)
+            bound(x, y, output, n_elements, BLOCK_SIZE=256)
+            return output
+
+        x = torch.randn(256, device="cuda")
+        y = torch.randn(256, device="cuda")
+        compiled = torch.compile(fn, backend="eager", fullgraph=True)
+        self.assertEqual(fn(x, y), compiled(x, y))
+
+    # ===================================================================
+    # CPython behavioral gaps — expectedFailure until implemented
+    # ===================================================================
+
+    # GAP 1: deque has only sq_item (int index), no mp_subscript.
+    # CPython: deque[slice] → TypeError "sequence index must be integer, not 'slice'"
+    # Dynamo: DequeVariable inherits BaseListVariable.mp_subscript_impl which accepts slices.
+    # TODO: DequeVariable should override mp_subscript_impl to reject slices, matching
+    # CPython's deque which only has sq_item (Modules/_collectionsmodule.c:1888).
+    @unittest.expectedFailure
+    def test_deque_slice_should_reject(self):
+        """deque does not support slicing in CPython — only sq_item (int index)."""
+
+        def fn(x):
+            d = collections.deque([x, x + 1, x + 2])
+            return operator.getitem(d, slice(0, 2))
+
+        x = torch.randn(4)
+        with self.assertRaises(TypeError):
+            self._compile(fn, x)
+
+    # TODO: deque int index works but through the wrong dispatch path.
+    # CPython: PyObject_GetItem Branch 2 → _PyIndex_Check(key) → PyNumber_AsSsize_t → sq_item.
+    # Dynamo: inherited BaseListVariable.mp_subscript_impl (Branch 1 path).
+    # Result is correct, dispatch path diverges. Fix when sq_item branch is implemented.
+    def test_deque_int_index(self):
+        def fn(x):
+            d = collections.deque([x, x + 1, x + 2])
+            return operator.getitem(d, 1)
+
+        x = torch.randn(4)
+        self.assertEqual(fn(x), self._compile(fn, x))
+
+    # GAP 2: dict_subscript calls _PyObject_HashFast → TypeError for unhashable keys.
+    # TODO: ConstDictVariable.mp_subscript_impl should check tp_hash and raise TypeError
+    # for unhashable keys, matching CPython's dict_subscript (Objects/dictobject.c:3680).
+    @unittest.expectedFailure
+    def test_dict_unhashable_key(self):
+        """dict[unhashable] should raise TypeError, not KeyError or silent failure."""
+
+        def fn(x):
+            d = {0: x, 1: x + 1}
+            return operator.getitem(d, [0])
+
+        x = torch.randn(4)
+        with self.assertRaises(TypeError):
+            self._compile(fn, x)
+
+    # TODO: str/bytes subscript works via constant fold fallback (base mp_subscript_impl
+    # raises Unsupported → _make_handler → operator.getitem("hello", 0) evaluates at
+    # Python level), not via a proper mp_subscript_impl override mirroring CPython's
+    # unicode_subscript / bytes_subscript. Should add dedicated overrides on
+    # ConstantVariable to match CPython's dispatch path.
+    # CPython: Objects/unicodeobject.c:13809 (unicode_subscript)
+    # CPython: Objects/bytesobject.c (bytes_subscript)
+
+    def test_str_subscript(self):
+        def fn(x):
+            s = "hello"
+            c = operator.getitem(s, 0)
+            return x + len(c)
+
+        x = torch.randn(4)
+        self.assertEqual(fn(x), self._compile(fn, x))
+
+    def test_bytes_subscript(self):
+        def fn(x):
+            b = b"hello"
+            val = operator.getitem(b, 0)
+            return x + val
+
+        x = torch.randn(4)
+        self.assertEqual(fn(x), self._compile(fn, x))
+
+    # ===================================================================
+    # Explicit __getitem__ dunder call path tests
+    # Exercises: obj.__getitem__(key) → LOAD_ATTR + CALL, which may
+    # route through call_method → mp_subscript_impl rather than
+    # vt_getitem → mp_subscript_impl.
+    # ===================================================================
+
+    def test_list_dunder_getitem(self):
+        def fn(x):
+            items = [x, x + 1, x + 2]
+            return items.__getitem__(1)
+
+        x = torch.randn(4)
+        self.assertEqual(fn(x), self._compile(fn, x))
+
+    def test_dict_dunder_getitem(self):
+        def fn(x):
+            d = {"a": x, "b": x + 1}
+            return d.__getitem__("a")
+
+        x = torch.randn(4)
+        self.assertEqual(fn(x), self._compile(fn, x))
+
+    def test_user_defined_dict_missing_dunder_getitem(self):
+        """__missing__ fallback must work via __getitem__ method call, not just operator.getitem."""
+
+        class MyDict(dict):
+            def __missing__(self, key):
+                return 42
+
+        def fn(x):
+            d = MyDict(a=1)
+            return x + d.__getitem__("b")
+
+        x = torch.randn(4)
+        self.assertEqual(fn(x), self._compile(fn, x))
+
+
+if __name__ == "__main__":
+    from torch._dynamo.test_case import run_tests
+
+    run_tests()
diff --git a/test/dynamo/test_global.py b/test/dynamo/test_global.py
index 74f50de6c13d9..119d56d674e24 100644
--- a/test/dynamo/test_global.py
+++ b/test/dynamo/test_global.py
@@ -1,5 +1,4 @@
 # Owner(s): ["module: dynamo"]
-from typing import Optional
 
 import torch
 import torch._dynamo.test_case
@@ -185,7 +184,7 @@ def fn(x):
     def test_store_global_inline_1(self):
         # Borrowed from test_python_autograd.py
         class Variable:
-            def __init__(self, value: torch.Tensor, name: Optional[str] = None):
+            def __init__(self, value: torch.Tensor, name: str | None = None):
                 self.value = value
                 self.name = name or fresh_name()
 
@@ -205,12 +204,12 @@ def fn(a, b):
     def test_store_global_inline_2(self):
         # Borrowed from test_python_autograd.py
         class Variable:
-            def __init__(self, value: torch.Tensor, name: Optional[str] = None):
+            def __init__(self, value: torch.Tensor, name: str | None = None):
                 self.value = value
                 self.name = name or fresh_name()
 
             @staticmethod
-            def constant(value: torch.Tensor, name: Optional[str] = None):
+            def constant(value: torch.Tensor, name: str | None = None):
                 return Variable(value, name)
 
         def fn(a, b):
@@ -232,13 +231,13 @@ def test_store_global_crossfile_inline(self):
         except ImportError:
             import mock_store_global_crossfile_inline
 
-        @torch.compile()
+        @torch.compile(backend="eager")
         def fn(x):
             mock_store_global_crossfile_inline.set_flag_true()
             mock_store_global_crossfile_inline.set_flag_false()
             return x + 1
 
-        @torch.compile()
+        @torch.compile(backend="eager")
         def fn_set_true(x):
             mock_store_global_crossfile_inline.set_flag_true()
             return x + 1
diff --git a/test/dynamo/test_graph_deduplication.py b/test/dynamo/test_graph_deduplication.py
index d0e712ffaa6cf..03daf70c5733f 100644
--- a/test/dynamo/test_graph_deduplication.py
+++ b/test/dynamo/test_graph_deduplication.py
@@ -620,7 +620,7 @@ def fn(x, y):
         add_node.args = (args[0], add_2)
         self.assertExpectedInline(
             _detect_cycles(mod.graph, {add_2: OrderedSet([add_2])}),
-            """cycle detected in path: deque([output, add_2, add_2])""",
+            """cycle detected in path: [output, add_2, sum_1, add, add_2]""",
         )
 
     def test_cycle_detection_two_node(self):
@@ -644,7 +644,7 @@ def fn(x, y):
                 mod.graph,
                 {add_2: OrderedSet([add_node]), add_node: OrderedSet([add_2])},
             ),
-            """cycle detected in path: deque([output, add_2, add, add_2])""",
+            """cycle detected in path: [output, add_2, sum_1, add, add_2]""",
         )
 
     def test_cycle_detection_arg_and_additional_deps(self):
@@ -665,7 +665,7 @@ def fn(x, y):
         add_node.args = (args[0], add_2)
         self.assertExpectedInline(
             _detect_cycles(mod.graph, {add_2: OrderedSet([add_node])}),
-            """cycle detected in path: deque([output, add_2, add, add_2])""",
+            """cycle detected in path: [output, add_2, sum_1, add, add_2]""",
         )
 
     def test_cycle_detection_simple(self):
@@ -676,7 +676,7 @@ def test_cycle_detection_simple(self):
         add_node.args = (args[0], add_2)
         self.assertExpectedInline(
             _detect_cycles(mod.graph, {}),
-            """cycle detected in path: deque([output, add_2, sum_1, add, add_2])""",
+            """cycle detected in path: [output, add_2, sum_1, add, add_2]""",
         )
 
     def test_cycle_detection_complex(self):
@@ -711,7 +711,7 @@ def fn(x, y):
         invoke_subgraph_node.args = (add_2, args[1])
         self.assertExpectedInline(
             _detect_cycles(mod.graph, {}),
-            """cycle detected in path: deque([output, add_2, add_1, sum_1, getitem, invoke_subgraph, add_2])""",
+            """cycle detected in path: [output, add_2, add_1, sum_1, getitem, invoke_subgraph, add_2]""",
         )
 
     def test_autocast_ordering(self):
@@ -1115,7 +1115,7 @@ def fn(x0, x1, x2, y0, y1, y2):
             x2 = inner_fn(x2, y2)
             return x0, x1, x2
 
-        fn_opt = torch.compile(fn, fullgraph=True)
+        fn_opt = torch.compile(fn, fullgraph=True, backend="eager")
         inps = [torch.rand(10, 10) for _ in range(6)]
         result_compiled = fn_opt(*inps)
         result_eager = fn(*inps)
@@ -1207,7 +1207,7 @@ def fn(x0, x1, x2, y0, y1, y2):
             x2 = inner_fn(x2, y2)
             return x0.sum() + x1.sum() + x2.sum()
 
-        fn_opt = torch.compile(fn, fullgraph=True)
+        fn_opt = torch.compile(fn, fullgraph=True, backend="eager")
         args = [torch.rand(10, 10) for _ in range(6)]
         for arg in args:
             torch._dynamo.mark_static_address(arg)
diff --git a/test/dynamo/test_guard_exclusion.py b/test/dynamo/test_guard_exclusion.py
new file mode 100644
index 0000000000000..cfa84c7e302d2
--- /dev/null
+++ b/test/dynamo/test_guard_exclusion.py
@@ -0,0 +1,566 @@
+# Owner(s): ["module: dynamo"]
+import torch
+import torch._dynamo
+import torch._dynamo.testing
+from torch.testing._internal.common_utils import run_tests, skipIfTorchDynamo, TestCase
+
+
+class GraphTracker:
+    """Backend that tracks which compiled graph (by compilation order) handles each call."""
+
+    def __init__(self):
+        self.graphs = []
+        self.call_log = []
+
+    def __call__(self, gm, example_inputs):
+        graph_id = len(self.graphs)
+        self.graphs.append(gm)
+
+        def wrapper(*args, **kwargs):
+            self.call_log.append(graph_id)
+            return gm.forward(*args, **kwargs)
+
+        return wrapper
+
+    @property
+    def frame_count(self):
+        return len(self.graphs)
+
+    def reset(self):
+        self.graphs.clear()
+        self.call_log.clear()
+
+
+@skipIfTorchDynamo("uses custom backend incompatible with PYTORCH_TEST_WITH_DYNAMO")
+@torch._dynamo.config.patch(automatic_dynamic_exclusion_guard=True)
+class TestGuardExclusion(TestCase):
+    def setUp(self):
+        super().setUp()
+        torch._dynamo.reset()
+
+    def tearDown(self):
+        super().tearDown()
+        torch._dynamo.reset()
+
+    @torch._dynamo.config.patch(
+        automatic_dynamic_shapes=True, assume_static_by_default=True
+    )
+    def test_automatic_dynamic_exclusive_guard_basic(self):
+        """
+        1. [3, 4] -> Graph 0 (static)
+        2. [5, 4] -> Graph 1 (dim 0 dynamic), exclusion rejects dim0==3
+        3. [7, 4] -> Graph 1 (reuse dynamic graph)
+        4. [3, 4] -> Graph 0 (exclusion triggers, reverts to static)
+        """
+
+        def foo(x):
+            return x * 2
+
+        tracker = GraphTracker()
+        opt = torch.compile(foo, backend=tracker)
+
+        x1 = torch.randn(3, 4)
+        result1 = opt(x1)
+        self.assertEqual(tracker.frame_count, 1)
+        self.assertEqual(tracker.call_log[-1], 0)
+
+        x2 = torch.randn(5, 4)
+        result2 = opt(x2)
+        self.assertEqual(tracker.frame_count, 2)
+        self.assertEqual(tracker.call_log[-1], 1)
+
+        # dynamic graph reuse
+        opt(torch.randn(7, 4))
+        self.assertEqual(tracker.frame_count, 2)
+        self.assertEqual(tracker.call_log[-1], 1)
+
+        # original shape reverts to Graph 0
+        x3 = torch.randn(3, 4)
+        result3 = opt(x3)
+        self.assertEqual(tracker.frame_count, 2)
+        self.assertEqual(tracker.call_log[-1], 0)
+
+        self.assertEqual(result1, x1 * 2)
+        self.assertEqual(result2, x2 * 2)
+        self.assertEqual(result3, x3 * 2)
+
+    @torch._dynamo.config.patch(
+        automatic_dynamic_shapes=True, assume_static_by_default=True
+    )
+    def test_accumulated_exclusion_does_not_shadow_intermediate_graph(self):
+        """
+        Tensor accumulation: dims become dynamic one at a time.
+        1. func(3, 4)  -> Graph 0: static (3, 4)
+        2. func(5, 4)  -> Graph 1: (s0, 4), excluded dim0=3
+        3. func(3, 19) -> Graph 2: (s0, s1), excluded dim1=4
+                          (dim0's exclusion is cleared since no dim transitioned)
+        4. func(5, 4)  -> should use Graph 1, not Graph 2
+        """
+
+        def foo(x):
+            return x * 2
+
+        tracker = GraphTracker()
+        opt = torch.compile(foo, backend=tracker)
+
+        # Call 1: shape [3, 4] -> compiles Graph 0 (static)
+        opt(torch.randn(3, 4))
+        self.assertEqual(tracker.frame_count, 1)
+        self.assertEqual(tracker.call_log[-1], 0)
+
+        # Call 2: shape [5, 4] -> compiles Graph 1 (s0, 4)
+        opt(torch.randn(5, 4))
+        self.assertEqual(tracker.frame_count, 2)
+        self.assertEqual(tracker.call_log[-1], 1)
+
+        # Call 3: shape [3, 19] -> Graph 1 exclusion rejects (size(0)==3),
+        # Graph 0 rejects (19!=4), recompiles Graph 2 (s0, s1)
+        opt(torch.randn(3, 19))
+        self.assertEqual(tracker.frame_count, 3)
+        self.assertEqual(tracker.call_log[-1], 2)
+
+        # Call 4: shape [5, 4] -> should still use Graph 1 (s0, 4).
+        # Graph 2's exclusion is dim1=4, so (5, 4) is rejected and
+        # falls through to Graph 1.
+        opt(torch.randn(5, 4))
+
+        self.assertEqual(
+            tracker.call_log[-1],
+            1,
+            "Input [5,4] should use Graph 1 (s0, 4), not Graph 2 (s0, s1). "
+            "Graph 2's exclusion must reject size(1)==4 independently, not "
+            "require size(0)==3 AND size(1)==4.",
+        )
+
+        # Call 5: shape [3, 4] -> should still use Graph 0 (static)
+        opt(torch.randn(3, 4))
+        self.assertEqual(
+            tracker.call_log[-1],
+            0,
+            "Input [3,4] should use Graph 0 (static)",
+        )
+
+    @torch._dynamo.config.patch(
+        automatic_dynamic_shapes=True, assume_static_by_default=True
+    )
+    def test_4d_progressive_dynamism_cascading(self):
+        """
+        4D tensor where dims become dynamic one at a time across recompilations.
+        Each new graph is more general, and exclusion guards ensure inputs cascade
+        to the most specialized graph.
+
+        Graph 0: (2, 3, 4, 5) static
+        Graph 1: (dyn, 3, 4, 5) excluded dim0=2
+        Graph 2: (dyn, dyn, 4, 5) excluded dim1=3
+        Graph 3: (dyn, dyn, dyn, 5) excluded dim2=4
+        """
+
+        def foo(x):
+            return x.sum()
+
+        tracker = GraphTracker()
+        opt = torch.compile(foo, backend=tracker)
+
+        # Graph 0: (2, 3, 4, 5) static
+        opt(torch.randn(2, 3, 4, 5))
+        self.assertEqual(tracker.frame_count, 1)
+
+        # Graph 1: dim0 changes -> (dyn, 3, 4, 5)
+        opt(torch.randn(7, 3, 4, 5))
+        self.assertEqual(tracker.frame_count, 2)
+
+        # Graph 2: dim1 also changes -> (dyn, dyn, 4, 5)
+        # Input (7, 8, 4, 5): Graph 1 rejects dim1=8≠3, Graph 0 rejects dim0=7≠2
+        opt(torch.randn(7, 8, 4, 5))
+        self.assertEqual(tracker.frame_count, 3)
+
+        # Graph 3: dim2 also changes -> (dyn, dyn, dyn, 5)
+        # Input (7, 8, 9, 5): Graph 2 rejects dim2=9≠4
+        opt(torch.randn(7, 8, 9, 5))
+        self.assertEqual(tracker.frame_count, 4)
+
+        # Now verify cascading: each input routes to the most specialized graph.
+        # (2, 3, 4, 5) -> Graph 0 (static, most specialized)
+        opt(torch.randn(2, 3, 4, 5))
+        self.assertEqual(tracker.call_log[-1], 0, "(2,3,4,5) -> Graph 0 (static)")
+
+        # (7, 3, 4, 5) -> Graph 1 (dyn, 3, 4, 5)
+        opt(torch.randn(7, 3, 4, 5))
+        self.assertEqual(tracker.call_log[-1], 1, "(7,3,4,5) -> Graph 1 (dyn,3,4,5)")
+
+        # (7, 8, 4, 5) -> Graph 2 (dyn, dyn, 4, 5)
+        opt(torch.randn(7, 8, 4, 5))
+        self.assertEqual(tracker.call_log[-1], 2, "(7,8,4,5) -> Graph 2 (dyn,dyn,4,5)")
+
+        # (7, 8, 9, 5) -> Graph 3 (dyn, dyn, dyn, 5)
+        opt(torch.randn(7, 8, 9, 5))
+        self.assertEqual(
+            tracker.call_log[-1], 3, "(7,8,9,5) -> Graph 3 (dyn,dyn,dyn,5)"
+        )
+
+        # (20, 30, 40, 5) -> Graph 3 (most general, no exclusion hit)
+        opt(torch.randn(20, 30, 40, 5))
+        self.assertEqual(
+            tracker.call_log[-1], 3, "(20,30,40,5) -> Graph 3 (most general)"
+        )
+
+        self.assertEqual(tracker.frame_count, 4, "No additional recompilations")
+
+    @torch._dynamo.config.patch(
+        automatic_dynamic_shapes=True, assume_static_by_default=True
+    )
+    def test_5d_two_rounds_of_dynamism(self):
+        """
+        5D tensor with two rounds of automatic_dynamic. Verify inputs route
+        to the most specialized graph after each round.
+
+        Graph 0: (2, 3, 4, 5, 6) static
+        Graph 1: (dyn, 3, 4, 5, 6) excluded dim0=2
+        Graph 2: (dyn, 3, dyn, 5, 6) excluded dim2=4
+        """
+
+        def foo(x):
+            return x * 2
+
+        tracker = GraphTracker()
+        opt = torch.compile(foo, backend=tracker)
+
+        # Graph 0: static
+        opt(torch.randn(2, 3, 4, 5, 6))
+        self.assertEqual(tracker.frame_count, 1)
+
+        # Graph 1: dim0 becomes dynamic
+        opt(torch.randn(8, 3, 4, 5, 6))
+        self.assertEqual(tracker.frame_count, 2)
+
+        # Graph 2: dim2 also becomes dynamic
+        opt(torch.randn(8, 3, 10, 5, 6))
+        self.assertEqual(tracker.frame_count, 3)
+
+        # Verify routing:
+        # Original static shape -> Graph 0
+        opt(torch.randn(2, 3, 4, 5, 6))
+        self.assertEqual(tracker.call_log[-1], 0, "Original -> Graph 0")
+
+        # dim0 differs, dim2 matches static -> Graph 1
+        opt(torch.randn(9, 3, 4, 5, 6))
+        self.assertEqual(tracker.call_log[-1], 1, "dim0 changed -> Graph 1")
+
+        # dim0 differs, dim2 differs -> Graph 2
+        opt(torch.randn(9, 3, 11, 5, 6))
+        self.assertEqual(tracker.call_log[-1], 2, "dim0+dim2 changed -> Graph 2")
+
+        # dim0 is original excluded value, dim2 differs -> Graph 2 should still
+        # accept because dim0's exclusion is None (already dynamic when snapshot taken)
+        opt(torch.randn(2, 3, 11, 5, 6))
+        self.assertEqual(
+            tracker.call_log[-1],
+            2,
+            "dim0=2 with dim2≠4 -> Graph 2 (dim0 exclusion is None)",
+        )
+
+        self.assertEqual(tracker.frame_count, 3, "No additional recompilations")
+
+    @torch._dynamo.config.patch(
+        automatic_dynamic_shapes=True, assume_static_by_default=True
+    )
+    def test_many_entries_wrong_graph_selection(self):
+        """
+        Convoluted scenario: 4D tensor, three rounds of dynamism creating 4 graphs.
+        Without exclusion guards, the most general graph would shadow all others.
+        Test that each input gets the best (most specialized) match.
+
+        Graph 0: (2, 3, 4, 5) static
+        Graph 1: (dyn, 3, 4, 5) excluded dim0=2
+        After Graph 1, (2, 8, 4, 5) triggers dim1 dynamic:
+        Graph 2: (dyn, dyn, 4, 5) excluded dim1=3
+        After Graph 2, (2, 8, 9, 5) triggers dim2 dynamic:
+        Graph 3: (dyn, dyn, dyn, 5) excluded dim2=4
+
+        Key: Graph 3 should NOT steal inputs that belong to Graph 0, 1, or 2.
+        """
+
+        def foo(x):
+            return x.relu()
+
+        tracker = GraphTracker()
+        opt = torch.compile(foo, backend=tracker)
+
+        # Build up 4 graphs progressively
+        opt(torch.randn(2, 3, 4, 5))  # Graph 0
+        opt(torch.randn(7, 3, 4, 5))  # Graph 1: dim0 dynamic
+        opt(torch.randn(7, 8, 4, 5))  # Graph 2: dim1 also dynamic
+        opt(torch.randn(7, 8, 9, 5))  # Graph 3: dim2 also dynamic
+        self.assertEqual(tracker.frame_count, 4)
+
+        # Now stress-test routing with various inputs:
+        test_cases = [
+            # (shape, expected_graph, description)
+            ((2, 3, 4, 5), 0, "exact original -> static Graph 0"),
+            ((7, 3, 4, 5), 1, "dim0 differs -> Graph 1"),
+            ((99, 3, 4, 5), 1, "dim0 differs (large) -> Graph 1"),
+            ((7, 8, 4, 5), 2, "dim0+dim1 differ -> Graph 2"),
+            ((7, 99, 4, 5), 2, "dim0+dim1 differ (large) -> Graph 2"),
+            ((7, 8, 9, 5), 3, "dim0+dim1+dim2 differ -> Graph 3"),
+            ((99, 99, 99, 5), 3, "all non-static dims differ -> Graph 3"),
+        ]
+
+        for shape, expected_graph, desc in test_cases:
+            opt(torch.randn(*shape))
+            self.assertEqual(tracker.call_log[-1], expected_graph, desc)
+
+        self.assertEqual(tracker.frame_count, 4, "No additional recompilations")
+
+    @torch._dynamo.config.patch(
+        automatic_dynamic_shapes=True, assume_static_by_default=True
+    )
+    def test_multi_dim_dynamic_and_semantics(self):
+        """
+        When multiple dims become dynamic at once, AND semantics is critical.
+        Graph 0: (4, 3, 234, 5) static
+        Graph 1: (s0, 3, s2, s3) dynamic on dims 0,2,3. excluded=(4, _, 234, 5)
+
+        OR semantics (wrong): rejects (4, 3, 100, 20) because dim0==4 matches.
+        AND semantics (correct): accepts (4, 3, 100, 20) because not ALL excluded
+        dims match (dim2=100≠234 and dim3=20≠5).
+        """
+
+        def foo(x):
+            return x * 2
+
+        tracker = GraphTracker()
+        opt = torch.compile(foo, backend=tracker)
+
+        opt(torch.randn(4, 3, 234, 5))  # Graph 0: static
+        opt(torch.randn(10, 3, 100, 20))  # Graph 1: dims 0,2,3 dynamic
+        self.assertEqual(tracker.frame_count, 2)
+
+        # Only the exact original shape should be excluded from Graph 1
+        opt(torch.randn(4, 3, 234, 5))
+        self.assertEqual(tracker.call_log[-1], 0, "Exact original -> Graph 0")
+
+        # Partial matches should NOT be excluded (AND semantics)
+        opt(torch.randn(4, 3, 100, 20))  # dim0=4 matches, dims 2,3 don't
+        self.assertEqual(tracker.call_log[-1], 1, "dim0=4 partial match -> Graph 1")
+
+        opt(torch.randn(10, 3, 234, 20))  # dim2=234 matches, dims 0,3 don't
+        self.assertEqual(tracker.call_log[-1], 1, "dim2=234 partial match -> Graph 1")
+
+        opt(torch.randn(10, 3, 234, 5))  # dim2=234, dim3=5 match, dim0 doesn't
+        self.assertEqual(tracker.call_log[-1], 1, "dim2+dim3 partial match -> Graph 1")
+
+        opt(torch.randn(4, 3, 234, 20))  # dim0=4, dim2=234 match, dim3 doesn't
+        self.assertEqual(tracker.call_log[-1], 1, "dim0+dim2 partial match -> Graph 1")
+
+        # Totally new shape, no exclusion hit
+        opt(torch.randn(99, 3, 88, 77))
+        self.assertEqual(tracker.call_log[-1], 1, "New shape -> Graph 1")
+
+        self.assertEqual(tracker.frame_count, 2, "No additional recompilations")
+
+    @torch._dynamo.config.patch(
+        automatic_dynamic_shapes=True, assume_static_by_default=True
+    )
+    def test_integer_input_exclusion_basic(self):
+        """
+        Integer inputs that become dynamic should also get exclusion guards.
+        1. foo(x, 3) -> Graph 0: static n=3
+        2. foo(x, 5) -> Graph 1: dynamic n, excluded should reject n==3
+        3. foo(x, 3) -> should use Graph 0 (static), not Graph 1
+        """
+
+        def foo(x, n):
+            return x * n
+
+        tracker = GraphTracker()
+        opt = torch.compile(foo, backend=tracker)
+
+        x = torch.randn(4)
+
+        opt(x, 3)
+        self.assertEqual(tracker.frame_count, 1)
+        self.assertEqual(tracker.call_log[-1], 0)
+
+        opt(x, 5)
+        self.assertEqual(tracker.frame_count, 2)
+        self.assertEqual(tracker.call_log[-1], 1)
+
+        opt(x, 3)
+        self.assertEqual(
+            tracker.call_log[-1],
+            0,
+            "Input n=3 should use Graph 0 (static), not Graph 1 (dynamic n).",
+        )
+
+    @torch._dynamo.config.patch(
+        automatic_dynamic_shapes=True, assume_static_by_default=True
+    )
+    def test_integer_input_exclusion_accumulation(self):
+        """
+        Same accumulation scenario as the tensor test but with integer inputs.
+        1. foo(x, 3, 4)  -> Graph 0: static (3, 4)
+        2. foo(x, 5, 4)  -> Graph 1: dynamic (s0, 4), exclusion rejects n0==3
+        3. foo(x, 3, 19) -> Graph 2: dynamic (s0, s1), exclusion should reject
+                            n1==4 independently, not require n0==3 AND n1==4
+        4. foo(x, 5, 4)  -> should use Graph 1, not Graph 2
+        """
+
+        def foo(x, n, m):
+            return x * n + m
+
+        tracker = GraphTracker()
+        opt = torch.compile(foo, backend=tracker)
+
+        x = torch.randn(4)
+
+        opt(x, 3, 4)
+        self.assertEqual(tracker.frame_count, 1)
+        self.assertEqual(tracker.call_log[-1], 0)
+
+        opt(x, 5, 4)
+        self.assertEqual(tracker.frame_count, 2)
+        self.assertEqual(tracker.call_log[-1], 1)
+
+        opt(x, 3, 19)
+        self.assertEqual(tracker.frame_count, 3)
+        self.assertEqual(tracker.call_log[-1], 2)
+
+        opt(x, 5, 4)
+        self.assertEqual(
+            tracker.call_log[-1],
+            1,
+            "Input (5, 4) should use Graph 1 (s0, 4), not Graph 2 (s0, s1).",
+        )
+
+        opt(x, 3, 4)
+        self.assertEqual(
+            tracker.call_log[-1],
+            0,
+            "Input (3, 4) should use Graph 0 (static)",
+        )
+
+    @torch._dynamo.config.patch(
+        automatic_dynamic_shapes=True, assume_static_by_default=True
+    )
+    def test_two_tensor_inputs_exclusion(self):
+        """
+        Multi-tensor exclusion: all pairs are flattened and guarded with
+        not-all semantics. The guard rejects only when ALL excluded values
+        across ALL tensors match simultaneously.
+
+        1. foo(x=[3,4], y=[5,6])   -> Graph 0: all static
+        2. foo(x=[3,10], y=[5,11]) -> Graph 1: x.dim1, y.dim1 dynamic
+                                      exclusion: Or(x.dim1!=4, y.dim1!=6)
+        3. foo(x=[3,10], y=[5,6])  -> Graph 1 (not all match -> passes)
+        4. foo(x=[3,4], y=[5,21])  -> Graph 1 (not all match -> passes)
+        5. foo(x=[3,4], y=[5,6])   -> Graph 0 (all match -> rejected)
+        6. foo(x=[3,10], y=[5,11]) -> Graph 1
+        """
+
+        def foo(x, y):
+            return x.sum() + y.sum()
+
+        tracker = GraphTracker()
+        opt = torch.compile(foo, backend=tracker)
+
+        opt(torch.randn(3, 4), torch.randn(5, 6))
+        self.assertEqual(tracker.frame_count, 1)
+        self.assertEqual(tracker.call_log[-1], 0)
+
+        opt(torch.randn(3, 10), torch.randn(5, 11))
+        self.assertEqual(tracker.frame_count, 2)
+        self.assertEqual(tracker.call_log[-1], 1)
+
+        # Only y.dim1 matches excluded value; combined Or guard passes.
+        opt(torch.randn(3, 10), torch.randn(5, 6))
+        self.assertEqual(tracker.frame_count, 2, "Should not recompile")
+        self.assertEqual(tracker.call_log[-1], 1)
+
+        # Only x.dim1 matches excluded value; combined Or guard passes.
+        opt(torch.randn(3, 4), torch.randn(5, 21))
+        self.assertEqual(tracker.frame_count, 2, "Should not recompile")
+        self.assertEqual(tracker.call_log[-1], 1)
+
+        # Both match excluded values; Or guard fails -> falls to Graph 0.
+        opt(torch.randn(3, 4), torch.randn(5, 6))
+        self.assertEqual(tracker.call_log[-1], 0)
+
+        # Neither matches; Or guard passes -> Graph 1.
+        opt(torch.randn(3, 10), torch.randn(5, 11))
+        self.assertEqual(tracker.call_log[-1], 1)
+
+    @torch._dynamo.config.patch(
+        automatic_dynamic_shapes=True, assume_static_by_default=True
+    )
+    def test_multi_tensor_and_scalar_accumulation(self):
+        """
+        3-dim tensors + 3 scalar inputs with cascading accumulation.
+        Each step transitions one input while the rest stay the same,
+        verifying that only the current transition's exclusion is emitted.
+
+        Graph 0: all static
+        Graph 1: x.dim2, y.dim2 dynamic    excl: Or(Ne(x.dim2,4), Ne(y.dim2,7))
+        Graph 2: + n dynamic                excl: Ne(n, 2)
+        Graph 3: + m dynamic                excl: Ne(m, 3)
+        Graph 4: + k dynamic                excl: Ne(k, 4)
+        """
+
+        def foo(x, y, n, m, k):
+            return x.sum() * n + y.sum() * m + k
+
+        tracker = GraphTracker()
+        opt = torch.compile(foo, backend=tracker)
+
+        # -- Compilation steps --
+
+        # Graph 0: all static
+        opt(torch.randn(2, 3, 4), torch.randn(5, 6, 7), 2, 3, 4)
+        self.assertEqual(tracker.frame_count, 1)
+        self.assertEqual(tracker.call_log[-1], 0)
+
+        # Graph 1: tensor dim2 changes
+        opt(torch.randn(2, 3, 10), torch.randn(5, 6, 11), 2, 3, 4)
+        self.assertEqual(tracker.frame_count, 2)
+        self.assertEqual(tracker.call_log[-1], 1)
+
+        # Graph 2: scalar n changes (tensor exclusions cleared)
+        opt(torch.randn(2, 3, 10), torch.randn(5, 6, 11), 8, 3, 4)
+        self.assertEqual(tracker.frame_count, 3)
+        self.assertEqual(tracker.call_log[-1], 2)
+
+        # Graph 3: scalar m changes (n exclusion cleared)
+        opt(torch.randn(2, 3, 10), torch.randn(5, 6, 11), 8, 9, 4)
+        self.assertEqual(tracker.frame_count, 4)
+        self.assertEqual(tracker.call_log[-1], 3)
+
+        # Graph 4: scalar k changes (m exclusion cleared)
+        opt(torch.randn(2, 3, 10), torch.randn(5, 6, 11), 8, 9, 15)
+        self.assertEqual(tracker.frame_count, 5)
+        self.assertEqual(tracker.call_log[-1], 4)
+
+        # -- Verification: each input routes to the correct graph --
+
+        # k=4 triggers Graph 4 exclusion -> Graph 3
+        opt(torch.randn(2, 3, 10), torch.randn(5, 6, 11), 8, 9, 4)
+        self.assertEqual(tracker.call_log[-1], 3, "k=4 should fall to Graph 3")
+
+        # m=3 also triggers Graph 3 exclusion -> Graph 2
+        opt(torch.randn(2, 3, 10), torch.randn(5, 6, 11), 8, 3, 4)
+        self.assertEqual(tracker.call_log[-1], 2, "m=3 should fall to Graph 2")
+
+        # n=2 also triggers Graph 2 exclusion -> Graph 1
+        opt(torch.randn(2, 3, 10), torch.randn(5, 6, 11), 2, 3, 4)
+        self.assertEqual(tracker.call_log[-1], 1, "n=2 should fall to Graph 1")
+
+        # tensor dims match original -> Graph 1 exclusion triggers -> Graph 0
+        opt(torch.randn(2, 3, 4), torch.randn(5, 6, 7), 2, 3, 4)
+        self.assertEqual(tracker.call_log[-1], 0, "Original sizes should use Graph 0")
+
+        # mixed: new tensor dims + new scalars -> Graph 4
+        opt(torch.randn(2, 3, 20), torch.randn(5, 6, 21), 50, 60, 70)
+        self.assertEqual(tracker.frame_count, 5, "Should not recompile")
+        self.assertEqual(tracker.call_log[-1], 4, "All-new values should use Graph 4")
+
+
+if __name__ == "__main__":
+    run_tests()
diff --git a/test/dynamo/test_guard_manager.py b/test/dynamo/test_guard_manager.py
index 1849950644fe4..73769c97f10f1 100644
--- a/test/dynamo/test_guard_manager.py
+++ b/test/dynamo/test_guard_manager.py
@@ -1066,19 +1066,23 @@ def hook(guard_wrapper, f_locals, builder):
 
 class RecursiveDictTagTests(torch._dynamo.test_case.TestCase):
     def setUp(self):
+        super().setUp()
         self._prev = torch._dynamo.config.use_recursive_dict_tags_for_guards
         torch._dynamo.config.use_recursive_dict_tags_for_guards = True
 
     def tearDown(self):
+        super().tearDown()
         torch._dynamo.config.use_recursive_dict_tags_for_guards = self._prev
 
 
 class TagSafetyChecks(RecursiveDictTagTests):
     def setUp(self):
+        super().setUp()
         self._prev = torch._dynamo.config.use_recursive_dict_tags_for_guards
         torch._dynamo.config.use_recursive_dict_tags_for_guards = True
 
     def tearDown(self):
+        super().tearDown()
         torch._dynamo.config.use_recursive_dict_tags_for_guards = self._prev
 
     def test_immutable_tag_safe(self):
@@ -1462,6 +1466,531 @@ def max_size_test(guard_wrapper, f_locals, builder):
                 opt_fn(x)
 
 
+class SourceCloneTests(torch._dynamo.test_case.TestCase):
+    def test_clone_identity_transform(self):
+        """Identity transform should produce a source with the same name."""
+        from torch._dynamo.source import AttrSource, GetItemSource, LocalSource
+
+        local = LocalSource("x")
+        attr = AttrSource(local, "weight")
+        item = GetItemSource(attr, 0)
+
+        for source in [local, attr, item]:
+            cloned = source.clone(lambda x: x)
+            self.assertEqual(cloned.name, source.name)
+
+    def test_clone_no_transform(self):
+        from torch._dynamo.source import AttrSource, LocalSource
+
+        local = LocalSource("x")
+        attr = AttrSource(local, "weight")
+
+        self.assertIs(local.clone(), local)
+        cloned_attr = attr.clone()
+        self.assertEqual(cloned_attr.name, attr.name)
+
+    def test_clone_parameterized_deep_chain(self):
+        """Replace leaf sources deep in a chain via a find->replace dictionary."""
+        from torch._dynamo.source import (
+            AttrSource,
+            ConstDictKeySource,
+            DictGetItemSource,
+            GetItemSource,
+            LocalSource,
+        )
+
+        # Build: L['x'].layers[0].weight
+        local = LocalSource("x")
+        attr1 = AttrSource(local, "layers")
+        item = GetItemSource(attr1, 0)
+        attr2 = AttrSource(item, "weight")
+
+        replacements = {local: LocalSource("y")}
+
+        def transform(s):
+            return replacements.get(s, s)
+
+        cloned = attr2.clone(transform)
+        self.assertEqual(cloned.name, "L['y'].layers[0].weight")
+
+        # Build: L['d'][list(dict.keys(L['d']))[0]]  (DictGetItemSource with Source key)
+        local_d = LocalSource("d")
+        key = ConstDictKeySource(local_d, 0)
+        dict_src = DictGetItemSource(local_d, key)
+
+        replacements = {local_d: LocalSource("other")}
+        cloned = dict_src.clone(transform)
+        self.assertEqual(cloned.name, "L['other'][list(dict.keys(L['other']))[0]]")
+
+    def test_clone_dict_get_item_source_with_constant_key(self):
+        from torch._dynamo.source import DictGetItemSource, LocalSource
+
+        local = LocalSource("d")
+        source = DictGetItemSource(local, "key")
+        cloned = source.clone()
+        self.assertEqual(cloned.name, source.name)
+
+        replacement = LocalSource("other_d")
+
+        def replace_local(s):
+            if isinstance(s, LocalSource) and s.local_name == "d":
+                return replacement
+            return s
+
+        cloned = source.clone(replace_local)
+        self.assertEqual(cloned.name, "L['other_d']['key']")
+
+    def test_clone_dict_get_item_source_with_source_key(self):
+        from torch._dynamo.source import (
+            ConstDictKeySource,
+            DictGetItemSource,
+            LocalSource,
+        )
+
+        local = LocalSource("d")
+        key_source = ConstDictKeySource(local, 0)
+        source = DictGetItemSource(local, key_source)
+
+        cloned = source.clone(lambda x: x)
+        self.assertEqual(cloned.name, source.name)
+
+    def test_clone_dict_subclass_get_item_source(self):
+        from torch._dynamo.source import DictSubclassGetItemSource, LocalSource
+
+        local = LocalSource("d")
+        source = DictSubclassGetItemSource(local, "key")
+
+        cloned = source.clone()
+        self.assertEqual(cloned.name, source.name)
+
+        cloned = source.clone(lambda x: x)
+        self.assertEqual(cloned.name, source.name)
+
+    def test_clone_get_item_source(self):
+        from torch._dynamo.source import GetItemSource, LocalSource
+
+        local = LocalSource("lst")
+        source = GetItemSource(local, 3)
+
+        cloned = source.clone()
+        self.assertEqual(cloned.name, source.name)
+
+        cloned = source.clone(lambda x: x)
+        self.assertEqual(cloned.name, source.name)
+
+
+class GuardCheckSpecTests(torch._dynamo.test_case.TestCase):
+    """Tests for the GuardCheckSpec get_metadata_fn/eval_fn handlers on GuardBuilder."""
+
+    def _get_handler(self, name):
+        from torch._dynamo.guards import GUARD_VALUE_DISPATCH
+
+        return GUARD_VALUE_DISPATCH[name]
+
+    def _make_guard(self, create_fn):
+        from torch._dynamo.source import LocalSource
+
+        return LocalSource("x").make_guard(create_fn)
+
+    def test_type_match(self):
+        from torch._dynamo.guards import GuardBuilder
+
+        guard = self._make_guard(GuardBuilder.TYPE_MATCH)
+        handler = self._get_handler("TYPE_MATCH")
+
+        expected = handler.get_metadata_fn(guard, 42)
+        self.assertIs(expected, int)
+        self.assertTrue(handler.eval_fn(100, expected))
+        self.assertFalse(handler.eval_fn("hello", expected))
+
+    def test_constant_match(self):
+        from torch._dynamo.guards import GuardBuilder
+
+        guard = self._make_guard(GuardBuilder.CONSTANT_MATCH)
+        handler = self._get_handler("CONSTANT_MATCH")
+
+        expected = handler.get_metadata_fn(guard, 42)
+        self.assertEqual(expected, 42)
+        self.assertTrue(handler.eval_fn(42, expected))
+        self.assertFalse(handler.eval_fn(99, expected))
+
+    def test_equals_match(self):
+        from torch._dynamo.guards import GuardBuilder
+
+        guard = self._make_guard(GuardBuilder.EQUALS_MATCH)
+        handler = self._get_handler("EQUALS_MATCH")
+
+        expected = handler.get_metadata_fn(guard, [1, 2, 3])
+        self.assertTrue(handler.eval_fn([1, 2, 3], expected))
+        self.assertFalse(handler.eval_fn([1, 2], expected))
+
+    def test_id_match(self):
+        from torch._dynamo.guards import GuardBuilder
+
+        guard = self._make_guard(GuardBuilder.ID_MATCH)
+        handler = self._get_handler("ID_MATCH")
+
+        obj = object()
+        expected = handler.get_metadata_fn(guard, obj)
+        self.assertIs(expected, obj)
+        self.assertTrue(handler.eval_fn(obj, expected))
+        self.assertFalse(handler.eval_fn(object(), expected))
+
+    def test_class_match(self):
+        from torch._dynamo.guards import GuardBuilder
+
+        guard = self._make_guard(GuardBuilder.CLASS_MATCH)
+        handler = self._get_handler("CLASS_MATCH")
+
+        expected = handler.get_metadata_fn(guard, dict)
+        self.assertIs(expected, dict)
+        self.assertTrue(handler.eval_fn(dict, expected))
+        self.assertFalse(handler.eval_fn(list, expected))
+
+    def test_module_match(self):
+        import types
+
+        from torch._dynamo.guards import GuardBuilder
+
+        guard = self._make_guard(GuardBuilder.MODULE_MATCH)
+        handler = self._get_handler("MODULE_MATCH")
+
+        mod = types.ModuleType("test_mod")
+        expected = handler.get_metadata_fn(guard, mod)
+        self.assertIs(expected, mod)
+        self.assertTrue(handler.eval_fn(mod, expected))
+
+        other = types.ModuleType("other_mod")
+        self.assertFalse(handler.eval_fn(other, expected))
+
+    def test_builtin_match(self):
+        from torch._dynamo.guards import GuardBuilder
+
+        guard = self._make_guard(GuardBuilder.BUILTIN_MATCH)
+        handler = self._get_handler("BUILTIN_MATCH")
+
+        expected = handler.get_metadata_fn(guard, len)
+        self.assertIs(expected, len)
+        self.assertTrue(handler.eval_fn(len, expected))
+        self.assertFalse(handler.eval_fn(print, expected))
+
+    def test_hasattr_present(self):
+        from torch._dynamo.guards import GuardBuilder
+
+        guard = self._make_guard(functools.partial(GuardBuilder.HASATTR, attr="weight"))
+        handler = self._get_handler("HASATTR")
+
+        class Obj:
+            weight = 1.0
+
+        expected = handler.get_metadata_fn(guard, Obj())
+        self.assertEqual(expected, ("weight", True))
+        self.assertTrue(handler.eval_fn(Obj(), expected))
+
+        class Empty:
+            pass
+
+        self.assertFalse(handler.eval_fn(Empty(), expected))
+
+    def test_hasattr_absent(self):
+        from torch._dynamo.guards import GuardBuilder
+
+        guard = self._make_guard(functools.partial(GuardBuilder.HASATTR, attr="bias"))
+        handler = self._get_handler("HASATTR")
+
+        class Obj:
+            weight = 1.0
+
+        obj = Obj()
+        expected = handler.get_metadata_fn(guard, obj)
+        self.assertEqual(expected, ("bias", False))
+        self.assertTrue(handler.eval_fn(obj, expected))
+        # Adding the attr should fail the "not hasattr" guard
+        obj.bias = 0.0  # type: ignore[attr-defined]
+        self.assertFalse(handler.eval_fn(obj, expected))
+
+    def test_sequence_length(self):
+        from torch._dynamo.guards import GuardBuilder
+
+        guard = self._make_guard(GuardBuilder.SEQUENCE_LENGTH)
+        handler = self._get_handler("SEQUENCE_LENGTH")
+
+        expected = handler.get_metadata_fn(guard, [1, 2, 3])
+        self.assertEqual(expected, 3)
+        self.assertTrue(handler.eval_fn([4, 5, 6], expected))
+        self.assertTrue(handler.eval_fn((7, 8, 9), expected))
+        self.assertFalse(handler.eval_fn([1], expected))
+
+    def test_dict_contains(self):
+        from torch._dynamo.guards import GuardBuilder
+
+        guard = self._make_guard(functools.partial(GuardBuilder.DICT_CONTAINS, key="a"))
+        handler = self._get_handler("DICT_CONTAINS")
+
+        d = {"a": 1, "b": 2}
+        expected = handler.get_metadata_fn(guard, d)
+        self.assertEqual(expected, "a")
+        self.assertTrue(handler.eval_fn({"a": 99}, expected))
+        self.assertFalse(handler.eval_fn({"b": 1}, expected))
+
+    def test_dict_not_contains(self):
+        from torch._dynamo.guards import GuardBuilder
+
+        guard = self._make_guard(
+            functools.partial(GuardBuilder.DICT_NOT_CONTAINS, key="a")
+        )
+        handler = self._get_handler("DICT_NOT_CONTAINS")
+
+        d = {"b": 2}
+        expected = handler.get_metadata_fn(guard, d)
+        self.assertEqual(expected, "a")
+        self.assertTrue(handler.eval_fn({"b": 1}, expected))
+        self.assertFalse(handler.eval_fn({"a": 1}, expected))
+
+    def test_not_present_in_generic_dict(self):
+        from torch._dynamo.guards import GuardBuilder
+
+        guard = self._make_guard(
+            functools.partial(GuardBuilder.NOT_PRESENT_IN_GENERIC_DICT, attr="hidden")
+        )
+        handler = self._get_handler("NOT_PRESENT_IN_GENERIC_DICT")
+
+        class Obj:
+            pass
+
+        obj = Obj()
+        expected = handler.get_metadata_fn(guard, obj)
+        self.assertEqual(expected, "hidden")
+        self.assertTrue(handler.eval_fn(obj, expected))
+        obj.hidden = 1  # type: ignore[attr-defined]
+        self.assertFalse(handler.eval_fn(obj, expected))
+
+    def test_closure_match(self):
+        from torch._dynamo.guards import GuardBuilder
+
+        guard = self._make_guard(GuardBuilder.CLOSURE_MATCH)
+        handler = self._get_handler("CLOSURE_MATCH")
+
+        def fn():
+            return 42
+
+        expected = handler.get_metadata_fn(guard, fn)
+        self.assertIs(expected, fn.__code__)
+        self.assertTrue(handler.eval_fn(fn, expected))
+
+        def other_fn():
+            return 99
+
+        self.assertFalse(handler.eval_fn(other_fn, expected))
+
+    def test_tensor_match(self):
+        from torch._dynamo.guards import extract_tensor_metadata, GuardBuilder
+
+        guard = self._make_guard(GuardBuilder.TENSOR_MATCH)
+        handler = self._get_handler("TENSOR_MATCH")
+
+        t = torch.randn(3, 4)
+        expected = handler.get_metadata_fn(guard, t)
+        self.assertEqual(expected, extract_tensor_metadata(t))
+        self.assertTrue(handler.eval_fn(t, expected))
+        self.assertTrue(handler.eval_fn(torch.randn(3, 4), expected))
+        self.assertFalse(handler.eval_fn(torch.randn(5, 4), expected))
+        self.assertFalse(
+            handler.eval_fn(torch.randn(3, 4, dtype=torch.float64), expected)
+        )
+        self.assertFalse(handler.eval_fn(42, expected))
+
+    def test_empty_nn_module_hooks_dict(self):
+        from torch._dynamo.guards import GuardBuilder
+
+        guard = self._make_guard(GuardBuilder.EMPTY_NN_MODULE_HOOKS_DICT)
+        handler = self._get_handler("EMPTY_NN_MODULE_HOOKS_DICT")
+
+        expected = handler.get_metadata_fn(guard, {})
+        self.assertIsNone(expected)
+        self.assertTrue(handler.eval_fn({}, expected))
+        self.assertFalse(handler.eval_fn({"hook": lambda: None}, expected))
+
+    def test_bool_match(self):
+        from torch._dynamo.guards import GuardBuilder
+
+        guard = self._make_guard(GuardBuilder.BOOL_MATCH)
+        handler = self._get_handler("BOOL_MATCH")
+
+        expected = handler.get_metadata_fn(guard, True)
+        self.assertTrue(handler.eval_fn(True, expected))
+        self.assertFalse(handler.eval_fn(False, expected))
+
+        expected_false = handler.get_metadata_fn(guard, False)
+        self.assertTrue(handler.eval_fn(False, expected_false))
+        self.assertFalse(handler.eval_fn(True, expected_false))
+
+    def test_none_match(self):
+        from torch._dynamo.guards import GuardBuilder
+
+        guard = self._make_guard(GuardBuilder.NONE_MATCH)
+        handler = self._get_handler("NONE_MATCH")
+
+        expected = handler.get_metadata_fn(guard, None)
+        self.assertIsNone(expected)
+        self.assertTrue(handler.eval_fn(None, expected))
+        self.assertFalse(handler.eval_fn(0, expected))
+        self.assertFalse(handler.eval_fn("", expected))
+
+    def test_function_match(self):
+        from torch._dynamo.guards import GuardBuilder
+
+        guard = self._make_guard(GuardBuilder.FUNCTION_MATCH)
+        handler = self._get_handler("FUNCTION_MATCH")
+
+        def my_fn():
+            return 42
+
+        expected = handler.get_metadata_fn(guard, my_fn)
+        self.assertIs(expected, my_fn)
+        self.assertTrue(handler.eval_fn(my_fn, expected))
+        self.assertFalse(handler.eval_fn(lambda: 42, expected))
+        self.assertFalse(handler.eval_fn(torch.add, expected))
+
+    def test_weakref_alive(self):
+        from torch._dynamo.guards import GuardBuilder
+
+        guard = self._make_guard(GuardBuilder.WEAKREF_ALIVE)
+        handler = self._get_handler("WEAKREF_ALIVE")
+
+        # The guard spec receives the dereferenced value (ref()), not the
+        # weakref itself. A live referent resolves to the object; a dead
+        # one resolves to None.
+        class C:
+            pass
+
+        obj = C()
+        ref = weakref.ref(obj)
+        expected = handler.get_metadata_fn(guard, ref())
+        self.assertIsNone(expected)
+        self.assertTrue(handler.eval_fn(ref(), expected))
+        # Delete the referent — weakref() now returns None
+        del obj
+        self.assertFalse(handler.eval_fn(ref(), expected))
+
+    def test_set_contains(self):
+        from torch._dynamo.guards import GuardBuilder
+
+        guard = self._make_guard(functools.partial(GuardBuilder.SET_CONTAINS, key="a"))
+        handler = self._get_handler("SET_CONTAINS")
+
+        s = {"a", "b", "c"}
+        expected = handler.get_metadata_fn(guard, s)
+        self.assertEqual(expected, "a")
+        self.assertTrue(handler.eval_fn({"a", "x"}, expected))
+        self.assertFalse(handler.eval_fn({"b", "c"}, expected))
+
+    def test_set_not_contains(self):
+        from torch._dynamo.guards import GuardBuilder
+
+        guard = self._make_guard(
+            functools.partial(GuardBuilder.SET_NOT_CONTAINS, key="z")
+        )
+        handler = self._get_handler("SET_NOT_CONTAINS")
+
+        s = {"a", "b"}
+        expected = handler.get_metadata_fn(guard, s)
+        self.assertEqual(expected, "z")
+        self.assertTrue(handler.eval_fn({"a", "b"}, expected))
+        self.assertFalse(handler.eval_fn({"a", "z"}, expected))
+
+    def test_not_none_match(self):
+        from torch._dynamo.guards import GuardBuilder
+
+        guard = self._make_guard(GuardBuilder.NOT_NONE_MATCH)
+        handler = self._get_handler("NOT_NONE_MATCH")
+
+        expected = handler.get_metadata_fn(guard, torch.randn(2))
+        self.assertIsNone(expected)
+        self.assertTrue(handler.eval_fn(torch.randn(3), expected))
+        self.assertTrue(handler.eval_fn(42, expected))
+        self.assertFalse(handler.eval_fn(None, expected))
+
+    def test_dispatch_key_set_match(self):
+        from torch._dynamo.guards import GuardBuilder
+
+        guard = self._make_guard(GuardBuilder.DISPATCH_KEY_SET_MATCH)
+        handler = self._get_handler("DISPATCH_KEY_SET_MATCH")
+
+        dks = torch._C._dispatch_keys(torch.randn(3))
+        expected = handler.get_metadata_fn(guard, dks)
+        self.assertEqual(expected, dks.raw_repr())
+        self.assertTrue(handler.eval_fn(dks, expected))
+        # Different tensor with same dispatch keys should match
+        dks2 = torch._C._dispatch_keys(torch.randn(5))
+        self.assertTrue(handler.eval_fn(dks2, expected))
+
+    def test_tuple_iterator_len(self):
+        from torch._dynamo.guards import GuardBuilder
+
+        guard = self._make_guard(GuardBuilder.TUPLE_ITERATOR_LEN)
+        handler = self._get_handler("TUPLE_ITERATOR_LEN")
+
+        it = iter((1, 2, 3))
+        expected = handler.get_metadata_fn(guard, it)
+        self.assertEqual(expected, 3)
+        self.assertTrue(handler.eval_fn(iter((4, 5, 6)), expected))
+        self.assertFalse(handler.eval_fn(iter((1,)), expected))
+
+    def test_range_iterator_match(self):
+        from torch._dynamo.guards import GuardBuilder
+
+        guard = self._make_guard(GuardBuilder.RANGE_ITERATOR_MATCH)
+        handler = self._get_handler("RANGE_ITERATOR_MATCH")
+
+        it = iter(range(1, 10, 2))
+        expected = handler.get_metadata_fn(guard, it)
+        self.assertTrue(handler.eval_fn(iter(range(1, 10, 2)), expected))
+        self.assertFalse(handler.eval_fn(iter(range(0, 10, 2)), expected))
+        self.assertFalse(handler.eval_fn(iter(range(1, 10, 3)), expected))
+
+    def test_nn_module(self):
+        from torch._dynamo.guards import GuardBuilder
+
+        guard = self._make_guard(GuardBuilder.NN_MODULE)
+        handler = self._get_handler("NN_MODULE")
+
+        mod = torch.nn.Linear(3, 4)
+        expected = handler.get_metadata_fn(guard, mod)
+        self.assertIs(expected, mod)
+        self.assertTrue(handler.eval_fn(mod, expected))
+        self.assertFalse(handler.eval_fn(torch.nn.Linear(3, 4), expected))
+
+    def test_mapping_keys_check(self):
+        import types
+
+        from torch._dynamo.guards import GuardBuilder
+
+        guard = self._make_guard(GuardBuilder.MAPPING_KEYS_CHECK)
+        handler = self._get_handler("MAPPING_KEYS_CHECK")
+
+        mp = types.MappingProxyType({"a": 1, "b": 2})
+        expected = handler.get_metadata_fn(guard, mp)
+        self.assertEqual(expected, ["a", "b"])
+        self.assertTrue(
+            handler.eval_fn(types.MappingProxyType({"a": 10, "b": 20}), expected)
+        )
+        self.assertFalse(handler.eval_fn(types.MappingProxyType({"x": 1}), expected))
+
+    def test_dict_keys_match(self):
+        from torch._dynamo.guards import GuardBuilder
+
+        guard = self._make_guard(GuardBuilder.DICT_KEYS_MATCH)
+        handler = self._get_handler("DICT_KEYS_MATCH")
+
+        d = {"a": 1, "b": 2, "c": 3}
+        expected = handler.get_metadata_fn(guard, d)
+        self.assertEqual(expected, ["a", "b", "c"])
+        self.assertTrue(handler.eval_fn({"a": 10, "b": 20, "c": 30}, expected))
+        self.assertFalse(handler.eval_fn({"a": 1, "b": 2}, expected))
+        self.assertFalse(handler.eval_fn({"x": 1, "y": 2, "z": 3}, expected))
+
+
 if __name__ == "__main__":
     from torch._dynamo.test_case import run_tests
 
diff --git a/test/dynamo/test_guard_serialization.py b/test/dynamo/test_guard_serialization.py
index 99727bb67ab33..55d83764f99af 100644
--- a/test/dynamo/test_guard_serialization.py
+++ b/test/dynamo/test_guard_serialization.py
@@ -1,6 +1,7 @@
 # Owner(s): ["module: dynamo"]
 
 import dataclasses
+import itertools
 import pickle
 import sys
 import tempfile
@@ -9,7 +10,6 @@
 import weakref
 from collections.abc import Iterator
 from typing import NamedTuple
-from unittest.mock import patch
 
 import torch
 import torch._dynamo.testing
@@ -154,7 +154,7 @@ def __torch_dispatch__(cls, func, types, args, kwargs):
         kwargs_a = pytree.tree_map_only(SubclassWithMeta, lambda x: x.a, kwargs)
         out_a = func(*args_a, **kwargs_a)
         if isinstance(out_a, torch.Tensor):
-            assert isinstance(args[0], SubclassWithMeta)
+            assert isinstance(args[0], SubclassWithMeta)  # noqa: S101
             return SubclassWithMeta(out_a, extra=args[0].extra)
         return out_a
 
@@ -164,13 +164,13 @@ def __tensor_flatten__(self):
 
     @staticmethod
     def __tensor_unflatten__(inner_tensors, meta, outer_size, outer_stride):
-        assert isinstance(meta, dict)
+        assert isinstance(meta, dict)  # noqa: S101
         a = inner_tensors["a"]
         # pull out extra from meta
         extra = meta["extra"]
         if type(a) is torch.Tensor:
-            assert outer_size is not None
-            assert outer_stride is not None
+            assert outer_size is not None  # noqa: S101
+            assert outer_stride is not None  # noqa: S101
         return SubclassWithMeta(a, extra, outer_size, outer_stride)
 
 
@@ -208,7 +208,7 @@ def __torch_dispatch__(cls, func, types, args, kwargs):
         )
         out_a = func(*args_a, **kwargs_a)
         if isinstance(out_a, torch.Tensor):
-            assert isinstance(args[0], SubclassWithCustomMetadataGuard)
+            assert isinstance(args[0], SubclassWithCustomMetadataGuard)  # noqa: S101
             return SubclassWithCustomMetadataGuard(out_a, extra=args[0].extra)
         return out_a
 
@@ -225,13 +225,13 @@ def __tensor_flatten__(self):
 
     @staticmethod
     def __tensor_unflatten__(inner_tensors, meta, outer_size, outer_stride):
-        assert isinstance(meta, dict)
+        assert isinstance(meta, dict)  # noqa: S101
         a = inner_tensors["a"]
         # pull out extra from meta
         extra = meta["extra"]
         if type(a) is torch.Tensor:
-            assert outer_size is not None
-            assert outer_stride is not None
+            assert outer_size is not None  # noqa: S101
+            assert outer_stride is not None  # noqa: S101
         return SubclassWithCustomMetadataGuard(a, extra, outer_size, outer_stride)
 
 
@@ -269,7 +269,7 @@ def __torch_dispatch__(cls, func, types, args, kwargs):
         )
         out_a = func(*args_a, **kwargs_a)
         if isinstance(out_a, torch.Tensor):
-            assert isinstance(args[0], SubclassWithSubclassInnerTensor)
+            assert isinstance(args[0], SubclassWithSubclassInnerTensor)  # noqa: S101
             return SubclassWithSubclassInnerTensor(out_a, extra=args[0].inner_sub.extra)
         return out_a
 
@@ -278,12 +278,12 @@ def __tensor_flatten__(self):
 
     @staticmethod
     def __tensor_unflatten__(inner_tensors, meta, outer_size, outer_stride):
-        assert meta is None
+        assert meta is None  # noqa: S101
         a = inner_tensors["a"]
         extra = inner_tensors["inner_sub"].extra
         if type(a) is torch.Tensor:
-            assert outer_size is not None
-            assert outer_stride is not None
+            assert outer_size is not None  # noqa: S101
+            assert outer_stride is not None  # noqa: S101
         return SubclassWithSubclassInnerTensor(a, extra, outer_size, outer_stride)
 
 
@@ -364,7 +364,8 @@ def _test_serialization(self, guard_type, fn, *args, **kwargs):
         finally:
             sys.settrace(None)
 
-        assert self._frame_state is not None
+        if self._frame_state is None:
+            raise AssertionError("Expected _frame_state to be set after tracing")
 
         # Set f_locals from regenerated kwargs to handle exhausted input iterators
         # NB: This is super janky and might cause unforeseen problems
@@ -623,11 +624,11 @@ def __tensor_flatten__(self):
 
             @staticmethod
             def __tensor_unflatten__(inner_tensors, meta, outer_size, outer_stride):
-                assert meta is None
+                assert meta is None  # noqa: S101
                 a = inner_tensors["a"]
                 if type(a) is torch.Tensor:
-                    assert outer_size is not None
-                    assert outer_stride is not None
+                    assert outer_size is not None  # noqa: S101
+                    assert outer_stride is not None  # noqa: S101
                 return LocalSubclass(a, outer_size, outer_stride)
 
         def fn(x):
@@ -789,22 +790,6 @@ def fn(x, y):
         # guard should fail for different y value
         self._test_check_fn(ref, loaded, {"x": torch.randn(3), "y": 6}, False)
 
-    def test_nn_module(self):
-        def fn(m, x):
-            return m(x)
-
-        m = GlobalModule()
-        x = torch.randn(3)
-
-        # config setting controls whether the NN_MODULE guard is installed
-        with patch("torch._dynamo.config.inline_inbuilt_nn_modules", False):
-            # we don't support NN_MODULE because it adds an ID_MATCH guard, and we don't
-            # support that in serialization
-            with self.assertRaisesRegex(
-                PackageError, "NN_MODULE guard cannot be serialized."
-            ):
-                self._test_serialization("NN_MODULE", fn, m, x)
-
     def test_class_match(self):
         def fn(x):
             # usage of this context manager installs a FUNCTION_MATCH guard
@@ -945,6 +930,35 @@ def _gen_kwargs(x=x):
             ref, loaded, {"x": torch.randn(4), "r": iter(range(2, 15, 4))}, False
         )
 
+    def test_count_iterator_match(self):
+        def fn(x, counter):
+            return x + next(counter)
+
+        x = torch.randn(3)
+
+        def _gen_kwargs(x=x):
+            return {"x": x, "counter": itertools.count(2, 3)}
+
+        ref, loaded = self._test_serialization(
+            "COUNT_ITERATOR_MATCH", fn, _gen_fn=_gen_kwargs
+        )
+
+        self._test_check_fn(
+            ref, loaded, {"x": x, "counter": itertools.count(2, 3)}, True
+        )
+        self._test_check_fn(
+            ref,
+            loaded,
+            {"x": torch.randn(4), "counter": itertools.count(2, 3)},
+            True,
+        )
+        self._test_check_fn(
+            ref, loaded, {"x": x, "counter": itertools.count(5, 3)}, False
+        )
+        self._test_check_fn(
+            ref, loaded, {"x": x, "counter": itertools.count(2, 4)}, False
+        )
+
     def test_dict_version(self):
         def fn(x):
             return pytree.tree_leaves(x)[0] + 1
@@ -1481,7 +1495,7 @@ def add(self, x):
                 return x + 1
 
         def foo(x: torch.Tensor, y: list[MyClass]):
-            assert len(y) == 1
+            assert len(y) == 1  # noqa: S101
             return x + 1
 
         ref, loaded = self._test_serialization(
@@ -1493,7 +1507,7 @@ def foo(x: torch.Tensor, y: list[MyClass]):
 
     def test_bound_methods_empty(self):
         def foo(x, y):
-            assert callable(y[0])
+            assert callable(y[0])  # noqa: S101
             return x + 1
 
         ref, loaded = self._test_serialization(
@@ -1506,6 +1520,7 @@ def foo(x, y):
             True,
         )
 
+    @torch._dynamo.config.patch(nested_graph_breaks=False)
     def test_ddp_module(self):
         import torch.distributed as dist
 
@@ -1661,19 +1676,17 @@ def foo(inputs):
         self._test_check_fn(ref, loaded, {"inputs": Inputs(x, weakref.ref(x))}, True)
 
     def test_unused_stream(self):
-        if not torch.cuda.is_available():
-            self.skipTest("CUDA is not available")
+        if not torch.accelerator.is_available():
+            self.skipTest("Accelerator is not available")
 
         def foo(inputs):
             return inputs.x + 1
 
         x = torch.randn(3, 2)
         ref, loaded = self._test_serialization(
-            "TENSOR_MATCH", foo, Inputs(x, torch.cuda.Stream())
-        )
-        self._test_check_fn(
-            ref, loaded, {"inputs": Inputs(x, torch.cuda.Stream())}, True
+            "TENSOR_MATCH", foo, Inputs(x, torch.Stream())
         )
+        self._test_check_fn(ref, loaded, {"inputs": Inputs(x, torch.Stream())}, True)
 
     def test_unused_process_group(self):
         import torch.distributed as dist
@@ -1732,7 +1745,7 @@ def forward(x):
         m.forward = forward
 
         def foo(f, x):
-            assert callable(f)
+            assert callable(f)  # noqa: S101
             return f(x)
 
         x = torch.randn(3, 2)
@@ -1762,7 +1775,9 @@ def foo(x):
 
         with torch.no_grad():
             compiled_fn = torch.compile(
-                foo, options={"guard_filter_fn": torch.compiler.skip_all_guards_unsafe}
+                foo,
+                backend="eager",
+                options={"guard_filter_fn": torch.compiler.skip_all_guards_unsafe},
             )
             compiled_fn(x)
 
@@ -1778,7 +1793,9 @@ def foo(x):
 
         with GlobalTorchFunctionMode():
             compiled_fn = torch.compile(
-                foo, options={"guard_filter_fn": torch.compiler.skip_all_guards_unsafe}
+                foo,
+                backend="eager",
+                options={"guard_filter_fn": torch.compiler.skip_all_guards_unsafe},
             )
             compiled_fn(x)
 
@@ -1852,6 +1869,17 @@ def test_source_serialization(self):
 
         self.assertEqual(pickle.dumps(src1), pickle.dumps(src2))
 
+    def test_source_serialization_init_false_fields(self):
+        # Test that source serialization handles fields that are not initialized
+        from torch._dynamo.source import DefaultsSource, LocalSource
+
+        base = LocalSource("x")
+        source = DefaultsSource(base=base, idx_key=0, is_kw=False)
+
+        # Round-trip through pickle should work even with init=False fields
+        restored = pickle.loads(pickle.dumps(source))
+        self.assertEqual(source, restored)
+
 
 class SimpleModule(torch.nn.Module):
     def __init__(self, c):
diff --git a/test/dynamo/test_higher_order_ops.py b/test/dynamo/test_higher_order_ops.py
index e10ed297f9781..683040499c228 100644
--- a/test/dynamo/test_higher_order_ops.py
+++ b/test/dynamo/test_higher_order_ops.py
@@ -38,6 +38,7 @@
     xfailIfTorchDynamo,
 )
 from torch.testing._internal.hop_db import hop_db
+from torch.testing._internal.inductor_utils import GPU_TYPE
 from torch.testing._internal.logging_utils import LoggingTestCase, make_logging_test
 from torch.testing._internal.triton_utils import (
     requires_cuda_and_triton,
@@ -47,7 +48,8 @@
 
 def count_ops(gm, args, freq, op):
     actual = [node.target for node in gm.graph.nodes].count(op)
-    assert actual == freq, f"expected={freq}, actual={actual}"
+    if actual != freq:
+        raise AssertionError(f"expected={freq}, actual={actual}")
     return gm
 
 
@@ -134,7 +136,7 @@ def default_args_generator(seed_value):
         yield new_args
 
 
-class HigherOrderOpTests(torch._dynamo.test_case.TestCaseWithNestedGraphBreaks):
+class HigherOrderOpTests(torch._dynamo.test_case.TestCase):
     def _assert_wrap_fallback(self, func, args, setup=lambda: None):
         counters.clear()
         backend = EagerAndRecordGraphs()
@@ -1157,7 +1159,7 @@ def restore_twotensor_impl():
         def _(pred, true_fn, false_fn, operands):
             nonlocal called
             called += 1
-            assert len(operands) == 1
+            assert len(operands) == 1  # noqa: S101
             a = cond_op(pred, true_fn, false_fn, (operands[0].a,))
             b = cond_op(pred, true_fn, false_fn, (operands[0].b,))
             return TwoTensor(a, b)
@@ -2581,10 +2583,6 @@ def f(x):
         # 3 args - 1 for input, and other 2 for the weight and bias
         self.assertTrue(len(wrap_node.args), 3)
 
-        # Check that the linear bias and weight are getattr in the outer graph
-        if not torch._dynamo.config.inline_inbuilt_nn_modules:
-            self.assertTrue(len(dict(backend.graphs[0].named_parameters())) == 2)
-
         # Check that the inner function has one op and its a linear op
         body_function = getattr(backend.graphs[0], wrap_node.args[0].name)
         self.assertEqual(op_count(body_function), 1)
@@ -2713,10 +2711,6 @@ def f(x):
         wrap_node = find_first_node(backend.graphs[0], wrap)
         self.assertTrue(len(wrap_node.args), 3)
 
-        # Check that the linear bias and weight are getattr in the outer graph
-        if not torch._dynamo.config.inline_inbuilt_nn_modules:
-            self.assertTrue(len(dict(backend.graphs[0].named_parameters())) == 2)
-
         # Check that the inner function has one op and its a linear op
         body_function = getattr(backend.graphs[0], wrap_node.args[0].name)
         self.assertEqual(op_count(body_function), 1)
@@ -3187,7 +3181,7 @@ def fn(pred, pytree_in):
                 torch.compile(fn, backend="eager")(pred, pytree_in)
 
     def test_cond_with_empty_operands(self):
-        @torch.compile(fullgraph=True)
+        @torch.compile(fullgraph=True, backend="eager")
         def fn(x, y, z):
             def true_fn():
                 return y + 2
@@ -3221,7 +3215,7 @@ def my_hop_fn(fn, *args, k=1, **kwargs):
 
         def my_hop_fn_2_impl(fn, *args, g=None):
             def wrapper(*args, **kwargs):
-                assert g is not None
+                assert g is not None  # noqa: S101
                 out = fn(*args)
                 if isinstance(out, tuple):
                     return (g(out[0]),)
@@ -3391,7 +3385,7 @@ def outer_body_fn(x):
         with self.assertRaisesRegex(RuntimeError, msg):
             fn_with_hints(x, y)
 
-    @requires_cuda_and_triton
+    @requires_gpu_and_triton
     def test_wrap_inductor_compiled_regions_option(self):
         """
         Test that wrap_inductor_compiled_regions option wraps compiled regions
@@ -3413,8 +3407,8 @@ def fn_wrapped(x, y):
         def fn_not_wrapped(x, y):
             return torch.matmul(x, y)
 
-        x = torch.randn(4, 4, device="cuda")
-        y = torch.randn(4, 4, device="cuda")
+        x = torch.randn(4, 4, device=GPU_TYPE)
+        y = torch.randn(4, 4, device=GPU_TYPE)
 
         # Test wrapped version - HOP should be visible in DebugMode
         with DebugMode() as debug_mode_wrapped:
@@ -3435,7 +3429,7 @@ def fn_not_wrapped(x, y):
         self.assertEqual(result_wrapped, expected)
         self.assertEqual(result_not_wrapped, expected)
 
-    @requires_cuda_and_triton
+    @requires_gpu_and_triton
     def test_wrap_inductor_compiled_regions_with_backward(self):
         """
         Test that wrap_inductor_compiled_regions works correctly with autograd.
@@ -3450,8 +3444,8 @@ def test_wrap_inductor_compiled_regions_with_backward(self):
         def fn(x, y):
             return torch.matmul(x, y)
 
-        x = torch.randn(4, 4, device="cuda", requires_grad=True)
-        y = torch.randn(4, 4, device="cuda", requires_grad=True)
+        x = torch.randn(4, 4, device=GPU_TYPE, requires_grad=True)
+        y = torch.randn(4, 4, device=GPU_TYPE, requires_grad=True)
 
         # Clone for eager comparison
         x_eager = x.detach().clone().requires_grad_(True)
@@ -3477,9 +3471,7 @@ def fn(x, y):
         self.assertEqual(y.grad, y_eager.grad)
 
 
-class HigherOrderOpVmapGuardTests(
-    torch._dynamo.test_case.TestCaseWithNestedGraphBreaks, LoggingTestCase
-):
+class HigherOrderOpVmapGuardTests(LoggingTestCase):
     @make_logging_test(recompiles=True)
     def test_vmap_grad_guard_ok(self, records):
         vmap = torch.vmap
@@ -3748,9 +3740,7 @@ def fn(x):
         self.assertGreater(len(records), 0)
 
 
-class FuncTorchHigherOrderOpTests(
-    torch._dynamo.test_case.TestCaseWithNestedGraphBreaks
-):
+class FuncTorchHigherOrderOpTests(torch._dynamo.test_case.TestCase):
     def tearDown(self):
         # Ensure that in the case of a test failure, the next test won't fail
         # because of a previous call to _vmap_increment_nesting that wasn't undone
@@ -3777,7 +3767,7 @@ def tearDown(self):
 
     def test_teardown_resets_nested_graph_breaks(self):
         expected_nested_state = getattr(
-            self, "prev_nested_graph_breaks", torch._dynamo.config.nested_graph_breaks
+            self, "_prior_nested_graph_breaks", torch._dynamo.config.nested_graph_breaks
         )
 
         def _check_flag():
@@ -4486,7 +4476,6 @@ def forward(self, L_x_: "f32[5]"):
 """,
         )
 
-    @config.patch(inline_inbuilt_nn_modules=True)
     def test_functional_call(self):
         def wrapper_fn(model, params, inputs, targets):
             prediction = torch.func.functional_call(model, params, (inputs,))
@@ -4503,10 +4492,9 @@ def wrapper_fn(model, params, inputs, targets):
             return
 
         actual = normalize_gm(wrapped_gm.print_readable(print_output=False))
-        if torch._dynamo.config.inline_inbuilt_nn_modules:
-            self.assertExpectedInline(
-                actual,
-                """\
+        self.assertExpectedInline(
+            actual,
+            """\
 class GraphModule(torch.nn.Module):
     def forward(self, L_model_parameters_weight_: "f32[3, 3]", L_model_parameters_bias_: "f32[3]", L_inputs_: "f32[64, 3]", L_targets_: "f32[64, 3]"):
         l_model_parameters_weight_ = L_model_parameters_weight_
@@ -4519,24 +4507,8 @@ def forward(self, L_model_parameters_weight_: "f32[3, 3]", L_model_parameters_bi
         mse_loss: "f32[]" = torch.nn.functional.mse_loss(prediction, l_targets_);  prediction = l_targets_ = None
         return (mse_loss,)
 """,
-            )
-        else:
-            self.assertExpectedInline(
-                actual,
-                """\
-class GraphModule(torch.nn.Module):
-    def forward(self, L_inputs_: "f32[64, 3]", L_targets_: "f32[64, 3]"):
-        l_inputs_ = L_inputs_
-        l_targets_ = L_targets_
-
-        prediction: "f32[64, 3]" = self.model(l_inputs_);  l_inputs_ = None
-
-        mse_loss: "f32[]" = torch.nn.functional.mse_loss(prediction, l_targets_);  prediction = l_targets_ = None
-        return (mse_loss,)
-""",
-            )
+        )
 
-    @config.patch(inline_inbuilt_nn_modules=True)
     def test_functional_call_sequential_params_and_buffers(self):
         # copied from test/test_stateless.py
         class MockModule(torch.nn.Module):
@@ -4566,8 +4538,7 @@ def wrapper_fn(model, params, buffers, inputs):
             return
 
         actual = normalize_gm(wrapped_gm.print_readable(print_output=False))
-        if torch._dynamo.config.inline_inbuilt_nn_modules:
-            expected = """\
+        expected = """\
 class GraphModule(torch.nn.Module):
     def forward(self, L_inputs_: "f32[1, 1]", L_model_modules_l1_parameters_weight_: "f32[1, 1]", L_model_modules_l1_parameters_bias_: "f32[1]", L_model_buffers_buffer_: "f32[1]"):
         l_inputs_ = L_inputs_
@@ -4578,49 +4549,11 @@ def forward(self, L_inputs_: "f32[1, 1]", L_model_modules_l1_parameters_weight_:
         add: "f32[1, 1]" = linear + l_model_buffers_buffer_;  linear = l_model_buffers_buffer_ = None
         return (add,)
 """
-            # We found Windows/Linux have some empty line difference, empty_line_normalizer will help fix it.
-            self.assertExpectedInline(
-                empty_line_normalizer(actual),
-                empty_line_normalizer(normalize_gm(expected)),
-            )
-        else:
-            self.assertExpectedInline(
-                actual,
-                """\
-class GraphModule(torch.nn.Module):
-    def forward(self, L_x_: "f32[1, 1]"):
-        l_x_ = L_x_
-
-        l__self___l1: "f32[1, 1]" = self.L__self___l1(l_x_);  l_x_ = None
-        l__self___buffer: "f32[1]" = self.L__self___buffer
-        add: "f32[1, 1]" = l__self___l1 + l__self___buffer;  l__self___l1 = l__self___buffer = None
-        return (add,)
-""",
-            )
-
-    @config.patch(inline_inbuilt_nn_modules=False)
-    def test_functional_call_disable_inline_nn_module(self):
-        counters.clear()
-
-        def wrapper_fn(model, params, inputs, targets):
-            prediction = torch.func.functional_call(model, params, (inputs,))
-            return torch.nn.functional.mse_loss(prediction, targets)
-
-        model = torch.nn.Linear(3, 3)
-        params = dict(model.named_parameters())
-        inputs = torch.randn(64, 3)
-        targets = torch.randn(64, 3)
-
-        actual = wrapper_fn(model, params, inputs, targets)
-        expected = torch.compile(wrapper_fn, backend="aot_eager", fullgraph=False)(
-            model, params, inputs, targets
-        )
-        self.assertEqual(len(counters["graph_break"]), 1)
-        self.assertIn(
-            "torch.func.functional_call capture is disabled",
-            next(iter(counters["graph_break"].keys())),
+        # We found Windows/Linux have some empty line difference, empty_line_normalizer will help fix it.
+        self.assertExpectedInline(
+            empty_line_normalizer(actual),
+            empty_line_normalizer(normalize_gm(expected)),
         )
-        self.assertEqual(actual, expected)
 
     def test_grad(self):
         counters.clear()
@@ -6110,7 +6043,7 @@ def wrapped_fn(x):
 
     def test_vmap_call_compiled_backward_fn(self):
         # See PyTorch issue #138422
-        @torch.compile
+        @torch.compile(backend="aot_eager")
         def f(x):
             return x**2
 
@@ -6128,7 +6061,7 @@ def get_vjp(v):
 
     def test_vjp_call_compiled_backward_fn(self):
         # See PyTorch issue #138422
-        @torch.compile
+        @torch.compile(backend="aot_eager")
         def f(x):
             return x**2
 
@@ -6146,7 +6079,7 @@ def get_vjp(v):
 
     def test_grad_call_compiled_backward_fn(self):
         # See PyTorch issue #138422
-        @torch.compile
+        @torch.compile(backend="aot_eager")
         def f(x):
             return x**2
 
@@ -6901,9 +6834,7 @@ def wrapper_fn(x):
         self.assertEqual(expected, actual)
 
 
-class ActivationCheckpointingTests(
-    torch._dynamo.test_case.TestCaseWithNestedGraphBreaks
-):
+class ActivationCheckpointingTests(torch._dynamo.test_case.TestCase):
     def _validate(self, fn, backend, *args, skip_check=False, fullgraph=True):
         cloned_args = []
         for arg in args:
@@ -7290,11 +7221,10 @@ def false_branch(x):
     # inductor
     "while_loop",  # LoweringException: AssertionError
     "flex_attention",  # LoweringException: AssertionError
-    "flex_attention_backward",  # AssertionError: Input shapes should have M >= 16, N >= 16 and K >= 16
 }
 
 
-class TestHigherOrderOpsOpInfo(torch._dynamo.test_case.TestCaseWithNestedGraphBreaks):
+class TestHigherOrderOpsOpInfo(torch._dynamo.test_case.TestCase):
     @requires_cuda_and_triton
     @parametrize("backend", ("aot_eager", "inductor"))
     @ops(
diff --git a/test/dynamo/test_hooks.py b/test/dynamo/test_hooks.py
index 5e0c90729d3f5..690441c53685e 100644
--- a/test/dynamo/test_hooks.py
+++ b/test/dynamo/test_hooks.py
@@ -1017,7 +1017,6 @@ def run(input):
                     x.grad = None
                     run(i).sum().backward()
 
-    @torch._dynamo.config.patch(inline_inbuilt_nn_modules=True)
     def test_no_recompile_on_same_hook(self):
         cnts = torch._dynamo.testing.CompileCounter()
 
diff --git a/test/dynamo/test_inline_and_install.py b/test/dynamo/test_inline_and_install.py
index e484ebaf9de51..157bea9b9c973 100644
--- a/test/dynamo/test_inline_and_install.py
+++ b/test/dynamo/test_inline_and_install.py
@@ -23,7 +23,6 @@ def make_dynamic_cls(cls):
         cls_prefix,
         suffix,
         (config, "install_free_tensors", True),
-        (config, "inline_inbuilt_nn_modules", True),
         xfail_prop="_expected_failure_inline_and_install",
     )
 
diff --git a/test/dynamo/test_install_free_tensors.py b/test/dynamo/test_install_free_tensors.py
index fd9e14c4c3f76..438ad5c58e2e1 100644
--- a/test/dynamo/test_install_free_tensors.py
+++ b/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 typing import Any
 
 import torch
 import torch._dynamo
@@ -80,11 +80,10 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
 
 
 class InstallParamsAsGraphAttrTests(torch._dynamo.test_case.TestCase):
-    @torch._dynamo.config.patch(inline_inbuilt_nn_modules=True)
     @torch._dynamo.config.patch(install_free_tensors=False)
     def check_num_inputs_and_equality_no_install(
         self,
-        fn_to_compile: Union[torch.nn.Module, Callable],
+        fn_to_compile: torch.nn.Module | Callable,
         expected_num_inline_inputs: int,
         example_inputs: Sequence[Any],
     ) -> None:
@@ -99,11 +98,10 @@ def check_num_inputs_and_equality_no_install(
         self.assertEqual(actual_num_inputs, expected_num_inline_inputs)
         self.assertEqual(opt_fn(*example_inputs), fn_to_compile(*example_inputs))
 
-    @torch._dynamo.config.patch(inline_inbuilt_nn_modules=True)
     @torch._dynamo.config.patch(install_free_tensors=True)
     def check_num_inputs_and_equality_install(
         self,
-        fn_to_compile: Union[torch.nn.Module, Callable],
+        fn_to_compile: torch.nn.Module | Callable,
         expected_num_installed_inputs: int,
         example_inputs: Sequence[Any],
     ) -> None:
@@ -333,7 +331,6 @@ def test_linear_explicit(
 
 
 class InstallParamsWhenExport(torch._dynamo.test_case.TestCase):
-    @torch._dynamo.config.patch(inline_inbuilt_nn_modules=True)
     @torch._dynamo.config.patch(install_free_tensors=True)
     def check_export_matches_expectation(
         self,
@@ -476,7 +473,6 @@ def fn(a: torch.Tensor) -> torch.Tensor:
         inp = torch.randn((5, 5))
         self.check_export_matches_expectation(fn, 1, (inp,))
 
-    @torch._dynamo.config.patch(inline_inbuilt_nn_modules=True)
     @torch._dynamo.config.patch(install_free_tensors=True)
     def test_modify_net_state(self) -> None:
         class Mod(torch.nn.Module):
@@ -508,7 +504,7 @@ def fn(x: list[torch.Tensor]):
         self.check_export_matches_expectation(fn, 2, (inp,))
 
     def test_nested_list_of_tensor(self) -> None:
-        def fn(x: list[Union[list[torch.Tensor], torch.Tensor]]):
+        def fn(x: list[list[torch.Tensor] | torch.Tensor]):
             return x[0][0] + x[1]  # type: ignore[index]
 
         inp = [[torch.tensor([1.3, 3.77, 0.1])], torch.tensor([8.7, 6.23, 9.9])]
diff --git a/test/dynamo/test_len_protocol.py b/test/dynamo/test_len_protocol.py
new file mode 100644
index 0000000000000..6b35b37473f45
--- /dev/null
+++ b/test/dynamo/test_len_protocol.py
@@ -0,0 +1,1297 @@
+# Owner(s): ["module: dynamo"]
+
+"""
+Comprehensive tests for len() builtin and __len__() method protocol in PyTorch Dynamo.
+
+Tests cover:
+- len(obj) builtin calls
+- obj.__len__() method calls
+- type(obj).__len__(obj) unbound method calls (marked as expectedFailure - Dynamo limitation)
+- Various container types: list, tuple, dict, set, frozenset, range, str, Tensor, nn.Module
+- Dict views: keys(), values(), items()
+- User-defined classes with __len__
+"""
+
+import collections
+import dataclasses
+import types
+
+import torch
+import torch._dynamo.test_case
+from torch.testing._internal.common_utils import make_dynamo_test
+
+
+class _BaseSequenceLen:
+    """Base class for testing len() on sequence types (list, tuple)"""
+
+    thetype = None  # Override in subclass
+
+    def setUp(self):
+        if self.thetype is None:
+            self.skipTest("Base class - not meant to be run directly")
+        self.old = torch._dynamo.config.enable_trace_unittest
+        torch._dynamo.config.enable_trace_unittest = True
+        super().setUp()
+
+    def tearDown(self):
+        torch._dynamo.config.enable_trace_unittest = self.old
+        return super().tearDown()
+
+    @make_dynamo_test
+    def test_len_basic(self):
+        seq = self.thetype([1, 2, 3])
+        self.assertEqual(len(seq), 3)
+        self.assertEqual(seq.__len__(), 3)
+
+    @make_dynamo_test
+    def test_len_empty(self):
+        seq = self.thetype([])
+        self.assertEqual(len(seq), 0)
+        self.assertEqual(seq.__len__(), 0)
+
+    @make_dynamo_test
+    def test_len_single_element(self):
+        seq = self.thetype([42])
+        self.assertEqual(len(seq), 1)
+        self.assertEqual(seq.__len__(), 1)
+
+    @make_dynamo_test
+    def test_len_nested(self):
+        inner = self.thetype([1, 2])
+        seq = self.thetype([inner, inner, inner])
+        self.assertEqual(len(seq), 3)
+        self.assertEqual(seq.__len__(), 3)
+
+    @make_dynamo_test
+    def test_len_with_tensors(self):
+        seq = self.thetype([torch.tensor(1), torch.tensor(2), torch.tensor(3)])
+        self.assertEqual(len(seq), 3)
+        self.assertEqual(seq.__len__(), 3)
+
+    @make_dynamo_test
+    def test_len_with_mixed_types(self):
+        seq = self.thetype([1, "hello", 3.14, torch.tensor(4)])
+        self.assertEqual(len(seq), 4)
+        self.assertEqual(seq.__len__(), 4)
+
+    @make_dynamo_test
+    def test_len_large(self):
+        seq = self.thetype(range(100))
+        self.assertEqual(len(seq), 100)
+        self.assertEqual(seq.__len__(), 100)
+
+
+class TestListLen(_BaseSequenceLen, torch._dynamo.test_case.TestCase):
+    """Tests for len() on list objects"""
+
+    thetype = list
+
+
+class TestTupleLen(_BaseSequenceLen, torch._dynamo.test_case.TestCase):
+    """Tests for len() on tuple objects"""
+
+    thetype = tuple
+
+
+class _BaseMappingLen:
+    """Base class for testing len() on mapping types (dict, OrderedDict)"""
+
+    def get_mapping(self, items):
+        """Override in subclass to return appropriate mapping type"""
+        return dict(items)
+
+    def setUp(self):
+        self.old = torch._dynamo.config.enable_trace_unittest
+        torch._dynamo.config.enable_trace_unittest = True
+        super().setUp()
+
+    def tearDown(self):
+        torch._dynamo.config.enable_trace_unittest = self.old
+        return super().tearDown()
+
+    @make_dynamo_test
+    def test_len_basic(self):
+        d = self.get_mapping({1: "a", 2: "b", 3: "c"}.items())
+        self.assertEqual(len(d), 3)
+        self.assertEqual(d.__len__(), 3)
+
+    @make_dynamo_test
+    def test_len_empty(self):
+        d = self.get_mapping({}.items())
+        self.assertEqual(len(d), 0)
+        self.assertEqual(d.__len__(), 0)
+
+    @make_dynamo_test
+    def test_len_single_entry(self):
+        d = self.get_mapping({"key": "value"}.items())
+        self.assertEqual(len(d), 1)
+        self.assertEqual(d.__len__(), 1)
+
+    @make_dynamo_test
+    def test_len_string_keys(self):
+        d = self.get_mapping({"one": 1, "two": 2, "three": 3}.items())
+        self.assertEqual(len(d), 3)
+        self.assertEqual(d.__len__(), 3)
+
+    @make_dynamo_test
+    def test_len_int_keys(self):
+        d = self.get_mapping({1: "one", 2: "two", 3: "three", 4: "four"}.items())
+        self.assertEqual(len(d), 4)
+        self.assertEqual(d.__len__(), 4)
+
+    @make_dynamo_test
+    def test_len_with_tensor_values(self):
+        d = self.get_mapping(
+            {"a": torch.tensor(1), "b": torch.tensor(2), "c": torch.tensor(3)}.items()
+        )
+        self.assertEqual(len(d), 3)
+        self.assertEqual(d.__len__(), 3)
+
+    @make_dynamo_test
+    def test_len_large(self):
+        d = self.get_mapping({i: i * 2 for i in range(20)}.items())
+        self.assertEqual(len(d), 20)
+        self.assertEqual(d.__len__(), 20)
+
+
+class TestDictLen(_BaseMappingLen, torch._dynamo.test_case.TestCase):
+    """Tests for len() on dict objects"""
+
+
+class TestOrderedDictLen(_BaseMappingLen, torch._dynamo.test_case.TestCase):
+    """Tests for len() on OrderedDict objects"""
+
+    def get_mapping(self, items):
+        return collections.OrderedDict(items)
+
+
+class TestDefaultDictLen(_BaseMappingLen, torch._dynamo.test_case.TestCase):
+    """Tests for len() on defaultdict objects"""
+
+    def get_mapping(self, items):
+        d = collections.defaultdict(int)
+        for k, v in items:
+            d[k] = v
+        return d
+
+
+class _BaseSetLen:
+    """Base class for testing len() on set types"""
+
+    __test__ = False  # Prevent pytest from collecting this as a test class
+
+    def get_set(self, items):
+        """Override in subclass to return appropriate set type"""
+        return set(items)
+
+    def setUp(self):
+        self.old = torch._dynamo.config.enable_trace_unittest
+        torch._dynamo.config.enable_trace_unittest = True
+        super().setUp()
+
+    def tearDown(self):
+        torch._dynamo.config.enable_trace_unittest = self.old
+        return super().tearDown()
+
+    @make_dynamo_test
+    def test_len_basic(self):
+        s = self.get_set([1, 2, 3])
+        self.assertEqual(len(s), 3)
+        self.assertEqual(s.__len__(), 3)
+
+    @make_dynamo_test
+    def test_len_empty(self):
+        s = self.get_set([])
+        self.assertEqual(len(s), 0)
+        self.assertEqual(s.__len__(), 0)
+
+    @make_dynamo_test
+    def test_len_single_element(self):
+        s = self.get_set([42])
+        self.assertEqual(len(s), 1)
+        self.assertEqual(s.__len__(), 1)
+
+    @make_dynamo_test
+    def test_len_with_strings(self):
+        s = self.get_set(["a", "b", "c", "d"])
+        self.assertEqual(len(s), 4)
+        self.assertEqual(s.__len__(), 4)
+
+    @make_dynamo_test
+    def test_len_with_duplicates(self):
+        # Set constructor deduplicates
+        s = self.get_set([1, 2, 2, 3, 3, 3])
+        self.assertEqual(len(s), 3)
+        self.assertEqual(s.__len__(), 3)
+
+    @make_dynamo_test
+    def test_len_large(self):
+        s = self.get_set(range(50))
+        self.assertEqual(len(s), 50)
+        self.assertEqual(s.__len__(), 50)
+
+
+class TestSetLen(_BaseSetLen, torch._dynamo.test_case.TestCase):
+    """Tests for len() on set objects"""
+
+    def get_set(self, items):
+        return set(items)
+
+
+class TestFrozenSetLen(_BaseSetLen, torch._dynamo.test_case.TestCase):
+    """Tests for len() on frozenset objects"""
+
+    def get_set(self, items):
+        return frozenset(items)
+
+
+class TestRangeLen(torch._dynamo.test_case.TestCase):
+    """Tests for len() on range objects"""
+
+    def setUp(self):
+        self.old = torch._dynamo.config.enable_trace_unittest
+        torch._dynamo.config.enable_trace_unittest = True
+        super().setUp()
+
+    def tearDown(self):
+        torch._dynamo.config.enable_trace_unittest = self.old
+        return super().tearDown()
+
+    @make_dynamo_test
+    def test_len_basic(self):
+        r = range(5)
+        self.assertEqual(len(r), 5)
+        self.assertEqual(r.__len__(), 5)
+        self.assertEqual(range.__len__(r), 5)
+
+    @make_dynamo_test
+    def test_len_with_start_stop(self):
+        r = range(5, 15)
+        self.assertEqual(len(r), 10)
+        self.assertEqual(r.__len__(), 10)
+        self.assertEqual(range.__len__(r), 10)
+
+    @make_dynamo_test
+    def test_len_with_step(self):
+        r = range(0, 10, 2)
+        self.assertEqual(len(r), 5)
+        self.assertEqual(r.__len__(), 5)
+        self.assertEqual(range.__len__(r), 5)
+
+    @make_dynamo_test
+    def test_len_negative_step(self):
+        r = range(10, 0, -1)
+        self.assertEqual(len(r), 10)
+        self.assertEqual(r.__len__(), 10)
+        self.assertEqual(range.__len__(r), 10)
+
+    @make_dynamo_test
+    def test_len_empty(self):
+        r = range(5, 5)
+        self.assertEqual(len(r), 0)
+        self.assertEqual(r.__len__(), 0)
+        self.assertEqual(range.__len__(r), 0)
+
+
+class TestStringLen(torch._dynamo.test_case.TestCase):
+    """Tests for len() on string objects"""
+
+    def setUp(self):
+        self.old = torch._dynamo.config.enable_trace_unittest
+        torch._dynamo.config.enable_trace_unittest = True
+        super().setUp()
+
+    def tearDown(self):
+        torch._dynamo.config.enable_trace_unittest = self.old
+        return super().tearDown()
+
+    @make_dynamo_test
+    def test_len_empty(self):
+        s = ""
+        self.assertEqual(len(s), 0)
+        self.assertEqual(s.__len__(), 0)
+        self.assertEqual(str.__len__(s), 0)
+
+    @make_dynamo_test
+    def test_len_single_char(self):
+        s = "a"
+        self.assertEqual(len(s), 1)
+        self.assertEqual(s.__len__(), 1)
+        self.assertEqual(str.__len__(s), 1)
+
+    @make_dynamo_test
+    def test_len_multiple_chars(self):
+        s = "hello"
+        self.assertEqual(len(s), 5)
+        self.assertEqual(s.__len__(), 5)
+        self.assertEqual(str.__len__(s), 5)
+
+    @make_dynamo_test
+    def test_len_with_spaces(self):
+        s = "hello world"
+        self.assertEqual(len(s), 11)
+        self.assertEqual(s.__len__(), 11)
+        self.assertEqual(str.__len__(s), 11)
+
+
+class TestTensorLen(torch._dynamo.test_case.TestCase):
+    """Tests for len() on torch.Tensor objects"""
+
+    def setUp(self):
+        self.old = torch._dynamo.config.enable_trace_unittest
+        torch._dynamo.config.enable_trace_unittest = True
+        super().setUp()
+
+    def tearDown(self):
+        torch._dynamo.config.enable_trace_unittest = self.old
+        return super().tearDown()
+
+    @make_dynamo_test
+    def test_len_1d(self):
+        t = torch.tensor([1, 2, 3, 4, 5])
+        self.assertEqual(len(t), 5)
+        self.assertEqual(t.__len__(), 5)
+        self.assertEqual(torch.Tensor.__len__(t), 5)
+
+    @make_dynamo_test
+    def test_len_2d(self):
+        t = torch.tensor([[1, 2], [3, 4], [5, 6]])
+        self.assertEqual(len(t), 3)
+        self.assertEqual(t.__len__(), 3)
+        self.assertEqual(torch.Tensor.__len__(t), 3)
+
+    @make_dynamo_test
+    def test_len_3d(self):
+        t = torch.randn(4, 5, 6)
+        self.assertEqual(len(t), 4)
+        self.assertEqual(t.__len__(), 4)
+        self.assertEqual(torch.Tensor.__len__(t), 4)
+
+    @make_dynamo_test
+    def test_len_empty(self):
+        t = torch.tensor([])
+        self.assertEqual(len(t), 0)
+        self.assertEqual(t.__len__(), 0)
+        self.assertEqual(torch.Tensor.__len__(t), 0)
+
+    @make_dynamo_test
+    def test_len_large_batch(self):
+        t = torch.randn(100, 5, 5)
+        self.assertEqual(len(t), 100)
+        self.assertEqual(t.__len__(), 100)
+        self.assertEqual(torch.Tensor.__len__(t), 100)
+
+    @make_dynamo_test
+    def test_len_different_dtypes(self):
+        t = torch.tensor([1, 2, 3], dtype=torch.float32)
+        self.assertEqual(len(t), 3)
+        self.assertEqual(t.__len__(), 3)
+        self.assertEqual(torch.Tensor.__len__(t), 3)
+
+
+class TestNNModuleLen(torch._dynamo.test_case.TestCase):
+    """Tests for len() on torch.nn module containers"""
+
+    def setUp(self):
+        self.old = torch._dynamo.config.enable_trace_unittest
+        torch._dynamo.config.enable_trace_unittest = True
+        super().setUp()
+
+        # Pre-construct nn.Module instances outside compiled regions
+        self.seq_3layers = torch.nn.Sequential(
+            torch.nn.Linear(10, 20), torch.nn.ReLU(), torch.nn.Linear(20, 5)
+        )
+        self.empty_seq = torch.nn.Sequential()
+        self.ml_2modules = torch.nn.ModuleList(
+            [torch.nn.Linear(10, 20), torch.nn.Linear(20, 30)]
+        )
+        self.empty_ml = torch.nn.ModuleList()
+        self.md_2modules = torch.nn.ModuleDict(
+            {"layer1": torch.nn.Linear(10, 20), "layer2": torch.nn.Linear(20, 5)}
+        )
+        self.empty_md = torch.nn.ModuleDict()
+        self.seq_5layers = torch.nn.Sequential(
+            *[torch.nn.Linear(10, 10) for _ in range(5)]
+        )
+
+    def tearDown(self):
+        torch._dynamo.config.enable_trace_unittest = self.old
+        return super().tearDown()
+
+    @make_dynamo_test
+    def test_len_sequential(self):
+        seq = self.seq_3layers
+        self.assertEqual(len(seq), 3)
+        self.assertEqual(seq.__len__(), 3)
+
+    @make_dynamo_test
+    def test_len_empty_sequential(self):
+        seq = self.empty_seq
+        self.assertEqual(len(seq), 0)
+        self.assertEqual(seq.__len__(), 0)
+
+    @make_dynamo_test
+    def test_len_module_list(self):
+        ml = self.ml_2modules
+        self.assertEqual(len(ml), 2)
+        self.assertEqual(ml.__len__(), 2)
+
+    @make_dynamo_test
+    def test_len_empty_module_list(self):
+        ml = self.empty_ml
+        self.assertEqual(len(ml), 0)
+        self.assertEqual(ml.__len__(), 0)
+
+    @make_dynamo_test
+    def test_len_module_dict(self):
+        md = self.md_2modules
+        self.assertEqual(len(md), 2)
+        self.assertEqual(md.__len__(), 2)
+
+    @make_dynamo_test
+    def test_len_empty_module_dict(self):
+        md = self.empty_md
+        self.assertEqual(len(md), 0)
+        self.assertEqual(md.__len__(), 0)
+
+    @make_dynamo_test
+    def test_len_sequential_with_multiple_layers(self):
+        seq = self.seq_5layers
+        self.assertEqual(len(seq), 5)
+        self.assertEqual(seq.__len__(), 5)
+
+
+class TestDictViewLen(torch._dynamo.test_case.TestCase):
+    """Tests for len() on dict view objects (keys, values, items)"""
+
+    def setUp(self):
+        self.old = torch._dynamo.config.enable_trace_unittest
+        torch._dynamo.config.enable_trace_unittest = True
+        super().setUp()
+
+    def tearDown(self):
+        torch._dynamo.config.enable_trace_unittest = self.old
+        return super().tearDown()
+
+    @make_dynamo_test
+    def test_len_dict_keys_view(self):
+        d = {"a": 1, "b": 2, "c": 3}
+        keys = d.keys()
+        self.assertEqual(len(keys), 3)
+        self.assertEqual(keys.__len__(), 3)
+
+    @make_dynamo_test
+    def test_len_dict_values_view(self):
+        d = {"a": 1, "b": 2, "c": 3}
+        values = d.values()
+        self.assertEqual(len(values), 3)
+        self.assertEqual(values.__len__(), 3)
+
+    @make_dynamo_test
+    def test_len_dict_items_view(self):
+        d = {"x": 10, "y": 20, "z": 30}
+        items = d.items()
+        self.assertEqual(len(items), 3)
+        self.assertEqual(items.__len__(), 3)
+
+    @make_dynamo_test
+    def test_len_dict_keys_empty(self):
+        d = {}
+        keys = d.keys()
+        self.assertEqual(len(keys), 0)
+        self.assertEqual(keys.__len__(), 0)
+
+    @make_dynamo_test
+    def test_len_dict_values_empty(self):
+        d = {}
+        values = d.values()
+        self.assertEqual(len(values), 0)
+        self.assertEqual(values.__len__(), 0)
+
+    @make_dynamo_test
+    def test_len_dict_items_empty(self):
+        d = {}
+        items = d.items()
+        self.assertEqual(len(items), 0)
+        self.assertEqual(items.__len__(), 0)
+
+    @make_dynamo_test
+    def test_len_dict_keys_single_entry(self):
+        """Test len() on dict.keys() with single entry"""
+        d = {"key": "value"}
+        keys = d.keys()
+        self.assertEqual(len(keys), 1)
+        self.assertEqual(keys.__len__(), 1)
+
+    @make_dynamo_test
+    def test_len_dict_values_single_entry(self):
+        """Test len() on dict.values() with single entry"""
+        d = {"key": "value"}
+        values = d.values()
+        self.assertEqual(len(values), 1)
+        self.assertEqual(values.__len__(), 1)
+
+    @make_dynamo_test
+    def test_len_dict_items_single_entry(self):
+        """Test len() on dict.items() with single entry"""
+        d = {"key": "value"}
+        items = d.items()
+        self.assertEqual(len(items), 1)
+        self.assertEqual(items.__len__(), 1)
+
+    @make_dynamo_test
+    def test_len_dict_keys_int_keys(self):
+        """Test len() on dict.keys() with integer keys"""
+        d = {1: "one", 2: "two", 3: "three", 4: "four"}
+        keys = d.keys()
+        self.assertEqual(len(keys), 4)
+        self.assertEqual(keys.__len__(), 4)
+
+    @make_dynamo_test
+    def test_len_dict_values_tensor_values(self):
+        """Test len() on dict.values() with tensor values"""
+        d = {"a": torch.tensor(1), "b": torch.tensor(2), "c": torch.tensor(3)}
+        values = d.values()
+        self.assertEqual(len(values), 3)
+        self.assertEqual(values.__len__(), 3)
+
+    @make_dynamo_test
+    def test_len_dict_items_mixed_types(self):
+        """Test len() on dict.items() with mixed key/value types"""
+        d = {"str": "value", 42: torch.tensor(1), (1, 2): "tuple_key"}
+        items = d.items()
+        self.assertEqual(len(items), 3)
+        self.assertEqual(items.__len__(), 3)
+
+    @make_dynamo_test
+    def test_len_dict_keys_large(self):
+        """Test len() on dict.keys() with large number of entries"""
+        d = {i: i * 2 for i in range(50)}
+        keys = d.keys()
+        self.assertEqual(len(keys), 50)
+        self.assertEqual(keys.__len__(), 50)
+
+    @make_dynamo_test
+    def test_len_dict_values_large(self):
+        """Test len() on dict.values() with large number of entries"""
+        d = {i: i * 2 for i in range(50)}
+        values = d.values()
+        self.assertEqual(len(values), 50)
+        self.assertEqual(values.__len__(), 50)
+
+    @make_dynamo_test
+    def test_len_dict_items_large(self):
+        """Test len() on dict.items() with large number of entries"""
+        d = {i: i * 2 for i in range(50)}
+        items = d.items()
+        self.assertEqual(len(items), 50)
+        self.assertEqual(items.__len__(), 50)
+
+
+# User-defined classes for TestUserDefinedLen
+class CustomList:
+    def __init__(self, items):
+        self.items = items
+
+    def __len__(self):
+        return len(self.items)
+
+
+class CustomContainer:
+    def __len__(self):
+        return 0
+
+
+class Container:
+    def __init__(self, size):
+        self.size = size
+
+    def __len__(self):
+        return self.size
+
+
+class FixedSize:
+    def __len__(self):
+        return 10
+
+
+class ListWrapper:
+    def __init__(self):
+        self.data = [1, 2, 3, 4, 5]
+
+    def __len__(self):
+        return len(self.data)
+
+
+class ListSubclassCustomLen(list):
+    def __len__(self):
+        return super().__len__() * 2
+
+
+class CustomMapping:
+    """A user-defined mapping (dict-like) class"""
+
+    def __init__(self, data):
+        self._data = dict(data) if not isinstance(data, dict) else data
+
+    def __len__(self):
+        return len(self._data)
+
+    def __getitem__(self, key):
+        return self._data[key]
+
+    def keys(self):
+        return self._data.keys()
+
+    def values(self):
+        return self._data.values()
+
+    def items(self):
+        return self._data.items()
+
+
+class CustomMappingSubclass(CustomMapping):
+    """A subclass of CustomMapping"""
+
+    def __len__(self):
+        # Custom len implementation (e.g., filtered length)
+        return super().__len__() + 1
+
+
+class SimpleDictLike:
+    """Minimal dict-like class with just __len__"""
+
+    def __init__(self, size):
+        self.size = size
+
+    def __len__(self):
+        return self.size
+
+
+class TupleSubclassCustomLen(tuple):
+    __slots__ = ()
+
+    def __len__(self):
+        return super().__len__() + 1
+
+
+class DictSubclassCustomLen(dict):
+    def __len__(self):
+        return super().__len__() - 1
+
+
+class SetSubclassCustomLen(set):
+    def __len__(self):
+        return 0
+
+
+class TestUserDefinedLen(torch._dynamo.test_case.TestCase):
+    """Tests for len() on user-defined classes with __len__"""
+
+    def setUp(self):
+        self.old = torch._dynamo.config.enable_trace_unittest
+        torch._dynamo.config.enable_trace_unittest = True
+        super().setUp()
+
+    def tearDown(self):
+        torch._dynamo.config.enable_trace_unittest = self.old
+        return super().tearDown()
+
+    @make_dynamo_test
+    def test_len_simple_custom_class(self):
+        obj = CustomList([1, 2, 3])
+        self.assertEqual(len(obj), 3)
+        self.assertEqual(obj.__len__(), 3)
+        self.assertEqual(CustomList.__len__(obj), 3)
+
+    @make_dynamo_test
+    def test_len_custom_class_empty(self):
+        obj = CustomContainer()
+        self.assertEqual(len(obj), 0)
+        self.assertEqual(obj.__len__(), 0)
+        self.assertEqual(CustomContainer.__len__(obj), 0)
+
+    @make_dynamo_test
+    def test_len_custom_class_with_properties(self):
+        obj = Container(42)
+        self.assertEqual(len(obj), 42)
+        self.assertEqual(obj.__len__(), 42)
+        self.assertEqual(Container.__len__(obj), 42)
+
+    @make_dynamo_test
+    def test_len_custom_class_constant_return(self):
+        obj = FixedSize()
+        self.assertEqual(len(obj), 10)
+        self.assertEqual(obj.__len__(), 10)
+        self.assertEqual(FixedSize.__len__(obj), 10)
+
+    @make_dynamo_test
+    def test_len_custom_class_with_list_attr(self):
+        obj = ListWrapper()
+        self.assertEqual(len(obj), 5)
+        self.assertEqual(obj.__len__(), 5)
+        self.assertEqual(ListWrapper.__len__(obj), 5)
+
+
+class TestSubclassOverloadedLen(torch._dynamo.test_case.TestCase):
+    """Tests for custom classes that inherit from builtins and overload __len__"""
+
+    def setUp(self):
+        self.old = torch._dynamo.config.enable_trace_unittest
+        torch._dynamo.config.enable_trace_unittest = True
+        super().setUp()
+
+    def tearDown(self):
+        torch._dynamo.config.enable_trace_unittest = self.old
+        return super().tearDown()
+
+    @make_dynamo_test
+    def test_list_subclass_custom_len(self):
+        obj = ListSubclassCustomLen([1, 2, 3])
+        self.assertEqual(len(obj), 6)
+        self.assertEqual(obj.__len__(), 6)
+        self.assertEqual(ListSubclassCustomLen.__len__(obj), 6)
+
+    @make_dynamo_test
+    def test_tuple_subclass_custom_len(self):
+        obj = TupleSubclassCustomLen([1, 2, 3])
+        self.assertEqual(len(obj), 4)
+        self.assertEqual(obj.__len__(), 4)
+        self.assertEqual(TupleSubclassCustomLen.__len__(obj), 4)
+
+    @make_dynamo_test
+    def test_dict_subclass_custom_len(self):
+        obj = DictSubclassCustomLen({"a": 1, "b": 2, "c": 3})
+        self.assertEqual(len(obj), 2)
+        self.assertEqual(obj.__len__(), 2)
+        self.assertEqual(DictSubclassCustomLen.__len__(obj), 2)
+
+    @make_dynamo_test
+    def test_set_subclass_custom_len(self):
+        obj = SetSubclassCustomLen([1, 2, 3])
+        self.assertEqual(len(obj), 0)
+        self.assertEqual(obj.__len__(), 0)
+        self.assertEqual(SetSubclassCustomLen.__len__(obj), 0)
+
+
+class DescriptorLenClass:
+    """Test class with __len__ as a regular instance method"""
+
+    def __len__(self):
+        """Regular instance method __len__ - should return 40"""
+        return 40
+
+
+class PartialLenClass:
+    """Class where __len__ is a lambda/callable object"""
+
+    def __init__(self):
+        self._items = [1, 2, 3, 4, 5]
+
+    def __len__(self):
+        return len(self._items)
+
+
+class StaticMethodLenClass:
+    """Test class where __len__ is a staticmethod (unusual, likely to fail)"""
+
+    @staticmethod
+    def __len__():
+        """Staticmethod __len__ - unusual pattern"""
+        return 10
+
+
+class ClassMethodLenClass:
+    """Test class where __len__ is a classmethod (unusual, likely to fail)"""
+
+    @classmethod
+    def __len__(cls):
+        """Classmethod __len__ - unusual pattern"""
+        return 20
+
+
+class CustomDescriptorLenClass:
+    """Test class where __len__ is a custom descriptor"""
+
+    class CustomDescriptorLen:
+        """Custom descriptor that implements __get__ method"""
+
+        def __get__(self, obj, objtype=None):
+            """Descriptor protocol: return a callable that returns len"""
+            if obj is None:
+                return self
+            return lambda: 50
+
+    __len__ = CustomDescriptorLen()
+
+
+class TestUserDefinedMappingLen(torch._dynamo.test_case.TestCase):
+    """Tests for len() on user-defined mapping (dict-like) classes"""
+
+    def setUp(self):
+        self.old = torch._dynamo.config.enable_trace_unittest
+        torch._dynamo.config.enable_trace_unittest = True
+        super().setUp()
+
+    def tearDown(self):
+        torch._dynamo.config.enable_trace_unittest = self.old
+        return super().tearDown()
+
+    @make_dynamo_test
+    def test_len_custom_mapping_basic(self):
+        """Test len() on a basic custom mapping class"""
+        m = CustomMapping({"a": 1, "b": 2, "c": 3})
+        self.assertEqual(len(m), 3)
+        self.assertEqual(m.__len__(), 3)
+
+    @make_dynamo_test
+    def test_len_custom_mapping_empty(self):
+        """Test len() on an empty custom mapping"""
+        m = CustomMapping({})
+        self.assertEqual(len(m), 0)
+        self.assertEqual(m.__len__(), 0)
+
+    @make_dynamo_test
+    def test_len_custom_mapping_single_entry(self):
+        """Test len() on a custom mapping with single entry"""
+        m = CustomMapping({"key": "value"})
+        self.assertEqual(len(m), 1)
+        self.assertEqual(m.__len__(), 1)
+
+    @make_dynamo_test
+    def test_len_custom_mapping_string_keys(self):
+        """Test len() on a custom mapping with string keys"""
+        m = CustomMapping({"one": 1, "two": 2, "three": 3})
+        self.assertEqual(len(m), 3)
+        self.assertEqual(m.__len__(), 3)
+
+    @make_dynamo_test
+    def test_len_custom_mapping_int_keys(self):
+        """Test len() on a custom mapping with int keys"""
+        m = CustomMapping({1: "one", 2: "two", 3: "three", 4: "four"})
+        self.assertEqual(len(m), 4)
+        self.assertEqual(m.__len__(), 4)
+
+    @make_dynamo_test
+    def test_len_custom_mapping_with_tensor_values(self):
+        """Test len() on a custom mapping with tensor values"""
+        m = CustomMapping(
+            {"a": torch.tensor(1), "b": torch.tensor(2), "c": torch.tensor(3)}
+        )
+        self.assertEqual(len(m), 3)
+        self.assertEqual(m.__len__(), 3)
+
+    @make_dynamo_test
+    def test_len_custom_mapping_large(self):
+        """Test len() on a large custom mapping"""
+        m = CustomMapping({i: i * 2 for i in range(20)})
+        self.assertEqual(len(m), 20)
+        self.assertEqual(m.__len__(), 20)
+
+    @make_dynamo_test
+    def test_len_custom_mapping_subclass(self):
+        """Test len() on a subclass of CustomMapping"""
+        m = CustomMappingSubclass({"a": 1, "b": 2})
+        # CustomMappingSubclass.__len__ returns len + 1
+        self.assertEqual(len(m), 3)
+        self.assertEqual(m.__len__(), 3)
+
+    @make_dynamo_test
+    def test_len_simple_dict_like(self):
+        """Test len() on a minimal dict-like class"""
+        m = SimpleDictLike(42)
+        self.assertEqual(len(m), 42)
+        self.assertEqual(m.__len__(), 42)
+
+
+class TestDescriptorLenImpl(torch._dynamo.test_case.TestCase):
+    """Test that len_impl handles descriptor-based __len__ correctly"""
+
+    def setUp(self):
+        self.old = torch._dynamo.config.enable_trace_unittest
+        torch._dynamo.config.enable_trace_unittest = True
+        super().setUp()
+
+    def tearDown(self):
+        torch._dynamo.config.enable_trace_unittest = self.old
+        return super().tearDown()
+
+    @make_dynamo_test
+    def test_regular_instance_len(self):
+        """Test regular instance method __len__"""
+        obj = DescriptorLenClass()
+
+        # Regular instance methods should work fine
+        self.assertEqual(len(obj), 40)
+        self.assertEqual(obj.__len__(), 40)
+
+    @make_dynamo_test
+    def test_partial_callable_len(self):
+        """Test lambda/callable as __len__"""
+        obj = PartialLenClass()
+
+        # Callable __len__ should work
+        self.assertEqual(len(obj), 5)
+        self.assertEqual(obj.__len__(), 5)
+
+    @make_dynamo_test
+    def test_staticmethod_len_works(self):
+        """Test that staticmethod as __len__ actually works (unusual but supported)
+
+        Staticmethods are descriptors that don't bind. CPython resolves the
+        descriptor and calls the underlying function without passing self.
+        """
+        obj = StaticMethodLenClass()
+
+        # Surprisingly, CPython's descriptor protocol makes this work
+        # staticmethod.__get__ returns the unwrapped function
+        self.assertEqual(len(obj), 10)
+
+    @make_dynamo_test
+    def test_classmethod_len_works(self):
+        """Test that classmethod as __len__ actually works (unusual but supported)
+
+        Classmethods are descriptors that bind to the class. CPython's
+        descriptor protocol handles this and passes the class instead of instance.
+        """
+        obj = ClassMethodLenClass()
+
+        # Surprisingly, CPython's descriptor protocol makes this work
+        # classmethod.__get__ returns a bound method with the class
+        self.assertEqual(len(obj), 20)
+
+    @make_dynamo_test
+    def test_custom_descriptor_len(self):
+        """Test custom descriptor with __get__ method as __len__
+
+        Custom descriptors implement the descriptor protocol via __get__.
+        The descriptor is resolved and the returned callable is used as __len__.
+        """
+        obj = CustomDescriptorLenClass()
+
+        # Custom descriptor's __get__ returns a callable that returns 50
+        self.assertEqual(len(obj), 50)
+
+
+class TestRaisesTypeError(torch._dynamo.test_case.TestCase):
+    """Tests for types that don't support len() - should raise TypeError like Python"""
+
+    def setUp(self):
+        self.old = torch._dynamo.config.enable_trace_unittest
+        torch._dynamo.config.enable_trace_unittest = True
+        super().setUp()
+
+    def tearDown(self):
+        torch._dynamo.config.enable_trace_unittest = self.old
+        return super().tearDown()
+
+    @make_dynamo_test
+    def test_len_slice_raises_type_error(self):
+        """slice objects do not support len() - should raise TypeError"""
+        s = slice(1, 5)
+        with self.assertRaises(TypeError):
+            len(s)
+
+    @make_dynamo_test
+    def test_len_slice_with_step_raises_type_error(self):
+        """slice with step also raises TypeError"""
+        s = slice(0, 10, 2)
+        with self.assertRaises(TypeError):
+            len(s)
+
+    @make_dynamo_test
+    def test_len_list_iterator_raises_type_error(self):
+        """list iterator does not support len() - should raise TypeError"""
+        it = iter([1, 2, 3])
+        with self.assertRaises(TypeError):
+            len(it)
+
+    @make_dynamo_test
+    def test_len_empty_list_iterator_raises_type_error(self):
+        """empty list iterator also raises TypeError"""
+        it = iter([])
+        with self.assertRaises(TypeError):
+            len(it)
+
+    @make_dynamo_test
+    def test_len_tuple_iterator_raises_type_error(self):
+        """tuple iterator does not support len() - should raise TypeError"""
+        it = iter((1, 2, 3))
+        with self.assertRaises(TypeError):
+            len(it)
+
+    @make_dynamo_test
+    def test_len_range_iterator_raises_type_error(self):
+        """range iterator does not support len() - should raise TypeError"""
+        it = iter(range(5))
+        with self.assertRaises(TypeError):
+            len(it)
+
+    @make_dynamo_test
+    def test_len_dict_iterator_raises_type_error(self):
+        """dict iterator (keys) does not support len() - should raise TypeError"""
+        d = {"a": 1, "b": 2}
+        it = iter(d)
+        with self.assertRaises(TypeError):
+            len(it)
+
+
+class TestDequeLen(torch._dynamo.test_case.TestCase):
+    """Tests for len() on collections.deque objects"""
+
+    def setUp(self):
+        self.old = torch._dynamo.config.enable_trace_unittest
+        torch._dynamo.config.enable_trace_unittest = True
+        super().setUp()
+
+    def tearDown(self):
+        torch._dynamo.config.enable_trace_unittest = self.old
+        return super().tearDown()
+
+    @make_dynamo_test
+    def test_len_basic(self):
+        d = collections.deque([1, 2, 3])
+        self.assertEqual(len(d), 3)
+        self.assertEqual(d.__len__(), 3)
+
+    @make_dynamo_test
+    def test_len_empty(self):
+        d = collections.deque([])
+        self.assertEqual(len(d), 0)
+        self.assertEqual(d.__len__(), 0)
+
+    @make_dynamo_test
+    def test_len_single_element(self):
+        d = collections.deque([42])
+        self.assertEqual(len(d), 1)
+        self.assertEqual(d.__len__(), 1)
+
+    @make_dynamo_test
+    def test_len_with_maxlen(self):
+        d = collections.deque([1, 2, 3, 4, 5], maxlen=3)
+        self.assertEqual(len(d), 3)
+        self.assertEqual(d.__len__(), 3)
+
+    @make_dynamo_test
+    def test_len_with_strings(self):
+        d = collections.deque(["a", "b", "c", "d"])
+        self.assertEqual(len(d), 4)
+        self.assertEqual(d.__len__(), 4)
+
+    @make_dynamo_test
+    def test_len_large(self):
+        d = collections.deque(range(50))
+        self.assertEqual(len(d), 50)
+        self.assertEqual(d.__len__(), 50)
+
+
+class TestMappingProxyLen(torch._dynamo.test_case.TestCase):
+    """Tests for len() on types.MappingProxyType objects"""
+
+    def setUp(self):
+        self.old = torch._dynamo.config.enable_trace_unittest
+        torch._dynamo.config.enable_trace_unittest = True
+        super().setUp()
+
+    def tearDown(self):
+        torch._dynamo.config.enable_trace_unittest = self.old
+        return super().tearDown()
+
+    @make_dynamo_test
+    def test_len_basic(self):
+        """Test len() on a basic MappingProxyType"""
+        d = types.MappingProxyType({"a": 1, "b": 2, "c": 3})
+        self.assertEqual(len(d), 3)
+        self.assertEqual(d.__len__(), 3)
+
+    @make_dynamo_test
+    def test_len_empty(self):
+        """Test len() on an empty MappingProxyType"""
+        d = types.MappingProxyType({})
+        self.assertEqual(len(d), 0)
+        self.assertEqual(d.__len__(), 0)
+
+    @make_dynamo_test
+    def test_len_single_entry(self):
+        """Test len() on a MappingProxyType with single entry"""
+        d = types.MappingProxyType({"key": "value"})
+        self.assertEqual(len(d), 1)
+        self.assertEqual(d.__len__(), 1)
+
+    @make_dynamo_test
+    def test_len_string_keys(self):
+        """Test len() on a MappingProxyType with string keys"""
+        d = types.MappingProxyType({"one": 1, "two": 2, "three": 3})
+        self.assertEqual(len(d), 3)
+        self.assertEqual(d.__len__(), 3)
+
+    @make_dynamo_test
+    def test_len_int_keys(self):
+        """Test len() on a MappingProxyType with int keys"""
+        d = types.MappingProxyType({1: "one", 2: "two", 3: "three", 4: "four"})
+        self.assertEqual(len(d), 4)
+        self.assertEqual(d.__len__(), 4)
+
+    @make_dynamo_test
+    def test_len_with_tensor_values(self):
+        """Test len() on a MappingProxyType with tensor values"""
+        d = types.MappingProxyType(
+            {"a": torch.tensor(1), "b": torch.tensor(2), "c": torch.tensor(3)}
+        )
+        self.assertEqual(len(d), 3)
+        self.assertEqual(d.__len__(), 3)
+
+    @make_dynamo_test
+    def test_len_large(self):
+        """Test len() on a large MappingProxyType"""
+        d = types.MappingProxyType({i: i * 2 for i in range(20)})
+        self.assertEqual(len(d), 20)
+        self.assertEqual(d.__len__(), 20)
+
+
+class MetaclassWithLen(type):
+    """A metaclass that defines __len__ on the class itself"""
+
+    def __len__(cls):
+        """Return the number of items defined in the metaclass"""
+        return 5
+
+
+class SimpleMetaclassClass(metaclass=MetaclassWithLen):
+    """A class using the MetaclassWithLen metaclass"""
+
+
+class TestMetaclassLen(torch._dynamo.test_case.TestCase):
+    """Tests for len() on metaclasses, classmethods, staticmethods, and properties"""
+
+    def setUp(self):
+        self.old = torch._dynamo.config.enable_trace_unittest
+        torch._dynamo.config.enable_trace_unittest = True
+        super().setUp()
+
+    def tearDown(self):
+        torch._dynamo.config.enable_trace_unittest = self.old
+        return super().tearDown()
+
+    @make_dynamo_test
+    def test_metaclass_len_basic(self):
+        """Test len() on a class with __len__ defined in metaclass"""
+        self.assertEqual(len(SimpleMetaclassClass), 5)
+
+    @make_dynamo_test
+    def test_metaclass_len_direct_call(self):
+        """Test direct call to __len__() on a class with metaclass-defined __len__"""
+        self.assertEqual(SimpleMetaclassClass.__len__(), 5)
+
+
+class CustomMutableMapping(collections.abc.MutableMapping):
+    """Custom mutable mapping implementation with __len__."""
+
+    def __init__(self, data=None):
+        self._data = data if data is not None else {}
+
+    def __len__(self):
+        return len(self._data)
+
+    def __getitem__(self, key):
+        return self._data[key]
+
+    def __setitem__(self, key, value):
+        self._data[key] = value
+
+    def __delitem__(self, key):
+        del self._data[key]
+
+    def __iter__(self):
+        return iter(self._data)
+
+
+@dataclasses.dataclass(frozen=True)
+class FrozenPoint:
+    """Frozen dataclass with __len__ method."""
+
+    x: float
+    y: float
+    z: float
+
+    def __len__(self):
+        return 3
+
+
+@dataclasses.dataclass(frozen=True)
+class FrozenData:
+    """Frozen dataclass with __len__ based on items."""
+
+    items: tuple
+
+    def __len__(self):
+        return len(self.items)
+
+
+class TestMutableMappingLen(torch._dynamo.test_case.TestCase):
+    """Tests for len() on mutable mapping types."""
+
+    def setUp(self):
+        self.old = torch._dynamo.config.enable_trace_unittest
+        torch._dynamo.config.enable_trace_unittest = True
+        super().setUp()
+
+    def tearDown(self):
+        torch._dynamo.config.enable_trace_unittest = self.old
+        return super().tearDown()
+
+    @make_dynamo_test
+    def test_len_custom_mutable_mapping(self):
+        """Test len() on custom mutable mapping."""
+        m = CustomMutableMapping({"x": 10, "y": 20, "z": 30})
+        self.assertEqual(len(m), 3)
+        self.assertEqual(m.__len__(), 3)
+
+    @make_dynamo_test
+    def test_len_custom_mutable_mapping_empty(self):
+        """Test len() on empty custom mutable mapping."""
+        m = CustomMutableMapping()
+        self.assertEqual(len(m), 0)
+
+
+class TestFrozenDataclassLen(torch._dynamo.test_case.TestCase):
+    """Tests for len() on frozen dataclasses."""
+
+    def setUp(self):
+        self.old = torch._dynamo.config.enable_trace_unittest
+        torch._dynamo.config.enable_trace_unittest = True
+        super().setUp()
+
+    def tearDown(self):
+        torch._dynamo.config.enable_trace_unittest = self.old
+        return super().tearDown()
+
+    @make_dynamo_test
+    def test_len_frozen_dataclass_with_len(self):
+        """Test len() on frozen dataclass with custom __len__."""
+        p = FrozenPoint(1.0, 2.0, 3.0)
+        self.assertEqual(len(p), 3)
+        self.assertEqual(p.__len__(), 3)
+
+    @make_dynamo_test
+    def test_len_frozen_dataclass_via_tuple(self):
+        """Test len() on frozen dataclass with __len__ based on contained data."""
+        obj = FrozenData((1, 2, 3, 4))
+        self.assertEqual(len(obj), 4)
+        self.assertEqual(obj.__len__(), 4)
+
+    @make_dynamo_test
+    def test_len_frozen_dataclass_consistency(self):
+        """Test that len() on frozen dataclass is consistent across multiple calls."""
+        obj = FrozenData(("a", "b", "c"))
+        # Call len() twice to ensure consistency
+        len1 = len(obj)
+        len2 = len(obj)
+        self.assertEqual(len1, 3)
+        self.assertEqual(len2, 3)
+
+
+if __name__ == "__main__":
+    from torch._dynamo.test_case import run_tests
+
+    run_tests()
diff --git a/test/dynamo/test_list.py b/test/dynamo/test_list.py
index 85415244db69c..379bcb037b75f 100644
--- a/test/dynamo/test_list.py
+++ b/test/dynamo/test_list.py
@@ -261,7 +261,8 @@ def test_pop(self):
         p = self.thetype("abcd")
         self.assertEqual(p.pop(), "d")
         self.assertEqual(p.pop(1), "b")
-        self.assertRaises(IndexError, p.pop, 10)
+        # The length of p is now 2, valid indices are 0, 1
+        self.assertRaises(IndexError, p.pop, 2)
 
         # Wrong number of arguments
         self.assertRaises(TypeError, p.pop, 2, 3)
diff --git a/test/dynamo/test_logging.py b/test/dynamo/test_logging.py
index 8029e1b3d1610..a06148d13ee08 100644
--- a/test/dynamo/test_logging.py
+++ b/test/dynamo/test_logging.py
@@ -817,6 +817,48 @@ def fn(x, lst):
 
         self.assertEqual(len(records), 0)
 
+    @make_logging_test(side_effects=True)
+    def test_side_effects_logs_fullgraph_graph_break(self, records):
+        """Test that side effects are logged even when fullgraph=True causes an error."""
+        my_list = [1, 2, 3]
+
+        @torch.compile(backend="eager", fullgraph=True)
+        def fn(x, lst):
+            lst.append(4)
+            # Force a graph break after the side effect
+            torch._dynamo.graph_break()
+            return x + len(lst)
+
+        with self.assertRaises(torch._dynamo.exc.Unsupported):
+            fn(torch.ones(1), my_list)
+
+        # Side effects should still be logged even though codegen never ran
+        self.assertGreater(len(records), 0)
+        self.assertIn("Mutating object of type list", records[0].getMessage())
+
+    @make_logging_test(side_effects=True)
+    def test_side_effects_logged_on_fullgraph_side_effect_error(self, records):
+        tracked_list = []
+        hop_list = []
+
+        def fn(x):
+            hop_list.append(1)
+            return x.sin()
+
+        @torch.compile(backend="eager", fullgraph=True)
+        def model(x, lst):
+            lst.append(1)
+            return torch.utils.checkpoint.checkpoint(fn, x, use_reentrant=False)
+
+        with self.assertRaisesRegex(
+            torch._dynamo.exc.Unsupported,
+            "HOP: Unsafe side effect",
+        ):
+            model(torch.ones(1), tracked_list)
+
+        self.assertGreaterEqual(len(records), 1)
+        self.assertIn("Mutating object of type list", records[0].getMessage())
+
     @make_settings_test("torch._dynamo.utils")
     def test_dump_compile_times(self, records):
         fn_opt = torch.compile(example_fn, backend="inductor")
@@ -1390,6 +1432,8 @@ def test_logs_out(self):
             env["TORCH_LOGS_OUT"] = file_path
             _, stderr = self.run_process_no_exception(
                 """\
+import warnings
+warnings.filterwarnings("ignore", category=FutureWarning, module="torch.utils._config_module")
 import torch
 @torch.compile(backend="eager")
 def fn(a):
@@ -1419,7 +1463,7 @@ def test_log_traced_frames(self, records):
         torch._dynamo.eval_frame.clear_dynamo_tls()
 
         # Test program
-        @torch.compile()
+        @torch.compile(backend="eager")
         def foo():
             x = torch.ones([10])
 
@@ -1429,10 +1473,20 @@ def bar():
                 z = y * x
                 return z
 
-            return bar(), bar
+            # force top-level trace of bar
+            try:
+                return bar(), bar
+            finally:
+                pass
 
         foo()
 
+        @torch.compile
+        def baz(x):
+            return x + 1
+
+        baz(torch.ones(3))
+
         # `_log_traced_frames` is registered as an atexit callback, so we invoke
         # it explicitly for testing.
         torch._dynamo.eval_frame._log_traced_frames()
@@ -1447,6 +1501,7 @@ def bar():
 TorchDynamo attempted to trace the following frames: [
   * foo test_logging.py:N
   * bar test_logging.py:N
+  * baz test_logging.py:N
 ]""",
         )
 
diff --git a/test/dynamo/test_misc.py b/test/dynamo/test_misc.py
index 98f59728e1769..893fe2417713e 100644
--- a/test/dynamo/test_misc.py
+++ b/test/dynamo/test_misc.py
@@ -44,6 +44,7 @@
 from torch import Tensor
 from torch._C import FileCheck
 from torch._dynamo import allow_in_graph
+from torch._dynamo.comptime import comptime
 from torch._dynamo.eval_frame import _debug_get_cache_entry_list
 from torch._dynamo.exc import Unsupported
 from torch._dynamo.source import ConstantSource, GetItemSource, LocalSource
@@ -80,7 +81,6 @@
     PLATFORM_SUPPORTS_FLASH_ATTENTION,
     SM80OrLater,
     TEST_CUDA,
-    TEST_MULTIGPU,
 )
 from torch.testing._internal.common_device_type import instantiate_device_type_tests
 from torch.testing._internal.common_methods_invocations import (
@@ -177,11 +177,11 @@ class UserDefineSetAttr:
     setup = False
 
     def __setattr__(self, key, value):
-        assert torch.compiler.is_dynamo_compiling() or UserDefineSetAttr.setup
+        assert torch.compiler.is_dynamo_compiling() or UserDefineSetAttr.setup  # noqa: S101
         super().__setattr__(f"pfx_{key}", value)
 
     def __getattr__(self, key, c=1):
-        assert torch.compiler.is_dynamo_compiling() or UserDefineSetAttr.setup
+        assert torch.compiler.is_dynamo_compiling() or UserDefineSetAttr.setup  # noqa: S101
         # c is added to force a guard on __defaults__ and checks the source for __getattr__
         if c:
             return self.__dict__[f"pfx_{key}"]
@@ -194,7 +194,7 @@ def test_get_cache_entry(self):
         def f(x):
             return x + 1
 
-        torch.compile(f)(torch.randn(5, 5, 5))
+        torch.compile(f, backend="eager")(torch.randn(5, 5, 5))
         entries = _debug_get_cache_entry_list(f)
         self.assertTrue(len(entries) > 0)
 
@@ -215,7 +215,7 @@ def h(x):
             torch._dynamo.graph_break()
             return x + 1
 
-        torch.compile(h)(torch.randn(3, 3))
+        torch.compile(h, backend="eager")(torch.randn(3, 3))
 
         entries = _debug_get_cache_entry_list(torch._dynamo.graph_break)
         self.assertEqual(len(entries), 0)
@@ -278,25 +278,25 @@ def func(a, b):
         torch._dynamo.decorators.mark_unbacked(b, 1)
 
         source_code = run_and_get_code(func, a, b)[1]
-
+        # Check that int64 indexing is used (either 1D [:] or 2D [:, None] form)
         self.assertTrue(
-            "xindex = xoffset + tl.arange(0, XBLOCK)[:].to(tl.int64)\\n"
-            in str(source_code)
+            "tl.arange(0, XBLOCK)[:].to(tl.int64)" in str(source_code)
+            or "tl.arange(0, XBLOCK)[:, None].to(tl.int64)" in str(source_code)
         )
+        # Check that 32-bit indexing is NOT used
         self.assertFalse(
-            "xindex = xoffset + tl.arange(0, XBLOCK)[:]\\n" in str(source_code)
+            "tl.arange(0, XBLOCK)[:]\n" in str(source_code)
+            and ".to(tl.int64)" not in str(source_code)
         )
 
         torch._dynamo.reset()
 
         with torch._inductor.config.patch(assume_32bit_indexing=True):
             source_code = run_and_get_code(func, a, b)[1]
+            # Check that int64 indexing is NOT used when assume_32bit_indexing=True
             self.assertFalse(
-                "xindex = xoffset + tl.arange(0, XBLOCK)[:].to(tl.int64)\\n"
-                in str(source_code)
-            )
-            self.assertTrue(
-                "xindex = xoffset + tl.arange(0, XBLOCK)[:]\\n" in str(source_code)
+                "tl.arange(0, XBLOCK)[:].to(tl.int64)" in str(source_code)
+                or "tl.arange(0, XBLOCK)[:, None].to(tl.int64)" in str(source_code)
             )
 
     def test_dynamo_side_effect(self):
@@ -322,7 +322,7 @@ def forward(self, x):
         with torch._dynamo.config.patch(
             replay_side_effects=False, side_effect_replay_policy="warn"
         ):
-            val = torch.compile(mod)(inp)
+            val = torch.compile(mod, backend="eager")(inp)
             # Verify new object is properly initialized
             self.assertIn("6", val[1]._tensors)
             self.assertEqual(val[1]._tensors["6"], inp + 2)
@@ -407,11 +407,14 @@ def __torch_dispatch__(self, func, types, args=(), kwargs=None):
 
         x = torch.ones(5)
         with YoloMode():
-            out = torch.compile(torch.add, backend=backend, fullgraph=True)(x, x)
-
+            out = torch.compile(torch.add, backend=backend, fullgraph=False)(x, x)
         self.assertEqual(out.sum().item(), 5.0)
         self.assertEqual(len(backend.graphs), 0)
 
+        with YoloMode():
+            with self.assertWarnsRegex(UserWarning, "found no compiled frames"):
+                torch.compile(torch.add, backend=backend, fullgraph=True)(x, x)
+
     def test_compile_non_infra_empty_with_disalloed_dispatch_mode(self):
         from torch.utils._python_dispatch import TorchDispatchMode
 
@@ -570,7 +573,7 @@ def _should_skip_dynamo(cls):
                 return False
 
             def __torch_dispatch__(self, func, types, args=(), kwargs=None):
-                out = torch.compile(func, backend=backend, fullgraph=True)(
+                out = torch.compile(func, backend=backend, fullgraph=False)(
                     *args, **kwargs
                 )
                 return out
@@ -806,7 +809,8 @@ def test_pt2_compliant_ops_are_allowed(self):
             torch.library.impl(
                 "mylib::bar", "CompositeImplicitAutograd", torch.sin, lib=lib
             )
-            assert torch.Tag.pt2_compliant_tag in torch.ops.mylib.bar.default.tags
+            if torch.Tag.pt2_compliant_tag not in torch.ops.mylib.bar.default.tags:
+                raise AssertionError("Expected pt2_compliant_tag in bar.default.tags")
 
             def f(x):
                 return torch.ops.mylib.bar(x)
@@ -832,7 +836,10 @@ def test_non_pt2_compliant_ops_graph_break(self):
             torch.library.impl(
                 "mylib::bar2", "CompositeImplicitAutograd", torch.sin, lib=lib
             )
-            assert torch.Tag.pt2_compliant_tag not in torch.ops.mylib.bar2.default.tags
+            if torch.Tag.pt2_compliant_tag in torch.ops.mylib.bar2.default.tags:
+                raise AssertionError(
+                    "Expected pt2_compliant_tag not in bar2.default.tags"
+                )
 
             def f(x):
                 return torch.ops.mylib.bar2(x)
@@ -1000,6 +1007,127 @@ def f(x):
         self.assertEqual(res.grad, ref.grad)
         self.assertEqual(res.foo, ref.foo)
 
+    def test_input_tensor_custom_attr_mutation(self):
+        def f(x, flag):
+            x.offloading_activation = flag
+            return x + 1
+
+        opt_f = torch.compile(f, backend="eager", fullgraph=True)
+        x = torch.ones(5)
+
+        res = opt_f(x, True)
+        self.assertEqual(res, torch.ones(5) + 1)
+        self.assertTrue(x.offloading_activation)
+
+    def test_intermediate_tensor_custom_attr_mutation(self):
+        def f(x, flag):
+            y = x + 1
+            y.offloading_activation = flag
+            return y
+
+        opt_f = torch.compile(f, backend="eager", fullgraph=True)
+        x = torch.ones(5)
+
+        res = opt_f(x, True)
+        self.assertEqual(res, torch.ones(5) + 1)
+        self.assertTrue(res.offloading_activation)
+
+    @unittest.skipIf(
+        not torch.cuda.is_available() or torch.cuda.get_device_capability() < (9, 0),
+        "requires Hopper+ (SM >= 9.0) for TMA",
+    )
+    @unittest.skipIf(
+        not torch.utils._triton.has_triton()
+        or not hasattr(__import__("triton"), "set_allocator"),
+        "requires triton with set_allocator support",
+    )
+    def test_triton_set_allocator_no_graph_break(self):
+        """set_allocator inside torch.compile does not graph break and
+        replays correctly at runtime (including cache hits)."""
+        import triton
+        import triton.language as tl
+        from triton.runtime._allocation import NullAllocator
+
+        @triton.jit
+        def tma_copy_kernel(
+            x_ptr,
+            out_ptr,
+            M,
+            N,
+            stride_m,
+            stride_n,
+            BLOCK_M: tl.constexpr,
+            BLOCK_N: tl.constexpr,
+        ):
+            pid = tl.program_id(0)
+            desc = tl.make_tensor_descriptor(
+                x_ptr,
+                shape=[M, N],
+                strides=[stride_m, stride_n],
+                block_shape=[BLOCK_M, BLOCK_N],
+            )
+            block = tl.load_tensor_descriptor(desc, [pid * BLOCK_M, 0])
+            out_desc = tl.make_tensor_descriptor(
+                out_ptr,
+                shape=[M, N],
+                strides=[stride_m, stride_n],
+                block_shape=[BLOCK_M, BLOCK_N],
+            )
+            tl.store_tensor_descriptor(out_desc, [pid * BLOCK_M, 0], block)
+
+        M, N, BLOCK_M, BLOCK_N = 128, 64, 64, 64
+
+        def run_kernel(x):
+            out = torch.empty_like(x)
+            tma_copy_kernel[(M // BLOCK_M,)](
+                x,
+                out,
+                M,
+                N,
+                x.stride(0),
+                x.stride(1),
+                BLOCK_M=BLOCK_M,
+                BLOCK_N=BLOCK_N,
+            )
+            return out
+
+        x = torch.randn(M, N, device="cuda")
+
+        from contextlib import contextmanager
+
+        from triton.runtime._allocation import _allocator
+
+        @contextmanager
+        def triton_allocator(allocator):
+            prev = _allocator.get()
+            triton.set_allocator(allocator)
+            try:
+                yield
+            finally:
+                triton.set_allocator(prev)
+
+        def fn_with_set_allocator(x):
+            triton.set_allocator(
+                lambda size, alignment, stream: torch.empty(
+                    size, device="cuda", dtype=torch.int8
+                )
+            )
+            return run_kernel(x)
+
+        opt_fn = torch.compile(
+            fn_with_set_allocator, backend="aot_eager", fullgraph=True
+        )
+
+        # set_allocator inside compiled region does NOT graph break
+        with triton_allocator(NullAllocator()):
+            out = opt_fn(x)
+            self.assertEqual(out, x)
+
+            # Verify set_allocator replays on cache hit (not just tracing)
+            triton.set_allocator(NullAllocator())
+            out2 = opt_fn(x)
+            self.assertEqual(out2, x)
+
     def test_closure_recompiles(self):
         cnt = CompileCounter()
 
@@ -1435,10 +1563,10 @@ def fn(x):
     def test_sys_modules(self):
         def fn(x, y):
             mod_a = sys.modules.get("aaaaaaaa")
-            assert mod_a is None
-            assert "bbbbbbbb" not in sys.modules
+            assert mod_a is None  # noqa: S101
+            assert "bbbbbbbb" not in sys.modules  # noqa: S101
 
-            assert "operator" in sys.modules
+            assert "operator" in sys.modules  # noqa: S101
             operator = sys.modules["operator"]
             builtins = sys.modules.get("builtins")
             operator2 = sys.modules.get("cccccccc", operator)
@@ -1500,7 +1628,7 @@ def fn(x):
         self.assertEqual(r1, r2)
 
     def test_tensor_hasattr(self):
-        @torch.compile(fullgraph=True)
+        @torch.compile(fullgraph=True, backend="eager")
         def fn(x):
             if hasattr(x, "test"):
                 return x + 2
@@ -1514,7 +1642,7 @@ def fn(x):
         self.assertEqual(torch.ones(2, 2) + 2, fn(inp))
 
     def test_tensor_call_obj_hasattr_view(self):
-        @torch.compile(fullgraph=True)
+        @torch.compile(fullgraph=True, backend="eager")
         def fn(x):
             output3 = getattr(x, "view", None)(10)
             return output3
@@ -1523,7 +1651,7 @@ def fn(x):
         self.assertEqual(x.view(10), fn(x))
 
     def test_mro_type_tensor_no_source(self):
-        @torch.compile(fullgraph=True)
+        @torch.compile(fullgraph=True, backend="eager")
         def fn(x):
             z = []
             input_type = type(torch.ones(2, 2))
@@ -1542,13 +1670,66 @@ def add_one(x):
         t = torch.nn.Parameter(torch.ones(1))
         t.add_one = add_one
 
-        @torch.compile(fullgraph=True)
+        @torch.compile(fullgraph=True, backend="eager")
         def fn(x):
             return t.add_one(t) + x
 
         result = fn(torch.ones(1))
         self.assertEqual(torch.ones(1) + 2, result)
 
+    def test_known_tensor_methods_traced(self):
+        # Verify that known tensor methods (in all_tensor_attrs) are still
+        # traced into the graph via the generic proxy path.
+        cnt = torch._dynamo.testing.CompileCounter()
+
+        @torch.compile(backend=cnt, fullgraph=True)
+        def fn(x):
+            return x.abs().cos()
+
+        result = fn(torch.randn(4))
+        self.assertEqual(cnt.frame_count, 1)
+        self.assertEqual(cnt.op_count, 2)
+
+    def test_tensor_subclass_method_traced(self):
+        # Methods defined on the actual tensor class (including dynamically
+        # added ones) should be proxied through the generic call_method path,
+        # not graph-broken.  This validates that the guard uses the concrete
+        # class_type rather than the static all_tensor_attrs dict.
+        def _dynamo_test_method(self):
+            return self + 1
+
+        with unittest.mock.patch.object(
+            torch.Tensor, "_dynamo_test_method", _dynamo_test_method, create=True
+        ):
+            cnt = CompileCounterWithBackend("eager")
+
+            @torch.compile(backend=cnt)
+            def fn(x):
+                y = x._dynamo_test_method()
+                return y + 1
+
+            result = fn(torch.randn(4))
+            self.assertEqual(cnt.frame_count, 1)
+            # Verify _dynamo_test_method appears as a call_method in the FX graph
+            call_method_targets = [
+                n.target for n in cnt.graphs[0].graph.nodes if n.op == "call_method"
+            ]
+            self.assertIn("_dynamo_test_method", call_method_targets)
+
+    def test_unknown_tensor_method_graph_break(self):
+        # Truly unknown methods raise AttributeError during tracing at
+        # LOAD_ATTR time (dynamic_getattr), ensuring dynamo does not
+        # silently proxy them into the compiled graph.
+        cnt = torch._dynamo.testing.CompileCounter()
+
+        @torch.compile(backend=cnt)
+        def fn(x):
+            y = x._nonexistent_test_method_xyz()
+            return y + 1
+
+        with self.assertRaises(AttributeError):
+            fn(torch.randn(4))
+
     def test_shape_unpack(self):
         def fn(x):
             a, b = x.size()
@@ -1615,7 +1796,7 @@ def forward(self, x):
         model = Model()
         x = torch.randn(4, 10)
         expected = model(x)
-        compiled_model = torch.compile(model, fullgraph=False)
+        compiled_model = torch.compile(model, fullgraph=False, backend="eager")
         actual = compiled_model(x)
         self.assertEqual(expected, actual)
 
@@ -1684,7 +1865,7 @@ def f2(x, index):
 
     @torch._dynamo.config.patch(capture_scalar_outputs=True)
     def test_arange_length_with_float32_dtype(self):
-        @torch.compile(fullgraph=True)
+        @torch.compile(fullgraph=True, backend="eager")
         def f(x):
             y = x.item()
             r = torch.arange(y, dtype=torch.float32)
@@ -1742,7 +1923,7 @@ def f(x):
             torch._check(y >= 0)
             # Cannot conditional on unbacked SymInt
             if y == 0:
-                assert False
+                assert False  # noqa: B011, S101
             else:
                 return torch.arange(0, y)
 
@@ -1913,16 +2094,16 @@ def f(x):
             f(x)
 
     def test_assert(self):
-        @torch.compile
+        @torch.compile(backend="eager")
         def fn1(x):
-            assert x.shape != x.shape
+            assert x.shape != x.shape  # noqa: S101
 
         with self.assertRaises(AssertionError):
             a = torch.randn(10)
             fn1(a)
 
         def fn2(x):
-            assert x.shape == x.shape
+            assert x.shape == x.shape  # noqa: S101
             return x.abs()
 
         torch._dynamo.testing.standard_test(self, fn=fn2, nargs=1, expected_ops=1)
@@ -2386,6 +2567,21 @@ def f():
         self.assertTrue(hasattr(result, "c"))
         self.assertEqual(result.c, torch.tensor(3.0))
 
+    def test_namedtuple___eq__(self):
+        class MyNamedTuple(typing.NamedTuple):
+            a: int
+            b: int
+
+        @torch.compile(backend="eager", fullgraph=True)
+        def f(x):
+            t1 = MyNamedTuple(a=1, b=2)
+            t2 = (1, 2)
+            return x.sin(), (t1 == t2)
+
+        x = torch.randn(2)
+        res = f(x)
+        self.assertTrue(res[1])
+
     def test_structseq1(self):
         def fn(x, y):
             return torch.return_types.max((x, y))
@@ -2393,7 +2589,7 @@ def fn(x, y):
         x = torch.randn(3, 2)
         y = torch.randn(2, 4)
         expected = fn(x, y)
-        fn_opt = torch.compile(fullgraph=True)(fn)
+        fn_opt = torch.compile(fullgraph=True, backend="eager")(fn)
         actual = fn_opt(x, y)
 
         self.assertEqual(actual, expected)
@@ -2405,7 +2601,7 @@ def fn(x, y):
         x = torch.randn(3, 2)
         y = torch.randn(2, 4)
         expected = fn(x, y)
-        fn_opt = torch.compile(fullgraph=True)(fn)
+        fn_opt = torch.compile(fullgraph=True, backend="eager")(fn)
         actual = fn_opt(x, y)
 
         self.assertEqual(actual, expected)
@@ -2457,7 +2653,7 @@ def fn(a):
         )
 
     def test_range_iter_guards(self):
-        @torch.compile()
+        @torch.compile(backend="eager")
         def func():
             @torch._dynamo.disable(recursive=False)
             def run(n):
@@ -2726,6 +2922,216 @@ def fn(x, foo):
         self.assertEqual(ref, res)
         self.assertEqual(foo1.field, foo2.field)
 
+    def test_dict_with_descriptor(self):
+        class MyDescriptor:
+            def __get__(self, obj, objtype=None):
+                if obj is None:
+                    return self
+                return obj.__dict__.get("_value", 0)
+
+            def __set__(self, obj, value):
+                obj.__dict__["_value"] = value
+
+        class MyClass:
+            prop = MyDescriptor()
+
+            def __init__(self):
+                self.prop = 42
+
+        def fn(obj):
+            obj.__dict__["extra"] = 99
+            return obj.prop + obj.__dict__["extra"]
+
+        opt_fn = torch.compile(fn, fullgraph=True, backend="eager")
+        obj = MyClass()
+        res = fn(obj)
+
+        obj2 = MyClass()
+        out = opt_fn(obj2)
+        self.assertEqual(res, out)
+        self.assertEqual(obj2.__dict__["_value"], 42)
+        self.assertEqual(obj2.__dict__["extra"], 99)
+
+    def test_dict_with_slots(self):
+        class SlottedClass:
+            __slots__ = ("x",)
+
+            def __init__(self, x):
+                self.x = x
+
+        def fn(obj):
+            # SlottedClass doesn't have __dict__, so this should fail or be handled
+            return obj.x * 2
+
+        opt_fn = torch.compile(fn, fullgraph=True, backend="eager")
+        obj = SlottedClass(5)
+        res = fn(obj)
+        out = opt_fn(obj)
+        self.assertEqual(res, out)
+
+    def test_dict_with_slots_and_dict(self):
+        class SlottedWithDict:
+            __slots__ = ("x", "__dict__")
+
+            def __init__(self, x):
+                self.x = x
+
+        def fn(obj):
+            obj.__dict__["custom"] = 100
+            return obj.x + obj.__dict__["custom"]
+
+        opt_fn = torch.compile(fn, fullgraph=True, backend="eager")
+        obj = SlottedWithDict(7)
+        res = fn(obj)
+
+        obj2 = SlottedWithDict(7)
+        out = opt_fn(obj2)
+        self.assertEqual(res, out)
+        self.assertEqual(obj2.__dict__["custom"], 100)
+
+    def test_dict_descriptor_interaction(self):
+        class DescriptorWithDict:
+            def __get__(self, obj, objtype=None):
+                if obj is None:
+                    return self
+                return obj.__dict__.get("_internal", "default")
+
+            def __set__(self, obj, value):
+                obj.__dict__["_internal"] = f"desc_{value}"
+
+        class MyClass:
+            desc = DescriptorWithDict()
+
+            def __init__(self):
+                pass
+
+        def fn(obj):
+            obj.desc = "hello"
+            obj.__dict__["_internal"] = "direct"
+            return obj.desc
+
+        opt_fn = torch.compile(fn, fullgraph=True, backend="eager")
+        obj = MyClass()
+        ref = fn(obj)
+
+        obj2 = MyClass()
+        out = opt_fn(obj2)
+        self.assertEqual(ref, out)
+        self.assertEqual(out, "direct")
+
+    def test_mutable_mapping_dict_update(self):
+        """Test that MutableMappingVariable handles __dict__ updates correctly."""
+        from collections import OrderedDict
+
+        class CustomMapping(OrderedDict):
+            pass
+
+        @torch.compile(fullgraph=True, backend="eager")
+        def fn(mapping):
+            mapping.__dict__["custom_attr"] = 42
+            return mapping.__dict__["custom_attr"]
+
+        mapping = CustomMapping()
+        out = fn(mapping)
+        self.assertEqual(out, 42)
+        self.assertIn("custom_attr", mapping.__dict__)
+        self.assertEqual(mapping.__dict__["custom_attr"], 42)
+
+    def test_mutable_mapping_dict_access_pattern(self):
+        """Test accessing attributes through __dict__ on MutableMapping subclasses."""
+        from collections import OrderedDict
+
+        class TrackedDict(OrderedDict):
+            pass
+
+        @torch.compile(fullgraph=True, backend="eager")
+        def fn(d):
+            # Pattern: check if attribute exists in __dict__
+            if "tracker" not in d.__dict__:
+                d.__dict__["tracker"] = []
+            d.__dict__["tracker"].append(1)
+            return len(d.__dict__["tracker"])
+
+        d = TrackedDict()
+        out = fn(d)
+        self.assertEqual(out, 1)
+        self.assertIn("tracker", d.__dict__)
+        self.assertEqual(d.__dict__["tracker"], [1])
+
+    def test_mutable_mapping_lazy_dict_initialization(self):
+        """Test lazy initialization pattern with MutableMapping __dict__."""
+        from collections import defaultdict
+
+        class LazyMapping(dict):
+            pass
+
+        @torch.compile(fullgraph=True, backend="eager")
+        def fn(mapping):
+            # Lazy initialization in __dict__
+            if not hasattr(mapping, "_cache"):
+                mapping.__dict__["_cache"] = {}
+            mapping.__dict__["_cache"]["key"] = "value"
+            return mapping.__dict__["_cache"]["key"]
+
+        m = LazyMapping()
+        out = fn(m)
+        self.assertEqual(out, "value")
+        self.assertIn("_cache", m.__dict__)
+        self.assertEqual(m.__dict__["_cache"]["key"], "value")
+
+    def test_mutable_mapping_dict_with_property_setter(self):
+        """Test MutableMapping with property setters that access __dict__."""
+        from collections import OrderedDict
+
+        class PropertyMapping(OrderedDict):
+            def __init__(self):
+                super().__init__()
+                self.__dict__["_value"] = 0
+
+            @property
+            def value(self):
+                return self.__dict__.get("_value", 0)
+
+            @value.setter
+            def value(self, v):
+                self.__dict__["_value"] = v
+
+        @torch.compile(fullgraph=True, backend="eager")
+        def fn(m):
+            m.value = 100
+            return m.value
+
+        m = PropertyMapping()
+        out = fn(m)
+        self.assertEqual(out, 100)
+        self.assertIn("_value", m.__dict__)
+        self.assertEqual(m.__dict__["_value"], 100)
+
+    def test_mutable_mapping_dict_multiple_accesses(self):
+        """Test multiple accesses and mutations to MutableMapping __dict__."""
+        from collections import OrderedDict
+
+        class MultiAccessMapping(OrderedDict):
+            pass
+
+        @torch.compile(fullgraph=True, backend="eager")
+        def fn(mapping):
+            # Multiple accesses to __dict__
+            mapping.__dict__["a"] = 1
+            mapping.__dict__["b"] = mapping.__dict__["a"] + 1
+            mapping.__dict__["c"] = mapping.__dict__["b"] + 1
+            return mapping.__dict__["a"] + mapping.__dict__["b"] + mapping.__dict__["c"]
+
+        m = MultiAccessMapping()
+        out = fn(m)
+        self.assertEqual(out, 6)  # 1 + 2 + 3
+        self.assertIn("a", m.__dict__)
+        self.assertIn("b", m.__dict__)
+        self.assertIn("c", m.__dict__)
+        self.assertEqual(m.__dict__["a"], 1)
+        self.assertEqual(m.__dict__["b"], 2)
+        self.assertEqual(m.__dict__["c"], 3)
+
     def test_get_attr_function(self):
         def fn(g, x):
             return g(x)
@@ -2809,6 +3215,55 @@ def fn(mod, x):
         opt_fn = torch.compile(fn, backend=cnts)
         self.assertTrue(same(opt_fn(mod, x), fn(mod, x)))
 
+    def test_nn_module_getattribute_simple_delegation(self):
+        # Test that nn.Module with __getattribute__ that overrides a
+        # single attribute name compiles without graph break.
+        class MyMod(torch.nn.Module):
+            def __init__(self) -> None:
+                super().__init__()
+                self.scale = 3.0
+
+            def __getattribute__(self, name):
+                if name == "my_scale":
+                    return super().__getattribute__("scale")
+                return super().__getattribute__(name)
+
+            def forward(self, x):
+                return x * self.my_scale
+
+        mod = MyMod()
+        x = torch.randn(2, 4)
+        expected = mod(x)
+        cnts = torch._dynamo.testing.CompileCounter()
+        opt_fn = torch.compile(mod, backend=cnts)
+        result = opt_fn(x)
+        self.assertTrue(same(result, expected))
+
+    def test_nn_module_getattribute_graph_break(self):
+        # __getattribute__ that Dynamo cannot trace produces correct results
+        # via eager fallback instead of crashing.
+        class MyMod(torch.nn.Module):
+            def __init__(self) -> None:
+                super().__init__()
+                self.linear = torch.nn.Linear(4, 4)
+
+            def __getattribute__(self, name):
+                if name == "my_attr":
+                    return eval("42")  # eval is untraceable
+                return super().__getattribute__(name)
+
+            def forward(self, x):
+                a = self.my_attr
+                return self.linear(x) + a
+
+        mod = MyMod()
+        x = torch.randn(2, 4)
+        expected = mod(x)
+        cnts = torch._dynamo.testing.CompileCounter()
+        opt_fn = torch.compile(mod, backend=cnts)
+        result = opt_fn(x)
+        self.assertTrue(same(result, expected))
+
     def test_constant_getattr(self):
         # https://github.com/pytorch/pytorch/issues/97480
         def fn():
@@ -2837,7 +3292,7 @@ def fn(cfg, x, y):
         self.assertEqual(cnts.op_count, 2)
 
     def test_data_access_in_inference_mode(self):
-        @torch.compile(fullgraph=True)
+        @torch.compile(fullgraph=True, backend="eager")
         def f(x):
             y = x.data
             return y
@@ -2858,12 +3313,12 @@ class MyDataClass:
 
         def fn(obj):
             class_fields = dataclasses.fields(obj)
-            assert len(class_fields)
-            assert all(field.default is None for field in class_fields[1:])
+            assert len(class_fields)  # noqa: S101
+            assert all(field.default is None for field in class_fields[1:])  # noqa: S101
             other_fields_are_none = all(
                 getattr(obj, field.name) is None for field in class_fields[1:]
             )
-            assert not other_fields_are_none
+            assert not other_fields_are_none  # noqa: S101
 
             if not hasattr(obj, "a"):
                 return -1
@@ -3180,7 +3635,7 @@ def fn(x, y):
         self.assertEqual(cnts.frame_count, 2)
 
     def test_numpy_readonly(self):
-        @torch.compile(fullgraph=True)
+        @torch.compile(fullgraph=True, backend="eager")
         def fn(x):
             return x
 
@@ -3289,7 +3744,7 @@ def test_numpy_as_global(self):
         global x
         x = np.arange(10)
 
-        @torch.compile(fullgraph=True)
+        @torch.compile(fullgraph=True, backend="eager")
         def fn(y):
             return y + x + x
 
@@ -3301,7 +3756,7 @@ def fn(y):
     def test_numpy_gt(self):
         x = np.arange(10)
 
-        @torch.compile
+        @torch.compile(backend="eager")
         def fn(y):
             return y >= 3
 
@@ -3312,7 +3767,7 @@ def fn(y):
     def test_numpy_min(self):
         x = np.arange(10)
 
-        @torch.compile
+        @torch.compile(backend="eager")
         def fn(y):
             return min(y, 3), min(y, y - 1)
 
@@ -3418,7 +3873,8 @@ def fn2():
 
         r2 = fn2()
         self.assertEqual(cnts.frame_count, 0)
-        assert r2 in ("L", "U")
+        if r2 not in ("L", "U"):
+            raise AssertionError(f"Expected r2 to be 'L' or 'U', got {r2}")
 
     def test_trace_ndarray_frame(self):
         def fn(x):
@@ -3516,7 +3972,7 @@ def fn(dt):
     # NB this may involve a graph break
     @torch._dynamo.config.patch(use_numpy_random_stream=True)
     def test_numpy_random_config_to_numpy(self):
-        @torch.compile
+        @torch.compile(backend="eager")
         def fn():
             return np.random.uniform(size=13)
 
@@ -3631,7 +4087,7 @@ def items(self):
                 for k, v in self._items.items():
                     yield (k, v)
 
-        @torch.compile(fullgraph=True)
+        @torch.compile(fullgraph=True, backend="eager")
         def to_weird_dict(td):
             return MyWeirdDict(**td)
 
@@ -3933,6 +4389,7 @@ def double_nested_call(x, y):
         o = torch.compile(foo, fullgraph=True, backend="eager")(x, y)
         self.assertEqual(o, x * y)
 
+    @torch._dynamo.config.patch(nested_graph_breaks=False)
     def test_module_deepcopy(self):
         m1 = torch.nn.Sequential(
             torch.nn.Linear(10, 10),
@@ -3963,6 +4420,187 @@ def fn(m, x):
         self.assertEqual(cnts.frame_count, 1)
         self.assertEqual(cnts.op_count, 4)
 
+    def test_id_of_container_as_dict_key(self):
+        MY_DICT = {"a": 1, "b": 2}
+
+        def fn(x):
+            memo = {}
+            memo[id(MY_DICT)] = True
+            if id(MY_DICT) in memo:
+                return x + 1.0
+            return x + 2.0
+
+        x = torch.randn(4)
+        correct = fn(x)
+        result = torch.compile(fn, fullgraph=True)(x)
+        self.assertEqual(result, correct)
+
+    def test_id_of_list_as_dict_key(self):
+        MY_LIST = [1.0, 2.0]
+
+        def fn(x):
+            memo = {}
+            memo[id(MY_LIST)] = True
+            if id(MY_LIST) in memo:
+                return x + 1.0
+            return x + 2.0
+
+        x = torch.randn(4)
+        correct = fn(x)
+        result = torch.compile(fn, fullgraph=True)(x)
+        self.assertEqual(result, correct)
+
+    def test_deepcopy_dict(self):
+        MY_DICT = {"a": 1, "b": 2.0, "c": None}
+
+        def fn(x):
+            d = copy.deepcopy(MY_DICT)
+            d["b"] = 3.0
+            return x + d["b"]
+
+        x = torch.randn(4)
+        correct = fn(x)
+        result = torch.compile(fn, fullgraph=True)(x)
+        self.assertEqual(result, correct)
+
+    def test_deepcopy_nested_dict(self):
+        NESTED = {"a": {"b": 1.0}, "c": [2.0, 3.0]}
+
+        def fn(x):
+            d = copy.deepcopy(NESTED)
+            return x + d["a"]["b"] + d["c"][0]
+
+        x = torch.randn(4)
+        correct = fn(x)
+        result = torch.compile(fn, fullgraph=True)(x)
+        self.assertEqual(result, correct)
+
+    def test_deepcopy_list(self):
+        MY_LIST = [1.0, 2.0, 3.0]
+
+        def fn(x):
+            lst = copy.deepcopy(MY_LIST)
+            lst[0] = 5.0
+            return x + lst[0]
+
+        x = torch.randn(4)
+        correct = fn(x)
+        result = torch.compile(fn, fullgraph=True)(x)
+        self.assertEqual(result, correct)
+
+    def test_deepcopy_user_defined_object(self):
+        class MyConfig:
+            def __init__(self, hidden_size=64):
+                self.hidden_size = hidden_size
+
+        class MyModule(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.config = MyConfig()
+                self.linear = torch.nn.Linear(64, 64)
+
+            def forward(self, x):
+                cfg = copy.deepcopy(self.config)
+                return self.linear(x) * cfg.hidden_size
+
+        m = MyModule()
+        x = torch.randn(2, 64)
+        result = torch.compile(m, fullgraph=True, backend="eager")(x)
+        self.assertEqual(m(x), result)
+
+    def test_deepcopy_user_defined_object_with_containers(self):
+        class Config:
+            def __init__(self):
+                self.sizes = [1, 2, 3]
+                self.mapping = {"a": 10, "b": 20}
+                self.flags = (True, False)
+
+        def fn(x, cfg):
+            c = copy.deepcopy(cfg)
+            c.sizes[0] = 99
+            c.mapping["a"] = 77
+            return x + c.sizes[0] + c.mapping["a"]
+
+        cfg = Config()
+        x = torch.randn(4)
+        correct = fn(x, cfg)
+        result = torch.compile(fn, fullgraph=True, backend="eager")(x, cfg)
+        self.assertEqual(result, correct)
+        # Verify deepcopy didn't mutate original
+        self.assertEqual(cfg.sizes[0], 1)
+        self.assertEqual(cfg.mapping["a"], 10)
+
+    def test_deepcopy_set(self):
+        MY_SET = {1, 2, 3}
+
+        def fn(x):
+            s = copy.deepcopy(MY_SET)
+            return x + len(s)
+
+        x = torch.randn(4)
+        correct = fn(x)
+        result = torch.compile(fn, fullgraph=True)(x)
+        self.assertEqual(result, correct)
+
+    def test_deepcopy_frozenset(self):
+        MY_FROZENSET = frozenset([1, 2, 3])
+
+        def fn(x):
+            s = copy.deepcopy(MY_FROZENSET)
+            return x + len(s)
+
+        x = torch.randn(4)
+        correct = fn(x)
+        result = torch.compile(fn, fullgraph=True)(x)
+        self.assertEqual(result, correct)
+
+    def test_deepcopy_user_defined_object_with_method(self):
+        class MyConfig:
+            def __init__(self, hidden_size=64):
+                self.hidden_size = hidden_size
+
+            def get_size(self):
+                return self.hidden_size
+
+        class MyModule(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.config = MyConfig()
+                self.linear = torch.nn.Linear(64, 64)
+
+            def forward(self, x):
+                cfg = copy.deepcopy(self.config)
+                return self.linear(x) * cfg.get_size()
+
+        m = MyModule()
+        x = torch.randn(2, 64)
+        correct = m(x)
+        result = torch.compile(m, fullgraph=True, backend="eager")(x)
+        self.assertEqual(result, correct)
+
+    def test_deepcopy_nested_user_defined_object(self):
+        class Inner:
+            def __init__(self, scale):
+                self.scale = scale
+
+        class Outer:
+            def __init__(self):
+                self.inner = Inner(2.0)
+                self.bias = 1.0
+
+        def fn(x, cfg):
+            c = copy.deepcopy(cfg)
+            c.inner.scale = 3.0
+            return x * c.inner.scale + c.bias
+
+        cfg = Outer()
+        x = torch.randn(4)
+        correct = fn(x, cfg)
+        result = torch.compile(fn, fullgraph=True, backend="eager")(x, cfg)
+        self.assertEqual(result, correct)
+        # Verify deepcopy didn't mutate original
+        self.assertEqual(cfg.inner.scale, 2.0)
+
     def test_global_state_guard_serialization(self):
         GlobalStateGuard = torch._C._dynamo.guards.GlobalStateGuard
         guards = GlobalStateGuard()
@@ -4224,7 +4862,7 @@ def test_thread_local_setattr(self):
 
         loc = local()
 
-        @torch.compile(fullgraph=True)
+        @torch.compile(fullgraph=True, backend="eager")
         def fn(x, l):
             l.x = x
             return x + 1
@@ -4417,8 +5055,8 @@ def fn1():
             v1 = torch.randn(10)
 
             def fn2(*args, **kwargs):
-                assert len(args) == 1
-                assert len(kwargs) == 1
+                assert len(args) == 1  # noqa: S101
+                assert len(kwargs) == 1  # noqa: S101
                 v2 = torch.randn(10) + args[0] + kwargs["b"]
 
                 def fn3(v3=torch.randn(10)):
@@ -4698,7 +5336,7 @@ def inner():
             # Calling `inner` so Dynamo won't skip this frame.
             return inner(), inner
 
-        @torch.compile
+        @torch.compile(backend="eager")
         def root():
             return get_inner()
 
@@ -4779,7 +5417,7 @@ def get_x():
 
         get_x, set_x = make_x_get_set()
 
-        @torch.compile(fullgraph=True)
+        @torch.compile(fullgraph=True, backend="eager")
         def fn(t):
             set_x(42)  # This sets the `x` created within `make_x_get_set`
             res = t + x  # This uses the `x` outside `make_x_get_set`.
@@ -4830,7 +5468,7 @@ def test_top_package_import(self):
         def fn(x):
             import torch.fx
 
-            assert not isinstance(x, torch.fx.Proxy)
+            assert not isinstance(x, torch.fx.Proxy)  # noqa: S101
             return torch.sin(x)
 
         x = torch.randn(4, 5)
@@ -5125,6 +5763,54 @@ def fn(x):
         # Graph breaks at manual_seed.
         self.assertEqual(len(counters["graph_break"]), 1)
 
+    def test_torch_generator_manual_seed(self):
+        from torch._dynamo.utils import counters
+
+        cnts = torch._dynamo.testing.CompileCounter()
+        counters.clear()
+
+        def fn(x, gen):
+            gen.manual_seed(3)
+            return x + 1
+
+        x = torch.randn(10)
+        ref = fn(x, torch.Generator())
+
+        opt_fn = torch.compile(fn, backend=cnts, fullgraph=False)
+        res = opt_fn(x, torch.Generator())
+
+        self.assertTrue(same(ref, res))
+        self.assertEqual(cnts.op_count, 1)
+        self.assertEqual(cnts.frame_count, 1)
+        self.assertEqual(len(counters["graph_break"]), 1)
+
+    def test_torch_generator_initial_seed(self):
+        from torch._dynamo.utils import counters
+
+        cnts = torch._dynamo.testing.CompileCounter()
+        counters.clear()
+
+        def fn(x):
+            return x + 1, torch.default_generator.initial_seed()
+
+        x = torch.randn(10)
+        ref = fn(x)
+
+        opt_fn = torch.compile(fn, backend=cnts, fullgraph=False)
+        res = opt_fn(x)
+
+        self.assertTrue(same(ref, res))
+        self.assertEqual(cnts.op_count, 1)
+        self.assertEqual(cnts.frame_count, 1)
+        self.assertEqual(len(counters["graph_break"]), 1)
+
+    def test_torch_generator_get_state_fullgraph(self):
+        def fn():
+            return torch.default_generator.get_state()
+
+        with self.assertRaises(torch._dynamo.exc.Unsupported):
+            torch.compile(fn, backend="eager", fullgraph=True)()
+
     def test_is_tensor_like(self):
         cnts = torch._dynamo.testing.CompileCounter()
 
@@ -5477,7 +6163,7 @@ def another_fn():
         def fn():
             return "another_fn" in str(another_fn)
 
-        opt_fn = torch.compile(fn, fullgraph=True)
+        opt_fn = torch.compile(fn, fullgraph=True, backend="eager")
         self.assertTrue(opt_fn())
 
     def test_repeat_interleave_graphbreaks(self):
@@ -6025,7 +6711,7 @@ def forward(self, x):
                     "While compiling, we found certain side effects happened in the model.forward. Here are the list of potential sources you can double check: [\"L['self']\"]"
                 ),
             ):
-                torch.compile(foo, fullgraph=True)(torch.randn(4, 4))
+                torch.compile(foo, fullgraph=True, backend="eager")(torch.randn(4, 4))
 
         with torch._dynamo.config.patch(
             replay_side_effects=False, side_effect_replay_policy="silent"
@@ -6033,7 +6719,9 @@ def forward(self, x):
             foo_v2_compile = Foo()
             foo_v2_eager = Foo()
             inp = torch.randn(4, 4)
-            res = torch.compile(foo_v2_compile, fullgraph=True)(torch.randn(4, 4))
+            res = torch.compile(foo_v2_compile, fullgraph=True, backend="eager")(
+                torch.randn(4, 4)
+            )
             self.assertEqual(foo_v2_compile.tensor, None)
             self.assertEqual(foo_v2_compile.const, 4)
             self.assertEqual(foo_v2_compile.bar.const, 4)
@@ -6057,7 +6745,7 @@ def forward(self, x):
             replay_side_effects=False, side_effect_replay_policy="error"
         ):
             foo = Foo()
-            torch.compile(foo, fullgraph=True)(torch.randn(4, 4))
+            torch.compile(foo, fullgraph=True, backend="eager")(torch.randn(4, 4))
 
     def test_list_append_return_none(self):
         def fn(x):
@@ -6072,6 +6760,29 @@ def fn(x):
         res2 = opt_fn(x)
         self.assertEqual(res, res2)
 
+    def test_function_return_none_creates_constant_variable(self):
+        """
+        Test that functions returning None properly return ConstantVariable.create(None)
+        instead of raw None, which would violate the stack's type contract.
+
+        Regression test for: Avoid using Optional[VariableTracker]
+        """
+
+        def gn(x):
+            return
+
+        torch._dynamo.config.reorderable_logging_functions.add(gn)
+
+        @torch.compile(backend="eager")
+        def fn(x):
+            x = x + 1
+            if gn(x) is None:
+                return x + 2
+            return x + 4
+
+        # If this doesn't crash, the test passes
+        fn(torch.ones(3))
+
     @patch.object(torch._dynamo.config, "capture_scalar_outputs", True)
     def test_tensor_ctor_list_of_tensor(self):
         def fn(x):
@@ -6088,7 +6799,7 @@ def fn(x):
     def test_tensor_types(self):
         def fn(dtype, tensor_type):
             x = torch.empty(4, dtype=dtype)
-            assert isinstance(x, tensor_type)
+            assert isinstance(x, tensor_type)  # noqa: S101
 
         opt_fn = torch.compile(fn, backend="eager")
         opt_fn(torch.float32, torch.FloatTensor)
@@ -6102,6 +6813,25 @@ def fn(dtype, tensor_type):
         opt_fn(torch.int16, torch.ShortTensor)
         opt_fn(torch.bool, torch.BoolTensor)
 
+    @unittest.skipIf(not torch.distributed.is_available(), "requires distributed")
+    def test_or_union_type_opaque_class(self):
+        # Test that or_ on opaque class types (e.g. Shard | _StridedShard)
+        # doesn't cause a graph break.
+        from torch.distributed.tensor import Shard
+        from torch.distributed.tensor.placement_types import _StridedShard
+
+        cnt = torch._dynamo.testing.CompileCounter()
+
+        @torch.compile(backend=cnt, fullgraph=True)
+        def fn(x):
+            _ = Shard | _StridedShard
+            return x + 1
+
+        x = torch.randn(4)
+        result = fn(x)
+        self.assertEqual(result, x + 1)
+        self.assertEqual(cnt.frame_count, 1)
+
     def test_nan(self):
         def f(x, n):
             return x * 2 + n
@@ -6278,7 +7008,7 @@ def forward(x, config):
         config = Config()
         x = torch.randn(2, 2)
 
-        compiled = torch.compile(forward, fullgraph=True)
+        compiled = torch.compile(forward, fullgraph=True, backend="eager")
         compiled(x, config)
 
     def test_nn_functional_reduction(self):
@@ -6405,7 +7135,7 @@ def __init__(self) -> None:
 
             def forward(self, idx, targets=None):
                 b, t = idx.size()
-                assert t <= self.block_size, (
+                assert t <= self.block_size, (  # noqa: S101
                     "Cannot forward, model block size is exhausted."
                 )
 
@@ -6556,6 +7286,49 @@ def test_recompile(foo, *, exp_frame_count):
         torch._dynamo.reset()
         test_recompile(foo_graph_break, exp_frame_count=2)
 
+    def test_multithread_compile_dynamic(self):
+        def f(x):
+            comptime.assert_static(x.shape[0])
+            return x * x
+
+        def _do_test(func):
+            success = True
+
+            def run(offset):
+                for i in range(20):
+                    print(func(torch.randn(i * 2 + offset)))
+
+            t1 = threading.Thread(target=run, args=[0])
+            t2 = threading.Thread(target=run, args=[1])
+
+            def exc_hook(x):
+                nonlocal success
+                success = False
+
+            try:
+                threading.excepthook = exc_hook
+                t1.start()
+                t2.start()
+
+                t1.join()
+                t2.join()
+            finally:
+                threading.excepthook = threading.__excepthook__
+            self.assertTrue(success)
+
+        _do_test(torch.compile(f, backend="eager", dynamic=False))
+        torch._dynamo.reset()
+
+        f_opt = torch.compile(f, backend="eager")
+
+        def g(x):
+            with torch._dynamo.config.patch(
+                automatic_dynamic_shapes=False, assume_static_by_default=True
+            ):
+                f_opt(x)
+
+        _do_test(g)
+
     def test_backend_match_guard_multi_threads(self):
         x = torch.randn([3, 4])
 
@@ -6736,11 +7509,7 @@ def body(x):
 
         mod = Module()
 
-        error_message = ""
-        if torch._dynamo.config.inline_inbuilt_nn_modules:
-            error_message = r"Higher Order Operator: torch\.ops\.higher_order\.map_impl"
-        else:
-            error_message = "Can't inplace modify module params/buffers"
+        error_message = r"Higher Order Operator: torch\.ops\.higher_order\.map_impl"
 
         with self.assertRaisesRegex(
             torch._dynamo.exc.UncapturedHigherOrderOpError, error_message
@@ -6866,7 +7635,10 @@ def test_duplicate_graph_break_log(self):
 
         @torch.compile(backend="eager")
         def f1(a, b):
-            f2(a, b)
+            try:
+                f2(a, b)
+            finally:
+                pass
 
         def f2(a, b):
             c = a + b
@@ -6875,7 +7647,10 @@ def f2(a, b):
 
         @torch.compile(backend="eager")
         def g1(a, b):
-            g2(a, b)
+            try:
+                g2(a, b)
+            finally:
+                pass
 
         def g2(a, b):
             c = a + b
@@ -6931,7 +7706,7 @@ def fn():
             return torch.LongTensor(y)
 
         ref = fn()
-        res = torch.compile(fullgraph=True)(fn)()
+        res = torch.compile(fullgraph=True, backend="eager")(fn)()
         self.assertEqual(ref, res)
 
     def test_object_classmethod(self):
@@ -7372,6 +8147,60 @@ def fn(x, obj):
         res = opt_fn(x, obj)
         self.assertTrue(same(ref, res))
 
+    def test_custom_instancecheck_does_not_cause_extra_init(self):
+        # When __new__ returns an object whose type is not a subclass of cls,
+        # CPython's type.__call__ skips __init__. The polyfill
+        # instantiate_user_defined_class_object must match this behavior even
+        # when the metaclass defines a custom __instancecheck__ that would
+        # return True for isinstance().
+        class Meta(type):
+            def __instancecheck__(cls, instance):
+                return isinstance(instance, Base) and instance.tag == cls._tag
+
+        class Base:
+            def __init__(self, tag="default"):
+                self.tag = tag
+
+        class Child(Base, metaclass=Meta):
+            _tag = "child"
+
+            def __new__(cls):
+                # Returns a Base (not a Child), like ByteStorage.__new__
+                return Base(tag="child")
+
+        def fn():
+            obj = Child()
+            return obj.tag
+
+        ref = fn()
+        self.assertEqual(ref, "child")
+
+        opt_fn = torch.compile(fn, backend="eager")
+        res = opt_fn()
+        self.assertEqual(res, "child")
+
+    def test_custom_instancecheck_init_not_called(self):
+        class AlwaysTrueMeta(type):
+            def __instancecheck__(cls, instance):
+                return True
+
+        class Child(metaclass=AlwaysTrueMeta):
+            def __new__(cls):
+                return object()
+
+            def __init__(self):
+                raise AssertionError("should NOT be called")
+
+        def fn():
+            return Child()
+
+        ref = fn()
+        self.assertIsInstance(ref, object)
+
+        opt_fn = torch.compile(fn, backend="eager")
+        res = opt_fn()
+        self.assertIsInstance(res, object)
+
     def test_variable_tracker_recursively_contains(self):
         # VariableTracker.recursively_contains should be updated correctly when mutation happens
         def fn(x):
@@ -7433,6 +8262,18 @@ def f4():
             self.assertEqual(f(), torch.zeros(2, 2))
             self.assertEqual(opt_f(), torch.ones(2, 2))
 
+    def test_is_exporting_not_true_during_compile(self):
+        def f():
+            if torch.compiler.is_exporting():
+                return torch.ones(2, 2)
+            else:
+                return torch.zeros(2, 2)
+
+        opt_f = torch.compile(f, backend="eager")
+
+        self.assertEqual(f(), torch.zeros(2, 2))
+        self.assertEqual(opt_f(), torch.zeros(2, 2))
+
     def test_torch_generator_set_state(self):
         def fn():
             default_state = torch.default_generator.get_state()
@@ -7681,31 +8522,6 @@ def guard_export_print(guards):
             # This guard was created
             self.assertTrue(guard.name != "nested_fn.__closure__[0].cell_contents")
 
-    @unittest.skipIf(not TEST_MULTIGPU, "need multiple GPU")
-    def test_symint_as_device_kwarg_multi_gpu(self):
-        def fn(rank):
-            # -2 to make device id smaller for easier testing on CI
-            return torch.ones(10, device=rank.size(0) - 2)
-
-        x = torch.randn(2)
-        out = fn(torch.randn(2))
-
-        guard_failure = None
-
-        def guard_failures(failure):
-            nonlocal guard_failure
-            guard_failure = failure
-
-        opt_fn = torch._dynamo.optimize(
-            "eager", guard_fail_fn=guard_failures, dynamic=True
-        )(fn)
-        self.assertEqual(out, opt_fn(x))
-
-        x = torch.randn(3)
-        self.assertEqual(fn(x), opt_fn(x))
-        self.assertTrue(guard_failure is not None)
-        self.assertIn("""tensor 'rank' size mismatch at index 0""", guard_failure[0])
-
     @unittest.skipIf(not TEST_CUDA and not TEST_XPU, "Test requires CUDA or XPU.")
     def test_symint_as_device_kwarg_non_strict_export(self):
         class Mod(torch.nn.Module):
@@ -7838,7 +8654,7 @@ def fn(
 
         tensor = torch.tensor([1, 2, 3, 5, 6, 7])
         mapping = torch.tensor([0, 3, 4, 5, 7])
-        opt = torch.compile(fn, fullgraph=True)
+        opt = torch.compile(fn, fullgraph=True, backend="eager")
         self.assertEqual(fn(tensor, mapping), opt(tensor, mapping))
 
     def test_torch_package_working_with_trace(self):
@@ -7881,7 +8697,7 @@ def injected(x):
 
         args = (torch.randn(3, 2),)
 
-        compiled_fn = torch.compile(fn)
+        compiled_fn = torch.compile(fn, backend="eager")
         _load_precompile_entry(
             fn.__code__,
             torch._dynamo.guards.GuardManagerWrapper(),
@@ -7906,7 +8722,7 @@ def injected(x):
 
         args = (torch.randn(3, 2),)
 
-        compiled_fn = torch.compile(fn)
+        compiled_fn = torch.compile(fn, backend="eager")
         _load_precompile_entry(
             fn.__code__,
             torch._dynamo.guards.GuardManagerWrapper(guard_manager),
@@ -7953,7 +8769,7 @@ def injected_str(x: str):
 
         args = (torch.randn(3, 2),)
 
-        compiled_fn = torch.compile(fn)
+        compiled_fn = torch.compile(fn, backend="eager")
         _load_precompile_entry(
             fn.__code__,
             torch._dynamo.guards.GuardManagerWrapper(guard_manager_bool),
@@ -8003,7 +8819,7 @@ def injected(x):
 
         args = (torch.randn(3, 2),)
 
-        compiled_fn = torch.compile(fn)
+        compiled_fn = torch.compile(fn, backend="eager")
         _load_precompile_entry(
             fn.__code__,
             torch._dynamo.guards.GuardManagerWrapper(),
@@ -8031,7 +8847,7 @@ def injected_bool(x: bool):
 
         args = (torch.randn(3, 2),)
 
-        compiled_fn = torch.compile(fn)
+        compiled_fn = torch.compile(fn, backend="eager")
         _load_precompile_entry(
             fn.__code__,
             torch._dynamo.guards.GuardManagerWrapper(guard_manager_bool),
@@ -8152,6 +8968,99 @@ def fn(x):
 
         self.assertTrue(same(ref, res))
 
+    def test_cast_with_different_module_types(self):
+        # typing.cast works correctly when used in a mixin pattern with
+        # different module types, producing correct results without
+        # graph breaks.
+        from typing import cast
+
+        class Mixin:
+            def get_self_as_module(self):
+                return cast(torch.nn.Module, self)
+
+        class ModuleA(Mixin, torch.nn.Module):
+            def forward(self, x):
+                self.get_self_as_module()
+                return x + 1
+
+        class ModuleB(Mixin, torch.nn.Module):
+            def forward(self, x):
+                self.get_self_as_module()
+                return x + 2
+
+        cnt = torch._dynamo.testing.CompileCounter()
+
+        @torch.compile(backend=cnt)
+        def fn(mod, x):
+            mod.get_self_as_module()
+            return x + 1
+
+        x = torch.randn(4)
+        ref_a = fn.__wrapped__(ModuleA(), x)
+        ref_b = fn.__wrapped__(ModuleB(), x)
+        res_a = fn(ModuleA(), x)
+        res_b = fn(ModuleB(), x)
+
+        self.assertEqual(ref_a, res_a)
+        self.assertEqual(ref_b, res_b)
+        self.assertEqual(cnt.frame_count, 2)
+
+    def test_cast_fullgraph_with_non_tensor(self):
+        # Verify typing.cast works with non-tensor values under fullgraph=True
+        from typing import cast
+
+        def fn(x):
+            val = cast(int, x.shape[0])
+            return x + val
+
+        opt_fn = torch.compile(backend="eager", fullgraph=True)(fn)
+
+        ref = fn(torch.ones(3, 4))
+        res = opt_fn(torch.ones(3, 4))
+
+        self.assertTrue(same(ref, res))
+
+    @torch._dynamo.config.patch(nested_graph_breaks=False)
+    def test_cast_no_recompile_after_graph_break(self):
+        # In FSDP, cast(nn.Module, self) can be called after a
+        # graph break. Without the polyfill + skip_code fix, PEP 523 compiles
+        # typing.cast as a standalone frame with TYPE_MATCH guards on val,
+        # causing recompilation when different module types pass through.
+        # https://github.com/pytorch/pytorch/blob/0feb90404fbeb9b1594ae194f8fd47bbe7f5f245/torch/distributed/fsdp/_fully_shard/_fully_shard.py#L376
+        from typing import cast
+
+        from torch._dynamo.utils import counters
+
+        counters.clear()
+
+        class Base(torch.nn.Module):
+            def get_state(self):
+                torch._dynamo.decorators.graph_break()
+                return cast(torch.nn.Module, self)
+
+        class ModuleA(Base):
+            pass
+
+        class ModuleB(Base):
+            pass
+
+        cnt = torch._dynamo.testing.CompileCounter()
+        a, b = ModuleA(), ModuleB()
+
+        @torch.compile(backend=cnt)
+        def fn(mod, x):
+            mod.get_state()
+            return x + 1
+
+        x = torch.randn(4)
+        fn(a, x)
+        fn(b, x)
+        self.assertEqual(cnt.frame_count, 1)
+        # 5 frames: fn (x2), get_state before graph_break (x2),
+        # get_state resume after graph_break (x1, no recompile).
+        # Without skip_code, typing.cast would add 2 more frames (7 total).
+        self.assertEqual(counters["frames"]["total"], 5)
+
     def test_T_tensor_attribute(self):
         def fn(x, y):
             a = x.T
@@ -8753,10 +9662,32 @@ def fn(x, y, z):
             )
             opt(*inputs)
 
+    def test_sym_max_creates_graph_node(self):
+        # Test that sym_max creates a graph node and works correctly
+        from torch._dynamo.testing import EagerAndRecordGraphs
+
+        def fn(x):
+            max_dim = torch.sym_max(x.size(0), x.size(1))
+            return x.sum() + max_dim
+
+        backend = EagerAndRecordGraphs()
+        compiled_fn = torch.compile(fn, backend=backend, fullgraph=True, dynamic=False)
+
+        x = torch.randn(4, 8)
+        result = compiled_fn(x)
+        expected = x.sum() + 8
+        self.assertEqual(result, expected)
+
+        # Verify sym_max appears in graph
+        self.assertGreater(len(backend.graphs), 0)
+        graph = backend.graphs[0]
+        sym_max_nodes = [n for n in graph.graph.nodes if n.target is torch.sym_max]
+        self.assertEqual(len(sym_max_nodes), 1, "sym_max should be in the graph")
+
     @torch._dynamo.config.patch(capture_scalar_outputs=True)
     @torch._dynamo.config.patch(assume_static_by_default=True)
     def test_symint_copy_into_unbacked_slice(self):
-        @torch.compile()
+        @torch.compile(backend="eager")
         def fn(a, x):
             u0 = torch.tensor(x[0].to(torch.int64).item()).item()
             B, H, T, D = a.shape
@@ -8775,7 +9706,7 @@ def fn(a, x):
 
     @torch._dynamo.config.patch(capture_scalar_outputs=True)
     def test_symint_fold_nontrivial_product_modulo(self):
-        @torch.compile(fullgraph=True)
+        @torch.compile(fullgraph=True, backend="eager")
         def f(x):
             u0, u1 = x.tolist()
             # The condition should fold to true.
@@ -9198,14 +10129,20 @@ def fn(x):
                     new_state = [False] * len(initial_state)
                     new_state[i] = True
                     write_state(new_state)
-                    assert read_state() == new_state
+                    if read_state() != new_state:
+                        raise AssertionError(f"Expected read_state() == {new_state}")
                     last_state.clear()
                     fn(y)
-                    assert last_state == new_state
+                    if last_state != new_state:
+                        raise AssertionError(f"Expected last_state == {new_state}")
                     if round == 0:
-                        assert cnt == i + 1
+                        if cnt != i + 1:
+                            raise AssertionError(f"Expected cnt == {i + 1}, got {cnt}")
                     else:
-                        assert cnt == len(initial_state)
+                        if cnt != len(initial_state):
+                            raise AssertionError(
+                                f"Expected cnt == {len(initial_state)}, got {cnt}"
+                            )
         finally:
             write_state(initial_state)
 
@@ -9228,16 +10165,21 @@ def fn(x):
         try:
             torch.set_grad_enabled(True)
             fn(torch.randn(10))
-            assert value is True
-            assert torch.is_grad_enabled() is False
+            if value is not True:
+                raise AssertionError(f"Expected value is True, got {value}")
+            if torch.is_grad_enabled() is not False:
+                raise AssertionError("Expected grad disabled after fn()")
 
             value = None
             torch.set_grad_enabled(True)
             fn(torch.randn(10))
-            assert value is True
-            assert torch.is_grad_enabled() is False
+            if value is not True:
+                raise AssertionError(f"Expected value is True, got {value}")
+            if torch.is_grad_enabled() is not False:
+                raise AssertionError("Expected grad disabled after fn()")
 
-            assert cnt.frame_count == 1
+            if cnt.frame_count != 1:
+                raise AssertionError(f"Expected frame_count 1, got {cnt.frame_count}")
         finally:
             torch.set_grad_enabled(prior)
 
@@ -9258,23 +10200,30 @@ def check_state():
         @torch.compile(backend=cnt, fullgraph=True)
         def fn(x):
             check_state()
-            assert torch.are_deterministic_algorithms_enabled() is True
+            assert torch.are_deterministic_algorithms_enabled() is True  # noqa: S101
             torch.use_deterministic_algorithms(False, warn_only=False)
             return x + 1
 
         def run_fn():
             torch.use_deterministic_algorithms(True, warn_only=True)
             fn(torch.randn(10))
-            assert value is True
-            assert warn_only is True
-            assert torch.are_deterministic_algorithms_enabled() is False
-            assert torch.is_deterministic_algorithms_warn_only_enabled() is False
+            if value is not True:
+                raise AssertionError(f"Expected value is True, got {value}")
+            if warn_only is not True:
+                raise AssertionError(f"Expected warn_only is True, got {warn_only}")
+            if torch.are_deterministic_algorithms_enabled() is not False:
+                raise AssertionError(
+                    "Expected deterministic algorithms disabled after fn()"
+                )
+            if torch.is_deterministic_algorithms_warn_only_enabled() is not False:
+                raise AssertionError("Expected warn_only disabled after fn()")
 
         try:
             run_fn()
             value, warn_only = None, None
             run_fn()
-            assert cnt.frame_count == 1
+            if cnt.frame_count != 1:
+                raise AssertionError(f"Expected frame_count 1, got {cnt.frame_count}")
         finally:
             torch.use_deterministic_algorithms(prior, warn_only=prior_warn_only)
 
@@ -9286,7 +10235,7 @@ def training_step(self):
                 self()
 
         model = MyModel()
-        compiled_model = torch.compile(model)
+        compiled_model = torch.compile(model, backend="eager")
 
         model.forward = compiled_model.dynamo_ctx(model.forward)
         model.training_step = compiled_model.dynamo_ctx(model.training_step)
@@ -9522,14 +10471,14 @@ def test_inline_dict_function(self):
         def _result_type_dict(dtype):
             return {bool: torch.float32}[dtype]
 
-        @torch.compile
+        @torch.compile(backend="eager")
         def f():
             return torch.ones(3, dtype=_result_type_dict(bool))
 
         self.assertEqual(f(), torch.ones(3, dtype=torch.float32))
 
     def test_inline_dict_function_passed_as_arg(self):
-        @torch.compile
+        @torch.compile(backend="eager")
         def fn(d, x, y):
             if d[x] is torch.float32:
                 return y.cos()
@@ -9571,7 +10520,7 @@ def fn(x, y):
         self.assertTrue(same(ref, res))
 
     def test_recursion_depth_guards(self):
-        @torch.compile(dynamic=True)
+        @torch.compile(dynamic=True, backend="eager")
         def foo(*args, **kwargs):
             if sum(args) == 0:
                 return 0
@@ -9920,6 +10869,30 @@ def fn():
         self.assertEqual(fn_out, compiled_out)
         self.assertEqual(fn_out, (True, False, True, False, True, False))
 
+    def test_class_hasattr_sourceless_descriptor(self):
+        """Test that hasattr on sourceless UserDefinedClassVariable does not graph break."""
+
+        class FlagDescriptor:
+            def __get__(self, instance, owner):
+                if hasattr(owner, "flag"):
+                    return 1
+                return 0
+
+        class WithFlag:
+            flag = True
+            prop = FlagDescriptor()
+
+        class WithoutFlag:
+            prop = FlagDescriptor()
+
+        def fn(x, obj):
+            return x + obj.prop
+
+        compiled_fn = torch.compile(backend="eager", fullgraph=True)(fn)
+        x = torch.randn(3)
+        self.assertEqual(fn(x, WithFlag()), compiled_fn(x, WithFlag()))
+        self.assertEqual(fn(x, WithoutFlag()), compiled_fn(x, WithoutFlag()))
+
     def test_torch_objects_as_keys(self):
         remap = {torch.float16: torch.float32}
 
@@ -10449,7 +11422,8 @@ def test(x):
         # custom strides (all integers)
         x2 = torch.randn(10000)
         x2 = x2.as_strided([4, 4, 4, 4], [1, 2, 4, 8])
-        assert g(x2) == (3, 2, 1, 0)
+        if g(x2) != (3, 2, 1, 0):
+            raise AssertionError(f"Expected g(x2) == (3, 2, 1, 0), got {g(x2)}")
 
         # custom unbacked strides with no ordering: ambiguity check should raise
         xs = torch.tensor([2, 3, 4])
@@ -10634,7 +11608,7 @@ def inner():
 
             return inner
 
-        @torch.compile
+        @torch.compile(backend="eager")
         def start():
             # Obtain the `inner` function, which holds reference to `x`.
             inner = fn()
@@ -10865,7 +11839,7 @@ def fn(target):
             target_dtype = target.dtype
             a = torch.zeros(2, 3, dtype=target_dtype)
             # Constant assert at trace time
-            assert isinstance(target_dtype, torch.dtype)
+            assert isinstance(target_dtype, torch.dtype)  # noqa: S101
             b = torch.zeros(2, 3, dtype=target_dtype)
             c = torch.zeros(2, 3, dtype=target_dtype)
             return a + b + c
@@ -11027,6 +12001,39 @@ def fn(x):
             self.assertEqual(list(eager), list(compiled))
             self.assertEqual(len(counters["graph_break"]), 0)
 
+    def test_itertools_count_from_uncompiled_region(self):
+        counters.clear()
+        counter = itertools.count()
+
+        def fn(x):
+            return x * (next(counter) + 1)
+
+        x = torch.randn([2, 5])
+        compiled_fn = torch.compile(fn, backend="eager", fullgraph=True)
+
+        self.assertEqual(compiled_fn(x), x)
+        self.assertEqual(compiled_fn(x), x * 2)
+        self.assertEqual(next(counter), 2)
+        self.assertEqual(len(counters["graph_break"]), 0)
+
+    def test_itertools_count_already_advanced(self):
+        counters.clear()
+        counter = itertools.count()
+        # Advance the counter before entering the compiled region
+        next(counter)  # 0
+        next(counter)  # 1
+
+        def fn(x):
+            return x * (next(counter) + 1)
+
+        x = torch.randn([2, 5])
+        compiled_fn = torch.compile(fn, backend="eager", fullgraph=True)
+
+        self.assertEqual(compiled_fn(x), x * 3)  # next(counter) = 2
+        self.assertEqual(compiled_fn(x), x * 4)  # next(counter) = 3
+        self.assertEqual(next(counter), 4)
+        self.assertEqual(len(counters["graph_break"]), 0)
+
     def test_itertools_infinite_cycle(self):
         counters.clear()
 
@@ -11341,7 +12348,7 @@ def fn(x, y):
             dc = TestDataClass(x, y)
             return inner_fn(dc)
 
-        fn_opt = torch.compile(fullgraph=True)(fn)
+        fn_opt = torch.compile(fullgraph=True, backend="eager")(fn)
         inps = (torch.ones(2, 2), torch.ones(2, 2))
         actual = fn_opt(*inps)
         expected = fn(*inps)
@@ -11359,7 +12366,7 @@ def fn(x, y):
             torch._dynamo.graph_break()
             return dc.x + dc.y
 
-        fn_opt = torch.compile()(fn)
+        fn_opt = torch.compile(backend="eager")(fn)
         inps = (torch.ones(2, 2), torch.ones(2, 2))
         actual = fn_opt(*inps)
         expected = fn(*inps)
@@ -11379,7 +12386,7 @@ def fn(y):
             z = dc.y + dc.ndc.x
             return z, dc
 
-        fn_opt = torch.compile()(fn)
+        fn_opt = torch.compile(backend="eager")(fn)
         inps = (torch.ones(2, 2),)
         actual = fn_opt(*inps)
         expected = fn(*inps)
@@ -11400,7 +12407,7 @@ def fn(x, y):
             dc = TestDataClass(x, y)
             return inner_fn(dc)
 
-        fn_opt = torch.compile(fullgraph=True)(fn)
+        fn_opt = torch.compile(fullgraph=True, backend="eager")(fn)
         inps = (torch.ones(2, 2), torch.ones(2, 2))
         actual = fn_opt(*inps)
         expected = fn(*inps)
@@ -11423,7 +12430,7 @@ def fn(x, y):
             dc = TestDataClass(x, y)
             return inner_fn(dc)
 
-        fn_opt = torch.compile(fullgraph=True)(fn)
+        fn_opt = torch.compile(fullgraph=True, backend="eager")(fn)
         inps = (torch.ones(2, 2), torch.ones(2, 2))
         actual = fn_opt(*inps)
         expected = fn(*inps)
@@ -11446,7 +12453,7 @@ def fn(x, y):
             dc = TestDataClass(x, y, z=5, a=2)
             return inner_fn(dc)
 
-        fn_opt = torch.compile(fullgraph=True)(fn)
+        fn_opt = torch.compile(fullgraph=True, backend="eager")(fn)
         inps = (torch.ones(2, 2), torch.ones(2, 2))
         actual = fn_opt(*inps)
         expected = fn(*inps)
@@ -11468,7 +12475,7 @@ def fn(x, y):
             dc = TestDataClass(x, y, z=5, a=2)
             return inner_fn(dc)
 
-        fn_opt = torch.compile(fullgraph=True)(fn)
+        fn_opt = torch.compile(fullgraph=True, backend="eager")(fn)
         inps = (torch.ones(2, 2), torch.ones(2, 2))
         actual = fn_opt(*inps)
         expected = fn(*inps)
@@ -11492,7 +12499,7 @@ def fn(x, y):
             dc = TestDataClass(x=3.2, y=2.5, z=5, a=2)
             return inner_fn(dc, x, y)
 
-        fn_opt = torch.compile(fullgraph=True)(fn)
+        fn_opt = torch.compile(fullgraph=True, backend="eager")(fn)
         inps = (torch.ones(2, 2), torch.ones(2, 2))
         actual = fn_opt(*inps)
         expected = fn(*inps)
@@ -11519,12 +12526,41 @@ def fn(x, y):
             dc = TestDataClass(b=TestDataClassInner(2.4, 4.4), z=5, a=2)
             return inner_fn(dc, x, y)
 
-        fn_opt = torch.compile(fullgraph=True)(fn)
+        fn_opt = torch.compile(fullgraph=True, backend="eager")(fn)
         inps = (torch.ones(2, 2), torch.ones(2, 2))
         actual = fn_opt(*inps)
         expected = fn(*inps)
         self.assertEqual(actual, expected)
 
+    def test_frozen_dataclass_in_compile(self):
+        from torch.utils._pytree import MappingKey, SequenceKey
+
+        def fn(x):
+            path = (MappingKey("a"), SequenceKey(0))
+            msg = f"path={path}"
+            return x * 2, msg
+
+        x = torch.randn(4, 4)
+        eager_result = fn(x)
+        compiled_result = torch.compile(fn, fullgraph=True)(x)
+        self.assertEqual(eager_result[0], compiled_result[0])
+        self.assertEqual(eager_result[1], compiled_result[1])
+
+    def test_frozen_dataclass_treespec_method_and_fields(self):
+        from torch.utils._pytree import tree_flatten
+
+        def fn(x):
+            d = {"a": x, "b": [x * 2, x * 3]}
+            flat, spec = tree_flatten(d)
+            is_leaf = spec.is_leaf()
+            return sum(flat), spec.num_leaves, spec.num_nodes, is_leaf
+
+        x = torch.randn(4)
+        eager_result = fn(x)
+        compiled_result = torch.compile(fn, fullgraph=True)(x)
+        for i in range(4):
+            self.assertEqual(eager_result[i], compiled_result[i])
+
     def test_shape_env_no_recording(self):
         main = ShapeEnv(should_record_events=False)
 
@@ -11823,7 +12859,7 @@ def test_shape_env_recorded_function_fallback(self):
         )
 
     def test_default_dtype_change(self):
-        @torch.compile
+        @torch.compile(backend="eager")
         def foo():
             def inner(a, b, res_dtype):
                 print(a, b, res_dtype)
@@ -11842,7 +12878,7 @@ def foo():
             x = np.arange(5)
             out = np.empty((x.shape[0], x.shape[0]))
             res_out = np.sin(x, out=out)
-            assert res_out is out
+            assert res_out is out  # noqa: S101
 
         foo()
 
@@ -11856,13 +12892,13 @@ def foo():
             out = np.empty((x.shape[0], x.shape[0]))
             res_out = np.sin(x, out=out)
             torch._dynamo.graph_break()
-            assert res_out is out
+            assert res_out is out  # noqa: S101
 
         foo()
 
     @wrapDeterministicFlagAPITest
     def test_backward_deterministic_mode_mismatch_warning(self):
-        @torch.compile
+        @torch.compile(backend="aot_eager")
         def func(a, b):
             return a + b
 
@@ -11897,7 +12933,7 @@ def pow(x):
             return x**2
 
         x = np.arange(8)
-        pow_opt = torch.compile(pow)
+        pow_opt = torch.compile(pow, backend="eager")
 
         actual, source_code = run_and_get_code(pow_opt, x)
         expect = pow(x)
@@ -12127,7 +13163,7 @@ def fn(x):
 
     @torch._dynamo.config.patch(assume_static_by_default=True)
     def test_mark_unbacked_strict(self):
-        @torch.compile()
+        @torch.compile(backend="eager")
         def fn(x, y):
             return torch.mul(x, y)
 
@@ -12146,14 +13182,14 @@ def test_infer_unbacked_size_gt_zero(self):
         def fn(x):
             y = torch.zeros(x.item())
             if y.size(0) < 0:
-                assert False
+                assert False  # noqa: B011, S101
             return y
 
         self.assertEqual(fn(torch.tensor([0])), torch.zeros(0))
 
     @torch.fx.experimental._config.patch(no_data_dependent_graph_break=True)
     def test_unbacked_strict_mode(self):
-        @torch.compile()
+        @torch.compile(backend="eager")
         def fn(x, y):
             if x.shape[0] == 5:
                 return torch.randn(5)
@@ -12414,62 +13450,6 @@ def fn(x, const):
         c2 = _debug_get_cache_entry_list(fn.__code__)
         self.assertIs(c1[1], c2[0])
 
-    @torch._dynamo.config.patch(inline_inbuilt_nn_modules=False)
-    @skipIfWindows(msg="TODO: (xuhancn) conform, AssertionError: False is not true")
-    def test_dynamo_cache_invalidate(self):
-        DeletedGuardManagerWrapper = torch._dynamo.guards.DeletedGuardManagerWrapper
-
-        class Mod(torch.nn.Module):
-            def __init__(self) -> None:
-                super(Mod, self).__init__()
-                self.fc = torch.nn.Linear(3, 3)
-
-            def forward(self, out):
-                return self.fc(out)
-
-        def fn(x, mod):
-            return mod(x)
-
-        opt_fn = torch.compile(fn, backend="eager")
-
-        m1 = Mod()
-        m2 = Mod()
-        m3 = Mod()
-        inp = torch.randn(3, 3)
-
-        # NOTE: assumes that each cache entry is guarded
-        # on unique Mod instance
-        opt_fn(inp, m1)
-        opt_fn(inp, m2)
-        opt_fn(inp, m3)
-
-        c1 = _debug_get_cache_entry_list(fn.__code__)
-        self.assertEqual(len(c1), 3)
-
-        # move cache entry to front
-        opt_fn(inp, m2)
-        c2 = _debug_get_cache_entry_list(fn.__code__)
-        self.assertIs(c1[1], c2[0])
-
-        # delete center of cache
-        del m3
-        c3 = _debug_get_cache_entry_list(fn.__code__)
-        self.assertEqual(len(c3), 3)
-        self.assertTrue(isinstance(c3[2].guard_manager, DeletedGuardManagerWrapper))
-
-        # delete end of cache
-        del m1
-        c4 = _debug_get_cache_entry_list(fn.__code__)
-        self.assertEqual(len(c4), 3)
-        self.assertTrue(isinstance(c4[1].guard_manager, DeletedGuardManagerWrapper))
-        self.assertTrue(isinstance(c4[2].guard_manager, DeletedGuardManagerWrapper))
-
-        del m2
-        c5 = _debug_get_cache_entry_list(fn.__code__)
-        self.assertTrue(isinstance(c5[0].guard_manager, DeletedGuardManagerWrapper))
-        self.assertTrue(isinstance(c5[1].guard_manager, DeletedGuardManagerWrapper))
-        self.assertTrue(isinstance(c5[2].guard_manager, DeletedGuardManagerWrapper))
-
     def test_inspect_signature_bind(self):
         import inspect
 
@@ -12886,7 +13866,7 @@ def test_hash_hop(self):
             "torch._higher_order_ops.associative_scan"
         )
 
-        @torch.compile(fullgraph=True)
+        @torch.compile(fullgraph=True, backend="eager")
         def fn(y, s):
             d = dict()
             d[s] = y
@@ -12895,7 +13875,7 @@ def fn(y, s):
         fn(torch.ones(2, 2, device="cpu"), associative_scan.AssociativeScanOp())
 
     def test_iter_type(self):
-        @torch.compile(fullgraph=True)
+        @torch.compile(fullgraph=True, backend="eager")
         def fn(y):
             x = iter([])
             if isinstance(x, list):
@@ -12967,6 +13947,136 @@ def fn(x):
         self.assertEqual(ref_t, res_t)
         self.assertEqual(ref_obj.field, res_obj.field)
 
+    def test_data_descriptor_priority_over_instance_dict(self):
+        # CPython: data descriptors on the type take priority over instance
+        # dict values. Verify Dynamo follows this ordering.
+        class Foo:
+            @property
+            def x(self):
+                return 10
+
+        foo = Foo()
+        # Manually put a different value in the instance dict.
+        # The property (data descriptor) should still win.
+        foo.__dict__["x"] = 999
+
+        def fn(t):
+            return t + foo.x
+
+        opt_fn = torch.compile(fn, backend="eager", fullgraph=True)
+        t = torch.randn(4)
+        ref = fn(t)
+        res = opt_fn(t)
+        self.assertEqual(ref, res)
+        self.assertEqual(ref, t + 10)
+
+    def test_instance_dict_priority_over_non_data_descriptor(self):
+        # CPython: instance dict values take priority over non-data
+        # descriptors (those with only __get__, no __set__/__delete__).
+        class Desc:
+            def __init__(self, val):
+                self.val = val
+
+            def __get__(self, obj, cls):
+                return self.val * 100
+
+        class Foo:
+            x = Desc(7)
+
+        foo = Foo()
+        # Instance dict value should shadow the non-data descriptor.
+        foo.__dict__["x"] = 10
+
+        def fn(t):
+            return t + foo.x
+
+        opt_fn = torch.compile(fn, backend="eager", fullgraph=True)
+        t = torch.randn(4)
+        ref = fn(t)
+        res = opt_fn(t)
+        self.assertEqual(ref, res)
+        # CPython: instance dict wins → foo.x is 10, not 700
+        self.assertEqual(ref, t + 10)
+
+    def test_user_defined_data_descriptor(self):
+        # A user-defined data descriptor (has __get__ + __set__) on the type
+        # should be invoked even when the same name exists in the instance dict.
+        class ValidatedAttr:
+            def __set_name__(self, owner, name):
+                self.storage_name = "_" + name
+
+            def __get__(self, obj, cls):
+                if obj is None:
+                    return self
+                return getattr(obj, self.storage_name)
+
+            def __set__(self, obj, value):
+                setattr(obj, self.storage_name, value)
+
+        class Foo:
+            x = ValidatedAttr()
+
+            def __init__(self, x):
+                self.x = x
+
+        foo = Foo(10)
+
+        def fn(t):
+            return t + foo.x
+
+        opt_fn = torch.compile(fn, backend="eager", fullgraph=True)
+        t = torch.randn(4)
+        ref = fn(t)
+        res = opt_fn(t)
+        self.assertEqual(ref, res)
+        self.assertEqual(ref, t + 10)
+
+    def test_property_setter(self):
+        class Box:
+            def __init__(self, value):
+                self._value = torch.tensor([value], dtype=torch.float32)
+
+            @property
+            def value(self):
+                return self._value + 1
+
+            @value.setter
+            def value(self, new_value):
+                self._value = new_value * 2
+
+        def fn(b):
+            b.value = 5
+            return b.value
+
+        b = Box(0)
+        opt_fn = torch.compile(fn, backend="eager", fullgraph=True)
+        ref = fn(Box(0))
+        res = opt_fn(b)
+        self.assertEqual(ref, res)
+
+    def test_property_setter_in_init(self):
+        class Clamped:
+            def __init__(self, val):
+                self.val = val
+
+            @property
+            def val(self):
+                return self._val
+
+            @val.setter
+            def val(self, v):
+                self._val = torch.clamp(v, 0, 100)
+
+        def fn(x):
+            obj = Clamped(x)
+            return obj.val
+
+        x = torch.tensor([200.0])
+        opt_fn = torch.compile(fn, backend="eager", fullgraph=True)
+        ref = fn(x)
+        res = opt_fn(x)
+        self.assertEqual(ref, res)
+
     def test_assert_size_stride(self):
         x = torch.randn(2, 3, 4)
         with self.assertRaisesRegex(
@@ -13274,7 +14384,7 @@ def inner():
         nonlocal_fn()
 
     def test_compare_tensor_with_none(self):
-        @torch.compile()
+        @torch.compile(backend="eager")
         def f(x):
             return torch.tensor(x == None)
 
@@ -13387,6 +14497,12 @@ def run(x):
         self.assertFalse(ne)
         self.assertEqual(len(counters["graph_break"]), 1)
 
+    @unittest.skipIf(sys.version_info < (3, 12), "Python 3.12+")
+    def test_CALL_INTRINSIC(self):
+        from torch.testing._internal.py312_intrinsics import Foo
+
+        Foo.test_default_update(self)
+
     @unittest.skipIf(sys.version_info < (3, 11), "Python 3.11+")
     def test_RAISE_VARARGS_0(self):
         def foo():
@@ -13476,7 +14592,11 @@ def test_guard_filter_fn_by_id(self):
         def guard_filter_fn(entries):
             return [entry.guard_type != "ID_MATCH" for entry in entries]
 
-        @torch.compile(fullgraph=True, options={"guard_filter_fn": guard_filter_fn})
+        @torch.compile(
+            fullgraph=True,
+            options={"guard_filter_fn": guard_filter_fn},
+            backend="eager",
+        )
         def fn(x):
             return id(x)
 
@@ -13493,7 +14613,11 @@ def guard_filter_fn(entries):
 
         global GLOBAL_INT
 
-        @torch.compile(fullgraph=True, options={"guard_filter_fn": guard_filter_fn})
+        @torch.compile(
+            fullgraph=True,
+            options={"guard_filter_fn": guard_filter_fn},
+            backend="eager",
+        )
         def fn(x):
             return x + GLOBAL_INT
 
@@ -13511,7 +14635,11 @@ def guard_filter_fn(entries):
                 not (entry.name == "y" and entry.value is None) for entry in entries
             ]
 
-        @torch.compile(fullgraph=True, options={"guard_filter_fn": guard_filter_fn})
+        @torch.compile(
+            fullgraph=True,
+            options={"guard_filter_fn": guard_filter_fn},
+            backend="eager",
+        )
         def fn(x, y):
             if y is not None:
                 x += y
@@ -13535,6 +14663,7 @@ def forward(self, x):
         mod = Mod()
         opt_mod = torch.compile(
             mod,
+            backend="eager",
             options={
                 "guard_filter_fn": torch.compiler.skip_guard_on_inbuilt_nn_modules_unsafe
             },
@@ -13561,6 +14690,7 @@ def forward(self, x):
         mod = Mod()
         opt_mod = torch.compile(
             mod,
+            backend="eager",
             options={
                 "guard_filter_fn": torch.compiler.skip_guard_on_all_nn_modules_unsafe
             },
@@ -13587,6 +14717,7 @@ def forward(self, x):
         mod = Mod()
         opt_mod = torch.compile(
             mod,
+            backend="eager",
             options={
                 "guard_filter_fn": torch.compiler.keep_tensor_guards_unsafe,
             },
@@ -13613,6 +14744,7 @@ def forward(self, x):
         mod = Mod()
         opt_mod = torch.compile(
             mod,
+            backend="eager",
             options={
                 "guard_filter_fn": torch.compiler.skip_guard_on_globals_unsafe,
             },
@@ -13640,6 +14772,39 @@ def f(x):
         res = opt_f(x)
         self.assertEqual(ref, res)
 
+    @torch._dynamo.config.patch(
+        capture_scalar_outputs=True, capture_dynamic_output_shape_ops=True
+    )
+    def test_new_tensor_break(self):
+        a = torch.tensor([1, 0, 0, 5])
+
+        cases = {
+            "scalar": lambda a: a.new_tensor([a.nonzero().squeeze(-1).numel()]),
+            "multi": lambda a: (
+                n := a.nonzero().squeeze(-1).numel(),
+                a.new_tensor([n, n + 1, n * 2]),
+            )[-1],
+            "mixed_shape": lambda a: (
+                n := a.nonzero().squeeze(-1).numel(),
+                a.new_tensor([n * a.shape[0], n + a.shape[0], a.shape[0] - n]),
+            )[-1],
+            "nested": lambda a: (
+                n := a.nonzero().squeeze(-1).numel(),
+                a.new_tensor([[n, n + 1], [n * 2, n - 1]]),
+            )[-1],
+            "with_zero": lambda a: (
+                n := a.nonzero().squeeze(-1).numel(),
+                a.new_tensor([0, n, n * n]),
+            )[-1],
+        }
+
+        for name, fn in cases.items():
+            with self.subTest(case=name):
+                self.assertEqual(
+                    torch.compile(fn, fullgraph=True, backend="eager")(a),
+                    fn(a),
+                )
+
     @torch._dynamo.config.patch(capture_scalar_outputs=True)
     def test_builtin_bool_on_tensor(self):
         def f_all(mask):
@@ -13722,6 +14887,36 @@ def f(*args, **kwargs):
         self.assertRaises(Unsupported, f, [])
         self.assertRaises(Unsupported, f, "1 + j")
 
+    def test_builtin_class_method_constant_fold(self):
+        @torch.compile(backend="eager", fullgraph=True)
+        def fn():
+            return (
+                bool.__new__(bool),
+                bool.__new__(bool, 1),
+                bool.__new__(bool, 0),
+                bool.from_bytes(b"\x00" * 8, "big"),
+                bool.from_bytes(b"abcd", "little"),
+                int.__new__(int),
+                int.__new__(int, 42),
+                int.from_bytes(b"\x00\x03", "big"),
+                int.from_bytes(b"\xff", byteorder="big", signed=True),
+                float.fromhex("0x1.ffffp10"),
+                float.hex(1.5),
+            )
+
+        res = fn()
+        self.assertIs(res[0], False)
+        self.assertIs(res[1], True)
+        self.assertIs(res[2], False)
+        self.assertIs(res[3], False)
+        self.assertIs(res[4], True)
+        self.assertEqual(res[5], 0)
+        self.assertEqual(res[6], 42)
+        self.assertEqual(res[7], 3)
+        self.assertEqual(res[8], -1)
+        self.assertEqual(res[9], float.fromhex("0x1.ffffp10"))
+        self.assertEqual(res[10], "0x1.8000000000000p+0")
+
     def test_guard_string_escaped(self):
         d = {frozenset({0}): {frozenset({0}): 1}}
 
@@ -13827,6 +15022,351 @@ def fn(x):
                     "to not return y. This is deprecated and will become an error in PyTorch 2.12.",
                 )
 
+    def test_make_contiguous_strides_for_under_compile(self):
+        # is_nested_int and sym_max must be traceable under Dynamo.
+        from torch._prims_common import make_contiguous_strides_for
+
+        def fn(x):
+            strides = make_contiguous_strides_for(x.shape)
+            return x.as_strided(x.shape, strides)
+
+        compiled_fn = torch.compile(fn, backend="eager", fullgraph=True, dynamic=True)
+        x = torch.randn(4, 8)
+        result = compiled_fn(x)
+        self.assertEqual(result, x)
+
+        x2 = torch.randn(7, 8)
+        result2 = compiled_fn(x2)
+        self.assertEqual(result2, x2)
+
+    def test_requires_grad_changes_dynamo_graph(self):
+        # requires_grad_() on a graph input graph-breaks, so no fullgraph
+        def fn(x):
+            x.requires_grad_()
+            if x.requires_grad:
+                return x * 2
+            return x + 1
+
+        x = torch.randn(3, 3)
+        opt_fn = torch.compile(fn)
+        result = opt_fn(x)
+        self.assertEqual(result, x * 2)
+
+    def test_requires_grad_backward_outside_compile(self):
+        # requires_grad_() on a graph input graph-breaks, but eager fallback
+        # produces correct results.
+        def fn(x):
+            x.requires_grad_()
+            return (x * 2).sum()
+
+        x_ref = torch.randn(3, 3)
+        x_test = x_ref.clone()
+
+        fn(x_ref).backward()
+        torch.compile(fn)(x_test).backward()
+
+        self.assertEqual(x_ref.grad, x_test.grad)
+
+    def test_detach_inplace_on_intermediate_updates_metadata(self):
+        def fn(x):
+            y = x * 2
+            y.detach_()
+            return y + 1, y.requires_grad, y.grad_fn is None
+
+        x = torch.randn(3, 3, requires_grad=True)
+        ref = fn(x.clone())
+        result = torch.compile(fn, backend="eager", fullgraph=True)(x.clone())
+
+        self.assertEqual(ref, result)
+        self.assertFalse(result[1])
+        self.assertTrue(result[2])
+
+    def test_requires_grad_on_intermediate(self):
+        def fn(x):
+            y = x * 2
+            y.requires_grad_()
+            return y
+
+        x = torch.randn(3, 3)
+
+        # fullgraph=True should error with actionable message
+        with self.assertRaisesRegex(
+            torch._dynamo.exc.Unsupported,
+            r"requires_grad_\(\)(.|\n)*\.detach\(\)",
+        ):
+            torch.compile(fn, fullgraph=True)(x)
+
+        # Without fullgraph, falls back to eager and is correct
+        result = torch.compile(fn)(x)
+        self.assertTrue(result.requires_grad)
+        self.assertEqual(fn(x), result)
+
+    def test_requires_grad_on_intermediate_derived_returned(self):
+        def fn(x):
+            y = x * 2
+            y.requires_grad_()
+            return y * 3
+
+        x = torch.randn(3, 3)
+
+        # Derived tensor also loses requires_grad — should error with message
+        with self.assertRaisesRegex(
+            torch._dynamo.exc.Unsupported,
+            r"requires_grad_\(\)(.|\n)*\.detach\(\)",
+        ):
+            torch.compile(fn, fullgraph=True)(x)
+
+        # Without fullgraph, falls back to eager and is correct
+        result = torch.compile(fn)(x)
+        ref = fn(x)
+        self.assertTrue(result.requires_grad)
+        self.assertEqual(ref, result)
+
+    def test_requires_grad_on_intermediate_partial_graph(self):
+        # When requires_grad_() on a source-less intermediate leaks as output,
+        # Dynamo should restart and graph break at requires_grad_(), capturing
+        # ops before it in a compiled graph (partial acceleration).
+        def fn(x):
+            a = x.sin()
+            b = a.cos()
+            b.requires_grad_()
+            return b
+
+        backend = torch._dynamo.testing.EagerAndRecordGraphs()
+        x = torch.randn(3, 3)
+        result = torch.compile(fn, backend=backend)(x)
+        self.assertEqual(result, fn(x))
+        self.assertTrue(result.requires_grad)
+        # The graph should capture the ops before requires_grad_()
+        self.assertEqual(len(backend.graphs), 1)
+        # Dynamic shapes adds shape guards to the graph, skip the exact check
+        if torch._dynamo.config.assume_static_by_default:
+            self.assertExpectedInline(
+                backend.graphs[0].code.strip(),
+                """\
+def forward(self, L_x_ : torch.Tensor):
+    l_x_ = L_x_
+    a = l_x_.sin();  l_x_ = None
+    b = a.cos();  a = None
+    return (b,)""",
+            )
+
+    @torch._dynamo.config.patch(trace_autograd_ops=True)
+    def test_requires_grad_on_intermediate_not_returned(self):
+        def fn(x):
+            y = x * 2
+            y.requires_grad_()
+            loss = (y * 3).sum()
+            loss.backward()
+            return y.grad
+
+        x = torch.randn(3, 3)
+
+        ref = fn(x.clone())
+        result = torch.compile(fn, fullgraph=True)(x.clone())
+        self.assertEqual(ref, result)
+
+    @torch._dynamo.config.patch(trace_autograd_ops=True)
+    def test_requires_grad_intermediate_backward_grad_used_in_compute(self):
+        # Use the grad result in further computation within compile
+        def fn(x):
+            y = x * 2
+            y.requires_grad_()
+            loss = (y**2).sum()
+            loss.backward()
+            return y.grad * 2 + 1
+
+        x = torch.randn(3, 3)
+
+        ref = fn(x.clone())
+        result = torch.compile(fn, fullgraph=True)(x.clone())
+        self.assertEqual(ref, result)
+
+    @torch._dynamo.config.patch(trace_autograd_ops=True)
+    def test_requires_grad_intermediate_chunked_loss_backward(self):
+        # Mirrors the TxtUnembedding pattern: forward compute, detach, make
+        # new leaf, chunked loss with per-chunk backward, then propagate
+        # accumulated grad back to the original input via h.backward().
+        def fn(x, targets):
+            # Forward computation before detach (e.g. transformer layers)
+            h = x * 2 + 1
+            x_detached = h.detach().requires_grad_()
+            chunksz = x_detached.shape[0] // 2
+            total_loss = torch.tensor(0.0)
+            for start in range(0, x_detached.shape[0], chunksz):
+                chunk = x_detached[start : start + chunksz]
+                chunk_targets = targets[start : start + chunksz]
+                logits = chunk @ torch.eye(chunk.shape[-1])
+                loss = torch.nn.functional.cross_entropy(logits, chunk_targets)
+                loss.backward()
+                total_loss = total_loss + loss.detach()
+            # Propagate chunked grad back through the forward computation
+            h.backward(x_detached.grad)
+            return x.grad, total_loss
+
+        x_ref = torch.randn(4, 8, requires_grad=True)
+        targets = torch.randint(0, 8, (4,))
+
+        x_test = x_ref.clone().detach().requires_grad_(True)
+        ref_grad, ref_loss = fn(x_ref, targets)
+        compiled_grad, compiled_loss = torch.compile(fn, fullgraph=True)(
+            x_test, targets
+        )
+        self.assertEqual(ref_grad, compiled_grad)
+        self.assertEqual(ref_loss, compiled_loss)
+        # Verify grad propagated to the input
+        self.assertEqual(x_ref.grad, x_test.grad)
+
+    @torch._dynamo.config.patch(trace_autograd_ops=True)
+    def test_requires_grad_intermediate_backward_and_return_detached(self):
+        # Returning a detached version of the tainted tensor is safe — detach()
+        # strips requires_grad so AOTAutograd functionalization can't lose anything.
+        def fn(x):
+            y = x * 2
+            y.requires_grad_()
+            out = y * 3
+            loss = out.sum()
+            loss.backward()
+            return y.grad, out.detach()
+
+        x = torch.randn(3, 3)
+
+        ref_grad, ref_out = fn(x.clone())
+        compiled_grad, compiled_out = torch.compile(fn, fullgraph=True)(x.clone())
+        self.assertEqual(ref_grad, compiled_grad)
+        self.assertEqual(ref_out, compiled_out)
+        self.assertFalse(compiled_out.requires_grad)
+
+    @torch._dynamo.config.patch(trace_autograd_ops=True)
+    def test_requires_grad_intermediate_metadata_checks(self):
+        # After requires_grad_() on an intermediate, requires_grad and is_leaf
+        # should report correctly and be usable in control flow.
+        def fn(x):
+            y = x * 2
+            y.requires_grad_()
+            if y.requires_grad and y.is_leaf:
+                loss = (y * 3).sum()
+                loss.backward()
+                return y.grad
+            return y
+
+        x = torch.randn(3, 3)
+        ref = fn(x.clone())
+        result = torch.compile(fn, fullgraph=True)(x.clone())
+        self.assertEqual(ref, result)
+
+    @torch._dynamo.config.patch(trace_autograd_ops=True)
+    def test_requires_grad_intermediate_side_effect_global(self):
+        # requires_grad_() on intermediate, then store grad in a global
+        saved = {}
+
+        def fn(x):
+            y = x * 2
+            y.requires_grad_()
+            loss = (y**2).sum()
+            loss.backward()
+            saved["grad"] = y.grad
+            return y.grad.clone()
+
+        x = torch.randn(3, 3)
+        ref = fn(x.clone())
+        saved_ref = saved["grad"].clone()
+        saved.clear()
+
+        result = torch.compile(fn, fullgraph=True)(x.clone())
+        self.assertEqual(ref, result)
+        self.assertEqual(saved_ref, saved["grad"])
+
+    def test_import_user_defined_module(self):
+        # testcase for https://github.com/pytorch/pytorch/issues/177682
+        # Bad import result for types.ModuleType subclass in sys.modules
+        class _ConfigModule(types.ModuleType):
+            x = 1
+
+        _ConfigModule.__module__ = __name__
+        sys.modules["my_config"] = _ConfigModule("my_config")
+
+        def fn():
+            import my_config  # noqa: F401
+
+            return torch.tensor(1)
+
+        compilefn = torch.compile(fn, fullgraph=True, backend="eager")
+
+        ret1 = fn()
+        ret2 = compilefn()
+        self.assertEqual(ret1, ret2)
+
+    def test_constant_subclass_guard_recompiles(self):
+        class MyInt(int):
+            def __eq__(self, other):
+                raise RuntimeError("should not be called during guard check")
+
+        class MyFloat(float):
+            def __eq__(self, other):
+                raise RuntimeError("should not be called during guard check")
+
+        class MyStr(str):
+            def __eq__(self, other):
+                raise RuntimeError("should not be called during guard check")
+
+        cnt = torch._dynamo.testing.CompileCounter()
+
+        # int subclass
+        @torch.compile(backend=cnt)
+        def f(x, y):
+            return x + y
+
+        r1 = f(torch.tensor(1), MyInt(5))
+        self.assertEqual(r1.item(), 6)
+        self.assertEqual(cnt.frame_count, 1)
+
+        r2 = f(torch.tensor(1), MyInt(10))
+        self.assertEqual(r2.item(), 11)
+        self.assertEqual(cnt.frame_count, 2)
+
+        r3 = f(torch.tensor(1), MyInt(5))
+        self.assertEqual(r3.item(), 6)
+        self.assertEqual(cnt.frame_count, 2)
+
+        # float subclass
+        cnt.clear()
+
+        @torch.compile(backend=cnt)
+        def g(x, y):
+            return x + y
+
+        r4 = g(torch.tensor(1.0), MyFloat(2.5))
+        self.assertEqual(r4.item(), 3.5)
+        self.assertEqual(cnt.frame_count, 1)
+
+        r5 = g(torch.tensor(1.0), MyFloat(3.5))
+        self.assertEqual(r5.item(), 4.5)
+        self.assertEqual(cnt.frame_count, 2)
+
+        r6 = g(torch.tensor(1.0), MyFloat(2.5))
+        self.assertEqual(r6.item(), 3.5)
+        self.assertEqual(cnt.frame_count, 2)
+
+        # str subclass
+        cnt.clear()
+
+        @torch.compile(backend=cnt, fullgraph=True)
+        def h(x, s):
+            return x + len(s)
+
+        r7 = h(torch.tensor(1), MyStr("abc"))
+        self.assertEqual(r7.item(), 4)
+        self.assertEqual(cnt.frame_count, 1)
+
+        r8 = h(torch.tensor(1), MyStr("abcde"))
+        self.assertEqual(r8.item(), 6)
+        self.assertEqual(cnt.frame_count, 2)
+
+        r9 = h(torch.tensor(1), MyStr("abc"))
+        self.assertEqual(r9.item(), 4)
+        self.assertEqual(cnt.frame_count, 2)
+
 
 class MiscTestsPyTree(torch._inductor.test_case.TestCase):
     @parametrize_pytree_module
@@ -13980,7 +15520,7 @@ def fn(x):
 
         x = torch.randn(3, 2)
         expected = fn(x)
-        fn_opt = torch.compile(fullgraph=True)(fn)
+        fn_opt = torch.compile(fullgraph=True, backend="eager")(fn)
         actual = fn_opt(x)
 
         self.assertEqual(actual, expected)
@@ -14031,7 +15571,7 @@ def fn(x, y):
         x = torch.randn(3, 2)
         y = torch.randn(3, 2)
         expected = fn(x, y)
-        fn_opt = torch.compile(fullgraph=True)(fn)
+        fn_opt = torch.compile(fullgraph=True, backend="eager")(fn)
         actual = fn_opt(x, y)
 
         self.assertEqual(actual, expected)
@@ -14073,7 +15613,7 @@ def fn(x, y):
         x = torch.randn(3, 2)
         y = torch.randn(3, 2)
         expected = fn(x, y)
-        fn_opt = torch.compile(fullgraph=True)(fn)
+        fn_opt = torch.compile(fullgraph=True, backend="eager")(fn)
         actual = fn_opt(x, y)
 
         self.assertEqual(actual, expected)
@@ -14085,7 +15625,7 @@ def fn(tree):
             return list(new_tree.keys()), list(new_tree.values())
 
         x = torch.randn(3, 2)
-        fn_opt = torch.compile(fullgraph=True)(fn)
+        fn_opt = torch.compile(fullgraph=True, backend="eager")(fn)
 
         tree1 = {"b": x + 2, "a": x, "c": x - 1}
         expected1 = fn(tree1)
@@ -14220,8 +15760,8 @@ def usual_function(inp1, inp2) -> torch.Tensor:
         inp1_usual = inp1_custom.detach().clone().requires_grad_(True)
 
         inp2 = torch.randn(2, 4)
-        c_custom_func = torch.compile(outer_function)
-        c_usual_func = torch.compile(usual_function)
+        c_custom_func = torch.compile(outer_function, backend="eager")
+        c_usual_func = torch.compile(usual_function, backend="eager")
 
         result_custom = c_custom_func(inp1_custom, inp2)
         result_custom.backward()
@@ -14415,21 +15955,6 @@ def f(rank):
         opt_out = torch.compile(backend="eager", dynamic=True, fullgraph=True)(f)(x)
         self.assertEqual(out, opt_out)
 
-    @unittest.skipIf(not TEST_MULTIGPU, "need multiple GPU")
-    def test_gpu_set_device(self, device):
-        def fn():
-            a = torch.ones(2, device=device)
-            torch.get_device_module(device).set_device(1)
-            return a + 1
-
-        with torch.get_device_module(device).device(0):
-            counter = CompileCounter()
-            opt_fn = torch.compile(fn, backend=counter)
-            res = opt_fn()
-            self.assertTrue(res.device.type in device)
-            self.assertEqual(res.device.index, 0)
-            self.assertEqual(counter.frame_count, 2)
-
     def test_torch_device_python_type(self, device):
         device_type = torch.device(device).type
         for device, device_type, index in [
@@ -14441,9 +15966,9 @@ def fn(target):
                 target_device = target.device
                 a = torch.zeros(2, 3, device=target_device)
                 # Constant assert at trace time
-                assert isinstance(target_device, torch.device)
-                assert target_device.type == device_type
-                assert target_device.index == index
+                assert isinstance(target_device, torch.device)  # noqa: S101
+                assert target_device.type == device_type  # noqa: S101
+                assert target_device.index == index  # noqa: S101
                 b = torch.zeros(2, 3, device=target_device)
                 c = torch.zeros(2, 3, device=target_device)
                 return a + b + c
@@ -14497,7 +16022,7 @@ def f(actions, n_act, epsilon=0.1):
         f(x, y)
 
     def test_full_graph_capture_scalar_outputs(self):
-        @torch.compile(fullgraph=True)
+        @torch.compile(fullgraph=True, backend="eager")
         def foo(a):
             return torch.randn(5) * a.item()
 
@@ -14514,10 +16039,10 @@ def fn(x):
 
         example_inputs = (torch.randn(1, 1, 1, 1),)
         # we expect to no longer raise here
-        torch.compile(fn, fullgraph=True)(*example_inputs)
+        torch.compile(fn, fullgraph=True, backend="eager")(*example_inputs)
 
     def test_dynamic_fill_diagonal_(self):
-        @torch.compile(dynamic=True)
+        @torch.compile(dynamic=True, backend="eager")
         def f(x):
             x.fill_diagonal_(True)
 
@@ -14532,7 +16057,7 @@ def __init__(self):
                     "Config", (), {"pad_val": 1123581321.0, "tolerance": 1e-6}
                 )
 
-            @torch.compile(fullgraph=True)
+            @torch.compile(fullgraph=True, backend="eager")
             def forward(self, input):
                 outputs = torch.where(
                     torch.abs(input - self.config.pad_val) < self.config.tolerance,
@@ -14577,7 +16102,9 @@ def symbool_mul_fn(x_bool, sentinel):
         sentinel = torch.tensor(2.0, requires_grad=True, device="cuda")
         eager_result_true = symbool_mul_fn(x_true, sentinel)
         eager_result_false = symbool_mul_fn(x_false, sentinel)
-        compiled_fn = torch.compile(symbool_mul_fn, fullgraph=True, dynamic=True)
+        compiled_fn = torch.compile(
+            symbool_mul_fn, fullgraph=True, dynamic=True, backend="eager"
+        )
         compiled_result_true = compiled_fn(x_true, sentinel)
         compiled_result_false = compiled_fn(x_false, sentinel)
         self.assertEqual(eager_result_true, compiled_result_true)
@@ -14626,7 +16153,7 @@ def fuzzed_program(arg_0, sentinel):
         args = (arg_0,) + (sentinel,)
         try:
             compiled_program = torch.compile(
-                fuzzed_program, fullgraph=True, dynamic=True
+                fuzzed_program, fullgraph=True, dynamic=True, backend="eager"
             )
             compiled_program(*args)
         except Exception as e:
@@ -14654,7 +16181,9 @@ def _random_resize(image: torch.Tensor):
 
             return (img_h_new, img_w_new), image_metanet
 
-        _random_resize_compiled = torch.compile(fullgraph=True)(_random_resize)
+        _random_resize_compiled = torch.compile(fullgraph=True, backend="eager")(
+            _random_resize
+        )
 
         # Test the function
         input_tensor = torch.rand(1, 3, 224, 224)
diff --git a/test/dynamo/test_model_output.py b/test/dynamo/test_model_output.py
index d2833e1a7195a..08d9c11b09906 100644
--- a/test/dynamo/test_model_output.py
+++ b/test/dynamo/test_model_output.py
@@ -166,12 +166,12 @@ class MyDataClass(ModelOutput):
 
         def fn(obj):
             class_fields = dataclasses.fields(obj)
-            assert len(class_fields)
-            assert all(field.default is None for field in class_fields[1:])
+            assert len(class_fields)  # noqa: S101
+            assert all(field.default is None for field in class_fields[1:])  # noqa: S101
             other_fields_are_none = all(
                 getattr(obj, field.name) is None for field in class_fields[1:]
             )
-            assert not other_fields_are_none
+            assert not other_fields_are_none  # noqa: S101
 
             total = getattr(obj, class_fields[0].name)
             for field in class_fields[1:]:
diff --git a/test/dynamo/test_modes.py b/test/dynamo/test_modes.py
index 7cdb67358eedd..0a3e01b4998d5 100644
--- a/test/dynamo/test_modes.py
+++ b/test/dynamo/test_modes.py
@@ -1,6 +1,7 @@
 # Owner(s): ["module: dynamo"]
 
 import operator
+import unittest
 from unittest.mock import patch
 
 import torch
@@ -18,6 +19,9 @@
     BaseTorchFunctionMode,
     TorchFunctionMode,
 )
+from torch.testing._internal.common_device_type import (
+    IS_FLEX_ATTENTION_CUDA_PLATFORM_SUPPORTED,
+)
 from torch.testing._internal.common_utils import skipIfXpu
 from torch.testing._internal.inductor_utils import GPU_TYPE
 from torch.testing._internal.triton_utils import requires_gpu
@@ -156,10 +160,12 @@ def tearDownClass(cls):
         super().tearDownClass()
 
     def setUp(self):
+        super().setUp()
         torch.set_default_device(None)
         torch._dynamo.reset()
 
     def tearDown(self):
+        super().tearDown()
         torch.set_default_device(None)
         torch._dynamo.reset()
 
@@ -595,6 +601,33 @@ def fn(x, y):
 
         self.assertEqual(expected, actual)
 
+    def test_torch_function_mode_mutated_state(self):
+        class CounterMode(TorchFunctionMode):
+            def __init__(self):
+                self.count = 0
+
+            def __torch_function__(self, func, types, args=(), kwargs=None):
+                self.count += 1
+                return func(*args, **(kwargs or {}))
+
+        cnt = torch._dynamo.testing.CompileCounter()
+
+        @torch.compile(backend=cnt)
+        def f(x):
+            return x.sin()
+
+        with CounterMode() as mode:
+            f(torch.randn(3))
+            first_count = mode.count
+            first_frame_count = cnt.frame_count
+
+            # count changed, should recompile
+            f(torch.randn(3))
+
+        self.assertEqual(first_frame_count, 1)
+        self.assertGreater(first_count, 0)
+        self.assertEqual(cnt.frame_count, 2)
+
     # Needs larger cache size since we recompile for each op
     @patch.object(torch._dynamo.config, "recompile_limit", 48)
     def test_builtin_equivalent_funcs(self):
@@ -768,7 +801,7 @@ def test_hop(self):
 
             with self.assertRaisesRegex(
                 torch._dynamo.exc.Unsupported,
-                "raised exception HopDetectionError([ConstantVariable(str: 'test')])",
+                "raised exception HopDetectionError\\('test'\\)",
             ):
                 # This runs in fullgraph already
                 with TestModeRaises():
@@ -790,7 +823,7 @@ def test_hop_eager(self):
         with torch.device(GPU_TYPE):
             with self.assertRaisesRegex(
                 torch._dynamo.exc.Unsupported,
-                "raised exception HopDetectionError([ConstantVariable(str: 'test')])",
+                "raised exception HopDetectionError\\('test'\\)",
             ):
                 with TestModeRaises():
                     flex_attention_eager(
@@ -799,6 +832,55 @@ def test_hop_eager(self):
                         torch.ones(2, 2, 2, 2),
                     )
 
+    @requires_gpu
+    def test_default_device_factory_functions(self):
+        """Test that factory functions respect default device in compiled code"""
+
+        @torch.compile(fullgraph=True)
+        def random_func(
+            x: torch.Tensor,
+        ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+            # Test various factory functions
+            rnd = torch.randint(0, 2**32, size=x.shape, dtype=torch.uint32)
+            zeros = torch.zeros_like(rnd, device="cpu")
+            zeros_matched = torch.zeros_like(rnd)
+            return x + rnd, rnd, zeros, zeros_matched
+
+        torch.set_default_device(device_type)
+        (result, rnd, zeros, zeros_matched) = random_func(torch.randn(()))
+
+        # Verify tensors are on the current accelerator
+        self.assertEqual(rnd.device.type, device_type)
+        self.assertEqual(result.device.type, device_type)
+        self.assertEqual(zeros.device.type, "cpu")
+        self.assertEqual(zeros_matched.device.type, rnd.device.type)
+
+        torch.set_default_device("cpu")
+        (result, rnd, zeros, zeros_matched) = random_func(torch.randn(()))
+
+        # Verify tensors are on cpu
+        self.assertEqual(rnd.device.type, "cpu")
+        self.assertEqual(result.device.type, "cpu")
+        self.assertEqual(zeros.device.type, "cpu")
+        self.assertEqual(zeros_matched.device.type, rnd.device.type)
+
+        torch.set_default_device(None)
+
+    @requires_gpu
+    def test_default_device_factory_functions_priority(self):
+        torch.set_default_device(device_type)
+
+        @torch.compile(fullgraph=True)
+        def with_explicit_device(x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+            rnd = torch.randint(
+                0, 2**32, size=x.shape, dtype=torch.uint32, device="cpu"
+            )
+            return x + rnd, rnd
+
+        (result, rnd) = with_explicit_device(torch.randn(()))
+        self.assertEqual(rnd.device.type, "cpu")
+        self.assertEqual(result.device.type, device_type)
+
 
 class InvokeSubgraphBackendTests(torch._dynamo.test_case.TestCase):
     @torch._dynamo.config.patch(force_compile_during_fx_trace=True)
@@ -904,6 +986,7 @@ def forward(self, arg0_1: "f32[3, 3]"):
 """,  # noqa: B950
         )
 
+    @torch._functorch.config.patch(guess_tangent_strides_as_outputs=True)
     @torch._dynamo.config.patch(force_compile_during_fx_trace=True)
     def test_invoke_subgraph_seq_nr(self):
         """
@@ -968,7 +1051,8 @@ def bw_compiler(gm, example_inputs):
             [
                 ["add"],  # seq_nr 21
                 [
-                    "clone",
+                    "copy",
+                    "empty_strided",
                     "getitem",
                     "getitem_1",
                     "getitem_2",
@@ -1026,30 +1110,23 @@ def forward(self, arg0_1: "f32[3, 3]"):
 
         self.assertExpectedInline(
             normalize_gm(bw_graph.print_readable(print_output=False)),
-            """
+            """\
 class GraphModule(torch.nn.Module):
     def forward(self, getitem_1: "f32[3, 3]", tangents_1: "f32[]"):
-        # Annotation: {'seq_nr': 15} No stacktrace found for following nodes
         expand: "f32[3, 3]" = torch.ops.aten.expand.default(tangents_1, [3, 3]);  tangents_1 = None
-
-        # Annotation: {'seq_nr': 14} No stacktrace found for following nodes
-        clone: "f32[3, 3]" = torch.ops.aten.clone.default(expand, memory_format = torch.contiguous_format);  expand = None
+        empty_strided: "f32[3, 3]" = torch.ops.aten.empty_strided.default([3, 3], [3, 1], dtype = torch.float32, layout = torch.strided, device = device(type='cpu'), pin_memory = False)
+        copy: "f32[3, 3]" = torch.ops.aten.copy.default(empty_strided, expand);  empty_strided = expand = None
         repeated_subgraph1 = self.repeated_subgraph1
-        invoke_subgraph_1 = torch.ops.higher_order.invoke_subgraph(repeated_subgraph1, 'invoke_subgraph_1', getitem_1, clone);  repeated_subgraph1 = getitem_1 = clone = None
+        invoke_subgraph_1 = torch.ops.higher_order.invoke_subgraph(repeated_subgraph1, 'invoke_subgraph_1', getitem_1, copy);  repeated_subgraph1 = getitem_1 = copy = None
         getitem_2: "f32[3, 3]" = invoke_subgraph_1[0];  invoke_subgraph_1 = None
         return (getitem_2,)
-
     class repeated_subgraph1(torch.nn.Module):
         def forward(self, arg0_1: "f32[3, 3]", arg1_1: "f32[3, 3]"):
-            # Annotation: {'seq_nr': 10} File: test_modes.py:921 in inner_fn, code: return y / 2
             div: "f32[3, 3]" = torch.ops.aten.div.Tensor(arg1_1, 2);  arg1_1 = None
-
-            # Annotation: {'test': 'test', 'seq_nr': 9} File: test_modes.py:920 in inner_fn, code: y = x.cos()
             sin: "f32[3, 3]" = torch.ops.aten.sin.default(arg0_1);  arg0_1 = None
             neg: "f32[3, 3]" = torch.ops.aten.neg.default(sin);  sin = None
             mul: "f32[3, 3]" = torch.ops.aten.mul.Tensor(div, neg);  div = neg = None
-            return (mul,)
-        """,  # noqa: B950
+            return (mul,)""",  # noqa: B950
             ignore_comments=True,
             ignore_empty_lines=True,
         )
@@ -1322,6 +1399,751 @@ def fw_compiler(gm, example_inputs):
             f"Expected 1 compilation, got {compile_counter.frame_count}",
         )
 
+    @torch._dynamo.config.patch(force_compile_during_fx_trace=True)
+    def test_aot_autograd_over_dynamo_train_step(self):
+        """Test a full training step with nn.Module traced by AOTAutograd over Dynamo.
+
+        This tests a realistic training scenario where:
+        1. We have an actual nn.Module with parameters
+        2. make_fx traces the forward and backward as a single graph
+        3. Inside the module, a torch.compile'd function with invoke_subgraph is called
+        4. We use torch.autograd.grad to compute gradients (not .backward())
+        5. Gradients are returned explicitly from the traced function
+        """
+        from torch._dynamo.testing import CompileCounterWithBackend
+        from torch.fx.experimental.proxy_tensor import make_fx
+
+        torch._dynamo.reset()
+
+        # Use a compile counter to track how many times Dynamo compiles
+        compile_counter = CompileCounterWithBackend("invoke_subgraph")
+
+        # Define a simple MLP module
+        class SimpleMLP(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.fc1 = torch.nn.Linear(4, 8)
+                self.fc2 = torch.nn.Linear(8, 4)
+
+            def forward(self, x):
+                x = self.fc1(x)
+                x = torch.relu(x)
+                x = self.fc2(x)
+                return x
+
+        model = SimpleMLP()
+
+        # Inner compiled function that will be called inside the training step
+        def inner_fn(x):
+            return x * 2 + 1
+
+        compiled_inner = torch.compile(inner_fn, backend=compile_counter)
+
+        # Training step function that uses the model and compiled inner function
+        def train_step(params, x):
+            # params is a flat list of parameters: [fc1.weight, fc1.bias, fc2.weight, fc2.bias]
+            fc1_weight, fc1_bias, fc2_weight, fc2_bias = params
+
+            # Manual forward pass using functional style
+            h = torch.nn.functional.linear(x, fc1_weight, fc1_bias)
+            h = torch.relu(h)
+            # Apply compiled inner function
+            h = compiled_inner(h)
+            out = torch.nn.functional.linear(h, fc2_weight, fc2_bias)
+
+            # Compute loss
+            loss = out.sum()
+
+            # Compute gradients using torch.autograd.grad
+            grads = torch.autograd.grad(loss, params)
+
+            return (loss, *grads)
+
+        # Prepare inputs
+        x = torch.randn(2, 4)
+        params = (
+            model.fc1.weight.detach().clone().requires_grad_(True),
+            model.fc1.bias.detach().clone().requires_grad_(True),
+            model.fc2.weight.detach().clone().requires_grad_(True),
+            model.fc2.bias.detach().clone().requires_grad_(True),
+        )
+
+        # Trace using make_fx
+        fx_graph = make_fx(train_step, _disable_torch_fn_metadata_mode=True)(params, x)
+
+        # Run the traced graph
+        result = fx_graph(params, x)
+        loss = result[0]
+        grads = result[1:]
+
+        # Basic sanity checks
+        self.assertEqual(loss.shape, ())
+        self.assertEqual(len(grads), 4)
+        self.assertEqual(grads[0].shape, params[0].shape)  # fc1.weight grad
+        self.assertEqual(grads[1].shape, params[1].shape)  # fc1.bias grad
+        self.assertEqual(grads[2].shape, params[2].shape)  # fc2.weight grad
+        self.assertEqual(grads[3].shape, params[3].shape)  # fc2.bias grad
+
+        # Check that we got a graph with invoke_subgraph
+        graph_code = fx_graph.print_readable(print_output=False)
+        self.assertIn("invoke_subgraph", graph_code)
+
+        # Check compile count - should be exactly 1 compilation
+        self.assertEqual(
+            compile_counter.frame_count,
+            1,
+            f"Expected 1 compilation, got {compile_counter.frame_count}",
+        )
+
+    @requires_gpu
+    def test_nested_compile_dynamic(self):
+        """Test that wrap_compiled_regions works with dynamic shapes."""
+
+        d_model = 64
+
+        class MMLayer(torch.nn.Module):
+            def __init__(self, d_model: int):
+                super().__init__()
+                self.linear = torch.nn.Linear(d_model, d_model)
+
+            def forward(self, x):
+                return self.linear(x)
+
+        torch._dynamo.reset()
+
+        layer = MMLayer(d_model).to(GPU_TYPE)
+        compiled_mm = torch.compile(
+            layer,
+            backend="inductor",
+            options={"wrap_inductor_compiled_regions": True},
+            dynamic=True,
+        )
+
+        x = torch.randn(2, d_model, device=GPU_TYPE)
+        result = compiled_mm(x)
+        self.assertEqual(result.shape, (2, d_model))
+        torch.testing.assert_close(result, layer(x))
+
+        # Different batch size reuses the same compiled code
+        x2 = torch.randn(5, d_model, device=GPU_TYPE)
+        result2 = compiled_mm(x2)
+        self.assertEqual(result2.shape, (5, d_model))
+        torch.testing.assert_close(result2, layer(x2))
+
+    @requires_gpu
+    def test_nested_compile_input_mutation(self):
+        """Test nested compile with input mutation inside a compiled region.
+
+        Uses regional_inductor with fx_traceback.annotate to compile
+        a layer that mutates a buffer.
+        """
+        import contextlib
+
+        import torch.fx.traceback as fx_traceback
+        from functorch.compile import nop
+        from torch._dynamo.backends.common import aot_autograd
+        from torch._dynamo.backends.debugging import invoke_subgraph_inner_compiler
+        from torch._export.utils import _compiling_state_context
+        from torch._functorch.aot_autograd import (
+            aot_compile_joint_with_descriptors,
+            aot_export_joint_with_descriptors,
+        )
+        from torch._guards import tracing, TracingContext
+        from torch._subclasses import FakeTensorMode
+        from torch.fx.passes.regional_inductor import regional_inductor
+
+        def regional_inductor_invoke_subgraph(gm, args):
+            out_gm = regional_inductor(gm, args)
+            return invoke_subgraph_inner_compiler(out_gm, args)
+
+        def aot_eager_regional_inductor_invoke_subgraph():
+            return aot_autograd(
+                fw_compiler=regional_inductor_invoke_subgraph,
+                bw_compiler=regional_inductor_invoke_subgraph,
+                keep_inference_input_mutations=True,
+            )
+
+        d_model = 64
+
+        class MutatingLayer(torch.nn.Module):
+            def __init__(self, d_model: int):
+                super().__init__()
+                self.linear = torch.nn.Linear(d_model, d_model)
+                self.register_buffer("call_count", torch.tensor(0, dtype=torch.int64))
+
+            def forward(self, x):
+                with fx_traceback.annotate({"compile_with_inductor": 0}):
+                    self.call_count.add_(1)
+                return self.linear(x)
+
+        class StackedMutating(torch.nn.Module):
+            def __init__(self, d_model: int, n_layers: int = 2):
+                super().__init__()
+                self.layers = torch.nn.ModuleList(
+                    [
+                        torch.compile(
+                            MutatingLayer(d_model),
+                            backend=aot_eager_regional_inductor_invoke_subgraph(),
+                        )
+                        for _ in range(n_layers)
+                    ]
+                )
+
+            def forward(self, x):
+                for layer in self.layers:
+                    x = layer(x)
+                return x
+
+        with (
+            torch._dynamo.config.patch(force_compile_during_fx_trace=True),
+            torch._inductor.config.patch(wrap_inductor_compiled_regions=True),
+            torch._functorch.config.patch(force_non_lazy_backward_lowering=True),
+        ):
+            torch._dynamo.reset()
+
+            model = StackedMutating(d_model, n_layers=2).to(GPU_TYPE)
+
+            x = torch.randn(2, d_model, device=GPU_TYPE, requires_grad=True)
+
+            fake_mode = FakeTensorMode(allow_non_fake_inputs=True)
+            saved_params = list(model.parameters())
+            saved_buffers = list(model.buffers())
+
+            with contextlib.ExitStack() as stack:
+                stack.enter_context(tracing(TracingContext(fake_mode)))
+                stack.enter_context(_compiling_state_context())
+                stack.enter_context(fake_mode)
+
+                joint_with_descriptors = aot_export_joint_with_descriptors(
+                    stack,
+                    model,
+                    args=(x,),
+                    kwargs={},
+                    keep_inference_input_mutations=True,
+                    _disable_torch_fn_metadata_mode=True,
+                )
+                gm = joint_with_descriptors.graph_module
+                print("=== input_mutation outer graph ===")
+                print(gm.print_readable(print_output=False))
+                compiled_fn = aot_compile_joint_with_descriptors(
+                    joint_with_descriptors,
+                    fw_compiler=nop,
+                    bw_compiler=nop,
+                )
+
+            def wrapped_fn(*args, **kwargs):
+                return compiled_fn(
+                    *saved_params,
+                    *saved_buffers,
+                    *args,
+                    **kwargs,
+                )
+
+            out = wrapped_fn(x)
+            out.sum().backward()
+
+    @requires_gpu
+    def test_nested_compile_output_aliases_input(self):
+        """Test nested compile where output is a view-alias of input.
+
+        Uses regional_inductor with fx_traceback.annotate to compile
+        a layer that returns a view of its input.
+        """
+        import contextlib
+
+        import torch.fx.traceback as fx_traceback
+        from functorch.compile import nop
+        from torch._dynamo.backends.common import aot_autograd
+        from torch._dynamo.backends.debugging import invoke_subgraph_inner_compiler
+        from torch._export.utils import _compiling_state_context
+        from torch._functorch.aot_autograd import (
+            aot_compile_joint_with_descriptors,
+            aot_export_joint_with_descriptors,
+        )
+        from torch._guards import tracing, TracingContext
+        from torch._subclasses import FakeTensorMode
+        from torch.fx.passes.regional_inductor import regional_inductor
+
+        def regional_inductor_invoke_subgraph(gm, args):
+            out_gm = regional_inductor(gm, args)
+            return invoke_subgraph_inner_compiler(out_gm, args)
+
+        def aot_eager_regional_inductor_invoke_subgraph():
+            return aot_autograd(
+                fw_compiler=regional_inductor_invoke_subgraph,
+                bw_compiler=regional_inductor_invoke_subgraph,
+                keep_inference_input_mutations=True,
+            )
+
+        d_model = 64
+
+        class ViewLayer(torch.nn.Module):
+            def __init__(self, d_model: int):
+                super().__init__()
+                self.linear = torch.nn.Linear(d_model, d_model)
+
+            def forward(self, x):
+                y = self.linear(x)
+                with fx_traceback.annotate({"compile_with_inductor": 0}):
+                    # View-alias: unsqueeze then squeeze back
+                    return y.unsqueeze(1).squeeze(1)
+
+        class StackedView(torch.nn.Module):
+            def __init__(self, d_model: int, n_layers: int = 2):
+                super().__init__()
+                self.layers = torch.nn.ModuleList(
+                    [
+                        torch.compile(
+                            ViewLayer(d_model),
+                            backend=aot_eager_regional_inductor_invoke_subgraph(),
+                        )
+                        for _ in range(n_layers)
+                    ]
+                )
+
+            def forward(self, x):
+                for layer in self.layers:
+                    x = layer(x)
+                return x
+
+        with (
+            torch._dynamo.config.patch(force_compile_during_fx_trace=True),
+            torch._inductor.config.patch(wrap_inductor_compiled_regions=True),
+            torch._functorch.config.patch(force_non_lazy_backward_lowering=True),
+        ):
+            torch._dynamo.reset()
+
+            model = StackedView(d_model, n_layers=2).to(GPU_TYPE)
+
+            x = torch.randn(2, d_model, device=GPU_TYPE, requires_grad=True)
+
+            fake_mode = FakeTensorMode(allow_non_fake_inputs=True)
+            saved_params = list(model.parameters())
+            saved_buffers = list(model.buffers())
+
+            with contextlib.ExitStack() as stack:
+                stack.enter_context(tracing(TracingContext(fake_mode)))
+                stack.enter_context(_compiling_state_context())
+                stack.enter_context(fake_mode)
+
+                joint_with_descriptors = aot_export_joint_with_descriptors(
+                    stack,
+                    model,
+                    args=(x,),
+                    kwargs={},
+                    keep_inference_input_mutations=True,
+                    _disable_torch_fn_metadata_mode=True,
+                )
+                gm = joint_with_descriptors.graph_module
+                print("=== output_aliases_input outer graph ===")
+                print(gm.print_readable(print_output=False))
+                compiled_fn = aot_compile_joint_with_descriptors(
+                    joint_with_descriptors,
+                    fw_compiler=nop,
+                    bw_compiler=nop,
+                )
+
+            def wrapped_fn(*args, **kwargs):
+                return compiled_fn(
+                    *saved_params,
+                    *saved_buffers,
+                    *args,
+                    **kwargs,
+                )
+
+            out = wrapped_fn(x)
+            out.sum().backward()
+
+    @unittest.skipUnless(
+        IS_FLEX_ATTENTION_CUDA_PLATFORM_SUPPORTED and not torch.version.hip,
+        "Requires CUDA with SM >= 8.0, Triton, and not ROCm",
+    )
+    def test_nested_compile_transformer_with_flex_attention_compiled_layers(
+        self,
+    ):
+        """Test a transformer model with 4 identical compiled layers using flex_attention.
+
+        This test:
+        1. Creates a small transformer model with 4 identical layers using flex_attention
+        2. Compiles each layer with the invoke_subgraph + regional_inductor backend
+        3. Passes the entire model through aot_export_joint_with_descriptors()
+        4. Verifies the final graph has invoke_subgraph HOPs
+        """
+        import torch.distributed as dist
+        import torch.fx.traceback as fx_traceback
+        from torch._dynamo.backends.common import aot_autograd
+        from torch._dynamo.backends.debugging import invoke_subgraph_inner_compiler
+        from torch._guards import tracing, TracingContext
+        from torch.distributed.device_mesh import init_device_mesh
+        from torch.distributed.tensor import (
+            distribute_module,
+            distribute_tensor,
+            DTensor,
+            Replicate,
+        )
+        from torch.fx.passes.regional_inductor import regional_inductor
+        from torch.nn.attention.flex_attention import create_block_mask, flex_attention
+        from torch.testing._internal.distributed.fake_pg import FakeStore
+
+        def regional_inductor_invoke_subgraph(gm, args):
+            out_gm = regional_inductor(gm, args)
+            return invoke_subgraph_inner_compiler(out_gm, args)
+
+        def aot_eager_regional_inductor_invoke_subgraph():
+            return aot_autograd(
+                fw_compiler=regional_inductor_invoke_subgraph,
+                bw_compiler=regional_inductor_invoke_subgraph,
+                # Keep input mutations in the graph
+                keep_inference_input_mutations=True,
+            )
+
+        d_model, n_heads, d_ff = 64, 4, 128
+        batch_size, seq_len = 2, 32
+
+        # dumb mask mod that closes over a tensor
+        mask_bias = torch.tensor(0, device=GPU_TYPE, dtype=torch.int32)
+
+        def mask_mod(b_idx, h_idx, q_idx, k_idx):
+            return (q_idx >= k_idx) | (mask_bias == 1)
+
+        # Create block mask from the mask_mod
+        block_mask = create_block_mask(
+            mask_mod=mask_mod,
+            B=batch_size,
+            H=None,  # Broadcast over heads
+            Q_LEN=seq_len,
+            KV_LEN=seq_len,
+            device=GPU_TYPE,
+        )
+
+        # Transformer layer with flex_attention
+        class TransformerLayer(torch.nn.Module):
+            def __init__(self, d_model: int, n_heads: int, d_ff: int):
+                super().__init__()
+                self.n_heads = n_heads
+                self.head_dim = d_model // n_heads
+                self.d_model = d_model
+
+                # Attention projections
+                self.q_proj = torch.nn.Linear(d_model, d_model)
+                self.k_proj = torch.nn.Linear(d_model, d_model)
+                self.v_proj = torch.nn.Linear(d_model, d_model)
+                self.out_proj = torch.nn.Linear(d_model, d_model)
+
+                # Layer norms
+                self.norm1 = torch.nn.LayerNorm(d_model)
+                self.norm2 = torch.nn.LayerNorm(d_model)
+
+                # FFN
+                self.ff = torch.nn.Sequential(
+                    torch.nn.Linear(d_model, d_ff),
+                    torch.nn.GELU(),
+                    torch.nn.Linear(d_ff, d_model),
+                )
+
+                # give the model a buffer mutation so we can test auto_functionalize too
+                self.register_buffer("_call_count", torch.tensor(0, dtype=torch.int64))
+
+            def forward(self, x):
+                # Mutate the buffer to trigger input mutation detection
+                self._call_count.add_(1)
+
+                batch_size, seq_len, _ = x.shape
+
+                # Pre-norm for attention
+                normed = self.norm1(x)
+
+                # Compute Q, K, V
+                q = self.q_proj(normed)
+                k = self.k_proj(normed)
+                v = self.v_proj(normed)
+
+                # Reshape for multi-head attention: (B, S, D) -> (B, H, S, D_head)
+                q = q.view(batch_size, seq_len, self.n_heads, self.head_dim).transpose(
+                    1, 2
+                )
+                k = k.view(batch_size, seq_len, self.n_heads, self.head_dim).transpose(
+                    1, 2
+                )
+                v = v.view(batch_size, seq_len, self.n_heads, self.head_dim).transpose(
+                    1, 2
+                )
+
+                with fx_traceback.annotate({"compile_with_inductor": 0}):
+                    attn_out_local = flex_attention(
+                        q.to_local(),
+                        k.to_local(),
+                        v.to_local(),
+                        block_mask=block_mask,
+                    )
+                    attn_out = DTensor.from_local(
+                        attn_out_local, device_mesh, [Replicate()]
+                    )
+
+                attn_out = (
+                    attn_out.transpose(1, 2).contiguous().view(batch_size, seq_len, -1)
+                )
+                attn_out = self.out_proj(attn_out)
+
+                x = x + attn_out
+
+                x = x + self.ff(self.norm2(x))
+
+                return x
+
+        # Transformer model with 4 identical layers
+        class SmallTransformer(torch.nn.Module):
+            def __init__(
+                self, d_model: int, n_heads: int, d_ff: int, n_layers: int = 4
+            ):
+                super().__init__()
+                self.layers = torch.nn.ModuleList(
+                    [
+                        torch.compile(
+                            TransformerLayer(d_model, n_heads, d_ff),
+                            backend=aot_eager_regional_inductor_invoke_subgraph(),
+                        )
+                        for _ in range(n_layers)
+                    ]
+                )
+
+            def forward(self, x):
+                for layer in self.layers:
+                    x = layer(x)
+                return x
+
+        fake_store = FakeStore()
+        dist.init_process_group("fake", store=fake_store, rank=0, world_size=2)
+        device_mesh = init_device_mesh(GPU_TYPE, (2,))
+
+        with (
+            # Needed when wrapping a compiled region with FX tracing
+            torch._dynamo.config.patch(force_compile_during_fx_trace=True),
+            # Needed because our inner compiled region uses inductor (for flex)
+            torch._inductor.config.patch(wrap_inductor_compiled_regions=True),
+            # AOTAutograd normally tries to "delay backward compilation to bw runtime",
+            # but for nested compile we actually need it to happen when we trace the fw.
+            torch._functorch.config.patch(force_non_lazy_backward_lowering=True),
+        ):
+            torch._dynamo.reset()
+
+            model = SmallTransformer(d_model, n_heads, d_ff, n_layers=4)
+            model = model.to(GPU_TYPE)
+
+            def replicate_all(name, module, device_mesh):
+                for param_name, param in module.named_parameters(recurse=False):
+                    dist_param = torch.nn.Parameter(
+                        distribute_tensor(param, device_mesh, [Replicate()])
+                    )
+                    module.register_parameter(param_name, dist_param)
+
+            model = distribute_module(model, device_mesh, replicate_all)
+
+            x = torch.randn(
+                batch_size,
+                seq_len,
+                d_model,
+                device=GPU_TYPE,
+                dtype=torch.float32,
+                requires_grad=True,
+            )
+            x = DTensor.from_local(x, device_mesh, [Replicate()])
+
+            import contextlib
+
+            from functorch.compile import nop
+            from torch._export.utils import _compiling_state_context
+            from torch._functorch.aot_autograd import (
+                aot_compile_joint_with_descriptors,
+                aot_export_joint_with_descriptors,
+            )
+            from torch._functorch.partitioners import default_partition
+            from torch._guards import detect_fake_mode
+            from torch._subclasses import FakeTensorMode
+
+            fake_mode = FakeTensorMode(allow_non_fake_inputs=True)
+
+            saved_params = list(model.parameters())
+            saved_buffers = list(model.buffers())
+
+            with contextlib.ExitStack() as stack:
+                stack.enter_context(tracing(TracingContext(fake_mode)))
+                stack.enter_context(_compiling_state_context())
+                if fake_mode is not None:
+                    stack.enter_context(fake_mode)
+
+                joint_with_descriptors = aot_export_joint_with_descriptors(
+                    stack,
+                    model,
+                    args=(x,),
+                    kwargs={},
+                    keep_inference_input_mutations=True,
+                    # see https://github.com/pytorch/pytorch/pull/172087
+                    _disable_torch_fn_metadata_mode=True,
+                )
+                gm = joint_with_descriptors.graph_module
+                fake_mode = detect_fake_mode()
+                num_fwd_outputs = (
+                    joint_with_descriptors._aot_state.fw_metadata.num_forward
+                )
+                fwd_gm, bwd_gm = default_partition(
+                    gm, None, num_fwd_outputs=num_fwd_outputs
+                )
+                compiled_fn = aot_compile_joint_with_descriptors(
+                    joint_with_descriptors,
+                    fw_compiler=nop,
+                    bw_compiler=nop,
+                )
+
+            def wrapped_fn(*args, **kwargs):
+                return compiled_fn(
+                    *saved_params,
+                    *saved_buffers,
+                    *args,
+                    **kwargs,
+                )
+
+            out = wrapped_fn(x)
+            out.sum().backward()
+
+        dist.destroy_process_group()
+
+    @unittest.skipUnless(
+        IS_FLEX_ATTENTION_CUDA_PLATFORM_SUPPORTED and not torch.version.hip,
+        "Requires CUDA with SM >= 8.0, Triton, and not ROCm",
+    )
+    def test_2tier_blockmask_tensor_closure_nested_compile_aot_export(self):
+        """2-tier AOT export with BlockMask whose mask_mod captures tensors.
+
+        Reproduces the sixlib/mango pattern where:
+        - BlockMask is pytree-registered (as in sixlib/attention_mask.py)
+        - BlockMask with tensor closure is created inside forward
+        - aot_export flattens intermediates via pytree, exposing the
+          mask_mod callable (with captured tensors) in the context
+
+        visibility is a model input -> FunctionalTensor during outer trace.
+        mask_mod closure captures derived slices (lower, upper).
+        """
+        import contextlib
+
+        import torch.fx.traceback as fx_traceback
+        from torch._functorch.aot_autograd import aot_export_joint_with_descriptors
+        from torch._subclasses import FakeTensorMode
+        from torch.nn.attention.flex_attention import (
+            BlockMask,
+            create_block_mask,
+            flex_attention,
+        )
+        from torch.utils._pytree import register_pytree_node, SUPPORTED_NODES
+
+        # Register BlockMask as pytree node (same as sixlib/attention_mask.py)
+        if BlockMask not in SUPPORTED_NODES:
+            register_pytree_node(
+                BlockMask,
+                BlockMask._flatten,
+                BlockMask._unflatten,
+                flatten_with_keys_fn=BlockMask._flatten_with_keys,
+                serialized_type_name="torch.nn.attention.flex_attention.BlockMask",
+            )
+
+        d_model, n_heads = 64, 4
+        batch_size, seq_len = 2, 32
+
+        class SimpleModel(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.q_proj = torch.nn.Linear(d_model, d_model)
+                self.k_proj = torch.nn.Linear(d_model, d_model)
+                self.v_proj = torch.nn.Linear(d_model, d_model)
+
+            def forward(self, x, visibility):
+                bs, sl, _ = x.shape
+                lower = visibility[:, 0, :]
+                upper = visibility[:, 1, :]
+
+                def mask_mod(b_idx, h_idx, q_idx, k_idx):
+                    return (lower[b_idx, q_idx] <= k_idx) & (
+                        k_idx <= upper[b_idx, q_idx]
+                    )
+
+                bm = create_block_mask(
+                    mask_mod,
+                    B=bs,
+                    H=None,
+                    Q_LEN=sl,
+                    KV_LEN=sl,
+                    device=x.device,
+                )
+
+                q = (
+                    self.q_proj(x)
+                    .view(bs, sl, n_heads, d_model // n_heads)
+                    .transpose(1, 2)
+                )
+                k = (
+                    self.k_proj(x)
+                    .view(bs, sl, n_heads, d_model // n_heads)
+                    .transpose(1, 2)
+                )
+                v = (
+                    self.v_proj(x)
+                    .view(bs, sl, n_heads, d_model // n_heads)
+                    .transpose(1, 2)
+                )
+
+                with fx_traceback.annotate({"compile_with_inductor": 1}):
+                    attn_out = flex_attention(q, k, v, block_mask=bm)
+
+                out = attn_out.transpose(1, 2).contiguous().view(bs, sl, -1)
+                # Return BlockMask alongside output so it appears in the
+                # output pytree spec — this matches how sixlib's decoder
+                # returns context alongside outputs.
+                return out, bm
+
+        with (
+            torch._dynamo.config.patch(force_compile_during_fx_trace=True),
+            torch._inductor.config.patch(wrap_inductor_compiled_regions=True),
+            torch._functorch.config.patch(force_non_lazy_backward_lowering=True),
+        ):
+            torch._dynamo.reset()
+
+            model = SimpleModel().to(GPU_TYPE)
+            x = torch.randn(
+                batch_size,
+                seq_len,
+                d_model,
+                device=GPU_TYPE,
+                dtype=torch.float32,
+                requires_grad=True,
+            )
+            visibility = torch.zeros(
+                batch_size,
+                2,
+                seq_len,
+                device=GPU_TYPE,
+                dtype=torch.int64,
+            )
+            visibility[:, 1, :] = seq_len - 1
+
+            fake_mode = FakeTensorMode(allow_non_fake_inputs=True)
+            with contextlib.ExitStack() as stack:
+                stack.enter_context(fake_mode)
+                result = aot_export_joint_with_descriptors(
+                    stack,
+                    model,
+                    args=(x, visibility),
+                    kwargs={},
+                    keep_inference_input_mutations=True,
+                    _disable_torch_fn_metadata_mode=True,
+                )
+
+            print("=== Outer (aot_export) graph ===")
+            print(result.graph_module.graph)
+            for name, submod in result.graph_module.named_modules():
+                if name and hasattr(submod, "graph"):
+                    print(f"\n=== Submodule: {name} ===")
+                    print(submod.graph)
+
 
 class TorchFunctionModeLifecycleTests(torch._dynamo.test_case.TestCase):
     def test_default_device_restored_after_mode_tests(self):
diff --git a/test/dynamo/test_modules.py b/test/dynamo/test_modules.py
index 43ee0588a3f47..13e8b5ef61044 100644
--- a/test/dynamo/test_modules.py
+++ b/test/dynamo/test_modules.py
@@ -2140,35 +2140,8 @@ def forward(self, x):
         ):
             x = torch.randn(*size, requires_grad=True)
             mod(x)
-            if torch._dynamo.config.inline_inbuilt_nn_modules:
-                self.assertEqual(cnts.frame_count, 1)
-            else:
-                self.assertEqual(cnts.frame_count, num_submodules)
-
-    @patch.object(torch._dynamo.config, "accumulated_recompile_limit", 2)
-    @patch.object(torch._dynamo.config, "inline_inbuilt_nn_modules", False)
-    def test_recompile_limit_on_freed_module(self):
-        class Mod(torch.nn.Module):
-            def __init__(self) -> None:
-                super().__init__()
-                self.lin = torch.nn.Linear(5, 5)
-
-            def forward(self, x):
-                return self.lin(x)
-
-        def fn(x, mod):
-            return mod(x)
-
-        cnts = torch._dynamo.testing.CompileCounterWithBackend("eager")
-        opt_mod = torch.compile(fn, backend=cnts)
-        for _ in range(8):
-            mod = Mod()
-            opt_mod(torch.randn(5, 5), mod)
-
-        # fn compiles twice
-        self.assertEqual(cnts.frame_count, 2)
+            self.assertEqual(cnts.frame_count, 1)
 
-    @patch.object(torch._dynamo.config, "inline_inbuilt_nn_modules", True)
     def test_inline_inbuilt_nn_modules(self):
         size = (10, 10)
         recompile_limit = 1
@@ -2249,10 +2222,7 @@ def forward(self, x):
             ]:
                 x = torch.randn(size)
                 mod(x)
-        if torch._dynamo.config.inline_inbuilt_nn_modules:
-            self.assertEqual(cnts.frame_count, 2)
-        else:
-            self.assertEqual(cnts.frame_count, 2 * num_submodules)
+        self.assertEqual(cnts.frame_count, 2)
 
     def test_recursion(self):
         mod = MockModule()
@@ -2449,56 +2419,6 @@ def new_forward_hook(
         self.assertEqual(compiled_func(inp), outer_func(inp))
         self.assertEqual(compiled_func(inp).item(), 16)
 
-    @patch.object(torch._dynamo.config, "guard_nn_modules", False)
-    @patch.object(torch._dynamo.config, "skip_nnmodule_hook_guards", True)
-    @patch.object(torch._dynamo.config, "inline_inbuilt_nn_modules", False)
-    def test_hooks_skip_guards(self):
-        class TestModule(torch.nn.Module):
-            def forward(self, x: torch.Tensor) -> torch.Tensor:
-                return 2 * x + 1
-
-        m = TestModule()
-
-        def forward_hook(
-            module: torch.nn.Module, inputs: tuple[torch.Tensor], output: torch.Tensor
-        ) -> torch.Tensor:
-            return 2 * output + 1
-
-        handle = m.register_forward_hook(forward_hook)
-
-        def outer_func(tensor):
-            x = tensor * 2 + 1
-            y = m(x)
-            return y
-
-        inp = torch.tensor(1.0, requires_grad=True)
-
-        failure_reason = None
-
-        def guard_fail_fn(failure):
-            nonlocal failure_reason
-            failure_reason = failure[0]
-
-        cc = torch._dynamo.testing.CompileCounterWithBackend("aot_eager")
-        compiled_func = torch._dynamo.optimize(
-            guard_fail_fn=guard_fail_fn,
-            backend=cc,
-        )(outer_func)
-
-        m = TestModule()
-        handle = m.register_forward_hook(forward_hook)
-        failure_reason = None
-        self.assertEqual(compiled_func(inp), outer_func(inp))
-        self.assertEqual(compiled_func(inp).item(), 15)
-        self.assertEqual(cc.frame_count, 1)
-        self.assertEqual(cc.op_count, 6)
-
-        # if we remove the hook, dynamo shouldn't notice
-        handle.remove()
-        self.assertNotEqual(compiled_func(inp), outer_func(inp))
-        self.assertEqual(compiled_func(inp).item(), 15)
-        self.assertEqual(cc.frame_count, 1)
-
     def _forward_hook_test_helper(self, model):
         forward_handles = {}
         compiled_activations = {}
@@ -2738,7 +2658,7 @@ def forward(self, x):
         model = ToyModel()
         x = torch.zeros((3, 4))
         obj = CustomClass(model)
-        out = torch.compile(obj, fullgraph=True)(x)
+        out = torch.compile(obj, fullgraph=True, backend="eager")(x)
         self.assertEqual(out, (x + 1) * (x + 1))
 
     def test_module_dict_iter_name(self):
@@ -2810,21 +2730,10 @@ def run():
             x = torch.randn(1, 3)
             return models(x)
 
-        run = torch.compile(run, fullgraph=True)
+        run = torch.compile(run, fullgraph=True, backend="eager")
         run()
         self.assertTrue(models[0].abc)
 
-    @torch._dynamo.config.patch(inline_inbuilt_nn_modules=False)
-    def test_assign_does_not_exist(self):
-        class MyModule(torch.nn.Module):
-            def forward(self, x):
-                self.text_encoding = x + 1
-                return self.text_encoding
-
-        mod = MyModule()
-        out = torch.compile(mod, fullgraph=True)(torch.randn(10))
-        assert mod.text_encoding is out
-
     def test_module_dict_iter_values(self):
         class MyModule(torch.nn.Module):
             def __init__(self) -> None:
@@ -2885,16 +2794,10 @@ def foo(mod, x):
 
         mod = Mod()
         foo(mod, torch.rand([4]))
-        if torch._dynamo.config.inline_inbuilt_nn_modules:
-            self.assertEqual(compiles_without_buffers, 1)
-        else:
-            self.assertEqual(compiles_without_buffers, 0)
+        self.assertEqual(compiles_without_buffers, 1)
 
         foo(mod, torch.rand([4], dtype=torch.half))
-        if torch._dynamo.config.inline_inbuilt_nn_modules:
-            self.assertEqual(compiles_without_buffers, 2)
-        else:
-            self.assertEqual(compiles_without_buffers, 1)
+        self.assertEqual(compiles_without_buffers, 2)
 
         class Mod2(Mod):
             def __setattr__(self, name, value):
@@ -2912,7 +2815,6 @@ def test_unspec_non_inlinable_module(self):
         expected = mod(x)
         self.assertEqual(actual, expected)
 
-    @torch._dynamo.config.patch("inline_inbuilt_nn_modules", True)
     def test_mark_static_previously_seen_tensor(self):
         # This test verifies that dynamo will mark
         # the buffers/params of a module as static
@@ -2958,7 +2860,6 @@ def fn(x, b, mod):
         fn(inp, buf, mod)
         self.assertEqual(num_compiles, 1)
 
-    @torch._dynamo.config.patch("inline_inbuilt_nn_modules", True)
     def test_mark_static_nn_module_tensor(self):
         # This test verifies that dynamo will mark
         # the nn module tensor attributes as static
@@ -3001,7 +2902,6 @@ def fn(x):
         fn(inp)
         self.assertEqual(num_compiles, 1)
 
-    @torch._dynamo.config.patch("inline_inbuilt_nn_modules", True)
     @torch._inductor.config.patch("freezing", True)
     @torch.no_grad()
     def test_mark_static_with_freezing(self):
@@ -3109,7 +3009,7 @@ def forward(self, x):
                 return x + 1
 
         model = Mod()
-        compiled_model = torch.compile(model)
+        compiled_model = torch.compile(model, backend="eager")
         compiled_model.foo = 42
         del compiled_model.foo
 
@@ -3241,8 +3141,26 @@ def fn(x):
             self.assertEqual(cnt.frame_count, 2)
 
         helper()
-        with torch._dynamo.config.patch(inline_inbuilt_nn_modules=True):
-            helper()
+
+    def test_monkeypatching_forward_inside_compiled_region(self):
+        class Mod(torch.nn.Module):
+            def forward(self, x):
+                return x - 1
+
+        @torch.compile(backend="eager", fullgraph=True)
+        def fn(mod, x, y):
+            def patch(x):
+                return x + y
+
+            mod.forward = patch
+            return mod(x)
+
+        inp0 = torch.ones(3)
+        inp1 = torch.ones(3)
+        mod = Mod()
+
+        self.assertEqual(fn(mod, inp0, inp1), inp0 + inp1)
+        self.assertEqual(mod(inp0), inp0 + inp1)
 
     def test_user_defined_nn_module_dynamic(self):
         class Conv2d(torch.nn.Conv2d):
@@ -3278,7 +3196,6 @@ def forward(self, x):
         # Must be 3 compilations. If not marked static there would be 2, because strides would be converted to symints.
         self.assertEqual(cnts.frame_count, 3)
 
-    @patch.object(torch._dynamo.config, "inline_inbuilt_nn_modules", True)
     def test_overridden_call(self):
         class OverRiddenCallModule(torch.nn.Module):
             def __init__(self):
@@ -3303,8 +3220,6 @@ def fn(x):
         res = opt_fn(x)
         self.assertEqual(ref, res)
 
-    @torch._dynamo.config.patch("skip_tensor_guards_with_matching_dict_tags", False)
-    @torch._dynamo.config.patch("inline_inbuilt_nn_modules", True)
     def test_param_requires_grad(self):
         def adjust_model(model):
             to_freeze = model.num_iter % 2 == 0
@@ -3350,6 +3265,94 @@ def forward(self, x):
         # Check that we have recompiled twice, which leads to 3 frames
         self.assertEqual(cnt.frame_count, 3)
 
+    @torch._dynamo.config.patch("use_recursive_dict_tags_for_guards", False)
+    def test_param_requires_grad_no_recursive_dict_tags(self):
+        class MyModule(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.linear = torch.nn.Linear(4, 4)
+
+            def forward(self, x):
+                return self.linear(x)
+
+        model = MyModule()
+        model.linear.weight.requires_grad_(False)
+
+        cnt = torch._dynamo.testing.CompileCounter()
+        opt_model = torch.compile(model, backend=cnt, fullgraph=True)
+
+        x = torch.randn(4, 4)
+        opt_model(x).sum().backward()
+        self.assertEqual(cnt.frame_count, 1)
+        self.assertIsNone(model.linear.weight.grad)
+
+        model.linear.weight.requires_grad_(True)
+        opt_model(x).sum().backward()
+        self.assertEqual(cnt.frame_count, 2)
+        self.assertIsNotNone(model.linear.weight.grad)
+
+    @torch._dynamo.config.patch(skip_tensor_guards_with_matching_dict_tags=True)
+    @torch._dynamo.config.patch("use_recursive_dict_tags_for_guards", True)
+    def test_param_dtype_change_recompiles_with_recursive_dict_tags(self):
+        class MyModule(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.scale = torch.nn.Parameter(torch.randn(4))
+
+            def forward(self, x):
+                return x * self.scale
+
+        model = MyModule()
+        x = torch.randn(4)
+
+        cnt = torch._dynamo.testing.CompileCounter()
+        compiled = torch.compile(model, backend=cnt, fullgraph=True)
+
+        self.assertTrue(torch._dynamo.testing.same(model(x), compiled(x)))
+        self.assertEqual(cnt.frame_count, 1)
+
+        model.to(dtype=torch.float64)
+
+        recompiled = torch.compile(model, backend=cnt, fullgraph=True)
+        result = recompiled(x)
+
+        self.assertEqual(result.dtype, torch.float64)
+        self.assertTrue(torch._dynamo.testing.same(model(x), result))
+        self.assertEqual(cnt.frame_count, 2)
+
+    def test_param_requires_grad_submodule(self):
+        class Inner(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.conv = torch.nn.Conv2d(3, 16, 3, 1)
+
+            def forward(self, x):
+                return self.conv(x)
+
+        class Outer(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.inner = Inner()
+
+            def forward(self, x):
+                return self.inner(x)
+
+        model = Outer()
+        model.inner.conv.weight.requires_grad_(False)
+
+        cnt = torch._dynamo.testing.CompileCounter()
+        opt_model = torch.compile(model, backend=cnt, fullgraph=True)
+
+        x = torch.randn(1, 3, 10, 10)
+        opt_model(x).sum().backward()
+        self.assertEqual(cnt.frame_count, 1)
+        self.assertIsNone(model.inner.conv.weight.grad)
+
+        model.inner.conv.weight.requires_grad_(True)
+        opt_model(x).sum().backward()
+        self.assertEqual(cnt.frame_count, 2)
+        self.assertIsNotNone(model.inner.conv.weight.grad)
+
     def test_branch_on_nn_module_custom_len(self):
         class Cache(torch.nn.Module):
             def __init__(self):
@@ -3495,18 +3498,26 @@ def forward(self, x):
         input_tensor = torch.randn(4, 10)
 
         eager_output = model(input_tensor)
-        assert hasattr(model.linear, "weight")
-        assert not hasattr(model.linear, "_tmp_weight")
+        if not hasattr(model.linear, "weight"):
+            raise AssertionError("Expected model.linear to have weight")
+        if hasattr(model.linear, "_tmp_weight"):
+            raise AssertionError("Expected model.linear to not have _tmp_weight")
 
         torch.manual_seed(0)
         model_to_compile = SimpleModel()
         model_to_compile.linear.register_forward_pre_hook(pre_forward_rename_hook)
         model_to_compile.linear.register_forward_hook(post_forward_restore_hook)
 
-        compiled_model = torch.compile(model_to_compile, fullgraph=True)
+        compiled_model = torch.compile(
+            model_to_compile, fullgraph=True, backend="eager"
+        )
         compiled_output = compiled_model(input_tensor)
-        assert hasattr(model.linear, "weight")
-        assert not hasattr(compiled_model.linear, "_tmp_weight")
+        if not hasattr(model.linear, "weight"):
+            raise AssertionError("Expected model.linear to have weight")
+        if hasattr(compiled_model.linear, "_tmp_weight"):
+            raise AssertionError(
+                "Expected compiled_model.linear to not have _tmp_weight"
+            )
         torch.testing.assert_close(eager_output, compiled_output)
 
     def test_submodule_forward_hooks_with_kwargs(self):
@@ -3545,7 +3556,9 @@ def noop_hook(module, args, kwargs, output):
             layer.linear.register_forward_hook(noop_hook, with_kwargs=True)
 
         for i, layer in model.layers.named_children():
-            model.layers.register_module(i, torch.compile(layer, fullgraph=True))
+            model.layers.register_module(
+                i, torch.compile(layer, fullgraph=True, backend="eager")
+            )
 
         output = model(inp)
         self.assertEqual(output_eager, output)
@@ -3600,7 +3613,9 @@ def scale_output(module, args, kwargs, output):
         self.assertEqual(eager_call_count, 7)
 
         for i, layer in model.layers.named_children():
-            model.layers.register_module(i, torch.compile(layer, fullgraph=True))
+            model.layers.register_module(
+                i, torch.compile(layer, fullgraph=True, backend="eager")
+            )
 
         call_count[0] = 0
         output = model(inp)
@@ -3609,6 +3624,67 @@ def scale_output(module, args, kwargs, output):
         self.assertEqual(compiled_call_count, eager_call_count)
         self.assertTrue(torch.allclose(output_eager, output))
 
+    @patch.object(torch._dynamo.config, "guard_nn_modules", True)
+    def test_dict_insertion_guard_method_func(self):
+        class SimpleModel(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.linear = torch.nn.Linear(10, 10)
+
+            def forward(self, x):
+                return self.linear(x)
+
+        def hook_function(module, args):
+            return (args[0] + 1.0,)
+
+        class HookHelper:
+            def hook_method(self, module, args):
+                return (args[0] + 2.0,)
+
+        model = SimpleModel()
+        helper = HookHelper()
+
+        model.register_forward_pre_hook(hook_function)
+        model.register_forward_pre_hook(helper.hook_method, prepend=True)
+
+        @torch.compile(fullgraph=True, backend="eager")
+        def runner_func(mod, x):
+            return mod(x)
+
+        input_tensor = torch.randn(1, 10)
+        # This would error before fixing guard orering on nn.Modules (https://github.com/pytorch/pytorch/issues/170429)
+        _ = runner_func(model, input_tensor)
+
+    def test_prepend_hook_ordering(self):
+        class HookedLinear(torch.nn.Linear):
+            def __init__(self, *args, **kwargs):
+                super().__init__(*args, **kwargs)
+                self.register_forward_pre_hook(self._hook_add)
+                self.register_forward_pre_hook(self._hook_mul, prepend=True)
+
+            @staticmethod
+            def _hook_add(module, args):
+                return (args[0] + 1,)
+
+            @staticmethod
+            def _hook_mul(module, args):
+                return (args[0] * 2,)
+
+        class Model(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.layer = HookedLinear(4, 4, bias=False)
+
+            def forward(self, x):
+                return self.layer(x)
+
+        model = Model()
+        x = torch.ones(1, 4)
+
+        eager = model(x)
+        compiled = torch.compile(model, backend="eager", fullgraph=True)(x)
+        self.assertEqual(eager, compiled)
+
 
 devices = ["cuda", "hpu", "xpu"]
 instantiate_device_type_tests(
diff --git a/test/dynamo/test_nb_bool.py b/test/dynamo/test_nb_bool.py
new file mode 100644
index 0000000000000..24887dada59ec
--- /dev/null
+++ b/test/dynamo/test_nb_bool.py
@@ -0,0 +1,278 @@
+# Owner(s): ["module: dynamo"]
+"""Tests for nb_bool / generic_bool: bool() via PyObject_IsTrue in Dynamo."""
+
+import collections
+import enum
+
+import torch
+
+
+class _Color(enum.Enum):
+    RED = 1
+    BLUE = 2
+
+
+import torch.nn
+from torch.testing._internal.common_utils import make_dynamo_test, run_tests, TestCase
+
+
+class NbBoolTests(TestCase):
+    # --- Scalar constants (ConstantVariable path) ---
+
+    @make_dynamo_test
+    def test_int(self):
+        self.assertEqual(bool(0), False)
+        self.assertEqual(bool(1), True)
+        self.assertEqual(bool(-1), True)
+
+    @make_dynamo_test
+    def test_float(self):
+        self.assertEqual(bool(0.0), False)
+        self.assertEqual(bool(-0.0), False)
+        self.assertEqual(bool(3.14), True)
+
+    @make_dynamo_test
+    def test_none(self):
+        self.assertEqual(bool(None), False)
+
+    @make_dynamo_test
+    def test_str(self):
+        self.assertEqual(bool(""), False)
+        self.assertEqual(bool("nonempty"), True)
+
+    @make_dynamo_test
+    def test_bytes(self):
+        self.assertEqual(bool(b""), False)
+        self.assertEqual(bool(b"hello"), True)
+
+    @make_dynamo_test
+    def test_bool(self):
+        self.assertEqual(False, False)
+        self.assertEqual(True, True)
+
+    @make_dynamo_test
+    def test_complex_zero(self):
+        self.assertEqual(bool(0j), False)
+
+    @make_dynamo_test
+    def test_complex_nonzero(self):
+        self.assertEqual(bool(1 + 2j), True)
+
+    @make_dynamo_test
+    def test_complex_real_nonzero_imag_zero(self):
+        self.assertEqual(bool(1 + 0j), True)
+
+    @make_dynamo_test
+    def test_complex_real_zero_imag_nonzero(self):
+        self.assertEqual(bool(0 + 1j), True)
+
+    # --- Containers (length fallback / _bool_from_length path) ---
+
+    @make_dynamo_test
+    def test_empty_list(self):
+        self.assertEqual(bool([]), False)
+
+    @make_dynamo_test
+    def test_nonempty_list(self):
+        self.assertEqual(bool([1, 2, 3]), True)
+
+    @make_dynamo_test
+    def test_empty_dict(self):
+        self.assertEqual(bool({}), False)
+
+    @make_dynamo_test
+    def test_nonempty_dict(self):
+        self.assertEqual(bool({"a": 1}), True)
+
+    @make_dynamo_test
+    def test_empty_tuple(self):
+        self.assertEqual(bool(()), False)
+
+    @make_dynamo_test
+    def test_nonempty_tuple(self):
+        self.assertEqual(bool((1,)), True)
+
+    @make_dynamo_test
+    def test_empty_set(self):
+        self.assertEqual(bool(set()), False)
+
+    @make_dynamo_test
+    def test_nonempty_set(self):
+        self.assertEqual(bool({1, 2}), True)
+
+    @make_dynamo_test
+    def test_empty_frozenset(self):
+        self.assertEqual(bool(frozenset()), False)
+
+    @make_dynamo_test
+    def test_nonempty_frozenset(self):
+        self.assertEqual(bool(frozenset({1})), True)
+
+    @make_dynamo_test
+    def test_empty_range(self):
+        self.assertEqual(bool(range(0)), False)
+
+    @make_dynamo_test
+    def test_nonempty_range(self):
+        self.assertEqual(bool(range(5)), True)
+
+    # --- dict subclasses ---
+
+    @make_dynamo_test
+    def test_empty_defaultdict(self):
+        d = collections.defaultdict(int)
+        self.assertEqual(bool(d), False)
+
+    @make_dynamo_test
+    def test_nonempty_defaultdict(self):
+        d = collections.defaultdict(int, {"x": 1})
+        self.assertEqual(bool(d), True)
+
+    @make_dynamo_test
+    def test_empty_counter(self):
+        c = collections.Counter()
+        self.assertEqual(bool(c), False)
+
+    @make_dynamo_test
+    def test_nonempty_counter(self):
+        c = collections.Counter("abc")
+        self.assertEqual(bool(c), True)
+
+    # --- Enum (UserDefinedClassVariable / ConstantVariable path) ---
+
+    @make_dynamo_test
+    def test_enum_member(self):
+        self.assertEqual(bool(_Color.RED), True)
+        self.assertEqual(bool(_Color.BLUE), True)
+
+    # --- UserDefinedObjectVariable tests (torch.compile path) ---
+
+    def test_user_defined_with_bool(self):
+        class MyObj:
+            def __init__(self, val):
+                self.val = val
+
+            def __bool__(self):
+                return self.val > 0
+
+        def fn(x, obj):
+            return x + 1 if bool(obj) else x - 1
+
+        x = torch.randn(4)
+        compiled = torch.compile(fn, backend="eager")
+        self.assertEqual(fn(x, MyObj(5)), compiled(x, MyObj(5)))
+        torch._dynamo.reset()
+        compiled = torch.compile(fn, backend="eager")
+        self.assertEqual(fn(x, MyObj(-1)), compiled(x, MyObj(-1)))
+
+    def test_user_defined_with_len(self):
+        class Container:
+            def __init__(self, items):
+                self.items = items
+
+            def __len__(self):
+                return len(self.items)
+
+        def fn(x, c):
+            return x + 1 if bool(c) else x - 1
+
+        x = torch.randn(4)
+        compiled = torch.compile(fn, backend="eager")
+        self.assertEqual(fn(x, Container([1, 2])), compiled(x, Container([1, 2])))
+        torch._dynamo.reset()
+        compiled = torch.compile(fn, backend="eager")
+        self.assertEqual(fn(x, Container([])), compiled(x, Container([])))
+
+    def test_user_defined_no_bool_no_len(self):
+        class Plain:
+            pass
+
+        def fn(x, obj):
+            return x + 1 if bool(obj) else x - 1
+
+        x = torch.randn(4)
+        compiled = torch.compile(fn, backend="eager", fullgraph=True)
+        self.assertEqual(fn(x, Plain()), compiled(x, Plain()))
+
+    def test_user_defined_bool_returns_non_bool_raises(self):
+        class BadBool:
+            def __bool__(self):
+                return 1  # noqa: PLE0305
+
+        def fn(x, obj):
+            return x + 1 if bool(obj) else x - 1
+
+        with self.assertRaises(TypeError):
+            bool(BadBool())
+        with self.assertRaises(TypeError):
+            torch.compile(fn, backend="eager")(torch.randn(4), BadBool())
+
+    # --- Metaclass with __bool__ (UserDefinedClassVariable path) ---
+
+    def test_metaclass_bool(self):
+        class Foo(type):
+            def __bool__(cls):
+                return False
+
+        class A(metaclass=Foo):
+            pass
+
+        class Bar(type):
+            pass
+
+        class B(metaclass=Bar):
+            pass
+
+        def fn(x):
+            # A's metaclass defines __bool__ returning False; B's does not (truthy).
+            return x + bool(A) + bool(B)
+
+        x = torch.randn(4)
+        compiled = torch.compile(fn, backend="eager", fullgraph=True)
+        self.assertEqual(fn(x), compiled(x))
+
+    # --- nn.Module (NNModuleVariable path) ---
+
+    def test_nn_module_nonempty(self):
+        mod = torch.nn.ModuleList([torch.nn.Linear(4, 4)])
+
+        def fn(x):
+            return x + 1 if bool(mod) else x - 1
+
+        x = torch.randn(4)
+        compiled = torch.compile(fn, backend="eager", fullgraph=True)
+        self.assertEqual(fn(x), compiled(x))
+
+    def test_nn_module_empty(self):
+        mod = torch.nn.ModuleList()
+
+        def fn(x):
+            return x + 1 if bool(mod) else x - 1
+
+        x = torch.randn(4)
+        compiled = torch.compile(fn, backend="eager", fullgraph=True)
+        self.assertEqual(fn(x), compiled(x))
+
+    # --- Tensor (TensorVariable path) ---
+
+    def test_tensor_nonzero(self):
+        def fn(x):
+            t = torch.tensor(1)
+            return x + 1 if bool(t) else x - 1
+
+        x = torch.randn(4)
+        compiled = torch.compile(fn, backend="eager", fullgraph=True)
+        self.assertEqual(fn(x), compiled(x))
+
+    def test_tensor_zero(self):
+        def fn(x):
+            t = torch.tensor(0)
+            return x + 1 if bool(t) else x - 1
+
+        x = torch.randn(4)
+        compiled = torch.compile(fn, backend="eager", fullgraph=True)
+        self.assertEqual(fn(x), compiled(x))
+
+
+if __name__ == "__main__":
+    run_tests()
diff --git a/test/dynamo/test_nb_float.py b/test/dynamo/test_nb_float.py
new file mode 100644
index 0000000000000..95e4df0446dee
--- /dev/null
+++ b/test/dynamo/test_nb_float.py
@@ -0,0 +1,321 @@
+# Owner(s): ["module: dynamo"]
+"""Tests for nb_float_impl: unified __float__ / float() protocol in Dynamo."""
+
+import torch
+import torch._dynamo.testing
+from torch.testing._internal.common_utils import (
+    make_dynamo_test,
+    run_tests,
+    skipIfCrossRef,
+    TestCase,
+)
+
+
+class NbFloatTests(TestCase):
+    # --- float / int / bool (ConstantVariable) ---
+
+    @make_dynamo_test
+    def test_float_float(self):
+        self.assertEqual(float(3.14), 3.14)  # noqa: UP018
+
+    @make_dynamo_test
+    def test_int_float(self):
+        self.assertEqual(float(5), 5.0)
+
+    @make_dynamo_test
+    def test_bool_float(self):
+        self.assertEqual(float(True), 1.0)
+        self.assertEqual(float(False), 0.0)
+
+    @make_dynamo_test
+    def test_float_dunder_float(self):
+        self.assertEqual((3.14).__float__(), 3.14)
+
+    @make_dynamo_test
+    def test_int_dunder_float(self):
+        self.assertEqual((5).__float__(), 5.0)
+
+    @make_dynamo_test
+    def test_bool_dunder_float(self):
+        self.assertEqual(True.__float__(), 1.0)
+
+    # --- TypeError for non-float types ---
+
+    def test_complex_float_raises(self):
+        def fn(x):
+            try:
+                return float(1 + 2j)
+            except TypeError as e:
+                return str(e)
+
+        result = torch.compile(fn, backend="eager", fullgraph=True)(torch.tensor(0))
+        eager_result = fn(torch.tensor(0))
+        self.assertEqual(result, eager_result)
+
+    def test_none_float_raises(self):
+        def fn(x):
+            try:
+                return float(None)
+            except TypeError as e:
+                return str(e)
+
+        result = torch.compile(fn, backend="eager", fullgraph=True)(torch.tensor(0))
+        eager_result = fn(torch.tensor(0))
+        self.assertEqual(result, eager_result)
+
+    def test_list_float_raises(self):
+        def fn(x):
+            try:
+                return float([1, 2])
+            except TypeError as e:
+                return str(e)
+
+        result = torch.compile(fn, backend="eager", fullgraph=True)(torch.tensor(0))
+        eager_result = fn(torch.tensor(0))
+        self.assertEqual(result, eager_result)
+
+    def test_dict_float_raises(self):
+        def fn(x):
+            try:
+                return float({})
+            except TypeError as e:
+                return str(e)
+
+        result = torch.compile(fn, backend="eager", fullgraph=True)(torch.tensor(0))
+        eager_result = fn(torch.tensor(0))
+        self.assertEqual(result, eager_result)
+
+    def test_set_float_raises(self):
+        def fn(x):
+            try:
+                return float(set())
+            except TypeError as e:
+                return str(e)
+
+        result = torch.compile(fn, backend="eager", fullgraph=True)(torch.tensor(0))
+        eager_result = fn(torch.tensor(0))
+        self.assertEqual(result, eager_result)
+
+    def test_tuple_float_raises(self):
+        def fn(x):
+            try:
+                return float(())
+            except TypeError as e:
+                return str(e)
+
+        result = torch.compile(fn, backend="eager", fullgraph=True)(torch.tensor(0))
+        eager_result = fn(torch.tensor(0))
+        self.assertEqual(result, eager_result)
+
+    # --- str parsing (constructor, not nb_float) ---
+
+    @make_dynamo_test
+    def test_str_float_parsing(self):
+        self.assertEqual(float("3.14"), 3.14)
+
+    @make_dynamo_test
+    def test_str_float_int(self):
+        self.assertEqual(float("123"), 123.0)
+
+    # --- UserDefinedObjectVariable with __float__ ---
+
+    def test_user_defined_float(self):
+        class MyFloat:
+            def __init__(self, v):
+                self.v = v
+
+            def __float__(self):
+                return self.v
+
+        obj = MyFloat(3.14)
+
+        def fn(x):
+            return float(obj)
+
+        self.assertEqual(
+            torch.compile(fn, backend="eager", fullgraph=True)(torch.tensor(0)), 3.14
+        )
+
+    def test_user_defined_dunder_float(self):
+        class MyFloat:
+            def __init__(self, v):
+                self.v = v
+
+            def __float__(self):
+                return self.v
+
+        obj = MyFloat(2.71)
+
+        def fn(x):
+            return obj.__float__()
+
+        self.assertEqual(
+            torch.compile(fn, backend="eager", fullgraph=True)(torch.tensor(0)), 2.71
+        )
+
+    def test_user_defined_no_float_raises(self):
+        class NoFloat:
+            pass
+
+        obj = NoFloat()
+
+        def fn(x):
+            try:
+                return float(obj)
+            except TypeError as e:
+                return str(e)
+
+        result = torch.compile(fn, backend="eager", fullgraph=True)(torch.tensor(0))
+        self.assertIn("float() argument must be a string", result)
+
+    def test_float_returning_non_float_raises(self):
+        class Bad:
+            def __float__(self):
+                return "not a float"  # noqa: PLE0305
+
+        obj = Bad()
+
+        def fn(x):
+            try:
+                return float(obj)
+            except TypeError as e:
+                return str(e)
+
+        result = torch.compile(fn, backend="eager", fullgraph=True)(torch.tensor(0))
+        self.assertIn("__float__ returned non-float", result)
+
+    def test_float_raising_exception_propagates(self):
+        class RaisingFloat:
+            def __float__(self):
+                raise ValueError("custom error")
+
+        obj = RaisingFloat()
+
+        def fn(x):
+            try:
+                return float(obj)
+            except ValueError as e:
+                return str(e)
+
+        result = torch.compile(fn, backend="eager", fullgraph=True)(torch.tensor(0))
+        self.assertEqual(result, "custom error")
+
+    def test_user_defined_staticmethod_float(self):
+        class StaticFloat:
+            @staticmethod
+            def __float__():
+                return 3.0
+
+        obj = StaticFloat()
+
+        def fn(x):
+            return float(obj)
+
+        self.assertEqual(
+            torch.compile(fn, backend="eager", fullgraph=True)(torch.tensor(0)), 3.0
+        )
+
+    def test_user_defined_classmethod_float(self):
+        class ClassFloat:
+            @classmethod
+            def __float__(cls):
+                return 4.0
+
+        obj = ClassFloat()
+
+        def fn(x):
+            return float(obj)
+
+        self.assertEqual(
+            torch.compile(fn, backend="eager", fullgraph=True)(torch.tensor(0)), 4.0
+        )
+
+    # --- nb_index fallback (PyNumber_Float step 3) ---
+
+    def test_index_fallback_for_float(self):
+        class HasIndex:
+            def __index__(self):
+                return 42
+
+        obj = HasIndex()
+
+        def fn(x):
+            return float(obj)
+
+        self.assertEqual(
+            torch.compile(fn, backend="eager", fullgraph=True)(torch.tensor(0)), 42.0
+        )
+
+    def test_index_fallback_no_float_no_index_raises(self):
+        class NoSlots:
+            pass
+
+        obj = NoSlots()
+
+        def fn(x):
+            try:
+                return float(obj)
+            except TypeError as e:
+                return str(e)
+
+        result = torch.compile(fn, backend="eager", fullgraph=True)(torch.tensor(0))
+        eager_result = fn(torch.tensor(0))
+        self.assertIn("float() argument must be a string", result)
+        self.assertEqual(result, eager_result)
+
+    # --- Tensor ---
+
+    @skipIfCrossRef
+    def test_tensor_int_dtype(self):
+        def fn(x):
+            return float(x)
+
+        result = torch.compile(fn, backend="eager", fullgraph=True)(torch.tensor(5))
+        self.assertEqual(result, 5.0)
+
+    @skipIfCrossRef
+    def test_tensor_float_dtype(self):
+        def fn(x):
+            return float(x)
+
+        result = torch.compile(fn, backend="eager", fullgraph=True)(torch.tensor(3.14))
+        self.assertAlmostEqual(result, 3.14, places=2)
+
+    @skipIfCrossRef
+    def test_tensor_complex_raises(self):
+        def fn(x):
+            try:
+                return float(x)
+            except RuntimeError as e:
+                return str(e)
+
+        result = torch.compile(fn, backend="eager", fullgraph=True)(
+            torch.tensor(1 + 2j)
+        )
+        eager_result = fn(torch.tensor(1 + 2j))
+        self.assertIn(
+            "value cannot be converted to type double without overflow", result
+        )
+        self.assertEqual(result, eager_result)
+
+    @skipIfCrossRef
+    def test_tensor_dunder_float(self):
+        def fn(x):
+            return x.__float__()
+
+        result = torch.compile(fn, backend="eager", fullgraph=True)(torch.tensor(5))
+        self.assertEqual(result, 5.0)
+
+    # --- SymNodeVariable ---
+
+    @skipIfCrossRef
+    def test_symnode_float(self):
+        def fn(x):
+            return float(x.size(0))
+
+        result = torch.compile(fn, backend="eager", fullgraph=True)(torch.randn(10, 20))
+        self.assertEqual(result, 10.0)
+
+
+if __name__ == "__main__":
+    run_tests()
diff --git a/test/dynamo/test_nb_index.py b/test/dynamo/test_nb_index.py
new file mode 100644
index 0000000000000..e06dc4addbdca
--- /dev/null
+++ b/test/dynamo/test_nb_index.py
@@ -0,0 +1,353 @@
+# Owner(s): ["module: dynamo"]
+"""Tests for nb_index_impl: unified __index__ / operator.index protocol in Dynamo."""
+
+import operator
+
+import torch
+import torch._dynamo.testing
+from torch.testing._internal.common_utils import make_dynamo_test, run_tests, TestCase
+
+
+class NbIndexTests(TestCase):
+    # --- int / bool (ConstantVariable) ---
+
+    @make_dynamo_test
+    def test_int_index(self):
+        self.assertEqual(operator.index(5), 5)
+
+    @make_dynamo_test
+    def test_bool_index(self):
+        self.assertEqual(operator.index(True), 1)
+
+    @make_dynamo_test
+    def test_int_dunder_index(self):
+        self.assertEqual((5).__index__(), 5)
+
+    # --- TypeError for non-index types ---
+
+    def test_str_index_raises(self):
+        def fn(x):
+            try:
+                return operator.index("hello")
+            except TypeError as e:
+                return str(e)
+
+        result = torch.compile(fn, backend="eager", fullgraph=True)(torch.tensor(0))
+        self.assertIn("cannot be interpreted as an integer", result)
+
+    def test_float_index_raises(self):
+        def fn(x):
+            try:
+                return operator.index(3.14)
+            except TypeError as e:
+                return str(e)
+
+        result = torch.compile(fn, backend="eager", fullgraph=True)(torch.tensor(0))
+        self.assertIn("cannot be interpreted as an integer", result)
+
+    def test_list_index_raises(self):
+        def fn(x):
+            try:
+                return operator.index([1, 2])
+            except TypeError as e:
+                return str(e)
+
+        result = torch.compile(fn, backend="eager", fullgraph=True)(torch.tensor(0))
+        self.assertIn("cannot be interpreted as an integer", result)
+
+    def test_none_index_raises(self):
+        def fn(x):
+            try:
+                return operator.index(None)
+            except TypeError as e:
+                return str(e)
+
+        result = torch.compile(fn, backend="eager", fullgraph=True)(torch.tensor(0))
+        self.assertIn("cannot be interpreted as an integer", result)
+
+    # --- UserDefinedObjectVariable with __index__ ---
+
+    def test_user_defined_index(self):
+        class MyInt:
+            def __init__(self, v):
+                self.v = v
+
+            def __index__(self):
+                return self.v
+
+        obj = MyInt(42)
+
+        def fn(x):
+            return operator.index(obj)
+
+        self.assertEqual(
+            torch.compile(fn, backend="eager", fullgraph=True)(torch.tensor(0)), 42
+        )
+
+    def test_user_defined_dunder_index(self):
+        class MyInt:
+            def __init__(self, v):
+                self.v = v
+
+            def __index__(self):
+                return self.v
+
+        obj = MyInt(7)
+
+        def fn(x):
+            return obj.__index__()
+
+        self.assertEqual(
+            torch.compile(fn, backend="eager", fullgraph=True)(torch.tensor(0)), 7
+        )
+
+    def test_user_defined_no_index_raises(self):
+        class NoIndex:
+            pass
+
+        obj = NoIndex()
+
+        def fn(x):
+            try:
+                return operator.index(obj)
+            except TypeError as e:
+                return str(e)
+
+        result = torch.compile(fn, backend="eager", fullgraph=True)(torch.tensor(0))
+        self.assertIn("cannot be interpreted as an integer", result)
+
+    # --- nb_index used in list/tuple subscript (PyNumber_AsSsize_t) ---
+
+    @make_dynamo_test
+    def test_list_subscript_with_bool(self):
+        lst = [10, 20, 30]
+        self.assertEqual(lst[True], 20)
+
+    @make_dynamo_test
+    def test_tuple_subscript_with_bool(self):
+        t = (10, 20, 30)
+        self.assertEqual(t[True], 20)
+
+    def test_list_subscript_with_user_defined_index(self):
+        class MyIdx:
+            def __index__(self):
+                return 1
+
+        idx = MyIdx()
+
+        def fn(x):
+            lst = [10, 20, 30]
+            return lst[idx]
+
+        self.assertEqual(
+            torch.compile(fn, backend="eager", fullgraph=True)(torch.tensor(0)), 20
+        )
+
+    def test_tuple_subscript_with_user_defined_index(self):
+        class MyIdx:
+            def __index__(self):
+                return 2
+
+        idx = MyIdx()
+
+        def fn(x):
+            t = (10, 20, 30)
+            return t[idx]
+
+        self.assertEqual(
+            torch.compile(fn, backend="eager", fullgraph=True)(torch.tensor(0)), 30
+        )
+
+    def test_list_subscript_with_no_index_raises(self):
+        class NoIndex:
+            pass
+
+        obj = NoIndex()
+
+        def fn(x):
+            lst = [10, 20, 30]
+            try:
+                return lst[obj]
+            except TypeError as e:
+                return str(e)
+
+        result = torch.compile(fn, backend="eager", fullgraph=True)(torch.tensor(0))
+        self.assertIn("list indices must be integers or slices", result)
+
+    def test_user_defined_staticmethod_index(self):
+        class StaticIdx:
+            @staticmethod
+            def __index__():
+                return 3
+
+        obj = StaticIdx()
+
+        def fn(x):
+            return operator.index(obj)
+
+        self.assertEqual(
+            torch.compile(fn, backend="eager", fullgraph=True)(torch.tensor(0)), 3
+        )
+
+    def test_user_defined_classmethod_index(self):
+        class ClassIdx:
+            @classmethod
+            def __index__(cls):
+                return 4
+
+        obj = ClassIdx()
+
+        def fn(x):
+            return operator.index(obj)
+
+        self.assertEqual(
+            torch.compile(fn, backend="eager", fullgraph=True)(torch.tensor(0)), 4
+        )
+
+    def test_custom_getitem_with_user_defined_index(self):
+        class MyIdx:
+            def __index__(self):
+                return 1
+
+        class MyContainer:
+            def __init__(self, data):
+                self.data = data
+
+            def __getitem__(self, idx):
+                return self.data[idx]
+
+        idx = MyIdx()
+        container = MyContainer([10, 20, 30])
+
+        def fn(x):
+            return container[idx]
+
+        self.assertEqual(
+            torch.compile(fn, backend="eager", fullgraph=True)(torch.tensor(0)), 20
+        )
+
+    def test_dict_getitem_with_non_indexable(self):
+        class NoIndex:
+            pass
+
+        obj = NoIndex()
+
+        def fn(x):
+            d = {obj: 42}
+            return d[obj]
+
+        self.assertEqual(
+            torch.compile(fn, backend="eager", fullgraph=True)(torch.tensor(0)), 42
+        )
+
+    def test_list_getitem_non_indexable_matches_cpython(self):
+        class NoIndex:
+            pass
+
+        obj = NoIndex()
+
+        def fn(x):
+            try:
+                return [10, 20][obj]
+            except TypeError:
+                return "caught"
+
+        # Both CPython and Dynamo should raise TypeError and catch it
+        eager_result = fn(torch.tensor(0))
+        compiled_result = torch.compile(fn, backend="eager", fullgraph=True)(
+            torch.tensor(0)
+        )
+        self.assertEqual(eager_result, "caught")
+        self.assertEqual(compiled_result, "caught")
+
+    def test_list_subscript_error_message_matches_cpython(self):
+        def fn(x):
+            try:
+                return [10, 20]["hello"]  # noqa: RUF016
+            except TypeError as e:
+                return str(e)
+
+        eager_result = fn(torch.tensor(0))
+        compiled_result = torch.compile(fn, backend="eager", fullgraph=True)(
+            torch.tensor(0)
+        )
+        self.assertIn("list indices must be integers or slices", eager_result)
+        self.assertIn("list indices must be integers or slices", compiled_result)
+
+    def test_index_returning_non_int_raises(self):
+        class Bad:
+            def __index__(self):
+                return "not an int"  # noqa: PLE0305
+
+        obj = Bad()
+
+        def fn(x):
+            try:
+                return operator.index(obj)
+            except TypeError as e:
+                return str(e)
+
+        result = torch.compile(fn, backend="eager", fullgraph=True)(torch.tensor(0))
+        self.assertIn("__index__ returned non-int", result)
+
+    def test_negative_index_via_user_defined(self):
+        class NegIdx:
+            def __index__(self):
+                return -1
+
+        idx = NegIdx()
+
+        def fn(x):
+            lst = [10, 20, 30]
+            return lst[idx]
+
+        self.assertEqual(
+            torch.compile(fn, backend="eager", fullgraph=True)(torch.tensor(0)), 30
+        )
+
+    def test_index_raising_exception_propagates(self):
+        class RaisingIdx:
+            def __index__(self):
+                raise ValueError("custom error")
+
+        obj = RaisingIdx()
+
+        def fn(x):
+            try:
+                return operator.index(obj)
+            except ValueError as e:
+                return str(e)
+
+        result = torch.compile(fn, backend="eager", fullgraph=True)(torch.tensor(0))
+        self.assertEqual(result, "custom error")
+
+    # --- Tensor __index__ ---
+
+    def test_tensor_int_index(self):
+        def fn(x):
+            return operator.index(x)
+
+        result = torch.compile(fn, backend="eager", fullgraph=True)(torch.tensor(5))
+        self.assertEqual(result, 5)
+
+    def test_tensor_float_index_raises(self):
+        def fn(x):
+            try:
+                return operator.index(x)
+            except TypeError as e:
+                return str(e)
+
+        result = torch.compile(fn, backend="eager", fullgraph=True)(torch.tensor(1.5))
+        self.assertIn("only integer tensors", result)
+
+    def test_list_subscript_with_tensor(self):
+        def fn(x):
+            lst = [10, 20, 30]
+            return lst[x]
+
+        result = torch.compile(fn, backend="eager", fullgraph=True)(torch.tensor(2))
+        self.assertEqual(result, 30)
+
+
+if __name__ == "__main__":
+    run_tests()
diff --git a/test/dynamo/test_nb_int.py b/test/dynamo/test_nb_int.py
new file mode 100644
index 0000000000000..431c1635a3ece
--- /dev/null
+++ b/test/dynamo/test_nb_int.py
@@ -0,0 +1,317 @@
+# Owner(s): ["module: dynamo"]
+"""Tests for nb_int_impl: unified __int__ / int() protocol in Dynamo."""
+
+import torch
+import torch._dynamo.testing
+from torch.testing._internal.common_utils import make_dynamo_test, run_tests, TestCase
+
+
+class NbIntTests(TestCase):
+    # --- int / bool (ConstantVariable) ---
+
+    @make_dynamo_test
+    def test_int_int(self):
+        self.assertEqual(5, 5)
+
+    @make_dynamo_test
+    def test_bool_int(self):
+        self.assertEqual(int(True), 1)
+        self.assertEqual(int(False), 0)
+
+    @make_dynamo_test
+    def test_int_dunder_int(self):
+        self.assertEqual((5).__int__(), 5)
+
+    @make_dynamo_test
+    def test_bool_dunder_int(self):
+        self.assertEqual(True.__int__(), 1)
+
+    # --- float (ConstantVariable) ---
+
+    @make_dynamo_test
+    def test_float_int(self):
+        self.assertEqual(int(3.14), 3)
+
+    @make_dynamo_test
+    def test_float_negative_int(self):
+        self.assertEqual(int(-2.9), -2)
+
+    @make_dynamo_test
+    def test_float_dunder_int(self):
+        self.assertEqual((3.14).__int__(), 3)
+
+    # --- TypeError for non-int types ---
+
+    def test_complex_int_raises(self):
+        def fn(x):
+            try:
+                return int(1 + 2j)
+            except TypeError as e:
+                return str(e)
+
+        result = torch.compile(fn, backend="eager", fullgraph=True)(torch.tensor(0))
+        eager_result = fn(torch.tensor(0))
+        self.assertEqual(result, eager_result)
+
+    def test_none_int_raises(self):
+        def fn(x):
+            try:
+                return int(None)
+            except TypeError as e:
+                return str(e)
+
+        result = torch.compile(fn, backend="eager", fullgraph=True)(torch.tensor(0))
+        eager_result = fn(torch.tensor(0))
+        self.assertEqual(result, eager_result)
+
+    def test_list_int_raises(self):
+        def fn(x):
+            try:
+                return int([1, 2])
+            except TypeError as e:
+                return str(e)
+
+        result = torch.compile(fn, backend="eager", fullgraph=True)(torch.tensor(0))
+        eager_result = fn(torch.tensor(0))
+        self.assertEqual(result, eager_result)
+
+    def test_dict_int_raises(self):
+        def fn(x):
+            try:
+                return int({})
+            except TypeError as e:
+                return str(e)
+
+        result = torch.compile(fn, backend="eager", fullgraph=True)(torch.tensor(0))
+        eager_result = fn(torch.tensor(0))
+        self.assertEqual(result, eager_result)
+
+    def test_set_int_raises(self):
+        def fn(x):
+            try:
+                return int(set())
+            except TypeError as e:
+                return str(e)
+
+        result = torch.compile(fn, backend="eager", fullgraph=True)(torch.tensor(0))
+        eager_result = fn(torch.tensor(0))
+        self.assertEqual(result, eager_result)
+
+    def test_tuple_int_raises(self):
+        def fn(x):
+            try:
+                return int(())
+            except TypeError as e:
+                return str(e)
+
+        result = torch.compile(fn, backend="eager", fullgraph=True)(torch.tensor(0))
+        eager_result = fn(torch.tensor(0))
+        self.assertEqual(result, eager_result)
+
+    # --- str parsing (constructor, not nb_int) ---
+
+    @make_dynamo_test
+    def test_str_int_parsing(self):
+        self.assertEqual(int("123"), 123)
+
+    @make_dynamo_test
+    def test_str_int_base(self):
+        self.assertEqual(int("ff", 16), 255)
+
+    # --- UserDefinedObjectVariable with __int__ ---
+
+    def test_user_defined_int(self):
+        class MyInt:
+            def __init__(self, v):
+                self.v = v
+
+            def __int__(self):
+                return self.v
+
+        obj = MyInt(42)
+
+        def fn(x):
+            return int(obj)
+
+        self.assertEqual(
+            torch.compile(fn, backend="eager", fullgraph=True)(torch.tensor(0)), 42
+        )
+
+    def test_user_defined_dunder_int(self):
+        class MyInt:
+            def __init__(self, v):
+                self.v = v
+
+            def __int__(self):
+                return self.v
+
+        obj = MyInt(7)
+
+        def fn(x):
+            return obj.__int__()
+
+        self.assertEqual(
+            torch.compile(fn, backend="eager", fullgraph=True)(torch.tensor(0)), 7
+        )
+
+    def test_user_defined_no_int_raises(self):
+        class NoInt:
+            pass
+
+        obj = NoInt()
+
+        def fn(x):
+            try:
+                return int(obj)
+            except TypeError as e:
+                return str(e)
+
+        result = torch.compile(fn, backend="eager", fullgraph=True)(torch.tensor(0))
+        self.assertIn("int() argument must be a string", result)
+
+    def test_int_returning_non_int_raises(self):
+        class Bad:
+            def __int__(self):
+                return "not an int"  # noqa: PLE0305
+
+        obj = Bad()
+
+        def fn(x):
+            try:
+                return int(obj)
+            except TypeError as e:
+                return str(e)
+
+        result = torch.compile(fn, backend="eager", fullgraph=True)(torch.tensor(0))
+        self.assertIn("__int__ returned non-int", result)
+
+    def test_int_raising_exception_propagates(self):
+        class RaisingInt:
+            def __int__(self):
+                raise ValueError("custom error")
+
+        obj = RaisingInt()
+
+        def fn(x):
+            try:
+                return int(obj)
+            except ValueError as e:
+                return str(e)
+
+        result = torch.compile(fn, backend="eager", fullgraph=True)(torch.tensor(0))
+        self.assertEqual(result, "custom error")
+
+    def test_user_defined_staticmethod_int(self):
+        class StaticInt:
+            @staticmethod
+            def __int__():
+                return 3
+
+        obj = StaticInt()
+
+        def fn(x):
+            return int(obj)
+
+        self.assertEqual(
+            torch.compile(fn, backend="eager", fullgraph=True)(torch.tensor(0)), 3
+        )
+
+    def test_user_defined_classmethod_int(self):
+        class ClassInt:
+            @classmethod
+            def __int__(cls):
+                return 4
+
+        obj = ClassInt()
+
+        def fn(x):
+            return int(obj)
+
+        self.assertEqual(
+            torch.compile(fn, backend="eager", fullgraph=True)(torch.tensor(0)), 4
+        )
+
+    # --- nb_index fallback (PyNumber_Long step 3) ---
+
+    def test_index_fallback_for_int(self):
+        class HasIndex:
+            def __index__(self):
+                return 42
+
+        obj = HasIndex()
+
+        def fn(x):
+            return int(obj)
+
+        self.assertEqual(
+            torch.compile(fn, backend="eager", fullgraph=True)(torch.tensor(0)), 42
+        )
+
+    def test_index_fallback_no_int_no_index_raises(self):
+        class NoSlots:
+            pass
+
+        obj = NoSlots()
+
+        def fn(x):
+            try:
+                return int(obj)
+            except TypeError as e:
+                return str(e)
+
+        result = torch.compile(fn, backend="eager", fullgraph=True)(torch.tensor(0))
+        eager_result = fn(torch.tensor(0))
+        self.assertIn("int() argument must be a string", result)
+        self.assertEqual(result, eager_result)
+
+    # --- Tensor ---
+
+    def test_tensor_int_dtype(self):
+        def fn(x):
+            return int(x)
+
+        result = torch.compile(fn, backend="eager", fullgraph=True)(torch.tensor(5))
+        self.assertEqual(result, 5)
+
+    def test_tensor_float_dtype(self):
+        def fn(x):
+            return int(x)
+
+        result = torch.compile(fn, backend="eager", fullgraph=True)(torch.tensor(3.7))
+        self.assertEqual(result, 3)
+
+    def test_tensor_complex_raises(self):
+        def fn(x):
+            try:
+                return int(x)
+            except RuntimeError as e:
+                return str(e)
+
+        result = torch.compile(fn, backend="eager", fullgraph=True)(
+            torch.tensor(1 + 2j)
+        )
+        eager_result = fn(torch.tensor(1 + 2j))
+        self.assertIn(
+            "value cannot be converted to type int64_t without overflow", result
+        )
+        self.assertEqual(result, eager_result)
+
+    def test_tensor_dunder_int(self):
+        def fn(x):
+            return x.__int__()
+
+        result = torch.compile(fn, backend="eager", fullgraph=True)(torch.tensor(5))
+        self.assertEqual(result, 5)
+
+    # --- SymNodeVariable ---
+
+    def test_symnode_int(self):
+        def fn(x):
+            return int(x.size(0))
+
+        result = torch.compile(fn, backend="eager", fullgraph=True)(torch.randn(10, 20))
+        self.assertEqual(result, 10)
+
+
+if __name__ == "__main__":
+    run_tests()
diff --git a/test/dynamo/test_nested_graph_breaks.py b/test/dynamo/test_nested_graph_breaks.py
index 52ec185b90b9b..54fdfd14426cd 100644
--- a/test/dynamo/test_nested_graph_breaks.py
+++ b/test/dynamo/test_nested_graph_breaks.py
@@ -73,7 +73,7 @@ def make_nested_cls(cls):
 global1, global2, global3, global4 = (torch.zeros(3),) * 4
 
 
-class NestedGraphBreakTests(torch._dynamo.test_case.TestCaseWithNestedGraphBreaks):
+class NestedGraphBreakTests(torch._dynamo.test_case.TestCase):
     def test_single_graph_break(self):
         # NOTE marking f1, f2, f3 as global
         # prevents them from being freevars
@@ -233,10 +233,10 @@ def test_supported_ctx_manager(self):
 
         @torch._dynamo.disable
         def check_disabled(value):
-            assert torch.is_grad_enabled() == value
+            assert torch.is_grad_enabled() == value  # noqa: S101
 
         def check(value):
-            assert torch.is_grad_enabled() == value
+            assert torch.is_grad_enabled() == value  # noqa: S101
 
         def f1(x):
             with torch.no_grad():
@@ -276,7 +276,7 @@ def test_inactive_ctx_manager(self):
         global check, f1, f2, f3
 
         def check(value):
-            assert torch.is_grad_enabled() == value
+            assert torch.is_grad_enabled() == value  # noqa: S101
 
         def f1(x, ctx1):
             x = x + 1
@@ -346,7 +346,7 @@ def f2(x):
     def test_ctx_manager_active_in_nested_call(self):
         @torch._dynamo.disable()
         def f1(x):
-            assert not torch.is_grad_enabled()
+            assert not torch.is_grad_enabled()  # noqa: S101
             return x + 1
 
         def f2(x):
@@ -370,7 +370,7 @@ def f3(x):
     def test_ctx_manager_active_in_nested_step_graph_break(self):
         def f1(x):
             torch._dynamo.step_unsupported()
-            assert not torch.is_grad_enabled()
+            assert not torch.is_grad_enabled()  # noqa: S101
             return x + 1
 
         def f2(x):
@@ -1038,6 +1038,27 @@ def f4(x):
         # multiplication by 32, 64, 128, 256
         self.assertEqual(cnts.op_count, 4)
 
+    def test_resume_closure_different_module_globals(self):
+        # Tests that resume functions with freevars (closures) from inlined
+        # frames get the correct f_globals. Without the factory fix,
+        # MAKE_FUNCTION inherits the root frame's globals, so the resume
+        # function for `inner` would not find HELPER_CONSTANT.
+        try:
+            from . import _test_nested_graph_breaks_helper
+        except ImportError:
+            import _test_nested_graph_breaks_helper
+
+        def outer(x):
+            return _test_nested_graph_breaks_helper.closure_with_graph_break(x) + 2
+
+        cnts = torch._dynamo.testing.CompileCounter()
+        opt_fn = torch._dynamo.optimize(backend=cnts)(outer)
+        x = torch.zeros(3)
+        res = outer(x)
+        ref = opt_fn(x)
+        self.assertEqual(ref, res)
+        self.assertEqual(cnts.frame_count, 2)
+
     def test_error_on_graph_break_nested(self):
         # error_on_graph_break in a nested frame
         cnts = torch._dynamo.testing.CompileCounter()
@@ -1104,6 +1125,256 @@ def f8(x):
         with self.assertRaises(torch._dynamo.exc.Unsupported):
             f8(inp)
 
+    def test_graph_break_in_wrapped_user_function(self):
+        def fn(x):
+            x = x + 1
+            torch._dynamo.graph_break()
+            assert torch.compiler.is_compiling()  # noqa: S101
+            assert not torch.is_grad_enabled()  # noqa: S101
+            return x + 2
+
+        cnts = torch._dynamo.testing.CompileCounter()
+
+        @torch.compile(backend=cnts)
+        def gn(x):
+            x = torch.no_grad()(fn)(x)
+            assert torch.compiler.is_compiling()  # noqa: S101
+            assert torch.is_grad_enabled()  # noqa: S101
+            return x
+
+        inp = torch.randn(3)
+        self.assertEqual(gn(inp), inp + 3)
+        self.assertEqual(cnts.frame_count, 2)
+
+    def test_graph_break_in_nested_wrapped_user_function(self):
+        # no_grad wraps enable_grad wraps no_grad wraps fn
+        def fn(x):
+            x = x + 1
+            torch._dynamo.graph_break()
+            assert torch.compiler.is_compiling()  # noqa: S101
+            assert not torch.is_grad_enabled()  # noqa: S101
+            return x + 2
+
+        cnts = torch._dynamo.testing.CompileCounter()
+
+        @torch.compile(backend=cnts)
+        def gn(x):
+            wrapped = torch.no_grad()(torch.enable_grad()(torch.no_grad()(fn)))
+            x = wrapped(x)
+            assert torch.compiler.is_compiling()  # noqa: S101
+            assert torch.is_grad_enabled()  # noqa: S101
+            return x
+
+        inp = torch.randn(3)
+        self.assertEqual(gn(inp), inp + 3)
+        self.assertEqual(cnts.frame_count, 2)
+
+    def test_graph_break_in_wrapped_user_function_with_multiple_breaks(self):
+        def fn(x):
+            x = x + 1
+            torch._dynamo.graph_break()
+            assert torch.compiler.is_compiling()  # noqa: S101
+            assert not torch.is_grad_enabled()  # noqa: S101
+            x = x + 2
+            torch._dynamo.graph_break()
+            assert torch.compiler.is_compiling()  # noqa: S101
+            assert not torch.is_grad_enabled()  # noqa: S101
+            return x + 3
+
+        cnts = torch._dynamo.testing.CompileCounter()
+
+        @torch.compile(backend=cnts)
+        def gn(x):
+            x = torch.no_grad()(fn)(x)
+            assert torch.compiler.is_compiling()  # noqa: S101
+            assert torch.is_grad_enabled()  # noqa: S101
+            return x
+
+        inp = torch.randn(3)
+        self.assertEqual(gn(inp), inp + 6)
+        self.assertEqual(cnts.frame_count, 3)
+
+    def test_graph_break_in_sequential_wrapped_user_functions(self):
+        def fn1(x):
+            x = x + 1
+            torch._dynamo.graph_break()
+            assert torch.compiler.is_compiling()  # noqa: S101
+            assert not torch.is_grad_enabled()  # noqa: S101
+            return x + 2
+
+        def fn2(x):
+            x = x + 3
+            torch._dynamo.graph_break()
+            assert torch.compiler.is_compiling()  # noqa: S101
+            assert not torch.is_grad_enabled()  # noqa: S101
+            return x + 4
+
+        cnts = torch._dynamo.testing.CompileCounter()
+
+        @torch.compile(backend=cnts)
+        def gn(x):
+            x = torch.no_grad()(fn1)(x)
+            assert torch.compiler.is_compiling()  # noqa: S101
+            assert torch.is_grad_enabled()  # noqa: S101
+            x = torch.no_grad()(fn2)(x)
+            assert torch.compiler.is_compiling()  # noqa: S101
+            assert torch.is_grad_enabled()  # noqa: S101
+            return x
+
+        inp = torch.randn(3)
+        self.assertEqual(gn(inp), inp + 10)
+        # 3 frames: before fn1's break, after fn1's break through fn2's break,
+        # after fn2's break to end
+        self.assertEqual(cnts.frame_count, 3)
+
+    def test_graph_break_in_wrapped_enable_grad(self):
+        @torch.no_grad()
+        def fn(x):
+            x = x + 1
+            # enable_grad inside no_grad
+            x = torch.enable_grad()(lambda y: y + 2)(x)
+            torch._dynamo.graph_break()
+            assert torch.compiler.is_compiling()  # noqa: S101
+            assert not torch.is_grad_enabled()  # noqa: S101
+            return x + 3
+
+        cnts = torch._dynamo.testing.CompileCounter()
+
+        @torch.compile(backend=cnts)
+        def gn(x):
+            return fn(x)
+
+        inp = torch.randn(3)
+        self.assertEqual(gn(inp), inp + 6)
+
+    def test_graph_break_in_wrapped_user_method(self):
+        class Foo:
+            def __init__(self):
+                self.a = 1
+                self.b = 2
+
+            def fn(self, x):
+                x = x + self.a
+                torch._dynamo.graph_break()
+                assert torch.compiler.is_compiling()  # noqa: S101
+                assert not torch.is_grad_enabled()  # noqa: S101
+                return x + self.b
+
+        obj = Foo()
+
+        cnts = torch._dynamo.testing.CompileCounter()
+
+        @torch.compile(backend=cnts)
+        def gn(x):
+            obj.fn = torch.no_grad()(obj.fn)
+            x = obj.fn(x)
+            assert torch.compiler.is_compiling()  # noqa: S101
+            assert torch.is_grad_enabled()  # noqa: S101
+            return x
+
+        inp = torch.randn(3)
+        self.assertEqual(gn(inp), inp + 3)
+        self.assertEqual(cnts.frame_count, 2)
+
+    def test_graph_break_in_wrapped_nested_function(self):
+        cnts = torch._dynamo.testing.CompileCounter()
+
+        @torch.compile(backend=cnts)
+        def gn(x):
+            a = 1
+            b = 2
+
+            @torch.no_grad()
+            def fn(x):
+                x = x + a
+                torch._dynamo.graph_break()
+                assert torch.compiler.is_compiling()  # noqa: S101
+                assert not torch.is_grad_enabled()  # noqa: S101
+                return x + b
+
+            x = fn(x)
+            assert torch.compiler.is_compiling()  # noqa: S101
+            assert torch.is_grad_enabled()  # noqa: S101
+            return x
+
+        inp = torch.randn(3)
+        self.assertEqual(gn(inp), inp + 3)
+        self.assertEqual(cnts.frame_count, 2)
+
+    def test_graph_break_in_wrapped_skipped_function(self):
+        from torch._dynamo.testing import _skipped_function_for_test_reconstruct
+
+        def fn(x):
+            x = x + 1
+            torch._dynamo.graph_break()
+            assert torch.compiler.is_compiling()  # noqa: S101
+            assert not torch.is_grad_enabled()  # noqa: S101
+            return x + 2
+
+        cnts = torch._dynamo.testing.CompileCounter()
+
+        @torch.compile(backend=cnts)
+        def gn(x):
+            x = torch.no_grad()(_skipped_function_for_test_reconstruct)(fn, x)
+            assert torch.compiler.is_compiling()  # noqa: S101
+            assert torch.is_grad_enabled()  # noqa: S101
+            return x
+
+        inp = torch.randn(3)
+        self.assertEqual(gn(inp), inp + 3)
+        self.assertEqual(cnts.frame_count, 2)
+
+    def test_step_graph_break_frame_values_not_corrupted(self):
+        """Bytecode generation bug in step_graph_break corrupted parent frame
+        locals when the parent had a non-empty operand stack (num_stack > 0).
+        """
+
+        def inner(x):
+            x = x + 1
+            x = x + 1
+            torch._dynamo.step_unsupported()
+            return x
+
+        cnts = torch._dynamo.testing.CompileCounter()
+
+        @torch.compile(backend=cnts)
+        def fn(x):
+            x = x + 1
+            y = (x, inner(x))
+            return x, y
+
+        x = torch.tensor([1.0, 2.0])
+        result = fn(x)
+        self.assertEqual(result[0], torch.tensor([2.0, 3.0]))
+        self.assertEqual(
+            result[1], (torch.tensor([2.0, 3.0]), torch.tensor([4.0, 5.0]))
+        )
+        self.assertEqual(cnts.frame_count, 1)
+        self.assertEqual(cnts.op_count, 3)
+
+    def test_contextmanager_graph_break_in_init(self):
+        """Graph break in _GeneratorContextManager.__init__ when the generator
+        function is @torch._disable_dynamo (the DDP pattern)."""
+        from contextlib import contextmanager
+
+        @contextmanager
+        @torch._disable_dynamo(recursive=False)
+        def my_ctx():
+            yield
+
+        cnts = torch._dynamo.testing.CompileCounter()
+
+        @torch.compile(backend=cnts)
+        def fn(x):
+            x = x + 1
+            with my_ctx():
+                x = x + 2
+            return x + 3
+
+        inp = torch.randn(3)
+        self.assertEqual(fn(inp), inp + 6)
+        self.assertEqual(cnts.frame_count, 1)
+
 
 if __name__ == "__main__":
     from torch._dynamo.test_case import run_tests
diff --git a/test/dynamo/test_package.py b/test/dynamo/test_package.py
index c6c9a8247515b..87b4c088a1a1a 100644
--- a/test/dynamo/test_package.py
+++ b/test/dynamo/test_package.py
@@ -492,7 +492,7 @@ def guard_filter_fn(guards):
             compiled_fn(*args)
 
         total_frames = torch._dynamo.convert_frame.FRAME_COUNTER
-        self._save_and_reload(expected_backends=8, expected_dynamo=1)
+        self._save_and_reload(expected_backends=9, expected_dynamo=1)
 
         compiled_fn = torch._dynamo.optimize(
             backend="inductor", guard_filter_fn=guard_filter_fn
@@ -677,6 +677,78 @@ def forward(self, x):
             compiled_fn(x)
             self.assertEqual(torch._dynamo.convert_frame.FRAME_COUNTER, total_frames)
 
+    class _tempTensorSamplerForQualName:
+        def __init__(self, val, mask, prob):
+            self.val = val
+            self.mask = mask
+            self.prob = prob
+
+        @classmethod
+        def class_method_that_is_used(cls, x):
+            prob = torch.sigmoid(x)
+            thresh = torch.rand(1, device=x.device)
+            mask = (prob > thresh).to(torch.bool)
+            return cls(x, mask, prob)
+
+        @classmethod
+        def class_method_that_is_not_used(cls, x):
+            prob = torch.sigmoid(x)
+            thresh = torch.rand(1, device=x.device)
+            mask = (prob > thresh).to(torch.bool)
+            return cls(x, mask, prob)
+
+        def instance_method_that_is_used(self, x):
+            return x / 2
+
+    class _tempNetForQualName(torch.nn.Module):
+        def __init__(self):
+            super().__init__()
+
+        def instance_method_without_args(self):
+            shape = [1, 2, 3, 4]
+            x = torch.randn(shape)
+            return x
+
+        def instance_method_with_args(self, x):
+            return x + 1
+
+        def forward(self, x):
+            x *= x
+            with torch.device(x.device):
+                y = self.instance_method_without_args()
+            # test classmethod called from class
+            sampler = (
+                TestPackage._tempTensorSamplerForQualName.class_method_that_is_used(x)
+            )
+            x = torch.where(torch.rand_like(x) < sampler.prob, sampler.val, x) + y.sum()
+            # test instance method called from instance
+            x = sampler.instance_method_that_is_used(x)
+            # test classmethod called from instance
+            another_sampler = sampler.class_method_that_is_not_used(x)
+            # test instance method called from instance
+            x = another_sampler.instance_method_that_is_used(x)
+            # test classmethod called from instance
+            x += y.sum()
+            x = self.instance_method_with_args(x)
+            return x
+
+    @parametrize("device", ("cpu", "cuda", "xpu"))
+    @torch._dynamo.config.patch(caching_precompile=True)
+    def test_classmethod_qualname(self, device):
+        if device == "cuda" and not HAS_CUDA_AND_TRITON:
+            raise unittest.SkipTest("Requires CUDA/Triton")
+        if device == "xpu" and not HAS_XPU_AND_TRITON:
+            raise unittest.SkipTest("Requires XPU/Triton")
+
+        x = torch.rand(10, device=device)
+        model = TestPackage._tempNetForQualName()
+        model.forward(x)
+        compiled_fn = torch.compile(
+            model.forward,
+            options=dict(guard_filter_fn=torch.compiler.skip_guard_on_globals_unsafe),
+        )
+        compiled_fn(x)
+
 
 if __name__ == "__main__":
     from torch._dynamo.test_case import run_tests
diff --git a/test/dynamo/test_pgo.py b/test/dynamo/test_pgo.py
index e7e572d6f1ab6..5c16a424d4bd9 100644
--- a/test/dynamo/test_pgo.py
+++ b/test/dynamo/test_pgo.py
@@ -11,7 +11,11 @@
 import torch._inductor.mock_cache as mock_cache
 import torch.compiler.config
 import torch.nested
-from torch._dynamo.testing import CompileCounter
+from torch._dynamo.testing import (
+    CompileCounter,
+    CompileCounterWithBackend,
+    normalize_gm,
+)
 from torch._inductor.cpp_builder import normalize_path_separator
 from torch._inductor.utils import clear_caches, fresh_cache
 from torch.testing._internal.common_utils import IS_WINDOWS
@@ -372,15 +376,18 @@ def func(x):
         path2 = normalize_path_separator(os.path.join(temp_dir2.name, "example.py"))
         cnts = CompileCounter()
 
-        assert path1 != path2
+        if path1 == path2:
+            raise AssertionError("Expected path1 != path2")
 
         def write_load_and_run(path):
             with open(path, "w") as file:
                 file.write(content)
             spec = importlib.util.spec_from_file_location("example", path1)
-            assert spec is not None
+            if spec is None:
+                raise AssertionError("Expected spec to not be None")
             module = importlib.util.module_from_spec(spec)
-            assert spec.loader is not None
+            if spec.loader is None:
+                raise AssertionError("Expected spec.loader to not be None")
             spec.loader.exec_module(module)
             module.run(cnts)
 
@@ -467,6 +474,53 @@ def t(x, y):
                 f(t(2, 2), t(2, 4))
                 self.assertEqual(cnts.frame_count, 2)
 
+    def test_factory_created_functions_separate_pgo_state(self):
+        cc_backend = CompileCounterWithBackend(backend="aot_eager")
+
+        def wrap_compile(func):
+            @torch.compile(backend=cc_backend, fullgraph=True)
+            def _func(x):
+                return func(x)
+
+            return _func
+
+        def impl_a(x):
+            return x.view(-1)
+
+        def impl_b(x):
+            return x.sum()
+
+        compiled_a = wrap_compile(impl_a)
+        compiled_b = wrap_compile(impl_b)
+
+        # Call with different shapes - they should NOT share PGO state
+        compiled_a(torch.randn(2, 3))  # shape (2, 3)
+        compiled_b(torch.randn(5, 7))  # shape (5, 7)
+
+        actual = normalize_gm(cc_backend.graphs[1].print_readable(print_output=False))
+
+        # If PGO state is properly separated, compiled_b should NOT have
+        # been triggered to use dynamic shapes based on compiled_a's shapes
+        # This should result in 2 separate static compilations
+        self.assertExpectedInline(
+            actual,
+            """\
+class GraphModule(torch.nn.Module):
+    def forward(self, L_x_: "f32[5, 7]"):
+        l_x_ = L_x_
+
+        sum_1: "f32[]" = l_x_.sum();  l_x_ = None
+        return (sum_1,)
+""",
+        )
+
+        # Verify no dynamic shapes were introduced incorrectly
+        # by checking that calling with same shapes doesn't recompile
+        cc_backend.clear()
+        compiled_a(torch.randn(2, 3))
+        compiled_b(torch.randn(5, 7))
+        self.assertEqual(cc_backend.frame_count, 0)  # No recompilation needed
+
 
 if __name__ == "__main__":
     from torch._dynamo.test_case import run_tests
diff --git a/test/dynamo/test_precompile_context.py b/test/dynamo/test_precompile_context.py
index af86220d0cdf1..805bd939eaebe 100644
--- a/test/dynamo/test_precompile_context.py
+++ b/test/dynamo/test_precompile_context.py
@@ -60,12 +60,16 @@ def simple_function(x):
         self.assertEqual(len(PrecompileContext._backend_artifacts_by_key), 1)
         for key in PrecompileContext._backend_artifacts_by_key:
             result = PrecompileContext.serialize_artifact_by_key(key)
-            assert isinstance(result, BackendCacheArtifact)
+            if not isinstance(result, BackendCacheArtifact):
+                raise AssertionError(
+                    f"Expected BackendCacheArtifact, got {type(result)}"
+                )
             self.assertEqual(result.key, key)
 
         # This should still work
         result, _ = PrecompileContext.create_cache_entries()
-        assert len(result) == 1
+        if len(result) != 1:
+            raise AssertionError(f"Expected len(result) == 1, got {len(result)}")
 
     @requires_triton()
     def test_editable(self):
@@ -94,13 +98,20 @@ def edit_fn(x):
         PrecompileContext.edit_artifact(key, edit_fn)
 
         result = PrecompileContext.serialize_artifact_by_key(key)
-        assert isinstance(result, BundledAOTAutogradCacheArtifact)
+        if not isinstance(result, BundledAOTAutogradCacheArtifact):
+            raise AssertionError(
+                f"Expected BundledAOTAutogradCacheArtifact, got {type(result)}"
+            )
         self.assertEqual(result.key, key)
 
         result, _ = PrecompileContext.create_cache_entries()
-        assert len(result) == 1
+        if len(result) != 1:
+            raise AssertionError(f"Expected len(result) == 1, got {len(result)}")
         aot_autograd_artifacts = next(iter(result.values())).backends
-        assert len(aot_autograd_artifacts) == 1
+        if len(aot_autograd_artifacts) != 1:
+            raise AssertionError(
+                f"Expected len(aot_autograd_artifacts) == 1, got {len(aot_autograd_artifacts)}"
+            )
         entry = next(iter(aot_autograd_artifacts.values())).content
         self.assertEqual(entry._my_private_field, 42)
 
diff --git a/test/dynamo/test_profiler.py b/test/dynamo/test_profiler.py
index 36c127d12ad54..0145b9f79bb83 100644
--- a/test/dynamo/test_profiler.py
+++ b/test/dynamo/test_profiler.py
@@ -259,7 +259,7 @@ def fn(x, y):
             return x * r
 
         with torch.profiler.profile() as prof:
-            fn_c = torch.compile(fn)
+            fn_c = torch.compile(fn, backend="eager")
 
             fn_c(
                 torch.randn(10),
diff --git a/test/dynamo/test_python_autograd.py b/test/dynamo/test_python_autograd.py
index a6117bb4093a7..c9b88e488e7c6 100644
--- a/test/dynamo/test_python_autograd.py
+++ b/test/dynamo/test_python_autograd.py
@@ -1,5 +1,5 @@
 # Owner(s): ["module: dynamo"]
-from typing import NamedTuple, Optional, TYPE_CHECKING
+from typing import NamedTuple, TYPE_CHECKING
 
 import torch
 import torch._dynamo
@@ -30,14 +30,14 @@ def fresh_name() -> str:
 
 
 class Variable:
-    def __init__(self, value: torch.Tensor, name: Optional[str] = None):
+    def __init__(self, value: torch.Tensor, name: str | None = None):
         self.value = value
         self.name = name or fresh_name()
 
     # We need to start with some tensors whose values were not computed
     # inside the autograd. This function constructs leaf nodes.
     @staticmethod
-    def constant(value: torch.Tensor, name: Optional[str] = None):
+    def constant(value: torch.Tensor, name: str | None = None):
         return Variable(value, name)
 
     def __repr__(self):
@@ -51,7 +51,7 @@ def __mul__(self, rhs: "Variable") -> "Variable":
     def __add__(self, rhs: "Variable") -> "Variable":
         return operator_add(self, rhs)
 
-    def sum(self, name: Optional[str] = None) -> "Variable":
+    def sum(self, name: str | None = None) -> "Variable":
         return operator_sum(self, name)
 
     def expand(self, sizes: list[int]) -> "Variable":
@@ -168,7 +168,7 @@ def propagate(dL_doutputs: list[Variable]):
     return r
 
 
-def operator_sum(self: Variable, name: Optional[str]) -> "Variable":
+def operator_sum(self: Variable, name: str | None) -> "Variable":
     r = Variable(torch.sum(self.value), name=name)
     # print(f'{r.name} = {self.name}.sum()')
 
@@ -184,7 +184,7 @@ def propagate(dL_doutputs: list[Variable]):
 
 
 def operator_expand(self: Variable, sizes: list[int]) -> "Variable":
-    assert self.value.dim() == 0  # only works for scalars
+    assert self.value.dim() == 0  # noqa: S101 - only works for scalars
     r = Variable(self.value.expand(sizes))
     # print(f'{r.name} = {self.name}.expand({sizes})')
 
diff --git a/test/dynamo/test_recompile_ux.py b/test/dynamo/test_recompile_ux.py
index 880d37434b31f..b90635f119090 100644
--- a/test/dynamo/test_recompile_ux.py
+++ b/test/dynamo/test_recompile_ux.py
@@ -42,7 +42,8 @@ def trigger():
         def compiler(gm, input):
             nonlocal attached
             f = gm.forward
-            assert not attached
+            if attached:
+                raise AssertionError("Expected not attached")
             # NB: making this a weakref.ref causes the cycle to no
             # longer be promptly GC'ed
             weakref.finalize(f, trigger)
@@ -320,6 +321,120 @@ def h(x):
         self.assertEqual(res, inp + 3)
 
 
+class RecompileLimitKwargTests(torch._dynamo.test_case.TestCase):
+    @staticmethod
+    def _num_cache_entries(code):
+        return len(torch._dynamo.eval_frame._debug_get_cache_entry_list(code))
+
+    def test_recompile_limit_basic(self):
+        cnt = torch._dynamo.testing.CompileCounter()
+
+        def f(x, y):
+            return x + y
+
+        opt_f = torch.compile(f, backend=cnt, recompile_limit=2)
+
+        opt_f(torch.randn(3), torch.randn(3))
+        self.assertEqual(self._num_cache_entries(f), 1)
+
+        opt_f(torch.randn(3, dtype=torch.float64), torch.randn(3, dtype=torch.float64))
+        self.assertEqual(self._num_cache_entries(f), 2)
+
+        # Third dtype should NOT trigger recompilation (recompile_limit=2)
+        opt_f(torch.randn(3, dtype=torch.float16), torch.randn(3, dtype=torch.float16))
+        self.assertEqual(self._num_cache_entries(f), 2)
+
+    def test_recompile_limit_none_uses_global(self):
+        cnt = torch._dynamo.testing.CompileCounter()
+
+        def f(x, y):
+            return x + y
+
+        # Without recompile_limit kwarg, uses global config (default 8)
+        opt_f = torch.compile(f, backend=cnt)
+
+        for i in range(10):
+            dtype = [
+                torch.float32,
+                torch.float64,
+                torch.float16,
+                torch.bfloat16,
+                torch.int32,
+                torch.int64,
+                torch.int16,
+                torch.int8,
+                torch.uint8,
+                torch.complex64,
+            ][i]
+            opt_f(torch.ones(3, dtype=dtype), torch.ones(3, dtype=dtype))
+
+        self.assertEqual(
+            self._num_cache_entries(f), torch._dynamo.config.recompile_limit
+        )
+
+    def test_recompile_limit_fullgraph_raises(self):
+        """With fullgraph=True, hitting the recompile_limit kwarg raises
+        FailOnRecompileLimitHit, consistent with the fullgraph contract."""
+        cnt = torch._dynamo.testing.CompileCounter()
+
+        def f(x):
+            return x.sin()
+
+        opt_f = torch.compile(f, backend=cnt, fullgraph=True, recompile_limit=1)
+
+        opt_f(torch.randn(3))
+        self.assertEqual(cnt.frame_count, 1)
+
+        with self.assertRaises(FailOnRecompileLimitHit):
+            opt_f(torch.randn(3, dtype=torch.float64))
+
+    @torch._dynamo.config.patch(automatic_dynamic_shapes=True)
+    def test_recompile_limit_resume_function_auto_dynamic(self):
+        """With automatic dynamic shapes and recompile_limit=2, the resume
+        function recompiles via dimension changes on a global tensor while
+        the main function gets cache hits. The resume function should stop
+        at 2 entries and fall back to eager."""
+        cnt = torch._dynamo.testing.CompileCounter()
+
+        y_holder = {"tensor": torch.randn(4, 8, 2)}
+
+        def f(x):
+            x.sin()
+            print("graph break")
+            return y_holder["tensor"].cos()
+
+        opt_f = torch.compile(f, backend=cnt, recompile_limit=2)
+
+        # Call 1: static compile
+        y_holder["tensor"] = torch.randn(4, 8, 2)
+        opt_f(torch.randn(4, 8, 2))
+
+        # Call 2: y dim0 changes -> f cache hit, resume recompiles
+        y_holder["tensor"] = torch.randn(5, 8, 2)
+        opt_f(torch.randn(4, 8, 2))
+        frame_count_after_2 = cnt.frame_count
+
+        # Call 3: y dim1 changes -> resume should NOT recompile
+        # (resume already has 2 entries = recompile_limit)
+        y_holder["tensor"] = torch.randn(5, 9, 2)
+        opt_f(torch.randn(4, 8, 2))
+        self.assertEqual(cnt.frame_count, frame_count_after_2)
+
+        # Verify f has 1 entry, resume has 2
+        num_f_entries = len(torch._dynamo.eval_frame._debug_get_cache_entry_list(f))
+        self.assertEqual(num_f_entries, 1)
+
+        from torch._dynamo.resume_execution import ContinueExecutionCache
+
+        resume_codes = list(ContinueExecutionCache.cache[f.__code__].values())
+        self.assertTrue(len(resume_codes) > 0, "No resume functions found")
+        for resume_code in resume_codes:
+            num_resume_entries = len(
+                torch._dynamo.eval_frame._debug_get_cache_entry_list(resume_code)
+            )
+            self.assertEqual(num_resume_entries, 2)
+
+
 if __name__ == "__main__":
     from torch._dynamo.test_case import run_tests
 
diff --git a/test/dynamo/test_recompiles.py b/test/dynamo/test_recompiles.py
index a8579cc4c6178..827062f17985e 100644
--- a/test/dynamo/test_recompiles.py
+++ b/test/dynamo/test_recompiles.py
@@ -8,54 +8,6 @@
 
 
 class RecompileTests(torch._dynamo.test_case.TestCase):
-    def test_inline_inbuilt_nn_modules_candidate(self):
-        def hook_flag_on(guard_manager, f_locals, builder):
-            self.assertTrue(
-                "[inline-inbuilt-nn-modules-candidate]" not in str(guard_manager)
-            )
-
-        def hook_flag_off(guard_manager, f_locals, builder):
-            self.assertTrue(
-                "[inline-inbuilt-nn-modules-candidate]" in str(guard_manager)
-            )
-
-        class SubMod(torch.nn.Module):
-            def __init__(self):
-                super().__init__()
-                self.linear = torch.nn.Linear(2, 2)
-
-            @torch.compile(backend="eager")
-            def forward(self, x):
-                return self.linear(x)
-
-        class Mod(torch.nn.Module):
-            def __init__(self):
-                super().__init__()
-                self.sm1 = SubMod()
-                self.sm2 = SubMod()
-
-            def forward(self, x):
-                return self.sm1(x) + self.sm2(x)
-
-        try:
-            from .utils import install_guard_manager_testing_hook
-        except ImportError:
-            from utils import install_guard_manager_testing_hook
-
-        with (
-            install_guard_manager_testing_hook(hook_flag_on),
-            dc.patch(inline_inbuilt_nn_modules=True),
-        ):
-            mod = Mod()
-            mod(torch.randn(2, 2))
-
-        with (
-            install_guard_manager_testing_hook(hook_flag_off),
-            dc.patch(inline_inbuilt_nn_modules=False),
-        ):
-            mod = Mod()
-            mod(torch.randn(2, 2))
-
     def test_automatic_dynamic_reduce_recompiles(self):
         # Test the counterfactual, lots of recompiles without this config
         def foo(x, y):
diff --git a/test/dynamo/test_reconstruct.py b/test/dynamo/test_reconstruct.py
index c65190572754d..6203d127fd13b 100644
--- a/test/dynamo/test_reconstruct.py
+++ b/test/dynamo/test_reconstruct.py
@@ -57,6 +57,89 @@ def f(d, t):
             opt_f(d_opt, t)
             self.assertEqual(d, d_opt)
 
+    def _compile_and_capture_side_effects(self, fn, *args):
+        """Compile fn and return side-effect metadata from bytecode hooks."""
+        captured = {}
+
+        def rewrite_hook(code, out_code):
+            return out_code.replace(co_name=f"{out_code.co_name}_hooked")
+
+        def inspect_hook(code, out_code):
+            captured["refs"] = (
+                torch._dynamo.convert_frame.get_compiled_code_side_effects(out_code)
+            )
+            captured["has_side_effects"] = (
+                torch._dynamo.convert_frame.compiled_code_has_side_effects(out_code)
+            )
+
+        torch._dynamo.reset()
+        rewrite_handle = torch._dynamo.convert_frame.register_bytecode_hook(
+            rewrite_hook
+        )
+        inspect_handle = torch._dynamo.convert_frame.register_bytecode_hook(
+            inspect_hook
+        )
+        try:
+            torch.compile(fn, backend="eager", fullgraph=True)(*args)
+        finally:
+            inspect_handle.remove()
+            rewrite_handle.remove()
+
+        return captured
+
+    def test_bytecode_hook_exposes_side_effect_refs(self):
+        def mutating_fn(x, lst):
+            lst.append(x + 1)
+            return x * 2
+
+        def pure_fn(x):
+            return x * 2
+
+        x = torch.randn(3)
+
+        mutated = self._compile_and_capture_side_effects(mutating_fn, x, [])
+        self.assertEqual(mutated["refs"], ("L['lst']",))
+        self.assertTrue(mutated["has_side_effects"])
+
+        pure = self._compile_and_capture_side_effects(pure_fn, x)
+        self.assertEqual(pure["refs"], ())
+        self.assertFalse(pure["has_side_effects"])
+
+    def test_side_effect_refs_dict_mutation(self):
+        def fn(x, d):
+            d["result"] = x + 1
+            return x * 2
+
+        result = self._compile_and_capture_side_effects(fn, torch.randn(3), {})
+        self.assertEqual(result["refs"], ("L['d']",))
+        self.assertTrue(result["has_side_effects"])
+
+    def test_side_effect_refs_tensor_in_container(self):
+        # Relevant to cudagraphs: a compiled function computes tensors and
+        # stores them into an external container as a side effect.
+        def fn(x, outputs):
+            y = x * 2
+            z = x + 3
+            outputs.append(y)
+            outputs.append(z)
+            return x
+
+        result = self._compile_and_capture_side_effects(fn, torch.randn(4), [])
+        self.assertEqual(result["refs"], ("L['outputs']",))
+        self.assertTrue(result["has_side_effects"])
+
+    def test_side_effect_refs_multiple_containers(self):
+        def fn(x, lst, d):
+            lst.append(x + 1)
+            d["out"] = x * 2
+            return x
+
+        result = self._compile_and_capture_side_effects(fn, torch.randn(3), [], {})
+        self.assertEqual(len(result["refs"]), 2)
+        self.assertIn("L['lst']", result["refs"])
+        self.assertIn("L['d']", result["refs"])
+        self.assertTrue(result["has_side_effects"])
+
     def test_ConstDict_pop_reconstruct(self):
         """
         If something is pop'ed from the dict, we reconstruct everything
@@ -309,16 +392,16 @@ def test_graph_break_in_wrapped_user_function(self):
         def fn(x):
             x = x + 1
             torch._dynamo.graph_break()
-            assert torch.compiler.is_compiling()
-            assert not torch.is_grad_enabled()
+            assert torch.compiler.is_compiling()  # noqa: S101
+            assert not torch.is_grad_enabled()  # noqa: S101
             return x + 2
 
         @torch.compile(backend="eager")
         def gn(x):
             x = torch.no_grad()(fn)(x)
             # reconstruction failure would cause a skipped frame
-            assert torch.compiler.is_compiling()
-            assert torch.is_grad_enabled()
+            assert torch.compiler.is_compiling()  # noqa: S101
+            assert torch.is_grad_enabled()  # noqa: S101
             return x
 
         inp = torch.randn(3)
@@ -333,8 +416,8 @@ def __init__(self):
             def fn(self, x):
                 x = x + self.a
                 torch._dynamo.graph_break()
-                assert torch.compiler.is_compiling()
-                assert not torch.is_grad_enabled()
+                assert torch.compiler.is_compiling()  # noqa: S101
+                assert not torch.is_grad_enabled()  # noqa: S101
                 return x + self.b
 
         obj = Foo()
@@ -344,8 +427,8 @@ def gn(x):
             obj.fn = torch.no_grad()(obj.fn)
             x = obj.fn(x)
             # reconstruction failure would cause a skipped frame
-            assert torch.compiler.is_compiling()
-            assert torch.is_grad_enabled()
+            assert torch.compiler.is_compiling()  # noqa: S101
+            assert torch.is_grad_enabled()  # noqa: S101
             return x
 
         inp = torch.randn(3)
@@ -361,14 +444,14 @@ def gn(x):
             def fn(x):
                 x = x + a
                 torch._dynamo.graph_break()
-                assert torch.compiler.is_compiling()
-                assert not torch.is_grad_enabled()
+                assert torch.compiler.is_compiling()  # noqa: S101
+                assert not torch.is_grad_enabled()  # noqa: S101
                 return x + b
 
             x = fn(x)
             # reconstruction failure would cause a skipped frame
-            assert torch.compiler.is_compiling()
-            assert torch.is_grad_enabled()
+            assert torch.compiler.is_compiling()  # noqa: S101
+            assert torch.is_grad_enabled()  # noqa: S101
             return x
 
         inp = torch.randn(3)
@@ -387,16 +470,16 @@ def test_graph_break_in_wrapped_skipped_function(self):
         def fn(x):
             x = x + 1
             torch._dynamo.graph_break()
-            assert torch.compiler.is_compiling()
-            assert not torch.is_grad_enabled()
+            assert torch.compiler.is_compiling()  # noqa: S101
+            assert not torch.is_grad_enabled()  # noqa: S101
             return x + 2
 
         @torch.compile(backend="eager")
         def gn(x):
             x = torch.no_grad()(_skipped_function_for_test_reconstruct)(fn, x)
             # reconstruction failure would cause a skipped frame
-            assert torch.compiler.is_compiling()
-            assert torch.is_grad_enabled()
+            assert torch.compiler.is_compiling()  # noqa: S101
+            assert torch.is_grad_enabled()  # noqa: S101
             return x
 
         inp = torch.randn(3)
@@ -471,7 +554,7 @@ def fn(x, construct_fn):
             print(l)
             x += 1
             # if reconstruction failed on the graph break, we should error here
-            assert torch.compiler.is_compiling()
+            assert torch.compiler.is_compiling()  # noqa: S101
             return l
 
         @torch.compile(backend="eager", fullgraph=True)
@@ -547,7 +630,7 @@ def hook(instructions: list[dis.Instruction]):
             dup_top_inst = create_dup_top().opname
             for i, inst in enumerate(instructions):
                 if inst.opname == "BUILD_LIST" and i + 2 < len(instructions):
-                    assert not (
+                    assert not (  # noqa: S101
                         instructions[i + 1].opname == dup_top_inst
                         and instructions[i + 2].opname == "STORE_FAST"
                     ), "found list stored as tmp"
@@ -559,6 +642,61 @@ def hook(instructions: list[dis.Instruction]):
                 opt_gn = torch.compile(gn, backend="eager", fullgraph=True)
                 opt_gn(torch.ones(3))
 
+    def test_opaque_reference_as_python_constant(self):
+        """TSOV.as_python_constant must succeed for reference-type opaque
+        objects. Without this, __eq__ between two opaque objects graph breaks.
+        """
+        import torch._library.opaque_object
+        import torch._opaque_base
+
+        class Config(torch._opaque_base.OpaqueBase):
+            def __init__(self, v):
+                self.v = v
+
+            def __bool__(self):
+                return True
+
+            def __eq__(self, other):
+                return isinstance(other, Config) and self.v == other.v
+
+            def __hash__(self):
+                return hash(self.v)
+
+        torch._library.opaque_object.register_opaque_type(Config, typ="reference")
+
+        cfg = Config(42)
+
+        def fn(x, cfg):
+            if cfg:
+                return x + 1
+            return x
+
+        opt = torch.compile(fn, backend="eager", fullgraph=True)
+        result = opt(torch.ones(4), cfg)
+        self.assertEqual(result, torch.ones(4) + 1)
+
+    def test_call_once_guard_allows_super_delegation(self):
+        """_add_call_once_guard must key on (id(self), id(original_method))
+        so that super().as_python_constant() between VT subclasses is not
+        mistaken for a self-referential call.
+        """
+        from torch._dynamo.variables.base import VariableTracker
+
+        class _Parent(VariableTracker):
+            def as_python_constant(self):
+                return 42
+
+        class _Child(_Parent):
+            def as_python_constant(self):
+                return super().as_python_constant()
+
+        child = _Child()
+        # With name-based keying, _Child and _Parent share the same key
+        # (id(self), "as_python_constant"), causing a false
+        # AsPythonConstantNotImplementedError("self-referential").
+        self.assertEqual(child.as_python_constant(), 42)
+        self.assertTrue(child.is_python_constant())
+
 
 if __name__ == "__main__":
     from torch._dynamo.test_case import run_tests
diff --git a/test/dynamo/test_regional_inductor.py b/test/dynamo/test_regional_inductor.py
index a94a42f82cc77..70d24836f3589 100644
--- a/test/dynamo/test_regional_inductor.py
+++ b/test/dynamo/test_regional_inductor.py
@@ -80,7 +80,8 @@ def regional_inductor_pickle(gm, *example_args):
             serialized = GraphPickler.dumps(result)
 
             fake_mode = detect_fake_mode(example_args)
-            assert fake_mode is not None
+            if fake_mode is None:
+                raise AssertionError("Expected fake_mode to be set")
             # Serialize and deserialize the result to confirm pickling works
             # Use a fresh tracing context on the new process
             context = torch._guards.TracingContext(fake_mode)
@@ -314,10 +315,10 @@ def verify_options(*args, **kwargs):
             captured_options.append(options)
 
             # Verify config is set as expected from explicit options
-            assert inductor_config.max_autotune, "max_autotune should be True"
-            assert not inductor_config.triton.cudagraphs, (
-                "triton.cudagraphs should be False"
-            )
+            if not inductor_config.max_autotune:
+                raise AssertionError("max_autotune should be True")
+            if inductor_config.triton.cudagraphs:
+                raise AssertionError("triton.cudagraphs should be False")
 
             return original_compile(*args, **kwargs)
 
@@ -524,8 +525,9 @@ def flex_attn_fn(x):
         )
 
         _, codes = run_fw_bw_and_get_code(lambda: compiled_module(x))
-        # flex in forward and flex_backward in backward
-        self.assertEqual(len(codes), 2)
+        # flex in forward and flex_backward in backward; contiguous partitioning
+        # may split non-contiguous annotated nodes into separate regions
+        self.assertGreaterEqual(len(codes), 2)
 
     def test_refcounts(self):
         """Tests that activations can be cleared before the end of graph"""
@@ -535,10 +537,12 @@ class RefcountCheckPassed(Exception):
 
         def regional_inductor_with_refcounting(gm, *example_args):
             fn = regional_inductor(gm, *example_args)
-            assert fn._boxed_call
+            if not fn._boxed_call:
+                raise AssertionError("Expected fn._boxed_call to be True")
 
             def run(args: Any) -> Any:
-                assert type(args) is list
+                if type(args) is not list:
+                    raise AssertionError(f"Expected args to be list, got {type(args)}")
 
                 # NOTE: sys.getrefcount adds a temporary reference to the object
                 # So sys.getrefcount(x) == 2 actually means we hold the single reference to x
@@ -1042,10 +1046,10 @@ def verify_options(*args, **kwargs):
             captured_options.append(options)
 
             # Verify config is set as expected from explicit options
-            assert torch._inductor.config.max_autotune, "max_autotune should be True"
-            assert not inductor_config.triton.cudagraphs, (
-                "triton.cudagraphs should be False"
-            )
+            if not torch._inductor.config.max_autotune:
+                raise AssertionError("max_autotune should be True")
+            if inductor_config.triton.cudagraphs:
+                raise AssertionError("triton.cudagraphs should be False")
 
             return original_compile(*args, **kwargs)
 
@@ -1311,10 +1315,12 @@ class RefcountCheckPassed(Exception):
 
         def regional_inductor_with_refcounting(gm, *example_args):
             fn = regional_inductor_invoke_subgraph(gm, *example_args)
-            assert fn._boxed_call
+            if not fn._boxed_call:
+                raise AssertionError("Expected fn._boxed_call to be True")
 
             def run(args: Any) -> Any:
-                assert type(args) is list
+                if type(args) is not list:
+                    raise AssertionError(f"Expected args to be list, got {type(args)}")
 
                 # NOTE: sys.getrefcount adds a temporary reference to the object
                 # So sys.getrefcount(x) == 2 actually means we hold the single reference to x
@@ -1519,5 +1525,152 @@ def fn(x):
         self.assertIsNotNone(post_compiled._graph_module)  # Now deserialized
 
 
+class RegionalInductorPartitionTests(torch._inductor.test_case.TestCase):
+    """Tests for _RegionScooper partitioning behavior.
+
+    Uses CapabilityBasedPartitioner per region ID. Nodes with the same region ID
+    are merged as aggressively as possible (only cycles prevent merging). Nodes
+    with different region IDs are never merged.
+    """
+
+    def _make_tag_node(self, g, inp, scalar, tagged):
+        node = g.call_function(torch.ops.aten.mul.Scalar, (inp, scalar))
+        if tagged:
+            node.meta["custom"] = {"compile_with_inductor": True}
+        node.meta["val"] = torch.empty(10)
+        return node
+
+    def _scoop_and_count(self, gm):
+        from torch.fx.passes.regional_inductor import _RegionScooper
+
+        with torch.fx.traceback.preserve_node_meta(enable=False):
+            partitioned = _RegionScooper.scoop_regions(gm)
+        return sum(
+            1
+            for node in partitioned.graph.nodes
+            if node.op == "call_module"
+            and node.target.startswith("__marked_inductor_submod")
+        )
+
+    def _build_chain(self, tags):
+        """Build a linear chain: x -> op0 -> op1 -> ... -> output"""
+        g = torch.fx.Graph()
+        current = g.placeholder("x")
+        for tagged in tags:
+            current = self._make_tag_node(g, current, 1.0, tagged)
+        g.output(current)
+        return torch.fx.GraphModule(torch.nn.Module(), g)
+
+    def test_linear_chain_partitioning(self):
+        """Contiguous runs of tagged nodes form separate partitions."""
+        cases = [
+            # (tags, expected_partitions)
+            ([False, True, True, True, False], 1),
+            ([True, True, False, True, True], 2),
+            ([True, False, True, False, True], 3),
+            ([True, False, True, False, True, False, True], 4),
+            ([False, False, False], 0),
+        ]
+        for tags, expected in cases:
+            with self.subTest(tags=tags):
+                gm = self._build_chain(tags)
+                self.assertEqual(self._scoop_and_count(gm), expected)
+
+    def test_parallel_branches_not_fused(self):
+        """Two adjacent independent tagged branches form 1 partition."""
+        g = torch.fx.Graph()
+        x = g.placeholder("x")
+        mul_a = self._make_tag_node(g, x, 2.0, tagged=True)
+        mul_b = self._make_tag_node(g, x, 3.0, tagged=True)
+        out = g.call_function(torch.ops.aten.add.Tensor, (mul_a, mul_b))
+        out.meta["val"] = torch.empty(10)
+        g.output(out)
+        gm = torch.fx.GraphModule(torch.nn.Module(), g)
+        self.assertEqual(self._scoop_and_count(gm), 1)
+
+    def test_parallel_branches_with_gap_same_region(self):
+        """Two independent tagged nodes separated by an untagged node but
+        sharing the same region ID are fused into 1 partition (no cycle).
+        """
+        g = torch.fx.Graph()
+        x = g.placeholder("x")
+        mul_a = self._make_tag_node(g, x, 2.0, tagged=True)
+        sin = g.call_function(torch.ops.aten.sin.default, (x,))
+        sin.meta["val"] = torch.empty(10)
+        mul_b = self._make_tag_node(g, x, 3.0, tagged=True)
+        out = g.call_function(torch.ops.aten.add.Tensor, (mul_a, mul_b))
+        out.meta["val"] = torch.empty(10)
+        g.output(out)
+        gm = torch.fx.GraphModule(torch.nn.Module(), g)
+        self.assertEqual(self._scoop_and_count(gm), 1)
+
+    def test_parallel_branches_with_gap_different_regions(self):
+        """Two independent tagged nodes with different region IDs produce
+        2 separate partitions regardless of graph topology.
+        """
+        g = torch.fx.Graph()
+        x = g.placeholder("x")
+        mul_a = self._make_tag_node(g, x, 2.0, tagged=True)
+        mul_a.meta["custom"] = {"compile_with_inductor": {"inductor_region": 0}}
+        sin = g.call_function(torch.ops.aten.sin.default, (x,))
+        sin.meta["val"] = torch.empty(10)
+        mul_b = self._make_tag_node(g, x, 3.0, tagged=True)
+        mul_b.meta["custom"] = {"compile_with_inductor": {"inductor_region": 1}}
+        out = g.call_function(torch.ops.aten.add.Tensor, (mul_a, mul_b))
+        out.meta["val"] = torch.empty(10)
+        g.output(out)
+        gm = torch.fx.GraphModule(torch.nn.Module(), g)
+        self.assertEqual(self._scoop_and_count(gm), 2)
+
+    def test_dependent_partitions_merged_across_gap(self):
+        """Two tagged nodes (same region) separated by an untagged node and
+        connected by a data dependency are merged into 1 partition.
+        """
+        g = torch.fx.Graph()
+        x = g.placeholder("x")
+        mul_a = self._make_tag_node(g, x, 2.0, tagged=True)
+        sin = g.call_function(torch.ops.aten.sin.default, (x,))
+        sin.meta["val"] = torch.empty(10)
+        # mul_b consumes mul_a, creating a data dependency across the gap
+        mul_b = self._make_tag_node(g, mul_a, 3.0, tagged=True)
+        g.output(mul_b)
+        gm = torch.fx.GraphModule(torch.nn.Module(), g)
+        self.assertEqual(self._scoop_and_count(gm), 1)
+
+    def test_different_annotations_not_merged(self):
+        """Two tagged nodes with different region IDs are NOT merged,
+        even if connected by a data dependency.
+        """
+        g = torch.fx.Graph()
+        x = g.placeholder("x")
+        mul_a = self._make_tag_node(g, x, 2.0, tagged=True)
+        mul_a.meta["custom"] = {"compile_with_inductor": {"inductor_region": 0}}
+        sin = g.call_function(torch.ops.aten.sin.default, (x,))
+        sin.meta["val"] = torch.empty(10)
+        mul_b = self._make_tag_node(g, mul_a, 3.0, tagged=True)
+        mul_b.meta["custom"] = {"compile_with_inductor": {"inductor_region": 1}}
+        g.output(mul_b)
+        gm = torch.fx.GraphModule(torch.nn.Module(), g)
+        self.assertEqual(self._scoop_and_count(gm), 2)
+
+    def test_chained_merges_across_multiple_gaps(self):
+        """Multiple tagged nodes (same region) with chained data dependencies
+        across untagged gaps are all merged into 1 partition.
+        """
+        g = torch.fx.Graph()
+        x = g.placeholder("x")
+        tagged_a = self._make_tag_node(g, x, 2.0, tagged=True)
+        # Untagged node consuming x (not tagged_a), just occupying a slot
+        filler_1 = g.call_function(torch.ops.aten.sin.default, (x,))
+        filler_1.meta["val"] = torch.empty(10)
+        tagged_b = self._make_tag_node(g, tagged_a, 3.0, tagged=True)
+        filler_2 = g.call_function(torch.ops.aten.sin.default, (x,))
+        filler_2.meta["val"] = torch.empty(10)
+        tagged_c = self._make_tag_node(g, tagged_b, 4.0, tagged=True)
+        g.output(tagged_c)
+        gm = torch.fx.GraphModule(torch.nn.Module(), g)
+        self.assertEqual(self._scoop_and_count(gm), 1)
+
+
 if __name__ == "__main__":
     run_tests()
diff --git a/test/dynamo/test_reorder_logs.py b/test/dynamo/test_reorder_logs.py
index 128211ac48a48..f297f43f4ce37 100644
--- a/test/dynamo/test_reorder_logs.py
+++ b/test/dynamo/test_reorder_logs.py
@@ -302,6 +302,35 @@ def f(x):
  For more details about this graph break, please visit: https://meta-pytorch.github.io/compile-graph-break-site/gb/gb0124.html""",  # noqa: B950
         )
 
+    def test_reorder_print_data_dependent_fstring(self):
+        """Print with data-dependent bool in f-string should graph break on print,
+        but work when print is reorderable."""
+
+        def f(x, mask):
+            make_causal = bool((mask == 0).all())
+            print(f"make_causal={make_causal}")
+            return x + 1
+
+        x = torch.randn(2, 3)
+        mask = torch.zeros(2, 3)
+
+        with self.assertRaisesRegex(
+            torch._dynamo.exc.Unsupported,
+            "Dynamo does not know how to trace builtin operator `print`",
+        ):
+            torch.compile(backend="eager", fullgraph=True)(f)(x, mask)
+
+        with torch._dynamo.config.patch(
+            reorderable_logging_functions={print}, capture_scalar_outputs=True
+        ):
+            opt_f = torch.compile(backend="eager", fullgraph=True)(f)
+            with patch("sys.stdout", new_callable=io.StringIO) as mock_stdout:
+                opt_out = opt_f(x, mask)
+                printed_output = mock_stdout.getvalue().strip()
+
+        self.assertTrue(same(opt_out, x + 1))
+        self.assertEqual(printed_output, "make_causal=True")
+
 
 if __name__ == "__main__":
     from torch._dynamo.test_case import run_tests
diff --git a/test/dynamo/test_repros.py b/test/dynamo/test_repros.py
index 65881c21b93c6..b63da1994dfa3 100644
--- a/test/dynamo/test_repros.py
+++ b/test/dynamo/test_repros.py
@@ -56,7 +56,11 @@
 )
 from torch._inductor.utils import fresh_cache
 from torch.nn import functional as F
-from torch.nn.attention.flex_attention import create_block_mask, flex_attention
+from torch.nn.attention.flex_attention import (
+    AuxRequest,
+    create_block_mask,
+    flex_attention,
+)
 from torch.profiler import profile, ProfilerActivity
 from torch.testing._internal.common_cuda import (
     PLATFORM_SUPPORTS_FLASH_ATTENTION,
@@ -187,7 +191,8 @@ def global_fn(x):
 
 def cat(tensors, dim=0):
     # from detectron2 wrappers.py
-    assert isinstance(tensors, (list, tuple))
+    if not isinstance(tensors, (list, tuple)):
+        raise AssertionError(f"Expected list or tuple, got {type(tensors)}")
     if len(tensors) == 1:
         return tensors[0]
     return torch.cat(tensors, dim)
@@ -198,9 +203,8 @@ def shapes_to_tensor(x, device=None):
     if torch.jit.is_scripting():
         return torch.as_tensor(x, device=device)
     if torch.jit.is_tracing():
-        assert all(isinstance(t, torch.Tensor) for t in x), (
-            "Shape should be tensor during tracing!"
-        )
+        if not all(isinstance(t, torch.Tensor) for t in x):
+            raise AssertionError("Shape should be tensor during tracing!")
         # as_tensor should not be used in tracing because it records a constant
         ret = torch.stack(x)
         if ret.device != device:  # avoid recording a hard-coded device if not necessary
@@ -250,7 +254,10 @@ def __init__(self, tensor: torch.Tensor):
             # Use reshape, so we don't end up creating a new tensor that does not depend on
             # the inputs (and consequently confuses jit)
             tensor = tensor.reshape((-1, 4)).to(dtype=torch.float32, device=device)
-        assert tensor.dim() == 2 and tensor.size(-1) == 4, tensor.size()
+        if not (tensor.dim() == 2 and tensor.size(-1) == 4):
+            raise AssertionError(
+                f"Expected 2D tensor with last dim 4, got {tensor.size()}"
+            )
         self.tensor = tensor
 
     def __len__(self) -> int:
@@ -499,9 +506,10 @@ def forward(
         real_seq_length = seq_length
 
         if past_key_value is not None:
-            assert len(past_key_value) == 2, (
-                f"past_key_value should have 2 past states: keys and values. Got {len(past_key_value)} past states"
-            )
+            if len(past_key_value) != 2:
+                raise AssertionError(
+                    f"past_key_value should have 2 past states: keys and values. Got {len(past_key_value)} past states"
+                )
             real_seq_length += (
                 past_key_value[0].shape[2] if query_length is None else query_length
             )
@@ -579,9 +587,11 @@ def forward_chunk(self, hidden_states):
 
 def apply_chunking_to_forward(forward_fn, *input_tensors):
     # simplified from HF model_utils.py
-    assert len(input_tensors) > 0
+    if len(input_tensors) == 0:
+        raise AssertionError("Expected at least one input tensor")
     tensor_shape = input_tensors[0].shape[1]
-    assert all(input_tensor.shape[1] == tensor_shape for input_tensor in input_tensors)
+    if not all(input_tensor.shape[1] == tensor_shape for input_tensor in input_tensors):
+        raise AssertionError("All input tensors must have the same shape[1]")
     num_args_in_forward_chunk_fn = len(inspect.signature(forward_fn).parameters)
     if num_args_in_forward_chunk_fn != len(input_tensors):
         raise ValueError
@@ -1016,12 +1026,12 @@ def tearDown(self) -> None:
     def test_compiled_module_truthiness(self):
         # Test with empty ModuleList
         original_empty = nn.ModuleList()
-        compiled_empty = torch.compile(original_empty)
+        compiled_empty = torch.compile(original_empty, backend="eager")
         self.assertEqual(bool(original_empty), bool(compiled_empty))
         self.assertFalse(bool(compiled_empty))
         # Test with non-empty ModuleList
         original_filled = nn.ModuleList([nn.Linear(10, 5)])
-        compiled_filled = torch.compile(original_filled)
+        compiled_filled = torch.compile(original_filled, backend="eager")
         self.assertEqual(bool(original_filled), bool(compiled_filled))
         self.assertTrue(bool(compiled_filled))
 
@@ -1034,6 +1044,40 @@ def guard_manager_clone_hook_fn(self, guard_manager_wrapper, f_locals, builder):
         if guard_manager_wrapper.diff_guard_root:
             self.assertTrue(guard_manager_wrapper.diff_guard_root.check(f_locals))
 
+    def test_swap_tensors_subclass_not_blocked_by_metaconverter_refcycle(self):
+        """Checks that MetaConverter doesn't create refcycles of itself that blocks swap_tensor on subclass tensors"""
+        was_gc_enabled = gc.isenabled()
+        gc.disable()
+        try:
+            from torch._subclasses.meta_utils import MetaConverter
+            from torch.utils.weak import WeakIdRef
+
+            w_a = torch.randn(4, 4)
+            w_b = w_a.clone()
+            # For custom tensor parameters, `nn.Parameter` returns the same custom tensor type
+            # and sets `_is_param`, which is what Module._apply uses to decide swap_tensors.
+            w = nn.Parameter(TwoTensor(w_a, w_b).requires_grad_())
+
+            def weak_id_ref_count(t: torch.Tensor) -> int:
+                return sum(isinstance(r, WeakIdRef) for r in weakref.getweakrefs(t))
+
+            baseline_weak_id_refs = weak_id_ref_count(w)
+            converter = MetaConverter()
+            _ = converter(w)
+            after_convert_weak_id_refs = weak_id_ref_count(w)
+            self.assertGreater(after_convert_weak_id_refs, baseline_weak_id_refs)
+
+            # If MetaConverter is collectable by refcount, the weakrefs go away immediately.
+            # If it's kept alive only by a refcycle, this weakref will persist while gc is
+            # disabled.
+            del _
+            del converter
+            self.assertEqual(weak_id_ref_count(w), baseline_weak_id_refs)
+        finally:
+            if was_gc_enabled:
+                gc.enable()
+            gc.collect()
+
     def test_do_paste_mask(self):
         torch._dynamo.utils.counters.clear()
         cnt = torch._dynamo.testing.CompileCounter()
@@ -1166,6 +1210,33 @@ def fn(x):
         self.assertEqual(eager_out, compiled_out)
         self.assertEqual(eager_out.stride(), compiled_out.stride())
 
+    # https://github.com/pytorch/pytorch/issues/174963
+    def test_roll_as_strided_channels_last(self):
+        def fn(x):
+            x = torch.roll(x, shifts=-1, dims=0)
+            return torch.as_strided(x, (5, 10), (4, 1), 0)
+
+        x = torch.arange(100).reshape(2, 5, 2, 5).to(memory_format=torch.channels_last)
+        eager_out = fn(x)
+        compiled_fn = torch.compile(fn, backend="aot_eager_decomp_partition")
+        compiled_out = compiled_fn(x)
+        self.assertEqual(eager_out, compiled_out)
+
+    # https://github.com/pytorch/pytorch/issues/166626
+    def test_inplace_add_from_meta_tensor_factory(self):
+        def fn(x):
+            log_det = torch.zeros(x.size(0), device=x.device)
+            log_det += torch.zeros(x.size(0), device="meta")
+            return log_det
+
+        x = torch.randn(2, 4)
+        eager_out = fn(x)
+        compiled_fn = torch.compile(fn, backend="eager", fullgraph=True)
+        compiled_out = compiled_fn(x)
+
+        self.assertEqual(eager_out.device, compiled_out.device)
+        self.assertEqual(eager_out, compiled_out)
+
     # https://github.com/pytorch/pytorch/issues/109053
     def test_view_dtype_overload(self):
         def f(x):
@@ -1298,7 +1369,7 @@ def forward(self, x) -> torch.Tensor:
 
         x = {MyEnum.A: torch.rand(8, 1)}
         model_pytorch = SomeModel()
-        model = torch.compile(model_pytorch)
+        model = torch.compile(model_pytorch, backend="eager")
         # Executing twice works
         model(x)
         y = model(x)
@@ -1332,12 +1403,8 @@ def test_reformer_eval(self):
     def test_reformer_train(self):
         with torch.enable_grad():
             cnt = self._reformer(nopython=False)
-        expected_op_count = (
-            """10""" if torch._dynamo.config.inline_inbuilt_nn_modules else """4"""
-        )
-
         self.assertExpectedInline(cnt.frame_count, """1""")
-        self.assertExpectedInline(cnt.op_count, expected_op_count)
+        self.assertExpectedInline(cnt.op_count, """10""")
 
     def test_longformer_chunk(self):
         input1 = torch.randn([1, 4096, 1])
@@ -1636,7 +1703,9 @@ def test_get_parameter_dtype(self):
         )
 
         def fn(model, x):
-            return x + torch.randn(10, dtype=get_parameter_dtype(model))
+            return x + torch.randn(
+                10, dtype=get_parameter_dtype(model), device=x.device
+            )
 
         cnt = torch._dynamo.testing.CompileCounter()
         opt_fn = torch._dynamo.optimize_assert(cnt)(fn)
@@ -1661,7 +1730,7 @@ def test_fn():
             a = torch.randn(4)
             x = torch.Size([1, 2, 3])
             # Checks that SizeVariable return torch.Size object
-            assert isinstance(x, torch.Size)
+            assert isinstance(x, torch.Size)  # noqa: S101
             # Causes graph breaks and checks reconstruction of SizeVariable
             # object
             self.assertIsInstance(x, torch.Size)
@@ -1699,7 +1768,7 @@ def test_indexing_with_list(self):
         def test_fn():
             def run_test(tensor, *idx):
                 npt = tensor.numpy()
-                assert npt[idx].shape == tensor[idx].shape
+                assert npt[idx].shape == tensor[idx].shape  # noqa: S101
 
             x = torch.arange(0, 10)
             cases = [
@@ -1719,13 +1788,13 @@ def run_test(tensor, *idx):
     def test_foreach_decomp_arg_names(self):
         # https://github.com/pytorch/pytorch/issues/138698
 
-        @torch.compile(fullgraph=True)
+        @torch.compile(fullgraph=True, backend="eager")
         def foreach_pow(**kwargs):
             return torch._foreach_pow(**kwargs)
 
         foreach_pow(self=[torch.ones(2, 2, device="cpu")], exponent=2.7)
 
-        @torch.compile(fullgraph=True)
+        @torch.compile(fullgraph=True, backend="eager")
         def foreach_lerp_(**kwargs):
             return torch._foreach_lerp_(**kwargs)
 
@@ -1941,6 +2010,38 @@ def fn(x):
         self.assertFalse(y.requires_grad)
         self.assertFalse(z.requires_grad)
 
+    def test_locals_traced_correctly_under_compile(self):
+        def fn(x):
+            if x.dim() > 2:
+                batch_size, seq_len, hidden_dim = x.shape
+                x = x.view(-1, hidden_dim)
+
+            x = x + 1
+
+            if "batch_size" in locals() and "seq_len" in locals():
+                x = x.view(batch_size, seq_len, -1)
+            return x
+
+        x = torch.randn(2, 3, 4)
+        opt_fn = torch.compile(fn, backend="eager")
+        self.assertTrue(same(fn(x), opt_fn(x)))
+
+    def test_vars_traced_correctly_under_compile(self):
+        def fn(x):
+            if x.dim() > 2:
+                batch_size, seq_len, hidden_dim = x.shape
+                x = x.view(-1, hidden_dim)
+
+            x = x + 1
+
+            if "batch_size" in vars() and "seq_len" in vars():
+                x = x.view(batch_size, seq_len, -1)
+            return x
+
+        x = torch.randn(2, 3, 4)
+        opt_fn = torch.compile(fn, backend="eager")
+        self.assertTrue(same(fn(x), opt_fn(x)))
+
     def test_abc_setattr(self):
         # tests that we correctly bail out of __setattr__ calls
 
@@ -2010,7 +2111,7 @@ def check_equal(condition, x, y):
                     # NumPy aggressively promotes to double, hence cast to output to correct dtype
                     expected = get_expected(condition, x, y)
                     result = torch.where(condition, x, y)
-                    assert torch.allclose(expected, result)
+                    assert torch.allclose(expected, result)  # noqa: S101
 
                 check_equal(condition, x, y)
                 check_equal(condition, y, x)
@@ -2251,7 +2352,7 @@ def test_isinstance_storage(self):
         def fn(x):
             f = bytearray([0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x10, 0x40])
             bools = torch.BoolStorage.from_buffer(f, "big")
-            assert isinstance(bools, torch.BoolStorage)
+            assert isinstance(bools, torch.BoolStorage)  # noqa: S101
             return x
 
         fn(torch.randn(3))
@@ -2725,7 +2826,7 @@ def f(x):
 
         x = torch.ones(2, requires_grad=True)
 
-        f_compiled = torch.compile(f)
+        f_compiled = torch.compile(f, backend="eager")
         with torch.no_grad():
             # compile an inference graph
             f_compiled(x)
@@ -2742,7 +2843,7 @@ def test_requires_grad_guards_with_grad_mode2(self):
         x_ref = x.detach().clone().requires_grad_(True)
 
         m = torch.nn.Linear(2, 2)
-        m_compiled = torch.compile(m)
+        m_compiled = torch.compile(m, backend="eager")
 
         with torch.no_grad():
             # compile an inference graph
@@ -3125,7 +3226,7 @@ def f(x):
 
     def test_error_return_without_exception_set(self):
         # https://github.com/pytorch/pytorch/issues/93781
-        @torch.compile
+        @torch.compile(backend="eager")
         def f():
             _generator_type = type(_ for _ in ())
 
@@ -3224,7 +3325,7 @@ def fn(a: CustomList2, b: CustomList2):
     def test_rewrite_assert_with_msg(self):
         def f(x):
             b = x.sin()
-            assert x[0] == 3, "First dim need to be 3"
+            assert x[0] == 3, "First dim need to be 3"  # noqa: S101
             return x.cos() + b
 
         args = (torch.Tensor([3, 4, 5]),)
@@ -3270,7 +3371,7 @@ def f(x):
     def test_rewrite_assert_dont_change_bytecode(self):
         def fn(x):
             with torch.no_grad():
-                assert x.max() < 5, f"invalid max {x.max()}"
+                assert x.max() < 5, f"invalid max {x.max()}"  # noqa: S101
                 x = torch.sin(x)
             return x
 
@@ -3281,7 +3382,7 @@ def fn(x):
     def test_rewrite_assert_without_msg(self):
         def f(x):
             b = x.sin()
-            assert x[0] == 3
+            assert x[0] == 3  # noqa: S101
             return x.cos() + b
 
         args = (torch.Tensor([3, 4, 5]),)
@@ -3294,7 +3395,7 @@ def f(x):
     def test_rewrite_assert_with_non_string_msg(self):
         def f(x):
             b = x.sin()
-            assert x[0] == 2, x
+            assert x[0] == 2, x  # noqa: S101
             return x.cos() + b
 
         torch._dynamo.utils.counters.clear()
@@ -3313,8 +3414,8 @@ def f(x):
     def test_rewrite_assert_noop(self):
         def f(x):
             b = x.sin()
-            assert True
-            assert x.dtype == torch.float32
+            assert True  # noqa: S101
+            assert x.dtype == torch.float32  # noqa: S101
             return x.cos() + b
 
         args = (torch.Tensor([3, 4, 5]),)
@@ -3726,7 +3827,7 @@ def fn(x, inp_list):
         ref2 = fn(x, inp_list2)
         ref3 = fn(x, inp_list3)
 
-        opt_fn = torch.compile(fn, fullgraph=True)
+        opt_fn = torch.compile(fn, fullgraph=True, backend="eager")
 
         opt_ret1 = opt_fn(x, inp_list1)
         opt_ret2 = opt_fn(x, inp_list2)
@@ -3987,7 +4088,7 @@ def f(x, y1, y2):
         f(torch.zeros(4), float, np.float16)
 
     def test_dedup_global(self):
-        @torch.compile()
+        @torch.compile(backend="eager")
         def f():
             return _GLOBAL_CPU_TENSOR + _GLOBAL_CPU_TENSOR
 
@@ -3997,7 +4098,7 @@ def test_randint_out_dynamic(self):
         def randint_fn(high, size, out):
             return torch.randint(high, size, out=out)
 
-        opt_model = torch.compile(randint_fn)
+        opt_model = torch.compile(randint_fn, backend="eager")
 
         out1 = torch.empty(10, dtype=torch.int32)
         opt_model(17, (10,), out1)
@@ -4100,7 +4201,7 @@ def fn(inp1, inp2, inp3, inp4, c):
         cnt = torch._dynamo.testing.CompileCounter()
         opt_fn = torch.compile(fn, backend=cnt)
         opt_fn(inp1, inp2, inp3, inp4, c)
-        self.assertEqual(cnt.frame_count, 3)
+        self.assertEqual(cnt.frame_count, 2)
 
     def test_torch_variable_type(self):
         # from torchvision
@@ -4273,16 +4374,18 @@ def fn(x):
         self.assertTrue(same(fn(x), opt_fn(x)))
 
     def test_add_sub_alpha_out(self):
-        inp = torch.randn(2, 3, 4)
-        other = 1
-        alpha = 2
+        test_cases = (
+            (torch.randn(2, 3, 4), 1, 2, torch.zeros(2, 3, 4)),
+            (2, 1.1, 0.4, torch.tensor(0.0)),
+        )
         for op in [torch.add, torch.sub]:
-            out = torch.zeros(2, 3, 4)
-            compile_out = torch.zeros(2, 3, 4)
-            op(inp, other, alpha=alpha, out=out)
-            compiled_fn = torch.compile(op, dynamic=True)
-            compiled_fn(inp, other, alpha=alpha, out=compile_out)
-            self.assertTrue(same(out, compile_out))
+            for inp, other, alpha, out in test_cases:
+                compiled_fn = torch.compile(op, dynamic=True, backend="eager")
+                eager_out = out.clone()
+                compiled_out = out.clone()
+                op(inp, other, alpha=alpha, out=eager_out)
+                compiled_fn(inp, other, alpha=alpha, out=compiled_out)
+                self.assertTrue(same(eager_out, compiled_out))
 
     def test_negative_shape_guard(self):
         def fn(x):
@@ -4301,7 +4404,7 @@ def fn(x):
 
     @torch._dynamo.config.patch(capture_scalar_outputs=True)
     def test_deferred_runtime_asserts(self):
-        @torch.compile(fullgraph=True)
+        @torch.compile(fullgraph=True, backend="eager")
         def f(x):
             y = x.item()
             torch._check(y >= 0)
@@ -4324,7 +4427,7 @@ def test_addr_alpha_beta_out(self):
         compile_out = torch.zeros(2, 3)
 
         torch.addr(inp, vec1, vec2, alpha=alpha, beta=beta, out=out)
-        compiled_fn = torch.compile(torch.addr, dynamic=True)
+        compiled_fn = torch.compile(torch.addr, dynamic=True, backend="eager")
         compiled_fn(inp, vec1, vec2, alpha=alpha, beta=beta, out=compile_out)
         self.assertTrue(same(out, compile_out))
 
@@ -4369,16 +4472,6 @@ def fn(x):
         self.assertEqual(fn(x), opt_fn(x))
         self.assertEqual(fn(x2), opt_fn(x2))
 
-    def test_inductor_no_recursionerror_on_for_loops(self):
-        def forward(x):
-            for _ in range(10000):
-                x = 1.0 * x
-            return x
-
-        self.assertTrue(
-            same(torch.compile(forward)(torch.tensor([1.0])), torch.tensor([1.0]))
-        )
-
     def test_user_defined_object_callable(self):
         # https://github.com/pytorch/pytorch/issues/114019
         class MyCallable:
@@ -4788,11 +4881,13 @@ def forward(self, x):
         model = SimpleModel().eval()
         input_tensor = torch.randn(1, 10, dtype=torch.float32)
         opt = torch.compile(model.eval(), backend="eager", fullgraph=True)
-        actual = opt(input_tensor)
         try:
             expected = model(input_tensor)
         except Exception as e:
-            raise unittest.SkipTest("eager failed, requires Python>=3.12") from e
+            raise unittest.SkipTest(
+                "eager failed, requires Python between 3.9 and 3.12"
+            ) from e
+        actual = opt(input_tensor)
         self.assertEqual(actual, expected)
 
     def test_invalid_seq_unpack(self):
@@ -4803,7 +4898,7 @@ def fn():
             return myfn((1, 2, 3))
 
         try:
-            torch.compile(fn)()
+            torch.compile(fn, backend="eager")()
         except ValueError:
             pass
         else:
@@ -4958,10 +5053,10 @@ def test_dont_aggressively_write_assert(self):
 
         @torch.compile(dynamic=True, backend=record_graph)
         def f(x):
-            assert x.shape[0] > 3
-            assert x[0].sum() > 0
-            assert 1 % (x.shape[0] // 2) != 0
-            assert 32 * (x.shape[0] // 2) ** 2 - 16 * (x.shape[0] // 2) != 0
+            assert x.shape[0] > 3  # noqa: S101
+            assert x[0].sum() > 0  # noqa: S101
+            assert 1 % (x.shape[0] // 2) != 0  # noqa: S101
+            assert 32 * (x.shape[0] // 2) ** 2 - 16 * (x.shape[0] // 2) != 0  # noqa: S101
             return x.cos()
 
         f(torch.ones(6, 4))
@@ -4996,7 +5091,7 @@ def forward(self, s77 : torch.SymInt, s27 : torch.SymInt, L_x_ : torch.Tensor):
 
         @torch.compile(dynamic=True, backend=record_graph)
         def f_fail(x):
-            assert x.shape[0] < 3
+            assert x.shape[0] < 3  # noqa: S101
 
         # We graph-break here, so the failure should be eager
         with self.assertRaisesRegex(AssertionError, ""):
@@ -5037,7 +5132,7 @@ def set(self, name: str, value: Any) -> None:
                 with warnings.catch_warnings(record=True):
                     data_len = len(value)
                 if len(self._fields):
-                    assert len(self) == data_len, (
+                    assert len(self) == data_len, (  # noqa: S101
                         f"Adding a field of length {data_len} to a Instances of length {len(self)}"
                     )
                 self._fields[name] = value
@@ -5047,8 +5142,8 @@ def get(self, name: str) -> Any:
 
             @staticmethod
             def cat(instance_lists: list["Instances"]) -> "Instances":
-                assert all(isinstance(i, Instances) for i in instance_lists)
-                assert len(instance_lists) > 0
+                assert all(isinstance(i, Instances) for i in instance_lists)  # noqa: S101
+                assert len(instance_lists) > 0  # noqa: S101
                 if len(instance_lists) == 1:
                     return instance_lists[0]
 
@@ -5057,7 +5152,7 @@ def cat(instance_lists: list["Instances"]) -> "Instances":
                     image_size, torch.Tensor
                 ):  # could be a tensor in tracing
                     for i in instance_lists[1:]:
-                        assert i.image_size == image_size
+                        assert i.image_size == image_size  # noqa: S101
                 ret = Instances(image_size)
                 for k in instance_lists[0]._fields:
                     values = [i.get(k) for i in instance_lists]
@@ -5334,17 +5429,11 @@ def fn(obj):
 
         obj = A()
 
-        try:
+        with self.assertRaisesRegex(RuntimeError, r"super\(\)"):
             fn(obj)
-        except Exception as e:
-            orig_str = str(e)
-        self.assertIn("no arguments", orig_str)
 
-        try:
+        with self.assertRaisesRegex(RuntimeError, r"super\(\)"):
             torch.compile(backend="eager")(fn)(obj)
-        except Exception as e:
-            compiled_str = str(e)
-        self.assertEqual(orig_str, compiled_str)
 
     def test_super_staticmethod(self):
         class Parent:
@@ -5455,7 +5544,7 @@ def fn(x):
         self.assertEqual(ref, res)
 
     def test_vc_bumped_in_inference_graph(self):
-        @torch.compile
+        @torch.compile(backend="eager")
         def f(x):
             return x.mul_(2)
 
@@ -5473,7 +5562,7 @@ def forward(self, x):
         def f(m):
             return callable(m)
 
-        res = torch.compile(f, fullgraph=True)(M())
+        res = torch.compile(f, fullgraph=True, backend="eager")(M())
         self.assertTrue(res)
 
     def test_stk_sdd_is_transposed(self):
@@ -5648,7 +5737,7 @@ def new_fn(x):
     def test_list_reverse(self):
         def ladder(x):
             trail = x.size(-1)
-            assert trail > 2
+            assert trail > 2  # noqa: S101
             weights = []
             for s in [trail, trail - 1, trail - 2]:
                 weights.append(torch.ones(s, s - 1))
@@ -5732,9 +5821,9 @@ def fn2():
             return torch.mm(x, y).sum()
 
         with fresh_cache():
-            torch.compile(fn)()
+            torch.compile(fn, backend="eager")()
 
-        torch.compile(fn2)()
+        torch.compile(fn2, backend="eager")()
 
     def test_jit_script_defaults(self):
         @torch.jit.script
@@ -5844,10 +5933,9 @@ def forward(self, x):
         opt_mod = torch.compile(mod, backend=compiler)
         opt_mod(torch.randn(2, 2))
 
-        with torch._dynamo.config.patch(inline_inbuilt_nn_modules=True):
-            mod = Mod()
-            opt_mod = torch.compile(mod, backend=compiler)
-            opt_mod(torch.randn(2, 2))
+        mod = Mod()
+        opt_mod = torch.compile(mod, backend=compiler)
+        opt_mod(torch.randn(2, 2))
 
         # an example similar to Pippy usecase
         mod = Mod()
@@ -5890,6 +5978,34 @@ def test_export_vs_dynamo_for_multiheadattention(self):
             self.assertEqual(len(compile_nodes), 0)
             self.assertEqual(len(export_nodes), 0)
 
+    def test_multiheadattention_tracing_slowpath_matches_fastpath_layout(self):
+        class MHAWithView(nn.Module):
+            def __init__(self) -> None:
+                super().__init__()
+                self.hidden_dim = 64
+                self.attention = nn.MultiheadAttention(
+                    self.hidden_dim, 8, batch_first=True
+                )
+
+            def forward(self, x):
+                attn_output, _ = self.attention(x, x, x)
+                return attn_output.view(-1, self.hidden_dim)
+
+        with torch.no_grad():
+            model = MHAWithView().eval()
+            x = torch.randn(4, 32, model.hidden_dim)
+            eager = model(x)
+
+            backend = EagerAndRecordGraphs()
+            compiled_model = torch.compile(model, backend=backend, fullgraph=True)
+            compiled = compiled_model(x)
+
+            compile_nodes = backend.graphs[0].graph.find_nodes(
+                op="call_function", target=torch._native_multi_head_attention
+            )
+            self.assertEqual(compiled, eager)
+            self.assertEqual(len(compile_nodes), 0)
+
     def test_negative_floor_div_solve(self):
         class CompiledClass(nn.Module):
             def __init__(self) -> None:
@@ -5916,7 +6032,7 @@ def random_op(tensor, args, kwargs):
             res = tensor.random_(*args, **kwargs)
             return res
 
-        random_op = torch.compile(random_op)
+        random_op = torch.compile(random_op, backend="eager")
         tensor = torch.randn([2, 3])
         random_op(tensor, [], {"from": -10, "to": 10})
         random_op(tensor, [-10], {"to": 10})
@@ -5928,7 +6044,7 @@ def uniform_op(tensor, args, kwargs):
             res = tensor.uniform_(*args, **kwargs)
             return res
 
-        uniform_op = torch.compile(uniform_op)
+        uniform_op = torch.compile(uniform_op, backend="eager")
         tensor = torch.randn([2, 3])
         uniform_op(tensor, [], {"from": -10, "to": 10})
         uniform_op(tensor, [-10], {"to": 10})
@@ -6040,7 +6156,7 @@ def f(x):
             return x + x.conj()
 
         x = torch.randn(4, dtype=torch.complex64, requires_grad=True)
-        out = torch.compile(f)(x)
+        out = torch.compile(f, backend="eager")(x)
         expected_complex = (2 * x.real).to(dtype=out.dtype)
 
         self.assertTrue(out.dtype == torch.complex64)
@@ -6105,13 +6221,44 @@ def f(x, l):
         self.assertEqual(result, result_test)
         self.assertEqual(x, x_test)
 
+    # https://github.com/pytorch/pytorch/issues/167009
+    def test_inbuilt_nn_module_forward_after_hook_graph_break(self):
+        # When a hook causes a graph break on an inbuilt nn.Module, the module's
+        # forward should still be traced after the graph break.
+
+        @torch._dynamo.disable
+        def my_hook(module, inp):
+            return inp
+
+        class Wrapper(nn.Module):
+            def __init__(self, lin):
+                super().__init__()
+                self.lin = lin
+
+            def forward(self, x):
+                return self.lin(x)
+
+        lin = nn.Linear(10, 5)
+        lin.register_forward_pre_hook(my_hook)
+        model = Wrapper(lin)
+
+        backend = EagerAndRecordGraphs()
+        torch._dynamo.reset()
+        compiled = torch.compile(model, backend=backend)
+        output = compiled(torch.randn(3, 10))
+
+        self.assertEqual(output.shape, torch.Size([3, 5]))
+        self.assertEqual(len(backend.graphs), 1)
+        graph_code = backend.graphs[0].print_readable(print_output=False)
+        self.assertIn("torch._C._nn.linear", graph_code)
+
     def test_aot_autograd_runtime_wrapper_prologue_profiled(self):
         # Names for prologue profiling event
         prologue_name = "AOTDispatcher Runtime Wrapper Prologue"
 
         # Simple linear op to compile
         mod = torch.nn.Linear(4, 4)
-        opt_mod = torch.compile(mod)
+        opt_mod = torch.compile(mod, backend="aot_eager")
         x = torch.randn(4, 4)
 
         # Run this test with grad and no-grad to test both boolean cases trace_joint
@@ -6182,96 +6329,6 @@ def fn(aot6_sub_58, aot6_mul_170):
         # No assert necessary since this used to crash.
         fn(aot6_sub_58, aot6_mul_170)
 
-    @torch._dynamo.config.patch(guard_nn_modules=False)
-    @torch._dynamo.config.patch(inline_inbuilt_nn_modules=False)
-    def test_inlining_cornercase(self):
-        """
-        nn.Modules can be mapped to either NNModuleVariable or UnspecializedNNModuleVariable. For NNModuleVariable, the
-        tensor attributes become part of the Dynamo graph. For unspecialized, they are lifted as inputs.
-
-        But there is a cornercase. Suppose you have NNModuleVariable with a submodule that is
-        UnspecializedNNModuleVariable. Today, Dynamo will still consider the submodule as specialized (courtesy of
-        guard.source().is_nn_module()). In retrospect, this is a mistake but there are dependencies of export and also
-        cudagraphs which make it harder to fix the corner case right away. The long term solution is
-        inline_inbuilt_nn_modules anyways, so we might have to live with this cornercase in the short term.
-
-        We are starting to annotate the source of each nn module more precisely - NNModuleVariable attribute is marked
-        as NNModuleSource, UnspecilaizedNNModuleVariable attribute is marked as UnspecializedNNModuleSource. But this
-        changes the behavior for the cornercase. And fails some tests which have unfortunately relied on this behavior.
-
-
-        To solve this, we tag the source only when inline_inbuilt_nn_module flag is turned on.
-
-        In this test, we purposely turn the flag off, testing that the tagging is disabled.
-        """
-
-        class SubMod(torch.nn.Module):
-            def __init__(self):
-                super().__init__()
-                self.linear = torch.nn.Linear(1, 1)
-                self.a = torch.randn(1, 1)
-                self.counter = 0
-                self.multipliers = [2.2, 3.3]
-
-            def forward(self, x):
-                self.counter += 1
-                return (
-                    self.linear(x) * self.a * self.multipliers[0] * self.multipliers[1]
-                )
-
-        class Mod(torch.nn.Module):
-            def __init__(self):
-                super().__init__()
-                self.submod = SubMod()
-
-            def forward(self, x):
-                return self.submod(x)
-
-        mod = Mod()
-        opt_mod = torch.compile(mod, backend="eager")
-
-        x = torch.randn(1, 1)
-        ref = mod(x)  # noqa: F841
-        res = opt_mod(x)  # noqa: F841
-
-        mod.submod.multipliers = [3.3, 4.4]
-        # Since guard_nn_modules is False, this will not recompile
-        with torch._dynamo.config.patch(error_on_recompile=True):
-            ref = mod(x)  # noqa: F841
-            res = opt_mod(x)  # noqa: F841
-
-    @torch._dynamo.config.patch(inline_inbuilt_nn_modules=False)
-    def test_nnmodule_variable_children_wrap_value(self):
-        """
-        tests wrap_values() in nn_module.py by calling children() on a submodule,
-        which triggers NNModuleVariable.call_method only when
-        inline_inbuilt_nn_modules=False.  This path was previously untested
-        causing #173924
-        """
-
-        class Parent(torch.nn.Module):
-            def __init__(self):
-                super().__init__()
-                self.container = torch.nn.Sequential(
-                    torch.nn.Linear(10, 10),
-                    torch.nn.ReLU(),
-                    torch.nn.Linear(10, 5),
-                )
-
-            def forward(self, x):
-                for child in self.container.children():
-                    x = child(x)
-                return x
-
-        model = Parent()
-        x = torch.randn(2, 10)
-
-        eager_result = model(x)
-        compiled_model = torch.compile(model, backend="eager", fullgraph=True)
-        compiled_result = compiled_model(x)
-
-        self.assertEqual(eager_result, compiled_result)
-
     def test_optimized_module_training(self):
         mod = torch.nn.Linear(3, 3)
         mod.eval()
@@ -6312,7 +6369,7 @@ def patch_init_for_class(cls):
 
         patch_init_for_class(MyModule)
         mod = MyModule()
-        opt_mod = torch.compile(mod)
+        opt_mod = torch.compile(mod, backend="eager")
 
         x = torch.rand(10)
         ref = mod(x)
@@ -6559,35 +6616,6 @@ def outer_func(x):
         res = compile_outer(x)
         self.assertEqual(ref, res)
 
-    # https://github.com/pytorch/pytorch/issues/136640
-    def test_inductor_dynamic_shapes_broadcasting(self) -> None:
-        def fn(x, y):
-            x_view = x.view(-1, 4)
-            y_view = y.view(-1, 4)
-            return x_view * y_view
-
-        x = torch.randn(4)
-        y = torch.randn(8)
-        out_ref = fn(x, y)
-        out_test = torch.compile(fn, dynamic=True)(x, y)
-        self.assertEqual(out_ref, out_test)
-
-    # https://github.com/pytorch/pytorch/issues/119162
-    def test_inductor_rng_default_dtype(self) -> None:
-        @torch.compile
-        def fn():
-            tmp = torch.randn(4, 4, dtype=torch.bfloat16)
-            return tmp
-
-        try:
-            old = torch.get_default_dtype()
-            torch.set_default_dtype(torch.bfloat16)
-            out = fn()
-        finally:
-            torch.set_default_dtype(old)
-        # output dtype should be float32
-        self.assertEqual(out.dtype, torch.bfloat16)
-
     @unittest.skipIf(not HAS_MSGSPEC, "missing msgspec package")
     def test_c_defined_metaclass(self):
         class User(msgspec.Struct):
@@ -6779,7 +6807,7 @@ def make_dist_and_execute(t, d):
     def test_bitwise_print_precedence(self):
         import math
 
-        @torch.compile(fullgraph=True, dynamic=True)
+        @torch.compile(fullgraph=True, dynamic=True, backend="eager")
         def f(x):
             torch._check(math.floor((x.size(0) | 3) * 4) == 12)
             return x.sin()
@@ -6787,7 +6815,7 @@ def f(x):
         f(torch.randn(2))
 
     def test_tensor_split_within_device_cm(self):
-        @torch.compile(fullgraph=True)
+        @torch.compile(fullgraph=True, backend="eager")
         def split(x):
             return x.split(4, 0)
 
@@ -6894,7 +6922,7 @@ def test_ones_out_dynamic(self):
         def ones_fn(size, out):
             return torch.ones(size, out=out)
 
-        opt_model = torch.compile(ones_fn)
+        opt_model = torch.compile(ones_fn, backend="eager")
 
         out1 = torch.empty(2, 3)
         opt_model((2, 3), out1)
@@ -6906,7 +6934,7 @@ def test_zeros_out_dynamic(self):
         def zeros_fn(size, out):
             return torch.zeros(size, out=out)
 
-        opt_model = torch.compile(zeros_fn)
+        opt_model = torch.compile(zeros_fn, backend="eager")
 
         out1 = torch.empty(2, 3)
         opt_model((2, 3), out1)
@@ -6918,7 +6946,7 @@ def test_empty_out_dynamic(self):
         def empty_fn(size, out):
             return torch.empty(size, out=out)
 
-        opt_model = torch.compile(empty_fn)
+        opt_model = torch.compile(empty_fn, backend="eager")
 
         out1 = torch.empty(2, 3)
         opt_model((2, 3), out1)
@@ -6944,7 +6972,7 @@ def forward(self, x):
         model = MyModel()
         inputs = torch.randn(2, 2)
         expected = model(inputs)
-        compiled_model = torch.compile(model)
+        compiled_model = torch.compile(model, backend="eager")
         actual = compiled_model(inputs)
         self.assertEqual(actual, expected)
 
@@ -7057,7 +7085,7 @@ def forward(self, x):
                     self.b_out,
                     use_reentrant=False,
                 )
-                assert isinstance(result, torch.Tensor)
+                assert isinstance(result, torch.Tensor)  # noqa: S101
                 return result
 
         x = torch.randn(100, SEQ_LEN, DIM)
@@ -7145,9 +7173,14 @@ def f(image_latent):
             ]
             return image_latent[torch.arange(B).unsqueeze(-1), indices][:, :num_ref]
 
+        # Generate input once to ensure consistency across runs
         torch.manual_seed(54321)
         torch.cuda.manual_seed_all(54321)
-        expected = f(torch.randn((2, 12, 16, 32, 32))).sum()
+        image_latent = torch.randn((2, 12, 16, 32, 32))
+
+        torch.manual_seed(54321)
+        torch.cuda.manual_seed_all(54321)
+        expected = f(image_latent).sum()
 
         # https://github.com/pytorch/pytorch/issues/147171
         with torch._inductor.config.patch(fallback_random=True):
@@ -7155,7 +7188,7 @@ def f(image_latent):
                 torch.manual_seed(54321)
                 torch.cuda.manual_seed_all(54321)
                 actual = torch.compile(backend=backend, fullgraph=True)(f)(
-                    torch.randn((2, 12, 16, 32, 32))
+                    image_latent
                 ).sum()
                 self.assertEqual(actual, expected)
 
@@ -7163,17 +7196,17 @@ def test_incompatible_configs(self):
         with torch._dynamo.config.patch(
             suppress_errors=False, fail_on_recompile_limit_hit=False
         ):
-            torch.compile(lambda: None)
+            torch.compile(lambda: None, backend="eager")
 
         with torch._dynamo.config.patch(
             suppress_errors=True, fail_on_recompile_limit_hit=False
         ):
-            torch.compile(lambda: None)
+            torch.compile(lambda: None, backend="eager")
 
         with torch._dynamo.config.patch(
             suppress_errors=False, fail_on_recompile_limit_hit=True
         ):
-            torch.compile(lambda: None)
+            torch.compile(lambda: None, backend="eager")
 
         with (
             torch._dynamo.config.patch(
@@ -7181,7 +7214,7 @@ def test_incompatible_configs(self):
             ),
             self.assertRaises(AssertionError),
         ):
-            torch.compile(lambda: None)
+            torch.compile(lambda: None, backend="eager")
 
     def test_str_isalnum(self):
         def f(x, c):
@@ -7515,7 +7548,7 @@ def forward(self, x):
             x = torch.randn(5)
 
             with self.assertRaises(AssertionError) as e:
-                torch.compile(model)(x)
+                torch.compile(model, backend="eager")(x)
 
         msg = str(e.exception)
         self.assertIn(
@@ -7592,9 +7625,10 @@ def fn(x):
             return x + 10
 
         torch._dynamo.reset()
-        assert not torch._C._dynamo.eval_frame._debug_get_cache_entry_list(
+        if torch._C._dynamo.eval_frame._debug_get_cache_entry_list(
             fn._torchdynamo_orig_callable.__code__
-        )
+        ):
+            raise AssertionError("Expected no cache entries after reset")
 
         # Step 1: Compile a static shapes graph
         x = torch.randn(10, 10)
@@ -7633,9 +7667,10 @@ def fn(x):
 
         def run():
             torch._dynamo.reset()
-            assert not torch._C._dynamo.eval_frame._debug_get_cache_entry_list(
+            if torch._C._dynamo.eval_frame._debug_get_cache_entry_list(
                 fn._torchdynamo_orig_callable.__code__
-            )
+            ):
+                raise AssertionError("Expected no cache entries after reset")
 
             # Step 1: Compile a static shapes graph
             x = torch.randn(10, 10)
@@ -7934,6 +7969,128 @@ def __torch_dispatch__(self, func, types, args, kwargs):
         with torch._dynamo.config.patch(error_on_recompile=True):
             linear(torch.randn(1, 2, device="cpu"))
 
+    def test_property_setter_with_dict_get_176608(self):
+        """
+        Test that property setters work correctly with __dict__.get() in compiled functions.
+        Regression test for https://github.com/pytorch/pytorch/issues/176608
+        """
+        from torch.compiler import disable
+
+        class Container:
+            def __init__(self):
+                self._len_value = 0
+
+            @property
+            def _len(self):
+                # Using __dict__.get instead of self._len_value triggers the bug
+                return self.__dict__.get("_len_value", 0)
+
+            @_len.setter
+            def _len(self, value):
+                self._len_value = value
+
+            def add(self, n):
+                self._len = self._len + n
+
+            @disable()
+            def __len__(self):
+                return self._len
+
+        c = Container()
+
+        @torch.compile(backend="eager")
+        def f(x):
+            c.add(x.shape[0])  # mutates c._len_value via property setter
+            return len(c)  # reads c._len via property getter -> __dict__.get
+
+        result = f(torch.randn(5))
+        self.assertEqual(result, 5)
+
+    def test_one_hot_bounds_check_compiled(self):
+        # https://github.com/pytorch/pytorch/issues/144211
+        # torch.compile(one_hot) should raise on out-of-bounds indices,
+        # not silently produce wrong results.
+        one_hot = torch.compile(torch.nn.functional.one_hot, fullgraph=True)
+
+        a = torch.arange(0, 5) % 3  # [0, 1, 2, 0, 1]
+        with self.assertRaises(RuntimeError):
+            one_hot(a, 1)
+
+        torch._dynamo.reset()
+        with self.assertRaises(RuntimeError):
+            one_hot(torch.tensor([-1, 0, 1]), 3)
+
+        torch._dynamo.reset()
+        expected = torch.nn.functional.one_hot(a, 3)
+        self.assertEqual(one_hot(a, 3), expected)
+
+    @unittest.expectedFailure
+    def test_method_dunder_dict_setitem(self):
+        # Reproducer for: getattr(obj, method_name).__dict__['key'] = value
+        # method.__dict__ is handled specially by CPython at C level (no
+        # tp_dictoffset, no Python-visible descriptor), which caused
+        # object.__getattribute__(method, "__dict__") to raise AttributeError.
+
+        @torch.compile(backend="eager", fullgraph=True)
+        def fn(x):
+            getattr(self, self._testMethodName).__dict__["slow_test"] = True
+            return x.sin()
+
+        x = torch.randn(2)
+        _ = fn(x)
+        self.assertTrue(getattr(self, self._testMethodName).__dict__.get("slow_test"))
+
+    def test_elementwise_dtypes_constant_fold(self):
+        from torch._prims_common import (
+            elementwise_dtypes,
+            ELEMENTWISE_TYPE_PROMOTION_KIND,
+        )
+
+        @torch.compile(fullgraph=True, backend="eager")
+        def fn(x):
+            dt, _ = elementwise_dtypes(
+                x, type_promotion_kind=ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT
+            )
+            return x.to(dt)
+
+        result = fn(torch.randn(3))
+        self.assertEqual(result.dtype, torch.float32)
+
+    def test_elementwise_dtypes_int_to_float(self):
+        from torch._prims_common import (
+            elementwise_dtypes,
+            ELEMENTWISE_TYPE_PROMOTION_KIND,
+        )
+
+        @torch.compile(fullgraph=True, backend="eager")
+        def fn(x):
+            dt, _ = elementwise_dtypes(
+                x, type_promotion_kind=ELEMENTWISE_TYPE_PROMOTION_KIND.INT_TO_FLOAT
+            )
+            return x.to(dt)
+
+        result = fn(torch.randint(0, 10, (3,)))
+        self.assertEqual(result.dtype, torch.float32)
+
+    def test_elementwise_dtypes_multi_args(self):
+        from torch._prims_common import (
+            elementwise_dtypes,
+            ELEMENTWISE_TYPE_PROMOTION_KIND,
+        )
+
+        @torch.compile(fullgraph=True, backend="eager")
+        def fn(x, y):
+            dt, _ = elementwise_dtypes(
+                x, y, type_promotion_kind=ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT
+            )
+            return x.to(dt)
+
+        result = fn(
+            torch.randn(3, dtype=torch.float16),
+            torch.randn(3, dtype=torch.float32),
+        )
+        self.assertEqual(result.dtype, torch.float32)
+
 
 class ReproTestsDevice(torch._dynamo.test_case.TestCase):
     def test_sub_alpha_scalar_repro(self, device):
@@ -7990,7 +8147,7 @@ def f():
 
     @skipIfHpu
     @unittest.skipIf(
-        TEST_WITH_ROCM or not PLATFORM_SUPPORTS_FLASH_ATTENTION,
+        not PLATFORM_SUPPORTS_FLASH_ATTENTION,
         "flash attention not supported",
     )
     def test_flash_attn_backward_mixed_strides(self, device):
@@ -8168,7 +8325,7 @@ def f(x, s0, s1, s2):
         s1 = 64
         s2 = 128
 
-        f_compiled = torch.compile(f, dynamic=True)
+        f_compiled = torch.compile(f, dynamic=True, backend="eager")
 
         with torch._dynamo.config.patch(assume_static_by_default=False):
             out_ref = f(x_ref, s0, s1, s2)
@@ -8357,7 +8514,7 @@ def test_data_dependent_error_log_no_print(self):
         capturedOutput = StringIO()
         sys.stderr = capturedOutput
 
-        @torch.compile(fullgraph=True)
+        @torch.compile(fullgraph=True, backend="eager")
         def func(a):
             if a.sum() > 0:
                 return a + 1
@@ -8479,7 +8636,7 @@ def fn(x):
         cnt = torch._dynamo.testing.CompileCounter()
         opt_fn = torch.compile(fn, backend=cnt)
         self.assertEqual(fn(x), opt_fn(x))
-        self.assertEqual(cnt.frame_count, 2)
+        self.assertEqual(cnt.frame_count, 1)
 
     def test_filter_warnings(self):
         x = torch.ones(2, 2, requires_grad=True)
@@ -8586,7 +8743,7 @@ def test_pytree_get_node_type_not_traced(self):
         def fn(x, y):
             # Call _get_node_type which is used internally by pytree operations
             node_type = _get_node_type([x, y])
-            assert node_type is list
+            assert node_type is list  # noqa: S101
             # Do some work with pytree structures
             data = {"a": x, "b": y}
             flat, spec = pytree.tree_flatten(data)
@@ -8619,7 +8776,7 @@ def fn(a, b):
             point = Point(a, b)
             # Call _get_node_type with a namedtuple instance
             node_type = _get_node_type(point)
-            assert node_type is namedtuple
+            assert node_type is namedtuple  # noqa: S101
             # Use pytree operations with namedtuples
             flat, spec = pytree.tree_flatten(point)
             result = flat[0] + flat[1]
@@ -8646,15 +8803,15 @@ def fn(x, y):
             # Test with various types
             # Tensors are leaves
             is_leaf_tensor = tree_is_leaf(x)
-            assert is_leaf_tensor is True
+            assert is_leaf_tensor is True  # noqa: S101
 
             # Lists are not leaves (they're in SUPPORTED_NODES)
             is_leaf_list = tree_is_leaf([x, y])
-            assert is_leaf_list is False
+            assert is_leaf_list is False  # noqa: S101
 
             # Dicts are not leaves
             is_leaf_dict = tree_is_leaf({"a": x, "b": y})
-            assert is_leaf_dict is False
+            assert is_leaf_dict is False  # noqa: S101
 
             return x + y
 
@@ -8708,6 +8865,40 @@ def f(x, policy):
                 ),
             )
 
+    def test_mro_source_cache_includes_attr_name(self):
+        # Base -> Mid -> A hierarchy: two class attributes with the same
+        # interned integer value share the same id().  The mro_source_cache
+        # must include the attribute name in its key; otherwise the second
+        # lookup returns the first attribute's source, installing a guard
+        # on the wrong key and missing mutations to the second attribute.
+        class Base:
+            x = 1
+            y = 1
+
+        class Mid(Base):
+            pass
+
+        class A(Mid):
+            pass
+
+        cnt = torch._dynamo.testing.CompileCounter()
+
+        @torch.compile(backend=cnt)
+        def fn(obj, t):
+            return t * obj.x + t * obj.y
+
+        obj = A()
+        t = torch.tensor([1.0])
+        result = fn(obj, t)
+        self.assertEqual(result, torch.tensor([2.0]))
+        self.assertEqual(cnt.frame_count, 1)
+
+        # Changing y on Base must trigger recompilation.
+        Base.y = 42
+        result = fn(obj, t)
+        self.assertEqual(result, torch.tensor([43.0]))
+        self.assertEqual(cnt.frame_count, 2)
+
     def test_pytree_tree_is_leaf_with_namedtuple(self):
         # Test that torch.utils._pytree.tree_is_leaf handles namedtuples correctly
         from collections import namedtuple
@@ -8723,11 +8914,11 @@ def fn(a, b):
             # Namedtuples are not leaves (they're in SUPPORTED_NODES)
             point = Point(a, b)
             is_leaf_namedtuple = tree_is_leaf(point)
-            assert is_leaf_namedtuple is False
+            assert is_leaf_namedtuple is False  # noqa: S101
 
             # But individual tensors are leaves
             is_leaf_tensor = tree_is_leaf(a)
-            assert is_leaf_tensor is True
+            assert is_leaf_tensor is True  # noqa: S101
 
             return a + b
 
@@ -8740,6 +8931,239 @@ def fn(a, b):
         # Should compile successfully with fullgraph=True
         self.assertEqual(cnt.frame_count, 1)
 
+    def test_data_attr_mutation_with_noop_add(self):
+        # Regression test: remove_no_ops incorrectly eliminated add(x, 0) -> x
+        # when x was subsequently mutated by set_, causing the return value to
+        # alias the mutated input instead of being an independent copy.
+        def fn(a, b):
+            a.data = b
+            b.data = torch.zeros_like(b)
+            return a + b
+
+        a = torch.tensor([True, False, True, False])
+        b = torch.tensor([False, False, True, True])
+        a_ = a.clone()
+        b_ = b.clone()
+        cfunc = torch.compile(fn, backend="inductor")
+        res1 = fn(a, b)
+        res2 = cfunc(a_, b_)
+        self.assertEqual(res1, res2)
+
+    def test_custom_op_mutation_with_noop_add(self):
+        @torch.library.custom_op("test_repros::mutate_tensor", mutates_args={"x"})
+        def mutate_tensor(x: torch.Tensor, src: torch.Tensor) -> None:
+            x.copy_(src)
+
+        def fn(b):
+            zeros = torch.zeros_like(b)
+            result = b + zeros
+            mutate_tensor(b, zeros)
+            return result
+
+        b = torch.tensor([4.0, 5.0, 6.0])
+        b_ = b.clone()
+        cfunc = torch.compile(fn, backend="inductor")
+        res1 = fn(b)
+        res2 = cfunc(b_)
+        self.assertEqual(res1, res2)
+
+    def test_getset_descriptor_objclass_identity(self):
+        # GetSetDescriptor.__objclass__ should preserve identity with the class
+        # under torch.compile. This is needed for inspect.getattr_static (and
+        # therefore inspect.signature) to work on callable class instances.
+        class Foo:
+            pass
+
+        desc = Foo.__dict__["__dict__"]
+
+        @torch.compile(backend="eager", fullgraph=True)
+        def f(x):
+            if desc.__objclass__ is Foo:
+                return x + 1.0
+            return x + 2.0
+
+        result = f(torch.tensor(0.0))
+        self.assertEqual(result.item(), 1.0)
+
+    def test_inspect_signature_callable_class(self):
+        # inspect.signature should work on callable class instances under
+        # torch.compile, needed by flex_attention's _get_mod_type.
+        class MyCallable:
+            def __call__(self, b, h, q_idx, kv_idx):
+                return q_idx >= kv_idx
+
+        obj = MyCallable()
+
+        @torch.compile(backend="eager", fullgraph=True)
+        def f(x):
+            sig = inspect.signature(obj)
+            num_params = sum(
+                1
+                for p in sig.parameters.values()
+                if p.default is inspect.Parameter.empty
+            )
+            return x + num_params
+
+        result = f(torch.tensor(0.0))
+        self.assertEqual(result.item(), 4.0)
+
+    def test_enum_with_class_values(self):
+        from enum import Enum
+
+        class AvgMetric:
+            def __init__(self):
+                self.sum = None
+                self.count = 0
+
+            def append(self, x):
+                if self.count > 0:
+                    self.sum = self.sum + x
+                else:
+                    self.sum = x.clone()
+                self.count += 1
+
+        class GlobalReduction(Enum):
+            AVG = AvgMetric
+
+        class ScalarLogger:
+            def __init__(self):
+                self.metrics = {}
+
+            def log(self, key, value, global_reduction):
+                if key not in self.metrics:
+                    self.metrics[key] = global_reduction.value()
+                self.metrics[key].append(value)
+
+        @torch.compile(backend="eager", fullgraph=True)
+        def fn(logger, x):
+            logger.log("test", x, GlobalReduction.AVG)
+            return x + 1
+
+        logger = ScalarLogger()
+        fn(logger, torch.tensor(1.0))
+
+    def test_class_attr_mutation_recompiles(self):
+        class GlobalState:
+            factor = 1.0
+
+        cnt = torch._dynamo.testing.CompileCounter()
+
+        @torch.compile(backend=cnt)
+        def fn(x):
+            return x * GlobalState.factor
+
+        x = torch.tensor([4.0])
+
+        GlobalState.factor = 1.0
+        result1 = fn(x)
+        self.assertEqual(result1, torch.tensor([4.0]))
+        self.assertEqual(cnt.frame_count, 1)
+
+        GlobalState.factor = 10.0
+        result2 = fn(x)
+        self.assertEqual(result2, torch.tensor([40.0]))
+        self.assertEqual(cnt.frame_count, 2)
+
+    @skipIfHpu
+    @requires_cuda
+    def test_deterministic_pad_replicate_compile(self, device):
+        from torch.testing._internal.common_utils import DeterministicGuard
+
+        pad = torch.nn.ReplicationPad1d(2).to(device)
+        compiled_pad = torch.compile(pad, backend="aot_eager", fullgraph=True)
+        x = torch.randn(3, 3, device=device, requires_grad=True)
+        with DeterministicGuard(True):
+            ref = pad(x)
+            res = compiled_pad(x)
+            self.assertEqual(ref, res)
+            grad = torch.autograd.grad(res.sum(), x)
+            ref_grad = torch.autograd.grad(ref.sum(), x)
+            self.assertEqual(grad, ref_grad)
+
+    @requires_cuda
+    @unittest.skipIf(
+        TEST_WITH_ROCM or not PLATFORM_SUPPORTS_FLASH_ATTENTION,
+        "flash attention not supported",
+    )
+    def test_flex_attention_guard_on_constant_func_defaults(self):
+        """
+        Dynamo must guard on mask_mod.__defaults__ so that when a
+        compiled function is re-invoked with a new BlockMask whose
+        mask_mod has the same __code__ but different __defaults__,
+        Dynamo recompiles instead of reusing the stale first graph.
+        """
+        from torch.utils._triton import has_triton
+
+        if not has_triton():
+            self.skipTest("requires triton")
+
+        @torch.compile(fullgraph=True)
+        def flex_chunk(q, k, v, block_mask, scale):
+            out, aux = flex_attention(
+                q,
+                k,
+                v,
+                block_mask=block_mask,
+                scale=scale,
+                return_aux=AuxRequest(lse=True),
+            )
+            return out, aux.lse
+
+        def merge(out, lse, new_out, new_lse):
+            lse, new_lse = lse.unsqueeze(-1), new_lse.unsqueeze(-1)
+            mx = torch.maximum(lse, new_lse)
+            e0, e1 = torch.exp(lse - mx), torch.exp(new_lse - mx)
+            d = e0 + e1
+            return (out * e0 + new_out * e1) / d, (mx + torch.log(d)).squeeze(-1)
+
+        @torch.compile(fullgraph=True)
+        def ref_attn(q, k, v, block_mask, scale):
+            return flex_attention(q, k, v, block_mask=block_mask, scale=scale)
+
+        torch.manual_seed(42)
+        B, H, S, D = 1, 1, 512, 16
+        device = "cuda"
+        NUM_CHUNKS = 4
+        chunk_size = S // NUM_CHUNKS
+
+        q = torch.randn(B, H, S, D, device=device)
+        k = torch.randn(B, H, S, D, device=device)
+        v = torch.randn(B, H, S, D, device=device)
+        scale = D**-0.5
+
+        merged_out = merged_lse = None
+        for step in range(NUM_CHUNKS):
+            kv_offset = step * chunk_size
+
+            def mask_mod(b, h, q_idx, kv_idx, _offset=kv_offset):
+                return q_idx >= kv_idx + _offset
+
+            bm = create_block_mask(
+                mask_mod, B=B, H=H, Q_LEN=S, KV_LEN=chunk_size, device=device
+            )
+            out, lse = flex_chunk(
+                q,
+                k[:, :, kv_offset : kv_offset + chunk_size],
+                v[:, :, kv_offset : kv_offset + chunk_size],
+                bm,
+                scale,
+            )
+            if merged_out is None:
+                merged_out, merged_lse = out, lse
+            else:
+                merged_out, merged_lse = merge(merged_out, merged_lse, out, lse)
+
+        def causal(b, h, q_idx, kv_idx):
+            return q_idx >= kv_idx
+
+        ref_bm = create_block_mask(causal, B=B, H=H, Q_LEN=S, KV_LEN=S, device=device)
+        ref_out = ref_attn(q, k, v, ref_bm, scale)
+
+        self.assertTrue(
+            (merged_out - ref_out).abs().max().item() < 1e-3,
+            "flex_attention mask_mod __defaults__ not properly guarded",
+        )
+
 
 instantiate_parametrized_tests(ReproTests)
 
diff --git a/test/dynamo/test_resume.py b/test/dynamo/test_resume.py
index 42103a7878e73..73cd3779868d1 100644
--- a/test/dynamo/test_resume.py
+++ b/test/dynamo/test_resume.py
@@ -26,10 +26,13 @@ def test_freevars(self):
         fn = fn_creator()
         opt_fn = torch.compile(fn, backend="eager")
         opt_fn(torch.randn(10))
-        codes = [v for k, v in list(globals().items()) if k.startswith("__resume_at")]
-        self.assertEqual(len(codes), 1)
+        entries = [v for k, v in list(globals().items()) if k.startswith("__resume_at")]
+        self.assertEqual(len(entries), 1)
+        # When freevars are present, install_resume_function_global stores a
+        # factory that closes over the code object (first closure cell).
+        code = entries[0].__closure__[0].cell_contents
         # co_freevars of resume functions, are sorted concatenation of the original function's co_freevars and co_cellvars
-        self.assertEqual(codes[0].co_freevars, ("var1", "var2"))
+        self.assertEqual(code.co_freevars, ("var1", "var2"))
 
 
 if __name__ == "__main__":
diff --git a/test/dynamo/test_streams.py b/test/dynamo/test_streams.py
index 2f71d40fde791..490ce1d1d0594 100644
--- a/test/dynamo/test_streams.py
+++ b/test/dynamo/test_streams.py
@@ -1,5 +1,4 @@
 # Owner(s): ["module: dynamo"]
-import functools
 import re
 import unittest
 import weakref
@@ -13,15 +12,9 @@
     store_user_object_weakrefs,
 )
 from torch._dynamo.testing import extract_graph, remove_trailing_space
-from torch.testing._internal.common_cuda import TEST_MULTIGPU
 from torch.testing._internal.common_utils import requires_cuda
 
 
-requires_multigpu = functools.partial(
-    unittest.skipIf, not TEST_MULTIGPU, "requires multiple cuda devices"
-)
-
-
 def remove_file_comment(gm_str: str) -> str:
     return remove_trailing_space(re.sub(r"File.*\n", "\n", gm_str))
 
@@ -75,13 +68,13 @@ def fn(x, y, s1, s2):
             """\
 class <lambda>(torch.nn.Module):
     def forward(self, arg0_1: "f32[2, 2]", arg1_1: "f32[2, 2]"):
-        # Annotation: {'stream': 0}
+        # Annotation: {'stream': 1}
         add: "f32[2, 2]" = torch.ops.aten.add.Tensor(arg0_1, arg1_1)
 
-        # Annotation: {'stream': 1}
+        # Annotation: {'stream': 2}
         add_1: "f32[2, 2]" = torch.ops.aten.add.Tensor(arg0_1, arg1_1);  arg0_1 = arg1_1 = None
 
-        # Annotation: {'stream': 1}
+        # Annotation: {'stream': 2}
         add_2: "f32[2, 2]" = torch.ops.aten.add.Tensor(add_1, 2);  add_1 = None
         add_3: "f32[2, 2]" = torch.ops.aten.add.Tensor(add_2, add);  add_2 = add = None
         return (add_3,)
@@ -164,38 +157,30 @@ def fn(x, s):
         self.assertEqual(s0, s1)
 
     @requires_cuda
-    @requires_multigpu()
-    def test_get_current_stream_return_different_device(self):
-        def fn(x, s0, s1):
-            with s1:
-                with s0:
-                    s = torch.accelerator.current_stream(torch.device("cuda:1"))
-            return s
+    def test_cuda_current_stream_attrs(self):
+        """Verify that torch.cuda.current_stream() attributes are accessible
+        under torch.compile and match eager behavior."""
 
-        s0 = torch.Stream(device="cuda:0")
-        s1 = torch.Stream(device="cuda:1")
-        inp = (torch.ones(2, 2) + 1, s0, s1)
-        fn_opt = torch.compile(fn, fullgraph=True)
-        s_act = fn_opt(*inp)
-        s_exp = fn(*inp)
-        self.assertEqual(s_act, s_exp)
+        def fn_cuda_stream(x):
+            return torch.cuda.current_stream().cuda_stream
+
+        x = torch.zeros(1, device="cuda")
+        compiled = torch.compile(fn_cuda_stream, backend="eager", fullgraph=True)
+        self.assertEqual(compiled(x), fn_cuda_stream(x))
 
     @requires_cuda
-    @requires_multigpu()
-    def test_get_current_stream_return_no_index(self):
-        def fn(x, s0, s1):
-            with s1:
-                with s0:
-                    s = torch.accelerator.current_stream(torch.device("cuda"))
-            return s
+    def test_cuda_current_stream_with_entered_stream(self):
+        """Verify that torch.cuda.current_stream().cuda_stream returns the
+        correct value when inside a stream context for a user-created stream."""
 
-        s0 = torch.Stream(device="cuda:0")
-        s1 = torch.Stream(device="cuda:1")
-        inp = (torch.ones(2, 2) + 1, s0, s1)
-        fn_opt = torch.compile(fn, fullgraph=True)
-        s_act = fn_opt(*inp)
-        s_exp = fn(*inp)
-        self.assertEqual(s_act, s_exp)
+        def fn(x, s):
+            with s:
+                return torch.cuda.current_stream().cuda_stream
+
+        s = torch.cuda.Stream()
+        x = torch.zeros(1, device="cuda")
+        compiled = torch.compile(fn, backend="eager", fullgraph=True)
+        self.assertEqual(compiled(x, s), fn(x, s))
 
     @requires_cuda
     def test_nested_stream_enter_exit(self):
@@ -231,13 +216,13 @@ def fn(x, y, s0, s1, s2):
             """\
 class <lambda>(torch.nn.Module):
     def forward(self, arg0_1: "f32[2, 2]", arg1_1: "f32[2, 2]"):
-        # Annotation: {'stream': 1}
+        # Annotation: {'stream': 2}
         add: "f32[2, 2]" = torch.ops.aten.add.Tensor(arg0_1, arg1_1)
 
-        # Annotation: {'stream': 2}
+        # Annotation: {'stream': 3}
         add_1: "f32[2, 2]" = torch.ops.aten.add.Tensor(arg0_1, arg1_1);  arg0_1 = arg1_1 = None
 
-        # Annotation: {'stream': 1}
+        # Annotation: {'stream': 2}
         add_2: "f32[2, 2]" = torch.ops.aten.add.Tensor(add, 2);  add = None
         return (add_1, add_2)
 """,
@@ -279,13 +264,13 @@ def fn(x, y):
             """\
 class <lambda>(torch.nn.Module):
     def forward(self, arg0_1: "f32[2, 2]", arg1_1: "f32[2, 2]"):
-        # Annotation: {'stream': 1}
+        # Annotation: {'stream': 2}
         add: "f32[2, 2]" = torch.ops.aten.add.Tensor(arg0_1, arg1_1)
 
-        # Annotation: {'stream': 0}
+        # Annotation: {'stream': 1}
         add_1: "f32[2, 2]" = torch.ops.aten.add.Tensor(arg0_1, arg1_1);  arg0_1 = arg1_1 = None
 
-        # Annotation: {'stream': 0}
+        # Annotation: {'stream': 1}
         add_2: "f32[2, 2]" = torch.ops.aten.add.Tensor(add_1, 2);  add_1 = None
         add_3: "f32[2, 2]" = torch.ops.aten.add.Tensor(add_2, add);  add_2 = add = None
         return (add_3,)
@@ -323,46 +308,18 @@ def fn(x, y):
             """\
 class <lambda>(torch.nn.Module):
     def forward(self, arg0_1: "f32[2, 2]", arg1_1: "f32[2, 2]"):
-        # Annotation: {'stream': 0}
+        # Annotation: {'stream': 1}
         add: "f32[2, 2]" = torch.ops.aten.add.Tensor(arg0_1, arg1_1)
 
-        # Annotation: {'stream': 2}
+        # Annotation: {'stream': 3}
         add_1: "f32[2, 2]" = torch.ops.aten.add.Tensor(arg0_1, arg1_1);  arg0_1 = arg1_1 = None
 
-        # Annotation: {'stream': 0}
+        # Annotation: {'stream': 1}
         add_2: "f32[2, 2]" = torch.ops.aten.add.Tensor(add, 2);  add = None
         return (add_1, add_2)
 """,
         )
 
-    @requires_cuda
-    @requires_multigpu()
-    def test_new_event_api(self) -> None:
-        from torch._dynamo.graph_bytecode_inputs import get_external_object_by_index
-        from torch._dynamo.variables.streams import new_event
-
-        def event_generation_backend(gm, *args, **kwargs):  # type: ignore[no-untyped-def]
-            e0_ind = new_event()
-            with torch.Stream(device="cuda:1"):
-                get_external_object_by_index(e0_ind).record()
-            e1_ind = new_event()
-            self.assertNotEqual(e0_ind, e1_ind)
-            self.assertNotEqual(
-                get_external_object_by_index(e0_ind),
-                get_external_object_by_index(e1_ind),
-            )
-            with gm.graph.inserting_after(next(iter(gm.graph.nodes))):
-                gm.graph.call_function(
-                    get_external_object_by_index, args=(1,), kwargs={}
-                )
-            return gm
-
-        @torch.compile(backend=event_generation_backend)
-        def fn(x):
-            return x + 1
-
-        fn(torch.ones(2, 2, device="cuda:0"))
-
     @requires_cuda
     def test_new_stream_api(self) -> None:
         from torch._dynamo.graph_bytecode_inputs import get_external_object_by_index
@@ -451,19 +408,19 @@ def fn(x, y):
             """\
 class <lambda>(torch.nn.Module):
     def forward(self, arg0_1: "f32[2, 2]", arg1_1: "f32[2, 2]"):
-        # Annotation: {'stream': 0}
+        # Annotation: {'stream': 1}
         add: "f32[2, 2]" = torch.ops.aten.add.Tensor(arg0_1, arg1_1)
 
-        # Annotation: {'stream': 2}
+        # Annotation: {'stream': 3}
         add_1: "f32[2, 2]" = torch.ops.aten.add.Tensor(arg1_1, arg1_1)
 
-        # Annotation: {'stream': 2}
+        # Annotation: {'stream': 3}
         add_2: "f32[2, 2]" = torch.ops.aten.add.Tensor(add_1, arg1_1);  arg1_1 = None
 
-        # Annotation: {'stream': 0}
+        # Annotation: {'stream': 1}
         add_3: "f32[2, 2]" = torch.ops.aten.add.Tensor(add_1, 2);  add_1 = None
 
-        # Annotation: {'stream': 0}
+        # Annotation: {'stream': 1}
         copy_: "f32[2, 2]" = torch.ops.aten.copy_.default(arg0_1, add);  arg0_1 = add = copy_ = None
         return (add_2, add_3)
 """,
@@ -499,13 +456,13 @@ def fn(x, y):
             """\
 class GraphModule(torch.nn.Module):
     def forward(self, primals_1: "f32[2, 2]", primals_2: "f32[2, 2]"):
-        # Annotation: {'stream': 1}
+        # Annotation: {'stream': 2}
         mul: "f32[2, 2]" = torch.ops.aten.mul.Tensor(primals_1, 2);  primals_1 = None
         add: "f32[2, 2]" = torch.ops.aten.add.Tensor(mul, primals_2)
 
-        # Annotation: {'stream': 0}
-        add_1: "f32[2, 2]" = torch.ops.aten.add.Tensor(mul, primals_2);  mul = primals_2 = None
-        return (add, add_1)
+        # Annotation: {'stream': 1}
+        add_1: "f32[2, 2]" = torch.ops.aten.add.Tensor(mul, primals_2);  primals_2 = None
+        return (add, add_1, mul, add_1)
 """,
         )
 
@@ -514,24 +471,39 @@ def forward(self, primals_1: "f32[2, 2]", primals_2: "f32[2, 2]"):
             print_graph(bw_graphs[0]),
             """\
 class GraphModule(torch.nn.Module):
-    def forward(self, tangents_1: "f32[2, 2]", tangents_2: "f32[2, 2]"):
-        # Annotation: {'stream': 0}
+    def forward(self, mul: "f32[2, 2]", add_1: "f32[2, 2]", tangents_1: "f32[2, 2]", tangents_2: "f32[2, 2]"):
+        # Annotation: {'stream': 1}
         mul_2: "f32[2, 2]" = torch.ops.aten.mul.Tensor(tangents_2, 2)
 
         #
         add_2: "f32[2, 2]" = torch.ops.aten.add.Tensor(tangents_2, tangents_1);  tangents_2 = None
 
-        # Annotation: {'stream': 1}
+        # Annotation: {'stream': 2}
         mul_3: "f32[2, 2]" = torch.ops.aten.mul.Tensor(tangents_1, 2);  tangents_1 = None
 
-        # Annotation: {'stream': 0}
-        add_3: "f32[2, 2]" = torch.ops.aten.add.Tensor(mul_2, mul_3);  mul_2 = None
+        # Annotation: {'stream': 1}
+        add_3: "f32[2, 2]" = torch.ops.aten.add.Tensor(mul_2, mul_3)
 
         # No stacktrace found for following nodes
-        record_event_default = torch.ops.streams.record_event.default(2, 0);  record_event_default = None
-        sync_dealloc_default = torch.ops.streams.sync_dealloc.default(2, 1, mul_3);  mul_3 = sync_dealloc_default = None
-        return (add_3, add_2)
-""",
+        subgraph_record_event_default = self.subgraph_record_event_default
+        control_deps = torch.ops.higher_order.control_deps((mul, add_1, mul_2, add_3, add_2), subgraph_record_event_default, add_1, add_3, add_2);  mul = add_1 = mul_2 = add_3 = add_2 = subgraph_record_event_default = None
+
+        #
+        getitem_2: "f32[2, 2]" = control_deps[3]
+
+        # Annotation: {'stream': 1}
+        getitem_1: "f32[2, 2]" = control_deps[2];  control_deps = None
+
+        # No stacktrace found for following nodes
+        sync_dealloc_default = torch.ops.streams.sync_dealloc.default(3, 2, mul_3);  mul_3 = sync_dealloc_default = None
+        return (getitem_1, getitem_2)
+
+    class subgraph_record_event_default(torch.nn.Module):
+        def forward(self, dep_0: "f32[2, 2]", dep_1: "f32[2, 2]", dep_2: "f32[2, 2]"):
+            # No stacktrace found for following nodes
+            record_event_default = torch.ops.streams.record_event.default(3, 1)
+            return (record_event_default, dep_0, dep_1, dep_2)
+""",  # noqa: B950
         )
 
     @requires_cuda
@@ -564,13 +536,13 @@ def fn(x, y):
             """\
 class GraphModule(torch.nn.Module):
     def forward(self, primals_1: "f32[2, 2]", primals_2: "f32[2, 2]"):
-        # Annotation: {'stream': 1}
+        # Annotation: {'stream': 2}
         mul: "f32[2, 2]" = torch.ops.aten.mul.Tensor(primals_1, 2);  primals_1 = None
         add: "f32[2, 2]" = torch.ops.aten.add.Tensor(mul, primals_2)
 
-        # Annotation: {'stream': 0}
-        add_1: "f32[2, 2]" = torch.ops.aten.add.Tensor(mul, primals_2);  mul = primals_2 = None
-        return (add, add_1)
+        # Annotation: {'stream': 1}
+        add_1: "f32[2, 2]" = torch.ops.aten.add.Tensor(mul, primals_2);  primals_2 = None
+        return (add, add_1, mul, add, add_1)
 """,
         )
 
@@ -579,28 +551,65 @@ def forward(self, primals_1: "f32[2, 2]", primals_2: "f32[2, 2]"):
             print_graph(bw_graphs[0]),
             """\
 class GraphModule(torch.nn.Module):
-    def forward(self, tangents_1: "f32[2, 2]", tangents_2: "f32[2, 2]"):
-        # Annotation: {'stream': 0}
+    def forward(self, mul: "f32[2, 2]", add: "f32[2, 2]", add_1: "f32[2, 2]", tangents_1: "f32[2, 2]", tangents_2: "f32[2, 2]"):
+        # Annotation: {'stream': 1}
         mul_2: "f32[2, 2]" = torch.ops.aten.mul.Tensor(tangents_2, 2)
 
         #
         add_2: "f32[2, 2]" = torch.ops.aten.add.Tensor(tangents_2, tangents_1);  tangents_2 = None
 
-        # Annotation: {'stream': 1}
+        # Annotation: {'stream': 2}
         mul_3: "f32[2, 2]" = torch.ops.aten.mul.Tensor(tangents_1, 2);  tangents_1 = None
 
         # No stacktrace found for following nodes
-        record_event_default = torch.ops.streams.record_event.default(2, 1);  record_event_default = None
-        wait_event_default = torch.ops.streams.wait_event.default(2, 0);  wait_event_default = None
+        subgraph_record_event_default = self.subgraph_record_event_default
+        control_deps = torch.ops.higher_order.control_deps((mul, add, mul_3, add_2), subgraph_record_event_default, add, mul_3, add_2);  add = mul_3 = add_2 = subgraph_record_event_default = None
 
-        # Annotation: {'stream': 0}
-        add_3: "f32[2, 2]" = torch.ops.aten.add.Tensor(mul_2, mul_3);  mul_2 = None
+        #
+        getitem_2: "f32[2, 2]" = control_deps[3]
+
+        # Annotation: {'stream': 2}
+        getitem_1: "f32[2, 2]" = control_deps[2]
 
         # No stacktrace found for following nodes
-        record_event_default_1 = torch.ops.streams.record_event.default(3, 0);  record_event_default_1 = None
-        sync_dealloc_default = torch.ops.streams.sync_dealloc.default(3, 1, mul_3);  mul_3 = sync_dealloc_default = None
-        return (add_3, add_2)
-""",
+        subgraph_wait_event_default = self.subgraph_wait_event_default
+        control_deps_1 = torch.ops.higher_order.control_deps((control_deps, mul, add_1, mul_2), subgraph_wait_event_default, add_1, mul_2);  control_deps = mul = add_1 = mul_2 = subgraph_wait_event_default = None
+
+        # Annotation: {'stream': 1}
+        getitem_4: "f32[2, 2]" = control_deps_1[2];  control_deps_1 = None
+
+        # Annotation: {'stream': 1}
+        add_3: "f32[2, 2]" = torch.ops.aten.add.Tensor(getitem_4, getitem_1);  getitem_4 = None
+
+        # No stacktrace found for following nodes
+        subgraph_record_event_default_1 = self.subgraph_record_event_default_1
+        control_deps_2 = torch.ops.higher_order.control_deps((add_3,), subgraph_record_event_default_1, add_3);  add_3 = subgraph_record_event_default_1 = None
+
+        # Annotation: {'stream': 1}
+        getitem_5: "f32[2, 2]" = control_deps_2[1];  control_deps_2 = None
+
+        # No stacktrace found for following nodes
+        sync_dealloc_default = torch.ops.streams.sync_dealloc.default(4, 2, getitem_1);  getitem_1 = sync_dealloc_default = None
+        return (getitem_5, getitem_2)
+
+    class subgraph_record_event_default(torch.nn.Module):
+        def forward(self, dep_0: "f32[2, 2]", dep_1: "f32[2, 2]", dep_2: "f32[2, 2]"):
+            # No stacktrace found for following nodes
+            record_event_default = torch.ops.streams.record_event.default(3, 2)
+            return (record_event_default, dep_0, dep_1, dep_2)
+
+    class subgraph_wait_event_default(torch.nn.Module):
+        def forward(self, dep_0: "f32[2, 2]", dep_1: "f32[2, 2]"):
+            # No stacktrace found for following nodes
+            wait_event_default = torch.ops.streams.wait_event.default(3, 1)
+            return (wait_event_default, dep_0, dep_1)
+
+    class subgraph_record_event_default_1(torch.nn.Module):
+        def forward(self, dep_0: "f32[2, 2]"):
+            # No stacktrace found for following nodes
+            record_event_default = torch.ops.streams.record_event.default(4, 1)
+            return (record_event_default, dep_0)
+""",  # noqa: B950
         )
 
     @requires_cuda
@@ -625,7 +634,7 @@ def fn(x) -> None:
 class <lambda>(torch.nn.Module):
     def forward(self, arg0_1: "f32[2, 2]"):
         #
-        record_event = torch.ops.streams.record_event.default(0, 1);  record_event = None
+        record_event = torch.ops.streams.record_event.default(1, 0);  record_event = None
 
         #
         add: "f32[2, 2]" = torch.ops.aten.add.Tensor(arg0_1, 1)
@@ -733,29 +742,59 @@ def fn(x):
             print_graph(bw_graphs[0]),
             """\
 class GraphModule(torch.nn.Module):
-    def forward(self, tangents_1: "f32[2, 2]", tangents_2: "f32[2, 2]"):
-        # Annotation: {'stream': 1}
-        ones: "f32[2, 2]" = torch.ops.aten.ones.default([2, 2], device = device(type='cuda', index=0), pin_memory = False)
+    def forward(self, getitem: "f32[2, 2]", mul: "f32[2, 2]", mul_1: "f32[2, 2]", mul_2: "f32[2, 2]", tangents_1: "f32[2, 2]", tangents_2: "f32[2, 2]"):
+        # Annotation: {'stream': 3}
+        mul_3: "f32[2, 2]" = torch.ops.aten.mul.Tensor(tangents_1, mul_1);  tangents_1 = None
 
         # Annotation: {'stream': 2}
-        mul_1: "f32[2, 2]" = torch.ops.aten.mul.Tensor(ones, 2)
-        mul_3: "f32[2, 2]" = torch.ops.aten.mul.Tensor(tangents_1, mul_1);  tangents_1 = mul_1 = None
-
-        # Annotation: {'stream': 1}
-        mul_4: "f32[2, 2]" = torch.ops.aten.mul.Tensor(tangents_2, ones);  tangents_2 = ones = None
+        mul_4: "f32[2, 2]" = torch.ops.aten.mul.Tensor(tangents_2, getitem);  tangents_2 = getitem = None
 
         # No stacktrace found for following nodes
-        record_event_default = torch.ops.streams.record_event.default(3, 1);  record_event_default = None
-        wait_event_default = torch.ops.streams.wait_event.default(3, 2);  wait_event_default = None
+        subgraph_record_event_default = self.subgraph_record_event_default
+        control_deps_2 = torch.ops.higher_order.control_deps((mul, mul_4), subgraph_record_event_default, mul, mul_4);  mul = mul_4 = subgraph_record_event_default = None
 
         # Annotation: {'stream': 2}
-        add: "f32[2, 2]" = torch.ops.aten.add.Tensor(mul_3, mul_4);  mul_3 = None
+        getitem_2: "f32[2, 2]" = control_deps_2[2]
 
         # No stacktrace found for following nodes
-        record_event_default_1 = torch.ops.streams.record_event.default(4, 2);  record_event_default_1 = None
-        sync_dealloc_default = torch.ops.streams.sync_dealloc.default(4, 1, mul_4);  mul_4 = sync_dealloc_default = None
-        return (add,)
-""",
+        subgraph_wait_event_default = self.subgraph_wait_event_default
+        control_deps_3 = torch.ops.higher_order.control_deps((control_deps_2, mul_1, mul_2, mul_3), subgraph_wait_event_default, mul_2, mul_3);  control_deps_2 = mul_1 = mul_2 = mul_3 = subgraph_wait_event_default = None
+
+        # Annotation: {'stream': 3}
+        getitem_4: "f32[2, 2]" = control_deps_3[2];  control_deps_3 = None
+
+        # Annotation: {'stream': 3}
+        add: "f32[2, 2]" = torch.ops.aten.add.Tensor(getitem_4, getitem_2);  getitem_4 = None
+
+        # No stacktrace found for following nodes
+        subgraph_record_event_default_1 = self.subgraph_record_event_default_1
+        control_deps_4 = torch.ops.higher_order.control_deps((add,), subgraph_record_event_default_1, add);  add = subgraph_record_event_default_1 = None
+
+        # Annotation: {'stream': 3}
+        getitem_5: "f32[2, 2]" = control_deps_4[1];  control_deps_4 = None
+
+        # No stacktrace found for following nodes
+        sync_dealloc_default = torch.ops.streams.sync_dealloc.default(5, 2, getitem_2);  getitem_2 = sync_dealloc_default = None
+        return (getitem_5,)
+
+    class subgraph_record_event_default(torch.nn.Module):
+        def forward(self, dep_0: "f32[2, 2]", dep_1: "f32[2, 2]"):
+            # No stacktrace found for following nodes
+            record_event_default = torch.ops.streams.record_event.default(4, 2)
+            return (record_event_default, dep_0, dep_1)
+
+    class subgraph_wait_event_default(torch.nn.Module):
+        def forward(self, dep_0: "f32[2, 2]", dep_1: "f32[2, 2]"):
+            # No stacktrace found for following nodes
+            wait_event_default = torch.ops.streams.wait_event.default(4, 3)
+            return (wait_event_default, dep_0, dep_1)
+
+    class subgraph_record_event_default_1(torch.nn.Module):
+        def forward(self, dep_0: "f32[2, 2]"):
+            # No stacktrace found for following nodes
+            record_event_default = torch.ops.streams.record_event.default(5, 3)
+            return (record_event_default, dep_0)
+""",  # noqa: B950
         )
 
     @requires_cuda
@@ -804,75 +843,169 @@ def fn(x):
             print_graph(bw_graphs[0]),
             """\
 class GraphModule(torch.nn.Module):
-    def forward(self, primals_1: "f32[2, 2]", mul: "f32[2, 2]", tangents_1: "f32[2, 2]", \
-tangents_2: "f32[2, 2]", tangents_3: "f32[2, 2]"):
-        # Annotation: {'stream': 1}
-        ones: "f32[2, 2]" = torch.ops.aten.ones.default([2, 2], device = device(type='cuda', index=0), pin_memory = False)
+    def forward(self, getitem: "f32[2, 2]", getitem_3: "f32[2, 2]", getitem_2: "f32[2, 2]", getitem_4: "f32[2, 2]", getitem_6: "f32[2, 2]", tangents_1: "f32[2, 2]", tangents_2: "f32[2, 2]", tangents_3: "f32[2, 2]"):
+        # Annotation: {'stream': 5}
+        mul_7: "f32[2, 2]" = torch.ops.aten.mul.Tensor(tangents_3, getitem_6);  tangents_3 = getitem_6 = None
+
+        # No stacktrace found for following nodes
+        subgraph_record_event_default = self.subgraph_record_event_default
+        control_deps_8 = torch.ops.higher_order.control_deps((mul_7,), subgraph_record_event_default, mul_7);  mul_7 = subgraph_record_event_default = None
+
+        # Annotation: {'stream': 5}
+        getitem_8: "f32[2, 2]" = control_deps_8[1]
+
+        # No stacktrace found for following nodes
+        subgraph_wait_event_default = self.subgraph_wait_event_default
+        control_deps_9 = torch.ops.higher_order.control_deps((control_deps_8,), subgraph_wait_event_default);  control_deps_8 = subgraph_wait_event_default = control_deps_9 = None
+
+        # Annotation: {'stream': 4}
+        add: "f32[2, 2]" = torch.ops.aten.add.Tensor(tangents_2, getitem_8);  tangents_2 = None
+
+        # No stacktrace found for following nodes
+        subgraph_record_event_default_4 = self.subgraph_record_event_default_4
+        control_deps_10 = torch.ops.higher_order.control_deps((add,), subgraph_record_event_default_4, add);  add = subgraph_record_event_default_4 = None
+
+        # Annotation: {'stream': 4}
+        getitem_9: "f32[2, 2]" = control_deps_10[1];  control_deps_10 = None
+
+        # No stacktrace found for following nodes
+        sync_dealloc_default = torch.ops.streams.sync_dealloc.default(10, 5, getitem_8);  getitem_8 = sync_dealloc_default = None
+
+        # Annotation: {'stream': 4}
+        mul_8: "f32[2, 2]" = torch.ops.aten.mul.Tensor(getitem_9, getitem_4);  getitem_4 = None
+
+        # No stacktrace found for following nodes
+        subgraph_record_event_default_1 = self.subgraph_record_event_default_1
+        control_deps_11 = torch.ops.higher_order.control_deps((mul_8,), subgraph_record_event_default_1, mul_8);  mul_8 = subgraph_record_event_default_1 = None
 
         # Annotation: {'stream': 4}
-        mul_5: "f32[2, 2]" = torch.ops.aten.mul.Tensor(ones, 4)
-        mul_7: "f32[2, 2]" = torch.ops.aten.mul.Tensor(tangents_3, mul_5);  tangents_3 = mul_5 = None
+        getitem_10: "f32[2, 2]" = control_deps_11[1]
+        mul_9: "f32[2, 2]" = torch.ops.aten.mul.Tensor(getitem_9, getitem_3);  getitem_9 = getitem_3 = None
+        mul_10: "f32[2, 2]" = torch.ops.aten.mul.Tensor(mul_9, getitem)
 
         # No stacktrace found for following nodes
-        record_event_default = torch.ops.streams.record_event.default(6, 4);  record_event_default = None
-        wait_event_default = torch.ops.streams.wait_event.default(6, 3);  wait_event_default = None
+        subgraph_wait_event_default_1 = self.subgraph_wait_event_default_1
+        control_deps_12 = torch.ops.higher_order.control_deps((control_deps_11,), subgraph_wait_event_default_1);  control_deps_11 = subgraph_wait_event_default_1 = control_deps_12 = None
 
         # Annotation: {'stream': 3}
-        add: "f32[2, 2]" = torch.ops.aten.add.Tensor(tangents_2, mul_7);  tangents_2 = None
+        mul_11: "f32[2, 2]" = torch.ops.aten.mul.Tensor(getitem_10, getitem_2);  getitem_2 = None
 
         # No stacktrace found for following nodes
-        record_event_default_4 = torch.ops.streams.record_event.default(10, 3);  record_event_default_4 = None
-        sync_dealloc_default = torch.ops.streams.sync_dealloc.default(10, 4, mul_7);  mul_7 = sync_dealloc_default = None
+        subgraph_record_event_default_5 = self.subgraph_record_event_default_5
+        control_deps_13 = torch.ops.higher_order.control_deps((mul_11,), subgraph_record_event_default_5, mul_11);  mul_11 = subgraph_record_event_default_5 = None
 
         # Annotation: {'stream': 3}
-        mul_3: "f32[2, 2]" = torch.ops.aten.mul.Tensor(ones, primals_1);  primals_1 = None
-        mul_8: "f32[2, 2]" = torch.ops.aten.mul.Tensor(add, mul_3);  mul_3 = None
+        getitem_11: "f32[2, 2]" = control_deps_13[1];  control_deps_13 = None
 
         # No stacktrace found for following nodes
-        record_event_default_1 = torch.ops.streams.record_event.default(7, 3);  record_event_default_1 = None
+        sync_dealloc_default_1 = torch.ops.streams.sync_dealloc.default(11, 4, getitem_10);  getitem_10 = sync_dealloc_default_1 = None
+        record_event_default_2 = torch.ops.streams.record_event.default(8, 3);  record_event_default_2 = None
+        wait_event_default_2 = torch.ops.streams.wait_event.default(8, 2);  wait_event_default_2 = None
 
         # Annotation: {'stream': 2}
-        mul_1: "f32[2, 2]" = torch.ops.aten.mul.Tensor(ones, 2)
-        mul_2: "f32[2, 2]" = torch.ops.aten.mul.Tensor(mul_1, mul);  mul = None
+        add_1: "f32[2, 2]" = torch.ops.aten.add.Tensor(tangents_1, getitem_11);  tangents_1 = None
 
-        # Annotation: {'stream': 3}
-        mul_9: "f32[2, 2]" = torch.ops.aten.mul.Tensor(add, mul_2);  add = mul_2 = None
-        mul_10: "f32[2, 2]" = torch.ops.aten.mul.Tensor(mul_9, ones);  mul_9 = None
+        # No stacktrace found for following nodes
+        subgraph_record_event_default_6 = self.subgraph_record_event_default_6
+        control_deps_14 = torch.ops.higher_order.control_deps((add_1,), subgraph_record_event_default_6, add_1);  add_1 = subgraph_record_event_default_6 = None
+
+        # Annotation: {'stream': 2}
+        getitem_12: "f32[2, 2]" = control_deps_14[1];  control_deps_14 = None
 
         # No stacktrace found for following nodes
-        wait_event_default_1 = torch.ops.streams.wait_event.default(7, 2);  wait_event_default_1 = None
+        sync_dealloc_default_2 = torch.ops.streams.sync_dealloc.default(12, 3, getitem_11);  getitem_11 = sync_dealloc_default_2 = None
 
         # Annotation: {'stream': 2}
-        mul_11: "f32[2, 2]" = torch.ops.aten.mul.Tensor(mul_8, mul_1);  mul_1 = None
+        mul_12: "f32[2, 2]" = torch.ops.aten.mul.Tensor(getitem_12, getitem);  getitem_12 = getitem = None
 
         # No stacktrace found for following nodes
-        record_event_default_5 = torch.ops.streams.record_event.default(11, 2);  record_event_default_5 = None
-        sync_dealloc_default_1 = torch.ops.streams.sync_dealloc.default(11, 3, mul_8);  mul_8 = sync_dealloc_default_1 = None
-        record_event_default_2 = torch.ops.streams.record_event.default(8, 2);  record_event_default_2 = None
-        wait_event_default_2 = torch.ops.streams.wait_event.default(8, 1);  wait_event_default_2 = None
+        subgraph_record_event_default_3 = self.subgraph_record_event_default_3
+        control_deps_15 = torch.ops.higher_order.control_deps((mul_12,), subgraph_record_event_default_3, mul_12);  mul_12 = subgraph_record_event_default_3 = None
 
-        # Annotation: {'stream': 1}
-        add_1: "f32[2, 2]" = torch.ops.aten.add.Tensor(tangents_1, mul_11);  tangents_1 = None
+        # Annotation: {'stream': 2}
+        getitem_13: "f32[2, 2]" = control_deps_15[1]
 
         # No stacktrace found for following nodes
-        record_event_default_6 = torch.ops.streams.record_event.default(12, 1);  record_event_default_6 = None
-        sync_dealloc_default_2 = torch.ops.streams.sync_dealloc.default(12, 2, mul_11);  mul_11 = sync_dealloc_default_2 = None
+        subgraph_wait_event_default_3 = self.subgraph_wait_event_default_3
+        control_deps_16 = torch.ops.higher_order.control_deps((control_deps_15, mul_9, mul_10), subgraph_wait_event_default_3, mul_10);  control_deps_15 = mul_9 = mul_10 = subgraph_wait_event_default_3 = None
 
-        # Annotation: {'stream': 1}
-        mul_12: "f32[2, 2]" = torch.ops.aten.mul.Tensor(add_1, ones);  add_1 = ones = None
+        # Annotation: {'stream': 4}
+        getitem_14: "f32[2, 2]" = control_deps_16[1];  control_deps_16 = None
+
+        # Annotation: {'stream': 4}
+        add_2: "f32[2, 2]" = torch.ops.aten.add.Tensor(getitem_14, getitem_13);  getitem_14 = None
 
         # No stacktrace found for following nodes
-        record_event_default_3 = torch.ops.streams.record_event.default(9, 1);  record_event_default_3 = None
-        wait_event_default_3 = torch.ops.streams.wait_event.default(9, 3);  wait_event_default_3 = None
+        subgraph_record_event_default_7 = self.subgraph_record_event_default_7
+        control_deps_17 = torch.ops.higher_order.control_deps((add_2,), subgraph_record_event_default_7, add_2);  add_2 = subgraph_record_event_default_7 = None
 
-        # Annotation: {'stream': 3}
-        add_2: "f32[2, 2]" = torch.ops.aten.add.Tensor(mul_10, mul_12);  mul_10 = None
+        # Annotation: {'stream': 4}
+        getitem_15: "f32[2, 2]" = control_deps_17[1];  control_deps_17 = None
 
         # No stacktrace found for following nodes
-        record_event_default_7 = torch.ops.streams.record_event.default(13, 3);  record_event_default_7 = None
-        sync_dealloc_default_3 = torch.ops.streams.sync_dealloc.default(13, 1, mul_12);  mul_12 = sync_dealloc_default_3 = None
-        return (add_2,)
-""",
+        sync_dealloc_default_3 = torch.ops.streams.sync_dealloc.default(13, 2, getitem_13);  getitem_13 = sync_dealloc_default_3 = None
+        return (getitem_15,)
+
+    class subgraph_record_event_default(torch.nn.Module):
+        def forward(self, dep_0: "f32[2, 2]"):
+            # No stacktrace found for following nodes
+            record_event_default = torch.ops.streams.record_event.default(6, 5)
+            return (record_event_default, dep_0)
+
+    class subgraph_wait_event_default(torch.nn.Module):
+        def forward(self):
+            # No stacktrace found for following nodes
+            wait_event_default = torch.ops.streams.wait_event.default(6, 4)
+            return wait_event_default
+
+    class subgraph_record_event_default_4(torch.nn.Module):
+        def forward(self, dep_0: "f32[2, 2]"):
+            # No stacktrace found for following nodes
+            record_event_default = torch.ops.streams.record_event.default(10, 4)
+            return (record_event_default, dep_0)
+
+    class subgraph_record_event_default_1(torch.nn.Module):
+        def forward(self, dep_0: "f32[2, 2]"):
+            # No stacktrace found for following nodes
+            record_event_default = torch.ops.streams.record_event.default(7, 4)
+            return (record_event_default, dep_0)
+
+    class subgraph_wait_event_default_1(torch.nn.Module):
+        def forward(self):
+            # No stacktrace found for following nodes
+            wait_event_default = torch.ops.streams.wait_event.default(7, 3)
+            return wait_event_default
+
+    class subgraph_record_event_default_5(torch.nn.Module):
+        def forward(self, dep_0: "f32[2, 2]"):
+            # No stacktrace found for following nodes
+            record_event_default = torch.ops.streams.record_event.default(11, 3)
+            return (record_event_default, dep_0)
+
+    class subgraph_record_event_default_6(torch.nn.Module):
+        def forward(self, dep_0: "f32[2, 2]"):
+            # No stacktrace found for following nodes
+            record_event_default = torch.ops.streams.record_event.default(12, 2)
+            return (record_event_default, dep_0)
+
+    class subgraph_record_event_default_3(torch.nn.Module):
+        def forward(self, dep_0: "f32[2, 2]"):
+            # No stacktrace found for following nodes
+            record_event_default = torch.ops.streams.record_event.default(9, 2)
+            return (record_event_default, dep_0)
+
+    class subgraph_wait_event_default_3(torch.nn.Module):
+        def forward(self, dep_0: "f32[2, 2]"):
+            # No stacktrace found for following nodes
+            wait_event_default = torch.ops.streams.wait_event.default(9, 4)
+            return (wait_event_default, dep_0)
+
+    class subgraph_record_event_default_7(torch.nn.Module):
+        def forward(self, dep_0: "f32[2, 2]"):
+            # No stacktrace found for following nodes
+            record_event_default = torch.ops.streams.record_event.default(13, 4)
+            return (record_event_default, dep_0)
+""",  # noqa: B950
         )
 
     @requires_cuda
@@ -900,7 +1033,7 @@ def fn(x):
             """\
 class <lambda>(torch.nn.Module):
     def forward(self, arg0_1: "f32[2, 2]"):
-        # Annotation: {'stream': 0}
+        # Annotation: {'stream': 1}
         add: "f32[2, 2]" = torch.ops.aten.add.Tensor(arg0_1, 2)
         copy_: "f32[2, 2]" = torch.ops.aten.copy_.default(arg0_1, add);  arg0_1 = add = None
         return (copy_,)
@@ -909,7 +1042,7 @@ def forward(self, arg0_1: "f32[2, 2]"):
 
     @requires_cuda
     def test_epilogue_copy_streams_external(self):
-        @torch.compile
+        @torch.compile(backend="eager")
         def fn(x):
             with torch.Stream(device="cuda:0"):
                 x.mul_(3)
@@ -924,115 +1057,1191 @@ def fn(x):
             extract_graph(fn, *inp)
 
     @requires_cuda
-    def test_epilogue_copy_stream_tracking(self):
-        """
-        Test that epilogue copies for mutated inputs use the correct stream.
-        This verifies that ViewAndMutationMeta.mutated_inp_stream_indices is
-        properly populated and used at runtime.
-        Uses a custom autograd.Function where the backward mutates a saved
-        tensor on a specific stream.
+    def test_control_deps_wrapping_record_event(self) -> None:
+        """Test wrapping record_event with control_deps on a two-stream graph.
+
+        The producer stream has work before the record_event, so control_deps
+        should capture those same-stream nodes as dependencies.
         """
 
-        class BwMutationWithStream(torch.autograd.Function):
-            @staticmethod
-            def forward(ctx, x, y):
-                ctx.save_for_backward(x)
-                ctx.s1 = torch.Stream(device="cuda:0")
-                ctx.s2 = torch.Stream(device="cuda:0")
-                # Do computation on stream s2
-                with ctx.s2:
-                    result = x * 2 + y
-                return result
+        def fn(x) -> torch.Tensor:
+            s1 = torch.Stream(device="cuda")
+            s2 = torch.Stream(device="cuda")
+            e = torch.Event()
 
-            @staticmethod
-            def backward(ctx, grad_output):
-                (x,) = ctx.saved_tensors
-                # Mutate saved tensor x on stream s1 in backward
-                with ctx.s1:
-                    x.mul_(2)
-                # Compute gradients on stream s2
-                with ctx.s2:
-                    grad_x = grad_output * 2
-                    grad_y = grad_output.clone()
-                return grad_x, grad_y, None, None
+            with s1:
+                y = x + 1
+                z = y * 2
+                e.record()
 
-        def fn(x, y):
-            result = BwMutationWithStream.apply(x, y)
-            return result
+            with s2:
+                e.wait()
+                w = z + 3
 
-        x = torch.ones(2, 2, requires_grad=True, device="cuda:0")
-        y = torch.ones(2, 2, requires_grad=True, device="cuda:0")
+            return w
+
+        inp = (torch.ones(2, 2, device="cuda"),)
         (
-            actual,
+            _,
             _,
             fw_graphs,
-            bw_graphs,
-        ) = extract_graph(fn, x.clone(), y.clone())
-        self.assertEqual(len(fw_graphs), 1)
-        # Forward graph should show computation on stream 1 (s2)
-        self.assertExpectedInline(
-            print_graph(fw_graphs[0]),
-            """\
-class GraphModule(torch.nn.Module):
-    def forward(self, primals_1: "f32[2, 2]", primals_2: "f32[2, 2]"):
-        # Annotation: {'stream': 1}
-        mul: "f32[2, 2]" = torch.ops.aten.mul.Tensor(primals_1, 2)
-        add: "f32[2, 2]" = torch.ops.aten.add.Tensor(mul, primals_2);  primals_2 = None
-        return (add, primals_1, mul)
-""",
+            _,
+        ) = extract_graph(fn, *inp)
+
+        gm = fw_graphs[0]
+        graph = gm.graph
+
+        import operator
+
+        from torch._functorch._aot_autograd.streams import (
+            wrap_all_sync_nodes_with_control_deps,
         )
-        # Run backward and check that the epilogue copy uses stream 0 (s1)
-        actual.sum().backward()
-        # The backward graph should show:
-        # 1. Mutation happening on stream 0 (s1)
-        # 2. Gradient computation on stream 1 (s2)
-        # 3. Epilogue copy for the mutated tensor on stream 0 (s1)
-        self.assertExpectedInline(
-            print_graph(bw_graphs[0]),
-            """\
-class GraphModule(torch.nn.Module):
-    def forward(self, primals_1: "f32[2, 2]", mul: "f32[2, 2]", tangents_1: "f32[2, 2]"):
-        # Annotation: {'stream': 1}
-        mul_2: "f32[2, 2]" = torch.ops.aten.mul.Tensor(tangents_1, 2)
+        from torch._inductor.fx_passes.control_dependencies import control_deps
 
-        # Annotation: {'stream': 1}
-        clone: "f32[2, 2]" = torch.ops.aten.clone.default(tangents_1);  tangents_1 = None
+        wrap_all_sync_nodes_with_control_deps(gm)
 
-        # Annotation: {'stream': 0} No stacktrace found for following nodes
-        copy_: "f32[2, 2]" = torch.ops.aten.copy_.default(primals_1, mul);  primals_1 = mul = copy_ = None
-        return (mul_2, clone)
-""",
+        # record_event has same-stream deps, and wait_event gets a
+        # cross-event dependency on the record's control_deps node.
+        control_deps_nodes = list(
+            graph.find_nodes(op="call_function", target=control_deps)
         )
+        self.assertEqual(len(control_deps_nodes), 2)
+        record_ctrl_node = control_deps_nodes[0]
+
+        # Verify getitem nodes were created for pass-through dependencies
+        getitem_nodes = [
+            n
+            for n in graph.nodes
+            if n.op == "call_function"
+            and n.target == operator.getitem
+            and n.args[0] is record_ctrl_node
+        ]
+        self.assertGreaterEqual(len(getitem_nodes), 1)
+
+        # Verify the wait's control_deps depends on the record's control_deps
+        wait_ctrl_node = control_deps_nodes[1]
+        self.assertIn(record_ctrl_node, wait_ctrl_node.args[0])
 
     @requires_cuda
-    def test_inductor_lowering(self):
-        with patch("torch._inductor.config.implicit_fallbacks", False):
+    def test_control_deps_wrapping_wait_event(self) -> None:
+        """Test wrapping wait_event with control_deps on a two-stream graph.
 
-            @torch.compile()
-            def fn(x):
-                e = torch.Event()
-                x += x + 1
+        The consumer stream has the wait_event, and there should be same-stream
+        nodes before it that become dependencies.
+        """
+
+        def fn(x) -> torch.Tensor:
+            s1 = torch.Stream(device="cuda")
+            s2 = torch.Stream(device="cuda")
+            e = torch.Event()
+
+            with s1:
+                y = x + 1
                 e.record()
-                return x
 
-            inp = (torch.ones(2, 2, device="cuda"),)
-            fn(*inp)
+            with s2:
+                a = x * 3
+                e.wait()
+                z = y * a
 
-    def test_is_marked_side_effectful(self):
-        self.assertIn(
-            torch.ops.streams.fork.default, torch.fx.node._side_effectful_functions
-        )
-        self.assertIn(
-            torch.ops.streams.join.default, torch.fx.node._side_effectful_functions
-        )
-        self.assertIn(
-            torch.ops.streams.wait_event.default,
-            torch.fx.node._side_effectful_functions,
-        )
+            return z
+
+        inp = (torch.ones(2, 2, device="cuda"),)
+        (
+            _,
+            _,
+            fw_graphs,
+            _,
+        ) = extract_graph(fn, *inp)
+
+        gm = fw_graphs[0]
+        graph = gm.graph
+
+        from torch._functorch._aot_autograd.streams import (
+            wrap_all_sync_nodes_with_control_deps,
+        )
+        from torch._inductor.fx_passes.control_dependencies import control_deps
+
+        wrap_all_sync_nodes_with_control_deps(gm)
+
+        # Both record_event (deps: [y]) and wait_event (deps: [a]) have
+        # same-stream deps, so both get wrapped.
+        control_deps_nodes = list(
+            graph.find_nodes(op="call_function", target=control_deps)
+        )
+        self.assertEqual(len(control_deps_nodes), 2)
+
+    @requires_cuda
+    def test_control_deps_prevents_invalid_reordering(self) -> None:
+        """
+        Test that control_deps creates proper data dependencies that prevent invalid reordering.
+
+        This test:
+        1. Creates a two-stream graph with control_deps wrapping
+        2. Manually moves a node to violate the ordering
+        3. Verifies graph.lint() catches the invalid ordering
+        """
+
+        def fn(x) -> torch.Tensor:
+            s1 = torch.Stream(device="cuda")
+            s2 = torch.Stream(device="cuda")
+            e = torch.Event()
+
+            with s1:
+                y = x + 1
+                z = y * 2
+                e.record()
+
+            with s2:
+                e.wait()
+                w = z + 3
+
+            return w
+
+        inp = (torch.ones(2, 2, device="cuda"),)
+        (
+            _,
+            _,
+            fw_graphs,
+            _,
+        ) = extract_graph(fn, *inp)
+
+        gm = fw_graphs[0]
+        graph = gm.graph
+
+        from torch._functorch._aot_autograd.streams import (
+            wrap_all_sync_nodes_with_control_deps,
+        )
+        from torch._inductor.fx_passes.control_dependencies import control_deps
+
+        wrap_all_sync_nodes_with_control_deps(gm)
+
+        control_deps_node = next(
+            iter(graph.find_nodes(op="call_function", target=control_deps))
+        )
+
+        # Find add_1 (the node after the sync that uses getitem output)
+        add_1_node = next(n for n in graph.nodes if n.name == "add_1")
+
+        # Verify valid ordering: control_deps before add_1
+        original_order = [n.name for n in graph.nodes if n.op == "call_function"]
+        self.assertLess(
+            original_order.index("control_deps"),
+            original_order.index("add_1"),
+            "control_deps should be before add_1 in valid ordering",
+        )
+
+        graph.lint()  # Should not raise
+
+        # Manually move add_1 BEFORE control_deps (violates dependency)
+        control_deps_node.prepend(add_1_node)
+
+        # Verify the order is now wrong (add_1 comes before control_deps)
+        violated_order = [n.name for n in graph.nodes if n.op == "call_function"]
+        self.assertLess(
+            violated_order.index("add_1"),
+            violated_order.index("control_deps"),
+            "add_1 should be before control_deps after manual move (violating order)",
+        )
+
+        # graph.lint() should catch the invalid ordering
+        with self.assertRaises(RuntimeError):
+            graph.lint()
+
+    @requires_cuda
+    def test_cross_event_deps_multiple_events(self) -> None:
+        """Stress test: multiple events across three streams.
+
+        Verifies that each wait_event's control_deps node depends on exactly
+        the matching record_event's control_deps node (by event index), not
+        on some other event's node.
+
+        Graph layout:
+          s1: x+1 -> record(e1) -> x+5 -> record(e2)
+          s2: wait(e1) -> x*3
+          s3: wait(e2) -> x*7
+        """
+
+        def fn(x) -> torch.Tensor:
+            s1 = torch.Stream(device="cuda")
+            s2 = torch.Stream(device="cuda")
+            s3 = torch.Stream(device="cuda")
+            e1 = torch.Event()
+            e2 = torch.Event()
+
+            with s1:
+                y = x + 1
+                e1.record()
+                z = x + 5
+                e2.record()
+
+            with s2:
+                e1.wait()
+                a = x * 3
+
+            with s3:
+                e2.wait()
+                b = x * 7
+
+            return a + b + y + z
+
+        inp = (torch.ones(2, 2, device="cuda"),)
+        # Patch out wrapping so we get the raw graph to manually wrap below.
+        with patch(
+            "torch._functorch._aot_autograd.graph_capture.wrap_all_sync_nodes_with_control_deps"
+        ):
+            (
+                _,
+                _,
+                fw_graphs,
+                _,
+            ) = extract_graph(fn, *inp)
+
+        gm = fw_graphs[0]
+        graph = gm.graph
+
+        from torch._functorch._aot_autograd.streams import (
+            wrap_all_sync_nodes_with_control_deps,
+        )
+        from torch._inductor.fx_passes.control_dependencies import control_deps
+
+        # Capture event indices before wrapping so we can verify correspondence.
+        # Each sync node has args (event_index, stream_index).
+        record_event_indices = []
+        wait_event_indices = []
+        for node in graph.nodes:
+            if node.op != "call_function":
+                continue
+            if node.target is torch.ops.streams.record_event.default:
+                record_event_indices.append(node.args[0])
+            elif node.target is torch.ops.streams.wait_event.default:
+                wait_event_indices.append(node.args[0])
+
+        self.assertEqual(len(record_event_indices), 2)
+        self.assertEqual(len(wait_event_indices), 2)
+        # Each wait corresponds to exactly one record by event index
+        self.assertEqual(set(record_event_indices), set(wait_event_indices))
+
+        wrap_all_sync_nodes_with_control_deps(gm)
+
+        # 2 records + 2 waits = 4 control_deps nodes. Each record has
+        # same-stream deps, and each wait gets a cross-event dep.
+        ctrl_nodes = list(graph.find_nodes(op="call_function", target=control_deps))
+        self.assertEqual(len(ctrl_nodes), 4)
+
+        # Graph order follows program order: both records (s1 block) before
+        # both waits (s2, s3 blocks).
+        record_e1_ctrl = ctrl_nodes[0]
+        record_e2_ctrl = ctrl_nodes[1]
+        wait_e1_ctrl = ctrl_nodes[2]
+        wait_e2_ctrl = ctrl_nodes[3]
+
+        # wait(e1) depends on record(e1), not record(e2)
+        self.assertIn(record_e1_ctrl, wait_e1_ctrl.args[0])
+        self.assertNotIn(record_e2_ctrl, wait_e1_ctrl.args[0])
+
+        # wait(e2) depends on record(e2), not record(e1)
+        self.assertIn(record_e2_ctrl, wait_e2_ctrl.args[0])
+        self.assertNotIn(record_e1_ctrl, wait_e2_ctrl.args[0])
+
+        graph.lint()
+
+    @requires_cuda
+    def test_cross_event_deps_event_reuse(self) -> None:
+        """Test that reusing an event updates the cross-event dependency.
+
+        When the same event is recorded twice with work in between, the wait
+        should depend on the LAST record's control_deps node (matching CUDA
+        semantics where record() overwrites the event).
+        """
+
+        def fn(x) -> torch.Tensor:
+            s1 = torch.Stream(device="cuda")
+            s2 = torch.Stream(device="cuda")
+            e = torch.Event()
+
+            with s1:
+                y = x + 1
+                e.record()  # first record
+                z = y * 2
+                e.record()  # second record, reuses same event
+
+            with s2:
+                e.wait()
+                w = x + 3
+
+            return w + z
+
+        inp = (torch.ones(2, 2, device="cuda"),)
+        (
+            _,
+            _,
+            fw_graphs,
+            _,
+        ) = extract_graph(fn, *inp)
+
+        gm = fw_graphs[0]
+        graph = gm.graph
+
+        from torch._functorch._aot_autograd.streams import (
+            wrap_all_sync_nodes_with_control_deps,
+        )
+        from torch._inductor.fx_passes.control_dependencies import control_deps
+
+        wrap_all_sync_nodes_with_control_deps(gm)
+
+        ctrl_nodes = list(graph.find_nodes(op="call_function", target=control_deps))
+        # 2 records + 1 wait = 3 control_deps nodes
+        self.assertEqual(len(ctrl_nodes), 3)
+
+        record1_ctrl = ctrl_nodes[0]
+        record2_ctrl = ctrl_nodes[1]
+        wait_ctrl = ctrl_nodes[2]
+
+        # The wait depends on the SECOND record (last write wins), not the first
+        self.assertIn(record2_ctrl, wait_ctrl.args[0])
+        self.assertNotIn(record1_ctrl, wait_ctrl.args[0])
+
+        graph.lint()
+
+    @requires_cuda
+    def test_epilogue_copy_stream_tracking(self):
+        """
+        Test that epilogue copies for mutated inputs use the correct stream.
+        This verifies that ViewAndMutationMeta.mutated_inp_stream_indices is
+        properly populated and used at runtime.
+        Uses a custom autograd.Function where the backward mutates a saved
+        tensor on a specific stream.
+        """
+
+        class BwMutationWithStream(torch.autograd.Function):
+            @staticmethod
+            def forward(ctx, x, y):
+                ctx.save_for_backward(x)
+                ctx.s1 = torch.Stream(device="cuda:0")
+                ctx.s2 = torch.Stream(device="cuda:0")
+                # Do computation on stream s2
+                with ctx.s2:
+                    result = x * 2 + y
+                return result
+
+            @staticmethod
+            def backward(ctx, grad_output):
+                (x,) = ctx.saved_tensors
+                # Mutate saved tensor x on stream s1 in backward
+                with ctx.s1:
+                    x.mul_(2)
+                # Compute gradients on stream s2
+                with ctx.s2:
+                    grad_x = grad_output * 2
+                    grad_y = grad_output.clone()
+                return grad_x, grad_y, None, None
+
+        def fn(x, y):
+            result = BwMutationWithStream.apply(x, y)
+            return result
+
+        x = torch.ones(2, 2, requires_grad=True, device="cuda:0")
+        y = torch.ones(2, 2, requires_grad=True, device="cuda:0")
+        (
+            actual,
+            _,
+            fw_graphs,
+            bw_graphs,
+        ) = extract_graph(fn, x.clone(), y.clone())
+        self.assertEqual(len(fw_graphs), 1)
+        # Forward graph should show computation on stream 1 (s2)
+        self.assertExpectedInline(
+            print_graph(fw_graphs[0]),
+            """\
+class GraphModule(torch.nn.Module):
+    def forward(self, primals_1: "f32[2, 2]", primals_2: "f32[2, 2]"):
+        # Annotation: {'stream': 2}
+        mul: "f32[2, 2]" = torch.ops.aten.mul.Tensor(primals_1, 2)
+        add: "f32[2, 2]" = torch.ops.aten.add.Tensor(mul, primals_2);  primals_2 = None
+        return (add, primals_1, mul)
+""",
+        )
+        # Run backward and check that the epilogue copy uses stream 0 (s1)
+        actual.sum().backward()
+        # The backward graph should show:
+        # 1. Mutation happening on stream 0 (s1)
+        # 2. Gradient computation on stream 1 (s2)
+        # 3. Epilogue copy for the mutated tensor on stream 0 (s1)
+        self.assertExpectedInline(
+            print_graph(bw_graphs[0]),
+            """\
+class GraphModule(torch.nn.Module):
+    def forward(self, primals_1: "f32[2, 2]", mul: "f32[2, 2]", tangents_1: "f32[2, 2]"):
+        # Annotation: {'stream': 2}
+        mul_2: "f32[2, 2]" = torch.ops.aten.mul.Tensor(tangents_1, 2)
+
+        # Annotation: {'stream': 2}
+        clone: "f32[2, 2]" = torch.ops.aten.clone.default(tangents_1);  tangents_1 = None
+
+        # Annotation: {'stream': 1} No stacktrace found for following nodes
+        copy_: "f32[2, 2]" = torch.ops.aten.copy_.default(primals_1, mul);  primals_1 = mul = copy_ = None
+        return (mul_2, clone)
+""",
+        )
+
+    @requires_cuda
+    def test_inductor_lowering(self):
+        with patch("torch._inductor.config.implicit_fallbacks", False):
+
+            @torch.compile()
+            def fn(x):
+                e = torch.Event()
+                x += x + 1
+                e.record()
+                return x
+
+            inp = (torch.ones(2, 2, device="cuda"),)
+            fn(*inp)
+
+    def test_is_marked_side_effectful(self):
+        self.assertIn(
+            torch.ops.streams.fork.default, torch.fx.node._side_effectful_functions
+        )
+        self.assertIn(
+            torch.ops.streams.join.default, torch.fx.node._side_effectful_functions
+        )
+        self.assertIn(
+            torch.ops.streams.wait_event.default,
+            torch.fx.node._side_effectful_functions,
+        )
         self.assertIn(
             torch.ops.streams.record_event.default,
             torch.fx.node._side_effectful_functions,
         )
+        self.assertIn(
+            torch.ops.streams.synchronize_event.default,
+            torch.fx.node._side_effectful_functions,
+        )
+
+    @requires_cuda
+    def test_backward_sync_control_deps_e2e(self) -> None:
+        """
+        End-to-end test verifying that backward sync nodes are wrapped with control_deps.
+
+        This test:
+        1. Creates a function with multiple streams
+        2. Compiles it and runs backward
+        3. Verifies control_deps nodes are present in the backward graph
+        """
+        from torch._inductor.fx_passes.control_dependencies import control_deps
+
+        def fn(x, y):
+            s0 = torch.Stream()
+            s1 = torch.Stream()
+            with s0:
+                z0 = 2 * x + y
+            with s1:
+                z1 = 3 * x + y
+            return z0, z1
+
+        inp = (
+            torch.ones(2, 2, device="cuda:0", requires_grad=True) + 1,
+            torch.ones(2, 2, device="cuda:0", requires_grad=True),
+        )
+
+        (
+            actual,
+            _,
+            fw_graphs,
+            bw_graphs,
+        ) = extract_graph(fn, *inp)
+
+        # Run backward to trigger backward graph capture
+        actual[0].sum().backward()
+
+        # Check that backward graph has control_deps nodes
+        self.assertGreaterEqual(len(bw_graphs), 1)
+        bw_graph = bw_graphs[0]
+
+        control_deps_nodes = list(
+            bw_graph.graph.find_nodes(op="call_function", target=control_deps)
+        )
+
+        # Should have at least one control_deps node wrapping sync operations
+        self.assertGreater(
+            len(control_deps_nodes),
+            0,
+            "Expected control_deps nodes in backward graph for stream synchronization",
+        )
+
+    def test_sync_dealloc_has_fake_impl(self):
+        """Test that sync_dealloc has a registered fake impl.
+
+        Without a fake impl, Inductor's backward compilation crashes when the
+        backward graph contains cross-stream sync_dealloc ops.
+        """
+        from torch._subclasses.fake_tensor import FakeTensorMode
+
+        with FakeTensorMode():
+            t = torch.randn(4)
+            # Should not raise "no fake impl registered"
+            torch.ops.streams.sync_dealloc.default(0, 1, t)
+
+    def test_record_stream_has_fake_impl(self):
+        """Test that record_stream's fake impl has the correct signature."""
+        from torch._subclasses.fake_tensor import FakeTensorMode
+
+        with FakeTensorMode():
+            t = torch.randn(4)
+            # Should not raise due to signature mismatch
+            torch.ops.streams.record_stream.default(t, 0)
+
+    @requires_cuda
+    def test_record_stream(self):
+        backend = torch._dynamo.testing.EagerAndRecordGraphs()
+
+        def fn(x):
+            s = torch.Stream()
+            x.record_stream(s)
+            return x
+
+        compiled = torch.compile(fn, backend=backend, fullgraph=True)
+        compiled(torch.randn(4, device="cuda"))
+
+        self.assertEqual(len(backend.graphs), 1)
+        found = any(
+            node.target is torch.ops.streams.record_stream
+            for node in backend.graphs[0].graph.nodes
+        )
+        self.assertTrue(found, "record_stream op not found in graph")
+
+    @requires_cuda
+    def test_event_record_after_input_mutation_errors(self):
+        def fn(x):
+            s = torch.Stream()
+            e = torch.Event()
+            with s:
+                x.add_(1)
+                e.record()
+            return e
+
+        with self.assertRaisesRegex(RuntimeError, "An event was recorded on a stream"):
+            torch.compile(fn, backend="eager", fullgraph=True)(
+                torch.ones(2, 2, device="cuda")
+            )
+
+    @requires_cuda
+    def test_event_record_after_input_mutation_stack_traces(self):
+        def fn(x):
+            s = torch.Stream()
+            e = torch.Event()
+            with s:
+                x.add_(1)
+                e.record()
+            return e
+
+        try:
+            torch.compile(fn, backend="eager", fullgraph=True)(
+                torch.ones(2, 2, device="cuda")
+            )
+            self.fail("Expected RuntimeError")
+        except RuntimeError as e:
+            msg = str(e)
+            self.assertIn("Input mutation occurred here:", msg)
+            self.assertIn("x.add_(1)", msg)
+            self.assertIn("Event record occurred here:", msg)
+            self.assertIn("e.record()", msg)
+
+    @requires_cuda
+    def test_event_record_after_input_mutation_record_event(self):
+        def fn(x):
+            s = torch.Stream()
+            with s:
+                x.add_(1)
+                e = s.record_event()
+            return e
+
+        with self.assertRaisesRegex(RuntimeError, "An event was recorded on a stream"):
+            torch.compile(fn, backend="eager", fullgraph=True)(
+                torch.ones(2, 2, device="cuda")
+            )
+
+    @requires_cuda
+    def test_event_record_after_input_mutation_through_view(self):
+        def fn(x):
+            s = torch.Stream()
+            e = torch.Event()
+            v = x.view(-1)
+            with s:
+                v.add_(1)
+                e.record()
+            return e
+
+        with self.assertRaisesRegex(RuntimeError, "An event was recorded on a stream"):
+            torch.compile(fn, backend="eager", fullgraph=True)(
+                torch.ones(2, 2, device="cuda")
+            )
+
+    @requires_cuda
+    def test_event_record_after_input_mutation_input_event(self):
+        def fn(x, e):
+            s = torch.Stream()
+            with s:
+                x.add_(1)
+                e.record()
+            return x
+
+        with self.assertRaisesRegex(RuntimeError, "An event was recorded on a stream"):
+            torch.compile(fn, backend="eager", fullgraph=True)(
+                torch.ones(2, 2, device="cuda"),
+                torch.Event(),
+            )
+
+    @requires_cuda
+    def test_event_record_before_input_mutation_no_error(self):
+        def fn(x):
+            s = torch.Stream()
+            e = torch.Event()
+            with s:
+                e.record()
+                x.add_(1)
+            return e
+
+        torch.compile(fn, backend="eager", fullgraph=True)(
+            torch.ones(2, 2, device="cuda")
+        )
+
+    @requires_cuda
+    def test_event_record_on_different_stream_no_error(self):
+        def fn(x):
+            s0 = torch.Stream()
+            s1 = torch.Stream()
+            e = torch.Event()
+            with s0:
+                x.add_(1)
+            with s1:
+                e.record()
+            return e
+
+        torch.compile(fn, backend="eager", fullgraph=True)(
+            torch.ones(2, 2, device="cuda")
+        )
+
+    @requires_cuda
+    def test_event_not_returned_no_error(self):
+        def fn(x):
+            s = torch.Stream()
+            e = torch.Event()
+            with s:
+                x.add_(1)
+                e.record()
+            return x
+
+        with self.assertRaisesRegex(RuntimeError, "An event was recorded on a stream"):
+            torch.compile(fn, backend="eager", fullgraph=True)(
+                torch.ones(2, 2, device="cuda")
+            )
+
+    @requires_cuda
+    @unittest.skip("https://github.com/pytorch/pytorch/issues/177771")
+    def test_cuda_event_record_on_stream(self):
+        """torch.cuda.Event should be accepted by torch.Stream.record_event (C++ type check)."""
+        s = torch.Stream(device="cuda")
+        e = torch.cuda.Event()
+        # This hits THPStream_record_event in Stream.cpp which does a type check
+        s.record_event(e)
+
+    @requires_cuda
+    def test_event_synchronize_tracing(self):
+        def fn(x):
+            e = torch.Event()
+            e.record()
+            x = x + 1
+            e.synchronize()
+            return x
+
+        inp = (torch.ones(2, 2, device="cuda"),)
+        (
+            _,
+            _,
+            fw_graphs,
+            _,
+        ) = extract_graph(fn, *inp)
+
+        self.assertExpectedInline(
+            print_graph(fw_graphs[0]),
+            """\
+class <lambda>(torch.nn.Module):
+    def forward(self, arg0_1: "f32[2, 2]"):
+        #
+        record_event = torch.ops.streams.record_event.default(1, 0);  record_event = None
+
+        #
+        add: "f32[2, 2]" = torch.ops.aten.add.Tensor(arg0_1, 1);  arg0_1 = None
+
+        # No stacktrace found for following nodes
+        subgraph_synchronize_event = self.subgraph_synchronize_event
+        control_deps = torch.ops.higher_order.control_deps((add,), subgraph_synchronize_event, add);  subgraph_synchronize_event = control_deps = None
+        return (add,)
+
+    class subgraph_synchronize_event(torch.nn.Module):
+        def forward(self, dep_0: "f32[2, 2]"):
+            #
+            synchronize_event_default = torch.ops.streams.synchronize_event.default(1)
+            return (synchronize_event_default, dep_0)
+""",  # noqa: B950
+        )
+
+    @requires_cuda
+    def test_event_synchronize_inductor_lowering(self):
+        with patch("torch._inductor.config.implicit_fallbacks", False):
+
+            @torch.compile()
+            def fn(x):
+                e = torch.Event()
+                x = x + 1
+                e.record()
+                e.synchronize()
+                return x
+
+            inp = (torch.ones(2, 2, device="cuda"),)
+            fn(*inp)
+
+    @requires_cuda
+    def test_control_deps_wrapping_synchronize_event(self) -> None:
+        """Test that synchronize_event threads recorded ops' values through.
+
+        After record_event wraps ops in control_deps and produces getitem
+        pass-throughs, synchronize_event must also thread those through so
+        that subsequent consumers depend on the synchronize.
+        """
+
+        def fn(x) -> torch.Tensor:
+            e = torch.Event()
+            y = x + 1
+            e.record()
+            e.synchronize()
+            # z uses y which was produced before the record — its value must
+            # be threaded through both record and synchronize control_deps.
+            z = y * 2
+            return z
+
+        inp = (torch.ones(2, 2, device="cuda"),)
+        # Patch out wrapping so we get the raw graph to manually wrap below.
+        with patch(
+            "torch._functorch._aot_autograd.graph_capture.wrap_all_sync_nodes_with_control_deps"
+        ):
+            (
+                _,
+                _,
+                fw_graphs,
+                _,
+            ) = extract_graph(fn, *inp)
+
+        gm = fw_graphs[0]
+        graph = gm.graph
+
+        import operator
+
+        from torch._functorch._aot_autograd.streams import (
+            set_stream,
+            wrap_all_sync_nodes_with_control_deps,
+        )
+        from torch._inductor.fx_passes.control_dependencies import control_deps
+
+        # extract_graph doesn't annotate streams, so set stream metadata on
+        # compute nodes to match the record_event's stream index.
+        record_node = next(
+            n
+            for n in graph.nodes
+            if n.op == "call_function"
+            and n.target is torch.ops.streams.record_event.default
+        )
+        stream_idx = record_node.args[1]
+        for n in graph.nodes:
+            if (
+                n.op == "call_function"
+                and "val" in n.meta
+                and n.target
+                not in (
+                    torch.ops.streams.record_event.default,
+                    torch.ops.streams.synchronize_event.default,
+                )
+            ):
+                set_stream(n, stream_idx)
+
+        wrap_all_sync_nodes_with_control_deps(gm)
+
+        ctrl_nodes = list(graph.find_nodes(op="call_function", target=control_deps))
+        # record_event + synchronize_event = 2 control_deps nodes
+        self.assertEqual(len(ctrl_nodes), 2)
+        record_ctrl = ctrl_nodes[0]
+        sync_ctrl = ctrl_nodes[1]
+
+        # synchronize_event's control_deps should depend on record's ctrl
+        self.assertIn(record_ctrl, sync_ctrl.args[0])
+
+        # The record should thread through the add (y = x + 1)
+        record_getitems = [
+            n
+            for n in graph.nodes
+            if n.op == "call_function"
+            and n.target == operator.getitem
+            and n.args[0] is record_ctrl
+        ]
+        self.assertGreaterEqual(len(record_getitems), 1)
+
+        # Those getitems should be passed through synchronize's control_deps
+        # as additional args (the passthrough deps)
+        sync_passthrough_args = sync_ctrl.args[2:]  # skip (deps_tuple, subgraph)
+        for getitem in record_getitems:
+            self.assertIn(
+                getitem,
+                sync_passthrough_args,
+                "record_event's getitem should be threaded through synchronize_event",
+            )
+
+        # The mul (z = y * 2) should consume a getitem from synchronize's
+        # control_deps, not directly from record's.
+        sync_getitems = [
+            n
+            for n in graph.nodes
+            if n.op == "call_function"
+            and n.target == operator.getitem
+            and n.args[0] is sync_ctrl
+        ]
+        self.assertGreaterEqual(len(sync_getitems), 1)
+
+        # Find the mul node and verify it uses a sync getitem
+        mul_nodes = [
+            n
+            for n in graph.nodes
+            if n.op == "call_function" and n.target == torch.ops.aten.mul.Tensor
+        ]
+        self.assertEqual(len(mul_nodes), 1)
+        mul_args = set(mul_nodes[0].args)
+        self.assertTrue(
+            mul_args & set(sync_getitems),
+            "mul should depend on synchronize_event's getitem, not record_event's",
+        )
+
+    @requires_cuda
+    def test_external_event_synchronize_threads_inputs(self) -> None:
+        """When the event was recorded externally, synchronize threads graph inputs through."""
+
+        def fn(x):
+            e = torch.Event()
+            y = x + 1
+            e.record()
+            e.synchronize()
+            z = y * 2
+            return z
+
+        inp = (torch.ones(2, 2, device="cuda"),)
+        # Patch out wrapping so we get the raw graph to manually wrap below.
+        with patch(
+            "torch._functorch._aot_autograd.graph_capture.wrap_all_sync_nodes_with_control_deps"
+        ):
+            (
+                _,
+                _,
+                fw_graphs,
+                _,
+            ) = extract_graph(fn, *inp)
+
+        gm = fw_graphs[0]
+        graph = gm.graph
+
+        from torch._functorch._aot_autograd.streams import (
+            set_stream,
+            wrap_all_sync_nodes_with_control_deps,
+        )
+
+        # Remove the record_event to simulate an externally-recorded event.
+        record_node = next(
+            n
+            for n in graph.nodes
+            if n.op == "call_function"
+            and n.target is torch.ops.streams.record_event.default
+        )
+        stream_idx = record_node.args[1]
+        graph.erase_node(record_node)
+
+        # Set stream metadata on compute nodes.
+        for n in graph.nodes:
+            if (
+                n.op == "call_function"
+                and "val" in n.meta
+                and n.target is not torch.ops.streams.synchronize_event.default
+            ):
+                set_stream(n, stream_idx)
+
+        wrap_all_sync_nodes_with_control_deps(gm)
+        gm.recompile()
+
+        self.assertExpectedInline(
+            print_graph(gm),
+            """\
+class <lambda>(torch.nn.Module):
+    def forward(self, arg0_1: "f32[2, 2]"):
+        # Annotation: {'stream': 0}
+        add: "f32[2, 2]" = torch.ops.aten.add.Tensor(arg0_1, 1)
+
+        # No stacktrace found for following nodes
+        subgraph_synchronize_event = self.subgraph_synchronize_event
+        control_deps = torch.ops.higher_order.control_deps((arg0_1, add), subgraph_synchronize_event, add);  arg0_1 = add = subgraph_synchronize_event = None
+
+        # Annotation: {'stream': 0}
+        getitem: "f32[2, 2]" = control_deps[1];  control_deps = None
+
+        # Annotation: {'stream': 0}
+        mul: "f32[2, 2]" = torch.ops.aten.mul.Tensor(getitem, 2);  getitem = None
+        return (mul,)
+
+    class subgraph_synchronize_event(torch.nn.Module):
+        def forward(self, dep_0: "f32[2, 2]"):
+            #
+            synchronize_event_default = torch.ops.streams.synchronize_event.default(1)
+            return (synchronize_event_default, dep_0)
+""",  # noqa: B950
+        )
+
+    @requires_cuda
+    def test_event_synchronize_control_deps_e2e(self):
+        """E2E: compute → record → synchronize → use result through torch.compile."""
+
+        def f(x):
+            e = torch.Event()
+            y = x + 1
+            e.record()
+            e.synchronize()
+            z = y * 2
+            return z
+
+        inp = torch.ones(2, 2, device="cuda")
+        eager_result = f(inp)
+        compiled_result = torch.compile(f)(inp)
+        self.assertEqual(eager_result, compiled_result)
+
+    @requires_cuda
+    def test_event_synchronize_e2e(self):
+        def f(a_list):
+            a_cpu_list = []
+            a_to_cpu_event_list = []
+            for a in a_list:
+                a_cpu = a.to(device="cpu", non_blocking=True)
+                e = torch.Event()
+                e.record()
+                a_cpu_list.append(a_cpu)
+                a_to_cpu_event_list.append(e)
+
+            for e in a_to_cpu_event_list:
+                e.synchronize()
+
+            return torch.cat(a_cpu_list)
+
+        f_compiled = torch.compile(f)
+        inputs = [
+            torch.rand(100, dtype=torch.float16, device="cuda") for _ in range(10)
+        ]
+        eager_result = f(inputs)
+        compiled_result = f_compiled(inputs)
+        self.assertEqual(eager_result, compiled_result)
+
+    @requires_cuda
+    def test_event_record_wait_on_default_stream(self):
+        e = torch.cuda.Event()
+
+        def f(x):
+            y = x + 1
+            e.record()
+            e.wait()
+            return y + 1
+
+        f_compiled = torch.compile(f)
+        x = torch.randn(10, device="cuda")
+        eager_result = f(x)
+        compiled_result = f_compiled(x)
+        self.assertEqual(eager_result, compiled_result)
+
+    @requires_cuda
+    def test_record_stream_inductor_output_code(self) -> None:
+        """Verify record_stream is ordered between the producing kernel and the
+        consuming kernel in inductor-generated wrapper code."""
+        from torch._inductor.utils import run_and_get_code
+        from torch.testing import FileCheck
+
+        def fn(x):
+            s = torch.Stream(device="cuda")
+            y = x + 1
+            y.record_stream(s)
+            z = y * 2
+            return z
+
+        compiled = torch.compile(fn, backend="inductor", fullgraph=True)
+        x = torch.randn(1024, device="cuda")
+        result, (code,) = run_and_get_code(compiled, x)
+        self.assertEqual(result, (x + 1) * 2)
+
+        # record_stream must appear after the kernel that produces the tensor
+        # and before the return.
+        FileCheck().check(".run(").check(
+            "torch.ops.streams.record_stream.default("
+        ).check("return").run(code)
+
+    @requires_cuda
+    def test_del_multi_stream_sync_dealloc(self):
+        def fn(x, y):
+            s = torch.Stream()
+            e = torch.Event()
+            z0 = x + 1
+            with s:
+                z = torch.add(x, y)
+                e.record()
+            e.wait()
+            del x
+            return z0, z
+
+        inp = (torch.ones(2, 2, device="cuda"), torch.ones(2, 2, device="cuda"))
+        expected = fn(*inp)
+        (
+            actual,
+            _,
+            fw_graphs,
+            _,
+        ) = extract_graph(fn, *inp)
+        self.assertEqual(len(fw_graphs), 1)
+        self.assertEqual(expected, actual)
+        graph_str = print_graph(fw_graphs[0])
+        self.assertIn("sync_dealloc", graph_str)
+        self.assertIn("record_event", graph_str)
+
+    @requires_cuda
+    def test_del_same_stream_no_sync_dealloc(self):
+        def fn(x, y):
+            s = torch.Stream()
+            e = torch.Event()
+            with s:
+                z = torch.add(x, y)
+                del x
+                e.record()
+            e.wait()
+            return z
+
+        inp = (torch.ones(2, 2, device="cuda"), torch.ones(2, 2, device="cuda"))
+        expected = fn(*inp)
+        (
+            actual,
+            _,
+            fw_graphs,
+            _,
+        ) = extract_graph(fn, *inp)
+        self.assertEqual(len(fw_graphs), 1)
+        self.assertEqual(expected, actual)
+        graph_str = print_graph(fw_graphs[0])
+        self.assertNotIn("sync_dealloc", graph_str)
+
+    @requires_cuda
+    def test_del_single_stream_no_sync_dealloc(self):
+        def fn(x, y):
+            z = torch.add(x, y)
+            del x
+            return z
+
+        inp = (torch.ones(2, 2, device="cuda"), torch.ones(2, 2, device="cuda"))
+        expected = fn(*inp)
+        (
+            actual,
+            _,
+            fw_graphs,
+            _,
+        ) = extract_graph(fn, *inp)
+        self.assertEqual(len(fw_graphs), 1)
+        self.assertEqual(expected, actual)
+        graph_str = print_graph(fw_graphs[0])
+        self.assertNotIn("sync_dealloc", graph_str)
+
+    @requires_cuda
+    def test_del_attr_multi_stream_sync_dealloc(self):
+        class Holder:
+            pass
+
+        def fn(x, y):
+            s = torch.Stream()
+            e = torch.Event()
+            h = Holder()
+            h.tensor = x
+            z0 = x + 1
+            with s:
+                z = torch.add(h.tensor, y)
+                e.record()
+            e.wait()
+            del h.tensor
+            return z0, z
+
+        inp = (torch.ones(2, 2, device="cuda"), torch.ones(2, 2, device="cuda"))
+        expected = fn(*inp)
+        (
+            actual,
+            _,
+            fw_graphs,
+            _,
+        ) = extract_graph(fn, *inp)
+        self.assertEqual(len(fw_graphs), 1)
+        self.assertEqual(expected, actual)
+        graph_str = print_graph(fw_graphs[0])
+        self.assertIn("sync_dealloc", graph_str)
+        self.assertIn("record_event", graph_str)
+
+    @requires_cuda
+    def test_del_subscr_multi_stream_sync_dealloc(self):
+        def fn(x, y):
+            s = torch.Stream()
+            e = torch.Event()
+            d = {"t": x}
+            z0 = x + 1
+            with s:
+                z = torch.add(d["t"], y)
+                e.record()
+            e.wait()
+            del d["t"]
+            return z0, z
+
+        inp = (torch.ones(2, 2, device="cuda"), torch.ones(2, 2, device="cuda"))
+        expected = fn(*inp)
+        (
+            actual,
+            _,
+            fw_graphs,
+            _,
+        ) = extract_graph(fn, *inp)
+        self.assertEqual(len(fw_graphs), 1)
+        self.assertEqual(expected, actual)
+        graph_str = print_graph(fw_graphs[0])
+        self.assertIn("sync_dealloc", graph_str)
+        self.assertIn("record_event", graph_str)
+
+    @requires_cuda
+    def test_stream_pointer_extraction_edge_cases(self):
+        def get_ptrs(stream_a, stream_b, default_stream):
+            return (
+                stream_a.cuda_stream,
+                stream_b.cuda_stream,
+                default_stream.cuda_stream,
+            )
+
+        s1, s2 = torch.cuda.Stream(), torch.cuda.Stream()
+        default_s = torch.cuda.default_stream()
+        expected_s1, expected_s2 = s1.cuda_stream, s2.cuda_stream
+
+        self.assertNotEqual(expected_s1, expected_s2)
+        self.assertGreater(expected_s1, 1000)
+
+        opt_get_ptrs = torch.compile(get_ptrs, backend="inductor")
+
+        s3 = torch.cuda.Stream()
+        with torch.cuda.stream(s3):
+            actual_s1, actual_s2, actual_default = opt_get_ptrs(s1, s2, default_s)
+
+        self.assertEqual(actual_s1, expected_s1)
+        self.assertEqual(actual_s2, expected_s2)
+        self.assertEqual(actual_default, default_s.cuda_stream)
 
 
 if __name__ == "__main__":
diff --git a/test/dynamo/test_structured_trace.py b/test/dynamo/test_structured_trace.py
index a3954f01e2045..d2e177bba1120 100644
--- a/test/dynamo/test_structured_trace.py
+++ b/test/dynamo/test_structured_trace.py
@@ -296,21 +296,24 @@ def test_compile_id_serialization_deserialization(self):
             frame_id=1,
             frame_compile_id=2,
         )
-        assert cid == torch._guards.CompileId.from_string(str(cid))
+        if cid != torch._guards.CompileId.from_string(str(cid)):
+            raise AssertionError("CompileId round-trip failed")
 
         cid = torch._guards.CompileId(
             compiled_autograd_id=1,
             frame_id=2,
             frame_compile_id=3,
         )
-        assert cid == torch._guards.CompileId.from_string(str(cid))
+        if cid != torch._guards.CompileId.from_string(str(cid)):
+            raise AssertionError("CompileId round-trip failed")
 
         cid = torch._guards.CompileId(
             compiled_autograd_id=1,
             frame_id=None,
             frame_compile_id=None,
         )
-        assert cid == torch._guards.CompileId.from_string(str(cid))
+        if cid != torch._guards.CompileId.from_string(str(cid)):
+            raise AssertionError("CompileId round-trip failed")
 
         for bad_cid in ["-/-", "-/1", "1/-", "!1/2", "!1/-/-"]:
             with self.assertRaises(ValueError):
@@ -328,9 +331,9 @@ def test_schedule(self):
 {"describe_tensor": {"id": 0, "ndim": 2, "dtype": "torch.float32", "device": "device(type='cuda', index=0)", "size": [1000, 1000], "dynamo_hint_overrides": {}, "is_leaf": true, "stride": [1000, 1], "storage": 0, "view_func": "VIEW_FUNC", "describer_id": "ID"}, "frame_id": 0, "frame_compile_id": 0, "attempt": 0}
 {"describe_source": {"describer_id": "ID", "id": 0, "source": "L['a']"}, "frame_id": 0, "frame_compile_id": 0, "attempt": 0}
 {"dynamo_output_graph": {"sizes": {"l_a_": [1000, 1000], "ones": [1000, 1000], "output": [1000, 1000]}}, "frame_id": 0, "frame_compile_id": 0, "attempt": 0, "has_payload": "HASH"}
+{"artifact": {"name": "aotautograd_cache_miss", "encoding": "json"}, "frame_id": 0, "frame_compile_id": 0, "attempt": 0, "has_payload": "HASH"}
 {"artifact": {"name": "before_pre_grad_graph", "encoding": "string"}, "frame_id": 0, "frame_compile_id": 0, "attempt": 0, "has_payload": "HASH"}
 {"artifact": {"name": "after_pre_grad_graph", "encoding": "string"}, "frame_id": 0, "frame_compile_id": 0, "attempt": 0, "has_payload": "HASH"}
-{"artifact": {"name": "aotautograd_cache_miss", "encoding": "json"}, "frame_id": 0, "frame_compile_id": 0, "attempt": 0, "has_payload": "HASH"}
 {"artifact": {"name": "aot_forward_graph_fw_metadata", "encoding": "string"}, "frame_id": 0, "frame_compile_id": 0, "attempt": 0, "has_payload": "HASH"}
 {"aot_inference_graph": {}, "frame_id": 0, "frame_compile_id": 0, "attempt": 0, "has_payload": "HASH"}
 {"artifact": {"name": "torch._functorch.config", "encoding": "string"}, "frame_id": 0, "frame_compile_id": 0, "attempt": 0, "has_payload": "HASH"}
@@ -362,9 +365,9 @@ def test_cudagraphs(self):
 {"describe_tensor": {"id": 0, "ndim": 2, "dtype": "torch.float32", "device": "device(type='cuda', index=0)", "size": [1000, 1000], "dynamo_hint_overrides": {}, "is_leaf": true, "stride": [1000, 1], "storage": 0, "view_func": "VIEW_FUNC", "describer_id": "ID"}, "frame_id": 0, "frame_compile_id": 0, "attempt": 0}
 {"describe_source": {"describer_id": "ID", "id": 0, "source": "L['a']"}, "frame_id": 0, "frame_compile_id": 0, "attempt": 0}
 {"dynamo_output_graph": {"sizes": {"l_a_": [1000, 1000], "ones": [1000, 1000], "output": [1000, 1000]}}, "frame_id": 0, "frame_compile_id": 0, "attempt": 0, "has_payload": "HASH"}
+{"artifact": {"name": "aotautograd_cache_miss", "encoding": "json"}, "frame_id": 0, "frame_compile_id": 0, "attempt": 0, "has_payload": "HASH"}
 {"artifact": {"name": "before_pre_grad_graph", "encoding": "string"}, "frame_id": 0, "frame_compile_id": 0, "attempt": 0, "has_payload": "HASH"}
 {"artifact": {"name": "after_pre_grad_graph", "encoding": "string"}, "frame_id": 0, "frame_compile_id": 0, "attempt": 0, "has_payload": "HASH"}
-{"artifact": {"name": "aotautograd_cache_miss", "encoding": "json"}, "frame_id": 0, "frame_compile_id": 0, "attempt": 0, "has_payload": "HASH"}
 {"artifact": {"name": "aot_forward_graph_fw_metadata", "encoding": "string"}, "frame_id": 0, "frame_compile_id": 0, "attempt": 0, "has_payload": "HASH"}
 {"aot_inference_graph": {}, "frame_id": 0, "frame_compile_id": 0, "attempt": 0, "has_payload": "HASH"}
 {"artifact": {"name": "torch._functorch.config", "encoding": "string"}, "frame_id": 0, "frame_compile_id": 0, "attempt": 0, "has_payload": "HASH"}
@@ -404,9 +407,9 @@ def fn(x, y):
 {"describe_tensor": {"id": 1, "ndim": 2, "dtype": "torch.float32", "device": "device(type='cpu')", "size": [1000, 1000], "dynamo_hint_overrides": {}, "is_leaf": true, "stride": [1000, 1], "storage": 1, "view_func": "VIEW_FUNC", "describer_id": "ID"}, "frame_id": 0, "frame_compile_id": 0, "attempt": 0}
 {"describe_source": {"describer_id": "ID", "id": 1, "source": "L['y']"}, "frame_id": 0, "frame_compile_id": 0, "attempt": 0}
 {"dynamo_output_graph": {"sizes": {"l_x_": [1000, 1000], "l_y_": [1000, 1000], "add": [1000, 1000]}}, "frame_id": 0, "frame_compile_id": 0, "attempt": 0, "has_payload": "HASH"}
+{"artifact": {"name": "aotautograd_cache_miss", "encoding": "json"}, "frame_id": 0, "frame_compile_id": 0, "attempt": 0, "has_payload": "HASH"}
 {"artifact": {"name": "before_pre_grad_graph", "encoding": "string"}, "frame_id": 0, "frame_compile_id": 0, "attempt": 0, "has_payload": "HASH"}
 {"artifact": {"name": "after_pre_grad_graph", "encoding": "string"}, "frame_id": 0, "frame_compile_id": 0, "attempt": 0, "has_payload": "HASH"}
-{"artifact": {"name": "aotautograd_cache_miss", "encoding": "json"}, "frame_id": 0, "frame_compile_id": 0, "attempt": 0, "has_payload": "HASH"}
 {"artifact": {"name": "aot_forward_graph_fw_metadata", "encoding": "string"}, "frame_id": 0, "frame_compile_id": 0, "attempt": 0, "has_payload": "HASH"}
 {"aot_inference_graph": {}, "frame_id": 0, "frame_compile_id": 0, "attempt": 0, "has_payload": "HASH"}
 {"artifact": {"name": "torch._functorch.config", "encoding": "string"}, "frame_id": 0, "frame_compile_id": 0, "attempt": 0, "has_payload": "HASH"}
@@ -426,9 +429,9 @@ def fn(x, y):
 {"describe_source": {"describer_id": "ID", "id": 0, "source": "L['x']"}, "frame_id": 0, "frame_compile_id": 1, "attempt": 0}
 {"create_symbol": {"symbol": "s48", "val": "1", "vr": "[-int_oo, int_oo]", "source": "L['y']", "user_stack": "STACK", "stack": "STACK"}, "frame_id": 0, "frame_compile_id": 1, "attempt": 0}
 {"dynamo_output_graph": {"sizes": {"l_x_": [1000, 1000], "add": [1000, 1000]}}, "frame_id": 0, "frame_compile_id": 1, "attempt": 0, "has_payload": "HASH"}
+{"artifact": {"name": "aotautograd_cache_miss", "encoding": "json"}, "frame_id": 0, "frame_compile_id": 1, "attempt": 0, "has_payload": "HASH"}
 {"artifact": {"name": "before_pre_grad_graph", "encoding": "string"}, "frame_id": 0, "frame_compile_id": 1, "attempt": 0, "has_payload": "HASH"}
 {"artifact": {"name": "after_pre_grad_graph", "encoding": "string"}, "frame_id": 0, "frame_compile_id": 1, "attempt": 0, "has_payload": "HASH"}
-{"artifact": {"name": "aotautograd_cache_miss", "encoding": "json"}, "frame_id": 0, "frame_compile_id": 1, "attempt": 0, "has_payload": "HASH"}
 {"artifact": {"name": "aot_forward_graph_fw_metadata", "encoding": "string"}, "frame_id": 0, "frame_compile_id": 1, "attempt": 0, "has_payload": "HASH"}
 {"aot_inference_graph": {}, "frame_id": 0, "frame_compile_id": 1, "attempt": 0, "has_payload": "HASH"}
 {"artifact": {"name": "torch._functorch.config", "encoding": "string"}, "frame_id": 0, "frame_compile_id": 1, "attempt": 0, "has_payload": "HASH"}
@@ -458,9 +461,9 @@ def test_example_fn(self):
 {"describe_tensor": {"id": 0, "ndim": 2, "dtype": "torch.float32", "device": "device(type='cpu')", "size": [1000, 1000], "dynamo_hint_overrides": {}, "is_leaf": true, "stride": [1000, 1], "storage": 0, "view_func": "VIEW_FUNC", "describer_id": "ID"}, "frame_id": 0, "frame_compile_id": 0, "attempt": 0}
 {"describe_source": {"describer_id": "ID", "id": 0, "source": "L['a']"}, "frame_id": 0, "frame_compile_id": 0, "attempt": 0}
 {"dynamo_output_graph": {"sizes": {"l_a_": [1000, 1000], "ones": [1000, 1000], "output": [1000, 1000], "ones_1": [1000, 1000], "output_1": [1000, 1000]}}, "frame_id": 0, "frame_compile_id": 0, "attempt": 0, "has_payload": "HASH"}
+{"artifact": {"name": "aotautograd_cache_miss", "encoding": "json"}, "frame_id": 0, "frame_compile_id": 0, "attempt": 0, "has_payload": "HASH"}
 {"artifact": {"name": "before_pre_grad_graph", "encoding": "string"}, "frame_id": 0, "frame_compile_id": 0, "attempt": 0, "has_payload": "HASH"}
 {"artifact": {"name": "after_pre_grad_graph", "encoding": "string"}, "frame_id": 0, "frame_compile_id": 0, "attempt": 0, "has_payload": "HASH"}
-{"artifact": {"name": "aotautograd_cache_miss", "encoding": "json"}, "frame_id": 0, "frame_compile_id": 0, "attempt": 0, "has_payload": "HASH"}
 {"artifact": {"name": "aot_forward_graph_fw_metadata", "encoding": "string"}, "frame_id": 0, "frame_compile_id": 0, "attempt": 0, "has_payload": "HASH"}
 {"aot_inference_graph": {}, "frame_id": 0, "frame_compile_id": 0, "attempt": 0, "has_payload": "HASH"}
 {"artifact": {"name": "torch._functorch.config", "encoding": "string"}, "frame_id": 0, "frame_compile_id": 0, "attempt": 0, "has_payload": "HASH"}
@@ -500,13 +503,16 @@ def test_example_training_fn(self):
 {"dynamo_cpp_guards_str": {}, "frame_id": 0, "frame_compile_id": 0, "attempt": 1, "has_payload": "HASH"}
 {"compilation_metrics": "METRICS", "frame_id": 0, "frame_compile_id": 0, "attempt": 1}
 {"dynamo_start": {"stack": "STACK"}, "frame_id": 1, "frame_compile_id": 0, "attempt": 0}
+{"artifact": {"name": "torch_dynamo_resume_in_example_training_fn_at_45_ORIGINAL_BYTECODE", "encoding": "string"}, "frame_id": 1, "frame_compile_id": 0, "attempt": 0, "has_payload": "HASH"}
 {"artifact": {"name": "dynamo_graph_break_reason", "encoding": "string"}, "frame_id": 1, "frame_compile_id": 0, "attempt": 0, "has_payload": "HASH"}
 {"describe_storage": {"id": 0, "describer_id": "ID", "size": 4000000}, "frame_id": 1, "frame_compile_id": 0, "attempt": 1}
 {"describe_tensor": {"id": 0, "ndim": 2, "dtype": "torch.float32", "device": "device(type='cpu')", "size": [1000, 1000], "dynamo_hint_overrides": {}, "is_leaf": true, "requires_grad": true, "stride": [1000, 1], "storage": 0, "view_func": "VIEW_FUNC", "describer_id": "ID"}, "frame_id": 1, "frame_compile_id": 0, "attempt": 1}
 {"describe_source": {"describer_id": "ID", "id": 0, "source": "L['___stack1']"}, "frame_id": 1, "frame_compile_id": 0, "attempt": 1}
+{"artifact": {"name": "torch_dynamo_resume_in_example_training_fn_at_45_MODIFIED_BYTECODE", "encoding": "string"}, "frame_id": 1, "frame_compile_id": 0, "attempt": 1, "has_payload": "HASH"}
 {"dynamo_cpp_guards_str": {}, "frame_id": 1, "frame_compile_id": 0, "attempt": 1, "has_payload": "HASH"}
 {"compilation_metrics": "METRICS", "frame_id": 1, "frame_compile_id": 0, "attempt": 1}
 {"dynamo_start": {"stack": "STACK"}, "frame_id": 2, "frame_compile_id": 0, "attempt": 0}
+{"artifact": {"name": "torch_dynamo_resume_in_example_training_fn_at_45_ORIGINAL_BYTECODE", "encoding": "string"}, "frame_id": 2, "frame_compile_id": 0, "attempt": 0, "has_payload": "HASH"}
 {"describe_storage": {"id": 0, "describer_id": "ID", "size": 4000000}, "frame_id": 2, "frame_compile_id": 0, "attempt": 0}
 {"describe_tensor": {"id": 0, "ndim": 2, "dtype": "torch.float32", "device": "device(type='cpu')", "size": [1000, 1000], "dynamo_hint_overrides": {}, "requires_grad": true, "stride": [1000, 1], "storage": 0, "view_func": "VIEW_FUNC", "describer_id": "ID"}, "frame_id": 2, "frame_compile_id": 0, "attempt": 0}
 {"describe_source": {"describer_id": "ID", "id": 0, "source": "L['___stack0']"}, "frame_id": 2, "frame_compile_id": 0, "attempt": 0}
@@ -515,9 +521,9 @@ def test_example_training_fn(self):
 {"describe_tensor": {"id": 0, "ndim": 2, "dtype": "torch.float32", "device": "device(type='cpu')", "size": [1000, 1000], "dynamo_hint_overrides": {}, "requires_grad": true, "stride": [1000, 1], "storage": 0, "view_func": "VIEW_FUNC", "describer_id": "ID"}, "frame_id": 2, "frame_compile_id": 0, "attempt": 1}
 {"describe_source": {"describer_id": "ID", "id": 0, "source": "L['___stack0']"}, "frame_id": 2, "frame_compile_id": 0, "attempt": 1}
 {"dynamo_output_graph": {"sizes": {"l_stack0_": [1000, 1000], "ones": [1000, 1000], "output": [1000, 1000], "sum_1": []}}, "frame_id": 2, "frame_compile_id": 0, "attempt": 1, "has_payload": "HASH"}
+{"artifact": {"name": "aotautograd_cache_miss", "encoding": "json"}, "frame_id": 2, "frame_compile_id": 0, "attempt": 1, "has_payload": "HASH"}
 {"artifact": {"name": "before_pre_grad_graph", "encoding": "string"}, "frame_id": 2, "frame_compile_id": 0, "attempt": 1, "has_payload": "HASH"}
 {"artifact": {"name": "after_pre_grad_graph", "encoding": "string"}, "frame_id": 2, "frame_compile_id": 0, "attempt": 1, "has_payload": "HASH"}
-{"artifact": {"name": "aotautograd_cache_miss", "encoding": "json"}, "frame_id": 2, "frame_compile_id": 0, "attempt": 1, "has_payload": "HASH"}
 {"aot_joint_graph": {}, "frame_id": 2, "frame_compile_id": 0, "attempt": 1, "has_payload": "HASH"}
 {"artifact": {"name": "torch._functorch.config", "encoding": "string"}, "frame_id": 2, "frame_compile_id": 0, "attempt": 1, "has_payload": "HASH"}
 {"artifact": {"name": "aot_forward_graph_fw_metadata", "encoding": "string"}, "frame_id": 2, "frame_compile_id": 0, "attempt": 1, "has_payload": "HASH"}
@@ -528,21 +534,21 @@ def test_example_training_fn(self):
 {"artifact": {"name": "inductor_post_grad_graph", "encoding": "string"}, "frame_id": 2, "frame_compile_id": 0, "attempt": 1, "has_payload": "HASH"}
 {"inductor_output_code": {"filename": "FILENAME", "file_path": "FILENAME"}, "frame_id": 2, "frame_compile_id": 0, "attempt": 1, "has_payload": "HASH"}
 {"artifact": {"name": "fx_graph_cache_miss", "encoding": "json"}, "frame_id": 2, "frame_compile_id": 0, "attempt": 1, "has_payload": "HASH"}
+{"artifact": {"name": "torch_dynamo_resume_in_example_training_fn_at_45_MODIFIED_BYTECODE", "encoding": "string"}, "frame_id": 2, "frame_compile_id": 0, "attempt": 1, "has_payload": "HASH"}
 {"dynamo_cpp_guards_str": {}, "frame_id": 2, "frame_compile_id": 0, "attempt": 1, "has_payload": "HASH"}
 {"compilation_metrics": "METRICS", "frame_id": 2, "frame_compile_id": 0, "attempt": 1}
-{"dynamo_start": {"stack": "STACK"}, "frame_id": 3, "frame_compile_id": 0, "attempt": 0}
-{"compilation_metrics": "METRICS", "frame_id": 3, "frame_compile_id": 0, "attempt": 0}
 {"artifact": {"name": "fx_graph_runnable", "encoding": "string"}, "frame_id": 2, "frame_compile_id": 0, "attempt": 1, "has_payload": "HASH"}
 {"artifact": {"name": "before_post_grad_graph", "encoding": "string"}, "frame_id": 2, "frame_compile_id": 0, "attempt": 1, "has_payload": "HASH"}
 {"artifact": {"name": "inductor_post_grad_graph", "encoding": "string"}, "frame_id": 2, "frame_compile_id": 0, "attempt": 1, "has_payload": "HASH"}
 {"inductor_output_code": {"filename": "FILENAME", "file_path": "FILENAME"}, "frame_id": 2, "frame_compile_id": 0, "attempt": 1, "has_payload": "HASH"}
 {"artifact": {"name": "fx_graph_cache_miss", "encoding": "json"}, "frame_id": 2, "frame_compile_id": 0, "attempt": 1, "has_payload": "HASH"}
 {"bwd_compilation_metrics": "METRICS", "frame_id": 2, "frame_compile_id": 0, "attempt": 1}
-{"dynamo_start": {"stack": "STACK"}, "frame_id": 4, "frame_compile_id": 0, "attempt": 0}
-{"describe_storage": {"id": 0, "describer_id": "ID", "size": 4000000}, "frame_id": 4, "frame_compile_id": 0, "attempt": 0}
-{"describe_tensor": {"id": 0, "ndim": 2, "dtype": "torch.float32", "device": "device(type='cpu')", "size": [1000, 1000], "dynamo_hint_overrides": {}, "requires_grad": true, "stride": [1000, 1], "storage": 0, "view_func": "VIEW_FUNC", "describer_id": "ID"}, "frame_id": 4, "frame_compile_id": 0, "attempt": 0}
-{"describe_source": {"describer_id": "ID", "id": 0, "source": "L['output']"}, "frame_id": 4, "frame_compile_id": 0, "attempt": 0}
-{"compilation_metrics": "METRICS", "frame_id": 4, "frame_compile_id": 0, "attempt": 0}
+{"dynamo_start": {"stack": "STACK"}, "frame_id": 3, "frame_compile_id": 0, "attempt": 0}
+{"artifact": {"name": "torch_dynamo_resume_in_example_training_fn_at_47_ORIGINAL_BYTECODE", "encoding": "string"}, "frame_id": 3, "frame_compile_id": 0, "attempt": 0, "has_payload": "HASH"}
+{"describe_storage": {"id": 0, "describer_id": "ID", "size": 4000000}, "frame_id": 3, "frame_compile_id": 0, "attempt": 0}
+{"describe_tensor": {"id": 0, "ndim": 2, "dtype": "torch.float32", "device": "device(type='cpu')", "size": [1000, 1000], "dynamo_hint_overrides": {}, "requires_grad": true, "stride": [1000, 1], "storage": 0, "view_func": "VIEW_FUNC", "describer_id": "ID"}, "frame_id": 3, "frame_compile_id": 0, "attempt": 0}
+{"describe_source": {"describer_id": "ID", "id": 0, "source": "L['output']"}, "frame_id": 3, "frame_compile_id": 0, "attempt": 0}
+{"compilation_metrics": "METRICS", "frame_id": 3, "frame_compile_id": 0, "attempt": 0}
 """,  # noqa: B950
         )
 
@@ -597,9 +603,9 @@ def throw(x):
 {"describe_tensor": {"id": 0, "ndim": 2, "dtype": "torch.float32", "device": "device(type='cpu')", "size": [1000, 1000], "dynamo_hint_overrides": {}, "is_leaf": true, "requires_grad": true, "stride": [1000, 1], "storage": 0, "view_func": "VIEW_FUNC", "describer_id": "ID"}, "frame_id": 0, "frame_compile_id": 0, "attempt": 0}
 {"describe_source": {"describer_id": "ID", "id": 0, "source": "L['a']"}, "frame_id": 0, "frame_compile_id": 0, "attempt": 0}
 {"dynamo_output_graph": {"sizes": {"l_a_": [1000, 1000], "output": [1000, 1000]}}, "frame_id": 0, "frame_compile_id": 0, "attempt": 0, "has_payload": "HASH"}
+{"artifact": {"name": "aotautograd_cache_miss", "encoding": "json"}, "frame_id": 0, "frame_compile_id": 0, "attempt": 0, "has_payload": "HASH"}
 {"artifact": {"name": "before_pre_grad_graph", "encoding": "string"}, "frame_id": 0, "frame_compile_id": 0, "attempt": 0, "has_payload": "HASH"}
 {"artifact": {"name": "after_pre_grad_graph", "encoding": "string"}, "frame_id": 0, "frame_compile_id": 0, "attempt": 0, "has_payload": "HASH"}
-{"artifact": {"name": "aotautograd_cache_miss", "encoding": "json"}, "frame_id": 0, "frame_compile_id": 0, "attempt": 0, "has_payload": "HASH"}
 {"aot_joint_graph": {}, "frame_id": 0, "frame_compile_id": 0, "attempt": 0, "has_payload": "HASH"}
 {"artifact": {"name": "torch._functorch.config", "encoding": "string"}, "frame_id": 0, "frame_compile_id": 0, "attempt": 0, "has_payload": "HASH"}
 {"artifact": {"name": "aot_forward_graph_fw_metadata", "encoding": "string"}, "frame_id": 0, "frame_compile_id": 0, "attempt": 0, "has_payload": "HASH"}
@@ -617,6 +623,7 @@ def throw(x):
 
     @requires_distributed()
     @requires_cuda_and_triton
+    @unittest.skip("https://github.com/pytorch/pytorch/issues/176188")
     def test_ddp_graphs(self):
         import torch._dynamo.convert_frame as convert_frame
 
@@ -646,67 +653,9 @@ def forward(self, x):
 
         dist.destroy_process_group()
 
-        if not torch._dynamo.config.inline_inbuilt_nn_modules:
-            self.assertExpectedInline(
-                self.buffer.getvalue(),
-                """\
-{"dynamo_start": {"stack": "STACK"}, "rank": 0, "frame_id": 0, "frame_compile_id": 0, "attempt": 0}
-{"artifact": {"name": "dynamo_graph_break_reason", "encoding": "string"}, "rank": 0, "frame_id": 0, "frame_compile_id": 0, "attempt": 0, "has_payload": "HASH"}
-{"dynamo_cpp_guards_str": {}, "rank": 0, "frame_id": 0, "frame_compile_id": 0, "attempt": 1, "has_payload": "HASH"}
-{"compilation_metrics": "METRICS", "rank": 0, "frame_id": 0, "frame_compile_id": 0, "attempt": 1}
-{"dynamo_start": {"stack": "STACK"}, "rank": 0, "frame_id": 1, "frame_compile_id": 0, "attempt": 0}
-{"artifact": {"name": "dynamo_graph_break_reason", "encoding": "string"}, "rank": 0, "frame_id": 1, "frame_compile_id": 0, "attempt": 0, "has_payload": "HASH"}
-{"dynamo_cpp_guards_str": {}, "rank": 0, "frame_id": 1, "frame_compile_id": 0, "attempt": 1, "has_payload": "HASH"}
-{"compilation_metrics": "METRICS", "rank": 0, "frame_id": 1, "frame_compile_id": 0, "attempt": 1}
-{"dynamo_start": {"stack": "STACK"}, "rank": 0, "frame_id": 2, "frame_compile_id": 0, "attempt": 0}
-{"artifact": {"name": "dynamo_graph_break_reason", "encoding": "string"}, "rank": 0, "frame_id": 2, "frame_compile_id": 0, "attempt": 0, "has_payload": "HASH"}
-{"dynamo_cpp_guards_str": {}, "rank": 0, "frame_id": 2, "frame_compile_id": 0, "attempt": 1, "has_payload": "HASH"}
-{"compilation_metrics": "METRICS", "rank": 0, "frame_id": 2, "frame_compile_id": 0, "attempt": 1}
-{"dynamo_start": {"stack": "STACK"}, "rank": 0, "frame_id": 3, "frame_compile_id": 0, "attempt": 0}
-{"describe_storage": {"id": 0, "describer_id": "ID", "size": 4194304}, "rank": 0, "frame_id": 3, "frame_compile_id": 0, "attempt": 0}
-{"describe_tensor": {"id": 0, "ndim": 2, "dtype": "torch.float32", "device": "device(type='cuda', index=0)", "size": [1024, 1024], "is_leaf": true, "stride": [1024, 1], "storage": 0, "view_func": "VIEW_FUNC", "describer_id": "ID"}, "rank": 0, "frame_id": 3, "frame_compile_id": 0, "attempt": 0}
-{"describe_source": {"describer_id": "ID", "id": 0, "source": "L['x']"}, "rank": 0, "frame_id": 3, "frame_compile_id": 0, "attempt": 0}
-{"dynamo_output_graph": {"sizes": {"l_x_": [1024, 1024], "l__self___layers_0": [1024, 1024], "l__self___layers_1": [1024, 1024]}}, "rank": 0, "frame_id": 3, "frame_compile_id": 0, "attempt": 0, "has_payload": "HASH"}
-{"optimize_ddp_split_graph": {}, "rank": 0, "frame_id": 3, "frame_compile_id": 0, "attempt": 0, "has_payload": "HASH"}
-{"optimize_ddp_split_child": {"name": "submod_0"}, "rank": 0, "frame_id": 3, "frame_compile_id": 0, "attempt": 0, "has_payload": "HASH"}
-{"optimize_ddp_split_child": {"name": "submod_1"}, "rank": 0, "frame_id": 3, "frame_compile_id": 0, "attempt": 0, "has_payload": "HASH"}
-{"describe_storage": {"id": 0, "describer_id": "ID", "size": 4194304}, "rank": 0, "frame_id": 3, "frame_compile_id": 0, "attempt": 0}
-{"describe_tensor": {"id": 0, "ndim": 2, "dtype": "torch.float32", "device": "device(type='cuda', index=0)", "size": [1024, 1024], "is_leaf": true, "stride": [1024, 1], "storage": 0, "view_func": "VIEW_FUNC", "describer_id": "ID"}, "rank": 0, "frame_id": 3, "frame_compile_id": 0, "attempt": 0}
-{"describe_source": {"describer_id": "ID", "id": 0, "source": "L['x']"}, "rank": 0, "frame_id": 3, "frame_compile_id": 0, "attempt": 0}
-{"artifact": {"name": "before_pre_grad_graph", "encoding": "string"}, "rank": 0, "frame_id": 3, "frame_compile_id": 0, "attempt": 0, "has_payload": "HASH"}
-{"artifact": {"name": "after_pre_grad_graph", "encoding": "string"}, "rank": 0, "frame_id": 3, "frame_compile_id": 0, "attempt": 0, "has_payload": "HASH"}
-{"aot_joint_graph": {}, "rank": 0, "frame_id": 3, "frame_compile_id": 0, "attempt": 0, "has_payload": "HASH"}
-{"artifact": {"name": "torch._functorch.config", "encoding": "string"}, "rank": 0, "frame_id": 3, "frame_compile_id": 0, "attempt": 0, "has_payload": "HASH"}
-{"artifact": {"name": "aot_forward_graph_fw_metadata", "encoding": "string"}, "rank": 0, "frame_id": 3, "frame_compile_id": 0, "attempt": 0, "has_payload": "HASH"}
-{"aot_forward_graph": {}, "rank": 0, "frame_id": 3, "frame_compile_id": 0, "attempt": 0, "has_payload": "HASH"}
-{"aot_backward_graph": {}, "rank": 0, "frame_id": 3, "frame_compile_id": 0, "attempt": 0, "has_payload": "HASH"}
-{"artifact": {"name": "fx_graph_runnable", "encoding": "string"}, "rank": 0, "frame_id": 3, "frame_compile_id": 0, "attempt": 0, "has_payload": "HASH"}
-{"artifact": {"name": "before_post_grad_graph", "encoding": "string"}, "rank": 0, "frame_id": 3, "frame_compile_id": 0, "attempt": 0, "has_payload": "HASH"}
-{"artifact": {"name": "inductor_post_grad_graph", "encoding": "string"}, "rank": 0, "frame_id": 3, "frame_compile_id": 0, "attempt": 0, "has_payload": "HASH"}
-{"inductor_output_code": {"filename": "FILENAME", "file_path": "FILENAME"}, "rank": 0, "frame_id": 3, "frame_compile_id": 0, "attempt": 0, "has_payload": "HASH"}
-{"artifact": {"name": "fx_graph_cache_miss", "encoding": "json"}, "rank": 0, "frame_id": 3, "frame_compile_id": 0, "attempt": 0, "has_payload": "HASH"}
-{"artifact": {"name": "aotautograd_cache_bypass", "encoding": "json"}, "rank": 0, "frame_id": 3, "frame_compile_id": 0, "attempt": 0, "has_payload": "HASH"}
-{"artifact": {"name": "before_pre_grad_graph", "encoding": "string"}, "rank": 0, "frame_id": 3, "frame_compile_id": 0, "attempt": 0, "has_payload": "HASH"}
-{"artifact": {"name": "after_pre_grad_graph", "encoding": "string"}, "rank": 0, "frame_id": 3, "frame_compile_id": 0, "attempt": 0, "has_payload": "HASH"}
-{"aot_joint_graph": {}, "rank": 0, "frame_id": 3, "frame_compile_id": 0, "attempt": 0, "has_payload": "HASH"}
-{"artifact": {"name": "torch._functorch.config", "encoding": "string"}, "rank": 0, "frame_id": 3, "frame_compile_id": 0, "attempt": 0, "has_payload": "HASH"}
-{"artifact": {"name": "aot_forward_graph_fw_metadata", "encoding": "string"}, "rank": 0, "frame_id": 3, "frame_compile_id": 0, "attempt": 0, "has_payload": "HASH"}
-{"aot_forward_graph": {}, "rank": 0, "frame_id": 3, "frame_compile_id": 0, "attempt": 0, "has_payload": "HASH"}
-{"aot_backward_graph": {}, "rank": 0, "frame_id": 3, "frame_compile_id": 0, "attempt": 0, "has_payload": "HASH"}
-{"artifact": {"name": "fx_graph_runnable", "encoding": "string"}, "rank": 0, "frame_id": 3, "frame_compile_id": 0, "attempt": 0, "has_payload": "HASH"}
-{"artifact": {"name": "before_post_grad_graph", "encoding": "string"}, "rank": 0, "frame_id": 3, "frame_compile_id": 0, "attempt": 0, "has_payload": "HASH"}
-{"artifact": {"name": "inductor_post_grad_graph", "encoding": "string"}, "rank": 0, "frame_id": 3, "frame_compile_id": 0, "attempt": 0, "has_payload": "HASH"}
-{"inductor_output_code": {"filename": "FILENAME", "file_path": "FILENAME"}, "rank": 0, "frame_id": 3, "frame_compile_id": 0, "attempt": 0, "has_payload": "HASH"}
-{"artifact": {"name": "fx_graph_cache_miss", "encoding": "json"}, "rank": 0, "frame_id": 3, "frame_compile_id": 0, "attempt": 0, "has_payload": "HASH"}
-{"artifact": {"name": "aotautograd_cache_bypass", "encoding": "json"}, "rank": 0, "frame_id": 3, "frame_compile_id": 0, "attempt": 0, "has_payload": "HASH"}
-{"dynamo_cpp_guards_str": {}, "rank": 0, "frame_id": 3, "frame_compile_id": 0, "attempt": 0, "has_payload": "HASH"}
-{"compilation_metrics": "METRICS", "rank": 0, "frame_id": 3, "frame_compile_id": 0, "attempt": 0}
-""",  # noqa: B950
-            )
-        else:
-            self.assertExpectedInline(
-                self.buffer.getvalue(),
-                """\
+        self.assertExpectedInline(
+            self.buffer.getvalue(),
+            """\
 {"dynamo_start": {"stack": "STACK"}, "rank": 0, "frame_id": 0, "frame_compile_id": 0, "attempt": 0}
 {"artifact": {"name": "dynamo_graph_break_reason", "encoding": "string"}, "rank": 0, "frame_id": 0, "frame_compile_id": 0, "attempt": 0, "has_payload": "HASH"}
 {"describe_storage": {"id": 0, "describer_id": "ID", "size": 4194304}, "rank": 0, "frame_id": 0, "frame_compile_id": 0, "attempt": 1}
@@ -723,6 +672,7 @@ def forward(self, x):
 {"dynamo_cpp_guards_str": {}, "rank": 0, "frame_id": 2, "frame_compile_id": 0, "attempt": 1, "has_payload": "HASH"}
 {"compilation_metrics": "METRICS", "rank": 0, "frame_id": 2, "frame_compile_id": 0, "attempt": 1}
 {"dynamo_start": {"stack": "STACK"}, "rank": 0, "frame_id": 3, "frame_compile_id": 0, "attempt": 0}
+{"artifact": {"name": "torch_dynamo_resume_in___init___at_103_ORIGINAL_BYTECODE", "encoding": "string"}, "frame_id": 3, "frame_compile_id": 0, "attempt": 0, "has_payload": "HASH"}
 {"compilation_metrics": "METRICS", "rank": 0, "frame_id": 3, "frame_compile_id": 0, "attempt": 0}
 {"dynamo_start": {"stack": "STACK"}, "rank": 0, "frame_id": 4, "frame_compile_id": 0, "attempt": 0}
 {"describe_storage": {"id": 0, "describer_id": "ID", "size": 4194304}, "rank": 0, "frame_id": 4, "frame_compile_id": 0, "attempt": 0}
@@ -788,11 +738,12 @@ def forward(self, x):
 {"dynamo_cpp_guards_str": {}, "rank": 0, "frame_id": 4, "frame_compile_id": 0, "attempt": 0, "has_payload": "HASH"}
 {"compilation_metrics": "METRICS", "rank": 0, "frame_id": 4, "frame_compile_id": 0, "attempt": 0}
 """,  # noqa: B950
-            )
+        )
 
         self.assertParses()
 
     @requires_tlparse
+    @unittest.skip("https://github.com/pytorch/pytorch/issues/176188")
     def test_graph_breaks(self):
         @torch.compile(backend="inductor")
         def fn(x):
@@ -809,13 +760,14 @@ def fn(x):
 {"dynamo_cpp_guards_str": {}, "frame_id": 0, "frame_compile_id": 0, "attempt": 1, "has_payload": "HASH"}
 {"compilation_metrics": "METRICS", "frame_id": 0, "frame_compile_id": 0, "attempt": 1}
 {"dynamo_start": {"stack": "STACK"}, "frame_id": 1, "frame_compile_id": 0, "attempt": 0}
+{"artifact": {"name": "torch_dynamo_resume_in_fn_at_808_ORIGINAL_BYTECODE", "encoding": "string"}, "frame_id": 1, "frame_compile_id": 0, "attempt": 0, "has_payload": "HASH"}
 {"describe_storage": {"id": 0, "describer_id": "ID", "size": 4}, "frame_id": 1, "frame_compile_id": 0, "attempt": 0}
 {"describe_tensor": {"id": 0, "ndim": 1, "dtype": "torch.float32", "device": "device(type='cpu')", "size": [1], "dynamo_hint_overrides": {}, "is_leaf": true, "stride": [1], "storage": 0, "view_func": "VIEW_FUNC", "describer_id": "ID"}, "frame_id": 1, "frame_compile_id": 0, "attempt": 0}
 {"describe_source": {"describer_id": "ID", "id": 0, "source": "L['x']"}, "frame_id": 1, "frame_compile_id": 0, "attempt": 0}
 {"dynamo_output_graph": {"sizes": {"l_x_": [1], "add": [1]}}, "frame_id": 1, "frame_compile_id": 0, "attempt": 0, "has_payload": "HASH"}
+{"artifact": {"name": "aotautograd_cache_miss", "encoding": "json"}, "frame_id": 1, "frame_compile_id": 0, "attempt": 0, "has_payload": "HASH"}
 {"artifact": {"name": "before_pre_grad_graph", "encoding": "string"}, "frame_id": 1, "frame_compile_id": 0, "attempt": 0, "has_payload": "HASH"}
 {"artifact": {"name": "after_pre_grad_graph", "encoding": "string"}, "frame_id": 1, "frame_compile_id": 0, "attempt": 0, "has_payload": "HASH"}
-{"artifact": {"name": "aotautograd_cache_miss", "encoding": "json"}, "frame_id": 1, "frame_compile_id": 0, "attempt": 0, "has_payload": "HASH"}
 {"artifact": {"name": "aot_forward_graph_fw_metadata", "encoding": "string"}, "frame_id": 1, "frame_compile_id": 0, "attempt": 0, "has_payload": "HASH"}
 {"aot_inference_graph": {}, "frame_id": 1, "frame_compile_id": 0, "attempt": 0, "has_payload": "HASH"}
 {"artifact": {"name": "torch._functorch.config", "encoding": "string"}, "frame_id": 1, "frame_compile_id": 0, "attempt": 0, "has_payload": "HASH"}
@@ -826,6 +778,7 @@ def fn(x):
 {"artifact": {"name": "inductor_post_grad_graph", "encoding": "string"}, "frame_id": 1, "frame_compile_id": 0, "attempt": 0, "has_payload": "HASH"}
 {"inductor_output_code": {"filename": "FILENAME", "file_path": "FILENAME"}, "frame_id": 1, "frame_compile_id": 0, "attempt": 0, "has_payload": "HASH"}
 {"artifact": {"name": "fx_graph_cache_miss", "encoding": "json"}, "frame_id": 1, "frame_compile_id": 0, "attempt": 0, "has_payload": "HASH"}
+{"artifact": {"name": "torch_dynamo_resume_in_fn_at_808_MODIFIED_BYTECODE", "encoding": "string"}, "frame_id": 1, "frame_compile_id": 0, "attempt": 0, "has_payload": "HASH"}
 {"dynamo_cpp_guards_str": {}, "frame_id": 1, "frame_compile_id": 0, "attempt": 0, "has_payload": "HASH"}
 {"compilation_metrics": "METRICS", "frame_id": 1, "frame_compile_id": 0, "attempt": 0}
 """,  # noqa: B950
@@ -974,9 +927,9 @@ def fn(a):
 {"describe_tensor": {"id": 0, "ndim": 1, "dtype": "torch.float32", "device": "device(type='cpu')", "size": [1], "dynamo_hint_overrides": {}, "is_leaf": true, "stride": [1], "storage": 0, "view_func": "VIEW_FUNC", "describer_id": "ID"}, "frame_id": 0, "frame_compile_id": 0, "attempt": 0}
 {"describe_source": {"describer_id": "ID", "id": 0, "source": "L['a']"}, "frame_id": 0, "frame_compile_id": 0, "attempt": 0}
 {"dynamo_output_graph": {"sizes": {"l_a_": [1], "sin": [1]}}, "frame_id": 0, "frame_compile_id": 0, "attempt": 0, "has_payload": "HASH"}
+{"artifact": {"name": "aotautograd_cache_miss", "encoding": "json"}, "frame_id": 0, "frame_compile_id": 0, "attempt": 0, "has_payload": "HASH"}
 {"artifact": {"name": "before_pre_grad_graph", "encoding": "string"}, "frame_id": 0, "frame_compile_id": 0, "attempt": 0, "has_payload": "HASH"}
 {"artifact": {"name": "after_pre_grad_graph", "encoding": "string"}, "frame_id": 0, "frame_compile_id": 0, "attempt": 0, "has_payload": "HASH"}
-{"artifact": {"name": "aotautograd_cache_miss", "encoding": "json"}, "frame_id": 0, "frame_compile_id": 0, "attempt": 0, "has_payload": "HASH"}
 {"artifact": {"name": "aot_forward_graph_fw_metadata", "encoding": "string"}, "frame_id": 0, "frame_compile_id": 0, "attempt": 0, "has_payload": "HASH"}
 {"aot_inference_graph": {}, "frame_id": 0, "frame_compile_id": 0, "attempt": 0, "has_payload": "HASH"}
 {"artifact": {"name": "torch._functorch.config", "encoding": "string"}, "frame_id": 0, "frame_compile_id": 0, "attempt": 0, "has_payload": "HASH"}
@@ -994,8 +947,6 @@ def fn(a):
 {"describe_tensor": {"id": 0, "ndim": 1, "dtype": "torch.float32", "device": "device(type='cpu')", "size": [1], "dynamo_hint_overrides": {}, "is_leaf": true, "stride": [1], "storage": 0, "view_func": "VIEW_FUNC", "describer_id": "ID"}, "frame_id": 0, "frame_compile_id": 0, "attempt": 0}
 {"describe_source": {"describer_id": "ID", "id": 0, "source": "L['a']"}, "frame_id": 0, "frame_compile_id": 0, "attempt": 0}
 {"dynamo_output_graph": {"sizes": {"l_a_": [1], "sin": [1]}}, "frame_id": 0, "frame_compile_id": 0, "attempt": 0, "has_payload": "HASH"}
-{"artifact": {"name": "before_pre_grad_graph", "encoding": "string"}, "frame_id": 0, "frame_compile_id": 0, "attempt": 0, "has_payload": "HASH"}
-{"artifact": {"name": "after_pre_grad_graph", "encoding": "string"}, "frame_id": 0, "frame_compile_id": 0, "attempt": 0, "has_payload": "HASH"}
 {"artifact": {"name": "aot_forward_graph_fw_metadata", "encoding": "string"}, "frame_id": 0, "frame_compile_id": 0, "attempt": 0, "has_payload": "HASH"}
 {"aot_inference_graph": {}, "frame_id": 0, "frame_compile_id": 0, "attempt": 0, "has_payload": "HASH"}
 {"artifact": {"name": "fx_graph_runnable", "encoding": "string"}, "frame_id": 0, "frame_compile_id": 0, "attempt": 0, "has_payload": "HASH"}
@@ -1093,7 +1044,7 @@ def fn(a):
                 '"attempt": 0, "has_payload": "HASH"}'
             ),
             (
-                '{"chromium_event": {}, "compiled_autograd_id": 0, "frame_id": 2, "frame_compile_id": 0, '
+                '{"chromium_event": {}, "compiled_autograd_id": 0, "frame_id": 1, "frame_compile_id": 0, '
                 '"attempt": 0, "has_payload": "HASH"}'
             ),
         ]
@@ -1135,13 +1086,10 @@ def backward(ctx, gO):
             '{"dynamo_start": {"stack": "STACK"}, "frame_id": 0, "frame_compile_id": 0, "attempt": 0}',
             '{"dynamo_start": {"stack": "STACK"}, "frame_id": 1, "frame_compile_id": 0, "attempt": 0}',
             '{"dynamo_start": {"stack": "STACK"}, "frame_id": 2, "frame_compile_id": 0, "attempt": 0}',
+            '{"dynamo_start": {"stack": "STACK"}, "compiled_autograd_id": 0, "frame_id": 3, "frame_compile_id": 0, "attempt": 0}',
             '{"dynamo_start": {"stack": "STACK"}, "compiled_autograd_id": 0, "frame_id": 4, "frame_compile_id": 0, "attempt": 0}',
             '{"dynamo_start": {"stack": "STACK"}, "compiled_autograd_id": 0, "frame_id": 5, "frame_compile_id": 0, "attempt": 0}',
-            '{"dynamo_start": {"stack": "STACK"}, "compiled_autograd_id": 0, "frame_id": 6, "frame_compile_id": 0, "attempt": 0}',
-            '{"dynamo_start": {"stack": "STACK"}, "compiled_autograd_id": 0, "frame_id": 7, "frame_compile_id": 0, "attempt": 0}',
-            '{"dynamo_start": {"stack": "STACK"}, "compiled_autograd_id": 0, "frame_id": 8, "frame_compile_id": 0, "attempt": 0}',
-            '{"dynamo_start": {"stack": "STACK"}, "compiled_autograd_id": 0, "frame_id": 9, "frame_compile_id": 0, "attempt": 0}',
-            '{"dynamo_start": {"stack": "STACK"}, "compiled_autograd_id": 1, "frame_id": 12, "frame_compile_id": 1, "attempt": 0}',
+            '{"dynamo_start": {"stack": "STACK"}, "compiled_autograd_id": 1, "frame_id": 6, "frame_compile_id": 0, "attempt": 0}',
         ]
         logs = self.buffer.getvalue()
         self.assertTrue(all(event in logs for event in expected))
diff --git a/test/dynamo/test_subclasses.py b/test/dynamo/test_subclasses.py
index fe4a5b1103468..039b4be9dc6db 100644
--- a/test/dynamo/test_subclasses.py
+++ b/test/dynamo/test_subclasses.py
@@ -74,7 +74,8 @@ def mk_basic(base_is_nt):
         # view is not a leaf and has the same requires grad as its basic case
         x, _ = get_jagged_tensor(((2, 3, 4), 3), None, requires_grad=True)
         x = x.clone() if base_is_nt else x
-        assert not x.is_leaf
+        if x.is_leaf:
+            raise AssertionError("Expected x to not be a leaf")
         return x.unsqueeze(-1)
 
     def mk_leaf(base_is_nt, requires_grad_1, requires_grad_2):
@@ -234,7 +235,8 @@ def __tensor_flatten__(self):
 
     @staticmethod
     def __tensor_unflatten__(inner_tensors, metadata, outer_size, outer_stride):
-        assert len(inner_tensors) == 2
+        if len(inner_tensors) != 2:
+            raise AssertionError(f"Expected 2 inner tensors, got {len(inner_tensors)}")
         return ScaledTensor(
             inner_tensors["_data"],
             inner_tensors["_scale"],
@@ -374,6 +376,41 @@ def __torch_dispatch__(cls, func, types, args, kwargs):
         return return_and_correct_aliasing(func, args, kwargs, out)
 
 
+class DeferredInitSubclass(torch.Tensor):
+    """
+    A traceable wrapper subclass that calls super().__init__() BEFORE
+    setting instance attributes (similar to torchao patterns).
+    """
+
+    @staticmethod
+    def __new__(cls, data, scale):
+        return torch.Tensor._make_wrapper_subclass(
+            cls, data.shape, dtype=data.dtype, device=data.device
+        )
+
+    def __init__(self, data, scale):
+        super().__init__()
+        self._data = data
+        self._scale = scale
+
+    def __tensor_flatten__(self):
+        return ["_data"], {"_scale": self._scale}
+
+    @classmethod
+    def __tensor_unflatten__(cls, inner_tensors, ctx, outer_size, outer_stride):
+        return cls(inner_tensors["_data"], ctx["_scale"])
+
+    @classmethod
+    def __torch_dispatch__(cls, func, types, args, kwargs=None):
+        if kwargs is None:
+            kwargs = {}
+        args_inner = pytree.tree_map_only(DeferredInitSubclass, lambda x: x._data, args)
+        out = func(*args_inner, **kwargs)
+        if isinstance(out, torch.Tensor):
+            return DeferredInitSubclass(out, args[0]._scale)
+        return out
+
+
 def func(a):
     return a.sin()
 
@@ -485,6 +522,110 @@ def fn(x):
         self.assertIsInstance(out, BaseTorchFunction)
         self.assertEqual(out, out_opt)
 
+    def test_torch_function_subclass_with_mode(self):
+        # Subclass __torch_function__ must still be inlined when a
+        # TorchFunctionMode is active, otherwise the runtime wrapper
+        # (DisableTorchFunctionSubclass) is not applied and the subclass
+        # dispatch fires twice.
+        class ScaledTensor(torch.Tensor):
+            @staticmethod
+            def __new__(cls, data):
+                return torch.Tensor._make_subclass(cls, data)
+
+            @classmethod
+            def __torch_function__(cls, func, types, args=(), kwargs=None):
+                kwargs = kwargs or {}
+                raw = super().__torch_function__(func, types, args, kwargs)
+                if isinstance(raw, torch.Tensor) and func is torch.add:
+                    return raw * 2.0
+                return raw
+
+        class NoopMode(torch.overrides.TorchFunctionMode):
+            def __torch_function__(self, func, types, args=(), kwargs=None):
+                return func(*args, **(kwargs or {}))
+
+        a = ScaledTensor(torch.tensor([1.0, 2.0]))
+        b = ScaledTensor(torch.tensor([3.0, 4.0]))
+
+        with NoopMode():
+            eager = torch.add(a, b)
+            torch._dynamo.reset()
+            compiled = torch.compile(torch.add, backend="eager")(a, b)
+        self.assertEqual(eager, compiled)
+
+    def test_torch_function_reentrant_dispatch(self):
+        class ScaledTensor(torch.Tensor):
+            @staticmethod
+            def __new__(cls, data):
+                return torch.Tensor._make_subclass(cls, data)
+
+            @classmethod
+            def __torch_function__(cls, func, types, args=(), kwargs=None):
+                kwargs = kwargs or {}
+                raw = super().__torch_function__(func, types, args, kwargs)
+                if isinstance(raw, torch.Tensor) and func is torch.add:
+                    return raw * 2.0
+                return raw
+
+        a = ScaledTensor(torch.tensor([1.0, 2.0]))
+        b = ScaledTensor(torch.tensor([3.0, 4.0]))
+
+        eager = torch.add(a, b)
+        torch._dynamo.reset()
+        compiled = torch.compile(torch.add, backend="eager")(a, b)
+        self.assertEqual(eager, compiled)
+
+    def test_tensorify_under_disabled_torch_function(self):
+        # Fixes #180906
+        # The checks tensorify_python_scalars works under dispatch
+        # as it relies on MetaProxy's __torch_function__ to intercept calls
+        from torch._subclasses.fake_tensor import FakeTensorMode
+        from torch.fx.experimental.symbolic_shapes import ShapeEnv
+        from torch.fx.passes._tensorify_python_scalars import tensorify_python_scalars
+
+        # Build a minimal graph containing _local_scalar_dense (i.e. .item())
+        # on a floating-point placeholder — just enough for tensorify to act.
+        shape_env = ShapeEnv()
+        with FakeTensorMode(shape_env=shape_env) as fake_mode:
+            graph = torch.fx.Graph()
+            x = graph.placeholder("x")
+            x.meta["val"] = torch.randn(4)
+
+            scale_ph = graph.placeholder("scale")
+            scale_ph.meta["val"] = torch.tensor(1.0)
+
+            item_node = graph.call_function(
+                torch.ops.aten._local_scalar_dense.default, (scale_ph,)
+            )
+            # tensorify needs a backed SymFloat with a sympy expression
+            item_node.meta["val"] = shape_env.create_unbacked_symfloat()
+
+            mul_node = graph.call_function(torch.ops.aten.mul.Tensor, (x, item_node))
+            mul_node.meta["val"] = torch.randn(4)
+            graph.output(mul_node)
+
+            gm = torch.fx.GraphModule(torch.nn.Module(), graph)
+
+            # Run tensorify_python_scalars with __torch_function__
+            # disabled — without the _EnableTorchFunction fix, this raises:
+            #   RuntimeError: prims::convert_element_type() Expected
+            #   a value of type 'Tensor' ... found type 'MetaProxy'.
+            with torch._C.DisableTorchFunctionSubclass():
+                tensorify_python_scalars(gm, shape_env, fake_mode)
+
+            # The pass should have inserted a convert_element_type node
+            # that upcasts the scale placeholder to float64.
+            convert_nodes = [
+                n
+                for n in gm.graph.nodes
+                if n.op == "call_function"
+                and n.target is torch.ops.prims.convert_element_type.default
+            ]
+            self.assertTrue(len(convert_nodes) > 0)
+            # The first convert_element_type is the float64 upcast;
+            # verify it carries correct metadata.
+            self.assertEqual(convert_nodes[0].meta["val"].dtype, torch.float64)
+
     def test_torch_function_state_graph_break(self):
         @torch.compile(backend="eager")
         def fn(x):
@@ -728,7 +869,7 @@ def fn(a, w):
             a.add_(w)
             return a
 
-        fn_opt = torch.compile(fn)
+        fn_opt = torch.compile(fn, backend="eager")
 
         res_exp = fn(x, wrapped)
         res_act = fn_opt(y, wrapped2)
@@ -1815,7 +1956,7 @@ def __metadata_guard__(cls, y):
         self.assertRaisesRegex(
             torch._dynamo.exc.InternalTorchDynamoError,
             "Tensor subclass method __metadata_guard__ must take exactly two subclass metadata arguments",
-            lambda: torch.compile(lambda x: x * x)(x),
+            lambda: torch.compile(lambda x: x * x, backend="eager")(x),
         )
 
     def test_tensor_subclass_ctx_custom_guards_error_not_classmethod(self):
@@ -1829,8 +1970,39 @@ def __metadata_guard__(self, x, y):
         self.assertRaisesRegex(
             torch._dynamo.exc.InternalTorchDynamoError,
             "Tensor subclass method __metadata_guard__ must be a classmethod",
-            lambda: torch.compile(lambda x: x * x)(x),
+            lambda: torch.compile(lambda x: x * x, backend="eager")(x),
+        )
+
+    def test_tensor_subclass_metadata_with_symint(self):
+        # TENSOR_SUBCLASS_METADATA_MATCH replaces SymInts in metadata with
+        # _AnyCompare sentinels so that (a) deepcopy doesn't pull in the
+        # ShapeEnv, (b) dynamic dims aren't over-guarded, and (c) unbacked
+        # SymInts don't cause errors.
+        from torch._subclasses.fake_tensor import FakeTensorMode
+
+        shape_env = ShapeEnv()
+        fake_mode = FakeTensorMode(shape_env=shape_env)
+        sym_ctx = StatelessSymbolicContext(
+            dynamic_sizes=[DimDynamic.DYNAMIC],
         )
+        fake_t = fake_mode.from_tensor(torch.randn(8), symbolic_context=sym_ctx)
+        sym_int = fake_t.shape[0]  # SymInt with hint 8
+
+        # Construct a real tensor whose metadata contains a SymInt.
+        x1 = CtxSubclassTensor(torch.randn(8, 8), sym_int)
+
+        @torch.compile(backend="eager", dynamic=True, fullgraph=True)
+        def f(x):
+            return x * x
+
+        f(x1)
+
+        # Without the fix, the guard stores the SymInt's hint value (8)
+        # and recompiles when it sees a different constant. With the fix,
+        # the SymInt position is replaced by _AnyCompare so the guard
+        # passes for any constant value.
+        x2 = CtxSubclassTensor(torch.randn(8, 8), 3)
+        _check_recompiles(self, f, (x1,), (x2,), False)
 
     def test_subclass_constructor_proxying(self):
         import dataclasses
@@ -1903,7 +2075,7 @@ def __torch_dispatch__(cls, func, types, args, kwargs):
                 )
                 return return_and_correct_aliasing(func, args, kwargs, out)
 
-        @torch.compile(fullgraph=True)
+        @torch.compile(fullgraph=True, backend="eager")
         def f1(x):
             meta = SubclassTensorArgs(
                 x.shape, x.device, SubclassTensorArgs(x.shape, x.device, None)
@@ -1914,7 +2086,7 @@ def f1(x):
         x = torch.randn(3, 3)
         f1(x)
 
-        @torch.compile(fullgraph=True)
+        @torch.compile(fullgraph=True, backend="eager")
         def f1(x):
             meta = SubclassTensorArgs2(
                 x.shape, x.device, SubclassTensorArgs2(x.shape, x.device, None)
@@ -1986,7 +2158,7 @@ def __torch_dispatch__(cls, func, types, args=(), kwargs=None):
                 )
                 return output
 
-        @torch.compile(dynamic=True)
+        @torch.compile(dynamic=True, backend="eager")
         def f(x):
             return x.unflatten(-1, [2, 5])
 
@@ -2164,11 +2336,10 @@ def f(x):
                 out_test = compiled_f(view)
                 self.assertEqual(out_ref, out_test)
 
-    @torch._dynamo.config.patch("inline_inbuilt_nn_modules", True)
     @parametrize("dynamic", [True, False])
     def test_mark_static_with_subclass_desugaring(self, dynamic):
         from collections.abc import Callable
-        from typing import Any, Optional
+        from typing import Any
 
         from torch._dynamo.decorators import mark_static_address
         from torch._inductor.compile_fx import compile_fx
@@ -2182,16 +2353,17 @@ def test_mark_static_with_subclass_desugaring(self, dynamic):
         def inner_compile(
             gm: torch.fx.GraphModule,
             example_inputs: list[torch.Tensor],
-            cudagraphs: Optional[BoxedBool] = None,
-            static_input_idxs: Optional[list[int]] = None,
+            cudagraphs: BoxedBool | None = None,
+            static_input_idxs: list[int] | None = None,
             is_backward: bool = False,
-            graph_id: Optional[int] = None,
+            graph_id: int | None = None,
             cpp_wrapper: bool = False,
             aot_mode: bool = False,
             is_inference: bool = False,
-            boxed_forward_device_index: Optional[BoxedDeviceIndex] = None,
-            layout_opt: Optional[bool] = None,
-            extern_node_serializer: Optional[Callable[[list[Any]], Any]] = None,
+            boxed_forward_device_index: BoxedDeviceIndex | None = None,
+            layout_opt: bool | None = None,
+            extern_node_serializer: Callable[[list[Any]], Any] | None = None,
+            **kwargs: Any,
         ):
             if dynamic:
                 self.assertEqual(static_input_idxs, [2, 3, 4])
@@ -2207,10 +2379,9 @@ def fn(t0, t1, t2):
 
         fn(torch.ones(4), x, torch.ones(4))
 
-    @torch._dynamo.config.patch("inline_inbuilt_nn_modules", True)
     def test_subclass_parameters_are_static_under_training(self):
         from collections.abc import Callable
-        from typing import Any, Optional
+        from typing import Any
 
         from torch._inductor.compile_fx import compile_fx
         from torch._inductor.cudagraph_utils import BoxedDeviceIndex
@@ -2219,16 +2390,17 @@ def test_subclass_parameters_are_static_under_training(self):
         def inner_compile(
             gm: torch.fx.GraphModule,
             example_inputs: list[torch.Tensor],
-            cudagraphs: Optional[BoxedBool] = None,
-            static_input_idxs: Optional[list[int]] = None,
+            cudagraphs: BoxedBool | None = None,
+            static_input_idxs: list[int] | None = None,
             is_backward: bool = False,
-            graph_id: Optional[int] = None,
+            graph_id: int | None = None,
             cpp_wrapper: bool = False,
             aot_mode: bool = False,
             is_inference: bool = False,
-            boxed_forward_device_index: Optional[BoxedDeviceIndex] = None,
-            layout_opt: Optional[bool] = None,
-            extern_node_serializer: Optional[Callable[[list[Any]], Any]] = None,
+            boxed_forward_device_index: BoxedDeviceIndex | None = None,
+            layout_opt: bool | None = None,
+            extern_node_serializer: Callable[[list[Any]], Any] | None = None,
+            **kwargs: Any,
         ):
             # Important bit: there are 3 params: linear.weight.a, linear.weight.b, linear.bias,
             # which are the first 3 args of the graph.
@@ -3053,7 +3225,7 @@ def __new__(cls, a, b, outer_size=None, outer_stride=None):
                 if outer_stride is None:
                     outer_stride = a.stride()
 
-                assert (
+                assert (  # noqa: S101
                     a.device == b.device
                     and a.layout == b.layout
                     and a.requires_grad == b.requires_grad
@@ -3072,14 +3244,14 @@ def __new__(cls, a, b, outer_size=None, outer_stride=None):
 
             @staticmethod
             def __tensor_unflatten__(inner_tensors, meta, outer_size, outer_stride):
-                assert meta is None
+                assert meta is None  # noqa: S101
                 a, b = inner_tensors["a"], inner_tensors["b"]
                 if type(a) is torch.Tensor:
-                    assert outer_size is not None
-                    assert outer_stride is not None
+                    assert outer_size is not None  # noqa: S101
+                    assert outer_stride is not None  # noqa: S101
                 return TT(a, b, outer_size, outer_stride)
 
-        @torch.compile(dynamic=True)
+        @torch.compile(dynamic=True, backend="eager")
         def f(x, y):
             tmp1 = x.sin()
             tmp2 = y.sin()
@@ -3328,6 +3500,33 @@ def forward(
 """,  # noqa: B950
         )
 
+    def test_deferred_init_subclass_init_not_traced(self):
+        """
+        Tracing a function that constructs a DeferredInitSubclass must not crash.
+
+        The bug: when Dynamo's frame hook intercepts __init__ as a root frame,
+        self is partially initialised (attributes not yet set by __init__).
+        Previously, wrap_tensor would call __tensor_flatten__ on this
+        partially-initialised self and raise AttributeError.
+
+        The fix skips tracing __init__ of traceable wrapper subclasses at the
+        frame level (convert_frame.py), so __init__ runs eagerly like
+        @torch._disable_dynamo would.
+        """
+        # Compile __init__ directly, simulating the root-frame interception
+        # scenario that occurs in practice (e.g. Diffusers + TorchAO + Dynamo).
+        compiled_init = torch.compile(
+            DeferredInitSubclass.__init__, backend="eager", fullgraph=False
+        )
+        data = torch.randn(4, 4)
+        shell = DeferredInitSubclass.__new__(DeferredInitSubclass, data, 2.0)
+
+        # Should not raise AttributeError from __tensor_flatten__ on partial self
+        compiled_init(shell, data, 2.0)
+
+        self.assertEqual(shell._data, data)
+        self.assertEqual(shell._scale, 2.0)
+
 
 instantiate_parametrized_tests(SubclassTests)
 
@@ -3690,7 +3889,7 @@ def fn(values, offsets, max_seqlen):
                 values, t.offsets(), max_seqlen=t._maybe_max_seqlen
             )
 
-        opt_fn = torch.compile(fn, fullgraph=True, dynamic=True)
+        opt_fn = torch.compile(fn, fullgraph=True, dynamic=True, backend="eager")
         values = torch.randn(10, 5)
         offsets = torch.tensor([0, 2, 4, 7, 10])
         max_seqlen = 5
diff --git a/test/dynamo/test_subgraphs.py b/test/dynamo/test_subgraphs.py
index 16c765bfb1409..7095b4f80db1f 100644
--- a/test/dynamo/test_subgraphs.py
+++ b/test/dynamo/test_subgraphs.py
@@ -119,7 +119,7 @@ def fn(a, b):
             c2 = b - a
             return indirectly_unsupported(c1, c2)
 
-        self._common(fn, 2, 3)
+        self._common(fn, 1, 3)
 
     def test_indirect_unsupported2(self):
         def fn(a, b):
@@ -129,14 +129,14 @@ def fn(a, b):
             c2 = b - a
             return local_const1 / (local_const2 - indirectly_unsupported(c1, c2))
 
-        self._common(fn, 3, 5)
+        self._common(fn, 2, 5)
 
     def test_indirect_unsupported3(self):
         def fn(a, b):
             args = [a - b, b - a]
             return indirectly_unsupported(*args)
 
-        self._common(fn, 2, 3)
+        self._common(fn, 1, 3)
 
     def test_stack_state1(self):
         def fn(a, b):
@@ -156,7 +156,7 @@ def fn(a, b):
             c2 = b - a
             return t1 / (t2 - indirectly_unsupported(c1, c2))
 
-        self._common(fn, 3, 7)
+        self._common(fn, 2, 7)
 
     def test_multigraph(self):
         def fn(a, b):
@@ -199,7 +199,7 @@ def fn(a, b):
             x = x + 2.0
             return x
 
-        self._common(fn, 3, 7)
+        self._common(fn, 2, 7)
 
     def test_resume3(self):
         def fn(a, b):
@@ -212,7 +212,7 @@ def fn(a, b):
             x = x + 2.0
             return x
 
-        self._common(fn, 3, 7)
+        self._common(fn, 2, 7)
 
     def test_resume4(self):
         def fn(a, b):
@@ -225,7 +225,7 @@ def fn(a, b):
             x = x + 2.0
             return x
 
-        self._common(fn, 3, 7)
+        self._common(fn, 2, 7)
 
     def test_resume5(self):
         def fn(a, b):
diff --git a/test/dynamo/test_tp_slots.py b/test/dynamo/test_tp_slots.py
new file mode 100644
index 0000000000000..6952b92058a27
--- /dev/null
+++ b/test/dynamo/test_tp_slots.py
@@ -0,0 +1,280 @@
+# Owner(s): ["module: dynamo"]
+
+"""Tests for CPython type slot detection in Dynamo.
+
+Tests that get_type_slots correctly identifies which protocol methods
+(sequence, mapping, number, type) are implemented by various types.
+"""
+
+import collections.abc
+import dataclasses
+import enum
+
+from torch._C._dynamo import (
+    get_type_slots,
+    has_slot,
+    PyMappingSlots,
+    PyNumberSlots,
+    PySequenceSlots,
+    PyTypeSlots,
+)
+from torch.testing._internal.common_utils import run_tests, TestCase
+
+
+class TestTypeSlots(TestCase):
+    """Test suite for type slot detection."""
+
+    def _get_slot_info(self, obj_type):
+        """Helper to get and unpack slot information."""
+        seq_slots, map_slots, num_slots, type_slots = get_type_slots(obj_type)
+        return seq_slots, map_slots, num_slots, type_slots
+
+    def test_dict_slots(self):
+        """Test that dict has mapping protocol but not sequence protocol."""
+        seq_slots, map_slots, num_slots, type_slots = self._get_slot_info(dict)
+
+        # dict should NOT have sq_length (sequence protocol)
+        self.assertFalse(has_slot(seq_slots, PySequenceSlots.SQ_LENGTH))
+
+        # dict SHOULD have mapping protocol
+        self.assertTrue(has_slot(map_slots, PyMappingSlots.MP_LENGTH))
+        self.assertTrue(has_slot(map_slots, PyMappingSlots.MP_SUBSCRIPT))
+        self.assertTrue(has_slot(map_slots, PyMappingSlots.MP_ASS_SUBSCRIPT))
+
+    def test_list_slots(self):
+        """Test that list has sequence protocol."""
+        seq_slots, map_slots, num_slots, type_slots = self._get_slot_info(list)
+
+        # list SHOULD have sequence protocol
+        self.assertTrue(has_slot(seq_slots, PySequenceSlots.SQ_LENGTH))
+        self.assertTrue(has_slot(seq_slots, PySequenceSlots.SQ_ITEM))
+        self.assertTrue(has_slot(seq_slots, PySequenceSlots.SQ_CONTAINS))
+
+        # list also has mapping protocol for compatibility (mp_length, mp_subscript, mp_ass_subscript)
+        self.assertTrue(has_slot(map_slots, PyMappingSlots.MP_LENGTH))
+        self.assertTrue(has_slot(map_slots, PyMappingSlots.MP_SUBSCRIPT))
+
+    def test_tuple_slots(self):
+        """Test that tuple has sequence protocol."""
+        seq_slots, map_slots, num_slots, type_slots = self._get_slot_info(tuple)
+
+        # tuple SHOULD have sequence protocol
+        self.assertTrue(has_slot(seq_slots, PySequenceSlots.SQ_LENGTH))
+        self.assertTrue(has_slot(seq_slots, PySequenceSlots.SQ_ITEM))
+        self.assertTrue(has_slot(seq_slots, PySequenceSlots.SQ_CONTAINS))
+
+    def test_set_slots(self):
+        """Test that set has sequence protocol for contains."""
+        seq_slots, map_slots, num_slots, type_slots = self._get_slot_info(set)
+
+        # set SHOULD have sq_contains
+        self.assertTrue(has_slot(seq_slots, PySequenceSlots.SQ_CONTAINS))
+
+        # set should NOT have mapping protocol
+        self.assertFalse(has_slot(map_slots, PyMappingSlots.MP_LENGTH))
+
+    def test_dict_subclass_slots(self):
+        """Test that dict subclasses have both mapping and sequence protocol."""
+
+        class DictSubclass(dict):
+            pass
+
+        seq_slots, map_slots, num_slots, type_slots = self._get_slot_info(DictSubclass)
+
+        # Dict subclasses expose both protocols for compatibility
+        self.assertTrue(has_slot(seq_slots, PySequenceSlots.SQ_LENGTH))
+        self.assertTrue(has_slot(map_slots, PyMappingSlots.MP_LENGTH))
+
+    def test_list_subclass_slots(self):
+        """Test that list subclasses have sequence protocol."""
+
+        class ListSubclass(list):
+            pass
+
+        seq_slots, map_slots, num_slots, type_slots = self._get_slot_info(ListSubclass)
+
+        # List subclasses should have sequence protocol
+        self.assertTrue(has_slot(seq_slots, PySequenceSlots.SQ_LENGTH))
+        self.assertTrue(has_slot(seq_slots, PySequenceSlots.SQ_ITEM))
+
+    def test_int_slots(self):
+        """Test that int has number protocol."""
+        seq_slots, map_slots, num_slots, type_slots = self._get_slot_info(int)
+
+        # int SHOULD have number protocol operations
+        self.assertTrue(has_slot(num_slots, PyNumberSlots.NB_ADD))
+        self.assertTrue(has_slot(num_slots, PyNumberSlots.NB_SUBTRACT))
+        self.assertTrue(has_slot(num_slots, PyNumberSlots.NB_MULTIPLY))
+
+    def test_str_slots(self):
+        """Test that str has sequence protocol."""
+        seq_slots, map_slots, num_slots, type_slots = self._get_slot_info(str)
+
+        # str SHOULD have sequence protocol
+        self.assertTrue(has_slot(seq_slots, PySequenceSlots.SQ_LENGTH))
+        self.assertTrue(has_slot(seq_slots, PySequenceSlots.SQ_ITEM))
+        self.assertTrue(has_slot(seq_slots, PySequenceSlots.SQ_CONTAINS))
+
+    def test_type_has_call_slot(self):
+        """Test that type objects have tp_call slot."""
+        seq_slots, map_slots, num_slots, type_slots = self._get_slot_info(type)
+
+        # type SHOULD have tp_call (for calling classes/types)
+        self.assertTrue(has_slot(type_slots, PyTypeSlots.TP_CALL))
+
+    def test_type_has_hash_slot(self):
+        """Test that most types have tp_hash slot."""
+        seq_slots, map_slots, num_slots, type_slots = self._get_slot_info(int)
+
+        # int SHOULD have tp_hash
+        self.assertTrue(has_slot(type_slots, PyTypeSlots.TP_HASH))
+
+    def test_custom_class_slots(self):
+        """Test that custom user-defined classes have minimal slots."""
+
+        class CustomClass:
+            pass
+
+        seq_slots, map_slots, num_slots, type_slots = self._get_slot_info(CustomClass)
+
+        # Custom classes should not have sequence/mapping/number protocol by default
+        self.assertFalse(has_slot(seq_slots, PySequenceSlots.SQ_LENGTH))
+        self.assertFalse(has_slot(map_slots, PyMappingSlots.MP_LENGTH))
+
+    def test_set_subclass_slots(self):
+        """Test that set subclasses inherit set protocol."""
+
+        class SetSubclass(set):
+            pass
+
+        seq_slots, map_slots, num_slots, type_slots = self._get_slot_info(SetSubclass)
+
+        # Set subclasses should have sq_contains
+        self.assertTrue(has_slot(seq_slots, PySequenceSlots.SQ_CONTAINS))
+
+    def test_tuple_subclass_slots(self):
+        """Test that tuple subclasses inherit sequence protocol."""
+
+        class TupleSubclass(tuple):
+            __slots__ = ()
+
+        seq_slots, map_slots, num_slots, type_slots = self._get_slot_info(TupleSubclass)
+
+        # Tuple subclasses should have sequence protocol
+        self.assertTrue(has_slot(seq_slots, PySequenceSlots.SQ_LENGTH))
+        self.assertTrue(has_slot(seq_slots, PySequenceSlots.SQ_ITEM))
+
+    def test_mutable_mapping_slots(self):
+        """Test that MutableMapping ABC has mapping protocol."""
+
+        class MyMapping(collections.abc.MutableMapping):
+            def __init__(self):
+                self._data = {}
+
+            def __getitem__(self, key):
+                return self._data[key]
+
+            def __setitem__(self, key, value):
+                self._data[key] = value
+
+            def __delitem__(self, key):
+                del self._data[key]
+
+            def __iter__(self):
+                return iter(self._data)
+
+            def __len__(self):
+                return len(self._data)
+
+        seq_slots, map_slots, num_slots, type_slots = self._get_slot_info(MyMapping)
+
+        # MutableMapping subclasses should have mapping protocol
+        self.assertTrue(has_slot(map_slots, PyMappingSlots.MP_LENGTH))
+        self.assertTrue(has_slot(map_slots, PyMappingSlots.MP_SUBSCRIPT))
+        self.assertTrue(has_slot(map_slots, PyMappingSlots.MP_ASS_SUBSCRIPT))
+
+    def test_frozen_dataclass_slots(self):
+        """Test that frozen dataclasses have standard object slots."""
+
+        @dataclasses.dataclass(frozen=True)
+        class FrozenData:
+            x: int
+            y: str
+
+        seq_slots, map_slots, num_slots, type_slots = self._get_slot_info(FrozenData)
+
+        # Frozen dataclasses should have basic type slots like hash
+        self.assertTrue(has_slot(type_slots, PyTypeSlots.TP_HASH))
+
+    def test_enum_slots(self):
+        """Test that Enum types have expected slots."""
+
+        class Color(enum.Enum):
+            RED = 1
+            GREEN = 2
+            BLUE = 3
+
+        seq_slots, map_slots, num_slots, type_slots = self._get_slot_info(Color)
+
+        # Enums should have type protocol slots
+        self.assertTrue(has_slot(type_slots, PyTypeSlots.TP_HASH))
+
+    def test_enum_member_slots(self):
+        """Test that individual enum members work correctly."""
+
+        class Status(enum.Enum):
+            PENDING = "pending"
+            ACTIVE = "active"
+
+        # Enum members are instances, test the enum class
+        seq_slots, map_slots, num_slots, type_slots = self._get_slot_info(Status)
+        self.assertTrue(has_slot(type_slots, PyTypeSlots.TP_HASH))
+
+    def test_metaclass_slots(self):
+        """Test that metaclass types have tp_call for instantiation."""
+
+        class MyMeta(type):
+            pass
+
+        class MyClass(metaclass=MyMeta):
+            pass
+
+        seq_slots, map_slots, num_slots, type_slots = self._get_slot_info(MyMeta)
+
+        # Metaclasses should have tp_call for creating instances
+        self.assertTrue(has_slot(type_slots, PyTypeSlots.TP_CALL))
+
+    def test_dict_subclass_with_custom_len(self):
+        """Test dict subclass with custom __len__ (the original bug case)."""
+
+        class DictWithCustomLen(dict):
+            def __len__(self):
+                return super().__len__() - 1
+
+        seq_slots, map_slots, num_slots, type_slots = self._get_slot_info(
+            DictWithCustomLen
+        )
+
+        # Should have both protocols (CPython compatibility)
+        self.assertTrue(has_slot(seq_slots, PySequenceSlots.SQ_LENGTH))
+        self.assertTrue(has_slot(map_slots, PyMappingSlots.MP_LENGTH))
+
+    def test_frozenset_slots(self):
+        """Test that frozenset has sequence protocol for contains."""
+        seq_slots, map_slots, num_slots, type_slots = self._get_slot_info(frozenset)
+
+        # frozenset SHOULD have sq_contains
+        self.assertTrue(has_slot(seq_slots, PySequenceSlots.SQ_CONTAINS))
+
+    def test_float_slots(self):
+        """Test that float has number protocol."""
+        seq_slots, map_slots, num_slots, type_slots = self._get_slot_info(float)
+
+        # float SHOULD have number protocol operations
+        self.assertTrue(has_slot(num_slots, PyNumberSlots.NB_ADD))
+        self.assertTrue(has_slot(num_slots, PyNumberSlots.NB_SUBTRACT))
+        self.assertTrue(has_slot(num_slots, PyNumberSlots.NB_MULTIPLY))
+
+
+if __name__ == "__main__":
+    run_tests()
diff --git a/test/dynamo/test_trace_rules.py b/test/dynamo/test_trace_rules.py
index 8833445bb495f..2496e0e9701d9 100644
--- a/test/dynamo/test_trace_rules.py
+++ b/test/dynamo/test_trace_rules.py
@@ -474,6 +474,8 @@ def test_no_special_handlers_for_torch_non_c_bindings(self):
             "handle_assert",  # No global state (constant)
             "handle_nested_tensor",  # No global state
             "handle_current_stream",  # Safely implemented
+            "handle_synchronize",  # Device type from function identity or arg
+            "handle_functorch_autograd_grad",  # Only inspects placeholder metadata
         )
         for fn in handlers:
             if isinstance(fn, staticmethod) or inspect.ismethod(fn):
diff --git a/test/dynamo/test_tree_map.py b/test/dynamo/test_tree_map.py
index 23a1f7f68aa20..dffa408a4b223 100644
--- a/test/dynamo/test_tree_map.py
+++ b/test/dynamo/test_tree_map.py
@@ -1,28 +1,18 @@
 # Owner(s): ["module: dynamo"]
 
-try:
-    import optree
-except ImportError:  # pragma: no cover
-    optree = None
-
 from collections import namedtuple
+from dataclasses import dataclass
 from typing import NamedTuple
 
 import torch
 import torch._dynamo
+import torch.utils._pytree as python_pytree
+from torch._dynamo.test_case import run_tests, TestCase
 from torch.testing._internal.common_utils import (
     instantiate_parametrized_tests,
     parametrize,
-    run_tests,
-    TestCase,
+    subtest,
 )
-from torch.utils import _pytree as pytree
-
-
-try:
-    import torch.utils._cxx_pytree as cxx_pytree
-except ImportError:  # pragma: no cover
-    cxx_pytree = None
 
 
 def _tensor_leaf(*values):
@@ -52,33 +42,56 @@ def _tuple_is_leaf(node):
     return isinstance(node, tuple)
 
 
+pytree_modules = {
+    "python": python_pytree,
+}
+if python_pytree._cxx_pytree_dynamo_traceable:
+    import torch.utils._cxx_pytree as cxx_pytree
+
+    pytree_modules["cxx"] = cxx_pytree
+    pytree_modules["native_optree"] = cxx_pytree.optree
+else:
+    cxx_pytree = None
+
+optree = cxx_pytree.optree if cxx_pytree is not None else None
+
+
 def _require_optree(test_case):
     if optree is None:
         test_case.skipTest("optree is unavailable")
 
 
-TREE_MAP_IMPLEMENTATIONS = []
-if optree is not None:
-    TREE_MAP_IMPLEMENTATIONS.append(("optree", optree.tree_map))
-TREE_MAP_IMPLEMENTATIONS.append(("pytree_python", pytree.tree_map))
-if cxx_pytree is not None:
-    TREE_MAP_IMPLEMENTATIONS.append(("pytree_cxx", cxx_pytree.tree_map))
+parametrize_pytree_module = parametrize(
+    "pytree_name,pytree",
+    [subtest((name, module), name=name) for name, module in pytree_modules.items()],
+)
+
+_PYTREE_MODULES_WITH_PATH = {"python"}
+
+parametrize_pytree_module_with_path = parametrize(
+    "pytree_name,pytree",
+    [
+        subtest((name, module), name=name)
+        for name, module in pytree_modules.items()
+        if name in _PYTREE_MODULES_WITH_PATH
+    ],
+)
 
 
 KWARG_CASES = [
     ("default", {}, None),
-    ("none_is_leaf", {"none_is_leaf": True}, {"optree"}),
+    ("none_is_leaf", {"none_is_leaf": True}, {"native_optree"}),
     ("is_leaf", {"is_leaf": _tuple_is_leaf}, None),
-    ("namespace", {"namespace": "torch"}, {"optree"}),
+    ("namespace", {"namespace": "torch"}, {"native_optree"}),
     (
         "namespace_and_none_is_leaf",
         {"namespace": "torch", "none_is_leaf": True},
-        {"optree"},
+        {"native_optree"},
     ),
     (
         "namespace_none_is_leaf_predicate",
         {"namespace": "torch", "none_is_leaf": True, "is_leaf": _tuple_is_leaf},
-        {"optree"},
+        {"native_optree"},
     ),
 ]
 
@@ -104,8 +117,8 @@ def _build_tree(offset: int) -> dict[str, object]:
 
 
 def _assert_trees_allclose(test_case: TestCase, ref, res) -> None:
-    ref_flat, ref_spec = pytree.tree_flatten(ref)
-    res_flat, res_spec = pytree.tree_flatten(res)
+    ref_flat, ref_spec = python_pytree.tree_flatten(ref)
+    res_flat, res_spec = python_pytree.tree_flatten(res)
     test_case.assertEqual(ref_spec, res_spec)
     for expected, actual in zip(ref_flat, res_flat):
         if isinstance(expected, torch.Tensor):
@@ -120,37 +133,35 @@ def setUp(self):
         super().setUp()
         torch._dynamo.reset()
 
-    def _run_tree_map(self, tree_map_impl, kwargs):
+    def _run_tree_map(self, pytree, kwargs):
         lhs = _build_tree(0)
         rhs = _build_tree(7)
 
         def fn(a, b):
-            return tree_map_impl(_combine_leaves, a, b, **kwargs)
+            return pytree.tree_map(_combine_leaves, a, b, **kwargs)
 
         compiled = torch.compile(fn, backend="eager", fullgraph=True)
         expected = fn(lhs, rhs)
         result = compiled(lhs, rhs)
         _assert_trees_allclose(self, expected, result)
 
-    @parametrize("tree_map_name,tree_map_impl", TREE_MAP_IMPLEMENTATIONS)
+    @parametrize_pytree_module
     @parametrize("kwargs_name,kwargs,allowed_impls", KWARG_CASES)
     def test_tree_map_variants(
         self,
-        tree_map_name: str,
-        tree_map_impl,
+        pytree_name: str,
+        pytree,
         kwargs_name: str,
         kwargs: dict,
         allowed_impls,
     ) -> None:
-        if tree_map_name == "pytree_cxx" and cxx_pytree is None:
-            self.skipTest("torch.utils._cxx_pytree is unavailable")
-        if allowed_impls is not None and tree_map_name not in allowed_impls:
+        if allowed_impls is not None and pytree_name not in allowed_impls:
             self.skipTest("kwargs unsupported for implementation")
-        self._run_tree_map(tree_map_impl, kwargs)
+        self._run_tree_map(pytree, kwargs)
 
     def test_tree_map_rejects_mismatched_container_types(self) -> None:
         def fn(a, b):
-            return pytree.tree_map(lambda u, v: u + v, a, b)
+            return python_pytree.tree_map(lambda u, v: u + v, a, b)
 
         lhs = [torch.ones(2), torch.ones(2)]
         rhs = (torch.ones(2), torch.ones(2))
@@ -167,7 +178,7 @@ def fn(a, b):
 
     def test_tree_map_is_leaf_handles_tensor_nodes(self) -> None:
         def fn(tree):
-            return pytree.tree_map(
+            return python_pytree.tree_map(
                 lambda pair: torch.stack(pair).sum(dim=0),
                 tree,
                 is_leaf=lambda node: isinstance(node, tuple),
@@ -188,7 +199,7 @@ def mapper(node):
             return node + 2
 
         def fn(arg):
-            return pytree.tree_map_only(torch.Tensor, mapper, arg)
+            return python_pytree.tree_map_only(torch.Tensor, mapper, arg)
 
         compiled = torch.compile(fn, backend="eager", fullgraph=True)
         expected = fn(tree)
@@ -200,7 +211,7 @@ def test_tree_map_only_multiple_trees_falls_back(self) -> None:
         rhs = {"a": torch.ones(2) * 3, "b": torch.ones(2) * 4}
 
         def fn(a, b):
-            return pytree.tree_map_only(torch.Tensor, lambda x, y: x + y, a, b)
+            return python_pytree.tree_map_only(torch.Tensor, lambda x, y: x + y, a, b)
 
         with self.assertRaisesRegex(TypeError, "callable"):
             fn(lhs, rhs)
@@ -223,7 +234,7 @@ def mapper(node):
             raise AssertionError("unexpected node passed to mapper")
 
         def fn(arg):
-            return pytree.tree_map_only((int, tuple), mapper, arg)
+            return python_pytree.tree_map_only((int, tuple), mapper, arg)
 
         compiled = torch.compile(fn, backend="eager", fullgraph=True)
         expected = fn(tree)
@@ -243,7 +254,7 @@ def mapper(node):
             return node * 2 if isinstance(node, torch.Tensor) else node
 
         def fn(arg):
-            return pytree.tree_map(mapper, arg, is_leaf=is_leaf)
+            return python_pytree.tree_map(mapper, arg, is_leaf=is_leaf)
 
         compiled = torch.compile(fn, backend="eager", fullgraph=True)
         expected = fn(tree)
@@ -260,7 +271,7 @@ def mapper(node):
             return node + 5 if isinstance(node, torch.Tensor) else node
 
         def fn(arg):
-            return pytree.tree_map_only(selector, mapper, arg)
+            return python_pytree.tree_map_only(selector, mapper, arg)
 
         compiled = torch.compile(fn, backend="eager", fullgraph=True)
         expected = fn(tree)
@@ -286,15 +297,15 @@ def fn(a, b):
         ):
             compiled(lhs, rhs)
 
-    @parametrize("tree_map_name,tree_map_impl", TREE_MAP_IMPLEMENTATIONS)
+    @parametrize_pytree_module
     def test_tree_map_none_nodes_default_behavior(
-        self, tree_map_name: str, tree_map_impl
+        self, pytree_name: str, pytree
     ) -> None:
-        if tree_map_name == "optree":
+        if pytree_name == "native_optree":
             self.skipTest("optree treats None as an internal node by default")
 
         def fn(a, b):
-            return tree_map_impl(lambda u, v: (u, v), a, b)
+            return pytree.tree_map(lambda u, v: (u, v), a, b)
 
         tree = {"k": None}
         compiled = torch.compile(fn, backend="eager", fullgraph=True)
@@ -359,9 +370,9 @@ def mapper(node):
         self.assertIs(result["nested"]["dtype"], torch.float64)
         self.assertEqual(result, expected)
 
-    @parametrize("tree_map_name,tree_map_impl", TREE_MAP_IMPLEMENTATIONS)
+    @parametrize_pytree_module
     def test_user_defined_object_treated_as_leaf(
-        self, tree_map_name: str, tree_map_impl
+        self, pytree_name: str, pytree
     ) -> None:
         """User-defined objects (not registered in pytree) should be treated as leaves."""
 
@@ -383,7 +394,7 @@ def mapper(node):
             return node
 
         def fn(arg):
-            return tree_map_impl(mapper, arg)
+            return pytree.tree_map(mapper, arg)
 
         compiled = torch.compile(fn, backend="eager", fullgraph=True)
 
@@ -404,10 +415,8 @@ def fn(arg):
         self.assertTrue(torch.allclose(result["custom"].value, obj1.value * 3))
         self.assertTrue(torch.allclose(result["tensor"], torch.ones(2)))
 
-    @parametrize("tree_map_name,tree_map_impl", TREE_MAP_IMPLEMENTATIONS)
-    def test_user_defined_object_multiple_trees(
-        self, tree_map_name: str, tree_map_impl
-    ) -> None:
+    @parametrize_pytree_module
+    def test_user_defined_object_multiple_trees(self, pytree_name: str, pytree) -> None:
         """User-defined objects should work correctly with multiple input trees."""
 
         class Point:
@@ -424,7 +433,7 @@ def mapper(a, b):
             return a + b
 
         def fn(t1, t2):
-            return tree_map_impl(mapper, t1, t2)
+            return pytree.tree_map(mapper, t1, t2)
 
         compiled = torch.compile(fn, backend="eager", fullgraph=True)
         result = compiled(tree1, tree2)
@@ -434,10 +443,8 @@ def fn(t1, t2):
         self.assertEqual(result["point"].y, 6)
         self.assertEqual(result["val"], 30)
 
-    @parametrize("tree_map_name,tree_map_impl", TREE_MAP_IMPLEMENTATIONS)
-    def test_dict_subclass_treated_as_leaf(
-        self, tree_map_name: str, tree_map_impl
-    ) -> None:
+    @parametrize_pytree_module
+    def test_dict_subclass_treated_as_leaf(self, pytree_name: str, pytree) -> None:
         """Dict subclasses (not registered in pytree) should be treated as leaves."""
 
         class MyDict(dict):
@@ -456,7 +463,7 @@ def mapper(node):
             return node
 
         def fn(arg):
-            return tree_map_impl(mapper, arg)
+            return pytree.tree_map(mapper, arg)
 
         compiled = torch.compile(fn, backend="eager", fullgraph=True)
         result = compiled(tree)
@@ -468,10 +475,8 @@ def fn(arg):
         # Regular dict should still be traversed
         self.assertEqual(result["regular"]["x"], 3)
 
-    @parametrize("tree_map_name,tree_map_impl", TREE_MAP_IMPLEMENTATIONS)
-    def test_list_subclass_treated_as_leaf(
-        self, tree_map_name: str, tree_map_impl
-    ) -> None:
+    @parametrize_pytree_module
+    def test_list_subclass_treated_as_leaf(self, pytree_name: str, pytree) -> None:
         """List subclasses (not registered in pytree) should be treated as leaves."""
 
         class MyList(list):
@@ -489,7 +494,7 @@ def mapper(node):
             return node
 
         def fn(arg):
-            return tree_map_impl(mapper, arg)
+            return pytree.tree_map(mapper, arg)
 
         compiled = torch.compile(fn, backend="eager", fullgraph=True)
         result = compiled(tree)
@@ -500,10 +505,8 @@ def fn(arg):
         # Regular list should be traversed
         self.assertEqual(result["regular"], [14, 15, 16])
 
-    @parametrize("tree_map_name,tree_map_impl", TREE_MAP_IMPLEMENTATIONS)
-    def test_tuple_subclass_treated_as_leaf(
-        self, tree_map_name: str, tree_map_impl
-    ) -> None:
+    @parametrize_pytree_module
+    def test_tuple_subclass_treated_as_leaf(self, pytree_name: str, pytree) -> None:
         """Tuple subclasses (not registered in pytree) should be treated as leaves."""
 
         class MyTuple(tuple):  # noqa: SLOT001
@@ -520,7 +523,7 @@ def mapper(node):
             return node
 
         def fn(arg):
-            return tree_map_impl(mapper, arg)
+            return pytree.tree_map(mapper, arg)
 
         compiled = torch.compile(fn, backend="eager", fullgraph=True)
         result = compiled(tree)
@@ -531,9 +534,9 @@ def fn(arg):
         # Regular tuple should be traversed
         self.assertEqual(result["regular"], (14, 15, 16))
 
-    @parametrize("tree_map_name,tree_map_impl", TREE_MAP_IMPLEMENTATIONS)
+    @parametrize_pytree_module
     def test_user_defined_object_nested_in_containers(
-        self, tree_map_name: str, tree_map_impl
+        self, pytree_name: str, pytree
     ) -> None:
         """User-defined objects nested inside containers should be leaves."""
 
@@ -552,7 +555,7 @@ def mapper(node):
             return node
 
         def fn(arg):
-            return tree_map_impl(mapper, arg)
+            return pytree.tree_map(mapper, arg)
 
         compiled = torch.compile(fn, backend="eager", fullgraph=True)
         result = compiled(tree)
@@ -561,9 +564,9 @@ def fn(arg):
         self.assertEqual(result["list_of_wrappers"][1].value, 20)
         self.assertEqual(result["nested"]["wrapper"].value, 30)
 
-    @parametrize("tree_map_name,tree_map_impl", TREE_MAP_IMPLEMENTATIONS)
+    @parametrize_pytree_module
     def test_user_defined_object_with_is_leaf_predicate(
-        self, tree_map_name: str, tree_map_impl
+        self, pytree_name: str, pytree
     ) -> None:
         """Test that is_leaf predicate interacts correctly with user-defined objects."""
 
@@ -586,7 +589,7 @@ def mapper(node):
             return node
 
         def fn(arg):
-            return tree_map_impl(mapper, arg, is_leaf=is_leaf_fn)
+            return pytree.tree_map(mapper, arg, is_leaf=is_leaf_fn)
 
         compiled = torch.compile(fn, backend="eager", fullgraph=True)
         result = compiled(tree)
@@ -602,7 +605,7 @@ def __init__(self, items):
                 self.items = list(items)
 
         # Register with pytree
-        pytree.register_pytree_node(
+        python_pytree.register_pytree_node(
             RegisteredContainer,
             lambda x: (x.items, None),
             lambda items, _: RegisteredContainer(items),
@@ -617,7 +620,7 @@ def mapper(node):
                 return node
 
             def fn(arg):
-                return pytree.tree_map(mapper, arg)
+                return python_pytree.tree_map(mapper, arg)
 
             compiled = torch.compile(fn, backend="eager", fullgraph=True)
             result = compiled(tree)
@@ -628,7 +631,38 @@ def fn(arg):
             self.assertTrue(torch.allclose(result.items[1], torch.zeros(2) + 1))
         finally:
             # Clean up registration
-            pytree._deregister_pytree_node(RegisteredContainer)
+            python_pytree._deregister_pytree_node(RegisteredContainer)
+
+    def test_registered_custom_type_multiple_trees_with_eq(self) -> None:
+        """Dataclass registrations compare class objects through the metaclass."""
+
+        @dataclass
+        class RegisteredContainer:
+            x: torch.Tensor
+            y: torch.Tensor
+
+        python_pytree.register_pytree_node(
+            RegisteredContainer,
+            lambda c: ([c.x, c.y], None),
+            lambda children, _: RegisteredContainer(*children),
+        )
+
+        try:
+            lhs = RegisteredContainer(torch.zeros(2, 3), torch.ones(2, 3))
+            rhs = RegisteredContainer(torch.ones(2, 3), torch.zeros(2, 3))
+
+            def fn(a, b):
+                return python_pytree.tree_map(lambda x, y: x + y, a, b)
+
+            compiled = torch.compile(fn, backend="eager", fullgraph=True)
+            expected = fn(lhs, rhs)
+            result = compiled(lhs, rhs)
+
+            self.assertIsInstance(result, RegisteredContainer)
+            self.assertEqual(result.x, expected.x)
+            self.assertEqual(result.y, expected.y)
+        finally:
+            python_pytree._deregister_pytree_node(RegisteredContainer)
 
     def test_registered_custom_type_falls_back_optree(self) -> None:
         """Custom types registered with optree should fall back to tracing."""
@@ -824,8 +858,8 @@ def fn(arg):
         finally:
             optree.unregister_pytree_node(NamespacedContainer, namespace="namespace_a")
 
-    @parametrize("tree_map_name,tree_map_impl", TREE_MAP_IMPLEMENTATIONS)
-    def test_dataclass_treated_as_leaf(self, tree_map_name: str, tree_map_impl) -> None:
+    @parametrize_pytree_module
+    def test_dataclass_treated_as_leaf(self, pytree_name: str, pytree) -> None:
         """Dataclasses should be treated as leaves (not registered by default)."""
         import dataclasses
 
@@ -842,7 +876,7 @@ def mapper(node):
             return node + 10
 
         def fn(arg):
-            return tree_map_impl(mapper, arg)
+            return pytree.tree_map(mapper, arg)
 
         compiled = torch.compile(fn, backend="eager", fullgraph=True)
         result = compiled(tree)
@@ -852,9 +886,9 @@ def fn(arg):
         self.assertEqual(result["point"].y, 4)
         self.assertEqual(result["val"], 13)
 
-    @parametrize("tree_map_name,tree_map_impl", TREE_MAP_IMPLEMENTATIONS)
+    @parametrize_pytree_module
     def test_user_defined_object_with_tensor_attribute(
-        self, tree_map_name: str, tree_map_impl
+        self, pytree_name: str, pytree
     ) -> None:
         """User-defined objects containing tensors should be treated as leaves."""
 
@@ -873,7 +907,7 @@ def mapper(node):
             return node
 
         def fn(arg):
-            return tree_map_impl(mapper, arg)
+            return pytree.tree_map(mapper, arg)
 
         compiled = torch.compile(fn, backend="eager", fullgraph=True)
         result = compiled(tree)
@@ -882,10 +916,8 @@ def fn(arg):
         self.assertTrue(torch.allclose(result["wrapper"].tensor, torch.ones(2, 2) * 2))
         self.assertTrue(torch.allclose(result["direct_tensor"], torch.ones(2)))
 
-    @parametrize("tree_map_name,tree_map_impl", TREE_MAP_IMPLEMENTATIONS)
-    def test_user_defined_object_no_fallback(
-        self, tree_map_name: str, tree_map_impl
-    ) -> None:
+    @parametrize_pytree_module
+    def test_user_defined_object_no_fallback(self, pytree_name: str, pytree) -> None:
         """Verify user-defined objects use fastpath without triggering fallback."""
         import logging
 
@@ -901,7 +933,7 @@ def mapper(node):
             return node + 1
 
         def fn(arg):
-            return tree_map_impl(mapper, arg)
+            return pytree.tree_map(mapper, arg)
 
         # Capture debug logs to ensure no fallback is triggered
         log_records = []
@@ -937,8 +969,8 @@ def fn(arg):
             logger.removeHandler(handler)
             logger.setLevel(old_level)
 
-    @parametrize("tree_map_name,tree_map_impl", TREE_MAP_IMPLEMENTATIONS)
-    def test_namedtuple_tree_map(self, tree_map_name: str, tree_map_impl) -> None:
+    @parametrize_pytree_module
+    def test_namedtuple_tree_map(self, pytree_name: str, pytree) -> None:
         """Test tree_map with namedtuple uses fast path."""
         Point = namedtuple("Point", ["x", "y"])
 
@@ -953,7 +985,7 @@ def mapper(node):
             return node
 
         def fn(arg):
-            return tree_map_impl(mapper, arg)
+            return pytree.tree_map(mapper, arg)
 
         compiled = torch.compile(fn, backend="eager", fullgraph=True)
         expected = fn(tree)
@@ -967,10 +999,8 @@ def fn(arg):
         self.assertTrue(torch.allclose(result["nested"][0].y, torch.ones(3) * 2 + 1))
         _assert_trees_allclose(self, expected, result)
 
-    @parametrize("tree_map_name,tree_map_impl", TREE_MAP_IMPLEMENTATIONS)
-    def test_namedtuple_tree_map_multiple_trees(
-        self, tree_map_name: str, tree_map_impl
-    ) -> None:
+    @parametrize_pytree_module
+    def test_namedtuple_tree_map_multiple_trees(self, pytree_name: str, pytree) -> None:
         """Test tree_map with multiple namedtuple trees."""
         Point = namedtuple("Point", ["x", "y"])
 
@@ -983,7 +1013,7 @@ def mapper(a, b):
             return a
 
         def fn(t1, t2):
-            return tree_map_impl(mapper, t1, t2)
+            return pytree.tree_map(mapper, t1, t2)
 
         compiled = torch.compile(fn, backend="eager", fullgraph=True)
         expected = fn(tree1, tree2)
@@ -994,8 +1024,8 @@ def fn(t1, t2):
         self.assertTrue(torch.allclose(result["point"].y, torch.ones(2) * 3))
         _assert_trees_allclose(self, expected, result)
 
-    @parametrize("tree_map_name,tree_map_impl", TREE_MAP_IMPLEMENTATIONS)
-    def test_namedtuple_with_is_leaf(self, tree_map_name: str, tree_map_impl) -> None:
+    @parametrize_pytree_module
+    def test_namedtuple_with_is_leaf(self, pytree_name: str, pytree) -> None:
         """Test tree_map with namedtuple and is_leaf predicate."""
 
         class Point(NamedTuple):
@@ -1016,7 +1046,7 @@ def mapper(node):
             return node
 
         def fn(arg):
-            return tree_map_impl(mapper, arg, is_leaf=is_leaf_fn)
+            return pytree.tree_map(mapper, arg, is_leaf=is_leaf_fn)
 
         compiled = torch.compile(fn, backend="eager", fullgraph=True)
         expected = fn(tree)
@@ -1029,5 +1059,203 @@ def fn(arg):
         _assert_trees_allclose(self, expected, result)
 
 
+@instantiate_parametrized_tests
+class TreeMapWithPathCompileTests(TestCase):
+    def setUp(self):
+        super().setUp()
+        torch._dynamo.reset()
+
+    @parametrize_pytree_module_with_path
+    def test_basic_nested_tree(self, pytree_name: str, pytree) -> None:
+        """Keypaths are correctly constructed for nested dicts, lists, and tuples."""
+        tree = {
+            "tensor": torch.ones(2),
+            "list": [torch.zeros(2), torch.ones(3)],
+            "tuple": (torch.ones(4),),
+        }
+
+        collected = []
+
+        def mapper(kp, x):
+            collected.append(kp)
+            return x + 1 if isinstance(x, torch.Tensor) else x
+
+        def fn(arg):
+            return pytree.tree_map_with_path(mapper, arg)
+
+        compiled = torch.compile(fn, backend="eager", fullgraph=True)
+
+        # Eager reference
+        collected.clear()
+        expected = fn(tree)
+        eager_keypaths = list(collected)
+
+        # Compiled
+        collected.clear()
+        result = compiled(tree)
+        compiled_keypaths = list(collected)
+
+        _assert_trees_allclose(self, expected, result)
+        self.assertEqual(eager_keypaths, compiled_keypaths)
+
+    @parametrize_pytree_module_with_path
+    def test_multiple_trees(self, pytree_name: str, pytree) -> None:
+        """tree_map_with_path with multiple input trees."""
+        tree1 = {"a": torch.ones(2), "b": [torch.zeros(3)]}
+        tree2 = {"a": torch.ones(2) * 2, "b": [torch.ones(3) * 3]}
+
+        def mapper(kp, x, y):
+            return x + y
+
+        def fn(t1, t2):
+            return pytree.tree_map_with_path(mapper, t1, t2)
+
+        compiled = torch.compile(fn, backend="eager", fullgraph=True)
+        expected = fn(tree1, tree2)
+        result = compiled(tree1, tree2)
+        _assert_trees_allclose(self, expected, result)
+
+    @parametrize_pytree_module_with_path
+    def test_is_leaf_predicate(self, pytree_name: str, pytree) -> None:
+        """is_leaf stops traversal and passes the subtree as a leaf."""
+        tree = {"a": [torch.ones(2), torch.zeros(2)]}
+
+        def is_leaf_fn(node):
+            return isinstance(node, list)
+
+        collected_keypaths = []
+
+        def mapper(kp, x):
+            collected_keypaths.append(kp)
+            if isinstance(x, list):
+                return [t * 2 for t in x]
+            return x
+
+        def fn(arg):
+            return pytree.tree_map_with_path(mapper, arg, is_leaf=is_leaf_fn)
+
+        compiled = torch.compile(fn, backend="eager", fullgraph=True)
+
+        collected_keypaths.clear()
+        fn(tree)
+        eager_keypaths = list(collected_keypaths)
+
+        collected_keypaths.clear()
+        compiled(tree)
+        compiled_keypaths = list(collected_keypaths)
+
+        # The list should be treated as a single leaf with keypath (MappingKey("a"),)
+        self.assertEqual(len(eager_keypaths), 1)
+        self.assertEqual(eager_keypaths, compiled_keypaths)
+
+    @parametrize_pytree_module_with_path
+    def test_namedtuple_uses_getattr_key(self, pytree_name: str, pytree) -> None:
+        """Namedtuple fields produce GetAttrKey in keypaths."""
+        Point = namedtuple("Point", ["x", "y"])
+        tree = {"point": Point(torch.ones(2), torch.zeros(2))}
+
+        collected = []
+
+        def mapper(kp, x):
+            collected.append(kp)
+            if isinstance(x, torch.Tensor):
+                return x + 1
+            return x
+
+        def fn(arg):
+            return pytree.tree_map_with_path(mapper, arg)
+
+        compiled = torch.compile(fn, backend="eager", fullgraph=True)
+
+        collected.clear()
+        expected = fn(tree)
+        eager_keypaths = list(collected)
+
+        collected.clear()
+        result = compiled(tree)
+        compiled_keypaths = list(collected)
+
+        _assert_trees_allclose(self, expected, result)
+        self.assertEqual(eager_keypaths, compiled_keypaths)
+
+    @parametrize_pytree_module_with_path
+    def test_deeply_nested_keypaths(self, pytree_name: str, pytree) -> None:
+        """Deeply nested structures produce correct multi-level keypaths."""
+        tree = {"outer": {"inner": [torch.ones(2)]}}
+
+        collected = []
+
+        def mapper(kp, x):
+            collected.append(kp)
+            return x * 2 if isinstance(x, torch.Tensor) else x
+
+        def fn(arg):
+            return pytree.tree_map_with_path(mapper, arg)
+
+        compiled = torch.compile(fn, backend="eager", fullgraph=True)
+
+        collected.clear()
+        expected = fn(tree)
+        eager_keypaths = list(collected)
+
+        collected.clear()
+        result = compiled(tree)
+        compiled_keypaths = list(collected)
+
+        self.assertEqual(len(eager_keypaths), 1)
+        self.assertEqual(len(eager_keypaths[0]), 3)  # outer -> inner -> [0]
+        self.assertEqual(eager_keypaths, compiled_keypaths)
+        _assert_trees_allclose(self, expected, result)
+
+    @parametrize_pytree_module_with_path
+    def test_keypath_values_used_in_computation(self, pytree_name: str, pytree) -> None:
+        """The map function can use keypath values to influence the result."""
+        from torch.utils._pytree import MappingKey
+
+        tree = {"scale2": torch.ones(2), "scale3": torch.ones(2)}
+
+        def mapper(kp, x):
+            # Use the dict key name from the keypath to determine scaling
+            key = kp[-1]
+            if isinstance(key, MappingKey) and key.key == "scale3":
+                return x * 3
+            return x * 2
+
+        def fn(arg):
+            return pytree.tree_map_with_path(mapper, arg)
+
+        compiled = torch.compile(fn, backend="eager", fullgraph=True)
+        expected = fn(tree)
+        result = compiled(tree)
+        _assert_trees_allclose(self, expected, result)
+
+    @parametrize_pytree_module_with_path
+    def test_user_defined_object_treated_as_leaf(
+        self, pytree_name: str, pytree
+    ) -> None:
+        """Unregistered user-defined objects are leaves in tree_map_with_path."""
+
+        class MyObj:
+            def __init__(self, val):
+                self.val = val
+
+        tree = {"obj": MyObj(42), "tensor": torch.ones(2)}
+
+        def mapper(kp, x):
+            if isinstance(x, MyObj):
+                return MyObj(x.val * 2)
+            return x + 1
+
+        def fn(arg):
+            return pytree.tree_map_with_path(mapper, arg)
+
+        compiled = torch.compile(fn, backend="eager", fullgraph=True)
+        result = compiled(tree)
+
+        self.assertIsInstance(result["obj"], MyObj)
+        self.assertEqual(result["obj"].val, 84)
+        self.assertTrue(torch.allclose(result["tensor"], torch.ones(2) + 1))
+
+
 if __name__ == "__main__":  # pragma: no cover
     run_tests()
diff --git a/test/dynamo/test_unittest.py b/test/dynamo/test_unittest.py
index de1364d424d7f..df1d1d7419b77 100644
--- a/test/dynamo/test_unittest.py
+++ b/test/dynamo/test_unittest.py
@@ -8,10 +8,12 @@
 
 class TestUnittest(torch._dynamo.test_case.TestCase):
     def setUp(self):
+        super().setUp()
         self._prev = torch._dynamo.config.enable_trace_unittest
         torch._dynamo.config.enable_trace_unittest = True
 
     def tearDown(self):
+        super().tearDown()
         torch._dynamo.config.enable_trace_unittest = self._prev
 
     @make_dynamo_test
diff --git a/test/dynamo/test_unspec.py b/test/dynamo/test_unspec.py
index 2085e46c5000e..c27fe14a0e9ff 100644
--- a/test/dynamo/test_unspec.py
+++ b/test/dynamo/test_unspec.py
@@ -1,4 +1,5 @@
 # Owner(s): ["module: dynamo"]
+import contextlib
 import math
 import random
 import unittest
@@ -314,6 +315,75 @@ def fn(x, rand1, rand2):
         y2 = opt_fn(inp, *get_rng())
         self.assertEqual(y1, y2)
 
+    def test_random_in_dynamo(self):
+        # test that system random calls still work even
+        # if Dynamo calls random methods.
+
+        exit_stack = contextlib.ExitStack()
+
+        def patch_fn_with_rng_burn(name):
+            orig_fn = eval(name)
+
+            def bad(*args, **kwargs):
+                # burn random call within dynamo
+                random.random()
+                return orig_fn(*args, **kwargs)
+
+            exit_stack.enter_context(unittest.mock.patch(name, bad))
+
+        x = torch.tensor([[1.0, 2.0], [3.0, 4.0]])
+
+        # we don't guard against random calls in eval_frame.py today
+        # patch_fn_with_rng_burn("torch._dynamo.eval_frame._maybe_set_eval_frame")
+        patch_fn_with_rng_burn("torch._dynamo.convert_frame._compile")
+        patch_fn_with_rng_burn(
+            "torch._dynamo.symbolic_convert.InstructionTranslator.run"
+        )
+
+        def f1(x):
+            # simple test
+            r1 = random.randint(1, 9)
+            y = x + random.uniform(10, 20)
+            r2 = random.randrange(0, 10)
+            return y + r1, r2
+
+        random.seed(1)
+        ref1 = f1(x)
+        opt_f1 = torch.compile(f1, backend="eager", fullgraph=True)
+        random.seed(1)
+        res1 = opt_f1(x)
+        self.assertEqual(ref1, res1)
+
+        def f2(x):
+            # test with graph breaks
+            r1 = random.randint(1, 9)
+            x = x + r1
+            torch._dynamo.graph_break()
+            r2 = random.randint(10, 19)
+            x = x + r2
+            return x, r1, r2
+
+        random.seed(2)
+        ref2 = f2(x)
+        opt_f2 = torch.compile(f2, backend="eager")
+        random.seed(2)
+        res2 = opt_f2(x)
+        self.assertEqual(ref2, res2)
+
+        def f3(x):
+            # test consecutive calls
+            return x + random.randint(1, 10)
+
+        random.seed(3)
+        ref3 = f3(x)
+        ref3_ = f3(x)
+        opt_f3 = torch.compile(f3, backend="eager", fullgraph=True)
+        random.seed(3)
+        res3 = opt_f3(x)
+        res3_ = opt_f3(x)
+        self.assertEqual(ref3, res3)
+        self.assertEqual(ref3_, res3_)
+
     def test_builtin_getitem(self):
         # builtin getitem args[0] is python list and args[1] is unspec
         def fn(x, idx):
@@ -420,7 +490,7 @@ def fn(x):
             comptime.assert_static(x.size(0))
             return x + 1
 
-        opt_fn = torch.compile(fn, dynamic=True, fullgraph=True)
+        opt_fn = torch.compile(fn, dynamic=True, fullgraph=True, backend="eager")
         opt_fn(torch.randn(12, 23))
 
     def test_shape_graph_break(self):
@@ -437,7 +507,7 @@ def fn(x):
 
     def test_isinstance_symint(self):
         def fn(x):
-            assert isinstance(x.size(0), int)
+            assert isinstance(x.size(0), int)  # noqa: S101
             return x * 2
 
         x = torch.randn(20)
@@ -467,7 +537,7 @@ def func(x):
             out = F.conv1d(x, kernel, padding=padding, stride=2)
             return out
 
-        opt_func = torch.compile(func)
+        opt_func = torch.compile(func, backend="eager")
 
         x = torch.randn(1, 1, 175)
         opt_func(x)  # passes
@@ -493,7 +563,9 @@ def test_rshift_dynamic(self):
         def shift_right(tensor: torch.Tensor) -> torch.Tensor:
             return (tensor >> 2).to(torch.long)
 
-        opt_fn = torch.compile(shift_right, fullgraph=True, dynamic=True)
+        opt_fn = torch.compile(
+            shift_right, fullgraph=True, dynamic=True, backend="eager"
+        )
         sample_input = torch.tensor([4, 4, 16, 32], dtype=torch.uint8)
         opt_fn(sample_input)
 
@@ -877,7 +949,7 @@ def test_mark_unbacked_hint_consistency(self):
         x = torch.randn(1)
         torch._dynamo.decorators.mark_unbacked(x, 0)
 
-        @torch.compile()
+        @torch.compile(backend="eager")
         def f(x):
             if guard_size_oblivious(x.size(0) != 1):
                 return x + 3
@@ -913,6 +985,31 @@ def forward(self, x: torch.Tensor, val: int) -> torch.Tensor:
         o1_2 = opt_model(x2, 2)
         self.assertEqual(o1_2_ref, o1_2)
 
+    def test_float_guard_source_on_recompile(self):
+        # Regression test: when a float attribute triggers recompilation and
+        # becomes dynamic, the guard produced by produce_guards_verbose should
+        # have a proper source annotation, not "(unknown source)".
+        cache = {}
+
+        class Module(torch.nn.Module):
+            def __init__(self, key: float):
+                super().__init__()
+                self.key = key
+                cache[key] = torch.randn(16)
+
+            def forward(self, x):
+                return x + cache[self.key]
+
+        x = torch.randn(16)
+        log_stream, ctx = logs_to_string("torch._dynamo.guards", "guards")
+        with ctx():
+            for key in [1.0, 2.0, 3.0]:
+                model = torch.compile(Module(key))
+                model(x)
+
+        guard_log = log_stream.getvalue()
+        self.assertNotIn("unknown source", guard_log)
+
 
 class UnspecTestsDevice(torch._dynamo.test_case.TestCase):
     def test_builtin_functions_on_device(self, device):
diff --git a/test/dynamo/test_user_defined_object.py b/test/dynamo/test_user_defined_object.py
new file mode 100644
index 0000000000000..eb920a800c4f7
--- /dev/null
+++ b/test/dynamo/test_user_defined_object.py
@@ -0,0 +1,695 @@
+# Owner(s): ["module: dynamo"]
+
+import dataclasses
+import types
+import unittest
+
+import torch
+import torch._dynamo.testing as dynamo_testing
+from torch._dynamo.test_case import run_tests, TestCase
+
+
+class SlotsOnly:
+    __slots__ = ("x", "y")
+
+    def __init__(self, x, y):
+        self.x = x
+        self.y = y
+
+
+class SlotsAndDict:
+    __slots__ = ("x", "__dict__")
+
+    def __init__(self, x):
+        self.x = x
+
+
+@dataclasses.dataclass(frozen=True, slots=True)
+class FrozenSlots:
+    x: int
+    y: int
+
+
+class SlotsAndSetattr:
+    __slots__ = ("x",)
+
+    def __init__(self, x):
+        self.x = x
+
+    def __setattr__(self, name, value):
+        object.__setattr__(self, name, value * 2)
+
+
+class SlotsAndDictAndSetattr:
+    __slots__ = ("x", "__dict__")
+
+    def __init__(self, x):
+        self.x = x
+
+    def __setattr__(self, name, value):
+        object.__setattr__(self, name, value * 2)
+
+
+class SlotsBase:
+    __slots__ = ("x",)
+
+    def __init__(self):
+        self.x = 0
+
+
+class SlotsDerived(SlotsBase):
+    __slots__ = ("y",)
+
+    def __init__(self):
+        super().__init__()
+        self.y = 0
+
+
+class Plain:
+    pass
+
+
+class SlotsChildOfPlain(Plain):
+    __slots__ = ("z",)
+
+    def __init__(self):
+        self.z = 0
+
+
+class Slots:
+    __slots__ = ("x",)
+
+
+class SlotsShadowed(SlotsBase):
+    x = 42  # class attribute shadows parent's slot descriptor
+
+
+class SlotsAndProperty:
+    __slots__ = ("_x",)
+
+    def __init__(self, x):
+        self._x = x
+
+    @property
+    def x(self):
+        return self._x
+
+    @x.setter
+    def x(self, value):
+        self._x = value * 2
+
+
+class TestSlotsAttrAssignment(TestCase):
+    """Tests for attribute assignment on objects with __slots__."""
+
+    def test_valid_slot_assignment(self):
+        # Case 1: assign to a declared slot — should succeed
+        def fn(t):
+            obj = SlotsOnly(1, 2)
+            obj.x = 99
+            return t + obj.x
+
+        dynamo_testing.standard_test(self, fn, nargs=1)
+
+    def test_invalid_slot_assignment_raises(self):
+        # Case 2: assign to an undeclared attr on a slotted object (no __dict__)
+        # should raise AttributeError in eager; compiled raises an exception too
+        def fn(t):
+            obj = SlotsOnly(1, 2)
+            obj.z = 99
+            return t + obj.x
+
+        compiled_fn = torch.compile(fn, backend="eager", fullgraph=True)
+        t = torch.ones(1)
+        self.assertRaises(AttributeError, fn, t)
+        self.assertRaises(Exception, compiled_fn, t)
+
+    def test_slots_with_dict_allows_arbitrary_attrs(self):
+        # Case 3: __slots__ includes __dict__ — arbitrary attr assignment should work
+        def fn(t):
+            obj = SlotsAndDict(1)
+            obj.extra = 42
+            return t + obj.x + obj.extra
+
+        dynamo_testing.standard_test(self, fn, nargs=1)
+
+    def test_frozen_dataclass_with_slots_construction(self):
+        # Case 4: frozen dataclass with slots uses object.__setattr__ in __init__
+        # to bypass the frozen __setattr__. Dynamo must allow this for slot descriptors.
+        def fn(t):
+            obj = FrozenSlots(3, 4)
+            return t + obj.x + obj.y
+
+        dynamo_testing.standard_test(self, fn, nargs=1)
+
+    def test_custom_setattr_with_slots(self):
+        # Case 5: __slots__ + custom __setattr__ — the custom __setattr__ is traced
+        def fn(t):
+            obj = SlotsAndSetattr(1)
+            obj.x = 10
+            return t + obj.x
+
+        dynamo_testing.standard_test(self, fn, nargs=1)
+
+    def test_slots_with_dict_valid_slot_assignment(self):
+        # Case 6: __slots__ + __dict__: assigning to a declared slot still works
+        def fn(t):
+            obj = SlotsAndDict(1)
+            obj.x = 99
+            return t + obj.x
+
+        dynamo_testing.standard_test(self, fn, nargs=1)
+
+    def test_slots_with_dict_undeclared_attr_goes_to_dict(self):
+        # Case 7: __slots__ + __dict__: assigning to an undeclared attr goes to
+        # __dict__ instead of raising AttributeError
+        def fn(t):
+            obj = SlotsAndDict(1)
+            obj.z = 42
+            return t + obj.z
+
+        dynamo_testing.standard_test(self, fn, nargs=1)
+
+    def test_custom_setattr_with_slots_and_dict(self):
+        # Case 8: __slots__ + __dict__ + custom __setattr__ — custom __setattr__
+        # is traced for both slot and non-slot attrs
+        def fn(t):
+            obj = SlotsAndDictAndSetattr(1)
+            obj.x = 10
+            obj.extra = 3
+            return t + obj.x + obj.extra
+
+        dynamo_testing.standard_test(self, fn, nargs=1)
+
+    def test_slots_inheritance_parent_and_child_slots(self):
+        # Subclass adds its own slot on top of parent's slot — both accessible
+        def fn(t):
+            obj = SlotsDerived()
+            obj.x = 1
+            obj.y = 2
+            return t + obj.x + obj.y
+
+        dynamo_testing.standard_test(self, fn, nargs=1)
+
+    def test_slots_child_inherits_dict_from_no_slots_parent(self):
+        # Subclass with __slots__ inheriting from a parent without __slots__
+        # gets __dict__ from the parent, so arbitrary attrs are allowed
+        def fn(t):
+            obj = SlotsChildOfPlain()
+            obj.z = 1
+            obj.extra = 42
+            return t + obj.z + obj.extra
+
+        dynamo_testing.standard_test(self, fn, nargs=1)
+
+    def test_object_setattr_bypasses_custom_setattr(self):
+        # object.__setattr__ skips the custom __setattr__ and writes directly to slot
+        def fn(t):
+            obj = SlotsAndSetattr(1)
+            object.__setattr__(obj, "x", 5)
+            return t + obj.x
+
+        dynamo_testing.standard_test(self, fn, nargs=1)
+
+    def test_getattr_default_on_unset_slot(self):
+        # getattr with a default on an unset slot returns the default
+        def fn(t):
+            obj = Slots()
+            val = getattr(obj, "x", 99)
+            return t + val
+
+        dynamo_testing.standard_test(self, fn, nargs=1)
+
+    def test_slot_read_after_delete_raises(self):
+        # Reading a slot after deletion raises AttributeError in both eager and compiled
+        def fn(t):
+            obj = Slots()
+            obj.x = 1
+            del obj.x
+            return t + obj.x
+
+        compiled_fn = torch.compile(fn, backend="eager", fullgraph=True)
+        t = torch.ones(1)
+        self.assertRaises(AttributeError, fn, t)
+        self.assertRaises(Exception, compiled_fn, t)
+
+    def test_slot_shadowed_by_class_attribute(self):
+        # Class attribute in subclass shadows parent slot descriptor:
+        # reads return the class attribute, writes raise AttributeError
+        def fn(t):
+            obj = SlotsShadowed()
+            return t + obj.x  # returns class attr 42
+
+        dynamo_testing.standard_test(self, fn, nargs=1)
+
+    def test_slot_assignment_with_object_as_argument(self):
+        # Slotted object passed as argument (not created inside fn)
+        def fn(t, obj):
+            obj.x = 10
+            return t + obj.x
+
+        compiled_fn = torch.compile(fn, backend="eager", fullgraph=True)
+        t = torch.ones(1)
+        obj = Slots()
+        self.assertEqual(fn(t.clone(), obj), compiled_fn(t.clone(), obj))
+
+    def test_slot_mutation_materialized_on_argument(self):
+        # Slot mutation on an object passed as argument must be visible after
+        # the compiled function returns (side effect materialization)
+        def fn(t, obj):
+            obj.x = 10
+            return t.sin()
+
+        compiled_fn = torch.compile(fn, backend="eager", fullgraph=True)
+        obj = Slots()
+        compiled_fn(torch.ones(1), obj)
+        self.assertEqual(obj.x, 10)
+
+    def test_slot_delete_materialized(self):
+        # del on a slot inside a compiled fn must be visible after the call returns
+        def fn(t, obj):
+            del obj.x
+            return t.sin()
+
+        compiled_fn = torch.compile(fn, backend="eager", fullgraph=True)
+        obj = Slots()
+        obj.x = 1
+        compiled_fn(torch.ones(1), obj)
+        self.assertFalse(hasattr(obj, "x"))
+
+    def test_hasattr_on_slotted_object(self):
+        # hasattr inside compiled code reflects actual slot state
+        def fn(t):
+            obj = Slots()
+            before = hasattr(obj, "x")  # False — slot not set
+            obj.x = 5
+            after = hasattr(obj, "x")  # True — slot is now set
+            return t + before + after
+
+        dynamo_testing.standard_test(self, fn, nargs=1)
+
+    def test_slots_with_property_setter(self):
+        # property setter is called instead of writing directly to the slot
+        def fn(t):
+            obj = SlotsAndProperty(1)
+            obj.x = 5  # calls setter: _x = 5 * 2 = 10
+            return t + obj.x  # calls getter: returns _x = 10
+
+        dynamo_testing.standard_test(self, fn, nargs=1)
+
+    def test_slot_assignment_no_recompile_same_type(self):
+        # Calling compiled fn repeatedly with the same slotted object type
+        # must not trigger recompilation
+        cnts = dynamo_testing.CompileCounter()
+
+        def fn(t, obj):
+            obj.x = 10
+            return t + obj.x
+
+        compiled_fn = torch.compile(fn, backend=cnts)
+        t = torch.ones(1)
+        compiled_fn(t, Slots())
+        compiled_fn(t, Slots())
+        compiled_fn(t, Slots())
+        self.assertEqual(cnts.frame_count, 1)
+
+    def test_slot_assignment_recompiles_on_type_change(self):
+        # Compiled fn sees slot assigned to int first, then float — guards recompile
+        cnts = dynamo_testing.CompileCounter()
+
+        def fn(t, a, obj):
+            obj.x = a
+            return t + obj.x
+
+        compiled_fn = torch.compile(fn, backend=cnts)
+        t = torch.ones(1)
+
+        compiled_fn(t, 1, Slots())
+        compiled_fn(t, 1, Slots())
+        self.assertEqual(cnts.frame_count, 1)  # same type, no recompile
+
+        x = t.clone()
+        res = compiled_fn(x, 1.0, Slots())
+        self.assertEqual(cnts.frame_count, 2)  # float instead of int — recompile
+        self.assertEqual(res, fn(x, 1.0, Slots()))
+
+
+class WithGetattribute:
+    # __slots__ = ("x", "_side_effects")
+
+    def __init__(self, x):
+        object.__setattr__(self, "x", x)
+        object.__setattr__(self, "_side_effects", set())
+
+    def __getattribute__(self, name):
+        effects = object.__getattribute__(self, "_side_effects")
+        effects.add(name)
+        return object.__getattribute__(self, name)
+
+
+class TestSlotsFromCPython(TestCase):
+    """Slot tests extracted from CPython's test_descr.py::test_slots."""
+
+    def setUp(self):
+        super().setUp()
+        self._u_prev = torch._dynamo.config.enable_trace_unittest
+        torch._dynamo.config.enable_trace_unittest = True
+
+    def tearDown(self):
+        super().tearDown()
+        torch._dynamo.config.enable_trace_unittest = self._u_prev
+
+    def test_slots_empty(self):
+        class C:
+            __slots__ = []
+
+        def fn(t):
+            x = C()
+            self.assertFalse(hasattr(x, "__dict__"))
+            self.assertFalse(hasattr(x, "foo"))
+            return t.sin()
+
+        dynamo_testing.standard_test(self, fn, nargs=1)
+
+    def test_slots_single(self):
+        class C:
+            __slots__ = ["a"]
+
+        def fn(t):
+            x = C()
+            self.assertFalse(hasattr(x, "__dict__"))
+            self.assertFalse(hasattr(x, "a"))
+            x.a = 1
+            self.assertEqual(x.a, 1)
+            x.a = None
+            self.assertEqual(x.a, None)
+            del x.a
+            self.assertFalse(hasattr(x, "a"))
+            return t.sin()
+
+        dynamo_testing.standard_test(self, fn, nargs=1)
+
+    def test_slots_multiple(self):
+        class C:
+            __slots__ = ["a", "b", "c"]
+
+        def fn(t):
+            x = C()
+            self.assertFalse(hasattr(x, "__dict__"))
+            self.assertFalse(hasattr(x, "a"))
+            self.assertFalse(hasattr(x, "b"))
+            self.assertFalse(hasattr(x, "c"))
+            x.a = 1
+            x.b = 2
+            x.c = 3
+            self.assertEqual(x.a, 1)
+            self.assertEqual(x.b, 2)
+            self.assertEqual(x.c, 3)
+            return t.sin()
+
+        dynamo_testing.standard_test(self, fn, nargs=1)
+
+    def test_slots_name_mangling(self):
+        class C:
+            __slots__ = ["__a"]
+
+            def __init__(self, value):
+                self.__a = value
+
+            def get(self):
+                return self.__a
+
+        def fn(t):
+            x = C(5)
+            self.assertFalse(hasattr(x, "__dict__"))
+            self.assertFalse(hasattr(x, "__a"))
+            self.assertEqual(x.get(), 5)
+            try:
+                x.__a = 6
+            except AttributeError:
+                pass
+            else:
+                self.fail("Double underscored names not mangled")
+            return t.sin()
+
+        dynamo_testing.standard_test(self, fn, nargs=1)
+
+    def test_slots_string_not_expanded(self):
+        # A single string is not expanded as a sequence
+        class C:
+            __slots__ = "abc"  # noqa: PLC0205
+
+        def fn(t):
+            c = C()
+            c.abc = 5
+            self.assertEqual(c.abc, 5)
+            return t.sin()
+
+        dynamo_testing.standard_test(self, fn, nargs=1)
+
+    def test_slots_tuple(self):
+        slots = ("foo", "bar")
+
+        class C:
+            __slots__ = slots
+
+        def fn(t):
+            x = C()
+            x.foo = 5
+            self.assertEqual(x.foo, 5)
+            self.assertIs(type(slots[0]), str)
+            return t.sin()
+
+        dynamo_testing.standard_test(self, fn, nargs=1)
+
+    def test_slots_get_unset_raises(self):
+        class X:
+            __slots__ = "a"  # noqa: PLC0205
+
+        def fn(t):
+            with self.assertRaises(AttributeError):
+                X().a
+            return t.sin()
+
+        dynamo_testing.standard_test(self, fn, nargs=1)
+
+    def test_slots_str_subclass(self):
+        # gh-98783: string subclass in __slots__
+        class SubStr(str):  # noqa: SLOT000
+            pass
+
+        class X:
+            __slots__ = (SubStr("x"),)
+
+        def fn(t):
+            X().x = 1
+            with self.assertRaises(AttributeError):
+                X().a
+            return t.sin()
+
+        dynamo_testing.standard_test(self, fn, nargs=1)
+
+    def test_slots_special_dict(self):
+        # __dict__ in __slots__ enables arbitrary attr assignment
+        class D:
+            __slots__ = ["__dict__"]
+
+        def fn(t):
+            a = D()
+            self.assertTrue(hasattr(a, "__dict__"))
+            self.assertFalse(hasattr(a, "__weakref__"))
+            a.foo = 42
+            self.assertEqual(a.__dict__, {"foo": 42})
+            return t.sin()
+
+        dynamo_testing.standard_test(self, fn, nargs=1)
+
+    def test_slots_special_weakref(self):
+        # __weakref__ in __slots__ — no __dict__, arbitrary attr raises
+        class W:
+            __slots__ = ["__weakref__"]
+
+        def fn(t):
+            a = W()
+            self.assertTrue(hasattr(a, "__weakref__"))
+            self.assertFalse(hasattr(a, "__dict__"))
+            with self.assertRaises(AttributeError):
+                a.foo = 42
+            return t.sin()
+
+        dynamo_testing.standard_test(self, fn, nargs=1)
+
+    def test_slots_special_inherit_dict_weakref(self):
+        # Inheriting from both __dict__ and __weakref__ slot classes
+        class D:
+            __slots__ = ["__dict__"]
+
+        class W:
+            __slots__ = ["__weakref__"]
+
+        class C1(W, D):
+            __slots__ = []
+
+        def fn(t):
+            a = C1()
+            self.assertTrue(hasattr(a, "__dict__"))
+            self.assertTrue(hasattr(a, "__weakref__"))
+            a.foo = 42
+            self.assertEqual(a.__dict__, {"foo": 42})
+            return t.sin()
+
+        dynamo_testing.standard_test(self, fn, nargs=1)
+
+    @unittest.expectedFailure
+    def test_slots_special2_classcell(self):
+        # Testing __classcell__ in __slots__
+        class Meta(type):
+            def __new__(metacls, name, bases, namespace, attr):
+                self.assertIn(attr, namespace)
+                return super().__new__(metacls, name, bases, namespace)
+
+        class C1:
+            def __init__(self):
+                self.b = 42
+
+        class C2(C1, metaclass=Meta, attr="__classcell__"):
+            __slots__ = ["__classcell__"]
+
+            def __init__(self):
+                super().__init__()
+
+        def fn(t):
+            self.assertIsInstance(
+                C2.__dict__["__classcell__"], types.MemberDescriptorType
+            )
+            c = C2()
+            self.assertEqual(c.b, 42)
+            self.assertFalse(hasattr(c, "__classcell__"))
+            c.__classcell__ = 42
+            self.assertEqual(c.__classcell__, 42)
+            with self.assertRaises(TypeError):
+
+                class C3:
+                    __classcell__ = 42
+                    __slots__ = ["__classcell__"]
+
+            return t.sin()
+
+        dynamo_testing.standard_test(self, fn, nargs=1)
+
+    def test_slots_multiple_inheritance(self):
+        # SF bug 575229: multiple inheritance w/ slots dumps core
+        class A:
+            __slots__ = ()
+
+        class B:
+            pass
+
+        class C(A, B):
+            __slots__ = ()
+
+        def fn(t):
+            self.assertTrue(hasattr(C, "__dict__"))
+            self.assertTrue(hasattr(C, "__weakref__"))
+            C().x = 2
+            return t.sin()
+
+        dynamo_testing.standard_test(self, fn, nargs=1)
+
+
+class TestUserDefinedClassDict(TestCase):
+    def test_class_dict_read(self):
+        class MyClass:
+            x = 3
+
+        def fn(t):
+            t = t + MyClass.__dict__["x"]
+            t = t + MyClass.__dict__.get("x", 0)
+            t = t + MyClass.__dict__.get("z", 99)
+            t = t + (1 if "x" in MyClass.__dict__ else 0)
+            t = t + (1 if "z" in MyClass.__dict__ else 0)
+            return t
+
+        dynamo_testing.standard_test(self, fn, nargs=1)
+
+    def test_class_dict_via_arg(self):
+        class MyClass:
+            x = 7
+
+        def fn(t, cls):
+            return t + cls.__dict__.get("x", 0)
+
+        cnt = dynamo_testing.CompileCounter()
+        compiled = torch.compile(fn, backend=cnt)
+        result = compiled(torch.tensor([0.0]), MyClass)
+        self.assertEqual(result, torch.tensor([7.0]))
+
+    def test_class_dict_mutation_recompiles(self):
+        # Mutating a class attribute between calls should trigger recompilation,
+        # and the compiled function should see the updated value.
+        class MyClass:
+            x = 1
+
+        def fn(t):
+            return t + MyClass.__dict__["x"]
+
+        cnt = dynamo_testing.CompileCounter()
+        compiled = torch.compile(fn, backend=cnt)
+
+        result1 = compiled(torch.tensor([0.0]))
+        self.assertEqual(result1, torch.tensor([1.0]))
+        self.assertEqual(cnt.frame_count, 1)
+
+        MyClass.x = 10
+        result2 = compiled(torch.tensor([0.0]))
+        self.assertEqual(result2, torch.tensor([10.0]))
+        # Should have recompiled due to guard failure
+        self.assertEqual(cnt.frame_count, 2)
+
+    def test_class_dict_add_key_recompiles(self):
+        # Adding a new attribute to the class should trigger recompilation
+        # when the compiled code checks for key presence.
+        class MyClass:
+            x = 1
+
+        def fn(t):
+            return t + (1 if "y" in MyClass.__dict__ else 0)
+
+        cnt = dynamo_testing.CompileCounter()
+        compiled = torch.compile(fn, backend=cnt)
+
+        result1 = compiled(torch.tensor([0.0]))
+        self.assertEqual(result1, torch.tensor([0.0]))
+        self.assertEqual(cnt.frame_count, 1)
+
+        MyClass.y = 99
+        result2 = compiled(torch.tensor([0.0]))
+        self.assertEqual(result2, torch.tensor([1.0]))
+        # Should have recompiled
+        self.assertEqual(cnt.frame_count, 2)
+
+    def test_class_dict_delete_key_recompiles(self):
+        # Deleting a class attribute should trigger recompilation.
+        class MyClass:
+            x = 5
+            y = 10
+
+        def fn(t):
+            return t + MyClass.__dict__.get("y", 0)
+
+        cnt = dynamo_testing.CompileCounter()
+        compiled = torch.compile(fn, backend=cnt)
+
+        result1 = compiled(torch.tensor([0.0]))
+        self.assertEqual(result1, torch.tensor([10.0]))
+        self.assertEqual(cnt.frame_count, 1)
+
+        del MyClass.y
+        result2 = compiled(torch.tensor([0.0]))
+        self.assertEqual(result2, torch.tensor([0.0]))
+        # Should have recompiled
+        self.assertEqual(cnt.frame_count, 2)
+
+
+if __name__ == "__main__":
+    run_tests()
diff --git a/test/dynamo/test_utils.py b/test/dynamo/test_utils.py
index 22c3d26207ef1..2d4c0f7dc13f2 100644
--- a/test/dynamo/test_utils.py
+++ b/test/dynamo/test_utils.py
@@ -1,5 +1,6 @@
 # Owner(s): ["module: dynamo"]
 import dataclasses
+import json
 import os
 import pprint
 import sys
@@ -84,7 +85,6 @@ def test_larger_multiplier_for_even_smaller_tensor(self):
     @dynamo_config.patch(
         {
             "log_compilation_metrics": True,
-            "inline_inbuilt_nn_modules": False,
         }
     )
     def test_graph_break_counting(self):
@@ -102,18 +102,18 @@ def run_forward_backward():
             loss = loss_fn(output, target)
             loss.backward()
 
-        @torch.compile
+        @torch.compile(backend="eager")
         def add(x, y):
             return x + y
 
-        @torch.compile
+        @torch.compile(backend="eager")
         def break_it(x):
             y = x.sum()
             if y > 0:
                 return x + y.item()
             return x - y.item()
 
-        @torch.compile
+        @torch.compile(backend="eager")
         def break_it2(x):
             y = x.sum()
             if y > 0:
@@ -194,7 +194,8 @@ def fn(x):
         @torch.compile(backend=my_backend)
         def fn(x):
             z = x + 1
-            y = break_it(z)
+            with torch._dynamo.disable_nested_graph_breaks():
+                y = break_it(z)
             return y * 2
 
         x = torch.randn(3)
@@ -298,6 +299,14 @@ def test_reinplace_counters_use_trigger_name_not_enum_value(self):
             "Should not use enum value (integer) in key, should use trigger.name instead",
         )
 
+    def test_get_dynamo_config_for_logging_ignores_logging_functions(self):
+        with dynamo_config.patch(ignore_logging_functions={print}):
+            result = utils._get_dynamo_config_for_logging()
+            parsed = json.loads(result)
+
+        self.assertIsInstance(parsed, dict)
+        self.assertNotIn("ignore_logging_functions", parsed)
+
 
 class TestModel(torch.nn.Module):
     def __init__(self):
@@ -330,7 +339,7 @@ def run_forward_backward(self):
     def warmup(self):
         # Helper to make sure any process-global lru_caches (e.g., torch_key())
         # have already executed. Just compile something.
-        @torch.compile
+        @torch.compile(backend="inductor")
         def add(x, y):
             return x + y
 
@@ -445,7 +454,6 @@ def backward(grad_output):
     @dynamo_config.patch(
         {
             "log_compilation_metrics": True,
-            "inline_inbuilt_nn_modules": False,
         }
     )
     @inductor_config.patch(
@@ -688,7 +696,7 @@ def filter_expected(s: str) -> str:
  'compile_time_autotune_time_us': None,
  'compiler_config': None,
  'compliant_custom_ops': set(),
- 'config_inline_inbuilt_nn_modules': False,
+ 'config_inline_inbuilt_nn_modules': True,
  'config_suppress_errors': False,
  'cuda_version': None,
  'cudagraph_skip_reason': None,
@@ -707,11 +715,11 @@ def filter_expected(s: str) -> str:
  'fail_user_frame_lineno': None,
  'frame_key': '1',
  'gc_time_us': 0,
- 'graph_input_count': 1,
- 'graph_node_count': 3,
+ 'graph_input_count': 3,
+ 'graph_node_count': 5,
  'graph_node_shapes': None,
  'graph_op_count': 1,
- 'guard_count': 10,
+ 'guard_count': 31,
  'has_guarded_code': True,
  'inductor_code_gen_cumulative_compile_time_us': 0,
  'inductor_compile_time_s': 0.0,
@@ -781,7 +789,7 @@ def filter_expected(s: str) -> str:
  'compile_time_autotune_time_us': None,
  'compiler_config': None,
  'compliant_custom_ops': set(),
- 'config_inline_inbuilt_nn_modules': False,
+ 'config_inline_inbuilt_nn_modules': True,
  'config_suppress_errors': False,
  'cuda_version': None,
  'cudagraph_skip_reason': None,
@@ -800,11 +808,11 @@ def filter_expected(s: str) -> str:
  'fail_user_frame_lineno': None,
  'frame_key': '1',
  'gc_time_us': 0,
- 'graph_input_count': 1,
- 'graph_node_count': 3,
+ 'graph_input_count': 3,
+ 'graph_node_count': 5,
  'graph_node_shapes': None,
  'graph_op_count': 1,
- 'guard_count': 10,
+ 'guard_count': 31,
  'has_guarded_code': True,
  'inductor_code_gen_cumulative_compile_time_us': 0,
  'inductor_compile_time_s': 0.0,
@@ -888,7 +896,7 @@ def filter_expected(s: str) -> str:
  'compile_time_autotune_time_us': None,
  'compiler_config': None,
  'compliant_custom_ops': None,
- 'config_inline_inbuilt_nn_modules': False,
+ 'config_inline_inbuilt_nn_modules': True,
  'config_suppress_errors': False,
  'cuda_version': None,
  'cudagraph_skip_reason': None,
@@ -981,7 +989,7 @@ def filter_expected(s: str) -> str:
  'compile_time_autotune_time_us': None,
  'compiler_config': None,
  'compliant_custom_ops': None,
- 'config_inline_inbuilt_nn_modules': False,
+ 'config_inline_inbuilt_nn_modules': True,
  'config_suppress_errors': False,
  'cuda_version': None,
  'cudagraph_skip_reason': None,
diff --git a/test/dynamo/test_wrap_inductor_compiled_regions.py b/test/dynamo/test_wrap_inductor_compiled_regions.py
index 69cce71ec7e88..0f2d335adfe30 100644
--- a/test/dynamo/test_wrap_inductor_compiled_regions.py
+++ b/test/dynamo/test_wrap_inductor_compiled_regions.py
@@ -38,7 +38,8 @@ def match_rng_op(node, op):
         return False
 
     if op is not None:
-        assert not isinstance(op, list)
+        if isinstance(op, list):
+            raise AssertionError("Expected op to not be a list")
         ops = [op]
     if freq is not None:
         freqs = [freq]
@@ -51,17 +52,20 @@ def match_rng_op(node, op):
                 if match_rng_op(node, op) or node.target == op:
                     actual_count += 1
             err_msg = f"In graph {gm}, expected {op} to have occurred {freq} times in the graph, but got {actual_count}."
-            assert actual_count == freq, err_msg
+            if actual_count != freq:
+                raise AssertionError(err_msg)
     else:
-        assert freqs_ge is not None
+        if freqs_ge is None:
+            raise AssertionError("Expected freqs_ge to not be None")
         for op, freq_ge in zip(ops, freqs_ge):
             actual_count = 0
             for node in gm.graph.nodes:
                 if match_rng_op(node, op) or node.target == op:
                     actual_count += 1
-            assert actual_count >= freq_ge, (
-                f"In graph {gm}, expected {op} to have occurred at least {freq_ge} times in the graph, but got {actual_count}."
-            )
+            if actual_count < freq_ge:
+                raise AssertionError(
+                    f"In graph {gm}, expected {op} to have occurred at least {freq_ge} times in the graph, but got {actual_count}."
+                )
     return gm
 
 
@@ -95,6 +99,33 @@ def fn(x, y):
         expected = torch.matmul(x, y)
         self.assertEqual(result, expected)
 
+    @requires_cuda_and_triton
+    def test_wrap_name_visible_in_debug_mode(self):
+        """Test that named compiled regions surface their name in DebugMode"""
+
+        @torch.compile(
+            backend="inductor",
+            options={"wrap_inductor_compiled_regions": True},
+            fullgraph=True,
+            name="flex_attention",
+        )
+        def fn(x, y):
+            return torch.matmul(x, y)
+
+        x = torch.randn(4, 4, device="cuda")
+        y = torch.randn(4, 4, device="cuda")
+
+        with DebugMode() as debug_mode:
+            result = fn(x, y)
+
+        debug_string = debug_mode.debug_string()
+
+        self.assertIn("inductor_compiled_code", debug_string)
+        self.assertIn("name=flex_attention", debug_string)
+
+        expected = torch.matmul(x, y)
+        self.assertEqual(result, expected)
+
     @requires_cuda_and_triton
     def test_wrap_disabled_not_visible_in_debug_mode(self):
         """Test that compiled regions are not wrapped when option is disabled"""
@@ -918,6 +949,60 @@ def checkpointed_fn(x, y):
         self.assertEqual(x.grad, x_eager.grad)
         self.assertEqual(y.grad, y_eager.grad)
 
+    @requires_cuda_and_triton
+    def test_sac_outer_compile_inner_name_visible_to_policy(self):
+        """Test that SAC policies can inspect torch.compile region names"""
+
+        @torch.compile(
+            backend="inductor",
+            options={"wrap_inductor_compiled_regions": True},
+            fullgraph=True,
+            name="flex_attention",
+        )
+        def inner_compiled_matmul(x, y):
+            return torch.matmul(x, y)
+
+        seen_region_names = []
+
+        def policy_fn(ctx, op, *args, **kwargs):
+            from torch._higher_order_ops.wrap import inductor_compiled_code
+
+            if op == inductor_compiled_code:
+                seen_region_names.append(kwargs.get("name"))
+            return CheckpointPolicy.PREFER_RECOMPUTE
+
+        def checkpointed_fn(x, y):
+            a = inner_compiled_matmul(x, y)
+            return torch.relu(a)
+
+        x = torch.randn(4, 4, device="cuda", requires_grad=True)
+        y = torch.randn(4, 4, device="cuda", requires_grad=True)
+
+        x_eager = x.detach().clone().requires_grad_(True)
+        y_eager = y.detach().clone().requires_grad_(True)
+
+        context_fn = functools.partial(create_selective_checkpoint_contexts, policy_fn)
+
+        output = checkpoint(
+            checkpointed_fn,
+            x,
+            y,
+            use_reentrant=False,
+            context_fn=context_fn,
+        )
+        loss = output.sum()
+        loss.backward()
+
+        a_eager = torch.matmul(x_eager, y_eager)
+        b_eager = torch.relu(a_eager)
+        loss_eager = b_eager.sum()
+        loss_eager.backward()
+
+        self.assertIn("flex_attention", seen_region_names)
+        self.assertEqual(output, b_eager)
+        self.assertEqual(x.grad, x_eager.grad)
+        self.assertEqual(y.grad, y_eager.grad)
+
     @requires_cuda_and_triton
     def test_wrap_no_dispatch_mode_no_hop_invoked(self):
         """
@@ -1064,20 +1149,75 @@ def checkpointed_flex_fn(q, k, v):
         torch.testing.assert_close(k.grad, k2.grad, rtol=1e-3, atol=1e-3)
         torch.testing.assert_close(v.grad, v2.grad, rtol=1e-3, atol=1e-3)
 
-    def test_fake_tensor_mode_raises_error(self):
-        """Test that running compiled code inside FakeTensorMode raises a clear error"""
+    def test_fake_tensor_mode_works(self):
+        """Test that running compiled code inside FakeTensorMode works with FX graph fallback"""
         from torch._subclasses.fake_tensor import FakeTensorMode
 
         with FakeTensorMode():
             model = torch.nn.Linear(4, 4)
             inp = torch.rand(4, 4)
 
+            # The FX graph is now serialized using SerializedGraphModule, so it
+            # survives cache serialization/deserialization and works with caching
             with inductor_config.patch({"wrap_inductor_compiled_regions": True}):
-                with self.assertRaisesRegex(
-                    RuntimeError,
-                    "Inductor compiled code cannot be run with FakeTensor inputs",
-                ):
-                    torch.compile(model)(inp)
+                # This should now work - the inductor_compiled_code HOP will
+                # use the stored FX graph to propagate fake tensors
+                result = torch.compile(model)(inp)
+                # Verify the result has the expected shape
+                self.assertEqual(result.shape, (4, 4))
+
+    def test_proxy_tensor_mode_works(self):
+        """Test that running compiled code inside ProxyTensorMode works with FX graph fallback"""
+        from torch._higher_order_ops.wrap import (
+            inductor_code_side_table,
+            inductor_compiled_code,
+            InductorCompiledCallable,
+        )
+        from torch.fx.experimental.proxy_tensor import make_fx
+
+        # Reset the side table for a clean test
+        inductor_code_side_table.reset_table()
+
+        # Create a simple FX graph
+        class SimpleModel(torch.nn.Module):
+            def forward(self, x):
+                return x + 1
+
+        model = SimpleModel()
+        gm = torch.fx.symbolic_trace(model)
+
+        # Create an InductorCompiledCallable
+        # The compiled callable should match FX graph's output convention
+        def simple_compiled(inputs):
+            return inputs[0] + 1
+
+        callable_obj = InductorCompiledCallable(simple_compiled, gm)
+
+        # Wrapper that uses the HOP
+        def wrapper(x):
+            return inductor_compiled_code(callable_obj, [x])
+
+        inp = torch.randn(4, 4)
+        traced = make_fx(wrapper)(inp)
+
+        # Verify the traced graph contains the inductor_compiled_code HOP
+        hop_found = False
+        for node in traced.graph.nodes:
+            if node.op == "call_function" and "inductor_compiled_code" in str(
+                node.target
+            ):
+                hop_found = True
+                # Verify the callable index is an int
+                self.assertIsInstance(node.args[0], int)
+                break
+        self.assertTrue(
+            hop_found, "inductor_compiled_code HOP not found in traced graph"
+        )
+
+        # Verify the traced graph can be executed and produces correct results
+        result = traced(inp)
+        expected = inp + 1
+        torch.testing.assert_close(result, expected)
 
 
 if __name__ == "__main__":
diff --git a/test/dynamo_expected_failures/CPython313-test_baseexception-ExceptionClassTests.test_interface_multi_arg b/test/dynamo_expected_failures/CPython313-test_class-ClassTests.testBadTypeReturned
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_baseexception-ExceptionClassTests.test_interface_multi_arg
rename to test/dynamo_expected_failures/CPython313-test_class-ClassTests.testBadTypeReturned
diff --git a/test/dynamo_expected_failures/CPython313-test_baseexception-ExceptionClassTests.test_interface_no_arg b/test/dynamo_expected_failures/CPython313-test_class-ClassTests.testBinaryOps
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_baseexception-ExceptionClassTests.test_interface_no_arg
rename to test/dynamo_expected_failures/CPython313-test_class-ClassTests.testBinaryOps
diff --git a/test/dynamo_expected_failures/CPython313-test_baseexception-ExceptionClassTests.test_interface_single_arg b/test/dynamo_expected_failures/CPython313-test_class-ClassTests.testClassCallRecursionLimit
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_baseexception-ExceptionClassTests.test_interface_single_arg
rename to test/dynamo_expected_failures/CPython313-test_class-ClassTests.testClassCallRecursionLimit
diff --git a/test/dynamo_expected_failures/CPython313-test_baseexception-UsageTests.test_raise_new_style_non_exception b/test/dynamo_expected_failures/CPython313-test_class-ClassTests.testConstructorErrorMessages
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_baseexception-UsageTests.test_raise_new_style_non_exception
rename to test/dynamo_expected_failures/CPython313-test_class-ClassTests.testConstructorErrorMessages
diff --git a/test/dynamo_expected_failures/CPython313-test_baseexception-UsageTests.test_raise_string b/test/dynamo_expected_failures/CPython313-test_class-ClassTests.testDel
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_baseexception-UsageTests.test_raise_string
rename to test/dynamo_expected_failures/CPython313-test_class-ClassTests.testDel
diff --git a/test/dynamo_expected_failures/CPython313-test_bool-BoolTest.test_bool_new b/test/dynamo_expected_failures/CPython313-test_class-ClassTests.testForExceptionsRaisedInInstanceGetattr2
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_bool-BoolTest.test_bool_new
rename to test/dynamo_expected_failures/CPython313-test_class-ClassTests.testForExceptionsRaisedInInstanceGetattr2
diff --git a/test/dynamo_expected_failures/CPython313-test_bool-BoolTest.test_from_bytes b/test/dynamo_expected_failures/CPython313-test_class-ClassTests.testHasAttrString
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_bool-BoolTest.test_from_bytes
rename to test/dynamo_expected_failures/CPython313-test_class-ClassTests.testHasAttrString
diff --git a/test/dynamo_expected_failures/CPython313-test_collections-TestCounter.test_conversions b/test/dynamo_expected_failures/CPython313-test_class-ClassTests.testHashComparisonOfMethods
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_collections-TestCounter.test_conversions
rename to test/dynamo_expected_failures/CPython313-test_class-ClassTests.testHashComparisonOfMethods
diff --git a/test/dynamo_expected_failures/CPython313-test_collections-TestCounter.test_copy_subclass b/test/dynamo_expected_failures/CPython313-test_class-ClassTests.testHashStuff
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_collections-TestCounter.test_copy_subclass
rename to test/dynamo_expected_failures/CPython313-test_class-ClassTests.testHashStuff
diff --git a/test/dynamo_expected_failures/CPython313-test_collections-TestCounter.test_order_preservation b/test/dynamo_expected_failures/CPython313-test_class-ClassTests.testMetaclassCallOptimization
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_collections-TestCounter.test_order_preservation
rename to test/dynamo_expected_failures/CPython313-test_class-ClassTests.testMetaclassCallOptimization
diff --git a/test/dynamo_expected_failures/CPython313-test_collections-TestCounter.test_repr_nonsortable b/test/dynamo_expected_failures/CPython313-test_class-ClassTests.testMisc
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_collections-TestCounter.test_repr_nonsortable
rename to test/dynamo_expected_failures/CPython313-test_class-ClassTests.testMisc
diff --git a/test/dynamo_expected_failures/CPython313-test_collections-TestNamedTuple.test_name_conflicts b/test/dynamo_expected_failures/CPython313-test_class-ClassTests.testObjectAttributeAccessErrorMessages
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_collections-TestNamedTuple.test_name_conflicts
rename to test/dynamo_expected_failures/CPython313-test_class-ClassTests.testObjectAttributeAccessErrorMessages
diff --git a/test/dynamo_expected_failures/CPython313-test_collections-TestNamedTuple.test_odd_sizes b/test/dynamo_expected_failures/CPython313-test_class-ClassTests.testPredefinedAttrs
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_collections-TestNamedTuple.test_odd_sizes
rename to test/dynamo_expected_failures/CPython313-test_class-ClassTests.testPredefinedAttrs
diff --git a/test/dynamo_expected_failures/CPython313-test_collections-TestUserObjects.test_dict_copy b/test/dynamo_expected_failures/CPython313-test_class-ClassTests.testSFBug532646
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_collections-TestUserObjects.test_dict_copy
rename to test/dynamo_expected_failures/CPython313-test_class-ClassTests.testSFBug532646
diff --git a/test/dynamo_expected_failures/CPython313-test_collections-TestUserObjects.test_list_copy b/test/dynamo_expected_failures/CPython313-test_class-ClassTests.testSetattrNonStringName
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_collections-TestUserObjects.test_list_copy
rename to test/dynamo_expected_failures/CPython313-test_class-ClassTests.testSetattrNonStringName
diff --git a/test/dynamo_expected_failures/CPython313-test_contextlib-ContextManagerTestCase.test_contextmanager_attribs b/test/dynamo_expected_failures/CPython313-test_class-ClassTests.testSetattrWrapperNameIntern
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_contextlib-ContextManagerTestCase.test_contextmanager_attribs
rename to test/dynamo_expected_failures/CPython313-test_class-ClassTests.testSetattrWrapperNameIntern
diff --git a/test/dynamo_expected_failures/CPython313-test_contextlib-ContextManagerTestCase.test_contextmanager_doc_attrib b/test/dynamo_expected_failures/CPython313-test_class-ClassTests.testTypeAttributeAccessErrorMessages
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_contextlib-ContextManagerTestCase.test_contextmanager_doc_attrib
rename to test/dynamo_expected_failures/CPython313-test_class-ClassTests.testTypeAttributeAccessErrorMessages
diff --git a/test/dynamo_expected_failures/CPython313-test_contextlib-ContextManagerTestCase.test_instance_docstring_given_cm_docstring b/test/dynamo_expected_failures/CPython313-test_class-ClassTests.testUnaryOps
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_contextlib-ContextManagerTestCase.test_instance_docstring_given_cm_docstring
rename to test/dynamo_expected_failures/CPython313-test_class-ClassTests.testUnaryOps
diff --git a/test/dynamo_expected_failures/CPython313-test_contextlib-TestAbstractContextManager.test_slots b/test/dynamo_expected_failures/CPython313-test_class-ClassTests.test_specialization_class_call_doesnt_crash
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_contextlib-TestAbstractContextManager.test_slots
rename to test/dynamo_expected_failures/CPython313-test_class-ClassTests.test_specialization_class_call_doesnt_crash
diff --git a/test/dynamo_expected_failures/CPython313-test_contextlib-TestExitStack.test_dont_reraise_RuntimeError b/test/dynamo_expected_failures/CPython313-test_descr-AAAPTypesLongInitTest.test_pytype_long_ready
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_contextlib-TestExitStack.test_dont_reraise_RuntimeError
rename to test/dynamo_expected_failures/CPython313-test_descr-AAAPTypesLongInitTest.test_pytype_long_ready
diff --git a/test/dynamo_expected_failures/CPython313-test_contextlib-TextExitStack.test_exit_exception_with_existing_context b/test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_abstractmethods
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_contextlib-TextExitStack.test_exit_exception_with_existing_context
rename to test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_abstractmethods
diff --git a/test/dynamo_expected_failures/CPython313-test_dict-DictTest.test_dictview_mixed_set_operations b/test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_altmro
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_dict-DictTest.test_dictview_mixed_set_operations
rename to test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_altmro
diff --git a/test/dynamo_expected_failures/CPython313-test_dict-DictTest.test_dictview_set_operations_on_items b/test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_attr_raise_through_property
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_dict-DictTest.test_dictview_set_operations_on_items
rename to test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_attr_raise_through_property
diff --git a/test/dynamo_expected_failures/CPython313-test_dict-DictTest.test_instance_dict_getattr_str_subclass b/test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_basic_inheritance
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_dict-DictTest.test_instance_dict_getattr_str_subclass
rename to test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_basic_inheritance
diff --git a/test/dynamo_expected_failures/CPython313-test_dict-DictTest.test_keys_contained b/test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_binary_operator_override
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_dict-DictTest.test_keys_contained
rename to test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_binary_operator_override
diff --git a/test/dynamo_expected_failures/CPython313-test_dict-DictTest.test_missing b/test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_bound_method_repr
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_dict-DictTest.test_missing
rename to test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_bound_method_repr
diff --git a/test/dynamo_expected_failures/CPython313-test_dict-DictTest.test_reverse_iterator_for_shared_shared_dicts b/test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_bpo25750
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_dict-DictTest.test_reverse_iterator_for_shared_shared_dicts
rename to test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_bpo25750
diff --git a/test/dynamo_expected_failures/CPython313-test_dict-DictTest.test_splittable_update b/test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_buffer_inheritance
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_dict-DictTest.test_splittable_update
rename to test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_buffer_inheritance
diff --git a/test/dynamo_expected_failures/CPython313-test_dict-DictTest.test_tuple_keyerror b/test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_builtin_bases
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_dict-DictTest.test_tuple_keyerror
rename to test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_builtin_bases
diff --git a/test/dynamo_expected_failures/CPython313-test_dict-SubclassMappingTests.test_bool b/test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_builtin_function_or_method
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_dict-SubclassMappingTests.test_bool
rename to test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_builtin_function_or_method
diff --git a/test/dynamo_expected_failures/CPython313-test_dict-SubclassMappingTests.test_read b/test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_carloverre
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_dict-SubclassMappingTests.test_read
rename to test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_carloverre
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-MiscTestCase.test__all__ b/test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_carloverre_multi_inherit_invalid
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-MiscTestCase.test__all__
rename to test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_carloverre_multi_inherit_invalid
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-OldTestFlag.test_and b/test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_carloverre_multi_inherit_valid
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-OldTestFlag.test_and
rename to test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_carloverre_multi_inherit_valid
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-OldTestFlag.test_bool b/test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_classic
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-OldTestFlag.test_bool
rename to test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_classic
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-OldTestFlag.test_iter b/test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_classmethods
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-OldTestFlag.test_iter
rename to test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_classmethods
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-OldTestFlag.test_member_contains b/test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_classmethods_in_c
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-OldTestFlag.test_member_contains
rename to test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_classmethods_in_c
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-OldTestFlag.test_member_iter b/test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_compattr
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-OldTestFlag.test_member_iter
rename to test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_compattr
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-OldTestFlag.test_or b/test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_copy_setstate
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-OldTestFlag.test_or
rename to test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_copy_setstate
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-OldTestFlag.test_programatic_function_from_dict b/test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_cycle_through_dict
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-OldTestFlag.test_programatic_function_from_dict
rename to test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_cycle_through_dict
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-OldTestFlag.test_programatic_function_iterable b/test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_delete_hook
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-OldTestFlag.test_programatic_function_iterable
rename to test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_delete_hook
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-OldTestFlag.test_programatic_function_string b/test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_dict_constructors
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-OldTestFlag.test_programatic_function_string
rename to test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_dict_constructors
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-OldTestFlag.test_programatic_function_string_list b/test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_dir
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-OldTestFlag.test_programatic_function_string_list
rename to test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_dir
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-OldTestFlag.test_programatic_function_string_with_start b/test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_dynamics
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-OldTestFlag.test_programatic_function_string_with_start
rename to test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_dynamics
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-OldTestFlag.test_xor b/test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_evil_type_name
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-OldTestFlag.test_xor
rename to test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_evil_type_name
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-OldTestIntFlag.test_and b/test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_file_fault
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-OldTestIntFlag.test_and
rename to test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_file_fault
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-OldTestIntFlag.test_bool b/test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_getattr_hooks
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-OldTestIntFlag.test_bool
rename to test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_getattr_hooks
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-OldTestIntFlag.test_global_repr_conform1 b/test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_gh55664
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-OldTestIntFlag.test_global_repr_conform1
rename to test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_gh55664
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-OldTestIntFlag.test_global_repr_keep b/test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_hash_inheritance
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-OldTestIntFlag.test_global_repr_keep
rename to test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_hash_inheritance
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-OldTestIntFlag.test_invert b/test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_init
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-OldTestIntFlag.test_invert
rename to test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_init
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-OldTestIntFlag.test_iter b/test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_instance_method_get_behavior
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-OldTestIntFlag.test_iter
rename to test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_instance_method_get_behavior
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-OldTestIntFlag.test_member_iter b/test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_ipow_returns_not_implemented
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-OldTestIntFlag.test_member_iter
rename to test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_ipow_returns_not_implemented
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-OldTestIntFlag.test_or b/test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_keyword_arguments
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-OldTestIntFlag.test_or
rename to test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_keyword_arguments
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-OldTestIntFlag.test_programatic_function_from_dict b/test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_keywords
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-OldTestIntFlag.test_programatic_function_from_dict
rename to test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_keywords
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-OldTestIntFlag.test_programatic_function_from_empty_list b/test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_load_attr_extended_arg
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-OldTestIntFlag.test_programatic_function_from_empty_list
rename to test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_load_attr_extended_arg
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-OldTestIntFlag.test_programatic_function_from_empty_tuple b/test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_metaclass
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-OldTestIntFlag.test_programatic_function_from_empty_tuple
rename to test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_metaclass
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-OldTestIntFlag.test_programatic_function_iterable b/test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_meth_class_get
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-OldTestIntFlag.test_programatic_function_iterable
rename to test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_meth_class_get
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-OldTestIntFlag.test_programatic_function_string b/test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_method_wrapper
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-OldTestIntFlag.test_programatic_function_string
rename to test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_method_wrapper
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-OldTestIntFlag.test_programatic_function_string_list b/test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_methods
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-OldTestIntFlag.test_programatic_function_string_list
rename to test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_methods
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-OldTestIntFlag.test_programatic_function_string_with_start b/test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_methods_in_c
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-OldTestIntFlag.test_programatic_function_string_with_start
rename to test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_methods_in_c
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-OldTestIntFlag.test_type b/test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_mixing_slot_wrappers
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-OldTestIntFlag.test_type
rename to test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_mixing_slot_wrappers
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-OldTestIntFlag.test_xor b/test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_module_subclasses
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-OldTestIntFlag.test_xor
rename to test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_module_subclasses
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestConvert.test_convert_raise b/test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_mutable_bases
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestConvert.test_convert_raise
rename to test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_mutable_bases
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestEmptyAndNonLatinStrings.test_non_latin_character_string b/test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_mutable_bases_catch_mro_conflict
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestEmptyAndNonLatinStrings.test_non_latin_character_string
rename to test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_mutable_bases_catch_mro_conflict
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestEmptyAndNonLatinStrings.test_non_latin_number_string b/test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_mutable_bases_with_failing_mro
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestEmptyAndNonLatinStrings.test_non_latin_number_string
rename to test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_mutable_bases_with_failing_mro
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestIntEnumClass.test_copy b/test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_mutable_names
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestIntEnumClass.test_copy
rename to test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_mutable_names
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestIntEnumClass.test_dir_on_class b/test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_no_ipow
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestIntEnumClass.test_dir_on_class
rename to test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_no_ipow
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestIntEnumClass.test_format b/test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_not_implemented
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestIntEnumClass.test_format
rename to test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_not_implemented
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestIntEnumClass.test_gnv_is_static b/test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_object_class
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestIntEnumClass.test_gnv_is_static
rename to test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_object_class
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestIntEnumClass.test_programmatic_function_from_dict b/test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_object_class_assignment_between_heaptypes_and_nonheaptypes
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestIntEnumClass.test_programmatic_function_from_dict
rename to test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_object_class_assignment_between_heaptypes_and_nonheaptypes
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestIntEnumClass.test_programmatic_function_iterable b/test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_object_new
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestIntEnumClass.test_programmatic_function_iterable
rename to test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_object_new
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestIntEnumClass.test_programmatic_function_string b/test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_object_new_and_init_with_parameters
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestIntEnumClass.test_programmatic_function_string
rename to test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_object_new_and_init_with_parameters
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestIntEnumClass.test_programmatic_function_string_list b/test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_overloading
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestIntEnumClass.test_programmatic_function_string_list
rename to test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_overloading
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestIntEnumClass.test_repr b/test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_properties
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestIntEnumClass.test_repr
rename to test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_properties
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestIntEnumClass.test_str b/test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_properties_plus
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestIntEnumClass.test_str
rename to test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_properties_plus
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestIntEnumFunction.test_copy b/test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_proxy_call
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestIntEnumFunction.test_copy
rename to test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_proxy_call
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestIntEnumFunction.test_dir_on_class b/test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_proxy_super
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestIntEnumFunction.test_dir_on_class
rename to test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_proxy_super
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestIntEnumFunction.test_format b/test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_python_lists
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestIntEnumFunction.test_format
rename to test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_python_lists
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestIntEnumFunction.test_gnv_is_static b/test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_qualname
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestIntEnumFunction.test_gnv_is_static
rename to test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_qualname
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestIntEnumFunction.test_programmatic_function_from_dict b/test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_recursive_call
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestIntEnumFunction.test_programmatic_function_from_dict
rename to test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_recursive_call
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestIntEnumFunction.test_programmatic_function_iterable b/test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_refleaks_in_classmethod___init__
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestIntEnumFunction.test_programmatic_function_iterable
rename to test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_refleaks_in_classmethod___init__
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestIntEnumFunction.test_programmatic_function_string b/test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_refleaks_in_staticmethod___init__
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestIntEnumFunction.test_programmatic_function_string
rename to test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_refleaks_in_staticmethod___init__
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestIntEnumFunction.test_programmatic_function_string_list b/test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_remove_subclass
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestIntEnumFunction.test_programmatic_function_string_list
rename to test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_remove_subclass
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestIntEnumFunction.test_repr b/test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_repr_as_str
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestIntEnumFunction.test_repr
rename to test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_repr_as_str
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestIntEnumFunction.test_shadowed_attr b/test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_repr_with_module_str_subclass
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestIntEnumFunction.test_shadowed_attr
rename to test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_repr_with_module_str_subclass
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestIntEnumFunction.test_str b/test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_rich_comparisons
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestIntEnumFunction.test_str
rename to test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_rich_comparisons
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestIntFlagClass.test_copy b/test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_set_and_no_get
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestIntFlagClass.test_copy
rename to test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_set_and_no_get
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestIntFlagClass.test_dir_on_class b/test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_set_class
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestIntFlagClass.test_dir_on_class
rename to test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_set_class
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestIntFlagClass.test_format b/test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_set_dict
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestIntFlagClass.test_format
rename to test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_set_dict
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestIntFlagClass.test_gnv_is_static b/test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_set_doc
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestIntFlagClass.test_gnv_is_static
rename to test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_set_doc
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestIntFlagClass.test_programmatic_function_from_dict b/test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_slices
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestIntFlagClass.test_programmatic_function_from_dict
rename to test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_slices
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestIntFlagClass.test_programmatic_function_iterable b/test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_slot_shadows_class_variable
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestIntFlagClass.test_programmatic_function_iterable
rename to test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_slot_shadows_class_variable
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestIntFlagClass.test_programmatic_function_string b/test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_slots_descriptor
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestIntFlagClass.test_programmatic_function_string
rename to test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_slots_descriptor
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestIntFlagClass.test_programmatic_function_string_list b/test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_slots_special2
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestIntFlagClass.test_programmatic_function_string_list
rename to test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_slots_special2
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestIntFlagClass.test_repr b/test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_slots_trash
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestIntFlagClass.test_repr
rename to test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_slots_trash
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestIntFlagClass.test_str b/test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_special_method_lookup
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestIntFlagClass.test_str
rename to test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_special_method_lookup
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestIntFlagFunction.test_copy b/test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_special_unbound_method_types
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestIntFlagFunction.test_copy
rename to test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_special_unbound_method_types
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestIntFlagFunction.test_dir_on_class b/test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_specialized_method_calls_check_types
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestIntFlagFunction.test_dir_on_class
rename to test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_specialized_method_calls_check_types
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestIntFlagFunction.test_format b/test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_specials
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestIntFlagFunction.test_format
rename to test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_specials
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestIntFlagFunction.test_gnv_is_static b/test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_staticmethods
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestIntFlagFunction.test_gnv_is_static
rename to test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_staticmethods
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestIntFlagFunction.test_programmatic_function_from_dict b/test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_staticmethods_in_c
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestIntFlagFunction.test_programmatic_function_from_dict
rename to test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_staticmethods_in_c
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestIntFlagFunction.test_programmatic_function_iterable b/test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_str_of_str_subclass
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestIntFlagFunction.test_programmatic_function_iterable
rename to test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_str_of_str_subclass
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestIntFlagFunction.test_programmatic_function_string b/test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_str_operations
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestIntFlagFunction.test_programmatic_function_string
rename to test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_str_operations
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestIntFlagFunction.test_programmatic_function_string_list b/test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_str_subclass_as_dict_key
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestIntFlagFunction.test_programmatic_function_string_list
rename to test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_str_subclass_as_dict_key
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestIntFlagFunction.test_repr b/test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_subclass_propagation
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestIntFlagFunction.test_repr
rename to test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_subclass_propagation
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestIntFlagFunction.test_str b/test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_subclass_right_op
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestIntFlagFunction.test_str
rename to test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_subclass_right_op
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestMinimalDateClass.test_copy b/test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_subtype_resurrection
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestMinimalDateClass.test_copy
rename to test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_subtype_resurrection
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestMinimalDateClass.test_dir_on_class b/test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_supers
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestMinimalDateClass.test_dir_on_class
rename to test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_supers
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestMinimalDateClass.test_gnv_is_static b/test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_type___getattribute__
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestMinimalDateClass.test_gnv_is_static
rename to test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_type___getattribute__
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestMinimalDateFunction.test_copy b/test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_uninitialized_modules
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestMinimalDateFunction.test_copy
rename to test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_uninitialized_modules
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestMinimalDateFunction.test_dir_on_class b/test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_vicious_descriptor_nonsense
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestMinimalDateFunction.test_dir_on_class
rename to test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_vicious_descriptor_nonsense
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestMinimalDateFunction.test_gnv_is_static b/test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_weakref_segfault
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestMinimalDateFunction.test_gnv_is_static
rename to test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_weakref_segfault
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestMinimalFloatClass.test_copy b/test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_weakrefs
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestMinimalFloatClass.test_copy
rename to test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_weakrefs
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestMinimalFloatClass.test_dir_on_class b/test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_wrapper_segfault
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestMinimalFloatClass.test_dir_on_class
rename to test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_wrapper_segfault
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestMinimalFloatClass.test_format b/test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_wrong_class_slot_wrapper
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestMinimalFloatClass.test_format
rename to test/dynamo_expected_failures/CPython313-test_descr-ClassPropertiesAndMethods.test_wrong_class_slot_wrapper
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestMinimalFloatClass.test_gnv_is_static b/test/dynamo_expected_failures/CPython313-test_descr-DictProxyTests.test_repr
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestMinimalFloatClass.test_gnv_is_static
rename to test/dynamo_expected_failures/CPython313-test_descr-DictProxyTests.test_repr
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestMinimalFloatClass.test_programmatic_function_from_dict b/test/dynamo_expected_failures/CPython313-test_descr-MiscTests.test_type_lookup_mro_reference
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestMinimalFloatClass.test_programmatic_function_from_dict
rename to test/dynamo_expected_failures/CPython313-test_descr-MiscTests.test_type_lookup_mro_reference
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestMinimalFloatClass.test_programmatic_function_iterable b/test/dynamo_expected_failures/CPython313-test_descr-MroTest.test_disappearing_custom_mro
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestMinimalFloatClass.test_programmatic_function_iterable
rename to test/dynamo_expected_failures/CPython313-test_descr-MroTest.test_disappearing_custom_mro
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestMinimalFloatClass.test_programmatic_function_string b/test/dynamo_expected_failures/CPython313-test_descr-MroTest.test_incomplete_extend
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestMinimalFloatClass.test_programmatic_function_string
rename to test/dynamo_expected_failures/CPython313-test_descr-MroTest.test_incomplete_extend
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestMinimalFloatClass.test_programmatic_function_string_list b/test/dynamo_expected_failures/CPython313-test_descr-MroTest.test_incomplete_set_bases_on_self
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestMinimalFloatClass.test_programmatic_function_string_list
rename to test/dynamo_expected_failures/CPython313-test_descr-MroTest.test_incomplete_set_bases_on_self
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestMinimalFloatClass.test_repr b/test/dynamo_expected_failures/CPython313-test_descr-MroTest.test_incomplete_super
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestMinimalFloatClass.test_repr
rename to test/dynamo_expected_failures/CPython313-test_descr-MroTest.test_incomplete_super
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestMinimalFloatClass.test_str b/test/dynamo_expected_failures/CPython313-test_descr-MroTest.test_reent_set_bases_on_base
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestMinimalFloatClass.test_str
rename to test/dynamo_expected_failures/CPython313-test_descr-MroTest.test_reent_set_bases_on_base
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestMinimalFloatFunction.test_copy b/test/dynamo_expected_failures/CPython313-test_descr-MroTest.test_reent_set_bases_on_direct_base
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestMinimalFloatFunction.test_copy
rename to test/dynamo_expected_failures/CPython313-test_descr-MroTest.test_reent_set_bases_on_direct_base
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestMinimalFloatFunction.test_dir_on_class b/test/dynamo_expected_failures/CPython313-test_descr-MroTest.test_reent_set_bases_tp_base_cycle
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestMinimalFloatFunction.test_dir_on_class
rename to test/dynamo_expected_failures/CPython313-test_descr-MroTest.test_reent_set_bases_tp_base_cycle
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestMinimalFloatFunction.test_format b/test/dynamo_expected_failures/CPython313-test_descr-MroTest.test_tp_subclasses_cycle_error_return_path
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestMinimalFloatFunction.test_format
rename to test/dynamo_expected_failures/CPython313-test_descr-MroTest.test_tp_subclasses_cycle_error_return_path
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestMinimalFloatFunction.test_gnv_is_static b/test/dynamo_expected_failures/CPython313-test_descr-MroTest.test_tp_subclasses_cycle_in_update_slots
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestMinimalFloatFunction.test_gnv_is_static
rename to test/dynamo_expected_failures/CPython313-test_descr-MroTest.test_tp_subclasses_cycle_in_update_slots
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestMinimalFloatFunction.test_programmatic_function_from_dict b/test/dynamo_expected_failures/CPython313-test_descr-OperatorsTest.test_complexes
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestMinimalFloatFunction.test_programmatic_function_from_dict
rename to test/dynamo_expected_failures/CPython313-test_descr-OperatorsTest.test_complexes
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestMinimalFloatFunction.test_programmatic_function_iterable b/test/dynamo_expected_failures/CPython313-test_descr-OperatorsTest.test_dicts
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestMinimalFloatFunction.test_programmatic_function_iterable
rename to test/dynamo_expected_failures/CPython313-test_descr-OperatorsTest.test_dicts
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestMinimalFloatFunction.test_programmatic_function_string b/test/dynamo_expected_failures/CPython313-test_descr-OperatorsTest.test_explicit_reverse_methods
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestMinimalFloatFunction.test_programmatic_function_string
rename to test/dynamo_expected_failures/CPython313-test_descr-OperatorsTest.test_explicit_reverse_methods
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestMinimalFloatFunction.test_programmatic_function_string_list b/test/dynamo_expected_failures/CPython313-test_descr-OperatorsTest.test_floats
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestMinimalFloatFunction.test_programmatic_function_string_list
rename to test/dynamo_expected_failures/CPython313-test_descr-OperatorsTest.test_floats
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestMinimalFloatFunction.test_repr b/test/dynamo_expected_failures/CPython313-test_descr-OperatorsTest.test_ints
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestMinimalFloatFunction.test_repr
rename to test/dynamo_expected_failures/CPython313-test_descr-OperatorsTest.test_ints
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestMinimalFloatFunction.test_str b/test/dynamo_expected_failures/CPython313-test_descr-OperatorsTest.test_lists
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestMinimalFloatFunction.test_str
rename to test/dynamo_expected_failures/CPython313-test_descr-OperatorsTest.test_lists
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestMixedDateClass.test_dir_on_class b/test/dynamo_expected_failures/CPython313-test_descr-OperatorsTest.test_spam_dicts
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestMixedDateClass.test_dir_on_class
rename to test/dynamo_expected_failures/CPython313-test_descr-OperatorsTest.test_spam_dicts
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestMixedDateClass.test_format b/test/dynamo_expected_failures/CPython313-test_descr-OperatorsTest.test_spam_lists
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestMixedDateClass.test_format
rename to test/dynamo_expected_failures/CPython313-test_descr-OperatorsTest.test_spam_lists
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestMixedDateClass.test_gnv_is_static b/test/dynamo_expected_failures/CPython313-test_descr-PicklingTests.test_issue24097
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestMixedDateClass.test_gnv_is_static
rename to test/dynamo_expected_failures/CPython313-test_descr-PicklingTests.test_issue24097
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestMixedDateClass.test_str b/test/dynamo_expected_failures/CPython313-test_descr-PicklingTests.test_object_reduce
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestMixedDateClass.test_str
rename to test/dynamo_expected_failures/CPython313-test_descr-PicklingTests.test_object_reduce
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestMixedDateFunction.test_dir_on_class b/test/dynamo_expected_failures/CPython313-test_descr-PicklingTests.test_pickle_slots
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestMixedDateFunction.test_dir_on_class
rename to test/dynamo_expected_failures/CPython313-test_descr-PicklingTests.test_pickle_slots
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestMixedDateFunction.test_format b/test/dynamo_expected_failures/CPython313-test_descr-PicklingTests.test_reduce
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestMixedDateFunction.test_format
rename to test/dynamo_expected_failures/CPython313-test_descr-PicklingTests.test_reduce
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestMixedDateFunction.test_gnv_is_static b/test/dynamo_expected_failures/CPython313-test_descr-PicklingTests.test_reduce_copying
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestMixedDateFunction.test_gnv_is_static
rename to test/dynamo_expected_failures/CPython313-test_descr-PicklingTests.test_reduce_copying
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestMixedDateFunction.test_str b/test/dynamo_expected_failures/CPython313-test_descr-PicklingTests.test_special_method_lookup
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestMixedDateFunction.test_str
rename to test/dynamo_expected_failures/CPython313-test_descr-PicklingTests.test_special_method_lookup
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestMixedFloatClass.test_dir_on_class b/test/dynamo_expected_failures/CPython313-test_descr-SharedKeyTests.test_subclasses
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestMixedFloatClass.test_dir_on_class
rename to test/dynamo_expected_failures/CPython313-test_descr-SharedKeyTests.test_subclasses
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestMixedFloatClass.test_format b/test/dynamo_expected_failures/CPython313-test_dict-DictTest.test_splittable_del
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestMixedFloatClass.test_format
rename to test/dynamo_expected_failures/CPython313-test_dict-DictTest.test_splittable_del
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestMixedFloatClass.test_gnv_is_static b/test/dynamo_expected_failures/CPython313-test_dict-DictTest.test_splittable_pop
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestMixedFloatClass.test_gnv_is_static
rename to test/dynamo_expected_failures/CPython313-test_dict-DictTest.test_splittable_pop
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestMixedFloatClass.test_programmatic_function_from_dict b/test/dynamo_expected_failures/CPython313-test_dict-DictTest.test_splittable_popitem
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestMixedFloatClass.test_programmatic_function_from_dict
rename to test/dynamo_expected_failures/CPython313-test_dict-DictTest.test_splittable_popitem
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestMixedFloatClass.test_programmatic_function_iterable b/test/dynamo_expected_failures/CPython313-test_dict-DictTest.test_splittable_setdefault
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestMixedFloatClass.test_programmatic_function_iterable
rename to test/dynamo_expected_failures/CPython313-test_dict-DictTest.test_splittable_setdefault
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestMixedFloatClass.test_programmatic_function_string b/test/dynamo_expected_failures/CPython313-test_enum-OldTestIntFlag.test_contains_tf
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestMixedFloatClass.test_programmatic_function_string
rename to test/dynamo_expected_failures/CPython313-test_enum-OldTestIntFlag.test_contains_tf
diff --git a/test/dynamo_expected_failures/CPython313-test_exceptions-ExceptionTests.test_raise_does_not_create_context_chain_cycle b/test/dynamo_expected_failures/CPython313-test_exceptions-ExceptionTests.test_raise_does_not_create_context_chain_cycle
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/test/dynamo_expected_failures/CPython313-test_exceptions-ExceptionTests.test_unicode_error_str_does_not_crash b/test/dynamo_expected_failures/CPython313-test_exceptions-ExceptionTests.test_unicode_error_str_does_not_crash
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/test/dynamo_expected_failures/CPython313-test_float-GeneralFloatCases.test_error_message b/test/dynamo_expected_failures/CPython313-test_float-GeneralFloatCases.test_error_message
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestMixedFloatClass.test_programmatic_function_string_list b/test/dynamo_expected_failures/CPython313-test_funcattrs-ArbitraryFunctionAttrTest.test_delete_unknown_attr
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestMixedFloatClass.test_programmatic_function_string_list
rename to test/dynamo_expected_failures/CPython313-test_funcattrs-ArbitraryFunctionAttrTest.test_delete_unknown_attr
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestMixedFloatClass.test_repr b/test/dynamo_expected_failures/CPython313-test_funcattrs-ArbitraryFunctionAttrTest.test_set_attr
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestMixedFloatClass.test_repr
rename to test/dynamo_expected_failures/CPython313-test_funcattrs-ArbitraryFunctionAttrTest.test_set_attr
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestMixedFloatClass.test_str b/test/dynamo_expected_failures/CPython313-test_funcattrs-BuiltinFunctionPropertiesTest.test_builtin__qualname__
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestMixedFloatClass.test_str
rename to test/dynamo_expected_failures/CPython313-test_funcattrs-BuiltinFunctionPropertiesTest.test_builtin__qualname__
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestMixedFloatFunction.test_dir_on_class b/test/dynamo_expected_failures/CPython313-test_funcattrs-BuiltinFunctionPropertiesTest.test_builtin__self__
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestMixedFloatFunction.test_dir_on_class
rename to test/dynamo_expected_failures/CPython313-test_funcattrs-BuiltinFunctionPropertiesTest.test_builtin__self__
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestMixedFloatFunction.test_format b/test/dynamo_expected_failures/CPython313-test_funcattrs-CellTest.test_comparison
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestMixedFloatFunction.test_format
rename to test/dynamo_expected_failures/CPython313-test_funcattrs-CellTest.test_comparison
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestMixedFloatFunction.test_gnv_is_static b/test/dynamo_expected_failures/CPython313-test_funcattrs-FunctionDictsTest.test_delete___dict__
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestMixedFloatFunction.test_gnv_is_static
rename to test/dynamo_expected_failures/CPython313-test_funcattrs-FunctionDictsTest.test_delete___dict__
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestMixedFloatFunction.test_programmatic_function_from_dict b/test/dynamo_expected_failures/CPython313-test_funcattrs-FunctionDictsTest.test_setting_dict_to_invalid
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestMixedFloatFunction.test_programmatic_function_from_dict
rename to test/dynamo_expected_failures/CPython313-test_funcattrs-FunctionDictsTest.test_setting_dict_to_invalid
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestMixedFloatFunction.test_programmatic_function_iterable b/test/dynamo_expected_failures/CPython313-test_funcattrs-FunctionDictsTest.test_setting_dict_to_valid
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestMixedFloatFunction.test_programmatic_function_iterable
rename to test/dynamo_expected_failures/CPython313-test_funcattrs-FunctionDictsTest.test_setting_dict_to_valid
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestMixedFloatFunction.test_programmatic_function_string b/test/dynamo_expected_failures/CPython313-test_funcattrs-FunctionDocstringTest.test_delete_docstring
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestMixedFloatFunction.test_programmatic_function_string
rename to test/dynamo_expected_failures/CPython313-test_funcattrs-FunctionDocstringTest.test_delete_docstring
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestMixedFloatFunction.test_programmatic_function_string_list b/test/dynamo_expected_failures/CPython313-test_funcattrs-FunctionDocstringTest.test_set_docstring_attr
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestMixedFloatFunction.test_programmatic_function_string_list
rename to test/dynamo_expected_failures/CPython313-test_funcattrs-FunctionDocstringTest.test_set_docstring_attr
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestMixedFloatFunction.test_repr b/test/dynamo_expected_failures/CPython313-test_funcattrs-FunctionPropertiesTest.test___builtins__
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestMixedFloatFunction.test_repr
rename to test/dynamo_expected_failures/CPython313-test_funcattrs-FunctionPropertiesTest.test___builtins__
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestMixedFloatFunction.test_str b/test/dynamo_expected_failures/CPython313-test_funcattrs-FunctionPropertiesTest.test___closure__
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestMixedFloatFunction.test_str
rename to test/dynamo_expected_failures/CPython313-test_funcattrs-FunctionPropertiesTest.test___closure__
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestMixedIntClass.test_dir_on_class b/test/dynamo_expected_failures/CPython313-test_funcattrs-FunctionPropertiesTest.test___code__
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestMixedIntClass.test_dir_on_class
rename to test/dynamo_expected_failures/CPython313-test_funcattrs-FunctionPropertiesTest.test___code__
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestMixedIntClass.test_format b/test/dynamo_expected_failures/CPython313-test_funcattrs-FunctionPropertiesTest.test___globals__
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestMixedIntClass.test_format
rename to test/dynamo_expected_failures/CPython313-test_funcattrs-FunctionPropertiesTest.test___globals__
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestMixedIntClass.test_gnv_is_static b/test/dynamo_expected_failures/CPython313-test_funcattrs-FunctionPropertiesTest.test___name__
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestMixedIntClass.test_gnv_is_static
rename to test/dynamo_expected_failures/CPython313-test_funcattrs-FunctionPropertiesTest.test___name__
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestMixedIntClass.test_programmatic_function_from_dict b/test/dynamo_expected_failures/CPython313-test_funcattrs-FunctionPropertiesTest.test___qualname__
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestMixedIntClass.test_programmatic_function_from_dict
rename to test/dynamo_expected_failures/CPython313-test_funcattrs-FunctionPropertiesTest.test___qualname__
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestMixedIntClass.test_programmatic_function_iterable b/test/dynamo_expected_failures/CPython313-test_funcattrs-FunctionPropertiesTest.test___type_params__
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestMixedIntClass.test_programmatic_function_iterable
rename to test/dynamo_expected_failures/CPython313-test_funcattrs-FunctionPropertiesTest.test___type_params__
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestMixedIntClass.test_programmatic_function_string b/test/dynamo_expected_failures/CPython313-test_funcattrs-FunctionPropertiesTest.test_blank_func_defaults
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestMixedIntClass.test_programmatic_function_string
rename to test/dynamo_expected_failures/CPython313-test_funcattrs-FunctionPropertiesTest.test_blank_func_defaults
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestMixedIntClass.test_programmatic_function_string_list b/test/dynamo_expected_failures/CPython313-test_funcattrs-FunctionPropertiesTest.test_cell_new
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestMixedIntClass.test_programmatic_function_string_list
rename to test/dynamo_expected_failures/CPython313-test_funcattrs-FunctionPropertiesTest.test_cell_new
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestMixedIntClass.test_repr b/test/dynamo_expected_failures/CPython313-test_funcattrs-FunctionPropertiesTest.test_copying___code__
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestMixedIntClass.test_repr
rename to test/dynamo_expected_failures/CPython313-test_funcattrs-FunctionPropertiesTest.test_copying___code__
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestMixedIntClass.test_str b/test/dynamo_expected_failures/CPython313-test_funcattrs-FunctionPropertiesTest.test_dir_includes_correct_attrs
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestMixedIntClass.test_str
rename to test/dynamo_expected_failures/CPython313-test_funcattrs-FunctionPropertiesTest.test_dir_includes_correct_attrs
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestMixedIntFlagClass.test_dir_on_class b/test/dynamo_expected_failures/CPython313-test_funcattrs-FunctionPropertiesTest.test_empty_cell
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestMixedIntFlagClass.test_dir_on_class
rename to test/dynamo_expected_failures/CPython313-test_funcattrs-FunctionPropertiesTest.test_empty_cell
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestMixedIntFlagClass.test_format b/test/dynamo_expected_failures/CPython313-test_funcattrs-FunctionPropertiesTest.test_func_default_args
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestMixedIntFlagClass.test_format
rename to test/dynamo_expected_failures/CPython313-test_funcattrs-FunctionPropertiesTest.test_func_default_args
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestMixedIntFlagClass.test_gnv_is_static b/test/dynamo_expected_failures/CPython313-test_funcattrs-FunctionPropertiesTest.test_invalid___code___assignment
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestMixedIntFlagClass.test_gnv_is_static
rename to test/dynamo_expected_failures/CPython313-test_funcattrs-FunctionPropertiesTest.test_invalid___code___assignment
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestMixedIntFlagClass.test_programmatic_function_from_dict b/test/dynamo_expected_failures/CPython313-test_funcattrs-FunctionPropertiesTest.test_set_cell
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestMixedIntFlagClass.test_programmatic_function_from_dict
rename to test/dynamo_expected_failures/CPython313-test_funcattrs-FunctionPropertiesTest.test_set_cell
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestMixedIntFlagClass.test_programmatic_function_iterable b/test/dynamo_expected_failures/CPython313-test_funcattrs-InstancemethodAttrTest.test___class__
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestMixedIntFlagClass.test_programmatic_function_iterable
rename to test/dynamo_expected_failures/CPython313-test_funcattrs-InstancemethodAttrTest.test___class__
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestMixedIntFlagClass.test_programmatic_function_string b/test/dynamo_expected_failures/CPython313-test_funcattrs-InstancemethodAttrTest.test___func__
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestMixedIntFlagClass.test_programmatic_function_string
rename to test/dynamo_expected_failures/CPython313-test_funcattrs-InstancemethodAttrTest.test___func__
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestMixedIntFlagClass.test_programmatic_function_string_list b/test/dynamo_expected_failures/CPython313-test_funcattrs-InstancemethodAttrTest.test___func___non_method
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestMixedIntFlagClass.test_programmatic_function_string_list
rename to test/dynamo_expected_failures/CPython313-test_funcattrs-InstancemethodAttrTest.test___func___non_method
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestMixedIntFlagClass.test_repr b/test/dynamo_expected_failures/CPython313-test_funcattrs-InstancemethodAttrTest.test___self__
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestMixedIntFlagClass.test_repr
rename to test/dynamo_expected_failures/CPython313-test_funcattrs-InstancemethodAttrTest.test___self__
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestMixedIntFlagClass.test_str b/test/dynamo_expected_failures/CPython313-test_funcattrs-StaticMethodAttrsTest.test_func_attribute
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestMixedIntFlagClass.test_str
rename to test/dynamo_expected_failures/CPython313-test_funcattrs-StaticMethodAttrsTest.test_func_attribute
diff --git a/test/dynamo_expected_failures/CPython313-test_functools-TestPartialC.test_arg_combinations b/test/dynamo_expected_failures/CPython313-test_functools-TestPartialC.test_arg_combinations
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/test/dynamo_expected_failures/CPython313-test_functools-TestPartialC.test_attributes b/test/dynamo_expected_failures/CPython313-test_functools-TestPartialC.test_attributes
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/test/dynamo_expected_failures/CPython313-test_functools-TestPartialC.test_basic_examples b/test/dynamo_expected_failures/CPython313-test_functools-TestPartialC.test_basic_examples
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/test/dynamo_expected_failures/CPython313-test_functools-TestPartialC.test_error_propagation b/test/dynamo_expected_failures/CPython313-test_functools-TestPartialC.test_error_propagation
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/test/dynamo_expected_failures/CPython313-test_functools-TestPartialC.test_keyword b/test/dynamo_expected_failures/CPython313-test_functools-TestPartialC.test_keyword
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/test/dynamo_expected_failures/CPython313-test_functools-TestPartialC.test_kw_combinations b/test/dynamo_expected_failures/CPython313-test_functools-TestPartialC.test_kw_combinations
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/test/dynamo_expected_failures/CPython313-test_functools-TestPartialC.test_kwargs_copy b/test/dynamo_expected_failures/CPython313-test_functools-TestPartialC.test_kwargs_copy
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/test/dynamo_expected_failures/CPython313-test_functools-TestPartialC.test_nested_optimization_bug b/test/dynamo_expected_failures/CPython313-test_functools-TestPartialC.test_nested_optimization_bug
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/test/dynamo_expected_failures/CPython313-test_functools-TestPartialC.test_no_side_effects b/test/dynamo_expected_failures/CPython313-test_functools-TestPartialC.test_no_side_effects
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/test/dynamo_expected_failures/CPython313-test_functools-TestPartialC.test_positional b/test/dynamo_expected_failures/CPython313-test_functools-TestPartialC.test_positional
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/test/dynamo_expected_failures/CPython313-test_functools-TestPartialC.test_protection_of_callers_dict_argument b/test/dynamo_expected_failures/CPython313-test_functools-TestPartialC.test_protection_of_callers_dict_argument
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/test/dynamo_expected_failures/CPython313-test_functools-TestPartialC.test_with_bound_and_unbound_methods b/test/dynamo_expected_failures/CPython313-test_functools-TestPartialC.test_with_bound_and_unbound_methods
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/test/dynamo_expected_failures/CPython313-test_functools-TestPartialMethod.test_repr b/test/dynamo_expected_failures/CPython313-test_functools-TestPartialMethod.test_repr
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/test/dynamo_expected_failures/CPython313-test_functools-TestPartialPy.test_nested_optimization b/test/dynamo_expected_failures/CPython313-test_functools-TestPartialPy.test_nested_optimization
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/test/dynamo_expected_failures/CPython313-test_functools-TestPartialPySubclass.test_nested_optimization b/test/dynamo_expected_failures/CPython313-test_functools-TestPartialPySubclass.test_nested_optimization
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/test/dynamo_expected_failures/CPython313-test_functools-TestSingleDispatch.test_assignment_behavior b/test/dynamo_expected_failures/CPython313-test_functools-TestSingleDispatch.test_assignment_behavior
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/test/dynamo_expected_failures/CPython313-test_functools-TestSingleDispatch.test_c3_abc b/test/dynamo_expected_failures/CPython313-test_functools-TestSingleDispatch.test_c3_abc
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/test/dynamo_expected_failures/CPython313-test_functools-TestSingleDispatch.test_c_classes b/test/dynamo_expected_failures/CPython313-test_functools-TestSingleDispatch.test_c_classes
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/test/dynamo_expected_failures/CPython313-test_functools-TestSingleDispatch.test_mro b/test/dynamo_expected_failures/CPython313-test_functools-TestSingleDispatch.test_mro
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/test/dynamo_expected_failures/CPython313-test_functools-TestSingleDispatch.test_mro_conflicts b/test/dynamo_expected_failures/CPython313-test_functools-TestSingleDispatch.test_mro_conflicts
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/test/dynamo_expected_failures/CPython313-test_functools-TestSingleDispatch.test_wrapping_attributes b/test/dynamo_expected_failures/CPython313-test_functools-TestSingleDispatch.test_wrapping_attributes
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/test/dynamo_expected_failures/CPython313-test_functools-TestTotalOrdering.test_total_ordering_no_overwrite b/test/dynamo_expected_failures/CPython313-test_functools-TestTotalOrdering.test_total_ordering_no_overwrite
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/test/dynamo_expected_failures/CPython313-test_functools-TestUpdateWrapper.test_default_update b/test/dynamo_expected_failures/CPython313-test_functools-TestUpdateWrapper.test_default_update
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/test/dynamo_expected_failures/CPython313-test_functools-TestUpdateWrapper.test_default_update_doc b/test/dynamo_expected_failures/CPython313-test_functools-TestUpdateWrapper.test_default_update_doc
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/test/dynamo_expected_failures/CPython313-test_functools-TestUpdateWrapper.test_missing_attributes b/test/dynamo_expected_failures/CPython313-test_functools-TestUpdateWrapper.test_missing_attributes
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/test/dynamo_expected_failures/CPython313-test_functools-TestWraps.test_default_update b/test/dynamo_expected_failures/CPython313-test_functools-TestWraps.test_default_update
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/test/dynamo_expected_failures/CPython313-test_functools-TestWraps.test_default_update_doc b/test/dynamo_expected_failures/CPython313-test_functools-TestWraps.test_default_update_doc
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/test/dynamo_expected_failures/CPython313-test_functools-TestWraps.test_missing_attributes b/test/dynamo_expected_failures/CPython313-test_functools-TestWraps.test_missing_attributes
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/test/dynamo_expected_failures/CPython313-test_heapq-TestErrorHandlingPython.test_iterable_args b/test/dynamo_expected_failures/CPython313-test_heapq-TestErrorHandlingPython.test_iterable_args
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/test/dynamo_expected_failures/CPython313-test_heapq-TestModules.test_c_functions b/test/dynamo_expected_failures/CPython313-test_heapq-TestModules.test_c_functions
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/test/dynamo_expected_failures/CPython313-test_int-IntStrDigitLimitsTests.test_disabled_limit b/test/dynamo_expected_failures/CPython313-test_int-IntStrDigitLimitsTests.test_disabled_limit
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/test/dynamo_expected_failures/CPython313-test_int-IntSubclassStrDigitLimitsTests.test_disabled_limit b/test/dynamo_expected_failures/CPython313-test_int-IntSubclassStrDigitLimitsTests.test_disabled_limit
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/test/dynamo_expected_failures/CPython313-test_int-IntSubclassStrDigitLimitsTests.test_max_str_digits_edge_cases b/test/dynamo_expected_failures/CPython313-test_int-IntSubclassStrDigitLimitsTests.test_max_str_digits_edge_cases
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/test/dynamo_expected_failures/CPython313-test_int-IntSubclassStrDigitLimitsTests.test_power_of_two_bases_unlimited b/test/dynamo_expected_failures/CPython313-test_int-IntSubclassStrDigitLimitsTests.test_power_of_two_bases_unlimited
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/test/dynamo_expected_failures/CPython313-test_int-IntSubclassStrDigitLimitsTests.test_sign_not_counted b/test/dynamo_expected_failures/CPython313-test_int-IntSubclassStrDigitLimitsTests.test_sign_not_counted
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/test/dynamo_expected_failures/CPython313-test_int-IntTestCases.test_int_subclass_with_index b/test/dynamo_expected_failures/CPython313-test_int-IntTestCases.test_int_subclass_with_index
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/test/dynamo_expected_failures/CPython313-test_int-IntTestCases.test_intconversion b/test/dynamo_expected_failures/CPython313-test_int-IntTestCases.test_intconversion
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/test/dynamo_expected_failures/CPython313-test_itertools-TestBasicOps.test_dropwhile b/test/dynamo_expected_failures/CPython313-test_itertools-TestBasicOps.test_dropwhile
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/test/dynamo_expected_failures/CPython313-test_itertools-TestBasicOps.test_starmap b/test/dynamo_expected_failures/CPython313-test_itertools-TestBasicOps.test_starmap
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/test/dynamo_expected_failures/CPython313-test_itertools-TestBasicOps.test_takewhile b/test/dynamo_expected_failures/CPython313-test_itertools-TestBasicOps.test_takewhile
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/test/dynamo_expected_failures/CPython313-test_itertools-TestExamples.test_stapmap b/test/dynamo_expected_failures/CPython313-test_itertools-TestExamples.test_stapmap
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/test/dynamo_expected_failures/CPython313-test_itertools-TestExamples.test_takewhile b/test/dynamo_expected_failures/CPython313-test_itertools-TestExamples.test_takewhile
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/test/dynamo_expected_failures/CPython313-test_itertools-TestGC.test_issue2246 b/test/dynamo_expected_failures/CPython313-test_itertools-TestGC.test_issue2246
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/test/dynamo_expected_failures/CPython313-test_itertools-TestGC.test_starmap b/test/dynamo_expected_failures/CPython313-test_itertools-TestGC.test_starmap
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/test/dynamo_expected_failures/CPython313-test_itertools-TestGC.test_takewhile b/test/dynamo_expected_failures/CPython313-test_itertools-TestGC.test_takewhile
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/test/dynamo_expected_failures/CPython313-test_itertools-TestVariousIteratorArgs.test_dropwhile b/test/dynamo_expected_failures/CPython313-test_itertools-TestVariousIteratorArgs.test_dropwhile
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/test/dynamo_expected_failures/CPython313-test_itertools-TestVariousIteratorArgs.test_islice b/test/dynamo_expected_failures/CPython313-test_itertools-TestVariousIteratorArgs.test_islice
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/test/dynamo_expected_failures/CPython313-test_itertools-TestVariousIteratorArgs.test_starmap b/test/dynamo_expected_failures/CPython313-test_itertools-TestVariousIteratorArgs.test_starmap
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/test/dynamo_expected_failures/CPython313-test_itertools-TestVariousIteratorArgs.test_takewhile b/test/dynamo_expected_failures/CPython313-test_itertools-TestVariousIteratorArgs.test_takewhile
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/test/dynamo_expected_failures/CPython313-test_math-MathTests.testRemainder b/test/dynamo_expected_failures/CPython313-test_math-MathTests.testRemainder
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/test/dynamo_expected_failures/CPython313-test_operator-COperatorTestCase.test___all__ b/test/dynamo_expected_failures/CPython313-test_operator-COperatorTestCase.test___all__
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/test/dynamo_expected_failures/CPython313-test_operator-COperatorTestCase.test_dunder_is_original b/test/dynamo_expected_failures/CPython313-test_operator-COperatorTestCase.test_dunder_is_original
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/test/dynamo_expected_failures/CPython313-test_operator-PyOperatorTestCase.test_dunder_is_original b/test/dynamo_expected_failures/CPython313-test_operator-PyOperatorTestCase.test_dunder_is_original
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/test/dynamo_expected_failures/CPython313-test_ordered_dict-CPythonSubclassMappingTests.test_bool b/test/dynamo_expected_failures/CPython313-test_ordered_dict-CPythonSubclassMappingTests.test_bool
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/test/dynamo_expected_failures/CPython313-test_ordered_dict-CPythonSubclassMappingTests.test_read b/test/dynamo_expected_failures/CPython313-test_ordered_dict-CPythonSubclassMappingTests.test_read
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/test/dynamo_expected_failures/CPython313-test_ordered_dict-CSimpleLRUCacheTests.test_add_after_full b/test/dynamo_expected_failures/CPython313-test_ordered_dict-CSimpleLRUCacheTests.test_add_after_full
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/test/dynamo_expected_failures/CPython313-test_ordered_dict-CSimpleLRUCacheTests.test_popitem b/test/dynamo_expected_failures/CPython313-test_ordered_dict-CSimpleLRUCacheTests.test_popitem
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/test/dynamo_expected_failures/CPython313-test_patma-TestInheritance.test_late_registration_mapping b/test/dynamo_expected_failures/CPython313-test_patma-TestInheritance.test_late_registration_mapping
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/test/dynamo_expected_failures/CPython313-test_patma-TestInheritance.test_late_registration_sequence b/test/dynamo_expected_failures/CPython313-test_patma-TestInheritance.test_late_registration_sequence
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/test/dynamo_expected_failures/CPython313-test_patma-TestPatma.test_patma_204 b/test/dynamo_expected_failures/CPython313-test_patma-TestPatma.test_patma_204
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/test/dynamo_expected_failures/CPython313-test_patma-TestPatma.test_patma_205 b/test/dynamo_expected_failures/CPython313-test_patma-TestPatma.test_patma_205
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/test/dynamo_expected_failures/CPython313-test_patma-TestPatma.test_patma_206 b/test/dynamo_expected_failures/CPython313-test_patma-TestPatma.test_patma_206
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/test/dynamo_expected_failures/CPython313-test_patma-TestPatma.test_patma_207 b/test/dynamo_expected_failures/CPython313-test_patma-TestPatma.test_patma_207
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/test/dynamo_expected_failures/CPython313-test_patma-TestPatma.test_patma_208 b/test/dynamo_expected_failures/CPython313-test_patma-TestPatma.test_patma_208
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/test/dynamo_expected_failures/CPython313-test_patma-TestPatma.test_patma_209 b/test/dynamo_expected_failures/CPython313-test_patma-TestPatma.test_patma_209
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/test/dynamo_expected_failures/CPython313-test_patma-TestPatma.test_patma_211 b/test/dynamo_expected_failures/CPython313-test_patma-TestPatma.test_patma_211
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/test/dynamo_expected_failures/CPython313-test_patma-TestPatma.test_patma_212 b/test/dynamo_expected_failures/CPython313-test_patma-TestPatma.test_patma_212
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/test/dynamo_expected_failures/CPython313-test_patma-TestPatma.test_patma_213 b/test/dynamo_expected_failures/CPython313-test_patma-TestPatma.test_patma_213
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/test/dynamo_expected_failures/CPython313-test_patma-TestPatma.test_patma_214 b/test/dynamo_expected_failures/CPython313-test_patma-TestPatma.test_patma_214
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/test/dynamo_expected_failures/CPython313-test_patma-TestPatma.test_patma_215 b/test/dynamo_expected_failures/CPython313-test_patma-TestPatma.test_patma_215
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/test/dynamo_expected_failures/CPython313-test_patma-TestPatma.test_patma_216 b/test/dynamo_expected_failures/CPython313-test_patma-TestPatma.test_patma_216
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/test/dynamo_expected_failures/CPython313-test_patma-TestPatma.test_patma_217 b/test/dynamo_expected_failures/CPython313-test_patma-TestPatma.test_patma_217
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/test/dynamo_expected_failures/CPython313-test_patma-TestPatma.test_patma_218 b/test/dynamo_expected_failures/CPython313-test_patma-TestPatma.test_patma_218
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/test/dynamo_expected_failures/CPython313-test_patma-TestPatma.test_patma_219 b/test/dynamo_expected_failures/CPython313-test_patma-TestPatma.test_patma_219
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/test/dynamo_expected_failures/CPython313-test_patma-TestPatma.test_patma_220 b/test/dynamo_expected_failures/CPython313-test_patma-TestPatma.test_patma_220
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/test/dynamo_expected_failures/CPython313-test_patma-TestPatma.test_patma_221 b/test/dynamo_expected_failures/CPython313-test_patma-TestPatma.test_patma_221
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/test/dynamo_expected_failures/CPython313-test_patma-TestPatma.test_patma_246 b/test/dynamo_expected_failures/CPython313-test_patma-TestPatma.test_patma_246
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/test/dynamo_expected_failures/CPython313-test_patma-TestPatma.test_patma_247 b/test/dynamo_expected_failures/CPython313-test_patma-TestPatma.test_patma_247
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/test/dynamo_expected_failures/CPython313-test_raise-TestRemovedFunctionality.test_strings b/test/dynamo_expected_failures/CPython313-test_raise-TestRemovedFunctionality.test_strings
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/test/dynamo_expected_failures/CPython313-test_raise-TestRemovedFunctionality.test_tuples b/test/dynamo_expected_failures/CPython313-test_raise-TestRemovedFunctionality.test_tuples
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/test/dynamo_expected_failures/CPython313-test_set-TestCopyingEmpty.test_deep_copy b/test/dynamo_expected_failures/CPython313-test_set-TestCopyingEmpty.test_deep_copy
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/test/dynamo_expected_failures/CPython313-test_set-TestCopyingNested.test_deep_copy b/test/dynamo_expected_failures/CPython313-test_set-TestCopyingNested.test_deep_copy
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/test/dynamo_expected_failures/CPython313-test_set-TestCopyingSingleton.test_deep_copy b/test/dynamo_expected_failures/CPython313-test_set-TestCopyingSingleton.test_deep_copy
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/test/dynamo_expected_failures/CPython313-test_set-TestCopyingTriple.test_deep_copy b/test/dynamo_expected_failures/CPython313-test_set-TestCopyingTriple.test_deep_copy
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/test/dynamo_expected_failures/CPython313-test_set-TestCopyingTuple.test_deep_copy b/test/dynamo_expected_failures/CPython313-test_set-TestCopyingTuple.test_deep_copy
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/test/dynamo_expected_failures/CPython313-test_set-TestFrozenSetSubclass.test_copy b/test/dynamo_expected_failures/CPython313-test_set-TestFrozenSetSubclass.test_copy
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/test/dynamo_expected_failures/CPython313-test_set-TestFrozenSetSubclass.test_intersection b/test/dynamo_expected_failures/CPython313-test_set-TestFrozenSetSubclass.test_intersection
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/test/dynamo_expected_failures/CPython313-test_set-TestFrozenSetSubclass.test_singleton_empty_frozenset b/test/dynamo_expected_failures/CPython313-test_set-TestFrozenSetSubclass.test_singleton_empty_frozenset
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/test/dynamo_expected_failures/CPython313-test_set-TestSetSubclass.test_copy b/test/dynamo_expected_failures/CPython313-test_set-TestSetSubclass.test_copy
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/test/dynamo_expected_failures/CPython313-test_set-TestSetSubclass.test_intersection b/test/dynamo_expected_failures/CPython313-test_set-TestSetSubclass.test_intersection
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/test/dynamo_expected_failures/CPython313-test_tuple-TupleTest.test_addmul b/test/dynamo_expected_failures/CPython313-test_tuple-TupleTest.test_addmul
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestMixedIntFlagFunction.test_dir_on_class b/test/dynamo_expected_failures/CPython313-test_types-ClassCreationTests.test_bad___prepare__
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestMixedIntFlagFunction.test_dir_on_class
rename to test/dynamo_expected_failures/CPython313-test_types-ClassCreationTests.test_bad___prepare__
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestMixedIntFlagFunction.test_format b/test/dynamo_expected_failures/CPython313-test_types-ClassCreationTests.test_get_original_bases
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestMixedIntFlagFunction.test_format
rename to test/dynamo_expected_failures/CPython313-test_types-ClassCreationTests.test_get_original_bases
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestMixedIntFlagFunction.test_gnv_is_static b/test/dynamo_expected_failures/CPython313-test_types-ClassCreationTests.test_metaclass_derivation
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestMixedIntFlagFunction.test_gnv_is_static
rename to test/dynamo_expected_failures/CPython313-test_types-ClassCreationTests.test_metaclass_derivation
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestMixedIntFlagFunction.test_programmatic_function_from_dict b/test/dynamo_expected_failures/CPython313-test_types-ClassCreationTests.test_metaclass_override_callable
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestMixedIntFlagFunction.test_programmatic_function_from_dict
rename to test/dynamo_expected_failures/CPython313-test_types-ClassCreationTests.test_metaclass_override_callable
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestMixedIntFlagFunction.test_programmatic_function_iterable b/test/dynamo_expected_failures/CPython313-test_types-ClassCreationTests.test_metaclass_override_function
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestMixedIntFlagFunction.test_programmatic_function_iterable
rename to test/dynamo_expected_failures/CPython313-test_types-ClassCreationTests.test_metaclass_override_function
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestMixedIntFlagFunction.test_programmatic_function_string b/test/dynamo_expected_failures/CPython313-test_types-ClassCreationTests.test_new_class_basics
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestMixedIntFlagFunction.test_programmatic_function_string
rename to test/dynamo_expected_failures/CPython313-test_types-ClassCreationTests.test_new_class_basics
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestMixedIntFlagFunction.test_programmatic_function_string_list b/test/dynamo_expected_failures/CPython313-test_types-ClassCreationTests.test_new_class_defaults
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestMixedIntFlagFunction.test_programmatic_function_string_list
rename to test/dynamo_expected_failures/CPython313-test_types-ClassCreationTests.test_new_class_defaults
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestMixedIntFlagFunction.test_repr b/test/dynamo_expected_failures/CPython313-test_types-ClassCreationTests.test_new_class_exec_body
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestMixedIntFlagFunction.test_repr
rename to test/dynamo_expected_failures/CPython313-test_types-ClassCreationTests.test_new_class_exec_body
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestMixedIntFlagFunction.test_str b/test/dynamo_expected_failures/CPython313-test_types-ClassCreationTests.test_new_class_meta
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestMixedIntFlagFunction.test_str
rename to test/dynamo_expected_failures/CPython313-test_types-ClassCreationTests.test_new_class_meta
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestMixedIntFunction.test_dir_on_class b/test/dynamo_expected_failures/CPython313-test_types-ClassCreationTests.test_new_class_meta_with_base
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestMixedIntFunction.test_dir_on_class
rename to test/dynamo_expected_failures/CPython313-test_types-ClassCreationTests.test_new_class_meta_with_base
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestMixedIntFunction.test_format b/test/dynamo_expected_failures/CPython313-test_types-ClassCreationTests.test_new_class_subclass
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestMixedIntFunction.test_format
rename to test/dynamo_expected_failures/CPython313-test_types-ClassCreationTests.test_new_class_subclass
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestMixedIntFunction.test_gnv_is_static b/test/dynamo_expected_failures/CPython313-test_types-ClassCreationTests.test_new_class_with_mro_entry
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestMixedIntFunction.test_gnv_is_static
rename to test/dynamo_expected_failures/CPython313-test_types-ClassCreationTests.test_new_class_with_mro_entry
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestMixedIntFunction.test_programmatic_function_from_dict b/test/dynamo_expected_failures/CPython313-test_types-ClassCreationTests.test_new_class_with_mro_entry_genericalias
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestMixedIntFunction.test_programmatic_function_from_dict
rename to test/dynamo_expected_failures/CPython313-test_types-ClassCreationTests.test_new_class_with_mro_entry_genericalias
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestMixedIntFunction.test_programmatic_function_iterable b/test/dynamo_expected_failures/CPython313-test_types-ClassCreationTests.test_new_class_with_mro_entry_multiple
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestMixedIntFunction.test_programmatic_function_iterable
rename to test/dynamo_expected_failures/CPython313-test_types-ClassCreationTests.test_new_class_with_mro_entry_multiple
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestMixedIntFunction.test_programmatic_function_string b/test/dynamo_expected_failures/CPython313-test_types-ClassCreationTests.test_new_class_with_mro_entry_multiple_2
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestMixedIntFunction.test_programmatic_function_string
rename to test/dynamo_expected_failures/CPython313-test_types-ClassCreationTests.test_new_class_with_mro_entry_multiple_2
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestMixedIntFunction.test_programmatic_function_string_list b/test/dynamo_expected_failures/CPython313-test_types-ClassCreationTests.test_new_class_with_mro_entry_none
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestMixedIntFunction.test_programmatic_function_string_list
rename to test/dynamo_expected_failures/CPython313-test_types-ClassCreationTests.test_new_class_with_mro_entry_none
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestMixedIntFunction.test_repr b/test/dynamo_expected_failures/CPython313-test_types-ClassCreationTests.test_one_argument_type
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestMixedIntFunction.test_repr
rename to test/dynamo_expected_failures/CPython313-test_types-ClassCreationTests.test_one_argument_type
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestMixedIntFunction.test_str b/test/dynamo_expected_failures/CPython313-test_types-ClassCreationTests.test_prepare_class
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestMixedIntFunction.test_str
rename to test/dynamo_expected_failures/CPython313-test_types-ClassCreationTests.test_prepare_class
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestMixedStrClass.test_dir_on_class b/test/dynamo_expected_failures/CPython313-test_types-ClassCreationTests.test_resolve_bases_with_mro_entry
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestMixedStrClass.test_dir_on_class
rename to test/dynamo_expected_failures/CPython313-test_types-ClassCreationTests.test_resolve_bases_with_mro_entry
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestMixedStrClass.test_format b/test/dynamo_expected_failures/CPython313-test_types-ClassCreationTests.test_tuple_subclass_as_bases
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestMixedStrClass.test_format
rename to test/dynamo_expected_failures/CPython313-test_types-ClassCreationTests.test_tuple_subclass_as_bases
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestMixedStrClass.test_gnv_is_static b/test/dynamo_expected_failures/CPython313-test_types-CoroutineTests.test_async_def
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestMixedStrClass.test_gnv_is_static
rename to test/dynamo_expected_failures/CPython313-test_types-CoroutineTests.test_async_def
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestMixedStrClass.test_programmatic_function_from_dict b/test/dynamo_expected_failures/CPython313-test_types-CoroutineTests.test_duck_coro
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestMixedStrClass.test_programmatic_function_from_dict
rename to test/dynamo_expected_failures/CPython313-test_types-CoroutineTests.test_duck_coro
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestMixedStrClass.test_programmatic_function_iterable b/test/dynamo_expected_failures/CPython313-test_types-CoroutineTests.test_duck_corogen
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestMixedStrClass.test_programmatic_function_iterable
rename to test/dynamo_expected_failures/CPython313-test_types-CoroutineTests.test_duck_corogen
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestMixedStrClass.test_programmatic_function_string b/test/dynamo_expected_failures/CPython313-test_types-CoroutineTests.test_duck_functional_gen
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestMixedStrClass.test_programmatic_function_string
rename to test/dynamo_expected_failures/CPython313-test_types-CoroutineTests.test_duck_functional_gen
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestMixedStrClass.test_programmatic_function_string_list b/test/dynamo_expected_failures/CPython313-test_types-CoroutineTests.test_duck_gen
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestMixedStrClass.test_programmatic_function_string_list
rename to test/dynamo_expected_failures/CPython313-test_types-CoroutineTests.test_duck_gen
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestMixedStrClass.test_repr b/test/dynamo_expected_failures/CPython313-test_types-CoroutineTests.test_gen
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestMixedStrClass.test_repr
rename to test/dynamo_expected_failures/CPython313-test_types-CoroutineTests.test_gen
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestMixedStrClass.test_str b/test/dynamo_expected_failures/CPython313-test_types-CoroutineTests.test_genfunc
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestMixedStrClass.test_str
rename to test/dynamo_expected_failures/CPython313-test_types-CoroutineTests.test_genfunc
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestMixedStrFunction.test_dir_on_class b/test/dynamo_expected_failures/CPython313-test_types-CoroutineTests.test_non_gen_values
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestMixedStrFunction.test_dir_on_class
rename to test/dynamo_expected_failures/CPython313-test_types-CoroutineTests.test_non_gen_values
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestMixedStrFunction.test_format b/test/dynamo_expected_failures/CPython313-test_types-CoroutineTests.test_returning_itercoro
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestMixedStrFunction.test_format
rename to test/dynamo_expected_failures/CPython313-test_types-CoroutineTests.test_returning_itercoro
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestMixedStrFunction.test_gnv_is_static b/test/dynamo_expected_failures/CPython313-test_types-CoroutineTests.test_wrapper_object
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestMixedStrFunction.test_gnv_is_static
rename to test/dynamo_expected_failures/CPython313-test_types-CoroutineTests.test_wrapper_object
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestMixedStrFunction.test_programmatic_function_from_dict b/test/dynamo_expected_failures/CPython313-test_types-CoroutineTests.test_wrong_args
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestMixedStrFunction.test_programmatic_function_from_dict
rename to test/dynamo_expected_failures/CPython313-test_types-CoroutineTests.test_wrong_args
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestMixedStrFunction.test_programmatic_function_iterable b/test/dynamo_expected_failures/CPython313-test_types-FunctionTests.test_function_type_defaults
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestMixedStrFunction.test_programmatic_function_iterable
rename to test/dynamo_expected_failures/CPython313-test_types-FunctionTests.test_function_type_defaults
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestMixedStrFunction.test_programmatic_function_string b/test/dynamo_expected_failures/CPython313-test_types-FunctionTests.test_function_type_wrong_defaults
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestMixedStrFunction.test_programmatic_function_string
rename to test/dynamo_expected_failures/CPython313-test_types-FunctionTests.test_function_type_wrong_defaults
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestMixedStrFunction.test_programmatic_function_string_list b/test/dynamo_expected_failures/CPython313-test_types-MappingProxyTests.test_chainmap
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestMixedStrFunction.test_programmatic_function_string_list
rename to test/dynamo_expected_failures/CPython313-test_types-MappingProxyTests.test_chainmap
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestMixedStrFunction.test_repr b/test/dynamo_expected_failures/CPython313-test_types-MappingProxyTests.test_constructor
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestMixedStrFunction.test_repr
rename to test/dynamo_expected_failures/CPython313-test_types-MappingProxyTests.test_constructor
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestMixedStrFunction.test_str b/test/dynamo_expected_failures/CPython313-test_types-MappingProxyTests.test_customdict
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestMixedStrFunction.test_str
rename to test/dynamo_expected_failures/CPython313-test_types-MappingProxyTests.test_customdict
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestPlainEnumClass.test_format b/test/dynamo_expected_failures/CPython313-test_types-MappingProxyTests.test_hash
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestPlainEnumClass.test_format
rename to test/dynamo_expected_failures/CPython313-test_types-MappingProxyTests.test_hash
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestPlainEnumClass.test_gnv_is_static b/test/dynamo_expected_failures/CPython313-test_types-MappingProxyTests.test_methods
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestPlainEnumClass.test_gnv_is_static
rename to test/dynamo_expected_failures/CPython313-test_types-MappingProxyTests.test_methods
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestPlainEnumClass.test_programmatic_function_from_dict b/test/dynamo_expected_failures/CPython313-test_types-MappingProxyTests.test_missing
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestPlainEnumClass.test_programmatic_function_from_dict
rename to test/dynamo_expected_failures/CPython313-test_types-MappingProxyTests.test_missing
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestPlainEnumClass.test_programmatic_function_iterable b/test/dynamo_expected_failures/CPython313-test_types-MappingProxyTests.test_reversed
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestPlainEnumClass.test_programmatic_function_iterable
rename to test/dynamo_expected_failures/CPython313-test_types-MappingProxyTests.test_reversed
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestPlainEnumClass.test_programmatic_function_string b/test/dynamo_expected_failures/CPython313-test_types-MappingProxyTests.test_union
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestPlainEnumClass.test_programmatic_function_string
rename to test/dynamo_expected_failures/CPython313-test_types-MappingProxyTests.test_union
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestPlainEnumClass.test_programmatic_function_string_list b/test/dynamo_expected_failures/CPython313-test_types-SimpleNamespaceTests.test_as_dict
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestPlainEnumClass.test_programmatic_function_string_list
rename to test/dynamo_expected_failures/CPython313-test_types-SimpleNamespaceTests.test_as_dict
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestPlainEnumClass.test_repr b/test/dynamo_expected_failures/CPython313-test_types-SimpleNamespaceTests.test_attrdel
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestPlainEnumClass.test_repr
rename to test/dynamo_expected_failures/CPython313-test_types-SimpleNamespaceTests.test_attrdel
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestPlainEnumClass.test_str b/test/dynamo_expected_failures/CPython313-test_types-SimpleNamespaceTests.test_attrget
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestPlainEnumClass.test_str
rename to test/dynamo_expected_failures/CPython313-test_types-SimpleNamespaceTests.test_attrget
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestPlainEnumFunction.test_format b/test/dynamo_expected_failures/CPython313-test_types-SimpleNamespaceTests.test_attrset
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestPlainEnumFunction.test_format
rename to test/dynamo_expected_failures/CPython313-test_types-SimpleNamespaceTests.test_attrset
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestPlainEnumFunction.test_gnv_is_static b/test/dynamo_expected_failures/CPython313-test_types-SimpleNamespaceTests.test_constructor
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestPlainEnumFunction.test_gnv_is_static
rename to test/dynamo_expected_failures/CPython313-test_types-SimpleNamespaceTests.test_constructor
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestPlainEnumFunction.test_programmatic_function_from_dict b/test/dynamo_expected_failures/CPython313-test_types-SimpleNamespaceTests.test_equal
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestPlainEnumFunction.test_programmatic_function_from_dict
rename to test/dynamo_expected_failures/CPython313-test_types-SimpleNamespaceTests.test_equal
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestPlainEnumFunction.test_programmatic_function_iterable b/test/dynamo_expected_failures/CPython313-test_types-SimpleNamespaceTests.test_fake_namespace_compare
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestPlainEnumFunction.test_programmatic_function_iterable
rename to test/dynamo_expected_failures/CPython313-test_types-SimpleNamespaceTests.test_fake_namespace_compare
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestPlainEnumFunction.test_programmatic_function_string b/test/dynamo_expected_failures/CPython313-test_types-SimpleNamespaceTests.test_nested
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestPlainEnumFunction.test_programmatic_function_string
rename to test/dynamo_expected_failures/CPython313-test_types-SimpleNamespaceTests.test_nested
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestPlainEnumFunction.test_programmatic_function_string_list b/test/dynamo_expected_failures/CPython313-test_types-SimpleNamespaceTests.test_pickle
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestPlainEnumFunction.test_programmatic_function_string_list
rename to test/dynamo_expected_failures/CPython313-test_types-SimpleNamespaceTests.test_pickle
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestPlainEnumFunction.test_repr b/test/dynamo_expected_failures/CPython313-test_types-SimpleNamespaceTests.test_recursive
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestPlainEnumFunction.test_repr
rename to test/dynamo_expected_failures/CPython313-test_types-SimpleNamespaceTests.test_recursive
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestPlainEnumFunction.test_str b/test/dynamo_expected_failures/CPython313-test_types-SimpleNamespaceTests.test_recursive_repr
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestPlainEnumFunction.test_str
rename to test/dynamo_expected_failures/CPython313-test_types-SimpleNamespaceTests.test_recursive_repr
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestPlainFlagClass.test_format b/test/dynamo_expected_failures/CPython313-test_types-SimpleNamespaceTests.test_replace
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestPlainFlagClass.test_format
rename to test/dynamo_expected_failures/CPython313-test_types-SimpleNamespaceTests.test_replace
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestPlainFlagClass.test_gnv_is_static b/test/dynamo_expected_failures/CPython313-test_types-SimpleNamespaceTests.test_replace_subclass
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestPlainFlagClass.test_gnv_is_static
rename to test/dynamo_expected_failures/CPython313-test_types-SimpleNamespaceTests.test_replace_subclass
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestPlainFlagClass.test_programmatic_function_from_dict b/test/dynamo_expected_failures/CPython313-test_types-SimpleNamespaceTests.test_repr
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestPlainFlagClass.test_programmatic_function_from_dict
rename to test/dynamo_expected_failures/CPython313-test_types-SimpleNamespaceTests.test_repr
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestPlainFlagClass.test_programmatic_function_iterable b/test/dynamo_expected_failures/CPython313-test_types-SimpleNamespaceTests.test_subclass
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestPlainFlagClass.test_programmatic_function_iterable
rename to test/dynamo_expected_failures/CPython313-test_types-SimpleNamespaceTests.test_subclass
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestPlainFlagClass.test_programmatic_function_string b/test/dynamo_expected_failures/CPython313-test_types-SimpleNamespaceTests.test_unbound
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestPlainFlagClass.test_programmatic_function_string
rename to test/dynamo_expected_failures/CPython313-test_types-SimpleNamespaceTests.test_unbound
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestPlainFlagClass.test_programmatic_function_string_list b/test/dynamo_expected_failures/CPython313-test_types-SimpleNamespaceTests.test_underlying_dict
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestPlainFlagClass.test_programmatic_function_string_list
rename to test/dynamo_expected_failures/CPython313-test_types-SimpleNamespaceTests.test_underlying_dict
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestPlainFlagClass.test_repr b/test/dynamo_expected_failures/CPython313-test_types-SubinterpreterTests.test_static_types_inherited_slots
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestPlainFlagClass.test_repr
rename to test/dynamo_expected_failures/CPython313-test_types-SubinterpreterTests.test_static_types_inherited_slots
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestPlainFlagClass.test_str b/test/dynamo_expected_failures/CPython313-test_types-TypesTests.test_call_unbound_crash
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestPlainFlagClass.test_str
rename to test/dynamo_expected_failures/CPython313-test_types-TypesTests.test_call_unbound_crash
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestPlainFlagFunction.test_format b/test/dynamo_expected_failures/CPython313-test_types-TypesTests.test_dunder_get_signature
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestPlainFlagFunction.test_format
rename to test/dynamo_expected_failures/CPython313-test_types-TypesTests.test_dunder_get_signature
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestPlainFlagFunction.test_gnv_is_static b/test/dynamo_expected_failures/CPython313-test_types-TypesTests.test_float__format__
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestPlainFlagFunction.test_gnv_is_static
rename to test/dynamo_expected_failures/CPython313-test_types-TypesTests.test_float__format__
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestPlainFlagFunction.test_programmatic_function_from_dict b/test/dynamo_expected_failures/CPython313-test_types-TypesTests.test_float__format__locale
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestPlainFlagFunction.test_programmatic_function_from_dict
rename to test/dynamo_expected_failures/CPython313-test_types-TypesTests.test_float__format__locale
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestPlainFlagFunction.test_programmatic_function_iterable b/test/dynamo_expected_failures/CPython313-test_types-TypesTests.test_float_to_string
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestPlainFlagFunction.test_programmatic_function_iterable
rename to test/dynamo_expected_failures/CPython313-test_types-TypesTests.test_float_to_string
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestPlainFlagFunction.test_programmatic_function_string b/test/dynamo_expected_failures/CPython313-test_types-TypesTests.test_format_spec_errors
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestPlainFlagFunction.test_programmatic_function_string
rename to test/dynamo_expected_failures/CPython313-test_types-TypesTests.test_format_spec_errors
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestPlainFlagFunction.test_programmatic_function_string_list b/test/dynamo_expected_failures/CPython313-test_types-TypesTests.test_int__format__
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestPlainFlagFunction.test_programmatic_function_string_list
rename to test/dynamo_expected_failures/CPython313-test_types-TypesTests.test_int__format__
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestPlainFlagFunction.test_repr b/test/dynamo_expected_failures/CPython313-test_types-TypesTests.test_int__format__locale
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestPlainFlagFunction.test_repr
rename to test/dynamo_expected_failures/CPython313-test_types-TypesTests.test_int__format__locale
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestPlainFlagFunction.test_str b/test/dynamo_expected_failures/CPython313-test_types-TypesTests.test_method_descriptor_crash
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestPlainFlagFunction.test_str
rename to test/dynamo_expected_failures/CPython313-test_types-TypesTests.test_method_descriptor_crash
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestSpecial.test_programmatic_function_string_list_with_start b/test/dynamo_expected_failures/CPython313-test_types-TypesTests.test_method_descriptor_types
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestSpecial.test_programmatic_function_string_list_with_start
rename to test/dynamo_expected_failures/CPython313-test_types-TypesTests.test_method_descriptor_types
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestSpecial.test_programmatic_function_string_with_start b/test/dynamo_expected_failures/CPython313-test_types-TypesTests.test_method_wrapper_types
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestSpecial.test_programmatic_function_string_with_start
rename to test/dynamo_expected_failures/CPython313-test_types-TypesTests.test_method_wrapper_types
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestSpecial.test_programmatic_function_type b/test/dynamo_expected_failures/CPython313-test_types-TypesTests.test_traceback_and_frame_types
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestSpecial.test_programmatic_function_type
rename to test/dynamo_expected_failures/CPython313-test_types-TypesTests.test_traceback_and_frame_types
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestSpecial.test_programmatic_function_type_from_subclass b/test/dynamo_expected_failures/CPython313-test_types-UnionTests.test_bad_instancecheck
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestSpecial.test_programmatic_function_type_from_subclass
rename to test/dynamo_expected_failures/CPython313-test_types-UnionTests.test_bad_instancecheck
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestSpecial.test_programmatic_function_type_from_subclass_with_start b/test/dynamo_expected_failures/CPython313-test_types-UnionTests.test_bad_subclasscheck
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestSpecial.test_programmatic_function_type_from_subclass_with_start
rename to test/dynamo_expected_failures/CPython313-test_types-UnionTests.test_bad_subclasscheck
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestSpecial.test_programmatic_function_type_with_start b/test/dynamo_expected_failures/CPython313-test_types-UnionTests.test_hash
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestSpecial.test_programmatic_function_type_with_start
rename to test/dynamo_expected_failures/CPython313-test_types-UnionTests.test_hash
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestStrEnumClass.test_copy b/test/dynamo_expected_failures/CPython313-test_types-UnionTests.test_instancecheck_and_subclasscheck_order
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestStrEnumClass.test_copy
rename to test/dynamo_expected_failures/CPython313-test_types-UnionTests.test_instancecheck_and_subclasscheck_order
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestStrEnumClass.test_dir_on_class b/test/dynamo_expected_failures/CPython313-test_types-UnionTests.test_or_type_operator_reference_cycle
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestStrEnumClass.test_dir_on_class
rename to test/dynamo_expected_failures/CPython313-test_types-UnionTests.test_or_type_operator_reference_cycle
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestStrEnumClass.test_format b/test/dynamo_expected_failures/CPython313-test_types-UnionTests.test_or_type_operator_with_Alias
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestStrEnumClass.test_format
rename to test/dynamo_expected_failures/CPython313-test_types-UnionTests.test_or_type_operator_with_Alias
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestStrEnumClass.test_gnv_is_static b/test/dynamo_expected_failures/CPython313-test_types-UnionTests.test_or_type_operator_with_Literal
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestStrEnumClass.test_gnv_is_static
rename to test/dynamo_expected_failures/CPython313-test_types-UnionTests.test_or_type_operator_with_Literal
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestStrEnumClass.test_programmatic_function_from_dict b/test/dynamo_expected_failures/CPython313-test_types-UnionTests.test_or_type_operator_with_NewType
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestStrEnumClass.test_programmatic_function_from_dict
rename to test/dynamo_expected_failures/CPython313-test_types-UnionTests.test_or_type_operator_with_NewType
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestStrEnumClass.test_programmatic_function_iterable b/test/dynamo_expected_failures/CPython313-test_types-UnionTests.test_or_type_operator_with_SpecialForm
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestStrEnumClass.test_programmatic_function_iterable
rename to test/dynamo_expected_failures/CPython313-test_types-UnionTests.test_or_type_operator_with_SpecialForm
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestStrEnumClass.test_programmatic_function_string b/test/dynamo_expected_failures/CPython313-test_types-UnionTests.test_or_type_operator_with_TypeVar
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestStrEnumClass.test_programmatic_function_string
rename to test/dynamo_expected_failures/CPython313-test_types-UnionTests.test_or_type_operator_with_TypeVar
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestStrEnumClass.test_programmatic_function_string_list b/test/dynamo_expected_failures/CPython313-test_types-UnionTests.test_or_type_operator_with_bad_module
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestStrEnumClass.test_programmatic_function_string_list
rename to test/dynamo_expected_failures/CPython313-test_types-UnionTests.test_or_type_operator_with_bad_module
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestStrEnumClass.test_repr b/test/dynamo_expected_failures/CPython313-test_types-UnionTests.test_or_type_operator_with_forward
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestStrEnumClass.test_repr
rename to test/dynamo_expected_failures/CPython313-test_types-UnionTests.test_or_type_operator_with_forward
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestStrEnumClass.test_str b/test/dynamo_expected_failures/CPython313-test_types-UnionTests.test_or_type_operator_with_genericalias
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestStrEnumClass.test_str
rename to test/dynamo_expected_failures/CPython313-test_types-UnionTests.test_or_type_operator_with_genericalias
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestStrEnumFunction.test_copy b/test/dynamo_expected_failures/CPython313-test_types-UnionTests.test_or_type_repr
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestStrEnumFunction.test_copy
rename to test/dynamo_expected_failures/CPython313-test_types-UnionTests.test_or_type_repr
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestStrEnumFunction.test_dir_on_class b/test/dynamo_expected_failures/CPython313-test_types-UnionTests.test_or_types_operator
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestStrEnumFunction.test_dir_on_class
rename to test/dynamo_expected_failures/CPython313-test_types-UnionTests.test_or_types_operator
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestStrEnumFunction.test_format b/test/dynamo_expected_failures/CPython313-test_types-UnionTests.test_union_args
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestStrEnumFunction.test_format
rename to test/dynamo_expected_failures/CPython313-test_types-UnionTests.test_union_args
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestStrEnumFunction.test_gnv_is_static b/test/dynamo_expected_failures/CPython313-test_types-UnionTests.test_union_copy
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestStrEnumFunction.test_gnv_is_static
rename to test/dynamo_expected_failures/CPython313-test_types-UnionTests.test_union_copy
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestStrEnumFunction.test_programmatic_function_from_dict b/test/dynamo_expected_failures/CPython313-test_types-UnionTests.test_union_of_unhashable
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestStrEnumFunction.test_programmatic_function_from_dict
rename to test/dynamo_expected_failures/CPython313-test_types-UnionTests.test_union_of_unhashable
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestStrEnumFunction.test_programmatic_function_iterable b/test/dynamo_expected_failures/CPython313-test_types-UnionTests.test_union_parameter_chaining
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestStrEnumFunction.test_programmatic_function_iterable
rename to test/dynamo_expected_failures/CPython313-test_types-UnionTests.test_union_parameter_chaining
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestStrEnumFunction.test_programmatic_function_string b/test/dynamo_expected_failures/CPython313-test_types-UnionTests.test_union_parameter_substitution
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestStrEnumFunction.test_programmatic_function_string
rename to test/dynamo_expected_failures/CPython313-test_types-UnionTests.test_union_parameter_substitution
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestStrEnumFunction.test_programmatic_function_string_list b/test/dynamo_expected_failures/CPython313-test_types-UnionTests.test_union_parameter_substitution_errors
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestStrEnumFunction.test_programmatic_function_string_list
rename to test/dynamo_expected_failures/CPython313-test_types-UnionTests.test_union_parameter_substitution_errors
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestStrEnumFunction.test_repr b/test/dynamo_expected_failures/CPython313-test_types-UnionTests.test_union_pickle
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestStrEnumFunction.test_repr
rename to test/dynamo_expected_failures/CPython313-test_types-UnionTests.test_union_pickle
diff --git a/test/dynamo_expected_failures/CPython313-test_userdict-UserDictTest.test_bool b/test/dynamo_expected_failures/CPython313-test_userdict-UserDictTest.test_bool
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/test/dynamo_expected_failures/CPython313-test_userlist-UserListTest.test_init b/test/dynamo_expected_failures/CPython313-test_userlist-UserListTest.test_init
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/test/dynamo_expected_failures/CPython313-test_userlist-UserListTest.test_userlist_copy b/test/dynamo_expected_failures/CPython313-test_userlist-UserListTest.test_userlist_copy
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestStrEnumFunction.test_shadowed_attr b/test/dynamo_expected_failures/TestCompileTransformsCPU.test_compile_vmap_hessian_cpu
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestStrEnumFunction.test_shadowed_attr
rename to test/dynamo_expected_failures/TestCompileTransformsCPU.test_compile_vmap_hessian_cpu
diff --git a/test/dynamo_expected_failures/TestComposabilityCPU.test_requires_grad_inside_transform_cpu b/test/dynamo_expected_failures/TestComposabilityCPU.test_requires_grad_inside_transform_cpu
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/test/dynamo_expected_failures/TestComposabilityCUDA.test_requires_grad_inside_transform_cuda b/test/dynamo_expected_failures/TestComposabilityCUDA.test_requires_grad_inside_transform_cuda
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/test/dynamo_expected_failures/TestJit.test_batchnorm b/test/dynamo_expected_failures/TestJit.test_batchnorm
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/test/dynamo_expected_failures/CPython313-test_enum-TestStrEnumFunction.test_str b/test/dynamo_expected_failures/TestNN.test_module_apply_inplace_op
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_enum-TestStrEnumFunction.test_str
rename to test/dynamo_expected_failures/TestNN.test_module_apply_inplace_op
diff --git a/test/dynamo_expected_failures/CPython313-test_exception_variations-ExceptStarTestCases.test_nested b/test/dynamo_expected_failures/TestNN.test_to
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_exception_variations-ExceptStarTestCases.test_nested
rename to test/dynamo_expected_failures/TestNN.test_to
diff --git a/test/dynamo_expected_failures/TestScript.test_parse_nested_names b/test/dynamo_expected_failures/TestScript.test_parse_nested_names
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/test/dynamo_expected_failures/TestScript.test_parse_tensor_constants b/test/dynamo_expected_failures/TestScript.test_parse_tensor_constants
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/test/dynamo_expected_failures/TestTestParametrizationDeviceTypeCPU.test_modules_decorator_applies_module_and_param_specific_decorators_cpu b/test/dynamo_expected_failures/TestTestParametrizationDeviceTypeCPU.test_modules_decorator_applies_module_and_param_specific_decorators_cpu
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/test/dynamo_expected_failures/TestTestParametrizationDeviceTypeCPU.test_ops_composition_names_cpu b/test/dynamo_expected_failures/TestTestParametrizationDeviceTypeCPU.test_ops_composition_names_cpu
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/test/dynamo_expected_failures/TestTestParametrizationDeviceTypeCPU.test_ops_decorator_applies_op_and_param_specific_decorators_cpu b/test/dynamo_expected_failures/TestTestParametrizationDeviceTypeCPU.test_ops_decorator_applies_op_and_param_specific_decorators_cpu
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/test/dynamo_expected_failures/TestTorch.test_as_subclass b/test/dynamo_expected_failures/TestTorch.test_as_subclass
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/test/dynamo_expected_failures/TestTorch.test_print b/test/dynamo_expected_failures/TestTorch.test_print
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/test/dynamo_expected_failures/TestTorchFunctionMode.test_disable_subclass_mode b/test/dynamo_expected_failures/TestTorchFunctionMode.test_disable_subclass_mode
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/test/dynamo_expected_failures/TestTorchFunctionMode.test_modes_return_notimplemented b/test/dynamo_expected_failures/TestTorchFunctionMode.test_modes_return_notimplemented
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/test/dynamo_expected_failures/TestTorchFunctionOverride.test_Tensor___cuda_array_interface_____get__ b/test/dynamo_expected_failures/TestTorchFunctionOverride.test_Tensor___cuda_array_interface_____get__
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/test/expect/HasDecompTest.test_aten_core_operators.expect b/test/expect/HasDecompTest.test_aten_core_operators.expect
index 103cb407e078b..63d248c5a1420 100644
--- a/test/expect/HasDecompTest.test_aten_core_operators.expect
+++ b/test/expect/HasDecompTest.test_aten_core_operators.expect
@@ -351,6 +351,8 @@ aten::lt.Tensor
 aten::lt.Tensor_out
 aten::lt_.Scalar
 aten::lt_.Tensor
+aten::max_pool2d_with_indices_backward
+aten::max_pool2d_with_indices_backward.grad_input
 aten::maximum
 aten::maximum.out
 aten::mean
diff --git a/test/expect/HasDecompTest.test_has_decomposition.expect b/test/expect/HasDecompTest.test_has_decomposition.expect
index 1080935c11bd0..61882630021d8 100644
--- a/test/expect/HasDecompTest.test_has_decomposition.expect
+++ b/test/expect/HasDecompTest.test_has_decomposition.expect
@@ -131,6 +131,7 @@ aten::_fill_mem_eff_dropout_mask_
 aten::_flash_attention_backward
 aten::_flash_attention_forward
 aten::_flash_attention_forward.quantized
+aten::_flash_attention_forward_no_dropout_inplace
 aten::_foobar
 aten::_foobar.out
 aten::_foreach_abs
@@ -196,6 +197,8 @@ aten::_foreach_clamp_min.Scalar_out
 aten::_foreach_clamp_min_.List
 aten::_foreach_clamp_min_.Scalar
 aten::_foreach_clamp_min_.ScalarList
+aten::_foreach_clone
+aten::_foreach_clone.out
 aten::_foreach_copy
 aten::_foreach_copy.out
 aten::_foreach_copy_
@@ -390,6 +393,7 @@ aten::_fw_primal_copy
 aten::_fw_primal_copy.out
 aten::_grid_sampler_2d_cpu_fallback
 aten::_grid_sampler_2d_cpu_fallback.out
+aten::_grid_sampler_2d_cpu_fallback_backward
 aten::_grouped_mm
 aten::_has_same_storage_numel
 aten::_histogramdd_bin_edges
@@ -506,6 +510,14 @@ aten::_pdist_backward
 aten::_pdist_backward.out
 aten::_pdist_forward
 aten::_pdist_forward.out
+aten::_philox_key_fold_in
+aten::_philox_key_split
+aten::_philox_normal
+aten::_philox_normal.out
+aten::_philox_normal_
+aten::_philox_uniform
+aten::_philox_uniform.out
+aten::_philox_uniform_
 aten::_pin_memory
 aten::_pin_memory.out
 aten::_reshape_alias_copy
@@ -648,6 +660,10 @@ aten::_upsample_bilinear2d_aa
 aten::_upsample_bilinear2d_aa.out
 aten::_upsample_bilinear2d_aa_backward
 aten::_upsample_bilinear2d_aa_backward.grad_input
+aten::_upsample_lanczos2d_aa
+aten::_upsample_lanczos2d_aa.out
+aten::_upsample_lanczos2d_aa_backward
+aten::_upsample_lanczos2d_aa_backward.grad_input
 aten::_upsample_nearest_exact1d_backward
 aten::_upsample_nearest_exact1d_backward.grad_input
 aten::_upsample_nearest_exact2d_backward
@@ -856,10 +872,6 @@ aten::hamming_window.periodic_alpha_beta
 aten::hamming_window.periodic_alpha_beta_out
 aten::hamming_window.periodic_alpha_out
 aten::hamming_window.periodic_out
-aten::hann_window
-aten::hann_window.out
-aten::hann_window.periodic
-aten::hann_window.periodic_out
 aten::hardshrink_backward
 aten::hardshrink_backward.grad_input
 aten::hash_tensor
@@ -953,8 +965,6 @@ aten::max_pool2d_backward
 aten::max_pool2d_backward.out
 aten::max_pool2d_with_indices
 aten::max_pool2d_with_indices.out
-aten::max_pool2d_with_indices_backward
-aten::max_pool2d_with_indices_backward.grad_input
 aten::max_pool3d_with_indices
 aten::max_pool3d_with_indices.out
 aten::max_pool3d_with_indices_backward
diff --git a/test/expect/TestFXAPIBackwardCompatibility.test_function_back_compat-fx_backcompat_function_signatures.expect b/test/expect/TestFXAPIBackwardCompatibility.test_function_back_compat-fx_backcompat_function_signatures.expect
index 985b353f759d4..4900a484c353e 100644
--- a/test/expect/TestFXAPIBackwardCompatibility.test_function_back_compat-fx_backcompat_function_signatures.expect
+++ b/test/expect/TestFXAPIBackwardCompatibility.test_function_back_compat-fx_backcompat_function_signatures.expect
@@ -1,4 +1,4 @@
-torch.fx._symbolic_trace.Tracer.__init__(self, autowrap_modules: Tuple[Callable] = (<module math>,), autowrap_functions: Tuple[Callable, ...] = (,), param_shapes_constant: bool = False) -> None
+torch.fx._symbolic_trace.Tracer.__init__(self, autowrap_modules: Tuple[Callable] = (<module math>,), autowrap_functions: Tuple[Callable[..., Any], ...] = (,), param_shapes_constant: bool = False) -> None
 torch.fx._symbolic_trace.Tracer.call_module(self, m: torch.nn.modules.module.Module, forward: Callable[..., Any], args: Tuple[Any, ...], kwargs: Dict[str, Any]) -> Any
 torch.fx._symbolic_trace.Tracer.create_arg(self, a: Any) -> 'Argument'
 torch.fx._symbolic_trace.Tracer.get_fresh_qualname(self, prefix: str) -> str
@@ -6,7 +6,7 @@ torch.fx._symbolic_trace.Tracer.is_leaf_module(self, m: torch.nn.modules.module.
 torch.fx._symbolic_trace.Tracer.path_of_module(self, mod: torch.nn.modules.module.Module) -> str
 torch.fx._symbolic_trace.Tracer.trace(self, root: Union[torch.nn.modules.module.Module, Callable[..., Any]], concrete_args: Optional[Dict[str, Any]] = None) -> torch.fx.graph.Graph
 torch.fx._symbolic_trace.symbolic_trace(root: Union[torch.nn.modules.module.Module, Callable[..., Any]], concrete_args: Optional[Dict[str, Any]] = None) -> torch.fx.graph_module.GraphModule
-torch.fx._symbolic_trace.wrap(fn_or_name: Union[str, Callable])
+torch.fx._symbolic_trace.wrap(fn_or_name: torch.fx.node.Target) -> torch.fx.node.Target
 torch.fx.graph.Graph.__init__(self, owning_module: Optional[GraphModule] = None, tracer_cls: Optional[Type[Tracer]] = None, tracer_extras: Optional[Dict[str, Any]] = None)
 torch.fx.graph.Graph.call_function(self, the_function: Callable[..., Any], args: Optional[Tuple[Argument, ...]] = None, kwargs: Optional[Dict[str, Argument]] = None, type_expr: Optional[Any] = None, name: Optional[str] = None) -> torch.fx.node.Node
 torch.fx.graph.Graph.call_method(self, method_name: str, args: Optional[Tuple[Argument, ...]] = None, kwargs: Optional[Dict[str, Argument]] = None, type_expr: Optional[Any] = None) -> torch.fx.node.Node
@@ -31,20 +31,20 @@ torch.fx.graph_module.GraphModule.delete_submodule(self, target: str) -> bool
 torch.fx.graph_module.GraphModule.recompile(self) -> torch.fx.graph.PythonCode
 torch.fx.graph_module.reduce_graph_module(body: Dict[Any, Any], import_block: str) -> torch.nn.modules.module.Module
 torch.fx.graph_module.reduce_package_graph_module(importer: Callable, body: Dict[Any, Any], generated_module_name: str) -> torch.nn.modules.module.Module
-torch.fx.interpreter.Interpreter.__init__(self, module: torch.nn.modules.module.Module, garbage_collect_values: bool = True, graph: Optional[torch.fx.graph.Graph] = None)
-torch.fx.interpreter.Interpreter.boxed_run(self, args_list)
+torch.fx.interpreter.Interpreter.__init__(self, module: torch.nn.modules.module.Module, garbage_collect_values: bool = True, graph: Optional[torch.fx.graph.Graph] = None) -> None
+torch.fx.interpreter.Interpreter.boxed_run(self, args_list: List[Any]) -> Any
 torch.fx.interpreter.Interpreter.call_function(self, target: 'Target', args: Tuple[torch.fx.node.Argument, ...], kwargs: Dict[str, Any]) -> Any
 torch.fx.interpreter.Interpreter.call_method(self, target: 'Target', args: Tuple[torch.fx.node.Argument, ...], kwargs: Dict[str, Any]) -> Any
 torch.fx.interpreter.Interpreter.call_module(self, target: 'Target', args: Tuple[torch.fx.node.Argument, ...], kwargs: Dict[str, Any]) -> Any
-torch.fx.interpreter.Interpreter.fetch_args_kwargs_from_env(self, n: torch.fx.node.Node) -> Tuple[Tuple, Dict]
-torch.fx.interpreter.Interpreter.fetch_attr(self, target: str)
+torch.fx.interpreter.Interpreter.fetch_args_kwargs_from_env(self, n: torch.fx.node.Node) -> Tuple[Tuple[Any, ...], Dict[str, Any]]
+torch.fx.interpreter.Interpreter.fetch_attr(self, target: str) -> Any
 torch.fx.interpreter.Interpreter.get_attr(self, target: 'Target', args: Tuple[torch.fx.node.Argument, ...], kwargs: Dict[str, Any]) -> Any
 torch.fx.interpreter.Interpreter.map_nodes_to_values(self, args: torch.fx.node.Argument, n: torch.fx.node.Node) -> torch.fx.node.Argument
 torch.fx.interpreter.Interpreter.output(self, target: 'Target', args: Tuple[torch.fx.node.Argument, ...], kwargs: Dict[str, Any]) -> Any
 torch.fx.interpreter.Interpreter.placeholder(self, target: 'Target', args: Tuple[torch.fx.node.Argument, ...], kwargs: Dict[str, Any]) -> Any
-torch.fx.interpreter.Interpreter.run(self, *args, initial_env: Optional[Dict[torch.fx.node.Node, Any]] = None, enable_io_processing: bool = True) -> Any
+torch.fx.interpreter.Interpreter.run(self, *args: Any, initial_env: Optional[Dict[torch.fx.node.Node, Any]] = None, enable_io_processing: bool = True) -> Any
 torch.fx.interpreter.Interpreter.run_node(self, n: torch.fx.node.Node) -> Any
-torch.fx.interpreter.Transformer.__init__(self, module)
+torch.fx.interpreter.Transformer.__init__(self, module: torch.fx.graph_module.GraphModule) -> None
 torch.fx.interpreter.Transformer.call_function(self, target: 'Target', args: Tuple[torch.fx.node.Argument, ...], kwargs: Dict[str, Any]) -> Any
 torch.fx.interpreter.Transformer.call_module(self, target: 'Target', args: Tuple[torch.fx.node.Argument, ...], kwargs: Dict[str, Any]) -> Any
 torch.fx.interpreter.Transformer.get_attr(self, target: 'Target', args: Tuple[torch.fx.node.Argument, ...], kwargs: Dict[str, Any]) -> torch.fx.proxy.Proxy
@@ -63,15 +63,15 @@ torch.fx.node.map_aggregate(a: torch.fx.node.Argument, fn: Callable[[torch.fx.no
 torch.fx.node.map_arg(a: torch.fx.node.Argument, fn: Callable[[torch.fx.node.Node], torch.fx.node.Argument]) -> torch.fx.node.Argument
 torch.fx.passes.reinplace.reinplace(gm, *sample_args)
 torch.fx.passes.runtime_assert.insert_deferred_runtime_asserts(gm: torch.fx.graph_module.GraphModule, shape_env: Any, name: str, export: bool = False) -> None
-torch.fx.passes.split_module.split_module(m: torch.fx.graph_module.GraphModule, root_m: torch.nn.modules.module.Module, split_callback: Callable[[torch.fx.node.Node], int], qualname_map: Optional[Dict[str, str]] = None, keep_original_order: Optional[bool] = False, keep_original_node_name: Optional[bool] = False, keep_original_input_name: bool = True, partition_affix: Optional[str] = None)
-torch.fx.proxy.Attribute.__init__(self, root: torch.fx.proxy.Proxy, attr: str)
-torch.fx.proxy.Proxy.__init__(self, node: torch.fx.node.Node, tracer: 'Optional[TracerBase]' = None)
-torch.fx.proxy.Proxy.keys(self)
+torch.fx.passes.split_module.split_module(m: torch.fx.graph_module.GraphModule, root_m: torch.nn.modules.module.Module, split_callback: Callable[[torch.fx.node.Node], int], qualname_map: Optional[Dict[str, str]] = None, keep_original_order: Optional[bool] = False, keep_original_node_name: Optional[bool] = False, keep_original_input_name: bool = True, partition_affix: Optional[str] = None, tuple_return: bool = False)
+torch.fx.proxy.Attribute.__init__(self, root: torch.fx.proxy.Proxy, attr: str) -> None
+torch.fx.proxy.Proxy.__init__(self, node: torch.fx.node.Node, tracer: 'Optional[TracerBase]' = None) -> None
+torch.fx.proxy.Proxy.keys(self) -> 'Proxy'
 torch.fx.proxy.TracerBase.create_arg(self, a: Any) -> torch.fx.node.Argument
 torch.fx.proxy.TracerBase.create_node(self, kind: str, target: torch.fx.node.Target, args: Tuple[torch.fx.node.Argument, ...], kwargs: Dict[str, torch.fx.node.Argument], name: Optional[str] = None, type_expr: Optional[Any] = None) -> torch.fx.node.Node
-torch.fx.proxy.TracerBase.create_proxy(self, kind: str, target: torch.fx.node.Target, args: Tuple[Any, ...], kwargs: Dict[str, Any], name: Optional[str] = None, type_expr: Optional[Any] = None, proxy_factory_fn: Callable[[torch.fx.node.Node], Proxy] = None)
+torch.fx.proxy.TracerBase.create_proxy(self, kind: str, target: torch.fx.node.Target, args: Tuple[Any, ...], kwargs: Dict[str, Any], name: Optional[str] = None, type_expr: Optional[Any] = None, proxy_factory_fn: Optional[Callable[[torch.fx.node.Node], Proxy]] = None) -> 'Proxy'
 torch.fx.proxy.TracerBase.iter(self, obj: 'Proxy') -> Iterator
-torch.fx.proxy.TracerBase.keys(self, obj: 'Proxy') -> Any
+torch.fx.proxy.TracerBase.keys(self, obj: 'Proxy') -> 'Proxy'
 torch.fx.proxy.TracerBase.proxy(self, node: torch.fx.node.Node) -> 'Proxy'
 torch.fx.proxy.TracerBase.to_bool(self, obj: 'Proxy') -> bool
-torch.fx.subgraph_rewriter.replace_pattern(gm: torch.fx.graph_module.GraphModule, pattern: Union[Callable, torch.fx.graph_module.GraphModule], replacement: Union[Callable, torch.fx.graph_module.GraphModule]) -> List[torch.fx.subgraph_rewriter.Match]
+torch.fx.subgraph_rewriter.replace_pattern(gm: torch.fx.graph_module.GraphModule, pattern: Union[Callable[..., Any], torch.fx.graph_module.GraphModule], replacement: Union[Callable[..., Any], torch.fx.graph_module.GraphModule]) -> List[torch.fx.subgraph_rewriter.Match]
diff --git a/test/export/test_converter.py b/test/export/test_converter.py
index d281507196f7f..0cb48529635cf 100644
--- a/test/export/test_converter.py
+++ b/test/export/test_converter.py
@@ -2,7 +2,7 @@
 
 import unittest
 from collections import OrderedDict
-from typing import Any, Optional
+from typing import Any
 
 import torch
 import torch.utils._pytree as pytree
@@ -37,10 +37,10 @@ def _check_equal_ts_ep_converter(
         self,
         M,
         tracing_inputs,
-        option: Optional[list[str]] = None,
+        option: list[str] | None = None,
         check_persistent=False,
         lifted_tensor_constants=None,
-        runtime_inputs: Optional[list[Any]] = None,
+        runtime_inputs: list[Any] | None = None,
     ) -> list[ExportedProgram]:
         # By default, it tests both jit.trace and jit.script.
         if option is None:
diff --git a/test/export/test_draft_export.py b/test/export/test_draft_export.py
index 30147b1a7c42f..fefd35ad99ead 100644
--- a/test/export/test_draft_export.py
+++ b/test/export/test_draft_export.py
@@ -3,8 +3,10 @@
 import re
 import tempfile
 import unittest
+from unittest.mock import patch
 
 import torch
+from torch._dynamo.exc import UserError, UserErrorType
 from torch._subclasses.fake_tensor import FakeTensorMode
 from torch.export import Dim, draft_export, export
 from torch.export._draft_export import FailureType
@@ -32,6 +34,42 @@ def setUp(self):
     def tearDown(self):
         return
 
+    def test_retry_on_constraint_violation_uses_dim_auto(self):
+        class M(torch.nn.Module):
+            def forward(self, x):
+                return x + 1
+
+        dim = Dim("d0", min=2, max=16)
+        dynamic_shapes = {"x": {0: dim}}
+        calls = []
+
+        import torch.export._draft_export as draft_export_mod
+
+        real_export = draft_export_mod._export
+
+        def patched_export(*args, **kwargs):
+            calls.append(copy.deepcopy(kwargs["dynamic_shapes"]))
+            if len(calls) == 1:
+                raise UserError(
+                    UserErrorType.CONSTRAINT_VIOLATION,
+                    "mocked constraint violation",
+                )
+            return real_export(*args, **kwargs)
+
+        with patch("torch.export._draft_export._export", side_effect=patched_export):
+            draft_export(
+                M(),
+                (torch.randn(3),),
+                dynamic_shapes=dynamic_shapes,
+            )
+
+        self.assertEqual(len(calls), 2)
+        first_dim = calls[0]["x"][0]
+        second_dim = calls[1]["x"][0]
+        self.assertEqual(getattr(second_dim, "min", None), first_dim.min)
+        self.assertEqual(getattr(second_dim, "max", None), first_dim.max)
+        self.assertIn("AUTO", repr(second_dim))
+
     def test_missing_meta_kernel_custom_op_basic(self):
         with torch.library._scoped_library("mylib", "FRAGMENT"):
 
@@ -389,7 +427,7 @@ def forward(self, x, y):
             for node in _ep.graph.nodes:
                 if bindings := node.meta.get("unbacked_bindings"):
                     unbacked_binding_symbols.update(bindings.keys())
-            self.assertEqual(len(unbacked_binding_symbols), 2)
+            self.assertEqual(len(unbacked_binding_symbols), 1)
 
     def test_offsets(self):
         class M(torch.nn.Module):
@@ -405,7 +443,7 @@ def forward(self, x):
     def test_shape_failure(self):
         class M(torch.nn.Module):
             def forward(self, a):
-                assert a.shape[0] == 3
+                assert a.shape[0] == 3  # noqa: S101
                 return a * a
 
         inp = (torch.ones(3, 3),)
diff --git a/test/export/test_experimental.py b/test/export/test_experimental.py
index cec235992a963..949c704e5903d 100644
--- a/test/export/test_experimental.py
+++ b/test/export/test_experimental.py
@@ -1,5 +1,6 @@
 # Owner(s): ["oncall: export"]
 # flake8: noqa
+import contextlib
 import copy
 import types
 import unittest
@@ -14,11 +15,19 @@
     dynamo_graph_capture_for_export,
 )
 from torch._dynamo.test_case import run_tests, TestCase
-from torch._functorch.aot_autograd import aot_export_module
+from torch._export.utils import _compiling_state_context
+from torch._functorch.aot_autograd import (
+    aot_export_joint_with_descriptors,
+    aot_export_module,
+)
+from torch._guards import tracing as torch_tracing, TracingContext
 from torch.export import export
 from torch.export.experimental import _export_forward_backward, _sticky_export
 from torch.export.graph_signature import OutputKind
 from torch.testing import FileCheck
+from torch.testing._internal.common_device_type import (
+    IS_FLEX_ATTENTION_CUDA_PLATFORM_SUPPORTED,
+)
 from torch.testing._internal.common_utils import TEST_CUDA
 from torch.utils import _pytree as pytree
 
@@ -59,7 +68,11 @@ def __unflatten__(cls, leaves, ctx: tuple[pytree.TreeSpec, pytree.TreeSpec]):
         """Reconstructs from (leaves, ctx)."""
         output = cls()
         summary_spec, tensor_spec = ctx
-        assert len(leaves) == summary_spec.num_leaves + tensor_spec.num_leaves
+        if len(leaves) != summary_spec.num_leaves + tensor_spec.num_leaves:
+            raise AssertionError(
+                f"Expected {summary_spec.num_leaves + tensor_spec.num_leaves} leaves, "
+                f"got {len(leaves)}"
+            )
         output._summaries = pytree.tree_unflatten(
             leaves[: summary_spec.num_leaves], summary_spec
         )
@@ -439,6 +452,78 @@ def forward(self, x):
         ):
             _dynamo_graph_capture_for_export(module)(x)
 
+    @unittest.skipUnless(
+        IS_FLEX_ATTENTION_CUDA_PLATFORM_SUPPORTED and not torch.version.hip,
+        "Requires CUDA with SM >= 8.0, Triton, and not ROCm",
+    )
+    def test_aot_export_flex_attention_callable_mask_mod(self):
+        """Test flex_attention AOT export with callable class as mask_mod.
+
+        _MaskModWrapper must delegate __eq__ to callable objects for TreeSpec
+        comparison in AOTAutograd's PytreeThunk.set() (utils.py:162).
+        """
+        from torch._functorch.aot_autograd import aot_export_module
+        from torch.nn.attention.flex_attention import create_block_mask, flex_attention
+
+        _register_blockmask_pytree()
+
+        class ComposedMaskMod:
+            def __init__(self, *mask_fns):
+                self.mask_fns = mask_fns
+
+            def __call__(self, b, h, q, k):
+                result = True
+                for fn in self.mask_fns:
+                    result = result & fn(b, h, q, k)
+                return result
+
+            def __eq__(self, other):
+                if not isinstance(other, ComposedMaskMod):
+                    return NotImplemented
+                return self.mask_fns == other.mask_fns
+
+            def __hash__(self):
+                return hash(self.mask_fns)
+
+        def causal_mask(b, h, q, k):
+            return q >= k
+
+        class FlexAttentionModel(torch.nn.Module):
+            def __init__(self, embed_dim: int, num_heads: int):
+                super().__init__()
+                self.num_heads = num_heads
+                self.head_dim = embed_dim // num_heads
+                self.q_proj = torch.nn.Linear(embed_dim, embed_dim)
+                self.k_proj = torch.nn.Linear(embed_dim, embed_dim)
+                self.v_proj = torch.nn.Linear(embed_dim, embed_dim)
+
+            def forward(self, x):
+                B, L, D = x.shape
+                q = self.q_proj(x).view(B, L, self.num_heads, self.head_dim)
+                k = self.k_proj(x).view(B, L, self.num_heads, self.head_dim)
+                v = self.v_proj(x).view(B, L, self.num_heads, self.head_dim)
+                q, k, v = q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2)
+
+                mask_mod = ComposedMaskMod(causal_mask)
+                block_mask = create_block_mask(
+                    mask_mod, B=B, H=self.num_heads, Q_LEN=L, KV_LEN=L, device=x.device
+                )
+                out = flex_attention(q, k, v, block_mask=block_mask)
+                return (out.transpose(1, 2).contiguous().view(B, L, D),)
+
+        embed_dim, num_heads, seq_len = 64, 2, 128
+        model = FlexAttentionModel(embed_dim, num_heads).cuda()
+        x = torch.randn(1, seq_len, embed_dim, device="cuda")
+
+        gm, signature = aot_export_module(model, [x], trace_joint=False)
+
+        # aot_export_module flattens params/buffers into the graph signature
+        params = [p for p in model.parameters()]
+        out_eager = model(x)[0]
+        out_export = gm(*params, x)[0]
+        self.assertEqual(out_eager.shape, out_export.shape)
+        self.assertTrue(torch.allclose(out_eager, out_export, atol=1e-5))
+
     def test_joint_dynamic(self) -> None:
         from torch.export import Dim
 
@@ -715,6 +800,34 @@ def forward(self, dct, lst, bleh):
 
         self.assertEqual(res_export, res_eager)
 
+    def test_single_op(self):
+        # from torch._dynamo.functional_export import dynamo_graph_capture_for_export
+        x, y = (torch.randn(2, 3), torch.randn(2, 3))
+        graph = dynamo_graph_capture_for_export(torch.ops.aten.add.Tensor)(x, y)
+        self.assertExpectedInline(
+            graph.code.strip("\r\n "),
+            """\
+def forward(self, other):
+    _fn_args = (self, other)
+    self, L_self_ , L_other_ , = self._dynamo_bytecode_flatten(*_fn_args)
+    l_self_ = L_self_
+    l_other_ = L_other_
+    add_tensor = torch.ops.aten.add.Tensor(self = l_self_, other = l_other_, alpha = 1);  l_self_ = l_other_ = None
+    return self._dynamo_bytecode_unflatten((add_tensor,), _fn_args)""",
+        )
+        out = torch.empty(10)
+        graph = dynamo_graph_capture_for_export(torch.ops.aten.range.out)(0, 9, out=out)
+        self.assertExpectedInline(
+            graph.code.strip("\r\n "),
+            """\
+def forward(self, start, end, out):
+    _fn_args = (start, end, out)
+    L_out_ , = self._dynamo_bytecode_flatten(*_fn_args)
+    l_out_ = L_out_
+    range_out = torch.ops.aten.range.out(start = 0, end = 9, step = 1, out = l_out_);  l_out_ = None
+    return self._dynamo_bytecode_unflatten((range_out,), _fn_args)""",
+        )
+
     def test_export_leaf(self):
         class Foo(torch.nn.Module):
             def forward(self, x):
@@ -1327,8 +1440,14 @@ def forward(self, arg0_1):
 
     @unittest.skipIf(not TEST_CUDA, "CUDA not available")
     def test_aot_export_blockmask_closure_spec_mismatch(self):
-        """BlockMasks with same closure code but different captured values must
-        produce different TreeSpecs, so pytree won't confuse them."""
+        """BlockMasks with same closure structure produce equal TreeSpecs.
+
+        Closure values are extracted into pytree leaves, so two BlockMasks
+        whose mask_mod closures have the same code + structure but different
+        captured values have the same spec (values differ in the leaves, not
+        the context).  BlockMasks with different closure *structure* (e.g.
+        different code) must still produce different specs.
+        """
         from torch.nn.attention.flex_attention import create_block_mask
 
         _register_blockmask_pytree()
@@ -1355,8 +1474,154 @@ def fn(b, h, q, k):
 
         # Same closure code + same captured value -> same spec
         self.assertEqual(spec_a, spec_a_same)
-        # Same closure code + different captured value -> different spec
-        self.assertNotEqual(spec_a, spec_b)
+        # Same closure code + different captured value -> same spec
+        # (values are in the leaves, not the context)
+        self.assertEqual(spec_a, spec_b)
+
+        # Different closure *code* -> different spec
+        def different_mask(b, h, q, k):
+            return q > k
+
+        mask_c = create_block_mask(
+            different_mask, B=1, H=1, Q_LEN=64, KV_LEN=64, device="cuda"
+        )
+        _, spec_c = pytree.tree_flatten(mask_c)
+        self.assertNotEqual(spec_a, spec_c)
+
+    @unittest.skipIf(not TEST_CUDA, "CUDA not available")
+    def test_blockmask_and_masks_closure_extraction(self):
+        """and_masks closure tensors are recursively extracted into pytree leaves.
+
+        and_masks(fn1, fn2) returns a closure capturing a tuple of functions.
+        _extract_closure_pytree must recursively process these functions
+        (extracting their closure tensors) rather than emitting the functions
+        themselves as leaves, since functions are not supported export input
+        types.
+        """
+        from torch.nn.attention.flex_attention import (
+            and_masks,
+            BlockMask,
+            create_block_mask,
+        )
+
+        _register_blockmask_pytree()
+
+        def causal(b, h, q_idx, kv_idx):
+            return q_idx >= kv_idx
+
+        offset = torch.tensor(3, device="cuda")
+
+        def offset_mask(b, h, q_idx, kv_idx):
+            return q_idx >= kv_idx + offset
+
+        mask_mod = and_masks(causal, offset_mask)
+        block_mask = create_block_mask(
+            mask_mod, B=1, H=None, Q_LEN=128, KV_LEN=128, device="cuda"
+        )
+
+        leaves, spec = pytree.tree_flatten(block_mask)
+
+        # 8 regular BlockMask tensor attrs + offset extracted from
+        # offset_mask's closure
+        n_regular = len(BlockMask._TENSOR_ATTRS)
+        self.assertEqual(len(leaves), n_regular + 1)
+        self.assertTrue(all(isinstance(l, torch.Tensor) for l in leaves))
+        self.assertIs(leaves[n_regular], offset)
+
+        # Leaves must all pass check_user_input_output
+        from torch._dynamo.eval_frame import check_user_input_output
+        from torch._dynamo.exc import UserErrorType
+
+        check_user_input_output(leaves, UserErrorType.INVALID_INPUT)
+
+        # Round-trip: unflatten should reconstruct a working mask_mod
+        restored = pytree.tree_unflatten(leaves, spec)
+        self.assertTrue(callable(restored.mask_mod))
+
+    def test_aot_export_closure_buffer_mutation(self):
+        class Mod(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.register_buffer("buf", torch.zeros(()))
+
+            def forward(self, x):
+                self.buf.add_(x.sum())
+                return x.sin()
+
+        def make_closure(mod):
+            def fn(x):
+                mod._buffers["buf"].add_(x.sum())
+                return x.sin()
+
+            return fn
+
+        class Wrapper(torch.nn.Module):
+            def __init__(self, fn, mod):
+                super().__init__()
+                self._parameters = mod._parameters
+                self._buffers = mod._buffers
+                self._modules = mod._modules
+                self._fn = fn
+
+            def forward(self, x):
+                return self._fn(x)
+
+        def run_export(capture_fn):
+            mod = Mod()
+            wrapped = Wrapper(make_closure(mod), mod)
+            x = torch.randn(4)
+            gm = capture_fn(wrapped)(x)
+
+            with contextlib.ExitStack() as stack:
+                stack.enter_context(
+                    torch_tracing(
+                        gm.meta.get(
+                            "tracing_context", TracingContext(gm.meta["fake_mode"])
+                        )
+                    )
+                )
+                stack.enter_context(_compiling_state_context())
+                stack.enter_context(gm.meta["fake_mode"])
+
+                jd = aot_export_joint_with_descriptors(
+                    stack,
+                    gm,
+                    args=(x,),
+                    kwargs={},
+                    keep_inference_input_mutations=True,
+                    disable_functionalization=True,
+                )
+            return jd.graph_module, wrapped, x
+
+        # Verify Dynamo-captured graph mutates the buffer via closure
+        mod = Mod()
+        wrapped = Wrapper(make_closure(mod), mod)
+        x = torch.randn(4)
+        gm = dynamo_graph_capture_for_export(wrapped)(x)
+        wrapped.buf.zero_()
+        gm(x)
+        self.assertEqual(wrapped.buf, x.sum())
+
+        # Verify joint graphs from both APIs match
+        joint_public, _, _ = run_export(dynamo_graph_capture_for_export)
+        joint_private, _, _ = run_export(_dynamo_graph_capture_for_export)
+        self.assertEqual(
+            str(joint_public.code).strip(), str(joint_private.code).strip()
+        )
+
+        # Verify numerical correctness of both joint graphs against eager
+        mod = Mod()
+        x = torch.randn(4)
+        eager_out = mod(x)
+        eager_buf = mod.buf.clone()
+
+        for label, joint_gm in [("public", joint_public), ("private", joint_private)]:
+            buf_input = torch.zeros(())
+            (exported_out,) = joint_gm(buf_input, x)
+            self.assertEqual(exported_out, eager_out, msg=f"{label}: output mismatch")
+            self.assertEqual(
+                buf_input, eager_buf, msg=f"{label}: buffer mutation mismatch"
+            )
 
 
 if __name__ == "__main__":
diff --git a/test/export/test_export.py b/test/export/test_export.py
index 7b51cda83e79a..baebcde04da07 100755
--- a/test/export/test_export.py
+++ b/test/export/test_export.py
@@ -48,6 +48,7 @@
 from torch._higher_order_ops.scan import scan
 from torch._higher_order_ops.while_loop import while_loop
 from torch._inductor.compile_fx import split_const_gm
+from torch._library.opaque_object import _OPAQUE_TYPES_BY_NAME
 from torch._subclasses import FakeTensorMode
 from torch.export import default_decompositions, Dim, export, unflatten
 from torch.export._trace import (
@@ -127,6 +128,8 @@
 from torch.export import export
 
 
+device_type = acc.type if (acc := torch.accelerator.current_accelerator()) else "cpu"
+
 torch.library.define("testlib::returns_tensor_symint", "(Tensor x) -> (Tensor, SymInt)")
 torch.library.define(
     "testlib::foo",
@@ -646,7 +649,7 @@ def forward(self, x):
         # should not be copied to other nodes
         counter = 0
         for node in new_ep.graph.nodes:
-            if "custom" in node.meta:
+            if "quantization_tag" in node.meta.get("custom", {}):
                 counter += 1
                 self.assertTrue(node.meta["custom"]["quantization_tag"] == "foo")
                 self.assertTrue(node.target == torch.ops.aten.linear.default)
@@ -1174,192 +1177,196 @@ def forward(self, x):
         export(Foo(), inputs, dynamic_shapes=dynamic_shapes)
 
     def test_dynamic_lstm(self):
-        from torch.export._patches import (
-            register_gru_while_loop_decomposition,
-            register_lstm_while_loop_decomposition,
-        )
-
-        # Test 1: Basic single-layer LSTM with dynamic sequence length
-        seqlen = 32
-        bs = 16
-        h = 512
-
-        class LSTM(torch.nn.Module):
-            def __init__(self, h):
-                super().__init__()
-                self.lstm = torch.nn.LSTM(h, h)
-
-            def forward(self, x, h0, c0):
-                out, (_, _) = self.lstm(x, (h0, c0))
-                return out
-
-        x = torch.randn(seqlen, bs, h)
-        h0, c0 = torch.randn(1, bs, h), torch.randn(1, bs, h)
-
-        model = LSTM(h)
-        eager_out = model(x, h0, c0)
-        dynamic_shapes = {
-            "x": {0: Dim.DYNAMIC},
-            "h0": None,
-            "c0": None,
-        }
+        # Disable mkldnn so eager uses the same pure-PyTorch LSTM/GRU as the
+        # while-loop decomposition, avoiding numerical divergence from the
+        # fused mkldnn kernel.
+        with torch.backends.mkldnn.flags(enabled=False):
+            from torch.export._patches import (
+                register_gru_while_loop_decomposition,
+                register_lstm_while_loop_decomposition,
+            )
 
-        with register_lstm_while_loop_decomposition():
-            ep = export(model, (x, h0, c0), dynamic_shapes=dynamic_shapes)
-        ep_out = ep.module()(x, h0, c0)
-        self.assertEqual(eager_out, ep_out)
-        # test dynamic output with different sequence length
-        x_ = torch.randn(64, bs, h)
-        ep_out_dynamic = ep.module()(x_, h0, c0)
-        self.assertEqual(ep_out_dynamic, model(x_, h0, c0))
+            # Test 1: Basic single-layer LSTM with dynamic sequence length
+            seqlen = 32
+            bs = 16
+            h = 512
 
-        # Test 2: Bidirectional LSTM
-        class BiLSTM(torch.nn.Module):
-            def __init__(self, h):
-                super().__init__()
-                self.lstm = torch.nn.LSTM(h, h, bidirectional=True)
+            class LSTM(torch.nn.Module):
+                def __init__(self, h):
+                    super().__init__()
+                    self.lstm = torch.nn.LSTM(h, h)
 
-            def forward(self, x, h0, c0):
-                out, (_, _) = self.lstm(x, (h0, c0))
-                return out
+                def forward(self, x, h0, c0):
+                    out, (_, _) = self.lstm(x, (h0, c0))
+                    return out
 
-        h0_bi = torch.randn(2, bs, h)
-        c0_bi = torch.randn(2, bs, h)
-        model_bi = BiLSTM(h)
-        eager_out_bi = model_bi(x, h0_bi, c0_bi)
-        dynamic_shapes_bi = {
-            "x": {0: Dim.DYNAMIC},
-            "h0": None,
-            "c0": None,
-        }
+            x = torch.randn(seqlen, bs, h)
+            h0, c0 = torch.randn(1, bs, h), torch.randn(1, bs, h)
 
-        with register_lstm_while_loop_decomposition():
-            ep_bi = export(
-                model_bi, (x, h0_bi, c0_bi), dynamic_shapes=dynamic_shapes_bi
-            )
-        ep_out_bi = ep_bi.module()(x, h0_bi, c0_bi)
-        self.assertEqual(eager_out_bi, ep_out_bi)
-        # test with different sequence length
-        ep_out_bi_dynamic = ep_bi.module()(x_, h0_bi, c0_bi)
-        self.assertEqual(ep_out_bi_dynamic, model_bi(x_, h0_bi, c0_bi))
+            model = LSTM(h)
+            eager_out = model(x, h0, c0)
+            dynamic_shapes = {
+                "x": {0: Dim.DYNAMIC},
+                "h0": None,
+                "c0": None,
+            }
 
-        # Test 3: Multi-layer LSTM
-        class MultiLayerLSTM(torch.nn.Module):
-            def __init__(self, h, num_layers=2):
-                super().__init__()
-                self.lstm = torch.nn.LSTM(h, h, num_layers=num_layers)
+            with register_lstm_while_loop_decomposition():
+                ep = export(model, (x, h0, c0), dynamic_shapes=dynamic_shapes)
+            ep_out = ep.module()(x, h0, c0)
+            self.assertEqual(eager_out, ep_out)
+            # test dynamic output with different sequence length
+            x_ = torch.randn(64, bs, h)
+            ep_out_dynamic = ep.module()(x_, h0, c0)
+            self.assertEqual(ep_out_dynamic, model(x_, h0, c0))
+
+            # Test 2: Bidirectional LSTM
+            class BiLSTM(torch.nn.Module):
+                def __init__(self, h):
+                    super().__init__()
+                    self.lstm = torch.nn.LSTM(h, h, bidirectional=True)
+
+                def forward(self, x, h0, c0):
+                    out, (_, _) = self.lstm(x, (h0, c0))
+                    return out
+
+            h0_bi = torch.randn(2, bs, h)
+            c0_bi = torch.randn(2, bs, h)
+            model_bi = BiLSTM(h)
+            eager_out_bi = model_bi(x, h0_bi, c0_bi)
+            dynamic_shapes_bi = {
+                "x": {0: Dim.DYNAMIC},
+                "h0": None,
+                "c0": None,
+            }
 
-            def forward(self, x, h0, c0):
-                out, (_, _) = self.lstm(x, (h0, c0))
-                return out
+            with register_lstm_while_loop_decomposition():
+                ep_bi = export(
+                    model_bi, (x, h0_bi, c0_bi), dynamic_shapes=dynamic_shapes_bi
+                )
+            ep_out_bi = ep_bi.module()(x, h0_bi, c0_bi)
+            self.assertEqual(eager_out_bi, ep_out_bi)
+            # test with different sequence length
+            ep_out_bi_dynamic = ep_bi.module()(x_, h0_bi, c0_bi)
+            self.assertEqual(ep_out_bi_dynamic, model_bi(x_, h0_bi, c0_bi))
+
+            # Test 3: Multi-layer LSTM
+            class MultiLayerLSTM(torch.nn.Module):
+                def __init__(self, h, num_layers=2):
+                    super().__init__()
+                    self.lstm = torch.nn.LSTM(h, h, num_layers=num_layers)
+
+                def forward(self, x, h0, c0):
+                    out, (_, _) = self.lstm(x, (h0, c0))
+                    return out
+
+            num_layers = 2
+            h0_multi = torch.randn(num_layers, bs, h)
+            c0_multi = torch.randn(num_layers, bs, h)
+            model_multi = MultiLayerLSTM(h, num_layers)
+            eager_out_multi = model_multi(x, h0_multi, c0_multi)
+            dynamic_shapes_multi = {
+                "x": {0: Dim.DYNAMIC},
+                "h0": None,
+                "c0": None,
+            }
 
-        num_layers = 2
-        h0_multi = torch.randn(num_layers, bs, h)
-        c0_multi = torch.randn(num_layers, bs, h)
-        model_multi = MultiLayerLSTM(h, num_layers)
-        eager_out_multi = model_multi(x, h0_multi, c0_multi)
-        dynamic_shapes_multi = {
-            "x": {0: Dim.DYNAMIC},
-            "h0": None,
-            "c0": None,
-        }
+            with register_lstm_while_loop_decomposition():
+                ep_multi = export(
+                    model_multi,
+                    (x, h0_multi, c0_multi),
+                    dynamic_shapes=dynamic_shapes_multi,
+                )
+            ep_out_multi = ep_multi.module()(x, h0_multi, c0_multi)
+            self.assertEqual(eager_out_multi, ep_out_multi)
+            ep_out_multi_dynamic = ep_multi.module()(x_, h0_multi, c0_multi)
+            self.assertEqual(ep_out_multi_dynamic, model_multi(x_, h0_multi, c0_multi))
+
+            # Test 4: batch_first=True
+            class BatchFirstLSTM(torch.nn.Module):
+                def __init__(self, h):
+                    super().__init__()
+                    self.lstm = torch.nn.LSTM(h, h, batch_first=True)
+
+                def forward(self, x, h0, c0):
+                    out, (_, _) = self.lstm(x, (h0, c0))
+                    return out
+
+            x_batch_first = torch.randn(bs, seqlen, h)
+            model_batch_first = BatchFirstLSTM(h)
+            eager_out_batch_first = model_batch_first(x_batch_first, h0, c0)
+            dynamic_shapes_batch_first = {
+                "x": {1: Dim.DYNAMIC},  # dynamic dimension is now dim 1
+                "h0": None,
+                "c0": None,
+            }
 
-        with register_lstm_while_loop_decomposition():
-            ep_multi = export(
-                model_multi,
-                (x, h0_multi, c0_multi),
-                dynamic_shapes=dynamic_shapes_multi,
+            with register_lstm_while_loop_decomposition():
+                ep_batch_first = export(
+                    model_batch_first,
+                    (x_batch_first, h0, c0),
+                    dynamic_shapes=dynamic_shapes_batch_first,
+                )
+            ep_out_batch_first = ep_batch_first.module()(x_batch_first, h0, c0)
+            self.assertEqual(eager_out_batch_first, ep_out_batch_first)
+            x_batch_first_dynamic = torch.randn(bs, 64, h)
+            ep_out_batch_first_dynamic = ep_batch_first.module()(
+                x_batch_first_dynamic, h0, c0
             )
-        ep_out_multi = ep_multi.module()(x, h0_multi, c0_multi)
-        self.assertEqual(eager_out_multi, ep_out_multi)
-        ep_out_multi_dynamic = ep_multi.module()(x_, h0_multi, c0_multi)
-        self.assertEqual(ep_out_multi_dynamic, model_multi(x_, h0_multi, c0_multi))
-
-        # Test 4: batch_first=True
-        class BatchFirstLSTM(torch.nn.Module):
-            def __init__(self, h):
-                super().__init__()
-                self.lstm = torch.nn.LSTM(h, h, batch_first=True)
-
-            def forward(self, x, h0, c0):
-                out, (_, _) = self.lstm(x, (h0, c0))
-                return out
-
-        x_batch_first = torch.randn(bs, seqlen, h)
-        model_batch_first = BatchFirstLSTM(h)
-        eager_out_batch_first = model_batch_first(x_batch_first, h0, c0)
-        dynamic_shapes_batch_first = {
-            "x": {1: Dim.DYNAMIC},  # dynamic dimension is now dim 1
-            "h0": None,
-            "c0": None,
-        }
-
-        with register_lstm_while_loop_decomposition():
-            ep_batch_first = export(
-                model_batch_first,
-                (x_batch_first, h0, c0),
-                dynamic_shapes=dynamic_shapes_batch_first,
+            self.assertEqual(
+                ep_out_batch_first_dynamic,
+                model_batch_first(x_batch_first_dynamic, h0, c0),
             )
-        ep_out_batch_first = ep_batch_first.module()(x_batch_first, h0, c0)
-        self.assertEqual(eager_out_batch_first, ep_out_batch_first)
-        x_batch_first_dynamic = torch.randn(bs, 64, h)
-        ep_out_batch_first_dynamic = ep_batch_first.module()(
-            x_batch_first_dynamic, h0, c0
-        )
-        self.assertEqual(
-            ep_out_batch_first_dynamic,
-            model_batch_first(x_batch_first_dynamic, h0, c0),
-        )
 
-        # Test 5: GRU with dynamic sequence length
-        class GRU(torch.nn.Module):
-            def __init__(self, h):
-                super().__init__()
-                self.gru = torch.nn.GRU(h, h)
+            # Test 5: GRU with dynamic sequence length
+            class GRU(torch.nn.Module):
+                def __init__(self, h):
+                    super().__init__()
+                    self.gru = torch.nn.GRU(h, h)
 
-            def forward(self, x, h0):
-                out, _ = self.gru(x, h0)
-                return out
+                def forward(self, x, h0):
+                    out, _ = self.gru(x, h0)
+                    return out
 
-        model_gru = GRU(h)
-        eager_out_gru = model_gru(x, h0)
-        dynamic_shapes_gru = {
-            "x": {0: Dim.DYNAMIC},
-            "h0": None,
-        }
+            model_gru = GRU(h)
+            eager_out_gru = model_gru(x, h0)
+            dynamic_shapes_gru = {
+                "x": {0: Dim.DYNAMIC},
+                "h0": None,
+            }
 
-        with register_gru_while_loop_decomposition():
-            ep_gru = export(model_gru, (x, h0), dynamic_shapes=dynamic_shapes_gru)
-        ep_out_gru = ep_gru.module()(x, h0)
-        self.assertEqual(eager_out_gru, ep_out_gru)
-        ep_out_gru_dynamic = ep_gru.module()(x_, h0)
-        self.assertEqual(ep_out_gru_dynamic, model_gru(x_, h0))
+            with register_gru_while_loop_decomposition():
+                ep_gru = export(model_gru, (x, h0), dynamic_shapes=dynamic_shapes_gru)
+            ep_out_gru = ep_gru.module()(x, h0)
+            self.assertEqual(eager_out_gru, ep_out_gru)
+            ep_out_gru_dynamic = ep_gru.module()(x_, h0)
+            self.assertEqual(ep_out_gru_dynamic, model_gru(x_, h0))
 
-        # Test 6: Bidirectional GRU
-        class BiGRU(torch.nn.Module):
-            def __init__(self, h):
-                super().__init__()
-                self.gru = torch.nn.GRU(h, h, bidirectional=True)
+            # Test 6: Bidirectional GRU
+            class BiGRU(torch.nn.Module):
+                def __init__(self, h):
+                    super().__init__()
+                    self.gru = torch.nn.GRU(h, h, bidirectional=True)
 
-            def forward(self, x, h0):
-                out, _ = self.gru(x, h0)
-                return out
+                def forward(self, x, h0):
+                    out, _ = self.gru(x, h0)
+                    return out
 
-        model_bigru = BiGRU(h)
-        eager_out_bigru = model_bigru(x, h0_bi)
-        dynamic_shapes_bigru = {
-            "x": {0: Dim.DYNAMIC},
-            "h0": None,
-        }
+            model_bigru = BiGRU(h)
+            eager_out_bigru = model_bigru(x, h0_bi)
+            dynamic_shapes_bigru = {
+                "x": {0: Dim.DYNAMIC},
+                "h0": None,
+            }
 
-        with register_gru_while_loop_decomposition():
-            ep_bigru = export(
-                model_bigru, (x, h0_bi), dynamic_shapes=dynamic_shapes_bigru
-            )
-        ep_out_bigru = ep_bigru.module()(x, h0_bi)
-        self.assertEqual(eager_out_bigru, ep_out_bigru)
-        ep_out_bigru_dynamic = ep_bigru.module()(x_, h0_bi)
-        self.assertEqual(ep_out_bigru_dynamic, model_bigru(x_, h0_bi))
+            with register_gru_while_loop_decomposition():
+                ep_bigru = export(
+                    model_bigru, (x, h0_bi), dynamic_shapes=dynamic_shapes_bigru
+                )
+            ep_out_bigru = ep_bigru.module()(x, h0_bi)
+            self.assertEqual(eager_out_bigru, ep_out_bigru)
+            ep_out_bigru_dynamic = ep_bigru.module()(x_, h0_bi)
+            self.assertEqual(ep_out_bigru_dynamic, model_bigru(x_, h0_bi))
 
     @testing.expectedFailureStrictV2
     def test_no_tensor_computation(self):
@@ -1435,7 +1442,8 @@ def forward(self, *args, **kwargs):
         class CacheModule(BaseModule):
             def __init__(self, cache: torch.Tensor):
                 super().__init__()
-                assert cache.ndim == 3
+                if cache.ndim != 3:
+                    raise AssertionError(f"Expected cache.ndim == 3, got {cache.ndim}")
                 self.cache = torch.nn.Parameter(cache, requires_grad=False)
 
             def forward(self, x: torch.Tensor) -> torch.Tensor:
@@ -2456,6 +2464,49 @@ def auto_dynamic_shapes_from_args(args):  # pyre-ignore
             dynamic_shapes=auto_dynamic_shapes_from_args(sample_input),
         ).run_decompositions({})
 
+    def test_where_decomp_non_strict_inference_mode_dynamic_shapes(self):
+        class TestModule(torch.nn.Module):
+            def forward(self, x):
+                return torch.where(x > 0)
+
+        test_module = TestModule()
+        sample_input = (torch.rand(2, 10),)
+        dynamic_shapes = ({0: Dim("batch_size", max=100)},)
+
+        with torch.inference_mode():
+            ep = torch.export.export(
+                test_module,
+                sample_input,
+                strict=False,
+                dynamic_shapes=dynamic_shapes,
+            ).run_decompositions({})
+
+        self.assertEqual(ep.module()(*sample_input), test_module(*sample_input))
+
+        larger_input = (torch.rand(4, 10),)
+        self.assertEqual(ep.module()(*larger_input), test_module(*larger_input))
+
+    def test_where_decomp_non_bool_input(self):
+        class TestModule(torch.nn.Module):
+            def forward(self, x):
+                return torch.where(x)
+
+        test_module = TestModule()
+        sample_input = (torch.tensor([[0.0, 1.0], [2.0, 0.0]]),)
+        dynamic_shapes = ({0: Dim("batch_size", max=100)},)
+
+        ep = torch.export.export(
+            test_module,
+            sample_input,
+            strict=False,
+            dynamic_shapes=dynamic_shapes,
+        ).run_decompositions({})
+
+        self.assertEqual(ep.module()(*sample_input), test_module(*sample_input))
+
+        larger_input = (torch.tensor([[0.0, 3.0], [4.0, 0.0], [5.0, 6.0]]),)
+        self.assertEqual(ep.module()(*larger_input), test_module(*larger_input))
+
     def test_basic_non_strict_fake_tensor(self):
         class Basic(torch.nn.Module):
             def __init__(self) -> None:
@@ -3667,9 +3718,9 @@ def forward(self, x):
     %add : [num_users=1] = call_function[target=torch.ops.aten.add.Tensor](args = (%mul, %p_p2), kwargs = {})
     %add_1 : [num_users=1] = call_function[target=torch.ops.aten.add.Tensor](args = (%add, %b_b1), kwargs = {})
     %sum_1 : [num_users=1] = call_function[target=torch.ops.aten.sum.default](args = (%add_1,), kwargs = {})
-    %access_subclass_inner_tensor_default_64 : [num_users=1] = call_function[target=torch.ops.export.access_subclass_inner_tensor.default](args = (%sum_1, a), kwargs = {})
-    %access_subclass_inner_tensor_default_69 : [num_users=1] = call_function[target=torch.ops.export.access_subclass_inner_tensor.default](args = (%access_subclass_inner_tensor_default_64, b), kwargs = {})
-    %add_2 : [num_users=1] = call_function[target=torch.ops.aten.add.Tensor](args = (%x, %access_subclass_inner_tensor_default_69), kwargs = {})
+    %access_subclass_inner_tensor_default_56 : [num_users=1] = call_function[target=torch.ops.export.access_subclass_inner_tensor.default](args = (%sum_1, a), kwargs = {})
+    %access_subclass_inner_tensor_default_61 : [num_users=1] = call_function[target=torch.ops.export.access_subclass_inner_tensor.default](args = (%access_subclass_inner_tensor_default_56, b), kwargs = {})
+    %add_2 : [num_users=1] = call_function[target=torch.ops.aten.add.Tensor](args = (%x, %access_subclass_inner_tensor_default_61), kwargs = {})
     return (add_2,)""",
         )
         ep = export(m, (ref_x,))
@@ -3721,9 +3772,9 @@ def forward(self, x):
     %add : [num_users=1] = call_function[target=torch.ops.aten.add.Tensor](args = (%mul, %p_bar_p2), kwargs = {})
     %add_1 : [num_users=1] = call_function[target=torch.ops.aten.add.Tensor](args = (%add, %b_bar_b1), kwargs = {})
     %sum_1 : [num_users=1] = call_function[target=torch.ops.aten.sum.default](args = (%add_1,), kwargs = {})
-    %access_subclass_inner_tensor_default_64 : [num_users=1] = call_function[target=torch.ops.export.access_subclass_inner_tensor.default](args = (%sum_1, a), kwargs = {})
-    %access_subclass_inner_tensor_default_69 : [num_users=1] = call_function[target=torch.ops.export.access_subclass_inner_tensor.default](args = (%access_subclass_inner_tensor_default_64, b), kwargs = {})
-    %add_2 : [num_users=1] = call_function[target=torch.ops.aten.add.Tensor](args = (%x, %access_subclass_inner_tensor_default_69), kwargs = {})
+    %access_subclass_inner_tensor_default_56 : [num_users=1] = call_function[target=torch.ops.export.access_subclass_inner_tensor.default](args = (%sum_1, a), kwargs = {})
+    %access_subclass_inner_tensor_default_61 : [num_users=1] = call_function[target=torch.ops.export.access_subclass_inner_tensor.default](args = (%access_subclass_inner_tensor_default_56, b), kwargs = {})
+    %add_2 : [num_users=1] = call_function[target=torch.ops.aten.add.Tensor](args = (%x, %access_subclass_inner_tensor_default_61), kwargs = {})
     return (add_2,)""",
         )
         ep = export(m, (ref_x,))
@@ -3759,9 +3810,9 @@ def forward(self, x):
     %mul : [num_users=1] = call_function[target=torch.ops.aten.mul.Tensor](args = (%p_p1, 2), kwargs = {})
     %add : [num_users=1] = call_function[target=torch.ops.aten.add.Tensor](args = (%mul, %p_p2), kwargs = {})
     %add_1 : [num_users=1] = call_function[target=torch.ops.aten.add.Tensor](args = (%add, 4), kwargs = {})
-    %access_subclass_inner_tensor_default_10 : [num_users=1] = call_function[target=torch.ops.export.access_subclass_inner_tensor.default](args = (%add_1, elem), kwargs = {})
-    %access_subclass_inner_tensor_default_13 : [num_users=1] = call_function[target=torch.ops.export.access_subclass_inner_tensor.default](args = (%access_subclass_inner_tensor_default_10, elem), kwargs = {})
-    %add_2 : [num_users=1] = call_function[target=torch.ops.aten.add.Tensor](args = (%x, %access_subclass_inner_tensor_default_13), kwargs = {})
+    %access_subclass_inner_tensor_default_8 : [num_users=1] = call_function[target=torch.ops.export.access_subclass_inner_tensor.default](args = (%add_1, elem), kwargs = {})
+    %access_subclass_inner_tensor_default_11 : [num_users=1] = call_function[target=torch.ops.export.access_subclass_inner_tensor.default](args = (%access_subclass_inner_tensor_default_8, elem), kwargs = {})
+    %add_2 : [num_users=1] = call_function[target=torch.ops.aten.add.Tensor](args = (%x, %access_subclass_inner_tensor_default_11), kwargs = {})
     return (add_2,)""",
         )
         ep = export(m, (ref_x,))
@@ -3835,13 +3886,13 @@ def forward(self, x):
     %x : [num_users=1] = placeholder[target=x]
     %mul : [num_users=1] = call_function[target=torch.ops.aten.mul.Tensor](args = (%p_p1, 2), kwargs = {})
     %add : [num_users=2] = call_function[target=torch.ops.aten.add.Tensor](args = (%mul, %p_p2), kwargs = {})
-    %access_subclass_inner_tensor_default_18 : [num_users=1] = call_function[target=torch.ops.export.access_subclass_inner_tensor.default](args = (%add, a), kwargs = {})
-    %access_subclass_inner_tensor_default_21 : [num_users=1] = call_function[target=torch.ops.export.access_subclass_inner_tensor.default](args = (%access_subclass_inner_tensor_default_18, elem), kwargs = {})
-    %add_1 : [num_users=1] = call_function[target=torch.ops.aten.add.Tensor](args = (%add, %access_subclass_inner_tensor_default_21), kwargs = {})
+    %access_subclass_inner_tensor_default_16 : [num_users=1] = call_function[target=torch.ops.export.access_subclass_inner_tensor.default](args = (%add, a), kwargs = {})
+    %access_subclass_inner_tensor_default_19 : [num_users=1] = call_function[target=torch.ops.export.access_subclass_inner_tensor.default](args = (%access_subclass_inner_tensor_default_16, elem), kwargs = {})
+    %add_1 : [num_users=1] = call_function[target=torch.ops.aten.add.Tensor](args = (%add, %access_subclass_inner_tensor_default_19), kwargs = {})
     %add_2 : [num_users=1] = call_function[target=torch.ops.aten.add.Tensor](args = (%add_1, 4), kwargs = {})
-    %access_subclass_inner_tensor_default_25 : [num_users=1] = call_function[target=torch.ops.export.access_subclass_inner_tensor.default](args = (%add_2, a), kwargs = {})
-    %access_subclass_inner_tensor_default_28 : [num_users=1] = call_function[target=torch.ops.export.access_subclass_inner_tensor.default](args = (%access_subclass_inner_tensor_default_25, elem), kwargs = {})
-    %add_3 : [num_users=1] = call_function[target=torch.ops.aten.add.Tensor](args = (%x, %access_subclass_inner_tensor_default_28), kwargs = {})
+    %access_subclass_inner_tensor_default_23 : [num_users=1] = call_function[target=torch.ops.export.access_subclass_inner_tensor.default](args = (%add_2, a), kwargs = {})
+    %access_subclass_inner_tensor_default_26 : [num_users=1] = call_function[target=torch.ops.export.access_subclass_inner_tensor.default](args = (%access_subclass_inner_tensor_default_23, elem), kwargs = {})
+    %add_3 : [num_users=1] = call_function[target=torch.ops.aten.add.Tensor](args = (%x, %access_subclass_inner_tensor_default_26), kwargs = {})
     return (add_3,)""",
         )
         ep = export(m, (ref_x,))
@@ -3876,9 +3927,9 @@ def forward(self, x):
     %mul : [num_users=1] = call_function[target=torch.ops.aten.mul.Tensor](args = (%p_p1, 2), kwargs = {})
     %add : [num_users=1] = call_function[target=torch.ops.aten.add.Tensor](args = (%x, %mul), kwargs = {})
     %add_1 : [num_users=1] = call_function[target=torch.ops.aten.add.Tensor](args = (%add, %p_p2), kwargs = {})
-    %access_subclass_inner_tensor_default_10 : [num_users=1] = call_function[target=torch.ops.export.access_subclass_inner_tensor.default](args = (%add_1, elem), kwargs = {})
-    %access_subclass_inner_tensor_default_13 : [num_users=1] = call_function[target=torch.ops.export.access_subclass_inner_tensor.default](args = (%access_subclass_inner_tensor_default_10, elem), kwargs = {})
-    %add_2 : [num_users=1] = call_function[target=torch.ops.aten.add.Tensor](args = (%access_subclass_inner_tensor_default_13, 4), kwargs = {})
+    %access_subclass_inner_tensor_default_8 : [num_users=1] = call_function[target=torch.ops.export.access_subclass_inner_tensor.default](args = (%add_1, elem), kwargs = {})
+    %access_subclass_inner_tensor_default_11 : [num_users=1] = call_function[target=torch.ops.export.access_subclass_inner_tensor.default](args = (%access_subclass_inner_tensor_default_8, elem), kwargs = {})
+    %add_2 : [num_users=1] = call_function[target=torch.ops.aten.add.Tensor](args = (%access_subclass_inner_tensor_default_11, 4), kwargs = {})
     return (add_2,)""",
         )
         ep = export(m, (ref_x,))
@@ -4090,7 +4141,7 @@ def forward(self, x):
 
         class Bar(torch.nn.Module):
             def forward(self, x):
-                assert x.shape[0] <= 32
+                assert x.shape[0] <= 32  # noqa: S101
                 return x + 2
 
         # static specialization
@@ -4172,7 +4223,7 @@ def test_dim_hint_range_violations(self):
         class Foo(torch.nn.Module):
             def forward(self, xs):
                 x, y = xs["data"][0]
-                assert y.shape[0] <= 32
+                assert y.shape[0] <= 32  # noqa: S101
                 return x[6:], y + 2
 
         x, y = torch.randn(8), torch.randn(8)
@@ -4213,8 +4264,8 @@ def forward(self, x):
         # multiple conflicts
         class Moo(torch.nn.Module):
             def forward(self, x, y):
-                assert x.shape[0] <= 32
-                assert y.shape[0] >= 128
+                assert x.shape[0] <= 32  # noqa: S101
+                assert y.shape[0] >= 128  # noqa: S101
                 return x + 2, y + 2
 
         inps = (torch.randn(16), torch.randn(256))
@@ -4301,11 +4352,17 @@ def _tuple_contains(_tuple, val):
             for node in graph.nodes:
                 # check node.users
                 for user in node.users.keys():
-                    assert _tuple_contains(user.args, node)
+                    if not _tuple_contains(user.args, node):
+                        raise AssertionError(
+                            f"node {node} not found in user {user}'s args"
+                        )
                 # check node.args
                 for arg in node.args:
                     if isinstance(arg, torch.fx.Node):
-                        assert _tuple_contains(arg.users, node)
+                        if not _tuple_contains(arg.users, node):
+                            raise AssertionError(
+                                f"node {node} not found in arg {arg}'s users"
+                            )
 
         # check set grad enabled
         class SetGradCase(torch.nn.Module):
@@ -5819,7 +5876,10 @@ def forward(self, x, y):
         inp = (torch.randn(4), torch.randn(4))
         dynamic_shapes = ((torch.export.Dim.DYNAMIC,), (torch.export.Dim.DYNAMIC,))
         ep = export(m, inp, dynamic_shapes=dynamic_shapes)
-        assert len(ep.range_constraints) == 1
+        if len(ep.range_constraints) != 1:
+            raise AssertionError(
+                f"Expected 1 range constraint, got {len(ep.range_constraints)}"
+            )
         vr = next(iter(ep.range_constraints.values()))
         self.assertEqual(vr.lower, 3)
 
@@ -6891,8 +6951,8 @@ def test_dont_duck_size_for_auto_dynamic(self):
         class Foo(torch.nn.Module):
             def forward(self, x, y):
                 # x: [s0, s1], y: [s0 + 1, 4]
-                assert y.shape[1] == 4
-                assert x.shape[0] == y.shape[0] - 1
+                assert y.shape[1] == 4  # noqa: S101
+                assert x.shape[0] == y.shape[0] - 1  # noqa: S101
                 return x * 2, y * 2
 
         # duck sizing would make all static based on these sample inputs
@@ -7206,16 +7266,16 @@ def test_more_multidimensional_slicing(self):
         # Output: a 3-tuple of indices
         @torch.library.custom_op("demo::indices3d", mutates_args=())
         def indices3d(t: torch.Tensor, x: torch.Tensor) -> tuple[int, int, int]:
-            assert t.ndim == 3
-            assert x.ndim == 1 and x.shape[0] == 3
+            assert t.ndim == 3  # noqa: S101
+            assert x.ndim == 1 and x.shape[0] == 3  # noqa: S101
             return tuple(x[i].item() for i in range(3))
 
         # The meta-kernel for this op constrains the indices in x
         # to be within bounds of t via torch._checks.
         @torch.library.register_fake("demo::indices3d")
         def _(t, x):
-            assert t.ndim == 3
-            assert x.ndim == 1 and x.shape[0] == 3
+            assert t.ndim == 3  # noqa: S101
+            assert x.ndim == 1 and x.shape[0] == 3  # noqa: S101
             sizes = tuple(torch.library.get_ctx().new_dynamic_size() for i in range(3))
             for i, size in enumerate(sizes):
                 torch._check(size >= 0)
@@ -10550,6 +10610,59 @@ def forward(self, x):
         )
         self.assertTrue(torch.allclose(core_aten_ep.module()(*inp), m(*inp)))
 
+    def test_export_decomps_isin_dynamic(self):
+        class M(torch.nn.Module):
+            def forward(self, elements, test_elements):
+                return torch.isin(elements, test_elements)
+
+        m = M()
+        elements = torch.tensor([1, 2, 3, 4, 5])
+        test_elements = torch.tensor([2, 4])
+        inp = (elements, test_elements)
+
+        ep = export(
+            m,
+            inp,
+            dynamic_shapes={
+                "elements": {0: Dim("n_elements")},
+                "test_elements": {0: Dim("n_test")},
+            },
+        )
+        decomposed = ep.run_decompositions()
+
+        self.assertEqual(decomposed.module()(*inp), m(*inp))
+
+    def test_where_broadcast_preserves_symint(self):
+        import torch.fx.experimental._config as config
+        from torch._dynamo.source import ConstantSource
+        from torch._subclasses.fake_tensor import FakeTensorMode
+        from torch.fx.experimental.symbolic_shapes import DimDynamic, ShapeEnv
+
+        with config.patch(backed_size_oblivious=True):
+            shape_env = ShapeEnv(specialize_zero_one=False)
+            mode = FakeTensorMode(shape_env=shape_env)
+            with mode:
+                s0 = shape_env.create_symintnode(
+                    shape_env.create_symbol(
+                        val=1,
+                        source=ConstantSource("s0"),
+                        dynamic_dim=DimDynamic.DYNAMIC,
+                        do_not_specialize_zero_one=True,
+                    ),
+                    hint=1,
+                )
+                t = torch.empty((s0, 8), device="meta")
+                cond = torch.empty((s0, 8), device="meta", dtype=torch.bool)
+                fill = torch.empty((1,), device="meta")
+
+                result = torch.ops.aten.where.self(cond, fill, t)
+
+            self.assertIsInstance(
+                result.shape[0],
+                torch.SymInt,
+                f"where output dim 0 should be symbolic but got {result.shape[0]}",
+            )
+
     def test_nonzero_2(self):
         class Module(torch.nn.Module):
             def forward(self, x):
@@ -10988,7 +11101,6 @@ def forward(self, x):
             export_res = decomposed_ep.module()(x)
             self.assertTrue(export_res.size() == exp_res.size())
 
-    @skipIfXpu
     def test_export_with_fake_tensor_inputs_on_cuda_devices(self):
         fake_mode = torch._subclasses.fake_tensor.FakeTensorMode()
 
@@ -11007,9 +11119,9 @@ def forward(self, x):
             model = Model()
 
         # Manually set the fake_device of fake tensors.
-        x.fake_device = torch.device("cuda:0")
+        x.fake_device = torch.device(f"{device_type}:0")
         for n, p in model.named_parameters():
-            p.fake_device = torch.device("cuda:0")
+            p.fake_device = torch.device(f"{device_type}:0")
 
         # Need to set all the requires_grad of tensors to False, because fake_tensor with CUDA device
         # doesn't quite work well with aot_autograd right now due to some logic fails
@@ -11789,8 +11901,10 @@ def forward(self, x):
         ep = export(n0, inp)
         epm = ep.module()
         ufm = torch.export.unflatten(ep)
-        assert torch.allclose(epm(*inp), eager)
-        assert torch.allclose(ufm(*inp), eager)
+        if not torch.allclose(epm(*inp), eager):
+            raise AssertionError("epm output does not match eager output")
+        if not torch.allclose(ufm(*inp), eager):
+            raise AssertionError("ufm output does not match eager output")
 
     def test_unflatten_random_dag_mutating_buf_4(self):
         class N3(torch.nn.Module):
@@ -11847,8 +11961,10 @@ def forward(self, x):
         ep = export(N0(), inp)
         epm = ep.module()
         ufm = torch.export.unflatten(ep)
-        assert torch.allclose(epm(*inp), eager)
-        assert torch.allclose(ufm(*inp), eager)
+        if not torch.allclose(epm(*inp), eager):
+            raise AssertionError("epm output does not match eager output")
+        if not torch.allclose(ufm(*inp), eager):
+            raise AssertionError("ufm output does not match eager output")
 
     def test_unflatten_random_dag_mutating_buf_6(self):
         class N5(torch.nn.Module):
@@ -11937,8 +12053,10 @@ def forward(self, x):
         ep = export(N0(), inp)
         epm = ep.module()
         ufm = torch.export.unflatten(ep)
-        assert torch.allclose(epm(*inp), eager)
-        assert torch.allclose(ufm(*inp), eager)
+        if not torch.allclose(epm(*inp), eager):
+            raise AssertionError("epm output does not match eager output")
+        if not torch.allclose(ufm(*inp), eager):
+            raise AssertionError("ufm output does not match eager output")
 
     def test_unflatten_random_dag_mutating_buf_9(self):
         class N8(torch.nn.Module):
@@ -12082,8 +12200,10 @@ def forward(self, x):
         ep = export(N0(), inp)
         epm = ep.module()
         ufm = torch.export.unflatten(ep)
-        assert torch.allclose(epm(*inp), eager)
-        assert torch.allclose(ufm(*inp), eager)
+        if not torch.allclose(epm(*inp), eager):
+            raise AssertionError("epm output does not match eager output")
+        if not torch.allclose(ufm(*inp), eager):
+            raise AssertionError("ufm output does not match eager output")
 
     def test_unflatten_random_dag_preserving_4(self):
         # {0: [1, 2, 3], 1: [2], 2: [], 3: []}
@@ -12132,8 +12252,10 @@ def forward(self, x):
         ep = export(N0(), inp, preserve_module_call_signature=fqns)
         epm = ep.module()
         ufm = torch.export.unflatten(ep)
-        assert torch.allclose(epm(*inp), eager)
-        assert torch.allclose(ufm(*inp), eager)
+        if not torch.allclose(epm(*inp), eager):
+            raise AssertionError("epm output does not match eager output")
+        if not torch.allclose(ufm(*inp), eager):
+            raise AssertionError("ufm output does not match eager output")
 
     def test_unflatten_random_dag_mutating_buf_preserving_4(self):
         # {0: [2, 3], 1: [2], 2: [3], 3: []}
@@ -12194,8 +12316,10 @@ def forward(self, x):
         ep = export(N0(), inp, preserve_module_call_signature=fqns)
         epm = ep.module()
         ufm = torch.export.unflatten(ep)
-        assert torch.allclose(epm(*inp), eager)
-        assert torch.allclose(ufm(*inp), eager)
+        if not torch.allclose(epm(*inp), eager):
+            raise AssertionError("epm output does not match eager output")
+        if not torch.allclose(ufm(*inp), eager):
+            raise AssertionError("ufm output does not match eager output")
 
     def test_unflatten_random_dag_mutating_buf_preserving_4_1(self):
         # {0: [2], 1: [3], 2: [3], 3: []}
@@ -12255,8 +12379,10 @@ def forward(self, x):
         ep = export(N0(), inp, preserve_module_call_signature=fqns)
         epm = ep.module()
         ufm = torch.export.unflatten(ep)
-        assert torch.allclose(epm(*inp), eager)
-        assert torch.allclose(ufm(*inp), eager)
+        if not torch.allclose(epm(*inp), eager):
+            raise AssertionError("epm output does not match eager output")
+        if not torch.allclose(ufm(*inp), eager):
+            raise AssertionError("ufm output does not match eager output")
 
     def test_unflatten_random_dag_mutating_buf_preserving_5(self):
         # {0: [1, 2, 3], 1: [3, 4], 2: [3, 4], 3: [4], 4: []}
@@ -12336,8 +12462,10 @@ def forward(self, x):
         ep = export(N0(), inp, preserve_module_call_signature=fqns)
         epm = ep.module()
         ufm = torch.export.unflatten(ep)
-        assert torch.allclose(epm(*inp), eager)
-        assert torch.allclose(ufm(*inp), eager)
+        if not torch.allclose(epm(*inp), eager):
+            raise AssertionError("epm output does not match eager output")
+        if not torch.allclose(ufm(*inp), eager):
+            raise AssertionError("ufm output does not match eager output")
 
     def test_unflatten_random_dag_mutating_buf_preserving_7(self):
         # {0: [3, 4, 5, 6], 1: [2, 3, 4, 5, 6], 2: [3, 4, 5], 3: [5, 6], 4: [6], 5: [6], 6: []}
@@ -12461,8 +12589,10 @@ def forward(self, x):
         ep = export(N0(), inp, preserve_module_call_signature=fqns)
         epm = ep.module()
         ufm = torch.export.unflatten(ep)
-        assert torch.allclose(epm(*inp), eager)
-        assert torch.allclose(ufm(*inp), eager)
+        if not torch.allclose(epm(*inp), eager):
+            raise AssertionError("epm output does not match eager output")
+        if not torch.allclose(ufm(*inp), eager):
+            raise AssertionError("ufm output does not match eager output")
 
     def test_unflatten_random_dag_mutating_buf_preserving_10(self):
         class N9(torch.nn.Module):
@@ -12654,8 +12784,10 @@ def forward(self, x):
         )
         epm = ep.module()
         ufm = torch.export.unflatten(ep)
-        assert torch.allclose(epm(*inp), eager)
-        assert torch.allclose(ufm(*inp), eager)
+        if not torch.allclose(epm(*inp), eager):
+            raise AssertionError("epm output does not match eager output")
+        if not torch.allclose(ufm(*inp), eager):
+            raise AssertionError("ufm output does not match eager output")
 
     def test_unflatten_random_dag_const_preserving_3(self):
         class N2(torch.nn.Module):
@@ -12698,8 +12830,10 @@ def forward(self, x):
         ep = export(N0(), inp, preserve_module_call_signature=fqns)
         epm = ep.module()
         ufm = torch.export.unflatten(ep)
-        assert torch.allclose(epm(*inp), eager)
-        assert torch.allclose(ufm(*inp), eager)
+        if not torch.allclose(epm(*inp), eager):
+            raise AssertionError("epm output does not match eager output")
+        if not torch.allclose(ufm(*inp), eager):
+            raise AssertionError("ufm output does not match eager output")
 
     def test_none_buffers(self):
         mod = torch.nn.InstanceNorm1d(1)
@@ -12787,8 +12921,8 @@ def forward(self, x, y):
     def test_symint_input_specialization(self):
         class M(torch.nn.Module):
             def forward(self, x, y):
-                assert x == 3
-                assert y.shape[0] == 4
+                assert x == 3  # noqa: S101
+                assert y.shape[0] == 4  # noqa: S101
                 return x * y
 
         inp = (3, torch.randn(4, 4))
@@ -12845,8 +12979,8 @@ def forward(self, x, y):
         # While tracing the range was found to be a subset of the original range
         class M(torch.nn.Module):
             def forward(self, x, y):
-                assert x > 3
-                assert x <= 5
+                assert x > 3  # noqa: S101
+                assert x <= 5  # noqa: S101
                 return x * y
 
         inp = (4, torch.randn(4, 4))
@@ -12863,8 +12997,8 @@ def forward(self, x, y):
         # While tracing the range was found to be bigger than the original range
         class M(torch.nn.Module):
             def forward(self, x, y):
-                assert x > 1
-                assert x < 20
+                assert x > 1  # noqa: S101
+                assert x < 20  # noqa: S101
                 return x * y
 
         inp = (4, torch.randn(4, 4))
@@ -12881,8 +13015,8 @@ def forward(self, x, y):
         # While tracing the range was found to be outside of the original range
         class M(torch.nn.Module):
             def forward(self, x, y):
-                assert x > 10
-                assert x < 20
+                assert x > 10  # noqa: S101
+                assert x < 20  # noqa: S101
                 return x * y
 
         inp = (14, torch.randn(4, 4))
@@ -12999,8 +13133,10 @@ def forward(self, x):
         ep = export(N0(), inp, preserve_module_call_signature=fqns)
         epm = ep.module()
         ufm = torch.export.unflatten(ep)
-        assert torch.allclose(epm(*inp), eager)
-        assert torch.allclose(ufm(*inp), eager)
+        if not torch.allclose(epm(*inp), eager):
+            raise AssertionError("epm output does not match eager output")
+        if not torch.allclose(ufm(*inp), eager):
+            raise AssertionError("ufm output does not match eager output")
 
     def test_unflatten_no_unroll(self):
         inp = (torch.ones(1),)
@@ -13761,6 +13897,12 @@ def forward(self, x, f):
                 return x + f.int_1 + f.int_2
 
         torch._library.opaque_object.register_opaque_type(MyInput, typ="value")
+        self.addCleanup(
+            lambda name=torch._library.opaque_object.get_opaque_type_name(MyInput): (
+                torch._C._unregister_opaque_type(name),
+                _OPAQUE_TYPES_BY_NAME.pop(name, None),
+            )
+        )
         ep = export(Foo(), (torch.randn(2, 2), MyInput(4, 4)), strict=False)
 
         inp = torch.ones(2, 2)
@@ -13989,7 +14131,7 @@ def test_python_asserts_with_sym_int(self):
         class Model(torch.nn.Module):
             def forward(self, x):
                 y = x + 1
-                assert y.max().item() > 0
+                assert y.max().item() > 0  # noqa: S101
                 return y
 
         model = Model()
@@ -15974,6 +16116,21 @@ def forward(self, x, y):
         for node in decomposed_program.graph.nodes:
             self.assertEqual(node.meta["custom"]["my_field"], "dummy")
 
+    def test_run_decompositions_leafspec_deepcopy(self):
+        # Regression test: run_decompositions internally deepcopies the module
+        # call graph which contains TreeSpec objects with LeafSpec nodes.
+        # and has uninitialized slots on Python 3.10
+        class SimpleModel(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.linear = torch.nn.Linear(10, 5)
+
+            def forward(self, x):
+                return self.linear(x)
+
+        ep = export(SimpleModel(), (torch.randn(1, 10),))
+        ep.run_decompositions({})
+
     @testing.expectedFailureStrictV2
     def test_run_decompositions_keep_tensor_constant_metadata(self):
         """Make sure the metadata of tensor constants are kept after run_decompositions."""
@@ -16831,9 +16988,9 @@ def trace_hop(proxy_mode, fn, x, schema):
             )
 
         def dense_hop(fn, x, schema):
-            assert isinstance(schema, pytree.TreeSpec)
+            assert isinstance(schema, pytree.TreeSpec)  # noqa: S101
             schema = pytree.tree_unflatten([], schema)
-            assert (
+            assert (  # noqa: S101
                 isinstance(schema, torch.FunctionSchema)
                 and schema == torch.ops.aten.sin.default._schema
             )
@@ -17442,6 +17599,25 @@ def forward(self, x, y):
         expected_mask = torch.ones(3, 5, dtype=torch.bool).triu(diagonal=2)
         self.assertEqual(eager_out, expected_mask)
 
+    def test_quantile_export(self):
+        class QuantilePair(torch.nn.Module):
+            def __init__(self, noise=0.1):
+                super().__init__()
+                self.noise = noise
+
+            def forward(self, x):
+                q = torch.tensor(
+                    [self.noise, 1.0 - self.noise],
+                    device=x.device,
+                    dtype=x.dtype,
+                )
+                return torch.quantile(x, q, dim=-1, keepdim=True)
+
+        model = QuantilePair(noise=0.1)
+        x = torch.randn(1, 3200)
+        ep = export(model, (x,))
+        self.assertEqual(ep.module()(x), model(x))
+
 
 @unittest.skipIf(not torchdynamo.is_dynamo_supported(), "dynamo isn't support")
 class TestOneOffModelExportResult(TestCase):
@@ -17480,6 +17656,7 @@ def forward(self, q, k, v):
             ep = torch.export.export(ScaledDotProductAttention(), (q, k, v))
             ep.run_decompositions()
 
+    @skipIfXpu(msg="scaled_dot_product_attention issue on xpu")
     @skipIfCrossRef
     @unittest.skipIf(
         not PLATFORM_SUPPORTS_FLASH_ATTENTION,
@@ -17504,9 +17681,9 @@ def forward(self, q, k, v):
                 )
                 return attn_output
 
-        q = torch.randn(1, 16, 16, 64, dtype=torch.bfloat16, device="cuda")
-        k = torch.randn(1, 16, 16, 64, dtype=torch.bfloat16, device="cuda")
-        v = torch.randn(1, 16, 16, 64, dtype=torch.bfloat16, device="cuda")
+        q = torch.randn(1, 16, 16, 64, dtype=torch.bfloat16, device=device_type)
+        k = torch.randn(1, 16, 16, 64, dtype=torch.bfloat16, device=device_type)
+        v = torch.randn(1, 16, 16, 64, dtype=torch.bfloat16, device=device_type)
 
         ep = torch.export.export(
             ScaledDotProductAttention(), (q, k, v)
@@ -18085,6 +18262,46 @@ def forward(self, x):
         exported_param_names = [name for name, _ in gm.named_parameters()]
         self.assertEqual(original_param_names, exported_param_names)
 
+    def test_run_decompositions_tensor_list_list(self):
+        """run_decompositions should work for custom ops with List[List[Tensor]] args.
+
+        The C++ Functionalize fallback kernel now handles nested list arguments
+        (Tensor[][]) by recursively unwrapping/wrapping functional tensors.
+        """
+        with torch.library._scoped_library("test_tll", "DEF") as lib:
+            lib.define(
+                "merge_op(Tensor[][] nested, Tensor[] flat, int[] weights)"
+                " -> (Tensor[], Tensor)",
+                tags=[torch.Tag.pt2_compliant_tag],
+            )
+
+            @torch.library.impl("test_tll::merge_op", "cpu", lib=lib)
+            def merge_op_cpu(nested, flat, weights):
+                out = [nested[i][0] + flat[i] for i in range(len(flat))]
+                combined = torch.cat([f.unsqueeze(0) for f in flat], dim=0).sum(dim=0)
+                return out, combined
+
+            @torch.library.register_fake("test_tll::merge_op")
+            def merge_op_fake(nested, flat, weights):
+                out = [torch.empty_like(flat[i]) for i in range(len(flat))]
+                combined = torch.empty_like(flat[0])
+                return out, combined
+
+            class M(torch.nn.Module):
+                def forward(self, a, b, c, d):
+                    return torch.ops.test_tll.merge_op([[a, b], [c, d]], [a, c], [1, 2])
+
+            m = M()
+            a, b, c, d = (torch.randn(3, 4) for _ in range(4))
+            ep = export(m, (a, b, c, d))
+            ep2 = ep.run_decompositions({})
+            eager_out = m(a, b, c, d)
+            decomp_out = ep2.module()(a, b, c, d)
+            self.assertEqual(len(eager_out[0]), len(decomp_out[0]))
+            for e, d_ in zip(eager_out[0], decomp_out[0]):
+                self.assertEqual(e, d_)
+            self.assertEqual(eager_out[1], decomp_out[1])
+
 
 @unittest.skipIf(not torchdynamo.is_dynamo_supported(), "dynamo doesn't support")
 class TestExportCustomClass(TorchTestCase):
@@ -18132,7 +18349,8 @@ def forward(self, x):
 
         constants = lift_constants_pass(ep.graph_module, ep.graph_signature, {})
         for k, v in constants.items():
-            assert k not in ep.constants
+            if k in ep.constants:
+                raise AssertionError(f"Key {k} already exists in ep.constants")
             ep._constants[k] = v
         serialized_vals = serialize(ep)
         deserialized_ep = deserialize(serialized_vals)
diff --git a/test/export/test_export_opinfo.py b/test/export/test_export_opinfo.py
index 361674a69c7ae..b33aeb45438a3 100644
--- a/test/export/test_export_opinfo.py
+++ b/test/export/test_export_opinfo.py
@@ -22,6 +22,7 @@
 )
 from torch.testing._internal.common_utils import (
     IS_FBCODE,
+    IS_WINDOWS,
     run_tests,
     skipIfRocm,
     TestCase,
@@ -152,6 +153,11 @@ class TestExportOnFakeCuda(TestCase):
     # We set CUDA_VISIBLE_DEVICES="" to simulate a CPU machine with cuda build
     # Running this on all ops in op_db is too slow, so we only run on a selected subset
     @onlyCUDA
+    @unittest.skipIf(
+        IS_WINDOWS,
+        'Subprocess with CUDA_VISIBLE_DEVICES="" imports op_db which triggers '
+        "get_device_capability(); 0 devices raises Invalid device id on Windows.",
+    )
     @ops(selected_op_db, allowed_dtypes=(torch.float,))
     def test_fake_export(self, device, dtype, op):
         test_script = f"""\
@@ -218,6 +224,10 @@ def forward(self, *args):
         self.assertEqual(r, "")
 
     @unittest.skipIf(not torch.backends.cuda.is_built(), "requires CUDA build")
+    @unittest.skipIf(
+        IS_WINDOWS,
+        "Failing on Windows, device_count() changes from 0 to 1 ",
+    )
     def test_preserve_original_behavior(self):
         test_script = f"""\
 import torch
diff --git a/test/export/test_lift_unlift.py b/test/export/test_lift_unlift.py
index 4ab1d17ca898f..09577dfb9a5bd 100644
--- a/test/export/test_lift_unlift.py
+++ b/test/export/test_lift_unlift.py
@@ -1,6 +1,6 @@
 # Owner(s): ["oncall: export"]
 from collections import OrderedDict
-from typing import Any, Optional
+from typing import Any
 
 import torch
 from torch._export.passes.lift_constants_pass import (
@@ -64,7 +64,7 @@ def call_function(self, target, args, out: str, module_fqn: str = ""):
         self.nodes[out] = node
 
     def constant(
-        self, name: str, value: Any, target: Optional[str] = None, module_fqn: str = ""
+        self, name: str, value: Any, target: str | None = None, module_fqn: str = ""
     ):
         if target is None:
             target = name
@@ -119,9 +119,11 @@ def create_input_specs(self):
     # NOTE: does not handle non-user-outputs atm
     def gen_graph_signature(self) -> ExportGraphSignature:
         output = [n for n in self.graph.nodes if n.op == "output"]
-        assert len(output) == 1
+        if len(output) != 1:
+            raise AssertionError(f"Expected 1 output node, got {len(output)}")
         output = output[0]
-        assert len(output.args) == 1, "multiple outputs NYI"
+        if len(output.args) != 1:
+            raise AssertionError(f"multiple outputs NYI, got {len(output.args)}")
 
         return ExportGraphSignature(
             input_specs=self.create_input_specs(),
diff --git a/test/export/test_package.py b/test/export/test_package.py
index b650cb0396a51..2f3abec29b269 100644
--- a/test/export/test_package.py
+++ b/test/export/test_package.py
@@ -67,13 +67,13 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
         x3 = torch.randn(5, 2)
 
         def spec(self, x):
-            assert x.shape[0] == 3
+            assert x.shape[0] == 3  # noqa: S101
 
         def spec2(self, x):
-            assert x.shape[0] == 4
+            assert x.shape[0] == 4  # noqa: S101
 
         def spec3(self, x):
-            assert x.shape[0] >= 5
+            assert x.shape[0] >= 5  # noqa: S101
             return {"x": (Dim("batch", min=5), Dim.STATIC)}
 
         package = _ExportPackage()
diff --git a/test/export/test_passes.py b/test/export/test_passes.py
index 9491ff9a88c91..56f56438776b2 100644
--- a/test/export/test_passes.py
+++ b/test/export/test_passes.py
@@ -576,7 +576,8 @@ def test_views_op_having_view_copy(self) -> None:
 
         for aten_schema in aten_schemas:
             val = aten_schema.split(".")
-            assert len(val) <= 2
+            if len(val) > 2:
+                raise AssertionError(f"Expected at most 2 parts, got {len(val)}: {val}")
             name = ""
             overload = ""
             if len(val) == 1:
diff --git a/test/export/test_schema.py b/test/export/test_schema.py
index f184fead8b413..24bd712793cfa 100644
--- a/test/export/test_schema.py
+++ b/test/export/test_schema.py
@@ -138,6 +138,8 @@ def test_schema_check(self):
             thrift_checksum_next="",
             thrift_schema="",
             thrift_schema_path="",
+            enum_converter_header="",
+            enum_converter_header_path="",
         )
         next_version, _ = check(commit)
         self.assertEqual(next_version, [4, 1])
@@ -177,6 +179,8 @@ def test_schema_check(self):
             thrift_checksum_next="",
             thrift_schema="",
             thrift_schema_path="",
+            enum_converter_header="",
+            enum_converter_header_path="",
         )
         next_version, _ = check(commit)
         self.assertEqual(next_version, [4, 1])
@@ -219,6 +223,8 @@ def test_schema_check(self):
             thrift_checksum_next="",
             thrift_schema="",
             thrift_schema_path="",
+            enum_converter_header="",
+            enum_converter_header_path="",
         )
         next_version, _ = check(commit)
         self.assertEqual(next_version, [3, 3])
@@ -284,6 +290,8 @@ def test_schema_check(self):
             thrift_checksum_next="",
             thrift_schema="",
             thrift_schema_path="",
+            enum_converter_header="",
+            enum_converter_header_path="",
         )
         next_version, _ = check(commit)
         self.assertEqual(next_version, [3, 3])
@@ -319,6 +327,8 @@ def test_schema_check(self):
             thrift_checksum_next="",
             thrift_schema="",
             thrift_schema_path="",
+            enum_converter_header="",
+            enum_converter_header_path="",
         )
         next_version, _ = check(commit)
         self.assertEqual(next_version, [3, 3])
@@ -361,6 +371,8 @@ def test_schema_check(self):
             thrift_checksum_next="",
             thrift_schema="",
             thrift_schema_path="",
+            enum_converter_header="",
+            enum_converter_header_path="",
         )
         next_version, _ = check(commit)
         self.assertEqual(next_version, [3, 3])
@@ -400,6 +412,8 @@ def test_schema_check(self):
             thrift_checksum_next="",
             thrift_schema="",
             thrift_schema_path="",
+            enum_converter_header="",
+            enum_converter_header_path="",
         )
         next_version, _ = check(commit)
         self.assertEqual(next_version, [4, 1])
diff --git a/test/export/test_serialize.py b/test/export/test_serialize.py
index 6e4d41fc1937a..ee991f090bea9 100644
--- a/test/export/test_serialize.py
+++ b/test/export/test_serialize.py
@@ -24,6 +24,11 @@
 
     from torch.library import wrap_triton
     from torch.utils._triton import has_triton
+else:
+
+    def has_triton():
+        return False
+
 
 import torch
 import torch._dynamo as torchdynamo
@@ -564,6 +569,59 @@ def forward(self, x):
             self.assertNotIn(name, seen)
             seen.add(name)
 
+    def test_multi_return_list_tensor_unused(self) -> None:
+        """
+        Make sure serialization handles unused List[Tensor] outputs in
+        multi-return ops. The getitem node for the unused list output is
+        removed by DCE, so the serializer must synthesize names.
+        """
+        with torch.library._scoped_library("mylib", "DEF") as lib:
+            lib.define(
+                "tensor_and_list(Tensor x) -> (Tensor, Tensor[])",
+            )
+
+            @torch.library.impl(lib, "tensor_and_list", "CPU")
+            def tensor_and_list_impl(x):
+                return x * 2, [x, x + 1]
+
+            @torch.library.impl(lib, "tensor_and_list", "Meta")
+            def tensor_and_list_meta(x):
+                return x * 2, [x, x + 1]
+
+            class MyModule(torch.nn.Module):
+                def forward(self, x):
+                    # Only use the first (Tensor) output; the List[Tensor] is unused
+                    return torch.ops.mylib.tensor_and_list(x)[0]
+
+            exported_module = export(MyModule(), (torch.ones(3),), strict=True)
+            # Simulate the scenario where DCE removes the getitem for the
+            # unused List[Tensor] output (as happens in real serialization
+            # pipelines like package_sigmoid_model).
+            exported_module.graph.eliminate_dead_code()
+            exported_module.graph_module.recompile()
+
+            serialized = ExportedProgramSerializer().serialize(exported_module)
+            node = serialized.exported_program.graph_module.graph.nodes[-1]
+            self.assertEqual(node.target, "torch.ops.mylib.tensor_and_list.default")
+            self.assertEqual(len(node.outputs), 2)
+
+            # First output is a single tensor
+            self.assertTrue(node.outputs[0].as_tensor.name != "")
+            # Second output is the unused list -- should still have tensor entries
+            self.assertTrue(len(node.outputs[1].as_tensors) > 0)
+            for t in node.outputs[1].as_tensors:
+                self.assertIn("_unused_", t.name)
+
+            # Roundtrip: deserialize and verify functional correctness
+            deserialized_ep = deserialize(serialize(exported_module))
+            inp = torch.randn(3)
+            self.assertTrue(
+                torch.allclose(
+                    exported_module.module()(inp),
+                    deserialized_ep.module()(inp),
+                )
+            )
+
     def test_rational_ranges(self) -> None:
         class M(torch.nn.Module):
             def forward(self, x):
@@ -574,7 +632,10 @@ def forward(self, x):
         )
 
         range_constraints = list(ep.range_constraints.keys())
-        assert len(range_constraints) == 1
+        if len(range_constraints) != 1:
+            raise AssertionError(
+                f"Expected 1 range constraint, got {len(range_constraints)}"
+            )
         symint = range_constraints[0]
 
         import sympy
@@ -654,7 +715,10 @@ def forward(self, x, y):
             args = (torch.randn(3, device=device), torch.randn(3, device=device))
             ep = torch.export.export(m, args=args)
             ep = ep.run_decompositions(decompose_custom_triton_ops=False)
-            assert torch.allclose(m(*args), ep.module()(*args))
+            if not torch.allclose(m(*args), ep.module()(*args)):
+                raise AssertionError(
+                    "Exported model output does not match eager output"
+                )
 
             serialized = ExportedProgramSerializer().serialize(ep)
 
@@ -774,7 +838,8 @@ def forward(self, x):
 
         ep = torch.export.export(m, args=args)
         ep = ep.run_decompositions(decompose_custom_triton_ops=False)
-        assert torch.allclose(eager_result, ep.module()(*args))
+        if not torch.allclose(eager_result, ep.module()(*args)):
+            raise AssertionError("Exported model output does not match eager result")
 
         # This should not raise - constexpr matching should work for bool values
         serialized = ExportedProgramSerializer().serialize(ep)
@@ -1228,6 +1293,7 @@ def _check_graph(pre_dispatch):
                 if isinstance(orig, torch.Tensor) and orig.dtype not in [
                     torch.float8_e4m3fn,
                     torch.float8_e5m2,
+                    torch.float8_e8m0fnu,
                 ]:
                     if orig.is_meta:
                         self.assertEqual(orig, loaded)
@@ -1485,7 +1551,7 @@ def forward(self, a, b, c) -> torch.Tensor:
     def test_sym_bool(self):
         class Module(torch.nn.Module):
             def forward(self, x, y):
-                assert x.size(0) in y
+                assert x.size(0) in y  # noqa: S101
                 return x + y
 
         f = Module()
@@ -1890,8 +1956,14 @@ def forward(self):
         roundtrip_ep = deserialize(serialize(ep))
         self.assertTrue(torch.allclose(ep.module()(), roundtrip_ep.module()()))
 
-    def test_serialize_float8(self):
-        for dtype in [torch.float8_e5m2, torch.float8_e4m3fn]:
+    def test_serialize_dtypes(self):
+        for dtype in [
+            torch.float8_e5m2,
+            torch.float8_e4m3fn,
+            torch.float8_e8m0fnu,
+            torch.uint32,
+            torch.uint64,
+        ]:
 
             class MyModule(torch.nn.Module):
                 def forward(self, x):
@@ -2227,8 +2299,12 @@ def forward(self, x):
         serialized_vals = serialize(ep)
 
         ep_str = serialized_vals.exported_program.decode("utf-8")
-        assert "class_fqn" in ep_str
-        assert custom_obj._type().qualified_name() in ep_str
+        if "class_fqn" not in ep_str:
+            raise AssertionError("'class_fqn' not found in serialized output")
+        if custom_obj._type().qualified_name() not in ep_str:
+            raise AssertionError(
+                f"'{custom_obj._type().qualified_name()}' not found in serialized output"
+            )
 
         deserialized_ep = deserialize(serialized_vals)
 
diff --git a/test/export/test_sparse.py b/test/export/test_sparse.py
index 975e9979982f5..addaf097b6f50 100644
--- a/test/export/test_sparse.py
+++ b/test/export/test_sparse.py
@@ -120,7 +120,7 @@ def assertEqualMeta(self, x, y):
                 x_meta1, y_meta1 = (x.ccol_indices(), y.ccol_indices())
                 x_meta2, y_meta2 = (x.row_indices(), y.row_indices())
             else:
-                assert 0  # unreachable
+                raise AssertionError(f"Unexpected layout: {x.layout}")
             self.assertEqual(x_meta1, y_meta1, exact_layout=True)
             self.assertEqual(x_meta2, y_meta2, exact_layout=True)
             self.assertEqual(x.values(), y.values(), exact_layout=True)
diff --git a/test/export/test_unflatten.py b/test/export/test_unflatten.py
index 13b234c173e5b..b62b31ea3ea8d 100644
--- a/test/export/test_unflatten.py
+++ b/test/export/test_unflatten.py
@@ -178,6 +178,48 @@ def forward(self, x):
             id(getattr(unflattened_module.sub_net, "2")),
         )
 
+    def test_unflatten_shared_submodule_reorder(self):
+        """Test that _reorder_submodules handles @N-suffixed FQNs correctly.
+
+        When modules are shared (aliased), PyTorch assigns @N suffixes to
+        duplicate entries in _modules. The fqn_order dict (built from
+        fqn_list) filters out @N entries, so _reorder_submodules must fall
+        back to the base FQN for ordering.
+        """
+
+        class Block(torch.nn.Module):
+            def __init__(self) -> None:
+                super().__init__()
+                self.linear = torch.nn.Linear(10, 10)
+
+            def forward(self, x):
+                return self.linear(x)
+
+        class Model(torch.nn.Module):
+            def __init__(self) -> None:
+                super().__init__()
+                shared_block = Block()
+                self.blocks = torch.nn.Sequential(
+                    shared_block,
+                    torch.nn.ReLU(),
+                    shared_block,
+                    torch.nn.ReLU(),
+                )
+
+            def forward(self, x):
+                return self.blocks(x)
+
+        eager_module = Model()
+        inps = (torch.rand(2, 10),)
+        export_module = export(eager_module, inps, {}, strict=True)
+        unflattened_module = unflatten(export_module)
+        self.compare_outputs(eager_module, unflattened_module, inps)
+        # Verify shared identity is preserved
+        self.assertEqual(
+            id(getattr(unflattened_module.blocks, "0")),
+            id(getattr(unflattened_module.blocks, "2")),
+        )
+
     def test_assert_tensor_metadata_stack(self):
         class N(torch.nn.Module):
             def __init__(self):
diff --git a/test/forward_backward_compatibility/check_forward_backward_compatibility.py b/test/forward_backward_compatibility/check_forward_backward_compatibility.py
index 5a962dfa57c05..41699c7b86ee6 100644
--- a/test/forward_backward_compatibility/check_forward_backward_compatibility.py
+++ b/test/forward_backward_compatibility/check_forward_backward_compatibility.py
@@ -301,7 +301,7 @@ def check_bc(existing_schemas):
             log.warning(
                 "Can NOT find backward compatible schemas after changes "
                 "for schema %s from the following candidates:\n[\n%s\n]",
-                str(existing_schema),
+                existing_schema,
                 "\n\t".join(str(s) for s in matching_new_schemas),
             )
             # TODO Print out more details about why candidates don't match.
@@ -346,7 +346,7 @@ def check_fc(existing_schemas):
             log.warning(
                 "Can NOT find forward compatible schemas after changes "
                 "for schema %s from the following candidates:\n[\n\t%s\n]",
-                str(existing_schema),
+                existing_schema,
                 "\n\t".join(str(s) for s in matching_new_schemas),
             )
             log.warning(
diff --git a/test/functorch/attn_ft.py b/test/functorch/attn_ft.py
index d2aabe51ec9f8..f36042fed98e6 100644
--- a/test/functorch/attn_ft.py
+++ b/test/functorch/attn_ft.py
@@ -47,7 +47,11 @@ def __init__(
         self.position_embedding_type = position_embedding_type
 
         if self.position_embedding_type is not None:
-            assert max_position_embeddings is not None
+            if max_position_embeddings is None:
+                raise AssertionError(
+                    "max_position_embeddings must not be None when "
+                    f"position_embedding_type is {self.position_embedding_type}"
+                )
             self.max_position_embeddings = max_position_embeddings
             self.distance_embedding = nn.Embedding(
                 2 * max_position_embeddings - 1, self.attention_head_size
@@ -111,7 +115,11 @@ def forward(
             # with the distance between them
             distance = query_sequence - key_sequence
 
-            assert key_sequence.size <= self.max_position_embeddings
+            if not (key_sequence.size <= self.max_position_embeddings):
+                raise AssertionError(
+                    f"key_sequence.size ({key_sequence.size}) exceeds "
+                    f"max_position_embeddings ({self.max_position_embeddings})"
+                )
 
             # we can then use that as an indirect index into the embedding table values to look up the features for that index
             # this is just a `gather` primitive op. The resulting tensor will
diff --git a/test/functorch/attn_positional.py b/test/functorch/attn_positional.py
index c75e8f3ec9498..4523447111e18 100644
--- a/test/functorch/attn_positional.py
+++ b/test/functorch/attn_positional.py
@@ -37,7 +37,11 @@ def __init__(
         self.position_embedding_type = position_embedding_type
 
         if self.position_embedding_type is not None:
-            assert max_position_embeddings is not None
+            if max_position_embeddings is None:
+                raise AssertionError(
+                    "max_position_embeddings must not be None when "
+                    f"position_embedding_type is {self.position_embedding_type}"
+                )
             self.max_position_embeddings = max_position_embeddings
             self.distance_embedding = nn.Embedding(
                 2 * max_position_embeddings - 1, self.attention_head_size
diff --git a/test/functorch/common_utils.py b/test/functorch/common_utils.py
index 4fa17b89f19ee..dd1258b574901 100644
--- a/test/functorch/common_utils.py
+++ b/test/functorch/common_utils.py
@@ -31,7 +31,8 @@ def loop(op, in_dims, out_dim, batch_size, *batched_args, **kwarg_values):
     for idx in range(batch_size):
         flat_args, args_spec = pytree.tree_flatten(batched_args)
         flat_dims, dims_spec = pytree.tree_flatten(in_dims)
-        assert args_spec == dims_spec
+        if args_spec != dims_spec:
+            raise AssertionError(f"args_spec {args_spec} != dims_spec {dims_spec}")
         new_args = [
             a.select(in_dim, idx) if in_dim is not None else a
             for a, in_dim in zip(flat_args, flat_dims)
@@ -80,9 +81,14 @@ def loop2(
     flat_args, args_spec = pytree.tree_flatten(batched_args)
     flat_dims1, dims_spec1 = pytree.tree_flatten(in_dims1)
     flat_dims2, dims_spec2 = pytree.tree_flatten(in_dims2)
-    assert args_spec == dims_spec1
-    assert args_spec == dims_spec2
-    assert len(flat_dims1) == len(flat_dims2)
+    if args_spec != dims_spec1:
+        raise AssertionError(f"args_spec {args_spec} != dims_spec1 {dims_spec1}")
+    if args_spec != dims_spec2:
+        raise AssertionError(f"args_spec {args_spec} != dims_spec2 {dims_spec2}")
+    if len(flat_dims1) != len(flat_dims2):
+        raise AssertionError(
+            f"len(flat_dims1) {len(flat_dims1)} != len(flat_dims2) {len(flat_dims2)}"
+        )
     for idx1 in range(batch_size1):
         out_split = []
         arg_split = [
@@ -162,7 +168,10 @@ def get_bdim_choices(num_tensors):
     options = (-1, None)
     choices.extend(itertools.product(options, repeat=num_tensors))
 
-    assert choices[-1] == (None,) * num_tensors
+    if choices[-1] != (None,) * num_tensors:
+        raise AssertionError(
+            f"Expected choices[-1] to be {(None,) * num_tensors}, got {choices[-1]}"
+        )
     return tuple(choices[:-1])
 
 
@@ -193,13 +202,18 @@ def get_bdim_choices_batch_norm(
                 continue
             choices.append(choice)
 
-    assert choices[-1] == (None,) * num_tensors
+    if choices[-1] != (None,) * num_tensors:
+        raise AssertionError(
+            f"Expected choices[-1] to be {(None,) * num_tensors}, got {choices[-1]}"
+        )
     return tuple(choices[:-1])
 
 
 def add_batch_dim(arg, bdim, batch_size=3):
-    assert bdim == 0 or bdim == -1
-    assert isinstance(arg, torch.Tensor)
+    if bdim != 0 and bdim != -1:
+        raise AssertionError(f"Expected bdim to be 0 or -1, got {bdim}")
+    if not isinstance(arg, torch.Tensor):
+        raise AssertionError(f"Expected arg to be a torch.Tensor, got {type(arg)}")
     if bdim == 0:
         shape = [1] * len(arg.shape)
         shape.insert(bdim, batch_size)
@@ -238,7 +252,10 @@ def is_batch_norm_training(op_name, kwarg_values):
     if len(is_training) == 0:
         return default_training
     else:
-        assert len(is_training) == 1
+        if len(is_training) != 1:
+            raise AssertionError(
+                f"Expected len(is_training) to be 1, got {len(is_training)}"
+            )
         return is_training[0]
 
 
@@ -260,8 +277,10 @@ def generate_vmap_inputs(
 
     @memoize
     def get_batched_arg(arg, bdim):
-        assert isinstance(arg, torch.Tensor)
-        assert bdim is not None
+        if not isinstance(arg, torch.Tensor):
+            raise AssertionError(f"Expected arg to be a torch.Tensor, got {type(arg)}")
+        if bdim is None:
+            raise AssertionError("Expected bdim to not be None")
         result, _ = add_batch_dim(arg, bdim, batch_size)
         return result
 
@@ -405,7 +424,10 @@ def make_batched(t):
             batch_size,
             compute_loop_out=False,
         ):
-            assert quantities[1] is None
+            if quantities[1] is not None:
+                raise AssertionError(
+                    f"Expected quantities[1] to be None, got {quantities[1]}"
+                )
             yield (quantities[0], expected_batched)
             yield (quantities[2], quantities[3])
 
@@ -429,7 +451,8 @@ def opinfo_in_dict(opinfo, d):
 def decorate(
     op_name, variant_name="", *, decorator=None, device_type=None, dtypes=None
 ):
-    assert decorator is not None
+    if decorator is None:
+        raise AssertionError("decorator must not be None")
     return DecorateMeta(
         op_name=op_name,
         variant_name=variant_name,
@@ -489,11 +512,13 @@ def skipOps(test_case_name, base_test_name, to_skip):
             if o.name == decorate_meta.op_name
             and o.variant_test_name == decorate_meta.variant_name
         ]
-        assert len(matching_opinfos) > 0, f"Couldn't find OpInfo for {decorate_meta}"
-        assert len(matching_opinfos) == 1, (
-            "OpInfos should be uniquely determined by their (name, variant_name). "
-            f"Got more than one result for ({decorate_meta.op_name}, {decorate_meta.variant_name})"
-        )
+        if len(matching_opinfos) == 0:
+            raise AssertionError(f"Couldn't find OpInfo for {decorate_meta}")
+        if len(matching_opinfos) != 1:
+            raise AssertionError(
+                "OpInfos should be uniquely determined by their (name, variant_name). "
+                f"Got more than one result for ({decorate_meta.op_name}, {decorate_meta.variant_name})"
+            )
         opinfo = matching_opinfos[0]
         decorators = list(opinfo.decorators)
         new_decorator = DecorateInfo(
@@ -523,15 +548,17 @@ def wrapped(
         dtypes=dtypes,
     ):
         name_parts = fn.__qualname__.split(".")
-        assert len(name_parts) == 2, (
-            "Decorator only applies to a test function of a test class"
-        )
+        if len(name_parts) != 2:
+            raise AssertionError(
+                "Decorator only applies to a test function of a test class"
+            )
         test_case_name, base_test_name = name_parts
         for module_cls in module_classes:
             matching_module_infos = [m for m in module_db if m.module_cls == module_cls]
-            assert len(matching_module_infos) == 1, (
-                f"Couldn't find single ModuleInfo for {module_cls}"
-            )
+            if len(matching_module_infos) != 1:
+                raise AssertionError(
+                    f"Couldn't find single ModuleInfo for {module_cls}"
+                )
             module_info = matching_module_infos[0]
             decorators = list(module_info.decorators)
             new_decorator = DecorateInfo(
@@ -574,7 +601,8 @@ def opsToleranceOverride(test_case_name, base_test_name, overrides):
             for o in all_opinfos
             if o.name == op_name and o.variant_test_name == variant_name
         ]
-        assert len(matching_opinfos) == 1, f"Couldn't find OpInfo for {override}"
+        if len(matching_opinfos) != 1:
+            raise AssertionError(f"Couldn't find OpInfo for {override}")
         opinfo = matching_opinfos[0]
         decorators = list(opinfo.decorators)
         decorators.append(
diff --git a/test/functorch/discover_coverage.py b/test/functorch/discover_coverage.py
index c0c0db8762bde..6dd371b6ec8d1 100644
--- a/test/functorch/discover_coverage.py
+++ b/test/functorch/discover_coverage.py
@@ -294,7 +294,10 @@ def get_num_usages(opname):
         if opname == "t":
             return 0
         result = [op[1] for op in data if op[0] == opname]
-        assert len(result) == 1
+        if len(result) != 1:
+            raise AssertionError(
+                f"Expected exactly 1 result for {opname}, got {len(result)}"
+            )
         return result[0]
 
     # get all operators that are not in the denylist
@@ -352,7 +355,10 @@ class Status(Enum):
 
 
 def is_decorateinfo_skip_or_xfail(decorateinfo):
-    assert len(decorateinfo.decorators) == 1
+    if len(decorateinfo.decorators) != 1:
+        raise AssertionError(
+            f"Expected exactly 1 decorator, got {len(decorateinfo.decorators)}"
+        )
     actual_decorator = decorateinfo.decorators[0]
     if isinstance(actual_decorator, toleranceOverride):
         return False
@@ -462,7 +468,8 @@ def transpose_statuses(for_subset=None, invert=False):
 
 
 def remove_torch(name):
-    assert name[:6] == "torch."
+    if name[:6] != "torch.":
+        raise AssertionError(f"Expected name to start with 'torch.', got {name!r}")
     return name[6:]
 
 
@@ -515,8 +522,15 @@ def get_jvp_coverage(subset=None):
     supports_forward_ad = {
         remove_torch(test) for test in list(supports_forwardad_ops_dct.keys())
     }
-    assert supports_forward_ad.issubset(supports_autograd)
-    assert supports_autograd.issubset(ops)
+    if not supports_forward_ad.issubset(supports_autograd):
+        raise AssertionError(
+            f"supports_forward_ad is not a subset of supports_autograd: "
+            f"{supports_forward_ad - supports_autograd}"
+        )
+    if not supports_autograd.issubset(ops):
+        raise AssertionError(
+            f"supports_autograd is not a subset of ops: {supports_autograd - ops}"
+        )
 
     failed_ops = get_skipped_or_xfailed_ops_for("test_jvp")
 
@@ -699,7 +713,10 @@ class Operator:
     def __init__(self, name):
         self.name = name
         self.opinfos = NAME_TO_OPINFO.get(name, None)
-        assert self.opinfos is None or len(self.opinfos) > 0
+        if self.opinfos is not None and len(self.opinfos) == 0:
+            raise AssertionError(
+                f"Operator {name!r} found in NAME_TO_OPINFO but has empty opinfos list"
+            )
 
     def has_opinfo(self):
         return self.opinfos is not None
diff --git a/test/functorch/test_aotdispatch.py b/test/functorch/test_aotdispatch.py
index 2375778a29e43..d71642d6f9b74 100644
--- a/test/functorch/test_aotdispatch.py
+++ b/test/functorch/test_aotdispatch.py
@@ -14,7 +14,7 @@
 from collections.abc import Callable
 from contextlib import ContextDecorator, ExitStack, nullcontext
 from functools import partial, wraps
-from typing import Any, Optional, Union
+from typing import Any
 from unittest.mock import patch
 
 from common_utils import (
@@ -59,6 +59,11 @@
     aot_export_module,
     SerializableAOTDispatchCompiler,
 )
+from torch._functorch.partitioners import (
+    _extract_fwd_bwd_modules,
+    _extract_fwd_bwd_outputs,
+    _extract_graph_with_inputs_outputs,
+)
 from torch._higher_order_ops.out_dtype import out_dtype
 from torch._inductor.codecache import compiled_fx_graph_hash
 from torch._inductor.custom_graph_pass import CustomPartitionerFn
@@ -171,9 +176,12 @@ def _pack_fp8_wrap(x):
         # Check only during compilation
         # Test calls hooks to get reference output
         ctx = torch._functorch._aot_autograd.graph_compile._get_saved_tensor_hook_context()
-        assert ctx["_fw_graph"] is not None
-        assert ctx["_bw_graph"] is not None
-        assert ctx["_node"] is not None
+        if ctx["_fw_graph"] is None:
+            raise AssertionError("Expected _fw_graph to not be None")
+        if ctx["_bw_graph"] is None:
+            raise AssertionError("Expected _bw_graph to not be None")
+        if ctx["_node"] is None:
+            raise AssertionError("Expected _node to not be None")
 
     return (x.dtype, x.to(torch.float8_e5m2))
 
@@ -188,9 +196,12 @@ def _unpack_fp8_wrap(x):
         # Check only during compilation
         # Test calls hooks to get reference output
         ctx = torch._functorch._aot_autograd.graph_compile._get_saved_tensor_hook_context()
-        assert ctx["_fw_graph"] is not None
-        assert ctx["_bw_graph"] is not None
-        assert ctx["_node"] is not None
+        if ctx["_fw_graph"] is None:
+            raise AssertionError("Expected _fw_graph to not be None")
+        if ctx["_bw_graph"] is None:
+            raise AssertionError("Expected _bw_graph to not be None")
+        if ctx["_node"] is None:
+            raise AssertionError("Expected _node to not be None")
     return tensor.to(dtype)
 
 
@@ -418,8 +429,8 @@ class TestAOTAutograd(AOTTestCase):
     def run_autograd(
         self,
         f: Callable,
-        fw_graph_cell: list[Optional[Callable]],
-        decompositions: Optional[dict],
+        fw_graph_cell: list[Callable | None],
+        decompositions: dict | None,
         keep_input_mutations: bool,
         dynamic: bool,
     ):
@@ -457,11 +468,11 @@ def run_autograd(
     def verify_aot_autograd(
         self,
         f,
-        inp_: Union[Callable, list[Any]],
+        inp_: Callable | list[Any],
         *,
         test_mutation: bool = False,
         keep_inp_mutations: bool = False,
-        decompositions: Optional[dict] = None,
+        decompositions: dict | None = None,
         dynamic: bool = False,
         # Only active when inp_ is Callable.
         # TODO: probably consolidate all tests to make inp a Callable.
@@ -1527,7 +1538,8 @@ def test_input_mutation_storage_resize_down_and_set_(self):
         class TracableCreateParameter(torch.autograd.Function):
             @staticmethod
             def forward(ctx, tensor, placeholder):
-                assert not tensor.requires_grad
+                if tensor.requires_grad:
+                    raise AssertionError("Expected tensor to not require grad")
                 return placeholder.set_(tensor)
 
             @staticmethod
@@ -2203,6 +2215,36 @@ def forward(self, primals_1, primals_2):
     return (view, add)""",
         )
 
+    def test_output_aliases_intermediate_no_grad_view(self):
+        # View of a differentiable intermediate created under no_grad().
+        # The view inherits requires_grad from the base but has no grad_fn,
+        # so it should not be treated as a differentiable output.
+        def f(a):
+            out = torch.mul(a, 3)
+            with torch.no_grad():
+                out_view = out.view(-1)
+            return out.sum(), out_view
+
+        inp = torch.ones(3, 3, requires_grad=True)
+        compiled_f = aot_function(f, nop)
+
+        ref_loss, ref_view = f(inp)
+        test_loss, test_view = compiled_f(inp)
+        # The differentiable output (sum) participates in backward.
+        self.assertEqual(ref_loss, test_loss)
+        self.assertTrue(test_loss.requires_grad)
+        # The no_grad view still has requires_grad (matching eager) but
+        # does not participate in backward.
+        self.assertEqual(ref_view, test_view)
+        self.assertTrue(test_view.requires_grad)
+
+        ref_loss.backward()
+        ref_grad = inp.grad.clone()
+        inp.grad = None
+        test_loss.backward()
+        test_grad = inp.grad
+        self.assertEqual(ref_grad, test_grad)
+
     def test_output_aliases_intermediate_returned_multiple_times(self):
         def f(a):
             out = torch.mul(a, 3)
@@ -3940,7 +3982,8 @@ def bn(x):
         aot_fn = make_fx(bn, decomposition_table=cudnn_batch_norm_decomp)(inp)
         res = aot_fn(inp)
         for a, b in zip(ref, res):
-            assert torch.allclose(a, b)
+            if not torch.allclose(a, b):
+                raise AssertionError(f"Tensors not allclose: {a} vs {b}")
 
     def test_output_op_depending_on_symint(self):
         """
@@ -4309,7 +4352,10 @@ def __init__(self) -> None:
             def forward(self, x):
                 y = self.buffer.add_(3)
                 y.resize_([20])
-                assert y.shape == self.buffer.shape
+                if y.shape != self.buffer.shape:
+                    raise AssertionError(
+                        f"Expected y.shape {y.shape} == buffer.shape {self.buffer.shape}"
+                    )
                 return x.sum() + self.buffer.sum()
 
         m = M().eval()
@@ -4491,14 +4537,20 @@ def fn(x, y):
 
         def make_assert_pack(dynamic):
             def pack(activation):
-                assert hasattr(activation, "_dynamo_weak_dynamic_indices") == dynamic
+                if hasattr(activation, "_dynamo_weak_dynamic_indices") != dynamic:
+                    raise AssertionError(
+                        f"Expected hasattr(..., '_dynamo_weak_dynamic_indices') to be {dynamic}"
+                    )
                 return activation
 
             return pack
 
         def make_assert_unpack(dynamic):
             def unpack(activation):
-                assert hasattr(activation, "_dynamo_weak_dynamic_indices") == dynamic
+                if hasattr(activation, "_dynamo_weak_dynamic_indices") != dynamic:
+                    raise AssertionError(
+                        f"Expected hasattr(..., '_dynamo_weak_dynamic_indices') to be {dynamic}"
+                    )
                 return activation
 
             return unpack
@@ -4577,14 +4629,20 @@ def test_mark_activations_dynamic_with_nested(self):
 
         def make_assert_pack(dynamic):
             def pack(activation):
-                assert hasattr(activation, "_dynamo_weak_dynamic_indices") == dynamic
+                if hasattr(activation, "_dynamo_weak_dynamic_indices") != dynamic:
+                    raise AssertionError(
+                        f"Expected hasattr(..., '_dynamo_weak_dynamic_indices') to be {dynamic}"
+                    )
                 return activation
 
             return pack
 
         def make_assert_unpack(dynamic):
             def unpack(activation):
-                assert hasattr(activation, "_dynamo_weak_dynamic_indices") == dynamic
+                if hasattr(activation, "_dynamo_weak_dynamic_indices") != dynamic:
+                    raise AssertionError(
+                        f"Expected hasattr(..., '_dynamo_weak_dynamic_indices') to be {dynamic}"
+                    )
                 return activation
 
             return unpack
@@ -4623,6 +4681,58 @@ def fn(x, y, nt):
         counters.clear()
         torch._dynamo.reset()
 
+    def test_compute_outer_size_and_stride_bug(self):
+        import dataclasses
+
+        from torch._functorch._aot_autograd.schemas import (
+            PlainTensorMeta,
+            SubclassCreationMeta,
+        )
+
+        # Simulate a 2D TwoTensor(4, 6), contiguous, with dynamic shapes.
+        # After make_runtime_safe(), outer_size=(None, None), outer_stride=(None, 1).
+        # The packed symints in all_args at curr_start_idx are: [4, 6, 6]
+        # (2 size symints for the two None entries, 1 stride symint for the first None).
+        meta = object.__new__(SubclassCreationMeta)
+        for f in dataclasses.fields(SubclassCreationMeta):
+            setattr(meta, f.name, f.default)
+        meta.flat_tensor_start_idx = 0
+        meta.arg_count = 2
+        meta.included_subclass_symints = True
+        meta.attrs = {
+            "a": PlainTensorMeta(unwrapped_idx=0),
+            "b": PlainTensorMeta(unwrapped_idx=1),
+        }
+        meta.outer_size = (None, None)
+        meta.outer_stride = (None, 1)
+        meta.meta = None
+        meta.original_subclass = None
+        meta.original_subclass_type = TwoTensor
+        meta.memory_format = None
+
+        # all_args: [tensor_a, tensor_b, 4, 6, 6]
+        # curr_start_idx = 2 (after the two inner tensors)
+        a = torch.randn(4, 6)
+        b = torch.randn(4, 6)
+        all_args = [a, b, 4, 6, 6]
+
+        outer_size, outer_stride = meta.compute_outer_size_and_stride(
+            all_args, curr_start_idx=2
+        )
+        self.assertEqual(outer_size, (4, 6))
+        self.assertEqual(outer_stride, (6, 1))
+
+    def test_subclass_dynamic_stride(self):
+        @torch.compile(backend="aot_eager", dynamic=True)
+        def f(x):
+            return x.sin()
+
+        a = torch.randn(4, 6)
+        b = torch.randn(4, 6)
+        inp = TwoTensor(a, b)
+        out = f(inp)
+        self.assertEqual(out.stride(), inp.stride())
+
 
 def extract_graph(fx_g, _, graph_cell):
     graph_cell[0] = fx_g
@@ -5800,9 +5910,12 @@ def compile_fn(x, _):
         )
         res = compiled_fn(res_a, res_b)
         res.sum().backward()
-        assert torch.allclose(ref, res, atol=1e-3, rtol=1e-3)
-        assert torch.allclose(ref_a.grad, res_a.grad, atol=1e-3, rtol=1e-3)
-        assert torch.allclose(ref_b.grad, res_b.grad, atol=1e-3, rtol=1e-3)
+        if not torch.allclose(ref, res, atol=1e-3, rtol=1e-3):
+            raise AssertionError("ref and res not allclose")
+        if not torch.allclose(ref_a.grad, res_a.grad, atol=1e-3, rtol=1e-3):
+            raise AssertionError("ref_a.grad and res_a.grad not allclose")
+        if not torch.allclose(ref_b.grad, res_b.grad, atol=1e-3, rtol=1e-3):
+            raise AssertionError("ref_b.grad and res_b.grad not allclose")
 
     def test_meta_tensor_inplace_op(self):
         # Following module results in inplace ops while tracing. The test checks
@@ -5825,7 +5938,10 @@ def forward(self, add_4):
         def check_meta_tensor(fx_g, _):
             for node in fx_g.graph.nodes:
                 if node.op != "output":
-                    assert "tensor_meta" in node.meta
+                    if "tensor_meta" not in node.meta:
+                        raise AssertionError(
+                            f"Expected 'tensor_meta' in node.meta for {node}"
+                        )
             return fx_g
 
         inp0 = torch.randn(16, 128, 768, requires_grad=True)
@@ -5852,6 +5968,92 @@ def f(x, mod_weight, mod_bias):
         self.assertEqual(get_num_ins_outs(fw_graph), (3, 6))
         self.assertEqual(get_num_ins_outs(bw_graph), (6, 3))
 
+    def test_extract_fwd_bwd_modules_omit_aot_autograd_runtime(self):
+        """Test splitting a backward graph into dI/dW subgraphs using
+        omit_aot_autograd_runtime and ignore_must_be_in_fw_bw."""
+        from torch._functorch.partitioners import is_sym_node
+
+        def f(x, w):
+            return ((x @ w).sum(),)
+
+        x = torch.randn(4, 4, requires_grad=True)
+        w = torch.randn(4, 4, requires_grad=True)
+        joint = aot_export_joint_simple(f, [x, w], trace_joint=True)
+        fw_g, bw_g = default_partition(joint, [x, w], num_fwd_outputs=1)
+
+        # Run reference full backward
+        fw_outs = fw_g(x.detach(), w.detach())
+        activations = fw_outs[1:]
+        grad_out = torch.ones_like(fw_outs[0])
+        full_bw = bw_g(*activations, grad_out)
+
+        # Re-partition bw_g into dI (grad_x) and dW (grad_w),
+        # mimicking the split_dI_dW pattern.
+        # NOTE: placeholder names are load-bearing for the partitioner —
+        # _is_primal / _is_tangent classify nodes by checking substrings of
+        # node.target (e.g. "tangents" in target → is_tangent). We must create
+        # new placeholder nodes (not just rename) so the target is also updated.
+        bw_gm = copy.deepcopy(bw_g)
+        for p in list(bw_gm.graph.find_nodes(op="placeholder")):
+            if p.name.startswith("tangent"):
+                with bw_gm.graph.inserting_before(p):
+                    new_p = bw_gm.graph.placeholder(f"not_tngnt{p.name[8:]}")
+                    new_p.meta.update(p.meta)
+                    p.replace_all_uses_with(new_p)
+                    bw_gm.graph.erase_node(p)
+        bw_gm.recompile()
+
+        num_di = 1  # first output = grad_x
+        di_outs, dw_outs, di_descs, _ = _extract_fwd_bwd_outputs(
+            bw_gm, num_fwd_outputs=num_di
+        )
+        args = list(bw_gm.graph.find_nodes(op="placeholder"))
+        di_graph = _extract_graph_with_inputs_outputs(
+            bw_gm.graph,
+            args,
+            di_outs,
+            di_descs,
+            "forward",
+            ignore_must_be_in_fw_bw=True,
+        )
+        di_names = {n.name for n in di_graph.nodes if n.op != "output"}
+
+        saved_values = [
+            n
+            for n in bw_gm.graph.nodes
+            if n.name in di_names
+            and not is_sym_node(n)
+            and "tensor_meta" in n.meta
+            and any(n2.name not in di_names for n2 in n.users)
+        ]
+
+        di_mod, dw_mod = _extract_fwd_bwd_modules(
+            bw_gm,
+            saved_values,
+            saved_sym_nodes=[],
+            num_fwd_outputs=num_di,
+            ignore_must_be_in_fw_bw=True,
+            omit_aot_autograd_runtime=True,
+        )
+
+        # Verify graph structure
+        def ph_names(gm):
+            return [n.name for n in gm.graph.find_nodes(op="placeholder")]
+
+        self.assertFalse(any(n.startswith("tangent") for n in ph_names(di_mod)))
+        self.assertFalse(any(n.startswith("tangent") for n in ph_names(dw_mod)))
+
+        # Verify numerics: dI + dW matches full backward
+        bw_inputs = [*activations, grad_out]
+        di_outs = di_mod(*bw_inputs)
+        grad_x = di_outs[:num_di]
+        dw_inputs = di_outs[num_di:]
+        grad_w = dw_mod(*dw_inputs)
+        if not isinstance(grad_w, (list, tuple)):
+            grad_w = [grad_w]
+        self.assertEqual(grad_x[0], full_bw[0])
+        self.assertEqual(grad_w[0], full_bw[1])
+
     @unittest.skipIf(not USE_NETWORKX, "networkx not available")
     def test_min_cut_partitioner_raise_getitems(self):
         def f(x):
@@ -6123,7 +6325,8 @@ def fn(x):
         aot_fn = aot_function(fn, nop)
         res = aot_fn(x)
 
-        assert torch.allclose(ref, res)
+        if not torch.allclose(ref, res):
+            raise AssertionError("ref and res not allclose")
 
     # https://github.com/pytorch/pytorch/issues/110666
     def test_generate_gives_inference_graph(self):
@@ -6676,6 +6879,206 @@ def is_must_save(node):
         finally:
             handle.destroy()
 
+    def test_static_input_indices_with_effect_tokens(self):
+        """Test that static_input_indices are correctly offset when effect tokens are prepended."""
+        from torch._functorch._aot_autograd.descriptors import PlainAOTInput
+        from torch._functorch._aot_autograd.graph_capture_wrappers import (
+            handle_effect_tokens_fn,
+        )
+        from torch._functorch._aot_autograd.schemas import ViewAndMutationMeta
+        from torch._higher_order_ops.effects import _register_effectful_op
+        from torch._library.effects import EffectType
+
+        @torch.library.custom_op("test::effectful_static_idx_op", mutates_args=())
+        def effectful_static_idx_op(x: torch.Tensor) -> torch.Tensor:
+            return x.clone()
+
+        @effectful_static_idx_op.register_fake
+        def _(x: torch.Tensor) -> torch.Tensor:
+            return torch.empty_like(x)
+
+        handle = _register_effectful_op(effectful_static_idx_op, EffectType.ORDERED)
+
+        try:
+            meta = ViewAndMutationMeta(
+                input_info=[],
+                output_info=[],
+                num_intermediate_bases=0,
+                keep_input_mutations=False,
+                traced_tangents=[],
+                subclass_inp_meta=[],
+                subclass_fw_graph_out_meta=[],
+                subclass_tangent_meta=[],
+                traced_tangents_descs=[],
+            )
+            meta.tokens = {EffectType.ORDERED: torch.tensor([])}
+            # Simulate: params at indices 1, 2 before token insertion
+            meta.static_input_indices = [1, 2]
+
+            args = [torch.randn(4, 3), torch.randn(2, 3), torch.randn(2)]
+            args_descs = [PlainAOTInput(i) for i in range(3)]
+
+            _, new_args, _ = handle_effect_tokens_fn(
+                lambda *a: (torch.randn(1), []),
+                args,
+                args_descs,
+                meta=meta,
+                trace_joint=False,
+            )
+
+            num_tokens = len(meta.tokens)
+            self.assertGreater(num_tokens, 0)
+            # After prepending tokens, args grow and indices must be offset
+            self.assertEqual(len(new_args), len(args) + num_tokens)
+            self.assertEqual(
+                meta.static_input_indices,
+                [1 + num_tokens, 2 + num_tokens],
+            )
+        finally:
+            handle.destroy()
+
+    def test_collect_metadata_subclass_fw_outs_follow_input_mutation_type(self):
+        from torch._functorch._aot_autograd.collect_metadata_analysis import (
+            run_functionalized_fw_and_collect_metadata,
+        )
+        from torch._functorch._aot_autograd.descriptors import PlainAOTInput
+        from torch._functorch._aot_autograd.schemas import SubclassCreationMeta
+
+        def f(x):
+            x.add_(1)
+            return [torch.sin(x)]
+
+        fake_mode = FakeTensorMode()
+        subclass_arg = TwoTensor(
+            fake_mode.from_tensor(torch.ones(2)),
+            fake_mode.from_tensor(torch.ones(2)),
+        )
+
+        keep_input_mutations_meta = run_functionalized_fw_and_collect_metadata(
+            f,
+            flat_args_descs=[PlainAOTInput(0)],
+            keep_input_mutations=True,
+            static_input_indices=[],
+        )(subclass_arg)
+        self.assertEqual(keep_input_mutations_meta.mutated_inp_runtime_indices, [])
+        self.assertEqual(len(keep_input_mutations_meta.subclass_fw_graph_out_meta), 1)
+        self.assertIsInstance(
+            keep_input_mutations_meta.subclass_fw_graph_out_meta[0],
+            SubclassCreationMeta,
+        )
+        self.assertEqual(
+            keep_input_mutations_meta.subclass_fw_graph_out_meta[
+                0
+            ].flat_tensor_start_idx,
+            0,
+        )
+
+        out_of_graph_mutation_meta = run_functionalized_fw_and_collect_metadata(
+            f,
+            flat_args_descs=[PlainAOTInput(0)],
+            keep_input_mutations=False,
+            static_input_indices=[],
+        )(subclass_arg)
+        self.assertEqual(out_of_graph_mutation_meta.mutated_inp_runtime_indices, [0])
+        self.assertEqual(len(out_of_graph_mutation_meta.subclass_fw_graph_out_meta), 2)
+        self.assertIsInstance(
+            out_of_graph_mutation_meta.subclass_fw_graph_out_meta[0],
+            SubclassCreationMeta,
+        )
+        self.assertIsInstance(
+            out_of_graph_mutation_meta.subclass_fw_graph_out_meta[1],
+            SubclassCreationMeta,
+        )
+        self.assertEqual(
+            out_of_graph_mutation_meta.subclass_fw_graph_out_meta[
+                0
+            ].flat_tensor_start_idx,
+            0,
+        )
+        self.assertEqual(
+            out_of_graph_mutation_meta.subclass_fw_graph_out_meta[
+                1
+            ].flat_tensor_start_idx,
+            2,
+        )
+
+    def test_collect_metadata_subclass_fw_outs_include_metadata_only_mutation(self):
+        from torch._functorch._aot_autograd.collect_metadata_analysis import (
+            run_functionalized_fw_and_collect_metadata,
+        )
+        from torch._functorch._aot_autograd.descriptors import PlainAOTInput
+        from torch._functorch._aot_autograd.schemas import (
+            PlainTensorMeta,
+            SubclassCreationMeta,
+        )
+
+        def f(x):
+            x.transpose_(0, 1)
+            return [TwoTensor(torch.sin(x), torch.cos(x))]
+
+        fake_mode = FakeTensorMode()
+        arg = fake_mode.from_tensor(torch.ones(2, 3))
+
+        metadata = run_functionalized_fw_and_collect_metadata(
+            f,
+            flat_args_descs=[PlainAOTInput(0)],
+            keep_input_mutations=True,
+            static_input_indices=[],
+        )(arg)
+
+        self.assertEqual(metadata.mutated_inp_runtime_indices, [0])
+        self.assertEqual(len(metadata.subclass_fw_graph_out_meta), 2)
+        self.assertIsInstance(metadata.subclass_fw_graph_out_meta[0], PlainTensorMeta)
+        self.assertEqual(metadata.subclass_fw_graph_out_meta[0].unwrapped_idx, 0)
+        self.assertIsInstance(
+            metadata.subclass_fw_graph_out_meta[1],
+            SubclassCreationMeta,
+        )
+        self.assertEqual(
+            metadata.subclass_fw_graph_out_meta[1].flat_tensor_start_idx,
+            1,
+        )
+
+    def test_collect_metadata_subclass_fw_outs_include_intermediate_bases(self):
+        from torch._functorch._aot_autograd.collect_metadata_analysis import (
+            run_functionalized_fw_and_collect_metadata,
+        )
+        from torch._functorch._aot_autograd.descriptors import PlainAOTInput
+        from torch._functorch._aot_autograd.schemas import PlainTensorMeta
+
+        def f(x):
+            y = x.a + x.b
+            return [y.view(-1), y.view(-1)]
+
+        fake_mode = FakeTensorMode()
+        subclass_arg = TwoTensor(
+            fake_mode.from_tensor(torch.ones(2, 3, requires_grad=True)),
+            fake_mode.from_tensor(torch.ones(2, 3, requires_grad=True)),
+        )
+
+        metadata = run_functionalized_fw_and_collect_metadata(
+            f,
+            flat_args_descs=[PlainAOTInput(0)],
+            keep_input_mutations=True,
+            static_input_indices=[],
+        )(subclass_arg)
+
+        self.assertEqual(metadata.num_intermediate_bases, 1)
+        self.assertEqual(len(metadata.subclass_fw_graph_out_meta), 3)
+        self.assertTrue(
+            all(
+                isinstance(out_meta, PlainTensorMeta)
+                for out_meta in metadata.subclass_fw_graph_out_meta
+            )
+        )
+        self.assertEqual(
+            [
+                out_meta.unwrapped_idx
+                for out_meta in metadata.subclass_fw_graph_out_meta
+            ],
+            [0, 1, 2],
+        )
+
 
 class TestAOTDispatch(AOTTestCase):
     # Tests to add cases for (non-exhaustive list, mostly for my notes):
@@ -7114,7 +7517,8 @@ def fn(x):
                     (ref_out[0] + ref_out[1]).sum().backward()
                     (out[0] + out[1]).sum().backward()
                     self.assertEqual(ref_x.grad, x.grad)
-                    assert torch.allclose(ref_x.grad, x.grad, atol=1e-3, rtol=1e-3)
+                    if not torch.allclose(ref_x.grad, x.grad, atol=1e-3, rtol=1e-3):
+                        raise AssertionError("ref_x.grad and x.grad not allclose")
 
     def test_aot_dispatch_output_requires_grad_in_no_grad_views(self):
         # view-type ops preserve requires_grad even in no_grad.
@@ -7248,9 +7652,16 @@ def forward(self, x, y):
         res = compiled_f(*cloned_inputs)
         res[0].sum().backward()
 
-        assert torch.allclose(ref[0], res[0])
-        assert torch.allclose(inputs[0].grad, cloned_inputs[0].grad)
-        assert torch.allclose(inputs[1].grad, cloned_inputs[1].grad)
+        if not torch.allclose(ref[0], res[0]):
+            raise AssertionError("ref[0] and res[0] not allclose")
+        if not torch.allclose(inputs[0].grad, cloned_inputs[0].grad):
+            raise AssertionError(
+                "inputs[0].grad and cloned_inputs[0].grad not allclose"
+            )
+        if not torch.allclose(inputs[1].grad, cloned_inputs[1].grad):
+            raise AssertionError(
+                "inputs[1].grad and cloned_inputs[1].grad not allclose"
+            )
 
     def test_aot_module_simplified_dynamic(self):
         class MockModule(torch.nn.Module):
@@ -7287,9 +7698,16 @@ def forward(self, x, y):
  - Eq(s70, 30)""",
         )
 
-        assert torch.allclose(ref[0], res[0])
-        assert torch.allclose(inputs[0].grad, cloned_inputs[0].grad)
-        assert torch.allclose(inputs[1].grad, cloned_inputs[1].grad)
+        if not torch.allclose(ref[0], res[0]):
+            raise AssertionError("ref[0] and res[0] not allclose")
+        if not torch.allclose(inputs[0].grad, cloned_inputs[0].grad):
+            raise AssertionError(
+                "inputs[0].grad and cloned_inputs[0].grad not allclose"
+            )
+        if not torch.allclose(inputs[1].grad, cloned_inputs[1].grad):
+            raise AssertionError(
+                "inputs[1].grad and cloned_inputs[1].grad not allclose"
+            )
 
     # https://github.com/pytorch/pytorch/issues/105327
     def test_lift_fresh_copy_in_graph(self):
@@ -7355,14 +7773,20 @@ def forward(self, x, y):
             if node.op != "call_function":
                 continue
             self.assertTrue(node.stack_trace is not None)
-            assert "test_aotdispatch.py" in node.stack_trace
+            if "test_aotdispatch.py" not in node.stack_trace:
+                raise AssertionError(
+                    f"Expected 'test_aotdispatch.py' in stack trace, got: {node.stack_trace}"
+                )
 
         def assert_compiler(gm: torch.fx.GraphModule, _):
             for node in gm.graph.nodes:
                 if node.op == "output" or node.op == "placeholder":
                     continue
                 self.assertTrue(node.stack_trace is not None)
-                assert "test_aotdispatch.py" in node.stack_trace
+                if "test_aotdispatch.py" not in node.stack_trace:
+                    raise AssertionError(
+                        f"Expected 'test_aotdispatch.py' in stack trace, got: {node.stack_trace}"
+                    )
             return gm.forward  # return a python callable
 
         x = torch.randn(128, 20, requires_grad=True)
@@ -7394,14 +7818,18 @@ def forward(self, x):
             if node.op != "call_function":
                 continue
             self.assertTrue(node.stack_trace is not None)
-            assert "test_aotdispatch.py" in node.stack_trace
+            if "test_aotdispatch.py" not in node.stack_trace:
+                raise AssertionError("Expected 'test_aotdispatch.py' in stack trace")
 
         def assert_compiler(gm: torch.fx.GraphModule, _):
-            assert torch.ops.aten.copy_.default in [x.target for x in gm.graph.nodes]
+            if torch.ops.aten.copy_.default not in [x.target for x in gm.graph.nodes]:
+                raise AssertionError("Expected aten.copy_.default in graph nodes")
             for node in gm.graph.nodes:
                 if node.target == torch.ops.aten.copy_.default:
-                    assert "stack_trace" in node.meta
-                    assert "x_view.mul_(2)" in node.meta["stack_trace"]
+                    if "stack_trace" not in node.meta:
+                        raise AssertionError("Expected 'stack_trace' in node.meta")
+                    if "x_view.mul_(2)" not in node.meta["stack_trace"]:
+                        raise AssertionError("Expected 'x_view.mul_(2)' in stack trace")
             return gm.forward  # return a python callable
 
         x = torch.randn(128, 20)
@@ -7628,10 +8056,10 @@ def fn(x):
         y2 = fn_comp(x2)
         with self.assertRaisesRegex(
             RuntimeError,
-            """
+            r"""
 During the backward, we encountered a tensor subclass where we guessed its
 metadata incorrectly.
-""",  # noqa: F541
+Expected a .* tangent but got a plain Tensor.""",
         ):
             y2.backward(gradient=torch.randn(2, 3))
 
@@ -8183,122 +8611,125 @@ def unpack_cpu(x):
     @unittest.skipIf(not torch.cuda.is_available(), "CUDA is unavailable")
     @unittest.skipIf(not SM80OrLater, "bfloat16, float8")
     def test_saved_tensors_hooks_params(self):
-        lib = torch.library.Library("_test_aotdispatch_lib", "FRAGMENT")
-        logged_shapes = []
-        logged_dtypes = []
-        lib.define("log(Tensor x) -> Tensor")
-
-        def log_impl(x):
-            logged_shapes.append(list(x.shape))
-            logged_dtypes.append(x.dtype)
-            return x.clone()
+        with torch.library._scoped_library("_test_aotdispatch_lib", "FRAGMENT") as lib:
+            logged_shapes = []
+            logged_dtypes = []
+            lib.define("log(Tensor x) -> Tensor")
+
+            def log_impl(x):
+                logged_shapes.append(list(x.shape))
+                logged_dtypes.append(x.dtype)
+                return x.clone()
 
-        def log_meta(x):
-            return x.clone()
+            def log_meta(x):
+                return x.clone()
 
-        for backend in ["CPU", "CUDA"]:
-            lib.impl(
-                "log",
-                log_impl,
-                backend,
-            )
-        lib.impl("log", log_meta, "Meta")
+            for backend in ["CPU", "CUDA"]:
+                lib.impl(
+                    "log",
+                    log_impl,
+                    backend,
+                )
+            lib.impl("log", log_meta, "Meta")
 
-        def pack_fp8_with_scale_and_log(x):
-            torch.ops._test_aotdispatch_lib.log(x)
-            return _pack_fp8_with_scale_wrap(x)
+            def pack_fp8_with_scale_and_log(x):
+                torch.ops._test_aotdispatch_lib.log(x)
+                return _pack_fp8_with_scale_wrap(x)
 
-        def unpack_fp8_with_scale_and_log(packed):
-            return _unpack_fp8_with_scale_wrap(packed)
+            def unpack_fp8_with_scale_and_log(packed):
+                return _unpack_fp8_with_scale_wrap(packed)
 
-        def m_inp_fn():
-            x = torch.ones(
-                2, 2, 2, device=device, dtype=torch.float64, requires_grad=True
-            )
-            torch._dynamo.mark_dynamic(x, 0)
-            torch._dynamo.mark_dynamic(x, 1)
-            return (x,)
+            def m_inp_fn():
+                x = torch.ones(
+                    2, 2, 2, device=device, dtype=torch.float64, requires_grad=True
+                )
+                torch._dynamo.mark_dynamic(x, 0)
+                torch._dynamo.mark_dynamic(x, 1)
+                return (x,)
 
-        class SAF0(torch.autograd.Function):
-            @staticmethod
-            def forward(ctx, x):
-                ctx.save_for_backward(x)
-                return x
+            class SAF0(torch.autograd.Function):
+                @staticmethod
+                def forward(ctx, x):
+                    ctx.save_for_backward(x)
+                    return x
 
-            @staticmethod
-            def backward(ctx, gx):
-                (saved_x,) = ctx.saved_tensors
-                return gx + saved_x
+                @staticmethod
+                def backward(ctx, gx):
+                    (saved_x,) = ctx.saved_tensors
+                    return gx + saved_x
 
-        class M(torch.nn.Module):
-            def __init__(self):
-                super().__init__()
-                self.fc1 = nn.Linear(2, 2)
-                self.relu = nn.ReLU()
-                self.fc2 = nn.Linear(2, 2)
+            class M(torch.nn.Module):
+                def __init__(self):
+                    super().__init__()
+                    self.fc1 = nn.Linear(2, 2)
+                    self.relu = nn.ReLU()
+                    self.fc2 = nn.Linear(2, 2)
 
-            def forward(self, x):
-                x = SAF0.apply(x)
-                x = x.to(dtype=torch.float32)
-                x = self.fc1(x)
-                x = self.relu(x)
-                x = self.fc2(x)
-                return x
+                def forward(self, x):
+                    x = SAF0.apply(x)
+                    x = x.to(dtype=torch.float32)
+                    x = self.fc1(x)
+                    x = self.relu(x)
+                    x = self.fc2(x)
+                    return x
 
-        def _reset_logged():
-            logged_shapes.clear()
-            logged_dtypes.clear()
+            def _reset_logged():
+                logged_shapes.clear()
+                logged_dtypes.clear()
 
-        device = torch.device("cuda:0")
-        m = M().to(device=device)
+            device = torch.device("cuda:0")
+            m = M().to(device=device)
 
-        def _test_m():
-            self._test_pack_hooks(
-                m,
-                m_inp_fn,
-                [
-                    (
+            def _test_m():
+                self._test_pack_hooks(
+                    m,
+                    m_inp_fn,
+                    [
                         (
-                            pack_fp8_with_scale_and_log,
-                            unpack_fp8_with_scale_and_log,
-                        ),
-                        True,
-                    )
-                ],
-                pre_compile_fn=_reset_logged,
-                backend="aot_eager",
-            )
-
-        with patch(
-            "torch._functorch.config.saved_tensors_hooks_filtering_mode", "donated"
-        ):
-            _reset_logged()
-            _test_m()
-            # Check that hooks were not applied to Parameters
-            # parameters excluded
-            self.assertFalse([2, 2] in logged_shapes)
-            self.assertTrue([2, 2, 2] in logged_shapes)
-            # input excluded
-            self.assertFalse(torch.float64 in logged_dtypes)
+                            (
+                                pack_fp8_with_scale_and_log,
+                                unpack_fp8_with_scale_and_log,
+                            ),
+                            True,
+                        )
+                    ],
+                    pre_compile_fn=_reset_logged,
+                    backend="aot_eager",
+                )
 
-        with patch(
-            "torch._functorch.config.saved_tensors_hooks_filtering_mode", "no_static"
-        ):
-            _reset_logged()
-            _test_m()
-            # Check that hooks were not applied to Parameters
-            # parameters excluded
-            self.assertFalse([2, 2] in logged_shapes)
-            self.assertTrue([2, 2, 2] in logged_shapes)
-            self.assertTrue(torch.float64 in logged_dtypes)
-
-        with patch("torch._functorch.config.saved_tensors_hooks_filtering_mode", "all"):
-            _reset_logged()
-            _test_m()
-            # Check that hooks were applied to all saved tensors
-            self.assertTrue([2, 2] in logged_shapes)
-            self.assertTrue([2, 2, 2] in logged_shapes)
-            self.assertTrue(torch.float64 in logged_dtypes)
+            with patch(
+                "torch._functorch.config.saved_tensors_hooks_filtering_mode", "donated"
+            ):
+                _reset_logged()
+                _test_m()
+                # Check that hooks were not applied to Parameters
+                # parameters excluded
+                self.assertFalse([2, 2] in logged_shapes)
+                self.assertTrue([2, 2, 2] in logged_shapes)
+                # input excluded
+                self.assertFalse(torch.float64 in logged_dtypes)
+
+            with patch(
+                "torch._functorch.config.saved_tensors_hooks_filtering_mode",
+                "no_static",
+            ):
+                _reset_logged()
+                _test_m()
+                # Check that hooks were not applied to Parameters
+                # parameters excluded
+                self.assertFalse([2, 2] in logged_shapes)
+                self.assertTrue([2, 2, 2] in logged_shapes)
+                self.assertTrue(torch.float64 in logged_dtypes)
+
+            with patch(
+                "torch._functorch.config.saved_tensors_hooks_filtering_mode", "all"
+            ):
+                _reset_logged()
+                _test_m()
+                # Check that hooks were applied to all saved tensors
+                self.assertTrue([2, 2] in logged_shapes)
+                self.assertTrue([2, 2, 2] in logged_shapes)
+                self.assertTrue(torch.float64 in logged_dtypes)
 
     @unittest.skipIf(not torch.cuda.is_available(), "CUDA is unavailable")
     @unittest.skipIf(not SM80OrLater, "bfloat16, float8")
@@ -8458,12 +8889,9 @@ def fn(x):
     xfail("nn.functional.gaussian_nll_loss"),
     xfail("tensor_split"),
     xfail("corrcoef"),
-    xfail("quantile"),
-    xfail("nanquantile"),
     skip("narrow"),
     xfail("istft"),
     xfail("linalg.eig"),
-    skip("as_strided_scatter"),
     skip("as_strided", "partial_views"),  # flaky
     # Given input size: (s0xs1x2). Calculated output size: ...
     skip("max_pool2d_with_indices_backward"),
@@ -8495,6 +8923,14 @@ def fn(x):
         # in AOTDispatcher to make them contiguous
         decorator=toleranceOverride({torch.float32: tol(atol=1e-02, rtol=1e-02)}),
     ),
+    decorate(
+        "cholesky_inverse",
+        # Numerical differences due to tangent stride differences between
+        # eager (.sum().backward() uses contiguous tangent) and compiled
+        # (tangent strides match output strides). With ill-conditioned inputs,
+        # matmul accumulates rounding errors differently for different layouts.
+        decorator=toleranceOverride({torch.float32: tol(atol=3e02, rtol=2e-03)}),
+    ),
     decorate(
         "nn.functional.interpolate",
         "bicubic",
@@ -8551,7 +8987,6 @@ def fn(x):
     xfail(
         "nn.functional.fractional_max_pool3d", ""
     ),  # rand() received an invalid combination of arguments - g...
-    xfail("trace", ""),  # Cannot call sizes() on tensor with symbolic sizes/strides
     decorate(
         "linalg.householder_product",
         decorator=unittest.skipIf(IS_MACOS and IS_X86, "flaky"),
@@ -8817,8 +9252,8 @@ def make_compiler(self, graph_cell):
     def run_autograd(
         self,
         f: Callable,
-        fw_graph_cell: list[Optional[Callable]],
-        decompositions: Optional[dict],
+        fw_graph_cell: list[Callable | None],
+        decompositions: dict | None,
         keep_input_mutations: bool,
         dynamic: bool,
     ):
@@ -8923,6 +9358,45 @@ def _inps():
         self.assertEqual(ref_inps_after_fw, inps_after_fw)
         self.assertEqual(ref_inps_after_bw, inps_after_bw)
 
+    def test_mutations_in_bw_requires_grad_input(self):
+        # Inplace mutation of a requires_grad=True forward input during the
+        # backward pass should not trigger the check_inplace error
+        class AF(torch.autograd.Function):
+            @staticmethod
+            def forward(ctx, dummy, inplace_tensor):
+                ctx.save_for_backward(inplace_tensor)
+                return dummy.clone()
+
+            @staticmethod
+            def backward(ctx, grad_output):
+                (inplace_tensor,) = ctx.saved_tensors
+                inplace_tensor.mul_(2.0)
+                return grad_output, None
+
+        def fn(dummy, inplace_tensor):
+            return AF.apply(dummy, inplace_tensor)
+
+        def _inps():
+            dummy = torch.zeros((2,), requires_grad=True)
+            # requires_grad=True is what triggered the bug
+            inplace_tensor = torch.ones((2,), requires_grad=True)
+            return dummy, inplace_tensor
+
+        inps = _inps()
+        out = fn(*inps)
+        ref_inps_after_fw = [x.clone().detach() for x in inps]
+        out.sum().backward()
+        ref_inps_after_bw = [x.clone().detach() for x in inps]
+
+        inps = _inps()
+        out = torch.compile(fn, backend="aot_eager", fullgraph=True)(*inps)
+        inps_after_fw = [x.clone().detach() for x in inps]
+        out.sum().backward()
+        inps_after_bw = [x.clone().detach() for x in inps]
+
+        self.assertEqual(ref_inps_after_fw, inps_after_fw)
+        self.assertEqual(ref_inps_after_bw, inps_after_bw)
+
     def test_mutation_of_input_in_fw_and_bw(self):
         class AF(torch.autograd.Function):
             @staticmethod
@@ -9041,8 +9515,8 @@ def compiler(gm, example_inputs):
     def run_autograd(
         self,
         f: Callable,
-        fw_graph_cell: list[Optional[Callable]],
-        decompositions: Optional[dict],
+        fw_graph_cell: list[Callable | None],
+        decompositions: dict | None,
         keep_input_mutations: bool,
         dynamic: bool,
     ):
@@ -9064,11 +9538,11 @@ def run_autograd(
     def verify_aot_autograd(
         self,
         f,
-        inp_: Union[Callable, list[Any]],
+        inp_: Callable | list[Any],
         *,
         test_mutation: bool = False,
         keep_inp_mutations: bool = False,
-        decompositions: Optional[dict] = None,
+        decompositions: dict | None = None,
         dynamic: bool = False,
         # Only active when inp_ is Callable.
         # TODO: probably consolidate all tests to make inp a Callable.
diff --git a/test/functorch/test_codegen_backward_epilogue.py b/test/functorch/test_codegen_backward_epilogue.py
new file mode 100644
index 0000000000000..ca60681037de0
--- /dev/null
+++ b/test/functorch/test_codegen_backward_epilogue.py
@@ -0,0 +1,190 @@
+# Owner(s): ["module: functorch"]
+
+"""
+Tests for codegen'ing _backward_epilogue_functional subclass wrapping.
+
+_backward_epilogue_functional wraps backward outputs (grad inputs) back
+into tensor subclasses using a codegen'd wrap_fn. At compile time,
+codegen_backward_subclass_wrap() generates straight-line wrapping code
+with all subclass types and attr names baked in, replacing the generic
+wrap_tensor_subclasses() loop.
+
+Tests verify that a "backward_subclass_wrapper" artifact is emitted via
+trace_structured, analogous to the "subclass_wrapper" artifact emitted
+by the forward path.
+"""
+
+import logging
+from contextlib import contextmanager
+
+import torch
+import torch._functorch.config
+from torch.testing._internal.common_utils import run_tests, TestCase
+from torch.testing._internal.two_tensor import TwoTensor
+
+
+trace_log = logging.getLogger("torch.__trace")
+
+
+class TestCodegenBackwardEpilogue(TestCase):
+    @contextmanager
+    def _capture_codegen_source(self, artifact_name):
+        """Capture codegen artifacts from the structured trace log."""
+        captured: list[str] = []
+
+        class _ArtifactHandler(logging.Handler):
+            def emit(self, record):
+                metadata = getattr(record, "metadata", {})
+                if (
+                    "artifact" in metadata
+                    and metadata["artifact"].get("name") == artifact_name
+                ):
+                    payload = getattr(record, "payload", None)
+                    if payload is not None:
+                        captured.append(payload)
+
+        handler = _ArtifactHandler()
+        handler.setLevel(logging.DEBUG)
+        old_level = trace_log.level
+        trace_log.setLevel(logging.DEBUG)
+        trace_log.addHandler(handler)
+        try:
+            yield captured
+        finally:
+            trace_log.removeHandler(handler)
+            trace_log.setLevel(old_level)
+
+    def test_simple_backward_wraps_grad_inputs(self):
+        """
+        f(TwoTensor) -> TwoTensor, backward should wrap grad inputs back
+        into TwoTensor via codegen'd epilogue.
+        """
+        with self._capture_codegen_source("backward_subclass_wrapper") as captured:
+
+            @torch.compile(backend="aot_eager")
+            def f(x):
+                return x * 2
+
+            a = torch.randn(4, requires_grad=True)
+            b = torch.randn(4, requires_grad=True)
+            tt = TwoTensor(a, b)
+
+            # Run ref (eager) and test (compiled) to compare gradients
+            tt_ref = TwoTensor(
+                a.clone().detach().requires_grad_(True),
+                b.clone().detach().requires_grad_(True),
+            )
+            out_ref = tt_ref * 2
+            out_ref.sum().backward()
+
+            out = f(tt)
+            out.sum().backward()
+
+        self.assertIsInstance(out, TwoTensor)
+        self.assertIsInstance(tt.grad, TwoTensor)
+        self.assertEqual(tt.grad.a, tt_ref.grad.a)
+        self.assertEqual(tt.grad.b, tt_ref.grad.b)
+
+        self.assertGreaterEqual(
+            len(captured),
+            1,
+            "Expected backward_subclass_wrapper codegen artifact to be emitted",
+        )
+
+    def test_multi_input_backward_wraps_grad_inputs(self):
+        """
+        f(TwoTensor, TwoTensor) -> TwoTensor, backward should wrap grad
+        inputs for each subclass input.
+        """
+        with self._capture_codegen_source("backward_subclass_wrapper") as captured:
+
+            @torch.compile(backend="aot_eager")
+            def f(x, y):
+                return x * y
+
+            a = torch.randn(4, requires_grad=True)
+            b = torch.randn(4, requires_grad=True)
+            c = torch.randn(4, requires_grad=True)
+            d = torch.randn(4, requires_grad=True)
+
+            tt1_ref = TwoTensor(
+                a.clone().detach().requires_grad_(True),
+                b.clone().detach().requires_grad_(True),
+            )
+            tt2_ref = TwoTensor(
+                c.clone().detach().requires_grad_(True),
+                d.clone().detach().requires_grad_(True),
+            )
+            out_ref = tt1_ref * tt2_ref
+            out_ref.sum().backward()
+
+            tt1 = TwoTensor(a, b)
+            tt2 = TwoTensor(c, d)
+            out = f(tt1, tt2)
+            out.sum().backward()
+
+        self.assertIsInstance(out, TwoTensor)
+        self.assertIsInstance(tt1.grad, TwoTensor)
+        self.assertIsInstance(tt2.grad, TwoTensor)
+        self.assertEqual(tt1.grad.a, tt1_ref.grad.a)
+        self.assertEqual(tt1.grad.b, tt1_ref.grad.b)
+        self.assertEqual(tt2.grad.a, tt2_ref.grad.a)
+        self.assertEqual(tt2.grad.b, tt2_ref.grad.b)
+
+        self.assertGreaterEqual(
+            len(captured),
+            1,
+            "Expected backward_subclass_wrapper codegen artifact to be emitted",
+        )
+
+    def test_nested_subclass_backward_wraps_grad_inputs(self):
+        """
+        Nested TwoTensor backward should recursively wrap grad inputs.
+        """
+        with self._capture_codegen_source("backward_subclass_wrapper") as captured:
+
+            @torch.compile(backend="aot_eager")
+            def f(x):
+                return x.sin()
+
+            a1 = torch.randn(4, requires_grad=True)
+            a2 = torch.randn(4, requires_grad=True)
+            a3 = torch.randn(4, requires_grad=True)
+            a4 = torch.randn(4, requires_grad=True)
+
+            inner_a_ref = TwoTensor(
+                a1.clone().detach().requires_grad_(True),
+                a2.clone().detach().requires_grad_(True),
+            )
+            inner_b_ref = TwoTensor(
+                a3.clone().detach().requires_grad_(True),
+                a4.clone().detach().requires_grad_(True),
+            )
+            tt_ref = TwoTensor(inner_a_ref, inner_b_ref)
+            out_ref = tt_ref.sin()
+            out_ref.sum().backward()
+
+            inner_a = TwoTensor(a1, a2)
+            inner_b = TwoTensor(a3, a4)
+            tt = TwoTensor(inner_a, inner_b)
+            out = f(tt)
+            out.sum().backward()
+
+        self.assertIsInstance(out, TwoTensor)
+        self.assertIsInstance(out.a, TwoTensor)
+        self.assertIsInstance(out.b, TwoTensor)
+        self.assertIsInstance(tt.grad, TwoTensor)
+        self.assertEqual(tt.grad.a.a, tt_ref.grad.a.a)
+        self.assertEqual(tt.grad.a.b, tt_ref.grad.a.b)
+        self.assertEqual(tt.grad.b.a, tt_ref.grad.b.a)
+        self.assertEqual(tt.grad.b.b, tt_ref.grad.b.b)
+
+        self.assertGreaterEqual(
+            len(captured),
+            1,
+            "Expected backward_subclass_wrapper codegen artifact to be emitted",
+        )
+
+
+if __name__ == "__main__":
+    run_tests()
diff --git a/test/functorch/test_codegen_backward_prologue.py b/test/functorch/test_codegen_backward_prologue.py
new file mode 100644
index 0000000000000..e53551146579a
--- /dev/null
+++ b/test/functorch/test_codegen_backward_prologue.py
@@ -0,0 +1,139 @@
+# Owner(s): ["module: functorch"]
+
+"""
+Tests for codegen'ing _backward_prologue_functional non-tangent subclass
+unwrapping.
+
+_backward_prologue_functional unwraps saved tensors (non-tangent subclass
+inputs) before passing them to the compiled backward. The tangent processing
+(process_runtime_tangent) is runtime-dependent and NOT a codegen candidate,
+but the non-tangent unwrapping is pure compile-time-determined subclass
+unwrapping, identical to the forward input unwrapping already codegen'd.
+
+Tests verify that a "backward_subclass_unwrap" artifact is emitted via
+trace_structured.
+"""
+
+import logging
+from contextlib import contextmanager
+
+import torch
+import torch._functorch.config
+from torch.testing._internal.common_utils import run_tests, TestCase
+from torch.testing._internal.two_tensor import TwoTensor
+
+
+trace_log = logging.getLogger("torch.__trace")
+
+
+class TestCodegenBackwardPrologue(TestCase):
+    @contextmanager
+    def _capture_codegen_source(self, artifact_name):
+        """Capture codegen artifacts from the structured trace log."""
+        captured: list[str] = []
+
+        class _ArtifactHandler(logging.Handler):
+            def emit(self, record):
+                metadata = getattr(record, "metadata", {})
+                if (
+                    "artifact" in metadata
+                    and metadata["artifact"].get("name") == artifact_name
+                ):
+                    payload = getattr(record, "payload", None)
+                    if payload is not None:
+                        captured.append(payload)
+
+        handler = _ArtifactHandler()
+        handler.setLevel(logging.DEBUG)
+        old_level = trace_log.level
+        trace_log.setLevel(logging.DEBUG)
+        trace_log.addHandler(handler)
+        try:
+            yield captured
+        finally:
+            trace_log.removeHandler(handler)
+            trace_log.setLevel(old_level)
+
+    def test_saved_subclass_tensors_unwrapped(self):
+        """
+        When subclass tensors are saved for backward, the prologue should
+        codegen their unwrapping (non-tangent path).
+        """
+        with self._capture_codegen_source("backward_subclass_unwrap") as captured:
+
+            @torch.compile(backend="aot_eager")
+            def f(x, y):
+                return x * y
+
+            a = torch.randn(4, requires_grad=True)
+            b = torch.randn(4, requires_grad=True)
+            c = torch.randn(4, requires_grad=True)
+            d = torch.randn(4, requires_grad=True)
+
+            tt1_ref = TwoTensor(
+                a.clone().detach().requires_grad_(True),
+                b.clone().detach().requires_grad_(True),
+            )
+            tt2_ref = TwoTensor(
+                c.clone().detach().requires_grad_(True),
+                d.clone().detach().requires_grad_(True),
+            )
+            out_ref = tt1_ref * tt2_ref
+            out_ref.sum().backward()
+
+            tt1 = TwoTensor(a, b)
+            tt2 = TwoTensor(c, d)
+            out = f(tt1, tt2)
+            out.sum().backward()
+
+        self.assertEqual(tt1.grad.a, tt1_ref.grad.a)
+        self.assertEqual(tt1.grad.b, tt1_ref.grad.b)
+        self.assertEqual(tt2.grad.a, tt2_ref.grad.a)
+        self.assertEqual(tt2.grad.b, tt2_ref.grad.b)
+
+        self.assertGreaterEqual(
+            len(captured),
+            1,
+            "Expected backward_subclass_unwrap codegen artifact to be emitted",
+        )
+
+    def test_mixed_subclass_and_plain_saved_tensors(self):
+        """
+        When both subclass and plain tensors are saved, the prologue should
+        codegen unwrapping only for the subclass ones.
+        """
+        with self._capture_codegen_source("backward_subclass_unwrap") as captured:
+
+            @torch.compile(backend="aot_eager")
+            def f(x, y):
+                return x * y
+
+            a = torch.randn(4, requires_grad=True)
+            b = torch.randn(4, requires_grad=True)
+
+            tt_ref = TwoTensor(
+                a.clone().detach().requires_grad_(True),
+                b.clone().detach().requires_grad_(True),
+            )
+            plain_ref = torch.randn(4, requires_grad=True)
+            out_ref = tt_ref * plain_ref
+            out_ref.sum().backward()
+
+            tt = TwoTensor(a, b)
+            plain = plain_ref.clone().detach().requires_grad_(True)
+            out = f(tt, plain)
+            out.sum().backward()
+
+        self.assertEqual(tt.grad.a, tt_ref.grad.a)
+        self.assertEqual(tt.grad.b, tt_ref.grad.b)
+        self.assertEqual(plain.grad, plain_ref.grad)
+
+        self.assertGreaterEqual(
+            len(captured),
+            1,
+            "Expected backward_subclass_unwrap codegen artifact to be emitted",
+        )
+
+
+if __name__ == "__main__":
+    run_tests()
diff --git a/test/functorch/test_codegen_debug_assert.py b/test/functorch/test_codegen_debug_assert.py
new file mode 100644
index 0000000000000..ca82d045ccc99
--- /dev/null
+++ b/test/functorch/test_codegen_debug_assert.py
@@ -0,0 +1,123 @@
+# Owner(s): ["module: functorch"]
+
+"""
+Tests for codegen'ing DebugAssertWrapper.
+
+The codegen'd assertion function emits checks only for the specific arg
+indices where requires_grad was False at compile time. Positions where
+requires_grad=True are safe and generate no runtime check, replacing
+the closure that iterated over all args.
+
+Enabled via torch._functorch.config.debug_assert = True.
+
+Tests verify that a "debug_assert_wrapper" artifact is emitted via
+trace_structured.
+"""
+
+import logging
+from contextlib import contextmanager
+
+import torch
+import torch._functorch.config
+from torch.testing._internal.common_utils import run_tests, TestCase
+
+
+trace_log = logging.getLogger("torch.__trace")
+
+
+class TestCodegenDebugAssert(TestCase):
+    @contextmanager
+    def _capture_codegen_source(self, artifact_name):
+        """Capture codegen artifacts from the structured trace log."""
+        captured: list[str] = []
+
+        class _ArtifactHandler(logging.Handler):
+            def emit(self, record):
+                metadata = getattr(record, "metadata", {})
+                if (
+                    "artifact" in metadata
+                    and metadata["artifact"].get("name") == artifact_name
+                ):
+                    payload = getattr(record, "payload", None)
+                    if payload is not None:
+                        captured.append(payload)
+
+        handler = _ArtifactHandler()
+        handler.setLevel(logging.DEBUG)
+        old_level = trace_log.level
+        trace_log.setLevel(logging.DEBUG)
+        trace_log.addHandler(handler)
+        try:
+            yield captured
+        finally:
+            trace_log.removeHandler(handler)
+            trace_log.setLevel(old_level)
+
+    @torch._functorch.config.patch(debug_assert=True)
+    def test_mixed_requires_grad(self):
+        """
+        With debug_assert=True, the wrapper should codegen assertions for
+        inputs compiled without requires_grad, and skip those with
+        requires_grad=True.
+        """
+        with self._capture_codegen_source("debug_assert_wrapper") as captured:
+
+            @torch.compile(backend="aot_eager")
+            def f(x, y):
+                return x + y
+
+            x = torch.randn(4, requires_grad=True)
+            y = torch.randn(4)
+            out = f(x, y)
+
+        self.assertEqual(out, x + y)
+
+        self.assertEqual(len(captured), 1)
+
+    @torch._functorch.config.patch(debug_assert=True)
+    def test_all_requires_grad(self):
+        """
+        All inputs with requires_grad=True. Codegen should emit no
+        assertions (all positions are safe), but the artifact should
+        still be emitted.
+        """
+        with self._capture_codegen_source("debug_assert_wrapper") as captured:
+
+            @torch.compile(backend="aot_eager")
+            def f(x, y, z):
+                return x + y + z
+
+            x = torch.randn(4, requires_grad=True)
+            y = torch.randn(4, requires_grad=True)
+            z = torch.randn(4, requires_grad=True)
+            out = f(x, y, z)
+
+        self.assertEqual(out, x + y + z)
+
+        self.assertEqual(len(captured), 1)
+
+    @torch._functorch.config.patch(debug_assert=True)
+    def test_some_no_grad_inputs(self):
+        """
+        Mix of requires_grad and non-requires_grad inputs going through
+        the training path. Codegen should emit assertions only for the
+        non-grad positions.
+        """
+        with self._capture_codegen_source("debug_assert_wrapper") as captured:
+
+            @torch.compile(backend="aot_eager")
+            def f(x, y, z):
+                return x * y + z
+
+            x = torch.randn(4, requires_grad=True)
+            y = torch.randn(4)
+            z = torch.randn(4)
+            out = f(x, y, z)
+
+        self.assertEqual(out, x * y + z)
+
+        self.assertEqual(len(captured), 1)
+
+
+if __name__ == "__main__":
+    run_tests()
diff --git a/test/functorch/test_codegen_dedup.py b/test/functorch/test_codegen_dedup.py
new file mode 100644
index 0000000000000..335572b4b040b
--- /dev/null
+++ b/test/functorch/test_codegen_dedup.py
@@ -0,0 +1,151 @@
+# Owner(s): ["module: functorch"]
+
+"""
+Tests for codegen'ing AOTDedupeWrapper.
+
+The codegen'd remove_dupe_args emits straight-line index selections like
+[args[0], args[2], args[5]] with all indices baked in as literals, replacing
+the closure-based zip + filter over keep_arg_mask.
+
+Strategy 2 (the dedup post_compile path) is triggered when duplicate args
+have mutations, so strategy 1 (leafification) can't handle them. Dynamo
+already deduplicates inputs, so these tests use aot_function directly.
+
+Tests verify that a "dedup_wrapper" artifact is emitted via trace_structured.
+"""
+
+import logging
+from contextlib import contextmanager
+
+import torch
+import torch._functorch.config
+from torch._functorch.aot_autograd import aot_function
+from torch.testing._internal.common_utils import run_tests, TestCase
+
+
+trace_log = logging.getLogger("torch.__trace")
+
+
+def _nop_compiler(gm, example_inputs):  # type: ignore[no-untyped-def]
+    return gm.forward
+
+
+class TestCodegenDedup(TestCase):
+    @contextmanager
+    def _capture_codegen_source(self, artifact_name):
+        """Capture codegen artifacts from the structured trace log."""
+        captured: list[str] = []
+
+        class _ArtifactHandler(logging.Handler):
+            def emit(self, record):
+                metadata = getattr(record, "metadata", {})
+                if (
+                    "artifact" in metadata
+                    and metadata["artifact"].get("name") == artifact_name
+                ):
+                    payload = getattr(record, "payload", None)
+                    if payload is not None:
+                        captured.append(payload)
+
+        handler = _ArtifactHandler()
+        handler.setLevel(logging.DEBUG)
+        old_level = trace_log.level
+        trace_log.setLevel(logging.DEBUG)
+        trace_log.addHandler(handler)
+        try:
+            yield captured
+        finally:
+            trace_log.removeHandler(handler)
+            trace_log.setLevel(old_level)
+
+    def test_duplicate_args_with_mutation(self):
+        """
+        When the same tensor is passed as two args and one position mutates,
+        strategy 2 dedup kicks in. The codegen should emit straight-line
+        arg selection.
+        """
+        with self._capture_codegen_source("dedup_wrapper") as captured:
+
+            def f(a, b):
+                b.mul_(2)
+                return a + b
+
+            compiled_f = aot_function(f, _nop_compiler)
+            x = torch.randn(4)
+            x_ref = x.clone()
+            out = compiled_f(x, x)
+
+        self.assertEqual(x, x_ref * 2)
+        self.assertEqual(out, x_ref * 4)
+
+        self.assertGreaterEqual(
+            len(captured),
+            1,
+            "Expected dedup_wrapper codegen artifact to be emitted",
+        )
+        source = captured[0]
+        self.assertIn("args[0]", source)
+        self.assertNotIn("args[1]", source)
+
+    def test_three_way_duplicate_with_mutation(self):
+        """
+        Three-way duplication where the last position mutates.
+        """
+        with self._capture_codegen_source("dedup_wrapper") as captured:
+
+            def f(a, b, c):
+                c.add_(1)
+                return a + b + c
+
+            compiled_f = aot_function(f, _nop_compiler)
+            x = torch.randn(4)
+            x_ref = x.clone()
+            out = compiled_f(x, x, x)
+
+        self.assertEqual(x, x_ref + 1)
+        self.assertEqual(out, (x_ref + 1) * 3)
+
+        self.assertGreaterEqual(
+            len(captured),
+            1,
+            "Expected dedup_wrapper codegen artifact to be emitted",
+        )
+        source = captured[0]
+        self.assertIn("args[0]", source)
+        self.assertNotIn("args[1]", source)
+        self.assertNotIn("args[2]", source)
+
+    def test_partial_duplicate_with_mutation(self):
+        """
+        Two args are duplicates (with mutation), third is distinct.
+        Codegen should select [args[0], args[2]] dropping the duplicate.
+        """
+        with self._capture_codegen_source("dedup_wrapper") as captured:
+
+            def f(a, b, c):
+                b.mul_(3)
+                return a + b + c
+
+            compiled_f = aot_function(f, _nop_compiler)
+            x = torch.randn(4)
+            y = torch.randn(4)
+            x_ref = x.clone()
+            y_ref = y.clone()
+            out = compiled_f(x, x, y)
+
+        self.assertEqual(x, x_ref * 3)
+        self.assertEqual(out, x_ref * 3 + x_ref * 3 + y_ref)
+
+        self.assertGreaterEqual(
+            len(captured),
+            1,
+            "Expected dedup_wrapper codegen artifact to be emitted",
+        )
+        source = captured[0]
+        self.assertIn("args[0]", source)
+        self.assertNotIn("args[1]", source)
+        self.assertIn("args[2]", source)
+
+
+if __name__ == "__main__":
+    run_tests()
diff --git a/test/functorch/test_control_flow.py b/test/functorch/test_control_flow.py
index 4cdc7c5b72b73..6ed0b5841b43d 100644
--- a/test/functorch/test_control_flow.py
+++ b/test/functorch/test_control_flow.py
@@ -7,11 +7,19 @@
 import torch.utils._pytree as pytree
 from functorch.experimental import control_flow
 from functorch.experimental.control_flow import cond
-from torch._dynamo.testing import EagerAndRecordGraphs, normalize_gm
+from torch._dynamo.testing import (
+    AotEagerAndRecordGraphs,
+    EagerAndRecordGraphs,
+    normalize_gm,
+)
 from torch._higher_order_ops.associative_scan import (
     _fake_associative_scan,
     associative_scan,
 )
+from torch._higher_order_ops.cudagraph_conditional_nodes import (
+    ControlFlowOpWarmupDispatchMode,
+    CUDAGraphCaptureControlFlowOpDispatchMode,
+)
 from torch._higher_order_ops.map import _fake_map
 from torch._higher_order_ops.scan import _fake_scan, scan
 from torch._higher_order_ops.schema import HopSchemaGenerator
@@ -34,12 +42,90 @@
     run_tests,
     skipIfCrossRef,
     skipIfTorchDynamo,
+    TEST_CUDA_GRAPH_CONDITIONAL_NODES,
     TEST_WITH_CROSSREF,
     TEST_WITH_TORCHDYNAMO,
     TestCase,
 )
 
 
+@contextlib.contextmanager
+def check_cudagraphs_not_skipped(test_case):
+    counters = torch._dynamo.utils.counters
+    old_cudagraph_skips = counters["inductor"]["cudagraph_skips"]
+    try:
+        yield
+    finally:
+        # reset before the assert, because otherwise the reset is
+        # skipped in case of assertion error
+        new_cudagraph_skips = counters["inductor"]["cudagraph_skips"]
+        counters["inductor"]["cudagraph_skips"] = old_cudagraph_skips
+        test_case.assertEqual(
+            counters["inductor"]["cudagraph_skips"], new_cudagraph_skips
+        )
+
+
+def _check_compile_cudagraph_backend(test_case, fn, args):
+    # test cudagraphs backend
+    cudagraphs_compiled_fn = torch.compile(fn, backend="cudagraphs")
+    # We run 3 times.
+    # This is what cuda graph trees does for the first 3 runs:
+    # 1) run in eager mode, for warmup.
+    # 2) do stream capture followed by graph replay.
+    # 3 and beyond) do graph replay
+    # So we need to get to iteration 3 to test all ways of running.
+    outputs = []
+    for i in range(3):
+        with check_cudagraphs_not_skipped(test_case):
+            torch.compiler.cudagraph_mark_step_begin()
+            outputs.append(
+                pytree.tree_map(
+                    lambda x: x.detach().clone() if isinstance(x, torch.Tensor) else x,
+                    cudagraphs_compiled_fn(*args),
+                )
+            )
+    eager_res = fn(*args)
+    for output in outputs:
+        test_case.assertEqual(eager_res, output)
+
+
+def _check_compile_many_backends_with_cudagraph(test_case, fn, args):
+    for backend in ["eager_no_compile", "eager", "aot_eager"]:
+        _check_compile_any_backend_with_cudagraph(test_case, fn, args, backend)
+
+
+def _check_compile_any_backend_with_cudagraph(test_case, fn, args, backend):
+    if backend == "eager_no_compile":
+        compiled_fn = fn
+    else:
+        compiled_fn = torch.compile(fn, backend=backend)
+    outputs = []
+
+    side_stream = torch.cuda.Stream()
+
+    with torch.cuda.stream(side_stream), ControlFlowOpWarmupDispatchMode():
+        warmup_output = compiled_fn(*args)
+    outputs.append(warmup_output)
+
+    graph = torch.cuda.CUDAGraph()
+    with (
+        torch.cuda.graph(graph, stream=side_stream),
+        CUDAGraphCaptureControlFlowOpDispatchMode(),
+    ):
+        captured_output = compiled_fn(*args)
+    with torch.cuda.stream(side_stream):
+        eager_res = fn(*args)
+    graph.replay()
+    outputs.append(
+        pytree.tree_map(
+            lambda x: x.clone() if isinstance(x, torch.Tensor) else x,
+            captured_output,
+        )
+    )
+    for output in outputs:
+        test_case.assertEqual(eager_res, output)
+
+
 # TODO: pull these helpers from AOTAutograd later
 def to_fun(t):
     if isinstance(t, torch.Tensor):
@@ -51,7 +137,8 @@ def from_fun(t):
     if not isinstance(t, FunctionalTensor):
         # quick sanity assert
         if isinstance(t, torch.Tensor):
-            assert not torch._is_functional_tensor(t)
+            if torch._is_functional_tensor(t):
+                raise AssertionError("Expected tensor to not be a functional tensor")
         return t
     torch._sync(t)
     return torch._from_functional_tensor(t.elem)
@@ -68,7 +155,8 @@ def to_fun_old(t):
 def from_fun_old(t):
     # quick sanity assert
     if isinstance(t, torch.Tensor):
-        assert torch._is_functional_tensor(t)
+        if not torch._is_functional_tensor(t):
+            raise AssertionError("Expected tensor to be a functional tensor")
         torch._sync(t)
         return torch._from_functional_tensor(t)
     return t
@@ -2300,6 +2388,44 @@ def test_scan_compile_cnt(self, reverse, device):
             )
             self.assertEqual(cnt.frame_count, 6)
 
+    def test_scan_operator_call_count(self):
+        from torch._higher_order_ops.scan import generic_scan, wrap_combine_fn_flat
+
+        counter = [0]
+
+        def counting_combine(carry, x):
+            counter[0] += 1
+            return carry + x, x
+
+        init = [torch.zeros(3)]
+        xs = [torch.ones(5, 3)]
+        combine_flat = functools.partial(
+            wrap_combine_fn_flat,
+            combine_fn=counting_combine,
+            spec_init=pytree.tree_flatten(init)[1],
+            spec_xs=pytree.tree_flatten(xs)[1],
+            num_init_leaves=len(init),
+            num_inp_leaves=len(xs),
+        )
+
+        counter[0] = 0
+        result = generic_scan(combine_flat, init, xs)  # noqa: F841
+        self.assertEqual(counter[0], 5)
+
+        # Single-element scan should call operator exactly once.
+        xs_one = [torch.ones(1, 3)]
+        combine_flat_one = functools.partial(
+            wrap_combine_fn_flat,
+            combine_fn=counting_combine,
+            spec_init=pytree.tree_flatten(init)[1],
+            spec_xs=pytree.tree_flatten(xs_one)[1],
+            num_init_leaves=len(init),
+            num_inp_leaves=len(xs_one),
+        )
+        counter[0] = 0
+        generic_scan(combine_flat_one, init, xs_one)
+        self.assertEqual(counter[0], 1)
+
     @skipIfTorchDynamo("don't test compile on compile")
     def test_scan_init_scanned_0(self):
         # Only init and no input
@@ -3707,9 +3833,10 @@ def _check_autograd(self, result, result_exp, autograd_param):
         result_exp_flatten = [r for r in result_exp_flatten if r.requires_grad]
 
         # Check the result and parameter lists
-        assert len(result_flatten) == len(result_exp_flatten), (
-            "The number of elements requiring gradients is different for the results and the expected results"
-        )
+        if len(result_flatten) != len(result_exp_flatten):
+            raise AssertionError(
+                "The number of elements requiring gradients is different for the results and the expected results"
+            )
 
         grad_exp_init = [torch.ones_like(el) for el in result_exp_flatten]
         expected_grads = torch.autograd.grad(
@@ -3994,7 +4121,8 @@ def test_associative_scan_non_contiguous_tensor(
             .view(10, 3)
             .t()
         )
-        assert not x.is_contiguous()
+        if x.is_contiguous():
+            raise AssertionError("Expected x to not be contiguous")
 
         kwargs = {
             "dim": 0,
@@ -5418,6 +5546,154 @@ def body_fn(c, a, b):
                 torch.randn(2, 3),
             )
 
+    @unittest.skipIf(
+        not TEST_CUDA_GRAPH_CONDITIONAL_NODES,
+        "CUDA 12.4 or greater is required for CUDA Graphs with conditional nodes",
+    )
+    def test_cond_traced_not_nested_cudagraphs(self):
+        def true_fn(x):
+            return x.sin()
+
+        def false_fn(x):
+            return x.cos()
+
+        def f(x, y):
+            return cond(y, true_fn, false_fn, [x])
+
+        x = torch.randn(4).cuda()
+        true_pred = torch.tensor(True).cuda()
+        false_pred = torch.tensor(False).cuda()
+
+        _check_compile_cudagraph_backend(self, f, [x, true_pred])
+        _check_compile_cudagraph_backend(self, f, [x, false_pred])
+        _check_compile_many_backends_with_cudagraph(
+            self,
+            f,
+            [x, true_pred],
+        )
+        _check_compile_many_backends_with_cudagraph(
+            self,
+            f,
+            [x, false_pred],
+        )
+
+    @unittest.skipIf(
+        not TEST_CUDA_GRAPH_CONDITIONAL_NODES,
+        "CUDA 12.4 or greater is required for CUDA Graphs with conditional nodes",
+    )
+    def test_cond_pin_memory_cudagraphs(self):
+        # Make sure that pinned host memory allocations get assigned
+        # to a private pool correctly during stream capture, even
+        # inside of conditional nodes.
+
+        # Ideally, we would call torch.Tensor.pin_memory() directly,
+        # but that is not allowed on fake tensors, so instead we call
+        # an op whose C++ implementation used pinned memory
+        # internally.
+        sizes = [3, 4, 3]
+
+        def true_fn(x):
+            return torch.split_with_sizes_copy(x, sizes)
+
+        def false_fn(x):
+            return torch.split_with_sizes_copy(2 * x, sizes)
+
+        def f(x, y):
+            return cond(y, true_fn, false_fn, [x])
+
+        x = torch.randn(10).cuda()
+        true_pred = torch.tensor(True).cuda()
+        false_pred = torch.tensor(False).cuda()
+
+        _check_compile_cudagraph_backend(self, f, [x, true_pred])
+        _check_compile_cudagraph_backend(self, f, [x, false_pred])
+        _check_compile_many_backends_with_cudagraph(
+            self,
+            f,
+            [x, true_pred],
+        )
+        _check_compile_many_backends_with_cudagraph(
+            self,
+            f,
+            [x, false_pred],
+        )
+
+    @unittest.skipIf(
+        not TEST_CUDA_GRAPH_CONDITIONAL_NODES,
+        "CUDA 12.4 or greater is required for CUDA Graphs with conditional nodes",
+    )
+    def test_cond_traced_triply_nested_cudagraphs(self):
+        def level3_true(x):
+            return x.sin()
+
+        def level3_false(x):
+            return x.cos()
+
+        def level2_true(x, p2):
+            return cond(p2, level3_true, level3_false, [x])
+
+        def level2_false(x, p2):
+            return cond(p2, lambda x: x + 1, lambda x: x - 1, [x])
+
+        def level1_true(x, p1, p2):
+            return cond(p1, level2_true, level2_false, [x, p2])
+
+        def level1_false(x, p1, p2):
+            return cond(p1, level2_false, level2_true, [x, p2])
+
+        def f(x, p0, p1, p2):
+            return cond(p0, level1_true, level1_false, [x, p1, p2])
+
+        x = torch.randn(4).cuda()
+
+        test_inputs = [
+            [
+                x,
+                torch.tensor(True).cuda(),
+                torch.tensor(True).cuda(),
+                torch.tensor(False).cuda(),
+            ],
+            [
+                x,
+                torch.tensor(False).cuda(),
+                torch.tensor(False).cuda(),
+                torch.tensor(True).cuda(),
+            ],
+        ]
+
+        for args in test_inputs:
+            _check_compile_cudagraph_backend(self, f, args)
+            _check_compile_many_backends_with_cudagraph(self, f, args)
+
+    @unittest.skipIf(
+        not TEST_CUDA_GRAPH_CONDITIONAL_NODES,
+        "CUDA 12.4 or greater is required for CUDA Graphs with conditional nodes",
+    )
+    def test_cond_traced_record_stream_reuse(self):
+        torch.cuda.memory._set_allocator_settings(
+            "graph_capture_record_stream_reuse:True"
+        )
+        try:
+            predicate = torch.tensor(True, device="cuda")
+
+            def true_fn():
+                return torch.zeros(8, device="cuda"), torch.zeros(8, device="cuda")
+
+            def false_fn():
+                return torch.zeros(8, device="cuda"), torch.zeros(8, device="cuda")
+
+            g = torch.cuda.CUDAGraph()
+            with self.assertRaisesRegex(
+                RuntimeError,
+                "graph_capture_record_stream_reuse:True",
+            ):
+                with torch.cuda.graph(g), CUDAGraphCaptureControlFlowOpDispatchMode():
+                    torch.cond(predicate, true_fn, false_fn, [])
+        finally:
+            torch.cuda.memory._set_allocator_settings(
+                "graph_capture_record_stream_reuse:False"
+            )
+
     def test_while_loop_nested_traced(self):
         fn, inp = WHILE_LOOP_TESTS["nested"]
         graphs = self._check_tracing(fn, inp)
@@ -5496,7 +5772,10 @@ def _wrap_with_functionalize(self, fn, func_type):
         elif func_type == "functorch":
             fn = torch.func.functionalize(fn)
         else:
-            assert func_type == "no"
+            if func_type != "no":
+                raise AssertionError(
+                    f"Expected func_type to be 'no', got {func_type!r}"
+                )
         return fn, mode
 
     @parametrize("func_type", ["no", "cpp", "python", "functorch"])
@@ -5669,10 +5948,9 @@ def test_while_loop_simple_with_linear_compile_check_graph(self):
         torch.compile(fn, backend=backend)(*inp)
         self.assertEqual(len(backend.graphs), 1)
         gm = backend.graphs[0]
-        if torch._dynamo.config.inline_inbuilt_nn_modules:
-            self.assertExpectedInline(
-                normalize_gm(gm.print_readable(print_output=False)),
-                """\
+        self.assertExpectedInline(
+            normalize_gm(gm.print_readable(print_output=False)),
+            """\
 class GraphModule(torch.nn.Module):
     def forward(self, L_iter_: "i64[]", L_x_: "f32[2, 2]", L_self_buffers_dec_: "i64[]", L_self_modules_linear_parameters_weight_: "f32[2, 2]", L_self_modules_linear_parameters_bias_: "f32[2]"):
         l_iter_ = L_iter_
@@ -5700,7 +5978,7 @@ def forward(self, child_2: "i64[]", child_3: "f32[2, 2]", l_self_buffers_dec__co
             child_4: "f32[2, 2]" = torch._C._nn.linear(child_3, l_self_modules_linear_parameters_weight__body_fn, l_self_modules_linear_parameters_bias__body_fn);  child_3 = l_self_modules_linear_parameters_weight__body_fn = l_self_modules_linear_parameters_bias__body_fn = None
             return (child, child_4)
 """,  # noqa: B950
-            )
+        )
 
     def test_while_loop_nested2_traced(self):
         fn, inp = WHILE_LOOP_TESTS["nested2"]
@@ -5834,6 +6112,21 @@ def f(x, pred, pred2):
             f(x, torch.tensor(True), torch.tensor(True)),
         )
 
+        if TEST_CUDA_GRAPH_CONDITIONAL_NODES:
+            x_cuda = x.cuda()
+            true_pred = torch.tensor(True).cuda()
+
+            _check_compile_cudagraph_backend(
+                self,
+                f,
+                [x_cuda, true_pred, true_pred],
+            )
+            _check_compile_many_backends_with_cudagraph(
+                self,
+                f,
+                [x_cuda, true_pred, true_pred],
+            )
+
     def test_cond_functionalized(self):
         def true_fn(x):
             y = x.sin()
@@ -5847,6 +6140,9 @@ def f(x):
             pred = x.shape[0] == 1
             return cond(pred, true_fn, false_fn, [x])
 
+        def f_(x, y):
+            return cond(y, true_fn, false_fn, [x])
+
         example_inputs = (torch.ones(4, 5),)
         functional_f = torch.func.functionalize(f)
         self.assertEqual(functional_f(*example_inputs), f(*example_inputs))
@@ -5865,6 +6161,13 @@ def f(x):
 
         self.assertEqual(graph_module(*example_inputs), f(*example_inputs))
 
+        if TEST_CUDA_GRAPH_CONDITIONAL_NODES:
+            pred = torch.tensor(example_inputs[0].shape[0] == 1, device="cuda")
+            _check_compile_cudagraph_backend(self, f_, [torch.ones(4, 5).cuda(), pred])
+            _check_compile_many_backends_with_cudagraph(
+                self, f_, [torch.ones(4, 5).cuda(), pred]
+            )
+
     def test_cond_accepts_torch_function_as_inputs(self):
         a = torch.randn(3, 4)
         b = torch.randn(3, 4)
@@ -6398,6 +6701,18 @@ def forward(self, arg0_1):
     return (mul,)""",
         )
 
+        if TEST_CUDA_GRAPH_CONDITIONAL_NODES:
+            _check_compile_cudagraph_backend(
+                self,
+                f,
+                [x.cuda(), torch.tensor(False).cuda(), torch.tensor(False).cuda()],
+            )
+            _check_compile_many_backends_with_cudagraph(
+                self,
+                f,
+                [x.cuda(), torch.tensor(False).cuda(), torch.tensor(False).cuda()],
+            )
+
     def test_raise_error_on_mismatch_type_size(self):
         def true_fn(x):
             return x.sin()
@@ -7175,7 +7490,8 @@ def forward(self, x_1):
         )
 
     def _check_closure_correctly_lifted(self, f, *, args, exp_res, exp_arg_num):
-        assert isinstance(args, (tuple, list))
+        if not isinstance(args, (tuple, list)):
+            raise AssertionError(f"Expected args to be tuple or list, got {type(args)}")
         self.assertEqual(f(*args), exp_res)
         gm = make_fx(f)(*args)
         self.assertEqual(gm(*args), exp_res)
@@ -8933,16 +9249,42 @@ def mutate_f(x):
             ):
                 torch.compile(fn)(f, x)
 
-            with self.assertRaisesRegex(
-                # Should be
-                # torch._dynamo.exc.Unsupported,
-                # "Encountered aliasing during higher order op tracing for HOP.*"
-                torch._dynamo.exc.UncapturedHigherOrderOpError,
-                r"Higher Order Operator: torch\.cond",
-            ):
-                with torch.inference_mode(inference_mode):
+            with torch.inference_mode(inference_mode):
+                if not inference_mode:
+                    with self.assertRaisesRegex(
+                        # Should be
+                        # torch._dynamo.exc.Unsupported,
+                        # "Encountered aliasing during higher order op tracing for HOP.*"
+                        torch._dynamo.exc.UncapturedHigherOrderOpError,
+                        r"Higher Order Operator: torch\.cond",
+                    ):
+                        torch.compile(fn)(f, x)
+                else:
                     torch.compile(fn)(f, x)
 
+    @requires_cuda
+    @parametrize("device", ["cuda", "cpu"])
+    def test_cond_input_mutation(self, device):
+        predicate_true = torch.tensor(True, device=device)
+        predicate_false = torch.tensor(False, device=device)
+        org_data = torch.ones(2, 2, device=device)
+
+        def fn(predicate, data):
+            return torch.cond(
+                predicate, lambda x: x + 1, lambda x: x.sin_().add_(2), [data]
+            )
+
+        with torch.no_grad():
+            expected = org_data.sin() + 2
+            data = org_data.clone()
+            output = torch.compile(fn)(predicate_false, data)
+            self.assertEqual(output, expected)
+            self.assertIsNot(output, data)
+
+            data = org_data.clone()
+            output = torch.compile(fn)(predicate_true, data)
+            self.assertEqual(output, org_data + 1)
+
     @skipIfTorchDynamo("Graph is not captured correctly when test with dynamo")
     def test_while_loop_unbacked_bindings(self):
         m, args = WHILE_LOOP_TESTS["pytree_int_carry"]
@@ -9341,6 +9683,238 @@ def false_fn(x, y):
         self.assertEqual(compiled_out, out)
 
 
+class TestAutoFunctionalizeControlFlow(TestCase):
+    def check(self, gen_fn, args, device, dynamic) -> torch.fx.GraphModule:
+        args = pytree.tree_map(lambda t: t.to(device=device), args)
+
+        def _clone(args):
+            return [
+                arg.clone() if isinstance(arg, torch.Tensor) else arg for arg in args
+            ]
+
+        def _new_fn():
+            mod_or_fn = gen_fn()
+            if isinstance(mod_or_fn, torch.nn.Module):
+                mod_or_fn.to(device)
+            return mod_or_fn
+
+        # Only support input mutation in inference
+        cloned_args = [_clone(args) for _ in range(3)]
+        with torch.no_grad():
+            exp = _new_fn()(*cloned_args[0])
+        backend = AotEagerAndRecordGraphs()
+        torch._dynamo.reset()
+        with torch.no_grad():
+            eager_out = torch.compile(
+                _new_fn(), backend=backend, fullgraph=True, dynamic=dynamic
+            )(*cloned_args[1])
+        torch._dynamo.reset()
+        with torch.no_grad():
+            inductor_out = torch.compile(
+                _new_fn(), backend="inductor", fullgraph=True, dynamic=dynamic
+            )(*cloned_args[2])
+
+        self.assertEqual(exp, eager_out)
+        self.assertEqual(exp, inductor_out)
+        self.assertEqual(cloned_args[0], cloned_args[1])
+        self.assertEqual(cloned_args[0], cloned_args[2])
+        return backend.fw_graphs[0]
+
+    @requires_cuda
+    @unittest.skipIf(not SM70OrLater, "triton")
+    @parametrize("device", ["cuda", "cpu"])
+    @parametrize("dynamic", [True, False])
+    def test_cond_auto_functionalize_input_mutation(self, device, dynamic):
+        class M(torch.nn.Module):
+            def forward(self, x, y):
+                def true_fn(x):
+                    x.add_(1)
+                    return x.sin()
+
+                x = x.clone()
+                ret = torch.cond(x.sum() > 0, true_fn, true_fn, (x,))
+                return y + ret
+
+        x, y = (
+            torch.randn(3, 4, requires_grad=True),
+            torch.randn(3, 4, requires_grad=True),
+        )
+        fw_gm = self.check(M, (x, y), device, dynamic)
+        if not TEST_WITH_CROSSREF and not dynamic and device == "cuda":
+            self.assertExpectedInline(
+                normalize_gm(fw_gm.print_readable(print_output=False)),
+                """\
+class <lambda>(torch.nn.Module):
+    def forward(self, arg0_1: "f32[3, 4]", arg1_1: "f32[3, 4]"):
+        clone: "f32[3, 4]" = torch.ops.aten.clone.default(arg0_1);  arg0_1 = None
+
+        sum_1: "f32[]" = torch.ops.aten.sum.default(clone)
+        gt: "b8[]" = torch.ops.aten.gt.Scalar(sum_1, 0);  sum_1 = None
+        auto_functionalized_subgraph_0 = self.auto_functionalized_subgraph_0
+        auto_functionalized_subgraph_1 = self.auto_functionalized_subgraph_1
+        _tree_spec_constant0 = self._tree_spec_constant0
+        auto_functionalized_v2 = torch.ops.higher_order.auto_functionalized_v2(torch.ops.higher_order.cond, pred = gt, true_fn = auto_functionalized_subgraph_0, false_fn = auto_functionalized_subgraph_1, _operand0_base_index = 0, _all_bases = [clone], _op_schema = _tree_spec_constant0);  gt = auto_functionalized_subgraph_0 = auto_functionalized_subgraph_1 = clone = _tree_spec_constant0 = None
+        getitem: "f32[3, 4]" = auto_functionalized_v2[0];  auto_functionalized_v2 = None
+
+        add: "f32[3, 4]" = torch.ops.aten.add.Tensor(arg1_1, getitem);  arg1_1 = getitem = None
+        return (add,)
+
+    class auto_functionalized_subgraph_0(torch.nn.Module):
+        def forward(self, arg0_1: "f32[3, 4]"):
+            add: "f32[3, 4]" = torch.ops.aten.add.Tensor(arg0_1, 1)
+            sin: "f32[3, 4]" = torch.ops.aten.sin.default(add)
+            copy_: "f32[3, 4]" = torch.ops.aten.copy_.default(arg0_1, add);  arg0_1 = add = copy_ = None
+            return (sin,)
+
+    class auto_functionalized_subgraph_1(torch.nn.Module):
+        def forward(self, arg0_1: "f32[3, 4]"):
+            add: "f32[3, 4]" = torch.ops.aten.add.Tensor(arg0_1, 1)
+            sin: "f32[3, 4]" = torch.ops.aten.sin.default(add)
+            copy_: "f32[3, 4]" = torch.ops.aten.copy_.default(arg0_1, add);  arg0_1 = add = copy_ = None
+            return (sin,)
+""",  # noqa: B950
+            )
+
+    @requires_cuda
+    @unittest.skipIf(not SM70OrLater, "triton")
+    @parametrize("device", ["cuda", "cpu"])
+    @parametrize("dynamic", [True, False])
+    def test_cond_auto_functionalize_buffer_mutation(self, device, dynamic):
+        class M(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.register_buffer(
+                    "buf", torch.ones(8, requires_grad=False, device=device)
+                )
+
+            def forward(self, p, x):
+                def true_fn(x):
+                    x.add_(1)
+                    self.buf.add_(1)
+                    return x + self.buf
+
+                x = x.clone()
+                out = torch.cond(p, true_fn, true_fn, (x,))
+                return x + self.buf + out
+
+        p, x = torch.tensor(True), torch.randn(1, requires_grad=True)
+        fw_gm = self.check(M, (p, x), device, dynamic)
+        if not TEST_WITH_CROSSREF and not dynamic and device == "cuda":
+            self.assertExpectedInline(
+                normalize_gm(fw_gm.print_readable(print_output=False)),
+                """\
+class <lambda>(torch.nn.Module):
+    def forward(self, arg0_1: "f32[1]", arg1_1: "b8[]", arg2_1: "f32[8]"):
+        clone: "f32[1]" = torch.ops.aten.clone.default(arg0_1);  arg0_1 = None
+
+        auto_functionalized_subgraph_0 = self.auto_functionalized_subgraph_0
+        auto_functionalized_subgraph_1 = self.auto_functionalized_subgraph_1
+        _tree_spec_constant0 = self._tree_spec_constant0
+        auto_functionalized_v2 = torch.ops.higher_order.auto_functionalized_v2(torch.ops.higher_order.cond, pred = arg1_1, true_fn = auto_functionalized_subgraph_0, false_fn = auto_functionalized_subgraph_1, _operand0_base_index = 0, _operand1_base_index = 1, _all_bases = [arg2_1, clone], _op_schema = _tree_spec_constant0);  arg1_1 = auto_functionalized_subgraph_0 = auto_functionalized_subgraph_1 = clone = _tree_spec_constant0 = None
+        getitem: "f32[8]" = auto_functionalized_v2[0]
+        getitem_1: "f32[8]" = auto_functionalized_v2[1]
+        getitem_2: "f32[1]" = auto_functionalized_v2[2];  auto_functionalized_v2 = None
+
+        add: "f32[8]" = torch.ops.aten.add.Tensor(getitem_2, getitem_1);  getitem_2 = None
+        add_1: "f32[8]" = torch.ops.aten.add.Tensor(add, getitem);  add = getitem = None
+
+        copy_: "f32[8]" = torch.ops.aten.copy_.default(arg2_1, getitem_1);  arg2_1 = getitem_1 = copy_ = None
+        return (add_1,)
+
+    class auto_functionalized_subgraph_0(torch.nn.Module):
+        def forward(self, arg0_1: "f32[8]", arg1_1: "f32[1]"):
+            add: "f32[1]" = torch.ops.aten.add.Tensor(arg1_1, 1)
+            add_1: "f32[8]" = torch.ops.aten.add.Tensor(arg0_1, 1)
+            add_2: "f32[8]" = torch.ops.aten.add.Tensor(add, add_1)
+            copy_: "f32[8]" = torch.ops.aten.copy_.default(arg0_1, add_1);  arg0_1 = add_1 = copy_ = None
+            copy__1: "f32[1]" = torch.ops.aten.copy_.default(arg1_1, add);  arg1_1 = add = copy__1 = None
+            return (add_2,)
+
+    class auto_functionalized_subgraph_1(torch.nn.Module):
+        def forward(self, arg0_1: "f32[8]", arg1_1: "f32[1]"):
+            add: "f32[1]" = torch.ops.aten.add.Tensor(arg1_1, 1)
+            add_1: "f32[8]" = torch.ops.aten.add.Tensor(arg0_1, 1)
+            add_2: "f32[8]" = torch.ops.aten.add.Tensor(add, add_1)
+            copy_: "f32[8]" = torch.ops.aten.copy_.default(arg0_1, add_1);  arg0_1 = add_1 = copy_ = None
+            copy__1: "f32[1]" = torch.ops.aten.copy_.default(arg1_1, add);  arg1_1 = add = copy__1 = None
+            return (add_2,)
+""",  # noqa: B950
+            )
+
+    @requires_cuda
+    @unittest.skipIf(not SM70OrLater, "triton")
+    @parametrize("device", ["cuda", "cpu"])
+    @parametrize("dynamic", [True, False])
+    def test_cond_auto_functionalize_union_input_mutation(self, device, dynamic):
+        class M(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.register_buffer("buf", torch.ones(4, 3, requires_grad=False))
+
+            def forward(self, x, y):
+                def true_fn(x):
+                    x.add_(1)
+                    return x.sin() @ self.buf
+
+                def false_fn(x):
+                    self.buf.add_(1)
+                    return x.sin() @ self.buf
+
+                x = x.clone()
+                ret = torch.cond(x.sum() > 0, true_fn, false_fn, (x,))
+                return y + ret + x.sum() + self.buf.sum()
+
+        x, y = (
+            torch.randn(3, 4, requires_grad=False),
+            torch.randn(1, requires_grad=False),
+        )
+        fw_gm = self.check(M, (x, y), device, dynamic)
+        if not TEST_WITH_CROSSREF and not dynamic and device == "cuda":
+            self.assertExpectedInline(
+                normalize_gm(fw_gm.print_readable(print_output=False)),
+                """\
+class <lambda>(torch.nn.Module):
+    def forward(self, arg0_1: "f32[3, 4]", arg1_1: "f32[4, 3]", arg2_1: "f32[1]"):
+        clone: "f32[3, 4]" = torch.ops.aten.clone.default(arg0_1);  arg0_1 = None
+
+        sum_1: "f32[]" = torch.ops.aten.sum.default(clone)
+        gt: "b8[]" = torch.ops.aten.gt.Scalar(sum_1, 0);  sum_1 = None
+        auto_functionalized_subgraph_0 = self.auto_functionalized_subgraph_0
+        auto_functionalized_subgraph_1 = self.auto_functionalized_subgraph_1
+        _tree_spec_constant0 = self._tree_spec_constant0
+        auto_functionalized_v2 = torch.ops.higher_order.auto_functionalized_v2(torch.ops.higher_order.cond, pred = gt, true_fn = auto_functionalized_subgraph_0, false_fn = auto_functionalized_subgraph_1, _operand0_base_index = 0, _operand1_base_index = 1, _all_bases = [arg1_1, clone], _op_schema = _tree_spec_constant0);  gt = auto_functionalized_subgraph_0 = auto_functionalized_subgraph_1 = clone = _tree_spec_constant0 = None
+        getitem: "f32[3, 3]" = auto_functionalized_v2[0]
+        getitem_1: "f32[4, 3]" = auto_functionalized_v2[1]
+        getitem_2: "f32[3, 4]" = auto_functionalized_v2[2];  auto_functionalized_v2 = None
+
+        add: "f32[3, 3]" = torch.ops.aten.add.Tensor(arg2_1, getitem);  arg2_1 = getitem = None
+        sum_2: "f32[]" = torch.ops.aten.sum.default(getitem_2);  getitem_2 = None
+        add_1: "f32[3, 3]" = torch.ops.aten.add.Tensor(add, sum_2);  add = sum_2 = None
+        sum_3: "f32[]" = torch.ops.aten.sum.default(getitem_1)
+        add_2: "f32[3, 3]" = torch.ops.aten.add.Tensor(add_1, sum_3);  add_1 = sum_3 = None
+
+        copy_: "f32[4, 3]" = torch.ops.aten.copy_.default(arg1_1, getitem_1);  arg1_1 = getitem_1 = copy_ = None
+        return (add_2,)
+
+    class auto_functionalized_subgraph_0(torch.nn.Module):
+        def forward(self, arg0_1: "f32[4, 3]", arg1_1: "f32[3, 4]"):
+            add: "f32[3, 4]" = torch.ops.aten.add.Tensor(arg1_1, 1)
+            sin: "f32[3, 4]" = torch.ops.aten.sin.default(add)
+            mm: "f32[3, 3]" = torch.ops.aten.mm.default(sin, arg0_1);  sin = arg0_1 = None
+            copy_: "f32[3, 4]" = torch.ops.aten.copy_.default(arg1_1, add);  arg1_1 = add = copy_ = None
+            return (mm,)
+
+    class auto_functionalized_subgraph_1(torch.nn.Module):
+        def forward(self, arg0_1: "f32[4, 3]", arg1_1: "f32[3, 4]"):
+            add: "f32[4, 3]" = torch.ops.aten.add.Tensor(arg0_1, 1)
+            sin: "f32[3, 4]" = torch.ops.aten.sin.default(arg1_1);  arg1_1 = None
+            mm: "f32[3, 3]" = torch.ops.aten.mm.default(sin, add);  sin = None
+            copy_: "f32[4, 3]" = torch.ops.aten.copy_.default(arg0_1, add);  arg0_1 = add = copy_ = None
+            return (mm,)
+""",  # noqa: B950
+            )
+
+
 _hop_schema_test_schema_types = [
     "bool",
     "int",
@@ -9672,11 +10246,115 @@ def body_fn(x, y, z, c):
         )
 
 
+class DynamicCondModel(torch.nn.Module):
+    def __init__(self, input_size=16, hidden_size=64, output_size=10):
+        super().__init__()
+        self.fc1_0 = torch.nn.Linear(input_size, hidden_size)
+        self.fc1_1 = torch.nn.Linear(input_size, 32)
+        self.fc1_2 = torch.nn.Linear(32, hidden_size)
+        self.relu = torch.nn.ReLU()
+        self.fc2 = torch.nn.Linear(hidden_size, output_size)
+
+    def forward(self, x):
+        def true_fn(x):
+            return self.fc1_0(x)
+
+        def false_fn(x):
+            x = self.fc1_1(x)
+            return self.fc1_2(x)
+
+        pred = x.sum() > 0
+        x = cond(pred, true_fn, false_fn, [x])
+
+        x = self.relu(x)
+        x = self.fc2(x)
+
+        return x
+
+
+@unittest.skipIf(
+    not TEST_CUDA_GRAPH_CONDITIONAL_NODES,
+    "CUDA 12.4 or greater is required for CUDA Graphs with conditional nodes",
+)
+class TestControlFlowNN(TestCase):
+    def test_cond_in_NN(self):
+        model = DynamicCondModel().cuda()
+
+        def autograd_test(x):
+            model.zero_grad(set_to_none=True)
+            output = model(x)
+            loss = output.sum()
+            loss.backward()
+            grads = [p.grad for p in model.parameters()]
+            return (output, loss, grads)
+
+        x = torch.randn(16, device="cuda")
+
+        _check_compile_many_backends_with_cudagraph(self, autograd_test, [x])
+        _check_compile_cudagraph_backend(self, autograd_test, [x])
+
+
+@unittest.skipIf(
+    not TEST_CUDA_GRAPH_CONDITIONAL_NODES,
+    "CUDA 12.4 or greater is required for CUDA Graphs with conditional nodes",
+)
+class TestControlFlowAndRNG(TestCase):
+    @parametrize("rng_func", ["custom_generator", "default_generator"])
+    def test_rng_with_conditional_nodes_errors(self, rng_func):
+        pred = torch.tensor(True, device="cuda")
+        x = torch.ones(10, dtype=torch.float32, device="cuda")
+
+        if rng_func == "custom_generator":
+            self.skipTest(
+                "randn() currently does not work with a generator argument in dynamo."
+            )
+            generator = torch.Generator("cuda")
+
+            def custom_generator(x):
+                return x + torch.randn(
+                    *x.shape, generator=generator, dtype=x.dtype, device=x.device
+                )
+
+            rng_func = custom_generator
+        elif rng_func == "default_generator":
+
+            def default_generator(x):
+                return x + torch.randn(*x.shape, dtype=x.dtype, device=x.device)
+
+            rng_func = default_generator
+
+        def func(pred, x):
+            return torch.cond(pred, rng_func, lambda x: 2 * x, [x])
+
+        compiled_func = torch.compile(func, backend="cudagraphs")
+        with self.assertRaisesRegex(
+            RuntimeError,
+            "RNG op during graph capture but generator is not registered",
+        ):
+            compiled_func(pred, x)
+
+    def test_rng_outside_conditional_nodes_does_not_error(self):
+        pred = torch.tensor(True, device="cuda")
+        x = torch.ones(10, dtype=torch.float32, device="cuda")
+
+        def func(pred, x):
+            y = torch.cond(pred, lambda t: 2 * t, lambda t: 3 * t, [x])
+            return y + torch.randn(*y.shape, dtype=y.dtype, device=y.device)
+
+        compiled_func = torch.compile(func, backend="cudagraphs")
+        for _ in range(3):
+            out = compiled_func(pred, x)
+            self.assertEqual(out.shape, x.shape)
+
+
 instantiate_parametrized_tests(TestHopSchema)
 instantiate_parametrized_tests(TestControlFlowTraced)
+instantiate_parametrized_tests(TestAutoFunctionalizeControlFlow)
 
 instantiate_parametrized_tests(TestControlFlow)
 instantiate_parametrized_tests(AssociativeScanTests)
 
+instantiate_parametrized_tests(TestControlFlowAndRNG)
+
 if __name__ == "__main__":
     run_tests()
diff --git a/test/functorch/test_control_flow_cuda_initialization.py b/test/functorch/test_control_flow_cuda_initialization.py
new file mode 100644
index 0000000000000..ae610e054152b
--- /dev/null
+++ b/test/functorch/test_control_flow_cuda_initialization.py
@@ -0,0 +1,91 @@
+# Owner(s): ["module: functorch"]
+import unittest
+
+import torch
+from torch.testing._internal.common_utils import (
+    run_tests,
+    TEST_CUDA_GRAPH_CONDITIONAL_NODES,
+    TestCase,
+)
+
+
+@unittest.skipIf(
+    not TEST_CUDA_GRAPH_CONDITIONAL_NODES,
+    "CUDA 12.4 or greater is required for CUDA Graphs with conditional nodes",
+)
+class TestControlFlowInCUDAGraphInitialization(TestCase):
+    # Duplicated from test_cuda_primary_ctx.py
+    CTX_ALREADY_CREATED_ERR_MSG = (
+        "Tests defined in TestControlFlowInCUDAGraphInitialization must be run in a process "
+        "where CUDA contexts are never created. Use either run_test.py or add "
+        "--subprocess to run each test in a different subprocess."
+    )
+
+    def setUp(self):
+        # Ensure context has not been created beforehand
+        self.assertFalse(
+            torch._C._cuda_hasPrimaryContext(0),
+            TestControlFlowInCUDAGraphInitialization.CTX_ALREADY_CREATED_ERR_MSG,
+        )
+
+    def _check_compile_cudagraphs(self, f, pred, *other_args):
+        f_compiled = torch.compile(f, backend="cudagraphs")
+
+        outputs = []
+
+        for p in [pred, torch.logical_not(pred)]:
+            for i in range(3):
+                torch.compiler.cudagraph_mark_step_begin()
+                outputs.append(f_compiled(pred, *other_args).clone())
+
+            # We compute the eager output only after running cudagraphs
+            # backend compiled function, in order to make sure that
+            # cudagraph trees warms up the conditional part of the code
+            # properly.
+            eager_output = f(pred, *other_args)
+
+            for output in outputs:
+                self.assertEqual(output, eager_output)
+
+    def test_cond_cudnn(self):
+        # Tests that cublasCreate() does not break stream capture
+        def f(pred, x, filters):
+            return torch.cond(
+                pred,
+                lambda y: torch.sum(y),
+                lambda y: torch.sum(torch.nn.functional.conv1d(y, filters)),
+                [x],
+            )
+
+        self.assertFalse(torch._C._cuda_hasPrimaryContext(0))
+
+        pred = torch.tensor(True, device="cuda")
+        x = torch.randn(33, 16, 30, device="cuda")
+        filters = torch.randn(20, 16, 5, device="cuda")
+
+        self._check_compile_cudagraphs(f, pred, x, filters)
+
+        self.assertTrue(torch._C._cuda_hasPrimaryContext(0))
+
+    def test_cond_stft(self):
+        # Tests that cufft plan creation does not break stream capture
+        def f(pred, x):
+            return torch.cond(
+                pred,
+                lambda y: torch.sum(y),
+                lambda y: torch.sum(torch.stft(y, 512, return_complex=False)),
+                [x],
+            )
+
+        self.assertFalse(torch._C._cuda_hasPrimaryContext(0))
+
+        pred = torch.tensor(True, device="cuda")
+        x = torch.ones(1024 * 1024, device="cuda")
+
+        self._check_compile_cudagraphs(f, pred, x)
+
+        self.assertTrue(torch._C._cuda_hasPrimaryContext(0))
+
+
+if __name__ == "__main__":
+    run_tests()
diff --git a/test/functorch/test_dims.py b/test/functorch/test_dims.py
index a0cd59c02665a..8cb755878dea4 100644
--- a/test/functorch/test_dims.py
+++ b/test/functorch/test_dims.py
@@ -415,9 +415,11 @@ def test_hello(self):
         )
 
         r = [id(x) for x in torch.rand_like(A[i, k]).dims]
-        assert id(i) in r and id(k) in r
+        if not (id(i) in r and id(k) in r):
+            raise AssertionError("Expected i and k to be in dims")
         r = [id(x) for x in torch.nn.functional.dropout(A[i, k]).dims]
-        assert id(i) in r and id(k) in r
+        if not (id(i) in r and id(k) in r):
+            raise AssertionError("Expected i and k to be in dims")
 
     def test_simple(self):
         i, j, k = dims()
@@ -503,11 +505,14 @@ def test_network(self):
 
     def test_dim_args(self):
         a = dimlists()
-        assert isinstance(a, DimList)
+        if not isinstance(a, DimList):
+            raise AssertionError(f"Expected DimList, got {type(a)}")
         a = dims()
         b = dimlists()
-        assert isinstance(a, Dim)
-        assert isinstance(b, DimList)
+        if not isinstance(a, Dim):
+            raise AssertionError(f"Expected Dim, got {type(a)}")
+        if not isinstance(b, DimList):
+            raise AssertionError(f"Expected DimList, got {type(b)}")
         self.assertEqual(str(a), "a")
         a, b = dims(sizes=[3, 4])
         self.assertEqual(a.size, 3)
@@ -604,7 +609,8 @@ def test_eq(self):
     def test_dims_with_size(self):
         x = dims(3)
         self.assertEqual(len(x), 3)
-        assert isinstance(x[0], Dim)
+        if not isinstance(x[0], Dim):
+            raise AssertionError(f"Expected Dim, got {type(x[0])}")
 
         class Foo:
             pass
diff --git a/test/functorch/test_eager_transforms.py b/test/functorch/test_eager_transforms.py
index 37bb013e5df82..6330640bf951a 100644
--- a/test/functorch/test_eager_transforms.py
+++ b/test/functorch/test_eager_transforms.py
@@ -76,7 +76,6 @@
     TEST_CUDA_MEM_LEAK_CHECK,
     TEST_WITH_TORCHDYNAMO,
     TestCase,
-    xfailIfTorchDynamo,
 )
 from torch.utils._pytree import tree_flatten, tree_map, tree_unflatten
 
@@ -216,7 +215,8 @@ def _get_weights_and_functional_call(net, mechanism):
     if mechanism == "make_functional":
         return make_functional(net)
     else:
-        assert mechanism == "functional_call"
+        if mechanism != "functional_call":
+            raise AssertionError(f"Expected 'functional_call', got {mechanism!r}")
         # this makes it so the function from make_functional and this call have the same signature
 
         def net_func(weights, data):
@@ -229,7 +229,8 @@ def _get_weights_and_functional_call_with_buffers(net, mechanism):
     if mechanism == "make_functional":
         return make_functional_with_buffers(net)
     else:
-        assert mechanism == "functional_call"
+        if mechanism != "functional_call":
+            raise AssertionError(f"Expected 'functional_call', got {mechanism!r}")
 
         # this makes it so the function from make_functional and this call have the same signature
         def net_func(weights, buffers, data):
@@ -452,9 +453,11 @@ def test_conj_bit(self):
         x = torch.tensor(1 + 1j)
 
         def foo(x):
-            assert not x.is_conj()
+            if x.is_conj():
+                raise AssertionError("Expected x to not be conj")
             y = x.conj()
-            assert y.is_conj()
+            if not y.is_conj():
+                raise AssertionError("Expected y to be conj")
             return y.abs()
 
         res = grad(foo)(x)
@@ -765,14 +768,18 @@ def test_vjp_outputs_can_any_pytree(self, device):
         # Check list output
         output, vjp_fn = vjp(lambda x: [x, x.sum()], x)
         (vjp_out,) = vjp_fn([t, t.sum()])
-        assert isinstance(output, list) and len(output) == 2
-        assert isinstance(vjp_out, torch.Tensor)
+        if not isinstance(output, list) or len(output) != 2:
+            raise AssertionError(f"Expected list of length 2, got {type(output)}")
+        if not isinstance(vjp_out, torch.Tensor):
+            raise AssertionError(f"Expected Tensor, got {type(vjp_out)}")
 
         # Check dict output
         output, vjp_fn = vjp(lambda x: {"x": x, "xsum": x.sum()}, x)
         (vjp_out,) = vjp_fn({"x": t, "xsum": t.sum()})
-        assert isinstance(output, dict) and len(output) == 2 and "xsum" in output
-        assert isinstance(vjp_out, torch.Tensor)
+        if not isinstance(output, dict) or len(output) != 2 or "xsum" not in output:
+            raise AssertionError(f"Expected dict with 'xsum', got {output}")
+        if not isinstance(vjp_out, torch.Tensor):
+            raise AssertionError(f"Expected Tensor, got {type(vjp_out)}")
 
         def composite_output(x):
             out = x.sum()
@@ -786,9 +793,12 @@ def composite_output(x):
                 (t.sum(), {"a": t, "out": [t, t.sum()]}),
             ]
         )
-        assert isinstance(output, list)
-        assert isinstance(output[0], tuple) and isinstance(output[0][1], dict)
-        assert isinstance(vjp_out, torch.Tensor)
+        if not isinstance(output, list):
+            raise AssertionError(f"Expected list, got {type(output)}")
+        if not isinstance(output[0], tuple) or not isinstance(output[0][1], dict):
+            raise AssertionError(f"Expected tuple with dict, got {output[0]}")
+        if not isinstance(vjp_out, torch.Tensor):
+            raise AssertionError(f"Expected Tensor, got {type(vjp_out)}")
 
     def test_vjp_pytree_error(self, device):
         def f(x):
@@ -1141,6 +1151,104 @@ def h(x):
         (z,) = torch.autograd.grad(y, x)
         self.assertEqual(z, 2)
 
+    @skipIfTorchDynamo("internal API test")
+    def test_pop_dynamic_layer_stack_to_depth_single(self, device):
+        ft = torch._C._functorch
+        ft._grad_increment_nesting()
+        self.assertEqual(ft.get_dynamic_layer_stack_depth(), 1)
+        ft.pop_dynamic_layer_stack_and_undo_to_depth(0)
+        self.assertEqual(ft.get_dynamic_layer_stack_depth(), 0)
+
+    @skipIfTorchDynamo("internal API test")
+    def test_pop_dynamic_layer_stack_to_depth_mixed(self, device):
+        ft = torch._C._functorch
+        ft._vmap_increment_nesting(3, "error")
+        ft._grad_increment_nesting()
+        ft._jvp_increment_nesting()
+        self.assertEqual(ft.get_dynamic_layer_stack_depth(), 3)
+        # Pop only jvp — must remove exactly one layer
+        ft.pop_dynamic_layer_stack_and_undo_to_depth(2)
+        self.assertEqual(ft.get_dynamic_layer_stack_depth(), 2)
+        # Pop remaining
+        ft.pop_dynamic_layer_stack_and_undo_to_depth(0)
+        self.assertEqual(ft.get_dynamic_layer_stack_depth(), 0)
+
+    def test_inference_mode_outside_grad(self, device):
+        x = torch.randn(3, device=device)
+        with torch.inference_mode():
+            y = grad(lambda x: (x**2).sum())(x)
+        self.assertEqual(y, 2 * x)
+
+    def test_inference_mode_nograd_outside_grad(self, device):
+        x = torch.randn(3, device=device)
+        with torch.inference_mode():
+            with torch.no_grad():
+                y = grad(lambda x: (x**2).sum())(x)
+        self.assertEqual(y, 2 * x)
+
+    def test_inference_mode_outside_vjp(self, device):
+        x = torch.randn(3, device=device)
+        with torch.inference_mode():
+            out, vjp_fn = vjp(lambda x: (x**2).sum(), x)
+            (y,) = vjp_fn(torch.tensor(1.0, device=device))
+        self.assertEqual(y, 2 * x)
+
+    def test_inference_mode_outside_jvp(self, device):
+        x = torch.randn(3, device=device)
+        t = torch.ones(3, device=device)
+        with torch.inference_mode():
+            _, y = jvp(lambda x: (x**2).sum(), (x,), (t,))
+        self.assertEqual(y, (2 * x * t).sum())
+
+    def test_inference_mode_outside_jacrev(self, device):
+        x = torch.randn(3, device=device)
+        with torch.inference_mode():
+            y = jacrev(lambda x: x**2)(x)
+        self.assertEqual(y, torch.diag(2 * x))
+
+    def test_inference_mode_outside_vmap_grad(self, device):
+        xs = torch.randn(5, 3, device=device)
+        with torch.inference_mode():
+            ys = vmap(grad(lambda x: (x**2).sum()))(xs)
+        self.assertEqual(ys, 2 * xs)
+
+    def test_inference_mode_outside_grad_vmap(self, device):
+        x = torch.randn(3, device=device)
+        with torch.inference_mode():
+            y = grad(lambda x: vmap(lambda x: (x**2).sum())(x).sum())(x)
+        self.assertEqual(y, 2 * x)
+
+    def test_inference_mode_nested_grad(self, device):
+        x = torch.randn([], device=device)
+        with torch.inference_mode():
+            y = grad(grad(lambda x: x**3))(x)
+        self.assertEqual(y, 6 * x)
+
+    def test_inference_mode_jacrev_grad(self, device):
+        x = torch.randn(3, device=device)
+        with torch.inference_mode():
+            H = jacrev(grad(lambda x: (x**3).sum()))(x)
+        self.assertEqual(H, torch.diag(6 * x))
+
+    def test_inference_mode_inside_grad(self, device):
+        def f(x):
+            with torch.inference_mode():
+                c = x**2
+            return x - c
+
+        x = torch.randn(3, device=device)
+        with torch.inference_mode():
+            y = grad(lambda x: f(x).sum())(x)
+        self.assertEqual(y, torch.ones_like(x))
+
+    def test_inference_mode_restored(self, device):
+        self.assertTrue(not torch.is_inference_mode_enabled())
+        with torch.inference_mode():
+            self.assertTrue(torch.is_inference_mode_enabled())
+            grad(lambda x: (x**2).sum())(torch.randn(3, device=device))
+            self.assertTrue(torch.is_inference_mode_enabled())
+        self.assertTrue(not torch.is_inference_mode_enabled())
+
 
 @markDynamoStrictTest
 class TestAutogradFunction(TestCase):
@@ -1515,7 +1623,8 @@ def setup_context(ctx, inputs, output):
 
             @staticmethod
             def vmap(info, in_dims, input):
-                assert in_dims == (0,)
+                if in_dims != (0,):
+                    raise AssertionError(f"Expected in_dims == (0,), got {in_dims}")
                 return torch.zeros(input.shape[1:], device=input.device), None
 
         B = 2
@@ -1535,7 +1644,8 @@ def setup_context(ctx, inputs, output):
 
             @staticmethod
             def vmap(info, in_dims, input):
-                assert in_dims == (0,)
+                if in_dims != (0,):
+                    raise AssertionError(f"Expected in_dims == (0,), got {in_dims}")
                 r = torch.zeros(input.shape[1:], device=input.device)
                 return (r, r), None
 
@@ -1644,7 +1754,8 @@ def setup_context(ctx, inputs, output):
 
                 @staticmethod
                 def vmap(info, in_dims, x, *, y):
-                    assert in_dims == (0,)
+                    if in_dims != (0,):
+                        raise AssertionError(f"Expected in_dims == (0,), got {in_dims}")
                     return x + y, 0
 
             x = torch.randn(3)
@@ -1719,7 +1830,10 @@ def _compare_expected_and_result(self, expected, result, mechanism):
             for r, e in zip(result, expected):
                 self.assertEqual(r, e, atol=0, rtol=1.5e-3)
         else:
-            assert mechanism == "functional_call"
+            if mechanism != "functional_call":
+                raise AssertionError(
+                    f"Expected mechanism 'functional_call', got '{mechanism}'"
+                )
             expected = {
                 k: tuple(d[k] for d in expected) for k, v in expected[0].items()
             }
@@ -1804,7 +1918,8 @@ def test_simple(self, device, jacapi):
         x = torch.randn(3, device=device)
         y = jacapi(torch.sin)(x)
         expected = torch.diagflat(x.cos())
-        assert torch.allclose(y, expected)
+        if not torch.allclose(y, expected):
+            raise AssertionError("torch.allclose failed: y and expected differ")
 
     @jacrev_and_jacfwd
     def test_simple_not_flat(self, device, jacapi):
@@ -1812,7 +1927,8 @@ def test_simple_not_flat(self, device, jacapi):
         y = jacapi(torch.sin)(x)
         expected = torch.diagflat(x.view(-1).cos())
         expected = expected.view(2, 3, 2, 3)
-        assert torch.allclose(y, expected)
+        if not torch.allclose(y, expected):
+            raise AssertionError("torch.allclose failed: y and expected differ")
 
     @jacrev_and_jacfwd
     def test_take(self, device, jacapi):
@@ -1847,7 +1963,8 @@ def test_vmap_on_jac_simple(self, device, jacapi):
         x = torch.randn(2, 3, device=device)
         y = vmap(jacapi(torch.sin))(x)
         expected = torch.stack([torch.diagflat(x[i].cos()) for i in range(2)])
-        assert torch.allclose(y, expected)
+        if not torch.allclose(y, expected):
+            raise AssertionError("torch.allclose failed: y and expected differ")
 
     @jacrev_and_jacfwd
     def test_nested_jac_simple(self, device, jacapi):
@@ -1857,7 +1974,8 @@ def foo(x):
         x = torch.randn(3, device=device)
         y = jacapi(jacapi(foo))(x)
         expected = torch.diagflat(-x.sin())
-        assert torch.allclose(y, expected)
+        if not torch.allclose(y, expected):
+            raise AssertionError("torch.allclose failed: y and expected differ")
 
     @jacrev_and_jacfwd
     def test_multiple_args(self, device, jacapi):
@@ -1865,7 +1983,8 @@ def test_multiple_args(self, device, jacapi):
         y = torch.randn(3, device=device)
         z = jacapi(torch.multiply, argnums=1)(x, y)
         expected = torch.diagflat(x)
-        assert torch.allclose(z, expected)
+        if not torch.allclose(z, expected):
+            raise AssertionError("torch.allclose failed: z and expected differ")
 
     @jacrev_and_jacfwd
     def test_multiple_outputs_multiple_argnums(self, device, jacapi):
@@ -2028,11 +2147,17 @@ def test_outputs_can_any_pytree(self, device, jacapi):
 
         # Check list output
         out = jacapi(lambda x: [x, x.sum()])(x)
-        assert isinstance(out, list) and len(out) == 2
+        if not (isinstance(out, list) and len(out) == 2):
+            raise AssertionError(
+                f"Expected list of length 2, got {type(out).__name__} of length {len(out)}"
+            )
 
         # Check dict output
         out = jacapi(lambda x: {"x": x, "xsum": x.sum()})(x)
-        assert isinstance(out, dict) and len(out) == 2 and "xsum" in out
+        if not (isinstance(out, dict) and len(out) == 2 and "xsum" in out):
+            raise AssertionError(
+                f"Expected dict of length 2 with 'xsum' key, got {type(out).__name__}"
+            )
 
         def composite_output(x):
             out = x.sum()
@@ -2041,8 +2166,12 @@ def composite_output(x):
             ]
 
         out = jacapi(composite_output)(x)
-        assert isinstance(out, list)
-        assert isinstance(out[0], tuple) and isinstance(out[0][1], dict)
+        if not isinstance(out, list):
+            raise AssertionError(f"Expected list, got {type(out).__name__}")
+        if not (isinstance(out[0], tuple) and isinstance(out[0][1], dict)):
+            raise AssertionError(
+                f"Expected (tuple, dict) structure, got ({type(out[0]).__name__}, {type(out[0][1]).__name__})"
+            )
 
     @jacrev_and_jacfwd
     def test_multiple_inputs_outputs_pytree(self, device, jacapi):
@@ -2141,9 +2270,12 @@ def test_argnums_tuple(self, device, jacapi):
         z = jacapi(torch.multiply, argnums=(0, 1))(x, y)
         expected0 = torch.diagflat(y)
         expected1 = torch.diagflat(x)
-        assert len(z) == 2
-        assert torch.allclose(z[0], expected0)
-        assert torch.allclose(z[1], expected1)
+        if len(z) != 2:
+            raise AssertionError(f"Expected len(z) == 2, got {len(z)}")
+        if not torch.allclose(z[0], expected0):
+            raise AssertionError("torch.allclose failed: z[0] and expected0 differ")
+        if not torch.allclose(z[1], expected1):
+            raise AssertionError("torch.allclose failed: z[1] and expected1 differ")
 
     @jacrev_and_jacfwd
     def test_argnums_effect_on_return(self, device, jacapi):
@@ -2151,16 +2283,21 @@ def test_argnums_effect_on_return(self, device, jacapi):
         y = torch.randn(3, device=device)
         z = jacapi(torch.multiply, argnums=(0,))(x, y)
         expected0 = torch.diagflat(y)
-        assert isinstance(z, tuple)
-        assert len(z) == 1
-        assert torch.allclose(z[0], expected0)
+        if not isinstance(z, tuple):
+            raise AssertionError(f"Expected tuple, got {type(z).__name__}")
+        if len(z) != 1:
+            raise AssertionError(f"Expected len(z) == 1, got {len(z)}")
+        if not torch.allclose(z[0], expected0):
+            raise AssertionError("torch.allclose failed: z[0] and expected0 differ")
 
         x = torch.randn(3, device=device)
         y = torch.randn(3, device=device)
         z = jacapi(torch.multiply, argnums=0)(x, y)
         expected0 = torch.diagflat(y)
-        assert isinstance(z, torch.Tensor)
-        assert torch.allclose(z, expected0)
+        if not isinstance(z, torch.Tensor):
+            raise AssertionError(f"Expected torch.Tensor, got {type(z).__name__}")
+        if not torch.allclose(z, expected0):
+            raise AssertionError("torch.allclose failed: z and expected0 differ")
 
     @jacrev_and_jacfwd
     def test_argnums_defaults_to_zero(self, device, jacapi):
@@ -2377,7 +2514,6 @@ def f(x):
         self.assertEqual(actual, expected)
 
     # https://github.com/pytorch/pytorch/issues/127036
-    @xfailIfTorchDynamo
     @parametrize("_preallocate_and_copy", (True, False))
     def test_chunk_jacrev_chunksize_one(self, device, _preallocate_and_copy):
         # With chunk_size=1, we shouldn't `vmap` and hence not be limited
@@ -2711,12 +2847,18 @@ def test_outputs_can_any_pytree(self, device):
         # Check list output
         out = jvp(lambda x: [x, x.sum()], (x,), (t,))
         for i in range(2):
-            assert isinstance(out[i], list) and len(out[i]) == 2
+            if not (isinstance(out[i], list) and len(out[i]) == 2):
+                raise AssertionError(
+                    f"Expected list of length 2, got {type(out[i]).__name__} of length {len(out[i])}"
+                )
 
         # Check dict output
         out = jvp(lambda x: {"x": x, "xsum": x.sum()}, (x,), (t,))
         for i in range(2):
-            assert isinstance(out[i], dict) and len(out[i]) == 2 and "xsum" in out[i]
+            if not (isinstance(out[i], dict) and len(out[i]) == 2 and "xsum" in out[i]):
+                raise AssertionError(
+                    f"Expected dict of length 2 with 'xsum' key, got {type(out[i]).__name__}"
+                )
 
         def composite_output(x):
             out = x.sum()
@@ -2726,8 +2868,12 @@ def composite_output(x):
 
         out = jvp(composite_output, (x,), (t,))
         for i in range(2):
-            assert isinstance(out[i], list)
-            assert isinstance(out[i][0], tuple) and isinstance(out[i][0][1], dict)
+            if not isinstance(out[i], list):
+                raise AssertionError(f"Expected list, got {type(out[i]).__name__}")
+            if not (isinstance(out[i][0], tuple) and isinstance(out[i][0][1], dict)):
+                raise AssertionError(
+                    f"Expected (tuple, dict) structure, got ({type(out[i][0]).__name__}, {type(out[i][0][1]).__name__})"
+                )
 
     def test_aux_tensor(self, device):
         x = torch.randn(3, device=device)
@@ -3226,16 +3372,34 @@ def backward(ctx, gy):
                 ctx_y = CtxWithSavedTensors(ctx, (y,))
                 # Can't use self.assertEqual because that relies on TLS
                 # that is not available in multithread autograd
-                assert len(ctx_y.saved_tensors) == 1
-                assert torch.allclose(ctx_y.saved_tensors[0], y)
+                if len(ctx_y.saved_tensors) != 1:
+                    raise AssertionError(
+                        f"Expected len(ctx_y.saved_tensors) == 1, got {len(ctx_y.saved_tensors)}"
+                    )
+                if not torch.allclose(ctx_y.saved_tensors[0], y):
+                    raise AssertionError(
+                        "torch.allclose failed: ctx_y.saved_tensors[0] and y differ"
+                    )
 
                 wrapped = CtxWithSavedTensors(ctx_y, (z,))
 
-                assert len(wrapped.saved_tensors) == 1
-                assert torch.allclose(wrapped.saved_tensors[0], z)
+                if len(wrapped.saved_tensors) != 1:
+                    raise AssertionError(
+                        f"Expected len(wrapped.saved_tensors) == 1, got {len(wrapped.saved_tensors)}"
+                    )
+                if not torch.allclose(wrapped.saved_tensors[0], z):
+                    raise AssertionError(
+                        "torch.allclose failed: wrapped.saved_tensors[0] and z differ"
+                    )
 
-                assert len(ctx_y.saved_tensors) == 1
-                assert torch.allclose(ctx_y.saved_tensors[0], y)
+                if len(ctx_y.saved_tensors) != 1:
+                    raise AssertionError(
+                        f"Expected len(ctx_y.saved_tensors) == 1, got {len(ctx_y.saved_tensors)}"
+                    )
+                if not torch.allclose(ctx_y.saved_tensors[0], y):
+                    raise AssertionError(
+                        "torch.allclose failed: ctx_y.saved_tensors[0] and y differ"
+                    )
 
                 return gy * wrapped.saved_tensors[0]
 
@@ -3259,7 +3423,10 @@ def backward(ctx, gy):
                 wrapped = torch._functorch.autograd_function.CtxWithSavedTensors(
                     ctx, override
                 )
-                assert wrapped.saved_tensors == override
+                if wrapped.saved_tensors != override:
+                    raise AssertionError(
+                        f"Expected wrapped.saved_tensors == {override}, got {wrapped.saved_tensors}"
+                    )
                 return gy
 
         out = A.apply(x)
@@ -3283,11 +3450,21 @@ def backward(ctx, gz):
                     ctx, override
                 )
 
-                assert wrapped.needs_input_grad[0] == ctx.needs_input_grad[0]
-                assert wrapped.needs_input_grad[1] == ctx.needs_input_grad[1]
+                if wrapped.needs_input_grad[0] != ctx.needs_input_grad[0]:
+                    raise AssertionError(
+                        f"needs_input_grad[0] mismatch: wrapped={wrapped.needs_input_grad[0]}, ctx={ctx.needs_input_grad[0]}"
+                    )
+                if wrapped.needs_input_grad[1] != ctx.needs_input_grad[1]:
+                    raise AssertionError(
+                        f"needs_input_grad[1] mismatch: wrapped={wrapped.needs_input_grad[1]}, ctx={ctx.needs_input_grad[1]}"
+                    )
                 wrapped.foo = "bar"
-                assert wrapped.foo == "bar"
-                assert ctx.foo == "bar"
+                if wrapped.foo != "bar":
+                    raise AssertionError(
+                        f"Expected wrapped.foo == 'bar', got {wrapped.foo!r}"
+                    )
+                if ctx.foo != "bar":
+                    raise AssertionError(f"Expected ctx.foo == 'bar', got {ctx.foo!r}")
                 return gz, gz
 
         out = A.apply(x, y)
@@ -4046,7 +4223,10 @@ def get_module_info(mod):
             if mechanism == "make_functional":
                 return make_functional(mod)
             else:
-                assert mechanism == "functional_call"
+                if mechanism != "functional_call":
+                    raise AssertionError(
+                        f"Expected mechanism 'functional_call', got '{mechanism}'"
+                    )
                 return mod, dict(mod.named_parameters())
 
         mod = Net()
@@ -4064,7 +4244,10 @@ def get_module_info(mod):
         if mechanism == "make_functional":
             func_mod(params, x)
         else:
-            assert mechanism == "functional_call"
+            if mechanism != "functional_call":
+                raise AssertionError(
+                    f"Expected mechanism 'functional_call', got '{mechanism}'"
+                )
             functional_call(func_mod, params, x)
 
         mod = func_mod.stateless_model if mechanism == "make_functional" else func_mod
@@ -4084,7 +4267,10 @@ def _update_params(self, params, grads, alpha, mechanism):
         if mechanism == "make_functional":
             return [(params[i] - alpha * grads[i]) for i in range(len(params))]
         else:
-            assert mechanism == "functional_call"
+            if mechanism != "functional_call":
+                raise AssertionError(
+                    f"Expected mechanism 'functional_call', got '{mechanism}'"
+                )
             return {k: params[k] - alpha * grads[k] for k in params}
 
     @parametrize("mechanism", ["make_functional", "functional_call"])
@@ -4486,7 +4672,8 @@ def slice_weights(batched_weights, index):
 
         weights0, spec0 = tree_flatten(weights0)
         weights1, spec1 = tree_flatten(weights1)
-        assert spec0 == spec1
+        if spec0 != spec1:
+            raise AssertionError(f"Expected spec0 == spec1, got {spec0} != {spec1}")
         expected_weights = tuple(
             torch.stack([w0, w1]) for w0, w1 in zip(weights0, weights1)
         )
diff --git a/test/functorch/test_leaf_function.py b/test/functorch/test_leaf_function.py
new file mode 100644
index 0000000000000..0f0e5e11d6842
--- /dev/null
+++ b/test/functorch/test_leaf_function.py
@@ -0,0 +1,2486 @@
+# Owner(s): ["oncall: pt2"]
+
+"""Tests for @leaf_function with make_fx, aot_function, and torch.compile."""
+
+import copy
+import re
+from functools import partial
+from unittest.mock import patch
+
+import torch
+import torch._dynamo.config as config
+import torch._dynamo.testing
+from functorch.compile import aot_function, nop
+from torch._dynamo.decorators import leaf_function
+from torch._dynamo.testing import normalize_gm
+from torch._higher_order_ops.invoke_leaf_function import invoke_leaf_function
+from torch.fx.experimental.proxy_tensor import make_fx
+from torch.testing._internal.common_utils import (
+    instantiate_parametrized_tests,
+    parametrize,
+    run_tests,
+    skipIfCrossRef,
+    skipIfTorchDynamo,
+    TestCase,
+)
+from torch.testing._internal.dynamo_pytree_test_utils import PytreeRegisteringTestCase
+
+
+def extract_graph(fx_g, _, graph_cell):
+    graph_cell[0] = fx_g
+    return fx_g
+
+
+@skipIfTorchDynamo("leaf_function tests manage their own compilation")
+class TestLeafFunctionMakeFx(TestCase):
+    def _has_invoke_leaf_function_node(self, gm):
+        for node in gm.graph.nodes:
+            if node.op == "call_function" and node.target is invoke_leaf_function:
+                return True
+        return False
+
+    def test_make_fx_simple(self):
+        @leaf_function
+        def my_fn(x, y):
+            if x.sum() > 0:
+                return (x + y,)
+            return (x - y,)
+
+        @my_fn.register_fake
+        def my_fn_fake(x, y):
+            return (x + y,)
+
+        def f(x, y):
+            return my_fn(x, y)
+
+        x = torch.randn(3, 3)
+        y = torch.randn(3, 3)
+        gm = make_fx(f)(x, y)
+
+        self.assertTrue(self._has_invoke_leaf_function_node(gm))
+        self.assertExpectedInline(
+            normalize_gm(gm.print_readable(print_output=False)),
+            """\
+class f(torch.nn.Module):
+    def forward(self, x_1: "f32[3, 3]", y_1: "f32[3, 3]"):
+        _opaque_obj0 = self._opaque_obj0
+        _opaque_obj1 = self._opaque_obj1
+        _tree_spec_constant0 = self._tree_spec_constant0
+        invoke_leaf_function = torch.ops.higher_order.invoke_leaf_function(_opaque_obj0, _opaque_obj1, _tree_spec_constant0, '', x_1, y_1, requires_grad_indices = '');  _opaque_obj0 = _opaque_obj1 = _tree_spec_constant0 = x_1 = y_1 = None
+        getitem: "f32[3, 3]" = invoke_leaf_function[0];  invoke_leaf_function = None
+        return (getitem,)
+""",  # noqa: B950
+        )
+
+        x2 = torch.randn(3, 3)
+        y2 = torch.randn(3, 3)
+        self.assertEqual(gm(x2, y2), f(x2, y2))
+
+    def test_make_fx_with_closure(self):
+        config = {"scale": 2.0}
+
+        @leaf_function
+        def scaled_fn(x):
+            return (x * config["scale"],)
+
+        @scaled_fn.register_fake
+        def scaled_fn_fake(x):
+            return (torch.empty_like(x),)
+
+        def f(x):
+            return scaled_fn(x)
+
+        x = torch.randn(3, 3)
+        gm = make_fx(f)(x)
+
+        self.assertTrue(self._has_invoke_leaf_function_node(gm))
+        self.assertExpectedInline(
+            normalize_gm(gm.print_readable(print_output=False)),
+            """\
+class f(torch.nn.Module):
+    def forward(self, x_1: "f32[3, 3]"):
+        _opaque_obj0 = self._opaque_obj0
+        _opaque_obj1 = self._opaque_obj1
+        _tree_spec_constant0 = self._tree_spec_constant0
+        invoke_leaf_function = torch.ops.higher_order.invoke_leaf_function(_opaque_obj0, _opaque_obj1, _tree_spec_constant0, '', x_1, requires_grad_indices = '');  _opaque_obj0 = _opaque_obj1 = _tree_spec_constant0 = x_1 = None
+        getitem: "f32[3, 3]" = invoke_leaf_function[0];  invoke_leaf_function = None
+        return (getitem,)
+""",  # noqa: B950
+        )
+
+        # Closure change reflected at runtime
+        config["scale"] = 3.0
+        x2 = torch.randn(3, 3)
+        self.assertEqual(gm(x2), f(x2))
+
+    def test_make_fx_data_dependent(self):
+        @leaf_function
+        def data_dep_fn(x):
+            nz = (x > 0).nonzero()
+            return (x.relu(), nz)
+
+        @data_dep_fn.register_fake
+        def data_dep_fn_fake(x):
+            return (x.relu(), (x > 0).nonzero())
+
+        def f(x):
+            return data_dep_fn(x)
+
+        x = torch.randn(4, 4)
+        gm = make_fx(f)(x)
+
+        self.assertTrue(self._has_invoke_leaf_function_node(gm))
+
+        x2 = torch.randn(4, 4)
+        eager_out = f(x2)
+        gm_out = gm(x2)
+        self.assertEqual(gm_out[0], eager_out[0])
+        self.assertEqual(gm_out[1], eager_out[1])
+
+    def test_make_fx_pytree_inputs(self):
+        @leaf_function
+        def pytree_fn(inputs):
+            return (inputs["x"] + inputs["y"],)
+
+        @pytree_fn.register_fake
+        def pytree_fn_fake(inputs):
+            return (inputs["x"] + inputs["y"],)
+
+        def f(x, y):
+            return pytree_fn({"x": x, "y": y})
+
+        x = torch.randn(3, 3)
+        y = torch.randn(3, 3)
+        gm = make_fx(f)(x, y)
+
+        self.assertTrue(self._has_invoke_leaf_function_node(gm))
+        self.assertExpectedInline(
+            normalize_gm(gm.print_readable(print_output=False)),
+            """\
+class f(torch.nn.Module):
+    def forward(self, x_1: "f32[3, 3]", y_1: "f32[3, 3]"):
+        _opaque_obj0 = self._opaque_obj0
+        _opaque_obj1 = self._opaque_obj1
+        _tree_spec_constant0 = self._tree_spec_constant0
+        invoke_leaf_function = torch.ops.higher_order.invoke_leaf_function(_opaque_obj0, _opaque_obj1, _tree_spec_constant0, '', x_1, y_1, requires_grad_indices = '');  _opaque_obj0 = _opaque_obj1 = _tree_spec_constant0 = x_1 = y_1 = None
+        getitem: "f32[3, 3]" = invoke_leaf_function[0];  invoke_leaf_function = None
+        return (getitem,)
+""",  # noqa: B950
+        )
+
+        x2 = torch.randn(3, 3)
+        y2 = torch.randn(3, 3)
+        self.assertEqual(gm(x2, y2), f(x2, y2))
+
+    def test_make_fx_no_fake_impl(self):
+        @leaf_function
+        def no_fake_fn(x):
+            return (x * 2,)
+
+        def f(x):
+            return no_fake_fn(x)
+
+        x = torch.randn(3, 3)
+        with self.assertRaisesRegex(Exception, "requires a fake implementation"):
+            make_fx(f)(x)
+
+    def test_make_fx_with_module(self):
+        @leaf_function
+        def mod_fn(mod, x):
+            return (mod(x),)
+
+        @mod_fn.register_fake
+        def mod_fn_fake(mod, x):
+            return (mod(x),)
+
+        linear = torch.nn.Linear(3, 3)
+
+        def f(w, b, x):
+            with torch.nn.utils.stateless._reparametrize_module(
+                linear, {"weight": w, "bias": b}
+            ):
+                return mod_fn(linear, x)
+
+        w = linear.weight.detach().clone()
+        b = linear.bias.detach().clone()
+        x = torch.randn(2, 3)
+        gm = make_fx(f)(w, b, x)
+
+        self.assertTrue(self._has_invoke_leaf_function_node(gm))
+
+        w2 = torch.randn(3, 3)
+        b2 = torch.randn(3)
+        x2 = torch.randn(2, 3)
+        self.assertEqual(gm(w2, b2, x2), f(w2, b2, x2))
+
+    def test_make_fx_over_invoke_leaf_function_hop(self):
+        """Tracing invoke_leaf_function HOP directly with make_fx should produce
+        a single opaque invoke_leaf_function node, not trace into its internals."""
+        import torch.utils._pytree as pytree
+        from torch._higher_order_ops.invoke_leaf_function import _LeafCallable
+
+        def real_fn(x, y):
+            return (x * y + x,)
+
+        def fake_fn(x, y):
+            return (x * y + x,)
+
+        real_fn_callable = _LeafCallable(real_fn)
+        fake_fn_callable = _LeafCallable(fake_fn)
+
+        def f(x, y):
+            _, input_spec = pytree.tree_flatten(((x, y), {}))
+            return invoke_leaf_function(
+                real_fn_callable, fake_fn_callable, input_spec, "", x, y
+            )
+
+        x = torch.randn(3, 3)
+        y = torch.randn(3, 3)
+        gm = make_fx(f)(x, y)
+
+        self.assertTrue(self._has_invoke_leaf_function_node(gm))
+        self.assertExpectedInline(
+            normalize_gm(gm.print_readable(print_output=False)),
+            """\
+class f(torch.nn.Module):
+    def forward(self, x_1: "f32[3, 3]", y_1: "f32[3, 3]"):
+        _opaque_obj0 = self._opaque_obj0
+        _opaque_obj1 = self._opaque_obj1
+        _tree_spec_constant0 = self._tree_spec_constant0
+        invoke_leaf_function = torch.ops.higher_order.invoke_leaf_function(_opaque_obj0, _opaque_obj1, _tree_spec_constant0, '', x_1, y_1, requires_grad_indices = '');  _opaque_obj0 = _opaque_obj1 = _tree_spec_constant0 = x_1 = y_1 = None
+        getitem: "f32[3, 3]" = invoke_leaf_function[0];  invoke_leaf_function = None
+        return (getitem,)
+""",  # noqa: B950
+        )
+
+        x2 = torch.randn(3, 3)
+        y2 = torch.randn(3, 3)
+        self.assertEqual(gm(x2, y2), f(x2, y2))
+
+    def test_make_fx_input_mutation(self):
+        @leaf_function(mutates_args={"buf"})
+        def mutate_fn(x, buf):
+            buf.add_(1)
+            return (x + buf,)
+
+        @mutate_fn.register_fake
+        def mutate_fn_fake(x, buf):
+            buf.add_(1)
+            return (x + buf,)
+
+        def f(x, buf):
+            return mutate_fn(x, buf)
+
+        x = torch.randn(3, 3)
+        buf = torch.randn(3, 3)
+
+        buf_eager = buf.clone()
+        eager_out = f(x, buf_eager)
+
+        buf_fx = buf.clone()
+        gm = make_fx(f)(x, buf_fx)
+        self.assertTrue(self._has_invoke_leaf_function_node(gm))
+
+        buf_gm = buf.clone()
+        gm_out = gm(x, buf_gm)
+        self.assertEqual(gm_out, eager_out)
+        self.assertEqual(buf_gm, buf_eager)
+
+
+@skipIfTorchDynamo("leaf_function tests manage their own compilation")
+class TestLeafFunctionAotFunction(TestCase):
+    def test_aot_function_simple(self):
+        @leaf_function
+        def my_fn(x, y):
+            return (x @ y,)
+
+        @my_fn.register_fake
+        def my_fn_fake(x, y):
+            return (x @ y,)
+
+        def f(x, y):
+            return my_fn(x, y)[0]
+
+        fw_graph_cell = [None]
+        bw_graph_cell = [None]
+
+        x = torch.randn(3, 3, requires_grad=True)
+        y = torch.randn(3, 3, requires_grad=True)
+        x_clone = x.clone().detach().requires_grad_(True)
+        y_clone = y.clone().detach().requires_grad_(True)
+
+        compiled_f = aot_function(
+            f,
+            fw_compiler=partial(extract_graph, graph_cell=fw_graph_cell),
+            bw_compiler=partial(extract_graph, graph_cell=bw_graph_cell),
+        )
+
+        out_eager = f(x, y)
+        out_compiled = compiled_f(x_clone, y_clone)
+        self.assertEqual(out_eager, out_compiled)
+
+        out_eager.sum().backward()
+        out_compiled.sum().backward()
+        self.assertEqual(x.grad, x_clone.grad)
+        self.assertEqual(y.grad, y_clone.grad)
+
+        self.assertExpectedInline(
+            normalize_gm(fw_graph_cell[0].print_readable(print_output=False)),
+            """\
+class GraphModule(torch.nn.Module):
+    def forward(self, primals_1: "f32[0]", primals_2: "f32[3, 3]", primals_3: "f32[3, 3]"):
+        _opaque_obj0 = self._opaque_obj0
+        _opaque_obj1 = self._opaque_obj1
+        _tree_spec_constant0 = self._tree_spec_constant0
+        with_effects = torch.ops.higher_order.with_effects(primals_1, torch.ops.higher_order.invoke_leaf_function, _opaque_obj0, _opaque_obj1, _tree_spec_constant0, '', primals_2, primals_3, requires_grad_indices = '0,1');  primals_1 = _opaque_obj0 = _opaque_obj1 = _tree_spec_constant0 = primals_2 = primals_3 = None
+
+        getitem: "f32[0]" = with_effects[0]
+        getitem_1: "f32[3, 3]" = with_effects[1];  with_effects = None
+        return (getitem, getitem_1)
+""",  # noqa: B950
+        )
+        self.assertExpectedInline(
+            normalize_gm(bw_graph_cell[0].print_readable(print_output=False)),
+            """\
+class GraphModule(torch.nn.Module):
+    def forward(self, tangents_1: "f32[3, 3]", tangents_token: "f32[0]"):
+        _opaque_obj2 = self._opaque_obj2
+        _opaque_obj3 = self._opaque_obj3
+        _tree_spec_constant1 = self._tree_spec_constant1
+        with_effects_1 = torch.ops.higher_order.with_effects(tangents_token, torch.ops.higher_order.invoke_leaf_function, _opaque_obj2, _opaque_obj3, _tree_spec_constant1, '', tangents_1, requires_grad_indices = '');  tangents_token = _opaque_obj2 = _opaque_obj3 = _tree_spec_constant1 = tangents_1 = None
+        getitem_2: "f32[0]" = with_effects_1[0]
+        getitem_3: "f32[3, 3]" = with_effects_1[1]
+        getitem_4: "f32[3, 3]" = with_effects_1[2];  with_effects_1 = None
+        return (getitem_3, getitem_4, getitem_2)
+""",  # noqa: B950
+        )
+
+    def test_aot_function_gradients(self):
+        @leaf_function
+        def my_fn(x, y):
+            return (x * y + x,)
+
+        @my_fn.register_fake
+        def my_fn_fake(x, y):
+            return (x * y + x,)
+
+        def f(x, y):
+            return my_fn(x, y)[0].sum()
+
+        x = torch.randn(3, 3, requires_grad=True)
+        y = torch.randn(3, 3, requires_grad=True)
+        x_clone = x.clone().detach().requires_grad_(True)
+        y_clone = y.clone().detach().requires_grad_(True)
+
+        compiled_f = aot_function(f, fw_compiler=nop, bw_compiler=nop)
+
+        out_eager = f(x, y)
+        out_compiled = compiled_f(x_clone, y_clone)
+        self.assertEqual(out_eager, out_compiled)
+
+        out_eager.backward()
+        out_compiled.backward()
+        self.assertEqual(x.grad, x_clone.grad)
+        self.assertEqual(y.grad, y_clone.grad)
+
+    def test_aot_function_closure(self):
+        config = {"use_double": True}
+
+        @leaf_function
+        def closure_fn(x, y):
+            if config["use_double"]:
+                return (x @ y * 2,)
+            return (x @ y * 3,)
+
+        @closure_fn.register_fake
+        def closure_fn_fake(x, y):
+            return (x @ y,)
+
+        def f(x, y):
+            return closure_fn(x, y)[0]
+
+        compiled_f = aot_function(f, fw_compiler=nop, bw_compiler=nop)
+
+        x_val = torch.randn(3, 3)
+        y_val = torch.randn(3, 3)
+
+        # Test with use_double=True
+        config["use_double"] = True
+        x = x_val.clone().requires_grad_(True)
+        y = y_val.clone().requires_grad_(True)
+        x_clone = x_val.clone().requires_grad_(True)
+        y_clone = y_val.clone().requires_grad_(True)
+
+        out_eager = f(x, y)
+        out_compiled = compiled_f(x_clone, y_clone)
+        self.assertEqual(out_eager, out_compiled)
+
+        out_eager.sum().backward()
+        out_compiled.sum().backward()
+        self.assertEqual(x.grad, x_clone.grad)
+        self.assertEqual(y.grad, y_clone.grad)
+
+        # Change closure and verify result changes
+        config["use_double"] = False
+        x2 = x_val.clone().requires_grad_(True)
+        y2 = y_val.clone().requires_grad_(True)
+        x2_clone = x_val.clone().requires_grad_(True)
+        y2_clone = y_val.clone().requires_grad_(True)
+
+        out_eager2 = f(x2, y2)
+        out_compiled2 = compiled_f(x2_clone, y2_clone)
+        self.assertEqual(out_eager2, out_compiled2)
+
+        out_eager2.sum().backward()
+        out_compiled2.sum().backward()
+        self.assertEqual(x2.grad, x2_clone.grad)
+        self.assertEqual(y2.grad, y2_clone.grad)
+
+        # Different closures => different results
+        self.assertNotEqual(out_eager, out_eager2)
+
+    def test_aot_function_pytree_inputs(self):
+        @leaf_function
+        def pytree_fn(inputs):
+            return (inputs["x"] + inputs["y"],)
+
+        @pytree_fn.register_fake
+        def pytree_fn_fake(inputs):
+            return (inputs["x"] + inputs["y"],)
+
+        def f(x, y):
+            return pytree_fn({"x": x, "y": y})[0].sum()
+
+        x = torch.randn(3, 3, requires_grad=True)
+        y = torch.randn(3, 3, requires_grad=True)
+        x_clone = x.clone().detach().requires_grad_(True)
+        y_clone = y.clone().detach().requires_grad_(True)
+
+        compiled_f = aot_function(f, fw_compiler=nop, bw_compiler=nop)
+
+        out_eager = f(x, y)
+        out_compiled = compiled_f(x_clone, y_clone)
+        self.assertEqual(out_eager, out_compiled)
+
+        out_eager.backward()
+        out_compiled.backward()
+        self.assertEqual(x.grad, x_clone.grad)
+        self.assertEqual(y.grad, y_clone.grad)
+
+    def test_aot_function_input_mutation(self):
+        @leaf_function(mutates_args={"buf"})
+        def mutate_fn(x, buf):
+            buf.add_(1)
+            return (x + buf,)
+
+        @mutate_fn.register_fake
+        def mutate_fn_fake(x, buf):
+            buf.add_(1)
+            return (x + buf,)
+
+        def f(x, buf):
+            return mutate_fn(x, buf)[0]
+
+        x = torch.randn(3, 3)
+        buf = torch.randn(3, 3)
+
+        buf_eager = buf.clone()
+        eager_out = f(x, buf_eager)
+
+        compiled_f = aot_function(f, nop)
+        buf_compiled = buf.clone()
+        compiled_out = compiled_f(x, buf_compiled)
+        self.assertEqual(compiled_out, eager_out)
+        self.assertEqual(buf_compiled, buf_eager)
+
+
+@skipIfTorchDynamo("leaf_function tests manage their own compilation")
+class TestLeafFunctionEscapedGradients(TestCase):
+    def test_aot_function_escaped_gradient_multiple_closures(self):
+        weight1 = torch.randn(3, 3, requires_grad=True)
+        weight2 = torch.randn(3, 3, requires_grad=True)
+
+        @leaf_function
+        def uses_multiple_closures(x):
+            return (x @ weight1 + x @ weight2,)
+
+        @uses_multiple_closures.register_fake
+        def uses_multiple_closures_fake(x):
+            return (torch.empty(x.shape[0], 3),)
+
+        def f(x):
+            return uses_multiple_closures(x)[0]
+
+        x = torch.randn(2, 3, requires_grad=True)
+
+        compiled_f = aot_function(f, fw_compiler=nop, bw_compiler=nop)
+        with config.patch(leaf_function_check_escaped_gradients=True):
+            with self.assertRaisesRegex(RuntimeError, "2 tensor"):
+                compiled_f(x)
+
+    def test_aot_function_escaped_gradient_disabled(self):
+        weight = torch.randn(3, 3, requires_grad=True)
+
+        @leaf_function
+        def uses_closure(x):
+            return (x @ weight,)
+
+        @uses_closure.register_fake
+        def uses_closure_fake(x):
+            return (torch.empty(x.shape[0], 3),)
+
+        def f(x):
+            return uses_closure(x)[0]
+
+        x = torch.randn(2, 3, requires_grad=True)
+
+        compiled_f = aot_function(f, fw_compiler=nop, bw_compiler=nop)
+        with config.patch(leaf_function_check_escaped_gradients=False):
+            result = compiled_f(x)
+            self.assertEqual(result.shape, (2, 3))
+
+    def test_aot_function_escaped_gradient_input_no_grad(self):
+        """No false positive when input doesn't require grad."""
+        closure_weight = torch.randn(3, 3, requires_grad=True)
+
+        @leaf_function
+        def uses_closure(x):
+            return (x @ closure_weight,)
+
+        @uses_closure.register_fake
+        def uses_closure_fake(x):
+            return (torch.empty(x.shape[0], 3),)
+
+        def f(x):
+            return uses_closure(x)[0]
+
+        x = torch.randn(2, 3, requires_grad=False)
+
+        compiled_f = aot_function(f, fw_compiler=nop, bw_compiler=nop)
+        with config.patch(leaf_function_check_escaped_gradients=True):
+            result = compiled_f(x)
+            self.assertEqual(result.shape, (2, 3))
+
+    def test_aot_function_escaped_gradient_actually_lost(self):
+        """Verify gradients don't flow to closure tensors."""
+        closure_weight = torch.randn(3, 3, requires_grad=True)
+
+        @leaf_function
+        def uses_closure(x):
+            return (x @ closure_weight,)
+
+        @uses_closure.register_fake
+        def uses_closure_fake(x):
+            return (torch.empty(x.shape[0], 3),)
+
+        def f(x):
+            return uses_closure(x)[0]
+
+        x = torch.randn(2, 3, requires_grad=True)
+
+        compiled_f = aot_function(f, fw_compiler=nop, bw_compiler=nop)
+        result = compiled_f(x)
+        result.sum().backward()
+
+        self.assertIsNotNone(x.grad)
+        self.assertIsNone(closure_weight.grad)
+
+    def test_make_fx_escaped_gradient_multiple_closures(self):
+        weight1 = torch.randn(3, 3, requires_grad=True)
+        weight2 = torch.randn(3, 3, requires_grad=True)
+
+        @leaf_function
+        def uses_multiple_closures(x):
+            return (x @ weight1 + x @ weight2,)
+
+        @uses_multiple_closures.register_fake
+        def uses_multiple_closures_fake(x):
+            return (torch.empty(x.shape[0], 3),)
+
+        def f(x):
+            return uses_multiple_closures(x)
+
+        x = torch.randn(2, 3, requires_grad=True)
+
+        with config.patch(leaf_function_check_escaped_gradients=True):
+            with self.assertRaisesRegex(RuntimeError, "2 tensor"):
+                make_fx(f)(x)
+
+    def test_make_fx_escaped_gradient_disabled(self):
+        weight = torch.randn(3, 3, requires_grad=True)
+
+        @leaf_function
+        def uses_closure(x):
+            return (x @ weight,)
+
+        @uses_closure.register_fake
+        def uses_closure_fake(x):
+            return (torch.empty(x.shape[0], 3),)
+
+        def f(x):
+            return uses_closure(x)
+
+        x = torch.randn(2, 3, requires_grad=True)
+
+        with config.patch(leaf_function_check_escaped_gradients=False):
+            gm = make_fx(f)(x)
+            x2 = torch.randn(2, 3)
+            result = gm(x2)
+            self.assertEqual(result[0].shape, (2, 3))
+
+
+@skipIfTorchDynamo("leaf_function tests manage their own compilation")
+class TestLeafFunctionMakeFxAndCompile(TestCase):
+    """Tests for @leaf_function when mixing torch.compile and make_fx."""
+
+    def test_not_called_during_compilation(self):
+        """The real leaf_fn body runs only at runtime, not during compilation."""
+        torch._dynamo.reset()
+        call_count = 0
+
+        @leaf_function
+        def my_leaf(x, y):
+            nonlocal call_count
+            call_count += 1
+            return (x + y,)
+
+        @my_leaf.register_fake
+        def my_leaf_fake(x, y):
+            return (torch.empty_like(x),)
+
+        def f(x, y):
+            return my_leaf(x, y)[0]
+
+        compiled_f = torch.compile(f, backend="eager", fullgraph=True)
+
+        x = torch.randn(3, 3)
+        y = torch.randn(3, 3)
+
+        # Compilation + first call
+        result = compiled_f(x, y)
+        self.assertEqual(call_count, 1)
+        self.assertEqual(result, x + y)
+
+        # Second call reuses compiled code, leaf_fn called again exactly once
+        x2 = torch.randn(3, 3)
+        y2 = torch.randn(3, 3)
+        result2 = compiled_f(x2, y2)
+        self.assertEqual(call_count, 2)
+        self.assertEqual(result2, x2 + y2)
+
+    @config.patch(force_compile_during_fx_trace=True)
+    def test_leaf_fn_only_in_compile(self):
+        """Leaf function only inside the torch.compile'd region."""
+        torch._dynamo.reset()
+
+        @leaf_function
+        def my_leaf(x, y):
+            return (x * y + x,)
+
+        @my_leaf.register_fake
+        def my_leaf_fake(x, y):
+            return (x * y + x,)
+
+        def inner(x, y):
+            return my_leaf(x, y)[0]
+
+        compiled = torch.compile(inner, backend="invoke_subgraph")
+
+        def outer(x, y):
+            z = x + 1
+            return compiled(z, y) - 1
+
+        x = torch.randn(3, 3)
+        y = torch.randn(3, 3)
+        traced = make_fx(
+            outer, tracing_mode="fake", _disable_torch_fn_metadata_mode=True
+        )(x, y)
+
+        self.assertExpectedInline(
+            normalize_gm(traced.print_readable(print_output=False)),
+            """\
+class outer(torch.nn.Module):
+    def forward(self, x_1: "f32[3, 3]", y_1: "f32[3, 3]"):
+        add: "f32[3, 3]" = torch.ops.aten.add.Tensor(x_1, 1);  x_1 = None
+        repeated_subgraph0 = self.repeated_subgraph0
+        invoke_subgraph = torch.ops.higher_order.invoke_subgraph(repeated_subgraph0, 'invoke_subgraph_0', None, add, y_1);  repeated_subgraph0 = add = y_1 = None
+        getitem: "f32[0]" = invoke_subgraph[0];  getitem = None
+        getitem_1: "f32[3, 3]" = invoke_subgraph[1];  invoke_subgraph = None
+        sub: "f32[3, 3]" = torch.ops.aten.sub.Tensor(getitem_1, 1);  getitem_1 = None
+        return sub
+
+    class repeated_subgraph0(torch.nn.Module):
+        def forward(self, arg0_1, arg1_1: "f32[3, 3]", arg2_1: "f32[3, 3]"):
+            _opaque_obj0 = self._opaque_obj0
+            _opaque_obj1 = self._opaque_obj1
+            _tree_spec_constant0 = self._tree_spec_constant0
+            with_effects = torch.ops.higher_order.with_effects(None, torch.ops.higher_order.invoke_leaf_function, _opaque_obj0, _opaque_obj1, _tree_spec_constant0, '', arg1_1, arg2_1, requires_grad_indices = '');  _opaque_obj0 = _opaque_obj1 = _tree_spec_constant0 = arg1_1 = arg2_1 = None
+            getitem: "f32[0]" = with_effects[0]
+            getitem_1: "f32[3, 3]" = with_effects[1];  with_effects = None
+            return (getitem, getitem_1)
+""",  # noqa: B950
+        )
+
+        x2 = torch.randn(3, 3)
+        y2 = torch.randn(3, 3)
+        expected = (x2 + 1) * y2 + (x2 + 1) - 1
+        self.assertEqual(traced(x2, y2), expected)
+
+    @config.patch(force_compile_during_fx_trace=True)
+    def test_leaf_fn_only_in_make_fx(self):
+        """Leaf function only in the make_fx-traced code, outside torch.compile."""
+        torch._dynamo.reset()
+
+        @leaf_function
+        def my_leaf(x):
+            return (x * 2,)
+
+        @my_leaf.register_fake
+        def my_leaf_fake(x):
+            return (torch.empty_like(x),)
+
+        def inner(x):
+            return x + 1
+
+        compiled = torch.compile(inner, backend="invoke_subgraph")
+
+        def outer(x):
+            y = compiled(x)
+            return my_leaf(y)
+
+        x = torch.randn(3, 3)
+        traced = make_fx(
+            outer, tracing_mode="fake", _disable_torch_fn_metadata_mode=True
+        )(x)
+
+        self.assertExpectedInline(
+            normalize_gm(traced.print_readable(print_output=False)),
+            """\
+class outer(torch.nn.Module):
+    def forward(self, x_1: "f32[3, 3]"):
+        repeated_subgraph0 = self.repeated_subgraph0
+        invoke_subgraph = torch.ops.higher_order.invoke_subgraph(repeated_subgraph0, 'invoke_subgraph_0', x_1);  repeated_subgraph0 = x_1 = None
+        getitem: "f32[3, 3]" = invoke_subgraph[0];  invoke_subgraph = None
+        _opaque_obj0 = self._opaque_obj0
+        _opaque_obj1 = self._opaque_obj1
+        _tree_spec_constant0 = self._tree_spec_constant0
+        invoke_leaf_function = torch.ops.higher_order.invoke_leaf_function(_opaque_obj0, _opaque_obj1, _tree_spec_constant0, '', getitem, requires_grad_indices = '');  _opaque_obj0 = _opaque_obj1 = _tree_spec_constant0 = getitem = None
+        getitem_1: "f32[3, 3]" = invoke_leaf_function[0];  invoke_leaf_function = None
+        return (getitem_1,)
+
+    class repeated_subgraph0(torch.nn.Module):
+        def forward(self, arg0_1: "f32[3, 3]"):
+            add: "f32[3, 3]" = torch.ops.aten.add.Tensor(arg0_1, 1);  arg0_1 = None
+            return (add,)
+""",  # noqa: B950
+        )
+
+        x2 = torch.randn(3, 3)
+        expected = ((x2 + 1) * 2,)
+        self.assertEqual(traced(x2), expected)
+
+    @config.patch(force_compile_during_fx_trace=True)
+    def test_leaf_fn_in_both_compile_and_make_fx(self):
+        """Same leaf function used inside torch.compile and in make_fx-traced code."""
+        torch._dynamo.reset()
+
+        @leaf_function
+        def my_leaf(x):
+            return (x * 2,)
+
+        @my_leaf.register_fake
+        def my_leaf_fake(x):
+            return (torch.empty_like(x),)
+
+        def inner(x):
+            return my_leaf(x)[0]
+
+        compiled = torch.compile(inner, backend="invoke_subgraph")
+
+        def outer(x):
+            y = compiled(x)
+            return my_leaf(y)
+
+        x = torch.randn(3, 3)
+        traced = make_fx(
+            outer, tracing_mode="fake", _disable_torch_fn_metadata_mode=True
+        )(x)
+
+        self.assertExpectedInline(
+            normalize_gm(traced.print_readable(print_output=False)),
+            """\
+class outer(torch.nn.Module):
+    def forward(self, x_1: "f32[3, 3]"):
+        repeated_subgraph0 = self.repeated_subgraph0
+        invoke_subgraph = torch.ops.higher_order.invoke_subgraph(repeated_subgraph0, 'invoke_subgraph_0', None, x_1);  repeated_subgraph0 = x_1 = None
+        getitem: "f32[0]" = invoke_subgraph[0];  getitem = None
+        getitem_1: "f32[3, 3]" = invoke_subgraph[1];  invoke_subgraph = None
+        _opaque_obj0 = self._opaque_obj0
+        _opaque_obj1 = self._opaque_obj1
+        _tree_spec_constant0 = self._tree_spec_constant0
+        invoke_leaf_function = torch.ops.higher_order.invoke_leaf_function(_opaque_obj0, _opaque_obj1, _tree_spec_constant0, '', getitem_1, requires_grad_indices = '');  _opaque_obj0 = _opaque_obj1 = _tree_spec_constant0 = getitem_1 = None
+        getitem_2: "f32[3, 3]" = invoke_leaf_function[0];  invoke_leaf_function = None
+        return (getitem_2,)
+
+    class repeated_subgraph0(torch.nn.Module):
+        def forward(self, arg0_1, arg1_1: "f32[3, 3]"):
+            _opaque_obj0 = self._opaque_obj0
+            _opaque_obj1 = self._opaque_obj1
+            _tree_spec_constant0 = self._tree_spec_constant0
+            with_effects = torch.ops.higher_order.with_effects(None, torch.ops.higher_order.invoke_leaf_function, _opaque_obj0, _opaque_obj1, _tree_spec_constant0, '', arg1_1, requires_grad_indices = '');  _opaque_obj0 = _opaque_obj1 = _tree_spec_constant0 = arg1_1 = None
+            getitem: "f32[0]" = with_effects[0]
+            getitem_1: "f32[3, 3]" = with_effects[1];  with_effects = None
+            return (getitem, getitem_1)
+""",  # noqa: B950
+        )
+
+
+@skipIfTorchDynamo("leaf_function tests manage their own compilation")
+class TestLeafFunctionDynamo(PytreeRegisteringTestCase):
+    def _assert_models_equal(
+        self,
+        model_expected,
+        model_test,
+        x_expected,
+        x_test,
+    ):
+        out_expected = model_expected(x_expected)
+        out_test = model_test(x_test)
+        self.assertEqual(out_expected, out_test)
+
+        loss_expected = out_expected.sum()
+        loss_test = out_test.sum()
+        loss_expected.backward()
+        loss_test.backward()
+        self.assertEqual(x_expected.grad, x_test.grad)
+
+        expected_grads = {
+            name: param.grad for name, param in model_expected.named_parameters()
+        }
+        test_grads = {name: param.grad for name, param in model_test.named_parameters()}
+
+        self.assertEqual(set(expected_grads.keys()), set(test_grads.keys()))
+        for name in expected_grads:
+            if expected_grads[name] is not None:
+                self.assertEqual(
+                    expected_grads[name],
+                    test_grads[name],
+                    msg=f"Gradient mismatch for parameter {name}",
+                )
+
+    def _test_leaf_function_helper(self, mod_class, args_fn, loss_fn):
+        import torch.utils._pytree as pytree
+        from torch._dynamo.testing import AotEagerAndRecordGraphs, EagerAndRecordGraphs
+
+        mod_eager = mod_class()
+        mod_compile_eager = mod_class()
+        mod_compile_eager.load_state_dict(dict(mod_eager.state_dict()))
+        mod_compile_aot = mod_class()
+        mod_compile_aot.load_state_dict(dict(mod_eager.state_dict()))
+
+        eager_backend = EagerAndRecordGraphs()
+        compiled_eager = torch.compile(
+            mod_compile_eager, backend=eager_backend, fullgraph=True
+        )
+
+        backend = AotEagerAndRecordGraphs()
+        compiled_aot = torch.compile(mod_compile_aot, backend=backend, fullgraph=True)
+
+        for _ in range(2):
+            mod_eager.zero_grad()
+            mod_compile_eager.zero_grad()
+            mod_compile_aot.zero_grad()
+
+            args = args_fn()
+            args_clone = pytree.tree_map(
+                lambda x: x.clone().detach().requires_grad_(x.requires_grad), args
+            )
+            args_clone2 = pytree.tree_map(
+                lambda x: x.clone().detach().requires_grad_(x.requires_grad), args
+            )
+
+            out_eager = mod_eager(*args)
+            loss_fn(out_eager).backward()
+
+            out_compile_eager = compiled_eager(*args_clone)
+            loss_fn(out_compile_eager).backward()
+
+            out_compile_aot = compiled_aot(*args_clone2)
+            loss_fn(out_compile_aot).backward()
+
+            self.assertEqual(out_eager, out_compile_eager)
+            self.assertEqual(out_eager, out_compile_aot)
+
+            for (name_eager, param_eager), (_, param_compile_eager), (
+                _,
+                param_compile_aot,
+            ) in zip(
+                mod_eager.named_parameters(),
+                mod_compile_eager.named_parameters(),
+                mod_compile_aot.named_parameters(),
+            ):
+                self.assertEqual(
+                    param_eager.grad,
+                    param_compile_eager.grad,
+                    msg=f"Gradient mismatch for {name_eager} between eager and compile_eager",
+                )
+                self.assertEqual(
+                    param_eager.grad,
+                    param_compile_aot.grad,
+                    msg=f"Gradient mismatch for {name_eager} between eager and compile_aot",
+                )
+
+            pytree.tree_map(
+                lambda x, compile_x: self.assertEqual(x.grad, compile_x.grad)
+                if isinstance(x, torch.Tensor) and x.requires_grad
+                else None,
+                args,
+                args_clone,
+            )
+            pytree.tree_map(
+                lambda x, compile_x: self.assertEqual(x.grad, compile_x.grad)
+                if isinstance(x, torch.Tensor) and x.requires_grad
+                else None,
+                args,
+                args_clone2,
+            )
+
+        def _normalize(gm):
+            s = normalize_gm(gm.print_readable(print_output=False))
+            # Normalize nn_module_index which varies depending on whether an
+            # accelerator is available (stream reserves index 0).
+            return re.sub(r"'', \d+, ", "'', 0, ", s)
+
+        return (
+            _normalize(eager_backend.graphs[0]),
+            _normalize(backend.fw_graphs[0]),
+            _normalize(backend.bw_graphs[0]),
+        )
+
+    def test_leaf_function_simple(self):
+        @leaf_function
+        def non_tracable_forward(mod, x):
+            if x.sum() > 0:
+                return (mod.linear(x),)
+            else:
+                return (mod.linear(x) + x,)
+
+        @non_tracable_forward.register_fake
+        def non_tracable_forward_fake(mod, x):
+            return (mod.linear(x),)
+
+        class NonTracable(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.linear = torch.nn.Linear(3, 3)
+
+            def forward(self, x):
+                return non_tracable_forward(self, x)
+
+        def args_fn():
+            return (torch.randn(3, 3, requires_grad=True),)
+
+        def loss_fn(out):
+            return out[0].sum()
+
+        dynamo_graph_str, fw_graph_str, bw_graph_str = self._test_leaf_function_helper(
+            NonTracable, args_fn, loss_fn
+        )
+        self.assertExpectedInline(
+            dynamo_graph_str,
+            """\
+class GraphModule(torch.nn.Module):
+    def forward(self, L_x_: "f32[3, 3]", L_self_modules_linear_parameters_weight_: "f32[3, 3]", L_self_modules_linear_parameters_bias_: "f32[3]"):
+        l_x_ = L_x_
+        l_self_modules_linear_parameters_weight_ = L_self_modules_linear_parameters_weight_
+        l_self_modules_linear_parameters_bias_ = L_self_modules_linear_parameters_bias_
+
+        real_fn : torch._higher_order_ops.invoke_leaf_function._LeafCallable = self.real_fn
+        fake_fn : torch._higher_order_ops.invoke_leaf_function._LeafCallable = self.fake_fn
+        input_spec : torch.utils._pytree.TreeSpec = self.input_spec
+        invoke_leaf_function = torch.ops.higher_order.invoke_leaf_function(real_fn, fake_fn, input_spec, '', 0, l_self_modules_linear_parameters_weight_, l_self_modules_linear_parameters_bias_, l_x_);  real_fn = fake_fn = input_spec = l_self_modules_linear_parameters_weight_ = l_self_modules_linear_parameters_bias_ = l_x_ = None
+        getitem: "f32[3, 3]" = invoke_leaf_function[0];  invoke_leaf_function = None
+        return (getitem,)
+""",  # noqa: B950
+        )
+        self.assertExpectedInline(
+            fw_graph_str,
+            """\
+class GraphModule(torch.nn.Module):
+    def forward(self, primals_1: "f32[0]", primals_2: "f32[3, 3]", primals_3: "f32[3, 3]", primals_4: "f32[3]"):
+        _opaque_obj0 = self._opaque_obj0
+        _opaque_obj1 = self._opaque_obj1
+        _tree_spec_constant0 = self._tree_spec_constant0
+        with_effects = torch.ops.higher_order.with_effects(primals_1, torch.ops.higher_order.invoke_leaf_function, _opaque_obj0, _opaque_obj1, _tree_spec_constant0, '', 0, primals_3, primals_4, primals_2, requires_grad_indices = '1,2,3');  primals_1 = _opaque_obj0 = _opaque_obj1 = _tree_spec_constant0 = primals_3 = primals_4 = primals_2 = None
+
+        getitem: "f32[0]" = with_effects[0]
+        getitem_1: "f32[3, 3]" = with_effects[1];  with_effects = None
+        return (getitem, getitem_1)
+""",  # noqa: B950
+        )
+        self.assertExpectedInline(
+            bw_graph_str,
+            """\
+class GraphModule(torch.nn.Module):
+    def forward(self, tangents_1: "f32[3, 3]", tangents_token: "f32[0]"):
+        _opaque_obj2 = self._opaque_obj2
+        _opaque_obj3 = self._opaque_obj3
+        _tree_spec_constant1 = self._tree_spec_constant1
+        with_effects_1 = torch.ops.higher_order.with_effects(tangents_token, torch.ops.higher_order.invoke_leaf_function, _opaque_obj2, _opaque_obj3, _tree_spec_constant1, '', tangents_1, requires_grad_indices = '');  tangents_token = _opaque_obj2 = _opaque_obj3 = _tree_spec_constant1 = tangents_1 = None
+        getitem_2: "f32[0]" = with_effects_1[0]
+        getitem_4: "f32[3, 3]" = with_effects_1[2]
+        getitem_5: "f32[3]" = with_effects_1[3]
+        getitem_6: "f32[3, 3]" = with_effects_1[4];  with_effects_1 = None
+        return (getitem_6, getitem_4, getitem_5, getitem_2)
+""",  # noqa: B950
+        )
+
+    def test_leaf_function_with_logging(self):
+        @leaf_function
+        def logging_forward(mod, x):
+            print("Processing input")
+            return (mod.linear(x),)
+
+        @logging_forward.register_fake
+        def logging_forward_fake(mod, x):
+            return (mod.linear(x),)
+
+        class LoggingModule(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.linear = torch.nn.Linear(3, 3)
+
+            def forward(self, x):
+                return logging_forward(self, x)
+
+        def args_fn():
+            return (torch.randn(3, 3, requires_grad=True),)
+
+        def loss_fn(out):
+            return out[0].sum()
+
+        with patch("builtins.print") as mock_print:
+            self._test_leaf_function_helper(LoggingModule, args_fn, loss_fn)
+            mock_print.assert_any_call("Processing input")
+            self.assertEqual(mock_print.call_count, 6)
+
+    def test_leaf_function_dynamic_autograd_module_config(self):
+        from torch._dynamo.testing import CompileCounterWithBackend
+
+        @leaf_function
+        def configurable_scale(mod, x):
+            if mod.use_double_scale:
+                return (mod.linear(x) * 2,)
+            else:
+                return (mod.linear(x) * 3,)
+
+        @configurable_scale.register_fake
+        def configurable_scale_fake(mod, x):
+            return (mod.linear(x),)
+
+        class ConfigurableModule(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.linear = torch.nn.Linear(3, 3)
+                self.use_double_scale = True
+
+            def forward(self, x):
+                return configurable_scale(self, x)
+
+        mod_eager = ConfigurableModule()
+        mod_compiled = ConfigurableModule()
+        mod_compiled.load_state_dict(dict(mod_eager.state_dict()))
+
+        counter = CompileCounterWithBackend("aot_eager")
+        compiled_fn = torch.compile(mod_compiled, backend=counter, fullgraph=True)
+
+        x_value = torch.randn(3, 3)
+
+        mod_eager.use_double_scale = True
+        mod_compiled.use_double_scale = True
+
+        x1 = x_value.clone().requires_grad_(True)
+        x1_clone = x_value.clone().requires_grad_(True)
+
+        out_eager_1 = mod_eager(x1)
+        out_eager_1[0].sum().backward()
+
+        out_compiled_1 = compiled_fn(x1_clone)
+        out_compiled_1[0].sum().backward()
+
+        self.assertEqual(out_eager_1, out_compiled_1)
+        self.assertEqual(x1.grad, x1_clone.grad)
+
+        mod_eager.zero_grad()
+        mod_compiled.zero_grad()
+
+        mod_eager.use_double_scale = False
+        mod_compiled.use_double_scale = False
+
+        x2 = x_value.clone().requires_grad_(True)
+        x2_clone = x_value.clone().requires_grad_(True)
+
+        out_eager_2 = mod_eager(x2)
+        out_eager_2[0].sum().backward()
+
+        out_compiled_2 = compiled_fn(x2_clone)
+        out_compiled_2[0].sum().backward()
+
+        self.assertEqual(out_eager_2, out_compiled_2)
+        self.assertEqual(x2.grad, x2_clone.grad)
+
+        self.assertNotEqual(x1.grad, x2.grad)
+
+        self.assertEqual(counter.frame_count, 1)
+
+    def test_leaf_function_dynamic_autograd_closure(self):
+        from torch._dynamo.testing import CompileCounterWithBackend
+
+        closure_config = {"use_double_scale": True}
+
+        @leaf_function
+        def configurable_scale(x, y):
+            if closure_config["use_double_scale"]:
+                return (x @ y * 2,)
+            else:
+                return (x @ y * 3,)
+
+        @configurable_scale.register_fake
+        def configurable_scale_fake(x, y):
+            return (x @ y,)
+
+        def fn(x, y):
+            return configurable_scale(x, y)
+
+        counter = CompileCounterWithBackend("aot_eager")
+        compiled_fn = torch.compile(fn, backend=counter, fullgraph=True)
+
+        x_value = torch.randn(3, 3)
+        y_value = torch.randn(3, 3)
+
+        closure_config["use_double_scale"] = True
+
+        x1 = x_value.clone().requires_grad_(True)
+        y1 = y_value.clone().requires_grad_(True)
+        x1_clone = x_value.clone().requires_grad_(True)
+        y1_clone = y_value.clone().requires_grad_(True)
+
+        out_eager_1 = fn(x1, y1)
+        out_eager_1[0].sum().backward()
+
+        out_compiled_1 = compiled_fn(x1_clone, y1_clone)
+        out_compiled_1[0].sum().backward()
+
+        self.assertEqual(out_eager_1, out_compiled_1)
+        self.assertEqual(x1.grad, x1_clone.grad)
+        self.assertEqual(y1.grad, y1_clone.grad)
+
+        closure_config["use_double_scale"] = False
+
+        x2 = x_value.clone().requires_grad_(True)
+        y2 = y_value.clone().requires_grad_(True)
+        x2_clone = x_value.clone().requires_grad_(True)
+        y2_clone = y_value.clone().requires_grad_(True)
+
+        out_eager_2 = fn(x2, y2)
+        out_eager_2[0].sum().backward()
+
+        out_compiled_2 = compiled_fn(x2_clone, y2_clone)
+        out_compiled_2[0].sum().backward()
+
+        self.assertEqual(out_eager_2, out_compiled_2)
+        self.assertEqual(x2.grad, x2_clone.grad)
+        self.assertEqual(y2.grad, y2_clone.grad)
+
+        self.assertNotEqual(x1.grad, x2.grad)
+        self.assertNotEqual(y1.grad, y2.grad)
+
+        self.assertEqual(counter.frame_count, 1)
+
+    def test_leaf_function_closure_constants_without_grad(self):
+        closure_scale = 2.0
+        closure_tensor = torch.tensor([1.0, 2.0, 3.0])
+
+        @leaf_function
+        def closure_forward(mod, x):
+            out = mod.linear(x) * closure_scale * mod.scale
+            out = out + closure_tensor + mod.offset
+            return (out,)
+
+        @closure_forward.register_fake
+        def closure_forward_fake(mod, x):
+            return (mod.linear(x) + mod.offset,)
+
+        class ClosureModule(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.linear = torch.nn.Linear(3, 3)
+                self.scale = 3.0
+                self.offset = torch.nn.Parameter(torch.ones(3))
+
+            def forward(self, x):
+                return closure_forward(self, x)
+
+        def args_fn():
+            return (torch.randn(3, 3, requires_grad=True),)
+
+        def loss_fn(out):
+            return out[0].sum()
+
+        dynamo_graph_str, fw_graph_str, bw_graph_str = self._test_leaf_function_helper(
+            ClosureModule, args_fn, loss_fn
+        )
+        self.assertExpectedInline(
+            dynamo_graph_str,
+            """\
+class GraphModule(torch.nn.Module):
+    def forward(self, L_x_: "f32[3, 3]", L_self_parameters_offset_: "f32[3]", L_self_modules_linear_parameters_weight_: "f32[3, 3]", L_self_modules_linear_parameters_bias_: "f32[3]"):
+        l_x_ = L_x_
+        l_self_parameters_offset_ = L_self_parameters_offset_
+        l_self_modules_linear_parameters_weight_ = L_self_modules_linear_parameters_weight_
+        l_self_modules_linear_parameters_bias_ = L_self_modules_linear_parameters_bias_
+
+        real_fn : torch._higher_order_ops.invoke_leaf_function._LeafCallable = self.real_fn
+        fake_fn : torch._higher_order_ops.invoke_leaf_function._LeafCallable = self.fake_fn
+        input_spec : torch.utils._pytree.TreeSpec = self.input_spec
+        invoke_leaf_function = torch.ops.higher_order.invoke_leaf_function(real_fn, fake_fn, input_spec, '', 0, l_self_parameters_offset_, l_self_modules_linear_parameters_weight_, l_self_modules_linear_parameters_bias_, l_x_);  real_fn = fake_fn = input_spec = l_self_parameters_offset_ = l_self_modules_linear_parameters_weight_ = l_self_modules_linear_parameters_bias_ = l_x_ = None
+        getitem: "f32[3, 3]" = invoke_leaf_function[0];  invoke_leaf_function = None
+        return (getitem,)
+""",  # noqa: B950
+        )
+        self.assertExpectedInline(
+            fw_graph_str,
+            """\
+class GraphModule(torch.nn.Module):
+    def forward(self, primals_1: "f32[0]", primals_2: "f32[3, 3]", primals_3: "f32[3]", primals_4: "f32[3, 3]", primals_5: "f32[3]"):
+        _opaque_obj0 = self._opaque_obj0
+        _opaque_obj1 = self._opaque_obj1
+        _tree_spec_constant0 = self._tree_spec_constant0
+        with_effects = torch.ops.higher_order.with_effects(primals_1, torch.ops.higher_order.invoke_leaf_function, _opaque_obj0, _opaque_obj1, _tree_spec_constant0, '', 0, primals_3, primals_4, primals_5, primals_2, requires_grad_indices = '1,2,3,4');  primals_1 = _opaque_obj0 = _opaque_obj1 = _tree_spec_constant0 = primals_3 = primals_4 = primals_5 = primals_2 = None
+
+        getitem: "f32[0]" = with_effects[0]
+        getitem_1: "f32[3, 3]" = with_effects[1];  with_effects = None
+        return (getitem, getitem_1)
+""",  # noqa: B950
+        )
+        self.assertExpectedInline(
+            bw_graph_str,
+            """\
+class GraphModule(torch.nn.Module):
+    def forward(self, tangents_1: "f32[3, 3]", tangents_token: "f32[0]"):
+        _opaque_obj2 = self._opaque_obj2
+        _opaque_obj3 = self._opaque_obj3
+        _tree_spec_constant1 = self._tree_spec_constant1
+        with_effects_1 = torch.ops.higher_order.with_effects(tangents_token, torch.ops.higher_order.invoke_leaf_function, _opaque_obj2, _opaque_obj3, _tree_spec_constant1, '', tangents_1, requires_grad_indices = '');  tangents_token = _opaque_obj2 = _opaque_obj3 = _tree_spec_constant1 = tangents_1 = None
+        getitem_2: "f32[0]" = with_effects_1[0]
+        getitem_4: "f32[3]" = with_effects_1[2]
+        getitem_5: "f32[3, 3]" = with_effects_1[3]
+        getitem_6: "f32[3]" = with_effects_1[4]
+        getitem_7: "f32[3, 3]" = with_effects_1[5];  with_effects_1 = None
+        return (getitem_7, getitem_4, getitem_5, getitem_6, getitem_2)
+""",  # noqa: B950
+        )
+
+    def test_leaf_function_pytree_inputs(self):
+        @leaf_function
+        def pytree_forward(mod, inputs):
+            if inputs["x"].sum() > 0:
+                return (mod.linear(inputs["x"]), inputs["y"] + 1)
+            return (mod.linear(inputs["x"]) + inputs["y"], inputs["y"] - 1)
+
+        @pytree_forward.register_fake
+        def pytree_forward_fake(mod, inputs):
+            return (mod.linear(inputs["x"]), inputs["y"])
+
+        class PytreeModule(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.linear = torch.nn.Linear(3, 3)
+
+            def forward(self, inputs):
+                return pytree_forward(self, inputs)
+
+        def args_fn():
+            return (
+                {
+                    "x": torch.randn(3, 3, requires_grad=True),
+                    "y": torch.randn(3, 3, requires_grad=True),
+                },
+            )
+
+        def loss_fn(out):
+            return out[0].sum() + out[1].sum()
+
+        self._test_leaf_function_helper(PytreeModule, args_fn, loss_fn)
+
+    def test_leaf_function_nested_annotations(self):
+        @leaf_function
+        def inner_leaf_forward(mod, x):
+            y = mod.linear(x)
+            return (y + x,)
+
+        @inner_leaf_forward.register_fake
+        def inner_leaf_forward_fake(mod, x):
+            return (mod.linear(x),)
+
+        class InnerLeaf(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.linear = torch.nn.Linear(3, 3)
+
+            def forward(self, x):
+                return inner_leaf_forward(self, x)
+
+        @leaf_function
+        def outer_leaf_forward(mod, x):
+            z = mod.linear(x)
+            return mod.inner(z + x)
+
+        @outer_leaf_forward.register_fake
+        def outer_leaf_forward_fake(mod, x):
+            return mod.inner(mod.linear(x))
+
+        class OuterLeaf(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.inner = InnerLeaf()
+                self.linear = torch.nn.Linear(3, 3)
+
+            def forward(self, x):
+                return outer_leaf_forward(self, x)
+
+        def args_fn():
+            return (torch.randn(3, 3, requires_grad=True),)
+
+        def loss_fn(out):
+            return out[0].sum()
+
+        dynamo_graph_str, fw_graph_str, bw_graph_str = self._test_leaf_function_helper(
+            OuterLeaf, args_fn, loss_fn
+        )
+        self.assertExpectedInline(
+            dynamo_graph_str,
+            """\
+class GraphModule(torch.nn.Module):
+    def forward(self, L_x_: "f32[3, 3]", L_self_modules_inner_modules_linear_parameters_weight_: "f32[3, 3]", L_self_modules_inner_modules_linear_parameters_bias_: "f32[3]", L_self_modules_linear_parameters_weight_: "f32[3, 3]", L_self_modules_linear_parameters_bias_: "f32[3]"):
+        l_x_ = L_x_
+        l_self_modules_inner_modules_linear_parameters_weight_ = L_self_modules_inner_modules_linear_parameters_weight_
+        l_self_modules_inner_modules_linear_parameters_bias_ = L_self_modules_inner_modules_linear_parameters_bias_
+        l_self_modules_linear_parameters_weight_ = L_self_modules_linear_parameters_weight_
+        l_self_modules_linear_parameters_bias_ = L_self_modules_linear_parameters_bias_
+
+        real_fn : torch._higher_order_ops.invoke_leaf_function._LeafCallable = self.real_fn
+        fake_fn : torch._higher_order_ops.invoke_leaf_function._LeafCallable = self.fake_fn
+        input_spec : torch.utils._pytree.TreeSpec = self.input_spec
+        invoke_leaf_function = torch.ops.higher_order.invoke_leaf_function(real_fn, fake_fn, input_spec, '', 0, l_self_modules_inner_modules_linear_parameters_weight_, l_self_modules_inner_modules_linear_parameters_bias_, l_self_modules_linear_parameters_weight_, l_self_modules_linear_parameters_bias_, l_x_);  real_fn = fake_fn = input_spec = l_self_modules_inner_modules_linear_parameters_weight_ = l_self_modules_inner_modules_linear_parameters_bias_ = l_self_modules_linear_parameters_weight_ = l_self_modules_linear_parameters_bias_ = l_x_ = None
+        getitem: "f32[3, 3]" = invoke_leaf_function[0];  invoke_leaf_function = None
+        return (getitem,)
+""",  # noqa: B950
+        )
+        self.assertExpectedInline(
+            fw_graph_str,
+            """\
+class GraphModule(torch.nn.Module):
+    def forward(self, primals_1: "f32[0]", primals_2: "f32[3, 3]", primals_3: "f32[3, 3]", primals_4: "f32[3]", primals_5: "f32[3, 3]", primals_6: "f32[3]"):
+        _opaque_obj0 = self._opaque_obj0
+        _opaque_obj1 = self._opaque_obj1
+        _tree_spec_constant0 = self._tree_spec_constant0
+        with_effects = torch.ops.higher_order.with_effects(primals_1, torch.ops.higher_order.invoke_leaf_function, _opaque_obj0, _opaque_obj1, _tree_spec_constant0, '', 0, primals_3, primals_4, primals_5, primals_6, primals_2, requires_grad_indices = '1,2,3,4,5');  primals_1 = _opaque_obj0 = _opaque_obj1 = _tree_spec_constant0 = primals_3 = primals_4 = primals_5 = primals_6 = primals_2 = None
+
+        getitem: "f32[0]" = with_effects[0]
+        getitem_1: "f32[3, 3]" = with_effects[1];  with_effects = None
+        return (getitem, getitem_1)
+""",  # noqa: B950
+        )
+        self.assertExpectedInline(
+            bw_graph_str,
+            """\
+class GraphModule(torch.nn.Module):
+    def forward(self, tangents_1: "f32[3, 3]", tangents_token: "f32[0]"):
+        _opaque_obj2 = self._opaque_obj2
+        _opaque_obj3 = self._opaque_obj3
+        _tree_spec_constant1 = self._tree_spec_constant1
+        with_effects_1 = torch.ops.higher_order.with_effects(tangents_token, torch.ops.higher_order.invoke_leaf_function, _opaque_obj2, _opaque_obj3, _tree_spec_constant1, '', tangents_1, requires_grad_indices = '');  tangents_token = _opaque_obj2 = _opaque_obj3 = _tree_spec_constant1 = tangents_1 = None
+        getitem_2: "f32[0]" = with_effects_1[0]
+        getitem_4: "f32[3, 3]" = with_effects_1[2]
+        getitem_5: "f32[3]" = with_effects_1[3]
+        getitem_6: "f32[3, 3]" = with_effects_1[4]
+        getitem_7: "f32[3]" = with_effects_1[5]
+        getitem_8: "f32[3, 3]" = with_effects_1[6];  with_effects_1 = None
+        return (getitem_8, getitem_4, getitem_5, getitem_6, getitem_7, getitem_2)
+""",  # noqa: B950
+        )
+
+    def test_leaf_function_data_dependent_nonzero(self):
+        @leaf_function
+        def nonzero_forward(mod, x):
+            out = mod.linear(x)
+            nonzero_indices = (out > 0).nonzero()
+            return (out, nonzero_indices)
+
+        @nonzero_forward.register_fake
+        def nonzero_forward_fake(mod, x):
+            out = mod.linear(x)
+            return out, (out > 0).nonzero()
+
+        class NonzeroModule(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.linear = torch.nn.Linear(3, 3)
+
+            def forward(self, x):
+                return nonzero_forward(self, x)
+
+        class OuterModule(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.pre_linear = torch.nn.Linear(3, 3)
+                self.nonzero_module = NonzeroModule()
+                self.scale = torch.nn.Parameter(torch.tensor(2.0))
+
+            def forward(self, x):
+                x = self.pre_linear(x)
+                x = torch.relu(x)
+                out, nonzero_indices = self.nonzero_module(x)
+                num_nonzero = nonzero_indices.shape[0]
+                scaled_out = out * self.scale + num_nonzero
+                return scaled_out, nonzero_indices
+
+        def args_fn():
+            return (torch.randn(3, 3, requires_grad=True),)
+
+        def loss_fn(out):
+            return out[0].sum()
+
+        self._test_leaf_function_helper(OuterModule, args_fn, loss_fn)
+
+    @skipIfCrossRef
+    def test_leaf_function_data_dependent_item(self):
+        @leaf_function
+        def item_forward(mod, x):
+            out = mod.linear(x)
+            scalar_value = out.sum().item()
+            return (out, scalar_value)
+
+        @item_forward.register_fake
+        def item_forward_fake(mod, x):
+            out = mod.linear(x)
+            return (out, out.sum().item())
+
+        class ItemModule(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.linear = torch.nn.Linear(3, 3)
+
+            def forward(self, x):
+                return item_forward(self, x)
+
+        def args_fn():
+            return (torch.randn(3, 3, requires_grad=True),)
+
+        def loss_fn(out):
+            return out[0].sum()
+
+        self._test_leaf_function_helper(ItemModule, args_fn, loss_fn)
+
+    @parametrize("backend", ["eager", "aot_eager"])
+    def test_leaf_function_multiple_compiled_submodules(self, backend):
+        @leaf_function
+        def leaf_forward(mod, x):
+            if x.sum() > 0:
+                return (mod.linear(x),)
+            else:
+                return (mod.linear(x) + x,)
+
+        @leaf_forward.register_fake
+        def leaf_forward_fake(mod, x):
+            return (mod.linear(x),)
+
+        class LeafModule(torch.nn.Module):
+            def __init__(self, in_features, out_features):
+                super().__init__()
+                self.linear = torch.nn.Linear(in_features, out_features)
+
+            def forward(self, x):
+                return leaf_forward(self, x)
+
+        class CompiledSubmodule1(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.pre_linear = torch.nn.Linear(4, 4)
+                self.leaf = LeafModule(4, 4)
+
+            def forward(self, x):
+                x = self.pre_linear(x)
+                x = torch.relu(x)
+                out = self.leaf(x)[0]
+                return out
+
+        class CompiledSubmodule2(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.leaf = LeafModule(4, 4)
+                self.post_linear = torch.nn.Linear(4, 4)
+
+            def forward(self, x):
+                out = self.leaf(x)[0]
+                out = self.post_linear(out)
+                return torch.sigmoid(out)
+
+        class CompiledSubmodule3(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.leaf1 = LeafModule(4, 4)
+                self.leaf2 = LeafModule(4, 4)
+
+            def forward(self, x):
+                out1 = self.leaf1(x)[0]
+                out2 = self.leaf2(x)[0]
+                return out1 + out2
+
+        class TopLevelModule(torch.nn.Module):
+            def __init__(self, compile_submodules=False):
+                super().__init__()
+                self.submodule1 = CompiledSubmodule1()
+                self.submodule2 = CompiledSubmodule2()
+                self.submodule3 = CompiledSubmodule3()
+                self.final_linear = torch.nn.Linear(4, 4)
+                self.compile_submodules = compile_submodules
+
+            def forward(self, x):
+                if self.compile_submodules:
+                    out1 = torch.compile(self.submodule1, backend=backend)(x)
+                    out2 = torch.compile(self.submodule2, backend=backend)(out1)
+                    out3 = torch.compile(self.submodule3, backend=backend)(out2)
+                else:
+                    out1 = self.submodule1(x)
+                    out2 = self.submodule2(out1)
+                    out3 = self.submodule3(out2)
+                final = self.final_linear(out3)
+                return final
+
+        model_eager = TopLevelModule(compile_submodules=False)
+        model_compiled = TopLevelModule(compile_submodules=True)
+        model_compiled.load_state_dict(model_eager.state_dict())
+
+        x = torch.randn(2, 4, requires_grad=True)
+        x_compiled = x.clone().detach().requires_grad_(True)
+
+        self._assert_models_equal(
+            model_eager,
+            model_compiled,
+            x,
+            x_compiled,
+        )
+
+    @parametrize("backend", ["eager", "aot_eager"])
+    @parametrize("do_compile", [False, True])
+    def test_leaf_function_with_graph_breaks(self, backend, do_compile):
+        @leaf_function
+        def leaf_forward(mod, x):
+            if x.sum() > 0:
+                return (mod.linear(x),)
+            else:
+                return (mod.linear(x) + 1,)
+
+        @leaf_forward.register_fake
+        def leaf_forward_fake(mod, x):
+            return (mod.linear(x),)
+
+        class LeafModule(torch.nn.Module):
+            def __init__(self, in_features, out_features):
+                super().__init__()
+                self.linear = torch.nn.Linear(in_features, out_features)
+
+            def forward(self, x):
+                return leaf_forward(self, x)
+
+        class TopLevelModule(torch.nn.Module):
+            def __init__(self, do_compile=False, backend="eager"):
+                super().__init__()
+                self.leaf1 = LeafModule(4, 4)
+                self.leaf2 = LeafModule(4, 4)
+                self.leaf3 = LeafModule(4, 4)
+                self.final_linear = torch.nn.Linear(4, 4)
+                self.do_compile = do_compile
+                self.backend = backend
+
+            def _forward(self, x):
+                out1 = self.leaf1(x)[0]
+                torch._dynamo.graph_break()
+                out2 = self.leaf2(out1)[0]
+                torch._dynamo.graph_break()
+                out3 = self.leaf3(out2)[0]
+                result = self.final_linear(out3)
+                return result
+
+            def forward(self, x):
+                if self.do_compile:
+                    return torch.compile(
+                        self._forward, backend=self.backend, fullgraph=False
+                    )(x)
+                else:
+                    return self._forward(x)
+
+        model_eager = TopLevelModule(do_compile=False)
+        model_test = TopLevelModule(do_compile=do_compile, backend=backend)
+        model_test.load_state_dict(model_eager.state_dict())
+
+        x = torch.randn(2, 4, requires_grad=True)
+        x_test = x.clone().detach().requires_grad_(True)
+
+        self._assert_models_equal(model_eager, model_test, x, x_test)
+
+    def test_leaf_function_with_module_in_pytree(self):
+        @leaf_function
+        def main_forward(modules_dict, x):
+            if x.sum() > 0:
+                return (modules_dict["first"](x) + modules_dict["second"](x),)
+            else:
+                return (modules_dict["first"](x) - modules_dict["second"](x),)
+
+        @main_forward.register_fake
+        def main_forward_fake(modules_dict, x):
+            return (modules_dict["first"](x) + modules_dict["second"](x),)
+
+        class HelperModule(torch.nn.Module):
+            def __init__(self, scale=1.0):
+                super().__init__()
+                self.linear = torch.nn.Linear(3, 3)
+                self.scale = scale
+
+            def forward(self, x):
+                return self.linear(x) * self.scale
+
+        class WrapperModule(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.helper1 = HelperModule(scale=1.0)
+                self.helper2 = HelperModule(scale=0.5)
+
+            def forward(self, x):
+                modules_dict = {"first": self.helper1, "second": self.helper2}
+                return main_forward(modules_dict, x)
+
+        def args_fn():
+            return (torch.randn(3, 3, requires_grad=True),)
+
+        def loss_fn(out):
+            return out[0].sum()
+
+        self._test_leaf_function_helper(WrapperModule, args_fn, loss_fn)
+
+    def test_leaf_function_with_module_as_kwarg(self):
+        @leaf_function
+        def main_forward(x, helper_mod=None):
+            if x.sum() > 0:
+                return (helper_mod(x),)
+            else:
+                return (helper_mod(x) + x,)
+
+        @main_forward.register_fake
+        def main_forward_fake(x, helper_mod=None):
+            return (helper_mod(x),)
+
+        class HelperModule(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.linear = torch.nn.Linear(3, 3)
+
+            def forward(self, x):
+                return self.linear(x)
+
+        class WrapperModule(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.helper = HelperModule()
+
+            def forward(self, x):
+                return main_forward(x, helper_mod=self.helper)
+
+        def args_fn():
+            return (torch.randn(3, 3, requires_grad=True),)
+
+        def loss_fn(out):
+            return out[0].sum()
+
+        self._test_leaf_function_helper(WrapperModule, args_fn, loss_fn)
+
+    def test_leaf_function_missing_fake_impl_error(self):
+        @leaf_function
+        def no_fake_impl_forward(mod, x):
+            return (mod.linear(x),)
+
+        class SimpleModule(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.linear = torch.nn.Linear(3, 3)
+
+            def forward(self, x):
+                return no_fake_impl_forward(self, x)
+
+        mod = SimpleModule()
+        x = torch.randn(3, 3)
+
+        with self.assertRaisesRegex(Exception, "requires a fake implementation"):
+            mod(x)
+
+        compiled_mod = torch.compile(mod, backend="eager", fullgraph=True)
+        with self.assertRaisesRegex(Exception, "requires a fake implementation"):
+            compiled_mod(x)
+
+    @parametrize("backend", ["eager", "aot_eager"])
+    def test_leaf_function_constant_tensor_closure_error(self, backend):
+        constant_weight = torch.randn(3, 3)
+
+        @leaf_function
+        def constant_closure_forward(x):
+            return (x @ constant_weight,)
+
+        @constant_closure_forward.register_fake
+        def constant_closure_forward_fake(x):
+            return (x @ constant_weight,)
+
+        class ConstantClosureModule(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+
+            def forward(self, x):
+                return constant_closure_forward(x)
+
+        mod = ConstantClosureModule()
+        x = torch.randn(3, 3, requires_grad=True)
+
+        result = mod(x)
+        expected = x @ constant_weight
+        self.assertEqual(result[0], expected)
+
+        compiled_mod = torch.compile(mod, backend=backend, fullgraph=True)
+        with self.assertRaisesRegex(
+            Exception, "Please convert all Tensors to FakeTensors"
+        ):
+            compiled_mod(x)
+
+    @parametrize("backend", ["eager", "aot_eager"])
+    def test_leaf_function_input_mutation_error(self, backend):
+        @leaf_function
+        def mutate_input(x):
+            x.add_(1)
+            return (x,)
+
+        @mutate_input.register_fake
+        def mutate_input_fake(x):
+            x.add_(1)
+            return (x,)
+
+        def fn(x):
+            return mutate_input(x)
+
+        x = torch.randn(3, 3)
+
+        x_eager = x.clone()
+        with self.assertRaisesRegex(RuntimeError, "Undeclared in-place mutation"):
+            fn(x_eager)
+
+        x = torch.randn(3, 3, requires_grad=True)
+
+        compiled_fn = torch.compile(fn, backend=backend, fullgraph=True)
+        with self.assertRaisesRegex(RuntimeError, "leaf Variable that requires grad"):
+            compiled_fn(x.clone().requires_grad_(True))
+
+    @parametrize("backend", ["eager", "aot_eager"])
+    def test_leaf_function_validation_dtype_mismatch(self, backend):
+        @leaf_function
+        def dtype_mismatch_forward(mod, x):
+            return (mod.linear(x),)
+
+        @dtype_mismatch_forward.register_fake
+        def dtype_mismatch_forward_fake(mod, x):
+            return (mod.linear(x).double(),)
+
+        class DtypeMismatchModule(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.linear = torch.nn.Linear(3, 3)
+
+            def forward(self, x):
+                return dtype_mismatch_forward(self, x)
+
+        mod = DtypeMismatchModule()
+        x = torch.randn(3, 3)
+
+        with config.patch(leaf_function_validate_outputs=True):
+            compiled_mod = torch.compile(mod, backend=backend)
+            with self.assertRaisesRegex(RuntimeError, "Dtype mismatch"):
+                compiled_mod(x)
+
+    @parametrize("backend", ["eager", "aot_eager"])
+    @parametrize("validate_outputs", [True, False])
+    def test_leaf_function_validation_shape_mismatch(self, backend, validate_outputs):
+        @leaf_function
+        def mismatched_forward(mod, x):
+            return (mod.linear(x),)
+
+        @mismatched_forward.register_fake
+        def mismatched_forward_fake(mod, x):
+            return (torch.zeros(x.shape[0], 6),)
+
+        class MismatchedModule(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.linear = torch.nn.Linear(3, 3)
+
+            def forward(self, x):
+                return mismatched_forward(self, x)
+
+        mod = MismatchedModule()
+        x = torch.randn(3, 3)
+
+        with config.patch(leaf_function_validate_outputs=validate_outputs):
+            compiled_mod = torch.compile(mod, backend=backend)
+            if validate_outputs:
+                with self.assertRaises((RuntimeError, AssertionError)):
+                    compiled_mod(x)
+            else:
+                result = compiled_mod(x)
+                self.assertEqual(result[0].shape, (3, 3))
+
+    def test_leaf_function_no_module_inputs(self):
+        @leaf_function
+        def my_custom_fn(inputs: dict[str, torch.Tensor], scale: float, offset: int):
+            x = inputs["x"]
+            y = inputs["y"]
+            if x.sum() > 0:
+                return (x * scale + y + offset, x.sum() + y.sum())
+            return (x * scale - y + offset, x.sum() - y.sum())
+
+        @my_custom_fn.register_fake
+        def my_custom_fn_fake(
+            inputs: dict[str, torch.Tensor], scale: float, offset: int
+        ):
+            x = inputs["x"]
+            y = inputs["y"]
+            return (x * scale + y + offset, x.sum() + y.sum())
+
+        class NoModuleInputsModule(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.scale = 2.0
+                self.offset = 1
+
+            def forward(self, x, y):
+                inputs = {"x": x, "y": y}
+                return my_custom_fn(inputs, self.scale, self.offset)
+
+        def args_fn():
+            return (
+                torch.randn(3, 3, requires_grad=True),
+                torch.randn(3, 3, requires_grad=True),
+            )
+
+        def loss_fn(out):
+            return out[0].sum() + out[1].sum()
+
+        self._test_leaf_function_helper(NoModuleInputsModule, args_fn, loss_fn)
+
+    @parametrize("backend", ["eager", "aot_eager"])
+    @parametrize("check_escaped_gradients", [True, False])
+    def test_leaf_function_escaped_gradient_multiple_tensors(
+        self, backend, check_escaped_gradients
+    ):
+        weight1 = torch.randn(3, 3, requires_grad=True)
+        weight2 = torch.randn(3, 3, requires_grad=True)
+
+        @leaf_function
+        def uses_multiple_closures(x):
+            return (x @ weight1 + x @ weight2,)
+
+        @uses_multiple_closures.register_fake
+        def uses_multiple_closures_fake(x):
+            return (torch.empty(x.shape[0], 3),)
+
+        def fn(x):
+            return uses_multiple_closures(x)
+
+        x = torch.randn(2, 3, requires_grad=True)
+
+        compiled_fn = torch.compile(fn, backend=backend, fullgraph=True)
+        with config.patch(
+            leaf_function_check_escaped_gradients=check_escaped_gradients
+        ):
+            if check_escaped_gradients:
+                with self.assertRaisesRegex(RuntimeError, "2 tensor"):
+                    compiled_fn(x)
+            else:
+                result = compiled_fn(x)
+                self.assertEqual(result[0].shape, (2, 3))
+
+    @parametrize("backend", ["eager", "aot_eager"])
+    @parametrize("check_escaped_gradients", [True, False])
+    def test_leaf_function_escaped_gradient_input_no_grad(
+        self, backend, check_escaped_gradients
+    ):
+        closure_weight = torch.randn(3, 3, requires_grad=True)
+
+        @leaf_function
+        def uses_closure(x):
+            return (x @ closure_weight,)
+
+        @uses_closure.register_fake
+        def uses_closure_fake(x):
+            return (torch.empty(x.shape[0], 3),)
+
+        def fn(x):
+            return uses_closure(x)
+
+        x = torch.randn(2, 3, requires_grad=False)
+
+        compiled_fn = torch.compile(fn, backend=backend, fullgraph=True)
+        with config.patch(
+            leaf_function_check_escaped_gradients=check_escaped_gradients
+        ):
+            result = compiled_fn(x)
+            self.assertEqual(result[0].shape, (2, 3))
+
+    @parametrize("backend", ["eager", "aot_eager"])
+    @parametrize("check_escaped_gradients", [True, False])
+    def test_leaf_function_escaped_gradient_mixed_inputs(
+        self, backend, check_escaped_gradients
+    ):
+        base1 = torch.randn(3, 3, requires_grad=True)
+        base2 = torch.randn(3, 4, requires_grad=True)
+        closure_weight1 = base1 * 2
+        closure_weight2 = base2 * 3
+
+        @leaf_function
+        def mixed_inputs(x, y):
+            out1 = x @ closure_weight1 + y
+            out2 = x @ closure_weight2
+            return (out1, out2)
+
+        @mixed_inputs.register_fake
+        def mixed_inputs_fake(x, y):
+            return (torch.empty(x.shape[0], 3), torch.empty(x.shape[0], 4))
+
+        def fn(x, y):
+            return mixed_inputs(x, y)
+
+        x = torch.randn(2, 3, requires_grad=True)
+        y = torch.randn(2, 3, requires_grad=False)
+
+        compiled_fn = torch.compile(fn, backend=backend, fullgraph=True)
+        with config.patch(
+            leaf_function_check_escaped_gradients=check_escaped_gradients
+        ):
+            if check_escaped_gradients:
+                with self.assertRaisesRegex(RuntimeError, "2 tensor"):
+                    compiled_fn(x, y)
+            else:
+                result = compiled_fn(x, y)
+                self.assertEqual(result[0].shape, (2, 3))
+                self.assertEqual(result[1].shape, (2, 4))
+
+    @parametrize("backend", ["eager", "aot_eager"])
+    def test_leaf_function_escaped_gradient_error_message_contains_tensor_info(
+        self, backend
+    ):
+        closure_weight = torch.randn(4, 5, dtype=torch.float32, requires_grad=True)
+
+        @leaf_function
+        def uses_closure(x):
+            return (x @ closure_weight,)
+
+        @uses_closure.register_fake
+        def uses_closure_fake(x):
+            return (torch.empty(x.shape[0], 5),)
+
+        def fn(x):
+            return uses_closure(x)
+
+        x = torch.randn(2, 4, requires_grad=True)
+
+        compiled_fn = torch.compile(fn, backend=backend, fullgraph=True)
+        with config.patch(leaf_function_check_escaped_gradients=True):
+            with self.assertRaisesRegex(RuntimeError, r"shape=\[4, 5\].*dtype="):
+                compiled_fn(x)
+
+    @parametrize("backend", ["eager", "aot_eager"])
+    def test_leaf_function_escaped_gradient_actually_lost(self, backend):
+        closure_weight = torch.randn(3, 3, requires_grad=True)
+
+        @leaf_function
+        def uses_closure(x):
+            return (x @ closure_weight,)
+
+        @uses_closure.register_fake
+        def uses_closure_fake(x):
+            return (torch.empty(x.shape[0], 3),)
+
+        def fn(x):
+            return uses_closure(x)
+
+        x = torch.randn(2, 3, requires_grad=True)
+
+        compiled_fn = torch.compile(fn, backend=backend, fullgraph=True)
+        result = compiled_fn(x)
+        loss = result[0].sum()
+        loss.backward()
+
+        self.assertIsNotNone(x.grad)
+        self.assertIsNone(closure_weight.grad)
+
+    def test_leaf_function_and_nonstrict_trace_mutually_exclusive(self):
+        from torch._dynamo.decorators import leaf_function, nonstrict_trace
+
+        with self.assertRaisesRegex(
+            ValueError,
+            "cannot be both marked as @leaf_function and @nonstrict_trace",
+        ):
+
+            @leaf_function
+            @nonstrict_trace
+            def bad_fn1(x):
+                return (x,)
+
+        with self.assertRaisesRegex(
+            ValueError,
+            "cannot be both marked as @leaf_function and @nonstrict_trace",
+        ):
+
+            @nonstrict_trace
+            @leaf_function
+            def bad_fn2(x):
+                return (x,)
+
+    @skipIfCrossRef
+    def test_leaf_function_no_return_value(self):
+        printed = []
+
+        @leaf_function
+        def fn_no_return(x):
+            print("processing")
+
+        @fn_no_return.register_fake
+        def fn_no_return_fake(x):
+            pass
+
+        class Mod(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.linear = torch.nn.Linear(3, 3)
+
+            def forward(self, x):
+                fn_no_return(x)
+                return (self.linear(x),)
+
+        def args_fn():
+            return (torch.randn(3, 3, requires_grad=True),)
+
+        def loss_fn(out):
+            return out[0].sum()
+
+        with patch("builtins.print", lambda *args, **kwargs: printed.append(args)):
+            eager_graph, fw_graph, bw_graph = self._test_leaf_function_helper(
+                Mod, args_fn, loss_fn
+            )
+        self.assertTrue(any("processing" in p for p in printed))
+        self.assertExpectedInline(
+            eager_graph,
+            """\
+class GraphModule(torch.nn.Module):
+    def forward(self, L_x_: "f32[3, 3]", L_self_modules_linear_parameters_weight_: "f32[3, 3]", L_self_modules_linear_parameters_bias_: "f32[3]"):
+        l_x_ = L_x_
+        l_self_modules_linear_parameters_weight_ = L_self_modules_linear_parameters_weight_
+        l_self_modules_linear_parameters_bias_ = L_self_modules_linear_parameters_bias_
+
+        real_fn : torch._higher_order_ops.invoke_leaf_function._LeafCallable = self.real_fn
+        fake_fn : torch._higher_order_ops.invoke_leaf_function._LeafCallable = self.fake_fn
+        input_spec : torch.utils._pytree.TreeSpec = self.input_spec
+        invoke_leaf_function = torch.ops.higher_order.invoke_leaf_function(real_fn, fake_fn, input_spec, '', l_x_);  real_fn = fake_fn = input_spec = invoke_leaf_function = None
+
+        linear: "f32[3, 3]" = torch._C._nn.linear(l_x_, l_self_modules_linear_parameters_weight_, l_self_modules_linear_parameters_bias_);  l_x_ = l_self_modules_linear_parameters_weight_ = l_self_modules_linear_parameters_bias_ = None
+        return (linear,)
+""",  # noqa: B950
+        )
+        self.assertExpectedInline(
+            fw_graph,
+            """\
+class GraphModule(torch.nn.Module):
+    def forward(self, primals_1: "f32[0]", primals_2: "f32[3, 3]", primals_3: "f32[3, 3]", primals_4: "f32[3]"):
+        _opaque_obj0 = self._opaque_obj0
+        _opaque_obj1 = self._opaque_obj1
+        _tree_spec_constant0 = self._tree_spec_constant0
+        with_effects = torch.ops.higher_order.with_effects(primals_1, torch.ops.higher_order.invoke_leaf_function, _opaque_obj0, _opaque_obj1, _tree_spec_constant0, '', primals_2, requires_grad_indices = '0');  primals_1 = _opaque_obj0 = _opaque_obj1 = _tree_spec_constant0 = None
+
+        getitem: "f32[0]" = with_effects[0];  with_effects = None
+
+        t: "f32[3, 3]" = torch.ops.aten.t.default(primals_3)
+        addmm: "f32[3, 3]" = torch.ops.aten.addmm.default(primals_4, primals_2, t);  primals_4 = t = None
+        return (getitem, addmm, primals_2, primals_3)
+""",  # noqa: B950
+        )
+        self.assertExpectedInline(
+            bw_graph,
+            """\
+class GraphModule(torch.nn.Module):
+    def forward(self, primals_2: "f32[3, 3]", primals_3: "f32[3, 3]", tangents_1: "f32[3, 3]"):
+        t: "f32[3, 3]" = torch.ops.aten.t.default(primals_3);  primals_3 = None
+        t_1: "f32[3, 3]" = torch.ops.aten.t.default(t);  t = None
+        mm: "f32[3, 3]" = torch.ops.aten.mm.default(tangents_1, t_1);  t_1 = None
+        t_2: "f32[3, 3]" = torch.ops.aten.t.default(tangents_1)
+        mm_1: "f32[3, 3]" = torch.ops.aten.mm.default(t_2, primals_2);  t_2 = primals_2 = None
+        t_3: "f32[3, 3]" = torch.ops.aten.t.default(mm_1);  mm_1 = None
+        sum_1: "f32[1, 3]" = torch.ops.aten.sum.dim_IntList(tangents_1, [0], True);  tangents_1 = None
+        view: "f32[3]" = torch.ops.aten.view.default(sum_1, [3]);  sum_1 = None
+        t_4: "f32[3, 3]" = torch.ops.aten.t.default(t_3);  t_3 = None
+        return (mm, t_4, view)
+""",  # noqa: B950
+        )
+
+    def test_leaf_function_output_structure_mismatch(self):
+        @leaf_function
+        def mismatched_fn(x):
+            return {"a": x, "b": x * 2}
+
+        @mismatched_fn.register_fake
+        def mismatched_fn_fake(x):
+            return (x, x * 2)
+
+        def fn(x):
+            return mismatched_fn(x)
+
+        x = torch.randn(3, 3)
+        with self.assertRaisesRegex(AssertionError, "output structure mismatch"):
+            torch.compile(fn, backend="eager")(x)
+
+    def test_leaf_function_nested_output(self):
+        @leaf_function
+        def nested_output_fn(linear1, linear2, linear3, x):
+            if x.sum() > 0:
+                return {
+                    "out": (linear1(x), linear2(x)),
+                    "extra": linear3(x),
+                    "count": 42,
+                }
+            else:
+                return {
+                    "out": (linear1(x) + 1, linear2(x) + 1),
+                    "extra": linear3(x) + 1,
+                    "count": 42,
+                }
+
+        @nested_output_fn.register_fake
+        def nested_output_fn_fake(linear1, linear2, linear3, x):
+            return {
+                "out": (linear1(x), linear2(x)),
+                "extra": linear3(x),
+                "count": 42,
+            }
+
+        class NestedOutputModule(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.linear1 = torch.nn.Linear(3, 3)
+                self.linear2 = torch.nn.Linear(3, 3)
+                self.linear3 = torch.nn.Linear(3, 3)
+
+            def forward(self, x):
+                result = nested_output_fn(self.linear1, self.linear2, self.linear3, x)
+                return (
+                    result["out"][0] * result["count"]
+                    + result["out"][1]
+                    + result["extra"]
+                )
+
+        def args_fn():
+            return (torch.randn(3, 3, requires_grad=True),)
+
+        def loss_fn(out):
+            return out.sum()
+
+        self._test_leaf_function_helper(NestedOutputModule, args_fn, loss_fn)
+
+    def test_leaf_function_custom_pytree_output(self):
+        class Point:
+            x: torch.Tensor
+            y: torch.Tensor
+
+            def __init__(self, x, y):
+                self.x = x
+                self.y = y
+
+        self.register_pytree_node(
+            Point,
+            lambda p: ((p.x, p.y), ()),
+            lambda xy, _: Point(xy[0], xy[1]),
+            serialized_type_name=f"{Point.__module__}.{Point.__qualname__}",
+        )
+
+        @leaf_function
+        def point_fn(linear1, linear2, x):
+            return (Point(linear1(x), linear2(x)), 0.5)
+
+        @point_fn.register_fake
+        def point_fn_fake(linear1, linear2, x):
+            return (Point(linear1(x), linear2(x)), 0.5)
+
+        class PointModule(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.linear1 = torch.nn.Linear(3, 3)
+                self.linear2 = torch.nn.Linear(3, 3)
+
+            def forward(self, x):
+                p, scale = point_fn(self.linear1, self.linear2, x)
+                return (p.x * scale, p.y * scale)
+
+        def args_fn():
+            return (torch.randn(3, 3, requires_grad=True),)
+
+        def loss_fn(out):
+            return out[0].sum() + out[1].sum()
+
+        self._test_leaf_function_helper(PointModule, args_fn, loss_fn)
+
+    def test_leaf_function_fake_requires_grad_ignored(self):
+        @leaf_function
+        def my_fn(x):
+            return (x * 2,)
+
+        @my_fn.register_fake
+        def my_fn_fake(x):
+            return (torch.empty_like(x).requires_grad_(False),)
+
+        from torch._dynamo.testing import EagerAndRecordGraphs
+
+        backend = EagerAndRecordGraphs()
+
+        @torch.compile(backend=backend, fullgraph=True)
+        def fn(x):
+            return my_fn(x)
+
+        x = torch.randn(3, 3, requires_grad=True)
+        out = fn(x)
+
+        self.assertTrue(out[0].requires_grad)
+        out[0].sum().backward()
+        self.assertIsNotNone(x.grad)
+
+        graph = backend.graphs[0]
+        for node in graph.graph.nodes:
+            if node.op == "call_function" and "invoke_leaf_function" in str(
+                node.target
+            ):
+                example_value = node.meta.get("example_value")
+                self.assertIsNotNone(example_value)
+                self.assertTrue(example_value[0].requires_grad)
+
+    @parametrize("backend", ["eager", "aot_eager"])
+    def test_leaf_function_input_mutation_non_grad(self, backend):
+        @leaf_function(mutates_args={"buf"})
+        def mutate_buffer(x, buf):
+            buf.add_(1)
+            return (x + buf,)
+
+        @mutate_buffer.register_fake
+        def mutate_buffer_fake(x, buf):
+            buf.add_(1)
+            return (x + buf,)
+
+        def fn(x, buf):
+            return mutate_buffer(x, buf)
+
+        x = torch.randn(3, 3)
+        buf = torch.randn(3, 3)
+
+        buf_eager = buf.clone()
+        result_eager = fn(x, buf_eager)
+        expected = x + buf + 1
+        self.assertEqual(result_eager[0], expected)
+        self.assertEqual(buf_eager, buf + 1)
+
+        compiled_fn = torch.compile(fn, backend=backend, fullgraph=True)
+        buf_compiled = buf.clone()
+        result_compiled = compiled_fn(x, buf_compiled)
+        self.assertEqual(result_compiled[0], expected)
+        self.assertEqual(buf_compiled, buf + 1)
+
+    @parametrize("backend", ["eager", "aot_eager"])
+    def test_leaf_function_input_mutation_mixed(self, backend):
+        @leaf_function(mutates_args={"buf"})
+        def mixed_fn(x, buf):
+            buf.mul_(2)
+            return (x * buf,)
+
+        @mixed_fn.register_fake
+        def mixed_fn_fake(x, buf):
+            buf.mul_(2)
+            return (x * buf,)
+
+        def fn(x, buf):
+            return mixed_fn(x, buf)
+
+        x = torch.randn(3, 3, requires_grad=True)
+        buf = torch.randn(3, 3)
+
+        buf_eager = buf.clone()
+        result_eager = fn(x, buf_eager)
+        expected = x * (buf * 2)
+        self.assertEqual(result_eager[0], expected)
+        self.assertEqual(buf_eager, buf * 2)
+
+        compiled_fn = torch.compile(fn, backend=backend, fullgraph=True)
+        buf_compiled = buf.clone()
+        result_compiled = compiled_fn(x, buf_compiled)
+        self.assertEqual(result_compiled[0], expected)
+        self.assertEqual(buf_compiled, buf * 2)
+
+    @parametrize("backend", ["eager", "aot_eager"])
+    def test_leaf_function_input_mutation_module_buffer(self, backend):
+        class MyModule(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.register_buffer("running_mean", torch.zeros(3))
+                self.linear = torch.nn.Linear(3, 3)
+
+            def forward(self, x):
+                return update_stats(self, x)
+
+        @leaf_function(mutates_args={"model.running_mean"})
+        def update_stats(model, x):
+            model.running_mean.add_(x.mean(dim=0))
+            return (model.linear(x),)
+
+        @update_stats.register_fake
+        def update_stats_fake(model, x):
+            model.running_mean.add_(x.mean(dim=0))
+            return (model.linear(x),)
+
+        mod = MyModule()
+        x = torch.randn(4, 3)
+
+        mod_eager = copy.deepcopy(mod)
+        result_eager = mod_eager(x)
+        expected_mean = torch.zeros(3) + x.mean(dim=0)
+        self.assertEqual(mod_eager.running_mean, expected_mean)
+
+        mod_compiled = copy.deepcopy(mod)
+        compiled_mod = torch.compile(mod_compiled, backend=backend, fullgraph=True)
+        result_compiled = compiled_mod(x)
+        self.assertEqual(result_compiled, result_eager)
+        self.assertEqual(mod_compiled.running_mean, expected_mean)
+
+    @parametrize("backend", ["eager", "aot_eager"])
+    def test_leaf_function_input_mutation_pytree(self, backend):
+        @leaf_function(mutates_args={"buffers"})
+        def update_buffers(x, buffers):
+            for buf in buffers:
+                buf.add_(1)
+            return (x + sum(buffers),)
+
+        @update_buffers.register_fake
+        def update_buffers_fake(x, buffers):
+            for buf in buffers:
+                buf.add_(1)
+            return (x + sum(buffers),)
+
+        def fn(x, buffers):
+            return update_buffers(x, buffers)
+
+        x = torch.randn(3, 3)
+        bufs = [torch.randn(3, 3), torch.randn(3, 3)]
+
+        bufs_eager = [b.clone() for b in bufs]
+        result_eager = fn(x, bufs_eager)
+        expected = x + (bufs[0] + 1) + (bufs[1] + 1)
+        self.assertEqual(result_eager[0], expected)
+        self.assertEqual(bufs_eager[0], bufs[0] + 1)
+        self.assertEqual(bufs_eager[1], bufs[1] + 1)
+
+        compiled_fn = torch.compile(fn, backend=backend, fullgraph=True)
+        bufs_compiled = [b.clone() for b in bufs]
+        result_compiled = compiled_fn(x, bufs_compiled)
+        self.assertEqual(result_compiled[0], expected)
+        self.assertEqual(bufs_compiled[0], bufs[0] + 1)
+        self.assertEqual(bufs_compiled[1], bufs[1] + 1)
+
+    @parametrize("backend", ["eager", "aot_eager"])
+    def test_leaf_function_input_mutation_pytree_fine_grained(self, backend):
+        @leaf_function(mutates_args={"buffers[0]"})
+        def update_first(x, buffers):
+            buffers[0].add_(1)
+            return (x + buffers[0] + buffers[1],)
+
+        @update_first.register_fake
+        def update_first_fake(x, buffers):
+            buffers[0].add_(1)
+            return (x + buffers[0] + buffers[1],)
+
+        def fn(x, buffers):
+            return update_first(x, buffers)
+
+        x = torch.randn(3, 3)
+        bufs = [torch.randn(3, 3), torch.randn(3, 3)]
+
+        bufs_eager = [b.clone() for b in bufs]
+        result_eager = fn(x, bufs_eager)
+        expected = x + (bufs[0] + 1) + bufs[1]
+        self.assertEqual(result_eager[0], expected)
+        self.assertEqual(bufs_eager[0], bufs[0] + 1)
+        self.assertEqual(bufs_eager[1], bufs[1])
+
+        compiled_fn = torch.compile(fn, backend=backend, fullgraph=True)
+        bufs_compiled = [b.clone() for b in bufs]
+        result_compiled = compiled_fn(x, bufs_compiled)
+        self.assertEqual(result_compiled[0], expected)
+        self.assertEqual(bufs_compiled[0], bufs[0] + 1)
+        self.assertEqual(bufs_compiled[1], bufs[1])
+
+    def test_leaf_function_mutates_args_invalid_parameter(self):
+        with self.assertRaisesRegex(ValueError, "refers to parameter 'buf'"):
+
+            @leaf_function(mutates_args={"buf"})
+            def bad_fn(x, buffers):
+                buffers.add_(1)
+                return (x + buffers,)
+
+        with self.assertRaisesRegex(ValueError, "refers to parameter 'mdl'"):
+
+            @leaf_function(mutates_args={"mdl.running_mean"})
+            def bad_fn2(x, model):
+                model.running_mean.add_(1)
+                return (x,)
+
+    def test_leaf_function_mutates_args_non_leaf_expression(self):
+        @leaf_function(mutates_args={"model"})
+        def bad_fn(x, model):
+            model.running_mean.add_(1)
+            return (x,)
+
+        @bad_fn.register_fake
+        def bad_fn_fake(x, model):
+            model.running_mean.add_(1)
+            return (x,)
+
+        class MyModule(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.register_buffer("running_mean", torch.zeros(3))
+
+            def forward(self, x):
+                return bad_fn(x, self)
+
+        mod = MyModule()
+        x = torch.randn(3)
+        compiled_fn = torch.compile(mod, backend="eager", fullgraph=True)
+        with self.assertRaisesRegex(
+            torch._dynamo.exc.UserError, "resolved to a non-leaf value"
+        ):
+            compiled_fn(x)
+
+
+instantiate_parametrized_tests(TestLeafFunctionDynamo)
+
+
+if __name__ == "__main__":
+    run_tests()
diff --git a/test/functorch/test_memory_efficient_fusion.py b/test/functorch/test_memory_efficient_fusion.py
index 96e0ccaa6f48a..b8807202ccea6 100644
--- a/test/functorch/test_memory_efficient_fusion.py
+++ b/test/functorch/test_memory_efficient_fusion.py
@@ -208,33 +208,36 @@ def check(f, t, delta, check_val=True, graph_input=False):
     old_num_nodes = len(fx_g.graph.nodes)
     new_num_nodes = len(new_graph.nodes)
     if delta == -1:
-        assert old_num_nodes >= new_num_nodes, (
-            f"number of nodes increased {old_num_nodes}, {new_num_nodes}"
-        )
+        if old_num_nodes < new_num_nodes:
+            raise AssertionError(
+                f"number of nodes increased {old_num_nodes}, {new_num_nodes}"
+            )
     else:
-        assert old_num_nodes == new_num_nodes + delta, (
-            f"number of nodes not the same {old_num_nodes - delta}, {new_num_nodes}\n {fx_g.graph} \n {new_graph}"
-        )
+        if old_num_nodes != new_num_nodes + delta:
+            raise AssertionError(
+                f"number of nodes not the same {old_num_nodes - delta}, {new_num_nodes}\n {fx_g.graph} \n {new_graph}"
+            )
 
     # a second pass should not reduce more nodes
     pass_2_graph = fx_graph_cse(new_graph)
     pass_2_num_nodes = len(pass_2_graph.nodes)
-    assert pass_2_num_nodes == new_num_nodes, (
-        f"second pass graph has less node {pass_2_num_nodes}, {new_num_nodes}\n {new_graph} \n {pass_2_graph}"
-    )
+    if pass_2_num_nodes != new_num_nodes:
+        raise AssertionError(
+            f"second pass graph has less node {pass_2_num_nodes}, {new_num_nodes}\n {new_graph} \n {pass_2_graph}"
+        )
 
     # check correctness
     if check_val:
         true_result = fx_g(t)
         our_result = new_g(t)
         if true_result is None:  # both return None
-            assert our_result is None, (
-                f"true result is None, CSE result is {our_result}"
-            )
+            if our_result is not None:
+                raise AssertionError(f"true result is None, CSE result is {our_result}")
         else:  # results returned are the same
-            assert torch.all(true_result == our_result), (
-                f"results are different {true_result}, {our_result}"
-            )  # check results are the same
+            if not torch.all(true_result == our_result):
+                raise AssertionError(
+                    f"results are different {true_result}, {our_result}"
+                )
 
 
 class NoChangeTestCase(TestCase):
@@ -301,7 +304,8 @@ def toy_backend(gm, _):
 
         t = torch.rand(3, 100)
         _ = fn(t, 50)
-        assert len(gms) == 1, gms
+        if len(gms) != 1:
+            raise AssertionError(f"Expected 1 graph module, got {len(gms)}: {gms}")
         fx_g = gms[0]
         check(fx_g, None, 0, check_val=False, graph_input=True)
 
diff --git a/test/functorch/test_minifier.py b/test/functorch/test_minifier.py
index 014117431ffa6..6ee5001f7c5ff 100644
--- a/test/functorch/test_minifier.py
+++ b/test/functorch/test_minifier.py
@@ -110,8 +110,12 @@ def pass_checker(fx_g, inps):
             return torch.isnan(fx_g(*inps)[0]).any()
 
         min_f, inps = minifier(failing_f, inps, pass_checker)
-        assert len(min_f.graph.nodes) == 3
-        assert len(inps) == 1
+        if len(min_f.graph.nodes) != 3:
+            raise AssertionError(
+                f"Expected 3 graph nodes, got {len(min_f.graph.nodes)}"
+            )
+        if len(inps) != 1:
+            raise AssertionError(f"Expected 1 input, got {len(inps)}")
 
 
 if __name__ == "__main__":
diff --git a/test/functorch/test_ops.py b/test/functorch/test_ops.py
index 861ac00de4b27..632a4a915a685 100644
--- a/test/functorch/test_ops.py
+++ b/test/functorch/test_ops.py
@@ -135,7 +135,8 @@ def normalize_op_input_output2(
         for i, arg in enumerate(flat_args)
         if diff_arg(arg, requires_grad=requires_grad)
     )
-    assert len(diff_argnums) > 0
+    if len(diff_argnums) == 0:
+        raise AssertionError("Expected at least one differentiable argument")
     primals = tuple(flat_args[i] for i in diff_argnums)
 
     @functools.wraps(f)
@@ -149,7 +150,8 @@ def wrapped(*primals):
             result = output_process_fn_grad(result)
         if isinstance(result, tuple):
             result = tuple(r for r in result if torch.is_floating_point(r))
-            assert len(result) > 0
+            if len(result) == 0:
+                raise AssertionError("Expected at least one floating point result")
         return result
 
     return wrapped, primals
@@ -166,7 +168,8 @@ def normalize_op_input_output3(
         for i, (arg, sample) in enumerate(zip(flat_args, flat_sample_args))
         if diff_arg(sample, requires_grad=True)
     )
-    assert len(diff_argnums) > 0
+    if len(diff_argnums) == 0:
+        raise AssertionError("Expected at least one differentiable argument")
     primals = tuple(flat_args[i] for i in diff_argnums)
 
     @functools.wraps(f)
@@ -180,7 +183,8 @@ def wrapped(*primals):
             result = output_process_fn_grad(result)
         if isinstance(result, tuple):
             result = tuple(r for r in result if torch.is_floating_point(r))
-            assert len(result) > 0
+            if len(result) == 0:
+                raise AssertionError("Expected at least one floating point result")
         return result
 
     return wrapped, primals
@@ -236,7 +240,10 @@ def get_vjp_fn_and_args_with_cotangents(f, sample, cotangents):
 
     @functools.wraps(f)
     def wrapped(*args):
-        assert len(args) == len(flat_args) + len(flat_cotangents)
+        if len(args) != len(flat_args) + len(flat_cotangents):
+            raise AssertionError(
+                f"Expected {len(flat_args) + len(flat_cotangents)} args, got {len(args)}"
+            )
         actual_args = args[: len(flat_args)]
         cotangents = args[len(flat_args) :]
         actual_args = tree_unflatten(actual_args, args_spec)
@@ -277,7 +284,8 @@ def wrapped(*args):
         cotangents = args[num_primals:]
         result, vjp_fn = vjp(fn, *primals)
         if isinstance(result, torch.Tensor):
-            assert len(cotangents) == 1
+            if len(cotangents) != 1:
+                raise AssertionError(f"Expected 1 cotangent, got {len(cotangents)}")
             cotangents = cotangents[0]
         return vjp_fn(cotangents)
 
@@ -537,7 +545,8 @@ def test_grad(self, device, dtype, op):
                 noncontig_kwargs = noncontig_sample.kwargs
 
             diff_argnums = tuple(i for i, arg in enumerate(args) if diff_arg(arg))
-            assert len(diff_argnums) > 0
+            if len(diff_argnums) == 0:
+                raise AssertionError("Expected at least one differentiable argument")
             diff_args = tuple(args[i] for i in diff_argnums)
 
             def wrapped_fn(*args, **kwargs):
@@ -972,9 +981,6 @@ def fn(inp, *args, **kwargs):
                     "masked.softmax",
                     device_type="cpu",
                 ),
-                xfail(
-                    "nanquantile", device_type="cpu"
-                ),  # vmap not implemented for at::equal.
                 xfail("native_layer_norm"),  # vmap: inplace into a regular tensor
                 # got a batched tensor as input while the running_mean or running_var,
                 # which will be updated in place, were not batched.
@@ -1026,9 +1032,6 @@ def fn(inp, *args, **kwargs):
                 xfail("normal"),  # calls random op
                 xfail("normal", "number_mean"),  # calls random op
                 xfail("pca_lowrank"),  # calls random op
-                xfail(
-                    "quantile", device_type="cpu"
-                ),  # Batching rule not implemented for `at::equal`
                 xfail(
                     "scatter_reduce", "prod"
                 ),  # vmap (looks like you are calling item/data-dependent)
@@ -1171,10 +1174,8 @@ def vjp_of_vjp(*args_and_cotangents):
             # TODO: implement batching rule
             skip("_batch_norm_with_update"),
             xfail("__getitem__", ""),  # dynamic error
-            xfail("nanquantile", device_type="cpu"),  # checks q via a .item() call
             xfail("nn.functional.gaussian_nll_loss"),  # checks var for if any value < 0
             xfail("narrow"),  # .item() call
-            xfail("quantile", device_type="cpu"),  # checks q via a .item() call
             xfail("view_as_complex"),  # Tensor must have a last dimension with stride 1
             # required rank 4 tensor to use channels_last format
             xfail("bfloat16"),
@@ -1878,7 +1879,10 @@ def push_vjp(primals, cotangents):
             def tree_map2(fn, first, second):
                 flat_first, spec_first = tree_flatten(first)
                 flat_second, spec_second = tree_flatten(second)
-                assert spec_first == spec_second
+                if spec_first != spec_second:
+                    raise AssertionError(
+                        f"Tree specs mismatch: {spec_first} != {spec_second}"
+                    )
                 flat_result = [fn(f, s) for f, s in zip(flat_first, flat_second)]
                 return tree_unflatten(flat_result, spec_first)
 
@@ -2498,7 +2502,8 @@ def compute_grad(cotangents):
 
     def test_vmapvmapjvp_linalg_solve(self):
         ops = [op for op in op_db if op.name == "linalg.solve"]
-        assert len(ops) > 0
+        if len(ops) == 0:
+            raise AssertionError("Expected at least one linalg.solve op")
 
         # this specializes a lot of code from the get_fallback_and_vmap_exhaustive test. If we need this more
         # generally, this could go for a refactor
@@ -2547,7 +2552,8 @@ def f(x):
                         (torch.randn_like(without_grad),),
                     )
             else:
-                assert grad_op == "vjp"
+                if grad_op != "vjp":
+                    raise AssertionError(f"Expected 'vjp', got {grad_op!r}")
                 with self.assertRaisesRegex(
                     RuntimeError,
                     "During a grad .* attempted to call in-place operation",
@@ -2583,7 +2589,8 @@ def f(x):
                         (torch.randn_like(without_grad),),
                     )
                 else:
-                    assert grad_op == "vjp"
+                    if grad_op != "vjp":
+                        raise AssertionError(f"Expected 'vjp', got {grad_op!r}")
                     vjp(f, torch.randn_like(without_grad))
 
     @parametrize("grad_op", ["jvp", "vjp"])
@@ -2616,7 +2623,8 @@ def f(x):
                         (torch.randn_like(without_grad),),
                     )
             else:
-                assert grad_op == "vjp"
+                if grad_op != "vjp":
+                    raise AssertionError(f"Expected 'vjp', got {grad_op!r}")
                 with self.assertRaisesRegex(
                     RuntimeError,
                     "During a grad .* attempted to call in-place operation",
diff --git a/test/functorch/test_subclass_codegen.py b/test/functorch/test_subclass_codegen.py
new file mode 100644
index 0000000000000..da9375a0a2393
--- /dev/null
+++ b/test/functorch/test_subclass_codegen.py
@@ -0,0 +1,212 @@
+# Owner(s): ["module: functorch"]
+
+import logging
+from contextlib import contextmanager
+
+import torch
+import torch._functorch.config
+from torch.testing._internal.common_utils import run_tests, TestCase
+from torch.testing._internal.two_tensor import TwoTensor
+
+
+trace_log = logging.getLogger("torch.__trace")
+
+
+class TestSubclassCodegen(TestCase):
+    @contextmanager
+    def _capture_wrapper_source(self):
+        """Capture subclass_wrapper artifacts from the structured trace log."""
+        captured: list[str] = []
+
+        class _ArtifactHandler(logging.Handler):
+            def emit(self, record):
+                metadata = getattr(record, "metadata", {})
+                if (
+                    "artifact" in metadata
+                    and metadata["artifact"].get("name") == "subclass_wrapper"
+                ):
+                    payload = getattr(record, "payload", None)
+                    if payload is not None:
+                        captured.append(payload)
+
+        handler = _ArtifactHandler()
+        handler.setLevel(logging.DEBUG)
+        old_level = trace_log.level
+        trace_log.setLevel(logging.DEBUG)
+        trace_log.addHandler(handler)
+        try:
+            yield captured
+        finally:
+            trace_log.removeHandler(handler)
+            trace_log.setLevel(old_level)
+
+    def test_compile_simple(self):
+        """torch.compile with TwoTensor produces correct wrapper and output."""
+        with self._capture_wrapper_source() as captured:
+
+            @torch.compile(backend="aot_eager")
+            def f(x):
+                return x * 2
+
+            a = torch.randn(4)
+            b = torch.randn(4)
+            tt = TwoTensor(a, b)
+            out = f(tt)
+
+        self.assertIsInstance(out, TwoTensor)
+        self.assertEqual(out.a, a * 2)
+        self.assertEqual(out.b, b * 2)
+
+        self.assertEqual(len(captured), 1)
+        self.assertExpectedInline(
+            captured[0],
+            """\
+def inner_fn(args):
+    unwrapped_args = []
+    _inp_0 = args[0]
+    assert type(_inp_0) is _expected_type_1, f'expected {_expected_type_1}, got {type(_inp_0)}'
+    unwrapped_args.append(_inp_0.a)
+    unwrapped_args.append(_inp_0.b)
+    unwrapped_args.extend(args[1:])
+    args.clear()
+    unwrapped_outs = compiled_fn(unwrapped_args)
+    _out_inner_2 = {'a': unwrapped_outs[0], 'b': unwrapped_outs[1]}
+    _out_5 = _subclass_type_3.__tensor_unflatten__(_out_inner_2, _meta_4, (4,), (1,))
+    return (_out_5,)""",
+        )
+
+    def test_compile_nested_subclass(self):
+        """torch.compile with nested TwoTensor produces correct recursive wrapper."""
+        with self._capture_wrapper_source() as captured:
+
+            @torch.compile(backend="aot_eager")
+            def f(x):
+                return x.sin().cos()
+
+            a1 = torch.randn(4)
+            a2 = torch.randn(4)
+            a3 = torch.randn(4)
+            a4 = torch.randn(4)
+            inner_a = TwoTensor(a1, a2)
+            inner_b = TwoTensor(a3, a4)
+            tt = TwoTensor(inner_a, inner_b)
+            out = f(tt)
+
+        self.assertIsInstance(out, TwoTensor)
+        self.assertIsInstance(out.a, TwoTensor)
+        self.assertIsInstance(out.b, TwoTensor)
+        self.assertEqual(out.a.a, inner_a.a.sin().cos())
+        self.assertEqual(out.b.b, inner_b.b.sin().cos())
+
+        self.assertEqual(len(captured), 1)
+        self.assertExpectedInline(
+            captured[0],
+            """\
+def inner_fn(args):
+    unwrapped_args = []
+    _inp_0 = args[0]
+    assert type(_inp_0) is _expected_type_1, f'expected {_expected_type_1}, got {type(_inp_0)}'
+    _inner_2 = _inp_0.a
+    unwrapped_args.append(_inner_2.a)
+    unwrapped_args.append(_inner_2.b)
+    _inner_3 = _inp_0.b
+    unwrapped_args.append(_inner_3.a)
+    unwrapped_args.append(_inner_3.b)
+    unwrapped_args.extend(args[1:])
+    args.clear()
+    unwrapped_outs = compiled_fn(unwrapped_args)
+    _out_inner_5 = {'a': unwrapped_outs[0], 'b': unwrapped_outs[1]}
+    _out_8 = _subclass_type_6.__tensor_unflatten__(_out_inner_5, _meta_7, (4,), (1,))
+    _out_inner_9 = {'a': unwrapped_outs[2], 'b': unwrapped_outs[3]}
+    _out_12 = _subclass_type_10.__tensor_unflatten__(_out_inner_9, _meta_11, (4,), (1,))
+    _out_inner_4 = {'a': _out_8, 'b': _out_12}
+    _out_15 = _subclass_type_13.__tensor_unflatten__(_out_inner_4, _meta_14, (4,), (1,))
+    return (_out_15,)""",
+        )
+
+    def test_trailing_args_forwarded(self):
+        """Extra trailing args (e.g. rng seed/offset) are forwarded to compiled_fn."""
+        # Build SubclassCreationMeta manually to avoid __post_init__ fake tensor check
+        from dataclasses import dataclass
+        from typing import Any
+
+        from torch._functorch._aot_autograd.schemas import PlainTensorMeta
+        from torch._functorch._aot_autograd.subclass_codegen import (
+            _codegen_subclass_wrapper_source,
+        )
+
+        @dataclass
+        class _TestSubclassMeta:
+            flat_tensor_start_idx: int
+            arg_count: int
+            included_subclass_symints: bool
+            attrs: dict
+            outer_size: tuple
+            outer_stride: tuple
+            meta: Any
+            original_subclass: Any
+            original_subclass_type: type
+
+        inp_meta = _TestSubclassMeta(
+            flat_tensor_start_idx=0,
+            arg_count=2,
+            included_subclass_symints=True,
+            attrs={
+                "a": PlainTensorMeta(unwrapped_idx=0),
+                "b": PlainTensorMeta(unwrapped_idx=1),
+            },
+            outer_size=(4,),
+            outer_stride=(1,),
+            meta=None,
+            original_subclass=None,
+            original_subclass_type=TwoTensor,
+        )
+        out_meta = _TestSubclassMeta(
+            flat_tensor_start_idx=0,
+            arg_count=2,
+            included_subclass_symints=True,
+            attrs={
+                "a": PlainTensorMeta(unwrapped_idx=0),
+                "b": PlainTensorMeta(unwrapped_idx=1),
+            },
+            outer_size=(4,),
+            outer_stride=(1,),
+            meta=None,
+            original_subclass=None,
+            original_subclass_type=TwoTensor,
+        )
+
+        source, globals_dict = _codegen_subclass_wrapper_source(
+            inp_metas=[inp_meta],
+            out_metas=[out_meta],
+            num_fw_outs_saved_for_bw=None,
+        )
+
+        received_args = []
+
+        def mock_compiled_fn(args):
+            received_args.extend(args)
+            return [args[0] * 2, args[1] * 2]
+
+        globals_dict["compiled_fn"] = mock_compiled_fn
+        local_dict = {}
+        exec(compile(source, "<test>", "exec"), globals_dict, local_dict)  # noqa: S102
+        wrapper = local_dict["inner_fn"]
+
+        a = torch.randn(4)
+        b = torch.randn(4)
+        tt = TwoTensor(a, b)
+        seed = torch.tensor(42)
+        offset = torch.tensor(0)
+        # Simulate FunctionalizedRngRuntimeWrapper appending rng state
+        wrapper([tt, seed, offset])
+
+        self.assertEqual(len(received_args), 4)
+        self.assertEqual(received_args[0], a)
+        self.assertEqual(received_args[1], b)
+        self.assertIs(received_args[2], seed)
+        self.assertIs(received_args[3], offset)
+
+
+if __name__ == "__main__":
+    run_tests()
diff --git a/test/functorch/test_vmap.py b/test/functorch/test_vmap.py
index 28572e528ebbd..e5672d814ae9e 100644
--- a/test/functorch/test_vmap.py
+++ b/test/functorch/test_vmap.py
@@ -71,6 +71,7 @@
     run_tests,
     skipIfTorchDynamo,
     subtest,
+    TEST_WITH_ROCM,
     TEST_WITH_TORCHDYNAMO,
     TestCase,
     unMarkDynamoStrictTest,
@@ -454,7 +455,8 @@ def f(t):
             return OrderedDict([("sin", y), ("cos", t.cos())])
 
         out = vmap(f)(x)
-        assert isinstance(out, OrderedDict)
+        if not isinstance(out, OrderedDict):
+            raise AssertionError(f"Expected OrderedDict, got {type(out)}")
         expected = f(x)
         self.assertEqual(out["sin"], expected["sin"])
         self.assertEqual(out["cos"], expected["cos"])
@@ -1162,41 +1164,62 @@ def _test_vmap_autocast(self, device):
         def func1(x, y, z, w):
             with torch.autocast(dtype=amp_dtype, device_type=device):
                 e_float16 = torch.matmul(x, y)
-                assert e_float16.dtype == amp_dtype, e_float16.dtype
+                if e_float16.dtype != amp_dtype:
+                    raise AssertionError(
+                        f"Expected dtype {amp_dtype}, got {e_float16.dtype}"
+                    )
                 f_float16 = torch.matmul(z, e_float16)
-                assert f_float16.dtype == amp_dtype, f_float16.dtype
+                if f_float16.dtype != amp_dtype:
+                    raise AssertionError(
+                        f"Expected dtype {amp_dtype}, got {f_float16.dtype}"
+                    )
             return torch.matmul(w, f_float16.float())
 
         expected = func1(a_float32, b_float32, c_float32, d_float32)
         out = vmap(func1)(a_float32, b_float32, c_float32, d_float32)
-        assert expected.allclose(out)
+        if not expected.allclose(out):
+            raise AssertionError("Expected func1 output to be close to vmap output")
 
         # Case 2, autocast decorator inside vmapped function
         @torch.autocast(dtype=amp_dtype, device_type=device)
         def func2(x, y, z, w):
             e_float16 = torch.matmul(x, y)
-            assert e_float16.dtype == amp_dtype, e_float16.dtype
+            if e_float16.dtype != amp_dtype:
+                raise AssertionError(
+                    f"Expected dtype {amp_dtype}, got {e_float16.dtype}"
+                )
             f_float16 = torch.matmul(z, e_float16)
-            assert f_float16.dtype == amp_dtype, f_float16.dtype
+            if f_float16.dtype != amp_dtype:
+                raise AssertionError(
+                    f"Expected dtype {amp_dtype}, got {f_float16.dtype}"
+                )
             return torch.matmul(w, f_float16)
 
         expected = func2(a_float32, b_float32, c_float32, d_float32)
         out = vmap(func2)(a_float32, b_float32, c_float32, d_float32)
-        assert expected.allclose(out)
+        if not expected.allclose(out):
+            raise AssertionError("Expected func2 output to be close to vmap output")
 
         # Case 3, autocast is outside vmapped function
         def func3(x, y, z, w):
             e_float16 = torch.matmul(x, y)
-            assert e_float16.dtype == amp_dtype, e_float16.dtype
+            if e_float16.dtype != amp_dtype:
+                raise AssertionError(
+                    f"Expected dtype {amp_dtype}, got {e_float16.dtype}"
+                )
             f_float16 = torch.matmul(z, e_float16)
-            assert f_float16.dtype == amp_dtype, f_float16.dtype
+            if f_float16.dtype != amp_dtype:
+                raise AssertionError(
+                    f"Expected dtype {amp_dtype}, got {f_float16.dtype}"
+                )
             return torch.matmul(w, f_float16)
 
         with torch.autocast(dtype=amp_dtype, device_type=device):
             expected = func3(a_float32, b_float32, c_float32, d_float32)
             out = vmap(func3)(a_float32, b_float32, c_float32, d_float32)
 
-        assert expected.allclose(out)
+        if not expected.allclose(out):
+            raise AssertionError("Expected func3 output to be close to vmap output")
 
     @unittest.skip("Somehow, vmap and autocast do not work on CPU")
     def test_vmap_autocast_cpu(self):
@@ -1281,14 +1304,19 @@ def reference_vmap(op, inputs, in_dims=0, out_dims=0, return_nt=False):
     if isinstance(in_dims, int):
         in_dims = (in_dims,) * len(inputs)
     bdim_sizes = [inp.size(dim) for inp, dim in zip(inputs, in_dims) if dim is not None]
-    assert all(bdim_size == bdim_sizes[0] for bdim_size in bdim_sizes)
+    if not all(bdim_size == bdim_sizes[0] for bdim_size in bdim_sizes):
+        raise AssertionError(f"Expected all bdim sizes to be equal, got {bdim_sizes}")
     bdim_size = bdim_sizes[0]
     results = tuple(op(*slice_inputs(inputs, in_dims, i)) for i in range(bdim_size))
 
-    assert len(results) > 0
+    if len(results) == 0:
+        raise AssertionError("Expected at least one result")
     op_has_single_return = not isinstance(results[0], tuple)
     if op_has_single_return:
-        assert all(isinstance(result, torch.Tensor) for result in results)
+        if not all(isinstance(result, torch.Tensor) for result in results):
+            raise AssertionError(
+                "Expected all results to be torch.Tensor for single-return op"
+            )
         if isinstance(out_dims, int):
             out_dims = (out_dims,) * 1
         if return_nt:
@@ -1296,9 +1324,11 @@ def reference_vmap(op, inputs, in_dims=0, out_dims=0, return_nt=False):
         else:
             return torch.stack(results, dim=out_dims[0])
 
-    assert all(isinstance(result, tuple) for result in results)
+    if not all(isinstance(result, tuple) for result in results):
+        raise AssertionError("Expected all results to be tuples")
     num_returns = len(results[0])
-    assert all(len(result) == num_returns for result in results)
+    if not all(len(result) == num_returns for result in results):
+        raise AssertionError(f"Expected all results to have {num_returns} elements")
     if isinstance(out_dims, int):
         out_dims = (out_dims,) * num_returns
     if return_nt:
@@ -2084,7 +2114,8 @@ def test_stride(self):
         x = torch.randn(B0, 2, 5, 7)
 
         def foo(x):
-            assert x.stride() == (7 * 5, 7, 1)
+            if x.stride() != (7 * 5, 7, 1):
+                raise AssertionError(f"Expected stride (35, 7, 1), got {x.stride()}")
             return x
 
         vmap(foo)(x)
@@ -2092,7 +2123,10 @@ def foo(x):
         x = torch.randn(2, B0, 5, 7).movedim(1, 0)
 
         def bar(x):
-            assert x.stride() == (7 * 5 * B0, 7, 1)
+            if x.stride() != (7 * 5 * B0, 7, 1):
+                raise AssertionError(
+                    f"Expected stride ({7 * 5 * B0}, 7, 1), got {x.stride()}"
+                )
             return x
 
         vmap(bar)(x)
@@ -2444,7 +2478,8 @@ def foo(x):
 
         # is_contiguous on empty tensor is True
         def bar(x):
-            assert x.is_contiguous()
+            if not x.is_contiguous():
+                raise AssertionError("Expected tensor to be contiguous")
             return x
 
         vmap(bar)(torch.randn(B0, 0, 3))
@@ -3304,9 +3339,11 @@ def test_conj_bit(self):
         x = torch.tensor([1 + 1j, 2 + 1j])
 
         def foo(x):
-            assert not x.is_conj()
+            if x.is_conj():
+                raise AssertionError("Expected x to not be conj")
             y = x.conj()
-            assert y.is_conj()
+            if not y.is_conj():
+                raise AssertionError("Expected y to be conj")
             return y
 
         res = vmap(foo)(x)
@@ -3682,7 +3719,8 @@ def vector_hessian_product(*vectors):
                     allow_unused=True,
                 )
                 outputs = tuple(out for out in outputs if out is not None)
-                assert len(outputs) > 0
+                if len(outputs) == 0:
+                    raise AssertionError("Expected at least one non-None output")
                 return outputs
 
             self._vmap_test(
@@ -4634,7 +4672,8 @@ def compute_A(out):
             return U @ diag_S @ Vh
 
         opinfos = [op for op in op_db if op.name == "linalg.svd"]
-        assert len(opinfos) > 0
+        if len(opinfos) == 0:
+            raise AssertionError("Expected at least one opinfo for linalg.svd")
 
         for op in opinfos:
             self.opinfo_vmap_test(
@@ -4661,7 +4700,8 @@ def compute_A(out):
             return Q @ diag_L @ Qh
 
         opinfos = [op for op in op_db if op.name == "linalg.eigh"]
-        assert len(opinfos) > 0
+        if len(opinfos) == 0:
+            raise AssertionError("Expected at least one opinfo for linalg.eigh")
 
         for op in opinfos:
             self.opinfo_vmap_test(
@@ -4788,6 +4828,60 @@ def test():
 
         check_vmap_fallback(self, test, Tensor.fill_)
 
+    @parametrize(
+        "op,msg,extra_positional_args,extra_kwargs",
+        [
+            subtest(
+                (
+                    Tensor.scatter_add_,
+                    "out-of-place operators instead of scatter_add_",
+                    (),
+                    {},
+                ),
+                name="scatter_add",
+            ),
+            subtest(
+                (
+                    Tensor.scatter_reduce_,
+                    "out-of-place operators instead of scatter_reduce_",
+                    ("sum",),
+                    {"include_self": True},
+                ),
+                name="scatter_reduce",
+            ),
+        ],
+    )
+    def test_scatter_inplace_self_not_batched(
+        self, device, op, msg, extra_positional_args, extra_kwargs
+    ):
+        x = torch.zeros(5, device=device)
+
+        def call_op(self, dim, index, src):
+            return op(
+                self,
+                dim,
+                index,
+                src,
+                *extra_positional_args,
+                **extra_kwargs,
+            )
+
+        with self.assertRaisesRegex(RuntimeError, msg):
+            vmap(call_op, in_dims=(None, None, None, 0))(
+                x,
+                0,
+                torch.tensor([0, 1], device=device),
+                torch.randn(2, 2, device=device),
+            )
+
+        with self.assertRaisesRegex(RuntimeError, msg):
+            vmap(call_op, in_dims=(None, None, 0, None))(
+                x,
+                0,
+                torch.tensor([[0, 1], [2, 3]], device=device),
+                torch.randn(2, device=device),
+            )
+
     @tf32_on_and_off(0.005)
     def test_conv_double_backward(self, device):
         images = torch.randn(2, 1, 5, 5, device=device)
@@ -4904,6 +4998,23 @@ def test_group_norm(self, device):
         bias = torch.randn(B, C)
         test(self, op, (x, 4, weight, bias), in_dims=(0, None, 0, 0))
 
+    def test_group_norm_layout_corruption(self, device):
+        # Regression test for https://github.com/pytorch/pytorch/issues/176432
+        def op_function(cotangent):
+            input = torch.tensor([[-6.517, -6.264]], device=device, requires_grad=True)
+            result = F.group_norm(input, num_groups=1)
+            return torch.autograd.grad(result, input, grad_outputs=cotangent)[0]
+
+        cotangent = torch.tensor([[-0.0236, -0.1431]], device=device)
+        result_in_dims_0 = vmap(op_function, in_dims=0)(
+            cotangent.unsqueeze(0).expand(2, -1, -1)
+        )
+        result_in_dims_neg1 = vmap(op_function, in_dims=-1)(
+            cotangent.unsqueeze(-1).expand(-1, -1, 2)
+        )
+
+        self.assertEqual(result_in_dims_0, result_in_dims_neg1)
+
     def test_index_put(self, device):
         def test(f, t, idx, values):
             base = f(t[0], idx[0], values[0])
@@ -5064,9 +5175,12 @@ def test_torch_return_types_returns(self, device):
                 vmap(torch.topk, (0, None, None))(t, 1, 0), torch.return_types.topk
             )
         )
-        self.assertTrue(
-            isinstance(vmap(torch.linalg.eig, (0))(t), torch.return_types.linalg_eig)
-        )
+        if not (TEST_WITH_ROCM and not torch.cuda.has_magma):
+            self.assertTrue(
+                isinstance(
+                    vmap(torch.linalg.eig, (0))(t), torch.return_types.linalg_eig
+                )
+            )
 
     def test_namedtuple_returns(self, device):
         Point = namedtuple("Point", ["x", "y"])
@@ -5325,7 +5439,10 @@ def _get_image(self, batched_input, batch_size, device):
             return torch.ones([batch_size, 3, 3, 14, 14], device=device)
         if batched_input == "last":
             return torch.ones([3, 3, 14, 14, batch_size], device=device)
-        assert batched_input == "none"
+        if batched_input != "none":
+            raise AssertionError(
+                f"Expected batched_input to be 'none', got '{batched_input}'"
+            )
         return torch.ones([3, 3, 14, 14], device=device)
 
     def _assert_all_slices_equal(self, tensor):
@@ -5337,7 +5454,10 @@ def _assert_all_slices_unique(self, tensor):
         slices_equal = vmap(vmap(lambda x, y: (x == y).all(), (0, None)), (None, 0))(
             tensor, tensor
         )
-        assert slices_equal.shape == (B0, B0)
+        if slices_equal.shape != (B0, B0):
+            raise AssertionError(
+                f"Expected shape ({B0}, {B0}), got {slices_equal.shape}"
+            )
         slices_equal.diagonal().zero_()
         self.assertEqual(slices_equal, torch.zeros_like(slices_equal))
 
@@ -5366,7 +5486,10 @@ def get_in_dim(batched_string):
                 return 0
             if batched_string == "last":
                 return -1
-            assert batched_string == "none"
+            if batched_string != "none":
+                raise AssertionError(
+                    f"Expected batched_string to be 'none', got '{batched_string}'"
+                )
             return None
 
         batched_strings = batched_strings + (
@@ -5500,7 +5623,10 @@ def op(t, ignored):
             self._assert_all_slices_unique(vmap_result)
             return
 
-        assert randomness == "same"
+        if randomness != "same":
+            raise AssertionError(
+                f"Expected randomness to be 'same', got '{randomness}'"
+            )
         self._assert_all_slices_equal(vmap_result)
 
     @parametrize("randomness", ["error", "same", "different"])
@@ -5530,7 +5656,10 @@ def op(t, ignored):
             self._assert_all_slices_unique(vmap_result)
             return
 
-        assert randomness == "same"
+        if randomness != "same":
+            raise AssertionError(
+                f"Expected randomness to be 'same', got '{randomness}'"
+            )
         self._assert_all_slices_equal(vmap_result)
 
     @parametrize("randomness", ["error", "same", "different"])
@@ -5579,7 +5708,10 @@ def op(t, ignored):
             self._assert_all_slices_unique(vmap_result)
             return
 
-        assert randomness == "same"
+        if randomness != "same":
+            raise AssertionError(
+                f"Expected randomness to be 'same', got '{randomness}'"
+            )
         self._assert_all_slices_equal(vmap_result)
 
     @parametrize("randomness", ["error", "same", "different"])
@@ -5623,7 +5755,10 @@ def op(t, ignored):
             self._assert_all_slices_unique(vmap_result)
             return
 
-        assert randomness == "same"
+        if randomness != "same":
+            raise AssertionError(
+                f"Expected randomness to be 'same', got '{randomness}'"
+            )
         self._assert_all_slices_equal(vmap_result)
 
     @parametrize("randomness", ["error", "same", "different"])
@@ -5673,7 +5808,10 @@ def test_like_functions(self, device, randomness, batched_input):
                     self.assertEqual(expected, vmap_result)
                 return
 
-            assert randomness == "same"
+            if randomness != "same":
+                raise AssertionError(
+                    f"Expected randomness to be 'same', got '{randomness}'"
+                )
             if batched_input != "none":
                 passed = passed[0]
             expected = op(passed, 0)
@@ -5870,7 +6008,12 @@ def test_random_binary_out_of_place(
                 self._assert_all_slices_unique(vmap_result)
                 self.assertEqual(vmap_result, expected)
             else:
-                assert batched_input == "none" and batched_other == "none"
+                if batched_input != "none" or batched_other != "none":
+                    raise AssertionError(
+                        f"Expected batched_input='none' and batched_other='none', "
+                        f"got batched_input='{batched_input}' and "
+                        f"batched_other='{batched_other}'"
+                    )
                 expected = op(input, other, always_batched)
                 self._assert_all_slices_equal(vmap_result)
                 for i in range(B0):
@@ -5957,7 +6100,8 @@ def flatten_input(input, batch_call, batch_location):
                 )
 
             ret = input.flatten(start_idx, end_idx)
-            assert ret.dim() == final_size
+            if ret.dim() != final_size:
+                raise AssertionError(f"Expected dim {final_size}, got {ret.dim()}")
             return ret
 
         def op(input, _):
@@ -6280,8 +6424,14 @@ def _create_nt(self, dims, device):
     # Creates an NT matching another NT's number of components and
     # shape / ragged structure for all dims specified to be -1.
     def _nt_from_similar(self, other, dims):
-        assert len(dims) == other.dim()
-        assert dims[0] == -1 or dims[0] == other.size(0)
+        if len(dims) != other.dim():
+            raise AssertionError(
+                f"Expected len(dims) == other.dim(), got {len(dims)} != {other.dim()}"
+            )
+        if dims[0] != -1 and dims[0] != other.size(0):
+            raise AssertionError(
+                f"Expected dims[0] to be -1 or {other.size(0)}, got {dims[0]}"
+            )
 
         ret_sizes = []
         for t in other.unbind():
@@ -6335,7 +6485,8 @@ def f(x, y):
 
     def test_nt_acts_as_dense_in_vmap(self, device):
         def f(x):
-            assert not x.is_nested
+            if x.is_nested:
+                raise AssertionError("Expected x to not be nested inside vmap")
             return x
 
         x = self._create_nt([5, None, 3], device=device)
diff --git a/test/functorch/test_vmap_registrations.py b/test/functorch/test_vmap_registrations.py
index adb66ac4d9709..d9d12526ea833 100644
--- a/test/functorch/test_vmap_registrations.py
+++ b/test/functorch/test_vmap_registrations.py
@@ -61,6 +61,7 @@
     "aten::cummaxmin_backward",
     "aten::data",
     "aten::diagflat",
+    "aten::dim",
     "aten::divide.out_mode",
     "aten::divide_.Scalar",
     "aten::dropout_",
@@ -82,6 +83,7 @@
     "aten::floor_divide_.Scalar",
     "aten::frobenius_norm",
     "aten::fused_moving_avg_obs_fake_quant",
+    "aten::get_device",
     "aten::get_gradients",
     "aten::greater_.Scalar",
     "aten::greater_.Tensor",
@@ -145,6 +147,7 @@
     "aten::norm_except_dim",
     "aten::not_equal_.Scalar",
     "aten::not_equal_.Tensor",
+    "aten::numel",
     "aten::one_hot",
     "aten::output_nr",
     "aten::pad_sequence",
@@ -203,6 +206,7 @@
     "aten::std_mean.names_dim",
     "aten::stft",
     "aten::stft.center",
+    "aten::storage_offset",
     "aten::stride.int",
     "aten::subtract.Scalar",
     "aten::subtract_.Scalar",
@@ -289,31 +293,34 @@ class TestFunctorchDispatcher(TestCase):
     def test_register_a_batching_rule_for_composite_implicit_autograd(
         self, registration
     ):
-        assert registration not in FuncTorchBatchedRegistrations, (
-            f"You've added a batching rule for a CompositeImplicitAutograd operator {registration}. "
-            "The correct way to add vmap support for it is to put it into BatchRulesDecomposition to "
-            "reuse the CompositeImplicitAutograd decomposition"
-        )
+        if registration in FuncTorchBatchedRegistrations:
+            raise AssertionError(
+                f"You've added a batching rule for a CompositeImplicitAutograd operator {registration}. "
+                "The correct way to add vmap support for it is to put it into BatchRulesDecomposition to "
+                "reuse the CompositeImplicitAutograd decomposition"
+            )
 
     @dispatch_registrations(
         "FuncTorchBatchedDecomposition", xfail_functorch_batched_decomposition
     )
     def test_register_functorch_batched_decomposition(self, registration):
-        assert registration in CompositeImplicitAutogradRegistrations, (
-            f"The registrations in BatchedDecompositions.cpp must be for CompositeImplicitAutograd "
-            f"operations. If your operation {registration} is not CompositeImplicitAutograd, "
-            "then please register it to the FuncTorchBatched key in another file."
-        )
+        if registration not in CompositeImplicitAutogradRegistrations:
+            raise AssertionError(
+                f"The registrations in BatchedDecompositions.cpp must be for CompositeImplicitAutograd "
+                f"operations. If your operation {registration} is not CompositeImplicitAutograd, "
+                "then please register it to the FuncTorchBatched key in another file."
+            )
 
     @dispatch_registrations(
         "CompositeImplicitAutograd", xfail_not_implemented, filter_vmap_implementable
     )
     def test_unimplemented_batched_registrations(self, registration):
-        assert registration in FuncTorchBatchedDecompositionRegistrations, (
-            f"Please check that there is an OpInfo that covers the operator {registration} "
-            "and add a registration in BatchedDecompositions.cpp. "
-            "If your operator isn't user facing, please add it to the xfail list"
-        )
+        if registration not in FuncTorchBatchedDecompositionRegistrations:
+            raise AssertionError(
+                f"Please check that there is an OpInfo that covers the operator {registration} "
+                "and add a registration in BatchedDecompositions.cpp. "
+                "If your operator isn't user facing, please add it to the xfail list"
+            )
 
 
 instantiate_parametrized_tests(TestFunctorchDispatcher)
diff --git a/test/fx/test_dce_pass.py b/test/fx/test_dce_pass.py
index 7fd3a6dbb0041..50768edcea4b1 100644
--- a/test/fx/test_dce_pass.py
+++ b/test/fx/test_dce_pass.py
@@ -1,7 +1,6 @@
 # Owner(s): ["module: fx"]
 import copy
 import unittest
-from typing import Optional
 
 import torch
 import torch.fx
@@ -42,7 +41,7 @@ def _run_dce_and_test(
         self,
         m: torch.nn.Module,
         expect_dce_changes: bool,
-        modules_to_be_leafs: Optional[set[type]] = None,
+        modules_to_be_leafs: set[type] | None = None,
         custom: bool = False,
     ):
         class TestTracer(torch.fx.Tracer):
diff --git a/test/fx/test_graph_pickler.py b/test/fx/test_graph_pickler.py
index e79a4ace9b024..0610b4a73598d 100644
--- a/test/fx/test_graph_pickler.py
+++ b/test/fx/test_graph_pickler.py
@@ -679,6 +679,33 @@ def inner_fn(b):
             self.assertIn("nested_closure", node.meta)
             self.assertEqual(node.meta["nested_closure"](3), 3 + 5 + 20)
 
+    def test_node_with_slice(self):
+        from torch._subclasses.fake_tensor import FakeTensorMode
+        from torch.fx.experimental.symbolic_shapes import ShapeEnv
+
+        def foo(x):
+            return x[0 : x.shape[0]]
+
+        gm = torch.fx.symbolic_trace(foo)
+
+        self.assertExpectedInline(
+            gm.code.strip(),
+            """\
+def forward(self, x):
+    getattr_1 = x.shape
+    getitem = getattr_1[0];  getattr_1 = None
+    getitem_1 = x[slice(0, getitem, None)];  x = getitem = None
+    return getitem_1""",
+        )
+        options = self.Options(node_metadata_key_filter=None)
+        serialized = self.GraphPickler.dumps(gm, options)
+
+        fake_mode = FakeTensorMode(shape_env=ShapeEnv())
+        deserialized = self.GraphPickler.loads(serialized, fake_mode)
+        deserialized.recompile()
+
+        self.assertEqual(gm.code, deserialized.code)
+
     def test_lambda_with_default_arguments(self):
         """
         Test that lambdas with default arguments can be serialized. Standard
@@ -831,6 +858,243 @@ def forward(self, x):
             self.assertIn("custom_key", pickle_data.meta)
 
 
+class TestNodeStateSerialization(TestCase):
+    def test_type_entry_preserved_in_getstate(self):
+        class M(torch.nn.Module):
+            def forward(self, x):
+                y = torch.neg(x)
+                return y + 1
+
+        gm = torch.fx.symbolic_trace(M())
+        node = next(n for n in gm.graph.nodes if n.op == "call_function")
+        node.type = torch.Tensor
+        state = node.__getstate__()
+        self.assertIs(state["type"], torch.Tensor)
+
+
+@unittest.skipUnless(HAS_DILL, "dill not available")
+class TestIgnoreRawNode(TestCase):
+    """Tests for the ignore_raw_node option in GraphPickler.Options."""
+
+    def setUp(self):
+        super().setUp()
+        from torch._subclasses.fake_tensor import FakeTensorMode
+        from torch.fx._graph_pickler import GraphPickler, Options
+        from torch.fx.experimental.symbolic_shapes import ShapeEnv
+
+        self.GraphPickler = GraphPickler
+        self.Options = Options
+        self.fake_mode = FakeTensorMode(shape_env=ShapeEnv())
+
+    def _make_graph_with_raw_node_in_meta(self):
+        """Return a graph module whose first call_function node has a raw
+        torch.fx.Node stored in its metadata under the key 'raw_ref'."""
+
+        class M(torch.nn.Module):
+            def forward(self, x):
+                return x + 1
+
+        gm = torch.fx.symbolic_trace(M())
+        call_node = next((n for n in gm.graph.nodes if n.op == "call_function"), None)
+        self.assertIsNotNone(call_node)
+        # Store a raw Node reference in meta – this is the problematic case.
+        call_node.meta["raw_ref"] = call_node
+        return gm
+
+    def test_raw_node_in_meta_raises_by_default(self):
+        """Pickling should raise AssertionError when a raw Node is in metadata
+        and ignore_raw_node is False (the default)."""
+        gm = self._make_graph_with_raw_node_in_meta()
+        with self.assertRaises(AssertionError) as cm:
+            self.GraphPickler.dumps(gm)
+        self.assertIn("raw Node", str(cm.exception))
+
+    def test_raw_node_in_meta_with_ignore_raw_node(self):
+        """With ignore_raw_node=True, pickling should succeed and the raw Node
+        should be replaced with None after round-trip deserialization."""
+        gm = self._make_graph_with_raw_node_in_meta()
+        options = self.Options(ignore_raw_node=True)
+        data = self.GraphPickler.dumps(gm, options)
+        restored = self.GraphPickler.loads(data, self.fake_mode)
+        self.assertIsInstance(restored, torch.fx.GraphModule)
+        call_node = next(
+            (n for n in restored.graph.nodes if n.op == "call_function"), None
+        )
+        self.assertIsNotNone(call_node)
+        self.assertIsNone(call_node.meta.get("raw_ref"))
+
+
+class _WeakrefTarget:
+    """A simple picklable class that supports weak references, for use in tests.
+
+    Plain dicts do not support weak references in Python, so tests must use
+    instances of a regular class instead.
+    """
+
+    def __init__(self, **kwargs):
+        self.__dict__.update(kwargs)
+
+
+@unittest.skipUnless(HAS_DILL, "dill not available")
+class TestWeakrefPickle(TestCase):
+    """Tests that weakref objects are properly serialized and reconstructed."""
+
+    def setUp(self):
+        super().setUp()
+        import weakref
+
+        from torch._subclasses.fake_tensor import FakeTensorMode
+        from torch.fx._graph_pickler import GraphPickler, Options
+        from torch.fx.experimental.symbolic_shapes import ShapeEnv
+
+        self.weakref = weakref
+        self.GraphPickler = GraphPickler
+        self.Options = Options
+        self.fake_mode = FakeTensorMode(shape_env=ShapeEnv())
+
+    def _make_graph_with_weakref_in_meta(self, ref_obj):
+        """Return a graph module with a weakref stored in node metadata."""
+
+        class M(torch.nn.Module):
+            def forward(self, x):
+                return x + 1
+
+        gm = torch.fx.symbolic_trace(M())
+        call_node = next((n for n in gm.graph.nodes if n.op == "call_function"), None)
+        self.assertIsNotNone(call_node)
+        call_node.meta["weak_ref"] = ref_obj
+        return gm
+
+    def test_alive_weakref_in_meta_is_reconstructed(self):
+        """An alive weakref.ref in node metadata should be reconstructed as a weakref."""
+        target = _WeakrefTarget(key="value")
+        weak = self.weakref.ref(target)
+        gm = self._make_graph_with_weakref_in_meta(weak)
+        # Also store a strong ref so the referent survives after unpickling
+        call_node = next((n for n in gm.graph.nodes if n.op == "call_function"), None)
+        call_node.meta["strong_ref"] = target
+
+        options = self.Options(node_metadata_key_filter=None)
+        data = self.GraphPickler.dumps(gm, options)
+        restored = self.GraphPickler.loads(data, self.fake_mode)
+
+        self.assertIsInstance(restored, torch.fx.GraphModule)
+        call_node = next(
+            (n for n in restored.graph.nodes if n.op == "call_function"), None
+        )
+        self.assertIsNotNone(call_node)
+        restored_ref = call_node.meta.get("weak_ref")
+        self.assertIsInstance(restored_ref, self.weakref.ref)
+        self.assertEqual(restored_ref().key, "value")
+
+    def test_dead_weakref_in_meta_unpickles_as_callable_none(self):
+        """A dead weakref should unpickle as a callable that returns None."""
+        target = _WeakrefTarget()
+        weak = self.weakref.ref(target)
+        gm = self._make_graph_with_weakref_in_meta(weak)
+        # Kill the referent so the weakref is dead at pickle time
+        del target
+
+        options = self.Options(node_metadata_key_filter=None)
+        data = self.GraphPickler.dumps(gm, options)
+        restored = self.GraphPickler.loads(data, self.fake_mode)
+
+        self.assertIsInstance(restored, torch.fx.GraphModule)
+        call_node = next(
+            (n for n in restored.graph.nodes if n.op == "call_function"), None
+        )
+        self.assertIsNotNone(call_node)
+        restored_ref = call_node.meta.get("weak_ref")
+        # Should be callable and return None, like a dead weakref
+        self.assertIsNotNone(restored_ref)
+        self.assertIsNone(restored_ref())
+
+    def test_keyed_ref_in_meta_is_reconstructed(self):
+        """A weakref.KeyedRef (from WeakValueDictionary) should be reconstructed."""
+        wvd = self.weakref.WeakValueDictionary()
+        target = _WeakrefTarget(val=42)
+        wvd["k"] = target
+        keyed_ref = wvd.data["k"]
+        self.assertIsInstance(keyed_ref, self.weakref.KeyedRef)
+
+        gm = self._make_graph_with_weakref_in_meta(keyed_ref)
+        # Also store a strong ref so the referent survives after unpickling
+        call_node = next((n for n in gm.graph.nodes if n.op == "call_function"), None)
+        call_node.meta["strong_ref"] = target
+
+        options = self.Options(node_metadata_key_filter=None)
+        data = self.GraphPickler.dumps(gm, options)
+        restored = self.GraphPickler.loads(data, self.fake_mode)
+
+        self.assertIsInstance(restored, torch.fx.GraphModule)
+        call_node = next(
+            (n for n in restored.graph.nodes if n.op == "call_function"), None
+        )
+        self.assertIsNotNone(call_node)
+        restored_ref = call_node.meta.get("weak_ref")
+        self.assertIsInstance(restored_ref, self.weakref.ref)
+        self.assertEqual(restored_ref().val, 42)
+
+    def test_weakref_in_module_dict_is_reconstructed(self):
+        """A weakref stored in the graph module's __dict__ should be reconstructed."""
+
+        class M(torch.nn.Module):
+            def forward(self, x):
+                return x + 1
+
+        gm = torch.fx.symbolic_trace(M())
+        target = _WeakrefTarget(key="value")
+        gm._weak_cache = self.weakref.ref(target)
+        # Also store a strong ref so the referent survives after unpickling
+        gm._strong_cache = target
+
+        options = self.Options(node_metadata_key_filter=None)
+        data = self.GraphPickler.dumps(gm, options)
+        restored = self.GraphPickler.loads(data, self.fake_mode)
+
+        self.assertIsInstance(restored, torch.fx.GraphModule)
+        restored_ref = restored._weak_cache
+        self.assertIsInstance(restored_ref, self.weakref.ref)
+        self.assertEqual(restored_ref().key, "value")
+
+    def test_weakref_and_strong_ref_share_same_object(self):
+        """When a weakref and a strong ref point to the same object, pickle's
+        memo should deduplicate them so they share identity after unpickling.
+        This also covers the case where the weakref is the first reference
+        pickle encounters — the memo must still work correctly."""
+
+        class M(torch.nn.Module):
+            def forward(self, x):
+                return x + 1
+
+        gm = torch.fx.symbolic_trace(M())
+        target = _WeakrefTarget(key="value")
+        weak = self.weakref.ref(target)
+
+        call_node = next((n for n in gm.graph.nodes if n.op == "call_function"), None)
+        self.assertIsNotNone(call_node)
+        # Put weakref first so it's the first reference pickle encounters
+        call_node.meta["weak_ref"] = weak
+        call_node.meta["strong_ref"] = target
+
+        options = self.Options(node_metadata_key_filter=None)
+        data = self.GraphPickler.dumps(gm, options)
+        restored = self.GraphPickler.loads(data, self.fake_mode)
+
+        self.assertIsInstance(restored, torch.fx.GraphModule)
+        restored_node = next(
+            (n for n in restored.graph.nodes if n.op == "call_function"), None
+        )
+        self.assertIsNotNone(restored_node)
+
+        restored_weak = restored_node.meta["weak_ref"]
+        restored_strong = restored_node.meta["strong_ref"]
+
+        self.assertIsInstance(restored_weak, self.weakref.ref)
+        # The weakref's referent and the strong ref should be the same object
+        self.assertIs(restored_weak(), restored_strong)
+
+
 if __name__ == "__main__":
     from torch.testing._internal.common_utils import run_tests
 
diff --git a/test/fx/test_source_matcher_utils.py b/test/fx/test_source_matcher_utils.py
index 1a5755c383d9c..b7a670f4f1933 100644
--- a/test/fx/test_source_matcher_utils.py
+++ b/test/fx/test_source_matcher_utils.py
@@ -480,6 +480,56 @@ def forward(self, x):
         input_node_names = {node.name for node in module_partitions[k][0].input_nodes}
         self.assertEqual(input_node_names, {"x"})
 
+    @unittest.skipIf(not is_dynamo_supported(), "Dynamo not supported")
+    def test_module_partitioner_linear_relu_linear_nn_module_stack(self):
+        class M(torch.nn.Module):
+            def __init__(self) -> None:
+                super().__init__()
+                self.linear1 = torch.nn.Linear(3, 3)
+                self.linear2 = torch.nn.Linear(3, 3)
+                self.relu = torch.nn.ReLU()
+                self.linear3 = torch.nn.Linear(3, 5)
+
+            def forward(self, x):
+                x = self.linear1(x)
+                x = self.linear2(x)
+                x = self.relu(x)
+                x = self.linear3(x)
+                return x
+
+        inputs = (torch.randn(3, 3),)
+        # strict=False does not populate source_fn_stack, so
+        # get_source_partitions should fall back to nn_module_stack.
+        gm = torch.export.export(M(), inputs, strict=False).module()
+        gm.graph.eliminate_dead_code()
+
+        module_partitions = get_source_partitions(
+            gm.graph, [torch.nn.Linear, torch.nn.ReLU]
+        )
+
+        self.assertEqual(len(module_partitions), 2)
+        self.assertEqual(len(module_partitions[torch.nn.Linear]), 3)
+        self.assertEqual(len(module_partitions[torch.nn.ReLU]), 1)
+
+        self.assertFalse(
+            check_subgraphs_connected(
+                module_partitions[torch.nn.Linear][0],
+                module_partitions[torch.nn.ReLU][0],
+            )
+        )
+        self.assertTrue(
+            check_subgraphs_connected(
+                module_partitions[torch.nn.Linear][1],
+                module_partitions[torch.nn.ReLU][0],
+            )
+        )
+        self.assertFalse(
+            check_subgraphs_connected(
+                module_partitions[torch.nn.Linear][2],
+                module_partitions[torch.nn.ReLU][0],
+            )
+        )
+
 
 instantiate_parametrized_tests(TestSourceMatcher)
 
diff --git a/test/higher_order_ops/test_inline_asm_elementwise.py b/test/higher_order_ops/test_inline_asm_elementwise.py
new file mode 100644
index 0000000000000..07201cc666854
--- /dev/null
+++ b/test/higher_order_ops/test_inline_asm_elementwise.py
@@ -0,0 +1,733 @@
+# Owner(s): ["module: higher order operators"]
+"""
+Tests for inline_asm_elementwise higher-order operator.
+
+Tests verify:
+1. Bitwise equivalence between eager (Jiterator) and compiled (Inductor) paths
+2. Correctness via approximate comparison with reference PyTorch ops
+"""
+
+import unittest
+from collections.abc import Callable
+from dataclasses import dataclass
+
+import torch
+from torch._higher_order_ops.inline_asm_elementwise import inline_asm_elementwise
+from torch.testing._internal.common_cuda import evaluate_gfx_arch_within, SM70OrLater
+from torch.testing._internal.common_utils import (
+    instantiate_parametrized_tests,
+    MI200_ARCH,
+    MI300_ARCH,
+    NAVI_ARCH,
+    parametrize,
+    run_tests,
+    TEST_CUDA,
+    TestCase,
+)
+
+
+@dataclass
+class AsmTestCase:
+    name: str
+    input_gen_fn: Callable
+    asm_str: str
+    constraints: str
+    dtype: torch.dtype
+    approx_fn: Callable
+    pack: int = 1
+    compile_only: bool = False
+    min_sm: int = 70
+
+
+TEST_CASES = [
+    # Basic float32 operations
+    AsmTestCase(
+        "identity_f32",
+        lambda: (torch.randn(100, device="cuda", dtype=torch.float32),),
+        "v_mov_b32 $0, $1" if torch.version.hip else "mov.f32 $0, $1;",
+        "=v, v" if torch.version.hip else "=f,f",
+        torch.float32,
+        lambda x: x,
+    ),
+    AsmTestCase(
+        "add_f32",
+        lambda: (
+            torch.randn(100, device="cuda", dtype=torch.float32),
+            torch.randn(100, device="cuda", dtype=torch.float32),
+        ),
+        "v_add_f32 $0, $1, $2" if torch.version.hip else "add.f32 $0, $1, $2;",
+        "=v, v, v" if torch.version.hip else "=f,f,f",
+        torch.float32,
+        lambda x, y: x + y,
+    ),
+    AsmTestCase(
+        "mul_f32",
+        lambda: (
+            torch.randn(100, device="cuda", dtype=torch.float32),
+            torch.randn(100, device="cuda", dtype=torch.float32),
+        ),
+        "v_mul_f32 $0, $1, $2" if torch.version.hip else "mul.f32 $0, $1, $2;",
+        "=v, v, v" if torch.version.hip else "=f,f,f",
+        torch.float32,
+        lambda x, y: x * y,
+    ),
+    AsmTestCase(
+        "fma_f32",
+        lambda: (
+            torch.randn(100, device="cuda", dtype=torch.float32),
+            torch.randn(100, device="cuda", dtype=torch.float32),
+            torch.randn(100, device="cuda", dtype=torch.float32),
+        ),
+        "v_fma_f32 $0, $1, $2, $3"
+        if torch.version.hip
+        else "fma.rn.f32 $0, $1, $2, $3;",
+        "=v, v, v, v" if torch.version.hip else "=f,f,f,f",
+        torch.float32,
+        lambda a, b, c: a * b + c,
+    ),
+    # Multi-line inline asm. PTX uses curly braces; AMDGCN uses newlines.
+    AsmTestCase(
+        "double_multiline",
+        lambda: (torch.randn(100, device="cuda", dtype=torch.float32),),
+        (
+            """
+            v_mov_b32 $0, $1
+            v_add_f32 $0, $0, $1
+            """
+            if torch.version.hip
+            else "{.reg .f32 tmp; mov.f32 tmp, $1; add.f32 $0, tmp, tmp;}"
+        ),
+        "=v, v" if torch.version.hip else "=f,f",
+        torch.float32,
+        lambda x: x * 2,
+    ),
+    # bf16/fp16 upcasting (compile-only: Jiterator can't handle dtype mismatch)
+    AsmTestCase(
+        "bf16_upcast",
+        lambda: (torch.randn(100, device="cuda", dtype=torch.bfloat16),),
+        "v_add_f32 $0, $1, $1" if torch.version.hip else "add.f32 $0, $1, $1;",
+        "=v, v" if torch.version.hip else "=f,f",
+        torch.float32,
+        lambda x: x.float() * 2,
+        compile_only=True,
+        min_sm=80,
+    ),
+    AsmTestCase(
+        "fp16_upcast",
+        lambda: (torch.randn(100, device="cuda", dtype=torch.float16),),
+        "v_add_f32 $0, $1, $1" if torch.version.hip else "add.f32 $0, $1, $1;",
+        "=v, v" if torch.version.hip else "=f,f",
+        torch.float32,
+        lambda x: x.float() * 2,
+        compile_only=True,
+    ),
+    # Integer operations
+    AsmTestCase(
+        "bitwise_and",
+        lambda: (
+            torch.randint(0, 2**16, (100,), device="cuda", dtype=torch.int32),
+            torch.randint(0, 2**16, (100,), device="cuda", dtype=torch.int32),
+        ),
+        "v_and_b32 $0, $1, $2" if torch.version.hip else "and.b32 $0, $1, $2;",
+        "=v, v, v" if torch.version.hip else "=r,r,r",
+        torch.int32,
+        lambda x, y: x & y,
+    ),
+    AsmTestCase(
+        "bitwise_or",
+        lambda: (
+            torch.randint(0, 2**16, (100,), device="cuda", dtype=torch.int32),
+            torch.randint(0, 2**16, (100,), device="cuda", dtype=torch.int32),
+        ),
+        "v_or_b32 $0, $1, $2" if torch.version.hip else "or.b32 $0, $1, $2;",
+        "=v, v, v" if torch.version.hip else "=r,r,r",
+        torch.int32,
+        lambda x, y: x | y,
+    ),
+    # Output dtype differs from input (compile-only: Jiterator returns input dtype)
+    # AMDGCN: v_bfe_u32 (bit-field extract) replaces PTX's multi-instruction
+    # shift-and-mask sequence in a single instruction.
+    AsmTestCase(
+        "exponent_extract",
+        lambda: (
+            torch.tensor([1.0, 2.0, 0.5, 16.0], device="cuda", dtype=torch.float32),
+        ),
+        (
+            "v_bfe_u32 $0, $1, 23, 8"
+            if torch.version.hip
+            else "{.reg .b32 t; mov.b32 t,$1; shr.u32 t,t,23; and.b32 $0,t,0xFF;}"
+        ),
+        "=v, v" if torch.version.hip else "=r,f",
+        torch.int32,
+        lambda x: ((x.view(torch.int32) >> 23) & 0xFF).to(torch.int32),
+        compile_only=True,
+    ),
+    # Truncate u32 -> u16 (compile-only).
+    # PTX: uses "h" (16-bit) output / "r" (32-bit) input constraints.
+    # AMDGCN: VGPRs are always 32-bit (no "h" equivalent), so we use "v"
+    # and extract the lower 16 bits via v_bfe_u32.
+    AsmTestCase(
+        "truncate_to_uint16",
+        lambda: (torch.randint(0, 256, (100,), device="cuda", dtype=torch.int32),),
+        "v_bfe_u32 $0, $1, 0, 16" if torch.version.hip else "cvt.u16.u32 $0, $1;",
+        "=v, v" if torch.version.hip else "=h,r",
+        torch.uint16,
+        lambda x: x.to(torch.uint16),
+        compile_only=True,
+    ),
+    # Broadcasting
+    AsmTestCase(
+        "broadcast_add",
+        lambda: (
+            torch.randn(4, 1, device="cuda", dtype=torch.float32),
+            torch.randn(1, 8, device="cuda", dtype=torch.float32),
+        ),
+        "v_add_f32 $0, $1, $2" if torch.version.hip else "add.f32 $0, $1, $2;",
+        "=v, v, v" if torch.version.hip else "=f,f,f",
+        torch.float32,
+        lambda x, y: x + y,
+    ),
+    # Non-contiguous
+    AsmTestCase(
+        "noncontiguous",
+        lambda: (torch.randn(8, 16, device="cuda", dtype=torch.float32).t(),),
+        "v_mov_b32 $0, $1" if torch.version.hip else "mov.f32 $0, $1;",
+        "=v, v" if torch.version.hip else "=f,f",
+        torch.float32,
+        lambda x: x,
+    ),
+    # fp16/bf16 native asm (compile-only: inductor computes in fp32, needs downcast)
+    # ROCm: Inductor feeds f32 values (upcasted for computation).  AMDGCN has no
+    # "h" constraint for 16-bit regs, so we add in f32 and convert to the target
+    # format.  PTX "h" constraints tell Triton to downcast before the asm.
+    AsmTestCase(
+        "add_fp16_native",
+        lambda: (
+            torch.randn(100, device="cuda", dtype=torch.float16),
+            torch.randn(100, device="cuda", dtype=torch.float16),
+        ),
+        (
+            "v_add_f32 $0, $1, $2\nv_cvt_f16_f32 $0, $0"
+            if torch.version.hip
+            else "add.f16 $0, $1, $2;"
+        ),
+        "=v,v,v" if torch.version.hip else "=h,h,h",
+        torch.float16,
+        lambda x, y: x + y,
+        compile_only=True,
+    ),
+    # AMDGCN: v_cvt_pk_bf16_f32 packs two f32 values into bf16 in a single
+    # 32-bit register.  We pass $0 twice — only the lower 16 bits (first
+    # bf16 slot) are used by Triton.
+    AsmTestCase(
+        "add_bf16_native",
+        lambda: (
+            torch.randn(100, device="cuda", dtype=torch.bfloat16),
+            torch.randn(100, device="cuda", dtype=torch.bfloat16),
+        ),
+        (
+            "v_add_f32 $0, $1, $2\nv_cvt_pk_bf16_f32 $0, $0, $0"
+            if torch.version.hip
+            else "add.bf16 $0, $1, $2;"
+        ),
+        "=v,v,v" if torch.version.hip else "=h,h,h",
+        torch.bfloat16,
+        lambda x, y: x + y,
+        compile_only=True,
+        min_sm=90,
+    ),
+    # pack=2: each asm invocation processes 2 elements (compile-only)
+    AsmTestCase(
+        "identity_pack2",
+        lambda: (torch.randn(128, device="cuda", dtype=torch.float32),),
+        (
+            """
+            v_mov_b32 $0, $2
+            v_mov_b32 $1, $3
+            """
+            if torch.version.hip
+            else "mov.b32 $0, $2; mov.b32 $1, $3;"
+        ),
+        "=v,=v,v,v" if torch.version.hip else "=r,=r,r,r",
+        torch.float32,
+        lambda x: x,
+        pack=2,
+        compile_only=True,
+    ),
+    AsmTestCase(
+        "add_pack2",
+        lambda: (
+            torch.randn(128, device="cuda", dtype=torch.float32),
+            torch.randn(128, device="cuda", dtype=torch.float32),
+        ),
+        (
+            """
+            v_add_f32 $0, $2, $4
+            v_add_f32 $1, $3, $5
+            """
+            if torch.version.hip
+            else "add.f32 $0, $2, $4; add.f32 $1, $3, $5;"
+        ),
+        "=v,=v,v,v,v,v" if torch.version.hip else "=f,=f,f,f,f,f",
+        torch.float32,
+        lambda x, y: x + y,
+        pack=2,
+        compile_only=True,
+    ),
+]
+TEST_CASE_NAMES = [tc.name for tc in TEST_CASES]
+
+
+@unittest.skipIf(not TEST_CUDA, "CUDA not available")
+@unittest.skipIf(not SM70OrLater, "Requires SM70+")
+@instantiate_parametrized_tests
+class TestInlineAsmElementwise(TestCase):
+    """Parametrized tests for inline_asm_elementwise."""
+
+    @parametrize(
+        "case_idx", list(range(len(TEST_CASES))), name_fn=lambda i: TEST_CASE_NAMES[i]
+    )
+    def test_eager_vs_compiled_bitwise(self, case_idx):
+        """Verify eager and compiled produce bitwise identical results."""
+        tc = TEST_CASES[case_idx]
+        if not torch.version.hip and torch.cuda.get_device_capability() < (
+            tc.min_sm // 10,
+            tc.min_sm % 10,
+        ):
+            self.skipTest(f"Requires SM{tc.min_sm}+")
+
+        # Native bf16 conversion instruction not available before gfx950.
+        if (
+            torch.version.hip
+            and tc.name == "add_bf16_native"
+            and evaluate_gfx_arch_within(
+                [
+                    *MI200_ARCH,
+                    *MI300_ARCH,
+                    *NAVI_ARCH,
+                ]
+            )
+        ):
+            self.skipTest("Requires gfx950+")
+
+        inputs = tc.input_gen_fn()
+
+        def fn(*args):
+            return inline_asm_elementwise(
+                *args,
+                asm_str=tc.asm_str,
+                constraints=tc.constraints,
+                dtype=tc.dtype,
+                pack=tc.pack,
+            )
+
+        torch._dynamo.reset()
+        compiled_result = torch.compile(fn, backend="inductor")(*inputs)
+
+        if tc.compile_only:
+            expected = tc.approx_fn(*inputs)
+            self.assertEqual(
+                compiled_result.float(), expected.float(), atol=1e-5, rtol=1e-5
+            )
+        else:
+            eager_result = fn(*inputs)
+            self.assertEqual(eager_result, compiled_result)
+
+    @parametrize(
+        "case_idx", list(range(len(TEST_CASES))), name_fn=lambda i: TEST_CASE_NAMES[i]
+    )
+    def test_correctness(self, case_idx):
+        """Verify result matches reference function."""
+        tc = TEST_CASES[case_idx]
+        if not torch.version.hip and torch.cuda.get_device_capability() < (
+            tc.min_sm // 10,
+            tc.min_sm % 10,
+        ):
+            self.skipTest(f"Requires SM{tc.min_sm}+")
+
+        # Native bf16 conversion instruction not available before gfx950.
+        if (
+            torch.version.hip
+            and tc.name == "add_bf16_native"
+            and evaluate_gfx_arch_within(
+                [
+                    *MI200_ARCH,
+                    *MI300_ARCH,
+                    *NAVI_ARCH,
+                ]
+            )
+        ):
+            self.skipTest("Requires gfx950+")
+
+        inputs = tc.input_gen_fn()
+
+        def fn(*args):
+            return inline_asm_elementwise(
+                *args,
+                asm_str=tc.asm_str,
+                constraints=tc.constraints,
+                dtype=tc.dtype,
+                pack=tc.pack,
+            )
+
+        if tc.compile_only:
+            torch._dynamo.reset()
+            result = torch.compile(fn, backend="inductor")(*inputs)
+        else:
+            result = fn(*inputs)
+        expected = tc.approx_fn(*inputs)
+
+        self.assertEqual(result.float(), expected.float(), atol=1e-5, rtol=1e-5)
+
+
+@unittest.skipIf(not TEST_CUDA, "CUDA not available")
+class TestInlineAsmElementwiseErrors(TestCase):
+    """Tests for error handling."""
+
+    def test_error_no_inputs(self):
+        with self.assertRaises(ValueError):
+            inline_asm_elementwise(
+                asm_str="v_mov_b32 $0, 1.0"
+                if torch.version.hip
+                else "mov.f32 $0, 1.0;",
+                constraints="=v" if torch.version.hip else "=f",
+                dtype=torch.float32,
+            )
+
+    def test_error_constraint_mismatch(self):
+        x = torch.randn(100, device="cuda", dtype=torch.float32)
+        y = torch.randn(100, device="cuda", dtype=torch.float32)
+        with self.assertRaises(ValueError):
+            inline_asm_elementwise(
+                x,
+                y,
+                asm_str="v_add_f32 $0, $1, $2"
+                if torch.version.hip
+                else "add.f32 $0, $1, $2;",
+                constraints="=v,v" if torch.version.hip else "=f,f",
+                dtype=torch.float32,
+            )
+
+    def test_error_mixed_dtypes(self):
+        x = torch.randn(100, device="cuda", dtype=torch.float32)
+        y = torch.randint(0, 10, (100,), device="cuda", dtype=torch.int32)
+        with self.assertRaises(ValueError):
+            inline_asm_elementwise(
+                x,
+                y,
+                asm_str="v_add_f32 $0, $1, $2"
+                if torch.version.hip
+                else "add.f32 $0, $1, $2;",
+                constraints="=v,v,v" if torch.version.hip else "=f,f,r",
+                dtype=torch.float32,
+            )
+
+    def test_error_cpu_tensor(self):
+        x = torch.randn(100, dtype=torch.float32)
+        with self.assertRaises(RuntimeError):
+            inline_asm_elementwise(
+                x,
+                asm_str="v_mov_b32 $0, $1" if torch.version.hip else "mov.f32 $0, $1;",
+                constraints="=v,v" if torch.version.hip else "=f,f",
+                dtype=torch.float32,
+            )
+
+
+@unittest.skipIf(not TEST_CUDA, "CUDA not available")
+@unittest.skipIf(not SM70OrLater, "Requires SM70+")
+class TestInlineAsmElementwiseEdgeCases(TestCase):
+    """Tests for edge cases."""
+
+    def test_empty_tensor(self):
+        x = torch.empty(0, device="cuda", dtype=torch.float32)
+        result = inline_asm_elementwise(
+            x,
+            asm_str="v_mov_b32 $0, $1" if torch.version.hip else "mov.f32 $0, $1;",
+            constraints="=v, v" if torch.version.hip else "=f,f",
+            dtype=torch.float32,
+        )
+        self.assertEqual(result.shape, torch.Size([0]))
+
+    def test_scalar_tensor(self):
+        x = torch.tensor(3.14, device="cuda", dtype=torch.float32)
+        result = inline_asm_elementwise(
+            x,
+            asm_str="v_mov_b32 $0, $1" if torch.version.hip else "mov.f32 $0, $1;",
+            constraints="=v, v" if torch.version.hip else "=f,f",
+            dtype=torch.float32,
+        )
+        self.assertEqual(result.shape, torch.Size([]))
+        self.assertEqual(result, x)
+
+    def test_4d_tensor(self):
+        x = torch.randn(2, 3, 4, 5, device="cuda", dtype=torch.float32)
+        result = inline_asm_elementwise(
+            x,
+            asm_str="v_mov_b32 $0, $1" if torch.version.hip else "mov.f32 $0, $1;",
+            constraints="=v, v" if torch.version.hip else "=f,f",
+            dtype=torch.float32,
+        )
+        self.assertEqual(result.shape, x.shape)
+        self.assertEqual(result, x)
+
+    def test_composition_with_pytorch_ops(self):
+        def fn(x, y):
+            z = x * 2
+            w = inline_asm_elementwise(
+                z,
+                y,
+                asm_str="v_add_f32 $0, $1, $2"
+                if torch.version.hip
+                else "add.f32 $0, $1, $2;",
+                constraints="=v, v, v" if torch.version.hip else "=f,f,f",
+                dtype=torch.float32,
+            )
+            return w + 1.0
+
+        x = torch.randn(100, device="cuda", dtype=torch.float32)
+        y = torch.randn(100, device="cuda", dtype=torch.float32)
+
+        eager_result = fn(x, y)
+        compiled_fn = torch.compile(fn, backend="inductor")
+        compiled_result = compiled_fn(x, y)
+
+        self.assertEqual(eager_result, compiled_result)
+        self.assertEqual(eager_result, x * 2 + y + 1.0)
+
+    def test_output_strides_mixed_inputs(self):
+        """Verify fake mode output strides match eager (TensorIterator) strides."""
+        from torch._subclasses.fake_tensor import FakeTensorMode
+
+        # Two inputs with different strides: one contiguous, one transposed.
+        # This exercises TensorIterator's slow path for stride computation.
+        x = torch.randn(8, 16, device="cuda", dtype=torch.float32)
+        y = torch.randn(16, 8, device="cuda", dtype=torch.float32).t()
+
+        eager_result = inline_asm_elementwise(
+            x,
+            y,
+            asm_str="v_add_f32 $0, $1, $2"
+            if torch.version.hip
+            else "add.f32 $0, $1, $2;",
+            constraints="=v, v, v" if torch.version.hip else "=f,f,f",
+            dtype=torch.float32,
+        )
+
+        with FakeTensorMode() as mode:
+            fake_x = mode.from_tensor(x)
+            fake_y = mode.from_tensor(y)
+            fake_result = inline_asm_elementwise(
+                fake_x,
+                fake_y,
+                asm_str="v_add_f32 $0, $1, $2"
+                if torch.version.hip
+                else "add.f32 $0, $1, $2;",
+                constraints="=v, v, v" if torch.version.hip else "=f,f,f",
+                dtype=torch.float32,
+            )
+
+        self.assertEqual(eager_result.shape, fake_result.shape)
+        self.assertEqual(eager_result.stride(), fake_result.stride())
+
+    def test_dynamic_shapes(self):
+        def fn(x, y):
+            return inline_asm_elementwise(
+                x,
+                y,
+                asm_str="v_add_f32 $0, $1, $2"
+                if torch.version.hip
+                else "add.f32 $0, $1, $2;",
+                constraints="=v, v, v" if torch.version.hip else "=f,f,f",
+                dtype=torch.float32,
+            )
+
+        compiled_fn = torch.compile(fn, backend="inductor", dynamic=True)
+
+        for size in [50, 100, 200]:
+            x = torch.randn(size, device="cuda", dtype=torch.float32)
+            y = torch.randn(size, device="cuda", dtype=torch.float32)
+            eager_result = fn(x, y)
+            compiled_result = compiled_fn(x, y)
+            self.assertEqual(eager_result, compiled_result)
+
+
+@unittest.skipIf(not TEST_CUDA, "CUDA not available")
+@unittest.skipIf(not SM70OrLater, "Requires SM70+")
+class TestInlineAsmPackPadding(TestCase):
+    """Test that pack padding works when block size < pack."""
+
+    def test_pack2_xblock1_padding(self):
+        """Force XBLOCK=1 with pack=2 so padding is needed."""
+        from torch._inductor.choices import InductorChoices
+        from torch._inductor.codegen.triton import FixedTritonConfig
+        from torch._inductor.utils import run_and_get_code
+        from torch.testing import FileCheck
+
+        class ForceXBlock1(InductorChoices):
+            def triton_kernel_kwargs(self, kernel_cls, features, groups, kernel_kwargs):
+                return {
+                    **kernel_kwargs,
+                    "fixed_config": FixedTritonConfig({"XBLOCK": 1}),
+                }
+
+        def fn(x):
+            return inline_asm_elementwise(
+                x,
+                asm_str=(
+                    """
+                    v_mov_b32 $0, $2
+                    v_mov_b32 $1, $3
+                    """
+                    if torch.version.hip
+                    else "mov.b32 $0, $2; mov.b32 $1, $3;"
+                ),
+                constraints="=v,=v,v,v" if torch.version.hip else "=r,=r,r,r",
+                dtype=torch.float32,
+                pack=2,
+            )
+
+        x = torch.randn(128, device="cuda", dtype=torch.float32)
+        with torch._inductor.virtualized.V.set_choices_handler(ForceXBlock1()):
+            torch._dynamo.reset()
+            result, (code,) = run_and_get_code(torch.compile(fn, backend="inductor"), x)
+
+        self.assertEqual(result, x)
+        # Verify padding helpers are emitted in the generated code
+        FileCheck().check("inline_asm_pack").check("inline_asm_unpack").run(code)
+
+    def test_pack4_xblock1_padding(self):
+        """Force XBLOCK=1 with pack=4 so padding is needed."""
+        from torch._inductor.choices import InductorChoices
+        from torch._inductor.codegen.triton import FixedTritonConfig
+        from torch._inductor.utils import run_and_get_code
+        from torch.testing import FileCheck
+
+        class ForceXBlock1(InductorChoices):
+            def triton_kernel_kwargs(self, kernel_cls, features, groups, kernel_kwargs):
+                return {
+                    **kernel_kwargs,
+                    "fixed_config": FixedTritonConfig({"XBLOCK": 1}),
+                }
+
+        def fn(x):
+            return inline_asm_elementwise(
+                x,
+                asm_str=(
+                    """
+                    v_mov_b32 $0, $4
+                    v_mov_b32 $1, $5
+                    v_mov_b32 $2, $6
+                    v_mov_b32 $3, $7
+                    """
+                    if torch.version.hip
+                    else "mov.b32 $0, $4; mov.b32 $1, $5; mov.b32 $2, $6; mov.b32 $3, $7;"
+                ),
+                constraints=(
+                    "=v,=v,=v,=v,v,v,v,v"
+                    if torch.version.hip
+                    else "=r,=r,=r,=r,r,r,r,r"
+                ),
+                dtype=torch.float32,
+                pack=4,
+            )
+
+        x = torch.randn(128, device="cuda", dtype=torch.float32)
+        with torch._inductor.virtualized.V.set_choices_handler(ForceXBlock1()):
+            torch._dynamo.reset()
+            result, (code,) = run_and_get_code(torch.compile(fn, backend="inductor"), x)
+
+        self.assertEqual(result, x)
+        FileCheck().check("inline_asm_pack").check("inline_asm_unpack").run(code)
+
+    def test_pack4_xblock2_partial_padding(self):
+        """XBLOCK=2 < pack=4, so 1 round of padding is needed (not 2)."""
+        from torch._inductor.choices import InductorChoices
+        from torch._inductor.codegen.triton import FixedTritonConfig
+        from torch._inductor.utils import run_and_get_code
+        from torch.testing import FileCheck
+
+        class ForceXBlock2(InductorChoices):
+            def triton_kernel_kwargs(self, kernel_cls, features, groups, kernel_kwargs):
+                return {
+                    **kernel_kwargs,
+                    "fixed_config": FixedTritonConfig({"XBLOCK": 2}),
+                }
+
+        def fn(x):
+            return inline_asm_elementwise(
+                x,
+                asm_str=(
+                    """
+                    v_mov_b32 $0, $4
+                    v_mov_b32 $1, $5
+                    v_mov_b32 $2, $6
+                    v_mov_b32 $3, $7
+                    """
+                    if torch.version.hip
+                    else "mov.b32 $0, $4; mov.b32 $1, $5; mov.b32 $2, $6; mov.b32 $3, $7;"
+                ),
+                constraints=(
+                    "=v,=v,=v,=v,v,v,v,v"
+                    if torch.version.hip
+                    else "=r,=r,=r,=r,r,r,r,r"
+                ),
+                dtype=torch.float32,
+                pack=4,
+            )
+
+        x = torch.randn(128, device="cuda", dtype=torch.float32)
+        with torch._inductor.virtualized.V.set_choices_handler(ForceXBlock2()):
+            torch._dynamo.reset()
+            result, (code,) = run_and_get_code(torch.compile(fn, backend="inductor"), x)
+
+        self.assertEqual(result, x)
+        FileCheck().check("inline_asm_pack").check("inline_asm_unpack").run(code)
+
+    def test_pack2_xblock1_yblock1_padding(self):
+        """Force XBLOCK=1, YBLOCK=1 with pack=2 on a 2D-tiled kernel."""
+        from torch._inductor.choices import InductorChoices
+        from torch._inductor.codegen.triton import FixedTritonConfig
+        from torch._inductor.utils import run_and_get_code
+        from torch.testing import FileCheck
+
+        class ForceXY1(InductorChoices):
+            def triton_kernel_kwargs(self, kernel_cls, features, groups, kernel_kwargs):
+                return {
+                    **kernel_kwargs,
+                    "fixed_config": FixedTritonConfig({"XBLOCK": 1, "YBLOCK": 1}),
+                }
+
+        def fn(x, y):
+            return inline_asm_elementwise(
+                x,
+                y,
+                asm_str=(
+                    """
+                    v_add_f32 $0, $2, $4
+                    v_add_f32 $1, $3, $5
+                    """
+                    if torch.version.hip
+                    else "add.f32 $0, $2, $4; add.f32 $1, $3, $5;"
+                ),
+                constraints="=v,=v,v,v,v,v" if torch.version.hip else "=f,=f,f,f,f,f",
+                dtype=torch.float32,
+                pack=2,
+            )
+
+        x = torch.randn(8, 16, device="cuda", dtype=torch.float32)
+        # Transposed input triggers 2D tiling (different stride patterns)
+        y = torch.randn(16, 8, device="cuda", dtype=torch.float32).T
+        with torch._inductor.virtualized.V.set_choices_handler(ForceXY1()):
+            torch._dynamo.reset()
+            result, (code,) = run_and_get_code(
+                torch.compile(fn, backend="inductor"), x, y
+            )
+
+        self.assertEqual(result, x + y)
+        FileCheck().check("YBLOCK").check("inline_asm_pack").check(
+            "inline_asm_unpack"
+        ).run(code)
+
+
+if __name__ == "__main__":
+    run_tests()
diff --git a/test/higher_order_ops/test_invoke_subgraph.py b/test/higher_order_ops/test_invoke_subgraph.py
index 97d5434c79e0d..708ffc54fa6d6 100644
--- a/test/higher_order_ops/test_invoke_subgraph.py
+++ b/test/higher_order_ops/test_invoke_subgraph.py
@@ -2,6 +2,7 @@
 # flake8: noqa: B950
 # flake8: noqa: E731
 
+import contextlib
 import unittest
 import unittest.mock as mock
 
@@ -339,11 +340,11 @@ def forward(self, L_x_: "f32[8]", L_y_: "f32[8]", L_mod_buffers_buf_: "f32[8]"):
 
         subgraph_0 = self.subgraph_0
         invoke_subgraph = torch.ops.higher_order.invoke_subgraph(subgraph_0, 'subgraph_0', l_mod_buffers_buf_, l_x_, l_y_);  subgraph_0 = None
-        getitem_8: "f32[8]" = invoke_subgraph[0];  invoke_subgraph = None
+        getitem: "f32[8]" = invoke_subgraph[0];  invoke_subgraph = None
         subgraph_1 = self.subgraph_0
         invoke_subgraph_1 = torch.ops.higher_order.invoke_subgraph(subgraph_1, 'subgraph_0', l_mod_buffers_buf_, l_x_, l_y_);  subgraph_1 = l_mod_buffers_buf_ = l_x_ = l_y_ = None
-        getitem_9: "f32[8]" = invoke_subgraph_1[0];  invoke_subgraph_1 = None
-        add: "f32[8]" = getitem_8 + getitem_9;  getitem_8 = getitem_9 = None
+        getitem_1: "f32[8]" = invoke_subgraph_1[0];  invoke_subgraph_1 = None
+        add: "f32[8]" = getitem + getitem_1;  getitem = getitem_1 = None
         return (add,)
 
     class subgraph_0(torch.nn.Module):
@@ -1064,16 +1065,16 @@ def forward(self, L_x_: "f32[8]"):
 
         subgraph_0 = self.subgraph_0
         invoke_subgraph = torch.ops.higher_order.invoke_subgraph(subgraph_0, 'subgraph_0', l_x_);  subgraph_0 = l_x_ = None
-        getitem_2: "f32[8]" = invoke_subgraph[0];  invoke_subgraph = None
-        detach: "f32[8]" = getitem_2.detach();  getitem_2 = None
+        getitem: "f32[8]" = invoke_subgraph[0];  invoke_subgraph = None
+        detach: "f32[8]" = getitem.detach();  getitem = None
         return (detach,)
 
     class subgraph_0(torch.nn.Module):
         def forward(self, l_x_: "f32[8]"):
             wrap_body_0 = self.wrap_body_0
             tag_activation_checkpoint = torch.ops.higher_order.tag_activation_checkpoint(wrap_body_0, l_x_, use_reentrant = False);  wrap_body_0 = l_x_ = None
-            getitem_2: "f32[8]" = tag_activation_checkpoint[0];  tag_activation_checkpoint = None
-            return (getitem_2,)
+            getitem: "f32[8]" = tag_activation_checkpoint[0];  tag_activation_checkpoint = None
+            return (getitem,)
 
         class wrap_body_0(torch.nn.Module):
             def forward(self, l_x_: "f32[8]"):
@@ -1130,12 +1131,12 @@ def forward(self, L_x_: "f32[8]", L_y_: "f32[8]"):
 
         subgraph_0 = self.subgraph_0
         invoke_subgraph = torch.ops.higher_order.invoke_subgraph(subgraph_0, 'subgraph_0', l_x_, l_y_);  subgraph_0 = l_x_ = None
-        getitem_4: "f32[8]" = invoke_subgraph[0];  invoke_subgraph = None
+        a: "f32[8]" = invoke_subgraph[0];  invoke_subgraph = None
 
         subgraph_1 = self.subgraph_1
-        invoke_subgraph_1 = torch.ops.higher_order.invoke_subgraph(subgraph_1, 'subgraph_1', getitem_4, l_y_);  subgraph_1 = getitem_4 = l_y_ = None
-        getitem_5: "f32[8]" = invoke_subgraph_1[0];  invoke_subgraph_1 = None
-        return (getitem_5,)
+        invoke_subgraph_1 = torch.ops.higher_order.invoke_subgraph(subgraph_1, 'subgraph_1', a, l_y_);  subgraph_1 = a = l_y_ = None
+        getitem_1: "f32[8]" = invoke_subgraph_1[0];  invoke_subgraph_1 = None
+        return (getitem_1,)
 
     class subgraph_0(torch.nn.Module):
         def forward(self, l_x_: "f32[8]", l_y_: "f32[8]"):
@@ -1252,20 +1253,20 @@ def forward(self, L_x_: "f32[8]", L_y_: "f32[8]"):
 
         subgraph_0 = self.subgraph_0
         invoke_subgraph = torch.ops.higher_order.invoke_subgraph(subgraph_0, 'subgraph_0', l_x_, l_y_);  subgraph_0 = l_x_ = None
-        getitem_25: "f32[8]" = invoke_subgraph[0];  invoke_subgraph = None
+        getitem: "f32[8]" = invoke_subgraph[0];  invoke_subgraph = None
         subgraph_1 = self.subgraph_0
-        invoke_subgraph_1 = torch.ops.higher_order.invoke_subgraph(subgraph_1, 'subgraph_0', getitem_25, l_y_);  subgraph_1 = getitem_25 = None
-        getitem_26: "f32[8]" = invoke_subgraph_1[0];  invoke_subgraph_1 = None
+        invoke_subgraph_1 = torch.ops.higher_order.invoke_subgraph(subgraph_1, 'subgraph_0', getitem, l_y_);  subgraph_1 = getitem = None
+        x: "f32[8]" = invoke_subgraph_1[0];  invoke_subgraph_1 = None
         subgraph_2 = self.subgraph_0
-        invoke_subgraph_2 = torch.ops.higher_order.invoke_subgraph(subgraph_2, 'subgraph_0', getitem_26, l_y_);  subgraph_2 = getitem_26 = None
-        getitem_27: "f32[8]" = invoke_subgraph_2[0];  invoke_subgraph_2 = None
+        invoke_subgraph_2 = torch.ops.higher_order.invoke_subgraph(subgraph_2, 'subgraph_0', x, l_y_);  subgraph_2 = x = None
+        x_1: "f32[8]" = invoke_subgraph_2[0];  invoke_subgraph_2 = None
         subgraph_3 = self.subgraph_0
-        invoke_subgraph_3 = torch.ops.higher_order.invoke_subgraph(subgraph_3, 'subgraph_0', getitem_27, l_y_);  subgraph_3 = getitem_27 = None
-        getitem_28: "f32[8]" = invoke_subgraph_3[0];  invoke_subgraph_3 = None
+        invoke_subgraph_3 = torch.ops.higher_order.invoke_subgraph(subgraph_3, 'subgraph_0', x_1, l_y_);  subgraph_3 = x_1 = None
+        x_2: "f32[8]" = invoke_subgraph_3[0];  invoke_subgraph_3 = None
         subgraph_4 = self.subgraph_0
-        invoke_subgraph_4 = torch.ops.higher_order.invoke_subgraph(subgraph_4, 'subgraph_0', getitem_28, l_y_);  subgraph_4 = getitem_28 = l_y_ = None
-        getitem_29: "f32[8]" = invoke_subgraph_4[0];  invoke_subgraph_4 = None
-        return (getitem_29,)
+        invoke_subgraph_4 = torch.ops.higher_order.invoke_subgraph(subgraph_4, 'subgraph_0', x_2, l_y_);  subgraph_4 = x_2 = l_y_ = None
+        x_3: "f32[8]" = invoke_subgraph_4[0];  invoke_subgraph_4 = None
+        return (x_3,)
 
     class subgraph_0(torch.nn.Module):
         def forward(self, l_x_: "f32[8]", l_y_: "f32[8]"):
@@ -2468,12 +2469,12 @@ def forward(self, L_x_: "f32[5]", L_y_: "f32[5]"):
 
         subgraph_0 = self.subgraph_0
         invoke_subgraph = torch.ops.higher_order.invoke_subgraph(subgraph_0, 'subgraph_0', x, y);  subgraph_0 = x = None
-        getitem_4: "f32[5]" = invoke_subgraph[0];  invoke_subgraph = None
+        z: "f32[5]" = invoke_subgraph[0];  invoke_subgraph = None
 
-        subgraph_1 = self.subgraph_1
-        invoke_subgraph_1 = torch.ops.higher_order.invoke_subgraph(subgraph_1, 'subgraph_1', getitem_4, y);  subgraph_1 = getitem_4 = y = None
-        getitem_5: "f32[5]" = invoke_subgraph_1[0];  invoke_subgraph_1 = None
-        return (getitem_5,)
+        subgraph_1 = self.subgraph_0
+        invoke_subgraph_1 = torch.ops.higher_order.invoke_subgraph(subgraph_1, 'subgraph_0', z, y);  subgraph_1 = z = y = None
+        getitem_1: "f32[5]" = invoke_subgraph_1[0];  invoke_subgraph_1 = None
+        return (getitem_1,)
 
     class subgraph_0(torch.nn.Module):
         def forward(self, x: "f32[5]", y: "f32[5]"):
@@ -2481,15 +2482,6 @@ def forward(self, x: "f32[5]", y: "f32[5]"):
 
             triton_kernel_wrapper_mutation = torch.ops.higher_order.triton_kernel_wrapper_mutation(kernel_idx = 0, constant_args_idx = 0, grid = [(5, 1, 1)], tma_descriptor_metadata = {}, kwargs = {'in_ptr0': x, 'in_ptr1': y, 'out_ptr': o});  x = y = triton_kernel_wrapper_mutation = None
 
-            sin: "f32[5]" = o.sin();  o = None
-            return (sin,)
-
-    class subgraph_1(torch.nn.Module):
-        def forward(self, z: "f32[5]", y: "f32[5]"):
-            o: "f32[5]" = torch.zeros_like(z)
-
-            triton_kernel_wrapper_mutation = torch.ops.higher_order.triton_kernel_wrapper_mutation(kernel_idx = 0, constant_args_idx = 1, grid = [(5, 1, 1)], tma_descriptor_metadata = {}, kwargs = {'in_ptr0': z, 'in_ptr1': y, 'out_ptr': o});  z = y = triton_kernel_wrapper_mutation = None
-
             sin: "f32[5]" = o.sin();  o = None
             return (sin,)
 """,
@@ -3102,176 +3094,1322 @@ def forward(self, l_x_: "f32[64, 1]"):
 
 
 @skipIfTorchDynamo("Not a torch._dynamo test")
-@parameterized_class(
-    [
-        {"strict": False},
-        {"strict": True},
-    ],
-    class_name_func=lambda cls,
-    _,
-    params: f"{cls.__name__}{'Strict' if params['strict'] else 'Nonstrict'}",
-)
-class TestInvokeSubgraphExport(TestCase):
-    def test_simple_func(self):
+class TestInvokeSubgraphReuse(TestCase):
+    @contextlib.contextmanager
+    def _count_speculate_calls(self):
+        count = 0
+        orig = torch._dynamo.variables.higher_order_ops.speculate_subgraph_with_auto_output_flattening
+
+        def _counting(*args, **kwargs):
+            nonlocal count
+            count += 1
+            return orig(*args, **kwargs)
+
+        with mock.patch.object(
+            torch._dynamo.variables.higher_order_ops,
+            "speculate_subgraph_with_auto_output_flattening",
+            _counting,
+        ):
+            yield lambda: count
+
+    def test_subgraph_reuse_skips_tracing(self):
         @nested_compile_region
         def gn(x, y):
             return torch.mul(x, y)
 
-        class M(torch.nn.Module):
-            def forward(self, x, y):
-                x = gn(x, y)
-                x = gn(x, y)
-                return x
+        def fn(x, y):
+            a = gn(x, y)
+            b = gn(x, y)
+            c = gn(x, y)
+            return a + b + c
 
-        x = torch.randn(8, requires_grad=True)
-        y = torch.randn(8, requires_grad=True)
+        x = torch.randn(8)
+        y = torch.randn(8)
 
-        ep = torch.export.export(M(), (x, y), strict=self.strict)
-        self.assertTrue(torch.allclose(ep.module()(x, y), M()(x, y)))
-        self.assertEqual(len(list(ep.graph_module.named_modules())), 2)
+        with self._count_speculate_calls() as count:
+            torch.compile(fn, backend="aot_eager", fullgraph=True)(x, y)
 
-        self.assertExpectedInline(
-            empty_line_normalizer(
-                normalize_gm(ep.graph_module.print_readable(print_output=False))
-            ),
-            """\
-class GraphModule(torch.nn.Module):
-    def forward(self, x: "f32[8]", y: "f32[8]"):
-        repeated_subgraph0 = self.repeated_subgraph0
-        invoke_subgraph = torch.ops.higher_order.invoke_subgraph(repeated_subgraph0, 'subgraph_0', x, y);  repeated_subgraph0 = x = None
-        getitem: "f32[8]" = invoke_subgraph[0];  invoke_subgraph = None
-        repeated_subgraph0_1 = self.repeated_subgraph0
-        invoke_subgraph_1 = torch.ops.higher_order.invoke_subgraph(repeated_subgraph0_1, 'subgraph_0', getitem, y);  repeated_subgraph0_1 = getitem = y = None
-        getitem_1: "f32[8]" = invoke_subgraph_1[0];  invoke_subgraph_1 = None
-        return (getitem_1,)
-    class repeated_subgraph0(torch.nn.Module):
-        def forward(self, arg0_1: "f32[8]", arg1_1: "f32[8]"):
-            mul: "f32[8]" = torch.ops.aten.mul.Tensor(arg0_1, arg1_1);  arg0_1 = arg1_1 = None
-            return (mul,)
-""",
-        )
+        self.assertEqual(count(), 1)
 
-    def test_unbacked(self):
+    def test_subgraph_reuse_different_shapes(self):
         @nested_compile_region
-        def gn(x, y):
-            b = x.item()
-            torch._check(b < y.shape[0])
-            return y[:b].clone()
+        def gn(x):
+            return x.sin()
 
-        class M(torch.nn.Module):
-            def forward(self, x, y):
-                res = []
-                for _ in range(10):
-                    res.append(gn(x, y))
-                return torch.cat(res)
+        def fn(x, y):
+            a = gn(x)
+            b = gn(y)
+            return a.sum() + b.sum()
 
-        x = torch.tensor(4)
+        x = torch.randn(4)
         y = torch.randn(8)
 
-        ep = torch.export.export(M(), (x, y), strict=self.strict)
-        ep = ep.run_decompositions()
+        with self._count_speculate_calls() as count:
+            res = torch.compile(fn, backend="aot_eager", fullgraph=True)(x, y)
 
-        self.assertTrue(torch.allclose(ep.module()(x, y), M()(x, y)))
-        self.assertEqual(len(list(ep.graph_module.named_modules())), 2)
+        # Different shapes → two separate traces
+        self.assertEqual(count(), 2)
+        self.assertEqual(res, fn(x, y))
+
+    def test_subgraph_reuse_module(self):
+        class Mod(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.c = 5
 
-    def test_pending_unbacked(self):
-        class M(torch.nn.Module):
             @nested_compile_region
-            def gn(self, x):
-                u = x[0].item()
-                return x * u
+            def forward(self, x, y):
+                return torch.mul(x, y).sin() + self.c
 
-            def forward(self, x):
-                for _ in range(4):
-                    x = self.gn(x)
-                return x
+        mod = Mod()
 
-        ep = torch.export.export(
-            M(),
-            (torch.randn(8),),
-            strict=self.strict,
-            dynamic_shapes={"x": {0: torch.export.Dim.DYNAMIC}},
-        )
-        ep = ep.run_decompositions()
+        def fn(x, y):
+            return mod(x, y) + mod(x, y)
 
-        self.assertEqual(len(list(ep.graph_module.named_modules())), 2)
+        x = torch.randn(8, requires_grad=True)
+        y = torch.randn(8, requires_grad=True)
+        ref = fn(x, y)
 
-        ep = torch.export.export(
-            M(),
-            (torch.randn(8, requires_grad=True),),
-            strict=self.strict,
-            dynamic_shapes={"x": {0: torch.export.Dim.DYNAMIC}},
-        )
-        ep = ep.run_decompositions()
+        x_clone = x.detach().clone().requires_grad_(True)
+        y_clone = y.detach().clone().requires_grad_(True)
 
-        self.assertEqual(len(list(ep.graph_module.named_modules())), 2)
+        with self._count_speculate_calls() as count:
+            res = torch.compile(fn, backend="aot_eager", fullgraph=True)(
+                x_clone, y_clone
+            )
 
-    def test_simple_method(self):
-        class M(torch.nn.Module):
-            @nested_compile_region
-            def gn(self, x, y):
-                return torch.mul(x, y)
+        # Second call reuses the first trace
+        self.assertEqual(count(), 1)
+        ref.sum().backward()
+        res.sum().backward()
+
+        self.assertEqual(ref, res)
+        self.assertEqual(x.grad, x_clone.grad)
+        self.assertEqual(y.grad, y_clone.grad)
+
+    def test_subgraph_reuse_module_different_instances(self):
+        class Mod(torch.nn.Module):
+            def __init__(self, c):
+                super().__init__()
+                self.c = c
 
+            @nested_compile_region
             def forward(self, x, y):
-                x = self.gn(x, y)
-                x = self.gn(x, y)
-                return x
+                return torch.mul(x, y).sin() + self.c
+
+        mod1 = Mod(5)
+        mod2 = Mod(5)
+
+        def fn(x, y):
+            return mod1(x, y) + mod2(x, y)
 
         x = torch.randn(8, requires_grad=True)
         y = torch.randn(8, requires_grad=True)
+        ref = fn(x, y)
 
-        ep = torch.export.export(M(), (x, y), strict=self.strict)
-        self.assertTrue(torch.allclose(ep.module()(x, y), M()(x, y)))
-        self.assertEqual(len(list(ep.graph_module.named_modules())), 2)
+        x_clone = x.detach().clone().requires_grad_(True)
+        y_clone = y.detach().clone().requires_grad_(True)
 
-    def test_multiple_module(self):
-        b = torch.randn(8)
+        with self._count_speculate_calls() as count:
+            res = torch.compile(fn, backend="aot_eager", fullgraph=True)(
+                x_clone, y_clone
+            )
 
-        class N(torch.nn.Module):
+        # mod1 and mod2 have the same structure and c value; source replacement
+        # means only one trace is needed.
+        self.assertEqual(count(), 1)
+        ref.sum().backward()
+        res.sum().backward()
+
+        self.assertEqual(ref, res)
+        self.assertEqual(x.grad, x_clone.grad)
+        self.assertEqual(y.grad, y_clone.grad)
+
+    def test_subgraph_reuse_tuple_output(self):
+        @nested_compile_region
+        def gn(x, y):
+            return torch.sin(x), torch.cos(y)
+
+        def fn(x, y):
+            a1, a2 = gn(x, y)
+            b1, b2 = gn(x, y)
+            return a1 + b1, a2 + b2
+
+        x = torch.randn(8, requires_grad=True)
+        y = torch.randn(8, requires_grad=True)
+        ref = fn(x, y)
+
+        x_clone = x.detach().clone().requires_grad_(True)
+        y_clone = y.detach().clone().requires_grad_(True)
+
+        with self._count_speculate_calls() as count:
+            res = torch.compile(fn, backend="aot_eager", fullgraph=True)(
+                x_clone, y_clone
+            )
+
+        # Second call reuses the first trace
+        self.assertEqual(count(), 1)
+        sum(r.sum() for r in ref).backward()
+        sum(r.sum() for r in res).backward()
+
+        self.assertEqual(ref[0], res[0])
+        self.assertEqual(ref[1], res[1])
+        self.assertEqual(x.grad, x_clone.grad)
+        self.assertEqual(y.grad, y_clone.grad)
+
+    def test_subgraph_reuse_mutated_attribute(self):
+        """Reuse must be skipped when a captured attribute is mutated between calls."""
+
+        class Mod(torch.nn.Module):
             def __init__(self):
                 super().__init__()
-                self.register_buffer("buf", b)
+                self.c = 5
 
             @nested_compile_region
-            def forward(self, x, y):
-                return x * y + self.buf
+            def forward(self, x):
+                return x * self.c
 
-        class M(torch.nn.Module):
+        mod = Mod()
+
+        def fn(x):
+            a = mod(x)
+            mod.c = 10
+            b = mod(x)
+            return a + b
+
+        x = torch.randn(8)
+        # Eager: first call uses c=5, then c is set to 10, second call uses c=10.
+        # Result = x*5 + x*10 = x*15.
+        mod.c = 5
+        ref = fn(x)
+        self.assertEqual(ref, x * 15)
+
+        # Compiled should produce the same result. If reuse incorrectly
+        # fires, both calls would use c=5, giving x*10 instead of x*15.
+        mod.c = 5
+        with self._count_speculate_calls() as count:
+            res = torch.compile(fn, backend="aot_eager", fullgraph=True)(x)
+        self.assertEqual(ref, res)
+        # c=5 and c=10 are distinct constants → two separate traces
+        self.assertEqual(count(), 2)
+
+    def test_subgraph_reuse_unrelated_attr_mutation(self):
+        """Reuse should still fire when a different attribute is mutated."""
+
+        class Mod(torch.nn.Module):
             def __init__(self):
                 super().__init__()
-                self.mod_list = torch.nn.ModuleList(N() for _ in range(10))
+                self.c = 5
+                self.d = 100
 
-            def forward(self, x, y):
-                for m in self.mod_list:
-                    x = m(x, y)
-                return x
+            @nested_compile_region
+            def forward(self, x):
+                # Only reads self.c, never self.d
+                return x * self.c
 
-        x = torch.randn(8, requires_grad=True)
-        y = torch.randn(8, requires_grad=True)
+        mod = Mod()
 
-        ep = torch.export.export(M(), (x, y), strict=self.strict)
-        self.assertTrue(torch.allclose(ep.module()(x, y), M()(x, y)))
-        self.assertEqual(len(list(ep.graph_module.named_modules())), 2)
+        def fn(x):
+            a = mod(x)
+            mod.d = 999  # unrelated attribute
+            b = mod(x)
+            return a + b
 
+        x = torch.randn(8)
+        ref = fn(x)
+        # Both calls use c=5, so result = x*5 + x*5 = x*10.
+        self.assertEqual(ref, x * 10)
+
+        with self._count_speculate_calls() as count:
+            res = torch.compile(fn, backend="aot_eager", fullgraph=True)(x)
+
+        self.assertEqual(ref, res)
+        # Mutating mod.d should not prevent reuse of the subgraph that only reads mod.c.
+        self.assertEqual(count(), 1)
+
+    def test_subgraph_reuse_same_class_attr_mutated(self):
+        """Reuse must be skipped when a captured attr changes between calls.
+
+        submod1 and submod2 are instances of the same class with the same
+        initial value for .c.  The first call traces submod1; the second call
+        to submod2 could reuse the cache entry via source replacement.  But
+        submod2.c is mutated between the two calls, so reuse must be skipped.
+        """
+
+        class Block(torch.nn.Module):
+            def __init__(self, c):
+                super().__init__()
+                self.c = c
+
+            @nested_compile_region
+            def forward(self, x):
+                return x * self.c
+
+        submod1 = Block(5)
+        submod2 = Block(5)  # same initial .c as submod1
+
+        def fn(x):
+            a = submod1(x)  # traces with c=5
+            submod2.c = 10  # mutate submod2.c
+            b = submod2(x)  # must NOT reuse the c=5 subgraph
+            return a + b
+
+        x = torch.randn(8)
+        ref = fn(x)
+        # a = x*5, b = x*10 → x*15
+        self.assertEqual(ref, x * 15)
+
+        submod2.c = 5  # reset for compiled run
+        with self._count_speculate_calls() as count:
+            res = torch.compile(fn, backend="aot_eager", fullgraph=True)(x)
+        self.assertEqual(ref, res)
+        # submod1 traces with c=5; submod2 has c mutated to 10 → two separate traces
+        self.assertEqual(count(), 2)
+
+    def test_subgraph_reuse_pre_existing_attr_guard(self):
+        """Guards installed before the subgraph trace must still block incorrect reuse.
+
+        If ``block.c`` is read in a conditional before the nested compile region,
+        its guard is installed before ``guards_before`` is snapshotted and won't
+        appear in the delta.  Reuse must still be rejected when a different module
+        with a different ``c`` is passed.
+        """
+
+        class Block(torch.nn.Module):
+            def __init__(self, c):
+                super().__init__()
+                self.c = c
 
-class NegativeTesting(TestCase):
-    def test_graph_break(self):
         @nested_compile_region
-        def gn(x):
-            torch._dynamo.graph_break()
-            return torch.cos(x)
+        def apply_block(mod, x):
+            return x * mod.c
+
+        block1 = Block(5)
+        block2 = Block(10)
 
         def fn(x):
-            return gn(x)
+            # The conditional installs EQUALS_MATCH on block1.c *before*
+            # the subgraph trace snapshots guards_before.
+            if block1.c == 5:
+                a = apply_block(block1, x)
+            else:
+                a = x
+            if block2.c == 10:
+                b = apply_block(block2, x)
+            else:
+                b = x
+            return a + b
 
-        x = torch.randn(8, 8, requires_grad=True)
+        x = torch.randn(8)
+        ref = fn(x)
+        self.assertEqual(ref, x * 5 + x * 10)
 
-        with self.assertRaisesRegex(
+        with self._count_speculate_calls() as count:
+            res = torch.compile(fn, backend="aot_eager", fullgraph=True)(x)
+        self.assertEqual(ref, res)
+        # block1.c=5 and block2.c=10 differ → two separate traces
+        self.assertEqual(count(), 2)
+
+    def test_subgraph_reuse_mutated_captured_variable(self):
+        """Reuse must be skipped when a captured (non-input) variable is mutated."""
+
+        class Config:
+            def __init__(self, c):
+                self.c = c
+
+        cfg = Config(5)
+
+        @nested_compile_region
+        def apply(x):
+            # cfg is captured from closure, not an explicit input
+            return x * cfg.c
+
+        def fn(x):
+            a = apply(x)
+            cfg.c = 10
+            b = apply(x)
+            return a + b
+
+        x = torch.randn(8)
+        cfg.c = 5
+        ref = fn(x)
+        self.assertEqual(ref, x * 15)
+
+        cfg.c = 5
+        with self._count_speculate_calls() as count:
+            res = torch.compile(fn, backend="aot_eager", fullgraph=True)(x)
+        self.assertEqual(ref, res)
+        # cfg.c=5 and cfg.c=10 are distinct → two separate traces
+        self.assertEqual(count(), 2)
+
+    def test_subgraph_reuse_synthetic_source(self):
+        """Reuse must handle TorchScriptObjectVariable with SyntheticLocalSource.
+
+        Hoisted opaque value types get a SyntheticLocalSource that can't be
+        resolved via VariableBuilder. On cache hit, stamp_out_subgraph must
+        call synthetic_graph_input to create a fresh graph input.
+        """
+        from test_opaque_obj_v2 import HoistedString, op_with_string
+
+        @nested_compile_region
+        def gn(x):
+            return op_with_string(x, HoistedString("double"))
+
+        def fn(x):
+            a = gn(x)
+            b = gn(x)
+            return a + b
+
+        x = torch.randn(8)
+        ref = fn(x)
+
+        with self._count_speculate_calls() as count:
+            res = torch.compile(fn, backend="aot_eager", fullgraph=True)(x)
+
+        self.assertEqual(ref, res)
+        self.assertEqual(count(), 1)
+
+    def test_subgraph_reuse_synthetic_source_different_args(self):
+        """Reuse when hoisted opaque ctor args differ across submodules."""
+        from test_opaque_obj_v2 import HoistedString, op_with_string
+
+        class Layer(torch.nn.Module):
+            def __init__(self, name):
+                super().__init__()
+                self.name = name
+
+            @nested_compile_region
+            def forward(self, x):
+                return op_with_string(x, HoistedString(self.name))
+
+        class Model(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.layer0 = Layer("double")
+                self.layer1 = Layer("square")
+                self.layer2 = Layer("double")
+
+            def forward(self, x):
+                x = self.layer0(x)
+                x = self.layer1(x)
+                x = self.layer2(x)
+                return x
+
+        model = Model()
+        x = torch.randn(8)
+        ref = model(x)
+
+        backend = EagerAndRecordGraphs()
+        with self._count_speculate_calls() as count:
+            res = torch.compile(model, backend=backend, fullgraph=True)(x)
+
+        self.assertEqual(ref, res)
+        self.assertEqual(count(), 1)
+
+        self.assertEqual(len(backend.graphs), 1)
+        if not TEST_WITH_CROSSREF:
+            self.assertExpectedInline(
+                normalize_gm(backend.graphs[0].print_readable(print_output=False)),
+                """\
+class GraphModule(torch.nn.Module):
+    def forward(self, L_x_: "f32[8]", SYNTHETIC_LOCAL_tmp_0_ : test_opaque_obj_v2_HoistedString, SYNTHETIC_LOCAL_tmp_2_ : test_opaque_obj_v2_HoistedString, SYNTHETIC_LOCAL_tmp_4_ : test_opaque_obj_v2_HoistedString):
+        l_x_ = L_x_
+        synthetic_local_tmp_0_ = SYNTHETIC_LOCAL_tmp_0_
+        synthetic_local_tmp_2_ = SYNTHETIC_LOCAL_tmp_2_
+        synthetic_local_tmp_4_ = SYNTHETIC_LOCAL_tmp_4_
+
+        subgraph_0 = self.subgraph_0
+        invoke_subgraph = torch.ops.higher_order.invoke_subgraph(subgraph_0, 'subgraph_0', l_x_, synthetic_local_tmp_0_);  subgraph_0 = l_x_ = synthetic_local_tmp_0_ = None
+        x: "f32[8]" = invoke_subgraph[0];  invoke_subgraph = None
+
+        subgraph_1 = self.subgraph_0
+        invoke_subgraph_1 = torch.ops.higher_order.invoke_subgraph(subgraph_1, 'subgraph_0', x, synthetic_local_tmp_2_);  subgraph_1 = x = synthetic_local_tmp_2_ = None
+        x_1: "f32[8]" = invoke_subgraph_1[0];  invoke_subgraph_1 = None
+
+        subgraph_2 = self.subgraph_0
+        invoke_subgraph_2 = torch.ops.higher_order.invoke_subgraph(subgraph_2, 'subgraph_0', x_1, synthetic_local_tmp_4_);  subgraph_2 = x_1 = synthetic_local_tmp_4_ = None
+        x_2: "f32[8]" = invoke_subgraph_2[0];  invoke_subgraph_2 = None
+        return (x_2,)
+
+    class subgraph_0(torch.nn.Module):
+        def forward(self, l_x_: "f32[8]", synthetic_local_tmp_0_ : test_opaque_obj_v2_HoistedString):
+            op_with_string_default: "f32[8]" = torch.ops.mylib.op_with_string.default(l_x_, synthetic_local_tmp_0_);  l_x_ = synthetic_local_tmp_0_ = None
+            return (op_with_string_default,)
+""",  # noqa: B950
+            )
+
+    def test_subgraph_reuse_different_list_lengths(self):
+        """Reuse must be skipped when list args have different lengths.
+
+        The first call passes lists of length 2; the second passes lists of
+        length 3.  The pytree treespec will differ, so the cache lookup must
+        fall through and trigger a second trace.
+        """
+
+        @nested_compile_region
+        def gn(xs, ys):
+            return [a + b for a, b in zip(xs, ys)]
+
+        def fn(xs1, ys1, xs2, ys2):
+            a = gn(xs1, ys1)
+            b = gn(xs2, ys2)
+            return a, b
+
+        xs1 = [torch.randn(4), torch.randn(4)]
+        ys1 = [torch.randn(4), torch.randn(4)]
+        xs2 = [torch.randn(4), torch.randn(4), torch.randn(4)]
+        ys2 = [torch.randn(4), torch.randn(4), torch.randn(4)]
+
+        ref = fn(xs1, ys1, xs2, ys2)
+
+        with self._count_speculate_calls() as count:
+            res = torch.compile(fn, backend="aot_eager", fullgraph=True)(
+                xs1, ys1, xs2, ys2
+            )
+
+        # Different list lengths → treespec mismatch → two separate traces
+        self.assertEqual(count(), 2)
+        for r, e in zip(res, ref):
+            for ri, ei in zip(r, e):
+                self.assertEqual(ri, ei)
+
+    def test_subgraph_reuse_different_constants_retrace(self):
+        """Constant args with different values each require a fresh trace.
+
+        Three calls with three distinct scalar constants → call_count == 3.
+        """
+
+        @nested_compile_region
+        def gn(x, scale):
+            return x * scale
+
+        def fn(x):
+            a = gn(x, 1)
+            b = gn(x, 2)
+            c = gn(x, 3)
+            return a, b, c
+
+        x = torch.randn(4)
+        ref = fn(x)
+
+        with self._count_speculate_calls() as count:
+            res = torch.compile(fn, backend="aot_eager", fullgraph=True)(x)
+
+        for r, e in zip(res, ref):
+            self.assertEqual(r, e)
+        # Three distinct constants → three separate traces
+        self.assertEqual(count(), 3)
+
+    def test_subgraph_reuse_tuple_destructure_with_intermediates(self):
+        class Layer(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.linear = torch.nn.Linear(8, 8, bias=False)
+
+            @nested_compile_region
+            def forward(self, x, residual):
+                h = self.linear(x)
+                h = h + residual
+                new_residual = h * 0.5
+                return h, new_residual
+
+        class Model(torch.nn.Module):
+            def __init__(self, num_layers):
+                super().__init__()
+                self.embed = torch.nn.Linear(8, 8, bias=False)
+                self.layers = torch.nn.ModuleList([Layer() for _ in range(num_layers)])
+
+            def forward(self, x, residual):
+                # embed gives x/residual requires_grad=True (same as layer outputs),
+                # so all layers see identical tensor metadata → single trace suffices.
+                x = self.embed(x)
+                residual = self.embed(residual)
+                for layer in self.layers:
+                    # Must support extra outputs
+                    hidden_states, residual = layer(x, residual)
+                    x = hidden_states
+                return x, residual
+
+        model = Model(3)
+        x = torch.randn(4, 8)
+        residual = torch.randn(4, 8)
+        ref = model(x, residual)
+
+        torch._dynamo.reset()
+
+        with self._count_speculate_calls() as count:
+            res = torch.compile(model, backend="aot_eager", fullgraph=True)(
+                x.clone(), residual.clone()
+            )
+        self.assertEqual(ref[0], res[0])
+        self.assertEqual(ref[1], res[1])
+        # All layers see identical tensor metadata (requires_grad=True throughout)
+        # so layers[1] and layers[2] reuse layers[0]'s trace.
+        self.assertEqual(count(), 1)
+
+    def test_subgraph_reuse_different_dynamic_symnodes(self):
+        @nested_compile_region
+        def gn(x, n):
+            return x * n
+
+        def fn(x):
+            a = gn(x, x.shape[0])
+            b = gn(x, x.shape[1])
+            return a, b
+
+        x = torch.ones(4, 6)
+        ref = fn(x)
+
+        with self._count_speculate_calls() as count:
+            res = torch.compile(fn, backend="aot_eager", dynamic=True)(x)
+
+        self.assertEqual(ref[0], res[0])
+        self.assertEqual(ref[1], res[1])
+        # s0 (dim 0) and s1 (dim 1) are distinct symbols → two separate traces
+        self.assertEqual(count(), 2)
+
+    def test_subgraph_reuse_cache_multiple_entries(self):
+        @nested_compile_region
+        def gn(x):
+            return x.sin()
+
+        def fn(x4, x8, x4_again):
+            a = gn(x4)
+            b = gn(x8)
+            c = gn(x4_again)
+            return a, b, c
+
+        x4 = torch.randn(4)
+        x8 = torch.randn(8)
+        x4_again = torch.randn(4)
+        ref = fn(x4, x8, x4_again)
+
+        with self._count_speculate_calls() as count:
+            res = torch.compile(fn, backend="aot_eager", fullgraph=True)(
+                x4, x8, x4_again
+            )
+
+        for r, e in zip(res, ref):
+            self.assertEqual(r, e)
+        self.assertEqual(count(), 2)
+
+    def test_subgraph_reuse_kwargs(self):
+        @nested_compile_region
+        def gn(x, *, scale=1.0):
+            return x * scale
+
+        def fn(x):
+            a = gn(x, scale=2.0)
+            b = gn(x, scale=2.0)
+            return a + b
+
+        x = torch.randn(8)
+        ref = fn(x)
+
+        with self._count_speculate_calls() as count:
+            res = torch.compile(fn, backend="aot_eager", fullgraph=True)(x)
+
+        self.assertEqual(ref, res)
+        self.assertEqual(count(), 1)
+
+    def test_subgraph_reuse_max_entries_raises(self):
+        """Exceeding max_reuse_entries raises RuntimeError."""
+
+        @nested_compile_region(max_reuse_entries=2)
+        def gn(x, c):
+            return x * c
+
+        def fn(x):
+            # Three distinct constants exceed the limit of 2
+            return gn(x, 1) + gn(x, 2) + gn(x, 3)
+
+        x = torch.randn(4)
+        with self.assertRaisesRegex(RuntimeError, "exceeded maximum reuse entries"):
+            torch.compile(fn, backend="aot_eager", fullgraph=True)(x)
+
+    def test_subgraph_reuse_module_different_instances_retrace(self):
+        """Different module instances with different weights require separate traces."""
+
+        class Mod(torch.nn.Module):
+            def __init__(self, c):
+                super().__init__()
+                self.c = c
+
+            @nested_compile_region
+            def forward(self, x):
+                return x * self.c
+
+        mod1 = Mod(5)
+        mod2 = Mod(10)
+
+        def fn(x):
+            return mod1(x) + mod2(x)
+
+        x = torch.randn(8)
+        ref = fn(x)
+
+        with self._count_speculate_calls() as count:
+            res = torch.compile(fn, backend="aot_eager", fullgraph=True)(x)
+
+        self.assertEqual(ref, res)
+        # mod1.c=5 and mod2.c=10 differ → two separate traces
+        self.assertEqual(count(), 2)
+
+    def test_subgraph_reuse_monkeypatch_forward(self):
+        """Monkeypatching forward with nested_compile_region should reuse cache.
+
+        When a user wraps a module's forward with nested_compile_region via
+        monkeypatching, each module instance gets a distinct bound method and
+        thus a distinct function object passed to nested_compile_region. The
+        cache key should be based on fn.__code__ rather than id(fn) so that
+        reuse still works across instances whose forward shares the same code.
+        """
+
+        class Mod(torch.nn.Module):
+            def __init__(self, c):
+                super().__init__()
+                self.c = c
+
+            def forward(self, x):
+                return x * self.c
+
+        def apply_nested_compile_region(mod):
+            # Grab the unbound function and create a fresh copy so that each
+            # module instance gets a function with a *different* id() but the
+            # *same* __code__ object, mimicking a monkeypatch/code-gen scenario.
+            import types
+
+            orig = type(mod).forward
+            fresh_fn = types.FunctionType(
+                orig.__code__,
+                orig.__globals__,
+                orig.__name__,
+                orig.__defaults__,
+                orig.__closure__,
+            )
+            mod.forward = nested_compile_region(fresh_fn).__get__(mod, type(mod))
+            return mod
+
+        mod1 = apply_nested_compile_region(Mod(5))
+        mod2 = apply_nested_compile_region(Mod(5))
+
+        # The inner functions passed to invoke_subgraph_placeholder are
+        # different objects but share the same __code__.
+        fn1 = mod1.forward.__func__.__marked_compile_region_fn__
+        fn2 = mod2.forward.__func__.__marked_compile_region_fn__
+        self.assertIsNot(fn1, fn2)
+        self.assertIs(fn1.__code__, fn2.__code__)
+
+        def fn(x):
+            return mod1(x) + mod2(x)
+
+        x = torch.randn(8)
+        ref = fn(x)
+
+        with self._count_speculate_calls() as count:
+            res = torch.compile(fn, backend="aot_eager", fullgraph=True)(x)
+
+        self.assertEqual(ref, res)
+        # Both modules have the same code and same c=5, so the second call
+        # should reuse the first trace despite having different function ids.
+        self.assertEqual(count(), 1)
+
+    def test_subgraph_reuse_module_apply(self):
+        """Using module.apply to wrap transformer layers with nested_compile_region."""
+
+        class Layer(torch.nn.Module):
+            def __init__(self, c):
+                super().__init__()
+                self.c = c
+
+            def forward(self, x):
+                return x * self.c
+
+        model = torch.nn.Sequential(Layer(5), Layer(5), Layer(5))
+        ref_model = torch.nn.Sequential(Layer(5), Layer(5), Layer(5))
+
+        def wrap_layer(mod):
+            if isinstance(mod, Layer):
+                fwd = type(mod).forward
+                if not hasattr(fwd, "__marked_compile_region_fn__"):
+                    type(mod).forward = nested_compile_region(fwd)
+
+        model.apply(wrap_layer)
+
+        x = torch.randn(8)
+        ref = ref_model(x)
+
+        with self._count_speculate_calls() as count:
+            res = torch.compile(model, backend="aot_eager", fullgraph=True)(x)
+
+        self.assertEqual(ref, res)
+        # All three layers share the same class and c=5, only one trace needed.
+        self.assertEqual(count(), 1)
+
+    def test_subgraph_reuse_class_level_wrap(self):
+        """Wrapping cls.forward with nested_compile_region for named_modules pattern."""
+
+        class Layer(torch.nn.Module):
+            def __init__(self, c):
+                super().__init__()
+                self.c = c
+
+            def forward(self, x):
+                return x * self.c
+
+        # Wrap at the class level — the pattern a user would use after
+        # iterating named_modules to wrap all layers of a given type.
+        Layer.forward = nested_compile_region(Layer.forward)
+
+        mod1 = Layer(5)
+        mod2 = Layer(5)
+
+        def fn(x):
+            return mod1(x) + mod2(x)
+
+        x = torch.randn(8)
+        ref = fn(x)
+
+        with self._count_speculate_calls() as count:
+            res = torch.compile(fn, backend="aot_eager", fullgraph=True)(x)
+
+        self.assertEqual(ref, res)
+        self.assertEqual(count(), 1)
+
+    def test_subgraph_reuse_module_instance_as_callable(self):
+        """Passing nn.Module instances directly to nested_compile_region."""
+
+        class Layer(torch.nn.Module):
+            def __init__(self, c):
+                super().__init__()
+                self.c = c
+
+            def forward(self, x):
+                return x * self.c
+
+        mod1 = Layer(5)
+        mod2 = Layer(5)
+
+        wrapped1 = nested_compile_region(mod1)
+        wrapped2 = nested_compile_region(mod2)
+
+        def fn(x):
+            return wrapped1(x) + wrapped2(x)
+
+        x = torch.randn(8)
+        ref = fn(x)
+
+        with self._count_speculate_calls() as count:
+            res = torch.compile(fn, backend="aot_eager", fullgraph=True)(x)
+
+        self.assertEqual(ref, res)
+        # Both modules have the same forward code and same c=5, so the
+        # second call should reuse the first trace.
+        self.assertEqual(count(), 1)
+
+
+@skipIfTorchDynamo("Not a torch._dynamo test")
+@parameterized_class(
+    [
+        {"strict": False},
+        {"strict": True},
+    ],
+    class_name_func=lambda cls,
+    _,
+    params: f"{cls.__name__}{'Strict' if params['strict'] else 'Nonstrict'}",
+)
+class TestInvokeSubgraphExport(TestCase):
+    def test_simple_func(self):
+        @nested_compile_region
+        def gn(x, y):
+            return torch.mul(x, y)
+
+        class M(torch.nn.Module):
+            def forward(self, x, y):
+                x = gn(x, y)
+                x = gn(x, y)
+                return x
+
+        x = torch.randn(8, requires_grad=True)
+        y = torch.randn(8, requires_grad=True)
+
+        ep = torch.export.export(M(), (x, y), strict=self.strict)
+        self.assertTrue(torch.allclose(ep.module()(x, y), M()(x, y)))
+        self.assertEqual(len(list(ep.graph_module.named_modules())), 2)
+
+        self.assertExpectedInline(
+            empty_line_normalizer(
+                normalize_gm(ep.graph_module.print_readable(print_output=False))
+            ),
+            """\
+class GraphModule(torch.nn.Module):
+    def forward(self, x: "f32[8]", y: "f32[8]"):
+        repeated_subgraph0 = self.repeated_subgraph0
+        invoke_subgraph = torch.ops.higher_order.invoke_subgraph(repeated_subgraph0, 'subgraph_0', x, y);  repeated_subgraph0 = x = None
+        getitem: "f32[8]" = invoke_subgraph[0];  invoke_subgraph = None
+        repeated_subgraph0_1 = self.repeated_subgraph0
+        invoke_subgraph_1 = torch.ops.higher_order.invoke_subgraph(repeated_subgraph0_1, 'subgraph_0', getitem, y);  repeated_subgraph0_1 = getitem = y = None
+        getitem_1: "f32[8]" = invoke_subgraph_1[0];  invoke_subgraph_1 = None
+        return (getitem_1,)
+    class repeated_subgraph0(torch.nn.Module):
+        def forward(self, arg0_1: "f32[8]", arg1_1: "f32[8]"):
+            mul: "f32[8]" = torch.ops.aten.mul.Tensor(arg0_1, arg1_1);  arg0_1 = arg1_1 = None
+            return (mul,)
+""",
+        )
+
+    def test_unbacked(self):
+        @nested_compile_region
+        def gn(x, y):
+            b = x.item()
+            torch._check(b < y.shape[0])
+            return y[:b].clone()
+
+        class M(torch.nn.Module):
+            def forward(self, x, y):
+                res = []
+                for _ in range(10):
+                    res.append(gn(x, y))
+                return torch.cat(res)
+
+        x = torch.tensor(4)
+        y = torch.randn(8)
+
+        ep = torch.export.export(M(), (x, y), strict=self.strict)
+        ep = ep.run_decompositions()
+
+        self.assertTrue(torch.allclose(ep.module()(x, y), M()(x, y)))
+        self.assertEqual(len(list(ep.graph_module.named_modules())), 2)
+
+    def test_pending_unbacked(self):
+        class M(torch.nn.Module):
+            @nested_compile_region
+            def gn(self, x):
+                u = x[0].item()
+                return x * u
+
+            def forward(self, x):
+                for _ in range(4):
+                    x = self.gn(x)
+                return x
+
+        ep = torch.export.export(
+            M(),
+            (torch.randn(8),),
+            strict=self.strict,
+            dynamic_shapes={"x": {0: torch.export.Dim.DYNAMIC}},
+        )
+        ep = ep.run_decompositions()
+
+        self.assertEqual(len(list(ep.graph_module.named_modules())), 2)
+
+        ep = torch.export.export(
+            M(),
+            (torch.randn(8, requires_grad=True),),
+            strict=self.strict,
+            dynamic_shapes={"x": {0: torch.export.Dim.DYNAMIC}},
+        )
+        ep = ep.run_decompositions()
+
+        self.assertEqual(len(list(ep.graph_module.named_modules())), 2)
+
+    def test_simple_method(self):
+        class M(torch.nn.Module):
+            @nested_compile_region
+            def gn(self, x, y):
+                return torch.mul(x, y)
+
+            def forward(self, x, y):
+                x = self.gn(x, y)
+                x = self.gn(x, y)
+                return x
+
+        x = torch.randn(8, requires_grad=True)
+        y = torch.randn(8, requires_grad=True)
+
+        ep = torch.export.export(M(), (x, y), strict=self.strict)
+        self.assertTrue(torch.allclose(ep.module()(x, y), M()(x, y)))
+        self.assertEqual(len(list(ep.graph_module.named_modules())), 2)
+
+    def test_multiple_module(self):
+        b = torch.randn(8)
+
+        class N(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.register_buffer("buf", b)
+
+            @nested_compile_region
+            def forward(self, x, y):
+                return x * y + self.buf
+
+        class M(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.mod_list = torch.nn.ModuleList(N() for _ in range(10))
+
+            def forward(self, x, y):
+                for m in self.mod_list:
+                    x = m(x, y)
+                return x
+
+        x = torch.randn(8, requires_grad=True)
+        y = torch.randn(8, requires_grad=True)
+
+        ep = torch.export.export(M(), (x, y), strict=self.strict)
+        self.assertTrue(torch.allclose(ep.module()(x, y), M()(x, y)))
+        self.assertEqual(len(list(ep.graph_module.named_modules())), 2)
+
+
+class NegativeTesting(TestCase):
+    def test_graph_break(self):
+        @nested_compile_region
+        def gn(x):
+            torch._dynamo.graph_break()
+            return torch.cos(x)
+
+        def fn(x):
+            return gn(x)
+
+        x = torch.randn(8, 8, requires_grad=True)
+
+        with self.assertRaisesRegex(
             RuntimeError,
             r"Higher Order Operator: torch\.ops\.higher_order\.invoke_subgraph",
         ):
             torch.compile(fn, backend="eager")(x)
 
 
+@skipIfTorchDynamo("Not a torch._dynamo test")
+class TestInlineInvokeSubgraph(TestCase):
+    def _assert_no_invoke_subgraph(self, fn, args):
+        """Compile fn and verify the backend receives no invoke_subgraph HOPs."""
+        backend = EagerAndRecordGraphs()
+        res = torch.compile(fn, backend=backend, fullgraph=True)(*args)
+        self.assertTrue(len(backend.graphs) > 0)
+        for gm in backend.graphs:
+            for node in gm.graph.nodes:
+                self.assertFalse(
+                    node.op == "call_function"
+                    and node.target is torch.ops.higher_order.invoke_subgraph,
+                )
+        return res
+
+    @torch._dynamo.config.patch(inline_invoke_subgraph=True)
+    def test_simple(self):
+        @nested_compile_region
+        def gn(x, y):
+            return torch.mul(x, y)
+
+        def fn(x, y):
+            return gn(x, y) + gn(x, y)
+
+        x = torch.randn(8)
+        y = torch.randn(8)
+        ref = fn(x, y)
+        res = self._assert_no_invoke_subgraph(fn, (x, y))
+        self.assertEqual(ref, res)
+
+    @torch._dynamo.config.patch(inline_invoke_subgraph=True)
+    def test_module(self):
+        class Mod(torch.nn.Module):
+            @nested_compile_region
+            def forward(self, x):
+                return x.sin() + x.cos()
+
+        mod = Mod()
+
+        def fn(x):
+            return mod(x) + mod(x)
+
+        x = torch.randn(8)
+        ref = fn(x)
+        res = self._assert_no_invoke_subgraph(fn, (x,))
+        self.assertEqual(ref, res)
+
+    @torch._dynamo.config.patch(inline_invoke_subgraph=True)
+    def test_backward(self):
+        @nested_compile_region
+        def gn(x, y):
+            return torch.mul(x, y).sin()
+
+        def fn(x, y):
+            return gn(x, y)
+
+        x = torch.randn(8, requires_grad=True)
+        y = torch.randn(8, requires_grad=True)
+        ref = fn(x, y)
+
+        x2 = x.detach().clone().requires_grad_(True)
+        y2 = y.detach().clone().requires_grad_(True)
+        res = torch.compile(fn, backend="inductor", fullgraph=True)(x2, y2)
+
+        ref.sum().backward()
+        res.sum().backward()
+
+        self.assertEqual(ref, res)
+        self.assertEqual(x.grad, x2.grad)
+        self.assertEqual(y.grad, y2.grad)
+
+
+@skipIfTorchDynamo("Not a torch._dynamo test")
+class TestInvokeSubgraphReuseHashFn(TestCase):
+    @contextlib.contextmanager
+    def _count_speculate_calls(self):
+        count = 0
+        orig = torch._dynamo.variables.higher_order_ops.speculate_subgraph_with_auto_output_flattening
+
+        def _counting(*args, **kwargs):
+            nonlocal count
+            count += 1
+            return orig(*args, **kwargs)
+
+        with mock.patch.object(
+            torch._dynamo.variables.higher_order_ops,
+            "speculate_subgraph_with_auto_output_flattening",
+            _counting,
+        ):
+            yield lambda: count
+
+    def test_reuse_hash_fn_module_distinct_hashes(self):
+        """nn.Module arg with hash fn returning different values per layer."""
+
+        def hash_fn(mod, x):
+            return mod.layer_id
+
+        @nested_compile_region(reuse_hash_fn=hash_fn)
+        def layer_fn(mod, x):
+            return x.sin() + mod.weight
+
+        class Layer(torch.nn.Module):
+            def __init__(self, layer_id):
+                super().__init__()
+                self.layer_id = layer_id
+                self.weight = torch.nn.Parameter(torch.randn(8))
+
+        class Model(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.layers = torch.nn.ModuleList([Layer(i) for i in range(4)])
+
+            def forward(self, x):
+                for layer in self.layers:
+                    x = layer_fn(layer, x)
+                return x
+
+        mod = Model()
+        x = torch.randn(8)
+        ref = mod(x)
+
+        with self._count_speculate_calls() as count:
+            res = torch.compile(mod, backend="aot_eager", fullgraph=True)(x)
+
+        # Each layer has a distinct layer_id → 4 separate traces
+        self.assertEqual(count(), 4)
+        self.assertEqual(ref, res)
+
+    def test_reuse_hash_fn_module_same_hash(self):
+        """nn.Module arg with hash fn returning same value → single trace."""
+
+        def hash_fn(mod, x):
+            return 0
+
+        @nested_compile_region(reuse_hash_fn=hash_fn)
+        def layer_fn(mod, x):
+            return x.sin() + mod.weight
+
+        class Layer(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.weight = torch.nn.Parameter(torch.randn(8))
+
+        class Model(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.layers = torch.nn.ModuleList([Layer() for _ in range(4)])
+
+            def forward(self, x):
+                for layer in self.layers:
+                    x = layer_fn(layer, x)
+                return x
+
+        mod = Model()
+        x = torch.randn(8)
+        ref = mod(x)
+
+        with self._count_speculate_calls() as count:
+            res = torch.compile(mod, backend="aot_eager", fullgraph=True)(x)
+
+        # All layers hash to 0 → single trace + 3 stamp-outs
+        self.assertEqual(count(), 1)
+        self.assertEqual(ref, res)
+
+    def test_reuse_hash_fn_tensor_shape(self):
+        """Hash fn that uses tensor shape to differentiate inputs."""
+
+        def hash_fn(x):
+            return x.shape[0]
+
+        @nested_compile_region(reuse_hash_fn=hash_fn)
+        def gn(x):
+            return x.sin()
+
+        def fn(x4, x8a, x8b):
+            # x8a and x8b have the same shape → same hash → reuse
+            a = gn(x4)
+            b = gn(x8a)
+            c = gn(x8b)
+            return a.sum() + b.sum() + c.sum()
+
+        x4 = torch.randn(4)
+        x8a = torch.randn(8)
+        x8b = torch.randn(8)
+        ref = fn(x4, x8a, x8b)
+
+        with self._count_speculate_calls() as count:
+            res = torch.compile(fn, backend="aot_eager", fullgraph=True)(x4, x8a, x8b)
+
+        # shape[0]=4 and shape[0]=8 → 2 traces, third call reuses shape=8
+        self.assertEqual(count(), 2)
+        self.assertEqual(ref, res)
+
+    def test_reuse_hash_fn_graph_break_raises(self):
+        """reuse_hash_fn with a graph break raises a clear error."""
+
+        def bad_hash_fn(mod, x):
+            torch._dynamo.graph_break()
+            return 0
+
+        @nested_compile_region(reuse_hash_fn=bad_hash_fn)
+        def layer_fn(mod, x):
+            return x.sin() + mod.weight
+
+        class Layer(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.weight = torch.nn.Parameter(torch.randn(8))
+
+        class Model(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.layers = torch.nn.ModuleList([Layer(), Layer()])
+
+            def forward(self, x):
+                for layer in self.layers:
+                    x = layer_fn(layer, x)
+                return x
+
+        mod = Model()
+        x = torch.randn(8)
+        with self.assertRaisesRegex(
+            RuntimeError, "reuse_hash_fn must be fully traceable"
+        ):
+            torch.compile(mod, backend="aot_eager", fullgraph=True)(x)
+
+    def test_reuse_hash_fn_if_cond_no_guard(self):
+        """if-condition in hash fn should not install guards on the module."""
+
+        def hash_fn(mod, x):
+            if mod.use_gelu:
+                return 1
+            return 0
+
+        @nested_compile_region(reuse_hash_fn=hash_fn)
+        def layer_fn(mod, x):
+            return x.sin() + mod.weight
+
+        class Layer(torch.nn.Module):
+            def __init__(self, use_gelu):
+                super().__init__()
+                self.use_gelu = use_gelu
+                self.weight = torch.nn.Parameter(torch.randn(8))
+
+        class Model(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.layers = torch.nn.ModuleList([Layer(True), Layer(False)])
+
+            def forward(self, x):
+                for layer in self.layers:
+                    x = layer_fn(layer, x)
+                return x
+
+        mod = Model()
+        x = torch.randn(8)
+        ref = mod(x)
+
+        cnt = torch._dynamo.testing.CompileCounter()
+        compiled = torch.compile(mod, backend=cnt)
+        res = compiled(x)
+        self.assertEqual(ref, res)
+
+        # Flip use_gelu on both layers — if hash fn guards leaked, this
+        # would cause a recompilation. With proper guard stripping it
+        # should not recompile (the outer frame guard count stays the same).
+        frame_count_before = cnt.frame_count
+        mod.layers[0].use_gelu = False
+        mod.layers[1].use_gelu = True
+        ref2 = mod(x)
+        res2 = compiled(x)
+        self.assertEqual(ref2, res2)
+        self.assertEqual(cnt.frame_count, frame_count_before)
+
+    def test_reuse_hash_fn_side_effect_allowed(self):
+        """Side effects (attribute mutation) should not block reuse with reuse_hash_fn."""
+
+        def hash_fn(mod, x):
+            return 0
+
+        @nested_compile_region(reuse_hash_fn=hash_fn)
+        def layer_fn(mod, x):
+            mod.call_count += 1
+            return x.sin() + mod.weight
+
+        class Layer(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.call_count = 0
+                self.weight = torch.nn.Parameter(torch.randn(8))
+
+        class Model(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.layers = torch.nn.ModuleList([Layer() for _ in range(3)])
+
+            def forward(self, x):
+                for layer in self.layers:
+                    x = layer_fn(layer, x)
+                return x
+
+        mod = Model()
+        x = torch.randn(8)
+
+        # Without reuse_hash_fn, the side effect (mod.call_count += 1)
+        # would prevent reuse. With reuse_hash_fn, it should still reuse.
+        with self._count_speculate_calls() as count:
+            torch.compile(mod, backend="aot_eager", fullgraph=True)(x)
+
+        self.assertEqual(count(), 1)
+
+    def test_reuse_hash_fn_unsupported_output_raises(self):
+        """Nested output (tuple of tuple of tensors) should raise with reuse_hash_fn."""
+
+        def hash_fn(x):
+            return 0
+
+        @nested_compile_region(reuse_hash_fn=hash_fn)
+        def gn(x):
+            return ((x.sin(), x.cos()),)
+
+        def fn(x):
+            return gn(x)
+
+        x = torch.randn(8)
+        with self.assertRaisesRegex(
+            RuntimeError, "reuse_hash_fn was provided but the subgraph is not eligible"
+        ):
+            torch.compile(fn, backend="aot_eager", fullgraph=True)(x)
+
+
 if __name__ == "__main__":
     run_tests()
diff --git a/test/higher_order_ops/test_local_map.py b/test/higher_order_ops/test_local_map.py
index 2483da3cc4903..2eb818f8831ae 100644
--- a/test/higher_order_ops/test_local_map.py
+++ b/test/higher_order_ops/test_local_map.py
@@ -6,7 +6,7 @@
 import unittest
 from collections.abc import Callable
 from contextlib import contextmanager, ExitStack
-from typing import Any, Optional
+from typing import Any
 
 import torch
 import torch._dynamo
@@ -756,8 +756,8 @@ def axis_size(axis_name):
 
         def _all_to_all(
             self: torch.Tensor,
-            output_split_sizes: Optional[list[int]],
-            input_split_sizes: Optional[list[int]],
+            output_split_sizes: list[int] | None,
+            input_split_sizes: list[int] | None,
             group_name: str,
         ):
             group_size = c10d._get_group_size_by_name(group_name)
@@ -781,8 +781,8 @@ class _AllToAll(torch.autograd.Function):
             def forward(
                 ctx: Any,
                 x: torch.Tensor,
-                output_split_sizes: Optional[list[int]],
-                input_split_sizes: Optional[list[int]],
+                output_split_sizes: list[int] | None,
+                input_split_sizes: list[int] | None,
                 axis_name: str,
             ):
                 group_name = _get_group_name_from_axis_name(axis_name)
@@ -932,6 +932,34 @@ def inputs_fn():
         self.assertEqual(mm_nodes[2].meta["custom"]["inside_local_map"], 1)
         self.assertEqual(mm_nodes[3].meta["custom"]["inside_local_map"], 0)
 
+    @unittest.skipIf(*get_skip_reasons())
+    def test_local_map_make_contiguous_strides_for(self):
+        # make_contiguous_strides_for inside local_map must compile with dynamic shapes.
+        from torch._prims_common import make_contiguous_strides_for
+
+        @local_map(
+            out_placements=((Shard(0), Replicate()),),
+            in_placements=((Shard(0), Replicate()),),
+            redistribute_inputs=True,
+            device_mesh=self.mesh,
+        )
+        def fn(x):
+            strides = make_contiguous_strides_for(x.shape)
+            return (x.as_strided(x.shape, strides),)
+
+        class MyModule(torch.nn.Module):
+            def forward(self, x):
+                return fn(x)
+
+        model = MyModule()
+        from torch._dynamo.testing import EagerAndRecordGraphs
+
+        backend = EagerAndRecordGraphs()
+        x = torch.randn(80, 80)
+        with enable_local_map_wrapping():
+            out = torch.compile(model, backend=backend)(x)
+        self.assertEqual(out[0].shape, x.shape)
+
 
 if __name__ == "__main__":
     run_tests()
diff --git a/test/higher_order_ops/test_print.py b/test/higher_order_ops/test_print.py
index 716b82ea4464e..a133044e081a6 100644
--- a/test/higher_order_ops/test_print.py
+++ b/test/higher_order_ops/test_print.py
@@ -1,5 +1,6 @@
 # Owner(s): ["module: higher order operators"]
 import io
+import unittest
 from unittest.mock import patch
 
 import torch
@@ -17,6 +18,17 @@
 )
 
 
+if torch.distributed.is_available():
+    from torch.distributed.tensor import DTensor, Replicate, Shard
+    from torch.testing._internal.distributed._tensor.common_dtensor import (
+        DTensorTestBase,
+        with_comms,
+    )
+else:
+    DTensorTestBase = TestCase  # type: ignore[assignment, misc]
+    with_comms = lambda fn: fn  # type: ignore[assignment]  # noqa: E731
+
+
 @instantiate_parametrized_tests
 class TestHopPrint(TestCase):
     def test_base_print(self):
@@ -672,5 +684,208 @@ def forward(self, arg1_1):
         )
 
 
+@unittest.skipIf(
+    not torch.distributed.is_available(), "torch.distributed not available"
+)
+class TestHopPrintDTensor(DTensorTestBase):
+    @property
+    def world_size(self) -> int:
+        return 4
+
+    @with_comms
+    def test_print_dtensor_basic(self):
+        """Sharded DTensor prints local shard on all ranks."""
+        device_mesh = self.build_device_mesh()
+        full_tensor = torch.arange(8, dtype=torch.float, device=self.device_type)
+        local_shard = full_tensor.chunk(self.world_size)[self.rank]
+        dtensor = DTensor.from_local(local_shard, device_mesh, [Shard(0)])
+
+        def f(x):
+            x = x + x
+            torch._higher_order_ops.print("tensor: {}", x)
+            return x
+
+        local_doubled = local_shard + local_shard
+        expected = f"[rank {self.rank}] tensor: {local_doubled}\n"
+
+        with patch("sys.stdout", new_callable=io.StringIO) as mock_stdout:
+            f(dtensor)
+            output = mock_stdout.getvalue()
+        self.assertEqual(output, expected)
+
+    @with_comms
+    def test_print_dtensor_replicate(self):
+        """Replicated DTensor prints full tensor on all ranks."""
+        device_mesh = self.build_device_mesh()
+        full_tensor = torch.tensor([1.0, 2.0, 3.0], device=self.device_type)
+        dtensor = DTensor.from_local(full_tensor, device_mesh, [Replicate()])
+
+        def f(x):
+            x = x * 2
+            torch._higher_order_ops.print("val: {}", x)
+            return x
+
+        expected = f"[rank {self.rank}] val: {full_tensor * 2}\n"
+
+        with patch("sys.stdout", new_callable=io.StringIO) as mock_stdout:
+            eager_result = f(dtensor)
+            eager_output = mock_stdout.getvalue()
+        self.assertEqual(eager_output, expected)
+
+        opt_f = torch.compile(f, backend="aot_eager", fullgraph=True)
+        with patch("sys.stdout", new_callable=io.StringIO) as mock_stdout:
+            compiled_result = opt_f(dtensor)
+            compiled_output = mock_stdout.getvalue()
+
+        self.assertEqual(compiled_result.to_local(), eager_result.to_local())
+        self.assertEqual(compiled_output, expected)
+
+    @with_comms
+    def test_print_dtensor_format_str(self):
+        """Test both positional and keyword sharded DTensor args in format strings."""
+        device_mesh = self.build_device_mesh()
+        full_tensor = torch.arange(4, dtype=torch.float, device=self.device_type)
+        local_shard = full_tensor.chunk(self.world_size)[self.rank]
+        dtensor = DTensor.from_local(local_shard, device_mesh, [Shard(0)])
+
+        def f_pos(x):
+            torch._higher_order_ops.print("pos: {}", x)
+
+        def f_kw(x):
+            torch._higher_order_ops.print("kw: {x}", x=x)
+
+        with patch("sys.stdout", new_callable=io.StringIO) as mock_stdout:
+            f_pos(dtensor)
+            pos_output = mock_stdout.getvalue()
+
+        with patch("sys.stdout", new_callable=io.StringIO) as mock_stdout:
+            f_kw(dtensor)
+            kw_output = mock_stdout.getvalue()
+
+        self.assertEqual(pos_output, f"[rank {self.rank}] pos: {local_shard}\n")
+        self.assertEqual(kw_output, f"[rank {self.rank}] kw: {local_shard}\n")
+
+    @with_comms
+    def test_print_dtensor_mixed_args(self):
+        """Mix sharded DTensor and scalar args in a single print call."""
+        device_mesh = self.build_device_mesh()
+        full_tensor = torch.arange(4, dtype=torch.float, device=self.device_type)
+        local_shard = full_tensor.chunk(self.world_size)[self.rank]
+        dtensor = DTensor.from_local(local_shard, device_mesh, [Shard(0)])
+
+        def f(x):
+            torch._higher_order_ops.print("dt: {} scalar: {}", x, 42)
+
+        with patch("sys.stdout", new_callable=io.StringIO) as mock_stdout:
+            f(dtensor)
+            output = mock_stdout.getvalue()
+
+        self.assertEqual(output, f"[rank {self.rank}] dt: {local_shard} scalar: 42\n")
+
+    @with_comms
+    def test_print_dtensor_multiple_prints(self):
+        """Multiple sharded DTensor prints with intermediate computations."""
+        device_mesh = self.build_device_mesh()
+        full_tensor = torch.arange(4, dtype=torch.float, device=self.device_type)
+        local_shard = full_tensor.chunk(self.world_size)[self.rank]
+        dtensor = DTensor.from_local(local_shard, device_mesh, [Shard(0)])
+
+        def f(x):
+            x1 = x + x
+            torch._higher_order_ops.print("after add: {}", x1)
+            x2 = x1 * x1
+            torch._higher_order_ops.print("after mul: {}", x2)
+            return x2
+
+        local_added = local_shard + local_shard
+        local_mulled = local_added * local_added
+        r = self.rank
+        expected = f"[rank {r}] after add: {local_added}\n[rank {r}] after mul: {local_mulled}\n"
+
+        with patch("sys.stdout", new_callable=io.StringIO) as mock_stdout:
+            f(dtensor)
+            output = mock_stdout.getvalue()
+        self.assertEqual(output, expected)
+
+    @with_comms
+    def test_print_dtensor_kwargs(self):
+        """Sharded DTensor print with kwargs."""
+        device_mesh = self.build_device_mesh()
+        full_tensor = torch.arange(4, dtype=torch.float, device=self.device_type)
+        local_shard = full_tensor.chunk(self.world_size)[self.rank]
+        dtensor = DTensor.from_local(local_shard, device_mesh, [Shard(0)])
+
+        def f(x):
+            x = x + 1
+            torch._higher_order_ops.print("result: {x} count: {n}", x=x, n=42)
+            return x
+
+        expected = f"[rank {self.rank}] result: {local_shard + 1} count: 42\n"
+
+        with patch("sys.stdout", new_callable=io.StringIO) as mock_stdout:
+            f(dtensor)
+            output = mock_stdout.getvalue()
+        self.assertEqual(output, expected)
+
+    @with_comms
+    @skipIfTorchDynamo("Skipped under Dynamo")
+    def test_print_dtensor_inductor_output_code(self):
+        """Verify inductor generated code contains print for replicated DTensor."""
+        device_mesh = self.build_device_mesh()
+        full_tensor = torch.arange(8, dtype=torch.float, device=self.device_type)
+        dtensor = DTensor.from_local(full_tensor, device_mesh, [Replicate()])
+
+        def f(x):
+            x = x + x
+            torch._higher_order_ops.print("val: {}", x)
+            return x
+
+        compiled_f = torch.compile(f, backend="inductor", fullgraph=True)
+        with patch("sys.stdout", new_callable=io.StringIO):
+            _, codes = run_and_get_code(compiled_f, dtensor)
+
+        merged_code = "\n".join(codes)
+        self.assertIn(
+            "print",
+            merged_code,
+            "Inductor output code should contain print call",
+        )
+        self.assertNotIn(
+            "torch.ops.higher_order.print",
+            merged_code,
+            "Inductor should use python print, not the HOP directly",
+        )
+
+    @with_comms
+    @skipIfTorchDynamo("Skipped under Dynamo")
+    def test_print_dtensor_compiled_sharded(self):
+        """Verify compiled sharded DTensor prints match eager output per rank."""
+        device_mesh = self.build_device_mesh()
+        full_tensor = torch.arange(8, dtype=torch.float, device=self.device_type)
+        local_shard = full_tensor.chunk(self.world_size)[self.rank]
+        dtensor = DTensor.from_local(local_shard, device_mesh, [Shard(0)])
+
+        def f(x):
+            x = x + x
+            torch._higher_order_ops.print("val: {}", x)
+            return x
+
+        local_doubled = local_shard + local_shard
+        expected = f"[rank {self.rank}] val: {local_doubled}\n"
+
+        with patch("sys.stdout", new_callable=io.StringIO) as mock_stdout:
+            eager_result = f(dtensor)
+            eager_output = mock_stdout.getvalue()
+        self.assertEqual(eager_output, expected)
+
+        opt_f = torch.compile(f, backend="aot_eager", fullgraph=True)
+        with patch("sys.stdout", new_callable=io.StringIO) as mock_stdout:
+            compiled_result = opt_f(dtensor)
+            compiled_output = mock_stdout.getvalue()
+
+        self.assertEqual(compiled_result.to_local(), eager_result.to_local())
+        self.assertEqual(compiled_output, expected)
+
+
 if __name__ == "__main__":
     run_tests()
diff --git a/test/inductor/extension_backends/triton/device_interface.py b/test/inductor/extension_backends/triton/device_interface.py
index 857885a662d5a..63d87575e3dfb 100644
--- a/test/inductor/extension_backends/triton/device_interface.py
+++ b/test/inductor/extension_backends/triton/device_interface.py
@@ -27,12 +27,15 @@ def __init__(
         def record(self, stream) -> None:
             if not self.enable_timing:
                 return
-            assert self.recorded_time is None
+            if self.recorded_time is not None:
+                raise AssertionError
             self.recorded_time = time.perf_counter_ns()
 
         def elapsed_time(self, end_event: DeviceInterface.Event) -> float:
-            assert self.recorded_time
-            assert end_event.recorded_time
+            if not self.recorded_time:
+                raise AssertionError
+            if not end_event.recorded_time:
+                raise AssertionError
             # convert to ms
             return (end_event.recorded_time - self.recorded_time) / 1000000
 
@@ -80,16 +83,18 @@ def device_count() -> int:
 
     @staticmethod
     def maybe_exchange_device(device: int) -> int:
-        assert device == 0, (
-            f"Only device index 0 is supported, tried to set index to {device}"
-        )
+        if device != 0:
+            raise AssertionError(
+                f"Only device index 0 is supported, tried to set index to {device}"
+            )
         return 0  # previous device is always 0
 
     @staticmethod
     def exchange_device(device: int) -> int:
-        assert device == 0, (
-            f"Only device index 0 is supported, tried to set index to {device}"
-        )
+        if device != 0:
+            raise AssertionError(
+                f"Only device index 0 is supported, tried to set index to {device}"
+            )
         return 0  # previous device is always 0
 
     @staticmethod
diff --git a/test/inductor/extension_backends/triton/extension_triton_heuristics.py b/test/inductor/extension_backends/triton/extension_triton_heuristics.py
index bfe558ae1708a..2fec4a35aea85 100644
--- a/test/inductor/extension_backends/triton/extension_triton_heuristics.py
+++ b/test/inductor/extension_backends/triton/extension_triton_heuristics.py
@@ -12,9 +12,11 @@ def _create_compile_meta(
         self,
         cfg: triton_heuristics.Config,
     ) -> dict[str, Any]:
-        assert EXTENSION_TRITON_META_FIELD in self.triton_meta
+        if EXTENSION_TRITON_META_FIELD not in self.triton_meta:
+            raise AssertionError
         compile_meta = super()._create_compile_meta(cfg)
-        assert EXTENSION_TRITON_META_FIELD in compile_meta
+        if EXTENSION_TRITON_META_FIELD not in compile_meta:
+            raise AssertionError
         return compile_meta
 
 
diff --git a/test/inductor/indirect_assert_helper.py b/test/inductor/indirect_assert_helper.py
index 6d1bc2b608fba..b5228d31eada1 100644
--- a/test/inductor/indirect_assert_helper.py
+++ b/test/inductor/indirect_assert_helper.py
@@ -52,17 +52,22 @@ def lower2(x):
     ]
 
     _, fn_name, dims, dyn_shape, one_size = sys.argv
-    assert fn_name in fns
-    assert one_size in ("True", "False")
+    if fn_name not in fns:
+        raise AssertionError(f"Unknown function: {fn_name}")
+    if one_size not in ("True", "False"):
+        raise AssertionError(f"one_size must be 'True' or 'False', got {one_size}")
     one_size = one_size == "True"
-    assert dims in ("2", "3")
+    if dims not in ("2", "3"):
+        raise AssertionError(f"dims must be '2' or '3', got {dims}")
     shape_x = [3, 2, 4] if dims == "3" else [3, 2]
     if one_size:
-        assert fn_name == "first_arg", (
-            "only first_arg can be tested for a special case of 1-size tensor"
-        )
+        if fn_name != "first_arg":
+            raise AssertionError(
+                "only first_arg can be tested for a special case of 1-size tensor"
+            )
         shape_x[0] = 1
-    assert dyn_shape in ("True", "False")
+    if dyn_shape not in ("True", "False"):
+        raise AssertionError(f"dyn_shape must be 'True' or 'False', got {dyn_shape}")
     dynamic_shapes = dyn_shape == "True"
 
     x = torch.randn(shape_x, device=GPU_TYPE)
diff --git a/test/inductor/pallas_expected_failures/CpuTests.test_cat_unbacked_2d_cpu b/test/inductor/pallas_expected_failures/CpuTests.test_cat_unbacked_2d_cpu
deleted file mode 100644
index 7c33cf72ea286..0000000000000
--- a/test/inductor/pallas_expected_failures/CpuTests.test_cat_unbacked_2d_cpu
+++ /dev/null
@@ -1,5 +0,0 @@
-ERROR
-  File "/tmp/tmp1dzn2qvd/o2/co2o56k2yvm6po5mwz3gfv6o2cefusoc3ubuhgvtvnh6thorgd3i.py", line 11, in pallas_fused_new_ones_4d751881_kernel
-    out_ptr0[...] = tmp0.reshape(u0, 3)
-                                 ^^
-NameError: name 'u0' is not defined
diff --git a/test/inductor/pallas_expected_failures/CpuTests.test_cat_unbacked_empty_1d_cpu b/test/inductor/pallas_expected_failures/CpuTests.test_cat_unbacked_empty_1d_cpu
deleted file mode 100644
index d047626943cdd..0000000000000
--- a/test/inductor/pallas_expected_failures/CpuTests.test_cat_unbacked_empty_1d_cpu
+++ /dev/null
@@ -1,5 +0,0 @@
-ERROR
-  File "/tmp/tmp9kov386o/x3/cx3prc5mbesouvvwujf53kn7wcjra7dl7db375ai33k25m6osbkb.py", line 11, in pallas_fused_new_ones_eddcd9fa_kernel
-    out_ptr0[...] = tmp0.reshape(u0)
-                                 ^^
-NameError: name 'u0' is not defined
diff --git a/test/dynamo_expected_failures/CPython313-test_exception_variations-ExceptStarTestCases.test_nested_else b/test/inductor/pallas_expected_failures/CpuTests.test_cumprod_backward_cpu
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_exception_variations-ExceptStarTestCases.test_nested_else
rename to test/inductor/pallas_expected_failures/CpuTests.test_cumprod_backward_cpu
diff --git a/test/dynamo_expected_failures/CPython313-test_exception_variations-ExceptStarTestCases.test_nested_else_mixed2 b/test/inductor/pallas_expected_failures/CpuTests.test_cumprod_backward_with_zeros_cpu
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_exception_variations-ExceptStarTestCases.test_nested_else_mixed2
rename to test/inductor/pallas_expected_failures/CpuTests.test_cumprod_backward_with_zeros_cpu
diff --git a/test/dynamo_expected_failures/CPython313-test_exception_variations-ExceptStarTestCases.test_nested_mixed1 b/test/inductor/pallas_expected_failures/CpuTests.test_index_reduce_on_view_input_cpu
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_exception_variations-ExceptStarTestCases.test_nested_mixed1
rename to test/inductor/pallas_expected_failures/CpuTests.test_index_reduce_on_view_input_cpu
diff --git a/test/inductor/pallas_expected_failures/CpuTests.test_lerp_cpu b/test/inductor/pallas_expected_failures/CpuTests.test_lerp_cpu
deleted file mode 100644
index fb04b580a778d..0000000000000
--- a/test/inductor/pallas_expected_failures/CpuTests.test_lerp_cpu
+++ /dev/null
@@ -1,5 +0,0 @@
-ERROR
-  File "/home/oulgen/py312/lib/python3.12/site-packages/jax/_src/numpy/ufuncs.py", line 1569, in subtract
-    out = lax.sub(*promote_args("subtract", x, y))
-          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
-TypeError: sub got incompatible shapes for broadcasting: (10, 3, 10), (3, 10, 10).
diff --git a/test/inductor/pallas_expected_failures/CpuTests.test_log_softmax_cpu b/test/inductor/pallas_expected_failures/CpuTests.test_log_softmax_cpu
deleted file mode 100644
index 7802101049de9..0000000000000
--- a/test/inductor/pallas_expected_failures/CpuTests.test_log_softmax_cpu
+++ /dev/null
@@ -1,3 +0,0 @@
-FAIL
-AssertionError: Tensor-likes are not close!
-Numerical accuracy issue in log_softmax
diff --git a/test/inductor/pallas_expected_failures/CpuTests.test_logsumexp_cpu b/test/inductor/pallas_expected_failures/CpuTests.test_logsumexp_cpu
deleted file mode 100644
index aedfa92ebe36c..0000000000000
--- a/test/inductor/pallas_expected_failures/CpuTests.test_logsumexp_cpu
+++ /dev/null
@@ -1,5 +0,0 @@
-FAIL
-           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
-  File "/home/oulgen/dev/pytorch/torch/testing/_internal/common_utils.py", line 4280, in assertEqual
-    raise error_metas.pop()[0].to_error(  # type: ignore[index]
-AssertionError: Tensor-likes are not close!
diff --git a/test/inductor/pallas_expected_failures/CpuTests.test_max_pool2d_with_indices_backward2_cpu b/test/inductor/pallas_expected_failures/CpuTests.test_max_pool2d_with_indices_backward2_cpu
deleted file mode 100644
index a2f73f3214dfe..0000000000000
--- a/test/inductor/pallas_expected_failures/CpuTests.test_max_pool2d_with_indices_backward2_cpu
+++ /dev/null
@@ -1,5 +0,0 @@
-ERROR
-  File "/home/oulgen/py312/lib/python3.12/site-packages/jax/_src/numpy/ufuncs.py", line 1238, in add
-    out = lax.add(x, y)
-          ^^^^^^^^^^^^^
-TypeError: add got incompatible shapes for broadcasting: (56,), (40,).
diff --git a/test/inductor/pallas_expected_failures/CpuTests.test_max_pool2d_with_indices_backward3_cpu b/test/inductor/pallas_expected_failures/CpuTests.test_max_pool2d_with_indices_backward3_cpu
deleted file mode 100644
index 4af5f1937326f..0000000000000
--- a/test/inductor/pallas_expected_failures/CpuTests.test_max_pool2d_with_indices_backward3_cpu
+++ /dev/null
@@ -1,5 +0,0 @@
-ERROR
-  File "/home/oulgen/py312/lib/python3.12/site-packages/jax/_src/numpy/ufuncs.py", line 1238, in add
-    out = lax.add(x, y)
-          ^^^^^^^^^^^^^
-TypeError: add got incompatible shapes for broadcasting: (38,), (37,).
diff --git a/test/inductor/pallas_expected_failures/CpuTests.test_max_pool2d_with_indices_backward4_cpu b/test/inductor/pallas_expected_failures/CpuTests.test_max_pool2d_with_indices_backward4_cpu
deleted file mode 100644
index 3308d592808d9..0000000000000
--- a/test/inductor/pallas_expected_failures/CpuTests.test_max_pool2d_with_indices_backward4_cpu
+++ /dev/null
@@ -1,5 +0,0 @@
-ERROR
-  File "/home/oulgen/py312/lib/python3.12/site-packages/jax/_src/numpy/ufuncs.py", line 1238, in add
-    out = lax.add(x, y)
-          ^^^^^^^^^^^^^
-TypeError: add got incompatible shapes for broadcasting: (4,), (3,).
diff --git a/test/inductor/pallas_expected_failures/CpuTests.test_max_pool2d_with_indices_backward_cpu b/test/inductor/pallas_expected_failures/CpuTests.test_max_pool2d_with_indices_backward_cpu
deleted file mode 100644
index 53aaef6eda9cd..0000000000000
--- a/test/inductor/pallas_expected_failures/CpuTests.test_max_pool2d_with_indices_backward_cpu
+++ /dev/null
@@ -1,5 +0,0 @@
-ERROR
-  File "/home/oulgen/py312/lib/python3.12/site-packages/jax/_src/numpy/ufuncs.py", line 1238, in add
-    out = lax.add(x, y)
-          ^^^^^^^^^^^^^
-TypeError: add got incompatible shapes for broadcasting: (14,), (18,).
diff --git a/test/inductor/pallas_expected_failures/CpuTests.test_permute1_cpu b/test/inductor/pallas_expected_failures/CpuTests.test_permute1_cpu
deleted file mode 100644
index aedfa92ebe36c..0000000000000
--- a/test/inductor/pallas_expected_failures/CpuTests.test_permute1_cpu
+++ /dev/null
@@ -1,5 +0,0 @@
-FAIL
-           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
-  File "/home/oulgen/dev/pytorch/torch/testing/_internal/common_utils.py", line 4280, in assertEqual
-    raise error_metas.pop()[0].to_error(  # type: ignore[index]
-AssertionError: Tensor-likes are not close!
diff --git a/test/inductor/pallas_expected_failures/CpuTests.test_prod_cpu b/test/inductor/pallas_expected_failures/CpuTests.test_prod_cpu
deleted file mode 100644
index 51668572442cb..0000000000000
--- a/test/inductor/pallas_expected_failures/CpuTests.test_prod_cpu
+++ /dev/null
@@ -1,3 +0,0 @@
-FAIL
-AssertionError: Tensor-likes are not close!
-Numerical accuracy issue in prod
diff --git a/test/inductor/pallas_expected_failures/CpuTests.test_scaled_dot_product_attention_cpu b/test/inductor/pallas_expected_failures/CpuTests.test_scaled_dot_product_attention_cpu
deleted file mode 100644
index aedfa92ebe36c..0000000000000
--- a/test/inductor/pallas_expected_failures/CpuTests.test_scaled_dot_product_attention_cpu
+++ /dev/null
@@ -1,5 +0,0 @@
-FAIL
-           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
-  File "/home/oulgen/dev/pytorch/torch/testing/_internal/common_utils.py", line 4280, in assertEqual
-    raise error_metas.pop()[0].to_error(  # type: ignore[index]
-AssertionError: Tensor-likes are not close!
diff --git a/test/inductor/pallas_expected_failures/CpuTests.test_sgn_extremal_cpu b/test/inductor/pallas_expected_failures/CpuTests.test_sgn_extremal_cpu
deleted file mode 100644
index aedfa92ebe36c..0000000000000
--- a/test/inductor/pallas_expected_failures/CpuTests.test_sgn_extremal_cpu
+++ /dev/null
@@ -1,5 +0,0 @@
-FAIL
-           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
-  File "/home/oulgen/dev/pytorch/torch/testing/_internal/common_utils.py", line 4280, in assertEqual
-    raise error_metas.pop()[0].to_error(  # type: ignore[index]
-AssertionError: Tensor-likes are not close!
diff --git a/test/inductor/pallas_expected_failures/CpuTests.test_slice_scatter_backward_with_overlapping_base_cpu b/test/inductor/pallas_expected_failures/CpuTests.test_slice_scatter_backward_with_overlapping_base_cpu
new file mode 100644
index 0000000000000..2c3f456e070c8
--- /dev/null
+++ b/test/inductor/pallas_expected_failures/CpuTests.test_slice_scatter_backward_with_overlapping_base_cpu
@@ -0,0 +1,2 @@
+ERROR
+ValueError: Incompatible shapes for broadcasting
diff --git a/test/inductor/pallas_expected_failures/CpuTests.test_softmax_cpu b/test/inductor/pallas_expected_failures/CpuTests.test_softmax_cpu
deleted file mode 100644
index 567b0ecec3d30..0000000000000
--- a/test/inductor/pallas_expected_failures/CpuTests.test_softmax_cpu
+++ /dev/null
@@ -1,3 +0,0 @@
-FAIL
-AssertionError: Tensor-likes are not close!
-Numerical accuracy issue
diff --git a/test/inductor/pallas_expected_failures/CpuTests.test_std_cpu b/test/inductor/pallas_expected_failures/CpuTests.test_std_cpu
deleted file mode 100644
index 567b0ecec3d30..0000000000000
--- a/test/inductor/pallas_expected_failures/CpuTests.test_std_cpu
+++ /dev/null
@@ -1,3 +0,0 @@
-FAIL
-AssertionError: Tensor-likes are not close!
-Numerical accuracy issue
diff --git a/test/inductor/pallas_expected_failures/CpuTests.test_var_mean_tile_reduction_False_cpu b/test/inductor/pallas_expected_failures/CpuTests.test_var_mean_tile_reduction_False_cpu
deleted file mode 100644
index c150a58f6b601..0000000000000
--- a/test/inductor/pallas_expected_failures/CpuTests.test_var_mean_tile_reduction_False_cpu
+++ /dev/null
@@ -1,3 +0,0 @@
-ERROR
-TypeError: cannot reshape array of shape (4,) (size 4) into shape (1, 2, 4, 1) (size 8)
-Reshape mismatch in welford_reduce with tile_reduction
diff --git a/test/inductor/pallas_expected_failures/CpuTests.test_var_mean_tile_reduction_True_cpu b/test/inductor/pallas_expected_failures/CpuTests.test_var_mean_tile_reduction_True_cpu
deleted file mode 100644
index c150a58f6b601..0000000000000
--- a/test/inductor/pallas_expected_failures/CpuTests.test_var_mean_tile_reduction_True_cpu
+++ /dev/null
@@ -1,3 +0,0 @@
-ERROR
-TypeError: cannot reshape array of shape (4,) (size 4) into shape (1, 2, 4, 1) (size 8)
-Reshape mismatch in welford_reduce with tile_reduction
diff --git a/test/inductor/pallas_expected_failures/SweepInputsGPUTest.test_cuda_transposed_double b/test/inductor/pallas_expected_failures/SweepInputsGPUTest.test_cuda_transposed_double
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/test/inductor/pallas_expected_failures/SweepInputsTpuTest.test_tpu_broadcast1_dense b/test/inductor/pallas_expected_failures/SweepInputsTpuTest.test_tpu_broadcast1_dense
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/test/inductor/pallas_expected_failures/SweepInputsTpuTest.test_tpu_broadcast1_double b/test/inductor/pallas_expected_failures/SweepInputsTpuTest.test_tpu_broadcast1_double
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/test/inductor/pallas_expected_failures/SweepInputsTpuTest.test_tpu_broadcast2_broadcast1 b/test/inductor/pallas_expected_failures/SweepInputsTpuTest.test_tpu_broadcast2_broadcast1
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/test/inductor/pallas_expected_failures/SweepInputsTpuTest.test_tpu_broadcast2_dense b/test/inductor/pallas_expected_failures/SweepInputsTpuTest.test_tpu_broadcast2_dense
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/test/inductor/pallas_expected_failures/SweepInputsTpuTest.test_tpu_broadcast2_double b/test/inductor/pallas_expected_failures/SweepInputsTpuTest.test_tpu_broadcast2_double
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/test/inductor/pallas_expected_failures/SweepInputsTpuTest.test_tpu_broadcast2_strided b/test/inductor/pallas_expected_failures/SweepInputsTpuTest.test_tpu_broadcast2_strided
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/test/inductor/pallas_expected_failures/SweepInputsTpuTest.test_tpu_broadcast3_double b/test/inductor/pallas_expected_failures/SweepInputsTpuTest.test_tpu_broadcast3_double
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/test/inductor/pallas_expected_failures/SweepInputsTpuTest.test_tpu_dense_broadcast1 b/test/inductor/pallas_expected_failures/SweepInputsTpuTest.test_tpu_dense_broadcast1
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/test/inductor/pallas_expected_failures/SweepInputsTpuTest.test_tpu_dense_broadcast2 b/test/inductor/pallas_expected_failures/SweepInputsTpuTest.test_tpu_dense_broadcast2
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/test/inductor/pallas_expected_failures/SweepInputsTpuTest.test_tpu_dense_double b/test/inductor/pallas_expected_failures/SweepInputsTpuTest.test_tpu_dense_double
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/test/inductor/pallas_expected_failures/SweepInputsTpuTest.test_tpu_dense_strided b/test/inductor/pallas_expected_failures/SweepInputsTpuTest.test_tpu_dense_strided
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/test/inductor/pallas_expected_failures/SweepInputsTpuTest.test_tpu_dense_transposed b/test/inductor/pallas_expected_failures/SweepInputsTpuTest.test_tpu_dense_transposed
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/test/inductor/pallas_expected_failures/SweepInputsTpuTest.test_tpu_double_broadcast1 b/test/inductor/pallas_expected_failures/SweepInputsTpuTest.test_tpu_double_broadcast1
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/test/inductor/pallas_expected_failures/SweepInputsTpuTest.test_tpu_double_broadcast2 b/test/inductor/pallas_expected_failures/SweepInputsTpuTest.test_tpu_double_broadcast2
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/test/inductor/pallas_expected_failures/SweepInputsTpuTest.test_tpu_double_broadcast3 b/test/inductor/pallas_expected_failures/SweepInputsTpuTest.test_tpu_double_broadcast3
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/test/inductor/pallas_expected_failures/SweepInputsTpuTest.test_tpu_double_dense b/test/inductor/pallas_expected_failures/SweepInputsTpuTest.test_tpu_double_dense
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/test/inductor/pallas_expected_failures/SweepInputsTpuTest.test_tpu_double_double b/test/inductor/pallas_expected_failures/SweepInputsTpuTest.test_tpu_double_double
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/test/inductor/pallas_expected_failures/SweepInputsTpuTest.test_tpu_double_int b/test/inductor/pallas_expected_failures/SweepInputsTpuTest.test_tpu_double_int
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/test/inductor/pallas_expected_failures/SweepInputsTpuTest.test_tpu_double_strided b/test/inductor/pallas_expected_failures/SweepInputsTpuTest.test_tpu_double_strided
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/test/inductor/pallas_expected_failures/SweepInputsTpuTest.test_tpu_double_transposed b/test/inductor/pallas_expected_failures/SweepInputsTpuTest.test_tpu_double_transposed
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/test/inductor/pallas_expected_failures/SweepInputsTpuTest.test_tpu_int_double b/test/inductor/pallas_expected_failures/SweepInputsTpuTest.test_tpu_int_double
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/test/inductor/pallas_expected_failures/SweepInputsTpuTest.test_tpu_strided_broadcast1 b/test/inductor/pallas_expected_failures/SweepInputsTpuTest.test_tpu_strided_broadcast1
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/test/inductor/pallas_expected_failures/SweepInputsTpuTest.test_tpu_strided_broadcast2 b/test/inductor/pallas_expected_failures/SweepInputsTpuTest.test_tpu_strided_broadcast2
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/test/inductor/pallas_expected_failures/SweepInputsTpuTest.test_tpu_strided_dense b/test/inductor/pallas_expected_failures/SweepInputsTpuTest.test_tpu_strided_dense
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/test/inductor/pallas_expected_failures/SweepInputsTpuTest.test_tpu_strided_double b/test/inductor/pallas_expected_failures/SweepInputsTpuTest.test_tpu_strided_double
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/test/inductor/pallas_expected_failures/SweepInputsTpuTest.test_tpu_strided_strided b/test/inductor/pallas_expected_failures/SweepInputsTpuTest.test_tpu_strided_strided
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/test/inductor/pallas_expected_failures/SweepInputsTpuTest.test_tpu_strided_transposed b/test/inductor/pallas_expected_failures/SweepInputsTpuTest.test_tpu_strided_transposed
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/test/inductor/pallas_expected_failures/SweepInputsTpuTest.test_tpu_transposed_broadcast1 b/test/inductor/pallas_expected_failures/SweepInputsTpuTest.test_tpu_transposed_broadcast1
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/test/inductor/pallas_expected_failures/SweepInputsTpuTest.test_tpu_transposed_broadcast2 b/test/inductor/pallas_expected_failures/SweepInputsTpuTest.test_tpu_transposed_broadcast2
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/test/inductor/pallas_expected_failures/SweepInputsTpuTest.test_tpu_transposed_dense b/test/inductor/pallas_expected_failures/SweepInputsTpuTest.test_tpu_transposed_dense
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/test/inductor/pallas_expected_failures/SweepInputsTpuTest.test_tpu_transposed_double b/test/inductor/pallas_expected_failures/SweepInputsTpuTest.test_tpu_transposed_double
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/test/inductor/pallas_expected_failures/SweepInputsTpuTest.test_tpu_transposed_strided b/test/inductor/pallas_expected_failures/SweepInputsTpuTest.test_tpu_transposed_strided
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/test/dynamo_expected_failures/CPython313-test_exception_variations-ExceptStarTestCases.test_try_except b/test/inductor/pallas_skip_tests/CpuTests.test_large_block_sizes_cpu
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_exception_variations-ExceptStarTestCases.test_try_except
rename to test/inductor/pallas_skip_tests/CpuTests.test_large_block_sizes_cpu
diff --git a/test/inductor/test_aot_inductor.py b/test/inductor/test_aot_inductor.py
index 0f0fcdb842085..fc77808446116 100644
--- a/test/inductor/test_aot_inductor.py
+++ b/test/inductor/test_aot_inductor.py
@@ -1,4 +1,5 @@
 # Owner(s): ["module: inductor"]
+import contextlib
 import itertools
 import logging
 import os
@@ -37,16 +38,22 @@
 from torch._utils_internal import full_aoti_runtime_assert
 from torch.export import Dim, export
 from torch.export.pt2_archive._package import load_pt2
+from torch.nn.attention import (
+    activate_flash_attention_impl,
+    restore_flash_attention_impl,
+)
 from torch.testing import FileCheck
 from torch.testing._internal import common_utils
 from torch.testing._internal.common_cuda import (
     CDNA2OrLater,
+    CDNA5OrLater,
     PLATFORM_SUPPORTS_FLASH_ATTENTION,
     PLATFORM_SUPPORTS_FP8,
     PLATFORM_SUPPORTS_FP8_GROUPED_GEMM,
     PLATFORM_SUPPORTS_MEM_EFF_ATTENTION,
     requires_triton_ptxas_compat,
     SM80OrLater,
+    SM90OrLater,
     tf32_on_and_off,
 )
 from torch.testing._internal.common_device_type import (
@@ -55,6 +62,10 @@
     e5m2_type,
     skipCUDAIf,
 )
+from torch.testing._internal.common_dtype import (
+    highest_precision_complex,
+    highest_precision_float,
+)
 from torch.testing._internal.common_quantization import (
     _group_quantize_tensor,
     skip_if_no_torchvision,
@@ -62,7 +73,9 @@
 )
 from torch.testing._internal.common_utils import (
     DeterministicGuard,
+    IS_ARM64,
     IS_CI,
+    IS_CPU_CAPABILITY_SVE256,
     IS_FBCODE,
     IS_MACOS,
     IS_WINDOWS,
@@ -78,6 +91,7 @@
     skipIfXpu,
     TEST_MPS,
     TEST_WITH_ROCM,
+    xfailIf,
 )
 from torch.testing._internal.custom_tensor import CustomTensorPlainOut
 from torch.testing._internal.inductor_utils import (
@@ -95,6 +109,36 @@
 )
 
 
+f8_msg = "FP8 is only supported on H100+, SM 8.9 and MI300+, XPU and CPU devices"
+
+
+@contextlib.contextmanager
+def caching_allocator_disabled():
+    if GPU_TYPE == "cuda":
+        from torch.cuda import (
+            caching_allocator_disabled as _cuda_caching_allocator_disabled,
+        )
+
+        with _cuda_caching_allocator_disabled():
+            yield
+    else:
+        yield
+
+
+@contextlib.contextmanager
+def use_fa3():
+    try:
+        activate_flash_attention_impl("FA3")
+    except (ModuleNotFoundError, RuntimeError) as err:
+        raise unittest.SkipTest(
+            "FA3 backend not available (flash_attn_interface missing)"
+        ) from err
+    try:
+        yield
+    finally:
+        restore_flash_attention_impl()
+
+
 if HAS_GPU:
     import triton  # @manual
     from triton import language as tl
@@ -115,6 +159,7 @@
         add_kernel_with_tma_2d_new_api,
         add_kernel_with_tma_2d_old_api,
         create_tensor_descriptor_shim,
+        masked_add_kernel_with_bool_tensor,
         mul2_inplace_kernel,
         strange_config_matmul_kernel,
         sub_kernel_autotuned,
@@ -252,6 +297,31 @@ def forward(self, x):
         inputs = (torch.randn(4, device=self.device),)
         self.check_model(Model(), inputs)
 
+    def test_triton_kernel_bool_tensor_arg(self):
+        if self.device != GPU_TYPE or self.device == "mps":
+            raise unittest.SkipTest("requires GPU")
+
+        class Model(torch.nn.Module):
+            def forward(self, x, y, mask):
+                out = torch.zeros_like(x)
+                n = x.numel()
+                masked_add_kernel_with_bool_tensor[(n,)](
+                    in_ptr0=x,
+                    in_ptr1=y,
+                    mask_ptr=mask,
+                    out_ptr=out,
+                    n_elements=n,
+                    BLOCK_SIZE=1024,
+                )
+                return out
+
+        n = 128
+        x = torch.randn(n, device=self.device)
+        y = torch.randn(n, device=self.device)
+        mask = torch.arange(n, device=self.device) < n // 2
+        inputs = (x, y, mask)
+        self.check_model(Model(), inputs)
+
     @unittest.skipIf(
         IS_FBCODE,
         "toolchain doesn't support ptx to fatbin",
@@ -858,6 +928,8 @@ def forward(self, y):
         IS_FBCODE,
         "Not yet runnable in fbcode when the model.so is newly generated while older PyTorch is used",
     )
+    @xfailIf(IS_ARM64 and IS_CPU_CAPABILITY_SVE256)
+    # see https://github.com/pytorch/pytorch/issues/177243
     @tf32_on_and_off(0.005)
     def test_deconv_freezing(self):
         dtypes = [torch.float]
@@ -1279,14 +1351,11 @@ def forward(self, x, y):
         example_inputs = (x, y)
         self.check_model(Model(), example_inputs, dynamic_shapes=dynamic_shapes)
 
-    @unittest.skipIf(
-        not PLATFORM_SUPPORTS_FP8,
-        "FP8 is only supported on H100+, SM 8.9 and MI300+ devices",
-    )
-    @skipIfXpu
+    @unittest.skipIf(not PLATFORM_SUPPORTS_FP8, f8_msg)
+    @skipIfMPS
     def test_fp8(self):
         # cuda only
-        if self.device != "cuda":
+        if self.device not in ("cuda", "xpu"):
             return
 
         class Model(torch.nn.Module):
@@ -1330,7 +1399,6 @@ def forward(self, x, weight, bias, scale_a, scale_b):
         not PLATFORM_SUPPORTS_FP8_GROUPED_GEMM,
         "scaled_grouped_mm is only supported on SM90 and MI300+ devices",
     )
-    @skipIfXpu
     def test_scaled_grouped_mm(self):
         # Test torch._scaled_grouped_mm AOTI lowering
         # cuda only
@@ -1392,11 +1460,8 @@ def forward(self, x, weight, scale_a, scale_b, offsets):
             (x_fp8, weight_fp8, scale_a, scale_b, offsets),
         )
 
-    @unittest.skipIf(
-        not PLATFORM_SUPPORTS_FP8,
-        "FP8 is only supported on H100+, SM 8.9 and MI300+ devices",
-    )
-    @skipIfXpu
+    @unittest.skipIf(not PLATFORM_SUPPORTS_FP8, f8_msg)
+    @skipIfMPS
     def test_fp8_view_of_param(self):
         # cuda only
         if self.device != GPU_TYPE:
@@ -1566,6 +1631,47 @@ def forward(self, a, b):
             dynamic_shapes=dynamic_shapes,
         )
 
+    @unittest.skipIf(
+        not IS_BIG_GPU, "Skipping triton backend only since not big GPU (not enough SM)"
+    )
+    def test_bmm_large_batch_dynamic(self):
+        if self.device == "cpu":
+            raise unittest.SkipTest("using triton backend only is not supported on CPU")
+
+        class Model(torch.nn.Module):
+            def forward(self, a, b):
+                return torch.bmm(a, b)
+
+        M, K, N = 64, 64, 64
+        dtype = torch.float16
+        model = Model()
+
+        # Compile with small batch, then run with batch > 65535 (CUDA grid.y limit)
+        compile_batch = 100
+        a = torch.randn(compile_batch, M, K, device=self.device, dtype=dtype)
+        b = torch.randn(compile_batch, K, N, device=self.device, dtype=dtype)
+        dim0_a = Dim("dim0_a", min=1, max=2**17)
+        dynamic_shapes = {"a": {0: dim0_a}, "b": {0: dim0_a}}
+        list_example_inputs = [(a, b)]
+
+        # Large batch exceeding CUDA grid.y limit of 65535
+        large_batch = 70000
+        list_example_inputs.append(
+            (
+                torch.randn(large_batch, M, K, device=self.device, dtype=dtype),
+                torch.randn(large_batch, K, N, device=self.device, dtype=dtype),
+            ),
+        )
+        self.check_model_with_multiple_inputs(
+            model,
+            list_example_inputs,
+            options={
+                "max_autotune": True,
+                "max_autotune_gemm_backends": "TRITON",
+            },
+            dynamic_shapes=dynamic_shapes,
+        )
+
     @skipIfWindows(msg="TODO: (xuhancn) confirm, Crash: access violation")
     def test_foreach_multiple_dynamic(self):
         class Model(torch.nn.Module):
@@ -1684,6 +1790,8 @@ def forward(self, x):
         with config.patch({"aot_inductor.use_runtime_constant_folding": True}):
             self.check_model(Model(self.device), example_inputs)
 
+    @xfailIf(IS_ARM64)
+    # see https://github.com/pytorch/pytorch/issues/177254
     @skipIfNoFBGEMM
     def test_quanatized_int8_linear(self):
         class Model(torch.nn.Module):
@@ -1793,23 +1901,24 @@ def forward(self, x, y, lst):
                 return cat * cat
 
         # Disable cuda caching allocator to check for IMA
-        torch.cuda.caching_allocator_enable(False)
-        model = Repro()
-        example_inputs = (
-            # s0, s1
-            torch.randn((100, 200), device=self.device),
-            # s2, s3
-            torch.randn((100, 3), device=self.device),
-            # u0, u1, u2, u3, u100
-            torch.tensor([200, 100, 0, 1, 300], device=self.device, dtype=torch.int),
-        )
-        spec = {
-            "x": (Dim.DYNAMIC, Dim.DYNAMIC),
-            "y": (Dim.DYNAMIC, Dim.DYNAMIC),
-            "lst": (Dim.STATIC,),
-        }
-        self.check_model(model, example_inputs, dynamic_shapes=spec)
-        torch.cuda.caching_allocator_enable(True)
+        with caching_allocator_disabled():
+            model = Repro()
+            example_inputs = (
+                # s0, s1
+                torch.randn((100, 200), device=self.device),
+                # s2, s3
+                torch.randn((100, 3), device=self.device),
+                # u0, u1, u2, u3, u100
+                torch.tensor(
+                    [200, 100, 0, 1, 300], device=self.device, dtype=torch.int
+                ),
+            )
+            spec = {
+                "x": (Dim.DYNAMIC, Dim.DYNAMIC),
+                "y": (Dim.DYNAMIC, Dim.DYNAMIC),
+                "lst": (Dim.STATIC,),
+            }
+            self.check_model(model, example_inputs, dynamic_shapes=spec)
 
     @skipIfMPS
     @config.patch({"unbacked_symint_fallback": 12})
@@ -1851,20 +1960,87 @@ def forward(self, x, y, lengths):
                 relevant_embeddings += ones
                 return relevant_embeddings * relevant_embeddings
 
-        torch.cuda.caching_allocator_enable(False)
-        model = Repro()
-        example_inputs = (
-            torch.randn((1000, INNER_DIM), device=self.device),
-            torch.randn((2000, INNER_DIM), device=self.device),
-            torch.ones(3000),
-        )
-        spec = {
-            "x": (Dim.DYNAMIC, Dim.STATIC),
-            "y": (Dim.DYNAMIC, Dim.STATIC),
-            "lengths": (Dim.DYNAMIC,),
-        }
-        self.check_model(model, example_inputs, dynamic_shapes=spec)
-        torch.cuda.caching_allocator_enable(True)
+        with caching_allocator_disabled():
+            model = Repro()
+            example_inputs = (
+                torch.randn((1000, INNER_DIM), device=self.device),
+                torch.randn((2000, INNER_DIM), device=self.device),
+                torch.ones(3000),
+            )
+            spec = {
+                "x": (Dim.DYNAMIC, Dim.STATIC),
+                "y": (Dim.DYNAMIC, Dim.STATIC),
+                "lengths": (Dim.DYNAMIC,),
+            }
+            self.check_model(model, example_inputs, dynamic_shapes=spec)
+
+    @skipIfMPS
+    @config.patch({"triton.autotune_at_compile_time": None})
+    @torch.fx.experimental._config.patch("backed_size_oblivious", True)
+    def test_slice_independent_backed_symints_no_unbacked(self):
+        # x[0:s1] where x.size(0) = s0-1 should produce Min(s1, s0-1),
+        # not an unbacked symint with a bad fallback value.
+        if self.device != GPU_TYPE:
+            raise unittest.SkipTest("requires triton")
+
+        INNER_DIM = 4224
+
+        class Repro(torch.nn.Module):
+            def forward(self, x, y):
+                x_trimmed = x[:-1]
+                sliced = x_trimmed[: y.size(0)]
+                reshaped = sliced.reshape(-1, 128, 33)
+                expanded = reshaped.unsqueeze(3).expand(-1, 128, 33, 8)
+                shifts = torch.arange(0, 64, 8, device=x.device, dtype=torch.int64)
+                return (expanded >> shifts) & 255
+
+        with caching_allocator_disabled():
+            model = Repro()
+            example_inputs = (
+                torch.randint(
+                    0, 256, (200, INNER_DIM), device=self.device, dtype=torch.int64
+                ),
+                torch.randn(50, 8, device=self.device),
+            )
+            spec = {
+                "x": (Dim.DYNAMIC, Dim.STATIC),
+                "y": (Dim.DYNAMIC, Dim.STATIC),
+            }
+            self.check_model(model, example_inputs, dynamic_shapes=spec)
+
+    @skipIfMPS
+    @config.patch({"triton.autotune_at_compile_time": None})
+    @torch.fx.experimental._config.patch("backed_size_oblivious", True)
+    def test_slice_negative_index_backed_symints_no_unbacked(self):
+        # x[-s1:] where x.size(0) = s0-1 should produce Max(s0-1 - s1, 0),
+        # not an unbacked symint with a bad fallback value.
+        if self.device != GPU_TYPE:
+            raise unittest.SkipTest("requires triton")
+
+        INNER_DIM = 4224
+
+        class Repro(torch.nn.Module):
+            def forward(self, x, y):
+                x_trimmed = x[:-1]
+                sliced = x_trimmed[-y.size(0) :]
+                reshaped = sliced.reshape(-1, 128, 33)
+                expanded = reshaped.unsqueeze(3).expand(-1, 128, 33, 8)
+                shifts = torch.arange(0, 64, 8, device=x.device, dtype=torch.int64)
+                return (expanded >> shifts) & 255
+
+        with caching_allocator_disabled():
+            model = Repro()
+            example_inputs = (
+                torch.randint(
+                    0, 256, (200, INNER_DIM), device=self.device, dtype=torch.int64
+                ),
+                torch.randn(50, 8, device=self.device),
+            )
+            spec = {
+                "x": (Dim.DYNAMIC, Dim.STATIC),
+                "y": (Dim.DYNAMIC, Dim.STATIC),
+            }
+            self.check_model(model, example_inputs, dynamic_shapes=spec)
 
     @config.patch({"triton.autotune_at_compile_time": None})
     def test_stride_with_unbacked_expr(self):
@@ -1927,6 +2103,7 @@ def forward(self, values, repeats, mask, embeddings, x, z, scalar):
         self.check_model(Repro(), example_inputs, dynamic_shapes=spec)
 
     @skipIfWindowsXPU(msg="crash on Windows XPU.")
+    @config.patch({"unbacked_symint_fallback": 128})
     def test_size_with_unbacked_add_expr_transitive(self):
         # Edge case with torch._check(expr1, expr2) + torch._check(expr2, unbacked).
         # When generating example input sizes for autotuning, it should coalesce
@@ -2029,7 +2206,9 @@ def forward(self, values, repeats, mask, embeddings, x, y, z, lst):
         }
         self.check_model(Repro(), example_inputs, dynamic_shapes=spec)
 
-    @skipIfXpu(msg="_scaled_dot_product_flash_attention is not supported on XPU yet")
+    @skipIfXpu(
+        msg="FlashAttentionForward headdim limitation on xpu - torch-xpu-ops: 2698"
+    )
     @unittest.skipIf(
         not PLATFORM_SUPPORTS_FLASH_ATTENTION, "Some archs don't support flash SDPA"
     )
@@ -2480,8 +2659,8 @@ def test_cond_cpu_predicate_cuda_operands(self, max_autotune):
         determined device from [predicate] + operands, causing CPU predicates
         to force CUDA outputs onto CPU during autotuning.
         """
-        if self.device != "cuda":
-            raise unittest.SkipTest("requires CUDA")
+        if self.device != "cuda" and self.device != "xpu":
+            raise unittest.SkipTest("requires CUDA or XPU")
 
         class Model(torch.nn.Module):
             def __init__(self, input_dim=4, hidden_dim=8):
@@ -2897,6 +3076,7 @@ def forward(self, x):
         torch._export.aot_compile(Model(), example_inputs)
 
     @skipCUDAIf(True, "Test for x86 backend")
+    @skipIfXpu(msg="Test for x86 backend")
     @unittest.skipIf(sys.platform == "darwin", "Skip MacOS")
     @unittest.skipIf(IS_FBCODE, "Need newer ideep")
     def test_buffer_mutation_and_force_mmap_weights(self):
@@ -3298,10 +3478,12 @@ def forward(self, x):
         mod(inp)
         mod2 = torch.fx.symbolic_trace(mod, concrete_args=[inp])
         so = torch._export.aot_compile(mod2, (inp,))
-        assert so is not None
+        if so is None:
+            raise AssertionError("Expected aot_compile to return non-None")
         # compile the 2nd time with cache hit
         so = torch._export.aot_compile(mod2, (inp,))
-        assert so is not None
+        if so is None:
+            raise AssertionError("Expected aot_compile to return non-None (cache hit)")
 
     def test_normal_functional(self):
         class Model(torch.nn.Module):
@@ -3502,6 +3684,7 @@ def forward(self, x):
         example_inputs = (torch.randn(3, 10, device=self.device),)
         self.check_model(Model(), example_inputs)
 
+    @skipIfRocmArch(NAVI_ARCH)  # regression on ROCm 7.2
     def test_repeated_calling(self):
         if self.device != "cuda":
             raise unittest.SkipTest("requires CUDA")
@@ -3519,18 +3702,20 @@ def forward(self, x):
                 torch.export.export(Model(), example_inputs, strict=True)
             )
         )
-        try:
-            torch.cuda.memory.empty_cache()
-            torch.cuda.memory._record_memory_history(context=None)
-            for _ in range(10):
-                optimized(*example_inputs)
-        finally:
-            torch.cuda.memory._record_memory_history(False)
-        segments = torch.cuda.memory._snapshot()["segments"]
-        self.assertTrue(
-            any(seg["requested_size"] == 400 for seg in segments),
-            f"Expected segment with size 400, got: {[s['requested_size'] for s in segments]}",
-        )
+        expected = torch.sin(example_inputs[0])
+
+        # Warm up to trigger any one-time allocations
+        result = optimized(*example_inputs)
+        torch.cuda.synchronize()
+
+        mem_before = torch.cuda.memory_allocated()
+        for _ in range(10):
+            result = optimized(*example_inputs)
+        torch.cuda.synchronize()
+        mem_after = torch.cuda.memory_allocated()
+
+        self.assertEqual(result, expected)
+        self.assertEqual(mem_before, mem_after)
 
     def test_view_outputs(self):
         class Model(torch.nn.Module):
@@ -3573,7 +3758,7 @@ def forward(self, x):
 
         # Call eval() here so that batch_norm won't update the running stats
         # Use float64 to avoid numeric difference failure
-        dtype = torch.float32 if self.device == "mps" else torch.float64
+        dtype = highest_precision_float(self.device)
         model = Model().to(device=self.device, dtype=dtype).eval()
         example_inputs = (torch.randn(4, 3, 64, 64, device=self.device, dtype=dtype),)
         self.check_model(model, example_inputs)
@@ -4747,6 +4932,67 @@ def forward(self, q, k, v, attn_bias):
         )
         self.check_model(Model(), example_inputs)
 
+    @unittest.skipIf(not SM90OrLater, "FA3 requires SM90+")
+    def test_varlen_attn_paged_kv_cache(self):
+        if self.device != GPU_TYPE:
+            raise unittest.SkipTest("requires GPU")
+
+        from torch.nn.attention.varlen import varlen_attn
+
+        batch_size, num_heads, head_dim = 2, 4, 64
+        page_size, max_pages_per_seq = 64, 2
+        cache_seqlen, new_seqlen = 32, 1
+
+        class Model(torch.nn.Module):
+            def forward(self, q, k_pages, v_pages, cu_seq_q, seqused_k, page_table):
+                return varlen_attn(
+                    query=q,
+                    key=k_pages,
+                    value=v_pages,
+                    cu_seq_q=cu_seq_q,
+                    cu_seq_k=None,
+                    max_q=new_seqlen,
+                    max_k=max_pages_per_seq * page_size,
+                    seqused_k=seqused_k,
+                    page_table=page_table,
+                )
+
+        total_new = new_seqlen * batch_size
+        total_pages = batch_size * max_pages_per_seq
+        example_inputs = (
+            torch.randn(
+                total_new, num_heads, head_dim, dtype=torch.bfloat16, device=self.device
+            ),
+            torch.randn(
+                total_pages,
+                page_size,
+                num_heads,
+                head_dim,
+                dtype=torch.bfloat16,
+                device=self.device,
+            ),
+            torch.randn(
+                total_pages,
+                page_size,
+                num_heads,
+                head_dim,
+                dtype=torch.bfloat16,
+                device=self.device,
+            ),
+            torch.arange(
+                0, total_new + 1, new_seqlen, dtype=torch.int32, device=self.device
+            ),
+            torch.full(
+                (batch_size,), cache_seqlen, dtype=torch.int32, device=self.device
+            ),
+            torch.arange(total_pages, dtype=torch.int32, device=self.device).reshape(
+                batch_size, max_pages_per_seq
+            ),
+        )
+
+        with use_fa3():
+            self.check_model(Model(), example_inputs)
+
     def test_aoti_runtime_asserts(self):
         from torch.export._draft_export import draft_export, FailureType
 
@@ -4939,10 +5185,7 @@ def forward(self, inputs):
 
         self.check_model(m, inputs)
 
-    @unittest.skipIf(
-        not PLATFORM_SUPPORTS_FP8,
-        "FP8 is only supported on H100+, SM 8.9 and MI300+ devices",
-    )
+    @unittest.skipIf(not PLATFORM_SUPPORTS_FP8, f8_msg)
     @patch.dict(os.environ, {"AOTI_RUNTIME_CHECK_INPUTS": "1"})
     def test_runtime_checks_fp8(self):
         # cuda only
@@ -5013,7 +5256,11 @@ def forward(self, x0, x1, x2):
             dtype=torch.complex64,
             device=self.device,
         )
-        x2 = torch.tensor(128, dtype=torch.complex128, device=self.device)
+        x2 = torch.tensor(
+            128,
+            dtype=highest_precision_complex(self.device),
+            device=self.device,
+        )
         inputs.append(x0)
         inputs.append(x1)
         inputs.append(x2)
@@ -5471,6 +5718,52 @@ def forward(self, a):
 
             self.check_model(Model(N, K, self.device), example_inputs)
 
+    @unittest.skipIf(
+        sys.platform not in ["linux", "win32"],
+        "enable_kernel_profile only supported on linux and win32",
+    )
+    def test_kernel_profile_repeated_kernel_numel(self):
+        # When enable_kernel_profile is on, each kernel call is wrapped in its
+        # own {} scope block. If the same kernel is called multiple times with
+        # symbolic numel, the int64_t declaration must be emitted each time
+        # (not just on first use) since prior declarations go out of scope.
+        class Model(torch.nn.Module):
+            def forward(self, x, y, z):
+                return torch.cat([x, y, z], dim=0)
+
+        example_inputs = tuple(torch.randn(4, 8, device=self.device) for _ in range(3))
+        dim0 = Dim("dim0", min=1, max=32)
+        dynamic_shapes = {"x": {0: dim0}, "y": {0: dim0}, "z": {0: dim0}}
+
+        with config.patch({"cpp.enable_kernel_profile": True}):
+            _, code = run_and_get_cpp_code(
+                AOTIRunnerUtil.compile,
+                Model(),
+                example_inputs,
+                dynamic_shapes=dynamic_shapes,
+            )
+            # When profiling is enabled, every kernel numel variable must have
+            # the int64_t type declaration since each kernel call lives in its
+            # own scope block. Verify no bare assignment (without int64_t)
+            # appears for numel variables.
+            if self.device == GPU_TYPE:
+                # Match bare numel assignments like "foo_xnumel = expr;"
+                # but not declarations like "int64_t foo_xnumel = expr;"
+                bare_numel_assign = re.compile(r"^\s*(\w+_[xr]numel)\s*=\s*.+;$")
+                for line in code.splitlines():
+                    m = bare_numel_assign.match(line)
+                    if m:
+                        self.fail(
+                            f"Found numel assignment without int64_t declaration "
+                            f"in profiling mode: {line.strip()}"
+                        )
+
+            self.check_model(
+                Model(),
+                example_inputs,
+                dynamic_shapes=dynamic_shapes,
+            )
+
     def test_aoti_user_defined_triton_kernel_profiling(self):
         if self.device != GPU_TYPE or self.device == "mps":
             raise unittest.SkipTest("requires GPU")
@@ -5632,11 +5925,8 @@ def forward(self, x):
                     2,
                 ).run(code)
 
-    @unittest.skipIf(
-        not PLATFORM_SUPPORTS_FP8,
-        "FP8 is only supported on H100+, SM 8.9 and MI300+ devices",
-    )
-    @skipIfXpu
+    @unittest.skipIf(not PLATFORM_SUPPORTS_FP8, f8_msg)
+    @skipIfMPS
     def test_aoti_debug_printer_fp8_dtype(self):
         if self.device != GPU_TYPE:
             raise unittest.SkipTest("requires GPU")
@@ -6789,7 +7079,6 @@ def forward(self, x):
         example_inputs = (torch.randn(500, device=self.device),)
         self.check_model(model, example_inputs)
 
-    @skipIfXpu
     def test_conv3d(self):
         if self.device != GPU_TYPE or not is_big_gpu():
             raise unittest.SkipTest("requires modern GPU to run max-autotune")
@@ -6886,6 +7175,8 @@ def test__weight_int4pack_mm(self, m, n, q_group, num_groups):
             raise unittest.SkipTest("requires GPU")
 
         if TEST_WITH_ROCM:
+            if CDNA5OrLater():
+                self.skipTest("int4 mm not yet implemented for gfx1250 (needs WMMA)")
             if not CDNA2OrLater():
                 self.skipTest("_int4_mm is supported only for CDNA2 or later")
 
@@ -6923,6 +7214,8 @@ def test__weight_int4pack_mm_with_scales_and_zeros(self, m, n, q_group, num_grou
             raise unittest.SkipTest("requires Intel GPU")
 
         if TEST_WITH_ROCM:
+            if CDNA5OrLater():
+                self.skipTest("int4 mm not yet implemented for gfx1250 (needs WMMA)")
             if not CDNA2OrLater():
                 self.skipTest("_int4_mm is supported only for CDNA2 or later")
 
@@ -6940,30 +7233,39 @@ def forward(self, a):
 
         def _group_quantize_tensor_xpu(w, n_bit=4, q_group_size=16):
             # w [k, n] = [32, 48]
-            assert w.dim() == 2
+            if w.dim() != 2:
+                raise AssertionError(f"Expected 2D tensor, got {w.dim()}D")
             # w [n, k] = [48, 32]
             w = w.transpose(0, 1).contiguous()
-            assert q_group_size > 1
-            assert w.shape[-1] % q_group_size == 0
+            if q_group_size <= 1:
+                raise AssertionError(f"Expected q_group_size > 1, got {q_group_size}")
+            if w.shape[-1] % q_group_size != 0:
+                raise AssertionError(
+                    f"w.shape[-1] ({w.shape[-1]}) must be divisible by q_group_size ({q_group_size})"
+                )
 
             # to_quant: [n * k / group_size, group_size]
             to_quant = w.reshape(-1, q_group_size)
-            assert torch.isnan(to_quant).sum() == 0
+            if torch.isnan(to_quant).sum() != 0:
+                raise AssertionError("to_quant contains NaN values")
 
             max_val = to_quant.amax(dim=1, keepdim=True)
             min_val = to_quant.amin(dim=1, keepdim=True)
             max_int = 2**n_bit - 1
             min_int = 0
             scales = (max_val - min_val).clamp(min=1e-6) / max_int
-            assert torch.isnan(scales).sum() == 0
+            if torch.isnan(scales).sum() != 0:
+                raise AssertionError("scales contains NaN values")
 
             zeros = min_int - min_val.div(scales).round()
             zeros = torch.clamp(zeros, min_int, max_int)
             zeros = zeros.to(torch.int8)
-            assert torch.isnan(zeros).sum() == 0
+            if torch.isnan(zeros).sum() != 0:
+                raise AssertionError("zeros contains NaN values")
 
             out = to_quant.div(scales).add(zeros).round().clamp_(min_int, max_int)
-            assert torch.isnan(out).sum() == 0
+            if torch.isnan(out).sum() != 0:
+                raise AssertionError("out contains NaN values")
 
             # [n, k]
             out = out.to(dtype=torch.int32).reshape(w.shape)
@@ -7088,34 +7390,22 @@ def forward(self, x, y, m):
         m = torch.tensor([4096], dtype=torch.int32, device=self.device)
 
         with config.patch("triton.autotune_with_sample_inputs", True):
-            if torch.version.hip:
-                # ROCm: Use dynamic grid checking (portable across different configs)
-                # Compile and get the generated code
-                _, src_code = run_and_get_cpp_code(
-                    torch._export.aot_compile, Model(), (x, y, m)
-                )
-                actual_grid, expected_grids = get_triton_grid_info(
-                    strange_config_matmul_kernel, 4096 * 2046, src_code
-                )
-                self.assertTrue(
-                    actual_grid is not None, "Could not find grid_0 in generated code"
-                )
-                self.assertIn(
-                    actual_grid,
-                    expected_grids,
-                    f"grid_0={actual_grid} not in expected {expected_grids} from kernel configs",
-                )
-
-            else:
-                # CUDA/XPU: Keep existing hardcoded values
-                # The tuned best config on XPU is different with CUDA.
-                if GPU_TYPE == "xpu":
-                    grid_0 = 32736
-                else:
-                    grid_0 = 1023
-                self.code_check_count(
-                    Model(), (x, y, m), f"uint32_t grid_0 = {grid_0}L;", 1
-                )
+            # Use dynamic grid checking (portable across different configs
+            # and Triton versions)
+            _, src_code = run_and_get_cpp_code(
+                torch._export.aot_compile, Model(), (x, y, m)
+            )
+            actual_grid, expected_grids = get_triton_grid_info(
+                strange_config_matmul_kernel, 4096 * 2046, src_code
+            )
+            self.assertTrue(
+                actual_grid is not None, "Could not find grid_0 in generated code"
+            )
+            self.assertIn(
+                actual_grid,
+                expected_grids,
+                f"grid_0={actual_grid} not in expected {expected_grids} from kernel configs",
+            )
 
     def test_triton_mutated_autotuning(self):
         if self.device != GPU_TYPE:
@@ -7158,34 +7448,22 @@ def forward(self, x, y, m):
         m = torch.tensor([4095], dtype=torch.int32, device=self.device)
 
         with config.patch("triton.autotune_with_sample_inputs", True):
-            if torch.version.hip:
-                # ROCm: Use dynamic grid checking (portable across different configs)
-                # Compile and get the generated code
-                _, src_code = run_and_get_cpp_code(
-                    torch._export.aot_compile, Model(), (x, y, m)
-                )
-                actual_grid, expected_grids = get_triton_grid_info(
-                    strange_config_matmul_kernel, 4096 * 2046, src_code
-                )
-                self.assertTrue(
-                    actual_grid is not None, "Could not find grid_0 in generated code"
-                )
-                self.assertIn(
-                    actual_grid,
-                    expected_grids,
-                    f"grid_0={actual_grid} not in expected {expected_grids} from kernel configs",
-                )
-
-            else:
-                # CUDA/XPU: Keep existing hardcoded values
-                # The tuned best config on XPU is different with CUDA.
-                if GPU_TYPE == "xpu":
-                    grid_0 = 32736
-                else:
-                    grid_0 = 1023
-                self.code_check_count(
-                    Model(), (x, y, m), f"uint32_t grid_0 = {grid_0}L;", 1
-                )
+            # Use dynamic grid checking (portable across different configs
+            # and Triton versions)
+            _, src_code = run_and_get_cpp_code(
+                torch._export.aot_compile, Model(), (x, y, m)
+            )
+            actual_grid, expected_grids = get_triton_grid_info(
+                strange_config_matmul_kernel, 4096 * 2046, src_code
+            )
+            self.assertTrue(
+                actual_grid is not None, "Could not find grid_0 in generated code"
+            )
+            self.assertIn(
+                actual_grid,
+                expected_grids,
+                f"grid_0={actual_grid} not in expected {expected_grids} from kernel configs",
+            )
 
     @patch.dict(os.environ, {"TRITON_DEBUG": "1"})
     def test_triton_dynamic_launcher_grid(self):
@@ -7645,8 +7923,8 @@ def test_codegen_int_array_var_fix_memory_leak(self):
         """
         Fix https://github.com/pytorch/pytorch/issues/167630
         """
-        if self.device != "cuda":
-            raise unittest.SkipTest("test is only for cuda")
+        if self.device not in ("cuda", "xpu"):
+            raise unittest.SkipTest("test is only for cuda or xpu")
 
         def make_mlp(in_dim=128, hidden=256, out_dim=64, depth=3):
             layers = []
@@ -7667,7 +7945,7 @@ def make_mlp(in_dim=128, hidden=256, out_dim=64, depth=3):
 
         allocated_memory = []
         for _ in range(3):
-            torch.cuda.reset_peak_memory_stats()
+            torch.accelerator.reset_peak_memory_stats()
 
             model = make_mlp(in_dim, hidden, out_dim, depth).to(self.device)
             example_inputs = (torch.randn(batch, in_dim, device=self.device),)
@@ -7678,10 +7956,10 @@ def make_mlp(in_dim=128, hidden=256, out_dim=64, depth=3):
             torch._inductor.aoti_compile_and_package(ep)
 
             del model, example_inputs, ep
-            torch.cuda.synchronize()
-            torch.cuda.empty_cache()
+            torch.accelerator.synchronize()
+            torch.accelerator.empty_cache()
             gc.collect()
-            allocated_memory.append(torch.cuda.memory_allocated())
+            allocated_memory.append(torch.accelerator.memory_allocated())
 
         self.assertTrue(allocated_memory[1] == allocated_memory[2])
 
@@ -7908,6 +8186,172 @@ def forward(self, x):
                 f"Failed to load package in subprocess: {result.stdout + result.stderr}",
             )
 
+    @unittest.skip(
+        "Skip this test, only for local test. SIGFPE is produced when viewing "
+        "empty tensor with dim=0. See D86125095 for model-side workaround."
+    )
+    def test_view_zero_dim_empty_tensor(self):
+        """
+        Regression test for viewing an empty tensor with dim=0 in AOTInductor.
+
+        When calling .view(s55, -1, 52) where:
+        - The tensor being viewed has size u1 (unbacked symint from .item())
+        - s55 is a backed symint from a DIFFERENT tensor's .size(0)
+        - s55 can be 0 at runtime
+
+        Previously, the AOTInductor C++ runtime crashed with SIGFPE because:
+        - The -1 computation is: (u1*520) / (s55*52)
+        - When s55=0, this is division by zero
+
+        Now, we expect a proper RuntimeError to be raised with a meaningful
+        error message instead of crashing the process.
+
+        Key pattern from exported_model_graph_aoti.txt line 1182:
+        view_90: "bf16[s55, ((10*u1)//s55), 52]" = _broadcast_impl.view(getitem_1, -1, 52)
+        - _broadcast_impl has shape [u1, 520] (u1 is unbacked from .item())
+        - getitem_1 is s55 (backed from embeddings_1.size(0))
+        """
+
+        class Model(torch.nn.Module):
+            def forward(self, batch_sizes, embeddings):
+                # batch_sizes: i64[s8] - batch size counts per item
+                # embeddings: bf16[s55, s12] - embeddings tensor (s55 can be 0)
+                # s8 and s55 are INDEPENDENT backed symints
+
+                # Get s55 from embeddings.size(0) - backed symint
+                s55 = embeddings.size(0)
+
+                # Get u1 via sum().item() - unbacked symint
+                # This mimics: item_1: "Sym(u1)" = sum_2.item()
+                u1 = batch_sizes.sum().item()
+                torch._check(u1 >= 0)
+
+                # Create values tensor with shape [s8, 520]
+                s8 = batch_sizes.size(0)
+                values = embeddings.new_zeros(s8, 520)
+
+                # Create data via repeat_interleave - result has shape [u1, 520]
+                # This mimics:
+                #   idx64: "i64[u1]" = torch.repeat_interleave(arange, batch_sizes, ...)
+                #   _broadcast_impl: "bf16[u1, 520]" = cat_2[idx64]
+                data = torch.repeat_interleave(values, batch_sizes, dim=0)
+
+                # The problematic view: uses s55 (backed from embeddings, can be 0)
+                # but the tensor data has size u1 (unbacked) - DECOUPLED from s55
+                # When s55=0: -1 = (u1*520) / (s55*52) = 0/0
+                # This mimics line 1182:
+                #   view_90 = _broadcast_impl.view(getitem_1, -1, 52)
+                result = data.view(s55, -1, 52)
+
+                return result
+
+        # Compile with non-empty tensors
+        # At compile time: s8=3, s55=5, u1=5 (sum of batch_sizes = 2+2+1 = 5)
+        compile_batch_sizes = torch.tensor(
+            [2, 2, 1], dtype=torch.int64, device=self.device
+        )
+        compile_embeddings = torch.randn(5, 104, device=self.device)
+        compile_inputs = (compile_batch_sizes, compile_embeddings)
+
+        # Define dynamic shapes - s55 can be 0 at runtime!
+        s8 = Dim("s8", min=1, max=1024)
+        s55 = Dim("s55", min=0, max=1024)  # Can be 0!
+        s12 = Dim("s12", min=1, max=1024)
+        dynamic_shapes = {
+            "batch_sizes": {0: s8},
+            "embeddings": {0: s55, 1: s12},
+        }
+
+        # Export and compile
+        model = Model().to(self.device)
+        ep = torch.export.export(
+            model, compile_inputs, dynamic_shapes=dynamic_shapes, strict=False
+        )
+
+        package_path = torch._inductor.aoti_compile_and_package(ep)
+
+        with zipfile.ZipFile(package_path, "r") as z:
+            for name in z.namelist():
+                if name.endswith("wrapper.cpp"):
+                    content = z.read(name).decode("utf-8", errors="replace")
+                    # Verify that division and modulo operations are guarded
+                    # The new implementation emits AOTI_TORCH_CHECK calls before
+                    # computing size expressions that contain division/modulo
+                    FileCheck().check("AOTI_TORCH_CHECK(").check("by zero").run(content)
+
+        optimized = torch._inductor.aoti_load_package(package_path)
+
+        # Run with s55=0 (embeddings is empty) and u1=0 (batch_sizes sum to 0)
+        # data has shape [0, 520] (u1=0)
+        # view(s55=0, -1, 52): -1 = 0*520 / (0*52) = 0/0 = division by zero
+        # sum=0, u1=0
+        run_batch_sizes = torch.tensor([0], dtype=torch.int64, device=self.device)
+        run_embeddings = torch.randn(0, 104, device=self.device)  # s55=0
+        run_inputs = (run_batch_sizes, run_embeddings)
+
+        # Should raise RuntimeError instead of crashing with SIGFPE.
+        # The generated AOTI_TORCH_CHECK logs the error (including "by zero") to stderr
+        # before returning an error code, which is then converted to a RuntimeError.
+        # We capture stderr at the file descriptor level to verify the message is logged.
+        import tempfile
+
+        # Save original stderr fd and redirect to temp file to capture C++ LOG(ERROR) output
+        original_stderr_fd = os.dup(2)
+        with tempfile.NamedTemporaryFile(mode="w+", delete=False) as temp_stderr:
+            temp_stderr_path = temp_stderr.name
+
+        try:
+            with open(temp_stderr_path, "w") as new_stderr:
+                os.dup2(new_stderr.fileno(), 2)
+                try:
+                    with self.assertRaises(RuntimeError):
+                        optimized(*run_inputs)
+                finally:
+                    # Flush and restore original stderr
+                    new_stderr.flush()
+                    os.dup2(original_stderr_fd, 2)
+                    os.close(original_stderr_fd)
+
+            # Read captured stderr and verify the "by zero" message was logged
+            with open(temp_stderr_path) as f:
+                stderr_content = f.read()
+            self.assertIn("by zero", stderr_content)
+        finally:
+            # Clean up temp file
+            if os.path.exists(temp_stderr_path):
+                os.unlink(temp_stderr_path)
+
+    def test_combo_kernel_grid_mixed_types(self):
+        if self.device != GPU_TYPE:
+            raise unittest.SkipTest("combo kernels require GPU")
+
+        class Model(torch.nn.Module):
+            def forward(self, a, b, c, d):
+                return a + b, c + d
+
+        # Non-contiguous (transposed) tensors force 2D Triton iteration,
+        # creating both x and y range trees in the combo kernel.
+        # Dynamic shapes on (a, b) make ynumel_0 dynamic (string in codegen),
+        # while static shapes on (c, d) keep ynumel_1 as int.
+        # This triggers GridExpr.maximum() with mixed [str, int] in C++ mode,
+        # which previously generated std::max(long, int) causing a template
+        # deduction error in the AOTInductor C++ wrapper.
+        example_inputs = (
+            torch.randn(30, 20, device=self.device),
+            torch.randn(20, 30, device=self.device).t(),
+            torch.randn(40, 30, device=self.device),
+            torch.randn(30, 40, device=self.device).t(),
+        )
+        dim0 = Dim("dim0", min=1, max=100)
+        dynamic_shapes = {
+            "a": {0: dim0, 1: None},
+            "b": {0: dim0, 1: None},
+            "c": None,
+            "d": None,
+        }
+        with config.patch({"combo_kernels": True}):
+            self.check_model(Model(), example_inputs, dynamic_shapes=dynamic_shapes)
+
 
 class AOTInductorLoggingTest(LoggingTestCase):
     @make_logging_test(dynamic=logging.DEBUG)
@@ -8040,8 +8484,6 @@ def fail_gpu(suffixes: tuple[str, ...], is_skip=False):
     "test_quantized_linear": fail_gpu(("cuda", "xpu")),
     "test_quanatized_int8_linear": fail_gpu(("cuda", "xpu")),
     "test_quantized_linear_bias_none": fail_gpu(("cuda", "xpu")),
-    # No scaled_dot_product_efficient_attention implementation for XPU yet.
-    "test_scaled_dot_product_efficient_attention": fail_gpu(("xpu",)),
 }
 
 MPS_TEST_FAILURES = {
diff --git a/test/inductor/test_aot_inductor_custom_ops.py b/test/inductor/test_aot_inductor_custom_ops.py
index 6f62362120b33..e83d11fe0af38 100644
--- a/test/inductor/test_aot_inductor_custom_ops.py
+++ b/test/inductor/test_aot_inductor_custom_ops.py
@@ -109,6 +109,64 @@ def _(x):
     return torch.randn(i0)
 
 
+# Custom op to test ScalarType argument serialization
+@torch.library.custom_op("aoti_custom_ops::fn_with_dtype_arg", mutates_args={})
+def fn_with_dtype_arg(x: torch.Tensor, dtype: torch.dtype) -> torch.Tensor:
+    return torch.zeros_like(x, dtype=dtype)
+
+
+@fn_with_dtype_arg.register_fake
+def _(x, dtype):
+    return torch.zeros_like(x, dtype=dtype)
+
+
+# Custom ops to test MemoryFormat and Layout argument serialization
+# Using lower-level torch.library API since custom_op decorator doesn't support these types
+_memory_format_test_lib = torch.library.Library("aoti_custom_ops", "FRAGMENT")  # noqa: TOR901
+_memory_format_test_lib.define(
+    "fn_with_memory_format_arg(Tensor x, MemoryFormat memory_format) -> Tensor"
+)
+
+
+def _fn_with_memory_format_arg_impl(x, memory_format):
+    return x.contiguous(memory_format=memory_format)
+
+
+_memory_format_test_lib.impl(
+    "fn_with_memory_format_arg", _fn_with_memory_format_arg_impl, "CPU"
+)
+_memory_format_test_lib.impl(
+    "fn_with_memory_format_arg", _fn_with_memory_format_arg_impl, "CUDA"
+)
+_memory_format_test_lib.impl(
+    "fn_with_memory_format_arg", _fn_with_memory_format_arg_impl, "XPU"
+)
+
+
+@torch.library.register_fake("aoti_custom_ops::fn_with_memory_format_arg")
+def fn_with_memory_format_arg_abstract(x, memory_format):
+    return x.contiguous(memory_format=memory_format)
+
+
+_layout_test_lib = torch.library.Library("aoti_custom_ops", "FRAGMENT")  # noqa: TOR901
+_layout_test_lib.define("fn_with_layout_arg(Tensor x, Layout layout) -> Tensor")
+
+
+def _fn_with_layout_arg_impl(x, layout):
+    # Layout argument is received correctly if we get here without error
+    return x.clone()
+
+
+_layout_test_lib.impl("fn_with_layout_arg", _fn_with_layout_arg_impl, "CPU")
+_layout_test_lib.impl("fn_with_layout_arg", _fn_with_layout_arg_impl, "CUDA")
+_layout_test_lib.impl("fn_with_layout_arg", _fn_with_layout_arg_impl, "XPU")
+
+
+@torch.library.register_fake("aoti_custom_ops::fn_with_layout_arg")
+def fn_with_layout_arg_abstract(x, layout):
+    return x.clone()
+
+
 class AOTInductorTestsTemplate:
     def test_custom_op_add(self) -> None:
         class M(torch.nn.Module):
@@ -413,7 +471,41 @@ def __torch_dispatch__(self, func, types, args=(), kwargs=None):
         self.assertEqual(len(inps), 0)
         self.assertTrue(sentinel_seen)
 
-    @skipIfXpu
+    def test_custom_op_with_dtype_arg_various_dtypes(self) -> None:
+        class Model(torch.nn.Module):
+            def forward(self, x):
+                a = torch.ops.aoti_custom_ops.fn_with_dtype_arg(x, torch.float32)
+                b = torch.ops.aoti_custom_ops.fn_with_dtype_arg(x, torch.float64)
+                c = torch.ops.aoti_custom_ops.fn_with_dtype_arg(x, torch.int32)
+                d = torch.ops.aoti_custom_ops.fn_with_dtype_arg(x, torch.uint16)
+                e = torch.ops.aoti_custom_ops.fn_with_dtype_arg(x, torch.bfloat16)
+                return a, b, c, d, e
+
+        m = Model().to(device=self.device)
+        args = (torch.randn(2, 2, device=self.device),)
+        self.check_model(m, args)
+
+    def test_custom_op_with_memory_format_arg(self) -> None:
+        class Model(torch.nn.Module):
+            def forward(self, x):
+                return torch.ops.aoti_custom_ops.fn_with_memory_format_arg(
+                    x, torch.channels_last
+                )
+
+        m = Model().to(device=self.device)
+        args = (torch.randn(1, 3, 4, 4, device=self.device),)
+        self.check_model(m, args)
+
+    def test_custom_op_with_layout_arg(self) -> None:
+        class Model(torch.nn.Module):
+            def forward(self, x):
+                return torch.ops.aoti_custom_ops.fn_with_layout_arg(x, torch.strided)
+
+        m = Model().to(device=self.device)
+        args = (torch.randn(4, 4, device=self.device),)
+        self.check_model(m, args)
+
+    @skipIfXpu(msg="compile error - torch-xpu-ops: 2609")
     @unittest.skipIf(IS_FBCODE, "unable to find library -laoti_custom_ops")
     def test_custom_op_square(self) -> None:
         class Model(torch.nn.Module):
diff --git a/test/inductor/test_aot_inductor_package.py b/test/inductor/test_aot_inductor_package.py
index 887e8ab30ee38..6ae1dc0208480 100644
--- a/test/inductor/test_aot_inductor_package.py
+++ b/test/inductor/test_aot_inductor_package.py
@@ -308,8 +308,11 @@ def forward(self, x, y):
                 if self.device == GPU_TYPE:
                     kernel_bin = get_kernel_bin_format(self.device)
                     self.assertTrue(not list(tmp_path.glob(f"*.{kernel_bin}")))
-                    # Check if .cubin.o files exist and use unique kernel names
-                    self.assertTrue(list(tmp_path.glob(f"triton_*.{kernel_bin}.o")))
+                    # Check that cubin binaries are embedded as object files.
+                    # Either individual per-kernel .o files or a single combined .o.
+                    individual_objs = list(tmp_path.glob(f"triton_*.{kernel_bin}.o"))
+                    combined_obj = list(tmp_path.glob("cubins_combined.o"))
+                    self.assertTrue(individual_objs or combined_obj)
 
                 # Check if the .so file was build successfully
                 so_path = build_path / "libaoti_model.so"
diff --git a/test/inductor/test_aot_inductor_utils.py b/test/inductor/test_aot_inductor_utils.py
index cb16f46a752b8..32cc645ff02d8 100644
--- a/test/inductor/test_aot_inductor_utils.py
+++ b/test/inductor/test_aot_inductor_utils.py
@@ -5,7 +5,7 @@
 import shutil
 import tempfile
 import types
-from typing import Any, Optional, TYPE_CHECKING, Union
+from typing import Any, TYPE_CHECKING
 
 import torch
 import torch._export
@@ -148,10 +148,10 @@ def legacy_run(
 
     @staticmethod
     def compile(
-        model: Union[torch.nn.Module, types.FunctionType],
+        model: torch.nn.Module | types.FunctionType,
         example_inputs: tuple[torch.Tensor, ...],
-        inductor_configs: Optional[dict[str, Any]] = None,
-        dynamic_shapes: Optional[Union[dict[str, Any], tuple[Any], list[Any]]] = None,
+        inductor_configs: dict[str, Any] | None = None,
+        dynamic_shapes: dict[str, Any] | tuple[Any] | list[Any] | None = None,
     ):
         if not isinstance(model, torch.nn.Module):
             # This should really be the default behavior of torch.export.export
@@ -176,10 +176,10 @@ def compile(
 
     @staticmethod
     def run(
-        model: Union[torch.nn.Module, types.FunctionType],
+        model: torch.nn.Module | types.FunctionType,
         example_inputs: tuple[torch.Tensor, ...],
-        inductor_configs: Optional[dict[str, Any]] = None,
-        dynamic_shapes: Optional[Union[dict[str, Any], tuple[Any], list[Any]]] = None,
+        inductor_configs: dict[str, Any] | None = None,
+        dynamic_shapes: dict[str, Any] | tuple[Any] | list[Any] | None = None,
     ):
         package_path = AOTIRunnerUtil.compile(
             model,
@@ -192,10 +192,10 @@ def run(
 
     @staticmethod
     def run_multiple(
-        model: Union[torch.nn.Module, types.FunctionType],
+        model: torch.nn.Module | types.FunctionType,
         list_example_inputs: list[tuple[torch.Tensor, ...]],
-        inductor_configs: Optional[dict[str, Any]] = None,
-        dynamic_shapes: Optional[Union[dict[str, Any], tuple[Any], list[Any]]] = None,
+        inductor_configs: dict[str, Any] | None = None,
+        dynamic_shapes: dict[str, Any] | tuple[Any] | list[Any] | None = None,
     ):
         package_path = AOTIRunnerUtil.compile(
             model,
diff --git a/test/inductor/test_aoti_cross_compile_windows.py b/test/inductor/test_aoti_cross_compile_windows.py
index 15904a8d9dcab..809cf74ed3d69 100644
--- a/test/inductor/test_aoti_cross_compile_windows.py
+++ b/test/inductor/test_aoti_cross_compile_windows.py
@@ -1,15 +1,18 @@
 # Owner(s): ["module: inductor"]
 import os
 import platform
+import shutil
 import tempfile
 import unittest
 from dataclasses import dataclass
 from pathlib import Path
-from typing import Any, Optional
+from typing import Any
+from unittest.mock import MagicMock, patch
 
 import torch
 import torch._inductor.config
 from torch._environment import is_fbcode
+from torch._inductor.cpp_builder import _ensure_mingw_cudart_import_lib
 from torch._inductor.test_case import TestCase
 from torch.testing._internal.inductor_utils import GPU_TYPE, HAS_GPU, requires_gpu
 
@@ -21,8 +24,8 @@ class ModelTestConfig:
     name: str
     model_class: type
     example_inputs: tuple[torch.Tensor, ...]
-    dynamic_shapes: Optional[dict[str, Any]] = None
-    inductor_configs: Optional[dict[str, Any]] = None
+    dynamic_shapes: dict[str, Any] | None = None
+    inductor_configs: dict[str, Any] | None = None
     rtol: float = 1e-4
     atol: float = 1e-4
 
@@ -34,8 +37,8 @@ class WindowsCrossCompilationTestFramework:
     Provides reusable logic for creating compile and load test methods.
     """
 
-    _base_path: Optional[Path] = None
-    _win_torch_libs_path: Optional[str] = None
+    _base_path: Path | None = None
+    _win_torch_libs_path: str | None = None
 
     @classmethod
     def base_path(cls) -> Path:
@@ -45,12 +48,12 @@ def base_path(cls) -> Path:
         return cls._base_path
 
     @classmethod
-    def set_base_path(cls, path: Optional[Path | str] = None) -> None:
+    def set_base_path(cls, path: Path | str | None = None) -> None:
         """Set the base path for package files."""
         cls._base_path = Path(path) if path else None
 
     @classmethod
-    def set_win_torch_libs_path(cls, path: Optional[str] = None) -> None:
+    def set_win_torch_libs_path(cls, path: str | None = None) -> None:
         """Set the path for Windows torch libs."""
         cls._win_torch_libs_path = path
 
@@ -318,6 +321,137 @@ def forward(self, x):
         )
 
 
+class TestEnsureMingwCudartImportLib(TestCase):
+    """Unit tests for _ensure_mingw_cudart_import_lib."""
+
+    def setUp(self):
+        super().setUp()
+        self.tmp_dir = tempfile.mkdtemp(prefix="test_mingw_cudart_")
+        self.cuda_home = os.path.join(self.tmp_dir, "cuda")
+        self.lib_dir = os.path.join(self.tmp_dir, "lib")
+        os.makedirs(os.path.join(self.cuda_home, "bin", "x64"), exist_ok=True)
+        os.makedirs(self.lib_dir, exist_ok=True)
+
+    def tearDown(self):
+        shutil.rmtree(self.tmp_dir, ignore_errors=True)
+        super().tearDown()
+
+    def _create_fake_dll(self, version: str = "130") -> str:
+        dll_path = os.path.join(self.cuda_home, "bin", "x64", f"cudart64_{version}.dll")
+        Path(dll_path).touch()
+        return dll_path
+
+    def _create_fake_cudart_lib(self) -> str:
+        lib_path = os.path.join(self.lib_dir, "cudart.lib")
+        Path(lib_path).touch()
+        return lib_path
+
+    def test_noop_when_no_windows_cuda_home(self):
+        """Should skip gracefully when WINDOWS_CUDA_HOME is not set."""
+        env = os.environ.copy()
+        env.pop("WINDOWS_CUDA_HOME", None)
+        with patch.dict(os.environ, env, clear=True):
+            result = _ensure_mingw_cudart_import_lib([self.lib_dir])
+        self.assertEqual(result, [])
+        self.assertFalse(os.path.exists(os.path.join(self.lib_dir, "libcudart.a")))
+
+    def test_noop_when_libcudart_a_already_exists(self):
+        """Should skip generation if libcudart.a already exists."""
+        existing = os.path.join(self.lib_dir, "libcudart.a")
+        Path(existing).touch()
+        with patch.dict(os.environ, {"WINDOWS_CUDA_HOME": self.cuda_home}):
+            result = _ensure_mingw_cudart_import_lib([self.lib_dir])
+        self.assertEqual(result, [])
+        self.assertTrue(os.path.exists(existing))
+
+    @patch("torch._inductor.cpp_builder._create_msvc_gs_stubs_lib", return_value=None)
+    def test_gs_stubs_fallback_when_no_dll_found(self, mock_stubs):
+        """Should attempt GS stubs fallback when no cudart64_*.dll is found."""
+        self._create_fake_cudart_lib()
+        with patch.dict(os.environ, {"WINDOWS_CUDA_HOME": self.cuda_home}):
+            _ensure_mingw_cudart_import_lib([self.lib_dir])
+        mock_stubs.assert_called_once_with(self.lib_dir)
+
+    @patch(
+        "torch._inductor.cpp_builder._create_msvc_gs_stubs_lib",
+        return_value="msvc_gs_stubs",
+    )
+    def test_gs_stubs_returned_when_no_dll(self, mock_stubs):
+        """Should return the stubs library name when DLL is unavailable."""
+        self._create_fake_cudart_lib()
+        with patch.dict(os.environ, {"WINDOWS_CUDA_HOME": self.cuda_home}):
+            result = _ensure_mingw_cudart_import_lib([self.lib_dir])
+        self.assertEqual(result, ["msvc_gs_stubs"])
+
+    def test_noop_when_no_writable_dir_with_cudart_lib(self):
+        """Should skip gracefully when no writable directory contains cudart.lib."""
+        self._create_fake_dll()
+        with patch.dict(os.environ, {"WINDOWS_CUDA_HOME": self.cuda_home}):
+            result = _ensure_mingw_cudart_import_lib([self.lib_dir])
+        self.assertEqual(result, [])
+        self.assertFalse(os.path.exists(os.path.join(self.lib_dir, "libcudart.a")))
+
+    @patch("subprocess.run")
+    def test_successful_generation(self, mock_run: MagicMock):
+        """Should call gendef and dlltool when all conditions are met."""
+        dll_path = self._create_fake_dll()
+        self._create_fake_cudart_lib()
+
+        with patch.dict(os.environ, {"WINDOWS_CUDA_HOME": self.cuda_home}):
+            result = _ensure_mingw_cudart_import_lib([self.lib_dir])
+
+        self.assertEqual(result, [])
+        self.assertEqual(mock_run.call_count, 2)
+
+        # Verify gendef call
+        gendef_call = mock_run.call_args_list[0]
+        gendef_cmd = gendef_call[0][0]
+        self.assertEqual(gendef_cmd[0], "gendef")
+        self.assertEqual(gendef_cmd[1], "-")
+        self.assertEqual(gendef_cmd[2], dll_path)
+
+        # Verify dlltool call
+        dlltool_call = mock_run.call_args_list[1]
+        dlltool_cmd = dlltool_call[0][0]
+        self.assertEqual(dlltool_cmd[0], "x86_64-w64-mingw32-dlltool")
+        self.assertIn("-d", dlltool_cmd)
+        self.assertIn("-l", dlltool_cmd)
+        self.assertIn("-D", dlltool_cmd)
+        self.assertIn("cudart64_130.dll", dlltool_cmd)
+
+    @patch(
+        "torch._inductor.cpp_builder._create_msvc_gs_stubs_lib",
+        return_value="msvc_gs_stubs",
+    )
+    @patch("subprocess.run", side_effect=FileNotFoundError("gendef not found"))
+    def test_gs_stubs_fallback_on_gendef_not_found(self, mock_run, mock_stubs):
+        """Should fall back to GS stubs when gendef is not installed."""
+        self._create_fake_dll()
+        self._create_fake_cudart_lib()
+
+        with patch.dict(os.environ, {"WINDOWS_CUDA_HOME": self.cuda_home}):
+            result = _ensure_mingw_cudart_import_lib([self.lib_dir])
+
+        self.assertEqual(result, ["msvc_gs_stubs"])
+        mock_stubs.assert_called_once()
+
+    @patch("subprocess.run")
+    def test_bin_x64_fallback_to_bin(self, mock_run: MagicMock):
+        """Should fall back to bin/ when bin/x64/ does not exist."""
+        shutil.rmtree(os.path.join(self.cuda_home, "bin", "x64"))
+        bin_dir = os.path.join(self.cuda_home, "bin")
+        dll_path = os.path.join(bin_dir, "cudart64_130.dll")
+        Path(dll_path).touch()
+        self._create_fake_cudart_lib()
+
+        with patch.dict(os.environ, {"WINDOWS_CUDA_HOME": self.cuda_home}):
+            _ensure_mingw_cudart_import_lib([self.lib_dir])
+
+        self.assertEqual(mock_run.call_count, 2)
+        gendef_cmd = mock_run.call_args_list[0][0][0]
+        self.assertIn(dll_path, gendef_cmd)
+
+
 if __name__ == "__main__":
     import sys
 
diff --git a/test/inductor/test_augmented_graph_helper.py b/test/inductor/test_augmented_graph_helper.py
index b9406b0cf8550..6e251aec40d6d 100644
--- a/test/inductor/test_augmented_graph_helper.py
+++ b/test/inductor/test_augmented_graph_helper.py
@@ -318,7 +318,8 @@ def test_cycle_through_merge(self):
         # Merging b4 and c4 would create cycle
         tracker4.merge_to_set(b4, c4)
 
-        self.assertTrue(tracker4.has_cycle())
+        # has_cycle() is not merge-aware
+        self.assertFalse(tracker4.has_cycle())
 
     def test_cycle_with_extra_deps(self):
         """Test cycle detection with extra dependencies."""
diff --git a/test/inductor/test_auto_chunker.py b/test/inductor/test_auto_chunker.py
index c131572a5bee3..50726e5ff0bf1 100644
--- a/test/inductor/test_auto_chunker.py
+++ b/test/inductor/test_auto_chunker.py
@@ -247,6 +247,134 @@ def step(func, xs, ys):
             f"Actual peak_memory {peak_memory}, expected bound {expected_bound}",
         )
 
+    @config.patch("auto_chunker.output_size_threshold", 1024)
+    @config.patch("auto_chunker.num_chunk", 2)
+    def test_propagate_tanh_neg(self):
+        M, K, N = 256, 4, 256
+        x = torch.randn(M, K, device=GPU_TYPE, requires_grad=True)
+        w = torch.randn(K, N, device=GPU_TYPE, requires_grad=True)
+
+        def f(x, w):
+            out = (x * 2) @ w
+            out = torch.tanh(out)
+            out = -out
+            loss = out.sum()
+            loss.backward()
+            return loss
+
+        expect = (f(x, w), x.grad, w.grad)
+        x.grad = None
+        w.grad = None
+        opt_f = torch.compile(f)
+        actual = (opt_f(x, w), x.grad, w.grad)
+
+        self.assertTrue(same(expect, actual, tol=1e-3))
+        self.assertEqual(metrics.num_auto_chunking, 1)
+
+    @config.patch("auto_chunker.output_size_threshold", 1024)
+    @config.patch("auto_chunker.num_chunk", 2)
+    def test_propagate_amax_unsqueeze(self):
+        M, K, N = 256, 4, 256
+        x = torch.randn(M, K, device=GPU_TYPE, requires_grad=True)
+        w = torch.randn(K, N, device=GPU_TYPE, requires_grad=True)
+
+        def f(x, w):
+            out = (x * 2) @ w
+            max_val = out.amax(dim=-1)
+            out = out - max_val.unsqueeze(-1)
+            out = torch.exp(out)
+            loss = out.sum()
+            loss.backward()
+            return loss
+
+        expect = (f(x, w), x.grad, w.grad)
+        x.grad = None
+        w.grad = None
+        opt_f = torch.compile(f)
+        actual = (opt_f(x, w), x.grad, w.grad)
+
+        self.assertTrue(same(expect, actual, tol=1e-3))
+        self.assertEqual(metrics.num_auto_chunking, 1)
+
+    @config.patch("auto_chunker.output_size_threshold", 1024)
+    @config.patch("auto_chunker.num_chunk", 2)
+    def test_propagate_gather(self):
+        M, K, N = 256, 4, 256
+        x = torch.randn(M, K, device=GPU_TYPE, requires_grad=True)
+        w = torch.randn(K, N, device=GPU_TYPE, requires_grad=True)
+        targets = torch.randint(0, N, (M,), device=GPU_TYPE)
+
+        def f(x, w, targets):
+            out = (x * 2) @ w
+            out = out.softmax(dim=-1)
+            selected = out.gather(1, targets.unsqueeze(1))
+            loss = selected.squeeze(1).sum()
+            loss.backward()
+            return loss
+
+        expect = (f(x, w, targets), x.grad, w.grad)
+        x.grad = None
+        w.grad = None
+        opt_f = torch.compile(f)
+        actual = (opt_f(x, w, targets), x.grad, w.grad)
+
+        self.assertTrue(same(expect, actual, tol=1e-3))
+        self.assertEqual(metrics.num_auto_chunking, 1)
+
+    @config.patch("auto_chunker.output_size_threshold", 1024)
+    @config.patch("auto_chunker.num_chunk", 2)
+    def test_propagate_scatter(self):
+        M, K, N = 256, 4, 256
+        x = torch.randn(M, K, device=GPU_TYPE, requires_grad=True)
+        w = torch.randn(K, N, device=GPU_TYPE, requires_grad=True)
+        targets = torch.randint(0, N, (M,), device=GPU_TYPE)
+
+        def f(x, w, targets):
+            out = (x * 2) @ w
+            out = out.softmax(dim=-1)
+            out = out.scatter(1, targets.unsqueeze(1), 0.0)
+            loss = out.sum()
+            loss.backward()
+            return loss
+
+        expect = (f(x, w, targets), x.grad, w.grad)
+        x.grad = None
+        w.grad = None
+        opt_f = torch.compile(f)
+        actual = (opt_f(x, w, targets), x.grad, w.grad)
+
+        self.assertTrue(same(expect, actual, tol=1e-3))
+        self.assertEqual(metrics.num_auto_chunking, 1)
+
+    @config.patch("auto_chunker.output_size_threshold", 1024)
+    @config.patch("auto_chunker.num_chunk", 2)
+    def test_propagate_manual_cross_entropy(self):
+        M, K, N = 256, 4, 256
+        x = torch.randn(M, K, device=GPU_TYPE, requires_grad=True)
+        w = torch.randn(K, N, device=GPU_TYPE, requires_grad=True)
+        targets = torch.randint(0, N, (M,), device=GPU_TYPE)
+
+        def f(x, w, targets):
+            logits = (x * 2) @ w
+            max_logits = logits.amax(dim=-1)
+            shifted = logits - max_logits.unsqueeze(-1)
+            exp_shifted = shifted.exp()
+            sum_exp = exp_shifted.sum(dim=-1)
+            log_probs = shifted - sum_exp.log().unsqueeze(-1)
+            target_log_probs = log_probs.gather(1, targets.unsqueeze(1))
+            loss = -target_log_probs.squeeze(1).sum() / M
+            loss.backward()
+            return loss
+
+        expect = (f(x, w, targets), x.grad, w.grad)
+        x.grad = None
+        w.grad = None
+        opt_f = torch.compile(f)
+        actual = (opt_f(x, w, targets), x.grad, w.grad)
+
+        self.assertTrue(same(expect, actual, tol=1e-3))
+        self.assertEqual(metrics.num_auto_chunking, 1)
+
     def test_set_num_chunk_with_compile_options(self):
         B = 32
         T = 1024
diff --git a/test/inductor/test_auto_functionalize.py b/test/inductor/test_auto_functionalize.py
index 95ac2d88b32a7..897e0c64a64a2 100644
--- a/test/inductor/test_auto_functionalize.py
+++ b/test/inductor/test_auto_functionalize.py
@@ -281,7 +281,8 @@ def foo_impl(x):
                 x = torch.randn(3)
                 expected = x.sin()
                 torch.ops.mylib.foo(x)
-                assert torch.allclose(x, expected)
+                if not torch.allclose(x, expected):
+                    raise AssertionError
 
                 @torch.compile(backend="aot_eager_decomp_partition", fullgraph=True)
                 def f(x):
@@ -476,6 +477,35 @@ def run_inductor(
 
         return [compiled_args, result, graph]
 
+    def test_cumulative_out_preserves_out_dtype_under_compile(self):
+        test_cases = [
+            (torch.cumsum, torch.tensor([0.9201, 0.1166], dtype=torch.float32)),
+            (torch.cumprod, torch.tensor([2.0, 0.6], dtype=torch.float32)),
+        ]
+
+        for op, x in test_cases:
+            with self.subTest(op=op.__name__):
+
+                def f(out, x):
+                    return op(x, 0, out=out)
+
+                eager_out = torch.tensor([0, 0], dtype=torch.int32)
+                aot_eager_out = eager_out.clone()
+                inductor_out = eager_out.clone()
+
+                eager_ret = f(eager_out, x)
+                aot_eager_ret = torch.compile(f, backend="aot_eager", fullgraph=True)(
+                    aot_eager_out, x
+                )
+                inductor_ret = torch.compile(f, backend="inductor", fullgraph=True)(
+                    inductor_out, x
+                )
+
+                self.assertEqual(aot_eager_out, eager_out)
+                self.assertEqual(inductor_out, eager_out)
+                self.assertEqual(aot_eager_ret, eager_ret)
+                self.assertEqual(inductor_ret, eager_ret)
+
     @torch._inductor.config.patch(enable_auto_functionalized_v2=True)
     def test_auto_functionalize_with_returns_v2(self):
         with torch.library._scoped_library("mylib", "FRAGMENT") as lib:
diff --git a/test/inductor/test_autoheuristic.py b/test/inductor/test_autoheuristic.py
index 270608086799b..0897662088ec3 100644
--- a/test/inductor/test_autoheuristic.py
+++ b/test/inductor/test_autoheuristic.py
@@ -1,6 +1,7 @@
 # Owner(s): ["module: inductor"]
 import os
 import unittest
+from unittest.mock import patch
 
 import torch
 import torch._inductor.config as inductor_config
@@ -10,11 +11,10 @@
 from torch._inductor.runtime.runtime_utils import cache_dir
 from torch._inductor.test_case import run_tests, TestCase
 from torch._inductor.utils import get_gpu_shared_memory
-from torch.testing._internal.common_utils import skipIfXpu
+from torch.testing._internal.common_utils import skipIfRocm, skipIfXpu
 from torch.testing._internal.inductor_utils import GPU_TYPE, HAS_GPU, IS_A100, IS_H100
 
 
-@skipIfXpu(msg="AutoHeuristic doesn't currently work on the XPU stack")
 class AutoHeuristicTest(TestCase):
     def count_lines_in_file(self, file_path):
         with open(file_path) as file:
@@ -35,12 +35,14 @@ def get_path_to_autoheuristic_log(self, name):
         path = cache_dir() + "/autoheuristic/" + device_name + "/" + name + ".txt"
         return path
 
+    @skipIfXpu(msg="AutoHeuristic doesn't currently work on the XPU stack")
     def test_autoheuristic_pad_mm_default(self):
         # this test ensures that data is not collected for pad_mm when autoheuristic config is set to its default value
         self.run_mm()
         self.assertFalse(os.path.exists(self.get_path_to_autoheuristic_log("pad_mm")))
 
-    @inductor_config.patch(autoheuristic_collect="foo")
+    @skipIfXpu(msg="AutoHeuristic doesn't currently work on the XPU stack")
+    @inductor_config.patch("autoheuristic_collect.pad_mm", False)
     def test_autoheuristic_pad_mm_off(self):
         # this test ensures that data is not collected for pad_mm when autoheuristic_collect does not contain "pad_mm"
         self.run_mm()
@@ -56,17 +58,21 @@ def assert_autoheuristic_collected_data(self):
         # 1 line for metadata, 1 line for header, 1 line per choice (orig, padded)
         self.assertEqual(num_lines, 4)
 
-    @inductor_config.patch(autoheuristic_collect="pad_mm")
+    @skipIfXpu(msg="AutoHeuristic doesn't currently work on the XPU stack")
+    @inductor_config.patch("autoheuristic_collect.pad_mm", True)
     def test_autoheuristic_pad_mm_collect_data(self):
-        # this test ensures that data is collected for pad_mm when autoheuristic_collect="pad_mm"
+        # this test ensures that data is collected for pad_mm when autoheuristic_collect.pad_mm=True
         self.assert_autoheuristic_collected_data()
 
-    @inductor_config.patch(autoheuristic_collect="foo,pad_mm")
+    @skipIfXpu(msg="AutoHeuristic doesn't currently work on the XPU stack")
+    @inductor_config.patch("autoheuristic_collect.pad_mm", True)
     def test_autoheuristic_pad_mm_collect_data2(self):
-        # this test ensures that data is collected for "pad_mm" when autoheuristic_collect contains "pad_mm"
+        # this test ensures that data is collected for pad_mm when autoheuristic_collect.pad_mm=True
         self.assert_autoheuristic_collected_data()
 
-    @inductor_config.patch(autoheuristic_collect="test")
+    @skipIfRocm(msg="AutoHeuristic doesn't currently work on the ROCm stack")
+    @skipIfXpu(msg="AutoHeuristic doesn't currently work on the XPU stack")
+    @patch.dict(os.environ, {"TORCHINDUCTOR_AUTOHEURISTIC_COLLECT": "test"})
     def test_autoheuristic(self):
         # test basic functionality of autoheuristic
         def fallback():
@@ -102,7 +108,9 @@ def feedback_fn(choice):
         self.assertEqual(num_lines, 5)
 
         shared_memory = get_gpu_shared_memory()
-        (fst, snd) = get_interface_for_device(GPU_TYPE).get_device_capability()
+        compute_cap = get_interface_for_device(GPU_TYPE).get_compute_capability()
+        # Convert single int compute capability (e.g., 90) to tuple (e.g., (9, 0))
+        fst, snd = compute_cap // 10, compute_cap % 10
 
         with open(path) as file:
             lines = file.readlines()
@@ -116,15 +124,17 @@ def feedback_fn(choice):
             self.assertEqual("5,b,2", lines[3].rstrip())
             self.assertEqual("5,c,3", lines[4].rstrip())
 
+    @skipIfXpu(msg="AutoHeuristic doesn't currently work on the XPU stack")
     @unittest.skipIf(not IS_A100, "heuristic only run on A100")
-    @inductor_config.patch(autoheuristic_use="pad_mm")
+    @inductor_config.patch("autoheuristic_use.pad_mm", True)
     def test_autoheuristic_a100(self):
         # Make sure heuristic does not break anything
         # TODO (AlnisM): Find a way to check whether heuristic is used
         self.run_mm()
 
+    @skipIfXpu(msg="AutoHeuristic doesn't currently work on the XPU stack")
     @unittest.skipIf(not IS_H100, "heuristic only run on H100")
-    @inductor_config.patch(autoheuristic_use="pad_mm")
+    @inductor_config.patch("autoheuristic_use.pad_mm", True)
     def test_autoheuristic_h100(self):
         # Make sure heuristic does not break anything
         # TODO (AlnisM): Find a way to check whether heuristic is used
@@ -145,12 +155,14 @@ def fn(a, b):
     # a choice made by the heuristic might be added to the list of choices
     # and if select_algorithm now creates a new precompile key, it will be
     # different from the precompile key created by autoheuristic
-    @inductor_config.patch(
-        autoheuristic_collect="mixed_mm",
-        autoheuristic_use="",
-        fx_graph_cache=False,
-        fx_graph_remote_cache=False,
+    @skipIfXpu(msg="AutoHeuristic doesn't currently work on the XPU stack")
+    @unittest.skip(
+        "mixed_mm autoheuristic collection is broken after mixed_mm special casing deletion (PR #147151)"
     )
+    @inductor_config.patch("autoheuristic_collect.mixed_mm", True)
+    @inductor_config.patch("autoheuristic_use.mixed_mm", False)
+    @inductor_config.patch(fx_graph_cache=False)
+    @inductor_config.patch(fx_graph_remote_cache=False)
     def test_global_feedback(self):
         self.run_mixed_mm()
         path = self.get_path_to_autoheuristic_log("mixed_mm")
@@ -161,12 +173,43 @@ def test_global_feedback(self):
         # 1 line for fallback + at least 1 config
         self.assertTrue(num_lines > 4)
 
-    @inductor_config.patch(autoheuristic_use="mixed_mm")
+    @skipIfXpu(msg="AutoHeuristic doesn't currently work on the XPU stack")
+    @inductor_config.patch("autoheuristic_use.mixed_mm", True)
     @unittest.skipIf(not IS_A100, "heuristic only run on A100")
     def test_mixed_mm_a100(self):
         self.run_mixed_mm()
         # TODO (AlnisM): Find a way to check whether heuristic is used
 
+    @skipIfXpu(msg="AutoHeuristic doesn't currently work on the XPU stack")
+    @unittest.skipIf(not IS_H100 and not IS_A100, "heuristic only run on H100")
+    @inductor_config.patch(deterministic=True)
+    @inductor_config.patch("autoheuristic_use.pad_mm", True)
+    def test_pad_mm_autoheuristic_deterministic_mode(self):
+        """Test that pad_mm AutoHeuristics works in deterministic mode."""
+        from torch._dynamo.utils import counters
+
+        counters.clear()
+
+        def f(a, b):
+            return torch.mm(a, b)
+
+        cf = torch.compile(f)
+        # Use shapes that would normally trigger padding but aren't well-aligned
+        a = torch.randn(2047, 2048, device=GPU_TYPE, dtype=torch.float16)
+        b = torch.randn(2048, 2049, device=GPU_TYPE, dtype=torch.float16)
+
+        # Run the compiled function
+        result = cf(a, b)
+
+        # Verify correctness with tolerance for potential padding differences
+        expected = torch.mm(a, b)
+        torch.testing.assert_close(result, expected, atol=0.1, rtol=0.05)
+
+        # In deterministic mode with AutoHeuristics enabled,
+        # we should not do any benchmarking (pad_mm_bench should stay 0)
+        # because AutoHeuristics makes the decision without benchmarking
+        self.assertEqual(counters["inductor"]["pad_mm_bench"], 0)
+
 
 if __name__ == "__main__":
     if HAS_GPU:
diff --git a/test/inductor/test_benchmark_fusion.py b/test/inductor/test_benchmark_fusion.py
index 499af078f1762..fb7fd68807169 100644
--- a/test/inductor/test_benchmark_fusion.py
+++ b/test/inductor/test_benchmark_fusion.py
@@ -195,7 +195,8 @@ def f(x):
             return y + 1
 
         x = torch.randn(1024, 1024, device=self.device)
-        self.common(f, (x,))
+        # Disable lowp check due to non-deterministic kernel fusion strategies affecting fp16 precision.
+        self.common(f, (x,), check_lowp=False)
 
 
 if HAS_GPU_AND_TRITON:
diff --git a/test/inductor/test_benchmarking.py b/test/inductor/test_benchmarking.py
index 7fac9fb705bf3..9732cb7b5043b 100644
--- a/test/inductor/test_benchmarking.py
+++ b/test/inductor/test_benchmarking.py
@@ -150,6 +150,77 @@ def capture_benchmark_gpu(self, _callable, **kwargs):
         self.assertEqual(captured_kwargs["warmup"], custom_warmup)
         self.assertEqual(captured_kwargs["rep"], custom_rep)
 
+    @unittest.skipIf(not HAS_CPU, "requires CPU")
+    @parametrize("benchmarker_cls", ALL_BENCHMARKER_CLASSES)
+    def test_benchmarker_cpu_override_dispatch(self, benchmarker_cls, device="cpu"):
+        # Registers a custom handler for 'cpu' and verifies dispatch uses it instead of the default path.
+        from torch._inductor.runtime import benchmarking as _bench
+
+        benchmarker = benchmarker_cls()
+
+        # Snapshot registry and restore at the end to avoid cross-test pollution.
+        orig = dict(_bench._BENCHMARK_DISPATCH)
+        try:
+            seen = {"cpu_override": 0}
+
+            def custom_cpu(self, fn, *, warmup, rep, **kw):
+                seen["cpu_override"] += 1
+                return "cpu-override"
+
+            # Override the built-in 'cpu' registration
+            _bench.register_benchmarker("cpu", custom_cpu, override=True)
+
+            # Ensure default CPU/GPU methods are NOT called if registry override works.
+            with (
+                patch.object(
+                    benchmarker_cls,
+                    "benchmark_cpu",
+                    side_effect=AssertionError(
+                        "benchmark_cpu should not be called when a custom 'cpu' handler is registered"
+                    ),
+                    create=True,
+                ),
+                patch.object(
+                    benchmarker_cls,
+                    "benchmark_gpu",
+                    side_effect=AssertionError(
+                        "benchmark_gpu should not be called for 'cpu' device"
+                    ),
+                    create=True,
+                ),
+            ):
+                (fn, fn_args, fn_kwargs), _ = self.make_params(device)
+                out = benchmarker.benchmark(fn, fn_args, fn_kwargs)
+                self.assertEqual(out, "cpu-override")
+                self.assertEqual(seen["cpu_override"], 1)
+        finally:
+            _bench._BENCHMARK_DISPATCH.clear()
+            _bench._BENCHMARK_DISPATCH.update(orig)
+
+    @unittest.skipIf(not HAS_CPU, "requires CPU")
+    @parametrize("benchmarker_cls", ALL_BENCHMARKER_CLASSES)
+    def test_benchmarker_cpu_override_runs_callable(
+        self, benchmarker_cls, device="cpu"
+    ):
+        from torch._inductor.runtime import benchmarking as _bench
+
+        benchmarker = benchmarker_cls()
+        orig = dict(_bench._BENCHMARK_DISPATCH)
+        try:
+            # Override CPU but still route to benchmark_cpu internally
+            def custom_cpu(self, f, *, warmup, rep, **kw):
+                # Just delegate to the original path; we want to ensure `f()` calls the user's fn.
+                return self.benchmark_cpu(f, warmup=warmup, rep=rep, **kw)
+
+            _bench.register_benchmarker("cpu", custom_cpu, override=True)
+            # Define a simple op and ensure it actually runs without TypeError
+            (fn, fn_args, fn_kwargs), _ = self.make_params(device)
+            out = benchmarker.benchmark(fn, fn_args, fn_kwargs, warmup=1, rep=1)
+            self.assertGreater(out, 0)
+        finally:
+            _bench._BENCHMARK_DISPATCH.clear()
+            _bench._BENCHMARK_DISPATCH.update(orig)
+
 
 if __name__ == "__main__":
     run_tests()
diff --git a/test/inductor/test_block_analysis.py b/test/inductor/test_block_analysis.py
index 3d2cb0373c430..83ec5cf20ae3c 100644
--- a/test/inductor/test_block_analysis.py
+++ b/test/inductor/test_block_analysis.py
@@ -118,13 +118,21 @@ def test_index_with_dynamic_shapes(self):
             )
             sizevars = V.graph.sizevars
             for expected, actual in zip((dims, strides, block_index_exprs), match):
-                assert isinstance(expected, (list, tuple)) and isinstance(
-                    actual, (list, tuple)
-                )
-                for expected_expr, actual_expr in zip(expected, actual):
-                    assert isinstance(expected_expr, sympy.Expr) and isinstance(
-                        actual_expr, sympy.Expr
+                if not (
+                    isinstance(expected, (list, tuple))
+                    and isinstance(actual, (list, tuple))
+                ):
+                    raise AssertionError(
+                        f"Expected list/tuple types, got {type(expected)} and {type(actual)}"
                     )
+                for expected_expr, actual_expr in zip(expected, actual):
+                    if not (
+                        isinstance(expected_expr, sympy.Expr)
+                        and isinstance(actual_expr, sympy.Expr)
+                    ):
+                        raise AssertionError(
+                            f"Expected sympy.Expr types, got {type(expected_expr)} and {type(actual_expr)}"
+                        )
                     self.assertTrue(
                         sizevars.statically_known_equals(
                             sizevars.remove_precomputed_replacements(expected_expr),
diff --git a/test/inductor/test_block_ptr_store_dtype.py b/test/inductor/test_block_ptr_store_dtype.py
new file mode 100644
index 0000000000000..e37446379e747
--- /dev/null
+++ b/test/inductor/test_block_ptr_store_dtype.py
@@ -0,0 +1,108 @@
+# Owner(s): ["module: inductor"]
+"""
+Unit test for block_ptr store dtype resolution with inplace buffers.
+
+Validates that codegen_block_ptr_store_line casts store values to match the
+block pointer element type (the actual tensor dtype), not the graph's
+intermediate buffer dtype, when storing into an inplace-mutated buffer.
+
+See T260349710 for the original issue: Inductor codegen for _to_copy ops
+generated .to(tl.float32) for block_ptr stores into bf16 gradient buffers,
+causing a dtype mismatch assertion in the Triton compiler.
+
+Usage:
+    buck test fbcode//caffe2/test/inductor:test_block_ptr_store_dtype
+"""
+
+import torch
+from torch._inductor.codegen.common import InplacedBuffer, REMOVED
+from torch._inductor.codegen.triton import triton_store_type
+from torch._inductor.test_case import run_tests, TestCase
+
+
+class TestBlockPtrStoreDtype(TestCase):
+    """Test dtype resolution for block_ptr stores with inplace buffers.
+
+    The logic under test (from codegen_block_ptr_store_line):
+
+        store_dtype = V.graph.get_dtype(name)
+        if name in self.args.inplace_buffers:
+            buf = self.args.inplace_buffers[name]
+            if not isinstance(buf, RemovedArg):
+                store_dtype = V.graph.get_dtype(buf.other_names[0])
+        value = f"{value}.to({triton_store_type(store_dtype)})"
+    """
+
+    @staticmethod
+    def _resolve_store_dtype(name, inplace_buffers, get_dtype):
+        """Reproduce the dtype resolution logic from codegen_block_ptr_store_line."""
+        from torch._inductor.codegen.common import RemovedArg
+
+        store_dtype = get_dtype(name)
+        if name in inplace_buffers:
+            buf = inplace_buffers[name]
+            if not isinstance(buf, RemovedArg):
+                store_dtype = get_dtype(buf.other_names[0])
+        return store_dtype
+
+    def test_non_inplace_uses_graph_dtype(self):
+        """Non-inplace buffer: store dtype matches graph dtype."""
+        dtypes = {"buf0": torch.float32}
+        result = self._resolve_store_dtype("buf0", {}, lambda n: dtypes[n])
+        self.assertEqual(result, torch.float32)
+        self.assertEqual(triton_store_type(result), "tl.float32")
+
+    def test_inplace_buffer_uses_input_dtype(self):
+        """Inplace buffer with dtype mismatch: store dtype matches the input
+        buffer's dtype (the actual tensor), not the graph output dtype.
+
+        This is the key scenario from T260349710: _to_copy produces fp32
+        intermediate values stored into a bf16 gradient buffer via inplace
+        mutation. The block pointer element type comes from the actual tensor
+        (bf16), so the cast must be .to(tl.bfloat16), not .to(tl.float32).
+        """
+        dtypes = {"buf0": torch.float32, "primals_1": torch.bfloat16}
+        inplace_buffers = {
+            "buf0": InplacedBuffer("in_out_ptr0", ["primals_1", "buf0"]),
+            "primals_1": InplacedBuffer("in_out_ptr0", ["primals_1", "buf0"]),
+        }
+        result = self._resolve_store_dtype("buf0", inplace_buffers, lambda n: dtypes[n])
+        self.assertEqual(result, torch.bfloat16)
+        self.assertEqual(triton_store_type(result), "tl.bfloat16")
+
+    def test_removed_inplace_falls_back_to_graph_dtype(self):
+        """Removed inplace buffer: falls back to graph dtype."""
+        dtypes = {"buf0": torch.float32}
+        inplace_buffers = {"buf0": REMOVED}
+        result = self._resolve_store_dtype("buf0", inplace_buffers, lambda n: dtypes[n])
+        self.assertEqual(result, torch.float32)
+        self.assertEqual(triton_store_type(result), "tl.float32")
+
+    def test_same_dtype_inplace_is_unchanged(self):
+        """Same-dtype inplace buffer: dtype is unchanged (both bf16)."""
+        dtypes = {"buf0": torch.bfloat16, "primals_1": torch.bfloat16}
+        inplace_buffers = {
+            "buf0": InplacedBuffer("in_out_ptr0", ["primals_1", "buf0"]),
+        }
+        result = self._resolve_store_dtype("buf0", inplace_buffers, lambda n: dtypes[n])
+        self.assertEqual(result, torch.bfloat16)
+        self.assertEqual(triton_store_type(result), "tl.bfloat16")
+
+    def test_chained_inplace_uses_original_input_dtype(self):
+        """Chained inplace mutations: uses the original input buffer's dtype."""
+        dtypes = {
+            "buf0": torch.float32,
+            "buf1": torch.float16,
+            "primals_1": torch.bfloat16,
+        }
+        # Chain: primals_1 -> buf1 -> buf0
+        inplace_buffers = {
+            "buf0": InplacedBuffer("in_out_ptr0", ["primals_1", "buf1", "buf0"]),
+        }
+        result = self._resolve_store_dtype("buf0", inplace_buffers, lambda n: dtypes[n])
+        self.assertEqual(result, torch.bfloat16)
+        self.assertEqual(triton_store_type(result), "tl.bfloat16")
+
+
+if __name__ == "__main__":
+    run_tests()
diff --git a/test/inductor/test_cache.py b/test/inductor/test_cache.py
index d7ac4df3bf078..aac5e1c1e4835 100644
--- a/test/inductor/test_cache.py
+++ b/test/inductor/test_cache.py
@@ -1,6 +1,7 @@
 # Owner(s): ["module: inductor"]
 from __future__ import annotations
 
+import json
 import pickle
 from concurrent.futures import ThreadPoolExecutor
 from inspect import isclass
@@ -12,6 +13,7 @@
 from typing_extensions import Self
 from unittest.mock import patch
 
+import torch._inductor.config as inductor_config
 from torch._inductor import cache as icache
 from torch._inductor.test_case import run_tests, TestCase
 from torch.testing._internal.common_utils import (
@@ -75,7 +77,8 @@ def cache_type_supports_key_and_value_types(
         key_type: type[icache.Key],
         value_type: type[icache.Value],
     ) -> bool:
-        assert len(cache_type.__orig_bases__) == 1
+        if len(cache_type.__orig_bases__) != 1:
+            raise AssertionError
         generic_base = cache_type.__orig_bases__[0]
         _key_type, _value_type = generic_base.__args__
         if ((_key_type != icache.Key) and (_key_type != key_type)) or (
@@ -814,5 +817,70 @@ def test_on_disk_cache_version_bump(
         self.assertEqual(cache.get(key), value)
 
 
+class ConfigSerializationTest(TestCase):
+    def test_callable_config_not_json_serializable_1(self):
+        # Repro: setting callable configs to a non-None value
+        # save_config_portable() return a dict that is not JSON-serializable.
+        with inductor_config.patch(
+            bucket_all_gathers_fx_bucket_size_determinator=lambda: None
+        ):
+            portable = inductor_config.save_config_portable(
+                ignore_private_configs=False
+            )
+            self.assertIn("bucket_all_gathers_fx_bucket_size_determinator", portable)
+            with self.assertRaises(TypeError, msg="not JSON serializable"):
+                json.dumps(portable)
+
+    def test_callable_config_not_json_serializable_2(self):
+        # save_config_portable calls the factory and then .uuid(),
+        # producing a JSON-serializable value for the cache key.
+        from torch._inductor.choices import InductorChoices
+
+        class ChoicesA(InductorChoices):
+            def uuid(self):
+                return "choices_a"
+
+        class ChoicesB(InductorChoices):
+            def uuid(self):
+                return "choices_b"
+
+        # None default stays as None in the config hash.
+        with inductor_config.patch(inductor_choices_class=None):
+            portable = inductor_config.save_config_portable(
+                ignore_private_configs=False
+            )
+            self.assertIsNone(portable["inductor_choices_class"])
+            json.dumps(portable)
+
+        # Factory returning ChoicesA → uuid string
+        with inductor_config.patch(inductor_choices_class=ChoicesA):
+            portable = inductor_config.save_config_portable(
+                ignore_private_configs=False
+            )
+            self.assertEqual(portable["inductor_choices_class"], "choices_a")
+            json_a = json.dumps(portable)
+
+        # Factory returning ChoicesB → different uuid string
+        with inductor_config.patch(inductor_choices_class=ChoicesB):
+            portable = inductor_config.save_config_portable(
+                ignore_private_configs=False
+            )
+            self.assertEqual(portable["inductor_choices_class"], "choices_b")
+            json_b = json.dumps(portable)
+
+        self.assertNotEqual(json_a, json_b)
+
+    def test_callable_config_without_uuid(self):
+        from torch._inductor.choices import InductorChoices
+
+        # A subclass without uuid() raises RuntimeError.
+        class PlainChoices(InductorChoices):
+            pass
+
+        with inductor_config.patch(inductor_choices_class=PlainChoices):
+            with self.assertRaises(RuntimeError):
+                inductor_config.save_config_portable(ignore_private_configs=False)
+
+
 if __name__ == "__main__":
     run_tests()
diff --git a/test/inductor/test_caching.py b/test/inductor/test_caching.py
index 1b8fc4a13bc45..0c1b3089fd782 100644
--- a/test/inductor/test_caching.py
+++ b/test/inductor/test_caching.py
@@ -11,7 +11,7 @@
 from random import Random
 from shutil import rmtree
 from threading import Event, Lock
-from typing import Any, TYPE_CHECKING, Union
+from typing import Any, TYPE_CHECKING
 from typing_extensions import TypeVar
 from unittest.mock import patch
 
@@ -635,13 +635,13 @@ def executor(self) -> Generator[ThreadPoolExecutor, None, None]:
         finally:
             executor.shutdown()
 
-    def is_lock(self, lock_or_flock: Union[Lock, FileLock]) -> bool:
+    def is_lock(self, lock_or_flock: Lock | FileLock) -> bool:
         return hasattr(lock_or_flock, "locked")
 
-    def is_flock(self, lock_or_flock: Union[Lock, FileLock]) -> bool:
+    def is_flock(self, lock_or_flock: Lock | FileLock) -> bool:
         return hasattr(lock_or_flock, "is_locked")
 
-    def lock_or_flock_locked(self, lock_or_flock: Union[Lock, FileLock]) -> bool:
+    def lock_or_flock_locked(self, lock_or_flock: Lock | FileLock) -> bool:
         if self.is_lock(lock_or_flock):
             return lock_or_flock.locked()
         elif self.is_flock(lock_or_flock):
@@ -691,7 +691,7 @@ def test_acquire_with_timeout(
         - Different lock types (Lock vs FileLock) behave consistently with their respective APIs
         """
 
-        def inner(lock_or_flock: Union[Lock, FileLock], timeout: int) -> None:
+        def inner(lock_or_flock: Lock | FileLock, timeout: int) -> None:
             if self.is_lock(lock_or_flock):
                 lock: Lock = lock_or_flock
                 if acquisition_mode == "safe":
@@ -718,12 +718,13 @@ def inner(lock_or_flock: Union[Lock, FileLock], timeout: int) -> None:
                 raise NotImplementedError
             self.assertFalse(self.lock_or_flock_locked(lock_or_flock))
 
-        assert lock_typename in ["Lock", "FileLock"]
+        if lock_typename not in ["Lock", "FileLock"]:
+            raise AssertionError(f"Unexpected lock_typename: {lock_typename}")
         flock_fpath: Path = (
             impls._OnDiskCacheImpl()._cache_dir
             / f"testing-locks-instance-{self.random_string}.lock"
         )
-        lock_or_flock: Union[Lock, FileLock] = (
+        lock_or_flock: Lock | FileLock = (
             Lock() if lock_typename == "Lock" else FileLock(str(flock_fpath))
         )
         lock_exception_type: type = (
@@ -741,7 +742,12 @@ def inner(lock_or_flock: Union[Lock, FileLock], timeout: int) -> None:
             raise NotImplementedError
 
         with self.executor() as executor:
-            assert lock_timeout in ["BLOCKING", "NON_BLOCKING", "BLOCKING_WITH_TIMEOUT"]
+            if lock_timeout not in [
+                "BLOCKING",
+                "NON_BLOCKING",
+                "BLOCKING_WITH_TIMEOUT",
+            ]:
+                raise AssertionError(f"Unexpected lock_timeout: {lock_timeout}")
             lock_or_flock_future: Future[None] = executor.submit(
                 inner,
                 lock_or_flock,
@@ -2171,9 +2177,10 @@ def _create_mock_match(self) -> Any:
 
         return mock_match
 
+    @patch("torch._prims_common.is_contiguous_or_false", return_value=True)
     @patch_on_disk_cache_base_dir
     @set_caching_module_enabled(True)
-    def test_should_pad_memoizer_caches_result(self) -> None:
+    def test_should_pad_memoizer_caches_result(self, mock_is_contiguous) -> None:
         """Test that the should_pad_memoizer caches function results.
 
         Verifies that when a function decorated with should_pad_memoizer.memoize
@@ -2215,10 +2222,11 @@ def mock_should_pad(
         self.assertTrue(result1)
         self.assertTrue(result2)
 
+    @patch("torch._prims_common.is_contiguous_or_false", return_value=True)
     @patch_on_disk_cache_base_dir
     @set_caching_module_enabled(True)
     def test_should_pad_memoizer_different_shapes_different_cache_entries(
-        self,
+        self, mock_is_contiguous
     ) -> None:
         """Test that different tensor shapes result in different cache entries.
 
@@ -2266,10 +2274,11 @@ def mock_should_pad(
         self.assertFalse(result_small)  # 8 <= 10
         self.assertTrue(result_large)  # 12 > 10
 
+    @patch("torch._prims_common.is_contiguous_or_false", return_value=True)
     @patch_on_disk_cache_base_dir
     @set_caching_module_enabled(True)
     def test_should_pad_memoizer_different_dtypes_different_cache_entries(
-        self,
+        self, mock_is_contiguous
     ) -> None:
         """Test that different tensor dtypes result in different cache entries.
 
@@ -2315,10 +2324,11 @@ def mock_should_pad(
         self.assertTrue(result_fp32)
         self.assertFalse(result_fp16)
 
+    @patch("torch._prims_common.is_contiguous_or_false", return_value=True)
     @patch_on_disk_cache_base_dir
     @set_caching_module_enabled(True)
     def test_should_pad_memoizer_different_ops_different_cache_entries(
-        self,
+        self, mock_is_contiguous
     ) -> None:
         """Test that different operations result in different cache entries.
 
@@ -2363,9 +2373,12 @@ def mock_should_pad(
         self.assertTrue(result_mm)
         self.assertFalse(result_addmm)
 
+    @patch("torch._prims_common.is_contiguous_or_false", return_value=True)
     @patch_on_disk_cache_base_dir
     @set_caching_module_enabled(True)
-    def test_should_pad_memoizer_replays_from_disk_cache(self) -> None:
+    def test_should_pad_memoizer_replays_from_disk_cache(
+        self, mock_is_contiguous
+    ) -> None:
         """Test that the memoizer replays results from disk cache after memory clear.
 
         Verifies that PersistentMemoizer correctly stores results to disk and
@@ -2408,9 +2421,12 @@ def mock_should_pad(
         self.assertEqual(call_count, 1)  # Function should NOT be called again
         self.assertTrue(result2)
 
+    @patch("torch._prims_common.is_contiguous_or_false", return_value=True)
     @patch_on_disk_cache_base_dir
     @set_caching_module_enabled(False)
-    def test_should_pad_memoizer_disabled_does_not_cache(self) -> None:
+    def test_should_pad_memoizer_disabled_does_not_cache(
+        self, mock_is_contiguous
+    ) -> None:
         """Test that the memoizer does not cache when caching is disabled.
 
         Verifies that when IS_CACHING_MODULE_ENABLED is False, the function
@@ -2449,9 +2465,10 @@ def mock_should_pad(
         self.assertTrue(result1)
         self.assertTrue(result2)
 
+    @patch("torch._prims_common.is_contiguous_or_false", return_value=True)
     @patch_on_disk_cache_base_dir
     @set_caching_module_enabled(True)
-    def test_should_pad_memoizer_with_input_tensor(self) -> None:
+    def test_should_pad_memoizer_with_input_tensor(self, mock_is_contiguous) -> None:
         """Test that the memoizer correctly handles the optional input tensor.
 
         Verifies that different input tensors (for addmm) result in different
@@ -2495,9 +2512,12 @@ def mock_should_pad(
         self.assertTrue(result_with_input)
         self.assertFalse(result_without_input)
 
+    @patch("torch._prims_common.is_contiguous_or_false", return_value=True)
     @patch_on_disk_cache_base_dir
     @set_caching_module_enabled(True)
-    def test_should_pad_params_encoder_produces_consistent_keys(self) -> None:
+    def test_should_pad_params_encoder_produces_consistent_keys(
+        self, mock_is_contiguous
+    ) -> None:
         """Test that the encoder produces consistent keys for the same inputs.
 
         Verifies that calling the encoder with the same tensor metadata produces
@@ -2525,9 +2545,12 @@ def test_should_pad_params_encoder_produces_consistent_keys(self) -> None:
         self.assertEqual(encoded1["mat1"]["shape"], tuple(mat1.shape))
         self.assertEqual(encoded1["mat2"]["shape"], tuple(mat2.shape))
 
+    @patch("torch._prims_common.is_contiguous_or_false", return_value=True)
     @patch_on_disk_cache_base_dir
     @set_caching_module_enabled(True)
-    def test_should_pad_memoizer_same_shape_different_data_uses_cache(self) -> None:
+    def test_should_pad_memoizer_same_shape_different_data_uses_cache(
+        self, mock_is_contiguous
+    ) -> None:
         """Test that tensors with the same metadata but different data share cache.
 
         Verifies that the memoizer caches based on tensor metadata (shape, stride,
diff --git a/test/inductor/test_ck_backend.py b/test/inductor/test_ck_backend.py
index 079be79fcc9d8..65da067671a0e 100644
--- a/test/inductor/test_ck_backend.py
+++ b/test/inductor/test_ck_backend.py
@@ -17,6 +17,7 @@
 from torch.testing._internal.common_utils import (
     instantiate_parametrized_tests,
     parametrize,
+    skipIfRocm,
 )
 from torch.testing._internal.inductor_utils import (
     _quantize_rowwise,
@@ -82,7 +83,8 @@ def mm(a, b):
         a = torch.randn(2240, 256, **tensor_options)
         b = torch.randn(256, 2048, **tensor_options)
 
-        assert "rocm" in dir(config)
+        if "rocm" not in dir(config):
+            raise AssertionError("'rocm' not found in dir(config)")
 
         with (
             config.patch(
@@ -132,7 +134,8 @@ def test_max_autotune_precompile_matmul_dynamic(
 
         torch._dynamo.mark_dynamic(a, 0)
 
-        assert "rocm" in dir(config)
+        if "rocm" not in dir(config):
+            raise AssertionError("'rocm' not found in dir(config)")
 
         with (
             config.patch(
@@ -178,7 +181,8 @@ def mm(a, b):
         a = torch.randn(2240, 256, **tensor_options)
         b = torch.randn(2048, 256, **tensor_options).transpose(0, 1)
 
-        assert "rocm" in dir(config)
+        if "rocm" not in dir(config):
+            raise AssertionError("'rocm' not found in dir(config)")
 
         with (
             config.patch(
@@ -210,7 +214,8 @@ def test_max_autotune_precompile_non_contiguous(self, max_autotune_gemm_backends
         a = torch.empty_strided((50257, 32768), (1, 50304), **tensor_options)
         b = torch.empty_strided((32768, 768), (768, 1), **tensor_options)
 
-        assert "rocm" in dir(config)
+        if "rocm" not in dir(config):
+            raise AssertionError("'rocm' not found in dir(config)")
 
         with (
             config.patch(
@@ -238,7 +243,10 @@ def mm(a, b):
     @unittest.skipIf(not torch.version.hip, "ROCM only")
     @unittest.mock.patch.dict(os.environ, _test_env)
     @parametrize("max_autotune_gemm_backends", ("CK", "ATen,Triton,CK"))
-    @parametrize("x_shape", ([4096, 2048], [2048], [4096, 1]))
+    @skipIfRocm(msg="Numerical accuracy errors in CK backend on gfx950 as of 06/03/26")
+    @parametrize(
+        "x_shape", ([4096, 2048], [2048], [4096, 1])
+    )  # NOTE: the first two shapes create "Tensor-likes are not close" errors
     def test_max_autotune_addmm(self, max_autotune_gemm_backends, x_shape):
         m, k, n = 4096, 224, 2048
         alpha, beta = 1.0, 1.0
@@ -248,7 +256,8 @@ def test_max_autotune_addmm(self, max_autotune_gemm_backends, x_shape):
         a = torch.randn(m, k, **tensor_options)
         b = torch.randn(k, n, **tensor_options)
 
-        assert "rocm" in dir(config)
+        if "rocm" not in dir(config):
+            raise AssertionError("'rocm' not found in dir(config)")
 
         with (
             config.patch(
@@ -319,7 +328,8 @@ def test_max_autotune_scaled_mm(
         # quantize input x
         x_fp8, x_inverse_scale = f_quantize(x, dtype_float8)
 
-        assert "rocm" in dir(config)
+        if "rocm" not in dir(config):
+            raise AssertionError("'rocm' not found in dir(config)")
 
         def linear(x_fp8, x_inverse_scale, w_t_fp8, w_inverse_scale, bias):
             y = torch._scaled_mm(
@@ -379,7 +389,8 @@ def test_max_autotune_conv2d(self, max_autotune_conv_backends):
         x_cl = x.to(memory_format=torch.channels_last)
         w_cl = w.to(memory_format=torch.channels_last)
 
-        assert "rocm" in dir(config)
+        if "rocm" not in dir(config):
+            raise AssertionError("'rocm' not found in dir(config)")
 
         with (
             config.patch(
@@ -423,7 +434,8 @@ def bmm(a, b):
         a = torch.randn(16, 2240, 256, **tensor_options)
         b = torch.randn(16, 2048, 256, **tensor_options).transpose(1, 2)
 
-        assert "rocm" in dir(config)
+        if "rocm" not in dir(config):
+            raise AssertionError("'rocm' not found in dir(config)")
 
         with (
             config.patch(
diff --git a/test/inductor/test_codecache.py b/test/inductor/test_codecache.py
index d0205cb295c06..d08d4d66fac15 100644
--- a/test/inductor/test_codecache.py
+++ b/test/inductor/test_codecache.py
@@ -10,7 +10,6 @@
 import textwrap
 import unittest
 from contextlib import contextmanager
-from typing import Optional, Union
 from typing_extensions import override
 from unittest import mock
 
@@ -149,6 +148,34 @@ def test_linemaps_empty(self):
         stack_frames = PyCodeCache.stack_frames_for_code(path, 0)
         self.assertEqual(stack_frames, None)
 
+    def test_load_by_key_path_can_defer_sys_modules_registration(self):
+        """
+        Benchmark-only modules should not need a global ``sys.modules`` entry,
+        but a later normal runtime load must still be able to publish the same
+        cached module there.
+        """
+        module_name = None
+        try:
+            PyCodeCache.cache_clear()
+            key, path = PyCodeCache.write("value = 1\n")
+
+            module = PyCodeCache.load_by_key_path(
+                key,
+                path,
+                set_sys_modules=False,
+            )
+            module_name = module.__name__
+            self.assertNotIn(module_name, sys.modules)
+
+            registered_module = PyCodeCache.load_by_key_path(key, path)
+            self.assertIs(registered_module, module)
+            self.assertIn(module_name, sys.modules)
+            self.assertIs(sys.modules[module_name], module)
+        finally:
+            if module_name is not None:
+                sys.modules.pop(module_name, None)
+            PyCodeCache.cache_clear()
+
     @unittest.skipIf(IS_FBCODE or IS_SANDCASTLE, "Skip in fbcode/sandcastle")
     def test_editable_cached_wrapper(self):
         with tempfile.TemporaryDirectory() as tmpdir:
@@ -1216,7 +1243,8 @@ def test_auto_functionalized_caching(self, variant):
         if variant == "v1":
             patch = torch._inductor.config.patch(enable_auto_functionalized_v2=False)
         else:
-            assert variant == "v2"
+            if variant != "v2":
+                raise AssertionError(f"Expected 'v2', got {variant!r}")
             patch = torch._inductor.config.patch(enable_auto_functionalized_v2=True)
 
         @torch.library.custom_op("mylib::sin_inplace", mutates_args=["x"])
@@ -1766,6 +1794,60 @@ def forward(self, x):
         self.assertEqual(counters["inductor"]["fxgraph_cache_miss"], 0)
         self.assertEqual(counters["inductor"]["fxgraph_cache_hit"], 2)
 
+    @config.patch({"fx_graph_cache": True})
+    @config.patch({"fx_graph_remote_cache": False})
+    def test_cache_guard_sqrt_no_recompilation(self):
+        """
+        Verify that guards containing OpaqueUnaryFn_sqrt (math.sqrt)
+        do not cause spurious recompilations when loaded from cache.
+
+        Previously, the guard printer emitted math.sqrt(...) for
+        OpaqueUnaryFn_sqrt, which when re-evaluated with SymInt inputs during
+        cache hit would force concretization/specialization of the symbol,
+        creating guards like `sym_float(size) == <concrete_value>` that didn't
+        exist in the original program. This caused a cache miss + recompilation
+        for every unique input size.
+
+        The fix emits torch._sym_sqrt(...) which propagates symbolically
+        without forcing specialization.
+
+        See https://github.com/pytorch/pytorch/issues/152435
+        """
+        import math
+
+        def func(x):
+            y = math.ceil((x.numel() // 5) / (math.ceil(math.sqrt(x.numel())))) > 64
+            if y:
+                return x * 5, y
+            else:
+                return x * 10, y
+
+        compiled_fn = torch.compile(func, dynamic=True)
+
+        # Warm up the cache with one size
+        compiled_fn(torch.rand(1000000))
+        self.assertEqual(counters["inductor"]["fxgraph_cache_miss"], 1)
+
+        # Simulate a new process loading from cache
+        self.reset()
+        counters.clear()
+
+        compiled_fn2 = torch.compile(func, dynamic=True)
+        # First call should hit the cache
+        compiled_fn2(torch.rand(2000000))
+        self.assertEqual(counters["inductor"]["fxgraph_cache_hit"], 1)
+        self.assertEqual(counters["inductor"]["fxgraph_cache_miss"], 0)
+
+        # Subsequent calls with different sizes should NOT recompile.
+        # Before the fix, each new size would create a spurious guard
+        # `sym_float(size) == <value>` causing a recompilation.
+        for size in [3000000, 5000000, 6000000, 7000000]:
+            compiled_fn2(torch.rand(size))
+
+        # All subsequent calls should reuse the already-loaded compiled graph
+        # without any additional cache misses or dynamo recompilations.
+        self.assertEqual(counters["inductor"]["fxgraph_cache_miss"], 0)
+
     @config.patch({"fx_graph_cache": True})
     @config.patch({"fx_graph_remote_cache": False})
     @config.patch({"freezing": True})
@@ -1920,7 +2002,8 @@ def f(x):
             )
             with fresh_cache():
                 gm, args, kwargs = self.capture(f)(x)
-                assert not kwargs
+                if kwargs:
+                    raise AssertionError
 
                 compiled_artifact = torch._inductor.standalone_compile(
                     gm, args, aot=is_aot
@@ -1961,7 +2044,8 @@ def f(x):
 
         with fresh_cache():
             gm, args, kwargs = self.capture(f)(x)
-            assert not kwargs
+            if kwargs:
+                raise AssertionError
 
             compiled_artifact = torch._inductor.standalone_compile(gm, args)
 
@@ -1996,6 +2080,45 @@ def backend(gm, args, **kwargs):
             compiled_out = torch.compile(f, fullgraph=True, backend=backend)(x)
             self.assertEqual(eager_out, compiled_out)
 
+    @config.patch({"fx_graph_cache": True})
+    @config.patch({"fx_graph_remote_cache": False})
+    @functorch_config.patch({"enable_autograd_cache": True})
+    @parametrize("donate", (False, True))
+    def test_donate_graph_module(self, donate: bool) -> None:
+        mod = torch.nn.Linear(1, 3)
+        x = torch.randn(4, 1)
+
+        def f(x):
+            with torch.no_grad():
+                return mod(x)
+
+        eager_out = f(x)
+
+        with fresh_cache():
+            gm, args, kwargs = self.capture(f)(x)
+            if kwargs:
+                raise AssertionError
+
+            # compile_fx mutates the graph module (e.g. adds
+            # mutation_region_id to node metadata). Use this as a
+            # fingerprint to detect mutation.
+            def has_mutation_region_ids(gm):
+                return any("mutation_region_id" in n.meta for n in gm.graph.nodes)
+
+            before = has_mutation_region_ids(gm)
+
+            compiled_artifact = torch._inductor.standalone_compile(
+                gm, args, donate_graph_module=donate
+            )
+            compiled_out = compiled_artifact(*args)
+            self.assertEqual(eager_out, compiled_out[0])
+
+            after = has_mutation_region_ids(gm)
+            if donate:
+                self.assertNotEqual(before, after)
+            else:
+                self.assertEqual(before, after)
+
     @config.patch({"fx_graph_cache": True})
     @config.patch({"fx_graph_remote_cache": False})
     @functorch_config.patch({"enable_autograd_cache": True})
@@ -2013,7 +2136,8 @@ def f(x):
             path = os.path.join(temp_dir, "new_dir")
             with fresh_cache():
                 gm, args, kwargs = self.capture(f)(x)
-                assert not kwargs
+                if kwargs:
+                    raise AssertionError
 
                 compiled_artifact = torch._inductor.standalone_compile(gm, args)
                 compiled_artifact.save(path=path, format="unpacked")
@@ -2046,7 +2170,8 @@ def f(x):
         with tempfile.TemporaryDirectory() as temp_dir:
             with fresh_cache():
                 gm, args, kwargs = self.capture(f)(x)
-                assert not kwargs
+                if kwargs:
+                    raise AssertionError
 
                 compiled_artifact = torch._inductor.standalone_compile(gm, args)
                 compiled_out = compiled_artifact(*args)
@@ -2064,7 +2189,8 @@ def f(x):
                     subdir_path = os.path.join(temp_dir, subdir)
                     for file in os.listdir(subdir_path):
                         file_path = os.path.join(subdir_path, file)
-                        assert os.path.isfile(file_path)
+                        if not os.path.isfile(file_path):
+                            raise AssertionError
                         with open(file_path) as f:
                             file_contents = f.read()
                         if device == GPU_TYPE:
@@ -2072,7 +2198,8 @@ def f(x):
                                 "2.0, tl.float32", "8.0, tl.float32"
                             )
                         else:
-                            assert device == "cpu"
+                            if device != "cpu":
+                                raise AssertionError(f"Expected 'cpu', got {device!r}")
                             file_contents = file_contents.replace(
                                 "auto tmp1 = static_cast<float>(2.0);",
                                 "auto tmp1 = static_cast<float>(8.0);",
@@ -2099,7 +2226,8 @@ def f(x):
             return x.sin() * 2
 
         gm, args, kwargs = self.capture(f)(x)
-        assert not kwargs
+        if kwargs:
+            raise AssertionError
 
         with tempfile.TemporaryDirectory() as temp_dir:
             path = normalize_path_separator(
@@ -2151,7 +2279,8 @@ def f(x):
         with fresh_cache():
             # captured graph is lambda s0, x: x * s0
             gm, args, kwargs = self.capture(f)(x)
-            assert not kwargs
+            if kwargs:
+                raise AssertionError
 
         compiled_artifact = torch._inductor.standalone_compile(
             gm, args, dynamic_shapes="from_graph", aot=is_aot
@@ -2245,7 +2374,8 @@ def f(x):
         with fresh_cache():
             # captured graph is lambda s0, x: x * s0
             gm, args, kwargs = self.capture(f)(x)
-            assert not kwargs
+            if kwargs:
+                raise AssertionError
 
         if config_patches:
             config_patches = {"fx_graph_cache": True}
@@ -2283,7 +2413,8 @@ def f(x):
         with fresh_cache():
             # static_gm is lambda x: x * 3
             static_gm, args, kwargs = self.capture(f, dynamic=False)(static_x)
-            assert not kwargs
+            if kwargs:
+                raise AssertionError
         compiled_artifact = torch._inductor.standalone_compile(
             static_gm, [static_x], dynamic_shapes=dynamic_shapes, aot=is_aot
         )
@@ -2381,7 +2512,7 @@ class TestCustomGraphPass(CustomGraphPass):
             def __call__(self, graph: torch.fx.graph.Graph) -> None:
                 return None
 
-            def uuid(self) -> Optional[Union[bytes, str]]:
+            def uuid(self) -> bytes | str | None:
                 return "uuid"
 
         def fn(a, b):
@@ -2434,7 +2565,7 @@ def __call__(
     ) -> tuple[torch.fx.GraphModule, torch.fx.GraphModule]:
         return gm, gm  # Dummy implementation
 
-    def uuid(self) -> Optional[Union[bytes, str]]:
+    def uuid(self) -> bytes | str | None:
         return self._uuid
 
 
@@ -2666,6 +2797,70 @@ def test_hash_private_config_changes(self):
             pickler.dumps(details3),
         )
 
+    def test_hash_provenance_tracking_level(self):
+        """
+        Test that provenance_tracking_level affects hashes.
+        """
+        with config.patch({"trace.provenance_tracking_level": 0}):
+            details1 = FxGraphHashDetails(None, [], {}, [])
+            details2 = FxGraphHashDetails(None, [], {}, [])
+
+        with config.patch({"trace.provenance_tracking_level": 1}):
+            details3 = FxGraphHashDetails(None, [], {}, [])
+
+        gm = torch.fx.GraphModule({}, torch.fx.Graph())
+        pickler = FxGraphCachePickler(gm)
+
+        self.assertEqual(
+            pickler.dumps(details1),
+            pickler.dumps(details2),
+        )
+        self.assertNotEqual(
+            pickler.dumps(details1),
+            pickler.dumps(details3),
+        )
+
+    def test_provenance_tracking_level_causes_cache_miss(self):
+        """
+        Test that changing provenance_tracking_level causes a cache miss.
+        """
+
+        class Mod(torch.nn.Module):
+            def __init__(self) -> None:
+                super().__init__()
+                self.param = torch.nn.Parameter(torch.rand(4, 4))
+
+            def forward(self, x):
+                return x @ self.param
+
+        mod = Mod()
+        mod_compiled = torch.compile(mod)
+        with torch.no_grad():
+            x = torch.rand(4, 4)
+            with config.patch({"trace.provenance_tracking_level": 0}):
+                # miss
+                mod_compiled(x)
+                self.assertEqual(counters["inductor"]["fxgraph_cache_bypass"], 0)
+                self.assertEqual(counters["inductor"]["fxgraph_cache_miss"], 1)
+                self.assertEqual(counters["inductor"]["fxgraph_cache_hit"], 0)
+                # hit
+                torch._dynamo.reset()
+                mod_compiled(x)
+                self.assertEqual(counters["inductor"]["fxgraph_cache_bypass"], 0)
+                self.assertEqual(counters["inductor"]["fxgraph_cache_miss"], 1)
+                self.assertEqual(counters["inductor"]["fxgraph_cache_hit"], 1)
+                torch._dynamo.reset()
+                counters.clear()
+
+            with config.patch({"trace.provenance_tracking_level": 1}):
+                # miss (provenance_tracking_level changed)
+                mod_compiled(x)
+                self.assertEqual(counters["inductor"]["fxgraph_cache_bypass"], 0)
+                self.assertEqual(counters["inductor"]["fxgraph_cache_miss"], 1)
+                self.assertEqual(counters["inductor"]["fxgraph_cache_hit"], 0)
+                torch._dynamo.reset()
+                counters.clear()
+
     def test_non_serializable_custom_passes_causes_cache_miss(self):
         class Mod(torch.nn.Module):
             def __init__(self) -> None:
@@ -2739,7 +2934,7 @@ def __init__(self):
             def __call__(self, graph: torch.fx.graph.Graph) -> None:
                 return None
 
-            def uuid(self) -> Optional[Union[bytes, str]]:
+            def uuid(self) -> bytes | str | None:
                 return self._uuid
 
         custom_pass = TestCustomGraphPass()
@@ -2775,7 +2970,7 @@ def __init__(self):
             def __call__(self, gm: torch.fx.GraphModule) -> None:
                 return None
 
-            def uuid(self) -> Optional[Union[bytes, str]]:
+            def uuid(self) -> bytes | str | None:
                 return self._uuid
 
         custom_pass = TestCustomGraphModulePass()
@@ -3001,26 +3196,39 @@ def test_cuda_compile_command(self):
         cmd_no_extra_args: str = cuda_compile_command(
             ["abc.cu", "def.cu"], "output", "so"
         )
-        assert "nvcc " in cmd_no_extra_args, cmd_no_extra_args
-        assert "abc.cu" in cmd_no_extra_args, cmd_no_extra_args
-        assert "def.cu" in cmd_no_extra_args, cmd_no_extra_args
-        assert "output" in cmd_no_extra_args, cmd_no_extra_args
+        if "nvcc " not in cmd_no_extra_args:
+            raise AssertionError(cmd_no_extra_args)
+        if "abc.cu" not in cmd_no_extra_args:
+            raise AssertionError(cmd_no_extra_args)
+        if "def.cu" not in cmd_no_extra_args:
+            raise AssertionError(cmd_no_extra_args)
+        if "output" not in cmd_no_extra_args:
+            raise AssertionError(cmd_no_extra_args)
         cmd_extra_args: str = cuda_compile_command(
             ["abc.cu", "def.cu"], "output", "so", ["-Wwhatever", "-nothing"]
         )
-        assert "nvcc " in cmd_extra_args, cmd_extra_args
-        assert " -Wwhatever" in cmd_extra_args, cmd_extra_args
-        assert " -nothing" in cmd_extra_args, cmd_extra_args
-        assert "abc.cu" in cmd_extra_args, cmd_extra_args
-        assert "def.cu" in cmd_extra_args, cmd_extra_args
-        assert "output " in cmd_extra_args, cmd_extra_args
+        if "nvcc " not in cmd_extra_args:
+            raise AssertionError(cmd_extra_args)
+        if " -Wwhatever" not in cmd_extra_args:
+            raise AssertionError(cmd_extra_args)
+        if " -nothing" not in cmd_extra_args:
+            raise AssertionError(cmd_extra_args)
+        if "abc.cu" not in cmd_extra_args:
+            raise AssertionError(cmd_extra_args)
+        if "def.cu" not in cmd_extra_args:
+            raise AssertionError(cmd_extra_args)
+        if "output " not in cmd_extra_args:
+            raise AssertionError(cmd_extra_args)
         with mock.patch("subprocess.check_output") as check_output_mock:
             CUDACodeCache.compile("test123.cu", "so", ["-Wsomething"])
             check_output_mock.assert_called()
             cmd_parts: list[str] = check_output_mock.call_args[0][0]
-            assert cmd_parts[0].endswith("nvcc"), cmd_parts
-            assert "-Wsomething" in cmd_parts, cmd_parts
-            assert "-DNDEBUG" in cmd_parts, cmd_parts
+            if not cmd_parts[0].endswith("nvcc"):
+                raise AssertionError(cmd_parts)
+            if "-Wsomething" not in cmd_parts:
+                raise AssertionError(cmd_parts)
+            if "-DNDEBUG" not in cmd_parts:
+                raise AssertionError(cmd_parts)
 
 
 @instantiate_parametrized_tests
@@ -3043,9 +3251,6 @@ def reset(self):
 
     @requires_cuda_and_triton
     @unittest.skipIf(not SM80OrLater, "Requires SM80+")
-    @unittest.skipIf(
-        TEST_WITH_ROCM, "Requires static cuda launcher, which does not support ROCM"
-    )
     @config.patch({"use_static_triton_launcher": True})
     @config.patch({"fx_graph_cache": True})
     @config.patch({"fx_graph_remote_cache": False})
@@ -3579,5 +3784,204 @@ def test_log_event_all_caches(self):
         )
 
 
+class TestAutotuneCacheExtraOptions(TestCase):
+    """
+    Unit tests for extra_options preservation in _load_cached_autotuning().
+
+    The extra_options field allows third-party backends to store custom tuned
+    options alongside the standard Triton config. These tests verify that
+    extra_options is correctly preserved when saving and loading from cache.
+    """
+
+    @requires_triton()
+    def test_load_cached_autotuning_preserves_extra_options_with_coordesc(self):
+        """
+        Test that extra_options is preserved when loading a config that was
+        found via coordinate descent tuning.
+        """
+        from torch._inductor.runtime.autotune_cache import _load_cached_autotuning
+
+        # Simulate a cached config with extra_options and found_by_coordesc=True
+        best_config = {
+            "BLOCK_M": 64,
+            "BLOCK_N": 64,
+            "num_warps": 4,
+            "num_stages": 2,
+            "configs_hash": "test_hash",
+            "found_by_coordesc": True,
+            "extra_options": {"custom_option": "value", "another_option": 42},
+        }
+
+        # Empty configs list since coordesc path creates a new config
+        configs = []
+        inductor_meta = {"coordinate_descent_tuning": True}
+
+        result = _load_cached_autotuning(
+            best_config.copy(), "test_hash", configs, inductor_meta
+        )
+
+        self.assertIsNotNone(result)
+        self.assertEqual(
+            result.extra_options, {"custom_option": "value", "another_option": 42}
+        )
+        self.assertTrue(result.found_by_coordesc)
+
+    @requires_triton()
+    def test_load_cached_autotuning_preserves_extra_options_with_matched_config(self):
+        """
+        Test that extra_options is preserved when loading by matching an
+        existing config from the configs list.
+        """
+        from triton import Config
+
+        from torch._inductor.runtime.autotune_cache import _load_cached_autotuning
+
+        # Create a config that matches what's in the cache
+        original_config = Config(
+            {"BLOCK_M": 64, "BLOCK_N": 64},
+            num_warps=4,
+            num_stages=2,
+        )
+        configs = [original_config]
+
+        # Simulate a cached config with extra_options
+        best_config = {
+            "BLOCK_M": 64,
+            "BLOCK_N": 64,
+            "num_warps": 4,
+            "num_stages": 2,
+            "configs_hash": "test_hash",
+            "extra_options": {"backend_specific": "option"},
+        }
+
+        inductor_meta = {"coordinate_descent_tuning": False}
+
+        result = _load_cached_autotuning(
+            best_config.copy(), "test_hash", configs, inductor_meta
+        )
+
+        self.assertIsNotNone(result)
+        self.assertEqual(result.extra_options, {"backend_specific": "option"})
+        # The returned config should be the matched one
+        self.assertIs(result, original_config)
+
+    @requires_triton()
+    def test_load_cached_autotuning_handles_none_extra_options(self):
+        """
+        Test that _load_cached_autotuning handles missing extra_options gracefully
+        (returns None for the attribute).
+        """
+        from triton import Config
+
+        from torch._inductor.runtime.autotune_cache import _load_cached_autotuning
+
+        # Create a config that matches what's in the cache
+        original_config = Config(
+            {"BLOCK_M": 64, "BLOCK_N": 64},
+            num_warps=4,
+            num_stages=2,
+        )
+        configs = [original_config]
+
+        # Simulate a cached config WITHOUT extra_options
+        best_config = {
+            "BLOCK_M": 64,
+            "BLOCK_N": 64,
+            "num_warps": 4,
+            "num_stages": 2,
+            "configs_hash": "test_hash",
+        }
+
+        inductor_meta = {"coordinate_descent_tuning": False}
+
+        result = _load_cached_autotuning(
+            best_config.copy(), "test_hash", configs, inductor_meta
+        )
+
+        self.assertIsNotNone(result)
+        # extra_options should be None when not in cache
+        self.assertIsNone(result.extra_options)
+
+    @requires_triton()
+    def test_autotune_cache_save_includes_extra_options(self):
+        """
+        Test that AutotuneCache.save() includes extra_options in the saved data.
+        """
+        from unittest.mock import MagicMock
+
+        from triton import Config
+
+        from torch._inductor.runtime.autotune_cache import AutotuneCache
+
+        # Create a config with extra_options
+        config_with_extra = Config(
+            {"BLOCK_M": 64, "BLOCK_N": 64},
+            num_warps=4,
+            num_stages=2,
+        )
+        config_with_extra.extra_options = {"custom_key": "custom_value"}
+
+        # Create an AutotuneCache instance with mocked caches
+        cache = AutotuneCache.__new__(AutotuneCache)
+        cache.configs_hash = "test_hash"
+
+        # Mock the local cache to capture what's being saved
+        mock_local_backend = MagicMock()
+        cache.local_cache = (mock_local_backend, "test_key")
+        cache.remote_cache = None
+
+        # Mock AutotuneCacheBundler and CacheArtifactManager
+        with (
+            mock.patch("torch._inductor.runtime.autotune_cache.AutotuneCacheBundler"),
+            mock.patch("torch._inductor.runtime.autotune_cache.CacheArtifactManager"),
+        ):
+            cache.save(config_with_extra, time_taken_ns=1000000)
+
+        # Verify extra_options was included in the saved data
+        mock_local_backend.put.assert_called_once()
+        saved_data = mock_local_backend.put.call_args[0][1]
+        self.assertIn("extra_options", saved_data)
+        self.assertEqual(saved_data["extra_options"], {"custom_key": "custom_value"})
+
+    @requires_triton()
+    def test_autotune_cache_save_omits_extra_options_when_none(self):
+        """
+        Test that AutotuneCache.save() does not include extra_options when it's None.
+        """
+        from unittest.mock import MagicMock
+
+        from triton import Config
+
+        from torch._inductor.runtime.autotune_cache import AutotuneCache
+
+        # Create a config without extra_options
+        config_without_extra = Config(
+            {"BLOCK_M": 64, "BLOCK_N": 64},
+            num_warps=4,
+            num_stages=2,
+        )
+
+        # Create an AutotuneCache instance with mocked caches
+        cache = AutotuneCache.__new__(AutotuneCache)
+        cache.configs_hash = "test_hash"
+
+        # Mock the local cache to capture what's being saved
+        mock_local_backend = MagicMock()
+        cache.local_cache = (mock_local_backend, "test_key")
+        cache.remote_cache = None
+
+        # Mock AutotuneCacheBundler and CacheArtifactManager
+        with (
+            mock.patch("torch._inductor.runtime.autotune_cache.AutotuneCacheBundler"),
+            mock.patch("torch._inductor.runtime.autotune_cache.CacheArtifactManager"),
+        ):
+            cache.save(config_without_extra, time_taken_ns=1000000)
+
+        # Verify extra_options was NOT included in the saved data
+        mock_local_backend.put.assert_called_once()
+        saved_data = mock_local_backend.put.call_args[0][1]
+        self.assertNotIn("extra_options", saved_data)
+
+
 if __name__ == "__main__":
     run_tests()
diff --git a/test/inductor/test_codegen_triton.py b/test/inductor/test_codegen_triton.py
index fa544073f5c9c..a44e63e0bc7e1 100644
--- a/test/inductor/test_codegen_triton.py
+++ b/test/inductor/test_codegen_triton.py
@@ -1,16 +1,30 @@
 # Owner(s): ["module: inductor"]
 import contextlib
+import unittest
 
 import sympy
 
 import torch
 import torch._inductor.config as inductor_config
 from torch._inductor.codegen import triton_utils
-from torch._inductor.codegen.common import SizeArg
+from torch._inductor.codegen.common import CSEVariable, SizeArg, TensorArg
+from torch._inductor.codegen.triton import (
+    _materialize_trunc_to_float_expr,
+    TritonKernelOverrides,
+)
+from torch._inductor.dtype_propagation import DtypePropagationOpsHandler, promote_types
 from torch._inductor.graph import GraphLowering
 from torch._inductor.test_case import TestCase as InductorTestCase
+from torch._inductor.utils import run_and_get_code
 from torch._inductor.virtualized import V
-from torch.testing._internal.inductor_utils import HAS_CPU, HAS_GPU
+from torch.testing._internal.inductor_utils import (
+    GPU_TYPE,
+    HAS_CPU,
+    HAS_GPU,
+    HAS_GPU_AND_TRITON,
+)
+from torch.utils._sympy.functions import FloorDiv, TruncToFloat, TruncToInt
+from torch.utils._sympy.value_ranges import ValueRanges
 
 
 class TestCodegenTriton(InductorTestCase):
@@ -56,10 +70,11 @@ def _check_divisibility(expected_divisible_indices, config):
             }:
                 self.assertEqual(expected_divisible_indices, config.divisibility_16)
             else:
-                assert (
+                if (
                     get_triton_attrs_descriptor_version()
-                    == TritonAttrsDescriptorVersion.V4_DICT
-                )
+                    != TritonAttrsDescriptorVersion.V4_DICT
+                ):
+                    raise AssertionError
                 self.assertIsInstance(config, dict)
 
                 for idx in expected_divisible_indices:
@@ -97,6 +112,151 @@ def _check_divisibility(expected_divisible_indices, config):
             ),
         )
 
+    def test_config_of_sizearg_with_check_constraint(self):
+        from torch.utils._sympy.functions import Mod
+
+        s2 = sympy.Symbol("s2", positive=True, integer=True)
+
+        self.assertFalse(
+            V.graph.sizevars.statically_known_multiple_of(s2, 16),
+        )
+
+        shape_env = V.graph.sizevars.shape_env
+        shape_env.axioms[sympy.Eq(Mod(s2, 16), 0)] = sympy.true
+
+        self.assertTrue(
+            V.graph.sizevars.statically_known_multiple_of(s2, 16),
+        )
+
+    def test_pow_uses_active_override_constant_lowering(self):
+        exponent = CSEVariable("ks0", ValueRanges.unknown(), torch.int64)
+
+        class TestTritonKernelOverrides(TritonKernelOverrides):
+            @classmethod
+            def constant(cls, value, dtype):
+                return f"custom_constant({value}, {dtype})"
+
+        self.assertEqual(
+            TestTritonKernelOverrides.pow(2, exponent),
+            "libdevice.pow(custom_constant(2, torch.float64), (ks0).to(tl.float64))",
+        )
+
+    def test_pow_preserves_integer_dtype_for_unsigned_scalar_exponents(self):
+        exponent = CSEVariable("ks0", ValueRanges.unknown(), torch.uint32)
+
+        self.assertEqual(
+            DtypePropagationOpsHandler().pow(2, exponent),
+            promote_types([2, exponent]),
+        )
+
+    def test_pow_uses_integer_helper_for_unsigned_scalar_exponents(self):
+        exponent = CSEVariable("ks0", ValueRanges.unknown(), torch.uint32)
+
+        class TestTritonKernelOverrides(TritonKernelOverrides):
+            @classmethod
+            def constant(cls, value, dtype):
+                return f"custom_constant({value}, {dtype})"
+
+        self.assertEqual(
+            TestTritonKernelOverrides.pow(3, exponent),
+            "triton_helpers.pow_integer(custom_constant(3, torch.uint32), ks0)",
+        )
+
+    def test_materialize_trunc_to_float_expr_preserves_integer_subexpressions(self):
+        s0 = sympy.Symbol("s0")
+
+        trunc_expr = TruncToInt(s0)
+        self.assertEqual(
+            _materialize_trunc_to_float_expr(trunc_expr, torch.float64),
+            TruncToFloat(s0),
+        )
+
+        integer_expr = FloorDiv(trunc_expr, sympy.Integer(5))
+        self.assertEqual(
+            _materialize_trunc_to_float_expr(integer_expr, torch.float64),
+            integer_expr,
+        )
+
+        predicate_expr = sympy.Eq(trunc_expr, sympy.Integer(9007199254740993))
+        self.assertEqual(
+            _materialize_trunc_to_float_expr(predicate_expr, torch.float64),
+            predicate_expr,
+        )
+
+        float_expr = sympy.Float(0.5) + trunc_expr
+        self.assertEqual(
+            _materialize_trunc_to_float_expr(float_expr, torch.float64),
+            sympy.Float(0.5) + TruncToFloat(s0),
+        )
+
+    @unittest.skipUnless(torch.version.hip is not None, "pointer_range_32 is HIP-only")
+    @unittest.skipUnless(HAS_GPU_AND_TRITON, "requires GPU and Triton")
+    def test_pointer_range_in_generated_code(self):
+        """Verify tt.pointer_range=32 appears in generated Triton code on HIP."""
+
+        def fn(x):
+            return x + 1
+
+        x = torch.randn(64, 64, device=GPU_TYPE, dtype=torch.bfloat16)
+        _, code = run_and_get_code(torch.compile(fn), x)
+        code_str = " ".join(code)
+        self.assertIn("tt.pointer_range", code_str)
+
+    def test_is_multiple_of_rules(self):
+        """Test structural divisibility rules in _is_multiple_of."""
+        from torch.utils._sympy.functions import FloorDiv, Mod
+
+        sv = V.graph.sizevars
+        shape_env = sv.shape_env
+
+        s1 = sympy.Symbol("s1", positive=True, integer=True)
+        s2 = sympy.Symbol("s2", positive=True, integer=True)
+        s3 = sympy.Symbol("s3", positive=True, integer=True)
+
+        # Product: any factor divisible → product divisible
+        self.assertTrue(sv.statically_known_multiple_of(16 * s1, 16))
+        self.assertTrue(sv.statically_known_multiple_of(4 * 4 * s1, 16))
+        shape_env.axioms[sympy.Eq(Mod(s1, 16), 0)] = sympy.true
+        self.assertTrue(sv.statically_known_multiple_of(s1 * s2, 16))
+        self.assertFalse(sv.statically_known_multiple_of(s2 * s3, 16))
+
+        # Sum: all terms divisible → sum divisible
+        self.assertFalse(sv.statically_known_multiple_of(s1 + s2, 16))
+        shape_env.axioms[sympy.Eq(Mod(s2, 16), 0)] = sympy.true
+        self.assertTrue(sv.statically_known_multiple_of(s1 + s2, 16))
+        self.assertTrue(sv.statically_known_multiple_of(s1 + 32, 16))
+        self.assertFalse(sv.statically_known_multiple_of(s1 + 3, 16))
+
+        # FloorDiv(a, b): a must be multiple of b*n
+        self.assertFalse(sv.statically_known_multiple_of(FloorDiv(s1, 3), 16))
+        shape_env.axioms[sympy.Eq(Mod(s3, 48), 0)] = sympy.true
+        self.assertTrue(sv.statically_known_multiple_of(FloorDiv(s3, 3), 16))
+
+        # Mod(a, b): both a and b must be multiples of n
+        self.assertTrue(sv.statically_known_multiple_of(Mod(s1, 48), 16))
+        s_nodiv = sympy.Symbol("s_nodiv", positive=True, integer=True)
+        self.assertFalse(sv.statically_known_multiple_of(Mod(s_nodiv, 32), 16))
+        self.assertFalse(sv.statically_known_multiple_of(Mod(s1, 7), 16))
+
+        # Axiom fallback: bare symbol resolved via statically_known_true
+        s4 = sympy.Symbol("s4", positive=True, integer=True)
+        self.assertFalse(sv.statically_known_multiple_of(s4, 8))
+        shape_env.axioms[sympy.Eq(Mod(s4, 8), 0)] = sympy.true
+        self.assertTrue(sv.statically_known_multiple_of(s4, 8))
+
+    def test_signature_of_fp8_dtypes(self):
+        """fp8 dtypes should produce correct Triton pointer signatures via _type_of."""
+        expected = {
+            torch.float8_e4m3fn: "*fp8e4nv",
+            torch.float8_e5m2: "*fp8e5",
+            torch.float8_e4m3fnuz: "*fp8e4b8",
+            torch.float8_e5m2fnuz: "*fp8e5b16",
+        }
+        for dtype, expected_sig in expected.items():
+            arg = TensorArg(name="x", buffer="buf0", dtype=dtype)
+            sig = triton_utils.signature_of(arg, size_dtype=None)
+            self.assertEqual(sig, expected_sig, f"wrong signature for {dtype}")
+
 
 if __name__ == "__main__":
     from torch._inductor.test_case import run_tests
diff --git a/test/inductor/test_combo_kernels.py b/test/inductor/test_combo_kernels.py
index f18ff7c8a8269..7ee8df912b710 100644
--- a/test/inductor/test_combo_kernels.py
+++ b/test/inductor/test_combo_kernels.py
@@ -2,6 +2,8 @@
 
 import contextlib
 import json
+import logging
+import re
 import sys
 import tempfile
 import unittest
@@ -19,7 +21,11 @@
     TestCase,
 )
 from torch.testing._internal.inductor_utils import GPU_TYPE, HAS_CPU, HAS_GPU_AND_TRITON
-from torch.testing._internal.triton_utils import requires_gpu_and_triton
+from torch.testing._internal.triton_utils import (
+    requires_cuda_and_triton,
+    requires_gpu_and_triton,
+    requires_xpu_and_triton,
+)
 
 
 aten = torch.ops.aten
@@ -710,7 +716,7 @@ def test_mutated(a, b, c, d):
         out_compiled = torch.compile(test_mutated)(*inps)
 
         self.assertEqual(out_eager, out_compiled)
-        self.assertTrue(torch._inductor.metrics.generated_kernel_count in [6, 9])
+        self.assertTrue(4 < torch._inductor.metrics.generated_kernel_count <= 10)
 
     @requires_gpu_and_triton
     def test_round_robin_dispatch(self):
@@ -995,7 +1001,7 @@ def fn(a0, a1, a2, b0, b1, b2):
         self.assertEqual(out_eager, out_compiled)
         self.assertTrue(5 <= torch._inductor.metrics.generated_kernel_count <= 6)
 
-    @requires_gpu_and_triton
+    @requires_cuda_and_triton
     @torch._dynamo.config.patch("automatic_dynamic_shapes", True)
     @torch._dynamo.config.patch("assume_static_by_default", True)
     @torch._inductor.config.patch("triton.autotune_at_compile_time", True)
@@ -1016,6 +1022,27 @@ def fn(x, y, z):
 
         self.assertEqual(out_eager, out_compiled)
 
+    @requires_xpu_and_triton
+    @torch._dynamo.config.patch("automatic_dynamic_shapes", True)
+    @torch._dynamo.config.patch("assume_static_by_default", True)
+    @torch._inductor.config.patch("triton.autotune_at_compile_time", True)
+    def test_dynamic_shapes_persistent_reduction_mixed_x_dim_xpu(self):
+        def fn(x, y, z):
+            return x.sum(1), y.mean(1), z.max(1)
+
+        inps = (
+            torch.rand(16, 128, device=GPU_TYPE),
+            torch.rand(32, 128, device=GPU_TYPE),
+            torch.rand(32, 256, device=GPU_TYPE),
+        )
+        torch._dynamo.mark_dynamic(inps[0], 0, min=1, max=256)
+        torch._dynamo.mark_dynamic(inps[1], 0, min=1, max=256)
+        torch._dynamo.mark_dynamic(inps[2], 0, min=1, max=256)
+        out_eager = fn(*inps)
+        out_compiled = torch.compile(fn)(*inps)
+
+        self.assertEqual(out_eager, out_compiled)
+
     @requires_gpu_and_triton
     def test_helper_fn_defined(self):
         def fn(x, y, z):
@@ -1173,6 +1200,146 @@ def fn(x, y):
         self.assertEqual(out_eager, out_compiled)
 
 
+class ComboKernelTestsMaxAutotune(TestCase):
+    def setUp(self):
+        super().setUp()
+        torch._inductor.metrics.reset()
+        self._test_stack = contextlib.ExitStack()
+        self._test_stack.enter_context(
+            torch._inductor.config.patch(
+                {
+                    "combo_kernels": True,
+                    "benchmark_combo_kernel": False,
+                    "combo_kernel_per_subkernel_blocks": True,
+                    "max_autotune": True,
+                    "autotune_local_cache": False,
+                }
+            )
+        )
+
+    def tearDown(self):
+        self._test_stack.close()
+        torch._inductor.metrics.reset()
+        super().tearDown()
+
+    @requires_gpu_and_triton
+    def test_combo_kernel_max_autotune(self):
+        def fn(a, b, c):
+            a1 = torch.nn.functional.relu(a)
+            b1 = torch.nn.functional.sigmoid(b)
+            c1 = torch.nn.functional.tanh(c)
+            return a1, b1, c1
+
+        inps = [
+            torch.rand(32, 1024, device=GPU_TYPE),
+            torch.rand(64, 512, device=GPU_TYPE),
+            torch.rand(16, 2048, device=GPU_TYPE),
+        ]
+
+        out_eager = fn(*inps)
+        fn_c = torch.compile(fn)
+
+        logger = logging.getLogger("torch._inductor.runtime.triton_heuristics")
+        with self.assertLogs(logger, level=logging.DEBUG) as cm:
+            out_compiled, code = run_and_get_code(fn_c, *inps)
+        chained_logs = [msg for msg in cm.output if "Combo sequential autotune" in msg]
+        self.assertGreater(
+            len(chained_logs),
+            0,
+            "_combo_sequential_autotune was not invoked",
+        )
+        self.assertEqual(out_eager, out_compiled)
+        self.assertEqual(torch._inductor.metrics.generated_kernel_count, 1)
+
+    @requires_gpu_and_triton
+    def test_combo_kernel_max_autotune_with_reduction(self):
+        def fn(x, y):
+            return x.sum(dim=-1), y.mean(dim=-1)
+
+        inps = [
+            torch.rand(128, 256, device=GPU_TYPE),
+            torch.rand(128, 256, device=GPU_TYPE),
+        ]
+
+        out_eager = fn(*inps)
+        fn_c = torch.compile(fn)
+
+        logger = logging.getLogger("torch._inductor.runtime.triton_heuristics")
+        with self.assertLogs(logger, level=logging.DEBUG) as cm:
+            out_compiled, code = run_and_get_code(fn_c, *inps)
+        chained_logs = [msg for msg in cm.output if "Combo sequential autotune" in msg]
+        self.assertGreater(
+            len(chained_logs),
+            0,
+            "_combo_sequential_autotune was not invoked",
+        )
+        self.assertEqual(out_eager, out_compiled)
+        self.assertEqual(torch._inductor.metrics.generated_kernel_count, 1)
+
+    @requires_gpu_and_triton
+    def test_combo_autotune_many_subkernels(self):
+        def fn(a, b, c, d, e, f):
+            return (
+                a * 2.0,
+                b + 1.0,
+                c.sin(),
+                d.cos(),
+                e.exp(),
+                f.neg(),
+            )
+
+        inps = [
+            torch.rand(8, 8192, device=GPU_TYPE),
+            torch.rand(128, 64, device=GPU_TYPE),
+            torch.rand(16, 4096, device=GPU_TYPE),
+            torch.rand(512, 16, device=GPU_TYPE),
+            torch.rand(32, 2048, device=GPU_TYPE),
+            torch.rand(256, 32, device=GPU_TYPE),
+        ]
+
+        out_eager = fn(*inps)
+        fn_c = torch.compile(fn)
+
+        logger = logging.getLogger("torch._inductor.runtime.triton_heuristics")
+        with self.assertLogs(logger, level=logging.DEBUG) as cm:
+            out_compiled, code = run_and_get_code(fn_c, *inps)
+
+        chained_logs = [msg for msg in cm.output if "Combo sequential autotune" in msg]
+        self.assertGreater(len(chained_logs), 0)
+        self.assertEqual(out_eager, out_compiled)
+        self.assertEqual(torch._inductor.metrics.generated_kernel_count, 1)
+
+    @requires_gpu_and_triton
+    def test_combo_kernel_per_subkernel_reduction_hint(self):
+        def fn(x, y):
+            return x.sum(dim=-1), y.sum(dim=0)
+
+        inps = [
+            torch.rand(128, 256, device=GPU_TYPE),
+            torch.rand(128, 256, device=GPU_TYPE),
+        ]
+
+        out_eager = fn(*inps)
+        out, code = run_and_get_code(torch.compile(fn), *inps)
+        self.assertEqual(out_eager, out)
+        # Verify per-subkernel reduction hints in generated code
+        found_hints = {}
+        for c in code:
+            for key in ["reduction_hint_0", "reduction_hint_1"]:
+                m = re.search(rf"'{key}':\s*'(\w+)'", c)
+                if m:
+                    found_hints[key] = m.group(1)
+
+        self.assertIn(
+            "reduction_hint_0", found_hints, "Missing per-subkernel reduction_hint_0"
+        )
+        self.assertIn(
+            "reduction_hint_1", found_hints, "Missing per-subkernel reduction_hint_1"
+        )
+        self.assertEqual(found_hints["reduction_hint_0"], "INNER")
+        self.assertEqual(found_hints["reduction_hint_1"], "OUTER")
+
+
 if __name__ == "__main__":
     from torch._dynamo.test_case import run_tests
 
diff --git a/test/inductor/test_comm_analysis.py b/test/inductor/test_comm_analysis.py
new file mode 100644
index 0000000000000..0ecf35dcec2de
--- /dev/null
+++ b/test/inductor/test_comm_analysis.py
@@ -0,0 +1,134 @@
+# Owner(s): ["module: inductor"]
+
+import sys
+import warnings
+
+import torch
+import torch.distributed as dist
+
+
+if not dist.is_available() or not dist.is_nccl_available():
+    print("c10d NCCL not available, skipping tests", file=sys.stderr)
+    sys.exit(0)
+
+try:
+    from torch.testing._internal.common_distributed import (
+        requires_nccl,
+        skip_if_lt_x_gpu,
+    )
+except ImportError:
+    print("common_distributed not importable, skipping tests", file=sys.stderr)
+    sys.exit(0)
+
+from torch.fx.experimental.proxy_tensor import make_fx
+from torch.testing._internal.common_utils import run_tests, TestCase
+
+
+def _get_all_gather_node(group_size, group_name):
+    """Trace a simple all_gather function and return the collective FX node."""
+
+    def func(inp, group_size, group_name):
+        out = torch.ops._c10d_functional.all_gather_into_tensor(
+            inp, group_size, group_name
+        )
+        wait = torch.ops._c10d_functional.wait_tensor(out)
+        return wait
+
+    gm = make_fx(func)(torch.ones(4, 4), group_size, group_name)
+    for n in gm.graph.nodes:
+        if n.op == "call_function" and "all_gather_into_tensor" in str(n.target):
+            return n
+    raise RuntimeError("No all_gather_into_tensor node found in traced graph")
+
+
+class TestNcclEstimateDeviceResolution(TestCase):
+    """
+    Tests for the device resolution fix in _nccl_estimate() inside
+    estimate_nccl_collective_runtime_from_fx_node.
+    """
+
+    def _init_pg(self, backend, world_size=2):
+        from torch.testing._internal.distributed.fake_pg import FakeStore
+
+        store = FakeStore()
+        dist.init_process_group(
+            backend=backend, rank=0, world_size=world_size, store=store
+        )
+        pg = dist.group.WORLD
+        group_name = "test_comm_analysis"
+        torch._C._distributed_c10d._register_process_group(group_name, pg)
+        return pg, group_name, pg.size()
+
+    def _init_pg_real_store(self, backend, world_size=1):
+        store = dist.HashStore()
+        dist.init_process_group(
+            backend=backend, rank=0, world_size=world_size, store=store
+        )
+        pg = dist.group.WORLD
+        group_name = "test_comm_analysis"
+        torch._C._distributed_c10d._register_process_group(group_name, pg)
+        return pg, group_name, pg.size()
+
+    def _destroy_pg(self):
+        dist.destroy_process_group()
+
+    def test_fake_backend_falls_back_to_analytical(self):
+        """FAKE backend: _nccl_estimate returns None, falls back to analytical formula."""
+        pg, group_name, group_size = self._init_pg("fake")
+        try:
+            node = _get_all_gather_node(group_size, group_name)
+            from torch._inductor.comm_analysis import (
+                estimate_nccl_collective_runtime_from_fx_node,
+            )
+
+            est_ms = estimate_nccl_collective_runtime_from_fx_node(
+                node, use_nccl_estimator=True
+            )
+            self.assertGreater(est_ms, 0)
+
+            est_ms_analytical = estimate_nccl_collective_runtime_from_fx_node(
+                node, use_nccl_estimator=False
+            )
+            self.assertEqual(est_ms, est_ms_analytical)
+        finally:
+            self._destroy_pg()
+
+    @requires_nccl()
+    @skip_if_lt_x_gpu(1)
+    def test_multi_backend_pg_resolves_to_nccl(self):
+        """
+        Multi-backend PG ("cpu:gloo,cuda:nccl"): We should resolve to the cuda device's backend.
+        """
+        torch.cuda.set_device(0)
+        pg, group_name, group_size = self._init_pg_real_store("cpu:gloo,cuda:nccl")
+        try:
+            from torch.distributed.distributed_c10d import _get_pg_default_device
+
+            with warnings.catch_warnings():
+                warnings.simplefilter("ignore", FutureWarning)
+                default_device = _get_pg_default_device(pg)
+            self.assertEqual(default_device, torch.device("cpu"))
+
+            nccl_backend = pg._get_backend(torch.device("cuda"))
+            self.assertTrue(nccl_backend.supports_time_estimate)
+
+            gloo_backend = pg._get_backend(torch.device("cpu"))
+            self.assertFalse(gloo_backend.supports_time_estimate)
+        finally:
+            self._destroy_pg()
+
+    @requires_nccl()
+    @skip_if_lt_x_gpu(1)
+    def test_single_nccl_backend_resolves_correctly(self):
+        """Single NCCL backend PG: cuda device resolves to NCCL with time estimation."""
+        torch.cuda.set_device(0)
+        pg, group_name, group_size = self._init_pg_real_store("nccl")
+        try:
+            backend = pg._get_backend(torch.device("cuda"))
+            self.assertTrue(backend.supports_time_estimate)
+        finally:
+            self._destroy_pg()
+
+
+if __name__ == "__main__":
+    run_tests()
diff --git a/test/inductor/test_compile_subprocess.py b/test/inductor/test_compile_subprocess.py
index e5b5567132bf7..9344dc777be89 100644
--- a/test/inductor/test_compile_subprocess.py
+++ b/test/inductor/test_compile_subprocess.py
@@ -76,6 +76,11 @@
     "test_weight_norm_conv2d": TestFailure(("cpu", "cuda"), is_skip=True),
 }
 
+if TEST_WITH_ROCM and not torch.cuda.has_magma:
+    test_failures["test_linalg_eig_stride_consistency"] = TestFailure(
+        ("cuda",), is_skip=True
+    )
+
 
 class TestSubprocess(TestCase):
     def setUp(self):
@@ -132,7 +137,8 @@ def save_output_code(code: str) -> None:
 
         with contextlib.ExitStack() as stack:
             # When this bug is fixed, remove the cache disabling below
-            assert torch._inductor.compile_fx_async.BUG_CACHES_DONT_WORK_WITH_ASYNC
+            if not torch._inductor.compile_fx_async.BUG_CACHES_DONT_WORK_WITH_ASYNC:
+                raise AssertionError
             stack.enter_context(
                 torch._inductor.config.patch(
                     autotune_local_cache=False, fx_graph_cache=False
@@ -201,7 +207,8 @@ def model_add(x, y):
         _AsyncFxCompile._reset_stats()
 
         with contextlib.ExitStack() as stack:
-            assert torch._inductor.compile_fx_async.BUG_CACHES_DONT_WORK_WITH_ASYNC
+            if not torch._inductor.compile_fx_async.BUG_CACHES_DONT_WORK_WITH_ASYNC:
+                raise AssertionError
             stack.enter_context(
                 torch._inductor.config.patch(
                     autotune_local_cache=False, fx_graph_cache=False
diff --git a/test/inductor/test_compile_worker.py b/test/inductor/test_compile_worker.py
index 270b15fdf49d8..1ef4fc9a3bc3d 100644
--- a/test/inductor/test_compile_worker.py
+++ b/test/inductor/test_compile_worker.py
@@ -163,6 +163,75 @@ def doit():
         t.quit()
 
 
+class TestSetTritonLibdevicePath(TestCase):
+    @config.patch({"compile_threads": 1, "eager_numerics.use_pytorch_libdevice": True})
+    def test_libdevice_path_no_subprocess(self):
+        """Test libdevice path is set with compile_threads=1 (no subprocess)."""
+        self._test_libdevice_path_with_compilation()
+
+    @config.patch("eager_numerics.use_pytorch_libdevice", True)
+    def test_libdevice_path_default_threads(self):
+        """Test libdevice path is set with default compile_threads (subprocess)."""
+        self._test_libdevice_path_with_compilation()
+
+    @config.patch(
+        {
+            "eager_numerics.use_pytorch_libdevice": True,
+            "eager_numerics.division_rounding": True,
+            "emulate_precision_casts": True,
+            "compile_threads": 1,
+        }
+    )
+    def test_pow_bitwise_precision(self):
+        """Test that compiled pow matches eager bitwise with system libdevice."""
+        import torch
+        from torch.utils.cpp_extension import CUDA_HOME
+
+        if not torch.cuda.is_available():
+            self.skipTest("CUDA not available")
+        if CUDA_HOME is None:
+            self.skipTest("CUDA_HOME not set")
+        expected = os.path.join(CUDA_HOME, "nvvm", "libdevice", "libdevice.10.bc")
+        if not os.path.isfile(expected):
+            self.skipTest(f"libdevice not found at {expected}")
+
+        torch._dynamo.reset()
+        torch.manual_seed(42)
+        base = torch.randn(1000, device="cuda", dtype=torch.float32).abs() + 1e-6
+        exp = torch.randn(1000, device="cuda", dtype=torch.float32)
+
+        eager_result = torch.pow(base, exp)
+        compiled_result = torch.compile(torch.pow)(base, exp)
+        self.assertEqual(eager_result, compiled_result, atol=0, rtol=0)
+
+    def _test_libdevice_path_with_compilation(self):
+        import torch
+        from torch.utils.cpp_extension import CUDA_HOME
+
+        if not torch.cuda.is_available():
+            self.skipTest("CUDA not available")
+
+        if CUDA_HOME is None:
+            self.skipTest("CUDA_HOME not set")
+
+        expected = os.path.join(CUDA_HOME, "nvvm", "libdevice", "libdevice.10.bc")
+        if not os.path.isfile(expected):
+            self.skipTest(f"libdevice not found at {expected}")
+
+        # Compile a simple function that uses pow (which uses libdevice)
+        @torch.compile
+        def fn(x):
+            return torch.pow(x, 2.0)
+
+        x = torch.randn(10, device="cuda", dtype=torch.float32)
+        fn(x)
+
+        # Verify libdevice path was set
+        from triton import knobs
+
+        self.assertEqual(knobs.nvidia.libdevice_path, expected)
+
+
 if __name__ == "__main__":
     from torch._inductor.test_case import run_tests
 
diff --git a/test/inductor/test_compiled_autograd.py b/test/inductor/test_compiled_autograd.py
index 3c945867ad63f..8e581956f60da 100644
--- a/test/inductor/test_compiled_autograd.py
+++ b/test/inductor/test_compiled_autograd.py
@@ -53,7 +53,7 @@
     HAS_CPU,
     HAS_CUDA_AND_TRITON,
     HAS_GPU,
-    HAS_XPU_AND_TRITON,
+    HAS_GPU_AND_TRITON,
 )
 from torch.testing._internal.logging_utils import logs_to_string
 from torch.testing._internal.triton_utils import (
@@ -69,7 +69,8 @@
 def make_compiler_fn(
     fullgraph=True, dynamic=True, backend="inductor", gm_hook=lambda gm: None
 ):
-    assert backend in ["inductor", "aot_eager", "eager", "ca_eager"]
+    if backend not in ["inductor", "aot_eager", "eager", "ca_eager"]:
+        raise AssertionError(f"unexpected backend: {backend}")
 
     def _compiler_fn(gm):
         """Same as torch.compile() but counts number of compiles"""
@@ -180,14 +181,14 @@ def run_as_subprocess(self, script) -> bytes:
     def test_hipify_not_loaded_with_import_torch(self):
         script = """
 import torch
-assert globals().get("hipify", False) is False
+assert globals().get("hipify", False) is False  # noqa: S101
 """
         self.run_as_subprocess(script)
 
     def test_hipify_not_loaded_with_import_cpp_extension(self):
         script = """
 import torch.utils.cpp_extension
-assert globals().get("hipify", False) is False
+assert globals().get("hipify", False) is False  # noqa: S101
 """
         self.run_as_subprocess(script)
 
@@ -208,7 +209,7 @@ def model(i):
         loss = out.sum()
         with torch._dynamo.compiled_autograd._enable(compiler_fn):
             loss.backward()
-        assert(w.grad is not None)
+        assert(w.grad is not None)  # noqa: S101
 
     inner()
     torch._dynamo.reset()
@@ -242,17 +243,21 @@ def test_reset(self):
         # state should be clean after reset
         compiled_autograd.reset()
 
-        assert compiled_autograd.compiled_autograd_enabled is False
+        if compiled_autograd.compiled_autograd_enabled is not False:
+            raise AssertionError
         (
             prior_compiler,
             prior_dynamic,
         ) = torch._C._dynamo.compiled_autograd.set_autograd_compiler(None, False)
-        assert prior_compiler is None
-        assert prior_dynamic is False
-        assert (
+        if prior_compiler is not None:
+            raise AssertionError
+        if prior_dynamic is not False:
+            raise AssertionError
+        if not (
             compiled_autograd.COMPILE_COUNTER is not None
             and next(compiled_autograd.COMPILE_COUNTER) == 0
-        )
+        ):
+            raise AssertionError
 
     def test_basic(self):
         def fn():
@@ -456,8 +461,10 @@ def test_reorder_post_hook2(self):
         y = torch.sigmoid(x)
         z = torch.tanh(y)
 
-        assert isinstance(z.grad_fn, torch.autograd.graph.Node)
-        assert isinstance(y.grad_fn, torch.autograd.graph.Node)
+        if not isinstance(z.grad_fn, torch.autograd.graph.Node):
+            raise AssertionError
+        if not isinstance(y.grad_fn, torch.autograd.graph.Node):
+            raise AssertionError
         handle_z = z.grad_fn.register_hook(lambda gI, gO: (gO[0] * 2,))
         handle_y = y.grad_fn.register_hook(lambda gI, gO: (gI[0] * 2,))
         z.sum().backward(retain_graph=True)
@@ -475,7 +482,8 @@ def test_reorder_post_hook3(self):
         x = torch.randn([1, 4, 32, 32])
         y = conv(x)
 
-        assert isinstance(y.grad_fn, torch.autograd.graph.Node)
+        if not isinstance(y.grad_fn, torch.autograd.graph.Node):
+            raise AssertionError
         # this hook will mul 2.0 to the conv weight gradient
         handle_y = y.grad_fn.register_hook(lambda gI, gO: (gI[0], gI[1] * 2, gI[2]))
         y.sum().backward(retain_graph=True)
@@ -1008,10 +1016,12 @@ def inner_compiler(gm_, example_inputs_):
                 placeholders = get_placeholders(gm_)
                 if is_bwd:
                     # boxed inputs
-                    assert isinstance(placeholders[0].meta["example_value"], list)
+                    if not isinstance(placeholders[0].meta["example_value"], list):
+                        raise AssertionError
                 else:
                     # not boxed inputs
-                    assert not isinstance(placeholders[0].meta["example_value"], list)
+                    if isinstance(placeholders[0].meta["example_value"], list):
+                        raise AssertionError
 
                 return gm_
 
@@ -1072,8 +1082,20 @@ def bytecode_hook(code, out_code):
                 call_op = "CALL"
 
             insts = list(dis.get_instructions(out_code))
+            # Find the CALL that invokes the compiled graph function
+            # (not an earlier CALL from e.g. store_user_object_weakrefs).
+            # The compiled fn is loaded via LOAD_GLOBAL __compiled_fn_*.
+            load_graph_idx = next(
+                i
+                for i, inst in enumerate(insts)
+                if inst.opname == "LOAD_GLOBAL"
+                and isinstance(inst.argval, str)
+                and inst.argval.startswith("__compiled_fn")
+            )
             call_graph_idx = next(
-                i for i, inst in enumerate(insts) if inst.opname == call_op
+                i
+                for i, inst in enumerate(insts)
+                if i > load_graph_idx and inst.opname == call_op
             )
             # pre-graph should alias: inputs_ref_0 = inputs[0]
             matches = [
@@ -1145,8 +1167,19 @@ def bytecode_hook(code, out_code):
                     call_op = "CALL"
 
                 insts = list(dis.get_instructions(out_code))
+                # Find the CALL that invokes the compiled graph function
+                # (not an earlier CALL from e.g. store_user_object_weakrefs).
+                load_graph_idx = next(
+                    i
+                    for i, inst in enumerate(insts)
+                    if inst.opname == "LOAD_GLOBAL"
+                    and isinstance(inst.argval, str)
+                    and inst.argval.startswith("__compiled_fn")
+                )
                 call_graph_idx = next(
-                    i for i, inst in enumerate(insts) if inst.opname == call_op
+                    i
+                    for i, inst in enumerate(insts)
+                    if i > load_graph_idx and inst.opname == call_op
                 )
                 # pre-graph should alias: inputs_ref_0 = inputs[0]
                 matches = [
@@ -1208,8 +1241,10 @@ def model(x):
 
                 result = model(x).sum()
                 gy, gz = torch.autograd.grad(result, inputs=[y, z])
-                assert y.grad is None
-                assert z.grad is None
+                if y.grad is not None:
+                    raise AssertionError
+                if z.grad is not None:
+                    raise AssertionError
                 yield gy
                 yield gz
 
@@ -1237,8 +1272,10 @@ def model(x):
 
                 result = model(x).sum()
                 gz = torch.autograd.grad(result, inputs=[z])
-                assert y.grad is None
-                assert z.grad is None
+                if y.grad is not None:
+                    raise AssertionError
+                if z.grad is not None:
+                    raise AssertionError
                 yield gz
 
         self.check_output_and_recompiles(fn)
@@ -1255,9 +1292,12 @@ def model(x):
                 x = torch.randn([1, 4])
                 result = model(x).sum()
                 out = result.backward()
-                assert out is None
-                assert y.grad is not None
-                assert z.grad is not None
+                if out is not None:
+                    raise AssertionError
+                if y.grad is None:
+                    raise AssertionError
+                if z.grad is None:
+                    raise AssertionError
                 yield y.grad
                 yield z.grad
                 y.grad = None
@@ -1287,10 +1327,14 @@ def model(x):
                 x = torch.randn([1, 4])
                 result = model(x).sum()
                 out = result.backward(inputs=[y, z])
-                assert out is None
-                assert y.grad is not None
-                assert z.grad is not None
-                assert a.grad is None
+                if out is not None:
+                    raise AssertionError
+                if y.grad is None:
+                    raise AssertionError
+                if z.grad is None:
+                    raise AssertionError
+                if a.grad is not None:
+                    raise AssertionError
                 yield y.grad
                 yield z.grad
                 y.grad = None
@@ -1496,8 +1540,10 @@ def fn():
                 x = torch.randn([10, 4])
                 result = opt_model(x).sum()
                 result.backward()
-                assert model[0].weight.grad is w_grad
-                assert model[0].bias.grad is b_grad
+                if model[0].weight.grad is not w_grad:
+                    raise AssertionError
+                if model[0].bias.grad is not b_grad:
+                    raise AssertionError
                 yield w_grad.clone()
                 yield b_grad.clone()
 
@@ -1820,9 +1866,10 @@ def test_custom_fn_output_metadata(self):
         def my_compiler_fn(gm):
             for node in gm.graph.nodes:
                 if isinstance(node.target, torch._ops.OpOverload):
-                    assert node.target._name != "aten::_to_copy", (
-                        "there should be no implicit copies (e.g. dtype casting)"
-                    )
+                    if node.target._name == "aten::_to_copy":
+                        raise AssertionError(
+                            "there should be no implicit copies (e.g. dtype casting)"
+                        )
 
             def inner_compiler(gm_, example_inputs_):
                 counters["compiled_autograd"]["compiles"] += 1
@@ -1934,7 +1981,7 @@ def backward(ctx, gO):
                 yield x.grad
 
         self.check_output_and_recompiles(
-            fn, count=[1, 3], compiler_fn=make_compiler_fn(fullgraph=False)
+            fn, count=[1, 2], compiler_fn=make_compiler_fn(fullgraph=False)
         )
 
     def test_custom_fn_compiled_fw_graph_break(self):
@@ -1988,9 +2035,9 @@ def backward(ctx, gO):
                 yield x.grad
 
         self.check_output_and_recompiles(
-            fn, count=[1, 3], compiler_fn=make_compiler_fn(fullgraph=False)
+            fn, count=[1, 2], compiler_fn=make_compiler_fn(fullgraph=False)
         )
-        self.assertEqual(counters["stats"]["unique_graphs"], 6)  # 3 fw, 3 bw
+        self.assertEqual(counters["stats"]["unique_graphs"], 5)
 
     def test_mismatch_fake_tensor_mode(self, dynamic_shape=False):
         """
@@ -2888,12 +2935,12 @@ def test_free_activation_memory(self):
 
 def main():
     device_interface = get_interface_for_device(GPU_TYPE)
-    assert(device_interface.memory_allocated() == 0)
+    assert(device_interface.memory_allocated() == 0)  # noqa: S101
 
     # Use an op to check that the memory is freed by the time the op is executed
     def assertion_impl(to_clone):
         mem_allocated = device_interface.memory_allocated()
-        assert mem_allocated < 4000000  # some activations should be freed
+        assert mem_allocated < 4000000  # noqa: S101  # some activations should be freed
         return to_clone.clone()
 
     with torch.library._scoped_library("test_compiled_autograd", "FRAGMENT") as lib:
@@ -2916,10 +2963,10 @@ def forward(activations):
 
         # allocate at least 4,000,000 bytes (1,000,000 * 4 bytes)
         activations = [torch.ones(1000000, dtype=torch.float32, device=GPU_TYPE)]
-        assert device_interface.memory_allocated() > 4000000
+        assert device_interface.memory_allocated() > 4000000  # noqa: S101
 
         out = compiled_fn(activations)
-        assert len(activations) == 0
+        assert len(activations) == 0  # noqa: S101
 
 main()
         """
@@ -2936,13 +2983,13 @@ def test_free_activation_memory_subclass(self):
 
 def main():
     device_interface = get_interface_for_device(GPU_TYPE)
-    assert device_interface.memory_allocated() == 0
+    assert device_interface.memory_allocated() == 0  # noqa: S101
 
     # Use an op to check that the memory is freed by the time the op is executed
     def assertion_impl(to_clone):
         mem_allocated = device_interface.memory_allocated()
-        assert mem_allocated < 1200000  # some activations should be freed
-        assert mem_allocated > 800000  # currently subclasses don't seem to be freed in inductor
+        assert mem_allocated < 1200000  # noqa: S101  # some activations should be freed
+        assert mem_allocated > 800000  # noqa: S101  # currently subclasses don't seem to be freed in inductor
         return to_clone.clone()
 
     with torch.library._scoped_library("test_compiled_autograd", "FRAGMENT") as lib:
@@ -2978,10 +3025,10 @@ def fn(inputs):
             torch.ones((1, 100000), device=GPU_TYPE),  # 400,000 bytes
         ]
         # 1,200,000 bytes (3 * 4 * 100,000 bytes)
-        assert device_interface.memory_allocated() > 1200000
+        assert device_interface.memory_allocated() > 1200000  # noqa: S101
 
         out = compiled_fn(activations)
-        assert len(activations) == 0
+        assert len(activations) == 0  # noqa: S101
 
 main()
         """
@@ -3208,11 +3255,13 @@ def test_logs(self):
 
         self.assertEqual(counters["compiled_autograd"]["captures"], 1)
         self.assertEqual(counters["compiled_autograd"]["compiles"], 1)
-        assert "torch::autograd::AccumulateGrad (NodeCall" in logs.getvalue()
-        assert (
+        if "torch::autograd::AccumulateGrad (NodeCall" not in logs.getvalue():
+            raise AssertionError
+        if (
             self.gen_cache_miss_log_prefix() + "torch::autograd::GraphRoot"
-            not in logs.getvalue()
-        )
+            in logs.getvalue()
+        ):
+            raise AssertionError
 
     def test_logs_aot_bwd_reuse(self):
         @torch.compile(backend="aot_eager")
@@ -3286,7 +3335,6 @@ def fn():
     @mock.patch(
         "torch._functorch.aot_autograd.AOT_COUNTER", new_callable=itertools.count
     )
-    @mock.patch("torch._dynamo.config.inline_inbuilt_nn_modules", True)
     def test_verbose_logs_aot_id(self, _):
         def fn():
             model = torch.nn.Sequential(
@@ -3454,9 +3502,10 @@ def fn():
         pattern1 = r"".join(patterns1)
         matches1 = re.findall(pattern1, all_logs)
         self.assertEqual(len(matches1), 1)
-        assert isinstance(
+        if not isinstance(
             matches1[0], str
-        )  # for a single match: matches1=['match'], for multiple matches: matches1=[('match1', 'match2')]...
+        ):  # for a single match: matches1=['match'], for multiple matches: matches1=[('match1', 'match2')]...
+            raise AssertionError
         self.assertEqual(len(matches1), len(patterns1))
 
     @skipIfWindows(msg="node name demangling inconsistent on windows")
@@ -3545,7 +3594,8 @@ def from_twotensor_backward(ctx, grad_a, grad_b):
                 "mylib::to_twotensor", TwoTensorMode, lib=lib
             )
             def _(_0, _1, _2, args, kwargs):
-                assert not kwargs
+                if kwargs:
+                    raise AssertionError
                 a, b = args
                 return TwoTensor(a.clone(), b.clone())
 
@@ -3553,7 +3603,8 @@ def _(_0, _1, _2, args, kwargs):
                 "mylib::from_twotensor", TwoTensor, lib=lib
             )
             def _(_0, _1, _2, args, kwargs):
-                assert not kwargs
+                if kwargs:
+                    raise AssertionError
                 (c,) = args
                 return c.a.clone(), c.b.clone()
 
@@ -4345,7 +4396,10 @@ def fn(fwd_compiler, ca_compiler):
                     third = torch.autograd.grad(second, x, create_graph=True)[0]
                     fourth = torch.autograd.grad(third, x, create_graph=True)[0]
             except RuntimeError as e:
-                assert "does not currently support higher order gradients" in str(e)
+                if "does not currently support higher order gradients" not in str(e):
+                    raise AssertionError(
+                        f"Expected error about higher order gradients, got: {e}"
+                    ) from e
                 return (first, second, third, fourth)
 
             return (first, second, third, fourth)
@@ -4448,7 +4502,8 @@ def backward(ctx, grad_output):
 
             def check(grad):
                 nonlocal pre_hook_storage_id
-                assert pre_hook_storage_id is None
+                if pre_hook_storage_id is not None:
+                    raise AssertionError
                 pre_hook_storage_id = id(grad.untyped_storage())
 
             var = torch.randn(2, 2, requires_grad=True)
@@ -4456,16 +4511,14 @@ def check(grad):
             output = StealableDenseOp.apply(var)
             output.backward(torch.ones_like(output))
 
-            assert var.grad is not None, "Grad should be defined"
-            assert torch.equal(var.grad, torch.ones_like(var) * 5), (
-                "Grad content should be as returned by backward"
-            )
-            assert var.grad.requires_grad is False, (
-                "Detached grad should not require grad"
-            )
-            assert id(var.grad.untyped_storage()) == pre_hook_storage_id, (
-                "Should be stolen"
-            )
+            if var.grad is None:
+                raise AssertionError("Grad should be defined")
+            if not torch.equal(var.grad, torch.ones_like(var) * 5):
+                raise AssertionError("Grad content should be as returned by backward")
+            if var.grad.requires_grad is not False:
+                raise AssertionError("Detached grad should not require grad")
+            if id(var.grad.untyped_storage()) != pre_hook_storage_id:
+                raise AssertionError("Should be stolen")
             yield var.grad
 
         self.check_output_and_recompiles(
@@ -4498,7 +4551,8 @@ def backward(ctx, grad_output):
 
             def check(grad):
                 nonlocal pre_hook_storages_id
-                assert pre_hook_storages_id is None
+                if pre_hook_storages_id is not None:
+                    raise AssertionError
                 pre_hook_storages_id = [
                     id(grad._indices().untyped_storage()),
                     id(grad._values().untyped_storage()),
@@ -4509,21 +4563,19 @@ def check(grad):
             output = StealableSparseOp.apply(var)
             output.backward(torch.ones_like(output))
 
-            assert var.grad is not None, "Grad should be defined"
-            assert var.grad.is_sparse, "Grad should be sparse"
+            if var.grad is None:
+                raise AssertionError("Grad should be defined")
+            if not var.grad.is_sparse:
+                raise AssertionError("Grad should be sparse")
             expected_dense_grad = torch.tensor([[5.0, 0.0], [0.0, 5.0]])
-            assert torch.equal(var.grad.to_dense(), expected_dense_grad), (
-                "Content should be equal after shallow copy"
-            )
-            assert var.grad.requires_grad is False, (
-                "Detached grad should not require grad"
-            )
-            assert (
-                id(var.grad._indices().untyped_storage()) == pre_hook_storages_id[0]
-            ), "Should be stolen"
-            assert (
-                id(var.grad._values().untyped_storage()) == pre_hook_storages_id[1]
-            ), "Should be stolen"
+            if not torch.equal(var.grad.to_dense(), expected_dense_grad):
+                raise AssertionError("Content should be equal after shallow copy")
+            if var.grad.requires_grad is not False:
+                raise AssertionError("Detached grad should not require grad")
+            if id(var.grad._indices().untyped_storage()) != pre_hook_storages_id[0]:
+                raise AssertionError("Should be stolen")
+            if id(var.grad._values().untyped_storage()) != pre_hook_storages_id[1]:
+                raise AssertionError("Should be stolen")
             yield var.grad
 
         self.check_output_and_recompiles(
@@ -4554,7 +4606,8 @@ def backward(ctx, grad_output):
 
             def check(grad):
                 nonlocal pre_hook_storages_id
-                assert pre_hook_storages_id is None
+                if pre_hook_storages_id is not None:
+                    raise AssertionError
                 pre_hook_storages_id = [
                     id(grad._indices().untyped_storage()),
                     id(grad._values().untyped_storage()),
@@ -4567,27 +4620,25 @@ def check(grad):
                 torch.ones_like(output), create_graph=True
             )  # grad mode == create_graph
 
-            assert var.grad is not None, "Grad should be defined"
-            assert var.grad.is_sparse, "Grad should be sparse"
+            if var.grad is None:
+                raise AssertionError("Grad should be defined")
+            if not var.grad.is_sparse:
+                raise AssertionError("Grad should be sparse")
             expected_dense_grad = torch.tensor([[5.0, 0.0], [0.0, 5.0]])
-            assert torch.equal(var.grad.to_dense(), expected_dense_grad), (
-                "Content should be equal after clone"
-            )
-            assert var.grad.requires_grad, (
-                "Grad should require grad for double backward"
-            )
-            assert (
-                id(var.grad._indices().untyped_storage()) != pre_hook_storages_id[0]
-            ), "Should be copied"
-            assert (
-                id(var.grad._values().untyped_storage()) != pre_hook_storages_id[1]
-            ), "Should be copied"
+            if not torch.equal(var.grad.to_dense(), expected_dense_grad):
+                raise AssertionError("Content should be equal after clone")
+            if not var.grad.requires_grad:
+                raise AssertionError("Grad should require grad for double backward")
+            if id(var.grad._indices().untyped_storage()) == pre_hook_storages_id[0]:
+                raise AssertionError("Should be copied")
+            if id(var.grad._values().untyped_storage()) == pre_hook_storages_id[1]:
+                raise AssertionError("Should be copied")
             yield var.grad
 
         self.check_output_and_recompiles(
             fn,
             compiler_fn=make_compiler_fn(fullgraph=False),
-            count=[1, 2],
+            count=[1, 1],
         )
 
     # Case 1.5.1: Dense variable gradient layout contract
@@ -4619,13 +4670,12 @@ def check(grad):
             output = NotStealableRefsOp.apply(var)
             output.backward(torch.ones_like(output))
 
-            assert var.grad is not None, "Grad should be defined"
-            assert torch.equal(var.grad, torch.ones_like(var) * 10.0), (
-                "Grad content should be as returned by backward"
-            )
-            assert (
-                grad_ref_holder[0].untyped_storage() is not var.grad.untyped_storage()
-            ), "Should be copied"
+            if var.grad is None:
+                raise AssertionError("Grad should be defined")
+            if not torch.equal(var.grad, torch.ones_like(var) * 10.0):
+                raise AssertionError("Grad content should be as returned by backward")
+            if grad_ref_holder[0].untyped_storage() is var.grad.untyped_storage():
+                raise AssertionError("Should be copied")
             yield var.grad
 
         self.check_output_and_recompiles(fn)
@@ -4650,9 +4700,8 @@ def backward(ctx, grad_output):
             # Create a non-contiguous variable
             base_tensor = torch.randn(4, 4)
             var = base_tensor[::2, ::2]
-            assert not var.is_contiguous(), (
-                "Variable should be non-contiguous for this test"
-            )
+            if var.is_contiguous():
+                raise AssertionError("Variable should be non-contiguous for this test")
             var.requires_grad_(True)
 
             grad_ref_holder = [None]
@@ -4666,18 +4715,18 @@ def check(grad):
             output = SimpleDenseGradOp.apply(var)
             output.backward(torch.ones_like(output))
 
-            assert var.grad is not None, "Grad should be defined"
+            if var.grad is None:
+                raise AssertionError("Grad should be defined")
             # The `clone_obey_contract` branch 2 (`new_grad.clone(at::MemoryFormat::Contiguous)`)
             # will make the resulting grad contiguous.
-            assert var.grad.is_contiguous(), (
-                "Resulting grad should be contiguous due to branch 2 of clone_obey_contract"
-            )
-            assert torch.equal(var.grad, torch.ones_like(var) * 7.0), (
-                "Grad content should be as returned by backward"
-            )
-            assert (
-                grad_ref_holder[0].untyped_storage() is not var.grad.untyped_storage()
-            ), "Should be copied"
+            if not var.grad.is_contiguous():
+                raise AssertionError(
+                    "Resulting grad should be contiguous due to branch 2 of clone_obey_contract"
+                )
+            if not torch.equal(var.grad, torch.ones_like(var) * 7.0):
+                raise AssertionError("Grad content should be as returned by backward")
+            if grad_ref_holder[0].untyped_storage() is var.grad.untyped_storage():
+                raise AssertionError("Should be copied")
             yield var.grad
 
         self.check_output_and_recompiles(
@@ -4707,12 +4756,14 @@ def backward(ctx, grad_output):
             expected_sum = (torch.ones_like(var) * 3.0) + initial_grad_ref.to_dense()
             output.backward(torch.ones_like(output))
 
-            assert var.grad is not None, "Grad should be defined"
-            assert not var.grad.is_sparse, "Resulting grad should be dense"
-            assert torch.equal(var.grad, expected_sum), "Grad content should be the sum"
-            assert var.grad is not initial_grad_ref, (
-                "Grad object should be replaced (out-of-place)"
-            )
+            if var.grad is None:
+                raise AssertionError("Grad should be defined")
+            if var.grad.is_sparse:
+                raise AssertionError("Resulting grad should be dense")
+            if not torch.equal(var.grad, expected_sum):
+                raise AssertionError("Grad content should be the sum")
+            if var.grad is initial_grad_ref:
+                raise AssertionError("Grad object should be replaced (out-of-place)")
             yield var.grad
 
         self.check_output_and_recompiles(
@@ -4740,12 +4791,16 @@ def backward(ctx, grad_output):
             expected_sum = initial_grad_ref + (torch.ones_like(var) * 3.0)
             output.backward(torch.ones_like(output))
 
-            assert var.grad is not None, "Grad should be defined"
-            assert not var.grad.is_sparse, "Resulting grad should be dense"
-            assert torch.equal(var.grad, expected_sum), "Grad content should be the sum"
-            assert var.grad is initial_grad_ref, (
-                "Grad object should be modified in-place (same object)"
-            )
+            if var.grad is None:
+                raise AssertionError("Grad should be defined")
+            if var.grad.is_sparse:
+                raise AssertionError("Resulting grad should be dense")
+            if not torch.equal(var.grad, expected_sum):
+                raise AssertionError("Grad content should be the sum")
+            if var.grad is not initial_grad_ref:
+                raise AssertionError(
+                    "Grad object should be modified in-place (same object)"
+                )
             yield var.grad
 
         self.check_output_and_recompiles(fn)
@@ -4788,14 +4843,16 @@ def backward(ctx, grad_output):
 
             output.backward(torch.ones_like(output))
 
-            assert var.grad is not None, "Grad should be defined"
-            assert var.grad.is_sparse, "Resulting grad should remain sparse"
-            assert torch.equal(var.grad.to_dense(), expected_sum_dense), (
-                "Grad content should be the sum of sparse grads"
-            )
-            assert var.grad is initial_grad_ref, (
-                "Grad object should be modified in-place (same object)"
-            )
+            if var.grad is None:
+                raise AssertionError("Grad should be defined")
+            if not var.grad.is_sparse:
+                raise AssertionError("Resulting grad should remain sparse")
+            if not torch.equal(var.grad.to_dense(), expected_sum_dense):
+                raise AssertionError("Grad content should be the sum of sparse grads")
+            if var.grad is not initial_grad_ref:
+                raise AssertionError(
+                    "Grad object should be modified in-place (same object)"
+                )
             yield var.grad
 
         self.check_output_and_recompiles(
@@ -4835,18 +4892,22 @@ def backward(ctx, grad_output):
 
             output.backward(torch.ones_like(output))
 
-            assert var.grad is not None, "Grad should be defined"
-            assert not var.grad.is_sparse, "Resulting grad should be dense"
-            assert torch.equal(var.grad, expected_sum), "Grad content should be the sum"
-            assert var.grad is initial_grad_ref, (
-                "Grad object should be modified in-place (same object)"
-            )
+            if var.grad is None:
+                raise AssertionError("Grad should be defined")
+            if var.grad.is_sparse:
+                raise AssertionError("Resulting grad should be dense")
+            if not torch.equal(var.grad, expected_sum):
+                raise AssertionError("Grad content should be the sum")
+            if var.grad is not initial_grad_ref:
+                raise AssertionError(
+                    "Grad object should be modified in-place (same object)"
+                )
             yield var.grad
 
         self.check_output_and_recompiles(
             fn,
             compiler_fn=make_compiler_fn(fullgraph=False),
-            count=[1, 2],
+            count=[1, 1],
         )
 
     # Case 3.1: Sparse variable_grad + Dense new_grad (reorder into Dense + Sparse)
@@ -4879,15 +4940,18 @@ def backward(ctx, grad_output):
 
             output.backward(torch.ones_like(output), create_graph=True)
 
-            assert var.grad is not None, "Grad should be defined"
-            assert not var.grad.is_sparse, "Resulting grad should be dense"
-            assert torch.equal(var.grad, expected_sum), "Grad content should be the sum"
-            assert var.grad is not initial_grad_ref, (
-                "Grad object should be replaced (out-of-place)"
-            )
-            assert var.grad.requires_grad, (
-                "Resulting grad should track history for double backward"
-            )
+            if var.grad is None:
+                raise AssertionError("Grad should be defined")
+            if var.grad.is_sparse:
+                raise AssertionError("Resulting grad should be dense")
+            if not torch.equal(var.grad, expected_sum):
+                raise AssertionError("Grad content should be the sum")
+            if var.grad is initial_grad_ref:
+                raise AssertionError("Grad object should be replaced (out-of-place)")
+            if not var.grad.requires_grad:
+                raise AssertionError(
+                    "Resulting grad should track history for double backward"
+                )
             yield var.grad
 
         self.check_output_and_recompiles(
@@ -4923,21 +4987,24 @@ def backward(ctx, grad_output):
 
             output.backward(torch.ones_like(output), create_graph=True)
 
-            assert var.grad is not None, "Grad should be defined"
-            assert not var.grad.is_sparse, "Resulting grad should be dense"
-            assert torch.equal(var.grad, expected_sum), "Grad content should be the sum"
-            assert var.grad is not initial_grad_ref, (
-                "Grad object should be replaced (out-of-place)"
-            )
-            assert var.grad.requires_grad, (
-                "Resulting grad should track history for double backward"
-            )
+            if var.grad is None:
+                raise AssertionError("Grad should be defined")
+            if var.grad.is_sparse:
+                raise AssertionError("Resulting grad should be dense")
+            if not torch.equal(var.grad, expected_sum):
+                raise AssertionError("Grad content should be the sum")
+            if var.grad is initial_grad_ref:
+                raise AssertionError("Grad object should be replaced (out-of-place)")
+            if not var.grad.requires_grad:
+                raise AssertionError(
+                    "Resulting grad should track history for double backward"
+                )
             yield var.grad
 
         self.check_output_and_recompiles(
             fn,
             compiler_fn=make_compiler_fn(fullgraph=False),
-            count=[1, 3],
+            count=[1, 2],
         )
 
     def test_torch_function_mode(self):
@@ -5077,7 +5144,8 @@ def lookup_backend(test_name):
     elif test_name in xfail_by_backend["eager"]:
         return "ca_eager"
     else:
-        assert test_name not in xfail_by_backend["ca_eager"]
+        if test_name in xfail_by_backend["ca_eager"]:
+            raise AssertionError(f"test {test_name} is in ca_eager xfail list")
         return "inductor"
 
 
@@ -5224,6 +5292,7 @@ def tearDown(self):
     "test_nested_checkpoint_same_graph_early_stop_False",  # dynamo disable
     "test_nested_checkpoint_same_graph_early_stop_True",  # dynamo disable
     "test_nested_checkpoint_set_early_stop",  # dynamo disable
+    "test_nested_checkpoint_set_early_stop_no_recompution_needed",  # TorchDispatchMode causes frame skip
     "test_nested_checkpoint_two_children_early_stop_False",  # dynamo disable
     "test_nested_checkpoint_two_children_early_stop_True",  # dynamo disable
     "test_custom_autograd_ac_early_stop",  # marked as skipped
@@ -5365,17 +5434,21 @@ def tearDown(self):
 # clear_saved_tensors_on_access is incompatible with compiled autograd
 skipped_tests.add("test_clear_saved_tensors_on_access")
 skipped_tests.add("test_clear_saved_tensors_on_access_double_access_error")
+skipped_tests.add("test_forward_traceback_preserves_exception_with_checkpoint")
+skipped_tests.add("test_checkpoint_error_suggests_mark_dynamic")
+skipped_tests.add("test_checkpoint_automatic_dynamic_graph_shadowing")
+skipped_tests.add("test_checkpoint_automatic_dynamic_mark_dynamic_workaround")
+skipped_tests.add("test_checkpoint_automatic_dynamic_lru_disabled_workaround")
 
 test_autograd = load_test_module("test_autograd")
 test_custom_ops = load_test_module("test_custom_ops")
 test_higher_order_ops = load_test_module("dynamo/test_higher_order_ops")
-if not HAS_XPU_AND_TRITON:
-    TestAutogradWithCompiledAutograd = wrap_test_class(test_autograd.TestAutograd)
+
+TestAutogradWithCompiledAutograd = wrap_test_class(test_autograd.TestAutograd)
 TestNestedCheckpointWithCompiledAutograd = wrap_test_class(
     test_autograd.TestNestedCheckpoint
 )
-if not HAS_XPU_AND_TRITON:
-    TestCustomOpWithCompiledAutograd = wrap_test_class(test_custom_ops.TestCustomOp)
+TestCustomOpWithCompiledAutograd = wrap_test_class(test_custom_ops.TestCustomOp)
 HigherOrderOpTestsWithCompiledAutograd = wrap_test_class(
     test_higher_order_ops.HigherOrderOpTests
 )
@@ -5386,7 +5459,7 @@ def tearDown(self):
     test_higher_order_ops.ActivationCheckpointingTests
 )
 
-if torch.distributed.is_available() and HAS_CUDA_AND_TRITON:
+if torch.distributed.is_available() and HAS_GPU_AND_TRITON:
     test_dtensor = load_test_module("distributed/tensor/test_dtensor_compile")
     TestDTensorCompileWithCompiledAutograd = wrap_test_class(
         test_dtensor.TestDTensorCompile
diff --git a/test/inductor/test_compiled_optimizers.py b/test/inductor/test_compiled_optimizers.py
index 595d1ccb3da66..197e3148dc419 100644
--- a/test/inductor/test_compiled_optimizers.py
+++ b/test/inductor/test_compiled_optimizers.py
@@ -58,18 +58,19 @@
     optim_db,
     optims,
 )
-from torch.testing._internal.common_utils import parametrize, skipIfWindows
+from torch.testing._internal.common_utils import (
+    parametrize,
+    skipIfRocm,
+    skipIfWindows,
+    skipIfXpu,
+)
 from torch.testing._internal.inductor_utils import (
     GPU_TYPE,
     HAS_CPU,
     HAS_GPU,
     has_triton,
 )
-from torch.testing._internal.triton_utils import (
-    requires_cuda_and_triton,
-    requires_gpu,
-    requires_gpu_and_triton,
-)
+from torch.testing._internal.triton_utils import requires_gpu, requires_gpu_and_triton
 
 
 def get_inputs(optim):
@@ -591,7 +592,7 @@ class CompiledOptimizerParityTests(TestCase):
     @optims(optim_db, dtypes=[torch.float32])
     @parametrize("use_closure", [True, False])
     def test_correctness(self, device, dtype, optim_info, use_closure):
-        torch.cuda.manual_seed_all(0)
+        torch.get_device_module(device).manual_seed_all(0)
         torch.manual_seed(0)
         random.seed(0)
         optim_cls = optim_info.optim_cls
@@ -931,7 +932,7 @@ def fn(xs, ys):
 
         self.assertLess(end - start, 90)
 
-    @requires_cuda_and_triton
+    @requires_gpu_and_triton
     def test_S429861(self):
         # Just verify we can compile this function without error
         try:
@@ -947,7 +948,7 @@ def test_S429861(self):
         from torch._inductor.utils import fresh_cache
 
         with fresh_cache():
-            kwargs = aot_graph_input_parser(forward)
+            kwargs = aot_graph_input_parser(forward, device=GPU_TYPE)
             torch.compile(forward)(**kwargs)
 
     @requires_gpu_and_triton
@@ -991,6 +992,90 @@ def loop():
             self.assertEqual(param, param_ref)
 
 
+@skipIfRocm(msg="ROCm may have different numerical behavior")
+@requires_gpu_and_triton
+class CompiledOptimizerBitwiseTests(TestCase):
+    """
+    Tests that compiled optimizers produce bitwise identical results to eager
+    when precision configs are enabled.
+
+    These tests verify that with the following config options:
+    - eager_numerics.division_rounding = True
+    - eager_numerics.use_pytorch_libdevice = True
+    - emulate_precision_casts = True
+
+    The compiled optimizer step produces results that are bitwise identical
+    to the eager optimizer step.
+    """
+
+    @staticmethod
+    def _test_optimizer_bitwise(
+        test_case,
+        optim_cls,
+        kernel_count=None,
+        num_steps=10,
+        **optim_kwargs,
+    ):
+        """Helper to test optimizer bitwise equality."""
+        torch._dynamo.reset()
+        torch._inductor.metrics.reset()
+        torch.manual_seed(42)
+
+        input = torch.ones([10, 10], device=GPU_TYPE)
+        model_eager = torch.nn.Sequential(
+            *[torch.nn.Linear(10, 10, device=GPU_TYPE) for _ in range(2)]
+        )
+        model_eager(input).sum().backward()
+
+        model_compiled = deepcopy(model_eager)
+        model_compiled(input).sum().backward()
+
+        opt_eager = optim_cls(model_eager.parameters(), **optim_kwargs)
+        opt_compiled = optim_cls(model_compiled.parameters(), **optim_kwargs)
+        compiled_step = compile_opt(opt_compiled)
+
+        with torch.set_grad_enabled(False):
+            for step in range(num_steps):
+                compiled_step()
+                opt_eager.step()
+
+                # Check bitwise equality
+                for i, (p_eager, p_compiled) in enumerate(
+                    zip(model_eager.parameters(), model_compiled.parameters())
+                ):
+                    test_case.assertEqual(
+                        p_eager,
+                        p_compiled,
+                        atol=0,
+                        rtol=0,
+                        msg=f"Step {step + 1}, param {i}: params differ",
+                    )
+
+        # Also check optimizer state
+        for p_eager, p_compiled in zip(
+            model_eager.parameters(), model_compiled.parameters()
+        ):
+            for key in opt_eager.state[p_eager]:
+                eager_val = opt_eager.state[p_eager][key]
+                compiled_val = opt_compiled.state[p_compiled][key]
+                if isinstance(eager_val, torch.Tensor):
+                    test_case.assertEqual(
+                        eager_val,
+                        compiled_val,
+                        atol=0,
+                        rtol=0,
+                        msg=f"State '{key}' differs",
+                    )
+
+        if kernel_count is not None and test_case.check_kernel_count:
+            if isinstance(kernel_count, types.LambdaType):
+                kernel_count(str(torch._inductor.metrics.generated_kernel_count))
+            else:
+                test_case.assertEqual(
+                    torch._inductor.metrics.generated_kernel_count, kernel_count
+                )
+
+
 for optim_cls, name, kwargs, scheduler_cls in COMPILED_OPT_KWARG_DB:
     setattr(
         CompiledOptimizerTests,
@@ -998,6 +1083,76 @@ def loop():
         make_test(optim_cls, scheduler_cls=scheduler_cls, **kwargs),
     )
 
+
+def _make_bitwise_test(optim_cls, kernel_count=None, **optim_kwargs):
+    @skipIfRocm(msg="ROCm may have different numerical behavior")
+    @skipIfXpu(msg="AttributeError, torch-xpu-ops: #2999")
+    @requires_gpu_and_triton
+    @config.patch(
+        {
+            "score_fusion_memory_threshold": 1,
+            "eager_numerics.division_rounding": True,
+            "eager_numerics.use_pytorch_libdevice": True,
+            "emulate_precision_casts": True,
+        }
+    )
+    def test_fn(self):
+        CompiledOptimizerBitwiseTests._test_optimizer_bitwise(
+            self, optim_cls, kernel_count=kernel_count, **optim_kwargs
+        )
+
+    return test_fn
+
+
+_BITWISE_CAPTURABLE_OPTIMS = (
+    Adam,
+    AdamW,
+    Adadelta,
+    Adamax,
+    ASGD,
+    NAdam,
+    RAdam,
+    RMSprop,
+    Rprop,
+)
+# SGD doesn't support capturable but has no item() calls
+# so it compiles without graph breaks and can be tested bitwise.
+_BITWISE_NON_CAPTURABLE_OPTIMS = (SGD,)
+
+for optim_cls, name, kwargs, scheduler_cls in COMPILED_OPT_KWARG_DB:
+    if (
+        kwargs.get("device") == GPU_TYPE
+        and "tensor_lr" not in name
+        and scheduler_cls is None
+        and (
+            (
+                optim_cls in _BITWISE_CAPTURABLE_OPTIMS
+                and kwargs.get("capturable", False)
+            )
+            or optim_cls in _BITWISE_NON_CAPTURABLE_OPTIMS
+        )
+    ):
+        bitwise_name = name.replace("test_", "test_bitwise_")
+        # Use the same kernel count as the non-bitwise test, including
+        # any overrides for specific test configurations.
+        if name in KERNEL_COUNT_OVERRIDES:
+            kernel_count = KERNEL_COUNT_OVERRIDES[name]
+        else:
+            kernel_count = (
+                KERNEL_COUNTS[optim_cls].multitensor
+                if kwargs.get("foreach", False)
+                else KERNEL_COUNTS[optim_cls].singletensor
+            )
+        optim_kwargs = {
+            k: v for k, v in kwargs.items() if k not in ("device", "kernel_count")
+        }
+        setattr(
+            CompiledOptimizerTests,
+            bitwise_name,
+            _make_bitwise_test(optim_cls, kernel_count=kernel_count, **optim_kwargs),
+        )
+
+
 instantiate_device_type_tests(
     CompiledOptimizerParityTests, globals(), allow_xpu=True, except_for="cpu"
 )
diff --git a/test/inductor/test_config.py b/test/inductor/test_config.py
index f617a4d12f11f..e78723e2df10b 100644
--- a/test/inductor/test_config.py
+++ b/test/inductor/test_config.py
@@ -7,7 +7,8 @@
 from torch._inductor import config
 from torch._inductor.pattern_matcher import PatternMatcherPass
 from torch._inductor.test_case import run_tests, TestCase
-from torch.testing._internal.inductor_utils import HAS_CPU, HAS_TRITON
+from torch.testing._internal.inductor_utils import GPU_TYPE, HAS_CPU, HAS_TRITON
+from torch.testing._internal.triton_utils import requires_gpu
 
 
 def dummy_fn(x):
@@ -327,6 +328,68 @@ def fn2(x, y):
         self.assertEqual(counters["inductor"]["fxgraph_cache_hit"], 0)
         self.assertEqual(counters["inductor"]["fxgraph_lookup_write_file"], 0)
 
+    @requires_gpu
+    @torch._inductor.config.patch(fx_graph_cache=False)
+    def test_config_read_in_backwards(self):
+        @torch.compile
+        def f(x, y):
+            z = x @ y
+            return z.sin().sum()
+
+        called = False
+
+        def my_pass(graph):
+            nonlocal called
+            called = True
+
+        x, y = (
+            torch.randn(3, 3, device=GPU_TYPE, requires_grad=True),
+            torch.randn(3, 3, device=GPU_TYPE),
+        )
+        z = f(x, y)
+        z.backward()
+        self.assertFalse(called)
+        torch._dynamo.reset()
+        z = f(x, y)
+        with torch._inductor.config.patch(post_grad_custom_pre_pass=my_pass):
+            z.backward()
+
+        self.assertTrue(called)
+
+        called = False
+        torch._dynamo.reset()
+        z = f(x, y)
+        with torch._inductor.config.patch(post_grad_custom_pre_pass=my_pass):
+            torch.autograd.grad(z, x)
+        self.assertTrue(called)
+
+    @torch._inductor.config.patch(fx_graph_cache=False)
+    def test_config_read_in_grad_fn(self):
+        @torch.compile
+        def f(x, y):
+            z = x @ y
+            return z.sin().sum()
+
+        called = False
+
+        def my_pass(graph):
+            nonlocal called
+            called = True
+
+        x, y = (
+            torch.randn(3, 3, requires_grad=True),
+            torch.randn(3, 3),
+        )
+
+        with torch._inductor.config.patch(post_grad_custom_pre_pass=my_pass):
+            z = f(x, y)
+        self.assertTrue(called)
+
+        # Make sure the context gets cleared after forward pass
+        called = False
+        z.grad_fn.apply(torch.tensor(0))
+        self.assertFalse(called)
+
 
 if __name__ == "__main__":
     run_tests()
diff --git a/test/inductor/test_control_deps.py b/test/inductor/test_control_deps.py
index 0b221f5c44ca2..9adc3458fce75 100644
--- a/test/inductor/test_control_deps.py
+++ b/test/inductor/test_control_deps.py
@@ -8,7 +8,7 @@
 from torch.testing._internal.common_utils import IS_LINUX
 from torch.testing._internal.inductor_utils import (
     GPU_TYPE,
-    HAS_CUDA_AND_TRITON,
+    HAS_GPU_AND_TRITON,
     requires_gpu,
 )
 
@@ -25,12 +25,14 @@ def fn(a, b):
 
         def add_control_deps(graph):
             nodes = [n for n in graph.nodes if n.op == "call_function"]
-            assert len(nodes) == 3
+            if len(nodes) != 3:
+                raise AssertionError(f"Expected 3 nodes, got {len(nodes)}")
             c_node = nodes[0]
             d_node = nodes[1]
             e_node = nodes[2]
 
-            assert d_node.target == torch.ops.aten.mm.default
+            if d_node.target != torch.ops.aten.mm.default:
+                raise AssertionError(f"Expected mm.default, got {d_node.target}")
 
             from torch.utils._ordered_set import OrderedSet
 
@@ -41,14 +43,24 @@ def add_control_deps(graph):
             sub_g = graph.find_nodes(
                 op="call_function", target=torch.ops.higher_order.control_deps
             )
-            assert len(sub_g) == 2
-
-            assert list(sub_g[0].meta["val"].shape) == [256, 256]
-            assert list(sub_g[1].meta["val"].shape) == [256, 256]
+            if len(sub_g) != 2:
+                raise AssertionError(f"Expected 2 control_deps nodes, got {len(sub_g)}")
+
+            if list(sub_g[0].meta["val"].shape) != [256, 256]:
+                raise AssertionError(
+                    f"Expected shape [256, 256], got {list(sub_g[0].meta['val'].shape)}"
+                )
+            if list(sub_g[1].meta["val"].shape) != [256, 256]:
+                raise AssertionError(
+                    f"Expected shape [256, 256], got {list(sub_g[1].meta['val'].shape)}"
+                )
 
             for attr in graph.find_nodes(op="get_attr"):
                 for n in getattr(graph.owning_module, attr.target).graph.nodes:
-                    assert list(n.meta["val"].shape) == [256, 256]
+                    if list(n.meta["val"].shape) != [256, 256]:
+                        raise AssertionError(
+                            f"Expected shape [256, 256], got {list(n.meta['val'].shape)}"
+                        )
 
             return graph
 
@@ -91,7 +103,8 @@ def add_control_deps(graph):
             mm_nodes = graph.find_nodes(
                 op="call_function", target=torch.ops.aten.mm.default
             )
-            assert len(mm_nodes) == 4, f"Expected 4 mm nodes, got {len(mm_nodes)}"
+            if len(mm_nodes) != 4:
+                raise AssertionError(f"Expected 4 mm nodes, got {len(mm_nodes)}")
 
             # Add control dep: mm3 depends on mm0's output
             # This should NOT extend mm0's buffer lifetime
@@ -135,13 +148,15 @@ def add_control_deps(graph):
             cat_nodes = graph.find_nodes(
                 op="call_function", target=torch.ops.aten.cat.default
             )
-            assert len(cat_nodes) == 1, f"Expected 1 cat node, got {len(cat_nodes)}"
+            if len(cat_nodes) != 1:
+                raise AssertionError(f"Expected 1 cat node, got {len(cat_nodes)}")
             cat_node = cat_nodes[0]
 
             # Verify it has nested args (list of tensors)
-            assert isinstance(cat_node.args[0], (list, tuple)), (
-                f"Expected nested args, got {type(cat_node.args[0])}"
-            )
+            if not isinstance(cat_node.args[0], (list, tuple)):
+                raise AssertionError(
+                    f"Expected nested args, got {type(cat_node.args[0])}"
+                )
 
             # Find a node that comes before cat to use as dependency
             add_nodes = graph.find_nodes(
@@ -159,7 +174,10 @@ def add_control_deps(graph):
             control_deps_nodes = graph.find_nodes(
                 op="call_function", target=torch.ops.higher_order.control_deps
             )
-            assert len(control_deps_nodes) == 1
+            if len(control_deps_nodes) != 1:
+                raise AssertionError(
+                    f"Expected 1 control_deps node, got {len(control_deps_nodes)}"
+                )
             return graph
 
         with torch._inductor.config.patch(
@@ -213,12 +231,14 @@ def add_control_deps(graph):
                 op="call_function",
                 target=torch.ops.higher_order.triton_kernel_wrapper_functional,
             )
-            assert triton_nodes
+            if not triton_nodes:
+                raise AssertionError("Expected triton_kernel_wrapper_functional nodes")
             # Find mul node (z = x * 2) to use as dependency
             mul_nodes = graph.find_nodes(
                 op="call_function", target=torch.ops.aten.mul.Tensor
             )
-            assert mul_nodes
+            if not mul_nodes:
+                raise AssertionError("Expected mul.Tensor nodes")
             deps_map = {triton_nodes[0]: OrderedSet([mul_nodes[0]])}
             torch._inductor.fx_passes.control_dependencies.preserve_node_ordering(
                 graph, deps_map
@@ -239,5 +259,5 @@ def add_control_deps(graph):
 
 
 if __name__ == "__main__":
-    if IS_LINUX and HAS_CUDA_AND_TRITON:
+    if IS_LINUX and HAS_GPU_AND_TRITON:
         run_tests(needs="filelock")
diff --git a/test/inductor/test_control_flow.py b/test/inductor/test_control_flow.py
index 4c11640c92487..680b503e3f815 100644
--- a/test/inductor/test_control_flow.py
+++ b/test/inductor/test_control_flow.py
@@ -14,7 +14,6 @@
     decorateIf,
     instantiate_parametrized_tests,
     parametrize,
-    skipIfXpu,
 )
 from torch.testing._internal.inductor_utils import GPU_TYPE, HAS_CPU, HAS_GPU
 from torch.testing._internal.triton_utils import requires_gpu
@@ -24,7 +23,8 @@ def _prepend_product_of_values(inputs, possible_values, num_to_prepend=1, device
     result = []
     if len(inputs) != 0:
         device = inputs[0].device
-    assert device
+    if not device:
+        raise AssertionError
     # iterate over the cartesian product of predicate values
     for values in itertools.product(*([possible_values] * num_to_prepend)):
         prepended = [torch.tensor(v, device=device) for v in values]
@@ -384,7 +384,6 @@ def test_cond_unbacked_symint_closure(self, device, dynamic):
             dynamic=dynamic,
         )
 
-    @skipIfXpu(msg="Remove this skip after issue #154949 resolved.")
     @requires_gpu
     def test_cond_control_flow_with_precomputed_size(self):
         class TestModel(torch.nn.Module):
diff --git a/test/inductor/test_cooperative_reductions.py b/test/inductor/test_cooperative_reductions.py
index b211e273c2a79..fa395dac4b591 100644
--- a/test/inductor/test_cooperative_reductions.py
+++ b/test/inductor/test_cooperative_reductions.py
@@ -19,7 +19,7 @@
 from torch.testing._internal.inductor_utils import GPU_TYPE, HAS_GPU
 
 
-class TestingHeuristics(InductorChoices):
+class _TestingHeuristics(InductorChoices):
     def __init__(self, *, cooperative: bool, persistent: bool, cfg: dict[str, int]):
         super().__init__()
         self.cooperative = cooperative
@@ -255,7 +255,7 @@ class MultiKernelCooperativeReductionTests(CooperativeReductionTests):
 class TestFixedConfigs(TestCase):
     def _check(self, fn, args, *, persistent=False, cooperative=True, cfg):
         expected = fn(*args)
-        heuristic = TestingHeuristics(
+        heuristic = _TestingHeuristics(
             persistent=persistent, cooperative=cooperative, cfg=cfg
         )
         with torch._inductor.virtualized.V.set_choices_handler(heuristic):
diff --git a/test/inductor/test_coordinate_descent_tuner.py b/test/inductor/test_coordinate_descent_tuner.py
index 7ad0994aeb9a3..c5b39f4491d0c 100644
--- a/test/inductor/test_coordinate_descent_tuner.py
+++ b/test/inductor/test_coordinate_descent_tuner.py
@@ -35,7 +35,8 @@ def mock_compare_config_prefer_larger_XBLOCK(
     self is the CoordescTuner object
     """
     if "XBLOCK" in candidate_config.kwargs:
-        assert "XBLOCK" in best_config.kwargs
+        if "XBLOCK" not in best_config.kwargs:
+            raise AssertionError
         if candidate_config.kwargs["XBLOCK"] < best_config.kwargs["XBLOCK"]:
             func(candidate_config)  # run func so the launcher will be created
             return False, best_timing * 1.1
diff --git a/test/inductor/test_cpp_wrapper_hipify.py b/test/inductor/test_cpp_wrapper_hipify.py
index 97522b872ff70..5a7d53960e08a 100644
--- a/test/inductor/test_cpp_wrapper_hipify.py
+++ b/test/inductor/test_cpp_wrapper_hipify.py
@@ -27,7 +27,10 @@
 
 class TestCppWrapperHipify(TestCase):
     def test_hipify_basic_declaration(self) -> None:
-        assert len(TEST_CODES) == len(HIP_CODES)
+        if len(TEST_CODES) != len(HIP_CODES):
+            raise AssertionError(
+                f"TEST_CODES length {len(TEST_CODES)} != HIP_CODES length {len(HIP_CODES)}"
+            )
         for i in range(len(TEST_CODES)):
             result = maybe_hipify_code_wrapper(TEST_CODES[i], True)
             expected = HIP_CODES[i]
@@ -118,7 +121,10 @@ def test_hipify_aoti_driver_header(self) -> None:
         self.assertEqual(result.rstrip(), expected.rstrip())
 
     def test_hipify_cross_platform(self) -> None:
-        assert len(TEST_CODES) == len(HIP_CODES)
+        if len(TEST_CODES) != len(HIP_CODES):
+            raise AssertionError(
+                f"TEST_CODES length {len(TEST_CODES)} != HIP_CODES length {len(HIP_CODES)}"
+            )
         for i in range(len(TEST_CODES)):
             hip_result = maybe_hipify_code_wrapper(TEST_CODES[i], True)
             result = maybe_hipify_code_wrapper(TEST_CODES[i])
diff --git a/test/inductor/test_cpu_cpp_wrapper.py b/test/inductor/test_cpu_cpp_wrapper.py
index c69226ff04d55..9f717654e58e8 100644
--- a/test/inductor/test_cpu_cpp_wrapper.py
+++ b/test/inductor/test_cpu_cpp_wrapper.py
@@ -104,7 +104,8 @@ def make_test_case(
         code_string_count = {}
 
     func = getattr(tests, test_name)
-    assert callable(func), "not a callable"
+    if not callable(func):
+        raise AssertionError("not a callable")
     func = slowTest(func) if slow else func
     new_test_name = f"{test_name}_separate" if test_build_separate else test_name
 
diff --git a/test/inductor/test_cpu_repro.py b/test/inductor/test_cpu_repro.py
index 576f8a00565f8..4295ade4df8d3 100644
--- a/test/inductor/test_cpu_repro.py
+++ b/test/inductor/test_cpu_repro.py
@@ -31,15 +31,19 @@
 from torch.fx.experimental.proxy_tensor import make_fx
 from torch.nn import functional as F
 from torch.testing._internal.common_utils import (
+    get_gcc_major_version,
     instantiate_parametrized_tests,
+    IS_ARM64,
+    IS_CPU_CAPABILITY_SVE256,
+    IS_CPU_EXT_SVE_SUPPORTED,
     IS_FBCODE,
     IS_MACOS,
     MI200_ARCH,
     parametrize,
-    skipIfRocm,
     skipIfRocmArch,
     slowTest,
     TEST_MKL,
+    xfailIf,
     xfailIfS390X,
 )
 from torch.utils._python_dispatch import TorchDispatchMode
@@ -82,7 +86,10 @@ def _can_check_vec_metrics():
 
 def check_metrics_vec_kernel_count(num_expected_vec_kernels):
     if _can_check_vec_metrics():
-        assert metrics.generated_cpp_vec_kernel_count == num_expected_vec_kernels
+        if metrics.generated_cpp_vec_kernel_count != num_expected_vec_kernels:
+            raise AssertionError(
+                f"Expected {num_expected_vec_kernels} vec kernels, got {metrics.generated_cpp_vec_kernel_count}"
+            )
 
 
 def simd_lengths_to_test():
@@ -145,7 +152,8 @@ def fn(x):
         self.assertEqual(len(actual), 1)
         torch.testing.assert_close(actual[0], expected[0])
 
-    @skipIfRocm
+    @torch._inductor.config.patch({"layout_optimization": True})
+    @patch("torch.cuda.is_available", lambda: False)
     def test_conv_stride_constraints(self):
         for fmt in [torch.contiguous_format, torch.channels_last]:
             # TorchDispatch doesn't work in our cuda invocation for some reason
@@ -703,6 +711,10 @@ def _test_lstm_packed(
             ]
         ),
     )
+    @unittest.skipIf(
+        IS_ARM64 and not IS_CPU_EXT_SVE_SUPPORTED,
+        "flaky on AArch64 (no SVE)",
+    )
     def test_lstm_packed(
         self,
         unbatched,
@@ -750,6 +762,10 @@ def test_lstm_packed(
         "unbatched, input_size, hidden_size, num_layers, bidirectional, bias, empty_state, batch_first, batch_size, seq_len",
         _test_lstm_packed_change_input_sizes_cpu_params,
     )
+    @unittest.skipIf(
+        IS_ARM64 and not IS_CPU_EXT_SVE_SUPPORTED,
+        "flaky on AArch64 (no SVE)",
+    )
     def test_lstm_packed_change_input_sizes_cpu(
         self,
         unbatched,
@@ -872,8 +888,14 @@ def test_pack_padded_sequence_lstm(self):
             [1, 2, 3, 4, 5, 1, 3, 2, 96, 5, 3, 1, 1, 2, 1, 2, 3, 6, 1, 2, 4, 6, 2, 1]
         )
 
-        assert sent_lens.shape[0] == batch_size
-        assert sent_lens.max().item() == max_lens
+        if sent_lens.shape[0] != batch_size:
+            raise AssertionError(
+                f"Expected sent_lens.shape[0] == {batch_size}, got {sent_lens.shape[0]}"
+            )
+        if sent_lens.max().item() != max_lens:
+            raise AssertionError(
+                f"Expected sent_lens.max().item() == {max_lens}, got {sent_lens.max().item()}"
+            )
 
         hidden_0 = hid_0.clone().requires_grad_(False)
         hidden_1 = hid_1.clone().requires_grad_(False)
@@ -1366,9 +1388,10 @@ def test_slice_scatter_default_end_value(self):
         # From HF AllenaiLongformerBase.
         def fn(query, key, window_overlap):
             batch_size, seq_len, num_heads, head_dim = query.size()
-            assert seq_len % (window_overlap * 2) == 0, (
-                f"Sequence length should be multiple of {window_overlap * 2}. Given {seq_len}"
-            )
+            if seq_len % (window_overlap * 2) != 0:
+                raise AssertionError(
+                    f"Sequence length should be multiple of {window_overlap * 2}. Given {seq_len}"
+                )
 
             chunks_count = torch.div(seq_len, window_overlap, rounding_mode="trunc") - 1
             diagonal_chunked_attention_scores = key
@@ -1432,7 +1455,8 @@ def fn(
                 ).to(dtype)
             return x
 
-        assert dtype in [torch.uint8, torch.int8]
+        if dtype not in [torch.uint8, torch.int8]:
+            raise AssertionError(f"Unexpected dtype: {dtype}")
         quant_min = 0 if dtype == torch.uint8 else -128
         quant_max = 255 if dtype == torch.uint8 else 127
 
@@ -1476,7 +1500,12 @@ def test_decomposed_dequant_relu_quant_uint8(self):
     def test_decomposed_dequant_relu_quant_int8(self):
         self._test_decomposed_dequant_relu_quant_helper(torch.int8)
 
-    def _test_dequant_quant_lowering_helper(self, dtype, dequant_out_dtype=None):
+    def _test_dequant_quant_lowering_helper(
+        self,
+        dtype,
+        input_dtype=torch.float32,
+        dequant_out_dtype=None,
+    ):
         def fn(
             x,
             scale,
@@ -1511,12 +1540,13 @@ def fn(
         use_quant_list = [False, True]
         use_tensor_overload_list = [False, True]
 
-        assert dtype in [
+        if dtype not in [
             torch.uint8,
             torch.int8,
             torch.float8_e4m3fn,
             torch.float8_e5m2,
-        ]
+        ]:
+            raise AssertionError(f"Unexpected dtype: {dtype}")
         quant_min = 0 if dtype == torch.uint8 else -128
         quant_max = 255 if dtype == torch.uint8 else 127
         if dtype in [torch.float8_e4m3fn, torch.float8_e5m2]:
@@ -1536,7 +1566,7 @@ def fn(
             use_tensor_overload_list,
         ):
             x = torch.clamp(
-                torch.randn((1, 7, 7, 9), dtype=torch.float32) * 100,
+                torch.randn((1, 7, 7, 9), dtype=input_dtype) * 100,
                 quant_min,
                 quant_max,
             )
@@ -1573,6 +1603,9 @@ def fn(
     @requires_vectorization
     def test_dequant_quant_lowering_uint8(self):
         self._test_dequant_quant_lowering_helper(torch.uint8)
+        self._test_dequant_quant_lowering_helper(
+            torch.uint8, input_dtype=torch.bfloat16
+        )
         self._test_dequant_quant_lowering_helper(
             torch.uint8, dequant_out_dtype=torch.bfloat16
         )
@@ -1580,6 +1613,7 @@ def test_dequant_quant_lowering_uint8(self):
     @requires_vectorization
     def test_dequant_quant_lowering_int8(self):
         self._test_dequant_quant_lowering_helper(torch.int8)
+        self._test_dequant_quant_lowering_helper(torch.int8, input_dtype=torch.bfloat16)
         self._test_dequant_quant_lowering_helper(
             torch.int8, dequant_out_dtype=torch.bfloat16
         )
@@ -1604,7 +1638,8 @@ def fn(x, scale, zero_point, quant_min, quant_max, dtype):
             )
             return max_pool2d_with_indices_default
 
-        assert dtype in [torch.uint8, torch.int8]
+        if dtype not in [torch.uint8, torch.int8]:
+            raise AssertionError(f"Unexpected dtype: {dtype}")
         quant_min = 0 if dtype == torch.uint8 else -128
         quant_max = 255 if dtype == torch.uint8 else 127
 
@@ -1672,7 +1707,8 @@ def fn(
                 )
             return y.contiguous()
 
-        assert dtype in [torch.uint8, torch.int8]
+        if dtype not in [torch.uint8, torch.int8]:
+            raise AssertionError(f"Unexpected dtype: {dtype}")
         quant_min = 0 if dtype == torch.uint8 else -128
         quant_max = 255 if dtype == torch.uint8 else 127
 
@@ -1748,7 +1784,8 @@ def fn(input, scales, zero_points, quant_min, quant_max, dtype):
 
         use_tensor_overload_list = [False, True]
         for use_tensor_overload in use_tensor_overload_list:
-            assert dtype in [torch.uint8, torch.int8]
+            if dtype not in [torch.uint8, torch.int8]:
+                raise AssertionError(f"Unexpected dtype: {dtype}")
             quant_min = 0 if dtype == torch.uint8 else -128
             quant_max = 255 if dtype == torch.uint8 else 127
             x = torch.clamp(
@@ -1765,7 +1802,10 @@ def fn(input, scales, zero_points, quant_min, quant_max, dtype):
                 torch._dynamo.reset()
                 metrics.reset()
                 self.common(fn, (x, scale, zero_point, quant_min, quant_max, dtype))
-                assert metrics.generated_cpp_vec_kernel_count == 1
+                if metrics.generated_cpp_vec_kernel_count != 1:
+                    raise AssertionError(
+                        f"Expected 1 vec kernel, got {metrics.generated_cpp_vec_kernel_count}"
+                    )
 
     @requires_vectorization
     def test_per_tensor_fake_quant_uint8(self):
@@ -1796,7 +1836,8 @@ def fn(
             )
             return input
 
-        assert dtype in [torch.uint8, torch.int8]
+        if dtype not in [torch.uint8, torch.int8]:
+            raise AssertionError(f"Unexpected dtype: {dtype}")
         quant_min = 0 if dtype == torch.uint8 else -128
         quant_max = 255 if dtype == torch.uint8 else 127
         x = torch.clamp(
@@ -1874,7 +1915,10 @@ def forward(self, input):
             torch._dynamo.reset()
             metrics.reset()
             self.common(m, (x,))
-            assert metrics.generated_cpp_vec_kernel_count == 1
+            if metrics.generated_cpp_vec_kernel_count != 1:
+                raise AssertionError(
+                    f"Expected 1 vec kernel, got {metrics.generated_cpp_vec_kernel_count}"
+                )
 
     @requires_vectorization
     def test_per_channel_fake_quant_int8(self):
@@ -1922,7 +1966,8 @@ def fn(
             x = x.view(batchsize, num_channels, height, width)
             return x
 
-        assert dtype in [torch.uint8, torch.int8]
+        if dtype not in [torch.uint8, torch.int8]:
+            raise AssertionError(f"Unexpected dtype: {dtype}")
         quant_min = 0 if dtype == torch.uint8 else -128
         quant_max = 255 if dtype == torch.uint8 else 127
 
@@ -1971,7 +2016,8 @@ def channel_shuffle(
             )
             return x.contiguous(memory_format=torch.channels_last)
 
-        assert dtype in [torch.uint8, torch.int8]
+        if dtype not in [torch.uint8, torch.int8]:
+            raise AssertionError(f"Unexpected dtype: {dtype}")
         quant_min = 0 if dtype == torch.uint8 else -128
         quant_max = 255 if dtype == torch.uint8 else 127
 
@@ -2038,7 +2084,8 @@ def fn(
             )
             return x
 
-        assert dtype in [torch.uint8, torch.int8]
+        if dtype not in [torch.uint8, torch.int8]:
+            raise AssertionError(f"Unexpected dtype: {dtype}")
         quant_min = 0 if dtype == torch.uint8 else -128
         quant_max = 255 if dtype == torch.uint8 else 127
 
@@ -2101,7 +2148,8 @@ def fn(x, y):
         fn_compiled = compile_fx_inner(fn_fx, [x1, y])
         fn(x2, y)
         fn_compiled([x3, y])
-        assert same(x2, x3)
+        if not same(x2, x3):
+            raise AssertionError("x2 and x3 are not the same")
 
     def test_int_div(self):
         def fn(x, y):
@@ -2130,9 +2178,11 @@ def fn(x1, x2):
         x1 = torch.randn((10, 20))
         x2 = torch.randn((10, 20))
         with set_num_threads(1):
-            assert same(x1 + x2, fn(x1, x2))
+            if not same(x1 + x2, fn(x1, x2)):
+                raise AssertionError("Results differ with 1 thread")
         with set_num_threads(4):
-            assert same(x1 + x2, fn(x1, x2))
+            if not same(x1 + x2, fn(x1, x2)):
+                raise AssertionError("Results differ with 4 threads")
 
     @patch("torch.cuda.is_available", lambda: False)
     def test_timed_cpu_only(self):
@@ -2383,7 +2433,10 @@ def fn(value, mask):
                         torch._dynamo.reset()
                         metrics.reset()
                         self.common(fn, (value, mask))
-                        assert metrics.generated_cpp_vec_kernel_count >= 1
+                        if metrics.generated_cpp_vec_kernel_count < 1:
+                            raise AssertionError(
+                                f"Expected >= 1 vec kernels, got {metrics.generated_cpp_vec_kernel_count}"
+                            )
 
     def test_channels_last_view_as_complex(self):
         # https://github.com/pytorch/pytorch/issues/122448#issuecomment-2046169554
@@ -2743,7 +2796,10 @@ def fn(a, dim, index, b):
         torch._dynamo.reset()
         metrics.reset()
         self.common(fn, inps)
-        assert metrics.generated_cpp_vec_kernel_count == 2
+        if metrics.generated_cpp_vec_kernel_count != 2:
+            raise AssertionError(
+                f"Expected 2 vec kernels, got {metrics.generated_cpp_vec_kernel_count}"
+            )
 
         with (
             set_num_threads(1),
@@ -2752,7 +2808,10 @@ def fn(a, dim, index, b):
             torch._dynamo.reset()
             metrics.reset()
             self.common(fn, inps)
-            assert metrics.generated_cpp_vec_kernel_count == 2
+            if metrics.generated_cpp_vec_kernel_count != 2:
+                raise AssertionError(
+                    f"Expected 2 vec kernels, got {metrics.generated_cpp_vec_kernel_count}"
+                )
 
     def test_large_mean(self):
         size = (30000, 100000)
@@ -2778,6 +2837,47 @@ def fn(x):
         actual = torch.compile(fn)(x)
         self.assertEqual(expected, actual, atol=1e-4, rtol=1e-4)
 
+    def test_two_step_variance(self):
+        M = 64
+        N = 1024
+
+        class L(torch.nn.Module):
+            def __init__(self, normalized_shape=N, eps=1e-5):
+                super().__init__()
+                self.layernorm = torch.nn.LayerNorm(normalized_shape, eps=eps)
+
+            def forward(self, x):
+                return self.layernorm(x)
+
+        mod = L().eval()
+        for mean, std in [
+            (0, 1e10),
+            (1e10, 10),
+            (1e10, 1e10),
+            (0, 1),
+            (0, 0.5),
+            (0.5, 1),
+            (0, 2),
+        ]:
+            x = torch.randn(M, N)
+            row_means = x.mean(dim=1, keepdim=True)
+            row_stds = x.std(dim=1, keepdim=True, unbiased=True)
+            x_norm = (x - row_means) / (row_stds + 1e-5)
+            x = x_norm * std + mean
+            input = (x,)
+            output_eager = mod(*input)
+            with torch.no_grad():
+                m = torch.compile(mod)
+                output_compiled = m(*input)
+            if mean == 1e10 or std == 1e10:
+                self.assertTrue(
+                    torch.allclose(output_eager, output_compiled, atol=1, rtol=1e-4)
+                )
+            else:
+                self.assertTrue(
+                    torch.allclose(output_eager, output_compiled, atol=1e-4, rtol=1e-4)
+                )
+
     @unittest.skipIf(IS_FBCODE, "Not yet runnable in fbcode")
     @requires_vectorization
     @patch("torch.cuda.is_available", lambda: False)
@@ -3049,10 +3149,13 @@ def fn(x):
                 metrics.reset()
                 self.common(fn, (x,))
                 check_metrics_vec_kernel_count(1)
-                assert (
+                if (
                     metrics.generated_kernel_count
                     - metrics.generated_cpp_vec_kernel_count
-                ) == 0
+                ) != 0:
+                    raise AssertionError(
+                        f"Expected 0, got {metrics.generated_kernel_count - metrics.generated_cpp_vec_kernel_count}"
+                    )
 
     @requires_vectorization
     def test_vec_remainder(self):
@@ -3162,7 +3265,10 @@ def fn(x):
         torch._dynamo.reset()
         metrics.reset()
         self.common(fn, (x,))
-        assert metrics.cpp_to_dtype_count == 0
+        if metrics.cpp_to_dtype_count != 0:
+            raise AssertionError(
+                f"Expected cpp_to_dtype_count == 0, got {metrics.cpp_to_dtype_count}"
+            )
         check_metrics_vec_kernel_count(1)
 
     def test_insert_to_dtype_count(self):
@@ -3175,7 +3281,10 @@ def fn(x):
         torch._dynamo.reset()
         metrics.reset()
         self.common(fn, (x,))
-        assert metrics.cpp_to_dtype_count == 2
+        if metrics.cpp_to_dtype_count != 2:
+            raise AssertionError(
+                f"Expected cpp_to_dtype_count == 2, got {metrics.cpp_to_dtype_count}"
+            )
         check_metrics_vec_kernel_count(1)
 
     def test_memory_copy_with_fusion(self):
@@ -3189,7 +3298,10 @@ def fn(x):
         torch._dynamo.reset()
         metrics.reset()
         self.common(fn, (x,))
-        assert metrics.cpp_to_dtype_count == 2
+        if metrics.cpp_to_dtype_count != 2:
+            raise AssertionError(
+                f"Expected cpp_to_dtype_count == 2, got {metrics.cpp_to_dtype_count}"
+            )
         check_metrics_vec_kernel_count(1)
 
     @requires_vectorization
@@ -3293,7 +3405,10 @@ def fn(x):
                 torch._dynamo.reset()
                 metrics.reset()
                 self.common(fn, (x,))
-                assert metrics.generated_cpp_vec_kernel_count == 1
+                if metrics.generated_cpp_vec_kernel_count != 1:
+                    raise AssertionError(
+                        f"Expected 1 vec kernel, got {metrics.generated_cpp_vec_kernel_count}"
+                    )
 
     @config.patch({"fx_graph_cache": False, "fx_graph_remote_cache": False})
     def test_outer_loop_fusion(self):
@@ -3376,6 +3491,8 @@ def fn(x):
                 3,
             )
 
+    @xfailIf(IS_ARM64 and not IS_CPU_CAPABILITY_SVE256)
+    # see https://github.com/pytorch/pytorch/issues/142231
     @config.patch({"fx_graph_cache": False, "fx_graph_remote_cache": False})
     def test_two_local_buffers_in_outer_loop_fusion(self):
         def fn(x):
@@ -3512,7 +3629,10 @@ def fn(x):
             torch._dynamo.reset()
             metrics.reset()
             self.common(fn, (x,))
-            assert metrics.generated_cpp_vec_kernel_count == 1
+            if metrics.generated_cpp_vec_kernel_count != 1:
+                raise AssertionError(
+                    f"Expected 1 vec kernel, got {metrics.generated_cpp_vec_kernel_count}"
+                )
 
     @requires_vectorization
     def test_argmax_argmin_with_nan_value(self):
@@ -3537,13 +3657,19 @@ def fn2(x):
             torch._dynamo.reset()
             metrics.reset()
             self.common(fn, (x,))
-            assert metrics.generated_cpp_vec_kernel_count == 1
+            if metrics.generated_cpp_vec_kernel_count != 1:
+                raise AssertionError(
+                    f"Expected generated_cpp_vec_kernel_count == 1, got {metrics.generated_cpp_vec_kernel_count}"
+                )
 
             # Test argmin
             torch._dynamo.reset()
             metrics.reset()
             self.common(fn2, (x,))
-            assert metrics.generated_cpp_vec_kernel_count == 1
+            if metrics.generated_cpp_vec_kernel_count != 1:
+                raise AssertionError(
+                    f"Expected generated_cpp_vec_kernel_count == 1, got {metrics.generated_cpp_vec_kernel_count}"
+                )
 
     # Currently, we enabled AVX2 and AVX512 for vectorization. If the platform is not
     # supported, the vectorization will not work and skip this test case. For ARM or
@@ -3594,7 +3720,10 @@ def fn(x1, x2):
                 torch._dynamo.reset()
                 metrics.reset()
                 self.common(fn, (x1, x2))
-                assert metrics.generated_cpp_vec_kernel_count == 0
+                if metrics.generated_cpp_vec_kernel_count != 0:
+                    raise AssertionError(
+                        f"Expected generated_cpp_vec_kernel_count == 0, got {metrics.generated_cpp_vec_kernel_count}"
+                    )
 
             with config.patch({"cpp.simdlen": None}):
                 torch._dynamo.reset()
@@ -3641,7 +3770,10 @@ def fn(a, b):
         for e in prof.profiler.function_events:
             if "cpp_fused_add_0" in e.name:
                 kernel_profile_events.append(e.name)
-        assert len(kernel_profile_events) > 0
+        if len(kernel_profile_events) <= 0:
+            raise AssertionError(
+                f"Expected kernel_profile_events to be non-empty, got {len(kernel_profile_events)}"
+            )
 
     @xfailIfS390X
     @requires_vectorization
@@ -3697,8 +3829,10 @@ def forward(self, x):
                     metrics.reset()
                     m = Model().eval() if eval_mode else Model()
                     self.common(m, (x,))
-                    check_metrics_vec_kernel_count(6)
+                    check_metrics_vec_kernel_count(8)
 
+    @xfailIf(IS_ARM64 and IS_CPU_CAPABILITY_SVE256)
+    # see https://github.com/pytorch/pytorch/issues/169958
     @requires_vectorization
     @config.patch("cpp.enable_tiling_heuristics", False)
     def test_transpose_copy(self):
@@ -3802,7 +3936,10 @@ def fn(x):
             x = torch.randn(100, 100).to(dtype)
             opt_fn = torch.compile(fn, backend="inductor")
             self.assertTrue(same(fn(x), opt_fn(x)))
-            assert metrics.cpp_to_dtype_count == 0
+            if metrics.cpp_to_dtype_count != 0:
+                raise AssertionError(
+                    f"Expected cpp_to_dtype_count == 0, got {metrics.cpp_to_dtype_count}"
+                )
             check_metrics_vec_kernel_count(1)
 
     @config.patch("cpp.enable_tiling_heuristics", False)
@@ -3943,6 +4080,8 @@ def fn(a, b):
         self.common(fn, (x, y))
         check_metrics_vec_kernel_count(2)
 
+    @xfailIf(IS_ARM64 and IS_CPU_CAPABILITY_SVE256)
+    # see https://github.com/pytorch/pytorch/issues/170877
     def test_transpose_mxn_16_16_bf16_fp16(self):
         def fn(a, b):
             c = a * b
@@ -4369,7 +4508,10 @@ def forward(self, x):
                 mod,
                 (v,),
             )
-            assert metrics.generated_kernel_count == 0
+            if metrics.generated_kernel_count != 0:
+                raise AssertionError(
+                    f"Expected generated_kernel_count == 0, got {metrics.generated_kernel_count}"
+                )
 
     @config.patch(implicit_fallbacks=True)
     def test_aten_normal_dtype(self):
@@ -4513,6 +4655,35 @@ def fn(x, weight, bias):
         torch.testing.assert_close(weight_cmp.grad, weight_ref.grad)
         torch.testing.assert_close(bias_cmp.grad, bias_ref.grad)
 
+    @torch._dynamo.config.patch(
+        capture_scalar_outputs=True, capture_dynamic_output_shape_ops=True
+    )
+    def test_backward_dynamic_item_from_size1_tensor(self):
+        def fn(x, w):
+            num = x.nonzero().numel()
+            num = x.new_tensor([num]).item()
+            y = w * x
+            return y.sum() / num
+
+        torch._dynamo.reset()
+        metrics.reset()
+
+        x = torch.tensor([0.0, 1.0, -1.0, 0.5])
+        w_ref = torch.tensor(0.0, requires_grad=True)
+        w_cmp = w_ref.detach().clone().requires_grad_(True)
+
+        eager_out = fn(x, w_ref)
+        eager_out.backward()
+
+        compiled = torch.compile(fn, backend="inductor", fullgraph=True, dynamic=True)
+        compiled_out = compiled(x, w_cmp)
+        compiled_out.backward()
+
+        self.assertEqual(eager_out.shape, torch.Size([]))
+        self.assertEqual(compiled_out.shape, eager_out.shape)
+        torch.testing.assert_close(compiled_out, eager_out)
+        torch.testing.assert_close(w_cmp.grad, w_ref.grad)
+
     @config.patch(emulate_precision_casts=True)
     def test_emulate_precision_casts_cpp_backend_no_error(self):
         """
@@ -4584,6 +4755,7 @@ def fn(x, y):
         y = torch.randint(0, 255, (3, 3), dtype=torch.uint8)
         self.common(fn, (x, y))
 
+    @xfailIf(IS_ARM64)  # see https://github.com/pytorch/pytorch/issues/168972
     def test_float32_to_uint8(self):
         # https://github.com/pytorch/pytorch/issues/156788
         @torch.compile
@@ -4712,7 +4884,7 @@ def forward(self, x):
             _, code = run_and_get_cpp_code(opt_m, x)
             self.assertTrue(same(m(x), opt_m(x)))
             # Two kernels: one for reduction, one pointwises
-            check_metrics_vec_kernel_count(2)
+            check_metrics_vec_kernel_count(4)
             FileCheck().check_count(
                 "Vectorized<float>::loadu(tmpbuf.data())", 0, exactly=True
             ).run(code)
@@ -4844,7 +5016,10 @@ def fn(x):
         metrics.reset()
         self.common(fn, (x,))
         # TODO(jgong5): change to 1 with vectorized uint32 load
-        assert metrics.generated_cpp_vec_kernel_count == 0
+        if metrics.generated_cpp_vec_kernel_count != 0:
+            raise AssertionError(
+                f"Expected generated_cpp_vec_kernel_count == 0, got {metrics.generated_cpp_vec_kernel_count}"
+            )
 
     def test_uint32_reduction_vec(self):
         def fn(x):
@@ -4854,7 +5029,10 @@ def fn(x):
         metrics.reset()
         self.common(fn, (x,))
         # TODO(jgong5): change to 1 with vectorized uint32/uint64 load
-        assert metrics.generated_cpp_vec_kernel_count == 0
+        if metrics.generated_cpp_vec_kernel_count != 0:
+            raise AssertionError(
+                f"Expected generated_cpp_vec_kernel_count == 0, got {metrics.generated_cpp_vec_kernel_count}"
+            )
 
     def test_int64_pointwise_vec(self):
         def fn(x):
@@ -4908,6 +5086,26 @@ def fn(x):
                     "at::vec::VectorizedN<int64_t,2>::loadu", 2, exactly=True
                 ).run(code)
 
+    @requires_vectorization
+    def test_indirect_assert_scalar_mask_tail_vec_no_crash(self):
+        # https://github.com/pytorch/pytorch/issues/178136
+        def fn(positions, cache):
+            x = cache[positions]
+            y = x[0].clone()
+            y[..., 1::3] = x[1, ..., 1::3]
+            y[..., 2::3] = x[2, ..., 2::3]
+            return y
+
+        positions = torch.tensor([[0, 0], [0, 0], [0, 0]], dtype=torch.int64)
+        cache = torch.arange(3, dtype=torch.float32).reshape(1, 3)
+
+        with config.patch({"cpp.enable_loop_tail_vec": True}):
+            expected = fn(positions, cache)
+            compiled_fn = torch.compile(fn, fullgraph=True)
+            actual = compiled_fn(positions, cache)
+
+        torch.testing.assert_close(actual, expected)
+
     def test_uint64_pointwise_vec(self):
         def fn(x):
             return x * x
@@ -4916,7 +5114,10 @@ def fn(x):
         metrics.reset()
         self.common(fn, (x,))
         # TODO(jgong5): change to 1 with vectorized uint64 load
-        assert metrics.generated_cpp_vec_kernel_count == 0
+        if metrics.generated_cpp_vec_kernel_count != 0:
+            raise AssertionError(
+                f"Expected generated_cpp_vec_kernel_count == 0, got {metrics.generated_cpp_vec_kernel_count}"
+            )
 
     def test_uint64_reduction_vec(self):
         def fn(x):
@@ -4926,7 +5127,10 @@ def fn(x):
         metrics.reset()
         self.common(fn, (x,))
         # TODO(jgong5): change to 1 with vectorized uint64 load
-        assert metrics.generated_cpp_vec_kernel_count == 0
+        if metrics.generated_cpp_vec_kernel_count != 0:
+            raise AssertionError(
+                f"Expected generated_cpp_vec_kernel_count == 0, got {metrics.generated_cpp_vec_kernel_count}"
+            )
 
     def test_convert_int8_to_half_vec(self):
         src_dtypes = [torch.int8, torch.uint8]
@@ -5036,6 +5240,10 @@ def fn(x):
                     "at::vec::VectorizedN<double,2>::loadu", 2, exactly=True
                 ).run(code)
 
+    @unittest.skipIf(
+        get_gcc_major_version() == 13,
+        "Fails under GCC 13 due to vector codegen (passes with GCC 11)",
+    )
     def test_convert_fp32_to_double_vec(self):
         def fn(x):
             return x.to(torch.double)
@@ -5126,6 +5334,52 @@ def fn(x):
         self.common(fn, (x,))
         check_metrics_vec_kernel_count(1)
 
+    def test_masked_handle_scalar_var(self):
+        def fn():
+            # Simplified reproducer of https://github.com/pytorch/pytorch/issues/173626
+            window_size = 7
+            shift_size = 3
+            height, width = 14, 14
+            img_mask = torch.zeros((1, height, width, 1), dtype=torch.float32)
+            h_slices = (
+                slice(0, -window_size),
+                slice(-window_size, -shift_size),
+                slice(-shift_size, None),
+            )
+            w_slices = (
+                slice(0, -window_size),
+                slice(-window_size, -shift_size),
+                slice(-shift_size, None),
+            )
+            count = 0
+            for h in h_slices:
+                for w in w_slices:
+                    img_mask[:, h, w, :] = count
+                    count += 1
+
+            masked_window = img_mask.view(
+                1,
+                height // window_size,
+                window_size,
+                width // window_size,
+                window_size,
+                1,
+            )
+            masked_window = (
+                masked_window.permute(0, 1, 3, 2, 4, 5)
+                .contiguous()
+                .view(-1, window_size * window_size)
+            )
+            attn_mask = masked_window.unsqueeze(1) - masked_window.unsqueeze(2)
+            res = attn_mask.masked_fill(attn_mask != 0, -100.0).masked_fill(
+                attn_mask == 0, 0.0
+            )
+            return res
+
+        metrics.reset()
+        self.common(fn, ())
+        check_metrics_vec_kernel_count(1)
+
     def test_highp_to_lowp_cse_var_cache_with_store(self):
         # Fix issue: https://github.com/pytorch/pytorch/issues/128263
         input = torch.randn(5, 128, dtype=torch.float32)
@@ -5221,6 +5475,9 @@ def forward(self, x):
         x = torch.randn(1, 4, 2, 2)
         self.common(fn, (x,))
 
+    @xfailIf(
+        IS_ARM64 and IS_CPU_EXT_SVE_SUPPORTED
+    )  # see https://github.com/pytorch/pytorch/issues/142134
     @parametrize("is_inference", (True, False))
     def test_disabled_amp(self, is_inference):
         class M(torch.nn.Module):
@@ -5677,6 +5934,87 @@ def func2(arg0, arg1):
             ):
                 check_use_full_bits(func, shapes, dtype, mixed, check_vecn)
 
+    @requires_vectorization
+    def test_full_bits_fp8_e4m3fn(self):
+        """
+        Test VecConvert<float,2,Float8_e4m3fn,1> and VecConvert<Float8_e4m3fn,1,float,2>
+        are emitted when fp8 and bf16 tensors are mixed in the same kernel.
+
+        A pure fp8 cast alone uses a 16-element loop (convert<float,1,...>).
+        Only when fp8 is combined with bf16 does the loop widen to 32 elements,
+        triggering convert<float,2,Float8_e4m3fn,1> and
+        convert<Float8_e4m3fn,1,float,2>.
+
+        Sub-cases:
+          func0 - fp8 dequant * bf16: tests convert<float,2,Float8_e4m3fn,1>
+          func1 - float * bf16 → fp8: tests convert<Float8_e4m3fn,1,float,2>
+        """
+        fp8_dtype = torch.float8_e4m3fn
+        deq_scale = 0.15
+
+        # func0: fp8 dequant * bf16
+        # Tests convert<float,2,Float8_e4m3fn,1>
+        def func0(arg0_fp8, arg1_bf16):
+            x = arg0_fp8.to(torch.float) * deq_scale
+            y = arg1_bf16.to(torch.float)
+            return x * y
+
+        # func1: float * bf16 -> fp8
+        # Tests convert<Float8_e4m3fn,1,float,2>
+        def func1(arg0_float, arg1_bf16):
+            y = arg1_bf16.to(torch.float)
+            z = arg0_float * y
+            return z.to(fp8_dtype)
+
+        large_shape = (10, 32, 20, 20)
+
+        # func0: verify fp8 dequant wide path
+        torch._dynamo.reset()
+        x_fp8 = torch.randn(large_shape).to(fp8_dtype)
+        x_bf16 = torch.randn(large_shape, dtype=torch.bfloat16)
+        _, code0 = run_and_get_cpp_code(torch.compile()(func0), x_fp8, x_bf16)
+        self.assertIn(
+            "at::vec::convert<float,2,at::Float8_e4m3fn,1>",
+            code0,
+            "Expected convert<float,2,at::Float8_e4m3fn,1> in generated code for func0",
+        )
+
+        # func1: verify fp8 quant wide path
+        torch._dynamo.reset()
+        x_float = torch.randn(large_shape)
+        x_bf16 = torch.randn(large_shape, dtype=torch.bfloat16)
+        _, code1 = run_and_get_cpp_code(torch.compile()(func1), x_float, x_bf16)
+        self.assertIn(
+            "at::vec::convert<at::Float8_e4m3fn,1,float,2>",
+            code1,
+            "Expected convert<at::Float8_e4m3fn,1,float,2> in generated code for func1",
+        )
+
+    @requires_vectorization
+    def test_bool_to_float8_e4m3fn(self):
+        """
+        Test that bool to float8_e4m3fn cast succeeds.
+        Issue: https://github.com/pytorch/pytorch/issues/178095
+        """
+
+        def fn(x):
+            return x.to(dtype=torch.float8_e4m3fn)
+
+        x = torch.ones(64, dtype=torch.bool)
+        self.common(fn, (x,))
+
+    @requires_vectorization
+    def test_bool_to_float8_e5m2(self):
+        """
+        Test that bool to float8_e5m2 cast succeeds.
+        """
+
+        def fn(x):
+            return x.to(dtype=torch.float8_e5m2)
+
+        x = torch.ones(64, dtype=torch.bool)
+        self.common(fn, (x,))
+
     @config.patch("cpp.simdlen", 256)
     @requires_vectorization
     def test_avx2_bool_constant_pad_nd(self):
@@ -5750,7 +6088,7 @@ def forward(self, context_layer, hidden_states):
         with torch.no_grad():
             metrics.reset()
             torch.compile(model)(*example_batch)
-            check_metrics_vec_kernel_count(3)
+            check_metrics_vec_kernel_count(5)
 
     def test_dropout(self):
         class Model(nn.Module):
@@ -5868,6 +6206,7 @@ def fn(x):
 
         torch.compile(fn)(torch.randn(2, 2))
 
+    @xfailIf(IS_ARM64)  # https://github.com/pytorch/pytorch/issues/176285
     @skipIfRocmArch(MI200_ARCH)
     @unittest.skipIf(not torch.backends.mkldnn.is_available(), "MKLDNN is not enabled")
     @requires_vectorization
@@ -5928,6 +6267,66 @@ def forward(self, x):
                 rtol = 1e-2 if dtype == torch.bfloat16 else 1e-5
                 torch.testing.assert_close(ref_res, res, atol=atol, rtol=rtol)
 
+    # https://github.com/pytorch/pytorch/issues/136640
+    def test_inductor_dynamic_shapes_broadcasting(self) -> None:
+        def fn(x, y):
+            x_view = x.view(-1, 4)
+            y_view = y.view(-1, 4)
+            return x_view * y_view
+
+        x = torch.randn(4)
+        y = torch.randn(8)
+        out_ref = fn(x, y)
+        out_test = torch.compile(fn, dynamic=True, backend="inductor")(x, y)
+        self.assertEqual(out_ref, out_test)
+
+    # https://github.com/pytorch/pytorch/issues/119162
+    def test_inductor_rng_default_dtype(self) -> None:
+        @torch.compile
+        def fn():
+            tmp = torch.randn(4, 4, dtype=torch.bfloat16)
+            return tmp
+
+        try:
+            old = torch.get_default_dtype()
+            torch.set_default_dtype(torch.bfloat16)
+            out = fn()
+        finally:
+            torch.set_default_dtype(old)
+        # output dtype should be float32
+        self.assertEqual(out.dtype, torch.bfloat16)
+
+    def test_inductor_no_recursionerror_on_for_loops(self):
+        def forward(x):
+            for _ in range(10000):
+                x = 1.0 * x
+            return x
+
+        self.assertTrue(
+            same(
+                torch.compile(forward, backend="inductor")(torch.tensor([1.0])),
+                torch.tensor([1.0]),
+            )
+        )
+
+    def test_mutation_transpose_reshape_ordering(self):
+        def fn(x, y):
+            return x.add_(y).reshape(-1, 1, 3)
+
+        torch.manual_seed(0)
+        x1 = torch.ones([1, 2, 3]).transpose(1, 2)
+        y1 = torch.randn(1, 3, 1)
+
+        x2 = x1.clone()
+        y2 = y1.clone()
+
+        cfunc = torch.compile(fn, backend="inductor")
+
+        out1 = fn(x1, y1)
+        out2 = cfunc(x2, y2)
+
+        self.assertTrue(torch.allclose(out1, out2, equal_nan=True))
+
 
 if __name__ == "__main__":
     from torch._inductor.test_case import run_tests
diff --git a/test/inductor/test_cpu_select_algorithm.py b/test/inductor/test_cpu_select_algorithm.py
index a76ea804cb612..befadb7a331a7 100644
--- a/test/inductor/test_cpu_select_algorithm.py
+++ b/test/inductor/test_cpu_select_algorithm.py
@@ -3,7 +3,6 @@
 import functools
 import sys
 import unittest
-from typing import Optional
 from unittest.mock import patch
 
 import torch
@@ -27,10 +26,13 @@
     _calculate_dynamic_per_channel_qparams,
 )
 from torch.testing._internal.common_utils import (
+    IS_ARM64,
+    IS_CPU_EXT_SVE_SUPPORTED,
     IS_MACOS,
     IS_WINDOWS,
     parametrize,
     TEST_MKL,
+    xfailIf,
 )
 
 
@@ -59,15 +61,17 @@ def skip_cache(self, choices, name, key, benchmark, hint_override=None):
         timings = benchmark(choices)
         for choice, timing in timings.items():
             if isinstance(choice, select_algorithm.ExternKernelCaller):
-                # we intentionally make ATEN kernel slower to cover the cases
+                # We intentionally make ATEN kernel slower to cover the cases
                 # where template kernels are always chosen with fusions applied
-                # and correctness checks at runtime.
-                timings[choice] = timing * 1000
+                # and correctness checks at runtime. On k8s ARC runner pods,
+                # CPU contention from parallel tests can make cpp template
+                # benchmarks 2-3x slower than normal, so the multiplier needs
+                # to be large enough to still exceed those inflated times.
+                timings[choice] = timing * 1000000
         return timings
 
     for patcher in [
         dynamo_config.patch(verbose=True),
-        dynamo_config.patch(inline_inbuilt_nn_modules=True),
         inductor_config.patch(
             debug=True,
             max_autotune=True,
@@ -100,7 +104,7 @@ def verify(dtype):
         yield atol, rtol
 
 
-def _get_epilogue(epilogue: str, other: Optional[torch.Tensor] = None):
+def _get_epilogue(epilogue: str, other: torch.Tensor | None = None):
     if epilogue == "none":
         return lambda x: x
     elif epilogue == "relu":
@@ -802,7 +806,8 @@ def forward(self, arg7_1):
         self.assertEqual(counters["inductor"]["cpp_epilogue_fusion_counter"], 0)
 
     @unittest.skipIf(
-        not torch.cpu._is_amx_tile_supported(), "AMX ISA support is required"
+        not isinstance(torch._inductor.cpu_vec_isa.pick_vec_isa(), VecAMX),
+        "AMX ISA support is required",
     )
     @inductor_config.patch({"freezing": True})
     @patches
@@ -1525,7 +1530,8 @@ def forward(self, x0, x1, other):
             )
 
     @unittest.skipIf(
-        not torch.cpu._is_amx_tile_supported(), "AMX ISA support is required"
+        not isinstance(torch._inductor.cpu_vec_isa.pick_vec_isa(), VecAMX),
+        "AMX ISA support is required",
     )
     @inductor_config.patch({"freezing": True})
     @patches
@@ -1631,6 +1637,9 @@ def forward(self, qx, x_scale, x_zp):
     )
     @parametrize("in_features", (128, 144, 1024))
     @parametrize("out_features", (64, 65, 1024))
+    @unittest.skipIf(
+        IS_ARM64 and not IS_CPU_EXT_SVE_SUPPORTED, "flaky on AArch64 (no SVE)"
+    )
     def test_int8_woq_mm(self, dtype, batch_size, mid_dim, in_features, out_features):
         def _convert_weight_to_int8pack(w):
             scale, zp = _calculate_dynamic_per_channel_qparams(
@@ -1693,6 +1702,9 @@ def forward(self, x, scale):
     )
     @parametrize("in_features", (128,))
     @parametrize("out_features", (64,))
+    @unittest.skipIf(
+        IS_ARM64 and not IS_CPU_EXT_SVE_SUPPORTED, "flaky on AArch64 (no SVE)"
+    )
     def test_int8_woq_mm_concat(
         self, dtype, batch_size, mid_dim, in_features, out_features
     ):
@@ -1758,7 +1770,8 @@ def forward(self, x, scale1, scale2, scale3):
             self._check_amx_counter(vec_amx)
 
     @unittest.skipIf(
-        not torch.cpu._is_amx_tile_supported(), "AMX ISA support is required"
+        not isinstance(torch._inductor.cpu_vec_isa.pick_vec_isa(), VecAMX),
+        "AMX ISA support is required",
     )
     @inductor_config.patch({"freezing": True})
     @patches
@@ -1842,9 +1855,8 @@ def forward(self, a):
 
         vec_amx = VecAMX()
         self._check_amx_counter(vec_amx)
-        if torch.cpu._is_amx_tile_supported():
-            # Only AMX ISA based micro-kernel is currently supported for da8w8
-            self.assertEqual(counters["inductor"]["cpp_templated_kernel_counter"], 1)
+        # Only AMX ISA based micro-kernel is currently supported for da8w8
+        self.assertEqual(counters["inductor"]["cpp_templated_kernel_counter"], 1)
 
     @inductor_config.patch({"freezing": True})
     @patches
@@ -1887,7 +1899,8 @@ def forward(self, x):
         )
 
     @unittest.skipIf(
-        not torch.cpu._is_amx_tile_supported(), "AMX ISA support is required"
+        not isinstance(torch._inductor.cpu_vec_isa.pick_vec_isa(), VecAMX),
+        "AMX ISA support is required",
     )
     @inductor_config.patch({"freezing": True})
     @patches
@@ -1958,7 +1971,8 @@ def forward(self, x):
             self.assertEqual(counters["inductor"]["select_algorithm_autotune"], 1)
 
     @unittest.skipIf(
-        not torch.cpu._is_amx_tile_supported(), "AMX ISA support is required"
+        not isinstance(torch._inductor.cpu_vec_isa.pick_vec_isa(), VecAMX),
+        "AMX ISA support is required",
     )
     @inductor_config.patch({"freezing": True})
     @inductor_config.patch({"cpp.use_small_dequant_buffer": True})
@@ -2001,8 +2015,10 @@ def forward(self, x):
             _target_code_check = f"constexpr int64_t Kc_blocks = {group_size // kr};"
             torch._C.FileCheck().check(_target_code_check).run(code)
 
+    @xfailIf(IS_ARM64)  # Int4 kernel numerical errors (5.4x rel diff, 5.8% mismatch)
     @unittest.skipIf(
-        not torch.cpu._is_amx_tile_supported(), "AMX ISA support is required"
+        not isinstance(torch._inductor.cpu_vec_isa.pick_vec_isa(), VecAMX),
+        "AMX ISA support is required",
     )
     @inductor_config.patch({"freezing": True})
     @patches
@@ -2071,8 +2087,10 @@ def forward(self, x):
             )
             self.assertEqual(counters["inductor"]["select_algorithm_autotune"], 1)
 
+    @xfailIf(IS_ARM64)  # Int4 kernel numerical errors (43.5x rel diff, 10.7% mismatch)
     @unittest.skipIf(
-        not torch.cpu._is_amx_tile_supported(), "AMX ISA support is required"
+        not isinstance(torch._inductor.cpu_vec_isa.pick_vec_isa(), VecAMX),
+        "AMX ISA support is required",
     )
     @inductor_config.patch({"freezing": True})
     @inductor_config.patch({"cpp.enable_concat_linear": True})
@@ -2286,7 +2304,8 @@ def __init__(self, bias=bias) -> None:
             def forward(self, x):
                 return self.mlp(x)
 
-        assert torch._inductor.config.freezing is False
+        if torch._inductor.config.freezing is not False:
+            raise AssertionError
 
         counters.clear()
         v = torch.randn(batch_size, in_features).to(dtype=dtype)
@@ -2434,6 +2453,7 @@ def forward(self, x):
             ["silu", "mul"],
         ),
     )
+    @parametrize("loop_ordering_after_fusion", (True, False))
     def test_grouped_linear_epilogue(
         self,
         batch_size,
@@ -2442,6 +2462,7 @@ def test_grouped_linear_epilogue(
         input_3d,
         bias,
         epilogue,
+        loop_ordering_after_fusion,
     ):
         class M(torch.nn.Module):
             def __init__(self, in_feature, out_feature, bias, epilogue):
@@ -2482,6 +2503,9 @@ def forward(self, x):
                 verify(dtype) as (atol, rtol),
                 torch.autocast(device_type="cpu", dtype=dtype),
                 torch.no_grad(),
+                inductor_config.patch(
+                    {"loop_ordering_after_fusion": loop_ordering_after_fusion}
+                ),
             ):
                 self.common(mod, (v,), atol=atol, rtol=rtol)
             self.assertEqual(counters["inductor"]["cpp_grouped_gemm_template"], 1)
@@ -2527,7 +2551,8 @@ def forward(self, x):
                 )
                 return self.relu(tmp)
 
-        assert torch._inductor.config.freezing is False
+        if torch._inductor.config.freezing is not False:
+            raise AssertionError
 
         counters.clear()
         v = torch.randn(batch_size, in_features).to(dtype=dtype)
@@ -2690,7 +2715,8 @@ def forward(self, x, y):
         u = torch.randn(bs, Mdim, Kdim).to(dtype=dtype)
         v = torch.randn(bs, Kdim, Ndim).to(dtype=dtype)
         mod = M().to(dtype=dtype).eval()
-        with verify(dtype) as (atol, rtol), torch.amp.autocast("cpu"):
+        # CPU autocast converts to bfloat16, so use bfloat16 tolerances
+        with verify(torch.bfloat16) as (atol, rtol), torch.amp.autocast("cpu"):
             self.common(mod, (u, v), atol=atol, rtol=rtol)
         self.assertEqual(counters["inductor"]["cpp_templated_kernel_counter"], 1)
 
@@ -2831,7 +2857,7 @@ def __init__(self):
                 super().__init__()
 
             def forward(self, x, w):
-                assert x.dim() == 2, f"Expected x to be 2D, got {x.dim()}D"
+                assert x.dim() == 2, f"Expected x to be 2D, got {x.dim()}D"  # noqa: S101
                 x_expanded = x.unsqueeze(0).expand(bs, -1, -1)
                 return x_expanded @ w
 
@@ -3057,7 +3083,8 @@ def forward(self, u, v):
             self.common(mod, (u, v))
 
     @unittest.skipIf(
-        not torch.cpu._is_amx_tile_supported(), "AMX ISA support is required"
+        not isinstance(torch._inductor.cpu_vec_isa.pick_vec_isa(), VecAMX),
+        "AMX ISA support is required",
     )
     @inductor_config.patch({"freezing": True})
     @patches
@@ -3095,7 +3122,8 @@ def forward(self, x):
                 exact_dtype=True,
             )
             # Check that only 2 kernels are in the generated code
-            assert code.count("AMXState amx_state") == 2
+            if code.count("AMXState amx_state") != 2:
+                raise AssertionError
 
 
 @dynamo_config.patch({"dynamic_shapes": True, "assume_static_by_default": False})
diff --git a/test/inductor/test_cuda_repro.py b/test/inductor/test_cuda_repro.py
index 2348e151dcac8..574ee73bae36f 100644
--- a/test/inductor/test_cuda_repro.py
+++ b/test/inductor/test_cuda_repro.py
@@ -44,8 +44,11 @@
     parametrize,
     skipIfRocm,
     skipIfRocmArch,
+    skipIfXpu,
+    TEST_CUDA,
     TEST_WITH_ASAN,
     TEST_WITH_ROCM,
+    TEST_XPU,
     xfailIfROCm,
 )
 from torch.testing._internal.inductor_utils import IS_BIG_GPU
@@ -89,14 +92,21 @@
 aten = torch.ops.aten
 
 
+device_type = (
+    acc.type
+    if (acc := torch.accelerator.current_accelerator(check_available=True))
+    else "cpu"
+)
+
+
 @instantiate_parametrized_tests
 class CudaReproTests(TestCase):
-    device = "cuda"
+    device = device_type
     common = check_model_cuda
 
     def test_mm_out_dtype_compile(self):
-        a = torch.randn(1, 3, device="cuda", dtype=torch.float16)
-        b = torch.randn(3, 2, device="cuda", dtype=torch.float16)
+        a = torch.randn(1, 3, device=device_type, dtype=torch.float16)
+        b = torch.randn(3, 2, device=device_type, dtype=torch.float16)
 
         def fn(x, y):
             return torch.mm(x, y, out_dtype=torch.float32)
@@ -135,7 +145,8 @@ def forward(
             (torch.Size([512, 768]), torch.float16),
         ]
         inps = [torch.zeros(())] + [
-            torch.ones(shape, dtype=dtype, device="cuda") for (shape, dtype) in inps
+            torch.ones(shape, dtype=dtype, device=device_type)
+            for (shape, dtype) in inps
         ]
         mod = make_fx(forward)(*inps)
         compiled = compile_fx_inner(mod, inps)
@@ -167,8 +178,8 @@ def forward(self, x: torch.Tensor):
                 return {"ten0": out0, "ten1": out1}
 
         torch.manual_seed(0)
-        model = ReproModule().cuda()
-        inputs = torch.randn(36, 9, 7, 16, device="cuda", requires_grad=True)
+        model = ReproModule().to(device_type)
+        inputs = torch.randn(36, 9, 7, 16, device=device_type, requires_grad=True)
 
         eager_out = model(inputs)
         compiled_model = torch.compile(
@@ -182,6 +193,7 @@ def forward(self, x: torch.Tensor):
         self.assertEqual(compiled_out["ten0"], eager_out["ten0"])
         self.assertEqual(compiled_out["ten1"], eager_out["ten1"])
 
+    @skipIfXpu(msg="RuntimeError, torch-xpu-ops: 2891")
     def test_effn_attn_bias_padding(self):
         batch_size, num_heads, seq_len, head_dim = 2, 32, 512, 128
 
@@ -207,15 +219,19 @@ def fn(
                 scale=None,
             )
 
-        query = torch.randn(batch_size, num_heads, seq_len, head_dim, device="cuda")
-        key = torch.randn(batch_size, num_heads, seq_len, head_dim, device="cuda")
-        value = torch.randn(batch_size, num_heads, seq_len, head_dim, device="cuda")
+        query = torch.randn(
+            batch_size, num_heads, seq_len, head_dim, device=device_type
+        )
+        key = torch.randn(batch_size, num_heads, seq_len, head_dim, device=device_type)
+        value = torch.randn(
+            batch_size, num_heads, seq_len, head_dim, device=device_type
+        )
 
-        input_tensor = torch.rand([2, 1, seq_len, 1], device="cuda")
+        input_tensor = torch.rand([2, 1, seq_len, 1], device=device_type)
 
         out, code = run_and_get_code(torch.compile(fn), query, key, value, input_tensor)
 
-        input_tensor2 = torch.rand([2, 32, seq_len, seq_len], device="cuda").copy_(
+        input_tensor2 = torch.rand([2, 32, seq_len, seq_len], device=device_type).copy_(
             input_tensor
         )
         # even though the last dim is broadcasted, needs stride 1 for alignment
@@ -230,6 +246,7 @@ def fn(
     # Greatest absolute difference: 0.07861328125 at index (14, 13, 1008, 36) (up to 1e-05 allowed)
     # Greatest relative difference: 2.90625 at index (14, 13, 1008, 36) (up to 0.016 allowed)
     @skipIfRocmArch(MI350_ARCH)
+    @skipIfXpu(msg="RuntimeError, torch-xpu-ops: 2697")
     def test_effn_attn_bias_padding_misaligned(self):
         seqlen_start = 1008
 
@@ -238,10 +255,18 @@ def test_effn_attn_bias_padding_misaligned(self):
             torch._dynamo.reset()
 
             bsz = 32
-            q = torch.randn(bsz, 16, seqlen, 64, dtype=torch.bfloat16, device="cuda")
-            k = torch.randn(bsz, 16, seqlen, 64, dtype=torch.bfloat16, device="cuda")
-            v = torch.randn(bsz, 16, seqlen, 64, dtype=torch.bfloat16, device="cuda")
-            mask = torch.ones([bsz, 1, seqlen, seqlen], dtype=torch.bool, device="cuda")
+            q = torch.randn(
+                bsz, 16, seqlen, 64, dtype=torch.bfloat16, device=device_type
+            )
+            k = torch.randn(
+                bsz, 16, seqlen, 64, dtype=torch.bfloat16, device=device_type
+            )
+            v = torch.randn(
+                bsz, 16, seqlen, 64, dtype=torch.bfloat16, device=device_type
+            )
+            mask = torch.ones(
+                [bsz, 1, seqlen, seqlen], dtype=torch.bool, device=device_type
+            )
             inputs = [q, k, v, mask]
 
             def f(q, k, v, mask):
@@ -266,8 +291,12 @@ def test_input_channels_last(self):
         m = torch.nn.Sequential(
             torch.nn.Conv2d(3, 3, 1, 1),
             ToTuple(),
-        ).cuda()
-        inp = torch.randn([2, 3, 16, 16]).to(memory_format=torch.channels_last).cuda()
+        ).to(device_type)
+        inp = (
+            torch.randn([2, 3, 16, 16])
+            .to(memory_format=torch.channels_last)
+            .to(device_type)
+        )
 
         self.common(
             m,
@@ -291,10 +320,12 @@ def forward(self, x, y):
                 return [permute, add]
 
         inps = [
-            rand_strided((12, 3, 512, 64), (64, 196608, 768, 1), torch.float32, "cuda"),
+            rand_strided(
+                (12, 3, 512, 64), (64, 196608, 768, 1), torch.float32, device_type
+            ),
             rand_strided((), (), torch.int64, "cpu"),
         ]
-        mod = make_fx(Repro().to(device="cuda"))(*inps)
+        mod = make_fx(Repro().to(device=device_type))(*inps)
         compiled = compile_fx_inner(mod, inps)
         compiled(inps)
 
@@ -305,7 +336,7 @@ def test_backward_context(self):
         def fn(x):
             return x * 3
 
-        x = torch.randn(4, device="cuda", requires_grad=True)
+        x = torch.randn(4, device=device_type, requires_grad=True)
         gO = torch.rand_like(x)
         opt_fn = torch.compile(fn)
         out = opt_fn(x)
@@ -317,7 +348,7 @@ def forward():
             randn = torch.ops.aten.randn.default(
                 [12, 64, 1, 64],
                 dtype=torch.float32,
-                device=torch.device(type="cuda", index=0),
+                device=torch.device(type=device_type, index=0),
                 pin_memory=False,
             )
             unsqueeze_default_2 = torch.ops.aten.unsqueeze.default(randn, -1)
@@ -325,7 +356,10 @@ def forward():
 
         mod = make_fx(forward)()
         compiled = compile_fx_inner(mod, ())
-        assert compiled([])[0].device.type == "cuda"
+        if compiled([])[0].device.type != device_type:
+            raise AssertionError(
+                f"Expected device type {device_type}, got {compiled([])[0].device.type!r}"
+            )
 
     @config.patch({"triton.cudagraphs": True})
     @dynamo_config.patch(automatic_dynamic_shapes=True)
@@ -340,7 +374,7 @@ def forward(self):
                     1,
                     dtype=torch.float32,
                     layout=torch.strided,
-                    device=torch.device(type="cuda", index=0),
+                    device=torch.device(type=device_type, index=0),
                     pin_memory=False,
                 )
                 full_1 = torch.ops.aten.full.default(
@@ -348,7 +382,7 @@ def forward(self):
                     0,
                     dtype=torch.int64,
                     layout=torch.strided,
-                    device=torch.device(type="cuda", index=0),
+                    device=torch.device(type=device_type, index=0),
                     pin_memory=False,
                 )
                 return (full_1, full)
@@ -363,8 +397,8 @@ def fn(x, y):
             return x + y
 
         inputs = (
-            rand_strided((5, 5, 5, 5), (0, 5, 0, 1), device="cuda"),
-            rand_strided((5, 5, 5, 5), (0, 5, 0, 1), device="cuda"),
+            rand_strided((5, 5, 5, 5), (0, 5, 0, 1), device=device_type),
+            rand_strided((5, 5, 5, 5), (0, 5, 0, 1), device=device_type),
         )
         self.assertTrue(same(fn(*inputs), inputs[0] + inputs[1]))
 
@@ -383,14 +417,14 @@ def fn(x, y):
                     return r, r.size(0)
 
                 inputs = (
-                    torch.randn((5, 5), device="cuda"),
-                    torch.randn((5, 5), device="cuda"),
+                    torch.randn((5, 5), device=device_type),
+                    torch.randn((5, 5), device=device_type),
                 )
                 self.assertTrue(same(fn(*inputs), (inputs[0] + inputs[1], 5)))
 
                 inputs = (
-                    torch.randn((6, 6), device="cuda"),
-                    torch.randn((6, 6), device="cuda"),
+                    torch.randn((6, 6), device=device_type),
+                    torch.randn((6, 6), device=device_type),
                 )
                 self.assertTrue(same(fn(*inputs), (inputs[0] + inputs[1], 6)))
 
@@ -409,17 +443,59 @@ def max(x):
             print(f"compile {ms_c=:.03f}, eager {ms_eager=:.03f}")
 
     def test_split_reduction_transposed(self):
-        x = torch.randn(4096, 8192, dtype=torch.bfloat16, device="cuda")
+        x = torch.randn(4096, 8192, dtype=torch.bfloat16, device=device_type)
         x = x.t().contiguous().t()
 
         self._test_split_reduction_impl(x)
 
     def test_split_reduction_channels_last(self):
-        x = torch.randn(4096, 8192, dtype=torch.bfloat16, device="cuda")
+        x = torch.randn(4096, 8192, dtype=torch.bfloat16, device=device_type)
         x = x.reshape([256, 256, 256, 2]).to(memory_format=torch.channels_last)
 
         self._test_split_reduction_impl(x)
 
+    def test_split_with_sizes_reshape_cat_cantsplit_regression(self):
+        class Repro(nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.embedding = nn.Embedding(num_embeddings=128, embedding_dim=32)
+                self.linear = nn.Linear(32, 96)
+                self.conv = nn.Conv2d(3, 16, 3, padding=1)
+
+            def forward(self, x, indices):
+                embedded = self.embedding(indices)
+                linear_out = self.linear(embedded)
+                conv_out = self.conv(x)
+                batch_size = conv_out.shape[0]
+                conv_flat = conv_out.view(batch_size, -1)
+                seq_out = linear_out[:, -1, :]
+                combined = torch.cat([seq_out, conv_flat], dim=1)
+                chunks = torch.ops.aten.split_with_sizes.default(
+                    combined, split_sizes=[32, 64, combined.size(1) - 96], dim=1
+                )
+                chunk0_reshaped = torch.ops.aten.reshape.default(
+                    chunks[0], (batch_size, 4, 8)
+                )
+                chunk1_reshaped = torch.ops.aten.reshape.default(
+                    chunks[1], (batch_size, 8, 8)
+                )
+                chunk2_reshaped = torch.ops.aten.reshape.default(
+                    chunks[2], (batch_size, -1, 8)
+                )
+                return torch.ops.aten.cat.default(
+                    [chunk0_reshaped, chunk1_reshaped, chunk2_reshaped], dim=1
+                )
+
+        model = Repro().to(device_type)
+        x = torch.randn(2, 3, 32, 32, dtype=torch.float32, device=device_type)
+        indices = torch.randint(0, 128, (2, 10), dtype=torch.long, device=device_type)
+
+        with torch.no_grad():
+            eager_out = model(x, indices)
+            compiled_out = torch.compile(model)(x, indices)
+
+        torch.testing.assert_close(compiled_out, eager_out)
+
     @config.patch({"emulate_precision_casts": True})
     def test_bool_emulate_low_precision(self):
         from torch import device
@@ -436,7 +512,7 @@ def forward():
                 pin_memory=False,
             )
             device_put_3 = torch.ops.prims.device_put.default(
-                full_1, device(type="cuda", index=0)
+                full_1, device(type=device_type, index=0)
             )
             full_1 = None
 
@@ -457,13 +533,13 @@ def forward():
             view_15 = torch.ops.aten.reshape.default(clone, [1536, 1536])
             clone = None
             scalar_tensor = torch.ops.aten.scalar_tensor.default(
-                -inf, dtype=torch.float16, device=device(type="cuda", index=0)
+                -inf, dtype=torch.float16, device=device(type=device_type, index=0)
             )
             scalar_tensor_1 = torch.ops.aten.scalar_tensor.default(
                 0.0,
                 dtype=torch.float16,
                 layout=torch.strided,
-                device=device(type="cuda", index=0),
+                device=device(type=device_type, index=0),
             )
             where = torch.ops.aten.where.self(view_15, scalar_tensor_1, scalar_tensor)
             view_15 = scalar_tensor_1 = scalar_tensor = None
@@ -479,7 +555,9 @@ def test_emulate_low_precision(self):
         def foo(x):
             return torch.nn.functional.gelu(x) * 10.0
 
-        inp = torch.rand([32], device="cuda", requires_grad=True, dtype=torch.bfloat16)
+        inp = torch.rand(
+            [32], device=device_type, requires_grad=True, dtype=torch.bfloat16
+        )
         out, codes = run_fw_bw_and_get_code(lambda: torch.compile(foo)(inp))
 
         # fwd, backward
@@ -522,8 +600,8 @@ def fn(x, y):
             return x + y
 
         inputs = (
-            rand_strided((5, 5, 5, 5), (0, 5, 0, 1), device="cuda"),
-            rand_strided((5, 5, 5, 5), (0, 5, 0, 1), device="cuda"),
+            rand_strided((5, 5, 5, 5), (0, 5, 0, 1), device=device_type),
+            rand_strided((5, 5, 5, 5), (0, 5, 0, 1), device=device_type),
         )
         self.assertTrue(same(fn(*inputs), inputs[0] + inputs[1]))
 
@@ -544,11 +622,11 @@ def forward(self, x):
 
         from copy import deepcopy
 
-        model = Repro().cuda()
+        model = Repro().to(device_type)
         model_ref = deepcopy(model)
         model_opt = torch.compile(model, backend="inductor")
 
-        input = torch.randn(10, 10, device="cuda", requires_grad=True)
+        input = torch.randn(10, 10, device=device_type, requires_grad=True)
 
         for _ in range(2):
             output_ref = model_ref(input)
@@ -571,7 +649,7 @@ def foo(x):
 
         foo_opt = torch.compile(foo, backend="inductor")
 
-        inpt = torch.randn(10, 10, device="cuda", requires_grad=True)
+        inpt = torch.randn(10, 10, device=device_type, requires_grad=True)
         # TODO: this is broken, fix later
         # out = foo_opt(inpt)
         # out.add_(2)
@@ -599,21 +677,22 @@ def forward(self, start_positions: torch.Tensor, x: torch.Tensor):
                 )
                 return cross_entropy
 
-        mod = Repro().cuda()
+        mod = Repro().to(device_type)
         opt_mod = torch.compile(mod, backend="inductor")
         mod.eval()
         opt_mod.eval()
 
         args = [
-            ((1,), (1,), torch.int64, "cuda", False),
-            ((1, 128, 768), (98304, 768, 1), torch.float32, "cuda", True),
+            ((1,), (1,), torch.int64, device_type, False),
+            ((1, 128, 768), (98304, 768, 1), torch.float32, device_type, True),
         ]
         args = [
             rand_strided(sh, st, dt, dev).requires_grad_(rg)
             for (sh, st, dt, dev, rg) in args
         ]
         with torch.cuda.amp.autocast(enabled=False):
-            assert same_two_models(mod, opt_mod, args), "Dynamo failed"
+            if not same_two_models(mod, opt_mod, args):
+                raise AssertionError("Dynamo failed")
 
     @config.patch(allow_buffer_reuse=False)
     def test_issue103461(self):
@@ -624,7 +703,7 @@ def forward(add_1):
             getitem_1 = var_mean[1]
             return getitem_1
 
-        x = torch.randn(1, 8, 768, device="cuda")
+        x = torch.randn(1, 8, 768, device=device_type)
         correct = forward(x)
         actual = torch.compile(forward, fullgraph=True)(x)
         self.assertEqual(actual, correct)
@@ -636,7 +715,7 @@ def forward(x):
                 0,
                 dtype=torch.float64,
                 layout=torch.strided,
-                device="cuda",
+                device=device_type,
                 pin_memory=False,
             )
             return x + full_10.to("cpu")
@@ -653,7 +732,11 @@ def test_autotune_inplace_kernel(self):
         https://github.com/triton-lang/triton/issues/781
         https://github.com/pytorch/torchdynamo/issues/1670
         """
-        from torch._C import _cuda_getCurrentRawStream as get_cuda_stream
+        if TEST_XPU:
+            from torch._C import _xpu_getCurrentRawStream as get_gpu_stream
+        else:
+            from torch._C import _cuda_getCurrentRawStream as get_gpu_stream
+
         from torch._inductor.runtime.hints import AttrsDescriptorWrapper, HeuristicType
         from torch._inductor.runtime.triton_heuristics import CachingAutotuner
         from torch._inductor.utils import triton_version_uses_attrs_dict
@@ -691,7 +774,7 @@ def decorator(fn):
                     "in_ptr0": "*fp32",
                     "xnumel": "i32",
                 },
-                "device": DeviceProperties.create(torch.device("cuda")),
+                "device": DeviceProperties.create(torch.device(device_type)),
                 "configs": [
                     AttrsDescriptorWrapper(divisible_by_16=(0, 1), equal_to_1=())
                 ],
@@ -710,17 +793,16 @@ def kernel(in_out_ptr0, in_ptr0, xnumel, XBLOCK: tl.constexpr):
             tl.store(in_out_ptr0 + offsets, output, mask=mask)
 
         xnumel = 384
-        in0 = rand_strided((xnumel,), (1,), device="cuda", dtype=torch.float32)
-        inout1 = rand_strided((xnumel,), (1,), device="cuda", dtype=torch.float32)
+        in0 = rand_strided((xnumel,), (1,), device=device_type, dtype=torch.float32)
+        inout1 = rand_strided((xnumel,), (1,), device=device_type, dtype=torch.float32)
         inout2 = inout1.clone()
 
-        stream0 = get_cuda_stream(0)
+        stream0 = get_gpu_stream(0)
         kernel.run(inout1, in0, xnumel, stream=stream0)
         kernel.run(inout2, in0, xnumel, stream=stream0)
 
-        assert same(inout1, inout2, tol=0.001, equal_nan=True), (
-            "failed autotune with inplace kernel"
-        )
+        if not same(inout1, inout2, tol=0.001, equal_nan=True):
+            raise AssertionError("failed autotune with inplace kernel")
 
     def test_sort_stride_issue(self):
         # This minified testcase comes from detectron2_maskrcnn_r_50_fpn
@@ -731,13 +813,14 @@ def forward(pred_objectness_logits_3_: torch.Tensor):
             getitem_12 = sort_3[0]
             return getitem_12
 
-        args = [((1, 100), (0, 1), torch.float16, "cuda", False)]
+        args = [((1, 100), (0, 1), torch.float16, device_type, False)]
         args = [
             rand_strided(sh, st, dt, dev).requires_grad_(rg)
             for (sh, st, dt, dev, rg) in args
         ]
         result = forward(*args)
-        assert same(result, torch.sort(args[0], descending=True, dim=1)[0])
+        if not same(result, torch.sort(args[0], descending=True, dim=1)[0]):
+            raise AssertionError
 
     def test_scalar_triton_index(self):
         # The indirect indexing via a scalar like below used to lead to
@@ -747,10 +830,11 @@ def fn(a):
             zero = torch.zeros((16,), device=a.device, dtype=torch.int64)
             return (a[zero],)
 
-        a = torch.randn((8,), dtype=torch.float32, device="cuda")
+        a = torch.randn((8,), dtype=torch.float32, device=device_type)
 
         fn_optimized = torch.compile(fn, backend="inductor")
-        assert same(fn(a), fn_optimized(a))
+        if not same(fn(a), fn_optimized(a)):
+            raise AssertionError
 
     def test_indirect_indexing_dense_mask(self):
         def fn(x, y):
@@ -763,17 +847,20 @@ def fn(x, y):
             out = torch.ops.aten.multiply(y, squeeze)
             return (out,)
 
-        a = torch.zeros((1, 128), dtype=torch.int64, device="cuda")
-        b = torch.zeros((1, 128), dtype=torch.int64, device="cuda")
+        a = torch.zeros((1, 128), dtype=torch.int64, device=device_type)
+        b = torch.zeros((1, 128), dtype=torch.int64, device=device_type)
 
         fn_optimized = torch.compile(fn, backend="inductor")
-        assert same(fn(a, b), fn_optimized(a, b))
+        if not same(fn(a, b), fn_optimized(a, b)):
+            raise AssertionError
 
     def test_simplify_dims(self):
         def fn(a):
             return (a + 1,)
 
-        self.common(fn, (torch.randn(2, 3, 10, 5, 6, device="cuda")[:, :, 2::2, :, :],))
+        self.common(
+            fn, (torch.randn(2, 3, 10, 5, 6, device=device_type)[:, :, 2::2, :, :],)
+        )
 
     @config.patch(permute_fusion=True)
     def test_permute_fusion(self):
@@ -786,8 +873,8 @@ def forward(self, view, reshape_2):
                 return (bmm,)
 
         args = [
-            ((1024, 642, 160), (102720, 160, 1), torch.float32, "cuda", True),
-            ((1024, 642, 20), (12840, 20, 1), torch.float32, "cuda", True),
+            ((1024, 642, 160), (102720, 160, 1), torch.float32, device_type, True),
+            ((1024, 642, 20), (12840, 20, 1), torch.float32, device_type, True),
         ]
         args = [
             rand_strided(sh, st, dt, dev).requires_grad_(rg)
@@ -807,17 +894,18 @@ def fn(x, y):
             aten.add_.Tensor(x, y, alpha=0.55)
             return (x,)
 
-        x1 = torch.zeros(2, 3, 4, 10, device="cuda")
-        x2 = torch.zeros(2, 3, 4, 10, device="cuda")
-        x3 = torch.zeros(2, 3, 4, 10, device="cuda")
-        y = torch.randn(2, 3, 4, 10, device="cuda").to(
+        x1 = torch.zeros(2, 3, 4, 10, device=device_type)
+        x2 = torch.zeros(2, 3, 4, 10, device=device_type)
+        x3 = torch.zeros(2, 3, 4, 10, device=device_type)
+        y = torch.randn(2, 3, 4, 10, device=device_type).to(
             memory_format=torch.channels_last
         )
         fn_fx = make_fx(fn)(x1, y)
         fn_compiled = compile_fx_inner(fn_fx, [x1, y])
         fn(x2, y)
         fn_compiled([x3, y])
-        assert same(x2, x3)
+        if not same(x2, x3):
+            raise AssertionError
 
     @config.patch({"triton.autotune_pointwise": True})
     def test_inplace_buffer_autotune(self):
@@ -825,15 +913,18 @@ def foo(x, y, z):
             a = x @ y
             return a.unsqueeze(0).unsqueeze(0) + z
 
-        x = torch.zeros(5, 5, device="cuda")
-        y = torch.zeros(5, 5, device="cuda")
-        z = torch.zeros(1, 1, 5, 5, device="cuda").to(memory_format=torch.channels_last)
+        x = torch.zeros(5, 5, device=device_type)
+        y = torch.zeros(5, 5, device=device_type)
+        z = torch.zeros(1, 1, 5, 5, device=device_type).to(
+            memory_format=torch.channels_last
+        )
         self.common(
             foo,
             (x, y, z),
             check_lowp=False,
         )
 
+    @skipIfXpu(msg="TypeError, torch-xpu-ops: 3004")
     def test_memory_history_inductor(self):
         def called_inside_compile(x, w, b):
             a = x @ w + b
@@ -844,16 +935,26 @@ def fn(x, w, b):
             x = called_inside_compile(x, w, b)
             return called_inside_compile(x, w, b)
 
-        w = torch.rand(3, 3, device="cuda")
-        b = torch.rand(3, device="cuda")
-        x = torch.rand(3, device="cuda")
+        w = torch.rand(3, 3, device=device_type)
+        b = torch.rand(3, device=device_type)
+        x = torch.rand(3, device=device_type)
+
+        def record_memory_history(value: bool):
+            if torch.xpu.is_available():
+                torch.xpu.memory._record_memory_history(value)
+            else:
+                torch.cuda.memory._record_memory_history(value)
+
         try:
-            torch.cuda.memory.empty_cache()
-            torch.cuda.memory._record_memory_history(True)
+            torch.accelerator.memory.empty_cache()
+            record_memory_history(True)
             r = fn(x, w, b)
         finally:
-            torch.cuda.memory._record_memory_history(False)
-        snapshot = str(torch.cuda.memory._snapshot())
+            record_memory_history(False)
+        if torch.xpu.is_available():
+            snapshot = str(torch.xpu.memory._snapshot())
+        else:
+            snapshot = str(torch.cuda.memory._snapshot())
         self.assertTrue("called_inside_compile" in snapshot)
 
     def test_negative_arange_dynamic_shapes(self):
@@ -886,10 +987,14 @@ def forward(self, enc_out: torch.Tensor, dec_in: torch.Tensor):
                 padmask = dec_in == 0
                 dec_mask = padmask.unsqueeze(-1) == padmask.unsqueeze(-2)
                 dec_mask = dec_mask.to(dtype=torch.float32)
-                dec_mask = dec_mask.tril(diagonal=0).cuda()
+                dec_mask = dec_mask.tril(diagonal=0).to(device_type)
 
-                q_pos = torch.arange(dec_in.size(1), dtype=torch.long, device="cuda")
-                k_pos = torch.arange(dec_in.size(1), dtype=torch.long, device="cuda")
+                q_pos = torch.arange(
+                    dec_in.size(1), dtype=torch.long, device=device_type
+                )
+                k_pos = torch.arange(
+                    dec_in.size(1), dtype=torch.long, device=device_type
+                )
                 rel_pos = k_pos[None, :] - q_pos[:, None]
                 values = rel_pos.abs().neg().unsqueeze(0).unsqueeze(0)
                 dec_bias = values * self.scales
@@ -900,20 +1005,20 @@ def forward(self, enc_out: torch.Tensor, dec_in: torch.Tensor):
                 out = self.dec_layer(out, enc_out, tgt_mask=dec_mask)
                 return self.head(out)
 
-        mod = Repro().cuda()
+        mod = Repro().to(device_type)
         opt_mod = torch.compile(mod, backend="inductor", dynamic=True)
         mod.eval()
         opt_mod.eval()
 
-        enc_out = torch.rand(1, 512, 256).cuda()
+        enc_out = torch.rand(1, 512, 256).to(device_type)
         dec_inputs = [
-            torch.randint(0, 512, (1, i + 1), dtype=torch.long).cuda() for i in range(8)
+            torch.randint(0, 512, (1, i + 1), dtype=torch.long).to(device_type)
+            for i in range(8)
         ]
 
         for dec_inp in dec_inputs:
-            assert same_two_models(mod, opt_mod, [enc_out, dec_inp], only_fwd=True), (
-                "Inductor with dynamic shapes failed"
-            )
+            if not same_two_models(mod, opt_mod, [enc_out, dec_inp], only_fwd=True):
+                raise AssertionError("Inductor with dynamic shapes failed")
 
     def test_issue97695_1input(self):
         def fn(arg3_1, relu, permute_1):
@@ -922,9 +1027,9 @@ def fn(arg3_1, relu, permute_1):
             return (cat_2,)
 
         args = [
-            ((96,), (1,), torch.float32, "cuda"),
-            ((10, 256), (256, 1), torch.float32, "cuda"),
-            ((256, 96), (1, 256), torch.float32, "cuda"),
+            ((96,), (1,), torch.float32, device_type),
+            ((10, 256), (256, 1), torch.float32, device_type),
+            ((256, 96), (1, 256), torch.float32, device_type),
         ]
         args = [rand_strided(sh, st, dt, dev) for (sh, st, dt, dev) in args]
         correct = fn(*args)
@@ -932,10 +1037,12 @@ def fn(arg3_1, relu, permute_1):
         mod = make_fx(fn, tracing_mode="real")(*args)
         compiled = compile_fx_inner(mod, args)
         ref = compiled(list(args))
-        assert same(ref, correct)
+        if not same(ref, correct):
+            raise AssertionError
 
         ref = torch.compile(fn, fullgraph=True)(*args)
-        assert same(ref, correct)
+        if not same(ref, correct):
+            raise AssertionError
 
     def test_issue_103924(self):
         class MyModule(torch.nn.Module):
@@ -950,7 +1057,7 @@ def forward(self, x):
                 inp = x / y[..., None]
                 return self.layer(inp)
 
-        x = torch.rand([4, 4], device="cuda")
+        x = torch.rand([4, 4], device=device_type)
         m = MyModule()
         opt_m = torch.compile(backend="inductor")(m)
         self.assertEqual(opt_m(x), m(x))
@@ -963,16 +1070,17 @@ def fn(arg3_1, arg3_2, relu, permute_1):
             return (cat_2,)
 
         args = [
-            ((96,), (1,), torch.float32, "cuda"),
-            ((96,), (1,), torch.float32, "cuda"),
-            ((10, 256), (256, 1), torch.float32, "cuda"),
-            ((256, 96), (1, 256), torch.float32, "cuda"),
+            ((96,), (1,), torch.float32, device_type),
+            ((96,), (1,), torch.float32, device_type),
+            ((10, 256), (256, 1), torch.float32, device_type),
+            ((256, 96), (1, 256), torch.float32, device_type),
         ]
         args = [rand_strided(sh, st, dt, dev) for (sh, st, dt, dev) in args]
         correct = fn(*args)
 
         ref = torch.compile(fn, fullgraph=True)(*args)
-        assert same(ref, correct)
+        if not same(ref, correct):
+            raise AssertionError
 
     def test_scatter_index_not_wrapped(self):
         src = torch.tensor([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], device=self.device)
@@ -994,7 +1102,9 @@ def test_normalize_norm_leq_one(self):
         def fn(x: torch.Tensor) -> torch.Tensor:
             return torch.nn.functional.normalize(x, dim=-1)
 
-        inp = torch.tensor([[3.799999, 0.0, 0.0]], device="cuda", dtype=torch.float32)
+        inp = torch.tensor(
+            [[3.799999, 0.0, 0.0]], device=device_type, dtype=torch.float32
+        )
         compiled = torch.compile(fn, backend="inductor", fullgraph=True)
         out = compiled(inp)
         norm = out.norm(dim=-1)
@@ -1006,7 +1116,7 @@ def test_libdevice_routing(self):
         def foo(x):
             return x.exp()
 
-        inp = torch.ones(64, device="cuda").to(torch.float64)
+        inp = torch.ones(64, device=device_type).to(torch.float64)
 
         out, code = run_and_get_code(torch.compile(foo), inp)
         FileCheck().check("libdevice.exp").run(code[0])
@@ -1020,7 +1130,7 @@ def foo(x):
         def foo(x):
             return x.sigmoid()
 
-        inp = torch.ones(64, device="cuda").to(torch.float64)
+        inp = torch.ones(64, device=device_type).to(torch.float64)
         out, code = run_and_get_code(torch.compile(foo), inp)
         FileCheck().check("libdevice.exp").run(code[0])
         self.assertEqual(foo(inp), out)
@@ -1028,11 +1138,11 @@ def foo(x):
     def test_uint_view_copy(self):
         @torch.compile
         def view_copy(target, source):
-            assert target.dtype == torch.bfloat16
-            assert source.dtype == torch.uint16
+            assert target.dtype == torch.bfloat16  # noqa: S101
+            assert source.dtype == torch.uint16  # noqa: S101
             target.view(torch.uint16).copy_(source)
 
-        target = torch.ones(1024, dtype=torch.bfloat16, device="cuda")
+        target = torch.ones(1024, dtype=torch.bfloat16, device=device_type)
         source = torch.full_like(target, 4, dtype=torch.uint16)
 
         out = target.view(torch.uint16).copy_(source).clone()
@@ -1046,7 +1156,7 @@ def forward(arg0_1):
                 1,
                 dtype=torch.float32,
                 layout=torch.strided,
-                device=torch.device(type="cuda", index=0),
+                device=torch.device(type=device_type, index=0),
                 pin_memory=False,
             )
             convert_element_type_1 = torch.ops.prims.convert_element_type.default(
@@ -1064,23 +1174,23 @@ def forward(arg0_1):
             )
             return [var_mean[0], var_mean[1], add_2]
 
-        emb = torch.randn([2050, 768], device="cuda")
+        emb = torch.randn([2050, 768], device=device_type)
         gm = make_fx(forward)(emb)
         opt = torch._inductor.compile_fx.compile_fx_inner(gm, [emb])
         opt([emb])
-        torch.cuda.synchronize()
+        torch.accelerator.synchronize()
 
     def test_deterministic_algorithms(self):
         N = 10000
 
         @torch.compile
         def fn(idx, values):
-            x = torch.zeros(1, device="cuda")
+            x = torch.zeros(1, device=device_type)
             x[idx] += values
             return x
 
-        idx = torch.zeros(N, dtype=torch.int64, device="cuda")
-        values = torch.randn(N, device="cuda")
+        idx = torch.zeros(N, dtype=torch.int64, device=device_type)
+        values = torch.randn(N, device=device_type)
 
         r0 = fn(idx, values)
         with DeterministicGuard(True):
@@ -1097,10 +1207,10 @@ def __init__(self) -> None:
                 self.linear = nn.Linear(4, 4)
 
             def forward(self, data):
-                data = data.to("cuda")
+                data = data.to(device_type)
                 return self.linear(data)
 
-        mod = Model().cuda().eval()
+        mod = Model().to(device_type).eval()
         with torch.no_grad():
             self.common(mod, (torch.randn(4, 4),))
 
@@ -1116,7 +1226,7 @@ def forward(self, x):
                 return self.dropout(y)
 
         mod = Repro()
-        x = torch.randn((512, 1, 4096), requires_grad=True, device="cuda")
+        x = torch.randn((512, 1, 4096), requires_grad=True, device=device_type)
         y = torch.compile(mod)(x)
         # Inductor claims the output layout of gelu's saved variable for
         # backwards will be (4096, 4096, 1) but in actuality it is (4096,
@@ -1143,9 +1253,9 @@ def forward(inductor_seeds, mul_4, view_15):
             sub_3 = torch.ops.aten.sub.Tensor(add_5, getitem_3)
             return (sub_3,)
 
-        buf0 = torch.zeros((37,), dtype=torch.int64, device="cuda")
-        buf1 = torch.zeros((2, 512, 768), device="cuda")
-        buf2 = torch.zeros((2, 512, 768), device="cuda")
+        buf0 = torch.zeros((37,), dtype=torch.int64, device=device_type)
+        buf1 = torch.zeros((2, 512, 768), device=device_type)
+        buf2 = torch.zeros((2, 512, 768), device=device_type)
         forward(buf0, buf1, buf2)
 
     def test_issue100806(self):
@@ -1165,7 +1275,7 @@ def forward(self, x):
                 x = self.relu(x)
                 return x
 
-        device = "cuda"
+        device = device_type
         batch_size = 2
         x = torch.randn(batch_size, 10).to(device)
         func = Model().to(device)
@@ -1185,8 +1295,8 @@ def fn(x, y):
             add = mean + y
             return add
 
-        x = torch.rand(4, 4, 4, 4, 4, 4, device="cuda")
-        y = torch.rand((), device="cuda")
+        x = torch.rand(4, 4, 4, 4, 4, 4, device=device_type)
+        y = torch.rand((), device=device_type)
         expect = fn(x, y)
 
         opt_fn = torch.compile(fn)
@@ -1205,8 +1315,8 @@ def test_bucketize_dynamic_dense(self):
         def fn(values, offsets):
             return torch.bucketize(values, offsets)
 
-        values = torch.rand((64, 64), device="cuda")
-        offsets = torch.tensor([0.05, 0.1, 0.5, 0.8, 0.85, 0.95], device="cuda")
+        values = torch.rand((64, 64), device=device_type)
+        offsets = torch.tensor([0.05, 0.1, 0.5, 0.8, 0.85, 0.95], device=device_type)
 
         expect = fn(values, offsets)
 
@@ -1244,9 +1354,9 @@ def fn(x: torch.Tensor, y: torch.Tensor, buckets: torch.Tensor) -> torch.Tensor:
             z = torch.mm(x, y)
             return torch.bucketize(z, buckets)
 
-        buckets = torch.arange(-100, 100, 10, device="cuda")
-        x = torch.randn(64, 64, device="cuda").clamp(-99, 99)
-        y = torch.randn(64, 64, device="cuda").clamp(-99, 99)
+        buckets = torch.arange(-100, 100, 10, device=device_type)
+        x = torch.randn(64, 64, device=device_type).clamp(-99, 99)
+        y = torch.randn(64, 64, device=device_type).clamp(-99, 99)
 
         opt_fn = torch.compile(fn, mode="max-autotune")
 
@@ -1259,7 +1369,7 @@ def test_float64_constants(self):
         def fn():
             # NOTE: tensors of all the same value are constant folded, so we
             # need a tensor with two distinct values
-            a = torch.tensor([1 / 10, 2 / 10], dtype=torch.float64, device="cuda")
+            a = torch.tensor([1 / 10, 2 / 10], dtype=torch.float64, device=device_type)
             return a * 2e50
 
         cfn = torch.compile(fn)
@@ -1299,10 +1409,15 @@ def fn(arg7_1, add_1, permute_2, select_scatter, slice_8):
             return (bmm,)
 
         args = []
-        args.append(torch.randn((2, 1, 4, 1200, 4), dtype=torch.float16, device="cuda"))
+        args.append(
+            torch.randn((2, 1, 4, 1200, 4), dtype=torch.float16, device=device_type)
+        )
         args.append(
             rand_strided(
-                (1, 4, 1000, 4), (16000, 4, 16, 1), dtype=torch.float16, device="cuda"
+                (1, 4, 1000, 4),
+                (16000, 4, 16, 1),
+                dtype=torch.float16,
+                device=device_type,
             )
         )
         args.append(
@@ -1310,7 +1425,7 @@ def fn(arg7_1, add_1, permute_2, select_scatter, slice_8):
                 (3, 1, 4, 1000, 4),
                 (16, 48000, 4, 48, 1),
                 dtype=torch.float16,
-                device="cuda",
+                device=device_type,
             )
         )
         args.append(
@@ -1318,7 +1433,7 @@ def fn(arg7_1, add_1, permute_2, select_scatter, slice_8):
                 (2, 1, 4, 1000, 4),
                 (16, 48000, 4, 48, 1),
                 dtype=torch.float16,
-                device="cuda",
+                device=device_type,
             )
         )
         args.append(
@@ -1326,7 +1441,7 @@ def fn(arg7_1, add_1, permute_2, select_scatter, slice_8):
                 (2, 1, 4, 1000, 4),
                 (19200, 19200, 4800, 4, 1),
                 dtype=torch.float16,
-                device="cuda",
+                device=device_type,
             )
         )
 
@@ -1334,7 +1449,8 @@ def fn(arg7_1, add_1, permute_2, select_scatter, slice_8):
         mod = make_fx(fn, tracing_mode="real")(*args)
         compiled = compile_fx_inner(mod, args)
         ref = compiled(list(args))
-        assert same(ref, correct)
+        if not same(ref, correct):
+            raise AssertionError
 
     @config.patch({"triton.cudagraphs": True})
     def test_index_put_inplace_cudagraph(self):
@@ -1342,9 +1458,9 @@ def fn(x, y, z):
             x = torch.zeros_like(x)
             return x.index_put_([y], z, True)
 
-        x = torch.zeros((512, 512), device="cuda", dtype=torch.bool)
-        y = torch.zeros((512,), device="cuda", dtype=torch.int64)
-        z = torch.ones((512, 512), device="cuda", dtype=torch.bool)
+        x = torch.zeros((512, 512), device=device_type, dtype=torch.bool)
+        y = torch.zeros((512,), device=device_type, dtype=torch.int64)
+        z = torch.ones((512, 512), device=device_type, dtype=torch.bool)
 
         opt_fn = torch.compile(fn, backend="inductor")
 
@@ -1365,9 +1481,9 @@ def fn(x, y, z):
                 x = torch.zeros_like(x)
                 return x.index_put([y], z, True)
 
-            x = torch.zeros((512, 512), device="cuda", dtype=torch.bool)
-            y = torch.zeros((512,), device="cuda", dtype=torch.int64)
-            z = torch.ones((512, 512), device="cuda", dtype=torch.bool)
+            x = torch.zeros((512, 512), device=device_type, dtype=torch.bool)
+            y = torch.zeros((512,), device=device_type, dtype=torch.int64)
+            z = torch.ones((512, 512), device=device_type, dtype=torch.bool)
 
             opt_fn = torch.compile(fn, backend="inductor")
 
@@ -1410,18 +1526,13 @@ def forward(self, x):
 
         cnts = torch._dynamo.testing.CompileCounterWithBackend("inductor")
 
-        model = Model().cuda().half()
+        model = Model().to(device_type).half()
         model = torch.compile(model, backend=cnts, dynamic=True)
 
-        with torch.backends.cuda.sdp_kernel(
-            enable_flash=True,
-            enable_math=False,
-            enable_mem_efficient=False,
-            enable_cudnn=False,
-        ):
-            input1 = torch.rand(5, 512, 1024, device="cuda", dtype=torch.float16)
-            input2 = torch.rand(5, 513, 1024, device="cuda", dtype=torch.float16)
-            input3 = torch.rand(5, 514, 1024, device="cuda", dtype=torch.float16)
+        with sdpa_kernel(SDPBackend.FLASH_ATTENTION):
+            input1 = torch.rand(5, 512, 1024, device=device_type, dtype=torch.float16)
+            input2 = torch.rand(5, 513, 1024, device=device_type, dtype=torch.float16)
+            input3 = torch.rand(5, 514, 1024, device=device_type, dtype=torch.float16)
 
             out1 = model(input1)
             out2 = model(input2)
@@ -1435,9 +1546,9 @@ def fn(x, y, z):
             x = torch.zeros_like(x)
             return x.index_put([y], z, True)
 
-        x = torch.zeros((512, 512), device="cuda", dtype=torch.int32)
-        y = torch.zeros((512,), device="cuda", dtype=torch.int64)
-        z = torch.ones((512, 512), device="cuda", dtype=torch.int32)
+        x = torch.zeros((512, 512), device=device_type, dtype=torch.int32)
+        y = torch.zeros((512,), device=device_type, dtype=torch.int64)
+        z = torch.ones((512, 512), device=device_type, dtype=torch.int32)
 
         opt_fn = torch.compile(fn, backend="inductor")
 
@@ -1449,13 +1560,32 @@ def fn(x, y, z):
 
         self.assertEqual(ref, res)
 
+    @torch._inductor.config.patch(emulate_precision_casts=True)
+    @torch._inductor.config.patch(pattern_matcher=False)
+    def test_emulate_precision_casts_convert_element_type(self):
+        torch.manual_seed(0)
+        torch.cuda.manual_seed_all(0)
+
+        x = torch.rand(1000, device=device_type, dtype=torch.float32)
+
+        def fn(x):
+            x_bf16 = x.to(torch.bfloat16)
+            x_fp32_casted = x_bf16.to(torch.float32)
+            return x_fp32_casted
+
+        opt_fn = torch.compile(fn, backend="inductor", fullgraph=True, dynamic=True)
+
+        expected = fn(x)
+        actual = opt_fn(x)
+        self.assertEqual(expected, actual)
+
     @torch._inductor.config.patch(emulate_precision_casts=True)
     def test_emulate_precision_casts_norm_rounding(self):
         torch.manual_seed(0)
         torch.cuda.manual_seed_all(0)
 
-        x = torch.rand(1000, device="cuda", dtype=torch.bfloat16)
-        scalar = torch.rand([], device="cuda", dtype=torch.float32)
+        x = torch.rand(1000, device=device_type, dtype=torch.bfloat16)
+        scalar = torch.rand([], device=device_type, dtype=torch.float32)
 
         def fn(inp, scale):
             y = inp.norm()
@@ -1468,10 +1598,11 @@ def fn(inp, scale):
 
         self.assertEqual(expected, actual)
 
+    @skipIfXpu(msg="AssertionError, torch-xpu-ops: #2554")
     @torch._inductor.config.patch(emulate_precision_casts=True)
     def test_emulate_precision_casts_min_pow_chain(self):
         torch.manual_seed(0)
-        torch.cuda.manual_seed_all(0)
+        torch.cuda.manual_seed_all(0) if TEST_CUDA else torch.xpu.manual_seed_all(0)
 
         with dynamo_config.patch(
             capture_scalar_outputs=True, capture_dynamic_output_shape_ops=True
@@ -1479,17 +1610,17 @@ def test_emulate_precision_casts_min_pow_chain(self):
             arg0 = torch.rand(
                 [383, 55, 2, 3],
                 dtype=torch.float16,
-                device="cuda",
+                device=device_type,
                 requires_grad=True,
             )
             arg1 = torch.rand(
-                [383, 55], dtype=torch.bfloat16, device="cuda", requires_grad=True
+                [383, 55], dtype=torch.bfloat16, device=device_type, requires_grad=True
             )
             arg2 = torch.rand(
-                [383, 55], dtype=torch.float32, device="cuda", requires_grad=True
+                [383, 55], dtype=torch.float32, device=device_type, requires_grad=True
             )
             arg3 = torch.rand(
-                [383, 55], dtype=torch.float32, device="cuda", requires_grad=True
+                [383, 55], dtype=torch.float32, device=device_type, requires_grad=True
             )
 
             def fn(a0, a1, a2, a3):
@@ -1528,22 +1659,37 @@ def test_emulate_precision_casts_mean_ratio_chain(self):
             capture_scalar_outputs=True, capture_dynamic_output_shape_ops=True
         ):
             arg0 = torch.rand(
-                [125070], dtype=torch.bfloat16, device="cuda", requires_grad=True
+                [125070], dtype=torch.bfloat16, device=device_type, requires_grad=True
             )
             arg1 = torch.rand(
-                [1895, 3, 11], dtype=torch.float16, device="cuda", requires_grad=True
+                [1895, 3, 11],
+                dtype=torch.float16,
+                device=device_type,
+                requires_grad=True,
             )
             arg2 = torch.rand(
-                [1895, 3, 11], dtype=torch.float32, device="cuda", requires_grad=True
+                [1895, 3, 11],
+                dtype=torch.float32,
+                device=device_type,
+                requires_grad=True,
             )
             arg3 = torch.rand(
-                [1895, 3, 11], dtype=torch.float32, device="cuda", requires_grad=True
+                [1895, 3, 11],
+                dtype=torch.float32,
+                device=device_type,
+                requires_grad=True,
             )
             arg4 = torch.rand(
-                [1895, 3, 11], dtype=torch.float32, device="cuda", requires_grad=True
+                [1895, 3, 11],
+                dtype=torch.float32,
+                device=device_type,
+                requires_grad=True,
             )
             arg5 = torch.rand(
-                [5, 379, 165], dtype=torch.float32, device="cuda", requires_grad=True
+                [5, 379, 165],
+                dtype=torch.float32,
+                device=device_type,
+                requires_grad=True,
             )
 
             def fn(a0, a1, a2, a3, a4, a5):
@@ -1575,14 +1721,14 @@ def fn(a0, a1, a2, a3, a4, a5):
     def test_dont_inplace_disjoint_accesses(self):
         # TODO - would not need mms if we could annotate donated buffer..
         def forward(  # noqa: F821, F722
-            arg0_1: "bf16[2048, 2048][2048, 1]cuda:0",  # noqa: F821, F722
-            arg1_1: "bf16[8, 4096, 2048][8388608, 2048, 1]cuda:0",  # noqa: F821, F722
-            arg2_1: "bf16[2048, 2048][2048, 1]cuda:0",  # noqa: F821, F722
-            arg3_1: "bf16[2048, 2048][2048, 1]cuda:0",  # noqa: F821, F722
-            arg4_1: "bf16[2048][1]cuda:0",  # noqa: F821, F722
-            arg5_1: "bf16[2048][1]cuda:0",  # noqa: F821, F722
-            arg6_1: "f32[4096, 128][128, 1]cuda:0",  # noqa: F821, F722
-            arg7_1: "f32[4096, 128][128, 1]cuda:0",  # noqa: F821, F722
+            arg0_1: f"bf16[2048, 2048][2048, 1]{device_type}:0",  # noqa: F821, F722
+            arg1_1: f"bf16[8, 4096, 2048][8388608, 2048, 1]{device_type}:0",  # noqa: F821, F722
+            arg2_1: f"bf16[2048, 2048][2048, 1]{device_type}:0",  # noqa: F821, F722
+            arg3_1: f"bf16[2048, 2048][2048, 1]{device_type}:0",  # noqa: F821, F722
+            arg4_1: f"bf16[2048][1]{device_type}:0",  # noqa: F821, F722
+            arg5_1: f"bf16[2048][1]{device_type}:0",  # noqa: F821, F722
+            arg6_1: f"f32[4096, 128][128, 1]{device_type}:0",  # noqa: F821, F722
+            arg7_1: f"f32[4096, 128][128, 1]{device_type}:0",  # noqa: F821, F722
         ):
             permute = torch.ops.aten.permute.default(arg0_1, [1, 0])
             arg0_1 = None
@@ -1744,7 +1890,7 @@ def forward(  # noqa: F821, F722
 
         from torch._dynamo.debug_utils import aot_graph_input_parser
 
-        kwargs = aot_graph_input_parser(forward)
+        kwargs = aot_graph_input_parser(forward, device=device_type)
         out, code = run_and_get_code(torch.compile(forward), **kwargs)
         # ignore tiny values.. prior to this fix absolute error was ~28
         self.assertEqual(forward(**kwargs), out, atol=0.01, rtol=2)
@@ -1752,8 +1898,8 @@ def forward(  # noqa: F821, F722
 
     # https://github.com/pytorch/pytorch/issues/104937
     def test_linear_with_zero_infeature_size(self):
-        m = nn.Linear(in_features=0, out_features=0, bias=True).to("cuda")
-        x = torch.rand(1, 1, 0, device="cuda")
+        m = nn.Linear(in_features=0, out_features=0, bias=True).to(device_type)
+        x = torch.rand(1, 1, 0, device=device_type)
         expect = m(x)
         opt_fn = torch.compile(m)
         actual = opt_fn(x)
@@ -1762,7 +1908,7 @@ def test_linear_with_zero_infeature_size(self):
     @config.patch(fallback_random=True)
     def test_multi_output_layout_fallback(self):
         mod = nn.RReLU(lower=3.2350976, upper=8.4220314, inplace=True)
-        inp = torch.rand([4, 4]).cuda()
+        inp = torch.rand([4, 4]).to(device_type)
         m = torch.compile(mod)
 
         with freeze_rng_state():
@@ -1777,8 +1923,8 @@ def test_sorted_masks(self):
         def foo(x, y):
             return (x + y).sum(dim=1)
 
-        x = torch.rand([255, 255], device="cuda")
-        y = torch.rand([255, 255], device="cuda")
+        x = torch.rand([255, 255], device=device_type)
+        y = torch.rand([255, 255], device=device_type)
 
         _, code = run_and_get_code(foo, x, y)
         FileCheck().check("tl.load").check_same("r0_mask").check_same("xmask").run(
@@ -1792,7 +1938,8 @@ def cat(inps):
 
         for dtype in [torch.uint8, torch.int8]:
             inps = [
-                torch.empty([256, 256], dtype=dtype, device="cuda") for _ in range(4)
+                torch.empty([256, 256], dtype=dtype, device=device_type)
+                for _ in range(4)
             ]
 
             out, code = run_and_get_code(cat, inps)
@@ -1831,20 +1978,20 @@ def fn(arg207_1, arg208_1, convert_element_type_40, expand, full, mul_3):
             return (sub_2,)
 
         args = [
-            torch.randn((8, 1024, 4, 4), device="cuda") > 0,  # torch.bool tensor
-            torch.randn((1, 1024, 1, 1), device="cuda"),
-            torch.randn((8, 1024, 4, 4), device="cuda"),
-            torch.randn((8, 1024, 1, 1), dtype=torch.float16, device="cuda").expand(
-                (8, 1024, 4, 4)
-            ),
-            torch.randn((), device="cuda"),
-            torch.randn((1024,), device="cuda"),
+            torch.randn((8, 1024, 4, 4), device=device_type) > 0,  # torch.bool tensor
+            torch.randn((1, 1024, 1, 1), device=device_type),
+            torch.randn((8, 1024, 4, 4), device=device_type),
+            torch.randn(
+                (8, 1024, 1, 1), dtype=torch.float16, device=device_type
+            ).expand((8, 1024, 4, 4)),
+            torch.randn((), device=device_type),
+            torch.randn((1024,), device=device_type),
         ]
         fn(*args)
-        torch.cuda.synchronize()  # shake out Triton Error [CUDA]: misaligned address
+        torch.accelerator.synchronize()  # shake out Triton Error [CUDA]: misaligned address
 
     def test_mutated_aligned_tensor(self):
-        t = torch.rand(4096, device="cuda", dtype=torch.float16)
+        t = torch.rand(4096, device=device_type, dtype=torch.float16)
 
         def foo(x):
             return x.add_(1)
@@ -1864,8 +2011,8 @@ def foo(x):
     def test_non_commutative_scan_op(self):
         from torch._higher_order_ops.associative_scan import associative_scan
 
-        a = torch.randn(1024, 8192, dtype=torch.float64, device="cuda")
-        b = torch.randn(1024, 8192, dtype=torch.float64, device="cuda")
+        a = torch.randn(1024, 8192, dtype=torch.float64, device=device_type)
+        b = torch.randn(1024, 8192, dtype=torch.float64, device=device_type)
 
         def baseline(v, u):
             A = []
@@ -1890,17 +2037,17 @@ def compiled_scan(a, b):
     def test_dynamic_persistent_reductions(self):
         @torch.compile(dynamic=True)
         def inner_reduce(x):
-            assert x.shape[1] <= 1024
+            assert x.shape[1] <= 1024  # noqa: S101
             return x.sum(1)
 
-        a = torch.randn(50, 600, device="cuda")
+        a = torch.randn(50, 600, device=device_type)
         out, code = run_and_get_code(inner_reduce, a)
         self.assertEqual(inner_reduce(a), out)
         self.assertTrue("for roffset" not in code)
 
         @torch.compile(dynamic=True)
         def outer_reduce(x):
-            assert x.shape[0] <= 64
+            assert x.shape[0] <= 64  # noqa: S101
             return x.sum(0)
 
         out, code = run_and_get_code(outer_reduce, a)
@@ -1950,7 +2097,7 @@ def forward(self, hidden_states: Tensor, attention_mask: Tensor) -> Tensor:
                 )
                 return attn_output
 
-        device = torch.device("cuda")
+        device = torch.device(device_type)
         num_attention_heads = 8
         hidden_size = 512
         attention_probs_dropout_prob = 0.0
@@ -1976,30 +2123,34 @@ def forward(self, hidden_states: Tensor, attention_mask: Tensor) -> Tensor:
     def test_non_contiguous_unaligned_input_indices(self):
         from torch._inductor.compile_fx import remove_unaligned_input_idxs
 
-        inputs = [torch.ones(2, 2, device="cuda"), torch.ones(2, 2, device="cuda")[1:]]
+        inputs = [
+            torch.ones(2, 2, device=device_type),
+            torch.ones(2, 2, device=device_type)[1:],
+        ]
         idxs = remove_unaligned_input_idxs(inputs, [1])
         self.assertEqual(idxs, [])
 
         inputs = [
-            torch.ones(2, 2, device="cuda"),
-            torch.ones(2, 2, device="cuda"),
-            torch.ones(2, 2, device="cuda")[1:],
+            torch.ones(2, 2, device=device_type),
+            torch.ones(2, 2, device=device_type),
+            torch.ones(2, 2, device=device_type)[1:],
         ]
         idxs = remove_unaligned_input_idxs(inputs, [0, 2])
         self.assertEqual(idxs, [0])
 
+    @skipIfXpu(msg="cudagraph is not supported on xpu")
     @config.patch("triton.cudagraphs", True)
     def test_unused_cpu_input_cudagraphs(self):
         def fn(x, y):
             return x.sin().sin().sin().sin().cos() + 1
 
         fx_graph = torch.fx.symbolic_trace(fn)
-        inp = [torch.randn(64, device="cuda"), torch.randn(64, device="cpu")]
+        inp = [torch.randn(64, device=device_type), torch.randn(64, device="cpu")]
         compiled_fn, (graph,) = run_and_get_graph_lowering(
             torch._inductor.compile, fx_graph, inp
         )
         self.assertEqual(graph.disable_cudagraphs_reason, None)
-        self.assertEqual(graph.device_types, {"cuda"})
+        self.assertEqual(graph.device_types, {device_type})
         self.assertEqual(compiled_fn(*inp), fn(*inp))
 
     def test_epilogue_fusion_with_view(self):
@@ -2015,8 +2166,8 @@ def forward(self, x):
                 x = x.view(x.size(0), -1)
                 return self.relu(self.linear(x))
 
-        m = ToyModel().to(device="cuda:0")
-        input_tensor = torch.randn(32, 3, 64, 64).to(device="cuda:0")
+        m = ToyModel().to(device=f"{device_type}:0")
+        input_tensor = torch.randn(32, 3, 64, 64).to(device=f"{device_type}:0")
         from torch._inductor.utils import fresh_cache
 
         with fresh_cache():
@@ -2025,6 +2176,7 @@ def forward(self, x):
             out2 = m(input_tensor)
             self.assertEqual(out, out2, atol=1e-3, rtol=1e-3)
 
+    @skipIfXpu(msg="cudagraph is not supported on xpu")
     @config.patch("triton.cudagraphs", True)
     def test_cpu_index(self):
         @torch.compile(fullgraph=True)
@@ -2032,25 +2184,25 @@ def fn(x):
             return x[torch.arange(32)]
 
         result, (graph,) = run_and_get_graph_lowering(
-            fn, torch.randn(64, device="cuda")
+            fn, torch.randn(64, device=device_type)
         )
         self.assertEqual(graph.disable_cudagraphs_reason, None)
-        self.assertEqual(graph.device_types, {"cuda"})
+        self.assertEqual(graph.device_types, {device_type})
 
-        inp = torch.randn(64, device="cuda", requires_grad=True)
+        inp = torch.randn(64, device=device_type, requires_grad=True)
         result, (graph,) = run_and_get_graph_lowering(fn, inp)
         self.assertEqual(graph.disable_cudagraphs_reason, None)
-        self.assertEqual(graph.device_types, {"cuda"})
+        self.assertEqual(graph.device_types, {device_type})
 
         result, (graph,) = run_and_get_graph_lowering(lambda: result.sum().backward())
         self.assertEqual(graph.disable_cudagraphs_reason, None)
-        self.assertEqual(graph.device_types, {"cuda"})
+        self.assertEqual(graph.device_types, {device_type})
 
     @unittest.skipIf(IS_FBCODE, "Not runnable in fbcode")
     def test_triton_interpret(self):
         import subprocess
 
-        script = """
+        script = f"""
 import os
 os.environ["TRITON_INTERPRET"] = "1"
 import torch
@@ -2060,8 +2212,9 @@ def foo(x):
     return x + 1
 
 # somehow gives different results.. still, check that it doesn't error
-foo(torch.rand([256], device="cuda"))
+foo(torch.rand([256], device=\"{device_type}\"))
 """
+
         subprocess.run([sys.executable, "-c", script], check=True)
 
     def test_reflection_pad_loop_order(self):
@@ -2071,8 +2224,8 @@ def fn(x, y):
             return a + b
 
         cfn = torch.compile(fn)
-        a = torch.rand((10, 10, 10), device="cuda")
-        b = torch.rand((10, 10, 10), device="cuda")
+        a = torch.rand((10, 10, 10), device=device_type)
+        b = torch.rand((10, 10, 10), device=device_type)
         expect = fn(a, b)
         actual, code = run_and_get_code(cfn, a, b)
         self.assertEqual(expect, actual)
@@ -2133,7 +2286,7 @@ def foo(inp):
             return cat_1
 
         for mark_dynamic in [False, True]:
-            inp = torch.rand((65536, 8192), dtype=torch.bfloat16, device="cuda")
+            inp = torch.rand((65536, 8192), dtype=torch.bfloat16, device=device_type)
             if mark_dynamic:
                 torch._dynamo.mark_dynamic(inp, 0)
             foo_c = torch.compile(foo)
@@ -2143,7 +2296,7 @@ def foo(inp):
         not SM90OrLater, "uses bfloat16 atomic add instrs which requires SM >= 90"
     )
     def test_float8_e8m0fnu(self):
-        device = "cuda"
+        device = device_type
         dtype = torch.float8_e8m0fnu
         hp_dtype = torch.float32  # and torch.bfloat16
 
@@ -2190,9 +2343,9 @@ def f(x, y):
             return torch.index_select(x, 0, y)
 
         x = torch.randn(
-            2000, 384, dtype=torch.bfloat16, device="cuda", requires_grad=True
+            2000, 384, dtype=torch.bfloat16, device=device_type, requires_grad=True
         )
-        y = torch.ones(713268, dtype=torch.int64, device="cuda")
+        y = torch.ones(713268, dtype=torch.int64, device=device_type)
         x_ref = x.clone().detach().requires_grad_(True)
         y_ref = y.clone().detach()
 
@@ -2225,15 +2378,15 @@ def f(arg0_1, arg1_1):
                     start=0,
                     step=1,
                     dtype=torch.int64,
-                    device="cuda",
+                    device=device_type,
                     requires_grad=False,
                 )
                 view_3 = torch.ops.aten.view.default(iota, [1, 36])
                 max_1 = torch.ops.aten.max.default(view_3)
                 return (max_1,)
 
-            x = torch.ones(1, 64, device="cuda", dtype=torch.int64)
-            y = torch.randn(64, 3072, device="cuda", dtype=torch.bfloat16)
+            x = torch.ones(1, 64, device=device_type, dtype=torch.int64)
+            y = torch.randn(64, 3072, device=device_type, dtype=torch.bfloat16)
             out = f(x, y)
             self.assertEqual(torch.compile(f)(x, y), out)
 
@@ -2249,9 +2402,9 @@ def f(x, y):
             return torch.index_select(x, 0, y)
 
         x = torch.randn(
-            2000, 384, dtype=torch.bfloat16, device="cuda", requires_grad=True
+            2000, 384, dtype=torch.bfloat16, device=device_type, requires_grad=True
         )
-        y = torch.ones(713268, dtype=torch.int64, device="cuda")
+        y = torch.ones(713268, dtype=torch.int64, device=device_type)
         x_ref = x.clone().detach().requires_grad_(True)
         y_ref = y.clone().detach()
 
@@ -2262,6 +2415,25 @@ def f(x, y):
 
         self.assertEqual(f(x_ref, y_ref), out)
 
+    @skipCUDAIf(
+        not SM90OrLater, "uses bfloat16 atomic add instrs which requires SM >= 90"
+    )
+    @unittest.skipIf(
+        config.is_fbcode(),
+        "bfloat16 atomic add is supported in fbcode, so we won't fallback",
+    )
+    def test_index_add_fallback_direct(self):
+        def f(x, idx, src):
+            return torch.index_add(x, 0, idx, src)
+
+        x = torch.randn(16, 256, dtype=torch.bfloat16, device=device_type)
+        idx = torch.randperm(8, device=device_type)
+        src = torch.randn(8, 256, dtype=torch.bfloat16, device=device_type)
+
+        out = f(x, idx, src)
+        compiled_out = torch.compile(f)(x, idx, src)
+        self.assertEqual(out, compiled_out)
+
     @requires_multigpu()
     def test_not_initializing_wrong_device(self):
         device_stats = torch.cuda.memory_stats("cuda:0")
@@ -2293,7 +2465,7 @@ def test_3d_tiling(self):
             1,
             2,
         )
-        GPU_TYPE = "cuda"
+        GPU_TYPE = device_type
 
         def get_input() -> torch.Tensor:
             device = torch.device(GPU_TYPE)
@@ -2315,7 +2487,7 @@ def test_repeated_masked_load(self):
         mem_eff_temporal_upsampling_interp_chunks = 2
         from functorch.einops import rearrange
 
-        x = torch.randn(1, 8, 12, 12, 4, dtype=torch.float16, device="cuda")
+        x = torch.randn(1, 8, 12, 12, 4, dtype=torch.float16, device=device_type)
         x = x.permute(0, 1, 4, 2, 3)  # make non-contiguous
         x = rearrange(x, "b c t h w -> b c t (h w)")
 
@@ -2377,8 +2549,8 @@ def forward(self, x):
 
                 return x
 
-        model = ToyModel().to("cuda")
-        input_tensor = torch.randn((2, 4)).to("cuda")
+        model = ToyModel().to(device_type)
+        input_tensor = torch.randn((2, 4)).to(device_type)
 
         compile_default = torch.compile(model, mode="default")
         compile_max_autotune = torch.compile(model, mode="max-autotune")
@@ -2405,7 +2577,7 @@ def forward(self, x):
                 x = self.adaptive_pool(x)
                 return x
 
-        model = Model().cuda()
+        model = Model().to(device_type)
         model.eval()
         test_cases = [
             (1, 3, 8, 8),
@@ -2417,7 +2589,7 @@ def forward(self, x):
 
         for batch, channels, h, w in test_cases:
             with self.subTest(input_shape=(batch, channels, h, w)):
-                input_tensor = torch.randn(batch, channels, h, w, device="cuda")
+                input_tensor = torch.randn(batch, channels, h, w, device=device_type)
 
                 # Test eager mode
                 with torch.no_grad():
@@ -2454,7 +2626,10 @@ def test_searchsorted_stride_permutations(
         class Foo(torch.nn.Module):
             def __init__(self, quantiles: torch.Tensor) -> None:
                 super().__init__()
-                assert quantiles.shape[0] > 0
+                if quantiles.shape[0] <= 0:
+                    raise AssertionError(
+                        f"Expected quantiles.shape[0] > 0, got {quantiles.shape[0]}"
+                    )
                 quantiles = quantiles.T
                 self.q = torch.nn.Parameter(quantiles, requires_grad=False)
 
@@ -2467,7 +2642,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
         numel = 1
         for dim in quantiles_shape:
             numel *= dim
-        data = torch.randn(numel, dtype=torch.float32, device="cuda")
+        data = torch.randn(numel, dtype=torch.float32, device=device_type)
 
         # Create tensor with specified shape and strides
         quantiles = torch.as_strided(
@@ -2477,7 +2652,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
         quantiles = torch.sort(quantiles, dim=0)[0]
 
         x_shape = (batch_size,) + quantiles_shape[1:]
-        x = torch.randn(*x_shape, dtype=torch.float32, device="cuda")
+        x = torch.randn(*x_shape, dtype=torch.float32, device=device_type)
 
         foo = Foo(quantiles)
         foo_compiled = torch.compile(Foo(quantiles), fullgraph=True)
@@ -2490,7 +2665,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
         self.assertEqual(eager, compiled)
 
     def test_identity_load(self):
-        device = "cuda"
+        device = device_type
 
         def f(x, y):
             y2 = torch.cat(
@@ -2525,9 +2700,12 @@ def f(x, y):
 
         self.assertEqual(eager_out, compile_out)
 
+    @skipIfXpu(
+        msg="Explicit attn_mask should not be set when is_causal=True - torch-xpu-ops: 2802"
+    )
     def test_qwen2_7b_sdpa_input_alignment_requires_recompile(self):
         # SDPA constraints ensures inputs have alignment (8).
-        device = "cuda"
+        device = device_type
 
         def forward(q_proj, k_proj, attn_mask):
             scale = 0.08838834764831845  # 1/sqrt(128)
@@ -2645,8 +2823,8 @@ def compiled_divide(x, y):
                 torch.float32,
                 torch.float64,
             ]:
-                y_ten = torch.tensor([y], dtype=y_dtype, device="cuda")
-                x_ten = torch.tensor([x], dtype=x_dtype, device="cuda")
+                y_ten = torch.tensor([y], dtype=y_dtype, device=device_type)
+                x_ten = torch.tensor([x], dtype=x_dtype, device=device_type)
 
                 torch._dynamo.reset()
                 compiled_div = Decimal(compiled_divide(x_ten, y_ten).item())
@@ -2654,6 +2832,7 @@ def compiled_divide(x, y):
 
                 self.assertEqual(eager_div, compiled_div)
 
+    @skipIfXpu(msg="triton dependency - torch-xpu-ops: 2554")
     @config.patch({"eager_numerics.division_rounding": False})
     @xfailIfROCm
     def test_truediv_base_not_bitwise_equivalent(self):
@@ -2661,8 +2840,8 @@ def test_truediv_base_not_bitwise_equivalent(self):
 
         y, x = 7.0, 11.0
 
-        y_ten = torch.tensor([y], dtype=torch.float32, device="cuda")
-        x_ten = torch.tensor([x], dtype=torch.float32, device="cuda")
+        y_ten = torch.tensor([y], dtype=torch.float32, device=device_type)
+        x_ten = torch.tensor([x], dtype=torch.float32, device=device_type)
 
         compile_out, code = run_and_get_code(
             torch.compile(lambda x, y: x / y),
@@ -2680,7 +2859,7 @@ def test_disabling_ftz_yields_subnormals(self):
         from decimal import Decimal
 
         x = -127.0
-        x_ten = torch.tensor([x], dtype=torch.float32, device="cuda")
+        x_ten = torch.tensor([x], dtype=torch.float32, device=device_type)
 
         def fn(x):
             return 2.0**x
@@ -2690,13 +2869,14 @@ def fn(x):
 
         self.assertTrue(compile_decimal > Decimal(0))
 
+    @skipIfXpu(msg="Decimal object comparison failed - torch-xpu-ops: 2810")
     @skipIfRocm(msg="ROCm preserves subnormals by default")
     @config.patch({"eager_numerics.disable_ftz": False})
     def test_not_disabling_ftz_yields_zero(self):
         from decimal import Decimal
 
         x = -128.0
-        x_ten = torch.tensor([x], dtype=torch.float32, device="cuda")
+        x_ten = torch.tensor([x], dtype=torch.float32, device=device_type)
 
         def fn(x):
             return 2.0**x
@@ -2706,10 +2886,36 @@ def fn(x):
 
         self.assertEqual(compile_decimal, Decimal(0))
 
+    @skipIfXpu(msg="AssertionError: torch-xpu-ops: #3006")
+    @config.patch(
+        {"triton.use_block_ptr": True, "triton.codegen_upcast_to_fp32": False}
+    )
+    def test_float16_reduction_with_int_output(self):
+        @torch.compile
+        def fn(input: torch.Tensor) -> torch.Tensor:
+            return torch.argmax(input, dim=0)
+
+        input = torch.randn(20, 20, device=device_type, dtype=torch.float16)
+        _, code = run_and_get_code(fn, input)
+        # There should not be any conversions to float16 in this code, since the input
+        # is already float16 and the output is int64.
+        self.assertNotIn(".to(tl.float16)", code[0])
+
+    @config.patch("eager_numerics.division_rounding", True)
+    def test_reciprocal_precision_rounding(self):
+        # Test that reciprocal matches eager when division_rounding is enabled.
+        # This requires OpDecompositions.reciprocal to use float32 constant so
+        # that div_rn can be applied (the dtype check requires both operands float32).
+        def fn(x):
+            return torch.reciprocal(x)
+
+        x = torch.randn(1000, device=device_type, dtype=torch.float32) + 0.1
+        self.common(fn, [x])
+
 
 if __name__ == "__main__":
     from torch._inductor.test_case import run_tests
-    from torch.testing._internal.inductor_utils import HAS_CUDA_AND_TRITON
+    from torch.testing._internal.inductor_utils import HAS_GPU_AND_TRITON
 
-    if HAS_CUDA_AND_TRITON and not TEST_WITH_ASAN:
+    if HAS_GPU_AND_TRITON and not TEST_WITH_ASAN:
         run_tests(needs="filelock")
diff --git a/test/inductor/test_cudagraph_trees.py b/test/inductor/test_cudagraph_trees.py
index 31706984a0419..8c040da118812 100644
--- a/test/inductor/test_cudagraph_trees.py
+++ b/test/inductor/test_cudagraph_trees.py
@@ -22,14 +22,14 @@
 from torch._inductor.codecache import FxGraphCache
 from torch._inductor.compile_fx import compile_fx_inner
 from torch._inductor.cudagraph_trees import cudagraphify_impl as tree_cudagraphify_impl
-from torch._inductor.cudagraph_utils import FunctionID, PlaceholderInfo
+from torch._inductor.cudagraph_utils import PlaceholderInfo
 from torch._inductor.test_case import TestCase as InductorTestCase
 from torch._inductor.utils import run_and_get_code
 from torch._ops import OpOverload
 from torch.fx.experimental.proxy_tensor import make_fx
 from torch.fx.immutable_collections import immutable_dict
 from torch.testing import FileCheck
-from torch.testing._internal.common_cuda import TEST_MULTIGPU
+from torch.testing._internal.common_cuda import blas_library_context, TEST_MULTIGPU
 from torch.testing._internal.common_utils import (
     instantiate_parametrized_tests,
     IS_ARM64,
@@ -65,6 +65,16 @@
 )
 from io import StringIO
 
+from torch._library.opaque_object import OpaqueBase, register_opaque_type
+
+
+class _CudagraphTestScaleFactor(OpaqueBase):
+    def __init__(self, factor):
+        self.factor = factor
+
+
+register_opaque_type(_CudagraphTestScaleFactor, typ="reference")
+
 
 def get_compile_fn(backend):
     if backend == "cudagraphs":
@@ -77,7 +87,8 @@ def get_num_partitions(code):
     """Get the number of cudagraph partitions from generated code."""
     code = "".join(code)
     found = re.search(r"partitions=\[(.*)\]", code)
-    assert found is not None, "Could not find partitions in generated code"
+    if found is None:
+        raise AssertionError("Could not find partitions in generated code")
     partitions = found.group(1)
     return len([p for p in partitions.split(",") if p])
 
@@ -745,8 +756,7 @@ def foo2(x):
             self.assertEqual(foo_opt(ones), foo(ones))
             # paths
             children = self.get_root_children()
-            # one root with two children
-            self.assertEqual(children, [2])
+            self.assertEqual(children, [1])
 
         def test_end_recording_early(self):
             def foo(x):
@@ -893,7 +903,8 @@ def foo(x):
             self.assertEqual(counters["aot_autograd"]["autograd_cache_hit"], 1)
 
             # we should not have cudagraph'd anything
-            assert self.get_manager() is None
+            if self.get_manager() is not None:
+                raise AssertionError
 
         @torch._functorch.config.patch("enable_autograd_cache", True)
         @torch._inductor.config.patch("fx_graph_cache", True)
@@ -981,7 +992,8 @@ def foo(x):
             out.backward(back_inp)
 
             # we should not have cudagraph'd anything
-            assert self.get_manager() is None
+            if self.get_manager() is not None:
+                raise AssertionError
 
         @torch._inductor.config.patch("triton.skip_cudagraph_warmup", True)
         @torch._functorch.config.patch("enable_autograd_cache", True)
@@ -1926,8 +1938,11 @@ def foo2(args):
             del x
             self.assertEqual(all_live_block_count(), 0)
 
-        @unittest.skipUnless(IS_X86 and IS_LINUX, "cpp contexts are linux only")
+        @unittest.skipUnless(
+            (IS_X86 or IS_ARM64) and IS_LINUX, "cpp contexts are linux x86/aarch64 only"
+        )
         @torch._inductor.config.patch("triton.cudagraph_trees_history_recording", True)
+        @blas_library_context("cublas")
         def test_workspace_allocation_error(self):
             torch._C._cuda_clearCublasWorkspaces()
 
@@ -2229,25 +2244,6 @@ def foo(mod, inp):
             node = self.get_manager().current_node
             self.assertEqual(len(list(node.path_live_weakrefs())), 1)
 
-        @torch._dynamo.config.patch("inline_inbuilt_nn_modules", False)
-        @torch._inductor.config.patch("triton.cudagraph_support_input_mutation", False)
-        def test_unstable_ptr(self):
-            import torch
-
-            @torch.compile(mode="reduce-overhead")
-            def foo(m, inp):
-                return m(inp)
-
-            def f():
-                l = []
-                m = torch.nn.Linear(20, 20).cuda()
-                for _ in range(4):
-                    inp = torch.rand([20, 20], device="cuda")
-                    foo(m, inp)
-                    m.weight.data = torch.rand([20, 20], device="cuda")
-
-            self.assertRaises(RuntimeError, f)
-
         @requires_multigpu()
         def test_manager_per_device(self):
             def test():
@@ -2308,6 +2304,63 @@ def foo(x):
 
             FileCheck().check("overwritten").check("x * x * x").run(repr(exc.exception))
 
+        def test_output_node_has_stack_traces_inference(self):
+            """Test that output_stack_traces on the output node provides
+            stack traces even when a post-grad pass strips them from arg nodes
+            in inference mode."""
+
+            def strip_stack_traces(graph):
+                for node in graph.nodes:
+                    if node.op not in ("placeholder", "output"):
+                        node.meta.pop("stack_trace", None)
+
+            with config.patch(post_grad_custom_post_pass=strip_stack_traces):
+
+                @torch.compile(mode="reduce-overhead")
+                def foo(x):
+                    return x * x * x
+
+                inp = torch.rand([4], device="cuda")
+                out = foo(inp).detach()
+                out2 = foo(inp).detach()
+
+                with self.assertRaises(Exception) as exc:
+                    out + out
+
+                self.assertIn("x * x * x", repr(exc.exception))
+
+        def test_output_node_has_stack_traces_training(self):
+            """Test that output_stack_traces on the output node provides
+            stack traces even when a post-grad pass strips them from arg nodes
+            in training mode (fwd/bwd graph split via partitioner)."""
+
+            def strip_stack_traces(graph):
+                for node in graph.nodes:
+                    if node.op not in ("placeholder", "output"):
+                        node.meta.pop("stack_trace", None)
+
+            with config.patch(post_grad_custom_post_pass=strip_stack_traces):
+
+                @torch.compile(mode="reduce-overhead")
+                def foo(x):
+                    return x * x * x
+
+                inp = torch.rand([4], device="cuda", requires_grad=True)
+                # Complete fwd+bwd to compile both graphs
+                torch.compiler.cudagraph_mark_step_begin()
+                foo(inp).sum().backward()
+
+                # Now trigger the dealloc error
+                torch.compiler.cudagraph_mark_step_begin()
+                out = foo(inp).detach()
+                torch.compiler.cudagraph_mark_step_begin()
+                out2 = foo(inp).detach()
+
+                with self.assertRaises(Exception) as exc:
+                    out + out
+
+                self.assertIn("x * x * x", repr(exc.exception))
+
         @unittest.skipIf(not torch.backends.cudnn.is_available(), "requires cudnn")
         def test_conv_benchmark(self):
             with torch.backends.cudnn.flags(
@@ -2442,9 +2495,12 @@ def foo(x):
             with warnings.catch_warnings(record=True) as w:
                 out = foo(torch.rand([4, 4], device="cuda", requires_grad=True))
 
-            FileCheck().check(
-                "Unable to hit fast path of CUDAGraphs because of pending"
-            ).run(str(w[0]))
+            # Match substring only; scan all warnings in case another warning fires first.
+            msgs = [str(x.message) for x in w]
+            self.assertTrue(
+                any("require backward" in m for m in msgs),
+                f"expected CUDAGraph pending-backward warning; got: {msgs}",
+            )
             self.assertTrue(self.get_manager().new_graph_id().id == 0)
 
         def test_mark_step(self):
@@ -2593,9 +2649,9 @@ def foo(x):
                 t = torch.rand([32], device="cuda")
                 self.assertEqual(foo(t), foo_c(t))
 
-            FileCheck().check("skipping cudagraphs due to cpp wrapper enabled").run(
-                log_stream.getvalue()
-            )
+            FileCheck().check(
+                "skipping cudagraphs due to cpp-wrapper does not support graph partition yet"
+            ).run(log_stream.getvalue())
             self.assertEqual(counters["inductor"]["cudagraph_skips"], 1)
 
         def test_storage_access_error(self):
@@ -2605,6 +2661,25 @@ def test_storage_access_error(self):
             with self.assertRaisesRegex(Exception, "custom error msg"):
                 device = x.untyped_storage()
 
+        def test_clear_storage_data_ptr_access_error(self):
+            x = torch.rand([4], device="cuda")
+            storage = x.untyped_storage()
+            storage_ptr = storage.data_ptr()
+            storage_impl_ptr = storage._cdata
+
+            storage.resize_(0)
+
+            torch._C._set_storage_data_ptr_access_error_msg(
+                storage_impl_ptr, "storage is dead"
+            )
+            with self.assertRaisesRegex(Exception, "storage is dead"):
+                storage.data_ptr()
+
+            torch._C._clear_storage_data_ptr_access_error_msg(storage_impl_ptr)
+            storage.resize_(4 * x.element_size())
+            # Should not raise
+            storage.data_ptr()
+
         def test_side_stream_memory_allocation(self):
             device = f"cuda:{self.device_idx}"
 
@@ -2648,49 +2723,6 @@ def multi_stream_allocation(args):
 
             self.assertEqual(self.get_manager().new_graph_id().id, 1)
 
-        @torch._dynamo.config.patch("inline_inbuilt_nn_modules", False)
-        @torch._inductor.config.patch("triton.cudagraph_support_input_mutation", False)
-        def test_static_inputs_address_mutation_log(self):
-            class Goo(torch.nn.Module):
-                def __init__(self) -> None:
-                    super().__init__()
-                    self.linear = torch.nn.Linear(2, 2, device="cuda")
-
-                def forward(self, x) -> torch.Tensor:
-                    return self.linear(x)
-
-            class Foo(torch.nn.Module):
-                def __init__(self) -> None:
-                    super().__init__()
-                    self.static_tensor = torch.zeros((2, 2), device="cuda")
-                    self.goo = Goo()
-
-                def forward(self, x) -> torch.Tensor:
-                    self.static_tensor.add_(torch.ones((2, 2), device="cuda"))
-                    return self.static_tensor + x + self.goo(x)
-
-            foo = Foo()
-            foo = torch.compile(foo, mode="reduce-overhead")
-            inp = torch.rand((2, 2), device="cuda")
-
-            for _ in range(3):
-                foo(inp)
-
-            # mutates static input tensors' addresses
-            foo.static_tensor = torch.ones((2, 2), device="cuda")
-            foo.goo.linear.bias = torch.nn.Parameter(torch.ones((2,), device="cuda"))
-
-            with self.assertRaisesRegex(
-                Exception,
-                r"(?s)static input data pointer changed.\n"
-                r"input name: primals_.*. data pointer changed from .* to .*. input stack trace:.*"
-                r"input name: primals_.*. data pointer changed from .* to .*. input stack trace:.*,"
-                r" in forward\n.* self.static_tensor.add\_\(torch.ones\(\(2, 2\), device=\"cuda\"\)\).*\n",
-            ):
-                self.curr_node().run(
-                    [foo.goo.linear.weight, foo.goo.linear.bias, foo.static_tensor, inp]
-                )
-
         def _run_iter(self, param, fn):
             fwd_output = fn(torch.ones(2, 2), param)
             fwd_output.sum().backward()
@@ -2757,7 +2789,6 @@ def run_test():
                 self.assertEqual(self.get_manager().new_graph_id().id, 4)
 
         @torch._dynamo.config.patch("error_on_recompile", True)
-        @torch._dynamo.config.patch("inline_inbuilt_nn_modules", True)
         def test_multi_dispatch_single_compile_param_inputs(self):
             # Verify that we can record multiple cudagraphs for a single
             # compiled function with param inputs
@@ -2768,7 +2799,6 @@ def fn(x, y):
             self.run_static_input_param_test(fn, 4)
 
         @torch._dynamo.config.patch("error_on_recompile", True)
-        @torch._dynamo.config.patch("inline_inbuilt_nn_modules", True)
         def test_multi_dispatch_single_compile_builtin_module(self):
             # Verify that we don't recompile when changing the param of a builtin module
             # and that we record another cudagraph
@@ -2776,7 +2806,6 @@ def test_multi_dispatch_single_compile_builtin_module(self):
             self._module_test(torch.nn.Linear(2, 3, device="cuda"))
 
         @torch._dynamo.config.patch("error_on_recompile", True)
-        @torch._dynamo.config.patch("inline_inbuilt_nn_modules", True)
         def test_multi_dispatch_single_compile_builtin_module_buffers(self):
             # Verify that we don't recompile when changing the buffer of a builtin module
             # and that we record another cudagraph
@@ -2788,7 +2817,6 @@ def test_multi_dispatch_single_compile_builtin_module_buffers(self):
 
         @torch._inductor.config.patch("triton.cudagraphs", True)
         @torch._dynamo.config.patch("error_on_recompile", True)
-        @torch._dynamo.config.patch("inline_inbuilt_nn_modules", True)
         def test_multi_dispatch_custom_module(self):
             # Test that we can correctly dispatch multiple graphs
             # if params of a custom module change
@@ -2805,7 +2833,6 @@ def forward(self, x):
             )
 
         @torch._dynamo.config.patch("error_on_recompile", True)
-        @torch._dynamo.config.patch("inline_inbuilt_nn_modules", True)
         def test_multi_dispatch_custom_module_buffer(self):
             # Test that we can correctly dispatch multiple graphs
             # if buffers of a custom module change
@@ -2829,7 +2856,6 @@ def forward(self, x):
 
         @torch._inductor.config.patch("triton.cudagraphs", True)
         @torch._dynamo.config.patch("error_on_recompile", True)
-        @torch._dynamo.config.patch("inline_inbuilt_nn_modules", True)
         def test_multi_dispatch_child_node(self):
             # Test that we can correctly dispatch multiple graphs if a child node
             # in the tree has stable input pointers change
@@ -2848,7 +2874,6 @@ def fn(x, p):
             self.run_static_input_param_test(fn, 5)
 
         @torch._dynamo.config.patch("error_on_recompile", True)
-        @torch._dynamo.config.patch("inline_inbuilt_nn_modules", True)
         def test_multi_dispatch_parent_node(self):
             def fn(x, p):
                 # Graph 1
@@ -2868,145 +2893,6 @@ def fn(x, p):
             self.run_static_input_param_test(fn, 6)
 
         @torch._dynamo.config.patch("error_on_recompile", True)
-        @torch._dynamo.config.patch("inline_inbuilt_nn_modules", False)
-        @torch._inductor.config.patch("triton.cudagraph_support_input_mutation", True)
-        @torch._inductor.config.patch("triton.cudagraph_unexpected_rerecord_limit", 0)
-        def test_fallback_to_eager_if_recompiling_too_many_times(self):
-            class Foo(torch.nn.Module):
-                def __init__(self) -> None:
-                    super().__init__()
-                    self.param = torch.nn.Parameter(torch.rand([2, 2], device="cuda"))
-
-                def forward(self, x):
-                    return x * self.param
-
-            log_stream, ctx = logs_to_string(
-                "torch._inductor.cudagraph_utils", "cudagraphs"
-            )
-            with ctx():
-                # We have 3 graphs here
-                #             None
-                #       /                           \
-                # (fwd w/ p1, Graph 0)            (bwd w/p2, Graph2)
-                # (bwd w/ p1, Graph 1)
-                # All other graphs are skipped because we hit the max recording limit
-                # (=0 for each node and function pair)
-                fn_compiled = torch.compile(Foo(), mode="reduce-overhead")
-                for _ in range(3):
-                    fn_compiled(torch.rand([2, 2], device="cuda")).sum().backward()
-                    fn_compiled.param.grad = None
-
-                # Change static tensor address
-                fn_compiled.param.data = torch.rand([2, 2], device="cuda")
-                fn_compiled(torch.rand([2, 2], device="cuda")).sum().backward()
-                self.assertEqual(self.get_manager().new_graph_id().id, 3)
-
-            FileCheck().check(
-                "skipping cudagraph due to function 0 exceeding max re-recording limit (=0) "
-                "on cudagraph node None due to static input data pointer changed."
-            ).run(log_stream.getvalue())
-            self.assertEqual(counters["inductor"]["cudagraph_skips"], 1)
-
-        @torch._dynamo.config.patch("error_on_recompile", True)
-        @torch._dynamo.config.patch("inline_inbuilt_nn_modules", False)
-        @torch._inductor.config.patch("triton.cudagraph_support_input_mutation", True)
-        @torch._inductor.config.patch("triton.cudagraph_unexpected_rerecord_limit", 0)
-        def test_fallback_to_eager_if_recompiling_too_many_times_warn_only_once(self):
-            class Foo(torch.nn.Module):
-                def __init__(self) -> None:
-                    super().__init__()
-                    self.param = torch.nn.Parameter(torch.rand([2, 2], device="cuda"))
-
-                def forward(self, x):
-                    return x * self.param
-
-            log_stream, ctx = logs_to_string(
-                "torch._inductor.cudagraph_utils", "cudagraphs"
-            )
-            with ctx():
-                with torch.device("cuda"):
-                    # We have 3 graphs here
-                    #             None
-                    #       /                           \
-                    # (fwd w/ p1, Graph 0)            (bwd w/p2, Graph2)
-                    # (bwd w/ p1, Graph 1)
-                    # All other graphs are skipped because we hit the max recording limit
-                    # (=0 for each node and function pair)
-                    fn_compiled = torch.compile(Foo(), mode="reduce-overhead")
-                    for _ in range(3):
-                        fn_compiled(torch.rand([2, 2], device="cuda")).sum().backward()
-                        fn_compiled.param.grad = None
-
-                    for _ in range(5):
-                        # Change static tensor address
-                        fn_compiled.param.data = torch.rand([2, 2], device="cuda")
-                        fn_compiled(torch.rand([2, 2], device="cuda")).sum().backward()
-                        fn_compiled.param.grad = None
-
-            FileCheck().check_count(
-                "skipping cudagraph due to function 0 exceeding max re-recording limit (=0) "
-                "on cudagraph node None due to static input data pointer changed.",
-                1,
-                exactly=True,
-            ).check_count(
-                "skipping cudagraph due to function 1 exceeding max re-recording limit (=0) "
-                "on cudagraph node None due to static input data pointer changed.",
-                1,
-                exactly=True,
-            ).run(log_stream.getvalue())
-            self.assertEqual(counters["inductor"]["cudagraph_skips"], 2)
-
-        @torch._dynamo.config.patch("inline_inbuilt_nn_modules", False)
-        @torch._inductor.config.patch("triton.cudagraph_support_input_mutation", True)
-        @torch._inductor.config.patch("triton.cudagraph_unexpected_rerecord_limit", 0)
-        def test_fallback_to_eager_if_recompiling_too_many_times_due_to_cudagraph_managed_tensor(
-            self,
-        ):
-            # By setting triton.cudagraph_support_input_mutation=True, we force re-record
-            # if cudagraph managed tensor addresses changed.
-            @torch.compile(mode="reduce-overhead")
-            def foo(x):
-                return x + 1
-
-            @torch.compile(mode="reduce-overhead")
-            def goo(x):
-                return x * 2
-
-            for _ in range(3):
-                torch.compiler.cudagraph_mark_step_begin()
-                inp = torch.rand((2, 3), device="cuda")
-                y = foo(inp)
-                z = goo(y)
-
-            log_stream, ctx = logs_to_string(
-                "torch._inductor.cudagraph_utils", "cudagraphs"
-            )
-            with ctx():
-                torch.compiler.cudagraph_mark_step_begin()
-                x = torch.rand(2, 3, device="cuda")
-                y = foo(x)
-                y_clone = y.clone()
-                z = goo(y_clone)
-
-            # eager function should run successfully
-            for _ in range(5):
-                torch.compiler.cudagraph_mark_step_begin()
-                x = torch.rand(2, 3, device="cuda")
-                y = foo(x)
-                y_clone = y.clone()
-                z = goo(y_clone)
-
-            FileCheck().check_count(
-                "skipping cudagraph due to function 1 exceeding max re-recording limit (=0) "
-                "on cudagraph node 0 due to cudagraph managed tensor data pointer changed",
-                1,
-                exactly=True,
-            ).run(log_stream.getvalue())
-            self.assertEqual(counters["inductor"]["cudagraph_skips"], 1)
-
-        @torch._dynamo.config.patch("inline_inbuilt_nn_modules", False)
-        @torch._dynamo.config.patch("error_on_recompile", True)
-        @torch._dynamo.config.patch("inline_inbuilt_nn_modules", True)
         @torch._inductor.config.patch("triton.cudagraph_unexpected_rerecord_limit", 1)
         def test_not_fallback_to_eager_if_have_not_recompiling_too_many_times(self):
             def fn(x, y):
@@ -3021,7 +2907,6 @@ def fn(x, y):
             self.assertEqual(counters["inductor"]["cudagraph_skips"], 0)
 
         @torch._dynamo.config.patch("error_on_recompile", True)
-        @torch._dynamo.config.patch("inline_inbuilt_nn_modules", True)
         def test_no_rerecord_with_mark_static_address(self):
             class Mod(torch.nn.Module):
                 def __init__(self):
@@ -3054,6 +2939,352 @@ def fn_eager(x, marked_static_y):
                 # One forward and one backward
                 self.assertEqual(self.get_manager().new_graph_id().id, 2)
 
+        def test_cudagraph_annotation_disable(self):
+            @torch._dynamo.override_cudagraphs(fwd=False, bwd=False)
+            def helper(x):
+                return x + 1
+
+            def fn(x):
+                return helper(x)
+
+            inp = torch.randn(4, 4, device="cuda")
+            fn_compiled = torch.compile(fn, mode="reduce-overhead")
+            result = fn_compiled(inp)
+
+            self.assertEqual(result, inp + 1)
+            self.assertIsNone(self.get_manager())
+
+        def test_cudagraph_annotation_disable_fwd_bwd(self):
+            @torch._dynamo.override_cudagraphs(fwd=False, bwd=False)
+            def helper(x):
+                return (x * x).sin()
+
+            def fn(x):
+                return helper(x)
+
+            inp = torch.randn(4, 4, device="cuda", requires_grad=True)
+            fn_compiled = torch.compile(fn, mode="reduce-overhead")
+            out = fn_compiled(inp)
+            out.sum().backward()
+
+            self.assertEqual(out, inp.detach().pow(2).sin())
+            self.assertIsNone(self.get_manager())
+
+        def test_cudagraph_annotation_disable_bwd_only(self):
+            @torch._dynamo.override_cudagraphs(bwd=False)
+            def helper(x):
+                return (x * x).sin()
+
+            def fn(x):
+                return helper(x)
+
+            fn_compiled = torch.compile(fn, mode="reduce-overhead")
+
+            # Run forward+backward pair for cudagraph warmup and recording
+            for _ in range(3):
+                inp = torch.randn(4, 4, device="cuda", requires_grad=True)
+                torch.compiler.cudagraph_mark_step_begin()
+                out = fn_compiled(inp)
+                out.sum().backward()
+
+            self.assertEqual(out, inp.detach().pow(2).sin())
+            # Forward is eligible for cudagraphs (manager created)
+            self.assertIsNotNone(self.get_manager())
+            # One graph for forward (backward is disabled)
+            self.assertEqual(self.get_manager().new_graph_id().id, 1)
+
+        def test_cudagraph_annotation_disable_fwd_only(self):
+            @torch._dynamo.override_cudagraphs(fwd=False)
+            def helper(x):
+                return (x * x).sin()
+
+            def fn(x):
+                return helper(x)
+
+            inp = torch.randn(4, 4, device="cuda", requires_grad=True)
+            fn_compiled = torch.compile(fn, mode="reduce-overhead")
+            out = fn_compiled(inp)
+            out.sum().backward()
+
+            self.assertEqual(out, inp.detach().pow(2).sin())
+            # fwd override also disables bwd since the BoxedBool is shared
+            self.assertIsNone(self.get_manager())
+
+        def test_cudagraph_annotation_graph_break_selective(self):
+            """Annotated function only disables cudagraphs for graph segments
+            that inline it. Other segments after a graph break are unaffected."""
+
+            @torch._dynamo.override_cudagraphs(fwd=False, bwd=False)
+            def disabled_fn(x):
+                return x + 1
+
+            def model(x):
+                # g1: disabled_fn inlined -> cudagraphs OFF
+                y = disabled_fn(x)
+                torch._dynamo.graph_break()
+                # g2: no annotation -> cudagraphs ON
+                z = y * 2
+                torch._dynamo.graph_break()
+                # g3: disabled_fn inlined -> cudagraphs OFF
+                w = disabled_fn(z)
+                return w
+
+            inp = torch.randn(4, 4, device="cuda")
+            compiled = torch.compile(model, mode="reduce-overhead")
+            # Warmup + record
+            out = compiled(inp)
+            out = compiled(inp)
+
+            expected = (inp + 1) * 2 + 1
+            self.assertEqual(out, expected)
+            # Only g2 uses cudagraphs (1 graph recorded)
+            self.assertEqual(self.get_manager().new_graph_id().id, 1)
+
+        def test_cudagraph_annotation_multiple_functions(self):
+            """Multiple annotated functions each disable their containing graph."""
+
+            @torch._dynamo.override_cudagraphs(fwd=False, bwd=False)
+            def no_cg_add(x):
+                return x + 1
+
+            @torch._dynamo.override_cudagraphs(fwd=False, bwd=False)
+            def no_cg_mul(x):
+                return x * 3
+
+            def model(x):
+                y = no_cg_add(x)
+                torch._dynamo.graph_break()
+                z = no_cg_mul(y)
+                return z
+
+            inp = torch.randn(4, 4, device="cuda")
+            compiled = torch.compile(model, mode="reduce-overhead")
+            out = compiled(inp)
+            out = compiled(inp)
+
+            expected = (inp + 1) * 3
+            self.assertEqual(out, expected)
+            # Both segments disabled, no cudagraphs
+            self.assertIsNone(self.get_manager())
+
+        def test_cudagraph_annotation_deeply_nested(self):
+            """Annotation on an innermost function propagates through
+            nested inlining to disable cudagraphs for the whole graph."""
+
+            @torch._dynamo.override_cudagraphs(fwd=False, bwd=False)
+            def inner(x):
+                return x + 1
+
+            def middle(x):
+                return inner(x) * 2
+
+            def outer(x):
+                return middle(x).sin()
+
+            def model(x):
+                return outer(x)
+
+            inp = torch.randn(4, 4, device="cuda")
+            compiled = torch.compile(model, mode="reduce-overhead")
+            out = compiled(inp)
+
+            self.assertEqual(out, ((inp + 1) * 2).sin())
+            self.assertIsNone(self.get_manager())
+
+        def test_cudagraph_annotation_context_manager(self):
+            """Context manager disables cudagraphs for all graphs within block,
+            including across graph breaks."""
+
+            def model(x):
+                y = x + 1  # g1: outside context, cudagraphs ON
+                torch._dynamo.graph_break()
+                with torch._dynamo.override_cudagraphs(fwd=False, bwd=False):
+                    z = y * 2  # g2: inside context, cudagraphs OFF
+                    torch._dynamo.graph_break()
+                    w = z - 1  # g3: still inside context, cudagraphs OFF
+                torch._dynamo.graph_break()
+                return w + 3  # g4: outside context, cudagraphs ON
+
+            inp = torch.randn(4, 4, device="cuda")
+            compiled = torch.compile(model, mode="reduce-overhead")
+            out = compiled(inp)
+            out = compiled(inp)
+
+            expected = ((inp + 1) * 2 - 1) + 3
+            self.assertEqual(out, expected)
+            # g1 and g4 use cudagraphs (2 graphs recorded)
+            self.assertEqual(self.get_manager().new_graph_id().id, 2)
+
+        def test_cudagraph_annotation_module_forward(self):
+            """Annotation on a module's forward method disables cudagraphs
+            when that module is called inside a compiled function."""
+
+            class MyModule(torch.nn.Module):
+                @torch._dynamo.override_cudagraphs(fwd=False, bwd=False)
+                def forward(self, x):
+                    return x + 1
+
+            mod = MyModule().cuda()
+
+            def model(x):
+                # g1: mod.forward has annotation -> cudagraphs OFF
+                y = mod(x)
+                return y * 2
+
+            inp = torch.randn(4, 4, device="cuda")
+            compiled = torch.compile(model, mode="reduce-overhead")
+            out = compiled(inp)
+            out = compiled(inp)
+
+            self.assertEqual(out, (inp + 1) * 2)
+            self.assertIsNone(self.get_manager())
+
+        @torch._inductor.config.patch("triton.cudagraphs", False)
+        def test_cudagraph_annotation_enable(self):
+            """override_cudagraphs(fwd=True, bwd=True) force-enables cudagraphs
+            even when config.triton.cudagraphs is False."""
+
+            @torch._dynamo.override_cudagraphs(fwd=True, bwd=True)
+            def helper(x):
+                return x + 1
+
+            def fn(x):
+                return helper(x)
+
+            fn_compiled = torch.compile(fn)
+
+            for _ in range(3):
+                inp = torch.randn(4, 4, device="cuda")
+                torch.compiler.cudagraph_mark_step_begin()
+                out = fn_compiled(inp)
+
+            self.assertEqual(out, inp + 1)
+            self.assertIsNotNone(self.get_manager())
+            self.assertEqual(self.get_manager().new_graph_id().id, 1)
+
+        @torch._inductor.config.patch("triton.cudagraphs", False)
+        def test_cudagraph_annotation_enable_fwd_bwd(self):
+            """override_cudagraphs(fwd=True, bwd=True) force-enables cudagraphs
+            for both forward and backward when config disables them."""
+
+            @torch._dynamo.override_cudagraphs(fwd=True, bwd=True)
+            def helper(x):
+                return (x * x).sin()
+
+            def fn(x):
+                return helper(x)
+
+            fn_compiled = torch.compile(fn)
+
+            for _ in range(3):
+                inp = torch.randn(4, 4, device="cuda", requires_grad=True)
+                torch.compiler.cudagraph_mark_step_begin()
+                out = fn_compiled(inp)
+                out.sum().backward()
+
+            self.assertEqual(out, inp.detach().pow(2).sin())
+            self.assertIsNotNone(self.get_manager())
+            # One forward and one backward
+            self.assertEqual(self.get_manager().new_graph_id().id, 2)
+
+        @torch._inductor.config.patch("triton.cudagraphs", False)
+        def test_cudagraph_annotation_enable_fwd_only(self):
+            """override_cudagraphs(fwd=True) force-enables cudagraphs for
+            forward. Backward is also enabled since the BoxedBool is shared."""
+
+            @torch._dynamo.override_cudagraphs(fwd=True)
+            def helper(x):
+                return (x * x).sin()
+
+            def fn(x):
+                return helper(x)
+
+            fn_compiled = torch.compile(fn)
+
+            for _ in range(3):
+                inp = torch.randn(4, 4, device="cuda", requires_grad=True)
+                torch.compiler.cudagraph_mark_step_begin()
+                out = fn_compiled(inp)
+                out.sum().backward()
+
+            self.assertEqual(out, inp.detach().pow(2).sin())
+            self.assertIsNotNone(self.get_manager())
+            # fwd override also enables bwd since the BoxedBool is shared
+            self.assertEqual(self.get_manager().new_graph_id().id, 2)
+
+        @torch._inductor.config.patch("triton.cudagraphs", False)
+        def test_cudagraph_annotation_enable_context_manager(self):
+            """Context manager force-enables cudagraphs for segments within
+            the block when config disables them."""
+
+            def model(x):
+                y = x + 1  # g1: no annotation, config off -> no cudagraphs
+                torch._dynamo.graph_break()
+                with torch._dynamo.override_cudagraphs(fwd=True, bwd=True):
+                    z = y * 2  # g2: annotation on -> cudagraphs
+                    torch._dynamo.graph_break()
+                    w = z - 1  # g3: annotation on -> cudagraphs
+                torch._dynamo.graph_break()
+                return w + 3  # g4: no annotation, config off -> no cudagraphs
+
+            compiled = torch.compile(model)
+            for _ in range(3):
+                inp = torch.randn(4, 4, device="cuda")
+                torch.compiler.cudagraph_mark_step_begin()
+                out = compiled(inp)
+
+            expected = ((inp + 1) * 2 - 1) + 3
+            self.assertEqual(out, expected)
+            # g2 and g3 use cudagraphs (2 graphs recorded)
+            self.assertEqual(self.get_manager().new_graph_id().id, 2)
+
+        def test_cudagraph_min_partition_size_skip_small(self):
+            """Partitions with fewer kernels than the threshold are skipped."""
+
+            # Simple function with few kernels
+            def fn(x):
+                return x * 2
+
+            inp = torch.randn(4, 4, device="cuda")
+
+            # Set threshold high enough to skip this simple function
+            with torch._inductor.config.patch(
+                {"triton.cudagraph_min_partition_size": 10}
+            ):
+                fn_compiled = torch.compile(fn, mode="reduce-overhead")
+                for _ in range(3):
+                    out = fn_compiled(inp)
+
+            self.assertEqual(out, inp * 2)
+            # Partition was too small, so no cudagraph recorded
+            self.assertIsNone(self.get_manager())
+
+        def test_cudagraph_min_partition_size_allow_large(self):
+            """Partitions with enough kernels pass the threshold."""
+
+            # Function with multiple operations -> more kernels
+            def fn(x):
+                y = x * 2
+                y = y + 1
+                y = y.sin()
+                y = y.cos()
+                y = y * x
+                return y @ y
+
+            inp = torch.randn(4, 4, device="cuda")
+
+            # Set threshold low enough to allow this function
+            with torch._inductor.config.patch(
+                {"triton.cudagraph_min_partition_size": 2}
+            ):
+                fn_compiled = torch.compile(fn, mode="reduce-overhead")
+                for _ in range(3):
+                    out = fn_compiled(inp)
+
+            # Partition was large enough, cudagraph should be recorded
+            self.assertIsNotNone(self.get_manager())
+            self.assertEqual(self.get_manager().new_graph_id().id, 1)
+
         def test_tensor_constant_mutation(self):
             class Foo(torch.nn.Module):
                 def __init__(self) -> None:
@@ -3105,20 +3336,8 @@ def forward(self, x) -> torch.Tensor:
             foo.static_tensor = torch.ones((2, 2), device="cuda")
             foo.goo.linear.bias = torch.nn.Parameter(torch.ones((2,), device="cuda"))
 
-            if torch._dynamo.config.inline_inbuilt_nn_modules:
-                for _ in range(3):
-                    foo(inp)
-            else:
-                # Run with specific function id to avoid dynamo recompiling
-                self.get_manager().run(
-                    [
-                        foo.goo.linear.weight,
-                        foo.goo.linear.bias,
-                        foo.static_tensor,
-                        inp,
-                    ],
-                    FunctionID(0),
-                )
+            for _ in range(3):
+                foo(inp)
 
             self.assertEqual(self.get_manager().new_graph_id().id, 2)
 
@@ -3311,7 +3530,7 @@ def foo(x, y):
                 foo(torch.ones([10], device="cuda"), torch.ones([20]))
 
             FileCheck().check_count(
-                "Recording function=partition_0, cuda_graph_id=0",
+                "Recording function=partition_0, mode=INFERENCE, cuda_graph_id=0,",
                 1,
                 exactly=True,
             ).run(log_stream.getvalue())
@@ -3515,6 +3734,185 @@ def forward(self, x):
             # 2 graph partitions lead to 2 fwd cudagraphs and 1 bwd cudagraphs
             self.assertEqual(self.get_manager().new_graph_id().id, 3)
 
+        @torch._inductor.config.patch("graph_partition", True)
+        def test_graph_partition_cpu_scalar_used_in_cpu_op(self):
+            class Mod(torch.nn.Module):
+                def __init__(self) -> None:
+                    super().__init__()
+                    self.linear = torch.nn.Linear(4, 4, device="cuda")
+                    self.cpu_scale = torch.nn.Parameter(torch.tensor(1.0))
+
+                def forward(self, x):
+                    y = self.linear(x)
+                    cpu_val = y.sum().cpu() * self.cpu_scale
+                    return y + cpu_val.cuda()
+
+            model = Mod()
+            x = torch.randn(4, 4, device="cuda")
+
+            compiled_model = torch.compile(model, mode="reduce-overhead")
+
+            # Verify graph_partition produces the expected partitions.
+            log_stream, ctx = logs_to_string("torch._inductor.scheduler", "cudagraphs")
+            with ctx():
+                compiled_model(x)
+            FileCheck().check("2 cudagraphable, 1 non-cudagraphable").run(
+                log_stream.getvalue()
+            )
+
+            criterion = torch.nn.CrossEntropyLoss()
+            optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
+
+            for _ in range(5):
+                output = compiled_model(x)
+                loss = criterion(output, torch.randint(0, 4, (4,)).cuda())
+                optimizer.zero_grad()
+                loss.backward()
+                optimizer.step()
+
+            eager_out = model(x)
+            compiled_out = compiled_model(x)
+            self.assertEqual(compiled_out, eager_out)
+
+        @torch._inductor.config.patch("graph_partition", True)
+        def test_graph_partition_cpu_scalar_as_output(self):
+            """
+            When a CPU scalar placeholder is moved to GPU by ConstructorMoverPass,
+            the forward graph's output must NOT replace the CPU placeholder with
+            the GPU copy.
+            """
+
+            class Mod(torch.nn.Module):
+                def __init__(self) -> None:
+                    super().__init__()
+                    self.linear = torch.nn.Linear(4, 4, device="cuda")
+                    self.cpu_scale = torch.nn.Parameter(torch.tensor(1.0))
+
+                def forward(self, x):
+                    return self.linear(x) * self.cpu_scale
+
+            model = Mod()
+            x = torch.randn(4, 4, device="cuda")
+
+            compiled_model = torch.compile(model, mode="reduce-overhead")
+            compiled_model(x)
+
+            criterion = torch.nn.CrossEntropyLoss()
+            optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
+
+            for _ in range(5):
+                output = compiled_model(x)
+                loss = criterion(output, torch.randint(0, 4, (4,)).cuda())
+                optimizer.zero_grad()
+                loss.backward()
+                optimizer.step()
+
+            eager_out = model(x)
+            compiled_out = compiled_model(x)
+            self.assertEqual(compiled_out, eager_out)
+
+        @torch._inductor.config.patch("graph_partition", True)
+        def test_graph_partition_saved_activation_not_static(self):
+            # When the forward is partitioned, saved activations produced
+            # by inline code between partitions (e.g., DeviceCopy) are
+            # NOT at fixed addresses. The backward must not mark them as
+            # static inputs, or it would re-record on every iteration.
+            # Primals (params/buffers) should still be marked static.
+            from unittest.mock import patch
+
+            from torch._inductor.utils import count_tangents, get_static_bw_input_idxs
+
+            bw_graph = None
+            orig_bw = torch._inductor.compile_fx.compile_fx_backward
+
+            def intercept_bw(gm, example_inputs, compiler_config_extra, **kwargs):
+                nonlocal bw_graph
+                bw_graph = gm
+                return orig_bw(gm, example_inputs, compiler_config_extra, **kwargs)
+
+            class Mod(torch.nn.Module):
+                def __init__(self) -> None:
+                    super().__init__()
+                    self.linear = torch.nn.Linear(16, 16)
+
+                def forward(self, x):
+                    a = x * 2
+                    # CPU round-trip creates a DeviceCopy partition boundary.
+                    # The .cuda() result is an activation saved for backward.
+                    b = a.cpu().cuda()
+                    c = b * b
+                    return self.linear(c)
+
+            model = Mod().cuda()
+            input_data = torch.randn(16, 16, device="cuda")
+            criterion = torch.nn.CrossEntropyLoss()
+            optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
+
+            with patch("torch._inductor.compile_fx.compile_fx_backward", intercept_bw):
+                compiled_model = torch.compile(model, mode="reduce-overhead")
+                output = compiled_model(input_data)
+                loss = criterion(output, torch.randint(0, 10, (16,)).cuda())
+                optimizer.zero_grad()
+                loss.backward()
+                optimizer.step()
+
+            self.assertIsNotNone(bw_graph)
+            # count_tangents marks ALL saved tensors as static (old behavior)
+            all_static = list(range(count_tangents(bw_graph)))
+            # get_static_bw_input_idxs only marks primals as static
+            primal_static = get_static_bw_input_idxs(bw_graph)
+            # With a partitioned forward, only primals should be static,
+            # so primal_static should be a strict subset of all_static.
+            self.assertTrue(len(primal_static) < len(all_static))
+            for idx in primal_static:
+                self.assertIn(idx, all_static)
+
+            # Run a few more iterations to confirm stability
+            for _ in range(4):
+                output = compiled_model(input_data)
+                loss = criterion(output, torch.randint(0, 10, (16,)).cuda())
+                optimizer.zero_grad()
+                loss.backward()
+                optimizer.step()
+
+        @torch._inductor.config.patch("graph_partition", True)
+        def test_graph_partition_no_partition_keeps_static(self):
+            # When graph_partition is enabled but the forward has no unsafe
+            # ops, forward_is_partitioned should be False and all saved
+            # tensors remain static in the backward.
+            from unittest.mock import patch
+
+            forward_partitioned = None
+            orig_bw = torch._inductor.compile_fx.compile_fx_backward
+
+            def intercept_bw(gm, example_inputs, compiler_config_extra, **kwargs):
+                nonlocal forward_partitioned
+                forward_partitioned = compiler_config_extra.forward_is_partitioned.value
+                return orig_bw(gm, example_inputs, compiler_config_extra, **kwargs)
+
+            class Mod(torch.nn.Module):
+                def __init__(self) -> None:
+                    super().__init__()
+                    self.linear = torch.nn.Linear(16, 16)
+
+                def forward(self, x):
+                    return self.linear(x * x + 1)
+
+            model = Mod().cuda()
+            input_data = torch.randn(16, 16, device="cuda")
+            criterion = torch.nn.CrossEntropyLoss()
+            optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
+
+            with patch("torch._inductor.compile_fx.compile_fx_backward", intercept_bw):
+                compiled_model = torch.compile(model, mode="reduce-overhead")
+                output = compiled_model(input_data)
+                loss = criterion(output, torch.randint(0, 10, (16,)).cuda())
+                optimizer.zero_grad()
+                loss.backward()
+                optimizer.step()
+
+            self.assertFalse(forward_partitioned)
+
         @torch._inductor.config.patch("graph_partition", True)
         def test_graph_partition_cpu_only(self):
             class Mod(torch.nn.Module):
@@ -4189,7 +4587,8 @@ def run(padded_size, original_size):
                 eager_out = f(padded_tensor, original_tensor, weight)
                 for _ in range(3):
                     compiled_out = compiled_f(padded_tensor, original_tensor, weight)
-                    assert torch.allclose(eager_out, compiled_out)
+                    if not torch.allclose(eager_out, compiled_out):
+                        raise AssertionError
 
             # although custom op `create_mask` happens at the middle of function, reorder
             # moves it to the front so we only have 1 partition. This leads to 1 cudagraph
@@ -4710,6 +5109,221 @@ def run(n):
             run(10)
             run(25)
 
+        @torch._inductor.config.patch("triton.skip_cudagraph_warmup", True)
+        def test_opaque_value_input_cudagraph(self):
+            """Opaque reference-type objects (e.g. DeviceMesh, ProcessGroup)
+            passed alongside tensors must be marked "static" so they are
+            excluded from the tensor-copy path (non_static_input_idx).
+            "Static" here just means "don't copy as a tensor", not that
+            the object is semantically immutable."""
+
+            sf = _CudagraphTestScaleFactor(3.0)
+
+            def foo(args):
+                obj = args[0]
+                x = args[1]
+                args.clear()
+                return (x * obj.factor,)
+
+            inp = torch.rand([4], device="cuda")
+            foo_cg = self.cudagraphify_impl(foo, [sf, inp], ())
+            result = foo_cg([sf, inp])
+            self.assertEqual(result[0], inp * 3.0)
+
+            result2 = foo_cg([sf, inp])
+            self.assertEqual(result2[0], inp * 3.0)
+
+    class TestCUDAGraphPolicy(TestCase):
+        def setUp(self):
+            super().setUp()
+            counters.clear()
+            self._stack = contextlib.ExitStack()
+            self._stack.enter_context(
+                config.patch(
+                    {
+                        "triton.cudagraphs": True,
+                        "triton.cudagraph_trees": True,
+                    }
+                )
+            )
+            torch._dynamo.reset()
+
+        def tearDown(self):
+            super().tearDown()
+            torch._dynamo.reset()
+            self._stack.close()
+
+        def test_policy_cudagraphify_called(self):
+            """Custom policy's cudagraphify is called instead of the default."""
+            from torch._inductor.cudagraph_utils import CUDAGraphPolicy
+
+            calls = []
+
+            class RecordingPolicy(CUDAGraphPolicy):
+                def cudagraphify(self, model, inputs, static_input_idxs, **kwargs):
+                    calls.append(
+                        {
+                            "static_input_idxs": static_input_idxs,
+                            "device_index": kwargs.get("device_index"),
+                        }
+                    )
+                    return model
+
+            def foo(x):
+                return x * x + 1
+
+            with config.patch("cudagraph_policy", RecordingPolicy()):
+                compiled = torch.compile(foo)
+                x = torch.randn(4, device="cuda")
+                compiled(x)
+                compiled(x)
+
+            self.assertGreater(len(calls), 0)
+
+        def test_should_wrap_false_skips_cudagraph(self):
+            """When should_wrap returns False, cudagraph wrapping is skipped."""
+            from torch._inductor.cudagraph_utils import CUDAGraphPolicy
+
+            class NoWrapPolicy(CUDAGraphPolicy):
+                def should_wrap(self, compiled_graph):
+                    return False
+
+            def foo(x):
+                return x + 1
+
+            with config.patch("cudagraph_policy", NoWrapPolicy()):
+                compiled = torch.compile(foo)
+                x = torch.randn(4, device="cuda")
+                result = compiled(x)
+                result = compiled(x)
+
+            self.assertEqual(result, x + 1)
+            self.assertGreater(counters["inductor"]["cudagraph_skips"], 0)
+
+        def test_wrap_output_called(self):
+            """Policy's wrap_output is called in BundledOutputCodeLoadable."""
+            from torch._inductor.cudagraph_utils import CUDAGraphPolicy
+
+            wrapped = []
+
+            class WrapPolicy(CUDAGraphPolicy):
+                def wrap_output(self, output_code):
+                    wrapped.append(type(output_code).__name__)
+                    return output_code
+
+            def foo(x):
+                return x * 2
+
+            with config.patch("cudagraph_policy", WrapPolicy()):
+                compiled = torch.compile(foo)
+                x = torch.randn(4, device="cuda")
+                result = compiled(x)
+                result = compiled(x)
+
+            self.assertEqual(result, x * 2)
+            self.assertGreater(len(wrapped), 0)
+            self.assertIn("CompiledFxGraph", wrapped)
+
+        def test_default_policy_matches_builtin(self):
+            """Default CUDAGraphPolicy produces same results as no policy."""
+            from torch._inductor.cudagraph_utils import CUDAGraphPolicy
+
+            def foo(x):
+                return x * x + x
+
+            x = torch.randn(4, device="cuda")
+
+            compiled_default = torch.compile(foo)
+            ref = compiled_default(x)
+            ref = compiled_default(x)
+
+            torch._dynamo.reset()
+
+            with config.patch("cudagraph_policy", CUDAGraphPolicy()):
+                compiled_policy = torch.compile(foo)
+                out = compiled_policy(x)
+                out = compiled_policy(x)
+
+            self.assertEqual(ref, out)
+
+        @torch._inductor.config.patch("graph_partition", True)
+        def test_default_policy_matches_builtin_partition(self):
+            """Default CUDAGraphPolicy matches builtin when graph_partition=True."""
+            from torch._inductor.cudagraph_utils import CUDAGraphPolicy
+
+            def foo(x):
+                return x * x + x
+
+            x = torch.randn(4, device="cuda")
+
+            compiled_default = torch.compile(foo)
+            ref = compiled_default(x)
+            ref = compiled_default(x)
+
+            torch._dynamo.reset()
+
+            with config.patch("cudagraph_policy", CUDAGraphPolicy()):
+                compiled_policy = torch.compile(foo)
+                out = compiled_policy(x)
+                out = compiled_policy(x)
+
+            self.assertEqual(ref, out)
+
+        @torch._inductor.config.patch("graph_partition", True)
+        def test_policy_cudagraphify_partition(self):
+            """Custom policy's cudagraphify is called when graph_partition is enabled."""
+            from torch._inductor.cudagraph_utils import CUDAGraphPolicy
+
+            calls = []
+
+            class RecordingPolicy(CUDAGraphPolicy):
+                def cudagraphify(self, model, inputs, static_input_idxs, **kwargs):
+                    calls.append(
+                        {
+                            "static_input_idxs": static_input_idxs,
+                            "device_index": kwargs.get("device_index"),
+                        }
+                    )
+                    return model
+
+            def foo(x):
+                return x * x + 1
+
+            with config.patch("cudagraph_policy", RecordingPolicy()):
+                compiled = torch.compile(foo)
+                x = torch.randn(4, device="cuda")
+                compiled(x)
+                compiled(x)
+
+            self.assertGreater(len(calls), 0)
+
+        def test_should_wrap_false_with_wrap_output(self):
+            """should_wrap=False skips inner wrapping; wrap_output does outer wrapping."""
+            from torch._inductor.cudagraph_utils import CUDAGraphPolicy
+
+            wrap_calls = []
+
+            class OuterOnlyPolicy(CUDAGraphPolicy):
+                def should_wrap(self, compiled_graph):
+                    return False
+
+                def wrap_output(self, output_code):
+                    wrap_calls.append(type(output_code).__name__)
+                    return output_code
+
+            def foo(x):
+                return x * x + 1
+
+            with config.patch("cudagraph_policy", OuterOnlyPolicy()):
+                compiled = torch.compile(foo)
+                x = torch.randn(4, device="cuda")
+                result = compiled(x)
+                result = compiled(x)
+
+            self.assertEqual(result, x * x + 1)
+            self.assertGreater(counters["inductor"]["cudagraph_skips"], 0)
+            self.assertGreater(len(wrap_calls), 0)
+
     class TestSAC(TestCase):
         def _make_observer_mode(self):
             class ObserverMode(TorchDispatchMode):
@@ -4879,7 +5493,7 @@ def _(mode, *args, **kwargs):
                             if isinstance(inp, torch._ops.OpOverload):
                                 op = inp
                                 break
-                        assert op is not None
+                        assert op is not None  # noqa: S101
                         if hop is run_and_save_rng_state:
                             mode.op_outputs[op].append(out[1])
                         else:
@@ -5113,6 +5727,7 @@ def fn(x, y):
             )
 
     instantiate_parametrized_tests(CudaGraphTreeTests)
+    instantiate_parametrized_tests(TestCUDAGraphPolicy)
     instantiate_parametrized_tests(TestSAC)
 
     # OpInfo-based test for index/scatter ops with cudagraphs
diff --git a/test/inductor/test_custom_lowering.py b/test/inductor/test_custom_lowering.py
index 33867186d5eb8..478b5768f5156 100644
--- a/test/inductor/test_custom_lowering.py
+++ b/test/inductor/test_custom_lowering.py
@@ -164,8 +164,10 @@ def foo(x: torch.Tensor) -> torch.Tensor:
         def foo_lowering(x):
             return x
 
-        assert torch.ops.helion_test.foo in custom_lowering_dict
-        assert torch.ops.helion_test.foo not in torch._inductor.lowering.lowerings
+        if torch.ops.helion_test.foo not in custom_lowering_dict:
+            raise AssertionError
+        if torch.ops.helion_test.foo in torch._inductor.lowering.lowerings:
+            raise AssertionError
 
     @requires_gpu()
     @skipIf(GPU_TYPE == "mps", "Not applicable to MPS")
diff --git a/test/inductor/test_custom_op_autotune.py b/test/inductor/test_custom_op_autotune.py
index a524d928875a2..d5b9572fa54f3 100644
--- a/test/inductor/test_custom_op_autotune.py
+++ b/test/inductor/test_custom_op_autotune.py
@@ -7,28 +7,55 @@
 Inductor benchmarks all variants and automatically selects the best performing one.
 """
 
+import unittest
+
 import torch
+import torch._inductor.runtime.benchmarking
+from torch._dynamo.utils import counters
 from torch._inductor import config
 from torch._inductor.kernel.custom_op import (
     CustomOpConfig,
     register_custom_op_autotuning,
 )
 from torch._inductor.test_case import run_tests, TestCase
-from torch.testing._internal.common_utils import skipIfXpu
-from torch.testing._internal.inductor_utils import HAS_GPU
+from torch.testing import FileCheck
+from torch.testing._internal.common_utils import (
+    instantiate_parametrized_tests,
+    IS_MACOS,
+    parametrize,
+    skipIfXpu,
+)
+from torch.testing._internal.inductor_utils import (
+    GPU_TYPE,
+    HAS_CPU,
+    HAS_GPU,
+    HAS_TRITON,
+    IS_BIG_GPU,
+)
 
 
 torch.set_float32_matmul_precision("high")
 
 
+@unittest.skipIf(IS_MACOS, "TODO: mac")
+@unittest.skipUnless(HAS_GPU and HAS_TRITON, "requires GPU and Triton")
 class TestCustomOpAutoTune(TestCase):
     """Test custom operation autotuning functionality."""
 
     def setUp(self) -> None:
         """Set up test environment with appropriate device and dtype."""
         super().setUp()
-        self.device = "cuda" if HAS_GPU else "cpu"
-        self.dtype = torch.float16 if self.device == "cuda" else torch.float32
+        torch._dynamo.reset()
+        self.device = GPU_TYPE if HAS_GPU else "cpu"
+        self.dtype = (
+            torch.float16
+            if self.device == "cuda" or self.device == "xpu"
+            else torch.float32
+        )
+        # Clear any previous lowering registrations to ensure test isolation
+        from torch._inductor.lowering import user_lowerings
+
+        user_lowerings.clear()
 
     def _run_autotune_test(self, op_object, inputs, expected, test_name):
         """Shared test infrastructure for autotuning tests."""
@@ -142,7 +169,6 @@ def _create_mlp_inputs(
         )
         return input_tensor, gate_weight, up_weight, down_weight
 
-    @skipIfXpu
     def test_rmsnorm_custom_op_autotune_with_dynamic_shape(self):
         """Test RMSNorm autotuning with multiple decomposition variants and dynamic shapes.
 
@@ -192,7 +218,8 @@ def _(input_tensor: torch.Tensor, weight: torch.Tensor, eps: float = 1e-8):
             configs=[CustomOpConfig(decomp) for decomp in decompositions],
             name="test_rmsnorm_autotuned",
             input_gen_fns={
-                "x": lambda x: torch.randn_like(x, device=self.device) * 0.02,
+                "input_tensor": lambda x: torch.randn_like(x, device=self.device)
+                * 0.02,
                 "weight": lambda weight: torch.ones_like(weight, device=self.device),
             },
         )
@@ -232,7 +259,6 @@ def _create_decompose_k_inputs(self, m=256, k=65536, n=1024):
         )
         return a, b, bias
 
-    @skipIfXpu
     def test_decompose_k_custom_op_autotune_dynamic_config_for_input_shape(self):
         """Test decompose_k autotuning with with epilogue fusion(matmul+bias+relu+scale) and
         dynamic config generation based on matmul input shapes.
@@ -353,7 +379,6 @@ def reference_model(a, b, bias):
                 msg=f"Failed for shape ({m}, {k}, {n})",
             )
 
-    @skipIfXpu
     def test_multi_parameter_tuning(self):
         """Test autotuning with multiple parameters for combinatorial parameter exploration.
 
@@ -453,7 +478,6 @@ def _(
             multi_param_op, (test_x, test_factor), expected_result, "MultiParam"
         )
 
-    @skipIfXpu
     def test_range_based_static_shape_no_cond_dispatch(self):
         """Test dispatch code generation for static vs dynamic shapes.
 
@@ -556,6 +580,946 @@ def test_model(x, weight):
             print("[Dynamic] No dispatch logic found (unexpected for dynamic shapes)")
         self.assertTrue(dispatch_dynamic, "Dynamic shapes should have dispatch logic")
 
+    @skipIfXpu
+    def test_benchmark_with_cudagraphs_uses_cuda_graph_benchmarking(self):
+        """Test that benchmark_with_cudagraphs flag causes CUDA graph benchmarking to be used."""
+        if self.device != "cuda":
+            self.skipTest("CUDA graph test requires CUDA device")
+
+        from unittest.mock import patch
+
+        test_op_name = f"test_lib::cudagraph_patch_{id(self)}"
+
+        def fast_decomposition(x: torch.Tensor, weight: torch.Tensor) -> torch.Tensor:
+            return x @ weight
+
+        @torch.library.custom_op(test_op_name, mutates_args=())
+        def cudagraph_patch_op(x: torch.Tensor, weight: torch.Tensor) -> torch.Tensor:
+            return x @ weight
+
+        @cudagraph_patch_op.register_fake
+        def _(x: torch.Tensor, weight: torch.Tensor):
+            return torch.empty(
+                x.shape[0], weight.shape[1], device=x.device, dtype=x.dtype
+            )
+
+        register_custom_op_autotuning(
+            cudagraph_patch_op,
+            configs=[CustomOpConfig(fast_decomposition)],
+            name="cudagraph_patch_autotuned",
+            benchmark_with_cudagraphs=True,
+            input_gen_fns={
+                "x": lambda t: torch.randn_like(t, device=self.device),
+                "weight": lambda t: torch.randn_like(t, device=self.device),
+            },
+        )
+
+        test_x = torch.randn(64, 256, device=self.device, dtype=self.dtype)
+        test_weight = torch.randn(256, 128, device=self.device, dtype=self.dtype)
+
+        @torch.compile
+        def test_model(x, weight):
+            return cudagraph_patch_op(x, weight)
+
+        cuda_graph_benchmark_called = False
+        original_benchmark_gpu_with_cuda_graph = torch._inductor.runtime.benchmarking.Benchmarker.benchmark_gpu_with_cuda_graph
+
+        def patched_benchmark_gpu_with_cuda_graph(self, fn):
+            nonlocal cuda_graph_benchmark_called
+            cuda_graph_benchmark_called = True
+            return original_benchmark_gpu_with_cuda_graph(self, fn)
+
+        torch._dynamo.reset()
+        with config.patch(max_autotune=True, fx_graph_cache=False):
+            with patch.object(
+                torch._inductor.runtime.benchmarking.Benchmarker,
+                "benchmark_gpu_with_cuda_graph",
+                patched_benchmark_gpu_with_cuda_graph,
+            ):
+                result = test_model(test_x, test_weight)
+
+        self.assertTrue(
+            cuda_graph_benchmark_called,
+            "benchmark_gpu_with_cuda_graph should have been called",
+        )
+        torch.testing.assert_close(result, test_x @ test_weight, rtol=1e-1, atol=1e-1)
+
+    @skipIfXpu
+    def test_min_speedup_threshold_api(self):
+        """Test that min_speedup_threshold parameter is accepted and compilation works."""
+        test_op_name = f"test_lib::min_speedup_{id(self)}"
+
+        def decomposition(x: torch.Tensor, weight: torch.Tensor) -> torch.Tensor:
+            return x @ weight
+
+        @torch.library.custom_op(test_op_name, mutates_args=())
+        def min_speedup_op(x: torch.Tensor, weight: torch.Tensor) -> torch.Tensor:
+            return x @ weight
+
+        @min_speedup_op.register_fake
+        def _(x: torch.Tensor, weight: torch.Tensor):
+            return torch.empty(
+                x.shape[0], weight.shape[1], device=x.device, dtype=x.dtype
+            )
+
+        # Test that API accepts min_speedup_threshold parameter
+        register_custom_op_autotuning(
+            min_speedup_op,
+            configs=[CustomOpConfig(decomposition)],
+            name="min_speedup_autotuned",
+            min_speedup_threshold=1.5,
+            input_gen_fns={
+                "x": lambda t: torch.randn_like(t, device=self.device),
+                "weight": lambda t: torch.randn_like(t, device=self.device),
+            },
+        )
+
+        test_x = torch.randn(64, 256, device=self.device, dtype=self.dtype)
+        test_weight = torch.randn(256, 128, device=self.device, dtype=self.dtype)
+
+        @torch.compile
+        def test_model(x, weight):
+            return min_speedup_op(x, weight)
+
+        torch._dynamo.reset()
+        with config.patch(max_autotune=True, fx_graph_cache=False):
+            result = test_model(test_x, test_weight)
+
+        torch.testing.assert_close(result, test_x @ test_weight, rtol=1e-1, atol=1e-1)
+
+    @skipIfXpu
+    def test_config_patching_in_generated_code(self):
+        """Test that coordinate_descent_tuning config_patches flows through to generated code."""
+        if self.device != "cuda":
+            self.skipTest(
+                "coordinate_descent_tuning test requires CUDA for Triton codegen"
+            )
+
+        test_op_name = f"test_lib::coord_descent_{id(self)}"
+
+        def decomposition(x: torch.Tensor, weight: torch.Tensor) -> torch.Tensor:
+            return torch.ones_like(x)
+
+        @torch.library.custom_op(test_op_name, mutates_args=())
+        def coord_op(x: torch.Tensor, weight: torch.Tensor) -> torch.Tensor:
+            return x @ weight
+
+        @coord_op.register_fake
+        def _(x: torch.Tensor, weight: torch.Tensor):
+            return torch.empty(
+                x.shape[0], weight.shape[1], device=x.device, dtype=x.dtype
+            )
+
+        # Register with config_patches containing coordinate_descent_tuning
+        register_custom_op_autotuning(
+            coord_op,
+            configs=[
+                CustomOpConfig(
+                    decomposition, config_patches={"coordinate_descent_tuning": True}
+                )
+            ],
+            name="coord_descent_autotuned",
+            input_gen_fns={
+                "x": lambda t: torch.randn_like(t, device=self.device),
+                "weight": lambda t: torch.randn_like(t, device=self.device),
+            },
+        )
+
+        test_x = torch.randn(1024, 1024, device=self.device, dtype=self.dtype)
+        test_weight = torch.randn(1024, 1024, device=self.device, dtype=self.dtype)
+
+        @torch.compile
+        def test_model(x, weight):
+            return coord_op(x, weight)
+
+        torch._dynamo.reset()
+
+        # Compile and capture generated code
+        with config.patch(max_autotune=True, fx_graph_cache=False):
+            result, code = torch._inductor.utils.run_and_get_code(
+                test_model, test_x, test_weight
+            )
+
+        # Check that coordinate_descent_tuning is enabled in the generated code's inductor_meta
+        FileCheck().check("'coordinate_descent_tuning': True").run("\n".join(code))
+
+    @skipIfXpu
+    @config.patch(
+        {
+            "test_configs.force_custom_op_decomposition": True,
+            "test_configs.force_no_impl_grouping": True,
+        }
+    )
+    def test_split_config_patching_in_generated_code(self):
+        """Test that coordinate_descent_tuning config_patches flows through to generated code."""
+        if self.device != "cuda":
+            self.skipTest(
+                "coordinate_descent_tuning test requires CUDA for Triton codegen"
+            )
+
+        test_op_name = f"test_lib::coord_descent_{id(self)}"
+
+        def decomposition(x: torch.Tensor, weight: torch.Tensor) -> torch.Tensor:
+            return torch.full(
+                (x.shape[0], weight.shape[1]),
+                x.shape[0],
+                dtype=weight.dtype,
+                device=weight.device,
+            )
+
+        @torch.library.custom_op(test_op_name, mutates_args=())
+        def coord_op(x: torch.Tensor, weight: torch.Tensor) -> torch.Tensor:
+            if x.shape[0] == 128:
+                return torch.empty(
+                    x.shape[0],
+                    weight.shape[1],
+                    dtype=weight.dtype,
+                    device=weight.device,
+                )
+            return x @ weight
+
+        @coord_op.register_fake
+        def _(x: torch.Tensor, weight: torch.Tensor):
+            return torch.empty(
+                x.shape[0], weight.shape[1], device=x.device, dtype=x.dtype
+            )
+
+        # Register with config_patches containing coordinate_descent_tuning
+        register_custom_op_autotuning(
+            coord_op,
+            configs=[
+                CustomOpConfig(
+                    decomposition, config_patches={"coordinate_descent_tuning": True}
+                )
+            ],
+            name="coord_descent_autotuned",
+            input_gen_fns={
+                "x": lambda t: torch.randn_like(t, device=self.device),
+                "weight": lambda t: torch.randn_like(t, device=self.device),
+            },
+            dispatch_on={"tensor_name": "x", "dim": 0},
+            split_points=[128, 512],
+        )
+
+        test_x = torch.randn(1024, 1024, device=self.device, dtype=self.dtype)
+        test_weight = torch.randn(1024, 1024, device=self.device, dtype=self.dtype)
+
+        @torch.compile
+        def test_model(x, weight):
+            return coord_op(x, weight)
+
+        torch._dynamo.mark_dynamic(test_x, 0)
+
+        result, code = torch._inductor.utils.run_and_get_code(
+            test_model, test_x, test_weight
+        )
+
+        #
+        with torch._dynamo.config.patch(error_on_recompile=True):
+            for i in range(10):
+                x = torch.randn(256 + i, 1024, device=self.device, dtype=self.dtype)
+                torch._dynamo.mark_dynamic(x, 0)
+                self.assertEqual(
+                    test_model(x, test_weight), decomposition(x, test_weight)
+                )
+
+        # Check that coordinate_descent_tuning is enabled in the generated code's inductor_meta.
+        # config_patches should flow to both benchmark graphs and final dispatch code.
+        code_with_coord = [c for c in code if "'coordinate_descent_tuning': True" in c]
+        self.assertEqual(
+            len(code_with_coord),
+            len(code),
+            f"Expected all {len(code)} code modules to have coordinate_descent_tuning, "
+            f"but only {len(code_with_coord)} have it",
+        )
+
+    @skipIfXpu
+    def test_benchmark_real_trace_symbolic(self):
+        """Verify benchmarking uses real values but tracing uses symbolic shapes."""
+        if self.device != "cuda":
+            self.skipTest("Test requires CUDA")
+
+        # Track shapes seen by the real op implementation
+        shapes_seen = []
+
+        test_op_name = f"test_lib::shape_tracker_{id(self)}"
+
+        def decomposition(x, weight):
+            return x @ weight
+
+        @torch.library.custom_op(test_op_name, mutates_args=())
+        def shape_tracker_op(x: torch.Tensor, weight: torch.Tensor) -> torch.Tensor:
+            # This runs during benchmarking with REAL values
+            shapes_seen.append(x.shape[0])
+            return x @ weight
+
+        @shape_tracker_op.register_fake
+        def _(x, weight):
+            return torch.empty(
+                x.shape[0], weight.shape[1], device=x.device, dtype=x.dtype
+            )
+
+        register_custom_op_autotuning(
+            shape_tracker_op,
+            configs=[CustomOpConfig(decomposition)],
+            name="shape_tracker_autotuned",
+            input_gen_fns={
+                "x": lambda t: torch.randn_like(t, device=self.device),
+                "weight": lambda t: torch.randn_like(t, device=self.device),
+            },
+            dispatch_on={"tensor_name": "x", "dim": 0},
+            split_points=[128, 512],
+        )
+
+        test_x = torch.randn(1024, 64, device=self.device, dtype=self.dtype)
+        test_weight = torch.randn(64, 32, device=self.device, dtype=self.dtype)
+
+        @torch.compile(dynamic=True)
+        def test_model(x, weight):
+            return shape_tracker_op(x, weight)
+
+        torch._dynamo.mark_dynamic(test_x, 0)
+        torch._dynamo.reset()
+        shapes_seen.clear()
+
+        with config.patch(max_autotune=True, fx_graph_cache=False):
+            result, code = torch._inductor.utils.run_and_get_code(
+                test_model, test_x, test_weight
+            )
+
+        # Verify we got concrete integers during benchmarking (not symbolic)
+        unique_shapes = sorted(set(shapes_seen))
+        for shape in unique_shapes:
+            self.assertIsInstance(shape, int, f"Expected int, got {type(shape)}")
+
+        # Verify we hit all 3 ranges during autotuning
+        ranges_hit = set()
+        for shape in shapes_seen:
+            if 1 <= shape <= 128:
+                ranges_hit.add("range_1_128")
+            elif 129 <= shape <= 512:
+                ranges_hit.add("range_129_512")
+            elif shape > 512:
+                ranges_hit.add("range_513_inf")
+
+        self.assertEqual(
+            len(ranges_hit),
+            3,
+            f"Expected 3 ranges hit during benchmarking, got {ranges_hit}",
+        )
+
+        # Verify tracing uses SYMBOLIC shapes in generated code
+        import re
+
+        has_symbolic = any(re.search(r"\bs\d+\b", c) for c in code)
+        self.assertTrue(has_symbolic, "Expected symbolic shapes in generated code")
+
+    @skipIfXpu
+    def test_torch_cond_with_shape_accessing_implementations(self):
+        """Test torch.cond dispatch with implementations that access tensor shapes.
+
+        Validates that implementations like decompose_k that access tensor shapes
+        (e.g., `m, k = mat1.shape`) work correctly with torch.cond dispatch.
+        The fix uses _build_cond_dispatch_graph to pre-trace each implementation.
+        """
+        test_op_name = f"test_lib::shape_access_cond_{id(self)}"
+
+        def shape_accessing_impl(
+            mat1: torch.Tensor, mat2: torch.Tensor
+        ) -> torch.Tensor:
+            m, k = mat1.shape  # Shape access that would break naive make_fx
+            n = mat2.shape[1]
+            k_splits = 4
+            if k % k_splits == 0:
+                k_parts = k // k_splits
+                a = torch.permute(mat1.reshape(m, k_splits, k_parts), (1, 0, 2))
+                b = mat2.reshape(k_splits, k_parts, n)
+                return torch.sum(torch.bmm(a, b), dim=0)
+            return mat1 @ mat2
+
+        def simple_impl(mat1: torch.Tensor, mat2: torch.Tensor) -> torch.Tensor:
+            return mat1 @ mat2
+
+        @torch.library.custom_op(test_op_name, mutates_args=())
+        def shape_access_op(mat1: torch.Tensor, mat2: torch.Tensor) -> torch.Tensor:
+            return mat1 @ mat2
+
+        @shape_access_op.register_fake
+        def _(mat1: torch.Tensor, mat2: torch.Tensor):
+            return torch.empty(
+                mat1.shape[0], mat2.shape[1], device=mat1.device, dtype=mat1.dtype
+            )
+
+        register_custom_op_autotuning(
+            shape_access_op,
+            configs=[CustomOpConfig(simple_impl), CustomOpConfig(shape_accessing_impl)],
+            name="shape_access_autotuned",
+            dispatch_on={"tensor_name": "mat1", "dim": 0},
+            split_points=[4, 16],
+            input_gen_fns={
+                "mat1": lambda t: torch.randn_like(t, device=self.device) * 0.1,
+                "mat2": lambda t: torch.randn_like(t, device=self.device) * 0.1,
+            },
+        )
+
+        test_mat1 = torch.randn(8, 64, device=self.device, dtype=self.dtype)
+        test_mat2 = torch.randn(64, 32, device=self.device, dtype=self.dtype)
+
+        @torch.compile(dynamic=True)
+        def test_model(mat1, mat2):
+            return shape_access_op(mat1, mat2)
+
+        torch._dynamo.reset()
+        with config.patch(max_autotune=True, fx_graph_cache=False):
+            result = test_model(test_mat1, test_mat2)
+
+        torch.testing.assert_close(result, test_mat1 @ test_mat2, rtol=1e-1, atol=1e-1)
+
+    @skipIfXpu
+    @unittest.skipIf(not IS_BIG_GPU, "Test requires large GPU memory")
+    def test_empty_config_generator_falls_back_to_triton(self):
+        """Test that empty config_generator falls back to normal mm autotuning.
+
+        When config_generator returns empty list, the user_lowering returns None
+        and graph.py falls back to the normal lowering (triton mm autotuning).
+        """
+        from torch._inductor.lowering import user_lowerings
+
+        # Config generator that returns empty - should trigger fallback
+        def empty_config_gen(fake_tensors):
+            return []
+
+        register_custom_op_autotuning(
+            torch.ops.aten.mm.default,
+            config_generator=empty_config_gen,
+            name="test_empty_fallback",
+        )
+
+        self.assertIn(torch.ops.aten.mm.default, user_lowerings)
+
+        # Use shapes that will trigger triton autotuning
+        test_a = torch.randn(64, 128, device=self.device, dtype=self.dtype)
+        test_b = torch.randn(128, 64, device=self.device, dtype=self.dtype)
+
+        @torch.compile
+        def test_model(a, b):
+            return torch.mm(a, b)
+
+        # Enable max_autotune with TRITON backend
+        with config.patch(
+            max_autotune=True,
+            max_autotune_gemm_backends="TRITON",
+            fx_graph_cache=False,
+        ):
+            result = test_model(test_a, test_b)
+
+        # Verify correctness
+        torch.testing.assert_close(
+            result, torch.mm(test_a, test_b), rtol=1e-1, atol=1e-1
+        )
+
+    @skipIfXpu
+    @config.patch({"test_configs.force_custom_op_decomposition": True})
+    def test_guard_safety_drops_unsafe_decomposition(self):
+        """Test that decompositions adding guards are replaced with fallback.
+
+        Compiles with m=8 (divisible by 4, satisfies the unsafe impl's guard),
+        then calls with m=7 (not divisible by 4). If the Mod guard leaked,
+        the m=7 call would either crash or produce incorrect results because
+        the compiled code assumes m is divisible by 4.
+        """
+        test_op_name = f"test_lib::guard_safety_{id(self)}"
+
+        def unsafe_impl(mat1: torch.Tensor, mat2: torch.Tensor) -> torch.Tensor:
+            m = mat1.shape[0]
+            if m % 4 == 0:
+                n = mat2.shape[1]
+                k = mat1.shape[1]
+                m_parts = m // 4
+                a = mat1.reshape(4, m_parts, k)
+                result = torch.bmm(a, mat2.unsqueeze(0).expand(4, -1, -1))
+                return result.reshape(m, n)
+            return mat1 @ mat2
+
+        def safe_impl(mat1: torch.Tensor, mat2: torch.Tensor) -> torch.Tensor:
+            return mat1 @ mat2
+
+        @torch.library.custom_op(test_op_name, mutates_args=())
+        def guard_safety_op(mat1: torch.Tensor, mat2: torch.Tensor) -> torch.Tensor:
+            return mat1 @ mat2
+
+        @guard_safety_op.register_fake
+        def _(mat1: torch.Tensor, mat2: torch.Tensor):
+            return torch.empty(
+                mat1.shape[0], mat2.shape[1], device=mat1.device, dtype=mat1.dtype
+            )
+
+        register_custom_op_autotuning(
+            guard_safety_op,
+            configs=[CustomOpConfig(unsafe_impl)],
+            name="guard_safety_autotuned",
+            dispatch_on={"tensor_name": "mat1", "dim": 0},
+            split_points=[4, 16],
+            input_gen_fns={
+                "mat1": lambda t: torch.randn_like(t, device=self.device) * 0.1,
+                "mat2": lambda t: torch.randn_like(t, device=self.device) * 0.1,
+            },
+        )
+
+        @torch.compile(dynamic=True)
+        def test_model(mat1, mat2):
+            return guard_safety_op(mat1, mat2)
+
+        torch._dynamo.reset()
+        counters.clear()
+
+        # First call: m=8 (divisible by 4) triggers compilation
+        mat1 = torch.randn(8, 64, device=self.device, dtype=self.dtype)
+        mat2 = torch.randn(64, 32, device=self.device, dtype=self.dtype)
+        with config.patch(max_autotune=True, fx_graph_cache=False):
+            result = test_model(mat1, mat2)
+        torch.testing.assert_close(result, mat1 @ mat2, rtol=1e-1, atol=1e-1)
+
+        # Verify the guard safety counter was incremented
+        self.assertGreater(
+            counters["inductor"]["custom_op_decomp_guard_skips"],
+            0,
+            "Expected custom_op_decomp_guard_skips counter to be incremented",
+        )
+
+        # Second call: m=7 (NOT divisible by 4). If the unsafe impl was kept
+        # and its Mod(m,4)==0 guard leaked, this would crash on the reshape
+        # (can't reshape [7, 64] into [4, m_parts, 64] when 7 isn't divisible by 4).
+        mat1_odd = torch.randn(7, 64, device=self.device, dtype=self.dtype)
+        mat2_odd = torch.randn(64, 32, device=self.device, dtype=self.dtype)
+        with config.patch(max_autotune=True, fx_graph_cache=False):
+            result_odd = test_model(mat1_odd, mat2_odd)
+        torch.testing.assert_close(
+            result_odd, mat1_odd @ mat2_odd, rtol=1e-1, atol=1e-1
+        )
+
+    def test_fallback_choice_reuse(self):
+        """Test that _create_fallback_choice reuses the same choice for the same op.
+
+        Since kwargs are passed at bind time rather than baked into the kernel,
+        the same ExternKernelChoice is reused across compilations.
+        """
+        from torch._inductor.kernel.custom_op import _create_fallback_choice
+
+        test_op_name = f"test_lib::fallback_reuse_{id(self)}"
+
+        @torch.library.custom_op(test_op_name, mutates_args=())
+        def reuse_test_op(x: torch.Tensor) -> torch.Tensor:
+            return x.clone()
+
+        @reuse_test_op.register_fake
+        def _(x: torch.Tensor):
+            return torch.empty_like(x)
+
+        op_overload = reuse_test_op._opoverload
+
+        choice1 = _create_fallback_choice(op_overload)
+        choice2 = _create_fallback_choice(op_overload)
+
+        self.assertIs(choice1, choice2)
+
+    @skipIfXpu
+    @config.patch({"test_configs.force_custom_op_decomposition": False})
+    def test_fallback_different_kwargs(self):
+        """Test that the same fallback ExternKernelChoice works with different kwargs.
+
+        Previously, kwargs were baked into the fallback wrapper at choice creation
+        time, so a second call with different kwargs would either fail (duplicate
+        extern kernel) or use stale kwargs. Now kwargs are passed at bind time,
+        so both calls in the same compiled function get their own correct kwargs.
+        """
+        from torch._inductor.utils import run_and_get_code
+
+        test_op_name = "test_lib::fallback_kwargs"
+
+        def scale_impl(x: torch.Tensor, scale: float = 1.0) -> torch.Tensor:
+            return x * scale
+
+        @torch.library.custom_op(test_op_name, mutates_args=())
+        def kwargs_op(x: torch.Tensor, scale: float = 1.0) -> torch.Tensor:
+            return x * scale
+
+        @kwargs_op.register_fake
+        def _(x: torch.Tensor, scale: float = 1.0):
+            return torch.empty_like(x)
+
+        register_custom_op_autotuning(
+            kwargs_op,
+            configs=[CustomOpConfig(scale_impl)],
+            name="fallback_kwargs_test",
+        )
+
+        x = torch.ones(4, 4, device=self.device, dtype=self.dtype)
+
+        @torch.compile
+        def test_model(x):
+            a = kwargs_op(x, scale=2.0)
+            b = kwargs_op(x, scale=5.0)
+            return a, b
+
+        torch._dynamo.reset()
+        (result_a, result_b), codes = run_and_get_code(test_model, x)
+
+        torch.testing.assert_close(result_a, x * 2.0, rtol=1e-3, atol=1e-3)
+        torch.testing.assert_close(result_b, x * 5.0, rtol=1e-3, atol=1e-3)
+
+        # Both calls should use the same fallback kernel with different kwargs in codegen
+        code = "\n".join(codes)
+        self.assertIn("scale=2.0", code)
+        self.assertIn("scale=5.0", code)
+
+    @skipIfXpu
+    @config.patch({"test_configs.force_custom_op_decomposition": True})
+    def test_config_overrides_runtime_kwargs(self):
+        """Test that CustomOpConfig params override runtime default kwargs.
+
+        The op has scale=1.0 as default, but the config specifies scale=7.0.
+        With force_custom_op_decomposition=True, the decomposition should be
+        called with scale=7.0 and this should be reflected in the output code.
+        """
+        from torch._inductor.utils import run_and_get_code
+
+        test_op_name = f"test_lib::config_override_{id(self)}"
+
+        def scale_impl(x: torch.Tensor, scale: float = 1.0) -> torch.Tensor:
+            return x * scale
+
+        @torch.library.custom_op(test_op_name, mutates_args=())
+        def config_override_op(x: torch.Tensor, scale: float = 1.0) -> torch.Tensor:
+            return x * scale
+
+        @config_override_op.register_fake
+        def _(x: torch.Tensor, scale: float = 1.0):
+            return torch.empty_like(x)
+
+        register_custom_op_autotuning(
+            config_override_op,
+            configs=[CustomOpConfig(scale_impl, scale=7.0)],
+            name=f"config_override_test_{id(self)}",
+        )
+
+        x = torch.ones(4, 4, device=self.device, dtype=self.dtype)
+
+        @torch.compile
+        def test_model(x):
+            return config_override_op(x)
+
+        torch._dynamo.reset()
+        result, codes = run_and_get_code(test_model, x)
+
+        # The config's scale=7.0 should override the runtime default scale=1.0
+        torch.testing.assert_close(result, x * 7.0, rtol=1e-3, atol=1e-3)
+        code = "\n".join(codes)
+        self.assertIn("7.0", code)
+
+    @skipIfXpu
+    def test_input_gen_fns_invoked(self):
+        """Test that input_gen_fns are actually called during benchmarking."""
+        test_op_name = f"test_lib::input_gen_test_{id(self)}"
+        gen_calls = []
+
+        def decomp(x: torch.Tensor, weight: torch.Tensor) -> torch.Tensor:
+            return x * weight
+
+        @torch.library.custom_op(test_op_name, mutates_args=())
+        def input_gen_op(x: torch.Tensor, weight: torch.Tensor) -> torch.Tensor:
+            return x * weight
+
+        @input_gen_op.register_fake
+        def _(x: torch.Tensor, weight: torch.Tensor):
+            return torch.empty_like(x)
+
+        register_custom_op_autotuning(
+            input_gen_op,
+            configs=[CustomOpConfig(decomp)],
+            name=f"input_gen_test_{id(self)}",
+            input_gen_fns={
+                "x": lambda t: (gen_calls.append("x"), torch.ones_like(t))[1],
+                "weight": lambda t: (gen_calls.append("weight"), torch.ones_like(t))[1],
+            },
+        )
+
+        x = torch.randn(4, 4, device=self.device, dtype=self.dtype)
+        w = torch.randn(4, 4, device=self.device, dtype=self.dtype)
+
+        @torch.compile
+        def test_model(x, w):
+            return input_gen_op(x, w)
+
+        torch._dynamo.reset()
+        test_model(x, w)
+
+        self.assertIn("x", gen_calls)
+        self.assertIn("weight", gen_calls)
+
+    @skipIfXpu
+    @parametrize("force_choice", [None, True, False])
+    def test_kwargs_codegen(self, force_choice):
+        """Test that kwargs are correctly passed through codegen for both fallback and decomposition.
+
+        This validates that when a custom op is called with non-default kwargs:
+        1. Fallback path: kwargs flow through ExternKernelCaller to FallbackKernel.create
+        2. Decomposition path: kwargs are passed via functools.partial to the decomposition
+
+        Tests all paths via force_custom_op_decomposition:
+        - None: normal autotuning
+        - True: force decomposition
+        - False: force fallback
+        """
+        test_op_name = f"test_lib::kwargs_codegen_{id(self)}_{force_choice}"
+
+        def scale_impl(x: torch.Tensor, scale: float = 1.0) -> torch.Tensor:
+            return x * scale
+
+        @torch.library.custom_op(test_op_name, mutates_args=())
+        def kwargs_op(x: torch.Tensor, scale: float = 1.0) -> torch.Tensor:
+            return x * scale
+
+        @kwargs_op.register_fake
+        def _(x: torch.Tensor, scale: float = 1.0):
+            return torch.empty_like(x)
+
+        register_custom_op_autotuning(
+            kwargs_op,
+            configs=[CustomOpConfig(scale_impl)],
+            name=f"kwargs_codegen_test_{force_choice}",
+        )
+
+        x = torch.ones(4, 4, device=self.device, dtype=self.dtype)
+        scale_value = 3.0
+        expected = x * scale_value
+
+        @torch.compile
+        def test_model(x):
+            return kwargs_op(x, scale=scale_value)
+
+        torch._dynamo.reset()
+        with config.patch({"test_configs.force_custom_op_decomposition": force_choice}):
+            result = test_model(x)
+
+        path_name = {None: "autotuned", True: "decomposition", False: "fallback"}[
+            force_choice
+        ]
+        torch.testing.assert_close(
+            result, expected, rtol=1e-3, atol=1e-3, msg=f"{path_name} path failed"
+        )
+
+    @skipIfXpu
+    @config.patch(
+        {
+            "test_configs.force_custom_op_decomposition": True,
+        }
+    )
+    def test_shape_dependent_computation(self):
+        """Test that decompositions using shape in computation (e.g., x * x.shape[0]) work correctly.
+
+        This validates that make_fx tracing uses symbolic inputs so that shape-dependent
+        computations produce symbolic results (e.g., x * s0) rather than concrete values
+        (e.g., x * 512). The symbolic tracing is essential for correct codegen with
+        dynamic shapes.
+
+        Key validation: compile ONCE, then run with MULTIPLE sizes. If tracing used
+        concrete values, the result would be wrong for different sizes.
+
+        NOTE: We use a single-input op to avoid shape compatibility guards between
+        two different dynamic tensors (which would trigger guard safety filtering).
+        """
+        test_op_name = f"test_lib::shape_compute_{id(self)}"
+
+        def shape_compute_impl(x: torch.Tensor) -> torch.Tensor:
+            """Scale by first dimension size - captures x * s0 with symbolic tracing."""
+            return x * x.shape[0]
+
+        @torch.library.custom_op(test_op_name, mutates_args=())
+        def shape_compute_op(x: torch.Tensor) -> torch.Tensor:
+            # Fallback: just returns x (no shape dependency)
+            return x.clone()
+
+        @shape_compute_op.register_fake
+        def _(x: torch.Tensor):
+            return torch.empty_like(x)
+
+        register_custom_op_autotuning(
+            shape_compute_op,
+            configs=[CustomOpConfig(shape_compute_impl)],
+            name="shape_compute_autotuned",
+            input_gen_fns={
+                "x": lambda t: torch.randn_like(t, device=self.device) * 0.1,
+            },
+        )
+
+        @torch.compile(dynamic=True)
+        def test_model(x):
+            return shape_compute_op(x)
+
+        # Compile once with initial size
+        test_x = torch.randn(8, 32, 64, device=self.device, dtype=self.dtype)
+        torch._dynamo.mark_dynamic(test_x, 0)
+
+        with config.patch(max_autotune=True, fx_graph_cache=False):
+            result = test_model(test_x)
+
+        # Expected: decomposition result (x * x.shape[0]), NOT fallback (x.clone())
+        expected = test_x * test_x.shape[0]
+        torch.testing.assert_close(result, expected, rtol=1e-5, atol=1e-5)
+
+        # Now test with DIFFERENT sizes using the SAME compiled model
+        # If tracing used concrete values (e.g., x * 8), these would fail
+        for first_dim in [16, 32, 64, 128]:
+            test_x = torch.randn(
+                first_dim, 32, 64, device=self.device, dtype=self.dtype
+            )
+
+            result = test_model(test_x)
+            expected = test_x * first_dim  # x * x.shape[0] should use actual size
+
+            torch.testing.assert_close(
+                result,
+                expected,
+                rtol=1e-5,
+                atol=1e-5,
+                msg=f"Failed for shape[0]={first_dim}: symbolic tracing may have captured concrete value",
+            )
+
+    @skipIfXpu
+    def test_partial_input_gen_fns(self):
+        """Test autotuning when input_gen_fns covers only some inputs.
+
+        The uncovered inputs should fall back to ir_node_to_tensor with concrete hints.
+        """
+        test_op_name = f"test_lib::partial_gen_{id(self)}"
+
+        def decomposition(x: torch.Tensor, weight: torch.Tensor) -> torch.Tensor:
+            return x @ weight
+
+        @torch.library.custom_op(test_op_name, mutates_args=())
+        def partial_gen_op(x: torch.Tensor, weight: torch.Tensor) -> torch.Tensor:
+            return x @ weight
+
+        @partial_gen_op.register_fake
+        def _(x: torch.Tensor, weight: torch.Tensor):
+            return torch.empty(
+                x.shape[0], weight.shape[1], device=x.device, dtype=x.dtype
+            )
+
+        # Only provide input_gen_fn for "x", not "weight"
+        register_custom_op_autotuning(
+            partial_gen_op,
+            configs=[CustomOpConfig(decomposition)],
+            name="partial_gen_autotuned",
+            input_gen_fns={
+                "x": lambda t: torch.randn_like(t, device=self.device),
+            },
+        )
+
+        test_x = torch.randn(64, 256, device=self.device, dtype=self.dtype)
+        test_weight = torch.randn(256, 128, device=self.device, dtype=self.dtype)
+
+        @torch.compile
+        def test_model(x, weight):
+            return partial_gen_op(x, weight)
+
+        torch._dynamo.reset()
+        with config.patch(max_autotune=True, fx_graph_cache=False):
+            result = test_model(test_x, test_weight)
+
+        torch.testing.assert_close(result, test_x @ test_weight, rtol=1e-1, atol=1e-1)
+
+    @skipIfXpu
+    def test_cudagraph_memory_cleanup(self):
+        """Test that CUDA graph destruction automatically cleans up cuBLAS workspaces."""
+        if self.device != "cuda":
+            self.skipTest("CUDA graph test requires CUDA device")
+
+        # Clear everything first
+        torch.cuda.synchronize()
+        torch._C._cuda_clearCublasWorkspaces()
+
+        # Create test tensors and establish baseline with some mm activity
+        a = torch.randn(256, 256, device=self.device, dtype=self.dtype)
+        b = torch.randn(256, 256, device=self.device, dtype=self.dtype)
+        _ = torch.mm(a, b)  # This creates cublas workspace on default stream
+        torch.cuda.synchronize()
+
+        baseline_memory = torch.cuda.memory_allocated()
+
+        # Warmup on the stream
+        _ = torch.mm(a, b)
+        torch.cuda.synchronize()
+
+        # Capture into CUDA graph
+        graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(graph):
+            c = torch.mm(a, b)
+        torch.cuda.synchronize()
+
+        memory_after_capture = torch.cuda.memory_allocated()
+        self.assertGreater(
+            memory_after_capture, baseline_memory, "Capture should allocate memory"
+        )
+
+        # Deleting the graph should automatically clean up cuBLAS workspaces
+        del graph, c
+        torch.cuda.synchronize()
+
+        memory_after_cleanup = torch.cuda.memory_allocated()
+
+        self.assertEqual(
+            memory_after_cleanup,
+            baseline_memory,
+            f"Memory leak detected: baseline={baseline_memory}, after_cleanup={memory_after_cleanup}",
+        )
+
+    @skipIfXpu
+    def test_cudagraph_memory_cleanup_benchmarker(self):
+        """Test that CUDA graph benchmarking cleans up memory without leaking."""
+        if self.device != "cuda":
+            self.skipTest("CUDA graph test requires CUDA device")
+
+        # Clear everything first
+        torch.cuda.synchronize()
+        torch._C._cuda_clearCublasWorkspaces()
+
+        # Create test tensors
+        a = torch.randn(256, 256, device=self.device, dtype=self.dtype)
+        b = torch.randn(256, 256, device=self.device, dtype=self.dtype)
+
+        # Use the actual benchmarking infrastructure with CUDA graph capture
+        benchmarker = torch._inductor.runtime.benchmarking.benchmarker
+
+        def mm_callable():
+            return torch.mm(a, b)
+
+        # This should capture into CUDA graph, benchmark, and clean up properly
+        _ = benchmarker.benchmark_gpu_with_cuda_graph(mm_callable)
+        torch.cuda.synchronize()
+
+        memory_after_first = torch.cuda.memory_allocated()
+
+        # Run benchmarking again - memory should not grow
+        for _ in range(3):
+            _ = benchmarker.benchmark_gpu_with_cuda_graph(mm_callable)
+
+        memory_after_many = torch.cuda.memory_allocated()
+
+        # Memory should not grow significantly across multiple benchmark runs
+        self.assertEqual(
+            memory_after_many,
+            memory_after_first,
+            f"Memory leak detected: after_first={memory_after_first}, after_many={memory_after_many}",
+        )
+
+
+instantiate_parametrized_tests(TestCustomOpAutoTune)
+
 
 if __name__ == "__main__":
-    run_tests()
+    from torch._inductor.utils import is_big_gpu
+
+    if HAS_GPU and HAS_CPU and is_big_gpu():
+        run_tests()
diff --git a/test/inductor/test_custom_op_out_lowering.py b/test/inductor/test_custom_op_out_lowering.py
new file mode 100644
index 0000000000000..96ffd12e6a4fe
--- /dev/null
+++ b/test/inductor/test_custom_op_out_lowering.py
@@ -0,0 +1,141 @@
+# Owner(s): ["module: inductor"]
+"""
+Tests for inductor lowering of functional custom ops to out-variant via ExternKernelOut.
+"""
+
+import torch
+from torch._C import FileCheck
+from torch._inductor.test_case import TestCase as InductorTestCase
+from torch._inductor.utils import run_and_get_code
+from torch.testing._internal.common_utils import (
+    instantiate_parametrized_tests,
+    parametrize,
+)
+from torch.testing._internal.inductor_utils import GPU_TYPE, HAS_GPU
+
+
+DEVICES = ("cpu", GPU_TYPE) if HAS_GPU else ("cpu",)
+
+
+@instantiate_parametrized_tests
+class TestCustomOpOutLowering(InductorTestCase):
+    """Tests for lowering functional custom ops to out-variant ExternKernelOut."""
+
+    def _register_add_one_ops(self, lib):
+        """Register a simple add_one op with functional + .out overloads."""
+        lib.define("add_one(Tensor x) -> Tensor")
+        lib.define(
+            "add_one.out(Tensor x, *, Tensor(a!) out) -> Tensor(a!)",
+            tags=(torch.Tag.out_variant,),
+        )
+
+        def _add_one_impl(x: torch.Tensor) -> torch.Tensor:
+            return x + 1
+
+        def _add_one_out_impl(x: torch.Tensor, *, out: torch.Tensor) -> torch.Tensor:
+            out.copy_(x + 1)
+            return out
+
+        lib.impl("add_one", _add_one_impl, "CompositeExplicitAutograd")
+        lib.impl("add_one.out", _add_one_out_impl, "CompositeExplicitAutograd")
+
+        @torch.library.register_fake("mylib::add_one", lib=lib)
+        def _add_one_fake(x):
+            return x.new_empty(x.shape)
+
+        return torch.ops.mylib.add_one, torch.ops.mylib.add_one.out
+
+    @parametrize("device", DEVICES)
+    def test_add_one_lowered_to_out(self, device):
+        """Test that a simple functional op gets lowered to its out-variant."""
+        with torch.library._scoped_library("mylib", "FRAGMENT") as lib:
+            self._register_add_one_ops(lib)
+
+            def f(x):
+                return torch.ops.mylib.add_one(x)
+
+            x = torch.randn(4, 4, device=device)
+            eager_out = f(x)
+
+            compiled_out, (code,) = run_and_get_code(
+                torch.compile(f, backend="inductor", fullgraph=True), x
+            )
+            self.assertEqual(compiled_out, eager_out)
+
+            FileCheck().check(".out(").check_not(".default(").run(code)
+
+    def _register_split_add_ops(self, lib):
+        """Register a split_add op returning two tensors with functional + .out overloads."""
+        lib.define("split_add(Tensor x, float a, float b) -> (Tensor, Tensor)")
+        lib.define(
+            "split_add.out(Tensor x, float a, float b, *, Tensor(a!) out0, Tensor(b!) out1) -> (Tensor(a!), Tensor(b!))",
+            tags=(torch.Tag.out_variant,),
+        )
+
+        def _split_add_impl(x, a, b):
+            return (x + a, x + b)
+
+        def _split_add_out_impl(x, a, b, *, out0, out1):
+            out0.copy_(x + a)
+            out1.copy_(x + b)
+            return (out0, out1)
+
+        lib.impl("split_add", _split_add_impl, "CompositeExplicitAutograd")
+        lib.impl("split_add.out", _split_add_out_impl, "CompositeExplicitAutograd")
+
+        @torch.library.register_fake("mylib::split_add", lib=lib)
+        def _split_add_fake(x, a, b):
+            return (x.new_empty(x.shape), x.new_empty(x.shape))
+
+        return torch.ops.mylib.split_add, torch.ops.mylib.split_add.out
+
+    @parametrize("device", DEVICES)
+    def test_multi_output_lowered_to_out(self, device):
+        """Test a two-output functional op gets lowered to its .out variant."""
+        with torch.library._scoped_library("mylib", "FRAGMENT") as lib:
+            func_op, out_op = self._register_split_add_ops(lib)
+
+            def f(x):
+                a, b = torch.ops.mylib.split_add(x, 1.0, 2.0)
+                return a + b
+
+            x = torch.randn(4, 4, device=device)
+            eager_out = f(x)
+
+            compiled_out, (code,) = run_and_get_code(
+                torch.compile(f, backend="inductor", fullgraph=True), x
+            )
+            self.assertEqual(compiled_out, eager_out)
+            FileCheck().check(".out(").check("out0=").check("out1=").run(code)
+
+    @parametrize("device", DEVICES)
+    def test_op_without_out_variant_falls_through(self, device):
+        """Test that ops without an out-variant fall through to FallbackKernel."""
+        with torch.library._scoped_library("mylib", "FRAGMENT") as lib:
+            lib.define("no_out_op(Tensor x) -> Tensor")
+
+            def _impl(x):
+                return x + 1
+
+            lib.impl("no_out_op", _impl, "CompositeExplicitAutograd")
+
+            @torch.library.register_fake("mylib::no_out_op", lib=lib)
+            def _fake(x):
+                return x.new_empty(x.shape)
+
+            def f(x):
+                return torch.ops.mylib.no_out_op(x)
+
+            x = torch.randn(4, 4, device=device)
+            eager_out = f(x)
+
+            compiled_out, (code,) = run_and_get_code(
+                torch.compile(f, backend="inductor", fullgraph=True), x
+            )
+            self.assertEqual(compiled_out, eager_out)
+
+
+if __name__ == "__main__":
+    from torch._inductor.test_case import run_tests
+
+    run_tests()
diff --git a/test/inductor/test_custom_post_grad_passes.py b/test/inductor/test_custom_post_grad_passes.py
index c7823845bd570..e964add7ad4fe 100644
--- a/test/inductor/test_custom_post_grad_passes.py
+++ b/test/inductor/test_custom_post_grad_passes.py
@@ -12,6 +12,7 @@
 from torch._inductor.custom_graph_pass import (
     CustomGraphModulePass,
     CustomGraphPass,
+    CustomInferenceAwareGraphPass,
     get_hash_for_files,
 )
 from torch._inductor.lowering import lowerings as L
@@ -274,9 +275,11 @@ def f(W, nested_seqs):
             ]
 
             f(W, nested_seqs)
-            assert saved_graph[0] is not None
+            if saved_graph[0] is None:
+                raise AssertionError("saved_graph[0] is None")
             matmuls = [n for n in saved_graph[0].nodes if n.target == torch.mm]
-            assert len(matmuls) == 1
+            if len(matmuls) != 1:
+                raise AssertionError(f"Expected 1 matmul, got {len(matmuls)}")
 
         inner_test()
 
@@ -309,6 +312,31 @@ def f(x):
             x = torch.randn(8, dtype=torch.float32)
             torch.testing.assert_close(torch.compile(f)(x), g(x))
 
+    def test_custom_pass_inference_flag(self):
+        class CustomPass(CustomInferenceAwareGraphPass):
+            def __init__(self):
+                super().__init__()
+                self.is_infer = False
+
+            def __call__(self, g: torch.fx.graph.Graph, is_inference: bool):
+                self.is_infer = is_inference
+
+            def uuid(self) -> bytes:
+                return get_hash_for_files((__file__,))
+
+        custom_pass = CustomPass()
+        with config.patch(post_grad_custom_post_pass=custom_pass):
+
+            def f(x):
+                return x.abs()
+
+            x = torch.randn(8, dtype=torch.float32)
+            torch.testing.assert_close(torch.compile(f)(x), f(x))
+            self.assertTrue(custom_pass.is_infer)
+            x.requires_grad_()
+            torch.testing.assert_close(torch.compile(f)(x), f(x))
+            self.assertFalse(custom_pass.is_infer)
+
 
 if __name__ == "__main__":
     if IS_LINUX and HAS_CPU and torch.backends.mkldnn.is_available():
diff --git a/test/inductor/test_cutedsl_grouped_mm.py b/test/inductor/test_cutedsl_grouped_mm.py
index bd7221adc4065..8b2c87c71d524 100644
--- a/test/inductor/test_cutedsl_grouped_mm.py
+++ b/test/inductor/test_cutedsl_grouped_mm.py
@@ -116,15 +116,18 @@ def test_grouped_gemm_assorted_layouts(
         elif layout_A == "view":
             A_storage = torch.randn(sum_M * K, device=device, dtype=dtype)
             A = A_storage.view(sum_M, K)
-            assert A._base is not None
-            assert A.shape == (sum_M, K)
+            if A._base is None:
+                raise AssertionError
+            if A.shape != (sum_M, K):
+                raise AssertionError
 
         B = torch.randn((G, K, N), dtype=dtype, device=device) * 0.01
 
         if layout_B == "broadcasted":
             # Broadcast B across groups (zero stride along G)
             B = B[0].expand(G, K, N)
-            assert B.stride(0) == 0
+            if B.stride(0) != 0:
+                raise AssertionError
 
         def grouped_gemm_fn(A_packed, B_batched, offs):
             return F.grouped_mm(A_packed, B_batched, offs=offs)
diff --git a/test/inductor/test_cutedsl_template.py b/test/inductor/test_cutedsl_template.py
index 4f72762728216..eb8f800a89c72 100644
--- a/test/inductor/test_cutedsl_template.py
+++ b/test/inductor/test_cutedsl_template.py
@@ -265,12 +265,13 @@ def cutedsl_add_lowering(a: TensorBox, b: TensorBox) -> TensorBox:
                 return default_lowering(a, b)
 
             # Use autotuning to select the best variant
-            return autotune_select_algorithm(
+            node, _ = autotune_select_algorithm(
                 "cutedsl_add_autotune",
                 choices,
                 [a, b],
                 a.get_layout(),
             )
+            return node
 
         with patch.dict(lowerings, {torch.ops.aten.add.Tensor: cutedsl_add_lowering}):
             # Test function
@@ -418,11 +419,13 @@ def test_cutedsl_op_overrides(self):
         mock_cse_a.__str__.return_value = "tensor_a"
         mock_cse_a.dtype = torch.float32
         mock_cse_a.bounds = ValueRanges.unknown()
+        mock_cse_a.shape = (32, 64)
 
         mock_cse_b = MagicMock(spec=CSEVariable)
         mock_cse_b.__str__.return_value = "tensor_b"
         mock_cse_b.dtype = torch.float32
         mock_cse_b.bounds = ValueRanges.unknown()
+        mock_cse_b.shape = (32, 64)
 
         mock_graph = MockGraphHandler()
         with V.set_graph_handler(mock_graph):
@@ -447,14 +450,20 @@ def test_cutedsl_op_overrides(self):
                 result = CuteDSLOpOverrides.sqrt(mock_cse_a)
                 self.assertIsInstance(result, CSEVariable)
 
-                with self.assertRaises(NotImplementedError):
-                    result = CuteDSLOpOverrides.maximum(mock_cse_a, mock_cse_b)
-                    result = CuteDSLOpOverrides.minimum(mock_cse_a, mock_cse_b)
+                result = CuteDSLOpOverrides.maximum(mock_cse_a, mock_cse_b)
+                self.assertIsInstance(result, CSEVariable)
+
+                result = CuteDSLOpOverrides.minimum(mock_cse_a, mock_cse_b)
+                self.assertIsInstance(result, CSEVariable)
 
-        scalar_result = CuteDSLOpOverrides._ensure_tensor_ssa("5.0", mock_cse_a)
+        scalar_result = CuteDSLOpOverrides._ensure_tensor_ssa(
+            "5.0", mock_cse_a, is_tensor=False
+        )
         self.assertEqual(scalar_result, "cute.full_like(tensor_a, 5.0)")
 
-        tensor_result = CuteDSLOpOverrides._ensure_tensor_ssa(mock_cse_a, mock_cse_b)
+        tensor_result = CuteDSLOpOverrides._ensure_tensor_ssa(
+            mock_cse_a, mock_cse_b, is_tensor=True
+        )
         self.assertEqual(tensor_result, "tensor_a")
 
     def test_cse_integration(self):
diff --git a/test/inductor/test_cutlass_backend.py b/test/inductor/test_cutlass_backend.py
index e281b7824fc6a..a57d26035594f 100644
--- a/test/inductor/test_cutlass_backend.py
+++ b/test/inductor/test_cutlass_backend.py
@@ -1,4 +1,5 @@
 # Owner(s): ["module: inductor"]
+import gc
 import itertools
 import logging
 import math
@@ -11,7 +12,6 @@
 from collections.abc import Callable
 from enum import Enum
 from pathlib import Path
-from typing import Optional
 
 from torch._dynamo.exc import BackendCompilerFailed
 from torch._inductor.codegen.cutlass.serialization import (
@@ -19,6 +19,7 @@
 )
 from torch._inductor.utils import clear_caches
 from torch.export import Dim
+from torch.testing._internal.common_utils import random_matrix_with_scaled_reduction_dim
 from torch.testing._internal.logging_utils import log_settings
 from torch.utils import _pytree as pytree
 
@@ -34,6 +35,7 @@
 from torch._dynamo import config as dynamo_config
 from torch._dynamo.utils import counters
 from torch._inductor import config
+from torch._inductor.codecache import XPUCodeCache
 from torch._inductor.codegen.cutlass.kernel import CUTLASSTemplateCaller
 from torch._inductor.codegen.cutlass.utils import _gen_ops_cached, get_max_alignment
 from torch._inductor.exc import InductorError
@@ -45,23 +47,32 @@
 from torch.testing import FileCheck
 from torch.testing._internal.common_cuda import (
     PLATFORM_SUPPORTS_FP8,
+    SM100OrLater,
     SM80OrLater,
     SM90OrLater,
 )
+from torch.testing._internal.common_device_type import skipCUDAIf, skipXPUIf
 from torch.testing._internal.common_utils import (
     IN_RE_WORKER,
     instantiate_parametrized_tests,
     IS_FBCODE,
     parametrize,
 )
+from torch.testing._internal.common_xpu import Xe2_Or_Later
 from torch.testing._internal.inductor_utils import (
     _quantize_rowwise,
     _quantize_tensorwise,
+    GPU_TYPE,
     HAS_CPU,
     HAS_CUDA_AND_TRITON,
 )
 
 
+# We don't need triton in this test suite.
+HAS_XPU = torch.xpu.is_available()
+HAS_CUDA = torch.cuda.is_available()
+HAS_GPU = HAS_CUDA or HAS_XPU
+
 torch.set_float32_matmul_precision("high")
 if HAS_CUDA_AND_TRITON:
     torch.cuda.memory._set_allocator_settings("expandable_segments:False")
@@ -123,7 +134,7 @@ def _check_if_instances_equal(op1, op2) -> bool:
 
 def gen_args(op, shape, dtype=torch.float16):
     if op in bin_ops_under_test:
-        return (torch.rand(*shape, device="cuda:0", dtype=dtype),)
+        return (torch.rand(*shape, device=f"{GPU_TYPE}:0", dtype=dtype),)
     else:
         return ()
 
@@ -159,9 +170,12 @@ def select_no_algorithm(*args, **kwargs):
 
 @instantiate_parametrized_tests
 class TestCutlassBackend(TestCase):
+    # device_type of each test case is necessary for skipCUDAIf decorator.
+    device_type = GPU_TYPE
+
     def setUp(self):
-        if not HAS_CUDA_AND_TRITON:
-            self.skipTest("CUDA and triton are not available")
+        if not HAS_GPU:
+            self.skipTest(f"{GPU_TYPE} and triton are not available")
         if torch.version.hip:
             self.skipTest("CUTLASS backend is not supported on HIP")
 
@@ -185,13 +199,19 @@ def setUp(self):
     def tearDown(self):
         super().tearDown()
         clear_caches()
+        if GPU_TYPE == "xpu":
+            for dll in XPUCodeCache.dll_cache.values():
+                dll.close()
+            XPUCodeCache.dll_cache.clear()
+            gc.collect()
+            torch.xpu.empty_cache()
 
     def run_evt_test(self, model, op, shape, num_fusions=1):
         M, N = shape
-        a = torch.ones(M, N).cuda().half()
-        b = torch.ones(N, N).cuda().half().t()
+        a = torch.ones(M, N).to(GPU_TYPE).half()
+        b = torch.ones(N, N).to(GPU_TYPE).half().t()
         extra_args = gen_args(op, (M, N))
-        model = model.cuda()
+        model = model.to(GPU_TYPE)
 
         result = torch.compile(model)(a, b, extra_args)
         ref_result = model(a, b, extra_args)
@@ -215,7 +235,8 @@ def test_check_paths(self):
         self.assertTrue(os.path.exists(cutlass_mock_pydot_path))
         self.assertTrue(os.path.exists(cutlass_mock_scipy_path))
 
-    @unittest.skipIf(not SM90OrLater, "need sm_90")
+    @skipXPUIf(not Xe2_Or_Later, "")
+    @skipCUDAIf(not SM90OrLater, "need sm_90")
     @mock.patch.dict(os.environ, {"PATH": _get_path_without_sccache()})
     def test_max_autotune_cutlass_threshold(self):
         """
@@ -225,15 +246,16 @@ def test_max_autotune_cutlass_threshold(self):
         def mm(a, b):
             return a @ b
 
-        a = torch.randn(100, 10).cuda().half()
-        b = torch.randn(100, 10).cuda().half().t()
+        a = torch.randn(100, 10).to(GPU_TYPE).half()
+        b = torch.randn(100, 10).to(GPU_TYPE).half().t()
 
         with config.patch(
             {
                 "max_autotune": True,
                 "max_autotune_gemm_backends": "CUTLASS",
                 "compile_threads": 4,
-                "cutlass.cutlass_backend_min_gemm_size": 100000,
+                # Make it slightly too large to be accepted (m*n*k)
+                "cutlass.cutlass_backend_min_gemm_size": 100 * 100 * 10 + 1,
                 "cutlass.cutlass_max_profiling_configs": 2,
             }
         ):
@@ -248,6 +270,7 @@ def mm(a, b):
 
                 self.assertEqual(choices, [])
 
+    @skipXPUIf(not Xe2_Or_Later, "")
     @mock.patch.dict(os.environ, {"PATH": _get_path_without_sccache()})
     def test_import_cutlass(self):
         from torch._inductor.codegen.cutlass.utils import try_import_cutlass
@@ -257,6 +280,7 @@ def test_import_cutlass(self):
         import cutlass_cppgen  # type: ignore[import-not-found]  # noqa: F401
         import cutlass_library  # noqa: F401
 
+    @skipXPUIf(not Xe2_Or_Later, "")
     def test_cutlass_key(self):
         from torch._inductor.codegen.cutlass.utils import try_import_cutlass
 
@@ -265,7 +289,8 @@ def test_cutlass_key(self):
 
         self.assertIsNotNone(cutlass_key())
 
-    @unittest.skipIf(not SM90OrLater, "need sm_90")
+    @skipXPUIf(not Xe2_Or_Later, "")
+    @skipCUDAIf(not SM90OrLater, "need sm_90")
     @mock.patch.dict(os.environ, {"PATH": _get_path_without_sccache()})
     def test_cutlass_backend_subproc_mm(self):
         """
@@ -277,8 +302,8 @@ def test_cutlass_backend_subproc_mm(self):
 
         M, N, K = 4096, 2048, 25728
 
-        a = torch.randn(M, K).cuda().half()
-        b = torch.randn(N, K).cuda().half().t()
+        a = torch.randn(M, K).to(GPU_TYPE).half()
+        b = torch.randn(N, K).to(GPU_TYPE).half().t()
 
         with config.patch(
             {
@@ -291,9 +316,17 @@ def test_cutlass_backend_subproc_mm(self):
         ):
             Y_compiled = torch.compile(torch.mm)(a, b)
             Y = torch.mm(a, b)
-            torch.testing.assert_close(Y_compiled, Y)
+            if GPU_TYPE == "xpu":
+                atol = 1e-3  # default is 1e-5
+                rtol = 1e-3  # default is 1e-3
+            else:
+                atol = None
+                rtol = None
 
-    @unittest.skipIf(not SM90OrLater, "need sm_90")
+            torch.testing.assert_close(Y_compiled, Y, atol=atol, rtol=rtol)
+
+    @skipXPUIf(not Xe2_Or_Later, "")
+    @skipCUDAIf(not SM90OrLater, "need sm_90")
     @mock.patch.dict(os.environ, {"PATH": _get_path_without_sccache()})
     @parametrize("dtype", (torch.float16, torch.bfloat16))
     def test_cutlass_backend_subproc_addmm(self, dtype):
@@ -302,10 +335,16 @@ def test_cutlass_backend_subproc_addmm(self, dtype):
         """
 
         M, N, K = 4096, 2048, 25728
-        dtype = torch.float16
 
-        a = torch.randn(M, K, dtype=dtype).cuda()
-        b = torch.randn(N, K, dtype=dtype).cuda().t()
+        # Scale inputs by 1/sqrt(K) so that the matmul output has O(1)
+        # magnitude, avoiding large accumulation errors in half precision
+        # that would require loose tolerances.
+        a = random_matrix_with_scaled_reduction_dim(
+            M, K, dtype=dtype, device=GPU_TYPE, reduction_dim=-1
+        )
+        b = random_matrix_with_scaled_reduction_dim(
+            N, K, dtype=dtype, device=GPU_TYPE, reduction_dim=-1
+        ).t()
 
         x_shapes = [
             (M, N),
@@ -330,12 +369,20 @@ def test_cutlass_backend_subproc_addmm(self, dtype):
                 torch._dynamo.reset()
                 clear_caches()
 
-                x = torch.randn(x_shape).cuda().to(dtype)
+                x = torch.randn(x_shape).to(GPU_TYPE).to(dtype)
                 Y_compiled = torch.compile(torch.addmm)(x, a, b, alpha=alpha, beta=beta)
                 Y = torch.addmm(x, a, b, alpha=alpha, beta=beta)
-                torch.testing.assert_close(Y_compiled, Y)
+                if GPU_TYPE == "xpu":
+                    atol = 1e-3
+                    rtol = 1e-3
+                else:
+                    # use default
+                    atol = None
+                    rtol = None
+                torch.testing.assert_close(Y_compiled, Y, atol=atol, rtol=rtol)
 
-    @unittest.skipIf(not SM90OrLater, "need sm_90")
+    @skipXPUIf(not Xe2_Or_Later, "")
+    @skipCUDAIf(not SM90OrLater, "need sm_90")
     @mock.patch.dict(os.environ, {"PATH": _get_path_without_sccache()})
     def test_cutlass_backend_subproc_bmm(self):
         """
@@ -344,8 +391,8 @@ def test_cutlass_backend_subproc_bmm(self):
 
         B, M, N, K = 10, 4096, 2048, 25728
 
-        a = torch.randn(B, M, K).cuda().half()
-        b = torch.randn(B, N, K).cuda().half().permute(0, 2, 1)
+        a = torch.randn(B, M, K).to(GPU_TYPE).half()
+        b = torch.randn(B, N, K).to(GPU_TYPE).half().permute(0, 2, 1)
 
         with config.patch(
             {
@@ -360,7 +407,8 @@ def test_cutlass_backend_subproc_bmm(self):
             Y = torch.bmm(a, b)
             torch.testing.assert_close(Y_compiled, Y)
 
-    @unittest.skipIf(not SM90OrLater, "need sm_90")
+    @skipXPUIf(not Xe2_Or_Later, "")
+    @skipCUDAIf(not SM90OrLater, "need sm_90")
     @parametrize("dynamic", (False, True))
     @mock.patch.dict(os.environ, {"PATH": _get_path_without_sccache()})
     def test_diff_matmul_share_same_kernel(self, dynamic):
@@ -376,9 +424,9 @@ def forward(self, a, b, c):
                 return ab, ac
 
         model = MyModel()
-        a = torch.randn(128, 16).cuda().half()
-        b = torch.randn(128, 16).cuda().half().t()
-        c = torch.randn(512, 16).cuda().half().t()
+        a = torch.randn(128, 16).to(GPU_TYPE).half()
+        b = torch.randn(128, 16).to(GPU_TYPE).half().t()
+        c = torch.randn(512, 16).to(GPU_TYPE).half().t()
 
         with config.patch(
             {
@@ -403,7 +451,8 @@ def forward(self, a, b, c):
                 2,
             ).run(codes[0])
 
-    @unittest.skipIf(not SM90OrLater, "need sm_90")
+    @skipXPUIf(not Xe2_Or_Later, "")
+    @skipCUDAIf(not SM90OrLater, "need sm_90")
     @mock.patch.dict(os.environ, {"PATH": _get_path_without_sccache()})
     def test_number_mm_precompiles(self):
         torch._dynamo.utils.counters.clear()
@@ -418,9 +467,9 @@ def forward(self, a, b, c):
                 return ab
 
         model = MyModel()
-        a = torch.randn(128, 16).cuda().half()
-        b = torch.randn(128, 16).cuda().half().t()
-        c = torch.randn(512, 16).cuda().half().t()
+        a = torch.randn(128, 16).to(GPU_TYPE).half()
+        b = torch.randn(128, 16).to(GPU_TYPE).half().t()
+        c = torch.randn(512, 16).to(GPU_TYPE).half().t()
 
         with config.patch(
             {
@@ -454,7 +503,8 @@ def forward(self, a, b, c):
             )
 
     # NOTE: right now tuned_mm doesn't support cutlass 2x, which is used by A100
-    @unittest.skipIf(not SM90OrLater, "need sm_90")
+    @skipXPUIf(not Xe2_Or_Later, "")
+    @skipCUDAIf(not SM90OrLater, "need sm_90")
     @parametrize("dynamic", (False, True))
     @parametrize("use_aoti", (False, True))
     @parametrize("dtype", (torch.float16, torch.bfloat16))
@@ -483,10 +533,13 @@ class MyModel(torch.nn.Module):
             def forward(self, a, b):
                 return a @ b
 
-        model = MyModel().cuda()
+        model = MyModel().to(GPU_TYPE)
 
         inputs = [
-            (torch.randn(M, K).cuda().to(dtype), torch.randn(K, N).cuda().to(dtype))
+            (
+                torch.randn(M, K).to(GPU_TYPE).to(dtype),
+                torch.randn(K, N).to(GPU_TYPE).to(dtype),
+            )
             for (M, N, K) in shapes
         ]
 
@@ -520,7 +573,8 @@ def forward(self, a, b):
 
             torch.testing.assert_close(actual, expected)
 
-    @unittest.skipIf(not SM90OrLater, "need sm_90")
+    @skipXPUIf(True, "XPU SYCL-TLA has not supported fp8 yet")
+    @skipCUDAIf(not SM90OrLater, "need sm_90")
     @parametrize("dynamic", (False, True))
     @parametrize("use_aoti", (False, True))
     @parametrize("dtype", (torch.float8_e4m3fn,))
@@ -549,7 +603,7 @@ def test_max_autotune_cutlass_backend_fp8_scaled_mm(
         for shape in shapes:
             M, N, K = shape
             output_dtype = torch.bfloat16
-            device = "cuda"
+            device = GPU_TYPE
 
             x = torch.randn(M, K, dtype=output_dtype, device=device)
             w = torch.randn(N, K, dtype=output_dtype, device=device)
@@ -587,7 +641,7 @@ def forward(self, x_fp8, x_inverse_scale, w_t_fp8, w_inverse_scale):
             if dynamic
             else None
         )
-        model = MyModel().cuda()
+        model = MyModel().to(GPU_TYPE)
 
         with (
             config.patch(
@@ -631,10 +685,12 @@ def test_cutlass_backend_fp8_scaled_mm_mixed_dtypes(self):
         m, k, n = 256, 256, 256
 
         # Create mixed FP8 dtypes: e4m3fn x e5m2
-        a8 = torch.randn(m, k, device="cuda", dtype=torch.float16).to(
+        a8 = torch.randn(m, k, device=GPU_TYPE, dtype=torch.float16).to(
             torch.float8_e4m3fn
         )
-        b8 = torch.randn(k, n, device="cuda", dtype=torch.float16).to(torch.float8_e5m2)
+        b8 = torch.randn(k, n, device=GPU_TYPE, dtype=torch.float16).to(
+            torch.float8_e5m2
+        )
         # _scaled_mm requires mat2 to be column-major
         b8 = b8.t().contiguous().t()
 
@@ -651,7 +707,8 @@ def scaled_mm_fn(a, b):
         expected = scaled_mm_fn(a8, b8)
         torch.testing.assert_close(actual, expected, rtol=1e-2, atol=0.05)
 
-    @unittest.skipIf(not SM90OrLater, "need sm_90")
+    @skipXPUIf(not Xe2_Or_Later, "")
+    @skipCUDAIf(not SM90OrLater, "need sm_90")
     @parametrize("dynamic", (False, True))
     @parametrize("use_aoti", (False, True))
     @parametrize("dtype", (torch.float16, torch.bfloat16))
@@ -671,7 +728,7 @@ class MyModel(torch.nn.Module):
             def forward(self, x, a, b):
                 return torch.addmm(x, a, b)
 
-        model = MyModel().cuda()
+        model = MyModel().to(GPU_TYPE)
         # M, N, K
         shapes = [
             (128, 128, 16),
@@ -691,9 +748,9 @@ def forward(self, x, a, b):
 
             inputs = [
                 (
-                    torch.randn(x_shape(M, N)).cuda().to(dtype),
-                    torch.randn(M, K).cuda().to(dtype),
-                    torch.randn(N, K).cuda().to(dtype).t(),
+                    torch.randn(x_shape(M, N)).to(GPU_TYPE).to(dtype),
+                    torch.randn(M, K).to(GPU_TYPE).to(dtype),
+                    torch.randn(N, K).to(GPU_TYPE).to(dtype).t(),
                 )
                 for (M, N, K) in shapes
             ]
@@ -732,6 +789,68 @@ def forward(self, x, a, b):
                 torch.testing.assert_close(actual, expected)
 
     @unittest.skipIf(not SM90OrLater, "need sm_90")
+    @mock.patch.dict(os.environ, {"PATH": _get_path_without_sccache()})
+    def test_max_autotune_cutlass_backend_addmm_input_reorder(self):
+        """
+        Without input_reorder on the CUTLASS call in tuned_addmm, the CUTLASS
+        C kernel's argument order doesn't match the benchmark tensor order,
+        causing illegal memory access during autotuning.
+
+        NOTE: Uses ATEN+CUTLASS (not CUTLASS-only) because the existing addmm test
+        has a pre-existing CUTLASS accuracy issue at all shapes. Also verifies that neither ATen nor CUTLASS
+        choices get inf timings (which would indicate benchmark failures).
+        """
+        from torch._inductor.select_algorithm import (
+            AlgorithmSelectorCache,
+            ExternKernelCaller,
+        )
+
+        original_benchmark_choices = AlgorithmSelectorCache.benchmark_choices
+        timings_by_type: dict[str, list[float]] = {"aten": [], "cutlass": []}
+
+        @classmethod
+        def tracking_benchmark_choices(cls, choices, autotune_args, **kwargs):
+            result = original_benchmark_choices.__func__(
+                cls, choices, autotune_args, **kwargs
+            )
+            for choice, timing in result.items():
+                if isinstance(choice, ExternKernelCaller):
+                    timings_by_type["aten"].append(timing)
+                else:
+                    timings_by_type["cutlass"].append(timing)
+            return result
+
+        AlgorithmSelectorCache.benchmark_choices = tracking_benchmark_choices
+        try:
+            M, K, N = 256, 3520, 2048
+            bias = torch.randn(N, device=GPU_TYPE, dtype=torch.bfloat16)
+            x = torch.randn(M, K, device=GPU_TYPE, dtype=torch.bfloat16)
+            w = torch.randn(K, N, device=GPU_TYPE, dtype=torch.bfloat16)
+
+            with config.patch(
+                {
+                    "max_autotune": True,
+                    "max_autotune_gemm_backends": "ATEN,CUTLASS",
+                    "cutlass.cutlass_max_profiling_configs": 2,
+                }
+            ):
+                expected = torch.addmm(bias, x, w)
+                actual = torch.compile(torch.addmm)(bias, x, w)
+                torch.testing.assert_close(actual, expected)
+
+            self.assertTrue(
+                all(t != float("inf") for t in timings_by_type["aten"]),
+                f"ATen benchmark failed: {timings_by_type['aten']}",
+            )
+            self.assertTrue(
+                all(t != float("inf") for t in timings_by_type["cutlass"]),
+                f"CUTLASS benchmark failed: {timings_by_type['cutlass']}",
+            )
+        finally:
+            AlgorithmSelectorCache.benchmark_choices = original_benchmark_choices
+
+    @skipXPUIf(not Xe2_Or_Later, "")
+    @skipCUDAIf(not SM90OrLater, "need sm_90")
     @parametrize("dynamic", (False, True))
     @parametrize("use_aoti", (False, True))
     @parametrize("dtype", (torch.float16, torch.bfloat16))
@@ -753,7 +872,7 @@ class MyModel(torch.nn.Module):
             def forward(self, a, b):
                 return torch.bmm(a, b)
 
-        model = MyModel().cuda()
+        model = MyModel().to(GPU_TYPE)
         # B, M, N, K
         shapes = [
             (10, 4096, 2048, 25728),
@@ -765,12 +884,18 @@ def forward(self, a, b):
         for B, M, N, K in shapes:
             if use_expand:
                 # Create A using unsqueeze and expand
-                A = torch.randn(M, K).cuda().to(dtype).unsqueeze(0).expand(B, -1, -1)
+                A = (
+                    torch.randn(M, K)
+                    .to(GPU_TYPE)
+                    .to(dtype)
+                    .unsqueeze(0)
+                    .expand(B, -1, -1)
+                )
             else:
                 # Original method
-                A = torch.randn(B, M, K).cuda().to(dtype)
+                A = torch.randn(B, M, K).to(GPU_TYPE).to(dtype)
 
-            B_tensor = torch.randn(B, N, K).cuda().to(dtype).permute(0, 2, 1)
+            B_tensor = torch.randn(B, N, K).to(GPU_TYPE).to(dtype).permute(0, 2, 1)
             inputs.append((A, B_tensor))
         dynamic_shapes = (
             {
@@ -797,7 +922,8 @@ def forward(self, a, b):
                 actual = [compiled_model(*input) for input in inputs]
             torch.testing.assert_close(actual, expected)
 
-    @unittest.skipIf(not SM90OrLater, "need sm_90")
+    @skipXPUIf(True, "streamk kernels not supported on xpu cutlass backend yet")
+    @skipCUDAIf(not SM90OrLater, "need sm_90")
     @mock.patch.dict(os.environ, {"PATH": _get_path_without_sccache()})
     def test_max_autotune_cutlass_backend_regular_mm_streamk(
         self, dynamic: bool = False, max_autotune_gemm_backends: str = "CUTLASS"
@@ -831,15 +957,16 @@ def test_max_autotune_cutlass_backend_regular_mm_streamk(
                     16384,
                 ),
             ):
-                a = torch.randn(M, K).cuda().half()
-                b = torch.randn(N, K).cuda().half().t()
+                a = torch.randn(M, K).to(GPU_TYPE).half()
+                b = torch.randn(N, K).to(GPU_TYPE).half().t()
                 Y_compiled = compiled_model(a, b)
                 Y = torch.mm(a, b)
                 # we need relaxed numerical limits due to the sheer size of the
                 # matmuls involved. Many small addition differences add up.
                 torch.testing.assert_close(Y_compiled, Y, atol=0.01, rtol=0.01)
 
-    @unittest.skipIf(not SM90OrLater, "need sm_90")
+    @skipXPUIf(True, "streamk kernels not supported on xpu cutlass backend yet")
+    @skipCUDAIf(not SM90OrLater, "need sm_90")
     def test_streamk_with_dynamic(
         self,
     ):
@@ -850,8 +977,8 @@ def test_streamk_with_dynamic(
         shape. Without a correct workspace, the kernel will fail at runtime.
         """
 
-        a = torch.randn(128, 16).cuda().half()
-        b = torch.randn(128, 16).cuda().half().t()
+        a = torch.randn(128, 16).to(GPU_TYPE).half()
+        b = torch.randn(128, 16).to(GPU_TYPE).half().t()
 
         with config.patch(
             {
@@ -863,7 +990,8 @@ def test_streamk_with_dynamic(
             with self.assertRaisesRegex(InductorError, r".*NoValidChoicesError.*"):
                 _ = torch.compile(torch.mm, dynamic=True)(a, b)
 
-    @unittest.skipIf(not SM90OrLater, "need sm_90")
+    @skipXPUIf(True, "streamk kernels not supported on xpu cutlass backend yet")
+    @skipCUDAIf(not SM90OrLater, "need sm_90")
     def test_streamk_with_static(
         self,
     ):
@@ -880,8 +1008,8 @@ def test_streamk_with_static(
 
         for shape in shapes:
             M, N, K = shape
-            a = torch.randn(M, K).cuda().half()
-            b = torch.randn(N, K).cuda().half().t()
+            a = torch.randn(M, K).to(GPU_TYPE).half()
+            b = torch.randn(N, K).to(GPU_TYPE).half().t()
 
             with config.patch(
                 {
@@ -899,8 +1027,8 @@ def _test_max_autotune_cutlass_backend_epilogue_fusion(
         max_autotune_gemm_backends: str = "CUTLASS",
         fp16=True,
         expected_fuse_count=0,
-        mm: Optional[Callable[[torch.Tensor, torch.Tensor], torch.Tensor]] = None,
-        batch_size: Optional[int] = None,
+        mm: Callable[[torch.Tensor, torch.Tensor], torch.Tensor] | None = None,
+        batch_size: int | None = None,
     ):
         # Note: The ops that are available
         # also depend on the alignment of the shapes
@@ -908,11 +1036,11 @@ def _test_max_autotune_cutlass_backend_epilogue_fusion(
         # it can happen that no Cutlass 3.x op is available
         # that allows fusions
         if batch_size is None:
-            a = torch.randn(256, 32).cuda()
-            b = torch.randn(256, 32).cuda().t()
+            a = torch.randn(256, 32).to(GPU_TYPE)
+            b = torch.randn(256, 32).to(GPU_TYPE).t()
         else:
-            a = torch.randn(batch_size, 256, 32).cuda()
-            b = torch.randn(batch_size, 256, 32).cuda().permute(0, 2, 1)
+            a = torch.randn(batch_size, 256, 32).to(GPU_TYPE)
+            b = torch.randn(batch_size, 256, 32).to(GPU_TYPE).permute(0, 2, 1)
         if fp16:
             a = a.half()
             b = b.half()
@@ -923,20 +1051,22 @@ def _test_max_autotune_cutlass_backend_epilogue_fusion(
                 "autotune_in_subproc": True,
                 "max_autotune_gemm_backends": max_autotune_gemm_backends,
                 "cutlass.cutlass_max_profiling_configs": 4,
-                "cuda.version": "12.2",  # required to enable the Kernels we need
             }
         ):
             counters["inductor"]["cutlass_epilogue_fusion_counter"] = 0
-            assert mm is not None
+            if mm is None:
+                raise AssertionError("mm is None")
             Y_compiled = torch.compile(mm, dynamic=dynamic)(a, b)
             Y = mm(a, b)
             actual_count = counters["inductor"]["cutlass_epilogue_fusion_counter"]
-            assert actual_count == expected_fuse_count, (
-                f"Expected fuse count of {expected_fuse_count} but got {actual_count}"
-            )
+            if actual_count != expected_fuse_count:
+                raise AssertionError(
+                    f"Expected fuse count of {expected_fuse_count} but got {actual_count}"
+                )
             torch.testing.assert_close(Y_compiled, Y, atol=1e-2, rtol=1e-2)
 
-    @unittest.skipIf(not SM90OrLater, "need sm_90")
+    @skipXPUIf(not Xe2_Or_Later, "")
+    @skipCUDAIf(not SM90OrLater, "need sm_90")
     def test_max_autotune_cutlass_backend_simple_fusion_fp16_fp32acc(self):
         def mm(a, b):
             return (a @ b) * 3.0
@@ -945,7 +1075,8 @@ def mm(a, b):
             fp16=True, expected_fuse_count=0, mm=mm
         )
 
-    @unittest.skipIf(not SM90OrLater, "need sm_90")
+    @skipXPUIf(not Xe2_Or_Later, "")
+    @skipCUDAIf(not SM90OrLater, "need sm_90")
     def test_max_autotune_cutlass_backend_chained_fusion_fp16_fp32acc(self):
         def mm(a, b):
             return (a @ b) * 3.3 - 1.234
@@ -954,7 +1085,8 @@ def mm(a, b):
             fp16=True, expected_fuse_count=0, mm=mm
         )
 
-    @unittest.skipIf(not SM90OrLater, "need sm_90")
+    @skipXPUIf(not Xe2_Or_Later, "")
+    @skipCUDAIf(not SM90OrLater, "need sm_90")
     def test_max_autotune_cutlass_backend_relu_fusion_fp16_fp32acc(self):
         def mm(a, b):
             return torch.nn.functional.relu((a @ b) * 3.3 - 1.234)
@@ -964,7 +1096,8 @@ def mm(a, b):
             fp16=True, expected_fuse_count=0, mm=mm
         )
 
-    @unittest.skipIf(not SM90OrLater, "need sm_90")
+    @skipXPUIf(not Xe2_Or_Later, "")
+    @skipCUDAIf(not SM90OrLater, "need sm_90")
     def test_max_autotune_cutlass_backend_relu6_fusion_fp16_fp32acc(self):
         def mm(a, b):
             return torch.clamp(torch.nn.functional.relu(a @ b), max=6.0)
@@ -974,7 +1107,8 @@ def mm(a, b):
             fp16=True, expected_fuse_count=0, mm=mm
         )
 
-    @unittest.skipIf(not SM90OrLater, "need sm_90")
+    @skipXPUIf(not Xe2_Or_Later, "")
+    @skipCUDAIf(not SM90OrLater, "need sm_90")
     def test_max_autotune_cutlass_backend_no_fusion_dtype_mismatch(self):
         def mm(a, b):
             # this should not be fused, since the output dtype is different from the matmul dtype
@@ -984,7 +1118,8 @@ def mm(a, b):
             fp16=True, expected_fuse_count=0, mm=mm
         )
 
-    @unittest.skipIf(not SM90OrLater, "need sm_90")
+    @skipXPUIf(not Xe2_Or_Later, "")
+    @skipCUDAIf(not SM90OrLater, "need sm_90")
     def test_max_autotune_cutlass_backend_shape_dependent_normalization_fusion(self):
         def mm(a, b):
             return (a @ b) / b.size(1)
@@ -993,8 +1128,9 @@ def mm(a, b):
             fp16=True, expected_fuse_count=0, mm=mm
         )
 
+    @skipXPUIf(True, "int_mm not supported on xpu cutlass backend")
     # TODO: Enable dynamic test cases when dynamic support is added.
-    @unittest.skipIf(not SM90OrLater, "need sm_90")
+    @skipCUDAIf(not SM90OrLater, "need sm_90")
     @parametrize("dynamic", (False,))
     @mock.patch.dict(os.environ, {"PATH": _get_path_without_sccache()})
     def test_max_autotune_cutlass_backend_int_mm(
@@ -1013,8 +1149,8 @@ def mm(a, b):
         # this combination, so it's excluded from the test).  Also,
         # for CUTLASS alignment requirements, number of columns in
         # both tensors has to be divisible by 16.
-        a = torch.randint(0, 5, (100, 16), dtype=torch.int8).cuda()
-        b = torch.randint(0, 5, (32, 16), dtype=torch.int8).cuda().T
+        a = torch.randint(0, 5, (100, 16), dtype=torch.int8).to(GPU_TYPE)
+        b = torch.randint(0, 5, (32, 16), dtype=torch.int8).to(GPU_TYPE).T
 
         with config.patch(
             {
@@ -1029,7 +1165,8 @@ def mm(a, b):
             torch.testing.assert_close(Y_compiled, Y)
 
     @mock.patch.dict(os.environ, {"PATH": _get_path_without_sccache()})
-    @unittest.skipIf(not SM90OrLater, "need sm_90")
+    @skipXPUIf(not Xe2_Or_Later, "")
+    @skipCUDAIf(not SM90OrLater, "need sm_90")
     def test_force_cutlass_backend_aoti_dynamic(self):
         class MyModel(torch.nn.Module):
             def forward(self, x, w):
@@ -1050,8 +1187,8 @@ def forward(self, x, w):
                 "w": {0: K, 1: N},
             }
 
-            x = torch.randn(M, K).cuda().half()
-            w = torch.randn(N, K).cuda().half().t()
+            x = torch.randn(M, K).to(GPU_TYPE).half()
+            w = torch.randn(N, K).to(GPU_TYPE).half().t()
 
             actual = AOTIRunnerUtil.run(
                 model,
@@ -1062,7 +1199,8 @@ def forward(self, x, w):
             torch.testing.assert_close(expected, actual)
 
     @mock.patch.dict(os.environ, {"PATH": _get_path_without_sccache()})
-    @unittest.skipIf(not SM90OrLater, "need sm_90")
+    @skipXPUIf(not Xe2_Or_Later, "")
+    @skipCUDAIf(not SM90OrLater, "need sm_90")
     def test_force_cutlass_backend_aoti_cexpr_codegen(self):
         class MyModel(torch.nn.Module):
             def forward(self, x, w):
@@ -1088,8 +1226,8 @@ def forward(self, x, w):
                 "w": None,
             }
 
-            x = torch.randn(M, K).cuda().half()
-            w = torch.randn(N, K).cuda().half().t()
+            x = torch.randn(M, K).to(GPU_TYPE).half()
+            w = torch.randn(N, K).to(GPU_TYPE).half().t()
 
             actual = AOTIRunnerUtil.run(
                 model,
@@ -1100,7 +1238,8 @@ def forward(self, x, w):
             torch.testing.assert_close(expected, actual)
 
     @mock.patch.dict(os.environ, {"PATH": _get_path_without_sccache()})
-    @unittest.skipIf(not SM90OrLater, "need sm_90")
+    @skipXPUIf(True, "streamk kernels not supported on xpu cutlass backend yet")
+    @skipCUDAIf(not SM90OrLater, "need sm_90")
     def test_aoti_workspace_ptr(self):
         class MyModel(torch.nn.Module):
             def forward(self, x, w):
@@ -1111,15 +1250,15 @@ def forward(self, x, w):
                 "max_autotune": True,
                 "autotune_in_subproc": False,
                 "max_autotune_gemm_backends": "CUTLASS",
-                "cutlass.cutlass_op_allowlist_regex": "128x256x64.*stream_k_warpspecialized_cooperative_epi_nosmem",
+                "cutlass.cutlass_op_allowlist_regex": "128x256x64.*stream_k",
                 "cutlass.cutlass_max_profiling_configs": 1,
             }
         ):
             model = MyModel()
             M, N, K = 200, 5216, 10_432
 
-            x = torch.randn(M, K).cuda().half()
-            w = torch.randn(N, K).cuda().half().t()
+            x = torch.randn(M, K).to(GPU_TYPE).half()
+            w = torch.randn(N, K).to(GPU_TYPE).half().t()
 
             actual = AOTIRunnerUtil.run(
                 model,
@@ -1145,10 +1284,10 @@ def mm(a, b):
             return torch.mm(a, b)
 
         m, n, k = 32, 8, 64
-        mask = torch.tensor([0, 0, 1, 1]).tile(m, k // 4).cuda().half()
-        a = torch.rand(m, k).cuda().half() * mask
+        mask = torch.tensor([0, 0, 1, 1]).tile(m, k // 4).to(GPU_TYPE).half()
+        a = torch.rand(m, k).to(GPU_TYPE).half() * mask
         a_sparse = to_sparse_semi_structured(a)
-        b = torch.rand(k, n).cuda().half()
+        b = torch.rand(k, n).to(GPU_TYPE).half()
 
         with config.patch(
             {
@@ -1166,7 +1305,8 @@ def mm(a, b):
         cache = torch._inductor.codecache.LocalCache().lookup(
             "sparse_semi_structured_mm"
         )
-        assert cache is not None
+        if cache is None:
+            raise AssertionError("cache is None")
         high = cache[
             f"[('cuda', 'torch.float16', {m}, {k // 2}, {k // 2}, 1, 0), "
             f"('cuda', 'torch.int16', {m}, {k // 16}, {k // 16}, 1, 0), "
@@ -1176,9 +1316,13 @@ def mm(a, b):
         for kernel, duration in high.items():
             if kernel.startswith("cutlass_gemm") and not math.isinf(duration):
                 cutlass_kernels_count += 1
-        assert cutlass_kernels_count > 0
+        if cutlass_kernels_count <= 0:
+            raise AssertionError(
+                f"Expected cutlass_kernels_count > 0, got {cutlass_kernels_count}"
+            )
 
-    @unittest.skipIf(not SM90OrLater, "need sm_90")
+    @skipXPUIf(not Xe2_Or_Later, "")
+    @skipCUDAIf(not SM90OrLater, "need sm_90")
     @mock.patch.dict(os.environ, {"PATH": _get_path_without_sccache()})
     def test_cutlass_backend_op_denylist(
         self,
@@ -1186,9 +1330,9 @@ def test_cutlass_backend_op_denylist(
         def my_addmm(x, a, b, alpha, beta):
             return torch.addmm(x, a, b, alpha=beta, beta=alpha)
 
-        x = torch.randn((128, 128)).cuda().half()
-        a = torch.randn(128, 128).cuda().half()
-        b = torch.randn(128, 128).cuda().half().t()
+        x = torch.randn((128, 128)).to(GPU_TYPE).half()
+        a = torch.randn(128, 128).to(GPU_TYPE).half()
+        b = torch.randn(128, 128).to(GPU_TYPE).half().t()
 
         with fresh_cache():
             with config.patch(
@@ -1210,20 +1354,29 @@ def my_addmm(x, a, b, alpha, beta):
                         torch.compile(my_addmm, dynamic=False)(x, a, b, 1.0, 2.0)
                     args, _ = sa.call_args
                     op_name, choices, _, __ = args
-                    assert op_name == "addmm"
+                    if op_name != "addmm":
+                        raise AssertionError(
+                            f"Expected op_name 'addmm', got {op_name!r}"
+                        )
                     cuda_template_count = 0
                     for choice in choices:
                         if isinstance(choice, CUTLASSTemplateCaller):
                             choice_info = choice.info_dict()
                             op_conf_name = choice_info.get("op_conf_name", "")
-                            assert isinstance(op_conf_name, str)
-                            assert "pingpong" not in op_conf_name, (
-                                "All pingpong Kernels should have been filtered"
-                            )
+                            if not isinstance(op_conf_name, str):
+                                raise AssertionError(
+                                    f"Expected op_conf_name to be str, got {type(op_conf_name)}"
+                                )
+                            if "pingpong" in op_conf_name:
+                                raise AssertionError(
+                                    "All pingpong Kernels should have been filtered"
+                                )
                             cuda_template_count += 1
-                    assert cuda_template_count > 0, "No CUTLASSTemplateCaller choices"
+                    if cuda_template_count <= 0:
+                        raise AssertionError("No CUTLASSTemplateCaller choices")
 
-    @unittest.skipIf(not SM90OrLater, "need sm_90")
+    @skipXPUIf(True, "Intel cutlass doesn't have pingpong kernels yet")
+    @skipCUDAIf(not SM90OrLater, "need sm_90")
     @mock.patch.dict(os.environ, {"PATH": _get_path_without_sccache()})
     def test_cutlass_backend_op_allowlist(
         self,
@@ -1231,9 +1384,9 @@ def test_cutlass_backend_op_allowlist(
         def addmm(x, a, b, alpha, beta):
             return torch.addmm(x, a, b, alpha=alpha, beta=beta)
 
-        x = torch.randn((128, 128)).cuda().half()
-        a = torch.randn(128, 128).cuda().half()
-        b = torch.randn(128, 128).cuda().half().t()
+        x = torch.randn((128, 128)).to(GPU_TYPE).half()
+        a = torch.randn(128, 128).to(GPU_TYPE).half()
+        b = torch.randn(128, 128).to(GPU_TYPE).half().t()
 
         with fresh_cache():
             with config.patch(
@@ -1255,20 +1408,29 @@ def addmm(x, a, b, alpha, beta):
                         torch.compile(addmm, dynamic=False)(x, a, b, 1.0, 1.0)
                     args, _ = sa.call_args
                     op_name, choices, _, __ = args
-                    assert op_name == "addmm"
+                    if op_name != "addmm":
+                        raise AssertionError(
+                            f"Expected op_name 'addmm', got {op_name!r}"
+                        )
                     cuda_template_count = 0
                     for choice in choices:
                         if isinstance(choice, CUTLASSTemplateCaller):
                             choice_info = choice.info_dict()
                             op_conf_name = choice_info.get("op_conf_name", "")
-                            assert isinstance(op_conf_name, str)
-                            assert "pingpong" in op_conf_name, (
-                                "Only pingpong Kernels should have been allowed"
-                            )
+                            if not isinstance(op_conf_name, str):
+                                raise AssertionError(
+                                    f"Expected op_conf_name to be str, got {type(op_conf_name)}"
+                                )
+                            if "pingpong" not in op_conf_name:
+                                raise AssertionError(
+                                    "Only pingpong Kernels should have been allowed"
+                                )
                             cuda_template_count += 1
-                    assert cuda_template_count > 0, "No CUTLASSTemplateCaller choices"
+                    if cuda_template_count <= 0:
+                        raise AssertionError("No CUTLASSTemplateCaller choices")
 
-    @unittest.skipIf(not SM90OrLater, "need sm_90")
+    @skipXPUIf(True, "fp8 not supported on xpu cutlass backend yet")
+    @skipCUDAIf(not SM90OrLater, "need sm_90")
     @mock.patch.dict(os.environ, {"PATH": _get_path_without_sccache()})
     def test_cutlass_backend_fp8_scaled_mm_fast_accum_filtering(
         self,
@@ -1276,7 +1438,7 @@ def test_cutlass_backend_fp8_scaled_mm_fast_accum_filtering(
         float8_dtype = torch.float8_e4m3fn
         # Only bf16 output type is supported for row-wise scaling, not fp32
         output_dtype: torch.dtype = torch.bfloat16
-        device = "cuda"
+        device = GPU_TYPE
         M, K, N = 128, 128, 128  # Matmul Y = X [M, K] x W [N, K]
         x = torch.randn(M, K, dtype=output_dtype, device=device)
         w = torch.randn(N, K, dtype=output_dtype, device=device)
@@ -1337,24 +1499,30 @@ def run_test(use_fast_accum):
                             if isinstance(choice, CUTLASSTemplateCaller):
                                 choice_info = choice.info_dict()
                                 op_conf_name = choice_info.get("op_conf_name", "")
-                                assert isinstance(op_conf_name, str)
-                                if use_fast_accum:
-                                    assert "fastaccum" in op_conf_name, (
-                                        "Only fastaccum Kernels should have been allowed"
+                                if not isinstance(op_conf_name, str):
+                                    raise AssertionError(
+                                        f"Expected op_conf_name to be str, got {type(op_conf_name)}"
                                     )
+                                if use_fast_accum:
+                                    if "fastaccum" not in op_conf_name:
+                                        raise AssertionError(
+                                            "Only fastaccum Kernels should have been allowed"
+                                        )
                                 else:
-                                    assert "fastaccum" not in op_conf_name, (
-                                        "fastaccum Kernels should have been filtered"
-                                    )
+                                    if "fastaccum" in op_conf_name:
+                                        raise AssertionError(
+                                            "fastaccum Kernels should have been filtered"
+                                        )
                                 cuda_template_count += 1
-                        assert cuda_template_count > 0, (
-                            "No CUTLASSTemplateCaller choices"
-                        )
+                        if cuda_template_count <= 0:
+                            raise AssertionError("No CUTLASSTemplateCaller choices")
 
-        run_test(True)
+        if not SM100OrLater:
+            run_test(True)
         run_test(False)
 
-    @unittest.skipIf(not SM90OrLater, "need sm_90")
+    @skipXPUIf(not Xe2_Or_Later, "")
+    @skipCUDAIf(not SM90OrLater, "need sm_90")
     @mock.patch.dict(os.environ, {"PATH": _get_path_without_sccache()})
     def test_cutlass_backend_shape_coverage_mm(
         self,
@@ -1368,19 +1536,25 @@ def test_cutlass_backend_shape_coverage_mm(
         """
 
         inputs = [
-            (torch.randn(128, 500).cuda().half(), torch.randn(500, 576).cuda().half()),
             (
-                torch.randn(500, 128).cuda().half(),
-                torch.randn(128, 576).cuda().half(),
+                torch.randn(128, 500).to(GPU_TYPE).half(),
+                torch.randn(500, 576).to(GPU_TYPE).half(),
             ),
-            (torch.randn(128, 250).cuda().half(), torch.randn(250, 576).cuda().half()),
             (
-                torch.randn(250, 128).cuda().half(),
-                torch.randn(128, 576).cuda().half(),
+                torch.randn(500, 128).to(GPU_TYPE).half(),
+                torch.randn(128, 576).to(GPU_TYPE).half(),
             ),
             (
-                torch.randn(125, 128).cuda().half(),
-                torch.randn(128, 576).cuda().half(),
+                torch.randn(128, 250).to(GPU_TYPE).half(),
+                torch.randn(250, 576).to(GPU_TYPE).half(),
+            ),
+            (
+                torch.randn(250, 128).to(GPU_TYPE).half(),
+                torch.randn(128, 576).to(GPU_TYPE).half(),
+            ),
+            (
+                torch.randn(125, 128).to(GPU_TYPE).half(),
+                torch.randn(128, 576).to(GPU_TYPE).half(),
             ),
         ]
 
@@ -1412,13 +1586,17 @@ def test_cutlass_backend_shape_coverage_mm(
                 )
                 args, _ = sa.call_args
                 op_name, choices, _, __ = args
-                assert op_name == "mm"
+                if op_name != "mm":
+                    raise AssertionError(f"Expected op_name 'mm', got {op_name!r}")
                 cuda_template_count = 0
                 for choice in choices:
                     if isinstance(choice, CUTLASSTemplateCaller):
                         choice_info = choice.info_dict()
                         op_conf_name = choice_info.get("op_conf_name", "")
-                        assert isinstance(op_conf_name, str)
+                        if not isinstance(op_conf_name, str):
+                            raise AssertionError(
+                                f"Expected op_conf_name to be str, got {type(op_conf_name)}"
+                            )
                         cuda_template_count += 1
 
                 self.assertGreater(
@@ -1428,7 +1606,8 @@ def test_cutlass_backend_shape_coverage_mm(
                     f"M={M}, N={N}, K={K}",
                 )
 
-    @unittest.skipIf(not SM80OrLater, "need sm_80")
+    @skipXPUIf(not Xe2_Or_Later, "")
+    @skipCUDAIf(not SM80OrLater, "need sm_80")
     @mock.patch.dict(os.environ, {"PATH": _get_path_without_sccache()})
     def test_get_max_alignment(self):
         l4 = FixedLayout(
@@ -1489,13 +1668,14 @@ def test_get_max_alignment(self):
             m4, 4, "Wrong max alignment. Should have been 4 (due to float32 dtype )."
         )
 
-    @unittest.skipIf(not SM90OrLater, "need sm_90")
+    @skipXPUIf(True, "To be enabled.")
+    @skipCUDAIf(not SM90OrLater, "need sm_90")
     @mock.patch.dict(os.environ, {"PATH": _get_path_without_sccache()})
     def test_standalone_runner(self):
         max_autotune_gemm_backends = "CUTLASS"
 
-        a = torch.randn(128, 16).cuda().half()
-        b = torch.randn(128, 16).cuda().half().t()
+        a = torch.randn(128, 16).to(GPU_TYPE).half()
+        b = torch.randn(128, 16).to(GPU_TYPE).half().t()
 
         with config.patch(
             {
@@ -1524,10 +1704,12 @@ def test_standalone_runner(self):
 
                 sources = ctx.sources
 
-            assert len(sources) >= 1
+            if len(sources) < 1:
+                raise AssertionError(f"Expected len(sources) >= 1, got {len(sources)}")
 
             # Get names for temporary source and executable files.
-            cu_file = NamedTemporaryFile("w", suffix=".cu", delete=False)  # noqa: SIM115
+            suffix = ".cpp" if GPU_TYPE == "xpu" else ".cu"
+            cu_file = NamedTemporaryFile("w", suffix=suffix, delete=False)  # noqa: SIM115
             cu_file.close()
             exe_file = NamedTemporaryFile("w", suffix="", delete=False)  # noqa: SIM115
             exe_file.close()
@@ -1567,7 +1749,8 @@ def test_standalone_runner(self):
             os.remove(cu_file.name)
             os.remove(exe_file.name)
 
-    @unittest.skipIf(not SM90OrLater, "need sm_90")
+    @skipXPUIf(not Xe2_Or_Later, "")
+    @skipCUDAIf(not SM90OrLater, "need sm_90")
     @mock.patch.dict(os.environ, {"PATH": _get_path_without_sccache()})
     def test_cutlass_backend_integration(self):
         """
@@ -1577,8 +1760,8 @@ def test_cutlass_backend_integration(self):
         def mm(a, b):
             return a @ b
 
-        a = torch.randn(128, 16).cuda().half()
-        b = torch.randn(128, 16).cuda().half().t()
+        a = torch.randn(128, 16).to(GPU_TYPE).half()
+        b = torch.randn(128, 16).to(GPU_TYPE).half().t()
 
         with config.patch(
             {
@@ -1605,11 +1788,13 @@ def mm(a, b):
             match = re.search(
                 r"Got cutlass configs: total number of ops: (\d+)", output
             )
-            assert match, "Expect to find the cutlass configs log"
+            if not match:
+                raise AssertionError("Expect to find the cutlass configs log")
             num_ops = int(match.group(1))
             self.assertTrue(num_ops > 0, "The number of ops should be greater than 0")
 
-    @unittest.skipIf(not SM90OrLater, "need sm_90")
+    @skipXPUIf(not Xe2_Or_Later, "")
+    @skipCUDAIf(not SM90OrLater, "need sm_90")
     def test_maybe_append_choice_caching(self):
         """
         Test if maybe_append_choice's caching leads to correct results and
@@ -1624,9 +1809,9 @@ def forward(self, A, B):
                     A = A @ B / 32
                 return A
 
-        model = TestModule().cuda()
-        A = torch.randn(1024, 1024, dtype=torch.bfloat16, device="cuda")
-        B = torch.randn(1024, 1024, dtype=torch.bfloat16, device="cuda").t()
+        model = TestModule().to(GPU_TYPE)
+        A = torch.randn(1024, 1024, dtype=torch.bfloat16, device=GPU_TYPE)
+        B = torch.randn(1024, 1024, dtype=torch.bfloat16, device=GPU_TYPE).t()
 
         expected = model(A, B)
 
@@ -1657,11 +1842,10 @@ def counting_render(self, *args, **kwargs):
 
         torch.testing.assert_close(actual, expected)
 
-        # Check render call count: render is called uniquely for each codegen
-        # and for each finalized codegen.
         self.assertEqual(render_call_count, NUM_ITERATIONS + 2)
 
-    @unittest.skipIf(not SM90OrLater, "need sm_90")
+    @skipXPUIf(not Xe2_Or_Later, "")
+    @skipCUDAIf(not SM90OrLater, "need sm_90")
     @mock.patch.dict(os.environ, {"PATH": _get_path_without_sccache()})
     def test_multiple_mm(self):
         """
@@ -1676,13 +1860,13 @@ def forward(self, a, b, c, d):
                 mm2 = c @ d
                 return mm1, mm2
 
-        model = MultipleMMModel().cuda()
+        model = MultipleMMModel().to(GPU_TYPE)
 
         # Create tensors with different shapes
-        a = torch.randn(128, 64).cuda().half()
-        b = torch.randn(32, 64).cuda().half().t()
-        c = torch.randn(256, 128).cuda().half()
-        d = torch.randn(64, 128).cuda().half().t()
+        a = torch.randn(128, 64).to(GPU_TYPE).half()
+        b = torch.randn(32, 64).to(GPU_TYPE).half().t()
+        c = torch.randn(256, 128).to(GPU_TYPE).half()
+        d = torch.randn(64, 128).to(GPU_TYPE).half().t()
 
         # Track render calls
         from torch._inductor.codegen.cutlass.gemm_template import CUTLASSGemmTemplate
@@ -1717,9 +1901,11 @@ def counting_render(self, *args, **kwargs):
                 torch.testing.assert_close(actual, expected)
 
         num_matmuls = 2
-        self.assertEqual(render_call_count, num_matmuls + num_matmuls * 2)
+        expected_count = num_matmuls + num_matmuls * 2
+        self.assertEqual(render_call_count, expected_count)
 
-    @unittest.skipIf(not SM90OrLater, "need sm_90")
+    @skipXPUIf(not Xe2_Or_Later, "")
+    @skipCUDAIf(not SM90OrLater, "need sm_90")
     @mock.patch.dict(os.environ, {"PATH": _get_path_without_sccache()})
     def test_multiple_mm_with_dynamic_shape(self):
         """
@@ -1729,8 +1915,8 @@ def test_multiple_mm_with_dynamic_shape(self):
         class MultipleMMDynamicModel(torch.nn.Module):
             def __init__(self) -> None:
                 super().__init__()
-                self.c = torch.randn(64, 256).cuda().half()
-                self.d = torch.randn(128, 256).cuda().half().t()
+                self.c = torch.randn(64, 256).to(GPU_TYPE).half()
+                self.d = torch.randn(128, 256).to(GPU_TYPE).half().t()
 
             def forward(self, a, b):
                 # dynamic shape matmul
@@ -1739,11 +1925,11 @@ def forward(self, a, b):
                 mm2 = self.c @ self.d
                 return mm1, mm2
 
-        model = MultipleMMDynamicModel().cuda()
+        model = MultipleMMDynamicModel().to(GPU_TYPE)
 
         # Create tensors with different shapes
-        a = torch.randn(128, 64).cuda().half()
-        b = torch.randn(32, 64).cuda().half().t()
+        a = torch.randn(128, 64).to(GPU_TYPE).half()
+        b = torch.randn(32, 64).to(GPU_TYPE).half().t()
 
         # Track render calls
         from torch._inductor.codegen.cutlass.gemm_template import CUTLASSGemmTemplate
@@ -1778,15 +1964,17 @@ def counting_render(self, *args, **kwargs):
                 torch.testing.assert_close(actual, expected)
 
         num_matmuls = 2
-        self.assertEqual(render_call_count, num_matmuls + num_matmuls * 2)
+        expected_count = num_matmuls + num_matmuls * 2
+        self.assertEqual(render_call_count, expected_count)
 
-    @unittest.skipIf(not SM90OrLater, "need sm_90")
+    @skipXPUIf(not Xe2_Or_Later, "")
+    @skipCUDAIf(not SM90OrLater, "need sm_90")
     @mock.patch.dict(os.environ, {"PATH": _get_path_without_sccache()})
     def test_cutlass_backend_matmul_same_tensor(self):
         max_autotune_gemm_backends = "CUTLASS"
 
         M = 128
-        A = torch.randn(M, M).cuda().half()
+        A = torch.randn(M, M).to(GPU_TYPE).half()
 
         with config.patch(
             {
@@ -1799,13 +1987,14 @@ def test_cutlass_backend_matmul_same_tensor(self):
 
             torch.testing.assert_close(A @ A.t(), compiled(A, A.t()))
 
-    @unittest.skipIf(not SM90OrLater, "need sm_90")
+    @skipXPUIf(not Xe2_Or_Later, "")
+    @skipCUDAIf(not SM90OrLater, "need sm_90")
     @mock.patch.dict(os.environ, {"PATH": _get_path_without_sccache()})
     def test_cutlass_backend_matmul_nonzero_offset(self):
         max_autotune_gemm_backends = "CUTLASS"
 
         M = 129
-        A = torch.randn(M, M - 1).cuda().half()
+        A = torch.randn(M, M - 1).to(GPU_TYPE).half()
 
         with config.patch(
             {
@@ -1819,7 +2008,8 @@ def test_cutlass_backend_matmul_nonzero_offset(self):
                 A[1:, :] @ A[1:, :].t(), compiled(A[1:, :], A[1:, :].t())
             )
 
-    @unittest.skipIf(not SM90OrLater, "need sm_90")
+    @skipXPUIf(not Xe2_Or_Later, "")
+    @skipCUDAIf(not SM90OrLater, "need sm_90")
     @mock.patch.dict(os.environ, {"PATH": _get_path_without_sccache()})
     def test_flexible_layout(self):
         class TestModel(torch.nn.Module):
@@ -1828,8 +2018,8 @@ def forward(self, B):
                 return A @ B.t()
 
         M = 1024
-        B = torch.randn(M, M).cuda().half()
-        model = TestModel().cuda()
+        B = torch.randn(M, M).to(GPU_TYPE).half()
+        model = TestModel().to(GPU_TYPE)
 
         with config.patch(
             {
@@ -1840,7 +2030,8 @@ def forward(self, B):
         ):
             _ = torch.compile(model)(B)
 
-    @unittest.skipIf(not SM90OrLater, "need sm_90")
+    @skipXPUIf(True, "evt not supported on xpu cutlass backend yet")
+    @skipCUDAIf(not SM90OrLater, "need sm_90")
     @mock.patch.dict(os.environ, {"PATH": _get_path_without_sccache()})
     @use_evt_config
     def test_evt_flexible_layout(self):
@@ -1850,8 +2041,8 @@ def forward(self, B):
                 return (A @ B.t()).relu()
 
         M = 1024
-        B = torch.randn(M, M).cuda().half()
-        model = TestModel().cuda().half()
+        B = torch.randn(M, M).to(GPU_TYPE).half()
+        model = TestModel().to(GPU_TYPE).half()
 
         with config.patch(
             {
@@ -1867,7 +2058,8 @@ def forward(self, B):
             1,
         )
 
-    @unittest.skipIf(not SM90OrLater, "need sm_90")
+    @skipXPUIf(not Xe2_Or_Later, "")
+    @skipCUDAIf(not SM90OrLater, "need sm_90")
     @mock.patch.dict(os.environ, {"PATH": _get_path_without_sccache()})
     def test_filtered_ops_cache(self):
         class TestModel(torch.nn.Module):
@@ -1878,8 +2070,8 @@ def forward(self, B):
                 return A
 
         M = 1024
-        B = torch.randn(M, M).cuda().half()
-        model = TestModel().cuda()
+        B = torch.randn(M, M).to(GPU_TYPE).half()
+        model = TestModel().to(GPU_TYPE)
 
         start_time = time.time()
         with config.patch(
@@ -1890,21 +2082,27 @@ def forward(self, B):
             }
         ):
             _ = torch.compile(model)(B)
-        self.assertTrue(time.time() - start_time < 60)
 
-    @unittest.skipIf(not SM90OrLater, "need sm_90")
+        if GPU_TYPE == "xpu":
+            time_limit = 100
+        else:
+            time_limit = 60
+        self.assertTrue(time.time() - start_time < time_limit)
+
+    @skipXPUIf(not Xe2_Or_Later, "")
+    @skipCUDAIf(not SM90OrLater, "need sm_90")
     @mock.patch.dict(os.environ, {"PATH": _get_path_without_sccache()})
     @parametrize("use_aoti", (False, True))
     def test_compilation_time(self, use_aoti):
         M = 1024
-        A = torch.randn(M, M).cuda().half()
-        B = torch.randn(M, M).cuda().half().t()
+        A = torch.randn(M, M).to(GPU_TYPE).half()
+        B = torch.randn(M, M).to(GPU_TYPE).half().t()
 
         class MyModel(torch.nn.Module):
             def forward(self, a, b):
                 return a @ b
 
-        model = MyModel().cuda()
+        model = MyModel().to(GPU_TYPE)
         expected = model(A, B)
 
         start_time = time.time()
@@ -1923,10 +2121,19 @@ def forward(self, a, b):
             else:
                 actual = torch.compile(model, fullgraph=True)(A, B)
 
-            torch.testing.assert_close(actual, expected)
-        self.assertTrue(time.time() - start_time < 50)
-
-    @unittest.skipIf(not SM90OrLater, "need sm_90")
+            if GPU_TYPE == "xpu":
+                atol = 1e-4  # default is 1e-5
+                rtol = 1e-3  # default is 1e-3
+                expected_time = 100
+            else:
+                atol = None
+                rtol = None
+                expected_time = 50
+            torch.testing.assert_close(actual, expected, atol=atol, rtol=rtol)
+        self.assertTrue(time.time() - start_time < expected_time)
+
+    @skipXPUIf(True, "evt not supported on xpu cutlass backend yet")
+    @skipCUDAIf(not SM90OrLater, "need sm_90")
     @use_evt_config
     @evt_all_ops
     @evt_all_shapes
@@ -1938,7 +2145,8 @@ def forward(self, a, b, extra_args):
 
         self.run_evt_test(TestModel(), op, shape)
 
-    @unittest.skipIf(not SM90OrLater, "need sm_90")
+    @skipXPUIf(True, "evt not supported on xpu cutlass backend yet")
+    @skipCUDAIf(not SM90OrLater, "need sm_90")
     @use_evt_config
     @evt_bin_ops
     def test_evt_broadcasting(self, op):
@@ -1949,10 +2157,10 @@ def forward(self, a, b, extra_args):
 
         M = 1024
         N = 512
-        a = torch.ones(M, N).cuda().half()
-        b = torch.ones(N, N).cuda().half().t()
+        a = torch.ones(M, N).to(GPU_TYPE).half()
+        b = torch.ones(N, N).to(GPU_TYPE).half().t()
         extra_args = gen_args(op, (M, N))
-        model = TestModel().cuda()
+        model = TestModel().to(GPU_TYPE)
 
         result = torch.compile(model)(a, b, extra_args)
         ref_result = model(a, b, extra_args)
@@ -1963,7 +2171,8 @@ def forward(self, a, b, extra_args):
         )
         torch.testing.assert_close(result, ref_result)
 
-    @unittest.skipIf(not SM90OrLater, "need sm_90")
+    @skipXPUIf(True, "evt not supported on xpu cutlass backend yet")
+    @skipCUDAIf(not SM90OrLater, "need sm_90")
     @use_evt_config
     @evt_un_ops
     def test_evt_activations(self, op):
@@ -1974,10 +2183,10 @@ def forward(self, a, b, extra_args):
 
         M = 1024
         N = 512
-        a = torch.ones(M, N).cuda().half()
-        b = torch.ones(N, N).cuda().half().t()
+        a = torch.ones(M, N).to(GPU_TYPE).half()
+        b = torch.ones(N, N).to(GPU_TYPE).half().t()
         extra_args = gen_args(op, (M, N))
-        model = TestModel().cuda()
+        model = TestModel().to(GPU_TYPE)
 
         result = torch.compile(model)(a, b, extra_args)
         ref_result = model(a, b, extra_args)
@@ -1988,14 +2197,15 @@ def forward(self, a, b, extra_args):
         )
         torch.testing.assert_close(result, ref_result)
 
-    @unittest.skipIf(not SM90OrLater, "need sm_90")
+    @skipXPUIf(True, "evt not supported on xpu cutlass backend yet")
+    @skipCUDAIf(not SM90OrLater, "need sm_90")
     @use_evt_config
     @evt_all_ops
     def test_evt_mixed_dtypes(self, op):
         M = 1024
         N = 256
 
-        fp32_tensor = torch.ones(M, N).cuda().float()
+        fp32_tensor = torch.ones(M, N).to(GPU_TYPE).float()
 
         class TestModel(torch.nn.Module):
             def forward(self, a, b, extra_args):
@@ -2004,9 +2214,9 @@ def forward(self, a, b, extra_args):
                 out1 = torch.add(out0, fp32_tensor)
                 return out1
 
-        model = TestModel().cuda()
-        a = torch.ones(M, N).cuda().half()
-        b = torch.ones(N, N).cuda().half().t()
+        model = TestModel().to(GPU_TYPE)
+        a = torch.ones(M, N).to(GPU_TYPE).half()
+        b = torch.ones(N, N).to(GPU_TYPE).half().t()
         extra_args = gen_args(op, (M, N), dtype=torch.float16)
 
         # baseline is cutlass kernel + triton
@@ -2029,7 +2239,8 @@ def forward(self, a, b, extra_args):
 
         torch.testing.assert_close(result, ref_result)
 
-    @unittest.skipIf(not SM90OrLater, "need sm_90")
+    @skipXPUIf(True, "evt not supported on xpu cutlass backend yet")
+    @skipCUDAIf(not SM90OrLater, "need sm_90")
     @use_evt_config
     @evt_all_ops
     def test_evt_multi_op(self, op):
@@ -2040,7 +2251,8 @@ def forward(self, a, b, extra_args):
 
         self.run_evt_test(TestModel(), op, (1024, 512))
 
-    @unittest.skipIf(not SM90OrLater, "need sm_90")
+    @skipXPUIf(True, "evt not supported on xpu cutlass backend yet")
+    @skipCUDAIf(not SM90OrLater, "need sm_90")
     @use_evt_config
     @evt_all_ops
     def test_evt_reuse_matmul_input(self, op):
@@ -2051,7 +2263,8 @@ def forward(self, a, b, extra_args):
 
         self.run_evt_test(TestModel(), op, (1024, 1024))  # shape needs to be square
 
-    @unittest.skipIf(not SM90OrLater, "need sm_90")
+    @skipXPUIf(True, "evt not supported on xpu cutlass backend yet")
+    @skipCUDAIf(not SM90OrLater, "need sm_90")
     @use_evt_config
     @evt_all_ops
     @parametrize(
@@ -2071,10 +2284,10 @@ def forward(self, a, b, extra_args):
         shapes = [(512, 512)] if not dynamic else [(1024, 64), (128, 256)]
         for i, shape in enumerate(shapes):
             M, N = shape
-            a = torch.ones(M, N).cuda().half()
-            b = torch.ones(N, N).cuda().half().t()
+            a = torch.ones(M, N).to(GPU_TYPE).half()
+            b = torch.ones(N, N).to(GPU_TYPE).half().t()
             extra_args = gen_args(op, (M, N))
-            model = TestModel().cuda()
+            model = TestModel().to(GPU_TYPE)
 
             result = torch.compile(model)(a, b, extra_args)
             ref_result = model(a, b, extra_args)
@@ -2087,7 +2300,8 @@ def forward(self, a, b, extra_args):
             )
             torch.testing.assert_close(result, ref_result)
 
-    @unittest.skipIf(not SM90OrLater, "need sm_90")
+    @skipXPUIf(True, "evt not supported on xpu cutlass backend yet")
+    @skipCUDAIf(not SM90OrLater, "need sm_90")
     @use_evt_config
     def test_evt_return_accumulator(self):
         op = torch.add
@@ -2099,10 +2313,10 @@ def forward(self, a, b, extra_args):
 
         M = 1024
         N = 512
-        a = torch.ones(M, N).cuda().half()
-        b = torch.ones(N, N).cuda().half().t()
+        a = torch.ones(M, N).to(GPU_TYPE).half()
+        b = torch.ones(N, N).to(GPU_TYPE).half().t()
         extra_args = gen_args(op, (M, N))
-        model = TestModel().cuda()
+        model = TestModel().to(GPU_TYPE)
 
         result = torch.compile(model)(a, b, extra_args)
         ref_result = model(a, b, extra_args)
@@ -2113,19 +2327,18 @@ def forward(self, a, b, extra_args):
         )
         torch.testing.assert_close(result, ref_result)
 
-    @mock.patch.dict(os.environ, {"PATH": _get_path_without_sccache()})
-    @parametrize("arch", ("90", "100"))
-    @parametrize("cuda_version", ("12.4", "12.8"))
-    def test_gemm_operation_serialization(self, arch: str, cuda_version: str):
+    def _test_gemm_operation_serialization(
+        self, arch: str, cuda_version: str, min_ops=1000
+    ):
         """
         Testing serialization for GEMM operations generated by CUTLASS.
         This should cover GroupedGemmOperation as well.
         """
-        full_ops = _gen_ops_cached(arch, cuda_version)
+        full_ops = _gen_ops_cached(arch, cuda_version, GPU_TYPE)
         ops = pytree.tree_flatten(full_ops)[0]
 
         # sanity check
-        self.assertGreater(len(ops), 1000, "Too few ops generated")
+        self.assertGreater(len(ops), min_ops, "Too few ops generated")
 
         # test if configuration name is unique
         op_config_names = [op.configuration_name() for op in ops]
@@ -2141,11 +2354,30 @@ def test_gemm_operation_serialization(self, arch: str, cuda_version: str):
         for op, deserialized_op in zip(ops, deserialized_ops, strict=False):
             self.assertTrue(_check_if_instances_equal(op, deserialized_op))
 
+    @unittest.skipIf(not HAS_CUDA, "CUDA not available")
+    @mock.patch.dict(os.environ, {"PATH": _get_path_without_sccache()})
+    @parametrize("arch", ("90", "100"))
+    @parametrize("cuda_version", ("12.4", "12.8"))
+    def test_gemm_operation_serialization_cuda(self, arch: str, cuda_version: str):
+        self._test_gemm_operation_serialization(arch, cuda_version)
+
+    @skipXPUIf(not Xe2_Or_Later, "")
+    @unittest.skipIf(not HAS_XPU, "XPU not available")
+    @mock.patch.dict(os.environ, {"PATH": _get_path_without_sccache()})
+    @parametrize("arch", ("Xe12", "Xe20"))
+    @parametrize("xpu_version", ("20250201", "20250301"))
+    def test_gemm_operation_serialization_xpu(self, arch: str, xpu_version: str):
+        from torch._inductor.codegen.cutlass.utils import _normalize_xpu_arch
+
+        arch = _normalize_xpu_arch(arch)
+        self._test_gemm_operation_serialization(arch, xpu_version, min_ops=40)
+
     @unittest.skipIf(
         torch.cuda.is_available() and not PLATFORM_SUPPORTS_FP8,
         "FP8 is only supported on H100+",
     )
-    @unittest.skipIf(not SM90OrLater, "need sm_90")
+    @skipXPUIf(True, "XPU SYCL-TLA has not supported fp8 yet")
+    @skipCUDAIf(not SM90OrLater, "need sm_90")
     @fp8_config
     @parametrize("float8_dtype", (torch.float8_e4m3fn,))
     @parametrize(
@@ -2171,7 +2403,7 @@ def test_fp8_rowwise_scaling(
     ):
         # Only bf16 output type is supported for row-wise scaling, not fp32
         output_dtype: torch.dtype = torch.bfloat16
-        device = "cuda"
+        device = GPU_TYPE
         M, K, N = shape  # Matmul Y = X [M, K] x W [N, K]
         x = torch.randn(M, K, dtype=input_dtype, device=device)
         w = torch.randn(N, K, dtype=input_dtype, device=device)
@@ -2222,7 +2454,8 @@ def linear(x_fp8, x_inverse_scale, w_t_fp8, w_inverse_scale, bias):
         torch.cuda.is_available() and not PLATFORM_SUPPORTS_FP8,
         "FP8 is only supported on H100+",
     )
-    @unittest.skipIf(not SM90OrLater, "need sm_90")
+    @skipXPUIf(True, "XPU SYCL-TLA has not supported fp8 yet")
+    @skipCUDAIf(not SM90OrLater, "need sm_90")
     @fp8_config
     @parametrize("float8_dtype", (torch.float8_e4m3fn,))
     @parametrize(
@@ -2252,7 +2485,7 @@ def test_fp8_rowwise_scaling_multiple_linear(
             self.skipTest("Accuracy issues when both AOTI and dynamic are enabled")
         # Only bf16 output type is supported for row-wise scaling, not fp32
         output_dtype: torch.dtype = torch.bfloat16
-        device = "cuda"
+        device = GPU_TYPE
         M, N = shape  # Matmul Y = X [M, K] x W [N, K]
         x = torch.randn(M, N, dtype=output_dtype, device=device)
         w1 = torch.randn(N, N, dtype=output_dtype, device=device)
@@ -2291,7 +2524,7 @@ def forward(self, x):
                 )
                 return y2
 
-        model = TestModule(w1, w2, float8_dtype).cuda()
+        model = TestModule(w1, w2, float8_dtype).to(GPU_TYPE)
 
         dynamic_shapes = (
             {
@@ -2319,7 +2552,8 @@ def forward(self, x):
         torch.cuda.is_available() and not PLATFORM_SUPPORTS_FP8,
         "FP8 is only supported on H100+",
     )
-    @unittest.skipIf(not SM90OrLater, "need sm_90")
+    @skipXPUIf(True, "XPU SYCL-TLA has not supported fp8 yet")
+    @skipCUDAIf(not SM90OrLater, "need sm_90")
     @fp8_config
     @parametrize("float8_dtype", (torch.float8_e4m3fn,))
     @parametrize(
@@ -2343,7 +2577,7 @@ def test_fp8_tensorwise_scaling(
         use_fast_accum: bool,
         input_dtype: torch.dtype,
     ):
-        device = "cuda"
+        device = GPU_TYPE
         M, K, N = shape  # Matmul Y = X [M, K] x W [N, K]
         output_dtype = input_dtype
         # input and output dtypes of _scaled_mm do not need to be the same, but
@@ -2396,7 +2630,8 @@ def linear(x_fp8, x_inverse_scale, w_t_fp8, w_inverse_scale, bias):
         # setting a small absolute tolerance in these tests
         torch.testing.assert_close(y_eager, y_compiled, rtol=1e-2, atol=0.05)
 
-    @unittest.skipIf(not SM90OrLater, "need sm_90")
+    @skipXPUIf(not Xe2_Or_Later, "")
+    @skipCUDAIf(not SM90OrLater, "need sm_90")
     def test_config_number_post_filtering(self) -> None:
         """
         Test if cutlass backend produces the same number of configs after filtering
@@ -2409,8 +2644,8 @@ def test_config_number_post_filtering(self) -> None:
 
         for layout in layouts:
             for dtype in dtypes:
-                a = torch.randn(128, 128, dtype=dtype).cuda()
-                b = torch.randn(128, 128, dtype=dtype).cuda()
+                a = torch.randn(128, 128, dtype=dtype).to(GPU_TYPE)
+                b = torch.randn(128, 128, dtype=dtype).to(GPU_TYPE)
                 if layout[0] == "c":
                     a = a.t()
                 if layout[1] == "c":
@@ -2440,8 +2675,10 @@ def test_config_number_post_filtering(self) -> None:
 
         # Check that all config counts are equal
         all_counts = list(config_counts.values())
+        # XPU has more configs for bf16 than f16.
+        expected_count = 2 if GPU_TYPE == "xpu" else 1
         self.assertTrue(
-            len(set(all_counts)) == 1,
+            len(set(all_counts)) == expected_count,
             f"Config counts should be equal across all layout/dtype combinations. "
             f"Got counts: {config_counts}",
         )
@@ -2451,5 +2688,5 @@ def test_config_number_post_filtering(self) -> None:
     from torch._inductor.utils import is_big_gpu
 
     # Set env to make it work in CI.
-    if HAS_CUDA_AND_TRITON and HAS_CPU and is_big_gpu():
+    if HAS_GPU and HAS_CPU and is_big_gpu():
         run_tests()
diff --git a/test/inductor/test_cutlass_evt.py b/test/inductor/test_cutlass_evt.py
index dd296b7f75ac7..af160fa696215 100644
--- a/test/inductor/test_cutlass_evt.py
+++ b/test/inductor/test_cutlass_evt.py
@@ -12,7 +12,7 @@
 from torch._inductor.ir import ComputedBuffer, FixedLayout, PermuteView, Pointwise
 from torch._inductor.scheduler import BaseSchedulerNode
 from torch._inductor.utils import OrderedSet
-from torch.testing._internal.common_cuda import SM90OrLater
+from torch.testing._internal.common_cuda import IS_SM100, IS_SM90, SM90OrLater
 from torch.testing._internal.inductor_utils import (
     HAS_CPU,
     HAS_CUDA_AND_TRITON,
@@ -66,8 +66,16 @@
         ),
     }
 
+    # For SM100
+    class MockMathInstruction:
+        def __init__(self):
+            self.opcode_class = cutlass_lib.library.OpcodeClass.TensorOp
+
     class MockTileDescription:
         threadblock_shape = (128, 128, 8)
+        # SM100 path
+        cluster_shape = (1, 1, 1)
+        math_instruction = MockMathInstruction()
 
     def _create_mock_buffer_name_map(example_tensors):
         name_to_buffer = {}
@@ -441,7 +449,7 @@ def fn(accum, bias):
 """,
         )
 
-    @unittest.skipIf(not SM90OrLater, "need sm_90")
+    @unittest.skipIf(not SM90OrLater, "need sm_90 or later")
     @unittest.skipIf(not try_import_cutlass(), "requires cutlass")
     def test_evt_codegen(self):
         _, _, code, _ = trace(
@@ -454,9 +462,10 @@ def test_evt_codegen(self):
             _create_mock_buffer_name_map(EXAMPLE_TENSORS),
             lambda x: x,  # static shapes
         )
-        self.assertExpectedInline(
-            code,
-            """\
+        if IS_SM90:
+            self.assertExpectedInline(
+                code,
+                """\
 
 using EpilogueDescriptor = cutlass::epilogue::collective::detail::EpilogueDescriptor<
   cute::Shape<_128, _128, _8>, cutlass::epilogue::collective::EpilogueTileAuto,
@@ -554,7 +563,110 @@ def test_evt_codegen(self):
 using StrideD = cute::Stride<int64_t, cute::Int<1>, cute::Int<0>>;
 
 """,
-        )
+            )
+
+        if IS_SM100:
+            self.assertExpectedInline(
+                code,
+                """\
+
+using EpilogueDescriptor = cutlass::epilogue::collective::detail::Sm100EpilogueDescriptor<
+  cutlass::arch::OpClassTensorOp,
+  cute::Shape<_128, _128, _8>, cutlass::epilogue::collective::EpilogueTileAuto,
+  float, float, float,
+  cutlass::epilogue::collective::EpilogueScheduleAuto, cute::Stride<int64_t, cute::Int<1>, cute::Int<0>>, cute::Stride<int64_t, cute::Int<1>, cute::Int<0>>,
+  false /* IsPerColScaleSupported */,
+  false /* IsBlockScaleSupported */
+>;
+
+using ElementC = float;
+using StrideC = cute::Stride<int64_t, cute::Int<1>, cute::Int<0>>;
+using TensorC = cutlass::epilogue::fusion::Sm90SrcFetch<float>;
+
+using Accum = cutlass::epilogue::fusion::Sm90AccFetch;
+
+using AuxDescriptor = cutlass::epilogue::collective::detail::Sm100AuxLoadDescriptor<EpilogueDescriptor, cute::Stride<int64_t, cute::Int<1>, cute::Int<0>>, float>;
+
+using Aux = cutlass::epilogue::fusion::Sm90AuxLoad<
+    AuxDescriptor::Stages, typename AuxDescriptor::EpilogueTile, float,
+    cute::Stride<int64_t, cute::Int<1>, cute::Int<0>>, typename AuxDescriptor::SmemLayoutAtom, typename AuxDescriptor::CopyOpS2R
+>;
+
+using Bias = cutlass::epilogue::fusion::Sm90ColBroadcast<
+    0 /*Stages*/, typename EpilogueDescriptor::TileShape, float, float,
+    cute::Stride<cute::Int<1>, cute::Int<0>, cute::Int<0>>
+>;
+
+using Compute0 = cutlass::epilogue::fusion::Sm90Compute<
+    cutlass::plus, float, float,
+    cutlass::FloatRoundStyle::round_to_nearest
+>;
+
+using EVTCompute0 = cutlass::epilogue::fusion::Sm90EVT<
+    Compute0,
+    Accum,
+    TensorC>;
+
+using Compute1 = cutlass::epilogue::fusion::Sm90Compute<
+    cutlass::plus, float, float,
+    cutlass::FloatRoundStyle::round_to_nearest
+>;
+
+using EVTCompute1 = cutlass::epilogue::fusion::Sm90EVT<
+    Compute1,
+    EVTCompute0,
+    Aux>;
+
+using FDescriptor = cutlass::epilogue::collective::detail::Sm100AuxStoreDescriptor<
+    EpilogueDescriptor, cute::Stride<int64_t, cute::Int<1>, cute::Int<0>>, float
+>;
+
+using F = cutlass::epilogue::fusion::Sm90AuxStore<
+    FDescriptor::Stages, typename FDescriptor::EpilogueTile, float,
+    cutlass::FloatRoundStyle::round_to_nearest, cute::Stride<int64_t, cute::Int<1>, cute::Int<0>>, typename FDescriptor::SmemLayoutAtom,
+    typename FDescriptor::CopyOpR2S
+>;
+
+using EVTF = cutlass::epilogue::fusion::Sm90EVT<
+    F,
+    EVTCompute1>;
+
+using Compute2 = cutlass::epilogue::fusion::Sm90Compute<
+    cutlass::epilogue::thread::ReLu, float, float,
+    cutlass::FloatRoundStyle::round_to_nearest
+>;
+
+using Compute3 = cutlass::epilogue::fusion::Sm90Compute<
+    cutlass::plus, float, float,
+    cutlass::FloatRoundStyle::round_to_nearest
+>;
+
+using Compute4 = cutlass::epilogue::fusion::Sm90Compute<
+    cutlass::plus, float, float,
+    cutlass::FloatRoundStyle::round_to_nearest
+>;
+
+using DagCompute4 = cutlass::epilogue::fusion::Sm90TopologicalVisitor<
+    float,
+    cute::tuple<
+        cute::seq<>,
+        cute::seq<>,
+        cute::seq<0>,
+        cute::seq<2, 1>,
+        cute::seq<3, 0>,
+    >,
+    EVTF,
+    Bias,
+    Compute2,
+    Compute3,
+    Compute4
+>;
+
+using ElementD = float;
+using StrideD = cute::Stride<int64_t, cute::Int<1>, cute::Int<0>>;
+
+""",  # noqa: B950
+            )
 
 
 if __name__ == "__main__":
diff --git a/test/inductor/test_cutlass_fallback.py b/test/inductor/test_cutlass_fallback.py
new file mode 100644
index 0000000000000..8de87bb02e574
--- /dev/null
+++ b/test/inductor/test_cutlass_fallback.py
@@ -0,0 +1,418 @@
+# Owner(s): ["module: inductor"]
+"""Tests for CUTLASS fallback and subprocess isolation (#171094)."""
+
+import os
+import unittest
+import unittest.mock as mock
+from unittest.mock import patch
+
+import torch
+from torch._inductor import config
+from torch._inductor.select_algorithm import AlgorithmSelectorCache, ExternKernelCaller
+from torch._inductor.test_case import run_tests, TestCase
+from torch._inductor.utils import clear_caches
+from torch.testing._internal.common_cuda import SM90OrLater
+from torch.testing._internal.inductor_utils import HAS_CUDA_AND_TRITON
+
+
+# Check if CUTLASS is available
+try:
+    from torch._inductor.codegen.cutlass.utils import try_import_cutlass
+
+    HAS_CUTLASS = try_import_cutlass()
+except ImportError:
+    HAS_CUTLASS = False
+
+
+def _get_path_without_sccache() -> str:
+    """Get the PATH environment variable without sccache."""
+    path_envs = os.environ.get("PATH", "").split(":")
+    path_envs = [env for env in path_envs if "/opt/cache/bin" not in env]
+    return ":".join(path_envs)
+
+
+class TestCutlassFallback(TestCase):
+    """Tests for CUTLASS fallback behavior when benchmarks fail."""
+
+    def setUp(self):
+        if not HAS_CUDA_AND_TRITON:
+            self.skipTest("CUDA and triton are not available")
+        if torch.version.hip:
+            self.skipTest("CUTLASS backend is not supported on HIP")
+
+        old_disable_fresh_cache_envvar = os.environ.get(
+            "INDUCTOR_TEST_DISABLE_FRESH_CACHE", ""
+        )
+        try:
+            os.environ["INDUCTOR_TEST_DISABLE_FRESH_CACHE"] = "1"
+            super().setUp()
+        finally:
+            os.environ["INDUCTOR_TEST_DISABLE_FRESH_CACHE"] = (
+                old_disable_fresh_cache_envvar
+            )
+        torch.random.manual_seed(1234)
+
+    def tearDown(self):
+        super().tearDown()
+        clear_caches()
+
+    @unittest.skipIf(not HAS_CUTLASS, "requires CUTLASS")
+    @unittest.skipIf(not SM90OrLater, "requires SM90+")
+    @mock.patch.dict(os.environ, {"PATH": _get_path_without_sccache()})
+    def test_fallback_to_aten_when_cutlass_benchmarks_fail(self):
+        """ATen fallback is used when all CUTLASS benchmarks return inf."""
+        from torch._inductor.codegen.cutlass.kernel import CUTLASSTemplateCaller
+
+        def mock_cuda_benchmark(*args, **kwargs):
+            return float("inf")
+
+        with config.patch(
+            {
+                "max_autotune": True,
+                "max_autotune_gemm_backends": "CUTLASS,ATen",
+                "cuda.cutlass_max_profiling_configs": 4,
+            }
+        ):
+            with patch.object(CUTLASSTemplateCaller, "benchmark", mock_cuda_benchmark):
+                torch._dynamo.reset()
+                clear_caches()
+
+                @torch.compile
+                def fn(a, b):
+                    return torch.mm(a, b)
+
+                # Use dimensions that trigger CUTLASS selection
+                M, N, K = 256, 2048, 3520
+                a = torch.randn(M, K, device="cuda", dtype=torch.bfloat16)
+                b = torch.randn(K, N, device="cuda", dtype=torch.bfloat16)
+
+                # Should not crash - should fall back to ATen
+                result = fn(a, b)
+
+                # Verify result is correct (comparing against eager execution)
+                expected = torch.mm(a, b)
+                torch.testing.assert_close(result, expected, rtol=1e-2, atol=1e-2)
+
+    @unittest.skipIf(not HAS_CUTLASS, "requires CUTLASS")
+    @unittest.skipIf(not SM90OrLater, "requires SM90+")
+    @mock.patch.dict(os.environ, {"PATH": _get_path_without_sccache()})
+    def test_extern_kernel_included_in_prescreening(self):
+        """Test that ExternKernelCaller choices are included in prescreening candidates."""
+        prescreening_candidates = []
+
+        original_prescreen = AlgorithmSelectorCache.prescreen_choices
+
+        @staticmethod
+        def capturing_prescreen(choices, name, inputs_key, prescreen_cache):
+            result = original_prescreen(choices, name, inputs_key, prescreen_cache)
+            if result:
+                prescreening_candidates.extend(result)
+            return result
+
+        with config.patch(
+            {
+                "max_autotune": True,
+                "max_autotune_gemm_backends": "CUTLASS,ATen",
+                "cuda.cutlass_max_profiling_configs": 40,
+                "cuda.cutlass_prescreening": True,
+            }
+        ):
+            with patch.object(
+                AlgorithmSelectorCache, "prescreen_choices", capturing_prescreen
+            ):
+                torch._dynamo.reset()
+                clear_caches()
+
+                @torch.compile
+                def fn(a, b):
+                    return torch.mm(a, b)
+
+                M, N, K = 256, 2048, 3520
+                a = torch.randn(M, K, device="cuda", dtype=torch.bfloat16)
+                b = torch.randn(K, N, device="cuda", dtype=torch.bfloat16)
+
+                _ = fn(a, b)
+
+                # Prescreening requires >= 10 CUTLASS candidates; if triggered,
+                # verify ATen was included
+                if prescreening_candidates:
+                    extern_in_prescreening = any(
+                        isinstance(c, ExternKernelCaller)
+                        for c in prescreening_candidates
+                    )
+                    self.assertTrue(
+                        extern_in_prescreening,
+                        "ExternKernelCaller should be included in prescreening candidates",
+                    )
+                else:
+                    # If prescreening wasn't triggered (< 10 CUTLASS candidates),
+                    # that's OK -- ATen fallback is in the main benchmark path
+                    pass
+
+    @unittest.skipIf(not HAS_CUTLASS, "requires CUTLASS")
+    @unittest.skipIf(not SM90OrLater, "requires SM90+")
+    @mock.patch.dict(os.environ, {"PATH": _get_path_without_sccache()})
+    def test_addmm_fallback_on_cutlass_failure(self):
+        """addmm falls back to ATen when CUTLASS benchmarks fail."""
+        from torch._inductor.codegen.cutlass.kernel import CUTLASSTemplateCaller
+
+        def mock_cuda_benchmark(*args, **kwargs):
+            return float("inf")
+
+        with config.patch(
+            {
+                "max_autotune": True,
+                "max_autotune_gemm_backends": "CUTLASS,ATen",
+                "cuda.cutlass_max_profiling_configs": 4,
+            }
+        ):
+            with patch.object(CUTLASSTemplateCaller, "benchmark", mock_cuda_benchmark):
+                torch._dynamo.reset()
+                clear_caches()
+
+                @torch.compile
+                def fn(bias, a, b):
+                    return torch.addmm(bias, a, b)
+
+                # Use dimensions from issue #171094
+                M, N, K = 256, 2048, 3520
+                a = torch.randn(M, K, device="cuda", dtype=torch.bfloat16)
+                b = torch.randn(K, N, device="cuda", dtype=torch.bfloat16)
+                bias = torch.randn(N, device="cuda", dtype=torch.bfloat16)
+
+                # Should not crash
+                result = fn(bias, a, b)
+
+                # Verify correctness
+                expected = torch.addmm(bias, a, b)
+                torch.testing.assert_close(result, expected, rtol=1e-2, atol=1e-2)
+
+    @unittest.skipIf(not HAS_CUTLASS, "requires CUTLASS")
+    @unittest.skipIf(not SM90OrLater, "requires SM90+")
+    @mock.patch.dict(os.environ, {"PATH": _get_path_without_sccache()})
+    def test_issue_171094_dimensions(self):
+        """Regression test with exact dimensions from #171094."""
+        with config.patch(
+            {
+                "max_autotune": True,
+                "max_autotune_gemm_backends": "CUTLASS,ATen",
+                "cuda.cutlass_max_profiling_configs": 4,
+            }
+        ):
+            torch._dynamo.reset()
+            clear_caches()
+
+            @torch.compile
+            def fn(bias, a, b):
+                return torch.addmm(bias, a, b)
+
+            # Exact dimensions from issue #171094
+            M, K, N = 256, 3520, 2048
+            a = torch.randn(M, K, device="cuda", dtype=torch.bfloat16)
+            b = torch.randn(K, N, device="cuda", dtype=torch.bfloat16)
+            bias = torch.randn(N, device="cuda", dtype=torch.bfloat16)
+
+            # This should complete without CUDA illegal memory access
+            result = fn(bias, a, b)
+
+            # Verify the result is correct
+            expected = torch.addmm(bias, a, b)
+            torch.testing.assert_close(result, expected, rtol=1e-2, atol=1e-2)
+
+
+class TestCutlassSubprocessRouting(TestCase):
+    """Tests for CUTLASS subprocess benchmarking routing and error recovery."""
+
+    def test_cutlass_forces_subprocess_benchmarking(self):
+        """make_benchmark_fn routes to subprocess when CUTLASS choices are present."""
+        from torch._inductor.codegen.cutlass.kernel import CUTLASSTemplateCaller
+
+        mock_cutlass = mock.create_autospec(CUTLASSTemplateCaller, instance=True)
+        mock_extern = mock.create_autospec(ExternKernelCaller, instance=True)
+
+        with config.patch({"autotune_in_subproc": False}):
+            fn = AlgorithmSelectorCache.make_benchmark_fn(
+                choices=[mock_extern, mock_cutlass],
+                input_nodes=[],
+                layout=mock.MagicMock(),
+                input_gen_fns=None,
+            )
+            self.assertEqual(fn.func, AlgorithmSelectorCache.benchmark_in_sub_process)
+
+    def test_no_subprocess_without_cutlass(self):
+        """make_benchmark_fn uses current process when no CUTLASS and autotune_in_subproc=False."""
+        mock_extern = mock.create_autospec(ExternKernelCaller, instance=True)
+
+        with config.patch({"autotune_in_subproc": False}):
+            fn = AlgorithmSelectorCache.make_benchmark_fn(
+                choices=[mock_extern],
+                input_nodes=[],
+                layout=mock.MagicMock(),
+                input_gen_fns=None,
+            )
+            self.assertEqual(
+                fn.func, AlgorithmSelectorCache.benchmark_in_current_process
+            )
+
+    def test_subprocess_restart_on_illegal_address(self):
+        """TuningProcessPool restarts subprocess on cudaErrorIllegalAddress."""
+        from torch._inductor.autotune_process import TuningProcessPool
+
+        mock_process = mock.MagicMock()
+        mock_process.get.side_effect = RuntimeError(
+            "cudaErrorIllegalAddress: an illegal memory access was encountered"
+        )
+
+        mock_queue = mock.MagicMock()
+        mock_queue.get.return_value = mock_process
+
+        pool = TuningProcessPool.__new__(TuningProcessPool)
+        pool.process_queue = mock_queue
+
+        mock_choice = mock.MagicMock()
+        mock_choice.bmreq = mock.MagicMock()
+
+        result = pool.target(mock_choice)
+        self.assertEqual(result, float("inf"))
+        mock_process.restart.assert_called_once()
+
+    def test_subprocess_restart_on_launch_failure(self):
+        """TuningProcessPool restarts subprocess on cudaErrorLaunchFailure."""
+        from torch._inductor.autotune_process import TuningProcessPool
+
+        mock_process = mock.MagicMock()
+        mock_process.get.side_effect = RuntimeError("cudaErrorLaunchFailure")
+
+        mock_queue = mock.MagicMock()
+        mock_queue.get.return_value = mock_process
+
+        pool = TuningProcessPool.__new__(TuningProcessPool)
+        pool.process_queue = mock_queue
+
+        mock_choice = mock.MagicMock()
+        mock_choice.bmreq = mock.MagicMock()
+
+        result = pool.target(mock_choice)
+        self.assertEqual(result, float("inf"))
+        mock_process.restart.assert_called_once()
+
+    def test_subprocess_no_restart_on_other_errors(self):
+        """TuningProcessPool does not restart on non-sticky CUDA errors."""
+        from torch._inductor.autotune_process import TuningProcessPool
+
+        mock_process = mock.MagicMock()
+        mock_process.get.side_effect = RuntimeError("some other error")
+
+        mock_queue = mock.MagicMock()
+        mock_queue.get.return_value = mock_process
+
+        pool = TuningProcessPool.__new__(TuningProcessPool)
+        pool.process_queue = mock_queue
+
+        mock_choice = mock.MagicMock()
+        mock_choice.bmreq = mock.MagicMock()
+
+        result = pool.target(mock_choice)
+        self.assertEqual(result, float("inf"))
+        mock_process.restart.assert_not_called()
+
+    def test_cutlass_ordered_after_triton_in_subprocess(self):
+        """benchmark_in_sub_process orders Triton before CUTLASS."""
+        from torch._inductor.codegen.cutlass.kernel import CUTLASSTemplateCaller
+
+        mock_cutlass = mock.create_autospec(CUTLASSTemplateCaller, instance=True)
+        mock_triton = mock.MagicMock()  # Not extern, not CUTLASS
+
+        captured_choices = []
+
+        def capture_benchmark(choices):
+            captured_choices.extend(choices)
+            return dict.fromkeys(choices, 1.0)
+
+        with (
+            mock.patch(
+                "torch._inductor.autotune_process.benchmark_in_sub_process",
+                side_effect=capture_benchmark,
+            ),
+            mock.patch.object(
+                AlgorithmSelectorCache,
+                "benchmark_in_current_process",
+                return_value={},
+            ),
+            mock.patch.object(
+                AlgorithmSelectorCache,
+                "_is_extern",
+                side_effect=lambda c: False,
+            ),
+        ):
+            AlgorithmSelectorCache.benchmark_in_sub_process(
+                choices=[mock_cutlass, mock_triton],
+                input_nodes=[],
+                layout=mock.MagicMock(),
+                input_gen_fns=None,
+            )
+
+            self.assertEqual(len(captured_choices), 2)
+            # Triton (non-CUTLASS) should come before CUTLASS
+            self.assertNotIsInstance(captured_choices[0], CUTLASSTemplateCaller)
+            self.assertIsInstance(captured_choices[1], CUTLASSTemplateCaller)
+
+    def test_prescreen_includes_extern_with_enough_cutlass(self):
+        """prescreen_choices includes ExternKernelCaller when >= 10 CUTLASS candidates."""
+        from torch._inductor.codegen.cutlass.kernel import CUTLASSTemplateCaller
+
+        def make_cutlass_mock(swizzle="2"):
+            m = mock.create_autospec(CUTLASSTemplateCaller, instance=True)
+            m.info_dict.return_value = {"swizzle": swizzle}
+            return m
+
+        cutlass_choices = [make_cutlass_mock() for _ in range(12)]
+        extern_choice = mock.create_autospec(ExternKernelCaller, instance=True)
+        choices = [extern_choice] + cutlass_choices
+
+        with config.patch(
+            {
+                "cuda.cutlass_prescreening": True,
+                "cuda.cutlass_max_profiling_swizzle_options": [1, 2],
+            }
+        ):
+            result = AlgorithmSelectorCache.prescreen_choices(
+                choices, "test_op", "test_key", {}
+            )
+
+        self.assertGreater(len(result), 0)
+        self.assertTrue(
+            any(isinstance(c, ExternKernelCaller) for c in result),
+            "ExternKernelCaller should be in prescreening candidates",
+        )
+        # Extern should be first in the list
+        self.assertIsInstance(result[0], ExternKernelCaller)
+
+    def test_prescreen_skipped_with_few_cutlass(self):
+        """prescreen_choices returns [] when < 10 CUTLASS candidates."""
+        from torch._inductor.codegen.cutlass.kernel import CUTLASSTemplateCaller
+
+        def make_cutlass_mock(swizzle="2"):
+            m = mock.create_autospec(CUTLASSTemplateCaller, instance=True)
+            m.info_dict.return_value = {"swizzle": swizzle}
+            return m
+
+        cutlass_choices = [make_cutlass_mock() for _ in range(5)]
+        extern_choice = mock.create_autospec(ExternKernelCaller, instance=True)
+        choices = [extern_choice] + cutlass_choices
+
+        with config.patch(
+            {
+                "cuda.cutlass_prescreening": True,
+                "cuda.cutlass_max_profiling_swizzle_options": [1, 2],
+            }
+        ):
+            result = AlgorithmSelectorCache.prescreen_choices(
+                choices, "test_op", "test_key", {}
+            )
+
+        self.assertEqual(result, [])
+
+
+if __name__ == "__main__":
+    run_tests()
diff --git a/test/inductor/test_debug_trace.py b/test/inductor/test_debug_trace.py
index 7b999a1e692d7..7a1793206f327 100644
--- a/test/inductor/test_debug_trace.py
+++ b/test/inductor/test_debug_trace.py
@@ -28,7 +28,8 @@
 
 
 def filesize(filename: Path):
-    assert filename.exists(), f"{filename} is missing"
+    if not filename.exists():
+        raise AssertionError(f"{filename} is missing")
     return os.stat(filename).st_size
 
 
@@ -70,7 +71,8 @@ def fn(a, b):
                 break
         self.assertTrue(m, "debug trace file path not found in logs")
         # For type checking, have to ensure it's not none.
-        assert m is not None
+        if m is None:
+            raise AssertionError
         filename = Path(m.group(1))
         self.assertTrue(filename.is_dir())
         self.assertGreater(filesize(filename / "fx_graph_readable.py"), 512)
@@ -87,6 +89,8 @@ def fn(a, b):
 op0.writes = [MemoryDep('buf0', c0, {c0: 256})]
 op0.unmet_dependencies = []
 op0.met_dependencies = [MemoryDep('arg0_1', c0, {c0: 256})]
+op0.min_input_distance = 0
+op0.max_input_distance = 0
 op0.outputs = [
     buf0: ComputedBuffer
     buf0.layout = FixedLayout('cpu', torch.float32, size=[16, 16], stride=[16, 1])
@@ -114,6 +118,8 @@ def body(self, ops):
 op1.writes = [MemoryDep('buf1', c0, {c0: 256})]
 op1.unmet_dependencies = [MemoryDep('buf0', c0, {c0: 256})]
 op1.met_dependencies = []
+op1.min_input_distance = 1
+op1.max_input_distance = 1
 op1.outputs = [
     buf1: ComputedBuffer
     buf1.layout = FixedLayout('cpu', torch.float32, size=[16, 16], stride=[16, 1])
@@ -141,6 +147,8 @@ def body(self, ops):
 op2.writes = [StarDep(name='buf2', mode=None)]
 op2.unmet_dependencies = [StarDep(name='buf1', mode=None)]
 op2.met_dependencies = [StarDep(name='arg1_1', mode=None)]
+op2.min_input_distance = 2
+op2.max_input_distance = 2
 op2.outputs = [
     buf2: ExternKernelOut
     buf2.layout = FixedLayout('cpu', torch.float32, size=[16, 16], stride=[16, 1])
@@ -158,6 +166,8 @@ def body(self, ops):
 op0_op1.writes = [MemoryDep('buf0', c0, {c0: 256}), MemoryDep('buf1', c0, {c0: 256})]
 op0_op1.unmet_dependencies = []
 op0_op1.met_dependencies = [MemoryDep('arg0_1', c0, {c0: 256})]
+op0_op1.min_input_distance = 0
+op0_op1.max_input_distance = 1
 op0_op1.outputs = [
     buf0: ComputedBuffer
     buf0.layout = FixedLayout('cpu', torch.float32, size=[16, 16], stride=[16, 1])
@@ -171,6 +181,8 @@ def body(self, ops):
 op0.writes = [MemoryDep('buf0', c0, {c0: 256})]
 op0.unmet_dependencies = []
 op0.met_dependencies = [MemoryDep('arg0_1', c0, {c0: 256})]
+op0.min_input_distance = 0
+op0.max_input_distance = 0
 op0.outputs = [
     buf0: ComputedBuffer
     buf0.layout = FixedLayout('cpu', torch.float32, size=[16, 16], stride=[16, 1])
@@ -197,6 +209,8 @@ def body(self, ops):
 op1.writes = [MemoryDep('buf1', c0, {c0: 256})]
 op1.unmet_dependencies = [MemoryDep('buf0', c0, {c0: 256})]
 op1.met_dependencies = []
+op1.min_input_distance = 1
+op1.max_input_distance = 1
 op1.outputs = [
     buf1: ComputedBuffer
     buf1.layout = FixedLayout('cpu', torch.float32, size=[16, 16], stride=[16, 1])
@@ -224,6 +238,8 @@ def body(self, ops):
 op2.writes = [StarDep(name='buf2', mode=None)]
 op2.unmet_dependencies = [StarDep(name='buf1', mode=None)]
 op2.met_dependencies = [StarDep(name='arg1_1', mode=None)]
+op2.min_input_distance = 2
+op2.max_input_distance = 2
 op2.outputs = [
     buf2: ExternKernelOut
     buf2.layout = FixedLayout('cpu', torch.float32, size=[16, 16], stride=[16, 1])
diff --git a/test/inductor/test_decompose_mem_bound_mm.py b/test/inductor/test_decompose_mem_bound_mm.py
index bbbff963d3489..79cb68b2d71ae 100644
--- a/test/inductor/test_decompose_mem_bound_mm.py
+++ b/test/inductor/test_decompose_mem_bound_mm.py
@@ -86,7 +86,8 @@ def compare_dict_tensors(self, ref_dict, res_dict, rtol=None, atol=None):
             return False
         for key1 in ref_dict:
             key2 = "_orig_mod." + key1
-            assert key2 in res_dict, f"{key1} does not exist in traced module"
+            if key2 not in res_dict:
+                raise AssertionError(f"{key1} does not exist in traced module")
             if not torch.allclose(
                 ref_dict[key1], res_dict[key2], rtol=self.rtol, atol=self.atol
             ):
diff --git a/test/inductor/test_deterministic.py b/test/inductor/test_deterministic.py
index b03e5b1ffad78..be14ab8a6697f 100644
--- a/test/inductor/test_deterministic.py
+++ b/test/inductor/test_deterministic.py
@@ -14,14 +14,13 @@
 from torch._inductor.utils import fresh_cache
 from torch.testing._internal.common_utils import (
     instantiate_parametrized_tests,
+    IS_FBCODE,
     parametrize,
-    skipIfXpu,
 )
 from torch.testing._internal.inductor_utils import (
     GPU_TYPE,
     HAS_GPU_AND_TRITON,
     IS_BIG_GPU,
-    IS_FBCODE,
 )
 
 
@@ -48,7 +47,6 @@ def test_use_deterministic_algorithsm(self):
         finally:
             torch.use_deterministic_algorithms(old_val, warn_only=True)
 
-    @skipIfXpu(msg="pad_mm is not enabled for XPU.")
     @parametrize("deterministic", [False, True])
     def test_mm_padding(self, deterministic):
         with inductor_config.patch(deterministic=deterministic):
@@ -115,6 +113,25 @@ def foo(x):
             else:
                 self.assertTrue(counters["inductor"]["coordesc_tuning_bench"] > 0)
 
+    def test_reorder_for_locality_preserves_randint_order(self):
+        with inductor_config.patch(fallback_random=True):
+
+            def fn():
+                torch.manual_seed(0)
+                out = torch.randint(0, 100, (4, 1), dtype=torch.int64)
+                _ = torch.randint(0, 100, (2, 1), dtype=torch.int64)
+                return out
+
+            compiled = torch.compile(fn, backend="inductor")
+
+            torch.manual_seed(0)
+            eager = fn()
+
+            torch.manual_seed(0)
+            compiled_out = compiled()
+
+            torch.testing.assert_close(eager, compiled_out)
+
     @unittest.skipIf(IS_FBCODE, "Skipping run2run determinism test in fbcode")
     @parametrize("model_name", ["GoogleFnet", "BertForMaskedLM", "DistillGPT2"])
     @parametrize("training_or_inference", ["training", "inference"])
diff --git a/test/inductor/test_device_assert.py b/test/inductor/test_device_assert.py
index cbeb7960f2f55..10de3a5b386f4 100644
--- a/test/inductor/test_device_assert.py
+++ b/test/inductor/test_device_assert.py
@@ -22,11 +22,11 @@ def test_assert_should_throw(self, backend):
         def func():
             a = torch.tensor([1.0, -2.0], device="cpu")
             result = torch.all(a > 0)
-            assert result, "should throw"
+            assert result, "should throw"  # noqa: S101
 
         def func_inline():
             a = torch.tensor([1.0, -2.0], device="cpu")
-            assert torch.all(a > 0), "should throw"
+            assert torch.all(a > 0), "should throw"  # noqa: S101
 
         with self.assertRaisesRegex(RuntimeError, "should throw"):
             torch._dynamo.reset()
@@ -43,11 +43,11 @@ def test_assert_should_not_throw(self, backend):
         def func():
             a = torch.tensor([1.0, 2.0], device="cpu")
             result = torch.all(a > 0)
-            assert result, "should throw"
+            assert result, "should throw"  # noqa: S101
 
         def func_inline():
             a = torch.tensor([1.0, 2.0], device="cpu")
-            assert torch.all(a > 0), "should throw"
+            assert torch.all(a > 0), "should throw"  # noqa: S101
 
         torch._dynamo.reset()
         f_c = torch.compile(func, backend=backend)
@@ -65,7 +65,7 @@ def test_assert_fusion(self):
         def func():
             a = torch.tensor([1.0, 2.0], device=device_type)
             result = torch.all(a > 0)
-            assert result, "should throw"
+            assert result, "should throw"  # noqa: S101
 
         torch._dynamo.reset()
         f_c = torch.compile(func, backend="inductor")
@@ -82,7 +82,7 @@ def test_run_assert_triton(self):
         def fn():
             a = torch.tensor([1.0, 2.0], device=device_type)
             result = torch.all(a > 0)
-            assert result, "should throw"
+            assert result, "should throw"  # noqa: S101
 
         def should_not_throw(fn):
             try:
diff --git a/test/inductor/test_distributed_patterns.py b/test/inductor/test_distributed_patterns.py
index 05ef76d6636ac..9a8f9a79ddf1f 100644
--- a/test/inductor/test_distributed_patterns.py
+++ b/test/inductor/test_distributed_patterns.py
@@ -93,14 +93,14 @@ def bw_post_hook(mod, gI, gO):
 
 def init_module_bw_hooks(allow_eager):
     def bw_pre_hook(mod, gO):
-        assert allow_eager or torch._dynamo.is_compiling()
-        assert mod.weight.size() == (10, 10)
+        assert allow_eager or torch._dynamo.is_compiling()  # noqa: S101
+        assert mod.weight.size() == (10, 10)  # noqa: S101
         mod.hook_count_pre.add_(1)
         return (torch.sin(gO[0] + 1.2),)
 
     def bw_post_hook(mod, gI, gO):
-        assert allow_eager or torch._dynamo.is_compiling()
-        assert mod.weight.size() == (10, 10)
+        assert allow_eager or torch._dynamo.is_compiling()  # noqa: S101
+        assert mod.weight.size() == (10, 10)  # noqa: S101
         mod.hook_count_post.add_(1)
         return (torch.sin(gI[0] + 3.4),)
 
@@ -214,7 +214,7 @@ def _test_storage_resize_nonzero(self, device):
         @torch.compile(fullgraph=True)
         def fn(x, out):
             y = torch.sin(x)
-            assert out.untyped_storage().size() == 0
+            assert out.untyped_storage().size() == 0  # noqa: S101
             out.untyped_storage().resize_(x.untyped_storage().size())
             out.copy_(y.cos())
 
diff --git a/test/inductor/test_dropout_align_random_eager.py b/test/inductor/test_dropout_align_random_eager.py
new file mode 100644
index 0000000000000..78746e0113377
--- /dev/null
+++ b/test/inductor/test_dropout_align_random_eager.py
@@ -0,0 +1,482 @@
+# Owner(s): ["module: inductor"]
+
+import struct
+import time
+
+import pytest
+
+import torch
+from torch._inductor import config
+from torch._inductor.test_case import run_tests, TestCase as InductorTestCase
+from torch._inductor.utils import run_and_get_code
+from torch.testing import FileCheck
+from torch.testing._internal.common_utils import IS_LINUX
+from torch.testing._internal.inductor_utils import (
+    GPU_TYPE,
+    HAS_CUDA_AND_TRITON,
+    requires_gpu,
+)
+
+
+# ───────────────────────────────────────────────────────────────
+# Global config
+# ───────────────────────────────────────────────────────────────
+BASE_SEED = 1234
+DROPOUT_P = 0.5
+FFN_DIM = 3072
+HIDDEN_DIM = 1024
+BATCH = 3
+SEQ_LEN = 512
+
+
+# ───────────────────────────────────────────────────────────────
+# Model under test
+# ───────────────────────────────────────────────────────────────
+class LinearBlock(torch.nn.Module):
+    def __init__(self, hidden_dim: int, ffn_dim: int, dropout: float = DROPOUT_P):
+        super().__init__()
+        self.net = torch.nn.Sequential(
+            torch.nn.Linear(hidden_dim, ffn_dim),
+            torch.nn.Dropout(dropout),
+            torch.nn.ReLU(inplace=False),
+            torch.nn.Linear(ffn_dim, hidden_dim),
+        )
+
+    def forward(self, x: torch.Tensor):
+        return self.net(x)
+
+
+class MultiDropoutBlock(torch.nn.Module):
+    """Block with multiple Dropout ops to stress RNG alignment."""
+
+    def __init__(self, hidden_dim: int, ffn_dim: int, dropout: float = DROPOUT_P):
+        super().__init__()
+        self.net = torch.nn.Sequential(
+            torch.nn.Linear(hidden_dim, ffn_dim),
+            torch.nn.Dropout(dropout),
+            torch.nn.ReLU(inplace=False),
+            torch.nn.Dropout(dropout),
+            torch.nn.Linear(ffn_dim, hidden_dim),
+            torch.nn.Dropout(dropout),
+        )
+
+    def forward(self, x: torch.Tensor):
+        return self.net(x)
+
+
+def build_models(dropout: float, *, mode=None, dynamic: bool = False):
+    eager = LinearBlock(HIDDEN_DIM, FFN_DIM, dropout)
+    compiled = LinearBlock(HIDDEN_DIM, FFN_DIM, dropout)
+    compiled.load_state_dict(eager.state_dict())
+    compiled = torch.compile(compiled, mode=mode, dynamic=dynamic)
+    return eager, compiled
+
+
+# ───────────────────────────────────────────────────────────────
+# Helpers
+# ───────────────────────────────────────────────────────────────
+def _set_seed(base: int = BASE_SEED):
+    torch.manual_seed(base)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed(base)
+
+
+def _sync(x: torch.Tensor):
+    if x.is_cuda:
+        torch.cuda.synchronize()
+
+
+def _timed_run(model, x, backward: bool = False):
+    _sync(x)
+    t0 = time.time()
+    y = model(x)
+    if backward:
+        (y.square().mean()).backward()
+    _sync(x)
+    return (time.time() - t0) * 1e3, y
+
+
+def _cuda_rng_u64_seed_off():
+    """Return (seed, offset) extracted from torch.cuda.get_rng_state()."""
+    st = torch.cuda.get_rng_state()
+    seed = struct.unpack("<Q", st[0:8].cpu().numpy().tobytes())[0]
+    off = struct.unpack("<Q", st[8:24].cpu().numpy().tobytes())[0]
+    return seed, off
+
+
+def set_seed(seed):
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed(seed)
+
+
+def dropout_parity(shape, p=0.3, dtype=torch.float32, seed=1234):
+    """Returns (masks_equal, eager_out, compiled_out)."""
+    DEVICE = torch.device("cuda")
+    torch._dynamo.reset()
+    x = torch.ones(shape, device=DEVICE, dtype=dtype)
+    drop_e = torch.nn.Dropout(p).to(DEVICE).train()
+    drop_c = torch.compile(torch.nn.Dropout(p).to(DEVICE).train())
+
+    set_seed(seed)
+    out_e = drop_e(x)
+    set_seed(seed)
+    out_c = drop_c(x)
+
+    masks_eq = torch.equal(out_e != 0, out_c != 0)
+    return masks_eq, out_e, out_c
+
+
+# ───────────────────────────────────────────────────────────────
+# Test class (Inductor idioms)
+# ───────────────────────────────────────────────────────────────
+@pytest.mark.skipif(
+    not (IS_LINUX and HAS_CUDA_AND_TRITON),
+    reason="Inductor CUDA dropout alignment tests require Linux and CUDA",
+)
+@config.patch(align_random_eager=True)
+class TestDropoutAlignRandomEager(InductorTestCase):
+    @requires_gpu()
+    def test_linear_block_compile_parity_forward(self):
+        device = torch.device(GPU_TYPE)
+
+        for training in (False, True):
+            eager, compiled = build_models(DROPOUT_P)
+            eager.to(device)
+            compiled.to(device)
+
+            if training:
+                eager.train()
+                compiled.train()
+            else:
+                eager.eval()
+                compiled.eval()
+
+            x = torch.randn(BATCH, SEQ_LEN, HIDDEN_DIM, device=device)
+
+            # same seed before both runs (align dropout masks)
+            _set_seed(BASE_SEED)
+            with torch.no_grad():
+                y_eager = eager(x)
+
+            _set_seed(BASE_SEED)
+            with torch.no_grad():
+                y_comp = compiled(x)
+
+            torch.testing.assert_close(y_eager, y_comp, rtol=0.0, atol=0.0)
+
+    @requires_gpu()
+    def test_linear_block_compile_parity_backward(self):
+        device = torch.device(GPU_TYPE)
+
+        eager, compiled = build_models(DROPOUT_P)
+        eager.to(device)
+        compiled.to(device)
+        eager.train()
+        compiled.train()
+
+        x = torch.randn(BATCH, SEQ_LEN, HIDDEN_DIM, device=device)
+
+        # eager fwd+bwd
+        _set_seed(BASE_SEED)
+        y_eager = eager(x)
+        (y_eager.square().mean()).backward()
+
+        # compiled fwd+bwd (Re-seed)
+        for p in compiled.parameters():
+            p.grad = None
+        _set_seed(BASE_SEED)
+        y_comp = compiled(x)
+        (y_comp.square().mean()).backward()
+
+        # outputs
+        torch.testing.assert_close(
+            y_eager.detach(), y_comp.detach(), rtol=1e-3, atol=1e-4
+        )
+        # grads
+        for p_ref, p_new in zip(eager.parameters(), compiled.parameters()):
+            self.assertIsNotNone(p_ref.grad)
+            self.assertIsNotNone(p_new.grad)
+            torch.testing.assert_close(p_ref.grad, p_new.grad, rtol=1e-3, atol=1e-5)
+
+    @requires_gpu()
+    def test_dropout_mask_parity_and_rng_offset_cuda(self):
+        device = torch.device(GPU_TYPE)
+        H, W = BATCH * SEQ_LEN, FFN_DIM
+
+        dtypes = [torch.float32, torch.float16, torch.bfloat16]
+        for dtype in dtypes:
+            if dtype is torch.bfloat16 and not torch.cuda.is_bf16_supported():
+                continue
+
+            x = torch.ones((H, W), device=device, dtype=dtype)
+
+            # Eager
+            _set_seed(BASE_SEED)
+            seed0_e, off0_e = _cuda_rng_u64_seed_off()
+            drop_e = torch.nn.Dropout(DROPOUT_P).to(device=device, dtype=dtype).train()
+            mask_e = drop_e(x) != 0
+            seed1_e, off1_e = _cuda_rng_u64_seed_off()
+            delta_e = off1_e - off0_e
+
+            # Compiled
+            _set_seed(BASE_SEED)
+            seed0_c, off0_c = _cuda_rng_u64_seed_off()
+            drop_c = torch.nn.Dropout(DROPOUT_P).to(device=device, dtype=dtype)
+            drop_c = torch.compile(drop_c)
+            drop_c.train()
+            mask_c = drop_c(x) != 0
+            seed1_c, off1_c = _cuda_rng_u64_seed_off()
+            delta_c = off1_c - off0_c
+
+            mismatch_ratio = (mask_e != mask_c).float().mean().item()
+            self.assertLessEqual(
+                mismatch_ratio,
+                1e-4,
+                msg=f"Dropout mask mismatch ratio too high: {mismatch_ratio:.8f}",
+            )
+            self.assertEqual(seed0_e, BASE_SEED)
+            self.assertEqual(seed0_c, BASE_SEED)
+            self.assertEqual(
+                delta_e,
+                delta_c,
+                msg=f"RNG offset delta mismatch: eager={delta_e}, compiled={delta_c}",
+            )
+
+    # ───────────────────────────────────────────────────────────
+    # multiple dropouts + multiple iterations
+    # ───────────────────────────────────────────────────────────
+    @requires_gpu()
+    def test_multi_dropout_multi_iterations_parity(self):
+        device = torch.device(GPU_TYPE)
+
+        eager = MultiDropoutBlock(HIDDEN_DIM, FFN_DIM, DROPOUT_P).to(device)
+        compiled = MultiDropoutBlock(HIDDEN_DIM, FFN_DIM, DROPOUT_P).to(device)
+        compiled.load_state_dict(eager.state_dict())
+        compiled = torch.compile(compiled)
+
+        eager.train()
+        compiled.train()
+
+        num_iters = 10
+        for i in range(num_iters):
+            seed = BASE_SEED + i
+            x = torch.randn(BATCH, SEQ_LEN, HIDDEN_DIM, device=device)
+
+            _set_seed(seed)
+            y_eager = eager(x)
+
+            _set_seed(seed)
+            y_comp = compiled(x)
+
+            mismatch_ratio = ((y_eager != 0) != (y_comp != 0)).float().mean().item()
+            self.assertLessEqual(mismatch_ratio, 1e-5)
+
+    # ───────────────────────────────────────────────────────────
+    # dynamic shapes test (a)
+    # ───────────────────────────────────────────────────────────
+    @requires_gpu()
+    def test_dropout_parity_dynamic_shapes(self):
+        device = torch.device(GPU_TYPE)
+
+        eager = LinearBlock(HIDDEN_DIM, FFN_DIM, DROPOUT_P).to(device)
+        compiled = LinearBlock(HIDDEN_DIM, FFN_DIM, DROPOUT_P).to(device)
+        compiled.load_state_dict(eager.state_dict())
+        compiled = torch.compile(compiled, dynamic=True)
+
+        eager.train()
+        compiled.train()
+
+        shapes = [
+            (BATCH, 512, HIDDEN_DIM),
+            (BATCH, 128, HIDDEN_DIM),
+        ]
+
+        for shape in shapes:
+            x = torch.randn(*shape, device=device)
+
+            _set_seed(BASE_SEED)
+            y_eager = eager(x)
+
+            _set_seed(BASE_SEED)
+            y_comp = compiled(x)
+
+            torch.testing.assert_close(y_eager, y_comp, rtol=1e-5, atol=1e-6)
+
+    # ───────────────────────────────────────────────────────────
+    # cudagraphs test via mode='reduce-overhead' (b)
+    # ───────────────────────────────────────────────────────────
+    @requires_gpu()
+    def test_dropout_parity_cudagraphs_reduce_overhead(self):
+        device = torch.device(GPU_TYPE)
+
+        eager = LinearBlock(HIDDEN_DIM, FFN_DIM, DROPOUT_P).to(device)
+        compiled = LinearBlock(HIDDEN_DIM, FFN_DIM, DROPOUT_P).to(device)
+        compiled.load_state_dict(eager.state_dict())
+        compiled = torch.compile(compiled, mode="reduce-overhead")
+
+        eager.train()
+        compiled.train()
+
+        x = torch.randn(BATCH, SEQ_LEN, HIDDEN_DIM, device=device)
+
+        _set_seed(BASE_SEED)
+        y_eager = eager(x)
+
+        _set_seed(BASE_SEED)
+        y_comp = compiled(x)
+
+        torch.testing.assert_close(y_eager, y_comp, rtol=0.0, atol=0.0)
+
+    # ───────────────────────────────────────────────────────────
+    # Codegen sanity: run_and_get_code + FileCheck
+    # ───────────────────────────────────────────────────────────
+    @requires_gpu()
+    def test_inductor_generated_code_contains_dropout(self):
+        device = torch.device(GPU_TYPE)
+        x = torch.randn(BATCH, SEQ_LEN, HIDDEN_DIM, device=device)
+
+        model = LinearBlock(HIDDEN_DIM, FFN_DIM, DROPOUT_P).to(device)
+        model.train()
+        compiled = torch.compile(model)
+
+        def fn(inp):
+            return compiled(inp)
+
+        _, codes = run_and_get_code(fn, x)
+        self.assertTrue(codes, msg="Expected inductor to generate at least one kernel")
+
+        # Minimal sanity check that generated code mentions dropout.
+        FileCheck().check("dropout").run(codes[0])
+
+    # ───────────────────────────────────────────────────────────
+    # Optional: perf smoke (GPU only)
+    # ───────────────────────────────────────────────────────────
+    @requires_gpu()
+    def test_perf_smoke_cuda(self):
+        device = torch.device(GPU_TYPE)
+        x = torch.randn(BATCH, SEQ_LEN, HIDDEN_DIM, device=device)
+
+        eager, compiled = build_models(DROPOUT_P)
+        eager.to(device)
+        compiled.to(device)
+        eager.eval()
+        compiled.eval()
+
+        # warm up
+        _timed_run(eager, x, backward=False)
+        _timed_run(compiled, x, backward=False)
+
+        t_eager, _ = _timed_run(eager, x, backward=False)
+        t_comp, _ = _timed_run(compiled, x, backward=False)
+
+        self.assertGreater(t_comp, 0)
+        self.assertGreater(t_eager, 0)
+
+    # ───────────────────────────────────────────────────────────
+    # Helper for primitive random parity (rand / randn / randint)
+    # ───────────────────────────────────────────────────────────
+    def _run_primitive_random_parity(self, kind, device, shape):
+        if kind == "rand":
+
+            def eager():
+                return torch.rand(shape, device=device)
+
+            compiled = torch.compile(eager)
+
+        elif kind == "randn":
+
+            def eager():
+                return torch.randn(shape, device=device)
+
+            compiled = torch.compile(eager)
+
+        elif kind == "randint":
+
+            def eager():
+                return torch.randint(0, 2**31 - 1, shape, device=device)
+
+            compiled = torch.compile(eager)
+
+        else:
+            raise AssertionError(f"unknown primitive random kind: {kind}")
+
+        _set_seed(BASE_SEED)
+        out_eager = eager()
+
+        _set_seed(BASE_SEED)
+        out_comp = compiled()
+
+        torch.testing.assert_close(out_eager, out_comp, rtol=0.0, atol=0.0)
+
+    # ───────────────────────────────────────────────────────────
+    # Primitive random fns: rand / randn / randint -> mark as XFAIL
+    # ───────────────────────────────────────────────────────────
+    @requires_gpu()
+    @pytest.mark.xfail(
+        reason="primitive torch.rand parity is tracked as future work",
+        strict=False,
+    )
+    def test_primitive_rand_parity(self):
+        device = torch.device(GPU_TYPE)
+        shape = (BATCH, SEQ_LEN, HIDDEN_DIM)
+        self._run_primitive_random_parity("rand", device, shape)
+
+    @requires_gpu()
+    @pytest.mark.xfail(
+        reason="primitive torch.randn parity is tracked as future work",
+        strict=False,
+    )
+    def test_primitive_randn_parity(self):
+        device = torch.device(GPU_TYPE)
+        shape = (BATCH, SEQ_LEN, HIDDEN_DIM)
+        self._run_primitive_random_parity("randn", device, shape)
+
+    @requires_gpu()
+    @pytest.mark.xfail(
+        reason="primitive torch.randint parity is tracked as future work",
+        strict=False,
+    )
+    def test_primitive_randint_parity(self):
+        device = torch.device(GPU_TYPE)
+        shape = (BATCH, SEQ_LEN, HIDDEN_DIM)
+        self._run_primitive_random_parity("randint", device, shape)
+
+    # ───────────────────────────────────────────────────────────
+    # nn.Dropout as primitive RNG consumer (should PASS)
+    # ───────────────────────────────────────────────────────────
+    @requires_gpu()
+    def test_primitive_nn_dropout_parity(self):
+        device = torch.device(GPU_TYPE)
+        shape = (BATCH, SEQ_LEN, HIDDEN_DIM)
+
+        x = torch.ones(shape, device=device)
+
+        drop_eager = torch.nn.Dropout(DROPOUT_P).to(device).train()
+        drop_compiled = torch.nn.Dropout(DROPOUT_P).to(device).train()
+        drop_compiled.load_state_dict(drop_eager.state_dict())
+        drop_compiled = torch.compile(drop_compiled)
+
+        _set_seed(BASE_SEED)
+        out_eager = drop_eager(x)
+
+        _set_seed(BASE_SEED)
+        out_comp = drop_compiled(x)
+
+        torch.testing.assert_close(out_eager, out_comp, rtol=0.0, atol=0.0)
+
+    # ───────────────────────────────────────────────────────────
+    # Large seed (>32-bit) packing truncation
+    # Seed and base are packed into int64 as (seed << 32) | base.
+    # Seeds > 2^32 overflow.
+    # ───────────────────────────────────────────────────────────
+    @requires_gpu()
+    def test_large_seed(self):
+        for seed in [2**33 + 1, 2**40 + 12345]:
+            with self.subTest(seed=seed):
+                masks_eq, _, _ = dropout_parity((1024,), seed=seed)
+                self.assertTrue(masks_eq, f"seed={seed}: mask mismatch")
+
+
+if __name__ == "__main__":
+    if IS_LINUX and HAS_CUDA_AND_TRITON:
+        run_tests(needs="filelock")
diff --git a/test/inductor/test_efficient_conv_bn_eval.py b/test/inductor/test_efficient_conv_bn_eval.py
index 86b6b6ac8a0d8..7d69e8a18191f 100644
--- a/test/inductor/test_efficient_conv_bn_eval.py
+++ b/test/inductor/test_efficient_conv_bn_eval.py
@@ -14,6 +14,7 @@
 sys.path.append(pytorch_test_dir)
 
 from torch._dynamo.utils import counters
+from torch._functorch import config as functorch_config
 from torch._inductor import config as inductor_config
 from torch._inductor.test_case import TestCase
 from torch.testing._internal.common_cuda import tf32_on_and_off
@@ -96,6 +97,7 @@ def forward(self, x):
 class EfficientConvBNEvalTemplate(TestCase):
     @tf32_on_and_off(0.003)
     @inductor_config.patch({"efficient_conv_bn_eval_fx_passes": True})
+    @functorch_config.patch({"enable_autograd_cache": False})
     def test_basic(self):
         def test_conv_bn_eval(
             test_class, use_bias, module, sync_bn, decompose_nn_module
diff --git a/test/inductor/test_extension_backend.py b/test/inductor/test_extension_backend.py
index 5b1e4062da603..08b458d761f46 100644
--- a/test/inductor/test_extension_backend.py
+++ b/test/inductor/test_extension_backend.py
@@ -103,7 +103,8 @@ def setUp(self):
         # this file, so we'll change the working directory temporarily
         self.old_working_dir = os.getcwd()
         os.chdir(os.path.dirname(os.path.abspath(__file__)))
-        assert self.module is not None
+        if self.module is None:
+            raise AssertionError
 
     def tearDown(self):
         super().tearDown()
diff --git a/test/inductor/test_flex_attention.py b/test/inductor/test_flex_attention.py
index 972e74c34dfa6..d4bab42d80173 100644
--- a/test/inductor/test_flex_attention.py
+++ b/test/inductor/test_flex_attention.py
@@ -14,7 +14,7 @@
 from contextlib import contextmanager
 from dataclasses import dataclass
 from itertools import product
-from typing import Optional, TypeVar, Union
+from typing import TypeVar
 from unittest import expectedFailure, mock, skip, skipUnless
 from unittest.mock import patch
 
@@ -70,6 +70,11 @@
     skipXPUIf,
 )
 from torch.testing._internal.common_quantized import _snr
+from torch.testing._internal.common_utils import (  # noqa: F401
+    MI200_ARCH,
+    skipIfRocm,
+    skipIfRocmArch,
+)
 from torch.testing._internal.inductor_utils import HAS_GPU
 from torch.utils._triton import has_triton, has_triton_tma_device
 
@@ -96,7 +101,7 @@
 
 
 def large_tensor_test_class(
-    size: str, device: Optional[Union[torch.device, str]] = None
+    size: str, device: torch.device | str | None = None
 ) -> Callable[[type[T]], type[T]]:
     def decorator(cls: type[T]) -> type[T]:
         for name, method in list(cls.__dict__.items()):
@@ -178,6 +183,21 @@ def create_attention(score_mod, block_mask, enable_gqa=False, kernel_options=Non
     )
 
 
+def flex_attention_fwd(q, k, v, score_mod, block_mask, scale):
+    # Uses the HOP directly because flex_attention_backward expects lse in log2
+    # scale, but the public flex_attention API converts lse to natural log.
+    out, lse, _ = flex_attention_hop(
+        q,
+        k,
+        v,
+        score_mod,
+        block_mask.as_tuple(),
+        scale,
+        {},
+    )
+    return out, lse
+
+
 def create_block_mask_test(score_mod, query, key):
     block_mask = create_block_mask(
         score_mod,
@@ -427,7 +447,7 @@ def query_key_value_clones(
     query: torch.Tensor,
     key: torch.Tensor,
     value: torch.Tensor,
-    dtype: Optional[torch.dtype] = None,
+    dtype: torch.dtype | None = None,
 ):
     """Clones the query, key, and value tensors and moves them to the specified dtype."""
     if dtype is None:
@@ -462,7 +482,7 @@ def _check_equal(
         ref_out: torch.Tensor,
         compiled_out: torch.Tensor,
         fudge_factor: float,
-        tensor_name: Optional[str] = None,
+        tensor_name: str | None = None,
         fudge_atol: float = 0,
     ):
         compiled_error = (golden_out - compiled_out).abs().mean()
@@ -548,11 +568,11 @@ def run_test(
         Q_H: int = H,
         Q_S: int = S,
         Q_D: int = D,
-        KV_B: Optional[int] = None,
-        KV_H: Optional[int] = None,
-        KV_S: Optional[int] = None,
-        V_D: Optional[int] = None,
-        block_mask: Optional[BlockMask] = None,
+        KV_B: int | None = None,
+        KV_H: int | None = None,
+        KV_S: int | None = None,
+        V_D: int | None = None,
+        block_mask: BlockMask | None = None,
     ):
         requires_grad = device in DEVICE_SUPPORTS_BACKWARDS
         if KV_B is None:
@@ -599,9 +619,12 @@ def run_test(
         ref_out = sdpa_partial(q_ref, k_ref, v_ref)
         compiled_out = compiled_sdpa(q, k, v)
 
-        assert isinstance(golden_out, torch.Tensor)
-        assert isinstance(ref_out, torch.Tensor)
-        assert isinstance(compiled_out, torch.Tensor)
+        if not isinstance(golden_out, torch.Tensor):
+            raise AssertionError(f"Expected torch.Tensor, got {type(golden_out)}")
+        if not isinstance(ref_out, torch.Tensor):
+            raise AssertionError(f"Expected torch.Tensor, got {type(ref_out)}")
+        if not isinstance(compiled_out, torch.Tensor):
+            raise AssertionError(f"Expected torch.Tensor, got {type(compiled_out)}")
 
         if not requires_grad:
             self._check_out(
@@ -636,7 +659,7 @@ def run_test(
 
     def preprocess_paged_attention(
         self,
-        score_mod: Optional[Callable],
+        score_mod: Callable | None,
         q: Tensor,
         k: Tensor,
         v: Tensor,
@@ -645,7 +668,8 @@ def preprocess_paged_attention(
         device: str,
         page_size: int = 128,
     ) -> tuple[Tensor, Tensor, BlockMask, _score_mod_signature]:
-        assert block_mask is not None, "Must provide block_mask"
+        if block_mask is None:
+            raise AssertionError("Must provide block_mask")
         Q_B, Q_H, Q_S, _ = q.shape
         KV_B, KV_H, KV_S, QK_D = k.shape
         _, _, _, V_D = v.shape
@@ -726,14 +750,14 @@ def make_reserve_tensor(fractions):
 
     def run_paged_attention(
         self,
-        score_mod: Optional[Callable],
+        score_mod: Callable | None,
         q: Tensor,
         k: Tensor,
         v: Tensor,
         dtype: torch.dtype,
         device: str,
-        block_mask: Optional[BlockMask] = None,
-        kernel_options: Optional[dict] = None,
+        block_mask: BlockMask | None = None,
+        kernel_options: dict | None = None,
     ) -> tuple[Tensor, Tensor]:
         B, Q_H, Q_S, KV_H, KV_S = (
             q.shape[0],
@@ -788,7 +812,7 @@ def run_paged_attention(
 
     def run_test_with_paged_attention(
         self,
-        score_mod: Optional[Callable],
+        score_mod: Callable | None,
         dtype: torch.dtype,
         device,
         Q_B: int = B,
@@ -799,9 +823,12 @@ def run_test_with_paged_attention(
         KV_H: int = H,
         KV_S: int = S,
         V_D: int = D,
-        block_mask: Optional[BlockMask] = None,
+        block_mask: BlockMask | None = None,
     ):
-        assert Q_H % KV_H == 0
+        if Q_H % KV_H != 0:
+            raise AssertionError(
+                f"Expected Q_H % KV_H == 0, got {Q_H} % {KV_H} = {Q_H % KV_H}"
+            )
         if device == "cpu" and dtype is torch.float16:
             dtype = torch.float32
 
@@ -1363,7 +1390,7 @@ def test_builtin_score_mods_different_block_size(
         device,
         dtype: torch.dtype,
         score_mod: Callable,
-        BLOCK_SIZE: Union[int, tuple[int, int]],
+        BLOCK_SIZE: int | tuple[int, int],
     ):
         block_mask = create_block_mask(
             noop_mask, B, H, S, S, BLOCK_SIZE=BLOCK_SIZE, device=device
@@ -1389,10 +1416,16 @@ def test_kv_batch_broadcast(
         score_mod: Callable,
     ):
         Hq, Hkv = head_dims
-        assert Hq % Hkv == 0
+        if Hq % Hkv != 0:
+            raise AssertionError(
+                f"Expected Hq % Hkv == 0, got {Hq} % {Hkv} = {Hq % Hkv}"
+            )
 
         Bq, Bkv = batch_dims
-        assert Bq > 1 and Bkv == 1
+        if not (Bq > 1 and Bkv == 1):
+            raise AssertionError(
+                f"Expected Bq > 1 and Bkv == 1, got Bq={Bq}, Bkv={Bkv}"
+            )
 
         block_mask = create_block_mask(noop_mask, Bq, 1, S, S, device=device)
 
@@ -1514,10 +1547,16 @@ def test_kv_batch_broadcast_causal_mask(
         score_mod: Callable,
     ):
         Hq, Hkv = head_dims
-        assert Hq % Hkv == 0
+        if Hq % Hkv != 0:
+            raise AssertionError(
+                f"Expected Hq % Hkv == 0, got {Hq} % {Hkv} = {Hq % Hkv}"
+            )
 
         Bq, Bkv = batch_dims
-        assert Bq > 1 and Bkv == 1
+        if not (Bq > 1 and Bkv == 1):
+            raise AssertionError(
+                f"Expected Bq > 1 and Bkv == 1, got Bq={Bq}, Bkv={Bkv}"
+            )
 
         def mask_mod(b, h, q, kv):
             return q >= kv
@@ -1534,7 +1573,6 @@ def mask_mod(b, h, q, kv):
     @dtypesIfCUDA(*device_configs["cuda"].dtypes_fast)
     @dtypesIfXPU(*device_configs["xpu"].dtypes_fast)
     @common_utils.parametrize("score_mod", test_score_mods)
-    @skip_on_rocm  # TODO: NaNs on ROCM
     def test_GQA(self, device, dtype: torch.dtype, score_mod: Callable):
         inputs = (
             score_mod,
@@ -1587,8 +1625,13 @@ def test_strided_inputs(self, device, dtype: torch.dtype, q_s, k_s, v_s, do_s):
         def coerce_to_strides(val, shape, strides):
             strides, offset = strides
             val_max = [x * (y - 1) for x, y in zip(strides, shape)]
-            assert sum(val_max) + offset < B * H * S * D * 2
-            assert strides[-1] == 1
+            if sum(val_max) + offset >= B * H * S * D * 2:
+                raise AssertionError(
+                    f"Expected sum(val_max) + offset < B * H * S * D * 2, "
+                    f"got {sum(val_max) + offset} >= {B * H * S * D * 2}"
+                )
+            if strides[-1] != 1:
+                raise AssertionError(f"Expected strides[-1] == 1, got {strides[-1]}")
             return torch.as_strided(val, shape, strides, offset).requires_grad_(
                 requires_grad
             )
@@ -1955,7 +1998,8 @@ def test_natten_2d(self, device, dtype):
         H = 32
         W = S // H
         WINDOW = 3
-        assert W * H == S
+        if W * H != S:
+            raise AssertionError(f"Expected W * H == S, got {W * H} != {S}")
 
         def get_x_y(idx):
             # This should be a floor divide, but we don't support that properly
@@ -2140,6 +2184,109 @@ def score_mod(score, batch, head, token_q, token_kv):
         for head_count in head_counts:
             run_with_head_count(compiled_fa, head_count)
 
+    @supported_platform
+    @skip_on_cpu
+    @dtypes(torch.bfloat16)
+    @dtypesIfCUDA(torch.bfloat16)
+    def test_compacted_block_mask_matches_reference_after_recompile(
+        self, device, dtype
+    ):
+        block_size = 128
+        seq_len = 512
+        head_dim = 64
+        batch_size = 1
+
+        def mask_fn(b, h, q_idx, kv_idx, document_id):
+            causal = q_idx >= kv_idx
+            same_document = document_id[b, q_idx] == document_id[b, kv_idx]
+            return causal | same_document
+
+        def make_block_mask(document_id, compact):
+            ref = create_block_mask(
+                functools.partial(mask_fn, document_id=document_id),
+                B=batch_size,
+                H=None,
+                Q_LEN=seq_len,
+                KV_LEN=seq_len,
+                device=device,
+                BLOCK_SIZE=block_size,
+            )
+            if not compact:
+                return ref
+
+            nums = [
+                ref.kv_num_blocks,
+                ref.full_kv_num_blocks,
+                ref.q_num_blocks,
+                ref.full_q_num_blocks,
+            ]
+            indices = [
+                ref.kv_indices,
+                ref.full_kv_indices,
+                ref.q_indices,
+                ref.full_q_indices,
+            ]
+            compact_indices = [
+                index[..., : max(int(num.max()), 1)].contiguous()
+                for num, index in zip(nums, indices, strict=True)
+            ]
+            return BlockMask(
+                seq_lengths=ref.seq_lengths,
+                kv_num_blocks=ref.kv_num_blocks,
+                kv_indices=compact_indices[0],
+                full_kv_num_blocks=ref.full_kv_num_blocks,
+                full_kv_indices=compact_indices[1],
+                q_num_blocks=ref.q_num_blocks,
+                q_indices=compact_indices[2],
+                full_q_num_blocks=ref.full_q_num_blocks,
+                full_q_indices=compact_indices[3],
+                BLOCK_SIZE=ref.BLOCK_SIZE,
+                mask_mod=ref.mask_mod,
+            )
+
+        document_id_1 = torch.arange(seq_len, dtype=torch.int64, device=device).repeat(
+            batch_size, 1
+        )
+        document_id_1[:, :400] = 0
+        document_id_2 = torch.arange(seq_len, dtype=torch.int64, device=device).repeat(
+            batch_size, 1
+        )
+
+        query = torch.randn(
+            (batch_size, 1, seq_len, head_dim), device=device, dtype=dtype
+        )
+        key = torch.randn_like(query)
+        value = torch.randn_like(query)
+
+        def run(mask, attention):
+            q, k, v = [
+                x.detach().clone().requires_grad_(True) for x in (query, key, value)
+            ]
+            out = attention(q, k, v, block_mask=mask)
+            return out, torch.autograd.grad(out.sum(), (q, k, v))
+
+        ref_mask_1 = make_block_mask(document_id_1, compact=False)
+        ref_mask_2 = make_block_mask(document_id_2, compact=False)
+        with warnings.catch_warnings():
+            warnings.simplefilter("ignore")
+            ref_out_1, ref_grads_1 = run(ref_mask_1, flex_attention)
+            ref_out_2, ref_grads_2 = run(ref_mask_2, flex_attention)
+
+        compact_mask_1 = make_block_mask(document_id_1, compact=True)
+        compact_mask_2 = make_block_mask(document_id_2, compact=True)
+
+        torch._dynamo.reset()
+        compiled_fa = torch.compile(flex_attention, fullgraph=True, dynamic=False)
+        out_1, grads_1 = run(compact_mask_1, compiled_fa)
+        out_2, grads_2 = run(compact_mask_2, compiled_fa)
+
+        torch.testing.assert_close(ref_out_1, out_1, atol=2e-2, rtol=2e-2)
+        torch.testing.assert_close(ref_out_2, out_2, atol=2e-2, rtol=2e-2)
+        for ref_grad, grad in zip(ref_grads_1, grads_1):
+            torch.testing.assert_close(ref_grad, grad, atol=3e-2, rtol=3e-2)
+        for ref_grad, grad in zip(ref_grads_2, grads_2):
+            torch.testing.assert_close(ref_grad, grad, atol=3e-2, rtol=3e-2)
+
     @supported_platform
     @dtypes(*device_configs["cpu"].dtypes_fast)
     @dtypesIfCUDA(*device_configs["cuda"].dtypes_fast)
@@ -2507,7 +2654,6 @@ def f(q, k1, k2, v1, v2):
 
     @supported_platform
     @skip_on_cpu
-    @skip_on_rocm  # TODO: Investigate
     def test_multiple_mask_calls(self, device):
         make_tensor = functools.partial(
             torch.randn,
@@ -3640,11 +3786,6 @@ def test_strided_backwards(self, device):
     def test_flex_attention_stride_ordering(self, device, mode, permute_order, shape):
         from torch._inductor.ir import get_stride_order
 
-        if torch.version.hip and mode == "paged_attention":
-            raise self.skipTest(
-                "TODO: figure out why mode_paged_attention_permute_order3_shape0 on MI200 caused mem fault"
-            )
-
         dtype = torch.float32
         # Setup
         requires_grad = device in DEVICE_SUPPORTS_BACKWARDS
@@ -4365,8 +4506,10 @@ def test_small_q_kv_len(self, device):
             query, key, value, return_lse=True, kernel_options=kernel_options
         )
 
-        assert torch.equal(out_eager, out_compiled)
-        assert torch.equal(lse_eager, lse_compiled)
+        if not torch.equal(out_eager, out_compiled):
+            raise AssertionError("out_eager and out_compiled are not equal")
+        if not torch.equal(lse_eager, lse_compiled):
+            raise AssertionError("lse_eager and lse_compiled are not equal")
 
         grads_eager = torch.autograd.grad(out_eager.sum(), (query, key, value))
         grads_compile = torch.autograd.grad(out_compiled.sum(), (query, key, value))
@@ -4426,9 +4569,10 @@ def forward(self, x, input_lengths):
             l = torch.randint(0, T, (B,), device=device)
             model(x, l)
 
-        assert counter.frame_count == 1, (
-            f"Expected 1 graph, but got {counter.frame_count} graphs"
-        )
+        if counter.frame_count != 1:
+            raise AssertionError(
+                f"Expected 1 graph, but got {counter.frame_count} graphs"
+            )
 
     @supported_platform
     @skip_on_cpu
@@ -4775,6 +4919,195 @@ def forward(self, arg0_1: "i32[]", arg1_1: "i32[]", arg2_1: "i32[]", arg3_1: "i3
             ),
         )
 
+    @supported_platform
+    @skip_on_cpu
+    def test_direct_backward_preserves_explicit_buffers(self, device):
+        mask_buffer = torch.full((), 128, device=device, dtype=torch.int32)
+
+        def score_mod(score, b, h, m, n):
+            return score
+
+        def mask_mod(b, h, m, n):
+            return m + mask_buffer >= n
+
+        block_mask = create_block_mask(
+            mask_mod,
+            B=2,
+            H=2,
+            Q_LEN=128,
+            KV_LEN=128,
+            device=device,
+        )
+        scale = 1.0 / 16**0.5
+
+        dtype = torch.float32
+        q = torch.randn(
+            (2, 2, 128, 16),
+            dtype=dtype,
+            device=device,
+            requires_grad=True,
+        )
+        k = torch.randn(
+            (2, 2, 128, 16),
+            dtype=dtype,
+            device=device,
+            requires_grad=True,
+        )
+        v = torch.randn(
+            (2, 2, 128, 16),
+            dtype=dtype,
+            device=device,
+            requires_grad=True,
+        )
+        q_ref, k_ref, v_ref = query_key_value_clones(q, k, v)
+        q_gold, k_gold, v_gold = query_key_value_clones(q, k, v, torch.float64)
+
+        sdpa_partial = create_attention(score_mod, block_mask)
+        golden_out = sdpa_partial(q_gold, k_gold, v_gold)
+        ref_out = sdpa_partial(q_ref, k_ref, v_ref)
+
+        backward_grad = torch.randn((2, 2, 128, 16), dtype=dtype, device=device)
+        golden_out.backward(backward_grad.to(torch.float64))
+        ref_out.backward(backward_grad)
+
+        out, logsumexp = flex_attention_fwd(
+            q,
+            k,
+            v,
+            score_mod,
+            block_mask,
+            scale,
+        )
+
+        @torch.compile(fullgraph=True)
+        def compiled_bw(query, key, value, fwd_out, lse, grad_out):
+            return torch.ops.higher_order.flex_attention_backward(
+                query,
+                key,
+                value,
+                fwd_out,
+                lse,
+                grad_out,
+                None,
+                score_mod,
+                None,
+                block_mask.as_tuple(),
+                scale,
+                {},
+                (),
+                (),
+            )
+
+        with torch.no_grad():
+            dq, dk, dv, _ = compiled_bw(
+                q,
+                k,
+                v,
+                out.detach(),
+                logsumexp.detach(),
+                backward_grad,
+            )
+
+        fudge_factor = 10.0
+        self._check_equal(q_gold.grad, q_ref.grad, dq, fudge_factor, "Grad_Query")
+        self._check_equal(k_gold.grad, k_ref.grad, dk, fudge_factor, "Grad_Key")
+        self._check_equal(v_gold.grad, v_ref.grad, dv, fudge_factor, "Grad_Value")
+
+    @supported_platform
+    @skip_on_cpu
+    def test_direct_backward_supports_symint_score_mod_buffers(self, device):
+        def score_mod(score, b, h, m, n, head_dim):
+            return score
+
+        def mask_mod(b, h, m, n):
+            return m >= n
+
+        block_mask = create_block_mask(
+            mask_mod,
+            B=2,
+            H=2,
+            Q_LEN=128,
+            KV_LEN=128,
+            device=device,
+        )
+        scale = 1.0 / 16**0.5
+        dtype = torch.float32
+        q = torch.randn((2, 2, 128, 16), dtype=dtype, device=device)
+        k = torch.randn((2, 2, 128, 16), dtype=dtype, device=device)
+        v = torch.randn((2, 2, 128, 16), dtype=dtype, device=device)
+        backward_grad = torch.randn((2, 2, 128, 16), dtype=dtype, device=device)
+
+        out, logsumexp = flex_attention_fwd(
+            q,
+            k,
+            v,
+            _identity,
+            block_mask,
+            scale,
+        )
+        static_head_dim = q.shape[-1]
+
+        @torch.compile(backend="aot_eager", fullgraph=True)
+        def compiled_literal_bw(query, key, value, fwd_out, lse, grad_out):
+            return torch.ops.higher_order.flex_attention_backward(
+                query,
+                key,
+                value,
+                fwd_out,
+                lse,
+                grad_out,
+                None,
+                score_mod,
+                None,
+                block_mask.as_tuple(),
+                scale,
+                {},
+                (static_head_dim,),
+                (),
+            )
+
+        @torch.compile(backend="aot_eager", fullgraph=True, dynamic=True)
+        def compiled_bw(query, key, value, fwd_out, lse, grad_out):
+            return torch.ops.higher_order.flex_attention_backward(
+                query,
+                key,
+                value,
+                fwd_out,
+                lse,
+                grad_out,
+                None,
+                score_mod,
+                None,
+                block_mask.as_tuple(),
+                scale,
+                {},
+                (query.shape[-1],),
+                (),
+            )
+
+        with torch.no_grad():
+            ref_dq, ref_dk, ref_dv, ref_buffer_grads = compiled_literal_bw(
+                q,
+                k,
+                v,
+                out.detach(),
+                logsumexp.detach(),
+                backward_grad,
+            )
+            compiled_dq, compiled_dk, compiled_dv, compiled_buffer_grads = compiled_bw(
+                q,
+                k,
+                v,
+                out.detach(),
+                logsumexp.detach(),
+                backward_grad,
+            )
+
+        torch.testing.assert_close(compiled_dq, ref_dq)
+        torch.testing.assert_close(compiled_dk, ref_dk)
+        torch.testing.assert_close(compiled_dv, ref_dv)
+        self.assertEqual(compiled_buffer_grads, ref_buffer_grads)
+
     @supported_platform
     def test_tensor_subclass_dispatch_order(self, device):
         """Test that tensor subclasses get proper dispatch priority over modes.
@@ -4789,7 +5122,8 @@ def test_tensor_subclass_dispatch_order(self, device):
         class AsStridedErrorTensor(torch.Tensor):
             @staticmethod
             def __new__(cls, elem):
-                assert isinstance(elem, torch.Tensor)
+                if not isinstance(elem, torch.Tensor):
+                    raise AssertionError(f"Expected torch.Tensor, got {type(elem)}")
                 return torch.Tensor._make_wrapper_subclass(
                     cls,
                     elem.shape,
@@ -4812,7 +5146,8 @@ def __tensor_flatten__(self):
 
             @staticmethod
             def __tensor_unflatten__(inner_tensors, meta, outer_size, outer_stride):
-                assert meta is None
+                if meta is not None:
+                    raise AssertionError(f"Expected meta to be None, got {meta}")
                 elem = inner_tensors["elem"]
                 return AsStridedErrorTensor(elem)
 
@@ -4980,7 +5315,8 @@ def __init__(self, N: int, R: int, H: int = 4, hidden: int = 32):
             def _init_tables(self, N: int, R: int) -> None:
                 P = N - R
                 S = int(P**0.5)
-                assert S * S == P
+                if S * S != P:
+                    raise AssertionError(f"Expected S * S == P, got {S * S} != {P}")
                 rng = torch.arange(-(S - 1), S, dtype=torch.float32)
                 dY, dX = torch.meshgrid(rng, rng, indexing="ij")
                 rel = torch.stack(
@@ -5323,21 +5659,21 @@ def make_tensor():
             *positional_args,
             **keyword_args,
         )
-        assert kernel_code is not None, "Failed to retrieve compiled kernel code"
-        assert "num_consumer_groups" in kernel_code[0], (
-            "num_consumer_groups missing in kernel definition"
-        )
-        assert "num_buffers_warp_spec" in kernel_code[0], (
-            "num_buffers_warp_spec missing in kernel definition"
-        )
+        if kernel_code is None:
+            raise AssertionError("Failed to retrieve compiled kernel code")
+        if "num_consumer_groups" not in kernel_code[0]:
+            raise AssertionError("num_consumer_groups missing in kernel definition")
+        if "num_buffers_warp_spec" not in kernel_code[0]:
+            raise AssertionError("num_buffers_warp_spec missing in kernel definition")
 
         # Validate correctness
         C1 = flex_compiled(q, k, v)
         C2 = flex_attention(q, k, v)
 
-        assert torch.allclose(C1, C2, atol=1e-2, rtol=1e-2), (
-            "Warp specialized kernel result differs from reference"
-        )
+        if not torch.allclose(C1, C2, atol=1e-2, rtol=1e-2):
+            raise AssertionError(
+                "Warp specialized kernel result differs from reference"
+            )
 
     @supported_platform
     @skip_on_cpu
@@ -5504,6 +5840,91 @@ def causal_mask_with_offset(b, h, q_idx, kv_idx, offset):
 
         self.assertEqual(out.shape, (B, H, S, D))
 
+    @supported_platform
+    def test_callable_class_mask_mod(self, device):
+        # Test that callable class instances work as mask_mod with flex_attention
+
+        B, H, S, D = 1, 8, 64, 64
+
+        class CausalOrWindowMask:
+            __name__ = "causal_or_window"
+
+            def __init__(self, window_size):
+                self.window_size = window_size
+
+            def __call__(self, b, h, q_idx, kv_idx):
+                causal = q_idx >= kv_idx
+                in_window = (kv_idx - q_idx) < self.window_size
+                return causal | in_window
+
+            def __eq__(self, other):
+                if not isinstance(other, CausalOrWindowMask):
+                    return NotImplemented
+                return self.window_size == other.window_size
+
+            def __hash__(self):
+                return hash(self.window_size)
+
+        mask_mod = CausalOrWindowMask(window_size=64)
+
+        query = torch.randn(B, H, S, D, device=device, dtype=torch.float16)
+        key = torch.randn(B, H, S, D, device=device, dtype=torch.float16)
+        value = torch.randn(B, H, S, D, device=device, dtype=torch.float16)
+
+        @torch.compile(fullgraph=True)
+        def f(q, k, v):
+            block_mask = create_block_mask(mask_mod, B, H, S, S, device=q.device)
+            return flex_attention(q, k, v, block_mask=block_mask)
+
+        _ = f(query, key, value)
+
+    @supported_platform
+    @skip_on_cpu
+    def test_flex_attention_always_freezes_layout(self, device):
+        """Test that flex attention always freezes FlexibleLayout inputs.
+
+        When always_freeze_layout=True on flex attention templates,
+        get_stride_and_maybe_freeze_layout should freeze FlexibleLayout
+        immediately rather than using layout constraints.
+        """
+        from torch._inductor import ir
+        from torch._inductor.select_algorithm import TritonTemplateKernel
+
+        B, H, S, D = 2, 4, 128, 64
+        dtype = torch.float16
+
+        query = torch.randn(B, H, S, D, device=device, dtype=dtype)
+        key = torch.randn(B, H, S, D, device=device, dtype=dtype)
+        value = torch.randn(B, H, S, D, device=device, dtype=dtype)
+
+        flexible_layout_called = False
+        orig_stride_call = TritonTemplateKernel.get_stride_and_maybe_freeze_layout
+
+        def tracking_get_stride(self, node):
+            nonlocal flexible_layout_called
+            flexible_layout = isinstance(node.data.layout, ir.FlexibleLayout)
+            result = orig_stride_call(self, node)
+            if flexible_layout:
+                flexible_layout_called = True
+                if not isinstance(node.data.layout, ir.FixedLayout):
+                    raise AssertionError(
+                        f"Expected FixedLayout, got {type(node.data.layout)}"
+                    )
+            return result
+
+        with patch.object(
+            TritonTemplateKernel,
+            "get_stride_and_maybe_freeze_layout",
+            tracking_get_stride,
+        ):
+            compiled_flex = torch.compile(flex_attention, fullgraph=True)
+            compiled_flex(query, key, value)
+
+        self.assertTrue(
+            flexible_layout_called,
+            "get_stride_and_maybe_freeze_layout should be called with FlexibleLayout nodes",
+        )
+
 
 class TestBlockMask(InductorTestCase):
     def setUp(self):
@@ -5534,7 +5955,7 @@ def causal_mask(b, h, q, kv):
 
     @supported_platform
     @common_utils.parametrize("BLOCK_SIZE", [32, 64, 128, 256, (32, 64), (64, 32)])
-    def test_block_size_changes(self, device, BLOCK_SIZE: Union[int, tuple[int, int]]):
+    def test_block_size_changes(self, device, BLOCK_SIZE: int | tuple[int, int]):
         B, H, Q_LEN, KV_LEN = 4, 2, 2048, 2048
 
         if isinstance(BLOCK_SIZE, int):
@@ -5558,41 +5979,77 @@ def causal_mask(b, h, q, kv):
             return (q + (offset[b] * 128)) >= kv
 
         block_mask = create_block_mask(causal_mask, 4, 2, 512, 512, device=device)
-        assert block_mask.kv_num_blocks.shape == (4, 2, 4)
-        assert block_mask.kv_indices.shape == (4, 2, 4, 4)
+        if block_mask.kv_num_blocks.shape != (4, 2, 4):
+            raise AssertionError(
+                f"Expected shape (4, 2, 4), got {block_mask.kv_num_blocks.shape}"
+            )
+        if block_mask.kv_indices.shape != (4, 2, 4, 4):
+            raise AssertionError(
+                f"Expected shape (4, 2, 4, 4), got {block_mask.kv_indices.shape}"
+            )
 
         # Index on batch dimension
         new_block_mask = block_mask[0]
-        assert new_block_mask.kv_num_blocks.shape == (1, 2, 4)
-        assert new_block_mask.kv_indices.shape == (1, 2, 4, 4)
+        if new_block_mask.kv_num_blocks.shape != (1, 2, 4):
+            raise AssertionError(
+                f"Expected shape (1, 2, 4), got {new_block_mask.kv_num_blocks.shape}"
+            )
+        if new_block_mask.kv_indices.shape != (1, 2, 4, 4):
+            raise AssertionError(
+                f"Expected shape (1, 2, 4, 4), got {new_block_mask.kv_indices.shape}"
+            )
 
         # Index on batch and head dimension
         new_block_mask = block_mask[0, 1]
-        assert new_block_mask.kv_num_blocks.shape == (
+        if new_block_mask.kv_num_blocks.shape != (
             1,
             1,
             4,
-        )
-        assert new_block_mask.kv_indices.shape == (1, 1, 4, 4)
+        ):
+            raise AssertionError(
+                f"Expected shape (1, 1, 4), got {new_block_mask.kv_num_blocks.shape}"
+            )
+        if new_block_mask.kv_indices.shape != (1, 1, 4, 4):
+            raise AssertionError(
+                f"Expected shape (1, 1, 4, 4), got {new_block_mask.kv_indices.shape}"
+            )
 
         # Index on batch and head dimension with -1 semantics
         new_block_mask = block_mask[-1, -2]
-        assert new_block_mask.kv_num_blocks.shape == (
+        if new_block_mask.kv_num_blocks.shape != (
             1,
             1,
             4,
-        )
-        assert new_block_mask.kv_indices.shape == (1, 1, 4, 4)
+        ):
+            raise AssertionError(
+                f"Expected shape (1, 1, 4), got {new_block_mask.kv_num_blocks.shape}"
+            )
+        if new_block_mask.kv_indices.shape != (1, 1, 4, 4):
+            raise AssertionError(
+                f"Expected shape (1, 1, 4, 4), got {new_block_mask.kv_indices.shape}"
+            )
 
         # slicing on batch and head dimension
         new_block_mask = block_mask[0:2, 1:2]
-        assert new_block_mask.kv_num_blocks.shape == (2, 1, 4)
-        assert new_block_mask.kv_indices.shape == (2, 1, 4, 4)
+        if new_block_mask.kv_num_blocks.shape != (2, 1, 4):
+            raise AssertionError(
+                f"Expected shape (2, 1, 4), got {new_block_mask.kv_num_blocks.shape}"
+            )
+        if new_block_mask.kv_indices.shape != (2, 1, 4, 4):
+            raise AssertionError(
+                f"Expected shape (2, 1, 4, 4), got {new_block_mask.kv_indices.shape}"
+            )
 
         # slicing on batch, head, and query dimension
         new_block_mask = block_mask[0:2, 1:2, torch.tensor([1], dtype=torch.int32)]
-        assert new_block_mask.kv_num_blocks.shape == (2, 1, 1)
-        assert new_block_mask.kv_indices.shape == (2, 1, 1, 4)
+        if new_block_mask.kv_num_blocks.shape != (2, 1, 1):
+            raise AssertionError(
+                f"Expected shape (2, 1, 1), got {new_block_mask.kv_num_blocks.shape}"
+            )
+        if new_block_mask.kv_indices.shape != (2, 1, 1, 4):
+            raise AssertionError(
+                f"Expected shape (2, 1, 1, 4), got {new_block_mask.kv_indices.shape}"
+            )
 
         # slicing on batch, head, and query dimension
         q_index = torch.tensor([0], dtype=torch.int32)
@@ -5609,8 +6066,10 @@ def causal_mask(b, h, q, kv):
         )
 
         if block_mask.full_kv_num_blocks is not None:
-            assert new_block_mask.full_kv_num_blocks is not None
-            assert new_block_mask.full_kv_indices is not None
+            if new_block_mask.full_kv_num_blocks is None:
+                raise AssertionError("Expected full_kv_num_blocks to not be None")
+            if new_block_mask.full_kv_indices is None:
+                raise AssertionError("Expected full_kv_indices to not be None")
             torch.testing.assert_close(
                 new_block_mask.full_kv_num_blocks,
                 block_mask.full_kv_num_blocks[:, :, q_index],
@@ -5651,22 +6110,50 @@ def causal_mask(b, h, q, kv):
             return (q + (offset[b] * 128)) >= kv
 
         block_mask = create_block_mask(causal_mask, 1, 1, 512, 512, device=device)
-        assert block_mask.kv_indices.device.type == device.type
-        assert block_mask.kv_num_blocks.device.type == device.type
-        assert block_mask.q_indices.device.type == device.type
-        assert block_mask.q_num_blocks.device.type == device.type
+        if block_mask.kv_indices.device.type != device.type:
+            raise AssertionError(
+                f"Expected device type {device.type}, got {block_mask.kv_indices.device.type}"
+            )
+        if block_mask.kv_num_blocks.device.type != device.type:
+            raise AssertionError(
+                f"Expected device type {device.type}, got {block_mask.kv_num_blocks.device.type}"
+            )
+        if block_mask.q_indices.device.type != device.type:
+            raise AssertionError(
+                f"Expected device type {device.type}, got {block_mask.q_indices.device.type}"
+            )
+        if block_mask.q_num_blocks.device.type != device.type:
+            raise AssertionError(
+                f"Expected device type {device.type}, got {block_mask.q_num_blocks.device.type}"
+            )
 
         block_mask = block_mask.to("cpu")
-        assert block_mask.kv_indices.is_cpu
-        assert block_mask.kv_num_blocks.is_cpu
-        assert block_mask.q_indices.is_cpu
-        assert block_mask.q_num_blocks.is_cpu
+        if not block_mask.kv_indices.is_cpu:
+            raise AssertionError("Expected kv_indices to be on CPU")
+        if not block_mask.kv_num_blocks.is_cpu:
+            raise AssertionError("Expected kv_num_blocks to be on CPU")
+        if not block_mask.q_indices.is_cpu:
+            raise AssertionError("Expected q_indices to be on CPU")
+        if not block_mask.q_num_blocks.is_cpu:
+            raise AssertionError("Expected q_num_blocks to be on CPU")
 
         block_mask = block_mask.to(device)
-        assert block_mask.kv_indices.device.type == device.type
-        assert block_mask.kv_num_blocks.device.type == device.type
-        assert block_mask.q_indices.device.type == device.type
-        assert block_mask.q_num_blocks.device.type == device.type
+        if block_mask.kv_indices.device.type != device.type:
+            raise AssertionError(
+                f"Expected device type {device.type}, got {block_mask.kv_indices.device.type}"
+            )
+        if block_mask.kv_num_blocks.device.type != device.type:
+            raise AssertionError(
+                f"Expected device type {device.type}, got {block_mask.kv_num_blocks.device.type}"
+            )
+        if block_mask.q_indices.device.type != device.type:
+            raise AssertionError(
+                f"Expected device type {device.type}, got {block_mask.q_indices.device.type}"
+            )
+        if block_mask.q_num_blocks.device.type != device.type:
+            raise AssertionError(
+                f"Expected device type {device.type}, got {block_mask.q_num_blocks.device.type}"
+            )
 
     @supported_platform
     def test_compiling_create_block_mask(self, device):
@@ -5921,9 +6408,7 @@ def _offsets_to_doc_ids_tensor(offsets):
                 torch.arange(len(counts), device=device, dtype=torch.int32), counts
             )
 
-        def length_to_offsets(
-            lengths: list[int], device: Union[str, torch.device]
-        ) -> Tensor:
+        def length_to_offsets(lengths: list[int], device: str | torch.device) -> Tensor:
             offsets = [0]
             offsets.extend(lengths)
             offsets = torch.tensor(offsets, device=device, dtype=torch.int32)
@@ -6166,10 +6651,14 @@ def causal(
         )
         out.sum().backward()
 
-        assert out.isfinite().all().item()
-        assert q.grad.isfinite().all().item()
-        assert k.grad.isfinite().all().item()
-        assert v.grad.isfinite().all().item()
+        if not out.isfinite().all().item():
+            raise AssertionError("out contains non-finite values")
+        if not q.grad.isfinite().all().item():
+            raise AssertionError("q.grad contains non-finite values")
+        if not k.grad.isfinite().all().item():
+            raise AssertionError("k.grad contains non-finite values")
+        if not v.grad.isfinite().all().item():
+            raise AssertionError("v.grad contains non-finite values")
 
     @supported_platform
     @skip_on_cpu
@@ -6196,10 +6685,13 @@ def causal(
             q, k, v, score_mod=score_mod, block_mask=block_mask
         )
         out.sum().backward()
-        assert out.isfinite().all().item()
-        assert q.grad.isfinite().all().item()
+        if not out.isfinite().all().item():
+            raise AssertionError("out contains non-finite values")
+        if not q.grad.isfinite().all().item():
+            raise AssertionError("q.grad contains non-finite values")
         # assert k.grad.isfinite().all().item()
-        assert v.grad.isfinite().all().item()
+        if not v.grad.isfinite().all().item():
+            raise AssertionError("v.grad contains non-finite values")
 
     @supported_platform
     @skip_on_cpu
@@ -6226,10 +6718,13 @@ def causal(
             q, k, v, score_mod=score_mod, block_mask=block_mask
         )
         out.sum().backward()
-        assert out.isfinite().all().item()
-        assert q.grad.isfinite().all().item()
+        if not out.isfinite().all().item():
+            raise AssertionError("out contains non-finite values")
+        if not q.grad.isfinite().all().item():
+            raise AssertionError("q.grad contains non-finite values")
         # assert k.grad.isfinite().all().item()
-        assert v.grad.isfinite().all().item()
+        if not v.grad.isfinite().all().item():
+            raise AssertionError("v.grad contains non-finite values")
 
     @supported_platform
     @skip_on_cpu
@@ -6510,7 +7005,7 @@ def _check_equal(
         ref_out: torch.Tensor,
         compiled_out: torch.Tensor,
         fudge_factor: float,
-        tensor_name: Optional[str] = None,
+        tensor_name: str | None = None,
     ):
         compiled_error = (golden_out - compiled_out).abs().mean()
         ref_error = (golden_out - ref_out).abs().mean()
@@ -6932,7 +7427,7 @@ class Params:
     seq_length: int
     head_dim: int
     dtype: torch.dtype
-    config_str: Optional[str] = None
+    config_str: str | None = None
 
     def __str__(self):
         return f"batch:{self.batch_size}_head:{self.num_heads}_seq_len:{self.seq_length}_headdim:{self.head_dim}_dtype:{str(self.dtype).split('.')[-1]}"
@@ -7553,8 +8048,10 @@ def bias_mod(score, b, h, q_idx, kv_idx):
         loss = torch.nn.functional.mse_loss(attn_output, random_target)
         loss.backward()
 
-        assert bias.grad, "No gradient computed for bias"
-        assert torch.any(bias.grad != 0), "Gradient for bias is 0"
+        if bias.grad is None:
+            raise AssertionError("No gradient computed for bias")
+        if not torch.any(bias.grad != 0):
+            raise AssertionError("Gradient for bias is 0")
 
     @skip_on_cpu
     def test_backprop_error_case(self, device):
@@ -7582,8 +8079,10 @@ def score_mod(score, b, h, q_idx, kv_idx):
 
         _ = test(x, y).mean().backward()
 
-        assert x.grad.norm() > 0
-        assert y.grad.norm() > 0
+        if not (x.grad.norm() > 0):
+            raise AssertionError(f"Expected x.grad.norm() > 0, got {x.grad.norm()}")
+        if not (y.grad.norm() > 0):
+            raise AssertionError(f"Expected y.grad.norm() > 0, got {y.grad.norm()}")
 
     @skip_on_cpu
     @common_utils.parametrize(
@@ -7742,22 +8241,39 @@ def causal_mask(b, h, q_idx, kv_idx):
                 self.assertIsInstance(log_data, list)
                 self.assertEqual(len(log_data), 2)
 
-                keys_seen = [next(iter(entry.keys())) for entry in log_data]
-
-                expected_fwd_key = "('forward', 1, 2, 2, 128, 128, 64, 64)"
-                expected_bwd_key = "('backward', 1, 2, 2, 128, 128, 64, 64)"
-
-                self.assertIn(expected_fwd_key, keys_seen)
-                self.assertIn(expected_bwd_key, keys_seen)
+                # Check that we have both forward and backward entries
+                kernel_types_seen = [entry["kernel_type"] for entry in log_data]
+                self.assertIn("forward", kernel_types_seen)
+                self.assertIn("backward", kernel_types_seen)
+
+                # Expected values from the test inputs
+                expected_query_shape = "[1, 2, 128, 64]"
+                expected_key_shape = "[1, 2, 128, 64]"
+                expected_value_shape = "[1, 2, 128, 64]"
+                expected_B = 1
+                expected_Hq = 2
+                expected_Hkv = 2
+                expected_seq_len_q = 128
+                expected_seq_len_kv = 128
+                expected_qk_head_dim = 64
+                expected_v_head_dim = 64
 
                 for entry in log_data:
                     self.assertIsInstance(entry, dict)
-                    self.assertEqual(len(entry), 1)
+                    # New format has: query_shape, key_shape, value_shape, kernel_type, choices
+                    self.assertIn("kernel_type", entry)
+                    self.assertIn("choices", entry)
+                    self.assertIn("query_shape", entry)
+                    self.assertIn("key_shape", entry)
+                    self.assertIn("value_shape", entry)
 
-                    dims_key = next(iter(entry.keys()))
-                    choices = entry[dims_key]
+                    # Check shape values
+                    self.assertEqual(entry["query_shape"], expected_query_shape)
+                    self.assertEqual(entry["key_shape"], expected_key_shape)
+                    self.assertEqual(entry["value_shape"], expected_value_shape)
 
-                    kernel_type = eval(dims_key)[0]
+                    kernel_type = entry["kernel_type"]
+                    choices = entry["choices"]
 
                     self.assertIsInstance(choices, list)
                     self.assertGreater(len(choices), 0)
@@ -7770,6 +8286,25 @@ def causal_mask(b, h, q_idx, kv_idx):
                             self.assertIn("num_warps", choice)
                             self.assertIn("num_stages", choice)
 
+                            # Check numerical values in each choice
+                            self.assertIn("B", choice)
+                            self.assertIn("Hq", choice)
+                            self.assertIn("Hkv", choice)
+                            self.assertIn("seq_len_q", choice)
+                            self.assertIn("seq_len_kv", choice)
+                            self.assertIn("qk_head_dim", choice)
+                            self.assertIn("v_head_dim", choice)
+
+                            self.assertEqual(choice["B"], expected_B)
+                            self.assertEqual(choice["Hq"], expected_Hq)
+                            self.assertEqual(choice["Hkv"], expected_Hkv)
+                            self.assertEqual(choice["seq_len_q"], expected_seq_len_q)
+                            self.assertEqual(choice["seq_len_kv"], expected_seq_len_kv)
+                            self.assertEqual(
+                                choice["qk_head_dim"], expected_qk_head_dim
+                            )
+                            self.assertEqual(choice["v_head_dim"], expected_v_head_dim)
+
                             if kernel_type == "forward":
                                 self.assertIn("BLOCK_M", choice)
                                 self.assertIn("BLOCK_N", choice)
diff --git a/test/inductor/test_flex_decoding.py b/test/inductor/test_flex_decoding.py
index f25e56930d968..d172e4b565187 100644
--- a/test/inductor/test_flex_decoding.py
+++ b/test/inductor/test_flex_decoding.py
@@ -6,7 +6,6 @@
 import unittest
 from collections import namedtuple
 from collections.abc import Callable
-from typing import Optional, Union
 from unittest import expectedFailure
 from unittest.mock import patch
 
@@ -329,7 +328,7 @@ def _check_equal(
         ref_out: torch.Tensor,
         compiled_out: torch.Tensor,
         fudge_factor: float,
-        tensor_name: Optional[str] = None,
+        tensor_name: str | None = None,
     ):
         compiled_error = (golden_out - compiled_out).abs().mean()
         ref_error = (golden_out - ref_out).abs().mean()
@@ -372,7 +371,7 @@ def _check_out(
 
     def run_test(
         self,
-        score_mod: Optional[Callable] = None,
+        score_mod: Callable | None = None,
         dtype: torch.dtype = torch.float16,
         Q_B: int = B,
         Q_H: int = Hq,
@@ -382,14 +381,16 @@ def run_test(
         KV_H: int = Hkv,
         KV_S: int = S,
         V_D: int = D,
-        block_mask: Optional[BlockMask] = None,
+        block_mask: BlockMask | None = None,
         device="cuda",
         kernel_options=None,
     ):
-        assert score_mod is not None or block_mask is not None, (
-            "Must provide score_mod or block_mask"
-        )
-        assert Q_H % KV_H == 0
+        if score_mod is None and block_mask is None:
+            raise AssertionError("Must provide score_mod or block_mask")
+        if Q_H % KV_H != 0:
+            raise AssertionError(
+                f"Expected Q_H % KV_H == 0, got {Q_H} % {KV_H} = {Q_H % KV_H}"
+            )
         if device == "cpu" and dtype is torch.float16:
             dtype = torch.float32
 
@@ -443,7 +444,7 @@ def run_test(
     def run_test_with_call(
         self,
         sdpa_call: Callable,
-        golden_call: Optional[Callable] = None,
+        golden_call: Callable | None = None,
         dtype: torch.dtype = torch.float16,
         Q_B: int = B,
         Q_H: int = Hq,
@@ -495,7 +496,7 @@ def run_test_with_call(
 
     def preprocess_paged_attention(
         self,
-        score_mod: Optional[Callable],
+        score_mod: Callable | None,
         q: Tensor,
         k: Tensor,
         v: Tensor,
@@ -504,7 +505,8 @@ def preprocess_paged_attention(
         page_size: int = 128,
         device="cuda",
     ):
-        assert block_mask is not None, "Must provide block_mask"
+        if block_mask is None:
+            raise AssertionError("Must provide block_mask")
         if device == "cpu" and dtype is torch.float16:
             dtype = torch.float32
         Q_B, Q_H, Q_S, _ = q.shape
@@ -577,12 +579,12 @@ def preprocess_paged_attention(
 
     def run_paged_attention(
         self,
-        score_mod: Optional[Callable],
+        score_mod: Callable | None,
         q: Tensor,
         k: Tensor,
         v: Tensor,
         dtype: torch.dtype = torch.float16,
-        block_mask: Optional[BlockMask] = None,
+        block_mask: BlockMask | None = None,
         device="cuda",
     ):
         Q_B, Q_H, KV_H = q.shape[0], q.shape[1], k.shape[1]
@@ -629,7 +631,7 @@ def run_paged_attention(
 
     def run_test_with_paged_attention(
         self,
-        score_mod: Optional[Callable],
+        score_mod: Callable | None,
         dtype: torch.dtype = torch.float16,
         Q_B: int = B,
         Q_H: int = Hq,
@@ -639,10 +641,13 @@ def run_test_with_paged_attention(
         KV_H: int = Hkv,
         KV_S: int = S,
         V_D: int = D,
-        block_mask: Optional[BlockMask] = None,
+        block_mask: BlockMask | None = None,
         device="cuda",
     ):
-        assert Q_H % KV_H == 0
+        if Q_H % KV_H != 0:
+            raise AssertionError(
+                f"Expected Q_H % KV_H == 0, got {Q_H} % {KV_H} = {Q_H % KV_H}"
+            )
         if device == "cpu" and dtype is torch.float16:
             dtype = torch.float32
         q = torch.randn(
@@ -691,8 +696,8 @@ def run_test_with_paged_attention(
 
     def run_test_with_call_paged_attention(
         self,
-        score_mod: Optional[Callable],
-        mask_mod: Optional[Callable],
+        score_mod: Callable | None,
+        mask_mod: Callable | None,
         sdpa_mask: Tensor,
         dtype: torch.dtype = torch.float16,
         Q_B: int = B,
@@ -790,7 +795,10 @@ def test_builtin_score_mods(
         self, device, dtype: torch.dtype, score_mod: Callable, head_dims
     ):
         Hq, Hkv = head_dims
-        assert Hq % Hkv == 0
+        if Hq % Hkv != 0:
+            raise AssertionError(
+                f"Expected Hq % Hkv == 0, got {Hq} % {Hkv} = {Hq % Hkv}"
+            )
         self.run_test(score_mod, dtype, Q_H=Hq, KV_H=Hkv, device=device)
         self.run_test_with_paged_attention(
             score_mod, dtype, Q_H=Hq, KV_H=Hkv, device=device
@@ -810,7 +818,10 @@ def test_paged_attention_page_size(
         page_size: int,
     ):
         Hq, Hkv = head_dims
-        assert Hq % Hkv == 0
+        if Hq % Hkv != 0:
+            raise AssertionError(
+                f"Expected Hq % Hkv == 0, got {Hq} % {Hkv} = {Hq % Hkv}"
+            )
 
         def generate_causal_offset(offset: torch.Tensor):
             def causal_offset_mask(b, h, q_idx, kv_idx):
@@ -846,7 +857,7 @@ def test_builtin_score_mods_different_block_size(
         device,
         dtype: torch.dtype,
         score_mod: Callable,
-        BLOCK_SIZE: Union[int, tuple[int, int]],
+        BLOCK_SIZE: int | tuple[int, int],
     ):
         block_mask = create_block_mask(
             noop_mask, B, 1, 1, S, BLOCK_SIZE=BLOCK_SIZE, device=device
@@ -882,7 +893,10 @@ def test_tma_decoding(self, device, dtype: torch.dtype):
     @common_utils.parametrize("head_dims", test_Hq_Hkv)
     def test_strided_inputs(self, device, dtype: torch.dtype, k_s, v_s, head_dims):
         Hq, Hkv = head_dims
-        assert Hq % Hkv == 0
+        if Hq % Hkv != 0:
+            raise AssertionError(
+                f"Expected Hq % Hkv == 0, got {Hq} % {Hkv} = {Hq % Hkv}"
+            )
         q1 = torch.randn((B * Hq * D), dtype=dtype, device=device)
         k1 = torch.randn((B * Hkv * S * D * 4), dtype=dtype, device=device)
         v1 = torch.randn((B * Hkv * S * D * 4), dtype=dtype, device=device)
@@ -894,14 +908,22 @@ def test_strided_inputs(self, device, dtype: torch.dtype, k_s, v_s, head_dims):
 
         k_strides, k_offset = k_s(B, Hkv, S, D)
         k_max = [x * (y - 1) for x, y in zip(k_strides, k_shape)]
-        assert sum(k_max) + k_offset < B * Hkv * S * D * 4
-        assert k_strides[-1] == 1
+        if sum(k_max) + k_offset >= B * Hkv * S * D * 4:
+            raise AssertionError(
+                f"Expected sum(k_max) + k_offset < {B * Hkv * S * D * 4}, got {sum(k_max) + k_offset}"
+            )
+        if k_strides[-1] != 1:
+            raise AssertionError(f"Expected k_strides[-1] == 1, got {k_strides[-1]}")
         k = torch.as_strided(k1, k_shape, k_strides, k_offset)
 
         v_strides, v_offset = v_s(B, Hkv, S, D)
         v_max = [x * (y - 1) for x, y in zip(v_strides, v_shape)]
-        assert sum(v_max) + v_offset < B * Hkv * S * D * 4
-        assert v_strides[-1] == 1
+        if sum(v_max) + v_offset >= B * Hkv * S * D * 4:
+            raise AssertionError(
+                f"Expected sum(v_max) + v_offset < {B * Hkv * S * D * 4}, got {sum(v_max) + v_offset}"
+            )
+        if v_strides[-1] != 1:
+            raise AssertionError(f"Expected v_strides[-1] == 1, got {v_strides[-1]}")
         v = torch.as_strided(v1, v_shape, v_strides, v_offset)
 
         score_mod = _generate_alibi_bias(8)
@@ -941,10 +963,16 @@ def test_kv_batch_broadcast(
         score_mod: Callable,
     ):
         Hq, Hkv = head_dims
-        assert Hq % Hkv == 0
+        if Hq % Hkv != 0:
+            raise AssertionError(
+                f"Expected Hq % Hkv == 0, got {Hq} % {Hkv} = {Hq % Hkv}"
+            )
 
         Bq, Bkv = batch_dims
-        assert Bq > 1 and Bkv == 1
+        if not (Bq > 1 and Bkv == 1):
+            raise AssertionError(
+                f"Expected Bq > 1 and Bkv == 1, got Bq={Bq}, Bkv={Bkv}"
+            )
 
         block_mask = create_block_mask(noop_mask, Bq, 1, 1, S, device=device)
 
@@ -1143,7 +1171,10 @@ def test_head_dependent_mask_mod(
         self, device, dtype: torch.dtype, score_mod, head_dims
     ):
         Hq, Hkv = head_dims
-        assert Hq % Hkv == 0
+        if Hq % Hkv != 0:
+            raise AssertionError(
+                f"Expected Hq % Hkv == 0, got {Hq} % {Hkv} = {Hq % Hkv}"
+            )
 
         def head_attention_mod(kv_head_num):
             head_type = torch.tensor(
diff --git a/test/inductor/test_flex_flash.py b/test/inductor/test_flex_flash.py
index 5415378ac22b4..dafb42fbdde77 100644
--- a/test/inductor/test_flex_flash.py
+++ b/test/inductor/test_flex_flash.py
@@ -17,6 +17,7 @@
 from torch._inductor.utils import run_and_get_code, run_fw_bw_and_get_code
 from torch.nn.attention.flex_attention import (
     _DEFAULT_SPARSE_BLOCK_SIZE,
+    AuxRequest,
     create_block_mask,
     flex_attention,
 )
@@ -24,6 +25,9 @@
 from torch.testing._internal.common_cuda import (
     IS_SM90,
     PLATFORM_SUPPORTS_FP8,
+    SM120OrLater,
+    xfailIfSM120OrLater,
+    xfailIfSM12X,
     xfailIfSM90,
 )
 from torch.testing._internal.common_device_type import (
@@ -179,6 +183,21 @@ def score_with_buffer(score, _b, h, _q_idx, _kv_idx):
     return score_with_buffer
 
 
+def create_mask_mod_scalar_tensor(device="cuda"):
+    """mask_mod that captures a 0-dim (scalar) tensor.
+
+    Regression test: loading a 0-dim tensor in CuTeDSL codegen produces a
+    constant index 0, which was incorrectly passed as a bare Python int to
+    ssa_to_indexable (expects TensorSSA).  See #177813.
+    """
+    offset = torch.tensor(5, dtype=torch.int32, device=device)
+
+    def mask_with_scalar_tensor(_b, _h, q_idx, kv_idx):
+        return (q_idx + offset) >= kv_idx
+
+    return mask_with_scalar_tensor
+
+
 def create_mask_mod_buffer(num_heads=4, dtype=torch.float16, device="cuda"):
     mask_bias = torch.randn(num_heads, device=device, dtype=dtype) * 0.1
 
@@ -326,22 +345,32 @@ def flash_vs_triton(q, k, v, score_mod=None, block_mask=None, rtol=2, *, dynamic
         kernel_options={"BACKEND": "TRITON"},
     )
 
-    assert out_flash.shape == out_ref_fp32.shape == out_triton.shape
-    assert not torch.isnan(out_flash).any()
-    assert not torch.isnan(out_triton).any()
-    assert not torch.isnan(out_ref_fp32).any()
-    assert torch.isfinite(out_flash).all()
-    assert torch.isfinite(out_triton).all()
-    assert torch.isfinite(out_ref_fp32).all()
+    if not (out_flash.shape == out_ref_fp32.shape == out_triton.shape):
+        raise AssertionError(
+            f"Shape mismatch: flash={out_flash.shape}, ref={out_ref_fp32.shape}, triton={out_triton.shape}"
+        )
+    if torch.isnan(out_flash).any():
+        raise AssertionError("out_flash contains NaN")
+    if torch.isnan(out_triton).any():
+        raise AssertionError("out_triton contains NaN")
+    if torch.isnan(out_ref_fp32).any():
+        raise AssertionError("out_ref_fp32 contains NaN")
+    if not torch.isfinite(out_flash).all():
+        raise AssertionError("out_flash contains non-finite values")
+    if not torch.isfinite(out_triton).all():
+        raise AssertionError("out_triton contains non-finite values")
+    if not torch.isfinite(out_ref_fp32).all():
+        raise AssertionError("out_ref_fp32 contains non-finite values")
 
     fwd_atol = 2 * (out_ref_fp32 + 0.3 - 0.3 - out_ref_fp32).abs().max().item()
 
     triton_error = (out_triton - out_ref_fp32).abs().max().item()
     flash_error = (out_flash - out_ref_fp32).abs().max().item()
 
-    assert flash_error <= rtol * triton_error + fwd_atol, (
-        f"Flash error {flash_error:.2e} exceeds {rtol}x Triton error {triton_error:.2e} + {fwd_atol:.2e}"
-    )
+    if flash_error > rtol * triton_error + fwd_atol:
+        raise AssertionError(
+            f"Flash error {flash_error:.2e} exceeds {rtol}x Triton error {triton_error:.2e} + {fwd_atol:.2e}"
+        )
 
     needs_backward = any(
         isinstance(t, torch.Tensor) and t.requires_grad for t in (q, k, v)
@@ -361,17 +390,21 @@ def flash_vs_triton(q, k, v, score_mod=None, block_mask=None, rtol=2, *, dynamic
         for grad_flash, grad_triton, grad_ref, atol in zip(
             grads_flash, grads_triton, grads_ref, atol_pack
         ):
-            assert torch.isfinite(grad_flash).all()
-            assert torch.isfinite(grad_triton).all()
-            assert torch.isfinite(grad_ref).all()
+            if not torch.isfinite(grad_flash).all():
+                raise AssertionError("grad_flash contains non-finite values")
+            if not torch.isfinite(grad_triton).all():
+                raise AssertionError("grad_triton contains non-finite values")
+            if not torch.isfinite(grad_ref).all():
+                raise AssertionError("grad_ref contains non-finite values")
 
             triton_error = (grad_triton - grad_ref).abs().max().item()
             flash_error = (
                 (grad_flash - grad_ref.to(grad_flash.dtype)).abs().max().item()
             )
-            assert flash_error <= rtol * triton_error + atol, (
-                f"Flash error {flash_error:.2e} exceeds {rtol}x Triton error {triton_error:.2e} + {atol:.2e}"
-            )
+            if flash_error > rtol * triton_error + atol:
+                raise AssertionError(
+                    f"Flash error {flash_error:.2e} exceeds {rtol}x Triton error {triton_error:.2e} + {atol:.2e}"
+                )
 
     return out_flash, out_triton, out_ref_fp32
 
@@ -583,6 +616,10 @@ def mask_case_name(case: MaskModCase):
 
 
 MASK_MOD_CASES = [
+    MaskModCase(
+        "mask_mod_scalar_tensor",
+        lambda _dtype, device: create_mask_mod_scalar_tensor(device=device),
+    ),
     MaskModCase("block_mask_causal", lambda _dtype, _device: _causal_mask),
     MaskModCase(
         "block_mask_causal_score_times_two",
@@ -767,7 +804,10 @@ def mask_case_name(case: MaskModCase):
 @unittest.skipIf(
     not ensure_flash_available(), "Flash attention (CUTE) library is not available"
 )
+@xfailIfSM12X
 class TestFlexFlash(InductorTestCase):
+    # `FlashAttentionForwardSm120` does not have `apply_score_mod`.
+    @xfailIfSM120OrLater
     @decorateIf(
         unittest.expectedFailure,
         lambda params: params["case"].requires_grad and IS_SM90,
@@ -807,16 +847,34 @@ def test_flash_attention_backward_deterministic_score_mod_cases(
             device=device,
             requires_grad=case.requires_grad,
         )
-        with DeterministicGuard(True):
-            flash_vs_triton(
-                q,
-                k,
-                v,
-                score_mod=case.score_mod_factory(dtype, device)
-                if case.score_mod_factory
-                else None,
+        if SM120OrLater:
+            cls, pattern = (
+                AttributeError,
+                r"'FlashAttentionForwardSm120' object has no attribute 'apply_score_mod'",
             )
+            if "gqa_basic" in case.name or "mqa_basic" in case.name:
+                cls, pattern = ValueError, "Operation creation failed"
+            with self.assertRaisesRegex(cls, pattern), DeterministicGuard(True):
+                flash_vs_triton(
+                    q,
+                    k,
+                    v,
+                    score_mod=case.score_mod_factory(dtype, device)
+                    if case.score_mod_factory
+                    else None,
+                )
+        else:
+            with DeterministicGuard(True):
+                flash_vs_triton(
+                    q,
+                    k,
+                    v,
+                    score_mod=case.score_mod_factory(dtype, device)
+                    if case.score_mod_factory
+                    else None,
+                )
 
+    @xfailIfSM120OrLater
     @dtypes(torch.float16, torch.bfloat16)
     @parametrize("case", MASK_MOD_CASES, name_fn=mask_case_name)
     def test_flash_attention_mask_mod_cases(self, device, dtype, case):
@@ -883,8 +941,12 @@ def test_flash_attention_backward_deterministic_block_mask_raises(
 
         with DeterministicGuard(True):
             with self.assertRaisesRegex(
-                BackendCompilerFailed,
-                "Deterministic backward for flex_attention with block_mask using the FLASH backend",
+                (BackendCompilerFailed if not SM120OrLater else AssertionError),
+                (
+                    "Deterministic backward for flex_attention with block_mask using the FLASH backend"
+                    if not SM120OrLater
+                    else "Block sparsity not supported on SM 12.0"
+                ),
             ):
                 out = compiled_fn(
                     q,
@@ -900,6 +962,20 @@ def test_flash_attention_backward_deterministic_block_mask_raises(
                 )
                 out.sum().backward()
 
+    @decorateIf(
+        unittest.expectedFailure,
+        lambda params: (
+            SM120OrLater
+            and params["case"].name
+            in {
+                "backward_block_mask_causal",
+                "backward_block_mask_causal_rel_bias",
+                "backward_block_mask_causal_score_squared",
+                "backward_block_mask_causal_score_times_two",
+                "mask_mod_view_buffer",
+            }
+        ),
+    )
     @dtypes(torch.float16, torch.bfloat16)
     @parametrize("case", DETERMINISTIC_MASK_MOD_CASES, name_fn=mask_case_name)
     def test_flash_attention_backward_deterministic_warn_only_block_mask(
@@ -948,6 +1024,27 @@ def test_flash_attention_backward_deterministic_warn_only_block_mask(
             ):
                 out.sum().backward()
 
+    @decorateIf(
+        unittest.expectedFailure,
+        lambda params: (
+            SM120OrLater
+            and params["case"].name
+            in {
+                "gqa_block_mask_causal",
+                "gqa_block_mask_causal_per_head",
+                "backward_gqa_block_mask_causal",
+                "backward_gqa_block_mask_causal_per_head",
+                "mqa_block_mask_causal",
+                "mqa_block_mask_causal_per_head",
+                "backward_mqa_block_mask_causal",
+                "backward_mqa_block_mask_causal_per_head",
+            }
+        ),
+    )
+    @decorateIf(
+        unittest.expectedFailure,
+        lambda params: SM120OrLater and params["case"].name.endswith("_dim128"),
+    )
     @dtypes(torch.float16, torch.bfloat16)
     @parametrize("case", GQA_MQA_BLOCK_MASK_CASES, name_fn=mask_case_name)
     def test_flash_attention_gqa_mqa_block_mask_cases(self, device, dtype, case):
@@ -985,6 +1082,7 @@ def test_flash_attention_gqa_mqa_block_mask_cases(self, device, dtype, case):
             ),
         )
 
+    @xfailIfSM120OrLater
     @dtypes(torch.float16, torch.bfloat16)
     def test_flash_attention_kernel_called(self, device, dtype):
         q, k, v = create_test_tensors(dtype=dtype, device=device)
@@ -1073,6 +1171,7 @@ def score_mod_with_capture(score, b, h, q_idx, kv_idx):
                 kernel_options={"BACKEND": "FLASH"},
             ).sum().backward()
 
+    @xfailIfSM120OrLater
     @dtypes(torch.float16, torch.bfloat16)
     def test_flash_attention_backward_kernel_called(self, device, dtype):
         q, k, v = create_test_tensors(dim=128, dtype=dtype, device=device)
@@ -1098,6 +1197,7 @@ def run_for_profile():
             f"Flash attention backward kernel not found. Kernels: {prof_result['kernel_names']}",
         )
 
+    @xfailIfSM120OrLater
     @dtypes(torch.float16, torch.bfloat16)
     def test_flash_attention_backward_forwards_deterministic_flag(self, device, dtype):
         q, k, v = create_test_tensors(dim=128, dtype=dtype, device=device)
@@ -1116,6 +1216,7 @@ def run_for_code():
             "Expected deterministic flag to be wired through flash backward",
         )
 
+    @xfailIfSM120OrLater
     @dtypes(torch.float16, torch.bfloat16)
     def test_flash_attention_generates_cute_hash(self, device, dtype):
         q, k, v = create_test_tensors(dtype=dtype, device=device)
@@ -1136,6 +1237,7 @@ def test_flash_attention_generates_cute_hash(self, device, dtype):
             "Generated code should set __cute_hash__ on score_mod for fast hashing",
         )
 
+    @xfailIfSM120OrLater
     @dtypes(torch.float16, torch.bfloat16)
     def test_flash_attention_fused_qkv_reinterpret_view(self, device, dtype):
         B, M, H, D = 2, 256, 4, 64
@@ -1155,6 +1257,7 @@ def fn(x, weight):
         out = compiled_fn(x, weight)
         self.assertEqual(out.shape, (B, H, M, D))
 
+    @xfailIfSM120OrLater
     @dtypes(torch.float16, torch.bfloat16)
     def test_gqa_expand_stride_zero_backward(self, device, dtype):
         """Test GQA backward with expand()-created K/V tensors (stride=0).
@@ -1188,8 +1291,17 @@ def test_gqa_expand_stride_zero_backward(self, device, dtype):
             _causal_mask, batch_size, n_heads, seqlen, seqlen, device=device
         )
 
-        flash_vs_triton(q, k, v, block_mask=block_mask)
+        if SM120OrLater:
+            # note: see [SM120 forward tile selection failure]
+            with self.assertRaisesRegex(
+                ValueError,
+                r"Block sparsity requires sparse_block_size[1]=64 to match tile_n",
+            ):
+                flash_vs_triton(q, k, v, block_mask=block_mask)
+        else:
+            flash_vs_triton(q, k, v, block_mask=block_mask)
 
+    @xfailIfSM120OrLater
     @dtypes(torch.float16, torch.bfloat16)
     def test_flash_backend_raises_on_grad_logsumexp(self, device, dtype):
         from torch._dynamo.exc import BackendCompilerFailed
@@ -1208,6 +1320,46 @@ def test_flash_backend_raises_on_grad_logsumexp(self, device, dtype):
         ):
             loss.backward()
 
+    @dtypes(torch.float16, torch.bfloat16)
+    def test_flash_backend_raises_on_return_max_scores(self, device, dtype):
+        q, k, v = create_test_tensors(dtype=dtype, device=device)
+        with self.assertRaisesRegex(
+            NotImplementedError,
+            r"Returning max scores is not supported with BACKEND='FLASH'",
+        ):
+            flex_attention(
+                q,
+                k,
+                v,
+                return_aux=AuxRequest(max_scores=True),
+                kernel_options={"BACKEND": "FLASH"},
+            )
+
+    # 'FlashAttentionForwardSm120' object has no attribute 'apply_score_mod'
+    @xfailIfSM120OrLater
+    @decorateIf(
+        unittest.expectedFailure,
+        lambda params: IS_SM90,
+    )
+    @torch._inductor.config.patch(max_autotune=True)
+    @dtypes(torch.float16, torch.bfloat16)
+    @parametrize(
+        "score_mod_fn",
+        [_times_two, _rel_bias],
+        name_fn=lambda fn: fn.__name__,
+    )
+    def test_max_autotune_score_mod(self, device, dtype, score_mod_fn):
+        q, k, v = create_test_tensors(
+            batch_size=2,
+            num_heads=4,
+            seq_len=257,
+            dim=64,
+            dtype=dtype,
+            device=device,
+            requires_grad=True,
+        )
+        flash_vs_triton(q, k, v, score_mod=score_mod_fn)
+
 
 instantiate_device_type_tests(TestFlexFlash, globals(), only_for="cuda")
 
@@ -1246,10 +1398,13 @@ def _run_dynamic_test(
     def _flash_triton_dynamic(self, q, k, v, **kwargs):
         flash_vs_triton(q, k, v, dynamic=True, **kwargs)
 
+    # sm120: AttributeError: 'NoneType' object has no attribute '_trait'
+    @xfailIfSM120OrLater
     def test_dynamic_seq_len_no_score_mod(self):
         """Test dynamic sequence lengths without score_mod."""
         self._run_dynamic_test(seq_lens=[128, 256, 512])
 
+    @xfailIfSM120OrLater
     def test_dynamic_seq_len_inline_literal(self):
         """Test dynamic sequence lengths with inline literal score_mod."""
 
@@ -1258,6 +1413,7 @@ def score_mod(score, _b, _h, _q, _k):
 
         self._run_dynamic_test(seq_lens=[128, 256, 512], score_mod=score_mod)
 
+    @xfailIfSM120OrLater
     def test_dynamic_seq_len_captured_tensor_buffer(self):
         """Test dynamic sequence lengths with captured tensor buffer (ALiBi-style)."""
         num_heads = 4
@@ -1270,6 +1426,7 @@ def alibi_score_mod(score, b, h, q_idx, kv_idx):
 
         self._run_dynamic_test(seq_lens=[128, 256, 512], score_mod=alibi_score_mod)
 
+    @xfailIfSM120OrLater
     def test_dynamic_seq_len_with_block_mask(self):
         """Test dynamic sequence lengths with block mask."""
 
@@ -1282,6 +1439,7 @@ def block_mask_factory(seq_len):
             seq_lens=[128, 256, 512], block_mask_factory=block_mask_factory
         )
 
+    @xfailIfSM120OrLater
     def test_dynamic_batch_size(self):
         """Test dynamic batch sizes."""
         for batch_size in [1, 2, 4, 8]:
@@ -1290,11 +1448,14 @@ def test_dynamic_batch_size(self):
             )
             self._flash_triton_dynamic(q, k, v)
 
+    @xfailIfSM120OrLater
     @xfailIfSM90
     def test_dynamic_backward(self):
         """Test backward with dynamic sequence lengths."""
         self._run_dynamic_test(seq_lens=[128, 256, 512], requires_grad=True)
 
+    # 'FlashAttentionForwardSm120' object has no attribute 'apply_score_mod'
+    @xfailIfSM120OrLater
     @xfailIfSM90
     def test_dynamic_backward_with_score_mod(self):
         """Test backward with score_mod and dynamic sequence lengths."""
@@ -1306,6 +1467,7 @@ def score_mod(score, _b, _h, _q, _k):
             seq_lens=[128, 256, 512], score_mod=score_mod, requires_grad=True
         )
 
+    @xfailIfSM120OrLater
     def test_dynamic_backward_with_block_mask(self):
         """Test backward with block mask and dynamic sequence lengths."""
         major, _ = torch.cuda.get_device_capability()
@@ -1323,6 +1485,7 @@ def block_mask_factory(seq_len):
             requires_grad=True,
         )
 
+    @xfailIfSM120OrLater
     def test_dynamic_gqa(self):
         """Test GQA with dynamic sequence lengths."""
         q_heads, kv_heads = 8, 2
@@ -1336,6 +1499,7 @@ def test_dynamic_gqa(self):
             )
             self._flash_triton_dynamic(q, k, v, score_mod=None, block_mask=None)
 
+    @xfailIfSM120OrLater
     def test_dynamic_mqa(self):
         """Test MQA with dynamic sequence lengths."""
         q_heads, kv_heads = 8, 1
@@ -1349,6 +1513,7 @@ def test_dynamic_mqa(self):
             )
             self._flash_triton_dynamic(q, k, v)
 
+    @xfailIfSM120OrLater
     def test_dynamic_non_divisible_seq_len(self):
         """Test non-block-divisible sequence lengths with dynamic shapes."""
         for seq_len in [127, 255, 383, 511, 513]:
@@ -1357,6 +1522,7 @@ def test_dynamic_non_divisible_seq_len(self):
             )
             self._flash_triton_dynamic(q, k, v)
 
+    @xfailIfSM120OrLater
     def test_dynamic_asymmetric_qkv_lengths(self):
         """Test asymmetric Q and KV lengths with dynamic shapes."""
         test_cases = [(256, 512), (512, 256), (128, 1024)]
@@ -1400,6 +1566,7 @@ def score_mod(score, _b, _h, _q, _k):
                 q, k, v, score_mod=score_mod, kernel_options={"BACKEND": "FLASH"}
             )
 
+    @xfailIfSM120OrLater
     def test_captured_float_works_with_static(self):
         """Test that captured Python float works with dynamic=False."""
         val = 2.0  # Captured float
@@ -1415,6 +1582,7 @@ def score_mod(score, _b, _h, _q, _k):
         )
         self.assertEqual(out.shape, q.shape)
 
+    @xfailIfSM120OrLater
     def test_dynamic_mask_from_input_lengths_single_graph(self):
         """Dynamic mask creation driven by input lengths should stay single-graph."""
         counter = CompileCounterWithBackend("inductor")
@@ -1461,6 +1629,7 @@ def forward(self, x, input_lengths):
             counter.frame_count, 1, f"Expected 1 graph, got {counter.frame_count}"
         )
 
+    @xfailIfSM120OrLater
     def test_dynamic_free_symbol_mask_single_graph(self):
         """Free-symbol dense mask under dynamic=True should not recompile."""
         counter = CompileCounterWithBackend("inductor")
@@ -1497,6 +1666,7 @@ def run(q, k, v, block_mask):
             counter.frame_count, 1, f"Expected 1 graph, got {counter.frame_count}"
         )
 
+    @xfailIfSM120OrLater
     def test_dynamic_max_autotune_with_block_mask(self):
         """Dynamic=True with max-autotune should succeed for FLASH backend."""
         q, k, v = create_test_tensors(
@@ -1523,6 +1693,7 @@ def test_dynamic_max_autotune_with_block_mask(self):
         )
         self.assertEqual(out.shape, q.shape)
 
+    @xfailIfSM120OrLater
     @xfailIfSM90
     def test_dynamic_captured_buffer_varying_heads(self):
         """Dynamic head_count with captured tensor buffer under FLASH/TRITON parity."""
@@ -1655,6 +1826,8 @@ def test_hierarchical_indexer_rank_mismatch(self):
             indexer([b])
         self.assertIn("Rank mismatch", str(ctx.exception))
 
+    # 'FlashAttentionForwardSm120' object has no attribute 'apply_score_mod'
+    @xfailIfSM120OrLater
     @unittest.skipIf(not torch.cuda.is_available(), "CUDA not available")
     @unittest.skipIf(
         not ensure_flash_available(), "Flash attention (CUTE) library not available"
@@ -1698,6 +1871,8 @@ def score_mod_2d(score, b, h, q_idx, kv_idx):
             f"Expected '{expected_pattern}' in generated code.\nExcerpt:\n{code_str[:2000]}",
         )
 
+    # 'FlashAttentionForwardSm120' object has no attribute 'apply_score_mod'
+    @xfailIfSM120OrLater
     @unittest.skipIf(not torch.cuda.is_available(), "CUDA not available")
     @unittest.skipIf(
         not ensure_flash_available(), "Flash attention (CUTE) library not available"
@@ -1743,6 +1918,7 @@ def score_mod_3d(score, b, h, q_idx, kv_idx):
             f"Expected '{expected_pattern}' in generated code.\nExcerpt:\n{code_str[:2000]}",
         )
 
+    @xfailIfSM120OrLater
     @unittest.skipIf(not torch.cuda.is_available(), "CUDA not available")
     @unittest.skipIf(
         not ensure_flash_available(), "Flash attention (CUTE) library not available"
diff --git a/test/inductor/test_foreach.py b/test/inductor/test_foreach.py
index b812a44ce7722..4e85153a6a0f5 100644
--- a/test/inductor/test_foreach.py
+++ b/test/inductor/test_foreach.py
@@ -157,6 +157,7 @@ def recipaddmul_op(x, y, z):
     torch._foreach_abs,
     torch._foreach_sqrt,
     torch._foreach_rsqrt,
+    torch._foreach_clone,
     *foreach_map_un_ops_under_test,
 ]
 
@@ -1417,6 +1418,26 @@ def fn(s, t1, t2):
         code = " ".join(code)
         self.assertIn("tl.fma", code, "Expected FMA to be used in generated code")
 
+    @requires_gpu
+    def test_foreach_reorder_loops_with_nested_foreach_snodes(self):
+        def fn(x0, x1, y0, y1):
+            xs, ys = [x0, x1], [y0, y1]
+            torch._foreach_add_(xs, ys)
+            zs = torch._foreach_mul(xs, 2)
+            return [z.sum() for z in zs]
+
+        args = (
+            torch.randint(0, 10, (128, 128), device=GPU_TYPE),
+            torch.randint(0, 10, (128,), device=GPU_TYPE),
+            torch.randint(0, 10, (128, 128), device=GPU_TYPE),
+            torch.randint(0, 10, (128,), device=GPU_TYPE),
+        )
+        args_clone = tuple(a.clone() for a in args)
+
+        expected = fn(*args)
+        actual = torch.compile(fn)(*args_clone)
+        self.assertEqual(actual, expected)
+
 
 if __name__ == "__main__":
     from torch._inductor.test_case import run_tests
diff --git a/test/inductor/test_fp8.py b/test/inductor/test_fp8.py
index aab220511c374..cf87d75e5beed 100644
--- a/test/inductor/test_fp8.py
+++ b/test/inductor/test_fp8.py
@@ -2,7 +2,6 @@
 
 import functools
 import unittest
-from typing import Union
 
 import torch
 from torch import Tensor
@@ -19,20 +18,21 @@
     PLATFORM_SUPPORTS_FP8,
     PLATFORM_SUPPORTS_MX_GEMM,
     SM100OrLater,
-    SM90OrLater,
 )
 from torch.testing._internal.common_device_type import (
     instantiate_device_type_tests,
     onlyCUDA,
     onlyOn,
+    skipCUDAIf,
 )
 from torch.testing._internal.common_quantized import ceil_div, to_blocked
-from torch.testing._internal.common_utils import parametrize, xfailIf
+from torch.testing._internal.common_utils import parametrize, skipIfXpu, xfailIf
 from torch.testing._internal.inductor_utils import (
     _quantize_blockwise,
     _quantize_rowwise,
     _quantize_tensorwise,
     _to_fp8_saturated,
+    GPU_TYPE,
     HAS_CPU,
     HAS_CUDA_AND_TRITON,
     is_big_gpu,
@@ -43,7 +43,7 @@
 torch.set_float32_matmul_precision("high")
 
 
-f8_msg = "FP8 is only supported on H100+, SM 8.9 and MI300+ and XPU devices"
+f8_msg = "FP8 is only supported on H100+, SM 8.9 and MI300+, XPU and CPU devices"
 
 
 def _is_cuda_device(device) -> bool:
@@ -55,8 +55,8 @@ def _is_cuda_device(device) -> bool:
 
 
 def _fix_fp8_dtype_for_rocm(
-    dtype: Union[torch.dtype, list[torch.dtype], tuple[torch.dtype]], device
-) -> Union[torch.dtype, list[torch.dtype], tuple[torch.dtype]]:
+    dtype: torch.dtype | list[torch.dtype] | tuple[torch.dtype], device
+) -> torch.dtype | list[torch.dtype] | tuple[torch.dtype]:
     # This function is used to change FP8 data types
     # with MI300 supported FP8 types if device is GPU:
     #    e4m3fn -> e4m3fnuz
@@ -82,6 +82,7 @@ def _fix_fp8_dtype_for_rocm(
 
 
 class TestFP8Types(TestCase):
+    @skipCUDAIf(not PLATFORM_SUPPORTS_FP8, f8_msg)
     @parametrize("float8_dtype", (torch.float8_e4m3fn, torch.float8_e5m2))
     def test_xblock_for_small_numel(self, float8_dtype: torch.dtype, device: str):
         """
@@ -93,8 +94,6 @@ def test_xblock_for_small_numel(self, float8_dtype: torch.dtype, device: str):
         We should not pick a XBLOCK larger than xnumel
         """
         float8_dtype = _fix_fp8_dtype_for_rocm(float8_dtype, device=device)
-        if _is_cuda_device(device) and not PLATFORM_SUPPORTS_FP8:
-            raise unittest.SkipTest(f8_msg)
 
         def f(x):
             return x.to(dtype=float8_dtype)
@@ -104,10 +103,9 @@ def f(x):
         actual = torch.compile(f)(x)
         torch.testing.assert_close(expected.half(), actual.half(), rtol=1e-2, atol=1e-2)
 
+    @skipCUDAIf(not PLATFORM_SUPPORTS_FP8, f8_msg)
     @parametrize("dtype", (torch.float16, torch.bfloat16))
     def test_eager_fallback(self, dtype: torch.dtype, device: torch.device):
-        if _is_cuda_device(device) and not PLATFORM_SUPPORTS_FP8:
-            raise unittest.SkipTest(f8_msg)
         weight_shape = (32, 16)
 
         e4m3_type = torch.float8_e4m3fn
@@ -146,14 +144,13 @@ def fp8_matmul_unwrapped(x):
         x = torch.rand(*x_shape, device=device, dtype=dtype).to(e4m3_type)
         y_fp8 = compiled_fp8_matmul(x)  # noqa: F841
 
+    @skipCUDAIf(not PLATFORM_SUPPORTS_FP8, f8_msg)
     @parametrize("dtype", (torch.float16, torch.bfloat16, torch.float))
     @parametrize("shape", ("15,3,13", "4,2048,4096"))
     @parametrize("dst_types", [(torch.float8_e4m3fn, torch.float8_e5m2)])
     def test_valid_cast(
         self, dtype: torch.dtype, shape: str, dst_types: tuple, device: torch.device
     ):
-        if _is_cuda_device(device) and not PLATFORM_SUPPORTS_FP8:
-            raise unittest.SkipTest(f8_msg)
         dst_types = _fix_fp8_dtype_for_rocm(dst_types, device=device)
         e4m3, e5m2 = dst_types
 
@@ -171,6 +168,9 @@ def fp8_cast(x):
         torch.testing.assert_close(y0_fp8, x, rtol=5e-1, atol=5e-1)
         torch.testing.assert_close(y1_fp8, x, rtol=5e-1, atol=5e-1)
 
+    @skipIfXpu(
+        msg="Conversions between float8_e5m2 and float8_e4m3fn is not supported, torch-xpu-ops: 2888"
+    )
     @unittest.skipIf(not PLATFORM_SUPPORTS_FP8, f8_msg)
     def test_bad_cast(self, device):
         def fp8_cast(x, dtype):
@@ -203,6 +203,7 @@ def fp8_cast(x, dtype):
             out = compiled_fp8_cast(x, torch.float8_e4m3fn)
             self.assertEqual(out.dtype, torch.float8_e4m3fn)
 
+    @skipCUDAIf(not PLATFORM_SUPPORTS_FP8, f8_msg)
     @parametrize("src_dtype", (torch.float16, torch.bfloat16, torch.float))
     @parametrize("dst_dtype", (torch.float8_e4m3fn, torch.float8_e5m2))
     @parametrize("shape", ("16,16,16", "4,2048,4096"))
@@ -213,8 +214,6 @@ def test_to_fp8_saturated(
         shape: str,
         device: torch.device,
     ):
-        if _is_cuda_device(device) and not PLATFORM_SUPPORTS_FP8:
-            raise unittest.SkipTest(f8_msg)
         dst_dtype = _fix_fp8_dtype_for_rocm(dst_dtype, device=device)
 
         def fp8_saturated(x, dtype):
@@ -230,16 +229,13 @@ def fp8_saturated(x, dtype):
 
         torch.testing.assert_close(y_compiled.half(), y.half(), rtol=5e-1, atol=5e-1)
 
+    @skipCUDAIf(not PLATFORM_SUPPORTS_FP8, f8_msg)
     @parametrize("float8_dtype", (torch.float8_e4m3fn, torch.float8_e5m2))
     @parametrize("shape", ("1,1,15", "1,10,15", "1,10,512", "1,10,4096", "4,2048,4096"))
     def test_amax_fp8_quant(
         self, float8_dtype: torch.dtype, shape: str, device: torch.device
     ):
         float8_dtype = _fix_fp8_dtype_for_rocm(float8_dtype, device=device)
-        if _is_cuda_device(device) and not PLATFORM_SUPPORTS_FP8:
-            raise unittest.SkipTest(
-                "FP8 is only supported on H100+ and sm_89 and MI300+ devices"
-            )
         shape = [int(dim) for dim in shape.split(",")]
         batch_size, sequence_length, hidden_size = shape
 
@@ -260,13 +256,12 @@ def amax_fp8(x: Tensor, scale: Tensor):
 
         torch.testing.assert_close(y_compiled.half(), y.half(), rtol=1e-2, atol=1e-2)
 
+    @skipCUDAIf(not PLATFORM_SUPPORTS_FP8, f8_msg)
     @parametrize("float8_dtype", (torch.float8_e4m3fn, torch.float8_e5m2))
     @parametrize("shape", ("1,1,15", "1,10,15", "1,10,512", "1,10,4096", "4,2048,4096"))
     def test_amax_along_with_fp8_quant(
         self, float8_dtype: torch.dtype, shape: str, device: torch.device
     ):
-        if _is_cuda_device(device) and not PLATFORM_SUPPORTS_FP8:
-            raise unittest.SkipTest(f8_msg)
         float8_dtype = _fix_fp8_dtype_for_rocm(float8_dtype, device=device)
         shape = [int(dim) for dim in shape.split(",")]
         batch_size, sequence_length, hidden_size = shape
@@ -293,6 +288,7 @@ def amax_fp8(x: Tensor, scale: Tensor, amax_buffer: Tensor):
             amax_buffer_compiled, amax_buffer, rtol=1e-2, atol=1e-2
         )
 
+    @skipCUDAIf(not PLATFORM_SUPPORTS_FP8, f8_msg)
     @parametrize("float8_dtype", (torch.float8_e4m3fn, torch.float8_e5m2))
     @parametrize("amax_keep_dim", (True, False))
     @parametrize("shape", ("1,1,15", "1,10,15", "1,10,512", "1,10,4096", "4,2048,4096"))
@@ -303,10 +299,6 @@ def test_layernorm_fp8_quant(
         shape: str,
         device: torch.device,
     ):
-        if _is_cuda_device(device) and not PLATFORM_SUPPORTS_FP8:
-            raise unittest.SkipTest(
-                "FP8 is only supported on H100+ and sm_89 and MI300+ devices"
-            )
         float8_dtype = _fix_fp8_dtype_for_rocm(float8_dtype, device=device)
         shape = [int(dim) for dim in shape.split(",")]
         batch_size, sequence_length, hidden_size = shape
@@ -342,7 +334,7 @@ def ln_fp8(x: Tensor, scale: Tensor, amax_buffer: Tensor):
             amax_buffer_compiled, amax_buffer, rtol=1e-2, atol=1e-2
         )
 
-    @onlyCUDA
+    @onlyOn(["cuda", "xpu"])
     @unittest.skipIf(not PLATFORM_SUPPORTS_FP8, f8_msg)
     @parametrize("float8_dtype", (torch.float8_e4m3fn, torch.float8_e5m2))
     @parametrize("shape", ("4,2048,4096",))
@@ -353,7 +345,7 @@ def test_layernorm_fp8_quant_benchmark(
         shape: str,
         keepdim: bool,
     ):
-        float8_dtype = _fix_fp8_dtype_for_rocm(float8_dtype, device="cuda")
+        float8_dtype = _fix_fp8_dtype_for_rocm(float8_dtype, device=GPU_TYPE)
         shape = [int(dim) for dim in shape.split(",")]
         batch_size, sequence_length, hidden_size = shape
 
@@ -384,11 +376,11 @@ def ln_fp8(x: Tensor, scale: Tensor, amax_buffer: Tensor):
         compiled_ln_fp8_quant = torch.compile(ln_fp8, backend="inductor")
 
         x_shape = (batch_size, sequence_length, hidden_size)
-        x = torch.rand(*x_shape, device="cuda", dtype=torch.half)
-        scale = torch.tensor(0.2, device="cuda", dtype=torch.float)
+        x = torch.rand(*x_shape, device=GPU_TYPE, dtype=torch.half)
+        scale = torch.tensor(0.2, device=GPU_TYPE, dtype=torch.float)
 
-        amax_buffer_compiled = torch.zeros((1), device="cuda", dtype=torch.half)
-        amax_buffer = torch.zeros((1), device="cuda", dtype=torch.half)
+        amax_buffer_compiled = torch.zeros((1), device=GPU_TYPE, dtype=torch.half)
+        amax_buffer = torch.zeros((1), device=GPU_TYPE, dtype=torch.half)
         _ = compiled_ln_fp8_quant(x, scale, amax_buffer_compiled)
         compiled_latency = utils.do_bench_using_profiling(
             functools.partial(compiled_ln_fp8_quant, x, scale, amax_buffer_compiled)
@@ -407,10 +399,7 @@ def ln_fp8(x: Tensor, scale: Tensor, amax_buffer: Tensor):
             f"LN only Inductor: {ln_latency}ms."
         )
 
-    @unittest.skipIf(
-        not SM90OrLater or torch.version.hip, "PDL requires NVIDIA SM 9.0+"
-    )
-    @onlyOn(["cuda", "xpu"])
+    @onlyOn(["cuda", "xpu", "cpu"])
     def test_scaled_mm_pdl_handles_none_bias(self, device):
         dtype_float8 = _fix_fp8_dtype_for_rocm(torch.float8_e4m3fn, device)
         M, K, N = 32, 64, 32
@@ -451,7 +440,7 @@ class TestFP8Lowering(TestCase):
     @parametrize(
         "persistent_matmul", [False, True] if has_triton_tma_device() else [False]
     )
-    @onlyOn(["cuda", "xpu"])
+    @onlyOn(["cuda", "xpu", "cpu"])
     def test_tensorwise_scaling(
         self,
         dtype: torch.dtype,
@@ -476,8 +465,8 @@ def test_tensorwise_scaling(
         if has_bias:
             bias = torch.randn(N, device=device, dtype=torch.bfloat16)
 
-        # if "xpu" in device and use_fast_accum:
-        self.skipTest("XPU does not support use_fast_accum=True for now")
+        if "xpu" in device and use_fast_accum:
+            self.skipTest("XPU does not support use_fast_accum=True for now")
 
         # quantize weight (prior to inference)
         w_fp8, w_inverse_scale = _quantize_tensorwise(w, dtype_float8)
@@ -528,7 +517,7 @@ def linear(x_fp8, x_inverse_scale, w_t_fp8, w_inverse_scale, bias):
                 self.assertEqual(y_eager, y_compiled, rtol=1e-2, atol=0.05)
 
     @unittest.skipIf(not PLATFORM_SUPPORTS_FP8, f8_msg)
-    @onlyOn(["cuda", "xpu"])
+    @onlyOn(["cuda", "xpu", "cpu"])
     def test_scaled_mm_preserves_strides(self, device):
         """Test that scaled_mm preserves stride ordering through a custom pass."""
 
@@ -707,7 +696,7 @@ def linear(x_fp8, x_inverse_scale, w_t_fp8, w_inverse_scale, bias):
             torch.testing.assert_close(y_eager, y_compiled, rtol=1e-2, atol=0.05)
 
     @unittest.skipIf(not PLATFORM_SUPPORTS_FP8, f8_msg)
-    @onlyOn(["cuda", "xpu"])
+    @onlyOn(["cuda", "xpu", "cpu"])
     @parametrize("shape", ("16,16,32", "16,32,32", "1024,1024,512"))
     @parametrize("has_bias", (False, True))
     @parametrize("use_fast_accum", (False, True))
@@ -1000,7 +989,7 @@ def linear(x_fp8, x_inverse_scale, w_t_fp8, w_inverse_scale, bias):
             torch.testing.assert_close(y_eager, y_compiled, rtol=1e-2, atol=0.05)
 
     @unittest.skipIf(not PLATFORM_SUPPORTS_FP8, f8_msg)
-    @onlyOn(["cuda", "xpu"])
+    @onlyOn(["cuda", "xpu", "cpu"])
     @parametrize("M", (1, 3, 33, 257, 1024))
     @parametrize("K", (16, 32, 1024))
     @parametrize("N", (16, 2048))
@@ -1045,9 +1034,20 @@ def linear(x_fp8, x_inverse_scale, w_t_fp8, w_inverse_scale, bias):
             w_inverse_scale,
             bias,
         )
+
+        # On gfx120x, autotuned kernels have issues with small M
+        compile_mode = "max-autotune"
+        if (
+            torch.version.hip is not None
+            and M < 16
+            and torch.cuda.is_available()
+            and "gfx120" in torch.cuda.get_device_properties(0).gcnArchName
+        ):
+            compile_mode = "default"
+
         with config.patch({"triton.enable_persistent_tma_matmul": persistent_matmul}):
             linear_compiled = torch.compile(
-                linear, backend="inductor", mode="max-autotune"
+                linear, backend="inductor", mode=compile_mode
             )
             y_compiled = linear_compiled(
                 x_fp8,
@@ -1289,7 +1289,8 @@ def forward(
                 with config.patch({"loop_index_inversion_in_fusion": False}):
                     torch.compile(forward)(A, B)
 
-                assert len(input_values) == 2
+                if len(input_values) != 2:
+                    raise AssertionError
                 for i in range(4):
                     self.assertEqual(
                         input_values[0][i],
@@ -1298,7 +1299,7 @@ def forward(
                     )
 
     @unittest.skipIf(not PLATFORM_SUPPORTS_FP8, f8_msg)
-    @onlyOn(["cuda", "xpu"])
+    @onlyOn(["cuda", "xpu", "cpu"])
     @parametrize("M", (1, 3, 33, 257, 1024))
     @parametrize("K", (16, 32, 1024))
     @parametrize("N", (16, 2048))
@@ -1344,9 +1345,20 @@ def linear(x_fp8, x_inverse_scale, w_t_fp8, w_inverse_scale, bias):
             w_inverse_scale,
             bias,
         )
+
+        # On gfx120x, autotuned kernels have issues with small M
+        compile_mode = "max-autotune"
+        if (
+            torch.version.hip is not None
+            and M < 16
+            and torch.cuda.is_available()
+            and "gfx120" in torch.cuda.get_device_properties(0).gcnArchName
+        ):
+            compile_mode = "default"
+
         with config.patch({"triton.enable_persistent_tma_matmul": persistent_matmul}):
             linear_compiled = torch.compile(
-                linear, backend="inductor", mode="max-autotune"
+                linear, backend="inductor", mode=compile_mode
             )
             y_compiled = linear_compiled(
                 x_fp8,
@@ -1359,7 +1371,7 @@ def linear(x_fp8, x_inverse_scale, w_t_fp8, w_inverse_scale, bias):
         self.assertEqual(y_compiled.dtype, dtype)
         torch.testing.assert_close(y_eager, y_compiled, rtol=1e-2, atol=0.07)
 
-    @onlyOn(["cuda", "xpu"])
+    @onlyOn(["cuda", "xpu", "cpu"])
     @unittest.skipIf(not PLATFORM_SUPPORTS_MX_GEMM, "Not supported on non B200")
     def test_mx_fp8_max_autotune(self, device):
         M, K, N = 128, 32, 128
@@ -1443,7 +1455,7 @@ def linear(x, w_t_fp8, w_inverse_scale, bias):
         )
 
     @unittest.skipIf(not PLATFORM_SUPPORTS_FP8, f8_msg)
-    @onlyOn(["cuda", "xpu"])
+    @onlyOn(["cuda", "xpu", "cpu"])
     def test_unacceptable_scale_dims_rowwise_scaling(self, device):
         dtype: torch.dtype = torch.bfloat16
         dtype_float8 = torch.float8_e4m3fn
diff --git a/test/inductor/test_fused_attention.py b/test/inductor/test_fused_attention.py
index b627c9617db01..6ec6c45c71795 100644
--- a/test/inductor/test_fused_attention.py
+++ b/test/inductor/test_fused_attention.py
@@ -15,7 +15,14 @@
     PLATFORM_SUPPORTS_FUSED_ATTENTION,
     SM80OrLater,
 )
-from torch.testing._internal.common_utils import IS_LINUX, skipIfXpu, TEST_WITH_ROCM
+from torch.testing._internal.common_utils import (
+    IS_ARM64,
+    IS_CPU_CAPABILITY_SVE256,
+    IS_LINUX,
+    skipIfXpu,
+    TEST_WITH_ROCM,
+    xfailIf,
+)
 from torch.testing._internal.inductor_utils import (
     GPU_TYPE,
     HAS_CPU,
@@ -85,9 +92,9 @@ def _check_common(
                     x.requires_grad = training
 
             if not self.use_static_shapes:
-                torch._dynamo.mark_dynamic(args2[0], 0)
-                torch._dynamo.mark_dynamic(args2[1], 0)
-                torch._dynamo.mark_dynamic(args2[2], 0)
+                for i in range(min(3, len(args2))):
+                    if isinstance(args2[i], torch.Tensor):
+                        torch._dynamo.mark_dynamic(args2[i], 0)
 
             dropout_arg = [training] if has_dropout else []
             torch.manual_seed(1234)
@@ -137,7 +144,7 @@ def dot_prod_attention(
             )
 
         for dtype in [torch.float, torch.half]:
-            atol = 0.001
+            atol = 0.0015
             rtol = 1.3e-6 if dtype == torch.float else 0.7
             if self.device in ["cpu", "xpu"] and dtype == torch.half:
                 atol = 2e-3
@@ -182,11 +189,6 @@ def dot_prod_attention(
                 )
 
     def _test_insignificant_strides(self):
-        if self.device == "xpu":
-            self.skipTest(
-                "The operator 'aten::_scaled_dot_product_efficient_attention'"
-                " is not currently implemented for the XPU device. "
-            )
         f32 = torch.float32
 
         # repro taken from https://github.com/pytorch/pytorch/issues/124289
@@ -363,10 +365,30 @@ def sfdp_pattern_5_v2(query, key, value):
             )
             return attn_weight @ value
 
+        def sfdp_pattern_5_v3(query, key, value):
+            # https://github.com/pytorch/pytorch/issues/174049.
+            attn_mask = torch.ones(
+                query.size(-2), key.size(-2), dtype=torch.bool, device=query.device
+            ).tril(diagonal=0)
+            attn_mask = attn_mask.masked_fill(
+                torch.logical_not(attn_mask), -float("inf")
+            )
+            attn_weight = torch.softmax(
+                (query @ key.transpose(-2, -1) / (math.sqrt(query.size(-1)) + 0.1))
+                + attn_mask,
+                dim=-1,
+            )
+            return attn_weight @ value
+
         self._check_common(sfdp_pattern_5_v1, contains=False)
         self._check_common(checkpoint_wrapper(sfdp_pattern_5_v1), contains=False)
         self._check_common(sfdp_pattern_5_v2, contains=False)
         self._check_common(checkpoint_wrapper(sfdp_pattern_5_v2), contains=False)
+        self._check_common(sfdp_pattern_5_v3, contains=False)
+        self._check_common(
+            checkpoint_wrapper(sfdp_pattern_5_v3),
+            contains=False,
+        )
 
     def _test_sdpa_rewriter_6(self):
         def sfdp_pattern_6(query, key, value, training):
@@ -383,10 +405,31 @@ def sfdp_pattern_6(query, key, value, training):
             attn_weight = torch.nn.functional.dropout(attn_weight, 0.5, training)
             return attn_weight @ value
 
+        def sfdp_pattern_6_v2(query, key, value, training):
+            attn_mask = torch.ones(
+                query.size(-2), key.size(-2), dtype=torch.bool, device=query.device
+            ).tril(diagonal=0)
+            attn_mask = attn_mask.masked_fill(
+                torch.logical_not(attn_mask), -float("inf")
+            )
+            attn_weight = torch.softmax(
+                (query @ key.transpose(-2, -1) / (math.sqrt(query.size(-1)) + 0.1))
+                + attn_mask,
+                dim=-1,
+            )
+            attn_weight = torch.nn.functional.dropout(attn_weight, 0.5, training)
+            return attn_weight @ value
+
         self._check_common(sfdp_pattern_6, contains=False, has_dropout=True)
         self._check_common(
             checkpoint_wrapper(sfdp_pattern_6), contains=False, has_dropout=True
         )
+        self._check_common(sfdp_pattern_6_v2, contains=False, has_dropout=True)
+        self._check_common(
+            checkpoint_wrapper(sfdp_pattern_6_v2),
+            contains=False,
+            has_dropout=True,
+        )
 
     def _test_sdpa_rewriter_7(self):
         def sfdp_pattern_7(query, key, value, training):
@@ -401,6 +444,18 @@ def sfdp_pattern_7(query, key, value, training):
             attn_weight = attn_weight.to(torch.float16)
             return attn_weight @ v
 
+        def sfdp_pattern_7_v2(query, key, value, training):
+            q = query.permute(0, 2, 1, 3)
+            k = key.permute(0, 2, 1, 3)
+            v = value.permute(0, 2, 1, 3)
+            div = q @ k.transpose(-2, -1) / (math.sqrt(q.size(-1)) + 0.1)
+            div = div.to(torch.float32)
+            attn_weight = torch.softmax(div, dim=-1)
+            # Set to False
+            attn_weight = torch.dropout(attn_weight, 0.00000000001, training)
+            attn_weight = attn_weight.to(torch.float16)
+            return attn_weight @ v
+
         args = (
             torch.randn((2, 8, 4, 16), device=self.device, dtype=torch.half),
             torch.randn((2, 8, 4, 16), device=self.device, dtype=torch.half),
@@ -414,6 +469,19 @@ def sfdp_pattern_7(query, key, value, training):
             override_check_equal=True,
             atol=2e-3,
         )
+        args = (
+            torch.randn((2, 8, 4, 16), device=self.device, dtype=torch.half),
+            torch.randn((2, 8, 4, 16), device=self.device, dtype=torch.half),
+            torch.randn((2, 8, 4, 16), device=self.device, dtype=torch.half),
+        )
+        self._check_common(
+            sfdp_pattern_7_v2,
+            args,
+            contains=False,
+            has_dropout=True,
+            override_check_equal=True,
+            atol=2e-3,
+        )
 
         args = (
             torch.randn((2, 8, 4, 16), device=GPU_TYPE, dtype=torch.half),
@@ -428,6 +496,19 @@ def sfdp_pattern_7(query, key, value, training):
             override_check_equal=True,
             atol=2e-3,
         )
+        args = (
+            torch.randn((2, 8, 4, 16), device=GPU_TYPE, dtype=torch.half),
+            torch.randn((2, 8, 4, 16), device=GPU_TYPE, dtype=torch.half),
+            torch.randn((2, 8, 4, 16), device=GPU_TYPE, dtype=torch.half),
+        )
+        self._check_common(
+            checkpoint_wrapper(sfdp_pattern_7_v2),
+            args,
+            contains=SM80OrLater,
+            has_dropout=True,
+            override_check_equal=True,
+            atol=2e-3,
+        )
 
     def _test_sdpa_rewriter_8(self):
         def sfdp_pattern_8(query, key, value):
@@ -440,12 +521,28 @@ def sfdp_pattern_8(query, key, value):
             attn_weight = attn_weight.to(torch.float16)
             return attn_weight @ v
 
+        def sfdp_pattern_8_v2(query, key, value):
+            q = query.permute(0, 2, 1, 3)
+            k = key.permute(0, 2, 1, 3)
+            v = value.permute(0, 2, 1, 3)
+            div = q @ k.transpose(-2, -1) / (math.sqrt(q.size(-1)) + 0.1)
+            div = div.to(torch.float32)
+            attn_weight = torch.softmax(div, dim=-1)
+            attn_weight = attn_weight.to(torch.float16)
+            return attn_weight @ v
+
         args = (
             torch.randn((2, 8, 4, 16), device=self.device, dtype=torch.half),
             torch.randn((2, 8, 4, 16), device=self.device, dtype=torch.half),
             torch.randn((2, 8, 4, 16), device=self.device, dtype=torch.half),
         )
         self._check_common(sfdp_pattern_8, args, atol=2e-3)
+        args = (
+            torch.randn((2, 8, 4, 16), device=self.device, dtype=torch.half),
+            torch.randn((2, 8, 4, 16), device=self.device, dtype=torch.half),
+            torch.randn((2, 8, 4, 16), device=self.device, dtype=torch.half),
+        )
+        self._check_common(sfdp_pattern_8_v2, args, atol=2e-3, contains=False)
 
         args = (
             torch.randn((2, 8, 4, 16), device=GPU_TYPE, dtype=torch.half),
@@ -453,6 +550,17 @@ def sfdp_pattern_8(query, key, value):
             torch.randn((2, 8, 4, 16), device=GPU_TYPE, dtype=torch.half),
         )
         self._check_common(checkpoint_wrapper(sfdp_pattern_8), args, atol=2e-3)
+        args = (
+            torch.randn((2, 8, 4, 16), device=GPU_TYPE, dtype=torch.half),
+            torch.randn((2, 8, 4, 16), device=GPU_TYPE, dtype=torch.half),
+            torch.randn((2, 8, 4, 16), device=GPU_TYPE, dtype=torch.half),
+        )
+        self._check_common(
+            checkpoint_wrapper(sfdp_pattern_8_v2),
+            args,
+            atol=2e-3,
+            contains=False,
+        )
 
     def _test_sdpa_rewriter_9(self):
         def sfdp_pattern_9(query, key, value, training):
@@ -468,6 +576,19 @@ def sfdp_pattern_9(query, key, value, training):
             attn_weight = attn_weight.to(torch.float16)
             return attn_weight @ v
 
+        def sfdp_pattern_9_v2(query, key, value, training):
+            q = query.permute(0, 2, 1, 3)
+            k = key.permute(0, 2, 1, 3)
+            v = value.permute(0, 2, 1, 3)
+            q = q / (math.sqrt(q.size(-1)) + 0.1)
+            div = q @ k.transpose(-2, -1)
+            div = div.to(torch.float32)
+            attn_weight = torch.softmax(div, dim=-1)
+            # very low dropout to make test pass
+            attn_weight = torch.dropout(attn_weight, 0.00000000001, training)
+            attn_weight = attn_weight.to(torch.float16)
+            return attn_weight @ v
+
         args = (
             torch.randn((2, 8, 4, 16), device=self.device, dtype=torch.half),
             torch.randn((2, 8, 4, 16), device=self.device, dtype=torch.half),
@@ -481,6 +602,19 @@ def sfdp_pattern_9(query, key, value, training):
             override_check_equal=True,
             atol=2e-3,
         )
+        args = (
+            torch.randn((2, 8, 4, 16), device=self.device, dtype=torch.half),
+            torch.randn((2, 8, 4, 16), device=self.device, dtype=torch.half),
+            torch.randn((2, 8, 4, 16), device=self.device, dtype=torch.half),
+        )
+        self._check_common(
+            sfdp_pattern_9_v2,
+            args,
+            contains=SM80OrLater,
+            has_dropout=True,
+            override_check_equal=True,
+            atol=2e-3,
+        )
         args = (
             torch.randn((2, 8, 4, 16), device=GPU_TYPE, dtype=torch.half),
             torch.randn((2, 8, 4, 16), device=GPU_TYPE, dtype=torch.half),
@@ -494,6 +628,19 @@ def sfdp_pattern_9(query, key, value, training):
             override_check_equal=True,
             atol=2e-3,
         )
+        args = (
+            torch.randn((2, 8, 4, 16), device=GPU_TYPE, dtype=torch.half),
+            torch.randn((2, 8, 4, 16), device=GPU_TYPE, dtype=torch.half),
+            torch.randn((2, 8, 4, 16), device=GPU_TYPE, dtype=torch.half),
+        )
+        self._check_common(
+            checkpoint_wrapper(sfdp_pattern_9_v2),
+            args,
+            contains=SM80OrLater,
+            has_dropout=True,
+            override_check_equal=True,
+            atol=2e-3,
+        )
 
     def _test_sdpa_rewriter_10(self):
         def sfdp_pattern_10(query, key, value):
@@ -507,12 +654,34 @@ def sfdp_pattern_10(query, key, value):
             attn_weight = attn_weight.to(torch.float16)
             return attn_weight @ v
 
+        def sfdp_pattern_10_v2(query, key, value):
+            q = query.permute(0, 2, 1, 3)
+            k = key.permute(0, 2, 1, 3)
+            v = value.permute(0, 2, 1, 3)
+            q = q / (math.sqrt(q.size(-1)) + 0.1)
+            div = q @ k.transpose(-2, -1)
+            div = div.to(torch.float32)
+            attn_weight = torch.softmax(div, dim=-1)
+            attn_weight = attn_weight.to(torch.float16)
+            return attn_weight @ v
+
         args = (
             torch.randn((2, 8, 4, 16), device=self.device, dtype=torch.half),
             torch.randn((2, 8, 4, 16), device=self.device, dtype=torch.half),
             torch.randn((2, 8, 4, 16), device=self.device, dtype=torch.half),
         )
         self._check_common(sfdp_pattern_10, args, atol=2e-3)
+        args = (
+            torch.randn((2, 8, 4, 16), device=self.device, dtype=torch.half),
+            torch.randn((2, 8, 4, 16), device=self.device, dtype=torch.half),
+            torch.randn((2, 8, 4, 16), device=self.device, dtype=torch.half),
+        )
+        self._check_common(
+            sfdp_pattern_10_v2,
+            args,
+            atol=2e-3,
+            contains=False,
+        )
 
         args = (
             torch.randn((2, 8, 4, 16), device=GPU_TYPE, dtype=torch.half),
@@ -520,6 +689,17 @@ def sfdp_pattern_10(query, key, value):
             torch.randn((2, 8, 4, 16), device=GPU_TYPE, dtype=torch.half),
         )
         self._check_common(checkpoint_wrapper(sfdp_pattern_10), args, atol=2e-3)
+        args = (
+            torch.randn((2, 8, 4, 16), device=GPU_TYPE, dtype=torch.half),
+            torch.randn((2, 8, 4, 16), device=GPU_TYPE, dtype=torch.half),
+            torch.randn((2, 8, 4, 16), device=GPU_TYPE, dtype=torch.half),
+        )
+        self._check_common(
+            checkpoint_wrapper(sfdp_pattern_10_v2),
+            args,
+            atol=2e-3,
+            contains=False,
+        )
 
     def _test_pattern_fails_with_tensor_factor(self):
         # https://github.com/pytorch/pytorch/issues/99124
@@ -888,6 +1068,37 @@ def dot_prod_attention(
                 value.permute([0, 2, 1, 3]),
             )
 
+        def dot_prod_attention_v2(
+            query: torch.Tensor,
+            key: torch.Tensor,
+            value: torch.Tensor,
+            causal_mask: torch.Tensor,
+        ) -> torch.Tensor:
+            query = query.permute([0, 2, 1, 3])
+            key = key.permute([0, 2, 1, 3])
+            value = value.permute([0, 2, 1, 3])
+            attn_weights = torch.matmul(query, key.permute(0, 1, 3, 2))
+            inv_scale = torch.full(
+                (),
+                math.sqrt(value.size(-1)) + 0.1,
+                dtype=query.dtype,
+                device=query.device,
+            )
+            attn_weights = attn_weights.div(inv_scale)
+            causal_mask_value = torch.full(
+                (), torch.finfo(query.dtype).min, dtype=query.dtype, device=query.device
+            )
+            attn_weights = torch.where(causal_mask, attn_weights, causal_mask_value)
+            return (
+                (
+                    torch.nn.functional.dropout(
+                        attn_weights.softmax(dim=-1), 0.0
+                    ).matmul(value)
+                ),
+                key.permute([0, 2, 1, 3]),
+                value.permute([0, 2, 1, 3]),
+            )
+
         tensor_shape = (4, 2, 16, 32)
         causal_mask = torch.ones(2, 2, dtype=torch.bool, device=self.device).tril(
             diagonal=0
@@ -906,6 +1117,15 @@ def dot_prod_attention(
             check_train=False,
         )
 
+        self._check_common(
+            dot_prod_attention_v2,
+            args1=args,
+            atol=2e-3,
+            contains=False,
+            has_dropout=False,
+            check_train=False,
+        )
+
         # also check batch_size=1 because the graph is slightly different
         tensor_shape = (1, 2, 16, 32)
         args = [
@@ -921,6 +1141,14 @@ def dot_prod_attention(
             has_dropout=False,
             check_train=False,
         )
+        self._check_common(
+            dot_prod_attention_v2,
+            args1=args,
+            contains=False,
+            atol=2e-3,
+            has_dropout=False,
+            check_train=False,
+        )
 
     def _test_sdpa_rewriter_19(self):
         def dot_prod_attention(
@@ -952,6 +1180,35 @@ def dot_prod_attention(
                 inplace=False,
             ).matmul(value)
 
+        def dot_prod_attention_v2(
+            query: torch.Tensor,
+            key: torch.Tensor,
+            value: torch.Tensor,
+            causal_mask: torch.Tensor,
+            attn_mask: torch.Tensor,
+            training,
+        ) -> torch.Tensor:
+            attn_weights = torch.matmul(query, key.permute(0, 1, 3, 2))
+            inv_scale = torch.full(
+                (),
+                math.sqrt(value.size(-1)) + 0.1,
+                dtype=attn_weights.dtype,
+                device=attn_weights.device,
+            )
+            attn_weights = attn_weights.div(inv_scale)
+            causal_mask_value = torch.full(
+                (), torch.finfo(query.dtype).min, dtype=query.dtype, device=query.device
+            )
+            attn_weights = torch.where(causal_mask, attn_weights, causal_mask_value)
+            attn_weights = attn_weights + attn_mask
+            attn_weights = attn_weights.softmax(dim=-1).type(value.dtype)
+            return torch.nn.functional.dropout(
+                attn_weights,
+                p=0.4,
+                training=training,
+                inplace=False,
+            ).matmul(value)
+
         tensor_shape = (4, 2, 16, 32)
         causal_mask = torch.ones(16, 16, dtype=torch.bool, device=self.device).tril(
             diagonal=0
@@ -971,6 +1228,13 @@ def dot_prod_attention(
             has_dropout=True,
             check_train=False,
         )
+        self._check_common(
+            dot_prod_attention_v2,
+            args1=args,
+            contains=False,
+            has_dropout=True,
+            check_train=False,
+        )
 
     def _test_sdpa_rewriter_20(self):
         def dot_prod_attention(
@@ -1181,6 +1445,71 @@ def dot_prod_attention(
             check_train=True,
         )
 
+    @torch._inductor.config.patch("cache_sdpa_constraint", True)
+    def _test_cache_sdpa_constraint_shared_kv_enabled(self):
+        """When cache_sdpa_constraint is True and the same tensor feeds both key
+        and value in SDPA (PMA pattern), the constraint cache should deduplicate
+        the buffer copy, producing fewer copy operations than when disabled."""
+
+        def pma_attention(query, kv):
+            return torch.nn.functional.scaled_dot_product_attention(query, kv, kv)
+
+        tensor_shape = (4, 2, 16, 32)
+        dtype = torch.float
+        query = torch.randn(tensor_shape, device=self.device, dtype=dtype)
+        kv = torch.randn(tensor_shape, device=self.device, dtype=dtype)
+
+        counters.clear()
+        torch._dynamo.reset()
+        result_eager = pma_attention(query, kv)
+        result_compiled, (source_code,) = run_and_get_code(
+            torch.compile(pma_attention, fullgraph=True),
+            query.clone(),
+            kv.clone(),
+        )
+        self.assertEqual(result_eager, result_compiled, atol=1e-3, rtol=0.2)
+        copy_count_cached = source_code.count("copy_")
+        return copy_count_cached
+
+    @torch._inductor.config.patch("cache_sdpa_constraint", False)
+    def _test_cache_sdpa_constraint_shared_kv_disabled(self):
+        """When cache_sdpa_constraint is False and the same tensor feeds both key
+        and value in SDPA (PMA pattern), the constraint should create separate
+        buffer copies, producing more copy operations than when enabled."""
+
+        def pma_attention(query, kv):
+            return torch.nn.functional.scaled_dot_product_attention(query, kv, kv)
+
+        tensor_shape = (4, 2, 16, 32)
+        dtype = torch.float
+        query = torch.randn(tensor_shape, device=self.device, dtype=dtype)
+        kv = torch.randn(tensor_shape, device=self.device, dtype=dtype)
+
+        counters.clear()
+        torch._dynamo.reset()
+        result_eager = pma_attention(query, kv)
+        result_compiled, (source_code,) = run_and_get_code(
+            torch.compile(pma_attention, fullgraph=True),
+            query.clone(),
+            kv.clone(),
+        )
+        self.assertEqual(result_eager, result_compiled, atol=1e-3, rtol=0.2)
+        copy_count_uncached = source_code.count("copy_")
+        return copy_count_uncached
+
+    def _test_cache_sdpa_constraint_shared_kv(self):
+        """Verify that cache_sdpa_constraint=True produces no more copy_
+        operations than cache_sdpa_constraint=False when the same tensor feeds
+        both key and value in SDPA."""
+        copy_count_cached = self._test_cache_sdpa_constraint_shared_kv_enabled()
+        copy_count_uncached = self._test_cache_sdpa_constraint_shared_kv_disabled()
+        self.assertLessEqual(
+            copy_count_cached,
+            copy_count_uncached,
+            f"Expected caching to produce <= copy_ ops (got {copy_count_cached}) "
+            f"vs no caching ({copy_count_uncached})",
+        )
+
     def _test_sdpa_rewriter_26(self):
         def dot_prod_attention(
             query: torch.Tensor,
@@ -1250,6 +1579,32 @@ def dot_prod_attention(
             check_train=True,
         )
 
+    def _test_sdpa_rewriter_28(self):
+        def dot_prod_attention(
+            qkv: torch.Tensor,
+            training: bool,
+        ) -> torch.Tensor:
+            q, k, v = qkv.permute(1, 0, 2, 4, 3).unbind(0)
+            scores = torch.matmul(q, k.transpose(-2, -1))
+            scores = scores.mul(0.2)
+            attn_weights = scores.softmax(dim=-1)
+            attn_weights = torch.nn.functional.dropout(
+                attn_weights, p=0.1, training=training
+            )
+            return attn_weights.matmul(v)
+
+        tensor_shape = (2, 3, 4, 16, 8)
+        args = [
+            torch.randn(tensor_shape, dtype=torch.half, device=self.device),
+        ]
+        self._check_common(
+            dot_prod_attention,
+            args1=args,
+            contains=False,
+            has_dropout=True,
+            check_train=True,
+        )
+
 
 if HAS_XPU_AND_TRITON or (HAS_CUDA_AND_TRITON and PLATFORM_SUPPORTS_FUSED_ATTENTION):
 
@@ -1332,6 +1687,12 @@ class SDPAPatternRewriterGpuTests(TestSDPAPatternRewriterTemplate):
         test_sdpa_rewriter_24_gpu = functools.partialmethod(
             TestSDPAPatternRewriterTemplate._test_sdpa_rewriter_24
         )
+        test_cache_sdpa_constraint_shared_kv_gpu = (
+            TestSDPAPatternRewriterTemplate._test_cache_sdpa_constraint_shared_kv
+        )
+        test_sdpa_rewriter_28_gpu = functools.partialmethod(
+            TestSDPAPatternRewriterTemplate._test_sdpa_rewriter_28
+        )
         if HAS_XPU_AND_TRITON:
             test_sdpa_rewriter_25_gpu = functools.partialmethod(
                 TestSDPAPatternRewriterTemplate._test_sdpa_rewriter_25
@@ -1393,7 +1754,10 @@ class SDPAPatternRewriterCpuTests(TestSDPAPatternRewriterTemplate):
             TestSDPAPatternRewriterTemplate._test_pattern_fails_with_reuse
         )
         test_sdpa_rewriter_2_cpu = TestSDPAPatternRewriterTemplate._test_sdpa_rewriter_2
-        test_sdpa_rewriter_5_cpu = TestSDPAPatternRewriterTemplate._test_sdpa_rewriter_5
+        # see https://github.com/pytorch/pytorch/issues/177244
+        test_sdpa_rewriter_5_cpu = xfailIf(IS_ARM64 and IS_CPU_CAPABILITY_SVE256)(
+            TestSDPAPatternRewriterTemplate._test_sdpa_rewriter_5
+        )
         test_pattern_fails_with_tensor_factor_cpu = (
             TestSDPAPatternRewriterTemplate._test_pattern_fails_with_tensor_factor
         )
@@ -1412,7 +1776,8 @@ class SDPAPatternRewriterCpuTests(TestSDPAPatternRewriterTemplate):
         test_sdpa_rewriter_13_cpu = functools.partialmethod(
             TestSDPAPatternRewriterTemplate._test_sdpa_rewriter_13, dtype=torch.float32
         )
-        test_sdpa_rewriter_14_cpu = functools.partialmethod(
+        # see https://github.com/pytorch/pytorch/issues/177244
+        test_sdpa_rewriter_14_cpu = xfailIf(IS_ARM64 and IS_CPU_CAPABILITY_SVE256)(
             TestSDPAPatternRewriterTemplate._test_sdpa_rewriter_14
         )
         test_sdpa_rewriter_15_cpu = functools.partialmethod(
@@ -1448,6 +1813,12 @@ class SDPAPatternRewriterCpuTests(TestSDPAPatternRewriterTemplate):
         test_sdpa_rewriter_24_cpu = functools.partialmethod(
             TestSDPAPatternRewriterTemplate._test_sdpa_rewriter_24
         )
+        test_cache_sdpa_constraint_shared_kv_cpu = (
+            TestSDPAPatternRewriterTemplate._test_cache_sdpa_constraint_shared_kv
+        )
+        test_sdpa_rewriter_28_cpu = functools.partialmethod(
+            TestSDPAPatternRewriterTemplate._test_sdpa_rewriter_28
+        )
 
     class SDPAPatternRewriterCpuDynamicTests(SDPAPatternRewriterCpuTests):
         use_static_shapes = False
diff --git a/test/inductor/test_fxir_backend.py b/test/inductor/test_fxir_backend.py
index 6c05e25ea85dc..3f95800007354 100644
--- a/test/inductor/test_fxir_backend.py
+++ b/test/inductor/test_fxir_backend.py
@@ -7,7 +7,6 @@
 import operator
 import unittest
 from collections.abc import Callable
-from typing import Optional
 
 import sympy
 
@@ -98,7 +97,7 @@ def _compile_and_check(
         args,
         expected_num_triton_kernels: int = 1,
         metadata_only: bool = False,
-        compile_kwargs: Optional[dict] = None,
+        compile_kwargs: dict | None = None,
     ):
         if compile_kwargs is None:
             compile_kwargs = {}
@@ -258,7 +257,8 @@ def foo(x, y):
 
         def get_offset(node: torch.fx.Node) -> int:
             (input_, shape, stride, offset) = node.args
-            assert isinstance(offset, int)
+            if not isinstance(offset, int):
+                raise AssertionError
             return offset
 
         # Check for 2 views, one of which is offset.
@@ -394,7 +394,7 @@ def foo(x, y):
 
         # Expect separate forward and backward graphs.
         (forward_gm, backward_gm) = self._compile_and_check(
-            foo, (x, y), expected_num_triton_kernels=3
+            foo, (x, y), expected_num_triton_kernels=4
         )
 
     def test_custom_compiler(self):
diff --git a/test/inductor/test_gpu_cpp_wrapper.py b/test/inductor/test_gpu_cpp_wrapper.py
index 6597749d435b7..bad8817c9d7f2 100644
--- a/test/inductor/test_gpu_cpp_wrapper.py
+++ b/test/inductor/test_gpu_cpp_wrapper.py
@@ -1,13 +1,24 @@
 # Owner(s): ["module: inductor"]
 import itertools
+import os
+import subprocess
 import sys
+import tempfile
 import unittest
 from typing import NamedTuple
 
 import torch
 from torch._inductor import config
+from torch._inductor.codegen.common import TritonScratchWorkspace
+from torch._inductor.codegen.cpp_wrapper_gpu import DeferredTritonCallWrapper
+from torch._inductor.codegen.cuda.device_op_overrides import CUDADeviceOpOverrides
 from torch._inductor.test_case import TestCase as InductorTestCase
-from torch.testing._internal.common_utils import slowTest
+from torch._inductor.utils import IndentedBuffer
+from torch.testing._internal.common_utils import (
+    instantiate_parametrized_tests,
+    parametrize,
+    slowTest,
+)
 from torch.testing._internal.inductor_utils import GPU_TYPE, RUN_GPU
 
 
@@ -59,6 +70,7 @@ def test_fn():
         comp = torch.compile(
             options={
                 "cpp_wrapper": True,
+                "cpp_wrapper_build_separate": True,
                 "aot_inductor.debug_intermediate_value_printer": "2",
             }
         )(test_fn)
@@ -77,6 +89,243 @@ def test_fn(x, s):
             _, code = test_torchinductor.run_and_get_cpp_code(compiled, x, 3)
         self.assertIn("torch.tensor(arg, device='cpu')", code)
 
+    def test_cpp_scratch_scales_with_grid_size_for_tma(self):
+        if GPU_TYPE != "cuda" or torch.version.hip:
+            self.skipTest("CUDA-only codegen test")
+
+        scratch_def, scratch_var = CUDADeviceOpOverrides().cpp_scratch(
+            0,
+            TritonScratchWorkspace(
+                size=256, generate_dtype_str=lambda: "at::ScalarType::Byte"
+            ),
+            prefix="global_scratch",
+        )
+        self.assertEqual(scratch_var, "global_scratch_scratch_0")
+        self.assertIn(
+            "static_cast<int64_t>(256) * grid_0 * grid_1 * grid_2", scratch_def[0]
+        )
+
+    def test_triton_wrapper_scales_scratch_with_num_ctas(self):
+        if GPU_TYPE != "cuda" or torch.version.hip:
+            self.skipTest("CUDA-only codegen test")
+
+        class FakeWrapper:
+            device = "cuda"
+
+            def __init__(self):
+                self.scratch_spaces = None
+
+            def generate_args_decl(
+                self,
+                prefix,
+                call_args,
+                arg_types,
+                arg_signatures,
+                is_triton_kernel=True,
+                scratch_spaces=None,
+            ):
+                self.scratch_spaces = scratch_spaces
+
+                return ""
+
+        wrapper = FakeWrapper()
+        prefix = IndentedBuffer()
+        params = {
+            "triton_meta": {"signature": {"x": "*fp32"}, "constants": {}},
+            "def_args": ["x"],
+            "call_args": ["x"],
+            "config": {"num_ctas": 8},
+            "num_warps": 4,
+            "shared_mem": 0,
+            "global_scratch": 256,
+        }
+
+        DeferredTritonCallWrapper(
+            wrapper_name="wrapper",
+            kernel_name="kernel",
+            kernel_name_to_body={},
+            arg_types=[torch.float32],
+        ).generate_launch_kernel(prefix, wrapper, "kernel_var", params)
+
+        self.assertEqual(wrapper.scratch_spaces, {"global_scratch": 256 * 8})
+
+    @parametrize("per_subkernel_blocks", [False, True])
+    def test_lazy_compile_combo_kernel_default_config(self, per_subkernel_blocks):
+        """Lazy compile should use default_config from combo_grid_meta for XBLOCK."""
+        if not RUN_GPU:
+            self.skipTest("GPU not available")
+
+        from unittest.mock import patch
+
+        from torch._inductor.codegen.triton_combo_kernel import (
+            DEFAULT_COMBO_BLOCK_SIZE_1D,
+        )
+        from torch._inductor.runtime import triton_lazy_compile as tlc
+
+        captured = {}
+        original = tlc.run_triton_kernel_with_autotune
+
+        def capture(pending_kernels, kernel_name, stream, args):
+            result = original(pending_kernels, kernel_name, stream, args)
+            if "triton_for_fused" in kernel_name:
+                captured[kernel_name] = result.xblock
+            return result
+
+        with patch.object(tlc, "run_triton_kernel_with_autotune", side_effect=capture):
+            params = [torch.randn(1024, device=self.device) for _ in range(4)]
+            grads = [torch.randn_like(p) for p in params]
+
+            @torch.compile(
+                options={
+                    "cpp_wrapper": True,
+                    "triton.autotune_at_compile_time": False,
+                    "combo_kernels": True,
+                    "combo_kernel_per_subkernel_blocks": per_subkernel_blocks,
+                }
+            )
+            def fn(params, grads):
+                torch._foreach_add_(params, grads, alpha=-0.1)
+
+            fn(params, grads)
+
+        self.assertTrue(len(captured) > 0, "No combo kernels were lazy-compiled")
+        for name, xblock in captured.items():
+            # When per_subkernel_blocks=False, default_config has a single XBLOCK
+            # that must be picked up correctly (not the hardcoded fallback of 128).
+            # When per_subkernel_blocks=True, default_config uses per-subkernel
+            # XBLOCK_N keys instead, so result.xblock is not used for grid
+            # computation; just verify compilation succeeded.
+            if not per_subkernel_blocks:
+                self.assertEqual(
+                    xblock,
+                    DEFAULT_COMBO_BLOCK_SIZE_1D,
+                    f"{name} got XBLOCK={xblock}, expected {DEFAULT_COMBO_BLOCK_SIZE_1D}",
+                )
+
+    def test_cudagraph_no_partition(self):
+        if not RUN_GPU:
+            self.skipTest("GPU not available")
+
+        def test_fn(x, s):
+            return (x + s).sum()
+
+        x = torch.randn(4, device=self.device)
+        s = 3
+        expected = test_fn(x, s)
+
+        comp = torch.compile(
+            options={
+                "cpp_wrapper": True,
+                "triton.cudagraphs": True,
+                "graph_partition": False,
+            }
+        )(test_fn)
+        for i in range(3):
+            res = comp(x, s)
+            self.assertEqual(res, expected)
+
+    def test_many_args_fold_expression_nesting(self):
+        if not RUN_GPU:
+            self.skipTest("GPU not available")
+        if GPU_TYPE == "xpu":
+            self.skipTest("ocloc backend compiler crashes with too many kernel args")
+
+        num_params = 130
+        params = [torch.randn(64, device=self.device) for _ in range(num_params)]
+        grads = [torch.randn_like(p) for p in params]
+        expected = [p.clone() + (-0.1) * g for p, g in zip(params, grads)]
+
+        @torch.compile(
+            options={
+                "cpp_wrapper": True,
+                "combo_kernels": True,
+                "combo_kernel_max_num_args": 1000,
+            }
+        )
+        def fn(params, grads):
+            torch._foreach_add_(params, grads, alpha=-0.1)
+
+        fn(params, grads)
+
+        for p, e in zip(params, expected):
+            self.assertEqual(p, e)
+
+
+instantiate_parametrized_tests(TestGpuWrapper)
+
+# Helper script for test_lazy_compile_kernel_name_collision_across_modules.
+# Run as a subprocess so dlopen truly re-runs .so static initializers.
+_LAZY_COMPILE_COLLISION_SCRIPT = """\
+import torch
+from torch._inductor import config
+
+config.cpp_wrapper = True
+config.triton.autotune_at_compile_time = False
+
+def fn(x, y, z, w):
+    a = x.sin()
+    torch._dynamo.graph_break()
+    b = (a * y).cos()
+    torch._dynamo.graph_break()
+    c = (b * z).sin()
+    torch._dynamo.graph_break()
+    d = (c * w).cos()
+    return d.sum()
+
+args = [torch.randn(32, device="cuda", requires_grad=True) for _ in range(4)]
+ref_args = [a.detach().clone().requires_grad_(True) for a in args]
+ref = fn(*ref_args)
+ref.backward()
+
+compiled_fn = torch.compile(fn)
+res = compiled_fn(*args)
+res.backward()
+
+assert torch.allclose(res.detach(), ref.detach()), f"Forward mismatch: {res} vs {ref}"
+for i, (a, r) in enumerate(zip(args, ref_args)):
+    assert torch.allclose(a.grad, r.grad), f"Grad mismatch for arg {i}"
+"""
+
+
+class TestLazyCompileKernelCollision(InductorTestCase):
+    device = GPU_TYPE
+
+    def test_lazy_compile_kernel_name_collision_across_modules(self):
+        """The collision manifests when a fresh process loads .so modules from
+        warm on-disk caches: AOTAutograd cache hits cause both forward and
+        backward .so to be loaded (static initializers register kernels in
+        _pending_kernels) before either executes.  If two modules share a
+        kernel name, the global dict collision corrupts the mapping.
+
+        This requires two process invocations because dlopen within a single
+        process reuses loaded libraries without re-running static initializers.
+        """
+        if not RUN_GPU:
+            self.skipTest("GPU not available")
+
+        with tempfile.TemporaryDirectory() as cache_dir:
+            env = {
+                **os.environ,
+                "TORCHINDUCTOR_CACHE_DIR": cache_dir,
+                "INDUCTOR_TEST_DISABLE_FRESH_CACHE": "1",
+            }
+            # First run: cold compile, populates on-disk caches.
+            r1 = subprocess.run(
+                [sys.executable, "-c", _LAZY_COMPILE_COLLISION_SCRIPT],
+                capture_output=True,
+                text=True,
+                env=env,
+            )
+            self.assertEqual(r1.returncode, 0, f"Cold run failed:\n{r1.stderr[-2000:]}")
+            # Second run: warm caches trigger the collision without the fix.
+            r2 = subprocess.run(
+                [sys.executable, "-c", _LAZY_COMPILE_COLLISION_SCRIPT],
+                capture_output=True,
+                text=True,
+                env=env,
+            )
+            self.assertEqual(r2.returncode, 0, f"Warm run failed:\n{r2.stderr[-2000:]}")
+
 
 class DynamicShapesGpuWrapperGpuTests(InductorTestCase):
     device = GPU_TYPE
@@ -108,15 +357,29 @@ def test_fn():
     "test_mm_plus_mm2_dynamic_shapes": test_torchinductor.TestFailure(
         ("gpu_wrapper",), is_skip=True
     ),
-    # ATen ops: scaled_dot_product_efficient_attention not implemented on XPU.
-    "test_scaled_dot_product_efficient_attention_xpu": test_torchinductor.TestFailure(
-        ("gpu_wrapper",), is_skip=False
-    ),
-    "test_scaled_dot_product_efficient_attention_xpu_dynamic_shapes": test_torchinductor.TestFailure(
-        ("gpu_wrapper",), is_skip=False
-    ),
 }
 
+# XPU: complex add decomposition can return NotImplemented in cpp_wrapper path,
+# which currently surfaces as InductorError in test_add_complex4_xpu_gpu_wrapper.
+# Keep this targeted skip to XPU only.
+if device_type == "xpu":
+    test_failures_gpu_wrapper["test_add_complex4_xpu"] = test_torchinductor.TestFailure(
+        ("gpu_wrapper",), is_skip=True
+    )
+    test_failures_gpu_wrapper["test_add_complex_xpu"] = test_torchinductor.TestFailure(
+        ("gpu_wrapper",), is_skip=True
+    )
+    test_failures_gpu_wrapper["test_adding_tensor_offsets_xpu"] = (
+        test_torchinductor.TestFailure(("gpu_wrapper",), is_skip=True)
+    )
+
+# Skip only on CUDA as wrapper dynamic shapes passes on ROCm.
+# Per https://github.com/pytorch/pytorch/pull/172780
+if not torch.version.hip:
+    test_failures_gpu_wrapper["test_mm_plus_mm3_dynamic_shapes"] = (
+        test_torchinductor.TestFailure(("gpu_wrapper",), is_skip=False)
+    )
+
 
 def make_test_case(
     name,
@@ -133,7 +396,8 @@ def make_test_case(
         code_string_count = {}
 
     func = getattr(tests, test_name)
-    assert callable(func), "not a callable"
+    if not callable(func):
+        raise AssertionError("not a callable")
     func = slowTest(func) if slow else func
 
     @config.patch(cpp_wrapper=True)
@@ -181,11 +445,6 @@ class BaseTest(NamedTuple):
         tests: InductorTestCase = test_torchinductor.GPUTests()
         check_code: bool = True
 
-    # XPU Not implemented yet
-    XPU_BASE_TEST_SKIP = [
-        "test_dynamic_shapes_persistent_reduction_mixed_x_dim",
-    ]
-
     # Maintain two separate test lists for cuda and cpp for now
     for item in [
         BaseTest("test_add_complex"),
@@ -286,36 +545,20 @@ class BaseTest(NamedTuple):
             )
         ],
         BaseTest("test_dtypeview_fusion"),
-        # skip if not enough SMs
+        # skip the next two tests if not enough SMs, logic for this is handled by TestSelectAlgorithm.setUp()
         BaseTest(
             "test_addmm",
             device=None,
             tests=test_select_algorithm.TestSelectAlgorithm(),
         ),
-        # skip if not enough SMs
         BaseTest(
             "test_linear_relu",
             device=None,
             tests=test_select_algorithm.TestSelectAlgorithm(),
         ),
     ]:
-        if item.device == "xpu" and item.name in XPU_BASE_TEST_SKIP:
-            continue
         make_test_case(item.name, item.device, item.tests, check_code=item.check_code)
 
-    from torch._inductor.utils import is_big_gpu
-
-    if GPU_TYPE in ("cuda", "xpu") and is_big_gpu():
-        skip_list = ["test_addmm", "test_linear_relu"]
-        # need to skip instead of omit, otherwise fbcode ci can be flaky
-        for test_name in skip_list:
-            test_failures_gpu_wrapper[f"{test_name}_{device_type}"] = (
-                test_torchinductor.TestFailure(("gpu_wrapper",), is_skip=True)
-            )
-            test_failures_gpu_wrapper[f"{test_name}_gpu_dynamic_shapes"] = (
-                test_torchinductor.TestFailure(("gpu_wrapper",), is_skip=True)
-            )
-
     test_torchinductor.copy_tests(
         GpuWrapperTemplate, TestGpuWrapper, "gpu_wrapper", test_failures_gpu_wrapper
     )
diff --git a/test/inductor/test_gpu_select_algorithm.py b/test/inductor/test_gpu_select_algorithm.py
index af101d1113325..1c43906fa49b9 100644
--- a/test/inductor/test_gpu_select_algorithm.py
+++ b/test/inductor/test_gpu_select_algorithm.py
@@ -52,7 +52,6 @@ def skip_cache(self, choices, name, key, benchmark, hint_override=None):
 
     for patcher in [
         dynamo_config.patch(verbose=True),
-        dynamo_config.patch(inline_inbuilt_nn_modules=True),
         inductor_config.patch(
             debug=True,
             max_autotune=True,
diff --git a/test/inductor/test_group_batch_fusion.py b/test/inductor/test_group_batch_fusion.py
index adccebe785916..670258df00197 100644
--- a/test/inductor/test_group_batch_fusion.py
+++ b/test/inductor/test_group_batch_fusion.py
@@ -324,7 +324,8 @@ def compare_dict_tensors(self, ref_dict, res_dict, rtol=1e-3, atol=1e-3):
             return False
         for key1 in ref_dict:
             key2 = "_orig_mod." + key1
-            assert key2 in res_dict, f"{key1} does not exist in traced module"
+            if key2 not in res_dict:
+                raise AssertionError(f"{key1} does not exist in traced module")
             if not torch.allclose(ref_dict[key1], res_dict[key2], rtol=rtol, atol=atol):
                 return False
         return True
@@ -684,7 +685,8 @@ def build_graph(self, desc):
         unsatisfied = 0
         while desc:
             unsatisfied += 1
-            assert unsatisfied <= len(desc)  # cycle or bad input?
+            if unsatisfied > len(desc):
+                raise AssertionError("cycle or bad input?")
             name, v = desc.popleft()
             args = tuple(lookup.get(n) for n in v)
             if None in args:
diff --git a/test/inductor/test_halide.py b/test/inductor/test_halide.py
index 97485e8f820b0..884d15869f69d 100644
--- a/test/inductor/test_halide.py
+++ b/test/inductor/test_halide.py
@@ -130,7 +130,7 @@ def generate(g):
                         out_ptr0.set_estimates([hl.Range(1024, 1024)])
 
                 __name__ == '__main__' and hl.main()
-                """
+                """,  # noqa: S101
             ),
         )
         a = torch.randn(1024)
@@ -204,7 +204,7 @@ def generate(g):
                         tmp1.compute_inline()
 
                 __name__ == '__main__' and hl.main()
-                """
+                """,  # noqa: S101
             ),
         )
         a = torch.randn(1024)
diff --git a/test/inductor/test_indexing.py b/test/inductor/test_indexing.py
index f875a7c7f5bac..373596d710204 100644
--- a/test/inductor/test_indexing.py
+++ b/test/inductor/test_indexing.py
@@ -6,6 +6,7 @@
 import sympy
 
 import torch
+from torch._dynamo.source import ConstantSource
 from torch._inductor.codegen.cpp import cexpr
 from torch._inductor.codegen.triton import texpr
 from torch._inductor.codegen.wrapper import pexpr
@@ -22,6 +23,7 @@
 from torch.testing._internal.inductor_utils import GPU_TYPE, HAS_CPU, HAS_GPU
 from torch.utils._sympy.functions import (
     FloorDiv,
+    Identity,
     Mod,
     ModularIndexing,
     PythonMod,
@@ -115,11 +117,63 @@ def test_indexing_simplification(self):
         self.assertEqual(sizevars.simplify_with_ranges(expr, var_ranges), expr)
         var_ranges = {i2: 784}
         expr = ModularIndexing(ModularIndexing(i2, 1, 28), 7, 4)
-        expected = FloorDiv(ModularIndexing(i2, 1, 28), 7)
+        # FloorDiv(ModularIndexing(b, d1, m), d2) simplifies to
+        # ModularIndexing(b, d1*d2, m//d2) when d2 | m
+        expected = ModularIndexing(i2, 7, 4)
         self.assertEqual(sizevars.simplify_with_ranges(expr, var_ranges), expected)
         expr = ModularIndexing(ModularIndexing(i2, 1, 28) + 1, 7, 4)
         self.assertEqual(sizevars.simplify_with_ranges(expr, var_ranges), expr)
 
+    def test_floordiv_modularindexing_simplification(self):
+        sizevars = SizeVarAllocator()
+        i0 = sympy.Symbol("i0", integer=True, nonneg=True)
+
+        # FloorDiv(ModularIndexing(b, d1, m), d2) -> ModularIndexing(b, d1*d2, m//d2)
+        # when d2 divides m
+        self.assertEqual(
+            sizevars.simplify_with_ranges(
+                FloorDiv(ModularIndexing(i0, 1, 8192), 128), {}
+            ),
+            ModularIndexing(i0, 128, 64),
+        )
+        self.assertEqual(
+            sizevars.simplify_with_ranges(FloorDiv(ModularIndexing(i0, 2, 120), 6), {}),
+            ModularIndexing(i0, 12, 20),
+        )
+        # Does NOT simplify when d2 does not divide m
+        expr = FloorDiv(ModularIndexing(i0, 1, 28), 5)
+        self.assertEqual(sizevars.simplify_with_ranges(expr, {}), expr)
+
+        # FloorDiv(base, divisor) -> 0 when 0 <= base < divisor
+        self.assertEqual(
+            sizevars.simplify_with_ranges(FloorDiv(ModularIndexing(i0, 1, 10), 10), {}),
+            sympy.S.Zero,
+        )
+
+    def test_remove_zero_terms_generalized(self):
+        sizevars = SizeVarAllocator()
+        i0 = sympy.Symbol("i0", integer=True, nonneg=True)
+        i1 = sympy.Symbol("i1", integer=True, nonneg=True)
+
+        # FloorDiv(v + 128*i1, 8192): gcd(128*i1, 8192) = 128
+        # Old rule fails (128 != 8192), new rule: v < 128 => drop v
+        self.assertEqual(
+            sizevars.simplify_with_ranges(FloorDiv(i0 + 128 * i1, 8192), {i0: 128}),
+            FloorDiv(128 * i1, 8192),
+        )
+        # v range equals gcd exactly — still safe since v < gcd (strict)
+        # v=127 max, 127 < 128
+        self.assertEqual(
+            sizevars.simplify_with_ranges(FloorDiv(i0 + 6 * i1, 18), {i0: 6}),
+            FloorDiv(6 * i1, 18),
+        )
+        # v range exceeds gcd — cannot simplify
+        expr = FloorDiv(i0 + 128 * i1, 8192)
+        self.assertEqual(
+            sizevars.simplify_with_ranges(expr, {i0: 129}),
+            expr,
+        )
+
     def test_indexing_join(self):
         sizevars = SizeVarAllocator()
         i0 = sympy.Symbol("i0", integer=True)
@@ -367,6 +421,29 @@ def test_print_round(self):
             texpr(expr), """libdevice.llrint((1/2)*x).to(tl.int64)"""
         )
 
+    def test_print_nan(self):
+        expr = sympy.nan
+        self.assertExpectedInline(pexpr(expr), """math.nan""")
+        self.assertExpectedInline(
+            cexpr(expr), """std::numeric_limits<double>::quiet_NaN()"""
+        )
+
+    def test_print_infinity(self):
+        expr = sympy.oo
+        self.assertExpectedInline(pexpr(expr), """math.inf""")
+        self.assertExpectedInline(
+            cexpr(expr),
+            """std::numeric_limits<double>::infinity()""",
+        )
+
+    def test_print_negative_infinity(self):
+        expr = -sympy.oo
+        self.assertExpectedInline(pexpr(expr), """-math.inf""")
+        self.assertExpectedInline(
+            cexpr(expr),
+            """-std::numeric_limits<double>::infinity()""",
+        )
+
     def test_print_integer(self):
         expr = sympy.S((-1) << 63)
         self.assertExpectedInline(cexpr(expr), f"""(-1{LONG_SUFFIX} << 63)""")
@@ -494,6 +571,185 @@ def test_print_Min_Max(self):
 instantiate_parametrized_tests(ExprPrinterTests)
 
 
+class TestEvaluateMinMax(InductorTestCase):
+    def test_evaluate_min_multiple(self):
+        """min(u0, k*u0) resolves via GCD: gcd(u0, k*u0)=u0.
+        UNSOUND: if u0 < 0 (e.g. u0=-1) true min is 10*u0=-10, not -1."""
+        sizevars = SizeVarAllocator()
+        u0 = sizevars.shape_env.create_unbacked_symint().node.expr
+        self.assertEqual(sizevars.evaluate_min(u0, 10 * u0), u0)
+        self.assertEqual(sizevars.evaluate_min(10 * u0, u0), u0)
+        self.assertEqual(sizevars.evaluate_max(u0, 10 * u0), 10 * u0)
+
+    def test_evaluate_max_concrete(self):
+        """works with concrete values even when negative.  Sound."""
+        sizevars = SizeVarAllocator()
+        self.assertEqual(sizevars.evaluate_max(-5, 3), 3)
+        self.assertEqual(sizevars.evaluate_min(-5, 3), -5)
+
+    def test_evaluate_min_product_gcd(self):
+        """evaluate_min(u0, u0*u1) resolves via GCD fallback: gcd(u0, u0*u1)=u0.
+        UNSOUND: when u1=0 the true min is 0, not u0; when u0<0 ordering inverts."""
+        sizevars = SizeVarAllocator()
+        u0 = sizevars.shape_env.create_unbacked_symint().node.expr
+        u1 = sizevars.shape_env.create_unbacked_symint().node.expr
+        self.assertEqual(sizevars.evaluate_min(u0, u0 * u1), u0)
+
+    def test_evaluate_min_product_with_both_gt_gcd(self):
+        """evaluate_min(u0, u0*u1) with u0>0, u1>0 resolves via GCD.
+        Sound: u1>=1 guarantees u0*u1 >= u0."""
+        sizevars = SizeVarAllocator()
+        u0 = sizevars.shape_env.create_unbacked_symint().node.expr
+        u1 = sizevars.shape_env.create_unbacked_symint().node.expr
+        sizevars.check(sympy.Gt(u0, 0))
+        sizevars.check(sympy.Gt(u1, 0))
+        self.assertEqual(sizevars.evaluate_min(u0, u0 * u1), u0)
+
+    def test_guard_or_false_le_unbacked_symint_with_check(self):
+        """guard_or_false(Le(u0, k*u0)) resolves after constraining u0 >= 0."""
+        sizevars = SizeVarAllocator()
+        u0 = sizevars.shape_env.create_unbacked_symint().node.expr
+        sizevars.check(sympy.Ge(u0, 0))
+        # Le(u0, 10*u0) => 9*u0 >= 0, provable with u0 >= 0
+        self.assertTrue(sizevars.guard_or_false(sympy.Le(u0, 10 * u0)))
+        # Le(10*u0, u0) => -9*u0 >= 0 => u0 <= 0, not provable with u0 >= 0
+        self.assertFalse(sizevars.guard_or_false(sympy.Le(10 * u0, u0)))
+        # Lt(u0, 10*u0) => 9*u0 > 0 => u0 > 0, not provable (u0 could be 0)
+        self.assertFalse(sizevars.guard_or_false(sympy.Lt(u0, 10 * u0)))
+
+
+class TestPrecomputedSizeHinting(InductorTestCase):
+    """Tests for optimization_hint and guarding_hint_or_throw with PRECOMPUTED_SIZE symbols."""
+
+    def test_optimization_hint_with_precomputed_size(self):
+        """Test that optimization_hint correctly resolves PRECOMPUTED_SIZE symbols.
+
+        When a complex expression is replaced with a precomputed size symbol (ps0, ps1, etc.),
+        optimization_hint must use inv_precomputed_replacements to resolve the symbol
+        back to its original expression before computing the hint.
+        """
+        sizevars = SizeVarAllocator()
+
+        # Create a backed symbol with a concrete hint value
+        s0 = sizevars.shape_env.create_symbol(168, source=ConstantSource("s0"))
+        sizevars.shape_env.var_to_val[s0] = sympy.Integer(168)
+        sizevars.backed_var_to_val[s0] = sympy.Integer(168)
+
+        # Create a complex expression that would be precomputed
+        complex_expr = s0 * 8  # Should evaluate to 168 * 8 = 1344
+
+        # Register the expression as a precomputed size (simulating what Inductor does)
+        ps_symbol = sizevars.lookup_precomputed_size(complex_expr)
+
+        # Verify the precomputed symbol was created
+        self.assertIn(complex_expr, sizevars.precomputed_replacements)
+        self.assertIn(ps_symbol, sizevars.inv_precomputed_replacements)
+
+        # Test optimization_hint resolves the ps symbol correctly
+        hint = sizevars.optimization_hint(ps_symbol)
+        expected = 168 * 8  # The concrete value of s0 * 8
+        self.assertEqual(hint, expected)
+
+    def test_guarding_hint_or_throw_with_precomputed_size(self):
+        """Test that guarding_hint_or_throw correctly resolves PRECOMPUTED_SIZE symbols."""
+        sizevars = SizeVarAllocator()
+
+        # Create a backed symbol with a concrete hint value
+        s0 = sizevars.shape_env.create_symbol(42, source=ConstantSource("s0"))
+        sizevars.shape_env.var_to_val[s0] = sympy.Integer(42)
+        sizevars.backed_var_to_val[s0] = sympy.Integer(42)
+
+        # Create a complex expression
+        complex_expr = s0 * 2
+
+        # Register as precomputed size
+        ps_symbol = sizevars.lookup_precomputed_size(complex_expr)
+
+        # Test guarding_hint_or_throw resolves correctly
+        hint = sizevars.guarding_hint_or_throw(ps_symbol)
+        expected = 42 * 2
+        self.assertEqual(hint, expected)
+
+    def test_optimization_hint_with_expression_containing_precomputed_size(self):
+        """Test optimization_hint with an expression that contains a PRECOMPUTED_SIZE symbol."""
+        sizevars = SizeVarAllocator()
+
+        # Create a backed symbol
+        s0 = sizevars.shape_env.create_symbol(10, source=ConstantSource("s0"))
+        sizevars.shape_env.var_to_val[s0] = sympy.Integer(10)
+        sizevars.backed_var_to_val[s0] = sympy.Integer(10)
+
+        # Register s0 * 5 as precomputed (ps0 = 50)
+        ps_symbol = sizevars.lookup_precomputed_size(s0 * 5)
+
+        # Create an expression using the precomputed symbol: ps0 + 3
+        expr = ps_symbol + 3
+
+        # optimization_hint should resolve ps0 -> s0*5 -> 50, then add 3 -> 53
+        hint = sizevars.optimization_hint(expr)
+        self.assertEqual(hint, 53)
+
+
+class TestOptimizationHintZeroDivision(InductorTestCase):
+    """Test that optimization_hint handles ZeroDivisionError from ModularIndexing with zero-valued unbacked symbols."""
+
+    def test_modular_indexing_with_zero_divisor(self):
+        sizevars = SizeVarAllocator()
+        u0 = sizevars.shape_env.create_unbacked_symint().node.expr
+        u1 = sizevars.shape_env.create_unbacked_symint().node.expr
+
+        # u0 + 1 ensures base != 0 after substitution; u1 is the divisor.
+        # With fallback=0: u0->0, u1->0, so (0+1) // 0 -> ZeroDivisionError.
+        # optimization_hint catches ZeroDivisionError and returns fallback.
+        expr = ModularIndexing(u0 + 1, u1, 4)
+        hint = sizevars.optimization_hint(expr, fallback=0)
+        self.assertEqual(hint, 0)
+
+    def test_floor_div_with_zero_divisor(self):
+        """optimization_hint should not crash when FloorDiv has an unbacked
+        symbol as divisor that gets substituted with 0."""
+        sizevars = SizeVarAllocator()
+        u0 = sizevars.shape_env.create_unbacked_symint().node.expr
+        u1 = sizevars.shape_env.create_unbacked_symint().node.expr
+
+        # With fallback=0: u0->0, u1->0, FloorDiv(0+1, 0) -> ZeroDivisionError.
+        # optimization_hint catches ZeroDivisionError and returns fallback.
+        expr = FloorDiv(u0 + 1, u1)
+        hint = sizevars.optimization_hint(expr, fallback=0)
+        self.assertEqual(hint, 0)
+
+    def test_modular_indexing_zero_divisor_nonzero_fallback(self):
+        """When fallback is nonzero, the hint should still not crash."""
+        sizevars = SizeVarAllocator()
+        u0 = sizevars.shape_env.create_unbacked_symint().node.expr
+        u1 = sizevars.shape_env.create_unbacked_symint().node.expr
+
+        # With fallback=8192: u0->8192, u1->8192
+        # (8192+1) // 8192 = 1, 1 % 4 = 1
+        expr = ModularIndexing(u0 + 1, u1, 4)
+        hint = sizevars.optimization_hint(expr, fallback=8192)
+        self.assertEqual(hint, 1)
+
+
+class TestOptimizationHintIdentityExpansion(InductorTestCase):
+    """Test that optimization_hint expands Identity wrappers after _sub_unbacked_exprs."""
+
+    def test_identity_wrapped_expr_resolves_to_int(self):
+        """An expression containing Identity-wrapped constants and an unbacked
+        symbol should resolve to a concrete int after substitution."""
+        sizevars = SizeVarAllocator()
+        u0 = sizevars.shape_env.create_unbacked_symint().node.expr
+
+        # Mirrors the real bug: -u0 * (-Identity(1) + Identity(0))
+        # simplifies to -u0 * (0 - 1) = u0.
+        # Without expand(identity=True) after _sub_unbacked_exprs,
+        # subs({u0: fallback}) leaves -Identity(1) + Identity(0) unexpanded,
+        # causing RuntimeError("Failed to realize expression to int").
+        expr = -u0 * (-Identity(sympy.Integer(1)) + Identity(sympy.Integer(0)))
+        hint = sizevars.optimization_hint(expr, fallback=42)
+        self.assertEqual(hint, 42)
+
+
 if __name__ == "__main__":
     from torch._inductor.test_case import run_tests
 
diff --git a/test/inductor/test_inductor_freezing.py b/test/inductor/test_inductor_freezing.py
index 1cbc6695c7fd1..299532a9cee75 100644
--- a/test/inductor/test_inductor_freezing.py
+++ b/test/inductor/test_inductor_freezing.py
@@ -17,13 +17,7 @@
 from torch._inductor.utils import override_lowering, run_and_get_code
 from torch.testing import FileCheck
 from torch.testing._internal.common_cuda import SM80OrLater, tf32_on_and_off
-from torch.testing._internal.common_utils import (
-    IS_FBCODE,
-    MI200_ARCH,
-    skipIfRocmArch,
-    skipIfXpu,
-    TEST_WITH_SLOW_GRADCHECK,
-)
+from torch.testing._internal.common_utils import IS_FBCODE, TEST_WITH_SLOW_GRADCHECK
 
 
 # Make the helper files in test/ importable
@@ -351,37 +345,15 @@ def foo(mod, inp):
                 mod2.b1 = torch.nn.Parameter(torch.rand([15], device=self.device))
                 mod2.b2 = torch.nn.Parameter(torch.rand([20], device=self.device))
 
-            # not fused
-            count = 3 if hasattr(mod2, "t3") else 2
-
+            # fused: weights share same dim 0 (in_features), different dim 1 is OK
             with torch.no_grad():
                 out_eager = mod2(inp)
                 out, code = run_and_get_code(foo, mod2, inp)
                 FileCheck().check_not(kernel_invoke).check_count(
-                    mm_invoke, count=count, exactly=True
+                    mm_invoke, count=1, exactly=True
                 ).run(code[0])
                 self.assertEqual(out_eager, out)
 
-    # With inlining of inbuilt nn modules, Dynamo traces the innards of inbuilt
-    # module and does not modify the eager module.
-    @torch._dynamo.config.patch(inline_inbuilt_nn_modules=False)
-    def test_error_on_eager(self):
-        mod = ConvBN(3, 32, kernel_size=3, stride=2).eval().to(self.device)
-
-        x = torch.rand(3, 3, 32, 32).to(self.device)
-
-        @torch.compile()
-        def foo(mod, x):
-            return mod(x)
-
-        with torch.no_grad():
-            foo(mod, x)
-
-        with self.assertRaisesRegex(
-            RuntimeError, "Trying to run Pytorch Eager Module after Dynamo Freezing"
-        ):
-            mod(x)
-
     def test_static_indices_cudagraph(self):
         if self.device != "cuda":
             return
@@ -748,7 +720,7 @@ def foo(mod, inp):
         self.assertEqual(eager, compiled)
         self.assertTrue(weight_ref() is None)
 
-    @skipIfRocmArch(MI200_ARCH)
+    @torch._inductor.config.patch(layout_optimization=True)
     def test_conv_with_as_strided(self):
         class Model(nn.Module):
             def __init__(self, groups):
@@ -790,7 +762,6 @@ def foo(mod, inp):
                 mod_eager = mod(x)
                 self.assertEqual(foo(mod, x), mod_eager)
 
-    @skipIfXpu
     @unittest.skipIf(IS_FBCODE, "Not yet runnable in fbcode")
     @unittest.skipIf(
         TEST_WITH_SLOW_GRADCHECK,
diff --git a/test/inductor/test_inductor_scheduler.py b/test/inductor/test_inductor_scheduler.py
index 73bca22a48541..0de5555d21e2b 100644
--- a/test/inductor/test_inductor_scheduler.py
+++ b/test/inductor/test_inductor_scheduler.py
@@ -4,6 +4,7 @@
 from unittest.mock import Mock
 
 import torch
+import torch._inductor.config as inductor_config
 import torch._inductor.metrics as metrics
 import torch.utils.flop_counter
 from torch._dynamo.utils import counters
@@ -23,7 +24,7 @@
     skipIfXpu,
     TestCase,
 )
-from torch.testing._internal.inductor_utils import IS_BIG_GPU
+from torch.testing._internal.inductor_utils import GPU_TYPE, HAS_GPU, IS_BIG_GPU
 from torch.utils._ordered_set import OrderedSet
 
 
@@ -290,7 +291,143 @@ def fn(x, idx):
         )
 
 
+class TestScoreFusionMemory(TestCase):
+    """
+    Tests for _score_fusion_memory_by_buffer_overlap.
+
+    These tests validate the fusion scoring logic that determines when nodes
+    should be fused together based on their memory access patterns.
+
+    Key scenarios:
+    1. Exact matches: read/write has exact matches → should fuse (1 kernel)
+    2. Large overlap (split/cat): reads on different offset but overlap is huge
+       → should fuse because the benefit is large (1 kernel)
+    3. Small overlap: reads on different offset but overlap is small → don't fuse (2 kernels)
+    """
+
+    @skipIf(not HAS_GPU, "GPU not available")
+    @inductor_config.patch("score_fusion_memory_threshold", 1)
+    @inductor_config.patch("min_overlap_ratio", 0.5)
+    def test_exact_same_reads_should_fuse(self) -> None:
+        """
+        Case 1: Exact matches in read/write → should fuse into 1 kernel.
+
+        Two operations reading from the exact same input tensor should be
+        fused together since they can share the data read from memory.
+        """
+
+        def exact_reads(x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+            # Both operations read the exact same input
+            out1 = x * 2
+            out2 = x + 1
+            return out1, out2
+
+        torch._dynamo.reset()
+        metrics.reset()
+
+        x = torch.randn(8, 512, device=GPU_TYPE, dtype=torch.float16)
+
+        compiled_fn = torch.compile(exact_reads, backend="inductor", fullgraph=True)
+        out1_eager, out2_eager = exact_reads(x)
+        out1_compiled, out2_compiled = compiled_fn(x)
+
+        self.assertTrue(torch.allclose(out1_eager, out1_compiled, atol=1e-3, rtol=1e-3))
+        self.assertTrue(torch.allclose(out2_eager, out2_compiled, atol=1e-3, rtol=1e-3))
+        # Should fuse into 1 kernel since both ops read exact same buffer
+        self.assertEqual(metrics.generated_kernel_count, 1)
+
+    @skipIf(not HAS_GPU, "GPU not available")
+    @inductor_config.patch("score_fusion_memory_threshold", 1)
+    @inductor_config.patch("min_overlap_ratio", 0.5)
+    def test_split_cat_large_overlap_should_fuse(self) -> None:
+        """
+        Case 2: Reads on different offset but overlap is huge (split/cat) → should fuse into 1 kernel.
+
+        Split operations read from the same input buffer at different offsets.
+        Since the overlap is large (same underlying buffer), fusing these
+        operations together saves reads and kernel launches.
+        """
+
+        def split_and_process(x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+            s1, s2, s3, s4 = torch.split(x, x.shape[-1] // 4, dim=-1)
+            out1 = torch.cat([s4, s3], dim=-1)
+            out2 = torch.cat([s2, s1], dim=-1)
+            return out1, out2
+
+        torch._dynamo.reset()
+        metrics.reset()
+
+        x = torch.randn(8, 512, device=GPU_TYPE, dtype=torch.float16)
+
+        compiled_fn = torch.compile(
+            split_and_process, backend="inductor", fullgraph=True
+        )
+        out1_eager, out2_eager = split_and_process(x)
+        out1_compiled, out2_compiled = compiled_fn(x)
+
+        self.assertTrue(torch.allclose(out1_eager, out1_compiled, atol=1e-3, rtol=1e-3))
+        self.assertTrue(torch.allclose(out2_eager, out2_compiled, atol=1e-3, rtol=1e-3))
+        # Should fuse into 1 kernel since all ops read from the same underlying buffer
+        self.assertEqual(metrics.generated_kernel_count, 1)
+
+    @skipIf(not HAS_GPU, "GPU not available")
+    @inductor_config.patch("score_fusion_memory_threshold", 1)
+    def test_partial_overlap_below_threshold(self) -> None:
+        """
+        Case 3: Partial overlap below the 0.5 threshold → should NOT fuse (2 kernels).
+
+        Similar to test_split_cat_large_overlap_should_fuse, but each operation
+        also reads from a separate large tensor, making the shared buffer portion
+        less than 50% of total reads.
+
+        Example scenario:
+        - Split x into 4 slices: s1, s2, s3, s4 (each 25% of x)
+        - op1 reads: s1 (from x, ~25%) + y (separate tensor, ~75%) → total 100%
+        - op2 reads: s2 (from x, ~25%) + z (separate tensor, ~75%) → total 100%
+        - Common buffer is x, but each op only reads 25% of their total from x
+        - overlap_ratio = 25% / 100% = 0.25 < 0.5 threshold → score = 0
+        - Result: 2 separate kernels (not fused)
+        """
+
+        def partial_overlap_split(
+            x: torch.Tensor, y: torch.Tensor, z: torch.Tensor
+        ) -> tuple[torch.Tensor, torch.Tensor]:
+            # Split x into 4 parts, use different slices in each output
+            s1, s2, _, _ = torch.split(x, x.shape[-1] // 4, dim=-1)
+            # op1 reads: s1 (small slice of x) + y (large separate tensor)
+            # op2 reads: s2 (small slice of x) + z (large separate tensor)
+            # The slices s1 and s2 come from the same buffer x,
+            # but each is only ~25% of total reads for that op
+            out1 = torch.cat([s1, y, y, y], dim=-1)
+            out2 = torch.cat([s2, z, z, z], dim=-1)
+            return out1, out2
+
+        torch._dynamo.reset()
+        metrics.reset()
+
+        # x is split into 4 parts (each 128 elements)
+        # y and z are 3x larger (384 elements each)
+        # So each op reads: 128 (from x slice) + 384 (from y or z) = 512 total
+        # overlap_ratio = 128 / 512 = 0.25 < 0.5 threshold
+        x = torch.randn(8, 512, device=GPU_TYPE, dtype=torch.float16)
+        y = torch.randn(8, 128, device=GPU_TYPE, dtype=torch.float16)
+        z = torch.randn(8, 128, device=GPU_TYPE, dtype=torch.float16)
+
+        compiled_fn = torch.compile(
+            partial_overlap_split, backend="inductor", fullgraph=True
+        )
+        out1_eager, out2_eager = partial_overlap_split(x, y, z)
+        out1_compiled, out2_compiled = compiled_fn(x, y, z)
+
+        self.assertTrue(torch.allclose(out1_eager, out1_compiled, atol=1e-3, rtol=1e-3))
+        self.assertTrue(torch.allclose(out2_eager, out2_compiled, atol=1e-3, rtol=1e-3))
+        # Should NOT fuse (2 kernels) because overlap_ratio = 0.25 < 0.5 threshold
+        # The _score_fusion_memory_by_buffer_overlap returns 0 for this case
+        self.assertEqual(metrics.generated_kernel_count, 2)
+
+
 instantiate_device_type_tests(TestScheduler, globals(), allow_xpu=True)
+instantiate_device_type_tests(TestScoreFusionMemory, globals(), allow_xpu=True)
 
 if __name__ == "__main__":
     run_tests()
diff --git a/test/inductor/test_inductor_utils.py b/test/inductor/test_inductor_utils.py
index 12468a09103b9..2871a579fe577 100644
--- a/test/inductor/test_inductor_utils.py
+++ b/test/inductor/test_inductor_utils.py
@@ -11,13 +11,19 @@
 
 log = logging.getLogger(__name__)
 
+device_type = (
+    acc.type
+    if (acc := torch.accelerator.current_accelerator(check_available=True))
+    else "cpu"
+)
+
 
 class TestBench(TestCase):
     @classmethod
     def setUpClass(cls):
         super().setUpClass()
-        x = torch.rand(1024, 10).cuda().half()
-        w = torch.rand(512, 10).cuda().half()
+        x = torch.rand(1024, 10).to(device_type).half()
+        w = torch.rand(512, 10).to(device_type).half()
         cls._bench_fn = functools.partial(torch.nn.functional.linear, x, w)
 
     def test_benchmarker(self):
@@ -32,4 +38,4 @@ def test_do_bench_using_profiling(self):
 
 
 if __name__ == "__main__":
-    run_tests("cuda")
+    run_tests(device_type)
diff --git a/test/inductor/test_inplace_padding.py b/test/inductor/test_inplace_padding.py
index 03f7e595a211a..c80671a1c4b9d 100644
--- a/test/inductor/test_inplace_padding.py
+++ b/test/inductor/test_inplace_padding.py
@@ -9,11 +9,11 @@
 from torch._inductor.test_case import run_tests, TestCase
 from torch._inductor.utils import run_and_get_code
 from torch.testing import FileCheck
-from torch.testing._internal.common_utils import serialTest
+from torch.testing._internal.common_utils import serialTest, skipIfXpu
 from torch.testing._internal.inductor_utils import (
     GPU_TYPE,
     HAS_GPU,
-    requires_cuda_with_enough_memory,
+    requires_gpu_with_enough_memory,
 )
 
 
@@ -112,6 +112,7 @@ def f(x):
         self.assertEqual(num_inplace_padding(), 1)
 
     @inductor_config.patch(cpp_wrapper=True)
+    @inductor_config.patch("triton.autotune_at_compile_time", True)
     def test_pad_non_zero_cpp_wrapper(self):
         def f(x):
             x = x + 1
@@ -210,7 +211,7 @@ def f(x, y):
 
         self.assertEqual(num_inplace_padding(), 0)
 
-    @requires_cuda_with_enough_memory(2e10)
+    @requires_gpu_with_enough_memory(2e10)
     @inductor_config.patch(force_shape_pad=True)
     @serialTest()
     def test_linear_and_cel(self):
@@ -241,7 +242,8 @@ def f(x, y):
 
         expect = (f(x, y), x.grad, linear.weight.grad, linear.bias.grad)
         actual = (opt_f(x, y), x.grad, linear.weight.grad, linear.bias.grad)
-        assert same(expect, actual, tol=1e-2), f"ref:\n{expect}\nact:\n{actual}"
+        if not same(expect, actual, tol=1e-2):
+            raise AssertionError(f"ref:\n{expect}\nact:\n{actual}")
 
         # We may disable inplace_padding via env-var to test perf.
         self.assertEqual(num_inplace_padding(), int(inductor_config.inplace_padding))
@@ -254,9 +256,10 @@ def f(x, y):
 
     # Enable Max-Autotune to repro this test failure:
     #   https://github.com/pytorch/pytorch/pull/140249#issuecomment-2556079406
-    @requires_cuda_with_enough_memory(2e10)
+    @requires_gpu_with_enough_memory(2e10)
     @inductor_config.patch(max_autotune=True)
     @serialTest()
+    @skipIfXpu(msg="AssertionError: torch-xpu-ops: #2997")
     def test_linear_and_cel_max_autotune(self):
         self.test_linear_and_cel()
 
diff --git a/test/inductor/test_inplacing_pass.py b/test/inductor/test_inplacing_pass.py
index cde22b7d79a7a..10de5116151d2 100644
--- a/test/inductor/test_inplacing_pass.py
+++ b/test/inductor/test_inplacing_pass.py
@@ -193,7 +193,8 @@ def get_not_inplaced_count(self, graph):
             ):
                 auto_functionalized_found = True
                 counter += len(node.meta["only_clone_these_tensors"])
-        assert auto_functionalized_found
+        if not auto_functionalized_found:
+            raise AssertionError
         return counter
 
     def test_view_inplaced_functionalize_v2(self):
diff --git a/test/inductor/test_kernel_optimization.py b/test/inductor/test_kernel_optimization.py
index dce810fd2cd14..9ea57d88664f4 100644
--- a/test/inductor/test_kernel_optimization.py
+++ b/test/inductor/test_kernel_optimization.py
@@ -34,7 +34,8 @@ def compare_dict_tensors(self, ref_dict, res_dict, rtol=1e-3, atol=1e-3):
             return False
         for key1 in ref_dict:
             key2 = "_orig_mod." + key1
-            assert key2 in res_dict, f"{key1} does not exist in traced module"
+            if key2 not in res_dict:
+                raise AssertionError(f"{key1} does not exist in traced module")
             if not torch.allclose(ref_dict[key1], res_dict[key2], rtol=rtol, atol=atol):
                 return False
         return True
diff --git a/test/inductor/test_lookup_table.py b/test/inductor/test_lookup_table.py
index 3e69afe3599d1..f113b345e1667 100644
--- a/test/inductor/test_lookup_table.py
+++ b/test/inductor/test_lookup_table.py
@@ -2,7 +2,7 @@
 import re
 import unittest
 from functools import partial
-from typing import Any, Optional, Union
+from typing import Any
 from unittest.mock import patch
 
 import torch
@@ -59,7 +59,7 @@ class MockMMKernelInputs(MMKernelInputs):
     def __init__(
         self,
         tensors: list[torch.Tensor],
-        scalars: Optional[dict[str, Union[float, int]]] = None,
+        scalars: dict[str, float | int] | None = None,
         mat1_idx: int = -2,
         mat2_idx: int = -1,
     ):
@@ -81,7 +81,7 @@ def mnk_hinted(self) -> tuple[int, int, int]:
         return self.mnk_symbolic()  # pyre-ignore
 
     @property
-    def device_type(self) -> Optional[str]:
+    def device_type(self) -> str | None:
         return self.tensors[0].device.type
 
 
@@ -105,10 +105,10 @@ def tearDown(self):
 
     def create_mock_mm_kernel_inputs(
         self,
-        shapes: Optional[list[tuple[int, ...]]] = None,
+        shapes: list[tuple[int, ...]] | None = None,
         device: torch.device = torch.device("cuda"),
         dtype: torch.dtype = torch.float32,
-        scalars: Optional[dict[str, Union[float, int]]] = None,
+        scalars: dict[str, float | int] | None = None,
     ) -> MockMMKernelInputs:
         """Create MockMMKernelInputs with real tensors"""
         if shapes is None:
@@ -209,7 +209,8 @@ def test_successful_lookup_with_template_filtering(self):
             result = test_choices.lookup_template_configs(
                 kernel_inputs, "mm", ["triton"]
             )
-            assert result is not None, "Result should not be None"
+            if result is None:
+                raise AssertionError("Result should not be None")
             self.assertEqual(len(result["triton"]), 2)
             for config in result["triton"]:
                 self.assertNotIn("template_id", config)
@@ -217,7 +218,8 @@ def test_successful_lookup_with_template_filtering(self):
 
             # Test tma template filtering
             result = test_choices.lookup_template_configs(kernel_inputs, "mm", ["tma"])
-            assert result is not None, "Result should not be None"
+            if result is None:
+                raise AssertionError("Result should not be None")
             self.assertEqual(len(result["tma"]), 1)
             self.assertNotIn("template_id", result["tma"][0])
             self.assertEqual(result["tma"][0]["BLOCK_M"], 256)
@@ -226,7 +228,8 @@ def test_successful_lookup_with_template_filtering(self):
             result = test_choices.lookup_template_configs(
                 kernel_inputs, "mm", ["decompose_k"]
             )
-            assert result is not None, "Result should not be None"
+            if result is None:
+                raise AssertionError("Result should not be None")
             self.assertEqual(len(result["decompose_k"]), 1)
             self.assertNotIn("template_id", result["decompose_k"][0])
             self.assertEqual(result["decompose_k"][0]["k_split"], 4)
@@ -292,8 +295,10 @@ def test_multiple_calls_work(self):
                 kernel_inputs, "mm", ["triton"]
             )
             result2 = test_choices.lookup_template_configs(kernel_inputs, "mm", ["tma"])
-            assert result1 is not None, "Result1 should not be None"
-            assert result2 is not None, "Result2 should not be None"
+            if result1 is None:
+                raise AssertionError("Result1 should not be None")
+            if result2 is None:
+                raise AssertionError("Result2 should not be None")
             self.assertEqual(len(result1["triton"]), 1)
             self.assertEqual(len(result2["tma"]), 1)
 
@@ -302,8 +307,10 @@ def test_multiple_calls_work(self):
                 kernel_inputs, "mm", ["triton"]
             )
             result4 = test_choices.lookup_template_configs(kernel_inputs, "mm", ["tma"])
-            assert result3 is not None, "Result3 should not be None"
-            assert result4 is not None, "Result4 should not be None"
+            if result3 is None:
+                raise AssertionError("Result3 should not be None")
+            if result4 is None:
+                raise AssertionError("Result4 should not be None")
             self.assertEqual(len(result3["triton"]), 1)
             self.assertEqual(len(result4["tma"]), 1)
 
@@ -326,7 +333,8 @@ def test_batch_lookup_mixed_entries(self):
             result = test_choices.lookup_template_configs(
                 kernel_inputs, "mm", ["triton", "tma", "decompose_k"]
             )
-            assert result is not None, "Result should not be None"
+            if result is None:
+                raise AssertionError("Result should not be None")
 
             # Should have entries for triton and tma, but not decompose_k
             self.assertIn("triton", result)
@@ -377,7 +385,8 @@ def test_template_hash_checking(self, config_hash, template_hash, expected_kept)
             )
 
             if expected_kept:
-                assert result is not None, "Result should not be None"
+                if result is None:
+                    raise AssertionError("Result should not be None")
                 self.assertIn("triton", result)
                 self.assertEqual(len(result["triton"]), 1)
                 # template_hash should be removed from returned config
@@ -412,7 +421,8 @@ def test_template_hash_checking_disabled(self):
             )
 
             # Should keep config even with mismatching hash since checking is disabled
-            assert result is not None, "Result should not be None"
+            if result is None:
+                raise AssertionError("Result should not be None")
             self.assertIn("triton", result)
             self.assertEqual(len(result["triton"]), 1)
             # template_hash should still be removed from returned config
@@ -445,7 +455,8 @@ def test_template_hash_mixed_scenarios(self):
                 kernel_inputs, "mm", ["triton"], template_hash_map
             )
 
-            assert result is not None, "Result should not be None"
+            if result is None:
+                raise AssertionError("Result should not be None")
             self.assertIn("triton", result)
             # Should keep 2 configs: the one with correct hash and the one without hash
             self.assertEqual(len(result["triton"]), 2)
@@ -496,7 +507,8 @@ def test_hash_checking_disabled_edge_cases(self, config_hash, description):
             )
 
             # Should keep config regardless of hash validity since checking is disabled
-            assert result is not None, f"Result should not be None for {description}"
+            if result is None:
+                raise AssertionError(f"Result should not be None for {description}")
             self.assertIn(
                 "triton", result, f"Should have triton result for {description}"
             )
@@ -587,9 +599,10 @@ def _get_device_key(device):
             )
 
         if expected_found:
-            assert result is not None, (
-                f"Result should not be None when expected_found={expected_found}"
-            )
+            if result is None:
+                raise AssertionError(
+                    f"Result should not be None when expected_found={expected_found}"
+                )
             self.assertIn("triton", result, "Should have triton result when found")
             self.assertEqual(len(result["triton"]), 1, "Should have exactly 1 config")
             self.assertEqual(
@@ -636,7 +649,8 @@ def test_device_key_priority(self):
             )
 
             # Should get device-specific config (BLOCK_M=256), not device-agnostic (BLOCK_M=128)
-            assert result is not None, "Result should not be None"
+            if result is None:
+                raise AssertionError("Result should not be None")
             self.assertIn("triton", result)
             self.assertEqual(len(result["triton"]), 1)
             self.assertEqual(
@@ -862,22 +876,23 @@ def test_no_lookup_table_entry_autotune_modes(self, max_autotune):
         # Inline validation function
         def validate_choices(choices):
             if max_autotune:
-                assert len(choices) > 2, (
-                    f"Max-autotune should have >2 choices, got {len(choices)}"
-                )
-                assert any(isinstance(c, ExternKernelCaller) for c in choices), (
-                    "Should have ExternKernelCaller"
-                )
-                assert any(isinstance(c, TritonTemplateCaller) for c in choices), (
-                    "Should have TritonTemplateCaller"
-                )
+                if len(choices) <= 2:
+                    raise AssertionError(
+                        f"Max-autotune should have >2 choices, got {len(choices)}"
+                    )
+                if not any(isinstance(c, ExternKernelCaller) for c in choices):
+                    raise AssertionError("Should have ExternKernelCaller")
+                if not any(isinstance(c, TritonTemplateCaller) for c in choices):
+                    raise AssertionError("Should have TritonTemplateCaller")
             else:
-                assert len(choices) == 1, (
-                    f"No max-autotune should have 1 choice, got {len(choices)}"
-                )
-                assert isinstance(choices[0], ExternKernelCaller), (
-                    f"Should be ExternKernelCaller, got {type(choices[0])}"
-                )
+                if len(choices) != 1:
+                    raise AssertionError(
+                        f"No max-autotune should have 1 choice, got {len(choices)}"
+                    )
+                if not isinstance(choices[0], ExternKernelCaller):
+                    raise AssertionError(
+                        f"Should be ExternKernelCaller, got {type(choices[0])}"
+                    )
             return choices
 
         add_preprocessing_fn(validate_choices)
@@ -891,6 +906,10 @@ def validate_choices(choices):
     @fresh_cache()
     def test_valid_lookup_table_entry(self, operation):
         """Test when there's a valid entry for the operation"""
+        if operation == "addmm" and torch.version.hip:
+            self.skipTest(
+                "skipping on ROCm since https://github.com/pytorch/pytorch/issues/179955 didn't skip as expected"
+            )
         k = 256 if operation == "mm_plus_mm" else 64
         tensors = self.create_tensors(operation, k=k)
 
@@ -971,13 +990,19 @@ def test_bias_addmm_lookup_table_entry(self):
             partial(verify_choice_names, pattern="bias_addmm", expected_count=1)
         )
 
-        # Run with original unexpanded bias
+        # Run with expanded bias (stride[0] == 0) so the inductor sees
+        # bias_addmm-eligible inputs and the lookup key matches.
+        # Limit backends to ATEN so only the lookup table entry is selected.
         with inductor_config.patch(
-            {"max_autotune_gemm": True, "triton.autotune_cublasLt": True}
+            {
+                "max_autotune_gemm": True,
+                "triton.autotune_cublasLt": True,
+                "max_autotune_gemm_backends": "ATEN",
+            }
         ):
             model = UnifiedModel("addmm")
             compiled_model = torch.compile(model.to(self.device), mode="max-autotune")
-            compiled_model(bias_unexpanded, tensors[1], tensors[2])
+            compiled_model(expanded_bias, tensors[1], tensors[2])
 
     @unittest.skipIf(not has_triton_tma_device(), "Need TMA support")
     @fresh_cache()
diff --git a/test/inductor/test_loop_ordering.py b/test/inductor/test_loop_ordering.py
index 5d3dc6763ed07..2cb41ece0c424 100644
--- a/test/inductor/test_loop_ordering.py
+++ b/test/inductor/test_loop_ordering.py
@@ -23,7 +23,7 @@
 from torch._inductor.utils import is_big_gpu, run_and_get_code, sympy_index_symbol
 from torch._inductor.virtualized import ops, V
 from torch.testing import FileCheck
-from torch.testing._internal.common_cuda import PLATFORM_SUPPORTS_FP8
+from torch.testing._internal.common_cuda import PLATFORM_SUPPORTS_FP8, SM90OrLater
 from torch.testing._internal.common_utils import (
     instantiate_parametrized_tests,
     parametrize,
@@ -84,7 +84,8 @@ def _get_snode_body_sym_prefix(snode):
             prefix = str(var)[0]
             break
 
-        assert prefix
+        if not prefix:
+            raise AssertionError
         return prefix
 
     @staticmethod
@@ -416,6 +417,7 @@ def f(x):
         self.assertEqual(1, metrics.generated_kernel_count)
 
     @unittest.skipIf(not PLATFORM_SUPPORTS_FP8, "FP8 requires H100+ and MI300+")
+    @unittest.skipIf(not SM90OrLater, "sm89 errors out on this test")
     def test_fp8_cast_and_t(self):
         """
         This test repros the not able to fuses issue in
@@ -561,6 +563,48 @@ def f(x):
             ms = do_bench(lambda: opt_f(x))
             print(f"{ms=:.3f}")
 
+    def test_factored_vs_expanded_pw_numel_in_fused_group(self):
+        """
+        Regression test for https://github.com/pytorch/pytorch/issues/181563
+
+        Conv2d(kernel=1, stride=4) decomposes its output spatial size as
+        ``1 + (T - 1)//4``, so when those ranges are multiplied together for
+        a node body via ``sympy_product`` the result stays factored
+        (``s*(a+1)*(b+1)``). The fused-group's ``pointwise_numel`` for the
+        same value goes through ``SizeVarAllocator.simplify`` which calls
+        ``sympy.expand`` and yields the expanded form
+        (``s + s*a + s*b + s*a*b``). Sympy's structural ``==`` says these are
+        not equal, so without the fix the early-return guard in
+        ``get_pw_red_splits`` is missed and we fall through to an assert
+        that can never hold for a non-reduction body in a group whose
+        ``red_numel > 1``.
+        """
+
+        class Mod(nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.proj = nn.Conv2d(8, 8, kernel_size=1, stride=4, bias=False)
+
+            def forward(self, x):
+                a = self.proj(x)
+                # Outer-only pointwise sharing `a`'s sym shape; returning it
+                # forces materialization as its own SchedulerNode rather than
+                # being inlined into the softmax kernel.
+                bias = a[:, 0]
+                bias = bias * 2 + 1
+                a = a.permute(0, 2, 3, 1).contiguous()
+                a = a + bias.unsqueeze(-1)
+                b = a.softmax(dim=-1)  # red_numel = 8 in the fused group
+                return b, bias
+
+        mod = Mod().to(GPU_TYPE)
+        x = torch.randn(2, 8, 17, 19, device=GPU_TYPE)
+        torch._dynamo.mark_dynamic(x, 2)
+        torch._dynamo.mark_dynamic(x, 3)
+
+        with torch.no_grad():
+            self.do_acc_test(mod, x)
+
     @inductor_config.patch(
         {
             "max_autotune": True,
@@ -584,11 +628,19 @@ def f(x, y):
 
         out, code = run_and_get_code(f, x, y)
 
-        # well when benchmark_kernel flag is on, we have one more .run
-        # call in the benchmarking code.
-        FileCheck().check("def call(").check_count(
-            ".run(", 1 + int(inductor_config.benchmark_kernel), exactly=True
-        ).run(code[0])
+        FileCheck().check("def call(").run(code[0])
+        # Prologue fused with mm: 1 kernel. Unfused: 2 kernels (expand+add + mm).
+        # With benchmark_kernel, add 1 for the benchmarking code path.
+        base_expected = 1 + int(inductor_config.benchmark_kernel)
+        run_count = code[0].count(".run(")
+        self.assertGreaterEqual(
+            run_count, base_expected, "Expected at least one kernel launch"
+        )
+        self.assertLessEqual(
+            run_count,
+            base_expected + 1,
+            "Prologue fusion produces 1 kernel; unfused produces 2",
+        )
 
     @inductor_config.patch(
         {
@@ -761,7 +813,8 @@ def test_remapped_reads(self):
         from torch._inductor import tiling_utils
 
         def fn(nodes):
-            assert len(nodes) == 1
+            if len(nodes) != 1:
+                raise AssertionError(f"Expected 1 node, got {len(nodes)}")
             fused_norm_read_writes = tiling_utils.extract_normalized_read_writes(
                 nodes[0]
             )
@@ -1123,7 +1176,8 @@ def T(self, layout: str):
                 .unsqueeze(0)
             )
         else:
-            assert layout == "NHWC"
+            if layout != "NHWC":
+                raise AssertionError(f"Unexpected layout: {layout}")
             return torch.rand([1, SIZE_A, SIZE_B, SIZE_C], device=GPU_TYPE).to(
                 memory_format=torch.channels_last
             )
@@ -1220,7 +1274,8 @@ def fn(nodes):
             self.assertTrue(len(nodes) == 1)
 
             coalesce_analysis = tiling_utils.analyze_memory_coalescing(nodes[0])
-            assert coalesce_analysis is not None
+            if coalesce_analysis is None:
+                raise AssertionError
 
             reads = coalesce_analysis.norm_read_writes.reads
             writes = coalesce_analysis.norm_read_writes.writes
@@ -1325,6 +1380,70 @@ def embedding_1d(idx, w):
         self.assertEqual(out, expected)
 
 
+class TestSplitIterationRanges(MockSchedulerTest):
+    """Unit tests for SIMDKernel._split_iteration_ranges."""
+
+    def test_exact_match(self):
+        """Groups exactly match lengths — no splitting needed."""
+        from torch._inductor.codegen.simd import SIMDKernel
+
+        new_ranges, getters = SIMDKernel._split_iteration_ranges(
+            [sympy.Integer(4), sympy.Integer(8)],
+            [[sympy.Integer(4)], [sympy.Integer(8)]],
+        )
+        self.assertEqual(len(new_ranges), 2)
+        self.assertEqual(len(new_ranges[0]), 1)
+        self.assertEqual(len(new_ranges[1]), 1)
+
+    def test_two_way_split(self):
+        """A single large dimension splits across two groups."""
+        from torch._inductor.codegen.simd import SIMDKernel
+
+        new_ranges, getters = SIMDKernel._split_iteration_ranges(
+            [sympy.Integer(4), sympy.Integer(8)],
+            [[sympy.Integer(32)], []],
+        )
+        # 32 should split into 4 * 8 across the two groups
+        self.assertEqual(len(new_ranges), 2)
+
+    def test_groups_exhausted_raises_cant_split(self):
+        """When all groups are consumed but sizes remain, CantSplit is raised."""
+        from torch._inductor.codegen.simd import CantSplit, SIMDKernel
+
+        # groups=[1, 2, 2] can only absorb 2 sizes of 2 (consuming groups 1 and 2),
+        # leaving the third size=2 with no group to map to.
+        with self.assertRaises(CantSplit):
+            SIMDKernel._split_iteration_ranges(
+                [sympy.Integer(1), sympy.Integer(2), sympy.Integer(2)],
+                [[], [sympy.Integer(2), sympy.Integer(2), sympy.Integer(2)]],
+            )
+
+    def test_single_group_multiple_sizes(self):
+        """Multiple sizes fitting within a single group."""
+        from torch._inductor.codegen.simd import SIMDKernel
+
+        # groups=[8], lengths=[[2, 2, 2], []] — all 3 sizes fit in group 0
+        new_ranges, getters = SIMDKernel._split_iteration_ranges(
+            [sympy.Integer(8)],
+            [[sympy.Integer(2), sympy.Integer(2), sympy.Integer(2)], []],
+        )
+        self.assertEqual(len(new_ranges), 1)
+        self.assertEqual(len(new_ranges[0]), 3)
+
+    def test_size_one_skipped(self):
+        """Dimensions of size 1 produce a zero-constant getter."""
+        from torch._inductor.codegen.simd import SIMDKernel
+
+        new_ranges, getters = SIMDKernel._split_iteration_ranges(
+            [sympy.Integer(4)],
+            [[sympy.Integer(1), sympy.Integer(4)], []],
+        )
+        # Size-1 dim should not consume any range
+        self.assertEqual(len(getters[0]), 2)
+        # The first getter should return 0 for any input
+        self.assertEqual(getters[0][0]([sympy.Integer(99)]), sympy.Integer(0))
+
+
 class TestIndexInversion(TestCase):
     @classmethod
     def setUpClass(cls):
@@ -1340,7 +1459,10 @@ def _check_expr(self, expr, reconstruction, val_range):
         import numpy as np
         from sympy import lambdify
 
-        assert len(expr.free_symbols) == 1
+        if len(expr.free_symbols) != 1:
+            raise AssertionError(
+                f"Expected 1 free symbol, got {len(expr.free_symbols)}"
+            )
         p0 = next(iter(expr.free_symbols))
 
         def floordiv_replacement(a, b):
diff --git a/test/inductor/test_max_autotune.py b/test/inductor/test_max_autotune.py
index d94dd81b4673a..dbcd721227a13 100644
--- a/test/inductor/test_max_autotune.py
+++ b/test/inductor/test_max_autotune.py
@@ -12,10 +12,11 @@
 import time
 import unittest
 from collections.abc import Callable
-from typing import Optional
 from unittest import mock
 from unittest.mock import patch
 
+import sympy
+
 import torch
 import torch._inductor.async_compile
 from torch import multiprocessing as mp, nn
@@ -39,7 +40,8 @@
 from torch._inductor.graph import GraphLowering
 from torch._inductor.ir import Buffer, ChoiceCaller, FixedLayout, FlexibleLayout
 from torch._inductor.kernel.mm_plus_mm import aten_mm_plus_mm
-from torch._inductor.runtime.triton_heuristics import CachingAutotuner
+from torch._inductor.runtime.hints import DeviceProperties
+from torch._inductor.runtime.triton_heuristics import CachingAutotuner, pointwise
 from torch._inductor.scheduler import Scheduler
 from torch._inductor.select_algorithm import (
     add_feedback_saver,
@@ -49,6 +51,7 @@
     clear_feedback_savers,
     clear_preprocessing_fns,
     ExternKernelCaller,
+    NoValidChoicesError,
     TritonTemplate,
     TritonTemplateCaller,
 )
@@ -60,10 +63,12 @@
     CUDAPersistentTMATemplateConfigHeuristic,
     GemmConfig,
     get_shared_memory_checker_opts,
+    ROCmMMTemplateConfigHeuristic,
     XPUMMTemplateConfigHeuristic,
     XPUPersistentTMATemplateConfigHeuristic,
 )
-from torch.testing._internal.common_cuda import PLATFORM_SUPPORTS_FP8
+from torch.testing._internal.common_cuda import PLATFORM_SUPPORTS_FP8, SM90OrLater
+from torch.testing._internal.common_device_type import largeTensorTest
 from torch.testing._internal.common_utils import (
     instantiate_parametrized_tests,
     IS_WINDOWS,
@@ -93,7 +98,7 @@
 from torch._inductor.virtualized import V
 from torch.fx.experimental.proxy_tensor import make_fx
 from torch.testing import FileCheck
-from torch.testing._internal.common_utils import MI300_ARCH, runOnRocmArch, skipIfXpu
+from torch.testing._internal.common_utils import skipIfXpu
 from torch.testing._internal.inductor_utils import (
     get_func_call,
     get_kernel_launch,
@@ -175,6 +180,53 @@ def mm_plus_mm(a, b, c, d):
         ):
             torch.compile(mm_plus_mm, dynamic=dynamic)(a, b, c, d)
 
+    def test_max_autotune_includes_max_autotune_pointwise_configs(self):
+        """
+        Verify that `max_autotune` includes all pointwise configs from
+        `max_autotune_pointwise` for 1D, 2D, and 3D pointwise kernels.
+        """
+        # Fake device properties for this unit test; no CUDA state is needed.
+        triton_meta = {
+            "device": DeviceProperties(
+                type="cuda",
+                index=0,
+                multi_processor_count=1,
+                cc=80,
+                major=8,
+                max_threads_per_block=1024,
+                warp_size=32,
+            )
+        }
+        inductor_meta_common = {"autotune_pointwise": False}
+
+        for size_hints in (
+            {"x": 2048},
+            {"x": 128, "y": 256},
+            {"x": 64, "y": 64, "z": 64},
+        ):
+            with self.subTest(size_hints=size_hints):
+                max_autotune_configs = pointwise(
+                    size_hints,
+                    triton_meta=triton_meta,
+                    inductor_meta={**inductor_meta_common, "max_autotune": True},
+                    return_configs=True,
+                )
+                max_autotune_pointwise_configs = pointwise(
+                    size_hints,
+                    triton_meta=triton_meta,
+                    inductor_meta={
+                        **inductor_meta_common,
+                        "max_autotune_pointwise": True,
+                    },
+                    return_configs=True,
+                )
+
+                self.assertEqual(
+                    len(max_autotune_configs),
+                    len(max_autotune_pointwise_configs),
+                    "max_autotune should include all pointwise configs from max_autotune_pointwise",
+                )
+
     @unittest.skipIf(
         not has_triton_tma_device(), "Need device-side TMA support in Triton"
     )
@@ -252,6 +304,183 @@ def mm(a, b):
 
         torch.testing.assert_close(c_actual, c_expected, atol=1e-2, rtol=1e-2)
 
+    def test_use_triton_tma_template_rejects_descriptor_shapes_exceeding_int32(self):
+        from torch._inductor.utils import use_triton_tma_template
+
+        int32_max = torch.iinfo(torch.int32).max
+        mat1 = mock.Mock()
+        mat1.get_size.return_value = [128, int32_max + 1]
+        mat2 = mock.Mock()
+        mat2.get_size.return_value = [int32_max + 1, 128]
+        output_layout = mock.Mock()
+        output_layout.size = [128, 128]
+
+        with (
+            config.patch({"triton.enable_persistent_tma_matmul": True}),
+            mock.patch("torch._inductor.utils.can_use_tma", return_value=True),
+        ):
+            self.assertFalse(
+                use_triton_tma_template(mat1, mat2, output_layout=output_layout)
+            )
+
+        mat1.get_size.return_value = [128, int32_max]
+        mat2.get_size.return_value = [int32_max, 128]
+
+        with (
+            config.patch({"triton.enable_persistent_tma_matmul": True}),
+            mock.patch("torch._inductor.utils.can_use_tma", return_value=True),
+        ):
+            self.assertTrue(
+                use_triton_tma_template(mat1, mat2, output_layout=output_layout)
+            )
+
+    def test_descriptor_shape_fits_in_int32_uses_expected_guarding(self):
+        from torch._inductor.utils import _descriptor_shape_fits_in_int32
+
+        gm = make_fx(lambda: torch.zeros(2, 3))()
+        graph = GraphLowering(gm)
+        size0 = sympy.Symbol("s0", integer=True, nonnegative=True)
+        size1 = sympy.Symbol("s1", integer=True, nonnegative=True)
+        condition = sympy.And(
+            sympy.Le(size0, torch.iinfo(torch.int32).max),
+            sympy.Le(size1, torch.iinfo(torch.int32).max),
+        )
+
+        with V.set_graph_handler(graph):
+            with (
+                mock.patch.object(
+                    V.graph.sizevars, "statically_known_true", return_value=True
+                ) as statically_known_true,
+                mock.patch.object(
+                    V.graph.sizevars, "guard_or_false", return_value=True
+                ) as guard_or_false,
+            ):
+                self.assertTrue(
+                    _descriptor_shape_fits_in_int32(
+                        [128, size0, size1], add_guards=False
+                    )
+                )
+                statically_known_true.assert_called_once_with(condition)
+                guard_or_false.assert_not_called()
+
+            with (
+                mock.patch.object(
+                    V.graph.sizevars, "statically_known_true", return_value=True
+                ) as statically_known_true,
+                mock.patch.object(
+                    V.graph.sizevars, "guard_or_false", return_value=True
+                ) as guard_or_false,
+            ):
+                self.assertTrue(
+                    _descriptor_shape_fits_in_int32(
+                        [128, size0, size1], add_guards=True
+                    )
+                )
+                guard_or_false.assert_called_once_with(condition)
+                statically_known_true.assert_not_called()
+
+    @unittest.skipIf(not torch.version.hip, "ROCM only")
+    @parametrize("a_transposed", (False, True))
+    @parametrize("b_transposed", (False, True))
+    @parametrize("dynamic", (False, True))
+    def test_max_autotune_regular_mm_persistent(
+        self,
+        a_transposed: bool,
+        b_transposed: bool,
+        dynamic: bool,
+    ):
+        def mm(a, b):
+            a = a.repeat(8, 8)
+            b = b.repeat(8, 8)
+
+            if a_transposed:
+                a = a.T
+            if b_transposed:
+                b = b.T
+
+            return torch.mm(a, b)
+
+        M, N, K = 21, 31, 11
+        a = (
+            torch.randn(*((K, M) if a_transposed else (M, K)))
+            .to(torch.float16)
+            .to(GPU_TYPE)
+        )
+        b = (
+            torch.randn(*((N, K) if b_transposed else (K, N)))
+            .to(torch.float16)
+            .to(GPU_TYPE)
+        )
+
+        with config.patch(
+            {
+                "max_autotune": True,
+                "triton.enable_persistent_tma_matmul": "1",
+                "triton.native_matmul": False,
+                "test_configs.autotune_choice_name_regex": "mm_persistent",
+            }
+        ):
+            c_actual, code = run_and_get_code(torch.compile(mm, dynamic=dynamic), a, b)
+            c_expected = mm(a, b)
+
+        # Verify that we are using the non-TMA persistent implementation
+        FileCheck().check("triton_tem_fused_mm").check("NUM_SMS").run(code[0])
+
+        torch.testing.assert_close(c_actual, c_expected, atol=1e-2, rtol=1e-2)
+
+    @unittest.skipIf(not torch.version.hip, "ROCM only")
+    @parametrize("a_transposed", (False, True))
+    @parametrize("b_transposed", (False, True))
+    @parametrize("dynamic", (False, True))
+    def test_max_autotune_regular_addmm_persistent(
+        self,
+        a_transposed: bool,
+        b_transposed: bool,
+        dynamic: bool,
+    ):
+        def addmm(x, a, b):
+            x = x.repeat(8)
+            a = a.repeat(8, 8)
+            b = b.repeat(8, 8)
+
+            if a_transposed:
+                a = a.T
+            if b_transposed:
+                b = b.T
+
+            return torch.addmm(x, a, b)
+
+        M, N, K = 21, 31, 11
+        a = (
+            torch.randn(*((K, M) if a_transposed else (M, K)))
+            .to(torch.float16)
+            .to(GPU_TYPE)
+        )
+        b = (
+            torch.randn(*((N, K) if b_transposed else (K, N)))
+            .to(torch.float16)
+            .to(GPU_TYPE)
+        )
+        x = torch.randn(N).to(torch.float16).to(GPU_TYPE)
+
+        with config.patch(
+            {
+                "max_autotune": True,
+                "triton.enable_persistent_tma_matmul": "1",
+                "triton.native_matmul": False,
+                "test_configs.autotune_choice_name_regex": "mm_persistent",
+            }
+        ):
+            c_actual, code = run_and_get_code(
+                torch.compile(addmm, dynamic=dynamic), x, a, b
+            )
+            c_expected = addmm(x, a, b)
+
+        # Verify that we are using the non-TMA persistent implementation
+        FileCheck().check("triton_tem_fused_addmm").check("NUM_SMS").run(code[0])
+
+        torch.testing.assert_close(c_actual, c_expected, atol=1e-2, rtol=1e-2)
+
     @unittest.skipIf(
         not has_triton_tma_device(), "Need device-side TMA support in Triton"
     )
@@ -327,6 +556,75 @@ def mock_get_tma_workspace_arg(*args, **kwargs):
             mm_tma_heuristic.mm_configs = original_tma_configs
             mm_heuristic.mm_configs = original_mm_configs
 
+    @unittest.skipIf(
+        not has_triton_tma_device(), "Need device-side TMA support in Triton"
+    )
+    def test_workspace_size_bytes_accounts_for_dtype(self):
+        """workspace_size passed to benchmark request must be in bytes, not elements."""
+        import sympy
+
+        from torch._inductor.codegen.common import WorkspaceZeroMode
+        from torch._inductor.utils import get_dtype_size
+
+        count = 1024
+        dtype = torch.float32
+        expected_bytes = count * get_dtype_size(dtype)
+
+        fake_ws = WorkspaceArg(
+            count=sympy.Integer(count),
+            zero_mode=WorkspaceZeroMode.UNINITIALIZED,
+            device=torch.device(GPU_TYPE),
+            outer_name="test_ws",
+            dtype=dtype,
+        )
+
+        captured_sizes = []
+        orig_init = TritonBenchmarkRequest.__init__
+
+        def spy_init(self, *args, **kwargs):
+            captured_sizes.append(kwargs.get("workspace_size"))
+            orig_init(self, *args, **kwargs)
+
+        M, K, N = 64, 64, 64
+        a = torch.randn(M, K, device=GPU_TYPE, dtype=torch.bfloat16)
+        b = torch.randn(K, N, device=GPU_TYPE, dtype=torch.bfloat16)
+
+        mm_tma_heuristic = CUDAPersistentTMATemplateConfigHeuristic()
+        mm_heuristic = CUDAMMTemplateConfigHeuristic()
+        original_tma_configs = mm_tma_heuristic.mm_configs
+        original_mm_configs = mm_heuristic.mm_configs
+
+        try:
+            mm_heuristic.mm_configs = []
+            mm_tma_heuristic.mm_configs = [GemmConfig(128, 128, 64, 4, 8, group_m=8)]
+
+            with (
+                config.patch(
+                    {
+                        "max_autotune_gemm": True,
+                        "max_autotune_gemm_backends": "TRITON",
+                        "triton.enable_persistent_tma_matmul": True,
+                    }
+                ),
+                fresh_cache(),
+                patch(
+                    "torch._inductor.template_heuristics.triton.get_tma_workspace_arg",
+                    return_value=fake_ws,
+                ),
+                patch.object(TritonBenchmarkRequest, "__init__", spy_init),
+            ):
+                torch._dynamo.reset()
+                torch.compile(torch.mm, mode="max-autotune-no-cudagraphs")(a, b)
+
+        finally:
+            mm_tma_heuristic.mm_configs = original_tma_configs
+            mm_heuristic.mm_configs = original_mm_configs
+
+        ws_sizes = [s for s in captured_sizes if s is not None]
+        self.assertTrue(len(ws_sizes) > 0, "No workspace benchmark requests created")
+        for size in ws_sizes:
+            self.assertEqual(size, expected_bytes)
+
     @unittest.skipIf(
         not has_triton_tma_device(), "Need device-side TMA support in Triton"
     )
@@ -475,7 +773,7 @@ def mm(a, b):
         a = a.repeat(8, 8)
         b = b.repeat(8, 8)
 
-        torch._dynamo.mark_dynamic(a, 0)
+        torch._dynamo.maybe_mark_dynamic(a, 0)
 
         with config.patch(
             {
@@ -490,6 +788,41 @@ def mm(a, b):
 
         torch.testing.assert_close(c_actual, c_expected, atol=1e-2, rtol=1e-2)
 
+    @unittest.skipIf(
+        not has_triton_tma_device(), "Need device-side TMA support in Triton"
+    )
+    def test_persistent_tma_epilogue_fusion_store_cache(self):
+        # Regression test: when epilogue fusion runs with TMA store, the
+        # store_cache must be updated so that a subsequent epilogue load from
+        # the same buffer hits the cache. Otherwise remove_kernel_local_buffers
+        # strips the buffer pointer from the kernel signature, causing a
+        # NameError at Triton compile time.
+        def f(a, b):
+            a = a.repeat(8, 8)
+            b = b.repeat(8, 8)
+            mm = torch.mm(a, b)
+            return mm.relu()
+
+        M, N, K = 21, 31, 11
+        a = torch.randn(M, K, dtype=torch.float16, device=GPU_TYPE)
+        b = torch.randn(K, N, dtype=torch.float16, device=GPU_TYPE)
+
+        with config.patch(
+            {
+                "max_autotune": True,
+                "epilogue_fusion": True,
+                "triton.enable_persistent_tma_matmul": "1",
+                "triton.native_matmul": False,
+                "triton.enable_template_tma_store": True,
+                "triton.disallow_failing_autotune_kernels_TESTING_ONLY": True,
+                "test_configs.autotune_choice_name_regex": "mm_persistent_tma",
+            }
+        ):
+            actual = torch.compile(f)(a, b)
+            expected = f(a, b)
+
+        torch.testing.assert_close(actual, expected, atol=1e-2, rtol=1e-2)
+
     @parametrize("dynamic", (False, True))
     def test_max_autotune_regular_mm_zero_size_input(self, dynamic: bool):
         """
@@ -506,6 +839,34 @@ def mm(a, b):
         with config.patch({"max_autotune": True}):
             torch.compile(mm, dynamic=dynamic)(a, b)
 
+    @fresh_cache()
+    def test_addmm_1d_bias_no_reinterpret_tensor(self):
+        """
+        Verify that aten addmm with 1D bias does not wrap bias in reinterpret_tensor.
+        This ensures cublasLt uses its optimized bias epilogue (requires dim==1).
+        """
+
+        def addmm(x, a, b):
+            return torch.addmm(x, a, b)
+
+        x = torch.randn(100).to(GPU_TYPE)
+        a = torch.randn(100, 10).to(GPU_TYPE)
+        b = torch.randn(10, 100).to(GPU_TYPE)
+
+        with config.patch(
+            {
+                "max_autotune": True,
+                "max_autotune_gemm_backends": "ATEN",
+            }
+        ):
+            Y_compiled, code = run_and_get_code(torch.compile(addmm), x, a, b)
+            Y = addmm(x, a, b)
+            torch.testing.assert_close(Y_compiled, Y, atol=1e-2, rtol=1e-2)
+
+            # Verify addmm is called without reinterpret_tensor on bias
+            FileCheck().check("addmm").run(code[0])
+            self.assertNotIn("addmm(reinterpret_tensor", code[0])
+
     @unittest.skipIf(
         not has_triton_tma_device(), "Need device-side TMA support in Triton"
     )
@@ -638,7 +999,7 @@ def addmm(x, a, b):
         a = a.repeat(8, 8)
         b = b.repeat(8, 8)
 
-        torch._dynamo.mark_dynamic(a, 0)
+        torch._dynamo.maybe_mark_dynamic(a, 0)
 
         with config.patch(
             {
@@ -1016,8 +1377,12 @@ def f(x, y):
         f_c = torch.compile(mode="max-autotune-no-cudagraphs")(f)
         _, code = run_and_get_code(f_c, inps[0], inps[1])
         self.assertEqual(f_c(*inps), f(*inps), atol=0.03, rtol=0.25)
+        # x has multi-consumers (cat + add), so cat skips pointwise_cat
+        # and uses ConcatKernel.  cos is no longer fused into cat, giving
+        # one extra kernel launch compared to the single-consumer case.
+        count = 3 if (using_triton_mm or config.triton.native_matmul) else 2
         FileCheck().check(get_func_call()).check_count(
-            get_kernel_launch(), 2, exactly=True
+            get_kernel_launch(), count, exactly=True
         ).run(code[0])
 
         def f(x, y):
@@ -1290,7 +1655,8 @@ def f(x, y):
 
             expect = (f(x, y), x.grad, linear.weight.grad, linear.bias.grad)
             actual = (opt_f(x, y), x.grad, linear.weight.grad, linear.bias.grad)
-            assert same(expect, actual, tol=1e-2), f"ref:\n{expect}\nact:\n{actual}"
+            if not same(expect, actual, tol=1e-2):
+                raise AssertionError(f"ref:\n{expect}\nact:\n{actual}")
 
     @unittest.skipIf(
         config.cpp_wrapper, "decompose_k not supported for cpp_wrapper yet"
@@ -1364,7 +1730,7 @@ def check_divisors(code):
                 ).run(code[0])
             else:
                 FileCheck().check("extern_kernels.bmm_dtype").check_regex(
-                    "triton_.*_fused_0.run"
+                    "triton_.*_fused_.*.run"
                 ).check("decompose_k").run(code[0])
                 check_divisors(code)
                 torch.testing.assert_close(out, a @ b, atol=atol, rtol=rtol)
@@ -1378,7 +1744,7 @@ def check_divisors(code):
                 ).run(code[0])
             else:
                 FileCheck().check("extern_kernels.bmm_dtype").check_regex(
-                    "triton_.*_fused_mm_0.run"
+                    "triton_.*_fused_.*.run"
                 ).check("decompose_k").run(code[0])
                 check_divisors(code)
                 torch.testing.assert_close(
@@ -1515,7 +1881,7 @@ def f(a, b):
                 out.backward()
 
                 FileCheck().check("extern_kernels.bmm_dtype").check_regex(
-                    "triton_.*_fused_0.run"
+                    "triton_.*_fused_.*.run"
                 ).check("decompose_k").check_regex(r"s[0-9]+ = s[0-9]+").check_regex(
                     r"256\*s[0-9]+"
                 ).check_regex("s[0-9]+ = 8").run(
@@ -1851,20 +2217,32 @@ def f(a, b):
 
         a = torch.randn(2, 3, 4, device=GPU_TYPE, dtype=torch.float16)
         b = torch.randn(2, 4, 5, device=GPU_TYPE, dtype=torch.float16)
+        expected = torch.bmm(a.float(), b.float())
+        with config.patch(
+            max_autotune=False,
+            max_autotune_gemm_backends="ATEN",
+        ):
+            compiled_f = torch.compile(f)
+            out, code = run_and_get_code(compiled_f, a, b)
+            FileCheck().check("extern_kernels.bmm_dtype").run(code[0])
+            self.assertEqual(out, expected, atol=1e-3, rtol=1e-3)
+
+    @unittest.skipIf(config.cpp_wrapper, "out_dtype override not supported for AOTI")
+    def test_triton_bmm_out_dtype(self):
+        def f(a, b, out_dtype=torch.float32):
+            return torch.bmm(a, b, out_dtype=out_dtype)
+
+        a = torch.randn(2, 3, 4, device=GPU_TYPE, dtype=torch.float16)
+        b = torch.randn(2, 4, 5, device=GPU_TYPE, dtype=torch.float16)
+        expected = torch.bmm(a.float(), b.float())
         with config.patch(
             max_autotune=True,
             max_autotune_gemm_backends="TRITON",
         ):
             compiled_f = torch.compile(f)
-            with self.assertRaisesRegex(
-                torch._inductor.exc.InductorError,
-                r"LoweringException: NoValidChoicesError: No choices to select",
-            ):
-                out, code = run_and_get_code(compiled_f, a, b)
-
-        compiled_f = torch.compile(f)
-        out, code = run_and_get_code(compiled_f, a, b)
-        FileCheck().check("extern_kernels.bmm_dtype").run(code[0])
+            out, code = run_and_get_code(compiled_f, a, b, out_dtype=torch.float32)
+            FileCheck().check("triton_tem_fused_bmm").run(code[0])
+            self.assertEqual(out, expected, atol=1e-3, rtol=1e-3)
 
     def test_triton_template_generated_code_cache_key(self):
         generate_and_load_args = len(
@@ -1972,7 +2350,7 @@ def func_test1(x, y, z, m):
                         'num_consumer_groups':0,'num_buffers_warp_spec':0,'epilogue_fn_hash':'identity','tma_store':False,
                         'transpose_discontiguous_tensor_descriptors_override':None,
                         'kwargs':{'EVEN_K':False,'USE_FAST_ACCUM':False,'ACC_TYPE':'tl.float32',
-                        'BLOCK_M':16,'BLOCK_N':32,'BLOCK_K':16,'GROUP_M':8,'ALLOW_TF32':True},
+                        'OUT_DTYPE':'tl.float32','BLOCK_M':16,'BLOCK_N':32,'BLOCK_K':16,'GROUP_M':8,'ALLOW_TF32':True},
                         'hint_override':None,'triton_meta':None}"""
 
                 expected = expected.replace("cuda", GPU_TYPE)
@@ -2013,7 +2391,7 @@ def func_test1(x, y, z, m):
                     'layout':"[[s77,s94],[s94,1],torch.float32,device(type='cuda',index=0),0]",'num_consumer_groups':0,
                     'num_buffers_warp_spec':0,'epilogue_fn_hash':'identity','tma_store':False,
                     'transpose_discontiguous_tensor_descriptors_override':None,
-                    'kwargs':{'EVEN_K':False,'USE_FAST_ACCUM':False,'ACC_TYPE':'tl.float32','BLOCK_M':16,'BLOCK_N':32,
+                    'kwargs':{'EVEN_K':False,'USE_FAST_ACCUM':False,'ACC_TYPE':'tl.float32','OUT_DTYPE':'tl.float32','BLOCK_M':16,'BLOCK_N':32,
                     'BLOCK_K':16,'GROUP_M':8,'ALLOW_TF32':True},'hint_override':None,'triton_meta':None}"""
                 expected = expected.replace("cuda", GPU_TYPE)
                 self.assertExpectedInline(
@@ -2178,7 +2556,6 @@ def misses():
             self.assertEqual(misses(), 4)
 
     @fresh_cache()
-    @skipIfXpu
     @unittest.skipIf(
         config.cpp_wrapper, "decompose_k not supported for cpp_wrapper yet"
     )
@@ -2227,9 +2604,6 @@ def test_max_autotune_decompose_k_envvars(
                     self.assertTrue(decompose_count > 0)
                     self.assertTrue(decompose_count <= num_decompose_k_splits)
 
-    @unittest.skipIf(
-        TEST_WITH_ROCM, "exhaustive currently only thoroughly tested on NVIDIA"
-    )
     @unittest.skipIf(
         config.triton.native_matmul,
         "native matmul takes different tuning configs",
@@ -2247,11 +2621,15 @@ def f(a, b):
         with mock.patch(
             "torch._inductor.template_heuristics.registry.get_template_heuristic"
         ) as config_mock:
-            config_heuristics = (
-                XPUMMTemplateConfigHeuristic()
-                if GPU_TYPE == "xpu"
-                else CUDAMMTemplateConfigHeuristic()
-            )
+            # Create heuristic instance and modify it before setting as mock return value
+            # On ROCm, use ROCmMMTemplateConfigHeuristic; on XPU use XPUMMTemplateConfigHeuristic;
+            # otherwise use CUDAMMTemplateConfigHeuristic
+            if GPU_TYPE == "xpu":
+                config_heuristics = XPUMMTemplateConfigHeuristic()
+            elif torch.version.hip:
+                config_heuristics = ROCmMMTemplateConfigHeuristic()
+            else:
+                config_heuristics = CUDAMMTemplateConfigHeuristic()
 
             # Traditionally, this would be set of all possible configs
             # We mock out the code path for the sake of the unit test
@@ -2530,10 +2908,340 @@ def mm_with_cat(a, b1, b2, d):
             # should force fallback to extern_kernels.bmm
             FileCheck().check_not("triton_tem").run(code[0])
 
+    @parametrize("always_freeze", [True, False])
+    def test_mm_layout_freezing_behavior(self, always_freeze):
+        """Test that mm layout freezing behavior depends on always_freeze_layout.
 
-@instantiate_parametrized_tests
-class TestTemplateConfigPruning(TestCase):
-    """Test class for pruning logic in GEMM autotuning."""
+        When always_freeze_layout=True, FlexibleLayout should be frozen to FixedLayout.
+        When always_freeze_layout=False (default), FlexibleLayout should remain flexible
+        and use layout constraints instead.
+        """
+        from torch._inductor import ir
+        from torch._inductor.kernel.mm import mm_template
+        from torch._inductor.select_algorithm import TritonTemplateKernel
+
+        flexible_layout_called = False
+        orig_stride_call = TritonTemplateKernel.get_stride_and_maybe_freeze_layout
+
+        def tracking_get_stride(self, node):
+            nonlocal flexible_layout_called
+            flexible_layout = isinstance(node.data.layout, ir.FlexibleLayout)
+            result = orig_stride_call(self, node)
+            if flexible_layout:
+                flexible_layout_called = True
+                if always_freeze:
+                    if not isinstance(node.data.layout, ir.FixedLayout):
+                        raise AssertionError(
+                            f"Expected FixedLayout, got {type(node.data.layout)}"
+                        )
+                else:
+                    if not isinstance(node.data.layout, ir.FlexibleLayout):
+                        raise AssertionError(
+                            f"Expected FlexibleLayout, got {type(node.data.layout)}"
+                        )
+            return result
+
+        def fn(x, y):
+            # Add intermediate computation to create FlexibleLayout buffer
+            x = x + 1
+            return torch.mm(x, y)
+
+        M, K, N = 256, 128, 256
+        x = torch.randn(M, K, device=GPU_TYPE, dtype=torch.float16)
+        y = torch.randn(K, N, device=GPU_TYPE, dtype=torch.float16)
+
+        # Save original value to restore later
+        original_always_freeze = mm_template.always_freeze_layout
+
+        with (
+            fresh_cache(),
+            patch.object(
+                TritonTemplateKernel,
+                "get_stride_and_maybe_freeze_layout",
+                tracking_get_stride,
+            ),
+            config.patch({"test_configs.max_mm_configs": 1}),
+        ):
+            torch._dynamo.reset()
+            mm_template.always_freeze_layout = always_freeze
+            try:
+                # Temporarily set always_freeze_layout on mm_template
+                mm_template.always_freeze_layout = always_freeze
+                compiled_f = torch.compile(fn, mode="max-autotune")
+                compiled_f(x, y)
+            finally:
+                # Restore original value
+                mm_template.always_freeze_layout = original_always_freeze
+
+        self.assertTrue(flexible_layout_called)
+
+    @config.patch(
+        {
+            "max_autotune": True,
+            "test_configs.max_mm_configs": 1,
+        }
+    )
+    def test_deffered_layout_constraint_reintepret(self):
+        batch, m, k, n = 4608, 40, 112, 1119
+
+        # Shape: batch x m x k (contiguous)
+        a = torch.randn(batch, m, k, dtype=torch.bfloat16, device=GPU_TYPE)
+
+        padded_batch_stride = k * n + 48
+        b = torch.empty_strided(
+            size=(batch, k, n),
+            stride=(padded_batch_stride, 1, k),
+            dtype=torch.bfloat16,
+            device=GPU_TYPE,
+        )
+        b.copy_(torch.randn_like(b))
+        c = torch.randn(batch, k, m, dtype=torch.bfloat16, device=GPU_TYPE)
+
+        def fn(a, b, c):
+            # Apply a pointwise op to b to make it FlexibleLayout in Inductor
+            # This ensures Inductor doesn't treat it as a fixed/external layout
+            # Ends up double padding
+            b_flex = b + 0
+            return (
+                torch.bmm(a, b_flex).to(torch.float32),
+                b + 1.0,
+                torch.bmm(b_flex.permute(0, 2, 1), c).to(torch.float32),
+            )
+
+        with (
+            mock.patch(
+                "torch._inductor.autotune_process.run_autotune_in_subprocess",
+                mock_benchmark_choice_wrapper(aten_time=1.0, triton_time=0.1),
+            ),
+            mock.patch.object(
+                AlgorithmSelectorCache,
+                "benchmark_choice",
+                mock_benchmark_choice_wrapper(aten_time=1.0, triton_time=0.1),
+            ),
+        ):
+            compiled_fn = torch.compile(fn)
+
+            # Previously would CUDA IMA
+            _, code = run_and_get_code(compiled_fn, a, b, c)
+
+            FileCheck().check("triton_tem_fused").run(code[0])
+
+    @fresh_cache()
+    @skipIfXpu
+    @unittest.skipIf(TEST_WITH_ROCM, "Test requires CUDA")
+    @unittest.skipIf(
+        not SM90OrLater, "Requires SM90+ (H100/B200) for sufficient GPU memory"
+    )
+    @largeTensorTest("10 GB", device=GPU_TYPE)
+    def test_max_autotune_mm_large_input_tensor_int64_indexing(self):
+        """
+        Test mm with input tensor exceeding 2^31 elements.
+        Regression test for https://github.com/pytorch/pytorch/issues/171389
+        When input tensor storage exceeds 2^31 elements, tl.arange() must be
+        cast to INDEX_DTYPE (int64) to avoid integer overflow in pointer arithmetic.
+        """
+
+        def mm(a, b):
+            return torch.mm(a, b)
+
+        M, K, N = 1280, 65536, 65536
+        a = torch.randn(M, K, device=GPU_TYPE, dtype=torch.float16)
+        b = torch.randn(K, N, device=GPU_TYPE, dtype=torch.float16)
+
+        self.assertTrue(
+            b.numel() > 2**31 - 1,
+            f"Test requires tensor with >2^31 elements, got {b.numel()}",
+        )
+
+        with config.patch(
+            {
+                "max_autotune": True,
+                "max_autotune_gemm_backends": "TRITON",
+                "test_configs.autotune_choice_name_regex": r"^triton_mm_",
+            }
+        ):
+            result = torch.compile(mm)(a, b)
+
+        torch.testing.assert_close(result, torch.mm(a, b), rtol=1e-2, atol=1e-2)
+
+    @fresh_cache()
+    @skipIfXpu
+    @unittest.skipIf(TEST_WITH_ROCM, "Test requires CUDA")
+    @largeTensorTest("6 GB", device=GPU_TYPE)
+    def test_max_autotune_mm_large_storage_offset_i64_indexing(self):
+        """
+        Test mm with input having dynamic storage offset exceeding i32 range.
+        When a dynamic-shaped input is a slice with offset proportional to
+        the dynamic dim (e.g., offset=70000*s6), the ks parameter in the
+        triton template signature must use i64 to avoid overflow in pointer
+        arithmetic like `A = arg_A + 70000*ks0`.
+        """
+
+        def mm(x, w):
+            batch = x.shape[0] // 8
+            a = x[7 * batch :]
+            return torch.mm(a, w)
+
+        K, N = 10000, 32
+        batch = 32768
+        x = torch.randn(8 * batch, K, device=GPU_TYPE, dtype=torch.bfloat16)
+        w = torch.randn(K, N, device=GPU_TYPE, dtype=torch.bfloat16)
+
+        expected_offset = 7 * batch * K
+        self.assertTrue(
+            expected_offset > 2**31 - 1,
+            f"Test requires offset > i32_max, got {expected_offset}",
+        )
+
+        torch._dynamo.mark_dynamic(x, 0)
+
+        with config.patch(
+            {
+                "max_autotune": True,
+                "max_autotune_gemm_backends": "TRITON",
+                "test_configs.autotune_choice_name_regex": r"^triton_mm_",
+            }
+        ):
+            result = torch.compile(mm)(x, w)
+
+        a = x[7 * batch :]
+        torch.testing.assert_close(result, torch.mm(a, w), rtol=1e-2, atol=1e-2)
+
+    @fresh_cache()
+    @skipIfXpu
+    @unittest.skipIf(TEST_WITH_ROCM, "Test requires CUDA")
+    @unittest.skipIf(
+        not SM90OrLater, "Requires SM90+ (H100/B200) for sufficient GPU memory"
+    )
+    @largeTensorTest("10 GB", device=GPU_TYPE)
+    def test_max_autotune_mm_large_output_tensor_int32_overflow(self):
+        """
+        Test mm with output tensor exceeding 2^32 elements.
+        Regression test for https://github.com/pytorch/pytorch/issues/171389
+        When M * N >= 2^32, the early exit check `if M * N == 0` can overflow
+        to 0 in int32 arithmetic, causing the kernel to return immediately
+        with all-zero output.
+        """
+
+        def mm(a, b):
+            return torch.mm(a, b)
+
+        M, K, N = 65536, 32, 65536
+        a = torch.randn(M, K, device=GPU_TYPE, dtype=torch.float16)
+        b = torch.randn(K, N, device=GPU_TYPE, dtype=torch.float16)
+
+        self.assertTrue(
+            M * N >= 2**32,
+            f"Test requires M*N >= 2^32 for overflow, got {M * N}",
+        )
+
+        with config.patch(
+            {
+                "max_autotune": True,
+                "max_autotune_gemm_backends": "TRITON",
+                "test_configs.autotune_choice_name_regex": r"^triton_mm_",
+            }
+        ):
+            result = torch.compile(mm)(a, b)
+
+        torch.testing.assert_close(result, torch.mm(a, b), rtol=1e-2, atol=1e-2)
+
+    @fresh_cache()
+    @skipIfXpu
+    @unittest.skipIf(TEST_WITH_ROCM, "Test requires CUDA")
+    @unittest.skipIf(
+        not SM90OrLater, "Requires SM90+ (H100/B200) for sufficient GPU memory"
+    )
+    @largeTensorTest("10 GB", device=GPU_TYPE)
+    def test_max_autotune_mm_persistent_tma_large_input_tensor_int64_indexing(self):
+        """
+        Test persistent TMA mm with input tensor exceeding 2^31 elements.
+        Regression test for https://github.com/pytorch/pytorch/issues/171389.
+        Triton TMA descriptors require 32-bit block offsets even when the
+        surrounding kernel uses int64 indexing.
+        """
+
+        def mm(a, b):
+            return torch.mm(a, b)
+
+        M, K, N = 1280, 65536, 65536
+        a = torch.randn(M, K, device=GPU_TYPE, dtype=torch.float16)
+        b = torch.randn(K, N, device=GPU_TYPE, dtype=torch.float16)
+
+        self.assertTrue(
+            b.numel() > 2**31 - 1,
+            f"Test requires tensor with >2^31 elements, got {b.numel()}",
+        )
+
+        with config.patch(
+            {
+                "max_autotune": True,
+                "max_autotune_gemm_backends": "TRITON",
+                "triton.enable_persistent_tma_matmul": "1",
+                "triton.native_matmul": False,
+                "test_configs.autotune_choice_name_regex": "mm_persistent_tma",
+            }
+        ):
+            result = torch.compile(mm)(a, b)
+
+        torch.testing.assert_close(result, torch.mm(a, b), rtol=1e-2, atol=1e-2)
+
+    @fresh_cache()
+    @config.patch(
+        {
+            "max_autotune": True,
+            "test_configs.max_mm_configs": 1,
+        }
+    )
+    def test_deferred_layout_constraint_reinterpret_3d(self):
+        batch, m, k, n = 32, 40, 1053, 40
+        batch_stride = 42176
+
+        a = torch.randn(batch, m, k, dtype=torch.bfloat16, device=GPU_TYPE)
+        b = torch.empty_strided(
+            size=(batch, k * n),
+            stride=(batch_stride, 1),
+            dtype=torch.bfloat16,
+            device=GPU_TYPE,
+        )
+        b.copy_(torch.randn_like(b))
+        c = torch.randn(batch, n, m, dtype=torch.bfloat16, device=GPU_TYPE)
+        idx0 = torch.tensor([0], device=GPU_TYPE)
+        idx1 = torch.tensor([0], device=GPU_TYPE)
+        value = torch.zeros(1, dtype=torch.bfloat16, device=GPU_TYPE)
+
+        def fn(a, b, c, idx0, idx1, value):
+            # Mirror the 2D regression first: a simple pointwise op gives us a
+            # rank-2 FlexibleLayout. Then force that 2D value to realize before
+            # reinterpreting it as the rank-3 view from the original repro.
+            b_flex_2d = b + 0
+            b_flex_2d = aten.index_put.default(b_flex_2d, (idx0, idx1), value, False)
+            b_view_3d = b_flex_2d.view(batch, k, n)
+            return (
+                torch.bmm(a, b_view_3d).to(torch.float32),
+                b_flex_2d + 1.0,
+                torch.bmm(b_view_3d, c).to(torch.float32),
+            )
+
+        with (
+            mock.patch(
+                "torch._inductor.autotune_process.run_autotune_in_subprocess",
+                mock_benchmark_choice_wrapper(aten_time=1.0, triton_time=0.1),
+            ),
+            mock.patch.object(
+                AlgorithmSelectorCache,
+                "benchmark_choice",
+                mock_benchmark_choice_wrapper(aten_time=1.0, triton_time=0.1),
+            ),
+        ):
+            compiled_fn = torch.compile(fn)
+            _, code = run_and_get_code(compiled_fn, a, b, c, idx0, idx1, value)
+            FileCheck().check("triton_tem_fused").run(code[0])
+
+
+@instantiate_parametrized_tests
+class TestTemplateConfigPruning(TestCase):
+    """Test class for pruning logic in GEMM autotuning."""
 
     @classmethod
     def setUpClass(cls):
@@ -2823,8 +3531,8 @@ def no_lookup(
             op: str,
             inputs: str,
             benchmark: Callable[[Any], dict[ChoiceCaller, float]],
-            hint_override: Optional[int] = None,
-        ) -> Optional[dict[ChoiceCaller, float]]:
+            hint_override: int | None = None,
+        ) -> dict[ChoiceCaller, float] | None:
             if benchmark is not None:
                 return benchmark(choices)
 
@@ -2850,12 +3558,14 @@ def fake_benchmark_fn(*args, **kwargs):
                     ):
                         asc("test_call", fake_choices, [], Mock())
             for fake_choice in fake_choices:
-                assert fake_choice.thread_id is not None, (
-                    "Expected all ChoiceCaller's precompile method to have been called"
-                )
-                assert fake_choice.thread_id != main_thread_id, (
-                    "Expected all ChoiceCaller's precompile method to have been called on separate thread"
-                )
+                if fake_choice.thread_id is None:
+                    raise AssertionError(
+                        "Expected all ChoiceCaller's precompile method to have been called"
+                    )
+                if fake_choice.thread_id == main_thread_id:
+                    raise AssertionError(
+                        "Expected all ChoiceCaller's precompile method to have been called on separate thread"
+                    )
         finally:
             V.set_debug_handler(old_debug_handler)
 
@@ -2878,7 +3588,6 @@ def fn(a, b, c):
         self.assertEqual(counters["inductor"]["select_algorithm_precompile"], 0)
 
     @config.patch(autotune_local_cache=False, autotune_remote_cache=False)
-    @runOnRocmArch(MI300_ARCH)
     @unittest.skipIf(config.triton.native_matmul, "native matmul has counter 0")
     def test_precompilations(self):
         def fn(a, b, c):
@@ -2887,7 +3596,9 @@ def fn(a, b, c):
             return (a @ b) @ c
 
         fn_c = torch.compile(mode="max-autotune-no-cudagraphs")(fn)
-        inputs = [torch.rand([256, 256], device=GPU_TYPE) for _ in range(3)]
+        # Scale down so float16 doesn't overflow: rand [0,1) -> (a@b)@c has elements in [0, 256^2).
+        # float16 max ~65504, so we keep values in a safe range (e.g. scale by 1/256).
+        inputs = [torch.rand([256, 256], device=GPU_TYPE) / 256.0 for _ in range(3)]
 
         torch.testing.assert_close(fn_c(*inputs), fn(*inputs), atol=1e-2, rtol=1e-2)
 
@@ -3172,7 +3883,7 @@ def __str__(self) -> str:
 
 
 class TestTuningProcess(TestCase):
-    def check_healthy(self, p: TuningProcess, device: Optional[int] = None):
+    def check_healthy(self, p: TuningProcess, device: int | None = None):
         result = random.random()
         bmreq = _TestBenchmarkRequest(result, device=device)
         p.put(bmreq.benchmark)
@@ -3450,6 +4161,11 @@ def setUpClass(cls):
             )
         )
 
+    @classmethod
+    def tearDownClass(cls):
+        cls._stack.close()
+        super().tearDownClass()
+
     def check_code(self, code_str, num_kernels, num_allocs, num_deallocs):
         FileCheck().check(get_func_call()).check_count(
             get_kernel_launch(),
@@ -3943,9 +4659,29 @@ def setUpClass(cls):
             )
         )
 
+    @classmethod
+    def tearDownClass(cls):
+        cls._stack.close()
+        super().tearDownClass()
+
     @contextlib.contextmanager
-    def get_common_patches(self, async_compile: bool, persistent_tma: bool):
-        common_patches = (
+    def get_common_patches(
+        self,
+        async_compile: bool,
+        persistent_tma: bool,
+        *,
+        aten_time: float | None = None,
+        triton_time: float | None = None,
+        mock_n_spills: int | None = None,
+        mock_fused_n_regs: int | None = None,
+        mock_unfused_n_regs: int | None = None,
+        epilogue_runtime: float | None = None,
+    ):
+        from torch._inductor.autotune_process import TritonBenchmarkRequest
+        from torch._inductor.runtime.triton_heuristics import CachingAutotuner
+        from torch._inductor.scheduler import BaseSchedulerNode
+
+        common_patches = [
             config.patch(
                 {
                     "triton.enable_persistent_tma_matmul": persistent_tma,
@@ -3959,13 +4695,96 @@ def get_common_patches(self, async_compile: bool, persistent_tma: bool):
                 autotune_select_algorithm_wrapper_return_multi(),
             ),
             fresh_cache(),
-        )
+        ]
+
+        if aten_time is not None and triton_time is not None:
+            common_patches.extend(
+                [
+                    mock.patch.object(
+                        AlgorithmSelectorCache,
+                        "benchmark_choice",
+                        mock_benchmark_choice_wrapper(aten_time, triton_time),
+                    ),
+                    mock.patch(
+                        "torch._inductor.autotune_process.run_autotune_in_subprocess",
+                        mock_benchmark_choice_wrapper(aten_time, triton_time),
+                    ),
+                ]
+            )
+
+        if mock_n_spills is not None or mock_fused_n_regs is not None:
+            original_precompile = CachingAutotuner.precompile
+
+            def mock_precompile(self, *args, **kwargs):
+                original_precompile(self, *args, **kwargs)
+                for launcher in self.launchers:
+                    if mock_n_spills is not None:
+                        launcher.n_spills = mock_n_spills
+                    if mock_fused_n_regs is not None:
+                        launcher.n_regs = mock_fused_n_regs
+
+            common_patches.append(
+                mock.patch.object(CachingAutotuner, "precompile", mock_precompile)
+            )
+
+        if mock_unfused_n_regs is not None:
+            original_bmreq_precompile = TritonBenchmarkRequest.precompile
+
+            def mock_bmreq_precompile(self):
+                original_bmreq_precompile(self)
+                self.n_regs = mock_unfused_n_regs
+
+            common_patches.append(
+                mock.patch.object(
+                    TritonBenchmarkRequest, "precompile", mock_bmreq_precompile
+                )
+            )
+
+        if epilogue_runtime is not None:
+            common_patches.append(
+                mock.patch.object(
+                    BaseSchedulerNode,
+                    "_get_estimated_runtime",
+                    lambda node: epilogue_runtime,
+                )
+            )
+
         with contextlib.ExitStack() as stack:
             for p in common_patches:
                 stack.enter_context(p)
 
             yield
 
+    def _get_mm_inputs(self):
+        """Common matmul inputs for epilogue fusion tests."""
+        a = torch.randn(512, 1024, device=GPU_TYPE, dtype=torch.bfloat16)
+        b = torch.randn(1024, 2048, device=GPU_TYPE, dtype=torch.bfloat16)
+        return a, b
+
+    def _get_mm_with_epilogue_fn(self):
+        """Common function: matmul with type cast and add epilogue."""
+
+        def f(a, b):
+            return (a @ b).to(torch.float32) + 1.0
+
+        return f
+
+    @contextlib.contextmanager
+    def _setup_mm_heuristic(self, use_async_compile: bool):
+        """Setup MM heuristic with single GemmConfig and handle cleanup."""
+        mm_heuristic = CUDAMMTemplateConfigHeuristic()
+        original_mm_configs = mm_heuristic.mm_configs
+        gemm_config = GemmConfig(64, 64, 32, 2, 4, group_m=8)
+        mm_heuristic.mm_configs = [gemm_config]
+
+        if use_async_compile:
+            torch._inductor.async_compile.AsyncCompile.wait_pool_ready()
+
+        try:
+            yield
+        finally:
+            mm_heuristic.mm_configs = original_mm_configs
+
     @unittest.skipIf(not has_triton_tma_device(), "Need TMA support in Triton")
     @skipIfXpu(msg="Bad tma config can be covered by XPU TMA")
     @parametrize("use_async_compile", (True, False))
@@ -4093,23 +4912,6 @@ def test_template_epilogue_fusion_static_analysis(
         1. Register spillage (n_spills <= 8 required for fusion)
         2. Runtime comparison (epilogue_runtime + ms_min_choice > choice_timings[choice])
         """
-
-        def f(a, b):
-            return (a @ b).to(torch.float32) + 1.0
-
-        a = torch.randn(512, 1024, device=GPU_TYPE, dtype=torch.bfloat16)
-        b = torch.randn(1024, 2048, device=GPU_TYPE, dtype=torch.bfloat16)
-
-        mm_heuristic = CUDAMMTemplateConfigHeuristic()
-        original_mm_configs = mm_heuristic.mm_configs
-
-        # Use a single GemmConfig for the regular MM template
-        gemm_config = GemmConfig(64, 64, 32, 2, 4, group_m=8)
-        mm_heuristic.mm_configs = [gemm_config]
-
-        if use_async_compile:
-            torch._inductor.async_compile.AsyncCompile.wait_pool_ready()
-
         if test_case == "spills_reject":
             mock_n_spills = 100
             triton_time = 0.1
@@ -4128,46 +4930,27 @@ def f(a, b):
         elif test_case == "accept_with_aten_faster":
             mock_n_spills = 0
             triton_time = 0.1
-            aten_time = 0.100000001
+            aten_time = 0.09999
             expect_fusion = True
         else:
             raise RuntimeError("Invalid test case")
 
-        from torch._inductor.runtime.triton_heuristics import CachingAutotuner
-
-        original_precompile = CachingAutotuner.precompile
-
-        def mock_precompile(self, *args, **kwargs):
-            original_precompile(self, *args, **kwargs)
-            # After precompile, patch launchers with mock n_spills
-            for launcher in self.launchers:
-                launcher.n_spills = mock_n_spills
+        f = self._get_mm_with_epilogue_fn()
+        a, b = self._get_mm_inputs()
 
-        try:
-            with (
-                self.get_common_patches(use_async_compile, False),
-                mock.patch.object(
-                    AlgorithmSelectorCache,
-                    "benchmark_choice",
-                    mock_benchmark_choice_wrapper(aten_time, triton_time),
-                ),
-                mock.patch.object(
-                    CachingAutotuner,
-                    "precompile",
-                    mock_precompile,
-                ),
-                mock.patch(
-                    "torch._inductor.autotune_process.run_autotune_in_subprocess",
-                    mock_benchmark_choice_wrapper(aten_time, triton_time),
-                ),
+        with self._setup_mm_heuristic(use_async_compile):
+            with self.get_common_patches(
+                use_async_compile,
+                False,
+                aten_time=aten_time,
+                triton_time=triton_time,
+                mock_n_spills=mock_n_spills,
             ):
                 compiled_f = torch.compile(f)
-                out, code = run_and_get_code(compiled_f, a, b)
+                _, code = run_and_get_code(compiled_f, a, b)
 
                 if expect_fusion:
-                    FileCheck().check("triton_tem_fused__to_copy_add_mm_0.run").run(
-                        code[0]
-                    )
+                    FileCheck().check("triton_tem_fused__to_copy_add_mm_0").run(code[0])
                 elif triton_time < aten_time:
                     FileCheck().check("triton_tem_fused_mm").check(
                         "triton_poi_fused__to_copy"
@@ -4177,14 +4960,6 @@ def mock_precompile(self, *args, **kwargs):
                         "triton_poi_fused__to_copy"
                     ).run(code[0])
 
-                # TODO (paulzhan): Accuracy fails with cpp_warpper, ok for now
-                # as path is only hit with Async Pipelined Autotuning
-                # mainly for training
-                if not config.cpp_wrapper:
-                    torch.testing.assert_close(out, f(a, b), atol=1e-2, rtol=1e-2)
-        finally:
-            mm_heuristic.mm_configs = original_mm_configs
-
     @unittest.skipIf(
         not HAS_CUDA_AND_TRITON, "Scheduler static analysis only tested on cuda"
     )
@@ -4194,7 +4969,8 @@ def mock_precompile(self, *args, **kwargs):
     def test_template_epilogue_fusion_extra_reads(
         self, fuse_epilogue: bool, use_async_compile: bool
     ):
-        # Function has 2 extra reads for epilogue, same size as write
+        """Test epilogue fusion with extra reads (bias and scale tensors)."""
+
         def fn(x, w, bias, scale):
             out = torch.mm(x, w)
             return out * scale + bias
@@ -4206,15 +4982,6 @@ def fn(x, w, bias, scale):
         bias = torch.randn(512, 2048, device=GPU_TYPE, dtype=torch.bfloat16)
         scale = torch.randn(512, 2048, device=GPU_TYPE, dtype=torch.bfloat16)
 
-        mm_heuristic = CUDAMMTemplateConfigHeuristic()
-        original_mm_configs = mm_heuristic.mm_configs
-
-        gemm_config = GemmConfig(64, 64, 32, 2, 4, group_m=8)
-        mm_heuristic.mm_configs = [gemm_config]
-
-        if use_async_compile:
-            torch._inductor.async_compile.AsyncCompile.wait_pool_ready()
-
         epilogue_runtime = 0.5
         aten_time = 1.0
         unfused_time = aten_time + epilogue_runtime
@@ -4228,46 +4995,199 @@ def fn(x, w, bias, scale):
         else:
             triton_time = unfused_time - estimated_fused + 0.01
 
-        from torch._inductor.scheduler import BaseSchedulerNode
-
-        def mock_get_estimated_runtime(node):
-            return epilogue_runtime
-
-        try:
-            with (
-                self.get_common_patches(use_async_compile, False),
-                mock.patch.object(
-                    AlgorithmSelectorCache,
-                    "benchmark_choice",
-                    mock_benchmark_choice_wrapper(aten_time, triton_time),
-                ),
-                mock.patch.object(
-                    BaseSchedulerNode,
-                    "_get_estimated_runtime",
-                    mock_get_estimated_runtime,
-                ),
-                mock.patch(
-                    "torch._inductor.autotune_process.run_autotune_in_subprocess",
-                    mock_benchmark_choice_wrapper(aten_time, triton_time),
-                ),
+        with self._setup_mm_heuristic(use_async_compile):
+            with self.get_common_patches(
+                use_async_compile,
+                False,
+                aten_time=aten_time,
+                triton_time=triton_time,
+                epilogue_runtime=epilogue_runtime,
             ):
                 compiled_fn = torch.compile(fn)
-                out, code = run_and_get_code(compiled_fn, x, w, bias, scale)
+                _, code = run_and_get_code(compiled_fn, x, w, bias, scale)
 
                 if fuse_epilogue:
                     FileCheck().check("triton_tem_fused_add_mm_mul").run(code[0])
                 else:
-                    # Aten wins in unfused case
                     FileCheck().check_not("triton_tem").check(
                         "triton_poi_fused_add_mul"
                     ).run(code[0])
 
-                if not config.cpp_wrapper:
-                    torch.testing.assert_close(
-                        out, fn(x, w, bias, scale), atol=1e-1, rtol=1e-1
-                    )
-        finally:
-            mm_heuristic.mm_configs = original_mm_configs
+    @unittest.skipIf(
+        not HAS_CUDA_AND_TRITON, "Scheduler static analysis only tested on cuda"
+    )
+    @skipIfRocm(msg="Scheduler static analysis needs investigation on ROCm")
+    @parametrize(
+        "test_case",
+        [
+            "occupancy_ratio_accept",  # ratio > 0.5, accept via Branch B
+            "occupancy_ratio_reject",  # ratio <= 0.5, reject
+        ],
+    )
+    @parametrize("use_async_compile", (True, False))
+    def test_template_epilogue_fusion_occupancy_ratio(
+        self, test_case: str, use_async_compile: bool
+    ):
+        """
+        Test occupancy ratio branch of _fuse_epilogue.
+
+        Occupancy calculation (assuming regs_per_sm = 65536):
+        blocks = regs_per_sm // (n_regs * threads_per_block)
+        threads_per_block = num_warps * warp_size = 4 * 32 = 128
+        """
+        triton_time = 0.1
+        epilogue_runtime = triton_time
+
+        if test_case == "occupancy_ratio_accept":
+            # blocks_unfused=5, blocks_fused=3, ratio=0.6 > 0.5 -> accept
+            # aten slightly faster to verify fusion picks triton even when aten wins
+            mock_unfused_n_regs, mock_fused_n_regs = 100, 160
+            aten_time = 0.09
+            expect_fusion = True
+        elif test_case == "occupancy_ratio_reject":
+            # blocks_unfused=8, blocks_fused=2, ratio=0.25 < 0.5 -> reject
+            mock_unfused_n_regs, mock_fused_n_regs = 64, 200
+            aten_time = 0.11
+            expect_fusion = False
+        else:
+            raise RuntimeError("Invalid test case")
+
+        f = self._get_mm_with_epilogue_fn()
+        a, b = self._get_mm_inputs()
+
+        with self._setup_mm_heuristic(use_async_compile):
+            with self.get_common_patches(
+                use_async_compile,
+                False,
+                aten_time=aten_time,
+                triton_time=triton_time,
+                mock_n_spills=0,
+                mock_fused_n_regs=mock_fused_n_regs,
+                mock_unfused_n_regs=mock_unfused_n_regs,
+                epilogue_runtime=epilogue_runtime,
+            ):
+                compiled_f = torch.compile(f)
+                _, code = run_and_get_code(compiled_f, a, b)
+
+                if expect_fusion:
+                    FileCheck().check("triton_tem_fused__to_copy_add_mm").run(code[0])
+                else:
+                    FileCheck().check("triton_tem_fused_mm").check(
+                        "triton_poi_fused__to_copy"
+                    ).run(code[0])
+
+    @unittest.skipIf(
+        not HAS_CUDA_AND_TRITON, "Scheduler static analysis only tested on cuda"
+    )
+    @skipIfRocm(msg="Scheduler static analysis needs investigation on ROCm")
+    @parametrize(
+        "test_case",
+        [
+            "memory_bound_accept",
+            "memory_bound_reject_low_occupancy",
+        ],
+    )
+    @parametrize("use_async_compile", (True, False))
+    def test_template_epilogue_fusion_dominating_epilogue(
+        self, test_case: str, use_async_compile: bool
+    ):
+        """
+        Test memory-bound epilogue branch of _fuse_epilogue (Branch C).
+
+        When Branches A and B fail (low occupancy), fusion can still be accepted
+        if the epilogue is memory-bound (ms2 > 2*ms1) AND blocks_fused > 1.
+        """
+        triton_time = 0.1
+
+        if test_case == "memory_bound_accept":
+            # blocks_fused=2 > 1, ms2=0.3 > 2*ms1=0.2 -> Branch C accepts
+            # aten slightly faster to verify fusion picks triton even when aten wins
+            mock_unfused_n_regs, mock_fused_n_regs = 64, 256
+            aten_time = 0.09
+            epilogue_runtime = 0.3
+            expect_fusion = True
+        elif test_case == "memory_bound_reject_low_occupancy":
+            # blocks_fused=1, ms2=0.3 > 2*ms1=0.2 BUT blocks_fused <= 1 -> reject
+            mock_unfused_n_regs, mock_fused_n_regs = 64, 512
+            aten_time = 0.11
+            epilogue_runtime = 0.3
+            expect_fusion = False
+        else:
+            raise RuntimeError("Invalid test case")
+
+        f = self._get_mm_with_epilogue_fn()
+        a, b = self._get_mm_inputs()
+
+        with self._setup_mm_heuristic(use_async_compile):
+            with self.get_common_patches(
+                use_async_compile,
+                False,
+                aten_time=aten_time,
+                triton_time=triton_time,
+                mock_n_spills=0,
+                mock_fused_n_regs=mock_fused_n_regs,
+                mock_unfused_n_regs=mock_unfused_n_regs,
+                epilogue_runtime=epilogue_runtime,
+            ):
+                compiled_f = torch.compile(f)
+                _, code = run_and_get_code(compiled_f, a, b)
+
+                if expect_fusion:
+                    FileCheck().check("triton_tem_fused__to_copy_add_mm").run(code[0])
+                else:
+                    FileCheck().check("triton_tem_fused_mm").check(
+                        "triton_poi_fused__to_copy"
+                    ).run(code[0])
+
+    @unittest.skipIf(
+        not HAS_CUDA_AND_TRITON, "Scheduler static analysis only tested on cuda"
+    )
+    @parametrize("use_async_compile", (True, False))
+    def test_epilogue_prologue_fusion_cache_preserved(self, use_async_compile: bool):
+        def f(a, b):
+            # Prologue: pointwise operation on input 'a' before matmul
+            a_transformed = a + 1.0
+            # Matmul
+            mm_result = a_transformed @ b
+            # Epilogue: pointwise operation on output after matmul
+            return mm_result + 2.0
+
+        torch._dynamo.reset()
+
+        # Use float32 to avoid low precision heuristic rejection
+        a = torch.randn(512, 1024, device=GPU_TYPE, dtype=torch.float32)
+        b = torch.randn(1024, 2048, device=GPU_TYPE, dtype=torch.float32)
+
+        triton_time = 0.1
+        aten_time = float("inf")
+        epilogue_runtime = 0.05
+
+        # Always allow prologue fusion heuristics
+        def always_allow_prologue(*args):
+            return True
+
+        with self._setup_mm_heuristic(use_async_compile):
+            with self.get_common_patches(
+                use_async_compile,
+                False,
+                aten_time=aten_time,
+                triton_time=triton_time,
+                mock_n_spills=0,
+                epilogue_runtime=epilogue_runtime,
+            ):
+                # Enable prologue fusion so both epilogue and prologue are considered
+                with config.patch(prologue_fusion=True):
+                    # Bypass prologue heuristics that might reject the fusion
+                    with mock.patch.object(
+                        Scheduler,
+                        "check_prologue_fusion_heuristics_fusable",
+                        always_allow_prologue,
+                    ):
+                        compiled_f = torch.compile(f)
+                        # If the bug exists, this will fail with:
+                        # "ValueError: min() arg is an empty sequence"
+                        # when get_min_choice() is called during prologue fusion
+                        run_and_get_code(compiled_f, a, b)
 
 
 def simple_fn():
@@ -4278,16 +5198,11 @@ class TestMaxAutotuneAsyncPipelined(TestMaxAutotune, TestEpilogueFusionStaticAna
     """Tests for AsyncPipelinedAutotuning path."""
 
     SKIP_TESTS = {
-        "test_max_autotune_decompose_k": "Subgraphs not supported with async pipelining",
         "test_inf_timing": "Logs not consistent with async pipelined autotuning",
         "test_non_contiguous_input_mm_plus_mm": "Flaky on trunk",
         "test_autotune_device_guard": "Flaky on trunk",
         "test_template_bad_epilogue_fusion": "Benchmarking path is different",
-        # Contiguous transform tests - SubgraphChoiceCaller not supported with async pipelining
-        "test_max_autotune_contiguous_transform_mm": "Subgraphs not supported with async pipelining",
-        "test_max_autotune_contiguous_transform_addmm": "Subgraphs not supported with async pipelining",
-        "test_max_autotune_contiguous_transform_non_contiguous_second_matrix": "Subgraphs not supported with async pipelining",
-        "test_max_autotune_contiguous_transform_with_epilogue": "Subgraphs not supported with async pipelining",
+        "test_persistent_tma_epilogue_fusion_store_cache": "Epilogue fusion disabled in async pipelining",
         # XPU specific skips due to lack of multiprocess tensor reduction support (issue #170636)
         "test_max_autotune_addmm_persistent_tma": "No XPU implementation for multiprocess tensor reduction",
         "test_max_autotune_regular_mm_persistent_tma": "No XPU implementation for multiprocess tensor reduction",
@@ -4326,7 +5241,7 @@ def tearDown(self):
         # Clear the AsyncAutotuner cache to prevent test pollution
         AsyncAutotuner.choice_hash_to_future.clear()
 
-    @config.patch(max_autotune=True)
+    @config.patch(max_autotune_gemm=True)
     def test_async_autotuner_cache_same_inputs(self):
         M, K, N = 128, 64, 256
         M2, K2, N2 = 256, 128, 64
@@ -4382,6 +5297,27 @@ def test_autotune_process_pool_inactivity_shutdown(self):
         self.assertIsNone(pool_instance._timer)
         self.assertTrue(AutotuneProcessPool._shutdown_for_inactivity)
 
+    @patch(
+        "torch._inductor.autotune_process.AUTOTUNE_POOL_INACTIVITY_TIMEOUT",
+        2,
+    )
+    def test_autotune_process_pool_inactivity_shutdown_warmup_only(self):
+        """Test that the pool shuts down from inactivity even when only warmup is called."""
+        AutotuneProcessPool.shutdown_instance()
+        AutotuneProcessPool._shutdown_for_inactivity = False
+
+        pool_instance = AutotuneProcessPool.get_instance()
+        warmup_future = pool_instance.warm_up()
+        warmup_future.result()
+
+        self.assertIsNotNone(pool_instance._pool)
+
+        time.sleep(5)
+
+        self.assertIsNone(pool_instance._pool)
+        self.assertIsNone(pool_instance._timer)
+        self.assertTrue(AutotuneProcessPool._shutdown_for_inactivity)
+
     @patch(
         "torch._inductor.autotune_process.AUTOTUNE_POOL_INACTIVITY_TIMEOUT",
         2,
@@ -4429,6 +5365,52 @@ def matmul_fn(a, b):
         cache_entries_after_second = len(AsyncAutotuner.choice_hash_to_future)
         self.assertEqual(cache_entries_after_second, 0)
 
+    @config.patch(max_autotune_gemm=True)
+    def test_triton_error_precompilation_and_autotuning(self):
+        """
+        Test error handling when do_autotuning throws NoValidChoicesError
+        for Triton choices. The fallback to extern kernels should still work.
+        """
+
+        def mock_do_autotuning(*args, **kwargs):
+            raise NoValidChoicesError("Simulated: all Triton choices failed")
+
+        a = torch.randn(64, 32, device=GPU_TYPE, dtype=torch.bfloat16)
+        b = torch.randn(32, 64, device=GPU_TYPE, dtype=torch.bfloat16)
+
+        def mm_func(a, b, epilogue):
+            if epilogue:
+                return torch.mm(a, b) + 1.0
+            else:
+                return torch.mm(a, b)
+
+        def test_aten_chosen():
+            for epilogue in (True, False):
+                torch._dynamo.reset()
+                compiled_fn = torch.compile(mm_func)
+                out, code = run_and_get_code(compiled_fn, a, b, epilogue)
+                FileCheck().check_not("triton_tem").run(code[0])
+
+        with mock.patch.object(
+            AlgorithmSelectorCache, "do_autotuning", mock_do_autotuning
+        ):
+            test_aten_chosen()
+
+        original_start = AsyncAutotuner.start
+        bmreq = _TestBenchmarkRequest(
+            exc=RuntimeError("Simulated benchmark failure in subprocess")
+        )
+        bmreq.module_cache_key = ""
+
+        def mock_start(choices, inputs_key):
+            for choice in choices:
+                if isinstance(choice, TritonTemplateCaller):
+                    choice.bmreq = bmreq
+            return original_start(choices, inputs_key)
+
+        with mock.patch.object(AsyncAutotuner, "start", mock_start):
+            test_aten_chosen()
+
 
 if __name__ == "__main__":
     from torch._inductor.utils import is_big_gpu
diff --git a/test/inductor/test_max_autotune_blackwell.py b/test/inductor/test_max_autotune_blackwell.py
index 1f497e280b599..2df07b9dd62f8 100644
--- a/test/inductor/test_max_autotune_blackwell.py
+++ b/test/inductor/test_max_autotune_blackwell.py
@@ -280,6 +280,19 @@ def addmm(x, a, b):
                 "triton.enable_template_tma_store": tma_store,
                 "test_configs.autotune_choice_name_regex": "blackwell_ws_persistent_device_tma",
                 "test_configs.autotune_choice_desc_regex": epilogue_subtile_regex,
+                # If we dynamically disable pipelining,
+                # triton_blackwell_ws_persistent_device_tma template will
+                # be picked and then cause mis-aligned memory access.
+                # If we don't dynamically disable pipelining,
+                # this template is skipped and triton_tem_fused_addmm_repeat_3
+                # is used.
+                #
+                # Fundamentally we should fix the triton_blackwell_ws_persistent_device_tma template
+                # The flag below work around the problem.
+                #
+                # This only happens in fbcode https://www.internalfb.com/diff/D101855575.
+                # Can not repro it in OSS build.
+                "triton.dynamic_disable_pipelining": False,
             }
         ):
             c_actual, code = run_and_get_code(
diff --git a/test/inductor/test_mem_estimation.py b/test/inductor/test_mem_estimation.py
index 2f0ccfe6b284d..8055f2b5c153c 100644
--- a/test/inductor/test_mem_estimation.py
+++ b/test/inductor/test_mem_estimation.py
@@ -4,7 +4,6 @@
 import weakref
 from collections import Counter
 from collections.abc import Callable
-from typing import Optional
 
 import torch
 from torch._inductor.fx_passes.memory_estimator import (
@@ -29,7 +28,7 @@ def device_filter(device):
 
 
 class FakeTensorMemoryProfilerMode(TorchDispatchMode):
-    def __init__(self, device_filter: Optional[Callable[[torch.device], bool]] = None):
+    def __init__(self, device_filter: Callable[[torch.device], bool] | None = None):
         # counter of storage ids to live references
         self.storage_count: dict[int, int] = Counter()
         # live fake tensors
@@ -275,9 +274,10 @@ def foo(primals_1):
             # Alternative schedule: change which operation is the last use of zeros
             # Original: zeros_like, add, sum (zeros freed after sum)
             # Alternative: zeros_like, sum, add (zeros freed after add)
-            assert len(compute_nodes) == 5, (
-                f"Expected 3 compute nodes, got {len(compute_nodes)}"
-            )
+            if len(compute_nodes) != 5:
+                raise AssertionError(
+                    f"Expected 3 compute nodes, got {len(compute_nodes)}"
+                )
             reordered_nodes = [
                 compute_nodes[0],  # zeros_like: zeros = torch.zeros_like(primals_1)
                 compute_nodes[2],  # sum: sum_result = zeros.sum() (zeros still alive)
diff --git a/test/inductor/test_memory.py b/test/inductor/test_memory.py
index dc81d1b022e33..eb362b6535a79 100644
--- a/test/inductor/test_memory.py
+++ b/test/inductor/test_memory.py
@@ -37,10 +37,10 @@ class Foo(torch.nn.Module):
 
     def __init__(self):
         super().__init__()
-        self.w1 = torch.nn.Parameter(torch.ones(1, 10))
-        self.w2 = torch.nn.Parameter(torch.ones(1, 1))
+        self.w1 = torch.nn.Parameter(torch.ones(2, 10))
+        self.w2 = torch.nn.Parameter(torch.ones(2, 2))
         self.w3 = torch.nn.Parameter(torch.ones(10, 1))
-        self.w4 = torch.nn.Parameter(torch.ones(1, 10))
+        self.w4 = torch.nn.Parameter(torch.ones(2, 10))
 
     def forward(self, x):
         t1 = torch.matmul(x, self.w1)
@@ -61,7 +61,7 @@ def setUp(self):
 
         self.model = Foo().to(GPU_TYPE)
         M = 4096 if torch.version.hip is not None else 2048
-        self.inputs = torch.ones((M, 1), device=GPU_TYPE)
+        self.inputs = torch.ones((M, 2), device=GPU_TYPE)
         self.orig_reorder_method = memory.reorder_for_peak_memory
 
     @mock.patch.object(config, "reorder_for_peak_memory", True)
@@ -432,7 +432,8 @@ def replace_foreach(gm):
             nodes = gm.find_nodes(
                 op="call_function", target=torch.ops.aten._foreach_add.Scalar
             )
-            assert len(nodes) == 1
+            if len(nodes) != 1:
+                raise AssertionError
             node = nodes[0]
             nodes[0].target = torch.ops.aten._foreach_add_.Scalar
             for inp, out in zip(node.args[0], list(node.users.keys())):
@@ -454,6 +455,104 @@ def replace_foreach(gm):
             code = run_and_get_triton_code(foo, inp, inp2)
             FileCheck().check("allocated=['buf0']").run(code)
 
+    @unittest.skipUnless(TRITON_AVAILABLE, "Triton is not available")
+    def test_torch_cond_ordering_consistency(self):
+        small_sz, large_sz = 256, 1024
+
+        class MultiCondModel(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.register_buffer("large_buffer", torch.zeros(large_sz))
+                self.register_buffer("small_buffer1", torch.zeros(small_sz))
+                self.register_buffer("small_buffer2", torch.zeros(small_sz))
+                self.register_buffer("counter", torch.tensor(0, dtype=torch.long))
+
+            def forward(self, x: torch.Tensor) -> torch.Tensor:
+                condition = self.counter % 2 == 0
+
+                def true_fn_large(buf):
+                    return buf.clone() * 2.0
+
+                def false_fn_large(buf):
+                    return buf.clone()
+
+                def true_fn_small(buf):
+                    return buf.clone() * 2.0
+
+                def false_fn_small(buf):
+                    return buf.clone()
+
+                result_large = torch.cond(
+                    condition,
+                    lambda: true_fn_large(self.large_buffer),
+                    lambda: false_fn_large(self.large_buffer),
+                )
+                result_small1 = torch.cond(
+                    condition,
+                    lambda: true_fn_small(self.small_buffer1),
+                    lambda: false_fn_small(self.small_buffer1),
+                )
+                result_small2 = torch.cond(
+                    condition,
+                    lambda: true_fn_small(self.small_buffer2),
+                    lambda: false_fn_small(self.small_buffer2),
+                )
+                return (
+                    x + result_large.sum() + result_small1.sum() + result_small2.sum()
+                )
+
+        def extract_cond_order(code: str) -> list[tuple[str, int]]:
+            """
+            Extract the order of torch.cond operations from generated code.
+            Returns list of (cond_name, buffer_size) tuples in execution order.
+            """
+            import re
+
+            cond_order = []
+            # Look for patterns like "cond" or "cond_1" in the generated code
+            # along with their buffer sizes
+            lines = code.split("\n")
+            for i, line in enumerate(lines):
+                # Match true_graph buffer allocations which indicate cond execution
+                match = re.search(r"true_graph_(\d+)_buf0\s*=.*\((\d+),", line)
+                if match:
+                    cond_idx = int(match.group(1))
+                    buf_size = int(match.group(2))
+                    cond_order.append((f"cond_{cond_idx}", buf_size))
+            return cond_order
+
+        model = MultiCondModel().to(GPU_TYPE)
+        x = torch.randn(10, device=GPU_TYPE)
+
+        # Compile with base settings (no reordering)
+        torch._dynamo.reset()
+        with config.patch({"reorder_for_peak_memory": False}):
+            compiled_base = torch.compile(model)
+            code_base = run_and_get_triton_code(compiled_base, x)
+
+        base_order = extract_cond_order(code_base)
+
+        # Compile with reorder_for_peak_memory=True
+        torch._dynamo.reset()
+        with config.patch({"reorder_for_peak_memory": True}):
+            compiled_peak_mem = torch.compile(model)
+            code_peak_mem = run_and_get_triton_code(compiled_peak_mem, x)
+
+        peak_mem_order = extract_cond_order(code_peak_mem)
+
+        if base_order and peak_mem_order:
+            self.assertEqual(
+                base_order,
+                peak_mem_order,
+                msg=(
+                    f"torch.cond operations were reordered by reorder_for_peak_memory!\n"
+                    f"Base order: {base_order}\n"
+                    f"Peak memory order: {peak_mem_order}\n"
+                    f"This can cause NCCL hangs when torch.cond contains collective operations "
+                    f"because different ranks may execute collectives in different orders."
+                ),
+            )
+
 
 if __name__ == "__main__":
     from torch._inductor.test_case import run_tests
diff --git a/test/inductor/test_metrics.py b/test/inductor/test_metrics.py
index 009db1c2af1c1..cc03b68414718 100644
--- a/test/inductor/test_metrics.py
+++ b/test/inductor/test_metrics.py
@@ -56,7 +56,8 @@ def triton_red_fused_add_sum_2(in_out_ptr0, in_ptr0, xnumel, rnumel, XBLOCK : tl
 class TestMetrics(TestCase):
     def test_parse_proper_kernel_fn_code(self):
         proper_kernel_fn_code = metrics._parse_proper_kernel_fn_code(example_kernel)
-        assert proper_kernel_fn_code.startswith("def ")
+        if not proper_kernel_fn_code.startswith("def "):
+            raise AssertionError
 
     def test_count_args(self):
         proper_kernel_fn_code = metrics._parse_proper_kernel_fn_code(example_kernel)
diff --git a/test/inductor/test_minifier.py b/test/inductor/test_minifier.py
index 2da1a1386212b..6c4e7bb992c43 100644
--- a/test/inductor/test_minifier.py
+++ b/test/inductor/test_minifier.py
@@ -241,9 +241,11 @@ def forward(self, inp, *, k):
         )
 
     def _aoti_check_relu_repro(self, res):
-        assert res is not None
+        if res is None:
+            raise AssertionError("res is None")
         ep_file_path = res.get_exported_program_path()
-        assert ep_file_path is not None
+        if ep_file_path is None:
+            raise AssertionError("ep_file_path is None")
         gm = export_load(ep_file_path).module(check_guards=False)
         self.assertExpectedInline(
             str(gm.code).strip(),
diff --git a/test/inductor/test_minifier_isolate.py b/test/inductor/test_minifier_isolate.py
index 61cf6e3961133..f1862b65f9bce 100644
--- a/test/inductor/test_minifier_isolate.py
+++ b/test/inductor/test_minifier_isolate.py
@@ -8,7 +8,6 @@
     IS_MACOS,
     skipIfRocm,
     skipIfWindows,
-    skipIfXpu,
     TEST_WITH_ASAN,
 )
 from torch.testing._internal.inductor_utils import GPU_TYPE
@@ -40,11 +39,13 @@ def test_after_aot_cpu_runtime_error(self):
         self._test_after_aot_runtime_error("cpu", "")
 
     @skipIfRocm
-    @skipIfXpu
     @requires_gpu
     @inductor_config.patch("triton.inject_relu_bug_TESTING_ONLY", "runtime_error")
     def test_after_aot_gpu_runtime_error(self):
-        self._test_after_aot_runtime_error(GPU_TYPE, "device-side assert")
+        expected_error = (
+            "injected assert fail" if GPU_TYPE == "xpu" else "device-side assert"
+        )
+        self._test_after_aot_runtime_error(GPU_TYPE, expected_error)
 
 
 if __name__ == "__main__":
diff --git a/test/inductor/test_minifier_utils.py b/test/inductor/test_minifier_utils.py
index 80c773830b4af..2ac067a9e6791 100644
--- a/test/inductor/test_minifier_utils.py
+++ b/test/inductor/test_minifier_utils.py
@@ -1,5 +1,8 @@
 # Owner(s): ["module: inductor"]
+from unittest.mock import patch
+
 import torch
+from torch._dynamo.exc import UserError, UserErrorType
 from torch._dynamo.repro.aoti import (
     AOTIMinifierError,
     export_for_aoti_minifier,
@@ -9,6 +12,26 @@
 
 
 class MinifierUtilsTests(TestCase):
+    def test_invalid_output_user_error_classification(self):
+        class SimpleModel(torch.nn.Module):
+            def forward(self, x):
+                return x + 1
+
+        with patch(
+            "torch.export.export",
+            side_effect=UserError(
+                UserErrorType.INVALID_OUTPUT,
+                "mocked invalid output user error",
+            ),
+        ):
+            gm = export_for_aoti_minifier(
+                SimpleModel(),
+                (torch.randn(2, 2),),
+                strict=True,
+                skip_export_error=False,
+            )
+            self.assertTrue(gm is None)
+
     def test_invalid_output(self):
         class SimpleModel(torch.nn.Module):
             def __init__(self):
diff --git a/test/inductor/test_mix_order_reduction.py b/test/inductor/test_mix_order_reduction.py
index 5dc30015976b6..460dfe4199cfe 100644
--- a/test/inductor/test_mix_order_reduction.py
+++ b/test/inductor/test_mix_order_reduction.py
@@ -6,16 +6,17 @@
 import torch
 import torch._inductor.config as inductor_config
 import torch.nn.functional as F
+from torch import nn
 from torch._dynamo.utils import same
 from torch._inductor import metrics, utils
 from torch._inductor.scheduler import MixOrderReduction
 from torch._inductor.test_case import run_tests, TestCase
 from torch.testing import FileCheck
+from torch.testing._internal.common_device_type import largeTensorTest
 from torch.testing._internal.common_utils import (
     instantiate_parametrized_tests,
-    isRocmArchAnyOf,
-    MI200_ARCH,
     parametrize,
+    skipIfXpu,
 )
 from torch.testing._internal.inductor_utils import GPU_TYPE, HAS_GPU
 
@@ -81,7 +82,7 @@ def outer_red():
                 if len(shape) == 3:
                     return reduction_fn(x, dim=(0, 1))
                 else:
-                    assert len(shape) == 2
+                    assert len(shape) == 2  # noqa: S101
                     return reduction_fn(x, dim=0)
 
             if swap:
@@ -160,6 +161,30 @@ def f(x, y):
         # shared memory.
         self.assertEqual(metrics.codegen_mix_order_reduction, 0)
 
+    @inductor_config.patch(split_reductions=False)
+    def test_fuse_non_contiguous_pointwise(self):
+        if not inductor_config.triton.mix_order_reduction:
+            self.skipTest("Mix order reduction not enabled")
+
+        # Regression: mix-order reduction can appear valid pre-fusion, but a pointwise
+        # fused into one side can change access patterns and break the contiguity
+        # invariant. This test builds a reduction + pointwise path plus a second
+        # reduction, matching the shape/ordering pattern seen in the E2E failure.
+
+        def f(x):
+            # First reduction (contiguous on its own).
+            r1 = x.sum(dim=1)
+            # Pointwise depends on both reduced and unreduced data, so fusing it
+            # with the reduction can change access strides.
+            y = r1 * x[:, 0]
+            # Second reduction across a different dimension to trigger mix-order logic.
+            r2 = x.sum(dim=0)
+            return y, r2
+
+        # Large, asymmetric shape encourages mix-order reduction heuristics.
+        x = torch.randn(32768, 768, dtype=torch.float, device=GPU_TYPE)
+        self.check_numeric(f, (x,))
+
     @inductor_config.patch(coordinate_descent_tuning=True)
     def test_XBLOCK_coordest_tuning(self):
         """
@@ -192,7 +217,13 @@ def f(x):
         self.check_numeric(f, (x,))
         # We don't do mix order reduction for split redutions
         # with more than 2 layers
-        self.assertEqual(metrics.codegen_mix_order_reduction, 0)
+        self.assertEqual(
+            metrics.codegen_mix_order_reduction,
+            1
+            if inductor_config.triton.cooperative_reductions
+            or inductor_config.triton.force_cooperative_reductions
+            else 0,
+        )
 
     def test_independent_split_size(self):
         """
@@ -268,6 +299,8 @@ def f(x, y):
     @parametrize("max_autotune", (False, True))
     @parametrize("initial_xblock", (1, 2))
     @parametrize("add_1dim", (False, True))
+    # The test OOM in CI sometimes. Ask for more memory to make it stable.
+    @largeTensorTest("16GB", device=GPU_TYPE, inductor=True)
     def test_rms_norm_bwd(
         self,
         wdtype,
@@ -508,9 +541,6 @@ def test_rms_norm_sharing_weights(self, split_reductions, dtype):
         if not inductor_config.triton.mix_order_reduction:
             self.skipTest("Mix order reduction not enabled")
 
-        if dtype is torch.bfloat16 and isRocmArchAnyOf(MI200_ARCH):
-            self.skipTest("Currently failing on rocm mi200")
-
         def f(xs, w, eps):
             ys = []
             for x in xs:
@@ -528,13 +558,22 @@ def f(xs, w, eps):
         eps = 1e-5
 
         ref = f(xs, w, eps)
+
+        # use float64 to compute ref_grads for precision
+        # and cast back to original dtype
+        xs_f64 = [x.to(torch.float64) for x in xs]
+        w_f64 = w.to(torch.float64)
+        dys_f64 = [dy.to(torch.float64) for dy in dys]
+        ref_f64 = f(xs_f64, w_f64, eps)
+        ref_grads_f64 = torch.autograd.grad(ref_f64, [*xs_f64, w_f64], dys_f64)
+        ref_grads = [g.to(dtype) for g in ref_grads_f64]
+
         act = torch.compile(
             f,
             options={
                 "split_reductions": split_reductions,
             },
         )(xs, w, eps)
-        ref_grads = torch.autograd.grad(ref, [*xs, w], dys)
         act_grads, (wrapper,) = utils.run_and_get_code(
             lambda: torch.autograd.grad(act, [*xs, w], dys)
         )
@@ -637,6 +676,7 @@ def f(x):
         # the other is the piontwise kernel
         self.assertTrue(2, metrics.generated_kernel_count)
 
+    @patch("torch._inductor.scheduler.MixOrderReduction.is_split_reduction")
     @patch("torch._inductor.scheduler.MixOrderReduction.get_numel_rnumel")
     @patch("torch._inductor.scheduler.MixOrderReduction.get_common_read")
     @patch("torch._inductor.scheduler.MixOrderReduction.has_mix_reduction_orders")
@@ -645,6 +685,7 @@ def test_mix_order_reduction_non_strict_mode(
         mock_has_mix_reduction_orders: mock.Mock,
         mock_get_common_read: mock.Mock,
         mock_get_numel_rnumel: mock.Mock,
+        mock_is_split_reduction: mock.Mock,
     ):
         """
         This tests whether we can skip some non-critical checks
@@ -677,6 +718,7 @@ def test_mix_order_reduction_non_strict_mode(
         from sympy import Integer
 
         mock_get_numel_rnumel.return_value = (Integer(1), Integer(1))
+        mock_is_split_reduction.return_value = False
 
         mock_node_1.read_writes = mock.Mock()
         mock_node_1.read_writes.reads = []
@@ -698,7 +740,11 @@ def test_mix_order_reduction_non_strict_mode(
             self.assertFalse(MixOrderReduction.can_fuse(mock_node_1, mock_node_2))
         with (
             V.set_graph_handler(graph),
-            inductor_config.patch({"triton.mix_order_reduction_non_strict_mode": True}),
+            inductor_config.patch(
+                {
+                    "triton.mix_order_reduction_non_strict_mode": True,
+                }
+            ),
         ):
             self.assertTrue(MixOrderReduction.can_fuse(mock_node_1, mock_node_2))
 
@@ -726,6 +772,299 @@ def f(x):
         compile_metrics = torch._dynamo.utils._compilation_metrics
         self.assertEqual(len(compile_metrics), 1, "Don't recompile")
 
+    @skipIfXpu(msg="https://github.com/intel/intel-xpu-backend-for-triton/issues/6398")
+    def test_additive_rnumel(self):
+        """
+        Fix https://github.com/pytorch/pytorch/issues/176375
+        """
+        x = torch.randn(32768, 300, device=GPU_TYPE)
+        y = torch.randn(32768, 200, device=GPU_TYPE)
+        w = torch.randn(550, device=GPU_TYPE)
+        torch._dynamo.mark_dynamic(x, 1)
+        torch._dynamo.mark_dynamic(y, 1)
+
+        def f(x, y, w):
+            z = torch.cat((x, y), dim=1)
+
+            # Slice w_pool to match dynamic dim. This avoids a guard that would
+            # resolve s0+s1 to a concrete value (640), which is essential for
+            # keeping rnumel symbolic in the generated code.
+            w = w[: z.shape[-1]]  # [s0+s1], no concrete-equality guard
+            z = z * w
+            scale = z.sum()
+            z = z + scale
+            return z.sum(dim=0), z.sum(dim=1)
+
+        ref = f(x, y, w)
+        act = torch.compile(f)(x, y, w)
+
+        torch.testing.assert_close(ref, act, atol=1e-3, rtol=1e-3)
+
+        self.assertEqual(
+            inductor_config.triton.mix_order_reduction,
+            metrics.codegen_mix_order_reduction,
+        )
+
+    def test_additive_num_splits(self):
+        """
+        When the `num_splits` is an additive expression, a pair of
+        parenthesis is required.
+        """
+        torch.set_float32_matmul_precision("high")
+        linear1 = nn.Linear(1000, 1000).to(GPU_TYPE)
+        norm = nn.LayerNorm(1000).to(GPU_TYPE)
+
+        def model(x):
+            return norm(linear1(x[:, :-1].reshape(-1, 1000)))
+
+        compiled_model = torch.compile(model)
+        x = torch.randn(32, 200, 1000, device=GPU_TYPE)
+        torch._dynamo.mark_dynamic(x, 1)
+        compiled_model(x).sum().backward()
+
+        act = linear1.weight.grad, linear1.bias.grad
+
+        linear1.zero_grad()
+        norm.zero_grad()
+        model(x).sum().backward()
+        ref = linear1.weight.grad, linear1.bias.grad
+
+        torch.testing.assert_close(ref, act, atol=1e-3, rtol=1e-3)
+
+    @largeTensorTest("36GB", device=GPU_TYPE, inductor=True)
+    def test_out_of_shared_memory(self):
+        """
+        Fix https://github.com/pytorch/pytorch/issues/175250
+        """
+        if not inductor_config.triton.mix_order_reduction:
+            self.skipTest("Mix order reduction not enabled")
+
+        NUM_HEADS = 32
+        NUM_KV_HEADS = 8
+        HEAD_DIM = 128
+        HIDDEN_SIZE = NUM_HEADS * HEAD_DIM * 2
+        SEQ_LEN = 8192 * 2
+
+        def rotate_half(x):
+            x1 = x[..., : x.shape[-1] // 2]
+            x2 = x[..., x.shape[-1] // 2 :]
+            return torch.cat((-x2, x1), dim=-1)
+
+        def apply_rotary_pos_emb(q, k, cos, sin):
+            cos = cos[:, None, :, :]
+            sin = sin[:, None, :, :]
+            return (q * cos) + (rotate_half(q) * sin), (k * cos) + (
+                rotate_half(k) * sin
+            )
+
+        @torch.compile
+        def forward(
+            x,
+            q_proj,
+            k_proj,
+            v_proj,
+            o_proj,
+            embed_norm,
+            hidden_norm,
+            cos,
+            sin,
+        ):
+            batch, seq_len, _ = x.shape
+
+            # Eagle3 first layer: split concatenated [embeds, hidden] input
+            mid = x.shape[2] // 2
+            embeds, hidden = x.split(mid, dim=-1)
+
+            # Dual RMSNorm (pow, sum, div, mul in backward)
+            embeds = embed_norm(embeds)
+            hidden = hidden_norm(hidden)
+            residual = hidden
+
+            # Recombine for attention input (2 * HIDDEN_SIZE)
+            x = torch.cat([embeds, hidden], dim=-1)
+
+            # Adding a graph break here "fixes" the issue
+            # by breaking up the fused op
+            # torch._dynamo.graph_break()
+
+            # Q/K/V projections from 2*hidden_size input
+            q = q_proj(x).view(batch, seq_len, NUM_HEADS, HEAD_DIM).transpose(1, 2)
+            k = k_proj(x).view(batch, seq_len, NUM_KV_HEADS, HEAD_DIM).transpose(1, 2)
+            v = v_proj(x).view(batch, seq_len, NUM_KV_HEADS, HEAD_DIM).transpose(1, 2)
+
+            q, k = apply_rotary_pos_emb(q, k, cos, sin)
+            k = torch.repeat_interleave(k, NUM_HEADS // NUM_KV_HEADS, dim=1)
+            v = torch.repeat_interleave(v, NUM_HEADS // NUM_KV_HEADS, dim=1)
+            out = q.contiguous() @ k.contiguous().transpose(-2, -1) @ v.contiguous()
+
+            out = out.transpose(1, 2).contiguous().reshape(batch, seq_len, -1)
+            return o_proj(out) + residual
+
+        # Layers
+        embed_norm = nn.RMSNorm(HIDDEN_SIZE).to(GPU_TYPE)
+        hidden_norm = nn.RMSNorm(HIDDEN_SIZE).to(GPU_TYPE)
+        # Q/K/V project from 2*HIDDEN_SIZE (concatenated embeds + hidden)
+        q_proj = nn.Linear(2 * HIDDEN_SIZE, NUM_HEADS * HEAD_DIM, bias=False).to(
+            GPU_TYPE
+        )
+        k_proj = nn.Linear(2 * HIDDEN_SIZE, NUM_KV_HEADS * HEAD_DIM, bias=False).to(
+            GPU_TYPE
+        )
+        v_proj = nn.Linear(2 * HIDDEN_SIZE, NUM_KV_HEADS * HEAD_DIM, bias=False).to(
+            GPU_TYPE
+        )
+        o_proj = nn.Linear(NUM_HEADS * HEAD_DIM, HIDDEN_SIZE, bias=False).to(GPU_TYPE)
+
+        # Block mask - simple causal only
+        def causal_mask(_b, _h, q, kv):
+            return q >= kv
+
+        # Rotary embeddings (precomputed, no grad needed)
+        inv_freq = 1.0 / (
+            500000.0
+            ** (
+                torch.arange(0, HEAD_DIM, 2, dtype=torch.float32, device=GPU_TYPE)
+                / HEAD_DIM
+            )
+        )
+        pos = torch.arange(1, SEQ_LEN + 1, dtype=torch.float32, device=GPU_TYPE)
+        freqs = torch.outer(pos, inv_freq)
+        emb = torch.cat((freqs, freqs), dim=-1).unsqueeze(0)
+        cos, sin = emb.cos(), emb.sin()
+
+        # Input: 2*HIDDEN_SIZE to match split [embeds, hidden]
+        x = torch.randn(
+            1, SEQ_LEN, 2 * HIDDEN_SIZE, device=GPU_TYPE, requires_grad=True
+        )
+
+        out = forward(
+            x,
+            q_proj,
+            k_proj,
+            v_proj,
+            o_proj,
+            embed_norm,
+            hidden_norm,
+            cos,
+            sin,
+        )
+        loss = out.sum()
+        loss.backward()
+        self.assertTrue(metrics.codegen_mix_order_reduction > 1)
+
+    @inductor_config.patch("triton.mix_order_reduction", True)
+    @inductor_config.patch("triton.mix_order_reduction_non_strict_mode", True)
+    def test_dimension_refactoring_mismatch(self):
+        """
+        This reproduces an issue where `simplify_and_reorder()` produces a different
+        dimension factorization than `_original_ranges` used during fusion decision.
+        For example, fusion might see (13, 8472) but codegen sees (26, 4236) after
+        the reduction split optimization adds a factor of 2 to the pointwise dimensions.
+
+        We skip fusing split reductions for node1 in this case.
+        """
+
+        if not inductor_config.triton.mix_order_reduction:
+            self.skipTest("Mix order reduction not enabled")
+
+        # Reproduce the RMSNorm backward pattern that triggered the bug.
+        # The key is:
+        # - Shape (M, N) = (13, 8472) where N=8472 is large enough to trigger split
+        # - RMSNorm backward creates reductions along both dimensions
+        # - The feature dimension reduction (8472) gets split with factor 2
+        # - Mix order reduction tries to fuse these, but groups don't match after split
+        def f(x, w, eps):
+            orig_dtype = x.dtype
+            x = x.float()
+            # RMSNorm forward: y = x * rsqrt(mean(x^2) + eps) * w
+            rsqrt = torch.rsqrt((x * x).sum(dim=-1) / x.shape[-1] + eps)
+            y = (x * rsqrt[:, None] * w).to(dtype=orig_dtype)
+            return y
+
+        def fwd_bwd(compiled_f):
+            x.grad = None
+            w.grad = None
+            out = compiled_f(x, w, eps)
+            out.backward(dy)
+            return x.grad, w.grad
+
+        # Use the exact shape from the bug report: (13, 8472)
+        # 8472 = 2 * 4236, so split with factor 2 gives sub-reductions of 4236
+        M, N = 13, 8472
+        x = torch.randn(M, N, dtype=torch.float32, device=GPU_TYPE, requires_grad=True)
+        w = torch.randn(N, dtype=torch.float32, device=GPU_TYPE, requires_grad=True)
+        dy = torch.randn_like(x)
+        eps = 1e-5
+
+        opt_f = torch.compile(f)
+
+        ref = fwd_bwd(f)
+        act = fwd_bwd(opt_f)
+        torch.testing.assert_close(ref, act, atol=1e-3, rtol=1e-3)
+        self.assertGreaterEqual(metrics.codegen_mix_order_reduction, 0)
+
+    def test_keepdim_shape_mismatch(self):
+        """
+        Test that MixOrderReduction correctly handles keepdim=True reductions.
+
+        This test reproduces a bug where the final reduction in MixOrderReduction
+        generates `view(nsplit, rnumel).sum(dim=0)` which produces shape [rnumel],
+        but the expected output should be [1, rnumel] when keepdim=True.
+
+        The error manifests as:
+        RuntimeError: Function CompiledFunctionBackward returned an invalid gradient
+        at index N - got [2048] but expected shape compatible with [1, 2048]
+        """
+        if not inductor_config.triton.mix_order_reduction:
+            self.skipTest("Mix order reduction not enabled")
+
+        # Create a model that produces reductions with keepdim=True in the backward pass
+        # This pattern is common in normalization layers like RMSNorm/LayerNorm
+        class KeepDimReductionModel(nn.Module):
+            def __init__(self, hidden_size):
+                super().__init__()
+                # Using shape [1, hidden_size] to ensure keepdim=True in backward
+                self.weight = nn.Parameter(torch.ones(1, hidden_size))
+                self.bias = nn.Parameter(torch.zeros(1, hidden_size))
+
+            def forward(self, x):
+                # x: [batch, hidden_size]
+                # Normalization-like operation that produces keepdim reductions in backward
+                mean = x.mean(dim=-1, keepdim=True)
+                var = ((x - mean) ** 2).mean(dim=-1, keepdim=True)
+                x_norm = (x - mean) / (var + 1e-5).sqrt()
+                return x_norm * self.weight + self.bias
+
+        M, N = 32768, 2048  # Large batch to trigger mix order reduction
+        model = KeepDimReductionModel(N).to(GPU_TYPE)
+
+        x = torch.randn(M, N, dtype=torch.float32, device=GPU_TYPE, requires_grad=True)
+        dy = torch.randn_like(x)
+
+        def fwd_bwd(model, x, dy):
+            x.grad = None
+            model.zero_grad()
+            out = model(x)
+            out.backward(dy)
+            return x.grad, model.weight.grad, model.bias.grad
+
+        # Reference (eager)
+        ref = fwd_bwd(model, x, dy)
+
+        # Compiled with mix order reduction
+        compiled_model = torch.compile(model)
+        act = fwd_bwd(compiled_model, x, dy)
+
+        # Verify numerical correctness
+        self.assertTrue(same(ref, act, tol=1e-3), f"ref:\n{ref}\nact:\n{act}")
+
+        # Verify mix order reduction was used
+        self.assertGreater(
+            metrics.codegen_mix_order_reduction,
+            0,
+            "Mix order reduction should be triggered",
+        )
+
 
 @inductor_config.patch(
     "triton.mix_order_reduction", not inductor_config.triton.mix_order_reduction
diff --git a/test/inductor/test_mkldnn_pattern_matcher.py b/test/inductor/test_mkldnn_pattern_matcher.py
index 5797b789e0dea..a342e67deecbd 100644
--- a/test/inductor/test_mkldnn_pattern_matcher.py
+++ b/test/inductor/test_mkldnn_pattern_matcher.py
@@ -184,7 +184,10 @@ def _test_common(
             maybe_autocast = torch.amp.autocast(device_type=device, dtype=torch.float16)
             atol, rtol = 5e-2, 5e-2
         else:
-            assert check_autocast == torch.float32
+            if check_autocast != torch.float32:
+                raise AssertionError(
+                    f"Expected check_autocast to be torch.float32, got {check_autocast}"
+                )
             maybe_autocast = contextlib.nullcontext()
         if check_quantization:
             raise NotImplementedError("not supported, please migrate to torchao")
@@ -251,7 +254,10 @@ def _test_code_common(
             for op in include_ops:
                 self.assertIn(op, source_code)
             if num_include_ops is not None:
-                assert len(include_ops) == len(num_include_ops)
+                if len(include_ops) != len(num_include_ops):
+                    raise AssertionError(
+                        f"len(include_ops)={len(include_ops)} != len(num_include_ops)={len(num_include_ops)}"
+                    )
                 for i in range(len(include_ops)):
                     self.assertEqual(
                         source_code.count(include_ops[i]), num_include_ops[i]
@@ -267,7 +273,8 @@ def _test_code_common(
 
 class TestPatternMatcherGeneric(TestPatternMatcherBase):
     def _test_conv_unary_base(self, dim=4):
-        assert dim == 4 or dim == 5
+        if dim != 4 and dim != 5:
+            raise AssertionError(f"Expected dim to be 4 or 5, got {dim}")
 
         class M(torch.nn.Module):
             def __init__(
@@ -360,7 +367,8 @@ def test_conv3d_unary(self, device):
         self._test_conv_unary_base(dim=5)
 
     def _test_conv_transpose_unary_base(self, dim=4):
-        assert dim == 4 or dim == 5
+        if dim != 4 and dim != 5:
+            raise AssertionError(f"Expected dim to be 4 or 5, got {dim}")
 
         class M(torch.nn.Module):
             def __init__(
@@ -453,7 +461,8 @@ def test_conv_transpose3d_unary(self, device):
         self._test_conv_transpose_unary_base(dim=5)
 
     def _test_conv_binary_base(self, dim=4):
-        assert dim == 4 or dim == 5
+        if dim != 4 and dim != 5:
+            raise AssertionError(f"Expected dim to be 4 or 5, got {dim}")
 
         class M(torch.nn.Module):
             def __init__(
@@ -554,7 +563,8 @@ def test_conv3d_binary(self, device):
         self._test_conv_binary_base(dim=5)
 
     def _test_conv_binary_broadcast_shapes_base(self, dim=4):
-        assert dim == 4 or dim == 5
+        if dim != 4 and dim != 5:
+            raise AssertionError(f"Expected dim to be 4 or 5, got {dim}")
         torch.manual_seed(12345)
 
         class M(torch.nn.Module):
@@ -1554,7 +1564,6 @@ def forward(self, x):
 
     @unittest.skipIf(not TEST_MKL, "Test requires MKL")
     @xfailIfACL
-    @torch._dynamo.config.patch("inline_inbuilt_nn_modules", True)
     def test_reproduce_121253_issue_addmm_fusion_check(self):
         class Mod(torch.nn.Module):
             def __init__(self, weight, bias, beta, alpha):
diff --git a/test/inductor/test_mmdecomp.py b/test/inductor/test_mmdecomp.py
index 22a5d83324597..1b25587b5398b 100644
--- a/test/inductor/test_mmdecomp.py
+++ b/test/inductor/test_mmdecomp.py
@@ -2,11 +2,10 @@
 
 import math
 import unittest
-from typing import Union
 
 import torch
 from torch._inductor import config
-from torch._inductor.decomposition import mm
+from torch._inductor.decomposition import bmm as decomp_bmm, mm
 from torch._subclasses.fake_tensor import FakeTensorMode
 from torch.fx.experimental.symbolic_shapes import (
     DimDynamic,
@@ -16,7 +15,12 @@
 from torch.testing._internal.common_cuda import SM80OrLater
 from torch.testing._internal.common_device_type import instantiate_device_type_tests
 from torch.testing._internal.common_nn import NNTestCase
-from torch.testing._internal.common_utils import IS_WINDOWS, parametrize, run_tests
+from torch.testing._internal.common_utils import (
+    IS_WINDOWS,
+    parametrize,
+    run_tests,
+    TEST_XPU,
+)
 from torch.testing._internal.inductor_utils import GPU_TYPE, HAS_GPU
 
 
@@ -33,7 +37,7 @@
 
 
 def rand_math_tensor(
-    shape: tuple[Union[int, list[int]]],
+    shape: tuple[int | list[int]],
     device: str,
     dtype: torch.dtype,
     requires_grad: bool = False,
@@ -133,7 +137,8 @@ def test_simple_mm(self, device, dtype):
 
     @unittest.skipIf(not HAS_GPU, "GPU tests require triton")
     @parametrize(
-        "dtype", [torch.float, torch.bfloat16] if SM80OrLater else [torch.float]
+        "dtype",
+        [torch.float, torch.bfloat16] if SM80OrLater or TEST_XPU else [torch.float],
     )
     @parametrize("bs", [1, 2, 4, 10])
     def test_batched_mm(self, device, dtype, bs):
@@ -168,6 +173,80 @@ def test_bmm_batch2_last_dim_size_is_one(self, device):
 
         run_comp_nocomp(torch_bmm, t1, t2, rtol=rtol, atol=atol)
 
+    @config.patch(coordinate_descent_tuning=False)
+    def test_bmm_outer_product_k_is_one(self, device):
+        t1 = torch.randn(32, 8, 1, device=device)
+        t2 = torch.randn(32, 1, 256, device=device)
+        expected = torch.bmm(t1, t2)
+
+        out = decomp_bmm(t1, t2)
+
+        self.assertIsNot(out, NotImplemented)
+        self.assertEqual(expected, out)
+
+    @unittest.skipIf(not HAS_GPU, "GPU tests require triton")
+    def test_bmm_outer_product_k_is_one_with_unbacked_k(self, device):
+        if device == "cpu":
+            self.skipTest("unbacked symints require GPU fake tensors")
+
+        shape_env = ShapeEnv()
+        with FakeTensorMode(shape_env=shape_env):
+            b, m, n = [shape_env.create_unbacked_symint() for _ in range(3)]
+            lhs_k_unbacked, rhs_k_unbacked = [
+                shape_env.create_unbacked_symint() for _ in range(2)
+            ]
+
+            lhs_static_k = torch.empty((b, m, 1), device=device)
+            rhs_static_k = torch.empty((b, 1, n), device=device)
+            lhs_unbacked_k = torch.empty((b, m, lhs_k_unbacked), device=device)
+            rhs_unbacked_k = torch.empty((b, rhs_k_unbacked, n), device=device)
+
+            self.assertIsNot(
+                decomp_bmm(lhs_static_k, rhs_unbacked_k),
+                NotImplemented,
+            )
+            self.assertIsNot(
+                decomp_bmm(lhs_unbacked_k, rhs_static_k),
+                NotImplemented,
+            )
+            self.assertIs(
+                decomp_bmm(lhs_unbacked_k, rhs_unbacked_k),
+                NotImplemented,
+            )
+
+    @config.patch(coordinate_descent_tuning=False)
+    def test_bmm_outer_product_permuted_inputs(self, device):
+        B, M, N = 4, 8, 16
+
+        cases = [
+            # LHS: batch dim permuted
+            (
+                torch.randn(M, B, 1, device=device).permute(1, 0, 2),
+                torch.randn(B, 1, N, device=device),
+            ),
+            # RHS: batch dim permuted
+            (
+                torch.randn(B, M, 1, device=device),
+                torch.randn(N, B, 1, device=device).permute(1, 2, 0),
+            ),
+            # Both permuted
+            (
+                torch.randn(M, B, 1, device=device).permute(1, 0, 2),
+                torch.randn(N, B, 1, device=device).permute(1, 2, 0),
+            ),
+            # LHS: fully transposed [1, M, B] -> [B, M, 1]
+            (
+                torch.randn(1, M, B, device=device).permute(2, 1, 0),
+                torch.randn(B, 1, N, device=device),
+            ),
+        ]
+
+        for t1, t2 in cases:
+            expected = torch.bmm(t1, t2)
+            out = decomp_bmm(t1, t2)
+            self.assertIsNot(out, NotImplemented)
+            self.assertEqual(expected, out, exact_stride=True)
+
     @unittest.skipIf(not HAS_GPU, "GPU tests require triton")
     @parametrize("dtype", [torch.float, torch.bfloat16, torch.int])
     def test_some(self, device, dtype):
@@ -274,7 +353,9 @@ def test_dynamic_shape_mm(self, device, dtype):
 
 
 device_types = ("cpu", GPU_TYPE)
-instantiate_device_type_tests(TestDecomp, globals(), only_for=device_types)
+instantiate_device_type_tests(
+    TestDecomp, globals(), only_for=device_types, allow_xpu=True
+)
 
 if __name__ == "__main__":
     # We don't support torch.compile() on Windows
diff --git a/test/inductor/test_move_constructors_to_gpu.py b/test/inductor/test_move_constructors_to_gpu.py
index f304196f4a2ce..a3c58cdd7f45c 100644
--- a/test/inductor/test_move_constructors_to_gpu.py
+++ b/test/inductor/test_move_constructors_to_gpu.py
@@ -30,7 +30,8 @@ def _check_fn(self, func, expect_cpu, *args):
         out_compiled, code = run_and_get_code(torch.compile(func), *args)
         self.assertEqual(out_eager, out_compiled)
 
-        assert len(code) == 1
+        if len(code) != 1:
+            raise AssertionError
         if expect_cpu:
             FileCheck().check("cpp_fused").run(code[0])
         else:
diff --git a/test/inductor/test_mps_basic.py b/test/inductor/test_mps_basic.py
index 5c1691e581ccc..5d1d68d391a66 100644
--- a/test/inductor/test_mps_basic.py
+++ b/test/inductor/test_mps_basic.py
@@ -74,15 +74,12 @@ def fn(x):
             return x.tanh()
 
         result = fn(x)
-        assert torch.allclose(result[0], torch.tensor(-1.0, device="mps")), (
-            "tanh(-100) should be -1"
-        )
-        assert torch.allclose(result[-1], torch.tensor(1.0, device="mps")), (
-            "tanh(100) should be +1"
-        )
-        assert not torch.isnan(result).any(), (
-            "tanh should not produce NaN for large values"
-        )
+        if not torch.allclose(result[0], torch.tensor(-1.0, device="mps")):
+            raise AssertionError("tanh(-100) should be -1")
+        if not torch.allclose(result[-1], torch.tensor(1.0, device="mps")):
+            raise AssertionError("tanh(100) should be +1")
+        if torch.isnan(result).any():
+            raise AssertionError("tanh should not produce NaN for large values")
 
     def test_floor(self):
         self.common(lambda x: x.floor(), (torch.rand(1024),))
@@ -191,6 +188,63 @@ def fn(a, b):
             ),
         )
 
+    def test_sdpa_split_qkv(self):
+        # regression test for metal compiler bug where fused (x / A) % B
+        # produces wrong results, causing incorrect reads from non-contiguous.
+        n_head, n_embd, seq_len = 6, 384, 1024
+        x = torch.randn(16, seq_len, n_embd, device="mps")
+        c_attn = torch.nn.Linear(n_embd, 3 * n_embd).to("mps").eval()
+        qkv = c_attn(x)
+        q, k, v = qkv.split(n_embd, dim=2)
+        q = q.view(16, seq_len, n_head, n_embd // n_head).transpose(1, 2)
+        k = k.view(16, seq_len, n_head, n_embd // n_head).transpose(1, 2)
+        v = v.view(16, seq_len, n_head, n_embd // n_head).transpose(1, 2)
+
+        def fn(q, k, v):
+            return torch.nn.functional.scaled_dot_product_attention(
+                q, k, v, is_causal=True
+            )
+
+        self.common(fn, (q, k, v), atol=1e-4, rtol=1e-4, check_lowp=False)
+
+    def test_nested_masked_cat(self):
+        # Regression test for YOLOv3 compilation failure on MPS.
+        # See https://github.com/pytorch/pytorch/actions/runs/23477894502
+        # YOLOv3 detection heads do view/permute/clone, then in-place slice
+        # assignment (sigmoid+grid, exp*anchor, sigmoid) followed by cat across
+        # scales. The slice_scatter decomposition fused with cat produces nested
+        # ops.masked calls in Metal codegen. Without depth-aware variable
+        # prefixes, inner scoped variables shadow outer ones, causing:
+        #   "variable 'tmp_scoped_1' declared with deduced type 'auto'
+        #    cannot appear in its own initializer"
+        na, no = 3, 5
+
+        def head(p, grid, anchor_wh):
+            bs, _, ny, nx = p.shape
+            p = p.view(bs, na, no, ny, nx).permute(0, 1, 3, 4, 2).contiguous()
+            io = p.clone()
+            io[..., :2] = torch.sigmoid(io[..., :2]) + grid
+            io[..., 2:4] = torch.exp(io[..., 2:4]) * anchor_wh
+            torch.sigmoid_(io[..., 4:])
+            return io.view(bs, -1, no)
+
+        def fn(p1, p2, grid1, grid2, anchor_wh1, anchor_wh2):
+            return torch.cat(
+                [head(p1, grid1, anchor_wh1), head(p2, grid2, anchor_wh2)], dim=1
+            )
+
+        self.common(
+            fn,
+            (
+                torch.randn(1, na * no, 4, 4, device="mps"),
+                torch.randn(1, na * no, 8, 8, device="mps"),
+                torch.randn(1, 1, 4, 4, 2, device="mps"),
+                torch.randn(1, 1, 8, 8, 2, device="mps"),
+                torch.randn(1, na, 1, 1, 2, device="mps"),
+                torch.randn(1, na, 1, 1, 2, device="mps"),
+            ),
+        )
+
 
 class MPSBasicTestsAOTI(TestCase):
     def check_model(self, m, inp, dynamic_shapes=None):
@@ -199,7 +253,8 @@ def check_model(self, m, inp, dynamic_shapes=None):
         path = torch._inductor.aoti_compile_and_package(ep)
         m = torch._inductor.aoti_load_package(path)
         res = m(*inp)
-        assert torch.allclose(res, res2)
+        if not torch.allclose(res, res2):
+            raise AssertionError
 
     def test_add_mps(self):
         class M(torch.nn.Module):
diff --git a/test/inductor/test_multi_kernel.py b/test/inductor/test_multi_kernel.py
index 55f54756913db..f8fc1c3df607c 100644
--- a/test/inductor/test_multi_kernel.py
+++ b/test/inductor/test_multi_kernel.py
@@ -70,6 +70,7 @@ def make_cpp_wrapper_test(orig_test, **extra_args):
     """
 
     @config.patch("cpp_wrapper", True)
+    @config.patch("triton.autotune_at_compile_time", True)
     def fn(self):
         # The same kernel may have been compiled by previous tests with
         # cpp_wrapper disabled. Clear the cache so we go ahead to re-compile
@@ -111,7 +112,7 @@ def test_softmax(self, expect_multi_kernel=True):
     # TODO: bobrenjc93 to fix multi-kernel for ROCM
     @skipIfRocm
     @unittest.skipIf(not IS_BIG_GPU, "templates require big gpu")
-    @skipIfXpu(msg="https://github.com/intel/torch-xpu-ops/issues/2295")
+    @skipIfXpu(msg="driver issue, torch-xpu-ops: 2295")
     def test_triton_gemm(self):
         def fn(x, y):
             return x @ y
@@ -136,7 +137,7 @@ def fn(x, y):
         self.assertEqual(ref, act)
         self.assertTrue(_contains_size_hint_multi_kernel_code(wrapper_code))
 
-    @skipIfXpu(msg="https://github.com/intel/torch-xpu-ops/issues/2295")
+    @skipIfXpu(msg="driver issue, torch-xpu-ops: 2295")
     @requires_triton()
     # TODO: bobrenjc93 to fix multi-kernel for ROCM
     @skipIfRocm
diff --git a/test/inductor/test_native_matmul.py b/test/inductor/test_native_matmul.py
index 1932634ac5d3d..56e393534f340 100644
--- a/test/inductor/test_native_matmul.py
+++ b/test/inductor/test_native_matmul.py
@@ -10,7 +10,6 @@
 from torch._inductor.test_case import run_tests, TestCase
 from torch._inductor.utils import run_and_get_triton_code
 from torch.testing import FileCheck
-from torch.testing._internal.common_utils import skipIfXpu
 from torch.testing._internal.inductor_utils import GPU_TYPE, HAS_GPU
 
 
@@ -105,9 +104,6 @@ def f(x, y):
         self._check_equal(f, (x, y))
         self._check_code(f, (x, y), 1, 1)
 
-    @skipIfXpu(
-        msg="Intel triton issue: https://github.com/intel/intel-xpu-backend-for-triton/issues/5394"
-    )
     def test_3mm_add(self):
         def f(x, y, z, w, r, t):
             return x @ y + z @ w + r @ t
diff --git a/test/inductor/test_needs_exact_strides.py b/test/inductor/test_needs_exact_strides.py
index 2d636db3f88f1..dc7c3d55f967e 100644
--- a/test/inductor/test_needs_exact_strides.py
+++ b/test/inductor/test_needs_exact_strides.py
@@ -13,13 +13,15 @@
     IS_LINUX,
     parametrize,
 )
-from torch.testing._internal.inductor_utils import HAS_CUDA_AND_TRITON
+from torch.testing._internal.inductor_utils import HAS_GPU_AND_TRITON
 
 
 class TestNeedsExactStrides(InductorTestCase):
     @parametrize("dtype", [torch.float, torch.float8_e8m0fnu])
     def test_custom_op(self, dtype):
-        device = "cuda"  # float8_e8m0fnu errors on "cpu"
+        device = (
+            torch.accelerator.current_accelerator()
+        )  # float8_e8m0fnu errors on "cpu"
         x = torch.ones(4, 4, 2, 2, device=device, dtype=torch.float8_e8m0fnu)
         other = torch.ones(4, 4, 2, 2, device=device, dtype=torch.float8_e8m0fnu)
 
@@ -71,8 +73,10 @@ def impl2(x):
             )
 
             def impl(x, y):
-                assert x.transpose(2, 3).is_contiguous()
-                assert y.is_contiguous()
+                if not x.transpose(2, 3).is_contiguous():
+                    raise AssertionError
+                if not y.is_contiguous():
+                    raise AssertionError
                 return x.float() + y.float()
 
             def meta(x, y):
@@ -92,11 +96,12 @@ def f(x, other):
                     f_compile(x, other)
                 except TestPassed:
                     pass
-                assert called
+                if not called:
+                    raise AssertionError
 
 
 instantiate_parametrized_tests(TestNeedsExactStrides)
 
 if __name__ == "__main__":
-    if IS_LINUX and HAS_CUDA_AND_TRITON:
+    if IS_LINUX and HAS_GPU_AND_TRITON:
         run_tests(needs="filelock")
diff --git a/test/inductor/test_nv_universal_gemm.py b/test/inductor/test_nv_universal_gemm.py
index 0a529baf79f83..9a6512923db7b 100644
--- a/test/inductor/test_nv_universal_gemm.py
+++ b/test/inductor/test_nv_universal_gemm.py
@@ -12,11 +12,79 @@
     NVUniversalGemmHeuristics,
 )
 from torch._inductor.test_case import run_tests, TestCase
-from torch._inductor.utils import ensure_nv_universal_gemm_available, run_and_get_code
+from torch._inductor.utils import (
+    ceildiv,
+    ensure_nv_universal_gemm_available,
+    ensure_nvmatmul_heuristics_available,
+    run_and_get_code,
+)
 from torch.testing._internal.common_utils import (
     instantiate_parametrized_tests,
     parametrize,
 )
+from torch.utils._ordered_set import OrderedSet
+
+
+def _round_up(x, multiple):
+    return ((x + multiple - 1) // multiple) * multiple
+
+
+def _prep_k(K, scale_size):
+    """Prepare K dimension for swizzle requirements (round up ceildiv to multiple of 4)."""
+    return _round_up(ceildiv(K, scale_size), 4)
+
+
+def _create_tensor_with_layout(layout, rows, cols, dtype, device="cuda"):
+    """Create a tensor with the specified layout and dtype.
+
+    Supports float16, bfloat16, float8_e4m3fn, and float4_e2m1fn_x2.
+    """
+    is_fp4 = dtype == torch.float4_e2m1fn_x2
+    is_fp8 = dtype in (torch.float8_e4m3fn, torch.float8_e5m2)
+
+    def _make_flat(n):
+        if is_fp4:
+            return torch.randint(0, 256, (n,), device=device, dtype=torch.uint8).view(
+                torch.float4_e2m1fn_x2
+            )
+        elif is_fp8:
+            return torch.randint(-1, 2, (n,), device=device).to(dtype)
+        else:
+            return torch.randn(n, device=device, dtype=dtype)
+
+    if layout == "contiguous":
+        if is_fp4:
+            return torch.randint(
+                0, 256, (rows, cols), device=device, dtype=torch.uint8
+            ).view(torch.float4_e2m1fn_x2)
+        elif is_fp8:
+            return torch.randint(-1, 2, (rows, cols), device=device).to(dtype)
+        else:
+            return torch.randn(rows, cols, device=device, dtype=dtype)
+    elif layout == "aligned_offset":
+        storage = _make_flat(rows * cols + 512)
+        offset = 16 // storage.element_size()
+        return torch.as_strided(storage[offset:], (rows, cols), (cols, 1))
+    elif layout == "view":
+        return _make_flat(rows * cols).view(rows, cols)
+    elif layout == "padded":
+        row_pitch = cols + 8
+        storage = _make_flat(rows * row_pitch)
+        return torch.as_strided(storage, (rows, cols), (row_pitch, 1))
+    else:
+        raise ValueError(f"Unknown layout: {layout}")
+
+
+def _nvgemm_config(**overrides):
+    """Standard NVGEMM test config. Always disables ATen fallback."""
+    cfg = {
+        "max_autotune": True,
+        "max_autotune_gemm_backends": "NVGEMM",
+        "nvgemm_max_profiling_configs": 3,
+        "autotune_fallback_to_aten": False,
+    }
+    cfg.update(overrides)
+    return cfg
 
 
 # TODO(nikhilap): Remove Blackwell restriction once cutlass_api includes H100 kernels
@@ -34,59 +102,31 @@ class TestNVUniversalGemm(TestCase):
         (
             ("contiguous", "contiguous"),
             ("aligned_offset", "contiguous"),
+            ("contiguous", "aligned_offset"),
             ("contiguous", "view"),
             ("aligned_offset", "view"),
             ("padded", "contiguous"),
+            ("contiguous", "padded"),
         ),
     )
     def test_matmul(self, dtype, layout_a, layout_b):
         """Test matmul with various dtypes and tensor layouts.
 
-        These layouts test various alignment scenarios:
-        - contiguous/view/aligned_offset: Standard aligned layouts
-        - padded: Non-16-byte-aligned stride, Inductor pads to aligned size
-        M=513 tests that non-divisible M dimension works (only N and K must be divisible by 16).
+        M=513 tests that non-divisible M dimension works
+        (only N and K must be divisible by 16).
         """
         m, n, k = 513, 512, 512
-        device = "cuda"
 
         def matmul(a, b):
             return a @ b
 
-        def create_tensor_with_layout(layout, rows, cols):
-            """Create a tensor with the specified layout."""
-            if layout == "contiguous":
-                return torch.randn(rows, cols, device=device, dtype=dtype)
-            elif layout == "aligned_offset":
-                # Allocate bigger buffer than needed, use 16-byte aligned offset
-                # offset=128 elements * 2 bytes = 256 bytes (16-byte aligned)
-                storage = torch.randn(rows * cols + 512, device=device, dtype=dtype)
-                offset = 128
-                return torch.as_strided(storage[offset:], (rows, cols), (cols, 1))
-            elif layout == "view":
-                storage = torch.randn(rows * cols, device=device, dtype=dtype)
-                return storage.view(rows, cols)
-            elif layout == "padded":
-                # Simulate row pitch > cols with non-16-byte-aligned stride
-                # row_stride = cols + 8 = 520, 520 * 2 bytes = 1040 bytes (not 16-byte aligned)
-                row_pitch = cols + 8
-                storage = torch.randn(rows * row_pitch, device=device, dtype=dtype)
-                return torch.as_strided(storage, (rows, cols), (row_pitch, 1))
-
-        a = create_tensor_with_layout(layout_a, m, k)
-        b = create_tensor_with_layout(layout_b, k, n)
-
+        a = _create_tensor_with_layout(layout_a, m, k, dtype)
+        b = _create_tensor_with_layout(layout_b, k, n, dtype)
         expected = matmul(a, b)
 
         torch._dynamo.reset()
 
-        with config.patch(
-            {
-                "max_autotune": True,
-                "max_autotune_gemm_backends": "NVGEMM",
-                "nvgemm_max_profiling_configs": 3,
-            }
-        ):
+        with config.patch(_nvgemm_config()):
             compiled_fn = torch.compile(matmul)
             result = compiled_fn(a, b)
 
@@ -106,7 +146,6 @@ def test_unaligned_base_pointer_rejected(self):
         def matmul(a, b):
             return a @ b
 
-        # Create tensor with unaligned base pointer
         # offset=117 elements * 2 bytes = 234 bytes (NOT 16-byte aligned)
         storage = torch.randn(m * k + 512, device=device, dtype=dtype)
         a = torch.as_strided(storage[117:], (m, k), (k, 1))
@@ -114,13 +153,7 @@ def matmul(a, b):
 
         torch._dynamo.reset()
 
-        with config.patch(
-            {
-                "max_autotune": True,
-                "max_autotune_gemm_backends": "NVGEMM",
-                "nvgemm_max_profiling_configs": 3,
-            }
-        ):
+        with config.patch(_nvgemm_config()):
             compiled_fn = torch.compile(matmul)
             with self.assertRaisesRegex(
                 Exception, "NoValidChoicesError|no valid choice"
@@ -132,34 +165,23 @@ def test_reinterpret_view_from_slice(self, dtype):
         """Test that sliced tensors (creating ReinterpretViews) work correctly.
 
         When tensors are slices of a shared buffer (e.g., from a fused projection),
-        they become ReinterpretViews with non-contiguous strides. NVIDIA Universal GEMM must
-        handle these correctly.
+        they become ReinterpretViews with non-contiguous strides.
         """
         m, n, k = 512, 512, 512
         device = "cuda"
 
         def fn(x, weight):
-            # Fused projection creates a single large output
             projected = x @ weight  # (m, 2*n)
-            # Slicing creates ReinterpretViews
             a, b = projected.split(n, dim=1)  # Each is (m, n)
             return a @ b.t()  # (m, m)
 
         x = torch.randn(m, k, device=device, dtype=dtype)
-        # Weight projects to 2*n so we can split
         weight = torch.randn(k, 2 * n, device=device, dtype=dtype)
-
         expected = fn(x, weight)
 
         torch._dynamo.reset()
 
-        with config.patch(
-            {
-                "max_autotune": True,
-                "max_autotune_gemm_backends": "NVGEMM",
-                "nvgemm_max_profiling_configs": 3,
-            }
-        ):
+        with config.patch(_nvgemm_config()):
             compiled_fn = torch.compile(fn)
             result = compiled_fn(x, weight)
 
@@ -169,8 +191,7 @@ def test_workspace_allocation(self):
         """Test that workspace allocation works correctly.
 
         Since no current CUTLASS kernels require a workspace, we mock the
-        kernel.get_workspace_size method to return a non-zero value. This
-        exercises the workspace allocation/deallocation code paths.
+        kernel.get_workspace_size method to return a non-zero value.
         """
         m, n, k = 512, 512, 512
         dtype = torch.bfloat16
@@ -181,12 +202,10 @@ def matmul(a, b):
 
         a = torch.randn(m, k, device=device, dtype=dtype)
         b = torch.randn(k, n, device=device, dtype=dtype)
-
         expected = matmul(a, b)
 
         torch._dynamo.reset()
 
-        # Patch cutlass_api.Kernel.get_workspace_size to return non-zero
         import cutlass_api
 
         def patched_get_workspace_size(self, args):
@@ -197,13 +216,7 @@ def patched_get_workspace_size(self, args):
             "get_workspace_size",
             patched_get_workspace_size,
         ):
-            with config.patch(
-                {
-                    "max_autotune": True,
-                    "max_autotune_gemm_backends": "NVGEMM",
-                    "nvgemm_max_profiling_configs": 3,
-                }
-            ):
+            with config.patch(_nvgemm_config()):
                 result, (code,) = run_and_get_code(
                     torch.compile(matmul),
                     a,
@@ -211,7 +224,6 @@ def patched_get_workspace_size(self, args):
                 )
 
         self.assertIn("workspace=workspace", code)
-
         torch.testing.assert_close(result, expected)
 
     @parametrize("dtype", (torch.float16, torch.bfloat16))
@@ -223,103 +235,172 @@ def test_bmm_non_standard_batch_stride(self, dtype):
         def bmm(a, b):
             return torch.bmm(a, b)
 
-        # Create tensors with non-largest batch stride by transposing
-        # a_base shape: (m, batch, k), stride: (batch*k, k, 1)
-        # After transpose: shape (batch, m, k), stride: (k, batch*k, 1)
-        # batch_stride = k = 128, but m*k = 64*128 = 8192, so batch_stride < m*k
+        # Transpose creates non-largest batch stride
         a_base = torch.randn(m, batch, k, device=device, dtype=dtype)
         a = a_base.transpose(0, 1)  # (batch, m, k) with stride (k, batch*k, 1)
 
         b_base = torch.randn(k, batch, n, device=device, dtype=dtype)
         b = b_base.transpose(0, 1)  # (batch, k, n) with stride (n, batch*n, 1)
 
-        # Verify batch stride is not largest (i.e., batch_stride_largest_or_zero would be False)
-        assert a.stride()[0] != a.shape[1] * a.shape[2], (
-            "Test setup error: a should have non-standard batch stride"
-        )
-        assert b.stride()[0] != b.shape[1] * b.shape[2], (
-            "Test setup error: b should have non-standard batch stride"
-        )
+        if a.stride()[0] == a.shape[1] * a.shape[2]:
+            raise AssertionError(
+                "Test setup error: a should have non-standard batch stride"
+            )
+        if b.stride()[0] == b.shape[1] * b.shape[2]:
+            raise AssertionError(
+                "Test setup error: b should have non-standard batch stride"
+            )
 
         expected = bmm(a, b)
 
         torch._dynamo.reset()
 
-        with config.patch(
-            {
-                "max_autotune": True,
-                "max_autotune_gemm_backends": "NVGEMM",
-                "nvgemm_max_profiling_configs": 3,
-            }
-        ):
+        with config.patch(_nvgemm_config()):
             compiled_fn = torch.compile(bmm)
             result = compiled_fn(a, b)
 
         torch.testing.assert_close(result, expected)
 
-    def test_scaled_gemm_mxfp8(self):
-        """Test MXFP8 scaled GEMM with NVGEMM backend.
-
-        Note: Invalid inputs (wrong shapes, dtypes, K not divisible by 16, etc.)
-        are caught early by Dynamo's _check_scaled_mm_sizes in torch/_meta_registrations.py.
-        NVGEMM can assume inputs are valid by the time they reach kernel selection.
-        """
-        from torch._inductor.utils import ceildiv
-
-        m, n, k = 256, 512, 1024
+    @parametrize(
+        "layout_a",
+        ("contiguous", "aligned_offset", "view"),
+    )
+    @parametrize(
+        "m,n,k",
+        (
+            (256, 512, 1024),
+            (256, 1024, 512),
+            (128, 256, 512),
+            (512, 256, 1024),
+        ),
+    )
+    def test_scaled_gemm_mxfp8(self, layout_a, m, n, k):
+        """Test MXFP8 scaled GEMM with NVGEMM backend."""
         block_size = 32
-        device = "cuda"
-
-        def _round_up(x, multiple):
-            return ((x + multiple - 1) // multiple) * multiple
-
-        def _prep_k(K, scale_size):
-            """Prepare K dimension for 32-4-4 swizzle requirements."""
-            return _round_up(ceildiv(K, scale_size), 4)
 
         def scaled_mm(a, b, scale_a, scale_b):
             return torch._scaled_mm(
                 a, b, scale_a=scale_a, scale_b=scale_b, out_dtype=torch.float32
             )
 
-        # Create FP8 tensors
-        a_fp8 = torch.randint(-1, 2, (m, k), device=device).to(torch.float8_e4m3fn)
-        # B is N x K, then transposed to K x N for scaled_mm
-        b_fp8 = torch.randint(-1, 2, (n, k), device=device).to(torch.float8_e4m3fn).T
+        a_fp8 = _create_tensor_with_layout(layout_a, m, k, torch.float8_e4m3fn)
+        b_fp8 = torch.randint(-1, 2, (n, k), device="cuda").to(torch.float8_e4m3fn).T
 
-        # Scale factors in float8_e8m0fnu (MXFP8 format)
-        # Shape: (M, prep_k(K, 32)) for A, (prep_k(K, 32), N) for B
-        scale_a = torch.rand(m, _prep_k(k, block_size), device=device).to(
+        scale_a = torch.rand(m, _prep_k(k, block_size), device="cuda").to(
             torch.float8_e8m0fnu
         )
-        scale_b = torch.rand(_prep_k(k, block_size), n, device=device).to(
+        scale_b = torch.rand(_prep_k(k, block_size), n, device="cuda").to(
             torch.float8_e8m0fnu
         )
 
-        # Get reference result from eager mode (ATen)
         expected = scaled_mm(a_fp8, b_fp8, scale_a, scale_b)
 
-        with config.patch(
-            {
-                "max_autotune": True,
-                "max_autotune_gemm_backends": "NVGEMM",
-                "nvgemm_max_profiling_configs": 3,
-                "autotune_fallback_to_aten": False,
-            }
-        ):
+        torch._dynamo.reset()
+
+        with config.patch(_nvgemm_config()):
             compiled_fn = torch.compile(scaled_mm)
             result = compiled_fn(a_fp8, b_fp8, scale_a, scale_b)
 
         torch.testing.assert_close(result, expected)
 
-    def test_grouped_gemm(self):
-        """Test grouped GEMM with NVGEMM backend.
+    @parametrize("out_dtype", (torch.float32, torch.bfloat16))
+    @parametrize(
+        "layout_a",
+        ("contiguous", "aligned_offset", "view"),
+    )
+    @parametrize(
+        "m,n,k",
+        (
+            (256, 512, 1024),
+            (256, 1024, 512),
+            (128, 256, 512),
+            (512, 256, 1024),
+        ),
+    )
+    def test_scaled_gemm_nvf4(self, out_dtype, layout_a, m, n, k):
+        """Test NVF4 (Float4 + Float8E4M3FN scales, block_size=16) with NVGEMM backend."""
+        packed_k = k // 2
+        block_size = 16
+
+        def scaled_mm(a, b, scale_a, scale_b):
+            return torch._scaled_mm(
+                a, b, scale_a=scale_a, scale_b=scale_b, out_dtype=out_dtype
+            )
+
+        a_fp4 = _create_tensor_with_layout(
+            layout_a, m, packed_k, torch.float4_e2m1fn_x2
+        )
+        b_fp4 = torch.randint(
+            0, 256, (n, packed_k), device="cuda", dtype=torch.uint8
+        ).view(torch.float4_e2m1fn_x2)
+        b_fp4_t = b_fp4.T
+
+        num_k_blocks = ceildiv(k, block_size)
+        padded_k_blocks = _round_up(num_k_blocks, 4)
+        block_size_mn = 128
+        scale_a_numel = block_size_mn * ceildiv(m, block_size_mn) * padded_k_blocks
+        scale_b_numel = block_size_mn * ceildiv(n, block_size_mn) * padded_k_blocks
+
+        scale_a = torch.rand(scale_a_numel, device="cuda").to(torch.float8_e4m3fn)
+        scale_b = torch.rand(scale_b_numel, device="cuda").to(torch.float8_e4m3fn)
 
-        This test runs the same shape twice with different offsets to verify that
-        different offset distributions produce correct results.
+        expected = scaled_mm(a_fp4, b_fp4_t, scale_a, scale_b)
 
-        Note: GroupedGemm currently only supports TN layout (column-major B).
-        B is created with shape (g, k, n) but column-major inner layout via permute.
+        torch._dynamo.reset()
+
+        with config.patch(
+            _nvgemm_config(
+                **{"test_configs.autotune_choice_desc_regex": "inductor_vendored"}
+            )
+        ):
+            compiled_fn = torch.compile(scaled_mm)
+            result = compiled_fn(a_fp4, b_fp4_t, scale_a, scale_b)
+
+        # a_fp4 and b_fp4_t could come with NaNs.
+        torch.testing.assert_close(result, expected, equal_nan=True)
+
+    @parametrize(
+        "layout_a",
+        ("contiguous", "aligned_offset", "view", "padded"),
+    )
+    def test_grouped_gemm(self, layout_a):
+        """Test grouped GEMM with NVGEMM backend and various A layouts.
+
+        GroupedGemm currently only supports TN layout (column-major B).
+        """
+        g, k, n = 4, 256, 256
+        dtype = torch.bfloat16
+        device = "cuda"
+
+        def grouped_mm(a, b, offsets):
+            return torch._grouped_mm(a, b, offs=offsets)
+
+        b = torch.randn(g, n, k, device=device, dtype=dtype).permute(0, 2, 1)
+
+        m_per_group = [64, 64, 64, 64]
+        total_m = sum(m_per_group)
+        offsets = torch.tensor(
+            [sum(m_per_group[: i + 1]) for i in range(g)],
+            device=device,
+            dtype=torch.int32,
+        )
+        a = _create_tensor_with_layout(layout_a, total_m, k, dtype, device)
+
+        expected = grouped_mm(a, b, offsets)
+
+        torch._dynamo.reset()
+
+        with config.patch(_nvgemm_config()):
+            compiled_fn = torch.compile(grouped_mm)
+            result = compiled_fn(a, b, offsets)
+
+        torch.testing.assert_close(result, expected)
+
+    def test_grouped_gemm_varying_offsets(self):
+        """Test that different offset distributions produce correct results.
+
+        Runs the same compiled function with two different offset distributions
+        (same total_m) to verify offsets are handled dynamically at runtime.
         """
         g, k, n = 4, 256, 256
         dtype = torch.bfloat16
@@ -331,34 +412,30 @@ def grouped_mm(a, b, offsets):
         b = torch.randn(g, n, k, device=device, dtype=dtype).permute(0, 2, 1)
 
         m_per_group_1 = [64, 64, 64, 64]
-        total_m_1 = sum(m_per_group_1)
+        total_m = sum(m_per_group_1)
         offsets_1 = torch.tensor(
             [sum(m_per_group_1[: i + 1]) for i in range(g)],
             device=device,
             dtype=torch.int32,
         )
-        a_1 = torch.randn(total_m_1, k, device=device, dtype=dtype)
+        a_1 = torch.randn(total_m, k, device=device, dtype=dtype)
 
         m_per_group_2 = [32, 96, 48, 80]
-        total_m_2 = sum(m_per_group_2)
-        assert total_m_1 == total_m_2, "Total M must match for cache key test"
+        if sum(m_per_group_2) != total_m:
+            raise AssertionError("Total M must match for cache key test")
         offsets_2 = torch.tensor(
             [sum(m_per_group_2[: i + 1]) for i in range(g)],
             device=device,
             dtype=torch.int32,
         )
-        a_2 = torch.randn(total_m_2, k, device=device, dtype=dtype)
+        a_2 = torch.randn(total_m, k, device=device, dtype=dtype)
 
         expected_1 = grouped_mm(a_1, b, offsets_1)
         expected_2 = grouped_mm(a_2, b, offsets_2)
 
-        with config.patch(
-            {
-                "max_autotune": True,
-                "max_autotune_gemm_backends": "NVGEMM",
-                "autotune_fallback_to_aten": False,
-            }
-        ):
+        torch._dynamo.reset()
+
+        with config.patch(_nvgemm_config()):
             compiled_fn = torch.compile(grouped_mm)
 
             result_1 = compiled_fn(a_1, b, offsets_1)
@@ -460,6 +537,72 @@ def test_filter_kernels_sorts_by_runtime(self):
                 self.assertEqual(len(result), 3)
 
 
+@unittest.skipIf(
+    not (
+        ensure_nv_universal_gemm_available()
+        and is_datacenter_blackwell_arch()
+        and ensure_nvmatmul_heuristics_available()
+    ),
+    "Requires cutlass_api, nvMatmulHeuristics, and Blackwell GPU",
+)
+class TestNVUniversalGemmHeuristicsIntegration(TestCase):
+    """Integration tests for nvMatmulHeuristics with real library calls."""
+
+    def test_fp4_heuristic_configs(self):
+        """Test that nvMatmulHeuristics returns configs for FP4 blockscaled GEMM."""
+        heuristics = NVUniversalGemmHeuristics()
+
+        m, n, k = 256, 512, 1024
+        configs = heuristics._get_heuristic_configs(
+            m,
+            n,
+            k,
+            dtype_a=torch.float4_e2m1fn_x2,
+            layout_a="row",
+            layout_b="col",
+            count=5,
+            valid_configs=OrderedSet(),
+            accumulator_type=torch.float32,
+            dtype_b=torch.float4_e2m1fn_x2,
+            out_dtype=torch.float32,
+        )
+
+        self.assertGreater(
+            len(configs), 0, "nvMatmulHeuristics returned no FP4 configs"
+        )
+        for cfg in configs:
+            self.assertGreater(cfg.tile_m, 0)
+            self.assertGreater(cfg.tile_n, 0)
+            self.assertGreater(cfg.estimated_runtime, 0)
+
+    def test_fp8_heuristic_configs(self):
+        """Test that nvMatmulHeuristics returns configs for FP8 GEMM."""
+        heuristics = NVUniversalGemmHeuristics()
+
+        m, n, k = 256, 512, 1024
+        configs = heuristics._get_heuristic_configs(
+            m,
+            n,
+            k,
+            dtype_a=torch.float8_e4m3fn,
+            layout_a="row",
+            layout_b="col",
+            count=5,
+            valid_configs=OrderedSet(),
+            accumulator_type=torch.float32,
+            dtype_b=torch.float8_e4m3fn,
+            out_dtype=torch.float32,
+        )
+
+        self.assertGreater(
+            len(configs), 0, "nvMatmulHeuristics returned no FP8 configs"
+        )
+        for cfg in configs:
+            self.assertGreater(cfg.tile_m, 0)
+            self.assertGreater(cfg.tile_n, 0)
+            self.assertGreater(cfg.estimated_runtime, 0)
+
+
 @unittest.skipIf(
     not (ensure_nv_universal_gemm_available() and is_datacenter_blackwell_arch()),
     "NVIDIA Universal GEMM (cutlass_api) library not available or not on Blackwell",
@@ -473,7 +616,6 @@ def test_unbacked_symint_rejected(self):
 
         def fn(x, w):
             nz = torch.nonzero(x)  # Creates unbacked symint for nz.size(0)
-            # Use unbacked symint as M dimension in matmul
             a = torch.ones(nz.size(0), w.size(0), dtype=w.dtype, device=w.device)
             return a @ w
 
@@ -482,13 +624,7 @@ def fn(x, w):
 
         torch._dynamo.reset()
 
-        with config.patch(
-            {
-                "max_autotune": True,
-                "max_autotune_gemm_backends": "NVGEMM",
-                "nvgemm_max_profiling_configs": 2,
-            }
-        ):
+        with config.patch(_nvgemm_config(nvgemm_max_profiling_configs=2)):
             compiled_fn = torch.compile(fn, dynamic=True)
             with self.assertRaisesRegex(
                 Exception, "NoValidChoicesError|no valid choice"
@@ -503,13 +639,7 @@ def matmul(a, b):
 
         torch._dynamo.reset()
 
-        with config.patch(
-            {
-                "max_autotune": True,
-                "max_autotune_gemm_backends": "NVGEMM",
-                "nvgemm_max_profiling_configs": 2,
-            }
-        ):
+        with config.patch(_nvgemm_config(nvgemm_max_profiling_configs=2)):
             compiled_fn = torch.compile(matmul, dynamic=True)
 
             shapes = [
diff --git a/test/inductor/test_op_dtype_prop.py b/test/inductor/test_op_dtype_prop.py
index 9cd9080a247cf..7645213ea797c 100644
--- a/test/inductor/test_op_dtype_prop.py
+++ b/test/inductor/test_op_dtype_prop.py
@@ -108,34 +108,32 @@ def func(a, b, c, d):
             fp32_cast_in_code = "to(tl.float32)" in code
             self.assertEqual(fp32_cast_in_code, upcast_to_fp32)
 
-        @requires_gpu()
-        @parametrize("input_shape", [(32, 32), (32, 128), (256, 32)])
-        @parametrize(
-            "reduction_func",
-            [
-                torch.prod,
-                torch.sum,
-                torch.argmax,
-                torch.argmin,
-                torch.min,
-                torch.max,
-            ],
-        )
-        @parametrize("input_dtype", [torch.float16, torch.bfloat16])
-        @config.patch("triton.use_block_ptr", True)
-        def test_low_precision_reduction(
-            self, input_shape, reduction_func, input_dtype
-        ):
-            @torch.compile
-            def func(a, b, c, d):
-                return reduction_func(a * b * c * d)
-
-            inps = (torch.rand(input_shape, device=GPU_TYPE, dtype=input_dtype),) * 4
-            with config.patch("triton.codegen_upcast_to_fp32", False):
-                func_opt = torch._dynamo.optimize("inductor")(func)
-                code = run_and_get_triton_code(func_opt, *inps)
-                self.assertTrue(".to(tl.float32)" in code)
-                self.assertEqual(func(*inps), func_opt(*inps))
+    @requires_gpu()
+    @parametrize("input_shape", [(32, 32), (32, 128), (256, 32)])
+    @parametrize(
+        "reduction_func",
+        [
+            torch.prod,
+            torch.sum,
+            torch.argmax,
+            torch.argmin,
+            torch.min,
+            torch.max,
+        ],
+    )
+    @parametrize("input_dtype", [torch.float16, torch.bfloat16])
+    @config.patch("triton.use_block_ptr", True)
+    def test_low_precision_reduction(self, input_shape, reduction_func, input_dtype):
+        @torch.compile
+        def func(a, b, c, d):
+            return reduction_func(a * b * c * d)
+
+        inps = (torch.rand(input_shape, device=GPU_TYPE, dtype=input_dtype),) * 4
+        with config.patch("triton.codegen_upcast_to_fp32", False):
+            func_opt = torch._dynamo.optimize("inductor")(func)
+            code = run_and_get_triton_code(func_opt, *inps)
+            self.assertTrue(".to(tl.float32)" in code)
+            self.assertEqual(func(*inps), func_opt(*inps))
 
     def test_op_dtype_support(self):
         """
diff --git a/test/inductor/test_pad_as_cat.py b/test/inductor/test_pad_as_cat.py
new file mode 100644
index 0000000000000..81f0d02c016e3
--- /dev/null
+++ b/test/inductor/test_pad_as_cat.py
@@ -0,0 +1,127 @@
+# Owner(s): ["module: inductor"]
+"""Tests for cat multi-consumer and pad-as-cat optimizations."""
+
+import torch
+from torch._dynamo.utils import counters
+from torch._inductor import metrics
+from torch._inductor.test_case import TestCase
+from torch._inductor.utils import run_and_get_code
+from torch.testing._internal.inductor_utils import GPU_TYPE, requires_gpu
+
+
+# required so that metrics.num_bytes_accessed is populated
+torch._logging.set_logs(inductor_metrics=True)
+
+
+class TestCatMultiConsumer(TestCase):
+    @torch._inductor.config.patch(fx_graph_cache=False)
+    @requires_gpu()
+    def test_cat_to_fp16(self):
+        """Multi-consumer cat avoids duplicate computation."""
+
+        def fn(x):
+            z = torch.cat([x, torch.zeros([6, 768], device=GPU_TYPE)], dim=0)
+            y = x.to(torch.float16)
+            return z, y
+
+        x = torch.randn(1024, 768, device=GPU_TYPE)
+        compiled = torch.compile(fn)
+        metrics.reset()
+        result = compiled(x)
+        ref = fn(x)
+
+        self.assertEqual(result[0], ref[0])
+        self.assertEqual(result[1], ref[1])
+
+        # Without the optimization x would be read twice (once by cat, once
+        # by to_fp16). With the optimization ConcatKernel shares x so it is
+        # read only once.
+        z = ref[0]
+        y = ref[1]
+        x_bytes = x.nelement() * x.element_size()
+        z_bytes = z.nelement() * z.element_size()
+        y_bytes = y.nelement() * y.element_size()
+        unoptimized_bytes = 2 * x_bytes + z_bytes + y_bytes
+        self.assertLess(
+            metrics.num_bytes_accessed,
+            unoptimized_bytes,
+            "Optimization should avoid reading x twice.",
+        )
+
+    @torch._inductor.config.patch(fx_graph_cache=False)
+    @requires_gpu()
+    def test_single_consumer_cat_unchanged(self):
+        """Single-consumer cat unchanged."""
+
+        def fn(x):
+            return torch.cat([x, torch.zeros([6, 768], device=GPU_TYPE)], dim=0)
+
+        x = torch.randn(1024, 768, device=GPU_TYPE)
+        compiled = torch.compile(fn)
+        metrics.reset()
+        result = compiled(x)
+        ref = fn(x)
+
+        self.assertEqual(result, ref)
+
+        # Single-consumer cat should use pointwise_cat which fuses the
+        # zeros fill into the cat kernel.  Total bytes: read x + write z.
+        x_bytes = x.nelement() * x.element_size()
+        z_bytes = ref.nelement() * ref.element_size()
+        expected_bytes = x_bytes + z_bytes
+        self.assertEqual(
+            metrics.num_bytes_accessed,
+            expected_bytes,
+            f"Expected {expected_bytes} bytes, got {metrics.num_bytes_accessed}.",
+        )
+
+
+class TestPadAsCat(TestCase):
+    @requires_gpu()
+    def test_mul_pad_addmm(self):
+        """Multi-consumer F.pad uses ConcatKernel zero-copy."""
+        counters.clear()
+
+        def fn(x, scale, bias, weight):
+            mul_result = x * scale
+            padded = torch.nn.functional.pad(mul_result, [0, 192])
+            mm_result = torch.addmm(bias, mul_result, weight)
+            return padded, mm_result
+
+        x = torch.randn(128, 2880, device=GPU_TYPE, dtype=torch.bfloat16)
+        scale = torch.randn(128, 2880, device=GPU_TYPE, dtype=torch.bfloat16)
+        bias = torch.randn(1024, device=GPU_TYPE, dtype=torch.bfloat16)
+        weight = torch.randn(2880, 1024, device=GPU_TYPE, dtype=torch.bfloat16)
+
+        compiled = torch.compile(fn)
+        result, (code,) = run_and_get_code(compiled, x, scale, bias, weight)
+        ref = fn(x, scale, bias, weight)
+
+        self.assertEqual(result[0], ref[0])
+        self.assertEqual(result[1], ref[1], atol=1e-2, rtol=1e-2)
+        self.assertIn("reinterpret_tensor", code)
+        self.assertGreater(counters["inductor"]["pad_rewritten_as_cat"], 0)
+
+    @requires_gpu()
+    def test_single_consumer_pad(self):
+        """Single-consumer F.pad is decomposed into cat, which fuses via pointwise_cat."""
+        counters.clear()
+
+        def fn(x, scale):
+            return torch.nn.functional.pad(x * scale, [0, 192])
+
+        x = torch.randn(128, 2880, device=GPU_TYPE)
+        scale = torch.randn(128, 2880, device=GPU_TYPE)
+
+        compiled = torch.compile(fn)
+        result = compiled(x, scale)
+        ref = fn(x, scale)
+
+        self.assertEqual(result, ref)
+        self.assertGreater(counters["inductor"]["pad_rewritten_as_cat"], 0)
+
+
+if __name__ == "__main__":
+    from torch._inductor.test_case import run_tests
+
+    run_tests()
diff --git a/test/inductor/test_pad_mm.py b/test/inductor/test_pad_mm.py
index 20e0359f50ec1..728b8635765e1 100644
--- a/test/inductor/test_pad_mm.py
+++ b/test/inductor/test_pad_mm.py
@@ -398,10 +398,10 @@ def test_pad_batch(self):
         mat2 = torch.ones((batch_size, k, n), device=GPU_TYPE, dtype=torch.float16)
         expected_alignment = get_alignment_size(mat1)
 
-        assert expected_alignment == 8, "Alignment for float16 should be 8"
-        assert can_pad(mat1, mat2, torch.ops.aten.bmm), (
-            "This should pass the common padding criteria"
-        )
+        if expected_alignment != 8:
+            raise AssertionError("Alignment for float16 should be 8")
+        if not can_pad(mat1, mat2, torch.ops.aten.bmm):
+            raise AssertionError("This should pass the common padding criteria")
 
         @torch.compile()
         def bmm(mat1, mat2):
@@ -414,9 +414,8 @@ def bmm(mat1, mat2):
             ".run(", 2, exactly=True
         ).check(f"empty_strided_{GPU_TYPE}((3, 8, 16)").run(code)
 
-        assert torch.allclose(res2, bmm_expected_result), (
-            "BMM results are not identical"
-        )
+        if not torch.allclose(res2, bmm_expected_result):
+            raise AssertionError("BMM results are not identical")
 
     @fresh_cache()
     def test_exclude_padding(self):
@@ -486,13 +485,14 @@ def test_pad_mm_bf16(self):
         mat2 = torch.ones((k, n), device=GPU_TYPE, dtype=torch.bfloat16)
         expected_alignment = get_alignment_size(mat1)
 
-        assert expected_alignment == 8, "Alignment for bfloat16 should be 8"
-        assert can_pad(mat1, mat2, torch.ops.aten.mm), (
-            "This should pass the common padding criteria"
-        )
-        assert should_pad_mm_bf16(mat1.dtype, m, n, k), (
-            "This should pass the should_pad_mm_bf16 padding criteria"
-        )
+        if expected_alignment != 8:
+            raise AssertionError("Alignment for bfloat16 should be 8")
+        if not can_pad(mat1, mat2, torch.ops.aten.mm):
+            raise AssertionError("This should pass the common padding criteria")
+        if not should_pad_mm_bf16(mat1.dtype, m, n, k):
+            raise AssertionError(
+                "This should pass the should_pad_mm_bf16 padding criteria"
+            )
 
         @torch.compile()
         def mm(mat1, mat2):
@@ -505,7 +505,8 @@ def mm(mat1, mat2):
             ".run(", 2, exactly=True
         ).check(f"empty_strided_{GPU_TYPE}((8, 16)").run(code)
 
-        assert torch.allclose(res2, mm_expected_result), "MM results are not identical"
+        if not torch.allclose(res2, mm_expected_result):
+            raise AssertionError("MM results are not identical")
 
     @fresh_cache()
     @inductor_config.patch(
@@ -531,7 +532,9 @@ def fn(x, y):
         ):
             opt_fn = torch.compile(fn, mode="max-autotune")
             ret, code = run_and_get_code(opt_fn, *args)
-        self.assertEqual(counters["inductor"]["pattern_matcher_count"], 1)
+        # xref: https://github.com/pytorch/pytorch/pull/172780
+        if not torch.version.hip:  # autotuning is not guaranteed to run on ROCm
+            self.assertEqual(counters["inductor"]["pattern_matcher_count"], 1)
 
         code = [c for c in code if "decompose_k" not in c]
         # The mm kernel should use a template (because we set max_autotune_gemm_backends = TRITON).
diff --git a/test/inductor/test_pallas.py b/test/inductor/test_pallas.py
index 6f8d92b0f203b..17baaae381eaf 100644
--- a/test/inductor/test_pallas.py
+++ b/test/inductor/test_pallas.py
@@ -5,6 +5,8 @@
 import sys
 import unittest
 
+import numpy as np
+
 import torch
 import torch._dynamo
 import torch._inductor.async_compile  # noqa: F401 required to warm up AsyncCompile pools
@@ -90,6 +92,9 @@ def skip_decorator(fn):
         decorator=skip_decorator,
     )
 
+    # Pallas does not support float64 or int64
+    test_class._unsupported_input_gen_types = {"double"}
+
     test_classes[test_class.__name__] = test_class
     # REMOVING THIS LINE WILL STOP TESTS FROM RUNNING
     globals()[test_class.__name__] = test_class
@@ -97,22 +102,27 @@ def skip_decorator(fn):
     return test_class
 
 
-def _skip_if(condition_fn, reason):
-    def decorator(fn):
-        @functools.wraps(fn)
-        def wrapper(self, *args, **kwargs):
-            if condition_fn(self):
-                self.skipTest(reason)
-            fn(self, *args, **kwargs)
+def _skip_if(condition_fn):
+    def skip(fn=None, *, reason=None):
+        def decorator(fn):
+            @functools.wraps(fn)
+            def wrapper(self, *args, **kwargs):
+                if condition_fn(self):
+                    self.skipTest(reason or f"Not yet working on {self.DEVICE}")
+                fn(self, *args, **kwargs)
 
-        return wrapper
+            return wrapper
 
-    return decorator
+        if fn is not None:
+            return decorator(fn)
+        return decorator
 
+    return skip
 
-skip_if_tpu = _skip_if(lambda self: self.DEVICE == "tpu", "Not yet working on TPU")
-skip_if_cpu = _skip_if(lambda self: self.DEVICE == "cpu", "Not yet working on CPU")
-skip_if_cuda = _skip_if(lambda self: self.DEVICE == "cuda", "Not yet working on GPU")
+
+skip_if_tpu = _skip_if(lambda self: self.DEVICE == "tpu")
+skip_if_cpu = _skip_if(lambda self: self.DEVICE == "cpu")
+skip_if_cuda = _skip_if(lambda self: self.DEVICE == "cuda")
 
 
 class PallasTestsMixin:
@@ -158,7 +168,9 @@ def _compile(self, fn):
             "tpu": "tpu_backend",
         }
         key = device_to_backend_key[self.DEVICE]
-        return torch.compile(fn, backend="inductor", options={key: "pallas"})
+        return torch.compile(
+            fn, backend="inductor", options={key: "pallas"}, dynamic=False
+        )
 
     def test_simple_add(self):
         """Test basic element-wise addition."""
@@ -166,13 +178,17 @@ def test_simple_add(self):
         def fn(a, b):
             return a + b
 
-        compiled = self._compile(fn)
-
-        a = torch.randn(1024, device=self.DEVICE)
-        b = torch.randn(1024, device=self.DEVICE)
-        result = compiled(a, b)
-        expected = fn(a, b)
-        self.assertEqual(result, expected)
+        shapes = [(1024,)]
+        if self.DEVICE != "cuda":
+            shapes += [(2048,), (2048, 128)]
+        for shape in shapes:
+            with self.subTest(shape=shape):
+                compiled = self._compile(fn)
+                a = torch.randn(shape, device=self.DEVICE)
+                b = torch.randn(shape, device=self.DEVICE)
+                result = compiled(a, b)
+                expected = fn(a, b)
+                self.assertEqual(result, expected)
 
     def test_simple_mul(self):
         """Test basic element-wise multiplication."""
@@ -180,13 +196,17 @@ def test_simple_mul(self):
         def fn(a, b):
             return a * b
 
-        compiled = self._compile(fn)
-
-        a = torch.randn(1024, device=self.DEVICE)
-        b = torch.randn(1024, device=self.DEVICE)
-        result = compiled(a, b)
-        expected = fn(a, b)
-        self.assertEqual(result, expected)
+        shapes = [(1024,)]
+        if self.DEVICE != "cuda":
+            shapes += [(2048,), (2048, 128)]
+        for shape in shapes:
+            with self.subTest(shape=shape):
+                compiled = self._compile(fn)
+                a = torch.randn(shape, device=self.DEVICE)
+                b = torch.randn(shape, device=self.DEVICE)
+                result = compiled(a, b)
+                expected = fn(a, b)
+                self.assertEqual(result, expected)
 
     def test_sin(self):
         """Test sin operation."""
@@ -194,12 +214,16 @@ def test_sin(self):
         def fn(x):
             return torch.sin(x)
 
-        compiled = self._compile(fn)
-
-        x = torch.randn(1024, device=self.DEVICE)
-        result = compiled(x)
-        expected = fn(x)
-        self.assertEqual(result, expected)
+        shapes = [(1024,)]
+        if self.DEVICE != "cuda":
+            shapes.append((2048,))
+        for shape in shapes:
+            with self.subTest(shape=shape):
+                compiled = self._compile(fn)
+                x = torch.randn(shape, device=self.DEVICE)
+                result = compiled(x)
+                expected = fn(x)
+                self.assertEqual(result, expected)
 
     def test_fused_ops(self):
         """Test fused operations (sin + add)."""
@@ -228,10 +252,9 @@ def fn(x):
         expected = fn(x)
         self.assertEqual(result, expected)
 
+    @skip_if_cuda(reason="sqrt primitive not implemented in Pallas Mosaic GPU")
     def test_sqrt(self):
         """Test sqrt operation."""
-        if self.DEVICE == "cuda":
-            self.skipTest("sqrt primitive not implemented in Pallas Mosaic GPU")
 
         def fn(x):
             return torch.sqrt(x)
@@ -358,26 +381,28 @@ def test_jax_jit_wrapper_is_emitted(self):
         self.assertNotIn("torch.", wrapper_block)
 
     def test_2d_tensor(self):
-        """Test with 2D tensors (though current implementation flattens)."""
+        """Test with 2D tensors."""
 
         def fn(x, y):
             return x + y
 
-        compiled = self._compile(fn)
-
-        x = torch.randn(32, 32, device=self.DEVICE)
-        y = torch.randn(32, 32, device=self.DEVICE)
-        result = compiled(x, y)
-        expected = fn(x, y)
-        self.assertEqual(result, expected)
-
-    @skip_if_tpu
+        shapes = [(32, 32)]
+        if self.DEVICE != "cuda":
+            shapes.append((16, 2048))
+        for shape in shapes:
+            with self.subTest(shape=shape):
+                compiled = self._compile(fn)
+                x = torch.randn(shape, device=self.DEVICE)
+                y = torch.randn(shape, device=self.DEVICE)
+                result = compiled(x, y)
+                expected = fn(x, y)
+                self.assertEqual(result, expected)
+
+    @skip_if_cuda(
+        reason="iteration variables not supported in Pallas GPU (Mosaic) backend"
+    )
     def test_different_shapes(self):
         """Test with different tensor shapes."""
-        if self.DEVICE == "cuda":
-            self.skipTest(
-                "iteration variables not supported in Pallas GPU (Mosaic) backend"
-            )
 
         def fn(x):
             return x * 2.0
@@ -424,31 +449,30 @@ def operate_on_tensor(x):
 
         compiled = self._compile(operate_on_tensor)
 
-        # Create a transposed (non-contiguous) view
-        x = torch.randn(64, 32, device=self.DEVICE)
-        x_t = x.t()  # Non-contiguous view
-        self.assertFalse(x_t.is_contiguous())
-
-        # With the simplified dlpack approach, non-contiguous tensors now work
-        result = compiled(x_t)
-        expected = operate_on_tensor(x_t)
-        self.assertEqual(result, expected)
-
-        # Contiguous tensors should also continue to work
-        x_t_contiguous = x_t.contiguous()
-        self.assertTrue(x_t_contiguous.is_contiguous())
-        result = compiled(x_t_contiguous)
-        expected = operate_on_tensor(x_t_contiguous)
-        self.assertEqual(result, expected)
-
-    @skip_if_tpu
+        for rows, cols in [(64, 32), (5, 8), (3215, 23), (8, 128), (128, 8)]:
+            with self.subTest(rows=rows, cols=cols):
+                # Create a transposed (non-contiguous) view
+                x = torch.randn(rows, cols, device=self.DEVICE)
+                x_t = x.t()  # Non-contiguous view
+                self.assertFalse(x_t.is_contiguous())
+
+                # With the simplified dlpack approach, non-contiguous tensors now work
+                result = compiled(x_t)
+                expected = operate_on_tensor(x_t)
+                self.assertEqual(result, expected)
+
+                # Contiguous tensors should also continue to work
+                x_t_contiguous = x_t.contiguous()
+                self.assertTrue(x_t_contiguous.is_contiguous())
+                result = compiled(x_t_contiguous)
+                expected = operate_on_tensor(x_t_contiguous)
+                self.assertEqual(result, expected)
+
+    @skip_if_cuda(reason="strided access not supported in Pallas GPU (Mosaic) backend")
     def test_strided_int_pallas(self):
         """Test strided access patterns with the Pallas backend."""
-        if self.DEVICE == "cuda":
-            self.skipTest("strided access not supported in Pallas GPU (Mosaic) backend")
 
         def fn(x):
-            # Access every other element (strided access)
             return x[::2] * 2.0
 
         compiled = self._compile(fn)
@@ -458,14 +482,11 @@ def fn(x):
         expected = fn(x)
         self.assertEqual(result, expected)
 
-    @skip_if_tpu
+    @skip_if_cuda(reason="strided access not supported in Pallas GPU (Mosaic) backend")
     def test_strided_offset_pallas(self):
         """Test strided access with offset."""
-        if self.DEVICE == "cuda":
-            self.skipTest("strided access not supported in Pallas GPU (Mosaic) backend")
 
         def fn(x):
-            # Access every other element starting from index 1
             return x[1::2] + 1.0
 
         compiled = self._compile(fn)
@@ -530,10 +551,20 @@ def test_stride_non_contiguous_2d_transpose(self):
         """Test 2D transposed input patterns."""
         compiled = self._compile(lambda x: x * 2.0 + 1.0)
 
-        base_2d = torch.randn(32, 32, device=self.DEVICE)
-        x = base_2d.t()  # (32, 32) with stride (1, 32)
-        self.assertFalse(x.is_contiguous())
-        self.assertEqual(compiled(x), x * 2.0 + 1.0)
+        for rows, cols in [
+            (32, 32),
+            (2048, 2048),
+            (64, 32),
+            (5, 8),
+            (3215, 23),
+            (8, 128),
+            (128, 8),
+        ]:
+            with self.subTest(rows=rows, cols=cols):
+                base_2d = torch.randn(rows, cols, device=self.DEVICE)
+                x = base_2d.t()
+                self.assertFalse(x.is_contiguous())
+                self.assertEqual(compiled(x), x * 2.0 + 1.0)
 
     def test_stride_non_contiguous_3d(self):
         """Test 3D non-contiguous input patterns."""
@@ -599,7 +630,7 @@ def test_stride_non_contiguous_unsqueeze(self):
         x = base_2d[::2, ::2].unsqueeze(0)
         self.assertEqual(compiled(x), x * 2.0 + 1.0)
 
-    @skip_if_tpu
+    @skip_if_tpu(reason="TPU doesn't support float 64")
     def test_stride_non_contiguous_dtypes(self):
         """Test non-contiguous patterns with various dtypes."""
         compiled = self._compile(lambda x: x * 2.0 + 1.0)
@@ -631,7 +662,6 @@ def test_stride_expanded_tensors(self):
         x = torch.randn(1, 1, 16, device=self.DEVICE).expand(4, 8, 16)
         self.assertEqual(compiled(x, x), x + x)
 
-    @skip_if_tpu
     def test_stride_multiple_inputs(self):
         """Test multiple strided inputs and broadcasting."""
         compiled = self._compile(lambda a, b, c: a * b + c)
@@ -657,6 +687,80 @@ def test_stride_multiple_inputs(self):
         compiled_bcast = self._compile(lambda x, y, s: x + y * s)
         self.assertEqual(compiled_bcast(x, y, s), x + y * s)
 
+    @skip_if_cuda
+    def test_scalar_scalar_ops(self):
+        """Test scalar-scalar operations."""
+
+        def test_add(a, b):
+            return a + b
+
+        def test_mul(a, b):
+            return a * b
+
+        def test_sub(a, b):
+            return a - b
+
+        def test_div(a, b):
+            return a / b
+
+        for fn in [test_add, test_mul, test_sub, test_div]:
+            with self.subTest(op=fn.__name__):
+                compiled = self._compile(fn)
+
+                # Test with 0-D tensors (scalars)
+                a = torch.tensor(3.5, dtype=torch.float32, device=self.DEVICE)
+                b = torch.tensor(2.0, dtype=torch.float32, device=self.DEVICE)
+
+                result = compiled(a, b)
+                expected = fn(a, b)
+                self.assertEqual(result, expected)
+
+                # Ensure result is also scalar
+                self.assertEqual(result.dim(), 0)
+                self.assertEqual(result.dtype, torch.float32)
+
+    def test_scalar_tensor_ops(self):
+        """Test scalar-tensor operations."""
+
+        def test_scalar_add_tensor(s, t):
+            return s + t
+
+        def test_tensor_add_scalar(t, s):
+            return t + s
+
+        def test_scalar_mul_tensor(s, t):
+            return s * t
+
+        def test_tensor_mul_scalar(t, s):
+            return t * s
+
+        shapes = [(16,), (8, 8), (4, 4, 4)]
+
+        for shape in shapes:
+            for fn in [
+                test_scalar_add_tensor,
+                test_tensor_add_scalar,
+                test_scalar_mul_tensor,
+                test_tensor_mul_scalar,
+            ]:
+                with self.subTest(op=fn.__name__, shape=shape):
+                    compiled = self._compile(fn)
+
+                    # Create 0-D scalar tensor
+                    scalar = torch.tensor(2.5, dtype=torch.float32, device=self.DEVICE)
+                    tensor = torch.randn(shape, dtype=torch.float32, device=self.DEVICE)
+
+                    if "scalar" in fn.__name__.split("_")[0]:
+                        result = compiled(scalar, tensor)
+                        expected = fn(scalar, tensor)
+                    else:
+                        result = compiled(tensor, scalar)
+                        expected = fn(tensor, scalar)
+
+                    self.assertEqual(result, expected)
+                    self.assertEqual(result.shape, shape)
+                    self.assertEqual(result.dtype, torch.float32)
+
     def test_non_power_of_2_sizes(self):
         """Test that non-power-of-2 tensor sizes work correctly.
 
@@ -691,11 +795,10 @@ def fn(x, y):
         expected = fn(x, y)
         self.assertEqual(result, expected)
 
-    @skip_if_tpu
+    @skip_if_tpu(reason="Cannot do int indexing on TPU")
+    @skip_if_cuda(reason="gather not supported in Pallas GPU (Mosaic) backend")
     def test_complex_indexing_gather(self):
         """Test complex indexing with gather-like operations."""
-        if self.DEVICE == "cuda":
-            self.skipTest("gather not supported in Pallas GPU (Mosaic) backend")
 
         def fn(x, indices):
             # Use indices to gather elements from x
@@ -712,15 +815,14 @@ def fn(x, indices):
         expected = fn(x, indices)
         self.assertEqual(result, expected)
 
-    @skip_if_tpu
+    @skip_if_tpu(reason="Cannot do int indexing on TPU")
+    # Pallas Mosaic backend doesn't support gather operations with array indices
+    # This limitation is in the Pallas/Mosaic lowering, not our implementation
+    @skip_if_cuda(
+        reason="Multi-dimensional gather not supported on Pallas Mosaic (CUDA) backend"
+    )
     def test_complex_indexing_2d(self):
         """Test complex indexing on 2D tensors with integer array indexing."""
-        if self.DEVICE == "cuda":
-            # Pallas Mosaic backend doesn't support gather operations with array indices
-            # This limitation is in the Pallas/Mosaic lowering, not our implementation
-            self.skipTest(
-                "Multi-dimensional gather not supported on Pallas Mosaic (CUDA) backend"
-            )
 
         def fn(x, row_indices):
             # Select specific rows using integer array indexing
@@ -742,13 +844,17 @@ def test_complex64_mul(self):
         def fn(a, b):
             return a * b
 
-        compiled = self._compile(fn)
-
-        a = torch.randn(128, dtype=torch.complex64, device=self.DEVICE)
-        b = torch.randn(128, dtype=torch.complex64, device=self.DEVICE)
-        result = compiled(a, b)
-        expected = fn(a, b)
-        self.assertEqual(result, expected)
+        sizes = [128]
+        if self.DEVICE != "cuda":
+            sizes.append(2048)
+        for size in sizes:
+            with self.subTest(size=size):
+                compiled = self._compile(fn)
+                a = torch.randn(size, dtype=torch.complex64, device=self.DEVICE)
+                b = torch.randn(size, dtype=torch.complex64, device=self.DEVICE)
+                result = compiled(a, b)
+                expected = fn(a, b)
+                self.assertEqual(result, expected)
 
     @skip_if_tpu
     def test_complex_conj(self):
@@ -855,13 +961,17 @@ def test_where(self):
         def fn(x, y):
             return torch.where(x > 0, x, y)
 
-        compiled = self._compile(fn)
-
-        x = torch.randn(1024, device=self.DEVICE)
-        y = torch.randn(1024, device=self.DEVICE)
-        result = compiled(x, y)
-        expected = fn(x, y)
-        self.assertEqual(result, expected)
+        shapes = [(1024,)]
+        if self.DEVICE != "cuda":
+            shapes += [(2048,), (2048, 128)]
+        for shape in shapes:
+            with self.subTest(shape=shape):
+                compiled = self._compile(fn)
+                x = torch.randn(shape, device=self.DEVICE)
+                y = torch.randn(shape, device=self.DEVICE)
+                result = compiled(x, y)
+                expected = fn(x, y)
+                self.assertEqual(result, expected)
 
     def test_clamp(self):
         """Test torch.clamp operation."""
@@ -869,12 +979,16 @@ def test_clamp(self):
         def fn(x):
             return torch.clamp(x, -1.0, 1.0)
 
-        compiled = self._compile(fn)
-
-        x = torch.randn(1024, device=self.DEVICE) * 2
-        result = compiled(x)
-        expected = fn(x)
-        self.assertEqual(result, expected)
+        shapes = [(1024,)]
+        if self.DEVICE != "cuda":
+            shapes += [(2048,), (2048, 128)]
+        for shape in shapes:
+            with self.subTest(shape=shape):
+                compiled = self._compile(fn)
+                x = torch.randn(shape, device=self.DEVICE) * 2
+                result = compiled(x)
+                expected = fn(x)
+                self.assertEqual(result, expected)
 
     def test_comparison_ops(self):
         """Test comparison operations."""
@@ -885,13 +999,17 @@ def fn(a, b):
             eq = a == b
             return gt.float() + lt.float() + eq.float()
 
-        compiled = self._compile(fn)
-
-        a = torch.randn(1024, device=self.DEVICE)
-        b = torch.randn(1024, device=self.DEVICE)
-        result = compiled(a, b)
-        expected = fn(a, b)
-        self.assertEqual(result, expected)
+        shapes = [(1024,)]
+        if self.DEVICE != "cuda":
+            shapes += [(2048,), (2048, 128)]
+        for shape in shapes:
+            with self.subTest(shape=shape):
+                compiled = self._compile(fn)
+                a = torch.randn(shape, device=self.DEVICE)
+                b = torch.randn(shape, device=self.DEVICE)
+                result = compiled(a, b)
+                expected = fn(a, b)
+                self.assertEqual(result, expected)
 
     def test_logical_ops(self):
         """Test logical operations."""
@@ -899,15 +1017,18 @@ def test_logical_ops(self):
         def fn(a, b):
             return torch.logical_and(a > 0, b > 0).float()
 
-        compiled = self._compile(fn)
+        shapes = [(1024,)]
+        if self.DEVICE != "cuda":
+            shapes += [(2048,), (2048, 128)]
+        for shape in shapes:
+            with self.subTest(shape=shape):
+                compiled = self._compile(fn)
+                a = torch.randn(shape, device=self.DEVICE)
+                b = torch.randn(shape, device=self.DEVICE)
+                result = compiled(a, b)
+                expected = fn(a, b)
+                self.assertEqual(result, expected)
 
-        a = torch.randn(1024, device=self.DEVICE)
-        b = torch.randn(1024, device=self.DEVICE)
-        result = compiled(a, b)
-        expected = fn(a, b)
-        self.assertEqual(result, expected)
-
-    @skip_if_tpu
     def test_sign(self):
         """Test sign operation."""
 
@@ -921,10 +1042,9 @@ def fn(x):
         expected = fn(x)
         self.assertEqual(result, expected)
 
+    @skip_if_cuda(reason="integer_pow primitive not implemented in Pallas Mosaic GPU")
     def test_reciprocal(self):
         """Test reciprocal operation."""
-        if self.DEVICE == "cuda":
-            self.skipTest("integer_pow primitive not implemented in Pallas Mosaic GPU")
 
         def fn(x):
             return torch.reciprocal(x)
@@ -949,7 +1069,9 @@ def fn(x):
         expected = fn(x)
         self.assertEqual(result, expected)
 
-    @skip_if_tpu
+    @skip_if_tpu(
+        reason="Pallas loweing crash: https://github.com/jax-ml/jax/issues/36149"
+    )
     def test_erf(self):
         """Test erf operation."""
 
@@ -963,7 +1085,9 @@ def fn(x):
         expected = fn(x)
         self.assertEqual(result, expected)
 
-    @skip_if_tpu
+    @skip_if_tpu(
+        reason="Pallas loweing crash: https://github.com/jax-ml/jax/issues/36149"
+    )
     def test_atan2(self):
         """Test atan2 operation."""
 
@@ -978,133 +1102,176 @@ def fn(a, b):
         expected = fn(a, b)
         self.assertEqual(result, expected)
 
-    @skip_if_tpu
     def test_sum_reduction(self):
         """Test sum reduction."""
 
         def fn(x):
             return x.sum()
 
-        compiled = self._compile(fn)
+        shapes = [(1024,)]
+        if self.DEVICE != "cuda":
+            shapes.append((2048,))
+        for shape in shapes:
+            with self.subTest(shape=shape):
+                compiled = self._compile(fn)
+                x = torch.randn(shape, device=self.DEVICE)
+                result = compiled(x)
+                expected = fn(x)
+                self.assertEqual(result, expected)
 
-        x = torch.randn(1024, device=self.DEVICE)
-        result = compiled(x)
-        expected = fn(x)
-        self.assertEqual(result, expected)
-
-    @skip_if_tpu
     def test_max_reduction(self):
         """Test max reduction."""
 
         def fn(x):
             return x.max()
 
-        compiled = self._compile(fn)
+        shapes = [(1024,)]
+        if self.DEVICE != "cuda":
+            shapes.append((2048,))
+        for shape in shapes:
+            with self.subTest(shape=shape):
+                compiled = self._compile(fn)
+                x = torch.randn(shape, device=self.DEVICE)
+                result = compiled(x)
+                expected = fn(x)
+                self.assertEqual(result, expected)
 
-        x = torch.randn(1024, device=self.DEVICE)
-        result = compiled(x)
-        expected = fn(x)
-        self.assertEqual(result, expected)
-
-    @skip_if_tpu
     def test_min_reduction(self):
         """Test min reduction."""
 
         def fn(x):
             return x.min()
 
-        compiled = self._compile(fn)
-
-        x = torch.randn(16, device=self.DEVICE)
-        result = compiled(x)
-        expected = fn(x)
-        self.assertEqual(result, expected)
-
-    @skip_if_tpu
+        shapes = [(16,)]
+        if self.DEVICE != "cuda":
+            shapes.append((2048,))
+        for shape in shapes:
+            with self.subTest(shape=shape):
+                compiled = self._compile(fn)
+                x = torch.randn(shape, device=self.DEVICE)
+                result = compiled(x)
+                expected = fn(x)
+                self.assertEqual(result, expected)
+
+    @skip_if_tpu(reason="reduce_prod primitive not implemented in Pallas TPU lowering")
+    @skip_if_cuda(reason="reduce_prod primitive not implemented in Pallas Mosaic GPU")
     def test_prod_reduction(self):
         """Test prod reduction."""
-        if self.DEVICE == "cuda":
-            self.skipTest("reduce_prod primitive not implemented in Pallas Mosaic GPU")
 
         def fn(x):
             # Use smaller values to avoid overflow
             return (x * 0.1).prod()
 
-        compiled = self._compile(fn)
-
-        x = torch.randn(16, device=self.DEVICE)
-        result = compiled(x)
-        expected = fn(x)
-        self.assertEqual(result, expected)
+        shapes = [(16,)]
+        if self.DEVICE != "cuda":
+            shapes.append((2048,))
+        for shape in shapes:
+            with self.subTest(shape=shape):
+                compiled = self._compile(fn)
+                x = torch.randn(shape, device=self.DEVICE)
+                result = compiled(x)
+                expected = fn(x)
+                self.assertEqual(result, expected)
 
     @skip_if_cuda
-    @skip_if_tpu
     def test_softmax_two_pass(self):
         """Test two-pass softmax (max reduction + sum reduction)."""
 
-        def fn(x):
-            return torch.softmax(x, dim=-1)
+        for shape in [(32, 64), (2048, 64)]:
+            with self.subTest(shape=shape):
+                torch._dynamo.reset()
 
-        compiled = self._compile(fn)
+                def fn(x):
+                    return torch.softmax(x, dim=-1)
 
-        x = torch.randn(32, 64, device=self.DEVICE)
-        result = compiled(x)
-        expected = fn(x)
-        self.assertEqual(result, expected)
+                compiled = self._compile(fn)
+                x = torch.randn(shape, device=self.DEVICE)
+                result = compiled(x)
+                expected = fn(x)
+                self.assertEqual(result, expected)
+
+    @skip_if_cuda
+    def test_non_stride1_reduction(self):
+        """Test reductions along non-innermost axis on square tensors.
+
+        On square tensors (e.g. 8x8), the reduction axis cannot be inferred
+        from shape alone since both dims have the same size. This verifies
+        that stride-based axis detection works for both dim=0 and dim=1.
+        """
+        x = torch.randn(8, 8, device=self.DEVICE)
+        for dim in [0, 1]:
+            with self.subTest(dim=dim):
+                torch._dynamo.reset()
+
+                def fn(x, dim=dim):
+                    return x.sum(dim)
+
+                compiled = self._compile(fn)
+                result = compiled(x)
+                expected = fn(x)
+                self.assertEqual(result, expected)
 
     @skip_if_cuda
-    @skip_if_tpu
     def test_rms_norm(self):
         """Test RMS normalization (mean-of-squares reduction + rsqrt)."""
 
-        def fn(x, weight):
-            variance = x.pow(2).mean(-1, keepdim=True)
-            x = x * torch.rsqrt(variance + 1e-6)
-            return x * weight
+        for rows, cols in [(32, 64), (2048, 64)]:
+            with self.subTest(rows=rows, cols=cols):
+                torch._dynamo.reset()
 
-        compiled = self._compile(fn)
+                def fn(x, weight):
+                    variance = x.pow(2).mean(-1, keepdim=True)
+                    x = x * torch.rsqrt(variance + 1e-6)
+                    return x * weight
 
-        x = torch.randn(32, 64, device=self.DEVICE)
-        weight = torch.randn(64, device=self.DEVICE)
-        result = compiled(x, weight)
-        expected = fn(x, weight)
-        self.assertEqual(result, expected)
+                compiled = self._compile(fn)
+                x = torch.randn(rows, cols, device=self.DEVICE)
+                weight = torch.randn(cols, device=self.DEVICE)
+                result = compiled(x, weight)
+                expected = fn(x, weight)
+                self.assertEqual(result, expected)
 
     @skip_if_cuda
-    @skip_if_tpu
     def test_welford(self):
         """Test Welford variance/mean computation (two-pass fallback)."""
 
-        def fn(x):
-            return torch.var_mean(x, dim=-1, keepdim=True)
+        for shape in [(32, 64), (2048, 64)]:
+            with self.subTest(shape=shape):
+                torch._dynamo.reset()
 
-        compiled = self._compile(fn)
+                def fn(x):
+                    return torch.var_mean(x, dim=-1, keepdim=True)
 
-        x = torch.randn(32, 64, device=self.DEVICE)
-        var_result, mean_result = compiled(x)
-        var_expected, mean_expected = fn(x)
-        self.assertEqual(mean_result, mean_expected)
-        self.assertEqual(var_result, var_expected)
+                compiled = self._compile(fn)
+                x = torch.randn(shape, device=self.DEVICE)
+                var_result, mean_result = compiled(x)
+                # Eager mode torch_tpu doesn't support lowering var_mean, so comparing with numpy
+                var_expected = np.var(x.cpu().numpy(), axis=-1, keepdims=True, ddof=1)
+                mean_expected = np.mean(x.cpu().numpy(), axis=-1, keepdims=True)
+                self.assertEqual(mean_result.cpu().numpy(), mean_expected)
+                self.assertEqual(var_result.cpu().numpy(), var_expected)
 
     @skip_if_cuda
-    @skip_if_tpu
     def test_layer_norm(self):
         """Test layer normalization (mean + variance reduction, normalize, scale + shift)."""
 
-        def fn(x, weight, bias):
-            mean = x.mean(-1, keepdim=True)
-            variance = (x - mean).pow(2).mean(-1, keepdim=True)
-            x = (x - mean) * torch.rsqrt(variance + 1e-6)
-            return x * weight + bias
+        for rows, cols in [(32, 64), (2048, 64)]:
+            with self.subTest(rows=rows, cols=cols):
+                torch._dynamo.reset()
 
-        compiled = self._compile(fn)
+                def fn(x, weight, bias):
+                    mean = x.mean(-1, keepdim=True)
+                    variance = (x - mean).pow(2).mean(-1, keepdim=True)
+                    x = (x - mean) * torch.rsqrt(variance + 1e-6)
+                    return x * weight + bias
 
-        x = torch.randn(32, 64, device=self.DEVICE)
-        weight = torch.randn(64, device=self.DEVICE)
-        bias = torch.randn(64, device=self.DEVICE)
-        result = compiled(x, weight, bias)
-        expected = fn(x, weight, bias)
-        self.assertEqual(result, expected)
+                compiled = self._compile(fn)
+                x = torch.randn(rows, cols, device=self.DEVICE)
+                weight = torch.randn(cols, device=self.DEVICE)
+                bias = torch.randn(cols, device=self.DEVICE)
+                result = compiled(x, weight, bias)
+                expected = fn(x, weight, bias)
+                self.assertEqual(result, expected)
 
     @skip_if_cuda
     @skip_if_tpu
@@ -1131,7 +1298,7 @@ def fn(x, cos, sin):
         self.assertEqual(result, expected)
 
     @skip_if_cuda
-    @skip_if_tpu
+    @skip_if_tpu  # stack+where fusion doesn't broadcast correctly on TPU yet
     def test_rope_interleaved(self):
         """Test Rotary Position Embedding with interleaved halves.
 
@@ -1156,7 +1323,7 @@ def fn(x, cos, sin):
         self.assertEqual(result, expected)
 
     @skip_if_cuda
-    @skip_if_tpu
+    @skip_if_tpu  # output last dim 10 not 128-aligned, Mosaic rejects it
     def test_chained_stride_slice(self):
         """Test that chained stride slices compose into a single strided access.
 
@@ -1174,11 +1341,75 @@ def fn(x):
         expected = fn(x)
         self.assertEqual(result, expected)
 
+    @skip_if_cuda
+    @skip_if_tpu  # store uses flatten+scatter, unsupported on Mosaic
+    def test_strided_multi_dim(self):
+        """Test strided access on multiple dimensions simultaneously."""
+
+        def fn(x):
+            return x[::2, ::3] + 1.0
+
+        compiled = self._compile(fn)
+
+        # 8 % 2 == 0 and 12 % 3 == 0
+        x = torch.randn(8, 12, device=self.DEVICE)
+        result = compiled(x)
+        expected = fn(x)
+        self.assertEqual(result, expected)
+
+    @skip_if_cuda
+    @skip_if_tpu  # falls back to flatten+gather, unsupported on Mosaic
+    def test_strided_non_divisible(self):
+        """Test strided access where dim is not divisible by stride.
+
+        Falls back to flatten+gather on CPU (blocks tiling).
+        """
+
+        def fn(x):
+            return x[::3] * 2.0
+
+        compiled = self._compile(fn)
+
+        # 16 % 3 != 0 → should fall back
+        x = torch.arange(16, dtype=torch.float32, device=self.DEVICE)
+        result = compiled(x)
+        expected = fn(x)
+        self.assertEqual(result, expected)
+
+    @skip_if_cuda
+    def test_strided_large_offset(self):
+        """Test strided access where offset >= stride (skip blocks)."""
+
+        def fn(x):
+            return x[5::2] + 1.0
+
+        compiled = self._compile(fn)
+
+        # offset=5, stride=2: skip=2, r=1 → reshape(5,2)[2:,1]
+        x = torch.arange(10, dtype=torch.float32, device=self.DEVICE)
+        result = compiled(x)
+        expected = fn(x)
+        self.assertEqual(result, expected)
+
+    @skip_if_cuda
+    @skip_if_tpu  # store uses scatter, unsupported on Mosaic
+    def test_strided_large_offset_2d(self):
+        """Test 2D strided access where offset >= stride on last dim."""
+
+        def fn(x):
+            return x[:, 5::2] * 2.0
+
+        compiled = self._compile(fn)
+
+        x = torch.randn(4, 8, device=self.DEVICE)
+        result = compiled(x)
+        expected = fn(x)
+        self.assertEqual(result, expected)
+
     @skip_if_tpu
+    @skip_if_cuda(reason="arange not supported in Pallas GPU (Mosaic) backend")
     def test_arange_multi_output(self):
         """Test arange with view and multiple outputs."""
-        if self.DEVICE == "cuda":
-            self.skipTest("arange not supported in Pallas GPU (Mosaic) backend")
 
         def fn(x):
             rng1 = torch.arange(8 * 8, dtype=torch.float32, device=x.device).view(8, 8)
@@ -1202,12 +1433,16 @@ def fn(x):
             # View float32 tensor as int32 (same byte size)
             return x.view(torch.int32)
 
-        compiled = self._compile(fn)
-
-        x = torch.randn(128, device=self.DEVICE, dtype=torch.float32)
-        result = compiled(x)
-        expected = fn(x)
-        self.assertEqual(result, expected)
+        sizes = [128]
+        if self.DEVICE != "cuda":
+            sizes.append(2048)
+        for size in sizes:
+            with self.subTest(size=size):
+                compiled = self._compile(fn)
+                x = torch.randn(size, device=self.DEVICE, dtype=torch.float32)
+                result = compiled(x)
+                expected = fn(x)
+                self.assertEqual(result, expected)
 
     def test_dtype_bitcast_float16_to_int16(self):
         """Test dtype bitcast from float16 to int16."""
@@ -1306,19 +1541,22 @@ def fn(a, b):
         self.assertEqual(result, expected)
 
     def test_warpgroup_size_2d_128x128(self):
-        """Test 2D tensor with 128x128 elements."""
+        """Test 2D tensor with 128x128 and tiling-exercising sizes."""
 
         def fn(x, y):
             return x + y
 
-        compiled = self._compile(fn)
-
-        # 128x128 = 16384 elements, multiple of 128
-        x = torch.randn(128, 128, device=self.DEVICE)
-        y = torch.randn(128, 128, device=self.DEVICE)
-        result = compiled(x, y)
-        expected = fn(x, y)
-        self.assertEqual(result, expected)
+        shapes = [(128, 128)]
+        if self.DEVICE != "cuda":
+            shapes.append((2048, 2048))
+        for shape in shapes:
+            with self.subTest(shape=shape):
+                compiled = self._compile(fn)
+                x = torch.randn(shape, device=self.DEVICE)
+                y = torch.randn(shape, device=self.DEVICE)
+                result = compiled(x, y)
+                expected = fn(x, y)
+                self.assertEqual(result, expected)
 
     def test_warpgroup_size_small_tensor(self):
         """Test with very small tensor (less than warpgroup size).
@@ -1496,7 +1734,6 @@ def fn(w_0, w_1, w_2, x):
         self.assertEqual(result, expected)
 
     @skip_if_cuda
-    @skip_if_tpu
     def test_nanogpt(self):
         """Test a minimal NanoGPT-style transformer block.
 
@@ -1563,6 +1800,815 @@ def transformer_block(x, w_q, w_k, w_v, w_proj, w_fc, w_out, mask):
         expected = transformer_block(x, w_q, w_k, w_v, w_proj, w_fc, w_out, mask)
         self.assertEqual(result, expected)
 
+    def _run_transformer_layer(
+        self,
+        seq_len,
+        hidden_dim,
+        num_heads,
+        head_dim,
+        ffn_dim,
+        atol=1e-5,
+        rtol=1.3e-6,
+    ):
+        """Run a Llama-style transformer layer forward pass and verify correctness.
+
+        Architecture: RMSNorm -> Multi-Head Attention -> Residual ->
+                      RMSNorm -> SwiGLU FFN -> Residual
+        """
+        torch._dynamo.reset()
+
+        def transformer_layer(
+            x,
+            rms_w1,
+            rms_w2,
+            w_q,
+            w_k,
+            w_v,
+            w_o,
+            w_gate,
+            w_up,
+            w_down,
+            mask,
+        ):
+            T, C = x.shape
+
+            # Pre-attention RMSNorm
+            variance = x.pow(2).mean(-1, keepdim=True)
+            h = x * torch.rsqrt(variance + 1e-6) * rms_w1
+
+            # Multi-head self-attention
+            q = (h @ w_q).view(T, num_heads, head_dim).permute(1, 0, 2)  # (H, T, D)
+            k = (h @ w_k).view(T, num_heads, head_dim).permute(1, 0, 2)
+            v = (h @ w_v).view(T, num_heads, head_dim).permute(1, 0, 2)
+
+            scale = 1.0 / (head_dim**0.5)
+            att = (q @ k.transpose(-2, -1)) * scale  # (H, T, T)
+            att = att + mask  # causal mask broadcasts (T, T) -> (H, T, T)
+            att = torch.softmax(att, dim=-1)
+            attn_out = (att @ v).permute(1, 0, 2).contiguous().view(T, C)  # (T, C)
+
+            x = x + (attn_out @ w_o)
+
+            # Pre-FFN RMSNorm
+            variance = x.pow(2).mean(-1, keepdim=True)
+            h = x * torch.rsqrt(variance + 1e-6) * rms_w2
+
+            # SwiGLU FFN
+            gate = torch.nn.functional.silu(h @ w_gate)
+            up = h @ w_up
+            x = x + ((gate * up) @ w_down)
+
+            return x
+
+        compiled = self._compile(transformer_layer)
+
+        # Scale weights by 1/sqrt(fan_in) to keep activations O(1) and
+        # avoid rsqrt amplification of reduction-order diffs.
+        s = hidden_dim**-0.5
+        w_q = torch.randn(hidden_dim, hidden_dim, device=self.DEVICE) * s
+        w_k = torch.randn(hidden_dim, hidden_dim, device=self.DEVICE) * s
+        w_v = torch.randn(hidden_dim, hidden_dim, device=self.DEVICE) * s
+        w_o = torch.randn(hidden_dim, hidden_dim, device=self.DEVICE) * s
+        w_gate = torch.randn(hidden_dim, ffn_dim, device=self.DEVICE) * s
+        w_up = torch.randn(hidden_dim, ffn_dim, device=self.DEVICE) * s
+        w_down = torch.randn(ffn_dim, hidden_dim, device=self.DEVICE) * s
+        rms_w1 = torch.ones(hidden_dim, device=self.DEVICE)
+        rms_w2 = torch.ones(hidden_dim, device=self.DEVICE)
+
+        # Causal mask (T, T) - broadcasts over heads
+        mask = torch.triu(
+            torch.full((seq_len, seq_len), float("-inf"), device=self.DEVICE),
+            diagonal=1,
+        )
+
+        x = torch.randn(seq_len, hidden_dim, device=self.DEVICE)
+
+        result = compiled(
+            x,
+            rms_w1,
+            rms_w2,
+            w_q,
+            w_k,
+            w_v,
+            w_o,
+            w_gate,
+            w_up,
+            w_down,
+            mask,
+        )
+        expected = transformer_layer(
+            x,
+            rms_w1,
+            rms_w2,
+            w_q,
+            w_k,
+            w_v,
+            w_o,
+            w_gate,
+            w_up,
+            w_down,
+            mask,
+        )
+        self.assertEqual(result, expected, atol=atol, rtol=rtol)
+
+    @skip_if_cuda
+    def test_transformer_layer_tiny(self):
+        """Test full Llama-style transformer layer at tiny dimensions."""
+        self._run_transformer_layer(
+            seq_len=32,
+            hidden_dim=64,
+            num_heads=2,
+            head_dim=32,
+            ffn_dim=256,
+        )
+
+    @skip_if_cuda
+    def test_transformer_layer_medium(self):
+        """Test full Llama-style transformer layer at Llama-7B dimensions."""
+        self._run_transformer_layer(
+            seq_len=128,
+            hidden_dim=4096,
+            num_heads=32,
+            head_dim=128,
+            ffn_dim=11008,
+            atol=2e-2,
+            rtol=1e-2,
+        )
+
+    @skip_if_cuda
+    def test_transformer_layer_large(self):
+        """Test full Llama-style transformer layer at Llama-405B dimensions."""
+        self._run_transformer_layer(
+            seq_len=32,
+            hidden_dim=16384,
+            num_heads=128,
+            head_dim=128,
+            ffn_dim=53248,
+            atol=2e-2,
+            rtol=1e-2,
+        )
+
+    @skip_if_cuda
+    def test_permute_contiguous_3d(self):
+        """Test that permute + contiguous on a 3D tensor produces correct results."""
+
+        def fn(x):
+            return x.permute(1, 0, 2).contiguous()
+
+        compiled = self._compile(fn)
+        x = torch.randn(2, 32, 32, device=self.DEVICE)
+        result = compiled(x)
+        expected = fn(x)
+        self.assertEqual(result, expected)
+
+    @skip_if_cuda
+    def test_transpose_contiguous_2d(self):
+        """Test that transpose + contiguous on a 2D tensor compiles and runs."""
+
+        def fn(x):
+            return x.transpose(0, 1).contiguous()
+
+        compiled = self._compile(fn)
+        x = torch.randn(2, 32, device=self.DEVICE)
+        result = compiled(x)
+        expected = fn(x)
+        self.assertEqual(result, expected)
+
+    @skip_if_cuda
+    def test_permute_contiguous_2d_asymmetric(self):
+        """Test transpose+contiguous on asymmetric 2D shapes."""
+
+        def fn(x):
+            return x.transpose(0, 1).contiguous()
+
+        compiled = self._compile(fn)
+        for shape in [(10000, 2), (2, 10000), (8, 128), (128, 8), (3, 256)]:
+            with self.subTest(shape=shape):
+                x = torch.randn(*shape, device=self.DEVICE)
+                result = compiled(x)
+                expected = fn(x)
+                self.assertEqual(result, expected)
+
+    @skip_if_cuda
+    def test_permute_contiguous_3d_all_perms(self):
+        """Test non-identity 3D permutations with distinct dim sizes."""
+        all_perms = [
+            (1, 0, 2),
+            (0, 2, 1),
+            (2, 1, 0),  # full-rank detection
+        ]
+        x = torch.randn(2, 1152, 2048, device=self.DEVICE)
+
+        for perm in all_perms:
+            with self.subTest(perm=perm):
+
+                def fn(x, p=perm):
+                    return x.permute(*p).contiguous()
+
+                compiled = self._compile(fn)
+                result = compiled(x)
+                expected = fn(x)
+                self.assertEqual(result, expected)
+
+    @skip_if_cuda
+    @skip_if_tpu
+    def test_permute_contiguous_3d_collapsed(self):
+        """Test 3D permutations that require collapsed-dim detection."""
+        all_perms = [
+            (2, 0, 1),
+            (1, 2, 0),
+        ]
+        x = torch.randn(2, 1152, 2048, device=self.DEVICE)
+
+        for perm in all_perms:
+            with self.subTest(perm=perm):
+
+                def fn(x, p=perm):
+                    return x.permute(*p).contiguous()
+
+                compiled = self._compile(fn)
+                result = compiled(x)
+                expected = fn(x)
+                self.assertEqual(result, expected)
+
+    @skip_if_cuda
+    def test_permute_contiguous_3d_large_notile(self):
+        """Test 3D perms with large shape but no tiling (dim=1024 exact fit)."""
+        perms = [(1, 0, 2), (0, 2, 1), (2, 1, 0)]
+
+        x = torch.randn(2, 1152, 1024, device=self.DEVICE)
+
+        for perm in perms:
+            with self.subTest(perm=perm):
+
+                def fn(x, p=perm):
+                    return x.permute(*p).contiguous()
+
+                compiled = self._compile(fn)
+                result = compiled(x)
+                expected = fn(x)
+                self.assertEqual(result, expected)
+
+    @skip_if_cuda
+    def test_permute_contiguous_3d_medium(self):
+        """Test 3D perms with medium shape that triggers tiling on last dim."""
+        perms = [(1, 0, 2), (0, 2, 1), (2, 1, 0)]
+
+        x = torch.randn(2, 8, 2048, device=self.DEVICE)
+
+        for perm in perms:
+            with self.subTest(perm=perm):
+
+                def fn(x, p=perm):
+                    return x.permute(*p).contiguous()
+
+                compiled = self._compile(fn)
+                result = compiled(x)
+                expected = fn(x)
+                self.assertEqual(result, expected)
+
+    @skip_if_cuda
+    def test_permute_contiguous_3d_small(self):
+        """Test 3D perms with small shapes that produce grid=(1,)."""
+        all_perms = [
+            (1, 0, 2),
+            (0, 2, 1),
+            (2, 1, 0),
+            (2, 0, 1),
+            (1, 2, 0),
+        ]
+        # (2,0,1) on (2,8,16) triggers a Mosaic "unsupported shape cast"
+        # bug on TPU due to internal tile padding (128x16 -> 16x128).
+        tpu_skip = {(2, 0, 1)} if self.DEVICE == "tpu" else set()
+        x = torch.randn(2, 8, 16, device=self.DEVICE)
+
+        for perm in all_perms:
+            if perm in tpu_skip:
+                continue
+            with self.subTest(perm=perm):
+
+                def fn(x, p=perm):
+                    return x.permute(*p).contiguous()
+
+                compiled = self._compile(fn)
+                result = compiled(x)
+                expected = fn(x)
+                self.assertEqual(result, expected)
+
+    @skip_if_cuda
+    @skip_if_tpu
+    def test_permute_contiguous_4d(self):
+        """Test all 23 non-identity 4D permutations with multi-tile grids."""
+        all_perms = [
+            (0, 1, 3, 2),
+            (0, 2, 1, 3),
+            (0, 2, 3, 1),
+            (0, 3, 1, 2),
+            (0, 3, 2, 1),
+            (1, 0, 2, 3),
+            (1, 0, 3, 2),
+            (1, 2, 0, 3),
+            (1, 2, 3, 0),
+            (1, 3, 0, 2),
+            (1, 3, 2, 0),
+            (2, 0, 1, 3),
+            (2, 0, 3, 1),
+            (2, 1, 0, 3),
+            (2, 1, 3, 0),
+            (2, 3, 0, 1),
+            (2, 3, 1, 0),
+            (3, 0, 1, 2),
+            (3, 0, 2, 1),
+            (3, 1, 0, 2),
+            (3, 1, 2, 0),
+            (3, 2, 0, 1),
+            (3, 2, 1, 0),
+        ]
+        needs_large_first = {
+            (0, 3, 2, 1),
+            (1, 3, 0, 2),
+            (1, 3, 2, 0),
+            (2, 3, 1, 0),
+            (3, 0, 2, 1),
+            (3, 1, 0, 2),
+            (3, 1, 2, 0),
+            (3, 2, 0, 1),
+            (3, 2, 1, 0),
+        }
+        shapes = {
+            "large_last": (2, 4, 128, 2048),
+            "large_first": (1152, 1152, 2, 4),
+        }
+
+        for perm in all_perms:
+            key = "large_first" if perm in needs_large_first else "large_last"
+            shape = shapes[key]
+            with self.subTest(perm=perm, shape=shape):
+                x = torch.randn(*shape, device=self.DEVICE)
+
+                def fn(x, p=perm):
+                    return x.permute(*p).contiguous()
+
+                compiled = self._compile(fn)
+                result = compiled(x)
+                expected = fn(x)
+                self.assertEqual(result, expected)
+
+    def _run_transformer(
+        self,
+        num_layers,
+        seq_len,
+        hidden_dim,
+        num_heads,
+        head_dim,
+        ffn_dim,
+        atol=1e-5,
+        rtol=1.3e-6,
+    ):
+        """Run a multi-layer Llama-style transformer and verify correctness."""
+        torch._dynamo.reset()
+
+        def transformer(x, mask, *layer_params):
+            T, C = x.shape
+            params_per_layer = (
+                9  # rms_w1, rms_w2, w_q, w_k, w_v, w_o, w_gate, w_up, w_down
+            )
+
+            for i in range(num_layers):
+                offset = i * params_per_layer
+                rms_w1 = layer_params[offset]
+                rms_w2 = layer_params[offset + 1]
+                w_q = layer_params[offset + 2]
+                w_k = layer_params[offset + 3]
+                w_v = layer_params[offset + 4]
+                w_o = layer_params[offset + 5]
+                w_gate = layer_params[offset + 6]
+                w_up = layer_params[offset + 7]
+                w_down = layer_params[offset + 8]
+
+                # Pre-attention RMSNorm
+                variance = x.pow(2).mean(-1, keepdim=True)
+                h = x * torch.rsqrt(variance + 1e-6) * rms_w1
+
+                # Multi-head self-attention
+                q = (h @ w_q).view(T, num_heads, head_dim).permute(1, 0, 2)
+                k = (h @ w_k).view(T, num_heads, head_dim).permute(1, 0, 2)
+                v = (h @ w_v).view(T, num_heads, head_dim).permute(1, 0, 2)
+
+                scale = 1.0 / (head_dim**0.5)
+                att = (q @ k.transpose(-2, -1)) * scale
+                att = att + mask
+                att = torch.softmax(att, dim=-1)
+                attn_out = (att @ v).permute(1, 0, 2).contiguous().view(T, C)
+
+                x = x + (attn_out @ w_o)
+
+                # Pre-FFN RMSNorm
+                variance = x.pow(2).mean(-1, keepdim=True)
+                h = x * torch.rsqrt(variance + 1e-6) * rms_w2
+
+                # SwiGLU FFN
+                gate = torch.nn.functional.silu(h @ w_gate)
+                up = h @ w_up
+                x = x + ((gate * up) @ w_down)
+
+            return x
+
+        compiled = self._compile(transformer)
+
+        s = hidden_dim**-0.5
+        all_params = []
+        for _ in range(num_layers):
+            all_params.extend(
+                [
+                    torch.ones(hidden_dim, device=self.DEVICE),  # rms_w1
+                    torch.ones(hidden_dim, device=self.DEVICE),  # rms_w2
+                    torch.randn(hidden_dim, hidden_dim, device=self.DEVICE) * s,  # w_q
+                    torch.randn(hidden_dim, hidden_dim, device=self.DEVICE) * s,  # w_k
+                    torch.randn(hidden_dim, hidden_dim, device=self.DEVICE) * s,  # w_v
+                    torch.randn(hidden_dim, hidden_dim, device=self.DEVICE) * s,  # w_o
+                    torch.randn(hidden_dim, ffn_dim, device=self.DEVICE) * s,  # w_gate
+                    torch.randn(hidden_dim, ffn_dim, device=self.DEVICE) * s,  # w_up
+                    torch.randn(ffn_dim, hidden_dim, device=self.DEVICE) * s,  # w_down
+                ]
+            )
+
+        mask = torch.triu(
+            torch.full((seq_len, seq_len), float("-inf"), device=self.DEVICE),
+            diagonal=1,
+        )
+        x = torch.randn(seq_len, hidden_dim, device=self.DEVICE)
+
+        result = compiled(x, mask, *all_params)
+        expected = transformer(x, mask, *all_params)
+        self.assertEqual(result, expected, atol=atol, rtol=rtol)
+
+    @skip_if_cuda
+    def test_transformer_tiny(self):
+        """Test a 4-layer Llama-style transformer at tiny dimensions."""
+        self._run_transformer(
+            num_layers=4,
+            seq_len=32,
+            hidden_dim=64,
+            num_heads=2,
+            head_dim=32,
+            ffn_dim=256,
+            atol=1e-2,
+            rtol=1e-2,
+        )
+
+    @skip_if_cuda
+    def test_transformer_medium(self):
+        """Test a 4-layer transformer at Llama-7B-like dimensions."""
+        self._run_transformer(
+            num_layers=4,
+            seq_len=128,
+            hidden_dim=4096,
+            num_heads=32,
+            head_dim=128,
+            ffn_dim=11008,
+            atol=0.1,
+            rtol=1e-2,
+        )
+
+    def _run_transformer_lm(
+        self,
+        num_layers,
+        seq_len,
+        vocab_size,
+        hidden_dim,
+        num_heads,
+        head_dim,
+        ffn_dim,
+        atol=1e-5,
+        rtol=1.3e-6,
+    ):
+        """Run a full Llama-style LM (embedding + layers + norm + lm_head)."""
+        torch._dynamo.reset()
+
+        def transformer_lm(
+            token_ids, embed_table, final_rms_w, lm_head_w, mask, *layer_params
+        ):
+            # Token embedding via gather
+            x = embed_table[token_ids]  # (T, C)
+            T, C = x.shape
+            params_per_layer = 9
+
+            for i in range(num_layers):
+                offset = i * params_per_layer
+                rms_w1 = layer_params[offset]
+                rms_w2 = layer_params[offset + 1]
+                w_q = layer_params[offset + 2]
+                w_k = layer_params[offset + 3]
+                w_v = layer_params[offset + 4]
+                w_o = layer_params[offset + 5]
+                w_gate = layer_params[offset + 6]
+                w_up = layer_params[offset + 7]
+                w_down = layer_params[offset + 8]
+
+                # Pre-attention RMSNorm
+                variance = x.pow(2).mean(-1, keepdim=True)
+                h = x * torch.rsqrt(variance + 1e-6) * rms_w1
+
+                # Multi-head self-attention
+                q = (h @ w_q).view(T, num_heads, head_dim).permute(1, 0, 2)
+                k = (h @ w_k).view(T, num_heads, head_dim).permute(1, 0, 2)
+                v = (h @ w_v).view(T, num_heads, head_dim).permute(1, 0, 2)
+
+                scale = 1.0 / (head_dim**0.5)
+                att = (q @ k.transpose(-2, -1)) * scale
+                att = att + mask
+                att = torch.softmax(att, dim=-1)
+                attn_out = (att @ v).permute(1, 0, 2).contiguous().view(T, C)
+
+                x = x + (attn_out @ w_o)
+
+                # Pre-FFN RMSNorm
+                variance = x.pow(2).mean(-1, keepdim=True)
+                h = x * torch.rsqrt(variance + 1e-6) * rms_w2
+
+                # SwiGLU FFN
+                gate = torch.nn.functional.silu(h @ w_gate)
+                up = h @ w_up
+                x = x + ((gate * up) @ w_down)
+
+            # Final RMSNorm
+            variance = x.pow(2).mean(-1, keepdim=True)
+            x = x * torch.rsqrt(variance + 1e-6) * final_rms_w
+
+            # LM head
+            logits = x @ lm_head_w  # (T, V)
+            return logits
+
+        compiled = self._compile(transformer_lm)
+
+        s = hidden_dim**-0.5
+        embed_table = torch.randn(vocab_size, hidden_dim, device=self.DEVICE) * s
+        final_rms_w = torch.ones(hidden_dim, device=self.DEVICE)
+        lm_head_w = torch.randn(hidden_dim, vocab_size, device=self.DEVICE) * s
+
+        all_params = []
+        for _ in range(num_layers):
+            all_params.extend(
+                [
+                    torch.ones(hidden_dim, device=self.DEVICE),
+                    torch.ones(hidden_dim, device=self.DEVICE),
+                    torch.randn(hidden_dim, hidden_dim, device=self.DEVICE) * s,
+                    torch.randn(hidden_dim, hidden_dim, device=self.DEVICE) * s,
+                    torch.randn(hidden_dim, hidden_dim, device=self.DEVICE) * s,
+                    torch.randn(hidden_dim, hidden_dim, device=self.DEVICE) * s,
+                    torch.randn(hidden_dim, ffn_dim, device=self.DEVICE) * s,
+                    torch.randn(hidden_dim, ffn_dim, device=self.DEVICE) * s,
+                    torch.randn(ffn_dim, hidden_dim, device=self.DEVICE) * s,
+                ]
+            )
+
+        mask = torch.triu(
+            torch.full((seq_len, seq_len), float("-inf"), device=self.DEVICE),
+            diagonal=1,
+        )
+        token_ids = torch.randint(0, vocab_size, (seq_len,), device=self.DEVICE)
+
+        result = compiled(
+            token_ids, embed_table, final_rms_w, lm_head_w, mask, *all_params
+        )
+        expected = transformer_lm(
+            token_ids, embed_table, final_rms_w, lm_head_w, mask, *all_params
+        )
+        self.assertEqual(result, expected, atol=atol, rtol=rtol)
+
+    @unittest.skip("numerical mismatch in embedding + RMSNorm fusion")
+    def test_transformer_lm_tiny(self):
+        """Test a full LM (embed + 4 layers + norm + lm_head) at tiny dims."""
+        self._run_transformer_lm(
+            num_layers=4,
+            seq_len=32,
+            vocab_size=256,
+            hidden_dim=64,
+            num_heads=2,
+            head_dim=32,
+            ffn_dim=256,
+        )
+
+    @skip_if_cuda(reason="scalar prefetch not supported in Pallas GPU (Mosaic) backend")
+    def test_embedding_lookup(self):
+        """Test simple embedding table lookup (integer indexing)."""
+
+        def fn(token_ids, embed_table):
+            return embed_table[token_ids]
+
+        compiled = self._compile(fn)
+        embed_table = torch.randn(256, 64, device=self.DEVICE)
+        token_ids = torch.randint(0, 256, (32,), device=self.DEVICE)
+        result = compiled(token_ids, embed_table)
+        expected = fn(token_ids, embed_table)
+        self.assertEqual(result, expected)
+
+    @skip_if_cuda(reason="scalar prefetch not supported in Pallas GPU (Mosaic) backend")
+    def test_indirect_access_basic(self):
+        """Test bare embedding lookup via indirect access detection."""
+
+        def fn(indices, table):
+            return table[indices]
+
+        compiled = self._compile(fn)
+        table = torch.randn(256, 64, device=self.DEVICE)
+        indices = torch.randint(0, 256, (32,), device=self.DEVICE)
+        result = compiled(indices, table)
+        expected = fn(indices, table)
+        self.assertEqual(result, expected)
+
+    @skip_if_cuda(reason="scalar prefetch not supported in Pallas GPU (Mosaic) backend")
+    def test_indirect_access_non_128_dim(self):
+        """Test indirect access with D not divisible by 128."""
+
+        def fn(indices, table):
+            return table[indices]
+
+        compiled = self._compile(fn)
+        table = torch.randn(256, 100, device=self.DEVICE)
+        indices = torch.randint(0, 256, (32,), device=self.DEVICE)
+        result = compiled(indices, table)
+        expected = fn(indices, table)
+        self.assertEqual(result, expected)
+
+    @skip_if_cuda(reason="scalar prefetch not supported in Pallas GPU (Mosaic) backend")
+    def test_indirect_access_large_vocab(self):
+        """Test indirect access with larger vocabulary."""
+
+        def fn(indices, table):
+            return table[indices]
+
+        compiled = self._compile(fn)
+        table = torch.randn(2048, 128, device=self.DEVICE)
+        indices = torch.randint(0, 2048, (64,), device=self.DEVICE)
+        result = compiled(indices, table)
+        expected = fn(indices, table)
+        self.assertEqual(result, expected)
+
+    @skip_if_cuda(reason="scalar prefetch not supported in Pallas GPU (Mosaic) backend")
+    def test_indirect_access_fused_add(self):
+        """Test embedding + pointwise add fused."""
+
+        def fn(indices, table, bias):
+            return table[indices] + bias
+
+        compiled = self._compile(fn)
+        table = torch.randn(256, 64, device=self.DEVICE)
+        indices = torch.randint(0, 256, (32,), device=self.DEVICE)
+        bias = torch.randn(32, 64, device=self.DEVICE)
+        result = compiled(indices, table, bias)
+        expected = fn(indices, table, bias)
+        self.assertEqual(result, expected)
+
+    @skip_if_cuda(reason="scalar prefetch not supported in Pallas GPU (Mosaic) backend")
+    def test_indirect_access_fused_mul(self):
+        """Test embedding + pointwise multiply with broadcast."""
+
+        def fn(indices, table, scale):
+            return table[indices] * scale
+
+        compiled = self._compile(fn)
+        table = torch.randn(256, 64, device=self.DEVICE)
+        indices = torch.randint(0, 256, (32,), device=self.DEVICE)
+        scale = torch.randn(64, device=self.DEVICE)
+        result = compiled(indices, table, scale)
+        expected = fn(indices, table, scale)
+        self.assertEqual(result, expected)
+
+    @skip_if_cuda(reason="scalar prefetch not supported in Pallas GPU (Mosaic) backend")
+    def test_indirect_access_fused_chain(self):
+        """Test embedding + chained add and multiply."""
+
+        def fn(indices, table, bias, scale):
+            x = table[indices] + bias
+            return x * scale
+
+        compiled = self._compile(fn)
+        table = torch.randn(256, 64, device=self.DEVICE)
+        indices = torch.randint(0, 256, (32,), device=self.DEVICE)
+        bias = torch.randn(32, 64, device=self.DEVICE)
+        scale = torch.randn(64, device=self.DEVICE)
+        result = compiled(indices, table, bias, scale)
+        expected = fn(indices, table, bias, scale)
+        self.assertEqual(result, expected)
+
+    @skip_if_cuda(reason="scalar prefetch not supported in Pallas GPU (Mosaic) backend")
+    def test_nn_embedding(self):
+        """Test nn.functional.embedding path through indirect access."""
+
+        def fn(indices, weight):
+            return torch.nn.functional.embedding(indices, weight)
+
+        compiled = self._compile(fn)
+        weight = torch.randn(256, 64, device=self.DEVICE)
+        indices = torch.randint(0, 256, (32,), device=self.DEVICE)
+        result = compiled(indices, weight)
+        expected = fn(indices, weight)
+        self.assertEqual(result, expected)
+
+    @skip_if_cuda(reason="scalar prefetch not supported in Pallas GPU (Mosaic) backend")
+    def test_indirect_access_single_token(self):
+        """Test indirect access with seq=1."""
+
+        def fn(indices, table):
+            return table[indices]
+
+        compiled = self._compile(fn)
+        table = torch.randn(256, 64, device=self.DEVICE)
+        indices = torch.randint(0, 256, (1,), device=self.DEVICE)
+        result = compiled(indices, table)
+        expected = fn(indices, table)
+        self.assertEqual(result, expected)
+
+    @skip_if_cuda(reason="scalar prefetch not supported in Pallas GPU (Mosaic) backend")
+    def test_indirect_access_duplicate_indices(self):
+        """Test indirect access with repeated indices."""
+
+        def fn(indices, table):
+            return table[indices]
+
+        compiled = self._compile(fn)
+        table = torch.randn(256, 64, device=self.DEVICE)
+        indices = torch.tensor([0, 1, 0, 1, 2, 2, 3, 3], device=self.DEVICE)
+        result = compiled(indices, table)
+        expected = fn(indices, table)
+        self.assertEqual(result, expected)
+
+    @skip_if_cuda
+    def test_transformer_with_final_norm_and_lm_head(self):
+        """Test multi-layer transformer + final RMSNorm + LM head (no embedding)."""
+        torch._dynamo.reset()
+        num_layers = 4
+        seq_len = 32
+        hidden_dim = 64
+        num_heads = 2
+        head_dim = 32
+        ffn_dim = 256
+        vocab_size = 256
+
+        def transformer_with_head(x, final_rms_w, lm_head_w, mask, *layer_params):
+            T, C = x.shape
+            params_per_layer = 9
+            for i in range(num_layers):
+                offset = i * params_per_layer
+                rms_w1 = layer_params[offset]
+                rms_w2 = layer_params[offset + 1]
+                w_q = layer_params[offset + 2]
+                w_k = layer_params[offset + 3]
+                w_v = layer_params[offset + 4]
+                w_o = layer_params[offset + 5]
+                w_gate = layer_params[offset + 6]
+                w_up = layer_params[offset + 7]
+                w_down = layer_params[offset + 8]
+                variance = x.pow(2).mean(-1, keepdim=True)
+                h = x * torch.rsqrt(variance + 1e-6) * rms_w1
+                q = (h @ w_q).view(T, num_heads, head_dim).permute(1, 0, 2)
+                k = (h @ w_k).view(T, num_heads, head_dim).permute(1, 0, 2)
+                v = (h @ w_v).view(T, num_heads, head_dim).permute(1, 0, 2)
+                scale = 1.0 / (head_dim**0.5)
+                att = (q @ k.transpose(-2, -1)) * scale
+                att = att + mask
+                att = torch.softmax(att, dim=-1)
+                attn_out = (att @ v).permute(1, 0, 2).contiguous().view(T, C)
+                x = x + (attn_out @ w_o)
+                variance = x.pow(2).mean(-1, keepdim=True)
+                h = x * torch.rsqrt(variance + 1e-6) * rms_w2
+                gate = torch.nn.functional.silu(h @ w_gate)
+                up = h @ w_up
+                x = x + ((gate * up) @ w_down)
+            variance = x.pow(2).mean(-1, keepdim=True)
+            x = x * torch.rsqrt(variance + 1e-6) * final_rms_w
+            return x @ lm_head_w
+
+        compiled = self._compile(transformer_with_head)
+        s = hidden_dim**-0.5
+        final_rms_w = torch.ones(hidden_dim, device=self.DEVICE)
+        lm_head_w = torch.randn(hidden_dim, vocab_size, device=self.DEVICE) * s
+        all_params = []
+        for _ in range(num_layers):
+            all_params.extend(
+                [
+                    torch.ones(hidden_dim, device=self.DEVICE),
+                    torch.ones(hidden_dim, device=self.DEVICE),
+                    torch.randn(hidden_dim, hidden_dim, device=self.DEVICE) * s,
+                    torch.randn(hidden_dim, hidden_dim, device=self.DEVICE) * s,
+                    torch.randn(hidden_dim, hidden_dim, device=self.DEVICE) * s,
+                    torch.randn(hidden_dim, hidden_dim, device=self.DEVICE) * s,
+                    torch.randn(hidden_dim, ffn_dim, device=self.DEVICE) * s,
+                    torch.randn(hidden_dim, ffn_dim, device=self.DEVICE) * s,
+                    torch.randn(ffn_dim, hidden_dim, device=self.DEVICE) * s,
+                ]
+            )
+        mask = torch.triu(
+            torch.full((seq_len, seq_len), float("-inf"), device=self.DEVICE),
+            diagonal=1,
+        )
+        x = torch.randn(seq_len, hidden_dim, device=self.DEVICE)
+        result = compiled(x, final_rms_w, lm_head_w, mask, *all_params)
+        expected = transformer_with_head(x, final_rms_w, lm_head_w, mask, *all_params)
+        self.assertEqual(result, expected)
+
     def test_warpgroup_size_2d_aligned_32x8(self):
         """Test 2D tensor with 32x8 = 256 elements (2 warpgroups)."""
 
diff --git a/test/inductor/test_pattern_matcher.py b/test/inductor/test_pattern_matcher.py
index 08575462bef5f..8ed4ce990d35a 100644
--- a/test/inductor/test_pattern_matcher.py
+++ b/test/inductor/test_pattern_matcher.py
@@ -3,7 +3,6 @@
 import os
 import unittest
 from collections.abc import Callable
-from typing import Optional
 
 import torch
 import torch._dynamo.config as dynamo_config
@@ -31,6 +30,11 @@
 from torch._inductor.test_case import run_tests, TestCase
 from torch._inductor.utils import run_and_get_code
 from torch._inductor.virtualized import V
+from torch._library.opaque_object import (
+    get_opaque_type_name,
+    OpaqueBase,
+    register_opaque_type,
+)
 from torch.fx.experimental.proxy_tensor import make_fx
 from torch.testing import FileCheck
 from torch.testing._internal.common_cuda import SM80OrLater, xfailIfSM89
@@ -48,6 +52,26 @@
 aten = torch.ops.aten
 
 
+class OpaqueScaleFactor(OpaqueBase):
+    def __init__(self, val):
+        self.val = val
+
+    def __eq__(self, other):
+        return isinstance(other, OpaqueScaleFactor) and self.val == other.val
+
+    def __hash__(self):
+        return hash(self.val)
+
+    def __fx_repr__(self):
+        return (
+            f"OpaqueScaleFactor({self.val!r})",
+            {"OpaqueScaleFactor": OpaqueScaleFactor},
+        )
+
+
+register_opaque_type(OpaqueScaleFactor, typ="value", hoist=True)
+
+
 @instantiate_parametrized_tests
 class TestPatternMatcher(TestCase):
     device_type = GPU_TYPE
@@ -137,11 +161,10 @@ def _test_fused_int_mm_mul_impl(self, fn, args, fused_int_mm_mul_expected=True):
                 ref[indices], test[indices]
             )  # also checks that dtype is correct
 
-    # @skipIfXpu
     @skipCUDAIf(not SM80OrLater, "need sm_80")
     @inductor_config.patch(
         {
-            "benchmark_epilogue_fusion": "False",
+            "benchmark_epilogue_fusion": False,
             "max_autotune_gemm_backends": "TRITON",
             "max_autotune_gemm": True,
         }
@@ -240,7 +263,7 @@ def f_replaced(x: torch.Tensor) -> torch.Tensor:
     @skipCUDAIf(not SM80OrLater, "need sm_80")
     @inductor_config.patch(
         {
-            "benchmark_epilogue_fusion": "False",
+            "benchmark_epilogue_fusion": False,
             "max_autotune_gemm_backends": "TRITON",
             "max_autotune_gemm": True,
         }
@@ -287,7 +310,7 @@ def fn2(a, b, c):
     @skipCUDAIf(not SM80OrLater, "need sm_80")
     @inductor_config.patch(
         {
-            "benchmark_epilogue_fusion": "False",
+            "benchmark_epilogue_fusion": False,
             "max_autotune_gemm_backends": "TRITON",
             "max_autotune_gemm": True,
         }
@@ -334,7 +357,8 @@ def _test_mixed_impl(
     @skipCUDAIf(not SM80OrLater, "need sm_80")
     @inductor_config.patch(
         {
-            "benchmark_epilogue_fusion": "False",
+            "benchmark_fusion": False,
+            "benchmark_epilogue_fusion": False,
             "max_autotune_gemm_backends": "TRITON",
             "max_autotune_gemm": True,
         }
@@ -393,7 +417,8 @@ def fn(a, b):
     @skipCUDAIf(not SM80OrLater, "need sm_80")
     @inductor_config.patch(
         {
-            "benchmark_epilogue_fusion": "False",
+            "benchmark_fusion": False,
+            "benchmark_epilogue_fusion": False,
             "max_autotune_gemm_backends": "TRITON",
             "max_autotune_gemm": True,
         }
@@ -424,7 +449,8 @@ def fn(a, b):
     @skipCUDAIf(not SM80OrLater, "need sm_80")
     @inductor_config.patch(
         {
-            "benchmark_epilogue_fusion": "False",
+            "benchmark_fusion": False,
+            "benchmark_epilogue_fusion": False,
             "max_autotune_gemm_backends": "TRITON",
             "max_autotune_gemm": True,
         }
@@ -473,7 +499,8 @@ def fn(a, b):
 
         with inductor_config.patch(
             {
-                "benchmark_epilogue_fusion": "False",
+                "benchmark_fusion": False,
+                "benchmark_epilogue_fusion": False,
                 "max_autotune_gemm_backends": "TRITON",
                 "max_autotune_gemm": True,
             }
@@ -1212,6 +1239,50 @@ def fn2(inp, a, b):
         _, (code) = run_and_get_code(fn2, args[0], args[1], args[2])
         FileCheck().check_not("extern_kernels.addmm(").run(code[0])
 
+    @parametrize("dtype", [torch.bfloat16, torch.float16])
+    @inductor_config.patch(
+        {
+            "fx_graph_remote_cache": False,
+            "keep_addmm_fused_for_half_dtypes": True,
+        }
+    )
+    def test_unfuse_bias_addmm_half_dtypes(self, dtype):
+        args = [
+            torch.randn(20, device=GPU_TYPE, dtype=dtype),
+            torch.randn(10, 15, device=GPU_TYPE, dtype=dtype),
+            torch.randn(15, 20, device=GPU_TYPE, dtype=dtype),
+        ]
+
+        # addmm with pointwise consumer should not be unfused for half dtypes
+        # to avoid precision loss from extra truncation at the mm output
+        @torch.compile()
+        def fn(inp, a, b):
+            return torch.nn.functional.gelu(torch.ops.aten.addmm(inp, a, b))
+
+        _, (code) = run_and_get_code(fn, args[0], args[1], args[2])
+        FileCheck().check("extern_kernels.addmm(").run(code[0])
+
+    @parametrize("dtype", [torch.bfloat16, torch.float16])
+    @inductor_config.patch(
+        {
+            "fx_graph_remote_cache": False,
+            "keep_addmm_fused_for_half_dtypes": False,
+        }
+    )
+    def test_unfuse_bias_addmm_half_dtypes_when_flag_disabled(self, dtype):
+        args = [
+            torch.randn(20, device=GPU_TYPE, dtype=dtype),
+            torch.randn(10, 15, device=GPU_TYPE, dtype=dtype),
+            torch.randn(15, 20, device=GPU_TYPE, dtype=dtype),
+        ]
+
+        @torch.compile()
+        def fn(inp, a, b):
+            return torch.nn.functional.gelu(torch.ops.aten.addmm(inp, a, b))
+
+        _, (code) = run_and_get_code(fn, args[0], args[1], args[2])
+        FileCheck().check_not("extern_kernels.addmm(").run(code[0])
+
     def test_addmm_alpha_beta_with_pointwise(self):
         # Test that addmm with alpha/beta != 1 is unfused correctly with pointwise ops
         # See https://github.com/pytorch/pytorch/issues/167313
@@ -1301,6 +1372,7 @@ def remap_fake_tensor(x):
             "max_autotune_gemm_backends": "TRITON",
         }
     )
+    @unittest.skipIf(not IS_BIG_GPU, "templates require big gpu")
     def test_original_aten_preserved_split_addmm(self):
         # addmm -> elementwise should be decomposed into mm -> add -> elementwise
         def fn(x, y, z):
@@ -1377,10 +1449,58 @@ def repl(inp, x1, x2):
                 actual, (code) = run_and_get_code(opt_fn, args[0], args[1], args[2])
                 # pattern should match
                 self.assertEqual(counter, 1)
-                torch.testing.assert_close(actual, expected)
+                self.assertEqual(actual, expected)
                 # addmm should be replaced
                 FileCheck().check_not("extern_kernels.addmm(").run(code[0])
 
+    @inductor_config.patch(fx_graph_remote_cache=False)
+    def test_replace_by_example_in_pre_grad(self):
+        counter = 0
+        test_pass = PatternMatcherPass()
+
+        args = [
+            torch.randn(20, device=GPU_TYPE),
+            torch.randn(10, 15, device=GPU_TYPE),
+            torch.randn(15, 20, device=GPU_TYPE),
+        ]
+
+        def fn(inp, a, b):
+            return torch.ops.aten.addmm(inp, a, b)
+
+        # This graph pattern should successfully match all of the above functions
+        @register_graph_pattern(
+            CallFunction(
+                torch.ops.aten.addmm,
+                Arg(),
+                Arg(),
+                Arg(),
+                beta=KeywordArg("beta"),
+                alpha=KeywordArg("alpha"),
+            ),
+            pass_dict=test_pass,
+        )
+        def addmm_replacement(match: Match, inp, mat1, mat2, beta, alpha):
+            nonlocal counter
+            counter += 1
+
+            def repl(inp, x1, x2):
+                return (x1 @ x2) * alpha + inp * beta
+
+            match.replace_by_example(repl, [inp, mat1, mat2])
+
+        with inductor_config.patch(
+            pre_grad_custom_pass=lambda *args: test_pass.apply(*args)
+        ):
+            counter = 0
+            expected = fn(*copy.deepcopy(args))
+            opt_fn = torch.compile(fn)
+            actual, (code) = run_and_get_code(opt_fn, args[0], args[1], args[2])
+            # pattern should match
+            self.assertEqual(counter, 1)
+            self.assertEqual(actual, expected)
+            # addmm should be replaced
+            FileCheck().check_not("extern_kernels.addmm(").run(code[0])
+
     def test_addmm_dtype_mismatch(self):
         a = torch.nn.Linear(1024, 1024, bias=False).to(GPU_TYPE)
         a = a.to(dtype=torch.float16)
@@ -1607,7 +1727,8 @@ def div_softmax(x, inv_scale):
 
     def test_mutation_op_matching(self):
         def check(type, func_name, args, kwargs, expect=True):
-            assert type in ["call_function", "call_method"]
+            if type not in ["call_function", "call_method"]:
+                raise AssertionError
             graph = torch.fx.Graph()
             getattr(graph, type)(func_name, args, kwargs)
             res = is_mutation_op(next(iter(graph.nodes)))
@@ -1706,7 +1827,7 @@ def static_scaled_int8_quant(
             result: torch.Tensor,
             input: torch.Tensor,
             scale: torch.Tensor,
-            azp: Optional[torch.Tensor] = None,
+            azp: torch.Tensor | None = None,
         ) -> None:
             # bogus implementation doesn't matter
             _result = torch.mul(input, scale).to(torch.int8)
@@ -1826,6 +1947,60 @@ def f(x):
         self.assertEqual(len(sigmoid_nodes), 1)
         self.assertTrue("original_aten" in sigmoid_nodes[0].meta)
 
+    def test_fwd_only_uses_get_decomp_fn(self):
+        """fwd_only traces the pattern graph using the table from get_decomp_fn."""
+
+        def fn(x):
+            return F.gelu(x)
+
+        x = torch.randn(4, 4, device=GPU_TYPE)
+
+        # Default: gelu decomposes into erf/mul/add primitives.
+        gm = fwd_only(fn, args=[x])
+        targets = {n.target for n in gm.graph.nodes if n.op == "call_function"}
+        self.assertNotIn(aten.gelu.default, targets)
+        self.assertIn(aten.erf.default, targets)
+
+        # Empty decomp table: gelu stays intact as aten.gelu.default.
+        gm_nodec = fwd_only(fn, args=[x], get_decomp_fn=dict)
+        targets_nodec = {
+            n.target for n in gm_nodec.graph.nodes if n.op == "call_function"
+        }
+        self.assertIn(aten.gelu.default, targets_nodec)
+        self.assertNotIn(aten.erf.default, targets_nodec)
+
+    def test_register_replacement_get_decomp_fn(self):
+        """A pattern registered with get_decomp_fn matches graphs traced with
+        the same decomposition table, and does not match graphs traced with a
+        different one."""
+
+        def gelu_pattern(x):
+            return F.gelu(x)
+
+        def gelu_double(x):
+            return F.gelu(x) * 2.0
+
+        x = torch.randn(4, 4, device=GPU_TYPE)
+
+        # Pattern traced without decompositions: stored as aten.gelu.default.
+        my_patterns = PatternMatcherPass()
+        register_replacement(
+            gelu_pattern,
+            gelu_double,
+            [x],
+            fwd_only,
+            my_patterns,
+            get_decomp_fn=dict,
+        )
+
+        # Graph where gelu is also not decomposed: pattern should match once.
+        gm_nodec = make_fx(gelu_pattern, {})(x)
+        self.assertEqual(my_patterns.apply(gm_nodec.graph), 1)
+
+        # Graph where gelu is decomposed (default decomps): no match.
+        gm_decomposed = fwd_only(gelu_pattern, args=[x])
+        self.assertEqual(my_patterns.apply(gm_decomposed.graph), 0)
+
     @inductor_config.patch(is_predispatch=True)
     def test_remove_noop_pass_with_remove_passes(self):
         def fn_with_noop(x):
@@ -2000,6 +2175,214 @@ def fn(x):
         self.assertEqual(result, x * 3)
         self.assertEqual(count, 1)
 
+    def test_register_replacement_single_tensor_input(self):
+        def pattern(x):
+            return x + 1
+
+        def replacement(x):
+            return x - 1
+
+        my_patterns = PatternMatcherPass()
+
+        # Single tensor should fail fast instead of reaching tracing logic.
+        example_input = torch.randn(4, 4, device=GPU_TYPE, requires_grad=True)
+        with self.assertRaisesRegex(
+            TypeError,
+            f"example_inputs must be a list or tuple, got {type(example_input)}",
+        ):
+            register_replacement(
+                pattern, replacement, example_input, fwd_only, my_patterns
+            )
+
+    def _inject_test_metadata(self, graph):
+        """Inject identifiable metadata on all call_function nodes for testing."""
+        for node in graph.nodes:
+            if node.op == "call_function":
+                node.meta["stack_trace"] = f"trace_for_{node.name}"
+                node.meta["nn_module_stack"] = {"test": ("m", "M")}
+                node.meta["_numeric_debug_handle"] = 42
+                node.meta["custom"] = {"test_key": "test_value"}
+
+    def test_metadata_propagation_register_replacement(self):
+        """Verify metadata from matched nodes transfers to replacement nodes."""
+
+        def pattern(x, y):
+            return x + y
+
+        def replacement(x, y):
+            return x * y
+
+        my_patterns = PatternMatcherPass()
+        inputs = [
+            torch.randn(4, 4, device=GPU_TYPE),
+            torch.randn(4, 4, device=GPU_TYPE),
+        ]
+        register_replacement(pattern, replacement, inputs, fwd_only, my_patterns)
+
+        def custom_pass(graph: torch.fx.Graph):
+            self._inject_test_metadata(graph)
+            # _transfer_meta runs inside replace_with_graph for
+            # each old->new pair to propagate metadata fields
+            my_patterns.apply(graph)
+
+        def fn(x, y):
+            return x + y
+
+        x = torch.randn(4, 4, device=GPU_TYPE)
+        y = torch.randn(4, 4, device=GPU_TYPE)
+
+        compiled_fn = torch.compile(
+            fn, options={"post_grad_custom_post_pass": custom_pass}
+        )
+        compiled_fn(x, y)
+
+    def test_metadata_propagation_register_replacement_multinode(self):
+        """Verify metadata propagation for multi-node patterns."""
+
+        def pattern(x, y):
+            tmp = x + y
+            return tmp * 2
+
+        def replacement(x, y):
+            return (x + y) * 3
+
+        my_patterns = PatternMatcherPass()
+        inputs = [
+            torch.randn(4, 4, device=GPU_TYPE),
+            torch.randn(4, 4, device=GPU_TYPE),
+        ]
+        register_replacement(pattern, replacement, inputs, fwd_only, my_patterns)
+
+        def custom_pass(graph: torch.fx.Graph):
+            self._inject_test_metadata(graph)
+            my_patterns.apply(graph)
+
+        def fn(x, y):
+            tmp = x + y
+            return tmp * 2
+
+        x = torch.randn(4, 4, device=GPU_TYPE)
+        y = torch.randn(4, 4, device=GPU_TYPE)
+
+        compiled_fn = torch.compile(
+            fn, options={"post_grad_custom_post_pass": custom_pass}
+        )
+        compiled_fn(x, y)
+
+    def test_metadata_propagation_graph_pattern_replace_by_example(self):
+        """Verify metadata propagation for replace_by_example (single-node match)."""
+
+        test_pass = PatternMatcherPass()
+
+        @register_graph_pattern(
+            CallFunction(aten.add.Tensor, KeywordArg("x"), KeywordArg("y")),
+            pass_dict=test_pass,
+        )
+        def add_to_mul(match: Match, x, y):
+            def repl(a, b):
+                return a * b
+
+            with V.fake_mode:
+                match.replace_by_example(repl, [x, y])
+
+        def custom_pass(graph: torch.fx.Graph):
+            self._inject_test_metadata(graph)
+            test_pass.apply(graph)
+
+        def fn(x, y):
+            return x + y
+
+        x = torch.randn(4, 4, device=GPU_TYPE)
+        y = torch.randn(4, 4, device=GPU_TYPE)
+
+        compiled_fn = torch.compile(
+            fn, options={"post_grad_custom_post_pass": custom_pass}
+        )
+        compiled_fn(x, y)
+
+    def test_metadata_propagation_replace_by_example_multinode(self):
+        """Verify metadata propagation for replace_by_example with a multi-node
+        match.  The output node should inherit metadata from the matched output
+        node via replace_with_graph's replace() inner function."""
+
+        test_pass = PatternMatcherPass()
+
+        @register_graph_pattern(
+            CallFunction(
+                aten.mul.Tensor,
+                CallFunction(aten.add.Tensor, KeywordArg("x"), KeywordArg("y")),
+                KeywordArg("z"),
+            ),
+            pass_dict=test_pass,
+        )
+        def add_mul_to_sub(match: Match, x, y, z):
+            def repl(a, b, c):
+                return (a - b) * c
+
+            with V.fake_mode:
+                match.replace_by_example(repl, [x, y, z])
+
+        def custom_pass(graph: torch.fx.Graph):
+            self._inject_test_metadata(graph)
+            test_pass.apply(graph)
+
+        def fn(x, y, z):
+            return (x + y) * z
+
+        x = torch.randn(4, 4, device=GPU_TYPE)
+        y = torch.randn(4, 4, device=GPU_TYPE)
+        z = torch.randn(4, 4, device=GPU_TYPE)
+
+        compiled_fn = torch.compile(
+            fn, options={"post_grad_custom_post_pass": custom_pass}
+        )
+        compiled_fn(x, y, z)
+
+    def test_metadata_propagation_lowering_pattern(self):
+        """Verify metadata propagation for LoweringPatternEntry.apply.
+
+        LoweringPatternEntry uses _transfer_meta to copy _COPY_META_FIELDS
+        and stack_trace from the matched node to the replacement.
+        """
+        from torch._inductor.pattern_matcher import _transfer_meta
+
+        test_pass = PatternMatcherPass()
+
+        counter = 0
+
+        @register_graph_pattern(
+            CallFunction(aten.add.Tensor, KeywordArg("x"), KeywordArg("y")),
+            pass_dict=test_pass,
+        )
+        def manual_lowering(match: Match, x, y):
+            nonlocal counter
+            # Manually exercise the LoweringPatternEntry code path:
+            # create a replacement node, propagate meta, and replace
+            node = match.output_node()
+            graph = match.graph
+            with graph.inserting_before(node):
+                replacement = graph.call_function(aten.mul.Tensor, (x, y))
+                _transfer_meta(replacement.meta, node)
+                node.replace_all_uses_with(replacement)
+            match.erase_nodes()
+            counter += 1
+
+        def custom_pass(graph: torch.fx.Graph):
+            self._inject_test_metadata(graph)
+            test_pass.apply(graph)
+
+        def fn(x, y):
+            return x + y
+
+        x = torch.randn(4, 4, device=GPU_TYPE)
+        y = torch.randn(4, 4, device=GPU_TYPE)
+
+        compiled_fn = torch.compile(
+            fn, options={"post_grad_custom_post_pass": custom_pass}
+        )
+        compiled_fn(x, y)
+        self.assertEqual(counter, 1)
+
 
 class TestPatternMatcherLogging(LoggingTestCase):
     device_type = GPU_TYPE
@@ -2080,6 +2463,98 @@ def fn(x):
             specific_record.getMessage(),
         )
 
+    @make_logging_test()
+    def test_pattern_match_debug_multiple_nodes(self, records):
+        def pattern_add(x, y):
+            return x + y
+
+        def replacement_add(x, y):
+            return x * y
+
+        def pattern_sub(x, y):
+            return x - y
+
+        def replacement_sub(x, y):
+            return x * y
+
+        my_patterns = PatternMatcherPass()
+        inputs = [
+            torch.randn(4, 4, device=GPU_TYPE),
+            torch.randn(4, 4, device=GPU_TYPE),
+        ]
+        register_replacement(
+            pattern_add, replacement_add, inputs, fwd_only, my_patterns
+        )
+        register_replacement(
+            pattern_sub, replacement_sub, inputs, fwd_only, my_patterns
+        )
+
+        def custom_pass(graph: torch.fx.Graph):
+            return my_patterns.apply(graph)
+
+        def fn(x, y):
+            return (x + y) + (x - y)
+
+        x = torch.randn(4, 4, device=GPU_TYPE)
+        y = torch.randn(4, 4, device=GPU_TYPE)
+
+        # Debug both "add" and "sub" nodes
+        with unittest.mock.patch.dict(
+            os.environ, {"TORCHINDUCTOR_PATTERN_MATCH_DEBUG": "add,sub"}
+        ):
+            compiled_fn = torch.compile(
+                fn, options={"post_grad_custom_post_pass": custom_pass}
+            )
+            _ = compiled_fn(x, y)
+
+        self.assertTrue(self.hasRecord(records, "Specific pattern match: add"))
+        self.assertTrue(self.hasRecord(records, "Specific pattern match: sub"))
+
+    @make_logging_test()
+    def test_pattern_match_debug_all_nodes(self, records):
+        def pattern_add(x, y):
+            return x + y
+
+        def replacement_add(x, y):
+            return x * y
+
+        def pattern_sub(x, y):
+            return x - y
+
+        def replacement_sub(x, y):
+            return x * y
+
+        my_patterns = PatternMatcherPass()
+        inputs = [
+            torch.randn(4, 4, device=GPU_TYPE),
+            torch.randn(4, 4, device=GPU_TYPE),
+        ]
+        register_replacement(
+            pattern_add, replacement_add, inputs, fwd_only, my_patterns
+        )
+        register_replacement(
+            pattern_sub, replacement_sub, inputs, fwd_only, my_patterns
+        )
+
+        def custom_pass(graph: torch.fx.Graph):
+            return my_patterns.apply(graph)
+
+        def fn(x, y):
+            return (x + y) + (x - y)
+
+        x = torch.randn(4, 4, device=GPU_TYPE)
+        y = torch.randn(4, 4, device=GPU_TYPE)
+
+        with unittest.mock.patch.dict(
+            os.environ, {"TORCHINDUCTOR_PATTERN_MATCH_DEBUG": "all"}
+        ):
+            compiled_fn = torch.compile(
+                fn, options={"post_grad_custom_post_pass": custom_pass}
+            )
+            _ = compiled_fn(x, y)
+        self.assertTrue(self.hasRecord(records, "Specific pattern match: add"))
+        self.assertTrue(self.hasRecord(records, "Specific pattern match: sub"))
+
     def test_gumbel_max_trick(self):
         counters.clear()
 
@@ -2117,6 +2592,135 @@ def sample(logits, temperature):
 
         self.assertTrue(counters["inductor"]["apply_gumbel_max_trick"] == 1)
 
+    def test_per_pattern_counter(self):
+        """Test that per-pattern counters track individual pattern matches"""
+        with inductor_config.patch(fx_graph_cache=False):
+            with unittest.mock.patch.dict(
+                os.environ, {"TORCHINDUCTOR_PATTERN_MATCH_DEBUG": "1"}
+            ):
+                counters.clear()
+
+                def fn(x, y):
+                    return torch.bmm(x, y)
+
+                x = torch.randn(4, 10, 10, device=GPU_TYPE)
+                y = torch.randn(4, 10, 10, device=GPU_TYPE)
+
+                compiled = torch.compile(fn)
+                compiled(x, y)
+
+                counter_key = "inductor_pattern_matcher_per_pattern"
+                per_pattern = counters.get(counter_key, None)
+
+                self.assertIsInstance(per_pattern, dict)
+                self.assertGreater(len(per_pattern), 0)
+                self.assertIn("CallFunction_aten.bmm.default", per_pattern)
+                self.assertEqual(per_pattern["CallFunction_aten.bmm.default"], 1)
+
+    def test_per_pattern_counter_accumulation(self):
+        """Test that per-pattern counters accumulate across compilations"""
+        with inductor_config.patch(fx_graph_cache=False):
+            with unittest.mock.patch.dict(
+                os.environ, {"TORCHINDUCTOR_PATTERN_MATCH_DEBUG": "1"}
+            ):
+                counter_key = "inductor_pattern_matcher_per_pattern"
+
+                counters.clear()
+
+                x = torch.randn(2, 10, 10, device=GPU_TYPE)
+                y = torch.randn(2, 10, 10, device=GPU_TYPE)
+
+                def fn1(a, b):
+                    return torch.bmm(a, b)
+
+                compiled1 = torch.compile(fn1)
+                compiled1(x, y)
+                count1 = sum(counters.get(counter_key, {}).values())
+
+                # Compile second function without clearing counters
+                def fn2(a, b):
+                    return torch.bmm(a, b) * 2
+
+                compiled2 = torch.compile(fn2)
+                compiled2(x, y)
+                accumulated_count = sum(counters.get(counter_key, {}).values())
+
+                # Verify accumulation
+                counters.clear()
+                torch._dynamo.reset()
+                compiled2 = torch.compile(fn2)
+                compiled2(x, y)
+                count2 = sum(counters.get(counter_key, {}).values())
+
+                self.assertEqual(accumulated_count, count1 + count2)
+
+    def test_opaque_obj_custom_op(self):
+        with torch.library._scoped_library("_test_pm", "FRAGMENT") as lib:
+            lib.define(
+                f"original_op(Tensor x, {get_opaque_type_name(OpaqueScaleFactor)} s) -> Tensor"
+            )
+            lib.impl("original_op", lambda x, s: x * s.val, "CompositeExplicitAutograd")
+
+            @torch.library.register_fake("_test_pm::original_op", lib=lib)
+            def _orig_fake(x, s):
+                return torch.empty_like(x)
+
+            lib.define(
+                f"replacement_op(Tensor x, {get_opaque_type_name(OpaqueScaleFactor)} s) -> Tensor"
+            )
+            lib.impl(
+                "replacement_op", lambda x, s: x + s.val, "CompositeExplicitAutograd"
+            )
+
+            @torch.library.register_fake("_test_pm::replacement_op", lib=lib)
+            def _repl_fake(x, s):
+                return torch.empty_like(x)
+
+            def pattern(x, factor):
+                return torch.ops._test_pm.original_op(x, factor)
+
+            def replacement(x, factor):
+                return torch.ops._test_pm.replacement_op(x, factor)
+
+            patterns = PatternMatcherPass()
+            register_replacement(
+                pattern,
+                replacement,
+                [torch.randn(4, 4), OpaqueScaleFactor(2.0)],
+                fwd_only,
+                patterns,
+            )
+
+            count = 0
+            post_pass_graph = None
+
+            def custom_pass(graph):
+                nonlocal count, post_pass_graph
+                count = patterns.apply(graph)
+                post_pass_graph = graph
+                return graph
+
+            def custom_backend(graph, example_inputs):
+                from torch._inductor.compile_fx import compile_fx
+
+                current_config = inductor_config.get_config_copy()
+                current_config["post_grad_custom_post_pass"] = custom_pass
+                return compile_fx(graph, example_inputs, config_patches=current_config)
+
+            @torch.compile(backend=custom_backend)
+            def f(x, s):
+                return torch.ops._test_pm.original_op(x, s)
+
+            inp = torch.randn(4, 4)
+            result = f(inp, OpaqueScaleFactor(4.0))
+            self.assertEqual(result, inp + 4.0)
+            self.assertEqual(count, 1)
+            op_targets = [
+                n.target for n in post_pass_graph.nodes if n.op == "call_function"
+            ]
+            self.assertNotIn(torch.ops._test_pm.original_op.default, op_targets)
+            self.assertIn(torch.ops._test_pm.replacement_op.default, op_targets)
+
 
 if __name__ == "__main__":
     if IS_LINUX and HAS_GPU:
diff --git a/test/inductor/test_perf.py b/test/inductor/test_perf.py
index cbb2a558b8865..4e168fd0b4c21 100644
--- a/test/inductor/test_perf.py
+++ b/test/inductor/test_perf.py
@@ -1,6 +1,7 @@
 # Owner(s): ["module: inductor"]
 import contextlib
 import re
+import unittest
 from unittest.mock import patch
 
 import functorch
@@ -28,6 +29,7 @@
 # performance for that setting.
 #
 # Defines all the kernels for tests
+from torch.testing._internal.common_utils import skipIfXpu, TEST_WITH_ROCM
 from torch.testing._internal.inductor_utils import GPU_TYPE, HAS_GPU_AND_TRITON
 from torch.testing._internal.triton_utils import requires_gpu_and_triton
 
@@ -337,6 +339,35 @@ def f(a, b):
         inp = (T(10), TI(10, mx=10))
         self.assertExpectedInline(count_numel(f, *inp), """30""")
 
+    @requires_gpu_and_triton
+    def test_delay_realize_cheap_outputs_shared_mask(self):
+        # Shared tril mask across multiple users gets eagerly materialized
+        # as an output buffer, inflating downstream read counts. With
+        # delay_realize_cheap_outputs, the mask stays inlined as index
+        # arithmetic, reducing memory traffic.
+        def f(x1, x2, x3):
+            mask = torch.tril(torch.ones(32, 32, device=x1.device))
+            return x1 + mask, x2 + mask, x3 + mask
+
+        inp = (
+            T(32, 32, grad=True),
+            T(32, 32, grad=True),
+            T(32, 32, grad=True),
+        )
+
+        # Without deferred realization: mask gets materialized as output buffer
+        with patch.object(config, "delay_realize_cheap_outputs", False):
+            metrics.reset()
+            torch.compile(f, backend=compile_but_use_eager)(*inp)
+            eager_bytes = metrics.num_bytes_accessed
+
+        # Default (deferred realization on): mask stays inlined
+        metrics.reset()
+        torch.compile(f, backend=compile_but_use_eager)(*inp)
+        deferred_bytes = metrics.num_bytes_accessed
+
+        self.assertLessEqual(deferred_bytes, eager_bytes)
+
 
 class FusionTests(TestCase):
     """
@@ -484,6 +515,26 @@ def f(a, b, c):
         inp = (T(10, 10), T(10, 10), T(10, 10))
         self.assertExpectedInline(count_numel(f, *inp), """500""")
 
+    @skipIfXpu(msg="copy_(cat()) fusion not supported on XPU")
+    @unittest.skipIf(TEST_WITH_ROCM, "copy_(cat()) fusion not supported on ROCm")
+    # TODO(ivankobzarev): enable copy_(cat()) fusion for CUDA 13+
+    @unittest.skipIf(
+        torch.version.cuda
+        and tuple(int(x) for x in torch.version.cuda.split(".")) >= (13, 0),
+        "copy_(cat()) fusion not supported on CUDA 13+",
+    )
+    def test_copy_cat_fusion(self):
+        """copy_(cat(...)) should fuse: no intermediate allocation for cat."""
+
+        def f(dst, a, b):
+            dst.copy_(torch.cat([a, b]))
+
+        dst = T(20)
+        inp = (dst, T(10), T(10))
+        # 10 (read a) + 10 (read b) + 20 (write dst) = 40
+        # Without fusion cat would allocate intermediate: 80
+        self.assertExpectedInline(count_numel(f, *inp), """40""")
+
     def test_reduction_pointwise_multi_level_reduction(self):
         hidden_size = 4096
         layer_norm = torch.nn.LayerNorm(hidden_size).to(GPU_TYPE).float()
@@ -505,7 +556,10 @@ def f(x, scale, amax_keep_dim):
         expected_numel = (
             1 + hidden_size * 2 + 4 * 2048 * hidden_size * 2 + 4 * 2048 * 2 + 1
         )
-        if config.triton.cooperative_reductions:
+        if (
+            config.triton.cooperative_reductions
+            or config.triton.force_cooperative_reductions
+        ):
             expected_numel = 134225922
 
         self.assertExpectedInline(count_numel(f, *inp, True), str(expected_numel))
diff --git a/test/inductor/test_profiler.py b/test/inductor/test_profiler.py
index 07db20e843780..2a90c55285f69 100644
--- a/test/inductor/test_profiler.py
+++ b/test/inductor/test_profiler.py
@@ -4,7 +4,6 @@
 import tempfile
 import unittest
 from collections.abc import Callable
-from typing import Optional
 
 import torch
 import torch._inductor.test_case
@@ -12,7 +11,7 @@
 from torch import _dynamo as torchdynamo
 from torch._inductor import config
 from torch.profiler import ProfilerActivity, record_function
-from torch.testing._internal.common_utils import skipIfXpu, TemporaryFileName
+from torch.testing._internal.common_utils import TemporaryFileName
 from torch.testing._internal.inductor_utils import (
     GPU_TYPE,
     HAS_GPU_AND_TRITON,
@@ -27,10 +26,6 @@
 
 
 class DynamoProfilerTests(torch._inductor.test_case.TestCase):
-    @skipIfXpu(
-        msg="AssertionError: False is not true, "
-        "https://github.com/intel/torch-xpu-ops/issues/2335"
-    )
     @unittest.skipIf(not HAS_TRITON, "requires cuda & triton")
     def test_inductor_profiling_triton_launch(self):
         # Verify that we get some sort of CPU-side indication of triton kernel launches
@@ -61,7 +56,7 @@ def fn(x, y):
         self.assertTrue(any((event.get("name") in valid_names) for event in events))
 
     def _test_profiling_kernel_names(
-        self, fn, args, kernel_name_str: str, check_fn: Optional[Callable] = None
+        self, fn, args, kernel_name_str: str, check_fn: Callable | None = None
     ):
         """
         We expect a record_function event to be added on the CPU side, surrounding
@@ -225,9 +220,6 @@ def fn(x, y):
         self.assertTrue(hooks_called["enter"])
         self.assertTrue(hooks_called["exit"])
 
-    @skipIfXpu(
-        msg="TypeError: list indices must be integers or slices, not str, https://github.com/intel/torch-xpu-ops/issues/2335"
-    )
     @unittest.skipIf(not HAS_TRITON, "requires cuda & triton")
     def test_pt2_triton_attributes(self):
         from torch._inductor.codecache import code_hash
diff --git a/test/inductor/test_provenance_tracing.py b/test/inductor/test_provenance_tracing.py
index 93397b6eae072..a5e3e8dfc4ad0 100644
--- a/test/inductor/test_provenance_tracing.py
+++ b/test/inductor/test_provenance_tracing.py
@@ -29,10 +29,7 @@
 from torch._inductor.virtualized import V
 from torch.testing._internal.common_utils import IS_MACOS
 from torch.testing._internal.inductor_utils import GPU_TYPE
-from torch.testing._internal.triton_utils import (
-    requires_cuda_and_triton,
-    requires_gpu_and_triton,
-)
+from torch.testing._internal.triton_utils import requires_gpu_and_triton
 
 
 try:
@@ -230,7 +227,8 @@ def _test_triton_kernel_to_post_grad_tracing(self, device):
                             filepath, expected_mapping
                         )
                     else:
-                        assert device == "cpu"
+                        if device != "cpu":
+                            raise AssertionError
                         # check the inductor kernel to post grad nodes mapping is expected for cpu
                         if backend == "aot_inductor":
                             expected_data = {
@@ -612,7 +610,7 @@ def test_tlparse_kernel_stack_traces(self):
     @torch._inductor.config.patch(
         {"trace.provenance_tracking_level": 2, "max_autotune_gemm_backends": "ATEN"}
     )
-    @requires_cuda_and_triton
+    @requires_gpu_and_triton
     def test_deferred_triton_kernels(self):
         def foo(m, inp):
             a = m(inp)
@@ -620,8 +618,8 @@ def foo(m, inp):
 
         foo_c = torch.compile(mode="max-autotune-no-cudagraphs")(foo)
 
-        m = torch.nn.Linear(512, 512, bias=True).half().cuda()
-        inp = torch.rand([1, 512]).half().cuda()
+        m = torch.nn.Linear(512, 512, bias=True).half().to(GPU_TYPE)
+        inp = torch.rand([1, 512]).half().to(GPU_TYPE)
 
         with self._setup_provenance_capture() as payload_buffer:
             with torch.no_grad():
@@ -901,16 +899,21 @@ def forward(self, x, a, b, c):
                 code
             )
 
-            if self.device == "cuda":
+            if self.device == "cuda" or self.device == "xpu":
+                device_type = torch.accelerator.current_accelerator().type
                 FileCheck().check(
-                    """KernelContextGuard _ctx("aoti_torch_cuda_mm_out", R"("""
-                ).check("AOTI_TORCH_ERROR_CODE_CHECK(aoti_torch_cuda_mm_out(").check(
+                    f"""KernelContextGuard _ctx("aoti_torch_{device_type}_mm_out", R"("""
+                ).check(
+                    f"AOTI_TORCH_ERROR_CODE_CHECK(aoti_torch_{device_type}_mm_out("
+                ).check(
                     """KernelContextGuard _ctx("triton_poi_fused_addmm_relu_sigmoid_0", R"("""
                 ).check("call_triton_poi_fused_addmm_relu_sigmoid_0(").check(
                     """KernelContextGuard _ctx("triton_poi_fused_mul_1", R"("""
                 ).check("call_triton_poi_fused_mul_1(").check(
-                    """KernelContextGuard _ctx("aoti_torch_cuda_mm_out", R"("""
-                ).check("AOTI_TORCH_ERROR_CODE_CHECK(aoti_torch_cuda_mm_out(").check(
+                    f"""KernelContextGuard _ctx("aoti_torch_{device_type}_mm_out", R"""
+                ).check(
+                    f"AOTI_TORCH_ERROR_CODE_CHECK(aoti_torch_{device_type}_mm_out("
+                ).check(
                     """ KernelContextGuard _ctx("triton_poi_fused_addmm_gelu_2", R"("""
                 ).check("call_triton_poi_fused_addmm_gelu_2(").run(code)
             else:
@@ -941,15 +944,17 @@ class TestProvenanceTracingKernelContextCpu(TestCase):
 
 
 @unittest.skipIf(sys.platform == "darwin", "No CUDA on MacOS")
-@unittest.skipIf(not torch.cuda.is_available(), "No CUDA")
+@unittest.skipIf(
+    not torch.cuda.is_available() and not torch.xpu.is_available(), "No CUDA and no XPU"
+)
 class TestProvenanceTracingKernelContextGpu(TestCase):
-    device = "cuda"
+    device = GPU_TYPE
 
 
 copy_tests(
     ProvenanceTracingKernelContextTemplate,
     TestProvenanceTracingKernelContextGpu,
-    "cuda",
+    GPU_TYPE,
 )
 
 
diff --git a/test/inductor/test_quantization.py b/test/inductor/test_quantization.py
index 0f137703d4f82..51f7f4f205092 100644
--- a/test/inductor/test_quantization.py
+++ b/test/inductor/test_quantization.py
@@ -47,6 +47,29 @@ def forward(self, x):
         return x
 
 
+class SharedOutputAndSavedModule(torch.nn.Module):
+    """A module where a 3D intermediate is both a user output and saved for backward.
+
+    This triggers the bug in T264303372: the activation quantization pass would
+    quantize both the user output and saved-for-backward positions of the same
+    tensor, creating duplicate backward placeholders that shift the stride mapping.
+    """
+
+    def __init__(self):
+        super().__init__()
+        self.scale = torch.nn.Parameter(torch.ones(8, 8))
+        self.W = torch.nn.Parameter(torch.randn(64, 32))
+
+    def forward(self, x):
+        h = x.view(x.shape[0], 8, 8)
+        attn = h * torch.sigmoid(h)
+        # attn is both returned (user output) and needed by backward (mul saves it)
+        scaled = attn * self.scale
+        flat = scaled.flatten(1)
+        result = flat @ self.W
+        return result, attn
+
+
 class LayernormNN(torch.nn.Module):
     def __init__(self):
         super().__init__()
@@ -68,7 +91,8 @@ def compare_dict_tensors(self, ref_dict, res_dict, rtol=1e-3, atol=1e-3):
             return False
         for key1 in ref_dict:
             key2 = "_orig_mod." + key1
-            assert key2 in res_dict, f"{key1} does not exist in traced module"
+            if key2 not in res_dict:
+                raise AssertionError(f"{key1} does not exist in traced module")
             # if both of them are None, continue
             if (
                 not isinstance(ref_dict[key1], torch.Tensor)
@@ -228,6 +252,37 @@ def test_activation_quantization_aten_without_scaling(self):
         self.assertTrue(torch.allclose(ref, res))
         counters.clear()
 
+    @requires_gpu()
+    @torch._inductor.config.patch(
+        pre_grad_fusion_options={},
+        post_grad_fusion_options={
+            "activation_quantization_aten_pass": {
+                "quant_type": "torch.float8_e5m2",
+                "use_scaling": True,
+                "size_in_mb": 0.0,
+                "exclude_primals": True,
+                "allowed_dtypes": "torch.bfloat16;torch.float32",
+            },
+        },
+    )
+    def test_activation_quantization_shared_output_and_saved(self):
+        """Test that activation quantization works when a tensor is both a user
+        output and a saved-for-backward activation (T264303372)."""
+        counters.clear()
+        module = SharedOutputAndSavedModule().to(GPU_TYPE)
+        x = torch.randn(8, 64, device=GPU_TYPE, requires_grad=True)
+        compiled = torch.compile(module)
+        result, attn = compiled(x)
+        loss = result.sum() + attn.sum()
+        loss.backward()
+        self.assertEqual(
+            counters["inductor"]["activation_quantization_fwd_aten_pass"], 1
+        )
+        self.assertEqual(
+            counters["inductor"]["activation_quantization_bwd_aten_pass"], 1
+        )
+        counters.clear()
+
 
 if __name__ == "__main__":
     if IS_LINUX and HAS_GPU:
diff --git a/test/inductor/test_scatter_optimization.py b/test/inductor/test_scatter_optimization.py
index 3e8561020fe68..a68565602e127 100644
--- a/test/inductor/test_scatter_optimization.py
+++ b/test/inductor/test_scatter_optimization.py
@@ -180,9 +180,8 @@ def f(m, x, label):
         ref_grad = ref_model.weight.grad
         opt_f(opt_model, x, label)
         act_grad = opt_model.weight.grad
-        assert torch.allclose(ref_grad, act_grad, atol=1e-3, rtol=1e-3), (
-            f"{ref_grad=}\n{act_grad=}"
-        )
+        if not torch.allclose(ref_grad, act_grad, atol=1e-3, rtol=1e-3):
+            raise AssertionError(f"{ref_grad=}\n{act_grad=}")
 
         self.check_metric()
 
diff --git a/test/inductor/test_segmented_tree.py b/test/inductor/test_segmented_tree.py
index 22d1c85027dc9..12503a774f9c4 100644
--- a/test/inductor/test_segmented_tree.py
+++ b/test/inductor/test_segmented_tree.py
@@ -46,7 +46,8 @@ class TestSegmentedTree(TestCase):
     def test_basic_construction(self):
         values = [1, 3, 5, 7, 9]
         tree = SegmentedTree(values, add_op, max_op, 0)
-        assert tree.summarize_range(0, 4) == 9
+        if tree.summarize_range(0, 4) != 9:
+            raise AssertionError(f"Expected 9, got {tree.summarize_range(0, 4)}")
 
     def test_empty_array(self):
         with self.assertRaises(ValueError):
@@ -61,9 +62,10 @@ def test_max_query_matches_naive(self, values):
             for end in range(start, len(values)):
                 expected = naive_range_max(values, start, end)
                 actual = tree.summarize_range(start, end)
-                assert actual == expected, (
-                    f"Range [{start}:{end}] expected {expected}, got {actual}"
-                )
+                if actual != expected:
+                    raise AssertionError(
+                        f"Range [{start}:{end}] expected {expected}, got {actual}"
+                    )
 
     @given(
         values=positive_integers, range_indices=st.data(), update_value=update_values
@@ -87,9 +89,10 @@ def test_range_update(self, values, range_indices, update_value):
             for j in range(i, len(values)):
                 expected = naive_range_max(naive_values, i, j)
                 actual = tree.summarize_range(i, j)
-                assert actual == expected, (
-                    f"After update, range [{i}:{j}] expected {expected}, got {actual}"
-                )
+                if actual != expected:
+                    raise AssertionError(
+                        f"After update, range [{i}:{j}] expected {expected}, got {actual}"
+                    )
 
     @given(values=positive_integers, range_data=st.data())
     def test_multiple_operations(self, values, range_data):
@@ -107,9 +110,10 @@ def test_multiple_operations(self, values, range_data):
             if operation_type == "query":
                 expected = naive_range_max(naive_values, start, end)
                 actual = tree.summarize_range(start, end)
-                assert actual == expected, (
-                    f"Range query [{start}:{end}] expected {expected}, got {actual}"
-                )
+                if actual != expected:
+                    raise AssertionError(
+                        f"Range query [{start}:{end}] expected {expected}, got {actual}"
+                    )
             else:  # update
                 update_value = range_data.draw(update_values)
                 tree.update_range(start, end, update_value)
@@ -120,40 +124,55 @@ def test_single_element_ranges(self):
         tree = SegmentedTree(values, add_op, max_op, 0)
 
         for i in range(len(values)):
-            assert tree.summarize_range(i, i) == values[i], (
-                f"Single element range at index {i} failed"
-            )
+            if tree.summarize_range(i, i) != values[i]:
+                raise AssertionError(f"Single element range at index {i} failed")
 
     def test_full_array_range(self):
         values = [1, 3, 5, 7, 9]
         tree = SegmentedTree(values, add_op, max_op, 0)
 
         # Test querying the entire array
-        assert tree.summarize_range(0, len(values) - 1) == max(values)
+        if tree.summarize_range(0, len(values) - 1) != max(values):
+            raise AssertionError(
+                f"Expected {max(values)}, got {tree.summarize_range(0, len(values) - 1)}"
+            )
 
         # Update the entire array and test again
         update_value = 10
         tree.update_range(0, len(values) - 1, update_value)
         expected = max([v + update_value for v in values])
-        assert tree.summarize_range(0, len(values) - 1) == expected
+        if tree.summarize_range(0, len(values) - 1) != expected:
+            raise AssertionError(
+                f"Expected {expected}, got {tree.summarize_range(0, len(values) - 1)}"
+            )
 
     def test_boundary_conditions(self):
         values = [1, 3, 5, 7, 9]
         tree = SegmentedTree(values, add_op, max_op, 0)
 
         # Test first element
-        assert tree.summarize_range(0, 0) == values[0]
+        if tree.summarize_range(0, 0) != values[0]:
+            raise AssertionError(
+                f"Expected {values[0]}, got {tree.summarize_range(0, 0)}"
+            )
 
         # Test last element
-        assert tree.summarize_range(len(values) - 1, len(values) - 1) == values[-1]
+        if tree.summarize_range(len(values) - 1, len(values) - 1) != values[-1]:
+            raise AssertionError(
+                f"Expected {values[-1]}, got {tree.summarize_range(len(values) - 1, len(values) - 1)}"
+            )
 
         # Test first two elements
-        assert tree.summarize_range(0, 1) == max(values[0:2])
+        if tree.summarize_range(0, 1) != max(values[0:2]):
+            raise AssertionError(
+                f"Expected {max(values[0:2])}, got {tree.summarize_range(0, 1)}"
+            )
 
         # Test last two elements
-        assert tree.summarize_range(len(values) - 2, len(values) - 1) == max(
-            values[-2:]
-        )
+        if tree.summarize_range(len(values) - 2, len(values) - 1) != max(values[-2:]):
+            raise AssertionError(
+                f"Expected {max(values[-2:])}, got {tree.summarize_range(len(values) - 2, len(values) - 1)}"
+            )
 
     def test_invalid_ranges(self):
         values = [1, 3, 5, 7, 9]
@@ -208,9 +227,10 @@ def test_overlapping_updates(self):
             for j in range(i, len(values)):
                 expected = naive_range_max(naive_values, i, j)
                 actual = tree.summarize_range(i, j)
-                assert actual == expected, (
-                    f"After overlapping updates, range [{i}:{j}] expected {expected}, got {actual}"
-                )
+                if actual != expected:
+                    raise AssertionError(
+                        f"After overlapping updates, range [{i}:{j}] expected {expected}, got {actual}"
+                    )
 
     def test_sequential_updates_and_queries(self):
         values = [2, 4, 6, 8, 10, 12, 14]
@@ -239,15 +259,17 @@ def test_sequential_updates_and_queries(self):
                     for j in range(i, len(values)):
                         expected = naive_range_max(naive_values, i, j)
                         actual = tree.summarize_range(i, j)
-                        assert actual == expected, (
-                            f"After update ({start}, {end}, {value}), query [{i}:{j}] expected {expected}, got {actual}"
-                        )
+                        if actual != expected:
+                            raise AssertionError(
+                                f"After update ({start}, {end}, {value}), query [{i}:{j}] expected {expected}, got {actual}"
+                            )
             else:  # query
                 _, start, end = op
                 expected = naive_range_max(naive_values, start, end)
-                assert tree.summarize_range(start, end) == expected, (
-                    f"Query [{start}:{end}] expected {expected}, got {tree.summarize_range(start, end)}"
-                )
+                if tree.summarize_range(start, end) != expected:
+                    raise AssertionError(
+                        f"Query [{start}:{end}] expected {expected}, got {tree.summarize_range(start, end)}"
+                    )
 
 
 if __name__ == "__main__":
diff --git a/test/inductor/test_select_algorithm.py b/test/inductor/test_select_algorithm.py
index 97cae123cfbb9..fd37f392bce69 100644
--- a/test/inductor/test_select_algorithm.py
+++ b/test/inductor/test_select_algorithm.py
@@ -3,11 +3,9 @@
 import functools
 import unittest.mock
 from collections.abc import Callable
-from typing import Any, Optional, Union
+from typing import Any
 from unittest.mock import patch
 
-import pytest
-
 import torch
 import torch._dynamo.config as dynamo_config
 import torch._inductor.config as inductor_config
@@ -27,7 +25,9 @@
 from torch._inductor.kernel_inputs import KernelInputs
 from torch._inductor.select_algorithm import (
     autotune_select_algorithm,
+    ExternalTritonTemplateKernel,
     ExternKernelChoice,
+    PartialRender,
     TritonTemplate,
     TritonTemplateKernel,
 )
@@ -72,9 +72,8 @@ def skip_cache(self, choices, name, key, benchmark, hint_override=None):
     def wrapped(*args, **kwargs):
         counters.clear()
         torch.manual_seed(12345)
-        assert torch.backends.cuda.matmul.fp32_precision != "tf32", (
-            "correctness testing is allergic to tf32"
-        )
+        if torch.backends.cuda.matmul.fp32_precision == "tf32":
+            raise AssertionError("correctness testing is allergic to tf32")
         return fn(*args, **kwargs)
 
     return wrapped
@@ -231,6 +230,23 @@ def foo(a, b):
         if not torch.version.hip:  # autotuning is not guaranteed to run on ROCm
             self.assertEqual(counters["inductor"]["select_algorithm_autotune"], 1)
 
+    @patches
+    def test_bmm_small_m(self):
+        # Verify BMM works when M < BLOCK_M. The triton_bmm template's
+        # tl.max_contiguous/tl.multiple_of hints must be guarded by
+        # M >= BLOCK_M to avoid out-of-bounds vectorized loads (see #179267).
+        @torch.compile
+        def foo(a, b):
+            return torch.bmm(a, b)
+
+        # M=2 is smaller than any BLOCK_M autotuning config (typically >= 16)
+        foo(
+            torch.randn(4, 2, 64, device=GPU_TYPE),
+            torch.randn(4, 64, 32, device=GPU_TYPE),
+        )
+        if not torch.version.hip:
+            self.assertEqual(counters["inductor"]["select_algorithm_autotune"], 1)
+
     @patches
     def test_mm_not_even_k(self):
         @torch.compile
@@ -293,10 +309,6 @@ def foo(a, b, c, d):
         if not torch.version.hip:  # autotuning is not guaranteed to run on ROCm
             self.assertEqual(counters["inductor"]["select_algorithm_autotune"], 1)
 
-    @pytest.mark.xfail(
-        condition=not torch.version.hip,
-        reason="C++ wrapper dynamic shapes fails on CUDA, fixed on ROCm",
-    )
     @patches
     def test_mm_plus_mm3(self):
         @torch.compile
@@ -481,9 +493,9 @@ class SingleMMConfigChoice(InductorChoices):
             def get_template_configs(
                 self,
                 kernel_inputs: KernelInputs,
-                templates: list[Union[KernelTemplate, ExternKernelChoice]],
+                templates: list[KernelTemplate | ExternKernelChoice],
                 op_name: str,
-                kwarg_overrides: Optional[dict[str, dict[str, Any]]] = None,
+                kwarg_overrides: dict[str, dict[str, Any]] | None = None,
             ):
                 return super().get_template_configs(
                     kernel_inputs, templates, op_name, kwarg_overrides
@@ -526,6 +538,53 @@ def foo(x, w, b):
         if not torch.version.hip:  # autotuning is not guaranteed to run on ROCm
             self.assertEqual(counters["inductor"]["select_algorithm_autotune"], 1)
 
+    @patches
+    def test_convolution_depthwise_conv1d(self):
+        @torch.compile
+        def foo(x, w, b):
+            return aten.convolution(
+                x,
+                w,
+                b,
+                stride=(1,),
+                padding=(4,),
+                dilation=(1,),
+                transposed=False,
+                output_padding=(0,),
+                groups=128,
+            )
+
+        foo(
+            torch.randn(2, 128, 202, device=GPU_TYPE),
+            torch.randn(128, 1, 9, device=GPU_TYPE),
+            torch.randn(128, device=GPU_TYPE),
+        )
+        if not torch.version.hip:
+            self.assertEqual(counters["inductor"]["select_algorithm_autotune"], 1)
+
+    @patches
+    def test_convolution_depthwise_conv1d_no_bias(self):
+        @torch.compile
+        def foo(x, w):
+            return aten.convolution(
+                x,
+                w,
+                None,
+                stride=(2,),
+                padding=(1,),
+                dilation=(1,),
+                transposed=False,
+                output_padding=(0,),
+                groups=64,
+            )
+
+        foo(
+            torch.randn(4, 64, 100, device=GPU_TYPE),
+            torch.randn(64, 1, 3, device=GPU_TYPE),
+        )
+        if not torch.version.hip:
+            self.assertEqual(counters["inductor"]["select_algorithm_autotune"], 1)
+
     def test_TritonTemplateCaller_str(self):
         """
         Make sure str(TritonTemplateCaller) does not raise exceptions.
@@ -550,6 +609,66 @@ def test_TritonTemplateCaller_str(self):
         self.assertEqual(caller_str, f"TritonTemplateCaller({module_path}, extra)")
 
 
+class TestSelectAlgorithmCleanup(TestCase):
+    def test_benchmark_only_clears_matching_precompile_cache_entry(self):
+        """
+        Autotune cleanup should release only the closure for the active site.
+        Clearing the whole precompile cache regresses compile-time reuse for
+        unrelated autotune sites.
+        """
+        cache = select_algorithm.AlgorithmSelectorCache()
+        cache.precompile_cache = {"keep": dict, "drop": dict}
+
+        with patch.object(cache, "make_benchmark_fn", return_value=lambda _choices: {}):
+            cache.benchmark([], [], unittest.mock.Mock(), None, precompile_key="drop")
+
+        self.assertIn("keep", cache.precompile_cache)
+        self.assertNotIn("drop", cache.precompile_cache)
+
+    def test_release_benchmark_artifacts_closes_and_clears_state(self):
+        """
+        Benchmark-only autotuners own temporary launchers and compile results.
+        Cleanup must close their module owners and clear the references that
+        otherwise keep Triton ``CompiledKernel`` objects alive.
+        """
+        from torch._inductor.runtime.triton_heuristics import CachingAutotuner
+
+        closed = []
+
+        class FakeKernel:
+            def __init__(self, name):
+                self.name = name
+
+            def close(self):
+                closed.append(self.name)
+
+        class FakeCompileResult:
+            def __init__(self, name):
+                self.kernel = FakeKernel(name)
+
+        def fake_launcher():
+            return None
+
+        launcher_kernel = FakeKernel("launcher")
+        fake_launcher.__self__ = launcher_kernel  # type: ignore[attr-defined]
+
+        autotuner = object.__new__(CachingAutotuner)
+        autotuner.launchers = [fake_launcher]
+        autotuner.compile_results = [FakeCompileResult("compile_result")]
+        autotuner.benchmark_failure_reasons = {fake_launcher: "failed"}
+        autotuner._cached_launcher = fake_launcher
+        autotuner._debug_call = object()
+
+        autotuner.release_benchmark_artifacts()
+
+        self.assertEqual(closed, ["launcher", "compile_result"])
+        self.assertEqual(autotuner.launchers, [])
+        self.assertEqual(autotuner.compile_results, [])
+        self.assertEqual(autotuner.benchmark_failure_reasons, {})
+        self.assertIsNone(autotuner._cached_launcher)
+        self.assertIsNone(autotuner._debug_call)
+
+
 class TestExternKernelCaller(TestCase):
     @requires_gpu()
     @patches
@@ -761,6 +880,73 @@ def patch_lowering(lowering_overrides) -> Callable[[], None]:
         yield
 
 
+class TestDtypeViewAutotuning(TestCase):
+    @requires_gpu()
+    @unittest.skipIf(
+        not hasattr(torch, "float4_e2m1fn_x2"),
+        "float4_e2m1fn_x2 dtype not available",
+    )
+    @patches
+    def test_benchmark_example_value_preserves_dtype_view(self):
+        """
+        Verify that benchmark_example_value preserves the dtype when the IR
+        node is a dtype view (e.g. uint8 storage viewed as float4_e2m1fn_x2).
+        """
+        from torch._inductor import ir
+
+        m, k = 256, 2048
+        device = torch.device(GPU_TYPE)
+
+        # We need a V.graph context for benchmark_example_value.
+        # Compile a trivial function to set up V.graph, then call the
+        # function under test within the compilation callback.
+        captured_results = {}
+
+        orig_lowering = torch._inductor.lowering.lowerings[aten.add.Tensor]
+
+        def patched_lowering(*args, **kwargs):
+            # We're inside compilation — V.graph is valid.
+            # Construct a dtype-viewed IR node and test benchmark_example_value.
+            base_layout = ir.FixedLayout(
+                device=device, dtype=torch.uint8, size=[m, k], stride=[k, 1]
+            )
+            base_buf = ir.Buffer(name="test_buf", layout=base_layout)
+
+            fp4_layout = ir.FixedLayout(
+                device=device,
+                dtype=torch.float4_e2m1fn_x2,
+                size=[m, k],
+                stride=[k, 1],
+            )
+            fp4_view = ir.ReinterpretView(data=base_buf, layout=fp4_layout)
+
+            example = select_algorithm.AlgorithmSelectorCache.benchmark_example_value(
+                fp4_view
+            )
+            captured_results["ir_dtype"] = fp4_view.get_dtype()
+            captured_results["example_dtype"] = example.dtype
+            captured_results["example_shape"] = tuple(example.shape)
+
+            return orig_lowering(*args, **kwargs)
+
+        torch._dynamo.reset()
+        with patch.dict(
+            torch._inductor.lowering.lowerings,
+            {aten.add.Tensor: patched_lowering},
+        ):
+            compiled_fn = torch.compile(lambda x: x + 1)
+            compiled_fn(torch.randn(4, device=device))
+
+        self.assertIn("ir_dtype", captured_results, "Patched lowering was not called")
+        self.assertEqual(
+            captured_results["example_dtype"],
+            torch.float4_e2m1fn_x2,
+            f"benchmark_example_value should preserve float4_e2m1fn_x2 dtype "
+            f"after unwrapping the view, but got {captured_results['example_dtype']}",
+        )
+        self.assertEqual(captured_results["example_shape"], (m, k))
+
+
 class TestTemplateRender(TestCase):
     @requires_gpu()
     @requires_triton()
@@ -829,7 +1015,8 @@ def add_override(a, b, alpha=None):
                 XBLOCK=XBLOCK,
                 triton_meta=custom_triton_meta,
             )
-            return autotune_select_algorithm("add", choices, [a, b], layout)
+            node, _ = autotune_select_algorithm("add", choices, [a, b], layout)
+            return node
 
         with patch_lowering(
             {
@@ -850,12 +1037,179 @@ def add(a, b):
             b = torch.zeros((XBLOCK,), device=GPU_TYPE)
 
             _result, kernels = run_and_get_kernels(add, a, b)
-            assert len(kernels) == 1
-            assert hook_identifier in kernels[0]
+            if len(kernels) != 1:
+                raise AssertionError
+            if hook_identifier not in kernels[0]:
+                raise AssertionError
             FileCheck().check("triton_meta=").check(str(custom_triton_meta)).run(
                 kernels[0]
             )
 
+    @requires_gpu()
+    @requires_triton()
+    @config.patch(cuda_backend="triton")
+    def test_external_template_prologue_epilogue_fusion(self):
+        """
+        Tests prologue fusion, epilogue fusion, and extra inputs through the
+        ExternalTritonTemplateKernel render()-based path.
+
+        Compiled function: relu(template_add(a, sigmoid(b))) * bias
+          - Prologue: sigmoid(b) fused into template as <LOAD_INPUT_B>
+          - Epilogue: relu(...) * bias fused into template as <STORE_OUTPUT_0>
+          - Extra inputs: bias is read by the epilogue but is not among the
+            template's original inputs, exercising kernel._extra_inputs
+        """
+        import torch._inductor.ir as ir
+        from torch._inductor.ir import OrderedSet
+        from torch._inductor.utils import Placeholder, run_and_get_code
+
+        XBLOCK = 128
+        render_called = [False]
+
+        # Template source with placeholders filled in by _render()
+        _MOCK_ADD_KERNEL_TEMPLATE = (
+            "import triton\n"
+            "import triton.language as tl\n"
+            "import torch\n"
+            "from torch._inductor.runtime import triton_helpers\n"
+            "\n"
+            "@triton.jit\n"
+            "def _mock_inner_add(A, B, {out_param}{extra_sig},"
+            " numel, XBLOCK: tl.constexpr):\n"
+            "    xoffset = tl.program_id(0) * XBLOCK\n"
+            "    xindex = xoffset + tl.arange(0, XBLOCK)\n"
+            "    xmask = xindex < numel\n"
+            "    a = tl.load(A + xindex, mask=xmask)\n"
+            "{prologue_load_b}"
+            "    _kernel_val_0 = a + b\n"
+            "    x_epilogue0_0 = xindex\n"
+            "    _tile_mask_0 = xmask\n"
+            "    <STORE_OUTPUT_0>\n"
+            "\n"
+            "def {kernel_name}(A, B, {out_param}{extra_sig}, numel):\n"
+            "    grid = ((numel + {xblock} - 1) // {xblock},)\n"
+            "    _mock_inner_add[grid]("
+            "A, B, {out_param}{extra_sig}, numel, XBLOCK={xblock})\n"
+            "    return {out_param}\n"
+        )
+
+        class _MockExternalTemplateBuffer(ir.TemplateBuffer):
+            def __init__(self, layout, inputs):
+                tb_self = self
+
+                def _make_kernel_render(out_node, hint_override=None):
+                    kernel = ExternalTritonTemplateKernel(tb_self)
+
+                    def render():
+                        return tb_self._render(kernel)
+
+                    return kernel, render
+
+                super().__init__(
+                    layout,
+                    inputs,
+                    _make_kernel_render,
+                    named_inputs={"A": inputs[0], "B": inputs[1]},
+                )
+                # Allow prologue fusion on input B (sigmoid(b) can be
+                # absorbed so the template reads b directly)
+                self.allowed_prologue_inps = OrderedSet([inputs[1].get_name()])
+                self.epilogue_fusable_outputs = {self.name: "result"}
+
+            def _render(self, kernel):
+                render_called[0] = True
+
+                # Set up all fusion hooks in one call
+                kernel._setup_fusion_hooks()
+
+                # --- Prologue handling for B ---
+                b_arg = self.inputs[1].get_name()
+                prologue_load_b = "    b = tl.load(B + xindex, mask=xmask)\n"
+                if kernel._prologue_source_buffers.get("B") is not None:
+                    b_arg = kernel._prologue_source_buffers["B"]
+                    prologue_load_b = (
+                        "    _prologue_B_xindex = xindex\n"
+                        "    _prologue_B_xmask = xmask\n"
+                        "    <LOAD_INPUT_B>\n"
+                        "    b = _prologue_B_result\n"
+                    )
+
+                call_args = [self.inputs[0].get_name(), b_arg]
+
+                # --- Epilogue handling ---
+                out_param = "result"
+                out_arg = self.name
+                for buf, param in kernel._extra_store_targets.items():
+                    out_param = param
+                    out_arg = buf
+                    break
+
+                call_args.append(out_arg)
+
+                # --- Extra inputs ---
+                extra_params = []
+                for buf_name, param_name in kernel._extra_inputs.items():
+                    extra_params.append(param_name)
+                    call_args.append(buf_name)
+
+                numel = self.get_size()[0]
+                call_args.append(str(numel))
+
+                extra_sig = ", " + ", ".join(extra_params) if extra_params else ""
+
+                source = _MOCK_ADD_KERNEL_TEMPLATE.format(
+                    out_param=out_param,
+                    extra_sig=extra_sig,
+                    prologue_load_b=prologue_load_b,
+                    kernel_name=str(Placeholder.KERNEL_NAME),
+                    xblock=XBLOCK,
+                )
+
+                kernel._call_preamble = []
+                kernel._call_args = call_args
+
+                return PartialRender(source, kernel.render_hooks)
+
+        def add_override(a, b, alpha=None):
+            layout = FixedLayout(a.get_device(), a.get_dtype(), a.get_size())
+            a = ir.ExternKernel.require_stride1(ir.ExternKernel.realize_input(a))
+            b = ir.ExternKernel.require_stride1(ir.ExternKernel.realize_input(b))
+            return ir.TensorBox.create(_MockExternalTemplateBuffer(layout, [a, b]))
+
+        # (override_fn, decompose, type_promotion, convert_input_to_bool)
+        with patch_lowering(
+            {
+                torch.ops.aten.add.Tensor: (
+                    add_override,
+                    True,
+                    ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT,
+                    False,
+                )
+            }
+        ):
+
+            @torch.compile
+            def f(a, b, bias):
+                # Use * for bias so it doesn't trigger add_override again
+                return torch.relu(a + torch.sigmoid(b)) * bias
+
+            a = torch.randn(32, device=GPU_TYPE)
+            b = torch.randn(32, device=GPU_TYPE)
+            bias = torch.randn(32, device=GPU_TYPE)
+
+            result, (code,) = run_and_get_code(f, a, b, bias)
+            expected = torch.relu(a + torch.sigmoid(b)) * bias
+            torch.testing.assert_close(result, expected)
+
+            # Verify render() was called (new protocol)
+            self.assertTrue(render_called[0])
+            # Verify template kernel was used
+            self.assertIn("_mock_inner_add", code)
+            # Verify epilogue fusion: relu fused via hook
+            self.assertIn("triton_helpers.maximum", code)
+            # Verify prologue fusion: sigmoid fused via hook
+            self.assertIn("tl.sigmoid", code)
+
 
 if __name__ == "__main__":
     if IS_LINUX and HAS_GPU and is_big_gpu():
diff --git a/test/inductor/test_snode_runtime.py b/test/inductor/test_snode_runtime.py
index cee78592153db..51be724876955 100644
--- a/test/inductor/test_snode_runtime.py
+++ b/test/inductor/test_snode_runtime.py
@@ -75,11 +75,13 @@ def tearDown(self):
         super().tearDown()
 
     def assertZero(self, x: float):
-        assert isinstance(x, float)
+        if not isinstance(x, float):
+            raise AssertionError(f"Expected float, got {type(x)}")
         super().assertEqual(x, 0.0, atol=0, rtol=0)
 
     def assertNotZero(self, x):
-        assert isinstance(x, float)
+        if not isinstance(x, float):
+            raise AssertionError(f"Expected float, got {type(x)}")
         super().assertNotEqual(x, 0.0, atol=0, rtol=0)
 
 
@@ -221,6 +223,97 @@ def f(x):
         self.assertNotZero(calculate_runtime(f, *inp))
 
 
+class InputDistanceTests(TestCase):
+    device = DEVICE
+
+    def _get_snodes(self, f, *args):
+        metrics.reset()
+        torch._logging.set_logs(inductor_metrics=True)
+        torch.compile(f, backend=compile_but_use_eager)(*args)
+        torch._logging.set_logs()
+        return [snode for snode, _ in metrics.nodes_num_elem]
+
+    def test_chain_with_reduction(self):
+        """
+        input -> sum (depth 0) -> sum (depth 1)
+        Reductions prevent full fusion, giving us distinct depth levels.
+        """
+
+        def f(x):
+            a = x.sum(dim=-1)
+            return a.sum(dim=-1)
+
+        snodes = self._get_snodes(f, T(10, 10, 10))
+        all_min = [s.min_input_distance for s in snodes]
+        all_max = [s.max_input_distance for s in snodes]
+        self.assertEqual(min(all_min), 0)
+        self.assertEqual(max(all_max), 1)
+
+    def test_fused_node_depth_range(self):
+        """
+        A reduction fused with its pointwise epilogue should have
+        min_input_distance=0 and max_input_distance=1.
+        """
+
+        def f(x):
+            a = x.sum(dim=-1)
+            return a.cos()
+
+        snodes = self._get_snodes(f, T(10, 10))
+        # The reduction and pointwise get fused
+        self.assertEqual(len(snodes), 1)
+        self.assertEqual(snodes[0].min_input_distance, 0)
+        self.assertEqual(snodes[0].max_input_distance, 1)
+
+    def test_extern_kernel_chain(self):
+        """
+        mm (depth 0, extern) -> cos+sum fused (depth 1)
+        """
+
+        def f(a, b):
+            c = torch.mm(a, b)
+            d = c.cos()
+            return d.sum(dim=-1)
+
+        snodes = self._get_snodes(f, T(10, 10), T(10, 10))
+        all_min = [s.min_input_distance for s in snodes]
+        all_max = [s.max_input_distance for s in snodes]
+        self.assertEqual(min(all_min), 0)
+        self.assertEqual(max(all_max), 1)
+
+    def test_foreach_basic(self):
+        """
+        foreach_add on graph inputs should have depth 0.
+        """
+
+        def f(xs, ys):
+            return torch._foreach_add(xs, ys)
+
+        xs = [T(10), T(20)]
+        ys = [T(10), T(20)]
+        snodes = self._get_snodes(f, xs, ys)
+        for s in snodes:
+            self.assertEqual(s.min_input_distance, 0)
+            self.assertEqual(s.max_input_distance, 0)
+
+    def test_foreach_after_extern(self):
+        """
+        mm (extern, depth 0) -> foreach_add (depth 1)
+        The extern kernel creates a fusion barrier so the foreach
+        has a real dependency chain.
+        """
+
+        def f(a, b, ys):
+            c = torch.mm(a, b)
+            return torch._foreach_add([c, c], ys)
+
+        snodes = self._get_snodes(f, T(10, 10), T(10, 10), [T(10, 10), T(10, 10)])
+        all_min = [s.min_input_distance for s in snodes]
+        all_max = [s.max_input_distance for s in snodes]
+        self.assertEqual(min(all_min), 0)
+        self.assertGreaterEqual(max(all_max), 1)
+
+
 @skipIf(not dist.is_available(), "requires distributed")
 class TestCommAnalysis(TestCase):
     device = DEVICE
diff --git a/test/inductor/test_split_cat_fx_aten_passes.py b/test/inductor/test_split_cat_fx_aten_passes.py
index 9b2f62e27488e..80a2df9f080f4 100644
--- a/test/inductor/test_split_cat_fx_aten_passes.py
+++ b/test/inductor/test_split_cat_fx_aten_passes.py
@@ -226,7 +226,8 @@ def compare_dict_tensors(self, ref_dict, res_dict, rtol=1e-3, atol=1e-3):
             return False
         for key1 in ref_dict:
             key2 = "_orig_mod." + key1
-            assert key2 in res_dict, f"{key1} does not exist in traced module"
+            if key2 not in res_dict:
+                raise AssertionError(f"{key1} does not exist in traced module")
             if not torch.allclose(ref_dict[key1], res_dict[key2], rtol=rtol, atol=atol):
                 return False
         return True
diff --git a/test/inductor/test_split_cat_fx_passes.py b/test/inductor/test_split_cat_fx_passes.py
index c1fc5ab8dd93f..aae07ba53d61e 100644
--- a/test/inductor/test_split_cat_fx_passes.py
+++ b/test/inductor/test_split_cat_fx_passes.py
@@ -179,15 +179,16 @@ def unequal_multi_split(x):
             return [torch.relu(s) for s in final_items]
 
         def unequal_multi_split_neg_index(x):
-            fs = torch.split(x, [10, 10, 12], dim=1)
+            fs = torch.split(x, [8, 8, 16], dim=1)
             item0 = fs[-3]
             item1 = fs[-2]
             item2 = fs[-1]
 
+            # Use different splits from `unequal_multi_split` to avoid cache hit.
             final_items = []
-            final_items.extend(item0.split([4, 6], 1))
-            final_items.extend(item1.split([6, 4], 1))
-            final_items.extend(item2.split([4, 4, 4], 1))
+            final_items.extend(item0.split([3, 5], 1))
+            final_items.extend(item1.split([5, 3], 1))
+            final_items.extend(item2.split([8, 4, 4], 1))
 
             return [torch.relu(s) for s in final_items]
 
@@ -241,47 +242,46 @@ def duplicate_getitems_neg_index(x):
             item1_2 = fs[-2]  # negative index
             item2 = fs[2]
 
+            # Use different splits from `duplicate_getitems` to avoid cache hit.
             final_items = []
-            final_items.extend(item0.split([4, 6], 1))
-            final_items.extend(item1_1.split([6, 4], 1))
+            final_items.extend(item0.split([3, 7], 1))
+            final_items.extend(item1_1.split([7, 3], 1))
             final_items.extend(item1_2)
             final_items.append(torch.sin(item2))
 
             return [torch.relu(s) for s in final_items]
 
         def split_getitem_gap(x):
+            # Clashes with `split_partial_getitem_cat`
             fs = torch.split(x, [4, 4, 24], dim=1)
             item0 = fs[0]
             item2 = fs[2]
 
-            final_items = [
-                item0,
-            ]
+            final_items = [item0]
             final_items.extend(item2.split((4, 4, 4, 4, 4, 4), 1))
 
             return torch.cat(final_items, dim=1)
 
-        def split_getitem_out_of_order(x):
-            fs = torch.split(x, [4, 4, 4, 20], dim=1)
+        def split_partial_getitem_cat(x):
+            # Use different splits from `split_getitem_gap` to avoid cache hit.
+            fs = torch.split(x, [2, 6, 24], dim=1)
             item0 = fs[0]
             item2 = fs[2]
-            item1 = fs[1]
-            item3 = fs[3]
 
-            final_items = [item0, item2, item1]
-            final_items.extend(item3.split((4, 4, 4, 4, 4), 1))
+            final_items = [item0]
+            final_items.extend(item2.split((4, 4, 4, 4, 4, 4), 1))
 
             return torch.cat(final_items, dim=1)
 
-        def split_partial_getitem_cat(x):
-            fs = torch.split(x, [4, 4, 24], dim=1)
+        def split_getitem_out_of_order(x):
+            fs = torch.split(x, [4, 4, 4, 20], dim=1)
             item0 = fs[0]
             item2 = fs[2]
+            item1 = fs[1]
+            item3 = fs[3]
 
-            final_items = [
-                item0,
-            ]
-            final_items.extend(item2.split((4, 4, 4, 4, 4, 4), 1))
+            final_items = [item0, item2, item1]
+            final_items.extend(item3.split((4, 4, 4, 4, 4), 1))
 
             return torch.cat(final_items, dim=1)
 
@@ -315,9 +315,9 @@ def next_split_getitem_partial_used(x):
             (duplicate_getitems, 1),
             (duplicate_getitems_neg_index, 1),
             (split_getitem_gap, 1),
+            (split_partial_getitem_cat, 1),
             (split_getitem_out_of_order, 1),
             (next_split_getitem_partial_used, 1),
-            (split_partial_getitem_cat, 1),
         ]:
             expected = fn(*args)
             actual = torch.compile(fn)(*args)
diff --git a/test/inductor/test_static_triton_launcher.py b/test/inductor/test_static_triton_launcher.py
index 84a4804c297d8..04811d3e8bf94 100644
--- a/test/inductor/test_static_triton_launcher.py
+++ b/test/inductor/test_static_triton_launcher.py
@@ -1,8 +1,9 @@
 # Owner(s): ["module: inductor"]
+import gc
 import os
 import random
 import tempfile
-from typing import Union
+import weakref
 from unittest import mock
 
 import torch
@@ -41,6 +42,41 @@ def launcher_with_native_code(*args, **kwargs):
     return launcher_with_native_code
 
 
+class TestStaticTritonLauncherUnit(TestCase):
+    def test_xpu_load_kernel_uses_existing_three_tuple_abi(self):
+        load_calls = []
+        kernel_capsule = object()
+
+        class FakeImpl:
+            @staticmethod
+            def _load_kernel(path, name, shared, device):
+                load_calls.append((path, name, shared, device))
+                return kernel_capsule, 7, 11
+
+        launcher = object.__new__(StaticallyLaunchedXpuKernel)
+        launcher.function = None
+        launcher.module = None
+        launcher.cubin_path = "/tmp/kernel.zebin"
+        launcher.cubin_raw = b"zebin"
+        launcher.name = "kernel"
+        launcher.shared = 13
+        launcher.C_impl = FakeImpl
+
+        launcher.load_kernel(3)
+
+        self.assertEqual(load_calls, [("/tmp/kernel.zebin", "kernel", 13, 3)])
+        self.assertIs(launcher.function, kernel_capsule)
+        self.assertEqual(launcher.n_regs, 7)
+        self.assertEqual(launcher.n_spills, 11)
+        self.assertIsNone(launcher.module)
+        self.assertIsNone(launcher.cubin_path)
+        self.assertIsNone(launcher.cubin_raw)
+
+        launcher.close()
+        self.assertIsNone(launcher.function)
+        self.assertIsNone(launcher.module)
+
+
 @requires_gpu_and_triton
 class TestStaticTritonLauncher(TestCase):
     def setUp(self):
@@ -82,7 +118,7 @@ def write_cubin_to_tmp(self, kernel: CompiledKernel) -> str:
     def _make_launcher(
         self,
         compiled_kernel: CompiledKernel,
-    ) -> Union[StaticallyLaunchedCudaKernel, StaticallyLaunchedXpuKernel]:
+    ) -> StaticallyLaunchedCudaKernel | StaticallyLaunchedXpuKernel:
         """
         Compiles a Triton kernel with the provided *args,
         writes its cubin to the temporary file, and returns the file path.
@@ -92,7 +128,8 @@ def _make_launcher(
         result = statically_launched_kernel_by_device(compiled_kernel, GPU_TYPE)
         # Test reload cubin from raw here
         old_cubin_path = result.cubin_path
-        assert old_cubin_path is not None
+        if old_cubin_path is None:
+            raise AssertionError
         result.cubin_path = None
         result.reload_cubin_from_raw(old_cubin_path)
         device_interface = get_interface_for_device(GPU_TYPE)
@@ -420,6 +457,56 @@ def kernel_many_args(out_tensor, {decl}):
         launcher.run(1, 1, 1, stream, buf1, *kernel_args)
         self.assertEqual(buf0, buf1)
 
+    def test_launcher_keeps_module_owner_alive_until_release(self):
+        """
+        Generated launchers capture ``runner=self.kernel.run``.  That closure
+        must keep the static kernel owner alive, otherwise its module could be
+        unloaded while the launcher can still be called.
+        """
+        unloaded_modules = []
+
+        class FakeImpl:
+            @staticmethod
+            def _unload_kernel(mod):
+                unloaded_modules.append(mod)
+
+        class FakeKernelOwner:
+            __del__ = StaticallyLaunchedCudaKernel.__del__
+
+            def __init__(self):
+                self.name = "fake_kernel"
+                self.module = 0xC0FFEE
+                self.function = 0xF00D
+                self.C_impl = FakeImpl
+
+            def close(self):
+                return StaticallyLaunchedCudaKernel.close(self)
+
+            def run(self, *_args, **_kwargs):
+                return None
+
+        owner = FakeKernelOwner()
+        owner_ref = weakref.ref(owner)
+
+        class FakeLauncher:
+            def __init__(self, runner):
+                self.runner = runner
+
+            def __call__(self, stream):
+                return self.runner(1, 1, 1, stream)
+
+        launcher = FakeLauncher(owner.run)
+
+        del owner
+        gc.collect()
+        self.assertIsNotNone(owner_ref())
+        self.assertEqual(unloaded_modules, [])
+
+        del launcher
+        gc.collect()
+        self.assertIsNone(owner_ref())
+        self.assertEqual(unloaded_modules, [0xC0FFEE])
+
 
 @requires_gpu_and_triton
 @torch._inductor.config.patch(
diff --git a/test/inductor/test_subgraph_choice.py b/test/inductor/test_subgraph_choice.py
index 408af8d379111..27100f7b66b8f 100644
--- a/test/inductor/test_subgraph_choice.py
+++ b/test/inductor/test_subgraph_choice.py
@@ -7,7 +7,6 @@
 from torch._inductor.lowering import register_lowering
 from torch._inductor.select_algorithm import autotune_select_algorithm
 from torch._inductor.test_case import run_tests, TestCase
-from torch.testing._internal.common_utils import skipIfXpu
 from torch.testing._internal.inductor_utils import GPU_TYPE, HAS_CPU, HAS_GPU
 
 
@@ -35,7 +34,6 @@ def _create_buffer(self, name, shape, dtype):
             layout=FixedLayout(torch.device(f"{GPU_TYPE}:0"), dtype=dtype, size=shape),
         )
 
-    @skipIfXpu
     def test_subgraph_decompose_k(self):
         from torch._inductor.kernel.mm import aten_mm
         from torch._inductor.kernel.mm_common import mm_args
@@ -74,9 +72,10 @@ def _(a, b):
 
             # Only return decomposeK case for codegen
             choices = [choices[1]]
-            return autotune_select_algorithm(
+            node, _ = autotune_select_algorithm(
                 "test_subgraph_choice", choices, [a, b], layout
             )
+            return node
 
         a_in = torch.randn(
             mat1_shape, dtype=torch.float16, device=torch.device(f"{GPU_TYPE}:0")
@@ -95,7 +94,6 @@ def func(mat1, mat2):
         # Check same results of compiled result and regular torch.mm
         torch.testing.assert_close(res, a_in @ b_in, atol=1e-1, rtol=1e-1)
 
-    @skipIfXpu
     def test_subgraph_freeze_layout(self):
         from torch._inductor.kernel.mm_common import mm_args
 
@@ -119,7 +117,8 @@ def _(a, b):
         def _(a, b):
             _, _, _, layout, mat1, mat2 = mm_args(a, b)
             mat1_layout = mat1.layout
-            assert isinstance(mat1_layout, FlexibleLayout)
+            if not isinstance(mat1_layout, FlexibleLayout):
+                raise AssertionError
             mat1_stride = mat1_layout.stride
 
             choices = []
@@ -138,21 +137,25 @@ def _(a, b):
             )
 
             choice = choices[0]
-            assert isinstance(mat1.layout, FixedLayout)
+            if not isinstance(mat1.layout, FixedLayout):
+                raise AssertionError
 
             # Creating the subgraph choice should have frozen the layout
             # We ensure padding so the stride should differ
-            assert mat1.layout.stride != mat1_stride
+            if mat1.layout.stride == mat1_stride:
+                raise AssertionError
 
             for example_stride, layout_stride in zip(
                 choice.example_inputs[0].stride(), mat1.layout.stride
             ):
                 # Example inputs should have same stride as current layout
-                assert example_stride == layout_stride
+                if example_stride != layout_stride:
+                    raise AssertionError
 
-            return autotune_select_algorithm(
+            node, _ = autotune_select_algorithm(
                 "test_subgraph_choice", choices, [a, b], layout
             )
+            return node
 
         def func(mat1, mat2):
             return torch.ops.mylib.matmul_decompose_padding((mat1 + 1.0), mat2)
diff --git a/test/inductor/test_torchinductor.py b/test/inductor/test_torchinductor.py
index 6c4f1b3f92890..68842494e1bb5 100644
--- a/test/inductor/test_torchinductor.py
+++ b/test/inductor/test_torchinductor.py
@@ -1,6 +1,7 @@
 # Owner(s): ["module: inductor"]
 # ruff: noqa: F841
 import contextlib
+import contextvars
 import copy
 import dataclasses
 import functools
@@ -83,7 +84,11 @@
     expectedFailureXPU,
     largeTensorTest,
 )
-from torch.testing._internal.common_dtype import all_types, get_all_dtypes
+from torch.testing._internal.common_dtype import (
+    all_types,
+    get_all_dtypes,
+    highest_precision_float,
+)
 from torch.testing._internal.common_quantization import (
     _dynamically_quantize_per_channel,
     _group_quantize_tensor_symmetric,
@@ -91,6 +96,8 @@
 from torch.testing._internal.common_utils import (
     DeterministicGuard,
     instantiate_parametrized_tests,
+    IS_ARM64,
+    IS_CPU_EXT_SVE_SUPPORTED,
     IS_FBCODE,
     IS_MACOS,
     IS_X86,
@@ -99,7 +106,6 @@
     NAVI_ARCH,
     parametrize,
     serialTest,
-    skipIfMPS,
     skipIfRocm,
     skipIfRocmArch,
     skipIfWindows,
@@ -107,6 +113,7 @@
     subtest,
     TEST_WITH_ASAN,
     TEST_WITH_ROCM,
+    xfailIf,
     xfailIfS390X,
 )
 from torch.testing._internal.logging_utils import logs_to_string
@@ -358,7 +365,8 @@ def tearDown(self):
         torch._dynamo.reset()
         if os.environ.get("ERROR_ON_SLOW") == "1":
             elapsed = time.perf_counter() - self._start
-            assert elapsed < 120
+            if elapsed >= 120:
+                raise AssertionError(f"Test took too long: {elapsed:.1f}s >= 120s")
 
 
 class ToTuple(torch.nn.Module):
@@ -412,10 +420,14 @@ def gather_leaf_tensors(args, kwargs):
     flat_diff_results = [
         r for r in flat_results if isinstance(r, torch.Tensor) and r.requires_grad
     ]
-    assert len(flat_diff_results) > 0
+    if len(flat_diff_results) <= 0:
+        raise AssertionError(
+            f"Expected len(flat_diff_results) > 0, got {len(flat_diff_results)}"
+        )
 
     leaf_tensors = gather_leaf_tensors(args, kwrags)
-    assert len(leaf_tensors) > 0
+    if len(leaf_tensors) <= 0:
+        raise AssertionError(f"Expected len(leaf_tensors) > 0, got {len(leaf_tensors)}")
     return torch.autograd.grad(
         flat_diff_results,
         leaf_tensors,
@@ -522,14 +534,21 @@ def run(*ex, **kwargs):
     #     for graph in exp[2]:
     #         print("Graph", graph)
     if check_has_compiled:
-        assert called, "Ran graph without calling compile_fx"
-    assert type(actual) is type(correct)
+        if not called:
+            raise AssertionError("Ran graph without calling compile_fx")
+    if type(actual) is not type(correct):
+        raise AssertionError(f"Expected type {type(correct)}, got {type(actual)}")
     if isinstance(actual, (tuple, list)):
-        assert len(actual) == len(correct)
-        assert all(
+        if len(actual) != len(correct):
+            raise AssertionError(f"Expected length {len(correct)}, got {len(actual)}")
+        if not all(
             type(actual_item) is type(correct_item)
             for actual_item, correct_item in zip(actual, correct)
-        )
+        ):
+            raise AssertionError(
+                f"Item type mismatch: expected types {[type(c) for c in correct]}, "
+                f"got {[type(a) for a in actual]}"
+            )
 
     correct_flat, correct_spec = tree_flatten(correct)
     actual_flat = pytree.tree_leaves(actual)
@@ -557,9 +576,10 @@ def reference_to_expect(actual_flat, correct_flat):
     if reference_in_float and exact_dtype:
         for expect_dtype, actual_result in zip(expect_dtypes, actual_flat):
             if expect_dtype is not None:
-                assert actual_result.dtype == expect_dtype, (
-                    f"dtype mismatch, expected {expect_dtype} but got {actual_result.dtype}"
-                )
+                if actual_result.dtype != expect_dtype:
+                    raise AssertionError(
+                        f"dtype mismatch, expected {expect_dtype} but got {actual_result.dtype}"
+                    )
 
     if reference_in_float:
         correct_flat = reference_to_expect(actual_flat, correct_flat)
@@ -607,18 +627,37 @@ def custom_assert_with_self(*args, **kwargs):
     else:
         for correct_val, actual_val in zip(correct_flat, actual_flat):
             if isinstance(correct_val, torch.Tensor):
-                assert correct_val.device == actual_val.device
-                assert correct_val.size() == actual_val.size()
+                if correct_val.device != actual_val.device:
+                    raise AssertionError(
+                        f"Expected device {correct_val.device}, got {actual_val.device}"
+                    )
+                if correct_val.size() != actual_val.size():
+                    raise AssertionError(
+                        f"Expected size {correct_val.size()}, got {actual_val.size()}"
+                    )
                 strides_equal, _ = torch._prims_common.check_significant_strides(
                     correct_val, actual_val
                 )
-                assert strides_equal
-                assert correct_val.layout == actual_val.layout
+                if not strides_equal:
+                    raise AssertionError(
+                        f"Significant strides mismatch: expected {correct_val.stride()}, "
+                        f"got {actual_val.stride()}"
+                    )
+                if correct_val.layout != actual_val.layout:
+                    raise AssertionError(
+                        f"Expected layout {correct_val.layout}, got {actual_val.layout}"
+                    )
                 if exact_dtype:
-                    assert correct_val.dtype == actual_val.dtype
+                    if correct_val.dtype != actual_val.dtype:
+                        raise AssertionError(
+                            f"Expected dtype {correct_val.dtype}, got {actual_val.dtype}"
+                        )
                 check_exact_stride = exact_stride and not has_zero_dim(correct_val)
                 if check_exact_stride:
-                    assert correct_val.stride() == actual_val.stride()
+                    if correct_val.stride() != actual_val.stride():
+                        raise AssertionError(
+                            f"Expected stride {correct_val.stride()}, got {actual_val.stride()}"
+                        )
 
     if check_gradient:
         actual = output_process_fn_grad(actual)
@@ -811,11 +850,19 @@ def _run_and_assert_no_indirect_indexing(
                 ("where" in code or ") ? (" in code) is has_wrapping,
                 msg=f"Wanted {has_wrapping=} but got\n{code}",
             )
+
+    def has_assert_in_code(code):
+        # Check for device_assert, or TORCH_CHECK that isn't AOTI_TORCH_CHECK
+        # AOTI_TORCH_CHECK is for lazy Triton compile infrastructure, not bounds checking
+        if "device_assert" in code:
+            return True
+        for line in code.split("\n"):
+            if "TORCH_CHECK" in line and "AOTI_TORCH_CHECK" not in line:
+                return True
+        return False
+
     test_case.assertTrue(
-        any(
-            ("device_assert" in code or "TORCH_CHECK" in code) is has_assert
-            for code in source_codes
-        )
+        any(has_assert_in_code(code) is has_assert for code in source_codes)
     )
     return result
 
@@ -847,9 +894,17 @@ class SweepInputs2:
     def kernel(a, b):
         return (a + b,)
 
+    # Input gen types unsupported by specific backends. Subclasses or wrapper
+    # test classes can set this to skip tests that use these input types.
+    _unsupported_input_gen_types: set[str] = set()
+
     @classmethod
     def gen_template(cls, name1, name2):
         def test(self):
+            unsupported = type(self)._unsupported_input_gen_types
+            for name in (name1, name2):
+                if name in unsupported:
+                    self.skipTest(f"{name} inputs not supported by this backend")
             check_model(
                 self,
                 cls.kernel,
@@ -1003,6 +1058,11 @@ def xfail_if_triton_cpu(fn):
     return fn
 
 
+def xfail_if_pallas(fn):
+    fn._expected_failure_pallas = True
+    return fn
+
+
 def skip_if_gpu_halide(fn):
     @functools.wraps(fn)
     def wrapper(self, *args, **kwargs):
@@ -2152,13 +2212,15 @@ def test_reduction_config_limit(self):
         triton reduction configs for large size hints. #128826 introduced a scaling XBLOCK
         feature to resolve the issue in reduction configs which may exceed the maxGridSize
         """
+        from torch._inductor.runtime.hints import get_warp_size
         from torch._inductor.runtime.runtime_utils import next_power_of_2
         from torch._inductor.runtime.triton_heuristics import triton_config_reduction
 
+        warp_size = get_warp_size(torch.device(self.device))
         size_hints = {"x": 67108864, "r0_": 8192}
         for _ in range(4):
             size_hints["x"] = next_power_of_2(size_hints["x"])
-            triton_config_reduction(size_hints, 1, 2048, 1, 8)
+            triton_config_reduction(size_hints, 1, 2048, 1, 8, warp_size=warp_size)
 
     def test_prod(self):
         def fn(a):
@@ -2431,9 +2493,14 @@ def fn(a, b, dim):
         _, code = run_and_get_code(cfn, a, b, 0)
 
         # Check everything is fused into a single kernel
-        FileCheck().check_not("run(").check_regex(
-            r"triton_.*\.run\(arg[01]_1, arg[12]_1, buf1,"
-        ).check_not("run(").run(code[0])
+        if config.cpp_wrapper:
+            FileCheck().check_regex(r"call_triton_.*\(arg[01]_1, arg[12]_1, buf1,").run(
+                code[0]
+            )
+        else:
+            FileCheck().check_not("run(").check_regex(
+                r"triton_.*\.run\(arg[01]_1, arg[12]_1, buf1,"
+            ).check_not("run(").run(code[0])
 
     @skip_if_halide  # scan ops
     # TODO: support lifted symints when dynamic
@@ -2544,6 +2611,8 @@ def fn(a, b_int8pack, b_scales, c):
     @xfail_if_triton_cpu
     @skipCUDAIf(True, "No _dyn_quant_pack_4bit_weight implementation on CUDA")
     @skipIfXpu(msg="No _dyn_quant_pack_4bit_weight implementation on XPU")
+    # Pallas codegen doesn't handle reduction axis after FloorDiv(ModularIndexing) simplification
+    @xfail_if_pallas
     def test__dyn_quant_pack_4bit_weight_fp32(self):
         q_group = 32
         k = 128
@@ -2580,6 +2649,8 @@ def fn(b, in_features, out_features):
     @skipCUDAIf(True, "No _dyn_quant_pack_4bit_weight implementation on CUDA")
     @skipIfXpu(msg="No _dyn_quant_pack_4bit_weight implementation on XPU")
     @skip_if_halide  # bf16
+    # Pallas codegen doesn't handle reduction axis after FloorDiv(ModularIndexing) simplification
+    @xfail_if_pallas
     def test__dyn_quant_pack_4bit_weight_bf16(self):
         k = 128
         n = 128
@@ -2618,6 +2689,8 @@ def fn(b, in_features, out_features):
 
     @xfail_if_mps_unimplemented
     @xfail_if_triton_cpu
+    # Pallas codegen doesn't handle reduction axis after FloorDiv(ModularIndexing) simplification
+    @xfail_if_pallas
     @skipCUDAIf(True, "No _dyn_quant_matmul_4bit implementation on CUDA")
     @skipIfXpu(msg="No _dyn_quant_matmul_4bit implementation on XPU")
     def test__dyn_quant_matmul_4bit_fp32_input(self):
@@ -2661,8 +2734,10 @@ def fn(a, q_group, in_features, out_features):
         self.common(fn, (a, q_group, in_features, out_features))
 
     @skipCPUIf(IS_MACOS, "fails on M1, mismatch in bf16 support reporting")
-    @xfail_if_mps_unimplemented
     @xfail_if_triton_cpu
+    @xfailIf(
+        IS_ARM64 and not IS_CPU_EXT_SVE_SUPPORTED
+    )  # see https://github.com/pytorch/pytorch/issues/170787
     @skipCUDAIf(True, "No _dyn_quant_matmul_4bit implementation on CUDA")
     @skipIfXpu(msg="No _dyn_quant_matmul_4bit implementation on XPU")
     @skip_if_halide  # bf16
@@ -2862,6 +2937,92 @@ def fn(x):
         a = torch.rand(())
         self.common(fn, (a,))
 
+    def test_cumprod_backward(self):
+        if self.device == "mps":
+            raise unittest.SkipTest(
+                "MPS inductor codegen bug with argmax: threadgroup_argmax"
+            )
+
+        # Regression test for https://github.com/pytorch/pytorch/issues/136263
+        # torch.compile used O(n^2) algorithm for cumprod backward with tensor
+        # subclasses (like FakeTensor), making it extremely slow.
+        def fn(x):
+            y = x.cumprod(dim=1)
+            return y.sum()
+
+        for dim_size in [8, 32]:
+            x = torch.rand(4, dim_size, 64, device=self.device, requires_grad=True)
+            x_ref = x.clone().detach().requires_grad_(True)
+
+            # Eager forward + backward
+            y_ref = fn(x_ref)
+            y_ref.backward()
+
+            # Compiled forward + backward
+            compiled_fn = torch.compile(fn, fullgraph=True)
+            y = compiled_fn(x)
+            y.backward()
+
+            # Check correctness
+            self.assertEqual(y, y_ref, atol=1e-4, rtol=1e-4)
+            self.assertEqual(x.grad, x_ref.grad, atol=1e-4, rtol=1e-4)
+
+    def test_cumprod_backward_with_zeros(self):
+        if self.device == "mps":
+            raise unittest.SkipTest(
+                "MPS inductor codegen bug with argmax: threadgroup_argmax"
+            )
+
+        # Test cumprod backward with zeros in the input
+        # This exercises the more complex O(n) algorithm path
+        def fn(x):
+            y = x.cumprod(dim=1)
+            return y.sum()
+
+        for dim_size in [8, 16]:
+            x = torch.rand(4, dim_size, 32, device=self.device, requires_grad=True)
+            # Insert some zeros to exercise the zero-handling path
+            x.data[:, dim_size // 2, :] = 0
+            x_ref = x.clone().detach().requires_grad_(True)
+
+            # Eager forward + backward
+            y_ref = fn(x_ref)
+            y_ref.backward()
+
+            # Compiled forward + backward
+            compiled_fn = torch.compile(fn, fullgraph=True)
+            y = compiled_fn(x)
+            y.backward()
+
+            # Check correctness
+            self.assertEqual(y, y_ref, atol=1e-4, rtol=1e-4)
+            self.assertEqual(x.grad, x_ref.grad, atol=1e-4, rtol=1e-4)
+
+    def test_view_dtype_bool(self):
+        # Regression test for boolean dtype handling in view.dtype lowering
+        # torch.iinfo doesn't support bool, so we need special handling
+        def fn(x, mask):
+            # Create a computation that involves boolean tensors and where
+            result = torch.where(mask, x, torch.zeros_like(x))
+            return result.sum()
+
+        x = torch.rand(4, 8, device=self.device, requires_grad=True)
+        mask = torch.rand(4, 8, device=self.device) > 0.5
+
+        x_ref = x.clone().detach().requires_grad_(True)
+
+        # Eager
+        y_ref = fn(x_ref, mask)
+        y_ref.backward()
+
+        # Compiled
+        compiled_fn = torch.compile(fn, fullgraph=True)
+        y = compiled_fn(x, mask)
+        y.backward()
+
+        self.assertEqual(y, y_ref)
+        self.assertEqual(x.grad, x_ref.grad)
+
     def test_cumsum_inf(self):
         def fn(x):
             return x.cumsum(-1)
@@ -2938,6 +3099,22 @@ def fn(a):
 
         self.common(fn, (torch.randint(4, (4,)),))
 
+    def test_comparison_scalar_type_promotion_bf16(self):
+        if self.device == "mps":
+            raise unittest.SkipTest(
+                "MPS codegen doesn't upcast bf16 constants to float32"
+            )
+
+        def fn(x):
+            return x == 3.14
+
+        x = torch.tensor(
+            [3.14, 1.0, 3.140625, 0.0], dtype=torch.bfloat16, device=self.device
+        )
+        expected = fn(x)
+        actual = torch.compile(fn)(x)
+        self.assertEqual(actual, expected)
+
     @skip_if_gpu_halide
     @xfail_if_triton_cpu
     def test_dist(self):
@@ -2949,7 +3126,6 @@ def fn(a, b):
 
         self.common(fn, (torch.randn(4, 4), torch.randn(4, 4)))
 
-    @xfail_if_mps
     @skip_if_halide  # different pow accuracies
     @xfail_if_triton_cpu
     def test_norm_constant_overflow(self):
@@ -3027,6 +3203,15 @@ def fn(x):
         self.common(fn, (make_arg(1, dtype=torch.float32),))
         self.common(fn, (make_arg(1, dtype=torch.int64),))
 
+    def test_arange7(self):
+        def fn(x):
+            # Test aten.arange.start_step lowering with integer dtypes
+            return x + torch.ops.aten.arange.start_step(
+                0, 10, 2, dtype=torch.int64, device=x.device
+            )
+
+        self.common(fn, (torch.zeros(5, dtype=torch.int64),), check_lowp=False)
+
     def test_linspace1(self):
         def fn(x):
             return torch.linspace(0.125, 0.875, 7, device=x.device) + x
@@ -3263,7 +3448,7 @@ def test_round_correctness(self):
         def fn(a):
             return torch.round(a)
 
-        dtype = torch.float64 if self.device != "mps" else torch.float32
+        dtype = highest_precision_float(self.device)
         self.common(
             fn,
             [torch.arange(-10, 10, 0.1, dtype=dtype)],
@@ -3401,6 +3586,8 @@ def fn(a):
             check_lowp=False,
         )
 
+    test_one_hot._expected_failure_halide = True
+
     def test_div1(self):
         def fn(a, b):
             return (
@@ -3483,6 +3670,7 @@ def fn(a, b):
         )
 
     @skip_if_triton_cpu  # divide by zero; cannot xfail because it crashes process
+    @skipIfXpu(msg="https://github.com/intel/intel-xpu-backend-for-triton/issues/6401")
     def test_div7(self):
         def fn(a, b):
             return (
@@ -3582,6 +3770,64 @@ def fn_int_input(a, i):
             fn_int_input, (make_tensor(10, device=self.device, dtype=torch.float32), 33)
         )
 
+    def test_floordiv_negative_int(self):
+        # Regression test for https://github.com/pytorch/pytorch/issues/158212
+        # Signed integer floor division must handle negative dividends correctly.
+        # The key case: (-1) // 2 should be -1 (Python floor), not 0 (C truncation).
+        def fn(a, b):
+            return a // b
+
+        for divisor in [2, 3, 7]:
+            a = torch.arange(-32, 32, device=self.device, dtype=torch.int32)
+            b = torch.full_like(a, divisor)
+            self.common(fn, (a, b))
+            b_neg = torch.full_like(a, -divisor)
+            self.common(fn, (a, b_neg))
+
+    @skip_if_cpu
+    def test_floordiv_div_by_zero_int(self):
+        # Integer floor division by zero is undefined behavior on CUDA/Triton.
+        # On CPU, integer division by zero correctly raises ZeroDivisionError.
+        # Eager (c10::div_floor_integer) and compiled (Triton floordiv) must
+        # both return 0 for elements where the divisor is zero on GPU.
+        def fn(a, b):
+            return torch.floor_divide(a, b)
+
+        for dtype in [torch.int32, torch.int64]:
+            # All-zero divisor: every element should be 0
+            for dividend in [0, 1, -1, 5, -5]:
+                a = torch.full((8,), dividend, device=self.device, dtype=dtype)
+                b = torch.full((8,), 0, device=self.device, dtype=dtype)
+                self.common(fn, (a, b))
+
+    def test_floordiv_float_accuracy(self):
+        # Triton uses an approximate reciprocal for fp32 division, so a naive
+        # floor(a / b) can be off by one when the true quotient is very close
+        # to an integer.  Verify the corrected lowering matches eager.
+        def fn(a, b):
+            return torch.floor_divide(a, b)
+
+        # The original repro: fp16 values whose exact quotient is 4.0, but
+        # Triton's approximate fp32 division yields 3.9999997…
+        a = torch.full((8,), 84.3125, dtype=torch.float16, device=self.device)
+        b = torch.full((8,), 21.078125, dtype=torch.float16, device=self.device)
+        self.common(fn, (a, b))
+
+        # Negative operands / mixed signs
+        for dtype in [torch.float16, torch.float32]:
+            for a_val, b_val in [
+                (7.0, 2.0),
+                (7.0, -2.0),
+                (-7.0, 2.0),
+                (-7.0, -2.0),
+                (6.0, 3.0),
+                (0.0, 5.0),
+                (0.0, -5.0),
+            ]:
+                a = torch.full((8,), a_val, dtype=dtype, device=self.device)
+                b = torch.full((8,), b_val, dtype=dtype, device=self.device)
+                self.common(fn, (a, b))
+
     def test_div_precision(self):
         # Reproducer for https://github.com/pytorch/pytorch/issues/101039
 
@@ -3689,8 +3935,6 @@ def forward(x, y):
         cf(x, 1e-5)
         cf(x, 1e-6)
 
-    # I think Triton CPU doesn't preserve -0.0 sign in tl.full
-    @xfail_if_triton_cpu
     def test_div_by_zero(self):
         def fn(x, runtime_zero, runtime_neg_zero):
             zero = torch.zeros_like(x)
@@ -3702,7 +3946,6 @@ def fn(x, runtime_zero, runtime_neg_zero):
                 x / -zero,
                 zero / zero,
                 x / runtime_zero,
-                # NOTE: -runtime_zero doesn't work as -(0.0) is broken in triton
                 x / runtime_neg_zero,
                 runtime_zero / runtime_neg_zero,
             )
@@ -3769,6 +4012,7 @@ def fn(a, b):
         self.assertEqual(actual, expect)
 
     @skip_if_halide  # only 32-bit indexing
+    @skip_if_pallas  # only 32-bit indexing
     @largeTensorTest("4GB", inductor=True)
     def test_large_pointwise(self):
         def fn(a):
@@ -3786,6 +4030,7 @@ def fn(a):
         self.assertTrue((actual == 2).all())
 
     @skip_if_halide  # only 32-bit indexing
+    @skip_if_pallas  # only 32-bit indexing
     @largeTensorTest("3GB", inductor=True)
     def test_large_offset_pointwise(self):
         # Test 64-bit indexing is used when input views a tensor that can be
@@ -3956,7 +4201,7 @@ def fn(a, b, c):
         cfn = torch.compile(backend="inductor")(fn)
         input = torch.tensor([2], dtype=torch.int32)
         mat = torch.tensor(np.random.randn(0, 0), dtype=torch.int32)
-        vec = torch.tensor([])
+        vec = torch.tensor([], dtype=torch.int32)
         with torch.no_grad():
             self.assertEqual(cfn(input, mat, vec), fn(input, mat, vec))
 
@@ -4412,8 +4657,7 @@ def foo(m, inp):
                 )
             else:
                 if config.cpp_wrapper:
-                    # with cpp_wrapper, the pattern is .run( and then .convolution(
-                    FileCheck().check(".run(").check("convolution(").run(code[0])
+                    FileCheck().check("call_triton").check("convolution(").run(code[0])
                 else:
                     # without cpp_wrapper, the pattern is .convolution( and then .run(
                     FileCheck().check("convolution(").check(".run(").run(code[0])
@@ -4534,7 +4778,7 @@ def fn(x, y):
 
         x1 = torch.randn(30, device=self.device)
         x2 = torch.randn(36, device=self.device)
-        dtype = torch.float64 if self.device != "mps" else torch.float32
+        dtype = highest_precision_float(self.device)
         y = torch.ones(1, dtype=dtype, device=self.device)
 
         self.assertEqual(torch.compile(fn)(x1, y), fn(x1, y))
@@ -4632,7 +4876,7 @@ def fn(a):
     def test_inductor_assert(self):
         @torch.compile(backend="inductor", dynamic=True)
         def fn(a):
-            assert a.shape[0] >= 2 and a.shape[1] >= 4
+            assert a.shape[0] >= 2 and a.shape[1] >= 4  # noqa: S101
             return a.cos()
 
         inp = torch.randn(2, 4, 6)
@@ -4708,7 +4952,7 @@ def fn(x):
         with config.patch({"triton.use_block_ptr": use_block_ptr}):
             self.common(fn, (torch.randn(1, 3, *[10] * dim),))
 
-    @skipIfMPS
+    @xfail_if_mps  # aten::full with zero-sized dim triggers AcceleratorError on MPS
     def test_max_unpool_empty_output(self):
         class Unpool1d(nn.Module):
             def __init__(self):
@@ -5002,8 +5246,8 @@ def fn(x, w):
 
         x = torch.randn([2, 1, 16, 20])
         w = torch.randn([1, 1, 5, 5])
-
-        torch._dynamo.mark_dynamic(x, 0)
+        # X get specialized ! dilation can't be dynamic.
+        torch._dynamo.maybe_mark_dynamic(x, 0)
 
         atol = None
         rtol = None
@@ -5293,15 +5537,15 @@ def run_weights_sharing_model(m, inp):
         threads = []
         compiled_m = torch.compile(model)
         for _ in range(1, numb_instance + 1):
+            ctx = contextvars.copy_context()
             thread = threading.Thread(
-                target=run_weights_sharing_model, args=(compiled_m, inp)
+                target=ctx.run, args=(run_weights_sharing_model, compiled_m, inp)
             )
             threads.append(thread)
             thread.start()
         for thread in threads:
             thread.join()
 
-    @unittest.skipIf(config.is_fbcode(), "fbcode triton error, needs debugging")
     @skip_if_triton_cpu("Flaky on Triton CPU")
     @skip_if_gpu_halide  # https://github.com/halide/Halide/issues/8311
     def test_adaptive_avg_pool2d_low_prec(self):
@@ -5809,10 +6053,12 @@ def fn(a):
 
         # test dtype separately
         out = fn(a)
-        assert out.dtype == torch.bfloat16
+        if out.dtype != torch.bfloat16:
+            raise AssertionError(f"Expected dtype torch.bfloat16, got {out.dtype}")
 
         out = torch.compile(fn)(a)
-        assert out.dtype == torch.bfloat16
+        if out.dtype != torch.bfloat16:
+            raise AssertionError(f"Expected dtype torch.bfloat16, got {out.dtype}")
 
     def test_repeat_interleave(self):
         def fn(x):
@@ -6038,6 +6284,24 @@ def test_layer_norm(self):
         if self.device != "cpu":
             assertGeneratedKernelCountEqual(self, 1)
 
+    def test_layer_norm_rejects_complex_inputs(self):
+        if self.device not in ("cpu", "cuda"):
+            raise unittest.SkipTest("Only validated on CPU/CUDA")
+
+        m = torch.nn.LayerNorm(10).to(self.device)
+        x = torch.randn(1, 1, 10, device=self.device, dtype=torch.complex64)
+
+        with self.assertRaises(RuntimeError):
+            m(x)
+
+        with self.assertRaises(RuntimeError) as compiled_error:
+            torch.compile(m)(x)
+
+        self.assertIn(
+            "native_layer_norm does not support complex inputs",
+            str(compiled_error.exception),
+        )
+
     @torch._functorch.config.patch("donated_buffer", True)
     def test_matmul_layer_norm(self):
         batch_size = 32
@@ -6196,6 +6460,94 @@ def fn():
         actual = compiled()
         self.assertEqual(actual, expected)
 
+    def test_complex_uniform_constant_folding(self):
+        # Fix https://github.com/pytorch/pytorch/issues/174891
+        # view.dtype with mismatched element sizes changes element count,
+        # so constant folding must not treat the result as uniform.
+        def fn(x):
+            mask = torch.ones(2, 1, dtype=torch.complex64, device=self.device)
+            return x + mask
+
+        x = torch.full((2, 2), 1.0, dtype=torch.complex64, device=self.device)
+        expected = fn(x)
+        compiled = torch.compile(fn, backend="inductor")
+        actual = compiled(x)
+        self.assertEqual(actual, expected)
+
+    def test_view_dtype_non_0d_larger_to_smaller_element_size(self):
+        # Non-0-d counterpart of test_view_dtype_0d_smaller_to_larger_element_size.
+        # element_size (8) > itemsize (4): complex64 -> float32.
+        import torch.fx as fx
+        from torch._inductor.fx_passes.joint_graph import UniformValueConstantFolder
+
+        graph = fx.Graph()
+
+        full_node = graph.call_function(
+            torch.ops.aten.full.default,
+            args=([2], 1 + 0j),
+            kwargs={
+                "dtype": torch.complex64,
+                "layout": torch.strided,
+                "device": self.device,
+                "pin_memory": False,
+            },
+        )
+        full_node.meta["val"] = torch.full(
+            [2], 1 + 0j, dtype=torch.complex64, device=self.device
+        )
+
+        view_node = graph.call_function(
+            torch.ops.aten.view.dtype, args=(full_node, torch.float32)
+        )
+        view_node.meta["val"] = torch.full(
+            [2], 1 + 0j, dtype=torch.complex64, device=self.device
+        ).view(torch.float32)
+
+        graph.output(view_node)
+        gm = fx.GraphModule(torch.nn.Module(), graph)
+
+        folder = UniformValueConstantFolder(gm)
+        folder.run()
+
+        self.assertNotIn(view_node, folder.node_replacements)
+
+    def test_view_dtype_non_0d_smaller_to_larger_element_size(self):
+        # Non-0-d counterpart of test_view_dtype_0d_smaller_to_larger_element_size.
+        # element_size (4) < itemsize (8): float32 -> complex64.
+        import torch.fx as fx
+        from torch._inductor.fx_passes.joint_graph import UniformValueConstantFolder
+
+        graph = fx.Graph()
+
+        full_node = graph.call_function(
+            torch.ops.aten.full.default,
+            args=([2], 1.0),
+            kwargs={
+                "dtype": torch.float32,
+                "layout": torch.strided,
+                "device": self.device,
+                "pin_memory": False,
+            },
+        )
+        full_node.meta["val"] = torch.full(
+            [2], 1.0, dtype=torch.float32, device=self.device
+        )
+
+        view_node = graph.call_function(
+            torch.ops.aten.view.dtype, args=(full_node, torch.complex64)
+        )
+        view_node.meta["val"] = torch.full(
+            [2], 1.0, dtype=torch.float32, device=self.device
+        ).view(torch.complex64)
+
+        graph.output(view_node)
+        gm = fx.GraphModule(torch.nn.Module(), graph)
+
+        folder = UniformValueConstantFolder(gm)
+        folder.run()
+
+        self.assertNotIn(view_node, folder.node_replacements)
+
     def test_uniform(self):
         def fn(x):
             return aten.uniform.default(x, 0, 1)
@@ -6222,7 +6574,11 @@ def fn(x, y):
             reference_in_float=False,
         )
 
-    @skipIfMPS
+    @unittest.skipIf(
+        TEST_WITH_ROCM and not torch.cuda.has_magma,
+        "ROCm hipsolver backend does not currently support eig",
+    )
+    @xfail_if_mps_unimplemented  # aten::linalg_eig not implemented for MPS
     def test_linalg_eig_stride_consistency(self):
         def fn(x):
             eigenvals, eigenvecs = torch.linalg.eig(x)
@@ -6308,7 +6664,7 @@ def test_polar(self):
         def fn(dist, angle):
             return torch.polar(dist, angle)
 
-        dtype = torch.float64 if self.device != "mps" else torch.float32
+        dtype = highest_precision_float(self.device)
         inp = (
             torch.tensor([1, 2], dtype=dtype),
             torch.tensor([np.pi / 2, 5 * np.pi / 4], dtype=dtype),
@@ -6477,7 +6833,6 @@ def fn(mask, value):
             )
             self.assertEqual(fn(*inputs), opt_fn(*inputs))
 
-    @xfail_if_mps  # 'NullHandler' object has no attribute 'wrapper_code'
     def test_masked_scatter(self):
         def fn(value, mask, source):
             return torch.masked_scatter(value, mask, source)
@@ -7233,10 +7588,11 @@ def fn(x):
             )
 
     @config.patch(force_disable_caches=True)
-    @skip_if_cpp_wrapper("run_and_get_kernels issue")
     def test_deterministic_codegen(self):
-        if "cpu" in str(self.device) and config.is_fbcode():
-            raise unittest.SkipTest("cpp packaging is wacky in fbcode")
+        if "cpu" in str(self.device) and config.cpp_wrapper:
+            raise unittest.SkipTest(
+                "run_and_get_kernels can't extract kernels from CPU cpp_wrapper code"
+            )
 
         @torch.compile(fullgraph=True)
         def a(x):
@@ -7283,10 +7639,11 @@ def c(x):
         self.assertEqual(coda_c0, coda_c2)
 
     @config.patch(force_disable_caches=True)
-    @skip_if_cpp_wrapper("run_and_get_kernels issue")
     def test_deterministic_codegen_on_graph_break(self):
-        if "cpu" in str(self.device) and config.is_fbcode():
-            raise unittest.SkipTest("cpp packaging is wacky in fbcode")
+        if "cpu" in str(self.device) and config.cpp_wrapper:
+            raise unittest.SkipTest(
+                "run_and_get_kernels can't extract kernels from CPU cpp_wrapper code"
+            )
 
         def a(x):
             return x.cos().sin().softmax(-1)
@@ -7299,8 +7656,15 @@ def b(x):
             return x
 
         x = torch.randn(16, 256, device=self.device)
-        _, (code0, code1) = _run_and_get_stripped_kernels(b, x)
-        self.assertEqual(code0, code1)
+        if config.cpp_wrapper and config.triton.autotune_at_compile_time is not True:
+            # With lazy compile, both graph segments produce identical code
+            # (no unique .cubin paths), so run_and_get_code deduplicates them
+            # and only 1 kernel is returned.
+            _, codes = _run_and_get_stripped_kernels(b, x)
+            self.assertEqual(len(codes), 1)
+        else:
+            _, (code0, code1) = _run_and_get_stripped_kernels(b, x)
+            self.assertEqual(code0, code1)
 
     @config.patch(
         {
@@ -7311,10 +7675,11 @@ def b(x):
             "triton.native_matmul": False,
         }
     )
-    @skip_if_cpp_wrapper("run_and_get_kernels issue")
     def test_deterministic_codegen_with_suffix(self):
-        if "cpu" in str(self.device) and config.is_fbcode():
-            raise unittest.SkipTest("cpp packaging is wacky in fbcode")
+        if "cpu" in str(self.device) and config.cpp_wrapper:
+            raise unittest.SkipTest(
+                "run_and_get_kernels can't extract kernels from CPU cpp_wrapper code"
+            )
 
         @torch.compile(fullgraph=True)
         def a(x):
@@ -8549,6 +8914,26 @@ def fn(a, b, c):
                 (a, b, c),
             )
 
+    @xfail_if_mps  # dtypes mismatch
+    def test_nll_loss_backward_1d_input(self):
+        # 1D input with 1D target used to crash the decomposition because
+        # target.unsqueeze(0) produced a 2D index for a 1D scatter.
+        def fn(a, b, c):
+            return aten.nll_loss_backward(
+                a, b, c, None, 0, 10, torch.tensor(1.0, device=self.device)
+            )
+
+        # 1D target — validation requires target.size(0) == self.size(0)
+        self.common(
+            fn,
+            (torch.randn(1), torch.randn(3), torch.tensor([2, 0, 1])),
+        )
+        # scalar target (the no_batch_dim path in the C++ kernel)
+        self.common(
+            fn,
+            (torch.randn(1), torch.randn(3), torch.tensor(2)),
+        )
+
     def test_isinf(self):
         def fn(x):
             return x.isinf(), x.isnan()
@@ -8848,6 +9233,31 @@ def fn(x, idx):
         self.common(fn, (a, idx))
         assertGeneratedKernelCountEqual(self, 2)
 
+    @skip_if_halide
+    @skip_if_pallas
+    def test_index_put_duplicate_indices_no_accumulate(self):
+        # https://github.com/pytorch/pytorch/issues/174074
+        def fn(x):
+            x[[0, 0]] += x[[0, 0]]
+            return x
+
+        # ncols=2048 forces multiple Triton blocks on GPU, exposing the
+        # inter-block race that small tensors mask.
+        x = torch.randn(1, 2048, dtype=torch.float32)
+        self.common(fn, (x,))
+
+    @skip_if_pallas
+    def test_index_put_duplicate_indices_accumulate(self):
+        # accumulate=True is well-defined with duplicate indices.
+        def fn(x):
+            idx = torch.tensor([0, 0], device=x.device)
+            vals = x[[0, 0]] * 2
+            x.index_put_((idx,), vals, accumulate=True)
+            return x
+
+        x = torch.randn(1, 2048, dtype=torch.float32)
+        self.common(fn, (x,))
+
     def test_adding_tensor_offsets(self):
         @torch.compile(fullgraph=True)
         def fn(x):
@@ -8978,7 +9388,6 @@ def fn(x):
         actual_out = compiled_fn(view)
         self.assertEqual(reference_out.stride(), actual_out.stride())
 
-    @skipIfMPS
     def test_nonzero_static_stride(self):
         from torch._subclasses import FakeTensorMode
 
@@ -9036,7 +9445,6 @@ def fn(x, a, b):
             ],
         )
 
-    @skipIfMPS  # MPS does not support float64
     def test_select_scatter_dtype_consistency(self):
         def fn(x, a):
             return (torch.select_scatter(x, a, 1, 0),)
@@ -9045,6 +9453,8 @@ def fn(x, a):
             torch.int64,
             torch.float64,
         ]:
+            if not self.is_dtype_supported(dtype):
+                continue
             self.common(
                 fn,
                 [
@@ -9069,6 +9479,30 @@ def fn(x, a):
             ],
         )
 
+    @skip_if_gpu_halide  # accuracy issue
+    def test_slice_scatter_backward_with_overlapping_base(self):
+        def fn(x, y):
+            return torch.slice_scatter(x, y, dim=1, start=0, end=6).sum(dim=1)
+
+        torch.manual_seed(0)
+        x = torch.randn(4, 13, 33, device=self.device, requires_grad=True)
+        y = torch.randn(4, 6, 33, device=self.device, requires_grad=True)
+        grad = torch.randn(4, 33, device=self.device)
+
+        x_ref = x.detach().clone().requires_grad_(True)
+        y_ref = y.detach().clone().requires_grad_(True)
+
+        out_ref = fn(x_ref, y_ref)
+        out_ref.backward(grad)
+
+        compiled_fn = torch.compile(fn, backend="inductor", fullgraph=True)
+        out = compiled_fn(x, y)
+        out.backward(grad)
+
+        self.assertEqual(out, out_ref)
+        self.assertEqual(x.grad, x_ref.grad)
+        self.assertEqual(y.grad, y_ref.grad)
+
     def test_slice_scatter2(self):
         def fn(a, b):
             return aten.slice_scatter(a, b, 0, 0, 9223372036854775807)
@@ -9157,7 +9591,6 @@ def forward(self, x, start_pos):
         else:
             assertGeneratedKernelCountEqual(self, 1)
 
-    @skipIfMPS
     def test_slice_scatter_dtype_consistency(self):
         # Test dtype consistency of slice_scatter
         def fn(x, y):
@@ -9167,6 +9600,8 @@ def fn(x, y):
             torch.int64,
             torch.float64,
         ]:
+            if not self.is_dtype_supported(dtype):
+                continue
             self.common(
                 fn,
                 [
@@ -9636,7 +10071,8 @@ def _(x):
             tags=[torch._C.Tag.needs_contiguous_strides],
         )
         def second_op(x: torch.Tensor) -> torch.Tensor:
-            assert x.is_contiguous()
+            if not x.is_contiguous():
+                raise AssertionError("Expected x to be contiguous")
             return torch.ones(2, 2)
 
         @second_op.register_fake
@@ -9691,10 +10127,6 @@ def f(a, b1, b2, c, d):
             f(*args)
             self.assertEqual(cloned_args, args)
 
-    @skip_if_cpp_wrapper(
-        "Without major redesign, cpp_wrapper will not support custom ops that are "
-        "defined in Python."
-    )
     @config.patch(implicit_fallbacks=True)
     def test_fallback_mutable_op_list_tensor(self):
         @torch.library.custom_op(
@@ -9795,10 +10227,6 @@ def f(a):
             self.assertEqual(cloned_args, args)
 
     @config.patch(implicit_fallbacks=True)
-    @skip_if_cpp_wrapper(
-        "Without major redesign, cpp_wrapper will not support custom ops that are "
-        "defined in Python."
-    )
     def test_fallback_mutable_op_list(self):
         with torch.library._scoped_library("mylib", "FRAGMENT") as m:
 
@@ -9825,7 +10253,8 @@ def f(a, b):
         @torch.library.custom_op("mylib::sin_out", mutates_args={"outs"})
         def sin_out(x: torch.Tensor, outs: list[torch.Tensor]) -> None:
             x_np = x.numpy()
-            assert len(outs) == 2
+            if len(outs) != 2:
+                raise AssertionError(f"Expected len(outs) == 2, got {len(outs)}")
             out_np0 = out[0].numpy()
             out_np1 = out[1].numpy()
             np.sin(x_np, out=out_np0)
@@ -10042,7 +10471,6 @@ def fn(x):
         self.common(fn, [torch.zeros([20, 20])], exact_stride=True)
 
     @config.patch(fallback_random=True)
-    @xfail_if_mps  # 100% are not close
     def test_like_rands_sliced(self):
         def fn(x):
             return (
@@ -10053,6 +10481,31 @@ def fn(x):
 
         self.common(fn, (torch.zeros([3, 4])[:, ::2].permute(1, 0),), exact_stride=True)
 
+    @config.patch(fallback_random=True)
+    def test_like_randn_non_contiguous_rng_consistency(self):
+        x = torch.zeros((3, 6), device=self.device).t()
+        self.assertFalse(x.is_contiguous())
+
+        randn_like_compiled = torch.compile(torch.randn_like, backend="inductor")
+
+        torch.manual_seed(0)
+        eager_out = torch.randn_like(x)
+        torch.manual_seed(0)
+        compiled_out = randn_like_compiled(x)
+
+        self.assertEqual(eager_out, compiled_out)
+
+    @config.patch(fallback_random=False)
+    def test_fast_like_rands_decomps_use_non_eager_path(self):
+        x = torch.zeros(3, 6, device=self.device)
+
+        randn_like_compiled = torch.compile(torch.randn_like, backend="inductor")
+        torch.manual_seed(0)
+        eager_randn_out = torch.randn_like(x)
+        torch.manual_seed(0)
+        compiled_randn_out = randn_like_compiled(x)
+        self.assertNotEqual(eager_randn_out, compiled_randn_out)
+
     @config.patch(check_stack_no_cycles_TESTING_ONLY=True)
     def test_check_stack_no_cycles(self):
         if config.cpp_wrapper and self.device != "cpu":
@@ -10074,7 +10527,8 @@ def fn(x):
     def test_like_rands2(self):
         # rand_like with kwargs `device` of str type
         d = self.device
-        assert isinstance(d, str)
+        if not isinstance(d, str):
+            raise AssertionError(f"Expected d to be str, got {type(d)}")
 
         @torch.compile
         def fn(x):
@@ -10210,11 +10664,12 @@ def fn(a, b, c):
                 indices,
             ],
         )
-        assertGeneratedKernelCountEqual(self, 1)
+        # Note: Kernel count varies by backend (CUDA ~3, ROCm ~2) due to fusion.
+        # Correctness is validated by self.common() above.
+        self.assertGreater(torch._inductor.metrics.generated_kernel_count, 0)
 
-    @expectedFailureXPU
     def test_max_pool2d_with_indices_backward5(self):
-        # Window size is too big. Should fallback
+        # Large window size - decomposition handles via scatter_add
         def fn(a, b, c):
             return aten.max_pool2d_with_indices_backward(
                 a, b, [13, 13], [1, 1], [2, 2], [1, 1], False, c
@@ -10238,11 +10693,15 @@ def fn(a, b, c):
                 indices,
             ],
         )
-        assertGeneratedKernelCountEqual(self, 0)
+        # Note: Kernel count varies by backend (CUDA ~3, ROCm ~2) due to fusion.
+        # Correctness is validated by self.common() above.
+        # MPS: decomposition falls back to native kernel, so no inductor kernels generated
+        if self.device != "mps":
+            self.assertGreater(torch._inductor.metrics.generated_kernel_count, 0)
 
     # From https://github.com/pytorch/pytorch/issues/93384
     def test_max_pool2d_with_indices_backward6(self):
-        # dilation is not 1. Should fallback
+        # dilation != 1 - decomposition handles all dilation cases
         def fn(a, b, c):
             return aten.max_pool2d_with_indices_backward(
                 a, b, [3, 2], [2, 1], [1, 1], [1, 2], False, c
@@ -10266,7 +10725,11 @@ def fn(a, b, c):
                 indices,
             ],
         )
-        assertGeneratedKernelCountEqual(self, 0)
+        # Note: Kernel count varies by backend (CUDA ~3, ROCm ~2) due to fusion.
+        # Correctness is validated by self.common() above.
+        # MPS: decomposition falls back to native kernel, so no inductor kernels generated
+        if self.device != "mps":
+            self.assertGreater(torch._inductor.metrics.generated_kernel_count, 0)
 
     def test_issue102546(self):
         def fn(x):
@@ -10490,6 +10953,10 @@ def fn(a):
 
     @xfail_if_mps
     def test_dropout2(self):
+        if is_mps_backend(self.device) and torch._inductor.config.align_random_eager:
+            raise AssertionError(
+                "MPS + align_random_eager: will pass but force failure for xfail_if_mps"
+            )
         n = 100000
         weight = torch.ones(
             n, device=self.device, dtype=torch.float32, requires_grad=True
@@ -10527,10 +10994,13 @@ def check(r, g):
         torch.manual_seed(1234)
         weight.grad.zero_()
         r2, (fw_code, bw_code) = run_fw_bw_and_get_code(lambda: run(ones))
-        if is_halide_backend(self.device):
+        if (
+            is_halide_backend(self.device)
+            and not torch._inductor.config.align_random_eager
+        ):
             self.assertEqual(fw_code.count("halide_helpers.rand"), 1)
             self.assertEqual(bw_code.count("halide_helpers.rand"), 0)
-        elif self.device == GPU_TYPE:
+        elif self.device == GPU_TYPE and not torch._inductor.config.align_random_eager:
             self.assertEqual(fw_code.count("tl.rand"), 1)
             self.assertEqual(bw_code.count("tl.rand"), 0)
         g2 = weight.grad.clone()
@@ -10550,6 +11020,10 @@ def check(r, g):
     @xfail_if_mps
     @config.patch(search_autotune_cache=False)
     def test_dropout3(self):
+        if is_mps_backend(self.device) and torch._inductor.config.align_random_eager:
+            raise AssertionError(
+                "MPS + align_random_eager + dynamic shapes: will pass but force failure for xfail_if_mps"
+            )
         m = torch.nn.Sequential(
             torch.nn.Linear(32, 32, bias=False),
             torch.nn.Dropout(),
@@ -10567,10 +11041,13 @@ def run(x):
             lambda: run(torch.randn([8, 32], device=self.device))
         )
 
-        if is_halide_backend(self.device):
+        if (
+            is_halide_backend(self.device)
+            and not torch._inductor.config.align_random_eager
+        ):
             self.assertEqual(fw_code.count("halide_helpers.rand"), 2)
             self.assertEqual(bw_code.count("halide_helpers.rand"), 0)
-        elif self.device == GPU_TYPE:
+        elif self.device == GPU_TYPE and not torch._inductor.config.align_random_eager:
             # the load_seed_offset arg can be 1 or non-1; depending on whether
             # the triton signature specializes on 1 vs non-1, you might get 1
             # or 2 kernels. In newer versions of triton, there's no specialization
@@ -10644,7 +11121,6 @@ def fn(x):
             ],
         )
 
-    @skipIfXpu(msg="Incorrect XPU reference")
     def test_argmax_argmin2(self):
         def fn(x):
             return (
@@ -10656,7 +11132,6 @@ def fn(x):
 
         self.common(fn, (torch.randn([144, 144]),))
 
-    @skipIfXpu(msg="Incorrect XPU reference")
     def test_argmax_argmin_with_duplicates(self):
         def fn(x):
             return (
@@ -10678,7 +11153,6 @@ def fn(x):
         t1 = torch.randint(8, size=(1028, 1028))
         self.common(fn, (t1,))
 
-    @skipIfXpu(msg="# Incorrect XPU reference ")
     @xfail_if_mps  # eager nan is wrong, see https://github.com/pytorch/pytorch/issues/130295
     @skip_if_halide  # nan behavior
     def test_argmax_argmin_with_nan(self):
@@ -10708,6 +11182,30 @@ def fn(x):
         t1[:, 100] = float("nan")
         self.common(fn, (t1,))
 
+    @requires_cuda
+    def test_max_min_bool(self):
+        # Regression test for https://github.com/pytorch/pytorch/issues/174069
+        # max/min on boolean tensors returned incorrect indices in Triton
+        def fn(x):
+            return (
+                x.max(0),
+                x.min(0),
+                x.max(1),
+                x.min(1),
+            )
+
+        # Unrolled reduction
+        t1 = torch.randint(2, size=(6, 6), dtype=torch.bool)
+        self.common(fn, (t1,))
+
+        # Persistent reduction
+        t1 = torch.randint(2, size=(32, 32), dtype=torch.bool)
+        self.common(fn, (t1,))
+
+        # Non-persistent reduction
+        t1 = torch.randint(2, size=(1028, 1028), dtype=torch.bool)
+        self.common(fn, (t1,))
+
     def test_conv_backward(self):
         def fn(rank4_inps, rank3_inps, rank5_inps):
             out1 = aten.convolution_backward(
@@ -10779,7 +11277,6 @@ def shrink_rank(x, rank):
                 [rank4_inps, rank3_inps, rank5_inps],
             )
 
-    @skipIfXpu(msg="Incorrect XPU reference")
     def test_argmax_argmin3(self):
         def fn(x):
             return (
@@ -11496,7 +11993,10 @@ def custom_pass(g: torch.fx.Graph) -> None:
             Applies `lambda x: x.t().contiguous().t()` to the output.
             """
             output_node = g.find_nodes(op="output")[0]
-            assert len(output_node.args) == 1
+            if len(output_node.args) != 1:
+                raise AssertionError(
+                    f"Expected len(output_node.args) == 1, got {len(output_node.args)}"
+                )
             output = output_node.args[0][0]
 
             with g.inserting_before(output_node):
@@ -11721,7 +12221,6 @@ def fn(n, a):
         self.assertEqual(x[0], -1)
         self.assertEqual(cnts.frame_count, frame_count + 1)
 
-    @config.patch({"triton.autotune_at_compile_time": False})
     @config.patch(profiler_mark_wrapper_call=True)
     def test_profiler_mark_wrapper_call(self):
         from torch.profiler import profile
@@ -11734,9 +12233,12 @@ def fn(a, b):
         b = torch.rand((100,), device=self.device)
         with profile() as prof:
             fn(a, b)
-        assert any(
+        if not any(
             "inductor_wrapper_call" in e.name for e in prof.profiler.function_events
-        )
+        ):
+            raise AssertionError(
+                "Expected 'inductor_wrapper_call' in profiler function events"
+            )
 
     def test_insignificant_strides(self):
         def f(x):
@@ -12395,6 +12897,14 @@ def _cases_resize_common():
                     x_strided = x[::2].reshape(25, 2).transpose(0, 1)
                     yield x_strided, y_size, memory_format
 
+    def test_resize_overlapping_strides(self):
+        # Resize on a stride-0 view should read logical elements, not raw storage.
+        def fn(x):
+            view = torch.as_strided(x, (100,), (0,))
+            return torch.ops.aten.resize(view, (50,))
+
+        self.common(fn, (torch.ones(10),))
+
     def test_resize(self):
         def fn(x, size, memory_format):
             # NOTE: Tensor.resize() =/= aten::resize()
@@ -12513,7 +13023,6 @@ def fn(q, k, v):
         )
 
     @xfail_if_mps_unimplemented
-    @expectedFailureXPU
     @unittest.skipIf(
         not PLATFORM_SUPPORTS_MEM_EFF_ATTENTION, "Some archs don't support mem eff SDPA"
     )
@@ -12742,6 +13251,24 @@ def fn(input, boundaries):
         # because bucketize is computationally expensive.
         FileCheck().check("def triton").check("def triton").run(code[0])
 
+    @requires_gpu()
+    @skip_if_gpu_halide
+    @skip_if_not_triton
+    def test_bucketize_nan_consistency(self):
+        # Regression test for https://github.com/pytorch/pytorch/issues/173133.
+        def fn(x, boundaries, right):
+            return torch.bucketize(torch.rsqrt(x), boundaries, right=right)
+
+        x = torch.tensor([-1.0], device=self.device)
+        boundaries = torch.tensor([0.2, 0.5, 0.8], device=self.device)
+        expected = torch.tensor(
+            [boundaries.numel()], device=self.device, dtype=torch.int64
+        )
+
+        for right in (False, True):
+            self.assertEqual(fn(x, boundaries, right), expected)
+            self.common(fn, (x, boundaries, right), check_lowp=False)
+
     @requires_gpu()
     @config.patch(assume_aligned_inputs=False)
     def test_config_option_dont_assume_alignment(self):
@@ -12912,7 +13439,6 @@ def fn(x):
 
     @requires_gpu()
     @skip_if_not_triton
-    @skip_if_cpp_wrapper("skip cpp_wrapper tests")
     @config.patch(implicit_fallbacks=True)
     def test_generated_code_has_size_stride_assert(self):
         def foo(x):
@@ -12949,7 +13475,6 @@ def fn(x):
 
     @requires_gpu()
     @skip_if_not_triton
-    @skip_if_cpp_wrapper("skip cpp_wrapper tests")
     @config.patch(implicit_fallbacks=True)
     def test_generated_code_has_alignment_assert(self):
         def foo(x):
@@ -13002,6 +13527,108 @@ def test_assert_alignment_op_name_fail(self):
         with self.assertRaisesRegex(AssertionError, "torch.ops.dummy.op_name"):
             assert_alignment(tensor, 0, "torch.ops.dummy.op_name")
 
+    @requires_gpu()
+    @skip_if_not_triton
+    @unittest.skipIf(
+        config.cpp_wrapper,
+        "Inductor does not generate size/stride asserts for cpp_wrapper",
+    )
+    def test_input_asserts_deferred_to_first_use(self):
+        def fn(x, y, z):
+            a = torch.mm(x, y)
+            b = torch.mm(a, z)
+            return b
+
+        x = torch.randn(16, 32, device=self.device)
+        y = torch.randn(32, 64, device=self.device)
+        z = torch.randn(64, 8, device=self.device)
+
+        _, code = run_and_get_code(torch.compile(fn), x, y, z)
+        # z's assert should appear after the first mm, not at the top
+        # with all the other asserts
+        FileCheck().check("def call").check_count(
+            "assert_size_stride", 2, exactly=True
+        ).check("extern_kernels.mm(").run(code[0])
+
+    @requires_gpu()
+    @skip_if_not_triton
+    @unittest.skipIf(
+        config.cpp_wrapper,
+        "Deferred alignment copies are not generated for cpp_wrapper",
+    )
+    def test_alignment_copy_deferred_to_first_use(self):
+        def fn(x, y, z):
+            a = torch.mm(x, y)
+            b = torch.mm(a, z)
+            return b
+
+        x = torch.randn(16, 32, device=self.device)
+        y = torch.randn(32, 64, device=self.device)
+
+        z = torch.randn(64, 8, device=self.device)
+
+        _, code = run_and_get_code(torch.compile(fn), x, y, z)
+        # z's alignment check should appear between the two mm calls:
+        # first mm (uses x, y) -> alignment clone (for z) -> second mm (uses z)
+        FileCheck().check("extern_kernels.mm(").check("copy_misaligned").check(
+            "extern_kernels.mm("
+        ).run(code[0])
+
+    @requires_gpu()
+    @skip_if_not_triton
+    @torch._inductor.config.patch(cpp_wrapper=True)
+    def test_alignment_copy_not_emitted_for_cpp_wrapper(self):
+        def fn(x, y):
+            return torch.mm(x, y)
+
+        x = torch.randn(16, 32, device=self.device)
+        y = torch.randn(32, 64, device=self.device)
+
+        _, code = run_and_get_code(torch.compile(fn), x, y)
+        # cpp_wrapper should NOT contain Python-syntax alignment copies
+        self.assertNotIn("copy_misaligned", code[0])
+
+    def test_copy_misaligned_returns_same_tensor_when_aligned(self):
+        import weakref
+
+        from torch._C._dynamo.guards import copy_misaligned
+
+        x = torch.randn(32, 32, device=self.device)
+        ref = weakref.ref(x)
+        result = copy_misaligned(x)
+        self.assertIs(result, x)
+        del x, result
+        self.assertIsNone(ref(), "aligned tensor should be freed")
+
+    def test_copy_misaligned_clones_when_misaligned(self):
+        import weakref
+
+        from torch._C._dynamo.guards import copy_misaligned
+
+        big = torch.randn(32 * 32 + 1, device=self.device)
+        x = big[1:].reshape(32, 32)
+        self.assertNotEqual(x.data_ptr() % 16, 0)
+        ref_orig = weakref.ref(x)
+        result = copy_misaligned(x)
+        self.assertIsNot(result, x)
+        self.assertEqual(result.data_ptr() % 16, 0)
+        self.assertEqual(result, x)
+        self.assertEqual(result.size(), x.size())
+        self.assertEqual(result.stride(), x.stride())
+        ref_clone = weakref.ref(result)
+        del x, result
+        self.assertIsNone(ref_clone(), "cloned tensor should be freed")
+        # orig kept alive by big's storage, that's fine
+        del big
+        self.assertIsNone(ref_orig(), "original tensor should be freed")
+
+    def test_copy_misaligned_empty_tensor(self):
+        from torch._C._dynamo.guards import copy_misaligned
+
+        x = torch.randn(0, device=self.device)
+        result = copy_misaligned(x)
+        self.assertEqual(result.size(), x.size())
+
     @torch._dynamo.config.patch(capture_dynamic_output_shape_ops=True)
     @torch._inductor.config.patch(implicit_fallbacks=True)
     def test_custom_op_unbacked_symints(self):
@@ -13073,10 +13700,6 @@ def fn(x):
 
     @requires_gpu()
     @config.patch(implicit_fallbacks=True)
-    @skip_if_cpp_wrapper(
-        "Without major redesign, cpp_wrapper will not support custom ops that are "
-        "defined in Python."
-    )
     @tf32_on_and_off(0.005)
     def test_mutable_custom_op_fixed_layout2(self):
         with torch.library._scoped_library("mylib", "DEF") as lib:
@@ -13096,7 +13719,10 @@ def _(x, is_compiling):
                 if is_compiling:
                     bar_strides.append(x.stride())
                 result = x.clone()
-                assert x.stride() == result.stride()
+                if x.stride() != result.stride():
+                    raise AssertionError(
+                        f"Expected stride {x.stride()}, got {result.stride()}"
+                    )
                 return result
 
             @torch.library.impl(lib, "bar", "Meta")
@@ -13132,10 +13758,6 @@ def fn(x):
                 self.assertNotEqual(bar_strides[0], expected_stride)
 
     @config.patch(implicit_fallbacks=True)
-    @skip_if_cpp_wrapper(
-        "Without major redesign, cpp_wrapper will not support custom ops that are "
-        "defined in Python."
-    )
     def test_mutable_custom_op_fixed_layout(self):
         with torch.library._scoped_library("mylib", "DEF") as lib:
             lib.define(
@@ -13211,10 +13833,6 @@ def baz_meta(x):
             net = torch.compile(model)
             out = net(input_t)
 
-    @skip_if_cpp_wrapper(
-        "Without major redesign, cpp_wrapper will not support custom ops that are "
-        "defined in Python."
-    )
     @config.patch(implicit_fallbacks=True)
     def test_custom_op_default_layout_constraint(self):
         with torch.library._scoped_library("mylib", "DEF") as lib:
@@ -13367,7 +13985,7 @@ def fn(tensor, index, source):
             return out
 
         device = "cpu"
-        dtype = torch.double if self.device != "mps" else torch.float32
+        dtype = highest_precision_float(self.device)
         tensor = torch.rand((1,), dtype=dtype, device=device)
         index = torch.tensor([0], dtype=torch.long, device=device)
         source = torch.rand((1,), dtype=dtype, device=device)
@@ -14189,14 +14807,15 @@ def f(x, shape):
         else:
             FileCheck().check("torch.ops.aten._fused_rms_norm.default(").run(code[0])
 
+        x = torch.randn(2, 3, device=self.device, requires_grad=True)
+        _, codes = run_fw_bw_and_get_code(lambda: f(x, [2, 3]))
+        self.assertEqual(len(codes), 2)
         if config.cpp_wrapper:
-            # arg type List[int] is not yet supported by custom_op_wrapper
-            pass
+            FileCheck().check(
+                "AOTI_TORCH_ERROR_CODE_CHECK(aoti_torch_cuda__fused_rms_norm("
+            ).run(codes[0])
         else:
-            x = torch.randn(2, 3, device=self.device, requires_grad=True)
-            _, codes = run_fw_bw_and_get_code(lambda: f(x, [2, 3]))
-            self.assertEqual(len(codes), 2)
-            FileCheck().check("torch.ops.aten._fused_rms_norm.default(").run(code[0])
+            FileCheck().check("torch.ops.aten._fused_rms_norm.default(").run(codes[0])
 
     def test_lite_regional_compile_flex_attention(self):
         if self.device != GPU_TYPE or self.device == "mps":
@@ -14365,11 +14984,11 @@ def forward(self, c):
                 r"raise RuntimeError\('u.* >= 0'\)"
             ).run(code[0])
 
-    @lowering.force_fallback(aten.sort.default)
     @unittest.skipIf(
         config.cpp_wrapper,
         "Inductor does not generate size/stride asserts for cpp_wrapper",
     )
+    @lowering.force_fallback(aten.sort.default)
     def test_size_asserts_for_multi_output_fallback(self):
         @torch.compile
         def f(x):
@@ -14654,9 +15273,18 @@ def __torch_dispatch__(self, func, types, args=(), kwargs=None):
                         gc.collect()
                         if func is aten.mm.out:
                             matmul_seen = True
-                            assert len(inps) == 0
-                            assert inp_refs[0]() is None
-                            assert inp_refs[1]() is None
+                            if len(inps) != 0:
+                                raise AssertionError(
+                                    f"Expected len(inps) == 0, got {len(inps)}"
+                                )
+                            if inp_refs[0]() is not None:
+                                raise AssertionError(
+                                    "Expected inp_refs[0]() to be None"
+                                )
+                            if inp_refs[1]() is not None:
+                                raise AssertionError(
+                                    "Expected inp_refs[1]() to be None"
+                                )
 
                         return func(*args, **kwargs)
 
@@ -14676,7 +15304,8 @@ def __torch_dispatch__(self, func, types, args=(), kwargs=None):
                 with TestRefMode():
                     fn_compiled(inps)
 
-                assert len(inps) == 0
+                if len(inps) != 0:
+                    raise AssertionError(f"Expected len(inps) == 0, got {len(inps)}")
 
     @torch._inductor.config.patch("graph_partition", True)
     def test_graph_partition_pad_dynamic(self):
@@ -14887,6 +15516,36 @@ def forward(self, x):
         inputs = (torch.randn(4, device=self.device),)
         self.common(Model(), inputs)
 
+    @requires_gpu_and_triton
+    def test_triton_kernel_bool_tensor_arg(self):
+        if self.device != GPU_TYPE or self.device == "mps":
+            raise unittest.SkipTest("requires GPU")
+
+        from torch.testing._internal.triton_utils import (
+            masked_add_kernel_with_bool_tensor,
+        )
+
+        class Model(torch.nn.Module):
+            def forward(self, x, y, mask):
+                out = torch.zeros_like(x)
+                n = x.numel()
+                masked_add_kernel_with_bool_tensor[(n,)](
+                    in_ptr0=x,
+                    in_ptr1=y,
+                    mask_ptr=mask,
+                    out_ptr=out,
+                    n_elements=n,
+                    BLOCK_SIZE=1024,
+                )
+                return out
+
+        n = 128
+        x = torch.randn(n, device=self.device)
+        y = torch.randn(n, device=self.device)
+        mask = torch.arange(n, device=self.device) < n // 2
+        inputs = (x, y, mask)
+        self.common(Model(), inputs)
+
     @skipIfXpu(
         msg="Profile not enabled on XPU CI, "
         "https://github.com/intel/torch-xpu-ops/issues/2334"
@@ -14953,6 +15612,10 @@ def fn(inp, repeats, output_size):
         args = (inp, repeats, output_size)
         self.assertEqual(fn(*args), torch.compile(fn)(*args))
 
+    @unittest.skipIf(
+        config.triton.autotune_at_compile_time is True,
+        "autotune_at_compile_time doesn't work for test with indexing",
+    )
     @parametrize("dtype", [torch.int32, torch.int64])
     @parametrize("nd", [1, 2])
     def test_repeat_interleave_Tensor_decomp(self, dtype, nd):
@@ -14970,9 +15633,7 @@ def f(input, repeats):
         output, (code,) = run_and_get_code(f_compiled, input, repeat)
         reference = f(input, repeat)
         self.assertEqual(output, reference)
-        # we don't lower when the cpp_wrapper is used because it cannot generate
-        # proper examples during autotune
-        can_lower = (not config.cpp_wrapper) and (input.device.type != "mps")
+        can_lower = input.device.type != "mps"
         has_lowered = not re.search(r"repeat_interleave.Tensor", code)
         self.assertEqual(has_lowered, can_lower)
 
@@ -15162,7 +15823,10 @@ def fn(x):
         self.assertTrue("'NaN or Inf found'" in code)
         torch.testing.assert_close(out, fn(x))
 
-    @skip_if_cpp_wrapper("skip cpp wrapper")
+    @unittest.skipIf(
+        config.triton.autotune_at_compile_time is True,
+        "autotune_at_compile_time doesn't work for test with indexing",
+    )
     @requires_gpu_and_triton
     def test_repeat_interleave_decomposition_has_clamp(self):
         repeat = torch.ones(2560, dtype=torch.int64, device=GPU_TYPE)
@@ -15253,7 +15917,6 @@ def foo(x):
             torch.compile(foo, fullgraph=True)(torch.ones(3, 3))
             self.assertTrue(len(log2), 1)
 
-    @skipIfMPS  # Accuracy issue on MPS
     def test_weight_norm_conv2d(self):
         """
         Verify fix for https://github.com/pytorch/pytorch/issues/165749
@@ -15273,7 +15936,7 @@ def test_weight_norm_conv2d(self):
 
         self.assertTrue(same((ref, ref_grad), (act, act_grad), tol=1e-3))
 
-    @skipIfMPS
+    @xfail_if_mps  # MPS codegen does not emit ReductionHint.OUTER
     def test_inner_reduction_detection(self):
         if self.device == "cpu":
             self.skipTest("Skip for CPU device")
@@ -15289,8 +15952,7 @@ def f(x):
         self.assertFalse("ReductionHint.INNER" in code)
 
     @skip_if_halide
-    @requires_cuda_and_triton
-    @skip_if_cpp_wrapper("skip cpp wrapper")
+    @requires_gpu_and_triton
     def test_triton_argmin_argmax_transpose_logical_index(self):
         def fn(x):
             x.tan_()
@@ -15327,7 +15989,7 @@ def fn(x):
         self.common(fn, (torch.randn(6, 4, device=GPU_TYPE).t().contiguous().t(),))
 
     @skip_if_halide
-    @requires_cuda_and_triton
+    @requires_gpu_and_triton
     def test_unbacked_float_item(self):
         def fn(x, max_val):
             return torch.clamp(x, 0, max_val.item())
@@ -15367,7 +16029,7 @@ def fn(a, b, c, d):
         if self.device.lower() == "cuda":
             self.assertEqual(torch._inductor.metrics.generated_kernel_count, 2)
 
-    @requires_cuda_and_triton
+    @requires_gpu_and_triton
     @config.patch(combo_kernels=True)
     @torch._dynamo.config.patch(assume_static_by_default=False)
     def test_combo_kernel_store_mask(self):
@@ -15562,6 +16224,7 @@ def fn(a, b, c):
             "coordinate_descent_tuning": True,
         }
     )
+    @xfailIf(not IS_BIG_GPU)  # templates require big gpu
     def test_pdl_template_and_delay(self):
         def fn(a, b):
             a = (a / (a**2).sum(-1, keepdim=True)) ** 2  # first kernel
@@ -15665,6 +16328,30 @@ def run_session(x_param, y_param, size, device):
             out2 = run_session(100, 16, 64, self.device)
             self.assertEqual(out2.device.type, self.device)
 
+    def test_index_reduce_on_view_input(self):
+        # Regression test for https://github.com/pytorch/pytorch/issues/144846
+        def fn(x, index, source):
+            return x.index_reduce_(2, index, source, "mean", include_self=False)
+
+        x_base = torch.randn(4, 34, 64, device=self.device)
+        index = torch.randint(0, 34, (64,), device=self.device)
+        source = torch.randn(4, 32, 64, device=self.device)
+
+        expected = fn(x_base.clone()[:, 2:, :], index, source)
+        result = torch.compile(fn)(x_base.clone()[:, 2:, :], index, source)
+        self.assertEqual(result, expected)
+
+    def test_pad_after_gelu(self):
+        # Regression test for Voxtral compilation on MPS
+        for dtype in [torch.float32, torch.float16, torch.bfloat16]:
+            if not self.is_dtype_supported(dtype):
+                continue
+            self.common(
+                lambda x: torch.nn.functional.pad(torch.nn.functional.gelu(x), [1, 0]),
+                (torch.randn(8, 16, dtype=dtype, device=self.device),),
+                check_lowp=False,
+            )
+
     # end of class CommonTemplate - add new tests here
 
 
@@ -15849,7 +16536,10 @@ def fn(a: torch.Tensor) -> torch.Tensor:
             if config.triton.multi_kernel:
                 self.assertEqual(len(kernels), 4)
                 expected_divisible[2] = expected_divisible.pop(1)
-            elif config.triton.cooperative_reductions:
+            elif (
+                config.triton.cooperative_reductions
+                or config.triton.force_cooperative_reductions
+            ):
                 self.assertEqual(len(kernels), 1)
                 expected_divisible = {
                     # one kernel, with extra workspace/semaphore args
@@ -16266,7 +16956,17 @@ def apply_rotary_pos_emb(
 
             # make sure that rope is fused into 1 kernel
             code = run_and_get_triton_code(compiled_fn, q, k, cos, sin, pos_ids)
-            self.assertEqual(code.count(".run("), 1)
+            if config.cpp_wrapper:
+                # Count actual kernel invocations, not function definitions
+                # Kernel invocations end with ');' and don't start with 'static'
+                import re
+
+                kernel_calls = len(
+                    re.findall(r"^\s+call_triton_.*\);$", code, re.MULTILINE)
+                )
+                self.assertEqual(kernel_calls, 1)
+            else:
+                self.assertEqual(code.count(".run("), 1)
 
         def test_numpy_autograd(self):
             def my_torch(x):
@@ -16433,6 +17133,45 @@ def f(a, b):
     tmp1 = tl.load(in_ptr1 + (x0), xmask, cache_modifier='.cg')""",
             )
 
+        @config.patch("triton.skip_l1_cache", True)
+        def test_skip_l1_cache_buf_read_counts_guard(self):
+            from torch._inductor.codegen import simd_kernel_features
+
+            class M(nn.Module):
+                def __init__(self, n: int):
+                    super().__init__()
+                    # nn.Parameter ensures we exercise the primals_* path.
+                    # We move the module to GPU later, so we don't need to
+                    # specify the device here.
+                    self.weight = nn.Parameter(torch.randn(n))
+
+                def forward(self, x: torch.Tensor) -> torch.Tensor:
+                    return x * self.weight
+
+            N = 512
+            m = torch.compile(M(N).to(device=GPU_TYPE))
+            x = torch.randn(N, device=GPU_TYPE)
+
+            # Monkeypatch SIMDKernelFeatures.buffer_read_counts to
+            # simulate missing entries for all buffers (including primals_*).
+            orig = simd_kernel_features.SIMDKernelFeatures.buffer_read_counts
+
+            def fake_buffer_read_counts(self_inner):
+                return {}
+
+            simd_kernel_features.SIMDKernelFeatures.buffer_read_counts = (
+                fake_buffer_read_counts
+            )
+            try:
+                # This used to raise InductorError(KeyError('primals_X'))
+                # when load() unconditionally indexed buffer_read_counts[name].
+                # With the fix, we should successfully generate Triton code.
+                code = run_and_get_triton_code(m, x)
+                lines = [line for line in code.split("\n") if "tl.load" in line]
+                self.assertTrue(lines)
+            finally:
+                simd_kernel_features.SIMDKernelFeatures.buffer_read_counts = orig
+
         @config.patch("triton.use_block_ptr", True)
         @config.patch("triton.coalesce_tiling_analysis", False)
         def test_evict_last_non_coalesced_loads_block_ptr(self):
@@ -16528,6 +17267,41 @@ def test_has_constant_mask(self, block_multiple, ynumel_exceed_ygrid_size):
                 self.assertTrue("ymask = yindex < ynumel" in code)
                 self.assertTrue("xmask = xindex < xnumel" in code)
 
+        @parametrize(
+            "rnumel",
+            [16, 32],
+        )
+        @config.patch("triton.persistent_reductions", True)
+        def test_has_constant_mask_small_persistent_reduction(self, rnumel):
+            from torch._inductor.runtime.hints import get_warp_size
+
+            def fn(x):
+                return x.sum(dim=-1)
+
+            x = torch.randn(1024, rnumel, device=GPU_TYPE)
+            opt_fn = torch.compile(fn)
+            code = run_and_get_triton_code(opt_fn, x)
+
+            device = torch.device(GPU_TYPE, 0)
+            warp_size = get_warp_size(device)
+
+            rblock = 1
+            while rblock < rnumel:
+                rblock *= 2
+
+            if rblock < warp_size:
+                self.assertTrue(
+                    "r0_index < r0_numel" in code or "rindex < rnumel" in code,
+                    f"Expected dynamic reduction mask for RBLOCK={rblock} < warp_size={warp_size}",
+                )
+            else:
+                self.assertTrue(
+                    "r0_mask = tl.full" in code or "rmask = tl.full" in code,
+                    f"Expected constant reduction mask for RBLOCK={rblock} >= warp_size={warp_size}",
+                )
+
+            self.assertEqual(fn(x), opt_fn(x))
+
         @config.patch("triton.native_matmul", False)
         def test_kernel_names_descriptive(self):
             @torch.compile(backend="inductor")
@@ -16590,6 +17364,9 @@ def test_funcs(func_and_kernel):
                 func_and_kernel_torch
             )
 
+        @skip_if_cpp_wrapper(
+            "profile_bandwidth is not supported with cpp_wrapper",
+        )
         @patch.object(config, "profile_bandwidth", True)
         def test_bandwidth_profiler(self):
             @torch.compile(backend="inductor")
@@ -16851,8 +17628,15 @@ def fn(inp, weight):
 
         @torch._functorch.config.patch("donated_buffer", True)
         # The inplace updating does not happen after we fused the
-        # layernorm backward
-        @torch._inductor.config.patch("triton.mix_order_reduction", False)
+        # layernorm backward, or if we do cooperative reductions (due to the change in
+        # optimal peak memory ordering).
+        @torch._inductor.config.patch(
+            {
+                "triton.cooperative_reductions": False,
+                "triton.force_cooperative_reductions": False,
+                "triton.mix_order_reduction": False,
+            }
+        )
         def test_donated_buffer_inplace(self):
             batch_size = 32
             seq_length = 50
@@ -16906,7 +17690,10 @@ def forward(self, input):
             class CausalSelfAttention(nn.Module):
                 def __init__(self, config):
                     super().__init__()
-                    assert config.n_embd % config.n_head == 0
+                    if config.n_embd % config.n_head != 0:
+                        raise AssertionError(
+                            f"n_embd ({config.n_embd}) must be divisible by n_head ({config.n_head})"
+                        )
                     # key, query, value projections for all heads, but in a batch
                     self.c_attn = nn.Linear(config.n_embd, 3 * config.n_embd)
                     # output projection
@@ -17011,9 +17798,10 @@ def __init__(self, config):
                 def forward(self, idx, targets):
                     device = idx.device
                     b, t = idx.size()
-                    assert t <= self.config.block_size, (
-                        f"Cannot forward sequence of length {t}, block size is only {self.config.block_size}"
-                    )
+                    if t > self.config.block_size:
+                        raise AssertionError(
+                            f"Cannot forward sequence of length {t}, block size is only {self.config.block_size}"
+                        )
                     pos = torch.arange(
                         0, t, dtype=torch.long, device=device
                     )  # shape (t)
@@ -17104,7 +17892,8 @@ def test_respect_scaled_grouped_mm_layout_tag(self):
             # Verify column-major layout
             def _is_column_major(x: torch.Tensor) -> bool:
                 """Check if tensor is column-major (stride(-2) == 1 and stride(-1) > 1)"""
-                assert x.ndim == 2 or x.ndim == 3, "input tensor must be 2D or 3D"
+                if not (x.ndim == 2 or x.ndim == 3):
+                    raise AssertionError("input tensor must be 2D or 3D")
                 return x.stride(-2) == 1 and x.stride(-1) > 1
 
             self.assertTrue(_is_column_major(B_t))
@@ -17201,6 +17990,135 @@ def fn(output, indices, values):
             expected = torch.zeros(512, 768, dtype=torch.bfloat16, device=GPU_TYPE)
             torch.testing.assert_close(result, fn(expected, indices, values))
 
+        @parametrize("dtype", (torch.float32, torch.float64))
+        def test_copysign(self, dtype):
+            def fn(x):
+                out = torch.copysign(x, -0.0)
+                return out
+
+            inp = torch.randn([6], dtype=dtype, device=GPU_TYPE)
+
+            result, code = run_and_get_code(torch.compile(fn), inp)
+            self.assertIn("0x80000000", code[0])
+            torch.testing.assert_close(result, fn(inp))
+
+        def test_3d_reductions_with_max_tiles_3(self):
+            # Inductor only supports at most two reduction iteration ranges, R0 and R1, which the
+            # reduction component of the kernel can be tiled across.
+            # When max_tiles>=3, SIMDScheduling.create_tiling would previously incorrectly allow the
+            # tiling of the kernel in three dimensions, despite there being no pointwise component
+            # of the kernel.
+
+            @torch._inductor.config.patch(
+                {
+                    "triton.prefer_nd_tiling": True,
+                    "triton.max_tiles": 3,
+                    "triton.tile_reductions": True,
+                    "triton.use_tensor_descriptor": True,
+                }
+            )
+            @torch.compile
+            def sum(x):
+                return torch.sum(x, [0, 1, 2])
+
+            shape = (2, 2, 4)
+            strides = (32, 8, 1)
+
+            arg0_1_orig = torch.arange(
+                math.prod(shape), device=GPU_TYPE, dtype=torch.int32
+            ).view(shape)
+            arg0_1 = torch.empty_strided(
+                shape, strides, device=GPU_TYPE, dtype=torch.int32
+            )
+            arg0_1.copy_(arg0_1_orig)
+
+            actual, code = run_and_get_code(sum, arg0_1)
+            expected = torch.sum(arg0_1, [0, 1, 2])
+
+            torch.testing.assert_close(actual=actual, expected=expected)
+
+            fc = FileCheck()
+            # There's no pointwise work to do, so xnumel should be 1...
+            fc.check("xnumel = 1")
+            fc.run(code[0])
+
+        @config.patch({"triton.decompose_sort_ops": True})
+        def test_median_decompose_sort_ops(self):
+            def fn_default(a):
+                return torch.median(a)
+
+            def fn_dim(a):
+                return torch.median(a, dim=1)
+
+            inp = torch.randn(8, 16, device=GPU_TYPE)
+            for fn in (fn_default, fn_dim):
+                torch._dynamo.reset()
+                result, code = run_and_get_code(torch.compile(fn), inp)
+                self.assertIn(
+                    "sort_with_index",
+                    " ".join(code),
+                    "Expected Triton sort codegen for median",
+                )
+                torch.testing.assert_close(result, fn(inp))
+
+        @config.patch({"triton.decompose_sort_ops": True})
+        def test_mode_decompose_sort_ops(self):
+            def fn(a):
+                return torch.mode(a, dim=1)
+
+            # Use integers so ties are common, exercising run-length logic
+            inp = torch.randint(
+                0, 5, size=[8, 16], dtype=torch.float32, device=GPU_TYPE
+            )
+            torch._dynamo.reset()
+            result, code = run_and_get_code(torch.compile(fn), inp)
+            self.assertIn(
+                "sort_with_index",
+                " ".join(code),
+                "Expected Triton sort codegen for mode",
+            )
+            expected = fn(inp)
+            torch.testing.assert_close(result[0], expected[0])
+            torch.testing.assert_close(result[1], expected[1])
+
+        @config.patch({"triton.decompose_sort_ops": True})
+        def test_topk_decompose_sort_ops(self):
+            def fn(a):
+                return torch.topk(a, 3, dim=-1)
+
+            def fn_largest_false(a):
+                return torch.topk(a, 3, dim=-1, largest=False)
+
+            inp = torch.randn(8, 16, device=GPU_TYPE)
+            for test_fn in (fn, fn_largest_false):
+                torch._dynamo.reset()
+                result, code = run_and_get_code(torch.compile(test_fn), inp)
+                self.assertIn(
+                    "sort_with_index",
+                    " ".join(code),
+                    "Expected Triton sort codegen for topk",
+                )
+                expected = test_fn(inp)
+                torch.testing.assert_close(result[0], expected[0])
+                torch.testing.assert_close(result[1], expected[1])
+
+        @config.patch({"triton.decompose_sort_ops": True})
+        def test_kthvalue_decompose_sort_ops(self):
+            def fn(a):
+                return torch.kthvalue(a, 3, dim=-1)
+
+            inp = torch.randn(8, 16, device=GPU_TYPE)
+            torch._dynamo.reset()
+            result, code = run_and_get_code(torch.compile(fn), inp)
+            self.assertIn(
+                "sort_with_index",
+                " ".join(code),
+                "Expected Triton sort codegen for kthvalue",
+            )
+            expected = fn(inp)
+            torch.testing.assert_close(result[0], expected[0])
+            torch.testing.assert_close(result[1], expected[1])
+
     class RNNTest(TestCase):
         device_type = GPU_TYPE
 
diff --git a/test/inductor/test_torchinductor_codegen_config_overrides.py b/test/inductor/test_torchinductor_codegen_config_overrides.py
index d8f1f06afd5df..930a74557aacd 100644
--- a/test/inductor/test_torchinductor_codegen_config_overrides.py
+++ b/test/inductor/test_torchinductor_codegen_config_overrides.py
@@ -1,7 +1,7 @@
 # Owner(s): ["module: inductor"]
 import importlib
 from collections.abc import Callable
-from typing import Any, Optional
+from typing import Any
 from unittest import skipIf
 
 import torch
@@ -30,8 +30,8 @@ def run_and_compare(
         self,
         func: Callable[..., Any],
         *args,
-        compile_kwargs: Optional[dict] = None,
-        config_patches: Optional[dict] = None,
+        compile_kwargs: dict | None = None,
+        config_patches: dict | None = None,
         atol: float | None = 1e-05,
         rtol: float | None = 1e-08,
     ):
@@ -59,7 +59,7 @@ def flatten_tensors(tensors):
 
         return result, code
 
-    def count_code(self, substr: str, code: list[str], expected: Optional[int]):
+    def count_code(self, substr: str, code: list[str], expected: int | None):
         count = sum(prog.count(substr) for prog in code)
         if expected is not None:
             self.assertEqual(count, expected)
@@ -119,6 +119,34 @@ def func(a, b):
         self.count_code("= tl.make_block_ptr(in_ptr", code, 2)
         self.count_code("= tl.load(block_ptr", code, 2)
 
+    @requires_gpu()
+    @skipIf(GPU_TYPE == "mps", "Triton is not available for MPS")
+    @parametrize("disable_welford_reduction", [True, False])
+    def test_disable_welford_reduction(self, disable_welford_reduction: bool):
+        def func(x):
+            return torch.var_mean(x, dim=1)
+
+        # Use a large number to force codegen to prefer welford reduction,
+        # in order to test effectiveness of config flag disable_welford_reduction
+        # This test should run fine on GPU as the configuration is not specific to MTIA backend.
+        x = torch.randn((4, 18000), device=torch.device(GPU_TYPE))
+        config_patches = {
+            "mtia.disable_welford_reduction": disable_welford_reduction,
+        }
+        _, code = self.run_and_compare(
+            func,
+            x,
+            config_patches=config_patches,
+            atol=1e-2,
+            rtol=1e-4,
+        )
+
+        welford_count = sum(prog.count("triton_helpers.welford") for prog in code)
+        if disable_welford_reduction:
+            self.assertEqual(welford_count, 0)
+        else:
+            self.assertGreater(welford_count, 0)
+
     @requires_gpu()
     @skipIf(GPU_TYPE == "mps", "Triton is not available for MPS")
     def test_kernel_fusion_thresholds(self):
diff --git a/test/inductor/test_torchinductor_codegen_dynamic_shapes.py b/test/inductor/test_torchinductor_codegen_dynamic_shapes.py
index c1c963b71496f..75c432f72d5b4 100644
--- a/test/inductor/test_torchinductor_codegen_dynamic_shapes.py
+++ b/test/inductor/test_torchinductor_codegen_dynamic_shapes.py
@@ -87,7 +87,8 @@ def run(*ex, **kwargs):
         code = run_and_get_triton_code(run, *example_inputs, **kwargs)
         self.assertTrue("def triton" in code, f"Failed to find triton kernel\n{code}")
 
-    assert called, "Ran graph without calling compile_fx"
+    if not called:
+        raise AssertionError("Ran graph without calling compile_fx")
 
     torch._dynamo.reset()
 
@@ -112,6 +113,9 @@ def run(*ex, **kwargs):
     "test_as_strided_on_views_dynamic_shapes": TestFailure(
         ("cpu", "cuda", "xpu"), is_skip=True
     ),
+    "test_resize_overlapping_strides_dynamic_shapes": TestFailure(
+        ("cpu",), is_skip=True
+    ),
     #
     # Failed to find dynamic for loop variable:
     #
@@ -127,6 +131,7 @@ def run(*ex, **kwargs):
     "test_arange3_dynamic_shapes": TestFailure(("cpu",)),
     "test_arange4_dynamic_shapes": TestFailure(("cpu",)),
     "test_arange6_dynamic_shapes": TestFailure(("cpu",)),
+    "test_arange7_dynamic_shapes": TestFailure(("cpu",)),
     "test_clamp_type_promotion_dynamic_shapes": TestFailure(("cpu",)),
     "test_conv2d_channels_last_dynamic_shapes": TestFailure(("cpu",)),
     "test_conv3d_dynamic_shapes": TestFailure(("cpu",)),
@@ -215,12 +220,6 @@ def run(*ex, **kwargs):
     "test_linspace4_dynamic_shapes": TestFailure(("cpu", "cuda", "xpu")),
     "test_logcumsumexp_dynamic_shapes": TestFailure(("cpu",)),
     "test_logcumsumexp_zero_dim_dynamic_shapes": TestFailure(("cpu",)),
-    "test_max_pool2d_with_indices_backward5_dynamic_shapes": TestFailure(
-        ("cpu", "cuda")
-    ),
-    "test_max_pool2d_with_indices_backward6_dynamic_shapes": TestFailure(
-        ("cpu", "cuda", "xpu")
-    ),
     "test_misaligned_address_issue1_dynamic_shapes": TestFailure(("cpu",)),
     "test_mm_views_dynamic_shapes": TestFailure(("cpu", "cuda", "xpu")),
     "test_new_empty_dynamic_shapes": TestFailure(("cpu", "cuda", "xpu")),
diff --git a/test/inductor/test_torchinductor_dynamic_shapes.py b/test/inductor/test_torchinductor_dynamic_shapes.py
index e579184978349..8d6adb5844720 100644
--- a/test/inductor/test_torchinductor_dynamic_shapes.py
+++ b/test/inductor/test_torchinductor_dynamic_shapes.py
@@ -27,8 +27,10 @@
 from torch.testing._internal.common_utils import (
     IS_ARM64,
     IS_FBCODE,
+    MI350_ARCH,
     parametrize,
     serialTest,
+    skipIfRocmArch,
     TEST_CUDA_MEM_LEAK_CHECK,
     TEST_WITH_ASAN,
 )
@@ -61,6 +63,8 @@
     # PDL tests are CUDA SM90+ only, skip on CPU
     "test_pdl_mutation_dynamic_shapes": TestFailure(("cpu",), is_skip=True),
     "test_pdl_template_and_delay_dynamic_shapes": TestFailure(("cpu",), is_skip=True),
+    # Bool argmax/argmin fix is Triton-only (see #174069), skip on CPU
+    "test_max_min_bool_dynamic_shapes": TestFailure(("cpu",), is_skip=True),
     # calling div on only symint args
     "test_AllenaiLongformerBase_repro_dynamic_shapes": TestFailure(
         ("cpu", "cuda", "xpu", "mps")
@@ -77,7 +81,6 @@
     "test_reduction3_dynamic_shapes": TestFailure(("mps",)),
     "test_reduction5_dynamic_shapes": TestFailure(("mps",)),
     "test_roll_dynamic_shapes": TestFailure(("mps",)),
-    "test_select_scatter_dtype_consistency_dynamic_shapes": TestFailure(("mps",)),
     "test_std_dynamic_shapes": TestFailure(("mps",)),
     "test_var_correction_dynamic_shapes": TestFailure(("mps",)),
     "test_var_mean_div_by_dynamic_shapes": TestFailure(("mps",)),
@@ -133,6 +136,16 @@ class DynamicShapesGPUTests(TestCase):
         DynamicShapesCommonTemplate, DynamicShapesGPUTests, GPU_TYPE, test_failures
     )
 
+    if HAS_GPU and hasattr(
+        DynamicShapesGPUTests, "test_conv_with_as_strided_dynamic_shapes_cuda"
+    ):
+        # gfx950 shows a deterministic numerical mismatch for this generated test.
+        DynamicShapesGPUTests.test_conv_with_as_strided_dynamic_shapes_cuda = (
+            skipIfRocmArch(MI350_ARCH)(
+                DynamicShapesGPUTests.test_conv_with_as_strided_dynamic_shapes_cuda
+            )
+        )
+
 
 class TestInductorDynamic(TestCase):
     compile_fn = partial(torch.compile, dynamic=True)
@@ -630,6 +643,7 @@ def f(x, w):
         torch.compile(fullgraph=True)(f)(x, w).sum().backward()
         self.assertEqual(orig_w, w.grad)
 
+    @onlyOn(GPU_TYPE)
     @torch._dynamo.config.patch(
         capture_scalar_outputs=True, capture_dynamic_output_shape_ops=True
     )
@@ -640,6 +654,7 @@ def test_embedding_backward_dynamic_shapes_large_grid(self, device):
         On CUDA: uses num_blocks only (not num_blocks * num_warps * warp_size).
         On ROCm: uses num_blocks * num_warps * warp_size (total threads limit).
         """
+        from torch._inductor.runtime.hints import get_warp_size
         from torch._inductor.runtime.triton_heuristics import (
             _check_max_grid_x,
             _num_warps,
@@ -647,13 +662,15 @@ def test_embedding_backward_dynamic_shapes_large_grid(self, device):
 
         size_hints = {"x": 600_000_000}
         x = 64
-        num_warps = _num_warps(8)
+        warp_size = get_warp_size(torch.device(device))
+        num_warps = _num_warps(8, warp_size=warp_size)
 
-        result_x, result_num_blocks = _check_max_grid_x(size_hints, x, num_warps)
+        result_x, result_num_blocks = _check_max_grid_x(
+            size_hints, x, num_warps, warp_size=warp_size
+        )
 
         max_grid_x = 2147483647
         if torch.version.hip:
-            warp_size = 64  # TODO: query warp size once #129663 is merged
             # ROCm limits total threads (num_blocks * num_warps * warp_size)
             self.assertLessEqual(
                 result_num_blocks * num_warps * warp_size,
@@ -1017,21 +1034,21 @@ def _test_wrapper_codegen_statically_known_int_or_none_in_context():
             batch_dim = input_layouts[0].size[0]
             if call_count == 1:
                 # testing fn_1
-                assert (
-                    PythonWrapperCodegen.statically_known_int_or_none(batch_dim) is None
-                ), "Should not be statically known on first call"
+                if (
+                    PythonWrapperCodegen.statically_known_int_or_none(batch_dim)
+                    is not None
+                ):
+                    raise AssertionError("Should not be statically known on first call")
             elif call_count == 2:
                 # testing fn_2
-                assert (
-                    PythonWrapperCodegen.statically_known_int_or_none(batch_dim) == 5
-                ), (
-                    "Should be limited to exactly 5 on second call due to multiple constraints"
-                )
+                if PythonWrapperCodegen.statically_known_int_or_none(batch_dim) != 5:
+                    raise AssertionError(
+                        "Should be limited to exactly 5 on second call due to multiple constraints"
+                    )
             elif call_count == 2:
                 # testing fn_3
-                assert (
-                    PythonWrapperCodegen.statically_known_int_or_none(batch_dim) == 5
-                ), "Should be exactly 5 on third call"
+                if PythonWrapperCodegen.statically_known_int_or_none(batch_dim) != 5:
+                    raise AssertionError("Should be exactly 5 on third call")
 
         class TestWrapperCodegen(PythonWrapperCodegen):
             def __init__(self, *args, **kwargs):
@@ -1188,7 +1205,7 @@ def reduce_bounded(x, y):
             """Reduce over a dimension with bounded size."""
             # x shape: [batch, features, reduction_dim]
             # reduction_dim is dynamic but bounded to max 128
-            assert x.shape[2] <= 64, f"Reduction dim {x.shape[2]} exceeds max 128"
+            assert x.shape[2] <= 64, f"Reduction dim {x.shape[2]} exceeds max 128"  # noqa: S101
 
             # Perform reduction (sum) over the last dimension
             result = torch.sum(x * y, dim=2)
@@ -1215,7 +1232,8 @@ def reduce_bounded(x, y):
         FileCheck().check_not("@triton_heuristics.persistent").run(source_codes[0])
         expected = reduce_bounded(x, y)
 
-        assert torch.allclose(result, expected, atol=1e-3, rtol=1e-3)
+        if not torch.allclose(result, expected, atol=1e-3, rtol=1e-3):
+            raise AssertionError
 
     def test_unspecialized_float_dynamic(self):
         def fn(x, y):
diff --git a/test/inductor/test_torchinductor_opinfo.py b/test/inductor/test_torchinductor_opinfo.py
index cefc17d7aa4eb..3ed8b24fc4d33 100644
--- a/test/inductor/test_torchinductor_opinfo.py
+++ b/test/inductor/test_torchinductor_opinfo.py
@@ -2,6 +2,7 @@
 import atexit
 import contextlib
 import functools
+import itertools
 import math
 import os
 import sys
@@ -31,6 +32,7 @@
 from torch.testing._internal.common_methods_invocations import op_db, skipOps
 from torch.testing._internal.common_utils import (
     IS_CI,
+    IS_LINUX,
     IS_MACOS,
     IS_WINDOWS,
     IS_X86,
@@ -107,11 +109,14 @@
 END = os.getenv("PYTORCH_TEST_RANGE_END", None)
 
 if START is not None or END is not None:
-    assert END is not None
-    assert START is not None
+    if END is None:
+        raise AssertionError("END must be set when START is set")
+    if START is None:
+        raise AssertionError("START must be set when END is set")
     START = int(START)
     END = int(END)
-    assert START < END
+    if START >= END:
+        raise AssertionError(f"START ({START}) must be less than END ({END})")
 else:
     START = 0
     END = len(op_db)
@@ -228,6 +233,11 @@ def format_op(op):
 
 inductor_skips["xpu"] = {}
 
+# torch-xpu-ops: #2956
+inductor_skips["xpu"]["lu"] = {f32}
+inductor_skips["xpu"]["nn.functional.linear"] = {f16}
+inductor_skips["xpu"]["masked.cumprod"] = {f16}
+
 inductor_expected_failures_single_sample = defaultdict(dict)
 
 inductor_expected_failures_single_sample["cpu"] = {
@@ -291,27 +301,6 @@ def format_op(op):
         i32,
         i64,
     },  # align with cuda.
-    ("linalg.pinv", "singular"): {f64},
-    # could not create a primitive
-    "addmv": {f64},
-    "fft.fft": {f16},
-    "fft.fft2": {f16},
-    "fft.fftn": {f16},
-    "fft.hfft": {f16},
-    "fft.hfft2": {f16},
-    "fft.hfftn": {f16},
-    "fft.rfft": {f16},
-    "fft.rfft2": {f16},
-    "fft.rfftn": {f16},
-    "fft.ifft": {f16},
-    "fft.ifft2": {f16},
-    "fft.ifftn": {f16},
-    "fft.ihfft": {f16},
-    "fft.ihfft2": {f16},
-    "fft.ihfftn": {f16},
-    "fft.irfft": {f16},
-    "fft.irfft2": {f16},
-    "fft.irfftn": {f16},
 }
 
 
@@ -987,8 +976,33 @@ def wrapper_noop_set_seed(op, *args, **kwargs):
 
 # TODO: Fix these so strides match.
 inductor_skip_exact_stride = {
+    "complex",
+    "empty_permuted",
+    "fft.irfftn",
+    "fft.irfft2",
+    "linalg.diagonal",
+    "linalg.eigvals",  # Fails for ROCM
+    "linalg.lu",
+    "linalg.lu_factor",
+    "linalg.lu_factor_ex",
     "linalg.matrix_norm",
+    "linalg.norm",
+    "linalg.norm.subgradients_at_zero",
+    "linalg.pinv.singular",
+    "linalg.svdvals",
+    "linalg.solve",
+    "linalg.solve_ex",
+    "linalg.qr",
+    "lu",
+    "matmul",
+    "__rmatmul__",
+    "nn.functional.adaptive_avg_pool1d",
+    "nn.functional.group_norm",
+    "nn.functional.linear",
+    "nn.functional.max_pool2d",
+    "nn.functional.unfold",
     "ormqr",
+    "pca_lowrank",
     "rot90",
     "sum",
     "tensordot",
@@ -1036,14 +1050,18 @@ def argsort_sort_assert_equal(
         exact_stride=False,
     ):
         if is_argsort:
-            assert isinstance(x, torch.Tensor)
-            assert isinstance(y, torch.Tensor)
+            if not isinstance(x, torch.Tensor):
+                raise AssertionError(f"Expected torch.Tensor, got {type(x)}")
+            if not isinstance(y, torch.Tensor):
+                raise AssertionError(f"Expected torch.Tensor, got {type(y)}")
         else:
             # The first tensor is the sorted values and can be asserted via
             # the usual means
             for t in (x, y):
-                assert isinstance(t, tuple)
-                assert len(t) == 2
+                if not isinstance(t, tuple):
+                    raise AssertionError(f"Expected tuple, got {type(t)}")
+                if len(t) != 2:
+                    raise AssertionError(f"Expected tuple of length 2, got {len(t)}")
 
             test_case_inst.assertEqual(
                 x[0],
@@ -1062,7 +1080,8 @@ def argsort_sort_assert_equal(
         if exact_dtype and (x.dtype != y.dtype):
             raise AssertionError(f"The dtypes do not match: {x.dtype} != {y.dtype}.")
 
-        assert x.shape == y.shape
+        if x.shape != y.shape:
+            raise AssertionError(f"Shape mismatch: {x.shape} != {y.shape}")
 
         if exact_stride and (x.stride() != y.stride()):
             raise AssertionError(
@@ -1075,7 +1094,8 @@ def el_to_indices(el):
             for cur_dim in reversed(range(x.dim())):
                 indices[cur_dim] = el % x.shape[cur_dim]
                 el //= x.shape[cur_dim]
-            assert None not in indices
+            if None in indices:
+                raise AssertionError("indices contains None")
             return indices
 
         def get_val_by_ids(t, ids):
@@ -1147,6 +1167,42 @@ def inner(self, device, dtype, op):
     return inner
 
 
+def _inductor_extra_samples(op_name, device, dtype, requires_grad):
+    """Extra sample inputs for inductor-specific coverage.
+
+    These exercise dynamo decomposition paths (e.g. tensor value/alpha triggering
+    fma) that the shared opinfo samples don't cover.
+    """
+    from torch.testing._internal.common_methods_invocations import (
+        make_tensor,
+        S,
+        SampleInput,
+    )
+
+    make_arg = partial(
+        make_tensor, device=device, dtype=dtype, requires_grad=requires_grad
+    )
+
+    if op_name in ("addcmul", "addcdiv") and (
+        dtype.is_floating_point or dtype.is_complex
+    ):
+        # Tensor value
+        args = tuple(
+            make_arg(shape, exclude_zero=True) for shape in ((S, S), (S, S), (S, S))
+        )
+        tensor_value = make_arg((), requires_grad=False)
+        return [SampleInput(*args, value=tensor_value)]
+
+    if op_name == "add" and (dtype.is_floating_point or dtype.is_complex):
+        # Tensor alpha
+        lhs = make_arg((S, S))
+        rhs = make_arg((S, S))
+        tensor_alpha = make_arg((), requires_grad=False)
+        return [SampleInput(lhs, args=(rhs,), kwargs={"alpha": tensor_alpha})]
+
+    return []
+
+
 @wrapper_noop_set_seed_decorator
 class TestInductorOpInfo(TestCase):
     def tearDown(self):
@@ -1176,11 +1232,13 @@ def tearDown(self):
     )
     @torch._inductor.config.patch("test_configs.runtime_triton_dtype_assert", True)
     @torch._inductor.config.patch("test_configs.static_cpp_dtype_assert", True)
+    @torch._inductor.config.patch("shape_padding", False)
     @collection_decorator
     def test_comprehensive(self, device, dtype, op):
         device_type = torch.device(device).type
 
-        assert device_type in (GPU_TYPE, "cpu")
+        if device_type not in (GPU_TYPE, "cpu"):
+            raise AssertionError(f"Unexpected device_type: {device_type}")
 
         torch._dynamo.reset()
         with torch.no_grad():
@@ -1216,8 +1274,11 @@ def test_comprehensive(self, device, dtype, op):
             # with open("test_output.txt", "a") as f:
             #     print(f"SKIPPING OP {op_name} on {device_type}", flush=True, file=f)
             #     print(f"SKIPPING OP {op_name} on {device_type}", flush=True)
-        elif dtype in inductor_expected_failures_single_sample[device_type].get(
-            op_name, set()
+        elif (
+            device_type == "cpu"
+            and IS_LINUX
+            and dtype
+            in inductor_expected_failures_single_sample[device_type].get(op_name, set())
         ) or dtype in inductor_gradient_expected_failures_single_sample[
             device_type
         ].get(op_name, set()):
@@ -1247,6 +1308,9 @@ def fn(*args, **kwargs):
             and dtype != torch.complex32
         )
         samples = op.sample_inputs(device, dtype, requires_grad=requires_grad)
+        extra = _inductor_extra_samples(op_name, device, dtype, requires_grad)
+        if extra:
+            samples = itertools.chain(samples, extra)
 
         if (
             dtype in inductor_one_sample.get(device_type, {}).get(op_name, {})
diff --git a/test/inductor/test_torchinductor_opinfo_properties.py b/test/inductor/test_torchinductor_opinfo_properties.py
new file mode 100644
index 0000000000000..362b11db6611b
--- /dev/null
+++ b/test/inductor/test_torchinductor_opinfo_properties.py
@@ -0,0 +1,914 @@
+# Owner(s): ["module: inductor"]
+"""
+OpInfo-based property tests for inductor numerical correctness.
+
+Tests three properties:
+1. Batch invariance - output shouldn't change based on batch size
+2. Run-to-run determinism - same input should give same output across
+   compilations, even with different autotuning choices
+3. Bitwise equivalence with torch eager mode
+
+Tests three compilation backends:
+1. aot_eager_decomp_partition - AOT autograd with eager execution
+2. inductor_default - Standard inductor compilation
+3. inductor_numerics - Inductor with strict numerics flags
+
+Focuses on ops commonly used in LLMs from unary_ufuncs, binary_ufuncs, and op_db.
+
+Note: How to modify this test
+-----------------------------
+To add a new op:
+    Add the op name to LLM_UNARY_OP_NAMES, LLM_BINARY_OP_NAMES, or LLM_OP_DB_NAMES
+    depending on which OpInfo database contains the op.
+
+To add an expected failure:
+    Add to the appropriate *_XFAILS dict (e.g., EAGER_EQUIV_XFAILS):
+        "backend": {"op_name": {dtype1, dtype2, ...}}
+    Use ALL (None) to match all dtypes.
+
+To remove a stale expected failure:
+    Run the tests - any expected failure that now passes will fail with
+    "XPASS" and tell you which entry to remove.
+
+To run a specific test:
+    pytest test/inductor/test_torchinductor_opinfo_properties.py -k "silu and eager"
+"""
+
+import sys
+import unittest
+
+import pytest
+
+import torch
+from torch._inductor.test_case import run_tests, TestCase
+from torch._inductor.utils import fresh_inductor_cache
+from torch.testing._internal.common_device_type import (
+    instantiate_device_type_tests,
+    onlyCUDA,
+    ops,
+)
+from torch.testing._internal.common_methods_invocations import (
+    binary_ufuncs,
+    op_db,
+    unary_ufuncs,
+)
+from torch.testing._internal.common_utils import (
+    IS_WINDOWS,
+    parametrize,
+    skipIfTorchDynamo,
+    TEST_WITH_ASAN,
+)
+from torch.testing._internal.inductor_utils import HAS_GPU
+
+
+# LLM-useful op names to filter from unary_ufuncs
+LLM_UNARY_OP_NAMES = {
+    # Basic math
+    "abs",
+    "neg",
+    "reciprocal",
+    # Exponential/log
+    "exp",
+    "exp2",
+    "expm1",
+    "log",
+    "log2",
+    "log10",
+    "log1p",
+    # Power/root
+    "sqrt",
+    "rsqrt",
+    # Trigonometric
+    "sin",
+    "cos",
+    "tan",
+    "tanh",
+    # Activations
+    "sigmoid",
+    "nn.functional.relu6",
+}
+
+# LLM-useful op names to filter from binary_ufuncs
+LLM_BINARY_OP_NAMES = {
+    "sub",
+    "mul",
+    "remainder",
+    "fmod",
+    "maximum",
+    "minimum",
+}
+
+# LLM-useful op names from op_db (nn.functional and others)
+LLM_OP_DB_NAMES = {
+    # Activations
+    "nn.functional.gelu",
+    "nn.functional.silu",
+    "nn.functional.leaky_relu",
+    "nn.functional.hardswish",
+    # Normalization
+    "nn.functional.layer_norm",
+    "nn.functional.rms_norm",
+    # Attention/linear
+    "nn.functional.linear",
+    "bmm",
+    # Softmax
+    "softmax",
+    "log_softmax",
+}
+
+# Filter OpInfos for LLM-useful ops
+llm_unary_ops = [op for op in unary_ufuncs if op.name in LLM_UNARY_OP_NAMES]
+llm_binary_ops = [op for op in binary_ufuncs if op.name in LLM_BINARY_OP_NAMES]
+llm_op_db_ops = [op for op in op_db if op.name in LLM_OP_DB_NAMES]
+
+# Combine all ops, avoiding duplicates by name
+_seen_names = set()
+llm_ops = []
+for op in llm_unary_ops + llm_binary_ops + llm_op_db_ops:
+    if op.name not in _seen_names:
+        _seen_names.add(op.name)
+        llm_ops.append(op)
+
+# Backends to test
+BACKENDS = [
+    "aot_eager_decomp_partition",
+    "inductor_default",
+    "inductor_numerics",
+]
+
+# Dtypes to test - common LLM dtypes
+DTYPES = [
+    torch.float32,
+    torch.float16,
+    torch.bfloat16,
+]
+
+# Number of samples for numerical testing
+NUM_SAMPLES = 65536
+
+
+def generate_exhaustive_16bit(dtype, device):
+    """Generate all 65536 possible values for a 16-bit float dtype.
+
+    For fp16 and bf16, there are exactly 2^16 = 65536 possible bit patterns.
+    This allows exhaustive testing of unary functions.
+    """
+    # Generate all 16-bit patterns as uint16
+    all_patterns = torch.arange(0, 65536, dtype=torch.int32, device=device)
+    # View as the target dtype
+    if dtype == torch.float16:
+        return all_patterns.to(torch.int16).view(torch.float16)
+    elif dtype == torch.bfloat16:
+        return all_patterns.to(torch.int16).view(torch.bfloat16)
+    else:
+        raise ValueError(f"Unsupported dtype for exhaustive 16-bit testing: {dtype}")
+
+
+def generate_sampled_fp32(num_samples, device, seed=42):
+    """Generate uniformly sampled fp32 bit patterns.
+
+    Samples random 32-bit patterns and interprets them as fp32.
+    This provides good coverage of the fp32 space including denormals,
+    infinities, and NaNs.
+    """
+    gen = torch.Generator(device=device)
+    gen.manual_seed(seed)
+    # Generate random 32-bit patterns
+    patterns = torch.randint(
+        0, 2**31, (num_samples,), dtype=torch.int32, device=device, generator=gen
+    )
+    # Also include negative patterns (high bit set)
+    signs = torch.randint(
+        0, 2, (num_samples,), dtype=torch.int32, device=device, generator=gen
+    )
+    patterns = patterns | (signs << 31)
+    return patterns.view(torch.float32)
+
+
+def generate_sampled_pairs(num_samples, dtype, device, seed=42):
+    """Generate pairs of sampled values for binary function testing.
+
+    For fp32: uses random bit patterns
+    For fp16/bf16: samples from all possible values
+    """
+    gen = torch.Generator(device=device)
+    gen.manual_seed(seed)
+
+    if dtype == torch.float32:
+        # Random bit patterns for fp32
+        patterns1 = torch.randint(
+            0, 2**31, (num_samples,), dtype=torch.int32, device=device, generator=gen
+        )
+        signs1 = torch.randint(
+            0, 2, (num_samples,), dtype=torch.int32, device=device, generator=gen
+        )
+        patterns1 = patterns1 | (signs1 << 31)
+        x = patterns1.view(torch.float32)
+
+        patterns2 = torch.randint(
+            0, 2**31, (num_samples,), dtype=torch.int32, device=device, generator=gen
+        )
+        signs2 = torch.randint(
+            0, 2, (num_samples,), dtype=torch.int32, device=device, generator=gen
+        )
+        patterns2 = patterns2 | (signs2 << 31)
+        y = patterns2.view(torch.float32)
+    else:
+        # Sample indices into the 16-bit space
+        indices1 = torch.randint(
+            0, 65536, (num_samples,), dtype=torch.int32, device=device, generator=gen
+        )
+        indices2 = torch.randint(
+            0, 65536, (num_samples,), dtype=torch.int32, device=device, generator=gen
+        )
+
+        if dtype == torch.float16:
+            x = indices1.to(torch.int16).view(torch.float16)
+            y = indices2.to(torch.int16).view(torch.float16)
+        elif dtype == torch.bfloat16:
+            x = indices1.to(torch.int16).view(torch.bfloat16)
+            y = indices2.to(torch.int16).view(torch.bfloat16)
+        else:
+            raise ValueError(f"Unsupported dtype: {dtype}")
+
+    return x, y
+
+
+# Inductor options for strict numerics matching eager behavior
+INDUCTOR_NUMERICS_OPTIONS = {
+    "deterministic": True,
+    "fallback_random": True,
+    "emulate_precision_casts": True,
+    "eager_numerics.division_rounding": True,
+    "eager_numerics.disable_ftz": True,
+    "eager_numerics.use_pytorch_libdevice": True,
+}
+
+
+def _slice_tensor_or_collection(value, batch_size, original_batch_size, input_ndim):
+    """Recursively slice tensors in a value (tensor, list, tuple, or other).
+
+    Only slices tensors that have the same number of dimensions as the input
+    to preserve broadcasting patterns.
+    """
+    if isinstance(value, torch.Tensor):
+        if (
+            value.dim() == input_ndim  # Same ndim to preserve broadcast pattern
+            and value.dim() > 0
+            and value.shape[0] == original_batch_size
+            and value.shape[0] > 1  # Don't slice broadcast dimensions
+        ):
+            return value[:batch_size]
+        return value
+    elif isinstance(value, (list, tuple)):
+        sliced = [
+            _slice_tensor_or_collection(v, batch_size, original_batch_size, input_ndim)
+            for v in value
+        ]
+        return type(value)(sliced)
+    else:
+        return value
+
+
+def slice_tensors_to_batch_size(sample_input, batch_size):
+    """Slice all tensors in a SampleInput to the given batch size.
+
+    For batch invariance testing, we need to slice all tensor inputs that have
+    matching batch dimensions. This function slices:
+    - sample_input.input if it's a tensor with dim > 0
+    - All tensor args that have the same size in dim 0 as the input (not broadcast dims)
+    - All tensor kwargs that have the same size in dim 0 as the input (not broadcast dims)
+
+    Tensors with size 1 in dim 0 are broadcast dimensions and should not be sliced.
+    Also handles nested lists/tuples of tensors (e.g., TensorList).
+    Only slices tensors with the same ndim as the input to preserve broadcast patterns.
+
+    Returns a tuple (sliced_input, sliced_args, sliced_kwargs), or None if slicing is not possible.
+    """
+    inp = sample_input.input
+    if not isinstance(inp, torch.Tensor) or inp.dim() == 0:
+        return None
+
+    original_batch_size = inp.shape[0]
+    input_ndim = inp.dim()
+    if batch_size > original_batch_size:
+        return None
+
+    # Slice the input
+    sliced_input = inp[:batch_size]
+
+    # Slice args - handles tensors and lists/tuples of tensors
+    sliced_args = tuple(
+        _slice_tensor_or_collection(arg, batch_size, original_batch_size, input_ndim)
+        for arg in sample_input.args
+    )
+
+    # Slice kwargs - handles tensors and lists/tuples of tensors
+    sliced_kwargs = {
+        key: _slice_tensor_or_collection(
+            val, batch_size, original_batch_size, input_ndim
+        )
+        for key, val in sample_input.kwargs.items()
+    }
+
+    return sliced_input, sliced_args, sliced_kwargs
+
+
+def sample_operates_on_batch_dim(op_name, sample_input):
+    """Check if a sample input operates on the batch dimension (dim 0).
+
+    For ops that normalize/reduce over a dimension, if that dimension is 0,
+    slicing the batch will change the result and batch invariance doesn't apply.
+    """
+    # Ops that take a 'dim' argument and normalize/reduce over it
+    dim_based_ops = {
+        "softmax",
+        "log_softmax",
+        "nn.functional.softmax",
+        "nn.functional.log_softmax",
+    }
+
+    if op_name not in dim_based_ops:
+        return False
+
+    # Get dim from args or kwargs
+    dim = None
+    if sample_input.args:
+        dim = sample_input.args[0]
+    if "dim" in sample_input.kwargs:
+        dim = sample_input.kwargs["dim"]
+
+    # If dim is 0 or -ndim (equivalent to 0), the op operates on the batch dimension
+    if dim is not None:
+        inp = sample_input.input
+        if isinstance(inp, torch.Tensor):
+            ndim = inp.dim()
+            # Normalize negative dim
+            if dim < 0:
+                dim = dim + ndim
+            return dim == 0
+
+    return False
+
+
+# Expected failures for bitwise equivalence tests.
+# Structure: backend -> op_name -> set of failing dtypes (None means all dtypes)
+# These track known numerical differences between eager and compiled execution.
+# The goal is to eventually fix these and remove entries.
+
+# Dtype shorthands
+fp32, fp16, bf16, ALL = torch.float32, torch.float16, torch.bfloat16, None
+
+EAGER_EQUIV_XFAILS = {
+    "aot_eager_decomp_partition": {
+        "nn.functional.gelu": {fp32},
+        "nn.functional.layer_norm": {fp32},
+        "nn.functional.rms_norm": {fp32},
+        "softmax": {fp32},
+        "log_softmax": {fp32},
+    },
+    "inductor_default": {
+        "reciprocal": {fp32},
+        "remainder": {ALL},
+        "sigmoid": {fp32},
+        "nn.functional.gelu": {fp32},
+        "nn.functional.layer_norm": {fp32},
+        "nn.functional.silu": {fp16, fp32},
+        "softmax": {fp32},
+        "log_softmax": {fp32},
+    },
+    "inductor_numerics": {
+        "remainder": {ALL},
+        "sigmoid": {fp32},
+        "sub": {ALL},
+        "nn.functional.gelu": {fp32},
+        "nn.functional.layer_norm": {fp32},
+        "softmax": {fp32},
+        "log_softmax": {fp32},
+    },
+}
+
+DETERMINISM_XFAILS = {}
+
+BATCH_INVARIANCE_XFAILS = {
+    "aot_eager_decomp_partition": {
+        "nn.functional.linear": {ALL},
+    },
+    "inductor_default": {
+        "nn.functional.linear": {ALL},
+    },
+    "inductor_numerics": {
+        "nn.functional.linear": {ALL},
+    },
+}
+
+UNARY_NUMERICAL_XFAILS = {
+    "inductor_default": {
+        "exp2": {bf16, fp32},
+        "expm1": {bf16, fp32},
+        "log1p": {fp32},
+        "reciprocal": {fp32},
+        "rsqrt": {bf16, fp32},
+        "sigmoid": {fp32},
+        "sin": {fp32},
+        "tan": {bf16, fp32},
+        "tanh": {fp32},
+    },
+    "inductor_numerics": {
+        "sigmoid": {fp32},
+    },
+}
+
+BINARY_NUMERICAL_XFAILS = {
+    "inductor_default": {
+        "fmod": {bf16, fp32},
+        "remainder": {ALL},
+    },
+    "inductor_numerics": {
+        "remainder": {ALL},
+    },
+}
+
+XFAIL_DICTS = {
+    "eager_equivalence": EAGER_EQUIV_XFAILS,
+    "determinism": DETERMINISM_XFAILS,
+    "batch_invariance": BATCH_INVARIANCE_XFAILS,
+    "unary_numerical": UNARY_NUMERICAL_XFAILS,
+    "binary_numerical": BINARY_NUMERICAL_XFAILS,
+}
+
+ROCM_EAGER_EQUIV_XFAILS = {
+    "aot_eager_decomp_partition": {
+        "nn.functional.gelu": {fp32},
+        "nn.functional.layer_norm": {fp32},
+        "nn.functional.rms_norm": {fp32},
+        "softmax": {fp32},
+        "log_softmax": {fp32},
+    },
+    "inductor_default": {
+        "sigmoid": {fp32},
+        "nn.functional.gelu": {fp32},
+        "nn.functional.layer_norm": {fp32},
+        "nn.functional.silu": {fp16, fp32},
+        "softmax": {fp32},
+        "log_softmax": {fp32},
+    },
+    "inductor_numerics": {
+        "sigmoid": {fp32},
+        "sub": {ALL},
+        "nn.functional.gelu": {fp32},
+        "nn.functional.layer_norm": {fp32},
+        "softmax": {fp32},
+        "log_softmax": {fp32},
+    },
+}
+
+ROCM_DETERMINISM_XFAILS = {}
+
+ROCM_BATCH_INVARIANCE_XFAILS = {
+    "aot_eager_decomp_partition": {
+        "nn.functional.linear": {ALL},
+    },
+    "inductor_default": {
+        "nn.functional.linear": {ALL},
+        "log1p": {fp32},
+    },
+    "inductor_numerics": {
+        "nn.functional.linear": {ALL},
+    },
+}
+
+ROCM_UNARY_NUMERICAL_XFAILS = {
+    "inductor_default": {
+        "log1p": {fp32},
+        "rsqrt": {bf16, fp32},
+        "sigmoid": {fp32},
+        "sin": {fp32},
+        "tan": {fp32},
+        "tanh": {fp32},
+    },
+    "inductor_numerics": {
+        "sigmoid": {fp32},
+    },
+}
+
+ROCM_BINARY_NUMERICAL_XFAILS = {}
+
+ROCM_XFAIL_DICTS = {
+    "eager_equivalence": ROCM_EAGER_EQUIV_XFAILS,
+    "determinism": ROCM_DETERMINISM_XFAILS,
+    "batch_invariance": ROCM_BATCH_INVARIANCE_XFAILS,
+    "unary_numerical": ROCM_UNARY_NUMERICAL_XFAILS,
+    "binary_numerical": ROCM_BINARY_NUMERICAL_XFAILS,
+}
+
+
+def is_expected_failure(device_type, op_name, backend, test_type, dtype=None):
+    """Check if a test is expected to fail."""
+    xfail_dicts = ROCM_XFAIL_DICTS if torch.version.hip is not None else XFAIL_DICTS
+    xfails = xfail_dicts.get(test_type, {}).get(backend, {}).get(op_name, set())
+    return dtype in xfails or ALL in xfails
+
+
+def compile_fn(fn, backend):
+    """Compile a function with the given backend."""
+    if backend == "aot_eager_decomp_partition":
+        return torch.compile(fn, backend="aot_eager_decomp_partition")
+    elif backend == "inductor_default":
+        return torch.compile(fn, backend="inductor")
+    elif backend == "inductor_numerics":
+        return torch.compile(fn, backend="inductor", options=INDUCTOR_NUMERICS_OPTIONS)
+    else:
+        raise ValueError(f"Unknown backend: {backend}")
+
+
+@unittest.skipIf(IS_WINDOWS, "Skipped on Windows")
+@unittest.skipIf(TEST_WITH_ASAN, "Skipped under ASAN")
+@unittest.skipIf(not HAS_GPU, "Requires GPU")
+class TestOpInfoProperties(TestCase):
+    """Test op properties under various inductor modes using OpInfo on CUDA."""
+
+    def setUp(self):
+        super().setUp()
+        torch._dynamo.reset()
+        patch = torch._dynamo.config.patch(
+            recompile_limit=sys.maxsize, accumulated_recompile_limit=sys.maxsize
+        )
+        patch.__enter__()
+        self.addCleanup(patch.__exit__, None, None, None)
+
+    def tearDown(self):
+        super().tearDown()
+        torch._dynamo.reset()
+
+    def _get_sample_inputs(self, op, device, dtype):
+        """Get sample inputs from OpInfo using reference_inputs for comprehensive coverage."""
+        # Use reference_inputs for more comprehensive test coverage
+        # Falls back to sample_inputs if reference_inputs_func is not defined
+        try:
+            samples = list(op.reference_inputs(device, dtype, requires_grad=False))
+        except Exception:
+            samples = list(op.sample_inputs(device, dtype, requires_grad=False))
+        return samples
+
+    def _run_with_expected_failure(
+        self, device_type, op_name, backend, test_type, dtype, test_fn
+    ):
+        """Run a test with expected failure handling.
+
+        Uses pytest.xfail for clear test output:
+        - If test is expected to fail and does fail: XFAIL (pytest.xfail called)
+        - If test is expected to fail but passes: FAILED (strict xpass behavior)
+        - If test is not expected to fail: runs normally (PASSED or FAILED)
+
+        Args:
+            test_fn: A callable that runs the actual test assertions
+        """
+        expected = is_expected_failure(device_type, op_name, backend, test_type, dtype)
+
+        if expected:
+            try:
+                test_fn()
+            except (AssertionError, RuntimeError):
+                pytest.xfail(f"Known failure: {op_name}/{backend}/{dtype}")
+            else:
+                self.fail(
+                    f"XPASS: {op_name}/{backend}/{dtype} - remove from {test_type} xfails"
+                )
+        else:
+            test_fn()
+
+    # =========================================================================
+    # Batch Invariance Tests
+    # =========================================================================
+
+    @onlyCUDA
+    @skipIfTorchDynamo("Test uses dynamo already")
+    @ops(llm_ops, allowed_dtypes=DTYPES)
+    @parametrize("backend", BACKENDS)
+    def test_batch_invariance(self, device, dtype, op, backend):
+        """Test batch invariance with exponentially decreasing batch sizes.
+
+        For each sample input, this test:
+        1. Runs the compiled op at full batch size
+        2. Runs at size/2, size/4, etc. (exponentially decreasing)
+        3. Verifies the sliced output matches the corresponding slice of the full output
+
+        All tensor inputs with matching batch dimensions are sliced together.
+
+        Note: We disable split-k GEMM accumulation to ensure consistent results
+        across batch sizes, as split-k can produce different rounding based on
+        how the work is partitioned.
+        """
+        torch._dynamo.reset()
+        device_type = torch.device(device).type
+
+        samples = self._get_sample_inputs(op, device, dtype)
+        if not samples:
+            self.skipTest(f"No samples for {op.name}")
+
+        fn = op.get_op()
+
+        # Disable split-k accumulation in GEMM to ensure batch-invariant results.
+        # Split-k can produce different rounding based on how work is partitioned.
+        # Disabling split-k requires the cuBLASLt backend, so we set it first.
+        prev_blas_library = torch.backends.cuda.preferred_blas_library()
+        torch.backends.cuda.preferred_blas_library("cublaslt")
+
+        # We need to save both the reduced precision and split_k settings.
+        prev_fp16 = (
+            torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction,
+            torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction_split_k,
+        )
+        prev_bf16 = (
+            torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction,
+            torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction_split_k,
+        )
+        # Disable both reduced precision and split-k
+        torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction = (
+            False,
+            False,
+        )
+        torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction = (
+            False,
+            False,
+        )
+
+        def run_test():
+            tested_any = False
+
+            for sample in samples:
+                # Skip if input is not a tensor or is 0-dim
+                if (
+                    not isinstance(sample.input, torch.Tensor)
+                    or sample.input.dim() == 0
+                ):
+                    continue
+
+                # Need at least size 4 in first dimension for meaningful test
+                full_size = sample.input.shape[0]
+                if full_size < 4:
+                    continue
+
+                # Skip broadcast/expanded tensors (stride 0 in batch dim)
+                # These don't have meaningful batch invariance since all rows are the same
+                if sample.input.stride()[0] == 0:
+                    continue
+
+                # Skip samples where the op normalizes/reduces over dim 0 (batch dimension)
+                # because slicing the batch changes the normalization result
+                if sample_operates_on_batch_dim(op.name, sample):
+                    continue
+
+                compiled_fn = compile_fn(fn, backend)
+
+                # Get reference output at full size
+                full_args = (sample.input,) + tuple(sample.args)
+                full_kwargs = sample.kwargs
+                full_out = compiled_fn(*full_args, **full_kwargs)
+
+                if not isinstance(full_out, torch.Tensor):
+                    continue
+
+                # Skip if output is 0-dim (scalar) - can't slice it
+                if full_out.dim() == 0:
+                    continue
+
+                # Skip if output's first dimension doesn't match input's batch size
+                # (e.g., due to broadcasting or reduction)
+                if full_out.shape[0] != full_size:
+                    continue
+
+                tested_any = True
+
+                # Test with exponentially decreasing sizes: size, size/2, size/4, ...
+                size = full_size
+                while size >= 1:
+                    # Slice all tensor inputs with matching batch dimensions
+                    sliced = slice_tensors_to_batch_size(sample, size)
+                    if sliced is None:
+                        break
+                    sliced_input, sliced_args, sliced_kwargs = sliced
+
+                    out = compiled_fn(sliced_input, *sliced_args, **sliced_kwargs)
+
+                    # Verify output matches the corresponding slice of full output (bitwise)
+                    self.assertEqual(out, full_out[:size], rtol=0, atol=0)
+
+                    # Step down exponentially
+                    size = size // 2
+
+            if not tested_any:
+                self.skipTest("No suitable samples found for batch invariance test")
+
+        try:
+            self._run_with_expected_failure(
+                device_type, op.name, backend, "batch_invariance", dtype, run_test
+            )
+        finally:
+            torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction = (
+                prev_fp16
+            )
+            torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction = (
+                prev_bf16
+            )
+            torch.backends.cuda.preferred_blas_library(prev_blas_library)
+
+    # =========================================================================
+    # Run-to-Run Determinism Tests
+    # =========================================================================
+
+    @onlyCUDA
+    @skipIfTorchDynamo("Test uses dynamo already")
+    @ops(llm_ops, allowed_dtypes=DTYPES)
+    @parametrize("backend", BACKENDS)
+    def test_determinism(self, device, dtype, op, backend):
+        """Test run-to-run determinism with fresh cache and benchmark perturbation.
+
+        For enhanced determinism testing, this test:
+        1. Resets dynamo and uses fresh inductor cache between compilations
+        2. Uses distort_benchmarking_result config to perturb autotuning choices
+        3. Enables deterministic algorithms to ensure reliable determinism
+        4. Tests multiple runs and verifies bitwise identical outputs
+        """
+        torch._dynamo.reset()
+        device_type = torch.device(device).type
+
+        samples = self._get_sample_inputs(op, device, dtype)
+        if not samples:
+            self.skipTest(f"No samples for {op.name}")
+
+        fn = op.get_op()
+
+        # Enable deterministic algorithms for reliable determinism testing
+        prev_deterministic = torch.are_deterministic_algorithms_enabled()
+        torch.use_deterministic_algorithms(True)
+
+        def run_test():
+            import torch._inductor.config as inductor_config
+
+            for sample in samples:
+                args = (sample.input,) + tuple(sample.args)
+                kwargs = sample.kwargs
+
+                # Run 1: Fresh cache, normal benchmarking
+                torch._dynamo.reset()
+                with fresh_inductor_cache():
+                    compiled_fn1 = compile_fn(fn, backend)
+                    out1 = compiled_fn1(*args, **kwargs)
+
+                # Run 2: Fresh cache, inverted benchmark results
+                torch._dynamo.reset()
+                with fresh_inductor_cache():
+                    with inductor_config.patch(
+                        {"test_configs.distort_benchmarking_result": "inverse"}
+                    ):
+                        compiled_fn2 = compile_fn(fn, backend)
+                        out2 = compiled_fn2(*args, **kwargs)
+
+                # Run 3: Fresh cache, random benchmark results
+                torch._dynamo.reset()
+                with fresh_inductor_cache():
+                    with inductor_config.patch(
+                        {"test_configs.distort_benchmarking_result": "random"}
+                    ):
+                        compiled_fn3 = compile_fn(fn, backend)
+                        out3 = compiled_fn3(*args, **kwargs)
+
+                # Bitwise identical
+                self.assertEqual(out1, out2, rtol=0, atol=0)
+                self.assertEqual(out2, out3, rtol=0, atol=0)
+
+        try:
+            self._run_with_expected_failure(
+                device_type, op.name, backend, "determinism", dtype, run_test
+            )
+        finally:
+            torch.use_deterministic_algorithms(prev_deterministic)
+
+    # =========================================================================
+    # Bitwise Equivalence with Eager Mode Tests
+    # =========================================================================
+
+    @onlyCUDA
+    @skipIfTorchDynamo("Test uses dynamo already")
+    @ops(llm_ops, allowed_dtypes=DTYPES)
+    @parametrize("backend", BACKENDS)
+    def test_eager_equivalence(self, device, dtype, op, backend):
+        """Test bitwise equivalence with eager execution."""
+        torch._dynamo.reset()
+        device_type = torch.device(device).type
+
+        samples = self._get_sample_inputs(op, device, dtype)
+        if not samples:
+            self.skipTest(f"No samples for {op.name}")
+
+        fn = op.get_op()
+
+        def run_test():
+            for sample in samples:
+                args = (sample.input,) + tuple(sample.args)
+                kwargs = sample.kwargs
+
+                # Eager reference
+                eager_out = fn(*args, **kwargs)
+
+                # Compiled output
+                compiled_fn = compile_fn(fn, backend)
+                compiled_out = compiled_fn(*args, **kwargs)
+
+                # Bitwise identical
+                self.assertEqual(eager_out, compiled_out, rtol=0, atol=0)
+
+        self._run_with_expected_failure(
+            device_type, op.name, backend, "eager_equivalence", dtype, run_test
+        )
+
+    # =========================================================================
+    # Exhaustive/Sampled Unary Ufunc Tests
+    # =========================================================================
+
+    @onlyCUDA
+    @skipIfTorchDynamo("Test uses dynamo already")
+    @ops(llm_unary_ops, allowed_dtypes=DTYPES)
+    @parametrize("backend", BACKENDS)
+    def test_unary_ufunc_numerical(self, device, dtype, op, backend):
+        """Test unary ufuncs with exhaustive (16-bit) or sampled (fp32) inputs.
+
+        For fp16 and bf16: exhaustively tests all 65536 possible bit patterns.
+        For fp32: tests 64k sampled random bit patterns.
+
+        Verifies bitwise equivalence between eager and compiled execution.
+        """
+        torch._dynamo.reset()
+        device_type = torch.device(device).type
+
+        fn = op.get_op()
+
+        # Generate test values based on dtype
+        if dtype in (torch.float16, torch.bfloat16):
+            # Exhaustive: all 65536 possible values
+            test_values = generate_exhaustive_16bit(dtype, device)
+        else:
+            # Sampled: 64k random bit patterns
+            test_values = generate_sampled_fp32(NUM_SAMPLES, device)
+
+        # Eager reference
+        eager_out = fn(test_values)
+
+        # Compiled output
+        compiled_fn = compile_fn(fn, backend)
+        compiled_out = compiled_fn(test_values)
+
+        def run_test():
+            # Bitwise identical
+            self.assertEqual(eager_out, compiled_out, rtol=0, atol=0)
+
+        self._run_with_expected_failure(
+            device_type, op.name, backend, "unary_numerical", dtype, run_test
+        )
+
+    # =========================================================================
+    # Sampled Binary Ufunc Tests
+    # =========================================================================
+
+    @onlyCUDA
+    @skipIfTorchDynamo("Test uses dynamo already")
+    @ops(llm_binary_ops, allowed_dtypes=DTYPES)
+    @parametrize("backend", BACKENDS)
+    def test_binary_ufunc_numerical(self, device, dtype, op, backend):
+        """Test binary ufuncs on 64k sampled value pairs.
+
+        For fp32: samples random 32-bit patterns.
+        For fp16/bf16: samples from all possible 16-bit values.
+
+        Verifies bitwise equivalence between eager and compiled execution.
+        """
+        torch._dynamo.reset()
+        device_type = torch.device(device).type
+
+        fn = op.get_op()
+
+        # Generate sampled pairs
+        x, y = generate_sampled_pairs(NUM_SAMPLES, dtype, device)
+
+        # Eager reference
+        eager_out = fn(x, y)
+
+        # Compiled output
+        compiled_fn = compile_fn(fn, backend)
+        compiled_out = compiled_fn(x, y)
+
+        def run_test():
+            # Bitwise identical
+            self.assertEqual(eager_out, compiled_out, rtol=0, atol=0)
+
+        self._run_with_expected_failure(
+            device_type, op.name, backend, "binary_numerical", dtype, run_test
+        )
+
+
+instantiate_device_type_tests(TestOpInfoProperties, globals(), except_for=["cpu"])
+
+if __name__ == "__main__":
+    run_tests()
diff --git a/test/inductor/test_torchinductor_strided_blocks.py b/test/inductor/test_torchinductor_strided_blocks.py
index f2c08596e4f3e..ec7961d7cde0b 100644
--- a/test/inductor/test_torchinductor_strided_blocks.py
+++ b/test/inductor/test_torchinductor_strided_blocks.py
@@ -2,11 +2,12 @@
 # ruff: noqa: F841
 import contextlib
 import dataclasses
+import functools
 import importlib
 import math
 import unittest
 from collections.abc import Callable
-from typing import Any, Optional, Union
+from typing import Any
 
 import torch
 import torch.utils._pytree as pytree
@@ -92,7 +93,7 @@ class BlockDescriptorTestBase(InductorTestCase):
     block_descriptor_constructor_str = "tl.make_block_ptr"
 
     def _discontiguous_tensor(
-        self, view_size: tuple[int, ...], device: Union[torch.device, str]
+        self, view_size: tuple[int, ...], device: torch.device | str
     ) -> torch.Tensor:
         """
         Create a padded tensor of the given size.
@@ -126,13 +127,13 @@ def _run_and_compare(
         self: InductorTestCase,
         func: Callable[..., Any],
         *args,
-        compile_kwargs: Optional[dict] = None,
-        expected_num_block_pointers: Optional[int] = None,
+        compile_kwargs: dict | None = None,
+        expected_num_block_pointers: int | None = None,
         expected_num_programs: int = 1,
         expected_num_triton_kernels: int = 1,
-        config_patches: Optional[dict] = None,
-        rtol: Optional[float] = None,
-        atol: Optional[float] = None,
+        config_patches: dict | None = None,
+        rtol: float | None = None,
+        atol: float | None = None,
     ):
         """
         Runs the module through Inductor, comparing to eager reference.
@@ -160,7 +161,7 @@ def flatten_tensors(tensors):
             }
             self.assertTrue(torch.allclose(ref, actual, **tol))
 
-        def count_code(substr: str, expected: Optional[int]):
+        def count_code(substr: str, expected: int | None):
             count = sum(prog.count(substr) for prog in code)
             if expected is not None:
                 self.assertEqual(count, expected)
@@ -247,8 +248,8 @@ def test_pointwise(
         self,
         full_size: tuple[int, ...],
         view_size: tuple[int, ...],
-        stride: Optional[tuple[int, ...]],
-        offset: Optional[int],
+        stride: tuple[int, ...] | None,
+        offset: int | None,
         require_block_ptr: bool,
         prefer_nd_tiling: bool,
     ):
@@ -552,7 +553,10 @@ def test_reduction(
 
         view = self._discontiguous_tensor(view_size, self.device)
 
-        if num_triton_kernels == 2 and config.triton.cooperative_reductions:
+        if num_triton_kernels == 2 and (
+            config.triton.cooperative_reductions
+            or config.triton.force_cooperative_reductions
+        ):
             # fewer kernels with cooperative reductions
             num_triton_kernels = 1
             num_block_pointers -= 2
@@ -812,6 +816,13 @@ def test_2d_reduction_odd_shapes(
         Tests 2D reduction kernels. These arise from "odd" shapes which are not
         expressible with a 1D block pointer.
         """
+        if reduction_op == torch.sum and (
+            config.triton.cooperative_reductions
+            or config.triton.force_cooperative_reductions
+        ):
+            num_triton_kernels = 1
+            num_block_pointers = 1
+
         if reduction_op == torch.sum and torch.version.hip is not None:
             view_size = (513, 513) if view_size == (129, 129) else view_size
         view = self._discontiguous_tensor(view_size, self.device)
@@ -881,6 +892,13 @@ def test_2d_welford_reduction(
         doesn't generate a block pointer. Since tiling welford reductions depends on
         the block pointer analysis, those cases would fall back to 1D.
         """
+        if (
+            config.triton.cooperative_reductions
+            or config.triton.force_cooperative_reductions
+        ):
+            expected_num_triton_kernels = 1
+            expected_num_block_pointers = 1
+
         if torch.version.hip is not None and expected_num_triton_kernels == 2:
             size = (256, 256)
         view = self._discontiguous_tensor(size, self.device)
@@ -913,11 +931,15 @@ def test_welford_non_block_pointer(
         view = self._discontiguous_tensor((259, 311), self.device)
 
         # We expect many block pointers for this one.
+        cooperative_reductions = (
+            config.triton.cooperative_reductions
+            or config.triton.force_cooperative_reductions
+        )
         result, (code,) = self._run_and_compare(
             torch.var_mean,
             view,
-            expected_num_block_pointers=6,
-            expected_num_triton_kernels=2,
+            expected_num_block_pointers=0 if cooperative_reductions else 6,
+            expected_num_triton_kernels=1 if cooperative_reductions else 2,
             config_patches={"triton.prefer_nd_tiling": True},
         )
 
@@ -1018,7 +1040,7 @@ def foo(*args):
 
         view_size = (5, 7)
         arg0 = self._discontiguous_tensor(view_size, self.device)
-        arg1 = torch.empty(view_size)
+        arg1 = torch.randn(view_size)
 
         # No guarantees on the number of kernels or pointers.
         result, (code,) = self._run_and_compare(
@@ -1056,6 +1078,104 @@ def test_enable_tiled_reductions(self, tile_reductions: bool):
         # Check the code for multiple Rn_BLOCK's
         self._assert_reduction_ndims(code, 2 if tile_reductions else 1)
 
+    # FIXME: fails for Triton CPU. Tiling does not contain YBLOCK.
+    @test_torchinductor.xfail_if_triton_cpu
+    @xfail_if_use_tensor_descriptor
+    def test_reduction_padded_output_tiling(self):
+        """
+        Test a [Y, X, R0] reduction with tiled output dimensions.
+        The key to elicit this test case is a padded output tensor.
+        """
+        x = torch.randn((9, 11, 2), device=self.device)
+
+        # We expect block pointers for the input and output.
+        result, (code,) = self._run_and_compare(
+            functools.partial(torch.amax, dim=-1),
+            x,
+            expected_num_block_pointers=2,
+            expected_num_triton_kernels=1,
+            config_patches={
+                "pad_outputs": True,
+                "padding_alignment_bytes": 32,
+                "padding_stride_threshold": 0,
+                "unroll_reductions_threshold": 1,
+                **tiled_reduction_config,
+            },
+        )
+
+        # Check the code for multiple output dims.
+        self._assert_pointwise_ndims(code, 2)
+        self._assert_reduction_ndims(code, 1)
+
+    @xfail_if_use_tensor_descriptor
+    @parametrize("unroll", (False, True))
+    def test_reduce_trailing_dims_discontiguous_input(self, unroll: bool):
+        """
+        Test a [Y, X, R0, R1] reduction where the input tensor is discontiguous, but we
+        only reduce over the last two dimensions.
+        """
+        view = self._discontiguous_tensor((7, 5, 3, 2), self.device)
+
+        # We expect block pointers for the inputs and output.
+        # Note there are more inputs if unrolled.
+        result, (code,) = self._run_and_compare(
+            functools.partial(
+                torch.amax,
+                dim=(-1, -2),
+            ),
+            view,
+            expected_num_block_pointers=7 if unroll else 2,
+            expected_num_triton_kernels=1,
+            config_patches={
+                "unroll_reductions_threshold": 1e4 if unroll else 1,
+                **tiled_reduction_config,
+            },
+        )
+
+        # Check the code for multiple pointwise dims.
+        self._assert_pointwise_ndims(code, 2)
+        self._assert_reduction_ndims(code, 0 if unroll else 2)
+
+    def test_2d_reduction_with_broadcast(self):
+        """
+        Tests 2D tiled reduction with a broadcasted 1D tensor.
+
+        This is a regression test for a bug where block pointers that only
+        advance in one reduction dimension (not both) would cause a KeyError
+        during codegen. The bug occurred because:
+        1. The broadcasted tensor only varies along the first reduction dimension (R0)
+        2. When building pointer_advancements, the block_ptr is only
+           added to R0_INDEX (non-zero advancement) but skipped for R1_INDEX
+           (zero/identity advancement)
+        3. During loop suffix generation, the code assumed if a block_ptr exists
+           in the outer loop's advancements, it must exist in the inner loop's too
+
+        The pattern: (x * y[:, None]).sum() where x is 2D and y is 1D.
+        """
+        # Use sizes that require looped (non-persistent) reductions
+        # to trigger 2D tiled reduction with R0 and R1 loops
+        M, N = 64, 128
+
+        def fn(x, y):
+            # y is 1D (M,), x is 2D (M, N)
+            # y[:, None] broadcasts to (M, N)
+            # The y block_ptr only advances with R0, not R1
+            return (x * y[:, None]).sum()
+
+        x = torch.randn(M, N, device=self.device)
+        y = torch.randn(M, device=self.device)
+
+        # This should compile without KeyError and produce correct results
+        result, (code,) = self._run_and_compare(
+            fn,
+            x,
+            y,
+            config_patches=tiled_reduction_config,
+        )
+
+        # Verify 2D reduction is used (R0_BLOCK and R1_BLOCK present)
+        self._assert_reduction_ndims(code, 2)
+
     def test_complex_reshape_block_ptr(self):
         def func(x, y):
             add_ = x + y
@@ -1345,7 +1465,7 @@ def test_boundary_check(self, block_multiple, ynumel_exceed_ygrid_size, include_
         class InputShape:
             x: int
             y: int
-            z: Optional[int] = None
+            z: int | None = None
 
             def to_list(self):
                 out = [self.y, self.x]
diff --git a/test/inductor/test_triton_extension_backend.py b/test/inductor/test_triton_extension_backend.py
index bb9004fab67ed..c3bd9bf14612c 100644
--- a/test/inductor/test_triton_extension_backend.py
+++ b/test/inductor/test_triton_extension_backend.py
@@ -2,9 +2,10 @@
 import functools
 import random
 import string
+import sys
 import unittest
 from pathlib import Path
-from typing import Any, Optional
+from typing import Any
 
 import torch
 import torch._dynamo
@@ -151,6 +152,17 @@ def foo(x):
         ).check("device_str='privateuseone'").run(code)
 
     def _register_custom_backend_with_heuristics(self, device):
+        path_to_ext_heuristics = str(
+            Path(__file__).parent / "extension_backends" / "triton"
+        )
+        # Add the path to sys.path in the parent process so that the
+        # ExtensionCachingAutotuner class (defined in extension_triton_heuristics)
+        # can be resolved when the compiled kernel is unpickled from the
+        # compile subprocess back into the parent process.
+        if path_to_ext_heuristics not in sys.path:
+            sys.path.append(path_to_ext_heuristics)
+            self.addCleanup(sys.path.remove, path_to_ext_heuristics)
+
         class ExtensionTritonKernel(codegen.triton.TritonKernel):
             @classmethod
             @functools.lru_cache(None)
@@ -158,13 +170,8 @@ def gen_common_triton_imports(cls) -> str:
                 default_imports = super().gen_common_triton_imports()
                 custom_imports = IndentedBuffer()
                 custom_imports.splice(default_imports)
-                path_to_ext_heuristics = (
-                    Path(__file__).parent / "extension_backends" / "triton"
-                )
 
-                custom_imports.splice(f"""
-                    import sys
-                    sys.path.append("{path_to_ext_heuristics}")
+                custom_imports.splice("""
                     import extension_triton_heuristics as triton_heuristics
                 """)
                 return custom_imports
@@ -186,13 +193,15 @@ def _get_triton_info_kernel_cls(cls) -> type[codegen.triton.TritonKernel]:
             @staticmethod
             def create(
                 is_subgraph: bool,
-                subgraph_name: Optional[str],
-                parent_wrapper: Optional[PythonWrapperCodegen],
-                partition_signatures: Optional[ir.GraphPartitionSignature] = None,
+                subgraph_name: str | None,
+                parent_wrapper: PythonWrapperCodegen | None,
+                partition_signatures: ir.GraphPartitionSignature | None = None,
             ):
                 if is_subgraph:
-                    assert subgraph_name is not None
-                    assert parent_wrapper is not None
+                    if subgraph_name is None:
+                        raise AssertionError
+                    if parent_wrapper is None:
+                        raise AssertionError
                     return PythonWrapperCodegen.create(
                         subgraph_name, parent_wrapper, partition_signatures
                     )
diff --git a/test/inductor/test_triton_helpers.py b/test/inductor/test_triton_helpers.py
index 8465209c22807..0767a0b946644 100644
--- a/test/inductor/test_triton_helpers.py
+++ b/test/inductor/test_triton_helpers.py
@@ -1,18 +1,21 @@
 # Owner(s): ["module: inductor"]
 
-"""Tests for exclusive_scan_decoupled_lookback_64 dtype fix.
-
-The fix ensures that `test_target` maintains consistent dtype in the
-`if flag == 2` branch by using `tl.full([], -1, index.dtype)` instead
-of the literal `-1` which would be int32.
-
-This is a compile-time check - Triton validates type consistency across
-all branches during compilation, so the test will fail to compile if
-the dtype mismatch exists.
+"""Tests for triton_helpers functions.
+
+Covers:
+- exclusive_scan_decoupled_lookback_64 dtype fix (D89705211): ensures
+  `test_target` maintains consistent dtype by using `tl.full([], -1, index.dtype)`
+  instead of the literal `-1`.
+- select_one bitcast fix for sub-32-bit dtypes (D93872067): ensures the
+  intermediate result from `tl.sum()` is truncated back to the original-width
+  integer type before the final bitcast, preventing size mismatch errors.
 """
 
 import torch
-from torch._inductor.runtime.triton_helpers import exclusive_scan_decoupled_lookback_64
+from torch._inductor.runtime.triton_helpers import (
+    exclusive_scan_decoupled_lookback_64,
+    select_one,
+)
 from torch._inductor.test_case import run_tests, TestCase
 from torch.testing._internal.inductor_utils import GPU_TYPE, HAS_GPU, requires_gpu
 
@@ -44,6 +47,19 @@ def test_kernel_exclusive_scan(
 
         tl.store(result_ptr, exclusive_prefix)
 
+    @triton.jit
+    def test_kernel_select_one(
+        x_ptr,
+        mask_ptr,
+        result_ptr,
+        BLOCK_SIZE: tl.constexpr,
+    ):
+        offsets = tl.arange(0, BLOCK_SIZE)
+        x = tl.load(x_ptr + offsets)
+        mask = tl.load(mask_ptr + offsets)
+        result = select_one(x, mask, dim=0)
+        tl.store(result_ptr, result)
+
 
 class ExclusiveScanDecoupledLookback64Test(TestCase):
     """Test cases for exclusive_scan_decoupled_lookback_64 dtype fix."""
@@ -97,6 +113,50 @@ def test_flag_2_branch_with_int32_index(self) -> None:
         torch.testing.assert_close(result, expected)
 
 
+class SelectOneTest(TestCase):
+    """Test cases for select_one bitcast fix with sub-32-bit dtypes.
+
+    The fix (D93872067) adds an intermediate .to(idtype) truncation before
+    the final bitcast in select_one. Without this fix, tl.sum() promotes
+    sub-32-bit unsigned integers (e.g. uint16) to int32, and the subsequent
+    bitcast from int32 to a 16-bit dtype fails with a size mismatch error.
+    """
+
+    def _run_select_one(self, dtype: torch.dtype) -> None:
+        device = torch.device(GPU_TYPE)
+        BLOCK_SIZE = 4
+
+        # Create input tensor and a one-hot mask selecting the element at index 2
+        x = torch.tensor([1.0, 2.0, 3.0, 4.0], dtype=dtype, device=device)
+        mask = torch.tensor([0, 0, 1, 0], dtype=torch.int32, device=device)
+        result = torch.zeros(1, dtype=dtype, device=device)
+
+        test_kernel_select_one[(1,)](x, mask, result, BLOCK_SIZE=BLOCK_SIZE)
+
+        expected = torch.tensor([3.0], dtype=dtype, device=device)
+        torch.testing.assert_close(result, expected)
+
+    @requires_gpu()
+    def test_select_one_bfloat16(self) -> None:
+        """Test select_one with bfloat16 (16-bit) — triggers the bitcast fix."""
+        self._run_select_one(torch.bfloat16)
+
+    @requires_gpu()
+    def test_select_one_float16(self) -> None:
+        """Test select_one with float16 (16-bit) — triggers the bitcast fix."""
+        self._run_select_one(torch.float16)
+
+    @requires_gpu()
+    def test_select_one_float32(self) -> None:
+        """Test select_one with float32 (32-bit) — baseline that always worked."""
+        self._run_select_one(torch.float32)
+
+    @requires_gpu()
+    def test_select_one_float64(self) -> None:
+        """Test select_one with float64 (64-bit) — baseline that always worked."""
+        self._run_select_one(torch.float64)
+
+
 if __name__ == "__main__":
     if HAS_GPU:
         run_tests()
diff --git a/test/inductor/test_triton_heuristics.py b/test/inductor/test_triton_heuristics.py
index d8ef6ed9fb6bd..b2e15c0ebcff8 100644
--- a/test/inductor/test_triton_heuristics.py
+++ b/test/inductor/test_triton_heuristics.py
@@ -23,7 +23,7 @@
     GPU_TYPE,
     HAS_GPU,
     HAS_GPU_AND_TRITON,
-    requires_cuda_with_enough_memory,
+    requires_gpu_with_enough_memory,
 )
 
 
@@ -45,12 +45,15 @@
 )
 from torch._inductor.runtime.triton_helpers import math as tl_math
 from torch._inductor.runtime.triton_heuristics import (
+    _check_max_grid_x,
+    _num_warps,
     autotune_hints_to_configs,
     CachingAutotuner,
     template,
     triton_config,
 )
 from torch._inductor.test_case import run_tests, TestCase
+from torch._inductor.utils import fresh_cache
 
 
 @triton.jit
@@ -129,11 +132,9 @@ def forward(primals_1, primals_2, primals_5):
         ]
         self.assertEqual(forward(*args), foo_c(*args))
 
-    # @skipIfXpu
     def test_artificial_zgrid(self):
         self._test_artificial_zgrid()
 
-    # @skipIfXpu
     @config.patch("cpp_wrapper", True)
     def test_artificial_grid_cpp_wrapper(self):
         self._test_artificial_zgrid()
@@ -179,7 +180,6 @@ def triton_(in_ptr0, out_ptr0, xnumel, XBLOCK: tl.constexpr):
             "inductor_meta": inductor_meta,
         }
 
-    # @skipIfXpu
     def test_pre_hook_assert(self):
         # assert if any of the configs passed to the CachingAutotuner have pre-hooks
         args = self._get_cos_kernel_caching_autotuner_args()
@@ -212,6 +212,8 @@ def mock_triton_config(
             num_stages=None,
             num_elements_per_warp=None,
             min_elem_per_thread=None,
+            *,
+            warp_size=32,
         ):
             seen_num_elements_per_warp.add(num_elements_per_warp)
             return None
@@ -273,7 +275,9 @@ def fn(x):
         res = torch.compile(fn)(x)
         self.assertEqual(ref, res)
 
-    @skipIfXpu(msg="https://github.com/intel/torch-xpu-ops/issues/2331")
+    @skipIfXpu(
+        msg="lack _get_exceeding_shared_memory_checker support - torch-xpu-ops: 2331"
+    )
     @skipUnless(HAS_GPU_AND_TRITON, "requires gpu and triton")
     @parametrize("do_pruning", [False, True])
     def test_prune_configs_over_shared_memory_limit(self, do_pruning):
@@ -304,6 +308,43 @@ def test_prune_configs_over_shared_memory_limit(self, do_pruning):
             )
             self.assertEqual(len(configs), expected_count)
 
+    @skipUnless(HAS_GPU_AND_TRITON, "requires gpu and triton")
+    def test_compile_time_autotune_not_repeated_at_runtime(self):
+        def fn(x):
+            return (x + 1).sum(dim=1)
+
+        x = torch.randn(2048, 2048, device=GPU_TYPE)
+        benchmark_calls = []
+
+        def fake_benchmark_all_configs(autotuner, *args, **kwargs):
+            benchmark_calls.append(autotuner.inductor_meta.get("kernel_name"))
+            return {
+                launcher: float(idx) for idx, launcher in enumerate(autotuner.launchers)
+            }
+
+        torch._dynamo.reset()
+        with (
+            fresh_cache(),
+            config.patch(
+                {
+                    "triton.autotune_at_compile_time": True,
+                    "max_autotune_pointwise": True,
+                    "compile_threads": 1,
+                }
+            ),
+            patch.object(
+                CachingAutotuner,
+                "benchmark_all_configs",
+                fake_benchmark_all_configs,
+            ),
+        ):
+            compiled = torch.compile(fn, backend="inductor")
+            self.assertEqual(compiled(x), fn(x))
+            self.assertEqual(len(benchmark_calls), 1)
+
+            self.assertEqual(compiled(x), fn(x))
+            self.assertEqual(len(benchmark_calls), 1)
+
 
 class TestArgumentCloneAndRestore(TestCase):
     # Our tensor is large enough. If a unexpected copy happens, the
@@ -336,7 +377,7 @@ def _do_test(self, gpu_tensor):
         old_storage_offset = gpu_tensor.storage_offset()
         gpu_tensor_clone = clone_preserve_strides(gpu_tensor)
 
-        peak_mem_before = torch.cuda.max_memory_allocated()
+        peak_mem_before = torch.get_device_module(GPU_TYPE).max_memory_allocated()
         cpu_copies = autotuner.copy_args_to_cpu_if_needed(gpu_tensor)
         self.assertTrue(len(cpu_copies) == 1)
 
@@ -363,21 +404,21 @@ def _do_test(self, gpu_tensor):
         # Avoid OOM in CI
         self.assertTrue(peak_mem_after < 1e10)
 
-    @requires_cuda_with_enough_memory(1e10)
+    @requires_gpu_with_enough_memory(1e10)
     def test_clone_contiguous_args(self):
         arg = self._create_tensor(pad=0)
         self.assertTrue(arg.is_contiguous())
         self.assertTrue(arg.storage_offset() == 0)
         self._do_test(arg)
 
-    @requires_cuda_with_enough_memory(1e10)
+    @requires_gpu_with_enough_memory(1e10)
     def test_clone_non_contiguous_args(self):
         arg = self._create_tensor(pad=1)
         self.assertFalse(arg.is_contiguous())
         self.assertTrue(arg.storage_offset() == 0)
         self._do_test(arg)
 
-    @requires_cuda_with_enough_memory(1e10)
+    @requires_gpu_with_enough_memory(1e10)
     def test_clone_args_with_non_zero_offset(self):
         arg = self._create_tensor(pad=1, with_offset=True)
         self.assertFalse(arg.is_contiguous())
@@ -496,6 +537,400 @@ def simple_model(x):
                 os.environ["TORCHINDUCTOR_DUMP_LAUNCH_TENSORS"] = old_dump_env
 
 
+class TestRecheckAutotuneCache(TestCase):
+    """Tests for CachingAutotuner.recheck_autotune_cache"""
+
+    @staticmethod
+    def _make_compile_result(cfg):
+        """Create a mock StaticTritonCompileResult with the given config."""
+        from torch._inductor.runtime.triton_heuristics import StaticTritonCompileResult
+
+        result = MagicMock(spec=StaticTritonCompileResult)
+        result.config = cfg
+        return result
+
+    @staticmethod
+    def _make_autotuner_with_results(configs, compile_results):
+        """
+        Create a CachingAutotuner and inject compile_results directly,
+        bypassing actual Triton compilation.
+        """
+        args = TestTritonHeuristics._get_cos_kernel_caching_autotuner_args()
+        args["configs"] = configs
+        autotuner = CachingAutotuner(**args)
+        autotuner.compile_results = compile_results
+        return autotuner
+
+    @skipUnless(HAS_GPU_AND_TRITON, "requires gpu and triton")
+    def test_recheck_single_config_enters_cache_hit_block(self):
+        """
+        When there is exactly 1 config and the autotune cache returns a hit,
+        recheck_autotune_cache should narrow compile_results to that single
+        matching result (not skip the block due to len(configs) == 1).
+        """
+        cfg = triton_config({"x": 16}, 64)
+        cfg.found_by_coordesc = False
+        compile_result = self._make_compile_result(cfg)
+
+        autotuner = self._make_autotuner_with_results([cfg], [compile_result])
+
+        # Cache returns the same config as the best
+        cached_cfg = triton_config({"x": 16}, 64)
+        cached_cfg.found_by_coordesc = True
+
+        with patch(
+            "torch._inductor.runtime.triton_heuristics.check_autotune_cache",
+            return_value=([cached_cfg], None, {"autotune_cache_state": "hit"}),
+        ):
+            autotuner.recheck_autotune_cache(reload_kernel_from_src=MagicMock())
+
+        # The compile_results should be narrowed to just the matching one
+        self.assertEqual(len(autotuner.compile_results), 1)
+        self.assertIs(autotuner.compile_results[0], compile_result)
+        # And found_by_coordesc must be propagated
+        self.assertTrue(autotuner.compile_results[0].config.found_by_coordesc)
+
+    @skipUnless(HAS_GPU_AND_TRITON, "requires gpu and triton")
+    def test_recheck_propagates_found_by_coordesc_true(self):
+        """
+        When the cached best config has found_by_coordesc=True,
+        it must be propagated to the compile result's config.
+        """
+        cfg_a = triton_config({"x": 16}, 64)
+        cfg_b = triton_config({"x": 256}, 64)
+        cfg_a.found_by_coordesc = False
+        cfg_b.found_by_coordesc = False
+        result_a = self._make_compile_result(cfg_a)
+        result_b = self._make_compile_result(cfg_b)
+
+        autotuner = self._make_autotuner_with_results(
+            [cfg_a, cfg_b], [result_a, result_b]
+        )
+
+        # Cache says cfg_b is the best, found via coordesc
+        cached_cfg = triton_config({"x": 256}, 64)
+        cached_cfg.found_by_coordesc = True
+
+        with patch(
+            "torch._inductor.runtime.triton_heuristics.check_autotune_cache",
+            return_value=([cached_cfg], None, {"autotune_cache_state": "hit"}),
+        ):
+            autotuner.recheck_autotune_cache(reload_kernel_from_src=MagicMock())
+
+        self.assertEqual(len(autotuner.compile_results), 1)
+        self.assertIs(autotuner.compile_results[0], result_b)
+        # The flag must be propagated from the cached config
+        self.assertTrue(autotuner.compile_results[0].config.found_by_coordesc)
+
+    @skipUnless(HAS_GPU_AND_TRITON, "requires gpu and triton")
+    def test_recheck_propagates_found_by_coordesc_false(self):
+        """
+        When the cached best config has found_by_coordesc=False, it must be
+        propagated so that coordinate descent can still run if enabled.
+        """
+        cfg = triton_config({"x": 16}, 64)
+        cfg.found_by_coordesc = True
+        compile_result = self._make_compile_result(cfg)
+
+        autotuner = self._make_autotuner_with_results([cfg], [compile_result])
+
+        cached_cfg = triton_config({"x": 16}, 64)
+        cached_cfg.found_by_coordesc = False
+
+        with patch(
+            "torch._inductor.runtime.triton_heuristics.check_autotune_cache",
+            return_value=([cached_cfg], None, {"autotune_cache_state": "hit"}),
+        ):
+            autotuner.recheck_autotune_cache(reload_kernel_from_src=MagicMock())
+
+        self.assertEqual(len(autotuner.compile_results), 1)
+        self.assertFalse(autotuner.compile_results[0].config.found_by_coordesc)
+
+    @skipUnless(HAS_GPU_AND_TRITON, "requires gpu and triton")
+    def test_recheck_no_cache_hit_leaves_results_unchanged(self):
+        """
+        When there's no autotune cache hit, compile_results should not change.
+        """
+        cfg_a = triton_config({"x": 16}, 64)
+        cfg_b = triton_config({"x": 256}, 64)
+        result_a = self._make_compile_result(cfg_a)
+        result_b = self._make_compile_result(cfg_b)
+
+        autotuner = self._make_autotuner_with_results(
+            [cfg_a, cfg_b], [result_a, result_b]
+        )
+
+        # Cache returns no hit (empty list)
+        with patch(
+            "torch._inductor.runtime.triton_heuristics.check_autotune_cache",
+            return_value=([], None, {"autotune_cache_state": "miss"}),
+        ):
+            autotuner.recheck_autotune_cache(reload_kernel_from_src=MagicMock())
+
+        # Nothing should change
+        self.assertEqual(len(autotuner.compile_results), 2)
+
+
+@triton.jit
+def hip_autotune_kernel(
+    in_ptr,
+    out_ptr,
+    numel,
+    XBLOCK: tl.constexpr,
+):
+    pid = tl.program_id(0)
+    offsets = pid * XBLOCK + tl.arange(0, XBLOCK)
+    mask = offsets < numel
+    data = tl.load(in_ptr + offsets, mask=mask)
+    tl.store(out_ptr + offsets, data * 2.0, mask=mask)
+
+
+@functools.lru_cache
+def get_hip_autotune_kernel_with_invalid_config():
+    # num_warps=32 with HIP warp_size=64 = 2048 threads, exceeds 1024 limit
+    return triton.autotune(
+        configs=[
+            triton.Config({"XBLOCK": 128}, num_warps=32, num_stages=1),  # invalid
+            triton.Config({"XBLOCK": 256}, num_warps=1, num_stages=1),  # valid
+        ],
+        key=["numel"],
+    )(hip_autotune_kernel)
+
+
+class TestHIPInvalidConfigHandling(TestCase):
+    @runOnRocm
+    @skipUnless(HAS_GPU_AND_TRITON, "requires GPU and Triton")
+    def test_benchmark_returns_inf_on_invalid_config(self):
+        from torch._inductor.runtime.benchmarking import TritonBenchmarker
+
+        benchmarker = TritonBenchmarker()
+
+        def failing_callable():
+            raise RuntimeError(
+                "Triton Error [HIP]: Code: 9, Message: invalid configuration argument"
+            )
+
+        result = benchmarker.benchmark(
+            fn=failing_callable,
+            device=GPU_TYPE,
+            is_vetted_benchmarking=True,
+        )
+        self.assertEqual(result, float("inf"))
+
+    @runOnRocm
+    @skipUnless(HAS_GPU_AND_TRITON, "requires GPU and Triton")
+    def test_autotune_skips_invalid_hip_config_and_succeeds(self):
+        numel = 1024 * 1024
+        x = torch.randn(numel, device=GPU_TYPE, dtype=torch.float32)
+        y = torch.empty_like(x)
+
+        kernel = get_hip_autotune_kernel_with_invalid_config()
+
+        def grid(meta):
+            return (triton.cdiv(numel, meta["XBLOCK"]),)
+
+        kernel[grid](x, y, numel)
+
+        expected = x * 2.0
+        torch.testing.assert_close(y, expected)
+
+
+class TestGridExprMaximum(TestCase):
+    def test_maximum_cpp_mode_casts_int_constants_to_long(self):
+        from torch._inductor.runtime.triton_heuristics import Grid1D
+
+        grid = Grid1D(inductor_meta={}, mode="cpp")
+        # Mixed str/int: int constants must be cast to (long) for std::max
+        result = grid.maximum(["ynumel_0", "ynumel_1", 4480])
+        self.assertIn("(long)4480", result)
+        self.assertIn("std::max", result)
+        # All strings: no cast needed
+        result = grid.maximum(["xnumel", "ynumel"])
+        self.assertNotIn("(long)", result)
+        # All ints: constant-folds
+        self.assertEqual(grid.maximum([10, 20, 5]), 20)
+
+
+class TestGrid2DWithYZOverflowZeroYnumel(TestCase):
+    """Regression test for https://github.com/pytorch/pytorch/issues/178530"""
+
+    def test_grid2d_yz_overflow_zero_ynumel_python(self):
+        from torch._inductor.runtime.triton_heuristics import Grid2DWithYZOverflow
+
+        grid = Grid2DWithYZOverflow(inductor_meta={}, mode="python")
+        grid.generate({"XBLOCK": 128, "YBLOCK": 128})
+        # ynumel=0 must not raise ZeroDivisionError
+        x, y, z = grid.eval_slow(
+            {"xnumel": 256, "ynumel": 0, "XBLOCK": 128, "YBLOCK": 128}
+        )
+        self.assertEqual(y, 0)
+        self.assertEqual(z, 0)
+
+    def test_grid2d_yz_overflow_zero_ynumel_cpp(self):
+        from torch._inductor.runtime.triton_heuristics import Grid2DWithYZOverflow
+
+        grid = Grid2DWithYZOverflow(inductor_meta={}, mode="cpp")
+        grid.generate({"XBLOCK": 128, "YBLOCK": 128})
+        # cpp mode: the generated expression should contain a zero-guard
+        self.assertIn("== 0", str(grid.y_grid))
+
+    def test_grid2d_yz_overflow_nonzero_ynumel_unchanged(self):
+        from torch._inductor.runtime.triton_heuristics import Grid2DWithYZOverflow
+
+        grid = Grid2DWithYZOverflow(inductor_meta={}, mode="python")
+        grid.generate({"XBLOCK": 128, "YBLOCK": 128})
+        # Normal case: ynumel > 0 still works correctly
+        x, y, z = grid.eval_slow(
+            {"xnumel": 256, "ynumel": 256, "XBLOCK": 128, "YBLOCK": 128}
+        )
+        self.assertEqual(x, 2)
+        self.assertEqual(y, 2)
+        self.assertEqual(z, 1)
+
+    def test_grid2d_yz_overflow_large_ynumel(self):
+        from torch._inductor.runtime.triton_heuristics import Grid2DWithYZOverflow
+
+        grid = Grid2DWithYZOverflow(inductor_meta={}, mode="python")
+        grid.generate({"XBLOCK": 128, "YBLOCK": 128})
+        # Large ynumel that requires overflow splitting across y and z
+        x, y, z = grid.eval_slow(
+            {"xnumel": 128, "ynumel": 128 * 131070, "XBLOCK": 128, "YBLOCK": 128}
+        )
+        self.assertEqual(x, 1)
+        # y * z must cover all y blocks
+        self.assertGreaterEqual(y * z, 131070)
+
+
+
+class TestWarpSizeUnification(TestCase):
+    """Tests for the unified warp_size threading through config helpers."""
+
+    def test_warp_size_or_default(self):
+        none_props = DeviceProperties(
+            type="cuda", index=0, multi_processor_count=80, cc=80, warp_size=None
+        )
+        with self.assertRaisesRegex(
+            RuntimeError, "cuda device properties must report warp_size"
+        ):
+            none_props.warp_size_or_default
+
+        cpu_props = DeviceProperties(
+            type="cpu", index=0, multi_processor_count=80, cc=80, warp_size=None
+        )
+        self.assertEqual(cpu_props.warp_size_or_default, 32)
+
+        w32 = DeviceProperties(
+            type="cuda", index=0, multi_processor_count=80, cc=80, warp_size=32
+        )
+        self.assertEqual(w32.warp_size_or_default, 32)
+
+        w64 = DeviceProperties(
+            type="hip", index=0, multi_processor_count=80, cc=80, warp_size=64
+        )
+        self.assertEqual(w64.warp_size_or_default, 64)
+
+    def test_num_warps_halves_on_wave64(self):
+        # wave64 (AMD CDNA/gfx9) halves the range so total threads match wave32.
+        self.assertEqual(_num_warps(8, max_num_warps=8, warp_size=64), 4)
+
+    def test_num_warps_does_not_halve_on_wave32(self):
+        # wave32 (NVIDIA and AMD RDNA) must not halve.
+        self.assertEqual(_num_warps(8, max_num_warps=8, warp_size=32), 8)
+        # Default warp_size is 32; confirm the default path matches wave32.
+        self.assertEqual(_num_warps(8, max_num_warps=8), 8)
+
+    def test_check_max_grid_x_respects_warp_size_hip(self):
+        # _check_max_grid_x only uses warp_size on the HIP path, where the
+        # bound is on total threads (num_blocks * num_warps * warp_size).
+        # Force the HIP branch so the assertion exercises that path even on
+        # a non-HIP build.
+        size_hints = {"x": 2**32}
+        with patch.object(torch.version, "hip", "6.0.0"):
+            x32, _ = _check_max_grid_x(size_hints, x=1, num_warps=8, warp_size=32)
+            x64, _ = _check_max_grid_x(size_hints, x=1, num_warps=8, warp_size=64)
+        # Doubling warp_size halves the per-block thread budget, so x must
+        # double to keep total threads under 2**31 - 1.
+        self.assertEqual(x64, 2 * x32)
+        self.assertLessEqual(x32, TRITON_MAX_BLOCK["X"])
+        self.assertLessEqual(x64, TRITON_MAX_BLOCK["X"])
+
+    def test_check_max_grid_x_ignores_warp_size_on_nvidia(self):
+        # NVIDIA bounds num_blocks directly, so warp_size must not change x.
+        size_hints = {"x": 2**32}
+        with patch.object(torch.version, "hip", None):
+            x32, _ = _check_max_grid_x(size_hints, x=1, num_warps=8, warp_size=32)
+            x64, _ = _check_max_grid_x(size_hints, x=1, num_warps=8, warp_size=64)
+        self.assertEqual(x32, x64)
+
+    def test_triton_config_uses_warp_size_for_min_elem(self):
+        # min_elem_per_thread * warp_size * num_warps drives the floor for
+        # block_size; doubling warp_size should at least double the XBLOCK.
+        size_hints = {"x": 4096}
+        cfg32 = triton_config(
+            size_hints, 128, num_warps=1, min_elem_per_thread=4, warp_size=32
+        )
+        cfg64 = triton_config(
+            size_hints, 128, num_warps=1, min_elem_per_thread=4, warp_size=64
+        )
+        self.assertEqual(cfg64.kwargs["XBLOCK"], 2 * cfg32.kwargs["XBLOCK"])
+
+
+class TestMakeLaunchersMemory(TestCase):
+    def test_failed_config_exception_not_retained(self):
+        """A failed launcher build must not retain its traceback frame chain."""
+        import contextlib
+        import gc
+        import types
+        import weakref
+
+        class _Result:
+            config = types.SimpleNamespace(num_stages=1, kwargs={})
+
+            def __init__(self, succeeds):
+                self.succeeds = succeeds
+
+            def make_launcher(self):
+                if self.succeeds:
+                    return object()
+                raise torch.cuda.OutOfMemoryError("out of resource (test)")
+
+        results = [_Result(True), _Result(False)]
+        fake_self = types.SimpleNamespace(
+            launchers=[],
+            compile_results=results,
+            triton_meta={"device": 0},
+            inductor_meta={},
+            get_device_interface=lambda: None,
+        )
+
+        class _Sentinel:
+            pass
+
+        def run():
+            l2_buffer = _Sentinel()
+            ref = weakref.ref(l2_buffer)
+            with patch(
+                "torch._dynamo.device_interface.DeviceGuard",
+                lambda *args, **kwargs: contextlib.nullcontext(),
+            ):
+                CachingAutotuner._make_launchers(fake_self)
+            return ref
+
+        gc.disable()
+        try:
+            ref = run()
+            self.assertIsNone(
+                ref(),
+                "_make_launchers retained the failed-config exception; its "
+                "traceback pins caller frames until cyclic GC.",
+            )
+        finally:
+            gc.enable()
+
+        self.assertEqual(len(fake_self.launchers), 1)
+
+
+
 if __name__ == "__main__":
     if IS_LINUX and HAS_GPU:
         run_tests()
diff --git a/test/inductor/test_triton_kernels.py b/test/inductor/test_triton_kernels.py
index 4b21bfb8692f4..7f5e1faf0af64 100644
--- a/test/inductor/test_triton_kernels.py
+++ b/test/inductor/test_triton_kernels.py
@@ -2,9 +2,12 @@
 # ruff: noqa: F841
 # flake8: noqa: E731
 # Skip do not assign a lambda expression, use a def
+import contextlib
 import functools
 import logging
 import os
+import sys
+from unittest import mock
 
 import torch
 import torch._dynamo.testing
@@ -28,6 +31,7 @@
 from torch.testing._internal import common_utils
 from torch.testing._internal.common_utils import parametrize, skipIfWindows, skipIfXpu
 from torch.testing._internal.inductor_utils import (
+    get_func_call,
     GPU_TYPE,
     HAS_CUDA_AND_TRITON,
     HAS_GPU,
@@ -44,6 +48,19 @@
 )
 
 
+@contextlib.contextmanager
+def _dump_launch_params(value: str):
+    launch_params_file = f"{os.path.abspath(sys.argv[0])}.launch_params"
+    if value == "1" and os.path.exists(launch_params_file):
+        os.remove(launch_params_file)
+    with mock.patch.dict(os.environ, {"TORCHINDUCTOR_DUMP_LAUNCH_PARAMS": value}):
+        try:
+            yield
+        finally:
+            if value == "1" and os.path.exists(launch_params_file):
+                os.remove(launch_params_file)
+
+
 if HAS_GPU:
     import triton
     from triton import language as tl
@@ -83,11 +100,20 @@ def _triton_get_ast_equal_to_str(params):
     USE_TF32 = torch.backends.cuda.matmul.fp32_precision == "tf32"
 
     if hasattr(triton, "constexpr_function"):
-
+        # Helper functions for triton kernels must be in globals.
         @triton.constexpr_function
         def log2(n):
             return len(bin(n)) - 3
 
+        _get_int_dtype_test = triton.constexpr_function(tl.core.get_int_dtype)
+
+        @triton.jit
+        def _dtype_helper_test(x):
+            idtype = _get_int_dtype_test(
+                bitwidth=x.dtype.primitive_bitwidth, signed=True
+            )
+            return x.to(idtype, bitcast=True)
+
 
 class KernelTests(torch._inductor.test_case.TestCase):
     def _kernel_launched_in_code(self, kernel_name: str, code: str) -> bool:
@@ -113,6 +139,10 @@ def f(x):
 
     @requires_gpu
     def test_triton_kernel_ill_formed(self):
+        if inductor_config.cpp_wrapper:
+            # With lazy compile, kernel errors happen at runtime, not codegen
+            self.skipTest("kernel errors happen at runtime instead of codegen time")
+
         @triton.jit
         def ill_formed_kernel(kernel):
             off_n = tl.program_id(0)
@@ -202,13 +232,13 @@ def test_triton_kernel_higher_order_func(self):
 
     @requires_gpu
     def test_triton_kernel_functionalize(self):
-        from functorch import make_fx
         from torch._higher_order_ops.triton_kernel_wrap import kernel_side_table
         from torch._subclasses.functional_tensor import (
             CppFunctionalizeAPI,
             FunctionalTensorMode,
             PythonFunctionalizeAPI,
         )
+        from torch.fx.experimental.proxy_tensor import make_fx
 
         kernel_side_table.reset_table()
 
@@ -326,10 +356,11 @@ def prep():
 
     @requires_gpu
     def test_triton_kernel_clone_wekdeps(self):
-        from functorch import make_fx
         from torch._higher_order_ops.triton_kernel_wrap import kernel_side_table
         from torch._inductor.choices import InductorChoices
         from torch._inductor.compile_fx import compile_fx_inner
+        from torch._inductor.virtualized import V
+        from torch.fx.experimental.proxy_tensor import make_fx
 
         TENSOR_SIZE = 30_000_000
         BLOCK_SIZE = 1024
@@ -392,7 +423,7 @@ def can_fuse(self, scheduler, node1, node2, shared_data_score) -> bool:
 
         gm = make_fx(f, tracing_mode="fake")(t)
 
-        with inductor_config.patch({"inductor_choices_class": NoFusionChoices}):
+        with V.set_choices_handler(NoFusionChoices()):
             log_stream, ctx = logs_to_string("torch._inductor.codecache", "output_code")
             with ctx():
                 compiled_gm = compile_fx_inner(gm, [t])
@@ -869,7 +900,8 @@ def call_triton(
         prev_c = CONSTANT_C
         # If the behavior of triton kernels change, this test will fail
         CONSTANT_C = tl.constexpr(10)
-        assert CONSTANT_C != prev_c
+        if CONSTANT_C == prev_c:
+            raise AssertionError
 
         t = torch.randn(5, device=GPU_TYPE)
         torch_result = call_triton(t)
@@ -941,10 +973,6 @@ def call_triton(x: torch.Tensor, y: torch.Tensor):
     @common_utils.parametrize("grid_type", [1, 2, 3])
     @common_utils.parametrize("tdlp", ["0", "1"])
     def test_triton_kernel_2d_autotune(self, grad, dynamic, backend, grid_type, tdlp):
-        import os
-
-        os.environ["TORCHINDUCTOR_DUMP_LAUNCH_PARAMS"] = tdlp
-
         def call_triton(x: torch.Tensor, y: torch.Tensor, output: torch.Tensor):
             x_elements = output.size()[0]
             y_elements = output.size()[1]
@@ -973,11 +1001,13 @@ def grid_fn(meta):
         output = torch.zeros_like(t1, requires_grad=grad)
 
         torch_result = call_triton(t1, t2, output)
-        compiled_func = torch.compile(
-            call_triton, backend=backend, fullgraph=True, dynamic=dynamic
-        )
-        output2 = torch.zeros_like(t1, requires_grad=grad)
-        self.assertEqual(compiled_func(t1, t2, output2), torch_result)
+
+        with _dump_launch_params(tdlp):
+            compiled_func = torch.compile(
+                call_triton, backend=backend, fullgraph=True, dynamic=dynamic
+            )
+            output2 = torch.zeros_like(t1, requires_grad=grad)
+            self.assertEqual(compiled_func(t1, t2, output2), torch_result)
 
     @requires_gpu
     @common_utils.parametrize("dynamic", [False, True])
@@ -1034,7 +1064,7 @@ def grid_fn(meta):
             functools.partial(call_triton_add, grid_type=2, num=200, autotuned=True),
             functools.partial(call_triton_add, grid_type=3, autotuned=True),
         ]
-        from functorch import make_fx
+        from torch.fx.experimental.proxy_tensor import make_fx
 
         tracing_mode = "symbolic" if dynamic else "fake"
 
@@ -1152,6 +1182,13 @@ def f(x):
             else f"empty_strided_{GPU_TYPE}((10, ), (1, ), torch.float32)"
         )
         num_bufs_allocated = code.count(code_string)
+        if (
+            inductor_config.cpp_wrapper
+            and inductor_config.triton.autotune_at_compile_time is not True
+        ):
+            # Lazy compile emits aoti_torch_empty_strided for scratch space
+            # allocation (global_scratch + profile_scratch) per unique kernel wrapper
+            num_bufs_allocated -= 4
         self.assertEqual(num_bufs_allocated, 2)
 
         # Check we're reusing buffers if not allocating.
@@ -1417,8 +1454,6 @@ def f(x):
     @common_utils.parametrize("dump_launch_params", ["0", "1"])
     @common_utils.parametrize("dynamic", [False, True])
     def test_triton_kernel_equal_to_1_arg(self, dynamic, dump_launch_params):
-        os.environ["TORCHINDUCTOR_DUMP_LAUNCH_PARAMS"] = dump_launch_params
-
         @triton.jit
         def add_kernel_half_n_elements(
             in_ptr0,
@@ -1447,23 +1482,25 @@ def f(x, y):
         x = torch.randn(2, device=GPU_TYPE)
         y = torch.randn(2, device=GPU_TYPE)
         eager_out = f(x, y)
-        compiled_out, sources = run_and_get_code(
-            torch.compile(f, dynamic=dynamic), x, y
-        )
 
-        if triton_version_uses_attrs_dict():
-            self.assertFalse("equal_to" in sources[0])
-        else:
-            if dynamic:
-                # when half_n_elements passed to the Triton kernel is
-                # dynamic, equal_to_1 specialization can't be enforced
+        with _dump_launch_params(dump_launch_params):
+            compiled_out, sources = run_and_get_code(
+                torch.compile(f, dynamic=dynamic), x, y
+            )
 
-                # also, equal_to_1 specialization doesn't occur (or appear in the signature)
-                # for newer versions of triton (i.e. the ones where triton_version_uses_attrs_dict() == True)
-                self.assertTrue(_triton_get_ast_equal_to_str(()) in sources[0])
+            if triton_version_uses_attrs_dict():
+                self.assertFalse("equal_to" in sources[0])
             else:
-                self.assertTrue(_triton_get_ast_equal_to_str((3,)) in sources[0])
-        self.assertEqual(compiled_out, eager_out)
+                if dynamic:
+                    # when half_n_elements passed to the Triton kernel is
+                    # dynamic, equal_to_1 specialization can't be enforced
+
+                    # also, equal_to_1 specialization doesn't occur (or appear in the signature)
+                    # for newer versions of triton (i.e. the ones where triton_version_uses_attrs_dict() == True)
+                    self.assertTrue(_triton_get_ast_equal_to_str(()) in sources[0])
+                else:
+                    self.assertTrue(_triton_get_ast_equal_to_str((3,)) in sources[0])
+            self.assertEqual(compiled_out, eager_out)
 
     @requires_gpu
     @common_utils.parametrize("dynamic", [False, True])
@@ -1559,6 +1596,50 @@ def f(x):
         self.assertIn("@triton.constexpr_function", triton_code)
         self.assertEqual(compiled_out, eager_out)
 
+    @unittest.skipIf(
+        not HAS_GPU or not hasattr(triton, "constexpr_function"),
+        "newer triton version required",
+    )
+    def test_triton_kernel_with_constexpr_dtype_annotations(self):
+        """
+        Test that constexpr functions with dtype type annotations work correctly.
+        This tests the fix for proper handling of:
+        1. Type annotations using triton.language.core.dtype
+        2. Dtype instances (int8, uint8, etc.) that lack __name__ attribute
+        3. Function name aliasing
+        """
+
+        @triton.jit
+        def kernel_with_dtype_annotation(out_ptr, n: tl.constexpr):
+            offs = tl.arange(0, n)
+            x = tl.full([n], 1.0, dtype=tl.float32)
+            y = _dtype_helper_test(x)
+            tl.store(out_ptr + offs, y.to(tl.float32, bitcast=True))
+
+        def f(n):
+            out = torch.empty(n, device=GPU_TYPE)
+            kernel_with_dtype_annotation[(1,)](out, n)
+            return out
+
+        n = 8
+        eager_out = f(n)
+        compiled_out, (triton_code,) = run_and_get_code(
+            torch.compile(f, fullgraph=True), n
+        )
+
+        # Verify the generated code has proper imports
+        self.assertIn("from triton.language.core import dtype as dtype", triton_code)
+        self.assertIn("@triton.constexpr_function", triton_code)
+        self.assertIn("_get_int_dtype_test = get_int_dtype", triton_code)
+        # Verify dtype instances are emitted correctly
+        self.assertIn("int8 = tl.int8", triton_code)
+        self.assertIn("@triton.jit", triton_code)
+        self.assertIn("def _dtype_helper_test", triton_code)
+
+        # Verify correctness
+        self.assertEqual(compiled_out, eager_out)
+        self.assertTrue(torch.all(compiled_out == 1.0).item())
+
     @requires_gpu
     def test_triton_kernel_with_imported_symbol_with_custom_name(self):
         @triton.jit
@@ -1768,6 +1849,11 @@ def f(x, y):
     @common_utils.parametrize("autotune", [False, True])
     @common_utils.parametrize("backend", ["eager", "aot_eager", "inductor"])
     def test_triton_kernel_special_params(self, autotune, backend):
+        if backend == "inductor" and inductor_config.cpp_wrapper:
+            raise unittest.SkipTest(
+                "cpp-wrapper doesn't support num_warps/num_stages as constexpr"
+            )
+
         @triton.jit
         def special_params_kernel(
             in_ptr,
@@ -1839,7 +1925,6 @@ def test_on_device_tma(self, dynamic, tma_version):
             self.skipTest("requires triton.tools.tensor_descriptor TMA support")
         if tma_version == "old" and not has_triton_experimental_host_tma():
             self.skipTest("requires triton.tools.experimental_descriptor TMA support")
-
         kernel = (
             add_kernel_on_device_tma_new_api
             if tma_version == "new"
@@ -2212,10 +2297,54 @@ def f(a):
 
         # 2 calls: one for two inputs (dedupped), one for the output
         if tma_version == "new":
-            self.assertEqual(code.count("TensorDescriptor.from_tensor("), 2)
+            if (
+                not inductor_config.cpp_wrapper
+                or inductor_config.triton.autotune_at_compile_time
+            ):
+                # Lazy kernel compilation implicitly calls TensorDescriptor.from_tensor
+                # when calling run_triton_kernel_with_autotune
+                self.assertEqual(code.count("TensorDescriptor.from_tensor("), 2)
         else:
             self.assertEqual(code.count("create_1d_tma_descriptor("), 2)
 
+    @requires_gpu
+    def test_wrap_tma_args_skips_constexpr(self):
+        if not has_triton_tensor_descriptor_host_tma():
+            self.skipTest("requires triton.tools.tensor_descriptor TMA support")
+
+        from unittest.mock import Mock
+
+        from triton.tools.tensor_descriptor import TensorDescriptor
+
+        from torch._inductor.runtime.triton_lazy_compile import _wrap_tma_args
+
+        # Simulate a signature where constexpr entries appear before TMA
+        # descriptor params. The signature includes constexpr entries but
+        # runtime args exclude them, so _wrap_tma_args must account for the
+        # index offset.
+        kernel_fn = Mock()
+        kernel_fn.triton_meta = {
+            "signature": {
+                "ADD_ALPHA": "constexpr",
+                "in_desc_ptr0": "tensordesc<fp32[256]>",
+                "in_desc_ptr1": "tensordesc<fp32[256]>",
+                "out_desc_ptr": "tensordesc<fp32[256]>",
+                "BLOCK_SIZE": "constexpr",
+            }
+        }
+
+        t0 = torch.randn(256, device=GPU_TYPE)
+        t1 = torch.randn(256, device=GPU_TYPE)
+        t2 = torch.zeros(256, device=GPU_TYPE)
+        args = [t0, t1, t2]
+
+        wrapped = _wrap_tma_args(args, kernel_fn)
+
+        self.assertEqual(len(wrapped), 3)
+        self.assertIsInstance(wrapped[0], TensorDescriptor)
+        self.assertIsInstance(wrapped[1], TensorDescriptor)
+        self.assertIsInstance(wrapped[2], TensorDescriptor)
+
     @requires_gpu
     @common_utils.parametrize("dynamic", [False, True])
     @common_utils.parametrize("backend", ["eager", "aot_eager"])
@@ -2690,7 +2819,10 @@ def fn(sz):
         self.assertEqual(actual, expected)
 
     @requires_gpu
-    @skipIfXpu(msg="`tl.inline_asm_elementwise` is not yet supported on Intel GPUs")
+    @skipIfXpu(
+        msg="`tl.inline_asm_elementwise` is not yet supported on Intel GPUs, "
+        "https://github.com/pytorch/pytorch/pull/167786"
+    )
     @inductor_config.patch({"triton.autotune_at_compile_time": True})
     @parametrize("quotes", ["single", "double"])
     def test_kernel_inline_asm(self, quotes):
@@ -2767,15 +2899,62 @@ def fn(a, b):
                 ) from e
             raise
 
+    @requires_gpu
+    def test_constexpr_handling(self):
+        @triton.jit
+        def copy_kernel(
+            src_ptr,
+            dst_ptr,
+            n_elements,
+            stride,
+            maybe_param,
+            BLOCK_SIZE: tl.constexpr,
+        ):
+            pid = tl.program_id(0)
+            offs = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+            mask = offs < n_elements
+            x = tl.load(src_ptr + offs * stride, mask=mask)
+            scale = tl.where(maybe_param != 0, 0.5, 1.0)
+            x = x * scale
+
+            tl.store(dst_ptr + offs * stride, x, mask=mask)
+
+        t = torch.randn(1024, device=GPU_TYPE)
+        out = torch.empty(1024, device=GPU_TYPE)
+
+        kwargs = {
+            "src_ptr": t,
+            "dst_ptr": out,
+            "n_elements": 1024,
+            "stride": 1,
+            "maybe_param": None,  # semantically wrong, but testing frontend specialization
+            "BLOCK_SIZE": 256,
+        }
+
+        ttir_module, _ = generate_ttir(copy_kernel, kwargs, tma_descriptor_metadata={})
+        ttir_str = str(ttir_module)
+
+        # `constexpr` and None values get inlined, and do not appear as function parameters.
+        self.assertIn("src_ptr", ttir_str)
+        self.assertIn("dst_ptr", ttir_str)
+        self.assertIn("n_elements", ttir_str)
+        self.assertIn("stride", ttir_str)
+        self.assertNotIn("BLOCK_SIZE", ttir_str)
+        self.assertNotIn("maybe_param", ttir_str)
+
 
 def make_mutation_test(fn):
     @requires_gpu
     def test_fn(self):
-        from torch._higher_order_ops.triton_kernel_wrap import identify_mutated_tensors
+        from torch._higher_order_ops.triton_kernel_wrap import identify_accessed_tensors
 
         kernel, inputs, tma_descriptor_metadata, outputs = fn()
+        tensor_accesses = identify_accessed_tensors(
+            kernel, inputs, tma_descriptor_metadata
+        )
+        mutated_tensor_names = [dep.name for dep in tensor_accesses.read_writes.writes]
         self.assertListEqual(
-            identify_mutated_tensors(kernel, inputs, tma_descriptor_metadata),
+            mutated_tensor_names,
             outputs,
         )
 
@@ -2798,6 +2977,26 @@ def helper_add_and_out(x, y, out_ptr):
 class MutationTests(torch._inductor.test_case.TestCase):
     # Tests injected below
 
+    # Test that a scalar args are not flagged as mutated when passed
+    # to a tt.call op.
+    @make_mutation_test
+    def test_scalar_via_nested_write():
+        @triton.jit
+        def inner(ptr, scalar_offset):
+            tl.store(ptr + scalar_offset, 1.0)
+
+        @triton.jit
+        def outer(out_ptr, n_elements):
+            inner(out_ptr, n_elements)
+
+        t = torch.randn(4)
+        return (
+            outer,
+            {"out_ptr": t, "n_elements": 4},
+            {},
+            ["out_ptr"],
+        )
+
     # Regression test for #169782
     @make_mutation_test
     def test_with_none_arg():
@@ -3150,7 +3349,7 @@ def add_4_times_kernel(
             in_ptr0,
             in_ptr1,
             out_ptr,
-            n_elements,
+            n_elements: "tl.constexpr",
             BLOCK_SIZE: "tl.constexpr",
         ):
             pid = tl.program_id(axis=0)
@@ -3159,7 +3358,7 @@ def add_4_times_kernel(
             mask = offsets < n_elements
             x = tl.load(in_ptr0 + offsets, mask=mask)
             y = tl.load(in_ptr1 + offsets, mask=mask)
-            output = tl.zeros((n_elements,), dtype=tl.float32)
+            output = tl.zeros((BLOCK_SIZE,), dtype=tl.float32)
             for _ in range(4):
                 output += x + y
             tl.store(out_ptr + offsets, output, mask=mask)
@@ -3220,7 +3419,7 @@ def add_4_times_kernel(
             in_ptr0,
             in_ptr1,
             out_ptr,
-            n_elements,
+            n_elements: "tl.constexpr",
             BLOCK_SIZE: "tl.constexpr",
         ):
             pid = tl.program_id(axis=0)
@@ -3229,7 +3428,7 @@ def add_4_times_kernel(
             mask = offsets < n_elements
             x = tl.load(in_ptr0 + offsets, mask=mask)
             y = tl.load(in_ptr1 + offsets, mask=mask)
-            output = tl.zeros((n_elements,), dtype=tl.float32)
+            output = tl.zeros((BLOCK_SIZE,), dtype=tl.float32)
             for _ in range(2):
                 for _ in range(2):
                     output += x + y
@@ -3256,7 +3455,7 @@ def add_4_times_kernel(
             in_ptr0,
             in_ptr1,
             out_ptr,
-            n_elements,
+            n_elements: "tl.constexpr",
             BLOCK_SIZE: "tl.constexpr",
         ):
             pid = tl.program_id(axis=0)
@@ -3265,8 +3464,8 @@ def add_4_times_kernel(
             mask = offsets < n_elements
             x = tl.load(in_ptr0 + offsets, mask=mask)
             y = tl.load(in_ptr1 + offsets, mask=mask)
-            output1 = tl.zeros((n_elements,), dtype=tl.float32)
-            output2 = tl.zeros((n_elements,), dtype=tl.float32)
+            output1 = tl.zeros((BLOCK_SIZE,), dtype=tl.float32)
+            output2 = tl.zeros((BLOCK_SIZE,), dtype=tl.float32)
             for _ in range(2):
                 for _ in range(2):
                     output1 += y
@@ -3654,6 +3853,75 @@ def test_add_kernel_on_device_tma_new_api():
             ["c_ptr"],
         )
 
+    @unittest.skipIf(
+        not has_triton_tensor_descriptor_host_tma(),
+        "requires TensorDescriptor API in Triton",
+    )
+    @requires_gpu
+    def test_descriptor_load_store_read_write_detection(self):
+        """
+        Regression test: tl.make_tensor_descriptor + desc.load()/desc.store()
+        generates tt.descriptor_load/tt.descriptor_store ops in TTIR. These
+        must be recognized by identify_accessed_tensors so that Inductor
+        correctly tracks read/write dependencies.
+        """
+        from torch._higher_order_ops.triton_kernel_wrap import identify_accessed_tensors
+
+        @triton.jit
+        def add_kernel_descriptor_method(
+            in_ptr0,
+            in_ptr1,
+            out_ptr,
+            M,
+            N,
+            BLOCK_M: "tl.constexpr",
+            BLOCK_N: "tl.constexpr",
+        ):
+            in0_desc = tl.make_tensor_descriptor(
+                in_ptr0,
+                [M, N],
+                [N, 1],
+                [BLOCK_M, BLOCK_N],
+            )
+            in1_desc = tl.make_tensor_descriptor(
+                in_ptr1,
+                [M, N],
+                [N, 1],
+                [BLOCK_M, BLOCK_N],
+            )
+            out_desc = tl.make_tensor_descriptor(
+                out_ptr,
+                [M, N],
+                [N, 1],
+                [BLOCK_M, BLOCK_N],
+            )
+            pid_m = tl.program_id(0)
+            pid_n = tl.program_id(1)
+            off_m = pid_m * BLOCK_M
+            off_n = pid_n * BLOCK_N
+            a = in0_desc.load([off_m, off_n])
+            b = in1_desc.load([off_m, off_n])
+            out_desc.store([off_m, off_n], a + b)
+
+        t = torch.randn(64, 64)
+        tensor_accesses = identify_accessed_tensors(
+            add_kernel_descriptor_method,
+            {
+                "in_ptr0": t,
+                "in_ptr1": t,
+                "out_ptr": t,
+                "M": 64,
+                "N": 64,
+                "BLOCK_M": 32,
+                "BLOCK_N": 32,
+            },
+            {},
+        )
+        read_names = [dep.name for dep in tensor_accesses.read_writes.reads]
+        write_names = [dep.name for dep in tensor_accesses.read_writes.writes]
+        self.assertListEqual(read_names, ["in_ptr0", "in_ptr1"])
+        self.assertListEqual(write_names, ["out_ptr"])
+
 
 if HAS_GPU:
     t = torch.randn(4)
@@ -3816,6 +4084,9 @@ def test_add_kernel_on_device_tma_new_api():
         while name in MutationTests.__dict__:
             name += "1"
 
+        if kernel.fn.__name__ == "add_kernel_2d_autotuned":
+            fn = unittest.skip("Fails with Triton update")(fn)
+
         setattr(MutationTests, name, fn)
 
 
@@ -4190,7 +4461,7 @@ def f(x, y):
         result = f(x, x)
         self.assertEqual(result, x + x)
 
-        from functorch import make_fx
+        from torch.fx.experimental.proxy_tensor import make_fx
 
         gm = make_fx(f, tracing_mode=tracing_mode)(x, x)
         self.assertEqual(gm(x, x), x + x)
@@ -4277,7 +4548,7 @@ def op_zeros(
         def foo(x, w):
             M, K = x.shape
             KB, N = w.shape
-            assert K == KB, f"incompatible dimensions {K}, {KB}"
+            assert K == KB, f"incompatible dimensions {K}, {KB}"  # noqa: S101
 
             z = torch.empty((M, N), device=x.device, dtype=x.dtype)
 
@@ -4369,7 +4640,8 @@ def add(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
         y = torch.ones((4096,), device=GPU_TYPE, dtype=torch.float16)
 
         # should always pass
-        assert add(x, y).mean() == 2, "Problem with add kernel"
+        if add(x, y).mean() != 2:
+            raise AssertionError("Problem with add kernel")
 
         # assert that the user_defined_* flags are properly set on the kernel before compilation
         self.assertEqual(isinstance(add_kernel, Autotuner), True)
@@ -4679,6 +4951,118 @@ def f(dst, src, add_float, N):
 
         self.assertEqual(counter.op_count, 2)
 
+    @requires_gpu
+    @common_utils.parametrize("backend", ["eager", "aot_eager", "inductor"])
+    def test_triton_kernel_prune_configs_by_called_twice(self, backend):
+        def early_config_prune(configs, named_args, **kwargs):
+            return [configs[0]]
+
+        @triton.autotune(
+            configs=[
+                triton.Config(kwargs={"BLOCK_SIZE": 128}),
+                triton.Config(kwargs={"BLOCK_SIZE": 256}),
+            ],
+            key=["N"],
+            prune_configs_by={"early_config_prune": early_config_prune},
+        )
+        @triton.jit
+        def add_kernel(
+            x_ptr,
+            y_ptr,
+            out_ptr,
+            N,
+            BLOCK_SIZE: tl.constexpr,
+        ):
+            offsets = tl.program_id(0) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+            mask = offsets < N
+            x = tl.load(x_ptr + offsets, mask=mask)
+            y = tl.load(y_ptr + offsets, mask=mask)
+            tl.store(out_ptr + offsets, x + y, mask=mask)
+
+        def call_kernel(x, y):
+            out = torch.empty_like(x)
+            n_elements = x.numel()
+            grid = lambda meta: (triton.cdiv(n_elements, meta["BLOCK_SIZE"]),)
+            add_kernel[grid](x, y, out, N=n_elements)
+            return out
+
+        @torch.compile(fullgraph=True, backend=backend)
+        def f(x, y):
+            out = call_kernel(x, y)
+            return call_kernel(out, y)
+
+        x = torch.randn(1024, device=GPU_TYPE)
+        y = torch.randn(1024, device=GPU_TYPE)
+
+        self.assertEqual(f(x, y), x + y + y)
+
+    @requires_gpu
+    @common_utils.parametrize("backend", ["eager", "aot_eager", "inductor"])
+    def test_triton_kernel_heuristics_and_prune_configs_by(self, backend):
+        def noop_prune(configs, named_args, **kwargs):
+            return configs
+
+        @triton.autotune(
+            configs=[
+                triton.Config({"BLOCK": 128}),
+                triton.Config({"BLOCK": 256}),
+            ],
+            key=["N"],
+            prune_configs_by={"early_config_prune": noop_prune},
+        )
+        @triton.heuristics({"EVEN": lambda args: args["N"] % 128 == 0})
+        @triton.jit
+        def kernel(x_ptr, out_ptr, N, BLOCK: tl.constexpr, EVEN: tl.constexpr):
+            pid = tl.program_id(0)
+            offs = pid * BLOCK + tl.arange(0, BLOCK)
+            mask = offs < N
+            tl.store(out_ptr + offs, tl.load(x_ptr + offs, mask=mask) * 2, mask=mask)
+
+        def f(x):
+            out = torch.empty_like(x)
+            kernel[(triton.cdiv(x.numel(), 128),)](x, out, x.numel())
+            return out
+
+        compiled_f = torch.compile(f, backend=backend, fullgraph=True)
+        x = torch.randn(1024, device=GPU_TYPE)
+        self.assertEqual(compiled_f(x), x * 2)
+
+    @requires_gpu
+    @common_utils.parametrize("backend", ["eager", "aot_eager", "inductor"])
+    def test_triton_kernel_run_with_prune_configs_by(self, backend):
+        def noop_prune(configs, named_args, **kwargs):
+            return configs
+
+        @triton.autotune(
+            configs=[
+                triton.Config({"BLOCK": 64}),
+                triton.Config({"BLOCK": 128}),
+            ],
+            key=["N"],
+            prune_configs_by={"early_config_prune": noop_prune},
+        )
+        @triton.jit
+        def kernel(x_ptr, out_ptr, N, BLOCK: tl.constexpr):
+            pid = tl.program_id(0)
+            offs = pid * BLOCK + tl.arange(0, BLOCK)
+            mask = offs < N
+            tl.store(out_ptr + offs, tl.load(x_ptr + offs, mask=mask) + 1, mask=mask)
+
+        def f(x):
+            out = torch.empty_like(x)
+            kernel.run(
+                x,
+                out,
+                x.numel(),
+                grid=(triton.cdiv(x.numel(), 64),),
+                warmup=False,
+            )
+            return out
+
+        compiled_f = torch.compile(f, backend=backend, fullgraph=True)
+        x = torch.randn(1024, device=GPU_TYPE)
+        self.assertEqual(compiled_f(x), x + 1)
+
     # see: https://github.com/triton-lang/triton/blob/67ea999935f4511a535a25bdecb27e79e3c3af41/python/test/unit/language/test_decorator.py#L31
     @requires_gpu
     @common_utils.parametrize("non_strict", [True, False])
@@ -4808,8 +5192,373 @@ def f(
                 compiled_f(dst, src, N=N)
 
 
+class TestUserKernelEpilogueFusion(torch._inductor.test_case.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        super().setUpClass()
+        cls._stack = contextlib.ExitStack()
+        torch._inductor.config.epilogue_fusion_user_defined_triton_kernel = True
+
+    def check_code(self, code_str, num_kernels, num_allocs, num_deallocs):
+        from torch._inductor import config
+
+        kernel_launch = "call_" if config.cpp_wrapper else ".run("
+        FileCheck().check(get_func_call()).check_count(
+            kernel_launch,
+            num_kernels,
+            exactly=True,
+        ).run(code_str)
+
+        if num_allocs is not None:
+            FileCheck().check(get_func_call()).check_count(
+                "empty_strided", num_allocs, exactly=True
+            ).run(code_str)
+
+        # skip the deallocation check when using cpp_wrapper; most deallocations happen
+        # outside of our control via RAIIAtenTensorHandle
+        if num_deallocs is not None and not config.cpp_wrapper:
+            FileCheck().check(get_func_call()).check_count(
+                "del", num_deallocs, exactly=True
+            ).run(code_str)
+
+    @requires_cuda_and_triton
+    def test_fusion_relu_epilogue(self):
+        @triton.jit
+        def add_kernel(in_ptr0, in_ptr1, out_ptr, n_elements, BLOCK_SIZE: tl.constexpr):
+            pid = tl.program_id(0)
+            offs = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+            mask = offs < n_elements
+            x = tl.load(in_ptr0 + offs, mask=mask)
+            y = tl.load(in_ptr1 + offs, mask=mask)
+            tl.store(out_ptr + offs, x + y, mask=mask)
+
+        def fn(a, b):
+            out = torch.empty_like(a)
+            grid = (triton.cdiv(a.numel(), 1024),)
+            add_kernel[grid](a, b, out, a.numel(), BLOCK_SIZE=1024)
+            return out.relu()
+
+        a = torch.tensor([-0.3] * 8, dtype=torch.float32, device="cuda")
+        b = torch.tensor(
+            [0.0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7], dtype=torch.float32, device="cuda"
+        )
+
+        out, code = run_and_get_code(torch.compile(fn), a, b)
+        self.assertEqual(out, fn(a, b), atol=0.05, rtol=0.05)
+        self.check_code(code[0], num_kernels=1, num_allocs=1, num_deallocs=2)
+
+    @requires_cuda_and_triton
+    def test_fusion_sigmoid_epilogue(self):
+        @triton.jit
+        def add_kernel(in_ptr0, out_ptr, n_elements, BLOCK_SIZE: tl.constexpr):
+            pid = tl.program_id(0)
+            offs = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+            mask = offs < n_elements
+            x = tl.load(in_ptr0 + offs, mask=mask)
+            tl.store(out_ptr + offs, x + x, mask=mask)
+
+        def fn(a):
+            out = torch.empty_like(a)
+            grid = (triton.cdiv(a.numel(), 1024),)
+            add_kernel[grid](a, out, a.numel(), BLOCK_SIZE=1024)
+            return out.sigmoid()
+
+        a = torch.tensor(
+            [0.0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7], dtype=torch.float32, device="cuda"
+        )
+
+        out, code = run_and_get_code(torch.compile(fn), a)
+        self.assertEqual(out, fn(a), atol=0.05, rtol=0.05)
+        self.check_code(code[0], num_kernels=1, num_allocs=1, num_deallocs=1)
+
+    @requires_cuda_and_triton
+    def test_fusion_custom_kernel_with_linebreaks(self):
+        # we do AST manipulation / string manipulation of the kernel source code
+        # so we wanna make sure to correctly handle edge cases with tricky line breaks
+        @triton.jit
+        def add_kernel(in_ptr0, out_ptr, n_elements, BLOCK_SIZE: tl.constexpr):
+            pid = tl.program_id(0)
+            offs = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+            mask = offs < n_elements
+            x = tl.load(in_ptr0 + offs, mask=mask)
+            tl.store(
+                out_ptr + offs,
+                # forcing the `value` arg to be written in multiple lines
+                (
+                    # forcing the `value` arg to be written in multiple lines
+                    x + x
+                ),
+                mask=mask,
+            )
+
+        def fn(a):
+            out = torch.empty_like(a)
+            grid = (triton.cdiv(a.numel(), 1024),)
+            add_kernel[grid](a, out, a.numel(), BLOCK_SIZE=1024)
+            return out.sigmoid()
+
+        a = torch.tensor(
+            [0.0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7], dtype=torch.float32, device="cuda"
+        )
+
+        out, code = run_and_get_code(torch.compile(fn), a)
+        self.assertEqual(out, fn(a), atol=0.05, rtol=0.05)
+        self.check_code(code[0], num_kernels=1, num_allocs=1, num_deallocs=1)
+
+    @requires_cuda_and_triton
+    def test_no_fusion_for_read_write_tensors(self):
+        # cannot epilogue-fuse this one because the kernel reads from `out_ptr` before writing to it
+        @triton.jit
+        def add_kernel(in_ptr0, in_ptr1, out_ptr, n_elements, BLOCK_SIZE: tl.constexpr):
+            pid = tl.program_id(0)
+            offs = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+            mask = offs < n_elements
+            x = tl.load(in_ptr0 + offs, mask=mask)
+            y = tl.load(in_ptr1 + offs, mask=mask)
+            original_out = tl.load(out_ptr + offs, mask=mask)
+            tl.store(out_ptr + offs, x + y + original_out, mask=mask)
+
+        def fn(a, b):
+            out = torch.empty_like(a)
+            grid = (triton.cdiv(a.numel(), 1024),)
+            add_kernel[grid](a, b, out, a.numel(), BLOCK_SIZE=1024)
+            return out.relu()
+
+        a = torch.tensor([-0.3] * 8, dtype=torch.float32, device="cuda")
+        b = torch.tensor(
+            [0.0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7], dtype=torch.float32, device="cuda"
+        )
+
+        _, code = run_and_get_code(torch.compile(fn), a, b)
+        # no fusion, so 2 kernels: add_kernel, relu
+        self.check_code(code[0], num_kernels=2, num_allocs=2, num_deallocs=3)
+
+    @requires_cuda_and_triton
+    def test_no_fusion_for_multiple_stores(self):
+        # cannot epilogue-fuse this one because the kernel stores twice (once directly and once indirectly via another triton func)
+        """
+        We defined these globally after this class, because triton doesn't like local jit functions that reference each other..
+
+        @triton.jit
+        def custom_store(ptr, val, mask):
+            tl.store(ptr, val, mask)
+
+        @triton.jit
+        def add_kernel_with_custom_store(in_ptr0, in_ptr1, out_ptr, n_elements, BLOCK_SIZE: tl.constexpr):
+            pid = tl.program_id(0)
+            offs = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+            mask = offs < n_elements
+            x = tl.load(in_ptr0 + offs, mask=mask)
+            y = tl.load(in_ptr1 + offs, mask=mask)
+            tl.store(out_ptr + offs, x + y, mask=mask)
+            custom_store(out_ptr + offs, x + y, mask=mask)
+        """
+
+        def fn(a, b):
+            out = torch.empty_like(a)
+            grid = (triton.cdiv(a.numel(), 1024),)
+            add_kernel_with_custom_store[grid](a, b, out, a.numel(), BLOCK_SIZE=1024)
+            return out.relu()
+
+        a = torch.tensor([-0.3] * 8, dtype=torch.float32, device="cuda")
+        b = torch.tensor(
+            [0.0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7], dtype=torch.float32, device="cuda"
+        )
+
+        _, code = run_and_get_code(torch.compile(fn), a, b)
+        # no fusion, so 2 kernels: add_kernel, relu
+        self.check_code(code[0], num_kernels=2, num_allocs=2, num_deallocs=3)
+
+    @requires_cuda_and_triton
+    def test_no_fusion_for_non_empty_outputs(self):
+        # cannot epilogue-fuse this one because the kernel writes to buffer initialized with non-UB values
+        @triton.jit
+        def add_kernel(in_ptr0, in_ptr1, out_ptr, n_elements, BLOCK_SIZE: tl.constexpr):
+            pid = tl.program_id(0)
+            offs = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+            mask = offs < n_elements
+            x = tl.load(in_ptr0 + offs, mask=mask)
+            y = tl.load(in_ptr1 + offs, mask=mask)
+            tl.store(out_ptr + offs // 2, x + y, mask=mask)
+
+        def fn(a, b):
+            out = torch.ones_like(a)
+            grid = (triton.cdiv(a.numel(), 1024),)
+            add_kernel[grid](a, b, out, a.numel(), BLOCK_SIZE=1024)
+            return out.relu()
+
+        a = torch.tensor([-0.3] * 8, dtype=torch.float32, device="cuda")
+        b = torch.tensor(
+            [0.0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7], dtype=torch.float32, device="cuda"
+        )
+
+        out, code = run_and_get_code(torch.compile(fn), a, b)
+        self.assertEqual(out, fn(a, b), atol=0.05, rtol=0.05)
+        # no fusion, so 3 kernels: ones, add_kernel, relu
+        self.check_code(code[0], num_kernels=3, num_allocs=2, num_deallocs=3)
+
+    @requires_cuda_and_triton
+    def test_no_fusion_for_concat_outputs(self):
+        @triton.jit
+        def add_kernel(in_ptr0, in_ptr1, out_ptr, n_elements, BLOCK_SIZE: tl.constexpr):
+            pid = tl.program_id(0)
+            offs = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+            mask = offs < n_elements
+            x = tl.load(in_ptr0 + offs, mask=mask)
+            y = tl.load(in_ptr1 + offs, mask=mask)
+            tl.store(out_ptr + offs, x + y, mask=mask)
+
+        def fn(a, b):
+            out_first_half = torch.empty(int(len(a) / 2), dtype=a.dtype, device="cuda")
+            out_second_half = torch.empty(int(len(a) / 2), dtype=a.dtype, device="cuda")
+            out = torch.concat((out_first_half, out_second_half))
+            grid = (triton.cdiv(a.numel(), 1024),)
+            add_kernel[grid](a, b, out, a.numel(), BLOCK_SIZE=1024)
+            return out.relu()
+
+        a = torch.tensor([-0.3] * 8, dtype=torch.float32, device="cuda")
+        b = torch.tensor(
+            [0.0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7], dtype=torch.float32, device="cuda"
+        )
+
+        out, code = run_and_get_code(torch.compile(fn), a, b)
+        self.assertEqual(out, fn(a, b), atol=0.05, rtol=0.05)
+        # no fusion, so 3 kernels: concat, add_kernel, relu
+        self.check_code(code[0], num_kernels=3, num_allocs=2, num_deallocs=3)
+
+    @requires_cuda_and_triton
+    def test_no_fusion_for_multiple_reads_on_mutated_tensor(self):
+        @triton.jit
+        def add_kernel(in_ptr0, in_ptr1, out_ptr, n_elements, BLOCK_SIZE: tl.constexpr):
+            pid = tl.program_id(0)
+            offs = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+            mask = offs < n_elements
+            x = tl.load(in_ptr0 + offs, mask=mask)
+            y = tl.load(in_ptr1 + offs, mask=mask)
+            tl.store(out_ptr + offs, x + y, mask=mask)
+
+        def fn(a, b):
+            out = torch.empty_like(a)
+            grid = (triton.cdiv(a.numel(), 1024),)
+            add_kernel[grid](a, b, out, a.numel(), BLOCK_SIZE=1024)
+            return out.relu(), out.sigmoid()
+
+        a = torch.tensor([-0.3] * 8, dtype=torch.float32, device="cuda")
+        b = torch.tensor(
+            [0.0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7], dtype=torch.float32, device="cuda"
+        )
+
+        out, code = run_and_get_code(torch.compile(fn), a, b)
+        self.assertEqual(out, fn(a, b), atol=0.05, rtol=0.05)
+        # two kernels: user defined `add_kernel`, and the generated one for both sigmoid and relu
+        self.check_code(code[0], num_kernels=2, num_allocs=3, num_deallocs=3)
+
+    @requires_cuda_and_triton
+    def test_no_fusion_for_atomic_store(self):
+        # on ROCm we skip this because `tl.atomic_xchg` fails to compile in ROCm
+        if torch.version.cuda is None:
+            return
+
+        @triton.jit
+        def add_kernel(in_ptr0, in_ptr1, out_ptr, n_elements, BLOCK_SIZE: tl.constexpr):
+            pid = tl.program_id(0)
+            offs = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+            mask = offs < n_elements
+            x = tl.load(in_ptr0 + offs, mask=mask)
+            y = tl.load(in_ptr1 + offs, mask=mask)
+            tl.atomic_xchg(out_ptr + offs, x + y, mask=mask)
+
+        def fn(a, b):
+            out = torch.empty_like(a)
+            grid = (triton.cdiv(a.numel(), 1024),)
+            add_kernel[grid](a, b, out, a.numel(), BLOCK_SIZE=1024)
+            return out.relu()
+
+        a = torch.tensor([-0.3] * 8, dtype=torch.float32, device="cuda")
+        b = torch.tensor(
+            [0.0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7], dtype=torch.float32, device="cuda"
+        )
+
+        out, code = run_and_get_code(torch.compile(fn), a, b)
+        self.assertEqual(out, fn(a, b), atol=0.05, rtol=0.05)
+        self.check_code(code[0], num_kernels=2, num_allocs=2, num_deallocs=3)
+
+    @requires_cuda_and_triton
+    def test_no_fusion_indexing_var_usage(self):
+        @triton.jit
+        def add_kernel(in_ptr0, in_ptr1, out_ptr, n_elements, BLOCK_SIZE: tl.constexpr):
+            pid = tl.program_id(0)
+            offs = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+            mask = offs < n_elements
+            x = tl.load(in_ptr0 + offs, mask=mask)
+            y = tl.load(in_ptr1 + offs, mask=mask)
+            tl.store(out_ptr + offs, x + y, mask=mask)
+
+        def fn(a, b):
+            out = torch.empty_like(a)
+            grid = (triton.cdiv(a.numel(), 1024),)
+            add_kernel[grid](a, b, out, a.numel(), BLOCK_SIZE=1024)
+            # this arange causes the final kernel to require indexing expressions not available in the user kernel
+            return out.relu() + torch.arange(out.shape[0], device="cuda")
+
+        a = torch.tensor([-0.3] * 8, dtype=torch.float32, device="cuda")
+        b = torch.tensor(
+            [0.0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7], dtype=torch.float32, device="cuda"
+        )
+
+        out, code = run_and_get_code(torch.compile(fn), a, b)
+        self.assertEqual(out, fn(a, b), atol=0.05, rtol=0.05)
+        self.check_code(code[0], num_kernels=2, num_allocs=2, num_deallocs=3)
+
+    @requires_cuda_and_triton
+    def test_no_fusion_non_unary_epilogue(self):
+        @triton.jit
+        def add_kernel(a_ptr, b_ptr, out_ptr, numel, BLOCK_SIZE: tl.constexpr):
+            pid = tl.program_id(0)
+            offs = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+            mask = offs < numel
+            a = tl.load(a_ptr + offs, mask=mask)
+            b = tl.load(b_ptr + offs, mask=mask)
+            out = a + b
+            tl.store(out_ptr + offs, out, mask=mask)
+
+        def fn(a, b, c):
+            out = torch.empty_like(a)
+            GRID = (a.numel(),)
+            add_kernel[GRID](a, b, out, a.numel(), 1)
+            return out + c
+
+        a = torch.randn(10, dtype=torch.float32, device="cuda")
+        b = torch.randn(10, dtype=torch.float32, device="cuda")
+        c = torch.randn(10, dtype=torch.float32, device="cuda")
+
+        out, code = run_and_get_code(torch.compile(fn), a, b, c)
+        self.assertEqual(out, fn(a, b, c))
+        self.check_code(code[0], num_kernels=2, num_allocs=2, num_deallocs=4)
+
+
+if HAS_CUDA_AND_TRITON:
+
+    @triton.jit
+    def custom_store(ptr, val, mask):
+        tl.store(ptr, val, mask)
+
+    @triton.jit
+    def add_kernel_with_custom_store(
+        in_ptr0, in_ptr1, out_ptr, n_elements, BLOCK_SIZE: tl.constexpr
+    ):
+        pid = tl.program_id(0)
+        offs = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+        mask = offs < n_elements
+        x = tl.load(in_ptr0 + offs, mask=mask)
+        y = tl.load(in_ptr1 + offs, mask=mask)
+        tl.store(out_ptr + offs, x + y, mask=mask)
+        custom_store(out_ptr + offs, x + y, mask=mask)
+
+
 common_utils.instantiate_parametrized_tests(KernelTests)
 common_utils.instantiate_parametrized_tests(CustomOpTests)
+common_utils.instantiate_parametrized_tests(TestUserKernelEpilogueFusion)
 
 
 if __name__ == "__main__":
diff --git a/test/inductor/test_triton_syntax.py b/test/inductor/test_triton_syntax.py
index 49f0bd06a8bc7..8a8a63d4cd2da 100644
--- a/test/inductor/test_triton_syntax.py
+++ b/test/inductor/test_triton_syntax.py
@@ -14,7 +14,7 @@ def test_triton_sqrt(self):
         import torch.nn as nn
 
         def newtonschulz5(G, steps: int, eps=1e-7):
-            assert len(G.shape) == 2
+            assert len(G.shape) == 2  # noqa: S101
             a, b, c = (3.4445, -4.7750, 2.0315)
             X = G.to(
                 torch.bfloat16
diff --git a/test/inductor/test_unbacked_symints.py b/test/inductor/test_unbacked_symints.py
index 85a66a1833037..565a632cb4741 100644
--- a/test/inductor/test_unbacked_symints.py
+++ b/test/inductor/test_unbacked_symints.py
@@ -311,6 +311,23 @@ def fn(x, a, b):
         expected = fn(*example_inputs)
         torch.testing.assert_close(actual, expected)
 
+    @skipGPUIf(not HAS_GPU, "requires gpu and triton")
+    @dynamo_config.patch(
+        {"capture_scalar_outputs": True, "capture_dynamic_output_shape_ops": True}
+    )
+    def test_repeat_interleave_with_unbacked_scalar(self, device):
+        def fn(x, repeats):
+            return x.repeat_interleave(repeats.item())
+
+        example_inputs = (
+            torch.arange(4, device=device),
+            torch.scalar_tensor(3, dtype=torch.int64, device=device),
+        )
+
+        actual = torch.compile(fn, fullgraph=True, dynamic=True)(*example_inputs)
+        expected = fn(*example_inputs)
+        torch.testing.assert_close(actual, expected)
+
     @skipGPUIf(not HAS_GPU, "requires gpu and triton")
     @dynamo_config.patch({"capture_scalar_outputs": True})
     @parametrize("dynamic", [False, True, None])
@@ -499,9 +516,9 @@ def fn(x):
         expected = fn(*example_inputs)
         torch.testing.assert_close(actual, expected)
 
+    @skipIfXpu(msg="FlashAttentionForward headdim limitation on xpu")
     @skipGPUIf(not HAS_GPU, "requires gpu and triton")
     @skipCUDAIf(not SM80OrLater, "Requires sm80 or later.")
-    @skipIfXpu(msg="_scaled_dot_product_flash_attention is not supported on XPU yet")
     @dynamo_config.patch({"capture_dynamic_output_shape_ops": True})
     def test_sdpfa(self, device):
         if device == "cpu":
@@ -526,9 +543,9 @@ def fn(x):
         x = torch.tensor([1.0, 0.0, 1.0, 0.0], device=device)
         torch.compile(fn, fullgraph=True)(x)
 
+    @skipIfXpu(msg="FlashAttentionForward headdim limitation on xpu")
     @skipGPUIf(not HAS_GPU, "requires gpu and triton")
     @skipCUDAIf(not SM80OrLater, "Requires sm80 or later.")
-    @skipIfXpu(msg="scaled_dot_product_attention is not supported on XPU yet")
     @dynamo_config.patch({"capture_dynamic_output_shape_ops": True})
     def test_sdfpa_unbacked_strides(self, device):
         if device == "cpu":
@@ -720,6 +737,124 @@ def fn(x):
         expected = fn(*example_inputs)
         torch.testing.assert_close(actual, expected)
 
+    @skipCPUIf(True, "Triton codegen bug only affects GPU")
+    @skipGPUIf(not HAS_GPU, "requires gpu and triton")
+    @dynamo_config.patch({"capture_scalar_outputs": True})
+    def test_triton_trunc_large_float_scalar_tensor(self, device):
+        import math
+
+        def fn(x):
+            r = math.sqrt(x.size(0))
+            r = r**70
+            return torch.tensor(math.trunc(r), dtype=torch.float64, device=device)
+
+        example_inputs = (torch.randn(4, device=device),)
+        torch._dynamo.mark_dynamic(example_inputs[0], 0)
+        actual = torch.compile(fn, fullgraph=True, dynamic=True)(*example_inputs)
+        expected = fn(*example_inputs)
+        torch.testing.assert_close(actual, expected)
+
+    @skipCPUIf(True, "Triton codegen bug only affects GPU")
+    @skipGPUIf(not HAS_GPU, "requires gpu and triton")
+    @dynamo_config.patch({"capture_scalar_outputs": True})
+    def test_triton_trunc_float_scalar_tensor_preserves_positive_zero(self, device):
+        import math
+
+        def fn(x):
+            r = math.sqrt(x.size(0)) - 2.5
+            return torch.signbit(
+                torch.tensor(math.trunc(r), dtype=torch.float64, device=device)
+            )
+
+        example_inputs = (torch.randn(4, device=device),)
+        torch._dynamo.mark_dynamic(example_inputs[0], 0)
+        actual = torch.compile(fn, fullgraph=True, dynamic=True)(*example_inputs)
+        expected = fn(*example_inputs)
+        self.assertEqual(actual, expected)
+
+    @skipCPUIf(True, "Triton codegen bug only affects GPU")
+    @skipGPUIf(not HAS_GPU, "requires gpu and triton")
+    @dynamo_config.patch({"capture_scalar_outputs": True})
+    def test_triton_pow_symbolic_int_exponent(self, device):
+        """
+        Test that symbolic integer scalar exponents are cast before libdevice.pow.
+
+        Reproducer from https://github.com/pytorch/pytorch/issues/177131:
+        capture_scalar_outputs=True can pass a scalar like num_bits as an int64 ks*
+        kernel argument, which previously led Triton to reject libdevice.pow(2.0, ks0).
+        """
+
+        def calculate_range(num_bits, device):
+            bit_range = 2**num_bits
+            q_max = torch.tensor(bit_range / 2 - 1, device=device)
+            q_min = torch.tensor(-bit_range / 2, device=device)
+            return q_min, q_max
+
+        def fn(x, num_bits):
+            q_min, q_max = calculate_range(num_bits, x.device)
+            bit_range = q_max - q_min
+            max_val = torch.max(torch.abs(x))
+            scale = max_val / (float(bit_range) / 2)
+            scale = torch.where(
+                scale == 0,
+                torch.tensor(1e-10, device=x.device, dtype=x.dtype),
+                scale,
+            )
+            x_q = torch.clamp(torch.round(x / scale), q_min, q_max)
+            x_dq = x_q * scale
+            return (x_dq - x).abs().pow(2.4).sum(), scale, scale
+
+        example_inputs = (torch.randn(1, 1, 128, device=device, dtype=torch.float16), 8)
+        actual = torch.compile(fn, fullgraph=True, dynamic=True)(*example_inputs)
+        expected = fn(*example_inputs)
+        torch.testing.assert_close(actual, expected)
+
+    @skipCPUIf(True, "Triton codegen bug only affects GPU")
+    @skipGPUIf(not HAS_GPU, "requires gpu and triton")
+    @dynamo_config.patch({"capture_scalar_outputs": True})
+    def test_triton_pow_symbolic_negative_int_exponent(self, device):
+        def fn(x, exponent_src):
+            exponent = exponent_src.max().to(torch.int64) - 1
+            return x * (2**exponent)
+
+        example_inputs = (
+            torch.randn(128, device=device, dtype=torch.float16),
+            torch.tensor([0], device=device, dtype=torch.int64),
+        )
+        actual = torch.compile(fn, fullgraph=True, dynamic=True)(*example_inputs)
+        expected = fn(*example_inputs)
+        torch.testing.assert_close(actual, expected)
+
+    @skipGPUIf(not HAS_GPU, "requires gpu and triton")
+    @dynamo_config.patch({"capture_scalar_outputs": True})
+    def test_slice_unbacked_bindings_with_later_constraint(self, device):
+        # Regression test for https://github.com/pytorch/pytorch/issues/166460
+        # When slicing with an unbacked symint end index (e.g. x[:total] where
+        # total = sum of data-dependent values), dynamo may allocate a fresh
+        # unbacked symbol for the output size because it can't prove bounds at
+        # trace time. Later operations (e.g. new_zeros with padding) may
+        # establish constraints that make the bounds provable at inductor time.
+        # Inductor must still define the unbacked symbol even when taking the
+        # efficient SliceView path.
+        def fn(x, sizes):
+            sizes_list = sizes.tolist()
+            total = sum(sizes_list)
+            num_padding = x.shape[0] - total
+            sliced = x[:total]
+            splits = torch.split(sliced, sizes_list, dim=0)
+            out = torch.cat([s * 2 for s in splits], dim=0)
+            out = torch.vstack((out, out.new_zeros((num_padding, out.shape[-1]))))
+            return out
+
+        example_inputs = (
+            torch.randn(64, 16, device=device, dtype=torch.bfloat16),
+            torch.tensor([8, 8], device=device, dtype=torch.int64),
+        )
+
+        actual = torch.compile(fn, fullgraph=True)(*example_inputs)
+        expected = fn(*example_inputs)
+        torch.testing.assert_close(actual, expected)
+
 
 instantiate_device_type_tests(TestUnbackedSymints, globals(), allow_xpu=True)
 
diff --git a/test/inductor/test_user_streams.py b/test/inductor/test_user_streams.py
new file mode 100644
index 0000000000000..eaba180530ada
--- /dev/null
+++ b/test/inductor/test_user_streams.py
@@ -0,0 +1,2394 @@
+# Owner(s): ["module: inductor"]
+"""Tests for user-annotated stream support in Inductor.
+
+This module tests the infrastructure for supporting user-annotated CUDA stream
+assignments on nodes, including stream utilities, event management, and codegen.
+"""
+
+from __future__ import annotations
+
+import re
+import unittest
+
+import torch
+import torch._inductor.metrics
+from torch._dynamo.testing import CompileCounterWithBackend, normalize_gm
+from torch._inductor import config as inductor_config
+from torch._inductor.codegen.wrapper import (
+    EnterCudaStreamContextLine,
+    EnterDeviceContextManagerWithStreamInfoLine,
+    ExitCudaStreamContextLine,
+    ExitDeviceContextManagerWithStreamInfoLine,
+)
+from torch._inductor.stream_constants import (
+    DEFAULT_STREAM,
+    DEFAULT_STREAM_IDX,
+    STREAM_NAME_TEMPLATE,
+)
+from torch._inductor.stream_utils import get_raw_stream_name, get_stream_name
+from torch._inductor.test_case import TestCase as InductorTestCase
+from torch._inductor.utils import IndentedBuffer
+from torch.testing import FileCheck
+from torch.testing._internal.common_cuda import SM90OrLater, TEST_CUDA
+from torch.testing._internal.common_utils import (
+    instantiate_parametrized_tests,
+    TEST_WITH_ROCM,
+)
+
+
+def _extract_wrapper_body(code):
+    """Extract and normalize the call method body from generated wrapper code.
+
+    Strips noise (comments, assert_size_stride, del statements, args.clear())
+    and normalizes triton kernel names, leaving just structural code:
+    stream declarations, context switches, event ops, kernel calls, buffer allocations.
+    """
+    lines = code.split("\n")
+
+    # Find the call function body
+    call_start = None
+    call_indent = 0
+    for i, line in enumerate(lines):
+        if "def call(" in line:
+            call_indent = len(line) - len(line.lstrip())
+            call_start = i + 1
+            break
+
+    if call_start is None:
+        return ""
+
+    # Extract body until next definition at same indent level or end
+    body_lines = []
+    for i in range(call_start, len(lines)):
+        line = lines[i]
+        stripped = line.strip()
+        if stripped:
+            current_indent = len(line) - len(line.lstrip())
+            if current_indent <= call_indent:
+                break
+        body_lines.append(line)
+
+    # Filter out noise
+    filtered = []
+    for line in body_lines:
+        stripped = line.strip()
+        if not stripped:
+            continue
+        if stripped.startswith("#"):
+            continue
+        if "assert_size_stride" in stripped:
+            continue
+        if stripped.startswith("del "):
+            continue
+        if "args.clear()" in stripped:
+            continue
+        # Strip inline comments (e.g., "# reuse")
+        line = re.sub(r"\s+#\s.*", "", line)
+        if not line.strip():
+            continue
+        filtered.append(line)
+
+    if not filtered:
+        return ""
+
+    # Dedent
+    min_indent = min(
+        len(line) - len(line.lstrip()) for line in filtered if line.strip()
+    )
+    dedented = [line[min_indent:] for line in filtered]
+    body = "\n".join(dedented)
+
+    # Normalize triton kernel names
+    body = re.sub(r"triton_\w+", "triton_kernel", body)
+
+    return body
+
+
+class TestStreamUtils(InductorTestCase):
+    """Tests for stream_utils module."""
+
+    def test_constants(self):
+        """Test stream utility constants are defined correctly."""
+        self.assertEqual(DEFAULT_STREAM, "default_stream")
+        self.assertEqual(DEFAULT_STREAM_IDX, 0)
+        self.assertEqual(STREAM_NAME_TEMPLATE, "stream{stream_idx:d}")
+
+    def test_get_stream_name_default(self):
+        """Test get_stream_name returns default stream for index 0."""
+        self.assertEqual(get_stream_name(0), DEFAULT_STREAM)
+
+    def test_get_stream_name_side_streams(self):
+        """Test get_stream_name returns formatted names for side streams."""
+        self.assertEqual(get_stream_name(1), "stream1")
+        self.assertEqual(get_stream_name(2), "stream2")
+        self.assertEqual(get_stream_name(10), "stream10")
+
+    def test_get_raw_stream_name(self):
+        self.assertEqual(get_raw_stream_name(0), "raw_stream0")
+        self.assertEqual(get_raw_stream_name(1), "raw_stream1")
+
+    def test_get_stream_name_cached(self):
+        """Test that get_stream_name results are cached."""
+        # Call twice and verify we get the same object (cached)
+        name1 = get_stream_name(5)
+        name2 = get_stream_name(5)
+        self.assertIs(name1, name2)
+
+
+class TestWrapperCodegenStreams(InductorTestCase):
+    """Tests for stream-related wrapper codegen classes."""
+
+    def test_enter_device_context_with_stream_info(self):
+        """Test EnterDeviceContextManagerWithStreamInfoLine has num_streams param."""
+        line = EnterDeviceContextManagerWithStreamInfoLine(
+            device_idx=0,
+            last_seen_device_guard_index=None,
+            num_streams=4,
+        )
+        self.assertEqual(line.num_streams, 4)
+        self.assertEqual(line.device_idx, 0)
+
+    def test_exit_device_context_with_stream_info(self):
+        """Test ExitDeviceContextManagerWithStreamInfoLine has num_streams param."""
+        line = ExitDeviceContextManagerWithStreamInfoLine(num_streams=4)
+        self.assertEqual(line.num_streams, 4)
+
+    def test_enter_cuda_stream_context(self):
+        """Test EnterCudaStreamContextLine creation."""
+        line = EnterCudaStreamContextLine(stream_idx=1)
+        self.assertEqual(line.stream_idx, 1)
+
+    def test_exit_cuda_stream_context(self):
+        """Test ExitCudaStreamContextLine creation."""
+        line = ExitCudaStreamContextLine()
+        self.assertIsNotNone(line)
+
+
+class TestStreamCodegen(InductorTestCase):
+    """End-to-end tests for stream code generation."""
+
+    def test_enter_cuda_stream_context_codegen(self):
+        """Test code generation for entering a CUDA stream context."""
+        code = IndentedBuffer()
+        code.writeline("def call(args):")
+        code.do_indent()
+        code.do_indent()  # Simulate being inside device guard
+
+        line = EnterCudaStreamContextLine(stream_idx=1)
+        line.codegen(code)
+
+        generated = code.getvalue()
+        # Should have stream context
+        self.assertIn("with torch.cuda.stream(stream1)", generated)
+
+    def test_exit_cuda_stream_context_codegen(self):
+        """Test code generation for exiting a CUDA stream context."""
+        code = IndentedBuffer()
+        code.do_indent()
+
+        line = ExitCudaStreamContextLine()
+        line.codegen(code)
+
+        # The exit just unindents, verify no error
+        self.assertIsNotNone(code.getvalue())
+
+
+@unittest.skipIf(not TEST_CUDA, "requires CUDA")
+class TestUserStreamCompile(InductorTestCase):
+    """End-to-end tests for torch.compile with user stream contexts."""
+
+    def test_compile_with_user_stream_context(self):
+        """Test that user code with stream context compiles and runs correctly."""
+        from torch._inductor.utils import run_and_get_code
+
+        def fn(x, y):
+            # Create a side stream
+            s = torch.cuda.Stream()
+            # Perform operation on default stream
+            z = x + y
+            # Perform operation on side stream
+            with torch.cuda.stream(s):
+                w = z * 2
+            # Synchronize before using result
+            s.synchronize()
+            return w + 1
+
+        x = torch.randn(1024, device="cuda")
+        y = torch.randn(1024, device="cuda")
+
+        # Get expected result from eager execution
+        expected = fn(x, y)
+
+        # Compile and run
+        compiled_fn = torch.compile(fn)
+        result, (code,) = run_and_get_code(compiled_fn, x, y)
+
+        # Verify correctness
+        self.assertEqual(result, expected)
+
+        # Verify generated code contains stream handling and synchronize survives
+        self.assertIn("torch.cuda.stream", code)
+        self.assertIn("synchronize_stream", code)
+
+    @unittest.skipIf(
+        not TEST_CUDA or torch.cuda.device_count() < 2,
+        "requires at least two CUDA devices",
+    )
+    def test_raw_stream_name_does_not_clobber_user_stream_on_cuda_1(self):
+        from torch._inductor.utils import run_and_get_code
+
+        device = torch.device("cuda:1")
+        aux = torch.cuda.Stream(device=device)
+        ev0 = torch.cuda.Event()
+        ev1 = torch.cuda.Event()
+
+        def fn(x, w):
+            ev0.record()
+            with torch.cuda.stream(aux):
+                ev0.wait()
+                a = torch.mm(x, w)
+                ev1.record()
+
+            ev1.wait()
+            c = (a.sin() * a.cos() + 0.1).relu()
+
+            with torch.cuda.stream(aux):
+                d = torch.mm(c, w.t())
+
+            return d
+
+        x = torch.randn(64, 128, device=device)
+        w = torch.randn(128, 64, device=device)
+
+        expected = fn(x, w)
+        torch.cuda.synchronize()
+
+        actual, (code,) = run_and_get_code(torch.compile(fn, fullgraph=True), x, w)
+        torch.cuda.synchronize()
+
+        self.assertEqual(actual, expected)
+        self.assertIn("stream1 = get_external_object_by_index", code)
+        self.assertIn("raw_stream1 = get_raw_stream(1)", code)
+        self.assertNotRegex(code, r"(?m)^\s*stream1 = get_raw_stream\(1\)")
+
+    def test_compile_preserves_stream_semantics(self):
+        """Test that compiled code preserves stream execution semantics."""
+        from torch._inductor.utils import run_and_get_code
+
+        def fn(x):
+            s = torch.cuda.Stream()
+            # Work on default stream
+            a = x * 2
+            # Work on side stream
+            with torch.cuda.stream(s):
+                b = x * 3
+            s.synchronize()
+            return a + b
+
+        x = torch.randn(1024, device="cuda")
+
+        expected = fn(x)
+        compiled_fn = torch.compile(fn)
+        result, (code,) = run_and_get_code(compiled_fn, x)
+
+        self.assertEqual(result, expected)
+
+        # Verify stream context and synchronize survive
+        self.assertIn("torch.cuda.stream", code)
+        self.assertIn("synchronize_stream", code)
+
+    def test_multiple_stream_contexts(self):
+        """Test compilation with multiple stream context switches."""
+        from torch._inductor.utils import run_and_get_code
+
+        def fn(x):
+            s1 = torch.cuda.Stream()
+            s2 = torch.cuda.Stream()
+
+            a = x + 1  # default stream
+
+            with torch.cuda.stream(s1):
+                b = x * 2  # stream 1
+
+            with torch.cuda.stream(s2):
+                c = x * 3  # stream 2
+
+            s1.synchronize()
+            s2.synchronize()
+            return a + b + c
+
+        x = torch.randn(1024, device="cuda")
+
+        expected = fn(x)
+        compiled_fn = torch.compile(fn)
+        result, (code,) = run_and_get_code(compiled_fn, x)
+
+        self.assertEqual(result, expected)
+
+        # Verify stream contexts and synchronization survive
+        self.assertGreaterEqual(code.count("torch.cuda.stream"), 1)
+        self.assertIn("synchronize_stream", code)
+
+    def test_nested_stream_contexts(self):
+        """Test compilation with nested stream contexts."""
+        from torch._inductor.utils import run_and_get_code
+
+        def fn(x):
+            s1 = torch.cuda.Stream()
+            s2 = torch.cuda.Stream()
+
+            with torch.cuda.stream(s1):
+                a = x * 2
+                with torch.cuda.stream(s2):
+                    b = x * 3
+                c = a + 1  # back on s1
+
+            s1.synchronize()
+            s2.synchronize()
+            return a + b + c
+
+        x = torch.randn(1024, device="cuda")
+
+        expected = fn(x)
+        compiled_fn = torch.compile(fn)
+        result, (code,) = run_and_get_code(compiled_fn, x)
+
+        self.assertEqual(result, expected)
+
+        # Verify nested stream contexts and synchronization survive
+        self.assertGreaterEqual(code.count("torch.cuda.stream"), 1)
+        self.assertIn("synchronize_stream", code)
+
+    def test_stream_context_with_data_dependency(self):
+        """Test stream contexts with data flowing between streams."""
+        from torch._inductor.utils import run_and_get_code
+
+        def fn(x):
+            s = torch.cuda.Stream()
+
+            # Compute on default stream
+            a = x * 2
+
+            # Use result on side stream
+            with torch.cuda.stream(s):
+                b = a + 1  # depends on 'a' from default stream
+
+            s.synchronize()
+            return b
+
+        x = torch.randn(1024, device="cuda")
+
+        expected = fn(x)
+        compiled_fn = torch.compile(fn)
+        result, (code,) = run_and_get_code(compiled_fn, x)
+
+        self.assertEqual(result, expected)
+
+        # Verify stream context and synchronize survive
+        self.assertIn("torch.cuda.stream", code)
+        self.assertIn("synchronize_stream", code)
+
+    def test_event_record_and_wait(self):
+        """Test compilation with explicit event record and wait."""
+        from torch._inductor.utils import run_and_get_code
+
+        def fn(x):
+            s = torch.cuda.Stream()
+            event = torch.cuda.Event()
+
+            # Compute on default stream
+            a = x * 2
+            # Record event on default stream
+            event.record()
+
+            with torch.cuda.stream(s):
+                # Wait for event before using 'a'
+                event.wait()
+                b = a + 1
+
+            s.synchronize()
+            return b
+
+        x = torch.randn(1024, device="cuda")
+
+        expected = fn(x)
+        compiled_fn = torch.compile(fn)
+        result, (code,) = run_and_get_code(compiled_fn, x)
+
+        self.assertEqual(result, expected)
+
+        # Verify event operations survive compilation as custom ops
+        self.assertIn("record_event", code)
+        self.assertIn("wait_event", code)
+
+    def test_event_record_on_stream(self):
+        """Test event recording on a specific stream."""
+        from torch._inductor.utils import run_and_get_code
+
+        def fn(x):
+            s1 = torch.cuda.Stream()
+            s2 = torch.cuda.Stream()
+            event = torch.cuda.Event()
+
+            with torch.cuda.stream(s1):
+                a = x * 2
+                # Record on s1
+                event.record(s1)
+
+            with torch.cuda.stream(s2):
+                # Wait for s1's work before proceeding
+                event.wait(s2)
+                b = a + 1
+
+            s1.synchronize()
+            s2.synchronize()
+            return b
+
+        x = torch.randn(1024, device="cuda")
+
+        expected = fn(x)
+        compiled_fn = torch.compile(fn)
+        result, (code,) = run_and_get_code(compiled_fn, x)
+
+        self.assertEqual(result, expected)
+
+        # Verify event operations survive compilation as custom ops
+        self.assertIn("record_event", code)
+        self.assertIn("wait_event", code)
+
+    def test_multiple_events_multiple_streams(self):
+        """Test multiple events synchronizing multiple streams."""
+        from torch._inductor.utils import run_and_get_code
+
+        def fn(x):
+            s1 = torch.cuda.Stream()
+            s2 = torch.cuda.Stream()
+            event1 = torch.cuda.Event()
+            event2 = torch.cuda.Event()
+
+            # Work on s1
+            with torch.cuda.stream(s1):
+                a = x * 2
+                event1.record(s1)
+
+            # Work on s2, depends on s1
+            with torch.cuda.stream(s2):
+                event1.wait(s2)
+                b = a + 1
+                event2.record(s2)
+
+            # Back to default stream, wait for s2
+            event2.wait()
+            c = b + x
+
+            s1.synchronize()
+            s2.synchronize()
+            return c
+
+        x = torch.randn(1024, device="cuda")
+
+        expected = fn(x)
+        compiled_fn = torch.compile(fn)
+        result, (code,) = run_and_get_code(compiled_fn, x)
+
+        self.assertEqual(result, expected)
+
+        # Verify multiple events and streams survive as custom ops
+        self.assertGreaterEqual(code.count("record_event"), 2)
+        self.assertGreaterEqual(code.count("wait_event"), 2)
+
+    def test_event_wait_without_record(self):
+        """Test that waiting on unrecorded event works (no-op)."""
+        from torch._inductor.utils import run_and_get_code
+
+        def fn(x):
+            s = torch.cuda.Stream()
+            event = torch.cuda.Event()
+
+            # Record the event first
+            event.record()
+
+            with torch.cuda.stream(s):
+                # Wait is valid after record
+                event.wait()
+                a = x * 2
+
+            s.synchronize()
+            return a
+
+        x = torch.randn(1024, device="cuda")
+
+        expected = fn(x)
+        compiled_fn = torch.compile(fn)
+        result, (code,) = run_and_get_code(compiled_fn, x)
+
+        self.assertEqual(result, expected)
+
+        # Verify event operations survive compilation as custom ops
+        self.assertIn("record_event", code)
+        self.assertIn("wait_event", code)
+
+    def test_stream_wait_event(self):
+        """Test stream.wait_event() method."""
+        from torch._inductor.utils import run_and_get_code
+
+        def fn(x):
+            s = torch.cuda.Stream()
+            event = torch.cuda.Event()
+
+            a = x * 2
+            event.record()
+
+            # Use stream.wait_event instead of event.wait
+            s.wait_event(event)
+            with torch.cuda.stream(s):
+                b = a + 1
+
+            s.synchronize()
+            return b
+
+        x = torch.randn(1024, device="cuda")
+
+        expected = fn(x)
+        compiled_fn = torch.compile(fn)
+        result, (code,) = run_and_get_code(compiled_fn, x)
+
+        self.assertEqual(result, expected)
+
+        # Verify stream.wait_event survives compilation as custom op
+        self.assertIn("wait_event", code)
+
+    def test_bidirectional_stream_sync(self):
+        """Test bidirectional synchronization between streams."""
+        from torch._inductor.utils import run_and_get_code
+
+        def fn(x):
+            s1 = torch.cuda.Stream()
+            s2 = torch.cuda.Stream()
+            event_s1 = torch.cuda.Event()
+            event_s2 = torch.cuda.Event()
+
+            # s1 does work
+            with torch.cuda.stream(s1):
+                a = x * 2
+                event_s1.record(s1)
+
+            # s2 waits for s1, does work, signals back
+            with torch.cuda.stream(s2):
+                event_s1.wait(s2)
+                b = a + 1
+                event_s2.record(s2)
+
+            # s1 waits for s2
+            with torch.cuda.stream(s1):
+                event_s2.wait(s1)
+                c = b * 2
+
+            s1.synchronize()
+            s2.synchronize()
+            return c
+
+        x = torch.randn(1024, device="cuda")
+
+        expected = fn(x)
+        compiled_fn = torch.compile(fn)
+        result, (code,) = run_and_get_code(compiled_fn, x)
+
+        self.assertEqual(result, expected)
+
+        # Verify bidirectional sync - multiple records and waits as custom ops
+        self.assertGreaterEqual(code.count("record_event"), 2)
+        self.assertGreaterEqual(code.count("wait_event"), 2)
+
+    def test_three_streams_pipeline(self):
+        """Test pipeline pattern with three streams."""
+        from torch._inductor.utils import run_and_get_code
+
+        def fn(x):
+            s1 = torch.cuda.Stream()
+            s2 = torch.cuda.Stream()
+            s3 = torch.cuda.Stream()
+            e1 = torch.cuda.Event()
+            e2 = torch.cuda.Event()
+
+            # Stage 1 on s1
+            with torch.cuda.stream(s1):
+                a = x * 2
+                e1.record(s1)
+
+            # Stage 2 on s2, depends on s1
+            with torch.cuda.stream(s2):
+                e1.wait(s2)
+                b = a + 1
+                e2.record(s2)
+
+            # Stage 3 on s3, depends on s2
+            with torch.cuda.stream(s3):
+                e2.wait(s3)
+                c = b * 3
+
+            s1.synchronize()
+            s2.synchronize()
+            s3.synchronize()
+            return c
+
+        x = torch.randn(1024, device="cuda")
+
+        expected = fn(x)
+        compiled_fn = torch.compile(fn)
+        result, (code,) = run_and_get_code(compiled_fn, x)
+
+        self.assertEqual(result, expected)
+
+        # Verify three-stage pipeline with 3 streams
+        self.assertGreaterEqual(code.count("torch.cuda.stream"), 3)
+        self.assertGreaterEqual(code.count("record_event"), 2)
+        self.assertGreaterEqual(code.count("wait_event"), 2)
+
+    def test_parallel_streams_join(self):
+        """Test parallel work on multiple streams joining at the end."""
+        from torch._inductor.utils import run_and_get_code
+
+        def fn(x):
+            s1 = torch.cuda.Stream()
+            s2 = torch.cuda.Stream()
+            s3 = torch.cuda.Stream()
+            e1 = torch.cuda.Event()
+            e2 = torch.cuda.Event()
+            e3 = torch.cuda.Event()
+
+            # Parallel work on three streams
+            with torch.cuda.stream(s1):
+                a = x * 2
+                e1.record(s1)
+
+            with torch.cuda.stream(s2):
+                b = x * 3
+                e2.record(s2)
+
+            with torch.cuda.stream(s3):
+                c = x * 4
+                e3.record(s3)
+
+            # Join all results on default stream
+            e1.wait()
+            e2.wait()
+            e3.wait()
+            result = a + b + c
+
+            s1.synchronize()
+            s2.synchronize()
+            s3.synchronize()
+            return result
+
+        x = torch.randn(1024, device="cuda")
+
+        expected = fn(x)
+        compiled_fn = torch.compile(fn)
+        result, (code,) = run_and_get_code(compiled_fn, x)
+
+        self.assertEqual(result, expected)
+
+        # Verify parallel streams joining
+        self.assertGreaterEqual(code.count("torch.cuda.stream"), 3)
+        self.assertGreaterEqual(code.count("record_event"), 3)
+        self.assertGreaterEqual(code.count("wait_event"), 3)
+
+    def test_fan_out_fan_in(self):
+        """Test fan-out from one stream to multiple, then fan-in."""
+        from torch._inductor.utils import run_and_get_code
+
+        def fn(x):
+            s1 = torch.cuda.Stream()
+            s2 = torch.cuda.Stream()
+            e_start = torch.cuda.Event()
+            e1 = torch.cuda.Event()
+            e2 = torch.cuda.Event()
+
+            # Initial work on default stream
+            a = x * 2
+            e_start.record()
+
+            # Fan out to s1 and s2
+            with torch.cuda.stream(s1):
+                e_start.wait(s1)
+                b = a + 1
+                e1.record(s1)
+
+            with torch.cuda.stream(s2):
+                e_start.wait(s2)
+                c = a + 2
+                e2.record(s2)
+
+            # Fan in on default stream
+            e1.wait()
+            e2.wait()
+            result = b + c
+
+            s1.synchronize()
+            s2.synchronize()
+            return result
+
+        x = torch.randn(1024, device="cuda")
+
+        expected = fn(x)
+        compiled_fn = torch.compile(fn)
+        result, (code,) = run_and_get_code(compiled_fn, x)
+
+        self.assertEqual(result, expected)
+
+        # Verify fan-out/fan-in pattern
+        self.assertGreaterEqual(code.count("record_event"), 3)
+        self.assertGreaterEqual(code.count("wait_event"), 4)
+
+    def test_four_streams_diamond(self):
+        """Test diamond pattern: one start, two parallel, one end."""
+        from torch._inductor.utils import run_and_get_code
+
+        def fn(x):
+            s1 = torch.cuda.Stream()
+            s2 = torch.cuda.Stream()
+            s3 = torch.cuda.Stream()
+            e_start = torch.cuda.Event()
+            e1 = torch.cuda.Event()
+            e2 = torch.cuda.Event()
+
+            # Start on default stream
+            a = x + 1
+            e_start.record()
+
+            # Parallel branches on s1 and s2
+            with torch.cuda.stream(s1):
+                e_start.wait(s1)
+                b = a * 2
+                e1.record(s1)
+
+            with torch.cuda.stream(s2):
+                e_start.wait(s2)
+                c = a * 3
+                e2.record(s2)
+
+            # Join on s3
+            with torch.cuda.stream(s3):
+                e1.wait(s3)
+                e2.wait(s3)
+                d = b + c
+
+            s1.synchronize()
+            s2.synchronize()
+            s3.synchronize()
+            return d
+
+        x = torch.randn(1024, device="cuda")
+
+        expected = fn(x)
+        compiled_fn = torch.compile(fn)
+        result, (code,) = run_and_get_code(compiled_fn, x)
+
+        self.assertEqual(result, expected)
+
+        # Verify diamond pattern
+        self.assertGreaterEqual(code.count("torch.cuda.stream"), 3)
+        self.assertGreaterEqual(code.count("record_event"), 3)
+        self.assertGreaterEqual(code.count("wait_event"), 4)
+
+    def test_stream_reuse_across_iterations(self):
+        """Test that streams can be reused across loop iterations."""
+        from torch._inductor.utils import run_and_get_code
+
+        def fn(x):
+            s = torch.cuda.Stream()
+            event = torch.cuda.Event()
+            result = x
+
+            for _ in range(3):
+                with torch.cuda.stream(s):
+                    result = result * 2
+                    event.record(s)
+                event.wait()
+
+            s.synchronize()
+            return result
+
+        x = torch.randn(1024, device="cuda")
+
+        expected = fn(x)
+        compiled_fn = torch.compile(fn)
+        result, (code,) = run_and_get_code(compiled_fn, x)
+
+        self.assertEqual(result, expected)
+
+        # Verify stream reuse in loop — events survive compilation
+        self.assertIn("torch.cuda.stream", code)
+        self.assertIn("record_event", code)
+        self.assertIn("wait_event", code)
+
+    def test_no_fusion_across_streams(self):
+        """Test that operations on different streams are not fused together."""
+        from torch._inductor.utils import run_and_get_code
+
+        def fn(x):
+            s1 = torch.cuda.Stream()
+            s2 = torch.cuda.Stream()
+            e1 = torch.cuda.Event()
+            e2 = torch.cuda.Event()
+
+            # These could be fused if on same stream, but should NOT be fused
+            # since they're on different streams
+            with torch.cuda.stream(s1):
+                # Multiple pointwise ops that would normally fuse
+                a = x * 2
+                b = a + 1
+                c = b * 3
+                e1.record(s1)
+
+            with torch.cuda.stream(s2):
+                # Another set of pointwise ops
+                d = x * 4
+                e = d + 2
+                f = e * 5
+                e2.record(s2)
+
+            e1.wait()
+            e2.wait()
+            return c + f
+
+        x = torch.randn(1024, device="cuda")
+
+        expected = fn(x)
+        compiled_fn = torch.compile(fn)
+        torch._inductor.metrics.reset()
+        result, (code,) = run_and_get_code(compiled_fn, x)
+
+        self.assertEqual(result, expected)
+
+        # 3 kernels: s1 pointwise, s2 pointwise, and the final add on
+        # the default stream (which is a third stream context).
+        self.assertEqual(torch._inductor.metrics.generated_kernel_count, 3)
+
+    def test_no_fusion_across_streams_with_dependency(self):
+        """Test no fusion when there's a data dependency across streams."""
+        from torch._inductor.utils import run_and_get_code
+
+        def fn(x):
+            s = torch.cuda.Stream()
+            event = torch.cuda.Event()
+
+            # Work on default stream
+            a = x * 2
+            b = a + 1
+            event.record()
+
+            # Work on side stream - depends on default stream
+            with torch.cuda.stream(s):
+                event.wait()
+                c = b * 3  # depends on b from default stream
+                d = c + 1
+
+            s.synchronize()
+            return d
+
+        x = torch.randn(1024, device="cuda")
+
+        expected = fn(x)
+        compiled_fn = torch.compile(fn)
+        torch._inductor.metrics.reset()
+        result, (code,) = run_and_get_code(compiled_fn, x)
+
+        self.assertEqual(result, expected)
+
+        self.assertEqual(torch._inductor.metrics.generated_kernel_count, 2)
+
+    def test_fusion_within_same_stream(self):
+        """Test that fusion still works for operations within the same stream."""
+        from torch._inductor.utils import run_and_get_code
+
+        def fn(x):
+            s = torch.cuda.Stream()
+
+            with torch.cuda.stream(s):
+                # Multiple pointwise ops on same stream - should fuse
+                a = x * 2
+                b = a + 1
+                c = b * 3
+                d = c + 2
+
+            s.synchronize()
+            return d
+
+        x = torch.randn(1024, device="cuda")
+
+        expected = fn(x)
+        compiled_fn = torch.compile(fn)
+        torch._inductor.metrics.reset()
+        result, (code,) = run_and_get_code(compiled_fn, x)
+
+        self.assertEqual(result, expected)
+
+        # All pointwise ops on same stream should fuse into 1 kernel
+        self.assertEqual(torch._inductor.metrics.generated_kernel_count, 1)
+
+    def test_no_fusion_simple_dependency_across_streams(self):
+        """Regression: a single pointwise consumed across a stream boundary must not fuse."""
+        from torch._inductor.utils import run_and_get_code
+
+        def fn(x):
+            s1 = torch.cuda.Stream()
+            s2 = torch.cuda.Stream()
+
+            with torch.cuda.stream(s1):
+                a = x + 1
+            e = s1.record_event()
+            s2.wait_event(e)
+            with torch.cuda.stream(s2):
+                b = a * 2
+            s1.synchronize()
+            s2.synchronize()
+            return b
+
+        x = torch.randn(1024, device="cuda")
+
+        expected = fn(x)
+        compiled_fn = torch.compile(fn)
+        torch._inductor.metrics.reset()
+        result, (code,) = run_and_get_code(compiled_fn, x)
+
+        self.assertEqual(result, expected)
+
+        # Must be 2 separate kernels on 2 streams, not fused into 1
+        self.assertEqual(torch._inductor.metrics.generated_kernel_count, 2)
+
+    @torch._inductor.config.patch(combo_kernels=True)
+    def test_no_combo_kernel_fusion_across_streams(self):
+        """Combo kernels must not group nodes on different streams."""
+        from torch._inductor.utils import run_and_get_code
+
+        def fn(x, y, z, w):
+            s = torch.cuda.Stream()
+            event = torch.cuda.Event()
+
+            # Independent pointwise ops on different streams at the same
+            # topological level — combo kernels must not merge them.
+            a = x + y
+            event.record()
+            with torch.cuda.stream(s):
+                event.wait()
+                b = z + w
+            s.synchronize()
+            return a, b
+
+        x = torch.randn(1024, device="cuda")
+        y = torch.randn(1024, device="cuda")
+        z = torch.randn(1024, device="cuda")
+        w = torch.randn(1024, device="cuda")
+
+        expected = fn(x, y, z, w)
+        compiled_fn = torch.compile(fn)
+        torch._inductor.metrics.reset()
+        result, (code,) = run_and_get_code(compiled_fn, x, y, z, w)
+
+        self.assertEqual(result, expected)
+        # 2 kernels: one per stream. Without the stream-aware fix, combo
+        # kernels would merge them into 1.
+        self.assertEqual(torch._inductor.metrics.generated_kernel_count, 2)
+
+    @torch._inductor.config.patch(combo_kernels=True)
+    def test_combo_kernel_fusion_within_same_stream(self):
+        """Combo kernels should still group independent nodes on the same stream."""
+        from torch._inductor.utils import run_and_get_code
+
+        def fn(x, y):
+            s = torch.cuda.Stream()
+
+            with torch.cuda.stream(s):
+                # Two independent pointwise ops on the same stream — eligible
+                # for combo kernel fusion.
+                a = x * 2
+                b = y * 3
+
+            s.synchronize()
+            return a + b
+
+        x = torch.randn(1024, device="cuda")
+        y = torch.randn(1024, device="cuda")
+
+        expected = fn(x, y)
+        compiled_fn = torch.compile(fn)
+        torch._inductor.metrics.reset()
+        result, (code,) = run_and_get_code(compiled_fn, x, y)
+
+        self.assertEqual(result, expected)
+        # With combo kernels, the two independent ops on the same stream
+        # should be combined, yielding fewer kernels than without.
+        self.assertLessEqual(torch._inductor.metrics.generated_kernel_count, 2)
+
+    def test_cross_stream_stride_copy(self):
+        """A contiguous copy forced by a non-contiguous slice across streams
+        must run on the consumer's stream, not the producer's."""
+        from torch._inductor.utils import run_and_get_code
+
+        def fn(x):
+            s1 = torch.cuda.Stream()
+            s2 = torch.cuda.Stream()
+
+            with torch.cuda.stream(s1):
+                a = x + 1
+                b = a[:, ::2]  # non-contiguous slice
+            e = s1.record_event()
+            s2.wait_event(e)
+            with torch.cuda.stream(s2):
+                c = b.contiguous()
+                d = c + 1
+            s2.synchronize()
+            return d
+
+        x = torch.randn(64, 64, device="cuda")
+
+        expected = fn(x)
+        compiled_fn = torch.compile(fn)
+        torch._inductor.metrics.reset()
+        result, (code,) = run_and_get_code(compiled_fn, x)
+
+        self.assertEqual(result, expected)
+        self.assertEqual(torch._inductor.metrics.generated_kernel_count, 2)
+
+        # Verify: s1 gets the pointwise (x+1), s2 gets the fused copy + add.
+        # The contiguous copy is fused into the s2 triton kernel (which reads
+        # from the s1 output buffer with strided indexing).  If the copy were
+        # incorrectly placed on s1, we'd see 2 kernels on s1 instead of 1.
+        wrapper = _extract_wrapper_body(code)
+        lines = wrapper.split("\n")
+        current_stream = None
+        stream_kernels: dict[str | None, list[str]] = {}
+        for line in lines:
+            stripped = line.strip()
+            if "with torch.cuda.stream(" in stripped:
+                if "stream1" in stripped:
+                    current_stream = "s1"
+                elif "stream2" in stripped:
+                    current_stream = "s2"
+                elif "default_stream" in stripped:
+                    current_stream = "default"
+            elif ".run(" in stripped:
+                stream_kernels.setdefault(current_stream, []).append(stripped)
+
+        self.assertEqual(
+            len(stream_kernels.get("s1", [])),
+            1,
+            f"Expected 1 kernel on s1, got: {stream_kernels}",
+        )
+        self.assertEqual(
+            len(stream_kernels.get("s2", [])),
+            1,
+            f"Expected 1 kernel on s2, got: {stream_kernels}",
+        )
+
+    def test_no_buffer_reuse_across_streams(self):
+        """Buffer produced on one stream must not be reused in-place on another."""
+        from torch._inductor.utils import run_and_get_code
+
+        def fn(x):
+            s1 = torch.cuda.Stream()
+            s2 = torch.cuda.Stream()
+            with torch.cuda.stream(s1):
+                a = x + 1
+            e = s1.record_event()
+            s2.wait_event(e)
+            with torch.cuda.stream(s2):
+                b = a + 2
+            s2.synchronize()
+            return b
+
+        x = torch.randn(1024, device="cuda")
+        expected = fn(x)
+        compiled_fn = torch.compile(fn)
+        result, (code,) = run_and_get_code(compiled_fn, x)
+        self.assertEqual(result, expected)
+        # The second kernel should allocate a fresh buffer, not reuse
+        # the one produced on the other stream
+        wrapper = _extract_wrapper_body(code)
+        self.assertIn("record_event", wrapper)
+        self.assertIn("wait_event", wrapper)
+        self.assertNotIn("buf0; del buf0", wrapper)
+
+    def test_stream_record_wait_event_not_dropped(self):
+        """stream.record_event() and stream.wait_event() must survive compilation."""
+        from torch._inductor.utils import run_and_get_code
+
+        def fn(x):
+            s1 = torch.cuda.Stream()
+            s2 = torch.cuda.Stream()
+            with torch.cuda.stream(s1):
+                a = x + 1
+            e = s1.record_event()
+            s2.wait_event(e)
+            with torch.cuda.stream(s2):
+                b = a * 2
+            s2.synchronize()
+            return b
+
+        x = torch.randn(1024, device="cuda")
+        expected = fn(x)
+        compiled_fn = torch.compile(fn)
+        result, (code,) = run_and_get_code(compiled_fn, x)
+        self.assertEqual(result, expected)
+        self.assertIn("record_event", code)
+        self.assertIn("wait_event", code)
+        self.assertIn("synchronize_stream", code)
+
+    def test_stream_synchronize_not_dropped(self):
+        """stream.synchronize() must survive compilation and appear in wrapper code."""
+        from torch._inductor.utils import run_and_get_code
+
+        def fn(x):
+            s = torch.cuda.Stream()
+            with torch.cuda.stream(s):
+                a = x + 1
+            s.synchronize()
+            b = a * 2
+            return b
+
+        x = torch.randn(1024, device="cuda")
+        expected = fn(x)
+        compiled_fn = torch.compile(fn)
+        result, (code,) = run_and_get_code(compiled_fn, x)
+        self.assertEqual(result, expected)
+        self.assertIn("synchronize_stream", code)
+
+    def test_stream_wait_stream_not_dropped(self):
+        """stream.wait_stream() must survive compilation and appear in wrapper code."""
+        from torch._inductor.utils import run_and_get_code
+
+        def fn(x):
+            s1 = torch.cuda.Stream()
+            s2 = torch.cuda.Stream()
+            with torch.cuda.stream(s1):
+                a = x + 1
+            s2.wait_stream(s1)
+            with torch.cuda.stream(s2):
+                b = a * 2
+            s2.synchronize()
+            return b
+
+        x = torch.randn(1024, device="cuda")
+        expected = fn(x)
+        compiled_fn = torch.compile(fn)
+        result, (code,) = run_and_get_code(compiled_fn, x)
+        self.assertEqual(result, expected)
+        self.assertIn("wait_stream", code)
+        self.assertIn("synchronize_stream", code)
+
+    def test_codegen_structure_single_stream(self):
+        """Verify wrapper structure for pointwise ops with one side stream."""
+        from torch._inductor.utils import run_and_get_code
+
+        def fn(x):
+            s = torch.cuda.Stream()
+            a = x * 2
+            with torch.cuda.stream(s):
+                b = x * 3
+            s.synchronize()
+            return a + b
+
+        x = torch.randn(1024, device="cuda")
+        expected = fn(x)
+        counter = CompileCounterWithBackend("inductor")
+        compiled_fn = torch.compile(fn, backend=counter)
+        result, (code,) = run_and_get_code(compiled_fn, x)
+        self.assertEqual(result, expected)
+
+        self.assertExpectedInline(
+            normalize_gm(counter.graphs[0].print_readable(print_output=False)),
+            """\
+class GraphModule(torch.nn.Module):
+    def forward(self, L_x_: "f32[1024]"):
+        l_x_ = L_x_
+
+        get_external_object_by_index = torch__dynamo_graph_bytecode_inputs_get_external_object_by_index(1);  get_external_object_by_index = None
+
+        a: "f32[1024]" = l_x_ * 2
+
+        b: "f32[1024]" = l_x_ * 3;  l_x_ = None
+
+        synchronize_stream = torch.ops.streams.synchronize_stream(1);  synchronize_stream = None
+
+        add: "f32[1024]" = a + b;  a = b = None
+        return (add,)
+""",  # noqa: B950
+        )
+
+        wrapper_body = _extract_wrapper_body(code)
+        self.assertExpectedInline(
+            wrapper_body,
+            """\
+arg0_1, = args
+with torch.cuda._DeviceGuard(0):
+    torch.cuda.set_device(0)
+    default_stream = torch.cuda.current_stream()
+    stream1 = get_external_object_by_index(1)
+    with torch.cuda.stream(stream1):
+        arg0_1 = copy_misaligned(arg0_1)
+        buf0 = empty_strided_cuda((1024, ), (1, ), torch.float32)
+        raw_stream = get_raw_stream(0)
+        triton_kernel.run(arg0_1, buf0, 1024, stream=raw_stream)
+    with torch.cuda.stream(default_stream):
+        buf3 = empty_strided_cuda((1024, ), (1, ), torch.float32)
+        raw_stream0 = get_raw_stream(0)
+        triton_kernel.run(arg0_1, buf0, buf3, 1024, stream=raw_stream0)
+        torch.ops.streams.synchronize_stream.default(1)
+    return (buf3, )""",
+        )
+
+    def test_codegen_structure_pipeline(self):
+        """Verify wrapper structure for two-stage matmul pipeline."""
+        from torch._inductor.utils import run_and_get_code
+
+        def fn(x, w1, w2):
+            s = torch.cuda.Stream()
+            event = torch.cuda.Event()
+            a = x @ w1
+            event.record()
+            with torch.cuda.stream(s):
+                event.wait()
+                b = a @ w2
+            s.synchronize()
+            return b
+
+        x = torch.randn(32, 32, device="cuda")
+        w1 = torch.randn(32, 32, device="cuda")
+        w2 = torch.randn(32, 32, device="cuda")
+        expected = fn(x, w1, w2)
+        counter = CompileCounterWithBackend("inductor")
+        compiled_fn = torch.compile(fn, backend=counter)
+        result, (code,) = run_and_get_code(compiled_fn, x, w1, w2)
+        self.assertEqual(result, expected)
+
+        self.assertExpectedInline(
+            normalize_gm(counter.graphs[0].print_readable(print_output=False)),
+            """\
+class GraphModule(torch.nn.Module):
+    def forward(self, L_x_: "f32[32, 32]", L_w1_: "f32[32, 32]", L_w2_: "f32[32, 32]"):
+        l_x_ = L_x_
+        l_w1_ = L_w1_
+        l_w2_ = L_w2_
+
+        get_external_object_by_index = torch__dynamo_graph_bytecode_inputs_get_external_object_by_index(1);  get_external_object_by_index = None
+
+        get_external_object_by_index_1 = torch__dynamo_graph_bytecode_inputs_get_external_object_by_index(2);  get_external_object_by_index_1 = None
+
+        a: "f32[32, 32]" = l_x_ @ l_w1_;  l_x_ = l_w1_ = None
+
+        get_external_object_by_index_2 = torch__dynamo_graph_bytecode_inputs_get_external_object_by_index(0);  get_external_object_by_index_2 = None
+        record_event = torch.ops.streams.record_event(2, 0);  record_event = None
+
+        wait_event = torch.ops.streams.wait_event(2, 1);  wait_event = None
+
+        b: "f32[32, 32]" = a @ l_w2_;  a = l_w2_ = None
+
+        synchronize_stream = torch.ops.streams.synchronize_stream(1);  synchronize_stream = None
+        return (b,)
+""",  # noqa: B950
+        )
+
+        wrapper_body = _extract_wrapper_body(code)
+        FileCheck().run(
+            """\
+# CHECK: with torch.cuda.stream(default_stream):
+# CHECK: copy_misaligned
+# CHECK: extern_kernels.mm(
+# CHECK: record_event
+# CHECK: with torch.cuda.stream(stream1):
+# CHECK: wait_event
+# CHECK: copy_misaligned
+# CHECK: extern_kernels.mm(
+# CHECK: with torch.cuda.stream(default_stream):
+# CHECK: synchronize_stream""",
+            wrapper_body,
+        )
+
+    def test_codegen_structure_three_stream_pipeline(self):
+        """Verify wrapper structure for three-stage matmul pipeline."""
+        from torch._inductor.utils import run_and_get_code
+
+        def fn(x, w1, w2, w3):
+            s1 = torch.cuda.Stream()
+            s2 = torch.cuda.Stream()
+            s3 = torch.cuda.Stream()
+            e1 = torch.cuda.Event()
+            e2 = torch.cuda.Event()
+            with torch.cuda.stream(s1):
+                a = x @ w1
+                e1.record(s1)
+            with torch.cuda.stream(s2):
+                e1.wait(s2)
+                b = a @ w2
+                e2.record(s2)
+            with torch.cuda.stream(s3):
+                e2.wait(s3)
+                c = b @ w3
+            s1.synchronize()
+            s2.synchronize()
+            s3.synchronize()
+            return c
+
+        x = torch.randn(32, 32, device="cuda")
+        w1 = torch.randn(32, 32, device="cuda")
+        w2 = torch.randn(32, 32, device="cuda")
+        w3 = torch.randn(32, 32, device="cuda")
+        expected = fn(x, w1, w2, w3)
+        counter = CompileCounterWithBackend("inductor")
+        compiled_fn = torch.compile(fn, backend=counter)
+        result, (code,) = run_and_get_code(compiled_fn, x, w1, w2, w3)
+        self.assertEqual(result, expected)
+
+        self.assertExpectedInline(
+            normalize_gm(counter.graphs[0].print_readable(print_output=False)),
+            """\
+class GraphModule(torch.nn.Module):
+    def forward(self, L_x_: "f32[32, 32]", L_w1_: "f32[32, 32]", L_w2_: "f32[32, 32]", L_w3_: "f32[32, 32]"):
+        l_x_ = L_x_
+        l_w1_ = L_w1_
+        l_w2_ = L_w2_
+        l_w3_ = L_w3_
+
+        get_external_object_by_index = torch__dynamo_graph_bytecode_inputs_get_external_object_by_index(1);  get_external_object_by_index = None
+
+        get_external_object_by_index_1 = torch__dynamo_graph_bytecode_inputs_get_external_object_by_index(2);  get_external_object_by_index_1 = None
+
+        get_external_object_by_index_2 = torch__dynamo_graph_bytecode_inputs_get_external_object_by_index(3);  get_external_object_by_index_2 = None
+
+        get_external_object_by_index_3 = torch__dynamo_graph_bytecode_inputs_get_external_object_by_index(4);  get_external_object_by_index_3 = None
+
+        get_external_object_by_index_4 = torch__dynamo_graph_bytecode_inputs_get_external_object_by_index(5);  get_external_object_by_index_4 = None
+
+        a: "f32[32, 32]" = l_x_ @ l_w1_;  l_x_ = l_w1_ = None
+
+        record_event = torch.ops.streams.record_event(4, 1);  record_event = None
+
+        wait_event = torch.ops.streams.wait_event(4, 2);  wait_event = None
+
+        b: "f32[32, 32]" = a @ l_w2_;  a = l_w2_ = None
+
+        record_event_1 = torch.ops.streams.record_event(5, 2);  record_event_1 = None
+
+        wait_event_1 = torch.ops.streams.wait_event(5, 3);  wait_event_1 = None
+
+        c: "f32[32, 32]" = b @ l_w3_;  b = l_w3_ = None
+
+        synchronize_stream = torch.ops.streams.synchronize_stream(1);  synchronize_stream = None
+
+        synchronize_stream_1 = torch.ops.streams.synchronize_stream(2);  synchronize_stream_1 = None
+
+        synchronize_stream_2 = torch.ops.streams.synchronize_stream(3);  synchronize_stream_2 = None
+        return (c,)
+""",  # noqa: B950
+        )
+
+        wrapper_body = _extract_wrapper_body(code)
+        FileCheck().run(
+            """\
+# CHECK: with torch.cuda.stream(stream1):
+# CHECK: copy_misaligned
+# CHECK: extern_kernels.mm(
+# CHECK: record_event
+# CHECK: with torch.cuda.stream(stream2):
+# CHECK: wait_event
+# CHECK: copy_misaligned
+# CHECK: extern_kernels.mm(
+# CHECK: record_event
+# CHECK: with torch.cuda.stream(stream3):
+# CHECK: wait_event
+# CHECK: copy_misaligned
+# CHECK: extern_kernels.mm(
+# CHECK: with torch.cuda.stream(default_stream):
+# CHECK: synchronize_stream
+# CHECK: synchronize_stream
+# CHECK: synchronize_stream""",
+            wrapper_body,
+        )
+
+    def test_codegen_structure_parallel_matmuls(self):
+        """Verify wrapper structure for parallel matmuls with join."""
+        from torch._inductor.utils import run_and_get_code
+
+        def fn(x, w1, w2):
+            s1 = torch.cuda.Stream()
+            s2 = torch.cuda.Stream()
+            e1 = torch.cuda.Event()
+            e2 = torch.cuda.Event()
+            with torch.cuda.stream(s1):
+                a = x @ w1
+                e1.record(s1)
+            with torch.cuda.stream(s2):
+                b = x @ w2
+                e2.record(s2)
+            e1.wait()
+            e2.wait()
+            c = a + b
+            s1.synchronize()
+            s2.synchronize()
+            return c
+
+        x = torch.randn(32, 32, device="cuda")
+        w1 = torch.randn(32, 32, device="cuda")
+        w2 = torch.randn(32, 32, device="cuda")
+        expected = fn(x, w1, w2)
+        counter = CompileCounterWithBackend("inductor")
+        compiled_fn = torch.compile(fn, backend=counter)
+        result, (code,) = run_and_get_code(compiled_fn, x, w1, w2)
+        self.assertEqual(result, expected)
+
+        self.assertExpectedInline(
+            normalize_gm(counter.graphs[0].print_readable(print_output=False)),
+            """\
+class GraphModule(torch.nn.Module):
+    def forward(self, L_x_: "f32[32, 32]", L_w1_: "f32[32, 32]", L_w2_: "f32[32, 32]"):
+        l_x_ = L_x_
+        l_w1_ = L_w1_
+        l_w2_ = L_w2_
+
+        get_external_object_by_index = torch__dynamo_graph_bytecode_inputs_get_external_object_by_index(1);  get_external_object_by_index = None
+
+        get_external_object_by_index_1 = torch__dynamo_graph_bytecode_inputs_get_external_object_by_index(2);  get_external_object_by_index_1 = None
+
+        get_external_object_by_index_2 = torch__dynamo_graph_bytecode_inputs_get_external_object_by_index(3);  get_external_object_by_index_2 = None
+
+        get_external_object_by_index_3 = torch__dynamo_graph_bytecode_inputs_get_external_object_by_index(4);  get_external_object_by_index_3 = None
+
+        a: "f32[32, 32]" = l_x_ @ l_w1_;  l_w1_ = None
+
+        record_event = torch.ops.streams.record_event(3, 1);  record_event = None
+
+        b: "f32[32, 32]" = l_x_ @ l_w2_;  l_x_ = l_w2_ = None
+
+        record_event_1 = torch.ops.streams.record_event(4, 2);  record_event_1 = None
+
+        wait_event = torch.ops.streams.wait_event(3, 0);  wait_event = None
+
+        wait_event_1 = torch.ops.streams.wait_event(4, 0);  wait_event_1 = None
+
+        c: "f32[32, 32]" = a + b;  a = b = None
+
+        synchronize_stream = torch.ops.streams.synchronize_stream(1);  synchronize_stream = None
+
+        synchronize_stream_1 = torch.ops.streams.synchronize_stream(2);  synchronize_stream_1 = None
+        return (c,)
+""",  # noqa: B950
+        )
+
+        wrapper_body = _extract_wrapper_body(code)
+        FileCheck().run(
+            """\
+# CHECK: with torch.cuda.stream(stream1):
+# CHECK: copy_misaligned
+# CHECK: extern_kernels.mm(
+# CHECK: record_event
+# CHECK: with torch.cuda.stream(stream2):
+# CHECK: copy_misaligned
+# CHECK: extern_kernels.mm(
+# CHECK: record_event
+# CHECK: with torch.cuda.stream(default_stream):
+# CHECK: wait_event
+# CHECK: triton_kernel.run(
+# CHECK: synchronize_stream
+# CHECK: synchronize_stream""",
+            wrapper_body,
+        )
+
+
+@unittest.skipUnless(TEST_CUDA, "requires CUDA")
+class TestStreamOrderingStress(InductorTestCase):
+    """Stress tests verifying that interleaved event record/wait ops
+    produce correct ordering under compilation.  Each test uses large
+    matmuls so there is real GPU work, and repeats many iterations so
+    that race conditions (if ordering is wrong) surface reliably."""
+
+    N = 4096  # matrix size — big enough for real GPU work
+    ITERS = 20  # repetitions per test
+
+    def _check_compiled_matches_eager(self, fn, *args):
+        """Run fn eagerly and compiled, assert results match over ITERS runs."""
+        compiled_fn = torch.compile(fn)
+        for _ in range(self.ITERS):
+            expected = fn(*args)
+            actual = compiled_fn(*args)
+            # Full device sync as a safety net to ensure all stream work
+            # is visible before comparing results.
+            torch.cuda.synchronize()
+            if not isinstance(expected, (tuple, list)):
+                expected, actual = [expected], [actual]
+            for e, a in zip(expected, actual):
+                self.assertEqual(a, e)
+
+    @staticmethod
+    def _heavy_matmul_chain(x, w, depth=8):
+        """Chain of matmuls to create substantial GPU work (~ms).
+        Used to widen the race window between streams so that missing
+        synchronization is observable."""
+        h = x
+        for _ in range(depth):
+            h = h @ w
+        return h
+
+    # ------------------------------------------------------------------
+    # 1. Race: producer does heavy work, consumer reads the result.
+    #    Without the event.wait() the consumer would launch immediately
+    #    and read stale memory because the producer chain hasn't finished.
+    # ------------------------------------------------------------------
+    def test_race_producer_consumer(self):
+        N = self.N
+
+        def fn(x, w):
+            s = torch.cuda.Stream()
+            e = torch.cuda.Event()
+
+            # Heavy producer on default stream — takes real GPU time
+            a = TestStreamOrderingStress._heavy_matmul_chain(x, w)
+            e.record()
+
+            with torch.cuda.stream(s):
+                e.wait()  # removing this would cause a race
+                b = a + 1
+
+            s.synchronize()
+            return b
+
+        x = torch.randn(N, N, device="cuda")
+        w = torch.eye(N, device="cuda") * 0.9  # use scaled identity for stability
+        self._check_compiled_matches_eager(fn, x, w)
+
+    # ------------------------------------------------------------------
+    # 2. Race: ping-pong where each direction has heavy work.
+    #    Both event.wait() calls are load-bearing.
+    # ------------------------------------------------------------------
+    def test_race_ping_pong(self):
+        N = self.N
+
+        def fn(x, w):
+            s1 = torch.cuda.Stream()
+            s2 = torch.cuda.Stream()
+            e1 = torch.cuda.Event()
+            e2 = torch.cuda.Event()
+
+            with torch.cuda.stream(s1):
+                a = TestStreamOrderingStress._heavy_matmul_chain(x, w)
+                e1.record(s1)
+
+            with torch.cuda.stream(s2):
+                e1.wait(s2)
+                b = TestStreamOrderingStress._heavy_matmul_chain(a, w)
+                e2.record(s2)
+
+            with torch.cuda.stream(s1):
+                e2.wait(s1)
+                c = b + a
+
+            s1.synchronize()
+            s2.synchronize()
+            return c
+
+        x = torch.randn(N, N, device="cuda")
+        w = torch.eye(N, device="cuda") * 0.9
+        self._check_compiled_matches_eager(fn, x, w)
+
+    # ------------------------------------------------------------------
+    # 3. Race: fan-out where the producer is slow.
+    #    All three consumers depend on the producer finishing.
+    # ------------------------------------------------------------------
+    def test_race_fan_out(self):
+        N = self.N
+
+        def fn(x, w):
+            s1 = torch.cuda.Stream()
+            s2 = torch.cuda.Stream()
+            s3 = torch.cuda.Stream()
+            e = torch.cuda.Event()
+            e1 = torch.cuda.Event()
+            e2 = torch.cuda.Event()
+            e3 = torch.cuda.Event()
+
+            # Slow producer
+            a = TestStreamOrderingStress._heavy_matmul_chain(x, w)
+            e.record()
+
+            with torch.cuda.stream(s1):
+                e.wait()
+                r1 = a * 2
+                e1.record(s1)
+
+            with torch.cuda.stream(s2):
+                e.wait()
+                r2 = a * 3
+                e2.record(s2)
+
+            with torch.cuda.stream(s3):
+                e.wait()
+                r3 = a * 4
+                e3.record(s3)
+
+            e1.wait()
+            e2.wait()
+            e3.wait()
+            result = r1 + r2 + r3
+
+            s1.synchronize()
+            s2.synchronize()
+            s3.synchronize()
+            return result
+
+        x = torch.randn(N, N, device="cuda")
+        w = torch.eye(N, device="cuda") * 0.9
+        self._check_compiled_matches_eager(fn, x, w)
+
+    # ------------------------------------------------------------------
+    # 4. Race: diamond pattern with heavy work on both branches.
+    #    The join must wait for both branches.
+    # ------------------------------------------------------------------
+    def test_race_diamond(self):
+        N = self.N
+
+        def fn(x, w):
+            s1 = torch.cuda.Stream()
+            s2 = torch.cuda.Stream()
+            e_fork = torch.cuda.Event()
+            e1 = torch.cuda.Event()
+            e2 = torch.cuda.Event()
+
+            base = x @ w
+            e_fork.record()
+
+            with torch.cuda.stream(s1):
+                e_fork.wait()
+                branch1 = TestStreamOrderingStress._heavy_matmul_chain(
+                    torch.relu(base), w
+                )
+                e1.record(s1)
+
+            with torch.cuda.stream(s2):
+                e_fork.wait()
+                branch2 = TestStreamOrderingStress._heavy_matmul_chain(
+                    torch.sigmoid(base), w
+                )
+                e2.record(s2)
+
+            e1.wait()
+            e2.wait()
+            result = branch1 + branch2
+
+            s1.synchronize()
+            s2.synchronize()
+            return result
+
+        x = torch.randn(N, N, device="cuda")
+        w = torch.eye(N, device="cuda") * 0.5
+        self._check_compiled_matches_eager(fn, x, w)
+
+    # ------------------------------------------------------------------
+    # 5. Race: 4-stage pipeline where each stage is heavy.
+    #    Every event.wait() is load-bearing.
+    # ------------------------------------------------------------------
+    def test_race_pipeline(self):
+        N = self.N
+
+        def fn(x, w):
+            streams = [torch.cuda.Stream() for _ in range(4)]
+            events = [torch.cuda.Event() for _ in range(3)]
+
+            with torch.cuda.stream(streams[0]):
+                h = TestStreamOrderingStress._heavy_matmul_chain(x, w, depth=4)
+                events[0].record(streams[0])
+
+            with torch.cuda.stream(streams[1]):
+                events[0].wait(streams[1])
+                h = TestStreamOrderingStress._heavy_matmul_chain(h, w, depth=4)
+                events[1].record(streams[1])
+
+            with torch.cuda.stream(streams[2]):
+                events[1].wait(streams[2])
+                h = TestStreamOrderingStress._heavy_matmul_chain(h, w, depth=4)
+                events[2].record(streams[2])
+
+            with torch.cuda.stream(streams[3]):
+                events[2].wait(streams[3])
+                h = h + x  # quick consumer — races if wait is missing
+
+            for s in streams:
+                s.synchronize()
+            return h
+
+        x = torch.randn(N, N, device="cuda")
+        w = torch.eye(N, device="cuda") * 0.9
+        self._check_compiled_matches_eager(fn, x, w)
+
+    # ------------------------------------------------------------------
+    # 6. Race: back-to-back sync, both directions carry heavy work
+    # ------------------------------------------------------------------
+    def test_race_back_to_back(self):
+        N = self.N
+
+        def fn(x, w):
+            s = torch.cuda.Stream()
+            e1 = torch.cuda.Event()
+            e2 = torch.cuda.Event()
+
+            a = TestStreamOrderingStress._heavy_matmul_chain(x, w)
+            e1.record()
+
+            with torch.cuda.stream(s):
+                e1.wait()
+                b = TestStreamOrderingStress._heavy_matmul_chain(a, w)
+                e2.record(s)
+
+            e2.wait()
+            c = b + 1
+
+            s.synchronize()
+            return c
+
+        x = torch.randn(N, N, device="cuda")
+        w = torch.eye(N, device="cuda") * 0.9
+        self._check_compiled_matches_eager(fn, x, w)
+
+    # ------------------------------------------------------------------
+    # 7. Race: triton kernel on user stream.
+    #    Without the triton stream fix the kernel launches on the default
+    #    stream and reads stale/in-progress data from the user stream.
+    # ------------------------------------------------------------------
+    def test_race_triton_on_user_stream(self):
+        N = self.N
+
+        def fn(x, w):
+            s = torch.cuda.Stream()
+            e = torch.cuda.Event()
+
+            with torch.cuda.stream(s):
+                # Heavy matmul chain produces data on user stream
+                a = TestStreamOrderingStress._heavy_matmul_chain(x, w)
+                # Triton pointwise on the same user stream — without fix
+                # this launches on the default stream
+                b = torch.relu(a)
+                e.record(s)
+
+            e.wait()
+            s.synchronize()
+            return b
+
+        x = torch.randn(N, N, device="cuda")
+        w = torch.eye(N, device="cuda") * 0.9
+        self._check_compiled_matches_eager(fn, x, w)
+
+
+@unittest.skipUnless(TEST_CUDA, "requires CUDA")
+class TestGenericStreamCompile(InductorTestCase):
+    """Tests for torch.compile with device-agnostic torch.Stream API."""
+
+    def test_generic_stream_basic(self):
+        """Test compilation with torch.Stream (device-agnostic API)."""
+        from torch._inductor.utils import run_and_get_code
+
+        # Create stream outside compiled function
+        stream = torch.Stream("cuda")
+        # Convert to cuda stream for use with torch.cuda.stream()
+        cuda_stream = torch.cuda.Stream(
+            stream_id=stream.stream_id,
+            device_index=stream.device_index,
+            device_type=stream.device_type,
+        )
+
+        def fn(x):
+            with torch.cuda.stream(cuda_stream):
+                a = x * 2
+                b = a + 1
+
+            cuda_stream.synchronize()
+            return b
+
+        x = torch.randn(1024, device="cuda")
+
+        expected = fn(x)
+        compiled_fn = torch.compile(fn)
+        result, (code,) = run_and_get_code(compiled_fn, x)
+
+        self.assertEqual(result, expected)
+
+        # Verify stream context is present
+        self.assertIn("torch.cuda.stream", code)
+
+    def test_generic_stream_with_event(self):
+        """Test compilation with torch.Stream and torch.Event."""
+        from torch._inductor.utils import run_and_get_code
+
+        # Create stream and event outside compiled function
+        stream = torch.Stream("cuda")
+        event = torch.Event("cuda")
+        cuda_stream = torch.cuda.Stream(
+            stream_id=stream.stream_id,
+            device_index=stream.device_index,
+            device_type=stream.device_type,
+        )
+
+        def fn(x):
+            # Work on default stream
+            a = x * 2
+            event.record()
+
+            with torch.cuda.stream(cuda_stream):
+                event.wait()
+                b = a + 1
+
+            cuda_stream.synchronize()
+            return b
+
+        x = torch.randn(1024, device="cuda")
+
+        expected = fn(x)
+        compiled_fn = torch.compile(fn)
+        result, (code,) = run_and_get_code(compiled_fn, x)
+
+        self.assertEqual(result, expected)
+
+        # Verify event operations survive compilation as custom ops
+        self.assertIn("record_event", code)
+
+    def test_generic_stream_multiple(self):
+        """Test compilation with multiple torch.Stream instances."""
+        from torch._inductor.utils import run_and_get_code
+
+        # Create streams and event outside compiled function
+        stream1 = torch.Stream("cuda")
+        stream2 = torch.Stream("cuda")
+        event = torch.Event("cuda")
+        cuda_stream1 = torch.cuda.Stream(
+            stream_id=stream1.stream_id,
+            device_index=stream1.device_index,
+            device_type=stream1.device_type,
+        )
+        cuda_stream2 = torch.cuda.Stream(
+            stream_id=stream2.stream_id,
+            device_index=stream2.device_index,
+            device_type=stream2.device_type,
+        )
+
+        def fn(x):
+            with torch.cuda.stream(cuda_stream1):
+                a = x * 2
+                event.record()
+
+            with torch.cuda.stream(cuda_stream2):
+                event.wait()
+                b = a + 1
+
+            cuda_stream1.synchronize()
+            cuda_stream2.synchronize()
+            return b
+
+        x = torch.randn(1024, device="cuda")
+
+        expected = fn(x)
+        compiled_fn = torch.compile(fn)
+        result, (code,) = run_and_get_code(compiled_fn, x)
+
+        self.assertEqual(result, expected)
+
+        # Verify stream handling and event ops survive
+        self.assertIn("torch.cuda.stream", code)
+        self.assertIn("record_event", code)
+        self.assertIn("wait_event", code)
+
+    def test_generic_event_record_on_stream(self):
+        """Test torch.Event.record() with explicit stream argument."""
+        from torch._inductor.utils import run_and_get_code
+
+        # Create stream and event outside compiled function
+        stream = torch.Stream("cuda")
+        event = torch.Event("cuda")
+        cuda_stream = torch.cuda.Stream(
+            stream_id=stream.stream_id,
+            device_index=stream.device_index,
+            device_type=stream.device_type,
+        )
+
+        def fn(x):
+            with torch.cuda.stream(cuda_stream):
+                a = x * 2
+                # Record event with explicit stream (using cuda_stream)
+                event.record(cuda_stream)
+
+            event.wait()
+            b = a + 1
+
+            cuda_stream.synchronize()
+            return b
+
+        x = torch.randn(1024, device="cuda")
+
+        expected = fn(x)
+        compiled_fn = torch.compile(fn)
+        result, (code,) = run_and_get_code(compiled_fn, x)
+
+        self.assertEqual(result, expected)
+
+        # Verify event operations survive compilation as custom ops
+        self.assertIn("record_event", code)
+
+
+@unittest.skipUnless(TEST_CUDA, "requires CUDA")
+class TestStreamIdentity(InductorTestCase):
+    """Verify that compiled code uses the user's original stream objects."""
+
+    def test_single_stream_identity(self):
+        """Codegen should retrieve the user's stream via get_external_object_by_index."""
+        from torch._inductor.utils import run_and_get_code
+
+        user_stream = torch.cuda.Stream()
+
+        def fn(x):
+            with torch.cuda.stream(user_stream):
+                return x * 2
+
+        x = torch.randn(1024, device="cuda")
+        result, (code,) = run_and_get_code(torch.compile(fn), x)
+
+        self.assertEqual(result, fn(x))
+        self.assertIn("get_external_object_by_index", code)
+        self.assertNotIn("torch.cuda.Stream(device=", code)
+
+    def test_multiple_stream_identity(self):
+        """Each stream context should retrieve a different user stream object."""
+        from torch._inductor.utils import run_and_get_code
+
+        stream_a = torch.cuda.Stream()
+        stream_b = torch.cuda.Stream()
+
+        def fn(x):
+            event = torch.cuda.Event()
+            with torch.cuda.stream(stream_a):
+                a = x * 2
+                event.record()
+            with torch.cuda.stream(stream_b):
+                event.wait()
+                b = a + 1
+            stream_b.synchronize()
+            return b
+
+        x = torch.randn(1024, device="cuda")
+        result, (code,) = run_and_get_code(torch.compile(fn), x)
+
+        self.assertEqual(result, fn(x))
+        # Should have two distinct get_external_object_by_index calls
+        matches = re.findall(r"get_external_object_by_index\((\d+)\)", code)
+        self.assertEqual(len(matches), 2)
+        self.assertNotEqual(matches[0], matches[1])
+        self.assertNotIn("torch.cuda.Stream(device=", code)
+
+
+@unittest.skipUnless(TEST_CUDA, "requires CUDA")
+class TestPDLWithMultiStream(InductorTestCase):
+    """Tests that PDL (Programmatic Dependent Launch) composes safely with
+    user-annotated multi-stream code under torch.compile.
+
+    PDL's GDC intrinsics are stream-local: gdc_wait/gdc_launch_dependents
+    only govern the overlap between consecutive kernels on the *same* CUDA
+    stream.  Cross-stream ordering is handled entirely by CUDA events at the
+    wrapper level.  These tests verify that enabling PDL in the presence of
+    multi-stream code doesn't break correctness, doesn't interfere with
+    stream-level invariants (no cross-stream fusion, event ops preserved),
+    and still applies within each stream's own kernel sequence.
+    """
+
+    @unittest.skipIf(not SM90OrLater or TEST_WITH_ROCM, "PDL requires NVIDIA sm90+")
+    @inductor_config.patch({"triton.enable_pdl": True})
+    def test_pdl_single_side_stream(self):
+        """PDL metadata is emitted for a kernel on a side stream."""
+        from torch._inductor.utils import run_and_get_code, run_and_get_triton_code
+
+        def fn(x):
+            s = torch.cuda.Stream()
+            with torch.cuda.stream(s):
+                a = x * 2
+                b = a + 1
+            s.synchronize()
+            return b
+
+        x = torch.randn(1024, device="cuda")
+        expected = fn(x)
+
+        compiled_fn = torch.compile(fn)
+        result, (wrapper_code,) = run_and_get_code(compiled_fn, x)
+        self.assertEqual(result, expected)
+
+        self.assertIn("torch.cuda.stream", wrapper_code)
+        self.assertIn("synchronize_stream", wrapper_code)
+
+        triton_code = run_and_get_triton_code(torch.compile(fn), x)
+        (
+            FileCheck()
+            .check("'launch_pdl': True")
+            .check("gdc_wait")
+            .check("gdc_launch")
+        ).run(triton_code)
+
+    @unittest.skipIf(not SM90OrLater or TEST_WITH_ROCM, "PDL requires NVIDIA sm90+")
+    @inductor_config.patch({"triton.enable_pdl": True})
+    def test_pdl_correctness_with_multiple_streams(self):
+        """Enabling PDL with independent side streams produces correct results."""
+        from torch._inductor.utils import run_and_get_triton_code
+
+        def fn(x):
+            s1 = torch.cuda.Stream()
+            s2 = torch.cuda.Stream()
+            e1 = torch.cuda.Event()
+            e2 = torch.cuda.Event()
+
+            with torch.cuda.stream(s1):
+                a = x * 2 + 1
+                e1.record(s1)
+
+            with torch.cuda.stream(s2):
+                b = x * 3 + 2
+                e2.record(s2)
+
+            e1.wait()
+            e2.wait()
+            return a + b
+
+        x = torch.randn(1024, device="cuda")
+        expected = fn(x)
+        compiled_fn = torch.compile(fn)
+        self.assertEqual(compiled_fn(x), expected)
+
+        triton_code = run_and_get_triton_code(torch.compile(fn), x)
+        # s1 kernel, s2 kernel, and default-stream add kernel
+        (
+            FileCheck()
+            # s1 kernel
+            .check("'launch_pdl': True")
+            .check("gdc_wait")
+            .check("gdc_launch")
+            # s2 kernel
+            .check("'launch_pdl': True")
+            .check("gdc_wait")
+            .check("gdc_launch")
+            # default stream add
+            .check("'launch_pdl': True")
+            .check("gdc_wait")
+            .check("gdc_launch")
+        ).run(triton_code)
+
+    @unittest.skipIf(not SM90OrLater or TEST_WITH_ROCM, "PDL requires NVIDIA sm90+")
+    @inductor_config.patch({"triton.enable_pdl": True})
+    def test_pdl_cross_stream_events_preserved(self):
+        """Event record/wait for cross-stream sync must survive with PDL on.
+
+        PDL is stream-local so it cannot replace event-based cross-stream
+        ordering.  Verify the events are still in the generated code."""
+        from torch._inductor.utils import run_and_get_code, run_and_get_triton_code
+
+        def fn(x):
+            s = torch.cuda.Stream()
+            event = torch.cuda.Event()
+
+            a = x * 2 + 1
+            event.record()
+
+            with torch.cuda.stream(s):
+                event.wait()
+                b = a + 3
+
+            s.synchronize()
+            return b
+
+        x = torch.randn(1024, device="cuda")
+        expected = fn(x)
+
+        compiled_fn = torch.compile(fn)
+        result, (code,) = run_and_get_code(compiled_fn, x)
+        self.assertEqual(result, expected)
+
+        self.assertIn("record_event", code)
+        self.assertIn("wait_event", code)
+
+        # Both kernels (default + side stream) get PDL intrinsics
+        triton_code = run_and_get_triton_code(torch.compile(fn), x)
+        (
+            FileCheck()
+            # default stream kernel
+            .check("'launch_pdl': True")
+            .check("gdc_wait")
+            .check("gdc_launch")
+            # side stream kernel
+            .check("'launch_pdl': True")
+            .check("gdc_wait")
+            .check("gdc_launch")
+        ).run(triton_code)
+
+    @unittest.skipIf(not SM90OrLater or TEST_WITH_ROCM, "PDL requires NVIDIA sm90+")
+    @inductor_config.patch({"triton.enable_pdl": True})
+    def test_pdl_same_stream_consecutive_kernels(self):
+        """Two consecutive kernels on the same side stream should both get PDL.
+
+        This is the case where PDL is actually useful: the second kernel can
+        overlap with the first via GDC intrinsics because they share a stream."""
+        from torch._inductor.utils import run_and_get_triton_code
+
+        def fn(x, y):
+            s = torch.cuda.Stream()
+            with torch.cuda.stream(s):
+                # Two separate fused groups on the same stream
+                a = x**2 + x
+                y.copy_(a)
+            s.synchronize()
+            return y
+
+        x = torch.randn(1024, device="cuda")
+        y = torch.empty(1024, device="cuda")
+        expected = fn(x, y.clone())
+        compiled_fn = torch.compile(fn)
+        self.assertEqual(compiled_fn(x, y.clone()), expected)
+
+        triton_code = run_and_get_triton_code(torch.compile(fn), x, y.clone())
+        (
+            FileCheck()
+            .check("'launch_pdl': True")
+            .check("gdc_wait")
+            .check("gdc_launch")
+        ).run(triton_code)
+
+    @unittest.skipIf(not SM90OrLater or TEST_WITH_ROCM, "PDL requires NVIDIA sm90+")
+    @inductor_config.patch({"triton.enable_pdl": True})
+    def test_pdl_no_fusion_across_streams(self):
+        """PDL must not cause cross-stream ops to be fused."""
+        from torch._inductor.utils import run_and_get_triton_code
+
+        def fn(x):
+            s1 = torch.cuda.Stream()
+            s2 = torch.cuda.Stream()
+            e1 = torch.cuda.Event()
+            e2 = torch.cuda.Event()
+
+            with torch.cuda.stream(s1):
+                a = x * 2
+                b = a + 1
+                e1.record(s1)
+
+            with torch.cuda.stream(s2):
+                c = x * 3
+                d = c + 2
+                e2.record(s2)
+
+            e1.wait()
+            e2.wait()
+            return b + d
+
+        x = torch.randn(1024, device="cuda")
+        expected = fn(x)
+
+        compiled_fn = torch.compile(fn)
+        torch._inductor.metrics.reset()
+        result = compiled_fn(x)
+        self.assertEqual(result, expected)
+
+        # 3 kernels: s1 pointwise, s2 pointwise, default stream add
+        self.assertEqual(torch._inductor.metrics.generated_kernel_count, 3)
+
+        # All 3 kernels should have PDL with GDC intrinsics
+        triton_code = run_and_get_triton_code(torch.compile(fn), x)
+        (
+            FileCheck()
+            # s1 kernel
+            .check("'launch_pdl': True")
+            .check("gdc_wait")
+            .check("gdc_launch")
+            # s2 kernel
+            .check("'launch_pdl': True")
+            .check("gdc_wait")
+            .check("gdc_launch")
+            # default stream add
+            .check("'launch_pdl': True")
+            .check("gdc_wait")
+            .check("gdc_launch")
+        ).run(triton_code)
+
+    @unittest.skipIf(not SM90OrLater or TEST_WITH_ROCM, "PDL requires NVIDIA sm90+")
+    @inductor_config.patch({"triton.enable_pdl": True})
+    def test_pdl_stress_multistream_correctness(self):
+        """Stress test: heavy work across streams with PDL must produce
+        correct results over many iterations to surface any races.
+
+        Uses 4096x4096 matmuls (matching TestStreamOrderingStress) so the
+        GPU work is long enough that a missing event.wait() would cause
+        the consumer to read stale data."""
+        from torch._inductor.utils import run_and_get_code
+
+        N = 4096
+        ITERS = 20
+
+        def fn(x, w):
+            s = torch.cuda.Stream()
+            e = torch.cuda.Event()
+
+            h = x
+            for _ in range(4):
+                h = h @ w
+            e.record()
+
+            with torch.cuda.stream(s):
+                e.wait()
+                out = torch.relu(h) + 1.0
+
+            s.synchronize()
+            return out
+
+        x = torch.randn(N, N, device="cuda")
+        w = torch.eye(N, device="cuda") * 0.9
+
+        # Verify codegen once before the stress loop
+        compiled_fn = torch.compile(fn)
+        result, (code,) = run_and_get_code(compiled_fn, x, w)
+        # Wrapper must have stream context and event sync
+        self.assertIn("torch.cuda.stream", code)
+        self.assertIn("wait_event", code)
+        # The relu+add pointwise kernel should have PDL
+        (
+            FileCheck()
+            .check("'launch_pdl': True")
+            .check("gdc_wait")
+            .check("gdc_launch")
+        ).run(code)
+
+        for _ in range(ITERS):
+            expected = fn(x, w)
+            actual = compiled_fn(x, w)
+            torch.cuda.synchronize()
+            self.assertEqual(actual, expected)
+
+    @unittest.skipIf(not SM90OrLater or TEST_WITH_ROCM, "PDL requires NVIDIA sm90+")
+    @inductor_config.patch({"triton.enable_pdl": True})
+    def test_pdl_mutation_across_streams(self):
+        """Buffer mutation on one stream, read on another, with PDL enabled.
+
+        The mutation is on a locally-created buffer (not an input) to avoid
+        the dynamo guard that forbids event.record() after input mutation."""
+        from torch._inductor.utils import run_and_get_code, run_and_get_triton_code
+
+        def fn(x):
+            s = torch.cuda.Stream()
+            event = torch.cuda.Event()
+
+            # Produce a new buffer (not input mutation) then record
+            a = x * 2
+            event.record()
+
+            with torch.cuda.stream(s):
+                event.wait()
+                # In-place add on side stream
+                a.add_(1)
+
+            s.synchronize()
+            return a
+
+        x = torch.randn(1024, device="cuda")
+
+        expected = fn(x)
+        compiled_fn = torch.compile(fn)
+        result, (code,) = run_and_get_code(compiled_fn, x)
+        self.assertEqual(result, expected)
+
+        self.assertIn("record_event", code)
+        self.assertIn("wait_event", code)
+
+        # Both kernels (default + side stream) get PDL intrinsics
+        triton_code = run_and_get_triton_code(torch.compile(fn), x)
+        (
+            FileCheck()
+            # default stream kernel
+            .check("'launch_pdl': True")
+            .check("gdc_wait")
+            .check("gdc_launch")
+            # side stream kernel
+            .check("'launch_pdl': True")
+            .check("gdc_wait")
+            .check("gdc_launch")
+        ).run(triton_code)
+
+
+@unittest.skipIf(not TEST_CUDA, "requires CUDA")
+@torch._inductor.config.patch({"triton.cudagraphs": True})
+class TestStreamCudagraphInteraction(InductorTestCase):
+    """Tests for user streams under cudagraph capture (reduce-overhead mode)."""
+
+    def test_implicit_current_stream_with_cudagraphs(self):
+        """Event record/wait with implicit current stream must work under cudagraph capture.
+
+        The implicit current stream resolves at runtime via torch.cuda.current_stream(),
+        which correctly returns the cudagraph capture stream during recording.
+        """
+        s1 = torch.cuda.Stream()
+        ev = torch.cuda.Event()
+        ev2 = torch.cuda.Event()
+
+        def fn(x, y):
+            ev.record()
+            with torch.cuda.stream(s1):
+                ev.wait()
+                z = x * 2
+                ev2.record()
+            ev2.wait()
+            return z + y
+
+        x = torch.randn(100, 100, device="cuda")
+        y = torch.randn(100, 100, device="cuda")
+
+        expected = fn(x, y)
+        compiled_fn = torch.compile(fn)
+        # Warmup + capture + replay
+        for _ in range(3):
+            result = compiled_fn(x, y)
+        self.assertEqual(result, expected)
+
+    def test_explicit_current_stream_with_cudagraphs(self):
+        """Passing torch.cuda.current_stream() explicitly must also work under capture.
+
+        The user writes ev.record(torch.cuda.current_stream()) which is
+        semantically identical to ev.record() — both should resolve to the
+        capture stream during cudagraph recording, not the stale default stream.
+        """
+        s1 = torch.cuda.Stream()
+        ev = torch.cuda.Event()
+        ev2 = torch.cuda.Event()
+
+        def fn(x, y):
+            cur = torch.cuda.current_stream()
+            ev.record(cur)
+            with torch.cuda.stream(s1):
+                ev.wait()
+                z = x * 2
+                ev2.record()
+            ev2.wait(cur)
+            return z + y
+
+        x = torch.randn(100, 100, device="cuda")
+        y = torch.randn(100, 100, device="cuda")
+
+        expected = fn(x, y)
+        compiled_fn = torch.compile(fn)
+        for _ in range(3):
+            result = compiled_fn(x, y)
+        self.assertEqual(result, expected)
+
+
+instantiate_parametrized_tests(TestStreamUtils)
+instantiate_parametrized_tests(TestWrapperCodegenStreams)
+instantiate_parametrized_tests(TestStreamCodegen)
+instantiate_parametrized_tests(TestUserStreamCompile)
+instantiate_parametrized_tests(TestStreamOrderingStress)
+instantiate_parametrized_tests(TestGenericStreamCompile)
+instantiate_parametrized_tests(TestStreamIdentity)
+instantiate_parametrized_tests(TestPDLWithMultiStream)
+instantiate_parametrized_tests(TestStreamCudagraphInteraction)
+
+
+if __name__ == "__main__":
+    from torch.testing._internal.common_utils import run_tests
+
+    run_tests()
diff --git a/test/inductor/test_utils.py b/test/inductor/test_utils.py
index 41bb05c4cf594..24f52580b7313 100644
--- a/test/inductor/test_utils.py
+++ b/test/inductor/test_utils.py
@@ -1,5 +1,6 @@
 # Owner(s): ["module: inductor"]
 
+import importlib.util
 import unittest
 
 from sympy import I, Max, Min, Symbol, sympify
@@ -87,6 +88,22 @@ def testSympySubsIdentityNonComparable(self):
         result = sympy_subs(expr, {q0: I})
         self.assertTrue(result.has(I))
 
+    def testIdentityComparisonNoRecursion(self):
+        self.assertTrue(Identity(sympify("0")) >= 0)
+        self.assertFalse(Identity(sympify("-6")) >= 0)
+        self.assertTrue(0 >= Identity(sympify("-6")))
+
+    def testIdentityComparableNumbersInMinMax(self):
+        expr = Identity(sympify("-6"))
+        self.assertTrue(expr.is_number)
+        self.assertTrue(expr.is_comparable)
+        self.assertEqual(Max(0, expr), 0)
+
+    def testIdentityRationalComparisonNoRecursion(self):
+        expr = Identity(sympify("1/7"))
+        self.assertTrue(expr >= 0)
+        self.assertTrue(Max(0, expr).has(expr))
+
     def test_sympy_str(self):
         self.assertEqual(sympy_str(sympify("a+b+c")), "a + b + c")
         self.assertEqual(sympy_str(sympify("a*b+c")), "c + a * b")
@@ -226,5 +243,47 @@ def test_get_device_tflops(self, dtype):
 
 instantiate_device_type_tests(TestUtils, globals(), allow_xpu=True)
 
+
+class TestFP4Support(TestCase):
+    """Tests for FP4 (float4_e2m1fn_x2) infrastructure support."""
+
+    @unittest.skipIf(
+        not torch.cuda.is_available()
+        or importlib.util.find_spec("cutlass_api") is None,
+        "requires CUDA and cutlass_api",
+    )
+    def test_ensure_fp4_dtype_registered(self):
+        """_ensure_fp4_dtype_registered should patch cutlass_api for FP4."""
+        from torch._inductor.utils import _ensure_fp4_dtype_registered
+
+        _ensure_fp4_dtype_registered()
+        import cutlass
+        import cutlass_api.utils
+
+        result = cutlass_api.utils.cutlass_type_from_torch_type(torch.float4_e2m1fn_x2)
+        self.assertEqual(result, cutlass.Float4E2M1FN)
+
+        result_fp32 = cutlass_api.utils.cutlass_type_from_torch_type(torch.float32)
+        self.assertEqual(result_fp32, cutlass.Float32)
+
+    def test_rand_strided_fp4(self):
+        """rand_strided should produce valid FP4 tensors."""
+        from torch._dynamo.testing import rand_strided
+
+        t = rand_strided((4, 8), (8, 1), dtype=torch.float4_e2m1fn_x2, device="cpu")
+        self.assertEqual(t.dtype, torch.float4_e2m1fn_x2)
+        self.assertEqual(t.shape, (4, 8))
+        self.assertEqual(t.stride(), (8, 1))
+
+    @unittest.skipIf(not torch.cuda.is_available(), "requires CUDA")
+    def test_rand_strided_fp4_cuda(self):
+        from torch._dynamo.testing import rand_strided
+
+        t = rand_strided((16, 32), (32, 1), dtype=torch.float4_e2m1fn_x2, device="cuda")
+        self.assertEqual(t.dtype, torch.float4_e2m1fn_x2)
+        self.assertEqual(t.shape, (16, 32))
+        self.assertTrue(t.is_cuda)
+
+
 if __name__ == "__main__":
     run_tests()
diff --git a/test/jit/test_autodiff_subgraph_slicing.py b/test/jit/test_autodiff_subgraph_slicing.py
index f128f9c7eec36..fe8a5e0e7523a 100644
--- a/test/jit/test_autodiff_subgraph_slicing.py
+++ b/test/jit/test_autodiff_subgraph_slicing.py
@@ -27,7 +27,8 @@
 )
 
 
-assert GRAPH_EXECUTOR is not None
+if GRAPH_EXECUTOR is None:
+    raise AssertionError("GRAPH_EXECUTOR is not set")
 
 
 @unittest.skipIf(
@@ -223,9 +224,21 @@ def func(
             output_ref = func(input0, input1)
             for _ in range(2):
                 output = jit_f(input0, input1)
-                assert output_ref[0].requires_grad == output[0].requires_grad
-                assert output_ref[1][0].requires_grad == output[1][0].requires_grad
-                assert output_ref[1][1].requires_grad == output[1][1].requires_grad
+                if output_ref[0].requires_grad != output[0].requires_grad:
+                    raise AssertionError(
+                        f"requires_grad mismatch for output[0]: "
+                        f"{output_ref[0].requires_grad} vs {output[0].requires_grad}"
+                    )
+                if output_ref[1][0].requires_grad != output[1][0].requires_grad:
+                    raise AssertionError(
+                        f"requires_grad mismatch for output[1][0]: "
+                        f"{output_ref[1][0].requires_grad} vs {output[1][0].requires_grad}"
+                    )
+                if output_ref[1][1].requires_grad != output[1][1].requires_grad:
+                    raise AssertionError(
+                        f"requires_grad mismatch for output[1][1]: "
+                        f"{output_ref[1][1].requires_grad} vs {output[1][1].requires_grad}"
+                    )
 
     @unittest.skip(
         "disable until we property handle tensor lists with undefined gradients"
diff --git a/test/jit/test_await.py b/test/jit/test_await.py
index 564b61d03c621..0f538fd9b909a 100644
--- a/test/jit/test_await.py
+++ b/test/jit/test_await.py
@@ -269,7 +269,7 @@ def delayed(x: Tensor) -> Tensor:
         def main(x: Tensor) -> Tensor:
             aw = torch.jit._awaitable(delayed, x)
             if torch.jit.is_scripting():
-                assert isinstance(aw, torch.jit._Await)
+                assert isinstance(aw, torch.jit._Await)  # noqa: S101
             return torch.jit._awaitable_wait(aw)
 
         inp = torch.eye(2)
diff --git a/test/jit/test_device_analysis.py b/test/jit/test_device_analysis.py
index 67dce398b58f1..2b5f1a6ea7d83 100644
--- a/test/jit/test_device_analysis.py
+++ b/test/jit/test_device_analysis.py
@@ -29,7 +29,8 @@ def setUpClass(cls):
     @staticmethod
     def node_output_device(graph):
         graph_out = list(graph.outputs())
-        assert len(graph_out) == 1
+        if len(graph_out) != 1:
+            raise AssertionError(f"Expected 1 graph output, got {len(graph_out)}")
         return graph_out[0].type().device()
 
     def prop_device_on_graph(self, graph, example_devices, in_shapes=None):
diff --git a/test/jit/test_dtype_analysis.py b/test/jit/test_dtype_analysis.py
index f2acd3fe5087f..0b2079e9998c2 100644
--- a/test/jit/test_dtype_analysis.py
+++ b/test/jit/test_dtype_analysis.py
@@ -109,7 +109,8 @@ def node_output_dtypes(graph):
     @staticmethod
     def node_output_dtype_single(graph):
         dtypes = TestDtypeBase.node_output_dtypes(graph)
-        assert len(dtypes) == 1
+        if len(dtypes) != 1:
+            raise AssertionError(f"Expected 1 output dtype, got {len(dtypes)}")
         return dtypes[0]
 
     def prop_dtype_on_graph(self, graph, example_inputs):
diff --git a/test/jit/test_exception.py b/test/jit/test_exception.py
index f46b6bc1f3275..894f2d23392dc 100644
--- a/test/jit/test_exception.py
+++ b/test/jit/test_exception.py
@@ -141,7 +141,7 @@ def foo(cond):
 
         @torch.jit.script
         def foo(cond):
-            assert bool(cond), "hi"
+            assert bool(cond), "hi"  # noqa: S101
 
         foo(torch.tensor(1))
         # we don't currently validate the name of the exception
diff --git a/test/jit/test_freezing.py b/test/jit/test_freezing.py
index d7fbbdab0096e..1b9fce7934d29 100644
--- a/test/jit/test_freezing.py
+++ b/test/jit/test_freezing.py
@@ -11,15 +11,17 @@
 import torch.nn.functional as F
 from torch.jit._recursive import wrap_cpp_module
 from torch.testing import FileCheck
-from torch.testing._internal.common_cuda import TEST_CUDA, TEST_CUDNN
+from torch.testing._internal.common_cuda import TEST_CUDA, TEST_CUDNN, tf32_on_and_off
 from torch.testing._internal.common_quantization import skipIfNoFBGEMM
 from torch.testing._internal.common_quantized import override_quantized_engine
 from torch.testing._internal.common_utils import (
+    IS_ARM64,
     raise_on_run_directly,
     set_default_dtype,
     skipCUDAMemoryLeakCheckIf,
     skipIfTorchDynamo,
     TEST_WITH_ROCM,
+    xfailIf,
 )
 from torch.testing._internal.jit_utils import JitTestCase
 from torch.utils import mkldnn as mkldnn_utils
@@ -1433,6 +1435,8 @@ def forward(self, x):
         self.assertTrue(fm.sub._has_method("method_a"))
         self.assertFalse(fm.sub._has_method("method_b"))
 
+    @xfailIf(IS_ARM64)
+    # see https://github.com/pytorch/pytorch/issues/177258
     @skipIfNoFBGEMM
     def test_module_with_shared_type_instances(self):
         class Child(nn.Module):
@@ -2964,6 +2968,7 @@ def test_conv_to_mkldnn_no_mkldnn(self):
             inp = torch.rand([4, 3, 4, 4])
             self.assertEqual(frozen(inp), mod(inp))
 
+    @tf32_on_and_off(0.005)
     @unittest.skipIf(not (TEST_CUDNN or TEST_WITH_ROCM), "requires CUDNN")
     def test_freeze_conv_relu_fusion(self):
         with set_default_dtype(torch.float):
diff --git a/test/jit/test_hooks.py b/test/jit/test_hooks.py
index 157dd588aa78c..b952ffc30c0ee 100644
--- a/test/jit/test_hooks.py
+++ b/test/jit/test_hooks.py
@@ -112,8 +112,8 @@ def test_hook_method_name_collision(self):
         m = ModuleForwardSingleInput("outer_mod_name", "inner_mod_name")
 
         def foo(self, input: Tuple[str]) -> Tuple[str]:
-            assert self.name == "inner_mod_name"
-            assert input[0] == "a_outermod"
+            assert self.name == "inner_mod_name"  # noqa: S101
+            assert input[0] == "a_outermod"  # noqa: S101
             return ("pre_hook_override_name",)
 
         m.submodule.register_forward_pre_hook(foo)
@@ -174,8 +174,8 @@ def pre_hook(self, input: Tuple[str]) -> Tuple[str]:
             return ("pre_hook_override_name",)
 
         def forward_hook(self, input: Tuple[str], output: str):
-            assert self.name == "outer_mod_name"
-            assert input == ("pre_hook_override_name",)
+            assert self.name == "outer_mod_name"  # noqa: S101
+            assert input == ("pre_hook_override_name",)  # noqa: S101
             output = output + "_fh"
             return output
 
@@ -197,7 +197,7 @@ def pre_hook(self, input: Tuple[str]) -> Tuple[str]:
             return ("pre_hook_override_name",)
 
         def forward_hook(self, input: Tuple[str], output: str):
-            assert input == ("pre_hook_override_name",)
+            assert input == ("pre_hook_override_name",)  # noqa: S101
             return output + "_fh"
 
         m_submod_forward_call.submodule.register_forward_pre_hook(pre_hook)
@@ -223,8 +223,8 @@ def test_hook_compilation_hint(self):
         m = ModuleForwardSingleInput("outer_mod_name", "inner_mod_name")
 
         def pre_hook(self, input: Tuple[str]) -> Tuple[str]:
-            assert self.name == "outer_mod_name"
-            assert input[4] == "a"  # out of bounds tuple range
+            assert self.name == "outer_mod_name"  # noqa: S101
+            assert input[4] == "a"  # noqa: S101 out of bounds tuple range
             return ("pre_hook_override_name",)
 
         m.register_forward_pre_hook(pre_hook)
diff --git a/test/jit/test_isinstance.py b/test/jit/test_isinstance.py
index 0781ad9c74765..ba07910aad704 100644
--- a/test/jit/test_isinstance.py
+++ b/test/jit/test_isinstance.py
@@ -19,88 +19,88 @@
 class TestIsinstance(JitTestCase):
     def test_int(self):
         def int_test(x: Any):
-            assert torch.jit.isinstance(x, int)
-            assert not torch.jit.isinstance(x, float)
+            assert torch.jit.isinstance(x, int)  # noqa: S101
+            assert not torch.jit.isinstance(x, float)  # noqa: S101
 
         x = 1
         self.checkScript(int_test, (x,))
 
     def test_float(self):
         def float_test(x: Any):
-            assert torch.jit.isinstance(x, float)
-            assert not torch.jit.isinstance(x, int)
+            assert torch.jit.isinstance(x, float)  # noqa: S101
+            assert not torch.jit.isinstance(x, int)  # noqa: S101
 
         x = 1.0
         self.checkScript(float_test, (x,))
 
     def test_bool(self):
         def bool_test(x: Any):
-            assert torch.jit.isinstance(x, bool)
-            assert not torch.jit.isinstance(x, float)
+            assert torch.jit.isinstance(x, bool)  # noqa: S101
+            assert not torch.jit.isinstance(x, float)  # noqa: S101
 
         x = False
         self.checkScript(bool_test, (x,))
 
     def test_list(self):
         def list_str_test(x: Any):
-            assert torch.jit.isinstance(x, List[str])
-            assert not torch.jit.isinstance(x, List[int])
-            assert not torch.jit.isinstance(x, Tuple[int])
+            assert torch.jit.isinstance(x, List[str])  # noqa: S101
+            assert not torch.jit.isinstance(x, List[int])  # noqa: S101
+            assert not torch.jit.isinstance(x, Tuple[int])  # noqa: S101
 
         x = ["1", "2", "3"]
         self.checkScript(list_str_test, (x,))
 
     def test_list_tensor(self):
         def list_tensor_test(x: Any):
-            assert torch.jit.isinstance(x, List[torch.Tensor])
-            assert not torch.jit.isinstance(x, Tuple[int])
+            assert torch.jit.isinstance(x, List[torch.Tensor])  # noqa: S101
+            assert not torch.jit.isinstance(x, Tuple[int])  # noqa: S101
 
         x = [torch.tensor([1]), torch.tensor([2]), torch.tensor([3])]
         self.checkScript(list_tensor_test, (x,))
 
     def test_dict(self):
         def dict_str_int_test(x: Any):
-            assert torch.jit.isinstance(x, Dict[str, int])
-            assert not torch.jit.isinstance(x, Dict[int, str])
-            assert not torch.jit.isinstance(x, Dict[str, str])
+            assert torch.jit.isinstance(x, Dict[str, int])  # noqa: S101
+            assert not torch.jit.isinstance(x, Dict[int, str])  # noqa: S101
+            assert not torch.jit.isinstance(x, Dict[str, str])  # noqa: S101
 
         x = {"a": 1, "b": 2}
         self.checkScript(dict_str_int_test, (x,))
 
     def test_dict_tensor(self):
         def dict_int_tensor_test(x: Any):
-            assert torch.jit.isinstance(x, Dict[int, torch.Tensor])
+            assert torch.jit.isinstance(x, Dict[int, torch.Tensor])  # noqa: S101
 
         x = {2: torch.tensor([2])}
         self.checkScript(dict_int_tensor_test, (x,))
 
     def test_tuple(self):
         def tuple_test(x: Any):
-            assert torch.jit.isinstance(x, Tuple[str, int, str])
-            assert not torch.jit.isinstance(x, Tuple[int, str, str])
-            assert not torch.jit.isinstance(x, Tuple[str])
+            assert torch.jit.isinstance(x, Tuple[str, int, str])  # noqa: S101
+            assert not torch.jit.isinstance(x, Tuple[int, str, str])  # noqa: S101
+            assert not torch.jit.isinstance(x, Tuple[str])  # noqa: S101
 
         x = ("a", 1, "b")
         self.checkScript(tuple_test, (x,))
 
     def test_tuple_tensor(self):
         def tuple_tensor_test(x: Any):
-            assert torch.jit.isinstance(x, Tuple[torch.Tensor, torch.Tensor])
+            assert torch.jit.isinstance(x, Tuple[torch.Tensor, torch.Tensor])  # noqa: S101
 
         x = (torch.tensor([1]), torch.tensor([[2], [3]]))
         self.checkScript(tuple_tensor_test, (x,))
 
     def test_optional(self):
         def optional_test(x: Any):
-            assert torch.jit.isinstance(x, Optional[torch.Tensor])
-            assert not torch.jit.isinstance(x, Optional[str])
+            assert torch.jit.isinstance(x, Optional[torch.Tensor])  # noqa: S101
+            assert not torch.jit.isinstance(x, Optional[str])  # noqa: S101
 
         x = torch.ones(3, 3)
         self.checkScript(optional_test, (x,))
 
     def test_optional_none(self):
         def optional_test_none(x: Any):
-            assert torch.jit.isinstance(x, Optional[torch.Tensor])
+            assert torch.jit.isinstance(x, Optional[torch.Tensor])  # noqa: S101
             # assert torch.jit.isinstance(x, Optional[str])
             # TODO: above line in eager will evaluate to True while in
             #       the TS interpreter will evaluate to False as the
@@ -111,28 +111,28 @@ def optional_test_none(x: Any):
 
     def test_list_nested(self):
         def list_nested(x: Any):
-            assert torch.jit.isinstance(x, List[Dict[str, int]])
-            assert not torch.jit.isinstance(x, List[List[str]])
+            assert torch.jit.isinstance(x, List[Dict[str, int]])  # noqa: S101
+            assert not torch.jit.isinstance(x, List[List[str]])  # noqa: S101
 
         x = [{"a": 1, "b": 2}, {"aa": 11, "bb": 22}]
         self.checkScript(list_nested, (x,))
 
     def test_dict_nested(self):
         def dict_nested(x: Any):
-            assert torch.jit.isinstance(x, Dict[str, Tuple[str, str, str]])
-            assert not torch.jit.isinstance(x, Dict[str, Tuple[int, int, int]])
+            assert torch.jit.isinstance(x, Dict[str, Tuple[str, str, str]])  # noqa: S101
+            assert not torch.jit.isinstance(x, Dict[str, Tuple[int, int, int]])  # noqa: S101
 
         x = {"a": ("aa", "aa", "aa"), "b": ("bb", "bb", "bb")}
         self.checkScript(dict_nested, (x,))
 
     def test_tuple_nested(self):
         def tuple_nested(x: Any):
-            assert torch.jit.isinstance(
+            assert torch.jit.isinstance(  # noqa: S101
                 x, Tuple[Dict[str, Tuple[str, str, str]], List[bool], Optional[str]]
             )
-            assert not torch.jit.isinstance(x, Dict[str, Tuple[int, int, int]])
-            assert not torch.jit.isinstance(x, Tuple[str])
-            assert not torch.jit.isinstance(x, Tuple[List[bool], List[str], List[int]])
+            assert not torch.jit.isinstance(x, Dict[str, Tuple[int, int, int]])  # noqa: S101
+            assert not torch.jit.isinstance(x, Tuple[str])  # noqa: S101
+            assert not torch.jit.isinstance(x, Tuple[List[bool], List[str], List[int]])  # noqa: S101
 
         x = (
             {"a": ("aa", "aa", "aa"), "b": ("bb", "bb", "bb")},
@@ -143,21 +143,21 @@ def tuple_nested(x: Any):
 
     def test_optional_nested(self):
         def optional_nested(x: Any):
-            assert torch.jit.isinstance(x, Optional[List[str]])
+            assert torch.jit.isinstance(x, Optional[List[str]])  # noqa: S101
 
         x = ["a", "b", "c"]
         self.checkScript(optional_nested, (x,))
 
     def test_list_tensor_type_true(self):
         def list_tensor_type_true(x: Any):
-            assert torch.jit.isinstance(x, List[torch.Tensor])
+            assert torch.jit.isinstance(x, List[torch.Tensor])  # noqa: S101
 
         x = [torch.rand(3, 3), torch.rand(4, 3)]
         self.checkScript(list_tensor_type_true, (x,))
 
     def test_tensor_type_false(self):
         def list_tensor_type_false(x: Any):
-            assert not torch.jit.isinstance(x, List[torch.Tensor])
+            assert not torch.jit.isinstance(x, List[torch.Tensor])  # noqa: S101
 
         x = [1, 2, 3]
         self.checkScript(list_tensor_type_false, (x,))
@@ -165,9 +165,9 @@ def list_tensor_type_false(x: Any):
     def test_in_if(self):
         def list_in_if(x: Any):
             if torch.jit.isinstance(x, List[int]):
-                assert True
+                assert True  # noqa: S101
             if torch.jit.isinstance(x, List[str]):
-                assert not True
+                assert not True  # noqa: S101
 
         x = [1, 2, 3]
         self.checkScript(list_in_if, (x,))
@@ -175,9 +175,9 @@ def list_in_if(x: Any):
     def test_if_else(self):
         def list_in_if_else(x: Any):
             if torch.jit.isinstance(x, Tuple[str, str, str]):
-                assert True
+                assert True  # noqa: S101
             else:
-                assert not True
+                assert not True  # noqa: S101
 
         x = ("a", "b", "c")
         self.checkScript(list_in_if_else, (x,))
@@ -187,7 +187,7 @@ def list_in_while_loop(x: Any):
             count = 0
             while torch.jit.isinstance(x, List[Dict[str, int]]) and count <= 0:
                 count = count + 1
-            assert count == 1
+            assert count == 1  # noqa: S101
 
         x = [{"a": 1, "b": 2}, {"aa": 11, "bb": 22}]
         self.checkScript(list_in_while_loop, (x,))
@@ -205,8 +205,8 @@ def type_refinement(obj: Any):
                 str_cat = ""
                 for val in obj.values():
                     str_cat = str_cat + val
-                assert "111222" == str_cat
-            assert hit
+                assert "111222" == str_cat  # noqa: S101
+            assert hit  # noqa: S101
 
         x = [torch.rand(3, 3), torch.rand(4, 3)]
         self.checkScript(type_refinement, (x,))
@@ -215,7 +215,7 @@ def type_refinement(obj: Any):
 
     def test_list_no_contained_type(self):
         def list_no_contained_type(x: Any):
-            assert torch.jit.isinstance(x, List)
+            assert torch.jit.isinstance(x, List)  # noqa: S101
 
         x = ["1", "2", "3"]
 
@@ -237,7 +237,7 @@ def list_no_contained_type(x: Any):
 
     def test_tuple_no_contained_type(self):
         def tuple_no_contained_type(x: Any):
-            assert torch.jit.isinstance(x, Tuple)
+            assert torch.jit.isinstance(x, Tuple)  # noqa: S101
 
         x = ("1", "2", "3")
 
@@ -259,7 +259,7 @@ def tuple_no_contained_type(x: Any):
 
     def test_optional_no_contained_type(self):
         def optional_no_contained_type(x: Any):
-            assert torch.jit.isinstance(x, Optional)
+            assert torch.jit.isinstance(x, Optional)  # noqa: S101
 
         x = ("1", "2", "3")
 
@@ -281,7 +281,7 @@ def optional_no_contained_type(x: Any):
 
     def test_dict_no_contained_type(self):
         def dict_no_contained_type(x: Any):
-            assert torch.jit.isinstance(x, Dict)
+            assert torch.jit.isinstance(x, Dict)  # noqa: S101
 
         x = {"a": "aa"}
 
@@ -303,19 +303,19 @@ def dict_no_contained_type(x: Any):
 
     def test_tuple_rhs(self):
         def fn(x: Any):
-            assert torch.jit.isinstance(x, (int, List[str]))
-            assert not torch.jit.isinstance(x, (List[float], Tuple[int, str]))
-            assert not torch.jit.isinstance(x, (List[float], str))
+            assert torch.jit.isinstance(x, (int, List[str]))  # noqa: S101
+            assert not torch.jit.isinstance(x, (List[float], Tuple[int, str]))  # noqa: S101
+            assert not torch.jit.isinstance(x, (List[float], str))  # noqa: S101
 
         self.checkScript(fn, (2,))
         self.checkScript(fn, (["foo", "bar", "baz"],))
 
     def test_nontuple_container_rhs_throws_in_eager(self):
         def fn1(x: Any):
-            assert torch.jit.isinstance(x, [int, List[str]])
+            assert torch.jit.isinstance(x, [int, List[str]])  # noqa: S101
 
         def fn2(x: Any):
-            assert not torch.jit.isinstance(x, {List[str], Tuple[int, str]})
+            assert not torch.jit.isinstance(x, {List[str], Tuple[int, str]})  # noqa: S101
 
         err_highlight = "must be a type or a tuple of types"
 
diff --git a/test/jit/test_misc.py b/test/jit/test_misc.py
index 4f9eb39ef714e..e271b00f8d5d6 100644
--- a/test/jit/test_misc.py
+++ b/test/jit/test_misc.py
@@ -108,7 +108,7 @@ def fn() -> torch.jit.Future[int]:
     def test_future_isinstance(self):
         @torch.jit.script
         def fn(x: Any) -> torch.jit.Future[int]:
-            assert isinstance(x, jit.Future[int])
+            assert isinstance(x, jit.Future[int])  # noqa: S101
             return x
 
         FileCheck().check("Future[int]").run(fn.graph)
diff --git a/test/jit/test_module_containers.py b/test/jit/test_module_containers.py
index 31254be34d671..67e5840ff1be4 100644
--- a/test/jit/test_module_containers.py
+++ b/test/jit/test_module_containers.py
@@ -279,42 +279,42 @@ def __init__(self) -> None:
                 self.moduledict = CustomModuleDict({"submod": self.submod})
 
             def forward(self, inputs):
-                assert self.modulelist[0] is self.submod, (
+                assert self.modulelist[0] is self.submod, (  # noqa: S101
                     "__getitem__ failing for ModuleList"
                 )
-                assert len(self.modulelist) == 1, "__len__ failing for ModuleList"
+                assert len(self.modulelist) == 1, "__len__ failing for ModuleList"  # noqa: S101
                 for module in self.modulelist:
-                    assert module is self.submod, "__iter__ failing for ModuleList"
+                    assert module is self.submod, "__iter__ failing for ModuleList"  # noqa: S101
 
-                assert self.sequential[0] is self.submod, (
+                assert self.sequential[0] is self.submod, (  # noqa: S101
                     "__getitem__ failing for Sequential"
                 )
-                assert len(self.sequential) == 1, "__len__ failing for Sequential"
+                assert len(self.sequential) == 1, "__len__ failing for Sequential"  # noqa: S101
                 for module in self.sequential:
-                    assert module is self.submod, "__iter__ failing for Sequential"
+                    assert module is self.submod, "__iter__ failing for Sequential"  # noqa: S101
 
-                assert self.moduledict["submod"] is self.submod, (
+                assert self.moduledict["submod"] is self.submod, (  # noqa: S101
                     "__getitem__ failing for ModuleDict"
                 )
-                assert len(self.moduledict) == 1, "__len__ failing for ModuleDict"
+                assert len(self.moduledict) == 1, "__len__ failing for ModuleDict"  # noqa: S101
 
                 # note: unable to index moduledict with a string variable currently
                 i = 0
                 for _ in self.moduledict:
                     i += 1
-                assert i == len(self.moduledict), "iteration failing for ModuleDict"
+                assert i == len(self.moduledict), "iteration failing for ModuleDict"  # noqa: S101
 
-                assert "submod" in self.moduledict, "__contains__ fails for ModuleDict"
+                assert "submod" in self.moduledict, "__contains__ fails for ModuleDict"  # noqa: S101
 
                 for key in self.moduledict:
-                    assert key == "submod", "keys() fails for ModuleDict"
+                    assert key == "submod", "keys() fails for ModuleDict"  # noqa: S101
 
                 for item in self.moduledict.items():
-                    assert item[0] == "submod", "items() fails for ModuleDict"
-                    assert item[1] is self.submod, "items() fails for ModuleDict"
+                    assert item[0] == "submod", "items() fails for ModuleDict"  # noqa: S101
+                    assert item[1] is self.submod, "items() fails for ModuleDict"  # noqa: S101
 
                 for value in self.moduledict.values():
-                    assert value is self.submod, "values() fails for ModuleDict"
+                    assert value is self.submod, "values() fails for ModuleDict"  # noqa: S101
 
                 return inputs
 
@@ -344,7 +344,7 @@ def __init__(self) -> None:
                 self.modulelist = CustomModuleList([self.submod])
 
             def forward(self, inputs):
-                assert len(self.modulelist) == 2, "__len__ failing for ModuleList"
+                assert len(self.modulelist) == 2, "__len__ failing for ModuleList"  # noqa: S101
                 return inputs
 
         m = MyModule()
@@ -362,8 +362,8 @@ def __init__(self) -> None:
                 )
 
             def forward(self, input):
-                assert self.moduledict["relu"] is self.relu
-                assert self.moduledict["tanh"] is self.tanh
+                assert self.moduledict["relu"] is self.relu  # noqa: S101
+                assert self.moduledict["tanh"] is self.tanh  # noqa: S101
                 return input
 
         m = MyModule()
@@ -376,7 +376,7 @@ def __init__(self) -> None:
                 self.moduledict = torch.nn.ModuleDict({"foo": None, "bar": None})
 
             def forward(self, input):
-                assert self.moduledict["blah"] == "blah", "this is a keyerror"
+                assert self.moduledict["blah"] == "blah", "this is a keyerror"  # noqa: S101
 
         with self.assertRaisesRegexWithHighlight(
             RuntimeError, "Key Error, blah", 'self.moduledict["blah"'
@@ -391,7 +391,7 @@ def __init__(self) -> None:
 
             def forward(self, input):
                 idx = "blah"
-                assert self.moduledict[idx] == "blah", "this is a string literal error"
+                assert self.moduledict[idx] == "blah", "this is a string literal error"  # noqa: S101
 
         with self.assertRaisesRegexWithHighlight(
             RuntimeError,
@@ -439,7 +439,7 @@ def __init__(self) -> None:
                 self.moduledict = CustomModuleDict()
 
             def forward(self, inputs):
-                assert "submod" not in self.moduledict, (
+                assert "submod" not in self.moduledict, (  # noqa: S101
                     "__contains__ fails for ModuleDict"
                 )
                 return inputs
diff --git a/test/jit/test_pdt.py b/test/jit/test_pdt.py
index eaff742f55591..ae48a0daa1df4 100644
--- a/test/jit/test_pdt.py
+++ b/test/jit/test_pdt.py
@@ -136,7 +136,7 @@ def test_nn_module_with_export_function(self):
         class TestModelWithExport(torch.nn.Module):
             @torch.jit.export
             def fn(self, x, y) -> Any:
-                assert not (isinstance(x, bool) and isinstance(y, bool))
+                assert not (isinstance(x, bool) and isinstance(y, bool))  # noqa: S101
                 if isinstance(x, int) and isinstance(y, int):
                     return x + y
                 elif isinstance(x, float) and isinstance(y, float):
@@ -669,7 +669,7 @@ def test_dict_int_list(a):
 
     def test_any(self):
         def test_multiple_types(a):
-            assert not isinstance(a, bool)
+            assert not isinstance(a, bool)  # noqa: S101
             return a
 
         def test_multiple_type_refinement(a):
@@ -718,7 +718,7 @@ def test_multiple_type_refinement(a):
     def test_class_as_profiled_types(self):
         class UserDefinedClass:
             def fn(self, b) -> Any:
-                assert b is not None
+                assert b is not None  # noqa: S101
                 if isinstance(b, int):
                     return b if b > 0 else -1
                 elif isinstance(b, float):
@@ -726,7 +726,7 @@ def fn(self, b) -> Any:
                 return 0
 
         def test_model(a, m):
-            assert not isinstance(a, bool)
+            assert not isinstance(a, bool)  # noqa: S101
             return m.fn(a)
 
         make_global(UserDefinedClass, test_model)
@@ -772,7 +772,7 @@ def fn(self, b):
                     return -1
 
         def test_model_with_args(a, m):
-            assert not isinstance(a, bool)
+            assert not isinstance(a, bool)  # noqa: S101
             return m.fn(a)
 
         make_global(ClassWithArgs, test_model_with_args)
diff --git a/test/jit/test_peephole.py b/test/jit/test_peephole.py
index 61c443fc1b659..3cdc09ba0f34b 100644
--- a/test/jit/test_peephole.py
+++ b/test/jit/test_peephole.py
@@ -341,7 +341,7 @@ def foo(x: List[int], b: List[int]):
                 x1 = True
             else:
                 x1 = len(b) != 4
-            assert x1 == False  # noqa: E712 TODO: canonicalize x is False to aten::eq
+            assert x1 == False  # noqa: E712, S101 TODO: canonicalize x is False to aten::eq
             return len(x), len(b)
 
         torch._C._jit_pass_peephole_list_idioms(foo.graph, refine_list_len=True)
@@ -355,7 +355,7 @@ def foo(x: List[int], b: List[int]):
                 x1 = False
             else:
                 x1 = len(b) != 4
-            assert x1 == False  # noqa: E712 TODO: canonicalize x is False to aten::eq
+            assert x1 == False  # noqa: E712, S101 TODO: canonicalize x is False to aten::eq
             return len(x), len(b)
 
         torch._C._jit_pass_peephole_list_idioms(foo.graph, refine_list_len=True)
@@ -369,7 +369,7 @@ def foo(x: List[int], b: List[int]):
                 x1 = True
             else:
                 x1 = len(b) == 4
-            assert x1 == False  # noqa: E712 TODO: canonicalize x is False to aten::eq
+            assert x1 == False  # noqa: E712, S101 TODO: canonicalize x is False to aten::eq
             return len(x), len(b)
 
         torch._C._jit_pass_peephole_list_idioms(foo.graph, refine_list_len=True)
@@ -383,7 +383,7 @@ def foo(x: List[int], b: List[int]):
                 x1 = True
             else:
                 x1 = len(b) != 4
-            assert x1 == False  # noqa: E712 TODO: canonicalize x is False to aten::eq
+            assert x1 == False  # noqa: E712, S101 TODO: canonicalize x is False to aten::eq
             return len(x), len(b)
 
         torch._C._jit_pass_peephole_list_idioms(foo.graph, refine_list_len=True)
@@ -398,7 +398,7 @@ def foo(x: List[int], b: List[int]):
                 x1 = len(b) != 4
             else:
                 x1 = True
-            assert x1 == False  # noqa: E712 TODO: canonicalize x is False to aten::eq
+            assert x1 == False  # noqa: E712, S101 TODO: canonicalize x is False to aten::eq
             return len(x), len(b)
 
         torch._C._jit_pass_peephole_list_idioms(foo.graph, refine_list_len=True)
@@ -413,7 +413,7 @@ def foo(x: List[int], b: List[int]):
                 x1 = len(b) != 4
             else:
                 x1 = True
-            assert not x1
+            assert not x1  # noqa: S101
             return len(x), len(b)
 
         torch._C._jit_pass_peephole_list_idioms(foo.graph, refine_list_len=True)
@@ -425,7 +425,7 @@ def foo(x: List[int], b: List[int]):
 
         @torch.jit.script
         def foo(x: List[int]):
-            assert len(x) == 4
+            assert len(x) == 4  # noqa: S101
             x.append(3)
             return len(x)
 
@@ -435,7 +435,7 @@ def foo(x: List[int]):
 
         @torch.jit.script
         def foo(x: List[int], y: List[int]):
-            assert len(x) == 4 or len(y) == 5
+            assert len(x) == 4 or len(y) == 5  # noqa: S101
             return len(x) + len(y)
 
         torch._C._jit_pass_peephole_list_idioms(foo.graph, refine_list_len=True)
@@ -515,7 +515,7 @@ def foo(x: int, cond: bool):
 
         @torch.jit.script
         def foo(x: int, y: int):
-            assert x == 4 or y == 5
+            assert x == 4 or y == 5  # noqa: S101
             return x + y
 
         torch._C._jit_pass_peephole_list_idioms(foo.graph, refine_list_len=True)
diff --git a/test/jit/test_python_bindings.py b/test/jit/test_python_bindings.py
index 2af4552dcebaa..dc6300f69198d 100644
--- a/test/jit/test_python_bindings.py
+++ b/test/jit/test_python_bindings.py
@@ -81,7 +81,8 @@ def test_graph_create(self):
     def test_add_input(self):
         gr = torch._C.Graph()
         foo_value = gr.addInput("foo")
-        assert foo_value in gr.inputs()
+        if foo_value not in gr.inputs():
+            raise AssertionError("foo_value not found in graph inputs")
 
     def test_canonicalize(self):
         ir = """
diff --git a/test/jit/test_python_ir.py b/test/jit/test_python_ir.py
index 593d98193b1ce..4b6d46fa6eed6 100644
--- a/test/jit/test_python_ir.py
+++ b/test/jit/test_python_ir.py
@@ -85,7 +85,10 @@ def foo(x):
         for mul in mul_constant_int:
             with g.insert_point_guard(mul):
                 outputs = g.insertGraph(unrolled_mul.graph, list(mul.inputs()))
-                assert len(outputs) == len(list(mul.outputs()))
+                if len(outputs) != len(list(mul.outputs())):
+                    raise AssertionError(
+                        f"Expected {len(list(mul.outputs()))} outputs, got {len(outputs)}"
+                    )
                 for new_out, old_out in zip(outputs, g.outputs()):
                     old_out.replaceAllUsesWith(new_out)
                 mul.destroy()
diff --git a/test/jit/test_recursive_script.py b/test/jit/test_recursive_script.py
index d6addfddca1a7..4399c26049922 100644
--- a/test/jit/test_recursive_script.py
+++ b/test/jit/test_recursive_script.py
@@ -36,7 +36,7 @@ def __init__(self) -> None:
                 self.x = None
 
             def forward(self):
-                assert self.x is None
+                assert self.x is None  # noqa: S101
 
         m = torch.jit.script(M())
         self.checkModule(M(), ())
diff --git a/test/jit/test_scriptmod_ann.py b/test/jit/test_scriptmod_ann.py
index 4541c24dc5e0c..754b3f4a4d474 100644
--- a/test/jit/test_scriptmod_ann.py
+++ b/test/jit/test_scriptmod_ann.py
@@ -32,7 +32,8 @@ def forward(self, x: int):
 
         with warnings.catch_warnings(record=True) as w:
             self.checkModule(M(), (1,))
-        assert len(w) == 0
+        if len(w) != 0:
+            raise AssertionError(f"Expected no warnings, got {len(w)}: {w}")
 
     def test_annotated_nonempty_container(self):
         class M(torch.nn.Module):
@@ -46,7 +47,8 @@ def forward(self, x: List[int]):
 
         with warnings.catch_warnings(record=True) as w:
             self.checkModule(M(), ([1, 2, 3],))
-        assert len(w) == 0
+        if len(w) != 0:
+            raise AssertionError(f"Expected no warnings, got {len(w)}: {w}")
 
     def test_annotated_empty_tensor(self):
         class M(torch.nn.Module):
@@ -60,7 +62,8 @@ def forward(self, x: torch.Tensor):
 
         with warnings.catch_warnings(record=True) as w:
             self.checkModule(M(), (torch.rand(2, 3),))
-        assert len(w) == 0
+        if len(w) != 0:
+            raise AssertionError(f"Expected no warnings, got {len(w)}: {w}")
 
     def test_annotated_with_jit_attribute(self):
         class M(torch.nn.Module):
@@ -74,7 +77,8 @@ def forward(self, x: List[int]):
 
         with warnings.catch_warnings(record=True) as w:
             self.checkModule(M(), ([1, 2, 3],))
-        assert len(w) == 0
+        if len(w) != 0:
+            raise AssertionError(f"Expected no warnings, got {len(w)}: {w}")
 
     def test_annotated_class_level_annotation_only(self):
         class M(torch.nn.Module):
@@ -90,7 +94,8 @@ def forward(self, y: List[int]):
 
         with warnings.catch_warnings(record=True) as w:
             self.checkModule(M(), ([1, 2, 3],))
-        assert len(w) == 0
+        if len(w) != 0:
+            raise AssertionError(f"Expected no warnings, got {len(w)}: {w}")
 
     def test_annotated_class_level_annotation_and_init_annotation(self):
         class M(torch.nn.Module):
@@ -106,7 +111,8 @@ def forward(self, y: List[int]):
 
         with warnings.catch_warnings(record=True) as w:
             self.checkModule(M(), ([1, 2, 3],))
-        assert len(w) == 0
+        if len(w) != 0:
+            raise AssertionError(f"Expected no warnings, got {len(w)}: {w}")
 
     def test_annotated_class_level_jit_annotation(self):
         class M(torch.nn.Module):
@@ -122,7 +128,8 @@ def forward(self, y: List[int]):
 
         with warnings.catch_warnings(record=True) as w:
             self.checkModule(M(), ([1, 2, 3],))
-        assert len(w) == 0
+        if len(w) != 0:
+            raise AssertionError(f"Expected no warnings, got {len(w)}: {w}")
 
     def test_annotated_empty_list(self):
         class M(torch.nn.Module):
diff --git a/test/jit/test_tensor_creation_ops.py b/test/jit/test_tensor_creation_ops.py
index 3933984179b5f..23379f3be6756 100644
--- a/test/jit/test_tensor_creation_ops.py
+++ b/test/jit/test_tensor_creation_ops.py
@@ -23,7 +23,7 @@ def randperm(x: int):
             perm = torch.randperm(x)
             # Have to perform assertion here because TorchScript returns dtypes
             # as integers, which are not comparable against eager torch.dtype.
-            assert perm.dtype == torch.int64
+            assert perm.dtype == torch.int64  # noqa: S101
 
         self.checkScript(randperm, (3,))
 
@@ -32,7 +32,7 @@ def randperm(x: int):
             perm = torch.randperm(x, dtype=torch.float)
             # Have to perform assertion here because TorchScript returns dtypes
             # as integers, which are not comparable against eager torch.dtype.
-            assert perm.dtype == torch.float
+            assert perm.dtype == torch.float  # noqa: S101
 
         self.checkScript(randperm, (3,))
 
@@ -41,7 +41,7 @@ def triu_indices(rows: int, cols: int):
             indices = torch.triu_indices(rows, cols)
             # Have to perform assertion here because TorchScript returns dtypes
             # as integers, which are not comparable against eager torch.dtype.
-            assert indices.dtype == torch.int64
+            assert indices.dtype == torch.int64  # noqa: S101
 
         self.checkScript(triu_indices, (3, 3))
 
@@ -50,7 +50,7 @@ def triu_indices(rows: int, cols: int):
             indices = torch.triu_indices(rows, cols, dtype=torch.int32)
             # Have to perform assertion here because TorchScript returns dtypes
             # as integers, which are not comparable against eager torch.dtype.
-            assert indices.dtype == torch.int32
+            assert indices.dtype == torch.int32  # noqa: S101
 
         self.checkScript(triu_indices, (3, 3))
 
@@ -59,7 +59,7 @@ def tril_indices(rows: int, cols: int):
             indices = torch.tril_indices(rows, cols)
             # Have to perform assertion here because TorchScript returns dtypes
             # as integers, which are not comparable against eager torch.dtype.
-            assert indices.dtype == torch.int64
+            assert indices.dtype == torch.int64  # noqa: S101
 
         self.checkScript(tril_indices, (3, 3))
 
@@ -68,7 +68,7 @@ def tril_indices(rows: int, cols: int):
             indices = torch.tril_indices(rows, cols, dtype=torch.int32)
             # Have to perform assertion here because TorchScript returns dtypes
             # as integers, which are not comparable against eager torch.dtype.
-            assert indices.dtype == torch.int32
+            assert indices.dtype == torch.int32  # noqa: S101
 
         self.checkScript(tril_indices, (3, 3))
 
diff --git a/test/jit/test_torchbind.py b/test/jit/test_torchbind.py
index aab330972b01e..3664ece0d38b4 100644
--- a/test/jit/test_torchbind.py
+++ b/test/jit/test_torchbind.py
@@ -266,9 +266,11 @@ def forward(self):
         inst = FooBar1234()
         scripted = torch.jit.script(inst)
         eic = self.getExportImportCopy(scripted)
-        assert eic() == "deserialized"
+        if eic() != "deserialized":
+            raise AssertionError(f"Expected 'deserialized', got {eic()!r}")
         for expected in ["deserialized", "was", "i"]:
-            assert eic.f.pop() == expected
+            if eic.f.pop() != expected:
+                raise AssertionError(f"Expected {expected!r}, got unexpected value")
 
     def test_torchbind_getstate(self):
         class FooBar4321(torch.nn.Module):
@@ -287,9 +289,11 @@ def forward(self):
         # values at instantiation in the test with some transformation, but
         # because it seems we serialize/deserialize multiple times, that
         # transformation isn't as you would it expect it to be.
-        assert eic() == 7
+        if eic() != 7:
+            raise AssertionError(f"Expected 7, got {eic()!r}")
         for expected in [7, 3, 3, 1]:
-            assert eic.f.pop() == expected
+            if eic.f.pop() != expected:
+                raise AssertionError(f"Expected {expected!r}, got unexpected value")
 
     def test_torchbind_deepcopy(self):
         class FooBar4321(torch.nn.Module):
@@ -303,9 +307,11 @@ def forward(self):
         inst = FooBar4321()
         scripted = torch.jit.script(inst)
         copied = copy.deepcopy(scripted)
-        assert copied.forward() == 7
+        if copied.forward() != 7:
+            raise AssertionError(f"Expected 7, got {copied.forward()!r}")
         for expected in [7, 3, 3, 1]:
-            assert copied.f.pop() == expected
+            if copied.f.pop() != expected:
+                raise AssertionError(f"Expected {expected!r}, got unexpected value")
 
     def test_torchbind_python_deepcopy(self):
         class FooBar4321(torch.nn.Module):
@@ -318,9 +324,11 @@ def forward(self):
 
         inst = FooBar4321()
         copied = copy.deepcopy(inst)
-        assert copied() == 7
+        if copied() != 7:
+            raise AssertionError(f"Expected 7, got {copied()!r}")
         for expected in [7, 3, 3, 1]:
-            assert copied.f.pop() == expected
+            if copied.f.pop() != expected:
+                raise AssertionError(f"Expected {expected!r}, got unexpected value")
 
     def test_torchbind_tracing(self):
         class TryTracing(torch.nn.Module):
diff --git a/test/jit/test_tracer.py b/test/jit/test_tracer.py
index a39542b5b21b9..fd7799dcb5f12 100644
--- a/test/jit/test_tracer.py
+++ b/test/jit/test_tracer.py
@@ -153,7 +153,8 @@ def f2(x):
             return (x,)
 
         jit_f2 = torch.jit.trace(f2, x)
-        assert f2(x) == jit_f2(x)  # fails
+        if f2(x) != jit_f2(x):
+            raise AssertionError(f"Traced output {jit_f2(x)} != eager output {f2(x)}")
 
     def test_trace_out_operator_with_two_output(self):
         example_input = torch.rand(2, 8)
@@ -1152,10 +1153,14 @@ def foo(x, w):
     def test_trace_detach_redispatch(self):
         def foo(x, w):
             y = torch.matmul(x, w)
-            assert y.requires_grad
+            if not y.requires_grad:
+                raise AssertionError("Expected y.requires_grad to be True")
             y = y.detach()
             # Make sure trace kernel redispatches to the right lower kernel.
-            assert not y.requires_grad
+            if y.requires_grad:
+                raise AssertionError(
+                    "Expected y.requires_grad to be False after detach"
+                )
             return y
 
         x, w = torch.rand(3, 4), torch.rand(4, 5, requires_grad=True)
@@ -1180,10 +1185,14 @@ def foo(x, w):
     def test_trace_detach_inplace_redispatch(self):
         def foo(x, w):
             y = torch.matmul(x, w)
-            assert y.requires_grad
+            if not y.requires_grad:
+                raise AssertionError("Expected y.requires_grad to be True")
             y.detach_()
             # Make sure trace kernel redispatches to the right lower kernel.
-            assert not y.requires_grad
+            if y.requires_grad:
+                raise AssertionError(
+                    "Expected y.requires_grad to be False after detach_"
+                )
             return y
 
         x, w = torch.rand(3, 4), torch.rand(4, 5, requires_grad=True)
@@ -1926,9 +1935,12 @@ def foo(bar, baz):
 
         traced = torch.jit.trace(foo, (torch.rand(3, 3), torch.rand(3, 3)))
         graph_str = str(traced.graph)
-        assert "bar" in graph_str
-        assert "baz" in graph_str
-        assert "quick_brown_fox" in graph_str
+        if "bar" not in graph_str:
+            raise AssertionError(f"'bar' not found in graph: {graph_str}")
+        if "baz" not in graph_str:
+            raise AssertionError(f"'baz' not found in graph: {graph_str}")
+        if "quick_brown_fox" not in graph_str:
+            raise AssertionError(f"'quick_brown_fox' not found in graph: {graph_str}")
 
     @skipIfTorchDynamo("Not a suitable test for TorchDynamo")
     def test_tracing_hooks(self):
diff --git a/test/jit_hooks/CMakeLists.txt b/test/jit_hooks/CMakeLists.txt
index ba32390b0b366..606625148f2e5 100644
--- a/test/jit_hooks/CMakeLists.txt
+++ b/test/jit_hooks/CMakeLists.txt
@@ -9,5 +9,5 @@ endif()
 find_package(Torch REQUIRED)
 
 add_executable(test_jit_hooks test_jit_hooks.cpp)
-set_property(TARGET test_jit_hooks PROPERTY CXX_STANDARD 17)
+set_property(TARGET test_jit_hooks PROPERTY CXX_STANDARD 20)
 target_link_libraries(test_jit_hooks "${TORCH_LIBRARIES}")
diff --git a/test/mobile/custom_build/CMakeLists.txt b/test/mobile/custom_build/CMakeLists.txt
index 52e713895ff8a..2629a8848d6b0 100644
--- a/test/mobile/custom_build/CMakeLists.txt
+++ b/test/mobile/custom_build/CMakeLists.txt
@@ -2,7 +2,7 @@ cmake_minimum_required(VERSION 3.15)
 
 project(custom_build_project)
 
-set(CMAKE_CXX_STANDARD 17 CACHE STRING "The C++ standard whose features are requested to build this target.")
+set(CMAKE_CXX_STANDARD 20 CACHE STRING "The C++ standard whose features are requested to build this target.")
 
 # Find torch library
 find_package(Torch REQUIRED)
diff --git a/test/mobile/model_test/android_api_module.py b/test/mobile/model_test/android_api_module.py
index a2beea339ff3d..45feff75c24ce 100644
--- a/test/mobile/model_test/android_api_module.py
+++ b/test/mobile/model_test/android_api_module.py
@@ -1,5 +1,3 @@
-from typing import Optional
-
 import torch
 from torch import Tensor
 
@@ -72,11 +70,11 @@ def tupleIntSumReturnTuple(
         return (input, sum)
 
     @torch.jit.script_method
-    def optionalIntIsNone(self, input: Optional[int]) -> bool:
+    def optionalIntIsNone(self, input: int | None) -> bool:
         return input is None
 
     @torch.jit.script_method
-    def intEq0None(self, input: int) -> Optional[int]:
+    def intEq0None(self, input: int) -> int | None:
         if input == 0:
             return None
         return input
diff --git a/test/nn/attention/_fa_test_common.py b/test/nn/attention/_fa_test_common.py
index c5cfdd8cc2816..43bba88c2af8c 100644
--- a/test/nn/attention/_fa_test_common.py
+++ b/test/nn/attention/_fa_test_common.py
@@ -8,7 +8,7 @@
 
 import itertools
 from contextlib import contextmanager
-from typing import NamedTuple, Optional
+from typing import NamedTuple
 
 import torch
 import torch.nn.functional as F
@@ -32,7 +32,7 @@ class SdpaShape(NamedTuple):
 
 
 @contextmanager
-def cuda_kernel_profiler(kernel_pattern: str, activities: Optional[list] = None):
+def cuda_kernel_profiler(kernel_pattern: str, activities: list | None = None):
     """
     Context manager to profile CUDA kernels and check if a pattern is found.
 
@@ -341,7 +341,7 @@ def _test_compiled_sdpa_backward_matches_math(
         dtype,
         is_causal,
         head_dim: int = 128,
-        min_atol: Optional[float] = None,
+        min_atol: float | None = None,
     ):
         """Test compiled flash attention backward numerical correctness against math backend."""
         shape = SdpaShape(2, 8, 512, head_dim)
diff --git a/test/nn/test_convolution.py b/test/nn/test_convolution.py
index 35f715a88c376..bb393f63fc38f 100644
--- a/test/nn/test_convolution.py
+++ b/test/nn/test_convolution.py
@@ -58,6 +58,8 @@ def _get_cudnn_version():
     GRADCHECK_NONDET_TOL,
     gradgradcheck,
     instantiate_parametrized_tests,
+    IS_ARM64,
+    IS_LINUX,
     MACOS_VERSION,
     MI300_ARCH,
     parametrize as parametrize_test,
@@ -3197,6 +3199,8 @@ def _make_noncontiguous(inp):
             gradgradcheck(convolution, inputs, nondet_tol=gradcheck_nondet_tol)
         )
 
+    @xfailIf(IS_LINUX and IS_ARM64)
+    # see https://github.com/pytorch/pytorch/issues/177245
     @onlyCPU
     def test_conv_contiguous_for_oneDNN(self):
         # See https://github.com/pytorch/pytorch/issues/80837.
@@ -3223,6 +3227,8 @@ def test_conv_contiguous_for_oneDNN(self):
                     y_ = conv(x2)
                     self.assertEqual(y, y_)
 
+    @xfailIf(IS_LINUX and IS_ARM64)
+    # see https://github.com/pytorch/pytorch/issues/177245
     @onlyCPU
     def test_conv_ic1_channels_last_for_oneDNN(self):
         # See https://github.com/pytorch/pytorch/issues/82060, N > 1 will call in OneDNN path.
@@ -3575,7 +3581,7 @@ def test_ConvTranspose3d_size_1_kernel(self, device):
     @dtypes(torch.float)
     @torch.backends.cudnn.flags(enabled=True, deterministic=True, benchmark=False)
     @torch.backends.miopen.flags(immediate=True)
-    @tf32_on_and_off(0.001)
+    @tf32_on_and_off(0.005)
     def test_Conv2d_naive_groups(self, device, dtype):
         # Check that grouped convolutions matches two half convolutions
         m = nn.Conv2d(4, 4, kernel_size=3, groups=2).to(device, dtype)
diff --git a/test/nn/test_dropout.py b/test/nn/test_dropout.py
index 9d26f3c9c16db..5110d87525634 100644
--- a/test/nn/test_dropout.py
+++ b/test/nn/test_dropout.py
@@ -11,7 +11,6 @@
 from torch.testing._internal.common_device_type import (
     dtypes,
     dtypesIfMPS,
-    expectedFailureMPS,
     expectedFailureXLA,
     instantiate_device_type_tests,
 )
@@ -216,7 +215,6 @@ def _test_dropoutNd_channel_zero(self, dropout, input):
     @expectedFailureXLA  # seems like freeze_rng_state is not honoured by XLA
     @dtypes(torch.double)
     @dtypesIfMPS(torch.float32)
-    @expectedFailureMPS
     def test_Dropout1d(self, device, dtype):
         with set_default_dtype(dtype):
             N, C, L = (
@@ -289,7 +287,6 @@ def test_Dropout2d(self, device):
         self._test_dropoutNd_channel_zero(nn.Dropout2d(p=0.5, inplace=True), input)
 
     @expectedFailureXLA  # seems like freeze_rng_state is not honoured by XLA
-    @expectedFailureMPS  # Failing on current pytorch MPS
     def test_Dropout3d(self, device):
         b = random.randint(1, 5)
         w = random.randint(1, 5)
diff --git a/test/nn/test_init.py b/test/nn/test_init.py
index 23cab3bb1fa80..7741cce27a36e 100644
--- a/test/nn/test_init.py
+++ b/test/nn/test_init.py
@@ -9,10 +9,11 @@
 import torch
 import torch.nn.functional as F
 import torch.nn.init as init
+from torch.testing._internal.common_device_type import instantiate_device_type_tests
 from torch.testing._internal.common_utils import (
+    parametrize as parametrize_test,
     run_tests,
     skipIfNoLapack,
-    skipIfTorchDynamo,
     slowTest,
     TEST_SCIPY,
     TestCase,
@@ -23,16 +24,39 @@
     from scipy import stats
 
 
+# (init_fn, shape, extra kwargs) covering every public init function.
+# Shapes satisfy each function's dimensionality requirements.
+# The 0-element variant is derived by replacing dim-0 with 0.
+ALL_INIT_FNS = [
+    (init.uniform_, (3, 5), {}),
+    (init.normal_, (3, 5), {}),
+    (init.trunc_normal_, (3, 5), {}),
+    (init.constant_, (3, 5), {"val": 0.3}),
+    (init.ones_, (3, 5), {}),
+    (init.zeros_, (3, 5), {}),
+    (init.eye_, (3, 5), {}),
+    (init.dirac_, (3, 16, 5), {}),
+    (init.xavier_uniform_, (3, 5), {}),
+    (init.xavier_normal_, (3, 5), {}),
+    (init.kaiming_uniform_, (3, 5), {}),
+    (init.kaiming_normal_, (3, 5), {}),
+    (init.orthogonal_, (3, 5), {}),
+    (init.sparse_, (3, 5), {"sparsity": 0.1}),
+]
+
+
 class TestNNInit(TestCase):
     def setUp(self):
         super().setUp()
         random.seed(123)
 
+    @torch._dynamo.disable
     def _is_normal(self, tensor, mean, std):
         samples = tensor.view(-1).tolist()
         p_value = stats.kstest(samples, "norm", args=(mean, std))[1]
         return p_value > 0.0001
 
+    @torch._dynamo.disable
     def _is_trunc_normal(self, tensor, mean, std, a, b):
         # scipy's trunc norm is suited for data drawn from N(0, 1),
         # so we need to transform our data to test it using scipy.
@@ -43,6 +67,7 @@ def _is_trunc_normal(self, tensor, mean, std, a, b):
         p_value = stats.kstest(z_samples, "truncnorm", args=(a0, b0))[1]
         return p_value > 0.0001
 
+    @torch._dynamo.disable
     def _is_uniform(self, tensor, a, b):
         samples = tensor.view(-1).tolist()
         p_value = stats.kstest(samples, "uniform", args=(a, (b - a)))[1]
@@ -110,7 +135,6 @@ def test_calculate_gain_only_accepts_valid_nonlinearities(self):
                 init.calculate_gain(random_string)
 
     @unittest.skipIf(not TEST_SCIPY, "Scipy not found.")
-    @skipIfTorchDynamo("scipy.kstest is failing under dynamo")
     def test_uniform(self):
         for dims in [1, 2, 4]:
             input_tensor = self._create_random_nd_tensor(dims, size_min=30, size_max=50)
@@ -121,7 +145,6 @@ def test_uniform(self):
                 raise AssertionError("Expected uniform distribution")
 
     @unittest.skipIf(not TEST_SCIPY, "Scipy not found.")
-    @skipIfTorchDynamo("scipy.kstest is failing under dynamo")
     def test_normal(self):
         for dims in [1, 2, 4]:
             input_tensor = self._create_random_nd_tensor(dims, size_min=30, size_max=50)
@@ -133,21 +156,6 @@ def test_normal(self):
                 raise AssertionError("Expected normal distribution")
 
     @unittest.skipIf(not TEST_SCIPY, "Scipy not found.")
-    @skipIfTorchDynamo("scipy.kstest is failing under dynamo")
-    def test_trunc_normal(self):
-        for dims in [1, 2, 4]:
-            input_tensor = self._create_random_nd_tensor(dims, size_min=30, size_max=50)
-            mean = self._random_float(-3, 3)
-            std = self._random_float(0.01, 1)
-            a = self._random_float(mean - 2 * std, mean)
-            b = self._random_float(mean, mean + 2 * std)
-            init.trunc_normal_(input_tensor, mean=mean, std=std, a=a, b=b)
-
-            if not self._is_trunc_normal(input_tensor, mean, std, a, b):
-                raise AssertionError("Expected truncated normal distribution")
-
-    @unittest.skipIf(not TEST_SCIPY, "Scipy not found.")
-    @skipIfTorchDynamo("scipy.kstest is failing under dynamo")
     def test_trunc_normal_generator(self):
         gen = torch.Generator()
         gen.manual_seed(42)
@@ -162,6 +170,20 @@ def test_trunc_normal_generator(self):
         if not self._is_trunc_normal(input_tensor, mean=0, std=1, a=0, b=1):
             raise AssertionError("Expected truncated normal distribution")
 
+    def test_init_methods_meta(self):
+        for fn, shape, kwargs in ALL_INIT_FNS:
+            with self.subTest(init_fn=fn.__name__):
+                t = torch.empty(shape, device="meta")
+                result = fn(t, **kwargs)
+                self.assertIs(result, t)
+
+    def test_init_methods_meta_zero_element(self):
+        for fn, shape, kwargs in ALL_INIT_FNS:
+            with self.subTest(init_fn=fn.__name__):
+                t = torch.empty((0,) + shape[1:], device="meta")
+                result = fn(t, **kwargs)
+                self.assertIs(result, t)
+
     def test_constant(self):
         for dims in [1, 2, 4]:
             input_tensor = self._create_random_nd_tensor(dims, size_min=1, size_max=5)
@@ -362,7 +384,6 @@ def test_xavier_uniform(self):
                     raise AssertionError("Expected uniform distribution")
 
     @unittest.skipIf(not TEST_SCIPY, "Scipy not found.")
-    @skipIfTorchDynamo("scipy.kstest is failing under dynamo")
     def test_xavier_normal(self):
         for use_gain in [True, False]:
             for dims in [2, 4]:
@@ -414,7 +435,6 @@ def test_kaiming_normal_warning_on_0element_tensor(self):
             _ = init.kaiming_normal_(tensor)
 
     @unittest.skipIf(not TEST_SCIPY, "Scipy not found.")
-    @skipIfTorchDynamo("scipy.kstest is failing under dynamo")
     def test_kaiming_uniform(self):
         for use_a in [True, False]:
             for dims in [2, 4]:
@@ -446,7 +466,6 @@ def test_kaiming_uniform(self):
                         raise AssertionError("Expected uniform distribution")
 
     @unittest.skipIf(not TEST_SCIPY, "Scipy not found.")
-    @skipIfTorchDynamo("scipy.kstest is failing under dynamo")
     def test_kaiming_normal(self):
         for use_a in [True, False]:
             for dims in [2, 4]:
@@ -484,7 +503,6 @@ def test_sparse_only_works_on_2d_inputs(self):
                 init.sparse_(tensor, sparsity)
 
     @unittest.skipIf(not TEST_SCIPY, "Scipy not found.")
-    @skipIfTorchDynamo("scipy.kstest is failing under dynamo")
     def test_sparse_default_std(self):
         for use_random_std in [True, False]:
             input_tensor = self._create_random_nd_tensor(2, size_min=30, size_max=35)
@@ -552,5 +570,183 @@ def fn():
             fn()
 
 
+class TestNNInitDeviceType(TestCase):
+    @torch._dynamo.disable
+    def _is_trunc_normal(self, tensor, mean, std, a, b):
+        z_samples = (tensor.view(-1) - mean) / std
+        z_samples = z_samples.tolist()
+        a0 = (a - mean) / std
+        b0 = (b - mean) / std
+        p_value = stats.kstest(z_samples, "truncnorm", args=(a0, b0))[1]
+        return p_value > 0.0001
+
+    @unittest.skipIf(not TEST_SCIPY, "Scipy not found.")
+    @parametrize_test("dims", [1, 2, 4])
+    @parametrize_test(
+        "dtype",
+        [torch.float32, torch.float64],
+    )
+    def test_trunc_normal_all_dtypes(self, device, dims, dtype):
+        size = [random.randint(30, 50) for _ in range(dims)]
+        input_tensor = torch.zeros(size, dtype=dtype, device=device)
+        mean = random.uniform(-3, 3)
+        std = random.uniform(0.01, 1)
+        a = random.uniform(mean - 2 * std, mean)
+        b = random.uniform(mean, mean + 2 * std)
+        init.trunc_normal_(input_tensor, mean=mean, std=std, a=a, b=b)
+
+        # Compare bounds at tensor precision: a and b get rounded when
+        # cast to reduced-precision dtypes.
+        lo = torch.tensor(a, dtype=dtype).item()
+        hi = torch.tensor(b, dtype=dtype).item()
+        self.assertTrue(
+            input_tensor.min().item() >= lo,
+            f"{dtype}: values below lower bound a={a}",
+        )
+        self.assertTrue(
+            input_tensor.max().item() <= hi,
+            f"{dtype}: values above upper bound b={b}",
+        )
+        self.assertTrue(
+            self._is_trunc_normal(input_tensor.float().cpu(), mean, std, a, b),
+            f"{dtype}: failed KS test against truncated normal",
+        )
+
+    # Reduced-precision KS test uses fixed wide params to avoid random
+    # intervals too narrow for the type's representable value count.
+    @unittest.skipIf(not TEST_SCIPY, "Scipy not found.")
+    @parametrize_test(
+        "dtype",
+        [torch.float16, torch.bfloat16],
+    )
+    def test_trunc_normal_low_precision(self, device, dtype):
+        n = 10000
+        mean, std, a, b = 0.0, 1.0, -2.0, 2.0
+        t = torch.empty(n, dtype=dtype, device=device)
+        init.trunc_normal_(t, mean=mean, std=std, a=a, b=b)
+
+        self.assertTrue(
+            t.min().item() >= a,
+            f"{dtype}: values below lower bound a={a}",
+        )
+        self.assertTrue(
+            t.max().item() <= b,
+            f"{dtype}: values above upper bound b={b}",
+        )
+        self.assertTrue(
+            self._is_trunc_normal(t.float().cpu(), mean, std, a, b),
+            f"{dtype}: failed KS test against truncated normal",
+        )
+
+    # Test that trunc_normal_ behaves well for narrow interval compared to std.
+    @unittest.skipIf(not TEST_SCIPY, "Scipy not found.")
+    @parametrize_test(
+        "dtype",
+        [torch.float32, torch.float64, torch.float16, torch.bfloat16],
+    )
+    def test_trunc_normal_narrow_interval(self, device, dtype):
+        n = 10000
+        mean, std, a, b = -3.0, 1.0, -0.5, 0.5
+        t = torch.empty(n, dtype=dtype, device=device)
+        init.trunc_normal_(t, mean=mean, std=std, a=a, b=b)
+
+        self.assertTrue(
+            t.min().item() >= a,
+            f"{dtype}: values below lower bound a={a}",
+        )
+        self.assertTrue(
+            t.max().item() <= b,
+            f"{dtype}: values above upper bound b={b}",
+        )
+        self.assertTrue(
+            self._is_trunc_normal(t.float().cpu(), mean, std, a, b),
+            f"{dtype}: failed KS test against truncated normal",
+        )
+
+    # Sanity check for trunc normal to ensure that we sample a decent
+    # subset of the value space.
+    @parametrize_test(
+        "dtype,min_unique",
+        [
+            (torch.float32, 5000),
+            (torch.float64, 5000),
+            (torch.float16, 3000),
+            (torch.bfloat16, 1000),
+        ],
+    )
+    def test_trunc_normal_unique_values(self, device, dtype, min_unique):
+        n = 10000
+        t = torch.empty(n, dtype=dtype, device=device)
+        init.trunc_normal_(t, mean=0.0, std=1.0, a=-2.0, b=2.0)
+
+        self.assertTrue(
+            t.min().item() >= -2.0,
+            f"{dtype}: values below lower bound",
+        )
+        self.assertTrue(
+            t.max().item() <= 2.0,
+            f"{dtype}: values above upper bound",
+        )
+
+        unique = t.unique().numel()
+        self.assertGreater(
+            unique,
+            min_unique,
+            f"{dtype}: only {unique} unique values, expected > {min_unique}",
+        )
+
+        self.assertFalse(
+            t.isinf().any().item(),
+            f"{dtype}: trunc_normal_ produced inf values",
+        )
+        self.assertFalse(
+            t.isnan().any().item(),
+            f"{dtype}: trunc_normal_ produced nan values",
+        )
+
+    # Sanity check that we don't round to the boundary by mistake.
+    @parametrize_test(
+        "dtype",
+        [torch.float32, torch.float64, torch.float16, torch.bfloat16],
+    )
+    def test_trunc_normal_no_boundary_values_small_std(self, device, dtype):
+        t = torch.empty(10000, dtype=dtype, device=device)
+        init.trunc_normal_(t, mean=0.0, std=0.1, a=-2.0, b=2.0)
+
+        at_lower = (t == -2.0).sum().item()
+        at_upper = (t == 2.0).sum().item()
+        self.assertEqual(
+            at_lower,
+            0,
+            f"{dtype}: {at_lower} values clamped to lower bound a=-2.0",
+        )
+        self.assertEqual(
+            at_upper,
+            0,
+            f"{dtype}: {at_upper} values clamped to upper bound b=2.0",
+        )
+
+    def _run_init_test(self, device, zero_element=False):
+        for fn, shape, kwargs in ALL_INIT_FNS:
+            if zero_element:
+                shape = (0,) + shape[1:]
+            with self.subTest(init_fn=fn.__name__):
+                t = torch.empty(shape, device=device)
+                result = fn(t, **kwargs)
+                self.assertIs(result, t)
+
+    def test_init_methods_no_error(self, device):
+        # Smoke test: every public init function must run without error on all
+        # devices (including meta where the ops are no-ops).
+        self._run_init_test(device)
+
+    def test_init_methods_zero_element(self, device):
+        # Every init function should handle 0-element tensors as a no-op.
+        self._run_init_test(device, zero_element=True)
+
+
+instantiate_device_type_tests(TestNNInitDeviceType, globals())
+
+
 if __name__ == "__main__":
     run_tests()
diff --git a/test/onnx/exporter/test_api.py b/test/onnx/exporter/test_api.py
index 50284c5a4225a..5d9d72ff5a778 100644
--- a/test/onnx/exporter/test_api.py
+++ b/test/onnx/exporter/test_api.py
@@ -66,7 +66,8 @@ def assert_export(
         self, *args, strategy: str | None = "TorchExportNonStrictStrategy", **kwargs
     ):
         onnx_program = torch.onnx.export(*args, **kwargs, dynamo=True, verbose=False)
-        assert onnx_program is not None
+        if onnx_program is None:
+            raise AssertionError("onnx_program is None")
         onnx_testing.assert_onnx_program(onnx_program, strategy=strategy)
         return onnx_program
 
diff --git a/test/onnx/exporter/test_capture_strategies.py b/test/onnx/exporter/test_capture_strategies.py
index fa17dfcb5ac0b..2fd61d6c35736 100644
--- a/test/onnx/exporter/test_capture_strategies.py
+++ b/test/onnx/exporter/test_capture_strategies.py
@@ -32,7 +32,8 @@ def forward(self, a, b):
 
         result = strategy_cls()(model, (a, b), kwargs=None, dynamic_shapes=None)
         ep = result.exported_program
-        assert ep is not None
+        if ep is None:
+            raise AssertionError("ep is None")
         torch.testing.assert_close(ep.module()(a, b), model(a, b))
 
     def test_draft_export_on_data_dependent_model(self):
@@ -53,7 +54,8 @@ def forward(self, a, b):
             expected_warning = "1 issue(s) found during export, and it was not able to soundly produce a graph."
             self.assertIn(expected_warning, str(cm.output))
         ep = result.exported_program
-        assert ep is not None
+        if ep is None:
+            raise AssertionError("ep is None")
         torch.testing.assert_close(ep.module()(a, b), model(a, b))
 
 
diff --git a/test/onnx/exporter/test_hf_models_e2e.py b/test/onnx/exporter/test_hf_models_e2e.py
index 2d1c4604b45d0..3cc7c8b02c78e 100644
--- a/test/onnx/exporter/test_hf_models_e2e.py
+++ b/test/onnx/exporter/test_hf_models_e2e.py
@@ -22,7 +22,8 @@ def export(self, model, args=(), kwargs=None, **options) -> torch.onnx.ONNXProgr
             verbose=False,
             **options,
         )
-        assert onnx_program is not None
+        if onnx_program is None:
+            raise AssertionError("onnx_program is None")
         return onnx_program
 
     def test_onnx_export_huggingface_llm_models_with_kv_cache(self):
diff --git a/test/onnx/exporter/test_input_observer.py b/test/onnx/exporter/test_input_observer.py
index 39386a0775fc2..68e74c40b7cd3 100644
--- a/test/onnx/exporter/test_input_observer.py
+++ b/test/onnx/exporter/test_input_observer.py
@@ -872,6 +872,291 @@ def forward(self, x=None, y=None):
         with self.assertRaises(RuntimeError):
             observer.infer_dynamic_shapes()
 
+    def test_infer_dynamic_shapes_missing(self):
+        class Model(torch.nn.Module):
+            def forward(
+                self,
+                input_ids=None,
+                pixel_values=None,
+                attention_mask=None,
+                position_ids=None,
+                past_key_values=None,
+                token_type_ids=None,
+                cache_position=None,
+            ):
+                return input_ids
+
+        inputs = [
+            dict(
+                input_ids=torch.ones((1, 28), dtype=torch.int64),
+                pixel_values=torch.ones((1, 3, 112, 112), dtype=torch.int64),
+                attention_mask=torch.ones((1, 28), dtype=torch.int64),
+                position_ids=torch.ones((1, 28), dtype=torch.int64),
+                token_type_ids=torch.ones((1, 28), dtype=torch.int64),
+                cache_position=torch.ones((28,), dtype=torch.int64),
+            ),
+            dict(
+                input_ids=torch.ones((1, 1), dtype=torch.int64),
+                attention_mask=torch.ones((1, 29), dtype=torch.int64),
+                position_ids=torch.ones((1, 1), dtype=torch.int64),
+                past_key_values=torch.rand((1, 1, 28, 32)),
+                token_type_ids=torch.ones((1, 1), dtype=torch.int64),
+                cache_position=torch.ones((1,), dtype=torch.int64),
+            ),
+            dict(
+                input_ids=torch.ones((1, 1), dtype=torch.int64),
+                attention_mask=torch.ones((1, 30), dtype=torch.int64),
+                position_ids=torch.ones((1, 1), dtype=torch.int64),
+                past_key_values=torch.rand((1, 1, 29, 32)),
+                token_type_ids=torch.ones((1, 1), dtype=torch.int64),
+                cache_position=torch.ones((1,), dtype=torch.int64),
+            ),
+        ]
+
+        model = Model()
+        observer = InputObserver(
+            value_if_missing=dict(pixel_values=torch.empty((0, 3, 112, 112)))
+        )
+        with observer(model):
+            for kwargs in inputs:
+                model(**kwargs)
+
+        shapes = observer.infer_dynamic_shapes(set_batch_dimension_for=True)
+        cst = torch.export.Dim.DYNAMIC
+        expected = {
+            "input_ids": {0: cst, 1: cst},
+            "pixel_values": {0: cst},
+            "attention_mask": {0: cst, 1: cst},
+            "position_ids": {0: cst, 1: cst},
+            "past_key_values": {0: cst, 2: cst},
+            "token_type_ids": {0: cst, 1: cst},
+            "cache_position": {0: cst},
+        }
+        self.assertEqual(expected, shapes)
+        kwargs = observer.infer_arguments()
+        self.assertEqual(list(expected), list(kwargs))
+        self.assertEqual((0, 3, 112, 112), kwargs["pixel_values"].shape)
+
+    def test_infer_dynamic_shapes_missing_args(self):
+        class Model(torch.nn.Module):
+            def forward(
+                self,
+                input_ids=None,
+                pixel_values=None,
+                attention_mask=None,
+                past_key_values=None,
+            ):
+                return input_ids
+
+        inputs = [
+            (
+                torch.ones((1, 28), dtype=torch.int64),
+                torch.ones((1, 3, 112, 112), dtype=torch.int64),
+                torch.ones((1, 28), dtype=torch.int64),
+            ),
+            (
+                torch.ones((1, 1), dtype=torch.int64),
+                None,
+                torch.ones((1, 29), dtype=torch.int64),
+                torch.rand((1, 1, 28, 32)),
+            ),
+            (
+                torch.ones((1, 1), dtype=torch.int64),
+                None,
+                torch.ones((1, 30), dtype=torch.int64),
+                torch.rand((1, 1, 29, 32)),
+            ),
+        ]
+
+        model = Model()
+        observer = InputObserver(
+            value_if_missing={1: torch.empty((0, 3, 112, 112), dtype=torch.int64)}
+        )
+        with observer(model):
+            for args in inputs:
+                model(*args)
+
+        shapes = observer.infer_dynamic_shapes(set_batch_dimension_for=True)
+        cst = torch.export.Dim.DYNAMIC
+        expected = ({0: cst, 1: cst}, {0: cst}, {0: cst, 1: cst}, {0: cst, 2: cst})
+        self.assertEqual(expected, shapes)
+        args = observer.infer_arguments()
+        self.assertEqual(len(expected), len(args))
+        self.assertEqual((0, 3, 112, 112), args[1].shape)
+
+    def test_infer_dynamic_shapes_missing_kwargs_nested(self):
+        class Model(torch.nn.Module):
+            def forward(
+                self,
+                input_ids=None,
+                pixel_values=None,
+                attention_mask=None,
+                position_ids=None,
+                past_key_values=None,
+                token_type_ids=None,
+                cache_position=None,
+            ):
+                return input_ids
+
+        inputs = [
+            dict(
+                input_ids=torch.ones((1, 28), dtype=torch.int64),
+                pixel_values=(
+                    torch.ones((1, 3, 112, 112), dtype=torch.int64),
+                    torch.ones((1, 3, 112, 112), dtype=torch.int64),
+                ),
+                attention_mask=torch.ones((1, 28), dtype=torch.int64),
+                position_ids=torch.ones((1, 28), dtype=torch.int64),
+                token_type_ids=torch.ones((1, 28), dtype=torch.int64),
+                cache_position=torch.ones((28,), dtype=torch.int64),
+            ),
+            dict(
+                input_ids=torch.ones((1, 1), dtype=torch.int64),
+                attention_mask=torch.ones((1, 29), dtype=torch.int64),
+                position_ids=torch.ones((1, 1), dtype=torch.int64),
+                past_key_values=torch.rand((1, 1, 28, 32)),
+                token_type_ids=torch.ones((1, 1), dtype=torch.int64),
+                cache_position=torch.ones((1,), dtype=torch.int64),
+            ),
+            dict(
+                input_ids=torch.ones((1, 1), dtype=torch.int64),
+                attention_mask=torch.ones((1, 30), dtype=torch.int64),
+                position_ids=torch.ones((1, 1), dtype=torch.int64),
+                past_key_values=torch.rand((1, 1, 29, 32)),
+                token_type_ids=torch.ones((1, 1), dtype=torch.int64),
+                cache_position=torch.ones((1,), dtype=torch.int64),
+            ),
+        ]
+
+        model = Model()
+        observer = InputObserver(
+            value_if_missing=dict(
+                pixel_values=(
+                    torch.empty((0, 3, 112, 112), dtype=torch.int64),
+                    torch.empty((0, 3, 112, 112), dtype=torch.int64),
+                )
+            )
+        )
+        with observer(model):
+            for kwargs in inputs:
+                model(**kwargs)
+
+        shapes = observer.infer_dynamic_shapes(set_batch_dimension_for=True)
+        cst = torch.export.Dim.DYNAMIC
+        expected = {
+            "input_ids": {0: cst, 1: cst},
+            "pixel_values": ({0: cst}, {0: cst}),
+            "attention_mask": {0: cst, 1: cst},
+            "position_ids": {0: cst, 1: cst},
+            "past_key_values": {0: cst, 2: cst},
+            "token_type_ids": {0: cst, 1: cst},
+            "cache_position": {0: cst},
+        }
+        self.assertEqual(expected, shapes)
+        kwargs = observer.infer_arguments()
+        self.assertEqual(list(expected), list(kwargs))
+        self.assertIsInstance(kwargs["pixel_values"], tuple)
+        self.assertEqual(2, len(kwargs["pixel_values"]))
+        self.assertEqual((0, 3, 112, 112), kwargs["pixel_values"][0].shape)
+        self.assertEqual((0, 3, 112, 112), kwargs["pixel_values"][1].shape)
+
+    def test_io_captured_kwargs_kwargs(self):
+        class Model(torch.nn.Module):
+            def forward(self, x, **kwargs):
+                return x + kwargs["y"]
+
+        inputs = [
+            dict(x=torch.randn((5, 6)), y=torch.randn((1, 6))),
+            dict(x=torch.randn((7, 7)), y=torch.randn((1, 7))),
+            dict(x=torch.randn((7, 8)), y=torch.randn((1, 8))),
+            dict(x=torch.randn((7, 9)), y=torch.randn((1, 9))),
+        ]
+
+        model = Model()
+        expected = [model(**kwargs) for kwargs in inputs]
+        observer = InputObserver()
+        with observer(model):
+            for kwargs in inputs:
+                model(**kwargs)
+        self.assertEqual(len(observer.info), 3)
+        for i in range(3):
+            self.assertEqual(len(observer.info.flat_outputs[i]), 1)
+            torch.testing.assert_close(expected[i], observer.info.flat_outputs[i][0])
+
+        cst = torch.export.Dim.DYNAMIC
+        ds = observer.infer_dynamic_shapes()
+        self.assertEqual(dict(x={0: cst, 1: cst}, kwargs=dict(y={1: cst})), ds)
+        args = observer.infer_arguments()
+        self.assertIsInstance(args, dict)
+        self.assertEqual(2, len(args))
+        self.assertEqual(["x", "y"], list(args))
+
+        dynamic_shapes = torch.export.AdditionalInputs()
+        for kwargs in inputs:
+            dynamic_shapes.add((), kwargs)
+        dss = dynamic_shapes.dynamic_shapes(model, (), inputs[0])
+        self.assertEqual({"x": (cst, cst), "kwargs": {"y": (None, cst)}}, dss)
+
+    def test_io_captured_kwargs_kwargs_with_args(self):
+        class Model(torch.nn.Module):
+            def forward(self, a, *args, **kwargs):
+                return a - args[0] * args[1] + kwargs["x"] - kwargs["y"]
+
+        inputs = [
+            (
+                (torch.randn((5, 6)), torch.randn((5, 6)), torch.randn((5, 6))),
+                dict(x=torch.randn((5, 6)), y=torch.randn((1, 6))),
+            ),
+            (
+                (torch.randn((7, 7)), torch.randn((7, 7)), torch.randn((7, 7))),
+                dict(x=torch.randn((7, 7)), y=torch.randn((1, 7))),
+            ),
+        ]
+
+        model = Model()
+        expected = [model(*args, **kwargs) for args, kwargs in inputs]
+        observer = InputObserver()
+        with observer(model):
+            for args, kwargs in inputs:
+                model(*args, **kwargs)
+        self.assertEqual(len(observer.info), 2)
+        for i in range(2):
+            self.assertEqual(len(observer.info.flat_outputs[i]), 1)
+            torch.testing.assert_close(expected[i], observer.info.flat_outputs[i][0])
+
+        cst = torch.export.Dim.DYNAMIC
+        ds = observer.infer_dynamic_shapes()
+        self.assertEqual(
+            {
+                "a": {0: cst, 1: cst},
+                "args": ({0: cst, 1: cst}, {0: cst, 1: cst}),
+                "kwargs": {"x": {0: cst, 1: cst}, "y": {1: cst}},
+            },
+            ds,
+        )
+
+        dynamic_shapes = torch.export.AdditionalInputs()
+        for args, kwargs in inputs:
+            dynamic_shapes.add(args, kwargs)
+        dss = dynamic_shapes.dynamic_shapes(model, *inputs[0])
+        self.assertEqual(
+            {
+                "a": (cst, cst),
+                "args": ((cst, cst), (cst, cst)),
+                "kwargs": {"x": (cst, cst), "y": (None, cst)},
+            },
+            dss,
+        )
+
+        with self.assertRaises(RuntimeError):
+            observer.infer_arguments()
+
+        args, kwargs = observer.infer_arguments(as_args_kwargs=True)
+        self.assertIsInstance(kwargs, dict)
+        self.assertEqual(["x", "y"], list(kwargs))
+        self.assertIsInstance(args, tuple)
+        self.assertEqual(len(args), 3)
+
 
 if __name__ == "__main__":
     common_utils.run_tests()
diff --git a/test/onnx/exporter/test_verification.py b/test/onnx/exporter/test_verification.py
index 662f249fab482..f296ce90adc5c 100644
--- a/test/onnx/exporter/test_verification.py
+++ b/test/onnx/exporter/test_verification.py
@@ -92,7 +92,8 @@ def forward(self, a, b):
         model = Model()
         args = (torch.tensor([1.0]), torch.tensor([2.0]))
         onnx_program = torch.onnx.export(model, args, dynamo=True, verbose=False)
-        assert onnx_program is not None
+        if onnx_program is None:
+            raise AssertionError("onnx_program is None")
         interpreter = _verification._VerificationInterpreter(onnx_program)
         results = interpreter.run(args)
         torch.testing.assert_close(results, model(*args))
@@ -113,7 +114,8 @@ def forward(self, a, b):
         model = Model()
         args = (torch.tensor([1.0]), torch.tensor([2.0]))
         onnx_program = torch.onnx.export(model, args, dynamo=True, verbose=False)
-        assert onnx_program is not None
+        if onnx_program is None:
+            raise AssertionError("onnx_program is None")
         verification_infos = _verification.verify_onnx_program(
             onnx_program, args, compare_intermediates=False
         )
@@ -128,7 +130,8 @@ def forward(self, a, b):
         model = Model()
         args = (torch.tensor([1.0]), torch.tensor([2.0]))
         onnx_program = torch.onnx.export(model, args, dynamo=True, verbose=False)
-        assert onnx_program is not None
+        if onnx_program is None:
+            raise AssertionError("onnx_program is None")
         verification_infos = _verification.verify_onnx_program(
             onnx_program, args, compare_intermediates=True
         )
diff --git a/test/onnx/internal/test_registration.py b/test/onnx/internal/test_registration.py
index 8d90553ac2181..0cad11b0096a8 100644
--- a/test/onnx/internal/test_registration.py
+++ b/test/onnx/internal/test_registration.py
@@ -180,7 +180,8 @@ def test(g, x):
 
         self.assertTrue(registration.registry.is_registered_op("test::test_op", 9))
         function_group = registration.registry.get_function_group("test::test_op")
-        assert function_group is not None
+        if function_group is None:
+            raise AssertionError("function_group is None")
         self.assertEqual(function_group.get(9), test)
 
     def test_onnx_symbolic_registers_function_applied_decorator_when_provided(self):
@@ -199,7 +200,8 @@ def test():
             return
 
         function_group = registration.registry.get_function_group("test::test_op")
-        assert function_group is not None
+        if function_group is None:
+            raise AssertionError("function_group is None")
         registered_function = function_group[9]
         self.assertFalse(wrapper_called)
         registered_function()
@@ -230,7 +232,8 @@ def test(g, x):
 
         self.assertTrue(registration.registry.is_registered_op("test::test_op", 9))
         function_group = registration.registry.get_function_group("test::test_op")
-        assert function_group is not None
+        if function_group is None:
+            raise AssertionError("function_group is None")
         self.assertEqual(function_group.get(9), test)
 
     def test_custom_onnx_symbolic_overrides_existing_function(self):
@@ -247,7 +250,8 @@ def test_custom():
             return "custom"
 
         function_group = registration.registry.get_function_group("test::test_op")
-        assert function_group is not None
+        if function_group is None:
+            raise AssertionError("function_group is None")
         self.assertEqual(function_group.get(9), test_custom)
 
 
diff --git a/test/onnx/model_defs/word_language_model.py b/test/onnx/model_defs/word_language_model.py
index be6c8fef006b0..22b0a1b2a2bd6 100644
--- a/test/onnx/model_defs/word_language_model.py
+++ b/test/onnx/model_defs/word_language_model.py
@@ -1,7 +1,6 @@
 # The model is from here:
 #   https://github.com/pytorch/examples/blob/master/word_language_model/model.py
 
-from typing import Optional
 
 import torch
 import torch.nn as nn
@@ -122,7 +121,7 @@ def repackage_hidden(h: tuple[Tensor, Tensor]):
         """Detach hidden states from their history."""
         return (h[0].detach(), h[1].detach())
 
-    def forward(self, input: Tensor, hidden: Optional[tuple[Tensor, Tensor]] = None):
+    def forward(self, input: Tensor, hidden: tuple[Tensor, Tensor] | None = None):
         emb = self.drop(self.encoder(input))
         output, hidden = self.rnn(emb, hidden)
         output = self.drop(output)
diff --git a/test/onnx/onnx_test_common.py b/test/onnx/onnx_test_common.py
index bdeeb40cb9f64..4a6b9f75d4de0 100644
--- a/test/onnx/onnx_test_common.py
+++ b/test/onnx/onnx_test_common.py
@@ -8,7 +8,7 @@
 import os
 import unittest
 from collections.abc import Callable, Collection, Iterable, Mapping, Sequence
-from typing import Any, Optional, Union
+from typing import Any
 
 import numpy as np
 import onnxruntime
@@ -24,11 +24,11 @@
 from torch.types import Number
 
 
-_NumericType = Union[Number, torch.Tensor, np.ndarray]
-_ModelType = Union[torch.nn.Module, Callable, torch_export.ExportedProgram]
-_InputArgsType = Optional[
-    Union[torch.Tensor, int, float, bool, Sequence[Any], Mapping[str, Any]]
-]
+_NumericType = Number | torch.Tensor | np.ndarray
+_ModelType = torch.nn.Module | Callable | torch_export.ExportedProgram
+_InputArgsType = (
+    torch.Tensor | int | float | bool | Sequence[Any] | Mapping[str, Any] | None
+)
 _OutputsType = Sequence[_NumericType]
 
 onnx_model_dir = os.path.join(
@@ -157,7 +157,7 @@ def _run_test(m, remained_onnx_input_idx, flatten=True, ignore_none=True):
 
 
 def run_ort(
-    onnx_model: Union[str, torch.onnx.ONNXProgram],
+    onnx_model: str | torch.onnx.ONNXProgram,
     pytorch_inputs: Sequence[_InputArgsType],
 ) -> _OutputsType:
     """Run ORT on the given ONNX model and inputs
@@ -268,13 +268,13 @@ class DecorateMeta:
     op_name: str
     variant_name: str
     decorator: Callable
-    opsets: Optional[Collection[Union[int, Callable[[int], bool]]]]
-    dtypes: Optional[Collection[torch.dtype]]
+    opsets: Collection[int | Callable[[int], bool]] | None
+    dtypes: Collection[torch.dtype] | None
     reason: str
     test_behavior: str
-    matcher: Optional[Callable[[Any], bool]] = None
+    matcher: Callable[[Any], bool] | None = None
     enabled_if: bool = True
-    model_type: Optional[pytorch_test_common.TorchModelType] = None
+    model_type: pytorch_test_common.TorchModelType | None = None
 
     def contains_opset(self, opset: int) -> bool:
         if self.opsets is None:
@@ -290,11 +290,11 @@ def xfail(
     variant_name: str = "",
     *,
     reason: str,
-    opsets: Optional[Collection[Union[int, Callable[[int], bool]]]] = None,
-    dtypes: Optional[Collection[torch.dtype]] = None,
-    matcher: Optional[Callable[[Any], bool]] = None,
+    opsets: Collection[int | Callable[[int], bool]] | None = None,
+    dtypes: Collection[torch.dtype] | None = None,
+    matcher: Callable[[Any], bool] | None = None,
     enabled_if: bool = True,
-    model_type: Optional[pytorch_test_common.TorchModelType] = None,
+    model_type: pytorch_test_common.TorchModelType | None = None,
 ):
     """Expects a OpInfo test to fail.
 
@@ -328,11 +328,11 @@ def skip(
     variant_name: str = "",
     *,
     reason: str,
-    opsets: Optional[Collection[Union[int, Callable[[int], bool]]]] = None,
-    dtypes: Optional[Collection[torch.dtype]] = None,
-    matcher: Optional[Callable[[Any], Any]] = None,
+    opsets: Collection[int | Callable[[int], bool]] | None = None,
+    dtypes: Collection[torch.dtype] | None = None,
+    matcher: Callable[[Any], Any] | None = None,
     enabled_if: bool = True,
-    model_type: Optional[pytorch_test_common.TorchModelType] = None,
+    model_type: pytorch_test_common.TorchModelType | None = None,
 ):
     """Skips a test case in OpInfo that we don't care about.
 
@@ -368,10 +368,10 @@ def skip_slow(
     variant_name: str = "",
     *,
     reason: str,
-    opsets: Optional[Collection[Union[int, Callable[[int], bool]]]] = None,
-    dtypes: Optional[Collection[torch.dtype]] = None,
-    matcher: Optional[Callable[[Any], Any]] = None,
-    model_type: Optional[pytorch_test_common.TorchModelType] = None,
+    opsets: Collection[int | Callable[[int], bool]] | None = None,
+    dtypes: Collection[torch.dtype] | None = None,
+    matcher: Callable[[Any], Any] | None = None,
+    model_type: pytorch_test_common.TorchModelType | None = None,
 ):
     """Skips a test case in OpInfo that is too slow.
 
@@ -423,14 +423,16 @@ def add_decorate_info(
             # Skip does not apply to this opset
             continue
         opinfo = ops_mapping.get((decorate_meta.op_name, decorate_meta.variant_name))
-        assert opinfo is not None, (
-            f"Couldn't find OpInfo for {decorate_meta}. Did you need to specify variant_name?"
-        )
-        assert decorate_meta.model_type is None, (
-            f"Tested op: {decorate_meta.op_name} in wrong position! "
-            "If model_type needs to be specified, it should be "
-            "put under SKIP_XFAIL_SUBTESTS_WITH_MATCHER_AND_MODEL_TYPE."
-        )
+        if opinfo is None:
+            raise AssertionError(
+                f"Couldn't find OpInfo for {decorate_meta}. Did you need to specify variant_name?"
+            )
+        if decorate_meta.model_type is not None:
+            raise AssertionError(
+                f"Tested op: {decorate_meta.op_name} in wrong position! "
+                "If model_type needs to be specified, it should be "
+                "put under SKIP_XFAIL_SUBTESTS_WITH_MATCHER_AND_MODEL_TYPE."
+            )
         decorators = list(opinfo.decorators)
         new_decorator = opinfo_core.DecorateInfo(
             decorate_meta.decorator,
@@ -468,28 +470,28 @@ def compare(other_opset: int):
 
 
 def reason_onnx_script_does_not_support(
-    operator: str, dtypes: Optional[Sequence[str]] = None
+    operator: str, dtypes: Sequence[str] | None = None
 ) -> str:
     """Formats the reason: ONNX script doesn't support the given dtypes."""
     return f"{operator} on {dtypes or 'dtypes'} not supported by ONNX script"
 
 
 def reason_onnx_runtime_does_not_support(
-    operator: str, dtypes: Optional[Sequence[str]] = None
+    operator: str, dtypes: Sequence[str] | None = None
 ) -> str:
     """Formats the reason: ONNX Runtime doesn't support the given dtypes."""
     return f"{operator} on {dtypes or 'dtypes'} not supported by ONNX Runtime"
 
 
 def reason_onnx_does_not_support(
-    operator: str, dtypes: Optional[Sequence[str]] = None
+    operator: str, dtypes: Sequence[str] | None = None
 ) -> str:
     """Formats the reason: ONNX doesn't support the given dtypes."""
     return f"{operator} on {dtypes or 'certain dtypes'} not supported by the ONNX Spec"
 
 
 def reason_dynamo_does_not_support(
-    operator: str, dtypes: Optional[Sequence[str]] = None
+    operator: str, dtypes: Sequence[str] | None = None
 ) -> str:
     """Formats the reason: Dynamo doesn't support the given dtypes."""
     return (
@@ -509,7 +511,7 @@ def reason_flaky() -> str:
 
 @contextlib.contextmanager
 def normal_xfail_skip_test_behaviors(
-    test_behavior: Optional[str] = None, reason: Optional[str] = None
+    test_behavior: str | None = None, reason: str | None = None
 ):
     """This context manager is used to handle the different behaviors of xfail and skip.
 
diff --git a/test/onnx/ops/test_ops.py b/test/onnx/ops/test_ops.py
index 6be99629f935f..5009be5a43740 100644
--- a/test/onnx/ops/test_ops.py
+++ b/test/onnx/ops/test_ops.py
@@ -167,7 +167,8 @@ def forward(self, x: torch.Tensor):
         onnx_program = torch.onnx.export(
             Model(), (torch.tensor(1),), dynamo=True, verbose=False
         )
-        assert onnx_program is not None
+        if onnx_program is None:
+            raise AssertionError("onnx_program is None")
         node = onnx_program.model.graph.node(0)
         self.assertEqual(node.op_type, "CustomOp")
         self.assertEqual(node.domain, "custom_domain")
@@ -208,7 +209,8 @@ def forward(self, x: torch.Tensor, y: torch.Tensor):
             dynamo=True,
             verbose=False,
         )
-        assert onnx_program is not None
+        if onnx_program is None:
+            raise AssertionError("onnx_program is None")
         node = onnx_program.model.graph.node(0)
         self.assertEqual(node.op_type, "CustomOp")
         self.assertEqual(node.domain, "custom_domain")
@@ -311,7 +313,8 @@ def forward(self, x: torch.Tensor):
         onnx_program = torch.onnx.export(
             Model(), (torch.tensor(1),), dynamo=True, verbose=False
         )
-        assert onnx_program is not None
+        if onnx_program is None:
+            raise AssertionError("onnx_program is None")
         node = onnx_program.model.graph.node(0)
         self.assertEqual(node.op_type, "CustomOp")
         self.assertEqual(node.domain, "custom_domain")
@@ -356,7 +359,8 @@ def forward(self, x: torch.Tensor, y: torch.Tensor):
             dynamo=True,
             verbose=False,
         )
-        assert onnx_program is not None
+        if onnx_program is None:
+            raise AssertionError("onnx_program is None")
         node = onnx_program.model.graph.node(0)
         self.assertEqual(node.op_type, "CustomOp")
         self.assertEqual(node.domain, "custom_domain")
@@ -427,7 +431,8 @@ def export(self, model, args=(), kwargs=None, **options) -> torch.onnx.ONNXProgr
             verbose=False,
             **options,
         )
-        assert onnx_program is not None
+        if onnx_program is None:
+            raise AssertionError("onnx_program is None")
         common_passes.CheckerPass()(onnx_program.model)
         return onnx_program
 
diff --git a/test/onnx/pytorch_test_common.py b/test/onnx/pytorch_test_common.py
index ffc359bd64020..0e63fb5d6832e 100644
--- a/test/onnx/pytorch_test_common.py
+++ b/test/onnx/pytorch_test_common.py
@@ -7,7 +7,6 @@
 import sys
 import unittest
 from enum import auto, Enum
-from typing import Optional
 
 import numpy as np
 import packaging.version
@@ -115,7 +114,7 @@ def wrapper(self, *args, **kwargs):
     return skip_dec
 
 
-def skipTraceTest(skip_before_opset_version: Optional[int] = None, reason: str = ""):
+def skipTraceTest(skip_before_opset_version: int | None = None, reason: str = ""):
     """Skip tracing test for opset version less than skip_before_opset_version.
 
     Args:
@@ -143,7 +142,7 @@ def wrapper(self, *args, **kwargs):
     return skip_dec
 
 
-def skipScriptTest(skip_before_opset_version: Optional[int] = None, reason: str = ""):
+def skipScriptTest(skip_before_opset_version: int | None = None, reason: str = ""):
     """Skip scripting test for opset version less than skip_before_opset_version.
 
     Args:
@@ -197,8 +196,8 @@ def wrapper(self, *args, **kwargs):
 
 def xfail_dynamic_fx_test(
     error_message: str,
-    model_type: Optional[TorchModelType] = None,
-    reason: Optional[str] = None,
+    model_type: TorchModelType | None = None,
+    reason: str | None = None,
 ):
     """Xfail dynamic exporting test.
 
@@ -275,7 +274,7 @@ def wrapper(self, *args, **kwargs):
     return skip_dec
 
 
-def xfail(error_message: str, reason: Optional[str] = None):
+def xfail(error_message: str, reason: str | None = None):
     """Expect failure.
 
     Args:
@@ -291,9 +290,10 @@ def inner(self, *args, **kwargs):
             try:
                 func(self, *args, **kwargs)
             except Exception as e:
-                assert error_message in str(e), (
-                    f"Expected error message: {error_message} NOT in {str(e)}"
-                )
+                if error_message not in str(e):
+                    raise AssertionError(
+                        f"Expected error message: {error_message} NOT in {str(e)}"
+                    ) from None
                 pytest.xfail(reason if reason else f"Expected failure: {error_message}")
             else:
                 pytest.fail("Unexpected success!")
@@ -340,7 +340,7 @@ def wrapper(self, *args, **kwargs):
 
 
 def xfail_if_model_type_is_exportedprogram(
-    error_message: str, reason: Optional[str] = None
+    error_message: str, reason: str | None = None
 ):
     """xfail test with models using ExportedProgram as input.
 
@@ -365,7 +365,7 @@ def wrapper(self, *args, **kwargs):
 
 
 def xfail_if_model_type_is_not_exportedprogram(
-    error_message: str, reason: Optional[str] = None
+    error_message: str, reason: str | None = None
 ):
     """xfail test without models using ExportedProgram as input.
 
diff --git a/test/onnx/test_onnx_opset.py b/test/onnx/test_onnx_opset.py
index f6e33fe599817..50cc9fdff40d0 100644
--- a/test/onnx/test_onnx_opset.py
+++ b/test/onnx/test_onnx_opset.py
@@ -19,11 +19,16 @@ def check_onnx_opset_operator(
     model, ops, opset_version=GLOBALS.export_onnx_opset_version
 ):
     # check_onnx_components
-    assert (
+    if not (
         model.producer_name == producer_name
         and model.producer_version == producer_version
         and model.opset_import[0].version == opset_version
-    )
+    ):
+        raise AssertionError(
+            f"Model metadata mismatch: producer_name={model.producer_name!r} (expected {producer_name!r}), "
+            f"producer_version={model.producer_version!r} (expected {producer_version!r}), "
+            f"opset_version={model.opset_import[0].version} (expected {opset_version})"
+        )
 
     # check the schema with the onnx checker
     onnx.checker.check_model(model)
@@ -35,17 +40,27 @@ def check_onnx_opset_operator(
     # At least the op_name should be specified,
     # but the op's attributes can optionally be
     # specified as well
-    assert len(ops) == len(graph.node)
+    if len(ops) != len(graph.node):
+        raise AssertionError(f"Expected {len(ops)} ops, got {len(graph.node)}")
     for i in range(len(ops)):
-        assert graph.node[i].op_type == ops[i]["op_name"]
+        if graph.node[i].op_type != ops[i]["op_name"]:
+            raise AssertionError(
+                f"Expected op {ops[i]['op_name']}, got {graph.node[i].op_type}"
+            )
         if "attributes" in ops[i]:
             attributes = ops[i]["attributes"]
-            assert len(attributes) == len(graph.node[i].attribute)
+            if len(attributes) != len(graph.node[i].attribute):
+                raise AssertionError(
+                    f"Expected {len(attributes)} attributes, got {len(graph.node[i].attribute)}"
+                )
             for j in range(len(attributes)):
                 for attribute_field in attributes[j]:
-                    assert attributes[j][attribute_field] == getattr(
-                        graph.node[i].attribute[j], attribute_field
-                    )
+                    actual = getattr(graph.node[i].attribute[j], attribute_field)
+                    if attributes[j][attribute_field] != actual:
+                        raise AssertionError(
+                            f"Attribute {attribute_field!r} mismatch on node {i}, attribute {j}: "
+                            f"expected {attributes[j][attribute_field]!r}, got {actual!r}"
+                        )
 
 
 def check_onnx_opsets_operator(
diff --git a/test/onnx/test_op_consistency.py b/test/onnx/test_op_consistency.py
index ee4742b25498e..073f503765ed8 100644
--- a/test/onnx/test_op_consistency.py
+++ b/test/onnx/test_op_consistency.py
@@ -25,7 +25,6 @@
 from __future__ import annotations
 
 import copy
-from typing import Optional
 
 import onnx_test_common
 import parameterized
@@ -218,7 +217,8 @@
 OP_WITH_SKIPPED_XFAIL_SUBTESTS = frozenset(meta.op_name for meta in SKIP_XFAIL_SUBTESTS)
 ALL_OPS_IN_DB = frozenset(op_info.name for op_info in OPS_DB)
 # Assert all ops in OPINFO_FUNCTION_MAPPING are in the OPS_DB
-assert TESTED_OPS.issubset(ALL_OPS_IN_DB), f"{TESTED_OPS - ALL_OPS_IN_DB} not in OPS_DB"
+if not TESTED_OPS.issubset(ALL_OPS_IN_DB):
+    raise AssertionError(f"{TESTED_OPS - ALL_OPS_IN_DB} not in OPS_DB")
 
 
 class SingleOpModel(torch.nn.Module):
@@ -235,14 +235,15 @@ def forward(self, *args):
 
 def _should_skip_xfail_test_sample(
     op_name: str, sample
-) -> tuple[Optional[str], Optional[str]]:
+) -> tuple[str | None, str | None]:
     """Returns a reason if a test sample should be skipped."""
     if op_name not in OP_WITH_SKIPPED_XFAIL_SUBTESTS:
         return None, None
     for decorator_meta in SKIP_XFAIL_SUBTESTS:
         # Linear search on ops_test_data.SKIP_XFAIL_SUBTESTS. That's fine because the list is small.
         if decorator_meta.op_name == op_name:
-            assert decorator_meta.matcher is not None, "Matcher must be defined"
+            if decorator_meta.matcher is None:
+                raise AssertionError("Matcher must be defined")
             if decorator_meta.matcher(sample):
                 return decorator_meta.test_behavior, decorator_meta.reason
     return None, None
@@ -281,7 +282,8 @@ class TestOnnxModelOutputConsistency(onnx_test_common._TestONNXRuntime):
     def test_output_match(self, device: str, dtype: torch.dtype, op):
         """Test the ONNX exporter."""
         # device is provided by instantiate_device_type_tests, but we only want to run in cpu.
-        assert device == "cpu"
+        if device != "cpu":
+            raise AssertionError(f"Expected device 'cpu', got {device!r}")
 
         samples = op.sample_inputs(
             device,
diff --git a/test/onnx/test_pytorch_onnx_onnxruntime.py b/test/onnx/test_pytorch_onnx_onnxruntime.py
index 459aca6f2b288..89fc795cd743c 100644
--- a/test/onnx/test_pytorch_onnx_onnxruntime.py
+++ b/test/onnx/test_pytorch_onnx_onnxruntime.py
@@ -9,7 +9,6 @@
 import os
 import unittest
 from collections import OrderedDict
-from typing import Optional, Union
 
 import numpy as np
 
@@ -95,8 +94,8 @@ def _init_test_rpn():
 
 def _construct_tensor_for_quantization_test(
     shape: tuple[int, ...],
-    offset: Optional[Union[int, float]] = None,
-    max_val: Optional[Union[int, float]] = None,
+    offset: int | float | None = None,
+    max_val: int | float | None = None,
 ) -> Tensor:
     """Helper function to generate weights and test inputs in a deterministic way.
 
@@ -361,7 +360,8 @@ def test_paste_mask_in_image(self):
             masks, boxes, [torch.tensor(o_im_s[0]), torch.tensor(o_im_s[1])]
         )
 
-        assert torch.all(out.eq(out_trace))
+        if not torch.all(out.eq(out_trace)):
+            raise AssertionError("Outputs do not match")
 
         masks2 = torch.rand(20, 1, 26, 26)
         boxes2 = torch.rand(20, 4)
@@ -375,7 +375,8 @@ def test_paste_mask_in_image(self):
             masks2, boxes2, [torch.tensor(o_im_s2[0]), torch.tensor(o_im_s2[1])]
         )
 
-        assert torch.all(out2.eq(out_trace2))
+        if not torch.all(out2.eq(out_trace2)):
+            raise AssertionError("Outputs do not match")
 
     def test_heatmaps_to_keypoints(self):
         maps = torch.rand(10, 1, 26, 26)
@@ -386,8 +387,10 @@ def test_heatmaps_to_keypoints(self):
         jit_trace = torch.jit.trace(heatmaps_to_keypoints, (maps, rois))
         out_trace = jit_trace(maps, rois)
 
-        assert torch.all(out[0].eq(out_trace[0]))
-        assert torch.all(out[1].eq(out_trace[1]))
+        if not torch.all(out[0].eq(out_trace[0])):
+            raise AssertionError("Outputs do not match at index 0")
+        if not torch.all(out[1].eq(out_trace[1])):
+            raise AssertionError("Outputs do not match at index 1")
 
         maps2 = torch.rand(20, 2, 21, 21)
         rois2 = torch.rand(20, 4)
@@ -396,8 +399,10 @@ def test_heatmaps_to_keypoints(self):
         out2 = heatmaps_to_keypoints(maps2, rois2)
         out_trace2 = jit_trace(maps2, rois2)
 
-        assert torch.all(out2[0].eq(out_trace2[0]))
-        assert torch.all(out2[1].eq(out_trace2[1]))
+        if not torch.all(out2[0].eq(out_trace2[0])):
+            raise AssertionError("Outputs do not match at index 0")
+        if not torch.all(out2[1].eq(out_trace2[1])):
+            raise AssertionError("Outputs do not match at index 1")
 
     def test_word_language_model_RNN_TANH(self):
         self.run_word_language_model("RNN_TANH")
@@ -509,8 +514,8 @@ def forward(self, x_in):
     def test_dict_output(self):
         class DictModelOutput(OrderedDict):
             tensor_out: Tensor
-            tuple_out: Optional[tuple[Tensor]] = None
-            list_out: Optional[list[Tensor]] = None
+            tuple_out: tuple[Tensor] | None = None
+            list_out: list[Tensor] | None = None
 
         class MyModel(torch.nn.Module):
             def forward(self, a, b, c, d):
@@ -581,8 +586,8 @@ class Model(torch.nn.Module):
             def forward(
                 self,
                 x,
-                y: Optional[Tensor] = None,
-                z: Optional[Tensor] = None,
+                y: Tensor | None = None,
+                z: Tensor | None = None,
             ):
                 if y is not None:
                     return x + y
@@ -610,8 +615,8 @@ class Model(torch.nn.Module):
             def forward(
                 self,
                 x,
-                y: Optional[Tensor] = torch.ones(2, 3),
-                z: Optional[Tensor] = torch.zeros(2, 3),
+                y: Tensor | None = torch.ones(2, 3),
+                z: Tensor | None = torch.zeros(2, 3),
             ):
                 if y is not None:
                     return x + y
@@ -634,8 +639,8 @@ class Model(torch.nn.Module):
             def forward(
                 self,
                 x,
-                y: Optional[Tensor] = torch.ones(2, 3),
-                z: Optional[Tensor] = torch.zeros(2, 3),
+                y: Tensor | None = torch.ones(2, 3),
+                z: Tensor | None = torch.zeros(2, 3),
             ):
                 if y is not None:
                     return x + y
@@ -669,7 +674,7 @@ def forward(
     @skipIfUnsupportedMinOpsetVersion(15)
     def test_all_optional_default_none(self):
         class Model(torch.nn.Module):
-            def forward(self, x: Optional[Tensor] = None, y: Optional[Tensor] = None):
+            def forward(self, x: Tensor | None = None, y: Tensor | None = None):
                 if x is not None:
                     return x
                 if y is not None:
@@ -694,8 +699,8 @@ def test_all_optional_default_tensor(self):
         class Model(torch.nn.Module):
             def forward(
                 self,
-                x: Optional[Tensor] = torch.ones(2, 3),
-                y: Optional[Tensor] = torch.zeros(2, 3),
+                x: Tensor | None = torch.ones(2, 3),
+                y: Tensor | None = torch.zeros(2, 3),
             ):
                 if x is not None:
                     return x
@@ -720,8 +725,8 @@ def test_all_optional_default_tensor_script(self):
         class Model(torch.nn.Module):
             def forward(
                 self,
-                x: Optional[Tensor] = torch.ones(2, 3),
-                y: Optional[Tensor] = torch.zeros(2, 3),
+                x: Tensor | None = torch.ones(2, 3),
+                y: Tensor | None = torch.zeros(2, 3),
             ):
                 if x is not None:
                     return x
@@ -818,7 +823,7 @@ def test_atleast_nd_single_tensor_input(
     @skipIfUnsupportedMinOpsetVersion(15)
     def test_mixed_optional(self):
         class Model(torch.nn.Module):
-            def forward(self, x, y: Optional[Tensor]):
+            def forward(self, x, y: Tensor | None):
                 if y is not None:
                     return x + y
                 return x
@@ -832,7 +837,7 @@ def forward(self, x, y: Optional[Tensor]):
     @skipIfUnsupportedMinOpsetVersion(15)
     def test_tuple_of_optional(self):
         class Model(torch.nn.Module):
-            def forward(self, x, y: tuple[Optional[Tensor], Optional[Tensor]]):
+            def forward(self, x, y: tuple[Tensor | None, Tensor | None]):
                 if y[0] is not None:
                     return x + y[0]
                 if y[1] is not None:
@@ -850,7 +855,7 @@ class Model(torch.nn.Module):
             def forward(
                 self,
                 x,
-                y: tuple[Optional[Tensor], Optional[Tensor]] = (
+                y: tuple[Tensor | None, Tensor | None] = (
                     torch.zeros(2, 3),
                     torch.zeros(2, 3),
                 ),
@@ -873,7 +878,7 @@ class Model(torch.nn.Module):
             def forward(
                 self,
                 x,
-                y: tuple[Optional[Tensor], Optional[Tensor]] = (
+                y: tuple[Tensor | None, Tensor | None] = (
                     torch.zeros(2, 3),
                     torch.zeros(2, 3),
                 ),
@@ -10816,7 +10821,8 @@ def forward(self, x):
         )
         ort_sess = verification._ort_session(model_onnx)
         ort_outs = verification._run_onnx(ort_sess, (x,))
-        assert not torch.all(torch.eq(x, torch.from_numpy(ort_outs[0])))
+        if torch.all(torch.eq(x, torch.from_numpy(ort_outs[0]))):
+            raise AssertionError("Outputs should not match input after dropout")
 
         script_model = torch.jit.script(model)
         output = model(x)
@@ -10831,7 +10837,8 @@ def forward(self, x):
             dynamo=False,
         )
         ort_outs = verification._run_onnx(ort_sess, (x,))
-        assert not torch.all(torch.eq(x, torch.from_numpy(ort_outs[0])))
+        if torch.all(torch.eq(x, torch.from_numpy(ort_outs[0]))):
+            raise AssertionError("Outputs should not match input after dropout")
 
     @skipIfUnsupportedMinOpsetVersion(12)
     def test_dropout_training_zero(self):
@@ -11234,7 +11241,7 @@ def get_embedding(embedding_dim: int):
                 )
                 return emb
 
-            def forward(self, input, incremental_state: Optional[Tensor] = None):
+            def forward(self, input, incremental_state: Tensor | None = None):
                 bsz, seq_len = input.shape[0], input.shape[1]
                 self.const = 3
                 if self.weights is None:
@@ -11295,7 +11302,7 @@ def get_embedding(embedding_dim: int):
                 )
                 return emb
 
-            def forward(self, input, incremental_state: Optional[Tensor] = None):
+            def forward(self, input, incremental_state: Tensor | None = None):
                 bsz, seq_len = input.shape[0], input.shape[1]
                 self.const = 1.5
                 self.weights = InnerModule.get_embedding(self.embedding_dim)
@@ -11350,13 +11357,13 @@ def __init__(self) -> None:
             def set_cell_anchors(self, anchors):
                 if self.conv.bias is not None:
                     b = self.conv.bias
-                    assert b is not None
+                    assert b is not None  # noqa: S101
                     self.conv.bias = anchors + b
                 elif self.conv.weight is not None:
                     self.conv.weight = torch.randn(3, 10)
                     self.conv.bias = self.conv.weight[:]
 
-            def forward(self, anchors) -> Optional[Tensor]:
+            def forward(self, anchors) -> Tensor | None:
                 self.set_cell_anchors(anchors)
                 return self.conv.bias
 
@@ -11403,7 +11410,7 @@ def set_cell_anchors(self, anchors):
                 self.conv.weight = torch.zeros(10, 3)
                 if self.conv.bias is not None:
                     w = self.conv.bias
-                    assert w is not None
+                    assert w is not None  # noqa: S101
                     self.conv.bias = anchors + w
                 else:
                     self.conv.bias = torch.ones(3, 10, 3)
@@ -11413,7 +11420,7 @@ def forward(self, feature_maps, anchors) -> tuple[Tensor, Tensor]:
                 result = []
                 if self.conv.bias is not None:
                     a = self.conv.bias
-                    assert a is not None
+                    assert a is not None  # noqa: S101
                     result += [a]
                 result += [feature_maps]
                 return result[0], result[1]
@@ -11443,7 +11450,7 @@ def set_cell_anchors(self, anchors):
                     # NOTE: `is not None` and `assert` is for passing torchscript.
                     if self.conv.bias is not None:
                         a = self.conv.bias
-                        assert a is not None
+                        assert a is not None  # noqa: S101
                         self.conv.bias = anchors + a
 
             def forward(self, anchors):
@@ -12661,12 +12668,14 @@ def forward(self, x, y):
         actual_std = np.std(ort_out)
         actual_mean = np.mean(ort_out)
 
-        assert abs(abs(actual_mean) - expected_mean) <= expected_mean * 0.1, (
-            "the gap of mean between ort outputs and expected one is unacceptable."
-        )
-        assert abs(abs(actual_std) - expected_std) <= expected_std * 0.1, (
-            "the gap of variance between ort outputs and expected one is unacceptable."
-        )
+        if not (abs(abs(actual_mean) - expected_mean) <= expected_mean * 0.1):
+            raise AssertionError(
+                "the gap of mean between ort outputs and expected one is unacceptable."
+            )
+        if not (abs(abs(actual_std) - expected_std) <= expected_std * 0.1):
+            raise AssertionError(
+                "the gap of variance between ort outputs and expected one is unacceptable."
+            )
 
     @skipScriptTest()
     @skipIfUnsupportedMinOpsetVersion(11)
@@ -12696,12 +12705,14 @@ def forward(self):
         actual_std = np.std(ort_out)
         actual_mean = np.mean(ort_out)
 
-        assert abs(abs(actual_mean) - expected_mean) <= expected_mean * 0.1, (
-            "the gap of mean between ort outputs and expected one is unacceptable."
-        )
-        assert abs(abs(actual_std) - expected_std) <= expected_std * 0.1, (
-            "the gap of variance between ort outputs and expected one is unacceptable."
-        )
+        if not (abs(abs(actual_mean) - expected_mean) <= expected_mean * 0.1):
+            raise AssertionError(
+                "the gap of mean between ort outputs and expected one is unacceptable."
+            )
+        if not (abs(abs(actual_std) - expected_std) <= expected_std * 0.1):
+            raise AssertionError(
+                "the gap of variance between ort outputs and expected one is unacceptable."
+            )
 
     @skipScriptTest()
     @skipIfUnsupportedMinOpsetVersion(11)
@@ -12744,15 +12755,12 @@ def forward(self, x, y):
         actual_max = np.max(ort_out)
         actual_mean = np.mean(ort_out)
 
-        assert actual_min >= expected_min, (
-            "the minimum value of ort outputs is out of scope."
-        )
-        assert actual_max <= expected_max, (
-            "the maximum value of ort outputs is out of scope."
-        )
-        assert abs(actual_mean - expected_mean) <= expected_mean * 0.05, (
-            "the mean value of ort outputs is out of scope."
-        )
+        if not (actual_min >= expected_min):
+            raise AssertionError("the minimum value of ort outputs is out of scope.")
+        if not (actual_max <= expected_max):
+            raise AssertionError("the maximum value of ort outputs is out of scope.")
+        if not (abs(actual_mean - expected_mean) <= expected_mean * 0.05):
+            raise AssertionError("the mean value of ort outputs is out of scope.")
 
     @skipIfUnsupportedMinOpsetVersion(13)
     def test_sequence_to_int(self):
@@ -13692,29 +13700,29 @@ def forward(self, input, grid):
                 )
 
     class IfNoneInput(torch.nn.Module):
-        def forward(self, x) -> Optional[Tensor]:
-            y: Optional[Tensor] = None
+        def forward(self, x) -> Tensor | None:
+            y: Tensor | None = None
             if x.size(0) > 1:
                 y = x
             return y
 
     class IfNoneOutput(torch.nn.Module):
-        def forward(self, x) -> Optional[Tensor]:
-            y: Optional[Tensor] = x
+        def forward(self, x) -> Tensor | None:
+            y: Tensor | None = x
             if x.size(0) > 1:
                 y = None
             return y
 
     class LoopNoneInput(torch.nn.Module):
-        def forward(self, x) -> Optional[Tensor]:
-            y: Optional[Tensor] = None
+        def forward(self, x) -> Tensor | None:
+            y: Tensor | None = None
             for _ in range(x.size(0)):
                 y = x
             return y
 
     class LoopNoneOutput(torch.nn.Module):
-        def forward(self, x) -> Optional[Tensor]:
-            y: Optional[Tensor] = x
+        def forward(self, x) -> Tensor | None:
+            y: Tensor | None = x
             for _ in range(x.size(0)):
                 y = None
             return y
@@ -13769,7 +13777,7 @@ def test_optional_output(self, module_class: type[torch.nn.Module], x_size: int)
     @skipIfUnsupportedMinOpsetVersion(16)
     def test_uninitialized_optional(self):
         class Module(torch.nn.Module):
-            def forward(self, y: Optional[Tensor]) -> Optional[Tensor]:
+            def forward(self, y: Tensor | None) -> Tensor | None:
                 if y is not None:
                     if y.shape[1] < 5:
                         if y.size(0) == 1:
diff --git a/test/onnx/test_pytorch_onnx_shape_inference.py b/test/onnx/test_pytorch_onnx_shape_inference.py
index e7c58e1ffdbe1..55363767b9288 100644
--- a/test/onnx/test_pytorch_onnx_shape_inference.py
+++ b/test/onnx/test_pytorch_onnx_shape_inference.py
@@ -367,7 +367,10 @@ def test_squeeze_after_dynamic_if(self):
         # Exporter will add "If" instead of raw "Squeeze" if it does not know
         # that if the dimension it is squeezing has size 1.
         squeezed = squeeze11(as_graphcontext(g), if_output, dim=0)
-        assert squeezed.node().kind() == "onnx::Squeeze"
+        if squeezed.node().kind() != "onnx::Squeeze":
+            raise AssertionError(
+                f"Expected 'onnx::Squeeze' but got {squeezed.node().kind()!r}"
+            )
         self.run_test(g, squeezed.node(), expect_tensor("Float", shape=(None, 5)))
 
 
@@ -402,7 +405,8 @@ def linalg_inv_settype(g, self):
         model_proto = onnx.load(io.BytesIO(f.getvalue()))
         model_value_info = model_proto.graph.value_info
         self.assertIsNotNone(model_value_info)
-        assert model_value_info
+        if not model_value_info:
+            raise AssertionError("model_value_info is empty")
         dims = model_value_info[0].type.tensor_type.shape.dim
         for i in range(len(dims)):
             # If node output has shape info, it should have dim_value
@@ -437,7 +441,8 @@ def linalg_inv_no_settype(g, self):
         model_proto = onnx.load(io.BytesIO(f.getvalue()))
         model_value_info = model_proto.graph.value_info
         self.assertIsNotNone(model_value_info)
-        assert model_value_info
+        if not model_value_info:
+            raise AssertionError("model_value_info is empty")
         dims = model_value_info[0].type.tensor_type.shape.dim
         for i in range(len(dims)):
             # If node output has shape info, it should have dim_value
@@ -476,7 +481,8 @@ def linalg_inv_settype(g, self):
         model_proto = onnx.load(io.BytesIO(f.getvalue()))
         model_value_info = model_proto.graph.value_info
         self.assertIsNotNone(model_value_info)
-        assert model_value_info
+        if not model_value_info:
+            raise AssertionError("model_value_info is empty")
         dims = model_value_info[0].type.tensor_type.shape.dim
         # The first axe should be dynamic as we defined when exporting
         self.assertTrue(dims[0].HasField("dim_param"))
@@ -522,12 +528,15 @@ def linalg_inv_settype(g, self):
             if node.op_type == "Inverse":
                 output_name = node.output[0]
                 break
-        assert output_name
+        if not output_name:
+            raise AssertionError("output_name not found")
         model_value_info = model_proto.graph.value_info
         self.assertIsNotNone(model_value_info)
-        assert model_value_info
+        if not model_value_info:
+            raise AssertionError("model_value_info is empty")
         for value_info in model_value_info:
-            assert value_info.name
+            if not value_info.name:
+                raise AssertionError("value_info.name is empty")
             if value_info.name == output_name:
                 dims = value_info.type.tensor_type.shape.dim
                 for i in range(len(dims)):
diff --git a/test/onnx/torchlib/ops_test_common.py b/test/onnx/torchlib/ops_test_common.py
index 93661dde12b91..d15fa66858469 100644
--- a/test/onnx/torchlib/ops_test_common.py
+++ b/test/onnx/torchlib/ops_test_common.py
@@ -13,7 +13,7 @@
 import unittest
 import warnings
 from collections.abc import Callable, Collection, Iterable, Mapping, Sequence
-from typing import Any, Optional, TypeVar
+from typing import Any, TypeVar
 
 import error_reproduction
 import numpy as np
@@ -68,15 +68,15 @@ class DecorateMeta:
     op_name: str
     variant_name: str
     decorator: Callable[..., Any]
-    dtypes: Optional[Collection[torch.dtype]]
-    device_type: Optional[str]
+    dtypes: Collection[torch.dtype] | None
+    device_type: str | None
     reason: str
     test_behavior: str
-    matcher: Optional[Callable[[Any], bool]] = None
+    matcher: Callable[[Any], bool] | None = None
     enabled_if: bool = True
     # The test_class_name to apply the decorator to. If None, the decorator is
     # applied to all test classes.
-    test_class_name: Optional[str] = None
+    test_class_name: str | None = None
 
 
 def xfail(
@@ -84,11 +84,11 @@ def xfail(
     variant_name: str = "",
     *,
     reason: str,
-    dtypes: Optional[Collection[torch.dtype]] = None,
-    device_type: Optional[str] = None,
-    matcher: Optional[Callable[[Any], Any]] = None,
+    dtypes: Collection[torch.dtype] | None = None,
+    device_type: str | None = None,
+    matcher: Callable[[Any], Any] | None = None,
     enabled_if: bool = True,
-    test_class_name: Optional[str] = None,
+    test_class_name: str | None = None,
 ) -> DecorateMeta:
     """Expects an OpInfo test to fail.
 
@@ -123,11 +123,11 @@ def skip(
     variant_name: str = "",
     *,
     reason: str,
-    dtypes: Optional[Collection[torch.dtype]] = None,
-    device_type: Optional[str] = None,
-    matcher: Optional[Callable[[Any], Any]] = None,
+    dtypes: Collection[torch.dtype] | None = None,
+    device_type: str | None = None,
+    matcher: Callable[[Any], Any] | None = None,
     enabled_if: bool = True,
-    test_class_name: Optional[str] = None,
+    test_class_name: str | None = None,
 ) -> DecorateMeta:
     """Skips an OpInfo test.
 
@@ -171,9 +171,10 @@ def add_decorate_info(
             # If the OpInfo doesn't exist and it is not enabled, we skip the OpInfo
             # because it could be an OpInfo that is in torch-nightly but not older versions.
             continue
-        assert opinfo is not None, (
-            f"Couldn't find OpInfo for {decorate_meta}. Did you need to specify variant_name?"
-        )
+        if opinfo is None:
+            raise AssertionError(
+                f"Couldn't find OpInfo for {decorate_meta}. Did you need to specify variant_name?"
+            )
         decorators = list(opinfo.decorators)
         new_decorator = opinfo_core.DecorateInfo(
             decorate_meta.decorator,
@@ -438,7 +439,8 @@ def dtype_op_schema_compatible(dtype: torch.dtype, schema) -> bool:
     # Since torch.float32 (tensor(float)) is in the allowed types, we return True.
 
     first_input_type_constraint = inputs[0].type_constraint
-    assert first_input_type_constraint is not None
+    if first_input_type_constraint is None:
+        raise AssertionError("first_input_type_constraint is None")
     allowed_types = first_input_type_constraint.allowed_types
     # Here we consider seq(tensor(float)) compatible with tensor(float) as well
     allowed_dtypes = {type_.dtype for type_ in allowed_types}
@@ -628,7 +630,7 @@ def _capture_graph_and_evaluate_torch_script_evaluator(
 
 @contextlib.contextmanager
 def normal_xfail_skip_test_behaviors(
-    test_behavior: Optional[str] = None, reason: Optional[str] = None
+    test_behavior: str | None = None, reason: str | None = None
 ):
     """This context manager is used to handle the different behaviors of xfail and skip.
 
diff --git a/test/onnx/torchlib/ops_test_data.py b/test/onnx/torchlib/ops_test_data.py
index d48ae8f1a28f9..916e180c0fdfd 100644
--- a/test/onnx/torchlib/ops_test_data.py
+++ b/test/onnx/torchlib/ops_test_data.py
@@ -39,7 +39,7 @@
 import copy
 import dataclasses
 import functools
-from typing import Any, Optional, TYPE_CHECKING
+from typing import Any, TYPE_CHECKING
 from typing_extensions import Self
 
 import numpy as np
@@ -71,9 +71,9 @@ class TorchLibOpInfo:
     # The torchlib ONNX Function to test
     op: Callable[..., Any]
     # The input wrangler function to adjust the input to fit the aten signature
-    input_wrangler: Optional[
-        Callable[[list[Any], dict[str, Any]], tuple[list[Any], dict[str, Any]]]
-    ] = None
+    input_wrangler: (
+        Callable[[list[Any], dict[str, Any]], tuple[list[Any], dict[str, Any]]] | None
+    ) = None
     # Whether the op is non-deterministic
     nondeterministic: bool = False
     # Whether to compare the shape only for the output[index]
@@ -113,11 +113,11 @@ def skip(
         variant_name: str = "",
         *,
         reason: str,
-        dtypes: Optional[Collection[torch.dtype]] = None,
-        device_type: Optional[str] = None,
-        matcher: Optional[Callable[[Any], Any]] = None,
+        dtypes: Collection[torch.dtype] | None = None,
+        device_type: str | None = None,
+        matcher: Callable[[Any], Any] | None = None,
         enabled_if: bool = True,
-        test_class_name: Optional[str] = None,
+        test_class_name: str | None = None,
     ) -> Self:
         """Skips an OpInfo test.
 
@@ -151,11 +151,11 @@ def xfail(
         variant_name: str = "",
         *,
         reason: str,
-        dtypes: Optional[Collection[torch.dtype]] = None,
-        device_type: Optional[str] = None,
-        matcher: Optional[Callable[[Any], Any]] = None,
+        dtypes: Collection[torch.dtype] | None = None,
+        device_type: str | None = None,
+        matcher: Callable[[Any], Any] | None = None,
         enabled_if: bool = True,
-        test_class_name: Optional[str] = None,
+        test_class_name: str | None = None,
     ) -> Self:
         """Expects an OpInfo test to fail.
 
@@ -710,7 +710,7 @@ def _where_input_wrangler(
 OP_WITH_SKIPPED_XFAIL_SUBTESTS = frozenset(meta.op_name for meta in SKIP_XFAIL_SUBTESTS)
 ALL_OPS_IN_DB = frozenset(op_info.name for op_info in OPS_DB)
 # Assert all ops in OPINFO_FUNCTION_MAPPING are in the OPS_DB
-assert TESTED_OPS.issubset(ALL_OPS_IN_DB), f"{TESTED_OPS - ALL_OPS_IN_DB} not in OPS_DB"
-assert NONDETERMINISTIC_OPS.issubset(TESTED_OPS), (
-    f"{NONDETERMINISTIC_OPS - TESTED_OPS} not in TESTED_OPS"
-)
+if not TESTED_OPS.issubset(ALL_OPS_IN_DB):
+    raise AssertionError(f"{TESTED_OPS - ALL_OPS_IN_DB} not in OPS_DB")
+if not NONDETERMINISTIC_OPS.issubset(TESTED_OPS):
+    raise AssertionError(f"{NONDETERMINISTIC_OPS - TESTED_OPS} not in TESTED_OPS")
diff --git a/test/onnx/torchlib/test_ops.py b/test/onnx/torchlib/test_ops.py
index f0144e3caa8de..7050a04c84e25 100644
--- a/test/onnx/torchlib/test_ops.py
+++ b/test/onnx/torchlib/test_ops.py
@@ -25,7 +25,7 @@
 from __future__ import annotations
 
 import os
-from typing import Optional, TYPE_CHECKING
+from typing import TYPE_CHECKING
 
 import error_reproduction
 import numpy as np
@@ -74,14 +74,15 @@ def dtypes_except(*dtypes: torch.dtype) -> Sequence[torch.dtype]:
 
 def _should_skip_xfail_test_sample(
     op_name: str, sample, dtype: torch.dtype, device_type: str
-) -> tuple[Optional[str], Optional[str]]:
+) -> tuple[str | None, str | None]:
     """Returns a reason if a test sample should be skipped."""
     if op_name not in ops_test_data.OP_WITH_SKIPPED_XFAIL_SUBTESTS:
         return None, None
     for decorator_meta in ops_test_data.SKIP_XFAIL_SUBTESTS:
         # Linear search on ops_test_data.SKIP_XFAIL_SUBTESTS. That's fine because the list is small.
         if decorator_meta.op_name == op_name:
-            assert decorator_meta.matcher is not None, "Matcher must be defined"
+            if decorator_meta.matcher is None:
+                raise AssertionError("Matcher must be defined")
             if not decorator_meta.enabled_if:
                 # Do not skip the test if the decorator meta is not enabled
                 continue
@@ -231,8 +232,13 @@ def run_test_output_match(
                 flattened_torch_outputs, _ = pytree.tree_flatten(torch_output)
                 flattened_function_outputs, _ = pytree.tree_flatten(function_output)
 
-                assert flattened_torch_outputs
-                assert len(flattened_torch_outputs) == len(flattened_function_outputs)
+                if not flattened_torch_outputs:
+                    raise AssertionError("flattened_torch_outputs is empty")
+                if len(flattened_torch_outputs) != len(flattened_function_outputs):
+                    raise AssertionError(
+                        f"Expected {len(flattened_torch_outputs)} outputs, "
+                        f"got {len(flattened_function_outputs)}"
+                    )
 
                 for j, (torch_output, function_output) in enumerate(
                     zip(flattened_torch_outputs, flattened_function_outputs)
diff --git a/test/profiler/test_memory_profiler.py b/test/profiler/test_memory_profiler.py
index 1a0c3347eb926..c9b01054929eb 100644
--- a/test/profiler/test_memory_profiler.py
+++ b/test/profiler/test_memory_profiler.py
@@ -5,7 +5,6 @@
 import textwrap
 import unittest
 from collections.abc import Callable, Iterator
-from typing import Optional
 
 import torch
 from torch._C._profiler import _EventType, _TensorMetadata
@@ -110,7 +109,7 @@ def gradient_detected(
         prof: torch.profiler.profile,
         ctx: _EventType,
         grad_tensor: torch.Tensor,
-        parameter: Optional[torch.Tensor] = None,
+        parameter: torch.Tensor | None = None,
     ) -> None:
         # This is not an exhaustive check, but for the purpose of unit testing
         # it is sufficient.
@@ -327,7 +326,7 @@ def formatSchemas(
     @staticmethod
     def _run_and_format_data_flow(
         inputs: dict[str, torch.Tensor],
-        f: Callable[..., Optional[dict[str, torch.Tensor]]],
+        f: Callable[..., dict[str, torch.Tensor] | None],
         indent: int = 12,
     ) -> str:
         with profile() as prof:
@@ -831,7 +830,7 @@ class TestMemoryProfilerE2E(TestCase):
     @staticmethod
     def _lookup_tensor_categories(
         t: torch.Tensor, memory_profile: _memory_profiler.MemoryProfile
-    ) -> dict[_memory_profiler.TensorAndID, Optional[_memory_profiler.Category]]:
+    ) -> dict[_memory_profiler.TensorAndID, _memory_profiler.Category | None]:
         storage = t.storage()
         if storage is None:
             raise ValueError("Cannot look up uninitialized Tensor.")
diff --git a/test/profiler/test_memory_viz.js b/test/profiler/test_memory_viz.js
new file mode 100644
index 0000000000000..46081fe489078
--- /dev/null
+++ b/test/profiler/test_memory_viz.js
@@ -0,0 +1,1131 @@
+// Test cases for process_alloc_data in torch/utils/viz/MemoryViz.js
+// Run: node test/profiler/test_memory_viz.js
+
+'use strict';
+
+// Polyfill for Node < 18
+if (!Array.prototype.findLastIndex) {
+  Array.prototype.findLastIndex = function(pred) {
+    for (let i = this.length - 1; i >= 0; i--) {
+      if (pred(this[i], i, this)) return i;
+    }
+    return -1;
+  };
+}
+if (!Array.prototype.at) {
+  Array.prototype.at = function(n) {
+    return n < 0 ? this[this.length + n] : this[n];
+  };
+}
+
+// ============================================================
+// Load process_alloc_data from the actual MemoryViz.js source
+// ============================================================
+
+const fs = require('fs');
+const path = require('path');
+const vm = require('vm');
+
+// Load process_alloc_data.js (ESM file) in Node.js CommonJS context.
+// Strip the `export` line since Node CommonJS doesn't support ESM syntax.
+const modPath = path.resolve(__dirname, '../../torch/utils/viz/process_alloc_data.js');
+let src = fs.readFileSync(modPath, 'utf-8');
+src = src.replace(/^export\s*\{[^}]*\};?\s*$/gm, '');
+const wrapper = `(function() { ${src}\nreturn { process_alloc_data, isPrivatePoolId, formatSize, formatAddr, elideRepeats }; })()`;
+const { process_alloc_data, isPrivatePoolId, formatSize, formatAddr, elideRepeats } = vm.runInThisContext(wrapper, { filename: modPath });
+
+// ============================================================
+// Test helpers
+// ============================================================
+
+function makeSnapshot({ traces = [], segments = [], categories = [] }) {
+  return {
+    device_traces: [traces],
+    segments,
+    categories,
+  };
+}
+
+let passed = 0;
+let failed = 0;
+
+function assert(condition, msg) {
+  if (!condition) {
+    failed++;
+    console.error(`  FAIL: ${msg}`);
+    console.trace();
+  } else {
+    passed++;
+  }
+}
+
+function assertEqual(actual, expected, msg) {
+  if (actual !== expected) {
+    failed++;
+    console.error(`  FAIL: ${msg} — expected ${expected}, got ${actual}`);
+  } else {
+    passed++;
+  }
+}
+
+function assertContains(str, substr, msg) {
+  if (!str.includes(substr)) {
+    failed++;
+    console.error(`  FAIL: ${msg} — "${substr}" not found in "${str.slice(0, 200)}..."`);
+  } else {
+    passed++;
+  }
+}
+
+// ============================================================
+// Tests
+// ============================================================
+
+function test_basic_alloc_free() {
+  console.log('test_basic_alloc_free');
+  const snapshot = makeSnapshot({
+    traces: [
+      { action: 'alloc', addr: 1000, size: 100, frames: [], stream: 0 },
+      { action: 'free_completed', addr: 1000, size: 100, frames: [], stream: 0 },
+    ],
+    segments: [{
+      device: 0, address: 0, total_size: 4096, segment_pool_id: [0, 0],
+      stream: 0, blocks: [],
+    }],
+  });
+  const result = process_alloc_data(snapshot, 0, false, 15000, false);
+  assertEqual(result.max_size, 100, 'peak should be 100');
+}
+
+function test_free_completed_is_matched() {
+  console.log('test_free_completed_is_matched');
+  // Only free_completed is matched (not free_requested). Verify alloc+free_completed works.
+  const snapshot = makeSnapshot({
+    traces: [
+      { action: 'alloc', addr: 1000, size: 200, frames: [], stream: 0 },
+      { action: 'free_completed', addr: 1000, size: 200, frames: [], stream: 0 },
+    ],
+    segments: [{
+      device: 0, address: 0, total_size: 4096, segment_pool_id: [0, 0],
+      stream: 0, blocks: [],
+    }],
+  });
+  const result = process_alloc_data(snapshot, 0, false, 15000, false);
+  assertEqual(result.max_size, 200, 'alloc+free_completed peak should be 200');
+}
+
+function test_pool_free_without_alloc_no_inflation() {
+  console.log('test_pool_free_without_alloc_no_inflation');
+  // Simulate: a private pool block was allocated BEFORE trace recording,
+  // then freed during recording. The trace only has free_completed.
+  // With include_private_inactive=true, the block should be pre-loaded into
+  // pool state and then freed — NOT treated as a new allocation.
+  const poolId = [1, 42];
+  const snapshot = makeSnapshot({
+    traces: [
+      // No alloc — it happened before recording started or got replaced in ring buffer
+      // Only the free_completed event is matched (free_requested is ignored).
+      // Note: real trace events never have segment_pool_id; pool is resolved via find_pool_id
+      // from the segment address ranges.
+      { action: 'free_completed', addr: 5000, size: 1000, frames: [], stream: 0 },
+    ],
+    segments: [{
+      device: 0, address: 4096, total_size: 8192, segment_pool_id: poolId,
+      stream: 0, blocks: [],
+    }],
+  });
+
+  const result = process_alloc_data(snapshot, 0, false, 15000, true);
+  // The free_completed creates an element in initially_allocated + actions.
+  // The fix pre-loads it into the pool (pool.active=1000, envelope grows to 1000).
+  // Then in actions, it's recognized as a free (pool.active goes to 0).
+  // Peak should be 1000 (the initial state), NOT 2000.
+  assert(result.max_size <= 1000,
+    `pool free-without-alloc should not inflate peak: got ${result.max_size}`);
+
+  // 1 element: the free_completed event that created an initially_allocated entry
+  assertEqual(result.elements_length, 1, 'should have 1 element');
+
+  // max_at_time should show the pool envelope at 1000, then dropping after the free
+  assert(result.max_at_time.length > 0, 'max_at_time should not be empty');
+  assertEqual(Math.max(...result.max_at_time), 1000,
+    'max_at_time peak should be 1000 (envelope high-water mark)');
+
+  // allocations_over_time should contain:
+  //   - pool envelope (elem = "pool:1,42")
+  //   - stripe for the block (elem = 0, the element index)
+  //   - summarized_mem
+  const envelopes = result.allocations_over_time.filter(
+    d => typeof d.elem === 'string' && d.elem.startsWith('pool:'));
+  assertEqual(envelopes.length, 1, 'should have 1 pool envelope');
+  assertEqual(envelopes[0].elem, 'pool:1,42,s0', 'envelope key matches pool id and stream');
+
+  const stripes = result.allocations_over_time.filter(
+    d => typeof d.elem === 'number' && d.opacity === 0.5);
+  assertEqual(stripes.length, 1, 'should have 1 pool stripe');
+  assertEqual(stripes[0].size, 1000, 'stripe size matches block size');
+}
+
+function test_pool_alloc_then_free_normal() {
+  console.log('test_pool_alloc_then_free_normal');
+  // Normal case: alloc + free both in trace, private pool, include_private_inactive=true
+  const poolId = [1, 7];
+  const snapshot = makeSnapshot({
+    traces: [
+      { action: 'alloc', addr: 2000, size: 500, frames: [], stream: 0 },
+      { action: 'free_completed', addr: 2000, size: 500, frames: [], stream: 0 },
+    ],
+    segments: [{
+      device: 0, address: 0, total_size: 8192, segment_pool_id: poolId,
+      stream: 0, blocks: [],
+    }],
+  });
+
+  const result = process_alloc_data(snapshot, 0, false, 15000, true);
+  assertEqual(result.max_size, 500, 'normal pool alloc/free peak should be 500');
+}
+
+function test_multiple_pool_frees_without_alloc() {
+  console.log('test_multiple_pool_frees_without_alloc');
+  // Multiple blocks freed from the same pool without matching allocs.
+  // This is the FSDP scenario with many free_storage calls.
+  const poolId = [1, 99];
+  const snapshot = makeSnapshot({
+    traces: [
+      // Each block has one free_completed (free_requested is ignored by the JS code)
+      { action: 'free_completed', addr: 1000, size: 500, frames: [], stream: 0 },
+      { action: 'free_completed', addr: 2000, size: 500, frames: [], stream: 0 },
+      { action: 'free_completed', addr: 3000, size: 500, frames: [], stream: 0 },
+    ],
+    segments: [{
+      device: 0, address: 0, total_size: 16384, segment_pool_id: poolId,
+      stream: 0, blocks: [],
+    }],
+  });
+
+  const result = process_alloc_data(snapshot, 0, false, 15000, true);
+  // 3 blocks of 500 each were initially allocated. Pool envelope = 1500.
+  // Then all 3 are freed. Peak should be 1500 (the initial state).
+  // BUG would give 3000+ (each free treated as a new alloc).
+  assert(result.max_size <= 1500,
+    `multiple pool frees should not inflate: got ${result.max_size}, expected <= 1500`);
+}
+
+function test_non_pool_free_without_alloc() {
+  console.log('test_non_pool_free_without_alloc');
+  // Non-pool block freed without matching alloc (ring buffer wrap).
+  // Should appear then disappear.
+  const snapshot = makeSnapshot({
+    traces: [
+      { action: 'free_completed', addr: 8000, size: 300, frames: [], stream: 0 },
+    ],
+    segments: [{
+      device: 0, address: 0, total_size: 16384, segment_pool_id: [0, 0],
+      stream: 0, blocks: [],
+    }],
+  });
+
+  const result = process_alloc_data(snapshot, 0, false, 15000, false);
+  // max_size is only updated inside the actions loop AFTER the free decrements total_mem,
+  // so for a free-without-alloc element, max_size ends up 0.
+  // The actual peak (300) is captured in max_at_time instead.
+  assertEqual(result.max_size, 0, 'non-pool free-without-alloc: max_size is 0 (peak is in max_at_time)');
+  assert(Math.max(...result.max_at_time) === 300,
+    'non-pool free-without-alloc: max_at_time peak should be 300');
+}
+
+function test_mixed_pool_and_nonpool() {
+  console.log('test_mixed_pool_and_nonpool');
+  // Mix of pool and non-pool allocations
+  const poolId = [1, 5];
+  const snapshot = makeSnapshot({
+    traces: [
+      // Non-pool alloc+free (addr within default pool segment)
+      { action: 'alloc', addr: 100, size: 200, frames: [], stream: 0 },
+      // Pool alloc+free (addr within private pool segment)
+      { action: 'alloc', addr: 5000, size: 400, frames: [], stream: 0 },
+      { action: 'free_completed', addr: 100, size: 200, frames: [], stream: 0 },
+      { action: 'free_completed', addr: 5000, size: 400, frames: [], stream: 0 },
+    ],
+    segments: [
+      { device: 0, address: 0, total_size: 4096, segment_pool_id: [0, 0], stream: 0, blocks: [] },
+      { device: 0, address: 4096, total_size: 8192, segment_pool_id: poolId, stream: 0, blocks: [] },
+    ],
+  });
+
+  const result = process_alloc_data(snapshot, 0, false, 15000, true);
+  // Peak: 200 (non-pool) + 400 (pool envelope) = 600
+  assertEqual(result.max_size, 600, 'mixed pool+nonpool peak should be 600');
+}
+
+function test_include_private_inactive_false_ignores_pools() {
+  console.log('test_include_private_inactive_false_ignores_pools');
+  // When include_private_inactive=false, pool blocks should be treated as regular
+  const poolId = [1, 10];
+  const snapshot = makeSnapshot({
+    traces: [
+      { action: 'alloc', addr: 5000, size: 800, frames: [], stream: 0 },
+      { action: 'free_completed', addr: 5000, size: 800, frames: [], stream: 0 },
+    ],
+    segments: [{
+      device: 0, address: 4096, total_size: 8192, segment_pool_id: poolId,
+      stream: 0, blocks: [],
+    }],
+  });
+
+  const result = process_alloc_data(snapshot, 0, false, 15000, false);
+  // Without pool logic, it's just a regular alloc+free. Peak = 800.
+  assertEqual(result.max_size, 800, 'with include_private_inactive=false, peak = 800');
+}
+
+// ============================================================
+// formatSize tests
+// ============================================================
+
+function test_formatSize_bytes() {
+  console.log('test_formatSize_bytes');
+  assertEqual(formatSize(0), '0.0B (0 bytes)', 'zero bytes');
+  assertEqual(formatSize(512), '512.0B (512 bytes)', '512 bytes');
+  assertEqual(formatSize(1), '1.0B (1 bytes)', '1 byte');
+}
+
+function test_formatSize_kib() {
+  console.log('test_formatSize_kib');
+  // 1024 bytes = 1.0 KiB
+  assertEqual(formatSize(1024), '1.0KiB (1024 bytes)', '1 KiB');
+  // 1536 = 1.5 * 1024
+  assertEqual(formatSize(1536), '1.5KiB (1536 bytes)', '1.5 KiB');
+}
+
+function test_formatSize_mib_gib() {
+  console.log('test_formatSize_mib_gib');
+  const mib = 1024 * 1024;
+  assertEqual(formatSize(mib), '1.0MiB (1048576 bytes)', '1 MiB');
+  const gib = 1024 * 1024 * 1024;
+  assertEqual(formatSize(gib), '1.0GiB (1073741824 bytes)', '1 GiB');
+}
+
+function test_formatSize_no_bytes() {
+  console.log('test_formatSize_no_bytes');
+  assertEqual(formatSize(1024, false), '1.0KiB', 'showBytes=false omits raw count');
+  assertEqual(formatSize(512, false), '512.0B', 'showBytes=false for small values');
+}
+
+// ============================================================
+// formatAddr tests
+// ============================================================
+
+function test_formatAddr_block_event() {
+  console.log('test_formatAddr_block_event');
+  const event = { action: 'alloc', addr: 0x7f4c00000, version: 3 };
+  assertEqual(formatAddr(event), "b'7f4c00000_3", 'block alloc address');
+}
+
+function test_formatAddr_segment_event() {
+  console.log('test_formatAddr_segment_event');
+  const event = { action: 'segment_alloc', addr: 0xabc, version: 0 };
+  assertEqual(formatAddr(event), "s'abc_0", 'segment alloc address');
+}
+
+function test_formatAddr_free_event() {
+  console.log('test_formatAddr_free_event');
+  const event = { action: 'free_completed', addr: 0xff, version: 5 };
+  assertEqual(formatAddr(event), "b'ff_5", 'free_completed is a block event');
+}
+
+// ============================================================
+// elideRepeats tests
+// ============================================================
+
+function test_elideRepeats_no_repeats() {
+  console.log('test_elideRepeats_no_repeats');
+  const result = elideRepeats(['a', 'b', 'c']);
+  assertEqual(result.join(','), 'a,b,c', 'no repeats passes through');
+}
+
+function test_elideRepeats_two_consecutive() {
+  console.log('test_elideRepeats_two_consecutive');
+  // Two consecutive duplicates are kept as-is (not collapsed)
+  const result = elideRepeats(['a', 'a', 'b']);
+  assertEqual(result.join(','), 'a,a,b', 'two consecutive kept verbatim');
+}
+
+function test_elideRepeats_three_or_more() {
+  console.log('test_elideRepeats_three_or_more');
+  // Three+ consecutive duplicates collapse to [frame, "<repeats N times>"]
+  const result = elideRepeats(['x', 'x', 'x', 'x']);
+  assertEqual(result.length, 2, 'collapsed to 2 entries');
+  assertEqual(result[0], 'x', 'first entry is the frame');
+  assertEqual(result[1], '<repeats 3 times>', 'second entry is repeat count');
+}
+
+function test_elideRepeats_mixed() {
+  console.log('test_elideRepeats_mixed');
+  // Realistic: a stack trace with a recursive section in the middle
+  const result = elideRepeats(['top', 'recurse', 'recurse', 'recurse', 'recurse', 'recurse', 'bottom']);
+  assertEqual(result.join(','), 'top,recurse,<repeats 4 times>,bottom', 'mixed with recursion');
+}
+
+function test_elideRepeats_empty() {
+  console.log('test_elideRepeats_empty');
+  const result = elideRepeats([]);
+  assertEqual(result.length, 0, 'empty input gives empty output');
+}
+
+// ============================================================
+// context_for_id tests
+// ============================================================
+
+function test_context_for_id_with_pool() {
+  console.log('test_context_for_id_with_pool');
+  // Alloc with a known private pool, stack frames, and stream
+  const poolId = [2, 7];
+  const snapshot = makeSnapshot({
+    traces: [
+      { action: 'alloc', addr: 0xabc000, size: 2048, version: 5,
+        frames: [{ filename: 'model.py', line: 42, name: 'forward' }],
+        stream: 3, timestamp: true, time_us: 1700000000000000 },
+    ],
+    segments: [{
+      device: 0, address: 0xa00000, total_size: 0x200000, segment_pool_id: poolId,
+      stream: 0, blocks: [],
+    }],
+  });
+
+  const result = process_alloc_data(snapshot, 0, false, 15000, false);
+  const ctx = result.context_for_id(0);
+
+  // Address should be formatted as hex with version
+  assertContains(ctx, "abc000", 'context includes hex address');
+  assertContains(ctx, '_5', 'context includes version');
+  // Size should be formatted
+  assertContains(ctx, '2.0KiB', 'context includes formatted size');
+  assertContains(ctx, '2048 bytes', 'context includes raw byte count');
+  // Pool ID resolved from segment
+  assertContains(ctx, 'pool_id (2, 7)', 'context includes resolved pool_id');
+  // Stream
+  assertContains(ctx, 'stream 3', 'context includes stream');
+  // Stack frame
+  assertContains(ctx, 'model.py:42:forward', 'context includes stack frame');
+}
+
+function test_context_for_id_unknown_pool() {
+  console.log('test_context_for_id_unknown_pool');
+  // Alloc at address outside any segment → pool_id unknown
+  const snapshot = makeSnapshot({
+    traces: [
+      { action: 'alloc', addr: 0xff0000, size: 512, version: 0,
+        frames: [{ filename: 'train.py', line: 10, name: 'step' }],
+        stream: 0, timestamp: null },
+    ],
+    segments: [{
+      device: 0, address: 0x100000, total_size: 0x100000, segment_pool_id: [0, 0],
+      stream: 0, blocks: [],
+    }],
+  });
+
+  const result = process_alloc_data(snapshot, 0, false, 15000, false);
+  const ctx = result.context_for_id(0);
+
+  assertContains(ctx, 'pool_id unknown', 'addr outside segments gives unknown pool');
+  assertContains(ctx, 'ff0000', 'context includes hex address');
+  assertContains(ctx, '512.0B', 'context includes size');
+}
+
+function test_context_for_id_free_without_alloc() {
+  console.log('test_context_for_id_free_without_alloc');
+  // free_completed without matching alloc → "alloc not recorded" message
+  const snapshot = makeSnapshot({
+    traces: [
+      { action: 'free_completed', addr: 0xdead00, size: 4096, version: 2,
+        frames: [{ filename: 'fsdp.py', line: 750, name: 'free_storage' }],
+        stream: 0, timestamp: null },
+    ],
+    segments: [{
+      device: 0, address: 0xd00000, total_size: 0x100000, segment_pool_id: [0, 0],
+      stream: 0, blocks: [],
+    }],
+  });
+
+  const result = process_alloc_data(snapshot, 0, false, 15000, false);
+  const ctx = result.context_for_id(0);
+
+  assertContains(ctx, 'alloc not recorded', 'free-without-alloc shows warning');
+  assertContains(ctx, 'dead00', 'context includes hex address');
+  assertContains(ctx, '4.0KiB', 'context includes size');
+  assertContains(ctx, 'fsdp.py:750:free_storage', 'context shows free stack trace');
+}
+
+// ============================================================
+// Post-PR#177717 tests: trace events carry segment_pool_id directly
+// ============================================================
+
+function test_post177717_pool_id_from_trace_event() {
+  console.log('test_post177717_pool_id_from_trace_event');
+  // After PR#177717, trace events include segment_pool_id. The code should
+  // use it directly instead of falling back to find_pool_id from segments.
+  // Here the addr is OUTSIDE any segment, but pool_id is on the event itself.
+  const poolId = [3, 15];
+  const snapshot = makeSnapshot({
+    traces: [
+      { action: 'alloc', addr: 0xff0000, size: 1024, frames: [], stream: 0,
+        segment_pool_id: poolId },
+      { action: 'free_completed', addr: 0xff0000, size: 1024, frames: [], stream: 0,
+        segment_pool_id: poolId },
+    ],
+    // No segment covers 0xff0000 — pool_id comes from the trace event
+    segments: [],
+  });
+
+  const result = process_alloc_data(snapshot, 0, false, 15000, true);
+  assertEqual(result.max_size, 1024, 'pool alloc/free with event-level pool_id');
+
+  const ctx = result.context_for_id(0);
+  assertContains(ctx, 'pool_id (3, 15)', 'pool_id resolved from trace event, not segment');
+}
+
+function test_post177717_pool_free_without_alloc_no_segment() {
+  console.log('test_post177717_pool_free_without_alloc_no_segment');
+  // Post-177717: free_completed has segment_pool_id on the event.
+  // The segment was unmapped (not in segments list), but pool_id is still known.
+  const poolId = [1, 42];
+  const snapshot = makeSnapshot({
+    traces: [
+      { action: 'free_completed', addr: 0xdead00, size: 2048, frames: [], stream: 0,
+        segment_pool_id: poolId },
+    ],
+    // Segment was unmapped — not present. Pool resolved from event.
+    segments: [],
+  });
+
+  const result = process_alloc_data(snapshot, 0, false, 15000, true);
+  // Should be pre-loaded into pool and then freed, not inflated
+  assert(result.max_size <= 2048,
+    `post-177717 pool free-without-alloc should not inflate: got ${result.max_size}`);
+
+  const ctx = result.context_for_id(0);
+  assertContains(ctx, 'pool_id (1, 42)', 'pool_id from event even without segment');
+  assertContains(ctx, 'alloc not recorded', 'still shows alloc not recorded');
+}
+
+function test_post177717_mixed_events_with_and_without_pool_id() {
+  console.log('test_post177717_mixed_events_with_and_without_pool_id');
+  // Some events have segment_pool_id (post-177717), others don't (pre-177717
+  // or default pool). Verify both paths work together.
+  const poolId = [2, 8];
+  const snapshot = makeSnapshot({
+    traces: [
+      // Default pool alloc — no segment_pool_id on event, resolved via segment
+      { action: 'alloc', addr: 100, size: 300, frames: [], stream: 0 },
+      // Private pool alloc — segment_pool_id on the event
+      { action: 'alloc', addr: 0xf000, size: 500, frames: [], stream: 0,
+        segment_pool_id: poolId },
+      { action: 'free_completed', addr: 100, size: 300, frames: [], stream: 0 },
+      { action: 'free_completed', addr: 0xf000, size: 500, frames: [], stream: 0,
+        segment_pool_id: poolId },
+    ],
+    segments: [
+      // Only covers addr=100 (default pool). addr=0xf000 has no segment.
+      { device: 0, address: 0, total_size: 4096, segment_pool_id: [0, 0],
+        stream: 0, blocks: [] },
+    ],
+  });
+
+  const result = process_alloc_data(snapshot, 0, false, 15000, true);
+  // Peak: 300 (non-pool) + 500 (pool envelope) = 800
+  assertEqual(result.max_size, 800, 'mixed pre/post-177717 events peak');
+
+  const ctx0 = result.context_for_id(0);
+  assertContains(ctx0, 'pool_id (0, 0)', 'default pool resolved from segment');
+
+  const ctx1 = result.context_for_id(1);
+  assertContains(ctx1, 'pool_id (2, 8)', 'private pool from event-level segment_pool_id');
+}
+
+
+// ============================================================
+// Pool grouping by (pool_id, stream) tests
+// ============================================================
+
+function test_pool_grouped_by_stream() {
+  console.log('test_pool_grouped_by_stream');
+  // Same pool_id but different streams should produce separate envelopes.
+  const poolId = [1, 5];
+  const snapshot = makeSnapshot({
+    traces: [
+      { action: 'alloc', addr: 1000, size: 500, frames: [], stream: 1 },
+      { action: 'alloc', addr: 2000, size: 300, frames: [], stream: 2 },
+      { action: 'free_completed', addr: 1000, size: 500, frames: [], stream: 1 },
+      { action: 'free_completed', addr: 2000, size: 300, frames: [], stream: 2 },
+    ],
+    segments: [
+      { device: 0, address: 0, total_size: 4096, segment_pool_id: poolId,
+        stream: 0, blocks: [] },
+    ],
+  });
+
+  const result = process_alloc_data(snapshot, 0, false, 15000, true);
+
+  const envelopes = result.allocations_over_time.filter(
+    d => typeof d.elem === 'string' && d.elem.startsWith('pool:'));
+  assertEqual(envelopes.length, 2, 'should have 2 pool envelopes (one per stream)');
+
+  const keys = envelopes.map(e => e.elem).sort();
+  assertEqual(keys[0], 'pool:1,5,s1', 'first envelope key includes stream 1');
+  assertEqual(keys[1], 'pool:1,5,s2', 'second envelope key includes stream 2');
+}
+
+// ============================================================
+// Initially Added Blocks Tests
+// ============================================================
+
+function test_default_pool_ghost_block() {
+  console.log('test_segment_snapshot_no_trace');
+  const poolId = [0, 0];
+  const snapshot = makeSnapshot({
+    traces: [],
+    segments: [{
+      device: 0, address: 4096, total_size: 8192, segment_pool_id: poolId,
+      stream: 0, blocks: [
+        { address: 5000, size: 1000, requested_size: 1000, state: 'active_allocated',
+          frames: [], version: 0 },
+      ],
+    }],
+  });
+
+  const result = process_alloc_data(snapshot, 0, false, 15000, false);
+  assertEqual(result.elements_length, 1,
+    'snapshot-only block should not be added (include_private_inactive=false)');
+
+}
+
+function test_segment_snapshot_with_trace_history() {
+  console.log('test_segment_snapshot_with_trace_history');
+  const poolId = [1, 42];
+  const snapshot = makeSnapshot({
+    traces: [
+      { action: 'alloc', addr: 5000, size: 1000, frames: [], stream: 0 },
+      { action: 'free_completed', addr: 5000, size: 1000, frames: [], stream: 0 },
+    ],
+    segments: [{
+      device: 0, address: 4096, total_size: 8192, segment_pool_id: poolId,
+      stream: 0, blocks: [
+        { address: 5000, size: 1000, requested_size: 1000, state: 'inactive',
+          frames: [], version: 0 },
+      ],
+    }],
+  });
+
+  const result = process_alloc_data(snapshot, 0, false, 15000, true);
+  assertEqual(result.elements_length, 1,
+    'only trace element present, snapshot block not duplicated');
+}
+
+function test_segment_snapshot_no_trace() {
+  console.log('test_segment_snapshot_no_trace');
+  const poolId = [1, 42];
+  const snapshot = makeSnapshot({
+    traces: [],
+    segments: [{
+      device: 0, address: 4096, total_size: 8192, segment_pool_id: poolId,
+      stream: 0, blocks: [
+        { address: 5000, size: 1000, requested_size: 1000, state: 'inactive',
+          frames: [], version: 0 },
+      ],
+    }],
+  });
+
+  const result = process_alloc_data(snapshot, 0, false, 15000, true);
+  assertEqual(result.elements_length, 0,
+    'snapshot-only block should not be added (include_private_inactive=true)');
+
+  const result2 = process_alloc_data(snapshot, 0, false, 15000, false);
+  assertEqual(result2.elements_length, 0,
+    'snapshot-only block should not be added (include_private_inactive=false)');
+}
+
+function test_ghost_blocks() {
+  console.log('test_ghost_blocks');
+  // Snapshot produced by (agent_space/test_ring_buffer_overflow.py):
+  //   pre_record = torch.empty(1024 * 1024, device="cuda", dtype=torch.uint8)  # 1 MiB
+  //   torch.cuda.memory._record_memory_history(max_entries=10)
+  //   early = torch.empty(2 * 1024 * 1024, device="cuda", dtype=torch.uint8)   # 2 MiB
+  //   for _ in range(15):  # overflow the 10-entry ring buffer
+  //     t = torch.empty(4 * 1024 * 1024, device="cuda", dtype=torch.uint8)     # 4 MiB
+  //     del t
+  //   snap = torch.cuda.memory._snapshot()
+  //
+  // pre_record: allocated before recording → no trace event at all
+  // early: alloc event evicted from ring buffer by churn → no trace event
+  // Both are active_allocated in segment snapshot but invisible in trace.
+  const snapshot = makeSnapshot({
+    traces: [
+      { action: 'free_completed', addr: 0x6a00000, size: 4194304, frames: [], stream: 0 },
+      { action: 'alloc', addr: 0x6a00000, size: 4194304, frames: [], stream: 0 },
+      { action: 'free_completed', addr: 0x6a00000, size: 4194304, frames: [], stream: 0 },
+      { action: 'alloc', addr: 0x6a00000, size: 4194304, frames: [], stream: 0 },
+      { action: 'free_completed', addr: 0x6a00000, size: 4194304, frames: [], stream: 0 },
+    ],
+    segments: [
+      { device: 0, address: 0x1e00000, total_size: 2097152, segment_pool_id: [0, 0],
+        stream: 0, blocks: [
+          { address: 0x1e00000, size: 1048576, requested_size: 1048576,
+            state: 'active_allocated', frames: [] },
+          { address: 0x1f00000, size: 1048576, requested_size: 1048576,
+            state: 'inactive', frames: [] },
+        ]},
+      { device: 0, address: 0x6800000, total_size: 20971520, segment_pool_id: [0, 0],
+        stream: 0, blocks: [
+          { address: 0x6800000, size: 2097152, requested_size: 2097152,
+            state: 'active_allocated', frames: [
+              { filename: 'test.py', line: 10, name: 'early_alloc' },
+            ]},
+          { address: 0x6a00000, size: 18874368, requested_size: 18874368,
+            state: 'inactive', frames: [] },
+        ]},
+    ],
+  });
+
+  // Ghost blocks (active_allocated not in trace) show on both tabs.
+  for (const include_private of [false, true]) {
+    const label = `include_private_inactive=${include_private}`;
+    const result = process_alloc_data(snapshot, 0, false, 15000, include_private);
+
+    // Trace creates 2 elements from alloc events + 1 from unmatched free.
+    // 2 active_allocated blocks from snapshot not in trace.
+    assertEqual(result.elements_length, 5,
+      `${label}: 3 trace elements + 2 snapshot blocks`);
+
+    const aot = result.allocations_over_time;
+    const ghosts = aot.filter(d => d.ghost === true);
+    assertEqual(ghosts.length, 2, `${label}: should have 2 ghost block entries`);
+
+    // Ghost block sizes match segment snapshot blocks
+    const ghost_sizes = ghosts.map(g => g.size).sort();
+    assertEqual(ghost_sizes[0], 1048576, `${label}: ghost block 1 MiB (pre_record)`);
+    assertEqual(ghost_sizes[1], 2097152, `${label}: ghost block 2 MiB (ring buffer overflow)`);
+
+    // context_for_id shows ghost explanation
+    const ghost_elem_ids = ghosts.map(g => g.elem);
+    for (const id of ghost_elem_ids) {
+      const ctx = result.context_for_id(id);
+      assertContains(ctx, '[Ghost block]', `${label}: context contains ghost label`);
+      assertContains(ctx, 'segment snapshot', `${label}: context explains source`);
+    }
+  }
+}
+
+function test_ghost_blocks_not_created_for_traced_addrs() {
+  console.log('test_ghost_blocks_not_created_for_traced_addrs');
+  // A block in the segment snapshot whose address DID appear in trace events
+  // should NOT be a ghost block.
+  const snapshot = makeSnapshot({
+    traces: [
+      { action: 'alloc', addr: 1000, size: 512, frames: [], stream: 0 },
+    ],
+    segments: [{
+      device: 0, address: 0, total_size: 4096, segment_pool_id: [0, 0],
+      stream: 0, blocks: [
+        { address: 1000, size: 512, requested_size: 512,
+          state: 'active_allocated', frames: [] },
+      ],
+    }],
+  });
+
+  const result = process_alloc_data(snapshot, 0, false, 15000, true);
+  const aot = result.allocations_over_time;
+  const ghosts = aot.filter(d => d.ghost === true);
+  assertEqual(ghosts.length, 0, 'no ghost blocks when addr is in trace');
+  assertEqual(result.elements_length, 1, 'only the trace element');
+}
+
+function test_ghost_blocks_default_pool_collected() {
+  console.log('test_ghost_blocks_default_pool_collected');
+  // Ghost blocks from default pool [0,0] should be collected
+  // from the segment snapshot when they have no trace events.
+  const snapshot = makeSnapshot({
+    traces: [],
+    segments: [{
+      device: 0, address: 0x1000, total_size: 8192, segment_pool_id: [0, 0],
+      stream: 0, blocks: [
+        { address: 0x1000, size: 2048, requested_size: 2048,
+          state: 'active_allocated', frames: [] },
+        { address: 0x1800, size: 1024, requested_size: 1024,
+          state: 'active_allocated', frames: [] },
+      ],
+    }],
+  });
+
+  // Shows on both tabs
+  for (const include_private of [false, true]) {
+    const label = `include_private_inactive=${include_private}`;
+    const result = process_alloc_data(snapshot, 0, false, 15000, include_private);
+    const ghosts = result.allocations_over_time.filter(d => d.ghost === true);
+    assertEqual(ghosts.length, 2, `${label}: 2 ghost blocks from default pool`);
+    const sizes = ghosts.map(g => g.size).sort();
+    assertEqual(sizes[0], 1024, `${label}: ghost 1024 bytes`);
+    assertEqual(sizes[1], 2048, `${label}: ghost 2048 bytes`);
+  }
+}
+
+function test_ghost_blocks_not_in_segment_mode() {
+  console.log('test_ghost_blocks_not_in_segment_mode');
+  // Ghost blocks should not be created in segment-level views
+  const snapshot = makeSnapshot({
+    traces: [],
+    segments: [{
+      device: 0, address: 0, total_size: 4096, segment_pool_id: [0, 0],
+      stream: 0, blocks: [
+        { address: 100, size: 512, requested_size: 512,
+          state: 'active_allocated', frames: [] },
+      ],
+    }],
+  });
+
+  const result_seg = process_alloc_data(snapshot, 0, true, 15000, false);
+  const ghosts_seg = result_seg.allocations_over_time.filter(d => d.ghost === true);
+  assertEqual(ghosts_seg.length, 0, 'no ghost blocks in segment_alloc mode');
+
+  // Even with include_private_inactive=true, segment modes should not create ghosts
+  const result_seg2 = process_alloc_data(snapshot, 0, true, 15000, true);
+  const ghosts_seg2 = result_seg2.allocations_over_time.filter(d => d.ghost === true);
+  assertEqual(ghosts_seg2.length, 0, 'no ghost blocks in segment_alloc mode (private pool tab)');
+}
+
+function test_ghost_blocks_private_pool() {
+  console.log('test_ghost_blocks_private_pool');
+  // Snapshot produced by (agent_space/test_ghost_blocks_private_pool.py):
+  //   pool = torch.cuda.MemPool()
+  //   with torch.cuda.use_mem_pool(pool):
+  //     pre_record_pool = torch.empty(1 MiB)     # ghost in private pool (0,1)
+  //   pre_record_default = torch.empty(2 MiB)     # ghost in default pool (0,0)
+  //   torch.cuda.memory._record_memory_history(max_entries=20)
+  //   with torch.cuda.use_mem_pool(pool):
+  //     traced_pool = torch.empty(3 MiB)          # traced in private pool (0,1)
+  //   traced_default = torch.empty(4 MiB)          # traced in default pool (0,0)
+  const snapshot = makeSnapshot({
+    traces: [
+      { action: 'alloc', addr: 0xe600000, size: 3145728, frames: [], stream: 0 },
+      { action: 'alloc', addr: 0x6a00000, size: 4194304, frames: [], stream: 0 },
+    ],
+    segments: [
+      // Private pool segment with ghost block
+      { device: 0, address: 0x1e00000, total_size: 2097152, segment_pool_id: [0, 1],
+        stream: 0, blocks: [
+          { address: 0x1e00000, size: 1048576, requested_size: 1048576,
+            state: 'active_allocated', frames: [] },
+        ]},
+      // Default pool segment with ghost block + traced block
+      { device: 0, address: 0x6800000, total_size: 20971520, segment_pool_id: [0, 0],
+        stream: 0, blocks: [
+          { address: 0x6800000, size: 2097152, requested_size: 2097152,
+            state: 'active_allocated', frames: [] },
+          { address: 0x6a00000, size: 4194304, requested_size: 4194304,
+            state: 'active_allocated', frames: [] },
+        ]},
+      // Private pool segment with traced block
+      { device: 0, address: 0xe600000, total_size: 20971520, segment_pool_id: [0, 1],
+        stream: 0, blocks: [
+          { address: 0xe600000, size: 3145728, requested_size: 3145728,
+            state: 'active_allocated', frames: [] },
+        ]},
+    ],
+  });
+
+  // With include_private_inactive=true, ghost blocks from BOTH pools are
+  // collected. Private pool ghosts go inside their pool envelope; default
+  // pool ghosts are rendered at the global bottom of the stacked area.
+  const result = process_alloc_data(snapshot, 0, false, 15000, true);
+
+  const aot = result.allocations_over_time;
+  const ghosts = aot.filter(d => d.ghost === true);
+  assertEqual(ghosts.length, 2, '2 ghost blocks (one from each pool)');
+
+  // Default pool ghost (2 MiB): at global bottom, spans full timeline
+  const default_ghost = ghosts.find(g => g.size === 2097152);
+  assert(default_ghost !== undefined, 'default pool ghost (2 MiB) exists');
+  assertEqual(default_ghost.offsets[0], 0, 'default ghost at offset 0');
+  assertEqual(default_ghost.timesteps[0], 0, 'default ghost starts at timestep 0');
+  assertEqual(default_ghost.timesteps.length, 2, 'default ghost has 2 timesteps');
+  assert(default_ghost.timesteps[1] > 0, 'default ghost ends after timestep 0');
+
+  // Private pool ghost (1 MiB): inside the pool (0,1) envelope
+  const pool_ghost = ghosts.find(g => g.size === 1048576);
+  assert(pool_ghost !== undefined, 'private pool ghost (1 MiB) exists');
+  assertEqual(pool_ghost.timesteps[0], 0, 'pool ghost starts at timestep 0');
+  assert(pool_ghost.timesteps.at(-1) > 0, 'pool ghost ends after timestep 0');
+
+  // Both ghosts should end at the same final timestep
+  assertEqual(default_ghost.timesteps[1], pool_ghost.timesteps.at(-1),
+    'both ghosts end at the same final timestep');
+
+  // Pool envelope only for (0,1) — default pool ghosts are at global bottom
+  const envelopes = aot.filter(d => typeof d.elem === 'string' && d.elem.startsWith('pool:'));
+  assertEqual(envelopes.length, 1, '1 pool envelope');
+  assertContains(envelopes[0].elem, '0,1', 'envelope for pool (0,1)');
+
+  // Envelope initial size at timestep 0 should already include the ghost block
+  const env = envelopes[0];
+  assertEqual(env.timesteps[0], 0, 'envelope starts at timestep 0');
+  assertEqual(env.size[0], 1048576, 'envelope initial size = ghost (1 MiB)');
+
+  // Ghost stripe should fit within the envelope
+  const env_offset = env.offsets[0];
+  const ghost_offset = pool_ghost.offsets[0];
+  assert(ghost_offset >= env_offset, 'ghost stripe offset >= envelope offset');
+  assert(ghost_offset + 1048576 <= env_offset + env.size[0],
+    'ghost stripe fits within envelope at timestep 0');
+
+  // Pool envelope max size should include both ghost (1 MiB) + traced (3 MiB) = 4 MiB
+  const env_max_size = Array.isArray(env.size)
+    ? Math.max(...env.size)
+    : env.size;
+  assertEqual(env_max_size, 1048576 + 3145728, 'pool envelope max = ghost + traced');
+
+  // With include_private_inactive=false, ghosts still exist but no pool envelope
+  const result_false = process_alloc_data(snapshot, 0, false, 15000, false);
+  const ghosts_false = result_false.allocations_over_time.filter(d => d.ghost === true);
+  assertEqual(ghosts_false.length, 2, '2 ghost blocks on regular tab too');
+  const envs_false = result_false.allocations_over_time.filter(
+    d => typeof d.elem === 'string' && d.elem.startsWith('pool:'));
+  assertEqual(envs_false.length, 0, 'no pool envelopes when include_private_inactive=false');
+}
+
+function test_ghost_stripe_offset_with_multiple_pools() {
+  console.log('test_ghost_stripe_offset_with_multiple_pools');
+  // When multiple private pools have initially_allocated blocks, pool envelopes
+  // are stacked. A ghost stripe must have its offset within its own envelope,
+  // not at the offset from when the stripe was first created (before other
+  // pools shifted it upward).
+  const snapshot = makeSnapshot({
+    traces: [
+      // Traced alloc in default pool so actions is non-empty
+      { action: 'alloc', addr: 0x100, size: 100, frames: [], stream: 0 },
+    ],
+    segments: [
+      // Default pool segment for the traced alloc
+      { device: 0, address: 0, total_size: 4096, segment_pool_id: [0, 0],
+        stream: 0, blocks: [] },
+      // Pool (0,3) segment 1: one ghost block.
+      // Added to initially_allocated first among private pools.
+      { device: 0, address: 0x4000, total_size: 4096, segment_pool_id: [0, 3],
+        stream: 0, blocks: [
+          { address: 0x4000, size: 3000, requested_size: 3000,
+            state: 'active_allocated', frames: [] },
+        ]},
+      // Pool (0,2): ghost block. Added to initially_allocated second.
+      { device: 0, address: 0x10000, total_size: 8192, segment_pool_id: [0, 2],
+        stream: 0, blocks: [
+          { address: 0x10000, size: 2000, requested_size: 2000,
+            state: 'active_allocated', frames: [] },
+        ]},
+      // Pool (0,3) segment 2: another ghost block.
+      // Added to initially_allocated LAST. After reverse(), processed FIRST.
+      // This creates pool (0,3) at offset 0. Then pool (0,2) is processed
+      // (stripe at offset 1000). Then pool (0,3) segment 1's ghost block
+      // grows pool (0,3) from 1000 to 4000, shifting pool (0,2)'s envelope
+      // up by 3000 but (without fix) not its stripe.
+      { device: 0, address: 0x5000, total_size: 4096, segment_pool_id: [0, 3],
+        stream: 0, blocks: [
+          { address: 0x5000, size: 1000, requested_size: 1000,
+            state: 'active_allocated', frames: [] },
+        ]},
+    ],
+  });
+
+  const result = process_alloc_data(snapshot, 0, false, 15000, true);
+  const aot = result.allocations_over_time;
+
+  const envelopes = aot.filter(d => typeof d.elem === 'string' && d.elem.startsWith('pool:'));
+  assertEqual(envelopes.length, 2, '2 pool envelopes');
+
+  // Find the pool (0,2) envelope and its ghost stripe
+  const env02 = envelopes.find(e => e.elem.includes('0,2'));
+  assert(env02 !== undefined, 'pool (0,2) envelope exists');
+
+  const ghosts = aot.filter(d => d.ghost === true);
+  assertEqual(ghosts.length, 3, '3 ghost blocks (2 in pool 0,3 + 1 in pool 0,2)');
+  // The ghost for pool (0,2) is the 2000-byte one
+  const ghost02 = ghosts.find(g => g.size === 2000);
+  assert(ghost02 !== undefined, 'pool (0,2) ghost exists');
+
+  // Ghost stripe offset must be within the envelope range
+  const env_offset = env02.offsets[0];
+  const env_size = env02.size[0];
+  const ghost_offset = ghost02.offsets[0];
+  assertEqual(ghost_offset, env_offset,
+    `ghost offset should equal envelope offset (single block in pool)`);
+  assert(ghost_offset + 2000 <= env_offset + env_size,
+    `ghost fits within envelope: ${ghost_offset}+2000 <= ${env_offset}+${env_size}`);
+}
+
+// ============================================================
+// Full snapshot integration test
+// ============================================================
+
+function test_full_snapshot_private_pools() {
+  console.log('test_full_snapshot_private_pools');
+  const snapshot = makeSnapshot({
+    traces: [
+      { action: 'segment_alloc', addr: 0x7f08cde00000, size: 2097152, frames: [], stream: 0 },
+      { action: 'alloc', addr: 0x7f08cde00000, size: 1024, frames: [], stream: 0 },
+      { action: 'segment_alloc', addr: 0x7f08d2800000, size: 2097152, frames: [], stream: 0 },
+      { action: 'alloc', addr: 0x7f08d2800000, size: 1024, frames: [], stream: 0 },
+      { action: 'free_requested', addr: 0x7f08cde00000, size: 1024, frames: [], stream: 0 },
+      { action: 'free_completed', addr: 0x7f08cde00000, size: 1024, frames: [], stream: 0 },
+      { action: 'free_requested', addr: 0x7f08d2800000, size: 1024, frames: [], stream: 0 },
+      { action: 'free_completed', addr: 0x7f08d2800000, size: 1024, frames: [], stream: 0 },
+      { action: 'alloc', addr: 0x7f08cde00000, size: 1024, frames: [], stream: 0 },
+      { action: 'alloc', addr: 0x7f08d2800000, size: 1024, frames: [], stream: 0 },
+      { action: 'alloc', addr: 0x7f08d2800400, size: 1024, frames: [], stream: 0 },
+      { action: 'alloc', addr: 0x7f08cde00400, size: 1024, frames: [], stream: 0 },
+    ],
+    segments: [
+      { device: 0, address: 0x7f08cde00000, total_size: 2097152,
+        segment_pool_id: [0, 0], stream: 0, blocks: [] },
+      { device: 0, address: 0x7f08d2800000, total_size: 2097152,
+        segment_pool_id: [0, 1], stream: 0, blocks: [] },
+    ],
+  });
+
+  const result = process_alloc_data(snapshot, 0, false, 15000, true);
+
+  assertEqual(result.elements_length, 6, 'should have 6 elements');
+  assertEqual(result.max_size, 4096, 'peak memory should be 4096');
+
+  const expected_max_at_time = [
+    1024, 2048, 2048, 2048, 2048, 2048, 2048, 2048, 2048,
+    1024, 2048, 2048, 3072, 3072, 3072, 3072, 4096,
+  ];
+  assertEqual(result.max_at_time.length, expected_max_at_time.length,
+    'max_at_time length');
+  for (let i = 0; i < expected_max_at_time.length; i++) {
+    assertEqual(result.max_at_time[i], expected_max_at_time[i],
+      `max_at_time[${i}]`);
+  }
+
+  const aot = result.allocations_over_time;
+  assertEqual(aot[0].elem, 0, 'first entry is element 0');
+  assertEqual(aot[0].size, 1024, 'element 0 size');
+
+  const envelopes = aot.filter(d => typeof d.elem === 'string' && d.elem.startsWith('pool:'));
+  assertEqual(envelopes.length, 1, 'should have 1 pool envelope');
+  assertEqual(envelopes[0].elem, 'pool:0,1,s0', 'envelope key matches pool (0,1)');
+
+  const stripes = aot.filter(d => typeof d.elem === 'number' && d.opacity === 0.5);
+  assertEqual(stripes.length, 3, 'should have 3 pool stripes (elements 1, 3, 4)');
+
+  const non_pool = aot.filter(d => typeof d.elem === 'number' && d.opacity === undefined);
+  assertEqual(non_pool.length, 3, 'should have 3 non-pool elements (0, 2, 5)');
+
+  const ctx0 = result.context_for_id(0);
+  assertContains(ctx0, '7f08cde00000', 'element 0 addr');
+  assertContains(ctx0, 'pool_id (0, 0)', 'element 0 pool');
+
+  const ctx1 = result.context_for_id(1);
+  assertContains(ctx1, '7f08d2800000', 'element 1 addr');
+  assertContains(ctx1, 'pool_id (0, 1)', 'element 1 pool');
+}
+
+function test_full_snapshot_no_private_pools() {
+  console.log('test_full_snapshot_no_private_pools');
+  const snapshot = makeSnapshot({
+    traces: [
+      { action: 'segment_alloc', addr: 0x7f08cde00000, size: 2097152, frames: [], stream: 0 },
+      { action: 'alloc', addr: 0x7f08cde00000, size: 1024, frames: [], stream: 0 },
+      { action: 'segment_alloc', addr: 0x7f08d2800000, size: 2097152, frames: [], stream: 0 },
+      { action: 'alloc', addr: 0x7f08d2800000, size: 1024, frames: [], stream: 0 },
+      { action: 'free_requested', addr: 0x7f08cde00000, size: 1024, frames: [], stream: 0 },
+      { action: 'free_completed', addr: 0x7f08cde00000, size: 1024, frames: [], stream: 0 },
+      { action: 'free_requested', addr: 0x7f08d2800000, size: 1024, frames: [], stream: 0 },
+      { action: 'free_completed', addr: 0x7f08d2800000, size: 1024, frames: [], stream: 0 },
+      { action: 'alloc', addr: 0x7f08cde00000, size: 1024, frames: [], stream: 0 },
+      { action: 'alloc', addr: 0x7f08d2800000, size: 1024, frames: [], stream: 0 },
+      { action: 'alloc', addr: 0x7f08d2800400, size: 1024, frames: [], stream: 0 },
+      { action: 'alloc', addr: 0x7f08cde00400, size: 1024, frames: [], stream: 0 },
+    ],
+    segments: [
+      { device: 0, address: 0x7f08cde00000, total_size: 2097152,
+        segment_pool_id: [0, 0], stream: 0, blocks: [] },
+      { device: 0, address: 0x7f08d2800000, total_size: 2097152,
+        segment_pool_id: [0, 1], stream: 0, blocks: [] },
+    ],
+  });
+
+  const result = process_alloc_data(snapshot, 0, false, 15000, false);
+
+  assertEqual(result.elements_length, 6, 'should have 6 elements');
+  assertEqual(result.max_size, 4096, 'peak memory should be 4096');
+
+  const expected_max_at_time = [
+    1024, 2048, 2048, 2048, 2048, 2048,
+    1024, 1024, 2048, 3072, 4096,
+  ];
+  assertEqual(result.max_at_time.length, expected_max_at_time.length,
+    'max_at_time length');
+  for (let i = 0; i < expected_max_at_time.length; i++) {
+    assertEqual(result.max_at_time[i], expected_max_at_time[i],
+      `max_at_time[${i}]`);
+  }
+
+  const aot = result.allocations_over_time;
+  const envelopes = aot.filter(d => typeof d.elem === 'string' && d.elem.startsWith('pool:'));
+  assertEqual(envelopes.length, 0, 'no pool envelopes');
+
+  const stripes = aot.filter(d => typeof d.elem === 'number' && d.opacity === 0.5);
+  assertEqual(stripes.length, 0, 'no pool stripes');
+
+  const regular = aot.filter(d => typeof d.elem === 'number');
+  assertEqual(regular.length, 6, 'all 6 elements are regular');
+
+  const ctx0 = result.context_for_id(0);
+  assertContains(ctx0, '7f08cde00000', 'element 0 addr');
+  assertContains(ctx0, 'Total memory used after allocation: 1.0KiB', 'element 0 total');
+  assertContains(ctx0, 'pool_id (0, 0)', 'element 0 pool');
+}
+
+// ============================================================
+// Run all tests
+// ============================================================
+
+test_basic_alloc_free();
+test_free_completed_is_matched();
+test_pool_free_without_alloc_no_inflation();
+test_pool_alloc_then_free_normal();
+test_multiple_pool_frees_without_alloc();
+test_non_pool_free_without_alloc();
+test_mixed_pool_and_nonpool();
+test_include_private_inactive_false_ignores_pools();
+test_formatSize_bytes();
+test_formatSize_kib();
+test_formatSize_mib_gib();
+test_formatSize_no_bytes();
+test_formatAddr_block_event();
+test_formatAddr_segment_event();
+test_formatAddr_free_event();
+test_elideRepeats_no_repeats();
+test_elideRepeats_two_consecutive();
+test_elideRepeats_three_or_more();
+test_elideRepeats_mixed();
+test_elideRepeats_empty();
+test_context_for_id_with_pool();
+test_context_for_id_unknown_pool();
+test_context_for_id_free_without_alloc();
+test_post177717_pool_id_from_trace_event();
+test_post177717_pool_free_without_alloc_no_segment();
+test_post177717_mixed_events_with_and_without_pool_id();
+test_pool_grouped_by_stream();
+test_segment_snapshot_with_trace_history();
+test_segment_snapshot_no_trace();
+test_default_pool_ghost_block();
+test_ghost_blocks();
+test_ghost_blocks_not_created_for_traced_addrs();
+test_ghost_blocks_default_pool_collected();
+test_ghost_blocks_not_in_segment_mode();
+test_ghost_blocks_private_pool();
+test_ghost_stripe_offset_with_multiple_pools();
+test_full_snapshot_private_pools();
+test_full_snapshot_no_private_pools();
+
+console.log(`\n${passed} passed, ${failed} failed`);
+process.exit(failed > 0 ? 1 : 0);
diff --git a/test/profiler/test_profiler.py b/test/profiler/test_profiler.py
index 868024270bf07..055f96bc8b79c 100644
--- a/test/profiler/test_profiler.py
+++ b/test/profiler/test_profiler.py
@@ -62,6 +62,7 @@
     IS_JETSON,
     IS_LINUX,
     IS_WINDOWS,
+    IS_X86,
     parametrize,
     run_tests,
     serialTest,
@@ -485,11 +486,11 @@ def join_threads(context: bool):
         self.assertEqual(len(observed_during_run), worker_threads)
         self.assertEqual(len(observed_during_run), len(set(observed_during_run)))
 
-    def payload(self, use_cuda=False):
-        x = torch.randn(10, 10)
+    def payload(self, use_cuda=False, tensor_size=10):
+        x = torch.randn(tensor_size, tensor_size)
         if use_cuda:
             x = x.cuda()
-        y = torch.randn(10, 10)
+        y = torch.randn(tensor_size, tensor_size)
         if use_cuda:
             y = y.cuda()
         z = torch.mm(x, y)
@@ -1439,6 +1440,32 @@ def test_profiler_strides(self):
                     if e["name"] == "aten::add":
                         self.assertEqual(args["Input Strides"], [[17, 1], [25, 2], []])
 
+    def test_profiler_strides_without_concrete_inputs(self):
+        torch._C._profiler._set_record_concrete_inputs_enabled_val(False)
+        try:
+            base_tensor = torch.randn(1024, dtype=torch.float32)
+            a = base_tensor.as_strided((16, 16), (17, 1), 0)
+            b = base_tensor.as_strided((16, 16), (25, 2), 272)
+            with _profile(record_shapes=True) as prof:
+                c = torch.add(a, b)
+
+            with TemporaryFileName(mode="w+") as fname:
+                prof.export_chrome_trace(fname)
+                with open(fname) as f:
+                    j = json.load(f)
+                    op_events = [
+                        e for e in j["traceEvents"] if e.get("cat", "") == "cpu_op"
+                    ]
+                    for e in op_events:
+                        args = e["args"]
+                        if e["name"] == "aten::add":
+                            self.assertIn("Input Strides", args)
+                            self.assertEqual(
+                                args["Input Strides"], [[17, 1], [25, 2], []]
+                            )
+        finally:
+            torch._C._profiler._set_record_concrete_inputs_enabled_val(True)
+
     def test_profiler_fwd_bwd_link(self):
         with _profile(use_kineto=True) as prof:
             t1, t2 = (
@@ -1525,7 +1552,7 @@ def workload() -> None:
             torch.cuda.synchronize()
             torch.add(t1, t2)
 
-        def trace_and_check(exp_config: Optional[_ExperimentalConfig]) -> None:
+        def trace_and_check(exp_config: _ExperimentalConfig | None) -> None:
             with _profile(
                 use_kineto=True,
                 use_device="cuda",
@@ -1538,10 +1565,12 @@ def trace_and_check(exp_config: Optional[_ExperimentalConfig]) -> None:
                 prof.export_chrome_trace(fname)
                 with open(fname) as f:
                     j = json.load(f)
-                    cats = {e.get("cat", None) for e in j["traceEvents"]}
+                    cat_or_name = "cat" if torch.version.cuda else "name"
+                    cats = {e.get(cat_or_name, None) for e in j["traceEvents"]}
+            event_name = "cuda_sync" if torch.version.cuda else "hipDeviceSynchronize"
             self.assertTrue(
-                "cuda_sync" in cats,
-                f"Expected to find cuda_sync event found = {cats}",
+                event_name in cats,
+                f"Expected to find {event_name} event found = {cats}",
             )
 
         print("Testing enable_cuda_sync_events in _ExperimentalConfig")
@@ -2424,7 +2453,7 @@ def validate_json(prof, disable_external_correlation):
                     disable_external_correlation=disable_external_correlation
                 ),
             ) as prof:
-                self.payload(use_cuda=True)
+                self.payload(use_cuda=True, tensor_size=256)
             validate_json(prof, disable_external_correlation)
 
     @skipIfTorchDynamo("profiler gets ignored if dynamo activated")
@@ -2697,6 +2726,140 @@ def test_public_api_post_processing_timeout_fails(self):
                 y = torch.randn(10, 10)
                 z = torch.mm(x, y)
 
+    @unittest.skipIf(not kineto_available(), "Kineto is required")
+    def test_profiler(self):
+        """Basic test for torch.profiler.profile API."""
+        use_cuda = torch.profiler.ProfilerActivity.CUDA in supported_activities()
+        activities = [ProfilerActivity.CPU]
+        if use_cuda:
+            activities.append(ProfilerActivity.CUDA)
+        with profile(activities=activities) as p:
+            self.payload(use_cuda=use_cuda)
+        events = p.events()
+        self.assertGreater(len(events), 0)
+        found_mm = False
+        for e in events:
+            if "aten::mm" in e.name:
+                found_mm = True
+        self.assertTrue(found_mm)
+        if use_cuda:
+            gpu_events = [e for e in events if e.device_type == DeviceType.CUDA]
+            self.assertGreater(len(gpu_events), 0, "No GPU events captured by profiler")
+
+    @unittest.skipIf(not torch.cuda.is_available(), "requires CUDA")
+    @unittest.skipIf(TEST_WITH_ROCM, "not supported on ROCm")
+    @unittest.skipIf(not kineto_available(), "Kineto is required")
+    def test_activity_filter_dict_syntax(self):
+        """Dict syntax collects only the requested activity types."""
+        with profile(
+            activities=[{ProfilerActivity.CUDA: ["GPU_MEMCPY", "CUDA_RUNTIME"]}],
+        ) as p:
+            x = torch.randn(10, 10).to("cuda")
+            y = torch.mm(x, x)
+        events = p.events()
+        self.assertGreater(len(events), 0)
+        print(events)
+        # Verify we got GPU_MEMCPY events (HtoD copy from .to("cuda")).
+        has_memcpy = any("Memcpy" in e.name for e in events)
+        self.assertTrue(has_memcpy, "Expected GPU_MEMCPY events")
+        # Verify we got CUDA_RUNTIME events (e.g. cudaLaunchKernel).
+        has_runtime = any("cuda" in e.name for e in events)
+        self.assertTrue(has_runtime, "Expected CUDA_RUNTIME events")
+        # OVERHEAD events (e.g. Lazy Function Loading) should NOT appear.
+        has_overhead = any("Lazy Function Loading" in e.name for e in events)
+        self.assertFalse(has_overhead)
+
+    @unittest.skipIf(not torch.cuda.is_available(), "requires CUDA")
+    @unittest.skipIf(not kineto_available(), "Kineto is required")
+    def test_activity_filter_mixed_syntax(self):
+        """Enum and dict entries can coexist for different activity groups."""
+        activities = [ProfilerActivity.CPU, {ProfilerActivity.CUDA: ["GPU_MEMCPY"]}]
+        with profile(activities=activities) as p:
+            with record_function("test_annotation"):
+                x = torch.randn(10, 10).to("cuda")
+                y = torch.mm(x, x)
+        self.assertGreater(len(p.events()), 0)
+
+    @unittest.skipIf(not kineto_available(), "Kineto is required")
+    def test_activity_filter_duplicate_raises(self):
+        """Same activity appearing more than once raises ValueError."""
+        with self.assertRaises(ValueError):
+            with profile(
+                activities=[
+                    ProfilerActivity.CPU,
+                    {ProfilerActivity.CPU: ["CPU_OP"]},
+                ],
+            ) as p:
+                pass
+
+    @unittest.skipIf(not kineto_available(), "Kineto is required")
+    def test_activity_filter_invalid_type_name(self):
+        """Invalid activity type name raises RuntimeError."""
+        with self.assertRaises(RuntimeError):
+            with profile(
+                activities=[{ProfilerActivity.CPU: ["NONEXISTENT_TYPE"]}],
+            ) as p:
+                x = torch.randn(10, 10)
+                y = torch.mm(x, x)
+
+    @unittest.skipIf(not torch.cuda.is_available(), "requires CUDA")
+    @unittest.skipIf(not kineto_available(), "Kineto is required")
+    def test_activity_filter_nonmember_type_name(self):
+        """Activity type name that is not a member of the requested activity group raises RuntimeError."""
+        with self.assertRaises(RuntimeError):
+            with profile(
+                activities=[{ProfilerActivity.CUDA: ["CPU_OP"]}],
+            ) as p:
+                x = torch.randn(10, 10)
+                y = torch.mm(x, x)
+
+    @unittest.skipIf(not torch.cuda.is_available(), "requires CUDA")
+    @unittest.skipIf(not kineto_available(), "Kineto is required")
+    def test_activity_filter_empty_list(self):
+        """Passing an empty list to activities means not collecting for the specified activity."""
+        with profile(
+            activities=[{ProfilerActivity.CUDA: []}],
+        ) as p:
+            x = torch.randn(10, 10).to("cuda")
+            y = torch.mm(x, x)
+        self.assertEqual(len(p.events()), 0)
+
+    @unittest.skipIf(not kineto_available(), "Kineto is required")
+    @unittest.skipIf(not TEST_CUDA, "CUDA is required")
+    def test_kineto_kernel_metadata_in_trace(self):
+        with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA]) as prof:
+            self.payload(use_cuda=True)
+
+        with TemporaryFileName(mode="w+") as fname:
+            prof.export_chrome_trace(fname)
+            with open(fname) as f:
+                trace = json.load(f)
+            events = trace["traceEvents"]
+            kernel_events = [e for e in events if e.get("cat", "") == "kernel"]
+            self.assertGreater(
+                len(kernel_events), 0, "Error: No kernel events in trace"
+            )
+            has_kernel_launch_metadata = False
+            for ke in kernel_events:
+                args = ke.get("args", {})
+                name = ke.get("name", "<unknown>")
+                for key in ["device", "stream", "correlation"]:
+                    self.assertIn(key, args, f"kernel '{name}' missing '{key}'")
+                # Some kernel events on ROCm (__amd_rocclr...) do not have grid/block metadata
+                # so we just validate that it shows up for at least one event
+                has_grid = "grid" in args
+                has_block = "block" in args
+                self.assertEqual(
+                    has_grid,
+                    has_block,
+                    f"kernel '{name}' should provide grid and block together",
+                )
+                has_kernel_launch_metadata |= has_grid
+            self.assertTrue(
+                has_kernel_launch_metadata,
+                "Error: No kernel events in trace contained grid/block metadata",
+            )
+
 
 class SimpleNet(nn.Module):
     def __init__(self) -> None:
@@ -3008,7 +3171,7 @@ def format_queue_depth(queue_depth_list, events):
 
     def test_utils_compute_queue_depth_when_no_cuda_events(self):
         # For traces with only cpu events, we expect empty queue depth list
-        x = torch.ones((1024, 1024))
+        x = torch.ones((100, 100))
         with profile() as prof:
             for _ in range(5):
                 x = x @ x
@@ -3058,7 +3221,7 @@ def test_utils_get_optimizable_events(self):
         )
 
     def test_profiler_name_pattern(self):
-        x = torch.ones((4096, 4096))
+        x = torch.ones((100, 100))
         with profile() as prof:
             for _ in range(5):
                 x = x @ x
@@ -3322,7 +3485,9 @@ def test_profiler_pattern_matcher_json_report(self):
                 actual_fields = sorted(event.keys())
                 self.assertEqual(expected_fields, actual_fields)
 
-    @unittest.skipIf(IS_ARM64 or not IS_LINUX, "x86 linux only cpp unwinding")
+    @unittest.skipIf(
+        not IS_LINUX or not (IS_X86 or IS_ARM64), "linux x86/aarch64 only cpp unwinding"
+    )
     def test_fuzz_symbolize(self):
         # generate some random addresses in the text section and make sure the
         # symbolizers do not throw exceptions/crash
@@ -3559,5 +3724,419 @@ def names(prof):
         self.assertEqual(n1, n2)
 
 
+class TestPrivateUse1ProfilerState(TestCase):
+    """Tests for PrivateUse1 profiler state selection logic."""
+
+    def test_kineto_privateuse1_state_with_use_kineto_true(self):
+        """Test that KINETO_PRIVATEUSE1 state is selected when use_kineto=True."""
+        from unittest.mock import patch
+
+        from torch._C._profiler import ProfilerState
+
+        with patch(
+            "torch.autograd.profiler._get_privateuse1_backend_name",
+            return_value="custom_backend",
+        ):
+            prof = _profile(
+                use_cpu=True,
+                use_device="custom_backend",
+                use_kineto=True,
+            )
+            self.assertEqual(prof.profiler_kind, ProfilerState.KINETO_PRIVATEUSE1)
+
+    def test_kineto_privateuse1_fallback_state_with_use_kineto_false(self):
+        """Test that KINETO_PRIVATEUSE1_FALLBACK is selected when use_kineto=False."""
+        from unittest.mock import patch
+
+        from torch._C._profiler import ProfilerState
+
+        with patch(
+            "torch.autograd.profiler._get_privateuse1_backend_name",
+            return_value="custom_backend",
+        ):
+            prof = _profile(
+                use_cpu=True,
+                use_device="custom_backend",
+                use_kineto=False,
+            )
+            self.assertEqual(
+                prof.profiler_kind, ProfilerState.KINETO_PRIVATEUSE1_FALLBACK
+            )
+
+    def test_privateuse1_fallback_requires_use_cpu(self):
+        """Test that KINETO_PRIVATEUSE1_FALLBACK requires use_cpu=True."""
+        from unittest.mock import patch
+
+        with patch(
+            "torch.autograd.profiler._get_privateuse1_backend_name",
+            return_value="custom_backend",
+        ):
+            # When use_kineto=False and use_cpu=False, should raise AssertionError
+            with self.assertRaises(AssertionError):
+                _profile(
+                    use_cpu=False,
+                    use_device="custom_backend",
+                    use_kineto=False,
+                )
+
+
+@unittest.skipIf(not kineto_available(), "Kineto is required")
+@unittest.skipIf(not torch.cuda.is_available(), "CUDA is required")
+class TestProfilerEventsParity(TestCase):
+    """Tests validating parity between events() and export_chrome_trace() JSON."""
+
+    def test_python_function_events_in_events(self):
+        class DummyModule(nn.Module):
+            def forward(self, x):
+                return x + 1
+
+        mod = DummyModule()
+        with profile(
+            activities=[ProfilerActivity.CPU],
+            with_stack=True,
+            experimental_config=_ExperimentalConfig(verbose=True),
+        ) as prof:
+            mod(torch.randn(4, 4))
+
+        events = prof.events()
+        python_events = [e for e in events if e.is_python_function]
+        self.assertGreater(len(python_events), 0)
+        for e in python_events:
+            self.assertIsInstance(e.name, str)
+            self.assertGreater(e.time_range.end - e.time_range.start, 0)
+
+        with TemporaryFileName(mode="w+") as fname:
+            prof.export_chrome_trace(fname)
+            with open(fname) as f:
+                trace = json.load(f)
+
+            json_py = [
+                e
+                for e in trace["traceEvents"]
+                if e.get("cat") == "python_function" and e.get("ph") == "X"
+            ]
+            self.assertEqual(len(python_events), len(json_py))
+
+            # Verify python_id/parent_id/module_id parity with JSON args
+            fe_mod = next((e for e in events if "DummyModule" in e.name), None)
+            self.assertIsNotNone(fe_mod)
+            self.assertGreater(fe_mod.python_id, 0)
+            self.assertGreaterEqual(fe_mod.python_module_id, 0)
+
+            json_mod = next(
+                (e for e in json_py if "DummyModule" in e.get("name", "")),
+                None,
+            )
+            self.assertIsNotNone(json_mod)
+            args = json_mod["args"]
+            self.assertEqual(fe_mod.python_id, args["Python id"])
+            self.assertEqual(fe_mod.python_parent_id, args["Python parent id"])
+            self.assertEqual(fe_mod.python_module_id, args["Python module id"])
+
+    def test_profiler_flow_events_parity(self):
+        """Verify that async CPU->GPU flow fields on events() match Chrome trace JSON."""
+        with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA]) as prof:
+            x = torch.randn(32, 32, device="cuda")
+            torch.mm(x, x)
+
+        # Collect async CPU->GPU flow info from events()
+        events_with_flow = [
+            e for e in prof.events() if e.flow_id is not None and e.flow_id != 0
+        ]
+        self.assertGreater(
+            len(events_with_flow), 0, "No flow events found via events()"
+        )
+
+        for e in events_with_flow:
+            self.assertIsInstance(e.flow_id, int)
+            self.assertIsInstance(e.flow_type, int)
+            self.assertIsInstance(e.flow_start, bool)
+
+        # Verify parity with Chrome trace JSON for async CPU->GPU flow
+        with TemporaryFileName(mode="w+") as fname:
+            prof.export_chrome_trace(fname)
+            with open(fname) as f:
+                j = json.load(f)
+
+            json_flow_events = [
+                e
+                for e in j["traceEvents"]
+                if e.get("ph") in ("s", "f") and e.get("cat") == "ac2g"
+            ]
+            json_flow_starts = {e["id"] for e in json_flow_events if e["ph"] == "s"}
+            json_flow_ends = {e["id"] for e in json_flow_events if e["ph"] == "f"}
+
+            # kLinkAsyncCpuGpu = 2
+            ac2g_events = [e for e in events_with_flow if e.flow_type == 2]
+            events_flow_starts = {e.flow_id for e in ac2g_events if e.flow_start}
+            events_flow_ends = {e.flow_id for e in ac2g_events if not e.flow_start}
+
+            self.assertEqual(
+                json_flow_starts,
+                events_flow_starts,
+                "Async CPU->GPU flow start IDs differ between events() and Chrome trace",
+            )
+            self.assertEqual(
+                json_flow_ends,
+                events_flow_ends,
+                "Async CPU->GPU flow end IDs differ between events() and Chrome trace",
+            )
+
+    def test_profiler_fwdbwd_flow_events_parity(self):
+        """Verify that fwd->bwd flow fields on events() match Chrome trace JSON."""
+        with profile(activities=[ProfilerActivity.CPU]) as prof:
+            t1 = torch.ones(1, requires_grad=True)
+            t2 = torch.ones(1, requires_grad=True)
+            z = torch.add(t1, t2)
+            y = torch.ones(1)
+            loss = torch.nn.functional.binary_cross_entropy_with_logits(z, y)
+            loss.backward()
+
+        fwdbwd_events = [
+            e for e in prof.events() if e.flow_type == 1 and e.flow_id != 0
+        ]
+        self.assertGreater(
+            len(fwdbwd_events), 0, "No fwdbwd flow events found via events()"
+        )
+
+        events_flow_starts = {e.flow_id for e in fwdbwd_events if e.flow_start}
+        events_flow_ends = {e.flow_id for e in fwdbwd_events if not e.flow_start}
+
+        with TemporaryFileName(mode="w+") as fname:
+            prof.export_chrome_trace(fname)
+            with open(fname) as f:
+                j = json.load(f)
+
+            json_flow_events = [
+                e
+                for e in j["traceEvents"]
+                if e.get("ph") in ("s", "f") and e.get("cat") == "fwdbwd"
+            ]
+            json_flow_starts = {e["id"] for e in json_flow_events if e["ph"] == "s"}
+            json_flow_ends = {e["id"] for e in json_flow_events if e["ph"] == "f"}
+
+            self.assertEqual(
+                json_flow_starts,
+                events_flow_starts,
+                "fwdbwd flow start IDs differ between events() and Chrome trace",
+            )
+            self.assertEqual(
+                json_flow_ends,
+                events_flow_ends,
+                "fwdbwd flow end IDs differ between events() and Chrome trace",
+            )
+
+    def test_profiler_timestamp_consistency(self):
+        """Verify that FunctionEvent timestamps can reconstruct Chrome trace ts values."""
+        with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA]) as prof:
+            x = torch.randn(32, 32, device="cuda")
+            torch.mm(x, x)
+
+        trace_start_ns = prof.profiler.kineto_results.trace_start_ns()
+
+        with TemporaryFileName(mode="w+") as fname:
+            prof.export_chrome_trace(fname)
+            with open(fname) as f:
+                j = json.load(f)
+
+            # Chrome trace is relative to a different base time which is not exposed in Python.
+            # It's probably not important to do so as we still have the relative differences
+            # in duration.
+            base_time_ns = j.get("baseTimeNanoseconds", 0)
+
+            # Grab mm timestamp from events() and json
+            fe_mm = next((e for e in prof.events() if e.name == "aten::mm"), None)
+            json_mm = next(
+                (
+                    e
+                    for e in j["traceEvents"]
+                    if e.get("name") == "aten::mm" and e.get("ph") == "X"
+                ),
+                None,
+            )
+
+            # Reconstruct Chrome trace ts from events():
+            # absolute_ns = mm_op_start_us * 1000 + trace_start_ns
+            # chrome_ts = (absolute_ns - base_time_ns) / 1000 -> realign with json timeframe
+            absolute_ns = int(fe_mm.time_range.start * 1000) + trace_start_ns
+            recovered_ts = (absolute_ns - base_time_ns) / 1000
+            self.assertEqual(
+                recovered_ts,
+                json_mm["ts"],
+                msg="Recovered Chrome trace ts doesn't match JSON for aten::mm",
+            )
+
+    def test_profiler_op_args_events_parity(self):
+        """Verify that cpu_op args on events() match Chrome trace JSON args."""
+        base_tensor = torch.randn(1024, dtype=torch.float32)
+        a = base_tensor.as_strided((16, 16), (17, 1), 0)
+        b = base_tensor.as_strided((16, 16), (25, 2), 272)
+        t1 = torch.ones((64, 32))
+        t2 = torch.ones((64, 32))
+        with profile(activities=[ProfilerActivity.CPU], record_shapes=True) as prof:
+            torch.add(a, b)
+            torch.cat([t1, t2])
+
+        fe_add = next((e for e in prof.events() if e.name == "aten::add"), None)
+        self.assertIsNotNone(fe_add)
+        fe_cat = next((e for e in prof.events() if e.name == "aten::cat"), None)
+        self.assertIsNotNone(fe_cat)
+
+        with TemporaryFileName(mode="w+") as fname:
+            prof.export_chrome_trace(fname)
+            with open(fname) as f:
+                j = json.load(f)
+
+            json_add = next(
+                (
+                    e
+                    for e in j["traceEvents"]
+                    if e.get("name") == "aten::add" and e.get("cat") == "cpu_op"
+                ),
+                None,
+            )
+            self.assertIsNotNone(json_add)
+            args = json_add["args"]
+            self.assertEqual(fe_add.structured_input_shapes, args["Input Dims"])
+            self.assertEqual(fe_add.structured_input_strides, args["Input Strides"])
+            self.assertEqual(fe_add.input_dtypes, args["Input type"])
+
+            # Test a case with TensorList inputs -- structured_input_shapes
+            # should handle TensorList nesting correctly.
+            json_cat = next(
+                (
+                    e
+                    for e in j["traceEvents"]
+                    if e.get("name") == "aten::cat" and e.get("cat") == "cpu_op"
+                ),
+                None,
+            )
+            self.assertIsNotNone(json_cat)
+            args_cat = json_cat["args"]
+            self.assertEqual(fe_cat.structured_input_shapes, args_cat["Input Dims"])
+            self.assertEqual(fe_cat.structured_input_strides, args_cat["Input Strides"])
+            self.assertEqual(fe_cat.input_dtypes, args_cat["Input type"])
+
+    def test_profiler_external_id_parity(self):
+        """Verify that FunctionEvent.external_id matches External id in Chrome trace JSON."""
+        from collections import Counter
+
+        with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA]) as prof:
+            with torch.profiler.record_function("test_region"):
+                x = torch.randn(32, 32, device="cuda")
+                y = torch.mm(x, x)
+                z = y + x
+                z.cpu()
+                torch.cuda.synchronize()
+
+        with TemporaryFileName(mode="w+") as fname:
+            prof.export_chrome_trace(fname)
+            with open(fname) as f:
+                j = json.load(f)
+
+            json_name_ext = Counter(
+                (e["name"], e["args"]["External id"])
+                for e in j["traceEvents"]
+                if e.get("args", {}).get("External id") is not None
+            )
+            events_name_ext = Counter(
+                (ev.name, ev.external_id) for ev in prof.events() if ev.external_id != 0
+            )
+
+            self.assertEqual(
+                events_name_ext,
+                json_name_ext,
+                "(name, external_id) pairs differ between events() and Chrome trace JSON",
+            )
+
+    def test_profiler_activity_type_parity(self):
+        """Verify activity_type on events() matches Chrome trace cat field."""
+        with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA]) as prof:
+            x = torch.randn(32, 32, device="cuda")
+            torch.mm(x, x)
+
+        events = prof.events()
+        for e in events:
+            self.assertIsInstance(e.activity_type, str)
+            self.assertGreater(len(e.activity_type), 0)
+
+        mm_event = next((e for e in events if e.name == "aten::mm"), None)
+        self.assertIsNotNone(mm_event)
+        self.assertEqual(mm_event.activity_type, "cpu_op")
+
+        with TemporaryFileName(mode="w+") as fname:
+            prof.export_chrome_trace(fname)
+            with open(fname) as f:
+                j = json.load(f)
+
+            json_name_cats = {
+                (e["name"], e["cat"])
+                for e in j["traceEvents"]
+                if e.get("ph") == "X" and "cat" in e
+            }
+            for e in events:
+                self.assertIn(
+                    (e.name, e.activity_type),
+                    json_name_cats,
+                    f"activity_type mismatch for {e.name}",
+                )
+
+    def test_structured_metadata_matches_chrome_trace(self):
+        from torch.autograd.profiler_util import _EVENT_METADATA_KEYS
+
+        with profile(
+            activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
+            experimental_config=torch._C._profiler._ExperimentalConfig(
+                expose_kineto_event_metadata=True
+            ),
+        ) as prof:
+            x = torch.randn(10, 10, device="cuda")
+            y = torch.mm(x, x)
+            z = x + y
+            z.cpu()
+
+        # Build lookup from External id -> FunctionEvent
+        fe_by_ext_id = {}
+        for fe in prof.events():
+            if fe.external_id != 0:
+                fe_by_ext_id[fe.external_id] = fe
+
+        with TemporaryFileName(mode="w+") as fname:
+            prof.export_chrome_trace(fname)
+            with open(fname) as f:
+                trace = json.load(f)
+
+            checked_kernel = 0
+            checked_memcpy = 0
+            for te in trace["traceEvents"]:
+                cat = te.get("cat", "")
+                args = te.get("args", {})
+                ext_id = args.get("External id")
+                if ext_id is None or ext_id not in fe_by_ext_id:
+                    continue
+                fe = fe_by_ext_id[ext_id]
+
+                if cat in ("kernel", "gpu_memcpy") and fe.event_metadata is not None:
+                    em = fe.event_metadata
+                    for kineto_key, (field_name, _) in _EVENT_METADATA_KEYS.items():
+                        if kineto_key in args:
+                            val = getattr(em, field_name)
+                            self.assertIsNotNone(
+                                val,
+                                f"'{fe.name}': {field_name} is None "
+                                f"but JSON has '{kineto_key}': {args[kineto_key]}",
+                            )
+                    if cat == "kernel":
+                        self.assertIsInstance(em.registers_per_thread, int)
+                        self.assertGreater(em.registers_per_thread, 0)
+                        checked_kernel += 1
+                    elif cat == "gpu_memcpy":
+                        if em.bytes is not None:
+                            self.assertGreater(em.bytes, 0)
+                        checked_memcpy += 1
+
+            self.assertGreater(checked_kernel, 0, "No kernel events were cross-checked")
+
+
 if __name__ == "__main__":
     run_tests()
diff --git a/test/profiler/test_python_tracer.py b/test/profiler/test_python_tracer.py
index f7732b0b3893f..930331cdbff86 100644
--- a/test/profiler/test_python_tracer.py
+++ b/test/profiler/test_python_tracer.py
@@ -3,8 +3,10 @@
 import json
 import subprocess
 import sys
+import threading
 import time
 
+import torch
 from torch.profiler import profile, ProfilerActivity
 from torch.testing._internal.common_utils import (
     run_tests,
@@ -104,6 +106,36 @@ def run():
             "Python replay stack is empty during pop operation" in result.stderr
         )
 
+    def test_concurrent_profiling(self):
+        """Repeatedly start/stop profiling while background threads are active.
+
+        On free-threaded Python (3.14t+), this exercises concurrent access to
+        the profiler's per-thread state without GIL protection. Without the
+        thread-safety fixes (setprofileAllThreads, per-thread ValueCache,
+        StopTheWorldGuard), this crashes from heap corruption due to data
+        races on the shared hash maps.
+        """
+        stop = threading.Event()
+
+        def work():
+            while not stop.is_set():
+                d = {str(i): list(range(i % 10)) for i in range(20)}
+                _ = sorted(d.items(), key=lambda x: len(x[1]))
+                torch.ones(10) + torch.zeros(10)
+
+        threads = [threading.Thread(target=work) for _ in range(8)]
+        for t in threads:
+            t.start()
+
+        try:
+            for _ in range(30):
+                with torch.profiler.profile(with_stack=True, with_modules=True):
+                    torch.ones(10)
+        finally:
+            stop.set()
+            for t in threads:
+                t.join(timeout=5)
+
 
 if __name__ == "__main__":
     run_tests()
diff --git a/test/python_native/test_dsl_registry.py b/test/python_native/test_dsl_registry.py
new file mode 100644
index 0000000000000..8b6012aa49788
--- /dev/null
+++ b/test/python_native/test_dsl_registry.py
@@ -0,0 +1,362 @@
+# Owner(s): ["module: dsl-native-ops"]
+
+from unittest.mock import Mock
+
+from torch._vendor.packaging.version import Version
+from torch.testing._internal.common_utils import run_tests, TestCase
+
+
+class TestDSLRegistry(TestCase):
+    """Comprehensive tests for DSL registry functionality."""
+
+    def setUp(self):
+        """Set up clean registry state for each test"""
+        # Import registry here to avoid import-time side effects
+        from torch._native.dsl_registry import DSLRegistry
+
+        # Save original registry state
+        self.original_registry = None
+        try:
+            from torch._native.dsl_registry import dsl_registry as original
+
+            self.original_modules = original._dsl_modules.copy()
+            self.original_registry = original
+        except ImportError:
+            pass
+
+        # Create isolated registry for testing
+        self.test_registry = DSLRegistry()
+
+    def tearDown(self):
+        """Restore original registry state"""
+        if self.original_registry is not None:
+            # Restore original DSL modules
+            self.original_registry._dsl_modules.clear()
+            self.original_registry._dsl_modules.update(self.original_modules)
+
+            # Clear any caches
+            if hasattr(self.original_registry.is_dsl_available, "cache_clear"):
+                self.original_registry.is_dsl_available.cache_clear()
+            if hasattr(self.original_registry.get_dsl_version, "cache_clear"):
+                self.original_registry.get_dsl_version.cache_clear()
+            if hasattr(self.original_registry.list_available_dsls, "cache_clear"):
+                self.original_registry.list_available_dsls.cache_clear()
+
+    def create_valid_mock_dsl(self, name="test_dsl", available=True, version="1.0.0"):
+        """Helper to create valid mock DSL for testing"""
+        mock = Mock()
+        mock.runtime_available.return_value = available
+        mock.runtime_version.return_value = Version(version) if version else None
+        mock.deregister_op_overrides = Mock()
+        mock.register_op_override = Mock()
+        return mock
+
+    def create_broken_mock_dsl(
+        self, break_method="runtime_available", error=ImportError
+    ):
+        """Helper to create mock DSL with broken methods for error testing"""
+        mock = self.create_valid_mock_dsl()
+        getattr(mock, break_method).side_effect = error("Simulated error")
+        return mock
+
+    # Phase 1: Basic Registry Operations (4 methods)
+
+    def test_register_dsl_basic(self):
+        """Test basic DSL registration functionality"""
+        mock_dsl = self.create_valid_mock_dsl()
+        self.test_registry.register_dsl("test_dsl", mock_dsl)
+
+        # Verify registration
+        self.assertIn("test_dsl", self.test_registry.list_all_dsls())
+        self.assertTrue(self.test_registry.is_dsl_available("test_dsl"))
+
+        # Verify the DSL module is stored correctly
+        self.assertEqual(len(self.test_registry.list_all_dsls()), 1)
+
+    def test_register_dsl_duplicates(self):
+        """Test duplicate DSL registration scenarios"""
+        scenarios = [
+            ("same_module", "re-registered with same module"),
+            ("different_module", "re-registered with different module"),
+        ]
+
+        for scenario, expected_log_contains in scenarios:
+            with self.subTest(scenario=scenario):
+                # Register initial DSL
+                mock_dsl1 = self.create_valid_mock_dsl()
+                dsl_name = f"duplicate_test_{scenario}"
+                self.test_registry.register_dsl(dsl_name, mock_dsl1)
+
+                # Register duplicate based on scenario
+                if scenario == "same_module":
+                    mock_dsl2 = mock_dsl1  # Same module object
+                else:
+                    mock_dsl2 = self.create_valid_mock_dsl()  # Different module object
+
+                # Capture logging
+                with self.assertLogs(
+                    "torch._native.dsl_registry", level="DEBUG"
+                ) as log_capture:
+                    self.test_registry.register_dsl(dsl_name, mock_dsl2)
+
+                # Verify appropriate logging
+                log_messages = " ".join(log_capture.output)
+                self.assertIn(expected_log_contains, log_messages)
+
+                # Verify DSL is still registered
+                self.assertIn(dsl_name, self.test_registry.list_all_dsls())
+
+    def test_get_dsl_version(self):
+        """Test DSL version querying with various conditions"""
+        test_cases = [("1.2.3", Version), (None, type(None)), ("error", type(None))]
+
+        for version_setup, expected_result_type in test_cases:
+            with self.subTest(version_setup=version_setup):
+                if version_setup == "error":
+                    mock_dsl = self.create_broken_mock_dsl(
+                        "runtime_version", RuntimeError
+                    )
+                else:
+                    mock_dsl = self.create_valid_mock_dsl(version=version_setup)
+
+                dsl_name = f"version_test_{version_setup}"
+                self.test_registry.register_dsl(dsl_name, mock_dsl)
+
+                result = self.test_registry.get_dsl_version(dsl_name)
+                self.assertIsInstance(result, expected_result_type)
+
+                if version_setup == "1.2.3":
+                    self.assertEqual(result, Version("1.2.3"))
+
+    def test_list_dsls_operations(self):
+        """Test list_all_dsls and list_available_dsls"""
+        # Register mix of available/unavailable DSLs
+        available_dsl = self.create_valid_mock_dsl("available", available=True)
+        unavailable_dsl = self.create_valid_mock_dsl("unavailable", available=False)
+
+        self.test_registry.register_dsl("available_dsl", available_dsl)
+        self.test_registry.register_dsl("unavailable_dsl", unavailable_dsl)
+
+        # Test list operations
+        all_dsls = self.test_registry.list_all_dsls()
+        available_dsls = self.test_registry.list_available_dsls()
+
+        # Verify results
+        self.assertEqual(set(all_dsls), {"available_dsl", "unavailable_dsl"})
+        self.assertEqual(set(available_dsls), {"available_dsl"})
+
+        # Verify available is subset of all
+        self.assertTrue(set(available_dsls).issubset(set(all_dsls)))
+
+    # Phase 2: Input Validation (2 methods)
+
+    def test_register_dsl_invalid_names(self):
+        """Test registration with invalid name inputs"""
+        test_cases = [
+            (None, TypeError),
+            (123, TypeError),
+            ([], TypeError),
+            ({}, TypeError),
+            ("", ValueError),
+            ("   ", ValueError),
+            ("\t\n", ValueError),
+        ]
+
+        mock_dsl = self.create_valid_mock_dsl()
+
+        for invalid_name, expected_exception in test_cases:
+            with self.subTest(invalid_name=repr(invalid_name)):
+                with self.assertRaises(expected_exception):
+                    self.test_registry.register_dsl(invalid_name, mock_dsl)
+
+    def test_register_dsl_valid_names(self):
+        """Test registration with valid name formats"""
+        valid_names = ["triton", "cutedsl", "my_dsl", "dsl_v2", "a", "test_dsl_123"]
+
+        for valid_name in valid_names:
+            with self.subTest(valid_name=valid_name):
+                mock_dsl = self.create_valid_mock_dsl()
+                self.test_registry.register_dsl(valid_name, mock_dsl)
+                self.assertIn(valid_name, self.test_registry.list_all_dsls())
+
+    # Phase 3: Error Handling (3 methods)
+
+    def test_is_dsl_available_errors(self):
+        """Test is_dsl_available when runtime_available() raises errors"""
+        error_types = [ImportError, ModuleNotFoundError, RuntimeError, AttributeError]
+
+        for error_type in error_types:
+            with self.subTest(error_type=error_type.__name__):
+                mock_dsl = self.create_broken_mock_dsl("runtime_available", error_type)
+                dsl_name = f"broken_dsl_{error_type.__name__}"
+                self.test_registry.register_dsl(dsl_name, mock_dsl)
+
+                result = self.test_registry.is_dsl_available(dsl_name)
+                self.assertEqual(result, False)  # All errors should result in False
+
+    def test_get_dsl_version_errors(self):
+        """Test get_dsl_version when runtime_version() raises errors"""
+        error_types = [ImportError, RuntimeError, AttributeError, TypeError]
+
+        for error_type in error_types:
+            with self.subTest(error_type=error_type.__name__):
+                mock_dsl = self.create_broken_mock_dsl("runtime_version", error_type)
+                dsl_name = f"broken_version_dsl_{error_type.__name__}"
+                self.test_registry.register_dsl(dsl_name, mock_dsl)
+
+                result = self.test_registry.get_dsl_version(dsl_name)
+                self.assertIsNone(result)
+
+    def test_nonexistent_dsl_queries(self):
+        """Test querying non-existent DSLs returns appropriate defaults"""
+        # Test with empty registry
+        self.assertFalse(self.test_registry.is_dsl_available("nonexistent"))
+        self.assertIsNone(self.test_registry.get_dsl_version("nonexistent"))
+        self.assertEqual(self.test_registry.list_all_dsls(), [])
+        self.assertEqual(self.test_registry.list_available_dsls(), [])
+
+        # Test with some DSLs registered but querying non-existent
+        mock_dsl = self.create_valid_mock_dsl()
+        self.test_registry.register_dsl("existing", mock_dsl)
+
+        self.assertFalse(self.test_registry.is_dsl_available("still_nonexistent"))
+        self.assertIsNone(self.test_registry.get_dsl_version("still_nonexistent"))
+
+    # Phase 4: Protocol & Integration (4 methods)
+
+    def test_dsl_protocol_interface(self):
+        """Test DSL modules implement complete protocol"""
+        # Test with actual registered DSLs from the global registry
+        if self.original_registry:
+            all_dsls = self.original_registry.list_all_dsls()
+
+            for dsl_name in all_dsls:
+                # Get the actual DSL module from registry
+                dsl_module = self.original_registry._dsl_modules.get(dsl_name)
+                if dsl_module:
+                    # Verify all required methods exist and are callable
+                    required_methods = [
+                        "runtime_available",
+                        "runtime_version",
+                        "deregister_op_overrides",
+                        "register_op_override",
+                    ]
+
+                    for method_name in required_methods:
+                        self.assertTrue(
+                            hasattr(dsl_module, method_name),
+                            f"DSL '{dsl_name}' missing method '{method_name}'",
+                        )
+                        method = getattr(dsl_module, method_name)
+                        self.assertTrue(
+                            callable(method),
+                            f"DSL '{dsl_name}' method '{method_name}' is not callable",
+                        )
+
+    def test_real_dsl_integration(self):
+        """Test integration with actual DSL modules"""
+        if not self.original_registry:
+            self.skipTest("Original registry not available")
+
+        dsl_names = ["triton", "cutedsl"]
+
+        for dsl_name in dsl_names:
+            with self.subTest(dsl_name=dsl_name):
+                # Verify DSL is registered in original registry
+                all_dsls = self.original_registry.list_all_dsls()
+                if dsl_name not in all_dsls:
+                    self.skipTest(
+                        f"DSL '{dsl_name}' not registered in original registry"
+                    )
+
+                # Test all registry operations work
+                availability = self.original_registry.is_dsl_available(dsl_name)
+                self.assertIsInstance(availability, bool)
+
+                version = self.original_registry.get_dsl_version(dsl_name)
+                self.assertTrue(version is None or isinstance(version, Version))
+
+                # Verify DSL appears in appropriate lists
+                if availability:
+                    self.assertIn(
+                        dsl_name, self.original_registry.list_available_dsls()
+                    )
+                self.assertIn(dsl_name, self.original_registry.list_all_dsls())
+
+    def test_common_utils_wrappers(self):
+        """Test common_utils wrapper functions work correctly"""
+        from torch.testing._internal.common_utils import (
+            get_all_dsls,
+            get_available_dsls,
+        )
+
+        # Compare wrapper results with direct registry calls
+        if self.original_registry:
+            self.assertEqual(get_all_dsls(), self.original_registry.list_all_dsls())
+            self.assertEqual(
+                get_available_dsls(), self.original_registry.list_available_dsls()
+            )
+
+    def test_skip_decorators(self):
+        """Test DSL skip decorators work with registry"""
+        from torch.testing._internal.common_utils import (
+            skipIfDSLUnavailable,
+            skipUnlessDSLAvailable,
+        )
+
+        # Test decorators are callable
+        self.assertTrue(callable(skipIfDSLUnavailable))
+        self.assertTrue(callable(skipUnlessDSLAvailable))
+
+        # Test decorator creation works
+        decorator1 = skipIfDSLUnavailable("nonexistent_dsl")
+        decorator2 = skipUnlessDSLAvailable("triton")
+
+        self.assertTrue(callable(decorator1))
+        self.assertTrue(callable(decorator2))
+
+    # Phase 5: Test Infrastructure (2 methods)
+
+    def test_registry_isolation(self):
+        """Test registry state can be saved and restored"""
+        # Verify we start with clean test registry
+        self.assertEqual(self.test_registry.list_all_dsls(), [])
+
+        # Register test DSL
+        mock_dsl = self.create_valid_mock_dsl()
+        self.test_registry.register_dsl("isolation_test", mock_dsl)
+
+        # Verify registration
+        self.assertIn("isolation_test", self.test_registry.list_all_dsls())
+
+        # Verify original registry is unaffected
+        if self.original_registry:
+            original_dsls = self.original_registry.list_all_dsls()
+            self.assertNotIn("isolation_test", original_dsls)
+
+    def test_mock_dsl_helpers(self):
+        """Test mock DSL creation utilities work correctly"""
+        # Test create_valid_mock_dsl
+        valid_mock = self.create_valid_mock_dsl("test", available=True, version="2.1.0")
+
+        self.assertTrue(callable(valid_mock.runtime_available))
+        self.assertTrue(callable(valid_mock.runtime_version))
+        self.assertTrue(callable(valid_mock.deregister_op_overrides))
+        self.assertTrue(callable(valid_mock.register_op_override))
+
+        # Test behavior
+        self.assertTrue(valid_mock.runtime_available())
+        self.assertEqual(valid_mock.runtime_version(), Version("2.1.0"))
+
+        # Test create_broken_mock_dsl
+        broken_mock = self.create_broken_mock_dsl("runtime_available", ImportError)
+
+        with self.assertRaises(ImportError):
+            broken_mock.runtime_available()
+
+        # Other methods should still work
+        self.assertIsInstance(broken_mock.runtime_version(), Version)
+
+
+if __name__ == "__main__":
+    run_tests()
diff --git a/test/python_native/test_native_dsl_ops.py b/test/python_native/test_native_dsl_ops.py
new file mode 100644
index 0000000000000..e01a2b608cb33
--- /dev/null
+++ b/test/python_native/test_native_dsl_ops.py
@@ -0,0 +1,565 @@
+# Owner(s): ["module: dsl-native-ops"]
+
+import importlib.util
+import os
+import subprocess
+import sys
+import textwrap
+import uuid
+from unittest.mock import patch
+
+from torch.testing._internal.common_utils import (
+    instantiate_parametrized_tests,
+    parametrize,
+    run_tests,
+    TestCase,
+)
+
+
+def _subprocess_lastline(script, env=None):
+    """Run script in a fresh interpreter and return the last line of stdout."""
+    result = subprocess.check_output(
+        [sys.executable, "-c", script],
+        cwd=os.path.dirname(os.path.realpath(__file__)),
+        text=True,
+    ).strip()
+    return result.rsplit("\n", 1)[-1]
+
+
+def _import_module_directly(module_name, file_name):
+    """Import a module directly without triggering package imports."""
+    test_dir = os.path.dirname(os.path.abspath(__file__))
+    pytorch_root = os.path.dirname(os.path.dirname(test_dir))
+    module_path = os.path.join(pytorch_root, "torch", "_native", file_name)
+
+    spec = importlib.util.spec_from_file_location(module_name, module_path)
+    module = importlib.util.module_from_spec(spec)
+    sys.modules[module_name] = module
+    spec.loader.exec_module(module)
+    return module
+
+
+class TestNativeDSLOps(TestCase):
+    """Tests for the torch._native DSL ops framework."""
+
+    def setUp(self):
+        """Clear all caches before each test to ensure test isolation."""
+        self._cache_functions_to_clear = [
+            (
+                "torch._native.common_utils",
+                ["check_native_jit_disabled", "check_native_version_skip"],
+            ),
+            (
+                "torch._native.triton_utils",
+                [
+                    "_version_is_sufficient",
+                    "check_native_jit_disabled",
+                    "check_native_version_skip",
+                ],
+            ),
+            (
+                "torch._native.cutedsl_utils",
+                [
+                    "_version_is_ok",
+                    "check_native_jit_disabled",
+                    "check_native_version_skip",
+                ],
+            ),
+        ]
+        self._clear_function_caches()
+
+    def _clear_function_caches(self):
+        """Helper method to clear function caches with error handling."""
+        for module_name, function_names in self._cache_functions_to_clear:
+            try:
+                module = __import__(module_name, fromlist=function_names)
+                for func_name in function_names:
+                    if hasattr(module, func_name):
+                        getattr(module, func_name).cache_clear()
+            except (AttributeError, ImportError):
+                # Some functions might not exist or be cached, ignore errors
+                pass
+
+    def test_consistent_helper_interface(self):
+        """Test all registered DSL utils expose consistent public APIs."""
+        from torch.testing._internal.common_utils import get_all_dsls
+
+        # Automatically discover all registered DSLs
+        dsl_names = get_all_dsls()
+        if not dsl_names:
+            # Fallback to hardcoded list if registry not available
+            dsl_names = ["triton", "cutedsl"]
+
+        modules_info = [
+            (f"{dsl}_utils.py", f"torch._native.{dsl}_utils") for dsl in dsl_names
+        ]
+
+        # Import modules directly to avoid dependency issues
+        modules = {}
+        for file_name, module_name in modules_info:
+            modules[module_name] = _import_module_directly(module_name, file_name)
+
+        required_methods = {
+            "runtime_available",
+            "runtime_version",
+            "register_op_override",
+            "deregister_op_overrides",
+        }
+
+        # Test each module has required methods and they're callable
+        public_apis = {}
+        for module_name, mod in modules.items():
+            with self.subTest(module=module_name, test="required_methods"):
+                public = {name for name in dir(mod) if not name.startswith("_")}
+                public_apis[module_name] = public
+
+                self.assertTrue(
+                    required_methods <= public,
+                    f"{module_name} missing: {required_methods - public}",
+                )
+
+                for method_name in required_methods:
+                    with self.subTest(module=module_name, method=method_name):
+                        self.assertTrue(callable(getattr(mod, method_name)))
+
+        # Test modules expose identical public APIs
+        api_sets = list(public_apis.values())
+        if len(api_sets) > 1:
+            for i, api_set in enumerate(api_sets[1:], 1):
+                self.assertEqual(
+                    api_sets[0],
+                    api_set,
+                    f"Module {i} should have identical public API to module 0",
+                )
+
+        # Test runtime functions return expected types
+        for module_name, mod in modules.items():
+            with self.subTest(module=module_name, test="runtime_functions"):
+                # runtime_available should return bool
+                self.assertIsInstance(mod.runtime_available(), bool)
+
+                # runtime_version should return Version or None
+                ver = mod.runtime_version()
+                if ver is not None:
+                    from torch._vendor.packaging.version import Version
+
+                    self.assertIsInstance(ver, Version)
+
+    def test_no_dsl_imports_after_import_torch(self):
+        """import torch must not transitively import DSL runtimes.
+
+        Note: cuda.bindings may appear because importlib.util.find_spec on
+        nested modules (e.g. cuda.bindings.driver) imports parent packages
+        as a side-effect.  We check only the primary DSL runtimes here.
+        """
+        script = textwrap.dedent("""\
+            import sys
+            import torch
+            dsl_modules = ["triton", "cutlass", "tvm_ffi"]
+            leaked = [m for m in dsl_modules if m in sys.modules]
+            print(repr(leaked))
+        """)
+        result = _subprocess_lastline(script)
+        self.assertEqual(result, "[]", f"DSL modules leaked on import torch: {result}")
+
+    def test_no_external_packaging_dependency(self):
+        """torch._native must not import the external `packaging` package.
+
+        It should use the vendored copy at torch._vendor.packaging instead.
+        This guards against ModuleNotFoundError in environments where the
+        external `packaging` is not installed (e.g. torchvision Windows CI).
+        """
+        script = textwrap.dedent("""\
+            import sys
+            # Remove external packaging from sys.modules if already loaded
+            for mod_name in list(sys.modules):
+                if mod_name == "packaging" or mod_name.startswith("packaging."):
+                    del sys.modules[mod_name]
+            # Block external packaging from being imported
+            import importlib.abc
+            import importlib.machinery
+            class BlockPackaging(importlib.abc.MetaPathFinder):
+                def find_module(self, fullname, path=None):
+                    if fullname == "packaging" or fullname.startswith("packaging."):
+                        return self
+                def load_module(self, fullname):
+                    raise ImportError(f"External {fullname} is blocked")
+            sys.meta_path.insert(0, BlockPackaging())
+            import torch
+            print("OK")
+        """)
+        result = _subprocess_lastline(script)
+        self.assertEqual(result, "OK")
+
+    @parametrize("env_value, expected", [(None, False), ("1", True)])
+    def test_check_native_jit_disabled_environment_variable(self, env_value, expected):
+        """Test TORCH_DISABLE_NATIVE_JIT environment variable behavior."""
+        from torch._native.common_utils import check_native_jit_disabled
+
+        if env_value is None:
+            os.environ.pop("TORCH_DISABLE_NATIVE_JIT", None)
+        else:
+            os.environ["TORCH_DISABLE_NATIVE_JIT"] = env_value
+
+        try:
+            # Clear cache so function re-reads environment variable
+            check_native_jit_disabled.cache_clear()
+            self.assertEqual(check_native_jit_disabled(), expected)
+        finally:
+            # Clean up environment variable
+            os.environ.pop("TORCH_DISABLE_NATIVE_JIT", None)
+
+    def test_unavailable_reason_missing(self):
+        """Nonexistent package -> _unavailable_reason returns a string."""
+        common_utils = _import_module_directly(
+            "torch._native.common_utils", "common_utils.py"
+        )
+        reason = common_utils._unavailable_reason(
+            [("nonexistent_pkg_xyz", "nonexistent_pkg_xyz")]
+        )
+        self.assertIsNotNone(reason)
+        self.assertIn("nonexistent_pkg_xyz", reason)
+
+    def test_available_version_parsing(self):
+        """Test _available_version parses various version formats and handles invalid ones."""
+        from torch._vendor.packaging.version import Version
+
+        common_utils = _import_module_directly(
+            "torch._native.common_utils", "common_utils.py"
+        )
+
+        # Test with real package that has clean version
+        ver = common_utils._available_version("typing_extensions")
+        self.assertIsInstance(ver, Version)
+
+        # Test various version format scenarios
+        version_scenarios = [
+            ("0.7.0rc1", Version("0.7.0rc1"), "pre-release version"),
+            ("3.1.0.post1", Version("3.1.0.post1"), "post-release version"),
+            ("2.4.0a1", Version("2.4.0a1"), "alpha version"),
+            ("1.2.3", Version("1.2.3"), "standard version"),
+            ("abc", None, "invalid version string"),
+        ]
+
+        for version_str, expected_result, description in version_scenarios:
+            with self.subTest(version=version_str, scenario=description):
+                with patch("importlib.metadata.version", return_value=version_str):
+                    result = common_utils._available_version("fake_package")
+                    self.assertEqual(
+                        result,
+                        expected_result,
+                        f"_available_version({version_str!r}) = {result}",
+                    )
+
+    def test_registry_mechanics(self):
+        """_get_or_create_library caches Library instances per (lib, dispatch_key)."""
+        import torch._native.registry as registry
+        import torch.library
+
+        # Save original state for restoration
+        original_libs = dict(registry._libs)
+        original_filter_state = (
+            set(registry._filter_state._dsl_names),
+            set(registry._filter_state._op_symbols),
+            set(registry._filter_state._dispatch_keys),
+        )
+
+        try:
+            key = ("_test_native_dsl_registry", "CPU")
+            registry._libs.pop(key, None)
+
+            lib1 = registry._get_or_create_library(*key)
+            self.assertIsInstance(lib1, torch.library.Library)
+            lib2 = registry._get_or_create_library(*key)
+            self.assertIs(lib1, lib2, "should return cached instance")
+
+            # Different dispatch key -> different Library
+            key2 = ("_test_native_dsl_registry", "CUDA")
+            registry._libs.pop(key2, None)
+            lib3 = registry._get_or_create_library(*key2)
+            self.assertIsNot(lib1, lib3)
+
+            # cleanup
+            registry._libs.pop(key, None)
+            registry._libs.pop(key2, None)
+        finally:
+            # Restore original registry state
+            registry._libs.clear()
+            registry._libs.update(original_libs)
+
+            # Restore filter state
+            filter_state = registry._filter_state
+            filter_state._dsl_names.clear()
+            filter_state._op_symbols.clear()
+            filter_state._dispatch_keys.clear()
+            filter_state._dsl_names.update(original_filter_state[0])
+            filter_state._op_symbols.update(original_filter_state[1])
+            filter_state._dispatch_keys.update(original_filter_state[2])
+
+    def test_deregister_op_overrides_functionality(self):
+        """Test deregister_op_overrides methods exist, are callable, and work correctly."""
+        modules_to_test = [
+            ("triton_utils.py", "torch._native.triton_utils"),
+            ("cutedsl_utils.py", "torch._native.cutedsl_utils"),
+        ]
+
+        # Use the preserve_filter_state context manager pattern
+        from torch._native.registry import _filter_state
+
+        original_filter_state = (
+            set(_filter_state._dsl_names),
+            set(_filter_state._op_symbols),
+            set(_filter_state._dispatch_keys),
+        )
+
+        try:
+            for file_name, module_name in modules_to_test:
+                with self.subTest(module=module_name):
+                    mod = _import_module_directly(module_name, file_name)
+
+                    # Test method exists and is callable
+                    self.assertTrue(hasattr(mod, "deregister_op_overrides"))
+                    self.assertTrue(callable(mod.deregister_op_overrides))
+
+                    # Test method can be called without error (should be no-op when no overrides registered)
+                    try:
+                        mod.deregister_op_overrides()
+                    except Exception as e:
+                        self.fail(
+                            f"deregister_op_overrides on {module_name} raised exception: {e}"
+                        )
+        finally:
+            # Restore original filter state
+            _filter_state._dsl_names.clear()
+            _filter_state._op_symbols.clear()
+            _filter_state._dispatch_keys.clear()
+            _filter_state._dsl_names.update(original_filter_state[0])
+            _filter_state._op_symbols.update(original_filter_state[1])
+            _filter_state._dispatch_keys.update(original_filter_state[2])
+
+    def test_register_op_skips_when_jit_disabled(self):
+        """register_op_override does not call through when TORCH_DISABLE_NATIVE_JIT=1."""
+        from torch._native import cutedsl_utils, triton_utils
+
+        # Test the actual environment variable behavior to ensure it works
+        # Set TORCH_DISABLE_NATIVE_JIT=1 and clear caches
+        with patch.dict(os.environ, {"TORCH_DISABLE_NATIVE_JIT": "1"}):
+            # Import and clear caches for both modules
+            from torch._native.common_utils import check_native_jit_disabled
+
+            check_native_jit_disabled.cache_clear()
+
+            # Import functions from each module and clear their caches too
+            triton_utils.check_native_jit_disabled.cache_clear()
+            cutedsl_utils.check_native_jit_disabled.cache_clear()
+
+            # Verify the function returns True
+            self.assertTrue(check_native_jit_disabled())
+
+            # Mock the registry calls to count how many times they would be called
+            with patch("torch._native.registry.register_op_override") as registry_mock:
+                # Use a unique operation name
+                unique_op = f"test_jit_disabled_{uuid.uuid4().hex[:8]}.Tensor"
+                triton_utils.register_op_override(
+                    "aten", unique_op, "CPU", lambda: None
+                )
+                cutedsl_utils.register_op_override(
+                    "aten", unique_op, "CPU", lambda: None
+                )
+                # Should not call the registry function at all since JIT is disabled
+                self.assertEqual(registry_mock.call_count, 0)
+
+    def test_version_skip_env_var_overrides(self):
+        """TORCH_NATIVE_SKIP_VERSION_CHECK=1 allows non-blessed versions."""
+        from torch._vendor.packaging.version import Version
+
+        fake_version = Version("99.99.99")
+
+        # Set the environment variable and clear caches
+        with patch.dict(os.environ, {"TORCH_NATIVE_SKIP_VERSION_CHECK": "1"}):
+            # Import fresh modules to avoid cached state
+            from torch._native import cutedsl_utils, triton_utils
+            from torch._native.common_utils import check_native_version_skip
+
+            # Clear all relevant caches to ensure clean state
+            check_native_version_skip.cache_clear()
+
+            # Clear module-specific caches for the imported modules
+            for module in [triton_utils, cutedsl_utils]:
+                for attr_name in dir(module):
+                    attr = getattr(module, attr_name)
+                    if hasattr(attr, "cache_clear"):
+                        attr.cache_clear()
+
+            with (
+                patch.object(
+                    triton_utils,
+                    "_check_runtime_available",
+                    return_value=(True, fake_version),
+                ),
+                patch.object(
+                    cutedsl_utils,
+                    "_check_runtime_available",
+                    return_value=(True, fake_version),
+                ),
+                patch.object(triton_utils, "_register_op_override_impl") as triton_mock,
+                patch.object(cutedsl_utils, "_register_op_override_impl") as cute_mock,
+            ):
+                # Use unique operation names to avoid conflicts
+                op_name = f"test_version_skip_{uuid.uuid4().hex[:8]}.Tensor"
+
+                # Call the register functions
+                triton_utils.register_op_override("aten", op_name, "CPU", lambda: None)
+                cutedsl_utils.register_op_override("aten", op_name, "CPU", lambda: None)
+
+                # Verify both implementation functions were called
+                self.assertEqual(
+                    triton_mock.call_count + cute_mock.call_count,
+                    2,
+                    f"Expected 2 calls but got triton: {triton_mock.call_count}, cutedsl: {cute_mock.call_count}",
+                )
+
+    @parametrize("env_value, expected", [(None, False), ("1", True)])
+    def test_check_native_version_skip_environment_variable(self, env_value, expected):
+        """Test TORCH_NATIVE_SKIP_VERSION_CHECK environment variable behavior."""
+        from torch._native.common_utils import check_native_version_skip
+
+        if env_value is None:
+            os.environ.pop("TORCH_NATIVE_SKIP_VERSION_CHECK", None)
+        else:
+            os.environ["TORCH_NATIVE_SKIP_VERSION_CHECK"] = env_value
+
+        try:
+            # Clear cache so function re-reads environment variable
+            check_native_version_skip.cache_clear()
+            self.assertEqual(check_native_version_skip(), expected)
+        finally:
+            # Clean up environment variable
+            os.environ.pop("TORCH_NATIVE_SKIP_VERSION_CHECK", None)
+
+    def test_dsl_registry_functionality(self):
+        """Test that DSL registry works correctly"""
+        from torch.testing._internal.common_utils import (
+            get_all_dsls,
+            get_available_dsls,
+            is_dsl_available,
+        )
+
+        # Test registry returns expected DSLs
+        all_dsls = get_all_dsls()
+        self.assertIsInstance(all_dsls, list)
+        self.assertIn("triton", all_dsls)
+        self.assertIn("cutedsl", all_dsls)
+
+        # Test available DSLs are subset of all DSLs
+        available_dsls = get_available_dsls()
+        self.assertIsInstance(available_dsls, list)
+        for dsl in available_dsls:
+            self.assertIn(dsl, all_dsls)
+
+        # Test availability check function
+        for dsl in all_dsls:
+            availability = is_dsl_available(dsl)
+            self.assertIsInstance(availability, bool)
+            # If DSL is in available list, it should return True
+            if dsl in available_dsls:
+                self.assertTrue(availability)
+
+    def test_dsl_test_helpers(self):
+        """Test that DSL test helper decorators work"""
+        from torch.testing._internal.common_utils import (
+            skipIfDSLUnavailable,
+            skipIfNoCuteDSL,
+            skipIfNoTritonDSL,
+            skipUnlessDSLAvailable,
+        )
+
+        # Test that decorators are callable
+        self.assertTrue(callable(skipIfNoTritonDSL))
+        self.assertTrue(callable(skipIfNoCuteDSL))
+        self.assertTrue(callable(skipIfDSLUnavailable))
+        self.assertTrue(callable(skipUnlessDSLAvailable))
+
+        # Test dynamic decorators can be called
+        try:
+            decorator1 = skipIfDSLUnavailable("nonexistent_dsl")
+            decorator2 = skipUnlessDSLAvailable("triton")
+            self.assertTrue(callable(decorator1))
+            self.assertTrue(callable(decorator2))
+        except Exception as e:
+            self.fail(f"Dynamic DSL decorators failed: {e}")
+
+    def test_cache_invalidation_after_re_registration(self):
+        """Test that caches are properly invalidated when DSLs are re-registered"""
+        from unittest.mock import Mock
+
+        from torch._native.dsl_registry import DSLRegistry
+
+        # Create a fresh registry for this test
+        registry = DSLRegistry()
+
+        # Create mock DSL modules
+        mock_dsl_1 = Mock()
+        mock_dsl_1.runtime_available.return_value = False  # Initially unavailable
+        mock_dsl_1.runtime_version.return_value = None
+
+        mock_dsl_2 = Mock()
+        mock_dsl_2.runtime_available.return_value = True  # Available
+        mock_dsl_2.runtime_version.return_value = None
+
+        # Register first DSL and cache results
+        registry.register_dsl("test_cache_dsl", mock_dsl_1)
+        initial_available = registry.is_dsl_available("test_cache_dsl")
+        initial_list = registry.list_available_dsls()
+
+        self.assertFalse(initial_available)
+        self.assertNotIn("test_cache_dsl", initial_list)
+
+        # Re-register with different module that is available
+        registry.register_dsl("test_cache_dsl", mock_dsl_2)
+
+        # Verify cache was invalidated and new results are returned
+        new_available = registry.is_dsl_available("test_cache_dsl")
+        new_list = registry.list_available_dsls()
+
+        self.assertTrue(
+            new_available, "Cache should be invalidated and return new result"
+        )
+        self.assertIn(
+            "test_cache_dsl",
+            new_list,
+            "Available DSLs list should reflect new registration",
+        )
+
+    def test_incomplete_protocol_implementation(self):
+        """Test that registration fails when module doesn't implement required protocol methods"""
+        from torch._native.dsl_registry import DSLRegistry
+
+        # Create a fresh registry for this test
+        registry = DSLRegistry()
+
+        # Create an object missing required protocol methods (not using Mock)
+        class IncompleteModule:
+            def runtime_available(self):
+                return True
+
+            # Missing: runtime_version, register_op_override, deregister_op_overrides
+
+        incomplete_module = IncompleteModule()
+
+        # Attempt to register should raise TypeError due to missing methods
+        with self.assertRaises(TypeError) as cm:
+            registry.register_dsl("incomplete_dsl", incomplete_module)
+
+        self.assertIn("missing required methods", str(cm.exception))
+        self.assertIn("runtime_version", str(cm.exception))
+        self.assertIn("register_op_override", str(cm.exception))
+
+        # Verify DSL was not registered
+        self.assertNotIn("incomplete_dsl", registry.list_all_dsls())
+
+
+instantiate_parametrized_tests(TestNativeDSLOps)
+
+if __name__ == "__main__":
+    run_tests()
diff --git a/test/python_native/test_registry.py b/test/python_native/test_registry.py
new file mode 100644
index 0000000000000..f58ea56ad3c9c
--- /dev/null
+++ b/test/python_native/test_registry.py
@@ -0,0 +1,258 @@
+# Owner(s): ["module: dsl-native-ops"]
+
+from unittest.mock import MagicMock, patch
+
+import torch._native.registry as registry_module
+from torch.testing._internal.common_utils import run_tests, skipIfTorchDynamo, TestCase
+
+
+@skipIfTorchDynamo("Registry tests don't need dynamo compilation")
+class TestRegistry(TestCase):
+    """Tests for the torch._native.registry module."""
+
+    def setUp(self):
+        """Clean up registry state before each test."""
+        self.registry = registry_module
+
+        # Store original state for restoration
+        self._original_libs = dict(self.registry._libs)
+        self._original_graphs = dict(self.registry._graphs)
+        self._original_dsl_name_to_lib_graph = {
+            k: list(v) for k, v in self.registry._dsl_name_to_lib_graph.items()
+        }
+        self._original_dispatch_key_to_lib_graph = {
+            k: list(v) for k, v in self.registry._dispatch_key_to_lib_graph.items()
+        }
+        self._original_op_symbol_to_lib_graph = {
+            k: list(v) for k, v in self.registry._op_symbol_to_lib_graph.items()
+        }
+
+        # Store original filter state
+        self._original_filter_state = (
+            set(self.registry._filter_state._dsl_names),
+            set(self.registry._filter_state._op_symbols),
+            set(self.registry._filter_state._dispatch_keys),
+        )
+
+        # Clear global state
+        self.registry._libs.clear()
+        self.registry._graphs.clear()
+        self.registry._dsl_name_to_lib_graph.clear()
+        self.registry._dispatch_key_to_lib_graph.clear()
+        self.registry._op_symbol_to_lib_graph.clear()
+
+        # Clear filter state to ensure clean start
+        self.registry._filter_state._dsl_names.clear()
+        self.registry._filter_state._op_symbols.clear()
+        self.registry._filter_state._dispatch_keys.clear()
+
+    def tearDown(self):
+        """Restore original registry state after each test."""
+        if hasattr(self, "registry"):
+            # Restore original state
+            self.registry._libs.clear()
+            self.registry._libs.update(self._original_libs)
+
+            self.registry._graphs.clear()
+            self.registry._graphs.update(self._original_graphs)
+
+            # Properly restore mapping dictionaries with new list instances
+            self.registry._dsl_name_to_lib_graph.clear()
+            for k, v in self._original_dsl_name_to_lib_graph.items():
+                self.registry._dsl_name_to_lib_graph[k] = list(v)
+
+            self.registry._dispatch_key_to_lib_graph.clear()
+            for k, v in self._original_dispatch_key_to_lib_graph.items():
+                self.registry._dispatch_key_to_lib_graph[k] = list(v)
+
+            self.registry._op_symbol_to_lib_graph.clear()
+            for k, v in self._original_op_symbol_to_lib_graph.items():
+                self.registry._op_symbol_to_lib_graph[k] = list(v)
+
+            # Restore filter state
+            self.registry._filter_state._dsl_names.clear()
+            self.registry._filter_state._op_symbols.clear()
+            self.registry._filter_state._dispatch_keys.clear()
+            self.registry._filter_state._dsl_names.update(
+                self._original_filter_state[0]
+            )
+            self.registry._filter_state._op_symbols.update(
+                self._original_filter_state[1]
+            )
+            self.registry._filter_state._dispatch_keys.update(
+                self._original_filter_state[2]
+            )
+
+    # Keep essential existing tests
+    def test_override_node_dataclass(self):
+        """Test _OverrideNode dataclass creation and defaults."""
+
+        def test_fn(x):
+            return x
+
+        node = self.registry._OverrideNode("test_dsl", "add.Tensor", "CPU", test_fn)
+        self.assertEqual(node.dsl_name, "test_dsl")
+        self.assertEqual(node.op_symbol, "add.Tensor")
+        self.assertEqual(node.dispatch_key, "CPU")
+        self.assertEqual(node.override_fn, test_fn)
+        self.assertFalse(node.unconditional_override)
+        self.assertTrue(node.active)
+
+    @patch("torch.library.Library")
+    def test_register_op_override_basic(self, mock_library_cls):
+        """Test basic register_op_override functionality."""
+
+        def impl_fn(x):
+            return x
+
+        mock_lib = MagicMock()
+        mock_library_cls.return_value = mock_lib
+
+        self.registry.register_op_override(
+            "test_backend", "aten", "add.Tensor", "CPU", impl_fn
+        )
+
+        key = ("add.Tensor", "CPU")
+        self.assertEqual(len(self.registry._graphs[key]), 1)
+        node = self.registry._graphs[key][0]
+        self.assertEqual(node.dsl_name, "test_backend")
+        self.assertEqual(node.override_fn, impl_fn)
+
+    @patch("torch.library.Library")
+    def test_deregister_op_overrides_basic(self, mock_library_cls):
+        """Test basic deregister_op_overrides functionality."""
+
+        def impl_fn(x):
+            return x
+
+        mock_lib = MagicMock()
+        mock_library_cls.return_value = mock_lib
+
+        # Register first
+        self.registry.register_op_override(
+            "test_backend", "aten", "mul.Tensor", "CPU", impl_fn
+        )
+
+        key = ("mul.Tensor", "CPU")
+        self.assertTrue(self.registry._graphs[key][0].active)
+
+        # Then deregister
+        self.registry.deregister_op_overrides(disable_dsl_names="test_backend")
+        self.assertFalse(self.registry._graphs[key][0].active)
+
+    # NEW FUNCTIONALITY TESTS - ONLY THE ESSENTIAL ONES
+
+    def test_reorder_graphs_from_user_function_basic(self):
+        """Test basic graph reordering functionality."""
+        # Set up test data
+        key = ("test_reorder.Tensor", "CPU")
+
+        def impl_fn(x):
+            return x
+
+        # Create nodes in specific order
+        nodes = [
+            self.registry._OverrideNode("dsl_c", "test_reorder.Tensor", "CPU", impl_fn),
+            self.registry._OverrideNode("dsl_a", "test_reorder.Tensor", "CPU", impl_fn),
+            self.registry._OverrideNode("dsl_b", "test_reorder.Tensor", "CPU", impl_fn),
+        ]
+        self.registry._graphs[key] = nodes
+
+        # Define alphabetical ordering function
+        def alphabetical_order(op_symbol, dispatch_key, graph):
+            return sorted(graph, key=lambda n: n.dsl_name)
+
+        # Apply reordering
+        self.registry.reorder_graphs_from_user_function(alphabetical_order)
+
+        # Verify alphabetical order
+        reordered_graph = self.registry._graphs[key]
+        actual_names = [node.dsl_name for node in reordered_graph]
+        self.assertEqual(actual_names, ["dsl_a", "dsl_b", "dsl_c"])
+
+    def test_reorder_graphs_from_user_function_error_handling(self):
+        """Test error handling in graph reordering."""
+        # Set up test data
+        key = ("test_error.Tensor", "CPU")
+
+        def impl_fn(x):
+            return x
+
+        node = self.registry._OverrideNode(
+            "test_dsl", "test_error.Tensor", "CPU", impl_fn
+        )
+        original_graph = [node]
+        self.registry._graphs[key] = original_graph.copy()
+
+        # Define failing ordering function
+        def failing_order_fn(op_symbol, dispatch_key, graph):
+            raise ValueError("Test exception")
+
+        # Should handle the exception gracefully
+        with self.assertLogs("torch._native.registry", level="WARNING") as log:
+            self.registry.reorder_graphs_from_user_function(failing_order_fn)
+
+        # Verify warning was logged and original graph preserved
+        self.assertEqual(len(log.records), 1)
+        self.assertIn("Graph transformation failed", log.records[0].getMessage())
+        self.assertEqual(self.registry._graphs[key], original_graph)
+
+    def test_get_user_ordering_fn_env_var_not_set(self):
+        """Test behavior when environment variable is not set."""
+        with patch.dict("os.environ", {}, clear=True):
+            from torch._native import get_user_ordering_fn
+
+            get_user_ordering_fn.cache_clear()
+            result = get_user_ordering_fn()
+            self.assertIsNone(result)
+
+    def test_get_user_ordering_fn_invalid_path(self):
+        """Test handling of invalid environment variable paths."""
+        with patch.dict(
+            "os.environ",
+            {"TORCH_PYTHON_NATIVE_USER_GRAPH_ORDER_FN": "nonexistent.module.function"},
+        ):
+            from torch._native import get_user_ordering_fn
+
+            get_user_ordering_fn.cache_clear()
+
+            with self.assertRaises(ValueError) as cm:
+                get_user_ordering_fn()
+            self.assertIn("Could not resolve", str(cm.exception))
+
+    def test_integration_reorder_and_register(self):
+        """Integration test: reorder then register functionality."""
+
+        def impl_fn1(x):
+            return x + 1
+
+        def impl_fn2(x):
+            return x + 2
+
+        # Register multiple overrides
+        self.registry.register_op_override(
+            "backend_z", "aten", "test.Tensor", "CPU", impl_fn1
+        )
+        self.registry.register_op_override(
+            "backend_a", "aten", "test.Tensor", "CPU", impl_fn2
+        )
+
+        key = ("test.Tensor", "CPU")
+
+        # Verify initial order
+        initial_names = [node.dsl_name for node in self.registry._graphs[key]]
+        self.assertEqual(initial_names, ["backend_z", "backend_a"])
+
+        # Reorder alphabetically
+        def alphabetical_order(op_symbol, dispatch_key, graph):
+            return sorted(graph, key=lambda n: n.dsl_name)
+
+        self.registry.reorder_graphs_from_user_function(alphabetical_order)
+
+        # Verify reordered
+        final_names = [node.dsl_name for node in self.registry._graphs[key]]
+        self.assertEqual(final_names, ["backend_a", "backend_z"])
+
+
+if __name__ == "__main__":
+    run_tests()
diff --git a/test/python_native/test_torch_backends.py b/test/python_native/test_torch_backends.py
new file mode 100644
index 0000000000000..e3e8f1e46a246
--- /dev/null
+++ b/test/python_native/test_torch_backends.py
@@ -0,0 +1,544 @@
+# Owner(s): ["module: dsl-native-ops"]
+
+import torch.backends.python_native as pn
+from torch.testing._internal.common_utils import run_tests, skipIfTorchDynamo, TestCase
+
+
+class RegistryTestMixin:
+    """Mixin for tests that need to preserve registry state."""
+
+    def setUp(self):
+        """Set up test state with registry preservation."""
+        super().setUp() if hasattr(super(), "setUp") else None
+
+        # Use the new _preserve_filter_state but setup manually for tearDown
+        filter_state = pn._get_filter_state()
+        self._original_filter_state = (
+            set(filter_state._dsl_names),
+            set(filter_state._op_symbols),
+            set(filter_state._dispatch_keys),
+        )
+
+        # Clear filter state for clean test start
+        filter_state._dsl_names.clear()
+        filter_state._op_symbols.clear()
+        filter_state._dispatch_keys.clear()
+
+        # Ensure all DSLs start enabled
+        try:
+            for dsl_name in pn.all_dsls:
+                dsl = getattr(pn, dsl_name)
+                dsl.enable()
+        except Exception:
+            pass
+
+    def tearDown(self):
+        """Restore original filter state."""
+        try:
+            if hasattr(self, "_original_filter_state"):
+                filter_state = pn._get_filter_state()
+
+                # Restore original filter state
+                filter_state._dsl_names.clear()
+                filter_state._op_symbols.clear()
+                filter_state._dispatch_keys.clear()
+
+                filter_state._dsl_names.update(self._original_filter_state[0])
+                filter_state._op_symbols.update(self._original_filter_state[1])
+                filter_state._dispatch_keys.update(self._original_filter_state[2])
+        except Exception:
+            pass
+        super().tearDown() if hasattr(super(), "tearDown") else None
+
+
+@skipIfTorchDynamo("Backend tests don't need dynamo compilation")
+class TestTorchBackendsPythonNative(RegistryTestMixin, TestCase):
+    """Tests for torch.backends.python_native user-facing API."""
+
+    def test_module_import(self):
+        """Test that torch.backends.python_native imports successfully."""
+        # Should not raise any exceptions
+        import torch.backends.python_native as pn_import
+
+        self.assertIsNotNone(pn_import)
+
+    def test_dsl_discovery(self):
+        """Test DSL discovery functionality."""
+        all_dsls = pn.all_dsls
+        available_dsls = pn.available_dsls
+
+        # Should return lists
+        self.assertIsInstance(all_dsls, list)
+        self.assertIsInstance(available_dsls, list)
+
+        # Available DSLs should be subset of all DSLs
+        self.assertTrue(set(available_dsls).issubset(set(all_dsls)))
+
+    def test_dsl_access(self):
+        """Test accessing DSL controllers."""
+        all_dsls = pn.all_dsls
+
+        for dsl_name in all_dsls:
+            with self.subTest(dsl_name=dsl_name):
+                # Should be able to access DSL controller
+                dsl = getattr(pn, dsl_name)
+                self.assertIsNotNone(dsl)
+
+                # Should have required attributes
+                self.assertEqual(dsl.name, dsl_name)
+                self.assertIsInstance(dsl.available, bool)
+                self.assertIsInstance(dsl.enabled, bool)
+
+    def test_dsl_properties(self):
+        """Test DSL controller properties."""
+        all_dsls = pn.all_dsls
+
+        for dsl_name in all_dsls:
+            with self.subTest(dsl_name=dsl_name):
+                dsl = getattr(pn, dsl_name)
+
+                # Test name property
+                self.assertEqual(dsl.name, dsl_name)
+
+                # Test available property (should not raise)
+                available = dsl.available
+                self.assertIsInstance(available, bool)
+
+                # Test version property (should not raise)
+                version = dsl.version
+                # Version can be None if DSL not available
+                self.assertTrue(version is None or hasattr(version, "major"))
+
+                # Test enabled property
+                enabled = dsl.enabled
+                self.assertIsInstance(enabled, bool)
+
+    def test_dsl_enable_disable(self):
+        """Test DSL enable/disable functionality."""
+        all_dsls = pn.all_dsls
+
+        for dsl_name in all_dsls:
+            with self.subTest(dsl_name=dsl_name):
+                dsl = getattr(pn, dsl_name)
+                original_state = dsl.enabled
+
+                try:
+                    # Test method-based disable/enable
+                    dsl.disable()
+                    self.assertEqual(dsl.enabled, False)
+
+                    dsl.enable()
+                    self.assertEqual(dsl.enabled, True)
+
+                    # Test property-based disable/enable
+                    dsl.disable()  # Use method to avoid flags_frozen issue
+                    self.assertEqual(dsl.enabled, False)
+
+                    dsl.enable()  # Use method to avoid flags_frozen issue
+                    self.assertEqual(dsl.enabled, True)
+
+                finally:
+                    # Restore original state
+                    if original_state:
+                        dsl.enable()
+                    else:
+                        dsl.disable()
+
+    def test_dsl_context_managers(self):
+        """Test DSL context manager functionality."""
+        all_dsls = pn.all_dsls
+
+        for dsl_name in all_dsls:
+            with self.subTest(dsl_name=dsl_name):
+                dsl = getattr(pn, dsl_name)
+                original_state = dsl.enabled
+
+                try:
+                    # Ensure DSL starts enabled
+                    dsl.enable()
+
+                    # Test disabled context manager
+                    with dsl.disabled():
+                        self.assertEqual(dsl.enabled, False)
+
+                    # Should be restored after context
+                    self.assertEqual(dsl.enabled, True)
+
+                finally:
+                    # Restore original state
+                    if original_state:
+                        dsl.enable()
+                    else:
+                        dsl.disable()
+
+    def test_nested_context_managers(self):
+        """Test nested DSL context managers."""
+        all_dsls = pn.all_dsls
+
+        if len(all_dsls) >= 2:
+            dsl1_name, dsl2_name = all_dsls[0], all_dsls[1]
+            dsl1 = getattr(pn, dsl1_name)
+            dsl2 = getattr(pn, dsl2_name)
+
+            original_state1 = dsl1.enabled
+            original_state2 = dsl2.enabled
+
+            try:
+                # Ensure both start enabled
+                dsl1.enable()
+                dsl2.enable()
+
+                with dsl1.disabled():
+                    self.assertEqual(dsl1.enabled, False)
+                    self.assertEqual(dsl2.enabled, True)
+
+                    with dsl2.disabled():
+                        self.assertEqual(dsl1.enabled, False)
+                        self.assertEqual(dsl2.enabled, False)
+
+                    # dsl2 should be restored
+                    self.assertEqual(dsl1.enabled, False)
+                    self.assertEqual(dsl2.enabled, True)
+
+                # Both should be restored
+                self.assertEqual(dsl1.enabled, True)
+                self.assertEqual(dsl2.enabled, True)
+
+            finally:
+                # Restore original states
+                if original_state1:
+                    dsl1.enable()
+                else:
+                    dsl1.disable()
+                if original_state2:
+                    dsl2.enable()
+                else:
+                    dsl2.disable()
+
+    def test_operation_discovery(self):
+        """Test operation discovery functionality."""
+        all_dsls = pn.all_dsls
+
+        for dsl_name in all_dsls:
+            with self.subTest(dsl_name=dsl_name):
+                operations = pn.get_dsl_operations(dsl_name)
+
+                # Should return a list
+                self.assertIsInstance(operations, list)
+
+                # Each operation should be a string
+                for op in operations:
+                    self.assertIsInstance(op, str)
+                    self.assertTrue(len(op) > 0)
+
+    def test_operation_control(self):
+        """Test operation-level control functionality."""
+        # Get an operation to test with
+        all_dsls = pn.all_dsls
+        test_operation = None
+
+        for dsl_name in all_dsls:
+            operations = pn.get_dsl_operations(dsl_name)
+            if operations:
+                test_operation = operations[0]
+                break
+
+        if test_operation:
+            # Test operation disable/enable (should not raise)
+            pn.disable_operations(test_operation)
+            pn.enable_operations(test_operation)
+
+            # Test multiple operations
+            pn.disable_operations(test_operation, "nonexistent_op")
+            pn.enable_operations(test_operation, "nonexistent_op")
+
+    def test_operation_context_manager(self):
+        """Test operation-level context manager."""
+        # Get an operation to test with
+        all_dsls = pn.all_dsls
+        test_operation = None
+
+        for dsl_name in all_dsls:
+            operations = pn.get_dsl_operations(dsl_name)
+            if operations:
+                test_operation = operations[0]
+                break
+
+        if test_operation:
+            # Should not raise exceptions
+            with pn.operations_disabled(test_operation):
+                pass  # Operation should be disabled in this context
+            # Operation should be re-enabled after context
+
+    def test_dispatch_key_control(self):
+        """Test dispatch key control functionality."""
+        # Test basic dispatch key control (should not raise)
+        pn.disable_dispatch_keys("CUDA", "CPU")
+        pn.enable_dispatch_keys("CUDA", "CPU")
+
+    def test_invalid_dsl_access(self):
+        """Test accessing invalid DSL names."""
+        with self.assertRaises(AttributeError):
+            _ = pn.nonexistent_dsl
+
+    def test_dsl_repr(self):
+        """Test DSL controller string representation."""
+        all_dsls = pn.all_dsls
+
+        for dsl_name in all_dsls:
+            with self.subTest(dsl_name=dsl_name):
+                dsl = getattr(pn, dsl_name)
+                repr_str = repr(dsl)
+
+                # Should contain DSL name
+                self.assertIn(dsl_name, repr_str)
+                # Should contain status info
+                self.assertTrue(
+                    any(status in repr_str for status in ["available", "unavailable"])
+                )
+                self.assertTrue(
+                    any(status in repr_str for status in ["enabled", "disabled"])
+                )
+
+    def test_module_dir(self):
+        """Test module __dir__ functionality."""
+        attrs = dir(pn)
+
+        # Should contain core attributes
+        expected_attrs = [
+            "available_dsls",
+            "all_dsls",
+            "get_dsl_operations",
+            "disable_operations",
+            "enable_operations",
+            "disable_dispatch_keys",
+            "enable_dispatch_keys",
+            "operations_disabled",
+        ]
+
+        for attr in expected_attrs:
+            self.assertIn(attr, attrs)
+
+        # Should contain DSL names
+        for dsl_name in pn.all_dsls:
+            self.assertIn(dsl_name, attrs)
+
+    def test_caching_functionality(self):
+        """Test caching integration and API contracts (not functools.lru_cache internals)."""
+        all_dsls = pn.all_dsls
+
+        if all_dsls:
+            dsl_name = all_dsls[0]
+
+            # Test operation caching - results should be identical
+            ops1 = pn.get_dsl_operations(dsl_name)
+            ops2 = pn.get_dsl_operations(dsl_name)
+            self.assertEqual(ops1, ops2)
+
+            # Test DSL controller caching - same object returned (API contract)
+            controller1 = getattr(pn, dsl_name)
+            controller2 = getattr(pn, dsl_name)
+            self.assertIs(controller1, controller2)
+
+            # Verify controller still works correctly
+            self.assertEqual(controller1.name, dsl_name)
+            self.assertIsInstance(controller1.enabled, bool)
+
+    def test_error_handling(self):
+        """Test error handling with invalid inputs."""
+        # Test with empty operation names - should not raise
+        pn.disable_operations("")
+        pn.enable_operations("")
+
+        # Test with invalid dispatch keys - should not raise
+        pn.disable_dispatch_keys("")
+        pn.enable_dispatch_keys("")
+
+        # Test invalid DSL access raises AttributeError
+        with self.assertRaises(AttributeError):
+            _ = pn.invalid_dsl_name
+
+        # Test get_dsl_operations with invalid DSL - should return empty list or raise ValueError
+        invalid_ops = pn.get_dsl_operations("invalid_dsl")
+        self.assertIsInstance(invalid_ops, list)
+        self.assertEqual(len(invalid_ops), 0)
+
+    def test_dispatch_behavior_verification(self):
+        """Test that DSL disable/enable actually affects operation dispatch."""
+        all_dsls = pn.all_dsls
+
+        for dsl_name in all_dsls:
+            with self.subTest(dsl_name=dsl_name):
+                dsl = getattr(pn, dsl_name)
+
+                if not dsl.available:
+                    continue  # Skip unavailable DSLs
+
+                operations = pn.get_dsl_operations(dsl_name)
+                if not operations:
+                    continue  # Skip DSLs with no operations
+
+                original_state = dsl.enabled
+
+                try:
+                    # Enable DSL first
+                    dsl.enable()
+                    self.assertTrue(
+                        dsl.enabled, f"{dsl_name} should be enabled after enable()"
+                    )
+
+                    # Disable DSL
+                    dsl.disable()
+                    self.assertFalse(
+                        dsl.enabled, f"{dsl_name} should be disabled after disable()"
+                    )
+
+                    # Verify operations are actually deregistered by checking registry state
+                    self.assertTrue(
+                        pn.is_dsl_disabled(dsl_name),
+                        f"{dsl_name} should be disabled in registry",
+                    )
+
+                    # Re-enable and verify registry state
+                    dsl.enable()
+                    self.assertTrue(
+                        dsl.enabled, f"{dsl_name} should be enabled after re-enable()"
+                    )
+                    self.assertFalse(
+                        pn.is_dsl_disabled(dsl_name),
+                        f"{dsl_name} should not be disabled in registry after re-enable",
+                    )
+
+                finally:
+                    # Restore original state
+                    if original_state:
+                        dsl.enable()
+                    else:
+                        dsl.disable()
+
+    def test_context_manager_dispatch_behavior(self):
+        """Test that context managers actually affect dispatch, not just state tracking."""
+        all_dsls = pn.all_dsls
+
+        for dsl_name in all_dsls:
+            with self.subTest(dsl_name=dsl_name):
+                dsl = getattr(pn, dsl_name)
+
+                if not dsl.available:
+                    continue  # Skip unavailable DSLs
+
+                original_state = dsl.enabled
+
+                try:
+                    # Start with DSL enabled
+                    dsl.enable()
+                    self.assertFalse(pn.is_dsl_disabled(dsl_name))
+
+                    # Use disabled context manager
+                    with dsl.disabled():
+                        # Verify DSL is actually disabled in registry during context
+                        self.assertTrue(
+                            pn.is_dsl_disabled(dsl_name),
+                            f"{dsl_name} should be disabled in registry during context",
+                        )
+
+                    # Verify DSL is re-enabled after context
+                    self.assertFalse(
+                        pn.is_dsl_disabled(dsl_name),
+                        f"{dsl_name} should be re-enabled in registry after context",
+                    )
+
+                finally:
+                    # Restore original state
+                    if original_state:
+                        dsl.enable()
+                    else:
+                        dsl.disable()
+
+    def test_operation_dispatch_behavior(self):
+        """Test that operation-level disable/enable affects registry state."""
+        all_dsls = pn.all_dsls
+        test_operation = None
+
+        # Find a real operation to test with
+        for dsl_name in all_dsls:
+            operations = pn.get_dsl_operations(dsl_name)
+            if operations:
+                test_operation = operations[0]
+                break
+
+        if test_operation:
+            # Verify operation starts enabled
+            initial_disabled = pn.is_operation_disabled(test_operation)
+
+            try:
+                # Disable operation and verify registry state
+                pn.disable_operations(test_operation)
+                self.assertTrue(
+                    pn.is_operation_disabled(test_operation),
+                    f"Operation {test_operation} should be disabled in registry",
+                )
+
+                # Re-enable operation and verify registry state
+                pn.enable_operations(test_operation)
+                # Note: enable_operations removes from disabled list
+                self.assertFalse(
+                    pn.is_operation_disabled(test_operation),
+                    f"Operation {test_operation} should not be disabled in registry after re-enable",
+                )
+
+            finally:
+                # Restore initial state
+                if initial_disabled:
+                    pn.disable_operations(test_operation)
+                else:
+                    pn.enable_operations(test_operation)
+
+
+class TestTorchBackendsPythonNativeIntegration(TestCase):
+    """Integration tests for torch.backends.python_native with actual DSLs."""
+
+    def test_real_dsl_integration(self):
+        """Test integration with real DSL modules if available."""
+        all_dsls = pn.all_dsls
+
+        for dsl_name in all_dsls:
+            with self.subTest(dsl_name=dsl_name):
+                dsl = getattr(pn, dsl_name)
+
+                if dsl.available:
+                    # Test that we can actually disable/enable real DSLs
+                    original_state = dsl.enabled
+
+                    try:
+                        # This should call actual DSL deregister functions
+                        dsl.disable()
+
+                        # This should call actual registry re-enable functions
+                        dsl.enable()
+
+                    finally:
+                        # Restore state
+                        if original_state:
+                            dsl.enable()
+                        else:
+                            dsl.disable()
+
+    def test_operations_with_real_registry(self):
+        """Test operation discovery with real registry."""
+        all_dsls = pn.all_dsls
+
+        for dsl_name in all_dsls:
+            with self.subTest(dsl_name=dsl_name):
+                operations = pn.get_dsl_operations(dsl_name)
+
+                # If DSL has operations, they should be valid
+                if operations:
+                    # Operations should be strings
+                    for op in operations:
+                        self.assertIsInstance(op, str)
+                        self.assertGreater(len(op), 0)
+
+
+if __name__ == "__main__":
+    run_tests()
diff --git a/test/quantization/ao_migration/common.py b/test/quantization/ao_migration/common.py
index acfc8065de846..35e49e38d1e87 100644
--- a/test/quantization/ao_migration/common.py
+++ b/test/quantization/ao_migration/common.py
@@ -1,5 +1,4 @@
 import importlib
-from typing import Optional
 
 from torch.testing._internal.common_utils import TestCase
 
@@ -9,8 +8,8 @@ def _test_function_import(
         self,
         package_name: str,
         function_list: list[str],
-        base: Optional[str] = None,
-        new_package_name: Optional[str] = None,
+        base: str | None = None,
+        new_package_name: str | None = None,
     ):
         r"""Tests individual function list import by comparing the functions
         and their hashes."""
@@ -25,14 +24,16 @@ def _test_function_import(
         for fn_name in function_list:
             old_function = getattr(old_location, fn_name)
             new_function = getattr(new_location, fn_name)
-            assert old_function == new_function, f"Functions don't match: {fn_name}"
-            assert hash(old_function) == hash(new_function), (
-                f"Hashes don't match: {old_function}({hash(old_function)}) vs. "
-                f"{new_function}({hash(new_function)})"
-            )
+            if old_function != new_function:
+                raise AssertionError(f"Functions don't match: {fn_name}")
+            if hash(old_function) != hash(new_function):
+                raise AssertionError(
+                    f"Hashes don't match: {old_function}({hash(old_function)}) vs. "
+                    f"{new_function}({hash(new_function)})"
+                )
 
     def _test_dict_import(
-        self, package_name: str, dict_list: list[str], base: Optional[str] = None
+        self, package_name: str, dict_list: list[str], base: str | None = None
     ):
         r"""Tests individual function list import by comparing the functions
         and their hashes."""
@@ -45,8 +46,10 @@ def _test_dict_import(
         for dict_name in dict_list:
             old_dict = getattr(old_location, dict_name)
             new_dict = getattr(new_location, dict_name)
-            assert old_dict == new_dict, f"Dicts don't match: {dict_name}"
+            if old_dict != new_dict:
+                raise AssertionError(f"Dicts don't match: {dict_name}")
             for key in new_dict:
-                assert old_dict[key] == new_dict[key], (
-                    f"Dicts don't match: {dict_name} for key {key}"
-                )
+                if old_dict[key] != new_dict[key]:
+                    raise AssertionError(
+                        f"Dicts don't match: {dict_name} for key {key}"
+                    )
diff --git a/test/quantization/core/experimental/test_bits.py b/test/quantization/core/experimental/test_bits.py
index e621ac042b7b3..b16546f8e6cb2 100644
--- a/test/quantization/core/experimental/test_bits.py
+++ b/test/quantization/core/experimental/test_bits.py
@@ -11,7 +11,8 @@
 
 class Int16Tensor(torch.Tensor):
     def __new__(cls, elem):
-        assert elem.dtype == torch.bits16
+        if elem.dtype != torch.bits16:
+            raise AssertionError(f"Expected dtype torch.bits16, got {elem.dtype}")
         return torch.Tensor._make_subclass(cls, elem, elem.requires_grad)
 
     def __init__(self, elem):
diff --git a/test/quantization/core/experimental/test_floatx.py b/test/quantization/core/experimental/test_floatx.py
index d234d857e84a1..75b542a78d095 100644
--- a/test/quantization/core/experimental/test_floatx.py
+++ b/test/quantization/core/experimental/test_floatx.py
@@ -330,7 +330,10 @@ def compare_binary_with_decimal(binary, decimal, number_name, dtype, device):
             tensor_int = torch.tensor([bits_int], dtype=torch.uint8, device=device)
             tensor_fp8 = tensor_int.view(dtype)
             if number_name == "nan":
-                assert tensor_fp8.isnan()
+                if not tensor_fp8.isnan():
+                    raise AssertionError(
+                        f"Expected NaN for {number_name}, got {tensor_fp8}"
+                    )
             else:
                 tensor_fp32 = tensor_fp8.float()
                 ref_tensor_fp32 = torch.tensor(
diff --git a/test/quantization/core/test_quantized_module.py b/test/quantization/core/test_quantized_module.py
index 1bdbdfe304fa6..6805cf2394852 100644
--- a/test/quantization/core/test_quantized_module.py
+++ b/test/quantization/core/test_quantized_module.py
@@ -157,7 +157,10 @@ def _test_linear_api_impl(self, qlinear_module, module_name, qlinear_op,
             loaded_dict = torch.load(b, weights_only=weights_only)
             for key in model_dict:
                 if isinstance(model_dict[key], torch._C.ScriptObject):
-                    assert isinstance(loaded_dict[key], torch._C.ScriptObject)
+                    if not isinstance(loaded_dict[key], torch._C.ScriptObject):
+                        raise AssertionError(
+                            f"Expected loaded_dict[{key}] to be ScriptObject, got {type(loaded_dict[key])}"
+                        )
                     w_model, b_model = torch.ops.quantized.linear_unpack(model_dict[key])
                     w_loaded, b_loaded = torch.ops.quantized.linear_unpack(loaded_dict[key])
                     self.assertEqual(w_model, w_loaded)
@@ -206,7 +209,8 @@ def _test_linear_api_impl(self, qlinear_module, module_name, qlinear_op,
 
         for name, _ in loaded_from_package.named_modules():
             # noop, just make sure attribute "_modules" is restored correctly during torch.package import
-            assert(name is not None)  # noqa: E275
+            if name is None:
+                raise AssertionError("name is None")
 
         # Test copy and deepcopy
         copied_linear = copy.copy(qlinear)
@@ -1643,7 +1647,10 @@ def test_linear_api(self, batch_size, in_features, out_features, use_bias, use_d
             loaded_dict = torch.load(b, weights_only=weights_only)
             for key in model_dict:
                 if isinstance(model_dict[key], torch._C.ScriptObject):
-                    assert isinstance(loaded_dict[key], torch._C.ScriptObject)
+                    if not isinstance(loaded_dict[key], torch._C.ScriptObject):
+                        raise AssertionError(
+                            f"Expected loaded_dict[{key}] to be ScriptObject, got {type(loaded_dict[key])}"
+                        )
                     w_model, b_model = torch.ops.quantized.linear_unpack(model_dict[key])
                     w_loaded, b_loaded = torch.ops.quantized.linear_unpack(loaded_dict[key])
                     self.assertEqual(w_model, w_loaded)
diff --git a/test/quantization/core/test_quantized_op.py b/test/quantization/core/test_quantized_op.py
index b78f3a56820ba..910fc677fdafb 100644
--- a/test/quantization/core/test_quantized_op.py
+++ b/test/quantization/core/test_quantized_op.py
@@ -25,7 +25,6 @@
 
 hu.assert_deadline_disabled()
 
-from typing import Optional
 
 import torch.backends.xnnpack
 from torch.ao.quantization import PerChannelMinMaxObserver
@@ -105,7 +104,10 @@ def avoid_vpmaddubsw_overflow_linear(
             x1 = X[i, k + 1] - X_min
             w0 = W[j, k] - 128 - W_min
             w1 = W[j, k + 1] - 128 - W_min
-            assert -(1 << 15) <= x0 * w0 + x1 * w1 < (1 << 15)
+            if not (-(1 << 15) <= x0 * w0 + x1 * w1 < (1 << 15)):
+                raise AssertionError(
+                    f"Product sum {x0 * w0 + x1 * w1} out of int16 range"
+                )
 
 
 # Reference quantized Linear operator
@@ -113,7 +115,8 @@ def qlinear_ref(X_q, X_scale, X_zp, W_q, W_scale, W_zp, b_q, Y_scale, Y_zp, dtyp
     X_q = np.reshape(X_q, (-1, X_q.shape[X_q.ndim - 1]))
     row_offsets_ref = X_q.sum(axis=1).astype(np.int32).reshape((-1, 1))
     col_offsets_ref = W_q.sum(axis=1).astype(np.int32).reshape((1, -1))
-    assert X_q.ndim == 2
+    if X_q.ndim != 2:
+        raise AssertionError(f"Expected X_q.ndim == 2, got {X_q.ndim}")
     batch_size, input_channels = X_q.shape
     Prod_XqWq_ref = (
         np.matmul(X_q.astype(np.int32), W_q.astype(np.int32).T)
@@ -164,7 +167,7 @@ def _get_random_tensor_and_q_params(shapes, rand_scale, torch_type):
         X_scale = 1e-10
     return X, X_scale, X_zero_point
 
-def _quantize_fp8e4m3(t: torch.Tensor, channelwise: bool, scale: Optional[torch.Tensor] = None):
+def _quantize_fp8e4m3(t: torch.Tensor, channelwise: bool, scale: torch.Tensor | None = None):
     quant_max = torch.finfo(torch.float8_e4m3fn).max
     eps = torch.Tensor([torch.finfo(torch.float32).eps])
     if channelwise:
@@ -2181,7 +2184,11 @@ def test_qtopk(self):
 
             quantized_out = torch.topk(qX, k, dim=dim, largest=large, sorted=sort)
 
-            assert len(unquantized_out) == len(quantized_out)
+            if len(unquantized_out) != len(quantized_out):
+                raise AssertionError(
+                    f"Expected len(unquantized_out) == len(quantized_out), "
+                    f"got {len(unquantized_out)} != {len(quantized_out)}"
+                )
             torch.testing.assert_close(quantized_out[0].dequantize(), unquantized_out[0])
             torch.testing.assert_close(quantized_out[1], unquantized_out[1])
 
@@ -3060,7 +3067,11 @@ def from_float(cls, other, qconfig=None):
                 self.assertEqual(
                     lstm_prepared[0].layers[0].layer_fw.cell.split_gates, split_gates
                 )
-                assert isinstance(lstm_prepared[0], torch.ao.nn.quantizable.LSTM)
+                if not isinstance(lstm_prepared[0], torch.ao.nn.quantizable.LSTM):
+                    raise AssertionError(
+                        f"Expected lstm_prepared[0] to be an instance of "
+                        f"torch.ao.nn.quantizable.LSTM, got {type(lstm_prepared[0])}"
+                    )
 
                 # Calibrate
                 y = lstm_prepared(x)
@@ -3070,7 +3081,11 @@ def from_float(cls, other, qconfig=None):
                 lstm_quantized = torch.ao.quantization.convert(
                     lstm_prepared, convert_custom_config_dict=custom_config_dict
                 )
-                assert type(lstm_quantized[0]) is torch.ao.nn.quantized.LSTM
+                if type(lstm_quantized[0]) is not torch.ao.nn.quantized.LSTM:
+                    raise AssertionError(
+                        f"Expected type(lstm_quantized[0]) to be "
+                        f"torch.ao.nn.quantized.LSTM, got {type(lstm_quantized[0])}"
+                    )
                 qy = lstm_quantized(qx)
 
                 snr = _snr(y, qy)
@@ -3100,9 +3115,9 @@ def forward(
                 query,
                 key,
                 value,
-                key_padding_mask: Optional[torch.Tensor] = None,
+                key_padding_mask: torch.Tensor | None = None,
                 need_weights: bool = True,
-                attn_mask: Optional[torch.Tensor] = None,
+                attn_mask: torch.Tensor | None = None,
             ):
                 return self.layer(query, key, value, key_padding_mask, need_weights, attn_mask)
 
@@ -4510,7 +4525,8 @@ def test_wrapped_quantized_linear_prepacked(self, m, n, k):
                        qparams=hu.qparams(dtypes=torch.qint8)))
     @override_qengines
     def test_qlinear_qnnpack_free_memory_and_unpack(self, W):
-        assert qengine_is_qnnpack
+        if not qengine_is_qnnpack:
+            raise AssertionError("Expected qengine_is_qnnpack to be True")
         W, (W_scale, W_zp, torch_type) = W
         qlinear_prepack = torch.ops.quantized.linear_prepack
         qlinear_unpack = torch.ops.quantized.linear_unpack
@@ -4872,7 +4888,8 @@ def _test_qlinear_fp8_helper(
 
                 self.assertEqual(x.dim(), qy.dim())
                 self.assertEqual(y_ref.float(), qy.float())
-                assert not torch.isnan(qy).any()
+                if torch.isnan(qy).any():
+                    raise AssertionError("Output qy contains NaN values")
 
     @unittest.skipIf(IS_FBCODE, "Skip pt2e ops in fbcode")
     @skipIfNoONEDNN
@@ -5347,8 +5364,8 @@ def _make_qconv_tensors(
         input_dtype=torch.quint8,
         weight_dtype=torch.qint8,
     ):
-        assert not (use_channelwise and use_transpose), \
-               "Cannot generate channelwise qconv_transpose_tensors "
+        if use_channelwise and use_transpose:
+            raise AssertionError("Cannot generate channelwise qconv_transpose_tensors ")
         input_channels = input_channels_per_group * groups
         output_channels = output_channels_per_group * groups
         # Padded input size should be at least as big as dilated kernel
@@ -5454,7 +5471,8 @@ def _test_qconv_impl(
             bias_float, requires_grad=False) if use_bias else None
         result_ref = conv_op(X)
         if post_op == 'relu':
-            assert not use_transpose, "Cannot fuse ReLU with ConvTranspose"
+            if use_transpose:
+                raise AssertionError("Cannot fuse ReLU with ConvTranspose")
             relu = torch.nn.ReLU()
             result_ref = relu(result_ref)
         elif post_op == 'add':
@@ -7136,20 +7154,25 @@ def _test_qconv_impl_cpu_tensor(
                 relu = torch.nn.ReLU()
                 result_ref = relu(result_ref)
         elif post_op.unary_attr == "relu":
-            assert not use_transpose, "Cannot fuse ReLU with ConvTranspose"
+            if use_transpose:
+                raise AssertionError("Cannot fuse ReLU with ConvTranspose")
             relu = torch.nn.ReLU()
             result_ref = relu(result_ref)
         elif post_op.unary_attr == "hardtanh":
-            assert not use_transpose, "Cannot fuse hardtanh with ConvTranspose"
-            assert len(post_op.scalars) == 2, "For post op hardtanh, expect 2 parameters passed in"
+            if use_transpose:
+                raise AssertionError("Cannot fuse hardtanh with ConvTranspose")
+            if len(post_op.scalars) != 2:
+                raise AssertionError("For post op hardtanh, expect 2 parameters passed in")
             hardtanh = torch.nn.Hardtanh(min_val=post_op.scalars[0], max_val=post_op.scalars[1])
             result_ref = hardtanh(result_ref)
         elif post_op.unary_attr == "hardswish":
-            assert not use_transpose, "Cannot fuse hardswish with ConvTranspose"
+            if use_transpose:
+                raise AssertionError("Cannot fuse hardswish with ConvTranspose")
             hardswish = torch.nn.Hardswish()
             result_ref = hardswish(result_ref)
         elif post_op.unary_attr == "swish":
-            assert not use_transpose, "Cannot fuse silu with ConvTranspose"
+            if use_transpose:
+                raise AssertionError("Cannot fuse silu with ConvTranspose")
             silu = torch.nn.SiLU()
             result_ref = silu(result_ref)
 
@@ -7882,8 +7905,8 @@ def _make_qconv_tensors_fp8(
         use_bias, use_channelwise, use_transpose, bfloat16_output,
         device=torch.device("cpu"),
     ):
-        assert not (use_channelwise and use_transpose), \
-               "Cannot generate channelwise qconv_transpose_tensors "
+        if use_channelwise and use_transpose:
+            raise AssertionError("Cannot generate channelwise qconv_transpose_tensors ")
         input_channels = input_channels_per_group * groups
         output_channels = output_channels_per_group * groups
         # Padded input size should be at least as big as dilated kernel
@@ -7987,20 +8010,25 @@ def _test_qconv_impl_cpu_tensor_fp8(
                 relu = torch.nn.ReLU()
                 result_ref = relu(result_ref)
         elif post_op.unary_attr == "relu":
-            assert not use_transpose, "Cannot fuse ReLU with ConvTranspose"
+            if use_transpose:
+                raise AssertionError("Cannot fuse ReLU with ConvTranspose")
             relu = torch.nn.ReLU()
             result_ref = relu(result_ref)
         elif post_op.unary_attr == "hardtanh":
-            assert not use_transpose, "Cannot fuse hardtanh with ConvTranspose"
-            assert len(post_op.scalars) == 2, "For post op hardtanh, expect 2 parameters passed in"
+            if use_transpose:
+                raise AssertionError("Cannot fuse hardtanh with ConvTranspose")
+            if len(post_op.scalars) != 2:
+                raise AssertionError("For post op hardtanh, expect 2 parameters passed in")
             hardtanh = torch.nn.Hardtanh(min_val=post_op.scalars[0], max_val=post_op.scalars[1])
             result_ref = hardtanh(result_ref)
         elif post_op.unary_attr == "hardswish":
-            assert not use_transpose, "Cannot fuse hardswish with ConvTranspose"
+            if use_transpose:
+                raise AssertionError("Cannot fuse hardswish with ConvTranspose")
             hardswish = torch.nn.Hardswish()
             result_ref = hardswish(result_ref)
         elif post_op.unary_attr == "swish":
-            assert not use_transpose, "Cannot fuse silu with ConvTranspose"
+            if use_transpose:
+                raise AssertionError("Cannot fuse silu with ConvTranspose")
             silu = torch.nn.SiLU()
             result_ref = silu(result_ref)
 
@@ -8084,13 +8112,15 @@ def _test_qconv_impl_cpu_tensor_fp8(
             self.assertTrue(result.dtype == qconv_output_dtype)
 
         self.assertEqual(result.float(), result_ref.float(), atol=1e-6, rtol=1e-5)
-        assert not torch.isnan(result).any()
+        if torch.isnan(result).any():
+            raise AssertionError("Output result contains NaN values")
 
     def _test_qconv_fp8_helper(self, nd, pointwise_post_op):
         # nd = 1,2,3 -> conv1d/2d/3d
         if pointwise_post_op.binary_attr != "none":
             # Only conv2d supports binary post op
-            assert nd == 2
+            if nd != 2:
+                raise AssertionError(f"Expected nd == 2, got {nd}")
         groups_list = [1, 3]
         input_channels_per_group = 2
         output_channels_per_group = 2
@@ -8880,13 +8910,21 @@ def test_quantize_tensor_with_min_max(self):
                     weight_quantized = quant_fn_with_rowwise_min_max(
                         weight, rowwise_min_max
                     )
-                    assert weight_quantized.dtype == torch.uint8
+                    if weight_quantized.dtype != torch.uint8:
+                        raise AssertionError(
+                            f"Expected weight_quantized.dtype == torch.uint8, "
+                            f"got {weight_quantized.dtype}"
+                        )
 
                     # Confirm that the quantization is matching the one without rowwise_min_max
                     weight_quantized_no_rowwise_min_max = quant_fn(weight)
-                    assert torch.equal(
+                    if not torch.equal(
                         weight_quantized, weight_quantized_no_rowwise_min_max
-                    )
+                    ):
+                        raise AssertionError(
+                            "weight_quantized does not equal "
+                            "weight_quantized_no_rowwise_min_max"
+                        )
 
                     # Confirtm that incorrect rowwise_min_max will result in different quantization output
                     incorrect_rowwise_min_max = torch.stack(
@@ -8895,10 +8933,18 @@ def test_quantize_tensor_with_min_max(self):
                     weight_incorrectly_quantized = quant_fn_with_rowwise_min_max(
                         weight, incorrect_rowwise_min_max
                     )
-                    assert weight_incorrectly_quantized.dtype == torch.uint8
-                    assert not torch.equal(
+                    if weight_incorrectly_quantized.dtype != torch.uint8:
+                        raise AssertionError(
+                            f"Expected weight_incorrectly_quantized.dtype == torch.uint8, "
+                            f"got {weight_incorrectly_quantized.dtype}"
+                        )
+                    if torch.equal(
                         weight_incorrectly_quantized, weight_quantized_no_rowwise_min_max
-                    )
+                    ):
+                        raise AssertionError(
+                            "weight_incorrectly_quantized should not equal "
+                            "weight_quantized_no_rowwise_min_max"
+                        )
 
 if __name__ == "__main__":
     raise_on_run_directly("test/test_quantization.py")
diff --git a/test/quantization/core/test_workflow_module.py b/test/quantization/core/test_workflow_module.py
index 431a8db871946..56b43430897e2 100644
--- a/test/quantization/core/test_workflow_module.py
+++ b/test/quantization/core/test_workflow_module.py
@@ -97,7 +97,8 @@ def test_per_tensor_observers(self, qdtype, qscheme, reduce_range):
                                                     reduce_range=reduce_range)]
 
         def _get_ref_params(reduce_range, qscheme, dtype, input_scale, min_val, max_val):
-            assert dtype in _INT_DTYPES, f"Not supported dtype: {dtype}, supported dtypes are {_INT_DTYPES}"
+            if dtype not in _INT_DTYPES:
+                raise AssertionError(f"Not supported dtype: {dtype}, supported dtypes are {_INT_DTYPES}")
             eps = torch.tensor([tolerance])
             if dtype in [torch.qint8, torch.int8]:
                 if reduce_range:
@@ -520,7 +521,8 @@ def _get_norm(delta_begin, delta_end, density, norm_type):
             norm = density * (integral_{begin, end} x^2)
                  = density * (end^3 - begin^3) / 3
             """
-            assert norm_type == "L2", "Only L2 norms are currently supported"
+            if norm_type != "L2":
+                raise AssertionError("Only L2 norms are currently supported")
             norm = 0.0
             if norm_type == "L2":
                 norm = (
@@ -581,7 +583,8 @@ def _compute_quantization_error(next_start_bin, next_end_bin, norm_type):
                     norm = norm + _get_norm(delta_begin, delta_end, density, norm_type)
             return norm
 
-        assert self.histogram.size()[0] == self.bins, "bins mismatch"
+        if self.histogram.size()[0] != self.bins:
+            raise AssertionError("bins mismatch")
         bin_width = (self.max_val - self.min_val) / self.bins
 
         # cumulative sum
@@ -867,8 +870,10 @@ def test_fq_module_per_channel(self, device, X):
         X.requires_grad_()
         fq_module = FakeQuantize(default_per_channel_weight_observer, quant_min, quant_max, ch_axis=axis).to(device)
         Y_prime = fq_module(X)
-        assert fq_module.scale is not None
-        assert fq_module.zero_point is not None
+        if fq_module.scale is None:
+            raise AssertionError("fq_module.scale should not be None")
+        if fq_module.zero_point is None:
+            raise AssertionError("fq_module.zero_point should not be None")
         Y = _fake_quantize_per_channel_affine_reference(X, fq_module.scale,
                                                         fq_module.zero_point, axis, quant_min, quant_max)
         np.testing.assert_allclose(Y.cpu().detach().numpy(), Y_prime.cpu().detach().numpy(), rtol=tolerance, atol=tolerance)
@@ -962,7 +967,10 @@ def forward(self, x):
                 quantizable_model = torch.ao.quantization.quantize_fx.prepare_qat_fx(fused_model,
                                                                                      qconfig_mapping,
                                                                                      example_inputs=None)
-                assert isinstance(quantizable_model.conv, torch.ao.nn.intrinsic.qat.ConvBn2d)
+                if not isinstance(quantizable_model.conv, torch.ao.nn.intrinsic.qat.ConvBn2d):
+                    raise AssertionError(
+                        f"Expected quantizable_model.conv to be ConvBn2d, got {type(quantizable_model.conv)}"
+                    )
 
                 # jit script
                 scripted_model = torch.jit.script(quantizable_model)
@@ -999,7 +1007,10 @@ def forward(self, x):
                 # convert to QAT
                 fused_model.qconfig = torch.ao.quantization.get_default_qconfig(qengine)
                 torch.ao.quantization.prepare_qat(fused_model, inplace=True)
-                assert isinstance(fused_model.conv, torch.ao.nn.intrinsic.qat.ConvBn2d)
+                if not isinstance(fused_model.conv, torch.ao.nn.intrinsic.qat.ConvBn2d):
+                    raise AssertionError(
+                        f"Expected fused_model.conv to be ConvBn2d, got {type(fused_model.conv)}"
+                    )
 
                 # Test jit script fails
                 # Prepared eager module fails due to observer hooks not being scriptable
diff --git a/test/quantization/core/test_workflow_ops.py b/test/quantization/core/test_workflow_ops.py
index 0ac6235bcec30..fd7e8516bf099 100644
--- a/test/quantization/core/test_workflow_ops.py
+++ b/test/quantization/core/test_workflow_ops.py
@@ -3,7 +3,6 @@
 
 import torch
 import math
-from typing import Union
 from torch.ao.quantization import (
     FakeQuantize,
     MovingAverageMinMaxObserver,
@@ -182,8 +181,8 @@ def _fake_quantize_learnable_per_channel_affine_grad_reference(
 
 def _get_tensor_min_max(
         X: torch.Tensor,
-        running_min: Union[float, torch.Tensor] = float("inf"),
-        running_max: Union[float, torch.Tensor] = float("-inf"),
+        running_min: float | torch.Tensor = float("inf"),
+        running_max: float | torch.Tensor = float("-inf"),
         averaging_const: float = 0.01,
         dtype: torch.dtype = torch.float32) -> tuple[float, float]:
     min_val_tensor = X.min().to(dtype=dtype)
@@ -567,8 +566,10 @@ def test_fq_module_per_tensor(self, device, X):
         X.requires_grad_()
         fq_module = torch.ao.quantization.default_fake_quant().to(device)
         Y_prime = fq_module(X)
-        assert fq_module.scale is not None
-        assert fq_module.zero_point is not None
+        if fq_module.scale is None:
+            raise AssertionError("fq_module.scale should not be None")
+        if fq_module.zero_point is None:
+            raise AssertionError("fq_module.zero_point should not be None")
         Y = _fake_quantize_per_tensor_affine_reference(X, fq_module.scale, fq_module.zero_point, quant_min, quant_max)
         np.testing.assert_allclose(Y.cpu().detach().numpy(), Y_prime.cpu().detach().numpy(), rtol=tolerance, atol=tolerance)
 
@@ -919,7 +920,10 @@ def _test_backward_per_channel_cachemask_impl(self, device):
             Y_prime.backward(dout)
             np.testing.assert_allclose(
                 dX.cpu().detach().numpy(), X.grad.cpu().detach().numpy(), rtol=tolerance, atol=tolerance)
-            assert X.grad.dtype == float_type
+            if X.grad.dtype != float_type:
+                raise AssertionError(
+                    f"Expected X.grad.dtype to be {float_type}, got {X.grad.dtype}"
+                )
 
 
     def test_backward_per_channel_cachemask_cpu(self):
diff --git a/test/quantization/eager/test_quantize_eager_ptq.py b/test/quantization/eager/test_quantize_eager_ptq.py
index 25d98c1c63b50..c15f0d33abdd0 100644
--- a/test/quantization/eager/test_quantize_eager_ptq.py
+++ b/test/quantization/eager/test_quantize_eager_ptq.py
@@ -996,7 +996,10 @@ def forward(self, x):
 
             @classmethod
             def from_float(cls, float_module):
-                assert hasattr(float_module, "qconfig")
+                if not hasattr(float_module, "qconfig"):
+                    raise AssertionError(
+                        f"float_module {type(float_module).__name__} must have qconfig"
+                    )
                 observed = cls(float_module.conv)
                 observed.qconfig = float_module.qconfig
                 return observed
@@ -1011,8 +1014,14 @@ def forward(self, x):
 
             @classmethod
             def from_observed(cls, observed_module):
-                assert hasattr(observed_module, "qconfig")
-                assert hasattr(observed_module, "activation_post_process")
+                if not hasattr(observed_module, "qconfig"):
+                    raise AssertionError(
+                        f"observed_module {type(observed_module).__name__} must have qconfig"
+                    )
+                if not hasattr(observed_module, "activation_post_process"):
+                    raise AssertionError(
+                        f"observed_module {type(observed_module).__name__} must have activation_post_process"
+                    )
                 observed_module.conv.activation_post_process = (
                     observed_module.activation_post_process
                 )
diff --git a/test/quantization/eager/test_quantize_eager_qat.py b/test/quantization/eager/test_quantize_eager_qat.py
index a6655798c5cff..1507246c3b798 100644
--- a/test/quantization/eager/test_quantize_eager_qat.py
+++ b/test/quantization/eager/test_quantize_eager_qat.py
@@ -106,7 +106,8 @@ def __init__(
             False,
             padding_mode,
         )
-        assert qconfig, "qconfig must be provided for QAT module"
+        if not qconfig:
+            raise AssertionError("qconfig must be provided for QAT module")
         self.qconfig = qconfig
         self.eps = eps
         self.momentum = momentum
@@ -241,17 +242,18 @@ def from_float(cls, mod, qconfig=None):
         Args: `mod` a float module, either produced by torch.ao.quantization utilities
         or directly from user
         """
-        assert type(mod) is cls._FLOAT_MODULE, (
-            "qat."
-            + cls.__name__
-            + ".from_float only works for "
-            + cls._FLOAT_MODULE.__name__
-        )
-        if not qconfig:
-            assert hasattr(mod, "qconfig"), (
-                "Input float module must have qconfig defined"
+        if type(mod) is not cls._FLOAT_MODULE:
+            raise AssertionError(
+                "qat."
+                + cls.__name__
+                + ".from_float only works for "
+                + cls._FLOAT_MODULE.__name__
             )
-            assert mod.qconfig, "Input float module must have a valid qconfig"
+        if not qconfig:
+            if not hasattr(mod, "qconfig"):
+                raise AssertionError("Input float module must have qconfig defined")
+            if not mod.qconfig:
+                raise AssertionError("Input float module must have a valid qconfig")
             qconfig = mod.qconfig
         conv, bn = mod[0], mod[1]
         qat_convbn = cls(
diff --git a/test/quantization/fx/test_model_report_fx.py b/test/quantization/fx/test_model_report_fx.py
index cab72394ae29d..d05c935190260 100644
--- a/test/quantization/fx/test_model_report_fx.py
+++ b/test/quantization/fx/test_model_report_fx.py
@@ -1675,7 +1675,10 @@ def test_all_outlier_report_gen(self):
                 # everything should be an outlier because we said that the max should be equal to the min for all of them
                 # however we will just test and say most should be in case we have several 0 channel values
                 outlier_info = module_dict[OutlierDetector.OUTLIER_KEY]
-                assert sum(outlier_info) >= len(outlier_info) / 2
+                if not sum(outlier_info) >= len(outlier_info) / 2:
+                    raise AssertionError(
+                        f"Expected at least half of channels to be outliers, got {sum(outlier_info)} out of {len(outlier_info)}"
+                    )
 
                 # ensure that the number of ratios and batches counted is the same as the number of params
                 self.assertEqual(len(module_dict[OutlierDetector.COMP_METRIC_KEY]), param_size)
@@ -1735,7 +1738,11 @@ def test_multiple_run_consistent_spike_outlier_report_gen(self):
                 # because we ran 30 times, we should have at least a couple be significant
                 # could be less because some channels could possibly be all 0
                 sufficient_batches_info = module_dict[OutlierDetector.IS_SUFFICIENT_BATCHES_KEY]
-                assert sum(sufficient_batches_info) >= len(sufficient_batches_info) / 2
+                if not sum(sufficient_batches_info) >= len(sufficient_batches_info) / 2:
+                    raise AssertionError(
+                        f"Expected at least half of channels to have sufficient batches, "
+                        f"got {sum(sufficient_batches_info)} out of {len(sufficient_batches_info)}"
+                    )
 
                 # half of them should be outliers, because we set a really high value every 2 channels
                 outlier_info = module_dict[OutlierDetector.OUTLIER_KEY]
@@ -1751,7 +1758,10 @@ def test_multiple_run_consistent_spike_outlier_report_gen(self):
                     # check that the non-zero channel count, at least 2 should be there
                     # for the first module
                     counts_info = module_dict[OutlierDetector.CONSTANT_COUNTS_KEY]
-                    assert sum(counts_info) >= 2
+                    if not sum(counts_info) >= 2:
+                        raise AssertionError(
+                            f"Expected at least 2 non-zero channel counts, got {sum(counts_info)}"
+                        )
 
                     # half of the recorded max values should be what we set
                     matched_max = sum(val == 3.28e8 for val in module_dict[OutlierDetector.MAX_VALS_KEY])
diff --git a/test/quantization/fx/test_numeric_suite_fx.py b/test/quantization/fx/test_numeric_suite_fx.py
index ed3eabd702690..af272830da4b2 100644
--- a/test/quantization/fx/test_numeric_suite_fx.py
+++ b/test/quantization/fx/test_numeric_suite_fx.py
@@ -1833,12 +1833,27 @@ def test_extend_logger_results_with_comparison(self):
             'cosine_similarity_int8_vs_fp32')
 
         for layer_results in results.values():
-            assert 'sqnr_int8_vs_fp32' in \
-                layer_results['weight']['int8'][0]
-            assert 'l2_error_int8_vs_fp32' in \
-                layer_results['weight']['int8'][0]
-            assert 'cosine_similarity_int8_vs_fp32' in \
-                layer_results['weight']['int8'][0]
+            if (
+                'sqnr_int8_vs_fp32'
+                not in layer_results['weight']['int8'][0]
+            ):
+                raise AssertionError(
+                    f"'sqnr_int8_vs_fp32' not found in layer results: {layer_results['weight']['int8'][0].keys()}"
+                )
+            if (
+                'l2_error_int8_vs_fp32'
+                not in layer_results['weight']['int8'][0]
+            ):
+                raise AssertionError(
+                    f"'l2_error_int8_vs_fp32' not found in layer results: {layer_results['weight']['int8'][0].keys()}"
+                )
+            if (
+                'cosine_similarity_int8_vs_fp32'
+                not in layer_results['weight']['int8'][0]
+            ):
+                raise AssertionError(
+                    f"'cosine_similarity_int8_vs_fp32' not found in layer results: {layer_results['weight']['int8'][0].keys()}"
+                )
 
     @skipIfNoFBGEMM
     def test_int8_shadows_fp32_simple(self):
diff --git a/test/quantization/fx/test_quantize_fx.py b/test/quantization/fx/test_quantize_fx.py
index d848e2bd8d580..8584b9f405d76 100644
--- a/test/quantization/fx/test_quantize_fx.py
+++ b/test/quantization/fx/test_quantize_fx.py
@@ -204,7 +204,6 @@
 import operator
 import unittest
 import io
-from typing import Optional
 from collections.abc import Callable
 
 class BinaryOp(torch.nn.Module):
@@ -3158,7 +3157,8 @@ def forward(self, x):
 
             @classmethod
             def from_float(cls, float_module):
-                assert hasattr(float_module, 'qconfig')
+                if not hasattr(float_module, 'qconfig'):
+                    raise AssertionError(f"float_module missing 'qconfig': {type(float_module)}")
                 observed = cls(float_module.linear)
                 observed.qconfig = float_module.qconfig
                 return observed
@@ -3173,8 +3173,10 @@ def forward(self, x):
 
             @classmethod
             def from_observed(cls, observed_module):
-                assert hasattr(observed_module, 'qconfig')
-                assert hasattr(observed_module, 'activation_post_process')
+                if not hasattr(observed_module, 'qconfig'):
+                    raise AssertionError(f"observed_module missing 'qconfig': {type(observed_module)}")
+                if not hasattr(observed_module, 'activation_post_process'):
+                    raise AssertionError(f"observed_module missing 'activation_post_process': {type(observed_module)}")
                 observed_module.linear.activation_post_process = \
                     observed_module.activation_post_process
                 quantized = cls(nnq.Linear.from_float(observed_module.linear))
@@ -3190,7 +3192,8 @@ def forward(self, x):
 
             @classmethod
             def from_observed(cls, observed_module):
-                assert hasattr(observed_module, 'qconfig')
+                if not hasattr(observed_module, 'qconfig'):
+                    raise AssertionError(f"observed_module missing 'qconfig': {type(observed_module)}")
                 observed_module.linear.qconfig = observed_module.qconfig
                 quantized = cls(nnqd.Linear.from_float(observed_module.linear))
                 return quantized
@@ -3328,7 +3331,8 @@ def forward(self, x):
 
             @classmethod
             def from_float(cls, float_module):
-                assert hasattr(float_module, 'qconfig')
+                if not hasattr(float_module, 'qconfig'):
+                    raise AssertionError(f"float_module missing 'qconfig': {type(float_module)}")
                 observed = cls(float_module.linear)
                 observed.qconfig = float_module.qconfig
                 return observed
@@ -3343,8 +3347,10 @@ def forward(self, x):
 
             @classmethod
             def from_observed(cls, observed_module):
-                assert hasattr(observed_module, 'qconfig')
-                assert hasattr(observed_module, 'activation_post_process')
+                if not hasattr(observed_module, 'qconfig'):
+                    raise AssertionError(f"observed_module missing 'qconfig': {type(observed_module)}")
+                if not hasattr(observed_module, 'activation_post_process'):
+                    raise AssertionError(f"observed_module missing 'activation_post_process': {type(observed_module)}")
                 observed_module.linear.activation_post_process = \
                     observed_module.activation_post_process
                 quantized = cls(nnq.Linear.from_float(observed_module.linear))
@@ -3412,7 +3418,8 @@ def forward(self, x):
 
             @classmethod
             def from_float(cls, float_module):
-                assert hasattr(float_module, 'qconfig')
+                if not hasattr(float_module, 'qconfig'):
+                    raise AssertionError(f"float_module missing 'qconfig': {type(float_module)}")
                 observed = cls(float_module.linear)
                 observed.qconfig = float_module.qconfig
                 return observed
@@ -3427,8 +3434,10 @@ def forward(self, x):
 
             @classmethod
             def from_observed(cls, observed_module):
-                assert hasattr(observed_module, 'qconfig')
-                assert hasattr(observed_module, 'activation_post_process')
+                if not hasattr(observed_module, 'qconfig'):
+                    raise AssertionError(f"observed_module missing 'qconfig': {type(observed_module)}")
+                if not hasattr(observed_module, 'activation_post_process'):
+                    raise AssertionError(f"observed_module missing 'activation_post_process': {type(observed_module)}")
                 observed_module.linear.activation_post_process = \
                     observed_module.activation_post_process
                 quantized = cls(nnq.Linear.from_float(observed_module.linear))
@@ -3812,9 +3821,12 @@ def forward(self, x):
         m = prepare_fx(model, qconfig_dict, example_inputs=example_inputs)
         m(*example_inputs)
         m = convert_fx(m)
-        assert hasattr(m, "mods1_0_packed_weight_0")
-        assert hasattr(m, "mods1_1_packed_weight_0")
-        assert hasattr(m, "mods2_packed_weight_0")
+        if not hasattr(m, "mods1_0_packed_weight_0"):
+            raise AssertionError(f"m missing 'mods1_0_packed_weight_0', attrs: {list(m.__dict__.keys())}")
+        if not hasattr(m, "mods1_1_packed_weight_0"):
+            raise AssertionError(f"m missing 'mods1_1_packed_weight_0', attrs: {list(m.__dict__.keys())}")
+        if not hasattr(m, "mods2_packed_weight_0"):
+            raise AssertionError(f"m missing 'mods2_packed_weight_0', attrs: {list(m.__dict__.keys())}")
 
     @skipIfNoFBGEMM
     def test_mul_add_fp16_config(self):
@@ -4916,7 +4928,8 @@ class UserObservedLSTM(torch.ao.nn.quantizable.LSTM):
             """
             @classmethod
             def from_float(cls, float_lstm):
-                assert isinstance(float_lstm, cls._FLOAT_MODULE)
+                if not isinstance(float_lstm, cls._FLOAT_MODULE):
+                    raise AssertionError(f"Expected instance of {cls._FLOAT_MODULE}, got {type(float_lstm)}")
                 # uint16, [-16, 16)
                 linear_output_obs_ctr = FixedQParamsObserver.with_args(scale=2 ** -11, zero_point=2 ** 15, dtype=torch.qint32)
                 # uint16, [0, 1)
@@ -4946,7 +4959,8 @@ class UserQuantizedLSTM(torch.ao.nn.quantized.LSTM):
             """
             @classmethod
             def from_observed(cls, observed_lstm):
-                assert isinstance(observed_lstm, cls._FLOAT_MODULE)
+                if not isinstance(observed_lstm, cls._FLOAT_MODULE):
+                    raise AssertionError(f"Expected instance of {cls._FLOAT_MODULE}, got {type(observed_lstm)}")
                 return torch.ao.quantization.fx.lstm_utils._get_reference_quantized_lstm_module(
                     observed_lstm=observed_lstm,
                     backend_config=my_backend_config,
@@ -5375,7 +5389,8 @@ def forward(self, x):
                               'activation_post_process_6',
                               'activation_post_process_7',
                               'activation_post_process_10']
-        assert name_list == expected_name_list
+        if name_list != expected_name_list:
+            raise AssertionError(f"name_list mismatch: {name_list} != {expected_name_list}")
 
     def test_conv_lowering(self):
         convs = {1: nn.Conv1d, 2: nn.Conv2d, 3: nn.Conv3d}
@@ -5777,7 +5792,7 @@ def _validate_qconfig_against_backend_config_constraints(
             qconfig: QConfig,
             backend_config: BackendConfig,
             satisfies_constraints: bool,
-            qconfig_name: Optional[str] = None):
+            qconfig_name: str | None = None):
         """
         Helper method to validate whether `qconfig` satisfies the constraints specified in `backend_config`.
         """
@@ -9431,7 +9446,8 @@ def _test_model_impl(
                      nprocs=world_size,  # noqa: F821
                      join=True)
         elif mode == 'qat':
-            assert prepared.training, 'prepared must be in training mode for qat'
+            if not prepared.training:
+                raise AssertionError("prepared must be in training mode for qat")
             optimizer = torch.optim.SGD(prepared.parameters(), lr=0.0001)
             criterion = nn.CrossEntropyLoss()
             train_one_epoch(prepared, criterion, optimizer, [(input_value, output_value)], torch.device('cpu'), 1)
@@ -9453,7 +9469,8 @@ def _test_model_impl(
 
         if is_not_tuple_out:
             diff_of_quant[mode][name] = (original_out - qgraph_out).abs().max()
-            assert torch.allclose(qgraph_out, qgraph_script), 'graph, scripted graph'
+            if not torch.allclose(qgraph_out, qgraph_script):
+                raise AssertionError("graph, scripted graph")
         else:
             print('tuple output')
 
@@ -9477,7 +9494,8 @@ def _test_model_impl(
                          nprocs=world_size,  # noqa: F821
                          join=True)
             elif mode == 'qat':
-                assert qeager.training, 'qeager should be in training mode for qat'
+                if not qeager.training:
+                    raise AssertionError("qeager should be in training mode for qat")
                 optimizer = torch.optim.SGD(qeager.parameters(), lr=0.0001)
                 train_one_epoch(qeager, criterion, optimizer, [(input_value, output_value)], torch.device('cpu'), 1)
             else:
@@ -9516,7 +9534,8 @@ def _test_building_block(self, quant_type, BB):
             data = self.img_data_2d
             is_qat = False
         else:
-            assert quant_type == QuantType.QAT
+            if quant_type != QuantType.QAT:
+                raise AssertionError(f"Expected QuantType.QAT, got {quant_type}")
             qconfig = default_qat_qconfig
             eager_prepare = prepare_qat
             graph_prepare = prepare_qat_fx
@@ -9657,8 +9676,8 @@ def __init__(self) -> None:
                     self.emb = torch.nn.EmbeddingBag(num_embeddings=10, embedding_dim=12, mode='sum')
                     self.linear = torch.nn.Linear(12, 1).to(dtype=torch.float)
 
-                def forward(self, input: torch.Tensor, offsets: Optional[torch.Tensor] = None,
-                            per_sample_weights: Optional[torch.Tensor] = None):
+                def forward(self, input: torch.Tensor, offsets: torch.Tensor | None = None,
+                            per_sample_weights: torch.Tensor | None = None):
                     x = self.emb(input, offsets, per_sample_weights)
                     x = self.linear(x)
                     return x
diff --git a/test/quantization/fx/test_subgraph_rewriter.py b/test/quantization/fx/test_subgraph_rewriter.py
index e410f93803d64..bdaa498fea192 100644
--- a/test/quantization/fx/test_subgraph_rewriter.py
+++ b/test/quantization/fx/test_subgraph_rewriter.py
@@ -454,8 +454,10 @@ def forward(self, x):
         symbolic_traced: torch.fx.GraphModule = symbolic_trace(module)
         for n, m in zip(symbolic_traced.graph.nodes, graph.nodes):
             if n.op == 'placeholder':
-                assert n.type is int
-                assert m.type is int
+                if n.type is not int:
+                    raise AssertionError(f"Expected n.type to be int, got {n.type}")
+                if m.type is not int:
+                    raise AssertionError(f"Expected m.type to be int, got {m.type}")
 
     def test_subgraph_writer_replace_consecutive_submodules(self):
 
diff --git a/test/quantization/jit/test_ondevice_quantization.py b/test/quantization/jit/test_ondevice_quantization.py
index 7fac580af2ef2..a92c73e0f821f 100644
--- a/test/quantization/jit/test_ondevice_quantization.py
+++ b/test/quantization/jit/test_ondevice_quantization.py
@@ -99,22 +99,21 @@ def get_linear_packed_param_fp_weight(node):
         ):
             raise ValueError("Quantized weight must be produced.")
         fp_weight = weight.inputsAt(0).node()
-        assert fp_weight.kind() == "prim::GetAttr", (
-            "Weight must be an attribute of the module."
-        )
+        if fp_weight.kind() != "prim::GetAttr":
+            raise AssertionError("Weight must be an attribute of the module.")
         fp_weight_name = fp_weight.s("name")
         return fp_weight_name
 
     @staticmethod
     def is_per_channel_quantized_packed_param(node):
-        assert node.kind() == "quantized::linear_prepack", (
-            "Node must corresponds to linear_prepack."
-        )
+        if node.kind() != "quantized::linear_prepack":
+            raise AssertionError("Node must corresponds to linear_prepack.")
         weight = node.inputsAt(0).node()
-        assert (
+        if not (
             weight.kind() != "aten::quantize_per_tensor"
             or weight.kind() != "aten::quantize_per_channel"
-        )
+        ):
+            raise AssertionError(f"Unexpected weight kind: {weight.kind()}")
         return weight.kind() != "aten::quantize_per_tensor"
 
 
diff --git a/test/quantization/jit/test_quantize_jit.py b/test/quantization/jit/test_quantize_jit.py
index 59e78f8694d8f..2a67a2ef622c0 100644
--- a/test/quantization/jit/test_quantize_jit.py
+++ b/test/quantization/jit/test_quantize_jit.py
@@ -453,9 +453,15 @@ def forward(self, x):
         qconfig_dict = {"": default_qconfig}
         m = prepare_jit(m, qconfig_dict)
         # for input and output of conv
-        assert len(attrs_with_prefix(m, "_observer_")) == 2
+        if len(attrs_with_prefix(m, "_observer_")) != 2:
+            raise AssertionError(
+                f"Expected 2 observers, got {len(attrs_with_prefix(m, '_observer_'))}"
+            )
         # for weight
-        assert len(attrs_with_prefix(m.conv, "_observer_")) == 1
+        if len(attrs_with_prefix(m.conv, "_observer_")) != 1:
+            raise AssertionError(
+                f"Expected 1 observer, got {len(attrs_with_prefix(m.conv, '_observer_'))}"
+            )
 
     def test_insert_observers_interface(self):
         @torch.jit.interface
@@ -551,13 +557,25 @@ def forward(self, x):
         qconfig_dict = {"sub.fc": default_qconfig}
         m = prepare_jit(m, qconfig_dict)
         # input and output of sub
-        assert len(attrs_with_prefix(m, "_observer_")) == 2
+        if len(attrs_with_prefix(m, "_observer_")) != 2:
+            raise AssertionError(
+                f"Expected 2 observers, got {len(attrs_with_prefix(m, '_observer_'))}"
+            )
         # not quantized
-        assert len(attrs_with_prefix(m.conv, "_observer_")) == 0
+        if len(attrs_with_prefix(m.conv, "_observer_")) != 0:
+            raise AssertionError(
+                f"Expected 0 observers, got {len(attrs_with_prefix(m.conv, '_observer_'))}"
+            )
         # no observers since we observe in the outer most call site
-        assert len(attrs_with_prefix(m.sub, "_observer_")) == 0
+        if len(attrs_with_prefix(m.sub, "_observer_")) != 0:
+            raise AssertionError(
+                f"Expected 0 observers, got {len(attrs_with_prefix(m.sub, '_observer_'))}"
+            )
         # weight of linear
-        assert len(attrs_with_prefix(m.sub.fc, "_observer_")) == 1
+        if len(attrs_with_prefix(m.sub.fc, "_observer_")) != 1:
+            raise AssertionError(
+                f"Expected 1 observer, got {len(attrs_with_prefix(m.sub.fc, '_observer_'))}"
+            )
 
     @unittest.skipUnless(
         "fbgemm" in torch.backends.quantized.supported_engines,
@@ -610,24 +628,45 @@ def attrs_with_prefix(module, prefix):
         m = torch.jit.script(ConvFunctionalReLU())
         m = prepare_jit(m, qconfig_dict)
         # observer for weight of conv
-        assert len(attrs_with_prefix(m.conv, "_observer_")) == 1
+        if len(attrs_with_prefix(m.conv, "_observer_")) != 1:
+            raise AssertionError(
+                f"Expected 1 observer, got {len(attrs_with_prefix(m.conv, '_observer_'))}"
+            )
         # observer for input of conv and output of relu
-        assert len(attrs_with_prefix(m, "_observer_")) == 2
+        if len(attrs_with_prefix(m, "_observer_")) != 2:
+            raise AssertionError(
+                f"Expected 2 observers, got {len(attrs_with_prefix(m, '_observer_'))}"
+            )
 
         m = torch.jit.script(ConvReLUModule())
         m = prepare_jit(m, qconfig_dict)
         # observer for input of conv and output of relu
-        assert len(attrs_with_prefix(m, "_observer_")) == 2
+        if len(attrs_with_prefix(m, "_observer_")) != 2:
+            raise AssertionError(
+                f"Expected 2 observers, got {len(attrs_with_prefix(m, '_observer_'))}"
+            )
         # observer for weight of conv
-        assert len(attrs_with_prefix(m.conv, "_observer_")) == 1
+        if len(attrs_with_prefix(m.conv, "_observer_")) != 1:
+            raise AssertionError(
+                f"Expected 1 observer, got {len(attrs_with_prefix(m.conv, '_observer_'))}"
+            )
         # observer for output of relu
-        assert len(attrs_with_prefix(m.relu, "_observer_")) == 0
+        if len(attrs_with_prefix(m.relu, "_observer_")) != 0:
+            raise AssertionError(
+                f"Expected 0 observers, got {len(attrs_with_prefix(m.relu, '_observer_'))}"
+            )
 
         m = torch.jit.script(AddReLUModule())
         qconfig_dict = {"": default_qconfig}
         m = prepare_jit(m, qconfig_dict)
-        assert len(attrs_with_prefix(m, "_observer")) == 3
-        assert len(attrs_with_prefix(m.relu, "_observer")) == 0
+        if len(attrs_with_prefix(m, "_observer")) != 3:
+            raise AssertionError(
+                f"Expected 3 observers, got {len(attrs_with_prefix(m, '_observer'))}"
+            )
+        if len(attrs_with_prefix(m.relu, "_observer")) != 0:
+            raise AssertionError(
+                f"Expected 0 observers, got {len(attrs_with_prefix(m.relu, '_observer'))}"
+            )
         FileCheck().check("aten::add_").check_not(
             'Observer = prim::GetAttr[name="_observer_'
         ).check("ReLU = prim::GetAttr").run(str(get_forward_graph(m._c)))
@@ -635,7 +674,10 @@ def attrs_with_prefix(module, prefix):
         m = torch.jit.script(AddFunctionalReLU())
         qconfig_dict = {"": default_qconfig}
         m = prepare_jit(m, qconfig_dict)
-        assert len(attrs_with_prefix(m, "_observer")) == 3
+        if len(attrs_with_prefix(m, "_observer")) != 3:
+            raise AssertionError(
+                f"Expected 3 observers, got {len(attrs_with_prefix(m, '_observer'))}"
+            )
         FileCheck().check("aten::add_").check_not(
             'Observer = prim::GetAttr[name="_observer_'
         ).check("CallFunction").check('Observer = prim::GetAttr[name="_observer_').run(
@@ -664,11 +706,12 @@ def forward(self, x):
             for x, obs in m.conv._modules._c.items()
             if x.startswith("_observer_")
         }
-        assert len(activation_dtypes) == 1, "Expected to have 1 activation dtype"
-        assert len(weight_dtypes) == 1, "Expected to have 1 weight dtype"
-        assert next(iter(activation_dtypes)) != next(iter(weight_dtypes)), (
-            "Expected activation dtype to "
-        )
+        if len(activation_dtypes) != 1:
+            raise AssertionError("Expected to have 1 activation dtype")
+        if len(weight_dtypes) != 1:
+            raise AssertionError("Expected to have 1 weight dtype")
+        if next(iter(activation_dtypes)) == next(iter(weight_dtypes)):
+            raise AssertionError("Expected activation dtype to ")
         " be different from wegiht dtype"
 
     def test_insert_observers_for_reused_weight(self):
@@ -681,7 +724,10 @@ def forward(self, x, y, weight):
         m = torch.jit.script(M()).eval()
         m = prepare_jit(m, {"": default_qconfig})
         # 3 for x, y, weight, one for output of each F.conv2d and one for output of add
-        assert len(attrs_with_prefix(m, "_observer")) == 6
+        if len(attrs_with_prefix(m, "_observer")) != 6:
+            raise AssertionError(
+                f"Expected 6 observers, got {len(attrs_with_prefix(m, '_observer'))}"
+            )
 
     def test_insert_observers_shared_class_type(self):
         class M(torch.nn.Module):
@@ -700,11 +746,14 @@ def forward(self, x):
         # make sure we didn't quantize the type twice
         conv1_observers = attrs_with_prefix(m.conv1, "_observer_")
         conv2_observers = attrs_with_prefix(m.conv2, "_observer_")
-        assert len(conv1_observers) == 1, "Expected to have 1 observer submodules"
-        assert len(conv2_observers) == 1, "Expected to have 1 observer submodules"
-        assert conv1_observers == conv2_observers, (
-            "Expect conv1 and conv2 to have same observers since the class type is shared"
-        )
+        if len(conv1_observers) != 1:
+            raise AssertionError("Expected to have 1 observer submodules")
+        if len(conv2_observers) != 1:
+            raise AssertionError("Expected to have 1 observer submodules")
+        if conv1_observers != conv2_observers:
+            raise AssertionError(
+                "Expect conv1 and conv2 to have same observers since the class type is shared"
+            )
 
     def test_insert_observers_for_general_ops(self):
         """Make sure we skip observers for ops that doesn't require
@@ -725,7 +774,10 @@ def forward(self, x):
         qconfig_dict = {"": default_qconfig}
         m = prepare_jit(m, qconfig_dict)
         # input and output of conv
-        assert len(attrs_with_prefix(m, "_observer_")) == 2
+        if len(attrs_with_prefix(m, "_observer_")) != 2:
+            raise AssertionError(
+                f"Expected 2 observers, got {len(attrs_with_prefix(m, '_observer_'))}"
+            )
         FileCheck().check('Observer = prim::GetAttr[name="_observer_').check(
             'prim::GetAttr[name="conv"]'
         ).check("prim::CallMethod").check(
@@ -757,7 +809,10 @@ def forward(self, x):
         qconfig_dict = {"": default_qconfig}
         m = prepare_jit(m, qconfig_dict)
         # input and output of conv
-        assert len(attrs_with_prefix(m, "_observer_")) == 3
+        if len(attrs_with_prefix(m, "_observer_")) != 3:
+            raise AssertionError(
+                f"Expected 3 observers, got {len(attrs_with_prefix(m, '_observer_'))}"
+            )
         FileCheck().check('Observer = prim::GetAttr[name="_observer_').check(
             'prim::GetAttr[name="conv1"]'
         ).check("prim::CallMethod").check(
@@ -790,7 +845,10 @@ def forward(self, x):
         qconfig_dict = {"": default_qconfig}
         m = prepare_jit(m, qconfig_dict)
         # input and output of conv
-        assert len(attrs_with_prefix(m, "_observer_")) == 3
+        if len(attrs_with_prefix(m, "_observer_")) != 3:
+            raise AssertionError(
+                f"Expected 3 observers, got {len(attrs_with_prefix(m, '_observer_'))}"
+            )
         FileCheck().check('Observer = prim::GetAttr[name="_observer_').check(
             'prim::GetAttr[name="conv1"]'
         ).check("prim::CallMethod").check(
@@ -836,15 +894,18 @@ def forward(self, x):
         # we want to test that channel_shuffle is going to pass
         # the observed property from the output of conv1 to input of conv2
         # so that we don't insert observers for input of conv2
-        assert (
+        if (
             len(
                 attrs_with_prefix(
                     m,
                     "_observer_",
                 )
             )
-            == 3
-        )
+            != 3
+        ):
+            raise AssertionError(
+                f"Expected 3 observers, got {len(attrs_with_prefix(m, '_observer_'))}"
+            )
 
     def test_insert_observers_for_if(self):
         class QuantProp(torch.nn.Module):
@@ -892,33 +953,45 @@ def forward(self, x):
             else:
                 m = torch.jit.script(M()).eval()
             m = prepare_jit(m, {"": default_qconfig})
-            assert (
+            if (
                 len(
                     attrs_with_prefix(
                         m,
                         "_observer_",
                     )
                 )
-                == result[tracing][0]
-            )
-            assert (
+                != result[tracing][0]
+            ):
+                raise AssertionError(
+                    f"Expected {result[tracing][0]} observers on m, "
+                    f"got {len(attrs_with_prefix(m, '_observer_'))}"
+                )
+            if (
                 len(
                     attrs_with_prefix(
                         m.quant_prop,
                         "_observer_",
                     )
                 )
-                == result[tracing][1]
-            )
-            assert (
+                != result[tracing][1]
+            ):
+                raise AssertionError(
+                    f"Expected {result[tracing][1]} observers on m.quant_prop, "
+                    f"got {len(attrs_with_prefix(m.quant_prop, '_observer_'))}"
+                )
+            if (
                 len(
                     attrs_with_prefix(
                         m.res,
                         "_observer_",
                     )
                 )
-                == result[tracing][2]
-            )
+                != result[tracing][2]
+            ):
+                raise AssertionError(
+                    f"Expected {result[tracing][2]} observers on m.res, "
+                    f"got {len(attrs_with_prefix(m.res, '_observer_'))}"
+                )
 
     def test_insert_observers_for_nested_if(self):
         class Res(torch.nn.Module):
@@ -956,7 +1029,11 @@ def forward(self, x):
             else:
                 m = torch.jit.script(M()).eval()
             m = prepare_jit(m, {"": default_qconfig})
-            assert len(attrs_with_prefix(m, "_observer_")) == result[tracing]
+            if len(attrs_with_prefix(m, "_observer_")) != result[tracing]:
+                raise AssertionError(
+                    f"Expected {result[tracing]} observers, "
+                    f"got {len(attrs_with_prefix(m, '_observer_'))}"
+                )
 
     def test_insert_observers_for_if_consistent_observation(self):
         """check quantization for if works as long as
@@ -1002,7 +1079,10 @@ def forward(self, x):
             else:
                 m = torch.jit.script(M(cond))
             m = prepare_jit(m, {"": default_qconfig})
-            assert len(attrs_with_prefix(m, "_observer_")) == 2
+            if len(attrs_with_prefix(m, "_observer_")) != 2:
+                raise AssertionError(
+                    f"Expected 2 observers, got {len(attrs_with_prefix(m, '_observer_'))}"
+                )
 
         for cond, tracing in options:
             if tracing:
@@ -1011,7 +1091,11 @@ def forward(self, x):
                 m = torch.jit.script(M2(cond))
             m = prepare_jit(m, {"": default_qconfig})
             num_observers = 2 if tracing and not cond else 3
-            assert len(attrs_with_prefix(m, "_observer_")) == num_observers
+            if len(attrs_with_prefix(m, "_observer_")) != num_observers:
+                raise AssertionError(
+                    f"Expected {num_observers} observers, "
+                    f"got {len(attrs_with_prefix(m, '_observer_'))}"
+                )
 
     def test_insert_quant_dequant(self):
         class M(torch.nn.Module):
@@ -1035,9 +1119,8 @@ def forward(self, x):
 
             m(data)
             m = convert_jit(m, debug=True)
-            assert len(m._modules._c.items()) == 1, (
-                "Expected to have single submodule of conv"
-            )
+            if len(m._modules._c.items()) != 1:
+                raise AssertionError("Expected to have single submodule of conv")
             # make sure the quantized model is executable
             m(data)
             quant_func = (
@@ -1068,34 +1151,32 @@ def forward(self, x):
             qconfig_dict = {"": qconfig}
             m = prepare_jit(m, qconfig_dict)
             # observers for input, output and value between conv1/conv2
-            assert len(attrs_with_prefix(m, "_observer_")) == 3, (
-                "Expected to have 3 observers"
-            )
+            if len(attrs_with_prefix(m, "_observer_")) != 3:
+                raise AssertionError("Expected to have 3 observers")
             # observer for weight
-            assert len(attrs_with_prefix(m.conv1, "_observer_")) == 1, (
-                "Expected to have 1 observers"
-            )
+            if len(attrs_with_prefix(m.conv1, "_observer_")) != 1:
+                raise AssertionError("Expected to have 1 observers")
             # observer for weight
-            assert len(attrs_with_prefix(m.conv2, "_observer_")) == 1, (
-                "Expected to have 1 observers"
-            )
+            if len(attrs_with_prefix(m.conv2, "_observer_")) != 1:
+                raise AssertionError("Expected to have 1 observers")
 
             data = torch.randn(1, 3, 10, 10, dtype=torch.float)
             m(data)
             m = convert_jit(m, debug=True)
             m(data)
-            assert m.conv1._c._type() == m.conv2._c._type()
+            if m.conv1._c._type() != m.conv2._c._type():
+                raise AssertionError(
+                    f"Expected conv1 and conv2 to share the same type, "
+                    f"got {m.conv1._c._type()} and {m.conv2._c._type()}"
+                )
 
             # check all observers have been removed
-            assert len(attrs_with_prefix(m, "_observer_")) == 0, (
-                "Expected to have 0 observers"
-            )
-            assert len(attrs_with_prefix(m.conv1, "_observer_")) == 0, (
-                "Expected to have 0 observers"
-            )
-            assert len(attrs_with_prefix(m.conv2, "_observer_")) == 0, (
-                "Expected to have 0 observers"
-            )
+            if len(attrs_with_prefix(m, "_observer_")) != 0:
+                raise AssertionError("Expected to have 0 observers")
+            if len(attrs_with_prefix(m.conv1, "_observer_")) != 0:
+                raise AssertionError("Expected to have 0 observers")
+            if len(attrs_with_prefix(m.conv2, "_observer_")) != 0:
+                raise AssertionError("Expected to have 0 observers")
 
             quant_func = (
                 "aten::quantize_per_channel"
@@ -1127,15 +1208,17 @@ def forward(self, x):
         data = torch.randn((2, 2))
         m = torch.jit.script(M())
         ref_res = m(data)
-        assert (
-            len([x for x, _ in m._modules._c.items() if x.startswith("relu")]) == 1
-        ), "Expected to have 1 relu modules after dedup module uses"
+        if len([x for x, _ in m._modules._c.items() if x.startswith("relu")]) != 1:
+            raise AssertionError(
+                "Expected to have 1 relu modules after dedup module uses"
+            )
         torch._C._jit_pass_dedup_module_uses(m._c)
         m = torch.jit._recursive.wrap_cpp_module(m._c)
         res = m(data)
-        assert (
-            len([x for x, _ in m._modules._c.items() if x.startswith("relu")]) == 2
-        ), "Expected to have 2 relu modules after dedup module uses"
+        if len([x for x, _ in m._modules._c.items() if x.startswith("relu")]) != 2:
+            raise AssertionError(
+                "Expected to have 2 relu modules after dedup module uses"
+            )
         self.assertEqual(res, ref_res)
 
     def test_replicate_dequantize(self):
@@ -1345,7 +1428,10 @@ def forward(self, x: torch.Tensor) -> list[torch.Tensor]:
         qconfig_dict = {"": default_qconfig}
         model = torch.jit.script(ComplexModel()).eval()
         model = prepare_jit(model, qconfig_dict)
-        assert len(attrs_with_prefix(model, "_observer")) == 3
+        if len(attrs_with_prefix(model, "_observer")) != 3:
+            raise AssertionError(
+                f"Expected 3 observers, got {len(attrs_with_prefix(model, '_observer'))}"
+            )
         model(data)
         model = convert_jit(model, debug=False)
         FileCheck().check("quantized::linear").check("quantized::linear").run(
@@ -3016,14 +3102,20 @@ def forward(self, x):
             m = prepare_dynamic_jit(model, {"": qconfig})
 
             # observer for weight
-            assert len(attrs_with_prefix(m.fc, "_observer_")) == 1
+            if len(attrs_with_prefix(m.fc, "_observer_")) != 1:
+                raise AssertionError(
+                    f"Expected 1 observer, got {len(attrs_with_prefix(m.fc, '_observer_'))}"
+                )
 
             if qconfig == float16_dynamic_qconfig:
                 observer_name = 'PlaceholderObserver = prim::GetAttr[name="_observer_'
                 FileCheck().check(observer_name).run(m.fc.graph)
             else:
                 # for input of FC for dynamic quant
-                assert len(attrs_with_prefix(m, "_observer_")) == 1
+                if len(attrs_with_prefix(m, "_observer_")) != 1:
+                    raise AssertionError(
+                        f"Expected 1 observer, got {len(attrs_with_prefix(m, '_observer_'))}"
+                    )
                 observer_name = 'Observer = prim::GetAttr[name="_observer_'
                 FileCheck().check(observer_name).check(
                     'prim::GetAttr[name="fc"]'
@@ -3052,13 +3144,25 @@ def forward(self, x):
         m = prepare_dynamic_jit(m, {"sub.fc": default_dynamic_qconfig})
 
         # input of sub for dynamic quant
-        assert len(attrs_with_prefix(m, "_observer_")) == 1
+        if len(attrs_with_prefix(m, "_observer_")) != 1:
+            raise AssertionError(
+                f"Expected 1 observer, got {len(attrs_with_prefix(m, '_observer_'))}"
+            )
         # not quantized
-        assert len(attrs_with_prefix(m.conv, "_observer_")) == 0
+        if len(attrs_with_prefix(m.conv, "_observer_")) != 0:
+            raise AssertionError(
+                f"Expected 0 observers, got {len(attrs_with_prefix(m.conv, '_observer_'))}"
+            )
         # no observers since we observe in the outer most call site
-        assert len(attrs_with_prefix(m.sub, "_observer_")) == 0
+        if len(attrs_with_prefix(m.sub, "_observer_")) != 0:
+            raise AssertionError(
+                f"Expected 0 observers, got {len(attrs_with_prefix(m.sub, '_observer_'))}"
+            )
         # weight of linear
-        assert len(attrs_with_prefix(m.sub.fc, "_observer_")) == 1
+        if len(attrs_with_prefix(m.sub.fc, "_observer_")) != 1:
+            raise AssertionError(
+                f"Expected 1 observer, got {len(attrs_with_prefix(m.sub.fc, '_observer_'))}"
+            )
         FileCheck().check('prim::GetAttr[name="sub').check("prim::CallMethod").check(
             'Observer = prim::GetAttr[name="_observer_'
         ).check("prim::CallMethod").check_not(
@@ -3084,9 +3188,8 @@ def forward(self, x):
                 else default_dynamic_qconfig
             )
             m = quantize_dynamic_jit(m, {"": qconfig}, debug=True)
-            assert len(m._modules._c.items()) == 2, (
-                "Expected to have two submodule of linear"
-            )
+            if len(m._modules._c.items()) != 2:
+                raise AssertionError("Expected to have two submodule of linear")
 
             wt_quant_func = (
                 "aten::quantize_per_channel"
diff --git a/test/run_test.py b/test/run_test.py
index 820206480be11..dc1203728d3b8 100755
--- a/test/run_test.py
+++ b/test/run_test.py
@@ -21,7 +21,7 @@
 from datetime import datetime
 from importlib.metadata import PackageNotFoundError, version
 from pathlib import Path
-from typing import Any, cast, NamedTuple, Optional, Union
+from typing import Any, cast, NamedTuple
 
 import torch
 import torch.distributed as dist
@@ -30,6 +30,7 @@
     get_report_path,
     IS_CI,
     IS_MACOS,
+    IS_WINDOWS,
     isRocmArchAnyOf,
     retry_shell,
     set_cwd,
@@ -104,6 +105,7 @@ def upload_adhoc_failure_json(*args, **kwargs):
 TEST_CONFIG = os.getenv("TEST_CONFIG", "")
 BUILD_ENVIRONMENT = os.getenv("BUILD_ENVIRONMENT", "")
 RERUN_DISABLED_TESTS = os.getenv("PYTORCH_TEST_RERUN_DISABLED_TESTS", "0") == "1"
+NUM_PYTEST_RERUNS = int(os.getenv("PYTORCH_NUM_PYTEST_RERUNS", "2"))
 DISTRIBUTED_TEST_PREFIX = "distributed"
 INDUCTOR_TEST_PREFIX = "inductor"
 IS_SLOW = "slow" in TEST_CONFIG or "slow" in BUILD_ENVIRONMENT
@@ -196,12 +198,17 @@ def __contains__(self, item):
     "test_jit_legacy",
     "test_cuda_nvml_based_avail",
     "test_jit_cuda_fuser",
+    "distributed/pipelining/test_dtensor_pp_integration",
 ]
 
 # Add architecture-specific blocklist entries
 if TEST_WITH_ROCM and isRocmArchAnyOf(("gfx1100",)):
     # Some autotune tests on gfx1100 are hanging, disable for now
     ROCM_BLOCKLIST.append("inductor/test_max_autotune")
+    # ROCm 7.2 gfx1100 started timing out due to these
+    ROCM_BLOCKLIST.append("inductor/test_torchinductor_dynamic_shapes")
+    ROCM_BLOCKLIST.append("inductor/test_torchinductor_opinfo")
+    ROCM_BLOCKLIST.append("inductor/test_ck_backend")
 
 S390X_BLOCKLIST = [
     # these tests fail due to various reasons
@@ -299,6 +306,7 @@ def __contains__(self, item):
     "test_autograd_fallback",
     "inductor/test_compiler_bisector",
     "test_privateuseone_python_backend",
+    "functorch/test_control_flow_cuda_initialization",
 ] + FSDP_TEST
 
 # Test files that should always be run serially with other test files,
@@ -983,6 +991,15 @@ def test_openreg(test_module, test_directory, options):
     if return_code != 0:
         return return_code
 
+    # Run the openreg C++ unit tests (gtest) built by cmake.
+    ortests_bin = os.path.join(
+        openreg_dir, "build", "third_party", "openreg", "ortests"
+    )
+    if os.path.isfile(ortests_bin):
+        return_code = shell([ortests_bin], cwd=openreg_dir)
+        if return_code != 0:
+            return return_code
+
     with extend_python_path([install_dir]):
         cmd = [
             sys.executable,
@@ -1252,9 +1269,9 @@ def get_pytest_args(options, is_cpp_test=False, is_distributed_test=False):
         # flakiness status. Default to 50 re-runs
         rerun_options = ["--flake-finder", f"--flake-runs={count}"]
     else:
-        # When under the normal mode, retry a failed test 2 more times. -x means stop at the first
-        # failure
-        rerun_options = ["-x", "--reruns=2"]
+        # When under the normal mode, retry a failed test NUM_PYTEST_RERUNS more times.
+        # -x means stop at the first failure. Set PYTORCH_NUM_PYTEST_RERUNS=0 to disable.
+        rerun_options = ["-x", f"--reruns={NUM_PYTEST_RERUNS}"]
 
     pytest_args = [
         "-vv",
@@ -1321,6 +1338,7 @@ def run_ci_sanity_check(test: ShardedTest, test_directory, options):
     "distributed/rpc/test_tensorpipe_agent": run_test_with_subprocess,
     "distributed/rpc/test_share_memory": run_test_with_subprocess,
     "distributed/rpc/cuda/test_tensorpipe_agent": run_test_with_subprocess,
+    "functorch/test_control_flow_cuda_initialization": run_test_with_subprocess,
     "doctests": run_doctests,
     "test_ci_sanity_check_fail": run_ci_sanity_check,
     "test_autoload_enable": test_autoload_enable,
@@ -1621,7 +1639,7 @@ def exclude_tests(
     return selected_tests
 
 
-def must_serial(file: Union[str, ShardedTest]) -> bool:
+def must_serial(file: str | ShardedTest) -> bool:
     if isinstance(file, ShardedTest):
         file = file.name
     return (
@@ -1786,7 +1804,50 @@ def get_selected_tests(options) -> list[str]:
 
     selected_tests = exclude_tests(options.exclude, selected_tests)
 
-    if sys.platform == "win32" and not options.ignore_win_blocklist:
+    if IS_WINDOWS and not options.ignore_win_blocklist:
+        from torch.testing._internal.common_cuda import SM120OrLater, SM89OrLater
+
+        # Disable tests on Windows for SM89 and later - tests failing in ci
+        # Enable tests after fixing the failures
+        if SM89OrLater:
+            WINDOWS_BLOCKLIST.extend(
+                [
+                    # Windows fatal exception / access violation
+                    "functorch/test_aotdispatch",
+                    "functorch/test_control_flow",
+                    "nn/test_convolution",
+                    "profiler/test_profiler",
+                    "test_modules",
+                    "test_expanded_weights",
+                    "test_jit",
+                    "test_nested_tensor",
+                    "test_nestedtensor",
+                    "test_nn",
+                    # DLL load failed errors, missing dependencies
+                    "test_custom_ops",
+                    "test_testing",
+                    # Features not supported on Windows ( e.g. rowwise scaling)
+                    "test_decomp",
+                    "test_transformers",
+                    "test_ops",
+                    # Output mismatch errors and long running tests
+                    "test_linalg",
+                    "test_matmul_cuda",
+                    "functorch/test_ops",
+                    "test_scaled_matmul_cuda",
+                ]
+            )
+
+        # Disable tests on Windows for SM120 and later - tests failing in ci
+        # Enable tests after fixing the failures
+        if SM120OrLater:
+            WINDOWS_BLOCKLIST.extend(
+                [
+                    # test_api fails on Windows SM120+. Triage pending.
+                    "cpp/test_api",
+                ]
+            )
+
         target_arch = os.environ.get("VSCMD_ARG_TGT_ARCH")
         if target_arch != "x64":
             WINDOWS_BLOCKLIST.append("cpp_extensions_aot_no_ninja")
@@ -1847,13 +1908,18 @@ def load_test_times_from_file(file: str) -> dict[str, Any]:
 
     with open(path) as f:
         test_times_file = cast(dict[str, Any], json.load(f))
-    job_name = os.environ.get("JOB_NAME")
+    raw_job_name = os.environ.get("JOB_NAME")
+    build_env = os.environ.get("BUILD_ENVIRONMENT")
+    job_name = raw_job_name
     if job_name is None or job_name == "":
         # If job name isn't available, use build environment as a backup
-        job_name = os.environ.get("BUILD_ENVIRONMENT")
+        job_name = build_env
     else:
         job_name = job_name.split(" / test (")[0]
     test_config = os.environ.get("TEST_CONFIG")
+    print_to_stderr(f"JOB_NAME={raw_job_name}")
+    print_to_stderr(f"BUILD_ENVIRONMENT={build_env}")
+    print_to_stderr(f"test-times lookup key={job_name}, test_config={test_config}")
     if test_config in test_times_file.get(job_name, {}):
         print_to_stderr("Found test times from artifacts")
         return test_times_file[job_name][test_config]
@@ -1927,7 +1993,7 @@ class TestFailure(NamedTuple):
 
 def run_test_module(
     test: ShardedTest, test_directory: str, options
-) -> Optional[TestFailure]:
+) -> TestFailure | None:
     try:
         maybe_set_hip_visible_devies()
 
@@ -1993,7 +2059,7 @@ def run_tests(
         ):
             shutil.copy(os.path.join(test_directory, conftest_file), cpp_file)
 
-    def handle_complete(failure: Optional[TestFailure]):
+    def handle_complete(failure: TestFailure | None):
         failed = failure is not None
         if IS_CI and options.upload_artifacts_while_running:
             parse_xml_and_upload_json()
diff --git a/test/slow_tests.json b/test/slow_tests.json
index abd30b98d47e1..975ad037c53f8 100644
--- a/test/slow_tests.json
+++ b/test/slow_tests.json
@@ -1,252 +1,277 @@
 {
-  "EndToEndLSTM (__main__.RNNTest)": 184.69266764322916,
-  "MultiheadAttention (__main__.ModulesTest)": 137.28699747721353,
-  "test_3mm_add (__main__.TestTritonDotReduction)": 100.44633229573567,
-  "test_AllenaiLongformerBase_repro_cpu_halide (__main__.HalideCpuTests)": 213.74633280436197,
-  "test_StridedShard_to_shard_order (__main__.Test_StridedShard_with_shard_order)": 283.25299580891925,
-  "test__adaptive_avg_pool2d (__main__.CPUReproTests)": 92.95833418104384,
-  "test_adaptive_max_pool2d1_cpu_halide (__main__.HalideCpuTests)": 116.62366739908855,
-  "test_after_aot_cpu_runtime_error (__main__.MinifierIsolateTests)": 62.366333855523,
-  "test_alexnet_prefix_cpu_halide (__main__.HalideCpuTests)": 177.8383331298828,
-  "test_aot_autograd_disable_functionalization_exhaustive_nn_functional_max_pool2d_cpu_float32 (__main__.TestEagerFusionOpInfoCPU)": 63.67699940999349,
-  "test_aot_autograd_disable_functionalization_symbolic_exhaustive_nn_functional_max_pool1d_cpu_float32 (__main__.TestEagerFusionOpInfoCPU)": 127.79400380452473,
-  "test_aot_autograd_disable_functionalization_symbolic_exhaustive_nn_functional_max_pool2d_cpu_float32 (__main__.TestEagerFusionOpInfoCPU)": 208.3969980875651,
-  "test_aot_autograd_disable_functionalization_symbolic_exhaustive_nn_functional_max_pool3d_cpu_float32 (__main__.TestEagerFusionOpInfoCPU)": 124.48566436767578,
-  "test_aot_autograd_disable_functionalization_symbolic_exhaustive_svd_cpu_float32 (__main__.TestEagerFusionOpInfoCPU)": 73.43719082786923,
-  "test_aot_autograd_symbolic_exhaustive_linalg_svd_cpu_float32 (__main__.TestEagerFusionOpInfoCPU)": 71.86966705322266,
-  "test_aot_autograd_symbolic_exhaustive_nn_functional_max_pool1d_cpu_float32 (__main__.TestEagerFusionOpInfoCPU)": 145.66200256347656,
-  "test_aot_autograd_symbolic_exhaustive_nn_functional_max_pool2d_cpu_float32 (__main__.TestEagerFusionOpInfoCPU)": 223.1346689860026,
-  "test_aot_autograd_symbolic_exhaustive_nn_functional_max_pool3d_cpu_float32 (__main__.TestEagerFusionOpInfoCPU)": 136.3513387044271,
-  "test_aot_autograd_symbolic_exhaustive_nn_functional_max_unpool2d_cpu_float32 (__main__.TestEagerFusionOpInfoCPU)": 66.18100102742513,
-  "test_aot_autograd_symbolic_exhaustive_nn_functional_unfold_cpu_float32 (__main__.TestEagerFusionOpInfoCPU)": 61.71500015258789,
-  "test_aot_autograd_symbolic_exhaustive_ormqr_cpu_float32 (__main__.TestEagerFusionOpInfoCPU)": 68.69900004069011,
-  "test_aot_autograd_symbolic_exhaustive_svd_cpu_float32 (__main__.TestEagerFusionOpInfoCPU)": 67.28733571370442,
-  "test_aot_autograd_symbolic_module_exhaustive_nn_TransformerDecoderLayer_cpu_float32 (__main__.TestEagerFusionModuleInfoCPU)": 149.20133209228516,
-  "test_associative_scan_partial_grad_combine_mode_generic_compile_mode_compile_dynamic_shape_reverse_False_cpu (__main__.AssociativeScanTests)": 91.24416859944661,
-  "test_associative_scan_partial_grad_combine_mode_generic_compile_mode_compile_dynamic_shape_reverse_True_cpu (__main__.AssociativeScanTests)": 108.57016626993816,
-  "test_avg_pool3d_backward2_cpu (__main__.CpuTests)": 616.325,
-  "test_avg_pool3d_backward2_cpu (__main__.CpuTritonTests)": 272.4786682128906,
-  "test_avg_pool3d_backward2_cuda (__main__.GPUTests)": 156.27499934605189,
-  "test_avg_pool3d_backward2_dynamic_shapes_cpu (__main__.DynamicShapesCodegenCpuTests)": 529.6212259928385,
-  "test_avg_pool3d_backward2_dynamic_shapes_cpu (__main__.DynamicShapesCpuTests)": 542.7406548394097,
-  "test_avg_pool3d_backward2_dynamic_shapes_cuda (__main__.DynamicShapesCodegenGPUTests)": 64.20183245340984,
-  "test_avg_pool3d_backward2_dynamic_shapes_cuda (__main__.DynamicShapesGPUTests)": 124.58166758219402,
-  "test_avg_pool3d_backward_cpu_halide (__main__.HalideCpuTests)": 63.05166753133138,
-  "test_backward_nn_functional_multi_head_attention_forward_cpu_float32 (__main__.TestCompositeComplianceCPU)": 94.64399973551433,
-  "test_backward_nn_functional_multi_head_attention_forward_cuda_float32 (__main__.TestCompositeComplianceCUDA)": 66.33616765340169,
-  "test_basic_cpu (__main__.EfficientConvBNEvalCpuTests)": 218.5576663547092,
-  "test_basic_cuda (__main__.EfficientConvBNEvalGpuTests)": 168.72183481852213,
-  "test_cat_2k_args (__main__.TestTEFuserStatic)": 64.58066653770705,
-  "test_checkpointing_without_reentrant_input_requires_grad_False (__main__.TestAutogradWithCompiledAutograd)": 331.3285556369358,
-  "test_checkpointing_without_reentrant_input_requires_grad_True (__main__.TestAutogradWithCompiledAutograd)": 418.45111083984375,
-  "test_collect_callgrind (__main__.TestBenchmarkUtils)": 317.6703355577257,
-  "test_comprehensive_diff_cuda_complex128 (__main__.TestDecompCUDA)": 105.86016591389973,
-  "test_comprehensive_diff_cuda_complex64 (__main__.TestDecompCUDA)": 101.96483357747395,
-  "test_comprehensive_diff_cuda_float32 (__main__.TestDecompCUDA)": 67.77633285522461,
-  "test_comprehensive_diff_cuda_float64 (__main__.TestDecompCUDA)": 69.17333221435547,
-  "test_comprehensive_grid_sampler_2d_cpu_bfloat16 (__main__.TestDecompCPU)": 102.28200022379558,
-  "test_comprehensive_grid_sampler_2d_cpu_float16 (__main__.TestDecompCPU)": 93.95999908447266,
-  "test_comprehensive_grid_sampler_2d_cpu_float32 (__main__.TestDecompCPU)": 430.5759989420573,
-  "test_comprehensive_grid_sampler_2d_cpu_float64 (__main__.TestDecompCPU)": 449.9323323567708,
-  "test_comprehensive_grid_sampler_2d_cuda_bfloat16 (__main__.TestDecompCUDA)": 253.6421686808268,
-  "test_comprehensive_grid_sampler_2d_cuda_float16 (__main__.TestDecompCUDA)": 271.6805013020833,
-  "test_comprehensive_grid_sampler_2d_cuda_float32 (__main__.TestDecompCUDA)": 1271.2601725260417,
-  "test_comprehensive_grid_sampler_2d_cuda_float32 (__main__.TestInductorOpInfoCUDA)": 77.94216664632161,
-  "test_comprehensive_grid_sampler_2d_cuda_float64 (__main__.TestDecompCUDA)": 1118.7591756184895,
-  "test_comprehensive_grid_sampler_2d_cuda_float64 (__main__.TestInductorOpInfoCUDA)": 82.65366744995117,
-  "test_comprehensive_linalg_lu_solve_cuda_float32 (__main__.TestInductorOpInfoCUDA)": 74.36716715494792,
-  "test_comprehensive_linalg_lu_solve_cuda_float64 (__main__.TestInductorOpInfoCUDA)": 71.19016647338867,
-  "test_comprehensive_linalg_solve_triangular_cuda_float32 (__main__.TestInductorOpInfoCUDA)": 73.25883356730144,
-  "test_comprehensive_linalg_solve_triangular_cuda_float64 (__main__.TestInductorOpInfoCUDA)": 75.40216700236003,
-  "test_comprehensive_linalg_svd_cuda_complex128 (__main__.TestDecompCUDA)": 66.99350102742513,
-  "test_comprehensive_linalg_svd_cuda_complex64 (__main__.TestDecompCUDA)": 69.02083269755046,
-  "test_comprehensive_logspace_cuda_float32 (__main__.TestInductorOpInfoCUDA)": 60.92988925509982,
-  "test_comprehensive_masked_norm_cuda_float16 (__main__.TestInductorOpInfoCUDA)": 113.01116689046223,
-  "test_comprehensive_masked_norm_cuda_float32 (__main__.TestInductorOpInfoCUDA)": 107.1844991048177,
-  "test_comprehensive_masked_norm_cuda_float64 (__main__.TestInductorOpInfoCUDA)": 108.23550160725911,
-  "test_comprehensive_nn_functional_conv_transpose3d_cuda_complex128 (__main__.TestDecompCUDA)": 65.7233346303304,
-  "test_comprehensive_nn_functional_conv_transpose3d_cuda_complex64 (__main__.TestDecompCUDA)": 64.3291670481364,
-  "test_comprehensive_nn_functional_gaussian_nll_loss_cpu_float32 (__main__.TestDecompCPU)": 62.670667012532554,
-  "test_comprehensive_nn_functional_gaussian_nll_loss_cpu_float64 (__main__.TestDecompCPU)": 62.265000661214195,
-  "test_comprehensive_nn_functional_gaussian_nll_loss_cuda_float32 (__main__.TestDecompCUDA)": 111.73116556803386,
-  "test_comprehensive_nn_functional_gaussian_nll_loss_cuda_float64 (__main__.TestDecompCUDA)": 113.69033304850261,
-  "test_comprehensive_nn_functional_grid_sample_cpu_float32 (__main__.TestDecompCPU)": 117.5066655476888,
-  "test_comprehensive_nn_functional_grid_sample_cpu_float64 (__main__.TestDecompCPU)": 118.56566619873047,
-  "test_comprehensive_nn_functional_grid_sample_cuda_bfloat16 (__main__.TestDecompCUDA)": 61.97433344523112,
-  "test_comprehensive_nn_functional_grid_sample_cuda_float32 (__main__.TestDecompCUDA)": 271.9861602783203,
-  "test_comprehensive_nn_functional_grid_sample_cuda_float64 (__main__.TestDecompCUDA)": 274.42466735839844,
-  "test_comprehensive_nn_functional_interpolate_bicubic_cuda_float32 (__main__.TestDecompCUDA)": 82.77866490681966,
-  "test_comprehensive_nn_functional_interpolate_bicubic_cuda_float32 (__main__.TestInductorOpInfoCUDA)": 84.92516581217448,
-  "test_comprehensive_nn_functional_interpolate_bicubic_cuda_float64 (__main__.TestDecompCUDA)": 82.4509989420573,
-  "test_comprehensive_nn_functional_interpolate_bicubic_cuda_float64 (__main__.TestInductorOpInfoCUDA)": 86.13049952189128,
-  "test_comprehensive_nn_functional_interpolate_trilinear_cuda_float32 (__main__.TestDecompCUDA)": 125.71083450317383,
-  "test_comprehensive_nn_functional_interpolate_trilinear_cuda_float64 (__main__.TestDecompCUDA)": 121.97066752115886,
-  "test_comprehensive_nn_functional_max_pool2d_cuda_float16 (__main__.TestInductorOpInfoCUDA)": 1262.5691528320312,
-  "test_comprehensive_nn_functional_max_pool2d_cuda_float32 (__main__.TestInductorOpInfoCUDA)": 1252.483642578125,
-  "test_comprehensive_nn_functional_max_pool2d_cuda_float64 (__main__.TestInductorOpInfoCUDA)": 1260.509501139323,
-  "test_comprehensive_nn_functional_max_pool3d_cuda_float32 (__main__.TestInductorOpInfoCUDA)": 604.4365030924479,
-  "test_comprehensive_nn_functional_max_pool3d_cuda_float64 (__main__.TestInductorOpInfoCUDA)": 614.1001586914062,
-  "test_comprehensive_nn_functional_max_unpool2d_cuda_float16 (__main__.TestInductorOpInfoCUDA)": 67.40933354695638,
-  "test_comprehensive_nn_functional_max_unpool2d_cuda_float32 (__main__.TestInductorOpInfoCUDA)": 66.7336654663086,
-  "test_comprehensive_nn_functional_max_unpool2d_cuda_float64 (__main__.TestInductorOpInfoCUDA)": 66.91466649373372,
-  "test_comprehensive_ormqr_cpu_complex64 (__main__.TestDecompCPU)": 65.80933380126953,
-  "test_comprehensive_ormqr_cuda_complex128 (__main__.TestDecompCUDA)": 117.8523343404134,
-  "test_comprehensive_ormqr_cuda_complex64 (__main__.TestDecompCUDA)": 145.19499842325845,
-  "test_comprehensive_ormqr_cuda_float32 (__main__.TestDecompCUDA)": 73.76366806030273,
-  "test_comprehensive_ormqr_cuda_float32 (__main__.TestInductorOpInfoCUDA)": 78.31683349609375,
-  "test_comprehensive_ormqr_cuda_float64 (__main__.TestDecompCUDA)": 76.74333190917969,
-  "test_comprehensive_ormqr_cuda_float64 (__main__.TestInductorOpInfoCUDA)": 70.06755489773221,
-  "test_comprehensive_svd_cuda_complex128 (__main__.TestDecompCUDA)": 80.16916529337566,
-  "test_comprehensive_svd_cuda_complex64 (__main__.TestDecompCUDA)": 76.32833353678386,
-  "test_comprehensive_svd_lowrank_cuda_complex128 (__main__.TestDecompCUDA)": 65.76883252461751,
-  "test_comprehensive_svd_lowrank_cuda_complex64 (__main__.TestDecompCUDA)": 67.16433334350586,
-  "test_constructor_autograd_SparseBSC_cuda (__main__.TestSparseAnyCUDA)": 88.39966583251953,
-  "test_constructor_autograd_SparseBSR_cuda (__main__.TestSparseAnyCUDA)": 83.59666442871094,
-  "test_conv1d_basic (__main__.TestXNNPACKConv1dTransformPass)": 83.94522179497613,
-  "test_conv1d_with_relu_fc (__main__.TestXNNPACKConv1dTransformPass)": 187.1534457736545,
-  "test_conv2d_binary_broadcast_shapes_cpu (__main__.TestPatternMatcherGenericCPU)": 91.1510009765625,
-  "test_conv3d_binary_broadcast_shapes_cpu (__main__.TestPatternMatcherGenericCPU)": 153.24633280436197,
-  "test_conv3d_binary_dynamic_shapes_cpu (__main__.TestDynamicPatternMatcherGenericCPU)": 62.278334299723305,
-  "test_conv_unary_fusion_nnc (__main__.TestMkldnnFusion)": 82.22033437093098,
-  "test_correctness_AdamW_use_closure_True_cuda_float32 (__main__.CompiledOptimizerParityTestsCUDA)": 86.71616617838542,
-  "test_correctness_Adam_use_closure_True_cuda_float32 (__main__.CompiledOptimizerParityTestsCUDA)": 83.59499867757161,
-  "test_count_nonzero_all (__main__.TestBool)": 676.7384372287327,
-  "test_ddp_uneven_inputs (__main__.TestDistBackendWithSpawn)": 448.8698336283366,
-  "test_diff_hyperparams_sharding_strategy_str_full_shard (__main__.TestFSDPUseOrigParamsMultipleParamGroups)": 62.429334004720054,
-  "test_diff_hyperparams_sharding_strategy_str_no_shard (__main__.TestFSDPUseOrigParamsMultipleParamGroups)": 61.66200065612793,
-  "test_dispatch_symbolic_meta_outplace_all_strides_nn_functional_gaussian_nll_loss_cuda_float32 (__main__.TestMetaCUDA)": 89.41600036621094,
-  "test_dtensor_op_db_nn_functional_gaussian_nll_loss_cpu_float32 (__main__.TestLocalDTensorOpsCPU)": 1525.8353271484375,
-  "test_dtensor_op_db_nn_functional_gaussian_nll_loss_cpu_float32 (__main__.TestMultiThreadedDTensorOpsCPU)": 90.60733540852864,
-  "test_error_detection_and_propagation (__main__.NcclErrorHandlingTest)": 66.99349848429362,
-  "test_fail_arithmetic_ops.py (__main__.TestTyping)": 71.01377783881293,
-  "test_fail_creation_ops.py (__main__.TestTyping)": 205.9894081398293,
-  "test_fn_fwgrad_bwgrad_cumprod_cuda_complex128 (__main__.TestFwdGradientsCUDA)": 79.61499913533528,
-  "test_fn_gradgrad_cumprod_cuda_complex128 (__main__.TestBwdGradientsCUDA)": 92.82350031534831,
-  "test_fuse_large_params_cpu (__main__.CpuTests)": 139.3203353881836,
-  "test_fuse_large_params_cuda (__main__.GPUTests)": 63.387572152274,
-  "test_fuse_large_params_dynamic_shapes_cpu (__main__.DynamicShapesCodegenCpuTests)": 171.7693345811632,
-  "test_fuse_large_params_dynamic_shapes_cpu (__main__.DynamicShapesCpuTests)": 168.58877902560764,
-  "test_fuse_large_params_dynamic_shapes_cuda (__main__.DynamicShapesCodegenGPUTests)": 126.14983367919922,
-  "test_fuse_large_params_dynamic_shapes_cuda (__main__.DynamicShapesGPUTests)": 120.24416605631511,
-  "test_grad_nn_Transformer_cuda_float64 (__main__.TestModuleCUDA)": 81.49983342488606,
-  "test_gradgrad_nn_TransformerDecoderLayer_cuda_float64 (__main__.TestModuleCUDA)": 192.15400187174478,
-  "test_gradgrad_nn_TransformerEncoder_eval_mode_cuda_float64 (__main__.TestModuleCUDA)": 113.96200052897136,
-  "test_gradgrad_nn_TransformerEncoder_train_mode_cuda_float64 (__main__.TestModuleCUDA)": 125.24116770426433,
-  "test_gradgrad_nn_Transformer_cuda_float64 (__main__.TestModuleCUDA)": 510.77049255371094,
-  "test_grid_sampler_2d_cpu_halide (__main__.HalideCpuTests)": 194.58466593424478,
-  "test_indirect_device_assert (__main__.TritonCodeGenTests)": 307.21116638183594,
-  "test_inductor_no_recursionerror_on_for_loops (__main__.ReproTests)": 65.38683425055609,
-  "test_inductor_no_recursionerror_on_for_loops_dynamic_shapes (__main__.DynamicShapesReproTests)": 70.7237786187066,
-  "test_inplace_gradgrad_cumprod_cuda_complex128 (__main__.TestBwdGradientsCUDA)": 91.1038335164388,
-  "test_inputs_overlapping_with_mutation_stress_dynamic_shapes (__main__.DynamicShapesAotAutogradFallbackTests)": 132.38899824354382,
-  "test_jit_cuda_archflags (__main__.TestCppExtensionJIT)": 78.77233250935872,
-  "test_linear_binary_cpp_wrapper (__main__.TestCppWrapper)": 120.73766835530598,
-  "test_linear_binary_dynamic_shapes_cpp_wrapper (__main__.DynamicShapesCppWrapperCpuTests)": 124.48300170898438,
-  "test_longformer_chunk_dynamic_shapes (__main__.DynamicShapesReproTests)": 104.92066701253255,
-  "test_lstm_cpu (__main__.TestMkldnnCPU)": 86.04666392008464,
-  "test_many_overlapping_inputs_does_not_explode_guards_dynamic_shapes (__main__.DynamicShapesReproTests)": 136.2572224934896,
-  "test_max_pool2d2_cpu_halide (__main__.HalideCpuTests)": 417.986328125,
-  "test_max_pool2d3_cpu_halide (__main__.HalideCpuTests)": 134.87866719563803,
-  "test_max_pool2d5_cpu_halide (__main__.HalideCpuTests)": 357.8183288574219,
-  "test_max_pool2d_with_indices_backward4_dynamic_shapes_cpu (__main__.DynamicShapesCodegenCpuTests)": 64.54022216796875,
-  "test_max_pool2d_with_indices_backward4_dynamic_shapes_cpu (__main__.DynamicShapesCpuTests)": 66.25211164686415,
-  "test_ordered_distribute_all_combination (__main__.DistributeWithDeviceOrderTest)": 158.06944274902344,
-  "test_ordered_distribute_all_combination (__main__.DistributeWithDeviceOrderTestWithLocalTensor)": 78.09822167290582,
-  "test_ordered_redistribute_with_partial (__main__.DistributeWithDeviceOrderTest)": 177.78066677517361,
-  "test_ordered_redistribute_with_partial (__main__.DistributeWithDeviceOrderTestWithLocalTensor)": 621.3526577419705,
-  "test_proper_exit (__main__.TestDataLoader)": 218.22477637396918,
-  "test_proper_exit (__main__.TestDataLoaderPersistentWorkers)": 219.7027791341146,
-  "test_python_ref__refs_special_zeta_cuda_float64 (__main__.TestCommonCUDA)": 67.61199951171875,
-  "test_python_ref_executor__refs_special_zeta_executor_aten_cuda_float64 (__main__.TestCommonCUDA)": 82.16616694132487,
-  "test_qlinear_add_int8_mixed_bf16_use_relu_False_is_qat_False_is_dynamic_False (__main__.TestPatternMatcher)": 95.37699890136719,
-  "test_qlinear_add_int8_mixed_bf16_use_relu_False_is_qat_False_is_dynamic_False_cpp_wrapper (__main__.TestCppWrapper)": 106.87333424886067,
-  "test_qlinear_add_int8_mixed_bf16_use_relu_False_is_qat_False_is_dynamic_False_dynamic_shapes_cpp_wrapper (__main__.DynamicShapesCppWrapperCpuTests)": 110.58233388264973,
-  "test_qlinear_add_int8_mixed_bf16_use_relu_False_is_qat_False_is_dynamic_True (__main__.TestPatternMatcher)": 93.18999989827473,
-  "test_qlinear_add_int8_mixed_bf16_use_relu_False_is_qat_False_is_dynamic_True_cpp_wrapper (__main__.TestCppWrapper)": 117.0056660970052,
-  "test_qlinear_add_int8_mixed_bf16_use_relu_False_is_qat_False_is_dynamic_True_dynamic_shapes_cpp_wrapper (__main__.DynamicShapesCppWrapperCpuTests)": 112.08600107828777,
-  "test_qlinear_add_int8_mixed_bf16_use_relu_False_is_qat_True_is_dynamic_False (__main__.TestPatternMatcher)": 90.89966583251953,
-  "test_qlinear_add_int8_mixed_bf16_use_relu_False_is_qat_True_is_dynamic_False_cpp_wrapper (__main__.TestCppWrapper)": 99.80333201090495,
-  "test_qlinear_add_int8_mixed_bf16_use_relu_False_is_qat_True_is_dynamic_False_dynamic_shapes_cpp_wrapper (__main__.DynamicShapesCppWrapperCpuTests)": 101.7693354288737,
-  "test_qlinear_add_int8_mixed_bf16_use_relu_False_is_qat_True_is_dynamic_True (__main__.TestPatternMatcher)": 106.9923324584961,
-  "test_qlinear_add_int8_mixed_bf16_use_relu_False_is_qat_True_is_dynamic_True_cpp_wrapper (__main__.TestCppWrapper)": 104.36800130208333,
-  "test_qlinear_add_int8_mixed_bf16_use_relu_False_is_qat_True_is_dynamic_True_dynamic_shapes_cpp_wrapper (__main__.DynamicShapesCppWrapperCpuTests)": 106.9153340657552,
-  "test_qlinear_add_int8_mixed_bf16_use_relu_True_is_qat_False_is_dynamic_False (__main__.TestPatternMatcher)": 100.8326644897461,
-  "test_qlinear_add_int8_mixed_bf16_use_relu_True_is_qat_False_is_dynamic_False_cpp_wrapper (__main__.TestCppWrapper)": 107.92199961344402,
-  "test_qlinear_add_int8_mixed_bf16_use_relu_True_is_qat_False_is_dynamic_False_dynamic_shapes_cpp_wrapper (__main__.DynamicShapesCppWrapperCpuTests)": 104.72100067138672,
-  "test_qlinear_add_int8_mixed_bf16_use_relu_True_is_qat_False_is_dynamic_True (__main__.TestPatternMatcher)": 93.16766611735027,
-  "test_qlinear_add_int8_mixed_bf16_use_relu_True_is_qat_False_is_dynamic_True_cpp_wrapper (__main__.TestCppWrapper)": 111.56466674804688,
-  "test_qlinear_add_int8_mixed_bf16_use_relu_True_is_qat_False_is_dynamic_True_dynamic_shapes_cpp_wrapper (__main__.DynamicShapesCppWrapperCpuTests)": 106.9739990234375,
-  "test_qlinear_add_int8_mixed_bf16_use_relu_True_is_qat_True_is_dynamic_False (__main__.TestPatternMatcher)": 88.35433451334636,
-  "test_qlinear_add_int8_mixed_bf16_use_relu_True_is_qat_True_is_dynamic_False_cpp_wrapper (__main__.TestCppWrapper)": 100.52966817220052,
-  "test_qlinear_add_int8_mixed_bf16_use_relu_True_is_qat_True_is_dynamic_False_dynamic_shapes_cpp_wrapper (__main__.DynamicShapesCppWrapperCpuTests)": 102.96500142415364,
-  "test_qlinear_add_int8_mixed_bf16_use_relu_True_is_qat_True_is_dynamic_True (__main__.TestPatternMatcher)": 98.01133473714192,
-  "test_qlinear_add_int8_mixed_bf16_use_relu_True_is_qat_True_is_dynamic_True_cpp_wrapper (__main__.TestCppWrapper)": 112.06033325195312,
-  "test_qlinear_add_int8_mixed_bf16_use_relu_True_is_qat_True_is_dynamic_True_dynamic_shapes_cpp_wrapper (__main__.DynamicShapesCppWrapperCpuTests)": 106.11466725667317,
-  "test_quick_core_backward__unsafe_masked_index_cpu_float64 (__main__.TestDecompCPU)": 616.5533447265625,
-  "test_quick_core_backward__unsafe_masked_index_cuda_float64 (__main__.TestDecompCUDA)": 1048.9943033854167,
-  "test_quick_core_backward__unsafe_masked_index_put_accumulate_cpu_float64 (__main__.TestDecompCPU)": 770.6356608072916,
-  "test_quick_core_backward__unsafe_masked_index_put_accumulate_cuda_float64 (__main__.TestDecompCUDA)": 1446.0608520507812,
-  "test_quick_core_backward_nn_functional_max_unpool3d_grad_cpu_float64 (__main__.TestDecompCPU)": 85.64499918619792,
-  "test_quick_core_backward_nn_functional_max_unpool3d_grad_cuda_float64 (__main__.TestDecompCUDA)": 278.70049540201825,
-  "test_quick_core_backward_roll_cpu_float64 (__main__.TestDecompCPU)": 127.52366892496745,
-  "test_quick_core_backward_roll_cuda_float64 (__main__.TestDecompCUDA)": 237.04200236002603,
-  "test_quick_core_backward_select_scatter_cpu_float64 (__main__.TestDecompCPU)": 70.95300038655598,
-  "test_quick_core_backward_select_scatter_cuda_float64 (__main__.TestDecompCUDA)": 146.1635004679362,
-  "test_quick_core_backward_split_cuda_float64 (__main__.TestDecompCUDA)": 72.04449971516927,
-  "test_quick_core_backward_split_with_sizes_copy_cpu_float64 (__main__.TestDecompCPU)": 99.4903335571289,
-  "test_quick_core_backward_split_with_sizes_copy_cuda_float64 (__main__.TestDecompCUDA)": 177.31633504231772,
-  "test_quick_core_backward_std_cpu_float64 (__main__.TestDecompCPU)": 72.39499918619792,
-  "test_quick_core_backward_std_cuda_float64 (__main__.TestDecompCUDA)": 142.12999725341797,
-  "test_register_spills_cuda (__main__.BenchmarkFusionGpuTest)": 114.4384994506836,
-  "test_run2run_determinism_model_name_BertForMaskedLM_training_or_inference_inference_precision_amp (__main__.DeterministicTest)": 70.59966659545898,
-  "test_run2run_determinism_model_name_BertForMaskedLM_training_or_inference_inference_precision_bfloat16 (__main__.DeterministicTest)": 61.23133405049642,
-  "test_run2run_determinism_model_name_BertForMaskedLM_training_or_inference_training_precision_amp (__main__.DeterministicTest)": 152.28532918294272,
-  "test_run2run_determinism_model_name_BertForMaskedLM_training_or_inference_training_precision_bfloat16 (__main__.DeterministicTest)": 151.11583201090494,
-  "test_run2run_determinism_model_name_BertForMaskedLM_training_or_inference_training_precision_float16 (__main__.DeterministicTest)": 170.14933268229166,
-  "test_run2run_determinism_model_name_BertForMaskedLM_training_or_inference_training_precision_float32 (__main__.DeterministicTest)": 141.25400034586588,
-  "test_run2run_determinism_model_name_DistillGPT2_training_or_inference_training_precision_amp (__main__.DeterministicTest)": 78.44533411661784,
-  "test_run2run_determinism_model_name_DistillGPT2_training_or_inference_training_precision_bfloat16 (__main__.DeterministicTest)": 91.19933319091797,
-  "test_run2run_determinism_model_name_DistillGPT2_training_or_inference_training_precision_float16 (__main__.DeterministicTest)": 80.00733184814453,
-  "test_run2run_determinism_model_name_DistillGPT2_training_or_inference_training_precision_float32 (__main__.DeterministicTest)": 76.68100102742513,
-  "test_run2run_determinism_model_name_GoogleFnet_training_or_inference_training_precision_amp (__main__.DeterministicTest)": 92.68900044759114,
-  "test_run2run_determinism_model_name_GoogleFnet_training_or_inference_training_precision_bfloat16 (__main__.DeterministicTest)": 92.34583409627278,
-  "test_run2run_determinism_model_name_GoogleFnet_training_or_inference_training_precision_float16 (__main__.DeterministicTest)": 92.69583257039388,
-  "test_run2run_determinism_model_name_GoogleFnet_training_or_inference_training_precision_float32 (__main__.DeterministicTest)": 93.89749908447266,
-  "test_runtime_checks_large_cpu (__main__.AOTInductorTestABICompatibleCpu)": 65.03716723124187,
-  "test_runtime_checks_large_cpu_with_stack_allocation (__main__.AOTInductorTestABICompatibleCpuWithStackAllocation)": 69.72733306884766,
-  "test_runtime_checks_large_cuda (__main__.AOTInductorTestABICompatibleGpu)": 171.71016947428384,
-  "test_sdpa_kernel_ctx_manager2_dynamic_shapes (__main__.DynamicShapesCtxManagerTests)": 107.34944322374132,
-  "test_shuffler_iterdatapipe (__main__.IntegrationTestDataLoaderDataPipe)": 129.4277776082357,
-  "test_slow_tasks (__main__.TestFunctionalAutogradBenchmark)": 118.2247797648112,
-  "test_sort_dynamic_shape_with_check_cuda (__main__.TestInductorDynamicCUDA)": 108.49850209554036,
-  "test_sort_stable_cpu (__main__.CpuTritonTests)": 1305.087646484375,
-  "test_sort_stable_cuda (__main__.GPUTests)": 102.32514299665179,
-  "test_split_cumsum_cpu (__main__.CpuTritonTests)": 90.37133534749348,
-  "test_svd_lowrank_cuda_complex128 (__main__.TestLinalgCUDA)": 285.8221740722656,
-  "test_tensor_split (__main__.TestVmapOperators)": 91.96030240144694,
-  "test_terminate_handler_on_crash (__main__.TestTorch)": 207.283443874783,
-  "test_terminate_signal (__main__.ForkTest)": 234.09955173068576,
-  "test_terminate_signal (__main__.ParallelForkServerShouldWorkTest)": 234.16177961561414,
-  "test_terminate_signal (__main__.SpawnTest)": 237.5818862915039,
-  "test_train_parity_with_activation_checkpointing (__main__.TestFullyShard1DTrainingCompose)": 64.86788728502061,
-  "test_triton_bsr_scatter_mm_blocksize_64_cuda_float32 (__main__.TestSparseCompressedTritonKernelsCUDA)": 88.8943322499593,
-  "test_triton_bsr_softmax_cuda_bfloat16 (__main__.TestSparseCompressedTritonKernelsCUDA)": 215.28750356038412,
-  "test_triton_bsr_softmax_cuda_float16 (__main__.TestSparseCompressedTritonKernelsCUDA)": 214.9933344523112,
-  "test_triton_bsr_softmax_cuda_float32 (__main__.TestSparseCompressedTritonKernelsCUDA)": 150.96166483561197,
-  "test_upsample_bicubic2d_cpu_halide (__main__.HalideCpuTests)": 97.6326675415039,
-  "test_variant_consistency_jit_nn_functional_max_pool2d_cpu_float32 (__main__.TestJitCPU)": 86.5586675008138,
-  "test_variant_consistency_jit_nn_functional_max_pool2d_cuda_float32 (__main__.TestJitCUDA)": 76.31433232625325,
-  "test_vec_compare_op_cpu_only (__main__.CPUReproTests)": 64.15555487738715,
-  "test_vmapjvpvjp_linalg_lstsq_grad_oriented_cpu_float32 (__main__.TestOperatorsCPU)": 86.58700052897136,
-  "test_vmapjvpvjp_linalg_lu_solve_cuda_float32 (__main__.TestOperatorsCUDA)": 70.91766738891602,
-  "test_vmapjvpvjp_max_pool2d_with_indices_backward_cpu_float32 (__main__.TestOperatorsCPU)": 64.271666208903,
-  "test_vmapjvpvjp_max_pool2d_with_indices_backward_cuda_float32 (__main__.TestOperatorsCUDA)": 81.8783327738444,
-  "test_vmapjvpvjp_nn_functional_conv2d_cpu_float32 (__main__.TestOperatorsCPU)": 65.68414279392788,
-  "test_vmapjvpvjp_nn_functional_max_pool2d_cpu_float32 (__main__.TestOperatorsCPU)": 70.83733367919922,
-  "test_vmapjvpvjp_nn_functional_max_pool2d_cuda_float32 (__main__.TestOperatorsCUDA)": 77.18000030517578,
-  "test_vmapjvpvjp_unbind_cpu_float32 (__main__.TestOperatorsCPU)": 67.37900034586589,
-  "test_vmapjvpvjp_unbind_cuda_float32 (__main__.TestOperatorsCUDA)": 80.73283386230469,
-  "test_vmapvjpvjp_meshgrid_list_of_tensors_cuda_float32 (__main__.TestOperatorsCUDA)": 90.59783426920573,
-  "test_vmapvjpvjp_meshgrid_variadic_tensors_cuda_float32 (__main__.TestOperatorsCUDA)": 83.13833363850911,
-  "test_vmapvjpvjp_nn_functional_bilinear_cuda_float32 (__main__.TestOperatorsCUDA)": 112.93466695149739
+  "EndToEndLSTM (__main__.RNNTest)": 193.73400370279947,
+  "MultiheadAttention (__main__.ModulesTest)": 144.85599772135416,
+  "test_3mm_add (__main__.TestTritonDotReduction)": 129.92483266194662,
+  "test__adaptive_avg_pool2d (__main__.CPUReproTests)": 90.46566390991211,
+  "test_aot_autograd_disable_functionalization_exhaustive_nn_functional_max_pool2d_cpu_float32 (__main__.TestEagerFusionOpInfoCPU)": 77.86966705322266,
+  "test_aot_autograd_disable_functionalization_symbolic_exhaustive_nn_functional_max_pool1d_cpu_float32 (__main__.TestEagerFusionOpInfoCPU)": 127.5673319498698,
+  "test_aot_autograd_disable_functionalization_symbolic_exhaustive_nn_functional_max_pool2d_cpu_float32 (__main__.TestEagerFusionOpInfoCPU)": 204.36866760253906,
+  "test_aot_autograd_disable_functionalization_symbolic_exhaustive_nn_functional_max_pool3d_cpu_float32 (__main__.TestEagerFusionOpInfoCPU)": 123.83066304524739,
+  "test_aot_autograd_disable_functionalization_symbolic_exhaustive_svd_cpu_float32 (__main__.TestEagerFusionOpInfoCPU)": 66.83466720581055,
+  "test_aot_autograd_exhaustive_nn_functional_max_pool2d_cpu_float32 (__main__.TestEagerFusionOpInfoCPU)": 72.18625852796767,
+  "test_aot_autograd_symbolic_exhaustive_linalg_svd_cpu_float32 (__main__.TestEagerFusionOpInfoCPU)": 77.01166534423828,
+  "test_aot_autograd_symbolic_exhaustive_masked_norm_cpu_float32 (__main__.TestEagerFusionOpInfoCPU)": 62.44844436645508,
+  "test_aot_autograd_symbolic_exhaustive_nn_functional_max_pool1d_cpu_float32 (__main__.TestEagerFusionOpInfoCPU)": 156.3863321940104,
+  "test_aot_autograd_symbolic_exhaustive_nn_functional_max_pool2d_cpu_float32 (__main__.TestEagerFusionOpInfoCPU)": 203.4173329671224,
+  "test_aot_autograd_symbolic_exhaustive_nn_functional_max_pool3d_cpu_float32 (__main__.TestEagerFusionOpInfoCPU)": 105.20900217692058,
+  "test_aot_autograd_symbolic_exhaustive_nn_functional_unfold_cpu_float32 (__main__.TestEagerFusionOpInfoCPU)": 65.38600158691406,
+  "test_aot_autograd_symbolic_exhaustive_ormqr_cpu_float32 (__main__.TestEagerFusionOpInfoCPU)": 83.66033426920573,
+  "test_aot_autograd_symbolic_exhaustive_svd_cpu_float32 (__main__.TestEagerFusionOpInfoCPU)": 101.71733093261719,
+  "test_aot_autograd_symbolic_module_exhaustive_nn_TransformerDecoderLayer_cpu_float32 (__main__.TestEagerFusionModuleInfoCPU)": 127.03733317057292,
+  "test_associative_scan_partial_grad_combine_mode_generic_compile_mode_compile_dynamic_shape_reverse_False_cpu (__main__.AssociativeScanTests)": 88.39666748046875,
+  "test_associative_scan_partial_grad_combine_mode_generic_compile_mode_compile_dynamic_shape_reverse_True_cpu (__main__.AssociativeScanTests)": 102.40783437093098,
+  "test_avg_pool3d_backward2_cpu (__main__.CpuTests)": 763.9162445068359,
+  "test_avg_pool3d_backward2_cuda (__main__.GPUTests)": 152.19875049591064,
+  "test_avg_pool3d_backward2_dynamic_shapes_cpu (__main__.DynamicShapesCodegenCpuTests)": 514.7272237141927,
+  "test_avg_pool3d_backward2_dynamic_shapes_cpu (__main__.DynamicShapesCpuTests)": 520.711890326606,
+  "test_avg_pool3d_backward2_dynamic_shapes_cuda (__main__.DynamicShapesCodegenGPUTests)": 64.15366808573405,
+  "test_avg_pool3d_backward2_dynamic_shapes_cuda (__main__.DynamicShapesGPUTests)": 128.34783426920572,
+  "test_backward_nn_functional_multi_head_attention_forward_cpu_float32 (__main__.TestCompositeComplianceCPU)": 71.84533437093098,
+  "test_backward_nn_functional_multi_head_attention_forward_cuda_float32 (__main__.TestCompositeComplianceCUDA)": 64.79450035095215,
+  "test_baddmm_search_space_EXHAUSTIVE (__main__.TestMaxAutotune)": 112.79600016276042,
+  "test_basic_cpu (__main__.EfficientConvBNEvalCpuTests)": 226.9800008138021,
+  "test_basic_cuda (__main__.EfficientConvBNEvalGpuTests)": 173.3750025431315,
+  "test_cat_2k_args (__main__.TestTEFuserDynamic)": 131.654914483428,
+  "test_checkpointing_without_reentrant_input_requires_grad_False (__main__.TestAutogradWithCompiledAutograd)": 347.29089016384546,
+  "test_checkpointing_without_reentrant_input_requires_grad_True (__main__.TestAutogradWithCompiledAutograd)": 448.10833062065973,
+  "test_collect_callgrind (__main__.TestBenchmarkUtils)": 285.1102227105035,
+  "test_compiled_dtensor_op_db_nn_functional_max_pool2d_cpu_float32 (__main__.TestCompiledDTensorOpsCPU)": 92.9769999186198,
+  "test_compiled_dtensor_op_db_nn_functional_poisson_nll_loss_cpu_float32 (__main__.TestCompiledDTensorOpsCPU)": 64.3003323872884,
+  "test_comprehensive_diff_cuda_complex128 (__main__.TestDecompCUDA)": 100.02783330281575,
+  "test_comprehensive_diff_cuda_complex64 (__main__.TestDecompCUDA)": 97.7408332824707,
+  "test_comprehensive_diff_cuda_float32 (__main__.TestDecompCUDA)": 67.63183339436848,
+  "test_comprehensive_diff_cuda_float64 (__main__.TestDecompCUDA)": 66.48316637674968,
+  "test_comprehensive_grid_sampler_2d_cpu_bfloat16 (__main__.TestDecompCPU)": 89.89199829101562,
+  "test_comprehensive_grid_sampler_2d_cpu_float16 (__main__.TestDecompCPU)": 90.3046646118164,
+  "test_comprehensive_grid_sampler_2d_cpu_float32 (__main__.TestDecompCPU)": 432.8856608072917,
+  "test_comprehensive_grid_sampler_2d_cpu_float64 (__main__.TestDecompCPU)": 431.4099934895833,
+  "test_comprehensive_grid_sampler_2d_cuda_bfloat16 (__main__.TestDecompCUDA)": 244.3134994506836,
+  "test_comprehensive_grid_sampler_2d_cuda_float16 (__main__.TestDecompCUDA)": 277.8193359375,
+  "test_comprehensive_grid_sampler_2d_cuda_float32 (__main__.TestDecompCUDA)": 1243.374491373698,
+  "test_comprehensive_grid_sampler_2d_cuda_float32 (__main__.TestInductorOpInfoCUDA)": 71.73899968465169,
+  "test_comprehensive_grid_sampler_2d_cuda_float64 (__main__.TestDecompCUDA)": 1143.0853271484375,
+  "test_comprehensive_grid_sampler_2d_cuda_float64 (__main__.TestInductorOpInfoCUDA)": 75.46933364868164,
+  "test_comprehensive_linalg_lu_cuda_float32 (__main__.TestInductorOpInfoCUDA)": 66.4915000597636,
+  "test_comprehensive_linalg_lu_cuda_float64 (__main__.TestInductorOpInfoCUDA)": 67.09683386484782,
+  "test_comprehensive_linalg_lu_factor_cuda_float64 (__main__.TestInductorOpInfoCUDA)": 62.22633425394694,
+  "test_comprehensive_linalg_lu_factor_ex_cuda_float32 (__main__.TestInductorOpInfoCUDA)": 62.32066790262858,
+  "test_comprehensive_linalg_lu_factor_ex_cuda_float64 (__main__.TestInductorOpInfoCUDA)": 62.156500498453774,
+  "test_comprehensive_linalg_lu_solve_cuda_float32 (__main__.TestInductorOpInfoCUDA)": 80.1408322652181,
+  "test_comprehensive_linalg_lu_solve_cuda_float64 (__main__.TestInductorOpInfoCUDA)": 79.4061648050944,
+  "test_comprehensive_linalg_matrix_norm_cuda_float32 (__main__.TestInductorOpInfoCUDA)": 95.98250071207683,
+  "test_comprehensive_linalg_matrix_norm_cuda_float64 (__main__.TestInductorOpInfoCUDA)": 94.55149968465169,
+  "test_comprehensive_linalg_norm_cuda_float16 (__main__.TestInductorOpInfoCUDA)": 98.94500128428142,
+  "test_comprehensive_linalg_norm_cuda_float32 (__main__.TestInductorOpInfoCUDA)": 115.80650011698405,
+  "test_comprehensive_linalg_norm_cuda_float64 (__main__.TestInductorOpInfoCUDA)": 115.1416654586792,
+  "test_comprehensive_linalg_norm_subgradients_at_zero_cuda_float16 (__main__.TestInductorOpInfoCUDA)": 87.03200038274129,
+  "test_comprehensive_linalg_norm_subgradients_at_zero_cuda_float32 (__main__.TestInductorOpInfoCUDA)": 110.46916611989339,
+  "test_comprehensive_linalg_norm_subgradients_at_zero_cuda_float64 (__main__.TestInductorOpInfoCUDA)": 111.24783261617024,
+  "test_comprehensive_linalg_solve_triangular_cuda_float32 (__main__.TestInductorOpInfoCUDA)": 83.31766764322917,
+  "test_comprehensive_linalg_solve_triangular_cuda_float64 (__main__.TestInductorOpInfoCUDA)": 80.25166447957356,
+  "test_comprehensive_linalg_svd_cuda_complex128 (__main__.TestDecompCUDA)": 69.97166570027669,
+  "test_comprehensive_linalg_svd_cuda_complex64 (__main__.TestDecompCUDA)": 68.84450022379558,
+  "test_comprehensive_linalg_svd_cuda_float32 (__main__.TestInductorOpInfoCUDA)": 96.60883108774821,
+  "test_comprehensive_linalg_svd_cuda_float64 (__main__.TestInductorOpInfoCUDA)": 93.79333368937175,
+  "test_comprehensive_linalg_vector_norm_cuda_float32 (__main__.TestInductorOpInfoCUDA)": 256.18433062235516,
+  "test_comprehensive_linalg_vector_norm_cuda_float64 (__main__.TestInductorOpInfoCUDA)": 248.46516863505045,
+  "test_comprehensive_logspace_cuda_float32 (__main__.TestInductorOpInfoCUDA)": 74.78166580200195,
+  "test_comprehensive_logspace_cuda_float64 (__main__.TestInductorOpInfoCUDA)": 69.59733327229817,
+  "test_comprehensive_masked_norm_cuda_float16 (__main__.TestInductorOpInfoCUDA)": 117.71599960327148,
+  "test_comprehensive_masked_norm_cuda_float32 (__main__.TestInductorOpInfoCUDA)": 113.31033579508464,
+  "test_comprehensive_masked_norm_cuda_float64 (__main__.TestInductorOpInfoCUDA)": 114.37183252970378,
+  "test_comprehensive_max_pool2d_with_indices_backward_cpu_float64 (__main__.TestDecompCPU)": 64.0938326517741,
+  "test_comprehensive_nn_functional_conv_transpose3d_cuda_complex128 (__main__.TestDecompCUDA)": 63.42050043741862,
+  "test_comprehensive_nn_functional_conv_transpose3d_cuda_complex64 (__main__.TestDecompCUDA)": 63.77633285522461,
+  "test_comprehensive_nn_functional_gaussian_nll_loss_cpu_float32 (__main__.TestDecompCPU)": 61.52166620890299,
+  "test_comprehensive_nn_functional_gaussian_nll_loss_cuda_float32 (__main__.TestDecompCUDA)": 118.71116765340169,
+  "test_comprehensive_nn_functional_gaussian_nll_loss_cuda_float64 (__main__.TestDecompCUDA)": 125.97533289591472,
+  "test_comprehensive_nn_functional_grid_sample_cpu_float32 (__main__.TestDecompCPU)": 110.10399881998698,
+  "test_comprehensive_nn_functional_grid_sample_cpu_float64 (__main__.TestDecompCPU)": 119.12533315022786,
+  "test_comprehensive_nn_functional_grid_sample_cuda_bfloat16 (__main__.TestDecompCUDA)": 63.85616683959961,
+  "test_comprehensive_nn_functional_grid_sample_cuda_float32 (__main__.TestDecompCUDA)": 263.6338348388672,
+  "test_comprehensive_nn_functional_grid_sample_cuda_float64 (__main__.TestDecompCUDA)": 260.0433349609375,
+  "test_comprehensive_nn_functional_interpolate_bicubic_cuda_float32 (__main__.TestDecompCUDA)": 80.7168337504069,
+  "test_comprehensive_nn_functional_interpolate_bicubic_cuda_float32 (__main__.TestInductorOpInfoCUDA)": 83.66033426920573,
+  "test_comprehensive_nn_functional_interpolate_bicubic_cuda_float64 (__main__.TestDecompCUDA)": 78.34583409627278,
+  "test_comprehensive_nn_functional_interpolate_bicubic_cuda_float64 (__main__.TestInductorOpInfoCUDA)": 83.96233240763347,
+  "test_comprehensive_nn_functional_interpolate_trilinear_cuda_float32 (__main__.TestDecompCUDA)": 120.8836669921875,
+  "test_comprehensive_nn_functional_interpolate_trilinear_cuda_float64 (__main__.TestDecompCUDA)": 124.52166748046875,
+  "test_comprehensive_nn_functional_max_pool2d_cuda_float16 (__main__.TestInductorOpInfoCUDA)": 819.3448282877604,
+  "test_comprehensive_nn_functional_max_pool2d_cuda_float32 (__main__.TestInductorOpInfoCUDA)": 800.9398295084635,
+  "test_comprehensive_nn_functional_max_pool2d_cuda_float64 (__main__.TestInductorOpInfoCUDA)": 798.9723409016927,
+  "test_comprehensive_nn_functional_max_pool3d_cuda_float32 (__main__.TestInductorOpInfoCUDA)": 650.4093424479166,
+  "test_comprehensive_nn_functional_max_pool3d_cuda_float64 (__main__.TestInductorOpInfoCUDA)": 672.7261657714844,
+  "test_comprehensive_nn_functional_max_unpool2d_cuda_float16 (__main__.TestInductorOpInfoCUDA)": 70.85333506266277,
+  "test_comprehensive_nn_functional_max_unpool2d_cuda_float32 (__main__.TestInductorOpInfoCUDA)": 73.07083384195964,
+  "test_comprehensive_nn_functional_max_unpool2d_cuda_float64 (__main__.TestInductorOpInfoCUDA)": 73.60183334350586,
+  "test_comprehensive_ormqr_cpu_complex128 (__main__.TestDecompCPU)": 63.091583251953125,
+  "test_comprehensive_ormqr_cuda_complex128 (__main__.TestDecompCUDA)": 119.43050130208333,
+  "test_comprehensive_ormqr_cuda_complex64 (__main__.TestDecompCUDA)": 129.18349838256836,
+  "test_comprehensive_ormqr_cuda_float32 (__main__.TestDecompCUDA)": 73.39383252461751,
+  "test_comprehensive_ormqr_cuda_float32 (__main__.TestInductorOpInfoCUDA)": 77.46083196004231,
+  "test_comprehensive_ormqr_cuda_float64 (__main__.TestDecompCUDA)": 100.95350011189778,
+  "test_comprehensive_ormqr_cuda_float64 (__main__.TestInductorOpInfoCUDA)": 76.22900009155273,
+  "test_comprehensive_svd_cuda_complex128 (__main__.TestDecompCUDA)": 76.60616556803386,
+  "test_comprehensive_svd_cuda_complex64 (__main__.TestDecompCUDA)": 72.88816579182942,
+  "test_comprehensive_svd_lowrank_cuda_complex128 (__main__.TestDecompCUDA)": 66.14150047302246,
+  "test_comprehensive_svd_lowrank_cuda_complex64 (__main__.TestDecompCUDA)": 71.52549870808919,
+  "test_constructor_autograd_SparseBSC_cuda (__main__.TestSparseAnyCUDA)": 103.31566747029622,
+  "test_constructor_autograd_SparseBSR_cuda (__main__.TestSparseAnyCUDA)": 90.61466598510742,
+  "test_conv1d_basic (__main__.TestXNNPACKConv1dTransformPass)": 85.5044453938802,
+  "test_conv1d_with_relu_fc (__main__.TestXNNPACKConv1dTransformPass)": 186.80122545030383,
+  "test_conv2d_binary_broadcast_shapes_cpu (__main__.TestPatternMatcherGenericCPU)": 73.75333404541016,
+  "test_conv3d_binary_broadcast_shapes_cpu (__main__.TestPatternMatcherGenericCPU)": 165.67066955566406,
+  "test_conv3d_binary_dynamic_shapes_cpu (__main__.TestDynamicPatternMatcherGenericCPU)": 69.34033330281575,
+  "test_conv3d_unary_dynamic_shapes_cpu (__main__.TestDynamicPatternMatcherGenericCPU)": 73.0816667344835,
+  "test_conv_bn_fuse_cpu (__main__.CpuTests)": 62.08125019073486,
+  "test_conv_unary_fusion_nnc (__main__.TestMkldnnFusion)": 82.00311109754774,
+  "test_correctness_AdamW_use_closure_True_cuda_float32 (__main__.CompiledOptimizerParityTestsCUDA)": 89.69633356730144,
+  "test_correctness_Adam_use_closure_True_cuda_float32 (__main__.CompiledOptimizerParityTestsCUDA)": 88.30749893188477,
+  "test_count_nonzero_all (__main__.TestBool)": 654.9627821180555,
+  "test_create_rand_mask_from_inputs_dynamic_shapes (__main__.DynamicShapesReproTests)": 92.5728333791097,
+  "test_dijkstra_expand_single_dim_strategy_to_mesh_hard_4d (__main__.TestDijkstraExpandSingleDimStrategy)": 303.05066935221356,
+  "test_dispatch_symbolic_meta_outplace_all_strides_nn_functional_gaussian_nll_loss_cuda_float32 (__main__.TestMetaCUDA)": 95.07450103759766,
+  "test_dtensor_op_db__native_batch_norm_legit_cpu_float32 (__main__.TestMultiThreadedDTensorOpsCPU)": 172.50000508626303,
+  "test_dtensor_op_db_native_batch_norm_cpu_float32 (__main__.TestMultiThreadedDTensorOpsCPU)": 200.29400126139322,
+  "test_dtensor_op_db_nn_functional_batch_norm_cpu_float32 (__main__.TestMultiThreadedDTensorOpsCPU)": 199.36066691080728,
+  "test_dtensor_op_db_nn_functional_binary_cross_entropy_with_logits_cpu_float32 (__main__.TestMultiThreadedDTensorOpsCPU)": 183.6556650797526,
+  "test_eager_sequence_nr_dynamic_shapes (__main__.DynamicShapesAotAutogradFallbackTests)": 195.9025821685791,
+  "test_fail_arithmetic_ops.py (__main__.TestTyping)": 85.94689008924696,
+  "test_fail_random.py (__main__.TestTyping)": 363.56556599934896,
+  "test_fail_torch_size.py (__main__.TestTyping)": 114.58616725787675,
+  "test_fn_fwgrad_bwgrad_cumprod_cuda_complex128 (__main__.TestFwdGradientsCUDA)": 83.33750025431316,
+  "test_fn_gradgrad_cumprod_cuda_complex128 (__main__.TestBwdGradientsCUDA)": 103.71716690063477,
+  "test_fuse_large_params_cpu (__main__.CpuTests)": 92.53525161743164,
+  "test_fuse_large_params_cuda (__main__.GPUTests)": 62.06562566757202,
+  "test_fuse_large_params_dynamic_shapes_cpu (__main__.DynamicShapesCodegenCpuTests)": 156.92011176215277,
+  "test_fuse_large_params_dynamic_shapes_cpu (__main__.DynamicShapesCpuTests)": 158.4638909233941,
+  "test_fuse_large_params_dynamic_shapes_cuda (__main__.DynamicShapesCodegenGPUTests)": 135.43400065104166,
+  "test_fuse_large_params_dynamic_shapes_cuda (__main__.DynamicShapesGPUTests)": 132.73733266194662,
+  "test_grad_nn_Transformer_cpu_float64 (__main__.TestModuleCPU)": 84.30703392028809,
+  "test_grad_nn_Transformer_cuda_float64 (__main__.TestModuleCUDA)": 80.69633356730144,
+  "test_gradgrad_nn_TransformerDecoderLayer_cuda_float64 (__main__.TestModuleCUDA)": 190.28999837239584,
+  "test_gradgrad_nn_TransformerEncoder_eval_mode_cuda_float64 (__main__.TestModuleCUDA)": 113.35149892171223,
+  "test_gradgrad_nn_TransformerEncoder_train_mode_cuda_float64 (__main__.TestModuleCUDA)": 123.2726656595866,
+  "test_gradgrad_nn_Transformer_cuda_float64 (__main__.TestModuleCUDA)": 494.573003133138,
+  "test_graph_based_redistribute_cost (__main__.DistributeWithDeviceOrderTest)": 67.4446652730306,
+  "test_graph_based_redistribute_cost (__main__.DistributeWithDeviceOrderTestWithLocalTensor)": 62.11899948120117,
+  "test_indirect_device_assert (__main__.TritonCodeGenTests)": 325.7254994710286,
+  "test_inplace_gradgrad_cumprod_cuda_complex128 (__main__.TestBwdGradientsCUDA)": 83.06583404541016,
+  "test_inputs_overlapping_with_mutation_stress_dynamic_shapes (__main__.DynamicShapesAotAutogradFallbackTests)": 122.00066714816623,
+  "test_jit_cuda_archflags (__main__.TestCppExtensionJIT)": 86.76616668701172,
+  "test_linalg_solve_triangular_large_cuda_complex128 (__main__.TestLinalgCUDA)": 1066.7661743164062,
+  "test_linalg_solve_triangular_large_cuda_complex64 (__main__.TestLinalgCUDA)": 72.56916681925456,
+  "test_linalg_solve_triangular_large_cuda_float64 (__main__.TestLinalgCUDA)": 151.44950103759766,
+  "test_linear_binary_cpp_wrapper (__main__.TestCppWrapper)": 140.42566935221353,
+  "test_linear_binary_dynamic_shapes_cpp_wrapper (__main__.DynamicShapesCppWrapperCpuTests)": 150.54800160725912,
+  "test_longformer_chunk_dynamic_shapes (__main__.DynamicShapesReproTests)": 104.55644395616319,
+  "test_low_memory_max_pool_dilation_1_dim_3_use_block_ptr_False_cpu_halide (__main__.HalideCpuTests)": 583.5033365885416,
+  "test_low_memory_max_pool_dilation_2_dim_3_use_block_ptr_False_cpu_halide (__main__.HalideCpuTests)": 516.3889973958334,
+  "test_lstm_cpu (__main__.TestMkldnnCPU)": 68.73433176676433,
+  "test_many_overlapping_inputs_does_not_explode_guards_dynamic_shapes (__main__.DynamicShapesReproTests)": 108.97577582465277,
+  "test_max_autotune_remote_caching_dynamic_True (__main__.TestMaxAutotuneRemoteCache)": 61.8836669921875,
+  "test_memory_format_operators_cuda (__main__.TestTorchDeviceTypeCUDA)": 125.83300018310547,
+  "test_ops_composition_names_cpu (__main__.TestTestParametrizationDeviceTypeCPU)": 157.14929997324944,
+  "test_ops_decorator_applies_op_and_param_specific_decorators_cpu (__main__.TestTestParametrizationDeviceTypeCPU)": 189.11739701827366,
+  "test_pattern_matcher_multi_user_cpu (__main__.CpuTritonTests)": 143.3016611735026,
+  "test_pq_vs_full_expansion_data_driven (__main__.TestDijkstraExpandSingleDimStrategy)": 881.4066772460938,
+  "test_proper_exit (__main__.TestDataLoader)": 219.32978057861328,
+  "test_proper_exit (__main__.TestDataLoaderPersistentWorkers)": 229.30333031548395,
+  "test_python_ref__refs_special_zeta_cuda_float64 (__main__.TestCommonCUDA)": 74.98866653442383,
+  "test_python_ref_executor__refs_special_zeta_executor_aten_cuda_float64 (__main__.TestCommonCUDA)": 89.93916702270508,
+  "test_quick_core_backward__unsafe_masked_index_cpu_float64 (__main__.TestDecompCPU)": 556.7093302408854,
+  "test_quick_core_backward__unsafe_masked_index_cuda_float64 (__main__.TestDecompCUDA)": 1110.48583984375,
+  "test_quick_core_backward__unsafe_masked_index_put_accumulate_cpu_float64 (__main__.TestDecompCPU)": 813.156005859375,
+  "test_quick_core_backward__unsafe_masked_index_put_accumulate_cuda_float64 (__main__.TestDecompCUDA)": 1398.7890014648438,
+  "test_quick_core_backward_nn_functional_max_unpool3d_grad_cpu_float64 (__main__.TestDecompCPU)": 84.88433329264323,
+  "test_quick_core_backward_nn_functional_max_unpool3d_grad_cuda_float64 (__main__.TestDecompCUDA)": 268.56183369954425,
+  "test_quick_core_backward_roll_cpu_float64 (__main__.TestDecompCPU)": 127.7530008951823,
+  "test_quick_core_backward_roll_cuda_float64 (__main__.TestDecompCUDA)": 241.10733795166016,
+  "test_quick_core_backward_select_scatter_cpu_float64 (__main__.TestDecompCPU)": 77.3393325805664,
+  "test_quick_core_backward_select_scatter_cuda_float64 (__main__.TestDecompCUDA)": 143.2455037434896,
+  "test_quick_core_backward_split_cuda_float64 (__main__.TestDecompCUDA)": 69.18416849772136,
+  "test_quick_core_backward_split_with_sizes_copy_cpu_float64 (__main__.TestDecompCPU)": 105.94433339436848,
+  "test_quick_core_backward_split_with_sizes_copy_cuda_float64 (__main__.TestDecompCUDA)": 174.2881647745768,
+  "test_quick_core_backward_std_cpu_float64 (__main__.TestDecompCPU)": 77.57733408610027,
+  "test_quick_core_backward_std_cuda_float64 (__main__.TestDecompCUDA)": 158.21883392333984,
+  "test_register_spills_cuda (__main__.BenchmarkFusionGpuTest)": 108.23566563924153,
+  "test_resize_as_cuda (__main__.GPUTests)": 89.16333452860515,
+  "test_resize_cuda (__main__.GPUTests)": 67.50366735458374,
+  "test_reveal_opt_size.py (__main__.TestTyping)": 60.49586756983772,
+  "test_run2run_determinism_model_name_BertForMaskedLM_training_or_inference_inference_precision_amp (__main__.DeterministicTest)": 104.33299763997395,
+  "test_run2run_determinism_model_name_BertForMaskedLM_training_or_inference_inference_precision_bfloat16 (__main__.DeterministicTest)": 64.46683311462402,
+  "test_run2run_determinism_model_name_BertForMaskedLM_training_or_inference_training_precision_amp (__main__.DeterministicTest)": 198.3463338216146,
+  "test_run2run_determinism_model_name_BertForMaskedLM_training_or_inference_training_precision_bfloat16 (__main__.DeterministicTest)": 153.71866353352866,
+  "test_run2run_determinism_model_name_BertForMaskedLM_training_or_inference_training_precision_float16 (__main__.DeterministicTest)": 167.72516632080078,
+  "test_run2run_determinism_model_name_BertForMaskedLM_training_or_inference_training_precision_float32 (__main__.DeterministicTest)": 188.2213338216146,
+  "test_run2run_determinism_model_name_DistillGPT2_training_or_inference_training_precision_amp (__main__.DeterministicTest)": 82.31683349609375,
+  "test_run2run_determinism_model_name_DistillGPT2_training_or_inference_training_precision_bfloat16 (__main__.DeterministicTest)": 115.22916666666667,
+  "test_run2run_determinism_model_name_DistillGPT2_training_or_inference_training_precision_float16 (__main__.DeterministicTest)": 84.5586649576823,
+  "test_run2run_determinism_model_name_DistillGPT2_training_or_inference_training_precision_float32 (__main__.DeterministicTest)": 122.09583409627278,
+  "test_run2run_determinism_model_name_GoogleFnet_training_or_inference_training_precision_amp (__main__.DeterministicTest)": 95.04883448282878,
+  "test_run2run_determinism_model_name_GoogleFnet_training_or_inference_training_precision_bfloat16 (__main__.DeterministicTest)": 94.20300165812175,
+  "test_run2run_determinism_model_name_GoogleFnet_training_or_inference_training_precision_float16 (__main__.DeterministicTest)": 118.3158327738444,
+  "test_run2run_determinism_model_name_GoogleFnet_training_or_inference_training_precision_float32 (__main__.DeterministicTest)": 119.6198336283366,
+  "test_runtime_checks_large_cpu (__main__.AOTInductorTestABICompatibleCpu)": 70.40016682942708,
+  "test_runtime_checks_large_cpu_with_stack_allocation (__main__.AOTInductorTestABICompatibleCpuWithStackAllocation)": 70.30955590142145,
+  "test_runtime_checks_large_cuda (__main__.AOTInductorTestABICompatibleGpu)": 171.0409952799479,
+  "test_sdpa_kernel_ctx_manager2_dynamic_shapes (__main__.DynamicShapesCtxManagerTests)": 106.9195556640625,
+  "test_searchsorted_cuda (__main__.GPUTests)": 82.54333392779033,
+  "test_shared_memory_pruning_addmm_bfloat16_mat1_transposed_False_mat2_transposed_False_use_tma_False (__main__.TestTemplateConfigPruning)": 143.74433390299478,
+  "test_shared_memory_pruning_addmm_bfloat16_mat1_transposed_False_mat2_transposed_True_use_tma_False (__main__.TestTemplateConfigPruning)": 143.3086675008138,
+  "test_shared_memory_pruning_addmm_bfloat16_mat1_transposed_True_mat2_transposed_False_use_tma_False (__main__.TestTemplateConfigPruning)": 143.81799825032553,
+  "test_shared_memory_pruning_addmm_bfloat16_mat1_transposed_True_mat2_transposed_True_use_tma_False (__main__.TestTemplateConfigPruning)": 144.1403350830078,
+  "test_shared_memory_pruning_addmm_float32_mat1_transposed_False_mat2_transposed_False_use_tma_False (__main__.TestTemplateConfigPruning)": 153.8959986368815,
+  "test_shared_memory_pruning_addmm_float32_mat1_transposed_False_mat2_transposed_True_use_tma_False (__main__.TestTemplateConfigPruning)": 159.3713353474935,
+  "test_shared_memory_pruning_addmm_float32_mat1_transposed_True_mat2_transposed_False_use_tma_False (__main__.TestTemplateConfigPruning)": 151.32366689046225,
+  "test_shared_memory_pruning_addmm_float32_mat1_transposed_True_mat2_transposed_True_use_tma_False (__main__.TestTemplateConfigPruning)": 151.53400166829428,
+  "test_shared_memory_pruning_mm_bfloat16_mat1_transposed_False_mat2_transposed_False_use_tma_False (__main__.TestTemplateConfigPruning)": 137.7143351236979,
+  "test_shared_memory_pruning_mm_bfloat16_mat1_transposed_False_mat2_transposed_True_use_tma_False (__main__.TestTemplateConfigPruning)": 137.6969985961914,
+  "test_shared_memory_pruning_mm_bfloat16_mat1_transposed_True_mat2_transposed_False_use_tma_False (__main__.TestTemplateConfigPruning)": 139.831662495931,
+  "test_shared_memory_pruning_mm_bfloat16_mat1_transposed_True_mat2_transposed_True_use_tma_False (__main__.TestTemplateConfigPruning)": 138.6189982096354,
+  "test_shared_memory_pruning_mm_float32_mat1_transposed_False_mat2_transposed_False_use_tma_False (__main__.TestTemplateConfigPruning)": 149.81566874186197,
+  "test_shared_memory_pruning_mm_float32_mat1_transposed_False_mat2_transposed_True_use_tma_False (__main__.TestTemplateConfigPruning)": 154.44766743977866,
+  "test_shared_memory_pruning_mm_float32_mat1_transposed_True_mat2_transposed_False_use_tma_False (__main__.TestTemplateConfigPruning)": 146.81933085123697,
+  "test_shared_memory_pruning_mm_float32_mat1_transposed_True_mat2_transposed_True_use_tma_False (__main__.TestTemplateConfigPruning)": 147.39400227864584,
+  "test_shuffler_iterdatapipe (__main__.IntegrationTestDataLoaderDataPipe)": 122.32033454047308,
+  "test_slow_tasks (__main__.TestFunctionalAutogradBenchmark)": 124.18311055501302,
+  "test_softmax_cpu_float64 (__main__.TestSparseCPU)": 60.396761218706764,
+  "test_sort_bool_cpu (__main__.CpuTritonTests)": 348.36634318033856,
+  "test_sort_dynamic_shape_with_check_cuda (__main__.TestInductorDynamicCUDA)": 112.24066670735677,
+  "test_sort_stable_cuda (__main__.GPUTests)": 126.56412506103516,
+  "test_sort_transpose_cpu (__main__.CpuTritonTests)": 376.54832967122394,
+  "test_svd_lowrank_cuda_complex128 (__main__.TestLinalgCUDA)": 327.46799723307294,
+  "test_tensorwise_scaling_acceptable_input_dims_M_1024_K_1024_N_2048_persistent_matmul_False_cpu (__main__.TestFP8LoweringCPU)": 75.4329337477684,
+  "test_terminate_handler_on_crash (__main__.TestTorch)": 207.00677956475153,
+  "test_terminate_signal (__main__.ForkTest)": 232.13122049967447,
+  "test_terminate_signal (__main__.ParallelForkServerShouldWorkTest)": 232.00266604953342,
+  "test_terminate_signal (__main__.SpawnTest)": 222.40377510918512,
+  "test_torch_size_tensor_index_scalar_constant_dynamic_shapes (__main__.DynamicShapesMiscTests)": 601.1639200846354,
+  "test_train_parity_multi_group (__main__.TestFullyShard1DTrainingCore)": 174.3316650390625,
+  "test_train_parity_multi_group_cpu_offload_eager (__main__.TestFullyShard1DTrainingCore)": 61.27133305867513,
+  "test_train_parity_multi_group_unshard_async_op (__main__.TestFullyShard1DTrainingCore)": 64.80400085449219,
+  "test_triton_bsr_scatter_mm_blocksize_64_cuda_float32 (__main__.TestSparseCompressedTritonKernelsCUDA)": 87.95033391316731,
+  "test_triton_bsr_softmax_cuda_bfloat16 (__main__.TestSparseCompressedTritonKernelsCUDA)": 254.8540013631185,
+  "test_triton_bsr_softmax_cuda_float16 (__main__.TestSparseCompressedTritonKernelsCUDA)": 256.1475016276042,
+  "test_triton_bsr_softmax_cuda_float32 (__main__.TestSparseCompressedTritonKernelsCUDA)": 194.4748331705729,
+  "test_tuning_pool_timeout (__main__.TestTuningProcessPool)": 126.16133371988933,
+  "test_unbacked_dtensor_op_db_clamp_cpu_float32 (__main__.TestUnbackedDTensorOpsCPU)": 96.5290018717448,
+  "test_unbind (__main__.TestVmapOperators)": 82.44380033289393,
+  "test_variant_consistency_jit_nn_functional_max_pool2d_cpu_float32 (__main__.TestJitCPU)": 88.58866628011067,
+  "test_variant_consistency_jit_nn_functional_max_pool2d_cuda_float32 (__main__.TestJitCUDA)": 78.71033223470052,
+  "test_vec_compare_op_cpu_only (__main__.CPUReproTests)": 62.92533323499892,
+  "test_views1_cuda (__main__.GPUTests)": 106.26166693369548,
+  "test_vmapjvpvjp_linalg_lstsq_grad_oriented_cpu_float32 (__main__.TestOperatorsCPU)": 113.2413330078125,
+  "test_vmapjvpvjp_linalg_lu_solve_cpu_float32 (__main__.TestOperatorsCPU)": 68.83388929013853,
+  "test_vmapjvpvjp_linalg_lu_solve_cuda_float32 (__main__.TestOperatorsCUDA)": 67.28550148010254,
+  "test_vmapjvpvjp_max_pool2d_with_indices_backward_cpu_float32 (__main__.TestOperatorsCPU)": 86.25900014241536,
+  "test_vmapjvpvjp_max_pool2d_with_indices_backward_cuda_float32 (__main__.TestOperatorsCUDA)": 78.42966588338216,
+  "test_vmapjvpvjp_nn_functional_conv2d_cpu_float32 (__main__.TestOperatorsCPU)": 69.55499903361003,
+  "test_vmapjvpvjp_nn_functional_max_pool2d_cpu_float32 (__main__.TestOperatorsCPU)": 89.97300211588542,
+  "test_vmapjvpvjp_nn_functional_max_pool2d_cuda_float32 (__main__.TestOperatorsCUDA)": 79.44366709391277,
+  "test_vmapjvpvjp_unbind_cpu_float32 (__main__.TestOperatorsCPU)": 66.21299997965495,
+  "test_vmapjvpvjp_unbind_cuda_float32 (__main__.TestOperatorsCUDA)": 81.87016677856445,
+  "test_vmapvjpvjp_meshgrid_list_of_tensors_cuda_float32 (__main__.TestOperatorsCUDA)": 82.90099970499675,
+  "test_vmapvjpvjp_meshgrid_variadic_tensors_cuda_float32 (__main__.TestOperatorsCUDA)": 81.67116800944011,
+  "test_vmapvjpvjp_nn_functional_bilinear_cuda_float32 (__main__.TestOperatorsCUDA)": 110.1931660970052
 }
\ No newline at end of file
diff --git a/test/test_accelerator.py b/test/test_accelerator.py
index 7daebc01adfe9..43622c98662a9 100644
--- a/test/test_accelerator.py
+++ b/test/test_accelerator.py
@@ -3,6 +3,7 @@
 import gc
 import sys
 import unittest
+from contextlib import nullcontext
 
 import torch
 from torch.testing._internal.common_utils import (
@@ -111,6 +112,21 @@ def test_stream_context_manager(self):
             self.assertEqual(torch.accelerator.current_stream(), s)
         self.assertEqual(torch.accelerator.current_stream(), prev_stream)
 
+    def test_stream_context_manager_reentrance(self):
+        prev_stream = torch.accelerator.current_stream()
+        s0 = torch.Stream()
+        with s0, s0:
+            self.assertEqual(torch.accelerator.current_stream(), s0)
+        self.assertEqual(torch.accelerator.current_stream(), prev_stream)
+        s1 = torch.Stream()
+        with s0:
+            self.assertEqual(torch.accelerator.current_stream(), s0)
+            with s1:
+                self.assertEqual(torch.accelerator.current_stream(), s1)
+                with s0:
+                    self.assertEqual(torch.accelerator.current_stream(), s0)
+        self.assertEqual(torch.accelerator.current_stream(), prev_stream)
+
     @unittest.skipIf(not TEST_MULTIACCELERATOR, "only one accelerator detected")
     def test_multi_device_stream_context_manager(self):
         src_device = 0
@@ -245,6 +261,36 @@ def test_get_memory_info(self):
         self.assertGreaterEqual(free_bytes, 0)
         self.assertGreaterEqual(total_bytes, 0)
 
+    def test_device_capability_supported_dtypes(self):
+        try:
+            caps = torch.accelerator.get_device_capability()
+        except RuntimeError:
+            self.skipTest("Backend doesn't support get_device_capability")
+
+        supported_dtypes = caps["supported_dtypes"]
+        self.assertIsInstance(supported_dtypes, set)
+        self.assertGreater(len(supported_dtypes), 0)
+
+        acc = torch.accelerator.current_accelerator()
+        reference_dtype = next(iter(supported_dtypes))
+
+        all_dtypes = [
+            getattr(torch, name)
+            for name in dir(torch)
+            if isinstance(getattr(torch, name), torch.dtype)
+        ]
+        for dtype in all_dtypes:
+            with self.subTest(dtype=dtype):
+                ctx = (
+                    nullcontext()
+                    if dtype in supported_dtypes
+                    else self.assertRaises((RuntimeError, TypeError))
+                )
+                with ctx:
+                    t = torch.empty(16, dtype=dtype, device=acc)
+                    t = t.to(reference_dtype)
+                    t = t.to(dtype)
+
 
 if __name__ == "__main__":
     run_tests()
diff --git a/test/test_as_strided.py b/test/test_as_strided.py
index 4f36722485f8e..4265ba9af8a46 100644
--- a/test/test_as_strided.py
+++ b/test/test_as_strided.py
@@ -1,7 +1,6 @@
 # Owner(s): ["oncall: pt2"]
 
 from collections import deque
-from typing import Optional
 
 import torch
 from torch.testing._internal.common_utils import run_tests, TestCase
@@ -154,7 +153,7 @@ def test_subset_property(self) -> None:
         Test that for sizes 2..10, each smaller tensor results in a strict
         subset of possible states compared to the next one.
         """
-        prev_states: Optional[set[tuple[tuple[int, ...], tuple[int, ...]]]] = None
+        prev_states: set[tuple[tuple[int, ...], tuple[int, ...]]] | None = None
         for size in range(2, 11):
             current_states = enumerate_reachable_states(size)
 
diff --git a/test/test_autocast.py b/test/test_autocast.py
index 0565626255af2..b262ed95dbbfc 100644
--- a/test/test_autocast.py
+++ b/test/test_autocast.py
@@ -544,6 +544,13 @@ def test_autocast_mixed_grad_contexts_cuda(self):
         """Test mixed grad contexts on CUDA"""
         self._test_autocast_mixed_grad_contexts_impl("cuda", torch.float16)
 
+    def test_autocast_called_with_non_callable(self):
+        """Test that autocast gives a clear error when misused as a function wrapper"""
+        x = torch.randn(2, 3)
+        msg = r"autocast\(\)\(func\) requires a callable, but got Tensor"
+        with self.assertRaisesRegex(TypeError, msg):
+            torch.autocast(device_type="cpu")(x)
+
 
 if __name__ == "__main__":
     run_tests()
diff --git a/test/test_autograd.py b/test/test_autograd.py
index eb2ad5cb91cb5..5aafc5e8de970 100644
--- a/test/test_autograd.py
+++ b/test/test_autograd.py
@@ -20,7 +20,7 @@
 import uuid
 import warnings
 import weakref
-from collections import OrderedDict
+from collections import defaultdict, OrderedDict
 from copy import deepcopy
 from functools import partial, reduce
 from itertools import product
@@ -58,6 +58,7 @@
     instantiate_device_type_tests,
     onlyCPU,
     onlyCUDA,
+    skipCUDAIf,
     skipMeta,
 )
 from torch.testing._internal.common_dtype import floating_types_and
@@ -82,6 +83,7 @@
     skipIfXpu,
     slowTest,
     TEST_WITH_TORCHDYNAMO,
+    TEST_XPU,
     TestCase,
 )
 from torch.utils._mode_utils import no_dispatch
@@ -93,6 +95,7 @@
     create_selective_checkpoint_contexts,
 )
 from torch.utils.flop_counter import FlopCounterMode
+from torch.utils.weak import WeakTensorKeyDictionary
 
 
 if TYPE_CHECKING:
@@ -5075,6 +5078,19 @@ def f(x):
             self.assertTrue(torch.autograd.is_view_replay_enabled())
         self.assertFalse(torch.autograd.is_view_replay_enabled())
 
+        prev = torch.autograd.is_view_replay_enabled()
+        ctx = torch.autograd._force_original_view_tracking(not prev)
+        # Construction eagerly sets state (function-form behavior).
+        self.assertEqual(torch.autograd.is_view_replay_enabled(), not prev)
+        with ctx:
+            self.assertEqual(torch.autograd.is_view_replay_enabled(), not prev)
+            out = f(x)
+            self.assertTrue(
+                ("ViewBackward" if not prev else "AsStridedBackward")
+                in str(out.grad_fn)
+            )
+        self.assertEqual(torch.autograd.is_view_replay_enabled(), prev)
+
         # Test as a function
         torch.autograd._force_original_view_tracking(False)
         out = f(x)
@@ -5086,6 +5102,20 @@ def f(x):
         self.assertTrue("ViewBackward" in str(out.grad_fn))
         self.assertTrue(torch.autograd.is_view_replay_enabled())
 
+        prev = torch.autograd.is_view_replay_enabled()
+
+        @torch.autograd._force_original_view_tracking(not prev)
+        def g(x):
+            return f(x)
+
+        # __call__ undoes the __init__ mutation, so ambient state is restored.
+        self.assertEqual(torch.autograd.is_view_replay_enabled(), prev)
+        out = g(x)
+        self.assertTrue(
+            ("ViewBackward" if not prev else "AsStridedBackward") in str(out.grad_fn)
+        )
+        self.assertEqual(torch.autograd.is_view_replay_enabled(), prev)
+
     def test_unsafe_set_version_counter(self):
         x = torch.ones(2, requires_grad=True).clone()
         x.add_(1)
@@ -5931,6 +5961,57 @@ def run_fn(a):
 
                 out.backward()
 
+    @unittest.skipIf(not TEST_CUDA, "test requires CUDA")
+    def test_forward_traceback_preserves_exception_with_checkpoint(self):
+        # Regression test: gatherForwardTraceback() must not clear a pending
+        # Python exception.  See combined_traceback.cpp for the fix.
+        #
+        # Ingredients: (1) CUDA memory history recording with context="all"
+        # so allocator callbacks fire on free, (2) a custom library op whose
+        # forward goes through THPFunction_apply (via Generated in
+        # torch._library.autograd), (3) non-reentrant checkpoint.
+        #
+        # During backward recomputation, _StopRecomputationError is raised
+        # from the checkpoint pack_hook inside _save_variables.  The exception
+        # stays pending in Python thread state while C++ stack-unwinds (the
+        # default python_error ctor does not persist).  Destroying the output
+        # THPObjectPtr during unwinding frees the recomputed output tensor's
+        # CUDA storage, triggering the allocator callback ->
+        # CapturedTraceback::gather() -> gatherForwardTraceback().  On
+        # Python < 3.13, without the PyErr_Fetch/PyErr_Restore fix, the
+        # PyDict_GetItemRef compat shim clears the pending exception ->
+        # SystemError.
+        with torch.library._scoped_library("_test_autograd", "FRAGMENT"):
+
+            @torch.library.custom_op("_test_autograd::sin_op", mutates_args=())
+            def sin_op(x: torch.Tensor) -> torch.Tensor:
+                return x.sin()
+
+            def setup_context(ctx, inputs, output):
+                (x,) = inputs
+                ctx.save_for_backward(x)
+
+            def backward(ctx, grad):
+                (x,) = ctx.saved_tensors
+                return grad * x.cos()
+
+            torch.library.register_autograd(
+                "_test_autograd::sin_op",
+                backward,
+                setup_context=setup_context,
+            )
+
+            def fn(x):
+                return torch.ops._test_autograd.sin_op(x)
+
+            try:
+                torch.cuda.memory._record_memory_history("all", stacks="python")
+                x = torch.randn(4, device="cuda", requires_grad=True)
+                y = checkpoint(fn, x, use_reentrant=False)
+                y.sum().backward()
+            finally:
+                torch.cuda.memory._record_memory_history(None)
+
     def test_no_grad_copy(self):
         # create autograd function that saves grad pointer as class static
         class MyFunc(Function):
@@ -7305,10 +7386,10 @@ def foo(x, y, z):
             x = torch.randn(3, 3, requires_grad=True)
             y = torch.randn(3, 3, requires_grad=True)
             z = torch.randn(3, 3, requires_grad=True)
-            if device_type == "cuda":
-                x = x.cuda()
-                y = y.cuda()
-                z = z.cuda()
+            if device_type in ("cuda", "xpu"):
+                x = x.to(device_type)
+                y = y.to(device_type)
+                z = z.to(device_type)
 
             with torch.autocast(
                 enabled=enabled, device_type=device_type, dtype=torch.bfloat16
@@ -7328,15 +7409,17 @@ def test_checkpointing_non_reentrant_autocast_cpu(self):
         self._test_checkpointing_non_reentrant_autocast(device_type="cpu")
 
     @unittest.skipIf(
-        not torch.cuda.is_available() or not torch.cuda.is_bf16_supported(),
-        "Test requires CUDA bf16 support",
+        (not torch.cuda.is_available() or not torch.cuda.is_bf16_supported())
+        and (not torch.xpu.is_available() or not torch.xpu.is_bf16_supported()),
+        "Test requires CUDA or XPU bf16 support",
     )
     def test_checkpointing_non_reentrant_autocast_gpu(self):
         """
         Test that autocast args/kwargs such as the dtype are preserved during
         non-reentrant checkpoint recomputation on GPU.
         """
-        self._test_checkpointing_non_reentrant_autocast(device_type="cuda")
+        device_type = "cuda" if torch.cuda.is_available() else "xpu"
+        self._test_checkpointing_non_reentrant_autocast(device_type=device_type)
 
     @unittest.skipIf(not torch.cuda.is_available(), "Test requires CUDA")
     @slowTest
@@ -7647,6 +7730,197 @@ def fn(x):
                 out = checkpoint(fn, a, use_reentrant=False, debug=True)
                 out.backward()
 
+    def test_checkpoint_error_suggests_mark_dynamic(self):
+        """CheckpointError messages should suggest mark_dynamic and lru config."""
+        from torch.utils.checkpoint import checkpoint
+
+        count = [0]
+
+        def save_2_tensors(a):
+            b = a**2
+            c = a**3
+            return b + c
+
+        def save_3_tensors(a):
+            b = a**2
+            c = a**3
+            d = a**4
+            return b + c + d
+
+        def get_non_det_fn(orig_fn, recompute_fn):
+            def fn(a):
+                if count[0] == 0:
+                    count[0] += 1
+                    return orig_fn(a)
+                return recompute_fn(a)
+
+            return fn
+
+        a = torch.randn(1, requires_grad=True)
+        fn = get_non_det_fn(orig_fn=save_3_tensors, recompute_fn=save_2_tensors)
+        with self.assertRaises(RuntimeError) as cm:
+            out = checkpoint(fn, a, use_reentrant=False)
+            out.backward()
+        error_msg = str(cm.exception)
+        self.assertIn(
+            "Use torch._dynamo.mark_dynamic() to explicitly mark varying dimensions",
+            error_msg,
+        )
+        self.assertIn(
+            "Call torch._C._dynamo.eval_frame._set_lru_cache(False) to disable LRU cache",
+            error_msg,
+        )
+
+    @torch._dynamo.config.patch(
+        automatic_dynamic_shapes=True,
+        assume_static_by_default=True,
+    )
+    def test_checkpoint_automatic_dynamic_graph_shadowing(self):
+        """Activation checkpointing + automatic dynamic shapes graph shadowing.
+
+        When a compiled function is used inside checkpoint and another call
+        with a different input size triggers automatic dynamic recompilation,
+        the new dynamic graph is inserted at the front of the Dynamo cache
+        and shadows the original static graph. During backward, checkpoint's
+        recompute may hit the dynamic graph instead of the static graph used
+        in the original forward, causing a saved-tensor mismatch.
+
+        This test exercises the scenario described in PR #174993: the forward
+        of a checkpointed layer uses a static graph (G0), then a subsequent
+        call with a different batch size triggers automatic dynamic (G1 at
+        cache front). During backward, checkpoint's recompute for the first
+        layer hits G1 instead of G0.
+        """
+        import torch._dynamo
+        from torch.utils.checkpoint import checkpoint
+
+        torch._dynamo.reset()
+
+        @torch.compile(backend="aot_eager")
+        def fn(x):
+            # Multiple ops so the partitioner has save/recompute choices.
+            # With static shapes the partitioner sees exact costs; with
+            # dynamic shapes it uses size_hint fallbacks, potentially
+            # making different save-vs-recompute decisions.
+            a = x.sin()
+            b = a.cos()
+            c = (a * b).exp()
+            return c
+
+        # Step 1: Checkpointed forward with batch_size=4 → static graph G0
+        x = torch.randn(4, 8, requires_grad=True)
+        cp_out = checkpoint(fn, x, use_reentrant=False)
+
+        # Step 2: Non-checkpointed call with batch_size=7 triggers automatic
+        # dynamic recompilation → dynamic graph G1 inserted at cache front,
+        # shadowing G0 for any future call (including checkpoint recompute).
+        y = torch.randn(7, 8, requires_grad=True)
+        plain_out = fn(y)
+
+        loss = cp_out.sum() + plain_out.sum()
+
+        # Step 3: Backward. Checkpoint recomputes fn(x) but now hits G1
+        # (dynamic) instead of G0 (static). If the graphs save different
+        # tensors, this raises CheckpointError with our diagnostic message.
+        try:
+            loss.backward()
+        except RuntimeError as e:
+            error_msg = str(e)
+            # If the mismatch triggers, verify the error message includes
+            # the automatic dynamic shapes suggestions.
+            self.assertIn(
+                "Use torch._dynamo.mark_dynamic() to explicitly mark varying "
+                "dimensions",
+                error_msg,
+            )
+            self.assertIn(
+                "Call torch._C._dynamo.eval_frame._set_lru_cache(False) to "
+                "disable LRU cache",
+                error_msg,
+            )
+
+        torch._dynamo.reset()
+
+    @torch._dynamo.config.patch(
+        automatic_dynamic_shapes=True,
+        assume_static_by_default=True,
+    )
+    def test_checkpoint_automatic_dynamic_mark_dynamic_workaround(self):
+        """mark_dynamic avoids the graph shadowing problem.
+
+        By marking the varying dimension as dynamic upfront, both the
+        checkpointed forward and the subsequent call compile a single
+        dynamic graph from the start — no static-to-dynamic transition,
+        no cache shadowing, no saved-tensor mismatch.
+        """
+        import torch._dynamo
+        from torch.utils.checkpoint import checkpoint
+
+        torch._dynamo.reset()
+
+        @torch.compile(backend="aot_eager")
+        def fn(x):
+            a = x.sin()
+            b = a.cos()
+            c = (a * b).exp()
+            return c
+
+        x = torch.randn(4, 8, requires_grad=True)
+        torch._dynamo.mark_dynamic(x, 0)
+        cp_out = checkpoint(fn, x, use_reentrant=False)
+
+        y = torch.randn(7, 8, requires_grad=True)
+        torch._dynamo.mark_dynamic(y, 0)
+        plain_out = fn(y)
+
+        loss = cp_out.sum() + plain_out.sum()
+        # Should succeed — both calls use the same dynamic graph.
+        loss.backward()
+
+        torch._dynamo.reset()
+
+    @torch._dynamo.config.patch(
+        automatic_dynamic_shapes=True,
+        assume_static_by_default=True,
+    )
+    def test_checkpoint_automatic_dynamic_lru_disabled_workaround(self):
+        """Disabling LRU cache reordering avoids the graph shadowing problem.
+
+        With LRU disabled, new dynamic graphs are appended to the back of
+        the cache instead of the front. The original static graph is still
+        checked first, so checkpoint's recompute hits the same graph that
+        the forward used.
+        """
+        import torch._dynamo
+        from torch.utils.checkpoint import checkpoint
+
+        torch._dynamo.reset()
+        # Disable LRU cache reordering so static graphs stay at the front.
+        torch._C._dynamo.eval_frame._set_lru_cache(False)
+
+        try:
+
+            @torch.compile(backend="aot_eager")
+            def fn(x):
+                a = x.sin()
+                b = a.cos()
+                c = (a * b).exp()
+                return c
+
+            x = torch.randn(4, 8, requires_grad=True)
+            cp_out = checkpoint(fn, x, use_reentrant=False)
+
+            y = torch.randn(7, 8, requires_grad=True)
+            plain_out = fn(y)
+
+            loss = cp_out.sum() + plain_out.sum()
+            # Should succeed — LRU disabled means static graph G0 stays
+            # first in cache, so checkpoint recompute still hits G0.
+            loss.backward()
+        finally:
+            torch._C._dynamo.eval_frame._set_lru_cache(True)
+            torch._dynamo.reset()
+
     def test_access_saved_tensor_twice_without_recomputation_works(self):
         count = [0]
 
@@ -7828,6 +8102,28 @@ def test_checkpointing_without_reentrant_correct_grad(self):
         self.assertEqual(b_grad, c_grad)
         self.assertEqual(b_grad, d_grad)
 
+    @unittest.skipIf(not torch.cuda.is_available(), "requires CUDA")
+    def test_checkpointing_without_reentrant_with_block_mask(self):
+        from torch.nn.attention.flex_attention import BlockMask, create_block_mask
+        from torch.utils._pytree import register_pytree_node, SUPPORTED_NODES
+
+        if BlockMask not in SUPPORTED_NODES:
+            register_pytree_node(
+                BlockMask,
+                BlockMask._flatten,
+                BlockMask._unflatten,
+                flatten_with_keys_fn=BlockMask._flatten_with_keys,
+                serialized_type_name="torch.nn.attention.flex_attention.BlockMask",
+            )
+
+        block_mask = create_block_mask(
+            lambda b, h, q, kv: q >= kv, B=1, H=1, Q_LEN=128, KV_LEN=128
+        )
+        x = torch.randn(4, 128, device="cuda")
+
+        result = checkpoint(lambda x, mask: x * 2, x, block_mask, use_reentrant=False)
+        self.assertEqual(result, x * 2)
+
     @skipIfXpu(msg="torch._C._scatter Not implemented on XPU, issue #143239")
     def test_checkpointing_without_reentrant_dataparallel(self):
         """
@@ -8304,7 +8600,39 @@ def unpack_hook(t):
         self.assertTrue(pow2_sv.data is c)
         self.assertIsNone(pow2_sv.unpack_hook)
 
-    def test_cant_create_saved_tensors(self):
+    def test_saved_tensor_constructor_with_hooks(self):
+        pack_count = [0]
+        unpack_count = [0]
+
+        def pack_hook(x):
+            pack_count[0] += 1
+            return x
+
+        def unpack_hook(x):
+            unpack_count[0] += 1
+            return x
+
+        a = torch.randn(5, requires_grad=True)
+        with torch.autograd.graph.saved_tensors_hooks(pack_hook, unpack_hook):
+            saved = torch._C._autograd._make_saved_tensor(a, is_output=False)
+            self.assertEqual(pack_count[0], 1)
+            self.assertEqual(unpack_count[0], 0)
+
+            unpacked = saved.unpack()
+            self.assertEqual(pack_count[0], 1)
+            self.assertEqual(unpack_count[0], 1)
+
+            unpacked2 = saved.unpack()
+            self.assertEqual(pack_count[0], 1)
+            self.assertEqual(unpack_count[0], 2)
+
+        # Check tensor equality outside the hooks context to avoid
+        # triggering additional pack hooks from assertEqual internals
+        self.assertEqual(unpacked, a)
+        self.assertEqual(unpacked2, a)
+
+    def test_saved_tensor_constructor_forbidden_without_flag(self):
+        a = torch.randn(5, requires_grad=True)
         with self.assertRaisesRegex(
             RuntimeError,
             "Trying to create a SavedTensor object from Python is forbidden",
@@ -10571,25 +10899,31 @@ def test(get_input, cuda, pin_memory):
                 )
                 test(lambda: x, cuda, pin_memory)
 
-    @unittest.skipIf(not TEST_CUDA, "test requires CUDA")
+    @unittest.skipIf(not TEST_CUDA and not TEST_XPU, "test requires CUDA or XPU")
     def test_graph_save_on_cpu_cuda(self):
+        device_type = torch.accelerator.current_accelerator().type
+
         def f(x):
             a = x + 1
             return a * a
 
         # with grad
-        a = torch.ones(1, requires_grad=True, device="cuda")
+        a = torch.ones(1, requires_grad=True, device=device_type)
         y = f(a)
-        memory_with_grad = torch.cuda.memory_allocated()
+        memory_with_grad = (
+            torch.cuda.memory_allocated() if TEST_CUDA else torch.xpu.memory_allocated()
+        )
 
         del a
         del y
 
         # without grad
-        a = torch.ones(1, requires_grad=True, device="cuda")
+        a = torch.ones(1, requires_grad=True, device=device_type)
         with torch.no_grad():
             y = f(a)
-        memory_without_grad = torch.cuda.memory_allocated()
+        memory_without_grad = (
+            torch.cuda.memory_allocated() if TEST_CUDA else torch.xpu.memory_allocated()
+        )
 
         self.assertGreater(memory_with_grad, memory_without_grad)
 
@@ -10598,15 +10932,20 @@ def f(x):
 
         # with hooks
         with torch.autograd.graph.save_on_cpu():
-            a = torch.ones(1, requires_grad=True, device="cuda")
+            a = torch.ones(1, requires_grad=True, device=device_type)
             y = f(a)
-            memory_with_hooks = torch.cuda.memory_allocated()
+            memory_with_hooks = (
+                torch.cuda.memory_allocated()
+                if TEST_CUDA
+                else torch.xpu.memory_allocated()
+            )
             self.assertEqual(memory_with_hooks, memory_without_grad)
 
-    @unittest.skipIf(not TEST_CUDA, "test requires CUDA")
+    @unittest.skipIf(not TEST_CUDA and not TEST_XPU, "test requires CUDA and XPU")
     def test_scalar_grad_mixed_device(self):
+        device_type = torch.accelerator.current_accelerator().type
         x = torch.tensor(1.0, requires_grad=True)
-        y = torch.randn(2, 2, device="cuda")
+        y = torch.randn(2, 2, device=device_type)
         out = x * y
         out.sum().backward()
 
@@ -11940,24 +12279,58 @@ def fn():
 
 # Generic device type autograd tests.
 class TestAutogradDeviceType(TestCase):
-    def test_min_max_median_backprops_to_all_values(self, device):
+    def test_min_max_aminmax_median_backprops_to_all_values(self, device):
+        # 1) Test min/max/median/nanmedian on both a non NaN and all NaN tensor
         for f in [torch.min, torch.max, torch.median, torch.nanmedian]:
-            x1 = torch.tensor(
-                [1.0, 0.0, 1.0, 0.0, 1.0, 0.0], device=device, requires_grad=True
-            )
-            x2 = torch.tensor(
-                [float("nan"), float("nan"), float("nan")], requires_grad=True
-            )
-            for x in [x1, x2]:
-                y = f(x)
+            with self.subTest(f=f):
+                x1 = torch.tensor(
+                    [1.0, 0.0, 1.0, 0.0, 1.0, 0.0], device=device, requires_grad=True
+                )
+                x2 = torch.tensor(
+                    [float("nan"), float("nan"), float("nan")],
+                    device=device,
+                    requires_grad=True,
+                )
+                for x in [x1, x2]:
+                    y = f(x)
+                    y.backward()
+                    self.assertEqual(x.grad.sum(), 1.0)
+                    self.assertEqual((x.grad == 1 / 3).sum(), 3)
+
+        # 2) Explicit amin/amax plus the two components of aminmax
+        def amin2(x):
+            return torch.aminmax(x)[0]  # min part
+
+        def amax2(x):
+            return torch.aminmax(x)[1]
+
+        for f in [torch.amin, torch.amax, amax2, amin2]:
+            with self.subTest(f=f):
+                x1 = torch.tensor(
+                    [1.0, 0.0, 1.0, 0.0, 1.0, 0.0],
+                    device=device,
+                    requires_grad=True,
+                )
+                y = f(x1)
                 y.backward()
-                self.assertEqual(x.grad.sum(), 1.0)
-                self.assertEqual((x.grad == 1 / 3).sum(), 3)
-
-    def test_scatter_index_reduce_amin_amax_backprops_to_all_values(self, device):
+                self.assertEqual(x1.grad.sum(), 1.0)
+                self.assertEqual((x1.grad == 1.0 / 3.0).sum(), 3)
+
+        # 3) Both min and max grads active simultaneously — exercises the add_ path
+        # in aminmax_backward when both grad_min and grad_max are defined.
+        # min ties at indices 1,3 → each gets 0.5; max ties at indices 0,2 → each gets 0.5
+        with self.subTest("aminmax_both_grads"):
+            x = torch.tensor([3.0, 1.0, 3.0, 1.0], device=device, requires_grad=True)
+            min_val, max_val = torch.aminmax(x)
+            (min_val + max_val).backward()
+            self.assertEqual(x.grad, torch.tensor([0.5, 0.5, 0.5, 0.5], device=device))
+
+    def test_scatter_index_reduce_amin_amax_aminmax_backprops_to_all_values(
+        self, device
+    ):
         # tests that gradients are evenly distributed when there are multiple max/min values
         # tested here instead of adding a SampleInput as the backward for this case is non-differentiable for gradgrad
-        # as is the case for test_min_max_median_backprops_to_all_values above
+        # as is the case for test_min_max_aminmax_median_backprops_to_all_values above
         fns = (torch.scatter_reduce, torch.index_reduce)
         reduces = ("amin", "amax")
         for fn, reduction in product(fns, reduces):
@@ -12234,7 +12607,7 @@ def forward(ctx, inp):
             @staticmethod
             def backward(ctx, grad):
                 # Reentrant backward in child will take much longer.
-                reentrant_root.backward()
+                reentrant_root.backward(grad)
                 return grad
 
         # Parent gpu graph.
@@ -12251,6 +12624,10 @@ def backward(ctx, grad):
         self.assertIsNone(t1.grad)
         self.assertIsNone(t3.grad)
 
+    @skipCUDAIf(
+        True,
+        "Flaky test: https://github.com/pytorch/pytorch/issues/86735",
+    )
     @onlyCUDA
     def test_reentrant_parent_error_on_cpu(self, device):
         def _get_cuda_memory_usage():
@@ -14843,6 +15220,74 @@ def fn2(x):
         self.assertEqual(counter[0], 1)
 
 
+@contextlib.contextmanager
+def _counter_op(name):
+    """Yields (op, counts, idx_log) where op is a custom op that counts
+    invocations. idx_log maps call_idx (passed by caller) to replay count."""
+    counts = [0]
+    idx_log: dict = {}
+    with torch.library._scoped_library("test_ckpt", "FRAGMENT"):
+
+        @torch.library.custom_op(f"test_ckpt::{name}", mutates_args=())
+        def op(x: torch.Tensor, call_idx: int) -> torch.Tensor:
+            counts[0] += 1
+            idx_log[call_idx] = idx_log.get(call_idx, 0) + 1
+            return x.sin()
+
+        def setup_context(ctx, inputs, output):
+            ctx.save_for_backward(inputs[0])
+
+        def backward(ctx, grad):
+            (x,) = ctx.saved_tensors
+            return grad * x.cos(), None
+
+        op.register_autograd(backward, setup_context=setup_context)
+
+        yield op, counts, idx_log
+
+
+class _AutoNamingMode(TorchDispatchMode):
+    """Test helper: names output tensors as ``fqn_op_count[_outputidx]``."""
+
+    def __init__(self):
+        from torch.utils.module_tracker import ModuleTracker
+
+        self._tracker = ModuleTracker()
+        self._func_counter: dict = defaultdict(int)
+        self.names = WeakTensorKeyDictionary()
+
+    def __enter__(self):
+        self._tracker.__enter__()
+        return super().__enter__()
+
+    def __exit__(self, *args):
+        self._tracker.__exit__(*args)
+        return super().__exit__(*args)
+
+    def __torch_dispatch__(self, func, types, args=(), kwargs=None):
+        out = func(*args, **(kwargs or {}))
+        parents = self._tracker.parents - {"Global"}
+        fqn = max(parents, key=len) if parents else "Global"
+        op_name = func.__name__ if hasattr(func, "__name__") else str(func)
+        key = (fqn, func)
+        count = self._func_counter[key]
+        self._func_counter[key] += 1
+        multi_output = (
+            isinstance(out, (tuple, list))
+            and sum(isinstance(o, torch.Tensor) for o in out) > 1
+        )
+        if isinstance(out, torch.Tensor):
+            self.names[out] = f"{fqn}_{op_name}_{count}"
+        elif isinstance(out, (tuple, list)):
+            for i, o in enumerate(out):
+                if isinstance(o, torch.Tensor):
+                    name = f"{fqn}_{op_name}_{count}"
+                    if multi_output:
+                        name += f"_{i}"
+                    self.names[o] = name
+        return out
+
+
 class TestSelectiveActivationCheckpoint(TestCase):
     @unittest.skipIf(not TEST_CUDA, "requires CUDA")
     def test_flops_and_mem(self):
@@ -15262,6 +15707,83 @@ def fn(x):
         with self.assertRaisesRegex(RuntimeError, "Trying to backward an extra time"):
             out.sum().backward(retain_graph=True)
 
+    @skipIfTorchDynamo("torch dispatch modes don't support compile")
+    def test_auto_naming_mode_names(self):
+        with _counter_op("my_op") as (my_op, my_count, idx_log):
+
+            class Block(torch.nn.Module):
+                def forward(self, x, counter):
+                    x = my_op(x, counter[0])
+                    counter[0] += 1
+                    x = my_op(x, counter[0])
+                    counter[0] += 1
+                    x = my_op(x, counter[0])
+                    counter[0] += 1
+                    return x
+
+            class Model(torch.nn.Module):
+                def __init__(self):
+                    super().__init__()
+                    self.layers = torch.nn.ModuleList([Block(), Block()])
+
+                def forward(self, x):
+                    counter = [0]
+                    for layer in self.layers:
+                        x = layer(x, counter)
+                    return x
+
+            mod = Model()
+            naming = _AutoNamingMode()
+
+            save_names = {
+                "Model.layers.0_my_op.default_1",
+                "Model.layers.1_my_op.default_0",
+                "Model.layers.1_my_op.default_2",
+            }
+
+            fwd_decisions: list = []
+            fwd_idx = [0]
+
+            def policy_fn(ctx, op, *args, **kwargs):
+                if ctx.is_recompute:
+                    decision = fwd_decisions[fwd_idx[0]]
+                    fwd_idx[0] += 1
+                    return decision
+                out = ctx.op_output
+                decision = CheckpointPolicy.PREFER_RECOMPUTE
+                if isinstance(out, torch.Tensor):
+                    name = naming.names.get(out)
+                    if name in save_names:
+                        decision = CheckpointPolicy.MUST_SAVE
+                fwd_decisions.append(decision)
+                return decision
+
+            x = torch.randn(4, requires_grad=True)
+            context_fn = functools.partial(
+                create_selective_checkpoint_contexts, policy_fn
+            )
+            with naming:
+                out = checkpoint(
+                    lambda x: mod(x),
+                    x,
+                    use_reentrant=False,
+                    context_fn=context_fn,
+                )
+                out.sum().backward()
+
+            self.assertEqual(
+                idx_log,
+                {
+                    0: 2,  # Model.layers.0_my_op.default_0 -> recomputed
+                    1: 1,  # Model.layers.0_my_op.default_1 -> saved
+                    2: 2,  # Model.layers.0_my_op.default_2 -> recomputed
+                    3: 1,  # Model.layers.1_my_op.default_0 -> saved
+                    4: 2,  # Model.layers.1_my_op.default_1 -> recomputed
+                    5: 1,  # Model.layers.1_my_op.default_2 -> saved
+                },
+            )
+            self.assertEqual(my_count[0], 9)
+
 
 class TestAutogradMultipleDispatch(TestCase):
     def test_autograd_multiple_dispatch_registrations(self, device):
@@ -15595,6 +16117,17 @@ def unwrap(t):
         )
         self.assertEqual(x.grad, 2 * torch.ones_like(x))
 
+    def test_trace_backward_nonsquare_171704(self):
+        # https://github.com/pytorch/pytorch/issues/171704
+        for shape in [(5, 2), (7, 3), (4, 1), (1, 5), (2, 5)]:
+            with self.subTest(shape=shape):
+                x = torch.randn(shape, dtype=torch.float64, requires_grad=True)
+                torch.trace(x).backward()
+                expected = torch.zeros(shape, dtype=torch.float64)
+                for i in range(min(shape)):
+                    expected[i, i] = 1.0
+                self.assertEqual(x.grad, expected)
+
 
 # Import test cases from below autograd/ here. These are found
 # implicitly by the loader, so Flake8 thinks they are unused, hence
diff --git a/test/test_binary_ufuncs.py b/test/test_binary_ufuncs.py
index af6478870bd46..29dff031d6e5a 100644
--- a/test/test_binary_ufuncs.py
+++ b/test/test_binary_ufuncs.py
@@ -50,7 +50,6 @@
 )
 from torch.testing._internal.common_methods_invocations import (
     binary_ufuncs,
-    binary_ufuncs_and_refs,
     generate_elementwise_binary_broadcasting_tensors,
     generate_elementwise_binary_extremal_value_tensors,
     generate_elementwise_binary_large_value_tensors,
@@ -505,7 +504,7 @@ def test_batch_vs_slicing(self, device, dtype, op):
     # NOTE: because the cross-product of all possible type promotion tests is huge, this
     #   just spot checks some handwritten cases.
     # NOTE: It may be possible to refactor this test into something simpler
-    @ops(binary_ufuncs_and_refs, dtypes=OpDTypes.none)
+    @ops(binary_ufuncs, dtypes=OpDTypes.none)
     def test_type_promotion(self, device, op):
         supported_dtypes = op.supported_dtypes(torch.device(device).type)
         make_lhs = partial(
diff --git a/test/test_bundled_inputs.py b/test/test_bundled_inputs.py
index 6b98b5c5ea18c..bf9d24f0b8bb3 100644
--- a/test/test_bundled_inputs.py
+++ b/test/test_bundled_inputs.py
@@ -4,7 +4,6 @@
 
 import io
 import textwrap
-from typing import Optional
 
 import torch
 import torch.utils.bundled_inputs
@@ -333,8 +332,8 @@ def test_dict_args(self):
         class MyModel(torch.nn.Module):
             def forward(
                 self,
-                arg1: Optional[dict[str, torch.Tensor]],
-                arg2: Optional[list[torch.Tensor]],
+                arg1: dict[str, torch.Tensor] | None,
+                arg2: list[torch.Tensor] | None,
                 arg3: torch.Tensor,
             ):
                 if arg1 is None:
diff --git a/test/test_cuda.py b/test/test_cuda.py
index 72a4e5e1296a6..41289162b9a70 100644
--- a/test/test_cuda.py
+++ b/test/test_cuda.py
@@ -3,6 +3,7 @@
 
 import contextlib
 import ctypes
+import functools
 import gc
 import json
 import os
@@ -19,6 +20,7 @@
 from copy import deepcopy
 from itertools import product
 from random import randint
+from unittest.mock import patch
 
 import psutil
 
@@ -36,8 +38,11 @@
 from torch.testing._internal.common_cuda import (
     _create_scaling_case,
     _get_torch_cuda_version,
+    blas_library_context,
     PLATFORM_SUPPORTS_GREEN_CONTEXT,
+    PLATFORM_SUPPORTS_WORKQUEUE_CONFIG,
     SM70OrLater,
+    SM89OrLater,
     TEST_CUDNN,
     TEST_MULTIGPU,
     tf32_on_and_off,
@@ -62,6 +67,7 @@
     gcIfJetson,
     get_cycles_per_ms,
     instantiate_parametrized_tests,
+    IS_ARM64,
     IS_FBCODE,
     IS_JETSON,
     IS_LINUX,
@@ -81,6 +87,7 @@
     skipCUDANonDefaultStreamIf,
     skipIfRocm,
     skipIfRocmArch,
+    skipIfRocmVersionLessThan,
     slowTest,
     subtest,
     TemporaryFileName,
@@ -97,7 +104,7 @@
 
 
 requiresCppContext = unittest.skipUnless(
-    IS_X86 and IS_LINUX, "cpp contexts are x86 linux only"
+    (IS_X86 or IS_ARM64) and IS_LINUX, "cpp contexts are linux x86/aarch64 only"
 )
 
 # load_tests from common_utils is used to automatically filter tests for
@@ -105,7 +112,8 @@
 load_tests = load_tests  # noqa: PLW0127
 
 try:
-    # import torchvision.models  # noqa: F401
+    import torchvision.models  # noqa: F401
+
     # from torchvision.models import resnet18  # noqa: F401
 
     HAS_TORCHVISION = True
@@ -129,6 +137,55 @@
 
 _cycles_per_ms = None
 
+_wait_for_cpu_kernel = None
+
+
+def skip_background_threads_on_windows(f):
+    @functools.wraps(f)
+    def wrapped(self, **kwargs):
+        if IS_WINDOWS and SM89OrLater and kwargs.get("use_background_threads"):
+            raise unittest.SkipTest("using background threads fails on Windows")
+        return f(self, **kwargs)
+
+    return wrapped
+
+
+def get_wait_for_cpu_kernel():
+    """Returns a compiled CUDA/HIP spin-wait kernel that blocks the GPU stream
+    until the host sets a pinned int32 flag to non-zero. On NVIDIA, requires
+    SM70+ for the relaxed-system-scope PTX load. On AMD, uses
+    __hip_atomic_load with __HIP_MEMORY_SCOPE_SYSTEM (the AMDGCN equivalent).
+
+    Usage::
+
+        kernel = get_wait_for_cpu_kernel()
+        flag = torch.zeros(1, dtype=torch.int32, device="cpu").pin_memory()
+        with torch.cuda.stream(s):
+            kernel(grid=(1, 1, 1), block=(1, 1, 1), args=[flag])
+        # stream s is now blocked until:
+        flag[0] = 1
+    """
+    global _wait_for_cpu_kernel
+    if _wait_for_cpu_kernel is None:
+        from torch.cuda import _compile_kernel
+
+        _wait_for_cpu_kernel = _compile_kernel(
+            r"""
+            __global__ void wait_for_cpu(int *pinned_cpu_flag) {
+                int flag = 0;
+                do {
+            #ifdef __HIP_PLATFORM_AMD__
+                    flag = __hip_atomic_load(pinned_cpu_flag, __ATOMIC_RELAXED, __HIP_MEMORY_SCOPE_SYSTEM);
+            #else
+                    asm volatile("ld.relaxed.sys.global.s32 %0, [%1];" : "=r"(flag) : "l"(pinned_cpu_flag) : "memory");
+            #endif
+                } while (flag == 0);
+            }
+            """,
+            "wait_for_cpu",
+        )
+    return _wait_for_cpu_kernel
+
 
 @unittest.skipIf(not TEST_CUDA, "CUDA not available, skipping tests")
 @torch.testing._internal.common_utils.markDynamoStrictTest
@@ -205,12 +262,21 @@ def empty_stats():
                 "active_requests.peak": 0,
             }
 
+        # Deltas from baseline captured after cleanup (not absolute totals), so
+        # ordering against other tests cannot leave host stats in a broken state.
+        baseline = {}
+
         def check_stats(expected):
             stats = torch.cuda.host_memory_stats()
             for k, v in expected.items():
-                if v != stats[k]:
-                    print(f"key: {k}, expected: {v}, stats: {stats[k]}")
-                self.assertEqual(v, stats[k])
+                want = baseline[k] + v
+                actual = stats[k]
+                if actual != want:
+                    print(
+                        f"key: {k}, baseline: {baseline[k]}, delta: {v}, "
+                        f"want: {want}, got: {actual}"
+                    )
+                self.assertEqual(actual, want)
 
         # Setup the test cleanly
         alloc1 = 10
@@ -219,11 +285,12 @@ def check_stats(expected):
         alloc2_aligned = 32
         expected = empty_stats()
 
-        # Reset any lingering state
         gc.collect()
         torch._C._host_emptyCache()
+        torch.cuda.reset_peak_host_memory_stats()
+        torch.cuda.reset_accumulated_host_memory_stats()
+        baseline.update(torch.cuda.host_memory_stats())
 
-        # Check that stats are empty
         check_stats(expected)
 
         # Make first allocation and check stats
@@ -290,6 +357,9 @@ def test_pinned_memory_empty_cache(self):
                 "pinned_use_cuda_host_register:False"
             )
 
+    # Pinned allocator background thread does not shut down cleanly on Windows
+    # Python process hangs
+    @unittest.skipIf(IS_WINDOWS and SM89OrLater, "Fails on windows with SM89+")
     def test_pinned_memory_use_background_threads(self):
         script = """
 import torch
@@ -329,6 +399,10 @@ def test_memory_allocation(self):
                 torch.cuda.caching_allocator_delete(mem)
                 self.assertEqual(torch.cuda.memory_allocated(), prev)
 
+    def test_caching_allocator_alloc_negative_size(self):
+        with self.assertRaisesRegex(ValueError, "Invalid memory size"):
+            torch.cuda.memory.caching_allocator_alloc(-1024)
+
     def test_memory_stats(self):
         gc.collect()
         torch.cuda.empty_cache()
@@ -432,9 +506,15 @@ def test_out_of_memory(self):
         tensor.fill_(1)
         self.assertTrue((tensor == 1).all())
 
+    # CUDA memory allocations on windows do not OOM on rtx even when they cross allowed memory
+    # Skip test until this is investigated
+    @unittest.skipIf(IS_WINDOWS and SM89OrLater, "Fails on windows with SM89+")
     @unittest.skipIf(
         TEST_CUDAMALLOCASYNC or IS_JETSON, "Segmentation fault (core dumped)"
     )
+    # On MI350 (gfx950) allocating ~1.02x of the reported free memory does not
+    # raise an OOM RuntimeError, so the assertRaisesRegex below fails.
+    @skipIfRocmArch(MI350_ARCH)
     @serialTest()
     def test_out_of_memory_retry(self):
         torch.cuda.empty_cache()
@@ -460,6 +540,7 @@ def test_out_of_memory_retry(self):
         IS_JETSON, "oom reporting has issues on jetson igx due to partial nvml support"
     )
     def test_set_per_process_memory_fraction(self):
+        torch.cuda.empty_cache()
         orig = torch.cuda.get_per_process_memory_fraction(0)
         torch.cuda.reset_peak_memory_stats(0)
         try:
@@ -613,6 +694,9 @@ def test_serialization_array_with_storage(self):
         q_copy[1].fill_(10)
         self.assertEqual(q_copy[3], torch.cuda.IntStorage(10).fill_(10))
 
+    @unittest.skipIf(
+        IS_WINDOWS and SM89OrLater, "preferred_blas_library not supported on Windows"
+    )
     @unittest.skipIf(IS_FBCODE or IS_SANDCASTLE, "Does not work in fbcode yet")
     @setBlasBackendsToDefaultFinally
     def test_preferred_blas_library_settings(self):
@@ -657,11 +741,13 @@ def _check_default():
             torch.backends.cuda.preferred_blas_library(1.0)
         # check env var override
         custom_envs = [
-            {"TORCH_BLAS_PREFER_CUBLASLT": "1"},
-            {"TORCH_BLAS_PREFER_HIPBLASLT": "1"},
+            ({"TORCH_BLAS_PREFER_CUBLASLT": "1"}, "_BlasBackend.Cublaslt"),
+            ({"TORCH_BLAS_PREFER_HIPBLASLT": "1"}, "_BlasBackend.Cublaslt"),
+            ({"TORCH_BLAS_PREFER_CUBLASLT": "0"}, "_BlasBackend.Cublas"),
+            ({"TORCH_BLAS_PREFER_HIPBLASLT": "0"}, "_BlasBackend.Cublas"),
         ]
         test_script = "import torch;print(torch.backends.cuda.preferred_blas_library())"
-        for env_config in custom_envs:
+        for env_config, expected in custom_envs:
             env = os.environ.copy()
             for key, value in env_config.items():
                 env[key] = value
@@ -670,12 +756,23 @@ def _check_default():
                 .decode("ascii")
                 .strip()
             )
-            self.assertEqual("_BlasBackend.Cublaslt", r)
+            self.assertEqual(expected, r)
 
+        # explicitly check default when no env vars are set
+        if not any(
+            os.environ.get(v)
+            for v in ("TORCH_BLAS_PREFER_CUBLASLT", "TORCH_BLAS_PREFER_HIPBLASLT")
+        ):
+            torch.backends.cuda.preferred_blas_library("default")
+            _check_default()
+
+    @unittest.skipIf(
+        IS_WINDOWS and SM89OrLater, "preferred_blas_library not supported on Windows"
+    )
     @unittest.skipIf(TEST_CUDAMALLOCASYNC, "temporarily disabled for async")
-    @setBlasBackendsToDefaultFinally
+    @serialTest()
+    @blas_library_context("cublas")
     def test_cublas_workspace_explicit_allocation(self):
-        torch.backends.cuda.preferred_blas_library("cublas")
         a = torch.randn(7, 7, device="cuda", requires_grad=False)
         if torch.version.hip:
             default_workspace_size = 1024 * 32 * 1024  # :1024:32  32MiB
@@ -683,14 +780,14 @@ def test_cublas_workspace_explicit_allocation(self):
             gcn_arch = str(
                 torch.cuda.get_device_properties(0).gcnArchName.split(":", 1)[0]
             )
-            if "gfx94" in gcn_arch or "gfx95" in gcn_arch:
+            if gcn_arch in ["gfx942", "gfx950", "gfx1250"]:
                 default_workspace_size = 1024 * 128 * 1024  # :1024:128
         else:
             default_workspace_size = (
                 4096 * 2 * 1024 + 16 * 8 * 1024
             )  # :4096:2:16:8  8MiB
-            # different size (32 MiB) expected on Hopper GPU
-            if torch.cuda.get_device_capability() == (9, 0):
+            # different size (32 MiB) expected on Hopper/Blackwell GPU
+            if torch.cuda.get_device_capability()[0] in (9, 10, 12):
                 default_workspace_size = 4096 * 8 * 1024
 
         def check_workspace_size(inp):
@@ -718,6 +815,26 @@ def check_workspace_size(inp):
 
         torch._C._cuda_clearCublasWorkspaces()
 
+    @unittest.skipIf(TEST_CUDAMALLOCASYNC, "temporarily disabled for async")
+    @unittest.skipIf(IS_FBCODE, "not enabled by default on fbcode")
+    @unittest.skipIf(TEST_WITH_ROCM, "not enabled by default on rocm")
+    @setBlasBackendsToDefaultFinally
+    @serialTest()
+    def test_cublas_unified_workspace(self):
+        torch.cuda.reset_peak_memory_stats()
+        # We run a cuBLAS matmul, guaranteeing a workspace present on the current stream
+        torch.backends.cuda.preferred_blas_library("cublas")
+        x = torch.randn(128, 128, device="cuda")
+        torch.matmul(x, x)
+        # Stash the maximum memory allocated after running matmuls
+        warmed_alloc = torch.cuda.max_memory_allocated()
+        torch.backends.cuda.preferred_blas_library("cublaslt")
+        torch.matmul(x, x)
+        lt_alloc = torch.cuda.max_memory_allocated()
+        # With unified workspaces, the peak memory allocation should not increase after
+        # switching to Lt, otherwise the temporary allocation would bump the peak
+        self.assertEqual(warmed_alloc, lt_alloc)
+
     def test_cublas_allow_tf32_get_set(self):
         skip_tf32_cublas = "TORCH_ALLOW_TF32_CUBLAS_OVERRIDE" in os.environ and int(
             os.environ["TORCH_ALLOW_TF32_CUBLAS_OVERRIDE"]
@@ -955,6 +1072,37 @@ def test_manual_seed(self):
             self.assertEqual(a, b)
             self.assertEqual(torch.cuda.initial_seed(), 2)
 
+    def test_lazy_call_reentrant_set_rng_state_does_not_deadlock(self):
+        # Separate process: a regression deadlocks the interpreter (non-reentrant lock).
+        # Happy path is usually a few seconds; allow margin for slow CI / CUDA init.
+        timeout_sec = 15
+        script = (
+            "import torch; "
+            "torch.cuda.init(); "
+            "state = torch.cuda.get_rng_state(); "
+            "torch.cuda._lazy_call(lambda: torch.cuda.set_rng_state(state)); "
+            "print('done')"
+        )
+        try:
+            proc = subprocess.run(
+                [sys.executable, "-c", script],
+                capture_output=True,
+                text=True,
+                timeout=timeout_sec,
+            )
+        except subprocess.TimeoutExpired as e:
+            self.fail(
+                f"lazy_call reentrancy subprocess did not finish within {timeout_sec}s "
+                "(likely deadlock in torch.cuda._lazy_call); "
+                f"cmd={e.cmd!r}"
+            )
+        self.assertEqual(
+            proc.returncode,
+            0,
+            msg=f"stdout={proc.stdout!r}\nstderr={proc.stderr!r}",
+        )
+        self.assertIn("done", proc.stdout)
+
     def test_specify_improper_device_name(self):
         with tempfile.TemporaryDirectory() as tmpdir:
             fname = os.path.join(tmpdir, "tempfile.pt")
@@ -1113,6 +1261,7 @@ def test_generic_stream_event(self):
         self.assertTrue(issubclass(type(cuda_event), torch.Event))
         self.assertTrue(torch.Event in type(cuda_event).mro())
 
+    @unittest.skip("Fails with Triton 3.7")
     def test_stream_event_compatibility(self):
         s1 = torch.cuda.Stream()
         s2 = torch.cuda.Stream()
@@ -2110,8 +2259,10 @@ def test_graph_is_current_stream_capturing(self):
             with torch.cuda.stream(s):
                 g = torch.cuda.CUDAGraph()
                 self.assertFalse(torch.cuda.is_current_stream_capturing())
+                self.assertFalse(s.is_capturing())
                 g.capture_begin()
                 self.assertTrue(torch.cuda.is_current_stream_capturing())
+                self.assertTrue(s.is_capturing())
                 g.capture_end()
 
     @unittest.skipIf(
@@ -2214,6 +2365,137 @@ def get_final_offsets_of_states(generator_state):
             # Compare the states generated outside and inside the graph
             self.assertEqual(random_values, graphed_random_values)
 
+    @skipIfRocmVersionLessThan((7, 14))
+    @unittest.skipIf(
+        not TEST_CUDA_GRAPH, "CUDA >= 11.0 or ROCM >= 5.3 required for graphs"
+    )
+    def test_graph_rng_after_failed_capture(self):
+        """Test that a stream can be captured again for RNG after a failed capture."""
+        if TEST_WITH_ROCM and self.expandable_segments:
+            self.skipTest(
+                "ROCm expandable segments has known issue with graph capture recovery - #179911"
+            )
+        stream = torch.cuda.Stream()
+        graph = torch.cuda.CUDAGraph()
+        x = torch.ones(1, device="cuda")
+
+        with torch.cuda.stream(stream):
+            graph.capture_begin()
+            with self.assertRaises(RuntimeError):
+                (x + 1).item()
+            with self.assertRaises(RuntimeError):
+                graph.capture_end()
+
+        torch.cuda.current_stream().wait_stream(stream)
+
+        result = torch.randn(4, device="cuda")
+        self.assertEqual(result.shape, (4,))
+
+        new_graph = torch.cuda.CUDAGraph()
+        buf = torch.empty(4, device="cuda")
+        with torch.cuda.stream(stream):
+            new_graph.capture_begin()
+            buf.copy_(torch.randn_like(buf))
+            new_graph.capture_end()
+        torch.cuda.current_stream().wait_stream(stream)
+        buf.zero_()
+        new_graph.replay()
+        torch.cuda.synchronize()
+        self.assertFalse(torch.allclose(buf, torch.zeros_like(buf)))
+
+    @skipIfRocmVersionLessThan((7, 14))
+    @unittest.skipIf(
+        not TEST_CUDA_GRAPH, "CUDA >= 11.0 or ROCM >= 5.3 required for graphs"
+    )
+    def test_graph_capture_error_releases_reserved_segments(self):
+        # A failed capture (capture_end raises) must still release the graph's
+        # private mempool on reset(); otherwise its reserved segments leak for
+        # the rest of the process. Unlike test_graph_rng_after_failed_capture,
+        # which asserts on active.all.current (no leaked live allocations), this
+        # checks memory_reserved(): a leaked pool can show zero active bytes yet
+        # still hold reserved segments, so only this assertion catches the leak.
+        torch.cuda.synchronize()
+        torch.cuda.empty_cache()
+        reserved_before = torch.cuda.memory_reserved()
+
+        s = torch.cuda.Stream()
+        s.wait_stream(torch.cuda.current_stream())
+        g = torch.cuda.CUDAGraph()
+        with torch.cuda.stream(s):
+            g.capture_begin()
+            x = torch.empty(8 * 1024 * 1024, device="cuda")
+            # A synchronizing op during capture invalidates it, but we must
+            # still call capture_end() so the stream's capture is ended (with an
+            # error); otherwise the stream is left mid-capture and teardown fails.
+            with self.assertRaises(RuntimeError):
+                torch.cuda.synchronize()
+            with self.assertRaises(RuntimeError):
+                g.capture_end()
+
+        del x
+        g.reset()
+        del g
+        torch.cuda.synchronize()
+        torch.cuda.empty_cache()
+        self.assertEqual(torch.cuda.memory_reserved(), reserved_before)
+
+    @unittest.skipIf(
+        not TEST_CUDA_GRAPH, "CUDA >= 11.0 or ROCM >= 5.3 required for graphs"
+    )
+    def test_graph_rng_concurrent_replay_on_different_streams(self):
+        """Concurrent replay of two graphs sharing a generator on different streams.
+
+        With per-generator state (old code), this would race on the shared
+        rng_state_offset_extragraph_ tensor. With per-(generator, capture_id)
+        state, each graph has its own tensors.
+        """
+        seed = 1234
+        shape = (64,)
+
+        torch.manual_seed(seed)
+        ref0 = torch.randn(shape, device="cuda")
+        ref1 = torch.randn(shape, device="cuda")
+
+        torch.manual_seed(seed)
+        g0 = torch.cuda.CUDAGraph()
+        g1 = torch.cuda.CUDAGraph()
+        s_cap = torch.cuda.Stream()
+        buf0 = torch.empty(shape, device="cuda")
+        buf1 = torch.empty(shape, device="cuda")
+
+        with torch.cuda.stream(s_cap):
+            g0.capture_begin()
+            buf0.copy_(torch.randn_like(buf0))
+            g0.capture_end()
+
+        torch.cuda.current_stream().wait_stream(s_cap)
+
+        with torch.cuda.stream(s_cap):
+            g1.capture_begin()
+            buf1.copy_(torch.randn_like(buf1))
+            g1.capture_end()
+
+        torch.cuda.current_stream().wait_stream(s_cap)
+
+        s0 = torch.cuda.Stream()
+        s1 = torch.cuda.Stream()
+
+        buf0.zero_()
+        buf1.zero_()
+
+        s0.wait_stream(torch.cuda.current_stream())
+        s1.wait_stream(torch.cuda.current_stream())
+
+        with torch.cuda.stream(s0):
+            g0.replay()
+        with torch.cuda.stream(s1):
+            g1.replay()
+
+        torch.cuda.synchronize()
+
+        self.assertEqual(buf0, ref0)
+        self.assertEqual(buf1, ref1)
+
     @unittest.skipIf(
         not TEST_CUDA_GRAPH, "CUDA >= 11.0 or ROCM >= 5.3 required for graphs"
     )
@@ -2224,14 +2506,16 @@ def clear_cuda_cache():
             torch.cuda.empty_cache()
 
         # Executes a simple graph task which includes capturing and executing a random number generation within a CUDA graph.
-        def simple_graph_task(graph):
+        def simple_graph_task(graph, default_generator, generator_states):
             s = torch.cuda.Stream()
             with torch.cuda.stream(s):
                 graph.capture_begin()
-                torch.rand(1, device="cuda")
+                for generator_state in generator_states:
+                    default_generator.graphsafe_set_state(generator_state)
+                    torch.rand(1, device="cuda")
                 graph.capture_end()
             torch.cuda.current_stream().wait_stream(s)
-            graph.replay()  # Replays the captured operations
+            graph.replay()
 
         def get_memory_stats():
             stats = torch.cuda.memory_stats()
@@ -2248,22 +2532,22 @@ def test(num_graphs, num_generators):
 
             # Allocate and manage generator states
             default_generator = torch.cuda.default_generators[0]
-            generators = [default_generator.graphsafe_get_state()]
+            generator_states = [default_generator.graphsafe_get_state()]
 
-            # Starts from 1 as one state is already added
             for _ in range(1, num_generators):
-                generators.append(default_generator.clone_state())
+                generator_states.append(default_generator.clone_state())
 
             for graph in graphs:
-                for generator_state in generators:
+                for generator_state in generator_states:
                     graph.register_generator_state(generator_state)
-                simple_graph_task(graph)
+                simple_graph_task(graph, default_generator, generator_states)
 
             # Assert conditions after graph tasks
             num_blocks, total_size = get_memory_stats()
-            # The allocated blocks should only be proportional to the number of generators
-            expected_blocks_diff = 2 * num_generators
-            expected_size_diff = 2 * 512 * num_generators  # Each block's size is 512
+            # Each (generator, graph) pair gets 2 tensors (seed + offset)
+            expected_captured_states = num_generators * num_graphs
+            expected_blocks_diff = 2 * expected_captured_states
+            expected_size_diff = 2 * 512 * expected_captured_states
 
             self.assertEqual(
                 (num_blocks - baseline_num_blocks),
@@ -2357,9 +2641,10 @@ def test_graph_debugdump(self):
             g.debug_dump(os.path.join(tempdir, "out_multi_stream.dot"))
 
     @unittest.skipIf(
-        not TEST_CUDA_GRAPH or TEST_WITH_ROCM,
-        "CUDA >= 11.0 required for external events in cuda graphs. rocm does not support external events",
+        not TEST_CUDA_GRAPH,
+        "CUDA >= 11.0 / ROCm >= 7.0 required for external events in cuda graphs",
     )
+    @skipIfRocmVersionLessThan((7, 0))
     def test_graph_timing(self):
         torch.cuda.empty_cache()
         x = torch.randn(10240000, device="cuda")
@@ -2448,9 +2733,8 @@ def test_graph_capture_oom(self):
         not TEST_CUDA_GRAPH, "CUDA >= 11.0 or ROCM >= 5.3 required for graphs"
     )
     @serialTest()
-    @setBlasBackendsToDefaultFinally
+    @blas_library_context("cublas")
     def test_repeat_graph_capture_cublas_workspace_memory(self):
-        torch.backends.cuda.preferred_blas_library("cublas")
         (x, y, z) = 1024, 512, 64
         a = torch.rand((x, y), device="cuda")
         b = torch.rand((y, z), device="cuda")
@@ -2458,19 +2742,19 @@ def test_repeat_graph_capture_cublas_workspace_memory(self):
         # warmup
         torch.mm(a, b)
 
-        free_bytes_before, total_bytes = torch.cuda.mem_get_info()
-        used_gb_before = (total_bytes - free_bytes_before) / 1e9
-
-        for _ in range(100):
+        for i in range(100):
             torch_graph = torch.cuda.CUDAGraph()
             with torch.cuda.graph(torch_graph):
                 torch.mm(a, b)
             torch_graph.replay()
+            if i == 0:
+                free_bytes_before, total_bytes = torch.cuda.mem_get_info()
+                used_gb_before = (total_bytes - free_bytes_before) / 1e9
 
         free_bytes_after, _ = torch.cuda.mem_get_info()
         used_gb_after = (total_bytes - free_bytes_after) / 1e9
 
-        self.assertFalse(used_gb_before + 0.1 < used_gb_after)
+        self.assertGreater(0.005 + used_gb_before, used_gb_after)
 
     @unittest.skipIf(
         not TEST_CUDA_GRAPH, "CUDA >= 11.0 or ROCM >= 5.3 required for graphs"
@@ -2937,9 +3221,12 @@ def test_graph_memory_stats_and_use_result_after_destroy_graph(self):
         elem = 4
 
         # this was annoying to write but stresses the expectations pretty rigorously
+        # For small_pool cases, delta_cudaMallocs and delta_cudaMalloc_bytes include
+        # an extra kSmallBuffer segment for the per-capture RNG state tensors, which
+        # are allocated on the default stream (separate from stream s).
         cases = (
-            (512 // elem, 1, kSmallBuffer, kSmallBuffer, "small_pool"),
-            (kSmallSize // elem, 2, 2 * kSmallBuffer, kSmallBuffer, "small_pool"),
+            (512 // elem, 2, 2 * kSmallBuffer, kSmallBuffer, "small_pool"),
+            (kSmallSize // elem, 3, 3 * kSmallBuffer, kSmallBuffer, "small_pool"),
             ((kSmallSize + 512) // elem, 1, kLargeBuffer, kLargeBuffer, "large_pool"),
             (
                 (kMinLargeAlloc - 512) // elem,
@@ -2987,9 +3274,9 @@ def test_graph_memory_stats_and_use_result_after_destroy_graph(self):
             g = torch.cuda.CUDAGraph()
             s.wait_stream(torch.cuda.current_stream())
             with torch.cuda.stream(s):
-                # Allocation stat estimates assume input is created on the same stream as capture_begin()
-                # (in other words, the same stream silo as the rng offset holder, which is not allocated from the
-                # capture's private pool).
+                # Per-capture RNG state tensors are allocated on the default stream
+                # (not the capture stream), so they occupy a separate segment from
+                # user tensors created here on stream s.
                 a = torch.ones((numel,), device="cuda")
 
                 precapture_stats = torch.cuda.memory_stats()
@@ -3999,6 +4286,7 @@ def check_output(script: str) -> str:
             )
             self.assertEqual(rc, "3")
 
+    @unittest.skipIf(TEST_WITH_ROCM, "Failed onr ROCm due to rocprofiler-sdk issue AIPROFSDK-840")
     @unittest.skipIf(not TEST_WITH_ROCM, "not relevant for CUDA testing")
     def test_hip_device_count(self):
         """Validate device_count works with both CUDA/HIP visible devices"""
@@ -4073,6 +4361,10 @@ def test_gds_fails_in_ci(self):
             with self.assertRaisesRegex(RuntimeError, error_msg):
                 torch.cuda.gds.GdsFile(f, os.O_CREAT | os.O_RDWR)
 
+    @unittest.skipIf(
+        IS_WINDOWS, "test relies on fork; Windows multiprocessing uses spawn"
+    )
+    @unittest.skipIf(sys.version_info >= (3, 14), "test fails on Python 3.14+")
     def test_is_pinned_no_context(self):
         test_script = """\
 import torch
@@ -4123,7 +4415,15 @@ def worker(conn):
 
 @unittest.skipIf(not TEST_CUDA, "CUDA not available, skipping tests")
 @torch.testing._internal.common_utils.markDynamoStrictTest
-class TestCudaMallocAsync(TestCase):
+class TestCudaAllocator(TestCase):
+    def tearDown(self):
+        super().tearDown()
+        # Regression check: no test in this class should leave the runtime
+        # expandable_segments knob mismatched against the suite's env-derived
+        # baseline. This is the only class that toggles it.
+        md = torch.cuda.memory._snapshot()["allocator_settings"]
+        self.assertEqual(md["expandable_segments"], EXPANDABLE_SEGMENTS)
+
     @unittest.skipIf(
         TEST_CUDAMALLOCASYNC, "setContextRecorder not supported by CUDAMallocAsync"
     )
@@ -4166,6 +4466,7 @@ def test_memory_snapshot(self):
                         self.assertTrue("test_cuda.py" in f2.read())
             del unused
             del x
+            torch._C._cuda_clearCublasWorkspaces()
             torch.cuda.empty_cache()
             ss = torch.cuda.memory._snapshot()
             self.assertTrue(
@@ -4176,7 +4477,40 @@ def test_memory_snapshot(self):
         finally:
             torch.cuda.memory._record_memory_history(None)
 
-    @unittest.skipUnless(IS_X86 and IS_LINUX, "x86 linux only cpp unwinding")
+    @unittest.skipIf(
+        TEST_CUDAMALLOCASYNC, "setContextRecorder not supported by CUDAMallocAsync"
+    )
+    def test_allocation_traceback(self):
+        def this_is_a_funky_allocator():
+            return torch.rand(311, 411, device="cuda")
+
+        try:
+            torch.cuda.memory.empty_cache()
+            torch.cuda.memory._record_memory_history("all", stacks="python")
+
+            x = this_is_a_funky_allocator()
+            frames = torch.cuda.memory._allocation_traceback(x.data_ptr())
+            self.assertIsInstance(frames, list)
+            self.assertGreater(len(frames), 0)
+            func_names = [f["name"] for f in frames]
+            self.assertIn("this_is_a_funky_allocator", func_names)
+
+            # freed pointer should return None
+            ptr = x.data_ptr()
+            del x
+            torch.cuda.empty_cache()
+            self.assertIsNone(torch.cuda.memory._allocation_traceback(ptr))
+        finally:
+            torch.cuda.memory._record_memory_history(None)
+
+    def test_allocation_traceback_no_recording(self):
+        torch.cuda.memory._record_memory_history(None)
+        x = torch.rand(64, device="cuda")
+        self.assertIsNone(torch.cuda.memory._allocation_traceback(x.data_ptr()))
+
+    @unittest.skipUnless(
+        (IS_X86 or IS_ARM64) and IS_LINUX, "linux x86/aarch64 only cpp unwinding"
+    )
     def test_direct_traceback(self):
         from torch._C._profiler import gather_traceback, symbolize_tracebacks  # @manual
 
@@ -4309,7 +4643,6 @@ def foo(p):
         TEST_CUDAMALLOCASYNC, "setContextRecorder not supported by CUDAMallocAsync"
     )
     @requiresCppContext
-    @skipIfRocmArch(MI350_ARCH)
     def test_memory_plots(self):
         for context, stacks in (
             ("all", "all" if IS_LINUX else "python"),
@@ -4403,6 +4736,10 @@ def forward(self, x):
 
         try:
             torch.cuda.memory.empty_cache()
+            # Region ids come from Dynamo's process-global frame counter, which
+            # reset() zeroes; do it before compiling so the region strings are
+            # deterministic.
+            torch._dynamo.reset()
             input_tensor = torch.randn(1, 10, device="cuda")
             # Create an instance of the model
             model = MyModel()
@@ -4426,11 +4763,19 @@ def forward(self, x):
                 if len(allocation_sequence) > 0 and allocation_sequence[-1] == context:
                     continue
                 allocation_sequence.append(context)
-            self.assertTrue(allocation_sequence == expected_allocation_sequence)
+            self.assertEqual(allocation_sequence, expected_allocation_sequence)
         except RuntimeError as e:
             pass
         finally:
             torch.cuda.memory._record_memory_history(None)
+            # Clean up the model and input tensor to fix memory leaks in other tests
+            del compiled_model, model, input_tensor
+            ss = None
+            torch._dynamo.reset()
+            # This test requires to run gc.collec() to fix other memory tests
+            gc.collect()
+            torch.cuda.synchronize()
+            torch.cuda.empty_cache()
 
     @unittest.skipIf(
         TEST_CUDAMALLOCASYNC, "setContextRecorder not supported by CUDAMallocAsync"
@@ -4636,6 +4981,11 @@ def alloc(n):
             alloc(120)
         finally:
             torch.cuda.memory.set_per_process_memory_fraction(orig)
+            # Test toggles expandable_segments internally; restore the
+            # suite's baseline so subsequent tests see consistent state.
+            torch.cuda.memory._set_allocator_settings(
+                f"expandable_segments:{EXPANDABLE_SEGMENTS}"
+            )
 
     @serialTest()
     def test_garbage_collect_expandable(self):
@@ -4668,6 +5018,11 @@ def alloc(n):
             alloc(80)
         finally:
             torch.cuda.memory.set_per_process_memory_fraction(orig)
+            # Test toggles expandable_segments internally; restore the
+            # suite's baseline so subsequent tests see consistent state.
+            torch.cuda.memory._set_allocator_settings(
+                f"expandable_segments:{EXPANDABLE_SEGMENTS}"
+            )
 
     def test_allocator_settings(self):
         def power2_div(size, div_factor):
@@ -4787,6 +5142,55 @@ def power2_div(size, div_factor):
                 "pinned_num_register_threads:1024"
             )
 
+        # Test throw_on_cudamalloc_oom config parsing - valid formats
+        torch.cuda.memory._set_allocator_settings("throw_on_cudamalloc_oom:True")
+        torch.cuda.memory._set_allocator_settings("throw_on_cudamalloc_oom:False")
+
+        # Test throw_on_cudamalloc_oom config parsing - invalid formats
+        with self.assertRaises(ValueError):
+            torch._C._accelerator_setAllocatorSettings("throw_on_cudamalloc_oom:maybe")
+
+    @unittest.skipIf(TEST_CUDAMALLOCASYNC, "throw_on_cudamalloc_oom not supported")
+    @serialTest()
+    def test_throw_on_cudamalloc_oom(self):
+        """Test that throw_on_cudamalloc_oom + per_process_memory_fraction works correctly."""
+        torch.cuda.empty_cache()
+        device = torch._C._cuda_getDevice()
+        torch._C._cuda_resetAccumulatedMemoryStats(device)
+
+        try:
+            # Test 1: With rejection disabled (default), allocations should succeed
+            torch._C._accelerator_setAllocatorSettings("throw_on_cudamalloc_oom:False")
+            x = torch.empty(10 * 1024 * 1024, dtype=torch.int8, device="cuda")
+            del x
+            torch.cuda.empty_cache()
+
+            # Test 2: With throw_on_cudamalloc_oom enabled and a tight memory
+            # fraction, allocations that exceed the fraction limit should be
+            # preemptively rejected with OutOfMemoryError.
+            # Both settings must go through _accelerator_setAllocatorSettings so
+            # they are read from CUDAAllocatorConfig.
+            fraction = 0.005
+            torch._C._accelerator_setAllocatorSettings(
+                f"throw_on_cudamalloc_oom:True,per_process_memory_fraction:{fraction}"
+            )
+
+            total_mem = torch.cuda.get_device_properties(0).total_memory
+            # Allocate the allowed threshold + 1 MiB to guarantee rejection
+            alloc_bytes = int(total_mem * fraction) + 1024 * 1024
+            with self.assertRaises(torch.cuda.OutOfMemoryError):
+                torch.empty(alloc_bytes, dtype=torch.int8, device="cuda")
+
+            # Check that rejection counter was incremented
+            stats = torch.cuda.memory_stats()
+            self.assertGreater(stats["num_oom_rejections"], 0)
+
+        finally:
+            torch.cuda.empty_cache()
+            torch._C._accelerator_setAllocatorSettings(
+                "throw_on_cudamalloc_oom:False,per_process_memory_fraction:1.0"
+            )
+
     def test_allocator_backend(self):
         def check_output(script: str) -> str:
             return (
@@ -4904,7 +5308,7 @@ def test_raises_oom(self, max_split_size_mb_setting):
             torch.empty(1024 * 1024 * 1024 * 1024, device="cuda")
 
     @unittest.skipIf(
-        not (IS_LINUX and os.uname().machine == "x86_64"), "cpp traces only on linux"
+        not ((IS_X86 or IS_ARM64) and IS_LINUX), "cpp traces are linux x86/aarch64 only"
     )
     @unittest.skipIf(
         TEST_CUDAMALLOCASYNC, "setContextRecorder not supported by CUDAMallocAsync"
@@ -4977,6 +5381,14 @@ def cb(device, alloc, device_alloc, device_free):
 
     def test_allocator_fuzz(self):
         # fuzz
+        if (
+            torch.version.hip
+            and "expandable_segments:True"
+            in torch._C._accelerator_getAllocatorSettings()
+        ):
+            raise unittest.SkipTest(
+                "ROCm needs https://github.com/ROCm/rocm-systems/pull/3023"
+            )
         state = random.getstate()
         random.seed(123)
         N = 10000
@@ -5022,7 +5434,11 @@ def test_temperature(self):
     @unittest.skipIf(not TEST_PYNVML, "pynvml/amdsmi is not available")
     def test_device_memory_used(self):
         """
-        Verify used device memory in bytes
+        Verify used device memory in bytes.
+        On Windows the NVML used value has been observed not to increase after
+        a CUDA allocation (delta 0); we only assert API sanity there (non-negative,
+        non-decreasing after alloc, <= total memory). Need to investigate expected behavior
+        with Windows WDDM
         """
         torch.cuda.synchronize()
         gc.collect()
@@ -5033,9 +5449,20 @@ def test_device_memory_used(self):
         torch.cuda.synchronize()
         torch.cuda.empty_cache()
         b = torch.cuda.device_memory_used()
-        mem_bytes = b - a
-        # test the order of magnitude
-        self.assertTrue(num_bytes // 32 <= mem_bytes <= num_bytes * 32)
+        if IS_WINDOWS:
+            # NVML used memory does not reflect CUDA allocations on WDDM; only check API sanity
+            self.assertGreaterEqual(a, 0, "device_memory_used should be non-negative")
+            self.assertGreaterEqual(b, 0, "device_memory_used should be non-negative")
+            self.assertGreaterEqual(
+                b, a, "used memory should not decrease after allocation"
+            )
+            total = torch.cuda.get_device_properties(0).total_memory
+            self.assertLessEqual(a, total, "used should not exceed total device memory")
+            self.assertLessEqual(b, total, "used should not exceed total device memory")
+        else:
+            mem_bytes = b - a
+            # test the order of magnitude
+            self.assertTrue(num_bytes // 32 <= mem_bytes <= num_bytes * 32)
 
     @unittest.skipIf(not TEST_PYNVML, "pynvml/amdsmi is not available")
     def test_power_draw(self):
@@ -5727,6 +6154,7 @@ def test_pin_memory_use(self, use_cuda_host_register):
         "use_memory, delete_memory",
         [(True, True), (True, False), (False, True), (False, False)],
     )
+    @skip_background_threads_on_windows
     def test_two_graphs(
         self, use_background_threads, use_cuda_host_register, use_memory, delete_memory
     ):
@@ -5830,6 +6258,43 @@ def test_mempool_id(self):
         # increments the id
         self.assertTrue(abs(pool2[1] - pool1[1]) > 0)
 
+    @unittest.skipIf(
+        TEST_CUDAMALLOCASYNC, "setContextRecorder not supported by CUDAMallocAsync"
+    )
+    def test_pool_id_in_snapshot(self):
+        try:
+            torch.cuda.memory.empty_cache()
+            torch.cuda.memory._record_memory_history("all")
+
+            pool = torch.cuda.MemPool()
+            with torch.cuda.use_mem_pool(pool):
+                x = torch.rand(64, device="cuda")
+
+            ss = torch.cuda.memory._snapshot()
+
+            # segment_pool_id should match the MemPool id
+            found_segment = False
+            for seg in ss["segments"]:
+                if seg["segment_pool_id"] == pool.id:
+                    found_segment = True
+                    break
+            self.assertTrue(found_segment)
+
+            # trace entries for this allocation should carry pool_id
+            found_trace = False
+            for trace in ss["device_traces"]:
+                for te in trace:
+                    if "pool_id" not in te:
+                        continue
+                    if te["pool_id"] == pool.id and te["action"] == "alloc":
+                        found_trace = True
+                        break
+            self.assertTrue(found_trace)
+
+            del x
+        finally:
+            torch.cuda.memory._record_memory_history(None)
+
     def get_dummy_allocator(self, check_vars):
         dummy_allocator_source_vars = """
         #include <torch/extension.h>
@@ -6055,67 +6520,68 @@ def test_mempool_limited_memory_with_allocator(self):
         nelem_1mb = 1024 * 1024 // 4
 
         self._setup_mempool_limited_memory_test(80)
-        # remaining free mem: 80 mb
-        # mempool_use [] 0 mb
-        # mempool_do_not_use [] 0 mb
-        # default pool [] 0 mb
-        with torch.cuda.use_mem_pool(pool_do_not_use):
-            a = torch.randn(40 * nelem_1mb, device="cuda")
-        with torch.cuda.use_mem_pool(pool_use):
-            b = torch.randn(40 * nelem_1mb, device="cuda")
-        a_dataptr = a.data_ptr()
-        b_dataptr = b.data_ptr()
-        # remaining free mem: 0 mb
-        # mempool_do_not_use [aaaa] 40 mb
-        # mempool_use [bbbb] 40 mb
-        # default pool [] 0 mb
-        with self.assertRaises(torch.OutOfMemoryError):
-            # out of memory
-            c = torch.randn(40 * nelem_1mb, device="cuda")
-
-        del a, b
-        # remaining free mem: 0 mb
-        # mempool_do_not_use [____] 40 mb
-        # mempool_use [____] 40 mb
-        # default pool [] 0 mb
-
-        # c should not oom and instead can use mempool_use as fallback
-        c = torch.randn(30 * nelem_1mb, device="cuda")
-        c_dataptr = c.data_ptr()
-        # remaining free mem: 0 mb
-        # mempool_do_not_use [____] 40 mb
-        # mempool_use [ccc_] 40 mb
-        # default pool [] 0 mb
-        with self.assertRaises(torch.OutOfMemoryError):
-            # out of memory since can't use mempool_do_not_use
-            d = torch.randn(30 * nelem_1mb, device="cuda")
-
-        del c
-        # remaining free mem: 0 mb
-        # mempool_do_not_use [____] 40 mb
-        # mempool_use [____] 40 mb
-        # default pool [] 0 mb
-
-        # expect that we used same memory address for both a and c
-        self.assertEqual(b_dataptr, c_dataptr)
-
-        # make sure we can still use mempool_use as intended after c is deleted
-        with torch.cuda.use_mem_pool(pool_use):
-            e = torch.randn(20 * nelem_1mb, device="cuda")
-        # remaining free mem: 0 mb
-        # mempool_do_not_use [____] 40 mb
-        # mempool_use [ee__] 40 mb
-        # default pool [] 0 mb
-
-        e_dataptr = e.data_ptr()
-        del e
-
-        self.assertEqual(e_dataptr, c_dataptr)
+        try:
+            # remaining free mem: 80 mb
+            # mempool_use [] 0 mb
+            # mempool_do_not_use [] 0 mb
+            # default pool [] 0 mb
+            with torch.cuda.use_mem_pool(pool_do_not_use):
+                a = torch.randn(40 * nelem_1mb, device="cuda")
+            with torch.cuda.use_mem_pool(pool_use):
+                b = torch.randn(40 * nelem_1mb, device="cuda")
+            a_dataptr = a.data_ptr()
+            b_dataptr = b.data_ptr()
+            # remaining free mem: 0 mb
+            # mempool_do_not_use [aaaa] 40 mb
+            # mempool_use [bbbb] 40 mb
+            # default pool [] 0 mb
+            with self.assertRaises(torch.OutOfMemoryError):
+                # out of memory
+                c = torch.randn(40 * nelem_1mb, device="cuda")
 
-        # pool's destructor calls emptyCache()
-        del pool_use, pool_do_not_use
+            del a, b
+            # remaining free mem: 0 mb
+            # mempool_do_not_use [____] 40 mb
+            # mempool_use [____] 40 mb
+            # default pool [] 0 mb
+
+            # c should not oom and instead can use mempool_use as fallback
+            c = torch.randn(30 * nelem_1mb, device="cuda")
+            c_dataptr = c.data_ptr()
+            # remaining free mem: 0 mb
+            # mempool_do_not_use [____] 40 mb
+            # mempool_use [ccc_] 40 mb
+            # default pool [] 0 mb
+            with self.assertRaises(torch.OutOfMemoryError):
+                # out of memory since can't use mempool_do_not_use
+                d = torch.randn(30 * nelem_1mb, device="cuda")
+
+            del c
+            # remaining free mem: 0 mb
+            # mempool_do_not_use [____] 40 mb
+            # mempool_use [____] 40 mb
+            # default pool [] 0 mb
+
+            # expect that we used same memory address for both a and c
+            self.assertEqual(b_dataptr, c_dataptr)
+
+            # make sure we can still use mempool_use as intended after c is deleted
+            with torch.cuda.use_mem_pool(pool_use):
+                e = torch.randn(20 * nelem_1mb, device="cuda")
+            # remaining free mem: 0 mb
+            # mempool_do_not_use [____] 40 mb
+            # mempool_use [ee__] 40 mb
+            # default pool [] 0 mb
+
+            e_dataptr = e.data_ptr()
+            del e
 
-        self._teardown_mempool_limited_memory_test()
+            self.assertEqual(e_dataptr, c_dataptr)
+
+            # pool's destructor calls emptyCache()
+            del pool_use, pool_do_not_use
+        finally:
+            self._teardown_mempool_limited_memory_test()
 
     @serialTest()
     def test_mempool_no_split(self):
@@ -6188,31 +6654,36 @@ def test_deleted_mempool_not_used_on_oom(self):
 
         # set 40 mb total available memory
         self._setup_mempool_limited_memory_test(40)
-
-        # Create many pools with use_on_oom=True, allocate memory, then delete the pools
-        for _ in range(10):
-            pool_use_on_oom = torch.cuda.MemPool(allocator.allocator(), use_on_oom=True)
-            with torch.cuda.use_mem_pool(pool_use_on_oom):
+        try:
+            # Create many pools with use_on_oom=True, allocate memory, then delete the pools
+            for _ in range(10):
+                pool_use_on_oom = torch.cuda.MemPool(
+                    allocator.allocator(), use_on_oom=True
+                )
+                with torch.cuda.use_mem_pool(pool_use_on_oom):
+                    a = torch.randn(40 * nelem_1mb, device="cuda")
+                del a
+                del pool_use_on_oom
+
+            # create new pool that we want to use_on_oom, all other pools should be deleted
+            # all available 40mb in use by mempool
+            new_pool_use_on_oom = torch.cuda.MemPool(
+                allocator.allocator(), use_on_oom=True
+            )
+            with torch.cuda.use_mem_pool(new_pool_use_on_oom):
                 a = torch.randn(40 * nelem_1mb, device="cuda")
             del a
-            del pool_use_on_oom
-
-        # create new pool that we want to use_on_oom, all other pools should be deleted
-        # all available 40mb in use by mempool
-        new_pool_use_on_oom = torch.cuda.MemPool(allocator.allocator(), use_on_oom=True)
-        with torch.cuda.use_mem_pool(new_pool_use_on_oom):
-            a = torch.randn(40 * nelem_1mb, device="cuda")
-        del a
 
-        # allocate tensors that will fallback to use_on_oom pool since all available 40mb in use by mempool
-        # tensors should only use valid pool and not deleted pools
-        b = torch.randn(20 * nelem_1mb, device="cuda")
-        c = torch.randn(20 * nelem_1mb, device="cuda")
+            # allocate tensors that will fallback to use_on_oom pool since all available 40mb in use by mempool
+            # tensors should only use valid pool and not deleted pools
+            b = torch.randn(20 * nelem_1mb, device="cuda")
+            c = torch.randn(20 * nelem_1mb, device="cuda")
 
-        del b
-        del c
-        del new_pool_use_on_oom
-        self._teardown_mempool_limited_memory_test()
+            del b
+            del c
+            del new_pool_use_on_oom
+        finally:
+            self._teardown_mempool_limited_memory_test()
 
     def test_mempool_multithread(self):
         pool_ids = []
@@ -6434,11 +6905,158 @@ def test_graph_capture_reclaim_4_streams(self):
             "graph_capture_record_stream_reuse:False"
         )
 
+    @unittest.skipIf(
+        not TEST_CUDA_GRAPH, "CUDA >= 11.0 or ROCM >= 5.3 required for graphs"
+    )
+    def test_graph_capture_reclaim_shared_pool(self):
+        torch.cuda.memory._set_allocator_settings(
+            "graph_capture_record_stream_reuse:True"
+        )
+        torch.cuda.empty_cache()
+        torch.cuda.synchronize()
+
+        shared_pool = torch.cuda.graph_pool_handle()
+        cap_stream = torch.cuda.Stream()
+        side_stream = torch.cuda.Stream()
+
+        g1 = torch.cuda.CUDAGraph()
+        g2 = torch.cuda.CUDAGraph()
+
+        numel = (8 * 1024 * 1024) // 4
+
+        with torch.cuda.stream(cap_stream):
+            g1.capture_begin(pool=shared_pool)
+            data = torch.empty(numel, device="cuda")
+            data_ptr = data.data_ptr()
+
+            side_stream.wait_stream(cap_stream)
+            with torch.cuda.stream(side_stream):
+                data.add_(1.0)
+                data.record_stream(side_stream)
+
+            cap_stream.wait_stream(side_stream)
+
+            del data
+            g1.capture_end()
+
+        torch.cuda.current_stream().wait_stream(cap_stream)
+        torch.cuda.synchronize()
+
+        with torch.cuda.stream(cap_stream):
+            g2.capture_begin(pool=shared_pool)
+            data2 = torch.empty(numel, device="cuda")
+            data2.fill_(42.0)
+            data2_ptr = data2.data_ptr()
+            g2.capture_end()
+
+        torch.cuda.current_stream().wait_stream(cap_stream)
+        torch.cuda.synchronize()
+
+        self.assertEqual(data_ptr, data2_ptr)
+
+        torch.cuda.memory._set_allocator_settings(
+            "graph_capture_record_stream_reuse:False"
+        )
+
+    @unittest.skipIf(
+        not TEST_CUDA_GRAPH, "CUDA >= 11.0 or ROCM >= 5.3 required for graphs"
+    )
+    @unittest.skipIf(TEST_WITH_ROCM, "ROCM does not support nvrtc")
+    @unittest.skipIf(
+        not SM70OrLater, "Compute capability >= SM70 required for relaxed ptx flag"
+    )
+    def test_graph_capture_pre_capture_stream_use(self):
+        # Tests that a block with pre-capture stream uses is correctly handled
+        # when freed during a subsequent capture on the same pool.
+        # Exercises the insert_events path in endAllocateToPool.
+        spin_wait_kernel = get_wait_for_cpu_kernel()
+
+        torch.cuda.memory._set_allocator_settings(
+            "graph_capture_record_stream_reuse:True"
+        )
+        torch.cuda.empty_cache()
+        torch.cuda.synchronize()
+
+        shared_pool = torch.cuda.graph_pool_handle()
+        cap_stream = torch.cuda.Stream()
+        side_stream = torch.cuda.Stream()
+        flag_cpu = torch.zeros(1, dtype=torch.int32, device="cpu").pin_memory()
+
+        g1 = torch.cuda.CUDAGraph()
+        g2 = torch.cuda.CUDAGraph()
+        g3 = torch.cuda.CUDAGraph()
+        g4 = torch.cuda.CUDAGraph()
+
+        numel = (8 * 1024 * 1024) // 4
+
+        # First capture: allocate data in the shared pool, keep it alive.
+        with torch.cuda.stream(cap_stream):
+            g1.capture_begin(pool=shared_pool)
+            data = torch.empty(numel, device="cuda")
+            data_ptr = data.data_ptr()
+            g1.capture_end()
+
+        torch.cuda.synchronize()
+
+        # Between captures: block side_stream with a spin-wait kernel
+        # (pre-capture stream use). The kernel holds the stream busy until
+        # we explicitly set the flag from the host.
+        with torch.cuda.stream(side_stream):
+            spin_wait_kernel(grid=(1, 1, 1), block=(1, 1, 1), args=[flag_cpu])
+            data.record_stream(side_stream)
+
+        # Second capture: free data during capture.
+        with torch.cuda.stream(cap_stream):
+            g2.capture_begin(pool=shared_pool)
+            del data
+            g2.capture_end()
+
+        # Trigger process_events. The spin kernel is still holding side_stream,
+        # so cudaEventQuery returns cudaErrorNotReady and the block stays pending.
+        torch.empty(1, device="cuda")
+
+        # Allocate from the same pool: block must NOT be reused yet.
+        with torch.cuda.stream(cap_stream):
+            g3.capture_begin(pool=shared_pool)
+            not_reused = torch.empty(numel, device="cuda")
+            not_reused_ptr = not_reused.data_ptr()
+            g3.capture_end()
+
+        self.assertNotEqual(data_ptr, not_reused_ptr)
+
+        # Release the spin kernel so side_stream can finish.
+        flag_cpu[0] = 1
+        torch.cuda.synchronize()
+
+        # Trigger process_events to reclaim the block.
+        torch.empty(1, device="cuda")
+
+        # Fourth capture: the block should now be reusable.
+        with torch.cuda.stream(cap_stream):
+            g4.capture_begin(pool=shared_pool)
+            reused = torch.empty(numel, device="cuda")
+            reused_ptr = reused.data_ptr()
+            g4.capture_end()
+
+        self.assertEqual(data_ptr, reused_ptr)
+
+        torch.cuda.memory._set_allocator_settings(
+            "graph_capture_record_stream_reuse:False"
+        )
+
+    @unittest.skipIf(
+        not EXPANDABLE_SEGMENTS,
+        "requires expandable_segments mode (run via test_cuda_expandable_segments.py)",
+    )
+    # expandable_segments not supported (PYTORCH_C10_DRIVER_API_SUPPORTED not defined for windows builds)
+    @unittest.skipIf(
+        IS_WINDOWS and SM89OrLater,
+        "expandable_segments not supported (PYTORCH_C10_DRIVER_API_SUPPORTED not defined for windows builds)",
+    )
     @skipIfRocm(msg="expandable_segments mode is not supported on ROCm")
     @unittest.skipIf(IS_FBCODE or IS_SANDCASTLE, "Load_inline doesn't work in fbcode")
     def test_mempool_expandable(self):
         torch.cuda.empty_cache()
-        torch.cuda.memory._set_allocator_settings("expandable_segments:True")
         allocator, _ = self.get_dummy_allocator(check_vars=False)
         pool = torch.cuda.MemPool(allocator.allocator())
 
@@ -6462,10 +7080,14 @@ def test_mempool_expandable(self):
             num_expandable_segments, 1, "Expected to have 1 expandable segment only"
         )
 
-        torch.cuda.memory._set_allocator_settings("expandable_segments:False")
-
     @serialTest()
     def test_mempool_ctx_multithread(self):
+        # Collect first: a tensor leaked by a prior test that is only reachable
+        # through a reference cycle keeps its segment active, so empty_cache
+        # can't reclaim it and the "empty pool" assertion below flakes.
+        gc.collect()
+        torch._C._cudnn_clear_dropout_state()
+        torch._C._cuda_clearCublasWorkspaces()
         torch.cuda.empty_cache()
         segments = torch.cuda.memory._snapshot()["segments"]
         self.assertEqual(len(segments), 0, "Expected empty pool in the beginning")
@@ -6635,6 +7257,203 @@ def get_total_allocated(segments):
             "Allocated memory should be same with or without traces",
         )
 
+    @unittest.skipIf(TEST_WITH_ROCM, "not enabled by default on rocm")
+    @serialTest()
+    def test_multi_threads_alloc_in_same_order(self):
+        def alloc_tensors(
+            stream: torch.cuda.Stream,
+            sizes: [int],
+            device: str,
+            alloc_results: [],
+            lock,
+            thread_id=0,
+        ):
+            while len(sizes) > 0:
+                with lock:
+                    if len(sizes) == 0:
+                        return
+                    sz = sizes.pop()
+                    alloc_results.append(
+                        (
+                            torch.empty(sz, dtype=torch.int8, device=device),
+                            thread_id,
+                            sz,
+                        )
+                    )
+
+        def alloc_with_threads(
+            thread_count: int, stream: torch.cuda.Stream, mem_sizes: [int], device: str
+        ):
+            tensors_list = []
+            threads = []
+            mem_sizes = mem_sizes[:]
+            lock = threading.Lock()
+            for i in range(thread_count):
+                t = threading.Thread(
+                    target=alloc_tensors,
+                    args=(stream, mem_sizes, "cuda", tensors_list, lock, i),
+                )
+                t.start()
+                threads.append(t)
+            for t in threads:
+                t.join()
+            tensors_ptrs = [
+                (t.data_ptr(), thread_id, sz) for (t, thread_id, sz) in tensors_list
+            ]
+            return tensors_ptrs
+
+        stream = torch.cuda.current_stream()
+        thread_count = 4
+        mem_sizes = [
+            s * 2 * 1024 * 1024 for s in range(1, 5) for _ in range(thread_count)
+        ]
+        tensor_ptrs_round_1 = alloc_with_threads(
+            thread_count, stream, mem_sizes, "cuda"
+        )
+        tensor_ptrs_round_2 = alloc_with_threads(
+            thread_count, stream, mem_sizes, "cuda"
+        )
+        ptrs_round_1 = torch.tensor([ptr for ptr, _, _ in tensor_ptrs_round_1])
+        ptrs_round_2 = torch.tensor([ptr for ptr, _, _ in tensor_ptrs_round_2])
+        self.assertTrue((ptrs_round_1 == ptrs_round_2).all().item())
+
+    @unittest.skipIf(TEST_WITH_ROCM, "not enabled by default on rocm")
+    @serialTest()
+    def test_nccl_mem_alloc_addresses_in_random_order(self):
+        """
+        Test NCCL mem allocator with non-increasing allocation addresses.
+
+        This test uses a custom allocator to simulate ncclMemAlloc scenario where
+        allocated memory addresses may not be monotonically increasing across
+        different ranks due to fragmentation, OS allocation policies, or prior
+        allocations.
+
+        The registration counter-based comparator ensures consistent block ordering
+        regardless of actual memory addresses, which is critical for NCCL
+        symmetric memory alignment.
+        """
+
+        from cuda.bindings import runtime
+
+        ALLOC_FN = ctypes.CFUNCTYPE(
+            ctypes.c_void_p, ctypes.c_size_t, ctypes.c_int, ctypes.c_void_p
+        )
+        FREE_FN = ctypes.CFUNCTYPE(
+            None, ctypes.c_void_p, ctypes.c_size_t, ctypes.c_int, ctypes.c_void_p
+        )
+
+        class AllocState:
+            def __init__(self, test_instance):
+                self.first_stream = None
+                self.second_stream = None
+                self.allocated_addrs = []
+                self.buffer_size = 1 * 1024 * 1024 * 1024  # 1GB
+                self.base_ptr = -1
+                err, base_ptr = runtime.cudaMalloc(self.buffer_size)
+                test_instance.assertEqual(
+                    err,
+                    runtime.cudaError_t.cudaSuccess,
+                    "init allocation for test_nccl_mem_alloc_addresses_in_random_order should be successful",
+                )
+                self.base_ptr = base_ptr
+                self.head = base_ptr
+                self.tail = base_ptr + self.buffer_size
+
+            def __del__(self):
+                if self.base_ptr > 0:
+                    runtime.cudaFree(self.base_ptr)
+
+        state = AllocState(self)
+
+        def my_alloc(size, device, stream, _runtime=runtime):
+            nonlocal state
+            if state.first_stream is None:
+                state.first_stream = stream
+            elif state.second_stream is None:
+                state.second_stream = stream
+
+            if state.first_stream == stream:
+                ptr = state.head
+                state.head += size
+            else:
+                state.tail -= size
+                ptr = state.tail
+            state.allocated_addrs.append(ptr)
+            return ptr
+
+        def my_free(ptr, size, device, stream, _runtime=runtime):
+            pass
+
+        # Must keep these alive for the lifetime of the allocator
+        c_alloc = ALLOC_FN(my_alloc)
+        c_free = FREE_FN(my_free)
+        alloc_ptr = ctypes.cast(c_alloc, ctypes.c_void_p).value
+        free_ptr = ctypes.cast(c_free, ctypes.c_void_p).value
+        allocator = torch._C._cuda_customAllocator(alloc_ptr, free_ptr)
+        pool = torch.cuda.MemPool(allocator)
+
+        first_stream = torch.cuda.Stream()
+        second_stream = torch.cuda.Stream()
+
+        tensor_sizes = [24 * 1024 * 1024, 32 * 1024 * 1024]
+        alloc_cases = []
+        case_no = 0
+        for idx in range(3):
+            for stream_idx, stream in enumerate([first_stream, second_stream]):
+                # for second stream in reverse order
+                reverse_order = stream_idx % 2 == 1
+                tensor_sizes.sort(reverse=reverse_order)
+                for size in tensor_sizes:
+                    alloc_cases.append(
+                        {"stream": stream, "size": size, "case_no": case_no}
+                    )
+                    case_no += 1
+
+        # Test allocation and deallocation patterns
+        def alloc_tensors():
+            all_tensors = []
+            for case in alloc_cases:
+                with torch.cuda.stream(case["stream"]):
+                    all_tensors.append(
+                        torch.empty(case["size"], dtype=torch.uint8, device="cuda")
+                    )
+            return all_tensors
+
+        with torch.cuda.use_mem_pool(pool):
+            first_round_tensors = alloc_tensors()
+        tensor_ptrs = [t.data_ptr() for t in first_round_tensors]
+        # for t in first_round_tensors:
+        #    del t
+        del first_round_tensors
+
+        with torch.cuda.use_mem_pool(pool):
+            second_round_tensors = alloc_tensors()
+        second_round_tensor_ptrs = [t.data_ptr() for t in second_round_tensors]
+        del second_round_tensors
+
+        mem_snapshot = pool.snapshot()
+        self.assertEqual(
+            len(mem_snapshot),
+            len(tensor_ptrs),
+            f"expected to have {len(tensor_ptrs)} segments, but actually got {len(mem_snapshot)}",
+        )
+
+        for idx, first_addr in enumerate(tensor_ptrs):
+            second_addr = second_round_tensor_ptrs[idx]
+            self.assertEqual(
+                second_addr,
+                first_addr,
+                "mem addr allocated for second round should be same with first round",
+            )
+            self.assertEqual(
+                second_addr,
+                state.allocated_addrs[idx],
+                f"{second_round_tensor_ptrs[idx]=} != {state.allocated_addrs[idx]=}",
+            )
+        del pool
+        del state
+        torch.cuda.empty_cache()
+
 
 @unittest.skipIf(not TEST_CUDA, "CUDA not available, skipping tests")
 @torch.testing._internal.common_utils.markDynamoStrictTest
@@ -6999,25 +7818,6 @@ def test_graph_grad_scaling(self, device, dtype, optim_info, foreach, fused):
             self.assertEqual(scaler._growth_tracker, growth_tracker)
 
 
-@unittest.skipIf(
-    not PLATFORM_SUPPORTS_GREEN_CONTEXT, "Green context not available, skipping tests"
-)
-class TestGreenContext(TestCase):
-    def test_greencontext_restores_stream(self):
-        # need to start on a side stream as we are comparing pointers and want to avoid
-        # two NULL streams...
-        s = torch.cuda.Stream()
-        with torch.cuda.stream(s):
-            start_stream = torch.cuda.current_stream()
-            ctx = torch.cuda.green_contexts.GreenContext.create(1, 0)
-            ctx.set_context()
-            context_stream = torch.cuda.current_stream()
-            ctx.pop_context()
-            end_stream = torch.cuda.current_stream()
-            self.assertEqual(start_stream.cuda_stream, end_stream.cuda_stream)
-            self.assertNotEqual(start_stream.cuda_stream, context_stream.cuda_stream)
-
-
 @unittest.skipIf(not TEST_CUDA, "CUDA not available, skipping tests")
 class TestGDS(TestCase):
     def _get_tmp_dir_fs_type(self):
@@ -7600,9 +8400,6 @@ def test_cuda_autocast_deprecated_warning(self):
                 _ = torch.ones(10)
 
 
-@unittest.skipIf(
-    os.environ.get("USE_LEGACY_DRIVER", None) == "1", "Doesn't work with older driver"
-)
 class TestCompileKernel(TestCase):
     @unittest.skipIf(not TEST_CUDA, "No CUDA")
     def test_compile_kernel(self):
@@ -7767,9 +8564,13 @@ def test_compile_kernel_large_shared_memory(self):
 
         # Test error handling with more than supported shared memory size
         if torch.version.hip:
-            max_smem = (
-                65536 if get_device_properties().gcnArchName != "gfx950" else 160 * 1024
-            )
+            gcn_arch = get_device_properties().gcnArchName.split(":", 1)[0]
+            if gcn_arch == "gfx1250":
+                max_smem = 320 * 1024
+            elif gcn_arch == "gfx950":
+                max_smem = 160 * 1024
+            else:
+                max_smem = 65536
         else:
             max_smem = get_device_properties().shared_memory_per_block_optin
         excessive_shared_mem = max_smem * 2
@@ -8125,26 +8926,14 @@ def test_compile_kernel_dlpack(self):
 
 @unittest.skipIf(not TEST_CUDA, "CUDA not available, skipping tests")
 class TestCudaDeviceParametrized(TestCase):
-    @unittest.skipIf(
-        TEST_WITH_ROCM, "ROCM does not support nvrtc or external cuda graph events"
-    )
+    @skipIfRocmVersionLessThan((7, 0))
     @skipCUDAIf(
         not SM70OrLater, "Compute capability >= SM70 required for relaxed ptx flag"
     )
     def test_graph_external_wait_and_record(self):
         torch.cuda.empty_cache()
 
-        kernel_source = r"""
-        __global__ void wait_for_cpu(int *pinned_cpu_flag) {
-            int flag = 0;
-            do {
-                    asm volatile("ld.relaxed.sys.global.s32 %0, [%1];" : "=r"(flag) : "l"(pinned_cpu_flag) : "memory");
-            } while (flag == 0);
-        }
-        """
-        from torch.cuda import _compile_kernel
-
-        spin_wait_kernel = _compile_kernel(kernel_source, "wait_for_cpu")
+        spin_wait_kernel = get_wait_for_cpu_kernel()
 
         x = torch.ones(4, device="cuda")
         x_cpu = torch.zeros(x.shape, device="cpu").pin_memory()
@@ -8217,6 +9006,13 @@ def test_graph_external_wait_and_record(self):
             torch.all(x_cpu == 1.0),
             "Copy should be done once end_event is synchronized",
         )
+        # On HIP graphs, the event-record node can become visible to
+        # event.synchronize() a few microseconds before the host-side
+        # stream state is updated; poll briefly so we exercise the
+        # invariant ("after sync the stream is drained") without flaking.
+        deadline = time.monotonic() + 0.1
+        while not work_stream.query() and time.monotonic() < deadline:
+            pass
         self.assertTrue(
             work_stream.query(),
             "end_event.synchronize() completing should imply that work_stream is done",
@@ -8349,12 +9145,195 @@ def test_fx_memory_profiler_augmentation(self):
                 self.assertIn("e = self.relu(d)", frame["fx_original_trace"])
 
 
+@unittest.skipIf(
+    not PLATFORM_SUPPORTS_GREEN_CONTEXT, "Green contexts are not supported"
+)
+class TestCudaGreenContexts(TestCase):
+    def setUp(self):
+        super().setUp()
+
+    def tearDown(self):
+        super().tearDown()
+
+    def test_greencontext_restores_stream(self):
+        # need to start on a side stream as we are comparing pointers and want to avoid
+        # two NULL streams...
+        s = torch.cuda.Stream()
+        with torch.cuda.stream(s):
+            start_stream = torch.cuda.current_stream()
+            ctx = torch.cuda.green_contexts.GreenContext.create(num_sms=1)
+            ctx.set_context()
+            context_stream = torch.cuda.current_stream()
+            ctx.pop_context()
+            end_stream = torch.cuda.current_stream()
+            self.assertEqual(start_stream.cuda_stream, end_stream.cuda_stream)
+            self.assertNotEqual(start_stream.cuda_stream, context_stream.cuda_stream)
+
+    @serialTest()
+    def test_greencontext_carveout(self):
+        # By default, everything is performed on the current device
+        a = torch.randn(4096, 4096, device="cuda", dtype=torch.bfloat16)
+        ctx = torch.cuda.green_contexts.GreenContext.create(num_sms=1)
+        ctx.set_context()
+        torch.matmul(a, a)
+        torch.cuda.synchronize()
+        t0 = time.perf_counter()
+        partial_res = torch.matmul(a, a)
+        torch.cuda.synchronize()
+        t1 = time.perf_counter()
+        ctx.pop_context()
+        torch.matmul(a, a)
+        torch.cuda.synchronize()
+        t2 = time.perf_counter()
+        full_res = torch.matmul(a, a)
+        torch.cuda.synchronize()
+        t3 = time.perf_counter()
+        self.assertEqual(partial_res, full_res)
+        self.assertGreater(t1 - t0, t3 - t2)
+
+    @serialTest()
+    def test_greencontext_stream_carveout(self):
+        # By default, everything is performed on the current device
+        a = torch.randn(4096, 4096, device="cuda", dtype=torch.bfloat16)
+        ctx = torch.cuda.green_contexts.GreenContext.create(num_sms=1)
+        ctx_stream = ctx.Stream()
+        with torch.cuda.stream(ctx_stream):
+            torch.matmul(a, a)
+            torch.cuda.synchronize()
+            t0 = time.perf_counter()
+            partial_res = torch.matmul(a, a)
+            torch.cuda.synchronize()
+            t1 = time.perf_counter()
+        torch.matmul(a, a)
+        torch.cuda.synchronize()
+        t2 = time.perf_counter()
+        full_res = torch.matmul(a, a)
+        torch.cuda.synchronize()
+        t3 = time.perf_counter()
+        self.assertEqual(partial_res, full_res)
+        self.assertGreater(t1 - t0, t3 - t2)
+
+    @serialTest()
+    def test_greencontext_graphs(self):
+        # By default, everything is performed on the current device
+        a = torch.randn(4096, 4096, device="cuda", dtype=torch.bfloat16)
+        ctx = torch.cuda.green_contexts.GreenContext.create(num_sms=1)
+        ctx.set_context()
+        partial_res = torch.matmul(a, a)
+        ctx.pop_context()
+        full_res = torch.matmul(a, a)
+        full_res.zero_()
+        partial_res.zero_()
+        torch.cuda.synchronize()
+
+        g = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(g):
+            ctx.set_context()
+            partial_res = torch.matmul(a, a)
+            ctx.pop_context()
+            full_res = torch.matmul(a, a)
+        g.replay()
+        self.assertEqual(partial_res, full_res)
+
+    @unittest.skipIf(
+        not PLATFORM_SUPPORTS_WORKQUEUE_CONFIG, "Workqueue config is not supported"
+    )
+    @serialTest()
+    def test_greencontext_workqueue_concurrency_limit(self):
+        # By default, everything is performed on the current device
+        n_streams = 4
+        GreenContext = torch.cuda.green_contexts.GreenContext
+        max_wq = GreenContext.max_workqueue_concurrency()
+        if max_wq < n_streams:
+            self.skipTest(
+                f"Device has {max_wq} workqueue(s), need >{n_streams} to test concurrency limiting"
+            )
+
+        wq_events = [torch.cuda.Event(enable_timing=True) for _ in range(n_streams * 2)]
+
+        def run_multi_stream_sleep(streams):
+            for i, s in enumerate(streams):
+                with torch.cuda.stream(s):
+                    # note we need to record timing events to ensure that the
+                    # workqueue is actually used
+                    wq_events[i * 2].record(s)
+                    # 100M cycles is enough on all currently supported GPUs
+                    # to ensure that the test runs significantly longer than
+                    # host overhead to 1) activate the different streams,
+                    # 2) record timing events, and 3) synchronize with the GPU.
+                    torch.cuda._sleep(100_000_000)
+                    wq_events[i * 2 + 1].record(s)
+            torch.cuda.synchronize()
+
+        # Note: in case we have lazy module loading, ensure that we called the
+        # sleep kernel at least once s.t. it is loaded in memory.
+        torch.cuda._sleep(1_000)
+
+        # Baseline: n_streams streams from default context, full workqueue concurrency
+        baseline_streams = [torch.cuda.Stream() for _ in range(n_streams)]
+        # note: torch.cuda.synchronize() will wait for all kernels on all streams
+        # so we can safely use CPU based timing here.
+        t0 = time.perf_counter()
+        run_multi_stream_sleep(baseline_streams)
+        t1 = time.perf_counter()
+        baseline_time = t1 - t0
+
+        # Green context with workqueue concurrency limited to 1
+        ctx = GreenContext.create(
+            workqueue_scope="balanced",
+            workqueue_concurrency_limit=1,
+        )
+        ctx.set_context()
+        green_streams = [ctx.Stream() for _ in range(n_streams)]
+        t2 = time.perf_counter()
+        run_multi_stream_sleep(green_streams)
+        t3 = time.perf_counter()
+        ctx.pop_context()
+        limited_time = t3 - t2
+
+        self.assertGreater(limited_time, baseline_time)
+
+
 instantiate_parametrized_tests(TestCuda)
-instantiate_parametrized_tests(TestCudaMallocAsync)
+instantiate_parametrized_tests(TestCudaAllocator)
 instantiate_parametrized_tests(TestCompileKernel)
 instantiate_parametrized_tests(TestCachingHostAllocatorCudaGraph)
 instantiate_device_type_tests(TestCudaOptims, globals())
 instantiate_device_type_tests(TestCudaDeviceParametrized, globals())
+instantiate_device_type_tests(TestCudaGreenContexts, globals(), except_for="cpu")
+
+
+class TestMemoryViz(TestCase):
+    def test_format_flamegraph_download_moves_temp_file(self):
+        # Regression test: format_flamegraph downloads flamegraph.pl into a temp
+        # file and moves it into place. Previously the temp file was created by a
+        # delete=True NamedTemporaryFile, so after os.rename moved it away the
+        # context manager raised FileNotFoundError on exit (a CI flake, since the
+        # download only happens when the cached script is absent).
+        from torch.cuda import _memory_viz
+
+        def fake_urlretrieve(url, filename):
+            with open(filename, "wb") as f:
+                f.write(b"#!/usr/bin/perl\n")
+            return filename, None
+
+        @contextlib.contextmanager
+        def fake_popen(*args, **kwargs):
+            proc = unittest.mock.MagicMock()
+            proc.stdout.read.return_value = "FLAMEGRAPH_OUTPUT"
+            proc.wait.return_value = 0
+            yield proc
+
+        with tempfile.TemporaryDirectory() as cache_dir:
+            script = os.path.join(cache_dir, "flamegraph.pl")
+            with (
+                patch("urllib.request.urlretrieve", fake_urlretrieve),
+                patch.object(_memory_viz.subprocess, "Popen", fake_popen),
+            ):
+                result = _memory_viz.format_flamegraph("stack 1\n", script)
+            self.assertEqual(result, "FLAMEGRAPH_OUTPUT")
+            self.assertTrue(os.path.exists(script))
+
 
 if __name__ == "__main__":
     run_tests()
diff --git a/test/test_cuda_expandable_segments.py b/test/test_cuda_expandable_segments.py
index 78e4cddab84ed..9b6e3da6081a4 100644
--- a/test/test_cuda_expandable_segments.py
+++ b/test/test_cuda_expandable_segments.py
@@ -1,20 +1,43 @@
 # Owner(s): ["module: cuda"]
 # run time cuda tests, but with the allocator using expandable segments
 
+import os
+
+
+# Must precede the test_cuda import: EXPANDABLE_SEGMENTS in
+# torch.testing._internal.common_utils reads PYTORCH_CUDA_ALLOC_CONF once at
+# import time. Setting it here lets the @skipIf(not EXPANDABLE_SEGMENTS, ...)
+# guards on expandable-only tests in test_cuda.py run in this runner. We
+# restore the prior env state below (after the import) so subprocesses
+# spawned by tests don't inherit expandable_segments:True, which would
+# conflict with backend:cudaMallocAsync and similar settings those
+# subprocesses set themselves.
+_PRIOR_ALLOC_CONF = os.environ.get("PYTORCH_CUDA_ALLOC_CONF")
+os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
+
 import pathlib
 import sys
 
 from test_cuda import (  # noqa: F401
     TestBlockStateAbsorption,
     TestCuda,
-    TestCudaMallocAsync,
+    TestCudaAllocator,
+    TestMemPool,
 )
 
 import torch
 from torch.testing._internal.common_cuda import IS_JETSON, IS_WINDOWS
-from torch.testing._internal.common_utils import run_tests, TEST_WITH_ROCM
+from torch.testing._internal.common_utils import run_tests
 
 
+# Restore prior env state. EXPANDABLE_SEGMENTS has already been resolved at
+# import time above; the runtime allocator is forced into expandable mode
+# via _set_allocator_settings in __main__ below.
+if _PRIOR_ALLOC_CONF is None:
+    os.environ.pop("PYTORCH_CUDA_ALLOC_CONF", None)
+else:
+    os.environ["PYTORCH_CUDA_ALLOC_CONF"] = _PRIOR_ALLOC_CONF
+
 REPO_ROOT = pathlib.Path(__file__).resolve().parent.parent
 sys.path.insert(0, str(REPO_ROOT))
 
@@ -25,16 +48,7 @@
 sys.path.remove(str(REPO_ROOT))
 
 if __name__ == "__main__":
-    if (
-        torch.cuda.is_available()
-        and not IS_JETSON
-        and not IS_WINDOWS
-        and not TEST_WITH_ROCM
-    ):
+    if torch.cuda.is_available() and not IS_JETSON and not IS_WINDOWS:
         get_disabled_tests(".")
-
         torch.cuda.memory._set_allocator_settings("expandable_segments:True")
-        TestCuda.expandable_segments = lambda _: True
-        TestBlockStateAbsorption.expandable_segments = lambda _: True
-
         run_tests()
diff --git a/test/test_cuda_graph_annotations.py b/test/test_cuda_graph_annotations.py
new file mode 100644
index 0000000000000..f29bf78a2ab34
--- /dev/null
+++ b/test/test_cuda_graph_annotations.py
@@ -0,0 +1,398 @@
+# Owner(s): ["module: cuda graphs"]
+
+"""Tests for CUDA graph kernel annotation via mark_kernels."""
+
+import unittest
+
+import torch
+from torch.cuda._graph_annotations import (
+    _is_tools_id_unavailable,
+    clear_kernel_annotations,
+    enable_annotations,
+    get_kernel_annotations,
+    mark_kernels,
+    remap_to_exec_graph,
+    resolve_pending_annotations,
+)
+from torch.testing._internal.common_utils import run_tests, TestCase
+
+
+TEST_CUDA = torch.cuda.is_available()
+
+try:
+    import cuda.bindings.runtime  # noqa: F401
+
+    TEST_CUDA_BINDINGS = True
+except ImportError:
+    TEST_CUDA_BINDINGS = False
+
+
+@unittest.skipUnless(TEST_CUDA, "CUDA not available")
+@unittest.skipUnless(TEST_CUDA_BINDINGS, "cuda.bindings not available")
+@unittest.skipIf(
+    _is_tools_id_unavailable(),
+    "cudaGraphNodeGetToolsId not available (needs cuda-compat >= 13.1)",
+)
+class TestMarkKernels(TestCase):
+    def setUp(self):
+        enable_annotations()
+        clear_kernel_annotations()
+
+    def tearDown(self):
+        clear_kernel_annotations()
+
+    def test_noop_outside_capture(self):
+        x = torch.randn(8, device="cuda")
+        with mark_kernels("test"):
+            _ = x + 1
+        self.assertEqual(len(get_kernel_annotations()), 0)
+
+    def test_single_scope(self):
+        graph = torch.cuda.CUDAGraph()
+        x = torch.randn(8, device="cuda")
+
+        with torch.cuda.graph(graph):
+            with mark_kernels("phase_a"):
+                _ = x + 1
+            resolve_pending_annotations()
+
+        annotations = get_kernel_annotations()
+        self.assertGreater(len(annotations), 0)
+        for anns in annotations.values():
+            for ann in anns:
+                self.assertEqual(ann, {"str": "phase_a"})
+
+    def test_multiple_scopes_no_overlap(self):
+        graph = torch.cuda.CUDAGraph()
+        x = torch.randn(8, device="cuda")
+
+        with torch.cuda.graph(graph):
+            with mark_kernels("scope_1"):
+                _ = x + 1
+            with mark_kernels("scope_2"):
+                _ = x * 2
+            resolve_pending_annotations()
+
+        annotations = get_kernel_annotations()
+        scope_1_ids = set()
+        scope_2_ids = set()
+        for tid, anns in annotations.items():
+            self.assertEqual(len(anns), 1)
+            if anns[0] == {"str": "scope_1"}:
+                scope_1_ids.add(tid)
+            elif anns[0] == {"str": "scope_2"}:
+                scope_2_ids.add(tid)
+
+        self.assertGreater(len(scope_1_ids), 0)
+        self.assertGreater(len(scope_2_ids), 0)
+        self.assertEqual(len(scope_1_ids & scope_2_ids), 0)
+
+    def test_dict_annotation(self):
+        graph = torch.cuda.CUDAGraph()
+        x = torch.randn(8, device="cuda")
+
+        annotation = {"name": "all_gather", "Group size": 2, "dtype": "bfloat16"}
+        with torch.cuda.graph(graph):
+            with mark_kernels(annotation):
+                _ = x + 1
+            resolve_pending_annotations()
+
+        annotations = get_kernel_annotations()
+        self.assertGreater(len(annotations), 0)
+        for anns in annotations.values():
+            self.assertEqual(anns[0]["name"], "all_gather")
+            self.assertEqual(anns[0]["Group size"], 2)
+
+    def test_clear_resets_state(self):
+        graph = torch.cuda.CUDAGraph()
+        x = torch.randn(8, device="cuda")
+
+        with torch.cuda.graph(graph):
+            with mark_kernels("test"):
+                _ = x + 1
+            resolve_pending_annotations()
+
+        self.assertGreater(len(get_kernel_annotations()), 0)
+        clear_kernel_annotations()
+        self.assertEqual(len(get_kernel_annotations()), 0)
+
+    def test_resolve_without_scopes_is_noop(self):
+        resolve_pending_annotations()
+        self.assertEqual(len(get_kernel_annotations()), 0)
+
+    def test_scope_with_no_kernels(self):
+        graph = torch.cuda.CUDAGraph()
+        x = torch.randn(8, device="cuda")
+
+        with torch.cuda.graph(graph):
+            _ = x + 1
+            with mark_kernels("empty"):
+                pass
+            _ = x * 2
+            resolve_pending_annotations()
+
+        for anns in get_kernel_annotations().values():
+            for ann in anns:
+                self.assertNotEqual(ann, "empty")
+
+    def test_only_annotates_scope_kernels(self):
+        graph = torch.cuda.CUDAGraph()
+        x = torch.randn(8, device="cuda")
+
+        with torch.cuda.graph(graph):
+            _ = x + 1
+            _ = x * 2
+            with mark_kernels("tagged"):
+                _ = x + 3
+            _ = x - 1
+            resolve_pending_annotations()
+
+        annotations = get_kernel_annotations()
+        total_annotated = sum(len(anns) for anns in annotations.values())
+        self.assertGreater(total_annotated, 0)
+        for anns in annotations.values():
+            for ann in anns:
+                self.assertEqual(ann, {"str": "tagged"})
+
+    def test_nested_scopes_innermost_wins(self):
+        """With nested string scopes, the innermost name wins."""
+        graph = torch.cuda.CUDAGraph()
+        x = torch.randn(8, device="cuda")
+
+        with torch.cuda.graph(graph):
+            with mark_kernels("outer"):
+                _ = x + 1  # outer only
+                with mark_kernels("inner"):
+                    _ = x * 2  # nested: inner should win
+                _ = x - 1  # outer only
+            resolve_pending_annotations()
+
+        annotations = get_kernel_annotations()
+        outer_ids = set()
+        inner_ids = set()
+        for tid, anns in annotations.items():
+            self.assertEqual(
+                len(anns), 1, f"toolsId {hex(tid)} has {len(anns)} annotations"
+            )
+            ann = anns[0]
+            self.assertIsInstance(ann, dict)
+            if ann["str"] == "outer":
+                outer_ids.add(tid)
+            elif ann["str"] == "inner":
+                inner_ids.add(tid)
+
+        self.assertGreater(len(outer_ids), 0, "Should have outer-only kernels")
+        self.assertGreater(len(inner_ids), 0, "Should have inner kernels")
+        self.assertEqual(len(outer_ids & inner_ids), 0)
+
+    def test_nested_dict_scopes_inner_wins_common_keys(self):
+        """With truly nested dict scopes, inner wins for common keys,
+        outer-only keys are preserved."""
+        graph = torch.cuda.CUDAGraph()
+        x = torch.randn(8, device="cuda")
+
+        outer_ann = {"name": "ag_collective", "stream": 71}
+        inner_ann = {
+            "name": "all_gather",
+            "stream": 62,
+            "In msg nelems": 1024,
+            "dtype": "bfloat16",
+        }
+
+        with torch.cuda.graph(graph):
+            with mark_kernels(outer_ann):
+                _ = x + 1  # outer only
+                with mark_kernels(inner_ann):
+                    _ = x * 2  # nested
+                _ = x - 1  # outer only
+            resolve_pending_annotations()
+
+        annotations = get_kernel_annotations()
+        outer_only_ids = set()
+        nested_ids = set()
+        for tid, anns in annotations.items():
+            self.assertEqual(len(anns), 1)
+            ann = anns[0]
+            self.assertIsInstance(ann, dict)
+            if ann["name"] == "ag_collective":
+                outer_only_ids.add(tid)
+            elif ann["name"] == "all_gather":
+                nested_ids.add(tid)
+                # Inner wins for common keys
+                self.assertEqual(ann["stream"], 62)
+                # Inner-only keys preserved
+                self.assertEqual(ann["In msg nelems"], 1024)
+                self.assertEqual(ann["dtype"], "bfloat16")
+
+        self.assertGreater(len(outer_only_ids), 0, "Should have outer-only kernels")
+        self.assertGreater(len(nested_ids), 0, "Should have nested kernels")
+
+    def test_same_range_scopes_inner_wins_common_keys(self):
+        """With same-range scopes (inner ctx exits first), inner wins
+        for common keys, outer-only keys are preserved."""
+        graph = torch.cuda.CUDAGraph()
+        x = torch.randn(8, device="cuda")
+
+        outer_ann = {"name": "ag_collective", "stream": 71}
+        inner_ann = {
+            "name": "all_gather",
+            "stream": 62,
+            "In msg nelems": 1024,
+            "dtype": "bfloat16",
+        }
+
+        with torch.cuda.graph(graph):
+            # Both scopes wrap the same kernels; inner exits first.
+            with mark_kernels(outer_ann):
+                with mark_kernels(inner_ann):
+                    _ = x + 1
+            resolve_pending_annotations()
+
+        annotations = get_kernel_annotations()
+        self.assertGreater(len(annotations), 0)
+        for anns in annotations.values():
+            self.assertEqual(len(anns), 1)
+            ann = anns[0]
+            self.assertIsInstance(ann, dict)
+            # Inner wins for common keys
+            self.assertEqual(ann["name"], "all_gather", "Inner name should win")
+            self.assertEqual(ann["stream"], 62, "Inner stream should win")
+            # Inner-only keys preserved
+            self.assertEqual(ann["In msg nelems"], 1024)
+            self.assertEqual(ann["dtype"], "bfloat16")
+
+    def test_remap_to_exec_graph(self):
+        from cuda.bindings import runtime as cuda_runtime
+
+        graph = torch.cuda.CUDAGraph()
+        x = torch.randn(8, device="cuda")
+
+        with torch.cuda.graph(graph):
+            with mark_kernels("test"):
+                _ = x + 1
+            resolve_pending_annotations()
+
+        annotations_before = dict(get_kernel_annotations())
+        self.assertGreater(len(annotations_before), 0)
+
+        exec_handle = cuda_runtime.cudaGraphExec_t(
+            init_value=graph.raw_cuda_graph_exec()
+        )
+        _, exec_graph_id = cuda_runtime.cudaGraphExecGetId(exec_handle)
+
+        remap_to_exec_graph(graph)
+
+        annotations_after = get_kernel_annotations()
+        self.assertEqual(len(annotations_after), len(annotations_before))
+        for tools_id in annotations_after:
+            self.assertEqual(tools_id >> 32, exec_graph_id)
+
+    def test_disabled_is_noop(self):
+        from torch.cuda._graph_annotations import disable_annotations
+
+        disable_annotations()
+
+        graph = torch.cuda.CUDAGraph()
+        x = torch.randn(8, device="cuda")
+
+        with torch.cuda.graph(graph):
+            with mark_kernels("should_not_appear"):
+                _ = x + 1
+            resolve_pending_annotations()
+
+        self.assertEqual(len(get_kernel_annotations()), 0)
+
+        # Re-enable for other tests
+        enable_annotations()
+
+    def test_enable_annotations_kwarg(self):
+        """enable_annotations on torch.cuda.graph auto-resolves annotations."""
+        from torch.cuda._graph_annotations import disable_annotations
+
+        # Start with annotations disabled to verify the kwarg enables them.
+        disable_annotations()
+        clear_kernel_annotations()
+
+        graph = torch.cuda.CUDAGraph()
+        x = torch.randn(8, device="cuda")
+
+        with torch.cuda.graph(graph, enable_annotations=True):
+            with mark_kernels("auto"):
+                _ = x + 1
+
+        annotations = get_kernel_annotations()
+        self.assertGreater(len(annotations), 0)
+        for anns in annotations.values():
+            for ann in anns:
+                self.assertEqual(ann, {"str": "auto"})
+
+    def test_enable_annotations_does_not_clear(self):
+        """Annotations from a previous graph survive a second capture."""
+        from torch.cuda._graph_annotations import disable_annotations
+
+        disable_annotations()
+        clear_kernel_annotations()
+
+        graph1 = torch.cuda.CUDAGraph()
+        x = torch.randn(8, device="cuda")
+
+        with torch.cuda.graph(graph1, enable_annotations=True):
+            with mark_kernels("first"):
+                _ = x + 1
+
+        first_count = len(get_kernel_annotations())
+        self.assertGreater(first_count, 0)
+
+        graph2 = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(graph2, enable_annotations=True):
+            with mark_kernels("second"):
+                _ = x * 2
+
+        # Both graphs' annotations should be present.
+        self.assertGreater(len(get_kernel_annotations()), first_count)
+
+    def test_enable_annotations_remaps_to_exec_graph(self):
+        """enable_annotations=True must remap toolsIds to the exec graph ID."""
+        from cuda.bindings import runtime as cuda_runtime
+
+        clear_kernel_annotations()
+
+        graph = torch.cuda.CUDAGraph()
+        x = torch.randn(8, device="cuda")
+
+        with torch.cuda.graph(graph, enable_annotations=True):
+            with mark_kernels("remap_test"):
+                _ = x + 1
+
+        exec_handle = cuda_runtime.cudaGraphExec_t(
+            init_value=graph.raw_cuda_graph_exec()
+        )
+        _, exec_graph_id = cuda_runtime.cudaGraphExecGetId(exec_handle)
+
+        annotations = get_kernel_annotations()
+        self.assertGreater(len(annotations), 0)
+        for tools_id in annotations:
+            graph_id = tools_id >> 32
+            self.assertEqual(
+                graph_id,
+                exec_graph_id,
+                f"toolsId 0x{tools_id:016x} has graph_id {graph_id}, "
+                f"expected exec_graph_id {exec_graph_id}",
+            )
+
+    def test_enable_annotations_false_does_not_auto_resolve(self):
+        """Without enable_annotations, pending scopes are not resolved."""
+        graph = torch.cuda.CUDAGraph()
+        x = torch.randn(8, device="cuda")
+
+        # enable_annotations=False (default): no auto-resolve.
+        with torch.cuda.graph(graph):
+            with mark_kernels("unresolved"):
+                _ = x + 1
+
+        # Annotations should be empty because resolve was never called.
+        self.assertEqual(len(get_kernel_annotations()), 0)
+
+
+if __name__ == "__main__":
+    run_tests()
diff --git a/test/test_cuda_multigpu.py b/test/test_cuda_multigpu.py
index aa447f06698bc..579ca1675f92b 100644
--- a/test/test_cuda_multigpu.py
+++ b/test/test_cuda_multigpu.py
@@ -11,7 +11,7 @@
 import threading
 import unittest
 from itertools import chain, repeat
-from typing import NamedTuple, Union
+from typing import NamedTuple
 
 import torch
 import torch.cuda.comm as comm
@@ -1006,7 +1006,7 @@ def test_rng_state_offset(self):
 
     # Verifies that mem_get_info works, including when called for a different device
     def test_mem_get_info(self):
-        def _test(device: Union[str, int, torch.device]):
+        def _test(device: str | int | torch.device):
             # Prevent PyTorch from reusing the allocated memory
             torch.cuda.empty_cache()
             torch.cuda.synchronize()
diff --git a/test/test_cuda_sanitizer.py b/test/test_cuda_sanitizer.py
index e5dae52354a72..35720176901d4 100644
--- a/test/test_cuda_sanitizer.py
+++ b/test/test_cuda_sanitizer.py
@@ -3,7 +3,6 @@
 import sys
 import textwrap
 import traceback
-from typing import Optional
 
 import torch
 import torch.cuda._sanitizer as csan
@@ -149,8 +148,8 @@ def setUp(self):
     def kernel_launch(
         self,
         stream: StreamId,
-        read_only: Optional[list[DataPtr]] = None,
-        read_write: Optional[list[DataPtr]] = None,
+        read_only: list[DataPtr] | None = None,
+        read_write: list[DataPtr] | None = None,
     ) -> list[csan.SynchronizationError]:
         if read_only is None:
             read_only = []
@@ -168,8 +167,8 @@ def kernel_launch(
     def assert_good_kernel_launch(
         self,
         stream: StreamId,
-        read_only: Optional[list[DataPtr]] = None,
-        read_write: Optional[list[DataPtr]] = None,
+        read_only: list[DataPtr] | None = None,
+        read_write: list[DataPtr] | None = None,
     ) -> None:
         self.assertEqual(self.kernel_launch(stream, read_only, read_write), [])
 
@@ -177,8 +176,8 @@ def assert_bad_kernel_launch(
         self,
         number_of_errors: int,
         stream: StreamId,
-        read_only: Optional[list[DataPtr]] = None,
-        read_write: Optional[list[DataPtr]] = None,
+        read_only: list[DataPtr] | None = None,
+        read_write: list[DataPtr] | None = None,
     ) -> None:
         errors = self.kernel_launch(stream, read_only, read_write)
         self.assertEqual(len(errors), number_of_errors)
diff --git a/test/test_custom_ops.py b/test/test_custom_ops.py
index c02b17c54cc94..72c7da5b0156a 100644
--- a/test/test_custom_ops.py
+++ b/test/test_custom_ops.py
@@ -53,8 +53,10 @@
     run_tests,
     scoped_load_inline,
     skipIfTorchDynamo,
+    skipIfXpu,
     subtest,
     TemporaryFileName,
+    TEST_XPU,
     TestCase,
 )
 from torch.testing._internal.custom_op_db import numpy_nonzero
@@ -68,6 +70,12 @@
 MyList = list
 MyTensor = torch.Tensor
 
+device_type = (
+    acc.type
+    if (acc := torch.accelerator.current_accelerator(check_available=True))
+    else "cpu"
+)
+
 
 def requires_compile(fun):
     fun = unittest.skipIf(IS_WINDOWS, "torch.compile doesn't work with windows")(fun)
@@ -1585,7 +1593,8 @@ def foo_backward(ctx, saved, grad0, grad1):
         with self.assertRaisesRegex(RuntimeError, "is not a Tensor"):
             op(x)
 
-    @unittest.skipIf(not TEST_CUDA, "requires CUDA")
+    @skipIfXpu(msg="Deprecated torch.custom_ops API")
+    @unittest.skipIf(not TEST_CUDA and not TEST_XPU, "requires CUDA or XPU")
     def test_impl_separate(self):
         @custom_ops.custom_op(f"{TestCustomOp.test_ns}::foo")
         def foo(x: torch.Tensor) -> torch.Tensor:
@@ -1595,7 +1604,7 @@ def foo(x: torch.Tensor) -> torch.Tensor:
         def foo_cpu(x):
             return x.sin()
 
-        @custom_ops.impl(f"{TestCustomOp.test_ns}::foo", device_types="cuda")
+        @custom_ops.impl(f"{TestCustomOp.test_ns}::foo", device_types=device_type)
         def foo_cuda(x):
             return x.cos()
 
@@ -1604,12 +1613,13 @@ def foo_cuda(x):
         result = op(x)
         self.assertEqual(result, foo_cpu(x))
 
-        x_cuda = x.cuda()
+        x_cuda = x.to(device_type)
         op = self.get_op(f"{self.test_ns}::foo")
         result = op(x_cuda)
         self.assertEqual(result, foo_cuda(x_cuda))
 
-    @unittest.skipIf(not TEST_CUDA, "requires CUDA")
+    @skipIfXpu(msg="Deprecated torch.custom_ops API")
+    @unittest.skipIf(not TEST_CUDA and not TEST_XPU, "requires CUDA or XPU")
     def test_impl_multiple(self):
         @custom_ops.custom_op(f"{TestCustomOp.test_ns}::foo")
         def foo(x: torch.Tensor) -> torch.Tensor:
@@ -1624,7 +1634,7 @@ def foo_impl(x):
         result = op(x)
         self.assertEqual(result, foo_impl(x))
 
-        x_cuda = x.cuda()
+        x_cuda = x.to(device_type)
         result = op(x_cuda)
         self.assertEqual(result, foo_impl(x_cuda))
 
@@ -2168,11 +2178,11 @@ def test_impl_device_cpu(self):
         self._test_impl_device("foo2", ["cpu"], "cpu")
         self._test_impl_device("foo3", ["cpu", "cuda"], "cpu")
 
-    @unittest.skipIf(not TEST_CUDA, "requires cuda")
+    @unittest.skipIf(not TEST_CUDA and not TEST_XPU, "requires cuda or xpu")
     def test_impl_device_cuda(self):
-        self._test_impl_device("foo4", "default", "cuda")
-        self._test_impl_device("foo5", ["cuda"], "cuda")
-        self._test_impl_device("foo6", ["cpu", "cuda"], "cuda")
+        self._test_impl_device("foo4", "default", device_type)
+        self._test_impl_device("foo5", [device_type], device_type)
+        self._test_impl_device("foo6", ["cpu", device_type], device_type)
 
     def test_impl_device_function(self):
         lib = self.lib()
@@ -2371,6 +2381,56 @@ def test_load_library(self):
         ):
             torch.ops.load_library("libnoexist.so")
 
+    def test_list_scalar_type(self):
+        lib = self.lib()
+        lib.define("scalar_list(Tensor x, ScalarType[] dts) -> Tensor")
+
+        received = None
+
+        @torch.library.impl(lib, "scalar_list", "CPU")
+        def _(x, dts):
+            nonlocal received
+            received = dts
+            return x.clone()
+
+        x = torch.randn(3)
+        torch.ops._test_custom_op.scalar_list(x, [torch.float32, torch.bfloat16])
+        self.assertEqual(received, [torch.float32, torch.bfloat16])
+
+    def test_list_layout(self):
+        lib = self.lib()
+        lib.define("layout_list(Tensor x, Layout[] layouts) -> Tensor")
+
+        received = None
+
+        @torch.library.impl(lib, "layout_list", "CPU")
+        def _(x, layouts):
+            nonlocal received
+            received = layouts
+            return x.clone()
+
+        x = torch.randn(3)
+        torch.ops._test_custom_op.layout_list(x, [torch.strided, torch.sparse_coo])
+        self.assertEqual(received, [torch.strided, torch.sparse_coo])
+
+    def test_list_memory_format(self):
+        lib = self.lib()
+        lib.define("memfmt_list(Tensor x, MemoryFormat[] fmts) -> Tensor")
+
+        received = None
+
+        @torch.library.impl(lib, "memfmt_list", "CPU")
+        def _(x, fmts):
+            nonlocal received
+            received = fmts
+            return x.clone()
+
+        x = torch.randn(3)
+        torch.ops._test_custom_op.memfmt_list(
+            x, [torch.contiguous_format, torch.channels_last]
+        )
+        self.assertEqual(received, [torch.contiguous_format, torch.channels_last])
+
 
 def op_with_incorrect_schema(testcase, name):
     lib = testcase.lib()
@@ -4500,6 +4560,32 @@ def test_library_get_kernel_invalid(self):
                 torch.library.get_kernel("test_invalid_kernel::cpu_only_op", "CUDA")
 
 
+class TestLibrarySourceLocation(TestCase):
+    def test_library_source_location(self):
+        # Library.__init__ uses sys._getframe(1) to capture the caller's
+        # filename and line number. Verify this works correctly by creating
+        # a Library and checking the source location in the error message
+        # that appears when a duplicate DEF library is created.
+        script = """\
+import torch
+lib1 = torch.library.Library("_test_loc", "DEF")
+lib1.define("foo(Tensor x) -> Tensor")
+try:
+    lib2 = torch.library.Library("_test_loc", "DEF")
+except RuntimeError as e:
+    print(str(e))
+"""
+        result = subprocess.run(
+            [sys.executable, "-c", script],
+            capture_output=True,
+            text=True,
+        )
+        self.assertEqual(result.returncode, 0, result.stderr)
+        # The error message should reference <string>:2, since
+        # lib1 = torch.library.Library(...) is on line 2 of the script.
+        self.assertIn("<string>:2", result.stdout)
+
+
 class MiniOpTestOther(CustomOpTestCaseBase):
     test_ns = "mini_op_test"
 
@@ -4795,10 +4881,10 @@ def test_optional(self):
 
     def test_mixed_types(self):
         result_type = tuple_to_list(Tuple[int, float])
-        self.assertEqual(result_type, list[typing.Union[int, float]])
+        self.assertEqual(result_type, list[int | float])
 
         result_type = tuple_to_list(Tuple[int, float, str])
-        self.assertEqual(result_type, list[typing.Union[int, float, str]])
+        self.assertEqual(result_type, list[int | float | str])
 
 
 class TestOpProfiles(TestCase):
diff --git a/test/test_dataloader.py b/test/test_dataloader.py
index 6ab1812fb02ff..df5c53b08d73d 100644
--- a/test/test_dataloader.py
+++ b/test/test_dataloader.py
@@ -25,6 +25,7 @@
 from torch.testing._internal.common_utils import (
     IS_CI,
     IS_JETSON,
+    IS_LINUX,
     IS_MACOS,
     IS_S390X,
     IS_SANDCASTLE,
@@ -40,7 +41,6 @@
     TEST_WITH_ASAN,
     TEST_WITH_TSAN,
     TestCase,
-    xfailIfLinux,
 )
 from torch.utils.data import (
     _utils,
@@ -1084,22 +1084,7 @@ def _test_worker_info_init_fn(worker_id):
         raise AssertionError("worker_info should have correct dataset copy")
     if hasattr(dataset, "value"):
         raise AssertionError("worker_info should have correct dataset copy")
-    # test that WorkerInfo attributes are read-only
-    try:
-        worker_info.id = 3999
-    except RuntimeError as e:
-        if str(e) != "Cannot assign attributes to WorkerInfo objects":
-            raise AssertionError(
-                "Expected RuntimeError for WorkerInfo attribute assignment"
-            ) from None
-    try:
-        worker_info.a = 3
-    except RuntimeError as e:
-        if str(e) != "Cannot assign attributes to WorkerInfo objects":
-            raise AssertionError(
-                "Expected RuntimeError for WorkerInfo attribute assignment"
-            ) from None
-    for k in ["id", "num_workers", "seed", "dataset"]:
+    for k in ["id", "num_workers", "seed", "dataset", "rng"]:
         if f"{k}=" not in repr(worker_info):
             raise AssertionError(f"Expected {k} in worker_info repr")
     dataset.value = [worker_id, os.getpid()]
@@ -1440,14 +1425,16 @@ def test_multiple_dataloaders(self):
             del loader1_it
             del loader2_it
 
-    # This case pass on Intel GPU, but currently expected failure on other device,
-    # please don't forget to remove this skip when remove the xfailIfLinux.
     @skipIfXpu
-    # This case passes on s390x too.
-    # please don't forget to remove this skip when remove the xfailIfLinux.
     @unittest.skipIf(IS_S390X, "Unexpectedly succeeds on s390x")
-    # https://github.com/pytorch/pytorch/issues/128551
-    @xfailIfLinux
+    # Test that DataLoader properly handles worker segfaults
+    # Note: This test has inconsistent behavior across Linux distributions:
+    # - Passes on RHEL 9.6 (segfault triggers correctly)
+    # - Fails on Ubuntu (process may not terminate as expected)
+    # Skipping on Linux due to kernel/distribution-dependent segfault behavior.
+    @unittest.skipIf(
+        IS_LINUX, "Segfault behavior is inconsistent across Linux distributions"
+    )
     def test_segfault(self):
         p = ErrorTrackingProcess(target=_test_segfault)
         p.start()
diff --git a/test/test_datapipe.py b/test/test_datapipe.py
index 12cff392ad8e8..393ebcbef3f3a 100644
--- a/test/test_datapipe.py
+++ b/test/test_datapipe.py
@@ -118,7 +118,7 @@ def create_temp_dir_and_files():
 
 
 def reset_after_n_next_calls(
-    datapipe: Union[IterDataPipe[T_co], MapDataPipe[T_co]], n: int
+    datapipe: IterDataPipe[T_co] | MapDataPipe[T_co], n: int
 ) -> tuple[list[T_co], list[T_co]]:
     """
     Given a DataPipe and integer n, iterate the DataPipe for n elements and store the elements into a list
@@ -2439,15 +2439,15 @@ def test_subtype(self):
                 self.assertFalse(issubtype(t1, t2))
 
         T = TypeVar("T", int, str)
-        S = TypeVar("S", bool, Union[str, int], tuple[int, T])  # type: ignore[valid-type]
+        S = TypeVar("S", bool, str | int, tuple[int, T])  # type: ignore[valid-type]
         types = (
-            (int, Optional[int]),
-            (list, Union[int, list]),
+            (int, Optional[int]),  # noqa: UP045
+            (list, Union[int, list]),  # noqa: UP007
             (tuple[int, str], S),
             (tuple[int, str], tuple),
             (T, S),
             (S, T_co),
-            (T, Union[S, set]),
+            (T, Union[S, set]),  # noqa: UP007
         )
         for sub, par in types:
             self.assertTrue(issubtype(sub, par))
@@ -2475,7 +2475,7 @@ def test_issubinstance(self):
 
         basic_data = (1, "1", True, 1.0, complex(1.0, 0.0))
         basic_type = (int, str, bool, float, complex)
-        S = TypeVar("S", bool, Union[str, int])
+        S = TypeVar("S", bool, str | int)
         for d in basic_data:
             self.assertTrue(issubinstance(d, Any))
             self.assertTrue(issubinstance(d, T_co))
diff --git a/test/test_decomp.py b/test/test_decomp.py
index 53ffd0e7b0bd5..64df1b88e44cf 100644
--- a/test/test_decomp.py
+++ b/test/test_decomp.py
@@ -228,12 +228,13 @@ def op_assert_ref(test_case, op, test_dtype, i, orig, decomp, ref, args, kwargs)
         (torch.float16, torch.ops.aten.reflection_pad3d_backward.default): 5e-3,
         (torch.bfloat16, torch.ops.aten.reflection_pad3d_backward.default): 5e-2,
         (torch.float16, torch.ops.aten._batch_norm_with_update.default): 2e-7,
-        (torch.bfloat16, torch.ops.aten._batch_norm_with_update.default): 2e-7,
+        (torch.bfloat16, torch.ops.aten._batch_norm_with_update.default): 5e-7,
         # see https://github.com/pytorch/pytorch/pull/96264
-        (torch.float16, torch.ops.aten.mv.default): 1e-5,
+        (torch.float16, torch.ops.aten.mv.default): 2e-5,
         (torch.bfloat16, torch.ops.aten.mv.default): 1e-5,
-        (torch.float16, torch.ops.aten.log_sigmoid_backward.default): 2e-5,
+        (torch.float16, torch.ops.aten.dot.default): 2e-6,
         (torch.float16, torch.ops.aten._softmax_backward_data.default): 3e-7,
+        (torch.bfloat16, torch.ops.aten._softmax_backward_data.default): 2e-7,
     }
     if ref.is_floating_point():
         orig_diff = (orig - ref).abs().max()
@@ -610,6 +611,29 @@ def test_quick_core_backward(self, device, dtype, op):
     def test_comprehensive(self, device, dtype, op):
         self.do_cross_ref(device, dtype, op, run_all=True)
 
+    def test_hann_window_decomp(self, device):
+        # Verify the hann_window decomp matches the native kernel for all four
+        # overloads: .default, .periodic, .out, .periodic_out.
+        from torch._decomp.decompositions import hann_window, hann_window_periodic
+
+        for n in (0, 1, 8, 9):
+            # .default (periodic=True)
+            ref = torch.hann_window(n, device=device)
+            res = hann_window(n, device=device)
+            self.assertEqual(ref, res)
+
+            # .periodic overload, explicit periodic flag
+            for periodic in (True, False):
+                ref = torch.hann_window(n, periodic, device=device)
+                res = hann_window_periodic(n, periodic, device=device)
+                self.assertEqual(ref, res)
+
+        # dtype forwarding
+        ref = torch.hann_window(8, dtype=torch.float64, device=device)
+        res = hann_window_periodic(8, dtype=torch.float64, device=device)
+        self.assertEqual(ref, res)
+        self.assertEqual(res.dtype, torch.float64)
+
     def test_uniform(self, device):
         size = (2, 3, 4, 5)
         dtype = torch.float32
@@ -1062,6 +1086,24 @@ def run_without_python_dispatcher(mode):
                     "only backwards is decomposed, but dtype doesn't support AD"
                 )
 
+    def test_binary_cross_entropy_with_logits_decomp(self, device):
+        op_config = {
+            "self": torch.randn([4, 5, 6], dtype=torch.bfloat16, device=device),
+            "target": torch.randn([4, 5, 6], dtype=torch.bfloat16, device=device),
+            "weight": torch.randn([6], dtype=torch.float32, device=device),
+            "reduction": 2,
+        }
+
+        ref = torch.ops.aten.binary_cross_entropy_with_logits.default(**op_config)
+
+        decomp_table = torch._inductor.decomposition.select_decomp_table()
+        bce_decomp = decomp_table[
+            torch.ops.aten.binary_cross_entropy_with_logits.default
+        ]
+        res = bce_decomp(**op_config)
+
+        torch.testing.assert_close(ref, res, check_dtype=True)
+
 
 instantiate_device_type_tests(TestDecomp, globals())
 
diff --git a/test/test_dlpack.py b/test/test_dlpack.py
index ad7889a360c78..e5cd153d5b270 100644
--- a/test/test_dlpack.py
+++ b/test/test_dlpack.py
@@ -10,8 +10,10 @@
     onlyCPU,
     onlyCUDA,
     onlyNativeDeviceTypes,
+    onlyOn,
     skipCUDAIfNotRocm,
     skipMeta,
+    skipXPUIf,
 )
 from torch.testing._internal.common_dtype import (
     all_mps_types_and,
@@ -101,26 +103,26 @@ def _dlpack_conversion_with_streams(self, stream, x):
             # DLPack protocol that establishes correct stream order
             # does not behave as expected on Jetson
             stream.synchronize()
-        stream = torch.cuda.Stream()
-        with torch.cuda.stream(stream):
+        stream = torch.Stream()
+        with stream:
             z = from_dlpack(x)
         stream.synchronize()
         return z
 
     @skipMeta
-    @onlyCUDA
+    @onlyOn(["xpu", "cuda"])
     @dtypes(*all_types_and_complex_and(torch.half, torch.bfloat16, torch.bool))
     def test_dlpack_conversion_with_streams(self, device, dtype):
         # Create a stream where the tensor will reside
-        stream = torch.cuda.Stream()
-        with torch.cuda.stream(stream):
+        stream = torch.Stream()
+        with stream:
             # Do an operation in the actual stream
             x = make_tensor((5,), dtype=dtype, device=device) + 1
         z = self._dlpack_conversion_with_streams(stream, x)
         self.assertEqual(z, x)
 
     @skipMeta
-    @onlyCUDA
+    @onlyOn(["xpu", "cuda"])
     @dtypes(
         torch.float8_e5m2,
         torch.float8_e5m2fnuz,
@@ -130,8 +132,8 @@ def test_dlpack_conversion_with_streams(self, device, dtype):
         torch.float4_e2m1fn_x2,
     )
     def test_dlpack_conversion_with_streams_narrow_precision(self, device, dtype):
-        stream = torch.cuda.Stream()
-        with torch.cuda.stream(stream):
+        stream = torch.Stream()
+        with stream:
             x = make_tensor((5,), dtype=torch.uint8, device=device) + 1
             x = x.view(dtype)
         z = self._dlpack_conversion_with_streams(stream, x)
@@ -361,7 +363,7 @@ def test_dlpack_invalid_cpu_stream(self):
             x.__dlpack__(stream=0)
 
     @skipMeta
-    @onlyCUDA
+    @onlyOn(["xpu", "cuda"])
     @deviceCountAtLeast(2)
     def test_dlpack_tensor_on_different_device(self, devices):
         dev0, dev1 = devices[:2]
@@ -372,7 +374,7 @@ def test_dlpack_tensor_on_different_device(self, devices):
         with self.assertRaisesRegex(
             BufferError, r"Can't export tensors on a different CUDA device"
         ):
-            with torch.cuda.device(dev1):
+            with torch.device(dev1):
                 x.__dlpack__()
 
     # TODO: add interchange tests once NumPy 1.22 (dlpack support) is required
@@ -521,7 +523,8 @@ def test_copy(self, device):
         self._test_from_dlpack(device, out_device="cpu", copy=True)
 
     @skipMeta
-    @onlyCUDA
+    @onlyOn(["xpu", "cuda"])
+    @skipXPUIf(True, "https://github.com/intel/torch-xpu-ops/issues/3077")
     def test_no_copy(self, device):
         # No copy, since tensor lives in the same device.
         self._test_from_dlpack(device)
@@ -576,6 +579,7 @@ def test_dlpack_unsupported_dtype_error(self, device):
 
     @skipMeta
     @onlyNativeDeviceTypes
+    @skipXPUIf(True, "https://github.com/intel/torch-xpu-ops/issues/3074")
     def test_dlpack_exchange_api(self, device):
         """Comprehensive test of all DLPack Exchange API functions using inline C++"""
         # Check that the C API capsule exists and get it
@@ -802,10 +806,13 @@ def test_dlpack_exchange_api(self, device):
             functions=["test_dlpack_exchange_api"],
             verbose=False,
             with_cuda=device.startswith("cuda"),
+            with_sycl=device.startswith("xpu"),
         )
 
         # Run the comprehensive C++ test
-        module.test_dlpack_exchange_api(tensor, api_capsule, device.startswith("cuda"))
+        module.test_dlpack_exchange_api(
+            tensor, api_capsule, device.startswith(("cuda", "xpu"))
+        )
 
     @skipMeta
     @onlyCUDA
@@ -858,7 +865,7 @@ def test_numpy_cross_device_transfer(self, device):
         self.assertEqual(np_array2[0], 999)
 
     @skipMeta
-    @onlyCUDA
+    @onlyOn(["xpu", "cuda"])
     @deviceCountAtLeast(2)
     def test_numpy_cross_device_multi_gpu(self, devices):
         """Test cross-device transfer to specific CUDA devices (cuda:0, cuda:1, etc)."""
@@ -883,7 +890,9 @@ def test_numpy_cross_device_multi_gpu(self, devices):
         self.assertNotEqual(t0.device, t1.device)
 
 
-instantiate_device_type_tests(TestTorchDlPack, globals(), allow_mps=True)
+instantiate_device_type_tests(
+    TestTorchDlPack, globals(), allow_mps=True, allow_xpu=True
+)
 
 if __name__ == "__main__":
     run_tests()
diff --git a/test/test_dynamic_shapes.py b/test/test_dynamic_shapes.py
index dc48ba1112d3b..8a0f6177781c6 100644
--- a/test/test_dynamic_shapes.py
+++ b/test/test_dynamic_shapes.py
@@ -29,14 +29,15 @@
     guard_bool,
     guard_float,
     guard_int,
+    guarding_hint_or_throw,
     GuardOnDataDependentSymNode,
     has_free_symbols,
     is_symbolic,
     ShapeEnv,
-    size_hint,
     StatelessSymbolicContext,
     statically_known_false,
     statically_known_true,
+    SYMPY_INTERP,
 )
 from torch.testing._internal.common_dtype import all_types_and
 from torch.testing._internal.common_utils import (
@@ -918,6 +919,46 @@ def test_non_overlapping_and_dense_unbacked(self):
             )
         )
 
+    def test_sympy_interp_is_non_overlapping_and_dense_flat_args(self):
+        # SYMPY_INTERP is used as the eval() namespace for guard code strings.
+        # Guard code prints IsNonOverlappingAndDenseIndicator(s0, s1, ..., st0, st1, ...)
+        # with flat args, so the SYMPY_INTERP function must accept flat args.
+        interp_fn = SYMPY_INTERP["IsNonOverlappingAndDenseIndicator"]
+
+        # 1D contiguous: sizes=(5,), strides=(1,)
+        self.assertEqual(interp_fn(5, 1), 1)
+        # 1D non-contiguous: sizes=(5,), strides=(2,)
+        self.assertEqual(interp_fn(5, 2), 0)
+        # 1D single element: sizes=(1,), strides=(42,)
+        self.assertEqual(interp_fn(1, 42), 1)
+
+        # 2D contiguous: sizes=(3, 4), strides=(4, 1)
+        self.assertEqual(interp_fn(3, 4, 4, 1), 1)
+        # 2D non-contiguous: sizes=(3, 4), strides=(5, 1) -- gap in memory
+        self.assertEqual(interp_fn(3, 4, 5, 1), 0)
+        # 2D transposed but still dense: sizes=(4, 3), strides=(1, 4)
+        self.assertEqual(interp_fn(4, 3, 1, 4), 1)
+
+        # 4D contiguous (the exact scenario from the MAST job failure):
+        # sizes=(2, 3, 4, 5), strides=(60, 20, 5, 1)
+        self.assertEqual(interp_fn(2, 3, 4, 5, 60, 20, 5, 1), 1)
+        # 4D non-contiguous:
+        self.assertEqual(interp_fn(2, 3, 4, 5, 100, 20, 5, 1), 0)
+
+    def test_sympy_interp_guard_eval_simulation(self):
+        # Simulate the actual guard eval() path: guard code strings use
+        # IsNonOverlappingAndDenseIndicator as a function name, and SYMPY_INTERP
+        # provides the binding in the eval namespace.
+        guard_code = "IsNonOverlappingAndDenseIndicator(2, 3, 4, 5, 60, 20, 5, 1) == 1"
+        result = eval(guard_code, SYMPY_INTERP)  # noqa: P204
+        self.assertTrue(result)
+
+        guard_code_false = (
+            "IsNonOverlappingAndDenseIndicator(2, 3, 4, 5, 100, 20, 5, 1) == 1"
+        )
+        result_false = eval(guard_code_false, SYMPY_INTERP)  # noqa: P204
+        self.assertFalse(result_false)
+
     def test_prims_is_non_overlapping_and_dense_or_false(self):
         shape_env = ShapeEnv()
         cf = torch._prims_common.is_non_overlapping_and_dense_or_false
@@ -1122,6 +1163,88 @@ def assert_not_optimized(sym):
         assert_not_optimized(torch.sym_max(mx3, s3 + s7))
         assert_not_optimized(torch.sym_max(mx3, s7 * 2))
 
+    def test_build_proxy_for_optimized_add(self):
+        """
+        Test that _build_proxy_for_sym_expr correctly handles flattened
+        sympy.Add expressions with more than 2 arguments.
+
+        When sympy.Add operations are optimized by _optimized_add
+        (in torch/fx/experimental/sym_node.py), they can be flattened into a
+        single Add node with 3 or more arguments (e.g., Add(s0, s1, s2)
+        instead of Add(Add(s0, s1), s2)).
+
+        _build_proxy_for_sym_expr must handle this by using torch.sym_sum
+        (which natively supports n-ary addition) instead of operator.add
+        (which is binary and would either raise TypeError or silently drop
+        arguments via FX codegen's 2-placeholder format string).
+        """
+        import torch.fx as fx
+        from torch.fx.experimental.proxy_tensor import (
+            _build_proxy_for_sym_expr,
+            _SympyExprTrackerValue,
+            PythonKeyTracer,
+            set_meta,
+        )
+        from torch.utils._thunk import Thunk
+
+        shape_env = ShapeEnv()
+        s0 = create_symint(shape_env, 2)
+        s1 = create_symint(shape_env, 3)
+        s2 = create_symint(shape_env, 4)
+
+        # This triggers _optimized_add and produces a flattened sympy.Add
+        # with 3 arguments: Add(s0_symbol, s1_symbol, s2_symbol)
+        optimized_sum = s0 + s1 + s2
+        self.assertTrue(optimized_sum.node._optimized_summation)
+        # Confirm the underlying sympy expression has 3 args (flattened).
+        self.assertEqual(len(optimized_sum.node.expr.args), 3)
+
+        # ------------------------------------------------------------------
+        # Case 1: out=None
+        #
+        # Previously this would call operator.add(a, b, c) and raise
+        # TypeError. Now uses torch.sym_sum which handles n-ary addition.
+        # ------------------------------------------------------------------
+        tracer_none = PythonKeyTracer()
+        for sym in [s0, s1, s2]:
+            tracer_none.sympy_expr_tracker[sym.node.expr] = _SympyExprTrackerValue(
+                proxy=sym, value=sym
+            )
+
+        result = _build_proxy_for_sym_expr(tracer_none, optimized_sum.node.expr)
+        self.assertEqual(int(result), 9)
+
+        # ------------------------------------------------------------------
+        # Case 2: out=optimized_sum  (the typical path from get_proxy_slot)
+        #
+        # Previously _sym_register would record operator.add with 3 args,
+        # and FX codegen's "{} + {}" format string would silently drop the
+        # third argument. Now uses torch.sym_sum which correctly handles
+        # all arguments.
+        # ------------------------------------------------------------------
+        tracer = PythonKeyTracer()
+        tracer.root = torch.nn.Module()
+        tracer.graph = fx.Graph(tracer_cls=PythonKeyTracer)
+
+        for sym, name in [(s0, "s0"), (s1, "s1"), (s2, "s2")]:
+            node = tracer.graph.placeholder(name)
+            proxy = fx.Proxy(node, tracer)
+            set_meta(proxy, sym)
+            tracer.sympy_expr_tracker[sym.node.expr] = _SympyExprTrackerValue(
+                proxy=proxy, value=sym
+            )
+            tracer.symnode_tracker[sym] = Thunk(lambda p=proxy: p)
+
+        _build_proxy_for_sym_expr(tracer, optimized_sum.node.expr, out=optimized_sum)
+
+        out_proxy = tracer.symnode_tracker[optimized_sum].force()
+        tracer.graph.output(out_proxy.node)
+
+        gm = fx.GraphModule(tracer.root, tracer.graph)
+
+        result = gm(2, 3, 4)  # should be 2 + 3 + 4 = 9
+        self.assertEqual(result, 9)
+
     def test_sym_max_multi_max_simplify(self):
         shape_env = ShapeEnv()
         u0 = shape_env.create_unbacked_symint()
@@ -1923,6 +2046,13 @@ def check(l, r):
         with self.assertRaises(ZeroDivisionError):
             y % x
 
+        # test pow operations preserve DynamicInt type
+        check(w**2, 1)  # DynamicInt ** int
+        check(2**z, 4)  # int ** DynamicInt
+        check(y**x, 1)  # DynamicInt ** DynamicInt
+        check(pow(z, 2), 4)  # pow(DynamicInt, int)
+        self.assertTrue(isinstance(pow(y, 3, 5), DynamicInt))  # pow with modulo
+
         # math, numpy
         self.assertEqual(math.cos(x), y)
         self.assertEqual(math.prod([z, z], start=z), 8)
@@ -3079,16 +3209,24 @@ def test_guards_gt_lt(self):
 
         guards = shape_env.produce_guards_expression([s0])
 
-        self.assertTrue(shape_env.evaluate_guards_expression(guards, [size_hint(s0)]))
-        self.assertFalse(shape_env.evaluate_guards_expression(guards, [size_hint(s1)]))
-        self.assertFalse(shape_env.evaluate_guards_expression(guards, [size_hint(s2)]))
+        self.assertTrue(
+            shape_env.evaluate_guards_expression(guards, [guarding_hint_or_throw(s0)])
+        )
+        self.assertFalse(
+            shape_env.evaluate_guards_expression(guards, [guarding_hint_or_throw(s1)])
+        )
+        self.assertFalse(
+            shape_env.evaluate_guards_expression(guards, [guarding_hint_or_throw(s2)])
+        )
 
     def test_guards_float_print(self):
         shape_env = ShapeEnv()
         s0 = create_symint(shape_env, 3)
         guard_bool(2 / s0 == 2 / 3)
         guards = shape_env.produce_guards_expression([s0])
-        self.assertTrue(shape_env.evaluate_guards_expression(guards, [size_hint(s0)]))
+        self.assertTrue(
+            shape_env.evaluate_guards_expression(guards, [guarding_hint_or_throw(s0)])
+        )
 
     @skipIfTorchDynamo("Not a TorchDynamo suitable test")
     @torch._dynamo.config.patch("capture_scalar_outputs", True)
@@ -3232,8 +3370,12 @@ def test_guards_float_div(self):
 
         self.assertIn("math.trunc(", guards)
         self.assertIn("float(", guards)
-        self.assertTrue(shape_env.evaluate_guards_expression(guards, [size_hint(s0)]))
-        self.assertFalse(shape_env.evaluate_guards_expression(guards, [size_hint(s1)]))
+        self.assertTrue(
+            shape_env.evaluate_guards_expression(guards, [guarding_hint_or_throw(s0)])
+        )
+        self.assertFalse(
+            shape_env.evaluate_guards_expression(guards, [guarding_hint_or_throw(s1)])
+        )
 
     @unittest.skipIf(
         TEST_XPU, "Skipped on XPU"
@@ -3612,6 +3754,58 @@ def func_negative(x):
             ):
                 func_negative(x)
 
+    def test_meta_copy(self):
+        """
+        Test that meta_copy_ does not raise when self and src have different
+        unbacked symint sizes.
+        """
+        from torch._meta_registrations import meta_copy_
+        from torch._subclasses.fake_tensor import FakeTensorMode
+
+        shape_env = ShapeEnv()
+        fake_mode = FakeTensorMode(shape_env=shape_env)
+        with fake_mode:
+            u0 = shape_env.create_unbacked_symint()
+            u1 = shape_env.create_unbacked_symint()
+
+            self_tensor = torch.empty((u0, 256))
+            src_tensor = torch.empty((u1, 256))
+
+            meta_copy_(self_tensor, src_tensor)
+
+    @torch._dynamo.config.patch("capture_scalar_outputs", True)
+    def test_unbacked_symbool_propagate_real_tensors(self):
+        """
+        Test that propagate_real_tensors properly handles SymBool from boolean .item() calls.
+
+        When tracing with propagate_real_tensors=True, if a boolean .item() returns False
+        during tracing, a runtime assertion should be generated that throws when the
+        boolean becomes True at runtime.
+
+        This is a regression test for a bug where int(real_t) was used instead of real_t,
+        causing boolean values to be incorrectly converted to integers and losing the
+        proper runtime assertion.
+        """
+
+        @torch.compile(backend="eager", fullgraph=True)
+        def f(x):
+            if x.eq(0.1).any().item():
+                return x
+            return x + 1
+
+        with torch._functorch.config.patch(fake_tensor_propagate_real_tensors=True):
+            # First call with 0.2 - item() returns False, traces the x + 1 branch
+            result1 = f(torch.ones(2) * 0.2)
+            torch.testing.assert_close(result1, torch.ones(2) * 1.2)
+
+            # Second call with 0.1 - item() returns True, should throw runtime assertion
+            # because the traced graph assumed False
+            with self.assertRaisesRegex(
+                RuntimeError,
+                r"Runtime assertion failed for expression Ne\(u0, 1\)",
+            ):
+                f(torch.ones(2) * 0.1)
+
 
 class TestUbackedOps(TestCase):
     @fresh_cache()
@@ -4608,6 +4802,32 @@ def f(x, start, end):
         res = f(x, start, 0)
         self.assertEqual(res.shape, torch.Size([0]))
 
+    @skipIfTorchDynamo("mark_unbacked not supported")
+    def test_unbacked_norm_no_dde(self):
+        def vector_norm(x):
+            return torch.linalg.vector_norm(x)
+
+        def dist_fn(x, y):
+            return torch.dist(x, y)
+
+        def norm_fro(x):
+            return torch.norm(x, p="fro")
+
+        x = torch.randn(4, 5)
+        y = torch.randn(4, 5)
+        for dim in range(x.ndim):
+            torch._dynamo.decorators.mark_unbacked(x, dim)
+            torch._dynamo.decorators.mark_unbacked(y, dim)
+
+        for fn, args in [
+            (vector_norm, (x,)),
+            (dist_fn, (x, y)),
+            (norm_fro, (x,)),
+        ]:
+            torch._dynamo.reset()
+            compiled = torch.compile(fn, fullgraph=True, backend="eager")
+            compiled(*args)
+
     @skipIfTorchDynamo()
     @torch.fx.experimental._config.patch("backed_size_oblivious", True)
     def test_backed_size_oblivious_expand(self):
@@ -4711,15 +4931,21 @@ def fn(x):
     def test_hint_override_consistent_stride1(self):
         @torch.compile(fullgraph=True, dynamic=True)
         def func(x):
-            a = torch.fx.experimental.symbolic_shapes.size_hint(x.size()[2])
-            b = torch.fx.experimental.symbolic_shapes.size_hint(x.stride()[1])
+            a = torch.fx.experimental.symbolic_shapes.guarding_hint_or_throw(
+                x.size()[2]
+            )
+            b = torch.fx.experimental.symbolic_shapes.guarding_hint_or_throw(
+                x.stride()[1]
+            )
             torch._check(a == b)
             torch._check(a == 6)
 
-            a = torch.fx.experimental.symbolic_shapes.size_hint(
+            a = torch.fx.experimental.symbolic_shapes.guarding_hint_or_throw(
                 x.size()[1] * x.size()[2]
             )
-            b = torch.fx.experimental.symbolic_shapes.size_hint(x.stride()[0])
+            b = torch.fx.experimental.symbolic_shapes.guarding_hint_or_throw(
+                x.stride()[0]
+            )
             torch._check(a == b)
             torch._check(a == 120)
 
@@ -4737,10 +4963,12 @@ def test_hint_override_consistent_stride2(self):
         @torch.compile(fullgraph=True, dynamic=True)
         def func(x):
             # only one of the sizes has hint overridden.
-            a = torch.fx.experimental.symbolic_shapes.size_hint(
+            a = torch.fx.experimental.symbolic_shapes.guarding_hint_or_throw(
                 x.size()[1] * x.size()[2]
             )
-            b = torch.fx.experimental.symbolic_shapes.size_hint(x.stride()[0])
+            b = torch.fx.experimental.symbolic_shapes.guarding_hint_or_throw(
+                x.stride()[0]
+            )
             torch._check(a == b)
             torch._check(a == 24)
 
@@ -4753,20 +4981,20 @@ def func(x):
         torch._dynamo.mark_dynamic(x, 2, hint_override=6)
         func(x)
 
-    def test_size_hint(self):
+    def test_optimization_hint(self):
         @torch.compile(fullgraph=True)
         def func(x):
             u0 = x.item()
             a = torch.ones([u0])
             torch._check(
-                torch.fx.experimental.symbolic_shapes.size_hint(
+                torch.fx.experimental.symbolic_shapes.optimization_hint(
                     a.size()[0], fallback=300
                 )
                 == 300
             )
             b = torch.ones([x.item() * 2])
             torch._check(
-                torch.fx.experimental.symbolic_shapes.size_hint(
+                torch.fx.experimental.symbolic_shapes.optimization_hint(
                     b.size()[0], fallback=300
                 )
                 == 600
@@ -4776,13 +5004,16 @@ def func(x):
 
         func(torch.tensor([33]))
 
-    def test_size_hint_no_fallback(self):
+    def test_guarding_hint_or_throw(self):
         @torch.compile(fullgraph=True)
         def func(x):
             u0 = x.item()
             a = torch.ones([u0])
             torch._check(
-                torch.fx.experimental.symbolic_shapes.size_hint(a.size()[0]) == 300
+                torch.fx.experimental.symbolic_shapes.guarding_hint_or_throw(
+                    a.size()[0]
+                )
+                == 300
             )
 
             return a * 10
@@ -4888,12 +5119,12 @@ def func(x):
         self.assertEqual(counter.frame_count, 2)
 
     @skipIfTorchDynamo("mark_unbacked is not traceable")
-    def test_shape_id_no_recompile_without_dynamic_indices(self):
+    def test_shape_id_no_recompile_without_unbacked_indices(self):
         """
-        Test that passing a tensor without _dynamo_dynamic_indices after
+        Test that passing a tensor without _dynamo_unbacked_indices after
         compiling with shape_ids does NOT trigger recompilation.
         The guard on shape_ids only applies when the runtime tensor
-        also has _dynamo_dynamic_indices.
+        also has _dynamo_unbacked_indices.
         """
         counter = CompileCounter()
 
@@ -4902,14 +5133,14 @@ def func(x):
 
         compiled_func = torch.compile(func, backend=counter)
 
-        # First call with shape_id (has _dynamo_dynamic_indices)
+        # First call with shape_id (has _dynamo_unbacked_indices)
         x1 = torch.rand(4, 3)
         torch._dynamo.decorators.mark_unbacked(x1, 0, shape_id="batch")
         compiled_func(x1)
         self.assertEqual(counter.frame_count, 1)
 
-        # Second call with regular tensor (no _dynamo_dynamic_indices)
-        # Should NOT recompile - guard passes when no _dynamo_dynamic_indices
+        # Second call with regular tensor (no _dynamo_unbacked_indices)
+        # Should NOT recompile - guard passes when no _dynamo_unbacked_indices
         x2 = torch.rand(4, 3)
         compiled_func(x2)
         self.assertEqual(counter.frame_count, 1)
@@ -4939,6 +5170,133 @@ def func(x):
         compiled_func(x2)
         self.assertEqual(counter.frame_count, 2)
 
+    @skipIfTorchDynamo("mark_unbacked is not traceable")
+    def test_unbacked_bounds_recompilation(self):
+        """
+        Test that changing _dynamo_unbacked_bounds triggers recompilation.
+        """
+        counter = CompileCounter()
+
+        def func(x):
+            return x + 1
+
+        compiled_func = torch.compile(func, backend=counter)
+
+        # First call with min/max bounds
+        x1 = torch.rand(4, 3)
+        torch._dynamo.decorators.mark_unbacked(x1, 0, min=1, max=100)
+        compiled_func(x1)
+        self.assertEqual(counter.frame_count, 1)
+
+        # Second call with same bounds - no recompilation
+        x2 = torch.rand(4, 3)
+        torch._dynamo.decorators.mark_unbacked(x2, 0, min=1, max=100)
+        compiled_func(x2)
+        self.assertEqual(counter.frame_count, 1)
+
+        # Third call without bounds - should recompile
+        x3 = torch.rand(4, 3)
+        torch._dynamo.decorators.mark_unbacked(x3, 0)
+        compiled_func(x3)
+        self.assertEqual(counter.frame_count, 2)
+
+        # Fourth call with different bounds - should recompile
+        x4 = torch.rand(4, 3)
+        torch._dynamo.decorators.mark_unbacked(x4, 0, min=1, max=200)
+        compiled_func(x4)
+        self.assertEqual(counter.frame_count, 3)
+
+    @skipIfTorchDynamo("mark_unbacked is not traceable")
+    def test_unbacked_bounds_no_recompile_without_unbacked_indices(self):
+        """
+        Test that passing a tensor without _dynamo_unbacked_indices after
+        compiling with bounds does NOT trigger recompilation.
+        The guard on bounds only applies when the runtime tensor
+        also has _dynamo_unbacked_indices.
+        """
+        counter = CompileCounter()
+
+        def func(x):
+            return x + 1
+
+        compiled_func = torch.compile(func, backend=counter)
+
+        # First call with bounds (has _dynamo_unbacked_indices)
+        x1 = torch.rand(4, 3)
+        torch._dynamo.decorators.mark_unbacked(x1, 0, min=1, max=100)
+        compiled_func(x1)
+        self.assertEqual(counter.frame_count, 1)
+
+        # Second call with regular tensor (no _dynamo_unbacked_indices)
+        # Should NOT recompile - guard passes when no _dynamo_unbacked_indices
+        x2 = torch.rand(4, 3)
+        compiled_func(x2)
+        self.assertEqual(counter.frame_count, 1)
+
+    @skipIfTorchDynamo("mark_unbacked is not traceable")
+    def test_unbacked_no_bounds_then_bounds(self):
+        """
+        Test that compiling without bounds then calling with bounds
+        triggers recompilation.
+        """
+        counter = CompileCounter()
+
+        def func(x):
+            return x + 1
+
+        compiled_func = torch.compile(func, backend=counter)
+
+        # First call with mark_unbacked but no bounds
+        x1 = torch.rand(4, 3)
+        torch._dynamo.decorators.mark_unbacked(x1, 0)
+        compiled_func(x1)
+        self.assertEqual(counter.frame_count, 1)
+
+        # Second call with same (no bounds) - no recompilation
+        x2 = torch.rand(4, 3)
+        torch._dynamo.decorators.mark_unbacked(x2, 0)
+        compiled_func(x2)
+        self.assertEqual(counter.frame_count, 1)
+
+        # Third call with min=1, max=100 - should recompile
+        # (compiled without bounds, now calling with bounds)
+        x3 = torch.rand(4, 3)
+        torch._dynamo.decorators.mark_unbacked(x3, 0, min=1, max=100)
+        compiled_func(x3)
+        self.assertEqual(counter.frame_count, 2)
+
+    @skipIfTorchDynamo("mark_unbacked is not traceable")
+    def test_unbacked_no_shape_id_then_shape_id(self):
+        """
+        Test that compiling without shape_id then calling with shape_id
+        triggers recompilation.
+        """
+        counter = CompileCounter()
+
+        def func(x):
+            return x + 1
+
+        compiled_func = torch.compile(func, backend=counter)
+
+        # First call with mark_unbacked but no shape_id
+        x1 = torch.rand(4, 3)
+        torch._dynamo.decorators.mark_unbacked(x1, 0)
+        compiled_func(x1)
+        self.assertEqual(counter.frame_count, 1)
+
+        # Second call with same (no shape_id) - no recompilation
+        x2 = torch.rand(4, 3)
+        torch._dynamo.decorators.mark_unbacked(x2, 0)
+        compiled_func(x2)
+        self.assertEqual(counter.frame_count, 1)
+
+        # Third call with shape_id="batch" - should recompile
+        # (compiled without shape_id, now calling with shape_id)
+        x3 = torch.rand(4, 3)
+        torch._dynamo.decorators.mark_unbacked(x3, 0, shape_id="batch")
+        compiled_func(x3)
+        self.assertEqual(counter.frame_count, 2)
+
     @skipIfTorchDynamo("mark_unbacked is not traceable")
     def test_unbacked_exec_fft_reshape_no_dde(self):
         """
diff --git a/test/test_fake_tensor.py b/test/test_fake_tensor.py
index 1ee0668c9dc63..313cc3fb96c5e 100644
--- a/test/test_fake_tensor.py
+++ b/test/test_fake_tensor.py
@@ -10,6 +10,8 @@
 import io
 import itertools
 import pickle
+import subprocess
+import sys
 import unittest
 import weakref
 from unittest.mock import patch
@@ -64,7 +66,9 @@
     skipIfCrossRef,
     skipIfTorchDynamo,
     skipIfWindows,
+    skipIfXpu,
     TemporaryFileName,
+    TEST_ACCELERATOR,
     TEST_WITH_TORCHDYNAMO,
     TestCase,
     xfailIfTorchDynamo,
@@ -82,6 +86,8 @@
 torch._dynamo.config.fake_tensor_cache_enabled = True
 torch._dynamo.config.fake_tensor_cache_crosscheck_enabled = True
 
+device_type = acc.type if (acc := torch.accelerator.current_accelerator()) else "cpu"
+
 
 def expectedFailurePropagateRealTensors(fn):
     fn._expected_failure_propagate_real_tensors = True
@@ -115,10 +121,9 @@ def test_basic(self):
             self.assertTrue(isinstance(z, FakeTensor))
 
     def test_custom_op_fallback(self):
-        from torch.library import impl, Library
+        from torch.library import _scoped_library, impl
 
-        try:
-            test_lib = Library("my_test_op", "DEF")  # noqa: TOR901
+        with _scoped_library("my_test_op", "DEF") as test_lib:
             test_lib.define("foo(Tensor self) -> Tensor")
 
             @impl(test_lib, "foo", "CPU")
@@ -133,9 +138,6 @@ def foo_impl(self):
                     x = mode.from_tensor(x)
                     torch.ops.my_test_op.foo(x)
 
-        finally:
-            test_lib._destroy()
-
     def test_parameter_instantiation(self):
         with FakeTensorMode():
             x = torch.rand([4])
@@ -191,6 +193,33 @@ def test_repr(self):
             x = torch.empty(2, 2, device="meta")
             self.assertEqual(repr(x), "FakeTensor(..., device='meta', size=(2, 2))")
 
+    def test_fake_device_property_normalization(self):
+        """Test that fake_device property normalizes device on assignment."""
+        with FakeTensorMode() as mode:
+            # Test CPU device - should remain unchanged
+            cpu_tensor = torch.empty(2, 2, device="meta")
+            fake_cpu = FakeTensor(mode, cpu_tensor, torch.device("cpu"))
+            self.assertEqual(fake_cpu.fake_device, torch.device("cpu"))
+
+            # Test device with explicit index - should remain unchanged
+            cuda_device_with_index = torch.device("cuda:0")
+            fake_cuda = FakeTensor(mode, cpu_tensor, cuda_device_with_index)
+            self.assertEqual(fake_cuda.fake_device, torch.device("cuda:0"))
+
+            # Test MPS device without index - should normalize to mps:0
+            mps_device = torch.device("mps")
+            fake_mps = FakeTensor(mode, cpu_tensor, mps_device)
+            self.assertEqual(fake_mps.fake_device, torch.device("mps:0"))
+
+            # Test property setter normalization with MPS
+            fake_tensor = FakeTensor(mode, cpu_tensor, torch.device("cpu"))
+            fake_tensor.fake_device = torch.device("mps")
+            self.assertEqual(fake_tensor.fake_device, torch.device("mps:0"))
+
+            # Test property setter normalization with CUDA
+            fake_tensor.fake_device = torch.device("cuda")
+            self.assertEqual(fake_tensor.fake_device, torch.device("cuda:0"))
+
     def test_convert_fake_to_real(self):
         x = torch.ones([20])
         with FakeTensorMode(allow_non_fake_inputs=True) as m:
@@ -323,6 +352,7 @@ def test_device_inplace_copy(self):
             if y.copy_(x).device.type != "cuda":
                 raise AssertionError("expected cuda device")
 
+    @unittest.skipIf(not RUN_CUDA, "requires cuda")
     def test_fake_device(self):
         t = torch.ones(3)
         t = t.view(1, 3)
@@ -573,6 +603,16 @@ def fn(x):
         out_eager = fn(torch.empty((0,)))
         self.checkMetaProps(out_fake, out_eager)
 
+    def test_as_strided_negative_stride_error(self):
+        error = (
+            r"as_strided: Negative strides are not supported at the "
+            r"moment, got strides: \[-?[0-9]+(, -?[0-9]+)*\]"
+        )
+        with FakeTensorMode():
+            x = torch.empty(0)
+            with self.assertRaisesRegex(RuntimeError, error):
+                torch.as_strided(x, (17, 18), (-80, 1), 1)
+
     @unittest.skipIf(not RUN_CUDA, "requires cuda")
     def test_cpu_fallback(self):
         with FakeTensorMode(allow_fallback_kernels=False):
@@ -1217,6 +1257,123 @@ def test_unbind_copy_out(self):
         self.assertEqual(out[1].dtype, eye.dtype)
         self.assertEqual(out[2].dtype, eye.dtype)
 
+    @unittest.skipIf(not torch.cuda._is_compiled(), "requires CUDA-compiled PyTorch")
+    def test_fake_device_guard_no_use_after_free(self):
+        # Regression test: when CUDA is compiled but no devices are visible,
+        # FakeTensorMode installs a FakeGuardImpl into the global
+        # device_guard_impl_registry.  The impl must outlive all threads that
+        # use it; previously it was stored as a thread-local unique_ptr, so
+        # the pointer became dangling after the thread exited, causing a
+        # segfault when a later thread called deviceCount() on it.
+        #
+        # CUDA_VISIBLE_DEVICES must be set before torch is imported, so we
+        # run the repro in a subprocess.
+
+        script = """\
+import os
+os.environ["CUDA_VISIBLE_DEVICES"] = ""
+import threading
+import torch
+
+def f(x):
+    return x * x
+
+def run_compile():
+    g = torch.compile(f, backend="eager")
+    for i in range(10):
+        g(torch.randn(i))
+
+threads = [threading.Thread(target=run_compile) for _ in range(2)]
+for t in threads:
+    t.start()
+for t in threads:
+    t.join()
+
+torch._dynamo.reset()
+
+threads = [threading.Thread(target=run_compile) for _ in range(2)]
+for t in threads:
+    t.start()
+for t in threads:
+    t.join()
+"""
+        result = subprocess.run(
+            [sys.executable, "-c", script],
+            capture_output=True,
+            timeout=60,
+        )
+        self.assertEqual(
+            result.returncode,
+            0,
+            msg=f"subprocess failed:\n{result.stderr.decode()}",
+        )
+
+    @unittest.skipIf(
+        TEST_ACCELERATOR, "Only execute when an accelerator is not present"
+    )
+    def test_avoid_device_init_without_backends(self):
+        fake_mode = FakeTensorMode()
+
+        self.assertTrue(
+            fake_mode.avoid_device_init,
+            "Expected avoid_device_init to return True when no backends are registered",
+        )
+
+    @parametrize("is_available", [False, True])
+    @skipIfTorchDynamo(
+        "TorchDynamo exposes https://github.com/pytorch/pytorch/issues/166696"
+    )
+    @unittest.skipIf(
+        TEST_ACCELERATOR, "Only execute when an accelerator is not present"
+    )
+    def test_avoid_device_init_with_privateuse1_backend(self, is_available):
+        class _DummyPrivateUse1Module:
+            @staticmethod
+            def is_available() -> bool:
+                return is_available
+
+        backend_name = "privateuseone"
+
+        try:
+            torch._register_device_module(backend_name, _DummyPrivateUse1Module)
+
+            fake_mode = FakeTensorMode()
+
+            self.assertEqual(fake_mode.avoid_device_init, not is_available)
+        finally:
+            delattr(torch, backend_name)
+            del sys.modules[f"torch.{backend_name}"]
+
+    def test_unique_output_dtype(self):
+        shape_env = ShapeEnv()
+        for input_dtype in [torch.float32, torch.float16, torch.bfloat16]:
+            x_real = torch.randn(10, dtype=input_dtype)
+            real_unique, real_inverse, real_counts = torch.unique(
+                x_real, return_inverse=True, return_counts=True
+            )
+            with FakeTensorMode(shape_env=shape_env):
+                x = torch.randn(10, dtype=input_dtype)
+                fake_unique, fake_inverse, fake_counts = torch.unique(
+                    x, return_inverse=True, return_counts=True
+                )
+                self.assertEqual(fake_unique.dtype, real_unique.dtype)
+                self.assertEqual(fake_inverse.dtype, real_inverse.dtype)
+                self.assertEqual(fake_counts.dtype, real_counts.dtype)
+
+        # Also test with dim argument
+        x_real = torch.randn(3, 4, dtype=torch.float32)
+        real_unique, real_inverse, real_counts = torch.unique(
+            x_real, dim=0, return_inverse=True, return_counts=True
+        )
+        with FakeTensorMode(shape_env=shape_env):
+            x = torch.randn(3, 4, dtype=torch.float32)
+            fake_unique, fake_inverse, fake_counts = torch.unique(
+                x, dim=0, return_inverse=True, return_counts=True
+            )
+            self.assertEqual(fake_unique.dtype, real_unique.dtype)
+            self.assertEqual(fake_inverse.dtype, real_inverse.dtype)
+            self.assertEqual(fake_counts.dtype, real_counts.dtype)
+
 
 instantiate_parametrized_tests(FakeTensorTest)
 
@@ -1617,6 +1774,7 @@ def test_cross_entropy_loss(self):
 
             self.assertEqual(ref.size(), meta_out.size())
 
+    @skipIfXpu(msg="MetadataMismatchError, torch-xpu-ops: 2802")
     @unittest.skipIf(
         not PLATFORM_SUPPORTS_FLASH_ATTENTION,
         "Does not support SDPA or pre-SM80 hardware",
@@ -1633,14 +1791,14 @@ def forward(self, arg1, arg2, arg3):
 
         args_new = [
             [
-                ((1, 48, 64, 64), (0, 4096, 64, 1), torch.float16, "cuda"),
-                ((1, 48, 64, 64), (0, 4096, 64, 1), torch.float16, "cuda"),
-                ((1, 48, 64, 64), (0, 4096, 64, 1), torch.float16, "cuda"),
+                ((1, 48, 64, 64), (0, 4096, 64, 1), torch.float16, device_type),
+                ((1, 48, 64, 64), (0, 4096, 64, 1), torch.float16, device_type),
+                ((1, 48, 64, 64), (0, 4096, 64, 1), torch.float16, device_type),
             ],
             [
-                ((4, 2, 16, 32), (1024, 512, 32, 1), torch.float16, "cuda"),
-                ((4, 2, 16, 32), (1024, 512, 32, 1), torch.float16, "cuda"),
-                ((4, 2, 16, 32), (1024, 512, 32, 1), torch.float16, "cuda"),
+                ((4, 2, 16, 32), (1024, 512, 32, 1), torch.float16, device_type),
+                ((4, 2, 16, 32), (1024, 512, 32, 1), torch.float16, device_type),
+                ((4, 2, 16, 32), (1024, 512, 32, 1), torch.float16, device_type),
             ],
         ]
         for args_list in args_new:
@@ -2834,5 +2992,69 @@ def test_fake_tensor_prefer_device_type_cpu_only(self):
                 self.assertTrue(isinstance(result, FakeTensor))
 
 
+class FakeTensorMetaDevicePropagation(TestCase):
+    @parametrize("device", ["cpu", "cuda"])
+    def test_inplace_add_with_meta_rhs_keeps_destination_device(self, device):
+        if device == "cuda" and not RUN_CUDA:
+            self.skipTest("requires cuda")
+
+        with FakeTensorMode():
+            log_det = torch.zeros(2, device=device)
+            log_det += torch.zeros(2, device="meta")
+
+            self.assertEqual(log_det.device.type, device)
+            self.assertTrue(isinstance(log_det, FakeTensor))
+
+
+instantiate_parametrized_tests(FakeTensorMetaDevicePropagation)
+
+
+class FakeTensorViewCopy(TestCase):
+    def test_expand_then_view_copy_matches_eager_mode(self):
+        x = torch.arange(7)
+        y = x.expand(12, 7)
+
+        # Eager baseline
+        eager = torch.ops.aten.view_copy.default(y, [84])
+        self.assertEqual(eager.shape, (84,))
+        self.assertIsNone(eager._base)  # non-aliasing
+
+        # FakeTensor behavior should match eager
+        with FakeTensorMode():
+            xf = torch.arange(7)
+            yf = xf.expand(12, 7)
+            fake = torch.ops.aten.view_copy.default(yf, [84])
+            self.assertEqual(fake.shape, (84,))
+            self.assertIsNone(fake._base)  # non-aliasing
+
+    def test_expand_then_view_still_not_allowed(self):
+        with FakeTensorMode():
+            xf = torch.arange(7)
+            yf = xf.expand(12, 7)
+            with self.assertRaisesRegex(
+                ValueError, "Cannot view a tensor with shape *"
+            ):
+                _ = yf.view(-1)
+
+    def test_expand_then_view_copy_unbacked_matches_eager_mode(self):
+        with torch.fx.experimental._config.patch(backed_size_oblivious=True):
+            with FakeTensorMode():
+                xf = torch.arange(7)
+                yf = xf.expand(12, 7)
+                fake = torch.ops.aten.view_copy.default(yf, [84])
+                self.assertEqual(fake.shape, (84,))
+                self.assertIsNone(fake._base)  # non-aliasing
+
+    def test_expand_then_view_unbacked_still_not_allowed(self):
+        with torch.fx.experimental._config.patch(backed_size_oblivious=True):
+            with FakeTensorMode():
+                xf = torch.arange(7)
+                yf = xf.expand(12, 7)
+                with self.assertRaisesRegex(
+                    ValueError, "Cannot view a tensor with shape *"
+                ):
+                    _ = yf.view(-1)
+
+
 if __name__ == "__main__":
     run_tests()
diff --git a/test/test_flop_counter.py b/test/test_flop_counter.py
index b3c498ce9eb9e..912883173b8fa 100644
--- a/test/test_flop_counter.py
+++ b/test/test_flop_counter.py
@@ -20,6 +20,7 @@
     TestCase,
 )
 from torch.testing._internal.triton_utils import requires_cuda_and_triton
+from torch.utils.flop_counter import sdpa_flop_count
 
 
 try:
@@ -49,6 +50,28 @@ def T(*shape, requires_grad=False):
     TEST_WITH_TORCHDYNAMO, "torchdynamo doesn't work with __torch_dispatch__ right now"
 )
 class TestFlopCounter(TestCase):
+    def test_sdpa_flop_count_gqa(self):
+        """sdpa_flop_count should handle GQA where KV heads < Q heads."""
+        # MHA: q_heads == kv_heads
+        q_shape = (2, 32, 128, 64)
+        k_shape = (2, 32, 128, 64)
+        v_shape = (2, 32, 128, 64)
+        mha_flops = sdpa_flop_count(q_shape, k_shape, v_shape)
+        self.assertTrue(mha_flops > 0)
+
+        # GQA: q_heads=32, kv_heads=8 (ratio 4:1)
+        k_gqa = (2, 8, 128, 64)
+        v_gqa = (2, 8, 128, 64)
+        gqa_flops = sdpa_flop_count(q_shape, k_gqa, v_gqa)
+        # Flops should be equal since KV heads are broadcast to match Q heads
+        self.assertEqual(gqa_flops, mha_flops)
+
+        # Incompatible: q_heads not a multiple of kv_heads
+        k_bad = (2, 5, 128, 64)
+        v_bad = (2, 5, 128, 64)
+        with self.assertRaises(AssertionError):
+            sdpa_flop_count(q_shape, k_bad, v_bad)
+
     def test_flop_counter_variety(self):
         mod = torch.nn.Linear(9, 10)
         with FlopCounterMode() as mode:
@@ -464,6 +487,70 @@ def get_flops(
         self.assertExpectedInline(str(flops_fw_bw_math), """805306368""")
         self.assertExpectedInline(str(flops_fw_bw_efficient), """939524096""")
 
+    @unittest.skipIf(not HAS_CUDA, "CUDA not available")
+    @unittest.skipIf(
+        not PLATFORM_SUPPORTS_FLASH_ATTENTION,
+        "Flash attention not supported (pre-SM80 hardware on CUDA)",
+    )
+    def test_sdpa_gqa(self):
+        """Test flop counting for grouped-query attention (GQA)."""
+        batch_size = 2
+        n_heads_q = 32
+        n_heads_kv = 8
+        seq_len = 128
+        head_dim = 64
+        dtype = torch.float16
+
+        query = torch.randn(
+            batch_size,
+            n_heads_q,
+            seq_len,
+            head_dim,
+            device="cuda",
+            dtype=dtype,
+        )
+        key = torch.randn(
+            batch_size,
+            n_heads_kv,
+            seq_len,
+            head_dim,
+            device="cuda",
+            dtype=dtype,
+        )
+        value = torch.randn(
+            batch_size,
+            n_heads_kv,
+            seq_len,
+            head_dim,
+            device="cuda",
+            dtype=dtype,
+        )
+
+        mode = FlopCounterMode()
+        with mode:
+            out = F.scaled_dot_product_attention(
+                query, key, value, dropout_p=0, is_causal=True, enable_gqa=True
+            )
+        gqa_flops = int(get_total_flops(mode))
+
+        # Compare with MHA (KV heads expanded to match Q heads)
+        key_expanded = key.repeat_interleave(n_heads_q // n_heads_kv, dim=1)
+        value_expanded = value.repeat_interleave(n_heads_q // n_heads_kv, dim=1)
+        mode_mha = FlopCounterMode()
+        with mode_mha:
+            out_mha = F.scaled_dot_product_attention(
+                query,
+                key_expanded,
+                value_expanded,
+                dropout_p=0,
+                is_causal=True,
+            )
+        mha_flops = int(get_total_flops(mode_mha))
+
+        # GQA flops should equal MHA flops (kernel broadcasts KV heads)
+        self.assertEqual(gqa_flops, mha_flops)
+        self.assertTrue(gqa_flops > 0)
+
     @unittest.skipIf(not HAS_CUDA, "CUDA not available")
     @unittest.skipIf(
         not PLATFORM_SUPPORTS_FLASH_ATTENTION
@@ -1087,6 +1174,94 @@ def test_scaled_mm(self):
 
         self.assertExpectedInline(get_total_flops(mode), """860160""")
 
+    @unittest.skipIf(not HAS_CUDA, "CUDA not available")
+    @unittest.skipIf(
+        not PLATFORM_SUPPORTS_FLASH_ATTENTION,
+        "Flash attention not supported (pre-SM80 hardware on CUDA)",
+    )
+    def test_varlen_attn(self):
+        import torch.nn.attention.varlen
+
+        n_heads = 8
+        head_dim = 64
+        dtype = torch.float16
+        seq_lens = [128, 64]
+        total_tokens = sum(seq_lens)
+        cu_seqs = torch.tensor(
+            [0] + list(torch.tensor(seq_lens).cumsum(0).tolist()),
+            dtype=torch.int32,
+            device="cuda",
+        )
+        max_s = max(seq_lens)
+
+        query = torch.randn(
+            total_tokens,
+            n_heads,
+            head_dim,
+            device="cuda",
+            dtype=dtype,
+            requires_grad=True,
+        )
+        key = torch.randn(
+            total_tokens,
+            n_heads,
+            head_dim,
+            device="cuda",
+            dtype=dtype,
+            requires_grad=True,
+        )
+        value = torch.randn(
+            total_tokens,
+            n_heads,
+            head_dim,
+            device="cuda",
+            dtype=dtype,
+            requires_grad=True,
+        )
+
+        mode = FlopCounterMode()
+        with mode:
+            out, _, _ = torch.ops.torch_attn._varlen_attn(
+                query,
+                key,
+                value,
+                cu_seqs,
+                cu_seqs,
+                max_s,
+                max_s,
+                is_causal=True,
+            )
+        fw_flops = int(get_total_flops(mode))
+        expected_fw = sum(
+            sdpa_flop_count(
+                (1, n_heads, s, head_dim),
+                (1, n_heads, s, head_dim),
+                (1, n_heads, s, head_dim),
+            )
+            for s in seq_lens
+        )
+        self.assertEqual(fw_flops, expected_fw)
+        # 2 bmms per sequence, each 2*h*s*d*s; total = 2048*(128^2 + 64^2) = 41943040
+        self.assertExpectedInline(str(fw_flops), """41943040""")
+
+        mode_bw = FlopCounterMode()
+        with mode_bw:
+            out, _, _ = torch.ops.torch_attn._varlen_attn(
+                query,
+                key,
+                value,
+                cu_seqs,
+                cu_seqs,
+                max_s,
+                max_s,
+                is_causal=True,
+            )
+            out.sum().backward()
+        fw_bw_flops = int(get_total_flops(mode_bw))
+        # fw=2 bmms, bw=5 bmms (flash recomputes scores), fw+bw = fw * 7/2
+        self.assertEqual(fw_bw_flops, fw_flops * 7 // 2)
+        self.assertExpectedInline(str(fw_bw_flops), """146800640""")
+
 
 if __name__ == "__main__":
     run_tests()
diff --git a/test/test_foreach.py b/test/test_foreach.py
index 5b6a84932f4e4..4a37d2b02da39 100644
--- a/test/test_foreach.py
+++ b/test/test_foreach.py
@@ -83,7 +83,7 @@ def __call__(self, inputs, is_cuda, expect_fastpath, **kwargs):
 
         # Skip profiler check for CUDA 12.6, 12.8 as the upgrade makes profiler results flaky
         # https://github.com/pytorch/pytorch/issues/148681. TODO: ADD IT BACK!!!
-        skip_profiler_check = _get_torch_cuda_version() in [(12, 6), (12, 8)]
+        skip_profiler_check = _get_torch_cuda_version() in [(12, 6), (12, 8)] or TEST_WITH_ROCM
         if (
             is_cuda
             and not skip_profiler_check
@@ -859,12 +859,12 @@ def test_binary_op_float_inf_nan(self, device, dtype, op):
 
     # note: Below three tests (postfixed with `_tensors_on_different_devices`)
     # checks whether foreach works with lists of tensors on different devices
-    # but tensors of the same index are on the same device, e.g., ['cuda', 'cpu].
+    # but tensors of the same index are on the same device, e.g., ['cuda', 'cpu'].
     @onlyCUDA
     @ops(foreach_unary_op_db)
     def test_unary_op_tensors_on_different_devices(self, device, dtype, op):
         method, ref, inplace_method, ref_inplace = self._get_funcs(op)
-        # tensors: ['cuda', 'cpu]
+        # tensors: ['cuda', 'cpu']
         tensors = next(
             iter(
                 op.sample_inputs(
@@ -876,26 +876,28 @@ def test_unary_op_tensors_on_different_devices(self, device, dtype, op):
             )
         ).input
         tensors[1] = tensors[1].to("cpu")
-        if not op.supports_out:
-            try:
-                actual = method((tensors,), False, False, zero_size=False)
-            except RuntimeError as e:
-                with self.assertRaisesRegex(type(e), str(e).splitlines()[0]):
-                    ref((tensors,))
-            else:
-                expected = ref((tensors,))
-                self.assertEqual(expected, actual)
 
         try:
-            inplace_method((tensors,), False, False, zero_size=False)
+            actual = method((tensors,), False, False, zero_size=False)
         except RuntimeError as e:
             with self.assertRaisesRegex(type(e), str(e).splitlines()[0]):
-                ref_inplace((tensors,))
+                ref((tensors,))
         else:
-            if not op.supports_out:
-                self.assertEqual(expected, tensors)
+            expected = ref((tensors,))
+            self.assertEqual(expected, actual)
+
+        # Some foreach functions (e.g. _foreach_clone) don't have an inplace variant, so
+        # we explicitly test for that here.
+        if not inplace_method.is_inplace:
+            self.assertIsNone(ref_inplace.func)
+        else:
+            try:
+                inplace_method((tensors,), False, False, zero_size=False)
+            except RuntimeError as e:
+                with self.assertRaisesRegex(type(e), str(e).splitlines()[0]):
+                    ref_inplace((tensors,))
             else:
-                self.assertEqual([torch.zeros_like(t) for t in tensors], tensors)
+                self.assertEqual(expected, tensors)
 
     @onlyCUDA
     @ops(filter(lambda op: op.supports_out, foreach_binary_op_db))
@@ -1057,11 +1059,11 @@ def test_big_num_tensors(self, device, dtype, op, use_cuda_graph, w_empty):
         # foreach_max cannot handle empty tensors as max requires an identity
         intersperse_empty_tensors = w_empty and op.name != "_foreach_max"
 
-        N = 600
+        N = 4000
         indices_with_empty_tensors = (
             set()
             if not intersperse_empty_tensors
-            else {200, 300, 301, 400, 401, 402, 404, 598}
+            else {200, 1500, 1501, 2800, 2801, 2802, 3500, 3998}
         )
         tensorlist = [
             make_tensor((2, 3), dtype=dtype, device=device, noncontiguous=False)
@@ -1074,7 +1076,7 @@ def test_big_num_tensors(self, device, dtype, op, use_cuda_graph, w_empty):
         import math
 
         if op.name == "_foreach_norm":
-            ords = [1, 2]
+            ords = [0, 1, 2]
             if not intersperse_empty_tensors:
                 # inf norm over an empty tensor is not defined by vector norm as it expects an identity
                 ords.append(math.inf)
@@ -1109,7 +1111,7 @@ def test_big_num_tensors(self, device, dtype, op, use_cuda_graph, w_empty):
     @ops(foreach_reduce_op_db)
     @parametrize("w_empty", (False, True))
     def test_foreach_reduce_large_input(self, device, dtype, op, w_empty):
-        # test inputs larger than kChunkSize (65536) * max_num_blocks (320)
+        # test inputs larger than kChunkSize (65536) * max_num_blocks (2240 for the 32kb config)
         N = 65536 * 320 * 2
         disable_fastpath = False
         kwargs = {}
diff --git a/test/test_functional_optim.py b/test/test_functional_optim.py
index e9c5566e26fd5..d1c17a9bc8be4 100644
--- a/test/test_functional_optim.py
+++ b/test/test_functional_optim.py
@@ -1,7 +1,6 @@
 # Owner(s): ["oncall: distributed"]
 
 import unittest
-from typing import Optional
 
 import torch
 import torch.distributed
@@ -56,14 +55,14 @@ def __init__(
         if len(params) == 0 and not _allow_empty_param_list:
             raise ValueError("optimizer got an empty parameter list")
 
-    def step_param(self, param: Tensor, grad: Optional[Tensor]):
+    def step_param(self, param: Tensor, grad: Tensor | None):
         # call the custom optimizer step_param implementation
         with torch.no_grad():
             raise RuntimeError(
                 "MyDummyFnOptimizer does not support step_param() as of now"
             )
 
-    def step(self, gradients: list[Optional[Tensor]]):
+    def step(self, gradients: list[Tensor | None]):
         # call the custom optimizer step implementation
         with torch.no_grad():
             raise RuntimeError("MyDummyFnOptimizer does not support step() as of now")
diff --git a/test/test_fx.py b/test/test_fx.py
index a9bbf8c22d699..fe3b6a702bc1b 100644
--- a/test/test_fx.py
+++ b/test/test_fx.py
@@ -37,7 +37,7 @@
 from torch.fx.operator_schemas import get_signature_for_torch_op
 from copy import deepcopy
 from collections import namedtuple
-from typing import Any, NamedTuple, Optional, Union
+from typing import Any, NamedTuple
 from collections.abc import Callable
 
 import torch
@@ -66,13 +66,17 @@
 from torch.fx._symbolic_trace import PHBase, PHWithMeta
 
 from torch.fx.proxy import TraceError
+from torch.testing._internal.common_cuda import blas_library_context
 from torch.testing._internal.common_utils import (
     find_library_location,
     IS_FBCODE,
     IS_MACOS,
+    IS_ARM64,
+    IS_LINUX,
     IS_WINDOWS,
     run_tests,
     skipIfTorchDynamo,
+    xfailIf,
 )
 from torch.testing._internal.jit_utils import JitTestCase
 
@@ -451,11 +455,11 @@ class NoMutableCallTracer(Tracer):
             def create_node(
                 self,
                 kind: str,
-                target: Union[str, Callable],
+                target: str | Callable,
                 args: tuple[Argument, ...],
                 kwargs: dict[str, Any],
-                name: Optional[str] = None,
-                type_expr: Optional[Any] = None,
+                name: str | None = None,
+                type_expr: Any | None = None,
             ) -> Node:
                 name = target if isinstance(target, str) else torch.typename(target)
                 if name[-1] == "_":
@@ -914,7 +918,7 @@ def lower_to_elementwise_interpreter(
 
             target_to_name = {operator.add: "add", operator.mul: "mul"}
 
-            output_node: Optional[Node] = None
+            output_node: Node | None = None
             # For each instruction, create a triple
             # (instruction_name : str, inputs : List[str], output : str)
             # to feed into the C++ interpreter
@@ -1076,11 +1080,11 @@ class TaggingTracer(Tracer):
             def create_node(
                 self,
                 kind: str,
-                target: Union[str, Callable],
+                target: str | Callable,
                 args: tuple[Argument, ...],
                 kwargs: dict[str, Any],
-                name: Optional[str] = None,
-                type_expr: Optional[Any] = None,
+                name: str | None = None,
+                type_expr: Any | None = None,
             ) -> Node:
                 n = super().create_node(kind, target, args, kwargs, name)
                 n.tag = "foo"
@@ -1956,8 +1960,8 @@ def forward(self, x):
 
         # Make sure we're testing all opcodes
         opcodes = set()
-        output_shape: Optional[torch.Shape] = None
-        output_stride: Optional[tuple[int]] = None
+        output_shape: torch.Shape | None = None
+        output_stride: tuple[int] | None = None
         for node in tc_traced.graph.nodes:
             opcodes.add(node.op)
             if node.op == "output":
@@ -2053,6 +2057,7 @@ def is_leaf_module(self, m, module_qualified_name):
                         f"got {tensor_meta[1].shape}"
                     )
 
+    @xfailIf(IS_ARM64 and IS_LINUX) # RuntimeError: label is too far
     def test_shape_prop_layout_3d(self):
         class ConvTest3d(torch.nn.Module):
             def __init__(self) -> None:
@@ -4193,11 +4198,11 @@ class TaggingTracer(Tracer):
             def create_node(
                 self,
                 kind: str,
-                target: Union[str, Callable],
+                target: str | Callable,
                 args: tuple[Argument, ...],
                 kwargs: dict[str, Any],
-                name: Optional[str] = None,
-                type_expr: Optional[Any] = None,
+                name: str | None = None,
+                type_expr: Any | None = None,
             ) -> Node:
                 n = super().create_node(kind, target, args, kwargs, name)
                 n.tag = "foo"
@@ -4426,8 +4431,10 @@ def fn(a, b, c, d):
         # recorver mutable checking flag
         torch.fx.proxy.TracerBase.check_mutable_operations = orig_tracer_mutable_flag
 
+    # This only fails on navi31
     @unittest.skipIf(not torch.cuda.is_available(), "CUDA not available")
     @torch.fx.experimental._config.patch("enrich_profiler_metadata", True)
+    @blas_library_context("cublaslt")
     def test_profiler_stack_trace_augmentation(self):
         """
         Test that map_recorded_events_to_aten_ops_with_stack_trace correctly
@@ -4476,8 +4483,29 @@ def forward(self, x):
         else:
             kernel_event = "cudaLaunchKernel"
             kernel_event_relu = "cudaLaunchKernel"
-
-        expected = f"""\
+        if IS_WINDOWS:
+            expected = f"""\
+event=aten::t node=t stack_trace=return F.linear(input, self.weight, self.bias)
+event=aten::transpose node=t stack_trace=return F.linear(input, self.weight, self.bias)
+event=aten::as_strided node=t stack_trace=return F.linear(input, self.weight, self.bias)
+event=aten::addmm node=addmm stack_trace=return F.linear(input, self.weight, self.bias)
+event=aten::expand node=addmm stack_trace=return F.linear(input, self.weight, self.bias)
+event=aten::as_strided node=addmm stack_trace=return F.linear(input, self.weight, self.bias)
+event={kernel_event} node=addmm stack_trace=return F.linear(input, self.weight, self.bias)
+event={kernel_event} node=addmm stack_trace=return F.linear(input, self.weight, self.bias)
+event=aten::relu node=relu stack_trace=return F.relu(input, inplace=self.inplace)
+event=aten::clamp_min node=relu stack_trace=return F.relu(input, inplace=self.inplace)
+event={kernel_event_relu} node=relu stack_trace=return F.relu(input, inplace=self.inplace)
+event=aten::t node=t_1 stack_trace=return F.linear(input, self.weight, self.bias)
+event=aten::transpose node=t_1 stack_trace=return F.linear(input, self.weight, self.bias)
+event=aten::as_strided node=t_1 stack_trace=return F.linear(input, self.weight, self.bias)
+event=aten::addmm node=addmm_1 stack_trace=return F.linear(input, self.weight, self.bias)
+event=aten::expand node=addmm_1 stack_trace=return F.linear(input, self.weight, self.bias)
+event=aten::as_strided node=addmm_1 stack_trace=return F.linear(input, self.weight, self.bias)
+event={kernel_event} node=addmm_1 stack_trace=return F.linear(input, self.weight, self.bias)
+event={kernel_event} node=addmm_1 stack_trace=return F.linear(input, self.weight, self.bias)"""
+        else:
+            expected = f"""\
 event=aten::t node=t stack_trace=x = self.linear1(x)
 event=aten::transpose node=t stack_trace=x = self.linear1(x)
 event=aten::as_strided node=t stack_trace=x = self.linear1(x)
@@ -4835,6 +4863,11 @@ def _annotation_type_to_stable_str(self, t, sig_str, recursive: bool = False):
 
         # Forward ref
         if isinstance(t, str):
+            # Normalize "X | None" string annotations to Optional format
+            if t.endswith(" | None"):
+                inner = t[: -len(" | None")]
+                result = f"Optional[{inner}]"
+                return result if recursive else f"'{result}'"
             if recursive:
                 return t
             else:
@@ -4874,7 +4907,7 @@ def _annotation_type_to_stable_str(self, t, sig_str, recursive: bool = False):
 
         if origin in {tuple, tuple}:
             return f"Tuple{contained_type_str}"
-        if origin == typing.Union:
+        if origin == typing.Union or isinstance(t, types.UnionType):
             # Annoying hack to detect Optional
             if len(contained) == 2 and (contained[0] is type(None)) ^ (
                 contained[1] is type(None)
diff --git a/test/test_fx_experimental.py b/test/test_fx_experimental.py
index c732c586d1171..0ff6bb62ab53d 100644
--- a/test/test_fx_experimental.py
+++ b/test/test_fx_experimental.py
@@ -45,6 +45,7 @@
 from torch.fx.passes import graph_manipulation
 from torch.fx.passes.param_fetch import lift_lowering_attrs_to_nodes
 from torch.fx.passes.shape_prop import ShapeProp
+from torch.fx._lazy_graph_module import _use_lazy_graph_module
 from torch.fx.passes.split_module import split_module
 from torch.fx.passes.annotate_getitem_nodes import annotate_getitem_nodes
 from torch.testing._internal.common_device_type import (
@@ -54,7 +55,14 @@
 )
 from torch.testing._internal.common_methods_invocations import op_db
 from torch.testing._internal.common_nn import module_tests, get_new_module_tests
-from torch.testing._internal.common_utils import TEST_Z3, run_tests, TestCase, TEST_WITH_CROSSREF
+from torch.testing._internal.common_utils import (
+    instantiate_parametrized_tests,
+    parametrize,
+    TEST_Z3,
+    run_tests,
+    TestCase,
+    TEST_WITH_CROSSREF,
+)
 from torch.testing._internal.jit_utils import JitTestCase
 import torch.utils._pytree as pytree
 
@@ -72,7 +80,7 @@
 )
 
 
-def symbolic_trace_with_rewrite(root: Union[torch.nn.Module, Callable]) -> GraphModule:
+def symbolic_trace_with_rewrite(root: torch.nn.Module | Callable) -> GraphModule:
     return GraphModule(
         root if isinstance(root, torch.nn.Module) else torch.nn.Module(),
         RewritingTracer().trace(root),
@@ -754,7 +762,8 @@ def forward(self, a, b):
         # Confirm that the output is correct
         self.assertEqual(traced(3, 3), m(3, 3))
 
-    def test_subgraph_creation(self):
+    @parametrize("use_lazy", [True, False])
+    def test_subgraph_creation(self, use_lazy):
         class MyModule(torch.nn.Module):
             def __init__(self) -> None:
                 super().__init__()
@@ -786,9 +795,10 @@ def mod_partition(node: Node):
             return partition
 
         # split module in module with submodules
-        module_with_submodules = split_module(
-            my_module_traced, my_module, mod_partition
-        )
+        with _use_lazy_graph_module(use_lazy):
+            module_with_submodules = split_module(
+                my_module_traced, my_module, mod_partition
+            )
 
         # Check that test_meta_info was still on all nodes.
         submodules = dict(module_with_submodules.named_modules())
@@ -897,6 +907,40 @@ def cb(_):
         else:
             raise RuntimeError("Expected the subgraph to have an output node.")
 
+    def test_split_module_tuple_return(self):
+        from torch._inductor.compile_fx import graph_returns_tuple
+
+        class M(torch.nn.Module):
+            def forward(self, x, y):
+                a = x + y
+                return a * x
+
+        gm = torch.fx.symbolic_trace(M())
+
+        # Assign ops to different partitions so a single-output submodule exists.
+        def partition_fn(node):
+            return 0 if node.target == operator.add else 1
+
+        # Without tuple_return: single-output submodules return a bare value.
+        sp = split_module(gm, None, partition_fn)
+        self.assertTrue(
+            any(
+                not graph_returns_tuple(submod)
+                for submod in sp.children()
+            ),
+            "expected at least one non-tuple-returning submodule",
+        )
+        x, y = torch.randn(4), torch.randn(4)
+        self.assertEqual(sp(x, y), gm(x, y))
+
+        # With tuple_return: all submodules return a tuple.
+        sp_boxed = split_module(gm, None, partition_fn, tuple_return=True)
+        self.assertTrue(
+            all(graph_returns_tuple(submod) for submod in sp_boxed.children()),
+            "all submodules should return a tuple with tuple_return=True",
+        )
+        self.assertEqual(sp_boxed(x, y), gm(x, y))
+
 
     def test_split_module_kwargs_expansion(self):
         class ModuleWithKwargsExpansion(torch.nn.Module):
@@ -1679,8 +1723,8 @@ def test_type_matches(self):
             (numbers.Number, int),
             (numbers.Number, float),
             (int, type(torch.float)),
-            (Union[int, float], int),
-            (Union[int, float], float),
+            (Union[int, float], int),  # noqa: UP007
+            (Union[int, float], float),  # noqa: UP007
             (list[int], int),
             (list[int], create_type_hint([int, int])),
             (list[int], create_type_hint((int, int))),
@@ -1700,8 +1744,8 @@ def test_type_matches(self):
             (torch.Tensor, torch.nn.Parameter),
             (list[torch.Tensor], create_type_hint((torch.nn.Parameter, torch.Tensor))),
             (list[torch.Tensor], create_type_hint((torch.Tensor, torch.nn.Parameter))),
-            (Optional[list[torch.Tensor]], list[torch.Tensor]),
-            (Optional[list[int]], list[int]),
+            (Optional[list[torch.Tensor]], list[torch.Tensor]),  # noqa: UP045
+            (Optional[list[int]], list[int]),  # noqa: UP045
         ] + [
             # pre-PEP585 signatures
             (typing.List[int], int),  # noqa: UP006
@@ -1721,8 +1765,8 @@ def test_type_matches(self):
             ),
             (typing.List[torch.Tensor], create_type_hint((torch.nn.Parameter, torch.Tensor))),  # noqa: UP006
             (typing.List[torch.Tensor], create_type_hint((torch.Tensor, torch.nn.Parameter))),  # noqa: UP006
-            (Optional[typing.List[torch.Tensor]], typing.List[torch.Tensor]),  # noqa: UP006
-            (Optional[typing.List[int]], typing.List[int]),  # noqa: UP006
+            (Optional[typing.List[torch.Tensor]], typing.List[torch.Tensor]),  # noqa: UP006, UP045
+            (Optional[typing.List[int]], typing.List[int]),  # noqa: UP006, UP045
         ]
 
         for sig_type, arg_type in should_be_equal:
@@ -1730,7 +1774,7 @@ def test_type_matches(self):
 
         should_fail = [
             (int, float),
-            (Union[int, float], str),
+            (Union[int, float], str),  # noqa: UP007
             (list[torch.Tensor], typing.List[int]),  # noqa: UP006
         ] + [
             # pre-PEP585 signatures
@@ -2152,6 +2196,7 @@ def test_z3str(self):
                 self.assertEqual(z3str(expr), expected)
 
 
+instantiate_parametrized_tests(TestFXExperimental)
 instantiate_device_type_tests(TestNormalizeOperators, globals())
 
 if __name__ == "__main__":
diff --git a/test/test_fx_passes.py b/test/test_fx_passes.py
index 4a48d0cd966aa..4182062bd8409 100644
--- a/test/test_fx_passes.py
+++ b/test/test_fx_passes.py
@@ -11,7 +11,7 @@
 
 from torch.fx.passes.infra.partitioner import CapabilityBasedPartitioner
 from torch.fx.passes.operator_support import OperatorSupport
-from torch.fx.passes.utils.fuser_utils import fuse_by_partitions
+from torch.fx.passes.utils.fuser_utils import fuse_by_partitions, topo_sort
 from torch.fx.passes.utils.matcher_utils import SubgraphMatcher
 
 from torch.testing._internal.common_utils import run_tests, parametrize, instantiate_parametrized_tests
@@ -428,6 +428,51 @@ def test_fuser_pass_deep_model(self):
                                                  allows_single_node_partition=True)
         partitions = partitioner.propose_partitions()
 
+    def test_topo_sort_stability(self):
+        """Test that topo_sort preserves relative order of independent nodes."""
+        m = TestModule()
+        gm = symbolic_trace(m)
+
+        nodes_by_name = {node.name: node for node in gm.graph.nodes}
+
+        # add_2 and linear depend on add_1 but are independent of each other.
+        # add_3 depends on both add_2 and linear (via add_1 -> linear -> add_3
+        # and add_1 -> add_2 -> add_3 paths in TestModule).
+        # Provide them in a specific order and verify it is preserved.
+        node_list = [
+            nodes_by_name["add_1"],
+            nodes_by_name["add_2"],
+            nodes_by_name["linear"],
+            nodes_by_name["add_3"],
+        ]
+        sorted_nodes = topo_sort(node_list)
+        sorted_names = [n.name for n in sorted_nodes]
+
+        # add_1 must come first (dependency of add_2, linear, add_3)
+        self.assertEqual(sorted_names[0], "add_1")
+        # add_2 and linear are independent; stable sort preserves input order
+        self.assertEqual(sorted_names[1], "add_2")
+        self.assertEqual(sorted_names[2], "linear")
+        # add_3 depends on both, so it comes last
+        self.assertEqual(sorted_names[3], "add_3")
+
+        # Now reverse the order of the independent nodes and verify stability
+        # is with respect to the new input order.
+        node_list_reversed = [
+            nodes_by_name["add_1"],
+            nodes_by_name["linear"],
+            nodes_by_name["add_2"],
+            nodes_by_name["add_3"],
+        ]
+        sorted_nodes_reversed = topo_sort(node_list_reversed)
+        sorted_names_reversed = [n.name for n in sorted_nodes_reversed]
+
+        self.assertEqual(sorted_names_reversed[0], "add_1")
+        # Now linear should come before add_2 since that's the input order
+        self.assertEqual(sorted_names_reversed[1], "linear")
+        self.assertEqual(sorted_names_reversed[2], "add_2")
+        self.assertEqual(sorted_names_reversed[3], "add_3")
+
     def test_partitioner_nested_getitem_chains(self):
         """Test that nested getitem chains are properly reassigned to producer's partition.
 
@@ -855,7 +900,7 @@ def setup(cls):
 
     @classmethod
     def tearDown(cls):
-        del cls.quantization
+        cls.quantization._destroy()
 
     @staticmethod
     def forward(self, arg0_1, arg1_1):
diff --git a/test/test_hop_infra.py b/test/test_hop_infra.py
index 6724c4ddd5360..58f1a3819a592 100644
--- a/test/test_hop_infra.py
+++ b/test/test_hop_infra.py
@@ -63,6 +63,7 @@ def test_all_hops_are_imported(self):
         FIXME_ALLOWLIST = {
             "autograd_function_apply",
             "run_with_rng_state",
+            "run_dtensor_rng_op",
             "graphsafe_run_with_rng_state",
             "map_impl",
             "_export_tracepoint",
diff --git a/test/test_indexing.py b/test/test_indexing.py
index 92aece3b070f9..87c3ddbc56e74 100644
--- a/test/test_indexing.py
+++ b/test/test_indexing.py
@@ -23,6 +23,7 @@
     onlyCPU,
     onlyNativeDeviceTypes,
     onlyOn,
+    skipMPS,
     skipXLA,
     skipXPUIf,
 )
@@ -31,6 +32,7 @@
     all_types_and,
     all_types_and_complex_and,
     all_types_complex_float8_and,
+    highest_precision_float,
 )
 from torch.testing._internal.common_utils import (
     DeterministicGuard,
@@ -1216,7 +1218,7 @@ def func1(x, i, v):
 
     @onlyNativeDeviceTypes
     def test_index_put_accumulate_duplicate_indices(self, device):
-        dtype = torch.float if device.startswith("mps") else torch.double
+        dtype = highest_precision_float(device)
         for i in range(1, 512):
             # generate indices by random walk, this will create indices with
             # lots of duplicates interleaved with each other
@@ -1816,17 +1818,20 @@ def test_index_limits(self, device):
 
     @parametrize("reduce", ["prod", "amin", "amax", "mean"])
     @dtypes(*all_types_and(torch.half, torch.bfloat16))
-    @expectedFailureMPS  # Unimplemented for MPS device
+    @dtypesIfMPS(
+        torch.half, torch.bfloat16, torch.float32, torch.int32, torch.int16, torch.int8
+    )
     def test_index_reduce(self, device, dtype, reduce):
         size = (3, 4, 5)
         index_dtypes = [torch.int, torch.long]
         include_selfs = [True, False]
+        noncontig_opts = [True, False]
         amin_init = float("inf") if dtype.is_floating_point else torch.iinfo(dtype).max
         amax_init = -float("inf") if dtype.is_floating_point else torch.iinfo(dtype).min
         reduction_init = {"prod": 1, "mean": 0, "amin": amin_init, "amax": amax_init}
 
         for dest_noncontig, src_noncontig, index_noncontig in product(
-            [True, False], repeat=3
+            noncontig_opts, repeat=3
         ):
             for idx_dtype, include_self in product(index_dtypes, include_selfs):
                 for dim in range(len(size)):
@@ -1884,7 +1889,16 @@ def test_index_reduce(self, device, dtype, reduce):
                             expected.div_(counts, rounding_mode="floor")
                     expected = expected.transpose(0, dim)
 
-                    self.assertEqual(dest, expected)
+                    # MPS uses atomics for index_reduce which causes
+                    # non-deterministic rounding for low-precision types
+                    kwargs = {}
+                    if (
+                        "mps" in device
+                        and dtype in [torch.bfloat16, torch.float16]
+                        and reduce in ["mean", "prod"]
+                    ):
+                        kwargs = {"atol": 0.02, "rtol": 0.1}
+                    self.assertEqual(dest, expected, **kwargs)
 
     @dtypes(*all_types_and_complex_and(torch.half, torch.bool, torch.bfloat16))
     @dtypesIfMPS(*all_mps_types_and(torch.bool, torch.cfloat))
@@ -2009,6 +2023,109 @@ def test_index_copy_deterministic(self, device: torch.device) -> None:
 
             self.assertEqual(x0, y0, atol=0, rtol=0)
 
+    @onlyNativeDeviceTypes
+    @skipMPS  # https://github.com/pytorch/pytorch/issues/161029
+    @dtypes(torch.float, torch.bfloat16, torch.half)
+    def test_index_add_large_index_vectorized_path(self, device, dtype):
+        # Tests that exercise both the VecSize=4 (vectorized) and VecSize=1
+        # (scalar) paths in the indexFuncLargeIndex CUDA kernel.
+        # VecSize=4 activates when: IndexIsMajor=true, sliceSize % 4 == 0,
+        # and both dst/src have stride-1 on their innermost dimension.
+        # numIndex must be > 16 to hit the LargeIndex kernel at all.
+        #
+        # Strategy: for each shape that hits the vectorized path, also run
+        # the same operation with sliceSize+1 (which forces the scalar path)
+        # and compare both against a CPU float reference.  We use UNIQUE
+        # indices to avoid non-deterministic atomic accumulation ordering.
+        num_idx = 64
+        num_dest = 100
+
+        def _run_and_check(dst_shape, dim, num_idx, num_dest_dim, tag):
+            dst = torch.zeros(dst_shape, dtype=dtype, device=device)
+            src_shape = list(dst_shape)
+            src_shape[dim] = num_idx
+            src = torch.randn(src_shape, dtype=dtype, device=device)
+            index = torch.randperm(num_dest_dim, device=device)[:num_idx]
+
+            # Reference: run the same op on CPU in float
+            dst_ref = dst.detach().cpu().float()
+            src_ref = src.detach().cpu().float()
+            idx_cpu = index.cpu()
+            dst_ref.index_add_(dim, idx_cpu, src_ref)
+
+            dst.index_add_(dim, index, src)
+
+            atol, rtol = (
+                (1e-2, 1e-2) if dtype in (torch.bfloat16, torch.half) else (1e-5, 1e-5)
+            )
+            self.assertEqual(
+                dst.cpu().float(),
+                dst_ref,
+                atol=atol,
+                rtol=rtol,
+                msg=f"Failed for {tag}, shape={dst_shape}, dim={dim}",
+            )
+
+        # --- VecSize=4 path: contiguous, dim=0, sliceSize divisible by 4 ---
+        for slice_size in [32, 128, 256]:
+            _run_and_check(
+                (num_dest, slice_size),
+                0,
+                num_idx,
+                num_dest,
+                f"vec4 sliceSize={slice_size}",
+            )
+
+        # --- VecSize=1 path: sliceSize NOT divisible by 4 ---
+        for slice_size in [3, 17, 33]:
+            _run_and_check(
+                (num_dest, slice_size),
+                0,
+                num_idx,
+                num_dest,
+                f"scalar sliceSize={slice_size}",
+            )
+
+        # --- VecSize=1 path: non-contiguous (stride != 1 on inner dim) ---
+        dst_base = torch.zeros(num_dest, 256, dtype=dtype, device=device)
+        src_base = torch.randn(num_idx, 256, dtype=dtype, device=device)
+        dst = dst_base[:, ::2]  # shape [100, 128], stride [256, 2]
+        src = src_base[:, ::2]
+        index = torch.randperm(num_dest, device=device)[:num_idx]
+
+        dst_ref = dst.detach().cpu().float().contiguous()
+        src_ref = src.detach().cpu().float().contiguous()
+        dst_ref.index_add_(0, index.cpu(), src_ref)
+        dst.index_add_(0, index, src)
+        atol, rtol = (
+            (1e-2, 1e-2) if dtype in (torch.bfloat16, torch.half) else (1e-5, 1e-5)
+        )
+        self.assertEqual(
+            dst.cpu().float(),
+            dst_ref,
+            atol=atol,
+            rtol=rtol,
+            msg="Failed for non-contiguous (stride-2)",
+        )
+
+        # --- VecSize=1 path: IndexIsMajor=false (dim=1 on row-major tensor) ---
+        _run_and_check(
+            (32, num_dest),
+            1,
+            num_idx,
+            num_dest,
+            "IndexIsMajor=false dim=1",
+        )
+
+        # --- VecSize=4 path: 3-D contiguous, dim=0, sliceSize divisible by 4 ---
+        _run_and_check(
+            (100, 8, 16),
+            0,
+            num_idx,
+            100,
+            "3D vec4 sliceSize=128",
+        )
+
     @onlyNativeDeviceTypes
     @expectedFailureMPS  # See https://github.com/pytorch/pytorch/issues/161029
     def test_index_add_deterministic(self, device: torch.device) -> None:
diff --git a/test/test_inspect_utils.py b/test/test_inspect_utils.py
new file mode 100644
index 0000000000000..b9480b62507b8
--- /dev/null
+++ b/test/test_inspect_utils.py
@@ -0,0 +1,150 @@
+# Owner(s): ["module: fx"]
+
+import inspect
+
+from torch.testing._internal.common_utils import run_tests, TestCase
+from torch.utils._inspect import _fast_bind
+
+
+class TestFastBind(TestCase):
+    def _assert_fast_bind_matches_sig_bind(self, sig, args, kwargs):
+        try:
+            ref = sig.bind(*args, **kwargs)
+        except TypeError:
+            # _fast_bind should raise the same error
+            with self.assertRaises(TypeError):
+                _fast_bind(sig, *args, **kwargs)
+            return
+        # Success path: compare BoundArguments
+        got = _fast_bind(sig, *args, **kwargs)
+        self.assertEqual(ref.arguments, got.arguments)
+        self.assertEqual(ref.args, got.args)
+        self.assertEqual(ref.kwargs, got.kwargs)
+
+        # Also validate default population matches
+        ref.apply_defaults()
+        got.apply_defaults()
+        self.assertEqual(ref.arguments, got.arguments)
+        self.assertEqual(ref.args, got.args)
+        self.assertEqual(ref.kwargs, got.kwargs)
+
+    def test_positional_or_keyword_and_defaults(self):
+        def f(a, b=1, c=2):
+            pass
+
+        sig = inspect.signature(f)
+        self._assert_fast_bind_matches_sig_bind(sig, (10,), {"c": 3})
+        self._assert_fast_bind_matches_sig_bind(sig, (10, 20), {})
+        self._assert_fast_bind_matches_sig_bind(sig, (), {"a": 1})
+
+    def test_missing_required(self):
+        def f(a, b):
+            pass
+
+        sig = inspect.signature(f)
+        self._assert_fast_bind_matches_sig_bind(sig, (), {})
+        self._assert_fast_bind_matches_sig_bind(sig, (1,), {})
+
+    def test_too_many_positional(self):
+        def f(a, b=1):
+            pass
+
+        sig = inspect.signature(f)
+        self._assert_fast_bind_matches_sig_bind(sig, (1, 2, 3), {})
+
+    def test_multiple_values_for_argument(self):
+        def f(a, b):
+            pass
+
+        sig = inspect.signature(f)
+        self._assert_fast_bind_matches_sig_bind(sig, (1,), {"a": 2, "b": 3})
+
+    def test_unexpected_keyword(self):
+        def f(a, b=1):
+            pass
+
+        sig = inspect.signature(f)
+        self._assert_fast_bind_matches_sig_bind(sig, (1,), {"c": 3})
+
+    def test_keyword_only(self):
+        def f(a, *, b, c=1):
+            pass
+
+        sig = inspect.signature(f)
+        self._assert_fast_bind_matches_sig_bind(sig, (1,), {"b": 2})
+        self._assert_fast_bind_matches_sig_bind(sig, (1, 2), {})
+        self._assert_fast_bind_matches_sig_bind(sig, (1,), {})
+
+    def test_positional_only(self):
+        # def f(x, /, y, *, z=1): ...
+        sig = inspect.Signature(
+            [
+                inspect.Parameter("x", inspect.Parameter.POSITIONAL_ONLY),
+                inspect.Parameter("y", inspect.Parameter.POSITIONAL_OR_KEYWORD),
+                inspect.Parameter("z", inspect.Parameter.KEYWORD_ONLY, default=1),
+            ]
+        )
+        self._assert_fast_bind_matches_sig_bind(sig, (1, 2), {"z": 3})
+        self._assert_fast_bind_matches_sig_bind(sig, (), {"x": 1, "y": 2})
+        self._assert_fast_bind_matches_sig_bind(sig, (1,), {"x": 2, "y": 3})
+
+    def test_from_keyword_positional_only_pattern(self):
+        # Mirrors a common TorchScript schema normalization pattern in operator_schemas.py
+        # where "from" is treated as positional-only.
+        sig = inspect.Signature(
+            [
+                inspect.Parameter("input", inspect.Parameter.POSITIONAL_ONLY),
+                inspect.Parameter(
+                    "from", inspect.Parameter.POSITIONAL_ONLY, default=0.0
+                ),
+                inspect.Parameter(
+                    "to", inspect.Parameter.POSITIONAL_OR_KEYWORD, default=1.0
+                ),
+                inspect.Parameter(
+                    "generator",
+                    inspect.Parameter.KEYWORD_ONLY,
+                    default=None,
+                ),
+            ]
+        )
+        self._assert_fast_bind_matches_sig_bind(sig, ("t",), {})
+        self._assert_fast_bind_matches_sig_bind(sig, ("t", 0.25), {})
+        self._assert_fast_bind_matches_sig_bind(
+            sig, ("t", 0.25, 0.75), {"generator": None}
+        )
+        self._assert_fast_bind_matches_sig_bind(sig, (), {"input": "t", "from": 0.1})
+
+    def test_varargs_and_varkw_fallback(self):
+        def f(a, *args, b=0, **kwargs):
+            pass
+
+        sig = inspect.signature(f)
+        # Extra positional should be captured by *args; extra keywords by **kwargs
+        self._assert_fast_bind_matches_sig_bind(sig, (1, 2, 3), {"b": 4, "x": 5})
+        self._assert_fast_bind_matches_sig_bind(sig, (), {"a": 1, "x": 2})
+
+    def test_zero_params(self):
+        def f():
+            pass
+
+        sig = inspect.signature(f)
+        self._assert_fast_bind_matches_sig_bind(sig, (), {})
+        self._assert_fast_bind_matches_sig_bind(sig, (1,), {})  # too many args
+        self._assert_fast_bind_matches_sig_bind(sig, (), {"x": 1})  # unexpected kwarg
+        self._assert_fast_bind_matches_sig_bind(
+            sig, (1,), {"b": 1}
+        )  # unexpected arg & kwarg
+
+    def test_mutable_defaults(self):
+        # This test case reproduces an issue where unhashable default values (like list)
+        # caused _fast_bind to fail because of lru_cache on _signature_metadata
+        def f(a, b=[]):  # noqa: B006
+            pass
+
+        sig = inspect.signature(f)
+        self._assert_fast_bind_matches_sig_bind(sig, (1,), {})
+        self._assert_fast_bind_matches_sig_bind(sig, (1, [2]), {})
+
+
+if __name__ == "__main__":
+    run_tests()
diff --git a/test/test_jit.py b/test/test_jit.py
index de35770ee6cfd..9519ed82f296b 100644
--- a/test/test_jit.py
+++ b/test/test_jit.py
@@ -103,11 +103,11 @@
 # Testing utils
 from torch.testing._internal import jit_utils
 from torch.testing._internal.common_jit import check_against_reference
-from torch.testing._internal.common_utils import run_tests, IS_WINDOWS, \
+from torch.testing._internal.common_utils import run_tests, IS_ARM64, IS_WINDOWS, \
     GRAPH_EXECUTOR, suppress_warnings, IS_SANDCASTLE, ProfilingMode, \
     TestCase, freeze_rng_state, slowTest, TemporaryFileName, \
     enable_profiling_mode_for_profiling_tests, TEST_MKL, set_default_dtype, num_profiled_runs, \
-    skipIfCrossRef, skipIfTorchDynamo
+    skipIfCrossRef, skipIfTorchDynamo, xfailIf
 from torch.testing._internal.jit_utils import JitTestCase, enable_cpu_fuser, disable_autodiff_subgraph_inlining, \
     _trace, do_input_map, get_execution_plan, make_global, \
     execWrapper, _inline_everything, _tmp_donotuse_dont_inline_everything, \
@@ -165,7 +165,8 @@ def doAutodiffCheck(testname):
     if "test_t_" in testname or testname == "test_t":
         return False
 
-    assert GRAPH_EXECUTOR
+    if not GRAPH_EXECUTOR:
+        raise AssertionError("GRAPH_EXECUTOR is not set")
     if GRAPH_EXECUTOR == ProfilingMode.SIMPLE:
         return False
 
@@ -209,7 +210,8 @@ def doAutodiffCheck(testname):
     return testname not in test_exceptions
 
 
-assert GRAPH_EXECUTOR
+if not GRAPH_EXECUTOR:
+    raise AssertionError("GRAPH_EXECUTOR is not set")
 # TODO: enable TE in PE when all tests are fixed
 torch._C._jit_set_texpr_fuser_enabled(GRAPH_EXECUTOR == ProfilingMode.PROFILING)
 torch._C._jit_set_profiling_executor(GRAPH_EXECUTOR != ProfilingMode.LEGACY)
@@ -976,7 +978,8 @@ def forward(self, x):
                 return x
 
         m = torch.jit.script(Seq())
-        assert m.graph  # ensure jit was able to compile
+        if not m.graph:
+            raise AssertionError("jit was unable to compile")  # ensure jit was able to compile
 
     def test_ModuleList(self):
         class Mod(nn.Module):
@@ -993,7 +996,8 @@ def forward(self, v):
                 return v
 
         m = torch.jit.script(Mod())
-        assert m.graph  # ensure jit was able to compile
+        if not m.graph:
+            raise AssertionError("jit was unable to compile")  # ensure jit was able to compile
 
     def test_disabled(self):
         torch.jit._state.disable()
@@ -1827,8 +1831,10 @@ def t(x, p: float, t: bool):
                         o_ref = t(x_ref, p, train)
                         o.sum().backward()
                         o_ref.sum().backward()
-                        assert o.equal(o_ref)
-                        assert x.grad.equal(x_ref.grad)
+                        if not o.equal(o_ref):
+                            raise AssertionError(f"Output mismatch: {o} vs {o_ref}")
+                        if not x.grad.equal(x_ref.grad):
+                            raise AssertionError(f"Grad mismatch: {x.grad} vs {x_ref.grad}")
 
     @slowTest
     @unittest.skipIf(GRAPH_EXECUTOR != ProfilingMode.LEGACY, 'Testing differentiable graph')
@@ -3420,7 +3426,7 @@ def __init__(self) -> None:
                 super().__init__()
                 self.foo = torch.jit.Attribute(0.1, float)
                 # we should be able to use self.foo as a float here
-                assert 0.0 < self.foo
+                assert 0.0 < self.foo  # noqa: S101
         M()
 
     def test_scriptable_fn_as_attr(self):
@@ -4314,20 +4320,20 @@ def fn2():
             a = SomeClass()
             a_copy = a
             a += 20
-            assert a is a_copy
+            assert a is a_copy  # noqa: S101
             b = SomeOutOfPlaceClass()
             b_copy = b
             b += 99
-            assert b is b_copy
+            assert b is b_copy  # noqa: S101
             c = [1, 2, 3]
             c_copy = c
             c *= 2
-            assert c is c_copy
+            assert c is c_copy  # noqa: S101
             c += [4, 5, 6]
             d = torch.ones(2, 2)
             d_copy = d
             d += torch.ones(2, 2)
-            assert d is d_copy
+            assert d is d_copy  # noqa: S101
             return a, b, c, d
 
         self.checkScript(fn2, [])
@@ -5872,7 +5878,7 @@ def fn(x, y, z):
                 m = x if not z else y
             while x < y > z:
                 q = x
-            assert 1 == 1, "hello"
+            assert 1 == 1, "hello"  # noqa: S101
             return x
 
         ast = torch.jit.frontend.get_jit_def(fn, fn.__name__)
@@ -6265,7 +6271,7 @@ def test_assertion_optional_refinement(self):
         @torch.jit.script
         def test(x, y):
             # type: (Optional[int], Optional[int]) -> int
-            assert x is not None and y is not None
+            assert x is not None and y is not None  # noqa: S101
             return x + y
 
         self.assertEqual(test(2, 2), 4)
@@ -7654,7 +7660,7 @@ def test_unspecialized_any_binding(self):
 
         @torch.jit.script
         def foo(x: Any):
-            assert isinstance(x, Dict[str, torch.Tensor])
+            assert isinstance(x, Dict[str, torch.Tensor])  # noqa: S101
 
         foo({"1": torch.tensor(3)})
         with self.assertRaises(Exception):
@@ -7899,13 +7905,13 @@ def assign_after_break_nested(y):
                 if y == 1:
                     x = 5
                     break
-                    assert 1 == 2
+                    assert 1 == 2  # noqa: S101
                 else:
                     x = x + 1
                     break
-                    assert 1 == 2
+                    assert 1 == 2  # noqa: S101
                 x = -30
-                assert 1 == 2
+                assert 1 == 2  # noqa: S101
             return x
 
         self.checkScript(assign_after_break_nested, (1,))
@@ -7976,7 +7982,7 @@ def test_will_break_after_guard(x):
                     break
                 a -= 1
                 break
-                assert 1 == 2
+                assert 1 == 2  # noqa: S101
                 a -= -100
             return a
 
@@ -8274,9 +8280,11 @@ def test_parse_scalar_tensor_constants(self):
                 for n in jit_graph.nodes()
                 if isinstance(n.output().type(), torch.TensorType)
             )
-            assert isinstance(node.output().type(), torch.TensorType)
+            if not isinstance(node.output().type(), torch.TensorType):
+                raise AssertionError(f"Expected TensorType, got {type(node.output().type())}")
             t = node.t("value")
-            assert isinstance(t, torch.Tensor)
+            if not isinstance(t, torch.Tensor):
+                raise AssertionError(f"Expected torch.Tensor, got {type(t)}")
             self.assertEqual(t.dtype, dtype)
             self.assertEqual(t.item(), value)
 
@@ -11521,7 +11529,8 @@ def test_empty_tuple_str(self):
         empty_tuple_type = torch._C.TupleType([])
         g = {'Tuple' : typing.Tuple}
         python_type = eval(empty_tuple_type.annotation_str, g)
-        assert python_type is typing.Tuple[()]
+        if python_type is not typing.Tuple[()]:
+            raise AssertionError(f"Expected Tuple[()], got {python_type}")
 
     def test_tuple_str(self):
         tuple1_type = torch._C.TupleType([torch._C.StringType.get()])
@@ -11537,7 +11546,8 @@ def test_none_type_str(self):
         none_type = torch._C.NoneType.get()
         g = {'NoneType' : type(None)}
         python_type = eval(none_type.annotation_str, g)
-        assert python_type is type(None)
+        if python_type is not type(None):
+            raise AssertionError(f"Expected NoneType, got {python_type}")
 
     @skipIfTorchDynamo("TorchDynamo fails with unknown reason")
     def test_zip_enumerate_modulelist(self):
@@ -12831,7 +12841,8 @@ def test_file_format_serialization(self):
 
         for i, offset in enumerate(parsed_serialized_offsets):
             data = reader.get_record(str(offset))
-            assert data == buffers[i]
+            if data != buffers[i]:
+                raise AssertionError(f"Record mismatch at offset {offset}: {data} != {buffers[i]}")
 
     def test_file_reader_no_memory_leak(self):
         num_iters = 10000
@@ -13453,9 +13464,9 @@ def getFooTest(self):
         def test_id_class_types():
             obj1 = FooTest(torch.tensor(3))
             obj2 = FooTest(torch.tensor(2))
-            assert obj1 is not obj2
-            assert id(obj1) != id(obj2)
-            assert id(obj1) != id(None)
+            assert obj1 is not obj2  # noqa: S101
+            assert id(obj1) != id(obj2)  # noqa: S101
+            assert id(obj1) != id(None)  # noqa: S101
             return True
 
         self.assertTrue(test_id_class_types())
@@ -14037,11 +14048,11 @@ def complicated(x):
             if x:
                 if x == 2:
                     return 1
-                    assert 1 == 2
+                    assert 1 == 2  # noqa: S101
                 else:
                     if x == 3:
                         return 2
-                        assert 1 == 2
+                        assert 1 == 2  # noqa: S101
                     else:
                         a = 2
                         b = 3
@@ -14049,7 +14060,7 @@ def complicated(x):
                 a = 4
                 b = 1
             return a + b
-            assert 1 == 2
+            assert 1 == 2  # noqa: S101
 
         for i in range(4):
             self.checkScript(complicated, (i,))
@@ -14222,9 +14233,11 @@ def _test_parse_generator(seed):
                 for n in jit_graph.nodes()
                 if isinstance(n.output().type(), torch._C._GeneratorType)
             )
-            assert isinstance(node.output().type(), torch._C._GeneratorType)
+            if not isinstance(node.output().type(), torch._C._GeneratorType):
+                raise AssertionError(f"Expected _GeneratorType, got {type(node.output().type())}")
             g = node.ival("value")
-            assert isinstance(g, torch.Generator)
+            if not isinstance(g, torch.Generator):
+                raise AssertionError(f"Expected torch.Generator, got {type(g)}")
             self.assertEqual(g.initial_seed(), seed)
 
         _test_parse_generator(2024)
@@ -14312,7 +14325,7 @@ def foo():
 
     def test_assert_is_scripting_metacompile(self):
         def foo():
-            assert not torch.jit.is_scripting(), "TestErrorMsg"
+            assert not torch.jit.is_scripting(), "TestErrorMsg"  # noqa: S101
             print("hello") + 2  # will not be compiled
 
         f = torch.jit.script(foo)
@@ -15783,6 +15796,8 @@ def split_two(tensor):
         y = torch.randn(3, 6)
         self.checkScript(split_two, [(x + y)])
 
+    @xfailIf(IS_ARM64)
+    # see https://github.com/pytorch/pytorch/issues/177255
     def test_conv_error(self):
         @torch.jit.script
         def fn(x, y):
@@ -16230,7 +16245,8 @@ def create_nn_module(*args, **kwargs):
 
 
 def post_add_test(test_name, skipTestIf, do_test, test_class):
-    assert not hasattr(test_class, test_name), 'Two tests have the same name: ' + test_name
+    if hasattr(test_class, test_name):
+        raise AssertionError('Two tests have the same name: ' + test_name)
 
     for skip in skipTestIf:
         do_test = skip(do_test)
diff --git a/test/test_jit_autocast.py b/test/test_jit_autocast.py
index b3f47ed5e7eb6..076d71a0f9901 100644
--- a/test/test_jit_autocast.py
+++ b/test/test_jit_autocast.py
@@ -2,12 +2,19 @@
 
 import torch
 from torch.cuda.amp import autocast
-from typing import Optional
 
 import sys
 import unittest
 from torch.testing._internal.common_cuda import TEST_CUDA
-from torch.testing._internal.common_utils import parse_cmd_line_args, run_tests, skipIfTorchDynamo
+from torch.testing._internal.common_utils import (
+    IS_ARM64,
+    IS_LINUX,
+    IS_CPU_EXT_SVE_SUPPORTED,
+    parse_cmd_line_args,
+    run_tests,
+    skipIfTorchDynamo,
+    xfailIf,
+)
 from torch.testing import FileCheck
 from jit.test_models import MnistNet
 
@@ -188,7 +195,7 @@ def fn(a, b):
     @unittest.skipIf(not TEST_CUDA, "No cuda")
     def test_fp32_set_opt_dtype_policy(self):
         @torch.jit.script
-        def fn(a, b, c, d, dtype: Optional[int]):
+        def fn(a, b, c, d, dtype: int | None):
             with autocast(enabled=True):
                 x = torch.softmax(a, 0)
                 y = torch.softmax(b, 0, None)
@@ -204,7 +211,7 @@ def fn(a, b, c, d, dtype: Optional[int]):
     @unittest.skipIf(not TEST_CUDA, "No cuda")
     def test_fp32_set_opt_dtype_policy_fp64(self):
         @torch.jit.script
-        def fn(a, b, c, d, dtype: Optional[int]):
+        def fn(a, b, c, d, dtype: int | None):
             with autocast(enabled=True):
                 x = torch.softmax(a, 0)
                 y = torch.softmax(b, 0, None)
@@ -809,6 +816,8 @@ def test_generate_autocast_jit_trace_model(model, x):
         for i in range(self.models.__len__()):
             test_generate_autocast_jit_trace_model(self.models[i], self.inputs[i])
 
+    @xfailIf(IS_ARM64 and IS_LINUX and not IS_CPU_EXT_SVE_SUPPORTED)
+    # see https://github.com/pytorch/pytorch/issues/177247
     def test_nchw_autocast_jit_trace_model(self):
         def test_nchw_autocast_jit_trace_model(model, x):
             model.eval()
@@ -823,6 +832,8 @@ def test_nchw_autocast_jit_trace_model(model, x):
         for i in range(self.models.__len__()):
             test_nchw_autocast_jit_trace_model(self.models[i], self.inputs[i])
 
+    @xfailIf(IS_ARM64 and IS_LINUX and not IS_CPU_EXT_SVE_SUPPORTED)
+    # see https://github.com/pytorch/pytorch/issues/177247
     def test_nhwc_autocast_jit_trace_model(self):
         def test_nhwc_autocast_jit_trace_model(model, x):
             model = model.to(memory_format=torch.channels_last)
diff --git a/test/test_jit_fuser_te.py b/test/test_jit_fuser_te.py
index c99c6ee2f4ce3..af6532d3ce1da 100644
--- a/test/test_jit_fuser_te.py
+++ b/test/test_jit_fuser_te.py
@@ -61,7 +61,6 @@
     skipIfTorchDynamo,
     slowTest,
     TEST_WITH_ASAN,
-    TEST_WITH_ROCM,
 )
 from torch.testing._internal.jit_metaprogramming_utils import create_traced_fn
 from torch.testing._internal.jit_utils import (
@@ -2406,32 +2405,34 @@ def eager(x):
 
     @skipIfTorchDynamo("too slow")
     @unittest.skipIf(TEST_WITH_ASAN, "takes 10+ minutes on asan")
-    @unittest.skipIf(TEST_WITH_ROCM, "Tensor-likes are not close for nans")
     def test_batch_norm(self):
         def test(fn, args):
             trace = torch.jit.trace(fn, args)
             self.assertAllFused(trace.graph_for(*args))
-            # TODO: Are `NaN`'s actually ok here or did this pass silently before, because `equal_nan=True` was the
-            #  default?
             torch.testing.assert_close(fn(*args), trace(*args), equal_nan=True)
 
-        def bn(i, x):
-            return torch.batch_norm(i, x, x, x, x, False, 0.1, 1e-4, False).relu()
+        def bn(i, x, rv):
+            return torch.batch_norm(i, x, x, x, rv, False, 0.1, 1e-4, False).relu()
 
-        def bn_no_weight(i, x):
-            return torch.batch_norm(i, None, x, x, x, False, 0.1, 1e-4, False).relu()
+        def bn_no_weight(i, x, rv):
+            return torch.batch_norm(i, None, x, x, rv, False, 0.1, 1e-4, False).relu()
 
-        def bn_no_bias(i, x):
-            return torch.batch_norm(i, x, None, x, x, False, 0.1, 1e-4, False).relu()
+        def bn_no_bias(i, x, rv):
+            return torch.batch_norm(i, x, None, x, rv, False, 0.1, 1e-4, False).relu()
 
-        def bn_neither(i, x):
-            return torch.batch_norm(i, None, None, x, x, False, 0.1, 1e-4, False).relu()
+        def bn_neither(i, x, rv):
+            return torch.batch_norm(
+                i, None, None, x, rv, False, 0.1, 1e-4, False
+            ).relu()
 
         for device in self.devices:
             i = torch.randn(4, 16, 32, 40, device=device)
             x = torch.randn(16, device=device)
+            rv = torch.randn(
+                16, device=device
+            ).abs()  # running_var must be non-negative
             for fn in [bn, bn_no_weight, bn_no_bias, bn_neither]:
-                test(fn, (i, x))
+                test(fn, (i, x, rv))
 
     def test_profiler(self):
         @torch.jit.script
diff --git a/test/test_linalg.py b/test/test_linalg.py
index 9a84211dde6c7..c720b52fcecb3 100644
--- a/test/test_linalg.py
+++ b/test/test_linalg.py
@@ -15,7 +15,6 @@
 from random import randrange
 from itertools import product
 from functools import reduce, partial
-from typing import Union, Optional
 from torch._prims_common import DimsType
 from packaging import version
 from torch.testing._internal.common_device_type import (
@@ -32,6 +31,7 @@
      runOnRocmArch, MI200_ARCH, MI300_ARCH, MI350_ARCH, NAVI_ARCH, TEST_CUDA)
 from torch.testing._internal.common_device_type import \
     (instantiate_device_type_tests, dtypes, has_cusolver, onlyCPU, skipIf, skipCUDAIfNoMagma, skipCPUIfNoLapack, precisionOverride,
+     skipCUDAIf,
      skipCUDAIfNoCusolver, skipCUDAIfNoMagmaAndNoCusolver, skipCUDAIfRocm, onlyNativeDeviceTypes, dtypesIfCUDA,
      onlyCUDA, skipMeta, skipCUDAIfNotRocm, dtypesIfMPS, largeTensorTest)
 from torch.testing import make_tensor
@@ -39,8 +39,8 @@
     all_types, all_types_and_complex_and, floating_and_complex_types, integral_types,
     floating_and_complex_types_and, floating_types_and, complex_types,
 )
-from torch.testing._internal.common_cuda import CDNA2OrLater, SM80OrLater, SM90OrLater, tf32_on_and_off, _get_magma_version, \
-    _get_torch_cuda_version, TEST_MULTIGPU
+from torch.testing._internal.common_cuda import CDNA2OrLater, CDNA5OrLater, SM80OrLater, SM90OrLater, tf32_on_and_off, _get_magma_version, \
+    _get_torch_cuda_version, TEST_MULTIGPU, PLATFORM_SUPPORTS_FP8, blas_library_context
 from torch.testing._internal.common_quantization import _group_quantize_tensor, _dynamically_quantize_per_channel, \
     _group_quantize_tensor_symmetric
 from torch.testing._internal.common_mkldnn import reduced_f32_on_and_off
@@ -65,6 +65,8 @@ def blaslt_supported_device():
                 archs.extend(['gfx110', 'gfx120'])
             if ROCM_VERSION >= (6, 5):
                 archs.append('gfx95')
+            if ROCM_VERSION >= (7, 14):
+                archs.append('gfx1250')
             for arch in archs:
                 if arch in torch.cuda.get_device_properties(0).gcnArchName:
                     return True
@@ -659,6 +661,22 @@ def run_test(shape, batch, contiguous):
         actual = torch.linalg.cholesky(A, upper=True)
         self.assertEqual(expected, actual)
 
+    @skipCUDAIfNoCusolver
+    @skipCPUIfNoLapack
+    @dtypes(*floating_and_complex_types())
+    def test_cholesky_upper_reconstructs(self, device, dtype):
+        batch_dims = (1,)
+        matrix_size = 65
+        A = torch.randn(
+            *(batch_dims + (matrix_size, matrix_size)), dtype=dtype, device=device
+        )
+        pd_matrix = A @ A.mH + torch.eye(matrix_size, dtype=dtype, device=device)
+        pd_matrix = pd_matrix.squeeze(0)
+        U = torch.linalg.cholesky(pd_matrix, upper=True)
+        self.assertEqual(U, torch.triu(U))
+        reconstructed = U.mH @ U
+        self.assertEqual(pd_matrix, reconstructed, atol=1e-4, rtol=1e-5)
+
     @skipCUDAIfNoCusolver
     @skipCPUIfNoLapack
     @dtypes(*floating_and_complex_types())
@@ -1042,6 +1060,24 @@ def test_det(self, device, dtype):
         with self.assertRaises(RuntimeError):
             op(t)
 
+    @skipCUDAIfNoCusolver
+    @skipCPUIfNoLapack
+    @dtypes(torch.double, torch.cdouble)
+    def test_det_backward(self, device, dtype):
+        # Regression test for #80761.
+        input = torch.tensor([[0.]], device=device, dtype=dtype, requires_grad=True)
+        self.assertTrue(torch.autograd.gradcheck(torch.det, inputs=input))
+
+        # When A has 0 elements (e.g. empty batch), backward should return a
+        # zeros tensor with the same shape as A, not an undefined tensor.
+        for shape in [(0, 3, 3), (2, 0, 0)]:
+            A = torch.randn(shape, device=device, dtype=dtype, requires_grad=True)
+            det = torch.linalg.det(A)
+            det.backward(torch.ones_like(det))
+            self.assertIsNotNone(A.grad)
+            self.assertEqual(A.grad.shape, A.shape)
+            self.assertEqual(A.grad, torch.zeros_like(A))
+
     @skipCUDAIfNoMagma
     @skipCPUIfNoLapack
     @dtypes(*floating_and_complex_types())
@@ -1079,7 +1115,7 @@ def run_test(shape, batch, uplo):
         for shape, batch, uplo in itertools.product(shapes, batches, uplos):
             run_test(shape, batch, uplo)
 
-    @skipCUDAIfNoMagma
+    @skipCUDAIfNoCusolver
     @skipCPUIfNoLapack
     @dtypes(*floating_and_complex_types())
     @precisionOverride({torch.float32: 1e-4, torch.complex64: 1e-4})
@@ -1097,7 +1133,7 @@ def run_test(shape, batch, uplo):
         for uplo in uplos:
             run_test(3, (2, 2), uplo)
 
-    @skipCUDAIfNoMagma
+    @skipCUDAIfNoCusolver
     @skipCPUIfNoLapack
     @dtypes(*floating_and_complex_types())
     def test_eigh_errors_and_warnings(self, device, dtype):
@@ -1172,7 +1208,7 @@ def test_eigh_svd_illcondition_matrix_input_should_not_crash(self, device, dtype
         self.assertEqual(eigh_out.eigenvalues.sort(descending=True).values[:2], [1.0e5, 511.0], atol=1.0, rtol=1.0e-2)
         self.assertEqual(svd_out.S[:2], [1.0e5, 511.0], atol=1.0, rtol=1.0e-2)
 
-    @skipCUDAIfNoMagma
+    @skipCUDAIfNoCusolver
     @skipCPUIfNoLapack
     @dtypes(*floating_and_complex_types())
     @precisionOverride({torch.float32: 1e-4, torch.complex64: 1e-4})
@@ -1197,7 +1233,7 @@ def run_test(shape, batch, uplo):
         for shape, batch, uplo in itertools.product(shapes, batches, uplos):
             run_test(shape, batch, uplo)
 
-    @skipCUDAIfNoMagma
+    @skipCUDAIfNoCusolver
     @skipCPUIfNoLapack
     @dtypes(*floating_and_complex_types())
     def test_eigvalsh_errors_and_warnings(self, device, dtype):
@@ -1477,8 +1513,8 @@ def forward(self, a):
                 return tensor_unbacked_size
 
         def test(
-            ord: Union[float, int],
-            dim: Optional[DimsType],
+            ord: float | int,
+            dim: DimsType | None,
             expect_numel_runtime_check: bool,
             expect_index_0_check: bool = False,
         ) -> None:
@@ -2258,7 +2294,11 @@ def test_norm_fastpaths(self, device):
         self.assertEqual(result, expected)
 
     @skipCPUIfNoLapack
-    @skipCUDAIfNoMagma
+    @skipCUDAIf(
+        not TEST_WITH_ROCM and _get_torch_cuda_version() < (12, 8) and not torch.cuda.has_magma,
+        "torch.linalg.eig requires MAGMA for CUDA versions < 12.8",
+    )
+    @skipCUDAIf(TEST_WITH_ROCM and not torch.cuda.has_magma, "ROCm hipsolver backend does not currently support eig")
     # NumPy computes only in float64 and complex128 precisions
     # for float32 or complex64 results might be very different from float64 or complex128
     @dtypes(torch.float64, torch.complex128)
@@ -2307,7 +2347,11 @@ def run_test(shape, *, symmetric=False):
             run_test(shape, symmetric=True)
 
     @onlyCUDA
-    @skipCUDAIfNoMagma
+    @skipCUDAIf(
+        not TEST_WITH_ROCM and _get_torch_cuda_version() < (12, 8) and not torch.cuda.has_magma,
+        "torch.linalg.eig requires MAGMA for CUDA versions < 12.8",
+    )
+    @skipCUDAIf(TEST_WITH_ROCM and not torch.cuda.has_magma, "ROCm hipsolver backend does not currently support eig")
     @dtypes(*floating_and_complex_types())
     def test_eig_identity(self, device, dtype):
 
@@ -2322,14 +2366,6 @@ def run_test(shape, *, symmetric=False):
 
             actual = torch.linalg.eig(a)
 
-            complementary_device = 'cpu'
-
-            # compare eigenvalues with CPU
-            expected = torch.linalg.eig(a.to(complementary_device))
-
-            # make sure all eigenvalues are returned
-            self.assertEqual(expected[0].shape, actual[0].shape)
-
             # set tolerance for correctness check
             if dtype in [torch.float32, torch.complex64]:
                 atol = 1e-3  # CuSolver gives less accurate results for single precision (1-2 larger than OOM NumPy)
@@ -2375,7 +2411,11 @@ def run_test(shape, *, symmetric=False):
 
 
     @onlyCUDA
-    @skipCUDAIfNoMagmaAndNoCusolver
+    @skipCUDAIf(
+        not TEST_WITH_ROCM and _get_torch_cuda_version() < (12, 8) and not torch.cuda.has_magma,
+        "torch.linalg.eig requires MAGMA for CUDA versions < 12.8",
+    )
+    @skipCUDAIf(TEST_WITH_ROCM and not torch.cuda.has_magma, "ROCm hipsolver backend does not currently support eig")
     @dtypes(*floating_and_complex_types())
     def test_eigvals_out_variants(self, device, dtype):
         from torch.testing._internal.common_utils import random_symmetric_matrix
@@ -2417,7 +2457,11 @@ def run_test(shape, *, symmetric=False):
 
 
     @onlyCUDA
-    @skipCUDAIfNoMagmaAndNoCusolver
+    @skipCUDAIf(
+        not TEST_WITH_ROCM and _get_torch_cuda_version() < (12, 8) and not torch.cuda.has_magma,
+        "torch.linalg.eig requires MAGMA for CUDA versions < 12.8",
+    )
+    @skipCUDAIf(TEST_WITH_ROCM and not torch.cuda.has_magma, "ROCm hipsolver backend does not currently support eig")
     @dtypes(*floating_and_complex_types())
     def test_eig_out_variants(self, device, dtype):
         from torch.testing._internal.common_utils import random_symmetric_matrix
@@ -2469,20 +2513,12 @@ def run_test(shape, *, symmetric=False):
             run_test(shape, symmetric=True)
 
 
-    @slowTest
     @onlyCUDA
-    @skipCUDAIfNoMagma
-    @dtypes(torch.float32)
-    def test_eig_check_magma(self, device, dtype):
-        # For CUDA inputs only matrices of size larger than 2048x2048 actually call MAGMA library
-        shape = (2049, 2049)
-        a = make_tensor(shape, dtype=dtype, device=device)
-        w, v = torch.linalg.eig(a)
-        # check correctness using eigendecomposition identity
-        self.assertEqual(a.to(v.dtype) @ v, w * v, atol=1e-3, rtol=1e-3)
-
-    @onlyCUDA
-    @skipIf(TEST_WITH_ROCM and not torch.cuda.has_magma, "ROCm hipsolver backend does not currently support eig")
+    @skipCUDAIf(
+        not TEST_WITH_ROCM and _get_torch_cuda_version() < (12, 8) and not torch.cuda.has_magma,
+        "torch.linalg.eig requires MAGMA for CUDA versions < 12.8",
+    )
+    @skipCUDAIf(TEST_WITH_ROCM and not torch.cuda.has_magma, "ROCm hipsolver backend does not currently support eig")
     @dtypes(torch.float32, torch.float64)
     def test_eig_cuda_complex_eigenvectors(self, device, dtype):
         """Test CUDA eigenvector decoding with known ground truth, including batching."""
@@ -2567,8 +2603,12 @@ def test_eig_cuda_complex_eigenvectors(self, device, dtype):
         rhs = vals_batch.unsqueeze(-2) * vecs_batch
         self.assertEqual(lhs, rhs, atol=1e-5, rtol=1e-5)
 
-    @skipCUDAIfNoMagma
     @skipCPUIfNoLapack
+    @skipCUDAIf(
+        not TEST_WITH_ROCM and _get_torch_cuda_version() < (12, 8) and not torch.cuda.has_magma,
+        "torch.linalg.eig requires MAGMA for CUDA versions < 12.8",
+    )
+    @skipCUDAIf(TEST_WITH_ROCM and not torch.cuda.has_magma, "ROCm hipsolver backend does not currently support eig")
     @dtypes(*floating_and_complex_types())
     def test_eig_errors_and_warnings(self, device, dtype):
         # eig requires the input to be at least 2 dimensional tensor
@@ -2630,7 +2670,11 @@ def test_eig_errors_and_warnings(self, device, dtype):
                 torch.linalg.eig(a, out=(out_w, out_v))
 
     @skipCPUIfNoLapack
-    @skipCUDAIfNoMagma
+    @skipCUDAIf(
+        not TEST_WITH_ROCM and _get_torch_cuda_version() < (12, 8) and not torch.cuda.has_magma,
+        "torch.linalg.eig requires MAGMA for CUDA versions < 12.8",
+    )
+    @skipCUDAIf(TEST_WITH_ROCM and not torch.cuda.has_magma, "ROCm hipsolver backend does not currently support eig")
     @dtypes(*floating_and_complex_types())
     def test_eig_with_nan(self, device, dtype):
         for val in [np.inf, np.nan]:
@@ -2642,7 +2686,11 @@ def test_eig_with_nan(self, device, dtype):
                     torch.linalg.eig(a)
 
     @skipCPUIfNoLapack
-    @skipCUDAIfNoMagma
+    @skipCUDAIf(
+        not TEST_WITH_ROCM and _get_torch_cuda_version() < (12, 8) and not torch.cuda.has_magma,
+        "torch.linalg.eig requires MAGMA for CUDA versions < 12.8",
+    )
+    @skipCUDAIf(TEST_WITH_ROCM and not torch.cuda.has_magma, "ROCm hipsolver backend does not currently support eig")
     # NumPy computes only in float64 and complex128 precisions
     # for float32 or complex64 results might be very different from float64 or complex128
     @dtypes(torch.float64, torch.complex128)
@@ -2688,8 +2736,12 @@ def run_test(shape, *, symmetric=False):
             run_test(shape, symmetric=True)
 
 
-    @skipCUDAIfNoMagma
     @skipCPUIfNoLapack
+    @skipCUDAIf(
+        not TEST_WITH_ROCM and _get_torch_cuda_version() < (12, 8) and not torch.cuda.has_magma,
+        "torch.linalg.eig requires MAGMA for CUDA versions < 12.8",
+    )
+    @skipCUDAIf(TEST_WITH_ROCM and not torch.cuda.has_magma, "ROCm hipsolver backend does not currently support eig")
     @dtypes(*floating_and_complex_types())
     def test_eigvals_errors_and_warnings(self, device, dtype):
         # eig requires the input to be at least 2 dimensional tensor
@@ -3116,8 +3168,12 @@ def mul_svd_factors(U, S, Vh):
                 S_s = torch.svd(A, compute_uv=False).S
                 self.assertEqual(S_s, S)
 
-    @skipCUDAIfNoMagmaAndNoCusolver
     @skipCPUIfNoLapack
+    @skipCUDAIf(
+        not TEST_WITH_ROCM and _get_torch_cuda_version() < (12, 8) and not torch.cuda.has_magma,
+        "torch.linalg.eig requires MAGMA for CUDA versions < 12.8",
+    )
+    @skipCUDAIf(TEST_WITH_ROCM and not torch.cuda.has_magma, "ROCm hipsolver backend does not currently support eig")
     @dtypes(torch.complex128)
     def test_invariance_error_spectral_decompositions(self, device, dtype):
         make_arg = partial(make_tensor, device=device, dtype=dtype, requires_grad=True)
@@ -3160,41 +3216,45 @@ def cholesky_solve_test_helper(self, A_dims, b_dims, upper, device, dtype):
         b = torch.randn(*b_dims, dtype=dtype, device=device)
         A = random_hermitian_pd_matrix(*A_dims, dtype=dtype, device=device)
         L = torch.cholesky(A, upper=upper)
-        return b, A, L
 
-    @skipCUDAIfNoMagma
+        # L should be col-major
+        self.assertTrue(L.mT.is_contiguous())
+
+        yield b, A, L  # L is col-major
+        yield b, A, L.contiguous()  # L is row-major
+
+    @skipCUDAIfNoCusolver
     @skipCPUIfNoLapack
     @dtypes(*floating_and_complex_types())
     @precisionOverride({torch.float32: 1e-3, torch.complex64: 1e-3,
                         torch.float64: 1e-8, torch.complex128: 1e-8})
     def test_cholesky_solve(self, device, dtype):
         for (k, n), upper in itertools.product(zip([2, 3, 5], [3, 5, 7]), [True, False]):
-            b, A, L = self.cholesky_solve_test_helper((n,), (n, k), upper, device, dtype)
-            x = torch.cholesky_solve(b, L, upper=upper)
-            self.assertEqual(b, np.matmul(A.cpu(), x.cpu()))
+            for b, A, L in self.cholesky_solve_test_helper((n,), (n, k), upper, device, dtype):
+                x = torch.cholesky_solve(b, L, upper=upper)
+                self.assertEqual(b, np.matmul(A.cpu(), x.cpu()))
 
-    @skipCUDAIfNoMagma
+    @skipCUDAIfNoCusolver
     @skipCPUIfNoLapack
     @dtypes(*floating_and_complex_types())
     @precisionOverride({torch.float32: 1e-3, torch.complex64: 1e-3,
                         torch.float64: 1e-8, torch.complex128: 1e-8})
     def test_cholesky_solve_batched(self, device, dtype):
         def cholesky_solve_batch_helper(A_dims, b_dims, upper):
-            b, A, L = self.cholesky_solve_test_helper(A_dims, b_dims, upper, device, dtype)
-            x_exp_list = []
-            for i in range(b_dims[0]):
-                x_exp_list.append(torch.cholesky_solve(b[i], L[i], upper=upper))
-            x_exp = torch.stack(x_exp_list)  # Stacked output
-            x_act = torch.cholesky_solve(b, L, upper=upper)  # Actual output
-            self.assertEqual(x_act, x_exp)  # Equality check
-            Ax = np.matmul(A.cpu(), x_act.cpu())
-            self.assertEqual(b, Ax)  # Correctness check
+            for b, A, L in self.cholesky_solve_test_helper(A_dims, b_dims, upper, device, dtype):
+                x_exp_list = []
+                for i in range(b_dims[0]):
+                    x_exp_list.append(torch.cholesky_solve(b[i], L[i], upper=upper))
+                x_exp = torch.stack(x_exp_list)  # Stacked output
+                x_act = torch.cholesky_solve(b, L, upper=upper)  # Actual output
+                self.assertEqual(x_act, x_exp)  # Equality check
+                Ax = np.matmul(A.cpu(), x_act.cpu())
+                self.assertEqual(b, Ax)  # Correctness check
 
         for upper, batchsize in itertools.product([True, False], [1, 3, 4]):
             cholesky_solve_batch_helper((5, batchsize), (batchsize, 5, 10), upper)
 
-    @slowTest
-    @skipCUDAIfNoMagma
+    @skipCUDAIfNoCusolver
     @skipCPUIfNoLapack
     @dtypes(*floating_and_complex_types())
     @precisionOverride({torch.float32: 1e-3, torch.complex64: 1e-3,
@@ -3202,12 +3262,12 @@ def cholesky_solve_batch_helper(A_dims, b_dims, upper):
     def test_cholesky_solve_batched_many_batches(self, device, dtype):
         for A_dims, b_dims in zip([(5, 256, 256), (5,)], [(5, 10), (512, 512, 5, 10)]):
             for upper in [True, False]:
-                b, A, L = self.cholesky_solve_test_helper(A_dims, b_dims, upper, device, dtype)
-                x = torch.cholesky_solve(b, L, upper)
-                Ax = torch.matmul(A, x)
-                self.assertEqual(Ax, b.expand_as(Ax))
+                for b, A, L in self.cholesky_solve_test_helper(A_dims, b_dims, upper, device, dtype):
+                    x = torch.cholesky_solve(b, L, upper)
+                    Ax = torch.matmul(A, x)
+                    self.assertEqual(Ax, b.expand_as(Ax))
 
-    @skipCUDAIfNoMagmaAndNoCusolver
+    @skipCUDAIfNoCusolver
     @skipCPUIfNoLapack
     @dtypes(*floating_and_complex_types())
     @precisionOverride({torch.float32: 1e-3, torch.complex64: 1e-3,
@@ -3238,7 +3298,7 @@ def run_test(A_dims, b_dims, upper):
             run_test((4, 4), (2, 1, 3, 4, 2), upper)  # broadcasting A
             run_test((1, 3, 1, 4, 4), (2, 1, 3, 4, 5), upper)  # broadcasting A & b
 
-    @skipCUDAIfNoMagma
+    @skipCUDAIfNoCusolver
     @skipCPUIfNoLapack
     @dtypes(*floating_and_complex_types())
     def test_cholesky_solve_out_errors_and_warnings(self, device, dtype):
@@ -3265,7 +3325,7 @@ def test_cholesky_solve_out_errors_and_warnings(self, device, dtype):
             self.assertEqual(len(w), 1)
             self.assertTrue("An output with one or more elements was resized" in str(w[-1].message))
 
-    @skipCUDAIfNoMagma
+    @skipCUDAIfNoCusolver
     @skipCPUIfNoLapack
     @dtypes(torch.double)
     def test_cholesky_solve_backward(self, device, dtype):
@@ -3750,11 +3810,14 @@ def test_tensorsolve_errors_and_warnings(self, device, dtype):
     @skipCUDAIfNoMagma
     @skipCPUIfNoLapack
     @dtypes(*floating_and_complex_types())
-    @precisionOverride({torch.float: 1e-3, torch.cfloat: 1e-3})
     def test_tensorinv(self, device, dtype):
+        make_fullrank = make_fullrank_matrices_with_distinct_singular_values
 
         def run_test(a_shape, ind):
-            a = torch.randn(a_shape, dtype=dtype, device=device)
+            n = 1
+            for s in a_shape[:ind]:
+                n *= s
+            a = make_fullrank(n, n, dtype=dtype, device=device).reshape(a_shape)
             a_numpy = a.cpu().numpy()
             result = torch.linalg.tensorinv(a, ind=ind)
             expected = np.linalg.tensorinv(a_numpy, ind=ind)
@@ -3767,7 +3830,15 @@ def run_test(a_shape, ind):
             self.assertEqual(ans, result)
 
         # compare to NumPy output
-        run_test((12, 3, 4), ind=1)
+        if not torch.version.hip:
+            # https://github.com/pytorch/pytorch/issues/174913
+            # Skip one config due to regression on ROCm 7.2 for hipSolver.
+            # Rather than skip entire unit test using @skipIfRocm
+            # This happened on MI355, MI300, and MI200.
+            # Mismatched elements: 1 / 144 (0.7%)
+            # Greatest absolute difference: 0.00130462646484375 at index (1, 3, 6) (up to 0.001 allowed)
+            # Greatest relative difference: 1.5133813576539978e-05 at index (1, 3, 6) (up to 1.3e-06 allowed)
+            run_test((12, 3, 4), ind=1)
         run_test((3, 8, 24), ind=2)
         run_test((18, 3, 3, 2), ind=1)
         run_test((1, 4, 2, 2), ind=2)
@@ -5108,9 +5179,6 @@ def test_matmul_small_brute_force_tunableop(self, device, dtype):
         # We set the TunableOp numerical check environment variable here because it is
         # possible to hit some invalid numerical solutions due to the small matrix sizes.
 
-        if torch.version.hip and isRocmArchAnyOf(MI350_ARCH) and dtype is torch.double:
-            self.skipTest("Currently hangs on rocm mi350")
-
         with self._tunableop_ctx():
             torch.cuda.tunable.set_rotating_buffer_size(0)
             # Numerical check adds significant overhead, unsure if this is needed
@@ -6080,6 +6148,7 @@ def test_blaslog_tunableop(self, device, dtype):
 
     @onlyCUDA
     @skipCUDAIfNotRocm
+    # Fails with triton 3.7
     @dtypes(torch.float)
     def test_mm_submatrix_offline_tunableop(self, device, dtype):
         import os
@@ -6210,8 +6279,12 @@ def test_mm_submatrix_offline_tunableop(self, device, dtype):
             # This stores total number of cumulative results
             total_num_results = new_results - ref_results
 
-            # There must be a new tuning results
-            self.assertEqual(total_num_results, 10)
+            preferred_blas = str(torch.backends.cuda.preferred_blas_library())
+            # With hipBLASLt preferred, the two linear/addmm+bias calls are
+            # tracked as GemmAndBias tunables (+2). With rocBLAS preferred,
+            # they fall back to regular GEMM signatures and don't add entries.
+            expected_num_results = 10 if preferred_blas == "_BlasBackend.Cublaslt" else 8
+            self.assertEqual(total_num_results, expected_num_results)
 
             results_filename = torch.cuda.tunable.get_filename()
             self.assertTrue(os.path.exists(results_filename))
@@ -6422,9 +6495,9 @@ def test_call_count_tunableop(self, device, dtype):
             # launched per PyTorch API. The kernels have string
             # that always starts with `Cijk*`
             mm_key = 'Cijk'
-            events = prof.key_averages()
+            events = prof.events()
             for evt in events:
-                if mm_key in evt.key:
+                if mm_key in evt.name:
                     self.assertEqual(evt.count, 1)
                     kernel_count = kernel_count + 1
 
@@ -6978,7 +7051,7 @@ def test_householder_product_errors_and_warnings(self, device):
                 torch.linalg.householder_product(reflectors, tau)
 
     @precisionOverride({torch.float32: 1e-2, torch.complex64: 1e-2})
-    @skipCUDAIfNoMagmaAndNoCusolver
+    @skipCUDAIfNoCusolver
     @skipIfTorchDynamo("Runtime error with torch._C._linalg.linalg_lu_factor")
     @skipCPUIfNoLapack
     @dtypes(*floating_and_complex_types())
@@ -7086,7 +7159,7 @@ def run_test(A, pivot, singular, fn):
                     f(torch.empty(1, 2, 2), pivot=False)
 
     @precisionOverride({torch.float32: 1e-2, torch.complex64: 1e-2})
-    @skipCUDAIfNoMagmaAndNoCusolver
+    @skipCUDAIfNoCusolver
     @skipCPUIfNoLapack
     @setLinalgBackendsToDefaultFinally
     @dtypes(*floating_and_complex_types())
@@ -7096,10 +7169,7 @@ def test_linalg_lu_solve(self, device, dtype):
         backends = ["default"]
 
         if torch.device(device).type == 'cuda':
-            if torch.cuda.has_magma:
-                backends.append("magma")
-            if has_cusolver():
-                backends.append("cusolver")
+            backends.append("cusolver")
 
         def gen_matrices():
             rhs = 3
@@ -7612,6 +7682,12 @@ def test_addmv(self, device, dtype):
         for m, v in itertools.product(ms, vs):
             self._test_addmm_addmv(torch.addmv, t, m, v, beta=0)
 
+    def test_addmv_mixed_types(self, device):
+        ten = torch.rand([2], dtype=torch.float64)
+        mat = torch.rand([2, 3], dtype=torch.complex128)
+        vec = torch.rand([3], dtype=torch.complex128)
+        self.assertRaises(RuntimeError, lambda: torch.addmv(ten, mat, vec))
+
     @dtypesIfCUDA(*floating_types_and(torch.bfloat16))
     @dtypes(torch.float, torch.double)
     def test_addmv_rowmajor_colmajor_incx_incy_lda(self, device, dtype):
@@ -7929,18 +8005,20 @@ def _gen_pair(m, k, n):
     @parametrize("use_transpose_a", [True, False])
     @parametrize("use_transpose_b", [True, False])
     @parametrize("non_contig_type", [0, 1, 2])
-    def test__int_mm_cpu(self, device, m, k, n, use_transpose_a, use_transpose_b, non_contig_type):
+    @parametrize("x_dtype", [torch.int8, torch.uint8])
+    def test__int_mm_cpu(self, device, m, k, n, use_transpose_a, use_transpose_b, non_contig_type, x_dtype):
         # non_contig_type:
         # 0: the whole data buffer is contiguous (can be transposed)
         # 1: stride of one dimension is 1, but the whole buffer is not contiguous
         # 2: Neither stride is 1
 
-        def genf_int_float(x, y, use_transpose, non_contig_type):
+        def genf_int_float(x, y, use_transpose, non_contig_type, dtype):
             if use_transpose:
                 x, y = y, x
             if non_contig_type != 0:
                 y = y * 2
-            x_int8 = torch.randint(-128, 127, (x, y), dtype=torch.int8, device=device)
+            dt_info = torch.iinfo(dtype)
+            x_int8 = torch.randint(dt_info.min, dt_info.max, (x, y), dtype=dtype, device=device)
             x_float = x_int8.to(torch.float32)
             if non_contig_type == 1:
                 x_int8 = x_int8[:, : y // 2]
@@ -7954,8 +8032,8 @@ def genf_int_float(x, y, use_transpose, non_contig_type):
 
         if non_contig_type != 0 and (m == 0 or k == 0):
             return
-        a_int8, a_float = genf_int_float(m, k, use_transpose_a, non_contig_type)
-        b_int8, b_float = genf_int_float(k, n, use_transpose_b, non_contig_type)
+        a_int8, a_float = genf_int_float(m, k, use_transpose_a, non_contig_type, x_dtype)
+        b_int8, b_float = genf_int_float(k, n, use_transpose_b, non_contig_type, torch.int8)
         c_int32 = torch._int_mm(a_int8, b_int8)
         self.assertTrue(c_int32.dtype is torch.int32)
         self.assertEqual(c_int32.device, torch.device(device))
@@ -8001,6 +8079,9 @@ def test__int4_mm(self, device, m, k, n):
         if self.device_type == 'cuda' and not SM80OrLater:
             self.skipTest("requires SM80 or later")
 
+        if TEST_WITH_ROCM and self.device_type == 'cuda' and CDNA5OrLater():
+            self.skipTest("int4 mm not yet implemented for gfx1250 (needs WMMA)")
+
         if TEST_WITH_ROCM and self.device_type == 'cuda' and not CDNA2OrLater():
             self.skipTest("_convert_weight_to_int4pack_cuda is supported only for CDNA2 or later")
 
@@ -8070,6 +8151,9 @@ def test_compile_int4_mm(self, device, m, k, n):
         if self.device_type == 'cuda' and not SM80OrLater:
             self.skipTest("requires SM80 or later")
 
+        if TEST_WITH_ROCM and self.device_type == 'cuda' and CDNA5OrLater():
+            self.skipTest("int4 mm not yet implemented for gfx1250 (needs WMMA)")
+
         if TEST_WITH_ROCM and self.device_type == 'cuda' and not CDNA2OrLater():
             self.skipTest("_convert_weight_to_int4pack_cuda supported only for CDNA2 or later")
 
@@ -9500,7 +9584,7 @@ def run_test(matsize, batchdims, mat_chars):
             run_test(matsize, batchdims, mat_chars=['sym', 'sym_pd', 'sym_psd'])
             run_test(matsize, batchdims, mat_chars=['sing', 'non_sing'])
 
-    @skipCUDAIfNoMagmaAndNoCusolver
+    @skipCUDAIfNoCusolver
     @skipCPUIfNoLapack
     @dtypes(*floating_and_complex_types())
     def test_cholesky_inverse(self, device, dtype):
@@ -9548,7 +9632,7 @@ def run_test(shape, batch, upper, contiguous):
         expected = torch.inverse(A)
         self.assertEqual(expected, out)
 
-    @skipCUDAIfNoMagma
+    @skipCUDAIfNoCusolver
     @skipCPUIfNoLapack
     @dtypes(*floating_and_complex_types())
     def test_cholesky_inverse_errors_and_warnings(self, device, dtype):
@@ -9745,7 +9829,7 @@ def lu_solve_test_helper(self, A_dims, b_dims, pivot, device, dtype):
         return b, A, LU_data, LU_pivots
 
     @skipCPUIfNoLapack
-    @skipCUDAIfNoMagmaAndNoCusolver
+    @skipCUDAIfNoCusolver
     @dtypes(*floating_and_complex_types())
     @precisionOverride({torch.float32: 1e-3, torch.complex64: 1e-3,
                         torch.float64: 1e-8, torch.complex128: 1e-8})
@@ -9761,7 +9845,7 @@ def sub_test(pivot):
             sub_test(False)
 
     @skipCPUIfNoLapack
-    @skipCUDAIfNoMagmaAndNoCusolver
+    @skipCUDAIfNoCusolver
     @dtypes(*floating_and_complex_types())
     @precisionOverride({torch.float32: 1e-3, torch.complex64: 1e-3,
                         torch.float64: 1e-8, torch.complex128: 1e-8})
@@ -9793,7 +9877,7 @@ def lu_solve_batch_test_helper(A_dims, b_dims, pivot):
 
     @slowTest
     @skipCPUIfNoLapack
-    @skipCUDAIfNoMagmaAndNoCusolver
+    @skipCUDAIfNoCusolver
     @dtypes(*floating_and_complex_types())
     def test_lu_solve_batched_many_batches(self, device, dtype):
         def run_test(A_dims, b_dims):
@@ -9806,7 +9890,7 @@ def run_test(A_dims, b_dims):
         run_test((262144, 5, 5), (262144, 5, 10))
 
     @skipCPUIfNoLapack
-    @skipCUDAIfNoMagmaAndNoCusolver
+    @skipCUDAIfNoCusolver
     @dtypes(*floating_and_complex_types())
     def test_lu_solve_batched_broadcasting(self, device, dtype):
         make_fullrank = make_fullrank_matrices_with_distinct_singular_values
@@ -9829,7 +9913,6 @@ def run_test(A_dims, b_dims, pivot=True):
         run_test((1, 3, 1, 4, 4), (2, 1, 3, 4, 5))  # broadcasting A & b
 
     @onlyCUDA
-    @skipCUDAIfNoMagma
     @dtypes(*floating_and_complex_types())
     # this tests https://github.com/pytorch/pytorch/issues/36921
     def test_lu_solve_large_matrices(self, device, dtype):
diff --git a/test/test_matmul_cuda.py b/test/test_matmul_cuda.py
index 1f1f1af0e7008..ca99d00e7068a 100644
--- a/test/test_matmul_cuda.py
+++ b/test/test_matmul_cuda.py
@@ -2,7 +2,6 @@
 
 import contextlib
 import os
-import time
 import unittest
 from itertools import product
 from functools import partial
@@ -19,8 +18,8 @@
 
 from torch.testing import make_tensor
 from torch.testing._internal.common_cuda import (
+    blas_library_context,
     PLATFORM_SUPPORTS_BF16,
-    PLATFORM_SUPPORTS_GREEN_CONTEXT,
     SM80OrLater,
     SM90OrLater,
     SM100OrLater,
@@ -47,6 +46,7 @@
     runOnRocmArch,
     serialTest,
     skipIfRocm,
+    skipIfRocmArch,
     TEST_CUDA,
     TEST_WITH_ROCM,
     TestCase,
@@ -77,15 +77,6 @@ def xfailIfSM100OrLaterNonRTXAndCondition(condition_fn):
     )
 
 
-@contextlib.contextmanager
-def blas_library_context(backend):
-    prev_backend = torch.backends.cuda.preferred_blas_library()
-    torch.backends.cuda.preferred_blas_library(backend)
-    try:
-        yield
-    finally:
-        torch.backends.cuda.preferred_blas_library(prev_backend)
-
 @contextlib.contextmanager
 def rocm_group_gemm_ck_env(value):
     var = "ROCM_ALLOW_GROUP_GEMM_CK"
@@ -250,6 +241,8 @@ def test_cublas_addmm_bias_shapes(self, size: int, dtype: torch.dtype, backend):
 
 
     @onlyCUDA
+    # Fails with triton 3.7
+    @skipIfRocmArch(NAVI_ARCH)
     @dtypes(torch.float16)
     # m == 4 chooses OUTPUT_TYPE reduction on H200
     # m == 8 chooses OUTPUT_TYPE reduction on A100
@@ -719,43 +712,56 @@ def test_grouped_gemm_compiled(self, op, a_row_major, b_row_major, max_autotune)
             self.assertEqual(C, C_ref)
 
     @skipCUDAIfNotRocm
+    # Fails with triton 3.7
     def test_grouped_gemm_rocm_ck_flag(self):
-        CK_HINT = "kernel_grouped_gemm_xdl_splitk"
+        CK_EQUAL_K_HINT = "kernel_grouped_gemm_xdl_splitk"
+        CK_UNEQUAL_K_HINT = "kernel_grouped_gemm_xdl_splitk"
         HIPBLASLT_HINT = "Cijk_Alik_Bljk_BBS_BH_Bias_HA_S_SAV_UserArgs"
 
-        def uses_ck(kernels: set[str]) -> bool:
-            return any(CK_HINT in k for k in kernels)
+        def has_ck_kernel(kernels: set[str], hint: str) -> bool:
+            return any(hint in k for k in kernels)
 
         def uses_hipblaslt(kernels: set[str]) -> bool:
             return any(HIPBLASLT_HINT in k for k in kernels)
 
-        def run_grouped_mm():
+        def run_grouped_mm(equal_k: bool):
             device = "cuda"
             dtype = torch.bfloat16
-            # row-major 3d-3d
-            G, M, N, K = 4, 16, 32, 64
-            a = torch.randn(G, M, K, device=device, dtype=dtype)
-            b = torch.randn(G, N, K, device=device, dtype=dtype)
-            # 3d-3d grouped GEMM: [G, M, K] @ [G, K, N]
-            out = F.grouped_mm(a, b.transpose(-2, -1), out_dtype=dtype)
-            return out
-
-        def collect_kernel_names():
+            if equal_k:
+                # 3d-3d grouped GEMM with identical K for all groups
+                G, M, N, K = 4, 16, 32, 64
+                a = torch.randn(G, M, K, device=device, dtype=dtype)
+                b = torch.randn(G, N, K, device=device, dtype=dtype)
+                return F.grouped_mm(a, b.transpose(-2, -1), out_dtype=dtype)
+
+            # 2d-2d grouped GEMM with non-uniform offs, i.e. per-group K is not equal
+            M, N = 16, 32
+            offs = torch.tensor([64, 136], device=device, dtype=torch.int32)
+            K_total = offs[-1].item()
+            a = torch.randn(M, K_total, device=device, dtype=dtype)
+            b = torch.randn(N, K_total, device=device, dtype=dtype)
+            return F.grouped_mm(a, b.transpose(-2, -1), offs=offs, out_dtype=dtype)
+
+        def collect_kernel_names(equal_k: bool):
             kernels = set()
             with profile(
                 activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
                 record_shapes=False,
                 with_stack=False,
             ) as prof:
-                run_grouped_mm()
+                run_grouped_mm(equal_k=equal_k)
             for evt in prof.key_averages(group_by_input_shape=False):
                 kernels.add(evt.key)
             return kernels
 
         with rocm_group_gemm_ck_env(None):
-            self.assertTrue(uses_hipblaslt(collect_kernel_names()))
+            self.assertTrue(uses_hipblaslt(collect_kernel_names(equal_k=True)))
         with rocm_group_gemm_ck_env("1"):
-            self.assertTrue(uses_ck(collect_kernel_names()))
+            ck_equal_kernels = collect_kernel_names(equal_k=True)
+            self.assertTrue(has_ck_kernel(ck_equal_kernels, CK_EQUAL_K_HINT))
+
+            ck_unequal_kernels = collect_kernel_names(equal_k=False)
+            self.assertTrue(has_ck_kernel(ck_unequal_kernels, CK_UNEQUAL_K_HINT))
 
     @onlyCUDA
     @parametrize("input_dtype", [torch.float32, torch.float16, torch.bfloat16])
@@ -1001,73 +1007,6 @@ def is_batched():
                     op(a, mismatch_batch_dim_b, out_dtype=torch.float32)
 
 
-    @unittest.skipIf(not PLATFORM_SUPPORTS_GREEN_CONTEXT, "Green contexts are not supported")
-    @serialTest()
-    def test_greencontext_carveout(self):
-        a = torch.randn(4096, 4096, device='cuda', dtype=torch.bfloat16)
-        ctx = torch.cuda.green_contexts.GreenContext.create(1, 0)
-        ctx.set_context()
-        torch.matmul(a, a)
-        torch.cuda.synchronize()
-        t0 = time.perf_counter()
-        partial_res = torch.matmul(a, a)
-        torch.cuda.synchronize()
-        t1 = time.perf_counter()
-        ctx.pop_context()
-        torch.matmul(a, a)
-        torch.cuda.synchronize()
-        t2 = time.perf_counter()
-        full_res = torch.matmul(a, a)
-        torch.cuda.synchronize()
-        t3 = time.perf_counter()
-        self.assertEqual(partial_res, full_res)
-        self.assertGreater(t1 - t0, t3 - t2)
-
-    @unittest.skipIf(not PLATFORM_SUPPORTS_GREEN_CONTEXT, "Green contexts are not supported")
-    @serialTest()
-    def test_greencontext_stream_carveout(self):
-        a = torch.randn(4096, 4096, device='cuda', dtype=torch.bfloat16)
-        ctx = torch.cuda.green_contexts.GreenContext.create(1, 0)
-        ctx_stream = ctx.Stream()
-        with torch.cuda.stream(ctx_stream):
-            torch.matmul(a, a)
-            torch.cuda.synchronize()
-            t0 = time.perf_counter()
-            partial_res = torch.matmul(a, a)
-            torch.cuda.synchronize()
-            t1 = time.perf_counter()
-        torch.matmul(a, a)
-        torch.cuda.synchronize()
-        t2 = time.perf_counter()
-        full_res = torch.matmul(a, a)
-        torch.cuda.synchronize()
-        t3 = time.perf_counter()
-        self.assertEqual(partial_res, full_res)
-        self.assertGreater(t1 - t0, t3 - t2)
-
-    @unittest.skipIf(not PLATFORM_SUPPORTS_GREEN_CONTEXT, "Green contexts are not supported")
-    @serialTest()
-    def test_greencontext_graphs(self):
-        a = torch.randn(4096, 4096, device='cuda', dtype=torch.bfloat16)
-        ctx = torch.cuda.green_contexts.GreenContext.create(1, 0)
-        ctx.set_context()
-        partial_res = torch.matmul(a, a)
-        ctx.pop_context()
-        full_res = torch.matmul(a, a)
-        full_res.zero_()
-        partial_res.zero_()
-        torch.cuda.synchronize()
-
-        g = torch.cuda.CUDAGraph()
-        with torch.cuda.graph(g):
-            ctx.set_context()
-            partial_res = torch.matmul(a, a)
-            ctx.pop_context()
-            full_res = torch.matmul(a, a)
-        g.replay()
-        self.assertEqual(partial_res, full_res)
-
-
 @unittest.skipIf(TEST_WITH_ROCM, "ROCm doesn't support CUTLASS")
 @unittest.skipIf(IS_WINDOWS, "Windows doesn't support CUTLASS extensions")
 @unittest.skipIf(not _IS_SM8X, "mixed dtypes linear only supported on SM 8.x")
@@ -1112,11 +1051,11 @@ def run_test(
             input_ref = input.reshape(-1, input.shape[-1])
 
             # First, test plain multiplication.
-            weight_ref = weight.T.to(input.dtype) * scale.view(1, n)
+            weight_ref = weight.T.to(torch.float32) * scale.float().view(1, n)
             weightq = (
                 pack_int4_to_int8(weight.T) if dtypeq == torch.quint4x2 else weight.T
             )
-            output_ref = torch.mm(input_ref, weight_ref).reshape(*input.shape[:-1], n)
+            output_ref = torch.mm(input_ref.float(), weight_ref).to(input.dtype).reshape(*input.shape[:-1], n)
             output = torch.ops.aten._mixed_dtypes_linear(
                 input,
                 quantized_weight_reorder_for_mixed_dtypes_linear_cutlass(
@@ -1127,12 +1066,12 @@ def run_test(
             torch.testing.assert_close(output, output_ref, rtol=rtol, atol=atol)
 
             # Second, test the linear operator itself.
-            weight_ref = weight.to(input.dtype) * scale.view(n, 1)
+            weight_ref = weight.to(torch.float32) * scale.float().view(n, 1)
             weightq = pack_int4_to_int8(weight) if dtypeq == torch.quint4x2 else weight
-            bias_ref = bias.view(1, n) if add_bias else None
+            bias_ref = bias.float().view(1, n) if add_bias else None
             output_ref = torch.nn.functional.linear(
-                input_ref, weight_ref, bias=bias_ref
-            ).reshape(*input.shape[:-1], n)
+                input_ref.float(), weight_ref, bias=bias_ref
+            ).to(input.dtype).reshape(*input.shape[:-1], n)
             if activation == "relu":
                 relu = torch.nn.ReLU()
                 output_ref = relu(output_ref)
diff --git a/test/test_meta.py b/test/test_meta.py
index 84b8967b1befb..ca5869722c98f 100644
--- a/test/test_meta.py
+++ b/test/test_meta.py
@@ -20,6 +20,7 @@
 from torch.testing._internal.common_utils import (
     TestCase,
     skipIfCrossRef,
+    skipIfTorchDynamo,
     suppress_warnings,
     TEST_WITH_TORCHDYNAMO,
     run_tests,
@@ -348,6 +349,9 @@ def test_tensor_outlives_converter(self):
 
 CHECK_STRIDES = {
     torch.Tensor.__getitem__,
+    aten._fft_c2c.default,
+    aten._fft_c2r.default,
+    aten._fft_r2c.default,
 }
 
 CHECK_ALL_STRIDES = {
@@ -356,9 +360,6 @@ def test_tensor_outlives_converter(self):
 
 CHECK_STRIDES_SKIPS = {
     aten._conj_physical.default,
-    aten._fft_c2c.default,
-    aten._fft_c2r.default,
-    aten._fft_r2c.default,
     aten._linalg_svd.default,
     aten.binary_cross_entropy.default,
     aten.complex.default,
@@ -1698,6 +1699,32 @@ def test_embedding_bag_dense_backward_per_sample_weights(self):
         )
         self.assertEqual(grad_weight.to('meta'), meta_grad_weight)
 
+    def _assert_fft_meta_stride_matches_eager(self, op, *args):
+        to_meta = MetaConverter()
+        meta_args = tree_map_only(torch.Tensor, to_meta, args)
+        ref_out = op(*args)
+        meta_out = op(*meta_args)
+        self.assertEqual(ref_out.size(), meta_out.size())
+        self.assertEqual(ref_out.stride(), meta_out.stride())
+
+    @onlyCUDA
+    @unittest.skipIf(torch.version.hip, "cuFFT-specific stride behavior")
+    def test_fft_multi_dim_cufft_stride_matches_meta(self, device):
+        self._assert_fft_meta_stride_matches_eager(
+            aten._fft_c2c.default,
+            torch.randn((5, 5, 5, 5, 5), device=device, dtype=torch.complex64),
+            [1, 2, 3, 4],
+            0,
+            True,
+        )
+        self._assert_fft_meta_stride_matches_eager(
+            aten._fft_c2r.default,
+            torch.randn((5, 5, 5, 5, 3), device=device, dtype=torch.complex64),
+            [0, 1, 2, 3, 4],
+            0,
+            5,
+        )
+
     # opinfo test is using aten.fill_, it's not testing aten.fill
     @onlyCUDA
     def test_fill_stride(self):
@@ -1869,8 +1896,87 @@ def fn(input, weight, bias, need_grad_input):
         else:
             self.assertEqual(out_dtype, [in_dtype,])
 
+class TestMetaKernelConv(TestCase):
+    @skipIfTorchDynamo("tests raw meta kernel, not dynamo")
+    def test_convolution_backward_meta_kernel_channels_last(self):
+        """Test the meta kernel directly (device='meta', no FakeTensorMode).
+        This exercises the @register_meta path used by torch.export, which
+        does NOT go through the FakeTensor intercept in fake_impls.py.
+        """
+        # channels_last grad_output + contiguous input/weight -> contiguous
+        grad_out = torch.empty(2, 3, 4, 4, device="meta").to(
+            memory_format=torch.channels_last
+        )
+        inp = torch.empty(2, 3, 4, 4, device="meta")
+        w = torch.empty(3, 3, 3, 3, device="meta")
+        gi, gw, _ = torch.ops.aten.convolution_backward(
+            grad_out,
+            inp,
+            w,
+            [3],
+            [1, 1],
+            [1, 1],
+            [1, 1],
+            False,
+            [0, 0],
+            1,
+            [True, True, True],
+        )
+        self.assertTrue(gi.is_contiguous())
+        self.assertTrue(gw.is_contiguous())
+
+        # contiguous grad_output + channels_last input -> channels_last
+        grad_out2 = torch.empty(2, 3, 4, 4, device="meta")
+        inp2 = torch.empty(2, 3, 4, 4, device="meta").to(
+            memory_format=torch.channels_last
+        )
+        gi2, gw2, _ = torch.ops.aten.convolution_backward(
+            grad_out2,
+            inp2,
+            w,
+            [3],
+            [1, 1],
+            [1, 1],
+            [1, 1],
+            False,
+            [0, 0],
+            1,
+            [True, True, True],
+        )
+        self.assertTrue(gi2.is_contiguous(memory_format=torch.channels_last))
+        self.assertTrue(gw2.is_contiguous(memory_format=torch.channels_last))
+
+
+
+
+class TestMetaKernelRegistrations(TestCase):
+    @skipIfTorchDynamo("tests raw meta kernel, not dynamo")
+    def test_make_dep_token(self):
+        cpu_result = torch.ops.aten._make_dep_token(device=torch.device("cpu"))
+        meta_result = torch.ops.aten._make_dep_token(device=torch.device("meta"))
+        self.assertEqual(cpu_result.shape, meta_result.shape)
+        self.assertEqual(cpu_result.dtype, meta_result.dtype)
+
+    @skipIfTorchDynamo("tests raw meta kernel, not dynamo")
+    def test_rrelu_backward_small_range(self):
+        from torch._decomp.decompositions import rrelu_with_noise_backward
+
+        x = torch.randn(5, requires_grad=True)
+        lower, upper = 0.125, 0.125 + torch.finfo(torch.float32).eps
+        noise = torch.rand(5)
+        grad = torch.ones(5)
+        cpp_result = torch.ops.aten.rrelu_with_noise_backward(
+            grad, x, noise, lower, upper, True, False
+        )
+        decomp_result = rrelu_with_noise_backward(
+            grad, x, noise, lower, upper, True, False
+        )
+        self.assertEqual(cpp_result, decomp_result)
+
+
 instantiate_device_type_tests(TestMeta, globals())
 
+
 def print_op_str_if_not_supported(op_str):
     op = OperatorName.parse(op_str)
     packet = getattr(torch.ops.aten, str(op.name))
diff --git a/test/test_mps.py b/test/test_mps.py
index 482a163986c17..9cdcd4b484b3c 100644
--- a/test/test_mps.py
+++ b/test/test_mps.py
@@ -21,6 +21,7 @@
 from collections import defaultdict
 from torch import inf
 from torch.nn import Buffer, Parameter
+from torch.export import Dim, export
 from torch.testing._internal import opinfo
 from torch.testing._internal.common_utils import \
     (gradcheck, gradgradcheck, parametrize, run_tests, TestCase, download_file, MACOS_VERSION, IS_CI,
@@ -809,6 +810,23 @@ def test_avg_pool2d_ceil_mode_mismatch(self):
             self.assertEqual(out_mps, out_cpu, msg=msg)
 
 
+    def test_channels_last_storage_offset(self):
+        # Regression test: channels_last tensors with non-zero storage_offset produced wrong
+        # results on MPS because the Placeholder path for NHWC ops ignored storage_offset.
+        # Build a channels_last tensor with offset>0 via slice+reshape+permute.
+        x = torch.randn(1, 65, 64, device="mps")[:, 1:].reshape(1, 8, 8, 64).permute(0, 3, 1, 2)
+        self.assertTrue(x.is_contiguous(memory_format=torch.channels_last))
+        self.assertGreater(x.storage_offset(), 0)
+
+        pool_mps = F.avg_pool2d(x, 2)
+        pool_cpu = F.avg_pool2d(x.cpu(), 2)
+        self.assertEqual(pool_mps.cpu(), pool_cpu)
+
+        bn_mps = torch.nn.BatchNorm2d(64).eval().to("mps")(x)
+        bn_cpu = torch.nn.BatchNorm2d(64).eval()(x.cpu())
+        self.assertEqual(bn_mps.cpu(), bn_cpu)
+
+
 class TestMPS(TestCaseMPS):
     def ulpAssertAllClose(self, output, reference, n_ulps):
         """
@@ -908,8 +926,6 @@ def test_exp1(self, device="mps", dtype=torch.float):
         # https://developer.apple.com/metal/Metal-Shading-Language-Specification.pdf Table 8.2
         self.ulpAssertAllClose(output.cpu(), output_cpu, n_ulps=4)
 
-
-
     def test_exp_strided_output(self):
         x = torch.rand((256, 10), device='mps')
         x_cpu = x.to("cpu")
@@ -958,6 +974,17 @@ def helper(val, shape, dtype):
         helper(0, [1024], torch.float32)
         helper(0.2, [2, 3], torch.float32)
         helper(0.2 + 0.5j, [2, 3], torch.complex64)
+        helper(2**63 + 100, [3], torch.uint64)
+
+    def test_fill_strided(self):
+        # Regression test: strided fill_ must convert byte strides to element strides
+        for dtype in [torch.cfloat, torch.float32, torch.int32, torch.uint16, torch.uint8]:
+            x_mps = torch.zeros(4, 4, device='mps', dtype=dtype)
+            x_cpu = x_mps.cpu()
+            # Every other row — non-contiguous view
+            for t in (x_mps, x_cpu):
+                t[::2].fill_(1)
+            self.assertEqual(x_mps, x_cpu)
 
     def test_fill_storage_offset(self):
         shape = [2, 10]
@@ -1207,6 +1234,27 @@ def test_bmm(self):
         self.assertEqual(output_cpu, output_mps)
         self.assertEqual(output_cpu.size(), output_mps.size())
 
+    def test_bmm_conj(self):
+        # bmm must respect the conjugate bit on input tensors.
+        # See https://github.com/pytorch/pytorch/issues/177474
+        a = torch.randn(4, 3, 5, dtype=torch.complex64, device="mps")
+        b = torch.randn(4, 5, 2, dtype=torch.complex64, device="mps")
+        result_mps = torch.bmm(a, b.conj())
+        result_cpu = torch.bmm(a.cpu(), b.cpu().conj())
+        self.assertEqual(result_cpu, result_mps)
+        result_mps = torch.bmm(a.conj(), b)
+        result_cpu = torch.bmm(a.cpu().conj(), b.cpu())
+        self.assertEqual(result_cpu, result_mps)
+
+    def test_addmm_conj(self):
+        # Regression test: addmm must respect the conjugate bit on the bias tensor.
+        bias = torch.randn(3, 2, dtype=torch.complex64, device="mps")
+        a = torch.randn(3, 5, dtype=torch.complex64, device="mps")
+        b = torch.randn(5, 2, dtype=torch.complex64, device="mps")
+        result_mps = torch.addmm(bias.conj(), a, b)
+        result_cpu = torch.addmm(bias.cpu().conj(), a.cpu(), b.cpu())
+        self.assertEqual(result_cpu, result_mps)
+
     @xfailIf(MACOS_VERSION < 15.0)
     @parametrize("dtype", [torch.float16, torch.bfloat16])
     def test_large_bmm(self, dtype):
@@ -1358,6 +1406,19 @@ def test_linear_non_contiguous(self):
         result_contig = torch.nn.functional.linear(input_s, weight_contiguous_equiv)
         self.assertEqual(result_contig, result_sliced)
 
+    def test_linear_backward_channels_last_grad(self):
+        # Regression test for https://github.com/pytorch/pytorch/issues/178222
+        # Linear backward crashed when grad_output had channels_last strides,
+        # because suggest_memory_format() returned ChannelsLast which was then
+        # applied to 2D weight grad and 1D bias grad tensors (requires rank 4).
+        x = torch.randn(2, 8, 3, 4, device='mps', requires_grad=True)
+        proj = torch.nn.Linear(4, 4, device='mps')
+        y = proj(x)
+        z = y.permute(0, 2, 3, 1).contiguous()
+        target = torch.randn_like(z)
+        loss = (z - target).pow(2).sum()
+        loss.backward()
+
     def _linear_helper(self, in_features, out_features, shape, bias=True, backward_pass=False):
         cpu_linear = torch.nn.Linear(in_features=in_features, out_features=out_features, device="cpu", bias=bias)
         mps_linear = torch.nn.Linear(in_features=in_features, out_features=out_features, device="mps", bias=bias)
@@ -1657,6 +1718,53 @@ def helper(shape):
         helper([10, 5, 10, 3])
         helper([10, 5, 10, 3, 20])
 
+        # Test scalar tensor (0-dimensional) - should preserve shape
+        x_mps = torch.tensor(0, device="mps")
+        x_cpu = x_mps.detach().clone().cpu()
+        mask_mps = torch.tensor(True, device="mps")
+        mask_cpu = mask_mps.detach().clone().cpu()
+        source_mps = torch.tensor([42], device="mps")
+        source_cpu = source_mps.detach().clone().cpu()
+        result_mps = x_mps.masked_scatter(mask_mps, source_mps)
+        result_cpu = x_cpu.masked_scatter(mask_cpu, source_cpu)
+        self.assertEqual(result_mps.shape, result_cpu.shape)
+        self.assertEqual(result_mps, result_cpu)
+
+    def test_masked_scatter_raises(self):
+        x_mps = torch.tensor([0, 0, 0], device="mps")
+        mask_mps = torch.tensor([1, 1, 1], dtype=torch.bool, device="mps")
+        source_mps = torch.tensor([2, 5], device="mps")
+
+        with self.assertRaisesRegex(RuntimeError, "source < number of ones in mask"):
+            x_mps.masked_scatter_(mask_mps, source_mps)
+
+    def test_masked_scatter_without_mask_indices(self):
+        x_mps = torch.tensor([1, 2, 3, 4], device="mps")
+        mask_mps = torch.tensor(0, dtype=torch.bool, device="mps")
+        source_mps = torch.tensor([2, 5], device="mps")
+
+        y_mps = x_mps.masked_scatter(mask_mps, source_mps)
+
+        self.assertEqual(x_mps, y_mps)
+
+    def test_masked_scatter_no_side_effects(self):
+        # Regression test for https://github.com/pytorch/pytorch/issues/175576
+        x_mps = torch.zeros(5, dtype=torch.int64, device="mps")
+        mask_mps = torch.tensor([1, 0, 1, 0, 1], dtype=torch.bool, device="mps")
+        source_mps = torch.tensor([[1, 2], [3, 4]], device="mps")
+
+        x_cpu = x_mps.detach().clone().cpu()
+        mask_cpu = mask_mps.detach().clone().cpu()
+        source_cpu = source_mps.detach().clone().cpu()
+
+        y_cpu = x_cpu.masked_scatter(mask_cpu, source_cpu)
+        y_mps = x_mps.masked_scatter(mask_mps, source_mps)
+
+        self.assertEqual(y_cpu, y_mps)
+        self.assertEqual(x_cpu, x_mps)
+        self.assertEqual(mask_cpu, mask_mps)
+        self.assertEqual(source_cpu, source_mps)
+
     def test_masked_fill(self):
         device = "mps"
         dtype = torch.float32
@@ -1919,6 +2027,19 @@ def test_batch_norm_backward_weight_bias_gradients(self):
         self.assertEqual(bn_cpu.weight.grad, bn_mps.weight.grad, atol=1e-5, rtol=1e-5)
         self.assertEqual(bn_cpu.bias.grad, bn_mps.bias.grad, atol=1e-5, rtol=1e-5)
 
+    def test_batch_norm_mixed_dtype(self):
+        # Regression test for https://github.com/pytorch/pytorch/issues/178770
+        # BatchNorm with float32 weights and float16 input should work
+        x = torch.rand(2, 32, 15, 15, device="mps", dtype=torch.float16, requires_grad=True)
+        model = nn.BatchNorm2d(32, device="mps", dtype=torch.float32).train()
+        y_mps = model(x)
+        # Compare against CPU using a copy of the model
+        import copy
+        y_cpu = copy.deepcopy(model).cpu()(x.detach().cpu())
+        self.assertEqual(y_cpu, y_mps.cpu(), atol=1e-3, rtol=1e-3)
+        y_mps.sum().backward()
+        self.assertIsNotNone(x.grad)
+
     def test_layer_norm_backward(self):
         inputs = torch.rand(4, 4, device="mps", requires_grad=True)
         x = torch.nn.LayerNorm(4).to("mps")
@@ -5998,7 +6119,7 @@ def run_cholesky_test(size, *batch_dims, upper=False, check_errors=False):
             input_mps = input_cpu.to('mps')
             output_cpu = torch.linalg.cholesky_ex(input_cpu, upper=upper)
             output_mps = torch.linalg.cholesky_ex(input_mps, upper=upper)
-            self.assertEqual(output_cpu, output_mps, atol=2e-5, rtol=1e-6)
+            self.assertEqual(output_cpu, output_mps, atol=3e-5, rtol=1e-6)
 
         # test with different even/odd matrix sizes
         matrix_sizes = [1, 2, 3, 4, 8, 17, 64, 128, 154]
@@ -8121,6 +8242,18 @@ def helper(shape, alpha, op_name, inplace):
         # Regression test for https://github.com/pytorch/pytorch/issues/160208
         self.assertEqual(torch.add(y, x, alpha=2).cpu(), torch.add(y.cpu(), x.cpu(), alpha=2))
 
+    def test_add_sub_alpha_cast(self):
+        # In-place with alpha when self is promoted (e.g. float16 + float32)
+        for op in [torch.Tensor.add_, torch.Tensor.sub_]:
+            for dtype_a, dtype_b in [(torch.float16, torch.float32), (torch.bfloat16, torch.float32)]:
+                a_mps = torch.arange(16, dtype=dtype_a, device='mps')
+                b_mps = torch.arange(16, dtype=dtype_b, device='mps')
+                a_cpu, b_cpu = a_mps.cpu(), b_mps.cpu()
+                alpha = torch.tensor(0.33, dtype=dtype_b)
+                op(a_mps, b_mps, alpha=alpha)
+                op(a_cpu, b_cpu, alpha=alpha)
+                self.assertEqual(a_mps, a_cpu)
+
     # Test add
     def test_add_scalars(self):
         def helper(alpha):
@@ -8434,6 +8567,12 @@ def randint(*shape, **kw):
         cond, a, b = rand(100) > 0.5, rand(100), rand(100)
         self.assertEqual(torch.where(cond, a, b), torch.where(cond.cpu(), a.cpu(), b.cpu()))
 
+    def test_unsigned_cast_div(self):
+        # See https://github.com/pytorch/pytorch/issues/176296
+        x = torch.tensor([0, 65535], dtype=torch.uint16, device="mps")
+        y = x / 64
+        self.assertEqual(y, torch.tensor([0., 1023.9844], device="mps"))
+
 
 class TestLargeTensors(TestCaseMPS):
     @serialTest()
@@ -9163,6 +9302,25 @@ def test_conv3d_backward_collision(self):
         # This used to crash with MPSNDArrayConvolutionA14.mm:4352: failed assertion
         y2.sum().backward()
 
+    def test_channels_last_channel_slice(self):
+        # Regression test for https://github.com/pytorch/pytorch/issues/180984
+        # A channel-slice view of a channels_last tensor has channels-last-like
+        # strides but is not packed NHWC in memory. MPSGraph ops that support NHWC
+        # are only work with packed NHWC buffer, giving wrong results.
+        shared = torch.randn(2, 4, 8, 8, device="mps").contiguous(memory_format=torch.channels_last)
+        mps_slice = shared[:, :2]
+
+        weight = torch.randn(3, 2, 3, 3, device="mps")
+        self.assertEqual(F.conv2d(mps_slice.cpu(), weight.cpu()), F.conv2d(mps_slice, weight).cpu())
+
+        self.assertEqual(F.avg_pool2d(mps_slice.cpu(), 2), F.avg_pool2d(mps_slice, 2).cpu())
+        self.assertEqual(F.adaptive_avg_pool2d(mps_slice.cpu(), 2), F.adaptive_avg_pool2d(mps_slice, 2).cpu())
+
+        bn = nn.BatchNorm2d(2).eval()
+        bn_mps = nn.BatchNorm2d(2).to("mps").eval()
+        bn_mps.load_state_dict(bn.state_dict())
+        self.assertEqual(bn(mps_slice.cpu()), bn_mps(mps_slice).cpu())
+
     # Regression test for https://github.com/pytorch/pytorch/issues/141471
     def test_conv3d_channels_last_3d(self):
         m_cpu = nn.Conv3d(16, 33, (3, 5, 2), stride=(2, 1, 1), padding=(4, 2, 0), device="cpu")
@@ -9627,12 +9785,39 @@ def weight_int8pack_mm(a, b_int8pack, b_scales):
             mean_err = ((res - ref).abs() / ref).mean()
             self.assertLess(mean_err, 0.05)
 
+    def test_loradown_regression_original_case(self):
+        a = torch.rand(2, 1025, device='mps', dtype=torch.half)
+        b = torch.rand(2, 1041, device='mps', dtype=torch.half)[:, :1025].t()
+        result = a @ b
+        self.assertEqual(result.shape, (2, 2))
+
+        self.assertFalse(torch.isnan(result).any())
+        self.assertFalse(torch.isinf(result).any())
+
+    @parametrize("padding", [0, 3, 4, 7, 8, 15, 16])
+    @parametrize("vector_dim", [2, 15, 16, 24])
+    def test_loradown_correctness_vs_cpu(self, padding, vector_dim):
+        torch.manual_seed(13)
+
+        base_size = 64
+        physical_size = base_size + padding
+
+        a_mps = torch.rand(vector_dim, base_size, device='mps', dtype=torch.half)
+        b_mps = torch.rand(vector_dim, physical_size, device='mps', dtype=torch.half)[:, :base_size].t()
+
+        a_cpu = a_mps.cpu()
+        b_cpu = b_mps.cpu()
+
+        result_cpu = (a_cpu @ b_cpu)
+        result_mps = (a_mps @ b_mps).cpu()
+
+        torch.testing.assert_close(result_mps, result_cpu, rtol=1e-3, atol=1e-3)
 
 class TestSDPA(TestCaseMPS):
-    def _compare_tensors(self, y, ref):
+    def _compare_tensors(self, y, ref, tol=0.01):
         denom = torch.maximum(ref.abs(), torch.tensor([1e-6], device=ref.device, dtype=ref.dtype))
         err = ((y - ref).abs() / denom).mean().item()
-        self.assertLess(err, 0.01)
+        self.assertLess(err, tol)
 
     def _test_sdpa_no_mask(
         self,
@@ -9677,6 +9862,38 @@ def test_sdpa_no_mask_causal_fp32(self):
     def test_sdpa_no_mask_causal_fp16(self):
         self._test_sdpa_no_mask(True, torch.float16)
 
+    def test_sdpa_export_dynamic_seq_len(self):
+        # Regression test for https://github.com/pytorch/pytorch/issues/177603
+        class M(nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.register_buffer("k", torch.zeros(1, 4, 64, 128))
+                self.register_buffer("v", torch.zeros(1, 4, 64, 128))
+
+            def forward(self, q, mask):
+                out, _ = torch.ops.aten._scaled_dot_product_attention_math_for_mps(
+                    q, self.k, self.v, mask, 0.0, False, None
+                )
+                return out
+
+        model = M().to(device="mps", dtype=torch.float32).eval()
+        seq = Dim("seq", min=1, max=64)
+
+        q = torch.randn(1, 4, 4, 128, device="mps", dtype=torch.float32)
+        mask = torch.zeros(4, 64, device="mps", dtype=torch.bool)
+
+        ep = export(
+            model,
+            (q, mask),
+            dynamic_shapes={"q": {2: seq}, "mask": {0: seq}},
+            strict=True,
+        )
+
+        q2 = torch.randn(1, 4, 7, 128, device="mps", dtype=torch.float32)
+        mask2 = torch.zeros(7, 64, device="mps", dtype=torch.bool)
+        out = ep.module()(q2, mask2)
+        self.assertEqual(out.shape, (1, 4, 7, 128))
+
     def test_sdpa_no_mask_causal_fp16_L7(self):
         self._test_sdpa_no_mask(True, torch.float16, 7)
 
@@ -9736,6 +9953,40 @@ def test_sdpa_full_mask(self, dtype):
         out_mps = F.scaled_dot_product_attention(q.to('mps'), k.to('mps'), v.to('mps'), attn_mask=mask.to('mps'))
         self._compare_tensors(out_mps.cpu(), out_cpu)
 
+    @parametrize("dtype", [torch.float32, torch.bfloat16])
+    def test_sdpa_math_mps_bool_mask_1pass(self, dtype):
+        torch.manual_seed(0)
+        q = torch.randn(1, 1, 8, 64, dtype=dtype, device='mps')
+        k = torch.randn(1, 1, 512, 64, dtype=dtype, device='mps')
+        v = torch.randn(1, 1, 512, 64, dtype=dtype, device='mps')
+        mask = torch.eye(8, 512, dtype=torch.bool, device='mps')
+
+        out_mps, _ = torch.ops.aten._scaled_dot_product_attention_math_for_mps(
+            q, k, v, mask
+        )
+        out_cpu = F.scaled_dot_product_attention(
+            q.cpu(), k.cpu(), v.cpu(), attn_mask=mask.cpu()
+        )
+        self._compare_tensors(out_mps.cpu(), out_cpu)
+
+    @parametrize("dtype", [torch.bfloat16, torch.float16, torch.float])
+    def test_sdpa_2pass(self, dtype):
+        # Regression test for https://github.com/pytorch/pytorch/issues/174861
+        q = torch.randn(1, 32, 1, 128, dtype=dtype)
+        k = torch.randn(1, 2, 1024, 128, dtype=dtype)
+        v = torch.randn(1, 2, 1024, 128, dtype=dtype)
+        sdpa_kwargs = {"enable_gqa": True}
+
+        out_cpu = F.scaled_dot_product_attention(q, k, v, **sdpa_kwargs)
+        out_mps = F.scaled_dot_product_attention(
+            q.to("mps"), k.to("mps"), v.to("mps"), **sdpa_kwargs
+        )
+
+        tol = 0.1 if dtype == torch.bfloat16 else 0.01
+
+        self.assertEqual(out_mps, out_cpu, atol=1e-3, rtol=1e-6)
+        self._compare_tensors(out_mps.cpu(), out_cpu, tol=tol)
+
     @parametrize("dtype", [torch.float16, torch.float32])
     def test_sdpa_3d_input(self, dtype):
         head_num, seq_len, embed_dim = 16, 16, 80
@@ -9950,6 +10201,26 @@ def test_fast_vector_attention_2pass(self, dtype: torch.dtype, layout: str, with
         q, k, v = self.generate_qkv(batch, NH, q_len, s_len, head_dim, layout, dtype)
         self.run_fast_attention_test(q, k, v, with_mask)
 
+    def test_fast_vector_permuted_inputs_regression(self):
+        # Regression test for https://github.com/pytorch/pytorch/issues/181133
+        torch.manual_seed(0)
+        B, S, H, D = 84, 3, 3, 64
+        dtype = torch.float32
+        q = torch.randn(B, S, H, D, dtype=dtype)
+        k = torch.randn(B, S, H, D, dtype=dtype)
+        v = torch.randn(B, S, H, D, dtype=dtype)
+
+        q_mps = q.to("mps").permute(0, 2, 1, 3)
+        k_mps = k.to("mps").permute(0, 2, 1, 3)
+        v_mps = v.to("mps").permute(0, 2, 1, 3)
+        q_cpu = q.permute(0, 2, 1, 3)
+        k_cpu = k.permute(0, 2, 1, 3)
+        v_cpu = v.permute(0, 2, 1, 3)
+
+        y_mps = F.scaled_dot_product_attention(q_mps, k_mps, v_mps)
+        y_cpu = F.scaled_dot_product_attention(q_cpu, k_cpu, v_cpu)
+        self._compare_tensors(y_mps.cpu(), y_cpu)
+
     @unittest.skip("Full attention fast kernel not implemented yet")
     @parametrize("dtype", [torch.float16, torch.float32])
     @parametrize("layout", ["contiguous", "mT"])
@@ -9964,6 +10235,26 @@ def test_fast_full_attention(self, dtype: torch.dtype, layout: str, head_dim: in
         q, k, v = self.generate_qkv(batch, NH, q_len, s_len, head_dim, layout, dtype)
         self.run_fast_attention_test(q, k, v, with_mask)
 
+    def test_sdpa_output_qv_shape_diff(self):
+        # This intentionally checks both output shape semantics and numerical consistency
+        # against CPU for Ev != E. It is stricter than a shape-only regression test.
+        # Regression test for https://github.com/pytorch/pytorch/issues/176767
+        test_cases = [
+            ((1, 16, 32), (1, 16, 32), (1, 16, 24)),
+            ((1, 1, 16, 32), (1, 1, 16, 32), (1, 1, 16, 24)),
+            ((2, 3, 1, 16, 32), (2, 3, 1, 16, 32), (2, 3, 1, 16, 24)),
+        ]
+
+        for q_shape, k_shape, v_shape in test_cases:
+            q = torch.randn(*q_shape, device="mps", dtype=torch.float32)
+            k = torch.randn(*k_shape, device="mps", dtype=torch.float32)
+            v = torch.randn(*v_shape, device="mps", dtype=torch.float32)
+
+            y_mps = torch.nn.functional.scaled_dot_product_attention(q, k, v)
+            y_cpu = torch.nn.functional.scaled_dot_product_attention(q.cpu(), k.cpu(), v.cpu())
+
+            self.assertEqual(y_cpu, y_mps)
+
 
 
 class TestSDPAMetaDispatchMode(TorchDispatchMode):
@@ -9992,8 +10283,10 @@ def format_res(res):
                 for t in pytree.tree_flatten(res)[0]
             ]
 
-        # Format the output so that we only look at the tensor metadata
-        self.test.assertEqual(format_res(res), format_res(meta_res))
+        # Format the output so that we only look at the tensor metadata.
+        # Only compare the first returned value for this op; the second output
+        # is not consumed and inconsistent across paths.
+        self.test.assertEqual(format_res(res)[0], format_res(meta_res)[0])
         return res
 
 
@@ -11505,6 +11798,18 @@ def test_nonzero_discontiguous(self):
         self.assertEqual(dst1, dst4, atol=0, rtol=0)
         self.assertEqual(strides, dst4.stride())
 
+    def test_nonzero_large(self):
+        # Regression test: with 2M elements and threadgroup size 1024, the
+        # prefix_sum_blocks kernel must handle ~1954 blocks via its
+        # multi-element-per-thread loop (each thread processes
+        # ceil(num_blocks/tg_size) blocks). This verifies correctness when
+        # the block count exceeds a single threadgroup.
+        x = torch.rand(2_000_000, device="mps")
+        x[x > 0.5] = 0
+        result = torch.nonzero(x)
+        expected = torch.nonzero(x.cpu())
+        self.assertEqual(result.cpu(), expected)
+
     def test_nonzero_non_diff(self):
         device = "mps"
         x = torch.randn(10, requires_grad=True, device=device)
@@ -12690,6 +12995,8 @@ def _compute_tolerances(self, op, dtype):
         # TODO: Rounding is broken for linspace, see https://github.com/pytorch/pytorch/issues/137635
         if op.name == 'linspace' and dtype in [torch.int8, torch.uint8, torch.int32, torch.int16, torch.int64]:
             return (1.0, 0.0)
+        if op.name == "index_reduce" and op.variant_test_name in ['mean', 'prod'] and dtype in [torch.float16, torch.bfloat16]:
+            return (0.01, 0.01)
         return (None, None)
 
     # Used for accept mode only
@@ -12708,9 +13015,10 @@ def _run_op(self, op, mps_sample, dtype=None):
                 # TODO: Handle list inputs later
                 if not isinstance(mps_out, torch.Tensor):
                     raise
-                if mps_sample.input.dtype not in [torch.float16, torch.bfloat16]:
-                    raise
-                dtype = torch.float32
+                if mps_sample.input.dtype in [torch.float16, torch.bfloat16]:
+                    dtype = torch.float32
+                elif mps_sample.input.dtype == torch.bool:
+                    dtype = torch.uint8
 
                 # Often CPU ops are not implemented for low precision dtypes
                 # In that case, upcast to higher precision and try again
@@ -12842,6 +13150,8 @@ def req_grad(t):
                 atol, rtol = 5e-3, 5e-3
             if op.name == "nn.functional.embedding_bag" and dtype == torch.float16:
                 atol, rtol = 5e-3, 5e-3
+            if op.name == "index_reduce" and op.variant_test_name in ['mean', 'prod'] and dtype in [torch.float16]:
+                atol, rtol = 0.02, 0.02
 
             if isinstance(cpu_sample.input, torch.Tensor):
                 equal_input_types = cpu_sample.input.dtype == mps_sample.input.dtype
@@ -12931,6 +13241,16 @@ def test_fmax_mixed_dtypes(self, device):
             # Broadcast
             self.assertEqual(op(x, y[0]), op(x.to("mps"), y.to("mps")[0]).cpu())
 
+    def test_mm_stride_zero(self):
+        # Regression test for https://github.com/pytorch/pytorch/issues/180201
+        # MPSGraph matrixMultiplication produces incorrect results with stride-0
+        # inputs on macOS < 26.4 (only every 16th row is correct).
+        expanded = torch.ones(1, 1, device="mps").expand(64, 1)
+        w = torch.randn(1, 16, device="mps")
+        result = expanded.mm(w)
+        expected = torch.ones(64, 1, device="mps").mm(w)
+        self.assertEqual(result, expected)
+
 
 class TestErrorInputs(TestCase):
     _ignore_not_implemented_error = True
@@ -13131,6 +13451,21 @@ def test_metal_error_checking(self):
         self.assertRaises(ValueError, lambda: lib.full(mps_tensor, threads=(3, max_thread_group_size),
                                                        group_size=(3, max_thread_group_size)))
 
+    def test_metal_randn(self):
+        lib = torch.mps.compile_shader("""
+            #include <c10/metal/random.h>
+            kernel void randn(device float* out, constant long2& seed_offset,
+                              uint idx [[thread_position_in_grid]]) {
+              out[idx] = c10::metal::randn(seed_offset.x, seed_offset.y + idx);
+            }
+        """)
+        N = 100_000
+        out = torch.empty(N, device="mps")
+        seed_offset = torch.tensor([42, 0], dtype=torch.long, device="mps")
+        lib.randn(out, seed_offset)
+        self.assertEqual(out.mean(), torch.tensor(0.0, device="mps"), atol=0.01, rtol=0)
+        self.assertEqual(out.std(), torch.tensor(1.0, device="mps"), atol=0.02, rtol=0)
+
     def test_metal_include(self):
         # Checks that includes embedding works
         lib = torch.mps.compile_shader("#include <c10/metal/special_math.h>")
@@ -13295,6 +13630,22 @@ def test_metal_error_buffer(self):
         with self.assertRaisesRegex(RuntimeError, "Index .* exceeds limit"):
             torch.mps.synchronize()
 
+    def test_metal_compiler_bug_workaround(self):
+        # (uint / 65536) % non_power of 2, gives wrong result without safe_mod
+        lib = torch.mps.compile_shader('''
+            #include <c10/metal/utils.h>
+
+            kernel void func(device int* out, uint idx [[thread_position_in_grid]]) {
+                out[idx] = c10::metal::safe_mod((idx / 65536), 6);
+            }
+        ''')
+        out = torch.empty(128, device='mps', dtype=torch.int32)
+        lib.func(out)
+        # Every value should be 0 since xindex/65536 == 0 for xindex in [0,127]
+        for i in [0, 5, 6, 7, 63, 64]:
+            self.assertEqual(out[i], 0)
+
+
 
 # TODO: Actually instantiate that test for the "mps" device to better reflect what it is doing.
 # This requires mps to be properly registered in the device generic test framework which is not the
diff --git a/test/test_native_functions.py b/test/test_native_functions.py
index 5a894c278fd04..198bdf8989101 100644
--- a/test/test_native_functions.py
+++ b/test/test_native_functions.py
@@ -1,6 +1,5 @@
 # Owner(s): ["module: unknown"]
 
-from typing import Optional
 import torch
 from torch.testing._internal.common_utils import TestCase, run_tests, skipIfTorchDynamo
 
@@ -8,12 +7,12 @@
 
 
 class FloatListWrapperModule(torch.nn.Module):
-    def forward(self, values, incr: Optional[list[float]]):
+    def forward(self, values, incr: list[float] | None):
         return torch._C._nn._test_optional_floatlist(values, incr)
 
 
 class IntListWrapperModule(torch.nn.Module):
-    def forward(self, values, incr: Optional[list[int]]):
+    def forward(self, values, incr: list[int] | None):
         return torch._C._nn._test_optional_intlist(values, incr)
 
 
diff --git a/test/test_nestedtensor.py b/test/test_nestedtensor.py
index 7491ce6107167..7a7cdd0ee91d1 100644
--- a/test/test_nestedtensor.py
+++ b/test/test_nestedtensor.py
@@ -10,7 +10,6 @@
 import tempfile
 import unittest
 from functools import partial
-from typing import Optional
 
 import numpy as np
 
@@ -3935,7 +3934,7 @@ def get_atol(true_value: torch.Tensor, computed_value: torch.Tensor) -> float:
 def get_tolerances(
     true_value: torch.Tensor,
     computed_value: torch.Tensor,
-    fudge_factor: Optional[float] = None,
+    fudge_factor: float | None = None,
 ) -> tuple[float, float]:
     """Returns the absolute and relative tolerances for comparing two tensors."""
     fudge_factor = fudge_factor if fudge_factor is not None else 1.0
@@ -8785,7 +8784,7 @@ def f(values, offsets):
     # clone() -> preserve format on an non-contiguous NJT with holes currently uses
     # unbind(), leading to data-dependent expression. Should be fixed via torch._check()
     XFailRule(
-        error_type=torch._dynamo.exc.Unsupported,
+        error_type=(torch._dynamo.exc.Unsupported, torch._dynamo.exc.UserError),
         # Ne(u1, u0) (unhinted: Ne(u1, u0)).  (Size-like symbols: u1, u0)
         error_msg="Could not guard on data-dependent expression",
         op_match_fn=lambda device, op: (op.full_name == "clone"),
diff --git a/test/test_nn.py b/test/test_nn.py
index 0695b021edbeb..51037a0f66966 100644
--- a/test/test_nn.py
+++ b/test/test_nn.py
@@ -2,6 +2,8 @@
 # ruff: noqa: F841
 
 import contextlib
+import ctypes
+import ctypes.util
 import math
 import random
 import unittest
@@ -32,10 +34,10 @@
 from torch.nn.parallel._functions import Broadcast
 from torch.testing._internal.common_dtype import integral_types, get_all_math_dtypes, floating_types
 from torch.testing._internal.common_utils import dtype_name, freeze_rng_state, run_tests, TestCase, \
-    skipIfNoLapack, skipIfRocm, MI300_ARCH, skipIfRocmArch, \
+    skipIfNoLapack, skipIfRocmVersionLessThan, MI300_ARCH, skipIfRocmArch, \
     TEST_NUMPY, TEST_SCIPY, TEST_WITH_CROSSREF, TEST_WITH_ROCM, \
     download_file, get_function_arglist, load_tests, skipIfMPS, \
-    IS_PPC, \
+    IS_PPC, IS_ARM64, IS_MACOS, IS_WINDOWS, IS_CPU_CAPABILITY_SVE256, IS_CPU_EXT_SVE_SUPPORTED, xfailIf, \
     parametrize as parametrize_test, subtest, instantiate_parametrized_tests, \
     skipIfTorchDynamo, gcIfJetson, set_default_dtype
 from torch.testing._internal.common_cuda import TEST_CUDA, TEST_MULTIGPU, TEST_CUDNN, \
@@ -45,7 +47,7 @@
     ctcloss_reference, get_new_module_tests, single_batch_reference_fn, _test_bfloat16_ops, _test_module_empty_input
 from torch.testing._internal.common_device_type import dtypesIfMPS, instantiate_device_type_tests, dtypes, \
     dtypesIfCUDA, precisionOverride, onlyCUDA, onlyCPU, \
-    skipCUDAIfRocm, skipCUDAIf, skipCUDAIfNotRocm, \
+    skipCUDAIf, skipCUDAIfNotRocm, \
     onlyNativeDeviceTypes, deviceCountAtLeast, largeTensorTest, expectedFailureMeta, expectedFailureMPS, \
     skipMeta, get_all_device_types
 
@@ -2223,6 +2225,7 @@ def test_threshold_bfloat16_half(self):
                 res_bf16 = F.threshold(x.to(dtype=dtype), threshold, 0).float()
                 self.assertEqual(res_bf16, expected)
 
+    @xfailIf(IS_ARM64 and not IS_CPU_EXT_SVE_SUPPORTED)  # SIGILL on AArch64 without SVE
     @unittest.skipUnless('fbgemm' in torch.backends.quantized.supported_engines,
                          'Linear_FP16_weight requires FBGEMM. FBGEMM is only optimized for CPUs'
                          ' with instruction set support avx2 or newer.')
@@ -2667,6 +2670,31 @@ def test_weighted_mse_loss(self):
         expected_loss = torch.tensor(0.25)
         self.assertTrue(torch.isclose(loss, expected_loss), f"Expected {expected_loss}, but got {loss}")
 
+    def test_mse_loss_mixed_dtype_grad(self):
+        devices = ['cpu'] + (['cuda'] if TEST_CUDA else [])
+        for device in devices:
+            for reduction in ['mean', 'sum', 'none']:
+                x = torch.tensor([1.0, 2.0], dtype=torch.float32, requires_grad=True, device=device)
+                y = torch.tensor([1.5, 2.5], dtype=torch.float64, requires_grad=True, device=device)
+                loss = F.mse_loss(x, y, reduction=reduction)
+                if reduction == 'mean':
+                    expected_loss = torch.tensor(0.25, dtype=torch.float64, device=device)
+                    expected_grad_x = torch.tensor([-0.5, -0.5], dtype=torch.float32, device=device)
+                    expected_grad_y = torch.tensor([0.5, 0.5], dtype=torch.float64, device=device)
+                elif reduction == 'sum':
+                    expected_loss = torch.tensor(0.5, dtype=torch.float64, device=device)
+                    expected_grad_x = torch.tensor([-1.0, -1.0], dtype=torch.float32, device=device)
+                    expected_grad_y = torch.tensor([1.0, 1.0], dtype=torch.float64, device=device)
+                else:  # 'none'
+                    expected_loss = torch.tensor([0.25, 0.25], dtype=torch.float64, device=device)
+                    expected_grad_x = torch.tensor([-1.0, -1.0], dtype=torch.float32, device=device)
+                    expected_grad_y = torch.tensor([1.0, 1.0], dtype=torch.float64, device=device)
+                self.assertEqual(loss, expected_loss, atol=1e-6, rtol=0)
+                grad = torch.ones_like(loss) if reduction == 'none' else None
+                loss.backward(grad)
+                self.assertEqual(x.grad, expected_grad_x, atol=1e-6, rtol=0)
+                self.assertEqual(y.grad, expected_grad_y, atol=1e-6, rtol=0)
+
     def test_weighted_l1_loss_with_weights(self):
         inputs = torch.tensor([1.0, 2.0, 3.0, 4.0], requires_grad=True)
         targets = torch.tensor([1.5, 2.5, 3.5, 4.5])
@@ -3986,22 +4014,26 @@ def test_rnn_check_device(self):
             hidden = torch.randn(correct_hidden_shape)
 
             # input and weights are not at the same device
-            with self.assertRaisesRegex(RuntimeError,
-                                        "Input and parameter tensors are not at the same device"):
+            rnn_param_device_msg = (
+                r"(?:Input and parameter tensors are not at the same device|"
+                r"Expected all tensors to be on the same device)"
+            )
+            with self.assertRaisesRegex(RuntimeError, rnn_param_device_msg):
                 model(input.to('cuda:0'))
-            with self.assertRaisesRegex(RuntimeError,
-                                        "Input and parameter tensors are not at the same device"):
+            with self.assertRaisesRegex(RuntimeError, rnn_param_device_msg):
                 model_cuda(input)
 
             # input and hiddens are not at the same device
-            with self.assertRaisesRegex(RuntimeError,
-                                        r"Input and hidden tensors are not at the same device"):
+            rnn_hidden_device_msg = (
+                r"(?:Input and hidden tensors are not at the same device|"
+                r"Expected all tensors to be on the same device)"
+            )
+            with self.assertRaisesRegex(RuntimeError, rnn_hidden_device_msg):
                 if mode == 'LSTM':
                     model(input, (hidden.to('cuda:0'), hidden.to('cuda:0')))
                 else:
                     model(input, (hidden.to('cuda:0')))
-            with self.assertRaisesRegex(RuntimeError,
-                                        r"Input and hidden tensors are not at the same device"):
+            with self.assertRaisesRegex(RuntimeError, rnn_hidden_device_msg):
                 if mode == 'LSTM':
                     model_cuda(input.to('cuda:0'), (hidden, hidden))
                 else:
@@ -4009,8 +4041,7 @@ def test_rnn_check_device(self):
 
             # hidden tensors are not at the same CUDA device
             if mode == 'LSTM':
-                with self.assertRaisesRegex(RuntimeError,
-                                            "Input and hidden tensors are not at the same device"):
+                with self.assertRaisesRegex(RuntimeError, rnn_hidden_device_msg):
                     model(input.to('cuda:0'), (hidden.to('cuda:0'), hidden.to('cuda:1')))
 
     @unittest.skipIf(not TEST_MULTIGPU, "multi-GPU not supported")
@@ -5457,6 +5488,13 @@ def func(x):
             # just run a single backward, as gradcheck/gradgradcheck is expensive here
             output.sum().backward()
 
+    # test for issue reported in https://github.com/pytorch/pytorch/issues/173799
+    def test_pdist_inf_nan_propagation(self):
+        x = torch.empty((2, 2), dtype=torch.float64).fill_(float("inf"))
+        for device in device_():
+            result = F.pdist(x.to(device), p=0.0)
+            self.assertTrue(torch.isnan(result).item(), "Expected NaN in pdist output")
+
     def test_cosine_embedding_loss_with_diff_type(self):
         for device in device_():
             input1 = torch.tensor([[2, 3, 4], [6, 2, 4]], dtype=torch.double, device=device)
@@ -6812,6 +6850,8 @@ def test_upsampling_small_scale(self):
         expected_out_t = torch.tensor([[[[2.5]]]])
         self.assertEqual(expected_out_t, out_t)
 
+    @xfailIf(IS_ARM64 and IS_CPU_EXT_SVE_SUPPORTED and not IS_CPU_CAPABILITY_SVE256)
+    # see https://github.com/pytorch/pytorch/issues/177250
     def test_upsampling_bfloat16(self, dtype=torch.bfloat16):
         def helper(size, scale_factor, mode, device, memory_format=torch.contiguous_format):
             input = torch.randn(size, device=device, dtype=dtype).to(memory_format=memory_format).detach().requires_grad_(True)
@@ -6945,6 +6985,77 @@ def helper(size, dtype, mode, device, is_channels_last):
                         helper(size, dtype, mode, device, is_channels_last)
 
 
+    @unittest.skipIf(IS_WINDOWS, "requires mmap/mprotect")
+    @unittest.skipUnless(TEST_NUMPY, "requires numpy")
+    def test_interpolate_uint8_overread(self):
+        # Regression test for vectorized resize overreads (NEON vld3_u8 and
+        # similar AVX2 paths). The vectorized block-of-4/8 loops may load
+        # more bytes than needed; on the last row this can read past the
+        # buffer. We detect this by placing tensor data right before an
+        # unmapped guard page so any overread triggers SIGBUS/SIGSEGV.
+
+        page_size = os.sysconf("SC_PAGE_SIZE")
+        libc = ctypes.CDLL(ctypes.util.find_library("c"))
+        libc.mmap.restype = ctypes.c_void_p
+        libc.mmap.argtypes = [
+            ctypes.c_void_p, ctypes.c_size_t, ctypes.c_int,
+            ctypes.c_int, ctypes.c_int, ctypes.c_longlong,
+        ]
+        libc.mprotect.restype = ctypes.c_int
+        libc.mprotect.argtypes = [ctypes.c_void_p, ctypes.c_size_t, ctypes.c_int]
+        libc.munmap.restype = ctypes.c_int
+        libc.munmap.argtypes = [ctypes.c_void_p, ctypes.c_size_t]
+
+        MAP_PRIVATE = 0x02
+        MAP_ANON = 0x1000 if IS_MACOS else 0x20
+        PROT_RW = 0x01 | 0x02
+        PROT_NONE = 0x00
+
+        for num_channels in (3, 4):  # 3ch → NEON on aarch64, 4ch → AVX2 on x86
+            for mode in ("bilinear", "bicubic"):
+                for w_in in range(4, 60):
+                    for w_out in [1, 2, 3]:
+                        h_in = 1
+                        tensor_bytes = h_in * w_in * num_channels
+                        num_data_pages = (tensor_bytes + page_size - 1) // page_size
+                        total_pages = num_data_pages + 1
+                        total_size = total_pages * page_size
+
+                        addr = libc.mmap(0, total_size, PROT_RW,
+                                         MAP_PRIVATE | MAP_ANON, -1, 0)
+                        self.assertTrue(
+                            addr not in (0, 2**64 - 1, -1), "mmap failed"
+                        )
+
+                        guard_start = addr + num_data_pages * page_size
+                        libc.mprotect(guard_start, page_size, PROT_NONE)
+
+                        data_start = guard_start - tensor_bytes
+                        ArrayType = ctypes.c_uint8 * tensor_bytes
+                        c_arr = ArrayType.from_address(data_start)
+                        np_arr = np.frombuffer(c_arr, dtype=np.uint8)
+                        np_arr[:] = 128
+
+                        flat = torch.from_numpy(np_arr)
+                        # channels-last strides: (H*W*C, 1, W*C, C)
+                        t = flat.as_strided(
+                            size=[1, num_channels, h_in, w_in],
+                            stride=[
+                                h_in * w_in * num_channels,
+                                1,
+                                w_in * num_channels,
+                                num_channels,
+                            ],
+                        )
+
+                        # Overread past the guard page will SIGBUS/SIGSEGV
+                        F.interpolate(
+                            t, size=(h_in, w_out), mode=mode,
+                            antialias=True, align_corners=False,
+                        )
+
+                        libc.munmap(addr, total_size)
+
     @set_default_dtype(torch.double)
     def test_interpolate(self):
         def _test_interpolate_non_integer_size_warning(in_t, out_size, dim, **kwargs):
@@ -7640,6 +7751,27 @@ def test_rnn_cell(cell_fn, gate_count):
         ]:
             test_rnn_cell(cell_fn, gate_count)
 
+    def test_conv3d_initialization_consistency(self):
+        # 1. Create a Conv3d layer
+        m = torch.nn.Conv3d(3, 6, kernel_size=3)
+
+        # 2. Test initialization in default (contiguous) format
+        torch.manual_seed(42)
+        m.reset_parameters()
+        weights_default = m.weight.clone().detach()
+
+        # 3. Test initialization in channels_last_3d format
+        m.to(memory_format=torch.channels_last_3d)
+        torch.manual_seed(42)
+        m.reset_parameters()
+        weights_channels_last = m.weight.clone().detach()
+
+        # 4. Verify both initializations are identical
+        self.assertEqual(
+            weights_default,
+            weights_channels_last,
+            msg="Conv3d initialization is inconsistent between memory formats"
+        )
 class TestFusionEval(TestCase):
     @set_default_dtype(torch.double)
     @given(X=hu.tensor(shapes=((5, 3, 5, 5),), dtype=np.double),
@@ -9362,6 +9494,49 @@ def group_norm_ref(X, gamma, beta, groups, channels, eps):
             Y_cpu = group_norm(X.cpu())
             self.assertEqual(Y_cpu, Y, rtol=0, atol=1e-5)
 
+    @onlyCUDA
+    @dtypes(torch.float32, torch.bfloat16)
+    def test_GroupNorm_backward_large_batch(self, device, dtype):
+        # Test GroupNorm backward with N > 128, which triggers
+        # GammaBetaBackwardCUDAKernel2. This kernel uses a (32, 16) block
+        # with WarpReduceSum after a shared memory transpose. On AMD
+        # wavefront-64, WarpReduceSum incorrectly summed across two tile
+        # columns, producing wrong dgamma/dbeta.
+        # bfloat16 tolerances are loose because the forward pass computes
+        # intermediate values (mean, var, x_norm) in bfloat16 vs float32 on CPU,
+        # causing accumulated differences in dgamma/dbeta.  The AMD wavefront-64
+        # bug this test targets produces ~100% error on all elements, so atol=1.0
+        # still catches it with wide margin.
+        rtol = 1.0 if dtype == torch.bfloat16 else 1e-3
+        atol = 1.0 if dtype == torch.bfloat16 else 1e-3
+        for N in [129, 256, 512]:
+            for C, G in [(32, 4), (64, 8), (128, 32)]:
+                x = torch.randn(N, C, 16, device=device, dtype=dtype, requires_grad=True)
+                gamma = torch.randn(C, device=device, dtype=dtype, requires_grad=True)
+                beta = torch.randn(C, device=device, dtype=dtype, requires_grad=True)
+
+                y = F.group_norm(x, G, gamma, beta)
+                grad = torch.randn_like(y)
+                y.backward(grad)
+                dgamma_gpu = gamma.grad.clone()
+                dbeta_gpu = beta.grad.clone()
+
+                # CPU reference in float32
+                x_cpu = x.detach().float().cpu().requires_grad_(True)
+                gamma_cpu = gamma.detach().float().cpu().requires_grad_(True)
+                beta_cpu = beta.detach().float().cpu().requires_grad_(True)
+                y_cpu = F.group_norm(x_cpu, G, gamma_cpu, beta_cpu)
+                y_cpu.backward(grad.float().cpu())
+
+                self.assertEqual(
+                    dgamma_gpu.float().cpu(), gamma_cpu.grad, atol=atol, rtol=rtol,
+                    msg=f"dgamma mismatch: N={N} C={C} G={G} dtype={dtype}",
+                )
+                self.assertEqual(
+                    dbeta_gpu.float().cpu(), beta_cpu.grad, atol=atol, rtol=rtol,
+                    msg=f"dbeta mismatch: N={N} C={C} G={G} dtype={dtype}",
+                )
+
     @expectedFailureMPS  # Double is not supported on MPS
     @onlyNativeDeviceTypes
     @dtypes(torch.float64, torch.complex128)
@@ -9914,11 +10089,7 @@ def test_linear_empty(self, device):
     def test_one_hot(self, device):
         # cuda throws device assert for invalid data
         # xla & mps ignore out of bound indices
-        if (
-            self.device_type != 'cuda'
-            and self.device_type != 'xla'
-            and self.device_type != 'mps'
-        ):
+        if self.device_type == 'cpu':
             with self.assertRaises(RuntimeError):
                 torch.nn.functional.one_hot(torch.tensor([3, 4, -1, 0], device=device), -1)
 
@@ -10512,7 +10683,7 @@ def test_upsamplingNearestExact3d_correctness(self, device, memory_format, isize
 
     @parametrize_test("antialias", [True, False])
     @parametrize_test("align_corners", [True, False])
-    @parametrize_test("mode", ["bilinear", "bicubic"])
+    @parametrize_test("mode", ["bilinear", "bicubic", "lanczos"])
     @parametrize_test("memory_format", [torch.contiguous_format, torch.channels_last])
     @expectedFailureMPS  # double device type
     @onlyNativeDeviceTypes
@@ -10520,6 +10691,14 @@ def test_upsamplingBiMode2d(self, device, antialias, align_corners, mode, memory
         # Forward AD does not support XLA because XLA tensors don't have storage
         check_forward_ad = torch.device(device).type != 'xla'
 
+        if mode == "lanczos":
+            if torch.device(device).type != "cpu":
+                raise SkipTest("Lanczos mode is only supported on CPU")
+            if not antialias:
+                raise SkipTest("Lanczos mode requires antialias=True")
+            if align_corners:
+                raise SkipTest("Lanczos mode does not support align_corners=True")
+
         kwargs = dict(mode=mode, align_corners=align_corners, antialias=antialias)
         # test float scale factor up & downsampling
         for scale_factor in [0.5, 1.5, 2]:
@@ -10582,7 +10761,7 @@ def test_upsamplingBiMode2d(self, device, antialias, align_corners, mode, memory
 
     @parametrize_test("antialias", [True, False])
     @parametrize_test("num_channels", [3, 5])
-    @parametrize_test("mode", ["nearest", "nearest-exact", "bilinear", "bicubic"])
+    @parametrize_test("mode", ["nearest", "nearest-exact", "bilinear", "bicubic", "lanczos"])
     @parametrize_test("dtype", integral_types() + floating_types())
     @skipIfMPS  # Error message is wrong for some dtypes
     @onlyNativeDeviceTypes
@@ -10597,6 +10776,14 @@ def test_upsamplingBiMode2d_nonsupported_dtypes(self, device, antialias, num_cha
             if dtype in (torch.uint8, ) + floating_types():
                 should_raise_runtime_error = False
 
+        elif mode == "lanczos":
+            if torch.device(device).type != "cpu":
+                raise SkipTest("Lanczos mode is only supported on CPU")
+            if not antialias:
+                raise SkipTest("Lanczos mode requires antialias=True")
+            if dtype in floating_types() or (device == "cpu" and dtype == torch.uint8):
+                should_raise_runtime_error = False
+
         elif mode in ("bilinear", "bicubic"):
             if dtype in floating_types() or (device == "cpu" and dtype == torch.uint8):
                 should_raise_runtime_error = False
@@ -10630,7 +10817,7 @@ def test_upsamplingBilinear2d_aa_correctness(self, device, memory_format):
     # Partially passes. NotImplementedError: aten::upsample_bicubic2d.out https://github.com/pytorch/pytorch/issues/77764
     @skipIfMPS
     @parametrize_test("memory_format", [torch.contiguous_format, torch.channels_last])
-    @parametrize_test("mode", ["bilinear", "bicubic"])
+    @parametrize_test("mode", ["bilinear", "bicubic", "lanczos"])
     @parametrize_test("antialias", [True, False])
     @parametrize_test("align_corners", [True, False])
     @parametrize_test("num_channels", [3, 5])
@@ -10655,13 +10842,19 @@ def test_upsamplingBiMode2d_consistency(
         if torch.device(device).type == "cuda":
             raise SkipTest("CUDA implementation is not yet supporting uint8")
 
+        if mode == "lanczos":
+            if not antialias:
+                raise SkipTest("Lanczos mode requires antialias=True")
+            if align_corners:
+                raise SkipTest("Lanczos mode does not support align_corners=True")
+
         torch.manual_seed(0)
 
-        # - input range is set to [30, 220] for bicubic mode, because the bicubic kernel may create
-        #   [intermediate] values outside of the [0, 255] range, which need
-        #   to be clipped in uint8 path, but not in float path. This isn't
-        #   an issue with bilinear kernel.
-        input_range = (30, 220) if mode == "bicubic" else (0, 256)
+        # - input range is set to [30, 220] for bicubic and lanczos modes,
+        #   because these kernels may create [intermediate] values outside of
+        #   the [0, 255] range, which need to be clipped in uint8 path, but
+        #   not in float path. This isn't an issue with bilinear kernel.
+        input_range = (30, 220) if mode in ("bicubic", "lanczos") else (0, 256)
         input_ui8 = torch.randint(*input_range, size=(batch_size, num_channels, 400, 400), dtype=torch.uint8, device=device)
         input_ui8 = input_ui8.contiguous(memory_format=memory_format)
 
@@ -10699,6 +10892,7 @@ def test_upsamplingBiMode2d_consistency(
         if mode == "bilinear":
             torch.testing.assert_close(output_f32, output_ui8.float(), rtol=0, atol=1)
         else:
+            # bicubic and lanczos
             diff = (output_f32 - output_ui8.float()).abs()
             self.assertLess(diff.max(), 15)
 
@@ -10767,6 +10961,78 @@ def test_upsamplingBicubic2d_aa_correctness(self, device, memory_format):
         t_out = F.interpolate(t_in, size=(2, 2), mode="bicubic", align_corners=False, antialias=True)
         self.assertEqual(expected_out, t_out)
 
+    @onlyCPU
+    @parametrize_test("memory_format", [torch.contiguous_format, torch.channels_last])
+    def test_upsamplingLanczos2d_aa_correctness(self, device, memory_format):
+        t_in = torch.arange(3 * 8 * 8, dtype=torch.float, device=device).reshape(1, 3, 8, 8)
+        t_in = t_in.contiguous(memory_format=memory_format)
+        # This expected result is obtained using PIL.Image.resize
+        # for c in range(3):
+        #   a_in = t_in.numpy()[0, c, ...]
+        #   pil_in = Image.fromarray(a_in)
+        #   pil_out = pil_in.resize((2, 2), resample=Image.LANCZOS)
+        expected_out = torch.tensor([
+            14.267621, 18.097038, 44.902962, 48.732376, 78.267616, 82.097038,
+            108.902962, 112.732384, 142.267624, 146.097031, 172.902969, 176.732376
+        ], device=device, dtype=t_in.dtype).reshape(1, 3, 2, 2)
+        t_out = F.interpolate(t_in, size=(2, 2), mode="lanczos", align_corners=False, antialias=True)
+        self.assertEqual(expected_out, t_out)
+
+    @onlyCPU
+    def test_upsamplingLanczos2d_errors(self, device):
+        # 3D input (1D spatial) not supported
+        x_3d = torch.randn(1, 3, 8, device=device)
+        with self.assertRaisesRegex(ValueError, "4-D tensor"):
+            F.interpolate(x_3d, size=(4,), mode="lanczos", antialias=True)
+
+        # 5D input (3D spatial) not supported
+        x_5d = torch.randn(1, 3, 8, 8, 8, device=device)
+        with self.assertRaisesRegex(ValueError, "4-D tensor"):
+            F.interpolate(x_5d, size=(4, 4, 4), mode="lanczos", antialias=True)
+
+        # antialias=False not supported
+        x_4d = torch.randn(1, 3, 8, 8, device=device)
+        with self.assertRaisesRegex(ValueError, "antialias=True"):
+            F.interpolate(x_4d, size=(4, 4), mode="lanczos", antialias=False)
+
+        # align_corners=True not supported
+        with self.assertRaisesRegex(ValueError, "align_corners=True"):
+            F.interpolate(x_4d, size=(4, 4), mode="lanczos", align_corners=True, antialias=True)
+
+    @onlyCPU
+    def test_upsamplingLanczos2d_identity(self, device):
+        x = torch.randn(1, 3, 8, 8, device=device)
+        out = F.interpolate(x, size=(8, 8), mode="lanczos", align_corners=False, antialias=True)
+        self.assertEqual(x, out)
+
+    @onlyCUDA
+    def test_upsamplingBicubic2d_many_channels(self, device):
+        # Exercises the parallelized batch/channel kernel for small spatial
+        # sizes with many channels, typical in VLM position embeddings.
+        for batch, channels, in_h, in_w, out_h, out_w in [
+            (64, 768, 16, 16, 6, 6),
+            (8, 1152, 32, 32, 14, 14),
+            (2, 256, 8, 8, 16, 16),
+        ]:
+            in_t = torch.randn(batch, channels, in_h, in_w, device=device)
+            # Compute via single-element batch to get reference without
+            # the parallel kernel (batch=1, channels loop is short)
+            expected_parts = []
+            for n in range(batch):
+                expected_parts.append(
+                    F.interpolate(
+                        in_t[n : n + 1],
+                        size=(out_h, out_w),
+                        mode="bicubic",
+                        align_corners=False,
+                    )
+                )
+            expected = torch.cat(expected_parts, dim=0)
+            out_t = F.interpolate(
+                in_t, size=(out_h, out_w), mode="bicubic", align_corners=False
+            )
+            self.assertEqual(out_t, expected)
+
     @expectedFailureMPS  # NotImplementedError: aten::upsample_trilinear3d.out https://github.com/pytorch/pytorch/issues/77764
     @parametrize_test("align_corners", [True, False])
     @parametrize_test("memory_format", [torch.contiguous_format, torch.channels_last_3d])
@@ -10806,16 +11072,19 @@ def test_upsamplingTrilinear3d(self, device, align_corners, memory_format):
             gradgradcheck(lambda x: F.interpolate(x, out_size, **kwargs), [input])
 
     @onlyCUDA
-    @skipCUDAIfRocm(msg="launch bounds error out on ROCM")
     @dtypes(torch.half, torch.bfloat16)
     @largeTensorTest('40GB')
     def test_upsampling_64bit_indexing_channels_last(self, device, dtype):
+        # Path 1 (NHWC): channels-last 1D grid. output.numel() = 2^31, below UINT32_MAX,
+        # but exercises the allclose correctness check between channels-last and contiguous.
         x = torch.rand((32, 64, 512, 512), dtype=dtype, device=device)
         out = torch.nn.functional.interpolate(x.to(memory_format=torch.channels_last), scale_factor=2, mode='nearest')
         out_ref = torch.nn.functional.interpolate(x, scale_factor=2, mode='nearest')
         del x
         self.assertTrue(torch.allclose(out, out_ref))
 
+        # Path 1 (NHWC): output.numel() = 17*256*1024*1024 ~ 4.56e9 > UINT32_MAX.
+        # Exercises the ROCm grid-stride clamp in the NHWC kernel.
         x = torch.ones((17, 256, 512, 512), dtype=dtype).cuda().to(memory_format=torch.channels_last)
         out = torch.nn.functional.interpolate(x, scale_factor=2, mode='nearest')
         self.assertEqual(out[0], out[-1])
@@ -11647,7 +11916,7 @@ def test_upsamplingNearest3d_launch_config(self, device):
         self.assertEqual(out_ref, out)
 
     @unittest.expectedFailure
-    @skipIfRocm
+    @skipIfRocmVersionLessThan((7, 14))
     @onlyCUDA
     def test_upsamplingNearest2d_launch_fail(self, device):
         m = nn.Upsample(scale_factor=2)
@@ -11655,14 +11924,6 @@ def test_upsamplingNearest2d_launch_fail(self, device):
         inp = torch.rand(1, 1, 2**15, 2**8, device=device)
         out = m(inp)
 
-    @onlyCUDA
-    @skipCUDAIfNotRocm
-    def test_upsamplingNearest2d_launch_rocm(self, device):
-        # test_upsamplingNearest2d_launch_fail should run OK on ROCm
-        m = nn.Upsample(scale_factor=2)
-        inp = torch.rand(1, 1, 2**15, 2**8, device=device)
-        out = m(inp)
-
     @onlyCUDA
     def test_CTCLoss_cudnn(self, device):
         def _helper(zero_infinity):
@@ -12711,7 +12972,9 @@ def test_cross_entropy_loss_2d_out_of_bounds_class_index(self):
         # CUDA says "device-side assert triggered"
         # ROCm says "unspecified launch failure", or HSA_STATUS_ERROR_EXCEPTION
         has_cuda_assert = 'CUDA error: device-side assert triggered' in stderr
-        has_hip_assert = 'launch failure' in stderr or 'HSA_STATUS_ERROR_EXCEPTION' in stderr
+        has_hip_assert = ('launch failure' in stderr
+                          or 'HSA_STATUS_ERROR_EXCEPTION' in stderr
+                          or 'illegal memory access' in stderr)
         self.assertTrue(has_cuda_assert or has_hip_assert,
                         f"Expected device assert error in stderr, got: {stderr}")
 
diff --git a/test/test_numa_binding.py b/test/test_numa_binding.py
index c38a477d56e48..53cdaa0248bb3 100644
--- a/test/test_numa_binding.py
+++ b/test/test_numa_binding.py
@@ -5,7 +5,7 @@
 import json
 import sys
 from dataclasses import dataclass
-from typing import Any, Optional
+from typing import Any
 from unittest import skipUnless
 from unittest.mock import mock_open, patch
 
@@ -230,7 +230,7 @@ def _mock_sched_getaffinity(self, pid: int) -> set[int]:
         return set(range(self._mock_num_logical_cpus))
 
     def _start_processes_for_str_entrypoint_and_get_command_args(
-        self, *, numa_options: Optional[NumaOptions], target_local_rank: int
+        self, *, numa_options: NumaOptions | None, target_local_rank: int
     ) -> tuple[str, ...]:
         with patch(
             "torch.distributed.elastic.multiprocessing.subprocess_handler.subprocess_handler.Popen"
@@ -256,8 +256,8 @@ def _start_processes_for_str_entrypoint_and_get_command_args(
             return call_args.kwargs["args"]
 
     def _start_processes_for_callable_entrypoint_and_get_sched_setaffinity_cpus(
-        self, *, numa_options: Optional[NumaOptions], target_local_rank: int
-    ) -> Optional[set[int]]:
+        self, *, numa_options: NumaOptions | None, target_local_rank: int
+    ) -> set[int] | None:
         target_sched_setaffinity_logical_cpu_indices = None
 
         def mock_sched_setaffinity(*args, **kwargs) -> None:
diff --git a/test/test_opaque_obj_v2.py b/test/test_opaque_obj_v2.py
index abe6fd404fa97..2d09fec46cab7 100644
--- a/test/test_opaque_obj_v2.py
+++ b/test/test_opaque_obj_v2.py
@@ -1,15 +1,18 @@
 # Owner(s): ["module: custom-operators"]
 
 import contextlib
+import enum
 import gc
 import random
+import re
+import unittest
 from contextlib import ExitStack
 from dataclasses import dataclass
-from typing import Optional
 
 import torch
+import torch.distributed as dist
 import torch.utils._pytree as pytree
-from torch._dynamo.test_case import run_tests, TestCase
+from torch._dynamo.functional_export import _dynamo_graph_capture_for_export
 from torch._dynamo.testing import (
     AotEagerAndRecordGraphs,
     CompileCounter,
@@ -20,14 +23,15 @@
 )
 from torch._dynamo.utils import counters, counters as dynamo_counters
 from torch._functorch import config as functorch_config
-from torch._functorch._aot_autograd.autograd_cache import AOTAutogradCache
 from torch._functorch.aot_autograd import (
     aot_compile_joint_with_descriptors,
     aot_export_joint_with_descriptors,
     aot_export_module,
 )
 from torch._inductor import config as inductor_config
-from torch._inductor.compile_fx import compile_fx
+from torch._inductor.compile_fx import compile_fx, compile_fx_inner
+from torch._inductor.test_case import run_tests, TestCase
+from torch._inductor.utils import fresh_inductor_cache, run_and_get_code
 from torch._library.effects import EffectType
 from torch._library.fake_class_registry import FakeScriptObject, maybe_to_fake_obj
 from torch._library.opaque_object import (
@@ -41,6 +45,7 @@
     register_opaque_type,
 )
 from torch._subclasses.fake_tensor import FakeTensorMode
+from torch.fx._graph_pickler import GraphPickler, Options
 from torch.fx.experimental.proxy_tensor import make_fx
 from torch.fx.experimental.symbolic_shapes import ShapeEnv
 from torch.fx.graph import _illegal_char_regex
@@ -48,6 +53,8 @@
     instantiate_parametrized_tests,
     parametrize,
 )
+from torch.testing._internal.inductor_utils import HAS_GPU
+from torch.utils._import_utils import import_dill
 
 
 class Color(OpaqueBase):
@@ -341,11 +348,7 @@ def _(x, s):
 )
 register_opaque_type(AddModule, typ="reference")
 register_opaque_type(ValueConfig, typ="value")
-register_opaque_type(
-    SizeStore,
-    typ="value",
-    members={"size": MemberType.USE_REAL, "increment_size": MemberType.USE_REAL},
-)
+register_opaque_type(SizeStore, typ="value")
 register_opaque_type(NestedValueSize, typ="value")
 register_opaque_type(OpaqueMultiplier, typ="reference")
 register_opaque_type(Color, typ="reference")
@@ -386,12 +389,13 @@ def __repr__(self):
         return f"TensorWithCounter({self.a}, {self.b}, {self._counter}, {self._size_store})"
 
     def __tensor_flatten__(self):
-        return ["a", "b"], (self._counter, self._size_store)
+        return ["a", "b", "_counter"], (self._size_store,)
 
     @staticmethod
     def __tensor_unflatten__(inner_tensors, ctx, outer_size, outer_stride):
         a, b = inner_tensors["a"], inner_tensors["b"]
-        counter, size_store = ctx
+        counter = inner_tensors["_counter"]
+        (size_store,) = ctx
         return TensorWithCounter(a, b, counter, size_store, outer_size, outer_stride)
 
     @classmethod
@@ -621,7 +625,7 @@ def process_nested_config_fake(
             lib=self.lib,
         )
         def process_multiple_sizes_impl(
-            x: torch.Tensor, config: Optional[list[SizeStore]]
+            x: torch.Tensor, config: list[SizeStore] | None
         ) -> torch.Tensor:
             if config is None:
                 return x.clone()
@@ -635,7 +639,7 @@ def process_multiple_sizes_impl(
             "_TestOpaqueObject::process_multiple_sizes", lib=self.lib
         )
         def process_multiple_sizes_fake(
-            x: torch.Tensor, config: Optional[list[SizeStore]]
+            x: torch.Tensor, config: list[SizeStore] | None
         ) -> torch.Tensor:
             return torch.empty_like(x)
 
@@ -808,6 +812,78 @@ def backward(ctx, grad_output: torch.Tensor):
             lib=self.lib,
         )
 
+        torch.library.define(
+            "_TestOpaqueObject::create_multiplier",
+            f"(Tensor scale) -> {opaque_multiplier_type}",
+            tags=torch.Tag.pt2_compliant_tag,
+            lib=self.lib,
+        )
+
+        @torch.library.impl(
+            "_TestOpaqueObject::create_multiplier",
+            "CompositeExplicitAutograd",
+            lib=self.lib,
+        )
+        def create_multiplier_impl(scale: torch.Tensor) -> OpaqueMultiplier:
+            return OpaqueMultiplier(scale.item())
+
+        @torch.library.register_fake(
+            "_TestOpaqueObject::create_multiplier", lib=self.lib
+        )
+        def create_multiplier_fake(scale: torch.Tensor) -> OpaqueMultiplier:
+            return OpaqueMultiplier(0.0)
+
+        counter_type = get_opaque_type_name(Counter)
+        torch.library.define(
+            "_TestOpaqueObject::create_counter",
+            f"(SymInt start, SymInt end) -> {counter_type}",
+            tags=torch.Tag.pt2_compliant_tag,
+            lib=self.lib,
+        )
+
+        @torch.library.impl(
+            "_TestOpaqueObject::create_counter",
+            "CompositeExplicitAutograd",
+            lib=self.lib,
+        )
+        def create_counter_impl(start: int, end: int) -> Counter:
+            return Counter(start, end)
+
+        @torch.library.register_fake("_TestOpaqueObject::create_counter", lib=self.lib)
+        def create_counter_fake(start: int, end: int) -> Counter:
+            return Counter(start, end)
+
+        @torch.library.custom_op(
+            "_TestOpaqueObject::create_counter2",
+            mutates_args=[],
+        )
+        def create_counter2_impl(start: int, end: int) -> tuple[Counter, Counter]:
+            return Counter(start, end), Counter(end, start)
+
+        @create_counter2_impl.register_fake
+        def create_counter2_fake(start: int, end: int) -> tuple[Counter, Counter]:
+            return Counter(start, end), Counter(end, start)
+
+        counter_type = get_opaque_type_name(Counter)
+        torch.library.define(
+            "_TestOpaqueObject::counter_start",
+            f"({counter_type} a) -> Tensor",
+            tags=torch.Tag.pt2_compliant_tag,
+            lib=self.lib,
+        )
+
+        @torch.library.impl(
+            "_TestOpaqueObject::counter_start",
+            "CompositeExplicitAutograd",
+            lib=self.lib,
+        )
+        def counter_start_impl(a: Counter) -> torch.Tensor:
+            return torch.scalar_tensor(a.start, dtype=torch.int64)
+
+        @torch.library.register_fake("_TestOpaqueObject::counter_start", lib=self.lib)
+        def counter_start_fake(a: Counter) -> torch.Tensor:
+            return torch.scalar_tensor(0, dtype=torch.int64)
+
         super().setUp()
 
     def tearDown(self):
@@ -834,6 +910,51 @@ def test_ops(self):
         size = torch.ops._TestOpaqueObject.queue_size(queue)
         self.assertEqual(size, 0)
 
+    def test_custom_op_reference_return(self):
+        @torch.compile(backend="inductor", fullgraph=True)
+        def f(x):
+            c1 = torch.ops._TestOpaqueObject.create_counter(0, 10)
+            c2, c3 = torch.ops._TestOpaqueObject.create_counter2(0, 10)
+            return (
+                c1,
+                torch.ops._TestOpaqueObject.increment_counter(c1, x),
+                c2,
+                torch.ops._TestOpaqueObject.increment_counter(c2, x),
+                c3,
+                torch.ops._TestOpaqueObject.increment_counter(c3, x),
+            )
+
+        x = torch.scalar_tensor(0, dtype=torch.int64)
+        c1, r1, c2, r2, c3, r3 = f(x)
+        self.assertIsInstance(c1, Counter)
+        self.assertIsInstance(c2, Counter)
+        self.assertIsInstance(c3, Counter)
+        self.assertEqual(c3.start, 11)
+        self.assertEqual(c3.end, 0)
+        self.assertEqual(x, torch.scalar_tensor(10, dtype=torch.int64))
+        self.assertEqual(r1, torch.scalar_tensor(1, dtype=torch.int64))
+        self.assertEqual(r2, torch.scalar_tensor(1, dtype=torch.int64))
+        self.assertEqual(r3, torch.scalar_tensor(11, dtype=torch.int64))
+
+    def test_custom_op_value_return(self):
+        with self.assertRaisesRegex(ValueError, "Return has unsupported type"):
+
+            @torch.library.custom_op(
+                "_TestOpaqueObject::create_value_config",
+                mutates_args=[],
+            )
+            def create_value_config(mode: str) -> ValueConfig:
+                return ValueConfig(mode)
+
+        with self.assertRaisesRegex(ValueError, "Return has unsupported type"):
+
+            @torch.library.custom_op(
+                "_TestOpaqueObject::create_value_config_tuple",
+                mutates_args=[],
+            )
+            def create_value_config_tuple(mode: str) -> tuple[ValueConfig, int]:
+                return ValueConfig(mode), 0
+
     def test_fake_script_object_isinstance_per_type(self):
         queue = OpaqueQueue([], torch.zeros(3))
         rng = RNGState(42)
@@ -898,6 +1019,132 @@ def forward(self, arg0_1, arg1_1):
     """,
         )
 
+    @parametrize("make_fx_tracing_mode", ["fake", "symbolic"])
+    def test_make_fx_value_type(self, make_fx_tracing_mode):
+        def f(x, cfg):
+            return torch.ops._TestOpaqueObject.process_with_config(x, cfg)
+
+        x = torch.randn(3, 3)
+        cfg = ValueConfig("square")
+        gm = make_fx(f, tracing_mode=make_fx_tracing_mode)(x, cfg)
+        self.assertEqual(gm(x, cfg), f(x, cfg))
+
+        self.assertExpectedInline(
+            gm.code.strip("\n"),
+            """\
+def forward(self, x_1, cfg_1):
+    process_with_config = torch.ops._TestOpaqueObject.process_with_config.default(x_1, ValueConfig(mode='square'));  x_1 = None
+    return process_with_config
+    """,  # noqa: B950
+        )
+
+    def test_subclass_opaque_output_reuses_input_proxy(self):
+        # Regression test: when a tensor subclass's __torch_dispatch__ wraps
+        # the output with the real OpaqueBase (not the FakeScriptObject proxy),
+        # the AOTAutograd forward graph should still reference the opaque via
+        # its input placeholder — not create a duplicate get_attr constant.
+        #
+        # This mirrors DTensor where C++ dispatch creates new DTensors storing
+        # the real DeviceMesh, so output flattening calls maybe_to_fake_obj
+        # and mints a fresh FakeScriptObject for the same underlying object.
+
+        class TensorWithRealCounter(torch.Tensor):
+            @staticmethod
+            def __new__(cls, data, counter):
+                return torch.Tensor._make_wrapper_subclass(
+                    cls,
+                    data.size(),
+                    strides=data.stride(),
+                    storage_offset=data.storage_offset(),
+                    device=data.device,
+                    dtype=data.dtype,
+                    layout=data.layout,
+                    requires_grad=data.requires_grad,
+                )
+
+            def __init__(self, data, counter):
+                self._data = data
+                self._counter = counter
+
+            def __repr__(self):
+                return "TensorWithRealCounter(...)"
+
+            def __tensor_flatten__(self):
+                return ["_data", "_counter"], ()
+
+            @staticmethod
+            def __tensor_unflatten__(inner_tensors, ctx, outer_size, outer_stride):
+                return TensorWithRealCounter(
+                    inner_tensors["_data"], inner_tensors["_counter"]
+                )
+
+            @classmethod
+            def __torch_dispatch__(cls, func, types, args, kwargs):
+                if kwargs is None:
+                    kwargs = {}
+
+                counter = None
+                for arg in torch.utils._pytree.tree_leaves(args):
+                    if isinstance(arg, TensorWithRealCounter):
+                        counter = arg._counter
+                        break
+
+                def unwrap(x):
+                    return x._data if isinstance(x, TensorWithRealCounter) else x
+
+                out = func(
+                    *torch.utils._pytree.tree_map(unwrap, args),
+                    **torch.utils._pytree.tree_map(unwrap, kwargs),
+                )
+
+                # Unwrap FakeScriptObject to real OpaqueBase, simulating what
+                # happens in DTensor's C++ dispatch path.
+                real_counter = counter
+                if isinstance(counter, FakeScriptObject):
+                    real_counter = counter.real_obj
+
+                return torch.utils._pytree.tree_map(
+                    lambda x: TensorWithRealCounter(x, real_counter)
+                    if isinstance(x, torch.Tensor)
+                    else x,
+                    out,
+                )
+
+        counter = Counter(start=0, end=10)
+        x = TensorWithRealCounter(torch.randn(4), counter)
+
+        backend = AotEagerAndRecordGraphs()
+        torch.compile(lambda x: x + 1, fullgraph=True, backend=backend)(x)
+
+        fw = backend.fw_graphs[0]
+        get_attr_nodes = [n for n in fw.graph.nodes if n.op == "get_attr"]
+        self.assertEqual(
+            get_attr_nodes,
+            [],
+            "Opaque output should reuse the input placeholder, not create a get_attr constant",
+        )
+
+    def test_guard_pickle_subclass_with_opaque_inner_attr(self):
+        # Regression test: the guard state pickler serializes tensor subclasses
+        # by iterating over __tensor_flatten__ inner attrs.  Opaque inner attrs
+        # (e.g. Counter) must be handled correctly — they are pickled by the
+        # normal pickle machinery alongside tensor inner attrs.
+        from torch._dynamo.guards import GuardsStatePickler
+
+        a = torch.randn(4)
+        b = torch.randn(4)
+        counter = Counter(start=0, end=10)
+        size = SizeStore(4)
+        x = TensorWithCounter(a, b, counter, size)
+
+        import io
+
+        buf = io.BytesIO()
+        pickler = GuardsStatePickler({id(x): x}, {}, {}, buf)
+        func, args = pickler.reducer_override(x)
+        obj = func(*args)
+        self.assertIsInstance(obj, torch.Tensor)
+
     @parametrize("make_fx_tracing_mode", ["fake", "symbolic"])
     def test_bad_fake(self, make_fx_tracing_mode):
         torch.library.define(
@@ -1001,6 +1248,25 @@ def forward(self, arg0_1, arg1_1, arg2_1):
                 "_TestOpaqueObject::noisy_inject", None
             )
 
+    def test_install_free_opaque_object(self):
+        rng = RNGState(0)
+        x = torch.ones(2, 3)
+
+        def fn(x):
+            return torch.ops._TestOpaqueObject.noisy_inject(x, rng)
+
+        gm = _dynamo_graph_capture_for_export(fn)(x)
+        self.assertExpectedInline(
+            gm.code.strip(),
+            """\
+def forward(self, x):
+    arg_0, = fx_pytree.tree_flatten_spec(([x], {}), self._in_spec)
+    l_flat_args_0_ = arg_0
+    l__self____export_root___closure___0_cell_contents = self.L__self____export_root___closure___0_cell_contents
+    res = torch.ops._TestOpaqueObject.noisy_inject(l_flat_args_0_, l__self____export_root___closure___0_cell_contents);  l_flat_args_0_ = l__self____export_root___closure___0_cell_contents = None
+    return pytree.tree_unflatten((res,), self._out_spec)""",  # noqa: B950
+        )
+
     def test_compile1(self):
         def foo(rng_state, x):
             if not isinstance(rng_state, RNGState):
@@ -1310,23 +1576,107 @@ def foo(counter, x):
                 NestedCounters(Counter(1, 5)), torch.ones(2, 3)
             )
 
-        config = ValueConfig("double")
+    def test_compile_fixed_stride_order(self):
+        hs_name = get_opaque_type_name(HoistedString)
 
-        def foo(mode, x):
-            return config.mode
+        torch.library.define(
+            "_TestOpaqueObject::stride_op",
+            f"(Tensor x, {hs_name} s) -> Tensor",
+            tags=(torch.Tag.needs_fixed_stride_order,),
+            lib=self.lib,
+        )
 
-        with self.assertRaisesRegex(
-            RuntimeError, "Attempted to access unregistered member on an OpaqueObject"
-        ):
-            torch.compile(foo, backend="eager")(config, torch.ones(2, 3))
+        @torch.library.impl(
+            "_TestOpaqueObject::stride_op",
+            "CompositeExplicitAutograd",
+            lib=self.lib,
+        )
+        def stride_op_impl(x: torch.Tensor, s: HoistedString) -> torch.Tensor:
+            return x * 2.0
 
-        def bar(mode, x):
-            config.print_mode()
+        @torch.library.register_fake("_TestOpaqueObject::stride_op", lib=self.lib)
+        def stride_op_fake(x, s):
+            return torch.empty_like(x)
 
-        with self.assertRaisesRegex(
-            RuntimeError, "Attempted to access unregistered member on an OpaqueObject"
-        ):
-            torch.compile(bar, backend="eager")(config, torch.ones(2, 3))
+        def fn(x, s):
+            return torch.ops._TestOpaqueObject.stride_op(x, s)
+
+        s = HoistedString("double")
+        x = torch.randn(4, 4)
+
+        compiled_fn = torch.compile(fn, backend="inductor")
+        result = compiled_fn(x, s)
+
+        expected = x * 2.0
+        self.assertEqual(result, expected)
+
+    def test_compile_exact_strides(self):
+        hs_name = get_opaque_type_name(HoistedString)
+
+        torch.library.define(
+            "_TestOpaqueObject::exact_op",
+            f"(Tensor x, {hs_name} s) -> Tensor",
+            tags=(torch.Tag.needs_exact_strides,),
+            lib=self.lib,
+        )
+
+        @torch.library.impl(
+            "_TestOpaqueObject::exact_op",
+            "CompositeExplicitAutograd",
+            lib=self.lib,
+        )
+        def exact_op_impl(x: torch.Tensor, s: HoistedString) -> torch.Tensor:
+            return x * 3.0
+
+        @torch.library.register_fake("_TestOpaqueObject::exact_op", lib=self.lib)
+        def exact_op_fake(x, s):
+            return torch.empty_like(x)
+
+        def fn(x, s):
+            return torch.ops._TestOpaqueObject.exact_op(x, s)
+
+        s = HoistedString("double")
+        x = torch.randn(4, 4)
+
+        compiled_fn = torch.compile(fn, backend="inductor")
+        result = compiled_fn(x, s)
+
+        expected = x * 3.0
+        self.assertEqual(result, expected)
+
+    def test_compile_contiguous_strides(self):
+        hs_name = get_opaque_type_name(HoistedString)
+
+        torch.library.define(
+            "_TestOpaqueObject::contig_op",
+            f"(Tensor x, {hs_name} s) -> Tensor",
+            tags=(torch.Tag.needs_contiguous_strides,),
+            lib=self.lib,
+        )
+
+        @torch.library.impl(
+            "_TestOpaqueObject::contig_op",
+            "CompositeExplicitAutograd",
+            lib=self.lib,
+        )
+        def contig_op_impl(x: torch.Tensor, s: HoistedString) -> torch.Tensor:
+            return x * 4.0
+
+        @torch.library.register_fake("_TestOpaqueObject::contig_op", lib=self.lib)
+        def contig_op_fake(x, s):
+            return torch.empty_like(x)
+
+        def fn(x, s):
+            return torch.ops._TestOpaqueObject.contig_op(x, s)
+
+        s = HoistedString("double")
+        x = torch.randn(4, 4)
+
+        compiled_fn = torch.compile(fn, backend="inductor")
+        result = compiled_fn(x, s)
+
+        expected = x * 4.0
+        self.assertEqual(result, expected)
 
     def test_export_joint(self):
         torch.library.define(
@@ -1558,6 +1908,20 @@ def foo(x, cfg):
         self.assertEqual(res, x + x)
         self.assertEqual(cnt.frame_count, 2)
 
+    @parametrize("backend", ["eager", "aot_eager", "inductor"])
+    def test_value_type_graph_output(self, backend):
+        def foo(x):
+            return x * x, ValueConfig("square")
+
+        x = torch.randn(3, 3)
+        opt_f = torch.compile(foo, fullgraph=True, backend=backend)
+        res = opt_f(x)
+        self.assertEqual(res[1], ValueConfig("square"))
+
+        gm = _dynamo_graph_capture_for_export(foo)(x)
+        res = gm(x)
+        self.assertEqual(res[1], ValueConfig("square"))
+
     def test_value_type_graph_input(self):
         # Even though cfg is an input, it should not be an input to the dynamo
         # graph. Instead it should directly put in the graph argument as a
@@ -1661,6 +2025,26 @@ def forward(self, arg0_1):
     return (add,)""",  # noqa: B950
         )
 
+    def test_value_type_unregistered_method(self):
+        # Unregistered methods on value types should inline (no error)
+        def foo(x):
+            cfg = ValueConfig("square")
+            return x + len(cfg.mode)
+
+        x = torch.randn(3)
+        backend = AotEagerAndRecordGraphs()
+        opt_f = torch.compile(foo, fullgraph=True, backend=backend)
+        res = opt_f(x)
+        self.assertEqual(res, foo(x))
+
+        self.assertExpectedInline(
+            backend.fw_graphs[0].code.strip(),
+            """\
+def forward(self, arg0_1):
+    add = torch.ops.aten.add.Tensor(arg0_1, 6);  arg0_1 = None
+    return (add,)""",
+        )
+
     def test_weakref_cleanup(self):
         def register_tmp_class():
             class TmpClass(OpaqueBase):
@@ -1848,6 +2232,653 @@ def forward(self, L_x_: "TensorWithCounter(f32[4, 4])"):
         # Recompile since SizeStore has changed
         self.assertEqual(cnt.frame_count, 3)
 
+    def test_tensor_subclass_with_opaque_attr_backward(self):
+        """Test opaque objects in tensor subclass are correctly remapped through backward."""
+
+        def fn(x):
+            return x * 2 + 1
+
+        a = torch.rand(4, 4, requires_grad=True)
+        b = torch.rand(4, 4, requires_grad=True)
+        counter = Counter(start=3, end=10)
+        size = SizeStore(4)
+        x = TensorWithCounter(a, b, counter, size)
+
+        cnt = CompileCounterWithBackend("aot_eager")
+        opt_fn = torch.compile(fn, backend=cnt, fullgraph=True)
+        out = opt_fn(x)
+
+        # Output should be TensorWithCounter with same opaque objects
+        self.assertTrue(isinstance(out, TensorWithCounter))
+        self.assertIs(out._counter, counter)
+        self.assertIs(out._size_store, size)
+
+        # Run backward
+        out.sum().backward()
+        self.assertIsNotNone(x.grad)
+        self.assertTrue(isinstance(x.grad, TensorWithCounter))
+
+        self.assertEqual(cnt.frame_count, 1)
+
+        # Call again with different instances that are __eq__ equal.
+        # Metadata guard uses ==, so equal values pass the guard.
+        # But they're different objects, so assertIs tests the fix.
+        a2 = torch.rand(4, 4, requires_grad=True)
+        b2 = torch.rand(4, 4, requires_grad=True)
+        counter2 = Counter(start=3, end=10)  # Same values, different instance
+        size2 = SizeStore(4)  # Same value, different instance
+        self.assertEqual(counter, counter2)  # Verify they're equal
+        self.assertIsNot(counter, counter2)  # But not the same object
+        x2 = TensorWithCounter(a2, b2, counter2, size2)
+
+        out2 = opt_fn(x2)
+        self.assertEqual(cnt.frame_count, 1)  # No recompilation
+
+        self.assertTrue(isinstance(out2, TensorWithCounter))
+        # Key checks: should use runtime instances, not traced instances
+        self.assertIs(out2._counter, counter2)
+        # Value-type opaques in ctx are baked at compile time
+        self.assertEqual(out2._size_store, size2)
+
+        out2.sum().backward()
+        self.assertIsNotNone(x2.grad)
+        self.assertTrue(isinstance(x2.grad, TensorWithCounter))
+        self.assertIs(x2.grad._counter, counter2)
+        self.assertEqual(x2.grad._size_store, size2)
+
+    def test_tangent_primal_proxy_collision_for_opaque_inner_attr(self):
+        """Regression test for tangent/primal proxy collision.
+
+        When a tensor subclass has an opaque inner attr, joint graph tracing
+        creates separate FakeScriptObject wrappers for the primal and tangent
+        that share the same underlying real object.  set_proxy_slot must map
+        the tangent wrapper to the *primal* proxy so that forward outputs
+        don't spuriously depend on tangent placeholders (which would crash the
+        partitioner with 'Node tangents_N was invalid, but is output').
+        """
+        from torch._library.fake_class_registry import FakeScriptObject
+        from torch.fx.experimental.proxy_tensor import (
+            _GraphAppendingTracerEx,
+            set_proxy_slot,
+        )
+
+        counter = Counter(start=3, end=10)
+        fso_primal = FakeScriptObject(counter, "Counter", counter)
+        fso_tangent = FakeScriptObject(counter, "Counter", counter)
+        # Sanity: different wrappers, same real_obj
+        self.assertIsNot(fso_primal, fso_tangent)
+        self.assertIs(
+            object.__getattribute__(fso_primal, "real_obj"),
+            object.__getattribute__(fso_tangent, "real_obj"),
+        )
+
+        graph = torch.fx.Graph()
+        tracer = _GraphAppendingTracerEx(graph)
+
+        primal_node = graph.placeholder("primals_1")
+        tangent_node = graph.placeholder("tangents_1")
+        primal_proxy = torch.fx.Proxy(primal_node, tracer)
+        tangent_proxy = torch.fx.Proxy(tangent_node, tracer)
+
+        # Register primal first, then tangent (mirrors joint graph tracing)
+        set_proxy_slot(fso_primal, tracer, primal_proxy)
+        set_proxy_slot(fso_tangent, tracer, tangent_proxy)
+
+        # Both wrappers should resolve to the primal proxy
+        self.assertIs(tracer.opaque_tracker[fso_primal].node, primal_node)
+        self.assertIs(tracer.opaque_tracker[fso_tangent].node, primal_node)
+
+    def test_opaque_produced_by_call_function_saved_for_backward(self):
+        """Test that an opaque object produced by a call_function node
+        (not a placeholder) is correctly saved for backward.
+
+        Without is_opaque_node() in the min_cut partitioner, this crashes
+        with 'AOT Autograd failed to partition' because the opaque node
+        is a non-tensor call_function that can't be saved or recomputed."""
+
+        def fn(x):
+            scale = torch.tensor(2.5)
+            multiplier = torch.ops._TestOpaqueObject.create_multiplier(scale)
+            return torch.ops._TestOpaqueObject.mul_with_scale(multiplier, x)
+
+        x = torch.randn(3, 3, requires_grad=True)
+        opt_fn = torch.compile(fn, backend="aot_eager", fullgraph=True)
+        out = opt_fn(x)
+        self.assertTrue(torch.allclose(out, x * 2.5))
+
+        out.backward(torch.ones_like(out))
+        self.assertIsNotNone(x.grad)
+        expected_grad = torch.ones_like(x) * 2.5
+        self.assertTrue(torch.allclose(x.grad, expected_grad))
+
+    def test_tensor_subclass_opaque_backward_compiled_autograd(self):
+        """Test opaque objects work with compiled autograd backward."""
+        import torch._dynamo.compiled_autograd
+
+        def fn(x):
+            return x * 2
+
+        a = torch.rand(4, 4, requires_grad=True)
+        b = torch.rand(4, 4, requires_grad=True)
+        counter = Counter(start=5, end=10)
+        size = SizeStore(4)
+        x = TensorWithCounter(a, b, counter, size)
+
+        opt_fn = torch.compile(fn, backend="aot_eager", fullgraph=True)
+        out = opt_fn(x)
+
+        self.assertTrue(isinstance(out, TensorWithCounter))
+        self.assertIs(out._counter, counter)
+        self.assertIs(out._size_store, size)
+
+        with torch._dynamo.compiled_autograd._enable(torch.compile(backend="eager")):
+            out.sum().backward()
+
+        self.assertIsNotNone(x.grad)
+        self.assertTrue(isinstance(x.grad, TensorWithCounter))
+
+        # Second invocation with different-but-equal opaques
+        a2 = torch.rand(4, 4, requires_grad=True)
+        b2 = torch.rand(4, 4, requires_grad=True)
+        counter2 = Counter(start=5, end=10)
+        size2 = SizeStore(4)
+        x2 = TensorWithCounter(a2, b2, counter2, size2)
+
+        out2 = opt_fn(x2)
+        self.assertTrue(isinstance(out2, TensorWithCounter))
+        self.assertIs(out2._counter, counter2)
+        self.assertEqual(out2._size_store, size2)
+
+        with torch._dynamo.compiled_autograd._enable(torch.compile(backend="eager")):
+            out2.sum().backward()
+
+        self.assertIsNotNone(x2.grad)
+        self.assertTrue(isinstance(x2.grad, TensorWithCounter))
+        self.assertIs(x2.grad._counter, counter2)
+        self.assertEqual(x2.grad._size_store, size2)
+
+    def test_tensor_subclass_shared_opaque_remapping(self):
+        """Test that shared opaques across inputs are correctly handled."""
+
+        def fn(x, y):
+            return x + y
+
+        counter = Counter(start=3, end=10)  # Shared opaque
+        size = SizeStore(4)
+
+        a1 = torch.rand(4, 4, requires_grad=True)
+        b1 = torch.rand(4, 4, requires_grad=True)
+        x = TensorWithCounter(a1, b1, counter, size)
+
+        a2 = torch.rand(4, 4, requires_grad=True)
+        b2 = torch.rand(4, 4, requires_grad=True)
+        y = TensorWithCounter(a2, b2, counter, size)  # Same counter
+
+        opt_fn = torch.compile(fn, backend="aot_eager", fullgraph=True)
+        out = opt_fn(x, y)
+
+        self.assertTrue(isinstance(out, TensorWithCounter))
+        self.assertIs(out._counter, counter)
+
+        # Second invocation with different-but-equal shared opaques
+        counter2 = Counter(start=3, end=10)
+        size2 = SizeStore(4)
+
+        a3 = torch.rand(4, 4, requires_grad=True)
+        b3 = torch.rand(4, 4, requires_grad=True)
+        x2 = TensorWithCounter(a3, b3, counter2, size2)
+
+        a4 = torch.rand(4, 4, requires_grad=True)
+        b4 = torch.rand(4, 4, requires_grad=True)
+        y2 = TensorWithCounter(a4, b4, counter2, size2)
+
+        out2 = opt_fn(x2, y2)
+        self.assertTrue(isinstance(out2, TensorWithCounter))
+        self.assertIs(out2._counter, counter2)
+        self.assertEqual(out2._size_store, size2)
+
+    def test_shared_opaque_identity_guard(self):
+        """When inputs share an opaque at trace time, breaking that sharing
+        should trigger recompilation.
+
+        Uses y + x so __torch_dispatch__ picks y's opaque (input 1), but
+        first-occurrence-wins remapping points to x's position (input 0).
+        With broken sharing, the output silently gets the wrong opaque.
+        """
+
+        def fn(x, y):
+            return y + x
+
+        counter = Counter(start=3, end=10)
+        size = SizeStore(4)
+
+        x = TensorWithCounter(torch.rand(4, 4), torch.rand(4, 4), counter, size)
+        y = TensorWithCounter(torch.rand(4, 4), torch.rand(4, 4), counter, size)
+
+        cnt = CompileCounterWithBackend("aot_eager")
+        opt_fn = torch.compile(fn, backend=cnt, fullgraph=True)
+        out = opt_fn(x, y)
+        self.assertEqual(cnt.frame_count, 1)
+        # Shared opaques: output should have the shared counter
+        self.assertIs(out._counter, counter)
+
+        # Same sharing pattern → no recompile
+        counter2 = Counter(start=3, end=10)
+        size2 = SizeStore(4)
+        x2 = TensorWithCounter(torch.rand(4, 4), torch.rand(4, 4), counter2, size2)
+        y2 = TensorWithCounter(torch.rand(4, 4), torch.rand(4, 4), counter2, size2)
+        out2 = opt_fn(x2, y2)
+        self.assertEqual(cnt.frame_count, 1)
+        self.assertIs(out2._counter, counter2)
+
+        # Sharing broken → should recompile.
+        # __torch_dispatch__ picks y's opaque, so the output should have
+        # counter3b. Without the guard, the remapping uses x's position
+        # and the output silently gets counter3a.
+        counter3a = Counter(start=3, end=10)
+        counter3b = Counter(start=3, end=10)
+        size3 = SizeStore(4)
+        x3 = TensorWithCounter(torch.rand(4, 4), torch.rand(4, 4), counter3a, size3)
+        y3 = TensorWithCounter(torch.rand(4, 4), torch.rand(4, 4), counter3b, size3)
+        out3 = opt_fn(x3, y3)
+        self.assertIs(out3._counter, counter3b)
+        self.assertEqual(cnt.frame_count, 2)
+
+    def test_shared_direct_opaque_identity_guard(self):
+        """When the same opaque is passed as multiple direct function
+        arguments, breaking the sharing should trigger recompilation.
+
+        Unlike test_shared_opaque_identity_guard where opaques are inside
+        tensor subclass metadata, here opaques are direct function arguments.
+        The tracer may conflate them when they're the same object, causing
+        counter_start(counter_b) to return counter_a's start instead.
+        """
+        get_start = torch.ops._TestOpaqueObject.counter_start
+
+        def fn(counter_a, counter_b, x):
+            return get_start(counter_a) + x, get_start(counter_b) + x
+
+        counter = Counter(start=3, end=10)
+        x = torch.rand(4, 4)
+
+        cnt = CompileCounterWithBackend("aot_eager")
+        opt_fn = torch.compile(fn, backend=cnt, fullgraph=True)
+        out_a, out_b = opt_fn(counter, counter, x)
+        self.assertEqual(cnt.frame_count, 1)
+        self.assertEqual(out_a, 3 + x)
+        self.assertEqual(out_b, 3 + x)
+
+        # Same sharing → no recompile
+        counter2 = Counter(start=3, end=10)
+        out_a2, out_b2 = opt_fn(counter2, counter2, x)
+        self.assertEqual(cnt.frame_count, 1)
+        self.assertEqual(out_a2, 3 + x)
+        self.assertEqual(out_b2, 3 + x)
+
+        # Broken sharing → should recompile.
+        # Both counters have start=3 so the existing guard_fn doesn't catch
+        # the difference — only an identity guard would.
+        counter3a = Counter(start=3, end=10)
+        counter3b = Counter(start=3, end=10)
+        out_a3, out_b3 = opt_fn(counter3a, counter3b, x)
+        self.assertEqual(cnt.frame_count, 2)
+        self.assertEqual(out_a3, 3 + x)
+        self.assertEqual(out_b3, 3 + x)
+
+    def test_tensor_subclass_shared_opaque_backward(self):
+        """Test backward with multiple subclass inputs sharing the same opaque.
+
+        This tests the case where grad_inputs share an opaque object.
+        The bug was that create_subclass_meta(grad_inputs) would incorrectly
+        build opaque_remappings between grad_inputs, causing an assertion
+        failure in wrap_tensor_subclasses.
+        """
+
+        def fn(x, y):
+            return x + y
+
+        # Create two TensorWithCounter inputs sharing the SAME opaque objects
+        counter = Counter(start=3, end=10)
+        size = SizeStore(4)
+
+        a1 = torch.rand(4, 4, requires_grad=True)
+        b1 = torch.rand(4, 4, requires_grad=True)
+        x = TensorWithCounter(a1, b1, counter, size)
+
+        a2 = torch.rand(4, 4, requires_grad=True)
+        b2 = torch.rand(4, 4, requires_grad=True)
+        y = TensorWithCounter(a2, b2, counter, size)  # Same counter and size
+
+        opt_fn = torch.compile(fn, backend="aot_eager", fullgraph=True)
+        out = opt_fn(x, y)
+
+        self.assertTrue(isinstance(out, TensorWithCounter))
+        self.assertIs(out._counter, counter)
+
+        out.sum().backward()
+        self.assertIsNotNone(x.grad)
+        self.assertIsNotNone(y.grad)
+
+        # Second invocation with different-but-equal shared opaques
+        counter2 = Counter(start=3, end=10)
+        size2 = SizeStore(4)
+
+        a3 = torch.rand(4, 4, requires_grad=True)
+        b3 = torch.rand(4, 4, requires_grad=True)
+        x2 = TensorWithCounter(a3, b3, counter2, size2)
+
+        a4 = torch.rand(4, 4, requires_grad=True)
+        b4 = torch.rand(4, 4, requires_grad=True)
+        y2 = TensorWithCounter(a4, b4, counter2, size2)
+
+        out2 = opt_fn(x2, y2)
+        self.assertTrue(isinstance(out2, TensorWithCounter))
+        self.assertIs(out2._counter, counter2)
+        self.assertEqual(out2._size_store, size2)
+
+        out2.sum().backward()
+        self.assertIsNotNone(x2.grad)
+        self.assertIsNotNone(y2.grad)
+        self.assertTrue(isinstance(x2.grad, TensorWithCounter))
+        self.assertIs(x2.grad._counter, counter2)
+        self.assertEqual(x2.grad._size_store, size2)
+
+    def test_deeply_nested_tensor_subclass_with_opaque(self):
+        """Test opaque objects in nested tensor subclasses are correctly remapped."""
+
+        # Outer subclass that wraps TensorWithCounter
+        class NestedSubclass(torch.Tensor):
+            @staticmethod
+            def __new__(cls, inner, scale):
+                # inner is a TensorWithCounter
+                kwargs = {
+                    "strides": inner.stride(),
+                    "storage_offset": inner.storage_offset(),
+                    "device": inner.device,
+                    "layout": inner.layout,
+                    "requires_grad": inner.requires_grad,
+                    "dtype": inner.dtype,
+                }
+                return torch.Tensor._make_wrapper_subclass(cls, inner.size(), **kwargs)
+
+            def __init__(self, inner, scale):
+                self._inner = inner
+                self._scale = scale
+
+            def __repr__(self):
+                return f"NestedSubclass({self._inner}, scale={self._scale})"
+
+            def __tensor_flatten__(self):
+                return ["_inner"], (self._scale,)
+
+            @staticmethod
+            def __tensor_unflatten__(inner_tensors, ctx, outer_size, outer_stride):
+                (scale,) = ctx
+                return NestedSubclass(inner_tensors["_inner"], scale)
+
+            @classmethod
+            def __torch_dispatch__(cls, func, types, args, kwargs):
+                if kwargs is None:
+                    kwargs = {}
+
+                def unwrap(x):
+                    return x._inner if isinstance(x, NestedSubclass) else x
+
+                def get_scale(x):
+                    if isinstance(x, NestedSubclass):
+                        return x._scale
+                    return None
+
+                scale = None
+                for arg in torch.utils._pytree.tree_leaves(args):
+                    scale = get_scale(arg)
+                    if scale is not None:
+                        break
+
+                unwrapped_args = torch.utils._pytree.tree_map(unwrap, args)
+                unwrapped_kwargs = torch.utils._pytree.tree_map(unwrap, kwargs)
+                out = func(*unwrapped_args, **unwrapped_kwargs)
+
+                def wrap(x):
+                    if isinstance(x, TensorWithCounter):
+                        return NestedSubclass(x, scale)
+                    return x
+
+                return torch.utils._pytree.tree_map(wrap, out)
+
+        def fn(x):
+            return x * 2 + 1
+
+        # Create nested structure: NestedSubclass -> TensorWithCounter -> tensors
+        a = torch.rand(4, 4, requires_grad=True)
+        b = torch.rand(4, 4, requires_grad=True)
+        counter = Counter(start=3, end=10)
+        size = SizeStore(4)
+        inner = TensorWithCounter(a, b, counter, size)
+        x = NestedSubclass(inner, scale=2.0)
+
+        cnt = CompileCounterWithBackend("aot_eager")
+        opt_fn = torch.compile(fn, backend=cnt, fullgraph=True)
+        out = opt_fn(x)
+
+        # Output should preserve nested structure with same opaque objects
+        self.assertTrue(isinstance(out, NestedSubclass))
+        self.assertTrue(isinstance(out._inner, TensorWithCounter))
+        self.assertIs(out._inner._counter, counter)
+        self.assertIs(out._inner._size_store, size)
+        self.assertEqual(out._scale, 2.0)
+
+        # Test backward
+        out.sum().backward()
+        self.assertIsNotNone(x.grad)
+        self.assertTrue(isinstance(x.grad, NestedSubclass))
+
+        self.assertEqual(cnt.frame_count, 1)
+
+        # Call again with different instances that are __eq__ equal.
+        a2 = torch.rand(4, 4, requires_grad=True)
+        b2 = torch.rand(4, 4, requires_grad=True)
+        counter2 = Counter(start=3, end=10)  # Same values, different instance
+        size2 = SizeStore(4)  # Same value, different instance
+        inner2 = TensorWithCounter(a2, b2, counter2, size2)
+        x2 = NestedSubclass(inner2, scale=2.0)
+
+        out2 = opt_fn(x2)
+        self.assertEqual(cnt.frame_count, 1)  # No recompilation
+
+        self.assertTrue(isinstance(out2, NestedSubclass))
+        # Key checks: should use runtime instances, not traced instances
+        self.assertIs(out2._inner._counter, counter2)
+        # Value-type opaques in ctx are baked at compile time, so identity
+        # is not preserved across re-invocations.
+        self.assertEqual(out2._inner._size_store, size2)
+        self.assertEqual(out2._scale, 2.0)
+
+        out2.sum().backward()
+        self.assertIsNotNone(x2.grad)
+        self.assertTrue(isinstance(x2.grad, NestedSubclass))
+        self.assertTrue(isinstance(x2.grad._inner, TensorWithCounter))
+        self.assertIs(x2.grad._inner._counter, counter2)
+        self.assertEqual(x2.grad._inner._size_store, size2)
+
+    def test_value_type_opaque_in_tensor_attrs_errors(self):
+        """Value-type opaques must go in ctx, not tensor attrs."""
+
+        class BadSubclass(torch.Tensor):
+            @staticmethod
+            def __new__(cls, data, size_store):
+                kwargs = {
+                    "strides": data.stride(),
+                    "storage_offset": data.storage_offset(),
+                    "device": data.device,
+                    "layout": data.layout,
+                    "requires_grad": data.requires_grad,
+                    "dtype": data.dtype,
+                }
+                return torch.Tensor._make_wrapper_subclass(cls, data.size(), **kwargs)
+
+            def __init__(self, data, size_store):
+                self._data = data
+                self._size_store = size_store
+
+            def __tensor_flatten__(self):
+                # BUG: SizeStore is value-type, should be in ctx not tensor attrs
+                return ["_data", "_size_store"], ()
+
+            @staticmethod
+            def __tensor_unflatten__(inner_tensors, ctx, outer_size, outer_stride):
+                return BadSubclass(inner_tensors["_data"], inner_tensors["_size_store"])
+
+            @classmethod
+            def __torch_dispatch__(cls, func, types, args, kwargs):
+                if kwargs is None:
+                    kwargs = {}
+
+                def unwrap(x):
+                    return x._data if isinstance(x, BadSubclass) else x
+
+                size_store = None
+                for arg in torch.utils._pytree.tree_leaves(args):
+                    if isinstance(arg, BadSubclass):
+                        size_store = arg._size_store
+                        break
+
+                unwrapped_args = torch.utils._pytree.tree_map(unwrap, args)
+                unwrapped_kwargs = torch.utils._pytree.tree_map(unwrap, kwargs)
+                out = func(*unwrapped_args, **unwrapped_kwargs)
+
+                def wrap(x):
+                    if isinstance(x, torch.Tensor):
+                        return BadSubclass(x, size_store)
+                    return x
+
+                return torch.utils._pytree.tree_map(wrap, out)
+
+        def fn(x):
+            return x * 2
+
+        x = BadSubclass(torch.rand(4, 4), SizeStore(4))
+        opt_fn = torch.compile(fn, backend="aot_eager", fullgraph=True)
+
+        with self.assertRaisesRegex(
+            RuntimeError,
+            "SizeStore.*tensor attrs.*Only tensors and reference-type opaques",
+        ):
+            opt_fn(x)
+
+    def test_multi_reference_opaque_identity(self):
+        """Test identity preservation with multiple reference-type opaques in tensor attrs."""
+
+        class Tag(OpaqueBase):
+            def __init__(self, name):
+                self.name = name
+
+            def __eq__(self, other):
+                return isinstance(other, Tag) and self.name == other.name
+
+            def __hash__(self):
+                return hash(self.name)
+
+        register_opaque_type(Tag, typ="reference")
+
+        class TwoRefSubclass(torch.Tensor):
+            @staticmethod
+            def __new__(cls, data, counter, tag):
+                kwargs = {
+                    "strides": data.stride(),
+                    "storage_offset": data.storage_offset(),
+                    "device": data.device,
+                    "layout": data.layout,
+                    "requires_grad": data.requires_grad,
+                    "dtype": data.dtype,
+                }
+                return torch.Tensor._make_wrapper_subclass(cls, data.size(), **kwargs)
+
+            def __init__(self, data, counter, tag):
+                self._data = data
+                self._counter = counter
+                self._tag = tag
+
+            def __repr__(self):
+                return f"TwoRefSubclass({self._data})"
+
+            def __tensor_flatten__(self):
+                return ["_data", "_counter", "_tag"], ()
+
+            @staticmethod
+            def __tensor_unflatten__(inner_tensors, ctx, outer_size, outer_stride):
+                return TwoRefSubclass(
+                    inner_tensors["_data"],
+                    inner_tensors["_counter"],
+                    inner_tensors["_tag"],
+                )
+
+            @classmethod
+            def __torch_dispatch__(cls, func, types, args, kwargs):
+                if kwargs is None:
+                    kwargs = {}
+
+                def unwrap(x):
+                    return x._data if isinstance(x, TwoRefSubclass) else x
+
+                counter = None
+                tag = None
+                for arg in torch.utils._pytree.tree_leaves(args):
+                    if isinstance(arg, TwoRefSubclass):
+                        counter = arg._counter
+                        tag = arg._tag
+                        break
+
+                unwrapped_args = torch.utils._pytree.tree_map(unwrap, args)
+                unwrapped_kwargs = torch.utils._pytree.tree_map(unwrap, kwargs)
+                out = func(*unwrapped_args, **unwrapped_kwargs)
+
+                def wrap(x):
+                    if isinstance(x, torch.Tensor):
+                        return TwoRefSubclass(x, counter, tag)
+                    return x
+
+                return torch.utils._pytree.tree_map(wrap, out)
+
+        def fn(x):
+            return x * 2 + 1
+
+        counter = Counter(start=3, end=10)
+        tag = Tag("test")
+        x = TwoRefSubclass(torch.rand(4, 4, requires_grad=True), counter, tag)
+
+        cnt = CompileCounterWithBackend("aot_eager")
+        opt_fn = torch.compile(fn, backend=cnt, fullgraph=True)
+        out = opt_fn(x)
+
+        self.assertIsInstance(out, TwoRefSubclass)
+        self.assertIs(out._counter, counter)
+        self.assertIs(out._tag, tag)
+
+        out.sum().backward()
+        self.assertIsNotNone(x.grad)
+        self.assertEqual(cnt.frame_count, 1)
+
+        # Second invocation with different-but-equal opaques
+        counter2 = Counter(start=3, end=10)
+        tag2 = Tag("test")
+        self.assertEqual(counter, counter2)
+        self.assertIsNot(counter, counter2)
+        x2 = TwoRefSubclass(torch.rand(4, 4, requires_grad=True), counter2, tag2)
+
+        out2 = opt_fn(x2)
+        self.assertEqual(cnt.frame_count, 1)  # No recompilation
+
+        self.assertIsInstance(out2, TwoRefSubclass)
+        self.assertIs(out2._counter, counter2)
+        self.assertIs(out2._tag, tag2)
+
+        out2.sum().backward()
+        self.assertIsNotNone(x2.grad)
+        self.assertIsInstance(x2.grad, TwoRefSubclass)
+        self.assertIs(x2.grad._counter, counter2)
+        self.assertIs(x2.grad._tag, tag2)
+
     def test_opaque_obj_saved_for_backward(self):
         """Test that opaque objects are correctly saved and passed to backward."""
         import torch._dynamo.compiled_autograd
@@ -2019,36 +3050,52 @@ def f(x):
     )
     def test_hoist_cache_hits(self):
         torch._dynamo.reset()
-        AOTAutogradCache.clear()
-        torch._inductor.codecache.FxGraphCache.clear()
         counters.clear()
 
-        # Because HoistedString should not be in the graph, the following
-        # two functions should share AOTAutogradCache and FXGraphCache entries
+        with fresh_inductor_cache():
+            # Because HoistedString should not be in the graph, the following
+            # two functions should share AOTAutogradCache and FXGraphCache entries
 
-        @torch.compile(fullgraph=True)
-        def f(x):
-            return op_with_string(x, HoistedString("double"))
+            @torch.compile(fullgraph=True)
+            def f(x):
+                return op_with_string(x, HoistedString("double"))
 
-        @torch.compile(fullgraph=True)
-        def g(x):
-            return op_with_string(x, HoistedString("square"))
+            @torch.compile(fullgraph=True)
+            def g(x):
+                return op_with_string(x, HoistedString("square"))
 
-        x = torch.tensor(3.0)
+            x = torch.tensor(3.0)
+
+            f(x)
+            self.assertEqual(counters["aot_autograd"]["autograd_cache_hit"], 0)
+            self.assertEqual(counters["aot_autograd"]["autograd_cache_miss"], 1)
+            self.assertEqual(counters["aot_autograd"]["autograd_cache_saved"], 1)
+            self.assertEqual(counters["inductor"]["fxgraph_cache_hit"], 0)
+            self.assertEqual(counters["inductor"]["fxgraph_cache_miss"], 1)
 
-        f(x)
-        self.assertEqual(counters["aot_autograd"]["autograd_cache_hit"], 0)
-        self.assertEqual(counters["aot_autograd"]["autograd_cache_miss"], 1)
-        self.assertEqual(counters["aot_autograd"]["autograd_cache_saved"], 1)
-        self.assertEqual(counters["inductor"]["fxgraph_cache_hit"], 0)
-        self.assertEqual(counters["inductor"]["fxgraph_cache_miss"], 1)
+            g(x)
+            self.assertEqual(counters["aot_autograd"]["autograd_cache_hit"], 1)
+            self.assertEqual(counters["aot_autograd"]["autograd_cache_miss"], 1)
+            self.assertEqual(counters["aot_autograd"]["autograd_cache_saved"], 1)
+            self.assertEqual(counters["inductor"]["fxgraph_cache_hit"], 1)
+            self.assertEqual(counters["inductor"]["fxgraph_cache_miss"], 1)
 
-        g(x)
-        self.assertEqual(counters["aot_autograd"]["autograd_cache_hit"], 1)
-        self.assertEqual(counters["aot_autograd"]["autograd_cache_miss"], 1)
-        self.assertEqual(counters["aot_autograd"]["autograd_cache_saved"], 1)
-        self.assertEqual(counters["inductor"]["fxgraph_cache_hit"], 1)
-        self.assertEqual(counters["inductor"]["fxgraph_cache_miss"], 1)
+    def test_hoisted_value_type_make_fx(self):
+        def foo(x, hoisted_str):
+            return op_with_string(x, hoisted_str)
+
+        x = torch.randn(3, 3)
+        gm = make_fx(foo)(x, HoistedString("double"))
+
+        self.assertExpectedInline(
+            gm.code.strip(),
+            """\
+def forward(self, x_1, hoisted_str_1):
+    op_with_string = torch.ops.mylib.op_with_string.default(x_1, hoisted_str_1);  x_1 = hoisted_str_1 = None
+    return op_with_string""",  # noqa: B950
+        )
+        self.assertEqual(gm(x, HoistedString("double")), x * 2)
+        self.assertEqual(gm(x, HoistedString("square")), x * x)
 
     def test_opaque_class_literal_attribute_inlined(self):
         """Test that literal attributes on opaque classes are inlined without source tracking.
@@ -2219,24 +3266,26 @@ def gn(scale_obj, x):
         res.sum().backward()
 
         actual = normalize_gm(backend.graphs[0].print_readable(print_output=False))
-        fx_class = _illegal_char_regex.sub("_", get_opaque_type_name(OpaqueMultiplier))
+        # Normalize module-qualified opaque type names since they differ
+        # depending on how the test is invoked (__main__ vs test_opaque_obj_v2).
+        actual = re.sub(r"\w+_OpaqueMultiplier", "OpaqueMultiplier", actual)
         self.assertExpectedInline(
             actual,
-            f"""\
+            """\
 class GraphModule(torch.nn.Module):
-    def forward(self, L_scale_obj_ : {fx_class}, L_x_: "f32[2, 2]"):
-        l_scale_obj_ = L_scale_obj_
+    def forward(self, L_x_: "f32[2, 2]", L_scale_obj_ : OpaqueMultiplier):
         l_x_ = L_x_
+        l_scale_obj_ = L_scale_obj_
 
         subgraph_0 = self.subgraph_0
         invoke_subgraph = torch.ops.higher_order.invoke_subgraph(subgraph_0, 'subgraph_0', l_scale_obj_, l_x_);  subgraph_0 = l_scale_obj_ = l_x_ = None
-        getitem_2: "f32[2, 2]" = invoke_subgraph[0];  invoke_subgraph = None
+        getitem: "f32[2, 2]" = invoke_subgraph[0];  invoke_subgraph = None
 
-        add: "f32[2, 2]" = getitem_2 + getitem_2;  getitem_2 = None
+        add: "f32[2, 2]" = getitem + getitem;  getitem = None
         return (add,)
 
     class subgraph_0(torch.nn.Module):
-        def forward(self, l_scale_obj_ : {fx_class}, l_x_: "f32[2, 2]"):
+        def forward(self, l_scale_obj_ : OpaqueMultiplier, l_x_: "f32[2, 2]"):
             result: "f32[2, 2]" = torch.ops._TestOpaqueObject.mul_with_scale(l_scale_obj_, l_x_);  l_scale_obj_ = l_x_ = None
 
             result_1: "f32[2, 2]" = result * 2;  result = None
@@ -2247,6 +3296,43 @@ def forward(self, l_scale_obj_ : {fx_class}, l_x_: "f32[2, 2]"):
         self.assertEqual(ref, res)
         self.assertEqual(ref.grad, res.grad)
 
+    def test_value_type_subclass(self):
+        """Subclasses of registered opaque types should be recognized as opaque."""
+
+        class ExtendedConfig(ValueConfig):
+            def __init__(self, mode: str, scale: float):
+                super().__init__(mode)
+                self.scale = scale
+
+            def __eq__(self, other):
+                return (
+                    isinstance(other, ExtendedConfig)
+                    and self.mode == other.mode
+                    and self.scale == other.scale
+                )
+
+            def __hash__(self):
+                return hash((self.mode, self.scale))
+
+            def __fx_repr__(self):
+                return (
+                    f"ExtendedConfig(mode={self.mode!r}, scale={self.scale!r})",
+                    {"ExtendedConfig": ExtendedConfig},
+                )
+
+        self.assertTrue(is_opaque_type(ExtendedConfig))
+        self.assertTrue(is_opaque_value_type(ExtendedConfig))
+
+        cfg = ExtendedConfig("square", 2.0)
+
+        def foo(x):
+            return torch.ops._TestOpaqueObject.process_with_config(x, cfg)
+
+        x = torch.randn(3, 3)
+        opt_f = torch.compile(foo, fullgraph=True, backend="aot_eager")
+        res = opt_f(x)
+        self.assertEqual(res, foo(x))
+
     def test_opaque_object_with_inductor_backend(self):
         """Test that opaque objects work correctly with inductor's get_attr handling."""
 
@@ -2272,6 +3358,635 @@ def forward(self, x):
         expected = x * float(Color.RED.value)
         self.assertTrue(torch.allclose(result, expected))
 
+    def test_captured_opaque_object_export(self):
+        """Test that opaque objects captured from closures work with torch.export.
+
+        When an opaque object is captured via a module hook closure, it gets a
+        CellContentsSource instead of GetItemSource. The export path needs to
+        handle this gracefully.
+        """
+        rng = RNGState(42)
+
+        class Model(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.linear = torch.nn.Linear(4, 4)
+
+            def forward(self, x):
+                torch.ops._TestOpaqueObject.noisy_inject(x, rng)
+                return self.linear(x)
+
+        ep = torch.export.export(Model(), (torch.randn(4, 4),), strict=True)
+        self.assertExpectedInline(
+            ep.graph_module.code.strip(),
+            """\
+def forward(self, p_linear_weight, p_linear_bias, obj_lifted_custom_0, x):
+    noisy_inject = torch.ops._TestOpaqueObject.noisy_inject.default(x, obj_lifted_custom_0);  obj_lifted_custom_0 = noisy_inject = None
+    linear = torch.ops.aten.linear.default(x, p_linear_weight, p_linear_bias);  x = p_linear_weight = p_linear_bias = None
+    return (linear,)""",  # noqa: B950
+        )
+
+    def test_hoist_no_recompile_on_different_string(self):
+        cnt = CompileCounter()
+
+        def f(x, label):
+            return op_with_string(x, HoistedString(label))
+
+        opt_f = torch.compile(f, backend=cnt, fullgraph=True)
+        x = torch.tensor(3.0)
+
+        res1 = opt_f(x, "double")
+        self.assertEqual(res1, f(x, "double"))
+        self.assertEqual(cnt.frame_count, 1)
+
+        res2 = opt_f(x, "square")
+        self.assertEqual(res2, f(x, "square"))
+        self.assertEqual(cnt.frame_count, 1)
+
+    def test_opaque_multi_output_not_tensor_irnode(self):
+        """OpaqueMultiOutput must not be classified as a tensor IR node.
+
+        _is_tensor_irnode is used to select nodes for stride constraint
+        functions (require_stride1, require_contiguous, etc). If
+        OpaqueMultiOutput passes the check, those functions would call
+        get_stride/get_dtype/make_loader on it and crash."""
+        from unittest.mock import MagicMock
+
+        from torch._inductor import ir
+        from torch._inductor.lowering import _is_tensor_irnode
+
+        opaque_multi = ir.OpaqueMultiOutput.__new__(ir.OpaqueMultiOutput)
+        self.assertIsInstance(opaque_multi, ir.IRNode)
+        self.assertFalse(_is_tensor_irnode(opaque_multi))
+
+        non_tensor = MagicMock(spec=ir.NonTensorObj)
+        non_tensor.__class__ = ir.NonTensorObj
+        self.assertFalse(_is_tensor_irnode(non_tensor))
+
+        self.assertFalse(_is_tensor_irnode(42))
+
+    @unittest.skipIf(not dist.is_available(), "requires distributed")
+    def test_fake_script_object_process_group_pybind(self):
+        """FakeScriptObject wrapping ProcessGroup must be unwrapped in the
+        pybind toIValue before casting to c10::intrusive_ptr<ProcessGroup>.
+
+        During Dynamo tracing with CooR, mesh_get_process_group returns a
+        FakeScriptObject-wrapped ProcessGroup. When that flows into a C++
+        custom op, toIValue needs to unwrap real_obj before the pybind
+        cast. This test exercises that path by calling
+        mesh_get_process_group (which returns a wrapped PG during tracing)
+        and passing it to another op that consumes the ProcessGroup."""
+        from torch.distributed.device_mesh import (
+            _register_distributed_opaque_types,
+            DeviceMesh,
+        )
+        from torch.testing._internal.distributed.fake_pg import FakeStore
+
+        already_initialized = dist.is_initialized()
+        if already_initialized:
+            dist.destroy_process_group()
+
+        dist.init_process_group("fake", store=FakeStore(), rank=0, world_size=2)
+        try:
+            _register_distributed_opaque_types()
+            mesh = DeviceMesh("cpu", torch.arange(2))
+
+            def f(mesh, x):
+                pg = torch.ops._dtensor.mesh_get_process_group(mesh, 0)
+                return x + pg.size()
+
+            x = torch.randn(4)
+            compiled_f = torch.compile(f, backend="aot_eager", fullgraph=True)
+            result = compiled_f(mesh, x)
+            expected = f(mesh, x)
+            self.assertEqual(result, expected)
+        finally:
+            dist.destroy_process_group()
+            if already_initialized:
+                dist.init_process_group("fake", store=FakeStore(), rank=0, world_size=2)
+
+    def test_enum_export(self):
+        class Direction(enum.Enum):
+            UP = 0
+            DOWN = 1
+
+        class Mod(torch.nn.Module):
+            def forward(self, x, d):
+                return x + d.value
+
+        ep = torch.export.export(Mod(), (torch.randn(4, 4), Direction.UP), strict=False)
+        self.assertEqual(
+            ep.module()(torch.ones(4, 4), Direction.UP),
+            torch.ones(4, 4) + Direction.UP.value,
+        )
+        self.assertExpectedInline(
+            normalize_gm(ep.graph_module.print_readable(False)),
+            """\
+class GraphModule(torch.nn.Module):
+    def forward(self, x: "f32[4, 4]", d):
+        add: "f32[4, 4]" = torch.ops.aten.add.Tensor(x, 0);  x = None
+        return (add,)
+""",
+        )
+
+        backend = EagerAndRecordGraphs()
+        opt_fn = torch.compile(Mod(), backend=backend)
+        x = torch.randn(4, 4)
+        res = opt_fn(x, Direction.UP)
+        self.assertEqual(
+            res,
+            x + Direction.UP.value,
+        )
+        self.assertExpectedInline(
+            normalize_gm(backend.graphs[0].print_readable(False)),
+            """\
+class GraphModule(torch.nn.Module):
+    def forward(self, L_x_: "f32[4, 4]"):
+        l_x_ = L_x_
+
+        add: "f32[4, 4]" = l_x_ + 0;  l_x_ = None
+        return (add,)
+""",
+        )
+
+    def test_enum_custom_op(self):
+        def get_color():
+            class Color(enum.Enum):
+                RED = 0
+                GREEN = 1
+                BLUE = 2
+
+            return Color
+
+        Color = get_color()
+
+        @torch.library.custom_op("test_enum::add_color", mutates_args=())
+        def add_color(x: torch.Tensor, c: Color) -> torch.Tensor:
+            return x + c.value
+
+        @add_color.register_fake
+        def _(x, c):
+            return torch.empty_like(x)
+
+        def fn(x, c):
+            return add_color(x, c)
+
+        x = torch.randn(4, 4)
+        ref = fn(x, Color.GREEN)
+        backend = EagerAndRecordGraphs()
+        opt_fn = torch.compile(fn, backend=backend, fullgraph=True)
+        res = opt_fn(x, Color.GREEN)
+        self.assertEqual(ref, res)
+        self.assertExpectedInline(
+            normalize_gm(backend.graphs[0].print_readable(False)),
+            """\
+class GraphModule(torch.nn.Module):
+    def forward(self, L_x_: "f32[4, 4]"):
+        l_x_ = L_x_
+
+        add_color_default: "f32[4, 4]" = torch.ops.test_enum.add_color.default(l_x_, Color.GREEN);  l_x_ = None
+        return (add_color_default,)
+""",
+        )
+
+    def test_subclass_parametrization_with_opaque_attrs(self):
+        """unwrap_tensor_subclass_parameters should handle non-tensor attrs."""
+        from torch._functorch._aot_autograd.subclass_parametrization import (
+            unwrap_tensor_subclass_parameters,
+        )
+
+        class M(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                a = torch.randn(3, 4)
+                b = torch.randn(3, 4)
+                self.p = torch.nn.Parameter(
+                    TensorWithCounter(a, b, Counter(1, 10), SizeStore(4))
+                )
+
+            def forward(self, x):
+                return x + self.p
+
+        m = M()
+        unwrap_tensor_subclass_parameters(m)
+        x = torch.randn(3, 4)
+        m(x)
+
+    def test_export_non_strict_with_opaque_attrs(self):
+        """torch.export with strict=False should handle subclasses with non-tensor attrs."""
+
+        class M(torch.nn.Module):
+            def forward(self, x):
+                return x * 2
+
+        a = torch.randn(4, 4)
+        b = torch.randn(4, 4)
+        x = TensorWithCounter(a, b, Counter(1, 10), SizeStore(4))
+        torch.export.export(M(), (x,), strict=False)
+
+    @unittest.skipIf(not dist.is_available(), "requires distributed")
+    def test_get_untyped_storages_with_opaque_attrs(self):
+        """get_untyped_storages should skip non-tensor attrs without warning."""
+        from torch.distributed._tools.common_utils import get_untyped_storages
+
+        a = torch.randn(4, 4)
+        b = torch.randn(4, 4)
+        x = TensorWithCounter(a, b, Counter(1, 10), SizeStore(4))
+        import warnings
+
+        with warnings.catch_warnings():
+            warnings.simplefilter("error")
+            storages = get_untyped_storages(x)
+        # Should get storages from the two inner tensors (a and b)
+        self.assertTrue(len(storages) > 0)
+
+    @functorch_config.patch({"enable_autograd_cache": True})
+    @inductor_config.patch(
+        {
+            "fx_graph_cache": True,
+            "fx_graph_remote_cache": False,
+        }
+    )
+    def test_subclass_opaque_attrs_cache_hit(self):
+        """opaque_attrs in SubclassCreationMeta should survive AOTAutograd cache round-trip."""
+        torch._dynamo.reset()
+        counters.clear()
+
+        with fresh_inductor_cache():
+
+            @torch.compile(fullgraph=True)
+            def fn(x):
+                return x * 2
+
+            a = torch.randn(4, 4)
+            b = torch.randn(4, 4)
+            counter = Counter(start=3, end=10)
+            size = SizeStore(4)
+            x = TensorWithCounter(a, b, counter, size)
+
+            out = fn(x)
+            self.assertIsInstance(out, TensorWithCounter)
+            self.assertIs(out._counter, counter)
+            self.assertIs(out._size_store, size)
+            self.assertEqual(counters["aot_autograd"]["autograd_cache_miss"], 1)
+            self.assertEqual(counters["aot_autograd"]["autograd_cache_hit"], 0)
+
+            # Reset dynamo to force a fresh compilation that should hit the cache
+            torch._dynamo.reset()
+
+            @torch.compile(fullgraph=True)
+            def fn2(x):
+                return x * 2
+
+            a2 = torch.randn(4, 4)
+            b2 = torch.randn(4, 4)
+            counter2 = Counter(start=3, end=10)
+            size2 = SizeStore(4)
+            x2 = TensorWithCounter(a2, b2, counter2, size2)
+
+            out2 = fn2(x2)
+            self.assertIsInstance(out2, TensorWithCounter)
+            # Runtime opaque values should be used, not the serialized ones
+            self.assertIs(out2._counter, counter2)
+            self.assertEqual(out2._size_store, size2)
+            self.assertEqual(counters["aot_autograd"]["autograd_cache_hit"], 1)
+
+    def test_op_passthrough_counter_in_tuple(self):
+        # When a fake kernel returns its Counter input directly, the getitem
+        # proxy's example_value is already a FakeScriptObject.
+        counter_type = get_opaque_type_name(Counter)
+        torch.library.define(
+            "_TestOpaqueObject::passthrough_counter",
+            f"({counter_type} c, Tensor x) -> ({counter_type}, Tensor)",
+            tags=torch.Tag.pt2_compliant_tag,
+            lib=self.lib,
+        )
+
+        @torch.library.impl(
+            "_TestOpaqueObject::passthrough_counter",
+            "CompositeExplicitAutograd",
+            lib=self.lib,
+        )
+        def passthrough_impl(c: Counter, x: torch.Tensor):
+            return c, x * c.start
+
+        @torch.library.register_fake(
+            "_TestOpaqueObject::passthrough_counter", lib=self.lib
+        )
+        def passthrough_fake(c: Counter, x: torch.Tensor):
+            return c, torch.empty_like(x)
+
+        def fn(c, x):
+            out_c, out_x = torch.ops._TestOpaqueObject.passthrough_counter(c, x)
+            return torch.ops._TestOpaqueObject.counter_start(out_c) + out_x
+
+        c = Counter(3, 10)
+        x = torch.randn(4)
+        ref = fn(c, x)
+        opt_fn = torch.compile(fn, fullgraph=True, backend="eager")
+        res = opt_fn(c, x)
+        self.assertEqual(ref, res)
+
+    @torch._dynamo.config.patch(skip_fwd_side_effects_in_bwd_under_checkpoint=True)
+    def test_script_object_intermediate_exposed_from_checkpoint(self):
+        # A TorchScriptObjectVariable created inside an AC region and accessed
+        # outside via a list side effect must be exposed as a subgraph output.
+        import torch.utils.checkpoint
+
+        def gn(x, results):
+            counter = torch.ops._TestOpaqueObject.create_counter(x.shape[0], x.shape[0])
+            results.append(counter)
+            return x * 2
+
+        def fn(x):
+            results = []
+            out = torch.utils.checkpoint.checkpoint(gn, x, results, use_reentrant=False)
+            return torch.ops._TestOpaqueObject.counter_start(results[0]) + out
+
+        x = torch.randn(3, 4)
+        ref = fn(x)
+        opt_fn = torch.compile(fn, fullgraph=True, backend="aot_eager_decomp_partition")
+        res = opt_fn(x)
+        self.assertEqual(ref, res)
+
+    def test_reference_type_opaque_object_state(self):
+        """When compile_fx_inner receives an FX graph whose placeholder meta['val']
+        is a raw opaque reference type (not wrapped in FakeScriptObject), inductor
+        handles it via OpaqueObjectState. This codepath is used by compile-on-one-rank
+        (COOR) where the FX graph is constructed with real objects."""
+        m = OpaqueMultiplier(2.0)
+        x = torch.ones(3)
+
+        graph = torch.fx.Graph()
+        m_node = graph.placeholder("m")
+        m_node.meta["val"] = m
+        fake_mode = FakeTensorMode()
+        x_node = graph.placeholder("x")
+        x_node.meta["val"] = fake_mode.from_tensor(x)
+        out = graph.call_function(
+            torch.ops._TestOpaqueObject.mul_with_scale.default, (m_node, x_node)
+        )
+        out.meta["val"] = fake_mode.from_tensor(x)
+        graph.output((out,))
+
+        gm = torch.fx.GraphModule({}, graph)
+        compiled = compile_fx_inner(gm, [m, x])
+        result = compiled([m, x])
+        self.assertEqual(result, (x * 2,))
+
+    @unittest.skipIf(not HAS_GPU, "Inductor+gpu needs triton and recent GPU arch")
+    def test_benchmark_harness_no_pickle_for_opaque_inputs(self):
+        """Opaque graph inputs must not be pickled in the benchmark harness."""
+        a = torch.randn(4, 4, device="cuda")
+        b = torch.randn(4, 4, device="cuda")
+        twc = TensorWithCounter(a, b, Counter(0, 10), SizeStore(4))
+
+        def fn(x):
+            return x + 1
+
+        compiled_fn = torch.compile(fn, backend="inductor", fullgraph=True)
+        _, codes = run_and_get_code(compiled_fn, twc)
+        self.assertGreater(len(codes), 0)
+        for code in codes:
+            self.assertNotIn("pickle", code)
+
+    def test_opaque_object_state_in_graph_output(self):
+        """When compile_fx_inner receives a graph where a raw opaque reference
+        type appears in the outputs (not just inputs), inductor must handle it.
+        This happens in CooR (compile-on-one-rank) precompile: aot_autograd
+        partitions the joint graph so the forward graph saves opaque objects
+        (e.g. DeviceMesh) for the backward graph by including them in forward
+        outputs.  Inductor must:
+          1. Accept OpaqueObjectState in the output type assertion
+             (GraphLowering.output).
+          2. Handle NonTensorObj graph inputs in the memory planner
+             (get_dep_size_hint) without calling get_numel/get_dtype on them."""
+        m = OpaqueMultiplier(2.0)
+        x = torch.ones(3)
+
+        graph = torch.fx.Graph()
+        m_node = graph.placeholder("m")
+        m_node.meta["val"] = m
+        fake_mode = FakeTensorMode()
+        x_node = graph.placeholder("x")
+        x_node.meta["val"] = fake_mode.from_tensor(x)
+        out = graph.call_function(
+            torch.ops._TestOpaqueObject.mul_with_scale.default, (m_node, x_node)
+        )
+        out.meta["val"] = fake_mode.from_tensor(x)
+        # Include the opaque object in the output tuple, simulating how
+        # aot_autograd's forward graph passes saved-for-backward objects
+        # through as outputs.
+        graph.output((out, m_node))
+
+        gm = torch.fx.GraphModule({}, graph)
+        compiled = compile_fx_inner(gm, [m, x])
+        result = compiled([m, x])
+        self.assertEqual(result[0], x * 2)
+        self.assertIs(result[1], m)
+
+    def test_reconstruct_fn_sets_meta_val(self):
+        """Opaque nodes created via reconstruct_fn have meta['val'] set.
+
+        When _try_reconstruct_opaque dispatches through a custom op via
+        Proxy.__torch_function__, the resulting node does not get
+        meta['val'] set automatically (unlike the __torch_dispatch__
+        path which calls track_tensor_tree).  The fix in
+        _try_reconstruct_opaque ensures set_meta is called so that
+        downstream consumers like the min-cut partitioner can classify
+        the node correctly via is_opaque_node().
+        """
+        from torch._functorch._aot_autograd.graph_compile import is_opaque_node
+
+        # Use the already-registered OpaqueMultiplier type.
+        # Register a reconstruct_fn that derives one multiplier from
+        # another via a custom op — mirrors how DeviceMesh submeshes
+        # are derived from a parent mesh via _get_submesh.
+        multiplier_type = get_opaque_type_name(OpaqueMultiplier)
+
+        self.lib.define(
+            f"derive_multiplier({multiplier_type} parent, float scale)"
+            f" -> {multiplier_type}",
+        )
+
+        @torch.library.impl(
+            "_TestOpaqueObject::derive_multiplier",
+            "CompositeExplicitAutograd",
+            lib=self.lib,
+        )
+        def derive_impl(parent, scale):
+            return OpaqueMultiplier(parent.multiplier * scale)
+
+        @torch.library.register_fake(
+            "_TestOpaqueObject::derive_multiplier", lib=self.lib
+        )
+        def derive_fake(parent, scale):
+            return OpaqueMultiplier(0.0)
+
+        parent = OpaqueMultiplier(3.0)
+        # The "child" is derived at eager time before tracing.
+        # It's NOT passed as an input — it's captured in a closure,
+        # simulating how DeviceMesh submeshes are captured in DTensor
+        # backward closures.  This forces _try_reconstruct_opaque to
+        # be called when make_fx encounters the child.
+        child = OpaqueMultiplier(parent.multiplier * 0.5)
+        child._parent = parent
+        child._scale = 0.5
+
+        from torch._library.opaque_object import _OPAQUE_TYPES
+
+        original_reconstruct_fn = _OPAQUE_TYPES[OpaqueMultiplier].reconstruct_fn
+
+        def multiplier_reconstruct_fn(obj, get_tracked_proxy, tracer):
+            if not hasattr(obj, "_parent"):
+                return None
+            parent_proxy = get_tracked_proxy(obj._parent)
+            if parent_proxy is None:
+                return None
+            return torch.ops._TestOpaqueObject.derive_multiplier(
+                parent_proxy, obj._scale
+            )
+
+        _OPAQUE_TYPES[OpaqueMultiplier].reconstruct_fn = multiplier_reconstruct_fn
+        try:
+            # child is captured from the closure, NOT an input
+            def fn(parent, x):
+                return torch.ops._TestOpaqueObject.mul_with_scale(child, x)
+
+            fake_mode = FakeTensorMode(shape_env=ShapeEnv())
+            with fake_mode:
+                fake_x = torch.randn(4)
+
+            gm = make_fx(fn, tracing_mode="fake")(parent, fake_x)
+
+            # Find the derive_multiplier node (created by reconstruct_fn)
+            derive_nodes = [
+                n
+                for n in gm.graph.nodes
+                if n.op == "call_function" and "derive_multiplier" in str(n.target)
+            ]
+            self.assertGreater(len(derive_nodes), 0, "No derive_multiplier node found")
+
+            for node in derive_nodes:
+                self.assertIn(
+                    "val",
+                    node.meta,
+                    f"Node {node.name} created via reconstruct_fn is missing "
+                    f"meta['val']. This would cause the partitioner to fail "
+                    f"with 'Expected {node.name} to be a tensor'.",
+                )
+                self.assertTrue(
+                    is_opaque_node(node),
+                    f"Node {node.name} should be classified as opaque",
+                )
+        finally:
+            _OPAQUE_TYPES[OpaqueMultiplier].reconstruct_fn = original_reconstruct_fn
+
+    def test_partitioner_must_save_opaque_node(self):
+        """Opaque nodes tagged MUST_SAVE go to saved_opaque_nodes.
+
+        When activation checkpointing tags an opaque node with
+        CheckpointPolicy.MUST_SAVE, the default_partition function
+        must route it to saved_opaque_nodes (not saved_values).
+        Otherwise the runtime assertion in save_from_forward will
+        fail because it expects all saved_values to be Tensors.
+        """
+        from functorch.compile import default_partition
+        from torch._functorch._aot_autograd.graph_compile import is_opaque_node
+        from torch.utils.checkpoint import CheckpointPolicy
+
+        m = OpaqueMultiplier(2.0)
+
+        fake_mode = FakeTensorMode(shape_env=ShapeEnv())
+
+        # Build a joint fwd+bwd graph where the opaque node is tagged
+        # MUST_SAVE — simulates activation checkpointing.
+        graph = torch.fx.Graph()
+        m_node = graph.placeholder("m")
+        m_node.meta["val"] = m
+        x_node = graph.placeholder("x")
+        with fake_mode:
+            x_node.meta["val"] = torch.randn(3)
+
+        mul_node = graph.call_function(
+            torch.ops._TestOpaqueObject.mul_with_scale.default,
+            (m_node, x_node),
+        )
+        with fake_mode:
+            mul_node.meta["val"] = torch.randn(3)
+
+        # Tag the opaque node as MUST_SAVE
+        m_node.meta["recompute"] = CheckpointPolicy.MUST_SAVE
+
+        # Create backward portion: identity backward
+        bw_node = graph.call_function(torch.ops.aten.mul.Scalar, (mul_node, 1.0))
+        with fake_mode:
+            bw_node.meta["val"] = torch.randn(3)
+
+        graph.output((mul_node, bw_node))
+        joint_gm = torch.fx.GraphModule({}, graph)
+
+        # Run default_partition — this should NOT raise even though
+        # the opaque node is tagged MUST_SAVE.
+        num_fwd_outputs = 1
+        fw_module, bw_module = default_partition(
+            joint_gm, [], num_fwd_outputs=num_fwd_outputs
+        )
+
+        # Verify the opaque node ended up in the forward graph as
+        # a pass-through (not as a saved tensor that would fail the
+        # save_from_forward assertion).
+        fw_opaque_nodes = [
+            n
+            for n in fw_module.graph.nodes
+            if n.op == "placeholder" and is_opaque_node(n)
+        ]
+        self.assertGreater(
+            len(fw_opaque_nodes),
+            0,
+            "Forward graph should have opaque placeholder nodes",
+        )
+
+    def test_hoisted_opaque_with_graphpickler(self):
+        if import_dill() is None:
+            self.skipTest("dill not available")
+
+        backend = EagerAndRecordGraphs()
+
+        def fn(x):
+            return torch.ops.mylib.op_with_string(x, HoistedString("double"))
+
+        compiled = torch.compile(fn, fullgraph=True, backend=backend)
+        compiled(torch.randn(3, 4))
+
+        self.assertEqual(len(backend.graphs), 1)
+        captured_gm = backend.graphs[0]
+
+        placeholders = [n for n in captured_gm.graph.nodes if n.op == "placeholder"]
+        opaque_ph = None
+        for ph in placeholders:
+            ev = ph.meta.get("example_value")
+            if isinstance(ev, FakeScriptObject):
+                opaque_ph = ph
+                break
+        self.assertIsNotNone(opaque_ph, "no placeholder with FakeScriptObject found")
+
+        fake_mode = FakeTensorMode(shape_env=ShapeEnv())
+        options = Options(ops_filter=None, node_metadata_key_filter=None)
+        data = GraphPickler.dumps(captured_gm, options)
+        restored_gm = GraphPickler.loads(data, fake_mode)
+
+        restored_phs = [n for n in restored_gm.graph.nodes if n.op == "placeholder"]
+        restored_opaque = None
+        for ph in restored_phs:
+            ev = ph.meta.get("example_value")
+            if isinstance(ev, FakeScriptObject):
+                restored_opaque = ph
+                break
+        restored_real = restored_opaque.meta["example_value"].real_obj
+        self.assertIsInstance(restored_real, HoistedString)
+        self.assertEqual(restored_real.val, "double")
+
 
 instantiate_parametrized_tests(TestOpaqueObject)
 
diff --git a/test/test_ops.py b/test/test_ops.py
index 83e7d0d311afc..579b592e883c3 100644
--- a/test/test_ops.py
+++ b/test/test_ops.py
@@ -39,6 +39,7 @@
 from torch.testing._internal.common_dtype import (
     all_types_and_complex_and,
     floating_and_complex_types_and,
+    highest_precision_float,
     integral_types_and,
 )
 from torch.testing._internal.common_methods_invocations import (
@@ -467,6 +468,7 @@ def test_reduction_ops_reduce(self, device, op):
     # resulting in possible equality check failures.
     # skip windows case on CPU due to https://github.com/pytorch/pytorch/issues/129947
     # XPU test will be enabled step by step. Skip the tests temporarily.
+    # MPS does not support double precision, so single precision has to be used instead.
     @skipXPU
     @onlyNativeDeviceTypesAnd(["hpu"])
     @suppress_warnings
@@ -483,7 +485,7 @@ def test_numpy_ref(self, device, dtype, op):
             raise unittest.SkipTest("XXX: raises tensor-likes are not close.")
 
         # Sets the default dtype to NumPy's default dtype of double
-        with set_default_dtype(torch.double):
+        with set_default_dtype(highest_precision_float(device)):
             for sample_input in op.reference_inputs(device, dtype):
                 self.compare_with_reference(
                     op, op.ref, sample_input, exact_dtype=(dtype is not torch.long)
@@ -635,9 +637,13 @@ def _ref_test_helper(
                 # precise dtypes -- they simply must be close
                 precise_dtype = dtype
             if prims.utils.is_float_dtype(dtype):
-                precise_dtype = torch.double
+                precise_dtype = highest_precision_float(device)
             if prims.utils.is_complex_dtype(dtype):
-                precise_dtype = torch.cdouble
+                precise_dtype = (
+                    torch.complex32
+                    if torch.device(device).type == "mps"
+                    else torch.cdouble
+                )
 
             # Checks if the results are close
             try:
diff --git a/test/test_optim.py b/test/test_optim.py
index 6de3cabefc417..23094907f944b 100644
--- a/test/test_optim.py
+++ b/test/test_optim.py
@@ -43,6 +43,7 @@
     markDynamoStrictTest,
     parametrize,
     run_tests,
+    serialTest,
     TEST_WITH_TORCHDYNAMO,
     TestCase,
 )
@@ -65,6 +66,38 @@ def drosenbrock(tensor):
     return torch.stack((-400 * x * (y - x**2) - 2 * (1 - x), 200 * (y - x**2)))
 
 
+def _bf16_state_init_hook(optimizer, args, kwargs):
+    """Step pre-hook that initializes Adam/AdamW states in bfloat16.
+
+    Pre-populates optimizer state before Adam's lazy initialization so that
+    ``_init_group`` finds non-empty state and skips its own fp32 allocation.
+    The fused CUDA kernel then dispatches to its mixed-precision path.
+    """
+    for group in optimizer.param_groups:
+        for p in group["params"]:
+            if p.grad is None:
+                continue
+            state = optimizer.state[p]
+            if len(state) == 0:
+                state["step"] = (
+                    torch.zeros((), dtype=torch.float32, device=p.device)
+                    if group.get("capturable") or group.get("fused")
+                    else torch.tensor(0.0, dtype=torch.float32)
+                )
+                state["exp_avg"] = torch.zeros_like(
+                    p, dtype=torch.bfloat16, memory_format=torch.preserve_format
+                )
+                state["exp_avg_sq"] = torch.zeros_like(
+                    p, dtype=torch.bfloat16, memory_format=torch.preserve_format
+                )
+                if group.get("amsgrad"):
+                    state["max_exp_avg_sq"] = torch.zeros_like(
+                        p,
+                        dtype=torch.bfloat16,
+                        memory_format=torch.preserve_format,
+                    )
+
+
 @markDynamoStrictTest
 class TestOptimRenewed(TestCase):
     """
@@ -964,6 +997,7 @@ def test_set_default_dtype_works_with_foreach(self, device, dtype, optim_info):
 
     @onlyCUDA
     @largeTensorTest("72GB", "cuda")
+    @serialTest()
     @optims(
         [optim for optim in optim_db if "foreach" in optim.supported_impls],
         dtypes=[torch.float16],
@@ -1129,6 +1163,7 @@ def test_fused_error_on_params_on_meta(self, device, dtype, optim_info):
 
     @onlyNativeDeviceTypes
     @largeTensorTest("64GB")
+    @serialTest()
     @optims(
         [optim for optim in optim_db if "fused" in optim.supported_impls],
         dtypes=[torch.float16],
@@ -2313,6 +2348,160 @@ def test_non_empty_state(self, device, dtype, optim_info):
             for state in optim.state.values():
                 self.assertGreater(len(state), 0)
 
+    @onlyCUDA
+    @parametrize("amsgrad", [False, True])
+    @optims(
+        [o for o in optim_db if o.optim_cls.__name__ in ["Adam", "AdamW"]],
+        dtypes=[torch.float32],
+    )
+    def test_fused_mixed_precision_state_init(self, device, dtype, optim_info, amsgrad):
+        optim_cls = optim_info.optim_cls
+        params = [torch.rand(20, 7, device=device, dtype=dtype) for _ in range(5)]
+        for p in params:
+            p.grad = torch.rand_like(p)
+
+        optim = optim_cls(params, lr=1e-3, fused=True, amsgrad=amsgrad)
+        optim.register_step_pre_hook(_bf16_state_init_hook)
+
+        optim.step()
+
+        for p in params:
+            self.assertEqual(p.dtype, torch.float32)
+            state = optim.state[p]
+            self.assertEqual(state["step"].dtype, torch.float32)
+            self.assertEqual(state["exp_avg"].dtype, torch.bfloat16)
+            self.assertEqual(state["exp_avg_sq"].dtype, torch.bfloat16)
+            if amsgrad:
+                self.assertEqual(state["max_exp_avg_sq"].dtype, torch.bfloat16)
+
+        # Second step: hook should be idempotent (skips already-populated state)
+        for p in params:
+            p.grad = torch.rand_like(p)
+        optim.step()
+
+        for p in params:
+            state = optim.state[p]
+            self.assertEqual(state["step"].dtype, torch.float32)
+            self.assertEqual(state["exp_avg"].dtype, torch.bfloat16)
+            self.assertEqual(state["exp_avg_sq"].dtype, torch.bfloat16)
+            if amsgrad:
+                self.assertEqual(state["max_exp_avg_sq"].dtype, torch.bfloat16)
+
+    @onlyCUDA
+    @parametrize("amsgrad", [False, True])
+    @optims(
+        [o for o in optim_db if o.optim_cls.__name__ in ["Adam", "AdamW"]],
+        dtypes=[torch.float32],
+    )
+    def test_fused_mixed_precision_hook_skips_existing_state(
+        self, device, dtype, optim_info, amsgrad
+    ):
+        optim_cls = optim_info.optim_cls
+
+        # Two param groups: group 1 gets f32 state pre-populated (hook should
+        # skip it), group 2 has no state (hook should initialize it in bf16).
+        # This exercises the fused kernel handling two groups whose states have
+        # different dtypes within the same optimizer.step() call.
+        g1_params = [torch.rand(10, 5, device=device, dtype=dtype) for _ in range(2)]
+        g2_params = [torch.rand(10, 5, device=device, dtype=dtype) for _ in range(2)]
+        for p in g1_params + g2_params:
+            p.grad = torch.rand_like(p)
+
+        optim = optim_cls(
+            [{"params": g1_params}, {"params": g2_params}],
+            lr=1e-3,
+            fused=True,
+            amsgrad=amsgrad,
+        )
+
+        for p in g1_params:
+            optim.state[p]["step"] = torch.zeros(
+                (), dtype=torch.float32, device=p.device
+            )
+            optim.state[p]["exp_avg"] = torch.zeros_like(p)
+            optim.state[p]["exp_avg_sq"] = torch.zeros_like(p)
+            if amsgrad:
+                optim.state[p]["max_exp_avg_sq"] = torch.zeros_like(p)
+
+        optim.register_step_pre_hook(_bf16_state_init_hook)
+        optim.step()
+
+        # Group 1: hook skipped (state was non-empty), dtypes stay f32.
+        for p in g1_params:
+            state = optim.state[p]
+            self.assertEqual(state["step"].dtype, torch.float32)
+            self.assertEqual(state["exp_avg"].dtype, torch.float32)
+            self.assertEqual(state["exp_avg_sq"].dtype, torch.float32)
+            if amsgrad:
+                self.assertEqual(state["max_exp_avg_sq"].dtype, torch.float32)
+
+        # Group 2: hook initialized state in bf16.
+        for p in g2_params:
+            state = optim.state[p]
+            self.assertEqual(state["step"].dtype, torch.float32)
+            self.assertEqual(state["exp_avg"].dtype, torch.bfloat16)
+            self.assertEqual(state["exp_avg_sq"].dtype, torch.bfloat16)
+            if amsgrad:
+                self.assertEqual(state["max_exp_avg_sq"].dtype, torch.bfloat16)
+
+    @onlyCUDA
+    @optims(
+        [o for o in optim_db if o.optim_cls.__name__ in ["Adam", "AdamW"]],
+        dtypes=[torch.float32],
+    )
+    def test_fused_mixed_precision_numerics(self, device, dtype, optim_info):
+        optim_inputs = optim_info.optim_inputs_func(device=device, dtype=dtype)
+        optim_cls = optim_info.optim_cls
+        for optim_input in optim_inputs:
+            kwargs = {**optim_input.kwargs, "fused": True}
+
+            params = [torch.rand(20, 7, device=device, dtype=dtype) for _ in range(10)]
+            for p in params:
+                p.grad = torch.rand_like(p)
+
+            params_c = [p.clone() for p in params]
+            for p, pc in zip(params, params_c):
+                pc.grad = p.grad.clone()
+
+            ref_optim = optim_cls(params, **kwargs)
+            bf16_optim = optim_cls(params_c, **kwargs)
+            bf16_optim.register_step_pre_hook(_bf16_state_init_hook)
+
+            # Simulate bf16 storage: after each ref step, quantize states to
+            # bf16 and back so the reference matches the mixed-precision kernel.
+            tracker = TensorTracker()
+            for i in range(7):
+                ref_optim.step()
+                bf16_optim.step()
+                for p in params:
+                    tracker.add(p)
+                    tracker.add(p.grad)
+                for d in ref_optim.state.values():
+                    exp_avg_bf16 = d["exp_avg"].to(torch.bfloat16)
+                    tracker.add(exp_avg_bf16)
+                    d["exp_avg"] = exp_avg_bf16.to(torch.float32)
+                    exp_avg_sq_bf16 = d["exp_avg_sq"].to(torch.bfloat16)
+                    tracker.add(exp_avg_sq_bf16)
+                    d["exp_avg_sq"] = exp_avg_sq_bf16.to(torch.float32)
+                    if "max_exp_avg_sq" in d:
+                        max_exp_avg_sq_bf16 = d["max_exp_avg_sq"].to(torch.bfloat16)
+                        tracker.add(max_exp_avg_sq_bf16)
+                        d["max_exp_avg_sq"] = max_exp_avg_sq_bf16.to(torch.float32)
+
+                for e, pc in enumerate(params_c):
+                    tracker.pop_check_set(pc, self)
+                    tracker.pop_check_set(pc.grad, self)
+
+                for p, pc in zip(params, params_c):
+                    self.assertEqual(p, pc)
+
+                for dc in bf16_optim.state.values():
+                    tracker.pop_check_set(dc["exp_avg"], self)
+                    tracker.pop_check_set(dc["exp_avg_sq"], self)
+                    if "max_exp_avg_sq" in dc:
+                        tracker.pop_check_set(dc["max_exp_avg_sq"], self)
+                self.assertTrue(tracker.all_popped())
+
     @parametrize("dtype", [torch.float32])
     def test_step_iteration(self, device, dtype):
         def _get_model_and_input_tensor(device, dtype):
diff --git a/test/test_proxy_tensor.py b/test/test_proxy_tensor.py
index 8296f386f0977..fdb10d06bfc72 100644
--- a/test/test_proxy_tensor.py
+++ b/test/test_proxy_tensor.py
@@ -1017,40 +1017,39 @@ def _test_dynamic(self, fn, trace_inputs, test_inputs, assert_eq=True):
 
 
     def test_debug_interpreter(self):
-        import torch.library
-        from torch.library import Library
+        from torch.library import _scoped_library
 
-        foo = Library("foo", "DEF")  # noqa: TOR901
-        foo.define("foo(Tensor self) -> Tensor")
+        with _scoped_library("foo", "DEF") as foo:
+            foo.define("foo(Tensor self) -> Tensor")
 
-        # Operator where meta and cpu disagree on strides
-        @torch.library.impl(foo, "foo", "CPU")
-        def foo_cpu(x):
-            return x.clone().T
+            # Operator where meta and cpu disagree on strides
+            @torch.library.impl(foo, "foo", "CPU")
+            def foo_cpu(x):
+                return x.clone().T
 
-        @torch.library.impl(foo, "foo", "Meta")
-        def foo_meta(x):
-            return x.clone()
+            @torch.library.impl(foo, "foo", "Meta")
+            def foo_meta(x):
+                return x.clone()
 
-        def f(x):
-            return torch.ops.foo.foo.default(x)
+            def f(x):
+                return torch.ops.foo.foo.default(x)
 
-        gm = make_fx(f, tracing_mode="symbolic")(torch.randn(2, 2))
-        from torch._functorch.compilers import DebugInterpreter
+            gm = make_fx(f, tracing_mode="symbolic")(torch.randn(2, 2))
+            from torch._functorch.compilers import DebugInterpreter
 
-        interp = DebugInterpreter(gm)
+            interp = DebugInterpreter(gm)
 
-        # input mismatch is caught (indicates guard problem)
-        self.assertRaisesRegex(
-            AssertionError, r"3 != 1",
-            lambda: interp.run(torch.randn(3, 3).T),
-        )
+            # input mismatch is caught (indicates guard problem)
+            self.assertRaisesRegex(
+                AssertionError, r"3 != 1",
+                lambda: interp.run(torch.randn(3, 3).T),
+            )
 
-        # Catch the incorrect meta
-        self.assertRaisesRegex(
-            AssertionError, r"\(3, 1\) != \(1, 3\)",
-            lambda: interp.run(torch.randn(3, 3))
-        )
+            # Catch the incorrect meta
+            self.assertRaisesRegex(
+                AssertionError, r"\(3, 1\) != \(1, 3\)",
+                lambda: interp.run(torch.randn(3, 3))
+            )
 
     def test_int_input(self):
         def f(x, y):
@@ -2029,8 +2028,6 @@ def f(t):
     xfail('cov'),
     xfail('nn.functional.gaussian_nll_loss'),
     xfail('corrcoef'),
-    xfail('quantile'),
-    xfail('nanquantile'),
 
     # Seems like it's creating a sparse tensor that isn't captured by tensor.is_sparse
     xfail('sparse.sampled_addmm'),
@@ -2061,11 +2058,9 @@ def f(t):
     xfail('geqrf', ''),  # aten.geqrf.default - couldn't find symbolic meta function/decomposition
     xfail('histogram', ''),  # Could not run 'aten::histogram.bin_ct' with arguments from the 'Meta' backend. This c...
     xfail('histogramdd', ''),  # aten._histogramdd_bin_edges.default - couldn't find symbolic meta function/decomposition
-    xfail('nanquantile', ''),  # Could not run 'aten::equal' with arguments from the 'Meta' backend.
     xfail('nn.functional.binary_cross_entropy', ''),  # aten.new_empty.default - couldn't find symbolic meta function/decom...
     xfail('nn.functional.cross_entropy', ''),  # aten.size.default - couldn't find symbolic meta function/decomposition
     xfail('nn.functional.ctc_loss'),  # aten._ctc_loss.Tensor - couldn't find symbolic meta function/decomposition
-    xfail('quantile', ''),  # Could not run 'aten::equal' with arguments from the 'Meta' backend.
 
     xfail('max_pool2d_with_indices_backward', ''),  # Expected a value of type 'List[int]' for argument 'kernel_size' but...
 }
diff --git a/test/test_public_bindings.py b/test/test_public_bindings.py
index 26e6f47713693..5cae005d8451f 100644
--- a/test/test_public_bindings.py
+++ b/test/test_public_bindings.py
@@ -184,7 +184,6 @@ def test_no_new_bindings(self):
             "PyTorchFileReader",
             "PyTorchFileWriter",
             "qscheme",
-            "read_vitals",
             "RRefType",
             "ScriptClass",
             "ScriptClassFunction",
@@ -211,7 +210,6 @@ def test_no_new_bindings(self):
             "set_flush_denormal",
             "set_num_interop_threads",
             "set_num_threads",
-            "set_vital",
             "Size",
             "StaticModule",
             "Stream",
@@ -232,7 +230,6 @@ def test_no_new_bindings(self):
             "Value",
             "set_autocast_gpu_dtype",
             "get_autocast_gpu_dtype",
-            "vitals_enabled",
             "wait",
             "Tag",
             "set_autocast_xla_enabled",
@@ -357,6 +354,7 @@ def onerror(modname):
             "torch.testing._internal.distributed.rpc.rpc_test",
             "torch.testing._internal.distributed.rpc.tensorpipe_rpc_agent_test_fixture",
             "torch.testing._internal.distributed.rpc_utils",
+            "torch.testing._internal.py312_intrinsics",
             "torch._inductor.codegen.cutlass.cuda_template",
             "torch._inductor.codegen.cutedsl._cutedsl_utils",
             "torch._inductor.codegen.cuda.gemm_template",
@@ -364,7 +362,10 @@ def onerror(modname):
             "torch._inductor.codegen.cpp_gemm_template",
             "torch._inductor.codegen.cpp_micro_gemm",
             "torch._inductor.codegen.cpp_template_kernel",
-            "torch._inductor.kernel.vendored_templates.cutedsl_grouped_gemm",  # depends on cutlass_cppgen
+            "torch._inductor.kernel.vendored_templates.cutedsl.kernels.cutedsl_grouped_gemm",  # depends on cutlass
+            "torch._inductor.kernel.vendored_templates.cutedsl.dense_blockscaled_gemm_persistent",  # depends on cutlass
+            "torch._inductor.kernel.vendored_templates.cutedsl.wrappers",  # depends on cutlass_api
+            "torch._inductor.kernel.vendored_templates.cutedsl.wrappers.dense_blockscaled_gemm_kernel",  # depends on cutlass_api
             "torch._inductor.runtime.triton_helpers",
             "torch.ao.pruning._experimental.data_sparsifier.lightning.callbacks.data_sparsity",
             "torch.backends._coreml.preprocess",
diff --git a/test/test_python_dispatch.py b/test/test_python_dispatch.py
index 0c69b1c61003e..041450e066ffd 100644
--- a/test/test_python_dispatch.py
+++ b/test/test_python_dispatch.py
@@ -26,7 +26,9 @@
 from torch.testing._internal.common_methods_invocations import op_db
 from torch.testing._internal.common_utils import (
     first_sample,
+    instantiate_parametrized_tests,
     IS_WINDOWS,
+    parametrize,
     run_tests,
     TEST_WITH_ROCM,
     TestCase,
@@ -219,9 +221,6 @@ def my_fallback(op, *args, **kwargs):
             self.assertEqual(c, a + b)
             self.assertTrue(is_called)
 
-    @unittest.skip(
-        "Causing flakiness, see https://github.com/pytorch/pytorch/issues/145108"
-    )
     def test_fallthrough_for_dense_key_with_meta_in_tls(self) -> None:
         # This tests that if meta is included in TlS dispatch key set,
         # then a meta kernel should be called regardless if a dense
@@ -244,6 +243,29 @@ def sum_meta(*args, **kwargs):
                 torch.ops.custom.sum.default(a)
                 self.assertTrue(meta_is_called)
 
+    def test_include_dispatch_key_guard_restores_tls_exactly(self) -> None:
+        before = torch._C._dispatch_tls_local_include_set().raw_repr()
+        with torch._C._IncludeDispatchKeyGuard(torch.DispatchKey.Meta):
+            pass
+        after = torch._C._dispatch_tls_local_include_set().raw_repr()
+        self.assertEqual(before, after)
+
+    @parametrize(
+        "key",
+        [
+            torch.DispatchKey.Meta,
+            torch.DispatchKey.CUDA,
+            torch.DispatchKey.CPU,
+        ],
+    )
+    def test_exclude_dispatch_key_guard_restores_tls_exactly(self, key) -> None:
+        keyset = torch._C.DispatchKeySet(key)
+        before = torch._C._dispatch_tls_local_exclude_set().raw_repr()
+        with torch._C._ExcludeDispatchKeyGuard(keyset):
+            pass
+        after = torch._C._dispatch_tls_local_exclude_set().raw_repr()
+        self.assertEqual(before, after)
+
     def test_dispatchkeyset_pickle(self) -> None:
         keyset = torch._C.DispatchKeySet(torch._C.DispatchKey.AutogradCPU)
         serialized = pickle.dumps(keyset)
@@ -689,6 +711,9 @@ def test_register_fallthrough(self):
             self.assertEqual(torch.mm(a, b).dtype, torch.bfloat16)
 
 
+instantiate_parametrized_tests(TestPythonRegistration)
+
+
 class TestPythonDispatch(TestCase):
     def test_basic(self) -> None:
         with capture_logs() as logs:
diff --git a/test/test_pytree.py b/test/test_pytree.py
index 410199a8e3c47..1b7902c020f27 100644
--- a/test/test_pytree.py
+++ b/test/test_pytree.py
@@ -1,8 +1,10 @@
 # Owner(s): ["module: pytree"]
 
+import copy
 import enum
 import inspect
 import os
+import pickle
 import re
 import subprocess
 import sys
@@ -11,7 +13,7 @@
 from collections import defaultdict, deque, namedtuple, OrderedDict, UserDict
 from dataclasses import dataclass, field
 from enum import auto
-from typing import Any, NamedTuple, Optional
+from typing import Any, NamedTuple
 
 import torch
 import torch.utils._pytree as python_pytree
@@ -817,6 +819,38 @@ def test_enum_treespec_roundtrip(self, pytree):
         deserialized_spec = pytree.treespec_loads(serialized)
         self.assertEqual(spec, deserialized_spec)
 
+    @parametrize_pytree_module
+    def test_treespec_deepcopy_roundtrip(self, pytree):
+        cases = [
+            1,
+            (1, 2),
+            [1, 2, 3],
+            {"a": 1, "b": 2},
+            (1, [2, {"a": 3}]),
+            {"a": [1, 2], "b": (3, 4)},
+        ]
+
+        for tree in cases:
+            treespec = pytree.tree_structure(tree)
+            reconstructed = copy.deepcopy(treespec)
+            self.assertEqual(treespec, reconstructed)
+
+    @parametrize_pytree_module
+    def test_treespec_pickle_roundtrip(self, pytree):
+        cases = [
+            1,
+            (1, 2),
+            [1, 2, 3],
+            {"a": 1, "b": 2},
+            (1, [2, {"a": 3}]),
+            {"a": [1, 2], "b": (3, 4)},
+        ]
+
+        for tree in cases:
+            treespec = pytree.tree_structure(tree)
+            reconstructed = pickle.loads(pickle.dumps(treespec))
+            self.assertEqual(treespec, reconstructed)
+
 
 class TestPythonPytree(TestCase):
     def test_deprecated_register_pytree_node(self):
@@ -1256,7 +1290,7 @@ def test_dataclass(self):
         class Data:
             a: torch.Tensor
             b: str = "moo"
-            c: Optional[str] = None
+            c: str | None = None
             d: str = field(init=False, default="")
 
         python_pytree.register_dataclass(Data)
diff --git a/test/test_reductions.py b/test/test_reductions.py
index 077f5e3fc405d..f50d53d496802 100644
--- a/test/test_reductions.py
+++ b/test/test_reductions.py
@@ -14,7 +14,8 @@
 from torch import inf, nan
 from torch.testing import make_tensor
 from torch.testing._internal.common_dtype import (
-    all_types_and_complex_and, get_all_math_dtypes, integral_types, complex_types, floating_types_and,
+    all_types_and_complex_and, get_all_math_dtypes, highest_precision_float,
+    integral_types, complex_types, floating_types_and,
     integral_types_and, floating_and_complex_types_and, all_types_and, all_types,
 )
 from torch.testing._internal.common_utils import (
@@ -427,8 +428,6 @@ def test_ref_duplicate_values(self, device, dtype, op: ReductionOpInfo):
          allowed_dtypes=[torch.float32, torch.complex64])
     def test_ref_extremal_values(self, device, dtype, op: ReductionOpInfo):
         """Compares op against reference for input tensors with extremal values"""
-        if "xpu" in device and dtype is torch.complex64 and op.name == "mean":
-            self.skipTest("accuracy issue with this test, see https://github.com/intel/torch-xpu-ops/issues/2300")
         t = make_tensor((5,), dtype=dtype, device=device, exclude_zero=True)
         extremals = [0, 1, nan, inf, -inf]
         for extremal in extremals:
@@ -1607,8 +1606,8 @@ def test_bucketization(self, device):
         self.assertEqual(torch.bucketize(values_0_el, boundaries), expected_result)
 
         # nan input
-        values_nan = torch.tensor([1.0, float('nan'), 2.0, float('nan')], device=device, dtype=torch.float64)
-        boundaries = torch.tensor([0.0, 1.0, 2.0, 3.0], device=device, dtype=torch.float64)
+        values_nan = torch.tensor([1.0, float('nan'), 2.0, float('nan')], device=device, dtype=highest_precision_float(device))
+        boundaries = torch.tensor([0.0, 1.0, 2.0, 3.0], device=device, dtype=highest_precision_float(device))
         expected_result = torch.tensor([1, 4, 2, 4], device=device)
         self.assertEqual(torch.searchsorted(boundaries, values_nan), expected_result)
         expected_result = torch.tensor([2, 4, 3, 4], device=device)
@@ -1617,7 +1616,7 @@ def test_bucketization(self, device):
 
         # type promotion and non contiguous tensors
         values_3d_permute = values_3d.permute(2, 1, 0).to(torch.int32)
-        boundaries_permute = values_3d.permute(2, 1, 0).to(torch.float64)
+        boundaries_permute = values_3d.permute(2, 1, 0).to(highest_precision_float(device))
         expected_result = torch.tensor([[[0, 0], [0, 1]], [[2, 0], [0, 1]], [[2, 0], [0, 0]]], device=device)
         if self.device_type != 'xla':
             self.assertWarnsRegex(
diff --git a/test/test_scaled_matmul_cuda.py b/test/test_scaled_matmul_cuda.py
index 94842c069a6f3..5aafa1399ff12 100644
--- a/test/test_scaled_matmul_cuda.py
+++ b/test/test_scaled_matmul_cuda.py
@@ -7,7 +7,6 @@
 import itertools
 import tempfile
 import unittest
-from typing import Optional
 
 import torch
 
@@ -47,11 +46,14 @@
 
 from torch.testing._internal.common_utils import (
     IS_WINDOWS,
+    MI350_ARCH,
     parametrize,
     run_tests,
+    runOnRocmArch,
     skipIfRocm,
     TEST_CUDA,
     TestCase,
+    skipIfXpu,
 )
 from torch.testing._internal.common_quantized import (
     _bfloat16_to_float4_e2m1fn_x2,
@@ -70,7 +72,7 @@
 if TEST_CUDA:
     _IS_SM8X = torch.cuda.get_device_capability(0)[0] == 8
 
-f8_msg = "FP8 is only supported on H100+, SM 8.9 and MI300+ and XPU devices"
+f8_msg = "FP8 is only supported on H100+, SM 8.9 and MI300+, XPU and CPU devices"
 f8_grouped_msg = "FP8 grouped is only supported on SM90 and MI300/MI350 devices"
 mx_skip_msg = "MX gemm is only supported on CUDA capability 10.0+"
 mxfp8_grouped_mm_skip_msg = "MXFP8 grouped GEMM is only supported when PyTorch is built with USE_MSLK=1 on SM100+"
@@ -295,7 +297,7 @@ def scaled_grouped_mm_wrap(
 
 
 
-def mm_float8_emulated(x, x_scale, y, y_scale, out_dtype, bias: Optional[torch.Tensor] = None) -> torch.Tensor:
+def mm_float8_emulated(x, x_scale, y, y_scale, out_dtype, bias: torch.Tensor | None = None) -> torch.Tensor:
     # naive implementation: dq -> op -> q
     x_fp32 = x.to(torch.float) / x_scale
     y_fp32 = y.to(torch.float) / y_scale
@@ -322,8 +324,8 @@ def addmm_float8_unwrapped(
     b_data: torch.Tensor,
     b_scale: torch.tensor,
     output_dtype: torch.dtype,
-    output_scale: Optional[torch.Tensor],
-    bias: Optional[torch.Tensor] = None,
+    output_scale: torch.Tensor | None,
+    bias: torch.Tensor | None = None,
 ) -> torch.Tensor:
     a_inverse_scale = a_scale.reciprocal()
     b_inverse_scale = b_scale.reciprocal()
@@ -632,7 +634,7 @@ class TestFP8Matmul(TestCase):
     def _test_tautological_mm(self, device: str = "cuda",
                               x_dtype: torch.dtype = e4m3_type,
                               y_dtype: torch.dtype = e4m3_type,
-                              out_dtype: Optional[torch.dtype] = None,
+                              out_dtype: torch.dtype | None = None,
                               x_cm: bool = True,
                               y_cm: bool = False,
                               size: int = 16) -> None:
@@ -654,30 +656,67 @@ def _test_tautological_mm(self, device: str = "cuda",
 
     # Skip on XPU due to known oneDNN accuracy issue (#169772)
     @skipXPU
-    def test_float8_basics(self, device) -> None:
+    @parametrize(
+        "test_case",
+        [
+            # test_case tuple schema:
+            # (case_name, x_dtype, y_dtype, out_dtype, size)
+            # "default" in case_name means out_dtype=None (backend default output dtype path).
+            ("e4m3_e4m3_default", e4m3_type, e4m3_type, None, 16),
+            ("e5m2_e5m2_default", e5m2_type, e5m2_type, None, 16),
+            ("e5m2_e5m2_f32", e5m2_type, e5m2_type, torch.float32, 16),
+            ("e4m3_e5m2_default", e4m3_type, e5m2_type, None, 32),
+            ("e5m2_e4m3_default", e5m2_type, e4m3_type, None, 48),
+            ("e4m3_e4m3_f16", e4m3_type, e4m3_type, torch.float16, 64),
+            ("e4m3_e4m3_f32", e4m3_type, e4m3_type, torch.float32, 96),
+            ("e4m3_e4m3_bf16", e4m3_type, e4m3_type, torch.bfloat16, 80),
+        ],
+        name_fn=lambda test_case: f"cuda_{test_case[0]}",
+    )
+    def test_float8_basics(self, device, test_case) -> None:
         if not _device_supports_scaled_mm_fp8(device):
             raise unittest.SkipTest(f8_msg)
-        self._test_tautological_mm(device, e4m3_type, e4m3_type, size=16)
+        _, x_dtype, y_dtype, out_dtype, size = test_case
+        expect_e5m2_cuda_error = x_dtype == e5m2_type and y_dtype == e5m2_type
         # According to https://docs.nvidia.com/cuda/cublas/#id99 8F_E5M2 MM is unsupported
         # supported on ROCm but fails on CUDA
-        ctx = self.assertRaises(ValueError) if torch.version.hip is None and "cuda" in device else contextlib.nullcontext()
+        ctx = (
+            self.assertRaises(ValueError)
+            if expect_e5m2_cuda_error and torch.version.hip is None and "cuda" in device
+            else contextlib.nullcontext()
+        )
         with ctx:
-            self._test_tautological_mm(device, e5m2_type, e5m2_type)
-
-        self._test_tautological_mm(device, e4m3_type, e5m2_type, size=32)
-        self._test_tautological_mm(device, e5m2_type, e4m3_type, size=48)
-
-        self._test_tautological_mm(device, size=64, out_dtype=torch.float16)
-        self._test_tautological_mm(device, size=96, out_dtype=torch.float32)
-        self._test_tautological_mm(device, size=80, out_dtype=torch.bfloat16)
+            self._test_tautological_mm(
+                device,
+                x_dtype=x_dtype,
+                y_dtype=y_dtype,
+                out_dtype=out_dtype,
+                size=size,
+            )
 
+    # Skip on XPU due to known oneDNN accuracy issue (#169772)
+    @skipXPU
+    def test_float8_basics_layout_permutations(self, device) -> None:
+        if not _device_supports_scaled_mm_fp8(device):
+            raise unittest.SkipTest(f8_msg)
         if torch.cuda.is_available():
             for (x_cm, y_cm) in itertools.product([True, False], repeat=2):
-                # Blackwell (SM_10) supports all possible layout permutations, while Hopper only TN
-                layouts_supported = (x_cm, y_cm) == (True, False) or torch.cuda.get_device_properties(0).major == 10
+                # SM 10 and 11 support all permutations, SM 12 TT and TN, SM 9 only TN
+                major, minor = torch.cuda.get_device_capability(0)
+                if major in (10, 11):
+                    layouts_supported = True
+                elif major == 12 and (minor == 1 or _get_torch_cuda_version() >= (13, 1)):
+                    layouts_supported = x_cm
+                else:
+                    layouts_supported = (x_cm, y_cm) == (True, False)
                 with contextlib.nullcontext() if layouts_supported else self.assertRaises(RuntimeError):
                     self._test_tautological_mm(device, size=64, out_dtype=torch.bfloat16, x_cm=x_cm, y_cm=y_cm)
 
+    # Skip on XPU due to known oneDNN accuracy issue (#169772)
+    @skipXPU
+    def test_float8_basics_invalid_out_dtype(self, device) -> None:
+        if not _device_supports_scaled_mm_fp8(device):
+            raise unittest.SkipTest(f8_msg)
         with self.assertRaises(
             AssertionError if (torch.version.hip or "xpu" in device or "cpu" in device)
             else RuntimeError
@@ -917,15 +956,15 @@ def _2d_to_blocked_scaled(X, K, G, offs, format):
         # Assert outputs are close.
         torch.testing.assert_close(y_lp, y_bf16, atol=8.0e-2, rtol=8.0e-2)
 
-
     @unittest.skipIf(not PLATFORM_SUPPORTS_FP8, f8_msg)
     @parametrize("base_dtype", [torch.float16, torch.bfloat16, torch.float32])
     @parametrize("x_cm", [True, False])
     @parametrize("y_cm", [True, False])
-    def test_scaled_mm_vs_emulated(self, base_dtype, x_cm, y_cm, device="cuda"):
+    def test_scaled_mm_vs_emulated(self, base_dtype, x_cm, y_cm, device):
         # Blackwell (SM_10) supports all possible layout permutations, while Hopper only TN
-        if (x_cm, y_cm) != (True, False) and torch.cuda.get_device_properties(0).major != 10:
-            raise unittest.SkipTest("Unsupported layout on the architecture")
+        if torch.cuda.is_available():
+            if (x_cm, y_cm) != (True, False) and torch.cuda.get_device_properties(0).major != 10:
+                raise unittest.SkipTest("Unsupported layout on the architecture")
         torch.manual_seed(42)
         input_dtype = e4m3_type
         output_dtype = base_dtype
@@ -1026,7 +1065,7 @@ def test_scaled_mm_change_stride(self, base_dtype, device="cuda"):
 
         torch.testing.assert_close(out_scaled_mm, out_emulated, atol=atol, rtol=rtol)
 
-    @onlyOn(["cuda", "xpu"])
+    @onlyOn(["cuda", "xpu", "cpu"])
     @skipCUDAIf(not PLATFORM_SUPPORTS_FP8, f8_msg)
     def test_float8_bias(self, device) -> None:
         (k, l, m) = (16, 48, 32)
@@ -1043,7 +1082,7 @@ def test_float8_bias(self, device) -> None:
         difference = torch.abs(out_fp32 - outb_fp32)
         self.assertEqual(difference, torch.tensor(4.0, device=device).expand_as(out_fp32))
 
-    @onlyOn(["cuda", "xpu"])
+    @onlyOn(["cuda", "xpu", "cpu"])
     @unittest.skipIf(not PLATFORM_SUPPORTS_FP8, f8_msg)
     @parametrize("bias", [True, False])
     def test_non_divisible_leading_dim(self, device, bias: bool) -> None:
@@ -1056,7 +1095,7 @@ def test_non_divisible_leading_dim(self, device, bias: bool) -> None:
             input_bias = torch.rand((16,), device=device).to(torch.bfloat16)
         _ = scaled_mm_wrap(x, y, scale_a, scale_b, bias=input_bias)
 
-    @onlyOn(["cuda", "xpu"])
+    @onlyOn(["cuda", "xpu", "cpu"])
     @unittest.skipIf(not PLATFORM_SUPPORTS_FP8, f8_msg)
     def test_float8_bias_relu_edgecase(self, device) -> None:
         (k, l, m) = (16, 48, 32)
@@ -1069,7 +1108,7 @@ def test_float8_bias_relu_edgecase(self, device) -> None:
         outb_fp32 = outb_fp8.to(torch.float32)
         self.assertEqual(outb_fp32, torch.tensor(-3.0, device=device).expand_as(outb_fp32))
 
-    @onlyOn(["cuda", "xpu"])
+    @onlyOn(["cuda", "xpu", "cpu"])
     @unittest.skipIf(not PLATFORM_SUPPORTS_FP8, f8_msg)
     def test_float32_output_errors_with_bias(self, device) -> None:
         (k, l, m) = (16, 48, 32)
@@ -1078,8 +1117,8 @@ def test_float32_output_errors_with_bias(self, device) -> None:
         scale_a = torch.tensor(1.0, device=device)
         scale_b = torch.tensor(1.0, device=device)
         bias = torch.full((m,), 4.0, device=device, dtype=torch.bfloat16)
-        # XPU supports the case when out_dtype is fp32 + bias. So we just test it with normal run.
-        if "xpu" not in device:
+        # XPU and CPU supports the case when out_dtype is fp32 + bias. So we just test it with normal run.
+        if "xpu" not in device and "cpu" not in device:
             self.assertRaisesRegex(
                 ValueError if torch.cuda.is_available() else RuntimeError,
                 "Bias is not supported when out_dtype is set to Float32",
@@ -1100,6 +1139,7 @@ def test_error_message_fp8_pre_sm89(self, device) -> None:
             lambda: scaled_mm_wrap(x, y, scale_a, scale_b, out_dtype=torch.float32),
         )
 
+    @skipIfXpu(msg="AssertionError, torch-xpu-ops: 2862")
     @unittest.skipIf(not PLATFORM_SUPPORTS_FP8, f8_msg)
     @unittest.skipIf(SM100OrLater, "fast_accum is SM90-only")
     def test_float8_scale_fast_accum(self, device) -> None:
@@ -1115,7 +1155,7 @@ def test_float8_scale_fast_accum(self, device) -> None:
         out_fp8_s = scaled_mm_wrap(x, y, scale_a=scale_a, scale_b=scale_b, out_dtype=e4m3_type, use_fast_accum=True)
         self.assertEqual(out_fp8, out_fp8_s)
 
-    @onlyOn(["cuda", "xpu"])
+    @onlyOn(["cuda", "xpu", "cpu"])
     @unittest.skipIf(not PLATFORM_SUPPORTS_FP8 or IS_WINDOWS, f8_msg)
     @skipCUDAIf(not SM89OrLater, "rowwise implementation is currently sm89-sm100 specific")
     @parametrize("use_fast_accum", [True, False])
@@ -1146,7 +1186,7 @@ def test_float8_rowwise_scaling_sanity(self, device, use_fast_accum: bool) -> No
             out_fp8.to(torch.float32), torch.full((M, N), K * (fill_value**2), device=device)
         )
 
-    @onlyOn(["cuda", "xpu"])
+    @onlyOn(["cuda", "xpu", "cpu"])
     @unittest.skipIf(not PLATFORM_SUPPORTS_FP8 or IS_WINDOWS, f8_msg)
     def test_float8_error_messages(self, device) -> None:
         M, K, N = (1024, 512, 2048)
@@ -1222,7 +1262,7 @@ def e5m2():
             (torch.cuda.is_available() and
              torch.cuda.get_device_capability() == (9, 0) and
              torch.version.cuda and
-             torch.version.cuda >= "12.9")):
+             torch.version.cuda >= "12.9") or (not torch.cuda.is_available() and torch.cpu.is_available())):
             out = e5m2()
             self.assertEqual(out, torch.ones_like(out) * 128.)
         else:
@@ -1621,20 +1661,21 @@ def test_scaled_mm_vs_emulated_block_wise_verify_small_shapes(
             output_dtype
         )
 
-    @skipIfRocm
     @onlyCUDA
     @unittest.skipIf(not PLATFORM_SUPPORTS_FP8 or IS_WINDOWS, f8_msg)
     @unittest.skipIf(IS_SM90, "DeepSeek style (1x128, 128x128) blockwise scaling works on SM90 (Hopper)")
     @unittest.skipIf(
-        _get_torch_cuda_version() < (12, 9),
+        not torch.version.hip and _get_torch_cuda_version() < (12, 9),
         "cuBLAS blockwise scaling added in CUDA 12.9",
     )
+    @runOnRocmArch(MI350_ARCH)
     @parametrize("output_dtype", [torch.bfloat16, ])
     @parametrize("lhs_block,rhs_block", [(1, 1), (128, 1), (1, 128)])
     @parametrize("M,N,K", [(256, 256, 256), (256, 256, 512)])
     def test_scaled_mm_deepseek_error_messages(
         self, output_dtype, lhs_block, rhs_block, M, N, K
     ):
+
         torch.manual_seed(42)
 
         x = torch.randn(M, K, device="cuda", dtype=output_dtype).pow(3)
@@ -1656,10 +1697,18 @@ def test_scaled_mm_deepseek_error_messages(
         else:
             rhs_recipe = ScalingType.BlockWise128x128
 
-        # Verify that actual F8 mm raises expected error on non-SM90
+        # Verify that actual F8 mm raises expected error
+        if torch.version.hip:
+            # ROCm does not yet support DeepSeek-style blockwise scaling
+            expected_error = NotImplementedError
+            expected_pattern = "1x128 and 128x128 scaling not available with ROCm"
+        else:
+            # CUDA non-SM90 should raise NotImplementedError
+            expected_error = NotImplementedError
+            expected_pattern = ".*DeepSeek.*scaling.*only supported in CUDA for SM90.*"
         with self.assertRaisesRegex(
-            NotImplementedError,
-            ".*DeepSeek.*scaling.*only supported in CUDA for SM90.*"
+            expected_error,
+            expected_pattern
         ):
             scaled_mm_wrap(
                 x_fp8,
@@ -2108,6 +2157,126 @@ def test_blockwise_mxfp8_nvfp4_mxfp4_numerics(self, test_case_name, fast_accum,
             if sqnr.item() <= approx_match_sqnr_target:
                 raise AssertionError(f"sqnr {sqnr.item()} should be > {approx_match_sqnr_target}")
 
+    @unittest.skipIf(not PLATFORM_SUPPORTS_MX_GEMM or IS_WINDOWS, mx_skip_msg)
+    def test_passed_swizzle_arrays(self, device) -> None:
+        # Ensure that incorrectly-sized swizzle arrays are caught
+        M, N, K = 128, 128, 128
+
+        # MXFP8: swizzle=[SWIZZLE_32_4_4]
+        x = torch.randn(M, K, device=device).to(torch.float8_e4m3fn)
+        w = torch.randn(N, K, device=device).to(torch.float8_e4m3fn)
+
+        x_scale = torch.full((M, K // 32), 1., dtype=torch.float8_e8m0fnu, device=device)
+        w_scale = torch.full((N, K // 32), 1., dtype=torch.float8_e8m0fnu, device=device)
+
+        # No swizzle passed - must fail on swizzle_a
+        with self.assertRaisesRegex(
+            ValueError,
+            "swizzle_a and swizzle_b must each have 1 value"
+            if torch.version.hip
+            else "swizzle_a must have 1 value, got 0",
+        ):
+            _ = torch.nn.functional.scaled_mm(
+                x,
+                w.t(),
+                x_scale,
+                ScalingType.BlockWise1x32,
+                w_scale,
+                ScalingType.BlockWise1x32,
+            )
+
+        # swizzle_a passed, not b, must fail on swizzle_b
+        with self.assertRaisesRegex(
+            ValueError,
+            "swizzle_a and swizzle_b must each have 1 value"
+            if torch.version.hip
+            else "swizzle_b must have 1 value, got 0",
+        ):
+            _ = torch.nn.functional.scaled_mm(
+                x,
+                w.t(),
+                x_scale,
+                ScalingType.BlockWise1x32,
+                w_scale,
+                ScalingType.BlockWise1x32,
+                swizzle_a=SwizzleType.SWIZZLE_32_4_4,
+            )
+        if torch.version.hip:
+            with self.assertRaisesRegex(
+                ValueError,
+                "swizzle_a and swizzle_b must both be NO_SWIZZLE",
+            ):
+                _ = torch.nn.functional.scaled_mm(
+                    x,
+                    w.t(),
+                    x_scale,
+                    ScalingType.BlockWise1x32,
+                    w_scale,
+                    ScalingType.BlockWise1x32,
+                    swizzle_a=SwizzleType.NO_SWIZZLE,
+                    swizzle_b=SwizzleType.SWIZZLE_32_4_4,
+                )
+
+        # NVFP4 two-level: swizzle=[SWIZZLE_32_4_4, NO_SWIZZLE]
+        x = _bfloat16_to_float4_e2m1fn_x2(x.to(torch.bfloat16))
+        w = _bfloat16_to_float4_e2m1fn_x2(w.to(torch.bfloat16))
+
+        x_scale = torch.full((M, K // 16), 1., dtype=torch.float8_e4m3fn, device=device)
+        w_scale = torch.full((N, K // 16), 1., dtype=torch.float8_e4m3fn, device=device)
+
+        global_scale = torch.full((1, ), 1., dtype=torch.float, device=device)
+
+        # No swizzles passed - must fail on swizzle_a
+        with self.assertRaisesRegex(
+            NotImplementedError if torch.version.hip else ValueError,
+            "NVFP4 scaling not supported on ROCM"
+            if torch.version.hip
+            else "swizzle_a must have 2 values, got 0",
+        ):
+            _ = torch.nn.functional.scaled_mm(
+                x,
+                w.t(),
+                [x_scale, global_scale],
+                [ScalingType.BlockWise1x16, ScalingType.TensorWise],
+                [w_scale, global_scale],
+                [ScalingType.BlockWise1x16, ScalingType.TensorWise],
+            )
+
+        # Not enough swizzles passed - must fail on swizzle_a
+        with self.assertRaisesRegex(
+            NotImplementedError if torch.version.hip else ValueError,
+            "NVFP4 scaling not supported on ROCM"
+            if torch.version.hip
+            else "swizzle_a must have 2 values, got 1",
+        ):
+            _ = torch.nn.functional.scaled_mm(
+                x,
+                w.t(),
+                [x_scale, global_scale],
+                [ScalingType.BlockWise1x16, ScalingType.TensorWise],
+                [w_scale, global_scale],
+                [ScalingType.BlockWise1x16, ScalingType.TensorWise],
+                swizzle_a=[SwizzleType.SWIZZLE_32_4_4, ],
+            )
+
+        # Not enough swizzles passed to b - must fail on swizzle_b
+        with self.assertRaisesRegex(
+            NotImplementedError if torch.version.hip else ValueError,
+            "NVFP4 scaling not supported on ROCM"
+            if torch.version.hip
+            else "swizzle_b must have 2 values, got 1",
+        ):
+            _ = torch.nn.functional.scaled_mm(
+                x,
+                w.t(),
+                [x_scale, global_scale],
+                [ScalingType.BlockWise1x16, ScalingType.TensorWise],
+                [w_scale, global_scale],
+                [ScalingType.BlockWise1x16, ScalingType.TensorWise],
+                swizzle_a=[SwizzleType.SWIZZLE_32_4_4, SwizzleType.NO_SWIZZLE],
+                swizzle_b=[SwizzleType.SWIZZLE_32_4_4, ],
+            )
+
 
     @unittest.skipIf(not PLATFORM_SUPPORTS_MX_GEMM or IS_WINDOWS, mx_skip_msg)
     @parametrize("recipe", ["mxfp8", "mxfp4" if torch.version.hip else "nvfp4"])
@@ -2416,7 +2585,7 @@ def test_blockwise_nvfp4_compile(self) -> None:
         torch.testing.assert_close(C, C_ref, atol=0, rtol=0)
 
 
-instantiate_device_type_tests(TestFP8Matmul, globals(), except_for="cpu", allow_xpu=True)
+instantiate_device_type_tests(TestFP8Matmul, globals(), allow_xpu=True)
 
 if __name__ == '__main__':
     TestCase._default_dtype_check_enabled = True
diff --git a/test/test_scatter_gather_ops.py b/test/test_scatter_gather_ops.py
index ce2c83c5ef56c..8950bda3c0c76 100644
--- a/test/test_scatter_gather_ops.py
+++ b/test/test_scatter_gather_ops.py
@@ -13,7 +13,7 @@
 from torch.testing._internal.common_dtype import \
     (get_all_dtypes,)
 
-from torch.testing._internal.common_cuda import CDNA3OrLater
+from torch.testing._internal.common_cuda import gfx_arch_supports_opportunistic_fastatomics, SM90OrLater
 
 # Protects against includes accidentally setting the default dtype
 if torch.get_default_dtype() is not torch.float32:
@@ -221,7 +221,7 @@ def _test_scatter_base(self, fn, *, device, dtype, is_scalar, reduction,
         else:
             # When we are running opportunistic_fastatomics, we will expect some floating point rounding
             # errors as the order of operation is not guaranteed.
-            if TEST_WITH_ROCM and CDNA3OrLater() \
+            if TEST_WITH_ROCM and gfx_arch_supports_opportunistic_fastatomics() \
                     and not torch.are_deterministic_algorithms_enabled():
                 self.assertEqual(actual, expected, atol=1e-9, rtol=1e-6)
             else:
diff --git a/test/test_serialization.py b/test/test_serialization.py
index eb963e0450f74..51ff9182fa9af 100644
--- a/test/test_serialization.py
+++ b/test/test_serialization.py
@@ -1218,6 +1218,10 @@ def test_serialization_4gb_file(self):
 
         with BytesIOContext() as f:
             torch.save(big_model.state_dict(), f)
+            # Delete the original model before loading another one to
+            # reduce peak memory from ~12GB to ~8GB
+            del big_model
+            gc.collect()
             f.seek(0)
             torch.load(f)
 
diff --git a/test/test_shape_ops.py b/test/test_shape_ops.py
index c8a06a49b5975..48834a80922e1 100644
--- a/test/test_shape_ops.py
+++ b/test/test_shape_ops.py
@@ -390,6 +390,13 @@ def test_clamp_propagates_nans(self, device):
                     op(X, min_val, max_val, out=Y_out)
                     self.assertEqual(Y_expected, torch.isnan(Y_out))
 
+    def test_clamp_scalar_nan_bounds(self, device):
+        x = torch.ones(3, device=device)
+        y = torch.clamp(x, None, float("nan"))
+        self.assertTrue(torch.isnan(y).all())
+        y = torch.clamp(x, float("nan"), None)
+        self.assertTrue(torch.isnan(y).all())
+
     def test_clamp_raises_arg_errors(self, device):
         X = torch.randn(100, dtype=torch.float, device=device)
         error_msg = "At least one of 'min' or 'max' must not be None"
diff --git a/test/test_sort_and_select.py b/test/test_sort_and_select.py
index 5be1758186467..72fe7a52b19e6 100644
--- a/test/test_sort_and_select.py
+++ b/test/test_sort_and_select.py
@@ -25,6 +25,7 @@
     integral_types,
 )
 from torch.testing._internal.common_utils import (
+    parametrize,
     run_tests,
     skipIfTorchDynamo,
     slowTest,
@@ -1329,6 +1330,90 @@ def test_sort_overflow(self, device, dtype):
         finally:
             torch.set_num_threads(prev_num_threads)
 
+    @onlyCUDA
+    @dtypes(torch.float16, torch.bfloat16, torch.float32)
+    @slowTest
+    @largeTensorTest("170GB", "cpu")
+    @largeTensorTest("72GB", "cuda")
+    @parametrize("test_case", ["random", "identical"])
+    def test_topk_large_k(self, device, dtype, test_case):
+        """Test topk with k > 2^32 (integer overflow bug fix).
+
+        Tests both random and edge case (all identical values) inputs.
+        The edge case will force billions of elements to enter the same radix bin.
+
+        Memory requirements (in float32 case):
+        - GPU: ~72 GB (data ~16GB + values ~16GB + indices ~32GB + other ~8GB)
+        - CPU: ~170 GB (indices copy ~32GB + torch.unique extra memory ~130GB + other ~8GB)
+        """
+        extra = random.randint(500, 2000)
+        n = 2**32 + extra
+        k = random.randint(2**32 + 100, n - 100)
+        largest = random.choice([True, False])
+
+        # check correctness chunkwise later to avoid OOM
+        chunk_size = 100_000_000
+        num_chunks = (k + chunk_size - 1) // chunk_size
+
+        # pre-allocate GPU tensors
+        data = torch.empty((1, n), device=device, dtype=dtype)
+        gpu_values = torch.empty((1, k), device=device, dtype=dtype)
+        gpu_indices = torch.empty((1, k), device=device, dtype=torch.long)
+        gpu_chunk_values = torch.empty(chunk_size, device=device, dtype=dtype)
+
+        # pre-allocate CPU tensors
+        cpu_indices_copy = torch.empty(k, dtype=torch.long, device="cpu")
+
+        random_constant = random.random()
+        if test_case == "random":
+            data.random_()
+        else:
+            data.fill_(random_constant)
+
+        torch.topk(
+            data, k, dim=1, largest=largest, sorted=False, out=(gpu_values, gpu_indices)
+        )
+
+        indices = gpu_indices
+        values = gpu_values
+
+        # all indices must be in valid range
+        self.assertGreaterEqual(indices.min().item(), 0)
+        self.assertLess(indices.max().item(), n)
+
+        # all indices must be unique
+        cpu_indices_copy.copy_(indices.squeeze(0))
+        total_unique = torch.unique(cpu_indices_copy).numel()
+        self.assertEqual(
+            total_unique, k, f"Duplicates found in topk test: {k - total_unique}"
+        )
+        # for random case, values must match at returned indices (use pre-allocated tensor)
+        if test_case == "random":
+            for i in range(num_chunks):
+                start = i * chunk_size
+                end = min((i + 1) * chunk_size, k)
+                chunk_size_actual = end - start
+
+                chunk_indices = indices[0, start:end]
+                chunk_values = values[0, start:end]
+
+                actual_values = torch.index_select(
+                    data[0], 0, chunk_indices, out=gpu_chunk_values[:chunk_size_actual]
+                )
+                self.assertEqual(
+                    chunk_values, actual_values, msg=f"Value mismatch in chunk {i + 1}"
+                )
+
+        # for identical case, all values must equal to constant
+        if test_case == "identical":
+            # use .all() instead of self.assertEqualBroadcasting because the latter
+            # allocates a full-size tensor with the same size as values, causing OOM
+            self.assertEqual(
+                (values == random_constant).all().item(),
+                True,
+                "Values not equal to random constant",
+            )
+
 
 instantiate_device_type_tests(TestSortAndSelect, globals())
 
diff --git a/test/test_sparse.py b/test/test_sparse.py
index fcc4a7c7613e5..b444c71131772 100644
--- a/test/test_sparse.py
+++ b/test/test_sparse.py
@@ -23,12 +23,13 @@
 from torch.testing._internal.common_device_type import \
     (instantiate_device_type_tests, ops, dtypes, dtypesIfCUDA, dtypesIfMPS, onlyCPU, onlyCUDA, precisionOverride,
      deviceCountAtLeast, OpDTypes, onlyNativeDeviceTypes, skipCUDAIf, expectedFailureMPS,
-     expectedFailureMPSComplex, largeTensorTest)
+     largeTensorTest)
 from torch.testing._internal.common_methods_invocations import \
     (op_db, reduction_ops, sparse_unary_ufuncs, sparse_masked_reduction_ops, binary_ufuncs)
 from torch.testing._internal.common_dtype import (
     all_types, all_types_and_complex, all_mps_types, all_types_and_complex_and, floating_and_complex_types,
-    floating_and_complex_types_and, integral_types, floating_types_and,
+    floating_and_complex_types_and, highest_precision_complex, highest_precision_float,
+    integral_types, floating_types_and,
 )
 from torch.testing._internal.opinfo.definitions.sparse import validate_sample_input_sparse
 from torch.testing._internal.opinfo.refs import (
@@ -333,7 +334,7 @@ def _test_print(self, device, dtype, coalesced):
             if values.dtype == torch.double:
                 dtypes.append(torch.float)
             else:
-                dtypes.append(torch.double if values.device != torch.device("mps:0") else torch.float32)
+                dtypes.append(highest_precision_float(values.device))
             for dtype in dtypes:
                 printed.append(f"########## {dtype} ##########")
                 x = sp_tensor.detach().to(dtype)
@@ -583,7 +584,7 @@ def fn(x):
             x.requires_grad_(True)
             gradcheck(fn, (x,))
 
-        values_types = [torch.double, torch.cdouble] if device != "mps:0" else [torch.float32, torch.complex64]
+        values_types = [highest_precision_float(device), highest_precision_complex(device)]
         for value_type in values_types:
             i = self.index_tensor([
                 [0, 1, 2, 2],
@@ -628,7 +629,7 @@ def fn(x):
     def test_to_sparse(self, device, dtype, coalesced):
         shape = [5, 2, 10, 4]
         max_nnz = 1
-        dtypes = [torch.double, torch.cdouble] if device != "mps:0" else [torch.float32, torch.complex64]
+        dtypes = [highest_precision_float(device), highest_precision_complex(device)]
         for value_type in dtypes:
             for dim, dim_sz in enumerate(shape, 1):
                 max_nnz *= dim_sz
@@ -963,7 +964,7 @@ def test_Sparse_to_Sparse_copy_(self, device, dtype, coalesced):
         x1.copy_(x2)
         self.assertEqual(x1_dtype, x1.dtype)
 
-        x2 = x2.to(torch.float64) if device != "mps:0" else x2.to(torch.float32)
+        x2 = x2.to(highest_precision_float(device))
         x1_dtype = x1.dtype
         x1.copy_(x2)
         self.assertEqual(x1_dtype, x1.dtype)
@@ -1525,12 +1526,16 @@ def test_shape(num_mats, dim_i, dim_j, dim_k, nnz):
         test_shape(10, 10, 100, 0, 20)
 
         a = torch.rand([10, 23, 32], dtype=dtype, device=device)
-        a[3] = torch.zeros(23, 32, dtype=dtype, device=device)
-        a[6] = torch.zeros(23, 32, dtype=dtype, device=device)
+        a[3].zero_()
+        a[6].zero_()
+        a[7].tril_()
+        a[8].zero_().diagonal().fill_(1)
         a = a.to_sparse()
         b = torch.rand([10, 32, 10], dtype=dtype, device=device)
-        b[4] = torch.zeros(32, 10, dtype=dtype, device=device)
-        b[6] = torch.zeros(32, 10, dtype=dtype, device=device)
+        b[4].zero_()
+        b[6].zero_()
+        b[7].tril_()
+        b[8].zero_().diagonal().fill_(1)
         ab = a.bmm(b)
         for mat_idx in range(ab.size(0)):
             ab_mat = ab[mat_idx]
@@ -1617,6 +1622,42 @@ def test_bmm_windows_error(self, device, dtype):
                 "bmm sparse-dense CUDA is not supported on Windows with cuda before 11.0"):
             ab = a.bmm(b)
 
+    @onlyCUDA
+    @unittest.skipIf(
+        IS_WINDOWS and TEST_CUDA,
+        "bmm sparse-dense CUDA is not yet supported in Windows, at least up to CUDA 10.1"
+    )
+    @dtypes(torch.double)
+    @dtypesIfMPS(torch.float32)
+    def test_bmm_coo_row_index_alignment(self, device, dtype):
+        # See https://github.com/pytorch/pytorch/issues/167901
+        # cuSPARSE SpMV has a bug before cuda 13.1 which assumes
+        # 16-byte alignment for COO row index array.
+        batch = 5
+        nrows = 291105
+
+        m1 = torch.randn(batch, nrows, 1, device=device, dtype=dtype)
+        m2 = torch.randn(batch, 1, 1, device=device, dtype=dtype)
+        ref = torch.bmm(m1, m2)
+
+        m1_sparse = m1.to_sparse()
+        row_indices = m1_sparse.indices()[1]
+        del m1
+
+        # Row indices ptr for odd matrices should be 16-byte aligned
+        for i in range(1, batch, 2):
+            self.assertTrue(
+                (row_indices[i * nrows].data_ptr() % 16) == 0
+            )
+        # Row indices ptr for even matrices should be 8-byte aligned
+        for i in range(0, batch, 2):
+            self.assertTrue(
+                (row_indices[i * nrows].data_ptr() % 16) == 8
+            )
+
+        res = torch.bmm(m1_sparse, m2)
+        self.assertEqual(ref, res)
+
     @onlyCPU
     @coalescedonoff
     @dtypes(torch.double, torch.cdouble)
@@ -1921,7 +1962,6 @@ def test_sparse_add_out_bfloat16(self, device, dtype, coalesced):
         self.assertEqual(res_fp32, res_bf16, atol=1e-2, rtol=0)
 
     @coalescedonoff
-    @expectedFailureMPSComplex
     @dtypes(torch.double, torch.cdouble)
     @dtypesIfMPS(torch.float32, torch.complex64)
     def test_norm(self, device, dtype, coalesced):
@@ -3332,12 +3372,14 @@ def test_is_nonzero(self, device):
             torch.sparse_coo_tensor(([0, 1],), torch.empty(2, 0), (4, 0), device=device).is_nonzero()
         self.assertTrue(torch.sparse_coo_tensor(([0],), 2.3 - 4.5j, (1,), dtype=torch.cfloat, device=device)
                         .is_nonzero())
-        self.assertTrue(torch.sparse_coo_tensor(([0],), 2.3 - 4.5j, (1,), dtype=torch.cdouble, device=device)
-                        .is_nonzero())
+        if device != 'mps:0':
+            self.assertTrue(torch.sparse_coo_tensor(([0],), 2.3 - 4.5j, (1,), dtype=torch.cdouble, device=device)
+                            .is_nonzero())
         self.assertFalse(torch.sparse_coo_tensor(([0],), 0. + 0j, (1,), dtype=torch.cfloat, device=device)
                          .is_nonzero())
-        self.assertFalse(torch.sparse_coo_tensor(([0],), 0. + 0j, (1,), dtype=torch.cdouble, device=device)
-                         .is_nonzero())
+        if device != 'mps:0':
+            self.assertFalse(torch.sparse_coo_tensor(([0],), 0. + 0j, (1,), dtype=torch.cdouble, device=device)
+                             .is_nonzero())
 
     @dtypes(torch.double, torch.cdouble)
     @dtypesIfMPS(torch.float32, torch.complex64)
diff --git a/test/test_sparse_csr.py b/test/test_sparse_csr.py
index e1f33e2b20381..7a65fce17578f 100644
--- a/test/test_sparse_csr.py
+++ b/test/test_sparse_csr.py
@@ -1030,6 +1030,16 @@ def _npref_block_addmm_addmv(c, a, b, alpha, beta):
 
 class TestSparseCSR(TestCase):
 
+    @onlyCPU
+    @dtypes(*all_types_and_complex_and(torch.half, torch.bool, torch.bfloat16))
+    def test_empty_plain_indices_with_stride_zero(self, device, dtype):
+        # Test that empty plain_indices with stride 0 works.
+        crow_indices = torch.tensor([0, 0], dtype=torch.int32, device=device)
+        col_indices = torch.as_strided(torch.empty((0,), device=device, dtype=torch.int32), (0,), (0,))
+        values = torch.empty(0, dtype=dtype, device=device)
+        t = torch.sparse_csr_tensor(crow_indices, col_indices, values, (1, 100), dtype=dtype, device=device)
+        self.assertEqual(t._nnz(), 0)
+
     def test_csr_stride(self):
         a = self.genSparseCSRTensor((3, 3), 3, dtype=torch.float, device=self.device_type, index_dtype=torch.int64)
 
@@ -1743,7 +1753,7 @@ def test_addmv(self, device, dtype, matrix_shape):
         mat[mat.real < 0] = 0
         sparse_mat = mat.to_sparse_csr()
         mvec = torch.randn((mat.size(1),), dtype=dtype, device=device)
-        avec = torch.randn((mat.size(0),), dtype=torch.float64, device=device)
+        avec = torch.randn((mat.size(0),), dtype=dtype, device=device)
         ref_output = torch.addmv(avec, mat, mvec)
         output = torch.addmv(avec, sparse_mat, mvec)
         self.assertEqual(ref_output, output)
diff --git a/test/test_sparse_semi_structured.py b/test/test_sparse_semi_structured.py
index 53c8409282a12..881470944aee3 100644
--- a/test/test_sparse_semi_structured.py
+++ b/test/test_sparse_semi_structured.py
@@ -4,45 +4,43 @@
 import random
 import unittest
 
+import pytest
 import torch
-from torch import nn
+import torch._dynamo.test_case
 import torch.nn.functional as F
-
+from torch import nn
 from torch.sparse import (
     SparseSemiStructuredTensor,
     SparseSemiStructuredTensorCUSPARSELT,
     SparseSemiStructuredTensorCUTLASS,
     to_sparse_semi_structured,
 )
-
 from torch.sparse._semi_structured_conversions import (
-    sparse_semi_structured_from_dense_cutlass,
-    _sparse_semi_structured_tile,
     _compute_compressed_swizzled_bitmask,
+    _sparse_semi_structured_tile,
+    sparse_semi_structured_from_dense_cutlass,
 )
-
 from torch.testing import make_tensor
-from torch.testing._internal.common_cuda import PLATFORM_SUPPORTS_FP8, xfailIfSM89PreCUDA13
+from torch.testing._internal.common_cuda import (
+    PLATFORM_SUPPORTS_FP8,
+    PLATFORM_SUPPORTS_FP8_SPARSE,
+    xfailIfSM89PreCUDA13,
+)
 from torch.testing._internal.common_device_type import (
     dtypes,
     instantiate_device_type_tests,
 )
-
 from torch.testing._internal.common_dtype import all_types_and_complex
-import torch._dynamo.test_case
 from torch.testing._internal.common_utils import (
+    IS_WINDOWS,
     parametrize,
     run_tests,
     subtest,
-    TestCase,
     TEST_WITH_ROCM,
-    IS_WINDOWS,
+    TestCase,
 )
-
 from torch.testing._internal.inductor_utils import HAS_GPU
 
-import pytest
-
 SEMI_STRUCTURED_SUPPORTED_BACKENDS = dict()
 
 _IS_SM8X = False
@@ -50,16 +48,28 @@
 _IS_HIPSPARSELT_AVAILABLE = False
 
 if torch.cuda.is_available():
-    _IS_SM8X = torch.version.cuda is not None and (torch.cuda.get_device_capability(0)[0] == 8)
-    _IS_SM9X = torch.version.cuda is not None and (torch.cuda.get_device_capability(0)[0] == 9)
-    _IS_HIPSPARSELT_AVAILABLE = torch.version.hip is not None and tuple(int(v) for v in torch.version.hip.split('.')[:2]) > (6, 4)
+    _IS_SM8X = torch.version.cuda is not None and (
+        torch.cuda.get_device_capability(0)[0] == 8
+    )
+    _IS_SM9X = torch.version.cuda is not None and (
+        torch.cuda.get_device_capability(0)[0] == 9
+    )
+    _IS_HIPSPARSELT_AVAILABLE = torch.version.hip is not None and tuple(
+        int(v) for v in torch.version.hip.split(".")[:2]
+    ) >= (7, 12)
     # CUTLASS kernels only work for Ampere
     if _IS_SM8X:
-        SEMI_STRUCTURED_SUPPORTED_BACKENDS["cutlass"] = SparseSemiStructuredTensorCUTLASS
+        SEMI_STRUCTURED_SUPPORTED_BACKENDS["cutlass"] = (
+            SparseSemiStructuredTensorCUTLASS
+        )
 
     # add cuSPASRELt tests if available
-    if torch.backends.cusparselt.is_available() and (_IS_SM8X or _IS_SM9X or _IS_HIPSPARSELT_AVAILABLE):
-        SEMI_STRUCTURED_SUPPORTED_BACKENDS["cusparselt"] = SparseSemiStructuredTensorCUSPARSELT
+    if torch.backends.cusparselt.is_available() and (
+        _IS_SM8X or _IS_SM9X or _IS_HIPSPARSELT_AVAILABLE
+    ):
+        SEMI_STRUCTURED_SUPPORTED_BACKENDS["cusparselt"] = (
+            SparseSemiStructuredTensorCUSPARSELT
+        )
 
 inference_dtypes = dtypes(torch.float16, torch.bfloat16, torch.int8)
 training_dtypes = dtypes(torch.float16, torch.bfloat16)
@@ -76,13 +86,16 @@
     },
 }
 
+
 def sparse24_largest_mask_2d(original):
     sparse = SparseSemiStructuredTensorCUTLASS.prune_dense_static_sort(original)
     return sparse.to_dense().bool()
 
+
 def sparsify24_dense(original):
     return sparse24_largest_mask_2d(original) * original
 
+
 def rand_sparse_semi_structured_mask(
     r, c, dtype=torch.float16, device="cuda", choice=None
 ):
@@ -100,15 +113,13 @@ def rand_sparse_semi_structured_mask(
         .contiguous()
     )
 
+
 def rand_sparse_semi_structured(r, c, dtype, device, choice=None):
-    pattern = '2by4' if dtype != torch.float32 else '1by2'
-    if pattern == '1by2':
+    pattern = "2by4" if dtype != torch.float32 else "1by2"
+    if pattern == "1by2":
         ksparse = 2
-        choices = [
-            [0, 1],
-            [1, 0]
-        ]
-    elif pattern == '2by4':
+        choices = [[0, 1], [1, 0]]
+    elif pattern == "2by4":
         ksparse = 4
         choices = [
             [1, 1, 0, 0],
@@ -116,7 +127,7 @@ def rand_sparse_semi_structured(r, c, dtype, device, choice=None):
             [1, 0, 0, 1],
             [0, 1, 1, 0],
             [0, 1, 0, 1],
-            [0, 0, 1, 1]
+            [0, 0, 1, 1],
         ]
     mask_entries = [choice or random.choice(choices) for i in range(r * c // ksparse)]
     mask = torch.tensor(mask_entries, dtype=torch.bool).view(r, c).to(device)
@@ -127,16 +138,16 @@ def rand_sparse_semi_structured(r, c, dtype, device, choice=None):
 
 
 def rand_sparse_semi_structured_all_patterns(r, c, dtype, device):
-    pattern = '2by4' if dtype != torch.float32 else '1by2'
-    if pattern == '1by2':
+    pattern = "2by4" if dtype != torch.float32 else "1by2"
+    if pattern == "1by2":
         ksparse = 2
         choices = [
             [[0, 0], [0, 1]],
             [[0, 1], [0, 1]],
             [[1, 0], [1, 0]],
-            [[1, 1], [1, 0]]
+            [[1, 1], [1, 0]],
         ]
-    elif pattern == '2by4':
+    elif pattern == "2by4":
         ksparse = 4
         choices = [
             [[0, 0, 0, 0], [0, 0, 1, 1]],
@@ -164,7 +175,7 @@ def rand_sparse_semi_structured_all_patterns(r, c, dtype, device):
     mask_inv = torch.tensor(mask_entries_inv, dtype=torch.bool).view(r, c).to(device)
     mask_val = torch.tensor(mask_entries_val, dtype=torch.bool).view(r, c).to(device)
     dense = make_tensor(r, c, dtype=dtype, device=device)
-    dense[dense == 0] = 1   # To prevent zeros except where mask below applied.
+    dense[dense == 0] = 1  # To prevent zeros except where mask below applied.
     dense_inv = dense.masked_fill(~mask_inv, 0)
     dense_val = dense_inv.masked_fill(~mask_val, 0)
 
@@ -172,10 +183,9 @@ def rand_sparse_semi_structured_all_patterns(r, c, dtype, device):
 
 
 class SparseSemiStructuredTensorCompileTest(torch._dynamo.test_case.TestCase):
-
     def setUp(self):
         if len(SEMI_STRUCTURED_SUPPORTED_BACKENDS) == 0:
-            self.skipTest('semi-structured sparsity has no available backend!')
+            self.skipTest("semi-structured sparsity has no available backend!")
         super().setUp()
 
     def tearDown(self):
@@ -209,32 +219,45 @@ def forward(self, x):
         mod_linear.weight = nn.Parameter(mod_linear.weight * mask)
 
         dense_result = model(input)
-        mod_linear.weight = nn.Parameter(SEMI_STRUCTURED_SUPPORTED_BACKENDS[backend].from_dense(mod_linear.weight))
+        mod_linear.weight = nn.Parameter(
+            SEMI_STRUCTURED_SUPPORTED_BACKENDS[backend].from_dense(mod_linear.weight)
+        )
         sparse_result = model(input)
 
         model = torch.compile(model, backend="inductor", fullgraph=True)
         sparse_compile_result = model(input)
 
         # test that sparse_compile_result and dense_result are numerically close
-        torch.testing.assert_close(dense_result, sparse_compile_result, rtol=1e-3, atol=1e-3)
+        torch.testing.assert_close(
+            dense_result, sparse_compile_result, rtol=1e-3, atol=1e-3
+        )
         # assert sparse and sparse_compile have the same strides,
         # as meta registrations may return contiguous tensors when the output is transposed
         # https://github.com/pytorch/pytorch/pull/114477
         if sparse_result.stride() != sparse_compile_result.stride():
-            raise AssertionError(f"stride mismatch: {sparse_result.stride()} != {sparse_compile_result.stride()}")
+            raise AssertionError(
+                f"stride mismatch: {sparse_result.stride()} != {sparse_compile_result.stride()}"
+            )
 
     @unittest.skipIf(IS_WINDOWS, "torch.compile not supported on windows")
-    @unittest.skipIf("cusparselt" not in SEMI_STRUCTURED_SUPPORTED_BACKENDS, "cusparselt not supported on this machine")
+    @unittest.skipIf(
+        "cusparselt" not in SEMI_STRUCTURED_SUPPORTED_BACKENDS,
+        "cusparselt not supported on this machine",
+    )
     @unittest.skipIf(TEST_WITH_ROCM, "Not supported on ROCm")
     def test_mlp_contiguous_relu_compile_cusparselt(self):
         """
         test for cuSPASRELt meta registrations (_cslt_sparse_mm) + torch.compile
         """
         for dense_input_shape in [(1, 128), (64, 128), (128, 128), (64, 128, 128)]:
-            SparseSemiStructuredTensorCompileTest._test_mlp_contiguous_relu_compile("cusparselt", dense_input_shape)
-
+            SparseSemiStructuredTensorCompileTest._test_mlp_contiguous_relu_compile(
+                "cusparselt", dense_input_shape
+            )
 
-    @unittest.skipIf("cutlass" not in SEMI_STRUCTURED_SUPPORTED_BACKENDS, "cutlass not supported on this machine")
+    @unittest.skipIf(
+        "cutlass" not in SEMI_STRUCTURED_SUPPORTED_BACKENDS,
+        "cutlass not supported on this machine",
+    )
     @unittest.skipIf(IS_WINDOWS, "torch.compile not supported on windows")
     @unittest.skipIf(TEST_WITH_ROCM, "Not supported on ROCm")
     def test_mlp_contiguous_relu_compile_cutlass(self):
@@ -242,18 +265,24 @@ def test_mlp_contiguous_relu_compile_cutlass(self):
         test for CUTLASS meta registrations (_sparse_semi_structured_addmm) + torch.compile
         """
         for dense_input_shape in [(1, 128), (64, 128), (128, 128), (64, 128, 128)]:
-            SparseSemiStructuredTensorCompileTest._test_mlp_contiguous_relu_compile("cutlass", dense_input_shape)
-
+            SparseSemiStructuredTensorCompileTest._test_mlp_contiguous_relu_compile(
+                "cutlass", dense_input_shape
+            )
 
     @unittest.skipIf(IS_WINDOWS, "torch.compile not supported on windows")
-    @unittest.skipIf("cusparselt" not in SEMI_STRUCTURED_SUPPORTED_BACKENDS, "cusparselt not supported on this machine")
+    @unittest.skipIf(
+        "cusparselt" not in SEMI_STRUCTURED_SUPPORTED_BACKENDS,
+        "cusparselt not supported on this machine",
+    )
     @unittest.skipIf(TEST_WITH_ROCM, "Not supported on ROCm")
     @unittest.skipIf(
         "RelWithAssert" in torch.__config__.show(),
         "failing in debug build, see https://github.com/pytorch/pytorch/pull/165158 for context",
     )
     def test_sp24_compile(self) -> None:
-        x = torch.randn([1024, 512], device="cuda", dtype=torch.float16, requires_grad=True)
+        x = torch.randn(
+            [1024, 512], device="cuda", dtype=torch.float16, requires_grad=True
+        )
 
         def fn(x):
             y = SparseSemiStructuredTensorCUSPARSELT.prune_dense_static_sort(x)
@@ -267,18 +296,18 @@ def fn(x):
         output = torch.compile(fn)(x)
         output.backward(output)
 
-class TestSparseSemiStructured(TestCase):
 
+class TestSparseSemiStructured(TestCase):
     def setUp(self):
         if len(SEMI_STRUCTURED_SUPPORTED_BACKENDS) == 0:
-            self.skipTest('semi-structured sparsity has no available backend!')
+            self.skipTest("semi-structured sparsity has no available backend!")
         if IS_WINDOWS:
             self.skipTest("torch.compile not supported on windows")
 
     @inference_dtypes
     @parametrize_backends
     def test_to_sparse_semi_structured(self, dtype, backend):
-        SparseSemiStructuredTensor._FORCE_CUTLASS = (backend == "cutlass")
+        SparseSemiStructuredTensor._FORCE_CUTLASS = backend == "cutlass"
         A = rand_sparse_semi_structured_mask(128, 256, dtype=dtype)
         A_sparse = to_sparse_semi_structured(A)
 
@@ -292,7 +321,9 @@ def test_to_sparse_semi_structured(self, dtype, backend):
         if not isinstance(A, torch.Tensor):
             raise AssertionError(f"A should be torch.Tensor, got {type(A)}")
         if not isinstance(A_sparse, SparseSemiStructuredTensor):
-            raise AssertionError(f"A_sparse should be SparseSemiStructuredTensor, got {type(A_sparse)}")
+            raise AssertionError(
+                f"A_sparse should be SparseSemiStructuredTensor, got {type(A_sparse)}"
+            )
 
     @inference_dtypes
     @parametrize_backends
@@ -301,7 +332,7 @@ def test_mm_sparse_first_NN(self, dense_input_shape, dtype, device, backend):
         """
         Ensure torch.mm(A_sparse, B) is correct for float16 and will throw error for int8
         """
-        SparseSemiStructuredTensor._FORCE_CUTLASS = (backend == "cutlass")
+        SparseSemiStructuredTensor._FORCE_CUTLASS = backend == "cutlass"
         A = rand_sparse_semi_structured_mask(256, 128, dtype=dtype)
         A_sparse = to_sparse_semi_structured(A)
 
@@ -310,16 +341,26 @@ def test_mm_sparse_first_NN(self, dense_input_shape, dtype, device, backend):
         # Currently we don't support int matmul on GPU, so evaluate on CPU and copy over
         if dtype is torch.int8:
             if backend == "cutlass":
-                with self.assertRaisesRegex(RuntimeError, "spgemm_cutlass_dispatch_layouts"):
+                with self.assertRaisesRegex(
+                    RuntimeError, "spgemm_cutlass_dispatch_layouts"
+                ):
                     sparse_result = torch.mm(A_sparse, B)
             else:
-                with self.assertRaisesRegex(RuntimeError,
-                                            "CUDA error: operation not supported when calling `cusparseLtMatmulDescriptorInit"):
+                if torch.version.hip:
+                    self.skipTest(
+                        "Skipping int8 sparse mm (NN, cuSPARSELt) test on ROCm"
+                    )
+                with self.assertRaisesRegex(
+                    RuntimeError,
+                    "CUDA error: operation not supported when calling `cusparseLtMatmulDescriptorInit",
+                ):
                     sparse_result = torch.mm(A_sparse, B)
         else:
             dense_result = torch.mm(A, B)
             sparse_result = torch.mm(A_sparse, B)
-            torch.testing.assert_close(dense_result, sparse_result, rtol=1e-3, atol=1e-3)
+            torch.testing.assert_close(
+                dense_result, sparse_result, rtol=1e-3, atol=1e-3
+            )
 
     @inference_dtypes
     @parametrize_backends
@@ -329,7 +370,7 @@ def test_mm_sparse_first_NT(self, dense_input_shape, dtype, device, backend):
         Ensure torch.mm(A_sparse, B.t()) is correct for float16/bfloat16
         and will throw an error for int8 + padding
         """
-        SparseSemiStructuredTensor._FORCE_CUTLASS = (backend == "cutlass")
+        SparseSemiStructuredTensor._FORCE_CUTLASS = backend == "cutlass"
         A = rand_sparse_semi_structured_mask(256, 128, dtype=dtype)
         A_sparse = to_sparse_semi_structured(A)
 
@@ -340,22 +381,34 @@ def test_mm_sparse_first_NT(self, dense_input_shape, dtype, device, backend):
             # padding with int8 throws an error because transposing B yields a contiguous output
             # and row-row 2:4 sparse @ dense with NN is not supported by cuSPARSELt or CUTLASS.
             if backend == "cutlass":
-                with self.assertRaisesRegex(RuntimeError, "spgemm_cutlass_dispatch_layouts"):
+                with self.assertRaisesRegex(
+                    RuntimeError, "spgemm_cutlass_dispatch_layouts"
+                ):
                     sparse_result = torch.mm(A_sparse, B.t())
             else:
-                with self.assertRaisesRegex(RuntimeError,
-                                            "CUDA error: operation not supported when calling `cusparseLtMatmulDescriptorInit"):
+                if torch.version.hip:
+                    self.skipTest(
+                        "Skipping int8 sparse mm (NT, cusparselt, shape=(1,128)) test on ROCm"
+                    )
+                with self.assertRaisesRegex(
+                    RuntimeError,
+                    "CUDA error: operation not supported when calling `cusparseLtMatmulDescriptorInit",
+                ):
                     sparse_result = torch.mm(A_sparse, B.t())
         elif dtype is torch.int8:
             # test transpose
             dense_result = torch.mm(A.cpu(), B.t().cpu()).to(device, dtype=torch.int8)
             sparse_result = torch.mm(A_sparse, B.t())
-            torch.testing.assert_close(dense_result, sparse_result, rtol=1e-3, atol=1e-3)
+            torch.testing.assert_close(
+                dense_result, sparse_result, rtol=1e-3, atol=1e-3
+            )
         else:
             # test transpose
             dense_result = torch.mm(A, B.t())
             sparse_result = torch.mm(A_sparse, B.t())
-            torch.testing.assert_close(dense_result, sparse_result, rtol=1e-3, atol=1e-3)
+            torch.testing.assert_close(
+                dense_result, sparse_result, rtol=1e-3, atol=1e-3
+            )
 
     @inference_dtypes
     @parametrize("dense_input_shape", [(1, 128), (64, 128), (128, 128)])
@@ -364,7 +417,7 @@ def test_mm_sparse_first_TN(self, dtype, dense_input_shape, device, backend):
         """
         Ensure torch.mm(A_sparse.t(), B) throws error
         """
-        SparseSemiStructuredTensor._FORCE_CUTLASS = (backend == "cutlass")
+        SparseSemiStructuredTensor._FORCE_CUTLASS = backend == "cutlass"
         if backend == "cutlass" and IS_WINDOWS:
             self.skipTest("CUTLASS not supported on Windows")
         A = rand_sparse_semi_structured_mask(128, 256, dtype=dtype)
@@ -385,7 +438,7 @@ def test_mm_sparse_second_NT(self, dense_input_shape, dtype, device, backend):
         """
         Ensure torch.mm(A, B_sparse.t()) is correct
         """
-        SparseSemiStructuredTensor._FORCE_CUTLASS = (backend == "cutlass")
+        SparseSemiStructuredTensor._FORCE_CUTLASS = backend == "cutlass"
         if backend == "cutlass" and IS_WINDOWS:
             self.skipTest("CUTLASS not supported on Windows")
         B = rand_sparse_semi_structured_mask(256, 128, dtype=dtype)
@@ -410,7 +463,7 @@ def test_mm_sparse_second_NN(self, dense_input_shape, dtype, device, backend):
         """
         Ensure torch.mm(A, B_sparse) throws error
         """
-        SparseSemiStructuredTensor._FORCE_CUTLASS = (backend == "cutlass")
+        SparseSemiStructuredTensor._FORCE_CUTLASS = backend == "cutlass"
         if backend == "cutlass" and IS_WINDOWS:
             self.skipTest("CUTLASS not supported on Windows")
         B = rand_sparse_semi_structured_mask(256, 128, dtype=dtype)
@@ -431,7 +484,7 @@ def test_linear(self, dense_input_shape, inference_mode, device, backend):
         """
         Test nn.Linear has the same numerics
         """
-        SparseSemiStructuredTensor._FORCE_CUTLASS = (backend == "cutlass")
+        SparseSemiStructuredTensor._FORCE_CUTLASS = backend == "cutlass"
         if backend == "cutlass" and IS_WINDOWS:
             self.skipTest("CUTLASS not supported on Windows")
         input = torch.rand((dense_input_shape), device=device).half()
@@ -456,7 +509,7 @@ def test_linear(self, dense_input_shape, inference_mode, device, backend):
     @parametrize("dense_input_shape", [(1, 128), (64, 128), (128, 128), (64, 128, 128)])
     @parametrize_backends
     def test_mlp(self, device, dense_input_shape, backend):
-        SparseSemiStructuredTensor._FORCE_CUTLASS = (backend == "cutlass")
+        SparseSemiStructuredTensor._FORCE_CUTLASS = backend == "cutlass"
         input = torch.rand(dense_input_shape, device=device).half()
         model = (
             nn.Sequential(
@@ -486,34 +539,44 @@ def test_mlp(self, device, dense_input_shape, backend):
 
     @parametrize_backends
     def test_values(self, backend):
-        SparseSemiStructuredTensor._FORCE_CUTLASS = (backend == "cutlass")
+        SparseSemiStructuredTensor._FORCE_CUTLASS = backend == "cutlass"
         if backend == "cutlass" and IS_WINDOWS:
             self.skipTest("CUTLASS not supported on Windows")
         A = rand_sparse_semi_structured_mask(128, 128)
         A_sparse = to_sparse_semi_structured(A)
         if A_sparse.values().shape != (128, 64):
-            raise AssertionError(f"values shape should be (128, 64), got {A_sparse.values().shape}")
+            raise AssertionError(
+                f"values shape should be (128, 64), got {A_sparse.values().shape}"
+            )
         if not (A_sparse.values() == 1).all():
             raise AssertionError("values should all be 1")
 
     @parametrize_backends
     def test_indices(self, backend):
-        SparseSemiStructuredTensor._FORCE_CUTLASS = (backend == "cutlass")
+        SparseSemiStructuredTensor._FORCE_CUTLASS = backend == "cutlass"
         if backend == "cutlass" and IS_WINDOWS:
             self.skipTest("CUTLASS not supported on Windows")
         A = rand_sparse_semi_structured_mask(128, 128)
         A_sparse = to_sparse_semi_structured(A)
         if A_sparse.indices().shape != (128, 8):
-            raise AssertionError(f"indices shape should be (128, 8), got {A_sparse.indices().shape}")
+            raise AssertionError(
+                f"indices shape should be (128, 8), got {A_sparse.indices().shape}"
+            )
 
     @inference_dtypes
     @parametrize_backends
     def test_min_sparse_shape(self, dtype, device, backend):
-        SparseSemiStructuredTensor._FORCE_CUTLASS = (backend == "cutlass")
-        config = SEMI_STRUCTURED_SUPPORTED_BACKENDS[backend]._DTYPE_SHAPE_CONSTRAINTS[dtype]
-        A = rand_sparse_semi_structured_mask(config.sparse_min_rows, config.sparse_min_cols, dtype=dtype, device=device)
+        SparseSemiStructuredTensor._FORCE_CUTLASS = backend == "cutlass"
+        config = SEMI_STRUCTURED_SUPPORTED_BACKENDS[backend]._DTYPE_SHAPE_CONSTRAINTS[
+            dtype
+        ]
+        A = rand_sparse_semi_structured_mask(
+            config.sparse_min_rows, config.sparse_min_cols, dtype=dtype, device=device
+        )
         A_sparse = to_sparse_semi_structured(A)
-        B = torch.rand((config.sparse_min_cols, config.dense_min_cols), device=device).to(dtype)
+        B = torch.rand(
+            (config.sparse_min_cols, config.dense_min_cols), device=device
+        ).to(dtype)
         if dtype == torch.int8:
             dense_res = torch.mm(A.cpu(), B.cpu()).to(device, dtype=torch.int8)
             # int8 sparse matmul not supported for R/R -> R layout, so we transpose one of the arguments to get R/C -> R
@@ -527,7 +590,7 @@ def test_min_sparse_shape(self, dtype, device, backend):
     @inference_dtypes
     @parametrize_backends
     def test_unsupported_shape(self, dtype, device, backend):
-        SparseSemiStructuredTensor._FORCE_CUTLASS = (backend == "cutlass")
+        SparseSemiStructuredTensor._FORCE_CUTLASS = backend == "cutlass"
         if backend == "cutlass" and IS_WINDOWS:
             self.skipTest("CUTLASS not supported on Windows")
         A = rand_sparse_semi_structured_mask(2, 2, dtype=dtype, device=device)
@@ -537,12 +600,15 @@ def test_unsupported_shape(self, dtype, device, backend):
     @dtypes(*all_types_and_complex())
     @parametrize_backends
     def test_unsupported_dtype(self, dtype, device, backend):
-        SparseSemiStructuredTensor._FORCE_CUTLASS = (backend == "cutlass")
+        SparseSemiStructuredTensor._FORCE_CUTLASS = backend == "cutlass"
         if backend == "cutlass" and IS_WINDOWS:
             self.skipTest("CUTLASS not supported on Windows")
         A = rand_sparse_semi_structured_mask(128, 128, dtype=dtype, device=device)
 
-        if dtype not in SEMI_STRUCTURED_SUPPORTED_BACKENDS[backend]._DTYPE_SHAPE_CONSTRAINTS:
+        if (
+            dtype
+            not in SEMI_STRUCTURED_SUPPORTED_BACKENDS[backend]._DTYPE_SHAPE_CONSTRAINTS
+        ):
             with self.assertRaisesRegex(RuntimeError, "Error original_tensor.dtype"):
                 A_sparse = to_sparse_semi_structured(A)
         else:
@@ -550,7 +616,7 @@ def test_unsupported_dtype(self, dtype, device, backend):
 
     @parametrize_backends
     def test_unsupported_dim(self, device, backend):
-        SparseSemiStructuredTensor._FORCE_CUTLASS = (backend == "cutlass")
+        SparseSemiStructuredTensor._FORCE_CUTLASS = backend == "cutlass"
         if backend == "cutlass" and IS_WINDOWS:
             self.skipTest("CUTLASS not supported on Windows")
         A = torch.rand(128, 128, 128, device=device, dtype=torch.float16)
@@ -577,18 +643,23 @@ def create_random_mask(shape) -> torch.Tensor:
             mask[line, col : col + 4] = torch.tensor(sparsity, dtype=torch.bool)
     return mask
 
-class TestSparseSemiStructuredTraining(TestCase):
 
+class TestSparseSemiStructuredTraining(TestCase):
     def setUp(self):
         if not _IS_SM8X:
-            self.skipTest("SparseSemiStructuredTensor training only supported on SM8x (Ampere)")
+            self.skipTest(
+                "SparseSemiStructuredTensor training only supported on SM8x (Ampere)"
+            )
 
         if IS_WINDOWS:
-            self.skipTest('CUTLASS not supported on windows')
-
+            self.skipTest("CUTLASS not supported on windows")
 
     @training_dtypes
     @unittest.skipIf(TEST_WITH_ROCM, "Not supported on ROCm")
+    @unittest.skipIf(
+        not torch.backends.cusparselt.is_available(),
+        "cuSPARSELt not available",
+    )
     @unittest.skipIf(
         "RelWithAssert" in torch.__config__.show(),
         "failing in debug build, see https://github.com/pytorch/pytorch/pull/165158 for context",
@@ -600,40 +671,57 @@ def test_prune_dense_static_sort(self, dtype) -> None:
         pruned = _sparse_semi_structured_tile(dense)
 
         # CUTLASS
-        reference_cutlass = SparseSemiStructuredTensorCUTLASS.prune_dense_static_sort(pruned, algorithm="largest_abs_values_greedy")
+        reference_cutlass = SparseSemiStructuredTensorCUTLASS.prune_dense_static_sort(
+            pruned, algorithm="largest_abs_values_greedy"
+        )
         torch.testing.assert_close(pruned, reference_cutlass.to_dense())
 
         packed_cutlass, meta_cutlass = sparse_semi_structured_from_dense_cutlass(pruned)
-        packed_t_cutlass, meta_t_cutlass = sparse_semi_structured_from_dense_cutlass(pruned.t().contiguous())
-        meta_cutlass = meta_cutlass.as_strided(reference_cutlass.meta.shape, reference_cutlass.meta.stride())
-        meta_t_cutlass = meta_t_cutlass.as_strided(reference_cutlass.meta_t.shape, reference_cutlass.meta_t.stride())
+        packed_t_cutlass, meta_t_cutlass = sparse_semi_structured_from_dense_cutlass(
+            pruned.t().contiguous()
+        )
+        meta_cutlass = meta_cutlass.as_strided(
+            reference_cutlass.meta.shape, reference_cutlass.meta.stride()
+        )
+        meta_t_cutlass = meta_t_cutlass.as_strided(
+            reference_cutlass.meta_t.shape, reference_cutlass.meta_t.stride()
+        )
         compressed_swizzled_bitmask = _compute_compressed_swizzled_bitmask(pruned)
-        compressed_swizzled_bitmask = compressed_swizzled_bitmask.as_strided(reference_cutlass.compressed_swizzled_bitmask.shape,
-                                                                             reference_cutlass.compressed_swizzled_bitmask.stride())
-        cutlass = SparseSemiStructuredTensorCUTLASS(dense.shape,
-                                                    packed_cutlass,
-                                                    meta_cutlass,
-                                                    packed_t_cutlass,
-                                                    meta_t_cutlass,
-                                                    compressed_swizzled_bitmask)
+        compressed_swizzled_bitmask = compressed_swizzled_bitmask.as_strided(
+            reference_cutlass.compressed_swizzled_bitmask.shape,
+            reference_cutlass.compressed_swizzled_bitmask.stride(),
+        )
+        cutlass = SparseSemiStructuredTensorCUTLASS(
+            dense.shape,
+            packed_cutlass,
+            meta_cutlass,
+            packed_t_cutlass,
+            meta_t_cutlass,
+            compressed_swizzled_bitmask,
+        )
         torch.testing.assert_close(reference_cutlass.to_dense(), cutlass.to_dense())
 
         # CUSPARSELT
-        reference_cusparselt = SparseSemiStructuredTensorCUSPARSELT.prune_dense_static_sort(pruned,
-                                                                                            algorithm="largest_abs_values_greedy")
+        reference_cusparselt = (
+            SparseSemiStructuredTensorCUSPARSELT.prune_dense_static_sort(
+                pruned, algorithm="largest_abs_values_greedy"
+            )
+        )
         torch.testing.assert_close(pruned, reference_cusparselt.to_dense())
 
         packed_cusparselt = torch._cslt_compress(pruned)
         packed_t_cusparselt = torch._cslt_compress(pruned.t().contiguous())
-        cusparselt = SparseSemiStructuredTensorCUSPARSELT(dense.shape,
-                                                          packed_cusparselt,
-                                                          None,
-                                                          packed_t_cusparselt,
-                                                          None,
-                                                          compressed_swizzled_bitmask)
-        torch.testing.assert_close(reference_cusparselt.to_dense(), cusparselt.to_dense())
-
-
+        cusparselt = SparseSemiStructuredTensorCUSPARSELT(
+            dense.shape,
+            packed_cusparselt,
+            None,
+            packed_t_cusparselt,
+            None,
+            compressed_swizzled_bitmask,
+        )
+        torch.testing.assert_close(
+            reference_cusparselt.to_dense(), cusparselt.to_dense()
+        )
 
     @training_dtypes
     @parametrize_backends
@@ -649,7 +737,9 @@ def test_pruning_algo_largest_abs_values_greedy(self, dtype, backend) -> None:
             dtype=dtype,
         )
         inp = F.pad(inp, (0, 128 - 4, 0, 128 - 4), "constant", 1)
-        sInp = SEMI_STRUCTURED_SUPPORTED_BACKENDS[backend].prune_dense_static_sort(inp, algorithm="largest_abs_values_greedy")
+        sInp = SEMI_STRUCTURED_SUPPORTED_BACKENDS[backend].prune_dense_static_sort(
+            inp, algorithm="largest_abs_values_greedy"
+        )
 
         mask = sInp.to_dense() / inp
         expected = [
@@ -659,9 +749,15 @@ def test_pruning_algo_largest_abs_values_greedy(self, dtype, backend) -> None:
             [1, 0, 0, 1],
         ]
         if mask[:4, :4].int().tolist() != expected:
-            raise AssertionError(f"mask[:4, :4] mismatch: {mask[:4, :4].int().tolist()} != {expected}")
+            raise AssertionError(
+                f"mask[:4, :4] mismatch: {mask[:4, :4].int().tolist()} != {expected}"
+            )
 
     @training_dtypes
+    @unittest.skipIf(
+        not torch.backends.cusparselt.is_available(),
+        "cuSPARSELt not available",
+    )
     def test_gemm(self, dtype) -> None:
         M, N, K = 32, 32, 64
         a = torch.randn([M, K], device="cuda", dtype=dtype)
@@ -677,7 +773,6 @@ def test_gemm(self, dtype) -> None:
         sp24_out = a_sparse @ b
         torch.testing.assert_close(ref_out, sp24_out, **atol_rtol_kw[dtype])
 
-
     @training_dtypes
     @parametrize_backends
     @unittest.skipIf(TEST_WITH_ROCM, "Not supported on ROCm")
@@ -695,15 +790,20 @@ def test_pack_both_ways_meta_correctness(self, dtype, backend) -> None:
         a = a.cuda().to(dtype)
         b = torch.randn([a.shape[1], 128], device="cuda", dtype=dtype)
 
-        a_sparse = SEMI_STRUCTURED_SUPPORTED_BACKENDS[backend].prune_dense_static_sort(a)
+        a_sparse = SEMI_STRUCTURED_SUPPORTED_BACKENDS[backend].prune_dense_static_sort(
+            a
+        )
 
         mask_dense = sparse24_largest_mask_2d(a).to(dtype)
 
         if backend == "cutlass":
             if not isinstance(a_sparse, SparseSemiStructuredTensorCUTLASS):
-                raise AssertionError(f"a_sparse should be SparseSemiStructuredTensorCUTLASS, got {type(a_sparse)}")
-            (packed, meta, packed_t, meta_t, bitmask) = torch._sparse_semi_structured_tile(
-                mask_dense, use_cutlass=True)
+                raise AssertionError(
+                    f"a_sparse should be SparseSemiStructuredTensorCUTLASS, got {type(a_sparse)}"
+                )
+            (packed, meta, packed_t, meta_t, bitmask) = (
+                torch._sparse_semi_structured_tile(mask_dense, use_cutlass=True)
+            )
 
             sparse_mask = SparseSemiStructuredTensorCUTLASS(
                 mask_dense.shape,
@@ -713,7 +813,9 @@ def test_pack_both_ways_meta_correctness(self, dtype, backend) -> None:
                 meta_t=meta_t,
                 compressed_swizzled_bitmask=bitmask,
             )
-            torch.testing.assert_close(a_sparse.meta.view(torch.short), sparse_mask.meta)
+            torch.testing.assert_close(
+                a_sparse.meta.view(torch.short), sparse_mask.meta
+            )
 
         ref_gemm = (mask_dense * a) @ b
         pack_gemm = a_sparse @ b
@@ -731,9 +833,7 @@ def test_pack_both_ways_id(self, dtype) -> None:
         a = torch.randn([N, N], dtype=dtype, device="cuda")
         b = torch.eye(N, dtype=dtype, device="cuda")
 
-        packed, meta, packed_t, meta_t = torch._sparse_semi_structured_tile(a)[
-            :4
-        ]
+        packed, meta, packed_t, meta_t = torch._sparse_semi_structured_tile(a)[:4]
         # Heuristic to ensure we pack the same values
         torch.testing.assert_close(
             packed.to(torch.float64).sum(), packed_t.to(torch.float64).sum()
@@ -746,7 +846,8 @@ def test_pack_both_ways_id(self, dtype) -> None:
         pack_gemm = torch._sparse_semi_structured_linear(b.t(), packed, meta).t()
         max_diff = (ref_gemm - pack_gemm).abs().argmax()
         torch.testing.assert_close(
-            ref_gemm, pack_gemm,
+            ref_gemm,
+            pack_gemm,
             **atol_rtol_kw[dtype],
             msg=f"packed is wrong at pos: ({max_diff // N}, {max_diff % N})",
         )
@@ -755,7 +856,8 @@ def test_pack_both_ways_id(self, dtype) -> None:
         max_diff = (ref_gemm - pack_gemm).abs().argmax()
 
         torch.testing.assert_close(
-            ref_gemm, pack_gemm,
+            ref_gemm,
+            pack_gemm,
             **atol_rtol_kw[dtype],
             msg=f"packed_t is wrong at pos: ({max_diff // N}, {max_diff % N})",
         )
@@ -789,9 +891,13 @@ def test_pack_both_ways_edge_case1(self, dtype) -> None:
             raise AssertionError(f"packed[0, 1] should be 0, got {packed[0, 1].item()}")
         # And first column in A.t
         if packed_t[0, 0].item() != 2:
-            raise AssertionError(f"packed_t[0, 0] should be 2, got {packed_t[0, 0].item()}")
+            raise AssertionError(
+                f"packed_t[0, 0] should be 2, got {packed_t[0, 0].item()}"
+            )
         if packed_t[0, 1].item() != 0:
-            raise AssertionError(f"packed_t[0, 1] should be 0, got {packed_t[0, 1].item()}")
+            raise AssertionError(
+                f"packed_t[0, 1] should be 0, got {packed_t[0, 1].item()}"
+            )
 
     @training_dtypes
     @unittest.skipIf(TEST_WITH_ROCM, "Not supported on ROCm")
@@ -854,7 +960,6 @@ def test_sp24_apply_dense(self, dtype) -> None:
         torch.testing.assert_close(dense, expected)
         torch.testing.assert_close(sparse.to_dense(), expected)
 
-
     @training_dtypes
     @unittest.skipIf(TEST_WITH_ROCM, "Not supported on ROCm")
     @unittest.skipIf(
@@ -867,7 +972,9 @@ def test_sp24_matmuls(self, dtype) -> None:
         b = torch.randn([K, N], device="cuda", dtype=dtype)
         a_m = sparse24_largest_mask_2d(a)
         b_m = sparse24_largest_mask_2d(b)
-        (packed, meta, packed_t, meta_t, bitmask) = torch._sparse_semi_structured_tile(a)
+        (packed, meta, packed_t, meta_t, bitmask) = torch._sparse_semi_structured_tile(
+            a
+        )
         a_s = SparseSemiStructuredTensorCUTLASS(
             a.shape,
             packed=packed,
@@ -876,7 +983,9 @@ def test_sp24_matmuls(self, dtype) -> None:
             meta_t=meta_t,
             compressed_swizzled_bitmask=bitmask,
         )
-        (packed, meta, packed_t, meta_t, bitmask) = torch._sparse_semi_structured_tile(b)
+        (packed, meta, packed_t, meta_t, bitmask) = torch._sparse_semi_structured_tile(
+            b
+        )
         b_s = SparseSemiStructuredTensorCUTLASS(
             b.shape,
             packed=packed,
@@ -891,11 +1000,13 @@ def test_sp24_matmuls(self, dtype) -> None:
         torch.testing.assert_close(
             a @ a_s.t(), a @ (a * a_m).t(), rtol=1e-1, atol=1.5e-1
         )
-        torch.testing.assert_close(
-            a_s.t() @ a, (a * a_m).t() @ a, rtol=1e-1, atol=1e-1
-        )
+        torch.testing.assert_close(a_s.t() @ a, (a * a_m).t() @ a, rtol=1e-1, atol=1e-1)
 
     @unittest.skipIf(TEST_WITH_ROCM, "Not supported on ROCm")
+    @unittest.skipIf(
+        not torch.backends.cusparselt.is_available(),
+        "cuSPARSELt not available",
+    )
     @unittest.skipIf(
         "RelWithAssert" in torch.__config__.show(),
         "failing in debug build, see https://github.com/pytorch/pytorch/pull/165158 for context",
@@ -910,6 +1021,10 @@ def test_sp24_matmuls_mat_vec(self) -> None:
             torch.testing.assert_close(a_s @ b, (a * a_m) @ b, **atol_rtol_kw[a.dtype])
 
     @unittest.skipIf(TEST_WITH_ROCM, "Not supported on ROCm")
+    @unittest.skipIf(
+        not torch.backends.cusparselt.is_available(),
+        "cuSPARSELt not available",
+    )
     @unittest.skipIf(
         "RelWithAssert" in torch.__config__.show(),
         "failing in debug build, see https://github.com/pytorch/pytorch/pull/165158 for context",
@@ -923,34 +1038,53 @@ def test_sp24_matmuls_bmm(self) -> None:
         with pytest.raises(NotImplementedError):
             torch.testing.assert_close(a_s @ b, (a * a_m) @ b, **atol_rtol_kw[a.dtype])
 
+
 class TestSparseSemiStructuredCUTLASS(TestCase):
     """
     This contains CUTLASS specific tests for
          - torch._sparse_semi_structured_linear
     """
+
     def setUp(self):
         SparseSemiStructuredTensor._FORCE_CUTLASS = True
         if "cutlass" not in SEMI_STRUCTURED_SUPPORTED_BACKENDS:
-            self.skipTest('CUTLASS not enabled')
+            self.skipTest("CUTLASS not enabled")
 
     def tearDown(self):
         SparseSemiStructuredTensor._FORCE_CUTLASS = False
         super().tearDown()
 
-    @unittest.skipIf(TEST_WITH_ROCM or IS_WINDOWS, "ROCm and Windows doesn't support CUTLASS")
+    @unittest.skipIf(
+        TEST_WITH_ROCM or IS_WINDOWS, "ROCm and Windows doesn't support CUTLASS"
+    )
     @inference_dtypes
     def test_linear_cutlass(self, device, dtype):
-
-        def run_test(batch_shape, m, n, k, device, dtype, dtype_out, add_bias, activation, rtol, atol):
+        def run_test(
+            batch_shape,
+            m,
+            n,
+            k,
+            device,
+            dtype,
+            dtype_out,
+            add_bias,
+            activation,
+            rtol,
+            atol,
+        ):
             weight = rand_sparse_semi_structured(m, k, dtype, device)
             input = make_tensor((*batch_shape, n, k), dtype=dtype, device=device)
-            bias = make_tensor((m,), dtype=dtype_out, device=device) if add_bias else None
+            bias = (
+                make_tensor((m,), dtype=dtype_out, device=device) if add_bias else None
+            )
 
             dtype_dense = torch.float32
             input_dense = input.to(dtype_dense)
             weight_dense = weight.to(dtype_dense)
             bias_dense = bias.to(dtype_dense) if add_bias else None
-            output0 = torch.nn.functional.linear(input_dense, weight_dense, bias=bias_dense)
+            output0 = torch.nn.functional.linear(
+                input_dense, weight_dense, bias=bias_dense
+            )
             if activation == "relu":
                 relu = torch.nn.ReLU()
                 output0 = relu(output0)
@@ -963,9 +1097,17 @@ def run_test(batch_shape, m, n, k, device, dtype, dtype_out, add_bias, activatio
             weight_sparse = compressed.values()
             meta = compressed.indices()
 
-            output1 = torch._sparse_semi_structured_linear(input, weight_sparse, meta, bias=bias, activation=activation,
-                                                           out_dtype=dtype_out if dtype == torch.int8 else None)
-            torch.testing.assert_close(output1.to(dtype_dense), output0, rtol=rtol, atol=atol)
+            output1 = torch._sparse_semi_structured_linear(
+                input,
+                weight_sparse,
+                meta,
+                bias=bias,
+                activation=activation,
+                out_dtype=dtype_out if dtype == torch.int8 else None,
+            )
+            torch.testing.assert_close(
+                output1.to(dtype_dense), output0, rtol=rtol, atol=atol
+            )
 
         if dtype == torch.float32:
             # Inputs are converted to TF32 internally for sparse GEMM,
@@ -974,32 +1116,51 @@ def run_test(batch_shape, m, n, k, device, dtype, dtype_out, add_bias, activatio
             torch.backends.cuda.matmul.allow_tf32 = True
 
         batch_shapes = [[], [3], [3, 1]]
-        dtype_out = {torch.int8: torch.int32, torch.half: torch.half, torch.bfloat16: torch.bfloat16, torch.float32: torch.float32}
+        dtype_out = {
+            torch.int8: torch.int32,
+            torch.half: torch.half,
+            torch.bfloat16: torch.bfloat16,
+            torch.float32: torch.float32,
+        }
         activations = [None, "relu", "silu"]
         rtol, atol = 1e-3, 1e-3
         if dtype == torch.bfloat16:
             rtol, atol = 5e-3, 5e-3
         elif dtype == torch.float32:
             rtol, atol = 1e-3, 75e-2
-        for batch_shape, m, n, k, add_bias, activation in \
-                itertools.product(batch_shapes, range(3), range(3), range(3), (False, True), activations):
+        for batch_shape, m, n, k, add_bias, activation in itertools.product(
+            batch_shapes, range(3), range(3), range(3), (False, True), activations
+        ):
             if activation == "silu" and dtype == torch.int8:
                 continue  # SiLU not supported for integer inputs
 
-            m = 2 ** m * 32
-            n = 2 ** n * 32
-            k = 2 ** k * 128
-            run_test(batch_shape, m, n, k, device, dtype, dtype_out[dtype], add_bias, activation, rtol, atol)
+            m = 2**m * 32
+            n = 2**n * 32
+            k = 2**k * 128
+            run_test(
+                batch_shape,
+                m,
+                n,
+                k,
+                device,
+                dtype,
+                dtype_out[dtype],
+                add_bias,
+                activation,
+                rtol,
+                atol,
+            )
 
         if dtype == torch.float32:
             torch.backends.cuda.matmul.allow_tf32 = orig
 
-
-    @unittest.skipIf(TEST_WITH_ROCM or IS_WINDOWS, "ROCm and Windows doesn't support CUTLASS")
+    @unittest.skipIf(
+        TEST_WITH_ROCM or IS_WINDOWS, "ROCm and Windows doesn't support CUTLASS"
+    )
     @parametrize("backend", ["cutlass"])
     @inference_dtypes
     def test_sparse_semi_structured_ops_cutlass(self, device, dtype, backend):
-        SparseSemiStructuredTensor._FORCE_CUTLASS = (backend == "cutlass")
+        SparseSemiStructuredTensor._FORCE_CUTLASS = backend == "cutlass"
         if backend == "cutlass" and IS_WINDOWS:
             self.skipTest("CUTLASS not supported on Windows")
 
@@ -1007,7 +1168,9 @@ def run_test(m, n, k, device, dtype, dtype_out, use_input, rtol, atol):
             mat1 = rand_sparse_semi_structured(m, k, dtype, device)
             # mat2 transposed as int8 case supports only row-major/column-major combination
             mat2 = make_tensor((n, k), dtype=dtype, device=device).t()
-            input = make_tensor((m,), dtype=dtype_out, device=device) if use_input else None
+            input = (
+                make_tensor((m,), dtype=dtype_out, device=device) if use_input else None
+            )
 
             if use_input:
                 if dtype.is_floating_point:
@@ -1024,7 +1187,9 @@ def run_test(m, n, k, device, dtype, dtype_out, use_input, rtol, atol):
                 output0 = torch.mm(mat1_dense, mat2_dense)
             else:
                 input_dense = input.to(dtype_dense)[:, None]
-                output0 = torch.addmm(input_dense, mat1_dense, mat2_dense, alpha=alpha, beta=beta)
+                output0 = torch.addmm(
+                    input_dense, mat1_dense, mat2_dense, alpha=alpha, beta=beta
+                )
 
             compressed = to_sparse_semi_structured(mat1)
 
@@ -1032,12 +1197,22 @@ def run_test(m, n, k, device, dtype, dtype_out, use_input, rtol, atol):
             mat1_meta = compressed.indices()
 
             if not use_input:
-                output1 = torch._sparse_semi_structured_mm(mat1_sparse, mat1_meta, mat2, out_dtype=dtype_out)
+                output1 = torch._sparse_semi_structured_mm(
+                    mat1_sparse, mat1_meta, mat2, out_dtype=dtype_out
+                )
             else:
                 output1 = torch._sparse_semi_structured_addmm(
-                    input, mat1_sparse, mat1_meta, mat2, alpha=alpha, beta=beta, out_dtype=dtype_out
+                    input,
+                    mat1_sparse,
+                    mat1_meta,
+                    mat2,
+                    alpha=alpha,
+                    beta=beta,
+                    out_dtype=dtype_out,
                 )
-            torch.testing.assert_close(output1.to(dtype_dense), output0, rtol=rtol, atol=atol)
+            torch.testing.assert_close(
+                output1.to(dtype_dense), output0, rtol=rtol, atol=atol
+            )
 
         if dtype == torch.float32:
             # Inputs are converted to TF32 internally for sparse GEMM,
@@ -1045,28 +1220,32 @@ def run_test(m, n, k, device, dtype, dtype_out, use_input, rtol, atol):
             orig = torch.backends.cuda.matmul.allow_tf32
             torch.backends.cuda.matmul.allow_tf32 = True
 
-        dtype_out = {torch.int8: torch.int32, torch.half: torch.half, torch.bfloat16: torch.bfloat16, torch.float32: torch.float32}
+        dtype_out = {
+            torch.int8: torch.int32,
+            torch.half: torch.half,
+            torch.bfloat16: torch.bfloat16,
+            torch.float32: torch.float32,
+        }
         rtol, atol = 1e-3, 1e-3
         if dtype == torch.bfloat16:
             rtol, atol = 5e-3, 5e-3
         elif dtype == torch.float32:
             rtol, atol = 1e-3, 75e-2
-        for m, n, k, use_input in \
-                itertools.product(range(3), range(3), range(3), (False, True)):
-            m = 2 ** m * 32
-            n = 2 ** n * 32
-            k = 2 ** k * 128
+        for m, n, k, use_input in itertools.product(
+            range(3), range(3), range(3), (False, True)
+        ):
+            m = 2**m * 32
+            n = 2**n * 32
+            k = 2**k * 128
             run_test(m, n, k, device, dtype, dtype_out[dtype], use_input, rtol, atol)
 
         if dtype == torch.float32:
             torch.backends.cuda.matmul.allow_tf32 = orig
 
-
     @unittest.skipIf(not HAS_GPU, "Inductor+gpu needs triton and recent GPU arch")
     @inference_dtypes
     @unittest.skipIf(TEST_WITH_ROCM, "Not supported on ROCm")
     def test_conversions(self, device, dtype):
-
         def run_test(r, c, device, dtype):
             dense_ref = rand_sparse_semi_structured(r, c, dtype, device)
 
@@ -1097,7 +1276,9 @@ def run_test(r, c, device, dtype):
     def test_conversions_all_patterns(self, device, dtype):
         r, c = 32, 128
 
-        dense_inv, dense_val = rand_sparse_semi_structured_all_patterns(r, c, dtype, device)
+        dense_inv, dense_val = rand_sparse_semi_structured_all_patterns(
+            r, c, dtype, device
+        )
 
         compressed = to_sparse_semi_structured(dense_inv)
         dense = compressed.to_dense()
@@ -1107,6 +1288,7 @@ def test_conversions_all_patterns(self, device, dtype):
 
 CUSPARSELT_MIXED_DTYPE_SUPPORT = [torch.float16, torch.bfloat16, torch.int32]
 
+
 def to_float8(x, dtype=torch.float8_e4m3fn):
     finfo = torch.finfo(dtype)
     # Calculate the scale as dtype max divided by absmax
@@ -1119,17 +1301,20 @@ def to_float8(x, dtype=torch.float8_e4m3fn):
     # as both required as inputs to torch._scaled_mm
     return x_scl_sat.to(dtype), scale.float().reciprocal()
 
+
 class TestSparseSemiStructuredCUSPARSELT(TestCase):
     """
     This contains cuSPARSELt specific tests for
         torch._cslt_compress
         torch._cslt_sparse_mm
     """
+
     def setUp(self):
         SparseSemiStructuredTensor._FORCE_CUTLASS = False
         if "cusparselt" not in SEMI_STRUCTURED_SUPPORTED_BACKENDS:
-            self.skipTest('cuSPARSELt not enabled')
+            self.skipTest("cuSPARSELt not enabled")
 
+    @unittest.skipIf(TEST_WITH_ROCM, "Not supported on ROCm")
     @unittest.skipIf(
         not PLATFORM_SUPPORTS_FP8,
         "FP8 is only supported on H100+, SM 8.9 and MI300+ devices",
@@ -1152,6 +1337,7 @@ def test_sparse_fp8fp8_mm(self, dense_input_shape, device):
         ):
             dense_result = torch.mm(A_fp8_sparse, B_fp8)
 
+    @unittest.skipIf(TEST_WITH_ROCM, "Not supported on ROCm")
     @unittest.skipIf(
         not PLATFORM_SUPPORTS_FP8,
         "FP8 is only supported on H100+, SM 8.9 and MI300+ devices",
@@ -1159,19 +1345,26 @@ def test_sparse_fp8fp8_mm(self, dense_input_shape, device):
     @xfailIfSM89PreCUDA13
     def test_sparse_semi_structured_scaled_mm_fp8(self, device) -> None:
         (k, l, m) = (32, 64, 32)
-        x = rand_sparse_semi_structured_mask(k, l, dtype=torch.float8_e4m3fn, device=device)
-        y = torch.full((m, l), .25, device=device, dtype=torch.float8_e4m3fn).t()
+        x = rand_sparse_semi_structured_mask(
+            k, l, dtype=torch.float8_e4m3fn, device=device
+        )
+        y = torch.full((m, l), 0.25, device=device, dtype=torch.float8_e4m3fn).t()
         scale_a = torch.tensor(1.0, device=device)
         scale_b = torch.tensor(1.0, device=device)
-        out_fp8 = torch._scaled_mm(x, y, scale_a=scale_a, scale_b=scale_b, out_dtype=torch.float8_e4m3fn)
+        out_fp8 = torch._scaled_mm(
+            x, y, scale_a=scale_a, scale_b=scale_b, out_dtype=torch.float8_e4m3fn
+        )
 
         x_sparse = to_sparse_semi_structured(x)
-        out_fp8_sparse = torch._scaled_mm(x_sparse, y, scale_a=scale_a, scale_b=scale_b, out_dtype=torch.float8_e4m3fn)
+        out_fp8_sparse = torch._scaled_mm(
+            x_sparse, y, scale_a=scale_a, scale_b=scale_b, out_dtype=torch.float8_e4m3fn
+        )
         # this fails on ROCm currently because hipblaslt doesn't have amax op
         out_fp32 = out_fp8.to(torch.float32)
         out_fp32_sparse = out_fp8_sparse.to(torch.float32)
         torch.testing.assert_close(out_fp32, out_fp32_sparse, rtol=1e-1, atol=1e-1)
 
+    @unittest.skipIf(TEST_WITH_ROCM, "Not supported on ROCm")
     @unittest.skipIf(
         not PLATFORM_SUPPORTS_FP8,
         "FP8 is only supported on H100+, SM 8.9 and MI300+ devices",
@@ -1206,7 +1399,9 @@ def test_cslt_sparse_mm_mixed_dtype(self, dense_input_shape, out_dtype, device):
 
         B = torch.rand(dense_input_shape, device=device).to(torch.int8)
 
-        dense_result = torch.mm(A.cpu().to(torch.int64), B.t().cpu().to(torch.int64)).to(device, dtype=out_dtype)
+        dense_result = torch.mm(
+            A.cpu().to(torch.int64), B.t().cpu().to(torch.int64)
+        ).to(device, dtype=out_dtype)
         sparse_result = torch._cslt_sparse_mm(A_compressed, B.t(), out_dtype=out_dtype)
         torch.testing.assert_close(dense_result, sparse_result, rtol=1e-3, atol=1e-3)
 
@@ -1216,7 +1411,7 @@ def test_cslt_sparse_mm_mixed_dtype(self, dense_input_shape, out_dtype, device):
     def test_cslt_sparse_mm_alpha(self, dtype, device):
         A = torch.Tensor([0, 0, 1, 1]).tile((128, 64)).to(dtype).cuda()
         B = torch.ones((256, 128), device=device).to(dtype)
-        alpha = torch.Tensor([2**(-i) for i in range(128)]).cuda()
+        alpha = torch.Tensor([2 ** (-i) for i in range(128)]).cuda()
         bias = torch.ones(128, device=device).to(dtype)
 
         A_compressed = torch._cslt_compress(A)
@@ -1233,40 +1428,52 @@ def test_cslt_sparse_mm_alpha(self, dtype, device):
     def test_cslt_sparse_mm_alpha_compile_autotune(self, device, out_dtype):
         A = torch.Tensor([0, 0, 1, 1]).tile((128, 64)).to(torch.int8).to(device)
         B = torch.ones((128, 256), device=device, dtype=torch.int8).t()
-        alpha = torch.Tensor([2**(-i) for i in range(128)]).cuda()
+        alpha = torch.Tensor([2 ** (-i) for i in range(128)]).cuda()
 
         A_compressed = torch._cslt_compress(A)
 
         cslt_sparse_mm_c = torch.compile(torch._cslt_sparse_mm, mode="max-autotune")
-        sparse_result = cslt_sparse_mm_c(A_compressed, B, alpha=alpha, out_dtype=out_dtype)
+        sparse_result = cslt_sparse_mm_c(
+            A_compressed, B, alpha=alpha, out_dtype=out_dtype
+        )
 
         # disable this otherwise inductor will attempt to reorder strides and pass a contiguous B
         @torch.compiler.disable
         def get_dense_result():
             alpha_scaled = torch.stack([alpha] * 128).t().cpu().float()
-            dense_result = alpha_scaled * torch.mm(A.to(torch.int64).cpu(), B.to(torch.int64).cpu())
+            dense_result = alpha_scaled * torch.mm(
+                A.to(torch.int64).cpu(), B.to(torch.int64).cpu()
+            )
             dense_result = dense_result.to(out_dtype)
             return dense_result
 
-        torch.testing.assert_close(sparse_result.cpu(), get_dense_result(), rtol=1e-3, atol=1e-3)
+        torch.testing.assert_close(
+            sparse_result.cpu(), get_dense_result(), rtol=1e-3, atol=1e-3
+        )
 
     @parametrize("out_dtype", [torch.float16, torch.bfloat16, torch.int32])
     @unittest.skipIf(TEST_WITH_ROCM, "Not supported on ROCm")
     def test_cslt_sparse_mm_alpha_mixed_dtype(self, out_dtype, device):
         A = torch.Tensor([0, 0, 10, 10]).tile((128, 64)).to(torch.int8).cuda()
         B = torch.ones((128, 256), device=device).to(torch.int8).t()
-        alpha = torch.Tensor([2**(-i) if out_dtype is not torch.int32 else 1
-                              for i in range(128)]).cuda()
+        alpha = torch.Tensor(
+            [2 ** (-i) if out_dtype is not torch.int32 else 1 for i in range(128)]
+        ).cuda()
 
         A_compressed = torch._cslt_compress(A)
-        sparse_result = torch._cslt_sparse_mm(A_compressed, B, alpha=alpha, out_dtype=out_dtype).cpu()
+        sparse_result = torch._cslt_sparse_mm(
+            A_compressed, B, alpha=alpha, out_dtype=out_dtype
+        ).cpu()
 
         alpha_scaled = torch.stack([alpha] * 128).t()
-        dense_result = alpha_scaled.cpu() * torch.mm(A.to(torch.int64).cpu(), B.to(torch.int64).cpu())
+        dense_result = alpha_scaled.cpu() * torch.mm(
+            A.to(torch.int64).cpu(), B.to(torch.int64).cpu()
+        )
         dense_result = dense_result.to(out_dtype)
 
         torch.testing.assert_close(sparse_result, dense_result, rtol=1e-3, atol=1e-3)
 
+    @unittest.skipIf(TEST_WITH_ROCM, "Not supported on ROCm")
     @inference_dtypes
     def test_cslt_sparse_mm_search(self, device, dtype):
         A = rand_sparse_semi_structured_mask(256, 128, dtype=dtype)
@@ -1280,6 +1487,7 @@ def test_cslt_sparse_mm_search(self, device, dtype):
         dense_result = dense_result.to(dtype)
         torch.testing.assert_close(sparse_result, dense_result, rtol=1e-3, atol=1e-3)
 
+    @unittest.skipIf(TEST_WITH_ROCM, "Not supported on ROCm")
     @inference_dtypes
     def test_csrc_cslt_sparse_mm_search(self, device, dtype):
         A = rand_sparse_semi_structured_mask(256, 128, dtype=dtype)
@@ -1287,30 +1495,184 @@ def test_csrc_cslt_sparse_mm_search(self, device, dtype):
         B = torch.ones((128, 128), device=device).to(dtype)
 
         A_compressed = torch._cslt_compress(A)
-        alg_id, split_k, split_k_mode, _ = torch._C._cusparselt.mm_search(A_compressed, B.t(), None, None, None, False)
-        sparse_result = torch._cslt_sparse_mm(A_compressed, B.t(),
-                                              alg_id=alg_id,
-                                              split_k=split_k,
-                                              split_k_mode=split_k_mode)
+        alg_id, split_k, split_k_mode, _ = torch._C._cusparselt.mm_search(
+            A_compressed, B.t(), None, None, None, False
+        )
+        sparse_result = torch._cslt_sparse_mm(
+            A_compressed,
+            B.t(),
+            alg_id=alg_id,
+            split_k=split_k,
+            split_k_mode=split_k_mode,
+        )
         dense_result = torch.mm(A.to(torch.float32), B.to(torch.float32))
         dense_result = dense_result.to(dtype)
         torch.testing.assert_close(sparse_result, dense_result, rtol=1e-3, atol=1e-3)
 
+    @unittest.skipIf(TEST_WITH_ROCM, "Not supported on ROCm")
     def test_cusparselt_backend(self):
         if not torch.backends.cusparselt.is_available():
             raise AssertionError("cusparselt backend should be available")
 
         # PyTorch CUDA 12.4+ using cuSPARSELt v0.6.2+
         if torch.backends.cusparselt.version() < 602:
-            raise AssertionError(f"cusparselt version should be >= 602, got {torch.backends.cusparselt.version()}")
+            raise AssertionError(
+                f"cusparselt version should be >= 602, got {torch.backends.cusparselt.version()}"
+            )
+
+    @inference_dtypes
+    def test_cslt_sparse_tensor_clone(self, dtype):
+        A = rand_sparse_semi_structured_mask(256, 128, dtype=dtype).cuda()
+        A_sparse = to_sparse_semi_structured(A)
+        A_clone = A_sparse.clone()
+        self.assertNotEqual(A_sparse.packed.data_ptr, A_clone.packed.data_ptr)
+        self.assertEqual(A_sparse.packed, A_clone.packed)
+
+    @dtypes(torch.float16, torch.bfloat16)
+    def test_semi_sparse_to_device_cpu(self, dtype):
+        """Test .to('cpu') converts sparse semi-structured tensor to dense on CPU."""
+        A = rand_sparse_semi_structured_mask(128, 128, dtype=dtype).cuda()
+        A_sparse = to_sparse_semi_structured(A)
+        A_dense = A_sparse.to_dense()
+
+        result = A_sparse.to("cpu")
+        self.assertEqual(result.device.type, "cpu")
+        self.assertFalse(isinstance(result, SparseSemiStructuredTensor))
+        torch.testing.assert_close(result, A_dense.cpu(), rtol=1e-3, atol=1e-3)
+
+    @dtypes(torch.float16, torch.bfloat16)
+    def test_semi_sparse_to_device_cuda(self, dtype):
+        """Test .to('cuda') / .to(device) raises NotImplementedError."""
+        A = rand_sparse_semi_structured_mask(128, 128, dtype=dtype).cuda()
+        A_sparse = to_sparse_semi_structured(A)
+
+        with self.assertRaises(NotImplementedError):
+            A_sparse.to(A_sparse.device, copy=True)
+
+    @dtypes(torch.float16, torch.bfloat16)
+    def test_semi_sparse_to_dtype(self, dtype):
+        """Test .to(dtype) raises NotImplementedError (only to('cpu') is supported)."""
+        A = rand_sparse_semi_structured_mask(128, 128, dtype=dtype).cuda()
+        A_sparse = to_sparse_semi_structured(A)
+
+        with self.assertRaises(NotImplementedError):
+            A_sparse.to(torch.float32)
+
+    @dtypes(torch.float16, torch.bfloat16)
+    def test_semi_sparse_to_device_and_dtype(self, dtype):
+        """Test .to(device, dtype) converts both device and dtype."""
+        A = rand_sparse_semi_structured_mask(128, 128, dtype=dtype).cuda()
+        A_sparse = to_sparse_semi_structured(A)
+        A_dense = A_sparse.to_dense()
+
+        result = A_sparse.to("cpu", torch.float32, copy=True)
+        self.assertEqual(result.device.type, "cpu")
+        self.assertEqual(result.dtype, torch.float32)
+        self.assertFalse(isinstance(result, SparseSemiStructuredTensor))
+        torch.testing.assert_close(
+            result, A_dense.cpu().to(torch.float32), rtol=1e-3, atol=1e-3
+        )
+
+    @dtypes(torch.float16, torch.bfloat16)
+    def test_semi_sparse_to_same_dtype_noop(self, dtype):
+        """Test .to(same_dtype) raises NotImplementedError (only to('cpu') is supported)."""
+        A = rand_sparse_semi_structured_mask(128, 128, dtype=dtype).cuda()
+        A_sparse = to_sparse_semi_structured(A)
+
+        with self.assertRaises(NotImplementedError):
+            A_sparse.to(dtype, copy=True)
+
+    @dtypes(torch.float16, torch.bfloat16)
+    def test_semi_sparse_to_kwargs(self, dtype):
+        """Test .to() with keyword arguments (device=, dtype=)."""
+        A = rand_sparse_semi_structured_mask(128, 128, dtype=dtype).cuda()
+        A_sparse = to_sparse_semi_structured(A)
+        A_dense = A_sparse.to_dense()
+
+        result = A_sparse.to(device="cpu", dtype=torch.float32)
+        self.assertEqual(result.device.type, "cpu")
+        self.assertEqual(result.dtype, torch.float32)
+        self.assertFalse(isinstance(result, SparseSemiStructuredTensor))
+        torch.testing.assert_close(
+            result, A_dense.cpu().to(torch.float32), rtol=1e-3, atol=1e-3
+        )
+
+    @unittest.skipIf(TEST_WITH_ROCM, "Not supported on ROCm")
+    def test_cslt_sparse_tensor_with_alg_id(self, device):
+        A = rand_sparse_semi_structured_mask(256, 128, dtype=torch.float16).cuda()
+        A_compressed = torch._cslt_compress(A)
+        B = torch.ones((128, 128), device=device).to(torch.float16)
+        alg_id = torch._cslt_sparse_mm_search(A_compressed, B.t())
+        A_sparse = to_sparse_semi_structured(A, alg_id=alg_id)
+        self.assertEqual(A_sparse.alg_id_cusparselt, alg_id)
+        dense_result = torch.mm(A, B).to(torch.float16)
+        sparse_result = torch.mm(A_sparse, B).to(torch.float16)
+        torch.testing.assert_close(sparse_result, dense_result, rtol=1e-3, atol=1e-3)
+
+
+    @unittest.skipIf(
+        not PLATFORM_SUPPORTS_FP8_SPARSE,
+        "FP8 sparse requires cuSPARSELt v0.6.2+ on SM 8.9+ or MI350+ (gfx950) on ROCm",
+    )
+    def test_cslt_compress_fp8(self, device):
+        A = rand_sparse_semi_structured_mask(256, 128, dtype=torch.float16)
+        A_fp8, _ = to_float8(A)
+        compressed = torch._cslt_compress(A_fp8)
+        self.assertEqual(compressed.dtype, torch.float8_e4m3fn)
+
+    @unittest.skipIf(
+        not PLATFORM_SUPPORTS_FP8_SPARSE,
+        "FP8 sparse requires cuSPARSELt v0.6.2+ on SM 8.9+ or MI350+ (gfx950) on ROCm",
+    )
+    @parametrize("dense_input_shape", [(256, 128)])
+    def test_cslt_sparse_mm_fp8_to_fp32(self, dense_input_shape, device):
+        A = rand_sparse_semi_structured_mask(256, 128, dtype=torch.float16)
+        B = torch.rand(dense_input_shape, device=device).to(torch.float16).t()
+
+        A_fp8, _ = to_float8(A)
+        B_fp8, _ = to_float8(B)
+
+        compressed = torch._cslt_compress(A_fp8)
+        sparse_result = torch._cslt_sparse_mm(compressed, B_fp8, out_dtype=torch.float32)
+
+        self.assertEqual(sparse_result.dtype, torch.float32)
+
+        dense_result = torch.mm(A_fp8.to(torch.float32), B_fp8.to(torch.float32))
+        torch.testing.assert_close(sparse_result, dense_result, rtol=1e-1, atol=1e-1)
+
+    @unittest.skipIf(
+        not PLATFORM_SUPPORTS_FP8_SPARSE,
+        "FP8 sparse requires cuSPARSELt v0.6.2+ on SM 8.9+ or MI350+ (gfx950) on ROCm",
+    )
+    @unittest.skipIf(not TEST_WITH_ROCM, "ROCm-specific out_dtype restriction")
+    @parametrize("out_dtype", [torch.float16, torch.bfloat16, torch.float8_e4m3fn])
+    def test_cslt_sparse_mm_fp8_unsupported_out_dtype_rocm(self, out_dtype, device):
+        A = rand_sparse_semi_structured_mask(256, 128, dtype=torch.float16)
+        B = torch.rand(128, 128, device=device).to(torch.float16).t()
+
+        A_fp8, _ = to_float8(A)
+        B_fp8, _ = to_float8(B)
+
+        compressed = torch._cslt_compress(A_fp8)
+        with self.assertRaisesRegex(
+            RuntimeError,
+            "Unsupported out_dtype passed, must be float32 for fp8 inputs on ROCm",
+        ):
+            torch._cslt_sparse_mm(compressed, B_fp8, out_dtype=out_dtype)
 
 if len(SEMI_STRUCTURED_SUPPORTED_BACKENDS) > 0:
     instantiate_device_type_tests(TestSparseSemiStructured, globals(), only_for="cuda")
 if "cutlass" in SEMI_STRUCTURED_SUPPORTED_BACKENDS:
-    instantiate_device_type_tests(TestSparseSemiStructuredCUTLASS, globals(), only_for="cuda")
-    instantiate_device_type_tests(TestSparseSemiStructuredTraining, globals(), only_for="cuda")
+    instantiate_device_type_tests(
+        TestSparseSemiStructuredCUTLASS, globals(), only_for="cuda"
+    )
+    instantiate_device_type_tests(
+        TestSparseSemiStructuredTraining, globals(), only_for="cuda"
+    )
 if "cusparselt" in SEMI_STRUCTURED_SUPPORTED_BACKENDS:
-    instantiate_device_type_tests(TestSparseSemiStructuredCUSPARSELT, globals(), only_for="cuda")
+    instantiate_device_type_tests(
+        TestSparseSemiStructuredCUSPARSELT, globals(), only_for="cuda"
+    )
 
 if __name__ == "__main__":
     run_tests()
diff --git a/test/test_spectral_ops.py b/test/test_spectral_ops.py
index 027f56d28f4fb..4a7dac9f0bdd8 100644
--- a/test/test_spectral_ops.py
+++ b/test/test_spectral_ops.py
@@ -20,7 +20,6 @@
     spectral_funcs, SpectralFuncType)
 from torch._prims_common import corresponding_complex_dtype
 
-from typing import Optional
 from packaging import version
 
 
@@ -599,7 +598,7 @@ def test_fft2_numpy(self, device, dtype):
                 else:
                     numpy_fn = getattr(np.fft, fname)
 
-                def fn(t: torch.Tensor, s: Optional[list[int]], dim: list[int] = (-2, -1), norm: Optional[str] = None):
+                def fn(t: torch.Tensor, s: list[int] | None, dim: list[int] = (-2, -1), norm: str | None = None):
                     return torch_fn(t, s, dim, norm)
 
                 torch_fns = (torch_fn, torch.jit.script(fn))
diff --git a/test/test_stateless_rng.py b/test/test_stateless_rng.py
new file mode 100644
index 0000000000000..08631a9eaa974
--- /dev/null
+++ b/test/test_stateless_rng.py
@@ -0,0 +1,567 @@
+# Owner(s): ["module: random"]
+
+import torch
+import torch.func._random as random
+from torch.testing._internal.common_device_type import (
+    dtypes,
+    instantiate_device_type_tests,
+)
+from torch.testing._internal.common_dtype import floating_types_and
+from torch.testing._internal.common_utils import parametrize, run_tests, TestCase
+
+
+all_floating_dtypes = floating_types_and(torch.half, torch.bfloat16)
+
+
+class TestStatelessRNGKey(TestCase):
+    def test_basic_shape_and_dtype(self, device):
+        key = random.key(42, device=device)
+        self.assertEqual(key.shape, (2,))
+        self.assertEqual(key.dtype, torch.uint64)
+        self.assertEqual(key.device, torch.device(device))
+
+    def test_different_seeds(self, device):
+        key1 = random.key(42, device=device)
+        key2 = random.key(43, device=device)
+        self.assertNotEqual(key1, key2)
+
+    def test_determinism(self, device):
+        key1 = random.key(42, device=device)
+        key2 = random.key(42, device=device)
+        self.assertEqual(key1, key2)
+
+    def test_error_unsupported_impl(self, device):
+        with self.assertRaisesRegex(
+            NotImplementedError, "does not support PRNG impl 'unsupported'"
+        ):
+            random.key(42, impl="unsupported", device=device)
+
+
+class TestStatelessRNGKeySplit(TestCase):
+    def test_basic_shape_and_dtype(self, device):
+        key = random.key(42, device=device)
+        splits = random.split(key, 4)
+        self.assertEqual(splits.shape, (4, 2))
+        self.assertEqual(splits.dtype, torch.uint64)
+        self.assertEqual(splits.device, key.device)
+
+    def test_single_split(self, device):
+        key = random.key(42, device=device)
+        splits = random.split(key, 1)
+        self.assertEqual(splits.shape, (1, 2))
+
+    def test_large_num_splits(self, device):
+        key = random.key(42, device=device)
+        splits = random.split(key, 10000)
+        self.assertEqual(splits.shape, (10000, 2))
+
+    def test_determinism(self, device):
+        key = random.key(42, device=device)
+        splits1 = random.split(key, 8)
+        splits2 = random.split(key, 8)
+        self.assertEqual(splits1, splits2)
+
+    def test_all_keys_unique(self, device):
+        key = random.key(42, device=device)
+        splits = random.split(key, 100)
+        unique_keys = torch.unique(splits, dim=0)
+        self.assertEqual(unique_keys.shape[0], 100)
+
+    def test_different_seeds_produce_different_outputs(self, device):
+        key1 = random.key(42, device=device)
+        key2 = random.key(43, device=device)
+        splits1 = random.split(key1, 4)
+        splits2 = random.split(key2, 4)
+        self.assertNotEqual(splits1, splits2)
+
+    def test_different_offsets_produce_different_outputs(self, device):
+        key1 = random.key(42, device=device)
+        key2 = random.fold_in(key1, 1)
+        splits1 = random.split(key1, 4)
+        splits2 = random.split(key2, 4)
+        self.assertNotEqual(splits1, splits2)
+
+    def test_offset_zero_vs_one_produce_different_splits(self, device):
+        key1 = random.key(42, device=device)
+        key2 = torch.tensor([42, 1], dtype=torch.uint64, device=device)
+        splits1 = random.split(key1, 4)
+        splits2 = random.split(key2, 4)
+        self.assertNotEqual(splits1, splits2)
+
+    def test_batched(self, device):
+        key = random.key(42, device=device)
+        keys = random.split(key, 4)  # (4, 2)
+        num_splits = 3
+        batched = random.split(keys, num_splits)  # (3, 4, 2)
+        self.assertEqual(batched.shape, (num_splits, 4, 2))
+        for k in range(4):
+            individual = random.split(keys[k], num_splits)
+            for s in range(num_splits):
+                self.assertEqual(batched[s][k], individual[s])
+
+    def test_multi_batch(self, device):
+        key = random.key(42, device=device)
+        keys = random.split(key, 12).reshape(3, 4, 2)
+        num_splits = 5
+        batched = random.split(keys, num_splits)  # (5, 3, 4, 2)
+        self.assertEqual(batched.shape, (num_splits, 3, 4, 2))
+        for i in range(3):
+            for j in range(4):
+                individual = random.split(keys[i][j], num_splits)
+                for s in range(num_splits):
+                    self.assertEqual(batched[s][i][j], individual[s])
+
+    def test_error_wrong_shape(self, device):
+        key = torch.tensor([42, 0, 1], dtype=torch.uint64, device=device)
+        with self.assertRaisesRegex(
+            RuntimeError, r"key must have shape \(\*batch, 2\)"
+        ):
+            random.split(key, 4)
+
+    def test_error_wrong_dtype(self, device):
+        key = torch.tensor([42, 0], dtype=torch.float32, device=device)
+        with self.assertRaisesRegex(RuntimeError, "key must have dtype uint64"):
+            random.split(key, 4)
+
+    def test_error_wrong_device(self, device):
+        key = random.key(42)  # CPU key
+        with self.assertRaisesRegex(
+            NotImplementedError,
+            "Could not run .* with arguments from the 'CPU' backend",
+        ):
+            random.split(key, 4)
+
+    def test_error_invalid_num_splits(self, device):
+        key = random.key(42, device=device)
+        with self.assertRaisesRegex(RuntimeError, "num_splits must be positive"):
+            random.split(key, 0)
+        with self.assertRaisesRegex(RuntimeError, "num_splits must be positive"):
+            random.split(key, -1)
+
+    def test_error_batched_last_dim_not_2(self, device):
+        key = torch.tensor([[42, 0, 1], [43, 0, 1]], dtype=torch.uint64, device=device)
+        with self.assertRaisesRegex(
+            RuntimeError, r"key must have shape \(\*batch, 2\)"
+        ):
+            random.split(key, 4)
+
+    def test_offset_overflow(self, device):
+        near_max = (1 << 64) - 1
+        key = torch.tensor([42, near_max], dtype=torch.uint64, device=device)
+        splits = random.split(key, 3)
+        # split_idx=1 wraps offset to 0, split_idx=2 wraps to 1
+        key0 = torch.tensor([42, 0], dtype=torch.uint64, device=device)
+        self.assertEqual(splits[1], random.fold_in(key0, 0))
+        self.assertEqual(splits[2], random.fold_in(key0, 1))
+
+
+class TestStatelessRNGKeyFoldIn(TestCase):
+    def test_basic_shape_and_dtype(self, device):
+        key = random.key(42, device=device)
+        result = random.fold_in(key, 7)
+        self.assertEqual(result.shape, (2,))
+        self.assertEqual(result.dtype, torch.uint64)
+        self.assertEqual(result.device, key.device)
+
+    def test_determinism(self, device):
+        key = random.key(42, device=device)
+        result1 = random.fold_in(key, 7)
+        result2 = random.fold_in(key, 7)
+        self.assertEqual(result1, result2)
+
+    def test_fold_in_produces_new_key_for_zero_data(self, device):
+        key = random.key(42, device=device)
+        folded = random.fold_in(key, 0)
+        self.assertNotEqual(folded, key)
+
+    def test_different_data_produces_different_outputs(self, device):
+        key = random.key(42, device=device)
+        result1 = random.fold_in(key, 0)
+        result2 = random.fold_in(key, 1)
+        self.assertNotEqual(result1, result2)
+
+    def test_consistency_with_split(self, device):
+        key = random.key(42, device=device)
+        splits = random.split(key, 10)
+        for i in range(10):
+            folded = random.fold_in(key, i)
+            self.assertEqual(folded, splits[i])
+
+    def test_batched(self, device):
+        key = random.key(42, device=device)
+        keys = random.split(key, 4)  # (4, 2)
+        data = 7
+        batched = random.fold_in(keys, data)  # (4, 2)
+        self.assertEqual(batched.shape, (4, 2))
+        for k in range(4):
+            individual = random.fold_in(keys[k], data)
+            self.assertEqual(batched[k], individual)
+
+    def test_multi_batch(self, device):
+        key = random.key(42, device=device)
+        keys = random.split(key, 12).reshape(3, 4, 2)
+        data = 7
+        batched = random.fold_in(keys, data)  # (3, 4, 2)
+        self.assertEqual(batched.shape, (3, 4, 2))
+        for i in range(3):
+            for j in range(4):
+                individual = random.fold_in(keys[i][j], data)
+                self.assertEqual(batched[i][j], individual)
+
+    def test_error_wrong_shape(self, device):
+        key = torch.tensor([42, 0, 1], dtype=torch.uint64, device=device)
+        with self.assertRaisesRegex(
+            RuntimeError, r"key must have shape \(\*batch, 2\)"
+        ):
+            random.fold_in(key, 0)
+
+    def test_error_wrong_dtype(self, device):
+        key = torch.tensor([42, 0], dtype=torch.float32, device=device)
+        with self.assertRaisesRegex(RuntimeError, "key must have dtype uint64"):
+            random.fold_in(key, 0)
+
+    def test_error_wrong_device(self, device):
+        key = random.key(42)  # CPU key
+        with self.assertRaisesRegex(
+            NotImplementedError,
+            "Could not run .* with arguments from the 'CPU' backend",
+        ):
+            random.fold_in(key, 0)
+
+    def test_error_batched_last_dim_not_2(self, device):
+        key = torch.tensor([[42, 0, 1], [43, 0, 1]], dtype=torch.uint64, device=device)
+        with self.assertRaisesRegex(
+            RuntimeError, r"key must have shape \(\*batch, 2\)"
+        ):
+            random.fold_in(key, 0)
+
+    def test_offset_overflow(self, device):
+        near_max = (1 << 64) - 1
+        key = torch.tensor([42, near_max], dtype=torch.uint64, device=device)
+        # fold_in(data=1) wraps offset to 0, so it should match fold_in on
+        # a key with offset=0 and data=0.
+        result = random.fold_in(key, 1)
+        key0 = torch.tensor([42, 0], dtype=torch.uint64, device=device)
+        self.assertEqual(result, random.fold_in(key0, 0))
+
+
+class TestStatelessRNGDistribution(TestCase):
+    def _gen(self, gen_fn_name, *args, **kwargs):
+        return getattr(random, gen_fn_name)(*args, **kwargs)
+
+    @parametrize("gen_fn_name", ["normal", "uniform"])
+    @dtypes(*all_floating_dtypes)
+    def test_basic_shape(self, device, dtype, gen_fn_name):
+        key = random.key(42, device=device)
+        result = self._gen(gen_fn_name, key, (100,), dtype=dtype)
+        self.assertEqual(result.shape, (100,))
+        self.assertEqual(result.dtype, dtype)
+
+    @parametrize("gen_fn_name", ["normal", "uniform"])
+    @dtypes(*all_floating_dtypes)
+    def test_determinism(self, device, dtype, gen_fn_name):
+        key = random.key(42, device=device)
+        a = self._gen(gen_fn_name, key, (1000,), dtype=dtype)
+        b = self._gen(gen_fn_name, key, (1000,), dtype=dtype)
+        self.assertEqual(a, b)
+
+    @parametrize("gen_fn_name", ["normal", "uniform"])
+    @dtypes(*all_floating_dtypes)
+    def test_different_keys(self, device, dtype, gen_fn_name):
+        key1 = random.key(42, device=device)
+        key2 = random.key(43, device=device)
+        a = self._gen(gen_fn_name, key1, (1000,), dtype=dtype)
+        b = self._gen(gen_fn_name, key2, (1000,), dtype=dtype)
+        self.assertNotEqual(a, b)
+
+    @parametrize("gen_fn_name", ["normal", "uniform"])
+    @dtypes(*all_floating_dtypes)
+    def test_batched_keys(self, device, dtype, gen_fn_name):
+        key = random.key(42, device=device)
+        keys = random.split(key, 4).unsqueeze(-2)  # (4, 1, 2)
+        result = self._gen(gen_fn_name, keys, (4, 100), dtype=dtype)
+        for i in range(4):
+            individual = self._gen(gen_fn_name, keys[i], (100,), dtype=dtype)
+            self.assertEqual(result[i], individual)
+
+    @parametrize("gen_fn_name", ["normal", "uniform"])
+    @dtypes(*all_floating_dtypes)
+    def test_batched_keys_large(self, device, dtype, gen_fn_name):
+        # Large event_numel to exercise the multi-key tiled kernel path.
+        key = random.key(42, device=device)
+        keys = random.split(key, 4).unsqueeze(-2)  # (4, 1, 2)
+        result = self._gen(gen_fn_name, keys, (4, 10000), dtype=dtype)
+        for i in range(4):
+            individual = self._gen(gen_fn_name, keys[i], (10000,), dtype=dtype)
+            self.assertEqual(result[i], individual)
+
+    @parametrize("gen_fn_name", ["normal", "uniform"])
+    @dtypes(*all_floating_dtypes)
+    def test_multi_batch(self, device, dtype, gen_fn_name):
+        key = random.key(42, device=device)
+        keys = random.split(key, 6).view(2, 3, 1, 2)
+        result = self._gen(gen_fn_name, keys, (2, 3, 50), dtype=dtype)
+        for i in range(2):
+            for j in range(3):
+                individual = self._gen(gen_fn_name, keys[i][j], (50,), dtype=dtype)
+                self.assertEqual(result[i][j], individual)
+
+    @parametrize("gen_fn_name", ["normal", "uniform"])
+    @dtypes(*all_floating_dtypes)
+    def test_key_broadcasting_semantics(self, device, dtype, gen_fn_name):
+        key = random.key(42, device=device)
+
+        # Broadcast key dim: size-1 dims replicate, other dims index keys.
+        keys = random.split(key, 3).unsqueeze(0).unsqueeze(-2)  # (1, 3, 1, 2)
+        result = self._gen(gen_fn_name, keys, (4, 3, 100), dtype=dtype)
+        for i in range(1, 4):
+            self.assertEqual(result[0], result[i])
+        for j in range(1, 3):
+            self.assertNotEqual(result[0][0], result[0][j])
+
+        # All-broadcast key matches unbatched.
+        batched = self._gen(gen_fn_name, key.view(1, 1, 2), (4, 100), dtype=dtype)
+        unbatched = self._gen(gen_fn_name, key, (400,), dtype=dtype)
+        self.assertEqual(batched.flatten(), unbatched)
+
+        # Multiple trailing size-1 dims to broadcast over.
+        keys = random.split(key, 4).view(4, 1, 1, 2)
+        result = self._gen(gen_fn_name, keys, (4, 10, 100), dtype=dtype)
+        for i in range(4):
+            individual = self._gen(gen_fn_name, keys[i], (10, 100), dtype=dtype)
+            self.assertEqual(result[i], individual)
+        keys_flat = random.split(key, 4).unsqueeze(-2)  # (4, 1, 2)
+        flat = self._gen(gen_fn_name, keys_flat, (4, 1000), dtype=dtype)
+        self.assertEqual(result.view(4, 1000), flat)
+
+        # No generation dims: every element gets its own key.
+        keys = random.split(key, 12).view(4, 3, 2)
+        result = self._gen(gen_fn_name, keys, (4, 3), dtype=dtype)
+        for i in range(4):
+            for j in range(3):
+                individual = self._gen(gen_fn_name, keys[i][j], (1,), dtype=dtype)
+                self.assertEqual(result[i][j], individual.squeeze())
+
+    @parametrize("gen_fn_name", ["normal", "uniform"])
+    def test_error_wrong_key_dtype(self, device, gen_fn_name):
+        key = torch.tensor([42, 0], dtype=torch.float32, device=device)
+        with self.assertRaisesRegex(RuntimeError, "key must have dtype uint64"):
+            self._gen(gen_fn_name, key, (100,))
+
+    @parametrize("gen_fn_name", ["normal", "uniform"])
+    def test_error_key_shape(self, device, gen_fn_name):
+        key = random.key(42, device=device)
+        # Last dim must be 2.
+        bad_key = torch.tensor([42, 0, 1], dtype=torch.uint64, device=device)
+        with self.assertRaisesRegex(
+            RuntimeError, r"key must have shape \(2,\) or \(\*batch, 2\)"
+        ):
+            self._gen(gen_fn_name, bad_key, (100,))
+        # Key batch ndim must equal output ndim (too few).
+        with self.assertRaisesRegex(
+            RuntimeError, "batched key must have ndim == output ndim \\+ 1"
+        ):
+            self._gen(gen_fn_name, random.split(key, 3), (3, 4, 100))
+        # Key batch ndim must equal output ndim (too many).
+        with self.assertRaisesRegex(
+            RuntimeError, "batched key must have ndim == output ndim \\+ 1"
+        ):
+            self._gen(gen_fn_name, random.split(key, 3).view(3, 1, 1, 2), (3, 100))
+        # Key batch dims must be broadcastable with output.
+        with self.assertRaisesRegex(
+            RuntimeError, "is not broadcastable with output shape"
+        ):
+            self._gen(gen_fn_name, random.split(key, 5).unsqueeze(-2), (3, 100))
+
+    @parametrize("gen_fn_name", ["normal", "uniform"])
+    @dtypes(*all_floating_dtypes)
+    def test_offset_shift_consistency(self, device, dtype, gen_fn_name):
+        seed = 42
+        n = 100
+        key0 = random.key(seed, device=device)
+        ref = self._gen(gen_fn_name, key0, (n,), dtype=dtype)
+
+        # as a key's offset shifts, we expect the stream to shift by
+        # the number of elements per philox call (2 for double; 4 otherwise)
+        for offset in range(1, 4):
+            key = torch.tensor([seed, offset], dtype=torch.uint64, device=device)
+            elems_per_call = 2 if dtype == torch.float64 else 4
+            expected_shift = offset * elems_per_call
+            result = self._gen(gen_fn_name, key, (n,), dtype=dtype)
+            self.assertEqual(ref[expected_shift:], result[:-expected_shift])
+
+    @parametrize("gen_fn_name", ["normal", "uniform"])
+    @dtypes(*all_floating_dtypes)
+    def test_offset_overflow(self, device, dtype, gen_fn_name):
+        seed = 42
+        n = 100
+        last_offset_before_wrap = (1 << 64) - 1
+        key = torch.tensor(
+            [seed, last_offset_before_wrap], dtype=torch.uint64, device=device
+        )
+        result = self._gen(gen_fn_name, key, (n,), dtype=dtype)
+
+        # ensure offset wraps around to 0 by comparing with 0-offset key results
+        key0 = random.key(seed, device=device)
+        result0 = self._gen(gen_fn_name, key0, (n,), dtype=dtype)
+        elems_per_call = 2 if dtype == torch.float64 else 4
+        self.assertEqual(result[elems_per_call:], result0[:-elems_per_call])
+
+    @parametrize("gen_fn_name", ["normal", "uniform"])
+    @dtypes(*all_floating_dtypes)
+    def test_small_output_sizes(self, device, dtype, gen_fn_name):
+        key = random.key(42, device=device)
+        large = self._gen(gen_fn_name, key, (100,), dtype=dtype)
+        for n in [0, 1, 2, 3, 4, 5, 7]:
+            result = self._gen(gen_fn_name, key, (n,), dtype=dtype)
+            self.assertEqual(result.shape, (n,))
+            # Determinism.
+            result2 = self._gen(gen_fn_name, key, (n,), dtype=dtype)
+            self.assertEqual(result, result2)
+            # Prefix consistency: first n elements of a larger output.
+            if n > 0:
+                self.assertEqual(result, large[:n])
+
+    @parametrize("gen_fn_name", ["normal", "uniform"])
+    @parametrize("layout", ["contiguous", "noncontiguous", "unaligned"])
+    @dtypes(*all_floating_dtypes)
+    def test_inplace(self, device, dtype, gen_fn_name, layout):
+        key = random.key(42, device=device)
+        if layout == "contiguous":
+            result = torch.empty(1000, dtype=dtype, device=device)
+        elif layout == "noncontiguous":
+            result = torch.empty(2000, dtype=dtype, device=device)[::2]
+        else:
+            # Contiguous but data pointer is not aligned to vectorized write width.
+            result = torch.empty(1001, dtype=dtype, device=device)[1:]
+        inplace_fn = getattr(random, gen_fn_name + "_")
+        out = inplace_fn(key, result)
+        self.assertIs(out, result)
+        functional = self._gen(gen_fn_name, key, (1000,), dtype=dtype)
+        self.assertEqual(result, functional)
+
+    @parametrize("gen_fn_name", ["normal", "uniform"])
+    @dtypes(*all_floating_dtypes)
+    def test_empty_output(self, device, dtype, gen_fn_name):
+        key = random.key(42, device=device)
+        result = self._gen(gen_fn_name, key, (0,), dtype=dtype)
+        self.assertEqual(result.shape, (0,))
+        self.assertEqual(result.dtype, dtype)
+        result = self._gen(gen_fn_name, key, (3, 0), dtype=dtype)
+        self.assertEqual(result.shape, (3, 0))
+        self.assertEqual(result.dtype, dtype)
+
+    # Distribution-specific tests
+
+    @dtypes(*all_floating_dtypes)
+    def test_standard_normal_statistics(self, device, dtype):
+        key = random.key(42, device=device)
+        result = random.normal(key, (100000,), dtype=dtype)
+        self.assertTrue(abs(result.mean().item()) < 0.05)
+        self.assertTrue(abs(result.std().item() - 1.0) < 0.05)
+
+    @dtypes(*all_floating_dtypes)
+    def test_custom_mean_std(self, device, dtype):
+        key = random.key(42, device=device)
+        result = random.normal(key, (100000,), mean=5.0, std=2.0, dtype=dtype)
+        self.assertTrue(abs(result.mean().item() - 5.0) < 0.1)
+        self.assertTrue(abs(result.std().item() - 2.0) < 0.1)
+
+    @dtypes(*all_floating_dtypes)
+    def test_standard_uniform_statistics(self, device, dtype):
+        key = random.key(42, device=device)
+        result = random.uniform(key, (100000,), dtype=dtype)
+        self.assertTrue(abs(result.mean().item() - 0.5) < 0.05)
+        self.assertTrue(result.min().item() >= 0.0)
+        self.assertTrue(result.max().item() < 1.0)
+
+    @dtypes(*all_floating_dtypes)
+    def test_custom_low_high(self, device, dtype):
+        key = random.key(42, device=device)
+        result = random.uniform(key, (100000,), low=2.0, high=5.0, dtype=dtype)
+        self.assertTrue(abs(result.mean().item() - 3.5) < 0.1)
+        self.assertTrue(result.min().item() >= 2.0)
+        self.assertTrue(result.max().item() <= 5.0)
+
+
+class TestStatelessRNGCompile(TestCase):
+    def test_split_fullgraph(self, device):
+        key = random.key(42, device=device)
+
+        @torch.compile(backend="aot_eager", fullgraph=True)
+        def f(key):
+            return random.split(key, 4)
+
+        self.assertEqual(f(key), random.split(key, 4))
+
+    def test_fold_in_fullgraph(self, device):
+        key = random.key(42, device=device)
+
+        @torch.compile(backend="aot_eager", fullgraph=True)
+        def f(key):
+            return random.fold_in(key, 7)
+
+        self.assertEqual(f(key), random.fold_in(key, 7))
+
+    def test_uniform_fullgraph(self, device):
+        key = random.key(42, device=device)
+
+        @torch.compile(backend="aot_eager", fullgraph=True)
+        def f(key):
+            return random.uniform(key, (100,))
+
+        self.assertEqual(f(key), random.uniform(key, (100,)))
+
+    def test_normal_fullgraph(self, device):
+        key = random.key(42, device=device)
+
+        @torch.compile(backend="aot_eager", fullgraph=True)
+        def f(key):
+            return random.normal(key, (100,))
+
+        self.assertEqual(f(key), random.normal(key, (100,)))
+
+    def test_batched_normal_fullgraph(self, device):
+        key = random.key(42, device=device)
+        keys = random.split(key, 4).unsqueeze(-2)  # (4, 1, 2)
+
+        @torch.compile(backend="aot_eager", fullgraph=True)
+        def f(keys):
+            return random.normal(keys, (4, 50))
+
+        self.assertEqual(f(keys), random.normal(keys, (4, 50)))
+
+    def test_split_then_normal_fullgraph(self, device):
+        key = random.key(42, device=device)
+
+        @torch.compile(backend="aot_eager", fullgraph=True)
+        def f(key):
+            keys = random.split(key, 4).unsqueeze(-2)
+            return random.normal(keys, (4, 100))
+
+        self.assertEqual(
+            f(key), random.normal(random.split(key, 4).unsqueeze(-2), (4, 100))
+        )
+
+    def test_fold_in_then_uniform_fullgraph(self, device):
+        key = random.key(42, device=device)
+
+        @torch.compile(backend="aot_eager", fullgraph=True)
+        def f(key):
+            k = random.fold_in(key, 3)
+            return random.uniform(k, (100,))
+
+        self.assertEqual(f(key), random.uniform(random.fold_in(key, 3), (100,)))
+
+
+instantiate_device_type_tests(TestStatelessRNGKey, globals(), only_for=("cuda",))
+instantiate_device_type_tests(TestStatelessRNGKeySplit, globals(), only_for=("cuda",))
+instantiate_device_type_tests(TestStatelessRNGKeyFoldIn, globals(), only_for=("cuda",))
+instantiate_device_type_tests(
+    TestStatelessRNGDistribution, globals(), only_for=("cuda",)
+)
+instantiate_device_type_tests(TestStatelessRNGCompile, globals(), only_for=("cuda",))
+
+
+if __name__ == "__main__":
+    run_tests()
diff --git a/test/test_static_runtime.py b/test/test_static_runtime.py
index 8933f57eacf90..002157b4a7d40 100644
--- a/test/test_static_runtime.py
+++ b/test/test_static_runtime.py
@@ -2,7 +2,6 @@
 # ruff: noqa: F841
 
 import unittest
-from typing import Optional
 
 import numpy as np
 import torch
@@ -12,7 +11,7 @@
 
 
 def linear_shim(
-    input: torch.Tensor, weight: torch.Tensor, bias: Optional[torch.Tensor] = None
+    input: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor | None = None
 ) -> torch.Tensor:
     output = input.matmul(weight.t())
     if bias is not None:
diff --git a/test/test_tensor_creation_ops.py b/test/test_tensor_creation_ops.py
index 287ee3cb7e421..87e7db57318b6 100644
--- a/test/test_tensor_creation_ops.py
+++ b/test/test_tensor_creation_ops.py
@@ -4138,7 +4138,8 @@ def _check(self, original, cvt=lambda t: t, is_alias=True, same_dtype=True, same
             3. Whether the result lives in the expected device
             4. Whether the result has its 'requires_grad' set or not
         """
-        result = torch.asarray(cvt(original), **kwargs)
+        converted_original = cvt(original)
+        result = torch.asarray(converted_original, **kwargs)
         self.assertTrue(isinstance(result, torch.Tensor))
 
         # 1. The storage pointers should be equal only if 'is_alias' is set
@@ -4169,8 +4170,10 @@ def _check(self, original, cvt=lambda t: t, is_alias=True, same_dtype=True, same
         if device.index is not None:
             self.assertEqual(device.index, result.device.index)
 
-        # 4. By default, 'requires_grad' is unset
-        self.assertEqual(result.requires_grad, kwargs.get("requires_grad", False))
+        # 4. By default, 'requires_grad' mirrors the original tensor's requires_grad, if
+        # present.
+        original_requires_grad = converted_original.requires_grad if isinstance(converted_original, torch.Tensor) else False
+        self.assertEqual(result.requires_grad, kwargs.get("requires_grad", original_requires_grad))
 
     def _test_alias_with_cvt(self, cvt, device, dtype, shape=(5, 5), only_with_dtype=False):
         original = make_tensor(shape, dtype=dtype, device=device)
@@ -4335,7 +4338,7 @@ def test_retain_autograd_history(self, device, dtype):
 
         def check(**kwargs):
             a = torch.asarray(cloned, **kwargs)
-            requires_grad = kwargs.get("requires_grad", False)
+            requires_grad = kwargs.get("requires_grad", cloned.requires_grad)
             self.assertEqual(a.requires_grad, requires_grad)
             # Autograd history shouldn't be retained when requires_grad is False
             self.assertEqual(a.grad_fn is None, not requires_grad)
diff --git a/test/test_tensorboard.py b/test/test_tensorboard.py
index 8b63014f46749..1f93a86eb2aab 100644
--- a/test/test_tensorboard.py
+++ b/test/test_tensorboard.py
@@ -740,6 +740,7 @@ def test_torchvision_smoke(self):
 
 class TestTensorBoardFigure(BaseTestCase):
     @skipIfNoMatplotlib
+    @skipIfTorchDynamo("dynamo fails to trace matplotlib WRITEABLE flag and slice.indices")
     def test_figure(self):
         writer = self.createSummaryWriter()
 
@@ -765,6 +766,7 @@ def test_figure(self):
         writer.close()
 
     @skipIfNoMatplotlib
+    @skipIfTorchDynamo("dynamo fails to trace matplotlib WRITEABLE flag and slice.indices")
     def test_figure_list(self):
         writer = self.createSummaryWriter()
 
diff --git a/test/test_testing.py b/test/test_testing.py
index c712071bdab3a..f7032a7dea64f 100644
--- a/test/test_testing.py
+++ b/test/test_testing.py
@@ -22,7 +22,7 @@
 from torch.testing._internal.common_utils import (
     IS_FBCODE, IS_JETSON, IS_MACOS, IS_SANDCASTLE, IS_WINDOWS, TestCase, run_tests, slowTest,
     parametrize, reparametrize, subtest, instantiate_parametrized_tests, dtype_name,
-    TEST_WITH_ROCM, decorateIf
+    TEST_WITH_ROCM, decorateIf, skipIfXpu
 )
 from torch.testing._internal.common_device_type import \
     (PYTORCH_TESTING_DEVICE_EXCEPT_FOR_KEY, PYTORCH_TESTING_DEVICE_ONLY_FOR_KEY, dtypes,
@@ -282,6 +282,38 @@ def test_isclose_nan_equality_shortcut(self, device, dtype):
 
         self._isclose_helper(tests, device, dtype, equal_nan=True, rtol=0, atol=0)
 
+    @onlyNativeDeviceTypes
+    @dtypes(torch.float16, torch.float32)
+    def test_isclose_equal_nan_broadcast(self, device, dtype):
+        # Regression test: isclose with equal_nan=True should handle broadcasting
+        # See https://github.com/pytorch/pytorch/issues/174985
+        nan = torch.nan
+
+        # One-sided broadcast: different sizes
+        a = torch.tensor([nan], device=device, dtype=dtype)
+        b = torch.tensor([nan, nan], device=device, dtype=dtype)
+        result = torch.isclose(a, b, equal_nan=True)
+        self.assertEqual(result, torch.tensor([True, True], device=device))
+
+        # One-sided broadcast: scalar-like vs vector
+        a = torch.tensor([nan], device=device, dtype=dtype)
+        b = torch.tensor([nan, 1.0, nan], device=device, dtype=dtype)
+        result = torch.isclose(a, b, equal_nan=True)
+        self.assertEqual(result, torch.tensor([True, False, True], device=device))
+
+        # Mutual broadcast
+        a = torch.tensor([[nan], [1.0]], device=device, dtype=dtype)  # [2, 1]
+        b = torch.tensor([[nan, 1.0]], device=device, dtype=dtype)    # [1, 2]
+        result = torch.isclose(a, b, equal_nan=True)
+        expected = torch.tensor([[True, False], [False, True]], device=device)
+        self.assertEqual(result, expected)
+
+        # Same shape (fast path) still works
+        a = torch.tensor([nan, 1.0], device=device, dtype=dtype)
+        b = torch.tensor([nan, 1.0], device=device, dtype=dtype)
+        result = torch.isclose(a, b, equal_nan=True)
+        self.assertEqual(result, torch.tensor([True, True], device=device))
+
     # The following tests (test_cuda_assert_*) are added to ensure test suite terminates early
     # when CUDA assert was thrown. Because all subsequent test will fail if that happens.
     # These tests are slow because it spawn another process to run test suite.
@@ -2364,6 +2396,7 @@ def _check_python_output(cls, program) -> str:
             # fail, so just set CWD to this script's directory
             cwd=os.path.dirname(os.path.realpath(__file__)),).decode("utf-8")
 
+    @skipIfXpu(msg="The test is flaky on XPU, see https://github.com/pytorch/pytorch/issues/110040")
     # The test is flaky on ROCm/XPU and has been open and close multiple times
     # https://github.com/pytorch/pytorch/issues/110040
     def test_circular_dependencies(self) -> None:
@@ -2386,7 +2419,7 @@ def test_circular_dependencies(self) -> None:
                            "torch.distributed._tools.sac_ilp",  # depends on pulp
                            "torch.csrc",  # files here are devtools, not part of torch
                            "torch.include",  # torch include files after install
-                           "torch._inductor.kernel.vendored_templates.cutedsl_grouped_gemm",  # depends on cutlass
+                           "torch._inductor.kernel.vendored_templates.cutedsl",  # depends on cutlass
                            ]
         if IS_WINDOWS or IS_MACOS or IS_JETSON:
             # Distributed should be importable on Windows(except nn.api.), but not on Mac
@@ -2402,6 +2435,10 @@ def test_circular_dependencies(self) -> None:
             ignored_modules.append("torch.testing._internal.common_fsdp")
             ignored_modules.append("torch.testing._internal.common_distributed")
 
+        if sys.version_info < (3, 12):
+            # depends on Python 3.12+ syntax
+            ignored_modules.append("torch.testing._internal.py312_intrinsics")
+
         torch_dir = os.path.dirname(torch.__file__)
         for base, _, files in os.walk(torch_dir):
             prefix = os.path.relpath(base, os.path.dirname(torch_dir)).replace(os.path.sep, ".")
@@ -2439,6 +2476,10 @@ def test_not_import_sympy(self) -> None:
                          "  - Use TYPE_CHECKING if you are using sympy + strings if you are using sympy on type annotations\n"
                          "  - Import things that depend on SymPy locally")
 
+    def test_not_import_triton(self) -> None:
+        out = self._check_python_output("import torch;import sys;print('triton' not in sys.modules)")
+        self.assertEqual(out.strip(), "True")
+
     @parametrize('path', ['torch', 'functorch'])
     def test_no_mutate_global_logging_on_import(self, path) -> None:
         # Calling logging.basicConfig, among other things, modifies the global
diff --git a/test/test_torch.py b/test/test_torch.py
index fa254541d1ece..aef6bdf4c1565 100644
--- a/test/test_torch.py
+++ b/test/test_torch.py
@@ -15,7 +15,6 @@
 import random
 import re
 import copy
-import os
 import tempfile
 import unittest
 import warnings
@@ -84,52 +83,6 @@
 
 AMPERE_OR_ROCM = TEST_WITH_ROCM or torch.cuda.is_tf32_supported()
 
-@contextlib.contextmanager
-def torch_vital_set(value):
-    stash = None
-    if 'TORCH_VITAL' in os.environ:
-        stash = os.environ['TORCH_VITAL']
-    os.environ['TORCH_VITAL'] = value
-    try:
-        yield
-    finally:
-        if stash:
-            os.environ['TORCH_VITAL'] = stash
-        else:
-            del os.environ['TORCH_VITAL']
-
-# Tests Vital Signs for Torch
-# FIXME: document or deprecate whatever this is
-class TestBasicVitalSigns(TestCase):
-    def test_basic_vitals(self):
-        with torch_vital_set(''):
-            self.assertFalse(torch.vitals_enabled())
-        with torch_vital_set('ON'):
-            self.assertTrue(torch.vitals_enabled())
-
-    def test_basic_vitals_read_write(self):
-        with torch_vital_set('ON'):
-            self.assertTrue(torch.vitals_enabled())
-            # This tests the code path of setting a vital
-            self.assertTrue(torch.set_vital('Dataloader', 'basic_unit_test', 'TEST_VALUE_STRING'))
-            self.assertIn('TEST_VALUE_STRING', torch.read_vitals())
-            self.assertIn('CUDA.used', torch.read_vitals())
-
-    def test_dataloader_vitals(self):
-        with torch_vital_set('ON'):
-            inps = torch.arange(10 * 5, dtype=torch.float32).view(10, 5)
-            tgts = torch.arange(10 * 5, dtype=torch.float32).view(10, 5)
-            dataset = torch.utils.data.TensorDataset(inps, tgts)
-            torch.utils.data.DataLoader(dataset, batch_size=2)
-            self.assertIn('Dataloader.enabled\t\t True', torch.read_vitals())
-
-# FIXME: document or deprecate whatever this is
-class TestVitalSignsCuda(TestCase):
-    @onlyCUDA
-    def test_cuda_vitals_gpu_only(self, device):
-        with torch_vital_set('ON'):
-            self.assertIn('CUDA.used\t\t true', torch.read_vitals())
-
 
 is_cuda_sm86 = torch.cuda.is_available() and torch.cuda.get_device_capability(0) == (8, 6)
 
@@ -873,8 +826,6 @@ def test_check_tensor_internal(self, device):
     def test_cpp_warnings_have_python_context(self, device):
         # Creates long string in advance to avoid a too-long Python line
         s = ".+Triggered internally at.+RangeFactories.+"
-        # nvfuser deprecation warning filter
-        warnings.filterwarnings("ignore", "torch::jit::fuser::cuda", UserWarning)
 
         def cpp_warn_fn():
             out = torch.empty((5,))
@@ -883,6 +834,8 @@ def cpp_warn_fn():
 
         # Checks eager-mode cpp warning
         with warnings.catch_warnings(record=True) as w:
+            # nvfuser deprecation warning filter
+            warnings.filterwarnings("ignore", "torch::jit::fuser::cuda", UserWarning)
             cpp_warn_fn()
             frameinfo = inspect.getframeinfo(inspect.currentframe())
             warning = w[0]
@@ -894,11 +847,15 @@ def cpp_warn_fn():
             # Checks the Python features of the warning
             # Note: the eager mode warning refers to the line in the function
             # that throws the warning.
-            self.assertEqual(frameinfo.lineno - 6, warning.lineno)
+            self.assertEqual(frameinfo.lineno - 8, warning.lineno)
             self.assertEqual(len(w), 1)
 
         # Checks jitted cpp warning
         with warnings.catch_warnings(record=True) as w:
+            # nvfuser deprecation warning filter
+            warnings.filterwarnings("ignore", "torch::jit::fuser::cuda", UserWarning)
+            # ignore all deprecation warnings
+            warnings.filterwarnings("ignore", category=DeprecationWarning)
             scripted_cpp_warn_fn = torch.jit.script(cpp_warn_fn)
             scripted_cpp_warn_fn()
             warning = w[0]
@@ -919,6 +876,10 @@ def warn_fn():
 
         # The jit mimics an eager-mode Python warning in this case
         with warnings.catch_warnings(record=True) as w:
+            # nvfuser deprecation warning filter
+            warnings.filterwarnings("ignore", "torch::jit::fuser::cuda", UserWarning)
+            # ignore all deprecation warnings
+            warnings.filterwarnings("ignore", category=DeprecationWarning)
             scripted_warn_fn = torch.jit.script(warn_fn)
             scripted_warn_fn()
             frameinfo = inspect.getframeinfo(inspect.currentframe())
@@ -927,7 +888,7 @@ def warn_fn():
             self.assertTrue(re.search('Warning!', str(warning.message)) is not None)
 
             # Checks the Python features of the warning
-            self.assertEqual(frameinfo.lineno - 6, warning.lineno)
+            self.assertEqual(frameinfo.lineno - 10, warning.lineno)
             self.assertEqual(len(w), 1)
 
     # FIXME: move to test_testing
@@ -2313,8 +2274,7 @@ def test_uniform_kstest(self, device, dtype):
                     self.assertTrue(res.statistic < 0.1)
 
     @skipIfNoSciPy
-    @dtypes(*floating_types_and(torch.half))
-    @dtypesIfCUDA(*floating_types_and(torch.half, torch.bfloat16))
+    @dtypes(*floating_types_and(torch.half, torch.bfloat16))
     def test_normal_kstest(self, device, dtype):
         from scipy import stats
         size = 1000
@@ -2324,6 +2284,74 @@ def test_normal_kstest(self, device, dtype):
                 res = stats.kstest(t.cpu().to(torch.double), 'norm', args=(mean, std))
                 self.assertTrue(res.statistic < 0.1)
 
+    @skipIfNoSciPy
+    @dtypes(torch.half, torch.bfloat16)
+    def test_normal_small_std_kstest(self, device, dtype):
+        from scipy import stats
+        size = 50000
+        for std in [0.005, 0.01]:
+            t = torch.empty(size, dtype=dtype, device=device).normal_(mean=0, std=std)
+            res = stats.kstest(t.cpu().to(torch.double), 'norm', args=(0, std))
+            self.assertTrue(
+                res.statistic < 0.01,
+                msg=f"KS statistic {res.statistic:.4f} for {dtype} std={std}",
+            )
+
+    @skipIfNoSciPy
+    @dtypes(torch.half, torch.bfloat16)
+    def test_normal_kurtosis(self, device, dtype):
+        from scipy import stats
+        size = 50000
+        for std in [0.01, 0.1, 1.0]:
+            t = torch.empty(size, dtype=dtype, device=device).normal_(mean=0, std=std)
+            kurtosis = stats.kurtosis(t.cpu().to(torch.double).numpy())
+            self.assertAlmostEqual(
+                kurtosis, 0.0, delta=0.15,
+                msg=f"Excess kurtosis {kurtosis:.3f} too far from 0 for {dtype} std={std}",
+            )
+
+    @dtypes(torch.bfloat16)
+    def test_normal_tail_reach(self, device, dtype):
+        # With the broken bf16 Box-Muller (256 uniform inputs), the maximum
+        # output is capped at sqrt(-2*log(1/256)) ≈ 3.33 sigma. A correct
+        # implementation generates in float and reaches 4+ sigma.
+        size = 1000000
+        t = torch.empty(size, dtype=dtype, device=device).normal_(mean=0, std=1.0)
+        max_sigma = t.float().abs().max().item()
+        self.assertGreater(
+            max_sigma, 3.5,
+            msg=f"Max |z| = {max_sigma:.2f} sigma, expected > 3.5 (tail truncation)",
+        )
+
+    @dtypes(torch.half, torch.bfloat16)
+    def test_normal_unique_values(self, device, dtype):
+        size = 10000
+        t = torch.empty(size, dtype=dtype, device=device).normal_(mean=0, std=0.01)
+        n_unique = t.unique().numel()
+        # bf16 has ~200 representable values in [-0.03, 0.03] (3-sigma range).
+        # fp16 has ~60 representable values in [-0.03, 0.03].
+        # A correct implementation generates in fp32 and casts, so we should
+        # hit most representable values. The broken implementation produces
+        # far fewer due to quantized Box-Muller inputs.
+        if dtype == torch.bfloat16:
+            self.assertGreater(n_unique, 100)
+        else:
+            self.assertGreater(n_unique, 50)
+
+    @dtypes(torch.half, torch.bfloat16)
+    def test_uniform_unique_values(self, device, dtype):
+        size = 10000
+        t = torch.empty(size, dtype=dtype, device=device).uniform_(-100, 100)
+        n_unique = t.unique().numel()
+        # With a correct implementation (generate in float, cast to dtype),
+        # bf16 gives ~1400 unique and fp16 gives ~5200 for this range/size.
+        # The broken implementation (generate directly in dtype) gives only
+        # 256 for bf16 and ~2000 for fp16 due to using too few random bits.
+        if dtype == torch.bfloat16:
+            self.assertGreater(n_unique, 1000)
+        else:
+            self.assertGreater(n_unique, 4000)
+
     @skipIfMPS
     @skipIfNoSciPy
     @dtypes(*floating_types_and(torch.half, torch.bfloat16))
@@ -6287,6 +6315,60 @@ def f(x):
         result = grad_f(x)
         self.assertEqual(result.tolist() , [1.0, 1.0, 1.0])
 
+    def test_swap_data_ptr_with_empty(self, device):
+        x = torch.randn(4, device=device)
+        y = torch.randn(4, device=device)
+        y_data = y.clone()
+        x_storage = x.untyped_storage()
+        y_storage = y.untyped_storage()
+        x_storage.resize_(0)
+        self.assertEqual(x_storage.nbytes(), 0)
+        x_storage._swap_data_ptr_(y_storage)
+        self.assertEqual(x, y_data)
+        self.assertEqual(y_storage.nbytes(), 0)
+
+    def test_swap_data_ptr_non_empty(self, device):
+        x = torch.randn(4, device=device)
+        y = torch.randn(4, device=device)
+        x_data = x.clone()
+        y_data = y.clone()
+        x_storage = x.untyped_storage()
+        y_storage = y.untyped_storage()
+        x_storage._swap_data_ptr_(y_storage)
+        self.assertEqual(x, y_data)
+        self.assertEqual(y, x_data)
+
+    def test_swap_data_ptr_preserves_views(self, device):
+        x = torch.randn(4, 4, device=device)
+        x_view = x[1:3]
+        y = torch.randn(4, 4, device=device)
+        y_data = y.clone()
+        x_storage = x.untyped_storage()
+        y_storage = y.untyped_storage()
+        x_storage.resize_(0)
+        x_storage._swap_data_ptr_(y_storage)
+        self.assertEqual(x, y_data)
+        self.assertEqual(x_view, y_data[1:3])
+
+    def test_swap_data_ptr_self(self, device):
+        x = torch.randn(4, device=device)
+        x_data = x.clone()
+        s = x.untyped_storage()
+        s._swap_data_ptr_(s)
+        self.assertEqual(x, x_data)
+
+    def test_swap_data_ptr_device_mismatch(self, device):
+        x = torch.randn(4, device=device)
+        y = torch.randn(4, device="meta")
+        with self.assertRaisesRegex(RuntimeError, "same device"):
+            x.untyped_storage()._swap_data_ptr_(y.untyped_storage())
+
+    def test_swap_data_ptr_nbytes_mismatch(self, device):
+        x = torch.randn(4, device=device)
+        y = torch.randn(8, device=device)
+        with self.assertRaisesRegex(RuntimeError, "same nbytes"):
+            x.untyped_storage()._swap_data_ptr_(y.untyped_storage())
+
 
 # Tests that compare a device's computation with the (gold-standard) CPU's.
 class TestDevicePrecision(TestCase):
@@ -10331,32 +10413,32 @@ def __del__(self):
 
     @skipIfTorchDynamo("Not a suitable test for TorchDynamo")
     def test_storage_preserve_nonhermetic_in_hermetic_context(self):
-        from torch.library import Library, impl
+        from torch.library import _scoped_library, impl
         global _my_storage
 
-        my_lib = Library("my_lib", "DEF")  # noqa: TOR901
-        my_lib.define('my_func() -> None')
+        with _scoped_library("my_lib", "DEF") as my_lib:
+            my_lib.define('my_func() -> None')
 
-        a = torch.tensor([1.])
-        _my_storage = a.untyped_storage()
+            a = torch.tensor([1.])
+            _my_storage = a.untyped_storage()
 
-        m, t = Tracker.make()
-        _my_storage._tracker = t
-        del t
+            m, t = Tracker.make()
+            _my_storage._tracker = t
+            del t
 
-        @impl(my_lib, 'my_func', '')
-        def my_func():
-            global _my_storage
-            del _my_storage
+            @impl(my_lib, 'my_func', '')
+            def my_func():
+                global _my_storage
+                del _my_storage
 
-        self.assertFalse(m[0])
-        torch.ops.my_lib.my_func()
-        self.assertFalse(m[0])
+            self.assertFalse(m[0])
+            torch.ops.my_lib.my_func()
+            self.assertFalse(m[0])
 
-        s = a.untyped_storage()
-        del a
-        del s
-        self.assertTrue(m[0])
+            s = a.untyped_storage()
+            del a
+            del s
+            self.assertTrue(m[0])
 
     # FIXME: move to test_autograd?
     @skipIfTorchDynamo("TorchDynamo does not work well with hooks")
@@ -10550,6 +10632,23 @@ def generate_inputs(num_batches):
         finally:
             torch.set_num_threads(num_threads)
 
+    @unittest.skipIf(not torch.cuda.is_available(), "CUDA not available")
+    def test_bmm_matmul_mixed_dtype_error(self):
+        a = torch.randn(2, 8, 8, device="cuda", dtype=torch.float16)
+        b = torch.randn(2, 8, 64, device="cuda", dtype=torch.float32)
+
+        with self.assertRaisesRegex(RuntimeError, "expected scalar type .* but found"):
+            torch.bmm(a, b)
+
+        with self.assertRaisesRegex(RuntimeError, "expected scalar type .* but found"):
+            torch.compile(lambda x, y: torch.bmm(x, y), fullgraph=True)(a, b)
+
+        with self.assertRaisesRegex(RuntimeError, "expected scalar type .* but found"):
+            torch.matmul(a, b)
+
+        with self.assertRaisesRegex(RuntimeError, "expected scalar type .* but found"):
+            torch.compile(lambda x, y: torch.matmul(x, y), fullgraph=True)(a, b)
+
     def test_conj_neg_tolist(self):
         x = torch.randn(2, dtype=torch.cfloat)
         y1 = x.conj()
@@ -10872,7 +10971,6 @@ class TestTensorDeviceOps(TestCase):
 # pytest will fail.
 add_neg_dim_tests()
 instantiate_device_type_tests(TestViewOps, globals())
-instantiate_device_type_tests(TestVitalSignsCuda, globals())
 instantiate_device_type_tests(TestTensorDeviceOps, globals())
 instantiate_device_type_tests(TestTorchDeviceType, globals())
 instantiate_device_type_tests(TestDevicePrecision, globals(), except_for='cpu')
diff --git a/test/test_torch_config_hash_determinism.py b/test/test_torch_config_hash_determinism.py
index 74865ebc89226..0f81ea386b5a6 100644
--- a/test/test_torch_config_hash_determinism.py
+++ b/test/test_torch_config_hash_determinism.py
@@ -19,6 +19,10 @@
 
 
 class TestConfigModule(TestCase):
+    # Config keys that legitimately contain absolute paths, for example,
+    # /opt/clang-15/bin/clang
+    KNOWN_PATH_CONFIGS = {"cpp.cxx"}
+
     def check_deterministic(self, key: str, value: object):
         if isinstance(value, (int, float, bool)) or value is None:
             return
@@ -79,6 +83,8 @@ def test_inductor_config_hash_portable_deterministic(self):
         torch_config = inductor_config.save_config_portable()
 
         for key, value in torch_config.items():
+            if key in self.KNOWN_PATH_CONFIGS:
+                continue
             self.check_deterministic(key, value)
 
     def test_inductor_config_hash_portable_without_ignore(self):
@@ -93,6 +99,8 @@ def test_inductor_config_hash_portable_without_ignore(self):
                     f"Detected path in config value '.*', key='{cutlass_dir_key}'",
                 ):
                     for key, value in changed_torch_config.items():
+                        if key in self.KNOWN_PATH_CONFIGS:
+                            continue
                         self.check_deterministic(key, value)
             finally:
                 inductor_config._cache_config_ignore_prefix.insert(idx, cutlass_dir_key)
diff --git a/test/test_transformers.py b/test/test_transformers.py
index 777b85cb173d3..ea89876790944 100644
--- a/test/test_transformers.py
+++ b/test/test_transformers.py
@@ -18,7 +18,6 @@
 import itertools
 import torch.optim as optim
 from torch.testing._internal.common_device_type import expectedFailureMPS, instantiate_device_type_tests, onlyCUDA, largeTensorTest
-from typing import Optional
 import torch.utils.cpp_extension
 from torch.testing._internal.common_nn import NNTestCase
 from torch.testing._internal.common_utils import (
@@ -53,9 +52,12 @@
     PLATFORM_SUPPORTS_MEM_EFF_ATTENTION,
     PLATFORM_SUPPORTS_FUSED_ATTENTION,
     PLATFORM_SUPPORTS_CUDNN_ATTENTION,
+    PLATFORM_SUPPORTS_CK_SDPA,
     tf32_on_and_off,
     tf32_enabled,
 )
+from torch.testing._internal.common_device_type import skipXPUIf
+from torch.testing._internal.common_xpu import PLATFORM_SUPPORTS_FLASH_ATTENTION_XPU
 
 if TEST_FAIRSEQ:
     import fairseq.models.transformer as fairseq_transformer
@@ -89,14 +91,13 @@ def use_deterministic_algorithims(mode: bool, warn_only: bool):
 isSM5xDevice = torch.cuda.is_available() and torch.cuda.get_device_capability()[0] == 5
 isLessThanSM80Device = torch.cuda.is_available() and torch.cuda.get_device_capability()[0] < 8
 
-TEST_WITH_CK = TEST_WITH_ROCM and torch.backends.cuda.preferred_rocm_fa_library() == torch.backends.cuda._ROCmFABackends['ck']
 
 def _check_equal(
     golden: torch.Tensor,
     reference: torch.Tensor,
     test: torch.Tensor,
     fudge_factor: float,
-    tensor_name: Optional[str] = None
+    tensor_name: str | None = None
 ) -> None:
     """
     Compare test tensor against golden and reference tensors.
@@ -151,8 +152,8 @@ def check_out_and_grad(
     grad_query_tuple: tuple[torch.Tensor, torch.Tensor, torch.Tensor],
     grad_key_tuple: tuple[torch.Tensor, torch.Tensor, torch.Tensor],
     grad_value_tuple: tuple[torch.Tensor, torch.Tensor, torch.Tensor],
-    grad_attn_mask_tuple: Optional[tuple[torch.Tensor, torch.Tensor, torch.Tensor]] = None,
-    fudge_factors: Optional[dict[str, float]] = None
+    grad_attn_mask_tuple: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,
+    fudge_factors: dict[str, float] | None = None
 ) -> None:
     """
     Check output and gradients of attention mechanism tensors.
@@ -379,6 +380,31 @@ def test_multiheadattention_fastpath_attn_mask(self, device, attn_mask_dim, key_
             # of NaNs for fully masked rows
             self.assertEqual(out, out_fp.nan_to_num())
 
+    @onlyCUDA
+    def test_multiheadattention_fastpath_fp16_head_dim_alignment(self, device):
+        previous_fastpath = torch.backends.mha.get_fastpath_enabled()
+        try:
+            torch.manual_seed(0)
+            mha = nn.MultiheadAttention(
+                96,
+                16,
+                dropout=0.2,
+                batch_first=True,
+                dtype=torch.float16,
+                device=device,
+            )
+            x = torch.randn(8, 32, 96, device=device, dtype=torch.float16) * 238.15560380049612
+            with torch.no_grad():
+                torch.backends.mha.set_fastpath_enabled(True)
+                mha.eval()
+                out, _ = mha(x, x, x, need_weights=False)
+                torch.backends.mha.set_fastpath_enabled(False)
+                out_ref, _ = mha(x, x, x, need_weights=False)
+            self.assertTrue(torch.isfinite(out).all())
+            self.assertEqual(out, out_ref, atol=2e-2, rtol=8e-2)
+        finally:
+            torch.backends.mha.set_fastpath_enabled(previous_fastpath)
+
     @parametrize("nhead", [1, 4, 8])
     def test_transformerencoderlayer_src_mask(self, device, nhead):
         batch_size = 2
@@ -859,7 +885,7 @@ def test_math_backend_high_precision(self):
         mask = None
 
         def scaled_dot_product_attention(
-            xq: torch.Tensor, xk: torch.Tensor, xv: torch.Tensor, mask: Optional[torch.Tensor], backend: SDPBackend
+            xq: torch.Tensor, xk: torch.Tensor, xv: torch.Tensor, mask: torch.Tensor | None, backend: SDPBackend
         ) -> torch.Tensor:
             n_rep = 1
             xq, xk, xv = (tensor.transpose(1, 2) for tensor in (xq, xk, xv))
@@ -1113,7 +1139,7 @@ def forward(
                 )[0]
 
     @tf32_on_and_off(0.003)
-    @parametrize("batch_size", [0, 5])
+    @parametrize("batch_size", [1, 5])
     @parametrize("input_dim,attn_mask_dim,is_causal",
                  [(3, None, False), (3, 2, False), (3, 2, True), (3, 3, False), (3, 3, True),
                   (4, None, False), (4, 2, False), (4, 2, True), (4, 4, False), (4, 4, True)],
@@ -1595,19 +1621,67 @@ def test_invalid_fused_inputs_broadcast(self, device, kernel: SDPBackend):
             self.assertRaises(RuntimeError, lambda: torch.nn.functional.scaled_dot_product_attention(
                 q, k, v, None, 0.0, False))
 
-    @onlyCUDA
-    @unittest.skipIf(not PLATFORM_SUPPORTS_FUSED_ATTENTION, "Does not support fused scaled dot product attention")
-    @parametrize("kernel", PLATFORM_SPECIFIC_SDPA)
-    def test_invalid_sequence_lengths(self, device, kernel: SDPBackend):
-        with sdpa_kernel(backends=[kernel]):
-            # Passing in a q,k,v with 0 length sequences will error
-            dtype = torch.float16
-            make_tensor = partial(torch.rand, device=device, dtype=dtype)
-            size = SdpaShape(2, 2, 0, 8)
-            q, k, v = make_tensor(size), make_tensor(size), make_tensor(size)
-            with self.assertWarnsRegex(UserWarning, "All fused kernels do not support zero seq_len_q or seq_len_kv."):
-                self.assertRaises(RuntimeError, lambda: torch.nn.functional.scaled_dot_product_attention(
-                    q, k, v, None, 0.0, False))
+    def test_zero_sequence_lengths(self, device):
+        dtype = torch.float16
+        make_tensor = partial(torch.rand, device=device, dtype=dtype)
+        size = SdpaShape(2, 2, 0, 8)
+        q, k, v = make_tensor(size), make_tensor(size), make_tensor(size)
+        out = torch.nn.functional.scaled_dot_product_attention(q, k, v, None, 0.0, False)
+        self.assertEqual(out.shape, (2, 2, 0, 8))
+        self.assertEqual(out, torch.zeros(2, 2, 0, 8, dtype=dtype, device=device))
+
+    @skipIfTorchDynamo("mark_dynamic is not traceable")
+    def test_sdpa_unbacked_no_dde(self, device):
+        """
+        Test that scaled_dot_product_attention with dynamic symbolic dimensions
+        does not raise GuardOnDataDependentSymNode when checking for empty tensors.
+        This tests the TORCH_GUARD_OR_FALSE fix in attention.cpp.
+        Regression test for D94287846.
+        """
+
+        def func(q, k, v):
+            return F.scaled_dot_product_attention(q, k, v)
+
+        dtype = torch.float32
+        q = torch.rand(4, 2, 8, 16, dtype=dtype, device=device)
+        k = torch.rand(4, 2, 8, 16, dtype=dtype, device=device)
+        v = torch.rand(4, 2, 8, 16, dtype=dtype, device=device)
+
+        torch._dynamo.mark_dynamic(q, 0)
+        torch._dynamo.mark_dynamic(k, 0)
+        torch._dynamo.mark_dynamic(v, 0)
+
+        torch._dynamo.reset()
+        compiled_func = torch.compile(func, fullgraph=True, backend="eager")
+        result = compiled_func(q, k, v)
+        expected = func(q, k, v)
+        self.assertEqual(result.shape, expected.shape)
+
+    @skipIfTorchDynamo("mark_dynamic is not traceable")
+    def test_sdpa_empty_unbacked_no_dde(self, device):
+        """
+        Test that scaled_dot_product_attention with zero sequence length and
+        dynamic dimensions does not raise GuardOnDataDependentSymNode.
+        This tests the early return path in attention.cpp with symbolic shapes.
+        Regression test for D94287846.
+        """
+
+        def func(q, k, v):
+            return F.scaled_dot_product_attention(q, k, v)
+
+        dtype = torch.float32
+        q = torch.rand(2, 2, 0, 8, dtype=dtype, device=device)
+        k = torch.rand(2, 2, 0, 8, dtype=dtype, device=device)
+        v = torch.rand(2, 2, 0, 8, dtype=dtype, device=device)
+
+        torch._dynamo.mark_dynamic(q, 2)
+        torch._dynamo.mark_dynamic(k, 2)
+        torch._dynamo.mark_dynamic(v, 2)
+
+        torch._dynamo.reset()
+        compiled_func = torch.compile(func, fullgraph=True, backend="eager")
+        result = compiled_func(q, k, v)
+        self.assertEqual(result.shape, (2, 2, 0, 8))
 
     @onlyCUDA
     @unittest.skipIf(not PLATFORM_SUPPORTS_FUSED_ATTENTION, "Does not support fused scaled dot product attention")
@@ -2116,6 +2190,46 @@ def test_scaled_dot_product_attention_fp16_overflow(self, device):
         y = torch.nn.functional.scaled_dot_product_attention(x, x, x)
         self.assertFalse(y.isnan().any().item())
 
+    @parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16])
+    def test_sdpa_zero_sized_tensors(self, device, dtype):
+        zero_size_configs = [
+            (0, 4, 10, 64),
+            (2, 4, 0, 64),
+            (2, 4, 10, 0),
+            (0, 0, 0, 0),
+        ]
+        for b, h, s, d in zero_size_configs:
+            q = torch.zeros(b, h, s, d, dtype=dtype, device=device, requires_grad=True)
+            k = torch.zeros(b, h, s, d, dtype=dtype, device=device, requires_grad=True)
+            v = torch.zeros(b, h, s, 32, dtype=dtype, device=device, requires_grad=True)
+            out = torch.nn.functional.scaled_dot_product_attention(q, k, v)
+            self.assertEqual(out.shape, (b, h, s, 32))
+            self.assertEqual(out, torch.zeros(b, h, s, 32, dtype=dtype, device=device))
+            out.sum().backward()
+            self.assertEqual(q.grad.shape, q.shape)
+            self.assertEqual(k.grad.shape, k.shape)
+            self.assertEqual(v.grad.shape, v.shape)
+
+    def test_sdpa_output_shape_uses_value_head_dim(self, device):
+        # Regression test for https://github.com/pytorch/pytorch/issues/176767
+        test_cases = [
+            ((1, 16, 32), (1, 16, 32), (1, 16, 24)),
+            ((1, 1, 16, 32), (1, 1, 16, 32), (1, 1, 16, 24)),
+            ((2, 3, 1, 16, 32), (2, 3, 1, 16, 32), (2, 3, 1, 16, 24)),
+        ]
+
+        for q_shape, k_shape, v_shape in test_cases:
+            q = torch.randn(*q_shape, device=device, dtype=torch.float32)
+            k = torch.randn(*k_shape, device=device, dtype=torch.float32)
+            v = torch.randn(*v_shape, device=device, dtype=torch.float32)
+
+            actual = torch.nn.functional.scaled_dot_product_attention(q, k, v)
+
+            expected_shape = list(q_shape)
+            expected_shape[-1] = v_shape[-1]
+            self.assertEqual(actual.shape, torch.Size(expected_shape))
+
+
 class TestSDPACpuOnly(NNTestCase):
     """ Used to test CPU only functionality of scaled_dot_product_attention """
 
@@ -2905,6 +3019,32 @@ def fn(q, k, v):
             for use_compile in [False, True]:
                 test_attention(SDPBackend.CUDNN_ATTENTION, list(permute_order) + [3], use_compile)
 
+    @unittest.skipIf(not PLATFORM_SUPPORTS_CUDNN_ATTENTION, "cudnn Attention is not supported on this system")
+    @parametrize("dtype", [torch.float16, torch.bfloat16])
+    def test_cudnn_attention_broadcast_stride_zero(self, device, dtype):
+        # Regression test: Q/K/V with zero strides (broadcast via as_strided)
+        # caused alloc_with_matching_layout to produce output with stride[-1] != 1,
+        # which cuDNN Frontend rejects.
+        N, n_head, L, S, E, Ev = 8, 2, 64, 16, 128, 64
+        q = torch.randn(N, n_head, L, E, dtype=dtype, device=device)
+        k = torch.randn(N, n_head, S, E, dtype=dtype, device=device)
+        v = torch.randn(N, n_head, S, Ev, dtype=dtype, device=device)
+
+        q_bc = torch.as_strided(q, size=q.shape, stride=(0, 0, E, 1))
+        k_bc = torch.as_strided(k, size=k.shape, stride=(0, 0, E, 1))
+        v_bc = torch.as_strided(v, size=v.shape, stride=(0, 0, Ev, 1))
+
+        with sdpa_kernel(SDPBackend.CUDNN_ATTENTION):
+            out = F.scaled_dot_product_attention(q_bc, k_bc, v_bc, is_causal=True)
+
+        self.assertEqual(out.shape, (N, n_head, L, Ev))
+        self.assertEqual(out.stride(-1), 1)
+
+        out_ref = F.scaled_dot_product_attention(
+            q_bc.contiguous(), k_bc.contiguous(), v_bc.contiguous(), is_causal=True
+        )
+        torch.testing.assert_close(out, out_ref, atol=3e-3, rtol=3e-3)
+
     @unittest.skipIf(not PLATFORM_SUPPORTS_CUDNN_ATTENTION, "cudnn Attention is not supported on this system")
     @unittest.skipIf(
         not (isSM90Device and torch.backends.cudnn.version() >= 91000),
@@ -2972,6 +3112,15 @@ def test_cudnn_attention_seqlen1_dropout_heuristic(self):
             out = torch.nn.functional.scaled_dot_product_attention(q, q, q, dropout_p=0.5)
             out.backward(grad)
 
+    @unittest.skipIf(not PLATFORM_SUPPORTS_CUDNN_ATTENTION, "cudnn Attention is not supported on this system")
+    def test_cudnn_attention_low_dropout(self):
+        q = torch.randn(2, 8, 128, 128, dtype=torch.half, device='cuda')
+        dropout_p = 0.00000000001
+        out1 = torch.nn.functional.scaled_dot_product_attention(q, q, q, dropout_p=dropout_p)
+        out2 = torch.nn.functional.scaled_dot_product_attention(q, q, q, dropout_p=0.)
+        with self.assertRaisesRegex(AssertionError, "AssertionError not raised"):
+            self.assertNotEqual(out1, out2)
+
     @skipIfRocm
     @unittest.skipIf(not PLATFORM_SUPPORTS_CUDNN_ATTENTION, "cudnn Attention is not supported on this system")
     def test_cudnn_attention_broken_166211(self):
@@ -2996,6 +3145,82 @@ def test_cudnn_attention_broken_166211(self):
             self.assertFalse(dk.isnan().any())
             self.assertFalse(dv.isnan().any())
 
+    @unittest.skipIf(not PLATFORM_SUPPORTS_CUDNN_ATTENTION, "cudnn Attention is not supported on this system")
+    def test_cudnn_attention_mask_broken_177842(self):
+        # https://github.com/pytorch/pytorch/issues/177842
+        q = torch.randn(1, 10, 8, 8, dtype=torch.bfloat16, device='cuda')
+        k = torch.randn(1, 10, 1, 8, dtype=torch.bfloat16, device='cuda')
+        v = torch.randn(1, 10, 1, 8, dtype=torch.bfloat16, device='cuda')
+
+        attention_mask_custom = torch.zeros(10, 10, dtype=torch.bool).to("cuda")
+        attention_mask_custom[:7, :7] = torch.tril(torch.ones(7, 7, dtype=torch.bool), diagonal=0)
+
+        with sdpa_kernel(SDPBackend.MATH):
+            attn_output_math = torch.nn.functional.scaled_dot_product_attention(
+                query=q.transpose(1, 2),
+                key=k.transpose(1, 2),
+                value=v.transpose(1, 2),
+                attn_mask=attention_mask_custom,
+                is_causal=False,
+                enable_gqa=True,
+            )
+        with sdpa_kernel(SDPBackend.CUDNN_ATTENTION):
+            attn_output_cudnn = torch.nn.functional.scaled_dot_product_attention(
+                query=q.transpose(1, 2),
+                key=k.transpose(1, 2),
+                value=v.transpose(1, 2),
+                attn_mask=attention_mask_custom,
+                is_causal=False,
+                enable_gqa=True,
+            )
+        self.assertEqual(attn_output_math, attn_output_cudnn, atol=5e-3, rtol=3e-3)
+
+        attention_mask_custom = torch.zeros(10, 10, dtype=torch.bool).to("cuda")
+        attention_mask_custom[:7, :7] = torch.triu(torch.ones(7, 7, dtype=torch.bool), diagonal=0)
+
+        with sdpa_kernel(SDPBackend.MATH):
+            attn_output_math = torch.nn.functional.scaled_dot_product_attention(
+                query=q.transpose(1, 2),
+                key=k.transpose(1, 2),
+                value=v.transpose(1, 2),
+                attn_mask=attention_mask_custom,
+                is_causal=False,
+                enable_gqa=True,
+            )
+        with sdpa_kernel(SDPBackend.CUDNN_ATTENTION):
+            attn_output_cudnn = torch.nn.functional.scaled_dot_product_attention(
+                query=q.transpose(1, 2),
+                key=k.transpose(1, 2),
+                value=v.transpose(1, 2),
+                attn_mask=attention_mask_custom,
+                is_causal=False,
+                enable_gqa=True,
+            )
+        self.assertEqual(attn_output_math, attn_output_cudnn, atol=5e-3, rtol=3e-3)
+
+        attention_mask_custom = torch.zeros(10, 10, dtype=torch.bool).to("cuda")
+        attention_mask_custom[:7, :10] = torch.ones(7, 10, dtype=torch.bool)
+
+        with sdpa_kernel(SDPBackend.MATH):
+            attn_output_math = torch.nn.functional.scaled_dot_product_attention(
+                query=q.transpose(1, 2),
+                key=k.transpose(1, 2),
+                value=v.transpose(1, 2),
+                attn_mask=attention_mask_custom,
+                is_causal=False,
+                enable_gqa=True,
+            )
+        with sdpa_kernel(SDPBackend.CUDNN_ATTENTION):
+            attn_output_cudnn = torch.nn.functional.scaled_dot_product_attention(
+                query=q.transpose(1, 2),
+                key=k.transpose(1, 2),
+                value=v.transpose(1, 2),
+                attn_mask=attention_mask_custom,
+                is_causal=False,
+                enable_gqa=True,
+            )
+        self.assertEqual(attn_output_math, attn_output_cudnn, atol=5e-3, rtol=3e-3)
+
     @unittest.skipIf(not PLATFORM_SUPPORTS_MEM_EFF_ATTENTION, "Fused SDPA was not built for this system")
     @parametrize("mask_dim", [1, 2, 3, 4])
     def test_mem_efficient_attention_mask_variants(self, device, mask_dim: list[int]):
@@ -3130,8 +3355,6 @@ def test_scaled_dot_product_attention_fused_kernels_packed(self, device, type: s
     @parametrize("type", ["nested"])
     @parametrize("is_contiguous", [True, False])
     def test_scaled_dot_product_attention_cudnn_nested(self, device, type: str, is_contiguous: bool):
-        if TEST_WITH_ROCM and type == 'nested':
-            self.skipTest("ROCM does not support efficient attention on nested tensors, for now")
         make_tensor = partial(rand_sdpa_tensor, type=type, device=device, dtype=torch.float16, packed=True)
 
         batch_size, seq_len, num_heads, head_dim = 8, 64, 16, 64
@@ -3343,12 +3566,12 @@ def test_fused_sdp_choice(self, device, type: str):
         if "cuda" in str(device):
             device_capability = torch.cuda.get_device_capability()
         prefer_cudnn = "TORCH_CUDNN_SDPA_PREFERRED" not in os.environ or bool(os.environ["TORCH_CUDNN_SDPA_PREFERRED"])
-        # cuDNN prioritization requires cuDNN >= 9.9.0 (90900) per sdp_utils.cpp:83
+        # cuDNN prioritization requires cuDNN > 9.15.0 (91500) per sdp_utils.cpp:83
         cudnn_version = torch.backends.cudnn.version() if torch.backends.cudnn.is_available() else 0
         is_hopper_or_newer = device_capability and (device_capability[0] == 9 or device_capability[0] == 10)
-        prefer_cudnn = prefer_cudnn and is_hopper_or_newer and cudnn_version >= 90900
+        prefer_cudnn = prefer_cudnn and is_hopper_or_newer and cudnn_version > 91500
 
-        # cuDNN is enabled by default on SM 9.0/10.0 with cuDNN >= 9.9.0 (per #162073)
+        # cuDNN is enabled by default on SM 9.0/10.0 with cuDNN > 9.15.0 (per #169849)
         # For older cuDNN versions or other architectures, Flash Attention is preferred
         if type != "nested" and PLATFORM_SUPPORTS_CUDNN_ATTENTION and prefer_cudnn:
             self.assertEqual(torch._fused_sdp_choice(query, key, value), SDPBackend.CUDNN_ATTENTION.value)
@@ -3429,7 +3652,6 @@ def compiled_func(order):
         reset_order = torch._C._get_sdp_priority_order()
         self.assertEqual(default_order, reset_order, "expected SDPA context manager to reset priority order.")
 
-    @skipIfRocm  # Missing deterministic algo
     @unittest.skipIf(not PLATFORM_SUPPORTS_FUSED_ATTENTION, "Fused SDPA was not built for this system")
     @parametrize("fused_kernel", PLATFORM_SPECIFIC_SDPA)
     @parametrize("warn_only", [True, False])
@@ -3765,10 +3987,12 @@ def _get_mem_eff_drop_mask(batch_size, n_heads, q_len, kv_len, p, seed, offset,
     @parametrize("scale", [None, "l1"])
     @parametrize("enable_gqa", [True, False])
     @parametrize("n_heads", [[16, 8], [10, 2]])
+    @parametrize("sdpa_backend", ["aotriton", "ck"] if PLATFORM_SUPPORTS_CK_SDPA else ["aotriton"])
     @tf32_enabled()
     def test_flash_attention_vs_math_ref_grads(self, device, batch_size: int, seq_len_q: int, seq_len_k: int,
-                                               head_dim: int, is_causal: bool, dropout_p: float, dtype: torch.dtype,
-                                               scale: str, enable_gqa: bool, n_heads: list[int]):
+                                               head_dim: int, is_causal: bool, dropout_p: float,
+                                               dtype: torch.dtype, scale: str, enable_gqa: bool,
+                                               n_heads: list[int], sdpa_backend: str):
         if isSM8XDevice or isSM120Device and head_dim in range(193, 256 + 1):
             self.skipTest("Flash attention on sm86, sm87, and sm89 for headdim > 192 currently disabled")
         if is_causal and seq_len_q != seq_len_k:
@@ -3778,8 +4002,14 @@ def test_flash_attention_vs_math_ref_grads(self, device, batch_size: int, seq_le
         if max(seq_len_q, seq_len_k) >= 2048 and torch.cuda.get_device_properties('cuda').total_memory < 40 * 2**30:
             unittest.skip("Reference implementation OOM")
             return
-        if TEST_WITH_CK and dropout_p != 0:
-            self.skipTest("CK does not support tensor format dropout masks")
+
+        # ROCm now supports 2 different backends for SDPA that require different set up.
+        TEST_WITH_CK = False
+        if TEST_WITH_ROCM:
+            torch.backends.cuda.preferred_rocm_fa_library(sdpa_backend)
+            # When no args are given to preferred_rocm_fa_library, it acts as a getter
+            TEST_WITH_CK = (torch.backends.cuda.preferred_rocm_fa_library() == torch._C._ROCmFABackend.Ck)
+
         if TEST_WITH_CK and head_dim > 128:
             self.skipTest("CK does not support head dims over 128")
 
@@ -3853,15 +4083,24 @@ def test_flash_attention_vs_math_ref_grads(self, device, batch_size: int, seq_le
             softmax_mask = self.convert_flash_attn_S_to_softmax(
                 dbug_mask, seq_len_q, seq_len_k, query_padding_mask, key_padding_mask,
                 causal=is_causal)[:, :, :seq_len_q, :seq_len_k]
+
+            # This is the default implementation for the mask but we need to match CK if we are using it
             dropout_mask = softmax_mask >= 0
+
+            # This logic matches how CK calculates the dropout mask.
+            # This is necessary because CK doesn't support passing in custom dropout masks
+            # So we use this logic to ensure we are comparing apples to apples.
+            if TEST_WITH_CK:
+                dropout_mask = (softmax_mask <= int((1.0 - dropout_p) * 255.0)).to(torch.float32)
+
             # High Precision Math Reference
             out_ref = torch.ops.aten._scaled_dot_product_attention_math(
                 query_ref, key_ref, value_ref, dropout_p=dropout_p, is_causal=is_causal,
                 scale=scale, dropout_mask=dropout_mask, enable_gqa=enable_gqa)[0]
             # Low Precision Math Reference
             out_lp_ref = torch.ops.aten._scaled_dot_product_attention_math(
-                query, key, value, dropout_p=dropout_p, is_causal=is_causal, scale=scale,
-                dropout_mask=dropout_mask, enable_gqa=enable_gqa)[0]
+                query, key, value, dropout_mask=dropout_mask, dropout_p=dropout_p,
+                is_causal=is_causal, scale=scale, enable_gqa=enable_gqa)[0]
 
         upstream_grad = torch.rand_like(out, requires_grad=False)
 
@@ -3881,6 +4120,7 @@ def test_flash_attention_vs_math_ref_grads(self, device, batch_size: int, seq_le
             'grad_value': 4,
         }
         if TEST_WITH_ROCM:
+
             fudge_factors['grad_value'] = 6.0
             if TEST_WITH_CK:
                 fudge_factors['out'] = 5.0
@@ -4134,9 +4374,6 @@ def test_fused_kernels_nested_broadcasting(
         rand_nested_tensor = partial(rand_sdpa_tensor, type="nested", device=device, dtype=dtype)
         batch, num_heads, head_dim = 32, 8, 64
         head_dim_v = 32 if is_efficient else head_dim
-        if TEST_WITH_ROCM and head_dim != head_dim_v:
-            self.skipTest("head_dim != head_dim_v unsupported on ROCm for now")
-            return
         seq_lens_q = (torch.randint(low=1, high=5, size=(1,)).item()
                       if expand_q_batch
                       else torch.randint(low=1, high=32, size=(batch,)).tolist())
@@ -4200,7 +4437,6 @@ def _broadcast(t, batch_broadcasted, num_heads_broadcasted):
 
         self.assertEqual(actual.contiguous(), math_ref.contiguous().to(dtype), atol=1.5e-3, rtol=1e-2)
 
-    @skipIfRocm(msg="Efficient Attention on ROCM does not support head_dim != head_dim_v for now.")
     @unittest.skipIf(not PLATFORM_SUPPORTS_MEM_EFF_ATTENTION, "Fused SDPA was not built for this system")
     def test_fused_kernels_nested_broadcasting_query_dense(self, device):
         rand_nested_tensor = partial(rand_sdpa_tensor, type="nested", device=device, dtype=torch.float32)
@@ -4352,7 +4588,6 @@ class TestSDPAXpuOnly(NNTestCase):
     Mostly migrate from TestSDPACudaOnly in test/test_transformers.py
     """
 
-    PLATFORM_SUPPORTS_XPU_FLASH_ATTENTION = torch.xpu.is_available() and torch._C._is_flash_attention_available()
 
     @parametrize("type", ["dense"])
     @parametrize("dropout", [0.0, 0.7])
@@ -4659,7 +4894,7 @@ def test_onednn_attention_mask_vs_math(
 
         self.assertEqual(actual.float(), math_ref, atol=tol.atol, rtol=tol.rtol)
 
-    @unittest.skipIf(not PLATFORM_SUPPORTS_XPU_FLASH_ATTENTION, "XPU Flash Attention is not supported")
+    @skipXPUIf(not PLATFORM_SUPPORTS_FLASH_ATTENTION_XPU, "XPU Flash Attention is not supported")
     @parametrize("dtype", [torch.float32, torch.float64])
     def test_flash_attention_unsupport_dtypes(self, device, dtype):
         make_tensor = partial(torch.rand, device=device, dtype=dtype, requires_grad=False)
@@ -4677,7 +4912,7 @@ def test_flash_attention_unsupport_dtypes(self, device, dtype):
             with self.assertRaisesRegex(RuntimeError, "No available kernel"):
                 F.scaled_dot_product_attention(q, k, v)
 
-    @unittest.skipIf(not PLATFORM_SUPPORTS_XPU_FLASH_ATTENTION, "XPU Flash Attention is not supported")
+    @skipXPUIf(not PLATFORM_SUPPORTS_FLASH_ATTENTION_XPU, "XPU Flash Attention is not supported")
     def test_flash_attention_unsupport_dropout(self, device):
         dtype = torch.bfloat16
         make_tensor = partial(torch.rand, device=device, dtype=dtype, requires_grad=False)
@@ -4695,7 +4930,7 @@ def test_flash_attention_unsupport_dropout(self, device):
             with self.assertRaisesRegex(RuntimeError, "No available kernel"):
                 F.scaled_dot_product_attention(q, k, v, dropout_p=0.1)
 
-    @unittest.skipIf(not PLATFORM_SUPPORTS_XPU_FLASH_ATTENTION, "XPU Flash Attention is not supported")
+    @skipXPUIf(not PLATFORM_SUPPORTS_FLASH_ATTENTION_XPU, "XPU Flash Attention is not supported")
     def test_flash_attention_headdim_size(self, device):
         dtype = torch.bfloat16
         make_tensor = partial(torch.rand, device=device, dtype=dtype, requires_grad=False)
@@ -4719,7 +4954,7 @@ def test_flash_attention_headdim_size(self, device):
             with self.assertRaisesRegex(RuntimeError, "No available kernel"):
                 F.scaled_dot_product_attention(q, k, v)
 
-    @unittest.skipIf(not PLATFORM_SUPPORTS_XPU_FLASH_ATTENTION, "XPU Flash Attention is not supported")
+    @skipXPUIf(not PLATFORM_SUPPORTS_FLASH_ATTENTION_XPU, "XPU Flash Attention is not supported")
     def test_flash_attention_fail_with_non_square_causal_attention(self, device):
         dtype = torch.bfloat16
         q_shape = SdpaShape(1, 1, 8, 16)
@@ -4733,7 +4968,7 @@ def test_flash_attention_fail_with_non_square_causal_attention(self, device):
                 self.assertRaises(RuntimeError, lambda: torch.nn.functional.scaled_dot_product_attention(
                     q, k, v, None, 0.0, is_causal=True))
 
-    @unittest.skipIf(not PLATFORM_SUPPORTS_XPU_FLASH_ATTENTION, "XPU Flash Attention is not supported")
+    @skipXPUIf(not PLATFORM_SUPPORTS_FLASH_ATTENTION_XPU, "XPU Flash Attention is not supported")
     @parametrize("fused_kernel", [SDPBackend.FLASH_ATTENTION])
     @parametrize("dtype", [torch.half, torch.bfloat16])
     @parametrize("batch_size", [1, 2, 4])
@@ -4848,8 +5083,8 @@ def run_test(
         make_q,
         make_kv,
         attn_bias=None,
-        forw_tolerances: Optional[Tolerances] = None,
-        grad_tolerances: Optional[Tolerances] = None,
+        forw_tolerances: Tolerances | None = None,
+        grad_tolerances: Tolerances | None = None,
         backend=None,
         causal_variant=None,
     ):
@@ -4899,6 +5134,7 @@ def run_test(
         "shape",
         [(16, 16, 128, 128, 16), (16, 16, 128, 256, 32), (16, 16, 256, 128, 32), (1, 1, 23, 56, 15)],
     )
+    @skipXPUIf(not PLATFORM_SUPPORTS_FLASH_ATTENTION_XPU, "XPU Flash Attention is not supported")
     def test_causal_variants(self, device, causal_variant: CausalVariant, shape: list[tuple[int]]):
         make_tensor = partial(
             torch.rand, device=device, dtype=torch.float16, requires_grad=True
@@ -4931,6 +5167,7 @@ def test_causal_variants(self, device, causal_variant: CausalVariant, shape: lis
         [(16, 16, 128, 128, 16), (16, 16, 128, 256, 32), (16, 16, 256, 128, 32), (1, 1, 23, 56, 15)],
     )
     @skipIfTorchDynamo("This function already calls torch.compile.")
+    @skipXPUIf(not PLATFORM_SUPPORTS_FLASH_ATTENTION_XPU, "XPU Flash Attention is not supported")
     def test_causal_variants_compile(self, device, causal_variant: CausalVariant, shape: list[tuple[int]]):
         cnts = CompileCounterWithBackend("aot_eager")
         make_tensor = partial(
diff --git a/test/test_typing.py b/test/test_typing.py
index 91c5869aca84e..5564d1979baaa 100644
--- a/test/test_typing.py
+++ b/test/test_typing.py
@@ -8,7 +8,7 @@
 import unittest
 from collections import defaultdict
 from threading import Lock
-from typing import IO, Optional
+from typing import IO
 
 from torch.testing._internal.common_utils import (
     instantiate_parametrized_tests,
@@ -52,7 +52,7 @@ def _strip_filename(msg: str) -> str:
 def _run_mypy() -> dict[str, list[str]]:
     """Clears the cache and run mypy before running any of the typing tests."""
     if os.path.isdir(CACHE_DIR):
-        shutil.rmtree(CACHE_DIR)
+        shutil.rmtree(CACHE_DIR, ignore_errors=True)
 
     rc: dict[str, list[str]] = {}
     for directory in (REVEAL_DIR, PASS_DIR, FAIL_DIR):
@@ -97,9 +97,7 @@ def get_test_cases(directory):
 """
 
 
-def _test_fail(
-    path: str, error: str, expected_error: Optional[str], lineno: int
-) -> None:
+def _test_fail(path: str, error: str, expected_error: str | None, lineno: int) -> None:
     if expected_error is None:
         raise AssertionError(_FAIL_MSG1.format(lineno, error))
     elif error not in expected_error:
@@ -161,7 +159,7 @@ def _test_reveal(path: str, reveal: str, expected_reveal: str, lineno: int) -> N
 @unittest.skipIf(NO_MYPY, reason="Mypy is not installed")
 class TestTyping(TestCase):
     _lock = Lock()
-    _cached_output: Optional[dict[str, list[str]]] = None
+    _cached_output: dict[str, list[str]] | None = None
 
     @classmethod
     def get_mypy_output(cls) -> dict[str, list[str]]:
diff --git a/test/test_unary_ufuncs.py b/test/test_unary_ufuncs.py
index 4d440ad51bcc3..6fa4bd4705701 100644
--- a/test/test_unary_ufuncs.py
+++ b/test/test_unary_ufuncs.py
@@ -1911,6 +1911,33 @@ def test_mvlgamma_integer_promotion(self, device, dtype):
         self.assertTrue(result.dtype.is_floating_point)
         self.assertTrue(torch.all(torch.isfinite(result)))
 
+    @onlyCUDA
+    @dtypes(torch.float32, torch.float16, torch.bfloat16)
+    def test_fp8_e4m3fn_conversion_subnormals(self, device, dtype):
+        # Regression test for ptxas codegen bug on sm_100 where FADD in the
+        # subnormal conversion path gets wrong source register for odd elements
+        # in the 8-wide unrolled vectorized_elementwise_kernel.
+        # e4m3fn subnormals: |x| < 2^-6
+        torch.manual_seed(0)
+        N = 2**20
+        x = (torch.randn(N, dtype=dtype, device=device) * 1e-3).clamp(-448, 448)
+        y = x.to(torch.float8_e4m3fn)
+        ref = x.cpu().float().to(torch.float8_e4m3fn)
+        self.assertEqual(y.cpu().view(torch.uint8), ref.view(torch.uint8))
+
+    @onlyCUDA
+    @dtypes(torch.float32, torch.float16, torch.bfloat16)
+    def test_fp8_e5m2_conversion_subnormals(self, device, dtype):
+        # Same regression test for e5m2.
+        # e5m2 subnormals: |x| < 2^-14
+        torch.manual_seed(0)
+        N = 2**20
+        x = (torch.randn(N, dtype=dtype, device=device) * 1e-4).clamp(-57344, 57344)
+        y = x.to(torch.float8_e5m2)
+        ref = x.cpu().float().to(torch.float8_e5m2)
+        self.assertEqual(y.cpu().view(torch.uint8), ref.view(torch.uint8))
+
+
 instantiate_device_type_tests(TestUnaryUfuncs, globals())
 
 if __name__ == "__main__":
diff --git a/test/test_utils_config_module.py b/test/test_utils_config_module.py
index 07892daab9981..97ff0effa450b 100644
--- a/test/test_utils_config_module.py
+++ b/test/test_utils_config_module.py
@@ -1,6 +1,9 @@
 # Owner(s): ["module: unknown"]
 import os
 import pickle
+import queue
+import threading
+import warnings
 from unittest.mock import patch
 
 
@@ -9,8 +12,6 @@
 os.environ["ENV_STR"] = "1234"
 os.environ["ENV_STR_EMPTY"] = ""
 
-from typing import Optional
-
 from torch.testing._internal import (
     fake_config_module as config,
     fake_config_module2 as config2,
@@ -24,7 +25,11 @@ class TestConfigModule(TestCase):
     def tearDown(self):
         # Config changes get persisted between test cases
         for k in config._config:
-            config._config[k].user_override = _UNSET_SENTINEL
+            config._config[k].user_override.set(_UNSET_SENTINEL)
+        config._hash_digest = None
+        # Reset deprecation warning flags
+        for k in config._config:
+            config._config[k]._deprecation_warned = False
 
     def test_base_value_loading(self):
         self.assertTrue(config.e_bool)
@@ -44,8 +49,8 @@ def test_base_value_loading(self):
             config.does_not_exist
 
     def test_type_loading(self):
-        self.assertEqual(config.get_type("e_optional"), Optional[bool])
-        self.assertEqual(config.get_type("e_none"), Optional[bool])
+        self.assertEqual(config.get_type("e_optional"), bool | None)
+        self.assertEqual(config.get_type("e_none"), bool | None)
 
     def test_overrides(self):
         config.e_bool = False
@@ -83,7 +88,7 @@ def test_none_override_semantics(self):
         config.e_bool = None
         self.assertIsNone(config.e_bool)
         for k in config._config:
-            config._config[k].user_override = _UNSET_SENTINEL
+            config._config[k].user_override.set(_UNSET_SENTINEL)
 
     def test_reference_semantics(self):
         config.e_list.append(2)
@@ -113,7 +118,10 @@ def test_multi_env(self):
         self.assertTrue(config2.e_env_force_multi)
 
     def test_save_config(self):
-        p = config.save_config()
+        with warnings.catch_warnings():
+            warnings.filterwarnings("ignore", category=FutureWarning)
+            p = config.save_config()
+
         self.assertDictEqual(
             pickle.loads(p),
             {
@@ -141,6 +149,8 @@ def test_save_config(self):
                 "e_env_default_str_empty": "",
                 "e_env_force": True,
                 "e_optional": True,
+                "e_deprecated": True,
+                "e_not_deprecated": False,
             },
         )
         config.e_bool = False
@@ -149,8 +159,47 @@ def test_save_config(self):
         self.assertTrue(config.e_bool)
         self.assertFalse(config.e_ignored)
 
+    def test_save_config_with_patch(self):
+        self.assertTrue(config.e_bool)
+        with config.patch(e_bool=False):
+            p = config.save_config()
+            self.assertDictEqual(
+                pickle.loads(p),
+                {
+                    "_cache_config_ignore_prefix": ["magic_cache_config"],
+                    "e_bool": False,
+                    "e_dict": {1: 2},
+                    "e_float": 1.0,
+                    "e_int": 1,
+                    "e_list": [1],
+                    "e_none": None,
+                    "e_set": {1},
+                    "e_string": "string",
+                    "e_tuple": (1,),
+                    "nested.e_bool": True,
+                    "_e_ignored": True,
+                    "e_compile_ignored": True,
+                    "magic_cache_config_ignored": True,
+                    "_save_config_ignore": ["e_ignored"],
+                    "e_config": True,
+                    "e_jk": True,
+                    "e_jk_false": False,
+                    "e_env_default": True,
+                    "e_env_default_FALSE": False,
+                    "e_env_default_str": "1234",
+                    "e_env_default_str_empty": "",
+                    "e_env_force": True,
+                    "e_optional": True,
+                    "e_deprecated": True,
+                    "e_not_deprecated": False,
+                },
+            )
+
     def test_save_config_portable(self):
-        p = config.save_config_portable()
+        with warnings.catch_warnings():
+            warnings.filterwarnings("ignore", category=FutureWarning)
+            p = config.save_config_portable()
+
         self.assertDictEqual(
             p,
             {
@@ -175,6 +224,8 @@ def test_save_config_portable(self):
                 "e_env_default_str_empty": "",
                 "e_env_force": True,
                 "e_optional": True,
+                "e_deprecated": True,
+                "e_not_deprecated": False,
             },
         )
         config.e_bool = False
@@ -215,20 +266,24 @@ def test_codegen_config_function(self):
         )
 
     def test_get_hash(self):
-        hash_value = b"\x87\xf7\xc6\x1di\x7f\x96-\x85\xdc\x04\xd5\xd0\xf6\x1c\x87"
-        self.assertEqual(
-            config.get_hash(),
-            hash_value,
-        )
-        # Test cached value
-        self.assertEqual(
-            config.get_hash(),
-            hash_value,
-        )
-        self.assertEqual(
-            config.get_hash(),
-            hash_value,
-        )
+        hash_value = b"#d\x8b\xd3\xbc'\xf5\x0c\xcd\xb6\x87zDw6g"
+
+        with warnings.catch_warnings():
+            warnings.filterwarnings("ignore", category=FutureWarning)
+            self.assertEqual(
+                config.get_hash(),
+                hash_value,
+            )
+            # Test cached value
+            self.assertEqual(
+                config.get_hash(),
+                hash_value,
+            )
+            self.assertEqual(
+                config.get_hash(),
+                hash_value,
+            )
+
         config._hash_digest = "fake"
         self.assertEqual(config.get_hash(), "fake")
 
@@ -247,7 +302,11 @@ def test_get_hash(self):
         )
 
     def test_dict_copy_semantics(self):
-        p = config.shallow_copy_dict()
+        with warnings.catch_warnings():
+            warnings.filterwarnings("ignore", category=FutureWarning)
+            p = config.shallow_copy_dict()
+            p2 = config.to_dict()
+
         self.assertDictEqual(
             p,
             {
@@ -276,9 +335,10 @@ def test_dict_copy_semantics(self):
                 "e_env_default_str_empty": "",
                 "e_env_force": True,
                 "e_optional": True,
+                "e_deprecated": True,
+                "e_not_deprecated": False,
             },
         )
-        p2 = config.to_dict()
         self.assertEqual(
             p2,
             {
@@ -307,9 +367,15 @@ def test_dict_copy_semantics(self):
                 "e_env_default_str_empty": "",
                 "e_env_force": True,
                 "e_optional": True,
+                "e_deprecated": True,
+                "e_not_deprecated": False,
             },
         )
-        p3 = config.get_config_copy()
+
+        with warnings.catch_warnings():
+            warnings.filterwarnings("ignore", category=FutureWarning)
+            p3 = config.get_config_copy()
+
         self.assertEqual(
             p3,
             {
@@ -338,6 +404,8 @@ def test_dict_copy_semantics(self):
                 "e_env_default_str_empty": "",
                 "e_env_force": True,
                 "e_optional": True,
+                "e_deprecated": True,
+                "e_not_deprecated": False,
             },
         )
 
@@ -359,6 +427,65 @@ def test_patch(self):
             with config.patch("does_not_exist"):
                 pass
 
+    def _assert_patch_reentrant_across_threads(self, fn):
+        barrier = threading.Barrier(2, timeout=5)
+        errors: queue.SimpleQueue[str] = queue.SimpleQueue()
+
+        def worker():
+            try:
+                fn(barrier)
+            except Exception as e:
+                errors.put(repr(e))
+
+        threads = [threading.Thread(target=worker) for _ in range(2)]
+        for thread in threads:
+            thread.start()
+        for thread in threads:
+            thread.join()
+
+        error_messages = []
+        while True:
+            try:
+                error_messages.append(errors.get_nowait())
+            except queue.Empty:
+                break
+
+        self.assertFalse(
+            error_messages, f"concurrent patch usage failed: {error_messages}"
+        )
+        self.assertTrue(config.e_bool)
+
+    def test_patch_context_manager_is_reentrant_across_threads(self):
+        patcher = config.patch(e_bool=False)
+
+        def fn(barrier):
+            with patcher:
+                self.assertFalse(config.e_bool)
+                barrier.wait()
+                self.assertFalse(config.e_bool)
+
+        self._assert_patch_reentrant_across_threads(fn)
+
+    def test_patch_context_manager_is_reentrant_when_nested(self):
+        patcher = config.patch(e_bool=False)
+
+        with patcher:
+            self.assertFalse(config.e_bool)
+            with patcher:
+                self.assertFalse(config.e_bool)
+            self.assertFalse(config.e_bool)
+
+        self.assertTrue(config.e_bool)
+
+    def test_patch_decorator_is_reentrant_across_threads(self):
+        @config.patch(e_bool=False)
+        def fn(barrier):
+            self.assertFalse(config.e_bool)
+            barrier.wait()
+            self.assertFalse(config.e_bool)
+
+        self._assert_patch_reentrant_across_threads(fn)
+
     def test_make_closur_patcher(self):
         revert = config._make_closure_patcher(e_bool=False)()
         self.assertFalse(config.e_bool)
@@ -378,7 +505,7 @@ def test_bad_jk_type(self):
             AssertionError,
             msg="AssertionError: justknobs only support booleans, thisisnotvalid is not a boolean",
         ):
-            _ConfigEntry(Config(default="bad", justknob="fake_knob"))
+            _ConfigEntry(Config(default="bad", justknob="fake_knob"), "test")
 
     def test_alias(self):
         self.assertFalse(config2.e_aliasing_bool)
@@ -398,15 +525,145 @@ def test_reference_is_default(self):
     def test_invalid_config_int(self):
         with self.assertRaises(AssertionError):
             _ConfigEntry(
-                Config(default=2, env_name_default="FAKE_DISABLE", value_type=int)
+                Config(default=2, env_name_default="FAKE_DISABLE", value_type=int),
+                "test",
             )
 
     def test_invalid_config_float(self):
         with self.assertRaises(AssertionError):
             _ConfigEntry(
-                Config(default=2, env_name_force="FAKE_DISABLE", value_type=float)
+                Config(default=2, env_name_force="FAKE_DISABLE", value_type=float),
+                "test",
+            )
+
+    def test_deprecated_config(self):
+        """Test that deprecated configs warn on access but still work normally."""
+        config._config["e_deprecated"]._deprecation_warned = False
+
+        with warnings.catch_warnings(record=True) as w:
+            warnings.simplefilter("always")
+
+            # Read access triggers warning
+            self.assertTrue(config.e_deprecated)
+            self.assertEqual(len(w), 1)
+            self.assertTrue(issubclass(w[0].category, FutureWarning))
+            msg = str(w[0].message)
+            self.assertIn("e_deprecated", msg)
+            self.assertIn("is no longer needed", msg)
+            self.assertIn("will be removed in a future version of PyTorch", msg)
+
+        # Write access also triggers warning
+        config._config["e_deprecated"]._deprecation_warned = False
+        with warnings.catch_warnings(record=True) as w:
+            warnings.simplefilter("always")
+            config.e_deprecated = False
+            self.assertEqual(len(w), 1)
+            self.assertFalse(config.e_deprecated)
+
+        # Functionality still works: patch, multiple accesses, etc.
+        config.e_deprecated = True
+        with config.patch(e_deprecated=False):
+            self.assertFalse(config.e_deprecated)
+        self.assertTrue(config.e_deprecated)
+
+    def test_deprecated_config_warn_once(self):
+        """Test that deprecation warning is only issued once per config."""
+        config._config["e_deprecated"]._deprecation_warned = False
+
+        with warnings.catch_warnings(record=True) as w:
+            warnings.simplefilter("always")
+            # First access
+            _ = config.e_deprecated
+            # Second access
+            _ = config.e_deprecated
+            config.e_deprecated = True  # Write - no warn
+            _ = config.e_deprecated  # Read again - no warn
+            self.assertEqual(len(w), 1)
+
+    def test_deprecated_alias_config(self):
+        """Test that aliased deprecated configs issue their own warning."""
+        config._config["e_deprecated_alias"]._deprecation_warned = False
+
+        with warnings.catch_warnings(record=True) as w:
+            warnings.simplefilter("always")
+            value = config.e_deprecated_alias
+            # Comes from e_not_deprecated = False
+            self.assertFalse(value)
+            self.assertEqual(len(w), 1)
+            self.assertIn("e_deprecated_alias", str(w[0].message))
+            self.assertIn("use something else instead", str(w[0].message))
+
+    def test_no_warnings_for_normal_user_code(self):
+        """Test that normal user operations don't produce deprecation warnings."""
+        import torch
+
+        with warnings.catch_warnings(record=True) as w:
+            warnings.simplefilter("always")
+
+            # Simple user code - compile and run a function
+            @torch.compile(backend="eager")
+            def fn(x):
+                return x + 1
+
+            _ = fn(torch.ones(3))
+
+            # No deprecation warnings should be issued from config system
+            deprecation_warnings = [
+                warning
+                for warning in w
+                if issubclass(warning.category, FutureWarning)
+                and "deprecated" in str(warning.message).lower()
+                and "config" in str(warning.message).lower()
+            ]
+            self.assertEqual(
+                len(deprecation_warnings),
+                0,
+                f"Unexpected config deprecation warnings: {[str(x.message) for x in deprecation_warnings]}",
             )
 
+    def test_patch_then_global(self):
+        self.assertTrue(config.e_bool)
+        with config.patch(e_bool=False):
+            self.assertFalse(config.e_bool)
+
+        config.e_bool = False
+        self.assertFalse(config.e_bool)
+
+    def test_is_default_patch(self):
+        self.assertTrue(config.e_bool)
+        with config.patch(e_bool=False):
+            self.assertFalse(config._is_default("e_bool"))
+
+    def test_dict_patch(self):
+        self.assertTrue(config.e_bool)
+        with config.patch(e_bool=False):
+            d = config._get_dict()
+            self.assertFalse(d["e_bool"])
+
+    def test_set_in_patch(self):
+        self.assertEqual(config.e_int, 1)
+        with config.patch(e_int=2):
+            self.assertEqual(config.e_int, 2)
+            config.e_int = 3
+            self.assertEqual(config.e_int, 3)
+        # Exiting the patch resets e_int to 1, losing the fact that we explicitly set it to 4 - see ConfigModule._do_setattr
+        self.assertEqual(config.e_int, 1)
+
+    def test_nested_patch(self):
+        self.assertEqual(config.e_int, 1)
+        self.assertEqual(config.e_string, "string")
+        with config.patch(e_int=2, e_string="inner"):
+            self.assertEqual(config.e_int, 2)
+            self.assertEqual(config.e_string, "inner")
+            with config.patch(e_int=3):
+                self.assertEqual(config.e_int, 3)
+                self.assertEqual(config.e_string, "inner")
+                config.e_int = 4
+                self.assertEqual(config.e_int, 4)
+            self.assertEqual(config.e_int, 2)
+        self.assertEqual(config.e_int, 1)
+        self.assertEqual(config.e_string, "string")
+
 
 if __name__ == "__main__":
     run_tests()
diff --git a/test/test_varlen_attention.py b/test/test_varlen_attention.py
index d01578cb5623f..f1283857a4d83 100644
--- a/test/test_varlen_attention.py
+++ b/test/test_varlen_attention.py
@@ -1,20 +1,76 @@
 # Owner(s): ["module: sdpa"]
 import unittest
 from collections import namedtuple
+from contextlib import contextmanager, nullcontext
 
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-from torch.nn.attention.varlen import varlen_attn
-from torch.testing._internal.common_cuda import PLATFORM_SUPPORTS_FLASH_ATTENTION
+from torch.nn.attention import (
+    activate_flash_attention_impl,
+    restore_flash_attention_impl,
+)
+from torch.nn.attention.varlen import varlen_attn, varlen_attn_out
+from torch.testing._internal.common_cuda import (
+    IS_SM90,
+    PLATFORM_SUPPORTS_CK_SDPA,
+    PLATFORM_SUPPORTS_FLASH_ATTENTION,
+    SM100OrLater,
+    SM120OrLater,
+    SM90OrLater,
+)
 from torch.testing._internal.common_device_type import instantiate_device_type_tests
 from torch.testing._internal.common_nn import NNTestCase
-from torch.testing._internal.common_utils import parametrize, run_tests
+from torch.testing._internal.common_utils import (
+    decorateIf,
+    parametrize,
+    run_tests,
+    TEST_WITH_ROCM,
+)
 from torch.utils._python_dispatch import TorchDispatchMode
 
 
+@contextmanager
+def use_fa3():
+    try:
+        activate_flash_attention_impl("FA3")
+    except (ModuleNotFoundError, RuntimeError) as err:
+        raise unittest.SkipTest(
+            "FA3 backend not available (flash_attn_interface missing)"
+        ) from err
+    try:
+        yield
+    finally:
+        restore_flash_attention_impl()
+
+
+@contextmanager
+def use_fa4():
+    try:
+        activate_flash_attention_impl("FA4")
+    except (ModuleNotFoundError, RuntimeError) as err:
+        raise unittest.SkipTest("FA4 backend not available") from err
+    try:
+        yield
+    finally:
+        restore_flash_attention_impl()
+
+
+def _use_backend(backend):
+    return {"fa2": nullcontext, "fa3": use_fa3, "fa4": use_fa4}[backend]()
+
+
+def _varlen_backends(*, include_fa4_paged_kv: bool) -> list[str]:
+    fa4_supported = (
+        SM100OrLater if include_fa4_paged_kv else SM90OrLater
+    ) and not SM120OrLater
+    return ["fa2"] + (["fa3"] if IS_SM90 else []) + (["fa4"] if fa4_supported else [])
+
+
 VarlenShape = namedtuple(
-    "VarlenShape", ["batch_size", "max_seq_len", "embed_dim", "num_heads"]
+    "VarlenShape",
+    ["batch_size", "max_seq_len", "embed_dim", "num_heads", "num_kv_heads"],
+    defaults=[None],
 )
 
 
@@ -32,31 +88,48 @@ def __torch_dispatch__(self, func, types, args=(), kwargs=None):
 
 class AttentionBlock(nn.Module):
     def __init__(
-        self, embed_dim: int, num_heads: int, device: torch.device, dtype: torch.dtype
+        self,
+        embed_dim: int,
+        num_heads: int,
+        device: torch.device,
+        dtype: torch.dtype,
+        num_kv_heads: int | None = None,
     ):
         super().__init__()
         self.embed_dim = embed_dim
         self.num_heads = num_heads
+        self.num_kv_heads = num_kv_heads if num_kv_heads is not None else num_heads
         self.head_dim = embed_dim // num_heads
 
-        self.qkv_proj = nn.Linear(
-            embed_dim, 3 * embed_dim, bias=False, device=device, dtype=dtype
+        self.q_proj = nn.Linear(
+            embed_dim,
+            num_heads * self.head_dim,
+            bias=False,
+            device=device,
+            dtype=dtype,
+        )
+        self.kv_proj = nn.Linear(
+            embed_dim,
+            2 * self.num_kv_heads * self.head_dim,
+            bias=False,
+            device=device,
+            dtype=dtype,
         )
         self.out_proj = nn.Linear(
             embed_dim, embed_dim, bias=False, device=device, dtype=dtype
         )
 
+    @property
+    def enable_gqa(self):
+        return self.num_kv_heads != self.num_heads
+
     def get_varlen_qkv(
         self,
         x_packed: torch.Tensor,
     ):
-        qkv = self.qkv_proj(x_packed)
-        q, k, v = qkv.chunk(3, dim=-1)
-
-        q = q.view(-1, self.num_heads, self.head_dim)
-        k = k.view(-1, self.num_heads, self.head_dim)
-        v = v.view(-1, self.num_heads, self.head_dim)
-
+        q = self.q_proj(x_packed).view(-1, self.num_heads, self.head_dim)
+        kv = self.kv_proj(x_packed).view(-1, 2, self.num_kv_heads, self.head_dim)
+        k, v = kv[:, 0], kv[:, 1]
         return q, k, v
 
     def forward_varlen(
@@ -79,6 +152,7 @@ def forward_varlen(
             max_len,
             scale=scale,
             window_size=window_size,
+            enable_gqa=self.enable_gqa,
         )
         attn_out = attn_out.view(-1, self.embed_dim)
 
@@ -93,8 +167,11 @@ def forward_sdpa(
     ):
         batch_size, seq_len, _ = x_padded.shape
 
-        qkv = self.qkv_proj(x_padded)
-        q, k, v = qkv.chunk(3, dim=-1)
+        q = self.q_proj(x_padded)
+        kv = self.kv_proj(x_padded)
+        k, v = kv.view(batch_size, seq_len, 2, self.num_kv_heads, self.head_dim).unbind(
+            dim=2
+        )
 
         padding_mask = (
             torch.arange(seq_len, device=x_padded.device)[None, :]
@@ -125,13 +202,33 @@ def forward_sdpa(
             attn_mask = attn_mask & window_mask[None, None, :, :]
 
         q = q.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
-        k = k.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
-        v = v.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
-
-        # Don't pass is_causal since we already incorporated it into attn_mask
-        attn_out = F.scaled_dot_product_attention(
-            q, k, v, attn_mask=attn_mask, scale=scale
+        k = k.view(batch_size, seq_len, self.num_kv_heads, self.head_dim).transpose(
+            1, 2
         )
+        v = v.view(batch_size, seq_len, self.num_kv_heads, self.head_dim).transpose(
+            1, 2
+        )
+
+        # Don't pass is_causal since we already incorporated it into attn_mask.
+        if self.enable_gqa:
+            # Force math backend for GQA
+            with torch.nn.attention.sdpa_kernel(torch.nn.attention.SDPBackend.MATH):
+                attn_out = F.scaled_dot_product_attention(
+                    q,
+                    k,
+                    v,
+                    attn_mask=attn_mask,
+                    scale=scale,
+                    enable_gqa=True,
+                )
+        else:
+            attn_out = F.scaled_dot_product_attention(
+                q,
+                k,
+                v,
+                attn_mask=attn_mask,
+                scale=scale,
+            )
 
         attn_out = (
             attn_out.transpose(1, 2)
@@ -196,7 +293,18 @@ class TestVarlenAttention(NNTestCase):
         not PLATFORM_SUPPORTS_FLASH_ATTENTION, "Flash Attention not supported"
     )
     @parametrize("dtype", [torch.bfloat16, torch.float16])
-    def test_basic_functionality(self, device, dtype):
+    @parametrize(
+        "sdpa_backend",
+        ["aotriton", "ck"] if PLATFORM_SUPPORTS_CK_SDPA else ["aotriton"],
+    )
+    @parametrize(
+        "backend",
+        _varlen_backends(include_fa4_paged_kv=False),
+    )
+    def test_basic_functionality(self, device, dtype, backend, sdpa_backend=None):
+        if TEST_WITH_ROCM:
+            torch.backends.cuda.preferred_rocm_fa_library(sdpa_backend)
+
         torch.manual_seed(42)
 
         shape = VarlenShape(batch_size=2, max_seq_len=512, embed_dim=1024, num_heads=16)
@@ -217,32 +325,59 @@ def test_basic_functionality(self, device, dtype):
             [0, shape.max_seq_len, total_tokens], device=device, dtype=torch.int32
         )
 
-        output = attention_block.forward_varlen(x_packed, cu_seq, shape.max_seq_len)
+        with _use_backend(backend):
+            output = attention_block.forward_varlen(x_packed, cu_seq, shape.max_seq_len)
 
-        self.assertEqual(output.shape, (total_tokens, shape.embed_dim))
-        self.assertEqual(output.device, torch.device(device))
-        self.assertEqual(output.dtype, dtype)
+            self.assertEqual(output.shape, (total_tokens, shape.embed_dim))
+            self.assertEqual(output.device, torch.device(device))
+            self.assertEqual(output.dtype, dtype)
 
-        varlen_grad_out = torch.ones_like(output)
+            # varlen_attn_out should produce the same result and write into the buffer
+            with torch.no_grad():
+                q, k, v = attention_block.get_varlen_qkv(x_packed)
+                expected = varlen_attn(
+                    q, k, v, cu_seq, cu_seq, shape.max_seq_len, shape.max_seq_len
+                )
+                out_buf = torch.empty_like(expected)
+                actual = varlen_attn_out(
+                    out_buf,
+                    q,
+                    k,
+                    v,
+                    cu_seq,
+                    cu_seq,
+                    shape.max_seq_len,
+                    shape.max_seq_len,
+                )
+                self.assertEqual(actual.data_ptr(), out_buf.data_ptr())
+                self.assertEqual(out_buf, expected)
 
-        varlen_grad = torch.autograd.grad(
-            outputs=output,
-            inputs=x_packed,
-            grad_outputs=varlen_grad_out,
-            retain_graph=True,
-            create_graph=False,
-            allow_unused=False,
-        )[0]
+            varlen_grad_out = torch.ones_like(output)
 
-        self.assertIsNotNone(varlen_grad)
-        self.assertEqual(varlen_grad.shape, x_packed.shape)
-        self.assertEqual(varlen_grad.dtype, x_packed.dtype)
+            varlen_grad = torch.autograd.grad(
+                outputs=output,
+                inputs=x_packed,
+                grad_outputs=varlen_grad_out,
+                retain_graph=True,
+                create_graph=False,
+                allow_unused=False,
+            )[0]
+
+            self.assertIsNotNone(varlen_grad)
+            self.assertEqual(varlen_grad.shape, x_packed.shape)
+            self.assertEqual(varlen_grad.dtype, x_packed.dtype)
 
     @unittest.skipIf(
         not PLATFORM_SUPPORTS_FLASH_ATTENTION, "Flash Attention not supported"
     )
+    @parametrize(
+        "sdpa_backend",
+        ["aotriton", "ck"] if PLATFORM_SUPPORTS_CK_SDPA else ["aotriton"],
+    )
     @parametrize("dtype", [torch.bfloat16, torch.float16])
-    def test_custom_op_compliance(self, device, dtype):
+    def test_custom_op_compliance(self, device, dtype, sdpa_backend=None):
+        if TEST_WITH_ROCM:
+            torch.backends.cuda.preferred_rocm_fa_library(sdpa_backend)
         torch.manual_seed(42)
 
         shape = VarlenShape(batch_size=2, max_seq_len=512, embed_dim=1024, num_heads=16)
@@ -295,11 +430,35 @@ def test_custom_op_compliance(self, device, dtype):
             test_utils=["test_schema", "test_faketensor"],
         )
 
+        # opcheck for _varlen_attn_out (no backward)
+        out_buf = torch.empty_like(q)
+        torch.library.opcheck(
+            torch.ops.torch_attn._varlen_attn_out,
+            (
+                out_buf,
+                q,
+                k,
+                v,
+                cu_seq,
+                cu_seq,
+                shape.max_seq_len,
+                shape.max_seq_len,
+                False,
+            ),
+            test_utils=["test_schema", "test_faketensor"],
+        )
+
     @unittest.skipIf(
         not PLATFORM_SUPPORTS_FLASH_ATTENTION, "Flash Attention not supported"
     )
+    @parametrize(
+        "sdpa_backend",
+        ["aotriton", "ck"] if PLATFORM_SUPPORTS_CK_SDPA else ["aotriton"],
+    )
     @parametrize("dtype", [torch.bfloat16, torch.float16])
-    def test_custom_op_registration(self, device, dtype):
+    def test_custom_op_registration(self, device, dtype, sdpa_backend=None):
+        if TEST_WITH_ROCM:
+            torch.backends.cuda.preferred_rocm_fa_library(sdpa_backend)
         torch.manual_seed(42)
 
         shape = VarlenShape(batch_size=2, max_seq_len=512, embed_dim=1024, num_heads=16)
@@ -321,7 +480,7 @@ def test_custom_op_registration(self, device, dtype):
         )
 
         compiled_forward = torch.compile(
-            attention_block.forward_varlen, backend="eager", fullgraph=True
+            attention_block.forward_varlen, backend="eager"
         )
         with OpLoggingMode() as mode:
             output = compiled_forward(x_packed, cu_seq, shape.max_seq_len)
@@ -344,9 +503,28 @@ def test_custom_op_registration(self, device, dtype):
         if not custom_ops_called:
             raise AssertionError("custom varlen attention ops should have been called")
 
+        # Also verify _varlen_attn_out dispatches correctly under compile
+        q, k, v = attention_block.get_varlen_qkv(x_packed.detach())
+
+        def run_varlen_out(q, k, v, cu_seq, max_len):
+            out_buf = torch.empty_like(q)
+            varlen_attn_out(out_buf, q, k, v, cu_seq, cu_seq, max_len, max_len)
+            return out_buf
+
+        compiled_out = torch.compile(run_varlen_out, backend="eager")
+        with OpLoggingMode() as out_mode:
+            compiled_out(q, k, v, cu_seq, shape.max_seq_len)
+
+        if not any("torch_attn._varlen_attn_out" in op for op in out_mode.called_ops):
+            raise AssertionError("custom _varlen_attn_out op should have been called")
+
     @unittest.skipIf(
         not PLATFORM_SUPPORTS_FLASH_ATTENTION, "Flash Attention not supported"
     )
+    @parametrize(
+        "sdpa_backend",
+        ["aotriton", "ck"] if PLATFORM_SUPPORTS_CK_SDPA else ["aotriton"],
+    )
     @parametrize("dtype", [torch.bfloat16, torch.float16])
     @parametrize("scale", [None, 0.1])
     @parametrize(
@@ -365,11 +543,27 @@ def test_custom_op_registration(self, device, dtype):
             (1025, 1025),
         ],
     )
-    def test_varlen_vs_sdpa(self, device, dtype, scale, window_size):
+    @parametrize(
+        "backend",
+        _varlen_backends(include_fa4_paged_kv=False),
+    )
+    @parametrize("enable_gqa", [False, True])
+    def test_varlen_vs_sdpa(
+        self, device, dtype, scale, window_size, backend, enable_gqa, sdpa_backend=None
+    ):
+        if TEST_WITH_ROCM:
+            torch.backends.cuda.preferred_rocm_fa_library(sdpa_backend)
+
         torch.manual_seed(42)
 
+        num_heads = 16
+        num_kv_heads = 4 if enable_gqa else num_heads
         shape = VarlenShape(
-            batch_size=4, max_seq_len=1024, embed_dim=1024, num_heads=16
+            batch_size=4,
+            max_seq_len=1024,
+            embed_dim=1024,
+            num_heads=num_heads,
+            num_kv_heads=num_kv_heads,
         )
 
         batch_data = create_variable_length_batch(shape, device, dtype)
@@ -381,26 +575,38 @@ def test_varlen_vs_sdpa(self, device, dtype, scale, window_size):
         x_padded_ref = batch_data["x_padded_ref"]
 
         golden_attention_block = AttentionBlock(
-            shape.embed_dim, shape.num_heads, device, torch.float32
+            shape.embed_dim,
+            shape.num_heads,
+            device,
+            torch.float32,
+            num_kv_heads=num_kv_heads,
         )
         attention_block = AttentionBlock(
-            shape.embed_dim, shape.num_heads, device, dtype
+            shape.embed_dim,
+            shape.num_heads,
+            device,
+            dtype,
+            num_kv_heads=num_kv_heads,
         )
         with torch.no_grad():
-            attention_block.qkv_proj.weight.copy_(
-                golden_attention_block.qkv_proj.weight.to(dtype)
+            attention_block.q_proj.weight.copy_(
+                golden_attention_block.q_proj.weight.to(dtype)
+            )
+            attention_block.kv_proj.weight.copy_(
+                golden_attention_block.kv_proj.weight.to(dtype)
             )
             attention_block.out_proj.weight.copy_(
                 golden_attention_block.out_proj.weight.to(dtype)
             )
 
-        varlen_output = attention_block.forward_varlen(
-            x_packed,
-            cu_seq,
-            max_len,
-            scale=scale,
-            window_size=window_size,
-        )
+        with _use_backend(backend):
+            varlen_output = attention_block.forward_varlen(
+                x_packed,
+                cu_seq,
+                max_len,
+                scale=scale,
+                window_size=window_size,
+            )
         sdpa_output = attention_block.forward_sdpa(
             x_padded,
             seq_lengths,
@@ -434,29 +640,30 @@ def test_varlen_vs_sdpa(self, device, dtype, scale, window_size):
 
             start_idx = end_idx
 
-        grad_out = torch.randn_like(varlen_output)
-        sdpa_grad_out = torch.zeros_like(sdpa_output)
-        golden_sdpa_grad_out = torch.zeros(
-            shape.batch_size,
-            max_len,
-            shape.embed_dim,
-            device=device,
-            dtype=torch.float32,
-        )
-        start_idx = 0
-        for i, seq_len in enumerate(seq_lengths):
-            end_idx = start_idx + seq_len
-            sdpa_grad_out[i, :seq_len] = grad_out[start_idx:end_idx]
-            golden_sdpa_grad_out[i, :seq_len] = grad_out[start_idx:end_idx].to(
-                torch.float32
+        with _use_backend(backend):
+            grad_out = torch.randn_like(varlen_output)
+            sdpa_grad_out = torch.zeros_like(sdpa_output)
+            golden_sdpa_grad_out = torch.zeros(
+                shape.batch_size,
+                max_len,
+                shape.embed_dim,
+                device=device,
+                dtype=torch.float32,
             )
-            start_idx = end_idx
+            start_idx = 0
+            for i, seq_len in enumerate(seq_lengths):
+                end_idx = start_idx + seq_len
+                sdpa_grad_out[i, :seq_len] = grad_out[start_idx:end_idx]
+                golden_sdpa_grad_out[i, :seq_len] = grad_out[start_idx:end_idx].to(
+                    torch.float32
+                )
+                start_idx = end_idx
 
-        varlen_grad = torch.autograd.grad(
-            outputs=varlen_output,
-            inputs=x_packed,
-            grad_outputs=grad_out,
-        )[0]
+            varlen_grad = torch.autograd.grad(
+                outputs=varlen_output,
+                inputs=x_packed,
+                grad_outputs=grad_out,
+            )[0]
 
         sdpa_grad = torch.autograd.grad(
             outputs=sdpa_output,
@@ -498,142 +705,460 @@ def test_varlen_vs_sdpa(self, device, dtype, scale, window_size):
         not PLATFORM_SUPPORTS_FLASH_ATTENTION, "Flash Attention not supported"
     )
     @parametrize("dtype", [torch.bfloat16, torch.float16])
+    @parametrize("num_splits", [1, None])
     @parametrize(
         "window_size",
         [
             (-1, -1),
             (-1, 0),
-            (4, 0),
-            (-1, 4),
-            (0, 0),
-            (1, 0),
-            (0, 1),
-            (0, -1),
-            (4, 4),
-            (8, 2),
             (1025, 1025),
+            (384, 0),  # edge case
         ],
     )
-    @parametrize("num_perms", [1, 3, 5])
-    def test_batch_invariance(self, device, dtype, window_size, num_perms):
+    @parametrize(
+        "backend",
+        ["fa2"] + (["fa3"] if IS_SM90 else []) + (["fa4"] if SM100OrLater else []),
+    )
+    def test_batch_invariance(
+        self, device, dtype, num_splits, window_size, backend, sdpa_backend=None
+    ):
+        if TEST_WITH_ROCM:
+            torch.backends.cuda.preferred_rocm_fa_library(sdpa_backend)
+
+        split_kwargs = {"num_splits": num_splits} if backend != "fa2" else {}
+
         torch.manual_seed(42)
 
-        batch_size, max_seq_len = 4, 128
+        num_heads, head_dim = 2, 128
+        target_seq_len = 512
+        extra_seq_len = 1024
 
-        seq_lengths = []
-        for _ in range(batch_size):
-            length = torch.randint(1, max_seq_len // 64 + 1, (1,)).item() * 64
-            seq_lengths.append(min(length, max_seq_len))
+        target_q = torch.randn(
+            target_seq_len, num_heads, head_dim, device=device, dtype=dtype
+        )
+        target_k = torch.randn(
+            target_seq_len, num_heads, head_dim, device=device, dtype=dtype
+        )
+        target_v = torch.randn(
+            target_seq_len, num_heads, head_dim, device=device, dtype=dtype
+        )
 
-        sequences_qkv = [
-            [
-                torch.testing.make_tensor(
-                    (seq_len, 2, 128), device=device, dtype=dtype, requires_grad=True
-                )
-                for _ in range(3)
-            ]
-            for seq_len in seq_lengths
-        ]
-        sequences_q, sequences_k, sequences_v = map(list, zip(*sequences_qkv))
-
-        q_packed_orig = torch.cat(sequences_q, dim=0)
-        k_packed_orig = torch.cat(sequences_k, dim=0)
-        v_packed_orig = torch.cat(sequences_v, dim=0)
-
-        seq_lens = torch.tensor(seq_lengths, device=device)
-        cu_seq_orig = torch.zeros(batch_size + 1, device=device, dtype=torch.int32)
-        cu_seq_orig[1:] = seq_lens.cumsum(0)
-
-        original_output = varlen_attn(
-            q_packed_orig,
-            k_packed_orig,
-            v_packed_orig,
-            cu_seq_orig,
-            cu_seq_orig,
-            max_seq_len,
-            max_seq_len,
-            window_size=window_size,
+        extra_q = torch.randn(
+            extra_seq_len, num_heads, head_dim, device=device, dtype=dtype
+        )
+        extra_k = torch.randn(
+            extra_seq_len, num_heads, head_dim, device=device, dtype=dtype
+        )
+        extra_v = torch.randn(
+            extra_seq_len, num_heads, head_dim, device=device, dtype=dtype
         )
 
-        original_grad_out = torch.randn_like(original_output)
-        original_grads = torch.autograd.grad(
-            outputs=original_output,
-            inputs=[q_packed_orig, k_packed_orig, v_packed_orig],
-            grad_outputs=original_grad_out,
+        cu_seq_solo = torch.tensor(
+            [0, target_seq_len], device=device, dtype=torch.int32
+        )
+        cu_seq_batch = torch.tensor(
+            [0, target_seq_len, target_seq_len + extra_seq_len],
+            device=device,
+            dtype=torch.int32,
         )
 
-        for _ in range(num_perms):
-            perm = torch.randperm(batch_size)
-            permuted_sequences_q = [sequences_q[perm[i]] for i in range(batch_size)]
-            permuted_sequences_k = [sequences_k[perm[i]] for i in range(batch_size)]
-            permuted_sequences_v = [sequences_v[perm[i]] for i in range(batch_size)]
+        all_q = torch.cat([target_q, extra_q], dim=0)
+        all_k = torch.cat([target_k, extra_k], dim=0)
+        all_v = torch.cat([target_v, extra_v], dim=0)
+
+        # fa4 is batch invariant (num_splits=1) by default
+        with _use_backend(backend), torch.no_grad():
+            solo_output = varlen_attn(
+                target_q,
+                target_k,
+                target_v,
+                cu_seq_solo,
+                cu_seq_solo,
+                target_seq_len,
+                target_seq_len,
+                window_size=window_size,
+                **split_kwargs,
+            )
 
-            q_packed_perm = torch.cat(permuted_sequences_q, dim=0)
-            k_packed_perm = torch.cat(permuted_sequences_k, dim=0)
-            v_packed_perm = torch.cat(permuted_sequences_v, dim=0)
+            batched_output = varlen_attn(
+                all_q,
+                all_k,
+                all_v,
+                cu_seq_batch,
+                cu_seq_batch,
+                extra_seq_len,
+                extra_seq_len,
+                window_size=window_size,
+                **split_kwargs,
+            )
 
-            permuted_seq_lens = torch.tensor(
-                [seq_lengths[perm[i]] for i in range(batch_size)], device=device
+            solo_out_buf = torch.empty_like(target_q)
+            varlen_attn_out(
+                solo_out_buf,
+                target_q,
+                target_k,
+                target_v,
+                cu_seq_solo,
+                cu_seq_solo,
+                target_seq_len,
+                target_seq_len,
+                window_size=window_size,
+                **split_kwargs,
             )
-            cu_seq_perm = torch.zeros(batch_size + 1, device=device, dtype=torch.int32)
-            cu_seq_perm[1:] = permuted_seq_lens.cumsum(0)
-
-            permuted_output = varlen_attn(
-                q_packed_perm,
-                k_packed_perm,
-                v_packed_perm,
-                cu_seq_perm,
-                cu_seq_perm,
-                max_seq_len,
-                max_seq_len,
+
+            batched_out_buf = torch.empty_like(all_q)
+            varlen_attn_out(
+                batched_out_buf,
+                all_q,
+                all_k,
+                all_v,
+                cu_seq_batch,
+                cu_seq_batch,
+                extra_seq_len,
+                extra_seq_len,
                 window_size=window_size,
+                **split_kwargs,
             )
+            if num_splits == 1:
+                self.assertEqual(solo_output, batched_output[:target_seq_len])
+                self.assertEqual(solo_out_buf, batched_out_buf[:target_seq_len])
+                self.assertEqual(solo_output, solo_out_buf)
+            else:
+                if backend == "fa3":
+                    self.assertNotEqual(solo_output, batched_output[:target_seq_len])
+                    self.assertNotEqual(solo_out_buf, batched_out_buf[:target_seq_len])
 
-            for i in range(batch_size):
-                orig_idx = perm[i].item()
+    @unittest.skipIf(
+        not PLATFORM_SUPPORTS_FLASH_ATTENTION, "Flash Attention not supported"
+    )
+    @decorateIf(
+        unittest.expectedFailure,
+        lambda params: params["backend"] != "fa2"
+        and any(kv_len < 128 for kv_len in params["actual_kv_lens"]),
+    )
+    @parametrize(
+        "sdpa_backend",
+        ["aotriton", "ck"] if PLATFORM_SUPPORTS_CK_SDPA else ["aotriton"],
+    )
+    @parametrize("dtype", [torch.bfloat16, torch.float16])
+    @parametrize(
+        "actual_kv_lens",
+        [
+            [32, 64, 96, 48],
+            [1, 1, 1, 1],
+            [128, 128, 128, 128],
+            [1, 128, 1, 128],
+            [127, 63, 33, 17],
+        ],
+    )
+    @parametrize("backend", _varlen_backends(include_fa4_paged_kv=False))
+    def test_seqused_k_kv_cache(
+        self, device, dtype, actual_kv_lens, backend, sdpa_backend=None
+    ):
+        if TEST_WITH_ROCM:
+            if sdpa_backend == "ck":
+                self.skipTest("CK backend does not support seqused_k")
+            torch.backends.cuda.preferred_rocm_fa_library(sdpa_backend)
 
-                orig_start = cu_seq_orig[orig_idx].item()
-                orig_end = cu_seq_orig[orig_idx + 1].item()
-                orig_seq_output = original_output[orig_start:orig_end]
+        torch.manual_seed(42)
 
-                perm_start = cu_seq_perm[i].item()
-                perm_end = cu_seq_perm[i + 1].item()
-                perm_seq_output = permuted_output[perm_start:perm_end]
+        batch_size = 4
+        num_heads = 8
+        head_dim = 64
+        cache_size = 128
 
-                self.assertEqual(orig_seq_output, perm_seq_output)
+        q_seqs = [
+            torch.randn(1, num_heads, head_dim, device=device, dtype=dtype)
+            for _ in range(batch_size)
+        ]
+        q_packed, cu_seq_q, max_q = pack_sequences(q_seqs, device)
 
-            permuted_grad_out = torch.zeros_like(permuted_output)
-            for i in range(batch_size):
-                orig_idx = perm[i].item()
-                orig_start = cu_seq_orig[orig_idx].item()
-                orig_end = cu_seq_orig[orig_idx + 1].item()
+        k_seqs = [
+            torch.randn(kv_len, num_heads, head_dim, device=device, dtype=dtype)
+            for kv_len in actual_kv_lens
+        ]
+        v_seqs = [
+            torch.randn(kv_len, num_heads, head_dim, device=device, dtype=dtype)
+            for kv_len in actual_kv_lens
+        ]
 
-                perm_start = cu_seq_perm[i].item()
-                perm_end = cu_seq_perm[i + 1].item()
+        k_cache_slots = []
+        v_cache_slots = []
+        for i in range(batch_size):
+            k_slot = torch.full(
+                (cache_size, num_heads, head_dim),
+                float("nan"),
+                device=device,
+                dtype=dtype,
+            )
+            v_slot = torch.full(
+                (cache_size, num_heads, head_dim),
+                float("nan"),
+                device=device,
+                dtype=dtype,
+            )
+            k_slot[: actual_kv_lens[i]] = k_seqs[i]
+            v_slot[: actual_kv_lens[i]] = v_seqs[i]
+            k_cache_slots.append(k_slot)
+            v_cache_slots.append(v_slot)
+
+        k_cache_packed = torch.cat(k_cache_slots, dim=0)
+        v_cache_packed = torch.cat(v_cache_slots, dim=0)
+        cu_seq_k_cache = torch.arange(
+            0,
+            (batch_size + 1) * cache_size,
+            cache_size,
+            device=device,
+            dtype=torch.int32,
+        )
+        seqused_k = torch.tensor(actual_kv_lens, device=device, dtype=torch.int32)
+
+        with _use_backend(backend), torch.no_grad():
+            output_cached = varlen_attn(
+                q_packed,
+                k_cache_packed,
+                v_cache_packed,
+                cu_seq_q,
+                cu_seq_k_cache,
+                max_q,
+                cache_size,
+                seqused_k=seqused_k,
+            )
 
-                permuted_grad_out[perm_start:perm_end] = original_grad_out[
-                    orig_start:orig_end
-                ]
+        k_real_packed, cu_seq_k_real, max_k_real = pack_sequences(k_seqs, device)
+        v_real_packed = torch.cat(v_seqs, dim=0)
+
+        with _use_backend(backend), torch.no_grad():
+            output_reference = varlen_attn(
+                q_packed,
+                k_real_packed,
+                v_real_packed,
+                cu_seq_q,
+                cu_seq_k_real,
+                max_q,
+                max_k_real,
+            )
 
-            permuted_grads = torch.autograd.grad(
-                outputs=permuted_output,
-                inputs=[q_packed_perm, k_packed_perm, v_packed_perm],
-                grad_outputs=permuted_grad_out,
+        self.assertFalse(output_cached.isnan().any())
+        self.assertEqual(output_cached, output_reference)
+
+        # varlen_attn_out with seqused_k should match
+        with _use_backend(backend), torch.no_grad():
+            out_buf = torch.empty_like(q_packed)
+            output_out = varlen_attn_out(
+                out_buf,
+                q_packed,
+                k_cache_packed,
+                v_cache_packed,
+                cu_seq_q,
+                cu_seq_k_cache,
+                max_q,
+                cache_size,
+                seqused_k=seqused_k,
             )
+            self.assertEqual(output_out.data_ptr(), out_buf.data_ptr())
+            self.assertEqual(out_buf, output_cached)
 
-            for original_grad, permuted_grad in zip(original_grads, permuted_grads):
-                for i in range(batch_size):
-                    orig_idx = perm[i].item()
+    @unittest.skipIf(
+        not PLATFORM_SUPPORTS_FLASH_ATTENTION, "Flash Attention not supported"
+    )
+    @unittest.skipIf(TEST_WITH_ROCM, "ROCm does not support block_table")
+    @parametrize("dtype", [torch.bfloat16, torch.float16])
+    @parametrize("page_size", [32, 64, 128, 256])
+    @parametrize("compile", [False, True])
+    @parametrize(
+        "actual_kv_lens",
+        [
+            [32, 64, 96, 48],
+            [1, 1, 1, 1],
+            [128, 128, 128, 128],
+            [1, 128, 1, 128],
+            [127, 63, 33, 17],
+        ],
+    )
+    @parametrize(
+        "backend",
+        _varlen_backends(include_fa4_paged_kv=True),
+    )
+    def test_block_table_kv_cache(
+        self, device, dtype, page_size, compile, actual_kv_lens, backend
+    ):
+        if backend == "fa2" and page_size % 256 != 0:
+            self.skipTest("FA2 paged KV requires page_size divisible by 256")
 
-                    orig_start = cu_seq_orig[orig_idx].item()
-                    orig_end = cu_seq_orig[orig_idx + 1].item()
-                    orig_seq_grad = original_grad[orig_start:orig_end]
+        torch.manual_seed(42)
 
-                    perm_start = cu_seq_perm[i].item()
-                    perm_end = cu_seq_perm[i + 1].item()
-                    perm_seq_grad = permuted_grad[perm_start:perm_end]
+        batch_size = 4
+        num_heads = 8
+        head_dim = 64
+        max_kv = max(actual_kv_lens)
+        max_pages_per_seq = (max_kv + page_size - 1) // page_size
+        cache_size = max_pages_per_seq * page_size
+        total_pages = batch_size * max_pages_per_seq
+
+        q_seqs = [
+            torch.randn(1, num_heads, head_dim, device=device, dtype=dtype)
+            for _ in range(batch_size)
+        ]
+        q_packed, cu_seq_q, max_q = pack_sequences(q_seqs, device)
+
+        k_pages = torch.randn(
+            total_pages, page_size, num_heads, head_dim, device=device, dtype=dtype
+        )
+        v_pages = torch.randn(
+            total_pages, page_size, num_heads, head_dim, device=device, dtype=dtype
+        )
+        block_table = torch.randperm(
+            total_pages, device=device, dtype=torch.int32
+        ).view(batch_size, max_pages_per_seq)
+        seqused_k = torch.tensor(actual_kv_lens, device=device, dtype=torch.int32)
+
+        idx = (
+            block_table.long()
+            .view(-1, 1, 1, 1)
+            .expand(-1, page_size, num_heads, head_dim)
+        )
+        k_gathered = k_pages.gather(0, idx).view(
+            batch_size, cache_size, num_heads, head_dim
+        )
+        v_gathered = v_pages.gather(0, idx).view(
+            batch_size, cache_size, num_heads, head_dim
+        )
+        k_seqs = [k_gathered[i, : actual_kv_lens[i]] for i in range(batch_size)]
+        v_seqs = [v_gathered[i, : actual_kv_lens[i]] for i in range(batch_size)]
+
+        k_real_packed, cu_seq_k_real, max_k_real = pack_sequences(k_seqs, device)
+        v_real_packed = torch.cat(v_seqs, dim=0)
+
+        attn_fn = torch.compile(varlen_attn, fullgraph=True) if compile else varlen_attn
+
+        # Reference: no block_table
+        with _use_backend(backend), torch.no_grad():
+            output_reference = varlen_attn(
+                q_packed,
+                k_real_packed,
+                v_real_packed,
+                cu_seq_q,
+                cu_seq_k_real,
+                max_q,
+                max_k_real,
+            )
+
+        cu_seq_k = torch.arange(
+            0,
+            (batch_size + 1) * cache_size,
+            cache_size,
+            device=device,
+            dtype=torch.int32,
+        )
+
+        # FA2 requires cu_seq_k for paged KV; FA3/FA4 pass None
+        cu_seq_k_paged = cu_seq_k if backend == "fa2" else None
+
+        with _use_backend(backend), torch.no_grad():
+            output_paged = attn_fn(
+                q_packed,
+                k_pages,
+                v_pages,
+                cu_seq_q,
+                cu_seq_k_paged,
+                max_q,
+                cache_size,
+                seqused_k=seqused_k,
+                block_table=block_table,
+            )
 
-                    self.assertEqual(orig_seq_grad, perm_seq_grad)
+        self.assertEqual(output_paged, output_reference)
+
+        # varlen_attn_out with paged KV cache should match
+        with _use_backend(backend), torch.no_grad():
+            out_buf = torch.empty_like(q_packed)
+            output_out = varlen_attn_out(
+                out_buf,
+                q_packed,
+                k_pages,
+                v_pages,
+                cu_seq_q,
+                cu_seq_k_paged,
+                max_q,
+                cache_size,
+                seqused_k=seqused_k,
+                block_table=block_table,
+            )
+            self.assertEqual(output_out.data_ptr(), out_buf.data_ptr())
+            self.assertEqual(out_buf, output_paged)
+
+        # compile the lower level aten op (FA3 only, will cause graph break)
+        if compile and backend != "fa2":
+            compiled_aten_op = torch.compile(
+                torch.ops.aten._flash_attention_forward_no_dropout_inplace
+            )
+            with _use_backend(backend), torch.no_grad():
+                out_buf = torch.empty_like(q_packed)
+                compiled_aten_op(
+                    out_buf,
+                    q_packed,
+                    k_pages,
+                    v_pages,
+                    cu_seq_q,
+                    None,
+                    max_q,
+                    cache_size,
+                    0.0,
+                    False,
+                    False,
+                    seqused_k=seqused_k,
+                    block_table=block_table,
+                )
+            self.assertEqual(out_buf, output_reference)
+
+    @unittest.skipIf(
+        not PLATFORM_SUPPORTS_FLASH_ATTENTION, "Flash Attention not supported"
+    )
+    @parametrize("dtype", [torch.bfloat16, torch.float16])
+    @parametrize(
+        "backend",
+        ["fa2"] + (["fa3"] if IS_SM90 else []) + (["fa4"] if SM100OrLater else []),
+    )
+    def test_enable_gqa(self, device, dtype, backend):
+        torch.manual_seed(42)
+
+        head_dim = 64
+        seq_len = 512
+        num_heads_q, num_heads_k = 16, 4
+        total_tokens = 2 * seq_len
+
+        q = torch.randn(total_tokens, num_heads_q, head_dim, device=device, dtype=dtype)
+        k = torch.randn(total_tokens, num_heads_k, head_dim, device=device, dtype=dtype)
+        v = torch.randn(total_tokens, num_heads_k, head_dim, device=device, dtype=dtype)
+        cu_seq = torch.tensor(
+            [0, seq_len, total_tokens], device=device, dtype=torch.int32
+        )
+
+        with self.assertRaisesRegex(ValueError, "enable_gqa=True"):
+            varlen_attn(q, k, v, cu_seq, cu_seq, seq_len, seq_len)
+
+        with self.assertRaisesRegex(ValueError, "enable_gqa=True"):
+            varlen_attn_out(
+                torch.empty_like(q), q, k, v, cu_seq, cu_seq, seq_len, seq_len
+            )
+
+        k_bad = torch.randn(total_tokens, 3, head_dim, device=device, dtype=dtype)
+        v_bad = torch.randn(total_tokens, 3, head_dim, device=device, dtype=dtype)
+        with self.assertRaisesRegex(ValueError, "multiple of kv heads"):
+            varlen_attn(
+                q, k_bad, v_bad, cu_seq, cu_seq, seq_len, seq_len, enable_gqa=True
+            )
+
+        with _use_backend(backend), torch.no_grad():
+            out = varlen_attn(
+                q, k, v, cu_seq, cu_seq, seq_len, seq_len, enable_gqa=True
+            )
+            out_buf = torch.empty_like(q)
+            varlen_attn_out(
+                out_buf, q, k, v, cu_seq, cu_seq, seq_len, seq_len, enable_gqa=True
+            )
+            self.assertEqual(out_buf, out)
 
 
 device_types = ("cuda",)
diff --git a/test/test_xpu.py b/test/test_xpu.py
index 9d54f838b0aa0..6d9f772655e49 100644
--- a/test/test_xpu.py
+++ b/test/test_xpu.py
@@ -23,6 +23,7 @@
 from torch.testing import make_tensor
 from torch.testing._internal.autocast_test_lists import AutocastTestLists, TestAutocast
 from torch.testing._internal.common_device_type import (
+    dtypes,
     instantiate_device_type_tests,
     OpDTypes,
     ops,
@@ -126,6 +127,8 @@ def test_get_device_properties(self):
         self.assertTrue(device_capability["max_work_group_size"] > 0)
         self.assertTrue(device_capability["max_num_sub_groups"] > 0)
         self.assertTrue(device_capability["local_mem_size"] > 0)
+        self.assertTrue(device_capability["memory_clock_rate"] > 0)
+        self.assertTrue(device_capability["memory_bus_width"] > 0)
         self.assertEqual(
             device_properties.driver_version, device_capability["driver_version"]
         )
@@ -244,6 +247,49 @@ def test_multi_process(model, input):
         rc = check_output(test_script).splitlines()[-1]
         self.assertEqual(rc, str(torch.xpu.device_count()))
 
+    @unittest.skipIf(
+        IS_WINDOWS, "Only for lazy initialization on Linux, not applicable on Windows."
+    )
+    def test_no_xpu_device_query_on_inductor_import(self):
+        """Validate that importing torch._inductor.lowering does not trigger XPU device queries"""
+
+        def check_output(script: str) -> str:
+            return (
+                subprocess.check_output([sys.executable, "-c", script])
+                .decode("ascii")
+                .strip()
+            )
+
+        test_script = """\
+import torch
+from unittest.mock import patch
+
+call_count = 0
+original_getDeviceCount = torch._C._xpu_getDeviceCount
+
+def counting_getDeviceCount():
+    global call_count
+    call_count += 1
+    return original_getDeviceCount()
+
+with patch.object(torch._C, '_xpu_getDeviceCount', counting_getDeviceCount):
+    import torch._inductor.lowering
+
+print(call_count)
+print(torch.xpu.is_initialized())
+"""
+        rc = check_output(test_script).splitlines()
+        self.assertEqual(
+            rc[0],
+            "0",
+            "Importing torch._inductor.lowering should not query XPU device count",
+        )
+        self.assertEqual(
+            rc[1],
+            "False",
+            "Importing torch._inductor.lowering should not initialize XPU",
+        )
+
     def test_streams(self):
         s0 = torch.xpu.Stream()
         torch.xpu.set_stream(s0)
@@ -1220,8 +1266,10 @@ def test_graph_is_current_stream_capturing(self):
         with torch.xpu.stream(s):
             g = torch.xpu.XPUGraph()
             self.assertFalse(torch.xpu.is_current_stream_capturing())
+            self.assertFalse(s.is_capturing())
             g.capture_begin()
             self.assertTrue(torch.xpu.is_current_stream_capturing())
+            self.assertTrue(s.is_capturing())
             g.capture_end()
 
     def test_graph_capture_simple(self):
@@ -1242,6 +1290,21 @@ def test_graph_capture_simple(self):
 
         self.assertEqual(b.sum().item(), 11000.0)
 
+    def test_accelerator_graph_simple(self):
+        s = torch.Stream()
+        g = torch.accelerator.Graph()
+
+        with s, g:
+            a = torch.full((1000,), 1, device="xpu")
+            b = a
+            for _ in range(10):
+                b = b + 1
+        torch.accelerator.current_stream().wait_stream(s)
+
+        g.replay()
+
+        self.assertEqual(b.sum().item(), 11000.0)
+
     def test_graphsafe_set_get_rng_state(self):
         # Define a function to create generator states, with optional graph registration
         def create_states(generator):
@@ -2727,6 +2790,78 @@ def convert_boolean_tensors(x):
 
             self.assertEqual(expect, actual)
 
+    @dtypes(torch.float16, torch.bfloat16, torch.float32)
+    def test_fused_rms_norm(self, device, dtype):
+        # Verify _fused_rms_norm is dispatched to XPU kernel (not fallback)
+        has_xpu_kernel = torch._C._dispatch_has_kernel_for_dispatch_key(
+            "aten::_fused_rms_norm",
+            torch._C._dispatch_key_name(torch._C.DispatchKey.XPU),
+        )
+        self.assertTrue(has_xpu_kernel, "_fused_rms_norm XPU kernel is not registered")
+        has_xpu_kernel = torch._C._dispatch_has_kernel_for_dispatch_key(
+            "aten::_fused_rms_norm_backward",
+            torch._C._dispatch_key_name(torch._C.DispatchKey.XPU),
+        )
+        self.assertTrue(
+            has_xpu_kernel, "_fused_rms_norm_backward XPU kernel is not registered"
+        )
+
+        shapes = [
+            (2, 16),  # small 2D
+            (4, 8, 32),  # 3D
+            (1, 1, 64),  # degenerate batch
+            (8, 128),  # typical sequence hidden
+            (2, 16, 512),  # typical LLM hidden dim
+            (4, 32, 1024),  # larger hidden dim
+            (1, 1, 4096),  # LLM-scale hidden
+            (3, 7, 17),  # non-power-of-2
+        ]
+        eps = 1e-5
+        atol_fwd = 1e-1 if dtype in [torch.float16, torch.bfloat16] else 1e-5
+        atol_bwd = 1e-1 if dtype in [torch.float16, torch.bfloat16] else 1e-5
+
+        for shape in shapes:
+            normalized_shape = list(shape[-1:])
+            x = torch.randn(*shape, dtype=dtype, device=device, requires_grad=True)
+            w = torch.randn(
+                *normalized_shape, dtype=dtype, device=device, requires_grad=True
+            )
+            grad_out = torch.randn(*shape, dtype=dtype, device=device)
+            x_cpu = x.detach().cpu().requires_grad_(True)
+            w_cpu = w.detach().cpu().requires_grad_(True)
+            grad_out_cpu = grad_out.detach().cpu()
+
+            # Forward
+            y, _ = torch.ops.aten._fused_rms_norm(x, normalized_shape, w, eps)
+            y_cpu, _ = torch.ops.aten._fused_rms_norm(
+                x_cpu, normalized_shape, w_cpu, eps
+            )
+            self.assertEqual(
+                y,
+                y_cpu,
+                atol=atol_fwd,
+                rtol=0,
+                msg=f"forward shape={shape}, dtype={dtype}",
+            )
+
+            # Backward
+            y.backward(grad_out)
+            y_cpu.backward(grad_out_cpu)
+            self.assertEqual(
+                x.grad.cpu(),
+                x_cpu.grad,
+                atol=atol_bwd,
+                rtol=0,
+                msg=f"x_grad shape={shape}, dtype={dtype}",
+            )
+            self.assertEqual(
+                w.grad.cpu(),
+                w_cpu.grad,
+                atol=atol_bwd,
+                rtol=0,
+                msg=f"w_grad shape={shape}, dtype={dtype}",
+            )
+
 
 instantiate_device_type_tests(TestXpuOps, globals(), only_for="xpu", allow_xpu=True)
 
@@ -2901,6 +3036,8 @@ def test_torch_config_for_xpu(self):
                 self.assertTrue(value.group(1) in ["ON", "1"])
             else:
                 self.assertTrue(value.group(1) in ["OFF", "0"])
+            value = re.search(r"SYCL_COMPILER_VERSION=([^,]+)", config)
+            self.assertEqual(value.group(1), torch.version.xpu)
         else:
             self.assertTrue(value.group(1) in ["OFF", "0"])
             self.assertFalse(torch.distributed.is_xccl_available())
diff --git a/test/torch_np/numpy_tests/core/test_multiarray.py b/test/torch_np/numpy_tests/core/test_multiarray.py
index 2f1fe86c2905d..c05e60ecf15d7 100644
--- a/test/torch_np/numpy_tests/core/test_multiarray.py
+++ b/test/torch_np/numpy_tests/core/test_multiarray.py
@@ -4443,6 +4443,10 @@ def test_index_getset(self):
         if it.index != it.base.size:
             raise AssertionError(f"index mismatch: {it.index} != {it.base.size}")
 
+    def test_flat_cumsum(self):
+        x = np.array([[1.0, 2.0], [3.0, 4.0]])
+        assert_array_equal(np.cumsum(x.flat), np.array([1.0, 3.0, 6.0, 10.0]))
+
 
 class TestResize(TestCase):
     @_no_tracing
@@ -6385,7 +6389,7 @@ def test_scalar_interface(self, val, iface, expected):
 
 
 class TestDelMisc(TestCase):
-    @xpassIfTorchDynamo_np  # (reason="TODO")
+    @xfail  # torch._numpy .flat returns ravel() instead of flatiter, so del is not supported
     def test_flat_element_deletion(self):
         it = np.ones(3).flat
         try:
@@ -6834,7 +6838,7 @@ def test_choose_mod_raise(self):
         np.choose(a, choices, out=out, mode="raise")
         assert_equal(out, np.array([[10, -10, 10], [-10, 10, -10], [10, -10, 10]]))
 
-    @xpassIfTorchDynamo_np  # (reason="XXX: ndarray.flat")
+    @xfail  # torch._numpy ndarray doesn't implement __array__
     def test_flatiter__array__(self):
         a = np.arange(9).reshape(3, 3)
         b = a.T.flat
diff --git a/test/torch_np/test_nep50_examples.py b/test/torch_np/test_nep50_examples.py
index a3ad346bf9f1c..964683bd74c81 100644
--- a/test/torch_np/test_nep50_examples.py
+++ b/test/torch_np/test_nep50_examples.py
@@ -31,6 +31,7 @@
 from torch._numpy.testing import assert_allclose
 from torch.testing._internal.common_utils import (
     instantiate_parametrized_tests,
+    IS_WINDOWS,
     parametrize,
     run_tests,
     TestCase,
@@ -216,11 +217,36 @@ def test_compare_ufuncs(self, name, scalar, array):
                 # TypeError: ufunc 'hypot' not supported for the input types
                 result_numpy = None
 
+            type_mismatch = False
+            expected_numpy_dtype = None
+            expected_torch_dtype = None
+
             if result is not None and result_numpy is not None:
-                if result.tensor.numpy().dtype != result_numpy.dtype:
-                    raise AssertionError(
-                        f"Expected result dtype == {result_numpy.dtype}, got {result.tensor.numpy().dtype}"
-                    )
+                expected_numpy_dtype = result_numpy.dtype
+                expected_torch_dtype = result.tensor.numpy().dtype
+                if IS_WINDOWS:
+                    if (
+                        array.tensor.numpy().dtype != _np.bool_
+                        and result.tensor.numpy().dtype != result_numpy.dtype
+                    ):
+                        type_mismatch = True
+
+                    if (
+                        array.tensor.numpy().dtype == _np.bool_
+                        and result_numpy.dtype == _np.int32
+                        and result.tensor.numpy().dtype != _np.int64
+                    ):
+                        expected_numpy_dtype = _np.int32
+                        expected_torch_dtype = tnp.int64
+                        type_mismatch = True
+                else:
+                    if result.tensor.numpy().dtype != result_numpy.dtype:
+                        type_mismatch = True
+
+            if type_mismatch:
+                raise AssertionError(
+                    f"Expected result numpy dtype == {expected_numpy_dtype}, torch dtype == {expected_torch_dtype}"
+                )
 
         finally:
             _np._set_promotion_state(state)
diff --git a/test/typing/pass/arithmetic_ops.py b/test/typing/pass/arithmetic_ops.py
index 14dda1cf39772..c2eb3c589526c 100644
--- a/test/typing/pass/arithmetic_ops.py
+++ b/test/typing/pass/arithmetic_ops.py
@@ -1,4 +1,4 @@
-from typing import TypeAlias, Union
+from typing import TypeAlias
 from typing_extensions import assert_type
 
 from torch import randn, Tensor
@@ -206,7 +206,7 @@
 assert_type(FLOAT ^ TENSOR, Tensor)  # type: ignore[operator]
 
 
-NUMBER: TypeAlias = Union[int, float, bool]
+NUMBER: TypeAlias = int | float | bool
 
 
 class Binary:
diff --git a/test/xpu/test_conv.py b/test/xpu/test_conv.py
index d25905f5804d1..3e069f2d92637 100644
--- a/test/xpu/test_conv.py
+++ b/test/xpu/test_conv.py
@@ -1357,6 +1357,39 @@ def test_channels_last_ouput_stride(self, device, dtype):
         # input NHWC, output NHWC
         assert_size_stride(out, (2, 512, 7, 7), (25088, 1, 3584, 512))
 
+    @dtypes(torch.float16, torch.bfloat16, torch.float32, torch.float64)
+    def test_conv_misaligned_input(self, device, dtype):
+        N, C, OUT_CHANNELS = 2, 3, 4
+
+        def make_misaligned_tensor(tensor, shape):
+            numel = math.prod(shape)
+            base = torch.empty(numel + 1, device=device, dtype=dtype)
+            tensor_device = base[1:].reshape(shape)
+            tensor_device.copy_(tensor.to(device))
+            self.assertTrue(tensor_device.data_ptr() % 64 != 0)
+            return tensor_device
+
+        for conv_fn, spatial in [
+            (F.conv1d, (16,)),
+            (F.conv2d, (8, 8)),
+            (F.conv3d, (4, 4, 4)),
+        ]:
+            kernel = (3,) * len(spatial)
+            input_shape = (N, C) + spatial
+            weight_shape = (OUT_CHANNELS, C) + kernel
+
+            input_cpu = torch.randn(input_shape, dtype=dtype)
+            weight_cpu = torch.randn(weight_shape, dtype=dtype)
+            bias_cpu = torch.randn(OUT_CHANNELS, dtype=dtype)
+
+            input_device = make_misaligned_tensor(input_cpu, input_shape)
+            weight_device = make_misaligned_tensor(weight_cpu, weight_shape)
+            bias_device = make_misaligned_tensor(bias_cpu, (OUT_CHANNELS,))
+
+            output_cpu = conv_fn(input_cpu, weight_cpu, bias_cpu)
+            output_device = conv_fn(input_device, weight_device, bias_device)
+            self.assertEqual(output_device.cpu(), output_cpu)
+
     @onlyXPU
     def test_onednn_allow_tf32_get_set(self):
         with torch.backends.mkldnn.flags(
diff --git a/test/xpu/test_gemm.py b/test/xpu/test_gemm.py
index 117d0133fc6df..4ceb858a4712c 100644
--- a/test/xpu/test_gemm.py
+++ b/test/xpu/test_gemm.py
@@ -26,6 +26,7 @@
     _dynamically_quantize_per_channel,
 )
 from torch.testing._internal.common_utils import (
+    DeterministicGuard,
     iter_indices,
     parametrize,
     run_tests,
@@ -175,6 +176,12 @@ def _test_addmm_addmv(
         self.assertEqual(res1, res2)
         self.assertEqual(res1, res3)
 
+        # Test inplace versions if they exist.
+        if hasattr(t, f.__name__ + "_"):
+            out_tensor = torch.broadcast_to(t, res1.shape).clone()
+            getattr(out_tensor, f.__name__ + "_")(m, v, alpha=alpha, beta=beta)
+            self.assertEqual(res1, out_tensor)
+
     def _test_addmm_impl(self, func, activation, device, dtype):
         M = torch.randn(10, 25, device="cpu", dtype=torch.float32).to(dtype).to(device)
         m1 = torch.randn(10, 50, device="cpu", dtype=torch.float32).to(dtype).to(device)
@@ -238,13 +245,13 @@ def maybe_transpose(cond, m):
                     activation=activation,
                 )
 
-    @precisionOverride({torch.float: 1e-4, torch.double: 1e-6, torch.half: 1e-1})
+    @precisionOverride({torch.float: 1e-4, torch.half: 1e-1})
     @dtypes(torch.float32, torch.half, torch.double, torch.complex64)
     @tf32_on_and_off(0.05)
     def test_addmm(self, device, dtype):
         self._test_addmm_impl(torch.addmm, None, device, dtype)
 
-    @precisionOverride({torch.float: 1e-4, torch.double: 1e-6, torch.half: 1e-1})
+    @precisionOverride({torch.float: 1e-4, torch.half: 1e-1})
     @dtypes(torch.float, torch.half, torch.double)
     def test_addmm_badmm_scalar_tnesor_input(self, device, dtype):
         input = torch.tensor(1).to(device=device, dtype=dtype)
@@ -638,7 +645,7 @@ def generate_tensor():
         for b1, b2, ref, out_tensor in generate_tensor():
             self._test_addbmm_baddbmm("addbmm", b1, b2, ref, out_tensor)
 
-    @precisionOverride({torch.half: 0.1, torch.bfloat16: 0.5, torch.float64: 1e-6})
+    @precisionOverride({torch.half: 0.1, torch.bfloat16: 0.5})
     @dtypes(torch.float64, torch.float32, torch.bfloat16, torch.half, torch.complex64)
     @tf32_on_and_off(0.01)
     def test_baddbmm(self, device, dtype):
@@ -864,6 +871,15 @@ def test_matmul_45724(self, device):
         torch.matmul(a, b, out=c)
         self.assertEqual(c, cpu_result)
 
+    @parametrize("shape", [513, 767])
+    @dtypes(torch.bfloat16, torch.half, torch.float, torch.double)
+    def test_matmul_deterministic_mode(self, device, shape, dtype):
+        with DeterministicGuard(True):
+            inp = torch.randn(shape, shape, device=device, dtype=dtype)
+            first = torch.matmul(inp, inp)
+            for _ in range(10):
+                self.assertEqual(first, torch.matmul(inp, inp), atol=0.0, rtol=0.0)
+
     @dtypes(
         torch.int16,
         torch.int32,
@@ -908,7 +924,6 @@ def test_baddbmm_nan_input_with_zero_beta(self, device, dtype):
                 y = torch.baddbmm(input, mat1, mat2, beta=0.0, out=out)
                 self.assertEqual(y_ref, y)
 
-    @precisionOverride({torch.double: 1e-6})
     @dtypes(torch.float, torch.double)
     @tf32_on_and_off(0.005)
     def test_addmm_sizes(self, device, dtype):
@@ -933,7 +948,6 @@ def test_addmm_sizes(self, device, dtype):
 
     @precisionOverride(
         {
-            torch.double: 1e-6,
             torch.float: 1e-4,
             torch.bfloat16: 5e-2,
             torch.half: 5e-2,
@@ -948,7 +962,6 @@ def test_addmm_gelu(self, device, dtype):
 
     @precisionOverride(
         {
-            torch.double: 1e-6,
             torch.float: 1e-4,
             torch.bfloat16: 5e-2,
             torch.half: 5e-2,
@@ -997,7 +1010,6 @@ def _test(row_major, incx, incy, lda_tail):
 
     @precisionOverride(
         {
-            torch.double: 1e-8,
             torch.float: 1e-4,
             torch.bfloat16: 0.6,
             torch.half: 1e-1,
diff --git a/third_party/BUILD b/third_party/BUILD
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/third_party/composable_kernel b/third_party/composable_kernel
index fcc9372c009c8..cd9e41c4eae94 160000
--- a/third_party/composable_kernel
+++ b/third_party/composable_kernel
@@ -1 +1 @@
-Subproject commit fcc9372c009c8e0a23fece77b582da83b04a654f
+Subproject commit cd9e41c4eae94f11e7d6e39751e8e13f46653030
diff --git a/third_party/cpp-httplib b/third_party/cpp-httplib
index bd95e67c23493..4d7c9a788de13 160000
--- a/third_party/cpp-httplib
+++ b/third_party/cpp-httplib
@@ -1 +1 @@
-Subproject commit bd95e67c234930cd6d6bb11309588c5462c63cec
+Subproject commit 4d7c9a788de136071ccf0dd4e96239151e2adadb
diff --git a/third_party/cpp-httplib.BUILD b/third_party/cpp-httplib.BUILD
deleted file mode 100644
index 3cd0c3dbe94ba..0000000000000
--- a/third_party/cpp-httplib.BUILD
+++ /dev/null
@@ -1,10 +0,0 @@
-load("@rules_cc//cc:defs.bzl", "cc_library")
-
-cc_library(
-    name = "cpp-httplib",
-    hdrs = ["httplib.h"],
-    includes = [
-        "/",
-    ],
-    visibility = ["//visibility:public"],
-)
diff --git a/third_party/cuda.BUILD b/third_party/cuda.BUILD
deleted file mode 100644
index 4767231b55830..0000000000000
--- a/third_party/cuda.BUILD
+++ /dev/null
@@ -1,82 +0,0 @@
-# Adopted from: https://github.com/tensorflow/runtime/blob/master/third_party/rules_cuda/private/BUILD.local_cuda
-# Library targets are created corresponding to BUILD.bazel's needs.
-
-cc_library(
-    name = "cuda_headers",
-    hdrs = glob([
-        "include/**",
-        "targets/x86_64-linux/include/**",
-    ]),
-    includes = [
-        "include",
-        "targets/x86_64-linux/include",
-    ],
-    visibility = ["//visibility:public"],
-)
-
-cc_library(
-    name = "cuda_driver",
-    srcs = ["lib64/stubs/libcuda.so"],
-    visibility = ["//visibility:public"],
-)
-
-cc_library(
-    name = "cuda",
-    srcs = ["targets/x86_64-linux/lib/libcudart.so"],
-    visibility = ["//visibility:public"],
-    deps = [":cuda_headers"],
-)
-
-cc_library(
-    name = "cufft",
-    srcs = ["targets/x86_64-linux/lib/libcufft.so"],
-    visibility = ["//visibility:public"],
-)
-
-cc_library(
-    name = "cublas",
-    srcs = [
-        "targets/x86_64-linux/lib/libcublasLt.so",
-        "targets/x86_64-linux/lib/libcublas.so",
-    ],
-    visibility = ["//visibility:public"],
-)
-
-cc_library(
-    name = "curand",
-    srcs = ["targets/x86_64-linux/lib/libcurand.so"],
-    visibility = ["//visibility:public"],
-)
-
-cc_library(
-    name = "cusolver",
-    srcs = ["targets/x86_64-linux/lib/libcusolver.so"],
-    visibility = ["//visibility:public"],
-)
-
-cc_library(
-    name = "cusparse",
-    srcs = ["targets/x86_64-linux/lib/libcusparse.so"],
-    visibility = ["//visibility:public"],
-)
-
-cc_library(
-    name = "cufile",
-    srcs = ["targets/x86_64-linux/lib/libcufile.so"],
-    visibility = ["//visibility:public"],
-)
-
-cc_library(
-    name = "nvrtc",
-    srcs = [
-        "targets/x86_64-linux/lib/libnvrtc.so",
-        "targets/x86_64-linux/lib/libnvrtc-builtins.so",
-    ],
-    visibility = ["//visibility:public"],
-)
-
-cc_library(
-    name = "nvToolsExt",
-    srcs = [ "lib64/libnvToolsExt.so"],
-    visibility = ["//visibility:public"],
-)
diff --git a/third_party/cudnn.BUILD b/third_party/cudnn.BUILD
deleted file mode 100644
index 82b4526c21e6c..0000000000000
--- a/third_party/cudnn.BUILD
+++ /dev/null
@@ -1,26 +0,0 @@
-# Adopted from: https://github.com/NVIDIA/TRTorch/blob/master/third_party/cudnn/local/BUILD
-
-cc_library(
-    name = "cudnn_headers",
-    hdrs = ["include/cudnn.h"] + glob([
-        "include/cudnn+.h",
-        "include/cudnn_*.h",
-    ]),
-    includes = ["include/"],
-    visibility = ["//visibility:private"],
-)
-
-cc_import(
-    name = "cudnn_lib",
-    shared_library = "lib64/libcudnn.so",
-    visibility = ["//visibility:private"],
-)
-
-cc_library(
-    name = "cudnn",
-    visibility = ["//visibility:public"],
-    deps = [
-        "cudnn_headers",
-        "cudnn_lib",
-    ],
-)
diff --git a/third_party/cudnn_frontend b/third_party/cudnn_frontend
index b8c0656e6f6c8..a91f0e04dcea1 160000
--- a/third_party/cudnn_frontend
+++ b/third_party/cudnn_frontend
@@ -1 +1 @@
-Subproject commit b8c0656e6f6c84fc194f4d57329b55d609eff596
+Subproject commit a91f0e04dcea10515f0f776fc5a89535e316a9c8
diff --git a/third_party/cudnn_frontend.BUILD b/third_party/cudnn_frontend.BUILD
deleted file mode 100644
index 0af69797a3e2d..0000000000000
--- a/third_party/cudnn_frontend.BUILD
+++ /dev/null
@@ -1,22 +0,0 @@
-# Adopted from: https://github.com/tensorflow/tensorflow/blob/master/third_party/cudnn_frontend.BUILD
-
-# Description:
-# The cuDNN Frontend API is a C++ header-only library that demonstrates how
-# to use the cuDNN C backend API.
-
-load("@rules_cc//cc:defs.bzl", "cc_library")
-
-package(
-    default_visibility = ["//visibility:public"],
-)
-
-licenses(["notice"])  # MIT
-
-exports_files(["LICENSE.txt"])
-
-cc_library(
-    name = "cudnn_frontend",
-    hdrs = glob(["include/**"]),
-    includes = ["include/"],
-    include_prefix = "third_party/cudnn_frontend",
-)
diff --git a/third_party/cutlass b/third_party/cutlass
index 0d2b201e8c1c4..da5e086dab31d 160000
--- a/third_party/cutlass
+++ b/third_party/cutlass
@@ -1 +1 @@
-Subproject commit 0d2b201e8c1c4a03efa6e9c468161916e2334725
+Subproject commit da5e086dab31d63815acafdac9a9c5893b1c69e2
diff --git a/third_party/cutlass.BUILD b/third_party/cutlass.BUILD
deleted file mode 100644
index 10100531d9be6..0000000000000
--- a/third_party/cutlass.BUILD
+++ /dev/null
@@ -1,26 +0,0 @@
-# Description:
-#   CUDA Templates for Linear Algebra Subroutines
-
-load("@rules_cc//cc:defs.bzl", "cc_library")
-
-cc_library(
-    name = "cutlass",
-    hdrs = glob([
-        "include/**/*.h",
-        "include/**/*.hpp",
-        "include/**/*.inl",
-        "tools/util/include/**/*.h",
-        "tools/util/include/**/*.hpp",
-        "tools/util/include/**/*.inl",
-    ]),
-    defines = [
-        "CUTLASS_ENABLE_TENSOR_CORE_MMA=1",
-        "CUTLASS_ENABLE_SM90_EXTENDED_MMA_SHAPES=1",
-        "CUTE_SM90_EXTENDED_MMA_SHAPES_ENABLED",
-    ],
-    includes = [
-        "include/",
-        "tools/util/include/",
-    ],
-    visibility = ["//visibility:public"],
-)
diff --git a/third_party/eigen.BUILD b/third_party/eigen.BUILD
deleted file mode 100644
index a6a735360633b..0000000000000
--- a/third_party/eigen.BUILD
+++ /dev/null
@@ -1,91 +0,0 @@
-# This is BUILD file is derived from https://github.com/tensorflow/tensorflow/blob/master/third_party/eigen.BUILD
-
-# Description:
-#   Eigen is a C++ template library for linear algebra: vectors,
-#   matrices, and related algorithms.
-
-load("@rules_cc//cc:defs.bzl", "cc_library")
-
-licenses([
-    # Note: Eigen is an MPL2 library that includes GPL v3 and LGPL v2.1+ code.
-    #       We've taken special care to not reference any restricted code.
-    "reciprocal",  # MPL2
-    "notice",  # Portions BSD
-])
-
-exports_files(["COPYING.MPL2"])
-
-# License-restricted (i.e. not reciprocal or notice) files inside Eigen/...
-EIGEN_RESTRICTED_FILES = [
-    "Eigen/src/OrderingMethods/Amd.h",
-    "Eigen/src/SparseCholesky/**",
-]
-
-# Notable transitive dependencies of restricted files inside Eigen/...
-EIGEN_RESTRICTED_DEPS = [
-    "Eigen/Eigen",
-    "Eigen/IterativeLinearSolvers",
-    "Eigen/MetisSupport",
-    "Eigen/Sparse",
-    "Eigen/SparseCholesky",
-    "Eigen/SparseLU",
-]
-
-EIGEN_FILES = [
-    "Eigen/**",
-    "unsupported/Eigen/CXX11/**",
-    "unsupported/Eigen/FFT",
-    "unsupported/Eigen/KroneckerProduct",
-    "unsupported/Eigen/src/FFT/**",
-    "unsupported/Eigen/src/KroneckerProduct/**",
-    "unsupported/Eigen/MatrixFunctions",
-    "unsupported/Eigen/SpecialFunctions",
-    "unsupported/Eigen/Splines",
-    "unsupported/Eigen/src/MatrixFunctions/**",
-    "unsupported/Eigen/src/SpecialFunctions/**",
-    "unsupported/Eigen/src/Splines/**",
-    "unsupported/Eigen/NonLinearOptimization",
-    "unsupported/Eigen/NumericalDiff",
-    "unsupported/Eigen/src/**",
-    "unsupported/Eigen/Polynomials",
-]
-
-# List of files picked up by glob but actually part of another target.
-EIGEN_EXCLUDE_FILES = ["Eigen/src/Core/arch/AVX/PacketMathGoogleTest.cc"]
-
-# Disallowed eigen modules/files in rNA:
-# * Using the custom STL and memory support, it is not needed and should
-#   not be used with c++17.
-# * We will only support the EulerAnglesZYX provided by //atg/geometry so
-#   just don't allow people to access the unsupported eigen module.
-EIGEN_DISALLOW_FILES = [
-    "Eigen/StlSupport/*.h",
-    "unsupported/Eigen/EulerAngles",
-    "unsupported/Eigen/src/EulerAngles/**",
-]
-
-# Files known to be under MPL2 license.
-EIGEN_MPL2_HEADER_FILES = glob(
-    EIGEN_FILES,
-    exclude = EIGEN_EXCLUDE_FILES +
-              EIGEN_RESTRICTED_FILES +
-              EIGEN_DISALLOW_FILES +
-              EIGEN_RESTRICTED_DEPS + [
-        # Guarantees any file missed by excludes above will not compile.
-        "Eigen/src/Core/util/NonMPL2.h",
-        "Eigen/**/CMakeLists.txt",
-    ],
-)
-
-cc_library(
-    name = "eigen",
-    hdrs = EIGEN_MPL2_HEADER_FILES,
-    defines = [
-        # This define (mostly) guarantees we don't link any problematic
-        # code. We use it, but we do not rely on it, as evidenced above.
-        "EIGEN_MPL2_ONLY",
-        "EIGEN_MAX_ALIGN_BYTES=64",
-    ],
-    includes = ["."],
-    visibility = ["//visibility:public"],
-)
diff --git a/third_party/eigen_pin.txt b/third_party/eigen_pin.txt
index 0062ac971805f..6b244dcd6960b 100644
--- a/third_party/eigen_pin.txt
+++ b/third_party/eigen_pin.txt
@@ -1 +1 @@
-5.0.0
+5.0.1
diff --git a/third_party/flash-attention b/third_party/flash-attention
index e2743ab5b3803..fec3a6a18460c 160000
--- a/third_party/flash-attention
+++ b/third_party/flash-attention
@@ -1 +1 @@
-Subproject commit e2743ab5b3803bb672b16437ba98a3b1d4576c50
+Subproject commit fec3a6a18460c1b40f097208d4c16fe8964a679d
diff --git a/third_party/fmt.BUILD b/third_party/fmt.BUILD
deleted file mode 100644
index ea8c566b98a53..0000000000000
--- a/third_party/fmt.BUILD
+++ /dev/null
@@ -1,9 +0,0 @@
-load("@rules_cc//cc:defs.bzl", "cc_library")
-
-cc_library(
-    name = "fmt",
-    hdrs = glob(["include/fmt/*.h",]),
-    defines = ["FMT_HEADER_ONLY=1"],
-    includes = ["include"],
-    visibility = ["//visibility:public"],
-)
diff --git a/third_party/gloo.BUILD b/third_party/gloo.BUILD
deleted file mode 100644
index daa17f15e7655..0000000000000
--- a/third_party/gloo.BUILD
+++ /dev/null
@@ -1,88 +0,0 @@
-load("@rules_cc//cc:defs.bzl", "cc_library")
-load("@pytorch//tools/rules:cu.bzl", "cu_library")
-load("@pytorch//third_party:substitution.bzl", "template_rule")
-load("@pytorch//tools/config:defs.bzl", "if_cuda")
-
-template_rule(
-    name = "gloo_config_cmake_macros",
-    src = "gloo/config.h.in",
-    out = "gloo/config.h",
-    substitutions = {
-        "@GLOO_VERSION_MAJOR@": "0",
-        "@GLOO_VERSION_MINOR@": "5",
-        "@GLOO_VERSION_PATCH@": "0",
-        "cmakedefine01 GLOO_USE_CUDA": "define GLOO_USE_CUDA 1",
-        "cmakedefine01 GLOO_USE_NCCL": "define GLOO_USE_NCCL 0",
-        "cmakedefine01 GLOO_USE_ROCM": "define GLOO_USE_ROCM 0",
-        "cmakedefine01 GLOO_USE_RCCL": "define GLOO_USE_RCCL 0",
-        "cmakedefine01 GLOO_USE_REDIS": "define GLOO_USE_REDIS 0",
-        "cmakedefine01 GLOO_USE_IBVERBS": "define GLOO_USE_IBVERBS 0",
-        "cmakedefine01 GLOO_USE_MPI": "define GLOO_USE_MPI 0",
-        "cmakedefine01 GLOO_USE_AVX": "define GLOO_USE_AVX 0",
-        "cmakedefine01 GLOO_USE_LIBUV": "define GLOO_USE_LIBUV 0",
-        # The `GLOO_HAVE_TRANSPORT_TCP_TLS` line should go above the `GLOO_HAVE_TRANSPORT_TCP` in order to properly substitute the template.
-        "cmakedefine01 GLOO_HAVE_TRANSPORT_TCP_TLS": "define GLOO_HAVE_TRANSPORT_TCP_TLS 1",
-        "cmakedefine01 GLOO_HAVE_TRANSPORT_TCP": "define GLOO_HAVE_TRANSPORT_TCP 1",
-        "cmakedefine01 GLOO_HAVE_TRANSPORT_IBVERBS": "define GLOO_HAVE_TRANSPORT_IBVERBS 0",
-        "cmakedefine01 GLOO_HAVE_TRANSPORT_UV": "define GLOO_HAVE_TRANSPORT_UV 0",
-    },
-)
-
-cc_library(
-    name = "gloo_headers",
-    hdrs = glob(
-        [
-            "gloo/*.h",
-            "gloo/common/*.h",
-            "gloo/rendezvous/*.h",
-            "gloo/transport/*.h",
-            "gloo/transport/tcp/*.h",
-            "gloo/transport/tcp/tls/*.h",
-        ],
-        exclude = [
-            "gloo/rendezvous/redis_store.h",
-        ],
-    ) + ["gloo/config.h"],
-    includes = [
-        ".",
-    ],
-)
-
-cu_library(
-    name = "gloo_cuda",
-    srcs = [
-        "gloo/cuda.cu",
-        "gloo/cuda_private.cu",
-    ],
-    visibility = ["//visibility:public"],
-    deps = [
-        ":gloo_headers",
-    ],
-    alwayslink = True,
-)
-
-cc_library(
-    name = "gloo",
-    srcs = glob(
-        [
-            "gloo/*.cc",
-            "gloo/common/*.cc",
-            "gloo/rendezvous/*.cc",
-            "gloo/transport/*.cc",
-            "gloo/transport/tcp/*.cc",
-        ],
-        exclude = [
-            "gloo/cuda*.cc",
-            "gloo/common/win.cc",
-            "gloo/rendezvous/redis_store.cc",
-        ]
-    ) + if_cuda(glob(["gloo/cuda*.cc"])),
-    copts = [
-        "-std=c++17",
-    ],
-    visibility = ["//visibility:public"],
-    deps = [":gloo_headers"] + if_cuda(
-        [":gloo_cuda"],
-        [],
-    ),
-)
diff --git a/third_party/ideep b/third_party/ideep
index 8e7ddd65df95f..e539e0f9774e2 160000
--- a/third_party/ideep
+++ b/third_party/ideep
@@ -1 +1 @@
-Subproject commit 8e7ddd65df95f13e41f0a40c820c5f35ae4a0ea3
+Subproject commit e539e0f9774e2018f0d56fe865da66581f692e3d
diff --git a/third_party/ideep.BUILD b/third_party/ideep.BUILD
deleted file mode 100644
index 882d5cb342a41..0000000000000
--- a/third_party/ideep.BUILD
+++ /dev/null
@@ -1,17 +0,0 @@
-load("@rules_cc//cc:defs.bzl", "cc_library")
-
-cc_library(
-    name = "ideep",
-    hdrs = glob([
-        "include/**/*.hpp",
-        "include/**/*.h",
-    ]),
-    defines = [
-        "IDEEP_USE_MKL",
-    ],
-    includes = [
-        "include/",
-    ],
-    visibility = ["//visibility:public"],
-    deps = ["@mkl_dnn//:mkl-dnn"],
-)
diff --git a/third_party/kineto b/third_party/kineto
index 7a731b6ae01cf..b2103f78d13fd 160000
--- a/third_party/kineto
+++ b/third_party/kineto
@@ -1 +1 @@
-Subproject commit 7a731b6ae01cfc2b1fc75d83a91f84e682e43fd7
+Subproject commit b2103f78d13fde4937af010c0ef8e24313568bc5
diff --git a/third_party/kineto.BUILD b/third_party/kineto.BUILD
deleted file mode 100644
index d8e484ae80b6b..0000000000000
--- a/third_party/kineto.BUILD
+++ /dev/null
@@ -1,10 +0,0 @@
-load("@rules_cc//cc:defs.bzl", "cc_library")
-
-cc_library(
-    name = "kineto",
-    hdrs = glob(["libkineto/include/*.h",]),
-    includes = [
-        "libkineto/include/",
-    ],
-    visibility = ["//visibility:public"],
-)
diff --git a/third_party/llvm-openmp/LICENSE.txt b/third_party/llvm-openmp/LICENSE.txt
new file mode 100644
index 0000000000000..90837782e958e
--- /dev/null
+++ b/third_party/llvm-openmp/LICENSE.txt
@@ -0,0 +1,312 @@
+macOS wheels comes bundled with libomp.dylib from LLVM project and uses
+pre-built binaries from conda-forge
+
+==============================================================================
+The LLVM Project is under the Apache License v2.0 with LLVM Exceptions:
+==============================================================================
+
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+    TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+    1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+    2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+    3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+    4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+    5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+    6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+    7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+    8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+    9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+    END OF TERMS AND CONDITIONS
+
+    APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+    Copyright [yyyy] [name of copyright owner]
+
+    Licensed under the Apache License, Version 2.0 (the "License");
+    you may not use this file except in compliance with the License.
+    You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+    Unless required by applicable law or agreed to in writing, software
+    distributed under the License is distributed on an "AS IS" BASIS,
+    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+    See the License for the specific language governing permissions and
+    limitations under the License.
+
+
+---- LLVM Exceptions to the Apache 2.0 License ----
+
+As an exception, if, as a result of your compiling your source code, portions
+of this Software are embedded into an Object form of such source code, you
+may redistribute such embedded portions in such Object form without complying
+with the conditions of Sections 4(a), 4(b) and 4(d) of the License.
+
+In addition, if you combine or link compiled forms of this Software with
+software that is licensed under the GPLv2 ("Combined Software") and if a
+court of competent jurisdiction determines that the patent provision (Section
+3), the indemnity provision (Section 9) or other Section of the License
+conflicts with the conditions of the GPLv2, you may retroactively and
+prospectively choose to deem waived or otherwise exclude such Section(s) of
+the License, but only in their entirety and only with respect to the Combined
+Software.
+
+==============================================================================
+Software from third parties included in the LLVM Project:
+==============================================================================
+The LLVM Project contains third party software which is under different license
+terms. All such code will be identified clearly using at least one of two
+mechanisms:
+1) It will be in a separate directory tree with its own `LICENSE.txt` or
+   `LICENSE` file at the top containing the specific license and restrictions
+   which apply to that software, or
+2) It will contain specific license and restriction terms at the top of every
+   file.
+
+==============================================================================
+Legacy LLVM License (https://llvm.org/docs/DeveloperPolicy.html#legacy):
+==============================================================================
+University of Illinois/NCSA
+Open Source License
+
+Copyright (c) 2003-2019 University of Illinois at Urbana-Champaign.
+All rights reserved.
+
+Developed by:
+
+    LLVM Team
+
+    University of Illinois at Urbana-Champaign
+
+    http://llvm.org
+
+Permission is hereby granted, free of charge, to any person obtaining a copy of
+this software and associated documentation files (the "Software"), to deal with
+the Software without restriction, including without limitation the rights to
+use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies
+of the Software, and to permit persons to whom the Software is furnished to do
+so, subject to the following conditions:
+
+    * Redistributions of source code must retain the above copyright notice,
+      this list of conditions and the following disclaimers.
+
+    * Redistributions in binary form must reproduce the above copyright notice,
+      this list of conditions and the following disclaimers in the
+      documentation and/or other materials provided with the distribution.
+
+    * Neither the names of the LLVM Team, University of Illinois at
+      Urbana-Champaign, nor the names of its contributors may be used to
+      endorse or promote products derived from this Software without specific
+      prior written permission.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS
+FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL THE
+CONTRIBUTORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS WITH THE
+SOFTWARE.
+
+================================================================================
+Binaries from conda-forge are licensed as follows
+================================================================================
+BSD-3-Clause license
+Copyright (c) 2015-2022, conda-forge contributors
+All rights reserved.
+
+Redistribution and use in source and binary forms, with or without
+modification, are permitted provided that the following conditions are met:
+
+  1. Redistributions of source code must retain the above copyright notice,
+     this list of conditions and the following disclaimer.
+  2. Redistributions in binary form must reproduce the above copyright
+     notice, this list of conditions and the following disclaimer in the
+     documentation and/or other materials provided with the distribution.
+  3. Neither the name of the copyright holder nor the names of its
+     contributors may be used to endorse or promote products derived from
+     this software without specific prior written permission.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
+ARE DISCLAIMED. IN NO EVENT SHALL THE REGENTS OR CONTRIBUTORS BE LIABLE FOR
+ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
+LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY
+OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH
+DAMAGE.
diff --git a/third_party/mkl-dnn.BUILD b/third_party/mkl-dnn.BUILD
deleted file mode 100644
index f62612ba18b86..0000000000000
--- a/third_party/mkl-dnn.BUILD
+++ /dev/null
@@ -1,164 +0,0 @@
-load("@rules_cc//cc:defs.bzl", "cc_library")
-load("@pytorch//third_party:substitution.bzl", "template_rule")
-
-_DNNL_RUNTIME_OMP = {
-    "#cmakedefine DNNL_CPU_THREADING_RUNTIME DNNL_RUNTIME_${DNNL_CPU_THREADING_RUNTIME}": "#define DNNL_CPU_THREADING_RUNTIME DNNL_RUNTIME_OMP",
-    "#cmakedefine DNNL_CPU_RUNTIME DNNL_RUNTIME_${DNNL_CPU_RUNTIME}": "#define DNNL_CPU_RUNTIME DNNL_RUNTIME_OMP",
-    "#cmakedefine DNNL_GPU_RUNTIME DNNL_RUNTIME_${DNNL_GPU_RUNTIME}": "#define DNNL_GPU_RUNTIME DNNL_RUNTIME_NONE",
-    "#cmakedefine DNNL_GPU_VENDOR DNNL_VENDOR_${DNNL_GPU_VENDOR}": "/* undef DNNL_GPU_VENDOR */",
-    "#cmakedefine DNNL_USE_RT_OBJECTS_IN_PRIMITIVE_CACHE": "/* undef DNNL_USE_RT_OBJECTS_IN_PRIMITIVE_CACHE */",
-    "#cmakedefine DNNL_WITH_SYCL": "/* #undef DNNL_WITH_SYCL */",
-    "#cmakedefine DNNL_WITH_LEVEL_ZERO": "/* #undef DNNL_WITH_LEVEL_ZERO */",
-    "#cmakedefine DNNL_SYCL_CUDA": "/* #undef DNNL_SYCL_CUDA */",
-    "#cmakedefine DNNL_SYCL_HIP": "/* #undef DNNL_SYCL_HIP */",
-    "#cmakedefine DNNL_SYCL_GENERIC": "/* #undef DNNL_SYCL_GENERIC */",
-    "#cmakedefine DNNL_ENABLE_STACK_CHECKER": "#undef DNNL_ENABLE_STACK_CHECKER",
-    "#cmakedefine DNNL_EXPERIMENTAL_UKERNEL": "/* undef DNNL_EXPERIMENTAL_UKERNEL */",
-    "#cmakedefine DNNL_EXPERIMENTAL": "#undef DNNL_EXPERIMENTAL",
-    "#cmakedefine ONEDNN_BUILD_GRAPH": "#undef ONEDNN_BUILD_GRAPH",
-    "#cmakedefine DNNL_EXPERIMENTAL_PROFILING": "#undef DNNL_EXPERIMENTAL_PROFILING",
-    "#cmakedefine DNNL_EXPERIMENTAL_LOGGING": "#undef DNNL_EXPERIMENTAL_LOGGING",
-    "#cmakedefine DNNL_EXPERIMENTAL_SYCL_KERNEL_COMPILER": "#undef DNNL_EXPERIMENTAL_SYCL_KERNEL_COMPILER",
-    "#cmakedefine DNNL_DISABLE_GPU_REF_KERNELS": "#undef DNNL_DISABLE_GPU_REF_KERNELS",
-    "#cmakedefine01 BUILD_TRAINING": "#define BUILD_TRAINING 1",
-    "#cmakedefine01 BUILD_INFERENCE": "#define BUILD_INFERENCE 0",
-    "#cmakedefine01 BUILD_PRIMITIVE_ALL": "#define BUILD_PRIMITIVE_ALL 1",
-    "#cmakedefine01 BUILD_BATCH_NORMALIZATION": "#define BUILD_BATCH_NORMALIZATION 0",
-    "#cmakedefine01 BUILD_BINARY": "#define BUILD_BINARY 0",
-    "#cmakedefine01 BUILD_CONCAT": "#define BUILD_CONCAT 0",
-    "#cmakedefine01 BUILD_CONVOLUTION": "#define BUILD_CONVOLUTION 0",
-    "#cmakedefine01 BUILD_DECONVOLUTION": "#define BUILD_DECONVOLUTION 0",
-    "#cmakedefine01 BUILD_ELTWISE": "#define BUILD_ELTWISE 0",
-    "#cmakedefine01 BUILD_GROUP_NORMALIZATION": "#define BUILD_GROUP_NORMALIZATION 0",
-    "#cmakedefine01 BUILD_INNER_PRODUCT": "#define BUILD_INNER_PRODUCT 0",
-    "#cmakedefine01 BUILD_LAYER_NORMALIZATION": "#define BUILD_LAYER_NORMALIZATION 0",
-    "#cmakedefine01 BUILD_LRN": "#define BUILD_LRN 0",
-    "#cmakedefine01 BUILD_MATMUL": "#define BUILD_MATMUL 0",
-    "#cmakedefine01 BUILD_POOLING": "#define BUILD_POOLING 0",
-    "#cmakedefine01 BUILD_PRELU": "#define BUILD_PRELU 0",
-    "#cmakedefine01 BUILD_REDUCTION": "#define BUILD_REDUCTION 0",
-    "#cmakedefine01 BUILD_REORDER": "#define BUILD_REORDER 0",
-    "#cmakedefine01 BUILD_RESAMPLING": "#define BUILD_RESAMPLING 0",
-    "#cmakedefine01 BUILD_RNN": "#define BUILD_RNN 0",
-    "#cmakedefine01 BUILD_SDPA": "#define BUILD_SDPA 0",
-    "#cmakedefine01 BUILD_SHUFFLE": "#define BUILD_SHUFFLE 0",
-    "#cmakedefine01 BUILD_SOFTMAX": "#define BUILD_SOFTMAX 0",
-    "#cmakedefine01 BUILD_SUM": "#define BUILD_SUM 0",
-    "#cmakedefine01 BUILD_PRIMITIVE_CPU_ISA_ALL": "#define BUILD_PRIMITIVE_CPU_ISA_ALL 1",
-    "#cmakedefine01 BUILD_SSE41": "#define BUILD_SSE41 0",
-    "#cmakedefine01 BUILD_AVX2": "#define BUILD_AVX2 0",
-    "#cmakedefine01 BUILD_AVX512": "#define BUILD_AVX512 0",
-    "#cmakedefine01 BUILD_AMX": "#define BUILD_AMX 0",
-    "#cmakedefine01 BUILD_PRIMITIVE_GPU_ISA_ALL": "#define BUILD_PRIMITIVE_GPU_ISA_ALL 1",
-    "#cmakedefine01 BUILD_XELP": "#define BUILD_XELP 0",
-    "#cmakedefine01 BUILD_XEHPG": "#define BUILD_XEHPG 0",
-    "#cmakedefine01 BUILD_XEHPC": "#define BUILD_XEHPC 0",
-    "#cmakedefine01 BUILD_XEHP": "#define BUILD_XEHP 0",
-    "#cmakedefine01 BUILD_XE2": "#define BUILD_XE2 0",
-    "#cmakedefine01 BUILD_XE3": "#define BUILD_XE3 0",
-    "#cmakedefine01 BUILD_GEMM_KERNELS_ALL": "#define BUILD_GEMM_KERNELS_ALL 0",
-    "#cmakedefine01 BUILD_GEMM_KERNELS_NONE": "#define BUILD_GEMM_KERNELS_NONE 0",
-    "#cmakedefine01 BUILD_GEMM_SSE41": "#define BUILD_GEMM_SSE41 0",
-    "#cmakedefine01 BUILD_GEMM_AVX2": "#define BUILD_GEMM_AVX2 0",
-    "#cmakedefine01 BUILD_GEMM_AVX512": "#define BUILD_GEMM_AVX512 0",
-}
-
-template_rule(
-    name = "include_dnnl_version",
-    src = "include/oneapi/dnnl/dnnl_version.h.in",
-    out = "include/oneapi/dnnl/dnnl_version.h",
-    substitutions = {
-        "@DNNL_VERSION_MAJOR@": "3",
-        "@DNNL_VERSION_MINOR@": "10",
-        "@DNNL_VERSION_PATCH@": "2",
-    },
-)
-
-template_rule(
-    name = "include_dnnl_config",
-    src = "include/oneapi/dnnl/dnnl_config.h.in",
-    out = "include/oneapi/dnnl/dnnl_config.h",
-    substitutions = _DNNL_RUNTIME_OMP,
-)
-
-template_rule(
-    name = "include_dnnl_version_hash",
-    src = "include/oneapi/dnnl/dnnl_version_hash.h.in",
-    out = "include/oneapi/dnnl/dnnl_version_hash.h",
-    substitutions = {"@DNNL_VERSION_HASH@": "f1d471933dc852f956fd05389f9313c7148783d5",}
-)
-
-cc_library(
-    name = "mkl-dnn",
-    srcs = glob([
-        "src/common/*.cpp",
-        "src/cpu/**/*.cpp",
-        "src/cpu/**/**/*.cpp",
-    ], exclude=[
-        "src/cpu/aarch64/**/*.cpp",
-        "src/cpu/rv64/**/*.cpp",
-        "src/cpu/sycl/**/*.cpp",
-        "src/cpu/ppc64/**/*.cpp",
-    ]),
-    hdrs = glob([
-        "include/oneapi/dnnl/*.h",
-        "include/oneapi/dnnl/*.hpp",
-        "include/*.h",
-        "include/*.hpp",
-        "src/cpu/**/*.hpp",
-        "src/cpu/**/*.h",
-        "src/cpu/**/**/*.h",
-        "src/common/*.hpp",
-        "src/common/**/**/*.h",
-        "third_party/xbyak/*.h",
-        "third_party/ittnotify/jitprofiling.h",
-        "third_party/spdlog/**/*.h",
-    ], exclude=[
-        "src/cpu/aarch64/**/*.hpp",
-        "src/cpu/aarch64/**/*.h",
-        "src/cpu/rv64/**/*.hpp",
-        "src/cpu/rv64/**/*.h",
-        "src/cpu/sycl/**/*.hpp",
-        "src/cpu/ppc64/**/*.hpp",
-    ]) + [
-        "include/oneapi/dnnl/dnnl_config.h",
-        "include/oneapi/dnnl/dnnl_version.h",
-        "include/oneapi/dnnl/dnnl_version_hash.h",
-    ],
-    copts = [
-        "-DDNNL_DLL",
-        "-DDNNL_DLL_EXPORTS",
-        "-DDNNL_ENABLE_CONCURRENT_EXEC",
-        "-D__STDC_CONSTANT_MACROS",
-        "-D__STDC_LIMIT_MACROS",
-        "-fno-strict-overflow",
-        "-fopenmp",
-    ] + select({
-        "@pytorch//tools/config:thread_sanitizer": ["-DDNNL_CPU_RUNTIME=0"],
-        "//conditions:default": ["-DDNNL_CPU_RUNTIME=2"],
-    }),
-    includes = [
-        "include/",
-        "include/oneapi/",
-        "include/oneapi/dnnl/",
-        "src/",
-        "src/common/",
-        "src/cpu/",
-        "third_party/",
-    ],
-    visibility = ["//visibility:public"],
-    linkopts = [
-        "-lgomp",
-    ],
-    deps = [
-        "@mkl",
-    ],
-    defines = [
-        "DNNL_ENABLE_MAX_CPU_ISA",
-        "DNNL_ENABLE_CONCURRENT_EXEC",
-        "DNNL_ENABLE_PRIMITIVE_CACHE",
-        "DNNL_ENABLE_CPU_ISA_HINTS",
-        "DNNL_EXPERIMENTAL_UKERNEL",
-        "ONEDNN_BUILD_GRAPH",
-    ],
-)
diff --git a/third_party/mkl.BUILD b/third_party/mkl.BUILD
deleted file mode 100644
index b7abb0e035ad9..0000000000000
--- a/third_party/mkl.BUILD
+++ /dev/null
@@ -1,18 +0,0 @@
-load("@rules_cc//cc:defs.bzl", "cc_library")
-
-cc_library(
-    name = "mkl",
-    srcs = [
-        "libmkl_avx2.so",
-        "libmkl_core.so",
-        "libmkl_def.so",
-        "libmkl_intel_lp64.so",
-        "libmkl_rt.so",
-        "libmkl_sequential.so",
-        "libmkl_vml_avx2.so",
-        "libmkl_vml_avx512.so",
-        "libmkl_vml_def.so",
-    ],
-    visibility = ["//visibility:public"],
-    deps = ["@mkl_headers"],
-)
diff --git a/third_party/mkl_headers.BUILD b/third_party/mkl_headers.BUILD
deleted file mode 100644
index 965801c91aa97..0000000000000
--- a/third_party/mkl_headers.BUILD
+++ /dev/null
@@ -1,8 +0,0 @@
-load("@rules_cc//cc:defs.bzl", "cc_library")
-
-cc_library(
-    name = "mkl_headers",
-    hdrs = glob(["include/*.h"]),
-    includes = ["include/"],
-    visibility = ["//visibility:public"],
-)
diff --git a/third_party/moodycamel.BUILD b/third_party/moodycamel.BUILD
deleted file mode 100644
index d3028205016fb..0000000000000
--- a/third_party/moodycamel.BUILD
+++ /dev/null
@@ -1,7 +0,0 @@
-load("@rules_cc//cc:defs.bzl", "cc_library")
-
-cc_library(
-    name = "moodycamel",
-    hdrs = glob(["**/*.h"]),
-    visibility = ["//visibility:public"],
-)
diff --git a/third_party/nlohmann.BUILD b/third_party/nlohmann.BUILD
deleted file mode 100644
index 64dfbbab2b6e9..0000000000000
--- a/third_party/nlohmann.BUILD
+++ /dev/null
@@ -1,18 +0,0 @@
-load("@rules_cc//cc:defs.bzl", "cc_library")
-
-cc_library(name = "nlohmann",
-    includes = ["include"],
-    deps = ["nlohmann-internal"],
-    visibility = ["//visibility:public"],
-)
-
-cc_import(name = "nlohmann-internal",
-     hdrs = glob(["include/**/*.hpp"]),
-     visibility = ["//visibility:private"],
-)
-
-cc_library(
-    name = "nlohmann_single_include",
-    hdrs = glob(["single_include/nlohmann/*.hpp"]),
-    visibility = ["//visibility:public"],
-)
diff --git a/third_party/onnx.BUILD b/third_party/onnx.BUILD
deleted file mode 100644
index c4134e98d0fe2..0000000000000
--- a/third_party/onnx.BUILD
+++ /dev/null
@@ -1,108 +0,0 @@
-load("@rules_proto//proto:defs.bzl", "proto_library")
-load("@rules_cc//cc:defs.bzl", "cc_library", "cc_proto_library")
-load("@rules_python//python:defs.bzl", "py_binary")
-
-py_binary(
-    name = "gen_proto",
-    srcs = ["onnx/gen_proto.py"],
-    data = [
-        "onnx/onnx.in.proto",
-        "onnx/onnx-operators.in.proto",
-        "onnx/onnx-data.in.proto",
-    ],
-)
-
-genrule(
-    name = "generate_onnx_proto",
-    outs = [
-        "onnx/onnx_onnx_torch-ml.proto",
-        "onnx/onnx-ml.pb.h",
-    ],
-    cmd = "$(location :gen_proto) -p onnx_torch -o $(@D)/onnx onnx -m >/dev/null && sed -i 's/onnx_onnx_torch-ml.pb.h/onnx\\/onnx_onnx_torch-ml.pb.h/g' $(@D)/onnx/onnx-ml.pb.h",
-    tools = [":gen_proto"],
-)
-
-genrule(
-    name = "generate_onnx_operators_proto",
-    outs = [
-        "onnx/onnx-operators_onnx_torch-ml.proto",
-        "onnx/onnx-operators-ml.pb.h",
-    ],
-    cmd = "$(location :gen_proto) -p onnx_torch -o $(@D)/onnx onnx-operators -m >/dev/null && sed -i 's/onnx-operators_onnx_torch-ml.pb.h/onnx\\/onnx-operators_onnx_torch-ml.pb.h/g' $(@D)/onnx/onnx-operators-ml.pb.h",
-    tools = [":gen_proto"],
-)
-
-genrule(
-    name = "generate_onnx_data_proto",
-    outs = [
-        "onnx/onnx-data_onnx_torch.proto",
-        "onnx/onnx-data.pb.h",
-    ],
-    cmd = "$(location :gen_proto) -p onnx_torch -o $(@D)/onnx onnx-data -m >/dev/null && sed -i 's/onnx-data_onnx_torch.pb.h/onnx\\/onnx-data_onnx_torch.pb.h/g' $(@D)/onnx/onnx-data.pb.h",
-    tools = [":gen_proto"],
-)
-
-cc_library(
-    name = "onnx",
-    srcs = glob(
-        [
-            "onnx/*.cc",
-            "onnx/common/*.cc",
-            "onnx/defs/**/*.cc",
-            "onnx/shape_inference/*.cc",
-            "onnx/version_converter/*.cc",
-        ],
-        exclude = [
-            "onnx/cpp2py_export.cc",
-        ],
-    ),
-    hdrs = glob([
-        "onnx/*.h",
-        "onnx/version_converter/*.h",
-        "onnx/common/*.h",
-        "onnx/defs/**/*.h",
-        "onnx/shape_inference/*.h",
-        "onnx/version_converter/adapters/*.h",
-    ]) + [
-        "onnx/onnx-ml.pb.h",
-        "onnx/onnx-operators-ml.pb.h",
-        "onnx/onnx-data.pb.h",
-    ],
-    defines = [
-        "ONNX_ML=1",
-        "ONNX_NAMESPACE=onnx_torch",
-    ],
-    includes = [
-        ".",
-        "onnx/",
-    ],
-    visibility = ["//visibility:public"],
-    deps = [
-        ":onnx_proto_lib",
-    ],
-)
-
-cc_library(
-    name = "onnx_proto_headers",
-    hdrs = glob([
-        "onnx/*_pb.h",
-    ]),
-    visibility = ["//visibility:public"],
-    deps = [
-        ":onnx_proto_lib",
-    ],
-)
-
-proto_library(
-    name = "onnx_proto",
-    srcs = [
-        "onnx/onnx-operators_onnx_torch-ml.proto",
-        "onnx/onnx_onnx_torch-ml.proto",
-        "onnx/onnx-data_onnx_torch.proto",
-    ],
-)
-
-cc_proto_library(
-    name = "onnx_proto_lib",
-    deps = [":onnx_proto"],
-)
diff --git a/third_party/sleef.BUILD b/third_party/sleef.BUILD
deleted file mode 100644
index f22a6e905e2be..0000000000000
--- a/third_party/sleef.BUILD
+++ /dev/null
@@ -1,487 +0,0 @@
-load("@rules_cc//cc:defs.bzl", "cc_binary", "cc_library")
-load("@pytorch//third_party:sleef.bzl", "sleef_cc_library")
-
-SLEEF_COPTS = [
-    "-DHAVE_MALLOC_USABLE_SIZE=1",
-    "-DHAVE_MMAP=1",
-    "-DHAVE_SHM_OPEN=1",
-    "-DHAVE_SHM_UNLINK=1",
-    "-DIDEEP_USE_MKL",
-    "-DDNNL_CPU_RUNTIME=TBB",
-    "-DONNX_ML=1",
-    "-DONNX_NAMESPACE=onnx",
-    "-D_FILE_OFFSET_BITS=64",
-    "-ffp-contract=off",
-    "-fno-math-errno",
-    "-fno-trapping-math",
-    "-DCAFFE2_USE_GLOO",
-    "-std=gnu99",
-]
-
-SLEEF_COMMON_TARGET_COPTS = [
-    "-DSLEEF_STATIC_LIBS=1",
-    "-DENABLE_ALIAS=1",
-]
-
-SLEEF_PRIVATE_HEADERS = glob([
-    "build/include/*.h",
-    "src/arch/*.h",
-    "src/common/*.h",
-    "src/libm/*.h",
-    "src/libm/include/*.h",
-])
-
-SLEEF_PUBLIC_HEADERS = [
-    ":sleef_h",
-]
-
-SLEEF_PRIVATE_INCLUDES = [
-    "-Iexternal/sleef/src/arch",
-    "-Iexternal/sleef/src/common",
-    "-Iexternal/sleef/src/libm",
-]
-
-SLEEF_PUBLIC_INCLUDES = [
-    "build/include",
-]
-
-SLEEF_VISIBILITY = [
-    "//visibility:public",
-]
-
-cc_binary(
-    name = "mkalias",
-    srcs = [
-        "src/libm/funcproto.h",
-        "src/libm/mkalias.c",
-    ],
-)
-
-genrule(
-    name = "alias_avx512f_h",
-    outs = ["alias_avx512f.h"],
-    cmd = "{ " + "; ".join([
-        "$(location :mkalias) -16 __m512 __m512i e avx512f",
-        "$(location :mkalias) 8 __m512d __m256i e avx512f",
-    ]) + "; } > $@",
-    tools = [":mkalias"],
-)
-
-cc_binary(
-    name = "mkdisp",
-    srcs = [
-        "src/libm/funcproto.h",
-        "src/libm/mkdisp.c",
-    ],
-    copts = SLEEF_COPTS,
-)
-
-genrule(
-    name = "dispavx_c",
-    srcs = ["src/libm/dispavx.c.org"],
-    outs = ["dispavx.c"],
-    cmd = "{ cat $(location src/libm/dispavx.c.org); $(location :mkdisp) 4 8 __m256d __m256 __m128i avx fma4 avx2; } > $@",
-    tools = [":mkdisp"],
-)
-
-genrule(
-    name = "dispsse_c",
-    srcs = ["src/libm/dispsse.c.org"],
-    outs = ["dispsse.c"],
-    cmd = "{ cat $(location src/libm/dispsse.c.org); $(location :mkdisp) 2 4 __m128d __m128 __m128i sse2 sse4 avx2128; } > $@",
-    tools = [":mkdisp"],
-)
-
-cc_binary(
-    name = "mkrename",
-    srcs = [
-        "src/libm/funcproto.h",
-        "src/libm/mkrename.c",
-    ],
-)
-
-genrule(
-    name = "renameavx_h",
-    outs = ["renameavx.h"],
-    cmd = "$(location :mkrename) cinz_ 4 8 avx > $@",
-    tools = [":mkrename"],
-)
-
-genrule(
-    name = "renameavx2_h",
-    outs = ["renameavx2.h"],
-    cmd = "$(location :mkrename) finz_ 4 8 avx2 > $@",
-    tools = [":mkrename"],
-)
-
-genrule(
-    name = "renameavx2128_h",
-    outs = ["renameavx2128.h"],
-    cmd = "$(location :mkrename) finz_ 2 4 avx2128 > $@",
-    tools = [":mkrename"],
-)
-
-genrule(
-    name = "renameavx512f_h",
-    outs = ["renameavx512f.h"],
-    cmd = "$(location :mkrename) finz_ 8 16 avx512f > $@",
-    tools = [":mkrename"],
-)
-
-genrule(
-    name = "renameavx512fnofma_h",
-    outs = ["renameavx512fnofma.h"],
-    cmd = "$(location :mkrename) cinz_ 8 16 avx512fnofma > $@",
-    tools = [":mkrename"],
-)
-
-genrule(
-    name = "renamefma4_h",
-    outs = ["renamefma4.h"],
-    cmd = "$(location :mkrename) finz_ 4 8 fma4 > $@",
-    tools = [":mkrename"],
-)
-
-genrule(
-    name = "renamepurec_scalar_h",
-    outs = ["renamepurec_scalar.h"],
-    cmd = "$(location :mkrename) cinz_ 1 1 purec > $@",
-    tools = [":mkrename"],
-)
-
-genrule(
-    name = "renamepurecfma_scalar_h",
-    outs = ["renamepurecfma_scalar.h"],
-    cmd = "$(location :mkrename) finz_ 1 1 purecfma > $@",
-    tools = [":mkrename"],
-)
-
-genrule(
-    name = "renamesse2_h",
-    outs = ["renamesse2.h"],
-    cmd = "$(location :mkrename) cinz_ 2 4 sse2 > $@",
-    tools = [":mkrename"],
-)
-
-genrule(
-    name = "renamesse4_h",
-    outs = ["renamesse4.h"],
-    cmd = "$(location :mkrename) cinz_ 2 4 sse4 > $@",
-    tools = [":mkrename"],
-)
-
-genrule(
-    name = "sleef_h",
-    srcs = [
-        "src/libm/sleeflibm_header.h.org.in",
-        "src/libm/sleeflibm_footer.h.org",
-    ],
-    outs = ["build/include/sleef.h"],
-    cmd = "{ " + "; ".join([
-        "cat $(location src/libm/sleeflibm_header.h.org.in)",
-        "$(location :mkrename) cinz_ 2 4 __m128d __m128 __m128i __m128i __SSE2__",
-        "$(location :mkrename) cinz_ 2 4 __m128d __m128 __m128i __m128i __SSE2__ sse2",
-        "$(location :mkrename) cinz_ 2 4 __m128d __m128 __m128i __m128i __SSE2__ sse4",
-        "$(location :mkrename) cinz_ 4 8 __m256d __m256 __m128i \"struct { __m128i x, y; }\" __AVX__",
-        "$(location :mkrename) cinz_ 4 8 __m256d __m256 __m128i \"struct { __m128i x, y; }\" __AVX__ avx",
-        "$(location :mkrename) finz_ 4 8 __m256d __m256 __m128i \"struct { __m128i x, y; }\" __AVX__ fma4",
-        "$(location :mkrename) finz_ 4 8 __m256d __m256 __m128i __m256i __AVX__ avx2",
-        "$(location :mkrename) finz_ 2 4 __m128d __m128 __m128i __m128i __SSE2__ avx2128",
-        "$(location :mkrename) finz_ 8 16 __m512d __m512 __m256i __m512i __AVX512F__",
-        "$(location :mkrename) finz_ 8 16 __m512d __m512 __m256i __m512i __AVX512F__ avx512f",
-        "$(location :mkrename) cinz_ 8 16 __m512d __m512 __m256i __m512i __AVX512F__ avx512fnofma",
-        "$(location :mkrename) cinz_ 1 1 double float int32_t int32_t __STDC__ purec",
-        "$(location :mkrename) finz_ 1 1 double float int32_t int32_t FP_FAST_FMA purecfma",
-        "cat $(location src/libm/sleeflibm_footer.h.org)",
-    ]) + "; } > $@",
-    tools = [":mkrename"],
-)
-
-cc_library(
-    name = "sleef",
-    srcs = [
-        "src/libm/rempitab.c",
-        "src/libm/sleefdp.c",
-        "src/libm/sleefsp.c",
-    ],
-    hdrs = SLEEF_PUBLIC_HEADERS,
-    copts = SLEEF_PRIVATE_INCLUDES + SLEEF_COPTS + SLEEF_COMMON_TARGET_COPTS + [
-        "-DDORENAME=1",
-        "-DENABLEFLOAT128=1",
-        "-Wno-unused-result",
-    ],
-    includes = SLEEF_PUBLIC_INCLUDES,
-    # -lgcc resolves
-    # U __addtf3
-    # U __eqtf2
-    # U __fixtfdi
-    # U __floatditf
-    # U __gttf2
-    # U __lttf2
-    # U __multf3
-    # U __subtf3
-    # in bazel-bin/external/sleef/_objs/sleef/sleefqp.pic.o
-    linkopts = [
-        "-lgcc",
-    ],
-    linkstatic = True,
-    visibility = SLEEF_VISIBILITY,
-    # The purpose of the lists in deps is to keep related pairs of
-    # libraries together. In particular, each pair that contains a *det*
-    # library originates with a sleef_cc_library().
-    deps = [
-        ":common",
-        ":dispavx",
-        ":dispsse",
-    ] + [
-        ":sleefavx",
-        ":sleefdetavx",
-    ] + [
-        ":sleefavx2",
-        ":sleefdetavx2",
-    ] + [
-        ":sleefavx2128",
-        ":sleefdetavx2128",
-    ] + [
-        ":sleefavx512f",
-        ":sleefdetavx512f",
-    ] + [
-        ":sleefavx512fnofma",
-        ":sleefdetavx512fnofma",
-    ] + [
-        ":sleeffma4",
-        ":sleefdetfma4",
-    ] + [
-        ":sleefsse2",
-        ":sleefdetsse2",
-    ] + [
-        ":sleefsse4",
-        ":sleefdetsse4",
-    ] + [
-        ":sleefpurec_scalar",
-        ":sleefdetpurec_scalar",
-    ] + [
-        ":sleefpurecfma_scalar",
-        ":sleefdetpurecfma_scalar",
-    ],
-    alwayslink = True,
-)
-
-cc_library(
-    name = "common",
-    srcs = SLEEF_PRIVATE_HEADERS + [
-        "src/common/common.c",
-    ],
-    copts = SLEEF_PRIVATE_INCLUDES + SLEEF_COPTS + [
-        "-Wno-unused-result",
-    ],
-    linkstatic = True,
-    visibility = SLEEF_VISIBILITY,
-    alwayslink = True,
-)
-
-cc_library(
-    name = "dispavx",
-    srcs = SLEEF_PRIVATE_HEADERS + SLEEF_PUBLIC_HEADERS + [
-        ":dispavx_c",
-    ],
-    copts = SLEEF_PRIVATE_INCLUDES + SLEEF_COPTS + SLEEF_COMMON_TARGET_COPTS + [
-        "-DENABLE_AVX2=1",
-        "-DENABLE_FMA4=1",
-        "-mavx",
-    ],
-    includes = SLEEF_PUBLIC_INCLUDES,
-    linkstatic = True,
-    visibility = SLEEF_VISIBILITY,
-    alwayslink = True,
-)
-
-cc_library(
-    name = "dispsse",
-    srcs = SLEEF_PRIVATE_HEADERS + SLEEF_PUBLIC_HEADERS + [
-        ":dispsse_c",
-    ],
-    copts = SLEEF_PRIVATE_INCLUDES + SLEEF_COPTS + SLEEF_COMMON_TARGET_COPTS + [
-        "-DENABLE_AVX2=1",
-        "-DENABLE_FMA4=1",
-        "-msse2",
-    ],
-    includes = SLEEF_PUBLIC_INCLUDES,
-    linkstatic = True,
-    visibility = SLEEF_VISIBILITY,
-    alwayslink = True,
-)
-
-sleef_cc_library(
-    name = "sleefavx512f",
-    srcs = SLEEF_PRIVATE_HEADERS + [
-        "src/libm/sleefsimddp.c",
-        "src/libm/sleefsimdsp.c",
-        ":alias_avx512f_h",
-        ":renameavx512f_h",
-    ],
-    copts = SLEEF_PRIVATE_INCLUDES + SLEEF_COPTS + SLEEF_COMMON_TARGET_COPTS + [
-        "-DDORENAME=1",
-        "-DALIAS_NO_EXT_SUFFIX=\\\"alias_avx512f.h\\\"",
-        "-DENABLE_AVX512F=1",
-        "-mavx512f",
-    ],
-    linkstatic = True,
-    visibility = SLEEF_VISIBILITY,
-    alwayslink = True,
-)
-
-sleef_cc_library(
-    name = "sleefavx512fnofma",
-    srcs = SLEEF_PRIVATE_HEADERS + [
-        "src/libm/sleefsimddp.c",
-        "src/libm/sleefsimdsp.c",
-        ":renameavx512fnofma_h",
-    ],
-    copts = SLEEF_PRIVATE_INCLUDES + SLEEF_COPTS + SLEEF_COMMON_TARGET_COPTS + [
-        "-DDORENAME=1",
-        "-DENABLE_AVX512FNOFMA=1",
-        "-mavx512f",
-    ],
-    linkstatic = True,
-    visibility = SLEEF_VISIBILITY,
-    alwayslink = True,
-)
-
-sleef_cc_library(
-    name = "sleefavx",
-    srcs = SLEEF_PRIVATE_HEADERS + [
-        "src/libm/sleefsimddp.c",
-        "src/libm/sleefsimdsp.c",
-        ":renameavx_h",
-    ],
-    copts = SLEEF_PRIVATE_INCLUDES + SLEEF_COPTS + SLEEF_COMMON_TARGET_COPTS + [
-        "-DDORENAME=1",
-        "-DENABLE_AVX=1",
-        "-mavx",
-    ],
-    linkstatic = True,
-    visibility = SLEEF_VISIBILITY,
-    alwayslink = True,
-)
-
-sleef_cc_library(
-    name = "sleefavx2",
-    srcs = SLEEF_PRIVATE_HEADERS + [
-        "src/libm/sleefsimddp.c",
-        "src/libm/sleefsimdsp.c",
-        ":renameavx2_h",
-    ],
-    copts = SLEEF_PRIVATE_INCLUDES + SLEEF_COPTS + SLEEF_COMMON_TARGET_COPTS + [
-        "-DDORENAME=1",
-        "-DENABLE_AVX2=1",
-        "-mavx2",
-        "-mfma",
-    ],
-    linkstatic = True,
-    visibility = SLEEF_VISIBILITY,
-    alwayslink = True,
-)
-
-sleef_cc_library(
-    name = "sleefavx2128",
-    srcs = SLEEF_PRIVATE_HEADERS + [
-        "src/libm/sleefsimddp.c",
-        "src/libm/sleefsimdsp.c",
-        ":renameavx2128_h",
-    ],
-    copts = SLEEF_PRIVATE_INCLUDES + SLEEF_COPTS + SLEEF_COMMON_TARGET_COPTS + [
-        "-DDORENAME=1",
-        "-DENABLE_AVX2128=1",
-        "-mavx2",
-        "-mfma",
-    ],
-    linkstatic = True,
-    visibility = SLEEF_VISIBILITY,
-    alwayslink = True,
-)
-
-sleef_cc_library(
-    name = "sleeffma4",
-    srcs = SLEEF_PRIVATE_HEADERS + [
-        "src/libm/sleefsimddp.c",
-        "src/libm/sleefsimdsp.c",
-        ":renamefma4_h",
-    ],
-    copts = SLEEF_PRIVATE_INCLUDES + SLEEF_COPTS + SLEEF_COMMON_TARGET_COPTS + [
-        "-DDORENAME=1",
-        "-DENABLE_FMA4=1",
-        "-mfma4",
-    ],
-    linkstatic = True,
-    visibility = SLEEF_VISIBILITY,
-    alwayslink = True,
-)
-
-sleef_cc_library(
-    name = "sleefsse2",
-    srcs = SLEEF_PRIVATE_HEADERS + [
-        "src/libm/sleefsimddp.c",
-        "src/libm/sleefsimdsp.c",
-        ":renamesse2_h",
-    ],
-    copts = SLEEF_PRIVATE_INCLUDES + SLEEF_COPTS + SLEEF_COMMON_TARGET_COPTS + [
-        "-DDORENAME=1",
-        "-DENABLE_SSE2=1",
-        "-msse2",
-    ],
-    linkstatic = True,
-    visibility = SLEEF_VISIBILITY,
-    alwayslink = True,
-)
-
-sleef_cc_library(
-    name = "sleefsse4",
-    srcs = SLEEF_PRIVATE_HEADERS + [
-        "src/libm/sleefsimddp.c",
-        "src/libm/sleefsimdsp.c",
-        ":renamesse4_h",
-    ],
-    copts = SLEEF_PRIVATE_INCLUDES + SLEEF_COPTS + SLEEF_COMMON_TARGET_COPTS + [
-        "-DDORENAME=1",
-        "-DENABLE_SSE4=1",
-        "-msse4.1",
-    ],
-    linkstatic = True,
-    visibility = SLEEF_VISIBILITY,
-    alwayslink = True,
-)
-
-sleef_cc_library(
-    name = "sleefpurec_scalar",
-    srcs = SLEEF_PRIVATE_HEADERS + [
-        "src/libm/sleefsimddp.c",
-        "src/libm/sleefsimdsp.c",
-        ":renamepurec_scalar_h",
-    ],
-    copts = SLEEF_PRIVATE_INCLUDES + SLEEF_COPTS + SLEEF_COMMON_TARGET_COPTS + [
-        "-DDORENAME=1",
-        "-DENABLE_PUREC_SCALAR=1",
-    ],
-    linkstatic = True,
-    visibility = SLEEF_VISIBILITY,
-    alwayslink = True,
-)
-
-sleef_cc_library(
-    name = "sleefpurecfma_scalar",
-    srcs = SLEEF_PRIVATE_HEADERS + [
-        "src/libm/sleefsimddp.c",
-        "src/libm/sleefsimdsp.c",
-        ":renamepurecfma_scalar_h",
-    ],
-    copts = SLEEF_PRIVATE_INCLUDES + SLEEF_COPTS + SLEEF_COMMON_TARGET_COPTS + [
-        "-DDORENAME=1",
-        "-DENABLE_PURECFMA_SCALAR=1",
-        "-mavx2",
-        "-mfma",
-    ],
-    linkstatic = True,
-    visibility = SLEEF_VISIBILITY,
-    alwayslink = True,
-)
diff --git a/third_party/substitution.bzl b/third_party/substitution.bzl
deleted file mode 100644
index 9310985e50ea5..0000000000000
--- a/third_party/substitution.bzl
+++ /dev/null
@@ -1,81 +0,0 @@
-# This Bazel rules file is derived from https://github.com/tensorflow/tensorflow/blob/master/third_party/common.bzl
-
-# Rule for simple expansion of template files. This performs a simple
-# search over the template file for the keys in substitutions,
-# and replaces them with the corresponding values.
-#
-# Typical usage:
-#   load("/tools/build_rules/template_rule", "template_rule")
-#   template_rule(
-#       name = "ExpandMyTemplate",
-#       src = "my.template",
-#       out = "my.txt",
-#       substitutions = {
-#         "$VAR1": "foo",
-#         "$VAR2": "bar",
-#       }
-#   )
-#
-# Args:
-#   name: The name of the rule.
-#   template: The template file to expand
-#   out: The destination of the expanded file
-#   substitutions: A dictionary mapping strings to their substitutions
-
-def template_rule_impl(ctx):
-    ctx.actions.expand_template(
-        template = ctx.file.src,
-        output = ctx.outputs.out,
-        substitutions = ctx.attr.substitutions,
-    )
-
-template_rule = rule(
-    attrs = {
-        "out": attr.output(mandatory = True),
-        "src": attr.label(
-            mandatory = True,
-            allow_single_file = True,
-        ),
-        "substitutions": attr.string_dict(mandatory = True),
-    },
-    # output_to_genfiles is required for header files.
-    output_to_genfiles = True,
-    implementation = template_rule_impl,
-)
-
-# Header template rule is an extension of template substitution rule
-# That also makes this header a valid dependency for cc_library
-# From https://stackoverflow.com/a/55407399
-def header_template_rule_impl(ctx):
-    ctx.actions.expand_template(
-        template = ctx.file.src,
-        output = ctx.outputs.out,
-        substitutions = ctx.attr.substitutions,
-    )
-    return [
-        # create a provider which says that this
-        # out file should be made available as a header
-        CcInfo(compilation_context = cc_common.create_compilation_context(
-
-            # pass out the include path for finding this header
-            system_includes = depset([ctx.attr.include, ctx.outputs.out.dirname, ctx.bin_dir.path]),
-
-            # and the actual header here.
-            headers = depset([ctx.outputs.out]),
-        )),
-    ]
-
-header_template_rule = rule(
-    attrs = {
-        "include": attr.string(),
-        "out": attr.output(mandatory = True),
-        "src": attr.label(
-            mandatory = True,
-            allow_single_file = True,
-        ),
-        "substitutions": attr.string_dict(mandatory = True),
-    },
-    # output_to_genfiles is required for header files.
-    output_to_genfiles = True,
-    implementation = header_template_rule_impl,
-)
diff --git a/third_party/tensorflow_cuda_bazel_build/cuda/build_defs.bzl b/third_party/tensorflow_cuda_bazel_build/cuda/build_defs.bzl
deleted file mode 100755
index a394b6ce92044..0000000000000
--- a/third_party/tensorflow_cuda_bazel_build/cuda/build_defs.bzl
+++ /dev/null
@@ -1,31 +0,0 @@
-# Macros for building CUDA code.
-def if_cuda(if_true, if_false = []):
-    """Shorthand for select()'ing on whether we're building with CUDA.
-
-    Returns a select statement which evaluates to if_true if we're building
-    with CUDA enabled.  Otherwise, the select statement evaluates to if_false.
-
-    """
-    return select({
-        "@local_config_cuda//cuda:using_clang": if_true,
-        "@local_config_cuda//cuda:using_nvcc": if_true,
-        "//conditions:default": if_false,
-    })
-
-def cuda_default_copts():
-    """Default options for all CUDA compilations."""
-    return if_cuda(["-x", "cuda", "-DGOOGLE_CUDA=1"] + [])
-
-def cuda_is_configured():
-    """Returns true if CUDA was enabled during the configure process."""
-    return True
-
-def if_cuda_is_configured(x):
-    """Tests if the CUDA was enabled during the configure process.
-
-    Unlike if_cuda(), this does not require that we are building with
-    --config=cuda. Used to allow non-CUDA code to depend on CUDA libraries.
-    """
-    if cuda_is_configured():
-        return x
-    return []
diff --git a/third_party/tensorpipe.BUILD b/third_party/tensorpipe.BUILD
deleted file mode 100644
index 5e5b69b4cb4ec..0000000000000
--- a/third_party/tensorpipe.BUILD
+++ /dev/null
@@ -1,178 +0,0 @@
-load("@rules_cc//cc:defs.bzl", "cc_library")
-load("@pytorch//third_party:substitution.bzl", "header_template_rule")
-
-LIBUV_COMMON_SRCS = [
-    "third_party/libuv/src/fs-poll.c",
-    "third_party/libuv/src/idna.c",
-    "third_party/libuv/src/inet.c",
-    "third_party/libuv/src/random.c",
-    "third_party/libuv/src/strscpy.c",
-    "third_party/libuv/src/strtok.c",
-    "third_party/libuv/src/threadpool.c",
-    "third_party/libuv/src/timer.c",
-    "third_party/libuv/src/uv-common.c",
-    "third_party/libuv/src/uv-data-getter-setters.c",
-    "third_party/libuv/src/version.c",
-]
-
-LIBUV_POSIX_SRCS = [
-    "third_party/libuv/src/unix/async.c",
-    "third_party/libuv/src/unix/core.c",
-    "third_party/libuv/src/unix/dl.c",
-    "third_party/libuv/src/unix/fs.c",
-    "third_party/libuv/src/unix/getaddrinfo.c",
-    "third_party/libuv/src/unix/getnameinfo.c",
-    "third_party/libuv/src/unix/loop.c",
-    "third_party/libuv/src/unix/loop-watcher.c",
-    "third_party/libuv/src/unix/pipe.c",
-    "third_party/libuv/src/unix/poll.c",
-    "third_party/libuv/src/unix/process.c",
-    "third_party/libuv/src/unix/random-devurandom.c",
-    "third_party/libuv/src/unix/signal.c",
-    "third_party/libuv/src/unix/stream.c",
-    "third_party/libuv/src/unix/tcp.c",
-    "third_party/libuv/src/unix/thread.c",
-    "third_party/libuv/src/unix/tty.c",
-    "third_party/libuv/src/unix/udp.c",
-]
-
-LIBUV_LINUX_SRCS = LIBUV_POSIX_SRCS + [
-    "third_party/libuv/src/unix/proctitle.c",
-    "third_party/libuv/src/unix/linux.c",
-    "third_party/libuv/src/unix/procfs-exepath.c",
-    "third_party/libuv/src/unix/random-getrandom.c",
-    "third_party/libuv/src/unix/random-sysctl-linux.c",
-]
-
-cc_library(
-    name = "libuv",
-    srcs = LIBUV_COMMON_SRCS + LIBUV_LINUX_SRCS,
-    includes = [
-        "third_party/libuv/include",
-        "third_party/libuv/src",
-    ],
-    hdrs = glob(
-        [
-            "third_party/libuv/include/*.h",
-            "third_party/libuv/include/uv/*.h",
-            "third_party/libuv/src/*.h",
-            "third_party/libuv/src/unix/*.h",
-        ],
-    ),
-    copts = ["-D_GNU_SOURCE"],
-    visibility = ["//visibility:public"],
-)
-
-cc_library(
-    name = "libnop",
-    srcs = [],
-    includes = ["third_party/libnop/include"],
-    hdrs = glob(["third_party/libnop/include/**/*.h"]),
-)
-
-header_template_rule(
-    name = "tensorpipe_cpu_config_header",
-    src = "tensorpipe/config.h.in",
-    out = "tensorpipe/config.h",
-    substitutions = {
-        "#cmakedefine01 TENSORPIPE_HAS_SHM_TRANSPORT": "#define TENSORPIPE_HAS_SHM_TRANSPORT 1",
-        "#cmakedefine01 TENSORPIPE_HAS_IBV_TRANSPORT": "#define TENSORPIPE_HAS_IBV_TRANSPORT 1",
-        "#cmakedefine01 TENSORPIPE_HAS_CMA_CHANNEL": "#define TENSORPIPE_HAS_CMA_CHANNEL 1",
-    },
-)
-
-header_template_rule(
-    name = "tensorpipe_cuda_config_header",
-    src = "tensorpipe/config_cuda.h.in",
-    out = "tensorpipe/config_cuda.h",
-    substitutions = {
-        "#cmakedefine01 TENSORPIPE_HAS_CUDA_IPC_CHANNEL": "#define TENSORPIPE_HAS_CUDA_IPC_CHANNEL 1",
-        "#cmakedefine01 TENSORPIPE_HAS_CUDA_GDR_CHANNEL": "#define TENSORPIPE_HAS_CUDA_GDR_CHANNEL 1",
-    },
-)
-
-# We explicitly list the CUDA headers & sources, and we consider everything else
-# as CPU (using a catch-all glob). This is both because there's fewer CUDA files
-# (thus making it easier to list them exhaustively) and because it will make it
-# more likely to catch a misclassified file: if we forget to mark a file as CUDA
-# we'll try to build it on CPU and that's likely to fail.
-
-TENSORPIPE_CUDA_HEADERS = [
-    "tensorpipe/tensorpipe_cuda.h",
-    "tensorpipe/channel/cuda_basic/*.h",
-    "tensorpipe/channel/cuda_gdr/*.h",
-    "tensorpipe/channel/cuda_ipc/*.h",
-    "tensorpipe/channel/cuda_xth/*.h",
-    "tensorpipe/common/cuda.h",
-    "tensorpipe/common/cuda_buffer.h",
-    "tensorpipe/common/cuda_lib.h",
-    "tensorpipe/common/cuda_loop.h",
-    "tensorpipe/common/nvml_lib.h",
-]
-
-TENSORPIPE_CUDA_SOURCES = [
-    "tensorpipe/channel/cuda_basic/*.cc",
-    "tensorpipe/channel/cuda_gdr/*.cc",
-    "tensorpipe/channel/cuda_ipc/*.cc",
-    "tensorpipe/channel/cuda_xth/*.cc",
-    "tensorpipe/common/cuda_buffer.cc",
-    "tensorpipe/common/cuda_loop.cc",
-]
-
-TENSORPIPE_CPU_HEADERS = glob(
-    [
-        "tensorpipe/*.h",
-        "tensorpipe/channel/*.h",
-        "tensorpipe/channel/*/*.h",
-        "tensorpipe/common/*.h",
-        "tensorpipe/core/*.h",
-        "tensorpipe/transport/*.h",
-        "tensorpipe/transport/*/*.h",
-    ],
-    exclude=TENSORPIPE_CUDA_HEADERS)
-
-TENSORPIPE_CPU_SOURCES = glob(
-    [
-        "tensorpipe/*.cc",
-        "tensorpipe/channel/*.cc",
-        "tensorpipe/channel/*/*.cc",
-        "tensorpipe/common/*.cc",
-        "tensorpipe/core/*.cc",
-        "tensorpipe/transport/*.cc",
-        "tensorpipe/transport/*/*.cc",
-    ],
-    exclude=TENSORPIPE_CUDA_SOURCES)
-
-cc_library(
-    name = "tensorpipe_cpu",
-    srcs = TENSORPIPE_CPU_SOURCES,
-    hdrs = TENSORPIPE_CPU_HEADERS + [":tensorpipe_cpu_config_header"],
-    includes = [
-        ".",
-    ],
-    copts = [
-        "-std=c++17",
-    ],
-    visibility = ["//visibility:public"],
-    deps = [
-        ":libnop",
-        ":libuv",
-    ],
-)
-
-cc_library(
-    name = "tensorpipe_cuda",
-    srcs = glob(TENSORPIPE_CUDA_SOURCES),
-    hdrs = glob(TENSORPIPE_CUDA_HEADERS) + [":tensorpipe_cuda_config_header"],
-    includes = [
-        ".",
-    ],
-    copts = [
-        "-std=c++17",
-    ],
-    visibility = ["//visibility:public"],
-    deps = [
-        ":tensorpipe_cpu",
-        "@cuda",
-    ],
-)
diff --git a/third_party/xnnpack.buck.bzl b/third_party/xnnpack.buck.bzl
index fbc2c56e5b307..4ea930d029ddf 100644
--- a/third_party/xnnpack.buck.bzl
+++ b/third_party/xnnpack.buck.bzl
@@ -43,6 +43,15 @@ def define_xnnpack(third_party, labels = [], XNNPACK_WINDOWS_AVX512F_ENABLED = F
         "-Wno-error=incompatible-pointer-types-discards-qualifiers",
     ]
 
+    WASM_EMSCRIPTEN_COMPILER_FLAGS = select({
+        "DEFAULT": [],
+        "ovr_config//runtime:wasm-emscripten": [
+            "-pthread",
+            "-matomics",
+            "-mbulk-memory",
+        ],
+    })
+
     XNN_COMMON_MICROKERNEL_EXPORTED_DEPS = [
         ":interface",
         third_party("FP16"),
@@ -73,7 +82,7 @@ def define_xnnpack(third_party, labels = [], XNNPACK_WINDOWS_AVX512F_ENABLED = F
         apple_sdks = (IOS, MACOSX),
         compiler_flags = [
             "-O2",
-        ],
+        ] + WASM_EMSCRIPTEN_COMPILER_FLAGS,
         labels = labels,
         fbandroid_link_whole = True,
         preferred_linkage = "static",
@@ -100,7 +109,7 @@ def define_xnnpack(third_party, labels = [], XNNPACK_WINDOWS_AVX512F_ENABLED = F
         apple_sdks = (IOS, MACOSX),
         compiler_flags = [
             "-O2",
-        ],
+        ] + WASM_EMSCRIPTEN_COMPILER_FLAGS,
         labels = labels,
         fbandroid_link_whole = True,
         preferred_linkage = "static",
@@ -127,7 +136,7 @@ def define_xnnpack(third_party, labels = [], XNNPACK_WINDOWS_AVX512F_ENABLED = F
             "-fno-fast-math",
             "-fno-math-errno",
             "-ffp-contract=off",
-        ],
+        ] + WASM_EMSCRIPTEN_COMPILER_FLAGS,
         labels = labels,
         fbandroid_link_whole = True,
         preferred_linkage = "static",
@@ -138,6 +147,29 @@ def define_xnnpack(third_party, labels = [], XNNPACK_WINDOWS_AVX512F_ENABLED = F
         deps = XNN_COMMON_MICROKERNEL_EXPORTED_DEPS,
     )
 
+    fb_xplat_cxx_library(
+        name = "ukernels_wasm",
+        srcs = select({
+            "DEFAULT": [],
+            "ovr_config//runtime:wasm-emscripten": prod_srcs_for_arch_wrapper("wasm"),
+        }),
+        headers = get_xnnpack_headers(),
+        header_namespace = "",
+        apple_sdks = (IOS, MACOSX),
+        compiler_flags = [
+            "-O2",
+            "-fno-fast-math",
+            "-fno-math-errno",
+            "-ffp-contract=off",
+        ] + WASM_EMSCRIPTEN_COMPILER_FLAGS,
+        labels = labels,
+        fbandroid_link_whole = True,
+        preferred_linkage = "static",
+        preprocessor_flags = XNN_COMMON_PREPROCESSOR_FLAGS,
+        visibility = ["PUBLIC"],
+        deps = XNN_COMMON_MICROKERNEL_EXPORTED_DEPS,
+    )
+
     fb_xplat_cxx_library(
         name = "ukernels_sse",
         srcs = select({
@@ -1819,6 +1851,7 @@ def define_xnnpack(third_party, labels = [], XNNPACK_WINDOWS_AVX512F_ENABLED = F
             "ovr_config//runtime:arm64-linux-ubuntu-neon": [":arm64_lib"],
             "ovr_config//runtime:fbcode-arm64": [":arm64_lib"],
             "ovr_config//runtime:platform010": [":x86_and_x86_64_lib"],
+            "ovr_config//runtime:wasm-emscripten": [":ukernels_wasm"],
         }),
     )
 
@@ -1843,7 +1876,7 @@ def define_xnnpack(third_party, labels = [], XNNPACK_WINDOWS_AVX512F_ENABLED = F
         compiler_flags = select({
             "DEFAULT": [],
             "ovr_config//os:macos": ["-fvisibility=default"],
-        }),
+        }) + WASM_EMSCRIPTEN_COMPILER_FLAGS,
         platforms = (APPLE, ANDROID, CXX, WINDOWS),
         preprocessor_flags = XNN_COMMON_PREPROCESSOR_FLAGS + [
             "-DXNN_NO_Q8_OPERATORS",
diff --git a/third_party/xpu.txt b/third_party/xpu.txt
index 8969f51e86378..26cdc6799705c 100644
--- a/third_party/xpu.txt
+++ b/third_party/xpu.txt
@@ -1 +1 @@
-de4f698b84142e660d5238e02e067182e39641ca
+62b793fed8e0a708551d451712af635b6255b322
diff --git a/tools/amd_build/build_amd.py b/tools/amd_build/build_amd.py
index 14c4470ee1db2..7f9289cae498d 100755
--- a/tools/amd_build/build_amd.py
+++ b/tools/amd_build/build_amd.py
@@ -109,6 +109,7 @@
     "torch/*",
     "tools/autograd/templates/python_variable_methods.cpp",
     "torch/csrc/stable/*",
+    "test/cpp/c10d/*",
 ]
 
 includes = [os.path.join(proj_dir, include) for include in includes]
@@ -225,8 +226,8 @@ def remove_hcc(line: str) -> str:
 
 if not buck_build:
     mslk_original = mslk_dir / "tuning_cache.cuh"
-
-    extra_files.append(mslk_original.as_posix())
+    if mslk_original.exists():
+        extra_files.append(mslk_original.as_posix())
 
 # TODO Remove once the following submodules are updated to use hipify v2
 hipify_v1_to_v2_files = [
@@ -300,6 +301,10 @@ def hipify_v1_to_v2(line: str) -> str:
     # only update the file if it changes or doesn't exist
     do_write = True
     src_lines = None
+
+    if not mslk_move_src.exists():
+        _error = f"Error: Source file {mslk_move_src} does not exist"
+        sys.exit(_error)
     with open(mslk_move_src) as src:
         src_lines = src.readlines()
     if os.path.exists(mslk_move_dst):
diff --git a/tools/autograd/derivatives.yaml b/tools/autograd/derivatives.yaml
index d550b6c590835..fa611a88889b0 100644
--- a/tools/autograd/derivatives.yaml
+++ b/tools/autograd/derivatives.yaml
@@ -769,13 +769,22 @@
 - name: grid_sampler_2d(Tensor input, Tensor grid, int interpolation_mode, int padding_mode, bool align_corners) -> Tensor
   input, grid: "grad.defined() ? grid_sampler_2d_backward(grad, input, grid, interpolation_mode, padding_mode, align_corners, grad_input_mask) : std::tuple<Tensor, Tensor>()"
 
+- name: grid_sampler_2d_backward(Tensor grad_output, Tensor input, Tensor grid, int interpolation_mode, int padding_mode, bool align_corners, bool[2] output_mask) -> (Tensor, Tensor)
+  grad_output, input, grid: grid_sampler_2d_double_backward(grads[0], grads[1], grad_output, input, grid, interpolation_mode, padding_mode, align_corners, grad_input_mask)
+
 - name: grid_sampler_3d(Tensor input, Tensor grid, int interpolation_mode, int padding_mode, bool align_corners) -> Tensor
   input, grid: "grad.defined() ? grid_sampler_3d_backward(grad, input, grid, interpolation_mode, padding_mode, align_corners, grad_input_mask) : std::tuple<Tensor, Tensor>()"
 
+- name: grid_sampler_3d_backward(Tensor grad_output, Tensor input, Tensor grid, int interpolation_mode, int padding_mode, bool align_corners, bool[2] output_mask) -> (Tensor, Tensor)
+  grad_output, input, grid: grid_sampler_3d_double_backward(grads[0], grads[1], grad_output, input, grid, interpolation_mode, padding_mode, align_corners, grad_input_mask)
+
 # See NOTE [ grid_sample CPU fallback ]
 - name: _grid_sampler_2d_cpu_fallback(Tensor input, Tensor grid, int interpolation_mode, int padding_mode, bool align_corners) -> Tensor
   input, grid: "grad.defined() ? _grid_sampler_2d_cpu_fallback_backward(grad, input, grid, interpolation_mode, padding_mode, align_corners) : std::tuple<Tensor, Tensor>()"
 
+- name: _grid_sampler_2d_cpu_fallback_backward(Tensor grad_output, Tensor input, Tensor grid, int interpolation_mode, int padding_mode, bool align_corners) -> (Tensor, Tensor)
+  grad_output, input, grid: grid_sampler_2d_double_backward(grads[0], grads[1], grad_output, input, grid, interpolation_mode, padding_mode, align_corners, grad_input_mask)
+
 - name: gt_.Scalar(Tensor(a!) self, Scalar other) -> Tensor(a!)
   self: zeros_like(self)
   result: self_t.zero_()
@@ -1201,6 +1210,11 @@
   self: scale_grad_by_count(restore_reduced_dims(grad, dim, keepdim), restore_reduced_dims(result, dim, keepdim) == self, dim)
   result: amaxamin_jvp(self_p, self_t, result, dim, keepdim)
 
+- name: aminmax(Tensor self, *, int? dim=None, bool keepdim=False) -> (Tensor min, Tensor max)
+  self: aminmax_backward(self, dim, keepdim, grad_min, grad_max, min, max)
+  min: aminmax_jvp(self_p, self_t, min, dim, keepdim)
+  max: aminmax_jvp(self_p, self_t, max, dim, keepdim)
+
 - name: mm(Tensor self, Tensor mat2) -> Tensor
   self: mm_mat1_backward(grad, mat2, self.sym_sizes(), self.sym_strides(), self.layout(), 1)
   mat2: mm_mat2_backward(grad, self, mat2.sym_sizes(), mat2.sym_strides(), mat2.layout(), 1)
@@ -2277,6 +2291,10 @@
   self: _upsample_bicubic2d_aa_backward_symint(grad, output_size, self.sym_sizes(), align_corners, scales_h, scales_w)
   result: auto_linear
 
+- name: _upsample_lanczos2d_aa(Tensor self, SymInt[2] output_size, bool align_corners, float? scales_h=None, float? scales_w=None) -> Tensor
+  self: _upsample_lanczos2d_aa_backward_symint(grad, output_size, self.sym_sizes(), align_corners, scales_h, scales_w)
+  result: auto_linear
+
 - name: upsample_trilinear3d(Tensor self, SymInt[3] output_size, bool align_corners, float? scales_d=None, float? scales_h=None, float? scales_w=None) -> Tensor
   self: upsample_trilinear3d_backward_symint(grad, output_size, self.sym_sizes(), align_corners, scales_d, scales_h, scales_w)
   result: auto_linear
@@ -2676,6 +2694,10 @@
   grad_output: _upsample_bicubic2d_aa_symint(grad, output_size, align_corners, scales_h, scales_w)
   result: auto_linear
 
+- name: _upsample_lanczos2d_aa_backward(Tensor grad_output, SymInt[2] output_size, SymInt[4] input_size, bool align_corners, float? scales_h=None, float? scales_w=None) -> Tensor
+  grad_output: _upsample_lanczos2d_aa_symint(grad, output_size, align_corners, scales_h, scales_w)
+  result: auto_linear
+
 - name: upsample_trilinear3d_backward(Tensor grad_output, SymInt[3] output_size, SymInt[5] input_size, bool align_corners, float? scales_d=None, float? scales_h=None, float? scales_w=None) -> Tensor
   grad_output: upsample_trilinear3d_symint(grad, output_size, align_corners, scales_d, scales_h, scales_w)
   result: auto_linear
@@ -2733,6 +2755,9 @@
 - name: cudnn_grid_sampler(Tensor self, Tensor grid) -> Tensor output
   self, grid: "grad.defined() ? cudnn_grid_sampler_backward(self, grid, grad) : std::tuple<Tensor, Tensor>()"
 
+- name: cudnn_grid_sampler_backward(Tensor self, Tensor grid, Tensor grad_output) -> (Tensor grad_self, Tensor grad_grid)
+  grad_output, self, grid: grid_sampler_2d_double_backward(grads[0], grads[1], grad_output, self, grid, 0, 0, true, grad_input_mask)
+
 - name: cudnn_affine_grid_generator(Tensor theta, int N, int C, int H, int W) -> Tensor grid
   theta: cudnn_affine_grid_generator_backward(grad, N, C, H, W)
 
@@ -2902,7 +2927,7 @@
   output_differentiability: [True, False]
   query, key, value: _scaled_dot_product_flash_attention_for_cpu_backward(grad, query, key, value, output, logsumexp, dropout_p, is_causal, attn_mask, scale)
 
-- name: _flash_attention_forward(Tensor query, Tensor key, Tensor value, Tensor? cum_seq_q, Tensor? cum_seq_k, SymInt max_q, SymInt max_k, float dropout_p, bool is_causal, bool return_debug_mask, *, float? scale=None, SymInt? window_size_left=None, SymInt? window_size_right=None, Tensor? seqused_k=None, Tensor? alibi_slopes=None) -> (Tensor output, Tensor softmax_logsumexp, Tensor rng_state, Tensor unused, Tensor debug_attn_mask)
+- name: _flash_attention_forward(Tensor query, Tensor key, Tensor value, Tensor? cum_seq_q, Tensor? cum_seq_k, SymInt max_q, SymInt max_k, float dropout_p, bool is_causal, bool return_debug_mask, *, float? scale=None, SymInt? window_size_left=None, SymInt? window_size_right=None, Tensor? seqused_k=None, Tensor? alibi_slopes=None, Tensor? block_table=None, int? num_splits=None) -> (Tensor output, Tensor softmax_logsumexp, Tensor rng_state, Tensor unused, Tensor debug_attn_mask)
   output_differentiability: [True, False, False, False, False]
   query, key, value: _flash_attention_backward_symint(grad, query, key, value, output, softmax_logsumexp, cum_seq_q, cum_seq_k, max_q, max_k, dropout_p, is_causal, rng_state, unused, scale, window_size_left, window_size_right)
 
@@ -3210,6 +3235,10 @@
   self: grad.reshape_symint(self.sym_sizes())
   result: auto_linear
 
+- name: narrow_copy(Tensor self, int dim, SymInt start, SymInt length) -> Tensor
+  self: slice_backward_wrapper(grad, self.sym_sizes(), dim, start, start + length, 1)
+  result: auto_linear
+
 # note(crcrpar): `torchgen/api/autograd` logic would unwantedly replace substrings of `self` and `other` of function names.
 - name: _foreach_div.List(Tensor[] self, Tensor[] other) -> Tensor[]
   self: div_tensor_self_backward(grads[i], other[i], self[i].scalar_type())
diff --git a/tools/autograd/gen_python_functions.py b/tools/autograd/gen_python_functions.py
index a18485a902d35..09688e561e525 100644
--- a/tools/autograd/gen_python_functions.py
+++ b/tools/autograd/gen_python_functions.py
@@ -92,10 +92,15 @@
 _SKIP_PYTHON_BINDINGS = [
     "alias",
     "contiguous",
+    "dim",
+    "get_device",
+    "is_contiguous",
     "is_cuda",
     "is_sparse",
     "is_sparse_csr",
+    "numel",
     "size",
+    "storage_offset",
     "stride",
     "sym_is_contiguous",
     "sym_size",
diff --git a/tools/autograd/gen_variable_type.py b/tools/autograd/gen_variable_type.py
index 7bacdd3461346..6e45f29a8232a 100644
--- a/tools/autograd/gen_variable_type.py
+++ b/tools/autograd/gen_variable_type.py
@@ -262,6 +262,7 @@
     "alias",
     "atan",
     "ldexp",
+    "linear",
     "log",
     "log10",
     "log1p",
@@ -376,6 +377,7 @@
     "_neg_view",
     "_reshape_alias",
     "_reshape_copy",
+    "narrow_copy",
     "_linalg_det",
     "lu_solve",
     "linalg_solve_triangular",
diff --git a/tools/autograd/templates/python_variable_methods.cpp b/tools/autograd/templates/python_variable_methods.cpp
index bfc5b80835c4b..2260f8cb2245f 100644
--- a/tools/autograd/templates/python_variable_methods.cpp
+++ b/tools/autograd/templates/python_variable_methods.cpp
@@ -867,7 +867,7 @@ static PyObject * THPVariable_is_contiguous(PyObject* self_, PyObject* args, PyO
   auto r = parser.parse(self_, args, kwargs, parsed_args);
 
   if(r.has_torch_function()){
-    return handle_torch_function(r, self_, args, kwargs, PyObject_Type(self_), "torch.Tensor");
+    return handle_torch_function(r, self_, args, kwargs, reinterpret_cast<PyObject*>(Py_TYPE(self_)), "torch.Tensor");
   }
 
   auto memory_format = r.memoryformat(0);
diff --git a/tools/bazel_tools/BUILD.bazel b/tools/bazel_tools/BUILD.bazel
deleted file mode 100644
index f4c37fb7389bd..0000000000000
--- a/tools/bazel_tools/BUILD.bazel
+++ /dev/null
@@ -1,5 +0,0 @@
-sh_binary(
-    name = "shellwrap",
-    srcs = ["shellwrap.sh"],
-    visibility = ["//visibility:public"],
-)
diff --git a/tools/bazel_tools/shellwrap.sh b/tools/bazel_tools/shellwrap.sh
deleted file mode 100755
index 712788ae09e06..0000000000000
--- a/tools/bazel_tools/shellwrap.sh
+++ /dev/null
@@ -1,58 +0,0 @@
-#!/bin/bash
-
-# This script is helpful in entering an interactive shell from a bazel build
-# before running a given bazel executable.
-# This can provide a quick way to explore the sandbox directory and filesystem.
-# Typical use is with
-#
-#     bazel run --run_under=//tools/bazel:shell_wrapper //:target
-#     OR
-#     bazel run --config=shell //:target
-
-shell='/bin/bash'
-rcfile='/tmp/pytorch_bazel_tools_shellwrap'
-while [[ $# -gt 0 ]] ; do
-    case "$1" in
-        --shell_bin_path)
-            # path for the shell executable
-            shell="$2"
-            shift 2
-            ;;
-        --rcfile)
-            # path for the file used to write the environment
-            rcfile="$2"
-            shift 2
-            ;;
-        *)
-            # remaining arguments are part of the command for execution
-            break
-            ;;
-    esac
-done
-
-if ! tty -s; then
-    echo 'A tty is not available.'
-    echo "Use \`bazel run\`, not \`bazel test\`."
-    exit 1
-fi
-
-NOCOLOR='\033[0m'
-YELLOW='\033[1;33m'
-
-# store the environment in a file
-export PYTORCH_SHELL_COMMAND=$*
-echo "alias run=\"$*\"" > "$rcfile"
-echo "PS1='\s-\v\$ '" >> "$rcfile"
-
-echo =====
-# print the execution command (command is yellow)
-echo -e "alias run=${YELLOW}$PYTORCH_SHELL_COMMAND${NOCOLOR}"
-echo =====
-
-echo "Entering interactive shell at the execution root:"
-
-# quote escape all the arguments to use as a single input string
-cmd="'$shell' --noprofile --rcfile '$rcfile'"
-
-# run the command in a script pseudo terminal and dump to null
-/usr/bin/script -c "$cmd" -q /dev/null
diff --git a/tools/clean.py b/tools/clean.py
new file mode 100644
index 0000000000000..e4f2b31adb8eb
--- /dev/null
+++ b/tools/clean.py
@@ -0,0 +1,23 @@
+import glob
+import os
+import shutil
+from pathlib import Path
+
+
+def clean():
+    """Clean, that is remove all files in .gitignore except in the NOT-CLEAN-FILES section."""
+    ignores = Path(".gitignore").read_text(encoding="utf-8")
+    for wildcard in filter(None, ignores.splitlines()):
+        if wildcard.strip().startswith("#"):
+            if "BEGIN NOT-CLEAN-FILES" in wildcard:
+                # Marker is found and stop reading .gitignore.
+                break
+            # Ignore lines which begin with '#'.
+        else:
+            # Don't remove absolute paths from the system
+            wildcard = wildcard.lstrip("./")
+            for filename in glob.iglob(wildcard):
+                try:
+                    os.remove(filename)
+                except OSError:
+                    shutil.rmtree(filename, ignore_errors=True)
diff --git a/tools/config/BUILD b/tools/config/BUILD
deleted file mode 100644
index ba13eda2bba7b..0000000000000
--- a/tools/config/BUILD
+++ /dev/null
@@ -1,41 +0,0 @@
-load("@bazel_skylib//lib:selects.bzl", "selects")
-
-config_setting(
-    name = "cuda",
-    define_values = {
-        "cuda": "true",
-    },
-)
-
-# Even when building with --config=cuda, host targets should be built with cuda disabled
-# as these targets will run on CI machines that have no GPUs.
-selects.config_setting_group(
-    name = "cuda_enabled_and_capable",
-    match_all = [
-        ":cuda",
-    ],
-)
-
-# Configures the system to build with cuda using clang.
-config_setting(
-    name = "cuda_clang",
-    define_values = {
-        "cuda_clang": "true",
-    },
-)
-
-# Indicates that cuda code should be compiled with nvcc
-# Mostly exists to support _analysis_ of tensorflow; more work is needed to actually make this
-# setting work.
-config_setting(
-    name = "cuda_nvcc",
-    define_values = {
-        "cuda_nvcc": "true",
-    },
-)
-
-config_setting(
-    name = "thread_sanitizer",
-    define_values = {"thread_sanitizer": "1"},
-    visibility = ["//visibility:public"],
-)
diff --git a/tools/config/defs.bzl b/tools/config/defs.bzl
deleted file mode 100644
index f8a1e9dc16f26..0000000000000
--- a/tools/config/defs.bzl
+++ /dev/null
@@ -1,65 +0,0 @@
-"""
- Macros for selecting with / without various GPU libraries.  Most of these are meant to be used
- directly by tensorflow in place of their build's own configure.py + bazel-gen system.
-"""
-
-load("@bazel_skylib//lib:selects.bzl", "selects")
-
-def if_cuda(if_true, if_false = []):
-    """Helper for selecting based on the whether CUDA is configured. """
-    return selects.with_or({
-        "@//tools/config:cuda_enabled_and_capable": if_true,
-        "//conditions:default": if_false,
-    })
-
-def if_tensorrt(if_true, if_false = []):
-    """Helper for selecting based on the whether TensorRT is configured. """
-    return select({
-        "//conditions:default": if_false,
-    })
-
-def if_rocm(if_true, if_false = []):
-    """Helper for selecting based on the whether ROCM is configured. """
-    return select({
-        "//conditions:default": if_false,
-    })
-
-def if_sycl(if_true, if_false = []):
-    """Helper for selecting based on the whether SYCL/ComputeCPP is configured."""
-
-    # NOTE: Tensorflow expects some strange behavior (see their if_sycl) if we
-    # actually plan on supporting this at some point.
-    return select({
-        "//conditions:default": if_false,
-    })
-
-def if_ccpp(if_true, if_false = []):
-    """Helper for selecting based on the whether ComputeCPP is configured. """
-    return select({
-        "//conditions:default": if_false,
-    })
-
-def cuda_default_copts():
-    return if_cuda(["-DGOOGLE_CUDA=1"])
-
-def cuda_default_features():
-    return if_cuda(["-per_object_debug_info", "-use_header_modules", "cuda_clang"])
-
-def rocm_default_copts():
-    return if_rocm(["-x", "rocm"])
-
-def rocm_copts(opts = []):
-    return rocm_default_copts() + if_rocm(opts)
-
-def cuda_is_configured():
-    # FIXME(dcollins): currently only used by tensorflow's xla stuff, which we aren't building.  However bazel
-    # query hits it so this needs to be defined.  Because bazel doesn't actually resolve config at macro expansion
-    # time, `select` can't be used here (since xla expects lists of strings and not lists of select objects).
-    # Instead, the xla build rules must be rewritten to use `if_cuda_is_configured`
-    return False
-
-def if_cuda_is_configured(x):
-    return if_cuda(x, [])
-
-def if_rocm_is_configured(x):
-    return if_rocm(x, [])
diff --git a/tools/create_worktree.py b/tools/create_worktree.py
new file mode 100755
index 0000000000000..cf8ef48c954bc
--- /dev/null
+++ b/tools/create_worktree.py
@@ -0,0 +1,286 @@
+#!/usr/bin/env python3
+"""Create or remove git worktrees with submodules cloned locally.
+
+This avoids fetching submodules from remote, which is slow for large repos
+like PyTorch. Instead, each submodule is cloned directly from the local
+checkout's git object store, so no network access is needed.
+
+Usage:
+    python tools/create_worktree.py                  # pytorch-worktree-1
+    python tools/create_worktree.py my-worktree      # custom name
+    python tools/create_worktree.py remove my-worktree  # force-remove
+"""
+
+import argparse
+import configparser
+import subprocess
+import sys
+from pathlib import Path
+
+
+def get_repo_root() -> Path:
+    result = subprocess.run(
+        ["git", "rev-parse", "--show-toplevel"],
+        capture_output=True,
+        text=True,
+        check=True,
+    )
+    return Path(result.stdout.strip())
+
+
+def get_existing_worktrees(repo_root: Path) -> list[str]:
+    result = subprocess.run(
+        ["git", "worktree", "list", "--porcelain"],
+        capture_output=True,
+        text=True,
+        check=True,
+        cwd=repo_root,
+    )
+    names = []
+    for line in result.stdout.splitlines():
+        if line.startswith("worktree "):
+            names.append(Path(line.split(" ", 1)[1]).name)
+    return names
+
+
+def next_worktree_name(repo_root: Path) -> str:
+    existing = get_existing_worktrees(repo_root)
+    n = 1
+    while True:
+        name = f"pytorch-worktree-{n}"
+        if name not in existing:
+            return name
+        n += 1
+
+
+def parse_gitmodules(root: Path) -> list[dict[str, str]]:
+    gitmodules = root / ".gitmodules"
+    if not gitmodules.exists():
+        return []
+    config = configparser.ConfigParser()
+    config.read(gitmodules)
+    modules = []
+    for section in config.sections():
+        if section.startswith('submodule "'):
+            path = config.get(section, "path")
+            url = config.get(section, "url")
+            modules.append({"path": path, "url": url})
+    return modules
+
+
+def get_submodule_commit(parent_repo: Path, submodule_path: str) -> str | None:
+    """Get the commit hash a parent repo expects for a submodule."""
+    result = subprocess.run(
+        ["git", "ls-tree", "HEAD", submodule_path],
+        capture_output=True,
+        text=True,
+        cwd=parent_repo,
+    )
+    if result.returncode != 0 or not result.stdout.strip():
+        return None
+    # format: <mode> <type> <hash>\t<path>
+    parts = result.stdout.strip().split()
+    return parts[2] if len(parts) >= 3 else None
+
+
+def resolve_git_dir(submodule_worktree: Path) -> Path | None:
+    """Resolve the actual git object directory for a submodule.
+
+    Submodules can have .git as either a directory (standalone clone) or a
+    file containing 'gitdir: <path>' (gitlink). This resolves to the actual
+    git directory in both cases.
+    """
+    git_path = submodule_worktree / ".git"
+    if not git_path.exists():
+        return None
+    if git_path.is_dir():
+        return git_path
+    # gitlink: read the target
+    content = git_path.read_text().strip()
+    if content.startswith("gitdir: "):
+        target = content[len("gitdir: ") :]
+        resolved = (submodule_worktree / target).resolve()
+        if resolved.exists():
+            return resolved
+    return None
+
+
+def clone_submodule_recursive(
+    worktree_root: Path,
+    source_root: Path,
+    submodule_path: str,
+    url: str,
+    depth: int = 0,
+) -> None:
+    prefix = "  " * depth
+    source_sub = source_root / submodule_path
+    worktree_sub = worktree_root / submodule_path
+
+    git_dir = resolve_git_dir(source_sub)
+    if git_dir is None:
+        print(f"{prefix}  skipping {submodule_path} (not initialized in source)")
+        return
+
+    commit = get_submodule_commit(source_root, submodule_path)
+
+    print(f"{prefix}  cloning {submodule_path}...")
+
+    # The worktree checkout creates empty dirs for submodule paths; remove them
+    # so git clone can create the directory itself.
+    if worktree_sub.exists():
+        try:
+            worktree_sub.rmdir()
+        except OSError:
+            pass
+
+    # Clone from the resolved git directory — works for both standalone .git
+    # dirs and gitlinks pointing into .git/modules/. No network needed.
+    subprocess.run(
+        ["git", "clone", str(git_dir), str(worktree_sub)],
+        check=True,
+        capture_output=True,
+        text=True,
+    )
+
+    # Checkout the exact commit the parent repo expects.
+    if commit:
+        subprocess.run(
+            ["git", "checkout", "--quiet", commit],
+            cwd=worktree_sub,
+            check=True,
+            capture_output=True,
+            text=True,
+        )
+
+    # Point the remote back to the real URL so future fetches work.
+    subprocess.run(
+        ["git", "remote", "set-url", "origin", url],
+        cwd=worktree_sub,
+        check=True,
+        capture_output=True,
+        text=True,
+    )
+
+    # Register nested submodules in this repo's config so that
+    # `git submodule status --recursive` recognizes them.
+    subprocess.run(
+        ["git", "submodule", "init"],
+        cwd=worktree_sub,
+        capture_output=True,
+        text=True,
+    )
+
+    # Recurse into nested submodules
+    nested = parse_gitmodules(worktree_sub)
+    for mod in nested:
+        try:
+            clone_submodule_recursive(
+                worktree_sub, source_sub, mod["path"], mod["url"], depth + 1
+            )
+        except subprocess.CalledProcessError as e:
+            print(
+                f"{prefix}    WARNING: nested {mod['path']}: {e}",
+                file=sys.stderr,
+            )
+
+
+def cmd_create(args: argparse.Namespace) -> None:
+    repo_root = get_repo_root()
+    parent_dir = Path(args.parent_dir) if args.parent_dir else repo_root.parent
+    name = args.name or next_worktree_name(repo_root)
+    worktree_path = parent_dir / name
+
+    print(f"Creating worktree at {worktree_path}")
+
+    cmd = ["git", "worktree", "add", str(worktree_path), args.commit]
+    subprocess.run(cmd, check=True, cwd=repo_root)
+
+    modules = parse_gitmodules(repo_root)
+    if modules:
+        print(f"Cloning {len(modules)} submodules from local checkout...")
+        for mod in modules:
+            try:
+                clone_submodule_recursive(
+                    worktree_path, repo_root, mod["path"], mod["url"]
+                )
+            except subprocess.CalledProcessError as e:
+                print(
+                    f"  WARNING: failed to clone {mod['path']}: {e}",
+                    file=sys.stderr,
+                )
+                if e.stderr:
+                    print(f"    {e.stderr.strip()}", file=sys.stderr)
+
+    print(f"\nWorktree ready at {worktree_path}")
+
+
+def cmd_remove(args: argparse.Namespace) -> None:
+    repo_root = get_repo_root()
+    parent_dir = Path(args.parent_dir) if args.parent_dir else repo_root.parent
+    worktree_path = parent_dir / args.name
+
+    print(f"Removing worktree at {worktree_path}")
+    subprocess.run(
+        ["git", "worktree", "remove", "--force", str(worktree_path)],
+        check=True,
+        cwd=repo_root,
+    )
+    print("Done.")
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(
+        description="Create or remove git worktrees with locally-cloned submodules."
+    )
+    subparsers = parser.add_subparsers(dest="command")
+
+    # "create" subcommand (also the default when no subcommand given)
+    create_parser = subparsers.add_parser(
+        "create", help="Create a new worktree (default)"
+    )
+    create_parser.add_argument(
+        "name",
+        nargs="?",
+        default=None,
+        help="Worktree directory name (default: pytorch-worktree-$n)",
+    )
+    create_parser.add_argument(
+        "--commit",
+        default="HEAD",
+        help="Commit/branch to check out in the worktree (default: HEAD)",
+    )
+    create_parser.add_argument(
+        "--parent-dir",
+        default=None,
+        help="Parent directory for the worktree (default: parent of repo root)",
+    )
+
+    # "remove" subcommand
+    remove_parser = subparsers.add_parser("remove", help="Force-remove a worktree")
+    remove_parser.add_argument(
+        "name",
+        help="Worktree directory name to remove",
+    )
+    remove_parser.add_argument(
+        "--parent-dir",
+        default=None,
+        help="Parent directory of the worktree (default: parent of repo root)",
+    )
+
+    args = parser.parse_args()
+
+    # Default to "create" when no subcommand is given
+    if args.command is None:
+        args.command = "create"
+        args.name = None
+        args.commit = "HEAD"
+        args.parent_dir = None
+
+    if args.command == "create":
+        cmd_create(args)
+    elif args.command == "remove":
+        cmd_remove(args)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/tools/dynamo/verify_dynamo.py b/tools/dynamo/verify_dynamo.py
index 5591baa5552ec..598b09f173dd9 100644
--- a/tools/dynamo/verify_dynamo.py
+++ b/tools/dynamo/verify_dynamo.py
@@ -6,7 +6,7 @@
 import warnings
 
 
-MIN_CUDA_VERSION = "11.6"
+MIN_CUDA_VERSION = "12.1"
 MIN_ROCM_VERSION = "5.4"
 MIN_PYTHON_VERSION = (3, 10)
 
diff --git a/tools/experimental/torchfuzz/checks.py b/tools/experimental/torchfuzz/checks.py
index 5b7b2e9da0e99..e1fc7452c1c63 100644
--- a/tools/experimental/torchfuzz/checks.py
+++ b/tools/experimental/torchfuzz/checks.py
@@ -25,6 +25,22 @@ def codegen(self, args_tuple: str) -> list[str]:
         ]
 
 
+class EagerVsFullGraphDynamicCompileWithBackwardCheck(Check):
+    """Check that runs eager then fullgraph+dynamic compilation with backward pass."""
+
+    def codegen(self, args_tuple: str) -> list[str]:
+        return [
+            f"args = {args_tuple}",
+            "result_original = fuzzed_program(*args)",
+            "result_original.sum().backward()",
+            "print('✅ eager + backward success')",
+            "compiled_program = torch.compile(fuzzed_program, fullgraph=True, dynamic=True)",
+            "result_compiled = compiled_program(*args)",
+            "result_compiled.sum().backward()",
+            "print('✅ compile + backward success')",
+        ]
+
+
 class EagerVsFullGraphDynamicCompileWithNumericsCheck(Check):
     """Check that runs eager and compiled, compares forward numerics."""
 
diff --git a/tools/experimental/torchfuzz/codegen.py b/tools/experimental/torchfuzz/codegen.py
index da41392623bd9..b8a9c313ff1e8 100644
--- a/tools/experimental/torchfuzz/codegen.py
+++ b/tools/experimental/torchfuzz/codegen.py
@@ -1,5 +1,6 @@
 # mypy: ignore-errors
 import os
+import random
 
 import torch
 
@@ -501,6 +502,176 @@ def epilogue_codegen(self):
         return []
 
 
+class StreamFuzzTemplate(DefaultFuzzTemplate):
+    """Template that wraps operations in random CUDA stream contexts.
+
+    Reuses the same operator set as DefaultFuzzTemplate but partitions non-leaf
+    operations across 2-3 CUDA streams, inserting proper wait_stream
+    synchronization between dependent operations on different streams.
+    """
+
+    def __init__(self):
+        super().__init__()
+        from torchfuzz.checks import EagerVsFullGraphDynamicCompileWithBackwardCheck
+
+        self.check = EagerVsFullGraphDynamicCompileWithBackwardCheck()
+
+    def imports_codegen(self):
+        return [
+            "import torch",
+        ]
+
+    def flags_codegen(self):
+        return [
+            "torch.set_default_device('cuda')",
+            "torch._dynamo.config.capture_scalar_outputs = True",
+        ]
+
+    def args_codegen(self, arg_operations):
+        """Generate args with requires_grad=True on float tensors.
+
+        This ensures the backward pass traces through stream-wrapped operations,
+        exercising Inductor's stream handling in the backward graph.
+        """
+        code_lines = super().args_codegen(arg_operations)
+        if arg_operations:
+            for i, (node_id, spec) in enumerate(arg_operations):
+                if isinstance(spec, TensorSpec) and spec.dtype in [
+                    torch.float32,
+                    torch.float64,
+                    torch.float16,
+                    torch.bfloat16,
+                ]:
+                    code_lines.append(f"arg_{i} = arg_{i}.requires_grad_(True)")
+        return code_lines
+
+    @staticmethod
+    def wrap_body_with_streams(
+        generated_code_lines: list[str],
+        graph: OperationGraph,
+    ) -> list[str]:
+        """Wrap generated function body lines with CUDA stream contexts.
+
+        Assigns each non-leaf operation to one of 2-3 random streams, wraps each
+        in ``with torch.cuda.stream(sN):``, and inserts ``wait_stream`` calls
+        between dependent operations on different streams.
+        """
+        topo_order = graph.get_topological_order()
+
+        # Identify leaf vs non-leaf node ids
+        leaf_ids = set()
+        non_leaf_ids = []
+        for nid in topo_order:
+            node = graph.nodes[nid]
+            if (
+                node.op_name == "arg"
+                or node.op_name.startswith("arg_")
+                or node.op_name == "constant"
+            ):
+                leaf_ids.add(nid)
+            else:
+                non_leaf_ids.append(nid)
+
+        if not non_leaf_ids:
+            return generated_code_lines
+
+        num_streams = random.randint(2, 3)
+        stream_names = [f"s{i + 1}" for i in range(num_streams)]
+
+        # Decide sync strategy: wait_stream or event-based (record + wait_event)
+        use_events = random.choice([True, False])
+        event_counter = 0
+
+        # Assign each non-leaf node to a random stream
+        node_stream: dict[str, str] = {}
+        for nid in non_leaf_ids:
+            node_stream[nid] = random.choice(stream_names)
+
+        # Build a mapping from node_id -> the original code lines for that node.
+        # Each node produces lines prefixed with "    " (4-space indent for the
+        # function body).  We identify nodes by their ``var_{node_id} =`` pattern.
+        node_lines: dict[str, list[str]] = {}
+        current_node: str | None = None
+        current_buf: list[str] = []
+
+        for line in generated_code_lines:
+            stripped = line.strip()
+            # Detect lines like "var_node_3 = ..." or "var_node_3, _ = ..."
+            matched_node = None
+            for nid in topo_order:
+                if stripped.startswith((f"var_{nid} =", f"var_{nid},")):
+                    matched_node = nid
+                    break
+
+            if matched_node is not None:
+                # Flush previous node buffer
+                if current_node is not None:
+                    node_lines[current_node] = current_buf
+                current_node = matched_node
+                current_buf = [line]
+            else:
+                current_buf.append(line)
+
+        # Flush last node
+        if current_node is not None:
+            node_lines[current_node] = current_buf
+
+        # Rebuild the body with stream contexts and synchronization
+        new_lines: list[str] = []
+
+        # Stream variable declarations at the top of the function body
+        for sname in stream_names:
+            new_lines.append(f"    {sname} = torch.cuda.Stream()")
+
+        for nid in topo_order:
+            lines_for_node = node_lines.get(nid, [])
+            if nid in leaf_ids:
+                # Leaf nodes (args) stay on the default stream
+                new_lines.extend(lines_for_node)
+                continue
+
+            stream = node_stream[nid]
+            node = graph.nodes[nid]
+
+            # Insert synchronization for cross-stream dependencies
+            waited: set[str] = set()
+            for dep_id in node.input_nodes:
+                if dep_id in node_stream and node_stream[dep_id] != stream:
+                    dep_stream = node_stream[dep_id]
+                    if dep_stream not in waited:
+                        if use_events:
+                            ename = f"e{event_counter}"
+                            event_counter += 1
+                            new_lines.append(f"    {ename} = torch.cuda.Event()")
+                            new_lines.append(f"    {ename}.record({dep_stream})")
+                            new_lines.append(f"    {stream}.wait_event({ename})")
+                        else:
+                            new_lines.append(f"    {stream}.wait_stream({dep_stream})")
+                        waited.add(dep_stream)
+
+            # Wrap the operation in a stream context
+            new_lines.append(f"    with torch.cuda.stream({stream}):")
+            for code_line in lines_for_node:
+                # Each line already has 4-space indent; add 4 more for the with block
+                new_lines.append("    " + code_line)
+
+        # Synchronize all streams before the return statement
+        if use_events:
+            for sname in stream_names:
+                ename = f"e{event_counter}"
+                event_counter += 1
+                new_lines.append(f"    {ename} = torch.cuda.Event()")
+                new_lines.append(f"    {ename}.record({sname})")
+                new_lines.append(f"    torch.cuda.current_stream().wait_event({ename})")
+        else:
+            for sname in stream_names:
+                new_lines.append(
+                    f"    torch.cuda.current_stream().wait_stream({sname})"
+                )
+
+        return new_lines
+
+
 class DTensorFuzzPlacementsTemplate(DTensorFuzzTemplate):
     """DTensor template with randomized placements (Replicate, Shard, Partial).
 
@@ -656,6 +827,8 @@ def convert_graph_to_python_code(
         fuzz_template = DTensorFuzzPlacementsTemplate()
     elif template == "unbacked":
         fuzz_template = UnbackedFuzzTemplate()
+    elif template == "streams":
+        fuzz_template = StreamFuzzTemplate()
     else:
         fuzz_template = DefaultFuzzTemplate()
 
@@ -730,6 +903,12 @@ def convert_graph_to_python_code(
         # Track this node's variable
         node_variables[node_id] = (output_var_name, output_spec)
 
+    # Wrap body with stream contexts if using the streams template
+    if template == "streams":
+        generated_code_lines = StreamFuzzTemplate.wrap_body_with_streams(
+            generated_code_lines, operation_graph
+        )
+
     # The final result comes from the root node
     root_node_id = operation_graph.root_node_id
     if root_node_id not in node_variables:
diff --git a/tools/experimental/torchfuzz/fuzzer.py b/tools/experimental/torchfuzz/fuzzer.py
index e1e1be477ea0e..d824fe0c6d0bb 100644
--- a/tools/experimental/torchfuzz/fuzzer.py
+++ b/tools/experimental/torchfuzz/fuzzer.py
@@ -262,7 +262,7 @@ def log(success: bool) -> None:
     )
     parser.add_argument(
         "--template",
-        choices=["default", "dtensor", "dtensor_placements", "unbacked"],
+        choices=["default", "dtensor", "dtensor_placements", "unbacked", "streams"],
         default="default",
         help="Template to use for code generation (default: default)",
     )
diff --git a/tools/experimental/torchfuzz/ops_fuzzer.py b/tools/experimental/torchfuzz/ops_fuzzer.py
index 9bea882bc5101..795c4fff60694 100644
--- a/tools/experimental/torchfuzz/ops_fuzzer.py
+++ b/tools/experimental/torchfuzz/ops_fuzzer.py
@@ -51,6 +51,10 @@ def _get_template_filtered_operators(
         from torchfuzz.codegen import UnbackedFuzzTemplate
 
         fuzz_template = UnbackedFuzzTemplate()
+    elif template == "streams":
+        from torchfuzz.codegen import StreamFuzzTemplate
+
+        fuzz_template = StreamFuzzTemplate()
     else:
         from torchfuzz.codegen import DefaultFuzzTemplate
 
@@ -252,6 +256,10 @@ def fuzz_spec(template: str = "default") -> Spec:
             from torchfuzz.codegen import UnbackedFuzzTemplate
 
             fuzz_template = UnbackedFuzzTemplate()
+        elif template == "streams":
+            from torchfuzz.codegen import StreamFuzzTemplate
+
+            fuzz_template = StreamFuzzTemplate()
         else:
             from torchfuzz.codegen import DefaultFuzzTemplate
 
diff --git a/tools/experimental/torchfuzz/tensor_fuzzer.py b/tools/experimental/torchfuzz/tensor_fuzzer.py
index 336b7a2ad1213..004464296cea6 100644
--- a/tools/experimental/torchfuzz/tensor_fuzzer.py
+++ b/tools/experimental/torchfuzz/tensor_fuzzer.py
@@ -1,6 +1,6 @@
 # mypy: ignore-errors
 import random
-from typing import NamedTuple, Union
+from typing import NamedTuple, TypeAlias
 
 import torch
 
@@ -30,8 +30,7 @@ class ScalarSpec(NamedTuple):
     )
 
 
-# Union type for specs
-Spec = Union[TensorSpec, ScalarSpec]
+Spec: TypeAlias = TensorSpec | ScalarSpec
 
 
 def fuzz_torch_tensor_type(template: str = "default") -> torch.dtype:
diff --git a/tools/experimental/torchfuzz/test_streams_template.py b/tools/experimental/torchfuzz/test_streams_template.py
new file mode 100644
index 0000000000000..b8fbf7e3fcb58
--- /dev/null
+++ b/tools/experimental/torchfuzz/test_streams_template.py
@@ -0,0 +1,107 @@
+#!/usr/bin/env python3
+"""Tests for the streams fuzzing template codegen."""
+
+import os
+import random
+import sys
+import unittest
+
+
+# Add parent directory to path so we can import torchfuzz as a module
+current_dir = os.path.dirname(os.path.abspath(__file__))
+parent_dir = os.path.dirname(current_dir)
+if parent_dir not in sys.path:
+    sys.path.insert(0, parent_dir)
+
+import torch
+from torchfuzz.codegen import convert_graph_to_python_code, StreamFuzzTemplate
+from torchfuzz.ops_fuzzer import fuzz_operation_graph, fuzz_spec
+
+
+class TestStreamsFuzzTemplate(unittest.TestCase):
+    def _generate_code(self, seed):
+        random.seed(seed)
+        torch.manual_seed(seed)
+        target_spec = fuzz_spec("streams")
+        graph = fuzz_operation_graph(
+            target_spec, max_depth=3, seed=seed, template="streams"
+        )
+        return convert_graph_to_python_code(graph, seed=seed, template="streams")
+
+    def test_template_inherits_default_ops(self):
+        template = StreamFuzzTemplate()
+        self.assertGreater(len(template.supported_ops), 0)
+        self.assertIn("torch.add", template.supported_ops)
+        self.assertIn("torch.matmul", template.supported_ops)
+
+    def test_template_uses_backward_check(self):
+        from torchfuzz.checks import EagerVsFullGraphDynamicCompileWithBackwardCheck
+
+        template = StreamFuzzTemplate()
+        self.assertIsInstance(
+            template.check, EagerVsFullGraphDynamicCompileWithBackwardCheck
+        )
+
+    def test_codegen_creates_streams(self):
+        code = self._generate_code(seed=999)
+        self.assertIn("torch.cuda.Stream()", code)
+
+    def test_codegen_has_stream_context(self):
+        code = self._generate_code(seed=999)
+        self.assertIn("with torch.cuda.stream(", code)
+
+    def test_codegen_has_final_sync(self):
+        code = self._generate_code(seed=999)
+        # Should sync all streams before return (wait_stream or wait_event)
+        self.assertTrue(
+            "torch.cuda.current_stream().wait_stream(" in code
+            or "torch.cuda.current_stream().wait_event(" in code
+        )
+
+    def test_codegen_has_backward(self):
+        code = self._generate_code(seed=999)
+        self.assertIn(".sum().backward()", code)
+
+    def test_codegen_has_requires_grad(self):
+        """Float tensor args should have requires_grad for backward testing."""
+        code = self._generate_code(seed=999)
+        self.assertIn("requires_grad_(True)", code)
+
+    def test_codegen_cross_stream_sync(self):
+        """Seeds with cross-stream deps should have inter-stream sync."""
+        # seed 999 produces a graph with ops on different streams
+        code = self._generate_code(seed=999)
+        # Should have either wait_stream or event-based sync between streams
+        has_wait_stream = ".wait_stream(s" in code
+        has_wait_event = ".wait_event(" in code
+        self.assertTrue(
+            has_wait_stream or has_wait_event,
+            "Expected cross-stream synchronization in generated code",
+        )
+
+    def test_codegen_event_based_sync(self):
+        """Some seeds should use event-based synchronization."""
+        found_events = False
+        for seed in range(50, 70):
+            code = self._generate_code(seed=seed)
+            if "torch.cuda.Event()" in code and ".wait_event(" in code:
+                found_events = True
+                # When events are used, should have record + wait pattern
+                self.assertIn(".record(", code)
+                break
+        self.assertTrue(found_events, "No seed in range produced event-based sync")
+
+    def test_codegen_deterministic(self):
+        """Same seed should produce identical code."""
+        code1 = self._generate_code(seed=42)
+        code2 = self._generate_code(seed=42)
+        self.assertEqual(code1, code2)
+
+    def test_codegen_is_valid_python(self):
+        """Generated code should be syntactically valid Python."""
+        code = self._generate_code(seed=999)
+        compile(code, "<test>", "exec")
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/tools/linter/adapters/actionlint_linter.py b/tools/linter/adapters/actionlint_linter.py
index 019f0fe896bcd..2884fe732625b 100644
--- a/tools/linter/adapters/actionlint_linter.py
+++ b/tools/linter/adapters/actionlint_linter.py
@@ -75,6 +75,10 @@ def check_file(
                 '"runs-on" section must be sequence node but got mapping node with "!!map" tag',
                 "-ignore",
                 'input "freethreaded" is not defined in action "actions/setup-python@v',
+                # GitHub increased workflow_dispatch limit to 25 inputs (Dec 2025).
+                # actionlint fixed this in v1.7.10; remove after upgrading from v1.7.7.
+                "-ignore",
+                'maximum number of inputs for "workflow_dispatch" event is 10 but',
                 file,
             ]
         )
diff --git a/tools/linter/adapters/bazel_linter.py b/tools/linter/adapters/bazel_linter.py
deleted file mode 100644
index 926628d3d76a9..0000000000000
--- a/tools/linter/adapters/bazel_linter.py
+++ /dev/null
@@ -1,197 +0,0 @@
-"""
-This linter ensures that users don't set a SHA hash checksum in Bazel for the http_archive.
-Although the security practice of setting the checksum is good, it doesn't work when the
-archive is downloaded from some sites like GitHub because it can change. Specifically,
-GitHub gives no guarantee to keep the same value forever. Check for more details at
-https://github.com/community/community/discussions/46034.
-"""
-
-from __future__ import annotations
-
-import argparse
-import json
-import re
-import shlex
-import subprocess
-import sys
-import xml.etree.ElementTree as ET
-from enum import Enum
-from typing import NamedTuple
-from urllib.parse import urlparse
-
-
-LINTER_CODE = "BAZEL_LINTER"
-SHA256_REGEX = re.compile(r"\s*sha256\s*=\s*['\"](?P<sha256>[a-zA-Z0-9]{64})['\"]\s*,")
-DOMAINS_WITH_UNSTABLE_CHECKSUM = {"github.com"}
-
-
-class LintSeverity(str, Enum):
-    ERROR = "error"
-    WARNING = "warning"
-    ADVICE = "advice"
-    DISABLED = "disabled"
-
-
-class LintMessage(NamedTuple):
-    path: str | None
-    line: int | None
-    char: int | None
-    code: str
-    severity: LintSeverity
-    name: str
-    original: str | None
-    replacement: str | None
-    description: str | None
-
-
-def is_required_checksum(urls: list[str | None]) -> bool:
-    if not urls:
-        return False
-
-    for url in urls:
-        if not url:
-            continue
-
-        parsed_url = urlparse(url)
-        if parsed_url.hostname in DOMAINS_WITH_UNSTABLE_CHECKSUM:
-            return False
-
-    return True
-
-
-def get_disallowed_checksums(
-    binary: str,
-) -> set[str]:
-    """
-    Return the set of disallowed checksums from all http_archive rules
-    """
-    # Use bazel to get the list of external dependencies in XML format
-    proc = subprocess.run(
-        [binary, "query", "kind(http_archive, //external:*)", "--output=xml"],
-        capture_output=True,
-        check=True,
-        text=True,
-    )
-
-    root = ET.fromstring(proc.stdout)
-
-    disallowed_checksums = set()
-    # Parse all the http_archive rules in the XML output
-    for rule in root.findall('.//rule[@class="http_archive"]'):
-        urls_node = rule.find('.//list[@name="urls"]')
-        if urls_node is None:
-            continue
-        urls = [n.get("value") for n in urls_node.findall(".//string")]
-
-        checksum_node = rule.find('.//string[@name="sha256"]')
-        if checksum_node is None:
-            continue
-        checksum = checksum_node.get("value")
-
-        if not checksum:
-            continue
-
-        if not is_required_checksum(urls):
-            disallowed_checksums.add(checksum)
-
-    return disallowed_checksums
-
-
-def check_bazel(
-    filename: str,
-    disallowed_checksums: set[str],
-) -> list[LintMessage]:
-    original = ""
-    replacement = ""
-
-    with open(filename) as f:
-        for line in f:
-            original += f"{line}"
-
-            m = SHA256_REGEX.match(line)
-            if m:
-                sha256 = m.group("sha256")
-
-                if sha256 in disallowed_checksums:
-                    continue
-
-            replacement += f"{line}"
-
-        if original == replacement:
-            return []
-
-        return [
-            LintMessage(
-                path=filename,
-                line=None,
-                char=None,
-                code=LINTER_CODE,
-                severity=LintSeverity.ADVICE,
-                name="format",
-                original=original,
-                replacement=replacement,
-                description="Found redundant SHA checksums. Run `lintrunner -a` to apply this patch.",
-            )
-        ]
-
-
-def main() -> None:
-    parser = argparse.ArgumentParser(
-        description="A custom linter to detect redundant SHA checksums in Bazel",
-        fromfile_prefix_chars="@",
-    )
-    parser.add_argument(
-        "--binary",
-        required=True,
-        help="bazel binary path",
-    )
-    parser.add_argument(
-        "filenames",
-        nargs="+",
-        help="paths to lint",
-    )
-    args = parser.parse_args()
-
-    try:
-        disallowed_checksums = get_disallowed_checksums(args.binary)
-    except subprocess.CalledProcessError as err:
-        err_msg = LintMessage(
-            path=None,
-            line=None,
-            char=None,
-            code=__file__,
-            severity=LintSeverity.ADVICE,
-            name="command-failed",
-            original=None,
-            replacement=None,
-            description=(
-                f"COMMAND (exit code {err.returncode})\n"
-                f"{shlex.join(err.cmd)}\n\n"
-                f"STDERR\n{err.stderr or '(empty)'}\n\n"
-                f"STDOUT\n{err.stdout or '(empty)'}"
-            ),
-        )
-        print(json.dumps(err_msg._asdict()))
-        return
-    except Exception as e:
-        err_msg = LintMessage(
-            path=None,
-            line=None,
-            char=None,
-            code=LINTER_CODE,
-            severity=LintSeverity.ERROR,
-            name="command-failed",
-            original=None,
-            replacement=None,
-            description=(f"Failed due to {e.__class__.__name__}:\n{e}"),
-        )
-        print(json.dumps(err_msg._asdict()), flush=True)
-        sys.exit(0)
-
-    for filename in args.filenames:
-        for lint_message in check_bazel(filename, disallowed_checksums):
-            print(json.dumps(lint_message._asdict()), flush=True)
-
-
-if __name__ == "__main__":
-    main()
diff --git a/tools/linter/adapters/codespell_linter.py b/tools/linter/adapters/codespell_linter.py
index 527e3d4fbfb61..dae71cde831f4 100644
--- a/tools/linter/adapters/codespell_linter.py
+++ b/tools/linter/adapters/codespell_linter.py
@@ -142,7 +142,7 @@ def check_dictionary(filename: str) -> list[LintMessage]:
         words_set = set(words)
         if len(words) != len(words_set):
             raise ValueError("The dictionary file contains duplicate entries.")
-        # pyrefly: ignore [no-matching-overload]
+        # pyrefly: ignore [bad-argument-type, no-matching-overload]
         uncased_words = list(map(str.lower, words))
         if uncased_words != sorted(uncased_words):
             raise ValueError(
diff --git a/tools/linter/adapters/docstring_linter-grandfather.json b/tools/linter/adapters/docstring_linter-grandfather.json
index f691f77008102..92fb43f07ede2 100644
--- a/tools/linter/adapters/docstring_linter-grandfather.json
+++ b/tools/linter/adapters/docstring_linter-grandfather.json
@@ -15,6 +15,14 @@
     "class MixedMMH100": 133,
     "def MixedMMH100.get_best_choices()": 86
   },
+  "torch/_inductor/autoheuristic/artifacts/_PadMMA100.py": {
+    "class PadMMA100": 249,
+    "def PadMMA100.get_best_choices()": 221
+  },
+  "torch/_inductor/autoheuristic/artifacts/_PadMMH200.py": {
+    "class PadMMH200": 159,
+    "def PadMMH200.get_best_choices()": 131
+  },
   "torch/_inductor/bounds.py": {
     "class ValueRangeAnalysis": 108
   },
@@ -31,7 +39,7 @@
   },
   "torch/_inductor/codegen/cpp.py": {
     "class CppKernelProxy": 616,
-    "class CppOverrides": 437,
+    "class CppOverrides": 496,
     "class CppScheduling": 819,
     "class CppVecKernel": 867,
     "class OuterLoopFusedSchedulerNode": 159,
@@ -79,7 +87,7 @@
   },
   "torch/_inductor/codegen/halide.py": {
     "class HalideKernel": 999,
-    "class HalideOverrides": 339,
+    "class HalideOverrides": 384,
     "class HalidePrinter": 128,
     "def HalideKernel.halide_kernel_meta()": 82
   },
diff --git a/tools/linter/adapters/gb_registry_linter.py b/tools/linter/adapters/gb_registry_linter.py
index 135d8b531bf12..41f09e22486b8 100644
--- a/tools/linter/adapters/gb_registry_linter.py
+++ b/tools/linter/adapters/gb_registry_linter.py
@@ -53,6 +53,49 @@ class LintMessage(NamedTuple):
     description: str | None
 
 
+def _is_noqa_suppressed(source_lines: list[str], lineno: int) -> bool:
+    if lineno <= 0 or lineno > len(source_lines):
+        return False
+    if source_lines[lineno - 1].rstrip().endswith(f"# noqa: {LINTER_CODE}"):
+        return True
+    if lineno > 1 and source_lines[lineno - 2].strip() == f"# noqa: {LINTER_CODE}":
+        return True
+    return False
+
+
+def _is_forbidden_raise(node: ast.Raise) -> bool:
+    if not isinstance(node.exc, ast.Call):
+        return False
+    if isinstance(node.exc.func, ast.Name):
+        return node.exc.func.id == "Unsupported"
+    if isinstance(node.exc.func, ast.Attribute):
+        return node.exc.func.attr == "Unsupported"
+    return False
+
+
+def _collect_forbidden_unsupported_raises(
+    dynamo_dir: Path,
+) -> list[tuple[Path, int, int]]:
+    forbidden_raises: list[tuple[Path, int, int]] = []
+
+    for py_file in dynamo_dir.rglob("*.py"):
+        source = py_file.read_text(encoding="utf-8")
+        source_lines = source.splitlines()
+        try:
+            tree = ast.parse(source, filename=str(py_file))
+        except SyntaxError:
+            continue
+
+        for node in ast.walk(tree):
+            if not isinstance(node, ast.Raise) or not _is_forbidden_raise(node):
+                continue
+            if _is_noqa_suppressed(source_lines, node.lineno):
+                continue
+            forbidden_raises.append((py_file, node.lineno, node.col_offset + 1))
+
+    return forbidden_raises
+
+
 def _collect_all_calls(
     dynamo_dir: Path,
 ) -> dict[str, list[tuple[dict[str, Any], Path]]]:
@@ -175,6 +218,26 @@ def check_registry_sync(dynamo_dir: Path, registry_path: Path) -> list[LintMessa
     """Check registry sync and return lint messages."""
     lint_messages = []
 
+    forbidden_raises = _collect_forbidden_unsupported_raises(dynamo_dir)
+    for path, line, char in forbidden_raises:
+        lint_messages.append(
+            LintMessage(
+                path=str(path),
+                line=line,
+                char=char,
+                code=LINTER_CODE,
+                severity=LintSeverity.ERROR,
+                name="Direct raise Unsupported",
+                original=None,
+                replacement=None,
+                description=(
+                    "Do not directly `raise Unsupported(...)` in `torch/_dynamo`. "
+                    "Use `unimplemented(...)` for graph breaks, or add `# noqa: GB_REGISTRY` "
+                    "for infra-only exceptions."
+                ),
+            )
+        )
+
     all_calls = _collect_all_calls(dynamo_dir)
 
     duplicates = []
diff --git a/tools/linter/adapters/pyrefly_linter.py b/tools/linter/adapters/pyrefly_linter.py
index 0291293087758..27ae0cc2f4c3b 100644
--- a/tools/linter/adapters/pyrefly_linter.py
+++ b/tools/linter/adapters/pyrefly_linter.py
@@ -5,7 +5,7 @@
 #   "numpy==2.1.0 ; python_version >= '3.12' and python_version <= '3.13'",
 #   "numpy==2.3.4 ; python_version >= '3.14'",
 #   "expecttest==0.3.0",
-#   "pyrefly==0.52.0",
+#   "pyrefly==0.58.0",
 #   "sympy==1.13.3",
 #   "types-requests==2.27.25",
 #   "types-pyyaml==6.0.2",
@@ -189,10 +189,12 @@ def check_files(
             )
         ]
 
-    # Parse JSON output from pyrefly
+    # Parse JSON output from pyrefly. In GitHub Actions, pyrefly appends
+    # ::error commands to stdout after the JSON, so use raw_decode to parse
+    # only the first JSON object and ignore trailing output.
     try:
         if stdout:
-            result = json.loads(stdout)
+            result, _ = json.JSONDecoder().raw_decode(stdout)
             errors = result.get("errors", [])
         else:
             errors = []
diff --git a/tools/linter/adapters/s3_init_config.json b/tools/linter/adapters/s3_init_config.json
index 9874f847a6736..67e50fd2c347d 100644
--- a/tools/linter/adapters/s3_init_config.json
+++ b/tools/linter/adapters/s3_init_config.json
@@ -7,11 +7,11 @@
         "Some binaries have custom builds; see https://github.com/pytorch/test-infra/blob/main/.github/workflows/clang-tidy-linux.yml and https://github.com/pytorch/test-infra/blob/main/.github/workflows/clang-tidy-macos.yml"
     ],
     "clang-format": {
-        "Darwin-arm": {
+        "Darwin-arm64": {
             "download_url": "https://oss-clang-format.s3.us-east-2.amazonaws.com/macos-arm/19.1.4/clang-format",
             "hash": "f0da3ecf0ab1e9b50e8c27bd2d7ca0baa619e2f4b824b35d79d46356581fa552"
         },
-        "Darwin-i386": {
+        "Darwin-x86_64": {
             "download_url": "https://oss-clang-format.s3.us-east-2.amazonaws.com/macos-i386/19.1.4/clang-format",
             "hash": "f5eb5037b9aa9d1d2de650fb2e0fe1a2517768a462fae8e98791a67b698302f4"
         },
@@ -25,11 +25,11 @@
         }
     },
     "clang-tidy": {
-        "Darwin-i386": {
+        "Darwin-x86_64": {
             "download_url": "https://oss-clang-format.s3.us-east-2.amazonaws.com/macos-i386/19.1.4/clang-tidy",
             "hash": "7b5da17d3f8b1c18c77d043999f05293f43402affb16de15dfcb276971984a3e"
         },
-        "Darwin-arm": {
+        "Darwin-arm64": {
             "download_url": "https://oss-clang-format.s3.us-east-2.amazonaws.com/macos-arm/19.1.4/clang-tidy",
             "hash": "04243f4044fe6d95f6d51d15be803331c3cbb61f2d8fcfeba5a5dec1e7ae6dfb"
         },
@@ -43,11 +43,11 @@
         }
     },
     "actionlint": {
-        "Darwin-i386": {
+        "Darwin-x86_64": {
             "download_url": "https://oss-clang-format.s3.us-east-2.amazonaws.com/actionlint/1.7.7/Darwin_amd64/actionlint",
             "hash": "996affd492c57441c5ecfe00dedaef1fde056872d242c0cf7cc15de058d59d03"
         },
-        "Darwin-arm": {
+        "Darwin-arm64": {
             "download_url": "https://oss-clang-format.s3.us-east-2.amazonaws.com/actionlint/1.7.7/Darwin_arm64/actionlint",
             "hash": "00aba386d026da33be6e85dd5a46d7af4dd9e4d6cbdb02335f4b267162fd2d9e"
         },
@@ -59,15 +59,5 @@
             "download_url": "https://oss-clang-format.s3.us-east-2.amazonaws.com/actionlint/1.7.7/Linux_arm64/actionlint",
             "hash": "446687e63fac45472b0a66bae28975c28678af062670af119c11a7087baf35cc"
         }
-    },
-    "bazel": {
-        "Darwin": {
-            "download_url": "https://raw.githubusercontent.com/bazelbuild/bazelisk/v1.16.0/bazelisk.py",
-            "hash": "1f6d76d023ddd5f1625f34d934418e7334a267318d084f31be09df8a8835ed16"
-        },
-        "Linux": {
-            "download_url": "https://raw.githubusercontent.com/bazelbuild/bazelisk/v1.16.0/bazelisk.py",
-            "hash": "1f6d76d023ddd5f1625f34d934418e7334a267318d084f31be09df8a8835ed16"
-        }
     }
 }
diff --git a/tools/linter/adapters/workflow_consistency_linter.py b/tools/linter/adapters/workflow_consistency_linter.py
index 54687a1ff4214..f931a117b6978 100644
--- a/tools/linter/adapters/workflow_consistency_linter.py
+++ b/tools/linter/adapters/workflow_consistency_linter.py
@@ -116,6 +116,15 @@ def get_jobs_with_sync_tag(
     # and ['with']['tests-to-include'], since dispatch filters differ
     if "tests-to-include" in job.get("with", {}):
         del job["with"]["tests-to-include"]
+    # and ['with']['build-environment'], since GPU-specific suffixes differ for ROCm
+    if (
+        "build-environment" in job.get("with", {})
+        and "rocm" in job["with"]["build-environment"]
+    ):
+        del job["with"]["build-environment"]
+    # and ['name'], since ROCm jobs append a GPU-specific suffix to the job name
+    if "name" in job and "rocm" in job.get("name", ""):
+        del job["name"]
 
     # normalize needs: remove helper job-filter so comparisons ignore it
     needs = job.get("needs")
diff --git a/tools/nightly.py b/tools/nightly.py
index 88ba7ccb2ba81..6c4a04971eaa6 100755
--- a/tools/nightly.py
+++ b/tools/nightly.py
@@ -288,7 +288,7 @@ def site_packages(self, python: Path | str | None = None) -> Path:
             python=python,
             capture_output=True,
         ).stdout
-        # pyrefly: ignore [no-matching-overload]
+        # pyrefly: ignore [bad-argument-type, no-matching-overload]
         candidates = list(map(Path, filter(None, map(str.strip, output.splitlines()))))
         candidates = [p for p in candidates if p.is_dir() and p.name == "site-packages"]
         if not candidates:
diff --git a/tools/optional_submodules.py b/tools/optional_submodules.py
index c3c4dacaa4032..5164acf69058b 100644
--- a/tools/optional_submodules.py
+++ b/tools/optional_submodules.py
@@ -1,3 +1,4 @@
+import os
 from pathlib import Path
 from subprocess import check_call
 
@@ -26,16 +27,18 @@ def _checkout_by_tag(repo: str, tag: str) -> None:
     )
 
 
-def read_nccl_pin() -> str:
+def read_nccl_pin(cuda_version: str = "") -> str:
     # Default NCCL version
     nccl_file = "nccl.txt"
 
     # If NCCL version diverges for different CUDA versions, uncomment the
     # following block and add the appropriate file (using CUDA 11 as an example)
 
-    # cuda_version = os.getenv("DESIRED_CUDA", os.getenv("CUDA_VERSION", ""))
-    # if cuda_version.startswith("11"):
-    #     nccl_file = "nccl-cu11.txt"
+    # 12.6 builds for sm50, needs a lower version
+    if not cuda_version:
+        cuda_version = os.getenv("DESIRED_CUDA", os.getenv("CUDA_VERSION", ""))
+    if cuda_version.startswith("12.6") or cuda_version == "cu126":
+        nccl_file = "nccl-cu126.txt"
 
     nccl_pin_path = repo_root / ".ci" / "docker" / "ci_commit_pins" / nccl_file
     return _read_file(nccl_pin_path)
diff --git a/tools/pyi/gen_pyi.py b/tools/pyi/gen_pyi.py
index 44f0360c1fb33..57c1d89f735d6 100644
--- a/tools/pyi/gen_pyi.py
+++ b/tools/pyi/gen_pyi.py
@@ -996,7 +996,7 @@ def add_docstr_to_hint(docstr: str, hint: str) -> str:
         hint = hint.removesuffix("...").rstrip()  # remove "..."
         content = hint + "\n" + textwrap.indent(f'r"""\n{docstr}\n"""', prefix="    ")
         # Remove trailing whitespace on each line
-        # pyrefly: ignore [no-matching-overload]
+        # pyrefly: ignore [bad-argument-type, no-matching-overload]
         return "\n".join(map(str.rstrip, content.splitlines())).rstrip()
 
     # attribute or property
@@ -1072,7 +1072,7 @@ def gen_pyi(
                         "dtype: _dtype | None = None",
                         "device: DeviceLikeType | None = None",
                         "copy: _bool | None = None",
-                        "requires_grad: _bool = False",
+                        "requires_grad: _bool | None = None",
                     ],
                     "Tensor",
                 )
diff --git a/tools/rules/BUILD b/tools/rules/BUILD
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/tools/rules/cu.bzl b/tools/rules/cu.bzl
deleted file mode 100644
index 8e137f63304cd..0000000000000
--- a/tools/rules/cu.bzl
+++ /dev/null
@@ -1,42 +0,0 @@
-load("@rules_cuda//cuda:defs.bzl", "cuda_library")
-
-NVCC_COPTS = [
-    "--expt-relaxed-constexpr",
-    "--expt-extended-lambda",
-    "--compiler-options=-Werror=all",
-    # The following warnings come from -Wall. We downgrade them from
-    # error to warnings here.
-    #
-    # sign-compare has a tremendous amount of violations in the
-    # codebase. It will be a lot of work to fix them, just disable it
-    # for now.
-    "--compiler-options=-Wno-sign-compare",
-    # We intentionally use #pragma unroll, which is compiler specific.
-    "--compiler-options=-Wno-error=unknown-pragmas",
-    "--compiler-options=-Werror=extra",
-    # The following warnings come from -Wextra. We downgrade them from
-    # error to warnings here.
-    #
-    # unused-parameter-compare has a tremendous amount of violations
-    # in the codebase. It will be a lot of work to fix them, just
-    # disable it for now.
-    "--compiler-options=-Wno-unused-parameter",
-    # missing-field-parameters has both a large number of violations
-    # in the codebase, but it also is used pervasively in the Python C
-    # API. There are a couple of catches though:
-    # * we use multiple versions of the Python API and hence have
-    #   potentially multiple different versions of each relevant
-    #   struct. They may have different numbers of fields. It will be
-    #   unwieldy to support multiple versions in the same source file.
-    # * Python itself for many of these structs recommends only
-    #   initializing a subset of the fields. We should respect the API
-    #   usage conventions of our dependencies.
-    #
-    # Hence, we just disable this warning altogether. We may want to
-    # clean up some of the clear-cut cases that could be risky, but we
-    # still likely want to have this disabled for the most part.
-    "-Wno-missing-field-initializers",
-]
-
-def cu_library(name, srcs, copts = [], **kwargs):
-    cuda_library(name, srcs = srcs, copts = NVCC_COPTS + copts, **kwargs)
diff --git a/tools/rules/workspace.bzl b/tools/rules/workspace.bzl
deleted file mode 100644
index 178cc256b1557..0000000000000
--- a/tools/rules/workspace.bzl
+++ /dev/null
@@ -1,54 +0,0 @@
-def _impl(repository_ctx):
-    archive = repository_ctx.attr.name + ".tar"
-    reference = Label("@%s_unpatched//:README" % repository_ctx.attr.name)
-    dirname = repository_ctx.path(reference).dirname
-    repository_ctx.execute(["tar", "hcf", archive, "-C", dirname, "."])
-    repository_ctx.extract(archive)
-    for patch in repository_ctx.attr.patches:
-        repository_ctx.patch(repository_ctx.path(patch), repository_ctx.attr.patch_strip)
-    build_file = repository_ctx.path(repository_ctx.attr.build_file)
-    repository_ctx.execute(["cp", build_file, "BUILD.bazel"])
-
-_patched_rule = repository_rule(
-    implementation = _impl,
-    attrs = {
-        "build_file": attr.label(),
-        "patch_strip": attr.int(),
-        "patches": attr.label_list(),
-    },
-)
-
-def new_patched_local_repository(name, path, **kwargs):
-    native.new_local_repository(
-        name = name + "_unpatched",
-        build_file_content = """
-pkg_tar(name = "content", srcs = glob(["**"]))
-""",
-        path = path,
-    )
-    _patched_rule(name = name, **kwargs)
-
-def _new_empty_repository_impl(repo_ctx):
-    build_file = repo_ctx.attr.build_file
-    build_file_content = repo_ctx.attr.build_file_content
-    if not (bool(build_file) != bool(build_file_content)):
-        fail("Exactly one of 'build_file' or 'build_file_content' is required")
-
-    if build_file_content:
-        repo_ctx.file("BUILD", build_file_content)
-    elif build_file:
-        repo_ctx.template("BUILD", repo_ctx.attr.build_file, {})
-
-new_empty_repository = repository_rule(
-    attrs = {
-        "build_file": attr.label(allow_files = True),
-        "build_file_content": attr.string(),
-    },
-    implementation = _new_empty_repository_impl,
-)
-
-"""Create an empty repository with the supplied BUILD file.
-
-This is mostly useful to create wrappers for specific target that we want
-to be used with the '@' syntax.
-"""
diff --git a/tools/rules_cc/cuda_support.patch b/tools/rules_cc/cuda_support.patch
index 1f48144512023..795143f5729f5 100644
--- a/tools/rules_cc/cuda_support.patch
+++ b/tools/rules_cc/cuda_support.patch
@@ -7,7 +7,7 @@ index ba992fc..e4e8364 100644
      "write_builtin_include_directory_paths",
  )
 +load("@rules_cuda//cuda:toolchain.bzl", "cuda_compiler_deps")
- 
+
  def _field(name, value):
      """Returns properly indented top level crosstool field."""
 @@ -397,7 +398,7 @@ def configure_unix_toolchain(repository_ctx, cpu_value, overriden_tools):
@@ -15,10 +15,10 @@ index ba992fc..e4e8364 100644
          repository_ctx,
          "BAZEL_CXXOPTS",
 -        "-std=c++0x",
-+        "-std=c++17",
++        "-std=c++20",
          False,
      ), ":")
- 
+
 @@ -463,7 +464,7 @@ def configure_unix_toolchain(repository_ctx, cpu_value, overriden_tools):
              )),
              "%{cc_compiler_deps}": get_starlark_list([":builtin_include_directory_paths"] + (
@@ -37,7 +37,7 @@ index c3cf3ba..1744eb4 100644
      "with_feature_set",
  )
 +load("@rules_cuda//cuda:toolchain.bzl", "cuda_toolchain_config")
- 
+
  all_compile_actions = [
      ACTION_NAMES.c_compile,
 @@ -580,7 +581,8 @@ def _impl(ctx):
@@ -53,7 +53,7 @@ index c3cf3ba..1744eb4 100644
 @@ -1152,10 +1154,15 @@ def _impl(ctx):
              unfiltered_compile_flags_feature,
          ]
- 
+
 +    cuda = cuda_toolchain_config(
 +        cuda_toolchain_info = ctx.attr._cuda_toolchain_info,
 +        compiler_path = ctx.attr.tool_paths["gcc"],
diff --git a/tools/setup_helpers/cmake.py b/tools/setup_helpers/cmake.py
index 88f0fe5d3094a..e3c5f2ce6e002 100644
--- a/tools/setup_helpers/cmake.py
+++ b/tools/setup_helpers/cmake.py
@@ -12,7 +12,6 @@
 import sysconfig
 from pathlib import Path
 from subprocess import CalledProcessError, check_call, check_output, DEVNULL
-from typing import cast
 
 from .cmake_utils import CMakeValue, get_cmake_cache_variables_from_file
 from .env import (
@@ -244,151 +243,34 @@ def generate(
         _mkdir_p(install_dir)
         _mkdir_p(self.build_dir)
 
-        # Store build options that are directly stored in environment variables
-        build_options: dict[str, CMakeValue] = {}
-
-        # Build options that do not start with "BUILD_", "USE_", or "CMAKE_" and are directly controlled by env vars.
-        # This is a dict that maps environment variables to the corresponding variable name in CMake.
-        additional_options = {
-            # Key: environment variable name. Value: Corresponding variable name to be passed to CMake. If you are
-            # adding a new build option to this block: Consider making these two names identical and adding this option
-            # in the block below.
-            "CUDNN_LIB_DIR": "CUDNN_LIBRARY",
-            "USE_CUDA_STATIC_LINK": "CAFFE2_STATIC_LINK_CUDA",
-        }
-        additional_options.update(
-            {
-                # Build options that have the same environment variable name and CMake variable name and that do not start
-                # with "BUILD_", "USE_", or "CMAKE_". If you are adding a new build option, also make sure you add it to
-                # CMakeLists.txt.
-                var: var
-                for var in (
-                    "UBSAN_FLAGS",
-                    "BLAS",
-                    "WITH_BLAS",
-                    "CUDA_HOST_COMPILER",
-                    "CUDA_NVCC_EXECUTABLE",
-                    "CUDA_SEPARABLE_COMPILATION",
-                    "CUDNN_LIBRARY",
-                    "CUDNN_INCLUDE_DIR",
-                    "CUDNN_ROOT",
-                    "EXPERIMENTAL_SINGLE_THREAD_POOL",
-                    "INSTALL_TEST",
-                    "JAVA_HOME",
-                    "INTEL_MKL_DIR",
-                    "INTEL_OMP_DIR",
-                    "MKL_THREADING",
-                    "MKLDNN_CPU_RUNTIME",
-                    "MSVC_Z7_OVERRIDE",
-                    "CAFFE2_USE_MSVC_STATIC_RUNTIME",
-                    "Numa_INCLUDE_DIR",
-                    "Numa_LIBRARIES",
-                    "ONNX_ML",
-                    "ONNX_NAMESPACE",
-                    "ATEN_THREADING",
-                    "WERROR",
-                    "OPENSSL_ROOT_DIR",
-                    "STATIC_DISPATCH_BACKEND",
-                    "SELECTED_OP_LIST",
-                    "TORCH_CUDA_ARCH_LIST",
-                    "TORCH_XPU_ARCH_LIST",
-                    "TRACING_BASED",
-                    "PYTHON_LIB_REL_PATH",
-                )
-            }
-        )
-
-        # Aliases which are lower priority than their canonical option
-        low_priority_aliases = {
-            "CUDA_HOST_COMPILER": "CMAKE_CUDA_HOST_COMPILER",
-            "CUDAHOSTCXX": "CUDA_HOST_COMPILER",
-            "CMAKE_CUDA_HOST_COMPILER": "CUDA_HOST_COMPILER",
-            "CMAKE_CUDA_COMPILER": "CUDA_NVCC_EXECUTABLE",
-            "CUDACXX": "CUDA_NVCC_EXECUTABLE",
-        }
-        for var, val in my_env.items():
-            # We currently pass over all environment variables that start with "BUILD_", "USE_", and "CMAKE_". This is
-            # because we currently have no reliable way to get the list of all build options we have specified in
-            # CMakeLists.txt. (`cmake -L` won't print dependent options when the dependency condition is not met.) We
-            # will possibly change this in the future by parsing CMakeLists.txt ourselves (then additional_options would
-            # also not be needed to be specified here).
-            true_var = additional_options.get(var)
-            if true_var is not None:
-                build_options[true_var] = val
-            elif var.startswith(("BUILD_", "USE_", "CMAKE_")) or var.endswith(
-                ("EXITCODE", "EXITCODE__TRYRUN_OUTPUT")
-            ):
-                build_options[var] = val
-
-            if var in low_priority_aliases:
-                key = low_priority_aliases[var]
-                if key not in build_options:
-                    build_options[key] = val
-
-        # The default value cannot be easily obtained in CMakeLists.txt. We set it here.
-        py_lib_path = sysconfig.get_path("purelib")
-        cmake_prefix_path = build_options.get("CMAKE_PREFIX_PATH")
-        if cmake_prefix_path:
-            build_options["CMAKE_PREFIX_PATH"] = (
-                py_lib_path + ";" + cast(str, cmake_prefix_path)
-            )
-        else:
-            build_options["CMAKE_PREFIX_PATH"] = py_lib_path
-
-        # Some options must be post-processed. Ideally, this list will be shrunk to only one or two options in the
-        # future, as CMake can detect many of these libraries pretty comfortably. We have them here for now before CMake
-        # integration is completed. They appear here not in the CMake.defines call below because they start with either
-        # "BUILD_" or "USE_" and must be overwritten here.
-        use_numpy = not check_negative_env_flag("USE_NUMPY")
-        build_options.update(
-            {
-                # Note: Do not add new build options to this dict if it is directly read from environment variable -- you
-                # only need to add one in `CMakeLists.txt`. All build options that start with "BUILD_", "USE_", or "CMAKE_"
-                # are automatically passed to CMake; For other options you can add to additional_options above.
-                "BUILD_PYTHON": build_python,
-                "BUILD_TEST": build_test,
-                # Most library detection should go to CMake script, except this one, which Python can do a much better job
-                # due to NumPy's inherent Pythonic nature.
-                "USE_NUMPY": use_numpy,
-            }
-        )
-
-        # Detect build dependencies from python lib path (in order to set *_HOME variables)
-        # NVSHMEM
-        nvshmem_py_dir = py_lib_path + "/nvidia/nvshmem"
-        if os.path.exists(nvshmem_py_dir):
-            build_options["NVSHMEM_PY_DIR"] = nvshmem_py_dir
+        # Environment variable forwarding (BUILD_*, USE_*, CMAKE_*, aliases,
+        # passthrough vars, CMAKE_PREFIX_PATH, low-priority aliases) is now
+        # handled by cmake/EnvVarForwarding.cmake, which is included early in
+        # the top-level CMakeLists.txt. Only options that require Python-side
+        # detection are passed here.
 
-        # Options starting with CMAKE_
-        cmake__options = {
+        build_options: dict[str, CMakeValue] = {
             "CMAKE_INSTALL_PREFIX": install_dir,
+            "BUILD_PYTHON": build_python,
+            "BUILD_TEST": build_test,
         }
 
-        # We set some CMAKE_* options in our Python build code instead of relying on the user's direct settings. Emit an
-        # error if the user also attempts to set these CMAKE options directly.
-        specified_cmake__options = set(build_options).intersection(cmake__options)
-        if len(specified_cmake__options) > 0:
-            eprint(
-                ", ".join(specified_cmake__options)
-                + " should not be specified in the environment variable. They are directly set by PyTorch build script."
-            )
-            sys.exit(1)
-        build_options.update(cmake__options)
-
+        use_numpy = not check_negative_env_flag("USE_NUMPY")
+        build_options["USE_NUMPY"] = use_numpy
         if use_numpy:
             try:
-                # This helps CMake find the correct include directory for NumPy
-                # This is especially useful in cross compiled environments
                 import numpy
 
-                Python_NumPy_INCLUDE_DIR = numpy.get_include()
-                build_options.update(
-                    dict(Python_NumPy_INCLUDE_DIR=Python_NumPy_INCLUDE_DIR)
-                )
+                build_options["Python_NumPy_INCLUDE_DIR"] = numpy.get_include()
             except ImportError:
-                # use_numpy is just a hint.... so we can fail silently here
                 pass
 
+        # NVSHMEM detection from Python lib path
+        py_lib_path = sysconfig.get_path("purelib")
+        nvshmem_py_dir = py_lib_path + "/nvidia/nvshmem"
+        if os.path.exists(nvshmem_py_dir):
+            build_options["NVSHMEM_PY_DIR"] = nvshmem_py_dir
+
         CMake.defines(
             args,
             Python_EXECUTABLE=sys.executable,
diff --git a/tools/stale_issues.py b/tools/stale_issues.py
new file mode 100644
index 0000000000000..7e341f4baee69
--- /dev/null
+++ b/tools/stale_issues.py
@@ -0,0 +1,316 @@
+import argparse
+import json
+import os
+import re
+import subprocess
+import sys
+import time
+from datetime import date, timedelta
+
+
+def parse_older_than(s):
+    """Parse a relative time string like '2 months' into a cutoff date."""
+    m = re.fullmatch(r"(\d+)\s*(days?|weeks?|months?|years?)", s.strip())
+    if not m:
+        raise argparse.ArgumentTypeError(
+            f"invalid time format: {s!r} (expected e.g. '30 days', '2 months', '1 year')"
+        )
+    n, unit = int(m.group(1)), m.group(2).rstrip("s")
+    today = date.today()
+    if unit == "day":
+        return today - timedelta(days=n)
+    elif unit == "week":
+        return today - timedelta(weeks=n)
+    elif unit == "month":
+        month = today.month - n
+        year = today.year + (month - 1) // 12
+        month = (month - 1) % 12 + 1
+        day = min(
+            today.day, [31, 29, 31, 30, 31, 30, 31, 31, 30, 31, 30, 31][month - 1]
+        )
+        return date(year, month, day)
+    elif unit == "year":
+        return date(today.year - n, today.month, min(today.day, 28))
+
+
+def gh_issue_list(search, label, limit):
+    """Fetch issues from gh issue list for a single label (or no label)."""
+    cmd = [
+        "gh",
+        "issue",
+        "list",
+        "-R",
+        "pytorch/pytorch",
+        "-S",
+        search,
+        "-L",
+        str(limit),
+        "--json",
+        "number,title,updatedAt,labels,url",
+    ]
+    if label:
+        cmd += ["-l", label]
+    result = subprocess.run(cmd, capture_output=True, text=True)
+    if result.returncode != 0:
+        print(result.stderr, file=sys.stderr)
+        sys.exit(1)
+    return json.loads(result.stdout)
+
+
+def gh_issue_count(search_query):
+    """Get total issue count via GitHub search API."""
+    result = subprocess.run(
+        [
+            "gh",
+            "api",
+            "search/issues",
+            "-q",
+            ".total_count",
+            "--method",
+            "GET",
+            "-f",
+            f"q={search_query}",
+            "-f",
+            "per_page=1",
+        ],
+        capture_output=True,
+        text=True,
+    )
+    return result.stdout.strip() if result.returncode == 0 else "?"
+
+
+def cmd_list(args):
+    search = "sort:updated-asc"
+    if args.older_than:
+        cutoff = parse_older_than(args.older_than)
+        search += f" updated:<{cutoff.isoformat()}"
+
+    labels = args.label or [None]
+    if len(labels) == 1:
+        issues = gh_issue_list(search, labels[0], args.limit)
+    else:
+        # Query per label and merge (gh -l does AND, we want OR)
+        seen = set()
+        issues = []
+        for label in labels:
+            for issue in gh_issue_list(search, label, args.limit):
+                if issue["number"] not in seen:
+                    seen.add(issue["number"])
+                    issues.append(issue)
+        issues.sort(key=lambda i: i["updatedAt"])
+        issues = issues[: args.limit]
+
+    for issue in issues:
+        issue_labels = ", ".join(l["name"] for l in issue["labels"])
+        print(
+            f"#{issue['number']:>6}  {issue['updatedAt'][:10]}  {issue['title'][:80]}"
+        )
+        print(f"         {issue['url']}")
+        if issue_labels:
+            print(f"         labels: {issue_labels}")
+
+    # Get total count via GitHub search API
+    base_query = "repo:pytorch/pytorch is:issue is:open"
+    if args.older_than:
+        cutoff = parse_older_than(args.older_than)
+        base_query += f" updated:<{cutoff.isoformat()}"
+    if not args.label:
+        total = gh_issue_count(base_query)
+    elif len(args.label) == 1:
+        total = gh_issue_count(base_query + f' label:"{args.label[0]}"')
+    else:
+        # Sum per-label counts (may slightly overcount shared issues)
+        total = 0
+        for label in args.label:
+            count = gh_issue_count(base_query + f' label:"{label}"')
+            try:
+                total += int(count)
+            except ValueError:
+                total = "?"
+                break
+        if isinstance(total, int):
+            total = f"~{total}"
+    print(f"\nShowing {len(issues)} of {total} issues.")
+
+
+def cmd_labels(args):
+    page = 1
+    labels = []
+    while True:
+        result = subprocess.run(
+            [
+                "gh",
+                "api",
+                "repos/pytorch/pytorch/labels",
+                "--method",
+                "GET",
+                "-f",
+                "per_page=100",
+                "-f",
+                f"page={page}",
+            ],
+            capture_output=True,
+            text=True,
+        )
+        if result.returncode != 0:
+            print(result.stderr, file=sys.stderr)
+            sys.exit(1)
+        batch = json.loads(result.stdout)
+        if not batch:
+            break
+        labels.extend(batch)
+        page += 1
+
+    labels.sort(key=lambda l: l["name"].lower())
+    for label in labels:
+        desc = f"  - {label['description']}" if label.get("description") else ""
+        print(f"{label['name']}{desc}")
+    print(f"\n{len(labels)} labels total.")
+
+
+SUBSCRIPTION_DIR = os.path.join(
+    os.environ.get("XDG_RUNTIME_DIR", os.path.join("/tmp", f"user-{os.getuid()}")),
+    "gh_subscriptions",
+)
+
+
+def gh_graphql(query):
+    result = subprocess.run(
+        ["gh", "api", "graphql", "-f", f"query={query}"],
+        capture_output=True,
+        text=True,
+    )
+    if result.returncode != 0:
+        print(result.stderr, file=sys.stderr)
+        sys.exit(1)
+    return json.loads(result.stdout)
+
+
+def subscription_state_path(issue_number):
+    return os.path.join(SUBSCRIPTION_DIR, f"{issue_number}.json")
+
+
+def cmd_subscription_save(args):
+    result = subprocess.run(
+        ["gh", "api", f"repos/pytorch/pytorch/issues/{args.issue}", "--jq", ".node_id"],
+        capture_output=True,
+        text=True,
+    )
+    if result.returncode != 0:
+        print(result.stderr, file=sys.stderr)
+        sys.exit(1)
+    node_id = result.stdout.strip()
+
+    data = gh_graphql(
+        f'{{ node(id: "{node_id}") {{ ... on Issue {{ viewerSubscription }} }} }}'
+    )
+    state = data["data"]["node"]["viewerSubscription"]
+
+    os.makedirs(SUBSCRIPTION_DIR, exist_ok=True)
+    with open(subscription_state_path(args.issue), "w") as f:
+        json.dump({"node_id": node_id, "state": state}, f)
+    print(f"#{args.issue}: {state} (node_id={node_id})")
+
+
+def cmd_subscription_restore(args):
+    path = subscription_state_path(args.issue)
+    if not os.path.exists(path):
+        print(f"No saved state for #{args.issue}", file=sys.stderr)
+        sys.exit(1)
+
+    with open(path) as f:
+        saved = json.load(f)
+    os.remove(path)
+
+    if saved["state"] == "SUBSCRIBED":
+        print(f"#{args.issue}: was already SUBSCRIBED, nothing to restore")
+        return
+
+    node_id = saved["node_id"]
+
+    # Poll until auto-subscribe has propagated, then unsubscribe.
+    for _ in range(30):
+        data = gh_graphql(
+            f'{{ node(id: "{node_id}") {{ ... on Issue {{ viewerSubscription }} }} }}'
+        )
+        if data["data"]["node"]["viewerSubscription"] == "SUBSCRIBED":
+            break
+        time.sleep(1)
+
+    gh_graphql(
+        f'mutation {{ updateSubscription(input: {{subscribableId: "{node_id}", '
+        f"state: UNSUBSCRIBED}}) {{ subscribable {{ ... on Issue {{ viewerSubscription }} }} }} }}"
+    )
+    print(f"#{args.issue}: restored to UNSUBSCRIBED")
+
+
+def cmd_collaborator_check(args):
+    result = subprocess.run(
+        [
+            "gh",
+            "api",
+            f"repos/pytorch/pytorch/collaborators/{args.username}",
+            "--silent",
+        ],
+        capture_output=True,
+        text=True,
+    )
+    if result.returncode == 0:
+        print(f"{args.username} is a collaborator")
+    else:
+        print(f"{args.username} is NOT a collaborator")
+        sys.exit(1)
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Tools for managing stale issues in pytorch/pytorch"
+    )
+    subparsers = parser.add_subparsers(dest="command", required=True)
+
+    list_parser = subparsers.add_parser(
+        "list", help="list least recently updated issues"
+    )
+    list_parser.add_argument(
+        "-L", "--limit", type=int, default=5, help="max issues to fetch (default: 5)"
+    )
+    list_parser.add_argument(
+        "-l", "--label", action="append", help="filter by label (multiple for OR)"
+    )
+    list_parser.add_argument(
+        "--older-than",
+        type=str,
+        default="3 months",
+        metavar="TIME",
+        help="only show issues not updated in this long, or '' for no cutoff (default: '3 months')",
+    )
+    list_parser.set_defaults(func=cmd_list)
+
+    labels_parser = subparsers.add_parser("labels", help="list all known labels")
+    labels_parser.set_defaults(func=cmd_labels)
+
+    collab_parser = subparsers.add_parser(
+        "collaborator-check", help="check if a user is a repo collaborator"
+    )
+    collab_parser.add_argument("username", type=str, help="GitHub username to check")
+    collab_parser.set_defaults(func=cmd_collaborator_check)
+
+    sub_parser = subparsers.add_parser(
+        "subscription", help="save/restore issue notification subscription state"
+    )
+    sub_sub = sub_parser.add_subparsers(dest="sub_command", required=True)
+    save_parser = sub_sub.add_parser("save", help="save current subscription state")
+    save_parser.add_argument("issue", type=int, help="issue number")
+    save_parser.set_defaults(func=cmd_subscription_save)
+    restore_parser = sub_sub.add_parser(
+        "restore", help="restore saved subscription state"
+    )
+    restore_parser.add_argument("issue", type=int, help="issue number")
+    restore_parser.set_defaults(func=cmd_subscription_restore)
+
+    args = parser.parse_args()
+    args.func(args)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/tools/stats/import_test_stats.py b/tools/stats/import_test_stats.py
index a7c661340d13e..3c3dfcdfd08f8 100644
--- a/tools/stats/import_test_stats.py
+++ b/tools/stats/import_test_stats.py
@@ -112,7 +112,7 @@ def process_disabled_test(the_response: dict[str, Any]) -> dict[str, Any]:
         return disabled_test_from_issues
 
     try:
-        url = "https://ossci-metrics.s3.amazonaws.com/disabled-tests-condensed.json"
+        url = "https://ossci-metrics.s3.amazonaws.com/disabled-tests-condensed.json?versionId=XRBoKk5TT4f6n48PbQ9OntiMMQveSs3J"
         return fetch_and_cache(dirpath, filename, url, process_disabled_test)
     except Exception:
         print("Couldn't download test skip set, leaving all tests enabled...")
diff --git a/tools/stats/monitor.py b/tools/stats/monitor.py
index 059854bd462b6..f2764aeb91134 100644
--- a/tools/stats/monitor.py
+++ b/tools/stats/monitor.py
@@ -459,6 +459,7 @@ def _initial_gpu_handler(self) -> None:
                 self._gpu_lib_detected = "amdsmi"
                 self._gpu_handles = amdsmi.amdsmi_get_processor_handles()
 
+            # pyrefly: ignore[bad-assignment]
             self._num_of_cpus = psutil.cpu_count(logical=True)
             # update summary info
             self._metadata.gpu_count = len(self._gpu_handles)
diff --git a/tools/test/test_gb_registry_linter.py b/tools/test/test_gb_registry_linter.py
index 2a4cc7e65be6c..780fa895f7b35 100644
--- a/tools/test/test_gb_registry_linter.py
+++ b/tools/test/test_gb_registry_linter.py
@@ -461,6 +461,93 @@ def test(self):
         )
         self.assertEqual(messages, [expected_msg])
 
+    def test_case8_forbid_direct_raise_unsupported(self):
+        callsite_content = """from torch._dynamo.exc import Unsupported
+
+def test(self):
+    raise Unsupported("testing")
+"""
+        with open(self.callsite_file, "w") as f:
+            f.write(callsite_content)
+
+        messages = check_registry_sync(self.test_data_dir, self.registry_path)
+        self.assertEqual(len(messages), 1)
+
+        message = messages[0]
+        self.assertEqual(message.path, str(self.callsite_file))
+        self.assertEqual(message.line, 4)
+        self.assertEqual(message.code, LINTER_CODE)
+        self.assertEqual(message.severity, LintSeverity.ERROR)
+        self.assertEqual(message.name, "Direct raise Unsupported")
+        self.assertIn("Use `unimplemented(...)`", message.description)
+
+    def test_case9_forbid_raise_exc_unsupported(self):
+        callsite_content = """from torch._dynamo import exc
+
+def test(self):
+    raise exc.Unsupported("testing")
+"""
+        with open(self.callsite_file, "w") as f:
+            f.write(callsite_content)
+
+        messages = check_registry_sync(self.test_data_dir, self.registry_path)
+        self.assertEqual(len(messages), 1)
+
+        message = messages[0]
+        self.assertEqual(message.path, str(self.callsite_file))
+        self.assertEqual(message.line, 4)
+        self.assertEqual(message.code, LINTER_CODE)
+        self.assertEqual(message.severity, LintSeverity.ERROR)
+        self.assertEqual(message.name, "Direct raise Unsupported")
+
+    def test_case10_allow_noqa_for_direct_raise_unsupported(self):
+        callsite_content = """from torch._dynamo.exc import Unsupported
+
+def test(self):
+    raise Unsupported("testing")  # noqa: GB_REGISTRY
+"""
+        with open(self.callsite_file, "w") as f:
+            f.write(callsite_content)
+
+        messages = check_registry_sync(self.test_data_dir, self.registry_path)
+        self.assertEqual(messages, [])
+
+    def test_case10_allow_previous_line_noqa_for_direct_raise_unsupported(self):
+        callsite_content = """from torch._dynamo.exc import Unsupported
+
+def test(self):
+    # noqa: GB_REGISTRY
+    raise Unsupported("testing")
+"""
+        with open(self.callsite_file, "w") as f:
+            f.write(callsite_content)
+
+        messages = check_registry_sync(self.test_data_dir, self.registry_path)
+        self.assertEqual(messages, [])
+
+    def test_case11_allow_noqa_in_exc_py(self):
+        callsite_content = """def test(self):
+    pass
+"""
+        with open(self.callsite_file, "w") as f:
+            f.write(callsite_content)
+
+        exc_file = self.test_data_dir / "exc.py"
+        exc_content = """class Unsupported(Exception):
+    pass
+
+def unimplemented():
+    raise Unsupported("testing")  # noqa: GB_REGISTRY
+
+def other():
+    raise Unsupported("not allowed")  # noqa: GB_REGISTRY
+"""
+        with open(exc_file, "w") as f:
+            f.write(exc_content)
+
+        messages = check_registry_sync(self.test_data_dir, self.registry_path)
+        self.assertEqual(messages, [])
+
 
 if __name__ == "__main__":
     unittest.main()
diff --git a/tools/testing/explicit_ci_jobs.py b/tools/testing/explicit_ci_jobs.py
index 0d25bc642678a..e7b7b803d46b7 100755
--- a/tools/testing/explicit_ci_jobs.py
+++ b/tools/testing/explicit_ci_jobs.py
@@ -7,92 +7,12 @@
 import subprocess
 import textwrap
 from pathlib import Path
-from typing import Any
-
-import yaml
 
 
 REPO_ROOT = Path(__file__).parents[2]
-CONFIG_YML = REPO_ROOT / ".circleci" / "config.yml"
 WORKFLOWS_DIR = REPO_ROOT / ".github" / "workflows"
 
 
-WORKFLOWS_TO_CHECK = [
-    "binary_builds",
-    "build",
-    "master_build",
-    # These are formatted slightly differently, skip them
-    # "scheduled-ci",
-    # "debuggable-scheduled-ci",
-    # "slow-gradcheck-scheduled-ci",
-    # "promote",
-]
-
-
-def add_job(
-    workflows: dict[str, Any],
-    workflow_name: str,
-    type: str,
-    job: dict[str, Any],
-    past_jobs: dict[str, Any],
-) -> None:
-    """
-    Add job 'job' under 'type' and 'workflow_name' to 'workflow' in place. Also
-    add any dependencies (they must already be in 'past_jobs')
-    """
-    if workflow_name not in workflows:
-        workflows[workflow_name] = {"when": "always", "jobs": []}
-
-    requires = job.get("requires")
-    if requires is not None:
-        for requirement in requires:
-            dependency = past_jobs[requirement]
-            add_job(
-                workflows,
-                dependency["workflow_name"],
-                dependency["type"],
-                dependency["job"],
-                past_jobs,
-            )
-
-    workflows[workflow_name]["jobs"].append({type: job})
-
-
-def get_filtered_circleci_config(
-    workflows: dict[str, Any], relevant_jobs: list[str]
-) -> dict[str, Any]:
-    """
-    Given an existing CircleCI config, remove every job that's not listed in
-    'relevant_jobs'
-    """
-    new_workflows: dict[str, Any] = {}
-    past_jobs: dict[str, Any] = {}
-    for workflow_name, workflow in workflows.items():
-        if workflow_name not in WORKFLOWS_TO_CHECK:
-            # Don't care about this workflow, skip it entirely
-            continue
-
-        for job_dict in workflow["jobs"]:
-            for type, job in job_dict.items():
-                if "name" not in job:
-                    # Job doesn't have a name so it can't be handled
-                    print("Skipping", type)
-                else:
-                    if job["name"] in relevant_jobs:
-                        # Found a job that was specified at the CLI, add it to
-                        # the new result
-                        add_job(new_workflows, workflow_name, type, job, past_jobs)
-
-                    # Record the job in case it's needed as a dependency later
-                    past_jobs[job["name"]] = {
-                        "workflow_name": workflow_name,
-                        "type": type,
-                        "job": job,
-                    }
-
-    return new_workflows
-
-
 def commit_ci(files: list[str], message: str) -> None:
     # Check that there are no other modified files than the ones edited by this
     # tool
@@ -114,9 +34,8 @@ def commit_ci(files: list[str], message: str) -> None:
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser(
-        description="make .circleci/config.yml only have a specific set of jobs and delete GitHub actions"
+        description="delete GitHub Actions workflow files that don't match a filter"
     )
-    parser.add_argument("--job", action="append", help="job name", default=[])
     parser.add_argument(
         "--filter-gha", help="keep only these github actions (glob match)", default=""
     )
@@ -127,16 +46,7 @@ def commit_ci(files: list[str], message: str) -> None:
     )
     args = parser.parse_args()
 
-    touched_files = [CONFIG_YML]
-    with open(CONFIG_YML) as f:
-        config_yml = yaml.safe_load(f.read())
-
-    config_yml["workflows"] = get_filtered_circleci_config(
-        config_yml["workflows"], args.job
-    )
-
-    with open(CONFIG_YML, "w") as f:
-        yaml.dump(config_yml, f)
+    touched_files: list[Path] = []
 
     if args.filter_gha:
         for relative_file in WORKFLOWS_DIR.iterdir():
@@ -146,13 +56,9 @@ def commit_ci(files: list[str], message: str) -> None:
                 path.resolve().unlink()
 
     if args.make_commit:
-        jobs_str = "\n".join([f" * {job}" for job in args.job])
         message = textwrap.dedent(
-            f"""
-        [skip ci][do not merge] Edit config.yml to filter specific jobs
-
-        Filter CircleCI to only run:
-        {jobs_str}
+            """
+        [skip ci][do not merge] Filter GitHub Actions workflow files
 
         See [Run Specific CI Jobs](https://github.com/pytorch/pytorch/blob/master/CONTRIBUTING.md#run-specific-ci-jobs) for details.
         """
diff --git a/tools/testing/modulefinder_determinator.py b/tools/testing/modulefinder_determinator.py
index e0ef858b96b21..8d87fa4e07f70 100644
--- a/tools/testing/modulefinder_determinator.py
+++ b/tools/testing/modulefinder_determinator.py
@@ -117,7 +117,7 @@ def test_impact_of_file(filename: str) -> str:
         CI - CI configuration files
     """
     parts = filename.split(os.sep)
-    if parts[0] in [".jenkins", ".circleci", ".ci"]:
+    if parts[0] in [".jenkins", ".ci"]:
         return "CI"
     if parts[0] in ["docs", "scripts", "CODEOWNERS", "README.md"]:
         return "NONE"
diff --git a/torch/_C/__init__.pyi.in b/torch/_C/__init__.pyi.in
index 565a2035fc663..caa44a14f7f06 100644
--- a/torch/_C/__init__.pyi.in
+++ b/torch/_C/__init__.pyi.in
@@ -147,6 +147,7 @@ class Stream:
     ) -> Self: ...
     def query(self) -> _bool: ...
     def synchronize(self) -> None: ...
+    def is_capturing(self) -> _bool: ...
     def wait_event(self, event: Event) -> None: ...
     def wait_stream(self, other: Stream) -> None: ...
     def record_event(self, event: Event | None = None) -> Event: ...
@@ -1997,6 +1998,7 @@ def _DTensor_compute_global_tensor_info(
 ) -> tuple[list[_int], list[_int]]: ...
 def _get_DTensor_sharding_propagator_cache_stats() -> tuple[_int, _int]: ...
 def _clear_DTensor_sharding_propagator_cache() -> None: ...
+def _reinit_DTensor_dispatch_logger() -> None: ...
 
 # Defined in torch/csrc/multiprocessing/init.cpp
 def _multiprocessing_init() -> None: ...
@@ -2036,6 +2038,10 @@ def _mtia_attachOutOfMemoryObserver(
     observer: Callable[[_int, _int, _int, _int], None],
 ) -> None: ...
 def _mtia_getDeviceCount() -> _int: ...
+def _mtia_getDevice() -> _int: ...
+def _mtia_setDevice(device: _int) -> None: ...
+def _mtia_exchangeDevice(device: _int) -> _int: ...
+def _mtia_maybeExchangeDevice(device: _int) -> _int: ...
 def _mtia_resetPeakMemoryStats(device: _int) -> None: ...
 def _mtia_graphPoolHandle() -> tuple[_int, _int]: ...
 
@@ -2082,6 +2088,7 @@ def _cuda_getDefaultStream(device: _int) -> tuple: ...
 def _cuda_getStreamFromExternal(data_ptr: _int, device_index: _int) -> tuple: ...
 def _cuda_getCurrentBlasHandle() -> _int: ...
 def _cuda_clearCublasWorkspaces() -> None: ...
+def _cudnn_clear_dropout_state() -> None: ...
 def _cuda_setDevice(device: _int) -> None: ...
 def _cuda_exchangeDevice(device: _int) -> _int: ...
 def _cuda_maybeExchangeDevice(device: _int) -> _int: ...
@@ -2101,6 +2108,7 @@ def _host_emptyCache() -> None: ...
 def _cuda_cudaCachingAllocator_raw_alloc(size: _int, cuda_stream: _int) -> _int: ...
 def _cuda_cudaCachingAllocator_raw_delete(ptr: _int) -> None: ...
 def _cuda_cudaCachingAllocator_enable(val: _bool) -> None: ...
+def _cuda_cudaCachingAllocator_is_enabled() -> _bool: ...
 def _cuda_beginAllocateToPool(device: _int, mempool_id: tuple[_int, _int]) -> None: ...
 def _cuda_beginAllocateCurrentThreadToPool(
     device: _int,
@@ -2158,6 +2166,7 @@ def _cuda_record_memory_history(
     global_record_annotations: _bool,
     skip_actions: list[str],
 ) -> None: ...
+def _cuda_allocationTraceback(data_ptr: _int) -> list[dict[str, Any]] | None: ...
 def _cuda_isHistoryEnabled() -> _bool: ...
 def _cuda_getAllocatorBackend() -> str: ...
 
@@ -2181,6 +2190,7 @@ def _construct_CUDA_Tensor_From_Storage_And_Metadata(
 ) -> Tensor: ...
 def _set_storage_access_error_msg(t: Tensor, s: str) -> None: ...
 def _set_storage_data_ptr_access_error_msg(storage_ptr: _int, s: str) -> None: ...
+def _clear_storage_data_ptr_access_error_msg(storage_ptr: _int) -> None: ...
 def _free_And_Remove_DeleterFn(storage_ptr: _int) -> None: ...
 def _has_Standard_Deleter(storage_ptr: _int) -> _bool: ...
 
@@ -2189,8 +2199,6 @@ class _cuda_CUDAAllocator: ...
 def _cuda_customAllocator(alloc_fn: _int, free_fn: _int) -> _cuda_CUDAAllocator: ...
 def _cuda_changeCurrentAllocator(allocator: _cuda_CUDAAllocator) -> None: ...
 def _cuda_getAllocator() -> _cuda_CUDAAllocator: ...
-def _cuda_lock_mutex() -> None: ...
-def _cuda_unlock_mutex() -> None: ...
 def _cuda_canDeviceAccessPeer(device: _int, peer_device: _int) -> _bool: ...
 def _cuda_jiterator_compile_and_launch_kernel(
     code_string: str,
@@ -2316,6 +2324,7 @@ def _is_flash_attention_available() -> _bool: ...
 def _can_use_cudnn_attention(params: _SDPAParams, debug: _bool) -> _bool: ...
 def _can_use_flash_attention(params: _SDPAParams, debug: _bool) -> _bool: ...
 def _can_use_mem_efficient_attention(params: _SDPAParams, debug: _bool) -> _bool: ...
+def _is_ck_sdpa_available() -> _bool: ...
 
 # Defined in torch/csrc/cuda/GdsFile.cpp
 def _gds_register_buffer(t: Storage) -> None: ...
@@ -2413,6 +2422,10 @@ class _CUDAGraph:
     def debug_dump(self, debug_path: str) -> None: ...
     def raw_cuda_graph(self) -> _int: ...
     def raw_cuda_graph_exec(self) -> _int: ...
+    @staticmethod
+    def get_currently_capturing_graph() -> _CUDAGraph: ...
+    def begin_capture_to_if_node(self, scalar_cuda_pred_tensor): ...
+    def end_capture_to_conditional_node(self): ...
 
 # Defined in torch/csrc/cuda/MemPool.cpp
 class _MemPool:
@@ -2481,6 +2494,8 @@ class _XpuDeviceProperties:
     gpu_eu_count: _int
     max_work_group_size: _int
     max_num_sub_groups: _int
+    memory_clock_rate: _int
+    memory_bus_width: _int
     sub_group_sizes: list[_int]
     local_mem_size: _int
     has_fp16: _bool
@@ -2594,12 +2609,32 @@ def _accelerator_exchangeDevice(device_index: _int) -> _int: ...
 def _accelerator_maybeExchangeDevice(device_index: _int) -> _int: ...
 def _accelerator_isAllocatorInitialized() -> _bool: ...
 def _accelerator_emptyCache() -> None: ...
+def _accelerator_emptyHostCache() -> None: ...
 def _accelerator_getDeviceStats(device_index: _int) -> dict[str, Any]: ...
 def _accelerator_resetAccumulatedStats(device_index: _int) -> None: ...
 def _accelerator_resetPeakStats(device_index: _int) -> None: ...
 def _accelerator_getMemoryInfo(device_index: _int) -> tuple[_int, _int]: ...
+def _accelerator_getAllocatorSettings() -> str: ...
 def _accelerator_setAllocatorSettings(env: str) -> None: ...
 
+class _acceleratorGraph:
+    def __new__(cls, _bool) -> Self: ...
+    def __init__(self, _bool) -> None: ...
+    def capture_begin(
+        self,
+        pool: tuple[_int, _int] | None = None,
+        capture_error_mode: Literal[
+            "default", "global", "thread_local", "relaxed"
+        ] = "default",
+    ) -> None: ...
+    def capture_end(self) -> None: ...
+    def instantiate(self) -> None: ...
+    def replay(self) -> None: ...
+    def reset(self) -> None: ...
+    def pool(self) -> tuple[_int, _int]: ...
+    def enable_debug_mode(self) -> None: ...
+    def debug_dump(self, path: str) -> None: ...
+
 # Defined in torch/csrc/jit/python/python_tracer.cpp
 class TracingState:
     def push_scope(self, scope_name: str) -> None: ...
diff --git a/torch/_C/_autograd.pyi b/torch/_C/_autograd.pyi
index 1ff5d847b61aa..733fb3410db9b 100644
--- a/torch/_C/_autograd.pyi
+++ b/torch/_C/_autograd.pyi
@@ -62,6 +62,7 @@ class _KinetoEvent:
     def duration_ns(self) -> int: ...
     def is_async(self) -> bool: ...
     def linked_correlation_id(self) -> int: ...
+    def external_id(self) -> int: ...
     def shapes(self) -> list[list[int]]: ...
     def dtypes(self) -> list[str]: ...
     def concrete_inputs(self) -> list[Any]: ...
@@ -78,8 +79,19 @@ class _KinetoEvent:
     def cuda_elapsed_us(self) -> int: ...
     def privateuse1_elapsed_us(self) -> int: ...
     def is_user_annotation(self) -> bool: ...
+    def is_python_function(self) -> bool: ...
     def is_hidden_event(self) -> bool: ...
     def metadata_json(self) -> str: ...
+    def activity_type(self) -> str: ...
+    def extra_meta(self) -> dict[str, str]: ...
+    def flow_id(self) -> int: ...
+    def flow_type(self) -> int: ...
+    def flow_start(self) -> bool: ...
+    def structured_input_shapes(self) -> list[list[int] | list[list[int]]]: ...
+    def structured_input_strides(self) -> list[list[int] | list[list[int]]]: ...
+    def python_id(self) -> int: ...
+    def python_parent_id(self) -> int: ...
+    def python_module_id(self) -> int: ...
 
 class _ProfilerResult:
     def events(self) -> list[_KinetoEvent]: ...
@@ -88,8 +100,14 @@ class _ProfilerResult:
     def experimental_event_tree(self) -> list[_ProfilerEvent]: ...
     def trace_start_ns(self) -> int: ...
 
-class SavedTensor: ...
+class SavedTensor:
+    def unpack(self) -> torch.Tensor: ...
 
+def _make_saved_tensor(
+    tensor: torch.Tensor,
+    is_output: bool,
+    is_inplace_on_view: bool = False,
+) -> SavedTensor: ...
 def _enable_profiler(
     config: ProfilerConfig,
     activities: set[ProfilerActivity],
@@ -97,6 +115,7 @@ def _enable_profiler(
 def _prepare_profiler(
     config: ProfilerConfig,
     activities: set[ProfilerActivity],
+    activity_filter: dict[ProfilerActivity, set[str]] = ...,
 ) -> None: ...
 def _toggle_collection_dynamic(
     enable: bool,
diff --git a/torch/_C/_distributed_c10d.pyi b/torch/_C/_distributed_c10d.pyi
index c0f2320028a44..972b88005db33 100644
--- a/torch/_C/_distributed_c10d.pyi
+++ b/torch/_C/_distributed_c10d.pyi
@@ -3,7 +3,7 @@
 from collections.abc import Callable
 from datetime import timedelta
 from enum import Enum
-from typing import Any, Optional, overload, Union
+from typing import Any, overload
 
 import torch
 from torch import Tensor
@@ -55,6 +55,7 @@ class Reducer:
         skip_all_reduce_unused_params: bool = ...,
         use_python_reducer: bool = ...,
         bucket_bytes_cap_list: list[int] = ...,
+        batched_grad_copy: bool = ...,
     ) -> None: ...
     def prepare_for_forward(self) -> None: ...
     def prepare_for_backward(self, output: list[Tensor]) -> None: ...
@@ -158,7 +159,7 @@ class AllreduceOptions:
     reduceOp: ReduceOp
     timeout: timedelta
     asyncOp: bool
-    sparseIndices: Optional[Tensor]
+    sparseIndices: Tensor | None
 
 class AllreduceCoalescedOptions(AllreduceOptions): ...
 
@@ -218,7 +219,7 @@ class Store:
     @overload
     def wait(self, keys: list[str], timeout: timedelta) -> None: ...
     def queue_pop(self, key: str, block: bool = True) -> bytes: ...
-    def queue_push(self, key: str, value: Union[bytes, str]) -> None: ...
+    def queue_push(self, key: str, value: bytes | str) -> None: ...
     def queue_len(self, key: str) -> int: ...
     def list_keys(self) -> list[str]: ...
 
@@ -375,18 +376,18 @@ class ProcessGroup:
     def split_group(
         self,
         new_ranks: list[int],
-        timeout: Optional[timedelta] = None,
-        opts: Optional[Backend.Options] = None,
+        timeout: timedelta | None = None,
+        opts: Backend.Options | None = None,
         group_name: GroupName | None = None,
-        group_desc: Optional[str] = None,
-    ) -> Optional[ProcessGroup]: ...
+        group_desc: str | None = None,
+    ) -> ProcessGroup | None: ...
     def merge_remote_group(
         self,
         store: Store,
         size: int,
         timeout: timedelta,
         group_name: GroupName | None = None,
-        group_desc: Optional[str] = None,
+        group_desc: str | None = None,
     ) -> ProcessGroup: ...
     def abort(self) -> None: ...
     def set_timeout(self, timeout: timedelta) -> None: ...
@@ -795,7 +796,9 @@ class _SymmetricMemory:
     @staticmethod
     def set_backend(name: str) -> None: ...
     @staticmethod
-    def get_backend(device: torch.device) -> Optional[str]: ...
+    def get_backend(device: torch.device) -> str | None: ...
+    @staticmethod
+    def is_symm_mem_tensor(tensor: torch.Tensor) -> bool: ...
     @staticmethod
     def get_mempool_allocator(device: torch.device) -> Any: ...
     signal_pad_size: int
@@ -893,3 +896,5 @@ class _Response:
 def _register_handler(
     name: str, handler: Callable[[_Request, _Response], None]
 ) -> None: ...
+def _set_comm_profiling_name(name: str) -> None: ...
+def _get_comm_profiling_name() -> str: ...
diff --git a/torch/_C/_dynamo/__init__.pyi b/torch/_C/_dynamo/__init__.pyi
index 67d515697cbe4..ab9470c2afac9 100644
--- a/torch/_C/_dynamo/__init__.pyi
+++ b/torch/_C/_dynamo/__init__.pyi
@@ -2,3 +2,68 @@ from . import compiled_autograd, eval_frame, guards  # noqa: F401
 
 def strip_function_call(name: str) -> str: ...
 def is_valid_var_name(name: str) -> bool | int: ...
+def get_type_slots(obj: type | object) -> tuple[int, int, int, int]: ...
+def has_slot(slots: int, slot_bit: int) -> bool: ...
+
+class PySequenceSlots:
+    SQ_LENGTH: int
+    SQ_CONCAT: int
+    SQ_REPEAT: int
+    SQ_ITEM: int
+    SQ_CONTAINS: int
+    SQ_ASS_ITEM: int
+    SQ_INPLACE_CONCAT: int
+    SQ_INPLACE_REPEAT: int
+
+class PyMappingSlots:
+    MP_LENGTH: int
+    MP_SUBSCRIPT: int
+    MP_ASS_SUBSCRIPT: int
+
+class PyNumberSlots:
+    NB_ADD: int
+    NB_SUBTRACT: int
+    NB_MULTIPLY: int
+    NB_REMAINDER: int
+    NB_POWER: int
+    NB_NEGATIVE: int
+    NB_POSITIVE: int
+    NB_ABSOLUTE: int
+    NB_BOOL: int
+    NB_INVERT: int
+    NB_LSHIFT: int
+    NB_RSHIFT: int
+    NB_AND: int
+    NB_XOR: int
+    NB_OR: int
+    NB_INT: int
+    NB_FLOAT: int
+    NB_INPLACE_ADD: int
+    NB_INPLACE_SUBTRACT: int
+    NB_INPLACE_MULTIPLY: int
+    NB_INPLACE_REMAINDER: int
+    NB_INPLACE_POWER: int
+    NB_INPLACE_LSHIFT: int
+    NB_INPLACE_RSHIFT: int
+    NB_INPLACE_AND: int
+    NB_INPLACE_XOR: int
+    NB_INPLACE_OR: int
+    NB_FLOOR_DIVIDE: int
+    NB_TRUE_DIVIDE: int
+    NB_INPLACE_FLOOR_DIVIDE: int
+    NB_INPLACE_TRUE_DIVIDE: int
+    NB_INDEX: int
+    NB_MATRIX_MULTIPLY: int
+    NB_INPLACE_MATRIX_MULTIPLY: int
+
+class PyTypeSlots:
+    TP_HASH: int
+    TP_ITER: int
+    TP_ITERNEXT: int
+    TP_CALL: int
+    TP_REPR: int
+    TP_RICHCOMPARE: int
+    TP_GETATTRO: int
+    TP_SETATTRO: int
+    TP_DESCR_GET: int
+    TP_DESCR_SET: int
diff --git a/torch/_C/_dynamo/eval_frame.pyi b/torch/_C/_dynamo/eval_frame.pyi
index 90f96ad3fb1f9..a57eae9fd8779 100644
--- a/torch/_C/_dynamo/eval_frame.pyi
+++ b/torch/_C/_dynamo/eval_frame.pyi
@@ -1,7 +1,7 @@
 import enum
 import types
 from collections.abc import Callable
-from typing import Optional, overload
+from typing import overload
 
 from torch._dynamo.guards import GuardManagerWrapper
 from torch._dynamo.types import DynamoCallback, DynamoGuardCompleteHook, DynamoGuardHook
@@ -17,8 +17,8 @@ def set_code_exec_strategy(
 ) -> None: ...
 def set_guard_error_hook(hook: DynamoGuardHook) -> None: ...
 def set_guard_complete_hook(
-    hook: Optional[DynamoGuardCompleteHook],
-) -> Optional[DynamoGuardCompleteHook]: ...
+    hook: DynamoGuardCompleteHook | None,
+) -> DynamoGuardCompleteHook | None: ...
 def raise_sigtrap() -> None: ...
 def set_c_recursion_limit(limit: int) -> None: ...
 def get_c_recursion_limit() -> int: ...
@@ -71,6 +71,19 @@ class _PyInterpreterFrame:
     closure: tuple[types.CellType]
 
 def _debug_get_cache_entry_list(code: types.CodeType) -> list[_CacheEntry]: ...
+def _get_frame_value_stack_with_depth(
+    frame: types.FrameType, depth: int
+) -> list[object]: ...
+def set_bytecode_debugger_callback(
+    callback: Callable[[types.CodeType], None] | None,
+) -> None: ...
+def get_bytecode_debugger_callback() -> Callable[[types.CodeType], None] | None: ...
+def register_breakpoint_code(code: types.CodeType) -> None: ...
+
+# Sentinel for NULL stack values returned by _get_frame_value_stack_at_depth
+class NullStackValue: ...
+
+NULL_STACK_VALUE: NullStackValue
 
 py_opcode_caches: list[int]
 
@@ -82,10 +95,5 @@ def _load_precompile_entry(
 ) -> None: ...
 def _reset_precompile_entries(code: types.CodeType) -> None: ...
 def _debug_get_precompile_entries(code: types.CodeType) -> list[_PrecompileEntry]: ...
-
-class _EvalFrameOverride(enum.IntEnum):
-    NONE = 0
-    SKIP = 1
-    ERROR = 2
-
-def set_eval_frame_override(override: _EvalFrameOverride) -> _EvalFrameOverride: ...
+def set_fullgraph_compiled_frame_count(value: int) -> int: ...
+def set_fullgraph_error_on_nested_compile(value: bool) -> bool: ...
diff --git a/torch/_C/_dynamo/guards.pyi b/torch/_C/_dynamo/guards.pyi
index e588cb437f91a..10845355e5bfa 100644
--- a/torch/_C/_dynamo/guards.pyi
+++ b/torch/_C/_dynamo/guards.pyi
@@ -1,7 +1,7 @@
 import enum
 import traceback
 from collections.abc import Callable
-from typing import Any, Optional, TypeAlias
+from typing import Any, TypeAlias
 
 import torch
 
@@ -23,7 +23,7 @@ class GuardDebugInfo:
     verbose_code_parts: list[str]
     result: bool
     num_guards_executed: int
-    user_stack: Optional[traceback.StackSummary]
+    user_stack: traceback.StackSummary | None
 
 class GuardManager:
     def check(self, value: Any) -> bool: ...
@@ -227,75 +227,75 @@ class GuardManager:
         self,
         user_lambda: Callable[..., Any],
         verbose_code_parts: list[str],
-        user_stack: Optional[traceback.StackSummary],
+        user_stack: traceback.StackSummary | None,
     ) -> None: ...
     def add_id_match_guard(
         self,
         id_val: int,
         verbose_code_parts: list[str],
-        user_stack: Optional[traceback.StackSummary],
+        user_stack: traceback.StackSummary | None,
     ) -> None: ...
     def add_equals_match_guard(
         self,
         equals_val: Any,
         verbose_code_parts: list[str],
-        user_stack: Optional[traceback.StackSummary],
+        user_stack: traceback.StackSummary | None,
     ) -> None: ...
     def add_global_state_guard(
         self,
         initial_state: Any,
         verbose_code_parts: list[str],
-        user_stack: Optional[traceback.StackSummary],
+        user_stack: traceback.StackSummary | None,
     ) -> None: ...
     def add_torch_function_mode_stack_guard(
         self,
         initial_stack: list[Any],
         verbose_code_parts: list[str],
-        user_stack: Optional[traceback.StackSummary],
+        user_stack: traceback.StackSummary | None,
     ) -> None: ...
     def add_mapping_keys_guard(
         self,
         value: Any,
         verbose_code_parts: list[str],
-        user_stack: Optional[traceback.StackSummary],
+        user_stack: traceback.StackSummary | None,
     ) -> None: ...
     def add_dict_length_check_guard(
         self,
         value: int,
         verbose_code_parts: list[str],
-        user_stack: Optional[traceback.StackSummary],
+        user_stack: traceback.StackSummary | None,
     ) -> None: ...
     def add_length_check_guard(
         self,
         value: int,
         verbose_code_parts: list[str],
-        user_stack: Optional[traceback.StackSummary],
+        user_stack: traceback.StackSummary | None,
     ) -> None: ...
     def add_true_match_guard(
         self,
         verbose_code_parts: list[str],
-        user_stack: Optional[traceback.StackSummary],
+        user_stack: traceback.StackSummary | None,
     ) -> None: ...
     def add_false_match_guard(
         self,
         verbose_code_parts: list[str],
-        user_stack: Optional[traceback.StackSummary],
+        user_stack: traceback.StackSummary | None,
     ) -> None: ...
     def add_none_match_guard(
         self,
         verbose_code_parts: list[str],
-        user_stack: Optional[traceback.StackSummary],
+        user_stack: traceback.StackSummary | None,
     ) -> None: ...
     def add_not_none_guard(
         self,
         verbose_code_parts: list[str],
-        user_stack: Optional[traceback.StackSummary],
+        user_stack: traceback.StackSummary | None,
     ) -> None: ...
     def add_dispatch_key_set_guard(
         self,
         dispatch_key: Any,
         verbose_code_parts: list[str],
-        user_stack: Optional[traceback.StackSummary],
+        user_stack: traceback.StackSummary | None,
     ) -> None: ...
     def add_tensor_match_guard(
         self,
@@ -304,7 +304,7 @@ class GuardManager:
         strides: list[int],
         tensor_name: str,
         verbose_code_parts: list[str],
-        user_stack: Optional[traceback.StackSummary],
+        user_stack: traceback.StackSummary | None,
         ptype: Any,
         dispatch_keys: Any,
     ) -> None: ...
@@ -312,62 +312,62 @@ class GuardManager:
         self,
         value: set[Any],
         verbose_code_parts: list[str],
-        user_stack: Optional[traceback.StackSummary],
+        user_stack: traceback.StackSummary | None,
     ) -> None: ...
     def add_no_hasattr_guard(
         self,
         attr_name: str,
         verbose_code_parts: list[str],
-        user_stack: Optional[traceback.StackSummary],
+        user_stack: traceback.StackSummary | None,
     ) -> None: ...
     def add_dict_contains_guard(
         self,
         contains: bool,
         key: Any,
         verbose_code_parts: list[str],
-        user_stack: Optional[traceback.StackSummary],
+        user_stack: traceback.StackSummary | None,
     ) -> None: ...
     def add_type_match_guard(
         self,
         value: int,
         verbose_code_parts: list[str],
-        user_stack: Optional[traceback.StackSummary],
+        user_stack: traceback.StackSummary | None,
     ) -> None: ...
     def add_dict_version_guard(
         self,
         value: Any,
         verbose_code_parts: list[str],
-        user_stack: Optional[traceback.StackSummary],
+        user_stack: traceback.StackSummary | None,
     ) -> None: ...
     def add_set_contains_guard(
         self,
         contains: bool,
         item: Any,
         verbose_code_parts: list[str],
-        user_stack: Optional[traceback.StackSummary],
+        user_stack: traceback.StackSummary | None,
     ) -> None: ...
     def add_dual_level_match_guard(
         self,
         level: int,
         verbose_code_parts: list[str],
-        user_stack: Optional[traceback.StackSummary],
+        user_stack: traceback.StackSummary | None,
     ) -> None: ...
     def add_float_is_nan_guard(
         self,
         verbose_code_parts: list[str],
-        user_stack: Optional[traceback.StackSummary],
+        user_stack: traceback.StackSummary | None,
     ) -> None: ...
     def add_complex_is_nan_guard(
         self,
         verbose_code_parts: list[str],
-        user_stack: Optional[traceback.StackSummary],
+        user_stack: traceback.StackSummary | None,
     ) -> None: ...
     def add_tuple_iterator_length_guard(
         self,
         length: int,
         type_id: int,
         verbose_code_parts: list[str],
-        user_stack: Optional[traceback.StackSummary],
+        user_stack: traceback.StackSummary | None,
     ) -> None: ...
     def add_range_iterator_match_guard(
         self,
@@ -376,12 +376,12 @@ class GuardManager:
         step: int,
         type_id: int,
         verbose_code_parts: list[str],
-        user_stack: Optional[traceback.StackSummary],
+        user_stack: traceback.StackSummary | None,
     ) -> None: ...
     def add_default_device_guard(
         self,
         verbose_code_parts: list[str],
-        user_stack: Optional[traceback.StackSummary],
+        user_stack: traceback.StackSummary | None,
     ) -> None: ...
     def mark_tag_safe(self) -> None: ...
     def mark_tag_safe_root(self) -> None: ...
@@ -392,7 +392,7 @@ class RootGuardManager(GuardManager):
         self,
         guard: LeafGuard,
         verbose_code_parts: list[str],
-        user_stack: Optional[traceback.StackSummary],
+        user_stack: traceback.StackSummary | None,
     ) -> None: ...
     def clone_manager(
         self, clone_filter_fn: Callable[[GuardManager], bool]
@@ -438,19 +438,19 @@ def install_object_aliasing_guard(
     x: GuardManager,
     y: GuardManager,
     verbose_code_parts: list[str],
-    user_stack: Optional[traceback.StackSummary],
+    user_stack: traceback.StackSummary | None,
 ) -> None: ...
 def install_no_tensor_aliasing_guard(
     guard_managers: list[GuardManager],
     tensor_names: list[str],
     verbose_code_parts: list[str],
-    user_stack: Optional[traceback.StackSummary],
+    user_stack: traceback.StackSummary | None,
 ) -> None: ...
 def install_storage_overlapping_guard(
     overlapping_guard_managers: list[GuardManager],
     non_overlapping_guard_managers: list[GuardManager],
     verbose_code_parts: list[str],
-    user_stack: Optional[traceback.StackSummary],
+    user_stack: traceback.StackSummary | None,
 ) -> None: ...
 def install_symbolic_shape_guard(
     guard_managers: list[GuardManager],
@@ -459,7 +459,7 @@ def install_symbolic_shape_guard(
     py_addr: int,
     py_addr_keep_alive: Any,
     verbose_code_parts: list[str],
-    user_stack: Optional[traceback.StackSummary],
+    user_stack: traceback.StackSummary | None,
 ) -> None: ...
 def profile_guard_manager(
     guard_manager: GuardManager,
@@ -476,7 +476,7 @@ class TensorGuards:
     ) -> None: ...
     def check(self, *args: Any) -> bool: ...
     def check_verbose(
-        self, *args: Any, tensor_check_names: Optional[list[str]] = None
+        self, *args: Any, tensor_check_names: list[str] | None = None
     ) -> bool | str: ...
 
 def assert_size_stride(
@@ -490,6 +490,7 @@ def assert_alignment(
     alignment: int,
     op_name: str | None = None,
 ) -> None: ...
+def copy_misaligned(item: torch.Tensor) -> torch.Tensor: ...
 def check_obj_id(obj: object, expected: int) -> bool: ...
 def check_type_id(obj: object, expected: int) -> bool: ...
 def dict_version(d: dict[Any, Any]) -> int: ...
diff --git a/torch/_C/_jit_tree_views.pyi b/torch/_C/_jit_tree_views.pyi
index cf4cffc05a9c3..bc7f7c3b7c1e9 100644
--- a/torch/_C/_jit_tree_views.pyi
+++ b/torch/_C/_jit_tree_views.pyi
@@ -1,4 +1,4 @@
-from typing import Any, Optional
+from typing import Any
 
 # Defined in torch/csrc/jit/python/python_tree_views.cpp
 
@@ -32,7 +32,7 @@ class Ident(TreeView):
     def name(self) -> str: ...
 
 class Param(TreeView):
-    def __init__(self, type: Optional[Any], name: Ident, kwarg_only: bool) -> None: ...
+    def __init__(self, type: Any | None, name: Ident, kwarg_only: bool) -> None: ...
 
 class Attribute(TreeView):
     def __init__(self, name: Ident, value: Any) -> None: ...
@@ -55,11 +55,11 @@ class Def(TreeView):
 
 class Property(TreeView):
     def __init__(
-        self, r: SourceRange, name: Ident, getter: Def, setter: Optional[Def]
+        self, r: SourceRange, name: Ident, getter: Def, setter: Def | None
     ) -> None: ...
     def name(self) -> Ident: ...
     def getter_name(self) -> str: ...
-    def setter_name(self) -> Optional[Ident]: ...
+    def setter_name(self) -> Ident | None: ...
 
 class ClassDef(TreeView):
     def __init__(
@@ -68,33 +68,31 @@ class ClassDef(TreeView):
 
 class Decl(TreeView):
     def __init__(
-        self, r: SourceRange, params: list[Param], return_type: Optional[Expr]
+        self, r: SourceRange, params: list[Param], return_type: Expr | None
     ) -> None: ...
 
 class Delete(Stmt):
     def __init__(self, range: SourceRange, targets: list[Expr]) -> None: ...
 
 class WithItem(Expr):
-    def __init__(
-        self, range: SourceRange, target: Expr, var: Optional[Any]
-    ) -> None: ...
+    def __init__(self, range: SourceRange, target: Expr, var: Any | None) -> None: ...
 
 class Assign(Stmt):
     def __init__(
-        self, lhs: list[Expr], rhs: Expr, type: Optional[Expr] = None
+        self, lhs: list[Expr], rhs: Expr, type: Expr | None = None
     ) -> None: ...
 
 class AugAssign(Stmt):
     def __init__(self, lhs: Expr, kind_str: str, rhs: Expr) -> None: ...
 
 class Return(Stmt):
-    def __init__(self, range: SourceRange, value: Optional[Expr]) -> None: ...
+    def __init__(self, range: SourceRange, value: Expr | None) -> None: ...
 
 class Raise(Stmt):
     def __init__(self, range: SourceRange, expr: Expr) -> None: ...
 
 class Assert(Stmt):
-    def __init__(self, range: SourceRange, test: Expr, msg: Optional[Expr]) -> None: ...
+    def __init__(self, range: SourceRange, test: Expr, msg: Expr | None) -> None: ...
 
 class Pass(Stmt):
     def __init__(self, range: SourceRange) -> None: ...
@@ -190,9 +188,9 @@ class SliceExpr(Expr):
     def __init__(
         self,
         range: SourceRange,
-        lower: Optional[Expr],
-        upper: Optional[Expr],
-        step: Optional[Expr],
+        lower: Expr | None,
+        upper: Expr | None,
+        step: Expr | None,
     ) -> None: ...
 
 class Starred(Expr):
diff --git a/torch/_C/_nn.pyi.in b/torch/_C/_nn.pyi.in
index 7be3dcff4da67..cf6e98fe42e30 100644
--- a/torch/_C/_nn.pyi.in
+++ b/torch/_C/_nn.pyi.in
@@ -129,6 +129,12 @@ def _upsample_bicubic2d_aa(
     align_corners: bool,
     scale_factors: Sequence[float] | None,
 ) -> Tensor: ...
+def _upsample_lanczos2d_aa(
+    input: Tensor,
+    output_size: Sequence[int] | None,
+    align_corners: bool,
+    scale_factors: Sequence[float] | None,
+) -> Tensor: ...
 def upsample_bicubic2d(
     input: Tensor,
     output_size: Sequence[int] | None,
diff --git a/torch/_C/_profiler.pyi b/torch/_C/_profiler.pyi
index ae8121e4b71d2..04c3f4a460086 100644
--- a/torch/_C/_profiler.pyi
+++ b/torch/_C/_profiler.pyi
@@ -27,6 +27,7 @@ class ProfilerState(Enum):
     KINETO = ...
     KINETO_GPU_FALLBACK = ...
     KINETO_PRIVATEUSE1_FALLBACK = ...
+    KINETO_PRIVATEUSE1 = ...
 
 class ActiveProfilerType(Enum):
     NONE = ...
diff --git a/torch/__init__.py b/torch/__init__.py
index 9f344606796b8..1852fe9c5a921 100644
--- a/torch/__init__.py
+++ b/torch/__init__.py
@@ -29,6 +29,7 @@
     get_origin as _get_origin,
     overload as _overload,
     TYPE_CHECKING,
+    TypeGuard as _TypeGuard,
     TypeVar as _TypeVar,
 )
 from typing_extensions import (
@@ -132,6 +133,7 @@ def _running_with_deploy() -> builtins.bool:
     "sym_min",
     "sym_not",
     "sym_sum",
+    "thread_safe_generator",
     "typename",
     "unravel_index",
     "use_deterministic_algorithms",
@@ -953,11 +955,17 @@ def sym_min(a, b):
         return builtins.min(a, b)  # type: ignore[call-overload]
 
 
-def sym_sum(args):
+def sym_sum(*args):
     """
     N-ary add which is faster to compute for long lists than iterated binary
     addition.  Only does something special for integers.
+
+    Accepts both ``sym_sum([a, b, c])`` and ``sym_sum(a, b, c)``.
     """
+    # Normalise: accept both sym_sum([a, b, c]) and sym_sum(a, b, c).
+    if len(args) == 1 and isinstance(args[0], (list, tuple)):
+        args = args[0]
+
     if overrides.has_torch_function(args):
         return overrides.handle_torch_function(sym_sum, args, args)
 
@@ -1164,7 +1172,7 @@ def is_tensor(obj: _Any, /) -> _TypeIs["torch.Tensor"]:
     return isinstance(obj, torch.Tensor)
 
 
-def is_storage(obj: _Any, /) -> builtins.bool:
+def is_storage(obj: _Any, /) -> _TypeGuard["TypedStorage | UntypedStorage"]:
     r"""Returns True if `obj` is a PyTorch storage object.
 
     Args:
@@ -2144,7 +2152,14 @@ def _dtype(self):
 from torch import amp as amp, random as random, serialization as serialization
 from torch._tensor_str import set_printoptions
 from torch.amp import autocast, GradScaler
-from torch.random import get_rng_state, initial_seed, manual_seed, seed, set_rng_state
+from torch.random import (
+    get_rng_state,
+    initial_seed,
+    manual_seed,
+    seed,
+    set_rng_state,
+    thread_safe_generator,
+)
 from torch.serialization import load, save
 
 
@@ -2211,7 +2226,10 @@ def _manager_path():
 
 
 __all__.extend(
-    name for name in dir(torch) if isinstance(getattr(torch, name), torch.dtype)
+    # pyrefly: ignore [unresolvable-dunder-all]
+    name
+    for name in dir(torch)
+    if isinstance(getattr(torch, name), torch.dtype)
 )
 
 ################################################################################
@@ -2389,11 +2407,12 @@ def compiled_with_cxx11_abi() -> builtins.bool:
 class _TorchCompileInductorWrapper:
     compiler_name = "inductor"
 
-    def __init__(self, mode, options, dynamic):
+    def __init__(self, mode, options, dynamic, name=None):
         from torch._inductor.compiler_bisector import CompilerBisector
 
         self.config: dict[str, _Any] = {}
         self.dynamic = dynamic
+        self.name = name
         self.apply_mode(mode)
         self.apply_options(options)
         self.apply_options(CompilerBisector.get_config_change("inductor"))
@@ -2420,6 +2439,7 @@ def __eq__(self, other):
             isinstance(other, _TorchCompileInductorWrapper)
             and self.config == other.config
             and self.dynamic == other.dynamic
+            and self.name == other.name
         )
 
     def apply_mode(self, mode: str | None):
@@ -2455,10 +2475,16 @@ def apply_options(self, options: dict[str, _Any] | None):
                     )
             self.config[attr_name] = val
 
-    def __call__(self, model_, inputs_):
+    def __call__(self, model_, inputs_, *, config_patches=None):
         from torch._inductor.compile_fx import compile_fx
 
-        return compile_fx(model_, inputs_, config_patches=self.config)
+        all_patches = {**self.config, **(config_patches or {})}
+        return compile_fx(
+            model_,
+            inputs_,
+            config_patches=all_patches,
+            compile_region_name=self.name,
+        )
 
     def get_compiler_config(self):
         from torch._inductor.compile_fx import get_patched_config_dict
@@ -2478,12 +2504,12 @@ def reset(self):
 class _TorchCompileAOTInductorWrapper(_TorchCompileInductorWrapper):
     compiler_name = "aotinductor"
 
-    def __init__(self, mode, options, dynamic):
-        super().__init__(mode, options, dynamic)
+    def __init__(self, mode, options, dynamic, name=None):
+        super().__init__(mode, options, dynamic, name)
         self.apply_options({"cpp_wrapper": True})
         self.apply_options({"aot_inductor.package": True})
 
-    def __call__(self, model_, inputs_):
+    def __call__(self, model_, inputs_, *, config_patches=None):
         from contextlib import nullcontext
         from unittest import mock
 
@@ -2501,7 +2527,7 @@ def __call__(self, model_, inputs_):
             ctx,
             torch._inductor.config.patch("enable_autograd_for_aot", True),
         ):
-            return super().__call__(model_, inputs_)
+            return super().__call__(model_, inputs_, config_patches=config_patches)
 
 
 class _TorchCompileWrapper:
@@ -2552,6 +2578,7 @@ def compile(
     backend: str | _Callable = "inductor",
     mode: str | None = None,
     options: dict[str, str | builtins.int | builtins.bool | _Callable] | None = None,
+    name: str | None = None,
     disable: builtins.bool = False,
 ) -> _Callable[_InputT, _RetT]: ...
 
@@ -2565,6 +2592,7 @@ def compile(
     backend: str | _Callable = "inductor",
     mode: str | None = None,
     options: dict[str, str | builtins.int | builtins.bool | _Callable] | None = None,
+    name: str | None = None,
     disable: builtins.bool = False,
 ) -> _Callable[[_Callable[_InputT, _RetT]], _Callable[_InputT, _RetT]]: ...
 
@@ -2577,7 +2605,9 @@ def compile(
     backend: str | _Callable = "inductor",
     mode: str | None = None,
     options: dict[str, str | builtins.int | builtins.bool | _Callable] | None = None,
+    name: str | None = None,
     disable: builtins.bool = False,
+    recompile_limit: builtins.int | None = None,
 ) -> (
     _Callable[[_Callable[_InputT, _RetT]], _Callable[_InputT, _RetT]]
     | _Callable[_InputT, _RetT]
@@ -2667,6 +2697,8 @@ def compile(
           - `torch.compiler.keep_tensor_guards_unsafe`
 
         - For inductor you can see the full list of configs that it supports by calling `torch._inductor.list_options()`
+       name (str or None): Optional identifier for the compiled region. When supported by downstream
+        tooling, this is surfaced on wrapped compiled-region higher-order operators and other debug metadata.
        disable (bool): Turn torch.compile() into a no-op for testing
 
     Example::
@@ -2704,6 +2736,7 @@ def fn(model: _Callable[_InputT, _RetT]) -> _Callable[_InputT, _RetT]:
                 backend=backend,
                 mode=mode,
                 options=options,
+                name=name,
                 disable=disable,
             )
 
@@ -2749,9 +2782,9 @@ def fn(model: _Callable[_InputT, _RetT]) -> _Callable[_InputT, _RetT]:
 
     if backend == "inductor":
         if use_aoti:
-            backend = _TorchCompileAOTInductorWrapper(mode, options, dynamic)
+            backend = _TorchCompileAOTInductorWrapper(mode, options, dynamic, name)
         else:
-            backend = _TorchCompileInductorWrapper(mode, options, dynamic)
+            backend = _TorchCompileInductorWrapper(mode, options, dynamic, name)
     else:
         backend = _TorchCompileWrapper(backend, mode, options, dynamic)
 
@@ -2761,6 +2794,7 @@ def fn(model: _Callable[_InputT, _RetT]) -> _Callable[_InputT, _RetT]:
         dynamic=dynamic,
         disable=disable,
         guard_filter_fn=guard_filter_fn,
+        recompile_limit=recompile_limit,
     )(model)  # type: ignore[return-value]
 
 
@@ -2878,6 +2912,13 @@ def __getattr__(name):
         if name in _lazy_modules:
             return importlib.import_module(f".{name}", __name__)
 
+        # set_vital
+        if name == "set_vital":
+            import warnings
+
+            warnings.warn(f"'{name}' is deprecated, please do not call", stacklevel=2)
+            return lambda *args: None
+
         raise AttributeError(f"module '{__name__}' has no attribute '{name}'")
 
 
@@ -2999,3 +3040,6 @@ def _as_tensor_fullprec(t):
 # an autoloaded backend are defined
 if _is_device_backend_autoload_enabled():
     _import_device_backends()
+
+# Register all registered custom / override ops in torch/_native
+import torch._native
diff --git a/torch/_custom_op/impl.py b/torch/_custom_op/impl.py
index 3e495ef785e83..ae01dab906226 100644
--- a/torch/_custom_op/impl.py
+++ b/torch/_custom_op/impl.py
@@ -59,9 +59,7 @@ def warn_deprecated():
     )
 
 
-def custom_op(
-    qualname: str, manual_schema: typing.Optional[str] = None
-) -> typing.Callable:
+def custom_op(qualname: str, manual_schema: str | None = None) -> typing.Callable:
     r"""
     This API is deprecated, please use torch.library.custom_op instead
     """
@@ -155,7 +153,7 @@ def __init__(
         self.__name__ = None  # mypy requires this
         # NB: Some of these impls are registered as kernels to DispatchKeys.
         # Modifying the _impls dict directly won't do anything in that case.
-        self._impls: dict[str, typing.Optional[FuncAndLocation]] = {}
+        self._impls: dict[str, FuncAndLocation | None] = {}
         # See NOTE [CustomOp autograd kernel indirection]
         self._registered_autograd_kernel_indirection = False
 
@@ -220,7 +218,7 @@ def __call__(self, *args, **kwargs):
 
     def impl(
         self,
-        device_types: typing.Union[str, typing.Iterable[str]],
+        device_types: str | typing.Iterable[str],
         _stacklevel=2,
     ) -> typing.Callable:
         r"""
diff --git a/torch/_decomp/__init__.py b/torch/_decomp/__init__.py
index 5ccd693eae534..2553f9e60de96 100644
--- a/torch/_decomp/__init__.py
+++ b/torch/_decomp/__init__.py
@@ -68,7 +68,7 @@ def _add_op_to_registry(registry, op, fn):
     If op is OpOverload, it will be added to the registry directly.
     If op is OpOverloadPacket, all the valid op_overloads in the packet will be added to the registry.
     """
-    overloads: list[Union[torch._ops.OperatorBase]] = []
+    overloads: list[torch._ops.OperatorBase] = []
     if isinstance(op, HigherOrderOperator):
         # There's no concept of overloads for HigherOrderOperator
         registry[op] = fn
@@ -229,7 +229,7 @@ def register(op):
 
 
 def get_decompositions(
-    aten_ops: Sequence[Union[torch._ops.OperatorBase, OpOverloadPacket]],
+    aten_ops: Sequence[torch._ops.OperatorBase | OpOverloadPacket],
     type: str = "post_autograd",
 ) -> dict[torch._ops.OperatorBase, Callable]:
     """
@@ -266,7 +266,7 @@ def get_decompositions(
 
 def remove_decompositions(
     decompositions: dict[torch._ops.OperatorBase, Callable],
-    aten_ops: Sequence[Union[OpOverload, OpOverloadPacket]],
+    aten_ops: Sequence[OpOverload | OpOverloadPacket],
 ) -> None:
     """
     Given a dictionary of decompositions obtained from get_decompositions(), removes
@@ -371,6 +371,7 @@ def _core_aten_decompositions_post_autograd() -> dict[
             aten.hardswish_backward,
             aten.hardtanh_,
             aten.hardtanh_backward,
+            aten.hann_window,
             aten.heaviside,
             aten.heaviside_,
             aten.huber_loss,
diff --git a/torch/_decomp/decompositions.py b/torch/_decomp/decompositions.py
index 082a9f363bb15..4eb171d50ec93 100644
--- a/torch/_decomp/decompositions.py
+++ b/torch/_decomp/decompositions.py
@@ -10,7 +10,7 @@
 from enum import Enum
 from functools import partial, reduce
 from itertools import chain, product
-from typing import Any, cast, Optional, Union
+from typing import Any, cast
 
 import torch
 import torch._meta_registrations
@@ -330,7 +330,7 @@ def rrelu_with_noise_backward(
     training: bool,
     self_is_result: bool,
 ) -> Tensor:
-    if training and upper - lower > 1e-6:
+    if training:
         return grad_output.mul(noise)
     else:
         negative_slope = (lower + upper) / 2
@@ -346,10 +346,10 @@ def log_sigmoid_backward(grad_output: Tensor, self: Tensor, buffer: Tensor) -> T
     in_negative = self < 0
     max_deriv = torch.where(in_negative, 1, 0)
     sign = torch.where(in_negative, 1, -1)
-    z = torch.exp(-torch.abs(self))
+    # On CPU, forward saves buffer = exp(-|x|); reuse it to match the native kernel.
+    # On CUDA/XPU, buffer is empty so we recompute.
+    z = buffer if buffer.numel() > 0 else torch.exp(-torch.abs(self))
     return grad_output * (max_deriv - sign * (z / (1 + z)))
-    # CPU has a special formula that uses buffer, but disabled for convenience sake
-    # return (max_deriv - sign * (buffer / (1 + buffer))) * grad_output
 
 
 @register_decomposition(aten.ldexp)
@@ -497,7 +497,7 @@ def _nll_loss_backward(
     grad_output: Tensor,
     self: Tensor,
     target: Tensor,
-    weight: Optional[Tensor],
+    weight: Tensor | None,
     reduction: int,
     ignore_index: int,
     total_weight: Tensor,
@@ -506,6 +506,12 @@ def _nll_loss_backward(
     if reduction == Reduction.MEAN.value:
         grad_output = grad_output / total_weight
 
+    # When self is 1D (no batch dim), the C++ kernel only uses target[0].
+    # Reduce target to a scalar so the subsequent unsqueeze produces a 1D
+    # tensor matching self's dimensionality.
+    if self.dim() == 1 and target.dim() > 0:
+        target = target[0]
+
     target = target.unsqueeze(channel_dim)
     safe_target = torch.where(target != ignore_index, target, 0)
     grad_input = torch.zeros_like(self)
@@ -554,7 +560,7 @@ def nll_loss_backward(
     grad_output: Tensor,
     self: Tensor,
     target: Tensor,
-    weight: Optional[Tensor],
+    weight: Tensor | None,
     reduction: int,
     ignore_index: int,
     total_weight: Tensor,
@@ -605,7 +611,7 @@ def nll_loss2d_backward(
     grad_output: Tensor,
     self: Tensor,
     target: Tensor,
-    weight: Optional[Tensor],
+    weight: Tensor | None,
     reduction: int,
     ignore_index: int,
     total_weight: Tensor,
@@ -646,7 +652,7 @@ def nll_loss2d_backward(
 def binary_cross_entropy(
     self: Tensor,
     target: Tensor,
-    weight: Optional[Tensor] = None,
+    weight: Tensor | None = None,
     reduction: int = Reduction.MEAN.value,
 ) -> Tensor:
     # We cannot currently model this without introducing data-dependent control flow
@@ -669,7 +675,7 @@ def binary_cross_entropy_backward(
     grad_output: Tensor,
     self: Tensor,
     target: Tensor,
-    weight: Optional[Tensor] = None,
+    weight: Tensor | None = None,
     reduction: int = Reduction.MEAN.value,
 ) -> Tensor:
     EPSILON = 1e-12
@@ -746,8 +752,8 @@ def slice_forward(
     # Tensor(a) self, int dim=0, SymInt? start=None, SymInt? end=None, SymInt step=1
     self: Tensor,
     dim: int = 0,
-    start: Optional[int] = None,
-    end: Optional[int] = None,
+    start: int | None = None,
+    end: int | None = None,
     step: int = 1,
 ):
     from torch.fx.experimental.symbolic_shapes import statically_known_true
@@ -797,7 +803,7 @@ def slice_forward(
 
 
 def _normalize_start_end(
-    x: Tensor, dim: int, start: Optional[int], end: Optional[int]
+    x: Tensor, dim: int, start: int | None, end: int | None
 ) -> tuple[int, int]:
     """
     Normalize start and end such that both are in the range
@@ -825,8 +831,8 @@ def slice_scatter(
     input: Tensor,
     src: Tensor,
     dim: int = 0,
-    start: Optional[int] = None,
-    end: Optional[int] = None,
+    start: int | None = None,
+    end: int | None = None,
     step: int = 1,
 ):
     dim = utils.canonicalize_dim(input.ndim, dim)
@@ -840,7 +846,7 @@ def slice_scatter(
     if start == 0 and end == dim_size and step == 1:
         return src.clone()
 
-    indices: list[Optional[Tensor]] = [None] * input.dim()
+    indices: list[Tensor | None] = [None] * input.dim()
     idx = torch.arange(dim_size, device=input.device)
     indices[dim] = (idx - start) // step
 
@@ -1151,7 +1157,7 @@ def unfold_backward(
 @register_decomposition(aten.logit_backward.default)
 @pw_cast_for_opmath
 def logit_backward(
-    grad_output: Tensor, self: Tensor, eps: Optional[float] = None
+    grad_output: Tensor, self: Tensor, eps: float | None = None
 ) -> Tensor:
     if eps is not None:
         lo = eps
@@ -1172,7 +1178,7 @@ def logit_backward(
 @register_decomposition(aten.dropout)
 @aten.dropout.default.py_impl(DispatchKey.CompositeImplicitAutograd)
 @aten.dropout.default.py_impl(DispatchKey.Autograd)
-def dropout(input: Tensor, p: float, train: Optional[bool]):
+def dropout(input: Tensor, p: float, train: bool | None):
     if train and p != 0:
         return aten.native_dropout(input, p, train)[0]
     else:
@@ -1181,7 +1187,7 @@ def dropout(input: Tensor, p: float, train: Optional[bool]):
 
 @register_decomposition(aten.native_dropout)
 @out_wrapper("out0", "out1")
-def native_dropout(input: Tensor, p: float, train: Optional[bool]):
+def native_dropout(input: Tensor, p: float, train: bool | None):
     if train and p != 0:
         if p == 1:
             return (torch.zeros_like(input), torch.zeros_like(input, dtype=torch.bool))
@@ -1393,7 +1399,7 @@ def _chunk_cat(
     tensors: list[Tensor],
     dim: int,
     num_chunks: int,
-    out: Optional[Tensor] = None,
+    out: Tensor | None = None,
 ) -> Tensor:
     dim = _preprocess_chunk_cat_inputs(tensors, dim, num_chunks)
     padded_tensors = _pad_chunk(tensors, dim, num_chunks)
@@ -1412,8 +1418,8 @@ def split_with_sizes_copy(
     self: Tensor,
     split_sizes: list[int],
     dim: int = 0,
-    out: Optional[list[Tensor]] = None,
-) -> Optional[list[Tensor]]:
+    out: list[Tensor] | None = None,
+) -> list[Tensor] | None:
     splits = aten.split_with_sizes(self, split_sizes, dim=dim)
     if out is None:
         return [s.clone(memory_format=torch.contiguous_format) for s in splits]
@@ -1573,13 +1579,13 @@ def native_group_norm_backward(
     input: Tensor,
     mean: Tensor,
     rstd: Tensor,
-    gamma: Optional[Tensor],
+    gamma: Tensor | None,
     N: int,
     C: int,
     HxW: int,
     group: int,
     output_mask: list[bool],
-) -> tuple[Optional[Tensor], Optional[Tensor], Optional[Tensor]]:
+) -> tuple[Tensor | None, Tensor | None, Tensor | None]:
     utils.check_same_device(
         grad_output, input, mean, rstd, allow_cpu_scalar_tensors=False
     )
@@ -1608,9 +1614,9 @@ def native_group_norm_backward(
     ds = torch.mul(grad_output, input).view(N, C, HxW).sum(dim=[2])
     db = grad_output.view(N, C, HxW).sum(dim=[2])
 
-    d_input: Optional[Tensor] = None
-    d_gamma: Optional[Tensor] = None
-    d_bias: Optional[Tensor] = None
+    d_input: Tensor | None = None
+    d_gamma: Tensor | None = None
+    d_bias: Tensor | None = None
     if output_mask[0]:
         s = 1.0 / (HxW * cpg)
         if gamma is not None:
@@ -1661,7 +1667,7 @@ def native_group_norm_backward_out(
     input: Tensor,
     mean: Tensor,
     rstd: Tensor,
-    gamma: Optional[Tensor],
+    gamma: Tensor | None,
     N: int,
     C: int,
     HxW: int,
@@ -1671,7 +1677,7 @@ def native_group_norm_backward_out(
     out0: torch.Tensor,
     out1: torch.Tensor,
     out2: torch.Tensor,
-) -> tuple[Optional[Tensor], Optional[Tensor], Optional[Tensor]]:
+) -> tuple[Tensor | None, Tensor | None, Tensor | None]:
     result = native_group_norm_backward(
         grad_output, input, mean, rstd, gamma, N, C, HxW, group, output_mask
     )
@@ -1684,7 +1690,7 @@ def native_group_norm_backward_out(
     return grad_input
 
 
-def _maybe_cast(x: Optional[Tensor], dtype) -> Optional[Tensor]:
+def _maybe_cast(x: Tensor | None, dtype) -> Tensor | None:
     if x is not None:
         return x.to(dtype)
     return x
@@ -1698,10 +1704,10 @@ def native_layer_norm_backward(
     normalized_shape: list[int],
     mean: Tensor,
     rstd: Tensor,
-    weight: Optional[Tensor],
-    bias: Optional[Tensor],
+    weight: Tensor | None,
+    bias: Tensor | None,
     output_mask: list[bool],
-) -> tuple[Optional[Tensor], Optional[Tensor], Optional[Tensor]]:
+) -> tuple[Tensor | None, Tensor | None, Tensor | None]:
     input_shape = input.shape
     input_ndim = input.dim()
     computation_dtype = utils.get_computation_dtype(input.dtype)
@@ -1751,9 +1757,9 @@ def native_layer_norm_backward(
     c3 = torch.mul(x_hat, c2)
 
     inner = a - b - c3
-    d_input: Optional[Tensor] = None
-    d_weight: Optional[Tensor] = None
-    d_bias: Optional[Tensor] = None
+    d_input: Tensor | None = None
+    d_weight: Tensor | None = None
+    d_bias: Tensor | None = None
     if output_mask[0]:
         d_input = (rstd / N) * inner
 
@@ -1784,14 +1790,14 @@ def native_layer_norm_backward_out(
     normalized_shape: list[int],
     mean: Tensor,
     rstd: Tensor,
-    weight: Optional[Tensor],
-    bias: Optional[Tensor],
+    weight: Tensor | None,
+    bias: Tensor | None,
     output_mask: list[bool],
     *,
     out0: torch.Tensor,
     out1: torch.Tensor,
     out2: torch.Tensor,
-) -> tuple[Optional[Tensor], Optional[Tensor], Optional[Tensor]]:
+) -> tuple[Tensor | None, Tensor | None, Tensor | None]:
     result = native_layer_norm_backward(
         grad_out, input, normalized_shape, mean, rstd, weight, bias, output_mask
     )
@@ -1808,8 +1814,8 @@ def native_layer_norm_backward_out(
 def _fused_rms_norm(
     input: Tensor,
     normalized_shape: list[int],
-    weight: Optional[Tensor],
-    eps: Optional[float],
+    weight: Tensor | None,
+    eps: float | None,
 ) -> tuple[Tensor, Tensor]:
     dims_to_reduce: list[int] = []
     for i in range(len(normalized_shape)):
@@ -1865,9 +1871,9 @@ def _fused_rms_norm_backward(
     input: Tensor,
     normalized_shape: list[int],
     rstd: Tensor,
-    weight: Optional[Tensor],
+    weight: Tensor | None,
     output_mask: list[bool],
-) -> tuple[Optional[Tensor], Optional[Tensor]]:
+) -> tuple[Tensor | None, Tensor | None]:
     input_shape = input.shape
     input_ndim = input.dim()
     computation_dtype = utils.get_computation_dtype(input.dtype)
@@ -1911,8 +1917,8 @@ def _fused_rms_norm_backward(
     else:
         grad_x_hat = grad_out_cast
 
-    d_input: Optional[Tensor] = None
-    d_weight: Optional[Tensor] = None
+    d_input: Tensor | None = None
+    d_weight: Tensor | None = None
 
     x_hat = input_cast * rstd
 
@@ -1937,15 +1943,15 @@ def _fused_rms_norm_backward(
 
 def native_batch_norm_helper(
     input: Tensor,
-    weight: Optional[Tensor],
-    bias: Optional[Tensor],
-    running_mean: Optional[Tensor],
-    running_var: Optional[Tensor],
+    weight: Tensor | None,
+    bias: Tensor | None,
+    running_mean: Tensor | None,
+    running_var: Tensor | None,
     training: bool,
     momentum: float,
     eps: float,
     functional: bool,
-) -> tuple[Tensor, Tensor, Tensor, Optional[Tensor], Optional[Tensor]]:
+) -> tuple[Tensor, Tensor, Tensor, Tensor | None, Tensor | None]:
     reduction_dims = [0] + list(range(2, input.dim()))
     computation_dtype = utils.get_computation_dtype(input.dtype)
     new_running_mean = running_mean
@@ -2024,10 +2030,10 @@ def native_batch_norm_helper(
 @out_wrapper("out", "save_mean", "save_invstd")
 def native_batch_norm(
     input: Tensor,
-    weight: Optional[Tensor],
-    bias: Optional[Tensor],
-    running_mean: Optional[Tensor],
-    running_var: Optional[Tensor],
+    weight: Tensor | None,
+    bias: Tensor | None,
+    running_mean: Tensor | None,
+    running_var: Tensor | None,
     training: bool,
     momentum: float,
     eps: float,
@@ -2052,10 +2058,10 @@ def native_batch_norm(
 @aten.native_batch_norm.default.py_impl(DispatchKey.CompositeImplicitAutograd)
 def native_batch_norm_decomposition(
     input: Tensor,
-    weight: Optional[Tensor],
-    bias: Optional[Tensor],
-    running_mean: Optional[Tensor],
-    running_var: Optional[Tensor],
+    weight: Tensor | None,
+    bias: Tensor | None,
+    running_mean: Tensor | None,
+    running_var: Tensor | None,
     training: bool,
     momentum: float,
     eps: float,
@@ -2100,8 +2106,8 @@ def unsafe_chunk_py_impl(tensor, chunks, dim=0) -> list[Tensor]:
 @register_decomposition(aten._native_batch_norm_legit_no_training.default)
 def _native_batch_norm_legit_no_training(
     input: Tensor,
-    weight: Optional[Tensor],
-    bias: Optional[Tensor],
+    weight: Tensor | None,
+    bias: Tensor | None,
     running_mean: Tensor,
     running_var: Tensor,
     momentum: float,
@@ -2122,8 +2128,8 @@ def _native_batch_norm_legit_no_training(
 @register_decomposition(aten._native_batch_norm_legit.default)
 def _native_batch_norm_legit(
     input: Tensor,
-    weight: Optional[Tensor],
-    bias: Optional[Tensor],
+    weight: Tensor | None,
+    bias: Tensor | None,
     running_mean: Tensor,
     running_var: Tensor,
     training: bool,
@@ -2139,8 +2145,8 @@ def _native_batch_norm_legit(
 @register_decomposition(aten._native_batch_norm_legit.no_stats)
 def _native_batch_norm_legit_no_stats(
     input: Tensor,
-    weight: Optional[Tensor],
-    bias: Optional[Tensor],
+    weight: Tensor | None,
+    bias: Tensor | None,
     training: bool,
     momentum: float,
     eps: float,
@@ -2154,8 +2160,8 @@ def _native_batch_norm_legit_no_stats(
 @register_decomposition(aten._native_batch_norm_legit_functional.default)
 def _native_batch_norm_legit_functional(
     input: Tensor,
-    weight: Optional[Tensor],
-    bias: Optional[Tensor],
+    weight: Tensor | None,
+    bias: Tensor | None,
     running_mean: Tensor,
     running_var: Tensor,
     training: bool,
@@ -2180,8 +2186,8 @@ def _native_batch_norm_legit_functional(
 
 def _get_batch_norm_reserve_tensor(
     input: Tensor,
-    weight: Optional[Tensor],
-    bias: Optional[Tensor],
+    weight: Tensor | None,
+    bias: Tensor | None,
     running_mean: Tensor,
     running_var: Tensor,
     eps: float,
@@ -2210,8 +2216,8 @@ def _get_batch_norm_reserve_tensor(
 @register_decomposition(aten._batch_norm_with_update.default)
 def _batch_norm_with_update(
     input: Tensor,
-    weight: Optional[Tensor],
-    bias: Optional[Tensor],
+    weight: Tensor | None,
+    bias: Tensor | None,
     running_mean: Tensor,
     running_var: Tensor,
     momentum: float,
@@ -2237,8 +2243,8 @@ def _batch_norm_with_update(
 @register_decomposition(aten._batch_norm_with_update_functional.default)
 def _batch_norm_with_update_functional(
     input: Tensor,
-    weight: Optional[Tensor],
-    bias: Optional[Tensor],
+    weight: Tensor | None,
+    bias: Tensor | None,
     running_mean: Tensor,
     running_var: Tensor,
     momentum: float,
@@ -2266,8 +2272,8 @@ def _batch_norm_with_update_functional(
 @register_decomposition(aten._batch_norm_no_update.default)
 def _batch_norm_no_update(
     input: Tensor,
-    weight: Optional[Tensor],
-    bias: Optional[Tensor],
+    weight: Tensor | None,
+    bias: Tensor | None,
     running_mean: Tensor,
     running_var: Tensor,
     momentum: float,
@@ -2306,14 +2312,14 @@ def _fused_dropout_decomposition(input, p, generator=None):
 @register_decomposition(aten._to_copy)
 @out_wrapper()
 def _to_copy(
-    x: Union[Tensor, NumberType],
+    x: Tensor | NumberType,
     *,
-    dtype: Optional[torch.dtype] = None,
+    dtype: torch.dtype | None = None,
     layout=None,
-    device: Optional[torch.device] = None,
+    device: torch.device | None = None,
     pin_memory: bool = False,
     non_blocking: bool = False,
-    memory_format: Optional[torch.memory_format] = None,
+    memory_format: torch.memory_format | None = None,
 ):
     if layout and layout != torch.strided:
         raise AssertionError(f"layout must be None or torch.strided, got {layout}")
@@ -2368,9 +2374,9 @@ def nop_decomposition(x):
 def cudnn_batch_norm(
     input: Tensor,
     weight: Tensor,
-    bias: Optional[Tensor],
-    running_mean: Optional[Tensor],
-    running_var: Optional[Tensor],
+    bias: Tensor | None,
+    running_mean: Tensor | None,
+    running_var: Tensor | None,
     training: bool,
     exponential_average_factor: float,
     epsilon: float,
@@ -2407,16 +2413,16 @@ def _broadcast_batch_norm_backward(x, broadcast_mask):
 def batch_norm_backward(
     grad_out: Tensor,
     input: Tensor,
-    weight: Optional[Tensor],
-    running_mean: Optional[Tensor],
-    running_var: Optional[Tensor],
-    save_mean: Optional[Tensor],
-    save_invstd: Optional[Tensor],
+    weight: Tensor | None,
+    running_mean: Tensor | None,
+    running_var: Tensor | None,
+    save_mean: Tensor | None,
+    save_invstd: Tensor | None,
     train: bool,
     eps: float,
     output_mask: list[bool],
     reserve: Tensor,
-) -> tuple[Tensor, Optional[Tensor], Optional[Tensor]]:
+) -> tuple[Tensor, Tensor | None, Tensor | None]:
     return native_batch_norm_backward(
         grad_out,
         input,
@@ -2435,15 +2441,15 @@ def batch_norm_backward(
 def native_batch_norm_backward(
     grad_out: Tensor,
     input: Tensor,
-    weight: Optional[Tensor],
-    running_mean: Optional[Tensor],
-    running_var: Optional[Tensor],
-    save_mean: Optional[Tensor],
-    save_invstd: Optional[Tensor],
+    weight: Tensor | None,
+    running_mean: Tensor | None,
+    running_var: Tensor | None,
+    save_mean: Tensor | None,
+    save_invstd: Tensor | None,
     train: bool,
     eps: float,
     output_mask: list[bool],
-) -> tuple[Tensor, Optional[Tensor], Optional[Tensor]]:
+) -> tuple[Tensor, Tensor | None, Tensor | None]:
     input_dtype = input.dtype
     if weight is not None:
         weight_dtype = weight.dtype
@@ -2545,11 +2551,11 @@ def native_batch_norm_backward(
 def native_batch_norm_backward_out(
     grad_out: Tensor,
     input: Tensor,
-    weight: Optional[Tensor],
-    running_mean: Optional[Tensor],
-    running_var: Optional[Tensor],
-    save_mean: Optional[Tensor],
-    save_invstd: Optional[Tensor],
+    weight: Tensor | None,
+    running_mean: Tensor | None,
+    running_var: Tensor | None,
+    save_mean: Tensor | None,
+    save_invstd: Tensor | None,
     train: bool,
     eps: float,
     output_mask: list[bool],
@@ -2557,7 +2563,7 @@ def native_batch_norm_backward_out(
     out0: torch.Tensor,
     out1: torch.Tensor,
     out2: torch.Tensor,
-) -> tuple[Tensor, Optional[Tensor], Optional[Tensor]]:
+) -> tuple[Tensor, Tensor | None, Tensor | None]:
     result = native_batch_norm_backward(
         grad_out,
         input,
@@ -2585,10 +2591,10 @@ def miopen_batch_norm_backward(
     input: Tensor,
     grad_output: Tensor,
     weight: Tensor,
-    running_mean: Optional[Tensor],
-    running_var: Optional[Tensor],
-    save_mean: Optional[Tensor],
-    save_var: Optional[Tensor],
+    running_mean: Tensor | None,
+    running_var: Tensor | None,
+    save_mean: Tensor | None,
+    save_var: Tensor | None,
     epsilon: float,
 ):
     return aten.native_batch_norm_backward(
@@ -2611,10 +2617,10 @@ def cudnn_batch_norm_backward(
     input: Tensor,
     grad_output: Tensor,
     weight: Tensor,
-    running_mean: Optional[Tensor],
-    running_var: Optional[Tensor],
-    save_mean: Optional[Tensor],
-    save_var: Optional[Tensor],
+    running_mean: Tensor | None,
+    running_var: Tensor | None,
+    save_mean: Tensor | None,
+    save_var: Tensor | None,
     epsilon: float,
     reserveSpace: Tensor,
 ):
@@ -3011,9 +3017,9 @@ def log_sigmoid_forward(self: Tensor) -> tuple[Tensor, Tensor]:
 @out_wrapper()
 def uniform(
     x: Tensor,
-    low: Union[bool, int, float] = 0.0,
-    high: Union[bool, int, float] = 1.0,
-    generator: Optional[torch.Generator] = None,
+    low: bool | int | float = 0.0,
+    high: bool | int | float = 1.0,
+    generator: torch.Generator | None = None,
 ):
     return prims._uniform_helper(
         x.shape,
@@ -3077,8 +3083,8 @@ def get_scale_value(scales, idx):
 @aten.upsample_nearest3d.vec.py_impl(DispatchKey.Autograd)
 def _upsample_nearest_vec(
     input: Tensor,
-    output_size: Optional[list[int]],
-    scale_factors: Optional[list[float]],
+    output_size: list[int] | None,
+    scale_factors: list[float] | None,
 ) -> Tensor:
     osize = upsample_compute_output_size(input.size(), output_size, scale_factors)
     scales = (
@@ -3098,8 +3104,8 @@ def _upsample_nearest_vec(
 @aten._upsample_nearest_exact3d.vec.py_impl(DispatchKey.Autograd)
 def _upsample_nearest_exact_vec(
     input: Tensor,
-    output_size: Optional[list[int]],
-    scale_factors: Optional[list[float]],
+    output_size: list[int] | None,
+    scale_factors: list[float] | None,
 ) -> Tensor:
     osize = upsample_compute_output_size(input.size(), output_size, scale_factors)
     scales = (
@@ -3153,7 +3159,7 @@ def _compute_upsample_nearest_indices(input, output_size, scales, exact=False):
 def upsample_nearest1d(
     input: Tensor,
     output_size: list[int],
-    scales: Optional[float] = None,
+    scales: float | None = None,
 ) -> Tensor:
     return _upsample_nearest(input, output_size, [scales])
 
@@ -3167,7 +3173,7 @@ def upsample_nearest1d(
 def upsample_nearest_exact1d(
     input: Tensor,
     output_size: list[int],
-    scales: Optional[float] = None,
+    scales: float | None = None,
 ) -> Tensor:
     return _upsample_nearest(input, output_size, [scales], exact=True)
 
@@ -3179,8 +3185,8 @@ def upsample_nearest_exact1d(
 def upsample_nearest2d(
     input: Tensor,
     output_size: list[int],
-    scales_h: Optional[float] = None,
-    scales_w: Optional[float] = None,
+    scales_h: float | None = None,
+    scales_w: float | None = None,
 ) -> Tensor:
     return _upsample_nearest(input, output_size, [scales_h, scales_w])
 
@@ -3194,8 +3200,8 @@ def upsample_nearest2d(
 def _upsample_nearest_exact2d(
     input: Tensor,
     output_size: list[int],
-    scales_h: Optional[float] = None,
-    scales_w: Optional[float] = None,
+    scales_h: float | None = None,
+    scales_w: float | None = None,
 ) -> Tensor:
     return _upsample_nearest(input, output_size, [scales_h, scales_w], exact=True)
 
@@ -3207,9 +3213,9 @@ def _upsample_nearest_exact2d(
 def upsample_nearest3d(
     input: Tensor,
     output_size: list[int],
-    scales_d: Optional[float] = None,
-    scales_h: Optional[float] = None,
-    scales_w: Optional[float] = None,
+    scales_d: float | None = None,
+    scales_h: float | None = None,
+    scales_w: float | None = None,
 ) -> Tensor:
     return _upsample_nearest(input, output_size, [scales_d, scales_h, scales_w])
 
@@ -3223,9 +3229,9 @@ def upsample_nearest3d(
 def _upsample_nearest_exact3d(
     input: Tensor,
     output_size: list[int],
-    scales_d: Optional[float] = None,
-    scales_h: Optional[float] = None,
-    scales_w: Optional[float] = None,
+    scales_d: float | None = None,
+    scales_h: float | None = None,
+    scales_w: float | None = None,
 ) -> Tensor:
     return _upsample_nearest(
         input, output_size, [scales_d, scales_h, scales_w], exact=True
@@ -3236,7 +3242,7 @@ def _upsample_nearest_exact3d(
 def _upsample_nearest(
     input: Tensor,
     output_size: list[int],
-    scales: list[Optional[float]],
+    scales: list[float | None],
     exact: bool = False,
 ) -> Tensor:
     spatial_indices = _compute_upsample_nearest_indices(
@@ -3947,6 +3953,18 @@ def upsample_bicubic2d_aa_vec(input, output_size, align_corners, scale_factors):
     )
 
 
+@register_decomposition(aten._upsample_lanczos2d_aa.vec)
+@aten._upsample_lanczos2d_aa.vec.py_impl(DispatchKey.CompositeImplicitAutograd)
+@aten._upsample_lanczos2d_aa.vec.py_impl(DispatchKey.Autograd)
+def upsample_lanczos2d_aa_vec(input, output_size, align_corners, scale_factors):
+    osize = upsample_compute_output_size(input.size(), output_size, scale_factors)
+    scale_h = get_scale_value(scale_factors, 0)
+    scale_w = get_scale_value(scale_factors, 1)
+    return torch.ops.aten._upsample_lanczos2d_aa(
+        input, osize, align_corners, scale_h, scale_w
+    )
+
+
 @register_decomposition(aten.upsample_bilinear2d.vec)
 @register_decomposition(aten.upsample_trilinear3d.vec)
 @aten.upsample_linear1d.vec.py_impl(DispatchKey.CompositeImplicitAutograd)
@@ -3967,7 +3985,7 @@ def upsample_linear1d(
     input: Tensor,
     output_size: list[int],
     align_corners: bool,
-    scales_w: Optional[float] = None,
+    scales_w: float | None = None,
 ) -> Tensor:
     return _upsample_linear(input, output_size, align_corners, [scales_w])
 
@@ -3981,8 +3999,8 @@ def upsample_bilinear2d(
     input: Tensor,
     output_size: list[int],
     align_corners: bool,
-    scales_h: Optional[float] = None,
-    scales_w: Optional[float] = None,
+    scales_h: float | None = None,
+    scales_w: float | None = None,
 ) -> Tensor:
     return _upsample_linear(input, output_size, align_corners, [scales_h, scales_w])
 
@@ -3995,9 +4013,9 @@ def upsample_trilinear3d(
     input: Tensor,
     output_size: list[int],
     align_corners: bool,
-    scales_d: Optional[float] = None,
-    scales_h: Optional[float] = None,
-    scales_w: Optional[float] = None,
+    scales_d: float | None = None,
+    scales_h: float | None = None,
+    scales_w: float | None = None,
 ) -> Tensor:
     return _upsample_linear(
         input, output_size, align_corners, [scales_d, scales_h, scales_w]
@@ -4042,7 +4060,7 @@ def _upsample_linear(
     input: Tensor,
     output_size: list[int],
     align_corners: bool,
-    scales: list[Optional[float]],
+    scales: list[float | None],
 ) -> Tensor:
     # get dimensions of original image
     n_channels = input.shape[1]
@@ -4193,7 +4211,7 @@ def _unsafe_masked_index_put_accumulate(x, mask, indices, values):
 def _nll_loss_forward(
     self: Tensor,
     target: Tensor,
-    weight: Optional[Tensor],
+    weight: Tensor | None,
     reduction: int,
     ignore_index: int,
 ) -> tuple[Tensor, Tensor]:
@@ -4249,7 +4267,7 @@ def _nll_loss_forward(
 def nll_loss_forward(
     self: Tensor,
     target: Tensor,
-    weight: Optional[Tensor],
+    weight: Tensor | None,
     reduction: int,
     ignore_index: int,
 ) -> tuple[Tensor, Tensor]:
@@ -4261,9 +4279,10 @@ def nll_loss_forward(
         )
 
     no_batch_dim = self.dim() == 1 and target.dim() == 0
-    if not (no_batch_dim or (self.shape[0] == target.shape[0])):
-        raise AssertionError(
-            f"size mismatch (got input: {self.shape}, target: {target.shape})"
+    if not no_batch_dim:
+        torch._check(
+            self.shape[0] == target.shape[0],
+            lambda: f"size mismatch (got input: {self.shape}, target: {target.shape})",
         )
 
     n_classes = self.shape[-1]
@@ -4282,7 +4301,7 @@ def nll_loss_forward(
 def nll_loss2d_forward(
     self: Tensor,
     target: Tensor,
-    weight: Optional[Tensor],
+    weight: Tensor | None,
     reduction: int,
     ignore_index: int,
 ) -> tuple[Tensor, Tensor]:
@@ -4624,6 +4643,9 @@ def binary_cross_entropy_with_logits(
     if weight is not None:
         loss = loss * weight
 
+    # this is to align the resulted data type with the in-place
+    # operation in binary_cross_entropy_with_logits of aten/src/ATen/native/Loss.cpp
+    loss = loss.to(target.dtype)
     return apply_loss_reduction(loss, reduction)
 
 
@@ -4794,8 +4816,8 @@ def upsample_bicubic2d_default(
     input: Tensor,
     output_size: tuple[int, int],
     align_corners: bool,
-    scale_h: Optional[float] = None,
-    scale_w: Optional[float] = None,
+    scale_h: float | None = None,
+    scale_w: float | None = None,
 ) -> Tensor:
     # get dimensions of original image
     _, _, in_h, in_w = input.shape
@@ -4884,9 +4906,9 @@ def get_x_interp(y):
 @pw_cast_for_opmath
 def upsample_bicubic2d_vec(
     a: Tensor,
-    output_size: Optional[tuple[int, int]],
+    output_size: tuple[int, int] | None,
     align_corners: bool,
-    scale_factors: Optional[tuple[float, float]] = None,
+    scale_factors: tuple[float, float] | None = None,
 ) -> Tensor:
     torch._check(
         bool(output_size) + bool(scale_factors) == 1,
@@ -5082,9 +5104,9 @@ def nansum(self, dim=None, keepdim=False, *, dtype=None):
 def arange_default(
     end: NumberType,
     *,
-    dtype: Optional[torch.dtype] = None,
+    dtype: torch.dtype | None = None,
     layout: torch.layout = torch.strided,
-    device: Optional[torch.device] = None,
+    device: torch.device | None = None,
     pin_memory: bool = False,
 ):
     return aten.arange.start_step(
@@ -5097,9 +5119,9 @@ def arange_start(
     start: NumberType,
     end: NumberType,
     *,
-    dtype: Optional[torch.dtype] = None,
+    dtype: torch.dtype | None = None,
     layout: torch.layout = torch.strided,
-    device: Optional[torch.device] = None,
+    device: torch.device | None = None,
     pin_memory: bool = False,
 ):
     return aten.arange.start_step(
@@ -5122,7 +5144,7 @@ def multi_margin_loss(
     target: Tensor,
     p: NumberType = 1,
     margin: NumberType = 1,
-    weight: Optional[Tensor] = None,
+    weight: Tensor | None = None,
     reduction: int = Reduction.MEAN.value,
 ) -> Tensor:
     input = torch.atleast_2d(input)
@@ -5232,8 +5254,8 @@ def scaled_dot_product_flash_attention_for_cpu(
     dropout_p: float = 0.0,
     is_causal: bool = False,
     *,
-    attn_mask: Optional[Tensor] = None,
-    scale: Optional[float] = None,
+    attn_mask: Tensor | None = None,
+    scale: float | None = None,
 ) -> tuple[Tensor, Tensor]:
     torch._check(
         torch.is_floating_point(query),
@@ -5334,8 +5356,8 @@ def sym_numel(t):
 def sum_default(
     self: Tensor,
     *,
-    dtype: Optional[torch.dtype] = None,
-    out: Optional[Tensor] = None,
+    dtype: torch.dtype | None = None,
+    out: Tensor | None = None,
 ) -> Tensor:
     if out is None:
         return aten.sum.dim_IntList(self, [], dtype=dtype)
@@ -5344,7 +5366,7 @@ def sum_default(
 
 
 @register_decomposition([aten.squeeze.default, aten.squeeze.dim])
-def squeeze_default(self: Tensor, dim: Optional[int] = None):
+def squeeze_default(self: Tensor, dim: int | None = None):
     # handle a scalar directly
     if not isinstance(self, torch.Tensor):
         return self
@@ -5377,7 +5399,9 @@ def isin(elements, test_elements, *, assume_unique=False, invert=False):
         else:
             return torch.eq(elements, test_elements)
 
-    if test_elements.numel() < 10.0 * pow(elements.numel(), 0.145):
+    from torch.fx.experimental.symbolic_shapes import guard_or_false
+
+    if guard_or_false(test_elements.numel() < 10.0 * pow(elements.numel(), 0.145)):
         return isin_default(elements, test_elements, invert=invert)
     else:
         return isin_sorting(
@@ -5389,7 +5413,7 @@ def isin(elements, test_elements, *, assume_unique=False, invert=False):
 def bernoulli(
     self: torch.Tensor,
     *,
-    generator: Optional[torch.Generator] = None,
+    generator: torch.Generator | None = None,
 ) -> torch.Tensor:
     if generator is None:
         raw_p = torch.rand(self.size(), dtype=torch.float32, device=self.device)
@@ -5462,6 +5486,219 @@ def resize_as(self, other, memory_format=None):
     return aten.resize(self, other.shape, memory_format=memory_format)
 
 
+@register_decomposition(aten.max_pool2d_with_indices_backward)
+def max_pool2d_with_indices_backward(
+    grad_output: Tensor,
+    self: Tensor,
+    kernel_size,
+    stride,
+    padding,
+    dilation,
+    ceil_mode: bool,
+    indices: Tensor,
+):
+    """
+    Decomposition of max_pool2d_with_indices_backward using scatter_add.
+
+    This replaces the native implementation with a high-level decomposition
+    that uses scatter_add for gradient accumulation. The scatter-based approach
+    provides automatic optimization opportunities for Inductor and handles all
+    pooling configurations without requiring specialized fallback paths.
+
+    Algorithm:
+        For each output gradient position, use the corresponding index from the
+        forward pass to scatter the gradient to the input position. When multiple
+        output positions select the same input position as max, scatter_add
+        automatically accumulates their gradients.
+
+    Complexity: O(B * C * H_out * W_out)
+        Independent of kernel size, unlike traditional O(B * C * H_in * W_in * K²)
+        approaches that iterate over input positions and kernel windows.
+
+    Known Limitations:
+        - FP16/BF16: Uses FP32 accumulation internally to preserve precision when
+          many gradients accumulate to the same position (overlapping pooling windows).
+          This adds slight overhead but ensures numerical stability.
+        - Deterministic mode: Falls back to native implementation to ensure
+          consistent results across runs
+
+    Args:
+        grad_output: Gradient w.r.t. pooling output [B, C, H_out, W_out]
+        self: Original input tensor (for shape) [B, C, H_in, W_in]
+        kernel_size: Pooling kernel size
+        stride: Pooling stride
+        padding: Pooling padding
+        dilation: Pooling dilation
+        ceil_mode: Whether to use ceil for output size calculation
+        indices: Indices from forward pass (per-channel linear positions)
+
+    Returns:
+        Gradient w.r.t. input [B, C, H_in, W_in]
+    """
+    # Use native kernel in deterministic mode
+    if torch.are_deterministic_algorithms_enabled():
+        return NotImplemented
+
+    # MPS: Use native kernel. scatter_add has correctness issues on macOS 14
+    # (#163327) and numerical differences on macOS 15+.
+    if grad_output.device.type == "mps":
+        return NotImplemented
+
+    # Get spatial dimensions
+    in_height = self.size(-2)
+    in_width = self.size(-1)
+    out_height = grad_output.size(-2)
+    out_width = grad_output.size(-1)
+
+    # Handle both 3D (C, H, W) and 4D (B, C, H, W) cases by treating 3D as 4D
+    is_batched = self.dim() == 4
+    if not is_batched:
+        self = self.unsqueeze(0)
+        grad_output = grad_output.unsqueeze(0)
+        indices = indices.unsqueeze(0)
+
+    batch_size = self.size(0)
+    channels = self.size(1)
+
+    # For FP16/BF16, use FP32 accumulation to avoid precision loss
+    # This is critical when many gradients accumulate to the same position
+    # (overlapping pooling windows with large kernels or stride < kernel_size)
+    use_fp32_accum = grad_output.dtype in (torch.float16, torch.bfloat16)
+    accum_dtype = torch.float32 if use_fp32_accum else grad_output.dtype
+
+    # Create grad_input with correct accumulation dtype from the start
+    grad_input_flat = torch.zeros(
+        batch_size * channels,
+        in_height * in_width,
+        dtype=accum_dtype,
+        device=grad_output.device,
+    )
+
+    # Reshape grad_output and indices to (B*C, H_out*W_out)
+    grad_output_flat = grad_output.reshape(
+        batch_size * channels, out_height * out_width
+    )
+    indices_flat = indices.reshape(batch_size * channels, out_height * out_width)
+
+    # Convert grad_output to accumulation dtype if needed
+    if use_fp32_accum:
+        grad_output_flat = grad_output_flat.to(torch.float32)
+
+    # Scatter gradients to input positions
+    grad_input_flat = grad_input_flat.scatter_add(1, indices_flat, grad_output_flat)
+
+    # Reshape back to original input shape
+    grad_input = grad_input_flat.reshape(batch_size, channels, in_height, in_width)
+
+    # Convert back to original dtype if we used FP32 accumulation
+    if use_fp32_accum:
+        grad_input = grad_input.to(grad_output.dtype)
+
+    # Preserve memory format from input (channels_last vs channels_first)
+    memory_format = utils.suggest_memory_format(self)
+    grad_input = grad_input.contiguous(memory_format=memory_format)
+
+    # Remove batch dimension for 3D case
+    if not is_batched:
+        grad_input = grad_input.squeeze(0)
+
+    return grad_input
+
+
+@register_decomposition([aten.hann_window.default, aten.hann_window.out])
+@out_wrapper()
+def hann_window(
+    window_length: int,
+    *,
+    dtype: torch.dtype | None = None,
+    layout: torch.layout | None = None,
+    device: torch.device | None = None,
+    pin_memory: bool | None = None,
+) -> Tensor:
+    """hann_window(window_length, *, dtype=None, layout=None, device=None, pin_memory=False) -> Tensor
+
+    Returns a Hann window of size :attr:`window_length` with ``periodic=True``.
+
+    Equivalent to :func:`torch.hann_window` with ``periodic=True``.
+
+    Args:
+        window_length (int): the size of returned window.
+
+    Keyword args:
+        dtype (:class:`torch.dtype`, optional): desired dtype. Default: global default.
+        layout (:class:`torch.layout`, optional): desired layout. Default: ``torch.strided``.
+        device (:class:`torch.device`, optional): desired device. Default: current device.
+        pin_memory (bool, optional): if ``True``, pins the returned tensor. Default: ``False``.
+    """
+    return aten.hann_window.periodic(
+        window_length,
+        True,
+        dtype=dtype,
+        layout=layout,
+        device=device,
+        pin_memory=pin_memory,
+    )
+
+
+@register_decomposition([aten.hann_window.periodic, aten.hann_window.periodic_out])
+@out_wrapper()
+def hann_window_periodic(
+    window_length: int,
+    periodic: bool = True,
+    *,
+    dtype: torch.dtype | None = None,
+    layout: torch.layout | None = None,
+    device: torch.device | None = None,
+    pin_memory: bool | None = None,
+) -> Tensor:
+    r"""hann_window(window_length, periodic=True, *, dtype=None, layout=None, device=None, pin_memory=False) -> Tensor
+
+    Returns a Hann window of size :attr:`window_length`.
+
+    .. math::
+        w[n] = 0.5 - 0.5 \cos\!\left(\frac{2\pi n}{N-1}\right)
+
+    where :math:`N` is ``window_length + 1`` when ``periodic=True`` (for spectral analysis),
+    or ``window_length`` when ``periodic=False`` (symmetric window).
+
+    Low-precision dtypes (``bfloat16``, ``float16``) are computed in ``float32`` then cast.
+
+    Args:
+        window_length (int): the size of returned window.
+        periodic (bool, optional): if ``True``, returns a periodic window for use with STFT.
+            Default: ``True``.
+
+    Keyword args:
+        dtype (:class:`torch.dtype`, optional): desired dtype. Default: global default.
+        layout (:class:`torch.layout`, optional): desired layout. Default: ``torch.strided``.
+        device (:class:`torch.device`, optional): desired device. Default: current device.
+        pin_memory (bool, optional): if ``True``, pins the returned tensor. Default: ``False``.
+    """
+    dtype = dtype if dtype is not None else torch.get_default_dtype()
+    if window_length == 0:
+        return torch.empty(
+            (0,), dtype=dtype, layout=layout, device=device, pin_memory=pin_memory
+        )
+    if window_length == 1:
+        return torch.ones(
+            (1,), dtype=dtype, layout=layout, device=device, pin_memory=pin_memory
+        )
+    compute_dtype = utils.get_computation_dtype(dtype)
+    n = window_length + 1 if periodic else window_length
+    t = torch.arange(
+        n,
+        dtype=compute_dtype,
+        layout=layout,
+        device=device,
+        pin_memory=pin_memory,
+    )
+    t = t * (2.0 * torch.pi / (n - 1))
+    t = torch.cos(t)
+    t = t * -0.5 + 0.5
+    window = t.narrow(0, 0, window_length) if periodic else t
+    return window.to(dtype)
+
+
 register_inplace(aten.addbmm_, aten.addbmm)
 register_inplace(aten.addmm_, aten.addmm)
 register_inplace(aten.addmv_, aten.addmv)
@@ -5488,3 +5725,22 @@ def resize_as(self, other, memory_format=None):
 register_inplace(aten.scatter_add_, aten.scatter_add)
 register_inplace(aten.scatter_reduce_, aten.scatter_reduce)
 register_inplace(aten.silu_, aten.silu)
+
+
+@aten.one_hot.default.py_impl(DispatchKey.CompositeImplicitAutograd)
+def one_hot(self: Tensor, num_classes: int = -1) -> Tensor:
+    if num_classes == -1:
+        num_classes = int(self.max().item()) + 1
+    # _assert_async is side-effectful and won't be DCE'd
+    aten._assert_async.msg(
+        torch.all(self >= 0),
+        "one_hot: Class values must be non-negative.",
+    )
+    aten._assert_async.msg(
+        torch.all(self < num_classes),
+        "one_hot: Class values must be smaller than num_classes.",
+    )
+    return (
+        self.unsqueeze(-1)
+        == torch.arange(num_classes, dtype=self.dtype, device=self.device)
+    ).to(torch.int64)
diff --git a/torch/_decomp/decompositions_for_jvp.py b/torch/_decomp/decompositions_for_jvp.py
index 86e7123ccebd4..45140652306c6 100644
--- a/torch/_decomp/decompositions_for_jvp.py
+++ b/torch/_decomp/decompositions_for_jvp.py
@@ -2,7 +2,6 @@
 # mypy: allow-untyped-defs
 import inspect
 from collections.abc import Callable
-from typing import Optional
 
 import torch
 import torch._decomp
@@ -136,10 +135,10 @@ def native_layer_norm_backward(
     normalized_shape: list[int],
     mean: Tensor,
     rstd: Tensor,
-    weight: Optional[Tensor],
-    bias: Optional[Tensor],
+    weight: Tensor | None,
+    bias: Tensor | None,
     output_mask: list[bool],
-) -> tuple[Optional[Tensor], Optional[Tensor], Optional[Tensor]]:
+) -> tuple[Tensor | None, Tensor | None, Tensor | None]:
     input_shape = input.shape
     input_ndim = input.dim()
 
@@ -177,13 +176,13 @@ def native_layer_norm_backward(
     inner = a - b - c3
 
     if output_mask[0]:
-        d_input: Optional[Tensor] = (rstd_ / N) * inner
+        d_input: Tensor | None = (rstd_ / N) * inner
     else:
         d_input = torch.zeros_like(input)  # should be None but doesn't work with vjp
 
     if output_mask[1] and weight is not None:
         if len(outer_dim_indices) > 0:
-            d_weight: Optional[Tensor] = torch.sum(
+            d_weight: Tensor | None = torch.sum(
                 grad_out * x_hat, outer_dim_indices, False
             )
         else:
@@ -195,7 +194,7 @@ def native_layer_norm_backward(
 
     if output_mask[2] and bias is not None:
         if len(outer_dim_indices) > 0:
-            d_bias: Optional[Tensor] = torch.sum(grad_out, outer_dim_indices, False)
+            d_bias: Tensor | None = torch.sum(grad_out, outer_dim_indices, False)
         else:
             d_bias = grad_out.clone()
     elif bias is not None:
@@ -217,15 +216,15 @@ def prod(x: list[int]):
 def native_batch_norm_backward(
     grad_out: Tensor,
     input: Tensor,
-    weight: Optional[Tensor],
-    running_mean: Optional[Tensor],
-    running_var: Optional[Tensor],
-    save_mean: Optional[Tensor],
-    save_invstd: Optional[Tensor],
+    weight: Tensor | None,
+    running_mean: Tensor | None,
+    running_var: Tensor | None,
+    save_mean: Tensor | None,
+    save_invstd: Tensor | None,
     train: bool,
     eps: float,
     output_mask: list[bool],
-) -> tuple[Tensor, Optional[Tensor], Optional[Tensor]]:
+) -> tuple[Tensor, Tensor | None, Tensor | None]:
     input_shape = input.shape
     input_rank = input.dim()
     if input_rank < 2:
@@ -309,15 +308,15 @@ def batch_norm_backward(
     grad_out: Tensor,
     input: Tensor,
     weight: Tensor,
-    running_mean: Optional[Tensor],
-    running_var: Optional[Tensor],
-    save_mean: Optional[Tensor],
-    save_var: Optional[Tensor],
+    running_mean: Tensor | None,
+    running_var: Tensor | None,
+    save_mean: Tensor | None,
+    save_var: Tensor | None,
     update: bool,
     eps: float,
     output_mask: list[bool],
     reserve: Tensor,
-) -> tuple[Tensor, Optional[Tensor], Optional[Tensor]]:
+) -> tuple[Tensor, Tensor | None, Tensor | None]:
     return native_batch_norm_backward(
         grad_out,
         input,
diff --git a/torch/_decomp/decompositions_for_rng.py b/torch/_decomp/decompositions_for_rng.py
index 27b29ca5f85aa..d1753b1baf5a8 100644
--- a/torch/_decomp/decompositions_for_rng.py
+++ b/torch/_decomp/decompositions_for_rng.py
@@ -78,7 +78,7 @@ def reset(self):
         self.seed = torch.tensor(())
         self.base_offset = torch.tensor(())
         self.relative_offset = 0
-        self.offset_advanced_alteast_once = False
+        self.offset_advanced_at_least_once = False
 
     def validate_state(self):
         if self.seed.numel() == 0 or self.base_offset.numel() == 0:
@@ -88,7 +88,7 @@ def validate_state(self):
             )
 
     def advance_offset(self, consumed_offset):
-        self.offset_advanced_alteast_once = True
+        self.offset_advanced_at_least_once = True
         self.relative_offset = self.relative_offset + consumed_offset
 
     def set_state(self, seed, base_offset, relative_offset=0):
@@ -205,7 +205,7 @@ def multiple_of_4(offset):
     @classmethod
     def get_updated_fwd_offset(cls):
         # Short circuit if no rand ops were observed
-        if not cls.fwd_state.offset_advanced_alteast_once:
+        if not cls.fwd_state.offset_advanced_at_least_once:
             return cls.fwd_state.base_offset
         return cls.multiple_of_4(
             cls.fwd_state.base_offset + cls.fwd_state.relative_offset
@@ -214,7 +214,7 @@ def get_updated_fwd_offset(cls):
     @classmethod
     def get_updated_bwd_offset(cls):
         # Short circuit if no rand ops were observed
-        if not cls.bwd_state.offset_advanced_alteast_once:
+        if not cls.bwd_state.offset_advanced_at_least_once:
             return cls.bwd_state.base_offset
         return cls.multiple_of_4(
             cls.bwd_state.base_offset + cls.bwd_state.relative_offset
diff --git a/torch/_dispatch/python.py b/torch/_dispatch/python.py
index b9bca98e7da65..ef6d1d652bf2d 100644
--- a/torch/_dispatch/python.py
+++ b/torch/_dispatch/python.py
@@ -2,7 +2,7 @@
 import unittest.mock
 from collections.abc import Callable, Generator, Iterator
 from contextlib import contextmanager
-from typing import TypeVar, Union
+from typing import TypeVar
 from typing_extensions import ParamSpec
 
 import torch
@@ -119,7 +119,7 @@ def _fmt(a: object) -> object:
 
 def make_crossref_functionalize(
     op: torch._ops.OpOverload[_P, _T], final_key: DispatchKey
-) -> Union[Callable[_P, _T], DispatchKey]:
+) -> Callable[_P, _T] | DispatchKey:
     from torch._subclasses.fake_tensor import FakeTensorMode
 
     # This case is pretty weird, suppress it for now
diff --git a/torch/_dynamo/CLAUDE.md b/torch/_dynamo/CLAUDE.md
new file mode 100644
index 0000000000000..e72321277fdb9
--- /dev/null
+++ b/torch/_dynamo/CLAUDE.md
@@ -0,0 +1,287 @@
+# torch/_dynamo
+
+TorchDynamo is a Python-level JIT compiler that captures PyTorch programs into
+FX graphs by symbolically executing Python bytecode. It hooks into CPython's
+PEP 523 frame evaluation API to intercept execution, traces operations into an
+FX graph, compiles the graph with a backend (e.g. Inductor), and generates new
+bytecode that calls the compiled code.
+
+## Architecture Overview
+
+The compilation pipeline, in execution order:
+
+1. **`eval_frame.py`** — Runtime entry point. `torch.compile()` wraps a
+   function in an `OptimizedModule`. At runtime, the C extension
+   (`torch._C._dynamo.eval_frame`) intercepts Python frames via PEP 523.
+2. **`convert_frame.py`** — `ConvertFrameAssert.__call__` checks caches,
+   handles recompilation limits, calls `_compile()` → `trace_frame()`.
+3. **`symbolic_convert.py`** — The heart of Dynamo. `InstructionTranslator`
+   symbolically executes bytecode instruction-by-instruction. Maintains a
+   symbolic `stack` (list of `VariableTracker`s) and `symbolic_locals` (dict of
+   name → `VariableTracker`). Opcodes are dispatched via a `dispatch_table`
+   built by `BytecodeDispatchTableMeta`.
+4. **`output_graph.py`** — `OutputGraph` owns the FX graph being built (via
+   `SubgraphTracer`), the `SideEffects` tracker, guards, shape environment,
+   and graph args. `compile_subgraph()` finalizes the graph, calls the backend,
+   and generates output bytecode.
+5. **`codegen.py`** — `PyCodegen` emits output bytecode: loads graph inputs
+   (via `Source.reconstruct()`), calls the compiled graph, unpacks outputs, and
+   replays side effects.
+6. **`resume_execution.py`** — Generates continuation functions for execution
+   after graph breaks.
+
+## Key Abstractions
+
+### VariableTracker (`variables/`)
+
+Every Python value encountered during tracing is wrapped in a `VariableTracker`
+subclass. Key interface: `as_python_constant()`, `as_proxy()`,
+`call_function()`, `call_method()`, `var_getattr()`, `reconstruct()`.
+
+Key fields: `source` (where the value came from, for guards) and
+`mutation_type` (whether/how mutations are tracked).
+
+**Factory**: `VariableTracker.build(tx, value, source=...)` dispatches to
+`VariableBuilder` (sourced values needing guards) or `SourcelessBuilder`
+(values created during tracing).
+
+Key subclass families in `variables/`: `TensorVariable` / `SymNodeVariable`
+(tensor.py), `ConstantVariable` (constant.py), `ListVariable` /
+`TupleVariable` (lists.py), `ConstDictVariable` (dicts.py), `SetVariable` (sets.py),
+`UserFunctionVariable` (functions.py), `BuiltinVariable` (builtin.py),
+`NNModuleVariable` (nn_module.py), `UserDefinedObjectVariable`
+(user_defined.py), `TorchHigherOrderOperatorVariable` (higher_order_ops.py),
+`LazyVariableTracker` (lazy.py). `VariableBuilder` and `SourcelessBuilder` are
+in builder.py.
+
+### Source (`source.py`)
+
+Tracks value provenance — how to access a value at runtime. Used for guard
+generation (`source.make_guard(GuardBuilder.XXX)`) and bytecode reconstruction
+(`source.reconstruct(codegen)`). Root sources: `LocalSource`, `GlobalSource`.
+Chained sources: `AttrSource`, `GetItemSource`, `NNModuleSource`, etc.
+
+### Guards (`guards.py`)
+
+Runtime conditions that must hold for cached compiled code to be reused.
+Install via `install_guard(source.make_guard(GuardBuilder.TYPE_MATCH))`.
+Common types: `TYPE_MATCH`, `ID_MATCH`, `EQUALS_MATCH`, `TENSOR_MATCH`,
+`SEQUENCE_LENGTH`. At finalization, `CheckFunctionManager` builds a tree of
+C++ `GuardManager` objects for fast runtime checking.
+
+### Side Effects (`side_effects.py`)
+
+Tracks mutations during tracing (attribute stores, list mutations, cell
+variable updates, tensor hooks) and replays them as bytecode after graph
+execution. The `MutationType` system (`variables/base.py`) controls what
+mutations are allowed: `ValueMutationNew/Existing`,
+`AttributeMutationNew/Existing`, or `None` (immutable). The `scope` field
+prevents cross-scope mutations inside higher-order operators.
+
+### Other key files
+
+- `trace_rules.py` — inline/skip/graph-break decisions per function
+- `exc.py` — exception hierarchy: `Unsupported` (graph break), `RestartAnalysis`
+  (restart tracing), `ObservedException` (user exceptions during tracing),
+  `BackendCompilerFailed`
+- `config.py` — configuration flags, supports `config.patch()` context
+  manager/decorator
+- `bytecode_transformation.py` / `bytecode_analysis.py` — low-level bytecode
+  manipulation, liveness analysis
+- `pgo.py` — profile-guided optimization for dynamic shapes
+- `polyfills/` — traceable replacements for stdlib functions
+- `repro/` — reproduction/minification tools
+
+## Graph Breaks
+
+Call `unimplemented()` (from `exc.py`) to trigger a graph break:
+
+```python
+from torch._dynamo.exc import unimplemented
+from torch._dynamo import graph_break_hints
+
+unimplemented(
+    gb_type="short_category_name",
+    context=f"dynamic details: {value}",
+    explanation="Human-readable explanation of why this breaks the graph.",
+    hints=[*graph_break_hints.SUPPORTABLE],
+)
+```
+
+- `gb_type`: Context-free category (no dynamic strings).
+- `context`: Developer-facing details (can be dynamic).
+- `explanation`: User-facing explanation (can be dynamic).
+- `hints`: Use constants from `graph_break_hints.py`: `SUPPORTABLE`,
+  `FUNDAMENTAL`, `DIFFICULT`, `DYNAMO_BUG`, `USER_ERROR`,
+  `CAUSED_BY_EARLIER_GRAPH_BREAK`.
+
+The `break_graph_if_unsupported` decorator on instruction handlers catches
+`Unsupported`, logs the graph break, updates the `SpeculationLog`, and restarts
+analysis. On the second pass, the partial graph is compiled at the break point
+and a resume function handles the rest.
+
+## Testing
+
+Tests live in `test/dynamo/`. Use `torch._dynamo.test_case.TestCase` as base
+class — it calls `torch._dynamo.reset()` in setUp/tearDown and patches config
+for strict error checking.
+
+```bash
+python test/dynamo/test_misc.py                       # whole file
+python test/dynamo/test_misc.py MiscTests.test_foo    # single test
+python test/dynamo/test_misc.py -k test_foo           # pattern match
+```
+
+### Common patterns
+
+The default backend to `torch.compile()` is `backend="eager"`.
+
+**CompileCounter** — count compilations and graph ops:
+```python
+cnt = torch._dynamo.testing.CompileCounter()
+
+@torch.compile(backend=cnt)
+def fn(x):
+    return x + 1
+
+fn(torch.randn(10))
+self.assertEqual(cnt.frame_count, 1)
+self.assertEqual(cnt.op_count, 1)
+```
+
+**fullgraph=True** — assert no graph breaks:
+```python
+torch.compile(fn, backend="eager", fullgraph=True)(x)
+```
+
+**EagerAndRecordGraphs** — inspect captured FX graphs:
+```python
+backend = torch._dynamo.testing.EagerAndRecordGraphs()
+torch.compile(fn, backend=backend)(x)
+graph = backend.graphs[0]
+```
+
+**normalize_gm + assertExpectedInline** — snapshot test graph output:
+```python
+from torch._dynamo.testing import normalize_gm
+self.assertExpectedInline(
+    normalize_gm(backend.graphs[0].print_readable(False)),
+    """\
+expected output here
+""",
+)
+```
+
+Call `torch._dynamo.reset()` within a test when testing multiple compilation
+scenarios in a single test method. The base class handles setUp/tearDown reset
+automatically.
+
+## Debugging
+
+### TORCH_LOGS
+
+```bash
+TORCH_LOGS="graph_breaks" python script.py       # see graph breaks
+TORCH_LOGS="guards,recompiles" python script.py  # see guards and recompilation reasons
+TORCH_LOGS="graph_code" python script.py         # see captured FX graph code
+TORCH_LOGS="+dynamo" python script.py            # full debug logging
+TORCH_LOGS="bytecode" python script.py           # see bytecode transformations
+```
+
+### Structured tracing (for production)
+
+```bash
+TORCH_TRACE=/path/to/dir python script.py  # explicit trace directory
+```
+
+Analyze with `tlparse`.
+
+### Compile-time profiling
+
+`TORCH_COMPILE_DYNAMO_PROFILER=1` prints per-function cumtime/tottime
+(cProfile-style) showing where Dynamo spends time during tracing. Set to a
+file path instead to save a profile loadable by `snakeviz`.
+
+### comptime.breakpoint() (`comptime.py`)
+
+Drops into pdb during **compilation** to inspect Dynamo state. Call
+`comptime.breakpoint()` in user code; in the pdb session use `ctx`
+(`ComptimeContext`) to call `print_locals()`, `print_bt()`, `print_graph()`,
+or `get_local("x").as_fake()`.
+
+### Bytecode Debugger (`bytecode_debugger.py`)
+
+pdb-like debugger for stepping through Dynamo-generated bytecode. Useful for
+debugging segfaults (no Python traceback) and codegen errors.
+
+```python
+with torch._dynamo.bytecode_debugger.debug():
+    my_compiled_fn(x)
+```
+
+**Programmatic breakpoints** (no graph break): call
+`torch._dynamo.bytecode_debugger.breakpoint()` in user code, or
+`codegen.extend_output(create_breakpoint())` in codegen. Auto-activates
+without an explicit `debug()` wrapper.
+
+**Segfault debugging**: `v` (verbose) then `c` (continue) — every instruction
+is printed with `flush=True` before execution, so the last line before a crash
+is the culprit. On exceptions, the debugger stops at the faulting instruction
+automatically.
+
+## C++ Runtime (`torch/csrc/dynamo/`)
+
+The C/C++ layer implements the PEP 523 frame evaluation hook, the cache, and
+the guard evaluation tree. Performance-critical runtime on every Python frame.
+
+### Frame Evaluation
+
+**`eval_frame.c`** — Installs a custom frame evaluation function via
+`_PyInterpreterState_SetEvalFrameFunc`. A thread-local callback controls
+behavior: `None` (disabled), `Py_False` (run-only / cache lookup), or a
+callable (full Dynamo).
+
+**`eval_frame_cpp.cpp`** — `dynamo__custom_eval_frame` is called for every
+frame: gets `ExtraState` from the code object, builds a `FrameLocalsMapping`
+(O(1) access to locals without dict materialization), evaluates guards via
+`run_root_guard_manager()` across all `CacheEntry`s (LRU ordered). On cache
+hit, executes compiled code via a shadow frame (`dynamo_eval_custom_code_impl`
+copies `localsplus` into a new frame with the compiled code object). On miss,
+calls the Python callback to trigger compilation.
+
+### Cache
+
+**`extra_state.cpp/.h`** — `ExtraState` is attached per code object via
+`_PyCode_SetExtra`. Contains a `cache_entry_list` (LRU linked list),
+`frame_state` (dynamic shapes detection), and `FrameExecStrategy`.
+
+**`cache_entry.cpp/.h`** — Each `CacheEntry` stores a `RootGuardManager*`
+(raw C++ pointer for fast guard eval), the compiled code object, and the
+backend.
+
+### Guard Evaluation Tree (`guards.cpp`)
+
+Guards are organized as a C++ tree (~7800 lines) mirroring the data access
+pattern. `RootGuardManager` is the root, receiving a `FrameLocalsMapping`.
+Each `GuardManager` node has leaf guards and child accessors.
+
+**LeafGuard** subclasses: `TYPE_MATCH` (Py_TYPE pointer comparison),
+`ID_MATCH`, `EQUALS_MATCH`, `TENSOR_MATCH` (dtype/device/shape/strides/dispatch
+keys in C++), `DICT_VERSION`, `GLOBAL_STATE` (grad mode, autocast, etc.).
+
+**GuardAccessor** subclasses define tree edges: `FrameLocalsGuardAccessor`
+(O(1) index), `GetAttrGuardAccessor`, `DictGetItemGuardAccessor`,
+`GlobalsGuardAccessor`, etc.
+
+Key optimizations: fail-fast accessor reordering, dict version tag matching to
+skip subtrees, `FrameLocalsMapping` avoids dict construction,
+`check_nopybind()` avoids pybind11 overhead.
+
+### Other C++ files
+
+- `framelocals_mapping.cpp` — O(1) frame locals/cells/freevars access
+- `cpython_defs.c` — copied CPython internals for frame manipulation
+- `init.cpp` — `torch._C._dynamo` module and pybind11 bindings
+- `debug_macros.h` — `DEBUG_TRACE`, `NULL_CHECK`, `INSPECT(...)` (drops into
+  pdb from C)
+- `compiled_autograd.cpp/.h` — compiled autograd engine
diff --git a/torch/_dynamo/__init__.py b/torch/_dynamo/__init__.py
index 89df282d37102..f258338c3baf9 100644
--- a/torch/_dynamo/__init__.py
+++ b/torch/_dynamo/__init__.py
@@ -12,6 +12,7 @@
 
 from . import (
     aot_compile,
+    bytecode_debugger,
     config,
     convert_frame,
     eval_frame,
@@ -38,6 +39,7 @@
     mark_static_address,
     maybe_mark_dynamic,
     nonstrict_trace,
+    override_cudagraphs,
     patch_dynamo_config,
     run,
     set_stance,
@@ -79,6 +81,7 @@
 __all__ = [
     "allow_in_graph",
     "assume_constant_result",
+    "bytecode_debugger",
     "config",
     "disable",
     "disable_nested_graph_breaks",
@@ -106,6 +109,7 @@
     "reset",
     "reset_recompile_user_contexts",
     "run",
+    "override_cudagraphs",
     "error_on_graph_break",
     "set_recursion_limit",
     "set_stance",
@@ -164,6 +168,12 @@ def reset() -> None:
 
         # Reset cudagraph trees unconditionally since they are global state
         # not tied to a specific backend instance
+        from torch._higher_order_ops.triton_kernel_wrap import kernel_side_table
+        from torch._higher_order_ops.wrap import inductor_code_side_table
+
+        kernel_side_table.reset_table()
+        inductor_code_side_table.reset_table()
+
         if torch.cuda.is_available():
             from torch._inductor.cudagraph_trees import reset_cudagraph_trees
 
diff --git a/torch/_dynamo/_trace_wrapped_higher_order_op.py b/torch/_dynamo/_trace_wrapped_higher_order_op.py
index 43e6f5be17f47..93772b05e3022 100644
--- a/torch/_dynamo/_trace_wrapped_higher_order_op.py
+++ b/torch/_dynamo/_trace_wrapped_higher_order_op.py
@@ -28,7 +28,7 @@
 compiled autograd do we inline into the function.
 """
 
-from typing import Any, Optional
+from typing import Any
 
 import torch
 import torch.utils._pytree as pytree
@@ -141,7 +141,7 @@ def __torch_function__(
         func: OpOverload,
         types: tuple[torch._C._TensorMeta, ...],
         args: tuple[object, ...] = (),
-        kwargs: Optional[dict[str, object]] = None,
+        kwargs: dict[str, object] | None = None,
     ) -> object:
         if func is torch.Tensor.__getitem__:
             tensor_to_index = args[0]
@@ -192,7 +192,7 @@ def _assert_meta(
 def inner_trace(
     mode: ProxyTorchDispatchMode,
     *args: Any,
-    bw_state: Optional[BackwardState] = None,
+    bw_state: BackwardState | None = None,
     **kwargs: Any,
 ) -> Any:
     def self_invoke(*args: Any, **dyn_kwargs: Any) -> Any:
diff --git a/torch/_dynamo/aot_compile.py b/torch/_dynamo/aot_compile.py
index e35a6eea57654..7ac01d6cb5488 100644
--- a/torch/_dynamo/aot_compile.py
+++ b/torch/_dynamo/aot_compile.py
@@ -3,7 +3,9 @@
 import inspect
 import io
 import logging
+import os
 import pickle
+import tempfile
 import types
 from collections.abc import Callable, Sequence
 from contextlib import AbstractContextManager, ExitStack, nullcontext
@@ -171,6 +173,20 @@ class AOTCompileSaveResult:
     serialized_data: bytes
 
 
+def atomic_write_binary(file_path: str, data: bytes):
+    dir_name = os.path.dirname(file_path) or "."
+
+    with tempfile.NamedTemporaryFile(
+        dir=dir_name, delete=False, mode="wb"
+    ) as temp_file:
+        temp_path = temp_file.name
+        temp_file.write(data)
+        temp_file.flush()
+        os.fsync(temp_file.fileno())
+
+    os.replace(temp_path, file_path)
+
+
 @dataclass
 class AOTCompiledFunction:
     _artifacts: CompileArtifacts
@@ -229,10 +245,9 @@ def source_info(self) -> "SourceInfo":
     def save_compiled_function(
         self, path: str, external_data: dict[str, Any] | None = None
     ) -> AOTCompileSaveResult:
-        with open(path, "wb") as f:
-            result = type(self).serialize(self, external_data)
-            f.write(result.serialized_data)
-            return result
+        result = type(self).serialize(self, external_data)
+        atomic_write_binary(path, result.serialized_data)
+        return result
 
     @classmethod
     def serialize(
@@ -302,6 +317,9 @@ def aot_compile_fullgraph(
     from torch._dynamo.guards import CheckFunctionManager
     from torch._dynamo.package import SourceInfo
     from torch._dynamo.utils import dynamo_timed, get_metrics_context
+    from torch._dynamo.variables.torch_function import (
+        torch_function_mode_stack_state_mgr,
+    )
     from torch._guards import TracingContext
 
     args, kwargs = example_inputs
@@ -317,6 +335,7 @@ def aot_compile_fullgraph(
         dynamo_timed("fullgraph_capture"),
         torch._functorch.config.patch(strict_autograd_cache=True),
         dynamic_ctx,
+        torch_function_mode_stack_state_mgr,
     ):
         capture_output = convert_frame.fullgraph_capture(model, args, kwargs)
         graph_capture_output = capture_output.graph_capture_output
@@ -368,12 +387,20 @@ def new_guard_filter_fn(
             compiled_fn = backend(
                 backend_input.graph_module, backend_input.example_inputs
             )
-            # If Inductor backend is used, grab the compiled_fn from PrecompileContext
+            # If Inductor backend or AOTAutograd-based backend is used,
+            # wrap the compiled_fn for serialization.
             # TODO: this should be replaced once we make the backend return the SerializableCallable directly.
-            if isinstance(backend, torch._TorchCompileInductorWrapper) or (
-                hasattr(backend, "compiler_fn")
-                and isinstance(
-                    backend.compiler_fn, torch._dynamo.backends.common.AotAutograd
+            if (
+                isinstance(backend, torch._TorchCompileInductorWrapper)
+                or (
+                    hasattr(backend, "compiler_fn")
+                    and isinstance(
+                        backend.compiler_fn, torch._dynamo.backends.common.AotAutograd
+                    )
+                )
+                or (
+                    hasattr(compiled_fn, "serialize")
+                    and compiled_fn.serialize is not None
                 )
             ):
                 compiled_fn = BundledAOTAutogradSerializableCallable(compiled_fn)
@@ -388,9 +415,17 @@ def new_guard_filter_fn(
                 + f"from backend {compiler_fn}) does not implement SerializableCallable."
             )
 
-        check_fn = graph_capture_output.build_guards(
-            fn.__code__, hooks=hooks, save=True, strict_error=True
-        )
+        # Temporarily restore the mode stack so guard expressions that
+        # reference modes can evaluate, matching the compile_inner path.
+        build_guards_ctx = ExitStack()
+        if torch_function_mode_stack_state_mgr.stack:
+            build_guards_ctx.enter_context(
+                torch_function_mode_stack_state_mgr.temp_restore_stack()
+            )
+        with build_guards_ctx:
+            check_fn = graph_capture_output.build_guards(
+                fn.__code__, hooks=hooks, save=True, strict_error=True
+            )
 
         assert check_fn.guards_state is not None
 
diff --git a/torch/_dynamo/backends/common.py b/torch/_dynamo/backends/common.py
index ad50cba3ac907..0b60e0d74a623 100644
--- a/torch/_dynamo/backends/common.py
+++ b/torch/_dynamo/backends/common.py
@@ -19,7 +19,7 @@
 import contextlib
 import functools
 import logging
-from collections.abc import Callable, Iterable
+from collections.abc import Callable, Iterable, Sequence
 from typing import Any
 from typing_extensions import ParamSpec, TypeVar
 from unittest.mock import patch
@@ -47,7 +47,7 @@ def __init__(self, **kwargs: Any) -> None:
         self.kwargs = kwargs
 
     def __call__(
-        self, gm: torch.fx.GraphModule, example_inputs: Iterable[Any], **kwargs: Any
+        self, gm: torch.fx.GraphModule, example_inputs: Sequence[Any], **kwargs: Any
     ) -> Callable[..., Any]:
         if kwargs:
             log.warning("aot_autograd-based backend ignoring extra kwargs %s", kwargs)
@@ -70,7 +70,6 @@ def __call__(
         if use_fallback:
             log.debug("Unable to use AOT Autograd because graph has mutation")
             counters["aot_autograd"]["not_ok"] += 1
-            # pyrefly: ignore [bad-return]
             return gm
 
         def wrap_bw_compiler(bw_compiler_fn: Callable[P, R]) -> Callable[..., R]:
@@ -150,7 +149,7 @@ def mem_efficient_fusion_kwargs(use_decomps: bool) -> dict[str, Any]:
     }
 
     if use_decomps:
-        # pyrefly: ignore [bad-typed-dict-key, unsupported-operation]
+        # pyrefly: ignore [bad-typed-dict-key]
         kwargs["decompositions"] = default_decompositions
 
     return kwargs
diff --git a/torch/_dynamo/backends/cudagraphs.py b/torch/_dynamo/backends/cudagraphs.py
index 8d4637dea6627..27f2e19b98f1c 100644
--- a/torch/_dynamo/backends/cudagraphs.py
+++ b/torch/_dynamo/backends/cudagraphs.py
@@ -24,7 +24,7 @@
 import functools
 from collections import defaultdict
 from collections.abc import Callable, Sequence
-from typing import Any, Optional
+from typing import Any
 
 import torch
 import torch.fx
@@ -103,7 +103,7 @@ def get_device_node_mapping(
 
 def check_for_mutation_ignore_cuda_graph_managed_tensor(
     aot_model: torch.fx.GraphModule, num_fixed: int
-) -> Optional[str]:
+) -> str | None:
     mutation_indices = find_input_mutations(aot_model.graph) - set(range(num_fixed))
     if not mutation_indices:
         return None
@@ -112,7 +112,7 @@ def check_for_mutation_ignore_cuda_graph_managed_tensor(
     return get_mutation_stack_trace(placeholders, mutation_indices)
 
 
-def check_for_skip(aot_model: torch.fx.GraphModule, num_fixed: int) -> Optional[str]:
+def check_for_skip(aot_model: torch.fx.GraphModule, num_fixed: int) -> str | None:
     if not config.cudagraph_backend_support_input_mutation:
         if mut_skip := check_for_mutation_ignore_cuda_graph_managed_tensor(
             aot_model, num_fixed
@@ -136,7 +136,7 @@ def get_device_index(gm: torch.fx.GraphModule) -> int:
     return device.index
 
 
-def get_stack_traces(gm: torch.fx.GraphModule) -> list[Optional[str]]:
+def get_stack_traces(gm: torch.fx.GraphModule) -> list[str | None]:
     output = output_node(gm)
     assert len(output.args) == 1
     args = output.args[0]
@@ -264,7 +264,6 @@ def cudagraphs_inner(
     """This isn't registered as a backend, but is used in some benchmarks"""
     assert isinstance(inputs, (list, tuple))
     if copy_inputs:
-        # pyrefly: ignore [bad-argument-type]
         static_inputs = [torch.zeros_like(x) for x in inputs]
     else:
         static_inputs = list(inputs)
diff --git a/torch/_dynamo/backends/debugging.py b/torch/_dynamo/backends/debugging.py
index fa2e54c42e010..eacb304abf7b1 100644
--- a/torch/_dynamo/backends/debugging.py
+++ b/torch/_dynamo/backends/debugging.py
@@ -28,7 +28,7 @@
 import logging
 from collections.abc import Callable, Iterable
 from importlib import import_module
-from typing import Any, Optional, TYPE_CHECKING, Union
+from typing import Any, TYPE_CHECKING
 
 import torch
 from functorch.compile import min_cut_rematerialization_partition
@@ -36,6 +36,7 @@
 from torch._dynamo.output_graph import GraphCompileReason
 from torch._functorch import config as functorch_config
 from torch._functorch.compilers import ts_compile
+from torch._inductor.output_code import OutputCode
 
 from .common import aot_autograd
 from .registry import CompiledFn, CompilerFn, register_debug_backend as register_backend
@@ -166,6 +167,7 @@ def invoke_subgraph_inner_compiler(
     from torch._higher_order_ops.invoke_subgraph import invoke_subgraph_infer
 
     @disable
+    # pyrefly: ignore [deprecated]
     @torch._dynamo.allow_in_graph
     def invoke_subgraph_wrapper_unboxed(*operands: Any) -> Any:
         return invoke_subgraph_infer(subgraph, *operands)
@@ -260,6 +262,54 @@ def invoke_subgraph(
     )(gm, fake_tensor_inputs)
 
 
+@dataclasses.dataclass
+class AOTEagerOutputCode(OutputCode):
+    """
+    An OutputCode that wraps a GraphModule for eager-mode execution.
+
+    This allows non-inductor backends (like aot_eager) to participate in
+    the bundled autograd cache and aot_compile serialization flow.
+    """
+
+    gm: torch.fx.GraphModule | None = None
+    _serialized_gm: bytes | None = dataclasses.field(default=None, init=False)
+
+    def __call__(self, inputs: Any) -> Any:
+        assert self.gm is not None
+        return self.gm.forward(inputs)
+
+    def prepare_for_serialization(self) -> None:
+        from torch.fx._graph_pickler import GraphPickler, Options
+
+        assert self.gm is not None
+        for node in self.gm.graph.nodes:
+            node.meta.pop("nn_module_stack", None)
+            node.meta.pop("source_fn_stack", None)
+            node.meta.pop("example_value", None)
+
+        self._serialized_gm = GraphPickler.dumps(self.gm, Options(ops_filter=None))
+        self.gm = None
+
+    def post_compile(self, *args: Any, **kwargs: Any) -> None:
+        if self.gm is None and self._serialized_gm is not None:
+            from torch._subclasses import FakeTensorMode
+            from torch.fx._graph_pickler import GraphPickler
+            from torch.fx.experimental.symbolic_shapes import ShapeEnv
+            from torch.fx.graph import _BoxedCodeGen
+
+            fake_mode = FakeTensorMode(shape_env=ShapeEnv())
+            gm = GraphPickler.loads(self._serialized_gm, fake_mode)
+            assert isinstance(gm, torch.fx.GraphModule)
+            self.gm = gm
+            assert isinstance(self.gm, torch.fx.GraphModule)
+            self.gm.graph.set_codegen(_BoxedCodeGen())
+            self.gm.recompile()
+            self._serialized_gm = None
+
+    def set_triton_bundle(self, triton_bundle: Any) -> None:
+        pass
+
+
 # used boxed call to discard inputs when they are no longer needed
 def boxed_nop(
     fx_g: torch.fx.GraphModule, example_inputs: list[torch.Tensor]
@@ -270,6 +320,11 @@ def boxed_nop(
     fx_g.graph.set_codegen(_BoxedCodeGen())
     fx_g.recompile()
 
+    if functorch_config.force_autograd_cache or functorch_config.bundled_autograd_cache:
+        result = AOTEagerOutputCode(gm=fx_g)
+        result._boxed_call = True  # type: ignore[attr-defined]
+        return result
+
     # Wrap the forward method in a function so we can set _boxed_call attribute
     forward_fn = fx_g.forward
 
@@ -312,7 +367,7 @@ def run(args: Any) -> Any:
 def fake_crossref_boxed_nop(
     fx_g: torch.fx.GraphModule,
     example_inputs: list[torch.Tensor],
-    ignore_op_fn: Optional[Callable[[torch._ops.OpOverload], bool]] = None,
+    ignore_op_fn: Callable[[torch._ops.OpOverload], bool] | None = None,
 ) -> Callable[..., Any]:
     from torch.fx.graph import _BoxedCodeGen
 
@@ -352,8 +407,8 @@ def get_nop_func() -> Callable[
 def aot_eager(
     gm: torch.fx.GraphModule,
     fake_tensor_inputs: list[torch.Tensor],
-    fw_compiler: Optional[Callable[..., Any]] = None,
-    bw_compiler: Optional[Callable[..., Any]] = None,
+    fw_compiler: Callable[..., Any] | None = None,
+    bw_compiler: Callable[..., Any] | None = None,
     **kwargs: Any,
 ) -> Callable[..., Any]:
     return aot_autograd(
@@ -545,9 +600,9 @@ class ExplainOutput:
     graph_break_count: int
     break_reasons: list[GraphCompileReason]
     op_count: int
-    ops_per_graph: Optional[list[list["Target"]]] = None
-    out_guards: Optional[list[_guards.Guard]] = None
-    compile_times: Optional[str] = None
+    ops_per_graph: list[list["Target"]] | None = None
+    out_guards: list[_guards.Guard] | None = None
+    compile_times: str | None = None
 
     def __str__(self) -> str:
         output = f"Graph Count: {self.graph_count}\n"
@@ -645,7 +700,7 @@ def fn(x):
         print(eb.output())
     """
 
-    def __init__(self, backend: Union[CompilerFn, str]) -> None:
+    def __init__(self, backend: CompilerFn | str) -> None:
         from .registry import lookup_backend
 
         self.backend = lookup_backend(backend)
diff --git a/torch/_dynamo/backends/distributed.py b/torch/_dynamo/backends/distributed.py
index 94164a070eb45..bbd8dc53eaa70 100644
--- a/torch/_dynamo/backends/distributed.py
+++ b/torch/_dynamo/backends/distributed.py
@@ -20,7 +20,7 @@
 import traceback
 from collections.abc import Callable
 from dataclasses import dataclass, field
-from typing import Any, Optional, TYPE_CHECKING
+from typing import Any, TYPE_CHECKING
 from unittest import mock
 
 import torch
@@ -84,7 +84,7 @@ def bucket_has_external_output(bucket: Bucket) -> bool:
 
 def pretty_print_buckets(buckets: list[Bucket], bucket_bytes_cap: int) -> None:
     headers = ("Index", "Size (b)", "Param Names")
-    rows: list[tuple[Optional[int], Optional[int], str]] = []
+    rows: list[tuple[int | None, int | None, str]] = []
     # pyrefly: ignore [implicit-any]
     extended_buckets = []
     for idx, bucket in enumerate(reversed(buckets)):
@@ -433,7 +433,7 @@ def __init__(
         self,
         bucket_bytes_cap: int,
         backend_compile_fn: CompilerFn,
-        first_bucket_cap: Optional[int] = None,
+        first_bucket_cap: int | None = None,
     ) -> None:
         if first_bucket_cap is not None:
             self.first_bucket_cap = first_bucket_cap
diff --git a/torch/_dynamo/backends/registry.py b/torch/_dynamo/backends/registry.py
index fcd8b93afd8c3..f8b311ceb52f8 100644
--- a/torch/_dynamo/backends/registry.py
+++ b/torch/_dynamo/backends/registry.py
@@ -62,7 +62,7 @@ def my_compiler_function(fx_graph, example_inputs):
 import logging
 from collections.abc import Callable, Sequence
 from importlib.metadata import EntryPoint
-from typing import Any, Optional, Protocol, Union
+from typing import Any, Protocol
 
 import torch
 from torch import fx
@@ -77,13 +77,13 @@ def __call__(self, *args: torch.Tensor) -> tuple[torch.Tensor, ...]: ...
 
 CompilerFn = Callable[[fx.GraphModule, list[torch.Tensor]], CompiledFn]
 
-_BACKENDS: dict[str, Optional[EntryPoint]] = {}
+_BACKENDS: dict[str, EntryPoint | None] = {}
 _COMPILER_FNS: dict[str, CompilerFn] = {}
 
 
 def register_backend(
-    compiler_fn: Optional[CompilerFn] = None,
-    name: Optional[str] = None,
+    compiler_fn: CompilerFn | None = None,
+    name: str | None = None,
     tags: Sequence[str] = (),
 ) -> Callable[..., Any]:
     """
@@ -116,7 +116,7 @@ def register_backend(
 )
 
 
-def lookup_backend(compiler_fn: Union[str, CompilerFn]) -> CompilerFn:
+def lookup_backend(compiler_fn: str | CompilerFn) -> CompilerFn:
     """Expand backend strings to functions"""
     if isinstance(compiler_fn, str):
         if compiler_fn not in _BACKENDS:
diff --git a/torch/_dynamo/backends/tvm.py b/torch/_dynamo/backends/tvm.py
index 88d16c95551a8..5aa20dd71f1fa 100644
--- a/torch/_dynamo/backends/tvm.py
+++ b/torch/_dynamo/backends/tvm.py
@@ -29,7 +29,7 @@
 from collections.abc import Callable
 from pathlib import Path
 from types import MappingProxyType
-from typing import Any, Optional
+from typing import Any
 
 import torch
 from torch import fx
@@ -47,7 +47,7 @@ def tvm(
     gm: fx.GraphModule,
     example_inputs: list[torch.Tensor],
     *,
-    options: Optional[MappingProxyType[str, Any]] = None,
+    options: MappingProxyType[str, Any] | None = None,
 ) -> Callable[..., Any]:
     if options is None:
         options = MappingProxyType({"scheduler": None, "trials": 20000, "opt_level": 3})
@@ -79,7 +79,7 @@ def tvm(
     opt_level = options.get("opt_level", 3)
 
     if scheduler == "auto_scheduler":
-        # pyrefly: ignore [import-error, missing-import]
+        # pyrefly: ignore [missing-import]
         from tvm import auto_scheduler
 
         with (
@@ -91,7 +91,7 @@ def tvm(
         ):
             lib = relay.build(mod, target=target, params=params)
     elif scheduler == "meta_schedule":
-        # pyrefly: ignore [import-error, missing-import]
+        # pyrefly: ignore [missing-import]
         from tvm import meta_schedule as ms
 
         with tempfile.TemporaryDirectory() as work_dir:
diff --git a/torch/_dynamo/bytecode_analysis.py b/torch/_dynamo/bytecode_analysis.py
index 592d959c01ed8..d872dcc4cd106 100644
--- a/torch/_dynamo/bytecode_analysis.py
+++ b/torch/_dynamo/bytecode_analysis.py
@@ -17,7 +17,7 @@
 import dis
 import itertools
 import sys
-from typing import Any, TYPE_CHECKING, Union
+from typing import Any, TYPE_CHECKING
 
 
 if TYPE_CHECKING:
@@ -203,8 +203,8 @@ class FixedPointBox:
 
 @dataclasses.dataclass
 class StackSize:
-    low: Union[int, float]
-    high: Union[int, float]
+    low: int | float
+    high: int | float
     fixed_point: FixedPointBox
 
     def zero(self) -> None:
@@ -227,7 +227,7 @@ def exn_tab_jump(self, depth: int) -> None:
             self.fixed_point.value = False
 
 
-def stacksize_analysis(instructions: list["Instruction"]) -> Union[int, float]:
+def stacksize_analysis(instructions: list["Instruction"]) -> int | float:
     assert instructions
     fixed_point = FixedPointBox()
     stack_sizes = {
diff --git a/torch/_dynamo/bytecode_debugger.py b/torch/_dynamo/bytecode_debugger.py
new file mode 100644
index 0000000000000..42cf59b58de38
--- /dev/null
+++ b/torch/_dynamo/bytecode_debugger.py
@@ -0,0 +1,1144 @@
+"""
+Bytecode debugger for Dynamo-optimized code.
+
+This module provides a pdb-like debugger for stepping through Python bytecode
+one instruction at a time, with the ability to inspect the value stack,
+locals, and globals.
+
+Usage:
+    >>> # xdoctest: +SKIP
+    >>> import torch
+    >>>
+    >>> @torch.compile
+    >>> def my_fn(x):
+    ...     return x + 1
+    >>>
+    >>> with torch._dynamo.bytecode_debugger.debug():
+    ...     my_fn(torch.randn(3))  # Debugger activates on Dynamo-generated code
+
+Programmatic breakpoints (for Dynamo developers):
+    In PyCodegen, use create_breakpoint() to insert a debugger stop:
+
+        from torch._dynamo.bytecode_transformation import create_breakpoint
+        codegen.extend_output(create_breakpoint())
+"""
+
+from __future__ import annotations
+
+import dis
+import sys
+import types
+from contextlib import contextmanager
+from dataclasses import dataclass, field
+from typing import Any, cast, TYPE_CHECKING
+
+
+if TYPE_CHECKING:
+    from collections.abc import Callable, Generator
+
+    from typing_extensions import Self
+
+from .bytecode_transformation import convert_instruction, Instruction, instruction_size
+
+
+# Python 3.12+ has sys.monitoring for efficient instruction-level tracing
+_HAS_SYS_MONITORING = hasattr(sys, "monitoring")
+
+
+# Sentinel for breakpoints inserted by PyCodegen via create_breakpoint()
+class _BreakpointMarker:
+    """Sentinel constant that signals the debugger to break.
+
+    Usage in PyCodegen:
+        from torch._dynamo.bytecode_transformation import create_breakpoint
+        codegen.extend_output(create_breakpoint())
+    """
+
+    __slots__ = ()
+    _instance: _BreakpointMarker | None = None
+
+    def __new__(cls) -> Self:
+        if cls._instance is None:
+            cls._instance = cast("_BreakpointMarker", object.__new__(cls))
+        return cls._instance
+
+    def __repr__(self) -> str:
+        return "<BREAKPOINT>"
+
+
+BREAKPOINT_MARKER = _BreakpointMarker()
+
+
+# Import NULL_STACK_VALUE sentinel from C++ module
+# This is returned by _get_frame_value_stack_with_depth for NULL stack slots
+from torch._C._dynamo.eval_frame import NULL_STACK_VALUE  # noqa: F401
+
+
+@dataclass
+class CodeInfo:
+    """Per-code-object data shared across all frames executing the same code."""
+
+    code: types.CodeType
+    instructions: list[Instruction]
+    offset_to_inst: dict[int, Instruction]
+    offset_to_index: dict[int, int]
+    index_width: int
+    offset_width: int
+    breakpoints: set[int] = field(default_factory=set)
+
+
+@dataclass
+class DebuggerState:
+    """Per-frame debugging state."""
+
+    code_info: CodeInfo
+    frame: types.FrameType
+    current_offset: int = -1  # -1 indicates first instruction not yet seen
+    current_stack_depth: int = 0  # Tracked dynamically
+    first_instruction_seen: bool = False
+    step_mode: bool = True
+    step_count: int = 0  # Remaining steps for "s [n]" command (0 = prompt each time)
+    user_locals: dict[str, Any] = field(
+        default_factory=dict
+    )  # User-defined variables from debugger
+    last_command: str = "s"  # Last command for repeat on empty input
+    list_index: int | None = (
+        None  # Current position for 'l' command (None = use current)
+    )
+
+    @property
+    def code(self) -> types.CodeType:
+        return self.code_info.code
+
+    @property
+    def instructions(self) -> list[Instruction]:
+        return self.code_info.instructions
+
+    @property
+    def offset_to_inst(self) -> dict[int, Instruction]:
+        return self.code_info.offset_to_inst
+
+    @property
+    def offset_to_index(self) -> dict[int, int]:
+        return self.code_info.offset_to_index
+
+    @property
+    def index_width(self) -> int:
+        return self.code_info.index_width
+
+    @property
+    def offset_width(self) -> int:
+        return self.code_info.offset_width
+
+    @property
+    def breakpoints(self) -> set[int]:
+        return self.code_info.breakpoints
+
+
+class _DebugContext:
+    """Internal debug context that manages the debugging session."""
+
+    def __init__(self) -> None:
+        self._code_info: dict[types.CodeType, CodeInfo] = {}
+        self._frame_states: dict[types.FrameType, DebuggerState] = {}
+        self._active = False
+        self._tracked_codes: set[types.CodeType] = set()
+        self._old_trace: Callable[..., Any] | None = None
+        self._prev_callback: Callable[..., Any] | None = None
+        self._return_from_frame: types.FrameType | None = None
+        self._stop_after_return: bool = False
+        self._next_in_frame: types.FrameType | None = None
+        self._next_count: int = 0
+        self._verbose: bool = False
+        self._stop_at_new_code: bool = True
+        self._quitting: bool = False
+        if _HAS_SYS_MONITORING:
+            self._tool_id = sys.monitoring.DEBUGGER_ID
+
+    def get_instructions(self, code: types.CodeType | None = None) -> list[Instruction]:
+        """Get the list of instructions for a tracked code object.
+
+        Args:
+            code: The code object to get instructions for. If None, returns
+                  instructions from the most recently tracked code object.
+
+        Returns:
+            List of instructions, or empty list if the code is not tracked.
+        """
+        if not self._code_info:
+            return []
+        if code is None:
+            info = next(reversed(self._code_info.values()))
+        else:
+            info = self._code_info.get(code)
+            if info is None:
+                return []
+        return info.instructions
+
+    def get_tracked_codes(self) -> list[types.CodeType]:
+        """Get all code objects that have been tracked by the debugger."""
+        return list(self._code_info.keys())
+
+    @staticmethod
+    def is_programmatic_breakpoint(inst: Instruction) -> bool:
+        return inst.opname == "LOAD_CONST" and inst.argval is BREAKPOINT_MARKER
+
+    def _get_or_create_code_info(self, code: types.CodeType) -> CodeInfo:
+        """Get or create CodeInfo for a code object."""
+        if code not in self._code_info:
+            # Use dis.get_instructions directly to preserve original offsets.
+            # cleaned_instructions strips EXTENDED_ARG and recalculates offsets,
+            # which would cause mismatches with offsets reported by callbacks.
+            instructions = [convert_instruction(i) for i in dis.get_instructions(code)]
+
+            # In 3.11+, instructions have inline cache entries that occupy
+            # bytecode space (e.g. BINARY_OP is 2 bytes + 2 bytes cache).
+            # dis.get_instructions only reports the instruction start offset,
+            # but sys.monitoring RAISE reports offsets within cache entries.
+            # Map the full byte range of each instruction so lookups succeed.
+            offset_to_inst: dict[int, Instruction] = {}
+            offset_to_index: dict[int, int] = {}
+            for i, inst in enumerate(instructions):
+                if inst.offset is not None:
+                    inst_size = instruction_size(inst)
+                    for off in range(inst.offset, inst.offset + inst_size, 2):
+                        offset_to_inst[off] = inst
+                        offset_to_index[off] = i
+
+            max_index = len(instructions) - 1 if instructions else 0
+            max_offset = max(
+                (inst.offset for inst in instructions if inst.offset is not None),
+                default=0,
+            )
+
+            # Pre-populate breakpoints from BREAKPOINT_MARKER instructions
+            programmatic_breakpoints: set[int] = {
+                i
+                for i, inst in enumerate(instructions)
+                if self.is_programmatic_breakpoint(inst)
+            }
+
+            self._code_info[code] = CodeInfo(
+                code=code,
+                instructions=instructions,
+                offset_to_inst=offset_to_inst,
+                offset_to_index=offset_to_index,
+                index_width=max(1, len(str(max_index))),
+                offset_width=max(1, len(str(max_offset))),
+                breakpoints=programmatic_breakpoints,
+            )
+        return self._code_info[code]
+
+    def _get_or_create_frame_state(
+        self, code: types.CodeType, frame: types.FrameType
+    ) -> DebuggerState:
+        """Get or create per-frame DebuggerState."""
+        if frame not in self._frame_states:
+            code_info = self._get_or_create_code_info(code)
+            self._frame_states[frame] = DebuggerState(code_info=code_info, frame=frame)
+        return self._frame_states[frame]
+
+    def _find_frame_for_code(self, code: types.CodeType) -> types.FrameType | None:
+        """Walk the stack to find the frame executing the given code."""
+        f: types.FrameType | None = sys._getframe()
+        while f is not None:
+            if f.f_code is code:
+                return f
+            f = f.f_back
+        return None
+
+    def _build_tracked_frame_stack(self, state: DebuggerState) -> list[DebuggerState]:
+        """Build list of tracked ancestor frames, outermost-first.
+
+        Walks f_back from the current execution frame to find all ancestor
+        frames whose code is tracked by the debugger. Returns the corresponding
+        DebuggerState objects with the execution state (state) last.
+        """
+        ancestors: list[DebuggerState] = []
+        f: types.FrameType | None = state.frame.f_back
+        while f is not None:
+            ancestor_state = self._frame_states.get(f)
+            if ancestor_state is not None:
+                ancestors.append(ancestor_state)
+            f = f.f_back
+        ancestors.reverse()
+        ancestors.append(state)
+        return ancestors
+
+    def _format_instruction(
+        self, state: DebuggerState, offset: int, mark_current: bool = True
+    ) -> str:
+        """Format an instruction for display."""
+        inst = state.offset_to_inst.get(offset)
+        iw, ow = state.index_width, state.offset_width
+        if inst is None:
+            return f"    [{offset:{ow}d}]: <unknown>"
+
+        index = state.offset_to_index.get(offset, -1)
+        marker = ">>>" if mark_current and offset == state.current_offset else "   "
+        bp_marker = (
+            "*" if state.offset_to_index.get(offset, -1) in state.breakpoints else " "
+        )
+        arg_str = f" {inst.argval}" if inst.arg is not None else ""
+        return f"{marker}{bp_marker} {index:{iw}d} [{offset:{ow}d}]: {inst.opname}{arg_str}"
+
+    def _format_header(self, state: DebuggerState) -> str:
+        """Format the header line for instruction listings."""
+        iw, ow = state.index_width, state.offset_width
+        return f"     {'#':>{iw}} [{'offset':>{ow}}]"
+
+    def _print_context(
+        self, state: DebuggerState, before: int = 3, after: int = 3
+    ) -> None:
+        """Print instructions around the current position."""
+        current_idx = state.offset_to_index.get(state.current_offset, 0)
+        start = max(0, current_idx - before)
+        end = min(len(state.instructions), current_idx + after + 1)
+
+        print(f"\nInstruction {current_idx} at offset {state.current_offset}:")
+        print(self._format_header(state))
+        for i in range(start, end):
+            inst = state.instructions[i]
+            if inst.offset is not None:
+                print(self._format_instruction(state, inst.offset))
+        print()
+
+    def _safe_stack_depth(self, state: DebuggerState, is_current_frame: bool) -> int:
+        """Compute a safe stack depth for reading frame values.
+
+        For parent frames suspended at a CALL instruction, CPython has already
+        popped (and DECREF'd) the call arguments.  The tracked depth still
+        includes those consumed entries, so reading them would dereference
+        dangling pointers.  Adjust by the current instruction's stack_effect.
+        """
+        depth = state.current_stack_depth
+        if not is_current_frame and state.current_offset >= 0:
+            inst = state.offset_to_inst.get(state.current_offset)
+            if inst is not None:
+                try:
+                    effect = dis.stack_effect(inst.opcode, inst.arg)
+                    if effect < 0:
+                        depth = max(0, depth + effect)
+                except (ValueError, TypeError):
+                    pass
+        return depth
+
+    def _print_stack(self, state: DebuggerState, is_current_frame: bool = True) -> None:
+        """Print the current value stack."""
+        from torch._C._dynamo.eval_frame import _get_frame_value_stack_with_depth
+
+        depth = self._safe_stack_depth(state, is_current_frame)
+        try:
+            stack = _get_frame_value_stack_with_depth(state.frame, depth)
+        except Exception as e:
+            print(f"\nStack: (error reading stack: {e})")
+            return
+        print("\nStack (TOS at end):")
+        if not stack:
+            print("  (empty)")
+        else:
+            for i, value in enumerate(stack):
+                if value is NULL_STACK_VALUE:
+                    addr = "0x0"
+                else:
+                    addr = f"0x{id(value):x}"
+                print(f"  [{i}] {addr} {value!r}")
+        print()
+
+    def _print_locals(self, state: DebuggerState) -> None:
+        """Print local variables."""
+        print("\nLocals:")
+        locals_dict = state.frame.f_locals
+        if not locals_dict:
+            print("  (none)")
+        else:
+            for name, value in locals_dict.items():
+                print(f"  {name} = {value!r}")
+        print()
+
+    def _print_globals(self, state: DebuggerState, pattern: str | None = None) -> None:
+        """Print global variables."""
+        print("\nGlobals:")
+        for name, value in state.frame.f_globals.items():
+            if pattern and pattern not in name:
+                continue
+            if name.startswith("__") and name.endswith("__"):
+                continue
+            if isinstance(value, types.ModuleType):
+                continue
+            print(f"  {name} = {value!r}")
+        print()
+
+    def _disassemble(self, state: DebuggerState) -> None:
+        """Print the full disassembly."""
+        print(f"\nDisassembly ({len(state.instructions)} instructions):")
+        print(self._format_header(state))
+        for inst in state.instructions:
+            if inst.offset is not None:
+                print(self._format_instruction(state, inst.offset, mark_current=False))
+        print()
+
+    def _print_help(self) -> None:
+        """Print help message."""
+        print("\nCommands:")
+        print("  s [n]       - Step n instructions (default 1), stepping into calls")
+        print("  n [n]       - Step n instructions (default 1), stepping over calls")
+        print("  c, cont     - Continue until breakpoint or exception")
+        print("  r           - Continue until current function's return instruction")
+        print(
+            "  v, verbose  - Toggle verbose mode (print each instruction before executing)"
+        )
+        print("                Use with 'c' to find segfaults - last printed = culprit")
+        print("  u [count]   - Move up the call stack (toward caller) for inspection")
+        print("  d [count]   - Move down the call stack (toward callee)")
+        print("  bt, w       - Print tracked frame backtrace")
+        print("  p <expr>    - Print expression")
+        print("  l [first[, last]]  - List instructions (like pdb)")
+        print("  l .         - List around current instruction")
+        print("  ll          - Disassemble all bytecode")
+        print("  stack       - Print value stack")
+        print("  locals      - Print local variables")
+        print("  globals     - Print global variables")
+        print("  b <n>       - Set breakpoint at instruction n (see # column)")
+        print("  b           - List breakpoints")
+        print("  cl <n>      - Clear breakpoint at instruction n")
+        print("  q, quit     - Exit debugger")
+        print("  <expr>      - Evaluate Python expression (like pdb)")
+        print()
+        print("Special variables: __stack__ (list of stack values, TOS at end)")
+        print("Note: Debugger stops on exceptions and shows the failing instruction.")
+        print()
+
+    def _get_stack_for_eval(
+        self, state: DebuggerState, is_current_frame: bool = True
+    ) -> list[Any]:
+        """Get the current stack for use in expression evaluation."""
+        from torch._C._dynamo.eval_frame import _get_frame_value_stack_with_depth
+
+        depth = self._safe_stack_depth(state, is_current_frame)
+        try:
+            return _get_frame_value_stack_with_depth(state.frame, depth)
+        except Exception:
+            return []
+
+    def _build_eval_locals(
+        self, state: DebuggerState, is_current_frame: bool = True
+    ) -> dict[str, Any]:
+        """Build locals dict for expression evaluation, including user-defined variables."""
+        eval_locals = dict(state.frame.f_locals)
+        eval_locals.update(state.user_locals)
+        eval_locals["__stack__"] = self._get_stack_for_eval(state, is_current_frame)
+        return eval_locals
+
+    def _interactive_prompt(self, state: DebuggerState) -> None:
+        """Interactive prompt during debugging."""
+        frame_stack = self._build_tracked_frame_stack(state)
+        view_index = len(frame_stack) - 1
+        view_state = frame_stack[view_index]
+
+        self._print_context(view_state)
+
+        while True:
+            try:
+                cmd = input("(bdb) ").strip()
+            except (EOFError, KeyboardInterrupt):
+                print("\nExiting debugger.")
+                self._quitting = True
+                raise KeyboardInterrupt from None
+
+            if not cmd:
+                cmd = state.last_command
+            else:
+                state.last_command = cmd
+
+            parts = cmd.split(maxsplit=1)
+            action = parts[0].lower()
+            arg = parts[1] if len(parts) > 1 else ""
+
+            if action in ("s", "step"):
+                state.step_mode = True
+                self._stop_at_new_code = True
+                # Parse optional count argument (e.g., "s 3" to step 3 times)
+                if arg:
+                    try:
+                        count = int(arg)
+                        state.step_count = (
+                            count - 1
+                        )  # -1 because we're about to execute one
+                    except ValueError:
+                        print(f"Invalid count: {arg}")
+                        continue
+                else:
+                    state.step_count = 0
+                return
+
+            elif action in ("n", "next"):
+                self._next_in_frame = view_state.frame
+                if arg:
+                    try:
+                        count = int(arg)
+                        self._next_count = count - 1
+                    except ValueError:
+                        print(f"Invalid count: {arg}")
+                        continue
+                else:
+                    self._next_count = 0
+                state.step_mode = False
+                self._stop_at_new_code = False
+                return
+
+            elif action in ("c", "cont", "continue"):
+                state.step_mode = False
+                self._stop_at_new_code = False
+                self._stop_after_return = False
+                return
+
+            elif action in ("r", "return"):
+                self._return_from_frame = state.frame
+                state.step_mode = False
+                self._stop_at_new_code = False
+                return
+
+            elif action in ("u", "up"):
+                count = 1
+                if arg:
+                    try:
+                        count = int(arg)
+                    except ValueError:
+                        print(f"Invalid count: {arg}")
+                        continue
+                new_index = view_index - count
+                if new_index < 0:
+                    new_index = 0
+                    print("Oldest tracked frame")
+                view_index = new_index
+                view_state = frame_stack[view_index]
+                print(f"> {view_state.code.co_name}")
+                self._print_context(view_state)
+
+            elif action in ("d", "down"):
+                count = 1
+                if arg:
+                    try:
+                        count = int(arg)
+                    except ValueError:
+                        print(f"Invalid count: {arg}")
+                        continue
+                max_index = len(frame_stack) - 1
+                new_index = view_index + count
+                if new_index > max_index:
+                    new_index = max_index
+                    print("Newest tracked frame")
+                view_index = new_index
+                view_state = frame_stack[view_index]
+                print(f"> {view_state.code.co_name}")
+                self._print_context(view_state)
+
+            elif action in ("bt", "w", "where"):
+                for i, fs in enumerate(frame_stack):
+                    marker = ">" if i == view_index else " "
+                    idx = fs.offset_to_index.get(fs.current_offset, -1)
+                    inst = fs.offset_to_inst.get(fs.current_offset)
+                    inst_str = inst.opname if inst else "<unknown>"
+                    print(
+                        f"  {marker} [{i}] {fs.code.co_name} instruction {idx} ({inst_str})"
+                    )
+
+            elif action in ("v", "verbose"):
+                self._verbose = not self._verbose
+                status = "enabled" if self._verbose else "disabled"
+                print(f"Verbose mode {status}.")
+
+            elif action in ("l", "list"):
+                # Like pdb:
+                # - 'l' continues from last position (or centers on current if first time)
+                # - 'l .' centers on current instruction
+                # - 'l N' lists 11 instructions starting at N
+                # - 'l first, last' lists range (if last < first, last is a count)
+                num_lines = 11  # pdb default
+                start: int | None = None
+                end: int | None = None
+
+                if arg == ".":
+                    # Reset to current instruction
+                    view_state.list_index = None
+                elif arg:
+                    # Check for range syntax: "first, last" or "first,last"
+                    if "," in arg:
+                        parts = arg.split(",", 1)
+                        try:
+                            first = int(parts[0].strip())
+                            second = int(parts[1].strip())
+                            if second < first:
+                                # second is a count
+                                start = first
+                                end = first + second
+                            else:
+                                start = first
+                                end = second + 1  # inclusive
+                        except ValueError:
+                            print(f"Invalid range: {arg}")
+                            continue
+                    else:
+                        # Single number: start at that instruction
+                        try:
+                            start = int(arg)
+                            end = start + num_lines
+                        except ValueError:
+                            print(f"Invalid argument: {arg}")
+                            continue
+
+                if start is None:
+                    if view_state.list_index is None:
+                        # Center on current instruction
+                        current_idx = view_state.offset_to_index.get(
+                            view_state.current_offset, 0
+                        )
+                        start = max(0, current_idx - num_lines // 2)
+                    else:
+                        start = view_state.list_index
+                    end = start + num_lines
+
+                assert end is not None
+                end = min(len(view_state.instructions), end)
+                if start >= len(view_state.instructions):
+                    print("(End of bytecode)")
+                else:
+                    print(self._format_header(view_state))
+                    for i in range(start, end):
+                        inst = view_state.instructions[i]
+                        if inst.offset is not None:
+                            print(self._format_instruction(view_state, inst.offset))
+                    print()
+                    view_state.list_index = end
+
+            elif action == "ll":
+                self._disassemble(view_state)
+
+            elif action == "locals":
+                self._print_locals(view_state)
+
+            elif action == "globals":
+                self._print_globals(view_state, arg if arg else None)
+
+            elif action == "stack":
+                is_current = view_state is frame_stack[-1]
+                self._print_stack(view_state, is_current_frame=is_current)
+
+            elif action == "b":
+                if arg:
+                    try:
+                        index = int(arg)
+                        num_instructions = len(view_state.instructions)
+                        if index < 0 or index >= num_instructions:
+                            print(
+                                f"Invalid instruction number: {index} "
+                                f"(must be 0-{num_instructions - 1})"
+                            )
+                        else:
+                            view_state.breakpoints.add(index)
+                            print(f"Breakpoint set at instruction {index}")
+                    except ValueError:
+                        print(f"Invalid instruction number: {arg}")
+                else:
+                    print("\nBreakpoints:")
+                    if not view_state.breakpoints:
+                        print("  (none)")
+                    else:
+                        for bp_index in sorted(view_state.breakpoints):
+                            if bp_index < len(view_state.instructions):
+                                inst = view_state.instructions[bp_index]
+                                if inst.offset is not None:
+                                    print(
+                                        f"  {self._format_instruction(view_state, inst.offset, mark_current=False)}"
+                                    )
+                    print()
+
+            elif action == "cl":
+                if arg:
+                    try:
+                        index = int(arg)
+                        num_instructions = len(view_state.instructions)
+                        if index < 0 or index >= num_instructions:
+                            print(
+                                f"Invalid instruction number: {index} "
+                                f"(must be 0-{num_instructions - 1})"
+                            )
+                        elif index in view_state.breakpoints:
+                            view_state.breakpoints.discard(index)
+                            print(f"Breakpoint cleared at instruction {index}")
+                        else:
+                            print(f"No breakpoint at instruction {index}")
+                    except ValueError:
+                        print(f"Invalid instruction number: {arg}")
+
+            elif action == "p":
+                if arg:
+                    try:
+                        is_current = view_state is frame_stack[-1]
+                        frame_globals = view_state.frame.f_globals
+                        eval_locals = self._build_eval_locals(
+                            view_state, is_current_frame=is_current
+                        )
+                        result = eval(arg, frame_globals, eval_locals)
+                        print(f"{arg} = {result!r}")
+                    except Exception as e:
+                        print(f"Error evaluating '{arg}': {e}")
+                else:
+                    print("Usage: p <expression>")
+
+            elif action in ("h", "help", "?"):
+                self._print_help()
+
+            elif action in ("q", "quit", "exit"):
+                print("Exiting debugger.")
+                self._quitting = True
+                raise KeyboardInterrupt
+
+            else:
+                # Try to execute as Python code (like pdb)
+                is_current = view_state is frame_stack[-1]
+                frame_globals = view_state.frame.f_globals
+                eval_locals = self._build_eval_locals(
+                    view_state, is_current_frame=is_current
+                )
+
+                try:
+                    # First try as expression (eval)
+                    result = eval(cmd, frame_globals, eval_locals)
+                    if result is not None:
+                        print(repr(result))
+                except SyntaxError:
+                    # If not a valid expression, try as statement (exec)
+                    try:
+                        keys_before = set(eval_locals.keys())
+                        ids_before = {k: id(v) for k, v in eval_locals.items()}
+                        exec(cmd, frame_globals, eval_locals)
+                        # Capture any new or modified variables into user_locals
+                        for key in eval_locals:
+                            is_new = key not in keys_before
+                            is_user_var = key in view_state.user_locals
+                            is_modified = id(eval_locals[key]) != ids_before.get(key)
+                            if (
+                                is_new or is_user_var or is_modified
+                            ) and key != "__stack__":
+                                view_state.user_locals[key] = eval_locals[key]
+                    except Exception as e:
+                        print(f"*** {type(e).__name__}: {e}")
+                except Exception as e:
+                    print(f"*** {type(e).__name__}: {e}")
+
+    def _handle_instruction(
+        self, code: types.CodeType, offset: int, frame: types.FrameType | None = None
+    ) -> None:
+        """Common instruction handling logic for both tracing backends."""
+        if frame is None:
+            frame = self._find_frame_for_code(code)
+            assert frame is not None
+        state = self._get_or_create_frame_state(code, frame)
+
+        # Update stack depth based on the previous instruction's effect.
+        # The callback is called BEFORE instruction at 'offset' executes,
+        # so we need to apply the effect of the previous instruction first.
+        previous_offset = state.current_offset
+        if previous_offset >= 0:
+            prev_inst = state.offset_to_inst.get(previous_offset)
+            if prev_inst is not None:
+                # Detect if a jump was taken: current offset != sequential next
+                prev_index = state.offset_to_index.get(previous_offset, -1)
+                expected_next_index = prev_index + 1
+                current_index = state.offset_to_index.get(offset, -1)
+                did_jump = current_index != expected_next_index
+
+                try:
+                    effect = dis.stack_effect(
+                        prev_inst.opcode, prev_inst.arg, jump=did_jump
+                    )
+                    state.current_stack_depth += effect
+                    if state.current_stack_depth < 0:
+                        raise RuntimeError(
+                            f"Stack depth went negative after {prev_inst.opname} "
+                            f"at offset {previous_offset}: depth={state.current_stack_depth}"
+                        )
+                except (ValueError, TypeError):
+                    pass
+
+        state.current_offset = offset
+        state.list_index = None  # Reset list position when stepping
+
+        # After 'r' command completes, resume stepping at the next instruction
+        if self._stop_after_return:
+            self._stop_after_return = False
+            state.step_mode = True
+            state.step_count = 0
+            self._stop_at_new_code = True
+
+        # After 'n', resume stepping when we reach that frame
+        if self._next_in_frame is not None and frame is self._next_in_frame:
+            if self._next_count > 0:
+                # More steps remaining — re-arm for the next instruction in this frame
+                self._next_count -= 1
+            else:
+                self._next_in_frame = None
+                state.step_mode = True
+                state.step_count = 0
+                self._stop_at_new_code = True
+            if _HAS_SYS_MONITORING:
+                sys.monitoring.restart_events()
+
+        # First instruction - print header and enter step mode (unless suppressed)
+        if not state.first_instruction_seen:
+            state.first_instruction_seen = True
+            if not self._stop_at_new_code:
+                state.step_mode = False
+            if state.step_mode:
+                print(f"\n=== Entering Dynamo-generated code: {code.co_name} ===")
+                self._print_help()
+            elif self._verbose:
+                print(f"\n=== Entering Dynamo-generated code: {code.co_name} ===")
+
+        # Verbose mode: print each instruction before executing (for segfault debugging)
+        if self._verbose:
+            inst = state.offset_to_inst.get(offset)
+            if inst:
+                idx = state.offset_to_index.get(offset, -1)
+                arg_str = f" {inst.argval}" if inst.arg is not None else ""
+                print(f"Running [{idx}] {inst.opname}{arg_str}", flush=True)
+
+        inst = state.offset_to_inst.get(offset)
+
+        # Check if current instruction has a breakpoint (by index)
+        current_index = state.offset_to_index.get(offset, -1)
+        hit_breakpoint = current_index in state.breakpoints
+
+        # 'r' command: stop at RETURN_VALUE/RETURN_CONST in the target frame
+        # so the user can inspect the return value on the stack before returning.
+        hit_return_target = False
+        if self._return_from_frame is not None and frame is self._return_from_frame:
+            if inst is not None and inst.opname in ("RETURN_VALUE", "RETURN_CONST"):
+                self._return_from_frame = None
+                # After stepping past the return, stop in the caller
+                self._stop_after_return = True
+                hit_return_target = True
+
+        # Check if we should stop
+        # If step_count > 0, we're in the middle of "N s" and should continue
+        if state.step_count > 0:
+            state.step_count -= 1
+            # But still stop for breakpoints and return targets
+            if hit_breakpoint or hit_return_target:
+                state.step_count = 0  # Cancel remaining steps
+            else:
+                return  # Continue without prompting
+
+        should_stop = state.step_mode or hit_breakpoint or hit_return_target
+        if should_stop:
+            if hit_breakpoint:
+                if inst is not None and self.is_programmatic_breakpoint(inst):
+                    print("Breakpoint hit (programmatic)")
+                else:
+                    print(f"Breakpoint hit at instruction {current_index}")
+            elif hit_return_target:
+                print(f"About to return from {code.co_name}")
+            self._interactive_prompt(state)
+
+    def _handle_return(
+        self, code: types.CodeType, retval: object, frame: types.FrameType | None = None
+    ) -> None:
+        """Common return handling logic."""
+        if self._quitting:
+            return
+        print(f"\n=== {code.co_name} returned: {retval!r} ===")
+
+        # For sys.monitoring, we don't get the frame directly, but since
+        # the frame that's currently returning is the one whose PY_RETURN just
+        # fired, code identity is sufficient to match.
+        if self._return_from_frame is not None and (
+            frame is self._return_from_frame
+            or (frame is None and code is self._return_from_frame.f_code)
+        ):
+            self._return_from_frame = None
+            self._stop_after_return = True
+            if _HAS_SYS_MONITORING:
+                sys.monitoring.restart_events()
+        if self._next_in_frame is not None and (
+            frame is self._next_in_frame
+            or (frame is None and code is self._next_in_frame.f_code)
+        ):
+            self._next_in_frame = None
+            self._stop_after_return = True
+            if _HAS_SYS_MONITORING:
+                sys.monitoring.restart_events()
+
+        # Clean up frame state to avoid holding dead frame references
+        if frame is not None:
+            self._frame_states.pop(frame, None)
+        else:
+            # sys.monitoring path: find and remove matching frame state by code
+            for f, s in list(self._frame_states.items()):
+                if s.code is code:
+                    self._frame_states.pop(f, None)
+                    break
+
+    def _handle_exception(
+        self, code: types.CodeType, offset: int, exception: BaseException
+    ) -> None:
+        """Common exception handling logic."""
+        if self._quitting:
+            return
+        frame = self._find_frame_for_code(code)
+        if frame is None:
+            return
+        state = self._frame_states.get(frame)
+        if state is None:
+            return
+
+        # If we're waiting for this frame to return, the return was interrupted
+        if self._return_from_frame is not None and (
+            frame is self._return_from_frame or code is self._return_from_frame.f_code
+        ):
+            self._return_from_frame = None
+
+        # If we're waiting for next-in-frame, the execution was interrupted
+        if self._next_in_frame is not None and (
+            frame is self._next_in_frame or code is self._next_in_frame.f_code
+        ):
+            self._next_in_frame = None
+            self._stop_at_new_code = True
+
+        state.current_offset = offset
+
+        inst = state.offset_to_inst.get(offset)
+        inst_str = inst.opname if inst else "<unknown>"
+        current_index = state.offset_to_index.get(offset, -1)
+
+        print(f"\n=== Exception raised at instruction {current_index}: {inst_str} ===")
+        print(f"=== {type(exception).__name__}: {exception} ===")
+        self._interactive_prompt(state)
+
+    # =========================================================================
+    # Python 3.12+ implementation using sys.monitoring
+    # =========================================================================
+
+    def _dynamo_code_callback(self, code: types.CodeType) -> None:
+        """Called by C++ before executing Dynamo-generated code."""
+        if code in self._tracked_codes:
+            return
+        self._tracked_codes.add(code)
+
+        if _HAS_SYS_MONITORING:
+            # Enable INSTRUCTION and PY_RETURN events for this specific code object
+            # RAISE must be a global event (cannot be set locally)
+            sys.monitoring.set_local_events(
+                self._tool_id,
+                code,
+                sys.monitoring.events.INSTRUCTION | sys.monitoring.events.PY_RETURN,
+            )
+        # For settrace, we enable opcode tracing when we see the frame
+
+        # Pre-create code info for this code
+        self._get_or_create_code_info(code)
+
+    def _monitoring_return_callback(
+        self, code: types.CodeType, instruction_offset: int, retval: object
+    ) -> None:
+        """Callback for PY_RETURN events (sys.monitoring)."""
+        self._handle_return(code, retval)
+
+    def _monitoring_instruction_callback(
+        self, code: types.CodeType, offset: int
+    ) -> object:
+        """Callback for INSTRUCTION events (sys.monitoring)."""
+        self._handle_instruction(code, offset)
+        return sys.monitoring.DISABLE
+
+    def _monitoring_raise_callback(
+        self, code: types.CodeType, offset: int, exception: BaseException
+    ) -> object:
+        """Callback for RAISE events (sys.monitoring)."""
+        # Only handle exceptions from tracked Dynamo-generated code
+        if code not in self._tracked_codes:
+            return None
+        self._handle_exception(code, offset, exception)
+        # Cannot return DISABLE for global events like RAISE
+        return None
+
+    # =========================================================================
+    # Python 3.11 and below implementation using sys.settrace
+    # =========================================================================
+
+    def _settrace_callback(
+        self, frame: types.FrameType, event: str, arg: Any
+    ) -> Callable[..., Any] | None:
+        """Trace function for sys.settrace."""
+        code = frame.f_code
+
+        # Only trace Dynamo-generated code
+        if code not in self._tracked_codes:
+            return self._settrace_callback
+
+        if event == "call":
+            # Enable opcode tracing for this frame
+            frame.f_trace_opcodes = True
+            return self._settrace_callback
+
+        elif event == "opcode":
+            # Get the current instruction offset
+            offset = frame.f_lasti
+            self._handle_instruction(code, offset, frame)
+            return self._settrace_callback
+
+        elif event == "return":
+            self._handle_return(code, arg, frame)
+            return self._settrace_callback
+
+        elif event == "exception":
+            # arg is (exception_type, exception_value, traceback)
+            offset = frame.f_lasti
+            exc_type, exc_value, exc_tb = arg
+            self._handle_exception(code, offset, exc_value)
+            return self._settrace_callback
+
+        return self._settrace_callback
+
+    # =========================================================================
+    # Context manager implementation
+    # =========================================================================
+
+    def __enter__(self) -> Self:
+        """Start the debug context."""
+        from torch._C._dynamo.eval_frame import (
+            get_bytecode_debugger_callback,
+            set_bytecode_debugger_callback,
+        )
+
+        self._active = True
+        self._prev_callback = get_bytecode_debugger_callback()
+        self._code_info.clear()
+        self._frame_states.clear()
+        self._tracked_codes.clear()
+
+        if _HAS_SYS_MONITORING:
+            # Python 3.12+: Use sys.monitoring
+            try:
+                sys.monitoring.use_tool_id(self._tool_id, "bytecode_debugger")
+            except ValueError:
+                pass
+
+            # Register callbacks (events enabled per-code via set_local_events)
+            sys.monitoring.register_callback(
+                self._tool_id,
+                sys.monitoring.events.INSTRUCTION,
+                self._monitoring_instruction_callback,
+            )
+            sys.monitoring.register_callback(
+                self._tool_id,
+                sys.monitoring.events.PY_RETURN,
+                self._monitoring_return_callback,
+            )
+            sys.monitoring.register_callback(
+                self._tool_id,
+                sys.monitoring.events.RAISE,
+                self._monitoring_raise_callback,
+            )
+            # RAISE must be a global event (cannot be local)
+            sys.monitoring.set_events(self._tool_id, sys.monitoring.events.RAISE)
+        else:
+            # Python 3.11 and below: Use sys.settrace
+            self._old_trace = sys.gettrace()
+            sys.settrace(self._settrace_callback)
+
+        # Set the C callback that will be called before executing Dynamo code
+        set_bytecode_debugger_callback(self._dynamo_code_callback)
+        return self
+
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: types.TracebackType | None,
+    ) -> bool:
+        """End the debug context."""
+        if not self._active:
+            return False
+        from torch._C._dynamo.eval_frame import set_bytecode_debugger_callback
+
+        self._active = False
+        prev = self._prev_callback
+        set_bytecode_debugger_callback(prev)
+
+        if _HAS_SYS_MONITORING:
+            if prev is None:
+                sys.monitoring.set_events(self._tool_id, 0)
+                try:
+                    sys.monitoring.free_tool_id(self._tool_id)
+                except ValueError:
+                    pass
+            else:
+                # Restore outer context's monitoring callbacks and events
+                assert isinstance(prev, types.MethodType)
+                outer = cast("_DebugContext", prev.__self__)
+                sys.monitoring.register_callback(
+                    self._tool_id,
+                    sys.monitoring.events.INSTRUCTION,
+                    outer._monitoring_instruction_callback,
+                )
+                sys.monitoring.register_callback(
+                    self._tool_id,
+                    sys.monitoring.events.PY_RETURN,
+                    outer._monitoring_return_callback,
+                )
+                sys.monitoring.register_callback(
+                    self._tool_id,
+                    sys.monitoring.events.RAISE,
+                    outer._monitoring_raise_callback,
+                )
+                sys.monitoring.set_events(self._tool_id, sys.monitoring.events.RAISE)
+                for code in outer._tracked_codes:
+                    sys.monitoring.set_local_events(
+                        self._tool_id,
+                        code,
+                        sys.monitoring.events.INSTRUCTION
+                        | sys.monitoring.events.PY_RETURN,
+                    )
+        else:
+            # Python 3.11 and below: Restore old trace
+            sys.settrace(self._old_trace)
+
+        return False  # Don't suppress exceptions
+
+
+@contextmanager
+def debug() -> Generator[_DebugContext, None, None]:
+    """
+    Context manager for debugging Dynamo-generated bytecode.
+
+    Any Dynamo-generated code executed within this context will trigger
+    the interactive bytecode debugger.
+
+    Example:
+        >>> # xdoctest: +SKIP
+        >>> import torch
+        >>>
+        >>> @torch.compile
+        >>> def my_fn(x):
+        ...     return x + 1
+        >>>
+        >>> with torch._dynamo.bytecode_debugger.debug():
+        ...     my_fn(torch.randn(3))
+    """
+    ctx = _DebugContext()
+    try:
+        with ctx:
+            yield ctx
+    except KeyboardInterrupt:
+        print("\n=== Debug session ended ===")
+
+
+def breakpoint() -> None:
+    """Programmatic breakpoint for user code.
+
+    Place this in code compiled by Dynamo. During tracing, Dynamo inserts a
+    BREAKPOINT_MARKER into the compiled bytecode. When the bytecode debugger
+    is active, execution pauses at this marker.
+    """
diff --git a/torch/_dynamo/bytecode_transformation.py b/torch/_dynamo/bytecode_transformation.py
index 8eab797523adb..b1c5551f106a9 100644
--- a/torch/_dynamo/bytecode_transformation.py
+++ b/torch/_dynamo/bytecode_transformation.py
@@ -22,7 +22,7 @@
 import types
 import uuid
 from collections.abc import Callable, Iterable, Iterator, Mapping, Sequence
-from typing import Any, cast, Optional, TYPE_CHECKING, Union
+from typing import Any, cast, Optional, TYPE_CHECKING
 
 from . import config
 from .bytecode_analysis import (
@@ -72,17 +72,17 @@ class Instruction:
 
     opcode: int
     opname: str
-    arg: Optional[int]
+    arg: int | None
     argval: Any
-    offset: Optional[int] = None
-    starts_line: Optional[int] = None
+    offset: int | None = None
+    starts_line: int | None = None
     is_jump_target: bool = False
 
     positions: Optional["dis.Positions"] = None
     # extra fields to make modification easier:
     target: Optional["Instruction"] = None
-    exn_tab_entry: Optional[InstructionExnTabEntry] = None
-    argrepr: Optional[str] = None
+    exn_tab_entry: InstructionExnTabEntry | None = None
+    argrepr: str | None = None
 
     def __hash__(self) -> int:
         return id(self)
@@ -165,9 +165,9 @@ def inst_has_op_bits(name: str) -> bool:
 def create_instruction(
     name: str,
     *,
-    arg: Optional[int] = None,
-    argval: Optional[Any] = _NotProvided,
-    target: Optional[Instruction] = None,
+    arg: int | None = None,
+    argval: Any | None = _NotProvided,
+    target: Instruction | None = None,
 ) -> Instruction:
     """
     At most one of `arg`, `argval`, and `target` can be not None/_NotProvided.
@@ -225,6 +225,24 @@ def create_load_const(val: Any, checked: bool = True) -> Instruction:
     return create_instruction("LOAD_CONST", argval=val)
 
 
+def create_breakpoint() -> list[Instruction]:
+    """
+    Create instructions that trigger the bytecode debugger to stop.
+
+    Usage:
+        codegen.extend_output(create_breakpoint())
+
+    When the bytecode debugger is active, execution will pause at this point.
+    At runtime, these instructions have no effect on program state.
+    """
+    from .bytecode_debugger import BREAKPOINT_MARKER
+
+    return [
+        create_instruction("LOAD_CONST", argval=BREAKPOINT_MARKER),
+        create_instruction("POP_TOP"),
+    ]
+
+
 def create_dup_top() -> Instruction:
     if sys.version_info >= (3, 11):
         return create_instruction("COPY", arg=1)
@@ -255,7 +273,7 @@ def create_rot_n(n: int) -> list[Instruction]:
 
 
 def add_push_null(
-    inst_or_insts: Union[Instruction, list[Instruction]],
+    inst_or_insts: Instruction | list[Instruction],
 ) -> list[Instruction]:
     """
     Appends or prepends a PUSH_NULL instruction to `inst_or_insts`,
@@ -314,7 +332,7 @@ def set_inst_bit(idx: int) -> None:
 
 
 def add_push_null_call_function_ex(
-    inst_or_insts: Union[Instruction, list[Instruction]],
+    inst_or_insts: Instruction | list[Instruction],
 ) -> list[Instruction]:
     """Like add_push_null, but the low bit of LOAD_ATTR/LOAD_SUPER_ATTR
     is not set, due to an expected CALL_FUNCTION_EX instruction.
@@ -502,7 +520,7 @@ def create_swap(n: int) -> list[Instruction]:
 
 
 def create_binary_slice(
-    start: Optional[int], end: Optional[int], store: bool = False
+    start: int | None, end: int | None, store: bool = False
 ) -> list[Instruction]:
     """
     BINARY_SLICE and STORE_SLICE (if `set` is True) for all Python versions
@@ -830,7 +848,7 @@ def assemble(instructions: list[Instruction], firstlineno: int) -> tuple[bytes,
 
 def _get_instruction_by_offset(
     offset_to_inst: dict[int, Instruction], offset: int
-) -> Optional[Instruction]:
+) -> Instruction | None:
     """
     Get the instruction located at a given offset, accounting for EXTENDED_ARGs
     """
@@ -1428,7 +1446,7 @@ def get_const_index(code_options: dict[str, Any], val: Any) -> int:
 def fix_vars(
     instructions: list[Instruction],
     code_options: dict[str, Any],
-    varname_from_oparg: Optional[Callable[..., Any]] = None,
+    varname_from_oparg: Callable[..., Any] | None = None,
 ) -> None:
     # compute instruction arg from argval if arg is not provided
     names = {name: idx for idx, name in enumerate(code_options["co_names"])}
@@ -1751,7 +1769,7 @@ def is_generator(code: types.CodeType) -> bool:
 
 def bytecode_from_template(
     fn: Callable[..., Any],
-    varname_map: Optional[Mapping[Any, Any]] = None,
+    varname_map: Mapping[Any, Any] | None = None,
     noreturn: bool = True,
     noprefix: bool = True,
 ) -> list[Instruction]:
diff --git a/torch/_dynamo/cache_size.py b/torch/_dynamo/cache_size.py
index d1a46742f37ac..265643953aa35 100644
--- a/torch/_dynamo/cache_size.py
+++ b/torch/_dynamo/cache_size.py
@@ -1,7 +1,7 @@
 import logging
 import weakref
 from dataclasses import dataclass
-from typing import Any, Optional
+from typing import Any
 
 from torch._guards import CompileId
 
@@ -101,7 +101,7 @@ def will_compilation_exceed_specific_limit(self, limit: int) -> bool:
 
 def _get_weakref_from_f_locals(
     frame: DynamoFrameType, local_name: str
-) -> Optional[weakref.ref[Any]]:
+) -> weakref.ref[Any] | None:
     obj = frame.f_locals.get(local_name, None)
     weak_id = None
     try:
diff --git a/torch/_dynamo/codegen.py b/torch/_dynamo/codegen.py
index ace7c40180405..3e0718a6fa28d 100644
--- a/torch/_dynamo/codegen.py
+++ b/torch/_dynamo/codegen.py
@@ -17,7 +17,7 @@
 import types
 from collections import Counter, deque
 from collections.abc import Callable, Iterable
-from typing import Any, Optional, TYPE_CHECKING, Union
+from typing import Any, TYPE_CHECKING, Union
 
 import torch.nn
 from torch.utils._ordered_set import OrderedSet
@@ -47,6 +47,7 @@
     LocalGeneratorObjectVariable,
 )
 from .variables.nn_module import NNModuleVariable
+from .variables.script_object import TorchScriptObjectVariable
 from .variables.tensor import (
     NumpyNdarrayVariable,
     SymNodeVariable,
@@ -76,21 +77,21 @@ class PyCodegen:
     def __init__(
         self,
         tx: "InstructionTranslatorBase",
-        root: Optional[torch.nn.Module] = None,
-        graph_output_var: Optional[str] = None,
-        tempvars: Optional[dict[Union[VariableTracker, Source], Any]] = None,
-        overridden_sources: Optional[dict[Source, Source]] = None,
+        root: torch.nn.Module | None = None,
+        graph_output_var: str | None = None,
+        tempvars: dict[VariableTracker | Source, Any] | None = None,
+        overridden_sources: dict[Source, Source] | None = None,
     ) -> None:
         self.root = root
-        self.top_of_stack: Optional[Union[VariableTracker, Source]] = None
-        self.uses: Counter[Union[VariableTracker, Source]] = collections.Counter()
+        self.top_of_stack: VariableTracker | Source | None = None
+        self.uses: Counter[VariableTracker | Source] = collections.Counter()
         self.graph_outputs: dict[int, GraphOutputEntry] = {}
         self._output: list[Instruction] = []
         # This determines which VariableTracker/Source should be stored as
         # locals, and maps the VariableTracker/Source to the local variable
         # name. Note that it could map to None initially, in which case we'll
         # overwrite it to map to real temporary names via `add_cache`.
-        self.tempvars: dict[Union[VariableTracker, Source], Any] = tempvars or {}
+        self.tempvars: dict[VariableTracker | Source, Any] = tempvars or {}
         self.tx = tx
         self.graph_output_var = graph_output_var
         self.code_options = self.tx.output.code_options
@@ -153,7 +154,7 @@ def add_push_null(
             self.clear_tos()
 
     def __call__(
-        self, value: Union[VariableTracker, Source, None], allow_cache: bool = True
+        self, value: VariableTracker | Source | None, allow_cache: bool = True
     ) -> None:
         """
         Generate code such that top-of-stack (TOS) is set to value.
@@ -329,6 +330,7 @@ def gen_fn() -> None:
                 SymNodeVariable,
                 UnspecializedPythonVariable,
                 NumpyNdarrayVariable,
+                TorchScriptObjectVariable,
             ),
         ):
             graph_outputs_key = self.add_graph_output(value)
@@ -398,12 +400,12 @@ def load_graph_output(self, index: int) -> None:
         output.append(self.create_load_const(index))
         output.append(self.create_binary_subscr())
 
-    def add_cache(self, value: Union[VariableTracker, Source]) -> None:
+    def add_cache(self, value: VariableTracker | Source) -> None:
         var = self.new_var()
         self.tempvars[value] = var
         self._output.append(self.create_store(var))
 
-    def foreach(self, items: Iterable[Union[VariableTracker, Source]]) -> None:
+    def foreach(self, items: Iterable[VariableTracker | Source]) -> None:
         for i in items:
             self(i)
 
diff --git a/torch/_dynamo/compiled_autograd.py b/torch/_dynamo/compiled_autograd.py
index 196c9043d5731..b6d84483007a6 100644
--- a/torch/_dynamo/compiled_autograd.py
+++ b/torch/_dynamo/compiled_autograd.py
@@ -21,7 +21,7 @@
 import time
 from collections import Counter, defaultdict
 from collections.abc import Callable, Generator, Sequence
-from typing import Any, Optional, TYPE_CHECKING, Union
+from typing import Any, TYPE_CHECKING
 
 import torch
 import torch.utils._pytree as pytree
@@ -91,7 +91,7 @@ def snapshot_cudagraph_enabled() -> bool:
     return torch._inductor.config.triton.cudagraphs
 
 
-def maybe_clone(x: Optional[torch.Tensor]) -> Optional[torch.Tensor]:
+def maybe_clone(x: torch.Tensor | None) -> torch.Tensor | None:
     if x is not None:
         return clone_preserve_strides(x)
     return x
@@ -223,6 +223,7 @@ def add(
         assert not hasattr(self, name)
         result = Op(name, fn, is_custom_function)
         if is_traceable:
+            # pyrefly: ignore [deprecated]
             setattr(self, name, torch._dynamo.allow_in_graph(result))
         else:
             # C++ autograd function was not marked as traceable
@@ -295,9 +296,9 @@ def __init__(self, compiler_fn: Callable[..., Any]) -> None:
         )
         self.fx_tracer = PythonKeyTracer()
         self.proxy_mode = ProxyTorchDispatchMode(self.fx_tracer, "symbolic")
-        self.hooks_proxy: Optional[Proxy] = None
+        self.hooks_proxy: Proxy | None = None
 
-    def wrap_fake(self, x: torch.Tensor, source: Optional[Source]) -> FakeTensor:
+    def wrap_fake(self, x: torch.Tensor, source: Source | None) -> FakeTensor:
         assert isinstance(x, torch.Tensor)
         return self.fake_tensor_mode.from_tensor(x, source=source)
 
@@ -309,7 +310,7 @@ def begin_capture(
         self,
         inputs: list[torch.Tensor],
         sizes: list[int],
-        scalars: list[Union[int, float]],
+        scalars: list[IntLikeType | FloatLikeType],
         origins: list[list[tuple[int, str]]],
         accumulate_grad: bool,
         check_nans: bool,
@@ -381,7 +382,8 @@ def begin_capture(
                 (proxies[i],),
                 {},
             )
-            self.symnode_proxy_lookup[symint.node] = proxies[i]
+            if not isinstance(symint, int):
+                self.symnode_proxy_lookup[symint.node] = proxies[i]
         proxies = self.bind_objects_to_proxies(sym_sizes, proxies, sizes_origins)
 
         for idx, val in enumerate(scalars):
@@ -448,7 +450,7 @@ def proxy_call_aot_backward(
         saved_tensors: Sequence[torch.Tensor],
         pctx: Any,
         ctx: Any,
-        maybe_backward_state_idx: Optional[int],
+        maybe_backward_state_idx: int | None,
         opaque_object_indices: list[int],
     ) -> Sequence[Any]:
         # The AOTBackward call consists of three things: the prologue, the
@@ -503,6 +505,7 @@ def call_aot_bwd_prologue(
 
         pgrads = self.fx_tracer.create_proxy(
             kind="call_function",
+            # pyrefly: ignore [bad-argument-type]
             target=call_aot_bwd_prologue,
             args=(
                 psaved_tensors,
@@ -549,7 +552,7 @@ def num_inputs(graph: torch.fx.Graph) -> int:
             args_idx = 0
             # pyrefly: ignore [implicit-any]
             value_remap = {}
-            poutputs: Optional[list[torch.fx.Proxy]] = None
+            poutputs: list[torch.fx.Proxy] | None = None
 
             # names of nodes must appear only once in the fx.Graph
             # dedup AOT backwards that appear multiple times
@@ -633,6 +636,7 @@ def make_subclass(*unwrapped_args: Any) -> Any:
 
             poutput = self.fx_tracer.create_proxy(
                 kind="call_function",
+                # pyrefly: ignore [bad-argument-type]
                 target=make_subclass,
                 args=tuple(punwrapped_args),
                 kwargs={},
@@ -658,9 +662,9 @@ def proxy_call_backward(
         saved_tensors: Sequence[torch.Tensor],
         backward_idx: int,
         ctx: torch.autograd.function.BackwardCFunction,
-        maybe_backward_state_idx: Optional[int],
+        maybe_backward_state_idx: int | None,
         opaque_object_indices: list[int],
-    ) -> tuple[Optional[torch.Tensor], ...]:
+    ) -> tuple[torch.Tensor | None, ...]:
         assert self.hooks_proxy is not None
         pctx = self.hooks_proxy[backward_idx]  # type: ignore[index]
         pinputs = self.to_proxy(inputs)
@@ -691,7 +695,7 @@ def proxy_call_backward(
 
         with disable_proxy_modes_tracing():
             # create fake Tensors
-            grad_ins: list[Optional[torch.Tensor]] = []
+            grad_ins: list[torch.Tensor | None] = []
             for idx, output_metadata in enumerate(output_metadatas):
                 if output_metadata is None or proxies[idx] is None:
                     grad_ins.append(None)
@@ -1416,7 +1420,7 @@ def bind_objects_to_proxies(
         self,
         objects: Sequence[Any],
         proxies: Any,
-        origins: Optional[list[tuple[int, str]]] = None,
+        origins: list[tuple[int, str]] | None = None,
     ) -> Sequence[Any]:
         if isinstance(proxies, torch.fx.Proxy):
             if origins:
@@ -1445,7 +1449,7 @@ def set_node_origin(
         self,
         node_name: str,
         nodecall_index: int,
-        pyobj: Optional[torch.autograd.Function],
+        pyobj: torch.autograd.Function | None,
     ) -> None:
         maybe_aot_id = ""
         if pyobj is not None:
@@ -1616,10 +1620,10 @@ def copy_slices_prologue(
 def copy_slices_epilogue(
     needs_input_grad: Sequence[bool],
     result: torch.Tensor,
-    res: Sequence[Optional[torch.Tensor]],
+    res: Sequence[torch.Tensor | None],
     grad_slice: torch.Tensor,
-) -> list[Optional[torch.Tensor]]:
-    grad_inputs: list[Optional[torch.Tensor]] = [None] * len(needs_input_grad)
+) -> list[torch.Tensor | None]:
+    grad_inputs: list[torch.Tensor | None] = [None] * len(needs_input_grad)
     for i in range(len(needs_input_grad)):
         if needs_input_grad[i]:
             if res[i] is None:
diff --git a/torch/_dynamo/comprehension_graph_break.py b/torch/_dynamo/comprehension_graph_break.py
new file mode 100644
index 0000000000000..d0e15cf5fa134
--- /dev/null
+++ b/torch/_dynamo/comprehension_graph_break.py
@@ -0,0 +1,694 @@
+from __future__ import annotations
+
+import copy
+import dataclasses
+import dis
+import functools
+import logging
+import sys
+from typing import Any, TYPE_CHECKING
+
+
+if TYPE_CHECKING:
+    import types
+
+    from collections.abc import Callable
+
+    from .symbolic_convert import InstructionTranslatorBase
+
+from .bytecode_transformation import (
+    create_copy,
+    create_dup_top,
+    create_instruction,
+    create_swap,
+    Instruction,
+    unique_id,
+)
+from .codegen import PyCodegen
+from .exc import unimplemented
+from .output_graph import GraphCompileReason, StackLocalsMetadata
+from .variables.misc import NullVariable, UnknownVariable
+
+
+log = logging.getLogger(__name__)
+
+
+@functools.cache
+def _get_comprehension_bytecode_prefix() -> list[str]:
+    """Get the bytecode instructions that precede BUILD_LIST in a list comprehension."""
+
+    assert sys.version_info >= (3, 12)
+
+    def fn() -> list[int]:
+        return [i for i in range(1)]  # noqa: C416
+
+    insts = [inst.opname for inst in dis.get_instructions(fn)]
+
+    start_idx = len(insts) - 1 - insts[::-1].index("LOAD_FAST_AND_CLEAR")
+    end_idx = insts.index("BUILD_LIST")
+
+    return insts[start_idx:end_idx]
+
+
+@functools.cache
+def _get_comprehension_result_patterns() -> dict[str, dict[str, Any]]:
+    """Discover bytecode patterns for comprehension result handling.
+
+    Analyzes sample functions to extract the opcode sequences that appear
+    after END_FOR for each result disposition (stored, discarded, returned, consumed).
+
+    Returns patterns with:
+        - pre_store_ops: opcodes between END_FOR and first STORE_FAST
+        - post_store_op: first opcode after all STORE_FASTs (for disambiguation)
+    """
+    assert sys.version_info >= (3, 12)
+
+    def fn_stored() -> list[int]:
+        result = [i for i in range(1)]  # noqa: C416
+        return result
+
+    def fn_discarded() -> int:
+        [i for i in range(1)]  # noqa: C416
+        return 1
+
+    def fn_returned() -> list[int]:
+        return [i for i in range(1)]  # noqa: C416
+
+    def fn_consumed() -> int:
+        return sum([i for i in range(1)])  # noqa: C416
+
+    def extract_pattern(fn: Callable[..., Any]) -> tuple[list[str], str | None]:
+        """Extract (pre_store_ops, post_store_op) from comprehension bytecode."""
+        target_line = list(dis.findlinestarts(fn.__code__))[1][1]
+        insts: list[str] = []
+        started = False
+        for instr in dis.get_instructions(fn):
+            if started and instr.starts_line:
+                break
+            pos = instr.positions
+            if pos and pos.lineno == target_line:
+                started = started or bool(instr.starts_line)
+                insts.append(instr.opname)
+
+        ops = insts[insts.index("END_FOR") + 1 :]
+        idx = 0
+
+        pre_store_ops = []
+        while idx < len(ops) and ops[idx] != "STORE_FAST":
+            pre_store_ops.append(ops[idx])
+            idx += 1
+
+        while idx < len(ops) and ops[idx] == "STORE_FAST":
+            idx += 1
+
+        return pre_store_ops, ops[idx] if idx < len(ops) else None
+
+    stored = extract_pattern(fn_stored)
+    discarded = extract_pattern(fn_discarded)
+    returned = extract_pattern(fn_returned)
+    consumed = extract_pattern(fn_consumed)
+
+    return {
+        "stored": {"pre_store_ops": stored[0], "post_store_op": stored[1]},
+        "discarded": {"pre_store_ops": discarded[0], "post_store_op": discarded[1]},
+        "returned": {"pre_store_ops": returned[0], "post_store_op": returned[1]},
+        "consumed": {"pre_store_ops": consumed[0], "post_store_op": []},
+    }
+
+
+@dataclasses.dataclass
+class ComprehensionAnalysis:
+    """Metadata about a comprehension's bytecode structure.
+
+    Attributes:
+        end_ip: Instruction pointer after all comprehension bytecode
+        result_var: Name of result variable, or None if result stays on stack
+        result_on_stack: True if result stays on stack (discarded, returned, or in expression)
+        iterator_vars: Variables from LOAD_FAST_AND_CLEAR (need restoration)
+        walrus_vars: Variables assigned via walrus operator (:=) inside comprehension
+        captured_vars: Variables read from outer scope via LOAD_FAST inside comprehension
+    """
+
+    end_ip: int
+    result_var: str | None
+    result_on_stack: bool
+    iterator_vars: list[str]
+    walrus_vars: list[str]
+    captured_vars: list[str]
+
+
+def _is_comprehension_start(tx: InstructionTranslatorBase) -> bool:
+    """Detect if we're at the start of a list/dict comprehension in 3.12+.
+
+    In Python 3.12+, comprehensions are inlined with a bytecode pattern that
+    precedes BUILD_LIST/BUILD_MAP.
+    """
+    assert sys.version_info >= (3, 12)
+
+    assert tx.instruction_pointer is not None
+    ip = tx.instruction_pointer - 1
+
+    pattern = _get_comprehension_bytecode_prefix()
+    prefix = [inst.opname for inst in tx.instructions[ip - len(pattern) : ip]]
+
+    return prefix == pattern
+
+
+def _find_comprehension_end_for_ip(tx: InstructionTranslatorBase) -> int:
+    """Find the instruction pointer of the outermost END_FOR for current comprehension."""
+    assert sys.version_info >= (3, 12)
+    assert tx.instruction_pointer is not None
+
+    nesting_depth = 0
+    for search_ip in range(tx.instruction_pointer, len(tx.instructions)):
+        inst = tx.instructions[search_ip]
+        if inst.opname == "FOR_ITER":
+            nesting_depth += 1
+        elif inst.opname == "END_FOR":
+            nesting_depth -= 1
+            if nesting_depth == 0:
+                return search_ip
+    return -1
+
+
+def _analyze_comprehension(tx: InstructionTranslatorBase) -> ComprehensionAnalysis:
+    """Analyze comprehension bytecode to determine result handling pattern."""
+    assert sys.version_info >= (3, 12)
+    assert tx.instruction_pointer is not None
+
+    patterns = _get_comprehension_result_patterns()
+    start_ip = tx.instruction_pointer - 1  # BUILD_LIST/BUILD_MAP
+
+    iterator_vars: list[str] = []
+    walrus_vars: list[str] = []
+    captured_vars: list[str] = []
+    defined_inside: set[str] = set()
+
+    # Collect iterator variables from LOAD_FAST_AND_CLEAR before BUILD_LIST/BUILD_MAP
+    iter_scan_ip = start_ip - 1
+    while iter_scan_ip >= 0:
+        inst = tx.instructions[iter_scan_ip]
+        if inst.opname == "LOAD_FAST_AND_CLEAR":
+            iterator_vars.insert(0, inst.argval)
+            iter_scan_ip -= 1
+        elif inst.opname in ("SWAP", "GET_ITER"):
+            iter_scan_ip -= 1
+        else:
+            break
+    defined_inside.update(iterator_vars)
+
+    end_for_ip = _find_comprehension_end_for_ip(tx)
+    if end_for_ip == -1:
+        unimplemented(
+            gb_type="Comprehension analysis failed: No END_FOR",
+            context="",
+            explanation="Could not find END_FOR instruction in comprehension bytecode.",
+            hints=[],
+        )
+
+    # Find first FOR_ITER to know where loop body starts
+    for_iter_ip = next(
+        i
+        for i in range(start_ip, end_for_ip)
+        if tx.instructions[i].opname == "FOR_ITER"
+    )
+
+    # Single pass through loop body to detect walrus vars and captured vars
+    for body_ip in range(for_iter_ip + 1, end_for_ip):
+        inst = tx.instructions[body_ip]
+
+        # Detect walrus pattern: COPY 1 followed by STORE_FAST
+        if inst.opname == "COPY" and inst.arg == 1 and body_ip + 1 < end_for_ip:
+            next_inst = tx.instructions[body_ip + 1]
+            if next_inst.opname == "STORE_FAST":
+                var_name = next_inst.argval
+                if var_name not in iterator_vars and var_name not in walrus_vars:
+                    walrus_vars.append(var_name)
+                    defined_inside.add(var_name)
+
+        # Track variables defined inside the loop
+        if inst.opname == "STORE_FAST":
+            defined_inside.add(inst.argval)
+
+        # Detect LOAD_FAST referencing outer variables
+        elif inst.opname.startswith("LOAD_FAST"):
+            var_names = (
+                inst.argval if isinstance(inst.argval, tuple) else (inst.argval,)
+            )
+            for var_name in var_names:
+                if var_name not in defined_inside and var_name not in captured_vars:
+                    captured_vars.append(var_name)
+
+    # Extract pre_store_ops: all opcodes from END_FOR+1 until first STORE_FAST
+    pre_store_ops: list[str] = []
+    scan_ip = end_for_ip + 1
+    while (
+        scan_ip < len(tx.instructions)
+        and tx.instructions[scan_ip].opname != "STORE_FAST"
+    ):
+        pre_store_ops.append(tx.instructions[scan_ip].opname)
+        scan_ip += 1
+
+    store_fast_ip = scan_ip
+
+    # Skip all STORE_FASTs to find post_store_op
+    while (
+        scan_ip < len(tx.instructions)
+        and tx.instructions[scan_ip].opname == "STORE_FAST"
+    ):
+        scan_ip += 1
+
+    post_store_op = (
+        tx.instructions[scan_ip].opname if scan_ip < len(tx.instructions) else None
+    )
+
+    def matches(name: str) -> bool:
+        pat = patterns[name]
+        return pre_store_ops == pat["pre_store_ops"] and (
+            post_store_op == pat["post_store_op"] or not pat["post_store_op"]
+        )
+
+    result_var: str | None = None
+    if matches("stored"):
+        result_var = tx.instructions[store_fast_ip].argval
+        result_on_stack = False
+    elif matches("discarded"):
+        result_var = None
+        result_on_stack = False
+        scan_ip = scan_ip + 1 if patterns["discarded"]["post_store_op"] else scan_ip
+    elif matches("returned") or pre_store_ops == patterns["consumed"]["pre_store_ops"]:
+        result_var = None
+        result_on_stack = True
+    else:
+        unimplemented(
+            gb_type="Comprehension analysis failed: No matches",
+            context=f"pre_store_ops={pre_store_ops}, post_store_op={post_store_op}",
+            explanation="Comprehension does not match any known bytecode pattern.",
+            hints=[],
+        )
+
+    return ComprehensionAnalysis(
+        end_ip=scan_ip,
+        result_var=result_var,
+        # pyrefly: ignore [unbound-name]
+        result_on_stack=result_on_stack,
+        iterator_vars=iterator_vars,
+        walrus_vars=walrus_vars,
+        captured_vars=captured_vars,
+    )
+
+
+def _handle_comprehension_graph_break(
+    tx: InstructionTranslatorBase, inst: Instruction
+) -> None:
+    """Handle list/dict comprehension graph break.
+
+    Builds a synthetic function wrapping the comprehension bytecode,
+    calls it via codegen_call_resume, then chains into the resume
+    function for the post-comprehension code.
+    """
+    assert sys.version_info >= (3, 12)
+    assert tx.instruction_pointer is not None
+
+    start_ip = tx.instruction_pointer - 1  # BUILD_LIST/BUILD_MAP
+    analysis = _analyze_comprehension(tx)
+    stack_pops = 1 + len(analysis.iterator_vars)
+    reason = GraphCompileReason("comprehension_graph_break", [tx.frame_summary()])
+    log.debug("comprehension triggered compile")
+
+    # --- Step 1: Compile the graph up to the comprehension ---
+
+    all_stack_locals_metadata = tx.output.compile_subgraph(
+        tx,
+        reason=reason,
+        stack_pops=stack_pops,
+    )
+    # Record which stack_pops items are NULL before popn loses the info.
+    # NULLs on the CPython stack can't be passed as function arguments.
+    stack_pops_null_mask = [
+        isinstance(tx.stack[len(tx.stack) - stack_pops + i], NullVariable)
+        for i in range(stack_pops)
+    ]
+
+    tx.popn(stack_pops)
+    meta = all_stack_locals_metadata[0]
+    cg = PyCodegen(tx.output.root_tx)
+
+    # Runtime stack after compile_subgraph:
+    #   cells, [frame_values], *(non-popped items), *(stack_pops items w/ NULLs)
+    # frame_values[0] = [frame N locals] (no stack items yet)
+
+    nonnull_count = sum(1 for m in stack_pops_null_mask if not m)
+
+    # live_stack_depth: stack items above cells/frame_values excluding NULLs
+    # that compile_subgraph didn't codegen (tracked in stack_null_idxes).
+    live_stack_depth = len(tx.stack) - len(meta.stack_null_idxes)
+
+    # --- Step 2: Pop stack_pops items and append non-nulls to frame_values[0] ---
+    # SWAP each item to TOS then LIST_APPEND or pop_null; fv_list stays at
+    # TOS throughout. Items append in TOS-first (reversed) order;
+    # _build_comprehension_fn compensates by loading in reverse.
+    cg.extend_output(
+        [
+            # frame_values[0] to TOS
+            *create_copy(live_stack_depth + stack_pops + 1),
+            cg.create_load_const(0),
+            cg.create_binary_subscr(),
+        ]
+    )
+    for i in reversed(range(stack_pops)):
+        cg.extend_output(create_swap(2))
+        if stack_pops_null_mask[i]:
+            cg.extend_output(cg.pop_null())
+        else:
+            cg.extend_output([create_instruction("LIST_APPEND", arg=1)])
+    cg.extend_output([create_instruction("POP_TOP")])
+
+    # Stack: cells, [frame_values], *(non-popped items)
+
+    # --- Step 3: Build comprehension function ---
+    new_code, fn_name = _build_comprehension_fn(
+        tx,
+        analysis,
+        start_ip,
+        stack_pops,
+        stack_pops_null_mask,
+        nonnull_count,
+        meta,
+    )
+
+    # --- Step 4: Extract [cells[0]] and [frame_values[0]] for codegen_call_resume ---
+    cg.extend_output(
+        [
+            *create_copy(live_stack_depth + 2),
+            cg.create_load_const(0),
+            cg.create_binary_subscr(),
+            create_instruction("BUILD_LIST", arg=1),
+            *create_copy(live_stack_depth + 2),
+            cg.create_load_const(0),
+            cg.create_binary_subscr(),
+            create_instruction("BUILD_LIST", arg=1),
+        ]
+    )
+
+    # Stack: ..., *(non-popped), [cells[0]], [frame_values[0]]
+
+    # --- Step 5: Call comprehension function via codegen_call_resume ---
+    tx.codegen_call_resume([new_code], [fn_name], cg)
+
+    # Stack: ..., *(non-popped), comp_result
+
+    # --- Step 6: Remove appended stack_pops items from frame_values[0] ---
+    if nonnull_count > 0:
+        frame_values_pos = live_stack_depth + 1 + 1  # +1 result, +1 frame_values
+        cg.extend_output(
+            [
+                *create_copy(frame_values_pos),
+                cg.create_load_const(0),
+                cg.create_binary_subscr(),
+                # frame_values[0] on TOS
+                create_dup_top(),
+                # frame_values[0], frame_values[0]
+                cg.create_load_const(-nonnull_count),
+                cg.create_load_const(None),
+                create_instruction("BUILD_SLICE", arg=2),
+                create_instruction("DELETE_SUBSCR"),
+                # del frame_values[0][-nonnull_count:]
+                create_instruction("POP_TOP"),
+            ]
+        )
+
+    # --- Step 7: Pass comprehension outputs to frame_values[0] ---
+    # Walrus vars first, then result_var.
+    vars_to_pass = analysis.walrus_vars + (
+        [analysis.result_var] if analysis.result_var else []
+    )
+
+    existing_vars: dict[str, int] = {}
+    for var_name in vars_to_pass:
+        tx.symbolic_locals[var_name] = UnknownVariable()
+        if var_name in meta.locals_names:
+            existing_vars[var_name] = meta.locals_names[var_name]
+        else:
+            meta.locals_names[var_name] = len(meta.locals_names)
+
+    fv_depth = live_stack_depth + 2  # comp_result + frame_values
+
+    # --- Walrus vars: extract from comp_result tuple ---
+    if analysis.walrus_vars:
+        # comp_result is (result, *walrus_vars).
+        cg.extend_output(
+            [
+                *create_copy(fv_depth),
+                cg.create_load_const(0),
+                cg.create_binary_subscr(),
+            ]
+        )
+        # Stack: ..., comp_tuple, fv0
+        for j, walrus_var in enumerate(analysis.walrus_vars):
+            cg.extend_output(
+                [
+                    *create_copy(2),
+                    cg.create_load_const(j + 1),
+                    cg.create_binary_subscr(),
+                ]
+            )
+            # Stack: ..., comp_tuple, fv0, walrus_value
+            if walrus_var in existing_vars:
+                # fv0[idx] = walrus_value
+                cg.extend_output(
+                    [
+                        *create_copy(2),  # copy fv0
+                        cg.create_load_const(existing_vars[walrus_var]),
+                        create_instruction("STORE_SUBSCR"),
+                    ]
+                )
+            else:
+                cg.extend_output([create_instruction("LIST_APPEND", arg=1)])
+            # Stack: ..., comp_tuple, fv0
+        cg.extend_output(
+            [
+                create_instruction("POP_TOP"),  # pop fv0
+                # Extract the result from the tuple.
+                cg.create_load_const(0),
+                cg.create_binary_subscr(),
+            ]
+        )
+        # Stack: ..., result
+
+    # --- Result: keep on stack, overwrite/append to fv[0], or discard ---
+    if analysis.result_on_stack:
+        tx.push(UnknownVariable())
+    elif analysis.result_var:
+        cg.extend_output(
+            [
+                *create_copy(fv_depth),
+                cg.create_load_const(0),
+                cg.create_binary_subscr(),
+                # Stack: ..., result, fv0
+            ]
+        )
+        if analysis.result_var in existing_vars:
+            cg.extend_output(
+                [
+                    cg.create_load_const(existing_vars[analysis.result_var]),
+                    create_instruction("STORE_SUBSCR"),
+                    # fv0[idx] = result
+                ]
+            )
+        else:
+            cg.extend_output(
+                [
+                    *create_swap(2),
+                    create_instruction("LIST_APPEND", arg=1),
+                    create_instruction("POP_TOP"),
+                ]
+            )
+    else:
+        cg.extend_output([create_instruction("POP_TOP")])
+
+    # Stack: cells, [frame_values], *(non-popped stack)
+    tx.output.add_output_instructions(cg.get_instructions())
+
+    # --- Step 8: Create resume function chain ---
+    resume_inst = tx.instructions[analysis.end_ip]
+    tx.output.add_output_instructions(
+        tx.create_call_resume_at(resume_inst, all_stack_locals_metadata)
+    )
+
+    tx.instruction_pointer = None
+
+
+def _build_comprehension_fn(
+    tx: InstructionTranslatorBase,
+    analysis: ComprehensionAnalysis,
+    start_ip: int,
+    stack_pops: int,
+    stack_pops_null_mask: list[bool],
+    nonnull_count: int,
+    meta: StackLocalsMetadata,
+) -> tuple[types.CodeType, str]:
+    """Build a synthetic function wrapping comprehension bytecode.
+
+    Uses the same calling convention as resume functions created by
+    create_resume / ContinueExecutionCache.generate: the first two args
+    are __nested_resume_fns and __nested_frame_values (ignored here),
+    followed by stack items and live locals.
+
+    Returns (code, name) where name is the global name for the function.
+    """
+    from .bytecode_transformation import transform_code_object
+    from .eval_frame import skip_code
+    from .resume_execution import CO_VARARGS, CO_VARKEYWORDS
+
+    # Args follow frame_values layout: locals first, then stack_pops items
+    # (appended to end of frame_values[0] by the caller).
+    # codegen_call_resume unpacks frame_values[0] as positional args.
+    argnames = tuple(k for k in meta.locals_names if k not in tx.cell_and_freevars())
+    args = (
+        ["__nested_resume_fns", "__nested_frame_values"]
+        + list(argnames)
+        + [f"___stack{i}" for i in range(nonnull_count)]
+    )
+
+    freevars = tuple(
+        sorted(list(tx.f_code.co_cellvars or []) + list(tx.f_code.co_freevars or []))
+    )
+
+    lineno = tx.lineno if tx.lineno is not None else tx.f_code.co_firstlineno
+    fn_name = unique_id(f"__comprehension_{tx.f_code.co_name}_at_{lineno}")
+
+    comprehension_body_vars = (
+        analysis.iterator_vars
+        + analysis.walrus_vars
+        + ([analysis.result_var] if analysis.result_var else [])
+        + analysis.captured_vars
+    )
+
+    def update(instructions: list[Instruction], code_options: dict[str, Any]) -> None:
+        code_options["co_name"] = fn_name
+        if sys.version_info >= (3, 11):
+            code_options["co_qualname"] = fn_name
+        code_options["co_firstlineno"] = lineno
+        code_options["co_cellvars"] = ()
+        code_options["co_freevars"] = freevars
+        code_options["co_argcount"] = len(args)
+        code_options["co_posonlyargcount"] = 0
+        code_options["co_kwonlyargcount"] = 0
+        code_options["co_varnames"] = tuple(
+            args + [v for v in comprehension_body_vars if v not in args]
+        )
+        code_options["co_flags"] = code_options["co_flags"] & ~(
+            CO_VARARGS | CO_VARKEYWORDS
+        )
+
+        prefix: list[Instruction] = []
+        if freevars:
+            prefix.append(create_instruction("COPY_FREE_VARS", arg=len(freevars)))
+        prefix.append(create_instruction("RESUME", arg=0))
+
+        # Push stack_pops items onto operand stack so the comprehension
+        # bytecode finds them where it expects (iterator + saved vars).
+        # NULL positions get PUSH_NULL, non-null get LOAD_FAST.
+        # Items were appended to frame_values[0] in TOS-first order,
+        # so load in reverse to reconstruct the original stack layout.
+        nonnull_i = nonnull_count - 1
+        for i in range(stack_pops):
+            if stack_pops_null_mask[i]:
+                prefix.append(create_instruction("PUSH_NULL"))
+            else:
+                prefix.append(
+                    create_instruction("LOAD_FAST", argval=f"___stack{nonnull_i}")
+                )
+                nonnull_i -= 1
+
+        comp_insts = _copy_comprehension_bytecode(tx, start_ip, analysis.end_ip)
+
+        # Epilogue: ensure result is on stack, pack walrus vars, return.
+        epilogue: list[Instruction] = []
+        if not analysis.result_on_stack:
+            if analysis.result_var:
+                epilogue.append(
+                    create_instruction("LOAD_FAST", argval=analysis.result_var)
+                )
+            else:
+                epilogue.append(create_instruction("LOAD_CONST", argval=None))
+        if analysis.walrus_vars:
+            for var_name in analysis.walrus_vars:
+                epilogue.append(create_instruction("LOAD_FAST", argval=var_name))
+            epilogue.append(
+                create_instruction(
+                    "BUILD_TUPLE",
+                    arg=1 + len(analysis.walrus_vars),
+                )
+            )
+        epilogue.append(create_instruction("RETURN_VALUE"))
+
+        instructions[:] = prefix + comp_insts + epilogue
+
+    new_code, _ = transform_code_object(tx.f_code, update)
+    skip_code(new_code)
+
+    # Install as global
+    tx.output.install_resume_function_global(fn_name, new_code, tx.f_globals)
+
+    return new_code, fn_name
+
+
+def _copy_comprehension_bytecode(
+    tx: InstructionTranslatorBase, start_ip: int, end_ip: int
+) -> list[Instruction]:
+    """Copy comprehension bytecode instructions, updating jump targets."""
+    inst_map: dict[Instruction, Instruction] = {}
+    copied_insts: list[Instruction] = []
+
+    for ip in range(start_ip, end_ip):
+        original_inst = tx.instructions[ip]
+        copied_inst = copy.copy(original_inst)
+        copied_inst.exn_tab_entry = None
+        inst_map[original_inst] = copied_inst
+        copied_insts.append(copied_inst)
+
+    for copied_inst in copied_insts:
+        if copied_inst.target is not None and copied_inst.target in inst_map:
+            copied_inst.target = inst_map[copied_inst.target]
+
+    return copied_insts
+
+
+def maybe_setup_comprehension_speculation(
+    tx: InstructionTranslatorBase, inst: Instruction
+) -> bool:
+    """
+    Handle comprehension start for Python 3.12+ BUILD_LIST/BUILD_MAP with argval 0.
+    Returns True if a graph break was triggered and the caller should return early.
+    """
+    if not (sys.version_info >= (3, 12) and inst.argval == 0):
+        return False
+
+    if not _is_comprehension_start(tx):
+        return False
+
+    can_speculate = (
+        all(b.can_restore() for b in tx.block_stack)
+        and not tx.one_graph
+        and not tx.error_on_graph_break
+        and not tx.is_tracing_resume_prologue
+        and not tx.active_generic_context_managers
+        and tx.output.current_tracer.parent is None
+    )
+
+    if can_speculate and tx.parent is not None:
+        can_speculate = tx._can_speculate_comprehension_nested()
+    # Only set up speculation at depth 0 (outermost comprehension)
+    if can_speculate and tx._comprehension_depth == 0:
+        speculation = tx.speculate()
+        if speculation.failed(tx):
+            _handle_comprehension_graph_break(tx, inst)
+            return True
+        tx.current_speculation = speculation
+    end_for_ip = _find_comprehension_end_for_ip(tx)
+    assert end_for_ip >= 0
+    tx._comprehension_end_for_ips.add(end_for_ip)
+    tx._comprehension_depth += 1
+    return False
diff --git a/torch/_dynamo/comptime.py b/torch/_dynamo/comptime.py
index c6302d5ce336e..811ed4a46e482 100644
--- a/torch/_dynamo/comptime.py
+++ b/torch/_dynamo/comptime.py
@@ -39,7 +39,7 @@ def my_model(x):
 import time
 import traceback
 from collections.abc import Callable, Sequence
-from typing import Any, Optional, TextIO, Union
+from typing import Any, TextIO
 
 import torch
 from torch._dynamo.symbolic_convert import InstructionTranslatorBase
@@ -66,7 +66,7 @@ class ComptimeVar:
     def __init__(self, v: VariableTracker) -> None:
         self.__variable = v
 
-    def as_proxy(self) -> Union[VariableTracker, Sequence[VariableTracker]]:
+    def as_proxy(self) -> VariableTracker | Sequence[VariableTracker]:
         """
         Returns an fx.Proxy (or tuple/list of fx.Proxy) representing
         this variable in the FX graph we are assembling to pass
@@ -86,7 +86,7 @@ def is_proxy(self) -> bool:
         """
         return self.__variable.is_proxy()
 
-    def as_fake(self) -> Union[FakeTensor, torch.SymInt]:
+    def as_fake(self) -> FakeTensor | torch.SymInt:
         """
         Returns a "fake" value (either a FakeTensor or a SymInt)
         representing the variable in question.  This only works
@@ -98,7 +98,7 @@ def as_fake(self) -> Union[FakeTensor, torch.SymInt]:
         """
         return self.__variable.as_proxy().node.meta["example_value"]
 
-    def size(self, dim: Optional[int] = None) -> Union[int, torch.SymInt]:
+    def size(self, dim: int | None = None) -> int | torch.SymInt:
         """
         Returns the size of the tensor (if dim is None) or the size
         at the dimension dim.  The returned size may be a SymInt.
@@ -214,9 +214,7 @@ def assert_static(self, val: ComptimeVar) -> None:
             "expected static but got dynamic (run with TORCH_LOGS=dynamic for more info)"
         )
 
-    def print_graph(
-        self, *, verbose: bool = True, file: Optional[TextIO] = None
-    ) -> None:
+    def print_graph(self, *, verbose: bool = True, file: TextIO | None = None) -> None:
         """
         Print the partially constructed FX graph that would be passed
         to the user compiler after compilation.
@@ -230,17 +228,15 @@ def parent(self) -> "ComptimeContext":
 
     def __get_tx(self, stacklevel: int) -> Any:
         tx = self.__tx
-        # pyrefly: ignore [bad-assignment]
+        # pyrefly: ignore [bad-assignment, non-convergent-recursion]
         for _ in range(stacklevel):
             tx = tx.parent  # type: ignore[assignment]
         return tx
 
-    def print(self, val: Any, *, file: Optional[TextIO] = None) -> None:
+    def print(self, val: Any, *, file: TextIO | None = None) -> None:
         print(repr(val), file=file)
 
-    def print_disas(
-        self, *, file: Optional[TextIO] = None, stacklevel: int = 0
-    ) -> None:
+    def print_disas(self, *, file: TextIO | None = None, stacklevel: int = 0) -> None:
         """
         Print the current series of opcodes being executed (not including
         parent frames), including where you are in the particular opcode
@@ -256,7 +252,7 @@ def print_disas(
         )
 
     def print_value_stack(
-        self, *, file: Optional[TextIO] = None, stacklevel: int = 0
+        self, *, file: TextIO | None = None, stacklevel: int = 0
     ) -> None:
         """
         Print the current Python value stack.  Note that this is NOT the same
@@ -272,9 +268,7 @@ def print_value_stack(
         for s in tx.stack:
             print(f"- {s.debug_repr()}", file=file)
 
-    def print_locals(
-        self, *, file: Optional[TextIO] = None, stacklevel: int = 0
-    ) -> None:
+    def print_locals(self, *, file: TextIO | None = None, stacklevel: int = 0) -> None:
         """
         Print all of the locals available in the current context.
         By default this view is very limited; you can get more information
@@ -284,7 +278,7 @@ def print_locals(
         for k, v in tx.symbolic_locals.items():
             print(f"{k} = {v.debug_repr()}", file=file)
 
-    def print_bt(self, *, file: Optional[TextIO] = None, stacklevel: int = 0) -> None:
+    def print_bt(self, *, file: TextIO | None = None, stacklevel: int = 0) -> None:
         """
         Print the user code backtrace, starting at the beginning of the
         frame Dynamo started evaluating.  Note that this MAY NOT go all
@@ -303,7 +297,7 @@ def print_bt(self, *, file: Optional[TextIO] = None, stacklevel: int = 0) -> Non
             file=file,
         )
 
-    def print_guards(self, *, file: Optional[TextIO] = None) -> None:
+    def print_guards(self, *, file: TextIO | None = None) -> None:
         """
         Print the currently installed guards for the Dynamo context.
         This does NOT include guards associated with variables that
@@ -328,7 +322,7 @@ def _i_will_not_complain_if_bc_breaks_InstructionTranslator(
         """
         return self.__tx
 
-    def sleep(self, sec: Union[int, float]) -> None:
+    def sleep(self, sec: int | float) -> None:
         time.sleep(sec)
 
 
@@ -437,7 +431,7 @@ def inner(inner_ctx: ComptimeContext) -> None:
         comptime(inner)
 
     @staticmethod
-    def sleep(sec: Union[int, float]) -> None:
+    def sleep(sec: int | float) -> None:
         comptime(lambda ctx: ctx.sleep(ctx.get_local("sec").as_python_constant()))
 
 
diff --git a/torch/_dynamo/config.py b/torch/_dynamo/config.py
index daca73b978c4f..367341dcb6403 100644
--- a/torch/_dynamo/config.py
+++ b/torch/_dynamo/config.py
@@ -17,7 +17,7 @@
 import tempfile
 from collections.abc import Callable
 from os.path import abspath, dirname
-from typing import Any, Literal, Optional, TYPE_CHECKING, Union
+from typing import Any, Literal, TYPE_CHECKING
 
 from torch._environment import is_fbcode
 from torch.utils._config_module import Config, get_tristate_env, install_config_module
@@ -31,7 +31,7 @@
 # Design doc: https://docs.google.com/document/d/1ZRfTWKa8eaPq1AxaiHrq4ASTPouzzlPiuquSBEJYwS8/edit#
 # the name of a file to write the logs to
 # [@compile_ignored: debug]
-log_file_name: Optional[str] = None
+log_file_name: str | None = None
 
 # [@compile_ignored: debug] Verbose will print full stack traces on warnings and errors
 verbose = os.environ.get("TORCHDYNAMO_VERBOSE", "0") == "1"
@@ -44,7 +44,7 @@
 #   - Individual IDs: "0,5,10"
 #   - Ranges: "10-20" (inclusive)
 #   - Comparisons: ">10", ">=10", "<5", "<=5"
-# Backends can be: "eager", "aot_eager", "inductor", "inductor:reduce-overhead", etc.
+# Backends can be: "eager", "aot_eager", "inductor", etc.
 # Examples:
 #   ">10:eager"                    - Run graphs with frame_id > 10 in dynamo eager backend
 #   "<=5:aot_eager;>5:inductor"    - First 6 graphs use aot_eager, rest use inductor
@@ -62,6 +62,17 @@
     "TORCH_COMPILE_OVERRIDE_INDUCTOR_CONFIGS", ""
 )
 
+# Override dynamo config for specific graphs (for debugging/bisecting).
+# Format: "filter1:config1;filter2:config2;..." where filter uses same syntax as
+# debug_backend_override, and config is "key=value" or "key=value,key2=value2".
+# Examples:
+#   "0-5:specialize_float=True"  - Specialize floats for graphs 0-5
+#   ">10:automatic_dynamic_shapes=False"  - Disable dynamic shapes for graphs > 10
+# [@compile_ignored: debug]
+debug_dynamo_config_override: str = os.environ.get(
+    "TORCH_COMPILE_OVERRIDE_DYNAMO_CONFIGS", ""
+)
+
 # Validate that fake_fn and real_fn in @leaf_function decorators produce outputs
 # with matching shapes and dtypes in eager mode. Helps catch mismatches early.
 # Disabled by default to avoid runtime overhead.
@@ -107,8 +118,9 @@
 # [@compile_ignored: runtime_behaviour] safeguarding to prevent horrible recomps
 accumulated_recompile_limit = 256
 
-# [@compile_ignored: runtime_behaviour] skip tracing recursively if cache limit is hit (deprecated: does not do anything)
-skip_code_recursive_on_recompile_limit_hit = True
+skip_code_recursive_on_recompile_limit_hit: bool = Config(
+    default=True, deprecated=True, deprecation_message="does not do anything"
+)
 
 # raise a hard error if cache limit is hit.  If you are on a model where you
 # know you've sized the cache correctly, this can help detect problems when
@@ -122,10 +134,12 @@
     alias="torch._dynamo.config.accumulated_recompile_limit"
 )
 
-# (deprecated: does not do anything)
 skip_code_recursive_on_cache_limit_hit: bool = Config(
-    alias="torch._dynamo.config.skip_code_recursive_on_recompile_limit_hit"
+    alias="torch._dynamo.config.skip_code_recursive_on_recompile_limit_hit",
+    deprecated=True,
+    deprecation_message="does not do anything",
 )
+
 fail_on_cache_limit_hit: bool = Config(
     alias="torch._dynamo.config.fail_on_recompile_limit_hit"
 )
@@ -166,6 +180,26 @@
 # Valid options: "dynamic", "unbacked"
 automatic_dynamic_shapes_mark_as: Literal["dynamic", "unbacked"] = "dynamic"
 
+# When True, adds exclusion guards for tensor dims and scalars that transition
+# from static to dynamic via automatic_dynamic_shapes.
+#
+# Invariant: when enabled, automatic_dynamic recompilation preserves graph
+# selection — inputs that matched a previous static cache entry will continue
+# to use that entry, not be intercepted by a newer dynamic entry. This holds
+# as long as recompilations are caused solely by the same variable being
+# observed with different static values (progressive dynamism). A recompilation
+# triggered by a different reason (e.g., a guard failure unrelated to shape
+# transitions) will clear the exclusion state for that entry.
+#
+# Mechanism: the exclusion guard rejects inputs matching the prior static
+# graph's sizes, so those inputs fall through to the more specialized static
+# graph instead of being captured by the newer dynamic graph.
+#
+# Scope: applies only to graph-input-level dimension and scalar transitions.
+# Does NOT handle data-dependent branching (if x.size(0) > k), graph breaks,
+# or other recompilation triggers where no dimension actually transitions.
+automatic_dynamic_exclusion_guard = False
+
 # log graph in/out metadata
 # This is only turned on for export today since we
 # know we are tracing a flat callable. later, this
@@ -374,15 +408,15 @@
 # Note that to avoid breaking the existing usage, mode 1 and mode 4 can be
 # specified with a boolean value. True is using ddp_optimizer and False is
 # no optimization.
-optimize_ddp: Union[
-    bool,
-    Literal[
+optimize_ddp: (
+    bool
+    | Literal[
         "ddp_optimizer",
         "python_reducer",
         "python_reducer_without_compiled_forward",
         "no_optimization",
-    ],
-] = True
+    ]
+) = True
 
 # By default, Dynamo emits runtime asserts (e.g. torch._check) in the graph.
 # In some cases those asserts could be performance costly
@@ -428,7 +462,7 @@
 skip_tensor_guards_with_matching_dict_tags = True
 
 # Skips guards on func.__defaults__ if the element to be guarded is a constant
-skip_guards_on_constant_func_defaults = True
+skip_guards_on_constant_func_defaults = False
 
 
 # The recursive-dict-tag guard relies on the class/function identity staying
@@ -451,7 +485,7 @@
 
 # Speedup guard execution of nested nn modules by recursively checking for dict
 # tags to avoid full guard execution.
-use_recursive_dict_tags_for_guards = True
+use_recursive_dict_tags_for_guards = False
 
 # Maximum number of objects for which we check dict pointers tags. This is
 # useful for regional compilation.
@@ -530,11 +564,13 @@
 inline_inbuilt_nn_modules = Config(  # type: ignore[var-annotated]
     default=True,
     justknob="pytorch/compiler:inline_inbuilt_nn_modules",
+    deprecated=True,
+    deprecation_message="does not do anything, inline_inbuilt_nn_modules is always True",
 )
 
 # Resume tracing in nested frames if a nested graph break occurs
 # Old behavior is to bubble up the graph break to the top level frame.
-nested_graph_breaks = False
+nested_graph_breaks: bool = False
 
 # If True, error if Dynamo attempts to trace more code while running compiled code in fullgraph=True.
 # If Dynamo determines that it should skip tracing the code (either at the C/C++ or Python level),
@@ -782,6 +818,11 @@ def default_debug_dir_root() -> str:
 # NCCL timeout.
 enable_compiler_collectives = os.environ.get("TORCH_COMPILER_COLLECTIVES", "0") == "1"
 
+# Allow for experimental support of compiled p2p ops
+enable_p2p_compilation = (
+    os.environ.get("TORCHDYNAMO_ENABLE_P2P_COMPILATION", "0") == "1"
+)
+
 # Enables a local, filesystem "profile" which can be used for automatic
 # dynamic decisions, analogous to profile-guided optimization.  This config
 # ONLY has an effect if torch.compiler.config.workflow_id is specified,
@@ -810,7 +851,7 @@ def default_debug_dir_root() -> str:
 )
 
 # Like above, but using remote cache
-automatic_dynamic_remote_pgo: Optional[bool] = get_tristate_env(
+automatic_dynamic_remote_pgo: bool | None = get_tristate_env(
     "TORCH_DYNAMO_AUTOMATIC_DYNAMIC_REMOTE_PGO"
 )
 
@@ -820,7 +861,7 @@ def default_debug_dir_root() -> str:
 )
 
 # Common prefix to append to the id of each compile run to filter out data
-pt2_compile_id_prefix: Optional[str] = os.environ.get("PT2_COMPILE_ID_PREFIX", None)
+pt2_compile_id_prefix: str | None = os.environ.get("PT2_COMPILE_ID_PREFIX", None)
 
 # Run GC at the end of compilation
 run_gc_after_compile = Config(  # type: ignore[var-annotated]
@@ -843,15 +884,22 @@ def default_debug_dir_root() -> str:
 # and AOTAutograd runtime wrapper.
 record_runtime_overhead = True
 
-enable_aot_compile = False
+# Flag to enable the use of torch.compile().aot_compile() API. Should be always True.
+enable_aot_compile = True
 
 # HACK: this is for testing custom ops profiling only
-_custom_ops_profile: Optional[Any] = None
+_custom_ops_profile: Any | None = None
 
 # Experimental flag to enable regional compile on invoke_subgraph HOP.
 # For testing only!
 enable_invoke_subgraph_regional_compile: bool = False
 
+# When True, run a post-tracing pass that inlines all invoke_subgraph HOPs
+# back into the parent graph, producing a flat FX graph. Useful when
+# downstream compilers (like vllm-compile) don't support HOPs or prefer a
+# flat graph.
+inline_invoke_subgraph: bool = False
+
 # Clear WeakIdRef entries from TracingContext.tensor_to_context and
 # MetaTensorDescriber.lookup_tensor at the end of compile. These weakrefs
 # can block torch.utils.swap_tensors from working after compile.
@@ -859,7 +907,7 @@ def default_debug_dir_root() -> str:
 #   don't clear for custom backends (to support standalone_compile, etc.)
 # - True: always clear regardless of backend
 # - False: never clear regardless of backend
-invalidate_compile_context_weakrefs: Optional[bool] = None
+invalidate_compile_context_weakrefs: bool | None = None
 
 if TYPE_CHECKING:
     from torch.utils._config_typing import *  # noqa: F401, F403
diff --git a/torch/_dynamo/convert_frame.py b/torch/_dynamo/convert_frame.py
index 7441dac7dfc97..d54af42415880 100644
--- a/torch/_dynamo/convert_frame.py
+++ b/torch/_dynamo/convert_frame.py
@@ -36,7 +36,7 @@
 import logging
 import os
 import pstats
-import random
+import random  # noqa: F401 -- eval_frame_cpp.cpp imports random at runtime; torch.package needs this to detect the dependency
 import subprocess
 import sys
 import tempfile
@@ -50,7 +50,7 @@
 from dataclasses import dataclass
 from pathlib import Path
 from types import CellType, CodeType, FunctionType, ModuleType
-from typing import Any, NoReturn, Optional, TypeVar, Union
+from typing import Any, NoReturn, TypeVar
 from typing_extensions import ParamSpec
 from weakref import ReferenceType
 
@@ -75,11 +75,11 @@
 from torch.fx.graph_module import _forward_from_src as original_forward_from_src
 from torch.monitor import _WaitCounter
 from torch.nn.parallel.distributed import DistributedDataParallel
-from torch.utils._ordered_set import OrderedSet
 from torch.utils._python_dispatch import (
     _disable_current_modes,
     any_torch_dispatch_mode_on_stack,
     is_in_any_mode_without_ignore_compile_internals,
+    is_traceable_wrapper_subclass,
 )
 from torch.utils._traceback import CapturedTraceback, format_traceback_short
 
@@ -99,6 +99,7 @@
     exceeds_recompile_limit,
     is_recompilation,
 )
+from .code_context import code_context
 from .eval_frame import (
     always_optimize_code_objects,
     Constraint,
@@ -121,6 +122,7 @@
     UncapturedHigherOrderOpError,
     unimplemented,
     Unsupported,
+    UserError,
 )
 from .graph_bytecode_inputs import reset_user_object_tracking
 from .guards import (
@@ -173,7 +175,7 @@
 from .variables.torch_function import torch_function_mode_stack_state_mgr
 
 
-np: Optional[ModuleType]
+np: ModuleType | None
 try:
     import numpy as np
 except ModuleNotFoundError:
@@ -208,7 +210,7 @@ class TODO_UNKNOWN:
 
 
 def _clear_fake_mode_weakrefs(
-    fake_mode: Optional[torch._subclasses.fake_tensor.FakeTensorMode],
+    fake_mode: torch._subclasses.fake_tensor.FakeTensorMode | None,
 ) -> None:
     """Clear WeakIdRef entries from a FakeTensorMode's describer."""
     if fake_mode is None:
@@ -241,12 +243,12 @@ def clear(self) -> None:
 input_codes = Tracker()
 output_codes = Tracker()
 
-initial_global_state: Optional[GlobalStateGuard] = None
+initial_global_state: GlobalStateGuard | None = None
 
 
 @functools.wraps(original_forward_from_src)
 def fx_forward_from_src_skip_result(
-    src: str, globals: dict[str, Any], co_fields: Optional[dict[str, str]] = None
+    src: str, globals: dict[str, Any], co_fields: dict[str, str] | None = None
 ) -> FunctionType:
     # we monkey patch FX to prevent infinite loop of trying to convert
     # our generated code
@@ -293,9 +295,11 @@ def preserve_global_state(fn: Callable[_P, _T]) -> Callable[_P, _T]:
     """
     Context manager to:
         1) Save/restore torch.is_grad_enabled() state
-        2) Save/restore python random state
-        3) Save/restore torch random state
-        4) Monkey patch torch.fx.graph_module._forward_from_src
+        2) Save/restore torch random state
+        3) Monkey patch torch.fx.graph_module._forward_from_src
+
+    NOTE: Python random state is preserved in eval_frame_cpp.cpp instead,
+    so that it wraps more of the compilation pipeline.
     """
 
     @functools.wraps(fn)
@@ -318,7 +322,6 @@ def _fn(*args: _P.args, **kwargs: _P.kwargs) -> _T:
             prior_mobile_allocator_state = (
                 torch._C._is_default_mobile_cpu_allocator_set()
             )
-            py_rng_state = random.getstate()
             prior_dtype = torch.get_default_dtype()
             torch_rng_state = torch.random.get_rng_state()
             cuda_rng_state = None
@@ -335,22 +338,17 @@ def _fn(*args: _P.args, **kwargs: _P.kwargs) -> _T:
             exit_stack.enter_context(
                 torch.fx._symbolic_trace._maybe_revert_all_patches()
             )
-            exit_stack.enter_context(torch_function_mode_stack_state_mgr)
             reset_user_object_tracking()
             try:
                 return fn(*args, **kwargs)
             finally:
                 cleanup.close()
-                assert torch._C._len_torch_function_stack() == 0, (
-                    "Torch function mode stack state changed while dynamo tracing, please report a bug"
-                )
                 exit_stack.close()
                 torch._C._set_grad_enabled(prior_grad_mode)
                 torch.autograd.grad_mode._enter_inference_mode(prior_inference_mode)
                 torch.use_deterministic_algorithms(
                     prior_deterministic, warn_only=prior_warn_only
                 )
-                random.setstate(py_rng_state)
                 torch.random.set_rng_state(torch_rng_state)
                 torch.set_default_dtype(prior_dtype)
                 curr_mobile_allocator_state = (
@@ -455,7 +453,7 @@ def has_tensor(obj: object) -> bool:
 def exception_handler(
     e: Exception,
     code: CodeType,
-    frame: Optional[DynamoFrameType] = None,
+    frame: DynamoFrameType | None = None,
     export: bool = False,
 ) -> None:
     record_filename = None
@@ -468,9 +466,7 @@ def exception_handler(
 
 
 FRAME_COUNTER = 0
-FRAME_COMPILE_COUNTER: typing.Counter[Union[int, FrameStateSizeEntry]] = (
-    collections.Counter()
-)
+FRAME_COMPILE_COUNTER: typing.Counter[int | FrameStateSizeEntry] = collections.Counter()
 
 
 def maybe_cprofile(func: Callable[_P, _T]) -> Callable[_P, _T]:
@@ -556,11 +552,11 @@ def profile_wrapper(*args: _P.args, **kwargs: _P.kwargs) -> _T:
 
 @dataclass
 class ConvertFrameBox:
-    error_on_graph_break: Optional[bool] = None
+    error_on_graph_break: bool | None = None
 
 
 def get_compile_id(
-    frame_state: dict[str, Union[int, FrameStateSizeEntry]],
+    frame_state: dict[str, int | FrameStateSizeEntry],
 ) -> CompileId:
     global FRAME_COUNTER
     if "_id" not in frame_state:
@@ -590,6 +586,7 @@ def __init__(
         export: bool = False,
         export_constraints: Any | None = None,
         package: CompilePackage | None = None,
+        recompile_limit: int | None = None,
     ) -> None:
         # assert export_constraints is None
         reset_graph_break_dup_checker()
@@ -598,6 +595,7 @@ def __init__(
         self._export = export
         self._export_constraints = export_constraints
         self._package = package
+        self._recompile_limit = recompile_limit
         self._box = ConvertFrameBox()
 
     @property
@@ -607,14 +605,15 @@ def _clone_with_backend(self) -> Callable[[CompilerFn], ConvertFrameAssert]:
             self._one_graph,
             self._export,
             self._export_constraints,
+            recompile_limit=self._recompile_limit,
         )
 
     def __call__(
         self,
         frame: DynamoFrameType,
-        cache_entry: Optional[CacheEntry],
+        cache_entry: CacheEntry | None,
         hooks: Hooks,
-        frame_state: dict[str, Union[int, FrameStateSizeEntry]],
+        frame_state: dict[str, int | FrameStateSizeEntry],
         *,
         skip: int = 0,
     ) -> ConvertFrameReturn:
@@ -725,7 +724,15 @@ def __call__(
             dynamo_tls.traced_frame_infos.append(info)
 
         try:
-            with compile_context(CompileContext(compile_id)):
+            compile_ctx = compile_context(CompileContext(compile_id))
+            # When recompile_limit is set, temporarily override the global
+            # config so the existing exceeds_recompile_limit check uses it.
+            recompile_ctx = (
+                config.patch(recompile_limit=self._recompile_limit)
+                if self._recompile_limit is not None
+                else contextlib.nullcontext()
+            )
+            with compile_ctx, recompile_ctx:
                 result = _compile(
                     frame.f_code,
                     frame.f_globals,
@@ -763,11 +770,17 @@ def convert_frame_assert(
     one_graph: bool = True,
     export: bool = False,
     export_constraints: Any | None = None,
-    package: Optional[CompilePackage] = None,
+    package: CompilePackage | None = None,
+    recompile_limit: int | None = None,
 ) -> ConvertFrameAssert:
     """Fully convert a frame into an FX graph, raising an exception if we fail."""
     return ConvertFrameAssert(
-        compiler_fn, one_graph, export, export_constraints, package
+        compiler_fn,
+        one_graph,
+        export,
+        export_constraints,
+        package,
+        recompile_limit,
     )
 
 
@@ -778,12 +791,47 @@ def convert_frame_assert(
 
 # we have to use `OrderedDict` to make `RemovableHandle` work.
 _bytecode_hooks: dict[int, BytecodeHook] = OrderedDict()
+_BYTECODE_HOOK_SIDE_EFFECTS_CONTEXT_KEY = "bytecode_hook_side_effects"
+
+
+def get_compiled_code_side_effects(
+    code: types.CodeType,
+) -> tuple[str, ...] | None:
+    """Return Dynamo's replayed Python side-effect sources for compiled code.
+
+    Returns ``None`` when ``code`` was not produced by Dynamo or no metadata was
+    attached to it.
+    """
+    if not code_context.has_context(code):
+        return None
+    side_effects = code_context.get_context(code).get(
+        _BYTECODE_HOOK_SIDE_EFFECTS_CONTEXT_KEY
+    )
+    if side_effects is None:
+        return None
+    return side_effects
+
+
+def compiled_code_has_side_effects(code: types.CodeType) -> bool:
+    """Return whether Dynamo recorded replayed Python side effects for compiled code."""
+    return bool(get_compiled_code_side_effects(code))
+
+
+def _copy_code_context(src_code: types.CodeType, dst_code: types.CodeType) -> None:
+    if not code_context.has_context(src_code):
+        return
+    src_context = code_context.get_context(src_code)
+    if _BYTECODE_HOOK_SIDE_EFFECTS_CONTEXT_KEY in src_context:
+        code_context.get_context(dst_code)[_BYTECODE_HOOK_SIDE_EFFECTS_CONTEXT_KEY] = (
+            src_context[_BYTECODE_HOOK_SIDE_EFFECTS_CONTEXT_KEY]
+        )
 
 
 def register_bytecode_hook(hook: BytecodeHook) -> RemovableHandle:
     """Register hooks for bytecode generated by Dynamo. The hook can do some
-    logging, as well as return a new code object to be used. Please refer
-    to `BytecodeHook` for the hook signature.
+    logging, as well as return a new code object to be used. Hooks can query
+    `get_compiled_code_side_effects(new_code)` for replayed Python side effect
+    sources. Please refer to `BytecodeHook` for the hook signature.
     """
     handle = RemovableHandle(_bytecode_hooks)
     _bytecode_hooks[handle.id] = hook
@@ -809,9 +857,9 @@ def trace_frame(
     *,
     export: bool = False,
     export_constraints: Any | None = None,
-    frame_state: Optional[dict[str, Union[int, FrameStateSizeEntry]]] = None,
-    distributed_state: Optional[DistributedState] = None,
-    package: Optional[CompilePackage] = None,
+    frame_state: dict[str, int | FrameStateSizeEntry] | None = None,
+    distributed_state: DistributedState | None = None,
+    package: CompilePackage | None = None,
 ) -> DynamoTracerOutput:
     from torch.fx.experimental.validator import bisect, translation_validation_enabled
 
@@ -889,14 +937,14 @@ class DynamoOutput:
 
     tracer_output: DynamoTracerOutput
     bytecode: types.CodeType
-    last_attempt_start_time: Optional[float]
+    last_attempt_start_time: float | None
 
     def build_guards(
         self,
         code: types.CodeType,
-        hooks: Optional[Hooks] = None,
+        hooks: Hooks | None = None,
         save: bool = False,
-        cache_entry: Optional[CacheEntry] = None,
+        cache_entry: CacheEntry | None = None,
         strict_error: bool = False,
     ) -> CheckFunctionManager:
         output_graph = self.tracer_output.output_graph
@@ -913,8 +961,8 @@ def build_guards(
 
     def graph_capture_output(
         self,
-        argdefs: Optional[tuple[Any, ...]] = None,
-        kwdefaults: Optional[dict[str, Any]] = None,
+        argdefs: tuple[Any, ...] | None = None,
+        kwdefaults: dict[str, Any] | None = None,
     ) -> GraphCaptureOutput:
         output_graph = self.tracer_output.output_graph
         assert output_graph is not None
@@ -959,9 +1007,9 @@ class GraphRuntimeEnv:
     bytecode: types.CodeType
     import_sources: dict[str, str]
     used_globals: dict[str, Any]
-    closure: Optional[tuple[Any, ...]]
-    argdefs: Optional[tuple[Any, ...]]
-    kwdefaults: Optional[dict[str, Any]] = None
+    closure: tuple[Any, ...] | None
+    argdefs: tuple[Any, ...] | None
+    kwdefaults: dict[str, Any] | None = None
     external_refs: set[str] = dataclasses.field(default_factory=set)
 
     def forward_callable(
@@ -969,7 +1017,7 @@ def forward_callable(
         backend_id: str,
         compiled_fn: Callable[..., Any],
         *,
-        extra_globals: Optional[dict[str, Any]] = None,
+        extra_globals: dict[str, Any] | None = None,
     ) -> Callable[..., Any]:
         import_sources = {
             alias: importlib.import_module(module_name)
@@ -1009,6 +1057,20 @@ def _check_external_refs(self, f_globals: dict[str, Any]) -> None:
             )
 
 
+def _safe_builtins_dict(builtins_dict: dict[str, Any]) -> dict[str, Any]:
+    """Filter a builtins dict to only picklable entries for serialization."""
+    import pickle
+
+    result = {}
+    for k, v in builtins_dict.items():
+        try:
+            pickle.dumps(v)
+            result[k] = v
+        except Exception:
+            pass
+    return result
+
+
 @dataclass
 class GraphCaptureOutput:
     """
@@ -1019,17 +1081,17 @@ class GraphCaptureOutput:
     import_sources: dict[str, str]
     traced_code: list[CodeType]
     bytecode: CodeType
-    closure: Optional[tuple[Any, ...]]
-    argdefs: Optional[tuple[Any, ...]]
-    kwdefaults: Optional[dict[str, Any]]
+    closure: tuple[Any, ...] | None
+    argdefs: tuple[Any, ...] | None
+    kwdefaults: dict[str, Any] | None
     f_globals: dict[str, Any]
 
     def build_guards(
         self,
         code: types.CodeType,
-        hooks: Optional[Hooks] = None,
+        hooks: Hooks | None = None,
         save: bool = False,
-        cache_entry: Optional[CacheEntry] = None,
+        cache_entry: CacheEntry | None = None,
         strict_error: bool = False,
     ) -> CheckFunctionManager:
         return CheckFunctionManager(
@@ -1058,6 +1120,17 @@ def get_runtime_env(self) -> GraphRuntimeEnv:
         # Scan bytecode for all external references
         external_refs = self._get_external_refs(self.bytecode)
 
+        # Best-effort serialization of builtins referenced by the bytecode.
+        # Similar to how guards prune __builtins_dict__ to only used entries.
+        import builtins as _builtins
+
+        for ref in external_refs:
+            if ref not in used_globals:
+                if ref.startswith("__builtins_dict__") and ref in self.f_globals:
+                    used_globals[ref] = _safe_builtins_dict(self.f_globals[ref])
+                elif hasattr(_builtins, ref):
+                    used_globals[ref] = getattr(_builtins, ref)
+
         return GraphRuntimeEnv(
             bytecode=self.bytecode,
             import_sources=self.import_sources,
@@ -1102,27 +1175,26 @@ class CaptureOutput:
 
     graph_capture_output: GraphCaptureOutput
     # BackendInput can be None when dynamo didn't compile any graph (no tensor op)
-    backend_input: Optional[BackendInput]
+    backend_input: BackendInput | None
 
     def forward_callable(
         self,
         *,
-        compiled_fn: Optional[Callable[..., Any]] = None,
-        extra_globals: Optional[dict[str, Any]] = None,
+        compiled_fn: Callable[..., Any] | None = None,
+        extra_globals: dict[str, Any] | None = None,
     ) -> Callable[..., Any]:
         runtime_env = self.graph_capture_output.get_runtime_env()
         assert self.backend_input is not None
         backend_id = self.backend_input.backend_id
-        # pyrefly: ignore [bad-assignment, not-callable]
         compiled_fn = compiled_fn or self.backend_input.graph_module
         return runtime_env.forward_callable(
             backend_id,
-            compiled_fn,  # pyrefly: ignore [bad-argument-type]
+            compiled_fn,
             extra_globals=extra_globals,
         )
 
 
-def get_traced_fn(mod: Any) -> tuple[FunctionType, Optional[object]]:
+def get_traced_fn(mod: Any) -> tuple[FunctionType, object | None]:
     """
     Utility function to get the function to trace, and optionally a bound self
     object, from a callable (nn.Module, function, or method).
@@ -1137,7 +1209,6 @@ def get_traced_fn(mod: Any) -> tuple[FunctionType, Optional[object]]:
 
         resolved_call = mod.__call__
         if hasattr(resolved_call, "__self__"):
-            # pyrefly: ignore [missing-attribute]
             resolved_call = resolved_call.__func__
 
         # Mirrored from NNModuleVariable.call_function:
@@ -1179,7 +1250,7 @@ def _get_signature(fn: Any) -> inspect.Signature:
 def _get_frame(
     mod: Any,
     args: tuple[Any, ...],
-    kwargs: Optional[dict[str, Any]] = None,
+    kwargs: dict[str, Any] | None = None,
 ) -> FrameInfo:
     """
     Create a frame to trace, given a model, args, and optional kwargs.
@@ -1219,9 +1290,9 @@ def _get_frame(
 def fullgraph_capture(
     mod: Any,
     args: tuple[Any, ...],
-    kwargs: Optional[dict[str, Any]] = None,
+    kwargs: dict[str, Any] | None = None,
     *,
-    constraints: Optional[list[Constraint]] = None,
+    constraints: list[Constraint] | None = None,
     _is_export_deprecated_do_not_use: bool = False,
 ) -> CaptureOutput:
     """
@@ -1262,19 +1333,19 @@ class FrameInfo:
     locals: dict[str, object]
     builtins: dict[str, object]
     closure: tuple[CellType]
-    argdefs: Optional[tuple[Any, ...]]
-    kwdefaults: Optional[dict[str, Any]]
+    argdefs: tuple[Any, ...] | None
+    kwdefaults: dict[str, Any] | None
 
 
 def _fullgraph_capture_frame(
     frame: FrameInfo,
     *,
-    constraints: Optional[list[Constraint]] = None,
+    constraints: list[Constraint] | None = None,
     _is_export_deprecated_do_not_use: bool = False,
 ) -> CaptureOutput:
     from torch._guards import TracingContext
 
-    backend_input: Optional[BackendInput] = None
+    backend_input: BackendInput | None = None
 
     def fullgraph_compiler(
         gm: torch.fx.GraphModule, example_inputs: list[torch.Tensor]
@@ -1297,7 +1368,6 @@ def fullgraph_compiler(
             frame.locals,
             frame.builtins,
             frame.closure,
-            # pyrefly: ignore [bad-argument-type]
             compiler_fn=fullgraph_compiler,
             export=_is_export_deprecated_do_not_use,
             export_constraints=constraints,  # type: ignore[arg-type]
@@ -1305,7 +1375,7 @@ def fullgraph_compiler(
             restart_reasons=set(),
         )
         # https://github.com/pytorch/pytorch/blob/main/torch/_dynamo/eval_frame.py#L831
-    except (Unsupported, UncapturedHigherOrderOpError) as e:
+    except (Unsupported, UncapturedHigherOrderOpError, UserError) as e:
         augment_exc_message(e)
         if config.verbose:
             raise
@@ -1357,9 +1427,9 @@ def compile_frame(  # type: ignore[return]
     *,
     export: bool = False,
     export_constraints: Any | None = None,
-    frame_state: Optional[dict[str, Union[int, FrameStateSizeEntry]]] = None,
-    distributed_state: Optional[DistributedState] = None,
-    package: Optional[CompilePackage] = None,
+    frame_state: dict[str, int | FrameStateSizeEntry] | None = None,
+    distributed_state: DistributedState | None = None,
+    package: CompilePackage | None = None,
     # pyrefly: ignore [bad-return]
 ) -> DynamoOutput:
     """
@@ -1374,8 +1444,8 @@ def compile_frame(  # type: ignore[return]
     def transform(
         instructions: list[Instruction], code_options: dict[str, object]
     ) -> DynamoTracerOutput:
-        tf_mode_stack: list[torch.overrides.TorchFunctionMode] = (
-            torch.overrides._get_current_function_mode_stack()
+        tf_mode_stack: list[torch.overrides.TorchFunctionMode] = list(
+            torch_function_mode_stack_state_mgr.stack
         )
         tracer_output = trace_frame(
             code,
@@ -1465,17 +1535,17 @@ def _compile(
     export: bool,
     export_constraints: Any | None,
     hooks: Hooks,
-    cache_entry: Optional[CacheEntry],
+    cache_entry: CacheEntry | None,
     cache_size: CacheSizeRelevantForFrame,
-    frame: Optional[DynamoFrameType] = None,
-    frame_state: Optional[dict[str, Union[int, FrameStateSizeEntry]]] = None,
+    frame: DynamoFrameType | None = None,
+    frame_state: dict[str, int | FrameStateSizeEntry] | None = None,
     *,
     compile_id: CompileId,
     skip: int = 0,
-    package: Optional[CompilePackage] = None,
+    package: CompilePackage | None = None,
     # Can be used to record things for the caller, both
     # in the case of normal and exception code paths
-    convert_frame_box: Optional[ConvertFrameBox] = None,
+    convert_frame_box: ConvertFrameBox | None = None,
 ) -> ConvertFrameReturn:
     from torch.fx.experimental.validator import (
         BisectValidationException,
@@ -1489,7 +1559,7 @@ def _compile(
     @compile_time_strobelight_meta(phase_name="compile_inner")
     def compile_inner(
         code: CodeType, one_graph: bool, hooks: Hooks
-    ) -> tuple[ConvertFrameReturn, Optional[DynamoTracerOutput]]:
+    ) -> tuple[ConvertFrameReturn, DynamoTracerOutput | None]:
         with contextlib.ExitStack() as stack:
             stack.enter_context(
                 torch._dynamo.callback_handler.install_callbacks(
@@ -1497,7 +1567,12 @@ def compile_inner(
                 )
             )
             stack.enter_context(CompileTimeInstructionCounter.record())
-            return _compile_inner(code, one_graph, hooks)
+            stack.enter_context(torch_function_mode_stack_state_mgr)
+            result = _compile_inner(code, one_graph, hooks)
+            assert torch._C._len_torch_function_stack() == 0, (
+                "Torch function mode stack state changed while dynamo tracing, please report a bug"
+            )
+            return result
 
         return (
             ConvertFrameReturn(),
@@ -1528,24 +1603,45 @@ def log_bytecode(
             code.co_firstlineno,
             code,
         )
-
+        # Dump the ORIGINAL bytecode of resumption frame into TORCH_TRACE
+        # log file, so that it is parsed by tlparse tool.
+        is_resumption_frame = "torch_dynamo_resume_in" in code.co_name
+        if is_resumption_frame:
+            torch._logging.trace_structured(
+                "artifact",
+                metadata_fn=lambda: {
+                    "name": code.co_name + "_ORIGINAL_BYTECODE",
+                    "encoding": "string",
+                },
+                payload_fn=lambda: dis.Bytecode(code).dis(),
+            )
         out_code = None
+        from .graph_id_filter import get_dynamo_config_override_for_compile_id
+
+        dynamo_config_override = get_dynamo_config_override_for_compile_id(
+            compile_id, config.debug_dynamo_config_override
+        )
         try:
-            dynamo_output = compile_frame(
-                code,
-                globals,
-                locals,
-                builtins,
-                closure,
-                compiler_fn,
-                one_graph,
-                restart_reasons,
-                export=export,
-                export_constraints=export_constraints,
-                frame_state=frame_state,
-                distributed_state=distributed_state,
-                package=package,
-            )
+            with (
+                config.patch(dynamo_config_override)
+                if dynamo_config_override
+                else contextlib.nullcontext()
+            ):
+                dynamo_output = compile_frame(
+                    code,
+                    globals,
+                    locals,
+                    builtins,
+                    closure,
+                    compiler_fn,
+                    one_graph,
+                    restart_reasons,
+                    export=export,
+                    export_constraints=export_constraints,
+                    frame_state=frame_state,
+                    distributed_state=distributed_state,
+                    package=package,
+                )
         except exc.SkipFrame as e:
             if one_graph:
                 log.debug("No graph captured with export/fullgraph=True")
@@ -1568,18 +1664,42 @@ def log_bytecode(
             code.co_firstlineno,
             out_code,
         )
+        # Dump the MODIFIED bytecode of resumption frame into TORCH_TRACE
+        # log file, so that it is parsed by tlparse tool.
+        if is_resumption_frame:
+            torch._logging.trace_structured(
+                "artifact",
+                metadata_fn=lambda: {
+                    "name": code.co_name + "_MODIFIED_BYTECODE",
+                    "encoding": "string",
+                },
+                payload_fn=lambda: dis.Bytecode(out_code).dis(),
+            )
+
+        assert tracer_output.output_graph is not None
+        output = tracer_output.output_graph
+        code_context.get_context(out_code)[_BYTECODE_HOOK_SIDE_EFFECTS_CONTEXT_KEY] = (
+            tuple(output.get_replayed_side_effect_source_refs())
+        )
 
         for idx, hook in enumerate(_bytecode_hooks.values()):
             with dynamo_timed(f"bytecode_hooks_{idx}", log_pt2_compile_event=True):
                 hook_output = hook(code, out_code)
                 if hook_output is not None:
+                    if hook_output is not out_code:
+                        _copy_code_context(out_code, hook_output)
                     out_code = hook_output
 
         orig_code_map[out_code] = code
         output_codes.add(out_code)
         dynamo_time_before_restart = last_attempt_start_time - start_time
-        assert tracer_output.output_graph is not None
-        output = tracer_output.output_graph
+
+        from .bytecode_debugger import BREAKPOINT_MARKER
+
+        if BREAKPOINT_MARKER in out_code.co_consts:
+            from torch._C._dynamo.eval_frame import register_breakpoint_code
+
+            register_breakpoint_code(out_code)
 
         # Tests for new code objects.
         # The rationale for these tests can be found in torch/csrc/dynamo/eval_frame.c
@@ -1630,7 +1750,16 @@ def count_args(code: CodeType) -> int:
         assert output.guards is not None
         CleanupManager.instance[out_code] = output.cleanups
         nonlocal cache_entry
-        with dynamo_timed("build_guards", log_pt2_compile_event=True):
+        # Temporarily restore the mode stack so guard expressions that
+        # reference modes can evaluate.  DisableTorchFunction prevents
+        # __torch_function__ dispatch during guard construction so modes
+        # with mutable state aren't triggered.
+        build_guards_ctx = contextlib.ExitStack()
+        if torch_function_mode_stack_state_mgr.stack:
+            build_guards_ctx.enter_context(
+                torch_function_mode_stack_state_mgr.temp_restore_stack()
+            )
+        with dynamo_timed("build_guards", log_pt2_compile_event=True), build_guards_ctx:
             check_fn = dynamo_output.build_guards(
                 code,
                 hooks=hooks,
@@ -1664,7 +1793,7 @@ def count_args(code: CodeType) -> int:
         return wrap_guarded_code(guarded_code), tracer_output
 
     metrics_context = get_metrics_context()
-    code_context = (
+    package_code_context = (
         package.code_context(code) if package is not None else contextlib.nullcontext()
     )
     with (
@@ -1680,7 +1809,7 @@ def count_args(code: CodeType) -> int:
             phase_name="entire_frame_compile",
             dynamo_compile_column_us="dynamo_cumulative_compile_time_us",
         ),
-        code_context,
+        package_code_context,
     ):
         restart_reasons: set[str] = set()
         if compile_pg := get_compile_pg():
@@ -1689,7 +1818,7 @@ def count_args(code: CodeType) -> int:
             distributed_state = None
 
         # Check recompilations
-        recompile_reason: Optional[str] = None
+        recompile_reason: str | None = None
         if is_recompilation(cache_size) and frame:
             reasons = get_and_maybe_log_recompilation_reasons(
                 cache_entry, frame, innermost_fn(compiler_fn)
@@ -1697,29 +1826,10 @@ def count_args(code: CodeType) -> int:
             recompile_reason = (
                 "Unable to find recompilation reasons" if not reasons else reasons[0]
             )
-        # Recheck for recompilation, for when inline_inbuilt_nn_modules is set to False
-        inline_inbuilt_nn_modules_candidate = False
-        if not config.inline_inbuilt_nn_modules and frame:
-            inbuilt_nn_reasons = get_and_maybe_log_recompilation_reasons(
-                cache_entry, frame, innermost_fn(compiler_fn), skip_logging=True
-            )
-            inbuilt_nn_recompile_reason = (
-                None if not inbuilt_nn_reasons else inbuilt_nn_reasons[0]
-            )
-
-            if (
-                inbuilt_nn_recompile_reason is not None
-                and "[inline-inbuilt-nn-modules-candidate]"
-                in inbuilt_nn_recompile_reason
-            ):
-                inline_inbuilt_nn_modules_candidate = True
-
-        # Set if the recompile is a candidate for inline_inbuilt_nn_modules
-        # regardless of whether inline_inbuilt_nn_modules is set or not
         metrics_context.update_outer(
             {
                 "recompile_reason": recompile_reason,
-                "inline_inbuilt_nn_modules_candidate": inline_inbuilt_nn_modules_candidate,
+                "inline_inbuilt_nn_modules_candidate": False,
             }
         )
 
@@ -1815,14 +1925,38 @@ def raise_unimplemented_cache_limit_exceeded() -> NoReturn:
         # torch/_dynamo/convert_frame.py:780 in <lambda>
         stack_trace = log_dynamo_start(code, skip)
         start_time_ns = time.time_ns()
-        fail_type: Optional[str] = None
-        fail_reason: Optional[str] = None
-        exception_stack_trace: Optional[list[str]] = None
-        fail_user_frame_filename: Optional[str] = None
-        fail_user_frame_lineno: Optional[int] = None
+        fail_type: str | None = None
+        fail_reason: str | None = None
+        exception_stack_trace: list[str] | None = None
+        fail_user_frame_filename: str | None = None
+        fail_user_frame_lineno: int | None = None
         torch._dynamo.utils.ReinplaceCounters.clear()
         guarded_code = None
         tracer_output = None
+
+        if (
+            config.debug_backend_override
+            or config.debug_dynamo_config_override
+            or config.debug_inductor_config_override
+        ):
+            # Eagerly validate config override strings before entering the
+            # compilation try/except so that typos surface as clean ValueErrors
+            # instead of being wrapped as InternalTorchDynamoError.
+            from .graph_id_filter import (
+                _validate_backend_names,
+                _validate_dynamo_config_keys,
+                _validate_inductor_config_keys,
+            )
+
+            if err := _validate_backend_names(config.debug_backend_override):
+                raise ValueError(err)
+            if err := _validate_dynamo_config_keys(config.debug_dynamo_config_override):
+                raise ValueError(err)
+            if err := _validate_inductor_config_keys(
+                config.debug_inductor_config_override
+            ):
+                raise ValueError(err)
+
         try:
             guarded_code, tracer_output = compile_inner(code, one_graph, hooks)
 
@@ -1871,6 +2005,7 @@ def raise_unimplemented_cache_limit_exceeded() -> NoReturn:
                 e,
                 (
                     Unsupported,
+                    UserError,
                     TorchRuntimeError,
                     BackendCompilerFailed,
                     AssertionError,
@@ -1995,7 +2130,7 @@ def raise_unimplemented_cache_limit_exceeded() -> NoReturn:
                 and tracer_output.output_graph is not None
                 and not tracer_output.output_graph.export
             ):
-                tracer_output.output_graph.tracing_context.guards_context.dynamo_guards.inner = OrderedSet()
+                tracer_output.output_graph.tracing_context.guards_context.dynamo_guards.clear()
 
             # Clear WeakIdRef entries that can block swap_tensors after compile.
             # Determine whether to clear based on config and backend type.
@@ -2022,28 +2157,33 @@ def __init__(
         self,
         compiler_fn: CompilerFn,
         hooks: Hooks,
-        package: Optional[CompilePackage] = None,
+        package: CompilePackage | None = None,
+        recompile_limit: int | None = None,
     ) -> None:
         self._torchdynamo_orig_backend = compiler_fn
         self._inner_convert = convert_frame_assert(
-            compiler_fn, one_graph=False, package=package
+            compiler_fn,
+            one_graph=False,
+            package=package,
+            recompile_limit=recompile_limit,
         )
         self._hooks = hooks
+        self._recompile_limit = recompile_limit
 
     @property
     def _clone_with_backend(self) -> Callable[[WrapBackendDebug], ConvertFrame]:
         return lambda backend: convert_frame(
-            # pyrefly: ignore [bad-argument-type]
             backend,
             self._hooks,
+            recompile_limit=self._recompile_limit,
         )
 
     def __call__(
         self,
         frame: DynamoFrameType,
-        cache_entry: Optional[CacheEntry],
+        cache_entry: CacheEntry | None,
         hooks: Hooks,
-        frame_state: dict[str, Union[int, FrameStateSizeEntry]],
+        frame_state: dict[str, int | FrameStateSizeEntry],
         skip: int = 0,
     ) -> ConvertFrameReturn:
         input_codes.add(frame.f_code)
@@ -2088,7 +2228,7 @@ def __call__(
             if isinstance(e, UncapturedHigherOrderOpError):
                 raise
 
-            soft_fail = isinstance(e, Unsupported)
+            soft_fail = isinstance(e, (Unsupported, UserError))
             code = frame.f_code
             # Log soft failure that was not already logged by symbolic_convert.
             # This happens e.g. for graph breaks that are raised in convert_frame.py
@@ -2159,10 +2299,13 @@ def __call__(
 def convert_frame(
     compiler_fn: CompilerFn,
     hooks: Hooks,
-    package: Optional[CompilePackage] = None,
+    package: CompilePackage | None = None,
+    recompile_limit: int | None = None,
 ) -> ConvertFrame:
     """Try to convert a frame into an FX graph, if error leave frame unmodified"""
-    return ConvertFrame(compiler_fn, hooks, package=package)
+    return ConvertFrame(
+        compiler_fn, hooks, package=package, recompile_limit=recompile_limit
+    )
 
 
 # TODO mlazos: add support for same args, or record them
@@ -2211,9 +2354,9 @@ class ConvertFrameProtocol(typing.Protocol):
     def __call__(
         self,
         frame: DynamoFrameType,
-        cache_entry: Optional[CacheEntry],
+        cache_entry: CacheEntry | None,
         hooks: Hooks,
-        frame_state: dict[str, Union[int, FrameStateSizeEntry]],
+        frame_state: dict[str, int | FrameStateSizeEntry],
         *,
         skip: int = 0,
     ) -> ConvertFrameReturn: ...
@@ -2228,13 +2371,13 @@ def __init__(self, callback: ConvertFrameProtocol, hooks: Hooks) -> None:
     def __call__(
         self,
         frame: DynamoFrameType,
-        cache_entry: Optional[CacheEntry],
-        frame_state: dict[str, Union[int, FrameStateSizeEntry]],
+        cache_entry: CacheEntry | None,
+        frame_state: dict[str, int | FrameStateSizeEntry],
     ) -> ConvertFrameReturn:
         assert frame_state is not None
         input_codes.add(frame.f_code)
 
-        is_skipfile = trace_rules.check(frame.f_code)
+        is_skipfile = trace_rules.check(frame.f_code, frame=frame)
         if sys.version_info >= (3, 13):
             has_started_execution = frame.f_lasti > first_real_inst_idx(frame.f_code)
         else:
@@ -2264,7 +2407,7 @@ def __call__(
             if log.isEnabledFor(logging.DEBUG):
                 if has_started_execution:
                     skip_reason = "traced frame already"
-                elif trace_rules.check(frame.f_code):
+                elif trace_rules.check(frame.f_code, frame=frame):
                     skip_reason = "in skipfiles"
                 elif should_skip_for_dispatch_mode:
                     skip_reason = "non-infra torch dispatch mode present, this is not supported today in torch.compile"
@@ -2287,6 +2430,21 @@ def __call__(
         ):
             # nametuple constructor/_make
             return ConvertFrameReturn()
+
+        if (
+            frame.f_code.co_name == "__init__"
+            and frame.f_code.co_argcount > 0
+            and frame.f_code.co_varnames
+            and is_traceable_wrapper_subclass(
+                frame.f_locals.get(frame.f_code.co_varnames[0])
+            )
+        ):
+            # Skip tracing __init__ of traceable wrapper subclasses: self is
+            # partially initialized at this point (attributes set by __init__
+            # don't exist yet), so faking it would call __tensor_flatten__ and
+            # crash. Run eagerly instead, matching @torch._disable_dynamo behavior.
+            return ConvertFrameReturn()
+
         if torch._dynamo.utils.get_optimize_ddp_mode() == "ddp_optimizer":
             ddp_module = DistributedDataParallel._get_active_ddp_module()
             if ddp_module:
diff --git a/torch/_dynamo/debug_utils.py b/torch/_dynamo/debug_utils.py
index ac862fb47ff73..e43f45a138e2d 100644
--- a/torch/_dynamo/debug_utils.py
+++ b/torch/_dynamo/debug_utils.py
@@ -34,7 +34,7 @@
 import textwrap
 from collections import Counter
 from importlib import import_module
-from typing import Any, Optional, TYPE_CHECKING, TypeVar
+from typing import Any, TYPE_CHECKING, TypeVar
 
 import torch
 import torch._prims_common as utils
@@ -544,14 +544,14 @@ def backend_accuracy_fails(
 
 
 def _stride_or_default(
-    stride: Optional[torch._prims_common.StrideType],
+    stride: torch._prims_common.StrideType | None,
     *,
     shape: torch._prims_common.ShapeType,
 ) -> torch._prims_common.StrideType:
     return stride if stride is not None else utils.make_contiguous_strides_for(shape)
 
 
-def _mk_defaulter(d: T) -> Callable[[Optional[T]], T]:
+def _mk_defaulter(d: T) -> Callable[[T | None], T]:
     return lambda x: x if x is not None else d
 
 
@@ -568,20 +568,32 @@ def __init__(self) -> None:
 
     def storage(
         self,
-        storage_hash: Optional[str],
+        storage_hash: str | None,
         nbytes: int,
         *,
-        device: Optional[torch._prims_common.DeviceLikeType] = None,
-        dtype_hint: Optional[torch.dtype] = None,
+        device: torch._prims_common.DeviceLikeType | None = None,
+        dtype_hint: torch.dtype | None = None,
     ) -> None:
         self.total += 1
 
-    def tensor(self, *args: Any, **kwargs: Any) -> Optional[torch.Tensor]:
+    def tensor(self, *args: Any, **kwargs: Any) -> torch.Tensor | None:
         pass
 
-    def symint(self, *args: Any, **kwargs: Any) -> Optional[int]:
+    def symint(self, *args: Any, **kwargs: Any) -> int | None:
         pass
 
+    def const(self, name: str) -> None:
+        pass
+
+    def unsupported(self, name: str) -> None:
+        pass
+
+    def generator(self, device_type: str, device_index: int) -> None:
+        pass
+
+    def opaque(self, script_class_name: str) -> None:
+        self.total += 1
+
 
 # TODO: Support bundling the entire repro into a zip file for ease of
 # transferring around
@@ -601,11 +613,11 @@ def __init__(
 
     def storage(
         self,
-        storage_hash: Optional[str],
+        storage_hash: str | None,
         nbytes: int,
         *,
-        device: Optional[torch._prims_common.DeviceLikeType] = None,
-        dtype_hint: Optional[torch.dtype] = None,
+        device: torch._prims_common.DeviceLikeType | None = None,
+        dtype_hint: torch.dtype | None = None,
     ) -> UntypedStorage:
         if self.pbar is not None:
             self.pbar.update(1)
@@ -632,12 +644,12 @@ def tensor(
         self,
         storage: UntypedStorage,
         shape: torch._prims_common.ShapeType,
-        stride: Optional[torch._prims_common.StrideType] = None,
+        stride: torch._prims_common.StrideType | None = None,
         *,
-        storage_offset: Optional[int] = None,
-        dtype: Optional[torch.dtype] = None,
-        requires_grad: Optional[bool] = None,
-        is_leaf: Optional[bool] = None,
+        storage_offset: int | None = None,
+        dtype: torch.dtype | None = None,
+        requires_grad: bool | None = None,
+        is_leaf: bool | None = None,
         **metadata: Any,
     ) -> torch.Tensor:
         stride = _stride_or_default(stride, shape=shape)
@@ -665,6 +677,20 @@ def symint(self, val: Any) -> Any:
         self.args.append(val)
         return val  # for BC
 
+    def const(self, name: str) -> None:
+        self.args.append(None)
+
+    def unsupported(self, name: str) -> None:
+        self.args.append(None)
+
+    def generator(self, device_type: str, device_index: int) -> torch._C.Generator:
+        gen = torch.cuda.default_generators[device_index].clone_state()
+        self.args.append(gen)
+        return gen
+
+    def opaque(self, script_class_name: str) -> None:
+        self.args.append(None)
+
 
 # Here is our writer strategy:
 #  1. We will stream all of the inputs to disk
@@ -679,7 +705,7 @@ def symint(self, val: Any) -> Any:
 
 
 class InputWriter:
-    def __init__(self, save_dir: Optional[str], *, stable_hash: bool = False) -> None:
+    def __init__(self, save_dir: str | None, *, stable_hash: bool = False) -> None:
         self._lines: list[str] = []
         # TODO: consider ensuring tensor and storage counters line up?
         self.storage_counter = itertools.count()
@@ -710,8 +736,8 @@ def storage(
         self,
         untyped_storage: UntypedStorage,
         *,
-        device_hint: Optional[torch._prims_common.DeviceLikeType] = None,
-        dtype_hint: Optional[torch.dtype] = None,
+        device_hint: torch._prims_common.DeviceLikeType | None = None,
+        dtype_hint: torch.dtype | None = None,
     ) -> str:
         ws = StorageWeakRef(untyped_storage)
         v = self.seen_storages.get(ws)
@@ -774,7 +800,9 @@ def tensor(self, name: str, t: torch.Tensor) -> None:
 
     def unsupported(self, name: str, arg: Any) -> None:
         # NB: Try hard not to /print/ a tensor, that will be very slow
-        self._lines.append(f"# {name} was unsupported type for dumping: {type(arg)}")
+        self._lines.append(
+            f"reader.unsupported({name!r})  # unsupported type for dumping: {type(arg)}"
+        )
         # Best effort dump as much useful stuff we can lol, in case you want
         # to repair the repro
         if isinstance(arg, (list, tuple)):
@@ -801,12 +829,21 @@ def symint(self, name: str, val: Any) -> None:
             val = val.node.hint
         self._lines.append(f"reader.symint({val!r})  # {name}")
 
+    def generator(self, name: str, arg: torch._C.Generator) -> None:
+        device = arg.device
+        self._lines.append(
+            f"reader.generator({device.type!r}, {device.index!r})  # {name}"
+        )
+
+    def opaque(self, name: str, script_class_name: str) -> None:
+        self._lines.append(f"reader.opaque({script_class_name!r})  # {name}")
+
 
 def aot_graph_input_parser(
     func: Callable[[list[Tensor]], list[Tensor]],
     device: str = "cuda",
-    sym_shapes: Optional[dict[str, int]] = None,
-    default_sym_shape: Optional[int] = None,
+    sym_shapes: dict[str, int] | None = None,
+    default_sym_shape: int | None = None,
 ) -> dict[str, Any]:
     """
     Takes in a function which has been printed with print_readable() and constructs kwargs to run it.
diff --git a/torch/_dynamo/decorators.py b/torch/_dynamo/decorators.py
index ed1354555a7d7..bf06b84d6b15b 100644
--- a/torch/_dynamo/decorators.py
+++ b/torch/_dynamo/decorators.py
@@ -8,10 +8,12 @@
 from collections.abc import Callable
 from dataclasses import dataclass
 from types import TracebackType
-from typing import Any, Optional, overload, TYPE_CHECKING, TypeVar, Union
+from typing import Any, overload, TYPE_CHECKING, TypeVar
 from typing_extensions import ParamSpec
 
 import torch
+import torch.utils._pytree as pytree
+from torch._opaque_base import OpaqueBase
 from torch.compiler import is_compiling
 from torch.utils._contextlib import _DecoratorContextManager
 from torch.utils._python_dispatch import is_traceable_wrapper_subclass
@@ -27,7 +29,6 @@
     RunOnlyContext,
     skip_code,
 )
-from .exc import IncorrectUsage
 from .external_utils import (
     get_nonrecursive_disable_wrapper,
     wrap_dunder_call_ctx_manager,
@@ -62,7 +63,7 @@
 F = TypeVar("F", bound=FuncType)
 
 
-def run(fn: Optional[Callable[_P, _R]] = None) -> Any:
+def run(fn: Callable[_P, _R] | None = None) -> Any:
     """Don't do any dynamic compiles, just use prior optimizations"""
     if fn is not None:
         fn = innermost_fn(fn)
@@ -103,7 +104,6 @@ def wrap(fn: Callable[_P, _R]) -> Callable[_P, _R]:
                 nonrecursive_disable_wrapper
             )
             nonrecursive_disable_wrapper._torchdynamo_disable_recursive = False  # type: ignore[attr-defined]
-            # pyrefly: ignore [bad-return]
             return nonrecursive_disable_wrapper
 
         if fn is None:
@@ -115,7 +115,7 @@ def wrap(fn: Callable[_P, _R]) -> Callable[_P, _R]:
 skip_code(_nonrecursive_disable_wrapper_code)
 
 
-def skip(fn: Optional[Callable[_P, _R]] = None) -> Callable[..., Any]:
+def skip(fn: Callable[_P, _R] | None = None) -> Callable[..., Any]:
     """
     Skip frames associated with the function code, but still process recursively
     invoked frames
@@ -143,7 +143,7 @@ def __init__(
         stance: str = "default",
         *,
         skip_guard_eval_unsafe: bool = False,
-        force_backend: Union[str, Callable[..., Any], None] = None,
+        force_backend: str | Callable[..., Any] | None = None,
     ) -> None:
         if force_backend is not None and stance != "default":
             raise RuntimeError("non-default stance cannot have force_backend set")
@@ -220,23 +220,69 @@ def _check_mutually_exclusive_decorators(fn: Callable, decorator_name: str) -> N
 
 
 def nonstrict_trace(traceable_fn: Callable[_P, _R]) -> Callable[_P, _R]:
-    # Like `allow_in_graph`, but with the following enhancements/differences:
-    #
-    # 1. Supports user-defined class as inputs, as long as the class has been
-    #    registered with pytree.
-    # 2. Reads to global/captured tensors forces the underlying graph to treat
-    #    those tensors as constant, and we _assume_ they will not be updated. This
-    #    is similar to FX tracing.
-    # 3. In the resulting Dynamo graph, the call to a `nonstrict_trace`-ed function
-    #    will be represented as a call to `torch._higher_order_ops.flat_apply`,
-    #    which takes in the `nonstrict_trace`-ed function and pytree-flattened
-    #    inputs.
-    # 4. Only the returned function is traceable, and the original function will
-    #    not be. Moreover, `nonstrict_trace` can be used inside a `torch.compile`
-    #    region.
-    #
-    # NOTE: like `allow_in_graph`, aliasing information is neither preserved
-    # between inputs themselves, nor between inputs and outputs.
+    """
+    Decorator to mark a function as nonstrict-traceable for dynamo.
+
+    A nonstrict-traced function appears as an opaque call in the dynamo graph.
+    Dynamo does not trace into the function body (hence the "nonstrict"), but
+    aot_autograd will trace into it.
+
+    This is similar to ``allow_in_graph`` but with enhanced support for:
+    - User-defined classes as inputs (must be registered with pytree)
+    - ``nn.Module`` as input arguments (parameters and buffers are tracked for autograd)
+    - Global/captured tensors treated as constants (assumed not updated during execution)
+
+    Note:
+        - With ``backend="eager"``, the original Python function runs directly.
+          With ``backend="aot_eager"``, the graph traced by aot_autograd runs.
+          With ``backend="inductor"``, the traced graph is compiled with inductor.
+
+        - Training is supported: you can call ``.backward()`` on outputs and gradients
+          will flow through the nonstrict-traced function.
+
+    Dangerous patterns (may cause silent incorrectness):
+        - Side effects between nonstric_traced fn and compiled region: The function should
+          not depend on variables mutated by other code inside the compiled function, and code
+          after the call should not depend on mutations made by it.
+
+        - Implicit inputs (closures/globals): Tensors captured from enclosing scopes
+          are treated as constants. Gradients will NOT flow back to them. Pass tensors
+          as explicit arguments if gradients are needed.
+
+    Restrictions:
+        - Both inputs and outputs must use pytree-compatible types. User-defined classes
+          must be registered via :func:`torch.utils._pytree.register_pytree_node`,
+          :func:`torch.utils._pytree.register_dataclass`, or
+          :func:`torch.utils._pytree.register_constant`. Tensors, Python primitives (int, float, bool, str),
+          symbolic types (SymInt, SymFloat, SymBool), and built-in containers (list,
+          tuple, dict) are already handled by default.
+        - Primitive values and container structure are specialized per call site:
+          each call site expects the same primitives and structure on every execution.
+
+    Example::
+
+        >>> import torch
+        >>> @torch._dynamo.nonstrict_trace
+        ... def traced_forward(model, x):
+        ...     # It's OK to have dynamo graph break within nonstrict_trace region
+        ...     torch._dynamo.graph_break()
+        ...     return model(x) + x
+        ...
+        >>> class MyModule(torch.nn.Module):
+        ...     def __init__(self):
+        ...         super().__init__()
+        ...         self.inner = torch.nn.Linear(10, 10)
+        ...
+        ...     def forward(self, x):
+        ...         return traced_forward(self.inner, x)
+        ...
+        >>> # Compile and run
+        >>> model = MyModule()
+        >>> opt_model = torch.compile(model, backend="aot_eager", fullgraph=True)
+        >>> out = opt_model(torch.randn(10, 10))
+        >>> out.sum().backward()  # Gradients flow through traced_forward
+
+    """
     assert callable(traceable_fn), "nonstrict_trace expects a callable"
 
     _check_mutually_exclusive_decorators(traceable_fn, "nonstrict_trace")
@@ -263,7 +309,84 @@ def deregister() -> None:
     return wrapped
 
 
-def leaf_function(fn: Callable[_P, _R]) -> Callable[_P, _R]:
+def _invoke_leaf_function_python(
+    real_impl: Callable[..., Any],
+    fake_impl: Callable[..., Any],
+    args: tuple[Any, ...],
+    kwargs: dict[str, Any],
+    mutates_args: frozenset[str] | None = None,
+) -> Any:
+    """Call invoke_leaf_function HOP directly from Python.
+
+    This enables @leaf_function to work with make_fx
+    without relying on Dynamo to intercept the call.
+    """
+    from torch._higher_order_ops.invoke_leaf_function import (
+        _LeafCallable,
+        convert_modules_to_states,
+        invoke_leaf_function,
+        make_leaf_function_wrappers,
+        store_makefx_modules,
+    )
+
+    captured_modules: list[torch.nn.Module] = []
+    seen_module_ids: dict[int, int] = {}  # id(module) -> position in captured_modules
+    for val in pytree.tree_flatten(
+        (args, kwargs), is_leaf=lambda x: isinstance(x, torch.nn.Module)
+    )[0]:
+        if isinstance(val, torch.nn.Module) and id(val) not in seen_module_ids:
+            seen_module_ids[id(val)] = len(captured_modules)
+            captured_modules.append(val)
+
+    global_indices = store_makefx_modules(captured_modules)
+    module_to_index = {
+        mod_id: global_indices[pos] for mod_id, pos in seen_module_ids.items()
+    }
+
+    processed = convert_modules_to_states((args, kwargs), module_to_index)
+    flat_args, input_spec = pytree.tree_flatten(processed)
+
+    # Single-element mutable list so the wrappers can write back the output
+    # TreeSpec. Read captured_out_spec[0] after the wrappers have been called.
+    captured_out_spec: list[pytree.TreeSpec | None] = [None]
+    wrapped_real, wrapped_fake = make_leaf_function_wrappers(
+        real_impl, fake_impl, captured_out_spec
+    )
+
+    real_fn_callable = _LeafCallable(wrapped_real)
+    fake_fn_callable = _LeafCallable(wrapped_fake)
+
+    mutated_flat_indices = ""
+    if mutates_args:
+        from torch._higher_order_ops.invoke_leaf_function import (
+            _resolve_mutated_flat_indices,
+        )
+
+        mutated_flat_indices = _resolve_mutated_flat_indices(
+            real_impl, mutates_args, len(flat_args), input_spec
+        )
+
+    flat_out = invoke_leaf_function(
+        real_fn_callable, fake_fn_callable, input_spec, mutated_flat_indices, *flat_args
+    )
+
+    assert captured_out_spec[0] is not None
+    return pytree.tree_unflatten(flat_out, captured_out_spec[0])
+
+
+@overload
+def leaf_function(fn: Callable[_P, _R]) -> Callable[_P, _R]: ...
+
+
+@overload
+def leaf_function(
+    *, mutates_args: set[str]
+) -> Callable[[Callable[_P, _R]], Callable[_P, _R]]: ...
+
+
+def leaf_function(
+    fn: Callable[_P, _R] | None = None, *, mutates_args: set[str] | None = None
+) -> Callable[_P, _R] | Callable[[Callable[_P, _R]], Callable[_P, _R]]:
     """
     Decorator to mark a function as a leaf function for :func:`torch.compile`.
 
@@ -326,17 +449,10 @@ def compute_and_log_stats(x):
         to trace through and potentially optimize the code.
 
     Usage:
-        **Supported Inputs**:
-        - Inputs must use pytree-compatible types: tensors, Python primitives
-        (int, float, bool, str), and built-in containers (list, tuple, dict).
-        User-defined classes must be registered as pytree nodes via
-        :func:`torch.utils.pytree.register_pytree_node`.
-        - :class:`torch.nn.Module` can also be passed as input; its parameters and buffers are
-        tracked for autograd. The module must exist outside the compile region.
-
-        **Supported Outputs**: Must be a tuple of tensors: ``return (tensor,)`` for one tensor,
-        ``return (a, b)`` for multiple. Primitive types (int, float, bool, str) can also be
-        included in outputs.
+        **Inputs and Outputs**:
+        - Both inputs and outputs must use pytree-compatible types.
+        Tensors, Python primitives (int, float, bool, str), and built-in containers
+        (list, tuple, dict) are supported by default.
 
         Note: We recommend leaf_functions only accept and return tensors. Though primitive
         types (int, float, bool, str) are supported in inputs and outputs, they may cause
@@ -367,7 +483,15 @@ def count_calls_fake(x):
             # At runtime: real function runs (counter increments to 1, 2, 3, ...)
             # But returned count is always 999 (from fake implementation at compile time)
 
-        **register_fake (required)**: Since the function body is not traced, you must
+        - User-defined classes must be registered via :func:`torch.utils._pytree.register_pytree_node`
+        or :func:`torch.utils._pytree.register_dataclass`.
+
+        - :class:`torch.nn.Module` can also be passed as input; its parameters and buffers
+        are tracked for autograd. The module must exist outside the compile region.
+        Running the module multiple times must not modify its parameters, buffers, or attributes.
+
+        **register_fake (required)**:
+        Since the function body is not traced, you must
         provide a shape-inference function via ``@fn.register_fake``. It runs at compile
         time with FakeTensor inputs (tensors with no data, only metadata) and must
         satisfy the following requirements:
@@ -438,16 +562,40 @@ def count_calls_fake(x):
           Logging/printing inside the leaf function is fine since it doesn't affect
           correctness.
 
-        - **In-place mutations on inputs**: In-place mutations on input tensors are
-          detected and will raise an error. Clone inputs before mutating.
+        - **In-place mutations on inputs**: Undeclared in-place mutations on input
+          tensors are detected and will raise an error. Either declare mutations via
+          ``mutates_args`` or clone inputs before mutating.
 
           Bad::
 
             @leaf_function
             def my_leaf_fn(x):
-                x.add_(1)  # Will raise: "In-place mutation detected"
+                x.add_(1)  # Will raise: "Undeclared in-place mutation"
                 return (x,)
 
+          Good::
+
+            @leaf_function(mutates_args={"buf"})
+            def my_leaf_fn(x, buf):
+                buf.add_(1)  # OK: declared in mutates_args
+                return (x + buf,)
+
+          For pytree args (lists, tuples, dicts of tensors), use the parameter name
+          to declare mutation on all contained tensors::
+
+            @leaf_function(mutates_args={"buffers"})
+            def my_leaf_fn(x, buffers):
+                for buf in buffers:
+                    buf.add_(1)  # OK: 'buffers' declared in mutates_args
+                return (x + sum(buffers),)
+
+          Or use bracket notation for fine-grained control::
+
+            @leaf_function(mutates_args={"buffers[0]"})
+            def my_leaf_fn(x, buffers):
+                buffers[0].add_(1)  # OK: 'buffers[0]' declared
+                return (x + buffers[1],)  # buffers[1] is not cloned
+
         - **Closures in fake implementation**: Tensors or modules captured from enclosing scopes
           in the fake implementation will cause compilation errors. The real function can
           close over them if they don't require gradient, but the fake implementation must
@@ -517,17 +665,57 @@ def my_leaf_fn_fake(x):
 
     Args:
         fn: The function being decorated.
+        mutates_args: Set of Python expressions (as strings) identifying arguments
+            that the function mutates in-place. Each string is evaluated against
+            the function's parameters. Examples: ``'buf'`` for a plain tensor,
+            ``'model.running_mean'`` for an nn.Module buffer,
+            ``'buffers'`` to mark all tensors in a list, ``'buffers[0]'`` for a
+            specific element, ``'state["key"]'`` for a dict entry.
     """
+    if fn is None:
+        return functools.partial(leaf_function, mutates_args=mutates_args)
+
     from . import trace_rules
 
     _check_mutually_exclusive_decorators(fn, "leaf_function")
 
+    if mutates_args:
+        import inspect
+        import re
+
+        params = set(inspect.signature(fn).parameters)
+        for expr in mutates_args:
+            root = re.split(r"[.\[]", expr, maxsplit=1)[0]
+            if root not in params:
+                raise ValueError(
+                    f"mutates_args expression '{expr}' refers to parameter '{root}', "
+                    f"which is not a parameter of '{fn.__name__}'. "
+                    f"Available parameters: {', '.join(params)}"
+                )
+
     @functools.wraps(fn)
     def inner(*args: _P.args, **kwargs: _P.kwargs) -> _R:
-        return fn(*args, **kwargs)
+        if inner._torchdynamo_leaf_fake_fn is None:  # type: ignore[attr-defined]
+            raise ValueError(
+                f"leaf_function '{getattr(fn, '__name__', fn)}' "
+                "requires a fake implementation. Please provide one using the @<func>.register_fake "
+                "decorator. See the leaf_function docstring for details."
+            )
+
+        return _invoke_leaf_function_python(
+            fn,  # pyrefly: ignore [bad-argument-type]
+            # pyrefly: ignore [bad-argument-type]
+            inner._torchdynamo_leaf_fake_fn,
+            args,
+            kwargs,
+            mutates_args=inner._torchdynamo_leaf_mutates_args,  # pyrefly: ignore [missing-attribute]
+        )  # type: ignore[attr-defined]
 
     inner._torchdynamo_leaf_real_fn = fn  # type: ignore[attr-defined]
     inner._torchdynamo_leaf_fake_fn = None  # type: ignore[attr-defined]
+    inner._torchdynamo_leaf_mutates_args = (  # pyrefly: ignore [missing-attribute]
+        frozenset(mutates_args) if mutates_args else frozenset()
+    )  # type: ignore[attr-defined]
 
     # Follow nonstrict_trace implementation
     wrapped_id = id(inner)
@@ -560,7 +748,7 @@ def inner(fn: Any) -> Any:
             != variables.TorchInGraphFunctionVariable
             and trace_rules.lookup(fn) != variables.TorchInGraphFunctionVariable
         ):
-            raise IncorrectUsage(
+            raise RuntimeError(
                 "disallow_in_graph is expected to be used on an already allowed callable (like torch.* ops). "
                 "Allowed callables means callables that TorchDynamo puts as-is in the extracted graph."
             )
@@ -844,9 +1032,16 @@ def _apply_func_to_inner_tensors_of_same_dim(
     attrs, _ctx = t.__tensor_flatten__()
     assert isinstance(t, torch.Tensor)
     for attr in attrs:
-        inner = getattr(t, attr)
-        if inner.dim() == t.dim():
-            func(inner, *args, **kwargs)
+        match getattr(t, attr):
+            case torch.Tensor() as inner:
+                if inner.dim() == t.dim():
+                    func(inner, *args, **kwargs)
+            case OpaqueBase():
+                pass
+            case unexpected:
+                raise AssertionError(
+                    f"expected Tensor or OpaqueBase, got {type(unexpected)}"
+                )
 
 
 @dataclass(frozen=True, slots=True)
@@ -865,11 +1060,13 @@ class directly; instead, use :func:`mark_dynamic`.
 @forbid_in_graph
 def mark_unbacked(
     t: Any,
-    index: Union[int, list[Any], tuple[Any]],
-    hint_override: Optional[int] = None,
+    index: int | list[Any] | tuple[Any],
+    hint_override: int | None = None,
     strict: bool = False,
-    specialize_on: Optional[list[Any]] = None,
-    shape_id: Optional[str] = None,
+    specialize_on: list[Any] | None = None,
+    shape_id: str | None = None,
+    min: int | None = None,
+    max: int | None = None,
 ) -> None:
     """
     Mark a tensor as having an unbacked dimension. This changes the semantics of operations:
@@ -894,6 +1091,10 @@ def mark_unbacked(
             All unbacked dimensions with the same shape_id will share the same unbacked symbol. This is useful when multiple tensors
             are known to have the same batch size at runtime. A runtime assertion is added
             to ensure this property at runtime.
+        min (Optional[int], default=None): Minimum value constraint for this dimension.
+            If provided, a runtime check will be added to ensure the dimension is >= min.
+        max (Optional[int], default=None): Maximum value constraint for this dimension.
+            If provided, a runtime check will be added to ensure the dimension is <= max.
     """
     if torch.distributed.is_available() and isinstance(
         t, torch.distributed.tensor.DTensor
@@ -927,6 +1128,12 @@ def mark_unbacked(
         if hint_override:
             t._dynamo_hint_overrides[index] = hint_override
 
+        if min is not None or max is not None:
+            if not hasattr(t, "_dynamo_unbacked_bounds"):
+                # pyrefly: ignore [implicit-any]
+                t._dynamo_unbacked_bounds = {}
+            t._dynamo_unbacked_bounds[index] = (min, max)
+
         if shape_id is not None:
             if not hasattr(t, "_dynamo_shape_ids"):
                 # pyrefly: ignore [implicit-any]
@@ -944,18 +1151,18 @@ def mark_unbacked(
 
     assert isinstance(index, (list, tuple))
     for i in index:
-        mark_unbacked(t, i, shape_id=shape_id)
+        mark_unbacked(t, i, shape_id=shape_id, min=min, max=max)
 
 
 @forbid_in_graph
 def mark_dynamic(
     t: Any,
-    index: Union[int, list[Any], tuple[Any]],
+    index: int | list[Any] | tuple[Any],
     *,
-    hint_override: Optional[int] = None,
-    min: Optional[int] = None,
-    max: Optional[int] = None,
-    specialize_on: Optional[list[Any]] = None,
+    hint_override: int | None = None,
+    min: int | None = None,
+    max: int | None = None,
+    specialize_on: list[Any] | None = None,
 ) -> None:
     """
     Mark a tensor as having a dynamic dim and set corresponding min and max range for the dim.
@@ -1042,7 +1249,7 @@ def mark_dynamic(
 
 
 @forbid_in_graph
-def maybe_mark_dynamic(t: Any, index: Union[int, list[Any], tuple[Any]]) -> None:
+def maybe_mark_dynamic(t: Any, index: int | list[Any] | tuple[Any]) -> None:
     """
     Mark a tensor as having a dynamic dim, but don't enforce it (i.e., if this
     dimension ends up getting specialized, don't error).
@@ -1065,9 +1272,7 @@ def maybe_mark_dynamic(t: Any, index: Union[int, list[Any], tuple[Any]]) -> None
         maybe_mark_dynamic(t, i)
 
 
-def mark_static(
-    t: Any, index: Optional[Union[int, list[Any], tuple[Any]]] = None
-) -> None:
+def mark_static(t: Any, index: int | list[Any] | tuple[Any] | None = None) -> None:
     """
     Mark a tensor as having a static dim or mark a nn module class as static.
 
@@ -1150,20 +1355,40 @@ def mark_static_address(t: Any, guard: bool = False) -> None:
         t._dynamo_static_input_type = "unguarded"  # type: ignore[attr-defined]
 
 
+def _patch_einops_symint_compat(einops_mod: Any) -> None:
+    """Backport the SymInt lru_cache fix from einops 0.7.0 into einops <= 0.6.1."""
+    for name in ("_reconstruct_from_shape", "_prepare_transformation_recipe"):
+        cached = getattr(einops_mod, name)
+        uncached = cached.__wrapped__
+
+        def make_wrapper(cached_fn: Any, uncached_fn: Any) -> Any:
+            @functools.wraps(cached_fn)
+            def wrapper(*args: Any, **kwargs: Any) -> Any:
+                try:
+                    return cached_fn(*args, **kwargs)
+                except TypeError:
+                    return uncached_fn(*args, **kwargs)
+
+            return wrapper
+
+        setattr(einops_mod, name, make_wrapper(cached, uncached))
+
+
 # One day, Dynamo will support tracing into einops directly (no allow_in_graph needed)
 # Note that PyTorch supports multiple versions of einops, so when that day comes,
 # we still need to be really careful about version matches.
 def _allow_in_graph_einops() -> None:
     import einops
 
-    if einops.__version__ >= "0.8.2":
-        if hasattr(einops, "einops") and hasattr(einops.einops, "get_backend"):
-            # trigger backend registration up front to avoid a later guard failure
-            # that would otherwise cause a recompilation
-            einops.rearrange(torch.randn(1), "i -> i")
-
-        # einops 0.8.2+ don't need explicit allow_in_graph calls
-        return
+    # There is a lru_cache logspam issue with einops when allow_in_graph is not
+    # used. Disabling this for now until the lru_cache issue is resolved.
+    # if einops.__version__ >= "0.8.2":
+    #     if hasattr(einops, "einops") and hasattr(einops.einops, "get_backend"):
+    #         # trigger backend registration up front to avoid a later guard failure
+    #         # that would otherwise cause a recompilation
+    #         einops.rearrange(torch.randn(1), "i -> i")
+    #     # einops 0.8.2+ don't need explicit allow_in_graph calls
+    #     return
 
     try:
         # requires einops > 0.6.1, torch >= 2.0
@@ -1173,7 +1398,10 @@ def _allow_in_graph_einops() -> None:
 
         # einops > 0.6.1 will call the op registration logic as it is imported.
     except ImportError:
-        # einops <= 0.6.1
+        # einops <= 0.6.1 doesn't handle unhashable SymInt in its lru_cache'd
+        # helpers. Backport the try/except TypeError fallback from einops 0.7.0+
+        # so allow_in_graph works during fake tensor validation.
+        _patch_einops_symint_compat(einops.einops)  # type: ignore[attr-defined]
         allow_in_graph(einops.rearrange)
         allow_in_graph(einops.reduce)
         if hasattr(einops, "repeat"):
@@ -1209,9 +1437,9 @@ def __enter__(self) -> None:
 
     def __exit__(
         self,
-        exc_type: Optional[type[BaseException]],
-        exc_val: Optional[BaseException],
-        exc_tb: Optional[TracebackType],
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: TracebackType | None,
     ) -> None:
         return self.config_patch.__exit__(exc_type, exc_val, exc_tb)
 
@@ -1262,7 +1490,7 @@ def _patch_dynamo_config_check(changes: dict[str, Any]) -> None:
 # Unlike config.patch, we also need to accept tuple as input in order to
 # deal with context manager reconstruction.
 def patch_dynamo_config(
-    arg1: Optional[Union[str, dict[str, Any], tuple[tuple[str, Any], ...]]] = None,
+    arg1: str | dict[str, Any] | tuple[tuple[str, Any], ...] | None = None,
     arg2: Any = None,
     **kwargs: Any,
 ) -> DynamoConfigPatchProxy:
@@ -1297,7 +1525,7 @@ def dont_skip_tracing(fn: None = None) -> DynamoConfigPatchProxy: ...
 def dont_skip_tracing(fn: Callable[_P, _R]) -> Callable[_P, _R]: ...
 
 
-def dont_skip_tracing(fn: Optional[Any] = None) -> Any:
+def dont_skip_tracing(fn: Any | None = None) -> Any:
     """
     Context manager/decorator to trace into functions intentionally marked by developers to be skipped
     when tracing.
@@ -1318,7 +1546,7 @@ def disable_nested_graph_breaks(fn: None = None) -> DynamoConfigPatchProxy: ...
 def disable_nested_graph_breaks(fn: Callable[_P, _R]) -> Callable[_P, _R]: ...
 
 
-def disable_nested_graph_breaks(fn: Optional[Any] = None) -> Any:
+def disable_nested_graph_breaks(fn: Any | None = None) -> Any:
     """
     Context manager/decorator to disable nested graph breaks when tracing
     this function and any nested functions. Used when nested graph breaks
@@ -1342,9 +1570,9 @@ def __enter__(self) -> None:
 
     def __exit__(
         self,
-        exc_type: Optional[type[BaseException]],
-        exc_val: Optional[BaseException],
-        exc_tb: Optional[TracebackType],
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: TracebackType | None,
     ) -> None:
         _set_error_on_graph_break(self.prev_error_on_graph_break)
 
@@ -1369,7 +1597,49 @@ def error_on_graph_break(
     return ErrorOnGraphBreakDecoratorContextManager(error_on_graph_break)
 
 
-def is_dynamo_disable_recursive(method: Callable[[Any], Any]) -> Optional[bool]:
+class CudagraphOverrideContextManager:
+    """Context manager that overrides cudagraph recording during tracing."""
+
+    def __init__(self, fwd: bool | None = None, bwd: bool | None = None) -> None:
+        self.fwd = fwd
+        self.bwd = bwd
+
+    __call__ = wrap_dunder_call_ctx_manager
+
+    def __enter__(self) -> None:
+        pass
+
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: TracebackType | None,
+    ) -> None:
+        pass
+
+
+def override_cudagraphs(
+    fwd: bool | None = None, bwd: bool | None = None
+) -> CudagraphOverrideContextManager:
+    """
+    Context manager/decorator to override cudagraph recording for compiled graphs.
+
+    When used as a context manager, overrides cudagraphs for all graph segments
+    within the block (including across graph breaks).
+
+    When used as a decorator, marks a function so that any compiled graph
+    inlining it will have cudagraphs overridden.
+
+    Args:
+        fwd: If False, disable cudagraphs for forward. If True, force enable.
+             If None, don't override.
+        bwd: If False, disable cudagraphs for backward. If True, force enable.
+             If None, don't override.
+    """
+    return CudagraphOverrideContextManager(fwd=fwd, bwd=bwd)
+
+
+def is_dynamo_disable_recursive(method: Callable[[Any], Any]) -> bool | None:
     """
     Check if a method is marked as `dynamo_disable` recursively. It returns:
     - True if disable(recursive=True)
diff --git a/torch/_dynamo/device_interface.py b/torch/_dynamo/device_interface.py
index 25581197600f0..7ff97f4ac5a08 100644
--- a/torch/_dynamo/device_interface.py
+++ b/torch/_dynamo/device_interface.py
@@ -20,12 +20,13 @@
 from collections import namedtuple
 from collections.abc import Callable, Iterable
 from dataclasses import dataclass
-from typing import Any, Literal, Optional, Union
+from typing import Any, Literal
 
 import torch
+from torch.utils._pallas import has_torch_tpu
 
 
-get_cuda_stream: Optional[Callable[[int], int]]
+get_cuda_stream: Callable[[int], int] | None
 if torch.cuda._is_compiled():
     from torch._C import _cuda_getCurrentRawStream as get_cuda_stream
 else:
@@ -181,7 +182,7 @@ class DeviceGuard:
     """
 
     def __init__(
-        self, device_interface: type[DeviceInterface], index: Optional[int]
+        self, device_interface: type[DeviceInterface], index: int | None
     ) -> None:
         self.device_interface = device_interface
         self.idx = index
@@ -241,7 +242,6 @@ def get_device_properties(device: torch.types.Device = None) -> Any:
     set_device = staticmethod(torch.cuda.set_device)
     device_count = staticmethod(torch.cuda.device_count)
     stream = staticmethod(torch.cuda.stream)  # type: ignore[assignment]
-    # pyrefly: ignore [bad-override]
     current_stream = staticmethod(torch.cuda.current_stream)
     set_stream = staticmethod(torch.cuda.set_stream)  # type: ignore[assignment]
     _set_stream_by_id = staticmethod(torch.cuda._set_stream_by_id)  # type: ignore[assignment]
@@ -259,7 +259,7 @@ def is_available() -> bool:
         return torch.cuda.is_available()
 
     @staticmethod
-    def get_compute_capability(device: torch.types.Device = None) -> Union[int, str]:
+    def get_compute_capability(device: torch.types.Device = None) -> int | str:
         if torch.version.hip is None:
             major, min = torch.cuda.get_device_capability(device)
             return major * 10 + min
@@ -290,7 +290,7 @@ def raise_if_triton_unavailable(device: torch.types.Device = None) -> None:
             raise RuntimeError("triton not built with the 'nvidia' backend")
 
 
-get_mtia_stream: Optional[Callable[[int], int]]
+get_mtia_stream: Callable[[int], int] | None
 if torch.mtia._is_compiled():
     from torch._C import _mtia_getCurrentRawStream as get_mtia_stream
 else:
@@ -338,7 +338,6 @@ def get_device_properties(device: torch.types.Device = None) -> Any:
     set_device = staticmethod(torch.mtia.set_device)  # type: ignore[assignment]
     device_count = staticmethod(torch.mtia.device_count)
     stream = staticmethod(torch.mtia.stream)  # type: ignore[assignment]
-    # pyrefly: ignore [bad-override]
     current_stream = staticmethod(torch.mtia.current_stream)
     set_stream = staticmethod(torch.mtia.set_stream)  # type: ignore[assignment]
     _set_stream_by_id = staticmethod(torch.mtia._set_stream_by_id)  # type: ignore[assignment]
@@ -373,7 +372,7 @@ def raise_if_triton_unavailable(device: torch.types.Device = None) -> None:
             raise RuntimeError("triton not built with the 'mtia' backend")
 
 
-get_xpu_stream: Optional[Callable[[int], int]]
+get_xpu_stream: Callable[[int], int] | None
 if torch.xpu._is_compiled():
     from torch._C import _xpu_getCurrentRawStream as get_xpu_stream
 else:
@@ -421,7 +420,6 @@ def get_device_properties(device: torch.types.Device = None) -> Any:
     set_device = staticmethod(torch.xpu.set_device)
     device_count = staticmethod(torch.xpu.device_count)  # type: ignore[has-type]
     stream = staticmethod(torch.xpu.stream)  # type: ignore[assignment]
-    # pyrefly: ignore [bad-override]
     current_stream = staticmethod(torch.xpu.current_stream)
     set_stream = staticmethod(torch.xpu.set_stream)  # type: ignore[assignment]
     _set_stream_by_id = staticmethod(torch.xpu._set_stream_by_id)  # type: ignore[assignment]
@@ -525,7 +523,7 @@ def raise_if_triton_unavailable(device: torch.types.Device = None) -> None:
 class MpsInterface(DeviceInterface):
     @staticmethod
     def is_bf16_supported(including_emulation: bool = False) -> bool:
-        return torch.backends.mps.is_macos_or_newer(14, 0)
+        return True
 
     @classmethod
     def is_dtype_supported(
@@ -533,7 +531,7 @@ def is_dtype_supported(
     ) -> bool:
         if dtype in [torch.float64, torch.complex128]:
             return False
-        return dtype != torch.bfloat16 or cls.is_bf16_supported(including_emulation)
+        return True
 
     @staticmethod
     def is_available() -> bool:
@@ -564,19 +562,61 @@ def current_device() -> int:
             return 0
 
 
+class TpuInterface(DeviceInterface):
+    @staticmethod
+    def is_bf16_supported(including_emulation: bool = False) -> bool:
+        return True
+
+    @classmethod
+    def is_dtype_supported(
+        cls, dtype: torch.dtype, including_emulation: bool = False
+    ) -> bool:
+        return dtype not in (
+            torch.float64,
+            torch.complex32,
+            torch.complex64,
+            torch.complex128,
+            torch.half,
+        )
+
+    @staticmethod
+    def is_available() -> bool:
+        return has_torch_tpu()
+
+    @staticmethod
+    def current_device() -> int:
+        return 0
+
+    @staticmethod
+    def get_compute_capability(device: torch.types.Device = None) -> str:
+        return ""
+
+    # pyrefly: ignore [bad-override]
+    class Worker:
+        @staticmethod
+        def get_device_properties(device: torch.types.Device = None) -> Any:
+            return namedtuple("TPUProperties", ["multi_processor_count"])(
+                1  # type: ignore[arg-type]
+            )
+
+        @staticmethod
+        def current_device() -> int:
+            return 0
+
+
 device_interfaces: dict[str, type[DeviceInterface]] = {}
 _device_initialized = False
 
 
 def register_interface_for_device(
-    device: Union[str, torch.device], device_interface: type[DeviceInterface]
+    device: str | torch.device, device_interface: type[DeviceInterface]
 ) -> None:
     if isinstance(device, torch.device):
         device = device.type
     device_interfaces[device] = device_interface
 
 
-def get_interface_for_device(device: Union[str, torch.device]) -> type[DeviceInterface]:
+def get_interface_for_device(device: str | torch.device) -> type[DeviceInterface]:
     if isinstance(device, torch.device):
         device = device.type
     if not _device_initialized:
@@ -608,5 +648,6 @@ def init_device_reg() -> None:
 
     register_interface_for_device("cpu", CpuInterface)
     register_interface_for_device("mps", MpsInterface)
+    register_interface_for_device("tpu", TpuInterface)
 
     _device_initialized = True
diff --git a/torch/_dynamo/distributed.py b/torch/_dynamo/distributed.py
index 490b6330fafa4..16c93592adfaa 100644
--- a/torch/_dynamo/distributed.py
+++ b/torch/_dynamo/distributed.py
@@ -14,18 +14,16 @@
 non-distributed scenarios.
 """
 
-from typing import Optional
-
 import torch.distributed as dist
 
 from . import config
 
 
-_COMPILE_PG: Optional[dist.ProcessGroup] = None
-_GUARD_PG: Optional[dist.ProcessGroup] = None
+_COMPILE_PG: dist.ProcessGroup | None = None
+_GUARD_PG: dist.ProcessGroup | None = None
 
 
-def get_compile_pg() -> Optional[dist.ProcessGroup]:
+def get_compile_pg() -> dist.ProcessGroup | None:
     if (
         config.enable_compiler_collectives
         and dist.is_available()
@@ -44,7 +42,7 @@ def get_compile_pg() -> Optional[dist.ProcessGroup]:
 
 # NB: Unlike get_compile_pg, this is only called when guard collectives were
 # explicitly requested
-def get_guard_pg() -> Optional[dist.ProcessGroup]:
+def get_guard_pg() -> dist.ProcessGroup | None:
     if dist.is_available() and dist.is_initialized():
         global _GUARD_PG
         if _GUARD_PG is None:
diff --git a/torch/_dynamo/eval_frame.py b/torch/_dynamo/eval_frame.py
index 6bbe8a49183cb..bcc300466dd89 100644
--- a/torch/_dynamo/eval_frame.py
+++ b/torch/_dynamo/eval_frame.py
@@ -43,7 +43,7 @@
 from dataclasses import dataclass
 from enum import Enum
 from os.path import dirname, join
-from typing import Any, NamedTuple, Optional, TYPE_CHECKING, Union
+from typing import Any, NamedTuple, TYPE_CHECKING
 from unittest.mock import patch
 
 import sympy
@@ -56,11 +56,11 @@
 
 # see discussion at https://github.com/pytorch/pytorch/issues/120699
 from torch._C._dynamo.eval_frame import (  # noqa: F401
-    _EvalFrameOverride,
     reset_code,
     set_code_exec_strategy,
     set_eval_frame,
-    set_eval_frame_override,
+    set_fullgraph_compiled_frame_count,
+    set_fullgraph_error_on_nested_compile,
     set_guard_complete_hook,
     set_guard_error_hook,
     set_skip_guard_eval_unsafe,
@@ -172,7 +172,7 @@ def _maybe_set_eval_frame(callback: DynamoCallback) -> DynamoCallback:
 class DynamoStance:
     stance: str = "default"
     skip_guard_eval_unsafe: bool = False
-    backend: Union[str, Callable[..., Any], None] = None
+    backend: str | Callable[..., Any] | None = None
 
 
 _stance = DynamoStance()
@@ -195,7 +195,7 @@ def _set_stance(stance: DynamoStance) -> DynamoStance:
 
 _set_stance._dynamo_forbidden = True  # type: ignore[attr-defined]
 
-_EXAMPLE_INPUTS: Optional[dict[str, list[Any]]] = None
+_EXAMPLE_INPUTS: dict[str, list[Any]] | None = None
 
 
 def get_example_inputs(key: str) -> list[Any]:
@@ -234,7 +234,6 @@ def _callback_from_stance(callback: DynamoCallback) -> DynamoCallback:
     if _stance.stance == "default":
         # force_backend
         if _stance.backend is not None and callback not in (False, None):
-            # pyrefly: ignore [bad-argument-type]
             callback = _create_wrapped_callback(get_compiler_fn(_stance.backend))
 
         return callback
@@ -279,7 +278,12 @@ def fail_callback(
             cache_entries = _debug_get_cache_entry_list(frame.f_code)
             if cache_entries:
                 reasons = get_and_maybe_log_recompilation_reasons(
-                    cache_entries[0], frame, innermost_fn(callback), skip_logging=True
+                    # pyrefly: ignore [bad-argument-type]
+                    cache_entries[0],
+                    frame,
+                    # pyrefly: ignore [bad-argument-type]
+                    innermost_fn(callback),
+                    skip_logging=True,
                 )
                 if reasons:
                     failures = textwrap.indent("\n".join(reasons), "- ")
@@ -341,7 +345,6 @@ def callback_fn(*args: Any, **kwargs: Any) -> convert_frame.ConvertFrameReturn:
                 )
             elif stance == "aot_eager_then_compile":
                 aot_eager_fn = get_compiler_fn("aot_eager")
-                # pyrefly: ignore [bad-argument-type]
                 return _create_wrapped_callback(aot_eager_fn)(*args, **kwargs)
 
         dynamism = track_dynamism_across_examples(example_inputs)
@@ -375,7 +378,7 @@ def _reset_guarded_backend_cache() -> None:
 
 
 def _debug_get_cache_entry_list(
-    code: Union[types.CodeType, Callable[..., Any]],
+    code: types.CodeType | Callable[..., Any],
 ) -> list[CacheEntry]:
     """
     Given a code object or a callable object, retrieve the cache entries
@@ -428,7 +431,7 @@ def __init__(self, mod: torch.nn.Module, dynamo_ctx: _TorchDynamoContext) -> Non
 
     def __len__(self) -> int:
         # Proxy the len call to the original module
-        # pyrefly: ignore [invalid-argument, unsafe-overlap]
+        # pyrefly: ignore [unsafe-overlap]
         if isinstance(self._orig_mod, Sized):
             return len(self._orig_mod)
         # Mimic python's default behavior for objects without a length
@@ -438,7 +441,6 @@ def _initialize(self) -> None:
         # Do this stuff in constructor to lower overhead slightly
         if isinstance(self.dynamo_ctx, DisableContext):
             # No need to check trace rules
-            # pyrefly: ignore [bad-argument-type]
             self.forward = self.dynamo_ctx(self._orig_mod.__call__)
         elif config.wrap_top_frame or (
             isinstance(self._orig_mod.forward, types.MethodType)
@@ -447,7 +449,6 @@ def _initialize(self) -> None:
             # This may be a torch.nn.* instance in trace_rules.py which
             # won't trigger a frame evaluation workaround to add an extra
             # frame we can capture
-            # pyrefly: ignore [bad-argument-type]
             self.forward = self.dynamo_ctx(external_utils.wrap_inline(self._orig_mod))
         else:
             # Invoke hooks outside of dynamo then pickup the inner frame
@@ -495,7 +496,7 @@ def _aot_compile(self, inputs: list[torch._dynamo.aot_compile.ModelInput]) -> No
 
         self.forward = aot_compile_module(model, inputs, hooks, backend)
 
-    def _save_aot_compiled_module(self, path: Optional[str] = None) -> bytes:
+    def _save_aot_compiled_module(self, path: str | None = None) -> bytes:
         if not config.enable_aot_compile:
             raise RuntimeError(
                 "AOT Compile is not enabled, please set torch._dynamo.config.enable_aot_compile=True"
@@ -725,12 +726,6 @@ def guard_collectives_hook(guard_eval_result: bool) -> bool:
 _not_set = object()
 
 
-def _get_eval_frame_override() -> _EvalFrameOverride:
-    if torch._dynamo.config.error_on_dynamo_callback_in_fullgraph_compiled_code:
-        return _EvalFrameOverride.ERROR
-    return _EvalFrameOverride.SKIP
-
-
 class _TorchDynamoContext:
     def __init__(
         self,
@@ -743,18 +738,18 @@ def __init__(
         first_ctx: bool = False,
         *,
         fullgraph: bool = False,
-        error_on_graph_break: Optional[bool] = None,
+        error_on_graph_break: bool | None = None,
         export: bool = False,
-        dynamic: Optional[bool] = None,
-        compiler_config: Optional[Any] = None,
-        package: Optional[CompilePackage] = None,
-        hooks: Optional[Hooks] = None,
+        dynamic: bool | None = None,
+        compiler_config: Any | None = None,
+        package: CompilePackage | None = None,
+        hooks: Hooks | None = None,
     ) -> None:
         super().__init__()
         assert callable(callback) or callback is False or callback is None
         self.callback: DynamoCallback = callback
         self._backend_ctx_ctor = backend_ctx_ctor
-        self.prior: Union[Unset, DynamoCallback] = unset
+        self.prior: Unset | DynamoCallback = unset
         self.first_ctx = first_ctx
         self.fullgraph = fullgraph
         self.error_on_graph_break = error_on_graph_break
@@ -784,7 +779,7 @@ def call_on_enter() -> Callable[[], None]:
 
         if backend_ctx_ctor is not contextlib.nullcontext:
             # this case is not common
-            def call_backend_ctx() -> functools.partial[Optional[bool]]:
+            def call_backend_ctx() -> functools.partial[bool | None]:
                 ctx = backend_ctx_ctor()
                 ctx.__enter__()
                 return functools.partial(ctx.__exit__, None, None, None)
@@ -807,10 +802,10 @@ def __enter__(self) -> None:
 
     def __exit__(
         self,
-        exc_type: Optional[type[BaseException]],
-        exc_val: Optional[BaseException],
-        exc_tb: Optional[types.TracebackType],
-    ) -> Optional[bool]:
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: types.TracebackType | None,
+    ) -> bool | None:
         assert self.prior is not unset
         set_eval_frame(None)
         set_skip_guard_eval_unsafe(self.prior_skip_guard_eval_unsafe)
@@ -887,6 +882,16 @@ def aot_compile(example_inputs: tuple[tuple[Any, ...], dict[str, Any]]) -> Any:
 
         # Optimize the forward method of torch.nn.Module object
         if isinstance(fn, torch.nn.Module):
+            if type(fn) is torch.jit._script.RecursiveScriptModule:
+                raise RuntimeError(
+                    "torch.compile does not support compiling torch.jit.script or "
+                    "torch.jit.freeze models directly.\n\n"
+                    "Workaround: compile the original eager module instead:\n"
+                    "  model = torch.nn.Linear(3, 3)\n"
+                    "  compiled_model = torch.compile(model)  # compile the eager module\n\n"
+                    "torch.jit.script and torch.jit.freeze are deprecated in favor of "
+                    "torch.compile. See https://pytorch.org/docs/main/jit.html for details."
+                )
             mod = fn
             new_mod = OptimizedModule(mod, self)
             # Save the function pointer to find the original callable while nesting
@@ -915,16 +920,19 @@ def aot_compile(example_inputs: tuple[tuple[Any, ...], dict[str, Any]]) -> Any:
             f"A callable function is expected, but {type(fn)} is provided."
         )
 
-        # NOTE [Top-level TorchInGraph functions]
+        # NOTE [Top-level TorchInGraph and polyfilled functions]
         # Some callables (e.g. torch.exp) are represented as TorchInGraphFunctionVariable
         # when traced inside a frame. When such a function is passed directly to
         # torch.compile, we detect it here so we can force it through wrap_inline.
+        # Similarly, functions registered via substitute_in_graph have a polyfill
+        # that Dynamo can trace, so they also need wrap_inline.
         from .variables import TorchInGraphFunctionVariable
 
         rule = trace_rules.lookup(fn)
         top_level_in_graph = isinstance(rule, type) and issubclass(
             rule, TorchInGraphFunctionVariable
         )
+        has_polyfill = trace_rules.is_polyfilled_callable(fn)
 
         try:
             filename = inspect.getsourcefile(fn)
@@ -933,7 +941,12 @@ def aot_compile(example_inputs: tuple[tuple[Any, ...], dict[str, Any]]) -> Any:
         if config.debug_force_nested_calls:
             fn = external_utils.wrap_inline(fn)
         elif config.wrap_top_frame or (
-            (filename is None or trace_rules.check(fn) or top_level_in_graph)
+            (
+                filename is None
+                or trace_rules.check(fn)
+                or top_level_in_graph
+                or has_polyfill
+            )
             and (
                 getattr(fn, "__name__", "")
                 not in ["_call_impl", "_wrapped_call_impl", "_lazy_forward"]
@@ -955,12 +968,21 @@ def do_nothing(*arg: Any, **kwargs: Any) -> None:
 
         @functools.wraps(fn)
         def compile_wrapper(*args: Any, **kwargs: Any) -> Any:
+            # NB: function calls here could change global state (e.g. random state)
+            # and that can result in different behavior between eager and compiled!
+            # In particular, we don't have control over internal functions like justknobs_check
+            # called in _maybe_set_eval_frame.
+            # Unlike in eval_frame_cpp.cpp/convert_frame.py, we don't attempt to restore global state
+            # due to additional overhead costs.
             prior = set_eval_frame(None)
-            prior_eval_frame_override: _EvalFrameOverride | None = None
+            prior_error_on_nested_compile: bool | None = None
+            fullgraph_count_enabled = False
             if self.fullgraph:
-                prior_eval_frame_override = set_eval_frame_override(
-                    _get_eval_frame_override()
+                prior_error_on_nested_compile = set_fullgraph_error_on_nested_compile(
+                    torch._dynamo.config.error_on_dynamo_callback_in_fullgraph_compiled_code
                 )
+                if not self.export:
+                    fullgraph_count_enabled = set_fullgraph_compiled_frame_count(0) < 0
             try:
                 # We shouldn't compile inside kernel invocation.
                 if tracing_context := torch._guards.TracingContext.try_get():
@@ -1020,9 +1042,11 @@ def compile_wrapper(*args: Any, **kwargs: Any) -> Any:
 
                 _maybe_set_eval_frame(_callback_from_stance(callback))
 
+                call_succeeded = False
                 try:
-                    return fn(*args, **kwargs)
-                except (Unsupported, UncapturedHigherOrderOpError) as e:
+                    result = fn(*args, **kwargs)
+                    call_succeeded = True
+                except (Unsupported, UncapturedHigherOrderOpError, UserError) as e:
                     if config.verbose:
                         raise
                     # strip internal tracebacks from causes
@@ -1039,10 +1063,21 @@ def compile_wrapper(*args: Any, **kwargs: Any) -> Any:
                 finally:
                     # Restore the dynamic layer stack depth if necessary.
                     set_eval_frame(None)
+                    if fullgraph_count_enabled and call_succeeded:
+                        count = set_fullgraph_compiled_frame_count(-1)
+                        if count == 0:
+                            torch._dynamo.utils.warn_once(
+                                "torch.compile with fullgraph=True found no compiled frames. "
+                                "The frame was likely skipped (e.g., a non-infra torch dispatch "
+                                "mode was active, dynamo was disabled, or the frame was skipped. "
+                                "This may become an error in the future."
+                            )
                     if prior_error_on_graph_break is not None:
                         _set_error_on_graph_break(prior_error_on_graph_break)
-                    if prior_eval_frame_override is not None:
-                        set_eval_frame_override(prior_eval_frame_override)
+                    if prior_error_on_nested_compile is not None:
+                        set_fullgraph_error_on_nested_compile(
+                            prior_error_on_nested_compile
+                        )
                     torch._C._functorch.pop_dynamic_layer_stack_and_undo_to_depth(
                         saved_dynamic_layer_stack_depth
                     )
@@ -1050,7 +1085,10 @@ def compile_wrapper(*args: Any, **kwargs: Any) -> Any:
                     set_skip_guard_eval_unsafe(prior_skip_guard_eval_unsafe)
                     for cleanup in cleanups:
                         cleanup()
+                return result
             finally:
+                if fullgraph_count_enabled:
+                    set_fullgraph_compiled_frame_count(-1)
                 _maybe_set_eval_frame(prior)
 
         # hooks to properly handle inlining
@@ -1125,15 +1163,13 @@ def __init__(
         first_ctx: bool = False,
         *,
         fullgraph: bool = False,
-        error_on_graph_break: Optional[bool] = None,
+        error_on_graph_break: bool | None = None,
         export: bool = False,
-        dynamic: Optional[bool] = None,
-        compiler_config: Optional[Any] = None,
-        rebuild_ctx: Optional[
-            Callable[[], Union[OptimizeContext, _NullDecorator]]
-        ] = None,
-        package: Optional[CompilePackage] = None,
-        hooks: Optional[Hooks] = None,
+        dynamic: bool | None = None,
+        compiler_config: Any | None = None,
+        rebuild_ctx: Callable[[], OptimizeContext | _NullDecorator] | None = None,
+        package: CompilePackage | None = None,
+        hooks: Hooks | None = None,
     ) -> None:
         def on_enter() -> None:
             install_generation_tagging_init()
@@ -1158,7 +1194,7 @@ def on_enter() -> None:
             if _dynamic is None:
                 _dynamic = not torch._dynamo.config.assume_static_by_default
 
-            def call_compiled_autograd() -> functools.partial[Optional[bool]]:
+            def call_compiled_autograd() -> functools.partial[bool | None]:
                 assert rebuild_ctx is not None
                 compiler_fn = rebuild_ctx()
                 ctx = torch._dynamo.compiled_autograd._enable(
@@ -1199,7 +1235,7 @@ def __reduce__(self) -> tuple[type[RunOnlyContext], tuple[Any, ...]]:
 
 
 class DisableContext(_TorchDynamoContext):
-    def __init__(self, msg: Optional[str] = None, wrapping: bool = True) -> None:
+    def __init__(self, msg: str | None = None, wrapping: bool = True) -> None:
         super().__init__(callback=None)
         self.msg = msg
         self.wrapping = wrapping
@@ -1229,7 +1265,6 @@ def __call__(self, fn: Callable[..., Any]) -> Callable[..., Any]:
             cls_obj.__call__ = self(cls_obj.__call__)
             if issubclass(cls_obj, torch.nn.Module):
                 # NN module variable tracker directly inlines the _call_impl. Disable it.
-                # pyrefly: ignore [missing-attribute]
                 cls_obj._call_impl = self(cls_obj._call_impl)
             return cls_obj
 
@@ -1297,12 +1332,12 @@ def _optimize_catch_errors(
         [], contextlib.AbstractContextManager[Any]
     ] = null_context,
     fullgraph: bool = False,
-    error_on_graph_break: Optional[bool] = None,
+    error_on_graph_break: bool | None = None,
     export: bool = False,
-    dynamic: Optional[bool] = None,
-    compiler_config: Optional[Any] = None,
-    rebuild_ctx: Optional[Callable[[], Union[OptimizeContext, _NullDecorator]]] = None,
-    package: Optional[CompilePackage] = None,
+    dynamic: bool | None = None,
+    compiler_config: Any | None = None,
+    rebuild_ctx: Callable[[], OptimizeContext | _NullDecorator] | None = None,
+    package: CompilePackage | None = None,
 ) -> OptimizeContext:
     return OptimizeContext(
         convert_frame.catch_errors_wrapper(compile_fn, hooks),
@@ -1320,7 +1355,7 @@ def _optimize_catch_errors(
 
 
 def get_compiler_fn(
-    compiler_fn: Union[str, Callable[..., Any], None],
+    compiler_fn: str | Callable[..., Any] | None,
 ) -> WrapBackendDebug:
     from .repro.after_dynamo import wrap_backend_debug
 
@@ -1349,7 +1384,7 @@ def __call__(self, fn: Callable[..., Any]) -> Callable[..., Any]:
 # Make dynamo graph to have same input/output spec as user code
 def argument_names(
     f_sig: inspect.Signature,
-    args: Union[list[Any], tuple[Any, ...]],
+    args: list[Any] | tuple[Any, ...],
     kwargs: dict[str, Any],
 ) -> list[str]:
     def signature_to_fullargspec(sig: inspect.Signature) -> inspect.FullArgSpec:
@@ -1473,8 +1508,8 @@ def check_for_incompatible_configs() -> None:
     )
 
 
-def optimize(*args: Any, **kwargs: Any) -> Union[OptimizeContext, _NullDecorator]:
-    def rebuild_ctx() -> Union[OptimizeContext, _NullDecorator]:
+def optimize(*args: Any, **kwargs: Any) -> OptimizeContext | _NullDecorator:
+    def rebuild_ctx() -> OptimizeContext | _NullDecorator:
         ca_kwargs_override = config.compiled_autograd_kwargs_override
         if ca_kwargs_override:
             # NOTE: The process of translating other `torch.compile` kwargs to `torch._dynamo.optimize` kwargs
@@ -1489,19 +1524,20 @@ def rebuild_ctx() -> Union[OptimizeContext, _NullDecorator]:
 
 
 def _optimize(
-    rebuild_ctx: Callable[[], Union[OptimizeContext, _NullDecorator]],
-    backend: Union[str, Callable[..., Any]] = "inductor",
+    rebuild_ctx: Callable[[], OptimizeContext | _NullDecorator],
+    backend: str | Callable[..., Any] = "inductor",
     *,
     nopython: bool = False,
-    error_on_graph_break: Optional[bool] = None,
-    guard_export_fn: Optional[Callable[[_guards.GuardsSet], None]] = None,
-    guard_fail_fn: Optional[Callable[[GuardFail], None]] = None,
+    error_on_graph_break: bool | None = None,
+    guard_export_fn: Callable[[_guards.GuardsSet], None] | None = None,
+    guard_fail_fn: Callable[[GuardFail], None] | None = None,
     guard_filter_fn: Callable[[Sequence[GuardFilterEntry]], Sequence[bool]]
     | None = None,
     disable: bool = False,
-    dynamic: Optional[bool] = None,
-    package: Optional[CompilePackage] = None,
-) -> Union[OptimizeContext, _NullDecorator]:
+    dynamic: bool | None = None,
+    package: CompilePackage | None = None,
+    recompile_limit: int | None = None,
+) -> OptimizeContext | _NullDecorator:
     """
     The main entrypoint of TorchDynamo.  Do graph capture and call
     backend() to optimize extracted graphs.
@@ -1559,6 +1595,7 @@ def toy_example(a, b): ...
             hooks=hooks,
             rebuild_ctx=rebuild_ctx,
             package=package,
+            recompile_limit=recompile_limit,
         )
 
     backend = get_compiler_fn(backend)
@@ -1579,10 +1616,10 @@ def toy_example(a, b): ...
 
     return _optimize_catch_errors(
         convert_frame.convert_frame(
-            # pyrefly: ignore [bad-argument-type]
             backend,
             hooks,
             package=package,
+            recompile_limit=recompile_limit,
         ),
         hooks,
         backend_ctx_ctor,
@@ -1686,7 +1723,7 @@ def __init__(
         matched_output_elements_positions: list[int],
         example_fake_inputs: list[torch.Tensor],
         flat_args_dynamic_dims: list[set[int]],
-        fake_mode: Optional[fake_tensor.FakeTensorMode] = None,
+        fake_mode: fake_tensor.FakeTensorMode | None = None,
     ) -> None:
         super().__init__(m)
 
@@ -1887,12 +1924,12 @@ def is_supported_type(val: Any) -> bool:
 def rewrite_signature(
     f_sig: inspect.Signature,
     graph: torch.fx.GraphModule,
-    fake_mode: Optional[fake_tensor.FakeTensorMode],
+    fake_mode: fake_tensor.FakeTensorMode | None,
     flat_args: list[Any],
     in_spec: pytree.TreeSpec,
     example_fake_inputs: list[Any],
     graph_captured_input: Iterable[Any],
-    graph_captured_output: Optional[Iterable[Any]],
+    graph_captured_output: Iterable[Any] | None,
     dynamo_traced_result: Any,
     flat_args_dynamic_dims: list[set[int]],
 ) -> torch.fx.GraphModule:
@@ -1925,8 +1962,8 @@ def check_optional_input_and_error(f_sig: inspect.Signature) -> None:
 
     def produce_matching(
         debug_type: str, sources: Iterable[Any], candidates: Iterable[Any]
-    ) -> list[Optional[int]]:
-        matched_elements_positions: list[Optional[int]] = []
+    ) -> list[int | None]:
+        matched_elements_positions: list[int | None] = []
         dict_of_source_vals = {}
         for i, val in enumerate(sources):
             dict_of_source_vals[id(val)] = i
@@ -1982,18 +2019,16 @@ def export(
     *extra_args: Any,
     aten_graph: bool = False,
     pre_dispatch: bool = False,
-    decomposition_table: Optional[
-        dict[torch._ops.OpOverload, Callable[..., Any]]
-    ] = None,
+    decomposition_table: dict[torch._ops.OpOverload, Callable[..., Any]] | None = None,
     tracing_mode: str = "symbolic",
-    dynamic_shapes: Optional[Union[dict[str, Any], tuple[Any], list[Any]]] = None,
+    dynamic_shapes: dict[str, Any] | tuple[Any] | list[Any] | None = None,
     specialize_float: bool = True,
     assume_static_by_default: bool = False,
     same_signature: bool = True,
     disable_constraint_solver: bool = False,
     prefer_deferred_runtime_asserts_over_guards: bool = False,
     _log_export_usage: bool = True,
-    constraints: Optional[list[Constraint]] = None,
+    constraints: list[Constraint] | None = None,
     **extra_kwargs: Any,
 ) -> Callable[..., ExportResult]:
     """
@@ -2085,7 +2120,7 @@ def inner(*args: Any, **kwargs: Any) -> ExportResult:
         graph = None
         out_guards = None
         graph_captured_input = None
-        graph_captured_result: Optional[tuple[torch.Tensor, ...]] = None
+        graph_captured_result: tuple[torch.Tensor, ...] | None = None
         fake_mode = None
         result_traced = None
 
@@ -2164,7 +2199,7 @@ def result_capturing_wrapper(*graph_inputs: Any) -> Any:
                     )
 
                     def fakify_with_ambient(
-                        path: KeyPath, t: Union[torch.Tensor, _IntWrapper, Any]
+                        path: KeyPath, t: torch.Tensor | _IntWrapper | Any
                     ) -> Any:
                         if isinstance(t, torch.Tensor):
                             # pyrefly: ignore [missing-attribute]
@@ -2430,13 +2465,14 @@ def rebuild_ctx() -> OptimizeContext:
 
 def _optimize_assert(
     rebuild_ctx: Callable[[], OptimizeContext],
-    backend: Union[str, Callable[..., Any], None],
+    backend: str | Callable[..., Any] | None,
     *,
     hooks: Hooks = Hooks(None, None, None),
     export: bool = False,
-    export_constraints: Optional[Any] = None,
-    dynamic: Optional[bool] = None,
-    package: Optional[CompilePackage] = None,
+    export_constraints: Any | None = None,
+    dynamic: bool | None = None,
+    package: CompilePackage | None = None,
+    recompile_limit: int | None = None,
 ) -> OptimizeContext:
     """
     Guarantees single-graph capture.
@@ -2463,11 +2499,11 @@ def _optimize_assert(
 
     return _optimize_catch_errors(
         convert_frame.convert_frame_assert(
-            # pyrefly: ignore [bad-argument-type]
             backend,
             export=export,
             export_constraints=export_constraints,
             package=package,
+            recompile_limit=recompile_limit,
         ),
         hooks,
         backend_ctx_ctor,
diff --git a/torch/_dynamo/exc.py b/torch/_dynamo/exc.py
index a8714ddf4d612..e2da1f37cb1d4 100644
--- a/torch/_dynamo/exc.py
+++ b/torch/_dynamo/exc.py
@@ -35,7 +35,7 @@
 from functools import lru_cache
 from pathlib import Path
 from traceback import extract_stack, format_exc, format_list, FrameSummary, StackSummary
-from typing import Any, NoReturn, Optional, TYPE_CHECKING
+from typing import Any, NoReturn, TYPE_CHECKING
 
 import torch._guards
 from torch._utils_internal import get_file_path_2
@@ -47,6 +47,7 @@
 if TYPE_CHECKING:
     import types
 
+    from torch._dynamo.variables import VariableTracker
     from torch._guards import CompileId
 
     from .output_graph import DynamoTracerOutput
@@ -80,7 +81,7 @@ class TorchDynamoException(RuntimeError):
 
     def __init__(self, *args: Any, **kwargs: Any) -> None:
         super().__init__(*args, **kwargs)
-        self._torch_dynamo_tracer_output: Optional[DynamoTracerOutput] = None
+        self._torch_dynamo_tracer_output: DynamoTracerOutput | None = None
         self.frame_exec_strategy: FrameExecStrategy | None = None
 
 
@@ -93,9 +94,9 @@ class ResumePrologueTracingError(TorchDynamoException):
 
 
 class RestartAnalysis(TorchDynamoException):
-    restart_reason: Optional[str]
+    restart_reason: str | None
 
-    def __init__(self, *args: Any, restart_reason: Optional[str] = None) -> None:
+    def __init__(self, *args: Any, restart_reason: str | None = None) -> None:
         self.restart_reason = restart_reason
         super().__init__(*args)
 
@@ -111,6 +112,14 @@ class AutogradGradRestartAnalysis(RestartAnalysis):
     """
 
 
+class RequiresGradRestartAnalysis(RestartAnalysis):
+    """Raised when a source-less requires_grad_() intermediate leaks as output.
+
+    On restart, requires_grad_() will graph break instead of being traced,
+    preserving partial acceleration for code before the call.
+    """
+
+
 class UnspecializeRestartAnalysis(RestartAnalysis):
     pass
 
@@ -163,7 +172,7 @@ def __init__(self) -> None:
 
 class ShortenTraceback(TorchDynamoException):
     def __init__(
-        self, *args: Any, first_useful_frame: Optional[types.FrameType], **kwargs: Any
+        self, *args: Any, first_useful_frame: types.FrameType | None, **kwargs: Any
     ) -> None:
         super().__init__(*args, **kwargs)
         self.first_useful_frame = first_useful_frame
@@ -183,7 +192,7 @@ def __init__(
         self,
         backend_fn: Any,
         inner_exception: Exception,
-        first_useful_frame: Optional[types.FrameType],
+        first_useful_frame: types.FrameType | None,
     ) -> None:
         self.backend_name = getattr(backend_fn, "__name__", "?")
         self.inner_exception = inner_exception
@@ -204,7 +213,7 @@ def __init__(
         gb_type: str = "",
         skip_frame: bool = False,
         *,
-        case_name: Optional[str] = None,
+        case_name: str | None = None,
         real_stack: StackSummary | None = None,
     ) -> None:
         super().__init__(msg)
@@ -213,7 +222,7 @@ def __init__(
         self.real_stack = real_stack
         self.msg = msg
         self.skip_frame = skip_frame
-        self.category: Optional[str] = None
+        self.category: str | None = None
         self.add_to_stats()
         self.gb_type: str | None = gb_type
         self.logged = False
@@ -259,9 +268,9 @@ class UserErrorType(Enum):
     UNSUPPORTED_ALIASED_MUTATED_DYNAMIC_INPUTS = auto()
 
 
-class UserError(Unsupported):
+class UserError(TorchDynamoException):
     def __init__(
-        self, error_type: UserErrorType, msg: str, case_name: Optional[str] = None
+        self, error_type: UserErrorType, msg: str, case_name: str | None = None
     ) -> None:
         """
         Type of errors that would be valid in Eager, but not supported in TorchDynamo.
@@ -278,8 +287,12 @@ def __init__(
             else:
                 msg += "\n"
             msg += exportdb_error_message(case_name)
-        super().__init__(msg, case_name if case_name else "UserError")
+        super().__init__(msg)
+        self.real_stack = torch._guards.TracingContext.extract_stack()
+        self.skip_frame = False
+        self.logged = False
         self.error_type = error_type
+        self.msg = msg
         self.message = msg
 
 
@@ -309,10 +322,6 @@ def __init__(self, msg: str, real_stack: StackSummary | None = None) -> None:
         )
 
 
-class IncorrectUsage(Exception):
-    pass
-
-
 # TODO: I'm a little uncertain about what error classification we should have
 # for this.  This is potentially a user error, but regressions in
 # specialization in PyTorch proper could also trigger this problem
@@ -327,7 +336,7 @@ class PackageError(TorchDynamoException):
 class ObservedException(TorchDynamoException):
     # An exception observed during the tracing. This exception is used by Dynamo to handle exceptions.
     def __init__(
-        self, *args: Any, real_stack: Optional[StackSummary] = None, **kwargs: Any
+        self, *args: Any, real_stack: StackSummary | None = None, **kwargs: Any
     ) -> None:
         super().__init__(*args, **kwargs)
         self.real_stack: StackSummary = (
@@ -339,12 +348,12 @@ def __init__(
 
 class ObservedUserStopIteration(ObservedException):
     # An UserStopIteration exception observed during the Dynamo tracing (e.g Dynamo tracing __next__)
-    value: Optional[Any]
+    value: Any | None
 
     # Reference `StopIteration_init` in CPython
     # https://github.com/python/cpython/blob/3.11/Objects/exceptions.c#L568-L584
     def __init__(
-        self, *args: Any, real_stack: Optional[StackSummary] = None, **kwargs: Any
+        self, *args: Any, real_stack: StackSummary | None = None, **kwargs: Any
     ) -> None:
         super().__init__("unhandled `raise StopIteration`", real_stack=real_stack)
         if len(args) > 0:
@@ -410,7 +419,7 @@ def get_dynamo_observed_exception(exc_type: type[Exception]) -> type[ObservedExc
         observed_exception_map[exc_type] = type(  # type: ignore[assignment]
             f"Observed{name}Error", (ObservedException,), {}
         )
-    # pyrefly: ignore [bad-index, index-error]
+    # pyrefly: ignore [bad-index]
     return observed_exception_map[exc_type]
 
 
@@ -418,18 +427,24 @@ def raise_observed_exception(
     exc_type: type[Exception],
     tx: InstructionTranslatorBase,
     *,
-    args: Optional[list[Any]] = None,
-    kwargs: Optional[dict[str, Any]] = None,
+    args: list[VariableTracker] | list[str] | None = None,
+    kwargs: dict[str, VariableTracker] | None = None,
 ) -> NoReturn:
     from .symbolic_convert import ExceptionVals
     from .variables.builder import SourcelessBuilder
 
+    if args:
+        args_ = [
+            SourcelessBuilder.create(tx, arg) if isinstance(arg, str) else arg
+            for arg in args
+        ]
+    else:
+        args_: list[VariableTracker] = []
+
     # CPython here raises an exception. Since there is no python code, we have to manually setup the exception
     # stack and raise the exception.
     exception_vt = SourcelessBuilder.create(tx, exc_type).call_function(
-        tx,  # pyrefly: ignore[bad-argument-type]
-        [SourcelessBuilder.create(tx, a) for a in args] if args else [],
-        kwargs or {},
+        tx, args_, kwargs or {}
     )
     assert isinstance(exception_vt, ExceptionVals)
     tx._attach_traceback_to_exception(exception_vt)
@@ -437,10 +452,15 @@ def raise_observed_exception(
     raised_exc = get_dynamo_observed_exception(exc_type)
     # Store the original exception arguments for better error messages
     if args:
-        raise raised_exc(*args)
+        raise raised_exc(*args_)
     raise raised_exc
 
 
+def raise_type_error(tx: InstructionTranslatorBase, msg: str) -> NoReturn:
+    """Raise a TypeError as an observed exception during tracing."""
+    raise_observed_exception(TypeError, tx, args=[msg])
+
+
 def handle_observed_exception(tx: Any) -> None:
     # This is essentially exception handling code, equivalent of this pseudo code
     #
@@ -570,7 +590,7 @@ def _load_gb_type_to_gb_id_map() -> dict[str, Any]:
     return mapping
 
 
-def get_gbid_documentation_link(gb_type: str) -> Optional[str]:
+def get_gbid_documentation_link(gb_type: str) -> str | None:
     """
     Retrieves the GBID documentation link for a given graph break type.
 
@@ -627,10 +647,13 @@ def unimplemented(
             past_real_stack = from_exc.real_stack
         if isinstance(from_exc, Unsupported):
             msg = f"{from_exc.msg}\n\n*** While handling this graph break, another graph break occurred: ***\n\n{msg}"
+            # noqa: GB_REGISTRY
             raise Unsupported(msg, gb_type, skip_frame, real_stack=past_real_stack)
+        # noqa: GB_REGISTRY
         raise Unsupported(
             msg, gb_type, skip_frame, real_stack=past_real_stack
         ) from from_exc
+    # noqa: GB_REGISTRY
     raise Unsupported(msg, gb_type, skip_frame)
 
 
@@ -710,7 +733,7 @@ def augment_exc_message(exc: Exception, msg: str = "\n", export: bool = False) -
 
 def get_exc_message(
     e: Exception, compile_id: CompileId
-) -> tuple[Optional[str], Optional[int]]:
+) -> tuple[str | None, int | None]:
     filename = None
     lineno = None
     if e.innermost_user_frame_summary is not None:  # type: ignore[attr-defined]
@@ -725,8 +748,8 @@ def get_stack_above_dynamo() -> StackSummary:
 
 
 def get_real_stack(
-    exc: Exception, frame: Optional[DynamoFrameType] = None
-) -> Optional[StackSummary]:
+    exc: Exception, frame: DynamoFrameType | None = None
+) -> StackSummary | None:
     real_stack = getattr(exc, "real_stack", None)
     if real_stack is None:
         return None
@@ -773,13 +796,13 @@ def filter_stack(stack: StackSummary) -> StackSummary:
     return user_stack
 
 
-def remove_resume_prefix(name: str) -> Optional[str]:
+def remove_resume_prefix(name: str) -> str:
     from .resume_execution import TORCH_DYNAMO_RESUME_IN_PREFIX
 
     match = re.match(f"{TORCH_DYNAMO_RESUME_IN_PREFIX}_(\\w+)_at_\\d+", name)
     if match:
         return match.group(1)
-    return None
+    return name
 
 
 def collapse_resume_frames(stack: StackSummary | list[FrameSummary]) -> StackSummary:
@@ -811,6 +834,7 @@ def collapse_resume_frames(stack: StackSummary | list[FrameSummary]) -> StackSum
             new_stack[-1] = frame
             frame.name = name
         else:
+            frame.name = name
             new_stack.append(frame)
 
     return new_stack
@@ -819,8 +843,8 @@ def collapse_resume_frames(stack: StackSummary | list[FrameSummary]) -> StackSum
 def format_error_msg_verbose(
     exc: Exception,
     code: types.CodeType,
-    record_filename: Optional[str] = None,
-    frame: Optional[DynamoFrameType] = None,
+    record_filename: str | None = None,
+    frame: DynamoFrameType | None = None,
 ) -> str:
     msg = (
         f"WON'T CONVERT {code.co_name} {code.co_filename} line {code.co_firstlineno}\n"
@@ -851,7 +875,7 @@ def format_frame_info(code: types.CodeType) -> str:
     )
 
 
-def format_skip_frame_message(code: Optional[types.CodeType], reason: str) -> str:
+def format_skip_frame_message(code: types.CodeType | None, reason: str) -> str:
     if code is not None:
         frame_info = format_frame_info(code)
         return (
@@ -868,8 +892,8 @@ def format_skip_frame_message(code: Optional[types.CodeType], reason: str) -> st
 def format_error_msg(
     exc: Exception,
     code: types.CodeType,
-    record_filename: Optional[str] = None,
-    frame: Optional[DynamoFrameType] = None,
+    record_filename: str | None = None,
+    frame: DynamoFrameType | None = None,
 ) -> str:
     if config.verbose:
         return format_error_msg_verbose(exc, code, record_filename, frame)
diff --git a/torch/_dynamo/external_utils.py b/torch/_dynamo/external_utils.py
index ef1d6cd3b11d1..fb3622cc0710c 100644
--- a/torch/_dynamo/external_utils.py
+++ b/torch/_dynamo/external_utils.py
@@ -23,7 +23,7 @@
 import functools
 import warnings
 from collections.abc import Callable
-from typing import Any, Optional, TYPE_CHECKING, TypeVar, Union
+from typing import Any, TYPE_CHECKING, TypeVar
 from typing_extensions import deprecated, ParamSpec
 
 import torch
@@ -72,7 +72,7 @@ def inner(*args: _P.args, **kwargs: _P.kwargs) -> _R:
 
 
 def call_hook(
-    hook: Callable[..., Optional[torch.Tensor]], *args: Any, **kwargs: Any
+    hook: Callable[..., torch.Tensor | None], *args: Any, **kwargs: Any
 ) -> torch.Tensor:
     """
     Used by compiled autograd to handle hook returning None.
@@ -97,7 +97,6 @@ def wrap(*args: _P.args, **kwargs: _P.kwargs) -> pytree.PyTree:
         args, kwargs = pytree.tree_map_only(
             torch.Tensor, lambda x: x.numpy(), (args, kwargs)
         )
-        # pyrefly: ignore [invalid-param-spec]
         out = f(*args, **kwargs)
         # pyrefly: ignore [missing-attribute]
         return pytree.tree_map_only(np.ndarray, lambda x: torch.as_tensor(x), out)
@@ -129,7 +128,7 @@ def call_backward(
     backward_c_function: torch.autograd.function.BackwardCFunction,
     saved_tensors: list[torch.Tensor],
     *args: Any,
-) -> Union[torch.Tensor, tuple[torch.Tensor, ...]]:
+) -> torch.Tensor | tuple[torch.Tensor, ...]:
     fake = FakeBackwardCFunction(backward_c_function, saved_tensors)
     grads = fake._forward_cls.backward(fake, *args)  # type: ignore[attr-defined]
 
@@ -244,7 +243,7 @@ def inner(*args: _P.args, **kwargs: _P.kwargs) -> _R:
 
 # Use only on ints marked dynamic via torch.empty(0, integer)
 # Currently only way to mark ints as dynamic: https://github.com/pytorch/pytorch/issues/129623
-def unwrap_maybe_dynamic_int(x: Union[torch.Tensor, int]) -> int:
+def unwrap_maybe_dynamic_int(x: torch.Tensor | int) -> int:
     if isinstance(x, torch.Tensor):
         # x.size() is expected to be [0, dynamic_int]
         return x.size(1)
diff --git a/torch/_dynamo/funcname_cache.py b/torch/_dynamo/funcname_cache.py
index f71cb5c6b02a3..6aa935cd67492 100644
--- a/torch/_dynamo/funcname_cache.py
+++ b/torch/_dynamo/funcname_cache.py
@@ -17,7 +17,6 @@
 """
 
 import tokenize
-from typing import Optional
 
 
 cache: dict[str, dict[int, str]] = {}
@@ -69,7 +68,7 @@ def _add_file(filename: str) -> None:
     cache[filename] = result
 
 
-def get_funcname(filename: str, lineno: int) -> Optional[str]:
+def get_funcname(filename: str, lineno: int) -> str | None:
     if filename not in cache:
         _add_file(filename)
     return cache[filename].get(lineno, None)
diff --git a/torch/_dynamo/functional_export.py b/torch/_dynamo/functional_export.py
index 97322a7690e2a..41de630917a5f 100644
--- a/torch/_dynamo/functional_export.py
+++ b/torch/_dynamo/functional_export.py
@@ -5,7 +5,7 @@
 import types
 from collections import namedtuple
 from collections.abc import Callable, Iterable, Sequence
-from typing import Any, Optional, TYPE_CHECKING, TypeVar, Union
+from typing import Any, Optional, TYPE_CHECKING, TypeVar
 
 import sympy
 
@@ -19,7 +19,6 @@
 from torch._dynamo.utils import dynamo_timed, get_metrics_context
 from torch._export.utils import _compiling_state_context
 from torch._guards import TracingContext
-from torch.export import _restore_state_dict
 from torch.export.dynamic_shapes import _RelaxedConstraint, Constraint
 from torch.fx import Node
 from torch.fx.experimental.symbolic_shapes import (
@@ -131,10 +130,8 @@ def _process_source_fn(source_fn_stack: Iterable[T]) -> Iterable[T]:
                 name, cls = item
                 if isinstance(name, str):
                     clean_name = clean_export_root_string(name)
-                    # pyrefly: ignore[bad-argument-type]
                     cleaned_stack.append((clean_name, cls))
                 else:
-                    # pyrefly: ignore[bad-argument-type]
                     cleaned_stack.append(item)
             else:
                 # pyrefly: ignore [bad-argument-type]
@@ -237,7 +234,8 @@ def __init__(
         flat_args_dynamic_dims: list[set[int]],
         graph_input_order: dict[int, int],
         graph_output_map: dict[int, tuple[str, Any]],
-        fake_mode: Optional[Any] = None,
+        fake_mode: Any | None = None,
+        graph_inputs: dict[int, Any] | None = None,
     ) -> None:
         super().__init__(module)
 
@@ -248,6 +246,7 @@ def __init__(
         self.graph_input_order = graph_input_order
         self.graph_output_map = graph_output_map
         self.fake_mode = fake_mode
+        self.graph_inputs = graph_inputs or {}
 
         # Get original placeholders and output
         self.placeholders = [n for n in module.graph.nodes if n.op == "placeholder"]
@@ -332,7 +331,25 @@ def placeholder(
 
             return new_arg
         else:
-            # Shouldn't happen if mapping is correct, but fallback
+            # Convert captured objects (e.g., opaque objects from closures) to
+            # get_attr nodes
+            placeholder_idx = self.placeholders.index(self.current_node)
+            if placeholder_idx in self.graph_inputs:
+                source = self.graph_inputs[placeholder_idx]
+                if not isinstance(source, torch._dynamo.source.GetItemSource):
+                    example_val = self.current_node.meta.get(
+                        "val"
+                    ) or self.current_node.meta.get("example_value")
+                    if example_val is not None:
+                        attr_name = f"_captured_{placeholder_idx}"
+                        if isinstance(example_val, torch.Tensor):
+                            self.module.register_buffer(attr_name, example_val)
+                        else:
+                            setattr(self.module, attr_name, example_val)
+                        result = self.tracer.create_proxy("get_attr", attr_name, (), {})
+                        result.node.meta = self.current_node.meta.copy()
+                        result.node.meta["val"] = example_val
+                        return result
             return super().placeholder(target, args, kwargs)
 
     def output(
@@ -383,16 +400,14 @@ def transform(self) -> torch.fx.GraphModule:
             if "dynamo_flat_name_to_original_fqn" in self.module.meta:
                 # pyrefly: ignore [bad-index]
                 result_gm.meta["dynamo_flat_name_to_original_fqn"] = self.module.meta[
-                    # pyrefly: ignore [bad-index, index-error]
-                    # pyrefly: ignore [bad-index, index-error]
+                    # pyrefly: ignore [bad-index]
                     "dynamo_flat_name_to_original_fqn"
                 ]
             # pyrefly: ignore [unsupported-operation]
             if "dynamo_compile_id" in self.module.meta:
                 # pyrefly: ignore [bad-index]
                 result_gm.meta["dynamo_compile_id"] = self.module.meta[
-                    # pyrefly: ignore [bad-index, index-error]
-                    # pyrefly: ignore [bad-index, index-error]
+                    # pyrefly: ignore [bad-index]
                     "dynamo_compile_id"
                 ]
 
@@ -406,7 +421,7 @@ def _suggest_or_raise_constraint_violation(
     graph_capture_output: CaptureOutput,
     args: Any,
     kwargs: Any,
-    dynamic_shapes: Optional[Union[dict[str, Any], tuple[Any], list[Any]]],
+    dynamic_shapes: dict[str, Any] | tuple[Any] | list[Any] | None,
 ) -> None:
     constraint_violation_error = None
     try:
@@ -590,13 +605,11 @@ def create_fx_graph_from_captured_output(
 
     graph_module = backend_input.graph_module
     if isinstance(root, torch.nn.Module):
-        graph_module._parameters = root._parameters.copy()
-        graph_module._buffers = root._buffers.copy()
+        graph_module._parameters = root._parameters
+        graph_module._buffers = root._buffers
         assert all(not hasattr(graph_module, m) for m in root._modules)
         graph_module._modules.update(root._modules)
-        graph_module._non_persistent_buffers_set = (
-            root._non_persistent_buffers_set.copy()
-        )
+        graph_module._non_persistent_buffers_set = root._non_persistent_buffers_set
         if sys.version_info >= (3, 14):
             import annotationlib  # added in 3.14
 
@@ -742,23 +755,51 @@ def gen_var_bindings(
     {", ".join(without_annotation)}, = self._dynamo_bytecode_flatten(*_fn_args)"""
 
     def generate_output(
-        self, output_args: torch.fx.node.Argument, *, descs: object | None = None
+        self,
+        output_args: torch.fx.node.Argument,
+        *,
+        descs: object | None = None,
+        repr_fn: Any | None = None,
     ) -> str:
+        if repr_fn is None:
+            repr_fn = repr
         # pyrefly: ignore [not-iterable]
-        returned = f"self._dynamo_bytecode_unflatten(({', '.join([str(a) for a in output_args])},), _fn_args)"
+        returned = f"self._dynamo_bytecode_unflatten(({', '.join([repr_fn(a) for a in output_args])},), _fn_args)"
         if self.wrap_tuple:
             returned = f"({returned},)"
         return f"return {returned}"
 
 
 def dynamo_graph_capture_for_export(
-    mod: Callable[..., Any],
-    constraints: Optional[list[Constraint]] = None,
-    restore_state_dict: bool = False,
+    fn: Callable[..., Any],
+    constraints: list[Constraint] | None = None,
 ) -> Callable[..., Any]:
+    if isinstance(fn, torch._ops.OpOverload):
+
+        def default_annotation(arg: torch.Argument) -> str:
+            if arg.has_default_value():
+                return f"={arg.default_value!r}"
+            return ""
+
+        has_kwarg_only = False
+        arg_list = []
+        for arg in fn._schema.arguments:
+            if arg.kwarg_only and not has_kwarg_only:
+                has_kwarg_only = True
+                arg_list.append("*")
+            arg_list.append(arg.name + default_annotation(arg))
+        func_str = f"""
+def op_overload_wrapper({", ".join(arg_list)}):
+    return op({", ".join([f"{arg.name}={arg.name}" for arg in fn._schema.arguments])})
+"""
+        out = {}
+        exec(func_str, {"op": fn}, out)
+        fn = out["op_overload_wrapper"]  # type: ignore[assignment]
+
     def inner(*args: Any, **kwargs: Any) -> Any:
         assert not torch._dynamo.config.install_free_tensors
         with (
+            _compiling_state_context(),
             torch._dynamo.config.patch(
                 replay_side_effects=False, side_effect_replay_policy="warn"
             ),
@@ -766,14 +807,12 @@ def inner(*args: Any, **kwargs: Any) -> Any:
             dynamo_timed("fullgraph_capture"),
         ):
             out = fullgraph_capture(
-                mod,
+                fn,
                 args,
                 kwargs,
                 constraints=constraints,
             )
-        graph_module = create_fx_graph_from_captured_output(out, mod, args, kwargs)
-        if restore_state_dict:
-            _restore_state_dict(mod, graph_module)
+        graph_module = create_fx_graph_from_captured_output(out, fn, args, kwargs)
         return graph_module
 
     return inner
@@ -782,8 +821,8 @@ def inner(*args: Any, **kwargs: Any) -> Any:
 def _dynamo_graph_capture_for_export(
     mod: Callable[..., Any],
     *,
-    constraints: Optional[list[Constraint]] = None,
-    dynamic_shapes: Optional[Union[dict[str, Any], tuple[Any], list[Any]]] = None,
+    constraints: list[Constraint] | None = None,
+    dynamic_shapes: dict[str, Any] | tuple[Any] | list[Any] | None = None,
 ) -> Callable[..., torch.fx.GraphModule]:
     """
     Improved dynamo graph capture using transformer approach with proper fake tensor handling.
@@ -817,8 +856,8 @@ def inner(*args: Any, **kwargs: Any) -> torch.fx.GraphModule:
             module_to_trace = ModuleToTrace(mod, in_spec)
             orig_callable = mod.forward if isinstance(mod, torch.nn.Module) else mod
 
-            constraints: Optional[list[Constraint]] = _constraints
-            dynamic_shapes: Optional[Union[dict[str, Any], tuple[Any], list[Any]]] = (
+            constraints: list[Constraint] | None = _constraints
+            dynamic_shapes: dict[str, Any] | tuple[Any] | list[Any] | None = (
                 _dynamic_shapes
             )
 
@@ -900,11 +939,14 @@ def inner(*args: Any, **kwargs: Any) -> torch.fx.GraphModule:
             ]
 
             # Create input order mapping from dynamo's internal order to user order
+            # Only process inputs that come from function arguments (GetItemSource).
+            # Skip inputs that come from other sources like closures (e.g., captured
+            # opaque objects like DeviceMesh).
             graph_input_order: dict[int, int] = {}
             for inp in graph_inputs:
                 source = graph_inputs[inp]
-                assert isinstance(source, torch._dynamo.source.GetItemSource), source
-                graph_input_order[source.index] = len(graph_input_order)
+                if isinstance(source, torch._dynamo.source.GetItemSource):
+                    graph_input_order[source.index] = len(graph_input_order)
 
             for real_idx, graph_idx in graph_input_order.items():
                 flat_inputs[real_idx] = example_inputs[graph_idx]
@@ -917,6 +959,7 @@ def inner(*args: Any, **kwargs: Any) -> torch.fx.GraphModule:
                 graph_input_order,
                 graph_output_map,
                 fake_mode,
+                graph_inputs,
             ).transform()
 
             # Set up PyTree codegen for proper input/output handling
@@ -929,9 +972,7 @@ def inner(*args: Any, **kwargs: Any) -> torch.fx.GraphModule:
             )
             transformed_graph.recompile()
 
-            clean_nn_module_stack_and_source_fn(
-                transformed_graph, torch._dynamo.config.inline_inbuilt_nn_modules
-            )
+            clean_nn_module_stack_and_source_fn(transformed_graph, True)
             clean_export_root(transformed_graph)
 
             transformed_graph.meta["module_call_specs"] = module_call_spec
diff --git a/torch/_dynamo/graph_break_registry.json b/torch/_dynamo/graph_break_registry.json
index 08bdbfcbc9dfa..46177b12200f7 100644
--- a/torch/_dynamo/graph_break_registry.json
+++ b/torch/_dynamo/graph_break_registry.json
@@ -322,6 +322,16 @@
     }
   ],
   "GB0031": [
+    {
+      "Gb_type": "Custom __getattribute__ in nn.Module attribute access",
+      "Context": "var_getattr {self} {name}",
+      "Explanation": "Dynamo could not trace through the custom `__getattribute__` method on this `nn.Module`.",
+      "Hints": [
+        "Simplify your `__getattribute__` implementation, ",
+        "or replace it with a targeted `@property`.",
+        "It may be possible to write Dynamo tracing rules for this code. Please report an issue to PyTorch if you encounter this graph break often and it is causing performance issues."
+      ]
+    },
     {
       "Gb_type": "Custom __getattribute__ in nn.Module attribute access",
       "Context": "var_getattr {self} {name}",
@@ -386,6 +396,16 @@
       ]
     }
   ],
+  "GB6999": [
+    {
+      "Gb_type": "unsupported variable type for __dict__ access",
+      "Context": "VariableTracker type: {type(vt)}",
+      "Explanation": "Dynamo does not know how to get __dict__ from {type(vt)}",
+      "Hints": [
+        "This is likely to be a Dynamo bug. Please report an issue to PyTorch."
+      ]
+    }
+  ],
   "GB0037": [
     {
       "Gb_type": "Dynamic shape operator (no meta kernel)",
@@ -557,6 +577,16 @@
       ]
     }
   ],
+  "GB4880": [
+    {
+      "Gb_type": "missing_mp_subscript",
+      "Context": "mp_subscript_impl not defined for {type(self).__name__}",
+      "Explanation": "Dynamo does not yet support subscripting '{self.python_type_name()}'.",
+      "Hints": [
+        "It may be possible to write Dynamo tracing rules for this code. Please report an issue to PyTorch if you encounter this graph break often and it is causing performance issues."
+      ]
+    }
+  ],
   "GB9741": [
     {
       "Gb_type": "skip frame due to being in functorh mode",
@@ -630,6 +660,16 @@
       "Hints": []
     }
   ],
+  "GB5195": [
+    {
+      "Gb_type": "elementwise_dtypes unsupported arg type",
+      "Context": "str(arg)",
+      "Explanation": "elementwise_dtypes requires tensor or constant arguments, got {type(arg).__name__}",
+      "Hints": [
+        "It may be possible to write Dynamo tracing rules for this code. Please report an issue to PyTorch if you encounter this graph break often and it is causing performance issues."
+      ]
+    }
+  ],
   "GB0054": [
     {
       "Gb_type": "Failed to construct Enum variable",
@@ -669,6 +709,18 @@
       ]
     }
   ],
+  "GB7299": [
+    {
+      "Gb_type": "_autograd_grad with lost grad_fn linkage",
+      "Context": "outputs lost autograd linkage during tracing",
+      "Explanation": "_autograd_grad() received tensors whose grad_fn was lost during tracing - this silently produces zero gradients.",
+      "Hints": [
+        "Compile the full transform instead of the returned ",
+        "closure: torch.compile(lambda x: torch.func.vjp(f, x))",
+        "It may be possible to write Dynamo tracing rules for this code. Please report an issue to PyTorch if you encounter this graph break often and it is causing performance issues."
+      ]
+    }
+  ],
   "GB0058": [
     {
       "Gb_type": "Failed to set tensor attribute",
@@ -694,6 +746,16 @@
       ]
     }
   ],
+  "GB6270": [
+    {
+      "Gb_type": "Unhandled tensor method",
+      "Context": "call_method {self} {name} {args} {kwargs}",
+      "Explanation": "Tensor method `{name}` is not defined on {check_type.__name__} and does not have an explicit handler in TensorVariable.",
+      "Hints": [
+        "It may be possible to write Dynamo tracing rules for this code. Please report an issue to PyTorch if you encounter this graph break often and it is causing performance issues."
+      ]
+    }
+  ],
   "GB0060": [
     {
       "Gb_type": "Failed to trace unittest method",
@@ -708,7 +770,7 @@
   "GB8656": [
     {
       "Gb_type": "Opaque object with custom __getattr__ not supported",
-      "Context": "{value_type.__name__} with custom __getattr__",
+      "Context": "{real_obj_type.__name__} with custom __getattr__",
       "Explanation": "Dynamo does not support opaque objects types with custom __getattr__ methods",
       "Hints": []
     }
@@ -866,6 +928,22 @@
       ]
     }
   ],
+  "GB7230": [
+    {
+      "Gb_type": "iter() with no arguments",
+      "Context": "iter()",
+      "Explanation": "iter() requires at least one argument",
+      "Hints": [
+        "Your code may result in an error when running in eager. Please double check that your code doesn't contain a similar error when actually running eager/uncompiled. You can do this by removing the `torch.compile` call, or by using `torch.compiler.set_stance(\"force_eager\")`. "
+      ]
+    },
+    {
+      "Gb_type": "iter() with no arguments",
+      "Context": "iter()",
+      "Explanation": "iter() requires at least one argument",
+      "Hints": []
+    }
+  ],
   "GB0073": [
     {
       "Gb_type": "Invoking an nn.Module inside a HigherOrderOperator",
@@ -882,6 +960,16 @@
       "Hints": []
     }
   ],
+  "GB6012": [
+    {
+      "Gb_type": "Cannot trace user-defined __len__",
+      "Context": "{self.python_type_name()}.__len__()",
+      "Explanation": "Dynamo cannot trace len() on {self.python_type_name()} because the __len__ method is either not traceable (e.g., defined in C or built-in) or returns a non-constant value.",
+      "Hints": [
+        "It may be possible to write Dynamo tracing rules for this code. Please report an issue to PyTorch if you encounter this graph break often and it is causing performance issues."
+      ]
+    }
+  ],
   "GB0075": [
     {
       "Gb_type": "LOAD_BUILD_CLASS bytecode not supported",
@@ -924,6 +1012,16 @@
       ]
     }
   ],
+  "GB3775": [
+    {
+      "Gb_type": "nb_int_impl not implemented",
+      "Context": "{type(self).__name__} has nb_int slot but no nb_int_impl override",
+      "Explanation": "The type {self.python_type_name()} has an nb_int C slot but the corresponding VariableTracker doesn't implement nb_int_impl.",
+      "Hints": [
+        "It may be possible to write Dynamo tracing rules for this code. Please report an issue to PyTorch if you encounter this graph break often and it is causing performance issues."
+      ]
+    }
+  ],
   "GB0079": [
     {
       "Gb_type": "Missing CALL_INTRINSIC_1 handler",
@@ -1030,6 +1128,14 @@
     }
   ],
   "GB0088": [
+    {
+      "Gb_type": "Observed exception",
+      "Context": "raised exception {curr_exc.debug_repr()}",
+      "Explanation": "observed_exn_gb_explanation",
+      "Hints": [
+        "Your code may result in an error when running in eager. Please double check that your code doesn't contain a similar error when actually running eager/uncompiled. You can do this by removing the `torch.compile` call, or by using `torch.compiler.set_stance(\"force_eager\")`. "
+      ]
+    },
     {
       "Gb_type": "Observed exception",
       "Context": "raised exception {curr_exc.python_type_name()}({curr_exc.args})",
@@ -1377,6 +1483,16 @@
       "Hints": []
     }
   ],
+  "GB7650": [
+    {
+      "Gb_type": "torch.Generator method",
+      "Context": "torch.Generator.{name}",
+      "Explanation": "torch.Generator.{name}() is a stateful RNG operation that cannot be soundly traced in the FX graph.",
+      "Hints": [
+        "This graph break is fundamental - it is unlikely that Dynamo will ever be able to trace through your code. Consider finding a workaround."
+      ]
+    }
+  ],
   "GB0118": [
     {
       "Gb_type": "Unimplemented next() call",
@@ -1446,6 +1562,16 @@
       ]
     }
   ],
+  "GB4695": [
+    {
+      "Gb_type": "torch.jit.script/freeze modules unsupported",
+      "Context": "str(value)",
+      "Explanation": "Dynamo does not support tracing into torch.jit.script or torch.jit.freeze modules because they execute in the TorchScript runtime, not Python. Replace the ScriptModule submodule with the original eager nn.Module.",
+      "Hints": [
+        "This graph break is fundamental - it is unlikely that Dynamo will ever be able to trace through your code. Consider finding a workaround."
+      ]
+    }
+  ],
   "GB0124": [
     {
       "Gb_type": "Unsupported Tensor.item() call with capture_scalar_outputs=False",
@@ -1642,6 +1768,22 @@
       ]
     }
   ],
+  "GB1380": [
+    {
+      "Gb_type": "Illegal __getitem__ invocation in strict mode",
+      "Context": "mp_subscript_impl {self} {key}",
+      "Explanation": "Dynamo currently does not support __getitem__ invocation in strict mode.",
+      "Hints": []
+    }
+  ],
+  "GB4240": [
+    {
+      "Gb_type": "unsupported __getitem__",
+      "Context": "mp_subscript_impl {self} {key}",
+      "Explanation": "Dynamo does not know how to handle __getitem__ on {self}",
+      "Hints": []
+    }
+  ],
   "GB0141": [
     {
       "Gb_type": "Unsupported call_id() without source",
@@ -1800,6 +1942,12 @@
     }
   ],
   "GB0153": [
+    {
+      "Gb_type": "Unsupported key type for nn.Module.__getitem__",
+      "Context": "mp_subscript_impl: {self} {key}",
+      "Explanation": "Dynamo does not support getitem on `nn.Module` with non-constant key.",
+      "Hints": []
+    },
     {
       "Gb_type": "Unsupported key type for nn.Module.__getitem__",
       "Context": "call_method: {self} {name} {args} {kwargs}",
@@ -1912,6 +2060,16 @@
       ]
     }
   ],
+  "GB4841": [
+    {
+      "Gb_type": "Unsupported python_type() call",
+      "Context": "{obj} does not implement python_type()",
+      "Explanation": "This VariableTracker does not implement python_type(), which is required for object protocol operations.",
+      "Hints": [
+        "This is likely to be a Dynamo bug. Please report an issue to PyTorch."
+      ]
+    }
+  ],
   "GB8335": [
     {
       "Gb_type": "Generator reconstruction with mutations",
@@ -2043,6 +2201,12 @@
     }
   ],
   "GB0170": [
+    {
+      "Gb_type": "Data-dependent branching",
+      "Context": "attempted to jump with {value}",
+      "Explanation": "\"Detected data-dependent branching (e.g. `if my_tensor.sum() > 0:`). \"             \"Dynamo does not support tracing dynamic control flow.\" + trace_info",
+      "Hints": []
+    },
     {
       "Gb_type": "Data-dependent branching",
       "Context": "attempted to jump with {value}",
@@ -2414,6 +2578,15 @@
     }
   ],
   "GB8843": [
+    {
+      "Gb_type": "Opaque object member with method-type USE_REAL returned a reference-type opaque object.",
+      "Context": "Opaque object type: {real_obj_type}. Method name: '{name}'",
+      "Explanation": "To properly guard reference-type opaque objects, we must lift them as inputs to the graph. In order to do this, they must all have a source, meaning they come from a global value or are an attribute of an input.",
+      "Hints": [
+        "Register member '{name}' with MemberType.INLINED in ",
+        "register_opaque_type({real_obj_type}, members=...)."
+      ]
+    },
     {
       "Gb_type": "Opaque object member with method-type USE_REAL returned a reference-type opaque object.",
       "Context": "Opaque object type: {value_type}. Method name: '{name}'",
@@ -2510,6 +2683,17 @@
       "Hints": []
     }
   ],
+  "GB2520": [
+    {
+      "Gb_type": "assume_constant_result argument conversion failed",
+      "Context": "function {name}, variable type {type(x).__name__}",
+      "Explanation": "Cannot convert argument of type {type(x).__name__} to a Python constant for function {name} marked with torch._dynamo.assume_constant_result. The variable tracker does not support constant conversion.",
+      "Hints": [
+        "Remove torch._dynamo.assume_constant_result from this function",
+        "Ensure all arguments passed to the function can be converted to constants"
+      ]
+    }
+  ],
   "GB0209": [
     {
       "Gb_type": "torch.autograd._unsafe_preserve_version_counter escaped from compiled region",
@@ -2530,6 +2714,16 @@
       ]
     }
   ],
+  "GB2977": [
+    {
+      "Gb_type": "wrap_jit_function: missing _torchdynamo_inline",
+      "Context": "type: {type(value).__name__}",
+      "Explanation": "Dynamo expected a JIT function with a _torchdynamo_inline attribute, but the object does not have one.",
+      "Hints": [
+        "It may be possible to write Dynamo tracing rules for this code. Please report an issue to PyTorch if you encounter this graph break often and it is causing performance issues."
+      ]
+    }
+  ],
   "GB0211": [
     {
       "Gb_type": "torch.nn.Module with a non-function custom __getattr__",
@@ -2813,6 +3007,16 @@
       ]
     }
   ],
+  "GB6025": [
+    {
+      "Gb_type": "Failed to trace list()",
+      "Context": "list({arg_types})",
+      "Explanation": "Dynamo does not know how to construct a list from argument types {arg_types}",
+      "Hints": [
+        "It may be possible to write Dynamo tracing rules for this code. Please report an issue to PyTorch if you encounter this graph break often and it is causing performance issues."
+      ]
+    }
+  ],
   "GB0236": [
     {
       "Gb_type": "Invalid input type for nonstrict_trace-ed function",
@@ -3021,6 +3225,16 @@
       ]
     }
   ],
+  "GB9860": [
+    {
+      "Gb_type": "P2POp mutation",
+      "Context": "object={self}, name={name}, value={value}",
+      "Explanation": "Dynamo does not support mutating torch.distributed.P2POp instances.",
+      "Hints": [
+        "Construct a new torch.distributed.P2POp instead of mutating an existing one inside torch.compile."
+      ]
+    }
+  ],
   "GB0253": [
     {
       "Gb_type": "call_method on generator",
@@ -3119,6 +3333,16 @@
     }
   ],
   "GB4927": [
+    {
+      "Gb_type": "autograd.grad consumed returned tensor's grad_fn",
+      "Context": "Leaked output tensors: {leaked_str}",
+      "Explanation": "torch.autograd.grad() consumes grad_fns that are needed by tensors returned from this compiled function. This would cause 'backward through graph a second time' errors.\n  The following returned tensors have consumed grad_fns: {leaked_str}",
+      "Hints": [
+        "Detach the problematic tensor(s) before returning: e.g. `{leaked[0]}.detach()`",
+        "Call .detach() on the tensor before returning.",
+        "If you need to backward through the returned tensor, use retain_graph=True in autograd.grad()."
+      ]
+    },
     {
       "Gb_type": "autograd.grad consumed returned tensor's grad_fn",
       "Context": "",
@@ -3327,6 +3551,16 @@
       ]
     }
   ],
+  "GB6604": [
+    {
+      "Gb_type": "Missing CALL_INTRINSIC_2 handler",
+      "Context": "CALL_INTRINSIC_2 operand: {inst.argval}",
+      "Explanation": "No handler implemented for CALL_INTRINSIC_2 {inst.argval} instruction.",
+      "Hints": [
+        "It may be possible to write Dynamo tracing rules for this code. Please report an issue to PyTorch if you encounter this graph break often and it is causing performance issues."
+      ]
+    }
+  ],
   "GB0276": [
     {
       "Gb_type": "Failed to make weakref to User Object",
@@ -3365,6 +3599,20 @@
       ]
     }
   ],
+  "GB7324": [
+    {
+      "Gb_type": "returning intermediate with requires_grad_()",
+      "Context": "graph output depends on source-less requires_grad_()",
+      "Explanation": "msg",
+      "Hints": [
+        "If you only need the tensor values without gradients, ",
+        "call .detach() before returning.",
+        "Consume the gradient inside the compiled function ",
+        "(call backward() and use .grad), ",
+        "or move requires_grad_() outside torch.compile."
+      ]
+    }
+  ],
   "GB0280": [
     {
       "Gb_type": "1-arg super not implemented",
@@ -3375,7 +3623,23 @@
       ]
     }
   ],
+  "GB3383": [
+    {
+      "Gb_type": "Pydantic dataclass constructor",
+      "Context": "{self.value}",
+      "Explanation": "Dynamo graph breaks on pydantic dataclass constructors because validation mutates the instance outside traced bytecode.",
+      "Hints": []
+    }
+  ],
   "GB0281": [
+    {
+      "Gb_type": "Invalid or non-const argument in nn.Module __getitem__",
+      "Context": "mp_subscript_impl: {self} {key}",
+      "Explanation": "Dynamo does not support calling method `__getitem__` of ``nn.Module`` {module} with a non-constant or non-(str, int) key.",
+      "Hints": [
+        "Use constant arguments of type str or int for __getitem__"
+      ]
+    },
     {
       "Gb_type": "Invalid or non-const argument in nn.Module __getitem__",
       "Context": "call_method: {self} {name} {args} {kwargs}",
@@ -3505,6 +3769,16 @@
       ]
     }
   ],
+  "GB7906": [
+    {
+      "Gb_type": "Missing len_impl",
+      "Context": "len({type(self).__name__})",
+      "Explanation": "Dynamo does not support len() on {type(self).__name__}. Add len_impl to this VariableTracker subclass.",
+      "Hints": [
+        "It may be possible to write Dynamo tracing rules for this code. Please report an issue to PyTorch if you encounter this graph break often and it is causing performance issues."
+      ]
+    }
+  ],
   "GB0291": [
     {
       "Gb_type": "logging.Logger method not supported for non-export cases",
@@ -3614,6 +3888,54 @@
       ]
     }
   ],
+  "GB9300": [
+    {
+      "Gb_type": "P2P compilation disabled for batch_isend_irecv",
+      "Context": "{self.fn}",
+      "Explanation": "P2P compilation is disabled.",
+      "Hints": [
+        "Set TORCHDYNAMO_ENABLE_P2P_COMPILATION=1 to enable."
+      ]
+    }
+  ],
+  "GB2696": [
+    {
+      "Gb_type": "P2P compilation disabled for isend/irecv",
+      "Context": "{self.fn}",
+      "Explanation": "P2P compilation is disabled.",
+      "Hints": [
+        "Set TORCHDYNAMO_ENABLE_P2P_COMPILATION=1 to enable."
+      ]
+    }
+  ],
+  "GB9277": [
+    {
+      "Gb_type": "unsupported function.__get__ call",
+      "Context": "call_function {self}, args: {args}, kwargs: {kwargs}",
+      "Explanation": "Dynamo only supports function.__get__(func, obj[, cls]) where func is a user-defined function.",
+      "Hints": [
+        "It may be possible to write Dynamo tracing rules for this code. Please report an issue to PyTorch if you encounter this graph break often and it is causing performance issues."
+      ]
+    }
+  ],
+  "GB1750": [
+    {
+      "Gb_type": "P2P compilation disabled for P2POp construction",
+      "Context": "torch.distributed.P2POp",
+      "Explanation": "P2P compilation is disabled.",
+      "Hints": [
+        "Set TORCHDYNAMO_ENABLE_P2P_COMPILATION=1 to enable."
+      ]
+    }
+  ],
+  "GB3343": [
+    {
+      "Gb_type": "Cannot reconstruct P2POpVariable",
+      "Context": "P2POpVariable.reconstruct",
+      "Explanation": "P2POpVariable does not support reconstruction.",
+      "Hints": []
+    }
+  ],
   "GB0301": [
     {
       "Gb_type": "HOP: non torch.Tensor leaf",
@@ -3719,6 +4041,14 @@
       ]
     }
   ],
+  "GB5458": [
+    {
+      "Gb_type": "Custom metaclass with __getattribute__",
+      "Context": "type({self.value}) = {metacls}",
+      "Explanation": "Dynamo does not trace attribute access on classes whose metaclass overrides __getattribute__",
+      "Hints": []
+    }
+  ],
   "GB0310": [
     {
       "Gb_type": "torch.cond: improper predicate",
@@ -3932,6 +4262,16 @@
       ]
     }
   ],
+  "GB1140": [
+    {
+      "Gb_type": "_autograd_grad with unsupported argument type",
+      "Context": "got {type(var).__name__}",
+      "Explanation": "_autograd_grad() received an argument of type {type(var).__name__} which is not supported. Expected tensor or sequence of tensors.",
+      "Hints": [
+        "Ensure outputs and inputs arguments are tensors or sequences of tensors."
+      ]
+    }
+  ],
   "GB0325": [
     {
       "Gb_type": "torch.map: kwargs not supported",
@@ -3952,6 +4292,16 @@
       ]
     }
   ],
+  "GB8435": [
+    {
+      "Gb_type": "missing_mp_subscript_impl",
+      "Context": "mp_subscript_impl not defined for {type(self).__name__}",
+      "Explanation": "'{self.python_type_name()}' subscript is not yet supported by Dynamo.",
+      "Hints": [
+        "It may be possible to write Dynamo tracing rules for this code. Please report an issue to PyTorch if you encounter this graph break often and it is causing performance issues."
+      ]
+    }
+  ],
   "GB0327": [
     {
       "Gb_type": "executorch_call_delegate: kwargs not supported",
@@ -4116,6 +4466,26 @@
       ]
     }
   ],
+  "GB4501": [
+    {
+      "Gb_type": "Reconstruction of FakeIdVariable",
+      "Context": "str(self.value)",
+      "Explanation": "A fake id produced by id() on a compile-time container cannot be reconstructed across a graph break.",
+      "Hints": [
+        "Avoid using id() on containers in code that may graph-break."
+      ]
+    }
+  ],
+  "GB4198": [
+    {
+      "Gb_type": "Attempted to copy.deepcopy a tensor",
+      "Context": "copy.deepcopy({self})",
+      "Explanation": "Dynamo does not support copy.deepcopy() on tensors.",
+      "Hints": [
+        "Avoid calling copy.deepcopy() on tensors inside compiled regions."
+      ]
+    }
+  ],
   "GB0344": [
     {
       "Gb_type": "wrap_with_autocast: expected constant arg",
@@ -4127,6 +4497,14 @@
     }
   ],
   "GB0345": [
+    {
+      "Gb_type": "strict_mode: improper args",
+      "Context": "args: {args}, kwargs: {kwargs}",
+      "Explanation": "strict_mode higher order op expects flat inputs (list/tuple/dict/set)",
+      "Hints": [
+        "Your code may result in an error when running in eager. Please double check that your code doesn't contain a similar error when actually running eager/uncompiled. You can do this by removing the `torch.compile` call, or by using `torch.compiler.set_stance(\"force_eager\")`. "
+      ]
+    },
     {
       "Gb_type": "strict_mode: improper args",
       "Context": "args: {args}, kwargs: {kwargs}",
@@ -4168,6 +4546,14 @@
       "Hints": []
     }
   ],
+  "GB8887": [
+    {
+      "Gb_type": "NamedTupleVariable.__setattr__ bad args",
+      "Context": "{len(args)} args and {len(kwargs)} kwargs",
+      "Explanation": "Expected exactly 2 positional args for __setattr__.",
+      "Hints": []
+    }
+  ],
   "GB0349": [
     {
       "Gb_type": "cannot unwrap variable for check_meta_consistency",
@@ -4266,6 +4652,19 @@
       ]
     }
   ],
+  "GB8985": [
+    {
+      "Gb_type": "requires_grad_() intermediate leaked as output",
+      "Context": "call_method {self} requires_grad_",
+      "Explanation": "An intermediate tensor with requires_grad_() called on it (or a tensor derived from it) is returned from the compiled region. Graph breaking here to preserve partial acceleration.",
+      "Hints": [
+        "Call .detach() before returning if you only need values.",
+        "Consume the gradient inside the compiled function ",
+        "(call backward() and use .grad), ",
+        "or move requires_grad_() outside torch.compile."
+      ]
+    }
+  ],
   "GB0358": [
     {
       "Gb_type": "optimizer: pending mutation on parameter",
@@ -4292,6 +4691,16 @@
       "Hints": []
     }
   ],
+  "GB7286": [
+    {
+      "Gb_type": "nb_float_impl not implemented",
+      "Context": "{type(self).__name__} has nb_float slot but no nb_float_impl override",
+      "Explanation": "The type {self.python_type_name()} has an nb_float C slot but the corresponding VariableTracker doesn't implement nb_float_impl.",
+      "Hints": [
+        "It may be possible to write Dynamo tracing rules for this code. Please report an issue to PyTorch if you encounter this graph break often and it is causing performance issues."
+      ]
+    }
+  ],
   "GB0361": [
     {
       "Gb_type": "triton kernel unsupported feature",
diff --git a/torch/_dynamo/graph_bytecode_inputs.py b/torch/_dynamo/graph_bytecode_inputs.py
index b70ab1002e051..921757b0645a1 100644
--- a/torch/_dynamo/graph_bytecode_inputs.py
+++ b/torch/_dynamo/graph_bytecode_inputs.py
@@ -30,6 +30,15 @@ def stash_graph_created_object(obj: Any) -> Any:
     return obj
 
 
+CURRENT_STREAM_INDEX = 0
+
+
+def set_external_object_by_index(index: int, value: Any) -> None:
+    """Update an entry in the external object registry at runtime."""
+    keep_alive.append(value)
+    index_to_external_object_weakref[index] = weakref.ref(value)
+
+
 def get_external_object_by_index(index: int) -> Any:
     assert index in index_to_external_object_weakref, (
         "Index not registered in index_to_user_object_weakref"
diff --git a/torch/_dynamo/graph_deduplication.py b/torch/_dynamo/graph_deduplication.py
index 68c6a996fce99..0fd2c2900148c 100644
--- a/torch/_dynamo/graph_deduplication.py
+++ b/torch/_dynamo/graph_deduplication.py
@@ -11,7 +11,6 @@
 import operator
 from collections import defaultdict, deque
 from collections.abc import Generator, Iterable
-from typing import Optional
 
 import torch
 import torch.fx
@@ -31,7 +30,7 @@
 
 log = logging.getLogger(__name__)
 
-last_node_to_additional_deps: Optional[dict[Node, OrderedSet[Node]]] = None
+last_node_to_additional_deps: dict[Node, OrderedSet[Node]] | None = None
 
 
 def apply_graph_deduplication(output_graph) -> dict[str, torch.fx.GraphModule]:  # type: ignore[no-untyped-def]
diff --git a/torch/_dynamo/graph_id_filter.py b/torch/_dynamo/graph_id_filter.py
index 59d7820fc5b90..173005605b2db 100644
--- a/torch/_dynamo/graph_id_filter.py
+++ b/torch/_dynamo/graph_id_filter.py
@@ -5,7 +5,8 @@
 import functools
 import logging
 import re
-from typing import Any, Generic, Optional, TYPE_CHECKING, TypeVar
+import warnings
+from typing import Any, Generic, TYPE_CHECKING, TypeVar
 
 
 if TYPE_CHECKING:
@@ -16,11 +17,6 @@
 
 log = logging.getLogger(__name__)
 
-# Valid modes for inductor backend
-_INDUCTOR_MODES = frozenset(
-    {"default", "reduce-overhead", "max-autotune", "max-autotune-no-cudagraphs"}
-)
-
 
 class GraphIdFilter:
     """
@@ -106,13 +102,13 @@ class _GraphRouterBase(Generic[T]):
 
     def __init__(self, config_str: str, rule_type: str) -> None:
         self._rules: list[tuple[GraphIdFilter, T]] = []
-        self._values: list[Optional[T]] = []
-        self._overflow_value: Optional[T] = None
+        self._values: list[T | None] = []
+        self._overflow_value: T | None = None
         self._rule_type = rule_type
         self._parse(config_str)
         self._precompute()
 
-    def _parse_value_str(self, value_str: str) -> Optional[T]:
+    def _parse_value_str(self, value_str: str) -> T | None:
         """Parse a value string into the appropriate type. Returns None to skip."""
         raise NotImplementedError
 
@@ -165,13 +161,13 @@ def _precompute(self) -> None:
         # For IDs > max_id, the result is constant (only unbounded conditions apply)
         self._overflow_value = self._match_rules(max_id + 1)
 
-    def _match_rules(self, graph_id: int) -> Optional[T]:
+    def _match_rules(self, graph_id: int) -> T | None:
         for f, value in self._rules:
             if graph_id in f:
                 return value
         return None
 
-    def get_value_for_graph(self, graph_id: int) -> Optional[T]:
+    def get_value_for_graph(self, graph_id: int) -> T | None:
         """Get the value for a given graph ID. Returns None if no rule matches."""
         if graph_id < len(self._values):
             return self._values[graph_id]
@@ -189,50 +185,47 @@ class GraphBackendRouter(_GraphRouterBase[Any]):
     The router parses a configuration string with rules in the format:
         "filter1:backend1;filter2:backend2;..."
 
-    Rules are evaluated in order, and the first matching rule wins.
+    If a graph ID matches multiple rules with different backends, a ValueError
+    is raised.
 
     Examples:
         "0-5:eager;>5:inductor"     - IDs 0-5 use eager, rest use inductor
         ">10:aot_eager"             - IDs > 10 use aot_eager
         "<=3:eager;4-10:aot_eager"  - IDs 0-3 use eager, 4-10 use aot_eager
 
-    Special backend values:
-        "eager"                     - Run in eager mode (no compilation)
-        "aot_eager"                 - AOT Autograd with eager execution
-        "aot_eager_decomp_partition" - AOT Autograd with partitioner and decomps
-        "inductor"                  - Default inductor backend
-        "inductor:reduce-overhead"  - Inductor with reduce-overhead mode (cudagraphs)
-        "inductor:max-autotune"     - Inductor with max-autotune mode
+    Supported backends include "eager", "aot_eager", "aot_eager_decomp_partition",
+    "inductor", and any other registered backend.
     """
 
     def __init__(self, config_str: str) -> None:
+        self._backend_names: dict[int, str] = {}
         super().__init__(config_str, "backend")
 
-    def _parse_value_str(self, value_str: str) -> Optional[Any]:
-        """Look up a backend, supporting 'backend:mode' format for inductor."""
-        import torch
-
+    def _parse_value_str(self, value_str: str) -> Any | None:
+        """Look up a backend by name."""
         from .backends.registry import lookup_backend
         from .eval_frame import cached_backends
 
-        backend: Any = None
-        if ":" in value_str:
-            parts = value_str.split(":", 1)
-            backend_name, mode = parts[0], parts[1]
-
-            if backend_name == "inductor" and mode in _INDUCTOR_MODES:
-                backend = torch._TorchCompileInductorWrapper(
-                    mode=mode, options=None, dynamic=None
-                )
-
-        if backend is None:
-            backend = lookup_backend(value_str)
+        backend = lookup_backend(value_str)
 
         # Register the backend so its reset() is called during torch._dynamo.reset()
         assert backend is not None, "Invalid override backend: " + value_str
         cached_backends.setdefault(id(backend), backend)
+        self._backend_names[id(backend)] = value_str
         return backend
 
+    def _match_rules(self, graph_id: int) -> Any | None:
+        """Match rules with conflict detection for overlapping filters."""
+        matches = {id(backend): backend for f, backend in self._rules if graph_id in f}
+        if len(matches) > 1:
+            names = [self._backend_names[bid] for bid in matches]
+            raise ValueError(
+                f"Conflicting backend override for graph {graph_id}: matched {names}"
+            )
+        if matches:
+            return next(iter(matches.values()))
+        return None
+
     def __repr__(self) -> str:
         if not self._rules:
             return "GraphBackendRouter(empty)"
@@ -246,12 +239,17 @@ class GraphConfigRouter(_GraphRouterBase[dict[str, Any]]):
     The router parses a configuration string with rules in the format:
         "filter1:config1;filter2:config2;..."
 
-    Rules are evaluated in order, and the first matching rule wins.
+    All matching rules are aggregated: configs from all matching rules are merged
+    into a single dict. Conflicting keys (same key, different value) raise an error.
     Config format is "key=value" or "key1=value1,key2=value2" for multiple settings.
 
     Examples:
         "0-5:triton.cudagraph_skip_dynamic_graphs=False"
         ">10:triton.cudagraphs=False,triton.cudagraph_trees=False"
+
+        With "0:a=1;>=0:b=2", graph 0 gets {"a": 1, "b": 2} (both rules match).
+        With "0:a=1;>=0:a=2", graph 0 raises an error (conflicting values for "a").
+        With "0:a=1;>=0:a=1", graph 0 gets {"a": 1} (same value is not a conflict).
     """
 
     def __init__(self, config_str: str) -> None:
@@ -274,7 +272,21 @@ def _parse_scalar_value(value_str: str) -> Any:
         except ValueError:
             return value_str
 
-    def _parse_value_str(self, value_str: str) -> Optional[dict[str, Any]]:
+    def _match_rules(self, graph_id: int) -> dict[str, Any] | None:
+        """Aggregate configs from all matching rules. Conflicts raise an error."""
+        result: dict[str, Any] = {}
+        for f, value in self._rules:
+            if graph_id in f:
+                for k, v in value.items():
+                    if k in result and result[k] != v:
+                        raise ValueError(
+                            f"Conflicting config override for graph {graph_id}: "
+                            f"key '{k}' has value {result[k]!r} and {v!r}"
+                        )
+                    result[k] = v
+        return result if result else None
+
+    def _parse_value_str(self, value_str: str) -> dict[str, Any] | None:
         """Parse a config string like 'key1=val1,key2=val2' into a dict."""
         result: dict[str, Any] = {}
         for item in value_str.split(","):
@@ -295,11 +307,11 @@ def __repr__(self) -> str:
 
 
 def _get_override_for_compile_id(
-    compile_id: Optional[CompileId],
+    compile_id: CompileId | None,
     config_str: str,
     create_router: Callable[[str], _GraphRouterBase[T]],
-    log_msg: str,
-) -> Optional[T]:
+    label: str,
+) -> T | None:
     """
     Get the override value for a given CompileId.
 
@@ -315,7 +327,7 @@ def _get_override_for_compile_id(
     router = create_router(config_str)
     value = router.get_value_for_graph(graph_id)
     if value is not None:
-        log.info(log_msg, compile_id, graph_id, value)
+        log.info("Overriding %s: %s", label, value)
     return value
 
 
@@ -326,13 +338,88 @@ def _create_backend_router(config_str: str) -> GraphBackendRouter:
 
 
 @functools.lru_cache
-def _create_config_router(config_str: str) -> GraphConfigRouter:
-    """Create and cache GraphConfigRouter instances based on config string."""
+def _validate_backend_names(config_str: str) -> str | None:
+    """Return an error message if any backend name is invalid, else None."""
+    if not config_str or not config_str.strip():
+        return None
+    from .backends.registry import lookup_backend
+
+    for rule_str in config_str.split(";"):
+        rule_str = rule_str.strip()
+        if not rule_str or ":" not in rule_str:
+            continue
+        backend_name = rule_str[rule_str.find(":") + 1 :].strip()
+        if not backend_name:
+            continue
+        try:
+            lookup_backend(backend_name)
+        except Exception:
+            return (
+                f"TORCH_COMPILE_OVERRIDE_BACKENDS: "
+                f"'{backend_name}' is not a valid backend, "
+                f"see `torch._dynamo.list_backends()` for available backends"
+            )
+    return None
+
+
+@functools.lru_cache
+def _validate_inductor_config_keys(config_str: str) -> str | None:
+    """Return an error message if any config key is invalid, else None."""
+    router = GraphConfigRouter(config_str)
+    from torch._inductor import config
+
+    for _, config_dict in router._rules:
+        for key in config_dict:
+            if not hasattr(config, key):
+                return (
+                    f"TORCH_COMPILE_OVERRIDE_INDUCTOR_CONFIGS: "
+                    f"'{key}' is not a valid torch._inductor.config option"
+                )
+    return None
+
+
+@functools.lru_cache
+def _validate_dynamo_config_keys(config_str: str) -> str | None:
+    """Return an error message if any config key is invalid, else None."""
+    router = GraphConfigRouter(config_str)
+    from torch._dynamo import config
+
+    for _, config_dict in router._rules:
+        for key in config_dict:
+            if not hasattr(config, key):
+                return (
+                    f"TORCH_COMPILE_OVERRIDE_DYNAMO_CONFIGS: "
+                    f"'{key}' is not a valid torch._dynamo.config option"
+                )
+    return None
+
+
+@functools.lru_cache
+def _create_inductor_config_router(config_str: str) -> GraphConfigRouter:
+    """Create and cache GraphConfigRouter for inductor config overrides."""
     return GraphConfigRouter(config_str)
 
 
+@functools.lru_cache
+def _create_dynamo_config_router(config_str: str) -> GraphConfigRouter:
+    """Create and cache GraphConfigRouter for dynamo config overrides.
+
+    Warns that dynamo config overrides are keyed by frame ID and some configs
+    can affect graph breaks, which may shift frame IDs.
+    """
+    router = GraphConfigRouter(config_str)
+    if not router.is_empty():
+        warnings.warn(
+            "TORCH_COMPILE_OVERRIDE_DYNAMO_CONFIGS is set. Dynamo config overrides are "
+            "keyed by frame ID. Some dynamo configs can affect graph breaks, "
+            "which may alter the number of frames and shift frame IDs, causing "
+            "overrides to target the wrong graphs.",
+        )
+    return router
+
+
 def get_backend_override_for_compile_id(
-    compile_id: Optional[CompileId],
+    compile_id: CompileId | None,
     config_str: str,
 ) -> Any:
     """
@@ -344,14 +431,14 @@ def get_backend_override_for_compile_id(
         compile_id,
         config_str,
         _create_backend_router,
-        "Graph %s (frame_id=%d) overridden to use backend: %s",
+        "torch.compile backend",
     )
 
 
 def get_inductor_config_override_for_compile_id(
-    compile_id: Optional[CompileId],
+    compile_id: CompileId | None,
     config_str: str,
-) -> Optional[dict[str, Any]]:
+) -> dict[str, Any] | None:
     """
     Get the inductor config override for a given CompileId.
 
@@ -360,8 +447,25 @@ def get_inductor_config_override_for_compile_id(
     return _get_override_for_compile_id(
         compile_id,
         config_str,
-        _create_config_router,  # type: ignore[arg-type]
-        "Graph %s (frame_id=%d) overridden with inductor config: %s",
+        _create_inductor_config_router,  # type: ignore[arg-type]
+        "inductor config",
+    )
+
+
+def get_dynamo_config_override_for_compile_id(
+    compile_id: CompileId | None,
+    config_str: str,
+) -> dict[str, Any] | None:
+    """
+    Get the dynamo config override for a given CompileId.
+
+    Returns a dict of config patches to apply, or None if no override applies.
+    """
+    return _get_override_for_compile_id(
+        compile_id,
+        config_str,
+        _create_dynamo_config_router,  # type: ignore[arg-type]
+        "dynamo config",
     )
 
 
diff --git a/torch/_dynamo/graph_region_tracker.py b/torch/_dynamo/graph_region_tracker.py
index 50874466bef89..6df9feb6bad38 100644
--- a/torch/_dynamo/graph_region_tracker.py
+++ b/torch/_dynamo/graph_region_tracker.py
@@ -23,7 +23,7 @@
 import pickle
 from collections import defaultdict, deque
 from dataclasses import fields
-from typing import Any, Optional, TYPE_CHECKING, TypeVar
+from typing import Any, TYPE_CHECKING, TypeVar
 
 import torch._logging
 import torch.fx
@@ -129,7 +129,7 @@ def _extract_args(arg: Any) -> Any:
 
 def _normalize_args(
     node: Node,
-) -> tuple[tuple[str, ...], tuple[Optional[Any], ...]]:
+) -> tuple[tuple[str, ...], tuple[Any | None, ...]]:
     flat_args, _ = tree_flatten(node.args)
     sorted_kwargs = sorted(node.kwargs.items(), key=operator.itemgetter(0))
     sorted_keys = tuple(sorted(node.kwargs.keys()))
@@ -174,8 +174,8 @@ def get_global_state_key() -> GlobalStateKey:
 # of a node
 class BackwardBfsArgIter:
     def __init__(self, origin: Node) -> None:
-        self._cur: Optional[Node] = origin
-        self._queue: deque[Optional[Node]] = deque()
+        self._cur: Node | None = origin
+        self._queue: deque[Node | None] = deque()
 
     @staticmethod
     def create(origin: Node) -> BackwardBfsArgIter:
@@ -186,7 +186,7 @@ def create(origin: Node) -> BackwardBfsArgIter:
         assert it.next()
         return it
 
-    def next(self) -> Optional[Node]:
+    def next(self) -> Node | None:
         ret = self._cur
         if not self._queue:
             self._cur = None
@@ -194,7 +194,7 @@ def next(self) -> Optional[Node]:
             self._cur = self._queue.popleft()
         return ret
 
-    def peek(self) -> Optional[Node]:
+    def peek(self) -> Node | None:
         return self._cur
 
     def add_children(self, node: Node) -> None:
@@ -232,7 +232,7 @@ def __init__(self) -> None:
         self.input_pickler = InputPickler()
 
     def _hash_node(
-        self, filename: str, lineno: int, instruction_pointer: Optional[int], node: Node
+        self, filename: str, lineno: int, instruction_pointer: int | None, node: Node
     ) -> str:
         from torch._inductor.codecache import sha256_hash
 
@@ -389,7 +389,7 @@ def __init__(
         self.ancestors = set(node_to_recursive_ancestors[node])
         self.region = region
 
-    def next_candidate(self) -> Optional[Node]:
+    def next_candidate(self) -> Node | None:
         return self.iter.next()
 
     def will_inclusion_create_cycle(self, node: Node) -> bool:
diff --git a/torch/_dynamo/graph_utils.py b/torch/_dynamo/graph_utils.py
index a7429e83174b8..68aea6f96e102 100644
--- a/torch/_dynamo/graph_utils.py
+++ b/torch/_dynamo/graph_utils.py
@@ -1,5 +1,4 @@
-from collections import deque
-from typing import Any, Optional
+from typing import Any
 
 import torch
 from torch.fx import Graph, map_arg, Node
@@ -35,51 +34,53 @@ def _get_flat_args_unique(
 def _detect_cycles(
     graph: Graph, node_to_additional_deps: dict[Node, OrderedSet[Node]]
 ) -> str:
-    current_path: deque[Node] = deque()
-    current_path_set: set[Node] = set()
-    pending: deque[tuple[Node, Node]] = deque()
+    # States: 0=Unvisited, 1=Visiting, 2=Visited(Safe)
+    state: dict[Node, int] = {}
 
-    def add_to_current_path(node: Node) -> None:
-        current_path.append(node)
-        current_path_set.add(node)
+    for root in reversed(graph.nodes):
+        if root in state:
+            continue
 
-    def pop_current_path() -> None:
-        node = current_path.pop()
-        current_path_set.remove(node)
+        # Stack holds (current_node, children_iterator).
+        # Using an iterator allows us to pause and resume processing a node's children.
+        stack = [(root, iter(_get_flat_args_unique(root, node_to_additional_deps)))]
+        state[root] = 1  # Visiting
 
-    def current_path_head() -> Node:
-        return current_path[-1]
+        while stack:
+            parent, children = stack[-1]
 
-    for origin in graph.find_nodes(op="output"):
-        current_path.clear()
-        current_path_set.clear()
-        add_to_current_path(origin)
-        for child in _get_flat_args_unique(origin, node_to_additional_deps):
-            pending.append((child, origin))
+            try:
+                child = next(children)
 
-        while pending:
-            cur_node, parent = pending.pop()
+                if not isinstance(child, Node):
+                    continue
 
-            # handle backtracking
-            while current_path and current_path_head() != parent:
-                pop_current_path()
+                child_state = state.get(child, 0)
 
-            if not isinstance(cur_node, Node):
-                continue
+                if child_state == 1:
+                    # Back-edge: child is on the current DFS path -> cycle
+                    cycle_path = [node for node, _ in stack] + [child]
+                    return f"cycle detected in path: {cycle_path}"
 
-            if cur_node in current_path_set:
-                current_path.append(cur_node)
-                return f"cycle detected in path: {current_path}"
+                if child_state == 0:
+                    state[child] = 1
+                    stack.append(
+                        (
+                            child,
+                            iter(_get_flat_args_unique(child, node_to_additional_deps)),
+                        )
+                    )
+                # child_state == 2 means already verified safe; skip.
 
-            add_to_current_path(cur_node)
-
-            for child in _get_flat_args_unique(cur_node, node_to_additional_deps):
-                pending.append((child, cur_node))
+            except StopIteration:
+                # All children processed — mark safe and pop.
+                stack.pop()
+                state[parent] = 2
 
     return "no cycle detected"
 
 
-def _graph_device_type(graph: Optional[Graph]) -> str:
+def _graph_device_type(graph: Graph | None) -> str:
     if graph is None:
         return "cpu"
 
diff --git a/torch/_dynamo/guards.py b/torch/_dynamo/guards.py
index 5bfa439bb188c..2687535b316c0 100644
--- a/torch/_dynamo/guards.py
+++ b/torch/_dynamo/guards.py
@@ -26,6 +26,7 @@
 import importlib
 import inspect
 import io
+import itertools
 import logging
 import math
 import pickle
@@ -38,7 +39,7 @@
 from contextlib import contextmanager
 from copy import deepcopy
 from inspect import currentframe
-from typing import Any, NamedTuple, NoReturn, Optional, TYPE_CHECKING, Union
+from typing import Any, NamedTuple, NoReturn, TYPE_CHECKING
 from typing_extensions import LiteralString, TypeAliasType, TypeVar
 from weakref import ReferenceType
 
@@ -178,19 +179,18 @@
     dataclass_fields,
     dict_keys,
     get_current_stream,
-    get_custom_getattr,
     get_torch_function_mode_stack,
     get_torch_function_mode_stack_at,
     guard_failures,
     istype,
     key_is_id,
     key_to_id,
+    normalize_count_iter,
     normalize_range_iter,
     orig_code_map,
     tensor_always_has_static_shape,
     tuple_iterator_getitem,
     tuple_iterator_len,
-    unpatched_nn_module_getattr,
     verify_guard_fn_signature,
 )
 
@@ -199,7 +199,8 @@
     from collections.abc import Callable
 
 
-guard_manager_testing_hook_fn: Optional[Callable[[Any, Any, Any], Any]] = None
+guard_manager_testing_hook_fn: Callable[[Any, Any, Any], Any] | None = None
+_COUNT_ITERATOR_TYPE = type(itertools.count())
 
 try:
     import numpy as np
@@ -270,21 +271,21 @@ class is stored in the Dynamo cache entry, so that the cache entry can
     the check_nopybind from C++.
     """
 
-    def __init__(self, root: Optional[RootGuardManager] = None) -> None:
+    def __init__(self, root: RootGuardManager | None = None) -> None:
         if root is None:
             self.root = RootGuardManager()
         else:
             self.root = root
 
-        self.diff_guard_root: Optional[RootGuardManager] = None
-        self.closure_vars: Optional[dict[str, Any]] = None
-        self.args: Optional[list[str]] = None
+        self.diff_guard_root: RootGuardManager | None = None
+        self.closure_vars: dict[str, Any] | None = None
+        self.args: list[str] | None = None
         self.code_parts: list[str] = []
-        self.verbose_code_parts: Optional[list[str]] = None
-        self.global_scope: Optional[dict[str, Any]] = None
-        self.guard_fail_fn: Optional[Callable[[GuardFail], None]] = None
-        self.cache_entry: Optional[CacheEntry] = None
-        self.extra_state: Optional[ExtraState] = None
+        self.verbose_code_parts: list[str] | None = None
+        self.global_scope: dict[str, Any] | None = None
+        self.guard_fail_fn: Callable[[GuardFail], None] | None = None
+        self.cache_entry: CacheEntry | None = None
+        self.extra_state: ExtraState | None = None
         self.id_matched_objs: dict[str, ReferenceType[object]] = {}
         self.no_tensor_aliasing_sources: list[str] = []
 
@@ -607,7 +608,7 @@ def get_guard_lines(self, guard: LeafGuard) -> list[str]:
         return parts
 
     def get_manager_line(
-        self, guard_manager: GuardManager, accessor_str: Optional[str] = None
+        self, guard_manager: GuardManager, accessor_str: str | None = None
     ) -> str:
         source = guard_manager.get_source()
         t = guard_manager.__class__.__name__
@@ -737,7 +738,7 @@ def uninteresting_files() -> set[str]:
     return {inspect.getfile(m) for m in mods}
 
 
-_CLOSURE_VARS: Optional[dict[str, object]] = None
+_CLOSURE_VARS: dict[str, object] | None = None
 
 
 def _get_closure_vars() -> dict[str, object]:
@@ -751,6 +752,7 @@ def _get_closure_vars() -> dict[str, object]:
             "___dict_version": dict_version,
             "___dict_contains": lambda a, b: dict.__contains__(b, a),
             "___tuple_iterator_len": tuple_iterator_len,
+            "___normalize_count_iter": normalize_count_iter,
             "___normalize_range_iter": normalize_range_iter,
             "___tuple_iterator_getitem": tuple_iterator_getitem,
             "___dataclass_fields": dataclass_fields,
@@ -778,7 +780,7 @@ def _ast_unparse(node: ast.AST) -> str:
 strip_function_call = torch._C._dynamo.strip_function_call
 
 
-def get_verbose_code_part(code_part: str, guard: Optional[Guard]) -> str:
+def get_verbose_code_part(code_part: str, guard: Guard | None) -> str:
     extra = ""
     if guard is not None:
         if guard.user_stack:
@@ -803,9 +805,9 @@ def get_verbose_code_part(code_part: str, guard: Optional[Guard]) -> str:
 
 
 def get_verbose_code_parts(
-    code_parts: Union[str, list[str]],
-    guard: Optional[Guard],
-    recompile_hint: Optional[str] = None,
+    code_parts: str | list[str],
+    guard: Guard | None,
+    recompile_hint: str | None = None,
 ) -> list[str]:
     if not isinstance(code_parts, list):
         code_parts = [code_parts]
@@ -832,7 +834,7 @@ def get_verbose_code_parts(
     return verbose_code_parts
 
 
-def _get_closure_var_hint(source: Optional[Source]) -> Optional[str]:
+def _get_closure_var_hint(source: Source | None) -> str | None:
     """
     Walk up the source chain to find a CellContentsSource ancestor.
     Returns a hint like 'guard on "varname".attr' or None if not found.
@@ -841,7 +843,7 @@ def _get_closure_var_hint(source: Optional[Source]) -> Optional[str]:
         return None
 
     full_name = source.name
-    current: Optional[Source] = source
+    current: Source | None = source
     while current is not None:
         if isinstance(current, CellContentsSource) and current.freevar_name:
             # Compute the path suffix by comparing names
@@ -854,7 +856,7 @@ def _get_closure_var_hint(source: Optional[Source]) -> Optional[str]:
     return None
 
 
-def convert_int_to_concrete_values(dim: Any) -> Optional[int]:
+def convert_int_to_concrete_values(dim: Any) -> int | None:
     if dim is None:
         return None
     if not is_symbolic(dim):
@@ -864,15 +866,15 @@ def convert_int_to_concrete_values(dim: Any) -> Optional[int]:
         return dim.node.maybe_as_int()
 
 
-def convert_to_concrete_values(size_or_stride: list[Any]) -> list[Optional[int]]:
+def convert_to_concrete_values(size_or_stride: list[Any]) -> list[int | None]:
     return [convert_int_to_concrete_values(dim) for dim in size_or_stride]
 
 
 def get_tensor_guard_code_part(
     value: torch.Tensor,
     name: str,
-    sizes: list[Optional[int]],
-    strides: list[Optional[int]],
+    sizes: list[int | None],
+    strides: list[int | None],
     pytype: type,
     dispatch_keys: DispatchKeySet,
 ) -> str:
@@ -910,13 +912,7 @@ def raise_local_type_error(obj: Any) -> NoReturn:
 
 
 def should_optimize_getattr_on_nn_module(value: Any) -> bool:
-    # If inline_inbuilt_nn_modules flag is True, Dynamo has already traced
-    # through the __getattr__, and therefore it is always safe to optimize
-    # getattr on nn modules.
-    return isinstance(value, torch.nn.Module) and (
-        config.inline_inbuilt_nn_modules
-        or get_custom_getattr(value) is unpatched_nn_module_getattr
-    )
+    return isinstance(value, torch.nn.Module)
 
 
 @dataclasses.dataclass(frozen=True)
@@ -928,14 +924,14 @@ class NNModuleAttrAccessorInfo:
     present_in_generic_dict: bool = False
 
     # Either the actual name or _parameters/_buffers/_modules
-    l1_key: Optional[str] = None
+    l1_key: str | None = None
 
     # Actual parameter/buffer/submodule name
-    l2_key: Optional[str] = None
+    l2_key: str | None = None
 
 
 def getitem_on_dict_manager(
-    source: Union[DictGetItemSource, DictSubclassGetItemSource],
+    source: DictGetItemSource | DictSubclassGetItemSource,
     base_guard_manager: DictGuardManager,
     base_example_value: Any,
     example_value: Any,
@@ -1010,19 +1006,166 @@ def code_framelocals_names_reversed_cached(code: types.CodeType) -> list[str]:
     return list(reversed(code_framelocals_names(code)))
 
 
+class GuardCheckSpec(NamedTuple):
+    """Lightweight guard evaluation for subgraph reuse during tracing.
+
+    WHY IS THIS NEEDED?
+    Normal guards are collected during tracing, installed via GuardBuilder, and
+    evaluated in C++ at runtime. GuardCheckSpec serves a different purpose: we
+    want to evaluate a guard *during tracing* to check whether a previously
+    traced partial subgraph can be reused. The guard is collected during the
+    first trace of subgraph and then evaluated with a new source value before we
+    begin tracing the second subgraph.
+
+    Existing C++ guards don't satisfy this need because they're designed for
+    runtime evaluation, not mid-trace comparisons with different source values.
+
+    NB - This separate spec table *increases maintenance burden*, but each
+    handler is small and the number of guard types is finite, so it's tractable.
+
+    get_metadata_fn: (guard, value) -> metadata  — captures the minimum
+        information from the guarded value at build time.
+    eval_fn: (value, metadata) -> bool  — re-evaluates the guard at lookup
+        time using a fresh value and the previously saved metadata.
+    """
+
+    get_metadata_fn: Any
+    eval_fn: Any
+
+
+SKIP_GUARD = object()
+
+
+# Used by TENSOR_MATCH guard check spec
+def extract_tensor_metadata(t: torch.Tensor) -> tuple[Any, ...]:
+    return (t.shape, t.stride(), t.dtype, t.device, t.requires_grad)
+
+
+# Used by TENSOR_SUBCLASS_METADATA_MATCH guard check spec
+def extract_subclass_metadata(guard: Any, value: Any) -> tuple[Any, ...]:
+    metadata = deepcopy(value.__tensor_flatten__()[1])
+    cls = type(value)
+    has_custom_guard = hasattr(value, "__metadata_guard__")
+    return (metadata, cls, has_custom_guard)
+
+
+# Used by TENSOR_SUBCLASS_METADATA_MATCH guard check spec
+def check_subclass_metadata(value: Any, metadata: tuple[Any, ...]) -> bool:
+    saved_metadata, cls, has_custom_guard = metadata
+    if has_custom_guard:
+        return cls.__metadata_guard__(saved_metadata, value.__tensor_flatten__()[1])
+    return value.__tensor_flatten__()[1] == saved_metadata
+
+
+# Used by DTENSOR_SPEC_MATCH guard check spec
+def extract_dtensor_spec(guard: Any, value: Any) -> Any:
+    return deepcopy(value)
+
+
+# Used by DTENSOR_SPEC_MATCH guard check spec
+def check_dtensor_spec(value: Any, metadata: Any) -> bool:
+    return value._check_equals(metadata, skip_shapes=True)
+
+
+# Used by OPAQUE_OBJ_GUARD_FN_MATCH guard check spec
+def extract_opaque_obj(guard: Any, value: Any) -> Any:
+    opaque_info = get_opaque_obj_info(type(value))
+    if not opaque_info or not opaque_info.guard_fn:
+        return None
+    return deepcopy(opaque_info.guard_fn(value))
+
+
+# Used by OPAQUE_OBJ_GUARD_FN_MATCH guard check spec
+def check_opaque_obj(value: Any, metadata: Any) -> bool:
+    if metadata is None:
+        return True
+    opaque_info = get_opaque_obj_info(type(value))
+    if not opaque_info or not opaque_info.guard_fn:
+        return metadata is None
+    return opaque_info.guard_fn(value) == metadata
+
+
+# Used by CLOSURE_MATCH guard check spec
+def extract_closure(guard: Any, value: Any) -> Any:
+    if type(value) is types.FunctionType and hasattr(value, "__code__"):
+        return value.__code__
+    return id(value)
+
+
+# Used by CLOSURE_MATCH guard check spec
+def check_closure(value: Any, metadata: Any) -> bool:
+    if type(value) is types.FunctionType and hasattr(value, "__code__"):
+        return value.__code__ is metadata
+    return id(value) == metadata
+
+
+def _constant_subclass_base_value(value: Any) -> Any:
+    """Extract the base constant value from a constant subclass instance."""
+    from .variables.user_defined import _CONSTANT_BASE_TYPES
+
+    for t in _CONSTANT_BASE_TYPES:
+        if isinstance(value, t):
+            return t(value)  # pyrefly: ignore[bad-argument-type]
+    raise TypeError(f"Not a constant subclass: {type(value)}")
+
+
+def register_guard_check_spec(
+    get_metadata_fn,
+    eval_fn,
+):
+    """Attach a GuardCheckSpec to a guard method for auto-dispatch."""
+    handler = GuardCheckSpec(get_metadata_fn=get_metadata_fn, eval_fn=eval_fn)
+
+    def decorator(fn):
+        fn.guard_check_spec = handler
+        return fn
+
+    return decorator
+
+
+def skip_guard_check_spec(fn):
+    """Mark a guard method as skipped during auto-cache dispatch."""
+    fn.guard_check_spec = SKIP_GUARD
+    return fn
+
+
+class UnsupportedGuardCheckSpec:
+    """Sentinel for guards with no check spec yet.
+
+    Raises NotImplementedError if get_metadata_fn/eval_fn are accidentally called.
+    """
+
+    def __init__(self, name: str) -> None:
+        self._name = name
+
+    @property
+    def get_metadata_fn(self):
+        raise NotImplementedError(f"Guard check spec not implemented for {self._name}")
+
+    @property
+    def eval_fn(self):
+        raise NotImplementedError(f"Guard check spec not implemented for {self._name}")
+
+
+def unsupported_guard_check_spec(fn):
+    """Mark a guard method as unsupported for auto-cache dispatch."""
+    fn.guard_check_spec = UnsupportedGuardCheckSpec(fn.__name__)
+    return fn
+
+
 class GuardBuilder(GuardBuilderBase):
     def __init__(
         self,
         f_code: types.CodeType,
         id_ref: Callable[[object, str], int],
         source_ref: Callable[[Source], str],
-        lookup_weakrefs: Callable[[object], Optional[weakref.ref[object]]],
+        lookup_weakrefs: Callable[[object], weakref.ref[object] | None],
         local_scope: dict[str, object],
         global_scope: dict[str, object],
         guard_manager: GuardManagerWrapper,
         check_fn_manager: CheckFunctionManager,
         save_guards: bool = False,
-        runtime_global_scope: Optional[dict[str, object]] = None,
+        runtime_global_scope: dict[str, object] | None = None,
         guard_filter_fn: Callable[[Sequence[GuardFilterEntry]], Sequence[bool]]
         | None = None,
     ) -> None:
@@ -1031,9 +1174,7 @@ def __init__(
         self.source_ref = source_ref
         self.lookup_weakrefs = lookup_weakrefs
         self.scope: dict[str, dict[str, object]] = {"L": local_scope, "G": global_scope}
-        self.src_get_value_cache: weakref.WeakKeyDictionary[Source, object] = (
-            weakref.WeakKeyDictionary()
-        )
+        self.src_get_value_cache: dict[Source, object] = {}
         self.runtime_global_scope = runtime_global_scope or global_scope
         self.scope["__builtins__"] = builtins.__dict__.copy()
         for (
@@ -1163,7 +1304,7 @@ def guard_on_dict_keys_and_order(self, value: dict[Any, Any], guard: Guard) -> N
                 )
 
     @staticmethod
-    def _get_generic_dict_manager_example_value(example_value: Any) -> Optional[Any]:
+    def _get_generic_dict_manager_example_value(example_value: Any) -> Any | None:
         # due to a bug in 3.13.0 (introduced by https://github.com/python/cpython/pull/116115,
         # reported in https://github.com/python/cpython/issues/125608,
         # fixed by https://github.com/python/cpython/pull/125611), we cannot take
@@ -1345,9 +1486,11 @@ def requires_key_order_guarding(self, source: Source) -> bool:
     def get_guard_manager_type(
         self,
         source: Source,
-        example_value: Optional[
-            Union[KeysView[Any], set[Any], frozenset[Any], dict[Any, Any]]
-        ],
+        example_value: KeysView[Any]
+        | set[Any]
+        | frozenset[Any]
+        | dict[Any, Any]
+        | None,
     ) -> GuardManagerType:
         guard_manager_enum = GuardManagerType.GUARD_MANAGER
         if self.requires_key_order_guarding(source):
@@ -1796,7 +1939,7 @@ def add_python_lambda_leaf_guard_to_root(
         self,
         code_parts: list[str],
         verbose_code_parts: list[str],
-        closure_vars: Optional[dict[str, object]] = None,
+        closure_vars: dict[str, object] | None = None,
         is_epilogue: bool = True,
     ) -> None:
         if closure_vars is None:
@@ -1832,7 +1975,7 @@ def add_python_lambda_leaf_guard_to_root(
     def get(
         self,
         guard_or_source: Guard | Source,
-        closure_vars: Optional[dict[str, Any]] = None,
+        closure_vars: dict[str, Any] | None = None,
     ) -> Any:
         if isinstance(guard_or_source, Source):
             src = guard_or_source
@@ -1848,7 +1991,7 @@ def get(
     # to call this before generating some code that makes use of 'guard',
     # because without this call, we won't actually bind the variable
     # you reference in the actual guard closure (oops!)
-    def arg_ref(self, guard: Union[str, Guard]) -> str:
+    def arg_ref(self, guard: str | Guard) -> str:
         name: str
         if isinstance(guard, str):
             name = guard
@@ -1881,20 +2024,22 @@ def _guard_on_attribute(
         new_guard.create(self)
 
     # Note: the order of the guards in this file matters since we sort guards on the same object by lineno
-    def HASATTR(self, guard: Guard) -> None:
+
+    @register_guard_check_spec(
+        get_metadata_fn=lambda guard, value: (
+            guard.create_fn.keywords["attr"],
+            hasattr(value, guard.create_fn.keywords["attr"]),
+        ),
+        eval_fn=lambda value, metadata: hasattr(value, metadata[0]) == metadata[1],
+    )
+    def HASATTR(self, guard: Guard, attr: str | None = None) -> None:
+        assert attr is not None
         source = guard.originating_source
         if isinstance(source, NNModuleSource):
             source = source.base
-        if isinstance(source, CodeSource):
-            # No need to guard that a function has a __code__ attribute
-            return
-        assert isinstance(source, AttrSource), f"invalid source {guard.name}"
-        base_source = source.base
-        base = base_source.name
-        attr = source.member
 
-        ref = self.arg_ref(base)
-        val = hasattr(self.get(base_source), attr)
+        ref = self.arg_ref(guard)
+        val = hasattr(self.get(guard), attr)
         code = None
         if val:
             code = f"hasattr({ref}, {attr!r})"
@@ -1905,27 +2050,28 @@ def HASATTR(self, guard: Guard) -> None:
             return
 
         self._set_guard_export_info(
-            guard, [code], provided_guarded_object=self.get(base_source)
+            guard, [code], provided_guarded_object=self.get(guard)
         )
 
-        base_manager = self.get_guard_manager_from_source(base_source)
+        base_manager = self.get_guard_manager_from_source(source)
         if val:
             # Just install a getattr manager. GetAttrGuardAccessor itself
             # acts as hasattr guard.
-            example_value = self.get(source)
-            base_example_value = self.get(base_source)
-            guard_manager_enum = self.get_guard_manager_type(source, example_value)
+            attr_source = AttrSource(source, attr)
+            example_value = self.get(attr_source)
+            base_example_value = self.get(guard)
+            guard_manager_enum = self.get_guard_manager_type(attr_source, example_value)
 
             # if the base value is nn.Module, check if we can speedup the
             # guard by going through __dict__ attrs.
             if should_optimize_getattr_on_nn_module(base_example_value):
                 self.getattr_on_nn_module(
-                    source,
+                    attr_source,
                     base_manager,
                     base_example_value,
                     example_value,
-                    base,
                     source.name,
+                    attr_source.name,
                     guard_manager_enum,
                 )
             else:
@@ -1941,8 +2087,12 @@ def HASATTR(self, guard: Guard) -> None:
             )
         self.already_added_code_parts.add(code)
 
+    @register_guard_check_spec(
+        get_metadata_fn=lambda guard, value: guard.create_fn.keywords["attr"],
+        eval_fn=lambda value, metadata: metadata not in value.__dict__,
+    )
     def NOT_PRESENT_IN_GENERIC_DICT(
-        self, guard: Guard, attr: Optional[Any] = None
+        self, guard: Guard, attr: Any | None = None
     ) -> None:
         assert attr is not None
         ref = self.arg_ref(guard)
@@ -1969,6 +2119,10 @@ def NOT_PRESENT_IN_GENERIC_DICT(
         )
         self.already_added_code_parts.add(code)
 
+    @register_guard_check_spec(
+        get_metadata_fn=lambda guard, value: type(value),
+        eval_fn=lambda value, metadata: type(value) is metadata,
+    )
     def TYPE_MATCH(self, guard: Guard) -> None:
         # ___check_type_id is same as `id(type(x)) == y`
         value = self.get(guard)
@@ -1995,6 +2149,9 @@ def TYPE_MATCH(self, guard: Guard) -> None:
             guard.user_stack,
         )
 
+    # Dict versions change on any mutation, so a version captured during one
+    # trace is meaningless for a later subgraph reuse check.
+    @unsupported_guard_check_spec
     def DICT_VERSION(self, guard: Guard) -> None:
         # ___check_dict_version is same as `dict_version(x) == y`
         ref = self.arg_ref(guard)
@@ -2009,27 +2166,53 @@ def DICT_VERSION(self, guard: Guard) -> None:
             val, get_verbose_code_parts(code, guard), guard.user_stack
         )
 
-    def DICT_CONTAINS(self, guard: Guard, key: str, invert: bool) -> None:
+    @register_guard_check_spec(
+        get_metadata_fn=lambda guard, value: guard.create_fn.keywords["key"],
+        eval_fn=lambda value, metadata: metadata in value,
+    )
+    def DICT_CONTAINS(self, guard: Guard, key: str) -> None:
+        dict_ref = self.arg_ref(guard)
+
+        code = f"___dict_contains({key!r}, {dict_ref})"
+        if code in self.already_added_code_parts:
+            return
+        self._set_guard_export_info(guard, [code])
+
+        self.get_guard_manager(guard).add_dict_contains_guard(
+            True,
+            key,
+            get_verbose_code_parts(code, guard),
+            guard.user_stack,
+        )
+        self.already_added_code_parts.add(code)
+
+    @register_guard_check_spec(
+        get_metadata_fn=lambda guard, value: guard.create_fn.keywords["key"],
+        eval_fn=lambda value, metadata: metadata not in value,
+    )
+    def DICT_NOT_CONTAINS(self, guard: Guard, key: str) -> None:
         dict_ref = self.arg_ref(guard)
 
-        maybe_not = "not " if invert else ""
-        code = f"{maybe_not}___dict_contains({key!r}, {dict_ref})"
+        code = f"not ___dict_contains({key!r}, {dict_ref})"
         if code in self.already_added_code_parts:
             return
         self._set_guard_export_info(guard, [code])
 
         self.get_guard_manager(guard).add_dict_contains_guard(
-            not invert,
+            False,
             key,
             get_verbose_code_parts(code, guard),
             guard.user_stack,
         )
         self.already_added_code_parts.add(code)
 
-    def SET_CONTAINS(self, guard: Guard, key: Any, invert: bool) -> None:
+    @register_guard_check_spec(
+        get_metadata_fn=lambda guard, value: guard.create_fn.keywords["key"],
+        eval_fn=lambda value, metadata: metadata in value,
+    )
+    def SET_CONTAINS(self, guard: Guard, key: Any) -> None:
         set_ref = self.arg_ref(guard)
         item = key
-        contains = not invert  # install_dict_contains_guard inverts "contains"
 
         code = f"set.__contains__({set_ref}, {item!r})"
         if code in self.already_added_code_parts:
@@ -2038,13 +2221,39 @@ def SET_CONTAINS(self, guard: Guard, key: Any, invert: bool) -> None:
         self._set_guard_export_info(guard, [code])
 
         self.get_guard_manager(guard).add_set_contains_guard(
-            contains,
+            True,
             item,
             get_verbose_code_parts(code, guard),
             guard.user_stack,
         )
         self.already_added_code_parts.add(code)
 
+    @register_guard_check_spec(
+        get_metadata_fn=lambda guard, value: guard.create_fn.keywords["key"],
+        eval_fn=lambda value, metadata: metadata not in value,
+    )
+    def SET_NOT_CONTAINS(self, guard: Guard, key: Any) -> None:
+        set_ref = self.arg_ref(guard)
+        item = key
+
+        code = f"not set.__contains__({set_ref}, {item!r})"
+        if code in self.already_added_code_parts:
+            return
+
+        self._set_guard_export_info(guard, [code])
+
+        self.get_guard_manager(guard).add_set_contains_guard(
+            False,
+            item,
+            get_verbose_code_parts(code, guard),
+            guard.user_stack,
+        )
+        self.already_added_code_parts.add(code)
+
+    @register_guard_check_spec(
+        get_metadata_fn=lambda guard, value: value,
+        eval_fn=lambda value, metadata: value == metadata,
+    )
     def BOOL_MATCH(self, guard: Guard) -> None:
         # checks val == True or val == False
         ref = self.arg_ref(guard)
@@ -2062,6 +2271,10 @@ def BOOL_MATCH(self, guard: Guard) -> None:
                 get_verbose_code_parts(code, guard), guard.user_stack
             )
 
+    @register_guard_check_spec(
+        get_metadata_fn=lambda guard, value: None,
+        eval_fn=lambda value, metadata: value is None,
+    )
     def NONE_MATCH(self, guard: Guard) -> None:
         # checks `val is None`
         ref = self.arg_ref(guard)
@@ -2074,7 +2287,11 @@ def NONE_MATCH(self, guard: Guard) -> None:
             get_verbose_code_parts(code, guard), guard.user_stack
         )
 
-    def ID_MATCH(self, guard: Guard, recompile_hint: Optional[str] = None) -> None:
+    @register_guard_check_spec(
+        get_metadata_fn=lambda guard, value: value,
+        eval_fn=lambda value, metadata: value is metadata,
+    )
+    def ID_MATCH(self, guard: Guard, recompile_hint: str | None = None) -> None:
         # TODO - Run a CI with the following uncommented to find the remaining places
         # val = self.get(guard)
         # if inspect.isclass(val):
@@ -2084,7 +2301,7 @@ def ID_MATCH(self, guard: Guard, recompile_hint: Optional[str] = None) -> None:
         return self.id_match_unchecked(guard, recompile_hint)
 
     def id_match_unchecked(
-        self, guard: Guard, recompile_hint: Optional[str] = None
+        self, guard: Guard, recompile_hint: str | None = None
     ) -> None:
         # ___check_obj_id is same as `id(x) == y`
         if isinstance(guard.originating_source, TypeSource):
@@ -2122,7 +2339,11 @@ def id_match_unchecked(
                 if weak_id is not None:
                     self.id_matched_objs[local_name] = weak_id
 
-    def NOT_NONE_MATCH(self, guard: Guard, value: Optional[Any] = None) -> None:
+    @register_guard_check_spec(
+        get_metadata_fn=lambda guard, value: None,
+        eval_fn=lambda value, metadata: value is not None,
+    )
+    def NOT_NONE_MATCH(self, guard: Guard, value: Any | None = None) -> None:
         ref = self.arg_ref(guard)
         val = self.get(guard)
         assert isinstance(val, torch.Tensor)
@@ -2133,6 +2354,10 @@ def NOT_NONE_MATCH(self, guard: Guard, value: Optional[Any] = None) -> None:
             get_verbose_code_parts(code, guard), guard.user_stack
         )
 
+    @register_guard_check_spec(
+        get_metadata_fn=lambda guard, value: value.raw_repr(),
+        eval_fn=lambda value, metadata: value.raw_repr() == metadata,
+    )
     def DISPATCH_KEY_SET_MATCH(self, guard: Guard) -> None:
         ref = self.arg_ref(guard)
         val = self.get(guard)
@@ -2145,6 +2370,8 @@ def DISPATCH_KEY_SET_MATCH(self, guard: Guard) -> None:
             guard.user_stack,
         )
 
+    # Global state guard — not source-specific, checked separately at runtime.
+    @skip_guard_check_spec
     def DUAL_LEVEL(self, guard: Guard) -> None:
         # Invalidate dual level if current dual level is different than the one
         # in the fx graph
@@ -2158,6 +2385,8 @@ def DUAL_LEVEL(self, guard: Guard) -> None:
             guard.user_stack,
         )
 
+    # Global state guard — not source-specific, checked separately at runtime.
+    @skip_guard_check_spec
     def FUNCTORCH_STACK_MATCH(self, guard: Guard) -> None:
         # Invalidate functorch code if current level is different than
         # the one when FX graph was generated
@@ -2177,6 +2406,8 @@ def fn(x: Any) -> bool:
             fn, get_verbose_code_parts(code, guard), guard.user_stack
         )
 
+    # Global state guard — not source-specific, checked separately at runtime.
+    @skip_guard_check_spec
     def AUTOGRAD_SAVED_TENSORS_HOOKS(self, guard: Guard) -> None:
         get_hooks = torch._functorch._aot_autograd.utils.top_saved_tensors_hooks
         are_inline_hooks = (
@@ -2185,7 +2416,7 @@ def AUTOGRAD_SAVED_TENSORS_HOOKS(self, guard: Guard) -> None:
 
         def hooks_ids_fn(
             hooks: tuple[Callable[[torch.Tensor], Any], Callable[[Any], torch.Tensor]],
-        ) -> Optional[tuple[int, ...]]:
+        ) -> tuple[int, ...] | None:
             if not are_inline_hooks(hooks):
                 return None
 
@@ -2205,9 +2436,26 @@ def fn(x: Any) -> bool:
             fn, get_verbose_code_parts(code, guard), guard.user_stack
         )
 
+    @register_guard_check_spec(
+        get_metadata_fn=extract_subclass_metadata,
+        eval_fn=check_subclass_metadata,
+    )
     def TENSOR_SUBCLASS_METADATA_MATCH(self, guard: Guard) -> None:
         value = self.get(guard)
-        original_metadata = deepcopy(self.get(guard).__tensor_flatten__()[1])
+
+        # Deepcopying SymInts result in an error from copying FakeTensors.
+        # Instead we just always assume the metadata is the same.
+        class _AnyCompare:
+            def __eq__(self, other: object) -> bool:
+                return True
+
+            def __ne__(self, other: object) -> bool:
+                return False
+
+        metadata = value.__tensor_flatten__()[1]
+        original_metadata = deepcopy(
+            pytree.tree_map_only(torch.SymInt, lambda _: _AnyCompare(), metadata)
+        )
         if hasattr(value, "__metadata_guard__"):
             verify_guard_fn_signature(value)
             cls = type(value)
@@ -2229,35 +2477,38 @@ def metadata_checker(x: Any) -> bool:
             guard.user_stack,
         )
 
+    @register_guard_check_spec(
+        get_metadata_fn=extract_dtensor_spec,
+        eval_fn=check_dtensor_spec,
+    )
     def DTENSOR_SPEC_MATCH(self, guard: Guard) -> None:
         # Copied from DTensor __metadata_guard__
         # TODO - Consider moving this to C++ if stable
-        value = deepcopy(self.get(guard))
+        expected = extract_dtensor_spec(guard, self.get(guard))
 
         def guard_fn(x: Any) -> bool:
-            return x._check_equals(value, skip_shapes=True)
+            return check_dtensor_spec(x, expected)
 
         code = f"__dtensor_spec_{id(guard_fn)}"
         self.get_guard_manager(guard).add_lambda_guard(
             guard_fn, get_verbose_code_parts(code, guard), guard.user_stack
         )
 
+    @register_guard_check_spec(
+        get_metadata_fn=extract_opaque_obj,
+        eval_fn=check_opaque_obj,
+    )
     def OPAQUE_OBJ_GUARD_FN_MATCH(self, guard: Guard) -> None:
         """Guard on the values returned by the opaque object's guard_fn."""
 
         value = self.get(guard)
-        opaque_info = get_opaque_obj_info(type(value))
+        original_values = extract_opaque_obj(guard, value)
 
-        if not opaque_info or not opaque_info.guard_fn:
+        if original_values is None:
             return
 
-        original_values = deepcopy(opaque_info.guard_fn(value))
-
         def opaque_guard_checker(x: Any) -> bool:
-            current_values = opaque_info.guard_fn(  # pyrefly: ignore[missing-attribute]
-                x
-            )
-            return current_values == original_values
+            return check_opaque_obj(x, original_values)
 
         global_name = f"___check_opaque_guard_fn_{id(opaque_guard_checker)}_c{CompileContext.current_compile_id()}"
         self.get_guard_manager(guard).add_lambda_guard(
@@ -2266,7 +2517,11 @@ def opaque_guard_checker(x: Any) -> bool:
             guard.user_stack,
         )
 
-    def EQUALS_MATCH(self, guard: Guard, recompile_hint: Optional[str] = None) -> None:
+    @register_guard_check_spec(
+        get_metadata_fn=lambda guard, value: value,
+        eval_fn=lambda value, metadata: value == metadata,
+    )
+    def EQUALS_MATCH(self, guard: Guard, recompile_hint: str | None = None) -> None:
         ref = self.arg_ref(guard)
         val = self.get(guard)
         if np:
@@ -2306,7 +2561,7 @@ def EQUALS_MATCH(self, guard: Guard, recompile_hint: Optional[str] = None) -> No
         )
 
         if torch.distributed.is_available():
-            from torch.distributed.device_mesh import DeviceMesh
+            from torch.distributed.device_mesh import _MeshLayout, DeviceMesh
             from torch.distributed.tensor.placement_types import (
                 _StridedShard,
                 Partial,
@@ -2320,6 +2575,7 @@ def EQUALS_MATCH(self, guard: Guard, recompile_hint: Optional[str] = None) -> No
                 Partial,
                 DeviceMesh,
                 _StridedShard,
+                _MeshLayout,
             )
 
         from torch.export.dynamic_shapes import _IntWrapper
@@ -2377,6 +2633,10 @@ def EQUALS_MATCH(self, guard: Guard, recompile_hint: Optional[str] = None) -> No
         self._set_guard_export_info(guard, code)
         return
 
+    @register_guard_check_spec(
+        get_metadata_fn=lambda guard, value: value,
+        eval_fn=lambda value, metadata: value == metadata,
+    )
     def CONSTANT_MATCH(self, guard: Guard) -> None:
         val = self.get(guard)
         if istype(val, bool):
@@ -2388,6 +2648,46 @@ def CONSTANT_MATCH(self, guard: Guard) -> None:
         else:
             self.EQUALS_MATCH(guard)
 
+    @register_guard_check_spec(
+        get_metadata_fn=lambda guard, value: _constant_subclass_base_value(value),
+        eval_fn=lambda value, metadata: _constant_subclass_base_value(value)
+        == metadata,
+    )
+    def CONSTANT_SUBCLASS_MATCH(self, guard: Guard) -> None:
+        """Guard for subclasses of constant types (int, float, str, etc.).
+
+        Extracts the base value using the base type's converter (e.g.,
+        int.__int__) to avoid calling user-overridden __eq__.
+        """
+        from .variables.user_defined import _CONSTANT_BASE_TYPES
+
+        val = self.get(guard)
+        ref = self.arg_ref(guard)
+
+        # Find the constant base type
+        base_type = None
+        for t in _CONSTANT_BASE_TYPES:
+            if isinstance(val, t):
+                base_type = t
+                break
+        assert base_type is not None
+
+        base_value = base_type(val)
+        code = [f"{base_type.__name__}({ref}) == {base_value!r}"]
+
+        def check_fn(x: Any) -> bool:
+            return base_type(x) == base_value
+
+        self.get_guard_manager(guard).add_lambda_guard(
+            check_fn,
+            get_verbose_code_parts(code, guard),
+            guard.user_stack,
+        )
+
+    @register_guard_check_spec(
+        get_metadata_fn=lambda guard, value: value,
+        eval_fn=lambda value, metadata: value is metadata,
+    )
     def NN_MODULE(self, guard: Guard) -> None:
         # don't support this in serialization because it uses unsupported ID_MATCH
         self.ID_MATCH(guard, "[inline-inbuilt-nn-modules-candidate]")
@@ -2408,11 +2708,19 @@ def NN_MODULE(self, guard: Guard) -> None:
                 ],
             )
 
+    @register_guard_check_spec(
+        get_metadata_fn=lambda guard, value: value,
+        eval_fn=lambda value, metadata: value is metadata,
+    )
     def FUNCTION_MATCH(self, guard: Guard) -> None:
         """things like torch.add and user defined functions"""
         # don't support this in serialization because it uses unsupported ID_MATCH
         return self.ID_MATCH(guard)
 
+    @register_guard_check_spec(
+        get_metadata_fn=lambda guard, value: value,
+        eval_fn=lambda value, metadata: value is metadata,
+    )
     def CLASS_MATCH(self, guard: Guard) -> None:
         """Equals ID_MATCH on classes - better readability than directly calling ID_MATCH"""
         val = self.get(guard)
@@ -2422,6 +2730,10 @@ def CLASS_MATCH(self, guard: Guard) -> None:
             )
         self.id_match_unchecked(guard)
 
+    @register_guard_check_spec(
+        get_metadata_fn=lambda guard, value: value,
+        eval_fn=lambda value, metadata: value is metadata,
+    )
     def MODULE_MATCH(self, guard: Guard) -> None:
         """Equals ID_MATCH on modules - better readability than directly calling ID_MATCH"""
         val = self.get(guard)
@@ -2431,17 +2743,26 @@ def MODULE_MATCH(self, guard: Guard) -> None:
             )
         self.id_match_unchecked(guard)
 
+    @register_guard_check_spec(
+        get_metadata_fn=extract_closure,
+        eval_fn=check_closure,
+    )
     def CLOSURE_MATCH(self, guard: Guard) -> None:
         """matches a closure by __code__ id."""
         # don't support this in serialization because it uses unsupported FUNCTION_MATCH
         val = self.get(guard)
         # Strictly only want user-defined functions
         if type(val) is types.FunctionType and hasattr(val, "__code__"):
-            self._guard_on_attribute(guard, "__code__", GuardBuilder.HASATTR)  # type: ignore[arg-type]
+            # No explicit HASATTR guard needed for __code__ — the getattr
+            # accessor installed by CONSTANT_MATCH implicitly guards hasattr.
             self._guard_on_attribute(guard, "__code__", GuardBuilder.CONSTANT_MATCH)  # type: ignore[arg-type]
         else:
             self.FUNCTION_MATCH(guard)
 
+    @register_guard_check_spec(
+        get_metadata_fn=lambda guard, value: value,
+        eval_fn=lambda value, metadata: value is metadata,
+    )
     def BUILTIN_MATCH(self, guard: Guard) -> None:
         if self.save_guards:
             # Record which builtin variables are used for pruning later.
@@ -2451,6 +2772,10 @@ def BUILTIN_MATCH(self, guard: Guard) -> None:
                 )
         return self.id_match_unchecked(guard)
 
+    @register_guard_check_spec(
+        get_metadata_fn=lambda guard, value: len(value),
+        eval_fn=lambda value, metadata: len(value) == metadata,
+    )
     def SEQUENCE_LENGTH(self, guard: Guard) -> None:
         # This guard is used to check length of PySequence objects like list,
         # tuple, collections.deque etc
@@ -2481,6 +2806,10 @@ def SEQUENCE_LENGTH(self, guard: Guard) -> None:
                 guard.user_stack,
             )
 
+    @register_guard_check_spec(
+        get_metadata_fn=lambda guard, value: tuple_iterator_len(value),
+        eval_fn=lambda value, metadata: tuple_iterator_len(value) == metadata,
+    )
     def TUPLE_ITERATOR_LEN(self, guard: Guard) -> None:
         ref = self.arg_ref(guard)
         value = self.get(guard)
@@ -2500,6 +2829,10 @@ def TUPLE_ITERATOR_LEN(self, guard: Guard) -> None:
             guard.user_stack,
         )
 
+    @register_guard_check_spec(
+        get_metadata_fn=lambda guard, value: normalize_range_iter(value),
+        eval_fn=lambda value, metadata: normalize_range_iter(value) == metadata,
+    )
     def RANGE_ITERATOR_MATCH(self, guard: Guard) -> None:
         ref = self.arg_ref(guard)
         value = self.get(guard)
@@ -2523,7 +2856,34 @@ def RANGE_ITERATOR_MATCH(self, guard: Guard) -> None:
             guard.user_stack,
         )
 
+    @register_guard_check_spec(
+        get_metadata_fn=lambda guard, value: (type(value), normalize_count_iter(value)),
+        eval_fn=lambda value, metadata: (
+            type(value) is metadata[0] and normalize_count_iter(value) == metadata[1]
+        ),
+    )
+    def COUNT_ITERATOR_MATCH(self, guard: Guard) -> None:
+        ref = self.arg_ref(guard)
+        value = self.get(guard)
+        count_type = type(value)
+        normalized_count_iter = normalize_count_iter(value)
+
+        def guard_fn(x: Any) -> bool:
+            return (
+                type(x) is count_type
+                and normalize_count_iter(x) == normalized_count_iter
+            )
+
+        code = [f"___normalize_count_iter({ref}) == {normalized_count_iter}"]
+        self._set_guard_export_info(guard, code)
+        self.get_guard_manager(guard).add_lambda_guard(
+            guard_fn, get_verbose_code_parts(code, guard), guard.user_stack
+        )
+
+    # Multi-source guard (two inputs aliasing) — not expressible as a
+    # single source → value check.
     # TODO(voz): Deduplicate w/ AOTAutograd dupe input guards
+    @skip_guard_check_spec
     def DUPLICATE_INPUT(self, guard: Guard, source_b: Source) -> None:
         if is_from_skip_guard_source(
             guard.originating_source
@@ -2570,6 +2930,10 @@ def DUPLICATE_INPUT(self, guard: Guard, source_b: Source) -> None:
                 guard.user_stack,
             )
 
+    @register_guard_check_spec(
+        get_metadata_fn=lambda guard, value: None,
+        eval_fn=lambda value, metadata: value is not None,
+    )
     def WEAKREF_ALIVE(self, guard: Guard) -> None:
         code = [f"{self.arg_ref(guard)} is not None"]
 
@@ -2578,6 +2942,10 @@ def WEAKREF_ALIVE(self, guard: Guard) -> None:
             get_verbose_code_parts(code, guard), guard.user_stack
         )
 
+    @register_guard_check_spec(
+        get_metadata_fn=lambda guard, value: list(value.keys()),
+        eval_fn=lambda value, metadata: list(value.keys()) == metadata,
+    )
     def MAPPING_KEYS_CHECK(self, guard: Guard) -> None:
         """Guard on the key order of types.MappingProxyType object"""
         ref = self.arg_ref(guard)
@@ -2590,6 +2958,10 @@ def MAPPING_KEYS_CHECK(self, guard: Guard) -> None:
             value, code, guard.user_stack
         )
 
+    @register_guard_check_spec(
+        get_metadata_fn=lambda guard, value: list(dict.keys(value)),
+        eval_fn=lambda value, metadata: list(dict.keys(value)) == metadata,
+    )
     def DICT_KEYS_MATCH(self, guard: Guard) -> None:
         """Insert guard to check that the keys of a dict are same"""
         ref = self.arg_ref(guard)
@@ -2615,6 +2987,10 @@ def DICT_KEYS_MATCH(self, guard: Guard) -> None:
         else:
             self.guard_on_dict_keys_and_ignore_order(value, guard)
 
+    @register_guard_check_spec(
+        get_metadata_fn=lambda guard, value: None,
+        eval_fn=lambda value, metadata: len(value) == 0,
+    )
     def EMPTY_NN_MODULE_HOOKS_DICT(self, guard: Guard) -> None:
         """Special guard to skip guards on empty hooks. This is controlled by skip_nnmodule_hook_guards"""
         if config.skip_nnmodule_hook_guards:
@@ -2622,15 +2998,23 @@ def EMPTY_NN_MODULE_HOOKS_DICT(self, guard: Guard) -> None:
             return
         self.SEQUENCE_LENGTH(guard)
 
+    # Global state guard — not source-specific, checked separately at runtime.
+    @skip_guard_check_spec
     def GRAD_MODE(self, guard: Guard) -> None:
         pass  # we always guard on this via GlobalStateGuard()
 
+    # Global state guard — not source-specific, checked separately at runtime.
+    @skip_guard_check_spec
     def DETERMINISTIC_ALGORITHMS(self, guard: Guard) -> None:
         pass  # we always guard on this via GlobalStateGuard()
 
+    # Global state guard — not source-specific, checked separately at runtime.
+    @skip_guard_check_spec
     def FSDP_TRAINING_STATE(self, guard: Guard) -> None:
         pass  # we always guard on this via GlobalStateGuard()
 
+    # Global state guard — not source-specific, checked separately at runtime.
+    @skip_guard_check_spec
     def GLOBAL_STATE(self, guard: Guard) -> None:
         output_graph = self.check_fn_manager.output_graph
         assert output_graph is not None
@@ -2645,6 +3029,8 @@ def GLOBAL_STATE(self, guard: Guard) -> None:
             global_state, code, guard.user_stack
         )
 
+    # Global state guard — not source-specific, checked separately at runtime.
+    @skip_guard_check_spec
     def TORCH_FUNCTION_STATE(self, guard: Guard) -> None:
         assert self.check_fn_manager.torch_function_mode_stack is not None
         self.check_fn_manager.torch_function_mode_stack_check_fn = (
@@ -2658,6 +3044,8 @@ def TORCH_FUNCTION_STATE(self, guard: Guard) -> None:
             guard.user_stack,
         )
 
+    # Global state guard — not source-specific, checked separately at runtime.
+    @skip_guard_check_spec
     def DEFAULT_DEVICE(self, guard: Guard) -> None:
         """Guard on CURRENT_DEVICE per torch.utils._device"""
         assert guard.source is GuardSource.GLOBAL
@@ -2672,6 +3060,7 @@ def DEFAULT_DEVICE(self, guard: Guard) -> None:
             get_verbose_code_parts(code, guard), guard.user_stack
         )
 
+    @skip_guard_check_spec
     def SHAPE_ENV(self, guard: Guard) -> None:
         from torch._dynamo.output_graph import OutputGraphCommon
 
@@ -2709,7 +3098,7 @@ def get_sources(t_id: int, dim: int) -> list[Source]:
                 source_pairs: list[tuple[Source, Source]] = []
                 derived_equalities: list[  # type: ignore[type-arg]
                     # pyrefly: ignore [implicit-any]
-                    tuple[Source, Union[Source, Symbol], Callable]
+                    tuple[Source, Source | Symbol, Callable]
                 ] = []
                 phantom_symbols: dict[str, Symbol] = {}
                 relaxed_sources: set[Source] = set()
@@ -2921,12 +3310,18 @@ def _get_code_parts(langs: tuple[str, ...]) -> list[_ShapeGuardsHelper]:
                 closure_vars={**SYMPY_INTERP, **_get_closure_vars()},
             )
 
-    def TENSOR_MATCH(self, guard: Guard, value: Optional[Any] = None) -> None:
+    @register_guard_check_spec(
+        get_metadata_fn=lambda guard, value: extract_tensor_metadata(value),
+        eval_fn=lambda value, metadata: (
+            isinstance(value, torch.Tensor)
+            and extract_tensor_metadata(value) == metadata
+        ),
+    )
+    def TENSOR_MATCH(self, guard: Guard, value: Any | None = None) -> None:
         if config._unsafe_skip_fsdp_module_guards and guard.is_fsdp_module():
             return
         # For tensors that are part of the Dynamo extracted Fx graph module, an
-        # ID_MATCH suffices. Once we turn on inline_inbuilt_nn_modules, these
-        # will be lifted as inputs and have a TENSOR_MATCH guard.
+        # ID_MATCH suffices.
         if match_on_id_for_tensor(guard):
             self.ID_MATCH(guard)
         else:
@@ -3113,10 +3508,13 @@ def TENSOR_MATCH(self, guard: Guard, value: Optional[Any] = None) -> None:
                         guard.user_stack,
                     )
 
-                # Guard on shape_ids when tensor has unbacked indices.
-                # shape_id is only set via mark_unbacked, which sets _dynamo_unbacked_indices.
-                # Empty dict is treated the same as not having the attribute.
-                if shape_ids := getattr(value, "_dynamo_shape_ids", None):
+                # Guard on shape_ids for tensors marked with mark_unbacked().
+                # - If the runtime tensor has _dynamo_unbacked_indices → check shape_ids match
+                # - If the runtime tensor doesn't have _dynamo_unbacked_indices → pass
+                # We must install guards even when shape_ids is None to detect runtime
+                # tensors that have the attribute when compile-time didn't.
+                if hasattr(value, "_dynamo_unbacked_indices"):
+                    shape_ids = getattr(value, "_dynamo_shape_ids", None)
                     code_part = f"((getattr({tensor_name}, '_dynamo_shape_ids', None) == {shape_ids!r}) if hasattr({tensor_name}, '_dynamo_unbacked_indices') else True)"  # noqa: B950
                     code.append(code_part)
                     self.get_guard_manager(guard).add_lambda_guard(
@@ -3128,7 +3526,27 @@ def TENSOR_MATCH(self, guard: Guard, value: Optional[Any] = None) -> None:
                         get_verbose_code_parts(code_part, guard),
                         guard.user_stack,
                     )
-                    # TODO we dont have guards on _dynamo_unbacked_indices like those of _dynamo_dynamic_indices this seems wrong!!
+
+                # Guard on unbacked_bounds for tensors marked with mark_unbacked().
+                # - If the runtime tensor has _dynamo_unbacked_indices → check bounds match
+                # - If the runtime tensor doesn't have _dynamo_unbacked_indices → pass
+                # We must install guards even when unbacked_bounds is None to detect runtime
+                # tensors that have the attribute when compile-time didn't.
+                if hasattr(value, "_dynamo_unbacked_indices"):
+                    unbacked_bounds = getattr(value, "_dynamo_unbacked_bounds", None)
+                    code_part = f"((getattr({tensor_name}, '_dynamo_unbacked_bounds', None) == {unbacked_bounds!r}) if hasattr({tensor_name}, '_dynamo_unbacked_indices') else True)"  # noqa: B950
+                    code.append(code_part)
+                    self.get_guard_manager(guard).add_lambda_guard(
+                        lambda x, expected=unbacked_bounds: (
+                            getattr(x, "_dynamo_unbacked_bounds", None) == expected
+                            if hasattr(x, "_dynamo_unbacked_indices")
+                            else True
+                        ),
+                        get_verbose_code_parts(code_part, guard),
+                        guard.user_stack,
+                    )
+
+                # TODO we dont have guards on _dynamo_unbacked_indices like those of _dynamo_dynamic_indices this seems wrong!!
 
             if len(code) > 0:
                 self._set_guard_export_info(guard, code)
@@ -3138,8 +3556,8 @@ def _set_guard_export_info(
         self,
         guard: Guard,
         code_list: list[str],
-        provided_guarded_object: Optional[Any] = None,
-        provided_func_name: Optional[str] = None,
+        provided_guarded_object: Any | None = None,
+        provided_func_name: str | None = None,
     ) -> None:
         # WARNING: It is important that cur_frame/caller do NOT stay in
         # the current frame, because they will keep things live longer
@@ -3315,7 +3733,7 @@ def populate_diff_guard_manager(self) -> None:
 class ShapeCodeParts:
     python_code_parts: _ShapeGuardsHelper
     verbose_code_parts: _ShapeGuardsHelper
-    cpp_code_parts: Optional[_CppShapeGuardsHelper]
+    cpp_code_parts: _CppShapeGuardsHelper | None
     python_fallback: bool
     shape_env_sources: list[Source]
 
@@ -3323,11 +3741,11 @@ class ShapeCodeParts:
 @dataclasses.dataclass
 class GuardsState:
     output_graph: OutputGraphGuardsState
-    shape_code_parts: Optional[ShapeCodeParts]
+    shape_code_parts: ShapeCodeParts | None
 
 
 class _Missing:
-    def __init__(self, reason: Optional[str] = None) -> None:
+    def __init__(self, reason: str | None = None) -> None:
         self._reason = reason
 
     def __repr__(self) -> str:
@@ -3407,12 +3825,9 @@ def _unpickle_traceable_wrapper_subclass(
         pytype: type,
         dispatch_keys_raw: int,
         ctx: Any,
-        inner_data: list[tuple[str, Callable[..., Any], tuple[Any, ...]]],
+        inner_data: list[tuple[str, Any]],
     ) -> torch.Tensor:
-        # Unpickle the inner tensor components. These could also be subclass instances.
-        inner_tensors = {}
-        for attr, unpickle_func, unpickle_func_args in inner_data:
-            inner_tensors[attr] = unpickle_func(*unpickle_func_args)
+        inner_tensors = dict(inner_data)
 
         outer_size, outer_stride = meta_tensor.shape, meta_tensor.stride()
         out = type(meta_tensor).__tensor_unflatten__(  # type: ignore[attr-defined]
@@ -3446,6 +3861,10 @@ def _unpickle_mapping_proxy(
     def _unpickle_dict_keys(cls, elems: list[Any]) -> Any:
         return dict.fromkeys(elems).keys()
 
+    @classmethod
+    def _unpickle_count_iter(cls, item: int, step: int) -> itertools.count[int]:
+        return itertools.count(item, step)
+
     @classmethod
     def _unpickle_fsdp_module_type(
         cls, original_type: type[torch.nn.Module]
@@ -3516,7 +3935,7 @@ def _unpickle_nested_function(
     # pyrefly: ignore [bad-override]
     def reducer_override(
         self, obj: Any
-    ) -> Union[tuple[Callable[..., Any], tuple[Any, ...]], Any]:
+    ) -> tuple[Callable[..., Any], tuple[Any, ...]] | Any:
         import sympy
 
         if id(obj) in self.empty_values:
@@ -3547,8 +3966,7 @@ def reducer_override(
                     inner = getattr(obj, attr)
                     if isinstance(inner, torch.Tensor):
                         self.guard_tree_values[id(inner)] = inner
-                    func, args_tuple = self.reducer_override(inner)
-                    inner_data.append((attr, func, args_tuple))
+                    inner_data.append((attr, inner))
 
                 return type(self)._unpickle_traceable_wrapper_subclass, (
                     torch.empty_like(obj, device="meta"),
@@ -3629,6 +4047,11 @@ def reducer_override(
         elif isinstance(obj, types.MappingProxyType):
             return type(self)._unpickle_mapping_proxy, (obj.copy(),)
 
+        elif type(obj) is _COUNT_ITERATOR_TYPE:
+            item, step = normalize_count_iter(obj)
+            if item is not NotImplemented and step is not NotImplemented:
+                return type(self)._unpickle_count_iter, (item, step)
+
         elif isinstance(obj, torch._dynamo.utils.dict_keys):
             return type(self)._unpickle_dict_keys, (list(obj),)
 
@@ -3797,12 +4220,12 @@ def __init__(
         self,
         f_code: types.CodeType,
         output_graph: OutputGraphCommon,
-        cache_entry: Optional[CacheEntry] = None,
-        guard_fail_fn: Optional[Callable[[GuardFail], None]] = None,
+        cache_entry: CacheEntry | None = None,
+        guard_fail_fn: Callable[[GuardFail], None] | None = None,
         guard_filter_fn: Callable[[Sequence[GuardFilterEntry]], Sequence[bool]]
         | None = None,
-        shape_code_parts: Optional[ShapeCodeParts] = None,
-        runtime_global_scope: Optional[dict[str, Any]] = None,
+        shape_code_parts: ShapeCodeParts | None = None,
+        runtime_global_scope: dict[str, Any] | None = None,
         save_guards: bool = False,
         strict_error: bool = False,
     ) -> None:
@@ -3812,7 +4235,7 @@ def __init__(
         existing_diff_guard_sources = (
             update_diff_guard_managers_for_existing_cache_entries(cache_entry)
         )
-        self.output_graph: Optional[OutputGraphCommon] = output_graph
+        self.output_graph: OutputGraphCommon | None = output_graph
         assert self.output_graph is not None
 
         # Only used for serialization.
@@ -3827,8 +4250,8 @@ def __init__(
         self.additional_used_local_vars: OrderedSet[str] = OrderedSet()
         self.additional_used_global_vars: OrderedSet[str] = OrderedSet()
         self.runtime_global_scope = runtime_global_scope
-        self.global_state: Optional[torch._C._dynamo.guards.GlobalStateGuard] = None
-        self.torch_function_mode_stack_check_fn: Optional[Callable[[], bool]] = None
+        self.global_state: torch._C._dynamo.guards.GlobalStateGuard | None = None
+        self.torch_function_mode_stack_check_fn: Callable[[], bool] | None = None
 
         if not justknobs_check("pytorch/compiler:guard_nn_modules"):
             log.warning("guard_nn_modules is turned off using justknobs killswitch")
@@ -3865,38 +4288,43 @@ def guard_filter_fn(guards: Sequence[GuardFilterEntry]) -> Sequence[bool]:
 
         sorted_guards = sorted(guards or (), key=Guard.sort_key)
 
-        if guard_filter_fn:
-            # If we're filtering guards, we need to build it an extra time first
-            # because filtering depends on the builder/guard_manager results
+        # Disable __torch_function__ dispatch during guard construction so
+        # modes with mutable state aren't triggered.  We exit the context
+        # before the guard sanity check so GlobalStateGuard.check() sees
+        # the true runtime state.
+        with torch._C.DisableTorchFunction():
+            if guard_filter_fn:
+                # If we're filtering guards, we need to build it an extra time first
+                # because filtering depends on the builder/guard_manager results
+                builder, guard_manager = self.build_guards(
+                    sorted_guards,
+                    existing_diff_guard_sources,
+                    f_code,
+                    output_graph,
+                    False,
+                )
+
+                filter_results = guard_filter_fn(
+                    [make_guard_filter_entry(guard, builder) for guard in sorted_guards]
+                )
+                assert len(filter_results) == len(sorted_guards)
+                assert all(type(x) is bool for x in filter_results)
+                sorted_guards = [
+                    guard for i, guard in enumerate(sorted_guards) if filter_results[i]
+                ]
+
+            # Redo the guards because filtering relies on the results from the last guard builder.
             builder, guard_manager = self.build_guards(
                 sorted_guards,
                 existing_diff_guard_sources,
                 f_code,
                 output_graph,
-                False,
+                save_guards,
+                guard_filter_fn=guard_filter_fn,
             )
 
-            filter_results = guard_filter_fn(
-                [make_guard_filter_entry(guard, builder) for guard in sorted_guards]
-            )
-            assert len(filter_results) == len(sorted_guards)
-            assert all(type(x) is bool for x in filter_results)
-            sorted_guards = [
-                guard for i, guard in enumerate(sorted_guards) if filter_results[i]
-            ]
-
-        # Redo the guards because filtering relies on the results from the last guard builder.
-        builder, guard_manager = self.build_guards(
-            sorted_guards,
-            existing_diff_guard_sources,
-            f_code,
-            output_graph,
-            save_guards,
-            guard_filter_fn=guard_filter_fn,
-        )
-
-        self.guard_manager = guard_manager
-        self.compile_check_fn(builder, sorted_guards, guard_fail_fn)
+            self.guard_manager = guard_manager
+            self.compile_check_fn(builder, sorted_guards, guard_fail_fn)
 
         # Keep track of weak references of objects with ID_MATCH guard. This
         # info is stored alongside optimized_code and guard_manager and is used to
@@ -3956,7 +4384,7 @@ def guard_filter_fn(guards: Sequence[GuardFilterEntry]) -> Sequence[bool]:
             # account for, we simply increment at the toplevel instead.
             CompileEventLogger.increment_toplevel("guard_latency_us", int(latency))
 
-        self.guards_state: Optional[bytes] = None
+        self.guards_state: bytes | None = None
         if save_guards:
             from torch._dynamo.output_graph import OutputGraphCommon
 
@@ -4195,7 +4623,7 @@ def compile_check_fn(
         self,
         builder: GuardBuilder,
         guards_out: list[Guard],
-        guard_fail_fn: Optional[Callable[[GuardFail], None]],
+        guard_fail_fn: Callable[[GuardFail], None] | None,
     ) -> None:
         # see parallel handling of ".0" / "___implicit0" in _eval_frame.c
         largs = builder.argnames
@@ -4217,7 +4645,7 @@ def compile_check_fn(
         self.torch_function_mode_stack = None
 
         def add_code_part(
-            code_part: str, guard: Optional[Guard], log_only: bool = False
+            code_part: str, guard: Guard | None, log_only: bool = False
         ) -> None:
             verbose_code_part = get_verbose_code_part(code_part, guard)
             guards_log.debug("%s", verbose_code_part)
@@ -4427,7 +4855,7 @@ def id_ref(self, obj: object, obj_str: str) -> int:
             pass  # cannot weakref bool object
         return id(obj)
 
-    def lookup_weakrefs(self, obj: object) -> Optional[weakref.ref[object]]:
+    def lookup_weakrefs(self, obj: object) -> weakref.ref[object] | None:
         """Lookup the _weakrefs created in id_ref function for ID_MATCH'd objects"""
         if id(obj) in self._weakrefs:
             return self._weakrefs[id(obj)]
@@ -4569,9 +4997,9 @@ def format_user_stack_trace(
 def get_guard_fail_reason_helper(
     guard_manager: GuardManagerWrapper,
     f_locals: dict[str, object],
-    compile_id: Optional[CompileId],
+    compile_id: CompileId | None,
     # pyrefly: ignore [implicit-any]
-    backend: Optional[Callable],
+    backend: Callable | None,
 ) -> str:
     """
     Return the reason why `guard_manager` failed.
@@ -4705,7 +5133,7 @@ def get_guard_fail_reason(
 
 
 def get_and_maybe_log_recompilation_reasons(
-    cache_entry: Optional[CacheEntry],
+    cache_entry: CacheEntry | None,
     frame: DynamoFrameType,
     # pyrefly: ignore [implicit-any]
     backend: Callable,
@@ -4774,7 +5202,7 @@ def get_and_maybe_log_recompilation_reasons(
 
 
 def update_diff_guard_managers_for_existing_cache_entries(
-    cache_entry: Optional[CacheEntry],
+    cache_entry: CacheEntry | None,
 ) -> OrderedSet[str]:
     first_cache_entry = cache_entry
 
@@ -4836,7 +5264,7 @@ def unique(seq: Sequence[T]) -> Generator[T, None, None]:
 
 def make_dupe_guard(
     obj_source: Source, dupe_source: Source | None
-) -> Optional[functools.partial[Any]]:
+) -> functools.partial[Any] | None:
     # Note - we may end up in a situation where we invoke something like
     # def fn(x, y)
     # with fn(x, x)
@@ -4880,12 +5308,47 @@ def install_guard(*guards: Guard, skip: int = 0) -> None:
     """
     from torch._guards import TracingContext
 
+    guards_context = TracingContext.get().guards_context
+    if guards_context.skip_install:
+        return
+
     collect_debug_stack = guards_log.isEnabledFor(
         logging.DEBUG
     ) or verbose_guards_log.isEnabledFor(logging.DEBUG)
-    add = TracingContext.get().guards_context.dynamo_guards.add
+    add = guards_context.dynamo_guards.add
     for guard in guards:
         assert isinstance(guard, Guard)
         if is_from_skip_guard_source(guard.originating_source):
             continue
         add(guard, collect_debug_stack=collect_debug_stack, skip=skip + 1)
+
+
+# ---------------------------------------------------------------------------
+# Guard check spec dispatch table — built by scanning GuardBuilder methods
+# for the guard_check_spec attribute set by the decorators above.
+GUARD_VALUE_DISPATCH: dict[str, GuardCheckSpec | object] = {
+    name: handler
+    for name in dir(GuardBuilder)
+    if (handler := getattr(getattr(GuardBuilder, name), "guard_check_spec", None))
+    is not None
+}
+
+# Validate that every guard method has an explicit guard_check_spec.
+# Guard methods are identified by: uppercase name, callable, and first
+# non-self parameter is 'guard'.
+MISSING_GUARD_SPECS = [
+    _name
+    for _name in dir(GuardBuilder)
+    if _name[0].isupper()
+    and not _name.startswith("_")
+    and callable(_method := getattr(GuardBuilder, _name, None))
+    and len(_params := list(inspect.signature(_method).parameters)) >= 2
+    and _params[1] == "guard"
+    and not hasattr(_method, "guard_check_spec")
+]
+if MISSING_GUARD_SPECS:
+    raise RuntimeError(
+        f"GuardBuilder methods missing guard_check_spec decorator: "
+        f"{', '.join(MISSING_GUARD_SPECS)}. Use @register_guard_check_spec, "
+        f"@skip_guard_check_spec, or @unsupported_guard_check_spec."
+    )
diff --git a/torch/_dynamo/metrics_context.py b/torch/_dynamo/metrics_context.py
index e20892b434029..41ca9a94a1e67 100644
--- a/torch/_dynamo/metrics_context.py
+++ b/torch/_dynamo/metrics_context.py
@@ -19,7 +19,7 @@
 import logging
 import time
 from collections.abc import Callable
-from typing import Any, Optional, TYPE_CHECKING, TypeAlias
+from typing import Any, TYPE_CHECKING, TypeAlias
 from typing_extensions import Self
 
 
@@ -54,7 +54,7 @@ def __iter__(self) -> Iterator[tuple[Any, int]]:
 
 
 OnExitType: TypeAlias = Callable[
-    [int, int, dict[str, Any], Optional[type[BaseException]], Optional[BaseException]],
+    [int, int, dict[str, Any], type[BaseException] | None, BaseException | None],
     None,
 ]
 
@@ -87,8 +87,8 @@ def __enter__(self) -> Self:
 
     def __exit__(
         self,
-        exc_type: Optional[type[BaseException]],
-        exc_value: Optional[BaseException],
+        exc_type: type[BaseException] | None,
+        exc_value: BaseException | None,
         _traceback: Any,
     ) -> None:
         """
@@ -218,7 +218,7 @@ def __init__(self, on_exit: OnExitType) -> None:
         self._start_time_ns: int = 0
 
     def increment(
-        self, metric: str, value: int, extra: Optional[dict[str, Any]] = None
+        self, metric: str, value: int, extra: dict[str, Any] | None = None
     ) -> None:
         """
         Increment a metric by a given amount.
diff --git a/torch/_dynamo/mutation_guard.py b/torch/_dynamo/mutation_guard.py
index 0467ea1ba1164..70f2d7dfc56fa 100644
--- a/torch/_dynamo/mutation_guard.py
+++ b/torch/_dynamo/mutation_guard.py
@@ -117,10 +117,8 @@ def is_dynamic_nn_module(obj: Any, is_export: bool) -> bool:
         return True
     if hasattr(obj, "torchdynamo_force_dynamic"):
         return obj.torchdynamo_force_dynamic
-    if (
-        isinstance(obj, torch.nn.Module)
-        and config.inline_inbuilt_nn_modules
-        and (not is_export or config.install_free_tensors)
+    if isinstance(obj, torch.nn.Module) and (
+        not is_export or config.install_free_tensors
     ):
         return True
 
diff --git a/torch/_dynamo/output_graph.py b/torch/_dynamo/output_graph.py
index 295cf03763dcc..5d292aaac616f 100644
--- a/torch/_dynamo/output_graph.py
+++ b/torch/_dynamo/output_graph.py
@@ -32,6 +32,7 @@
 import sys
 import time
 import traceback
+import types
 import warnings
 import weakref
 from collections.abc import Callable, Generator, Sequence
@@ -119,6 +120,7 @@
     AttrSource,
     BackwardStateSource,
     ConstantSource,
+    DictGetItemSource,
     GetItemSource,
     GlobalStateSource,
     is_constant_source,
@@ -178,6 +180,7 @@
     from torch._dynamo.dynamo_profiler import DynamoProfilerState
     from torch._dynamo.package import CompilePackage
     from torch._dynamo.symbolic_convert import InstructionTranslatorBase
+    from torch._inductor import _CudagraphAnnotation
     from torch.multiprocessing.reductions import StorageWeakRef
 
 log = logging.getLogger(__name__)
@@ -201,12 +204,13 @@ def _wrap_with_inductor_config(
 ) -> Callable[..., Any]:
     """
     Wrap a compiler function to apply inductor config patches during compilation.
+
+    Passes config_patches as a keyword argument so that compile_fx can
+    propagate them to backward compilation via its inner_compile wrapping.
     """
-    from torch._inductor import config as inductor_config
 
     def wrapped(gm: Any, example_inputs: Any) -> Any:
-        with inductor_config.patch(config_patches):
-            return compiler_fn(gm, example_inputs)
+        return compiler_fn(gm, example_inputs, config_patches=config_patches)
 
     # Preserve function metadata for logging
     wrapped.__name__ = getattr(compiler_fn, "__name__", "<wrapped>")
@@ -401,7 +405,7 @@ class OutputGraphGuardsState:
     input_source_to_sizes_strides: dict[Source, dict[str, Any]]
     dual_level: int
     functorch_layers: list[torch._functorch.pyfunctorch.FuncTorchInterpreter]
-    current_device: Optional[torch.device]
+    current_device: torch.device | None
     global_state_guard: torch._C._dynamo.guards.GlobalStateGuard
     _guards: torch._guards.GuardsSet
     _aotautograd_guards: list[torch._guards.GuardEnvExpr]
@@ -411,7 +415,7 @@ class OutputGraphGuardsState:
     export: bool = False
     skip_guards_check: bool = False
     export_constraints: bool = False
-    name_of_builtins_dict_key_in_fglobals: Optional[str] = None
+    name_of_builtins_dict_key_in_fglobals: str | None = None
 
     @property
     def shape_env(self) -> ShapeEnv:
@@ -478,12 +482,12 @@ class ExportMetaData:
     # 3) constants
     output_return_type: dict[int, tuple[str, Any]] = dc_field(default_factory=dict)
     # output spec of the traced function
-    out_spec: Union[torch.utils._pytree.TreeSpec, torch.utils._pytree.LeafSpec] = (
+    out_spec: torch.utils._pytree.TreeSpec | torch.utils._pytree.LeafSpec = (
         torch.utils._pytree._LEAF_SPEC
     )
     module_call_spec: dict[
         str,
-        dict[str, Union[torch.utils._pytree.TreeSpec, torch.utils._pytree.LeafSpec]],
+        dict[str, torch.utils._pytree.TreeSpec | torch.utils._pytree.LeafSpec],
     ] = dc_field(default_factory=dict)
 
 
@@ -534,10 +538,10 @@ class OutputGraphCommon(OutputGraphGuardsState):
     def __init__(
         self,
         output_graph_guards_state: OutputGraphGuardsState,
-        import_sources: Optional[dict[str, str]] = None,
-        shape_env: Optional[ShapeEnv] = None,
-        export_metadata: Optional[ExportMetaData] = None,
-        tracked_fakes_id_to_source: Optional[dict[int, list[Source]]] = None,
+        import_sources: dict[str, str] | None = None,
+        shape_env: ShapeEnv | None = None,
+        export_metadata: ExportMetaData | None = None,
+        tracked_fakes_id_to_source: dict[int, list[Source]] | None = None,
     ) -> None:
         super().__init__(
             output_graph_guards_state.local_scope,
@@ -596,7 +600,7 @@ class OutputGraph(OutputGraphCommon):
     def __init__(
         self,
         code_options: dict[str, Any],
-        compiler_fn: Optional[CompilerFn],
+        compiler_fn: CompilerFn | None,
         root_tx: "InstructionTranslatorBase",
         export: bool,
         export_constraints: Sequence[_ConstraintTarget],
@@ -650,7 +654,13 @@ def __init__(
             "co_firstlineno": f_code.co_firstlineno,
         }
 
+        # Cudagraph annotation is set during inlining in
+        # InliningInstructionTranslator.build_inline_tracer when a decorated
+        # function is inlined.
+        self.cudagraph_annotation: _CudagraphAnnotation | None = None
+
         self.region_tracker = GraphRegionTracker()
+        self._emit_debugger_breakpoint: bool = False
 
         # tracked_fakes says where any tensor that was wrapped to fake came
         # from.  It is similar to GraphArg, in that all GraphArgs will get
@@ -689,10 +699,9 @@ def __init__(
             )
         self.tracing_context: TracingContext = TracingContext(fake_mode)
         self.tracing_context.traced_code.append(f_code)
+        self.tracing_context.cudagraph_annotation = self.cudagraph_annotation
         self.traced_code = self.tracing_context.traced_code
-        self.dynamo_compile_id: Optional[CompileId] = (
-            CompileContext.current_compile_id()
-        )
+        self.dynamo_compile_id: CompileId | None = CompileContext.current_compile_id()
         self.init_ambient_guards()
 
         # Map each tensor id to a list of sources. This is necessary because
@@ -704,11 +713,26 @@ def __init__(
             collections.defaultdict(list)
         )
         # Stores the full fqn of a param or buffer to the relevant source.
-        self.param_name_to_source: Optional[dict[str, Source]] = {}
+        self.param_name_to_source: dict[str, Source] | None = {}
         self.side_effects = SideEffects(self)
         # Cached variable trackers. This makes symbolic analysis of LOAD_GLOBAL
         # and LOAD_ATTR for same python objects free.
         self.variable_tracker_cache: dict[Source, VariableTracker] = {}
+        # Cache for sources resolved via MRO walk, keyed by id(obj).
+        # When the same descriptor (e.g. property) is reached from multiple
+        # subclasses, we reuse the first source to avoid redundant guards.
+        # We thought of rolling this in variable_tracker_cache but here
+        # different sources point to the same object, we also don't want it to
+        # go through the side effects cache because even though these objects
+        # are same, we dont want OBJECT_ALIASING guards on them. For these
+        # objects, we have DICT_CONTAINS absent guards on the mro walk, so there
+        # is no need of the OBJECT_ALIASING guards.
+        self.mro_source_cache: dict[tuple[int, str], DictGetItemSource] = {}
+        # Tracks (id(klass), attr_name) pairs that already have a
+        # DICT_CONTAINS absent guard installed during MRO walks.  When
+        # multiple subclasses share the same intermediate MRO class, we
+        # only need to guard the absence once per (class, attr) pair.
+        self.guarded_mro_absent_keys: set[tuple[int, str]] = set()
         # Cache for inspect.signature results: function -> VariableTracker
         self.signature_cache: dict[Any, VariableTracker] = {}
         self.unique_var_id = itertools.count()
@@ -718,15 +742,23 @@ def __init__(
         # and restore_graphstate
         self.timestamp = 0
 
+        # Maps stream id (id(stream_value)) → user stack trace for input
+        # mutations on that stream.  Used to error when an event records on a
+        # stream that already has an input mutation (the epilogue copy_()
+        # wouldn't be captured).  We key by id() of the underlying
+        # torch.Stream so we can peek lazy variables without realizing them.
+        self._input_mutation_streams: dict[int, traceback.StackSummary] = {}
+        self._last_checked_input_versions: dict[int, int] | None = None
+
         # A list of register_finalizer_fns to apply to the output graph module
         self.register_finalizer_fns: list[Callable[[fx.GraphModule], None]] = []
 
         # Not checkpointed
-        self.compiler_fn: Optional[CompilerFn] = compiler_fn
+        self.compiler_fn: CompilerFn | None = compiler_fn
         self.root_tx = root_tx
 
         # Profiler state for tracking function trace timings
-        self.profiler_state: Optional[DynamoProfilerState] = None
+        self.profiler_state: DynamoProfilerState | None = None
 
         self.package = package
         # Given a source, what are the user stacks of all locations that
@@ -749,6 +781,10 @@ def __init__(
         # This returns false if TF Overall (both mode and subclass) is disabled OR that TF Mode stack is empty
         self.torch_function_mode_enabled = torch._C._is_torch_function_mode_enabled()
 
+        # Used to wrap the compiled graph at runtime with
+        # DisableTorchFunctionSubclass to prevent double dispatch.
+        self.torch_function_subclass_inlined = False
+
         # Tracks if the output graph has a user defined allowed function in the
         # graph. This is used later to determine if we should fallback to eager
         # for certain exceptions. THe idea is that if the user has applied
@@ -794,10 +830,21 @@ def __init__(
         # Bytecode to insert right before we call the graph
         self.pregraph_bytecode: list[Instruction] = []
 
+        # Maps SyntheticLocalSource → (fn, args, arg_sources) for hoisted
+        # graph inputs created by synthetic_graph_input, so invoke_subgraph
+        # reuse can recreate them on cache hit. arg_sources allows stamp-out
+        # to resolve new ctor arg values via source replacement.
+        self.synthetic_source_ctor_info: dict[
+            SyntheticLocalSource,
+            tuple[
+                Callable[..., Any], tuple[Any, ...], tuple[Source | None, ...] | None
+            ],
+        ] = {}
+
         # Use to pass values to backward hooks when using compiled autograd
         self.backward_state: dict[str, VariableTracker] = {}
-        self.backward_state_proxy: Optional[torch.fx.Proxy] = None
-        self.backward_state_var: Optional[str] = None
+        self.backward_state_proxy: torch.fx.Proxy | None = None
+        self.backward_state_var: str | None = None
 
         # pyrefly: ignore [bad-override]
         self.name_of_builtins_dict_key_in_fglobals: str = (
@@ -812,7 +859,7 @@ def __init__(
         # that are set for the entire duration of the compiled region.
         # This is an invariant today because we graph break on the saved_tensor_hook
         # context manager inside a compiled region
-        self.saved_tensors_hooks_subgraph_names: Optional[list[str]] = (
+        self.saved_tensors_hooks_subgraph_names: list[str] | None = (
             self.maybe_install_saved_tensors_hooks_subgraphs()
         )
 
@@ -826,6 +873,62 @@ def __init__(
         self.used_inlined_inbuilt_modules_names: OrderedSet[str] = OrderedSet()
 
         self.attr_source_cache: dict[tuple[Source, str], AttrSource] = {}
+        self._cached_replayed_side_effect_source_refs: tuple[str, ...] | None = None
+
+    def get_replayed_side_effect_source_refs(
+        self, *, populate_export_metadata: bool = False
+    ) -> list[str]:
+        """Return Python-side effect sources that Dynamo replays outside the FX graph."""
+        if (
+            not populate_export_metadata
+            and not self.side_effects.id_to_variable
+            and self._cached_replayed_side_effect_source_refs is not None
+        ):
+            return list(self._cached_replayed_side_effect_source_refs)
+
+        from torch.export._trace import _ExportModuleSpecTrackerDict
+
+        potential_side_effects = []
+        for var in self.side_effects._get_modified_vars():
+            if hasattr(var, "mutation_type"):
+                mut_type = var.mutation_type
+                # Skip codegen-specific mutations that never materialize as
+                # externally visible Python side effects.
+                if isinstance(
+                    mut_type, (AttributeMutationExisting, ValueMutationExisting)
+                ):
+                    if isinstance(var, UserDefinedDictVariable) and isinstance(
+                        var.value, _ExportModuleSpecTrackerDict
+                    ):
+                        if populate_export_metadata:
+                            assert var._base_vt is not None
+                            for (
+                                k,
+                                v,
+                            ) in (
+                                var._base_vt.items.items()  # pyrefly: ignore[missing-attribute]
+                            ):
+                                # pyrefly: ignore [implicit-any]
+                                specs = {}
+                                # pyrefly: ignore[missing-attribute]
+                                for k_spec, val in v.items.items():
+                                    specs[k_spec.vt.as_python_constant()] = (
+                                        val.as_python_constant()
+                                    )
+                                assert ["in_spec", "out_spec"] == list(specs.keys())
+                                self.export_metadata.module_call_spec[
+                                    # pyrefly: ignore[missing-attribute]
+                                    k.vt.as_python_constant()
+                                ] = specs
+                    # export uses tracepoint pass to dump submodule inp/out spec
+                    # into global state, so we filter it here
+                    if not (
+                        isinstance(var, UserDefinedDictVariable)
+                        and isinstance(var.value, _ExportModuleSpecTrackerDict)
+                    ):
+                        potential_side_effects.append(var)
+
+        return [_get_source_debug_name(var.source) for var in potential_side_effects]
 
     def get_chained_attr_source(self, base: Source, path: str) -> AttrSource:
         parts = path.split(".")
@@ -967,7 +1070,7 @@ def init_ambient_guards(self) -> None:
                 )
             )
 
-    def maybe_install_saved_tensors_hooks_subgraphs(self) -> Optional[list[str]]:
+    def maybe_install_saved_tensors_hooks_subgraphs(self) -> list[str] | None:
         if torch._dynamo.compiled_autograd.in_compiled_autograd_region:
             return None
 
@@ -997,7 +1100,10 @@ def maybe_install_saved_tensors_hooks_subgraphs(self) -> Optional[list[str]]:
         return [pack_subgraph_name, unpack_subgraph_name]
 
     def synthetic_graph_input(
-        self, fn: Callable[..., Any], args: tuple[Any, ...]
+        self,
+        fn: Callable[..., Any],
+        args: tuple[Any, ...],
+        ctor_arg_sources: tuple[Source | None, ...] | None = None,
     ) -> VariableTracker:
         """
         call fn(*args) before the graph runs and turn the result into a fake input.
@@ -1011,11 +1117,18 @@ def synthetic_graph_input(
                 fn.__name__,
             )
         )
-        cg.foreach(map(variables.ConstantVariable.create, args))
+        if ctor_arg_sources:
+            cg.foreach(
+                src if src is not None else variables.ConstantVariable.create(val)
+                for val, src in zip(args, ctor_arg_sources)
+            )
+        else:
+            cg.foreach(map(variables.ConstantVariable.create, args))
         cg.call_function(len(args), False)
         cg.store(varname)
         self.pregraph_bytecode.extend(cg.get_instructions())
         source = SyntheticLocalSource(varname)
+        self.synthetic_source_ctor_info[source] = (fn, args, ctor_arg_sources)
         result = VariableTracker.build(self.root_tx, example_value, source)
         # Realize the VT because we will delete the guards on it in the next line.
         result = result.realize()
@@ -1044,6 +1157,75 @@ def is_root_tracer(self) -> bool:
         # Helper to tell if we are inside the higher order operator tracing.
         return len(self.tracers) == 1
 
+    def check_input_mutation_on_current_stream(
+        self, tx: "InstructionTranslatorBase"
+    ) -> None:
+        """Record which stream index has input mutations by comparing current
+        tensor versions against the versions captured at graph input creation."""
+        if not hasattr(tx, "symbolic_stream_state"):
+            return
+        if not tx.symbolic_stream_state.in_stream_context():
+            return
+
+        tracer = self.root_tracer
+        if self._last_checked_input_versions is None:
+            self._last_checked_input_versions = dict(
+                enumerate(tracer._input_versions_at_beginning)
+            )
+
+        cur_stream_index = tx.symbolic_stream_state.cur_stream_id()
+        input_idx = 0
+        for node in tracer.graph.nodes:
+            if node.op != "placeholder":
+                break
+            example_value = node.meta.get("example_value")
+            if not isinstance(example_value, torch.Tensor):
+                continue
+            prev_version = self._last_checked_input_versions.get(input_idx)
+            cur_version = example_value._version
+            if prev_version is not None and cur_version > prev_version:
+                if cur_stream_index not in self._input_mutation_streams:
+                    self._input_mutation_streams[cur_stream_index] = (
+                        TracingContext.extract_stack()
+                    )
+                self._last_checked_input_versions[input_idx] = cur_version
+            input_idx += 1
+
+    _EVENT_INPUT_MUTATION_FIX = (
+        "To fix this, either:\n"
+        "  1. Move the input mutation after the event.record() call.\n"
+        "  2. Record the event outside the compiled region:\n"
+        "       compiled_fn(x)\n"
+        "       event.record(stream)  # after torch.compile returns\n"
+        "  3. Insert a graph break before recording:\n"
+        "       torch._dynamo.graph_break()\n"
+        "       event.record(stream)\n"
+        "  4. Record the event on a stream that has no input mutations."
+    )
+
+    def check_event_record_after_input_mutation(self, stream_index: int) -> None:
+        """Error if an event is being recorded on a stream that already has
+        an input mutation. Called at record time so ordering is naturally
+        respected — records before mutations won't trigger this."""
+        if stream_index not in self._input_mutation_streams:
+            return
+
+        mutation_stack = self._input_mutation_streams[stream_index]
+        record_stack = TracingContext.extract_stack()
+
+        msg = (
+            "An event was recorded on a stream where a graph input was "
+            "previously mutated. The input mutation is applied via copy_() "
+            "in the runtime epilogue after the graph executes, so the event "
+            "would not capture the mutation, leading to incorrect "
+            "synchronization.\n\n"
+            "Input mutation occurred here:\n"
+            f"{''.join(mutation_stack.format())}\n"
+            "Event record occurred here:\n"
+            f"{''.join(record_stack.format())}\n" + self._EVENT_INPUT_MUTATION_FIX
+        )
+        raise RuntimeError(msg)
+
     @property
     def graph(self) -> torch.fx.Graph:
         return self.current_tracer.graph
@@ -1083,9 +1265,9 @@ def remove_node(self, *args: Any, **kwargs: Any) -> None:
     @contextlib.contextmanager
     def subtracer(
         self,
-        source_target: Optional[Target],
+        source_target: Target | None,
         prior_tracer: Optional["SubgraphTracer"],
-        description: Optional[str] = None,
+        description: str | None = None,
     ) -> Generator["SubgraphTracer", None, None]:
         new_scope_ctx = enter_new_scope()
         try:
@@ -1138,7 +1320,7 @@ def aotautograd_guards(self) -> list[torch._guards.GuardEnvExpr]:
         return self.tracing_context.guards_context.aotautograd_guards
 
     def save_global_state(
-        self, out: Optional[dict[str, tuple[Callable[..., Any], bool]]] = None
+        self, out: dict[str, tuple[Callable[..., Any], bool]] | None = None
     ) -> None:
         """
         Saves to out if it is provided. Else saves to the tracing context's global_state.
@@ -1154,17 +1336,20 @@ def save_global_state(
 
         global_state["grad_enabled"] = (torch.set_grad_enabled, torch.is_grad_enabled())
 
+        gpu_type = (
+            acc.type if (acc := torch.accelerator.current_accelerator()) else "cuda"
+        )
         global_state["autocast_enabled"] = (
-            functools.partial(torch.set_autocast_enabled, "cuda"),
-            torch.is_autocast_enabled("cuda"),
+            functools.partial(torch.set_autocast_enabled, gpu_type),
+            torch.is_autocast_enabled(gpu_type),
         )
         global_state["autocast_cpu_enabled"] = (
             functools.partial(torch.set_autocast_enabled, "cpu"),
             torch.is_autocast_enabled("cpu"),
         )
         global_state["autocast_gpu_dtype"] = (  # type:ignore[assignment]
-            functools.partial(torch.set_autocast_dtype, "cuda"),
-            torch.get_autocast_dtype("cuda"),
+            functools.partial(torch.set_autocast_dtype, gpu_type),
+            torch.get_autocast_dtype(gpu_type),
         )
         global_state["autocast_cpu_dtype"] = (  # type:ignore[assignment]
             functools.partial(torch.set_autocast_dtype, "cpu"),
@@ -1185,6 +1370,19 @@ def pop_tx(self) -> "InstructionTranslatorBase":
     def current_tx(self) -> "InstructionTranslatorBase":
         return self.root_tx if not self._current_tx else self._current_tx[-1]
 
+    def resolve_source_value(self, source: Source) -> Any:
+        """
+        Resolve the runtime value a Source points to using root_tx's frame.
+
+        Useful to inspect the python object associated with the source during
+        debugging. Will also be used by invoke subgraph caching later on.
+        """
+        return source.get_value(
+            {"G": self.root_tx.f_globals, "L": self.root_tx.f_locals},
+            {},
+            {},  # type: ignore[arg-type]
+        )
+
     def count_calls(self) -> int:
         return count_calls(self.graph)
 
@@ -1194,9 +1392,9 @@ def is_empty_graph(self) -> bool:
     def has_outputs(self) -> bool:
         return len([x for x in self.graph.nodes if x.op == "output"]) > 0
 
-    def get_submodule(self, keys: str) -> Union[torch.nn.Module, Any]:
+    def get_submodule(self, keys: str) -> torch.nn.Module | Any:
         assert keys
-        obj: Union[torch.nn.Module, dict[str, torch.nn.Module]] = self.nn_modules
+        obj: torch.nn.Module | dict[str, torch.nn.Module] = self.nn_modules
         for k in keys.split("."):
             if isinstance(obj, dict):
                 obj = obj[k]
@@ -1263,7 +1461,7 @@ def register_static_attr_and_return_proxy(
 
     def register_attr_or_module(
         self,
-        target: Union[torch.nn.Module, torch.Tensor, Any],
+        target: torch.nn.Module | torch.Tensor | Any,
         *names: Any,
         **options: Any,
     ) -> VariableTracker:
@@ -1363,6 +1561,27 @@ def wrap_name(module_key: str) -> VariableTracker:
                 )
 
             # HACKY CODE REGION END
+        elif is_opaque_type(type(target)):
+            # HACKY CODE REGION BEGIN
+            # Same as SymInt/SymFloat above: piggybacking on self.nn_modules
+            # to store opaque objects as graph attributes.
+
+            tracer = self.current_tracer
+            if not self.is_root_tracer():
+                tracer = self.root_tracer
+
+            def wrap_name(module_key: str) -> VariableTracker:
+                fake_script_obj = torch._library.fake_class_registry.maybe_to_fake_obj(
+                    self.fake_mode, target
+                )
+                proxy = tracer.create_proxy("get_attr", module_key, (), {})
+                set_example_value(proxy.node, fake_script_obj)
+                return torch._dynamo.variables.script_object.TorchScriptObjectVariable.create(
+                    proxy, fake_script_obj, **options
+                )
+
+            # HACKY CODE REGION END
+
         else:
 
             def wrap_name(module_key: str) -> VariableTracker:
@@ -1655,8 +1874,9 @@ def compile_subgraph(
         # i.e. last element corresponds to root frame (1),
         # first element corresponds to current frame (N)
         all_stack_values = []
+        # pyrefly: ignore [implicit-any]
         all_stack_locals_metas = []
-        cur_tx: Optional[InstructionTranslatorBase] = tx
+        cur_tx: InstructionTranslatorBase | None = tx
         while cur_tx is not None:
             # this should have been checked by the caller
             assert all(block.can_restore() for block in cur_tx.block_stack)
@@ -1678,6 +1898,12 @@ def compile_subgraph(
 
         self.add_output_instructions(prefix_insts)
 
+        if self._emit_debugger_breakpoint:
+            from .bytecode_transformation import create_breakpoint
+
+            self.add_output_instructions(create_breakpoint())
+            self._emit_debugger_breakpoint = False
+
         assert not (self.pregraph_bytecode and self.export), (
             "export does not support pregraph_bytecode"
         )
@@ -1807,13 +2033,11 @@ def compile_subgraph(
             for val, count in pass1.uses.items():
                 # If it's already a local source, no need to cache it
                 if count > 1 and not istype(val, (SyntheticLocalSource, LocalSource)):
-                    # pyrefly: ignore [unsupported-operation]
                     tempvars[val] = None
             pass2 = PyCodegen(
                 self.root_tx,
                 root,
                 graph_output_var,
-                # pyrefly: ignore [bad-argument-type]
                 tempvars=tempvars,
                 overridden_sources=overridden_sources,
             )
@@ -1851,7 +2075,7 @@ def compile_subgraph(
                         elif (
                             vt.source is not None
                             and (source := getattr(vt.source, "base", None))  # type: ignore[assignment]
-                            and source.is_input
+                            and getattr(source, "is_input", False)
                         ):
                             self.export_metadata.output_return_type[idx] = (
                                 "input",
@@ -2010,43 +2234,9 @@ def compile_subgraph(
             )
 
         if torch._dynamo.config.side_effect_replay_policy in ["warn", "error"]:
-            from torch.export._trace import _ExportModuleSpecTrackerDict
-
-            potential_side_effects = []
-            for var in self.side_effects._get_modified_vars():
-                if hasattr(var, "mutation_type"):
-                    mut_type = var.mutation_type
-                    # Make sure to skip codegen specific mutations
-                    if isinstance(
-                        mut_type, (AttributeMutationExisting, ValueMutationExisting)
-                    ):
-                        if isinstance(var, UserDefinedDictVariable) and isinstance(
-                            var.value, _ExportModuleSpecTrackerDict
-                        ):
-                            for k, v in var.items.items():
-                                # pyrefly: ignore [implicit-any]
-                                specs = {}
-                                # pyrefly: ignore[missing-attribute]
-                                for k_spec, val in v.items.items():
-                                    specs[k_spec.vt.as_python_constant()] = (
-                                        val.as_python_constant()
-                                    )
-                                assert ["in_spec", "out_spec"] == list(specs.keys())
-                                self.export_metadata.module_call_spec[
-                                    # pyrefly: ignore[missing-attribute]
-                                    k.vt.as_python_constant()
-                                ] = specs
-                        # export uses tracepoint pass to dump submodule inp/out spec
-                        # into global state, so we filter it here
-                        if not (
-                            isinstance(var, UserDefinedDictVariable)
-                            and isinstance(var.value, _ExportModuleSpecTrackerDict)
-                        ):
-                            potential_side_effects.append(var)
-            side_effect_refs = [
-                _get_source_debug_name(var.source) for var in potential_side_effects
-            ]
-
+            side_effect_refs = self.get_replayed_side_effect_source_refs(
+                populate_export_metadata=True
+            )
             if side_effect_refs:
                 if torch._dynamo.config.side_effect_replay_policy == "warn":
                     warnings.warn(
@@ -2065,7 +2255,7 @@ def codegen_cells(self, tx: "InstructionTranslatorBase", cg: PyCodegen) -> None:
         # no need to codegen if reason.graph_break is False (since we won't resume)
         if self.compile_subgraph_reason.graph_break:
             tx_cnt = 0
-            cur_tx: Optional[InstructionTranslatorBase] = tx
+            cur_tx: InstructionTranslatorBase | None = tx
             while cur_tx is not None:
                 # NOTE: we generate cells in the same order as resume_execution.py: sorted freevars + cellvars
                 # Emitting `LOAD_FAST/LOAD_CLOSURE` with names in `co_freevars`
@@ -2178,12 +2368,12 @@ def bypass_package(self, reason: str = "", **kwargs: Any) -> None:
         self.package.bypass_current_entry()
         self.package = None
 
-    def get_graph_sizes_structured(self) -> dict[str, list[Union[int, str]]]:
-        ret: dict[str, list[Union[int, str]]] = {}
+    def get_graph_sizes_structured(self) -> dict[str, list[int | str]]:
+        ret: dict[str, list[int | str]] = {}
         for node in self.graph.nodes:
             example_value = node.meta.get("example_value", None)
             if isinstance(example_value, torch._subclasses.FakeTensor):
-                size = example_value.size()
+                size = example_value.shape
                 ret[node.name] = [s if isinstance(s, int) else repr(s) for s in size]
         return ret
 
@@ -2193,7 +2383,7 @@ def get_graph_sizes(self, name: str) -> str:
         for node in self.graph.nodes:
             example_value = node.meta.get("example_value", None)
             if isinstance(example_value, torch._subclasses.FakeTensor):
-                size = example_value.size()
+                size = example_value.shape
                 graph_sizes_str += f"{node.name}: {tuple(size)}\n"
                 concrete_size = []
                 has_symint = False
@@ -2295,11 +2485,80 @@ def _validate_outputs_safe_for_autograd_nodes(
             # if any node was consumed by autograd.grad
             reachable_grad_fns = collect_reachable_grad_fns([(fake_tensor, None)])
             if reachable_grad_fns & self.autograd_grad_consumed_grad_fns:
-                # Set the flag to graph break at autograd.grad on retry
-                tx.speculation_log.graph_break_on_autograd_grad = True
-                raise exc.AutogradGradRestartAnalysis(
-                    restart_reason="autograd.grad consumed grad_fns of returned tensors"
+                # Record info about the leaked tensor for the error message
+                tensor_name = str(var.source) if var.source else var.proxy.node.name
+                tx.speculation_log.autograd_grad_leaked_tensors.append(tensor_name)
+
+        if tx.speculation_log.autograd_grad_leaked_tensors:
+            # Set the flag to graph break at autograd.grad on retry
+            tx.speculation_log.graph_break_on_autograd_grad = True
+            raise exc.AutogradGradRestartAnalysis(
+                restart_reason="autograd.grad consumed grad_fns of returned tensors"
+            )
+
+    def _check_requires_grad_intermediate_outputs(
+        self, rv: list["VariableTracker"], tx: "InstructionTranslatorBase"
+    ) -> None:
+        """Skip frame if a source-less requires_grad_() intermediate leaks as output.
+
+        AOTAutograd's functionalization drops requires_grad_() on intermediates,
+        so returning them (or tensors derived from them) produces wrong results.
+        We detect this via FX graph reachability: find the requires_grad_() nodes
+        for source-less intermediates, then check if any output is downstream.
+        """
+        from .variables.tensor import TensorVariable
+
+        # Collect FX nodes for source-less requires_grad_() intermediates
+        tainted_nodes: set[torch.fx.Node] = set()
+        for v in self.leaf_var_creation_order:
+            if isinstance(v, TensorVariable) and not v.source:
+                tainted_nodes.add(v.as_proxy().node)
+
+        if not tainted_nodes:
+            return
+
+        # Propagate taint forward through the FX graph
+        for node in self.graph.nodes:
+            if node in tainted_nodes:
+                continue
+            if any(inp in tainted_nodes for inp in node.all_input_nodes):
+                tainted_nodes.add(node)
+
+        # Check leaked outputs: tainted + requires_grad means the output
+        # carries autograd state that AOTAutograd would silently drop.
+        # Detached outputs (requires_grad=False) are fine — no autograd to lose.
+        for var in rv:
+            if (
+                isinstance(var, TensorVariable)
+                and var.requires_grad
+                and var.as_proxy().node in tainted_nodes
+            ):
+                msg = (
+                    "An intermediate tensor that had requires_grad_() called "
+                    "on it (or a tensor derived from it) is being returned "
+                    "from the compiled region. AOTAutograd's functionalization "
+                    "drops the requires_grad_() effect on graph outputs, "
+                    "producing wrong results. If you only need the tensor "
+                    "values without gradients, call .detach() before returning."
                 )
+                if tx.one_graph:
+                    unimplemented(
+                        gb_type="returning intermediate with requires_grad_()",
+                        context="graph output depends on source-less requires_grad_()",
+                        explanation=msg,
+                        hints=[
+                            "If you only need the tensor values without gradients, "
+                            "call .detach() before returning.",
+                            "Consume the gradient inside the compiled function "
+                            "(call backward() and use .grad), "
+                            "or move requires_grad_() outside torch.compile.",
+                        ],
+                    )
+                else:
+                    tx.speculation_log.graph_break_on_requires_grad_ = True
+                    raise exc.RequiresGradRestartAnalysis(
+                        restart_reason="source-less requires_grad_() intermediate leaked as output"
+                    )
 
     def compile_and_call_fx_graph(
         self,
@@ -2328,6 +2587,11 @@ def compile_and_call_fx_graph(
             assert isinstance(rv, list)
             assert isinstance(root, FakeRootModule)
 
+            # Error on source-less requires_grad_() outputs.
+            # Must run before autograd validation since detaching resolves the
+            # "consumed grad_fn" conflict for backward-consumed intermediates.
+            self._check_requires_grad_intermediate_outputs(rv, tx)
+
             # Check if autograd.grad is used with outputs that require grad
             # This would cause double backward issues in aot_autograd
             self._validate_outputs_safe_for_autograd_nodes(rv, tx)
@@ -2397,23 +2661,27 @@ def compile_and_call_fx_graph(
                 # If dynamo produces a graph with parameters, skip package stuff
                 # Bypass output graph
                 self.bypass_package(
-                    "Graph contains named parameters: either inline_inbuilt_nn_modules=False or there are static addresses.",
-                    inline_builtin_nn_modules=torch._dynamo.config.inline_inbuilt_nn_modules,
+                    "Graph contains named parameters due to static addresses.",
                     gm=gm.print_readable(
                         print_output=False, include_stride=True, include_device=True
                     ),
                 )
 
             if self.package is not None:
-                gm._backend_id = name
+                gm._backend_id = name  # pyrefly: ignore[bad-argument-type]
 
+            # pyrefly: ignore[bad-argument-type]
             gm.compile_subgraph_reason = self.compile_subgraph_reason
+            # pyrefly: ignore[bad-argument-type]
             gm.meta["dynamo_flat_name_to_original_fqn"] = (
                 self.dynamo_flat_name_to_original_fqn.copy()
             )
             gm.meta["dynamo_compile_id"] = self.dynamo_compile_id
             gm.meta["backend_id"] = name
 
+            if self.cudagraph_annotation is not None:
+                gm.meta["cudagraph_annotation"] = self.cudagraph_annotation
+
             graph_code_log.debug(
                 "%s",
                 lazy_format_graph_code(
@@ -2485,6 +2753,20 @@ def compile_and_call_fx_graph(
             if self.package is not None:
                 self.package.add_backend_id(name, compiled_fn)
 
+            # If __torch_function__ subclass dispatch was inlined during
+            # tracing, wrap the compiled graph to disable __torch_function__
+            # at runtime, preventing double dispatch (the C++ dispatcher
+            # would otherwise re-trigger __torch_function__ on subclass
+            # inputs that the graph already handles).
+            if self.torch_function_subclass_inlined:
+                real_compiled_fn = compiled_fn
+
+                def _tf_disabled_wrapper(*args, **kwargs):
+                    with torch._C.DisableTorchFunctionSubclass():
+                        return real_compiled_fn(*args, **kwargs)
+
+                compiled_fn = _tf_disabled_wrapper
+
             compiled_fn = disable(
                 compiled_fn, reason="do not trace Dynamo-compiled graph"
             )
@@ -2657,6 +2939,13 @@ def _call_user_compiler(
             if hasattr(compiler_fn, "__name__")
             else "<unknown compiler_fn>"
         )
+        if config.inline_invoke_subgraph:
+            from torch._higher_order_ops.passes.inline_invoke_subgraph import (
+                inline_invoke_subgraph,
+            )
+
+            gm = inline_invoke_subgraph(gm)
+
         try:
             _step_logger()(logging.INFO, f"calling compiler function {name}")
             if config.verify_correctness:
@@ -2701,7 +2990,7 @@ def _call_user_compiler(
             },
         )
 
-        # pyrefly: ignore [unbound-name, bad-return]
+        # pyrefly: ignore [bad-return]
         return compiled_fn
 
     def dedup_pass(self) -> dict[str, torch.fx.GraphModule]:
@@ -2808,7 +3097,7 @@ def is_symnode_compute_node(node: fx.Node) -> bool:
                 ):
                     self.remove_node(node)
 
-        def placeholder_binds_symbol(node: fx.Node) -> Optional[sympy.Symbol]:
+        def placeholder_binds_symbol(node: fx.Node) -> sympy.Symbol | None:
             arg = node.meta["grapharg"]
             example = arg.example
             if isinstance(example, torch.SymInt) and isinstance(
@@ -2828,7 +3117,7 @@ def remove_unused(node: fx.Node) -> None:
         used_symbols: set[sympy.Symbol] = set()
 
         def update_used_symbols(
-            used_symbols: set[sympy.Symbol], fake: Union[torch.SymInt, torch.Tensor]
+            used_symbols: set[sympy.Symbol], fake: torch.SymInt | torch.Tensor
         ) -> None:
             used_symbols |= free_symbols(fake)
 
@@ -2925,6 +3214,34 @@ def add_output_instructions(self, prefix: list[Instruction]) -> None:
         self.output_instructions.extend(prefix)
         self.should_exit = True
 
+    def install_resume_function_global(
+        self,
+        name: str,
+        code: types.CodeType,
+        f_globals: dict[str, Any],
+    ) -> None:
+        """Install a resume function as a global.
+
+        When the code has freevars, installs a factory that creates the
+        function with correct globals and closure (since MAKE_FUNCTION
+        inherits the current frame's globals, which is wrong for resume
+        functions from inlined frames). Otherwise installs the function
+        directly.
+        """
+        if code.co_freevars:
+
+            def _make_fn(
+                closure: tuple[types.CellType, ...],
+            ) -> types.FunctionType:
+                return types.FunctionType(code, f_globals, name, None, closure)
+
+            self.install_global_unsafe(name, _make_fn)
+        else:
+            self.install_global_unsafe(
+                name,
+                types.FunctionType(code, f_globals, name),
+            )
+
     def install_global_unsafe(self, name: str, value: Any) -> None:
         """
         WARNING: prefer the safer `install_global_by_id/install_global`.
@@ -2976,8 +3293,13 @@ def cleanup(self) -> None:
                 del node.meta["grapharg"]
         self.real_value_cache.clear()
         self.input_name_to_proxy.clear()
+        self._cached_replayed_side_effect_source_refs = tuple(
+            self.get_replayed_side_effect_source_refs()
+        )
         self.side_effects.clear()
         self.variable_tracker_cache.clear()
+        self.mro_source_cache.clear()
+        self.guarded_mro_absent_keys.clear()
         self.signature_cache.clear()
         self.register_finalizer_fns.clear()
         self.dynamo_flat_name_to_original_fqn.clear()
@@ -3041,15 +3363,15 @@ def register_leaf_name(leaf_name: str) -> None:
 class DynamoTracerOutput:
     error_on_graph_break: bool
     is_tracing_resume_prologue: bool
-    output_graph: Optional[OutputGraph]
+    output_graph: OutputGraph | None
     # output_graph_for_cleanup is set even when there's an error, to allow
     # cleanup of graph nodes to break reference cycles
-    output_graph_for_cleanup: Optional[OutputGraph]
-    closure: Optional[tuple[Any, ...]]
+    output_graph_for_cleanup: OutputGraph | None
+    closure: tuple[Any, ...] | None
     f_globals: dict[str, Any]
 
     def __init__(
-        self, tracer: "InstructionTranslatorBase", error: Optional[Any] = None
+        self, tracer: "InstructionTranslatorBase", error: Any | None = None
     ) -> None:
         self.error_on_graph_break = tracer.error_on_graph_break
         self.is_tracing_resume_prologue = tracer.is_tracing_resume_prologue
@@ -3146,7 +3468,6 @@ def encountered_non_compliant_op(target: torch._ops.OpOverload, msg: str) -> Non
                 hints=[],
             )
 
-        # pyrefly: ignore [unbound-name]
         op = getattr(target, overload)
         if torch.Tag.pt2_compliant_tag in op.tags:
             encountered_compliant_op(op)
@@ -3154,7 +3475,6 @@ def encountered_non_compliant_op(target: torch._ops.OpOverload, msg: str) -> Non
             encountered_non_compliant_op(
                 op,
                 f"Encountered the torch.ops.OpOverloadPacket {target} "
-                # pyrefly: ignore [unbound-name]
                 f"which resolves to the overload ({overload}) that is "
                 f"not PT2 compliant.",
             )
@@ -3197,8 +3517,8 @@ def __init__(
         output_graph: "OutputGraph",
         parent: Optional["SubgraphTracer"] = None,
         is_export: bool = False,
-        source_target: Optional[Target] = None,
-        description: Optional[str] = None,
+        source_target: Target | None = None,
+        description: str | None = None,
     ) -> None:
         super().__init__()
         self.output_graph = weakref.proxy(output_graph)
@@ -3238,7 +3558,7 @@ def __init__(
         # 2. when we track_produced_symints for intermediate results
         # bound_symbols always map the symbol to the proxy whose
         # tracer is the current tracer that's readily accessible in current tracer's graph.
-        self.bound_symbols: dict[sympy.Symbol, Union[torch.fx.Proxy, LazyProxy]] = {}
+        self.bound_symbols: dict[sympy.Symbol, torch.fx.Proxy | LazyProxy] = {}
 
         # Maps _DynamicScalar object ids to allocated SymInt nodes, for symbol reuse
         self.dynamic_scalar_nodes: dict[int, torch.SymInt] = {}
@@ -3254,6 +3574,14 @@ def __init__(
         # This is set when enable_side_effects_in_hop=True for HOPs like invoke_subgraph
         # and checkpoint (when skip_fwd_side_effects_in_bwd_under_checkpoint config is True).
         self.allow_side_effects_in_hop = False
+        # User code stack at the point where an externally-visible side effect
+        # was first detected and allowed (not raised).  None means no side
+        # effect; non-None means one occurred.
+        self.side_effect_stack: traceback.StackSummary | None = None
+        # All sources accessed via VariableBuilder (and cell loads) during
+        # this subtracer's trace. Used by invoke_subgraph reuse to collect
+        # guards and detect mutations on captured variables.
+        self.traced_sources: OrderedSet[Source] = OrderedSet()
 
         # True if this tracer is currently tracing (reconstructing) into a Python generator
         self.is_reconstructing_generator = False
@@ -3261,9 +3589,9 @@ def __init__(
         self.debug_level: int = parent.debug_level + 1 if parent is not None else 0
 
         self._cur_code = None
-        self._orig_gm_meta: Optional[list[Any]] = None
-        self._orig_gm_lineno_map: Optional[dict[int, Optional[int]]] = None
-        self._orig_gm_firstlineno: Optional[int] = None
+        self._orig_gm_meta: list[Any] | None = None
+        self._orig_gm_lineno_map: dict[int, int | None] | None = None
+        self._orig_gm_firstlineno: int | None = None
         # Each SubgraphTracer is associated with a source target, which indicates
         # which operator this subgraph is attached to. We compute a source_fn_stack
         # based on the source target. For the root tracer, it's set to [].
@@ -3285,10 +3613,10 @@ def __init__(
                 "Inference mode is supposed to be disabled during compilation. Please open an issue."
             )
 
-        self.tracked_tensor_or_symint_vt: OrderedSet[VariableTracker] = OrderedSet()
+        self.tracked_proxyable_vt: OrderedSet[VariableTracker] = OrderedSet()
 
-    def record_tensor_or_symint_vt(self, vt: VariableTracker) -> None:
-        self.tracked_tensor_or_symint_vt.add(vt)
+    def record_proxyable_vt(self, vt: VariableTracker) -> None:
+        self.tracked_proxyable_vt.add(vt)
 
     # preserve original meta if it is available
     def _maybe_preserve_original_meta(
@@ -3319,9 +3647,9 @@ def create_proxy(
         target: Any,
         args: Any,
         kwargs: Any,
-        name: Optional[str] = None,
-        type_expr: Optional[Any] = None,
-        proxy_factory_fn: Optional[Callable[[fx.Node], fx.Proxy]] = None,
+        name: str | None = None,
+        type_expr: Any | None = None,
+        proxy_factory_fn: Callable[[fx.Node], fx.Proxy] | None = None,
     ) -> fx.Proxy:
         _t0 = time.time_ns()
         try:
@@ -3339,9 +3667,9 @@ def _create_proxy(
         target: Any,
         args: Any,
         kwargs: Any,
-        name: Optional[str] = None,
-        type_expr: Optional[Any] = None,
-        proxy_factory_fn: Optional[Callable[[fx.Node], fx.Proxy]] = None,
+        name: str | None = None,
+        type_expr: Any | None = None,
+        proxy_factory_fn: Callable[[fx.Node], fx.Proxy] | None = None,
     ) -> fx.Proxy:
         # NOTE: [Nested SubgraphTracer and free_variable handling]
         # --------------------------------------------------------
@@ -3461,7 +3789,6 @@ def get_trace_call_log_str() -> str:
             rv.node.meta["source_fn_stack"] = self.source_fn_stack + [stack]
         elif kind == "call_module":
             if self.parent is not None:
-                # TODO can remove once inline_inbuilt_nn_modules is always True
                 unimplemented(
                     gb_type="Invoking an nn.Module inside a higher order operator",
                     context=f"Higher order op name: {self.source_target}",
@@ -3495,7 +3822,6 @@ def get_trace_call_log_str() -> str:
                     ]
                 elif kind == "call_module":
                     if self.parent is not None:
-                        # TODO can remove once inline_inbuilt_nn_modules is always True
                         unimplemented(
                             gb_type="Invoking an nn.Module inside a HigherOrderOperator",
                             context="",
@@ -3547,8 +3873,8 @@ def create_node(
         target: Target,
         args: Any = None,
         kwargs: Any = None,
-        name: Optional[str] = None,
-        type_expr: Optional[Any] = None,
+        name: str | None = None,
+        type_expr: Any | None = None,
     ) -> fx.Node:
         check_pt2_compliant_op(self.output_graph, kind, target, args, kwargs)
         if self.parent is not None:
@@ -3595,7 +3921,7 @@ def create_graph_input(
         type_expr: Any,
         example_value: Any,
         before: bool = False,
-        source: Optional[Source] = None,
+        source: Source | None = None,
     ) -> fx.Proxy:
         if isinstance(example_value, torch.Tensor):
             self._input_versions_at_beginning.append(example_value._version)
@@ -3714,9 +4040,7 @@ def create_graph_input(
             return proxy
 
     # See NOTE: [Nested SubgraphTracer and free_variable handling] for more details
-    def lift_tracked_freevar_to_input(
-        self, proxy: fx.Proxy
-    ) -> Union[LazyProxy, fx.Proxy]:
+    def lift_tracked_freevar_to_input(self, proxy: fx.Proxy) -> LazyProxy | fx.Proxy:
         # You're doing something wrong if we are the root SubgraphTracer because
         # Dynamo adds tensors to graph inputs before creating a proxy for them.
         assert self.parent is not None, (
@@ -3792,7 +4116,7 @@ def maybe_lift_tracked_freevar_to_input(self, arg: Any) -> Any:
     # for symbols that're not going to be used, the LazyProxy will be turned into a proxy
     # when it's lifted as input to subgraph.
     def track_produced_symints(
-        self, example_value: Any, e_proxy: Union[LazyProxy, torch.fx.Proxy]
+        self, example_value: Any, e_proxy: LazyProxy | torch.fx.Proxy
     ) -> None:
         # When binding the symbols in an example_value, we bind the symbols
         # to the proxy's associated Tracer instead of current tracer.
@@ -3916,15 +4240,15 @@ def _proxy_with_example_value(
 
     # See Note [Auto lift basic free symbols when create_graph_input]
     def _lift_basic_symbols(
-        self, example_value: Union[torch.SymInt, torch.Tensor], src: Optional[Source]
+        self, example_value: torch.SymInt | torch.Tensor, src: Source | None
     ) -> None:
         # The before arg is for inserting symints in the sizes/strides of a tensor
         # before the tensor. This ordering ensures that when we look at the tensor's
         # symbols, they're already lifted/tracked. E.g. this assumption is used
         # in insert_deferred_runtime_asserts.
         def _lift_symbols_in_symint(
-            s: Union[int, torch.SymInt],
-            source: Optional[Source],
+            s: int | torch.SymInt,
+            source: Source | None,
             before: bool = False,
         ) -> None:
             if not is_symbolic(s):
diff --git a/torch/_dynamo/package.py b/torch/_dynamo/package.py
index f53cce2dc653f..33ce8f55b922a 100644
--- a/torch/_dynamo/package.py
+++ b/torch/_dynamo/package.py
@@ -16,6 +16,7 @@
 import hashlib
 import importlib
 import inspect
+import itertools
 import json
 import logging
 import os
@@ -78,10 +79,10 @@ class SerializedCode:
     co_firstlineno: int
     co_cellvars: tuple[str, ...]
     co_freevars: tuple[str, ...]
-    co_linetable: Optional[bytes] = None
-    co_qualname: Optional[str] = None
-    co_exceptiontable: Optional[bytes] = None
-    co_lnotab: Optional[str] = None
+    co_linetable: bytes | None = None
+    co_qualname: str | None = None
+    co_exceptiontable: bytes | None = None
+    co_lnotab: str | None = None
 
     @classmethod
     @_code_cache
@@ -214,7 +215,7 @@ class _DynamoCodeCacheEntry:
     guarded_codes: list[_GuardedCodeCacheEntry]
     import_sources: dict[str, str]
     backend_ids: list[_BackendId]
-    code_source: Optional[str]
+    code_source: str | None
     install_to_global: bool
     has_compile_id: bool = False
     bypassed: bool = False
@@ -274,11 +275,11 @@ def _get_code_source(code: types.CodeType) -> tuple[str, str]:
             if not hasattr(toplevel, part):
                 _raise_resolution_error(code, toplevel)
             toplevel = getattr(toplevel, part)
-            if inspect.isfunction(toplevel):
+            if inspect.isfunction(toplevel) or inspect.ismethod(toplevel):
                 break
     seen = set()
 
-    def _find_code_source(obj: Any) -> Optional[str]:
+    def _find_code_source(obj: Any) -> str | None:
         nonlocal toplevel
         nonlocal seen
         if obj in seen:
@@ -294,6 +295,11 @@ def _find_code_source(obj: Any) -> Optional[str]:
                 if (res := _find_code_source(const)) is not None:
                     return f".co_consts[{i}]{res}"
 
+        if inspect.ismethod(obj):
+            if (res := _find_code_source(obj.__func__)) is not None:
+                toplevel = obj
+                return f".__func__{res}"
+
         if inspect.isfunction(obj):
             if (res := _find_code_source(obj.__code__)) is not None:
                 toplevel = obj
@@ -307,6 +313,7 @@ def _find_code_source(obj: Any) -> Optional[str]:
                     if not (
                         inspect.isfunction(cell_contents)
                         or inspect.iscode(cell_contents)
+                        or inspect.ismethod(cell_contents)
                     ):
                         continue
                     if (res := _find_code_source(cell_contents)) is not None:
@@ -316,15 +323,26 @@ def _find_code_source(obj: Any) -> Optional[str]:
         if sys.version_info < (3, 11):
             if inspect.ismodule(obj):
                 for value in obj.__dict__.values():
-                    if not (inspect.isfunction(value) or inspect.isclass(value)):
+                    if not (
+                        inspect.isfunction(value)
+                        or inspect.isclass(value)
+                        or inspect.ismethod(value)
+                    ):
                         continue
                     if (res := _find_code_source(value)) is not None:
                         return res
 
             if inspect.isclass(obj):
-                for name, value in obj.__dict__.items():
-                    value = getattr(obj, name)
-                    if not (inspect.isfunction(value) or inspect.isclass(value)):
+                for name in itertools.chain(obj.__dict__.keys(), dir(obj)):
+                    try:
+                        value = getattr(obj, name)
+                    except AttributeError:
+                        continue
+                    if not (
+                        inspect.isfunction(value)
+                        or inspect.isclass(value)
+                        or inspect.ismethod(value)
+                    ):
                         continue
                     if (res := _find_code_source(value)) is not None:
                         if value.__name__ != name:
@@ -335,7 +353,6 @@ def _find_code_source(obj: Any) -> Optional[str]:
     code_source = _find_code_source(toplevel)
     if code_source is None:
         _raise_resolution_error(code, toplevel)
-    # pyrefly: ignore [missing-attribute]
     return toplevel.__qualname__, code_source.strip(".")
 
 
@@ -349,9 +366,9 @@ class SystemInfo:
 
     python_version: str
     torch_version: str
-    toolkit_version: Optional[str]
-    triton_version: Optional[tuple[int, int]]
-    gpu_name: Optional[str]
+    toolkit_version: str | None
+    triton_version: tuple[int, int] | None
+    gpu_name: str | None
     CHECK_GPUS = ("cuda", "xpu")
 
     @classmethod
@@ -426,8 +443,8 @@ class _DynamoCacheEntry:
     source_info: SourceInfo
     device_type: str
     system_info: SystemInfo = dataclasses.field(default_factory=SystemInfo.current)
-    fn_name: Optional[str] = None
-    fn_first_lineno: Optional[str] = None
+    fn_name: str | None = None
+    fn_first_lineno: str | None = None
 
     @property
     def backend_ids(self) -> set[_BackendId]:
@@ -538,6 +555,7 @@ def _compile_frame_context(
     # under the same cache entry, so we don't have recompile ids
     # i.e. If cold start had 0/0, 0/1, 1/0, 1/1, these would be
     # collapsed into 0/0, 1/0 on warm.
+    # pyrefly: ignore [deprecated]
     @contextlib.contextmanager
     def _ctx() -> Iterator[None]:
         increment_frame()
@@ -578,14 +596,14 @@ class CompilePackage:
 
     def __init__(
         self,
-        fn: Optional[Callable[..., Any]],
-        dynamo: Optional[_DynamoCacheEntry] = None,
+        fn: Callable[..., Any] | None,
+        dynamo: _DynamoCacheEntry | None = None,
         ignore_inlined_sources: bool = False,
     ) -> None:
         self._innermost_fn = None
         self._codes: dict[types.CodeType, _DynamoCodeCacheEntry] = {}
 
-        self._current_entry: Optional[_DynamoCodeCacheEntry] = None
+        self._current_entry: _DynamoCodeCacheEntry | None = None
         self._installed_globals: dict[types.ModuleType, list[str]] = {}
         # device_type that model compiled with.
         self._device_type = "cpu"
@@ -606,7 +624,7 @@ def is_initialized(self) -> bool:
     def initialize(
         self,
         fn: Any,
-        dynamo: Optional[_DynamoCacheEntry] = None,
+        dynamo: _DynamoCacheEntry | None = None,
         ignore_inlined_sources: bool = False,
     ) -> None:
         from .eval_frame import innermost_fn
@@ -627,11 +645,9 @@ def initialize(
                             f"Source code changes detected for {code.module} (line {code.firstlineno} - line {code.lastlineno})"
                         )
 
-                # pyrefly: ignore [bad-assignment]
                 self._source_info = dynamo.source_info
 
             main, *codes = dynamo.codes
-            # pyrefly: ignore [bad-assignment]
             self._codes = {self._innermost_fn.__code__: main}
             for code in codes:
                 self._codes[SerializedCode.to_code_object(code.python_code)] = code
@@ -639,15 +655,14 @@ def initialize(
             self._add_function(
                 self._innermost_fn.__code__, self._innermost_fn.__module__
             )
-        # pyrefly: ignore [bad-assignment]
         self._initialized = True
 
     def _add_function(
         self,
         python_code: types.CodeType,
         python_module: str,
-        function_name: Optional[_FunctionId] = None,
-        code_source: Optional[str] = None,
+        function_name: _FunctionId | None = None,
+        code_source: str | None = None,
         install_to_global: bool = False,
     ) -> None:
         if python_code not in self._codes:
@@ -732,7 +747,7 @@ def add_inlined_source(self, sources: list[types.CodeType]) -> None:
                 continue
             self._source_info.add_code(code)
 
-    def update_device_type(self, graph: Optional[torch.fx.Graph]) -> None:
+    def update_device_type(self, graph: torch.fx.Graph | None) -> None:
         self._device_type = _graph_device_type(graph)
 
     def bypass_current_entry(self) -> None:
@@ -743,12 +758,12 @@ def add_resume_function(
         self,
         python_code: types.CodeType,
         python_module: str,
-        function_name: Optional[str],
+        function_name: str,
     ) -> None:
         self._add_function(
             python_code,
             python_module,
-            function_name=_FunctionId(function_name) if function_name else None,
+            function_name=_FunctionId(function_name),
             install_to_global=True,
         )
         self._resume_codes.add(python_code)
@@ -757,7 +772,7 @@ def add_import_source(self, alias: str, module_name: str) -> None:
         assert self._current_entry is not None
         self._current_entry.import_sources[alias] = module_name
 
-    def add_backend_id(self, backend_id: str, backend: Optional[Any] = None) -> None:
+    def add_backend_id(self, backend_id: str, backend: Any | None = None) -> None:
         assert self._current_entry is not None
         assert backend_id.startswith("__compiled_fn_")  # sanity check
         backend_id = _BackendId(backend_id)
@@ -783,7 +798,6 @@ def uninstall(self) -> None:
             for name in names:
                 module.__dict__.pop(name)
 
-        # pyrefly: ignore [bad-assignment]
         self._installed_globals = {}
 
         _reset_precompile_entries(self._innermost_fn.__code__)
@@ -815,8 +829,29 @@ def install(self, backends: dict[_BackendId, Any]) -> None:
                 target_code = code
                 if entry.install_to_global:
                     for function_name in entry.function_names:
-                        fn = types.FunctionType(code, module.__dict__, function_name)
-                        self._install_global(module, function_name, fn)
+                        if code.co_freevars:
+                            # Resume functions with freevars need a factory
+                            # that takes a closure tuple, matching
+                            # install_resume_function_global in output_graph.py.
+                            f_globals = module.__dict__
+                            fn_name = function_name
+
+                            def _make_fn(
+                                closure: tuple[types.CellType, ...],
+                                _code: types.CodeType = code,
+                                _globals: dict[str, Any] = f_globals,
+                                _name: str = fn_name,
+                            ) -> types.FunctionType:
+                                return types.FunctionType(
+                                    _code, _globals, _name, None, closure
+                                )
+
+                            self._install_global(module, function_name, _make_fn)
+                        else:
+                            fn = types.FunctionType(
+                                code, module.__dict__, function_name
+                            )
+                            self._install_global(module, function_name, fn)
                 if entry.code_source:
                     target_code = _lookup_code(entry)
 
@@ -1126,7 +1161,7 @@ def save(self, package: CompilePackage) -> None:
         logger.info("Saving CompilePackage for %s", package.source_id)
         super().save_package(package, key)
 
-    def load(self, fn: Callable[..., Any]) -> Optional[PrecompileCacheEntry]:
+    def load(self, fn: Callable[..., Any]) -> PrecompileCacheEntry | None:
         """
         Loads a package from a given path and returns it plus a list of deserialized backends
         """
@@ -1146,9 +1181,7 @@ def load(self, fn: Callable[..., Any]) -> Optional[PrecompileCacheEntry]:
         counters["dynamo_cache"]["dynamo_cache_miss"] += 1
         return None
 
-    def load_and_install_package(
-        self, fn: Callable[..., Any]
-    ) -> Optional[CompilePackage]:
+    def load_and_install_package(self, fn: Callable[..., Any]) -> CompilePackage | None:
         """
         Load directly into a package and install backends
         """
diff --git a/torch/_dynamo/pgo.py b/torch/_dynamo/pgo.py
index 5ade3fe375dfa..80bb6ddfda1cc 100644
--- a/torch/_dynamo/pgo.py
+++ b/torch/_dynamo/pgo.py
@@ -22,7 +22,7 @@
 import re
 import zlib
 from collections import defaultdict
-from typing import Optional, TYPE_CHECKING, TypeVar, Union
+from typing import TYPE_CHECKING, TypeVar
 from typing_extensions import override, Self
 
 import torch._dynamo.config
@@ -124,6 +124,20 @@ def _hash_containing_file(filepath: str) -> str:
         return hash
 
 
+def _get_closure_content(content: types.CellType) -> object:
+    if callable(content) and hasattr(content, "__code__"):
+        return content.__code__
+    return None
+
+
+def _get_cell_hash_content(content: types.CellType) -> object:
+    # Safely extract cell contents without blowing up the entire tuple hash
+    try:
+        return _get_closure_content(content.cell_contents)
+    except ValueError:
+        return None
+
+
 @dataclasses.dataclass(frozen=True)
 class CodeId:
     filename: str
@@ -136,6 +150,7 @@ class CodeId:
     # self.filename is kept in the object to give readable information/pointer to the actual file, in a local
     # code state it will refer to the first seen file path.
     file_hash: str
+    closure_hash: int | None = None
 
     # Exclude file name.
     def __eq__(self, other: object) -> bool:
@@ -145,22 +160,30 @@ def __eq__(self, other: object) -> bool:
             self.file_hash == other.file_hash
             and self.firstlineno == other.firstlineno
             and self.name == other.name
+            and self.closure_hash == other.closure_hash
         )
 
     # Ensure if two CodeIds are the same, then they have the same hash by excluding filename.
     def __hash__(self) -> int:
-        return hash((self.file_hash, self.name, self.firstlineno))
+        return hash((self.file_hash, self.name, self.firstlineno, self.closure_hash))
 
     def __str__(self) -> str:
         return f"hash({self.file_hash}){self.filename}:{self.firstlineno}:{self.name}"
 
     @staticmethod
-    def make(code: types.CodeType) -> CodeId:
+    def make(
+        code: types.CodeType, closure: tuple[types.CellType, ...] | None = None
+    ) -> CodeId:
+        closure_hash = None
+        if closure:
+            closure_hash = hash(tuple(_get_cell_hash_content(c) for c in closure))
+
         return CodeId(
             code.co_filename,
             code.co_firstlineno,
             code.co_name,
             _hash_containing_file(code.co_filename),
+            closure_hash,
         )
 
 
@@ -171,8 +194,8 @@ class CodeState:
     )
 
 
-_INIT_CODE_STATE: Optional[defaultdict[CodeId, CodeState]] = None
-_CODE_STATE: Optional[defaultdict[CodeId, CodeState]] = None
+_INIT_CODE_STATE: defaultdict[CodeId, CodeState] | None = None
+_CODE_STATE: defaultdict[CodeId, CodeState] | None = None
 _LOGGED_DYNAMIC_ALLOWLIST: bool = False
 _KNOWN_DYNAMIC_SOURCES: set[str] = set()
 
@@ -227,19 +250,23 @@ class AutoDynamic(enum.Enum):
 
 @dataclasses.dataclass
 class FrameStateSizeEntry:
-    scalar: Union[int, AutoDynamic, AutoUnset] = dataclasses.field(default=auto_unset)
+    scalar: int | AutoDynamic | AutoUnset = dataclasses.field(default=auto_unset)
     # NB: We don't have cases where we have a known dimensionality but
     # we know NOTHING about the individual sizes
-    size: Union[AutoDynamic, AutoUnset, tuple[Union[int, AutoDynamic], ...]] = (
+    size: AutoDynamic | AutoUnset | tuple[int | AutoDynamic, ...] = dataclasses.field(
+        default=auto_unset
+    )
+    stride: AutoDynamic | AutoUnset | tuple[int | AutoDynamic | InferStride, ...] = (
         dataclasses.field(default=auto_unset)
     )
-    stride: Union[
-        AutoDynamic, AutoUnset, tuple[Union[int, AutoDynamic, InferStride], ...]
-    ] = dataclasses.field(default=auto_unset)
+    excluded_sizes: tuple[int | None, ...] | None = dataclasses.field(
+        default=None, compare=False
+    )
+    excluded_scalar: int | None = dataclasses.field(default=None, compare=False)
 
     def render(self) -> str:
         # Special cases
-        def render_single(s: Union[int, AutoDynamic, AutoUnset, InferStride]) -> str:
+        def render_single(s: int | AutoDynamic | AutoUnset | InferStride) -> str:
             if s is auto_dynamic:
                 return "?"
             elif s is auto_unset:
@@ -250,7 +277,7 @@ def render_single(s: Union[int, AutoDynamic, AutoUnset, InferStride]) -> str:
             else:
                 return str(s)
 
-        def render_tuple(ss: tuple[Union[int, AutoDynamic, InferStride], ...]) -> str:
+        def render_tuple(ss: tuple[int | AutoDynamic | InferStride, ...]) -> str:
             return "[" + ", ".join(render_single(s) for s in ss) + "]"
 
         # Common cases
@@ -309,7 +336,7 @@ def is_stride_dynamic(self, dim: int) -> bool:
         return self.stride[dim] is auto_dynamic
 
     @staticmethod
-    def _munge_symint(xs: tuple[int, ...]) -> tuple[Union[AutoDynamic, int], ...]:
+    def _munge_symint(xs: tuple[int, ...]) -> tuple[AutoDynamic | int, ...]:
         return tuple(auto_dynamic if isinstance(x, torch.SymInt) else x for x in xs)
 
     @classmethod
@@ -335,7 +362,7 @@ def make_size(cls, size: tuple[int, ...]) -> FrameStateSizeEntry:
         )
 
     @staticmethod
-    def _merge_atom(x: _T, y: _T) -> Union[AutoDynamic, _T]:
+    def _merge_atom(x: _T, y: _T) -> AutoDynamic | _T:
         if x is auto_unset:
             return y
         if y is auto_unset:
@@ -347,9 +374,9 @@ def _merge_atom(x: _T, y: _T) -> Union[AutoDynamic, _T]:
     @classmethod
     def _merge_atom_tup(
         cls,
-        xs: Union[AutoDynamic, AutoUnset, tuple[_T, ...]],
-        ys: Union[AutoDynamic, AutoUnset, tuple[_T, ...]],
-    ) -> Union[AutoDynamic, AutoUnset, tuple[Union[AutoDynamic, _T], ...]]:
+        xs: AutoDynamic | AutoUnset | tuple[_T, ...],
+        ys: AutoDynamic | AutoUnset | tuple[_T, ...],
+    ) -> AutoDynamic | AutoUnset | tuple[AutoDynamic | _T, ...]:
         if xs is auto_unset:
             return ys
         if ys is auto_unset:
@@ -361,8 +388,33 @@ def _merge_atom_tup(
         return tuple(cls._merge_atom(x, y) for x, y in zip(xs, ys))
 
     def __ior__(self, other: Self) -> Self:
+        # Record current static sizes before merge. For dims that become
+        # dynamic, the exclusion guard will reject these values so inputs
+        # fall through to the earlier, more specialized cache entry.
+        # Already-dynamic dims become None and are ignored by the guard.
+        # When no dim transitions, clear stale excluded_sizes so later
+        # compilations don't inherit exclusions from earlier transitions.
+        new_size = self._merge_atom_tup(self.size, other.size)
+        if isinstance(self.size, tuple):
+            if new_size != self.size:
+                self.excluded_sizes = tuple(
+                    s if type(s) is int else None for s in self.size
+                )
+            elif self.excluded_sizes is not None:
+                self.excluded_sizes = None
+        self.size = new_size
+        # Same idea for scalars: record the static value about to become dynamic.
+        # Re-derive like excluded_sizes: only set when transitioning from a
+        # concrete int, clear when already dynamic.
+        if (
+            type(self.scalar) is int
+            and type(other.scalar) is int
+            and self.scalar != other.scalar
+        ):
+            self.excluded_scalar = self.scalar
+        elif self.scalar is auto_dynamic and self.excluded_scalar is not None:
+            self.excluded_scalar = None
         self.scalar = self._merge_atom(self.scalar, other.scalar)
-        self.size = self._merge_atom_tup(self.size, other.size)
         self.stride = self._merge_atom_tup(self.stride, other.stride)
         return self
 
@@ -374,7 +426,7 @@ def update_automatic_dynamic(
     *,
     is_unspecialized_nn_module: bool = False,
 ) -> FrameStateSizeEntry:
-    code_id = CodeId.make(tx.f_code)
+    code_id = CodeId.make(tx.f_code, tx.closure)
     frame_state = get_code_state()[code_id]
     if torch._dynamo.config.automatic_dynamic_shapes:
         is_update = name in frame_state.automatic_dynamic
@@ -410,7 +462,7 @@ def update_automatic_dynamic(
                 )
 
         def log_tup(
-            tup_name: str, short_reason: str, long_reason: str, i: Optional[int] = None
+            tup_name: str, short_reason: str, long_reason: str, i: int | None = None
         ) -> None:
             entry_tup = (
                 getattr(entry, tup_name) if i is None else getattr(entry, tup_name)[i]
@@ -532,7 +584,7 @@ def format_cache_key(key: str) -> str:
     return f"{key}:{rank}:{tag}"
 
 
-def get_cache_key() -> Optional[str]:
+def get_cache_key() -> str | None:
     # TODO: info versions of these logs that log only once
     if torch.compiler.config.force_disable_caches:
         warn_once(
@@ -558,7 +610,7 @@ def get_cache_key() -> Optional[str]:
     return None
 
 
-def get_extra_cache_key(sticky_key: str) -> Optional[str]:
+def get_extra_cache_key(sticky_key: str) -> str | None:
     if torch.compiler.config.force_disable_caches:
         warn_once(
             "dynamo_pgo force disabled by torch.compiler.config.force_disable_caches"
@@ -569,7 +621,7 @@ def get_extra_cache_key(sticky_key: str) -> Optional[str]:
 
 
 # This solely controls local PGO
-def code_state_path(cache_key: str) -> Optional[str]:
+def code_state_path(cache_key: str) -> str | None:
     if not torch._dynamo.config.automatic_dynamic_local_pgo:
         log.debug("automatic_dynamic_local_pgo not enabled")
         return None
@@ -603,7 +655,7 @@ def should_use_remote_dynamo_pgo_cache() -> bool:
     )
 
 
-def get_remote_cache() -> Optional[RemoteCache[JsonDataTy]]:
+def get_remote_cache() -> RemoteCache[JsonDataTy] | None:
     from torch._inductor.remote_cache import create_cache
 
     if not should_use_remote_dynamo_pgo_cache():
@@ -647,7 +699,7 @@ def _collect_missing_sources(all_sources: OrderedSet[str]) -> OrderedSet[str]:
 
 def log_frame_dynamic_whitelist(f_code: types.CodeType) -> None:
     global _KNOWN_DYNAMIC_SOURCES
-    code_id = CodeId.make(f_code)
+    code_id = CodeId.make(f_code, None)
     frame_state = get_code_state()[code_id]
     all_dynamic_sources = _collect_dynamic_sources(frame_state)
     frame_whitelist = ",".join(all_dynamic_sources)
@@ -738,7 +790,7 @@ def hit(key: str, ty: str) -> defaultdict[CodeId, CodeState]:
     return _CODE_STATE
 
 
-def get_local_code_state(cache_key: str) -> Optional[defaultdict[CodeId, CodeState]]:
+def get_local_code_state(cache_key: str) -> defaultdict[CodeId, CodeState] | None:
     global _CODE_STATE
     path = code_state_path(cache_key)
     if path is not None and os.path.exists(path):
@@ -768,8 +820,8 @@ def get_local_code_state(cache_key: str) -> Optional[defaultdict[CodeId, CodeSta
 def lookup_remote_cache_entry(
     remote_cache: RemoteCache[JsonDataTy],
     cache_key: str,
-    event_name: Optional[str] = None,
-) -> Optional[defaultdict[CodeId, CodeState]]:
+    event_name: str | None = None,
+) -> defaultdict[CodeId, CodeState] | None:
     code_state = None
     try:
         cache_data = remote_cache.get(cache_key)
@@ -802,7 +854,7 @@ def lookup_remote_cache_entry(
     return code_state
 
 
-def get_remote_code_state(cache_key: str) -> Optional[defaultdict[CodeId, CodeState]]:
+def get_remote_code_state(cache_key: str) -> defaultdict[CodeId, CodeState] | None:
     global _CODE_STATE
     remote_cache = get_remote_cache()
     if remote_cache is not None:
@@ -907,7 +959,7 @@ def put_code_state() -> None:
             put_remote_code_state(extra_write_key)
 
 
-def write_local_impl(cache_key: str, pickled_code: bytes) -> Optional[tuple[str, int]]:
+def write_local_impl(cache_key: str, pickled_code: bytes) -> tuple[str, int] | None:
     path = code_state_path(cache_key)
 
     if path is None:
diff --git a/torch/_dynamo/polyfills/__init__.py b/torch/_dynamo/polyfills/__init__.py
index d3c3160166b33..c206572aee856 100644
--- a/torch/_dynamo/polyfills/__init__.py
+++ b/torch/_dynamo/polyfills/__init__.py
@@ -11,7 +11,8 @@
 from collections.abc import Callable, Hashable, Iterable, Iterator, Mapping, Sequence
 from itertools import repeat as _repeat
 from operator import eq, ne
-from typing import Any, TYPE_CHECKING, TypeVar
+from typing import Any, TYPE_CHECKING, TypeGuard, TypeVar
+from typing_extensions import TypeIs
 
 import torch
 
@@ -62,6 +63,16 @@ def __enter__(self) -> None:
         pass
 
 
+# Used by WrappedUserFunctionVariable and similar to inline decorated function
+# calls with bytecode backing. Without this, the context enter/exit happens in
+# Python-level VT code, so a nested graph break inside `fn` would skip applying
+# the context in the compiled fn/resume. By inlining through this polyfill, the
+# `with` statement has real bytecode that the resume function can continue from.
+def _fn_with_ctx(ctx: Any, fn: Callable[..., T], *args: Any, **kwargs: Any) -> T:
+    with ctx:
+        return fn(*args, **kwargs)
+
+
 def index(
     iterator: Iterator[T], item: T, start: int = 0, end: int | None = None
 ) -> int:
@@ -85,11 +96,11 @@ def radians(x: float) -> float:
     return math.pi / 180.0 * x
 
 
-def impl_IS_MAPPING(a: object) -> bool:
+def impl_IS_MAPPING(a: object) -> TypeIs[Mapping[Any, Any]]:
     return isinstance(a, Mapping)
 
 
-def impl_MATCH_SEQUENCE(a: object) -> bool:
+def impl_MATCH_SEQUENCE(a: object) -> TypeGuard[Sequence[Any]]:
     return isinstance(a, Sequence) and not isinstance(a, (str, bytes, bytearray))
 
 
@@ -316,7 +327,7 @@ def set_union(
         set_update(union_set, set2)
 
     # frozenset also uses this function
-    # pyrefly: ignore[not-callable]
+    # pyrefly: ignore [bad-argument-count, not-callable]
     return cls(union_set)
 
 
@@ -404,9 +415,12 @@ def instantiate_user_defined_class_object(
 ) -> T:
     obj = cls.__new__(cls, *args, **kwargs)
 
-    # Only call __init__ if the object is an instance of the class
+    # Only call __init__ if the object's type is a subclass of cls.
+    # CPython uses PyType_IsSubtype(Py_TYPE(obj), type) at the C level, which does NOT
+    # go through metaclass __instancecheck__. Using isinstance() here would be wrong
+    # for classes with custom __instancecheck__ (e.g. torch.ByteStorage).
     # Reference: https://github.com/python/cpython/blob/3.12/Objects/typeobject.c#L1670-L1673
-    if isinstance(obj, cls):
+    if issubclass(type(obj), cls):
         obj.__init__(*args, **kwargs)
     return obj
 
@@ -486,14 +500,45 @@ def foreach_map_fn(*args: Any) -> Any:
 
 def foreach_lerp_inplace(
     self,
-    end: list[torch.Tensor] | tuple[torch.Tensor, ...] | None,
-    weight: Sequence[bool | complex | float | int],
+    end: list[torch.Tensor] | tuple[torch.Tensor, ...],
+    weight: float | int | torch.Tensor,
 ) -> None:
-    # decompose foreach lerp into constituent ops, prevents a graph break due to
-    # converting a value to a scalar when arg[2] is a single tensor
-    result = torch._foreach_sub(end, self)
-    result = torch._foreach_mul(result, weight)
-    return torch._foreach_add_(self, result)
+    # Decompose lerp via addcmul_ for FMA.  Uses the same dual-formula
+    # approach as CUDA's native lerp to get bitwise identical results:
+    #   |w| <  0.5  (low):  fma(w, diff, start)
+    #   |w| >= 0.5  (high): fma(-(1-w), diff, end)
+    # For tensor weights (e.g. 0-dim tensor from tensor betas in Adam) the
+    # low formula is always used because the native lerp_scalar lowering
+    # would crash on float(weight) for symbolic expressions.
+    diff = torch._foreach_sub(end, self)
+    if isinstance(weight, torch.Tensor):
+        # Select base and weight for the dual formula before a single addcmul:
+        #   low  (|w| <  0.5): fma(w,      diff, self)
+        #   high (|w| >= 0.5): fma(-(1-w), diff, end)
+        mask = weight.abs() >= 0.5
+        neg_omw = -(1.0 - weight)
+        w = torch.where(mask, neg_omw, weight)
+        bases = [torch.where(mask, e, s) for s, e in zip(self, end)]
+        w_list = [w] * len(diff)
+        torch._foreach_addcmul_(bases, w_list, diff)
+        for s, b in zip(self, bases):
+            s.copy_(b)
+    else:
+        abs_weight = weight if weight >= 0 else -weight
+        if abs_weight >= 0.5:
+            # High formula: end + (-(1-w)) * diff  →  fma(-(1-w), diff, end)
+            # Compute 1-w in target dtype to match CUDA rounding.
+            d0 = self[0]
+            neg_omw = -(1.0 - torch.tensor(weight, dtype=d0.dtype, device=d0.device))
+            neg_omw_list = [neg_omw] * len(diff)
+            for s, e in zip(self, end):
+                s.copy_(e)
+            torch._foreach_addcmul_(self, neg_omw_list, diff)
+        else:
+            # Low formula: start + w * diff  →  fma(w, diff, start)
+            weights = [torch.full_like(d, weight) for d in diff]
+            torch._foreach_addcmul_(self, weights, diff)
+    return self
 
 
 def foreach_pow_scalar(
@@ -502,12 +547,6 @@ def foreach_pow_scalar(
     return torch._foreach_pow([scalar for _ in exps], exps)
 
 
-def addcmul_inplace(
-    self, tensor1: torch.Tensor, tensor2: torch.Tensor, value: Any
-) -> None:
-    return self.add_(tensor1 * tensor2 * value)
-
-
 def predicate(obj: object) -> bool:
     # This will cause the rest of dynamo to handle the if statement correctly, so we don't have to rewrite it here.
     # We can't just use bool() here since we can't trace into that in general.
@@ -525,20 +564,47 @@ def cmp_eq(a: object, b: object) -> bool:
     # slow in some corner cases.
     # if a is b:
     #     return True
-    result = a.__eq__(b)
+    if isinstance(a, type):
+        # Default metaclass equality is identity-based. Preserve the reflected
+        # operand fallback without tracing through type.__eq__.
+        if type(a).__eq__ is type.__eq__:
+            result = True if a is b else NotImplemented
+        else:
+            result = type(a).__eq__(a, b)
+    else:
+        result = a.__eq__(b)
     if result is NotImplemented:
-        result = b.__eq__(a)
+        if isinstance(b, type):
+            if type(b).__eq__ is type.__eq__:
+                result = True if a is b else NotImplemented
+            else:
+                result = type(b).__eq__(b, a)
+        else:
+            result = b.__eq__(a)
     return result is not NotImplemented and result
 
 
 def cmp_ne(a: object, b: object) -> bool:
-    # Check if __ne__ is overridden
-    if isinstance(type(a).__ne__, types.FunctionType):
+    if isinstance(a, type):
+        if type(a).__ne__ is type.__ne__:
+            result = False if a is b else NotImplemented
+        else:
+            result = type(a).__ne__(a, b)
+        if result is not NotImplemented:
+            return result
+    elif isinstance(type(a).__ne__, types.FunctionType):
         result = a.__ne__(b)
         if result is not NotImplemented:
             return result
         # Fall through to try b.__ne__(a) or cmp_eq
-    if isinstance(type(b).__ne__, types.FunctionType):
+    if isinstance(b, type):
+        if type(b).__ne__ is type.__ne__:
+            result = False if a is b else NotImplemented
+        else:
+            result = type(b).__ne__(b, a)
+        if result is not NotImplemented:
+            return result
+    elif isinstance(type(b).__ne__, types.FunctionType):
         result = b.__ne__(a)
         if result is not NotImplemented:
             return result
diff --git a/torch/_dynamo/polyfills/builtins.py b/torch/_dynamo/polyfills/builtins.py
index 0c32e8ba5d524..386a96f6db4ef 100644
--- a/torch/_dynamo/polyfills/builtins.py
+++ b/torch/_dynamo/polyfills/builtins.py
@@ -7,6 +7,7 @@
 import builtins
 import functools
 import operator
+import typing
 from collections.abc import Callable
 from typing import TYPE_CHECKING, TypeVar
 
@@ -20,6 +21,7 @@
 __all__ = [
     "all",
     "any",
+    "cast",
     "enumerate",
     "sum",
 ]
@@ -120,3 +122,8 @@ def sequence_protocol(iterable):  # type: ignore[no-untyped-def]
             raise TypeError("iter(v, w): v must be a callable")
 
         return _CallableIterator(fn, sentinel)
+
+
+@substitute_in_graph(typing.cast, can_constant_fold_through=True)
+def cast(typ: type, val: _T) -> _T:  # type: ignore[type-var]
+    return val
diff --git a/torch/_dynamo/polyfills/copy.py b/torch/_dynamo/polyfills/copy.py
new file mode 100644
index 0000000000000..0188945f038cb
--- /dev/null
+++ b/torch/_dynamo/polyfills/copy.py
@@ -0,0 +1,36 @@
+"""
+Python polyfills for copy
+"""
+
+from __future__ import annotations
+
+from typing import TypeVar
+
+from ..decorators import substitute_in_graph
+
+
+__all__ = [
+    "reduce_ex_user_defined_object",
+]
+
+T = TypeVar("T")
+
+
+@substitute_in_graph(object.__reduce_ex__, skip_signature_check=True)  # type: ignore[arg-type]
+def reduce_ex_user_defined_object(obj: T, protocol: int, /) -> tuple:  # type: ignore[type-arg]
+    """Traceable polyfill for object.__reduce_ex__ on user-defined objects.
+
+    Returns the same tuple that CPython's _common_reduce produces:
+    (copyreg.__newobj__, (cls,), obj.__dict__, None, None).
+    copy._reconstruct then calls cls.__new__(cls) and updates __dict__.
+    """
+    import copyreg
+
+    cls = type(obj)
+    return (
+        copyreg.__newobj__,  # pyrefly: ignore[missing-attribute]
+        (cls,),
+        obj.__dict__,
+        None,
+        None,
+    )
diff --git a/torch/_dynamo/polyfills/itertools.py b/torch/_dynamo/polyfills/itertools.py
index 7b72d755315b0..c21c33516db90 100644
--- a/torch/_dynamo/polyfills/itertools.py
+++ b/torch/_dynamo/polyfills/itertools.py
@@ -7,7 +7,7 @@
 import itertools
 import operator
 from collections.abc import Callable
-from typing import Optional, overload, TYPE_CHECKING, TypeAlias, TypeVar
+from typing import overload, TYPE_CHECKING, TypeAlias, TypeVar
 
 from ..decorators import substitute_in_graph
 
@@ -25,9 +25,11 @@
     "dropwhile",
     "filterfalse",
     "islice",
+    "pairwise",
+    "starmap",
+    "takewhile",
     "tee",
     "zip_longest",
-    "pairwise",
 ]
 
 
@@ -49,9 +51,9 @@ def chain(*iterables: Iterable[_T]) -> Iterator[_T]:
 @substitute_in_graph(itertools.accumulate, is_embedded_type=True)  # type: ignore[arg-type]
 def accumulate(
     iterable: Iterable[_T],
-    func: Optional[Callable[[_T, _T], _T]] = None,
+    func: Callable[[_T, _T], _T] | None = None,
     *,
-    initial: Optional[_T] = None,
+    initial: _T | None = None,
 ) -> Iterator[_T]:
     # call iter outside of the generator to match cypthon behavior
     iterator = iter(iterable)
@@ -115,6 +117,8 @@ def _cycle(iterator: Iterator[_T]) -> Iterator[_T]:
 @substitute_in_graph(itertools.dropwhile, is_embedded_type=True)  # type: ignore[arg-type]
 def dropwhile(predicate: _Predicate[_T], iterable: Iterable[_T], /) -> Iterator[_T]:
     # dropwhile(lambda x: x < 5, [1, 4, 6, 3, 8]) -> 6 3 8
+    if not callable(predicate):
+        raise TypeError(f"'{type(predicate).__name__}' object is not callable")
 
     iterator = iter(iterable)
     for x in iterator:
@@ -125,6 +129,63 @@ def dropwhile(predicate: _Predicate[_T], iterable: Iterable[_T], /) -> Iterator[
     yield from iterator
 
 
+# Reference: https://docs.python.org/3/library/itertools.html#itertools.takewhile
+@substitute_in_graph(itertools.takewhile, is_embedded_type=True)  # type: ignore[arg-type]
+def takewhile(predicate: _Predicate[_T], iterable: Iterable[_T], /) -> Iterator[_T]:
+    # takewhile(lambda x: x<5, [1,4,6,3,8]) → 1 4
+    if not callable(predicate):
+        raise TypeError(f"'{type(predicate).__name__}' object is not callable")
+
+    for x in iterable:
+        if not predicate(x):
+            break
+        yield x
+
+
+@overload
+def starmap(
+    function: Callable[[], _U],
+    iterable: Iterable[tuple[()]],
+    /,
+) -> itertools.starmap[_U]: ...
+
+
+@overload
+def starmap(
+    function: Callable[[_T], _U],
+    iterable: Iterable[tuple[_T]],
+    /,
+) -> itertools.starmap[_U]: ...
+
+
+@overload
+def starmap(
+    function: Callable[[_T, _T1], _U],
+    iterable: Iterable[tuple[_T, _T1]],
+    /,
+) -> itertools.starmap[_U]: ...
+
+
+@overload
+def starmap(
+    function: Callable[[_T, _T1, _T2], _U],
+    iterable: Iterable[tuple[_T, _T1, _T2]],
+    /,
+) -> itertools.starmap[_U]: ...
+
+
+# Reference: https://docs.python.org/3/library/itertools.html#itertools.starmap
+@substitute_in_graph(itertools.starmap, is_embedded_type=True)  # type: ignore[arg-type]
+# pyrefly: ignore [implicit-any]
+def starmap(function: Callable[..., _T], iterable: Iterable, /) -> Iterable[_T]:
+    # starmap(pow, [(2,5), (3,2), (10,3)]) → 32 9 1000
+    if not callable(function):
+        raise TypeError(f"'{type(function).__name__}' object is not callable")
+
+    for args in iterable:
+        yield function(*args)
+
+
 @substitute_in_graph(itertools.filterfalse, is_embedded_type=True)  # type: ignore[arg-type]
 def filterfalse(function: _Predicate[_T], iterable: Iterable[_T], /) -> Iterator[_T]:
     it = iter(iterable)
diff --git a/torch/_dynamo/polyfills/loader.py b/torch/_dynamo/polyfills/loader.py
index 9b24295ada3c9..cfce3edfb0ab0 100644
--- a/torch/_dynamo/polyfills/loader.py
+++ b/torch/_dynamo/polyfills/loader.py
@@ -17,6 +17,7 @@
 POLYFILLED_MODULE_NAMES: tuple[str, ...] = (
     "_collections",
     "builtins",
+    "copy",
     "functools",
     "itertools",
     "operator",
diff --git a/torch/_dynamo/polyfills/os.py b/torch/_dynamo/polyfills/os.py
index 0414bf425ab82..595b7f3636a0e 100644
--- a/torch/_dynamo/polyfills/os.py
+++ b/torch/_dynamo/polyfills/os.py
@@ -18,7 +18,6 @@
 @substitute_in_graph(os.fspath, can_constant_fold_through=True)
 def fspath(path: AnyStr | os.PathLike[AnyStr]) -> AnyStr:
     if isinstance(path, (str, bytes)):
-        # pyrefly: ignore [bad-return]
         return path
 
     path_type = type(path)
diff --git a/torch/_dynamo/polyfills/pytree.py b/torch/_dynamo/polyfills/pytree.py
index 4af119ce85cf3..d834ce1770667 100644
--- a/torch/_dynamo/polyfills/pytree.py
+++ b/torch/_dynamo/polyfills/pytree.py
@@ -1,3 +1,5 @@
+# Owner(s): ["module: pytree"]
+
 """
 Python polyfills for torch.utils.pytree
 """
@@ -230,7 +232,6 @@ def helper(treespec: PyTreeSpec) -> str:
                 or optree.is_namedtuple_class(treespec.type)
                 or optree.is_structseq_class(treespec.type)
             ):
-                # pyrefly: ignore [bad-return]
                 return treespec._unflatten_func(
                     treespec._metadata,
                     children_representations,
diff --git a/torch/_dynamo/polyfills/sys.py b/torch/_dynamo/polyfills/sys.py
index 4475b7b41ef62..d04a8b2febc64 100644
--- a/torch/_dynamo/polyfills/sys.py
+++ b/torch/_dynamo/polyfills/sys.py
@@ -26,10 +26,14 @@ def getrecursionlimit() -> int:
     return sys.getrecursionlimit()
 
 
-if hasattr(sys, "get_int_max_str_digits"):
+if sys.version_info >= (3, 11):
 
     @substitute_in_graph(sys.get_int_max_str_digits, can_constant_fold_through=True)
     def get_int_max_str_digits() -> int:
         return sys.get_int_max_str_digits()
 
-    __all__ += ["get_int_max_str_digits"]
+    @substitute_in_graph(sys.set_int_max_str_digits, can_constant_fold_through=True)
+    def set_int_max_str_digits(maxdigits: int) -> None:
+        sys.set_int_max_str_digits(maxdigits)
+
+    __all__ += ["get_int_max_str_digits", "set_int_max_str_digits"]
diff --git a/torch/_dynamo/precompile_context.py b/torch/_dynamo/precompile_context.py
index bae360041b58c..63569dc5b74ec 100644
--- a/torch/_dynamo/precompile_context.py
+++ b/torch/_dynamo/precompile_context.py
@@ -5,7 +5,7 @@
 from collections import defaultdict
 from collections.abc import Callable
 from dataclasses import dataclass
-from typing import Any, Generic, Optional, TypeVar
+from typing import Any, Generic, TypeVar
 
 import torch
 from torch._dynamo.package import (
@@ -134,7 +134,7 @@ def edit_artifact(cls, key: str, edit_fn: Callable[..., Any]) -> None:
         artifact.edit_contents(edit_fn)
 
     @classmethod
-    def serialize_artifact_by_key(cls, key: str) -> Optional[BackendCacheArtifact[Any]]:
+    def serialize_artifact_by_key(cls, key: str) -> BackendCacheArtifact[Any] | None:
         """
         Return the backend cache artifact with the associated key
         """
diff --git a/torch/_dynamo/replay_record.py b/torch/_dynamo/replay_record.py
index 5d01217fdbb61..d65f5ef4e3046 100644
--- a/torch/_dynamo/replay_record.py
+++ b/torch/_dynamo/replay_record.py
@@ -17,7 +17,7 @@
 from dataclasses import field
 from io import BufferedReader, BufferedWriter
 from types import CellType, CodeType, ModuleType
-from typing import Any, IO, Union
+from typing import Any, IO
 from typing_extensions import Self
 
 from torch.utils._import_utils import import_dill
@@ -52,12 +52,12 @@ class ExecutionRecord:
     builtins: dict[str, Any] = field(default_factory=dict)
     code_options: dict[str, Any] = field(default_factory=dict)
 
-    def dump(self, f: Union[IO[str], BufferedWriter]) -> None:
+    def dump(self, f: IO[str] | BufferedWriter) -> None:
         assert dill is not None, "replay_record requires `pip install dill`"
         dill.dump(self, f)
 
     @classmethod
-    def load(cls, f: Union[IO[bytes], BufferedReader]) -> Self:
+    def load(cls, f: IO[bytes] | BufferedReader) -> Self:
         assert dill is not None, "replay_record requires `pip install dill`"
         return dill.load(f)
 
diff --git a/torch/_dynamo/repro/after_aot.py b/torch/_dynamo/repro/after_aot.py
index 7dfb3a4cdbdb2..ca1883194feac 100644
--- a/torch/_dynamo/repro/after_aot.py
+++ b/torch/_dynamo/repro/after_aot.py
@@ -30,10 +30,11 @@
 import subprocess
 import sys
 import textwrap
+import typing
 import uuid
 from importlib import import_module
 from tempfile import TemporaryFile
-from typing import Any, IO, Optional, TYPE_CHECKING, Union
+from typing import Any, IO, TYPE_CHECKING
 from typing_extensions import Unpack
 
 import sympy
@@ -98,6 +99,39 @@ class TritonConstexpr:  # type: ignore[no-redef]
 from .. import config
 
 
+def _find_repeat_interleave_constraints(
+    gm: torch.fx.GraphModule,
+) -> list[tuple[str, str]]:
+    """
+    Find repeat_interleave operations with output_size constraints.
+
+    Returns list of (repeats_placeholder_name, output_size_placeholder_name) pairs.
+    These represent constraints where sum(repeats) must equal output_size.
+    """
+    constraints = []
+    for node in gm.graph.nodes:
+        if (
+            node.op != "call_function"
+            or "repeat_interleave" not in str(node.target)
+            or not node.args
+        ):
+            continue
+
+        output_size_node = node.kwargs.get("output_size")
+        repeats_node = node.args[0]
+
+        # Both must be FX nodes (not constants) and direct placeholders
+        if (
+            isinstance(repeats_node, torch.fx.Node)
+            and isinstance(output_size_node, torch.fx.Node)
+            and repeats_node.op == "placeholder"
+            and output_size_node.op == "placeholder"
+        ):
+            constraints.append((str(repeats_node.target), str(output_size_node.target)))
+
+    return constraints
+
+
 if TYPE_CHECKING:
     from collections.abc import Callable, Sequence
 
@@ -121,6 +155,7 @@ def _extract_distributed_info(
     Returns a dict mapping group names to dicts with 'size' and 'rank' keys.
     Example: {'tp': {'size': 4, 'rank': 0}, 'dp': {'size': 2, 'rank': 0}}
     """
+    from torch.distributed import GroupName
     from torch.fx.operator_schemas import normalize_function
 
     group_info: dict[str, dict[str, int]] = {}
@@ -143,9 +178,10 @@ def _extract_distributed_info(
             continue
         _, kwargs = opt_args_kwargs
 
-        group_name = kwargs.get("group_name")
-        if group_name is None:
+        group_name_ = kwargs.get("group_name")
+        if not isinstance(group_name_, str):
             continue
+        group_name = typing.cast(GroupName, group_name_)
 
         if group_name in group_info:
             continue
@@ -213,7 +249,7 @@ def generate_standalone_repro(
     gm: torch.fx.GraphModule,
     args: Sequence[Any],
     *,
-    save_path: Optional[str] = None,
+    save_path: str | None = None,
 ) -> str:
     """
     Generate a self-contained repro script from an FX graph.
@@ -253,11 +289,16 @@ def wrap_compiler_debug(
     def debug_wrapper(
         gm: torch.fx.GraphModule,
         example_inputs: Sequence[InputType],
+        compile_region_name: str | None = None,
         **kwargs: Unpack[_CompileFxKwargs],
     ) -> OutputCode:
         from torch._subclasses import FakeTensorMode
 
-        compiler_fn = functools.partial(unconfigured_compiler_fn, **kwargs)
+        compiler_fn = functools.partial(
+            unconfigured_compiler_fn,
+            compile_region_name=compile_region_name,
+            **kwargs,
+        )
 
         from torch._functorch.aot_autograd import get_aot_graph_name
 
@@ -442,7 +483,6 @@ def generate_custom_triton_kernel(kernel: Any) -> str:
         config_strs = []
         # pyrefly: ignore [missing-attribute]
         for kernel_config in kernel.configs:
-            # pyrefly: ignore [bad-argument-type]
             config_strs.append(f"""triton.Config(
                     {str(kernel_config.kwargs)},
                     num_warps={kernel_config.num_warps},
@@ -473,7 +513,7 @@ def generate_compiler_repro_string(
     args: Sequence[Any],
     *,
     stable_output: bool = False,
-    save_dir: Optional[str] = None,
+    save_dir: str | None = None,
     stable_hash: bool = False,
     has_distributed_ops: bool = False,
 ) -> str:
@@ -547,11 +587,8 @@ def generate_compiler_repro_string(
     )
 
     def get_fn_name(kernel: Any) -> str:
-        fn_name = (
-            # pyrefly: ignore [missing-attribute]
-            kernel._fn_name if isinstance(kernel, JITFunction) else kernel.fn._fn_name
-        )
-        return fn_name.split(".")[-1]
+        fn: Any = kernel if isinstance(kernel, JITFunction) else kernel.fn
+        return fn.__name__.split(".")[-1]
 
     def write_kernel_dependencies(
         kernel: Any,
@@ -565,7 +602,7 @@ def write_kernel_dependencies(
             return result
 
         fn_globals = getattr(jit_fn.fn, "__globals__", {})
-        src = jit_fn.src  # pyrefly: ignore [missing-attribute]
+        src = jit_fn.src
         full_src = src if src.strip().startswith("def ") else "def " + src
 
         referenced_names: set[str] = set()
@@ -633,9 +670,7 @@ def write_kernel_dependencies(
             model_str += "ERROR: Repro will not work as intended, "
             model_str += f"User defined triton kernel exception: {e}\n"
 
-    # pyrefly: ignore [unbound-name]
     if len(kernel_side_table.constant_args) > 0:
-        # pyrefly: ignore [unbound-name]
         model_str += f"{kernel_side_table_prefix}.constant_args={kernel_side_table.constant_args}\n"
 
     model_str += NNModuleToString.convert(gm)
@@ -647,22 +682,23 @@ def write_kernel_dependencies(
     # Extract from graph placeholders and their corresponding arguments
     placeholder_targets = fx_placeholder_targets(gm)
     for placeholder, arg in zip(placeholder_targets, args):
-        # pyrefly: ignore [unbound-name]
         if isinstance(arg, (int, torch.SymInt)):
             writer.symint(placeholder, arg)
-        # pyrefly: ignore [unbound-name]
         elif isinstance(arg, torch.Tensor):
             # TODO: improve these names with FQN
             writer.tensor(placeholder, arg)
         elif arg is None:
             writer.const(placeholder)
+        elif isinstance(arg, FakeScriptObject):
+            writer.opaque(placeholder, arg.script_class_name)
+        elif isinstance(arg, torch._C.Generator):
+            writer.generator(placeholder, arg)
         else:
             writer.unsupported(placeholder, arg)
 
         # Extract symbolic variables from the same arguments
 
         if (
-            # pyrefly: ignore [unbound-name]
             isinstance(arg, torch.SymInt)
             # By checking sympy.Symbol, we are excluding any symbolic expressions.
             # TODO: we may need to solve expressions to extract symbol definitions.
@@ -670,12 +706,10 @@ def write_kernel_dependencies(
             and arg.node.hint is not None
         ):
             used_syms[str(arg.node)] = arg.node.hint
-        # pyrefly: ignore [unbound-name]
         elif isinstance(arg, torch.Tensor):
             # Extract symbolic variables from tensor shapes and strides
             for dim in arg.shape:
                 if (
-                    # pyrefly: ignore [unbound-name]
                     isinstance(dim, torch.SymInt)
                     and isinstance(dim.node.expr, sympy.Symbol)
                     and dim.node.hint is not None
@@ -683,7 +717,6 @@ def write_kernel_dependencies(
                     used_syms[str(dim.node)] = dim.node.hint
             for stride in arg.stride():
                 if (
-                    # pyrefly: ignore [unbound-name]
                     isinstance(stride, torch.SymInt)
                     and isinstance(stride.node.expr, sympy.Symbol)
                     and stride.node.hint is not None
@@ -692,7 +725,6 @@ def write_kernel_dependencies(
             # Extract symbols from storage nbytes (can be a symbolic expression)
             storage = arg.untyped_storage()
             nbytes = storage.nbytes()
-            # pyrefly: ignore [unbound-name]
             if isinstance(nbytes, torch.SymInt):
                 expr = nbytes.node.expr
                 shape_env = nbytes.node.shape_env
@@ -709,6 +741,32 @@ def write_kernel_dependencies(
         )
         model_str = f"{hint_lines}\n\n{model_str}"
 
+    # Add fixup code for repeat_interleave constraints
+    # When inputs are regenerated randomly, sum(repeats) != output_size
+    # This fixup adjusts the repeats tensor to satisfy the constraint
+    constraints = _find_repeat_interleave_constraints(gm)
+    if constraints:
+        placeholder_to_idx = {name: idx for idx, name in enumerate(placeholder_targets)}
+        for repeats_name, output_size_name in constraints:
+            repeats_idx = placeholder_to_idx.get(repeats_name)
+            output_size_idx = placeholder_to_idx.get(output_size_name)
+            if repeats_idx is not None and output_size_idx is not None:
+                # Guard with hasattr since NopInputReader doesn't have args
+                writer._lines.append(
+                    "# Fixup: ensure sum(repeats) == output_size for repeat_interleave"
+                )
+                writer._lines.append("if hasattr(reader, 'args'):")
+                writer._lines.append(f"    _repeats = reader.args[{repeats_idx}]")
+                writer._lines.append(
+                    f"    _output_size = reader.args[{output_size_idx}]"
+                )
+                writer._lines.append(
+                    "    if isinstance(_repeats, torch.Tensor) and _repeats.dtype == torch.int64:"
+                )
+                writer._lines.append("        _n = _repeats.numel()")
+                writer._lines.append("        _repeats.fill_(_output_size // _n)")
+                writer._lines.append("        _repeats[:_output_size % _n] += 1")
+
     load_args_lines = writer.lines()
     load_args_code = "\n".join(load_args_lines)
     model_str += load_args_code + "\n"
@@ -724,11 +782,11 @@ def save_graph_repro(
     compiler_name: str,
     *,
     stable_output: bool = False,
-    save_dir: Optional[str] = None,
+    save_dir: str | None = None,
     command: str = "run",
-    accuracy: Optional[Union[str, bool]] = None,
-    tracing_mode: Optional[str] = None,
-    check_str: Optional[str] = None,
+    accuracy: str | bool | None = None,
+    tracing_mode: str | None = None,
+    check_str: str | None = None,
     stable_hash: bool = False,
 ) -> None:
     if any(
@@ -795,7 +853,7 @@ def dump_compiler_graph_state(
     args: Sequence[Any],
     compiler_name: str,
     *,
-    accuracy: Optional[Union[str, bool]] = None,
+    accuracy: str | bool | None = None,
 ) -> None:
     subdir = os.path.join(minifier_dir(), "checkpoints")
     if not os.path.exists(subdir):
@@ -840,11 +898,11 @@ def isolate_fails(
     fx_g: torch.fx.GraphModule,
     args: Sequence[Any],
     compiler_name: str,
-    env: Optional[dict[str, Any]] = None,
-    save_dir: Optional[str] = None,
-    accuracy: Optional[Union[bool, str]] = None,
-    tracing_mode: Optional[str] = None,
-    check_str: Optional[str] = None,
+    env: dict[str, Any] | None = None,
+    save_dir: str | None = None,
+    accuracy: bool | str | None = None,
+    tracing_mode: str | None = None,
+    check_str: str | None = None,
 ) -> bool:
     if env is None:
         env = {}
@@ -905,7 +963,7 @@ def isolate_fails(
 
 
 def inductor_fails(
-    fx_g: torch.fx.GraphModule, args: Sequence[Any], check_str: Optional[str] = None
+    fx_g: torch.fx.GraphModule, args: Sequence[Any], check_str: str | None = None
 ) -> bool:
     has_cuda = False
     for arg in args:
@@ -945,7 +1003,7 @@ def sync() -> None:
 def inductor_accuracy_fails(
     fx_g: torch.fx.GraphModule,
     args: Sequence[Any],
-    check_str: Optional[str] = None,
+    check_str: str | None = None,
     *,
     require_fp64: bool = False,
     ignore_non_fp: bool = False,
@@ -971,7 +1029,7 @@ def inductor_accuracy_fails(
 
 def repro_common(
     options: Any, mod: nn.Module, load_args: Any
-) -> tuple[torch.fx.GraphModule, Sequence[Any]]:
+) -> tuple[torch.fx.GraphModule, list[Any]]:
     # Invariant for graphs we generate with the repro script
     assert not any(mod.named_parameters())
     for n, b in mod.named_buffers():
@@ -1007,7 +1065,6 @@ def repro_common(
 
     # Turn mod into a GraphModule the slow way
     # TODO: speed this up
-    # pyrefly: ignore [bad-argument-type]
     mod = make_fx(mod, tracing_mode=options.tracing_mode)(*args)
 
     # pyrefly: ignore [bad-assignment]
@@ -1116,7 +1173,7 @@ def save_hook(name: str, val: Any) -> None:
         compiled(new_args)  # type: ignore[arg-type]
         assert not new_args
 
-    def compare_tuples(tuple1: tuple[Any], tuple2: tuple[Any]) -> Optional[str]:
+    def compare_tuples(tuple1: tuple[Any], tuple2: tuple[Any]) -> str | None:
         diff_indices = [i for i in range(len(tuple1)) if tuple1[i] != tuple2[i]]
         diff_values = [(tuple1[i], tuple2[i]) for i in diff_indices]
 
@@ -1221,7 +1278,7 @@ def repro_get_args(
     options: Any, mod: nn.Module, load_args: Any
 ) -> tuple[torch.fx.GraphModule, list[Any]]:
     mod, args = repro_common(options, mod, load_args)
-    return mod, args  # type: ignore[return-value]
+    return mod, args
 
 
 def repro_run(options: Any, mod: nn.Module, load_args: Any) -> None:
@@ -1265,11 +1322,11 @@ def run_repro(
     load_args: Any,
     *,
     command: str = "run",
-    accuracy: Union[bool, str] = "",
-    save_dir: Optional[str] = None,
-    tracing_mode: Optional[str] = None,
-    patch_code: Optional[str] = None,
-    check_str: Optional[str] = None,
+    accuracy: bool | str = "",
+    save_dir: str | None = None,
+    tracing_mode: str | None = None,
+    patch_code: str | None = None,
+    check_str: str | None = None,
     **kwargs: Any,
 ) -> Any:
     for k in kwargs:
diff --git a/torch/_dynamo/repro/after_dynamo.py b/torch/_dynamo/repro/after_dynamo.py
index 5f4f29829647b..95c2ac5ac2aa6 100644
--- a/torch/_dynamo/repro/after_dynamo.py
+++ b/torch/_dynamo/repro/after_dynamo.py
@@ -26,7 +26,7 @@
 import textwrap
 from collections.abc import Callable, Sequence
 from importlib import import_module
-from typing import Any, Optional, Union
+from typing import Any
 
 import torch
 import torch.fx as fx
@@ -82,7 +82,7 @@ def _accuracy_fails(
 
 class WrapBackendDebug:
     def __init__(
-        self, unconfigured_compiler_fn: CompilerFn, compiler_name: Optional[str]
+        self, unconfigured_compiler_fn: CompilerFn, compiler_name: str | None
     ) -> None:
         functools.wraps(unconfigured_compiler_fn)(self)
         self._torchdynamo_orig_backend = unconfigured_compiler_fn
@@ -159,7 +159,7 @@ def add_paths(exc: Exception) -> None:
 
 
 def wrap_backend_debug(
-    unconfigured_compiler_fn: CompilerFn, compiler_name: Optional[str]
+    unconfigured_compiler_fn: CompilerFn, compiler_name: str | None
 ) -> WrapBackendDebug:
     """
     A minifier decorator that wraps the TorchDynamo produced Fx graph modules.
@@ -180,11 +180,11 @@ def wrap_backend_debug(
 def generate_dynamo_fx_repro_string(
     gm: torch.fx.GraphModule,
     args: Sequence[Any],
-    compiler_name: Optional[str],
+    compiler_name: str | None,
     check_accuracy: bool = False,
     *,
     stable_output: bool = False,
-    save_dir: Optional[str] = None,
+    save_dir: str | None = None,
     command: str = "run",
 ) -> str:
     """
@@ -236,7 +236,7 @@ def generate_dynamo_fx_repro_string(
 def dump_backend_repro_as_file(
     gm: torch.fx.GraphModule,
     args: Sequence[Any],
-    compiler_name: Optional[str],
+    compiler_name: str | None,
     check_accuracy: bool = False,
 ) -> None:
     """
@@ -269,7 +269,7 @@ def dump_backend_repro_as_file(
 def dump_backend_state(
     gm: torch.fx.GraphModule,
     args: Sequence[Any],
-    compiler_name: Optional[str],
+    compiler_name: str | None,
     check_accuracy: bool = False,
 ) -> None:
     """
@@ -290,7 +290,7 @@ def dump_backend_state(
 
 
 def dump_to_minify_after_dynamo(
-    gm: torch.fx.GraphModule, args: Sequence[Any], compiler_name: Optional[str]
+    gm: torch.fx.GraphModule, args: Sequence[Any], compiler_name: str | None
 ) -> None:
     # TODO: factor this out
     subdir = os.path.join(minifier_dir(), "checkpoints")
@@ -315,7 +315,7 @@ def dump_to_minify_after_dynamo(
 
 @register_debug_backend  # type: ignore[arg-type]
 def dynamo_minifier_backend(
-    gm: fx.GraphModule, example_inputs: Sequence[Any], compiler_name: Optional[str]
+    gm: fx.GraphModule, example_inputs: Sequence[Any], compiler_name: str | None
 ) -> fx.GraphModule:
     from functorch.compile import minifier
 
@@ -357,7 +357,7 @@ def dynamo_minifier_backend(
 
 @register_debug_backend  # type: ignore[arg-type]
 def dynamo_accuracy_minifier_backend(
-    gm: fx.GraphModule, example_inputs: Sequence[Any], compiler_name: Optional[str]
+    gm: fx.GraphModule, example_inputs: Sequence[Any], compiler_name: str | None
 ) -> fx.GraphModule:
     from functorch.compile import minifier
 
@@ -471,7 +471,6 @@ def repro_minify(options: Any, mod: torch.nn.Module, load_args: Any) -> None:
         compiler_fn,
         compiler_name=options.backend,  # type: ignore[call-arg]
     )
-    # pyrefly: ignore [bad-argument-type]
     opt_mod = torch._dynamo.optimize(dynamo_minifier_backend)(mod)
 
     with torch.amp.autocast("cuda", enabled=options.autocast):
@@ -479,7 +478,6 @@ def repro_minify(options: Any, mod: torch.nn.Module, load_args: Any) -> None:
 
 
 def repro_run(options: Any, mod: torch.nn.Module, load_args: Any) -> None:
-    # pyrefly: ignore [bad-argument-type]
     opt_mod = torch._dynamo.optimize(options.backend)(mod)
 
     if options.accuracy != "":
@@ -518,8 +516,8 @@ def run_repro(
     load_args: Any,
     *,
     command: str = "run",
-    accuracy: Union[bool, str] = "",
-    save_dir: Optional[str] = None,
+    accuracy: bool | str = "",
+    save_dir: str | None = None,
     autocast: bool = False,
     backend: str = "inductor",
     **kwargs: Any,
diff --git a/torch/_dynamo/repro/aoti.py b/torch/_dynamo/repro/aoti.py
index 35eb0bd18beec..156055509e3ad 100644
--- a/torch/_dynamo/repro/aoti.py
+++ b/torch/_dynamo/repro/aoti.py
@@ -26,7 +26,7 @@
 import textwrap
 from collections.abc import Sequence
 from importlib import import_module
-from typing import Any, IO, Optional, Union
+from typing import Any, IO
 
 import torch
 from torch._dynamo.debug_utils import (
@@ -53,7 +53,7 @@
 
 
 class AOTIMinifierError(Exception):
-    def __init__(self, original_exception: Union[str, Exception]) -> None:
+    def __init__(self, original_exception: str | Exception) -> None:
         additional_message = "This error is caused by a bug in the AOTI minifier, please report a bug to PyTorch"
         full_message = f"{additional_message}: {str(original_exception)}"
         super().__init__(full_message)
@@ -64,7 +64,7 @@ def dump_to_minify(
     exported_program: ExportedProgram,
     compiler_name: str,
     command: str = "minify",
-    options: Optional[dict[str, Any]] = None,
+    options: dict[str, Any] | None = None,
 ) -> None:
     """
     If command is "minify":
@@ -134,15 +134,15 @@ def save_graph_repro_ep(
     fd: IO[Any],
     compiler_name: str,
     *,
-    exported_program: Optional[ExportedProgram] = None,
-    gm: Optional[torch.nn.Module] = None,
-    args: Optional[tuple[Any]] = None,
-    config_patches: Optional[dict[str, str]] = None,
+    exported_program: ExportedProgram | None = None,
+    gm: torch.nn.Module | None = None,
+    args: tuple[Any] | None = None,
+    config_patches: dict[str, str] | None = None,
     stable_output: bool = False,
-    save_dir: Optional[str] = None,
+    save_dir: str | None = None,
     command: str = "run",
-    accuracy: Optional[Union[str, bool]] = None,
-    check_str: Optional[str] = None,
+    accuracy: str | bool | None = None,
+    check_str: str | None = None,
     module_in_comment: bool = False,
     strict: bool = False,
 ) -> None:
@@ -193,8 +193,8 @@ def dump_compiler_graph_state(
     args: Sequence[Any],
     compiler_name: str,
     *,
-    config_patches: Optional[dict[str, str]] = None,
-    accuracy: Optional[Union[str, bool]] = None,
+    config_patches: dict[str, str] | None = None,
+    accuracy: str | bool | None = None,
     strict: bool = False,
 ) -> None:
     subdir = os.path.join(minifier_dir(), "checkpoints")
@@ -235,9 +235,9 @@ def dump_compiler_graph_state(
 def generate_compiler_repro_exported_program(
     exported_program: ExportedProgram,
     *,
-    options: Optional[dict[str, str]] = None,
+    options: dict[str, str] | None = None,
     stable_output: bool = False,
-    save_dir: Optional[str] = None,
+    save_dir: str | None = None,
 ) -> str:
     model_str = textwrap.dedent(
         f"""
@@ -272,7 +272,7 @@ def generate_compiler_repro_exported_program(
     return model_str
 
 
-def repro_load_args(load_args: Any, save_dir: Optional[str]) -> tuple[Any]:
+def repro_load_args(load_args: Any, save_dir: str | None) -> tuple[Any]:
     if not hasattr(load_args, "_version"):
         log.warning(
             "load_args does not have a _version attribute, please file a bug to PyTorch "
@@ -311,7 +311,7 @@ def repro_common(
 def repro_get_args(
     options: Any,
     exported_program: ExportedProgram,
-    config_patches: Optional[dict[str, Any]],
+    config_patches: dict[str, Any] | None,
 ) -> tuple[torch.fx.GraphModule, Any, Any]:
     mod, args, kwargs = repro_common(options, exported_program)
     return mod, args, kwargs
@@ -320,7 +320,7 @@ def repro_get_args(
 def repro_run(
     options: Any,
     exported_program: ExportedProgram,
-    config_patches: Optional[dict[str, Any]],
+    config_patches: dict[str, Any] | None,
 ) -> None:
     from torch._inductor import _aoti_compile_and_package_inner
 
@@ -353,7 +353,7 @@ def export_for_aoti_minifier(
     tuple_inputs: tuple[Any],
     strict: bool = False,
     skip_export_error: bool = True,
-) -> Optional[torch.nn.Module]:
+) -> torch.nn.Module | None:
     # Some graphs cannot be used for AOTI/export (illegal graphs), these should be
     # considered as graphs that don't fail in the minifier, so the minifier keeps searching.
     # In these case, we return None. Otherwise, we return the exported graph module.
@@ -391,7 +391,7 @@ def export_for_aoti_minifier(
 def repro_minify(
     options: Any,
     exported_program: ExportedProgram,
-    config_patches: Optional[dict[str, Any]],
+    config_patches: dict[str, Any] | None,
 ) -> None:
     from functorch.compile import minifier
     from torch._inductor import _aoti_compile_and_package_inner
@@ -420,7 +420,7 @@ def repro_minify(
     def module_fails(
         gm: torch.fx.GraphModule,
         flat_example_inputs: list[Any],
-        check_str: Optional[str] = None,
+        check_str: str | None = None,
     ) -> bool:
         # Need to export first so the in_spec and out_spec are populated
         tuple_inputs = tuple(flat_example_inputs)
@@ -474,12 +474,12 @@ def module_fails(
 def run_repro(
     exported_program: ExportedProgram,
     *,
-    config_patches: Optional[dict[str, str]] = None,
+    config_patches: dict[str, str] | None = None,
     command: str = "run",
-    accuracy: Union[bool, str] = "",
-    save_dir: Optional[str] = None,
-    tracing_mode: Optional[str] = None,
-    check_str: Optional[str] = None,
+    accuracy: bool | str = "",
+    save_dir: str | None = None,
+    tracing_mode: str | None = None,
+    check_str: str | None = None,
     minifier_export_mode: str = "python",
     skip_export_error: bool = True,
     **more_kwargs: Any,
diff --git a/torch/_dynamo/resume_execution.py b/torch/_dynamo/resume_execution.py
index 5be66b7ab3a77..cc3d0068e3152 100644
--- a/torch/_dynamo/resume_execution.py
+++ b/torch/_dynamo/resume_execution.py
@@ -19,7 +19,7 @@
 import types
 from collections.abc import Callable, Iterable
 from contextlib import AbstractContextManager
-from typing import Any, cast, Optional
+from typing import Any, cast
 
 from .bytecode_transformation import (
     add_push_null,
@@ -67,7 +67,7 @@ def _initial_push_null(insts: list[Instruction]) -> None:
 def _bytecode_from_template_with_split(
     template: Callable[..., Any],
     stack_index: int,
-    varname_map: Optional[dict[str, Any]] = None,
+    varname_map: dict[str, Any] | None = None,
 ) -> tuple[list[Instruction], list[Instruction]]:
     template_code = bytecode_from_template(template, varname_map=varname_map)
     template_code.append(create_instruction("POP_TOP"))
@@ -116,7 +116,7 @@ def _try_except_tf_mode_template(dummy: Any, stack_var_name: Any) -> None:
 @dataclasses.dataclass(frozen=True)
 class ReenterWith:
     stack_index: int
-    target_values: Optional[tuple[Any, ...]] = None
+    target_values: tuple[Any, ...] | None = None
 
     def try_except_torch_function_mode(
         self, code_options: dict[str, Any], cleanup: list[Instruction]
@@ -191,7 +191,7 @@ def _template(ctx: AbstractContextManager[Any], dummy: Any) -> None:
 
     def __call__(
         self, code_options: dict[str, Any], cleanup: list[Instruction]
-    ) -> tuple[list[Instruction], Optional[Instruction]]:
+    ) -> tuple[list[Instruction], Instruction | None]:
         """
         Codegen based off of:
         with ctx(args):
diff --git a/torch/_dynamo/side_effects.py b/torch/_dynamo/side_effects.py
index bd6c6669cceba..0d7d0eb325bdc 100644
--- a/torch/_dynamo/side_effects.py
+++ b/torch/_dynamo/side_effects.py
@@ -24,17 +24,20 @@
 import collections
 import contextlib
 import inspect
+import logging
 import textwrap
 import traceback
 import warnings
 import weakref
 from collections.abc import Generator, MutableMapping
 from types import CellType
-from typing import Any, Optional, TYPE_CHECKING
+from typing import Any, TYPE_CHECKING
 
 import torch
 import torch.nn
 from torch._dynamo.variables.misc import AutogradFunctionContextVariable
+from torch.utils._ordered_set import OrderedSet
+from torch.utils._pytree import is_structseq_class
 
 from . import config, graph_break_hints, utils, variables
 from .bytecode_transformation import (
@@ -45,8 +48,8 @@
 )
 from .codegen import PyCodegen
 from .exc import collapse_resume_frames, get_stack_above_dynamo, unimplemented
-from .source import GlobalSource, LocalCellSource, Source, TempLocalSource
-from .utils import is_frozen_dataclass, nn_module_new, object_new
+from .source import AttrSource, GlobalSource, LocalCellSource, Source, TempLocalSource
+from .utils import is_frozen_dataclass, is_namedtuple_cls, nn_module_new, object_new
 from .variables.base import (
     AttributeMutation,
     AttributeMutationExisting,
@@ -115,27 +118,26 @@ class SideEffects:
     def __init__(
         self,
         output_graph: "OutputGraph",
-        id_to_variable: Optional[dict[int, VariableTracker]] = None,
-        store_attr_mutations: Optional[
-            dict[VariableTracker, dict[str, VariableTracker]]
-        ] = None,
+        id_to_variable: dict[int, VariableTracker] | None = None,
+        store_attr_mutations: dict[VariableTracker, dict[str, VariableTracker]]
+        | None = None,
         mutation_user_stacks: dict[VariableTracker, list[traceback.StackSummary]]
         | None = None,
-        keepalive: Optional[list[Any]] = None,
-        save_for_backward: Optional[
-            list[tuple[AutogradFunctionContextVariable, list[VariableTracker]]]
-        ] = None,
-        tensor_hooks: Optional[
-            dict[
-                int,
-                tuple[
-                    "variables.TensorVariable",
-                    VariableTracker,
-                    "variables.RemovableHandleVariable",
-                    str,
-                ],
-            ]
-        ] = None,
+        keepalive: list[Any] | None = None,
+        save_for_backward: list[
+            tuple[AutogradFunctionContextVariable, list[VariableTracker]]
+        ]
+        | None = None,
+        tensor_hooks: dict[
+            int,
+            tuple[
+                "variables.TensorVariable",
+                VariableTracker,
+                "variables.RemovableHandleVariable",
+                str,
+            ],
+        ]
+        | None = None,
     ) -> None:
         super().__init__()
         self.output_graph_weakref = weakref.ref(output_graph)
@@ -150,7 +152,7 @@ def __init__(
         self._has_existing_dict_mutation = False
         # Track Compiled Autograd final callbacks that must be called at the end of Compiled Autograd backward graph.
         # Only applicable if this graph is created from Dynamo tracing in Compiled Autograd.
-        self.ca_final_callbacks_var: Optional[ListVariable] = None
+        self.ca_final_callbacks_var: ListVariable | None = None
 
         # Tracks VariableTracker objects whose mutations can be skipped.
         # For normal mutated variables, Dynamo generates code to replay/reconstruct
@@ -160,6 +162,9 @@ def __init__(
         # Used for temporary mutations in contexts like torch.func.functional_call,
         # where module parameters/buffers are modified but later restored.
         self.ignore_mutation_on_these_variables: set[VariableTracker] = set()
+        # Sources mutated during tracing: AttrSource for attribute
+        # mutations, var.source for value mutations (list/dict/etc).
+        self.mutated_sources: OrderedSet[Source] = OrderedSet()
 
     def ignore_mutations_on(self, var: VariableTracker) -> None:
         """Mutations to this variable will be executed but not not tracked,
@@ -191,7 +196,7 @@ def __eq__(self, other: object) -> bool:
             and self.tensor_hooks == other.tensor_hooks
         )
 
-    def diff(self, other: "SideEffects") -> Optional[str]:
+    def diff(self, other: "SideEffects") -> str | None:
         if self.id_to_variable != other.id_to_variable:
             sk_itv = self.id_to_variable.keys()
             ok_itv = other.id_to_variable.keys()
@@ -257,6 +262,19 @@ def is_reconstructing_generator(self) -> bool:
             and output_graph.current_tx.output.current_tracer.is_reconstructing_generator
         )
 
+    def _maybe_record_side_effect(self, item: VariableTracker) -> None:
+        """Record the first externally-visible side effect on the current tracer."""
+        if item.mutation_type is not None and not is_side_effect_safe(
+            item.mutation_type
+        ):
+            output_graph = self.output_graph_weakref()
+            if output_graph:
+                tracer = output_graph.current_tx.output.current_tracer
+                if tracer.side_effect_stack is None:
+                    tracer.side_effect_stack = (
+                        torch._guards.TracingContext.extract_stack()
+                    )
+
     def check_allowed_side_effect(self, item: VariableTracker) -> bool:
         from torch._dynamo.variables.misc import AutogradFunctionContextVariable
 
@@ -265,8 +283,10 @@ def check_allowed_side_effect(self, item: VariableTracker) -> bool:
         if isinstance(item, AutogradFunctionContextVariable):
             return True
         if self.should_allow_externally_visible_side_effects_in_subtracer():
+            self._maybe_record_side_effect(item)
             return True
         if self.should_allow_side_effects_in_hop():
+            self._maybe_record_side_effect(item)
             return True
         if self.is_reconstructing_generator():
             # This is missing the case where one mutates a tensor. See
@@ -307,6 +327,9 @@ def store_attr(
         self.store_attr_mutations[item][name] = value
         # Capture user stack for this mutation
         self._capture_user_stack(item)
+        item_source = getattr(item, "source", None)
+        if item_source is not None:
+            self.mutated_sources.add(AttrSource(item_source, name))
 
     def load_attr(
         self,
@@ -341,6 +364,16 @@ def store_cell(self, cellvar: VariableTracker, value: VariableTracker) -> None:
 
     def load_cell(self, cellvar: VariableTracker) -> VariableTracker:
         assert isinstance(cellvar, variables.CellVariable)
+        # Track the cell_contents source during subgraph tracing so that
+        # mutations (e.g. nonlocal counter = 3) are detected by the reuse
+        # mechanism via set intersection with mutated_sources.
+        output_graph = self.output_graph_weakref()
+        if output_graph:
+            cell_source = getattr(cellvar, "source", None)
+            if cell_source is not None:
+                output_graph.current_tx.output.current_tracer.traced_sources.add(
+                    AttrSource(cell_source, "cell_contents")
+                )
         if self.has_pending_mutation_of_attr(cellvar, "cell_contents"):
             return self.load_attr(cellvar, "cell_contents", check=False)
         if cellvar.pre_existing_contents:
@@ -398,9 +431,7 @@ def is_modified(self, item: VariableTracker) -> bool:
 
         if isinstance(item, variables.UserDefinedObjectVariable):
             # Checks if the underlying dict or tuple vt has been modified
-            return item in self.store_attr_mutations or item.is_underlying_vt_modified(
-                self
-            )
+            return item in self.store_attr_mutations or item.is_base_vt_modified(self)
 
         if self.is_attribute_mutation(item):
             return item in self.store_attr_mutations
@@ -491,7 +522,10 @@ def get_variable_cls(self, user_cls: type) -> type:
         elif issubclass(user_cls, (set, frozenset)):
             variable_cls = variables.UserDefinedSetVariable
         elif issubclass(user_cls, tuple):
-            variable_cls = variables.UserDefinedTupleVariable
+            if is_namedtuple_cls(user_cls):
+                variable_cls = variables.UserDefinedTupleVariable.get_vt_cls(user_cls)
+            else:
+                variable_cls = variables.UserDefinedTupleVariable
         elif issubclass(user_cls, list):
             variable_cls = variables.UserDefinedListVariable
         elif issubclass(user_cls, MutableMapping):
@@ -500,6 +534,11 @@ def get_variable_cls(self, user_cls: type) -> type:
             variable_cls = FrozenDataClassVariable
         elif issubclass(user_cls, BaseException):
             variable_cls = variables.UserDefinedExceptionObjectVariable
+        elif issubclass(
+            user_cls,
+            variables.user_defined._CONSTANT_BASE_TYPES,
+        ):
+            variable_cls = variables.UserDefinedConstantVariable
         elif variables.InspectVariable.is_matching_class(user_cls):
             variable_cls = variables.InspectVariable
         assert issubclass(variable_cls, variables.UserDefinedObjectVariable)
@@ -518,6 +557,10 @@ def get_example_value(
         else:
             if isinstance(base_cls_vt, variables.BuiltinVariable):
                 base_cls = base_cls_vt.fn
+            elif isinstance(base_cls_vt, variables.DictBuiltinVariable):
+                base_cls = dict
+            elif isinstance(base_cls_vt, variables.ListBuiltinVariable):
+                base_cls = list
             elif isinstance(base_cls_vt, variables.UserDefinedClassVariable):
                 base_cls = base_cls_vt.value
             else:
@@ -526,14 +569,28 @@ def get_example_value(
             assert variables.UserDefinedClassVariable.is_supported_new_method(
                 base_cls.__new__
             )
-            # TODO(anijain2305) - Consider adding get_example_value method to
-            # each VT to get an example value for all args. As we expand the
-            # scope to other __new__ methods, we might need to call __new__ with
-            # init_args (like functools.partial)
-            # init_args = [arg.get_example_value() for arg in init_args]
-            # obj = base_cls.__new__(user_cls, *init_args)
-
-            obj = base_cls.__new__(user_cls)
+            if is_structseq_class(user_cls):
+                # Structseq tp_new requires a sequence argument and rejects
+                # tuple.__new__, so create a dummy with None placeholders.
+                obj = user_cls([None] * user_cls.n_fields)
+            elif init_args and issubclass(
+                user_cls,
+                variables.user_defined._CONSTANT_BASE_TYPES,
+            ):
+                example_args = [arg.as_python_constant() for arg in init_args]
+                try:
+                    obj = base_cls.__new__(  # pyrefly: ignore[bad-specialization]
+                        user_cls, *example_args
+                    )
+                except Exception:
+                    # __new__ can raise (e.g., exceeding int str digit limits).
+                    # Fall back to creating without args — the example value is
+                    # only used for tracing, not for correctness.
+                    obj = base_cls.__new__(  # pyrefly: ignore[bad-specialization]
+                        user_cls
+                    )
+            else:
+                obj = base_cls.__new__(user_cls)
         return obj
 
     def track_new_user_defined_object(
@@ -578,7 +635,7 @@ def track_cell_new(
         return variable
 
     def track_cell_existing(
-        self, source: Optional[Source], cell: CellType, contents: VariableTracker
+        self, source: Source | None, cell: CellType, contents: VariableTracker
     ) -> VariableTracker:
         variable = variables.CellVariable(
             # We don't support mutation to cell without source because we need
@@ -657,7 +714,7 @@ def is_live(var: VariableTracker) -> bool:
         # Recursively visit Variables and see if any of them have been mutated.
         init_live_vars = []
         # gather stack/symbolic_locals for all tx's up the chain
-        cur_tx: Optional[InstructionTranslatorBase] = tx
+        cur_tx: InstructionTranslatorBase | None = tx
         while cur_tx is not None:
             init_live_vars.extend([cur_tx.stack, cur_tx.symbolic_locals])
             if cur_tx.parent is not None:
@@ -705,11 +762,9 @@ def mutation(self, var: VariableTracker) -> None:
 
         if isinstance(var.mutation_type, ValueMutationExisting):
             var.mutation_type.is_modified = True
-        if (
-            var.source
-            and isinstance(var, variables.ConstDictVariable)
-            and not isinstance(var, variables.SetVariable)
-        ):
+        if var.source is not None:
+            self.mutated_sources.add(var.source)
+        if var.source and isinstance(var, variables.ConstDictVariable):
             self._has_existing_dict_mutation = True
 
     def has_existing_dict_mutation(self) -> bool:
@@ -730,6 +785,19 @@ def codegen_save_tempvars(self, cg: PyCodegen) -> None:
                 assert var.source is not None
                 continue
 
+            # Namedtuples/structseqs with no pending mutations should skip
+            # codegen_save_tempvars so that restore_stack handles them. In
+            # export, restore_stack uses value_from_source=False which makes
+            # child tensors become graph outputs. If we processed them here,
+            # add_cache would assign a TempLocalSource and restore_stack would
+            # load from cache with value_from_source=True, hiding the tensors
+            # from export.
+            if isinstance(
+                var,
+                (variables.NamedTupleVariable, variables.StructSequenceVariable),
+            ) and not self.has_pending_mutation(var):
+                continue
+
             if isinstance(var, variables.CellVariable):
                 # Cells created in the root frame are created either by
                 # `MAKE_CELL` or by them being in `co_cellvars`, so we only emit
@@ -802,6 +870,26 @@ def load_new_method() -> None:
                 cg.add_cache(var)
                 var.source = TempLocalSource(cg.tempvars[var])
 
+                # For frozen dataclasses, we must emit object.__setattr__
+                # immediately after __new__ — before any other code can
+                # access the object.  The suffix-based codegen in
+                # codegen_update_mutated runs too late: if intervening code
+                # calls __repr__ (e.g. f-strings), the attributes won't be
+                # set yet.
+                if (
+                    isinstance(var, variables.FrozenDataClassVariable)
+                    and var in self.store_attr_mutations
+                ):
+                    for name, value in self.store_attr_mutations[var].items():
+                        cg.load_import_from("builtins", "object")
+                        cg.load_method("__setattr__")
+                        cg(var.source)
+                        cg(variables.ConstantVariable(name))
+                        cg(value)
+                        cg.extend_output(
+                            [*create_call_method(3), create_instruction("POP_TOP")]
+                        )
+
         for ctx, args in self.save_for_backward:
             cg(ctx.source)
             cg.load_method("save_for_backward")
@@ -958,6 +1046,24 @@ def _format_side_effect_message(self, var: VariableTracker) -> str:
 
         return log_str
 
+    def _emit_side_effect_messages(self, side_effect_messages: list[str]) -> None:
+        if not side_effect_messages:
+            return
+
+        for msg in side_effect_messages:
+            side_effects_log.debug(msg)
+
+        torch._logging.trace_structured(
+            "artifact",
+            metadata_fn=lambda: {
+                "name": "dynamo_side_effects",
+                "encoding": "string",
+            },
+            payload_fn=lambda: "\n\n========================================\n\n".join(
+                side_effect_messages
+            ),
+        )
+
     def codegen_update_mutated(
         self, cg: PyCodegen, log_side_effects: bool = False
     ) -> None:
@@ -968,8 +1074,6 @@ def _maybe_log_side_effect(var: VariableTracker) -> None:
             if config.side_effect_replay_policy != "silent" and log_side_effects:
                 msg = self._format_side_effect_message(var)
                 side_effect_messages.append(msg)
-                # Log individual side effects for granular debugging
-                side_effects_log.debug(msg)
 
         suffixes = []
         for var in self._get_modified_vars():
@@ -1026,7 +1130,7 @@ def _maybe_log_side_effect(var: VariableTracker) -> None:
                 )
                 _maybe_log_side_effect(var)
 
-            elif isinstance(var, variables.ConstDictVariable):
+            elif isinstance(var, (variables.ConstDictVariable, variables.SetVariable)):
                 # Reconstruct works as follow:
                 # (1) Skip codegen if there are no new items
                 # (2) codegen(...) each pair of key/value
@@ -1099,10 +1203,20 @@ def _maybe_log_side_effect(var: VariableTracker) -> None:
                     _maybe_log_side_effect(var)
 
             elif self.is_attribute_mutation(var):
+                # FrozenDataClassVariable attributes were emitted in
+                # codegen_save_tempvars right after __new__. Skip here to
+                # avoid double-emitting.
+                if isinstance(var.mutation_type, AttributeMutationNew) and isinstance(
+                    var, variables.FrozenDataClassVariable
+                ):
+                    continue
+
                 if isinstance(
                     var,
                     variables.UserDefinedDictVariable,
-                ) and self.is_modified(var._dict_vt):
+                ) and self.is_modified(
+                    var._base_vt  # pyrefly: ignore[bad-argument-type]
+                ):
                     # Do dict related update manually here. The store_attr
                     # mutations will be applied later.
                     varname_map = {}
@@ -1134,7 +1248,7 @@ def _maybe_log_side_effect(var: VariableTracker) -> None:
                         ]
                     )
 
-                    cg(var._dict_vt, allow_cache=False)  # Don't codegen via source
+                    cg(var._base_vt, allow_cache=False)  # Don't codegen via source
                     cg.extend_output(
                         [
                             create_instruction(
@@ -1153,11 +1267,15 @@ def _maybe_log_side_effect(var: VariableTracker) -> None:
                             create_instruction("POP_TOP"),
                         ]
                     )
-                    _maybe_log_side_effect(var._dict_vt)
+                    _maybe_log_side_effect(
+                        var._base_vt  # pyrefly: ignore[bad-argument-type]
+                    )
                 elif isinstance(
                     var,
                     variables.UserDefinedListVariable,
-                ) and self.is_modified(var._list_vt):
+                ) and self.is_modified(
+                    var._base_vt  # pyrefly: ignore[bad-argument-type]
+                ):
                     # Update the list to the updated items. Be careful in
                     # calling the list methods and not the overridden methods.
                     varname_map = {}
@@ -1173,7 +1291,7 @@ def _maybe_log_side_effect(var: VariableTracker) -> None:
                         ]
                     )
 
-                    cg(var._list_vt, allow_cache=False)  # Don't codegen via source
+                    cg(var._base_vt, allow_cache=False)  # Don't codegen via source
                     cg.extend_output(
                         [
                             create_instruction(
@@ -1192,7 +1310,9 @@ def _maybe_log_side_effect(var: VariableTracker) -> None:
                             create_instruction("POP_TOP"),
                         ]
                     )
-                    _maybe_log_side_effect(var._list_vt)
+                    _maybe_log_side_effect(
+                        var._base_vt  # pyrefly: ignore[bad-argument-type]
+                    )
 
                 # Applying mutations involves two steps: 1) Push all
                 # reconstructed objects onto the stack.  2) Call STORE_ATTR to
@@ -1282,6 +1402,15 @@ def _maybe_log_side_effect(var: VariableTracker) -> None:
                     cg.call_function(1, False)
                     cg.pop_top()
                 _maybe_log_side_effect(var)
+            elif isinstance(var, variables.CountIteratorVariable):
+                for _ in range(var.advance_count):
+                    cg.add_push_null(
+                        lambda: cg.load_import_from(utils.__name__, "iter_next")
+                    )
+                    cg(var.source)  # type: ignore[attr-defined]
+                    cg.call_function(1, False)
+                    cg.pop_top()
+                _maybe_log_side_effect(var)
             elif isinstance(var, variables.RandomVariable):
                 # set correct random seed state
                 def gen_fn() -> None:
@@ -1307,17 +1436,16 @@ def gen_fn() -> None:
 
         # Send batched structured trace for all side effects in this compilation
         if log_side_effects and side_effect_messages:
-            combined_msg = "\n\n========================================\n\n".join(
-                side_effect_messages
-            )
-            torch._logging.trace_structured(
-                "artifact",
-                metadata_fn=lambda: {
-                    "name": "dynamo_side_effects",
-                    "encoding": "string",
-                },
-                payload_fn=lambda: combined_msg,
-            )
+            self._emit_side_effect_messages(side_effect_messages)
+
+    def log_side_effects_summary(self) -> None:
+        if config.side_effect_replay_policy == "silent":
+            return
+        if not side_effects_log.isEnabledFor(logging.DEBUG):
+            return
+        for var in self._get_modified_vars():
+            msg = self._format_side_effect_message(var)
+            side_effects_log.debug(msg)
 
     def is_empty(self) -> bool:
         return not (
diff --git a/torch/_dynamo/source.py b/torch/_dynamo/source.py
index 5309463e77834..3f9453f95c585 100644
--- a/torch/_dynamo/source.py
+++ b/torch/_dynamo/source.py
@@ -21,7 +21,7 @@
 import enum
 import functools
 from collections.abc import Callable
-from typing import Any, Optional, TYPE_CHECKING, Union
+from typing import Any, TYPE_CHECKING
 
 from torch import device as device_type
 from torch._guards import (
@@ -118,7 +118,7 @@ def is_constant_source(source: Source) -> bool:
     return False
 
 
-def _get_source_debug_name(source: Optional[Source]) -> str:
+def _get_source_debug_name(source: Source | None) -> str:
     if source is None:
         return "<unknown source>"
     else:
@@ -152,7 +152,7 @@ class LocalSource(Source):
 
     # Whether we know this input is dynamic (based on example_inputs)
     # For non tensors, we simply look at the first index of the tuple
-    dynamism: Optional[frozenset[str]] = None
+    dynamism: frozenset[str] | None = None
 
     # Whether the item at this source is the _content_ of a cell that is
     # dereferenced from the root frame, i.e., it's a part of the `co_cellvars`
@@ -456,7 +456,7 @@ class UnspecializedParamBufferSource(AttrSource):
 # present within the final view shape metadata.
 @dataclass_with_cached_hash(frozen=True)
 class EphemeralSource(Source):
-    desc: Optional[str] = None
+    desc: str | None = None
 
     @property
     def guard_source(self) -> GuardSource:
@@ -503,7 +503,7 @@ def method_name(self) -> str:
 @dataclass_with_cached_hash(frozen=True)
 class TensorPropertySource(ChainedSource):
     prop: TensorProperty
-    idx: Optional[int] = None  # None for STORAGE_OFFSET
+    idx: int | None = None  # None for STORAGE_OFFSET
 
     def __post_init__(self) -> None:
         assert self.base is not None
@@ -639,7 +639,7 @@ def _name_template(self) -> str:
 
 @dataclass_with_cached_hash(frozen=True)
 class DefaultsSource(ChainedSource):
-    idx_key: Union[int, str]
+    idx_key: int | str
     is_kw: bool = False
     field: str = dataclasses.field(init=False, repr=False, compare=False)
     _name: str = dataclasses.field(init=False, repr=False, compare=False)
@@ -1155,7 +1155,7 @@ def guard_source(self) -> GuardSource:
 @functools.lru_cache
 def get_local_source_name(
     source: Source, *, only_allow_input: bool = False
-) -> Optional[str]:
+) -> str | None:
     if isinstance(source, ChainedSource):
         return get_local_source_name(source.base, only_allow_input=only_allow_input)
     if not isinstance(source, LocalSource):
@@ -1206,9 +1206,11 @@ def is_from_closure_source(source: Source) -> bool:
 
 @functools.lru_cache
 def is_from_source(source: Source, target: Source) -> bool:
+    if source == target:
+        return True
     if isinstance(source, ChainedSource):
         return is_from_source(source.base, target)
-    return source == target
+    return False
 
 
 @functools.lru_cache
diff --git a/torch/_dynamo/symbolic_convert.py b/torch/_dynamo/symbolic_convert.py
index 8d28398fe965d..5017c2ae54cd9 100644
--- a/torch/_dynamo/symbolic_convert.py
+++ b/torch/_dynamo/symbolic_convert.py
@@ -45,7 +45,7 @@
 import types
 import weakref
 from collections import deque
-from typing import Any, cast, NoReturn, Optional, TYPE_CHECKING, TypeAlias, Union
+from typing import Any, cast, NoReturn, TYPE_CHECKING, TypeAlias, TypeVar
 from typing_extensions import TypeIs
 
 import torch
@@ -90,6 +90,7 @@
 )
 from .code_context import code_context
 from .codegen import PyCodegen
+from .comprehension_graph_break import maybe_setup_comprehension_speculation
 from .exc import (
     augment_exc_message_with_hop_name,
     BackendCompilerFailed,
@@ -100,6 +101,7 @@
     StepUnsupported,
     unimplemented,
     Unsupported,
+    UserError,
 )
 from .funcname_cache import get_funcname
 from .guards import GuardBuilder, install_guard
@@ -116,7 +118,6 @@
     ContinueExecutionCache,
     IS_TRACING_RESUME_PROLOGUE_VARNAME,
     ReenterWith,
-    TORCH_DYNAMO_RESUME_IN_PREFIX,
 )
 from .source import (
     AttrSource,
@@ -142,15 +143,15 @@
 )
 from .variables.base import typestr, ValueMutationNew, VariableTracker
 from .variables.builder import FrameStateSizeEntry, VariableBuilder, wrap_fx_proxy
-from .variables.builtin import BuiltinVariable
-from .variables.constant import CONSTANT_VARIABLE_NONE, ConstantVariable
+from .variables.builtin import BuiltinVariable, DictBuiltinVariable
+from .variables.constant import ConstantVariable
 from .variables.ctx_manager import (
     ContextWrappingVariable,
     GenericContextWrappingVariable,
     WithEnterFunctionVariable,
     WithExitFunctionVariable,
 )
-from .variables.dicts import ConstDictVariable, SetVariable
+from .variables.dicts import ConstDictVariable
 from .variables.functions import (
     BaseUserFunctionVariable,
     LocalGeneratorFunctionVariable,
@@ -180,8 +181,9 @@
     UnknownVariable,
 )
 from .variables.nn_module import NNModuleVariable, UnspecializedNNModuleVariable
+from .variables.sets import SetVariable
 from .variables.streams import SymbolicStreamState
-from .variables.tensor import supported_comparison_ops, SymNodeVariable
+from .variables.tensor import supported_comparison_ops, SymNodeVariable, TensorVariable
 from .variables.torch_function import (
     SymbolicTorchFunctionState,
     TorchFunctionModeVariable,
@@ -220,11 +222,11 @@
     tx, [handle_contains(tx, [*reversed(args)], {})], {}
 )
 
-ExceptionVals: TypeAlias = Union[
-    variables.ExceptionVariable,
-    UserDefinedExceptionClassVariable,
-    UserDefinedExceptionObjectVariable,
-]
+ExceptionVals: TypeAlias = (
+    variables.ExceptionVariable
+    | UserDefinedExceptionClassVariable
+    | UserDefinedExceptionObjectVariable
+)
 
 
 @functools.cache
@@ -245,8 +247,8 @@ class SpeculationEntry:
     inst: Instruction  # for debugging only
     _failed: bool = False
     # TX error_on_graph_break setting at the time of failure
-    error_on_graph_break: Optional[bool] = None
-    reason: Optional[GraphCompileReason] = None
+    error_on_graph_break: bool | None = None
+    reason: GraphCompileReason | None = None
 
     def fail_and_restart_analysis(self, error_on_graph_break: bool) -> None:
         """
@@ -283,6 +285,12 @@ class SpeculationLog:
     # If True, graph break at autograd.grad instead of tracing it.
     # Set when we detect that autograd.grad consumed grad_fns that are returned.
     graph_break_on_autograd_grad: bool = False
+    # Names of output tensors whose grad_fns were consumed by autograd.grad.
+    autograd_grad_leaked_tensors: list[str] = dataclasses.field(default_factory=list)
+    # If True, graph break at requires_grad_() on source-less intermediates
+    # instead of tracing it. Set when we detect that such an intermediate
+    # leaks as a graph output with requires_grad=True.
+    graph_break_on_requires_grad_: bool = False
 
     def restart(self) -> None:
         self.index = 0
@@ -356,7 +364,7 @@ def render(self) -> str:
 class DistributedState:
     compile_pg: Any
     local_state: LocalState
-    all_states: Optional[list[LocalState]] = None
+    all_states: list[LocalState] | None = None
 
 
 class TensorifyState:
@@ -425,9 +433,7 @@ class BlockStackEntry:
     inst: Instruction
     target: Instruction | None
     stack_index: int
-    with_context: Optional[
-        Union[ContextWrappingVariable, GenericContextWrappingVariable]
-    ] = None
+    with_context: ContextWrappingVariable | GenericContextWrappingVariable | None = None
 
     def can_restore(self) -> bool:
         return self.with_context is not None
@@ -508,110 +514,6 @@ def fn(x: Any) -> None:
     return insts[begin_idx + 1 : end_idx + 1]
 
 
-@functools.cache
-def _get_comprehension_bytecode_prefix() -> list[str]:
-    """Get the bytecode instructions that precede BUILD_LIST in a list comprehension."""
-
-    assert sys.version_info >= (3, 12)
-
-    def fn() -> list[int]:
-        return [i for i in range(1)]  # noqa: C416
-
-    insts = [inst.opname for inst in dis.get_instructions(fn)]
-
-    start_idx = len(insts) - 1 - insts[::-1].index("LOAD_FAST_AND_CLEAR")
-    end_idx = insts.index("BUILD_LIST")
-
-    return insts[start_idx:end_idx]
-
-
-@functools.cache
-def _get_comprehension_result_patterns() -> dict[str, dict[str, Any]]:
-    """Discover bytecode patterns for comprehension result handling.
-
-    Analyzes sample functions to extract the opcode sequences that appear
-    after END_FOR for each result disposition (stored, discarded, returned, consumed).
-
-    Returns patterns with:
-        - pre_store_ops: opcodes between END_FOR and first STORE_FAST
-        - post_store_op: first opcode after all STORE_FASTs (for disambiguation)
-    """
-    assert sys.version_info >= (3, 12)
-
-    def fn_stored() -> list[int]:
-        result = [i for i in range(1)]  # noqa: C416
-        return result
-
-    def fn_discarded() -> int:
-        [i for i in range(1)]  # noqa: C416
-        return 1
-
-    def fn_returned() -> list[int]:
-        return [i for i in range(1)]  # noqa: C416
-
-    def fn_consumed() -> int:
-        return sum([i for i in range(1)])  # noqa: C416
-
-    def extract_pattern(fn: Callable[..., Any]) -> tuple[list[str], Optional[str]]:
-        """Extract (pre_store_ops, post_store_op) from comprehension bytecode."""
-        target_line = list(dis.findlinestarts(fn.__code__))[1][1]
-        insts: list[str] = []
-        started = False
-        for instr in dis.get_instructions(fn):
-            if started and instr.starts_line:
-                break
-            pos = instr.positions
-            if pos and pos.lineno == target_line:
-                started = started or bool(instr.starts_line)
-                insts.append(instr.opname)
-
-        ops = insts[insts.index("END_FOR") + 1 :]
-        idx = 0
-
-        pre_store_ops = []
-        while idx < len(ops) and ops[idx] != "STORE_FAST":
-            pre_store_ops.append(ops[idx])
-            idx += 1
-
-        while idx < len(ops) and ops[idx] == "STORE_FAST":
-            idx += 1
-
-        return pre_store_ops, ops[idx] if idx < len(ops) else None
-
-    stored = extract_pattern(fn_stored)
-    discarded = extract_pattern(fn_discarded)
-    returned = extract_pattern(fn_returned)
-    consumed = extract_pattern(fn_consumed)
-
-    return {
-        "stored": {"pre_store_ops": stored[0], "post_store_op": stored[1]},
-        "discarded": {"pre_store_ops": discarded[0], "post_store_op": discarded[1]},
-        "returned": {"pre_store_ops": returned[0], "post_store_op": returned[1]},
-        "consumed": {"pre_store_ops": consumed[0], "post_store_op": []},
-    }
-
-
-@dataclasses.dataclass
-class ComprehensionAnalysis:
-    """Metadata about a comprehension's bytecode structure.
-
-    Attributes:
-        end_ip: Instruction pointer after all comprehension bytecode
-        result_var: Name of result variable, or None if result stays on stack
-        result_on_stack: True if result stays on stack (discarded, returned, or in expression)
-        iterator_vars: Variables from LOAD_FAST_AND_CLEAR (need restoration)
-        walrus_vars: Variables assigned via walrus operator (:=) inside comprehension
-        captured_vars: Variables read from outer scope via LOAD_FAST inside comprehension
-    """
-
-    end_ip: int
-    result_var: Optional[str]
-    result_on_stack: bool
-    iterator_vars: list[str]
-    walrus_vars: list[str]
-    captured_vars: list[str]
-
-
 def _detect_and_normalize_assert_statement(
     self: InstructionTranslatorBase,
     truth_fn: Callable[[object], bool],
@@ -648,7 +550,7 @@ def _detect_and_normalize_assert_statement(
                 ].argval
             else:
                 error_msg = "assertion error"
-            self.push(ConstantVariable.create(error_msg))
+            self.push(VariableTracker.build(self, error_msg))
             return True
 
     return False
@@ -687,19 +589,128 @@ def _detect_and_normalize_assert_statement(
 #   3. Any compile_subgraph call should be preceded immediately by a log in the form of "... triggered compile".
 
 
+def get_node_source_info(n: torch.fx.Node) -> str:
+    """Extract the innermost user source location from an FX node's stack trace."""
+    st = n.meta.get("stack_trace", "") or getattr(n, "stack_trace", "")
+    if not st:
+        return ""
+    trace_lines = st.strip().split("\n")
+    last_file_idx = -1
+    for i, line in enumerate(trace_lines):
+        if line.strip().startswith("File "):
+            last_file_idx = i
+    if last_file_idx < 0:
+        return ""
+    file_line = trace_lines[last_file_idx].strip()
+    code = ""
+    if last_file_idx + 1 < len(trace_lines):
+        code = trace_lines[last_file_idx + 1].strip()
+    return f"{file_line}" + (f", code: {code}" if code else "")
+
+
+def format_tensor_computation_trace(value: VariableTracker, max_lines: int = 20) -> str:
+    """Walk the FX graph backwards from a TensorVariable to show how it was
+    computed, with user source locations. Graph inputs (placeholders) are always
+    shown first, followed by operations in dataflow order ending at the branch
+    condition."""
+    if not isinstance(value, TensorVariable):
+        return ""
+    try:
+        node = value.proxy.node
+
+        # Collect operations and placeholders separately so placeholders (root
+        # cause) always appear first regardless of traversal order.
+        op_blocks: list[list[str]] = []
+        placeholder_blocks: list[list[str]] = []
+        visited: set[str] = set()
+
+        def fmt_arg(a: object) -> str:
+            return a.name if isinstance(a, torch.fx.Node) else repr(a)
+
+        def walk(n: torch.fx.Node) -> None:
+            if n.name in visited:
+                return
+            visited.add(n.name)
+            source = get_node_source_info(n)
+
+            if n.op == "placeholder":
+                block = []
+                if source:
+                    block.append(f"# {source}")
+                block.append(f"{n.name}: graph input ({n.target})")
+                placeholder_blocks.append(block)
+                return
+
+            if n.op == "call_function" and len(op_blocks) < max_lines:
+                target_name = getattr(n.target, "__name__", str(n.target))
+                args_str = ", ".join(fmt_arg(a) for a in n.args)
+                block = []
+                if source:
+                    block.append(f"# {source}")
+                block.append(f"{n.name} = {target_name}({args_str})")
+                op_blocks.append(block)
+
+                for a in n.args:
+                    if isinstance(a, torch.fx.Node):
+                        walk(a)
+
+        walk(node)
+
+        if not op_blocks and not placeholder_blocks:
+            return ""
+
+        # Placeholders first (root cause), then ops in dataflow order (reversed
+        # from the DFS collection order) ending at the branch condition.
+        all_lines: list[str] = []
+        for block in placeholder_blocks + list(reversed(op_blocks)):
+            all_lines.extend(block)
+            all_lines.append("")
+
+        return (
+            "\n\n  The branch condition involves a tensor computed as follows:\n"
+            + "\n".join(f"    {line}" for line in all_lines)
+        )
+    except Exception:
+        log.debug("format_tensor_computation_trace failed", exc_info=True)
+        return ""
+
+
 def generic_jump(
     truth_fn: Callable[[object], bool], push: bool
 ) -> Callable[[InstructionTranslatorBase, Instruction], None]:
     def raise_jump_graph_break(value: VariableTracker) -> NoReturn:
+        trace_info = format_tensor_computation_trace(value)
+        hints: list[str] = []
+        if isinstance(value, TensorVariable):
+            try:
+                example = value.proxy.node.meta.get("example_value")
+                if (
+                    example is not None
+                    and example.dim() == 0
+                    and example.dtype
+                    in (
+                        torch.bool,
+                        torch.int32,
+                        torch.int64,
+                    )
+                ):
+                    hints.append(
+                        "The branch condition uses a scalar integer tensor. "
+                        "Consider rewriting the computation to use plain Python "
+                        "ints (e.g. use int attributes instead of tensor buffers) "
+                        "so the condition becomes a shape guard instead of "
+                        "data-dependent branching."
+                    )
+            except Exception:
+                pass
+        hints.extend(graph_break_hints.FUNDAMENTAL)
+        hints.append("Use `torch.cond` to express dynamic control flow.")
         unimplemented(
             gb_type="Data-dependent branching",
             context=f"attempted to jump with {value}",
             explanation="Detected data-dependent branching (e.g. `if my_tensor.sum() > 0:`). "
-            "Dynamo does not support tracing dynamic control flow.",
-            hints=[
-                *graph_break_hints.FUNDAMENTAL,
-                "Use `torch.cond` to express dynamic control flow.",
-            ],
+            "Dynamo does not support tracing dynamic control flow." + trace_info,
+            hints=hints,
         )
 
     def jump_graph_break(
@@ -713,7 +724,7 @@ def jump_graph_break(
         exc = None
         try:
             raise_jump_graph_break(value)
-        except Unsupported as e:
+        except (Unsupported, UserError) as e:
             exc = e
 
         assert exc is not None
@@ -763,6 +774,7 @@ def inner(self: InstructionTranslatorBase, inst: Instruction) -> None:
         value: VariableTracker = self.pop()
         if (
             config.rewrite_assert_with_torch_assert
+            and sys.flags.optimize == 0
             and _detect_and_normalize_assert_statement(self, truth_fn, push)
         ):
             error_msg: VariableTracker = self.pop()
@@ -772,6 +784,7 @@ def inner(self: InstructionTranslatorBase, inst: Instruction) -> None:
                     return self.jump(inst)
                 elif self.should_compile_partial_graph():
                     jump_graph_break(self, inst, value)
+                    return
                 else:
                     unimplemented(
                         gb_type="Data-dependent assertion failed (cannot compile partial graph)",
@@ -871,16 +884,19 @@ def inner(self: InstructionTranslatorBase, inst: Instruction) -> None:
                     x = None
 
             # __bool__ or __len__ is function
-            if isinstance(x, UserMethodVariable):
+            if isinstance(x, (GetAttrVariable, UserMethodVariable)):
                 result = x.call_function(self, [], {})  # type: ignore[arg-type, assignment]
                 method_name = getattr(getattr(x, "fn", None), "__name__", None)
                 if result.is_python_constant():
                     result_value = result.as_python_constant()
                     if method_name == "__bool__" and not isinstance(result_value, bool):
-                        msg = variables.ConstantVariable.create(
-                            f"__bool__ should return bool, returned {type(result_value).__name__}"
+                        exc.raise_observed_exception(
+                            TypeError,
+                            self,
+                            args=[
+                                f"__bool__ should return bool, returned {type(result_value).__name__}"
+                            ],
                         )
-                        exc.raise_observed_exception(TypeError, self, args=[msg])
                     if isinstance(result_value, (bool, int)) and truth_fn(result_value):
                         if push:
                             self.push(value)
@@ -954,7 +970,7 @@ def _reconstruct_block_stack(
     in the compiled bytecode."""
     # Reconstruct the context variable CLASS in the block stack
     all_txes: list[InstructionTranslatorBase] = []
-    cur_tx: Optional[InstructionTranslatorBase] = tx
+    cur_tx: InstructionTranslatorBase | None = tx
     while cur_tx is not None:
         all_txes.append(cur_tx)
         cur_tx = cur_tx.parent
@@ -999,7 +1015,7 @@ def wrapper(self: InstructionTranslatorBase, inst: Instruction) -> None:
                 return handle_graph_break(self, inst, speculation.reason)
             try:
                 return inner_fn(self, inst)
-            except Unsupported as excp:
+            except (Unsupported, UserError) as excp:
                 if self.active_generic_context_managers:
                     # raise original graph break if fullgraph/error_on_graph_break=True
                     if self.one_graph or self.error_on_graph_break:
@@ -1007,7 +1023,8 @@ def wrapper(self: InstructionTranslatorBase, inst: Instruction) -> None:
 
                     # We don't support graph break under GenericContextWrappingVariable,
                     # If there is, we roll back to the checkpoint and fall back.
-                    excp.remove_from_stats()
+                    if isinstance(excp, Unsupported):
+                        excp.remove_from_stats()
                     unimplemented(
                         gb_type="Graph break under GenericContextWrappingVariable",
                         context=f"Active generic context managers: {self.active_generic_context_managers}",
@@ -1019,7 +1036,7 @@ def wrapper(self: InstructionTranslatorBase, inst: Instruction) -> None:
                         from_exc=excp,
                     )
 
-                if excp.skip_frame:
+                if getattr(excp, "skip_frame", False):
                     raise
 
                 if not self.should_compile_partial_graph():
@@ -1034,9 +1051,13 @@ def wrapper(self: InstructionTranslatorBase, inst: Instruction) -> None:
                     exc=excp,
                 )
 
-                excp.remove_from_stats()
-                excp.add_to_stats("graph_break")
-                speculation.reason = GraphCompileReason(excp.msg, excp.real_stack)
+                if isinstance(excp, Unsupported):
+                    excp.remove_from_stats()
+                    excp.add_to_stats("graph_break")
+                speculation.reason = GraphCompileReason(
+                    getattr(excp, "msg", str(excp)),
+                    getattr(excp, "real_stack", [self.frame_summary()]),
+                )
             finally:
                 self.current_instruction_push = prev_push
             speculation.fail_and_restart_analysis(self.error_on_graph_break)
@@ -1130,7 +1151,6 @@ def _missing(opname: str, *args: Any) -> None:
             op: getattr(cls, opname, functools.partial(_missing, opname))
             for opname, op in dis.opmap.items()
         }
-        # pyrefly: ignore [missing-attribute]
         cls.dispatch_table = [dispatch_table.get(i) for i in range(2**8)]
 
 
@@ -1152,7 +1172,7 @@ class ExceptionStack:
     #  + POP_EXCEPT := pops TOS from the *exception stack*
 
     _exc_stack: list[ExceptionVals] = dataclasses.field(default_factory=list)
-    _current_exception: Optional[ExceptionVals] = dataclasses.field(default=None)
+    _current_exception: ExceptionVals | None = dataclasses.field(default=None)
 
     def clear_current_exception(self) -> None:
         self._current_exception = None
@@ -1178,7 +1198,8 @@ def _set_context_recursive(
         if len(self._exc_stack) + prev_idx > 0:
             prev = self._exc_stack[prev_idx]
             self._set_context_recursive(prev, prev_idx - 1)
-            val.set_context(prev)  # type: ignore[union-attr, arg-type]
+            if prev is not val:
+                val.set_context(prev)  # type: ignore[union-attr, arg-type]
         return val
 
     def _break_context_reference_cycle(self, val: ExceptionVals) -> None:
@@ -1194,7 +1215,7 @@ def _break_context_reference_cycle(self, val: ExceptionVals) -> None:
                 break
 
             if context is val:
-                o.set_context(CONSTANT_VARIABLE_NONE)  # type: ignore[union-attr, arg-type]
+                o.set_context(ConstantVariable.create(None))  # type: ignore[union-attr, arg-type]
                 break
 
             o = context  # type: ignore[assignment]
@@ -1241,32 +1262,32 @@ class InstructionTranslatorBase(
     symbolic_globals: dict[str, VariableTracker]
     symbolic_torch_function_state: SymbolicTorchFunctionState
     symbolic_stream_state: SymbolicStreamState
-    post_prune_cell_and_freevars: Optional[dict[str, VariableTracker]]
+    post_prune_cell_and_freevars: dict[str, VariableTracker] | None
     stack: list[VariableTracker]
-    instruction_pointer: Optional[int]
+    instruction_pointer: int | None
     current_instruction: Instruction
     current_instruction_push: bool
     block_stack: list[BlockStackEntry]
     lineno: int
-    kw_names: Optional[ConstantVariable]
+    kw_names: ConstantVariable | None
     accept_prefix_inst: bool
     prefix_insts: list[Instruction]
     inline_depth: int
     inconsistent_side_effects: bool
-    current_speculation: Optional[SpeculationEntry]
+    current_speculation: SpeculationEntry | None
     dispatch_table: list[Any]
     exn_vt_stack: ExceptionStack
-    exec_recorder: Optional[ExecutionRecorder]
-    strict_checks_fn: Optional[Callable[[VariableTracker], bool]]
-    start_point: Optional[int]
+    exec_recorder: ExecutionRecorder | None
+    strict_checks_fn: Callable[[VariableTracker], bool] | None
+    start_point: int | None
     # Does this function make no inlined function calls?
     has_no_inlined_calls: bool
-    parent: Optional[InstructionTranslatorBase]
+    parent: InstructionTranslatorBase | None
     # Does this tx currently have a child tx tracing?
     # Used to correctly implement should_compile_partial_graph
     is_child_tracer_active: bool
     debug_locals: list[tuple[VariableTracker, list[VariableTracker]]]
-    package: Optional[CompilePackage]
+    package: CompilePackage | None
     latest_bytecode_queue: deque[str]
     # Store the latest bytecode before graph_break() call by user
 
@@ -1299,7 +1320,7 @@ def maybe_has_backedge(self) -> bool:
         # negatives so that Dynamo does not skip the whole frame.
 
         # If any parent tx has a backedge, then return True
-        cur_tx: Optional[InstructionTranslatorBase] = self
+        cur_tx: InstructionTranslatorBase | None = self
         while cur_tx is not None:
             cur_offset = cur_tx.current_instruction.offset
             assert cur_tx.instruction_pointer is not None
@@ -1386,7 +1407,7 @@ def inline_user_function_return(
         else:
             return InliningInstructionTranslator.inline_call(self, fn, args, kwargs)
 
-    def get_line_of_code_header(self, lineno: Optional[int] = None) -> str:
+    def get_line_of_code_header(self, lineno: int | None = None) -> str:
         if lineno is None:
             lineno = self.lineno
         inline_depth_str = (
@@ -1465,7 +1486,7 @@ def step(self) -> bool:
             return True
         except (ReturnValueOp, YieldValueOp):
             return False
-        except (Unsupported, StepUnsupported) as e:
+        except (Unsupported, UserError, StepUnsupported) as e:
             # More restrictive condition than should_compile_partial_graph:
             # if this condition is true, then we SHOULD NOT attempt to find
             # a previous checkpoint to resume from and try to resume - we should
@@ -1477,7 +1498,7 @@ def step(self) -> bool:
                 self.one_graph
                 or self.error_on_graph_break
                 or self.is_tracing_resume_prologue
-                or (isinstance(e, Unsupported) and e.skip_frame)
+                or getattr(e, "skip_frame", False)
             ):
                 if isinstance(e, StepUnsupported):
                     unimplemented(
@@ -1508,7 +1529,6 @@ def step(self) -> bool:
                         ],
                         skip_frame=True,
                     )
-                assert isinstance(e, Unsupported)
                 e.skip_frame = True
                 raise
             reason = (
@@ -1704,9 +1724,8 @@ def step_graph_break(self, continue_inst: Instruction) -> None:
                         *create_copy(2),
                         cg.create_load_const(0),
                         cg.create_binary_subscr(),
-                        create_dup_top(),
                         *create_binary_slice(num_stack, None),
-                        *create_swap(2),
+                        *create_copy(3),
                         cg.create_load_const(0),
                         create_instruction("STORE_SUBSCR"),
                     ]
@@ -1811,7 +1830,7 @@ def run(self) -> None:
                 # This code path happens e.g. for bytecodes we don't support
                 # or when we are unable to resume from a graph break.
                 if (
-                    isinstance(e, Unsupported)
+                    isinstance(e, (Unsupported, UserError))
                     and isinstance(self, InstructionTranslator)
                     and not self.error_on_graph_break
                     and not self.one_graph
@@ -1828,6 +1847,9 @@ def run(self) -> None:
                         exc=e,
                     )
 
+                if not isinstance(e, exc.RestartAnalysis):
+                    self.output.side_effects.log_side_effects_summary()
+
                 if hasattr(e, "msg") and "Data-dependent" in e.msg:
                     readable_graph = torch.fx.GraphModule(
                         self.output.nn_modules, self.output.graph
@@ -1841,6 +1863,8 @@ def run(self) -> None:
             except Exception as e:
                 if self.exec_recorder:
                     e.exec_record = self.exec_recorder.get_record()  # type: ignore[attr-defined]
+                if not isinstance(e, exc.RestartAnalysis):
+                    self.output.side_effects.log_side_effects_summary()
 
                 raise
             finally:
@@ -1858,11 +1882,11 @@ def run(self) -> None:
                     # twice is not an issue (second stop is a no op).
                     self.output.mark_bytecode_tracing_stop()
 
-    def push(self, val: Optional[VariableTracker]) -> None:
-        assert val is None or isinstance(val, VariableTracker), (
+    def push(self, val: VariableTracker) -> None:
+        assert isinstance(val, VariableTracker), (
             f"push expects VariableTracker, got {typestr(val)}"
         )
-        self.stack.append(val)  # type: ignore[arg-type]
+        self.stack.append(val)
 
     def push_many(self, vals: list[VariableTracker]) -> None:
         for val in vals:
@@ -1929,8 +1953,54 @@ def STORE_FAST(self, inst: Instruction) -> None:
             self.is_tracing_resume_prologue = val
 
     def DELETE_FAST(self, inst: Instruction) -> None:
+        var = self.symbolic_locals.get(inst.argval)
+        if isinstance(var, TensorVariable):
+            self._maybe_emit_sync_dealloc(var)
         del self.symbolic_locals[inst.argval]
 
+    def _maybe_emit_sync_dealloc(self, var: TensorVariable) -> None:
+        from .variables.streams import get_current_stream, new_event
+
+        device = var.device
+        if device is None or device.type not in ("cuda", "xpu"):
+            return
+
+        node = var.proxy.node
+        alloc_stream = node.meta.get("custom", {}).get("stream", None)
+
+        users = set(node.users.keys())
+        if not users:
+            return
+        last_user = None
+        for n in self.output.graph.nodes:
+            if n in users:
+                last_user = n
+        if last_user is None or last_user.op == "output":
+            return
+
+        last_use_stream = last_user.meta.get("custom", {}).get("stream", None)
+        if alloc_stream == last_use_stream:
+            return
+
+        if alloc_stream is None:
+            alloc_stream = get_current_stream(device)
+        if last_use_stream is None:
+            last_use_stream = get_current_stream(device)
+
+        event_idx = new_event()
+        self.output.create_proxy(
+            "call_function",
+            torch.ops.streams.record_event,
+            (event_idx, last_use_stream),
+            {},
+        )
+        self.output.create_proxy(
+            "call_function",
+            torch.ops.streams.sync_dealloc,
+            (event_idx, alloc_stream, var.as_proxy()),
+            {},
+        )
+
     def STORE_DEREF(self, inst: Instruction) -> None:  # type: ignore[override]
         assert inst.argval in self.cell_and_freevars()
         cell = self.symbolic_locals[inst.argval]
@@ -2126,8 +2196,8 @@ def IMPORT_NAME(self, inst: Instruction) -> None:
             self.exec_recorder.add_local_mod(recorded_name, value)
 
         # pyrefly: ignore [unbound-name]
-        if istype(value, (types.ModuleType, DummyModule)):
-            # pyrefly: ignore [unbound-name]
+        if isinstance(value, (types.ModuleType, DummyModule)):
+            # pyrefly: ignore [unbound-name, bad-argument-type]
             self.push(PythonModuleVariable(value, source=source))
         else:
             unimplemented(
@@ -2168,7 +2238,7 @@ def load_builtin_from_argval(self, argval: Any) -> VariableTracker:
             return VariableTracker.build(self, val, var_source)
         else:
             assert is_builtin_constant(val)
-            return ConstantVariable.create(value=val)
+            return VariableTracker.build(self, val)
 
     def load_builtin(self, inst: Instruction) -> None:
         self.push(self.load_builtin_from_argval(inst.argval))
@@ -2211,20 +2281,6 @@ def SETUP_FINALLY(self, inst: Instruction) -> None:
         assert inst.target is not None
         self.block_stack.append(BlockStackEntry(inst, inst.target, len(self.stack)))
 
-    def BEGIN_FINALLY(self, inst: Instruction) -> None:
-        self.push(None)
-
-    def WITH_CLEANUP_START(self, inst: Instruction) -> None:
-        exit, exc = self.popn(2)
-        assert exc is None
-        self.push(exc)
-
-        self.push(exit.call_function(self, [CONSTANT_VARIABLE_NONE] * 3, {}))
-
-    def WITH_CLEANUP_FINISH(self, inst: Instruction) -> None:
-        self.popn(2)
-        self.push(None)
-
     def FOR_ITER(self, inst: Instruction) -> None:
         it = self.pop().realize()
         self.push(it)
@@ -2240,7 +2296,7 @@ def FOR_ITER(self, inst: Instruction) -> None:
                 # and performs the action of END_FOR as part of FOR_ITER. We jump
                 # to the END_FOR and run it, so we need to make sure 2 values are
                 # on the stack for it to pop.
-                self.push(CONSTANT_VARIABLE_NONE)
+                self.push(ConstantVariable.create(None))
             else:
                 # pop the iterator in Python < 3.12
                 self.pop()
@@ -2271,7 +2327,7 @@ def _attach_traceback_to_exception(self, exc: ExceptionVals) -> None:
             # pyrefly: ignore [bad-argument-type]
             self,
             "__setattr__",
-            [ConstantVariable("__traceback__"), new_tb],
+            [VariableTracker.build(self, "__traceback__"), new_tb],
             {},
         )
 
@@ -2290,7 +2346,7 @@ def _raise_exception_variable(self, val: VariableTracker) -> NoReturn:
             and isinstance(val, variables.ExceptionVariable)
             and val.exc_type is StopIteration
         ):
-            val = variables.BuiltinVariable(RuntimeError).call_function(self, [], {})  # type: ignore[arg-type]
+            val = VariableTracker.build(self, RuntimeError).call_function(self, [], {})  # type: ignore[arg-type]
 
         # Capture the python_stack when the exception is first raised.
         # This preserves the original exception location even if the exception
@@ -2312,21 +2368,23 @@ def _raise_exception_variable(self, val: VariableTracker) -> NoReturn:
             # Pass the stored python_stack to preserve the original exception location
             python_stack = getattr(val, "python_stack", None)
             raise observed_exception_type(
-                f"raised exception {val}", real_stack=python_stack
+                f"raised exception {val.debug_repr()}", real_stack=python_stack
             )
 
-        unimplemented(
-            gb_type="Failed to raise exception",
-            context=str(exc),
-            explanation="Attempted to raise a non-Exception type/value.",
-            hints=[*graph_break_hints.USER_ERROR],
+        exc.raise_observed_exception(
+            TypeError,
+            self,
+            args=[
+                f"exceptions must derive from BaseException, not {val.python_type_name()}",
+            ],
         )
 
     def RAISE_VARARGS(self, inst: Instruction) -> None:
         if inst.arg == 0:
             if not len(self.exn_vt_stack):
-                msg = ConstantVariable("No active exception to reraise")
-                exc.raise_observed_exception(RuntimeError, self, args=[msg])
+                exc.raise_observed_exception(
+                    RuntimeError, self, args=["No active exception to reraise"]
+                )
 
             # re-raise the previous exception. Here CPython refers to the exception
             # on top of the exception stack
@@ -2354,7 +2412,12 @@ def RAISE_VARARGS(self, inst: Instruction) -> None:
                 curr_exc = self.exn_vt_stack.get_current_exception()
                 self._attach_traceback_to_exception(curr_exc)
                 cause = self._create_exception_type(from_vt)
-                curr_exc.call_setattr(self, ConstantVariable("__cause__"), cause)  # type: ignore[arg-type, union-attr, assignment]
+                curr_exc.call_method(
+                    self,  # pyrefly: ignore [bad-argument-type]
+                    "__setattr__",
+                    [VariableTracker.build(self, "__cause__"), cause],
+                    {},
+                )  # type: ignore[arg-type, union-attr, assignment]
 
     def CLEANUP_THROW(self, inst: Instruction) -> None:
         # https://github.com/python/cpython/pull/96010
@@ -2447,7 +2510,7 @@ def bubble_exception_to_interpreter() -> None:
             assert isinstance(raised_exception, dynamo_exc)  # sanity check
             unimplemented(
                 gb_type="Observed exception",
-                context=f"raised exception {curr_exc.python_type_name()}({curr_exc.args})",  # type: ignore[union-attr]
+                context=f"raised exception {curr_exc.debug_repr()}",
                 explanation=observed_exn_gb_explanation,
                 hints=[
                     *graph_break_hints.USER_ERROR,
@@ -2469,7 +2532,7 @@ def bubble_exception_to_interpreter() -> None:
                 # 2) if 'lasti' is true, then push the offset that the exception was raised at
                 if exn_tab_entry.lasti:
                     self.push(
-                        variables.ConstantVariable(self.current_instruction.offset)
+                        VariableTracker.build(self, self.current_instruction.offset)
                     )
 
                 # 3) push the exception to the stack
@@ -2544,9 +2607,9 @@ def bubble_exception_to_interpreter() -> None:
                     self.push(variables.BuiltinVariable(old_exception.exc_type))
                 else:
                     # Push empty exception tb, value, type
-                    self.push(variables.CONSTANT_VARIABLE_NONE)
-                    self.push(variables.CONSTANT_VARIABLE_NONE)
-                    self.push(variables.CONSTANT_VARIABLE_NONE)
+                    self.push(ConstantVariable.create(None))
+                    self.push(ConstantVariable.create(None))
+                    self.push(ConstantVariable.create(None))
 
                 # Push new exception - tb, val, type
                 # Traceback is currently mapped to UnknownVariable
@@ -2587,7 +2650,7 @@ def PUSH_EXC_INFO(self, inst: Instruction) -> None:
 
         val = self.pop()
         if len(self.exn_vt_stack) == 0:
-            prev_exc: VariableTracker = CONSTANT_VARIABLE_NONE
+            prev_exc: VariableTracker = ConstantVariable.create(None)
         else:
             prev_exc = self.exn_vt_stack[-1]
         self.push(prev_exc)
@@ -2691,7 +2754,7 @@ def check_if_exc_matches(self) -> bool:
                 return True
             elif isinstance(exc_instance, variables.BuiltinVariable) and issubclass(
                 exc_instance.fn,
-                # pyrefly: ignore [invalid-argument, missing-attribute]
+                # pyrefly: ignore [invalid-argument]
                 expected_type.fn,
             ):
                 return True
@@ -2699,7 +2762,7 @@ def check_if_exc_matches(self) -> bool:
         return False
 
     def CHECK_EXC_MATCH(self, inst: Instruction) -> None:
-        self.push(variables.ConstantVariable(self.check_if_exc_matches()))
+        self.push(VariableTracker.build(self, self.check_if_exc_matches()))
 
     def JUMP_IF_NOT_EXC_MATCH(self, inst: Instruction) -> None:
         if not self.check_if_exc_matches():
@@ -2712,7 +2775,7 @@ def COMPARE_OP(self, inst: Instruction) -> None:
             self.push(compare_op_handlers[inst.argval](self, self.popn(2), {}))
 
     def GET_ITER(self, inst: Instruction) -> None:
-        self.call_function(BuiltinVariable(iter), [self.pop()], {})
+        self.call_function(VariableTracker.build(self, iter), [self.pop()], {})
 
     @break_graph_if_unsupported(
         push=True,
@@ -2758,18 +2821,14 @@ def CALL_FUNCTION_EX(self, inst: Instruction) -> None:
             assert isinstance(null, NullVariable)
 
         if not isinstance(
-            # pyrefly: ignore [unbound-name]
             argsvars,
             BaseListVariable,
-            # pyrefly: ignore [unbound-name]
         ) and argsvars.has_force_unpack_var_sequence(self):
-            # pyrefly: ignore [unbound-name]
             argsvars = TupleVariable(argsvars.force_unpack_var_sequence(self))
 
         # Unpack for cases like fn(**obj) where obj is a map
-        # pyrefly: ignore [unbound-name]
         if isinstance(kwargsvars, UserDefinedObjectVariable):
-            kwargsvars = BuiltinVariable.call_custom_dict(self, dict, kwargsvars)  # type: ignore[arg-type]
+            kwargsvars = DictBuiltinVariable.call_custom_dict(self, dict, kwargsvars)  # type: ignore[arg-type]
 
         # pyrefly: ignore [unbound-name]
         if not isinstance(argsvars, BaseListVariable) or not isinstance(
@@ -2786,9 +2845,9 @@ def CALL_FUNCTION_EX(self, inst: Instruction) -> None:
             )
 
         # Map to a dictionary of str -> VariableTracker
-        # pyrefly: ignore [bad-assignment, missing-attribute, unbound-name]
+        # pyrefly: ignore [bad-assignment, unbound-name]
         kwargsvars = kwargsvars.keys_as_python_constant()
-        # pyrefly: ignore [bad-argument-type, missing-attribute, unbound-name]
+        # pyrefly: ignore [bad-argument-type, unbound-name]
         self.call_function(fn, argsvars.items, kwargsvars)
 
     @break_graph_if_unsupported(
@@ -2835,21 +2894,20 @@ def LOAD_METHOD(self, inst: Instruction) -> None:
             self.PUSH_NULL(inst)
             self.push(obj)
         else:
-            self.push(obj)
-            self.push(None)
+            raise AssertionError(
+                "LOAD_METHOD should have been rewritten to LOAD_ATTR. We should never reach here."
+            )
 
     def CALL_METHOD(self, inst: Instruction) -> None:
-        args = self.popn(inst.argval)
-        dummy = self.pop()
-        assert dummy is None
-        fn = self.pop()
-        self.call_function(fn, args, {})
+        raise AssertionError(
+            "CALL_METHOD should have been rewritten to CALL_FUNCTION. This function should never be called."
+        )
 
     def _load_attr(self, attr: Any) -> None:
         obj = self.pop()
-        result = BuiltinVariable(getattr).call_function(
+        result = VariableTracker.build(self, getattr).call_function(
             self,  # type: ignore[arg-type]
-            [obj, ConstantVariable.create(attr)],
+            [obj, VariableTracker.build(self, attr)],
             {},
         )
         self.push(result)
@@ -2868,20 +2926,30 @@ def LOAD_ATTR(self, inst: Instruction) -> None:
     )
     def STORE_ATTR(self, inst: Instruction) -> None:
         val, obj = self.popn(2)
-        BuiltinVariable(setattr).call_function(
+        VariableTracker.build(self, setattr).call_function(
             self,  # type: ignore[arg-type]
-            [obj, ConstantVariable.create(inst.argval), val],
+            [obj, VariableTracker.build(self, inst.argval), val],
             {},
         )
 
     def DELETE_ATTR(self, inst: Instruction) -> None:
         obj = self.pop()
-        BuiltinVariable(delattr).call_function(
+        self._maybe_sync_dealloc_attr(obj, inst.argval)
+        VariableTracker.build(self, delattr).call_function(
             self,  # type: ignore[arg-type]
-            [obj, ConstantVariable.create(inst.argval)],
+            [obj, VariableTracker.build(self, inst.argval)],
             {},
         )
 
+    def _maybe_sync_dealloc_attr(self, obj: VariableTracker, name: str) -> None:
+        # Only check side_effects — a pure dict lookup with no observable
+        # side effects. We intentionally avoid var_getattr here because it
+        # can trigger __getattr__, add graph nodes, or cause graph breaks.
+        if self.output.side_effects.has_pending_mutation_of_attr(obj, name):
+            attr_var = self.output.side_effects.load_attr(obj, name)
+            if isinstance(attr_var, TensorVariable):
+                self._maybe_emit_sync_dealloc(attr_var)
+
     @staticmethod
     def codegen_return_with_pops(
         inst: Instruction, num_stack: int
@@ -3088,21 +3156,12 @@ def create_resume(
             )
 
         # add resume function to the global scope
-        if new_code.co_freevars:
-            # expose code object for debugging purposes
-            self.output.install_global_unsafe(resume_name, new_code)
-            package_name = None
-        else:
-            # This is safe: we pre-generate a unique name
-            self.output.install_global_unsafe(
-                resume_name,
-                types.FunctionType(new_code, self.f_globals, resume_name),
-            )
-            package_name = resume_name
-
+        self.output.install_resume_function_global(
+            resume_name, new_code, self.f_globals
+        )
         if self.package is not None:
             self.package.add_resume_function(
-                new_code, self.f_globals["__name__"], package_name
+                new_code, self.f_globals["__name__"], resume_name
             )
 
         counters["resumes"][new_code.co_name] += 1
@@ -3148,7 +3207,7 @@ def create_call_resume_at(
         txes: list[InstructionTranslatorBase] = []
         idxes: list[int] = []
         resume_insts: list[Instruction] = []
-        cur_tx: Optional[InstructionTranslatorBase] = self
+        cur_tx: InstructionTranslatorBase | None = self
         idx = 0
         while cur_tx is not None:
             if cur_tx is self:
@@ -3317,7 +3376,15 @@ def codegen_call_resume(
                         cg.create_binary_subscr(),
                     ]
                 )
-                cg.make_function_with_closure(name, code)
+                # Call the factory function (stored under resume_name) to
+                # create the resume function with correct globals and closure.
+                cg.extend_output(
+                    [
+                        cg.create_load_global(name, add=True),
+                        *create_swap(2),
+                        *create_call_function(1, push_null=True),
+                    ]
+                )
             else:
                 cg.extend_output(cg.load_function_name(name, False, 0))
             cg.extend_output(create_swap(2))
@@ -3338,7 +3405,15 @@ def codegen_call_resume(
                     cg.create_binary_subscr(),
                 ]
             )
-            cg.make_function_with_closure(resume_names[-1], resume_codes[-1])
+            # Call the factory function to create the resume function with
+            # correct globals and closure.
+            cg.extend_output(
+                [
+                    cg.create_load_global(resume_names[-1], add=True),
+                    *create_swap(2),
+                    *create_call_function(1, push_null=True),
+                ]
+            )
             cg.extend_output(
                 [
                     *create_rot_n(3),
@@ -3423,8 +3498,32 @@ def STORE_SUBSCR(self, inst: Instruction) -> None:
 
     def DELETE_SUBSCR(self, inst: Instruction) -> None:
         obj, key = self.popn(2)
+        # Check for tensor items using side-effect-free internal lookups
+        # only. We avoid call_method("__getitem__") because it can execute
+        # user code and add unwanted graph nodes.
+        self._maybe_sync_dealloc_subscr(obj, key)
         obj.call_method(self, "__delitem__", [key], {})
 
+    def _maybe_sync_dealloc_subscr(
+        self, obj: VariableTracker, key: VariableTracker
+    ) -> None:
+        from .variables.dicts import ConstDictVariable
+        from .variables.lists import BaseListVariable
+
+        item_var = None
+        try:
+            if isinstance(obj, BaseListVariable):
+                item_var = obj.getitem_const(
+                    self,  # pyrefly: ignore [bad-argument-type]
+                    key,
+                )
+            elif isinstance(obj, ConstDictVariable):
+                item_var = obj.maybe_getitem_const(key)
+        except Exception:
+            pass
+        if isinstance(item_var, TensorVariable):
+            self._maybe_emit_sync_dealloc(item_var)
+
     def BUILD_TUPLE(self, inst: Instruction) -> None:
         items = self.popn(inst.argval)
         self.push(TupleVariable(items))
@@ -3440,47 +3539,12 @@ def _can_speculate_comprehension_nested(self) -> bool:
         """
         return False
 
-    def _maybe_setup_comprehension_speculation(self, inst: Instruction) -> bool:
-        """
-        Handle comprehension start for Python 3.12+ BUILD_LIST/BUILD_MAP with argval 0.
-        Returns True if a graph break was triggered and the caller should return early.
-        """
-        if not (sys.version_info >= (3, 12) and inst.argval == 0):
-            return False
-
-        if not self._is_comprehension_start():
-            return False
-
-        can_speculate = (
-            all(b.can_restore() for b in self.block_stack)
-            and not self.one_graph
-            and not self.error_on_graph_break
-            and not self.is_tracing_resume_prologue
-            and not self.active_generic_context_managers
-            and self.output.current_tracer.parent is None
-        )
-
-        if can_speculate and self.parent is not None:
-            can_speculate = self._can_speculate_comprehension_nested()
-        # Only set up speculation at depth 0 (outermost comprehension)
-        if can_speculate and self._comprehension_depth == 0:
-            speculation = self.speculate()
-            if speculation.failed(self):
-                self._handle_comprehension_graph_break(inst)
-                return True
-            self.current_speculation = speculation
-        end_for_ip = self._find_comprehension_end_for_ip()
-        assert end_for_ip >= 0
-        self._comprehension_end_for_ips.add(end_for_ip)
-        self._comprehension_depth += 1
-        return False
-
     def BUILD_LIST(self, inst: Instruction) -> None:
-        if self._maybe_setup_comprehension_speculation(inst):
+        if maybe_setup_comprehension_speculation(self, inst):
             return
 
         items = self.popn(inst.argval)
-        self.push(ListVariable(items, mutation_type=ValueMutationNew()))
+        self.push(VariableTracker.build(self, items))
 
     def BUILD_SET(self, inst: Instruction) -> None:
         if config.inject_BUILD_SET_unimplemented_TESTING_ONLY:
@@ -3516,25 +3580,28 @@ def BUILD_TUPLE_UNPACK(self, inst: Instruction) -> None:
     BUILD_TUPLE_UNPACK_WITH_CALL = BUILD_TUPLE_UNPACK
 
     def BUILD_MAP(self, inst: Instruction) -> None:
-        if self._maybe_setup_comprehension_speculation(inst):
+        if maybe_setup_comprehension_speculation(self, inst):
             return
 
         items = self.popn(inst.argval * 2)
         d = dict(zip(items[::2], items[1::2]))
-        self.push(ConstDictVariable(d, mutation_type=ValueMutationNew()))
+        self.push(VariableTracker.build(self, d))
 
     def BUILD_MAP_UNPACK(self, inst: Instruction) -> None:
         items = self.popn(inst.argval)
         # ensure everything is a dict
-        items = [BuiltinVariable(dict).call_function(self, [x], {}) for x in items]  # type: ignore[arg-type]
+        items = [
+            VariableTracker.build(self, dict).call_function(self, [x], {})
+            for x in items
+        ]  # type: ignore[arg-type]
         result: dict[Any, Any] = {}
         for x in items:
             assert isinstance(x, ConstDictVariable)
             result.update(x.items)
         self.push(
-            ConstDictVariable(
+            VariableTracker.build(
+                self,
                 result,
-                mutation_type=ValueMutationNew(),
             )
         )
 
@@ -3550,9 +3617,9 @@ def BUILD_CONST_KEY_MAP(self, inst: Instruction) -> None:
         assert len(keys) == len(values)
 
         self.push(
-            ConstDictVariable(
+            VariableTracker.build(
+                self,
                 dict(zip(keys, values)),
-                mutation_type=ValueMutationNew(),
             )
         )
 
@@ -3601,7 +3668,7 @@ def MAKE_FUNCTION(self, inst: Instruction) -> None:
             # MAKE_FUNCTION behavior actually changed in 3.11, see
             # https://github.com/python/cpython/pull/93189/
             assert hasattr(code.value, "co_qualname")  # type: ignore[attr-defined]
-            fn_name = ConstantVariable.create(value=code.value.co_qualname)  # type: ignore[attr-defined]
+            fn_name = VariableTracker.build(self, code.value.co_qualname)  # type: ignore[attr-defined]
         defaults = None
         closure = None
         annotations = None
@@ -3619,17 +3686,27 @@ def MAKE_FUNCTION(self, inst: Instruction) -> None:
                 if flags & 0x01:
                     defaults = self.pop()
 
-        self.push(
-            NestedUserFunctionVariable(
-                fn_name,
-                code,
-                self.f_globals,
-                defaults,
-                kwdefaults,
-                annotations,
-                closure,
-            )
+        fn = NestedUserFunctionVariable(
+            fn_name,
+            code,
+            self.f_globals,
+            defaults,
+            kwdefaults,
+            closure,
         )
+        if annotations:
+            assert isinstance(annotations, TupleVariable)
+            # Convert the attribute to a dictionary before assigning it
+            # https://github.com/python/cpython/blob/28fb13cb33d569720938258db68956b5f9c9eb40/Objects/funcobject.c#L574-L594
+            items = annotations.items
+            ann = ConstDictVariable(
+                dict(zip(items[::2], items[1::2], strict=True)),
+                mutation_type=ValueMutationNew(),
+            )
+            fn.get_dict_vt(self).setitem(  # pyrefly: ignore[bad-argument-type]
+                "__annotations__", ann
+            )
+        self.push(fn)
 
     def UNPACK_SEQUENCE(self, inst: Instruction) -> None:
         seq = self.pop()
@@ -3649,17 +3726,14 @@ def UNPACK_SEQUENCE(self, inst: Instruction) -> None:
                 "(i.e. `a, b, c = d`).",
                 hints=[*graph_break_hints.USER_ERROR],
             )
-        # pyrefly: ignore [unbound-name]
         if len(val) != inst.argval:
             unimplemented(
                 gb_type="Length mismatch when unpacking object for UNPACK_SEQUENCE",
-                # pyrefly: ignore [unbound-name]
                 context=f"expected length: {inst.argval}, actual: {len(val)}",
                 explanation=f"{seq} unpacked to a list for the UNPACK_SEQUENCE bytecode "
                 "(i.e. `a, b, c = d`) with unexpected length.",
                 hints=[*graph_break_hints.DYNAMO_BUG],
             )
-        # pyrefly: ignore [unbound-name]
         for i in reversed(val):
             self.push(i)
 
@@ -3749,11 +3823,11 @@ def DUP_TOP_TWO(self, inst: Instruction) -> None:
 
     def _convert_value(self, value: VariableTracker, flag: int) -> VariableTracker:
         if flag == 1:
-            return BuiltinVariable(str).call_function(self, [value], {})  # type: ignore[arg-type]
+            return VariableTracker.build(self, str).call_function(self, [value], {})  # type: ignore[arg-type]
         elif flag == 2:
-            return BuiltinVariable(repr).call_function(self, [value], {})  # type: ignore[arg-type]
+            return VariableTracker.build(self, repr).call_function(self, [value], {})  # type: ignore[arg-type]
         elif flag == 3:
-            return BuiltinVariable(ascii).call_function(self, [value], {})  # type: ignore[arg-type]
+            return VariableTracker.build(self, ascii).call_function(self, [value], {})  # type: ignore[arg-type]
         return value
 
     def _format_value(self, fmt_spec: VariableTracker, flags: int) -> None:
@@ -3772,7 +3846,9 @@ def _format_value(self, fmt_spec: VariableTracker, flags: int) -> None:
 
         value = self._convert_value(value, flags & 0x03)
 
-        fmt_var = ConstantVariable.create("{:" + fmt_spec.as_python_constant() + "}")
+        fmt_var = VariableTracker.build(
+            self, "{:" + fmt_spec.as_python_constant() + "}"
+        )
 
         self.call_function(BuiltinVariable(str.format), [fmt_var, value], {})
 
@@ -3782,8 +3858,7 @@ def FORMAT_VALUE(self, inst: Instruction) -> None:
         if (flags & 0x04) == 0x04:
             fmt_spec = self.pop()
         else:
-            fmt_spec = ConstantVariable.create("")
-
+            fmt_spec = VariableTracker.build(self, "")
         return self._format_value(fmt_spec, flags)
 
     def BUILD_STRING(self, inst: Instruction) -> None:
@@ -3867,7 +3942,9 @@ def LIST_EXTEND(self, inst: Instruction) -> None:
         obj.call_method(self, "extend", [v], {})  # type: ignore[arg-type]
 
     def LIST_TO_TUPLE(self, inst: Instruction) -> None:
-        self.push(BuiltinVariable(tuple).call_function(self, [self.pop()], {}))  # type: ignore[arg-type]
+        self.push(
+            VariableTracker.build(self, tuple).call_function(self, [self.pop()], {})
+        )  # type: ignore[arg-type]
 
     def STOPITERATION_ERROR(self, inst: Instruction) -> None:
         # wrap the generator body in a try: ... except StopIteration: ... which
@@ -3878,13 +3955,13 @@ def STOPITERATION_ERROR(self, inst: Instruction) -> None:
         val = self.stack[-1]
         assert self._isinstance_exception(val)
         if val.exc_type is StopIteration:  # type: ignore[union-attr]
-            new_val = variables.BuiltinVariable(RuntimeError).call_function(
+            new_val = VariableTracker.build(self, RuntimeError).call_function(
                 self,  # type: ignore[arg-type]
-                [ConstantVariable("generator raised StopIteration")],
+                [VariableTracker.build(self, "generator raised StopIteration")],
                 {},
             )
-            new_val.call_setattr(self, ConstantVariable("__context__"), val)  # type: ignore[attr-defined]
-            new_val.call_setattr(self, ConstantVariable("__cause__"), val)  # type: ignore[attr-defined]
+            new_val.call_setattr(self, VariableTracker.build(self, "__context__"), val)  # type: ignore[attr-defined]
+            new_val.call_setattr(self, VariableTracker.build(self, "__cause__"), val)  # type: ignore[attr-defined]
             self.stack[-1] = new_val
 
     def DICT_MERGE(self, inst: Instruction) -> None:
@@ -3904,7 +3981,7 @@ def GEN_START(self, inst: Instruction) -> None:
     def GET_LEN(self, inst: Instruction) -> None:
         tos = self.stack[-1]
         if tos.is_python_constant():
-            self.push(ConstantVariable.create(len(tos.as_python_constant())))
+            self.push(VariableTracker.build(self, len(tos.as_python_constant())))
         else:
             self.push(tos.call_method(self, "__len__", [], {}))
 
@@ -3934,7 +4011,7 @@ def MATCH_SEQUENCE(self, inst: Instruction) -> None:
 
     def MATCH_CLASS(self, inst: Instruction) -> None:
         subject, cls, names = self.popn(3)
-        arg = ConstantVariable.create(inst.arg)
+        arg = VariableTracker.build(self, inst.arg)
         self.push(
             self.inline_user_function_return(
                 VariableTracker.build(self, impl_MATCH_CLASS),
@@ -3946,7 +4023,7 @@ def MATCH_CLASS(self, inst: Instruction) -> None:
         if sys.version_info < (3, 11):
             # for versions < 3.11, also push the boolean result
             tos = self.stack[-1]
-            self.push(ConstantVariable.create(not istype(tos, ConstantVariable)))
+            self.push(VariableTracker.build(self, not istype(tos, ConstantVariable)))
 
     def MATCH_KEYS(self, inst: Instruction) -> None:
         keys = self.stack[-1]
@@ -3963,7 +4040,7 @@ def MATCH_KEYS(self, inst: Instruction) -> None:
         if sys.version_info < (3, 11):
             # for versions < 3.11, also push the boolean result
             tos = self.stack[-1]
-            self.push(ConstantVariable.create(not istype(tos, ConstantVariable)))
+            self.push(VariableTracker.build(self, not istype(tos, ConstantVariable)))
 
     def LOAD_ASSERTION_ERROR(self, inst: Instruction) -> None:
         self.push(self.load_builtin_from_argval("AssertionError"))
@@ -3985,6 +4062,7 @@ def LOAD_BUILD_CLASS(self, inst: Instruction) -> None:
     BINARY_REMAINDER = stack_op(operator.mod)
     BINARY_ADD = stack_op(operator.add)
     BINARY_SUBTRACT = stack_op(operator.sub)
+
     BINARY_SUBSCR = break_graph_if_unsupported(
         push=True,
         msg_prefix="Encountered graph break when attempting to trace BINARY_SUBSCR: a binary subscript, e.g. x[attr]",
@@ -4116,7 +4194,7 @@ def BEFORE_WITH(self, inst: Instruction) -> None:
 
     def enter_ctx(
         self,
-        ctx: Union[ContextWrappingVariable, GenericContextWrappingVariable],
+        ctx: ContextWrappingVariable | GenericContextWrappingVariable,
         inst: Instruction,
     ) -> VariableTracker:
         if (
@@ -4256,6 +4334,11 @@ def CALL_INTRINSIC_1(self, inst: Instruction) -> None:
         elif inst.argval == 6:
             # INTRINSIC_LIST_TO_TUPLE
             self.push(TupleVariable(self.pop().force_unpack_var_sequence(self)))
+        elif inst.argval == 7:
+            # INTRINSIC_TYPEVAR
+            v = self.pop().as_python_constant()
+            tv = variables.TypingVariable(TypeVar(v))
+            self.push(tv)
         else:
             unimplemented(
                 gb_type="Missing CALL_INTRINSIC_1 handler",
@@ -4264,6 +4347,28 @@ def CALL_INTRINSIC_1(self, inst: Instruction) -> None:
                 hints=[*graph_break_hints.SUPPORTABLE],
             )
 
+    def CALL_INTRINSIC_2(self, inst: Instruction) -> None:
+        arg2 = self.pop()
+        arg1 = self.pop()
+        if inst.argval == 4:
+            # INTRINSIC_SET_FUNCTION_TYPE_PARAMS
+            # same as => arg1.__type_params__ = arg2
+            assert isinstance(arg1, BaseUserFunctionVariable)
+            arg1.call_method(
+                self,  # pyrefly: ignore[bad-argument-type]
+                "__setattr__",
+                [ConstantVariable.create("__type_params__"), arg2],
+                {},
+            )
+            self.push(arg1)
+        else:
+            unimplemented(
+                gb_type="Missing CALL_INTRINSIC_2 handler",
+                context=f"CALL_INTRINSIC_2 operand: {inst.argval}",
+                explanation=f"No handler implemented for CALL_INTRINSIC_2 {inst.argval} instruction.",
+                hints=[*graph_break_hints.SUPPORTABLE],
+            )
+
     def END_SEND(self, inst: Instruction) -> None:
         tos = self.pop()
         self.pop()
@@ -4302,8 +4407,13 @@ def SET_FUNCTION_ATTRIBUTE(self, inst: Instruction) -> None:
 
             # maybe use Format.VALUE_WITH_FAKE_GLOBALS instead?
             # https://docs.python.org/3/library/annotationlib.html#annotationlib.Format.VALUE_WITH_FAKE_GLOBALS
-            attr = attr.call_function(self, [ConstantVariable.create(1)], {})
-            fn.annotations = attr
+            attr = attr.call_function(self, [VariableTracker.build(self, 1)], {})
+            fn.call_method(
+                self,  # pyrefly: ignore[bad-argument-type]
+                "__setattr__",
+                [ConstantVariable.create("__annotations__"), attr],
+                {},
+            )
         elif flags & 0x08:
             fn.closure = attr
         elif flags & 0x04:
@@ -4315,7 +4425,9 @@ def SET_FUNCTION_ATTRIBUTE(self, inst: Instruction) -> None:
                 dict(zip(items[::2], items[1::2], strict=True)),
                 mutation_type=ValueMutationNew(),
             )
-            fn.annotations = ann
+            fn.get_dict_vt(self).setitem(  # pyrefly: ignore[bad-argument-type]
+                "__annotations__", ann
+            )
         elif flags & 0x02:
             fn.kwdefaults = attr
         elif flags & 0x01:
@@ -4324,10 +4436,11 @@ def SET_FUNCTION_ATTRIBUTE(self, inst: Instruction) -> None:
         self.push(fn)
 
     def CONVERT_VALUE(self, inst: Instruction) -> None:
-        self.push(self._convert_value(self.pop(), inst.argval))
+        assert inst.arg is not None
+        self.push(self._convert_value(self.pop(), inst.arg))
 
     def FORMAT_SIMPLE(self, inst: Instruction) -> None:
-        self._format_value(ConstantVariable.create(""), 0)
+        self._format_value(VariableTracker.build(self, ""), 0)
 
     def FORMAT_WITH_SPEC(self, inst: Instruction) -> None:
         self._format_value(self.pop(), 0)
@@ -4376,7 +4489,7 @@ def LOAD_SPECIAL(self, inst: Instruction) -> None:
             self.PUSH_NULL(inst)
 
     def LOAD_SMALL_INT(self, inst: Instruction) -> None:
-        self.push(ConstantVariable.create(inst.argval))
+        self.push(VariableTracker.build(self, inst.argval))
 
     # See
     # https://github.com/python/cpython/blob/7519ac294fc5c4fd7fb9cb8dc0edc960688cf887/Python/pylifecycle.c#L814
@@ -4404,7 +4517,7 @@ def is_non_empty_graph(self) -> bool:
         return False
 
     def format_frame_summary(
-        self, additional_stack_frames: Optional[list[Any]] = None
+        self, additional_stack_frames: list[Any] | None = None
     ) -> str:
         if additional_stack_frames is None:
             additional_stack_frames = []
@@ -4435,7 +4548,7 @@ def store_global_weakref_by_id(self, prefix: str, value: Any) -> str:
         return global_name
 
     @property
-    def fake_mode(self) -> Optional[FakeTensorMode]:
+    def fake_mode(self) -> FakeTensorMode | None:
         return self.output.tracing_context.fake_mode
 
     @contextlib.contextmanager
@@ -4462,549 +4575,8 @@ def speculate(self) -> SpeculationEntry:
             self.instructions[self.instruction_pointer - 1],
         )
 
-    def _is_comprehension_start(self) -> bool:
-        """Detect if we're at the start of a list/dict comprehension in 3.12+.
-
-        In Python 3.12+, comprehensions are inlined with a bytecode pattern that
-        precedes BUILD_LIST/BUILD_MAP.
-        """
-        assert sys.version_info >= (3, 12)
-
-        assert self.instruction_pointer is not None
-        ip = self.instruction_pointer - 1
-
-        pattern = _get_comprehension_bytecode_prefix()
-        prefix = [inst.opname for inst in self.instructions[ip - len(pattern) : ip]]
-
-        return prefix == pattern
-
-    def _find_comprehension_end_for_ip(self) -> int:
-        """Find the instruction pointer of the outermost END_FOR for current comprehension."""
-        assert sys.version_info >= (3, 12)
-        assert self.instruction_pointer is not None
-
-        nesting_depth = 0
-        for search_ip in range(self.instruction_pointer, len(self.instructions)):
-            inst = self.instructions[search_ip]
-            if inst.opname == "FOR_ITER":
-                nesting_depth += 1
-            elif inst.opname == "END_FOR":
-                nesting_depth -= 1
-                if nesting_depth == 0:
-                    return search_ip
-        return -1
-
-    def _analyze_comprehension(self) -> ComprehensionAnalysis:
-        """Analyze comprehension bytecode to determine result handling pattern."""
-        assert sys.version_info >= (3, 12)
-        assert self.instruction_pointer is not None
-
-        patterns = _get_comprehension_result_patterns()
-        start_ip = self.instruction_pointer - 1  # BUILD_LIST/BUILD_MAP
-
-        iterator_vars: list[str] = []
-        walrus_vars: list[str] = []
-        captured_vars: list[str] = []
-        defined_inside: set[str] = set()
-
-        # Collect iterator variables from LOAD_FAST_AND_CLEAR before BUILD_LIST/BUILD_MAP
-        iter_scan_ip = start_ip - 1
-        while iter_scan_ip >= 0:
-            inst = self.instructions[iter_scan_ip]
-            if inst.opname == "LOAD_FAST_AND_CLEAR":
-                iterator_vars.insert(0, inst.argval)
-                iter_scan_ip -= 1
-            elif inst.opname in ("SWAP", "GET_ITER"):
-                iter_scan_ip -= 1
-            else:
-                break
-        defined_inside.update(iterator_vars)
-
-        end_for_ip = self._find_comprehension_end_for_ip()
-        if end_for_ip == -1:
-            unimplemented(
-                gb_type="Comprehension analysis failed: No END_FOR",
-                context="",
-                explanation="Could not find END_FOR instruction in comprehension bytecode.",
-                hints=[],
-            )
-
-        # Find first FOR_ITER to know where loop body starts
-        for_iter_ip = next(
-            i
-            for i in range(start_ip, end_for_ip)
-            if self.instructions[i].opname == "FOR_ITER"
-        )
-
-        # Single pass through loop body to detect walrus vars and captured vars
-        for body_ip in range(for_iter_ip + 1, end_for_ip):
-            inst = self.instructions[body_ip]
-
-            # Detect walrus pattern: COPY 1 followed by STORE_FAST
-            if inst.opname == "COPY" and inst.arg == 1 and body_ip + 1 < end_for_ip:
-                next_inst = self.instructions[body_ip + 1]
-                if next_inst.opname == "STORE_FAST":
-                    var_name = next_inst.argval
-                    if var_name not in iterator_vars and var_name not in walrus_vars:
-                        walrus_vars.append(var_name)
-                        defined_inside.add(var_name)
-
-            # Track variables defined inside the loop
-            if inst.opname == "STORE_FAST":
-                defined_inside.add(inst.argval)
-
-            # Detect LOAD_FAST referencing outer variables
-            elif inst.opname.startswith("LOAD_FAST"):
-                var_names = (
-                    inst.argval if isinstance(inst.argval, tuple) else (inst.argval,)
-                )
-                for var_name in var_names:
-                    if var_name not in defined_inside and var_name not in captured_vars:
-                        captured_vars.append(var_name)
-
-        # Extract pre_store_ops: all opcodes from END_FOR+1 until first STORE_FAST
-        pre_store_ops: list[str] = []
-        scan_ip = end_for_ip + 1
-        while (
-            scan_ip < len(self.instructions)
-            and self.instructions[scan_ip].opname != "STORE_FAST"
-        ):
-            pre_store_ops.append(self.instructions[scan_ip].opname)
-            scan_ip += 1
-
-        store_fast_ip = scan_ip
-
-        # Skip all STORE_FASTs to find post_store_op
-        while (
-            scan_ip < len(self.instructions)
-            and self.instructions[scan_ip].opname == "STORE_FAST"
-        ):
-            scan_ip += 1
-
-        post_store_op = (
-            self.instructions[scan_ip].opname
-            if scan_ip < len(self.instructions)
-            else None
-        )
-
-        def matches(name: str) -> bool:
-            pat = patterns[name]
-            return pre_store_ops == pat["pre_store_ops"] and (
-                post_store_op == pat["post_store_op"] or not pat["post_store_op"]
-            )
-
-        result_var: Optional[str] = None
-        if matches("stored"):
-            result_var = self.instructions[store_fast_ip].argval
-            result_on_stack = False
-        elif matches("discarded"):
-            result_var = None
-            result_on_stack = False
-            scan_ip = scan_ip + 1 if patterns["discarded"]["post_store_op"] else scan_ip
-        elif (
-            matches("returned")
-            or pre_store_ops == patterns["consumed"]["pre_store_ops"]
-        ):
-            result_var = None
-            result_on_stack = True
-        else:
-            unimplemented(
-                gb_type="Comprehension analysis failed: No matches",
-                context=f"pre_store_ops={pre_store_ops}, post_store_op={post_store_op}",
-                explanation="Comprehension does not match any known bytecode pattern.",
-                hints=[],
-            )
-
-        return ComprehensionAnalysis(
-            end_ip=scan_ip,
-            result_var=result_var,
-            # pyrefly: ignore [unbound-name]
-            result_on_stack=result_on_stack,
-            iterator_vars=iterator_vars,
-            walrus_vars=walrus_vars,
-            captured_vars=captured_vars,
-        )
-
-    def _handle_comprehension_graph_break(self, inst: Instruction) -> None:
-        """Handle list/dict comprehension graph break.
-
-        Builds a synthetic function wrapping the comprehension bytecode,
-        calls it via codegen_call_resume, then chains into the resume
-        function for the post-comprehension code.
-        """
-        assert sys.version_info >= (3, 12)
-
-        analysis = self._analyze_comprehension()
-
-        # Validate: can't handle captured vars in resume functions due to nested sources
-        if self.f_code.co_name.startswith(TORCH_DYNAMO_RESUME_IN_PREFIX):
-            if analysis.captured_vars:
-                unimplemented(
-                    gb_type="Comprehension graph break in resume function with captured variables",
-                    context=str(analysis.captured_vars),
-                    explanation="Cannot use comprehension optimization inside a resume "
-                    "function when there are captured variables. This can cause issues "
-                    "with deeply nested source chains.",
-                    hints=[],
-                )
-
-        assert self.instruction_pointer is not None
-        start_ip = self.instruction_pointer - 1  # BUILD_LIST/BUILD_MAP
-        stack_pops = 1 + len(analysis.iterator_vars)
-        reason = GraphCompileReason("comprehension_graph_break", [self.frame_summary()])
-        log.debug("comprehension triggered compile")
-
-        # --- Step 1: Compile the graph up to the comprehension ---
-
-        all_stack_locals_metadata = self.output.compile_subgraph(
-            self,
-            reason=reason,
-            stack_pops=stack_pops,
-        )
-        # Record which stack_pops items are NULL before popn loses the info.
-        # NULLs on the CPython stack can't be passed as function arguments.
-        stack_pops_null_mask = [
-            isinstance(self.stack[len(self.stack) - stack_pops + i], NullVariable)
-            for i in range(stack_pops)
-        ]
-
-        self.popn(stack_pops)
-        meta = all_stack_locals_metadata[0]
-        cg = PyCodegen(self.output.root_tx)
-
-        # Runtime stack after compile_subgraph:
-        #   cells, [frame_values], *(non-popped items), *(stack_pops items w/ NULLs)
-        # frame_values[0] = [frame N locals] (no stack items yet)
-
-        nonnull_count = sum(1 for m in stack_pops_null_mask if not m)
-
-        # live_stack_depth: stack items above cells/frame_values excluding NULLs
-        # that compile_subgraph didn't codegen (tracked in stack_null_idxes).
-        live_stack_depth = len(self.stack) - len(meta.stack_null_idxes)
-
-        # --- Step 2: Pop stack_pops items and append non-nulls to frame_values[0] ---
-        # SWAP each item to TOS then LIST_APPEND or pop_null; fv_list stays at
-        # TOS throughout. Items append in TOS-first (reversed) order;
-        # _build_comprehension_fn compensates by loading in reverse.
-        cg.extend_output(
-            [
-                # frame_values[0] to TOS
-                *create_copy(live_stack_depth + stack_pops + 1),
-                cg.create_load_const(0),
-                cg.create_binary_subscr(),
-            ]
-        )
-        for i in reversed(range(stack_pops)):
-            cg.extend_output(create_swap(2))
-            if stack_pops_null_mask[i]:
-                cg.extend_output(cg.pop_null())
-            else:
-                cg.extend_output([create_instruction("LIST_APPEND", arg=1)])
-        cg.extend_output([create_instruction("POP_TOP")])
-
-        # Stack: cells, [frame_values], *(non-popped items)
-
-        # --- Step 3: Build comprehension function ---
-        new_code, fn_name = self._build_comprehension_fn(
-            analysis,
-            start_ip,
-            stack_pops,
-            stack_pops_null_mask,
-            nonnull_count,
-            meta,
-        )
-
-        # --- Step 4: Extract [cells[0]] and [frame_values[0]] for codegen_call_resume ---
-        cg.extend_output(
-            [
-                *create_copy(live_stack_depth + 2),
-                cg.create_load_const(0),
-                cg.create_binary_subscr(),
-                create_instruction("BUILD_LIST", arg=1),
-                *create_copy(live_stack_depth + 2),
-                cg.create_load_const(0),
-                cg.create_binary_subscr(),
-                create_instruction("BUILD_LIST", arg=1),
-            ]
-        )
-
-        # Stack: ..., *(non-popped), [cells[0]], [frame_values[0]]
-
-        # --- Step 5: Call comprehension function via codegen_call_resume ---
-        self.codegen_call_resume([new_code], [fn_name], cg)
-
-        # Stack: ..., *(non-popped), comp_result
-
-        # --- Step 6: Remove appended stack_pops items from frame_values[0] ---
-        if nonnull_count > 0:
-            frame_values_pos = live_stack_depth + 1 + 1  # +1 result, +1 frame_values
-            cg.extend_output(
-                [
-                    *create_copy(frame_values_pos),
-                    cg.create_load_const(0),
-                    cg.create_binary_subscr(),
-                    # frame_values[0] on TOS
-                    create_dup_top(),
-                    # frame_values[0], frame_values[0]
-                    cg.create_load_const(-nonnull_count),
-                    cg.create_load_const(None),
-                    create_instruction("BUILD_SLICE", arg=2),
-                    create_instruction("DELETE_SUBSCR"),
-                    # del frame_values[0][-nonnull_count:]
-                    create_instruction("POP_TOP"),
-                ]
-            )
-
-        # --- Step 7: Pass comprehension outputs to frame_values[0] ---
-        # Walrus vars first, then result_var.
-        vars_to_pass = analysis.walrus_vars + (
-            [analysis.result_var] if analysis.result_var else []
-        )
-
-        existing_vars: dict[str, int] = {}
-        for var_name in vars_to_pass:
-            self.symbolic_locals[var_name] = UnknownVariable()
-            if var_name in meta.locals_names:
-                existing_vars[var_name] = meta.locals_names[var_name]
-            else:
-                meta.locals_names[var_name] = len(meta.locals_names)
-
-        fv_depth = live_stack_depth + 2  # comp_result + frame_values
-
-        # --- Walrus vars: extract from comp_result tuple ---
-        if analysis.walrus_vars:
-            # comp_result is (result, *walrus_vars).
-            cg.extend_output(
-                [
-                    *create_copy(fv_depth),
-                    cg.create_load_const(0),
-                    cg.create_binary_subscr(),
-                ]
-            )
-            # Stack: ..., comp_tuple, fv0
-            for j, walrus_var in enumerate(analysis.walrus_vars):
-                cg.extend_output(
-                    [
-                        *create_copy(2),
-                        cg.create_load_const(j + 1),
-                        cg.create_binary_subscr(),
-                    ]
-                )
-                # Stack: ..., comp_tuple, fv0, walrus_value
-                if walrus_var in existing_vars:
-                    # fv0[idx] = walrus_value
-                    cg.extend_output(
-                        [
-                            *create_copy(2),  # copy fv0
-                            cg.create_load_const(existing_vars[walrus_var]),
-                            create_instruction("STORE_SUBSCR"),
-                        ]
-                    )
-                else:
-                    cg.extend_output([create_instruction("LIST_APPEND", arg=1)])
-                # Stack: ..., comp_tuple, fv0
-            cg.extend_output(
-                [
-                    create_instruction("POP_TOP"),  # pop fv0
-                    # Extract the result from the tuple.
-                    cg.create_load_const(0),
-                    cg.create_binary_subscr(),
-                ]
-            )
-            # Stack: ..., result
-
-        # --- Result: keep on stack, overwrite/append to fv[0], or discard ---
-        if analysis.result_on_stack:
-            self.push(UnknownVariable())
-        elif analysis.result_var:
-            cg.extend_output(
-                [
-                    *create_copy(fv_depth),
-                    cg.create_load_const(0),
-                    cg.create_binary_subscr(),
-                    # Stack: ..., result, fv0
-                ]
-            )
-            if analysis.result_var in existing_vars:
-                cg.extend_output(
-                    [
-                        cg.create_load_const(existing_vars[analysis.result_var]),
-                        create_instruction("STORE_SUBSCR"),
-                        # fv0[idx] = result
-                    ]
-                )
-            else:
-                cg.extend_output(
-                    [
-                        *create_swap(2),
-                        create_instruction("LIST_APPEND", arg=1),
-                        create_instruction("POP_TOP"),
-                    ]
-                )
-        else:
-            cg.extend_output([create_instruction("POP_TOP")])
-
-        # Stack: cells, [frame_values], *(non-popped stack)
-        self.output.add_output_instructions(cg.get_instructions())
-
-        # --- Step 8: Create resume function chain ---
-        resume_inst = self.instructions[analysis.end_ip]
-        self.output.add_output_instructions(
-            self.create_call_resume_at(resume_inst, all_stack_locals_metadata)
-        )
-
-        self.instruction_pointer = None
-
-    def _build_comprehension_fn(
-        self,
-        analysis: ComprehensionAnalysis,
-        start_ip: int,
-        stack_pops: int,
-        stack_pops_null_mask: list[bool],
-        nonnull_count: int,
-        meta: StackLocalsMetadata,
-    ) -> tuple[types.CodeType, str]:
-        """Build a synthetic function wrapping comprehension bytecode.
-
-        Uses the same calling convention as resume functions created by
-        create_resume / ContinueExecutionCache.generate: the first two args
-        are __nested_resume_fns and __nested_frame_values (ignored here),
-        followed by stack items and live locals.
-
-        Returns (code, name) where name is the global name for the function.
-        """
-        from .bytecode_transformation import transform_code_object
-        from .eval_frame import skip_code
-        from .resume_execution import CO_VARARGS, CO_VARKEYWORDS
-
-        # Args follow frame_values layout: locals first, then stack_pops items
-        # (appended to end of frame_values[0] by the caller).
-        # codegen_call_resume unpacks frame_values[0] as positional args.
-        argnames = tuple(
-            k for k in meta.locals_names if k not in self.cell_and_freevars()
-        )
-        args = (
-            ["__nested_resume_fns", "__nested_frame_values"]
-            + list(argnames)
-            + [f"___stack{i}" for i in range(nonnull_count)]
-        )
-
-        freevars = tuple(
-            sorted(
-                list(self.f_code.co_cellvars or [])
-                + list(self.f_code.co_freevars or [])
-            )
-        )
-
-        lineno = self.lineno if self.lineno is not None else self.f_code.co_firstlineno
-        fn_name = unique_id(f"__comprehension_{self.f_code.co_name}_at_{lineno}")
-
-        comprehension_body_vars = (
-            analysis.iterator_vars
-            + analysis.walrus_vars
-            + ([analysis.result_var] if analysis.result_var else [])
-            + analysis.captured_vars
-        )
-
-        def update(
-            instructions: list[Instruction], code_options: dict[str, Any]
-        ) -> None:
-            code_options["co_name"] = fn_name
-            if sys.version_info >= (3, 11):
-                code_options["co_qualname"] = fn_name
-            code_options["co_firstlineno"] = lineno
-            code_options["co_cellvars"] = ()
-            code_options["co_freevars"] = freevars
-            code_options["co_argcount"] = len(args)
-            code_options["co_posonlyargcount"] = 0
-            code_options["co_kwonlyargcount"] = 0
-            code_options["co_varnames"] = tuple(
-                args + [v for v in comprehension_body_vars if v not in args]
-            )
-            code_options["co_flags"] = code_options["co_flags"] & ~(
-                CO_VARARGS | CO_VARKEYWORDS
-            )
-
-            prefix: list[Instruction] = []
-            if freevars:
-                prefix.append(create_instruction("COPY_FREE_VARS", arg=len(freevars)))
-            prefix.append(create_instruction("RESUME", arg=0))
-
-            # Push stack_pops items onto operand stack so the comprehension
-            # bytecode finds them where it expects (iterator + saved vars).
-            # NULL positions get PUSH_NULL, non-null get LOAD_FAST.
-            # Items were appended to frame_values[0] in TOS-first order,
-            # so load in reverse to reconstruct the original stack layout.
-            nonnull_i = nonnull_count - 1
-            for i in range(stack_pops):
-                if stack_pops_null_mask[i]:
-                    prefix.append(create_instruction("PUSH_NULL"))
-                else:
-                    prefix.append(
-                        create_instruction("LOAD_FAST", argval=f"___stack{nonnull_i}")
-                    )
-                    nonnull_i -= 1
-
-            comp_insts = self._copy_comprehension_bytecode(start_ip, analysis.end_ip)
-
-            # Epilogue: ensure result is on stack, pack walrus vars, return.
-            epilogue: list[Instruction] = []
-            if not analysis.result_on_stack:
-                if analysis.result_var:
-                    epilogue.append(
-                        create_instruction("LOAD_FAST", argval=analysis.result_var)
-                    )
-                else:
-                    epilogue.append(create_instruction("LOAD_CONST", argval=None))
-            if analysis.walrus_vars:
-                for var_name in analysis.walrus_vars:
-                    epilogue.append(create_instruction("LOAD_FAST", argval=var_name))
-                epilogue.append(
-                    create_instruction(
-                        "BUILD_TUPLE",
-                        arg=1 + len(analysis.walrus_vars),
-                    )
-                )
-            epilogue.append(create_instruction("RETURN_VALUE"))
-
-            instructions[:] = prefix + comp_insts + epilogue
-
-        new_code, _ = transform_code_object(self.f_code, update)
-        skip_code(new_code)
-
-        # Install as global
-        if new_code.co_freevars:
-            self.output.install_global_unsafe(fn_name, new_code)
-        else:
-            self.output.install_global_unsafe(
-                fn_name,
-                types.FunctionType(new_code, self.f_globals, fn_name),
-            )
-
-        return new_code, fn_name
-
-    def _copy_comprehension_bytecode(
-        self, start_ip: int, end_ip: int
-    ) -> list[Instruction]:
-        """Copy comprehension bytecode instructions, updating jump targets."""
-        inst_map: dict[Instruction, Instruction] = {}
-        copied_insts: list[Instruction] = []
-
-        for ip in range(start_ip, end_ip):
-            original_inst = self.instructions[ip]
-            copied_inst = copy.copy(original_inst)
-            copied_inst.exn_tab_entry = None
-            inst_map[original_inst] = copied_inst
-            copied_insts.append(copied_inst)
-
-        for copied_inst in copied_insts:
-            if copied_inst.target is not None and copied_inst.target in inst_map:
-                copied_inst.target = inst_map[copied_inst.target]
-
-        return copied_insts
-
     def _make_frame_loc(
-        self, filename: str, lineno: Optional[int], fallback_lineno: int
+        self, filename: str, lineno: int | None, fallback_lineno: int
     ) -> tuple[str, int]:
         if lineno is None or lineno < 0:
             return (filename, fallback_lineno)
@@ -5016,7 +4588,7 @@ def _get_frame_loc_chain(
         frame_loc_chain_list: list[tuple[str, int]] = []
 
         if config.nested_graph_breaks:
-            current_tx: Optional[InstructionTranslatorBase] = self.parent
+            current_tx: InstructionTranslatorBase | None = self.parent
             while current_tx is not None:
                 parent_frame_loc = self._make_frame_loc(
                     current_tx.f_code.co_filename,
@@ -5034,7 +4606,7 @@ def log_graph_break(
         self,
         code_options: dict[str, Any],
         reason: str,
-        exc: Unsupported | StepUnsupported,
+        exc: Unsupported | UserError | StepUnsupported,
     ) -> None:
         if exc.logged:
             return
@@ -5141,7 +4713,9 @@ def log_graph_break(
         exc.logged = True
 
     @staticmethod
-    def raise_loop_graph_break(code: types.CodeType, exc: Unsupported) -> NoReturn:
+    def raise_loop_graph_break(
+        code: types.CodeType, exc: Unsupported | UserError
+    ) -> NoReturn:
         unimplemented(
             gb_type="graph break in loop",
             context=f"frame skipped: {format_frame_info(code)}",
@@ -5169,12 +4743,12 @@ def __init__(
         inline_depth: int,
         speculation_log: SpeculationLog,
         exn_vt_stack: ExceptionStack,
-        distributed_state: Optional[DistributedState],
+        distributed_state: DistributedState | None,
         # This determines whether to use the execution recorder.
-        closure: Optional[tuple[types.CellType]] = None,
-        package: Optional[CompilePackage] = None,
+        closure: tuple[types.CellType] | None = None,
+        package: CompilePackage | None = None,
         # Pre-computed indexof for cache reuse
-        indexof: Optional[dict[Instruction, int]] = None,
+        indexof: dict[Instruction, int] | None = None,
     ) -> None:
         super().__init__()
         self.speculation_log = speculation_log
@@ -5270,15 +4844,15 @@ def __init__(
 
         self.inline_depth = inline_depth
         self.inconsistent_side_effects = False
-        self._constants_cache: list[
-            Optional[Union[ConstantVariable, SliceVariable]]
-        ] = [None] * len(f_code.co_consts)
+        self._constants_cache: list[ConstantVariable | SliceVariable | None] = [
+            None
+        ] * len(f_code.co_consts)
 
-        self.is_trace_bytecode_log_enabled: Optional[bool] = (
+        self.is_trace_bytecode_log_enabled: bool | None = (
             trace_bytecode_log.isEnabledFor(logging.DEBUG)
         )
-        self.is_trace_source_log_enabled: Optional[bool] = (
-            trace_source_log.isEnabledFor(logging.DEBUG)
+        self.is_trace_source_log_enabled: bool | None = trace_source_log.isEnabledFor(
+            logging.DEBUG
         )
         linecache.lazycache(f_code.co_filename, f_globals)
 
@@ -5304,7 +4878,7 @@ def __init__(
         f_locals: dict[str, Any],
         f_globals: dict[str, Any],
         f_builtins: dict[str, Any],
-        closure: Optional[tuple[Any, ...]],
+        closure: tuple[Any, ...] | None,
         torch_function_mode_stack: Any,
         code_options: dict[str, Any],
         compiler_fn: Any,
@@ -5314,8 +4888,8 @@ def __init__(
         frame_state: Any,
         speculation_log: SpeculationLog,
         exn_vt_stack: ExceptionStack,
-        distributed_state: Optional[DistributedState],
-        package: Optional[CompilePackage],
+        distributed_state: DistributedState | None,
+        package: CompilePackage | None,
     ) -> None:
         _step_logger()(
             logging.INFO,
@@ -5537,7 +5111,6 @@ def _return(self, inst: Instruction) -> None:
             and not self.error_on_graph_break
             and not self.is_tracing_resume_prologue
         ):
-            # TODO graph break if one_graph is set - this might break things
             raise exc.SkipFrame(
                 "No ops traced for the FX graph. `torch.compile` will skip the frame and fall back to eager.\n"
                 f"Frame info: {format_frame_info(self.f_code)}"
@@ -5648,7 +5221,7 @@ def profile_inline_call(
 class InliningInstructionTranslator(InstructionTranslatorBase):
     """Trace and inline a called method"""
 
-    symbolic_result: Optional[VariableTracker]
+    symbolic_result: VariableTracker | None
     # pyrefly: ignore [bad-override]
     parent: InstructionTranslatorBase
 
@@ -5825,23 +5398,30 @@ def get_trace_call_log_str() -> str:
 
         # Detect inline GraphModule calls in order to propagate node metadata,
         # by checking if the first argument (self) is a variable tracking a GraphModule.
-        if args and isinstance(args[0], NNModuleVariable):
-            module = parent.output.get_submodule(args[0].module_key)
-            if isinstance(module, torch.fx.GraphModule):
-                # The inline call might not actually be a call to `forward`,
-                # but it is enough to add a context for `forward` in case it is called.
-                code_context.get_context(module.forward.__code__)[
-                    "orig_graphmodule"
-                ] = weakref.ref(module)
-        # When we have inline_nn_module turned on, modules resolve to UnspecializedNNModuleVariable
-        if args and isinstance(args[0], UnspecializedNNModuleVariable):
-            module = args[0].value
-            if isinstance(module, torch.fx.GraphModule):
-                # The inline call might not actually be a call to `forward`,
-                # but it is enough to add a context for `forward` in case it is called.
-                code_context.get_context(module.forward.__code__)[
-                    "orig_graphmodule"
-                ] = weakref.ref(module)
+        # For unrealized lazy VTs, use peek_type to skip non-Module args
+        # without realizing them (which would install unnecessary guards).
+        if args:
+            arg0 = args[0]
+            should_check = True
+            if isinstance(arg0, LazyVariableTracker) and not arg0.is_realized():
+                if issubclass(arg0.peek_type(), torch.nn.Module):
+                    arg0 = arg0.realize()
+                else:
+                    should_check = False
+
+            if should_check:
+                if isinstance(arg0, NNModuleVariable):
+                    module = parent.output.get_submodule(arg0.module_key)
+                    if isinstance(module, torch.fx.GraphModule):
+                        code_context.get_context(module.forward.__code__)[
+                            "orig_graphmodule"
+                        ] = weakref.ref(module)
+                elif isinstance(arg0, UnspecializedNNModuleVariable):
+                    module = arg0.value
+                    if isinstance(module, torch.fx.GraphModule):
+                        code_context.get_context(module.forward.__code__)[
+                            "orig_graphmodule"
+                        ] = weakref.ref(module)
 
         assert not isinstance(func, SkipFunctionVariable)
         tracer: InliningInstructionTranslator
@@ -5885,7 +5465,7 @@ def inline_call_(self) -> VariableTracker:
             log.debug(msg)
             # bubble up the exception to the parent frame.
             raise
-        except Unsupported as e:
+        except (Unsupported, UserError) as e:
             # If this graph break has skip_frame set, unset it
             # since it refers to the current frame and not the parent.
             e.skip_frame = False
@@ -5902,7 +5482,7 @@ def inline_call_(self) -> VariableTracker:
 
         if self.output.should_exit:
             # graph break
-            return CONSTANT_VARIABLE_NONE  # return dummy variable
+            return ConstantVariable.create(None)  # return dummy variable
 
         assert self.symbolic_result is not None
 
@@ -6008,7 +5588,7 @@ def __init__(
         self.one_graph = parent.one_graph
 
     @property
-    def fake_mode(self) -> Optional[FakeTensorMode]:
+    def fake_mode(self) -> FakeTensorMode | None:
         return self.parent.fake_mode
 
     def run_ctx_mgr(self) -> Any:
@@ -6161,7 +5741,7 @@ def YIELD_VALUE(self, inst: Instruction) -> None:
         self.generated_items.append(top)
         if len(self.generated_items) > MAX_ITERATOR_LIMIT:
             raise exc.InfiniteGeneratorError
-        self.push(CONSTANT_VARIABLE_NONE)
+        self.push(ConstantVariable.create(None))
         if (
             config.enable_faithful_generator_behavior
             or self.is_generator_from_ctx_manager
@@ -6174,7 +5754,7 @@ def GET_YIELD_FROM_ITER(self, inst: Instruction) -> None:
         tos = self.stack[-1]
         if not isinstance(tos, ListIteratorVariable):
             self.pop()
-            res = BuiltinVariable(iter).call_function(self, [tos], {})  # type: ignore[arg-type]
+            res = VariableTracker.build(self, iter).call_function(self, [tos], {})  # type: ignore[arg-type]
             self.push(res)
 
     def RETURN_VALUE(self, inst: Instruction) -> None:
diff --git a/torch/_dynamo/test_case.py b/torch/_dynamo/test_case.py
index c152a58d7855a..f8047cc1b043e 100644
--- a/torch/_dynamo/test_case.py
+++ b/torch/_dynamo/test_case.py
@@ -17,7 +17,7 @@
 import sys
 import unittest
 from collections.abc import Callable
-from typing import Any, Union
+from typing import Any
 
 import torch
 import torch.testing
@@ -36,7 +36,7 @@
 log = logging.getLogger(__name__)
 
 
-def run_tests(needs: Union[str, tuple[str, ...]] = ()) -> None:
+def run_tests(needs: str | tuple[str, ...] = ()) -> None:
     from torch.testing._internal.common_utils import run_tests
 
     if TEST_WITH_TORCHDYNAMO or TEST_WITH_CROSSREF:
@@ -60,6 +60,7 @@ def run_tests(needs: Union[str, tuple[str, ...]] = ()) -> None:
                 importlib.import_module(need)
             except ImportError:
                 return
+
     run_tests()
 
 
@@ -85,6 +86,8 @@ def setUpClass(cls) -> None:
 
     def setUp(self) -> None:
         self._prior_is_grad_enabled = torch.is_grad_enabled()
+        self._prior_nested_graph_breaks = config.nested_graph_breaks
+        config.nested_graph_breaks = True
         super().setUp()
         reset()
         utils.counters.clear()
@@ -94,7 +97,7 @@ def setUp(self) -> None:
     def tearDown(self) -> None:
         trace_log.removeHandler(self.handler)
         for k, v in utils.counters.items():
-            print(k, v.most_common())
+            log.debug("%s %s", k, v.most_common())
         reset()
         utils.counters.clear()
         torch._C._autograd._saved_tensors_hooks_enable()
@@ -102,6 +105,7 @@ def tearDown(self) -> None:
         if self._prior_is_grad_enabled is not torch.is_grad_enabled():
             log.warning("Running test changed grad mode")
             torch.set_grad_enabled(self._prior_is_grad_enabled)
+        config.nested_graph_breaks = self._prior_nested_graph_breaks
 
     def assertEqual(self, x: Any, y: Any, *args: Any, **kwargs: Any) -> None:  # type: ignore[override]
         if (
@@ -116,21 +120,6 @@ def assertEqual(self, x: Any, y: Any, *args: Any, **kwargs: Any) -> None:  # typ
     # graph break tests
 
 
-# NB: multiple inheritance with LoggingTestCase is possible - this should be fine
-# since there is no overlap in overridden methods.
-class TestCaseWithNestedGraphBreaks(TestCase):
-    def setUp(self) -> None:
-        super().setUp()
-        self.prev_nested_graph_breaks = torch._dynamo.config.nested_graph_breaks
-        # pyrefly: ignore [bad-assignment]
-        torch._dynamo.config.nested_graph_breaks = True
-
-    def tearDown(self) -> None:
-        super().tearDown()
-        # pyrefly: ignore [bad-assignment]
-        torch._dynamo.config.nested_graph_breaks = self.prev_nested_graph_breaks
-
-
 class CPythonTestCase(TestCase):
     """
     Test class for CPython tests located in "test/dynamo/CPython/Py_version/*".
@@ -184,7 +173,7 @@ class CPythonTestCase(TestCase):
     def compile_fn(
         self,
         fn: Callable[..., Any],
-        backend: Union[str, Callable[..., Any]],
+        backend: str | Callable[..., Any],
         nopython: bool,
     ) -> Callable[..., Any]:
         # We want to compile only the test function, excluding any setup code
diff --git a/torch/_dynamo/test_minifier_common.py b/torch/_dynamo/test_minifier_common.py
index 07c0c172342ef..40a1eaa811176 100644
--- a/torch/_dynamo/test_minifier_common.py
+++ b/torch/_dynamo/test_minifier_common.py
@@ -24,7 +24,7 @@
 import tempfile
 import traceback
 from collections.abc import Sequence
-from typing import Any, Optional, Union
+from typing import Any
 from unittest.mock import patch
 
 import torch
@@ -47,7 +47,7 @@ def _get_module(self, t: str) -> str:
         r = re.sub(r"\n{3,}", "\n\n", r)
         return r.strip()
 
-    def get_exported_program_path(self) -> Optional[str]:
+    def get_exported_program_path(self) -> str | None:
         # Extract the exported program file path from AOTI minifier's repro.py
         # Regular expression pattern to match the file path
         pattern = r'torch\.export\.load\(\s*["\'](.*?)["\']\s*\)'
@@ -109,7 +109,7 @@ def _gen_codegen_fn_patch_code(self, device: str, bug_type: str) -> str:
 """
 
     def _maybe_subprocess_run(
-        self, args: Sequence[Any], *, isolate: bool, cwd: Optional[str] = None
+        self, args: Sequence[Any], *, isolate: bool, cwd: str | None = None
     ) -> subprocess.CompletedProcess[bytes]:
         from torch._inductor.cpp_builder import normalize_path_separator
 
@@ -180,7 +180,7 @@ def _maybe_subprocess_run(
     # minifier launcher script, if `code` outputted it.
     def _run_test_code(
         self, code: str, *, isolate: bool
-    ) -> tuple[subprocess.CompletedProcess[bytes], Union[str, Any]]:
+    ) -> tuple[subprocess.CompletedProcess[bytes], str | Any]:
         proc = self._maybe_subprocess_run(
             ["python3", "-c", code], isolate=isolate, cwd=self.DEBUG_DIR
         )
@@ -201,7 +201,7 @@ def _run_minifier_launcher(
         isolate: bool,
         *,
         minifier_args: Sequence[Any] = (),
-        repro_after: Optional[str] = None,
+        repro_after: str | None = None,
     ) -> tuple[subprocess.CompletedProcess[bytes], str]:
         self.assertIsNotNone(repro_dir)
         launch_file = _as_posix_path(os.path.join(repro_dir, "minifier_launcher.py"))
@@ -282,11 +282,11 @@ def _run_full_test(
         self,
         run_code: str,
         repro_after: str,
-        expected_error: Optional[str],
+        expected_error: str | None,
         *,
         isolate: bool,
         minifier_args: Sequence[Any] = (),
-    ) -> Optional[MinifierTestResult]:
+    ) -> MinifierTestResult | None:
         if isolate:
             repro_level = 3
         elif expected_error is None or expected_error == "AccuracyError":
diff --git a/torch/_dynamo/testing.py b/torch/_dynamo/testing.py
index 3928e1e799bd0..22badc5814b6e 100644
--- a/torch/_dynamo/testing.py
+++ b/torch/_dynamo/testing.py
@@ -24,7 +24,7 @@
 import types
 import unittest
 from collections.abc import Callable, Generator, Sequence
-from typing import Any, Optional, overload, TypeVar, Union
+from typing import Any, overload, TypeVar
 from typing_extensions import ParamSpec
 from unittest.mock import patch
 
@@ -45,7 +45,7 @@
 from .utils import CompileCounterInt, same
 
 
-np: Optional[types.ModuleType] = None
+np: types.ModuleType | None = None
 try:
     import numpy as np
 except ModuleNotFoundError:
@@ -60,7 +60,7 @@
 _P = ParamSpec("_P")
 
 
-def clone_me(x: Optional[torch.Tensor]) -> Optional[torch.Tensor]:
+def clone_me(x: torch.Tensor | None) -> torch.Tensor | None:
     if x is None:
         return None
     return x.detach().clone().requires_grad_(x.requires_grad)
@@ -151,11 +151,11 @@ def reduce_to_scalar_loss(out: torch.Tensor) -> torch.Tensor: ...
 
 @overload
 def reduce_to_scalar_loss(
-    out: Union[list[Any], tuple[Any, ...], dict[Any, Any]],
+    out: list[Any] | tuple[Any, ...] | dict[Any, Any],
 ) -> float: ...
 
 
-def reduce_to_scalar_loss(out: Any) -> Union[torch.Tensor, float]:
+def reduce_to_scalar_loss(out: Any) -> torch.Tensor | float:
     """Reduce the output of a model to get scalar loss"""
     if isinstance(out, torch.Tensor):
         # Mean does not work on integer tensors
@@ -233,7 +233,7 @@ def insert_nops(instructions: list[Any], code_options: Any) -> None:
 
 class CompileCounter:
     def __init__(self) -> None:
-        self.frame_count: Union[int, CompileCounterInt] = 0
+        self.frame_count: int | CompileCounterInt = 0
         self.clear()
 
     def __call__(
@@ -255,7 +255,7 @@ def clear(self) -> None:
 
 class CompileCounterWithBackend:
     def __init__(self, backend: str) -> None:
-        self.frame_count: Union[int, CompileCounterInt] = 0
+        self.frame_count: int | CompileCounterInt = 0
         self.backend = backend
         self.graphs: list[torch.fx.GraphModule] = []
         self.clear()
@@ -388,8 +388,8 @@ def standard_test(
     self: Any,
     fn: Callable[..., Any],
     nargs: int,
-    expected_ops: Optional[int] = None,
-    expected_ops_dynamic: Optional[int] = None,
+    expected_ops: int | None = None,
+    expected_ops_dynamic: int | None = None,
     expected_frame_count: int = 1,
 ) -> None:
     if not config.assume_static_by_default and expected_ops_dynamic is not None:
@@ -440,7 +440,7 @@ def rand_strided(
     size: Sequence[int],
     stride: Sequence[int],
     dtype: torch.dtype = torch.float32,
-    device: Union[str, torch.device] = "cpu",
+    device: str | torch.device = "cpu",
     extra_size: int = 0,
 ) -> torch.Tensor:
     needed_size = extra_size
@@ -450,7 +450,11 @@ def rand_strided(
             sum((shape - 1) * stride for shape, stride in zip(size, stride)) + 1
         )
     if dtype.is_floating_point:
-        if dtype.itemsize == 1:
+        if dtype == torch.float4_e2m1fn_x2:
+            buffer = torch.randint(
+                0, 256, (needed_size,), dtype=torch.uint8, device=device
+            ).view(torch.float4_e2m1fn_x2)
+        elif dtype.itemsize == 1:
             """
             normal distribution kernel is not implemented for fp8..
             Workaround that by creating a fp16 tensor and then cast.
@@ -490,7 +494,7 @@ def make_test_cls_with_patches(
     cls_prefix: str,
     fn_suffix: str,
     *patches: Any,
-    xfail_prop: Optional[str] = None,
+    xfail_prop: str | None = None,
     decorator: Callable[[Callable[..., Any]], Callable[..., Any]] = lambda x: x,
 ) -> type:
     DummyTestClass = type(f"{cls_prefix}{cls.__name__}", cls.__bases__, {})
@@ -600,6 +604,7 @@ def _testing_capture_invoke_subgraph_inductor_compile_gms() -> Generator[
         # captured_gms will contain the list of captured graph modules
     """
     global _testing_invoke_subgraph_inductor_compile_captured_gms
+    # pyrefly: ignore [implicit-any]
     _testing_invoke_subgraph_inductor_compile_captured_gms = []
     try:
         yield _testing_invoke_subgraph_inductor_compile_captured_gms
diff --git a/torch/_dynamo/trace_rules.py b/torch/_dynamo/trace_rules.py
index 17cd668bb8eaa..9f757639f1b7d 100644
--- a/torch/_dynamo/trace_rules.py
+++ b/torch/_dynamo/trace_rules.py
@@ -34,11 +34,12 @@
 import re
 import sys
 import types
+import typing
 import unittest
 from collections import defaultdict
 from collections.abc import Callable, Iterator
 from pathlib import Path
-from typing import Any, cast, Optional, Union
+from typing import Any
 
 import torch
 import torch._inductor.test_operators
@@ -58,9 +59,12 @@
 )
 from .variables import (
     BuiltinVariable,
+    DictBuiltinVariable,
     FunctionalCallVariable,
     FunctorchHigherOrderVariable,
     InspectSignatureVariable,
+    IterBuiltinVariable,
+    ListBuiltinVariable,
     LocalGeneratorFunctionVariable,
     LocalGeneratorObjectVariable,
     NestedUserFunctionVariable,
@@ -77,7 +81,7 @@
 from .variables.base import VariableTracker
 
 
-np: Optional[types.ModuleType] = None
+np: types.ModuleType | None = None
 try:
     import numpy as np
 except ModuleNotFoundError:
@@ -153,11 +157,9 @@
 """
 manual_torch_name_rule_map: dict[
     str,
-    Union[
-        type[TorchInGraphFunctionVariable],
-        type[SkipFunctionVariable],
-        type[UserFunctionVariable],
-    ],
+    type[TorchInGraphFunctionVariable]
+    | type[SkipFunctionVariable]
+    | type[UserFunctionVariable],
 ] = {
     "torch.onnx.is_in_onnx_export": TorchInGraphFunctionVariable,
     "torch.onnx.operators.shape_as_tensor": TorchInGraphFunctionVariable,
@@ -170,11 +172,14 @@
     "torch.distributed.get_rank": TorchInGraphFunctionVariable,
     "torch.distributed.get_world_size": TorchInGraphFunctionVariable,
     "torch.distributed.tensor._api.DTensor#from_local": TorchInGraphFunctionVariable,
+    "torch.distributed.device_mesh.DeviceMesh#__init__": SkipFunctionVariable,
     "torch.distributed.distributed_c10d._get_group_size_by_name": TorchInGraphFunctionVariable,
     "torch.distributed.distributed_c10d._resolve_group_name_by_ranks_and_tag": TorchInGraphFunctionVariable,
     "torch.distributed.distributed_c10d._get_group_tag": TorchInGraphFunctionVariable,
     "torch.distributed.distributed_c10d.get_process_group_ranks": TorchInGraphFunctionVariable,
+    "torch.distributed.destroy_process_group": SkipFunctionVariable,
     "torch._utils.is_compiling": TorchInGraphFunctionVariable,
+    "torch._utils._chunk_or_narrow_cat": UserFunctionVariable,
     "torch.fx._symbolic_trace.is_fx_tracing": TorchInGraphFunctionVariable,
     "torch.fx._symbolic_trace.is_fx_symbolic_tracing": TorchInGraphFunctionVariable,
     "torch._dynamo.external_utils.is_compiling": TorchInGraphFunctionVariable,
@@ -188,17 +193,14 @@
     "torch.to_dlpack": SkipFunctionVariable,
     "torch._check": TorchInGraphFunctionVariable,
     # We graph break on RNG state setters or getters like
-    # `torch.get_rng_state` or `torch.set_rng_state`. These functions
-    # are not aten operations and therefore they are completely ignored
-    # by the AOT dispatcher. As a result, the AOT graph does not have
-    # these setter or getter functions, producing an incorrect graph
-    # when it comes to rng states.
-    "torch.default_generator#get_state": SkipFunctionVariable,
-    "torch._C.Generator#get_state": SkipFunctionVariable,
+    # `torch.get_rng_state`, `torch.set_rng_state`, and
+    # `torch.Generator.manual_seed`. These functions are not aten
+    # operations and therefore they are completely ignored by the AOT
+    # dispatcher. As a result, the AOT graph does not have these setter
+    # or getter functions, producing an incorrect graph when it comes
+    # to rng states.
     "torch.get_rng_state": SkipFunctionVariable,
     "torch.cuda.get_rng_state": SkipFunctionVariable,
-    "torch.default_generator#set_state": SkipFunctionVariable,
-    "torch._C.Generator#set_state": SkipFunctionVariable,
     "torch.set_rng_state": SkipFunctionVariable,
     "torch.cuda.set_rng_state": SkipFunctionVariable,
     # https://github.com/pytorch/pytorch/issues/107187
@@ -234,18 +236,13 @@
     "torch.Tensor#split": TorchInGraphFunctionVariable,
     "torch.cuda.set_device": SkipFunctionVariable,
     "torch.cuda.current_device": TorchInGraphFunctionVariable,
-    "torch._C.autocast_decrement_nesting": SkipFunctionVariable,
-    "torch._C.autocast_increment_nesting": SkipFunctionVariable,
     "torch.autograd.grad": TorchInGraphFunctionVariable,
     "torch.autograd.backward": SkipFunctionVariable,
-    "torch._C.clear_autocast_cache": SkipFunctionVariable,
     "torch.distributions.constraints.is_dependent": SkipFunctionVariable,
     "torch.jit.isinstance": SkipFunctionVariable,
     "torch._C.set_anomaly_enabled": SkipFunctionVariable,
-    "torch._C.set_autocast_cache_enabled": SkipFunctionVariable,
     "torch._C.set_autocast_cpu_dtype": SkipFunctionVariable,
     "torch._C.set_autocast_cpu_enabled": SkipFunctionVariable,
-    "torch._C.set_autocast_enabled": SkipFunctionVariable,
     "torch._C.set_autocast_gpu_dtype": SkipFunctionVariable,
     "torch._C.set_autocast_ipu_dtype": SkipFunctionVariable,
     "torch._C.set_autocast_ipu_enabled": SkipFunctionVariable,
@@ -355,10 +352,14 @@
     "torch._dynamo.dont_skip_tracing": UserFunctionVariable,
     "torch._dynamo.mark_static": UserFunctionVariable,
     "torch._dynamo.nonstrict_trace": UserFunctionVariable,
+    "torch._dynamo.bytecode_debugger.breakpoint": UserFunctionVariable,
     "torch._dynamo.patch_dynamo_config": UserFunctionVariable,
     "torch._dynamo.error_on_graph_break": UserFunctionVariable,
+    "torch._dynamo.override_cudagraphs": UserFunctionVariable,
     "torch.fx.experimental.symbolic_shapes.guard_size_oblivious": TorchInGraphFunctionVariable,
-    "torch.fx.experimental.symbolic_shapes.size_hint": TorchInGraphFunctionVariable,
+    "torch.fx.experimental.symbolic_shapes.is_nested_int": UserFunctionVariable,
+    "torch.fx.experimental.symbolic_shapes.guarding_hint_or_throw": TorchInGraphFunctionVariable,
+    "torch.fx.experimental.symbolic_shapes.optimization_hint": TorchInGraphFunctionVariable,
     "torch.fx.experimental.symbolic_shapes.guard_or_true": TorchInGraphFunctionVariable,
     "torch.fx.experimental.symbolic_shapes.guard_or_false": TorchInGraphFunctionVariable,
     "torch.fx.experimental.symbolic_shapes.statically_known_true": TorchInGraphFunctionVariable,
@@ -367,6 +368,7 @@
     "torch.fx.experimental.symbolic_shapes.sym_or": TorchInGraphFunctionVariable,
     "torch.fx.experimental.symbolic_shapes.guard_scalar": TorchInGraphFunctionVariable,
     "torch.fx.experimental.symbolic_shapes.has_static_value": TorchInGraphFunctionVariable,
+    "torch.fx.experimental.symbolic_shapes.has_free_unbacked_symbols": TorchInGraphFunctionVariable,
     "torch.cuda._get_device_properties": TorchInGraphFunctionVariable,
     "torch.utils.hooks.BackwardHook": TorchInGraphFunctionVariable,
     "torch.set_default_device": UserFunctionVariable,
@@ -392,6 +394,27 @@
     "inspect.signature": InspectSignatureVariable,
 }
 
+# Keep this in sync with the stateful generator methods exposed from
+# torch/csrc/Generator.cpp.
+_GENERATOR_METHODS_THAT_GRAPH_BREAK = (
+    "clone_state",
+    "get_offset",
+    "get_state",
+    "graphsafe_get_state",
+    "graphsafe_set_state",
+    "initial_seed",
+    "manual_seed",
+    "seed",
+    "set_offset",
+    "set_state",
+)
+
+for generator_prefix in ("torch.default_generator", "torch._C.Generator"):
+    for method_name in _GENERATOR_METHODS_THAT_GRAPH_BREAK:
+        manual_torch_name_rule_map[f"{generator_prefix}#{method_name}"] = (
+            SkipFunctionVariable
+        )
+
 
 # In graph functions (including constant folding) that are C bindings
 torch_c_binding_in_graph_functions = dict.fromkeys(
@@ -463,6 +486,7 @@
         "torch._C._accelerator_getAccelerator",
         "torch._C._accelerator_getDeviceIndex",
         "torch._C._accelerator_getStream",
+        "torch._C._accelerator_getAllocatorSettings",
         "torch._C._accelerator_setAllocatorSettings",
         "torch._C._accelerator_setStream",
         "torch._C._accelerator_synchronizeDevice",
@@ -471,6 +495,8 @@
         "torch._C._add_docstr",
         "torch._C._are_functorch_transforms_active",
         "torch._C._autograd_init",
+        "torch._C._autograd._saved_tensors_hooks_disable",
+        "torch._C._autograd._saved_tensors_hooks_enable",
         "torch._C._awaitable_nowait",
         "torch._C._awaitable_wait",
         "torch._C._awaitable",
@@ -542,7 +568,6 @@
         "torch._C._cuda_isHistoryEnabled",
         "torch._C._cuda_isInBadFork",
         "torch._C._cuda_jiterator_compile_and_launch_kernel",
-        "torch._C._cuda_lock_mutex",
         "torch._C._cuda_maybeExchangeDevice",
         "torch._C._cuda_memorySnapshot",
         "torch._C._cuda_memoryStats",
@@ -561,7 +586,6 @@
         "torch._C._cuda_setStream",
         "torch._C._cuda_sleep",
         "torch._C._cuda_synchronize",
-        "torch._C._cuda_unlock_mutex",
         "torch._C._cudnn_set_conv_benchmark_empty_cache",
         "torch._C._cudnn.getCompileVersion",
         "torch._C._cudnn.getRuntimeVersion",
@@ -651,6 +675,9 @@
         "torch._C._from_dlpack",
         "torch._C._functionality_to_backend_keys",
         "torch._C._functionalization_reapply_views_tls",
+        "torch._C._functorch._grad_decrement_nesting",
+        "torch._C._functorch._grad_increment_nesting",
+        "torch._C._functorch.set_inplace_requires_grad_allowed",
         "torch._C._fuse_to_static_module",
         "torch._C._gather_out",
         "torch._C._gather",
@@ -737,6 +764,7 @@
         "torch._C._initExtension",
         "torch._C._is_alias_of",
         "torch._C._is_any_autocast_enabled",
+        "torch._C._is_autocast_available",
         "torch._C._is_cached_tensor",
         "torch._C._is_flash_attention_available",
         "torch._C._is_fwd_grad_enabled",
@@ -1074,6 +1102,7 @@
         "torch._C._nn._test_warn_in_autograd",
         "torch._C._nn._upsample_bicubic2d_aa",
         "torch._C._nn._upsample_bilinear2d_aa",
+        "torch._C._nn._upsample_lanczos2d_aa",
         "torch._C._nn._upsample_nearest_exact1d",
         "torch._C._nn._upsample_nearest_exact2d",
         "torch._C._nn._upsample_nearest_exact3d",
@@ -1391,6 +1420,9 @@
         "torch._C._xpu_resetPeakMemoryStats",
         "torch._C._xpu_setStream",
         "torch._C._xpu_synchronize",
+        "torch._C.autocast_decrement_nesting",
+        "torch._C.autocast_increment_nesting",
+        "torch._C.clear_autocast_cache",
         "torch._C.fork",
         "torch._C.get_autocast_cpu_dtype",
         "torch._C.get_autocast_dtype",
@@ -1416,10 +1448,9 @@
         "torch._C.parse_ir",
         "torch._C.parse_schema",
         "torch._C.parse_type_comment",
-        "torch._C.read_vitals",
-        "torch._C.set_vital",
+        "torch._C.set_autocast_cache_enabled",
+        "torch._C.set_autocast_enabled",
         "torch._C.unify_type_list",
-        "torch._C.vitals_enabled",
         "torch._C.wait",
         "torch._cast_Byte",
         "torch._cast_Char",
@@ -1498,6 +1529,7 @@
         "torch._foreach_clamp_max",
         "torch._foreach_clamp_min_",
         "torch._foreach_clamp_min",
+        "torch._foreach_clone",
         "torch._foreach_copy_",
         "torch._foreach_cos_",
         "torch._foreach_cos",
@@ -1580,6 +1612,7 @@
         "torch._functionalize_replace",
         "torch._functionalize_sync",
         "torch._functionalize_was_storage_changed",
+        "torch._fused_adagrad_",
         "torch._fused_adam_",
         "torch._fused_adamw_",
         "torch._fused_dropout",
@@ -2215,6 +2248,7 @@
         "torch.select",
         "torch.selu_",
         "torch.selu",
+        "torch.set_autocast_dtype",
         "torch.sgn",
         "torch.sigmoid_",
         "torch.sigmoid",
@@ -2438,8 +2472,6 @@
         "torch.accelerator.set_stream",
         "torch.accelerator.synchronize",
         "torch.align_tensors",
-        "torch.amp.autocast_mode._enter_autocast",
-        "torch.amp.autocast_mode._exit_autocast",
         "torch.amp.autocast_mode.autocast_decorator",
         "torch.amp.autocast_mode.custom_bwd",
         "torch.amp.autocast_mode.custom_fwd",
@@ -2613,6 +2645,7 @@
         "torch.cuda.get_device_properties",
         "torch.cuda.get_gencode_flags",
         "torch.cuda.get_sync_debug_mode",
+        "torch.cuda.synchronize",
         "torch.cuda.graphs.graph_pool_handle",
         "torch.cuda.graphs.is_current_stream_capturing",
         "torch.cuda.graphs.make_graphed_callables",
@@ -2625,7 +2658,6 @@
         "torch.cuda.jiterator._create_multi_output_jit_fn",
         "torch.cuda.memory_usage",
         "torch.cuda.memory._dump_snapshot",
-        "torch.cuda.memory._free_mutex",
         "torch.cuda.memory._get_current_allocator",
         "torch.cuda.memory._host_allocator",
         "torch.cuda.memory._record_memory_history_impl",
@@ -3092,11 +3124,11 @@ class FunctionIdSet:
     added to the graph and what will cause a graph break.
     """
 
-    function_ids: Optional[set[int]] = None
-    function_names: Optional[dict[int, str]] = None
+    function_ids: set[int] | None = None
+    function_names: dict[int, str] | None = None
 
     def __init__(
-        self, lazy_initializer: Callable[[], Union[dict[int, str], set[int]]]
+        self, lazy_initializer: Callable[[], dict[int, str] | set[int]]
     ) -> None:
         self.lazy_initializer = lazy_initializer
 
@@ -3168,12 +3200,6 @@ def _builtin_function_ids() -> dict[int, str]:
             if not k.startswith("_") and callable(v)
         }
     )
-    rv.update(
-        {
-            id(cast): "typing.cast",
-            id(copy.deepcopy): "copy.deepcopy",
-        }
-    )
     return rv
 
 
@@ -3317,6 +3343,13 @@ def is_numpy_type_info(obj: Any) -> bool:
     linecache,
 )
 
+# Builtin modules that should be skipped at the top-level (PEP 523 frame
+# evaluation) but inlined when called from code dynamo is already tracing.
+# For example, copy.deepcopy should be inlined when the user calls it inside
+# a compiled function, but copy module frames should be skipped when they
+# appear as top-level frames (e.g. called internally by autograd).
+BUILTIN_INLINE_WHEN_CALLED: set[str] = set()
+
 # third party libraries skiplist is defined by str, because users may not use these libraries.
 # we should use lazy import & skip in the future.
 THIRDPARTY_SKIPLIST = (
@@ -3354,7 +3387,7 @@ def _strip_init_py(s: str) -> str:
     return _as_posix_path(s)
 
 
-def _module_dir(m: types.ModuleType) -> Optional[str]:
+def _module_dir(m: types.ModuleType) -> str | None:
     # Protect against a module not exporting __file__ - this can happen for
     # frozen modules, for example.
     file = getattr(m, "__file__", None)
@@ -3629,6 +3662,8 @@ def get_mod_skiplist() -> set[str]:
 ]
 SKIP_DIRS.extend(map(_as_posix_path, filter(None, map(_module_dir, BUILTIN_SKIPLIST))))
 
+BUILTIN_INLINE_WHEN_CALLED.update(filter(None, (_module_dir(copy),)))
+
 SKIP_DIRS_RE = re.compile(r"match nothing^")
 
 # Skip fbcode paths(including torch.package paths) containing
@@ -3696,26 +3731,29 @@ def add(import_name: str) -> None:
 @dataclasses.dataclass
 class SkipResult:
     skipped: bool
-    reason: Optional[str]
+    reason: str | None
 
 
-def check_file(filename: Optional[str], is_inlined_call: bool = False) -> SkipResult:
+def check_file(filename: str | None, is_inlined_call: bool = False) -> SkipResult:
     """Should skip this file?"""
     if filename is None:
-        return SkipResult(True, "filename is None")
+        return SkipResult(
+            True, "cannot determine source file (likely a C extension or builtin)"
+        )
     filename = _as_posix_path(filename)
     if filename in FORCE_SKIP_FILES:
-        return SkipResult(True, "FORCE_SKIP_FILES")
+        return SkipResult(True, f"file is force-skipped ({filename})")
 
-    if any(filename.startswith(d) for d in get_legacy_mod_inlinelist()):
-        return SkipResult(
-            False,
-            "LEGACY_MOD_INLINELIST",
-        )
+    for d in get_legacy_mod_inlinelist():
+        if filename.startswith(d):
+            return SkipResult(False, f"file matches LEGACY_MOD_INLINELIST ({d})")
     if is_inlined_call and is_torch_inline_allowed(filename):
+        return SkipResult(False, f"file matches MOD_INLINELIST ({filename})")
+    if is_inlined_call and any(
+        filename.startswith(d) for d in BUILTIN_INLINE_WHEN_CALLED
+    ):
         return SkipResult(
-            False,
-            "MOD_INLINELIST",
+            False, f"file matches BUILTIN_INLINE_WHEN_CALLED ({filename})"
         )
     if (
         is_fbcode()
@@ -3723,10 +3761,7 @@ def check_file(filename: Optional[str], is_inlined_call: bool = False) -> SkipRe
         and bool(FBCODE_SKIP_DIRS_RE.match(filename))
         and not bool(FBCODE_INLINE_FILES_IN_SKIPPED_DIRS_RE.match(filename))
     ):
-        return SkipResult(
-            True,
-            "FBCODE_SKIP_DIRS",
-        )
+        return SkipResult(True, "file matches FBCODE_SKIP_DIRS")
 
     if (
         is_fbcode()
@@ -3735,7 +3770,7 @@ def check_file(filename: Optional[str], is_inlined_call: bool = False) -> SkipRe
         and bool(FBCODE_SKIP_TORCHREC_DIRS_RE.match(filename))
         and not bool(FBCODE_INLINE_FILES_IN_SKIPPED_DIRS_RE.match(filename))
     ):
-        return SkipResult(True, "FBCODE_SKIP_TORCHREC_DIRS")
+        return SkipResult(True, "file matches FBCODE_SKIP_TORCHREC_DIRS")
 
     unittest_dir = _module_dir(unittest)
     if (
@@ -3743,22 +3778,24 @@ def check_file(filename: Optional[str], is_inlined_call: bool = False) -> SkipRe
         and filename.startswith(unittest_dir)
         and not torch._dynamo.config.enable_trace_unittest
     ):
-        return SkipResult(True, "unittest")
+        return SkipResult(True, "file is in unittest directory")
 
     if bool(SKIP_DIRS_RE.match(filename)):
-        return SkipResult(True, "SKIP_DIRS")
+        matched_dir = next((d for d in SKIP_DIRS if filename.startswith(d)), filename)
+        return SkipResult(True, f"file is under skip directory ({matched_dir})")
 
-    if any(filename.startswith(d) for d in get_mod_skiplist()):
-        return SkipResult(True, "MOD_SKIPLIST")
+    for d in get_mod_skiplist():
+        if filename.startswith(d):
+            return SkipResult(True, f"file matches MOD_SKIPLIST ({d})")
     return SkipResult(False, "inlined by default")
 
 
 @dataclasses.dataclass
 class FunctionInfo:
-    py_obj: Optional[object]
-    name: Optional[str]
+    py_obj: object | None
+    name: str | None
     filename: str
-    code: Optional[types.CodeType]
+    code: types.CodeType | None
 
 
 """
@@ -3797,6 +3834,7 @@ def f3(x, y):
 _force_inline_flag = False
 
 
+# pyrefly: ignore [deprecated]
 @contextlib.contextmanager
 def _force_inline() -> Iterator[None]:
     """
@@ -3806,7 +3844,7 @@ def _force_inline() -> Iterator[None]:
     When active, check_verbose() will skip all inline/skip decision logic and
     always return SkipResult(False, ...), meaning functions will be inlined.
 
-    See _make_inlined() in higher_order_ops.py which uses this to ensure that
+    See _make_inlined() in utils.py which uses this to ensure that
     a python function is fully traced to produce the needed variable trackers.
     """
     global _force_inline_flag
@@ -3818,13 +3856,40 @@ def _force_inline() -> Iterator[None]:
         _force_inline_flag = old_val
 
 
-def check_verbose(obj: Any, is_inlined_call: bool = False) -> SkipResult:
+def check_verbose(
+    obj: Any, is_inlined_call: bool = False, frame: Any | None = None
+) -> SkipResult:
     if _force_inline_flag:
         return SkipResult(
             False,
             "don't skip because we're inside _force_inline() context",
         )
 
+    # For eval frame callback (not inlined calls), allow tracing inbuilt
+    # nn.Module.forward methods when the module has hooks. Any hook can cause
+    # a graph break (via @torch._dynamo.disable, print, unsupported ops, etc.),
+    # which skips the entire _call_impl frame. By allowing forward to be traced,
+    # Dynamo can capture the module's operations in a new graph after the break.
+    if (
+        not is_inlined_call
+        and frame is not None
+        and isinstance(obj, types.CodeType)
+        and obj.co_name == "forward"
+    ):
+        from .utils import nnmodule_has_hooks
+
+        module = frame.f_locals.get("self")
+        if (
+            module is not None
+            and isinstance(module, torch.nn.Module)
+            and module.__class__.__module__.startswith(("torch.nn.", "torch.ao."))
+            and nnmodule_has_hooks(module, check_forward_hooks=True)
+        ):
+            return SkipResult(
+                False,
+                "inbuilt nn.Module.forward allowed - module has hooks",
+            )
+
     if isinstance(
         obj,
         (
@@ -3856,6 +3921,12 @@ def check_verbose(obj: Any, is_inlined_call: bool = False) -> SkipResult:
         assert filename is not None
         fi = FunctionInfo(obj, None, filename, None)
 
+    # typing.cast is a polyfilled no-op, but unlike C builtins it has a code
+    # object that PEP 523 can intercept as a standalone frame after a graph
+    # break. Skip it at the top level to avoid installing unnecessary guards.
+    if fi.code is not None and fi.code is typing.cast.__code__:
+        return SkipResult(True, "typing.cast is a no-op, skip at top level")
+
     # Consulte the central trace rules defined in torch._dynamo.trace_rules.
     reasons: set[str] = set()
     rule = lookup_inner(fi.py_obj, fi.name, fi.filename, is_inlined_call, reasons)
@@ -3868,25 +3939,33 @@ def check_verbose(obj: Any, is_inlined_call: bool = False) -> SkipResult:
             PolyfilledFunctionVariable,
         ),
     ):
-        return SkipResult(
-            False,
-            f"inlined according trace_rules.lookup {reasons.pop()}",
-        )
+        return SkipResult(False, reasons.pop())
     elif issubclass(rule, TorchInGraphFunctionVariable):
-        return SkipResult(
-            False,
-            f"registered in torch_obj_rule {reasons.pop()}",
-        )
+        return SkipResult(False, reasons.pop())
     else:
         assert rule == SkipFunctionVariable, rule
-        return SkipResult(
-            True,
-            f"skipped according trace_rules.lookup {reasons.pop()}",
-        )
+        return SkipResult(True, reasons.pop())
 
 
-def check(obj: Any, is_inlined_call: bool = False) -> bool:
-    return check_verbose(obj, is_inlined_call).skipped
+def check(obj: Any, is_inlined_call: bool = False, frame: Any | None = None) -> bool:
+    return check_verbose(obj, is_inlined_call, frame).skipped
+
+
+def get_skip_reason(obj: Any) -> str:
+    """Compute a descriptive skip reason for a callable. Only called on graph break."""
+    if is_callable_disallowed(obj):
+        return _disallowed_callable_ids.get_name(id(obj), repr(obj))
+
+    filename = getfile(obj)
+    if filename is not None:
+        skip_result = check_file(filename)
+        if skip_result.reason is not None:
+            return skip_result.reason
+
+    module = getattr(obj, "__module__", None) or ""
+    return (
+        f"cannot determine source file for {module} (likely a C extension or builtin)"
+    )
 
 
 # skip common third party libs
@@ -3901,7 +3980,7 @@ def is_torch_inline_allowed(filename: str) -> bool:
 
 
 @functools.cache
-def dynamo_dir() -> Optional[str]:
+def dynamo_dir() -> str | None:
     import torch._dynamo
 
     return _module_dir(torch._dynamo)
@@ -3919,8 +3998,14 @@ def is_torch(filename: str) -> bool:
 Main entry point for looking up the trace rule (the Dynamo variable) for a given callable object.
 """
 
+BUILTIN_CALLABLES = {
+    dict: DictBuiltinVariable,
+    iter: IterBuiltinVariable,
+    list: ListBuiltinVariable,
+}
+
 
-def lookup_callable(obj: Callable[..., Any]) -> Optional[type[VariableTracker]]:
+def lookup_callable(obj: Callable[..., Any]) -> type[VariableTracker] | None:
     if not hashable(obj):
         return None
     # Custom allow/disallow in graph takes precedence over the general lookup.
@@ -3930,6 +4015,8 @@ def lookup_callable(obj: Callable[..., Any]) -> Optional[type[VariableTracker]]:
         return TorchInGraphFunctionVariable
     if is_polyfilled_callable(obj):
         return PolyfilledFunctionVariable
+    if obj in BUILTIN_CALLABLES:
+        return BUILTIN_CALLABLES[obj]
     if is_builtin_callable(obj):
         return BuiltinVariable
     return None
@@ -3941,18 +4028,18 @@ def lookup_callable(obj: Callable[..., Any]) -> Optional[type[VariableTracker]]:
 """
 
 
-def lookup(obj: Any) -> Optional[type[VariableTracker]]:
+def lookup(obj: Any) -> type[VariableTracker] | None:
     return lookup_inner(obj)
 
 
 # also takes config.dont_skip_tracing into account
 def lookup_inner(
     obj: Any,
-    name: Optional[str] = None,
-    filename: Optional[str] = None,
+    name: str | None = None,
+    filename: str | None = None,
     is_direct_call: bool = True,
-    reasons: Union[None, set[str]] = None,
-) -> Optional[type[VariableTracker]]:
+    reasons: None | set[str] = None,
+) -> type[VariableTracker] | None:
     result = _lookup_inner(
         obj,
         name=name,
@@ -3986,11 +4073,11 @@ def lookup_inner(
 
 def _lookup_inner(
     obj: Any,
-    name: Optional[str] = None,
-    filename: Optional[str] = None,
+    name: str | None = None,
+    filename: str | None = None,
     is_direct_call: bool = True,
-    reasons: Optional[set[str]] = None,
-) -> Optional[type[VariableTracker]]:
+    reasons: set[str] | None = None,
+) -> type[VariableTracker] | None:
     # Step 1: lookup obj's tracing rule in `torch_name_rule_map`.
     # The rules defined in `torch_name_rule_map` mainly includes two parts:
     # - Manually defined rules for any functions.
diff --git a/torch/_dynamo/types.py b/torch/_dynamo/types.py
index 8236d8b229be2..bbf76c0ee0550 100644
--- a/torch/_dynamo/types.py
+++ b/torch/_dynamo/types.py
@@ -14,7 +14,7 @@
 import dataclasses
 import types
 from collections.abc import Callable
-from typing import Any, NamedTuple, Optional, Protocol, Union
+from typing import Any, NamedTuple, Protocol
 
 # CacheEntry has a `guard_manager` field for the guard, and a `code` field for the code object.
 from torch._C._dynamo.eval_frame import (
@@ -55,9 +55,9 @@ class GuardFn(Protocol):
     code_parts: list[str]
     verbose_code_parts: list[str]
     global_scope: dict[str, object]
-    guard_fail_fn: Optional[Callable[[GuardFail], None]]
-    cache_entry: Optional[CacheEntry]
-    extra_state: Optional[ExtraState]
+    guard_fail_fn: Callable[[GuardFail], None] | None
+    cache_entry: CacheEntry | None
+    extra_state: ExtraState | None
 
     # maps locals of user function to bool
     def __call__(self, f_locals: dict[str, object]) -> bool: ...
@@ -82,7 +82,7 @@ class ConvertFrameReturn:
     )
     # also apply frame_exec strategy to future frames with same code
     apply_to_code: bool = True
-    guarded_code: Optional[GuardedCode] = None
+    guarded_code: GuardedCode | None = None
 
 
 def wrap_guarded_code(guarded_code: GuardedCode) -> ConvertFrameReturn:
@@ -96,12 +96,12 @@ class DynamoCallbackFn(Protocol):
     def __call__(
         self,
         frame: DynamoFrameType,
-        cache_entry: Optional[CacheEntry],
+        cache_entry: CacheEntry | None,
         frame_state: FrameState,
     ) -> ConvertFrameReturn: ...
 
 
-DynamoCallback = Union[DynamoCallbackFn, None, bool]
+DynamoCallback = DynamoCallbackFn | None | bool
 
 
 class DynamoGuardHook(Protocol):
@@ -137,4 +137,4 @@ def __call__(self, record: Any) -> None: ...
 class BytecodeHook(Protocol):
     def __call__(
         self, code: types.CodeType, new_code: types.CodeType
-    ) -> Optional[types.CodeType]: ...
+    ) -> types.CodeType | None: ...
diff --git a/torch/_dynamo/utils.py b/torch/_dynamo/utils.py
index 210e289be2ff1..366a5a9e6b6ce 100644
--- a/torch/_dynamo/utils.py
+++ b/torch/_dynamo/utils.py
@@ -56,12 +56,10 @@
     Generic,
     Literal,
     NoReturn,
-    Optional,
     overload,
     TypeAlias,
     TypeGuard,
     TypeVar,
-    Union,
 )
 from typing_extensions import ParamSpec, TypeIs
 
@@ -162,6 +160,16 @@
 
 
 T = TypeVar("T")
+T1 = TypeVar("T1")
+T2 = TypeVar("T2")
+T3 = TypeVar("T3")
+T4 = TypeVar("T4")
+T5 = TypeVar("T5")
+T6 = TypeVar("T6")
+T7 = TypeVar("T7")
+T8 = TypeVar("T8")
+T9 = TypeVar("T9")
+T10 = TypeVar("T10")
 R = TypeVar("R")
 _P = ParamSpec("_P")
 
@@ -237,8 +245,8 @@ def log(cls) -> None:
 
 
 def tabulate(
-    rows: Union[list[tuple[str, Any]], list[list[Any]]],
-    headers: Union[tuple[str, ...], list[str]],
+    rows: list[tuple[str, Any]] | list[list[Any]],
+    headers: tuple[str, ...] | list[str],
 ) -> str:
     try:
         import tabulate
@@ -267,7 +275,7 @@ def reset_frame_count() -> None:
     curr_frame = 0
 
 
-_recompile_user_contexts: Optional[list[Callable[[], str]]] = None
+_recompile_user_contexts: list[Callable[[], str]] | None = None
 
 
 def register_hook_for_recompile_user_context(hook: Callable[[], str]) -> None:
@@ -284,7 +292,7 @@ def register_hook_for_recompile_user_context(hook: Callable[[], str]) -> None:
     _recompile_user_contexts.append(hook)
 
 
-def get_hook_for_recompile_user_context() -> Optional[list[Callable[[], str]]]:
+def get_hook_for_recompile_user_context() -> list[Callable[[], str]] | None:
     return _recompile_user_contexts
 
 
@@ -307,13 +315,11 @@ def increment_op_count(cnt: int) -> None:
 def calculate_time_spent() -> dict[str, float]:
     total_by_key = {}
     for phase, timing in cumulative_time_spent_ns.items():
-        # pyrefly: ignore [unsupported-operation]
         total_by_key[phase] = timing / 1e9
 
     total_by_key["total_wall_time"] = total_by_key.get(
         "entire_frame_compile", 0
     ) + total_by_key.get("entire_backward_compile", 0)
-    # pyrefly: ignore [bad-return]
     return total_by_key
 
 
@@ -415,7 +421,7 @@ class CompileEventLogger:
     def log_instant_event(
         event_name: str,
         metadata: dict[str, Any],
-        time_ns: Optional[int] = None,
+        time_ns: int | None = None,
         log_level: CompileEventLogLevel = CompileEventLogLevel.CHROMIUM,
     ) -> None:
         if time_ns is None:
@@ -648,7 +654,7 @@ def compilation_metric(overwrite: bool = False, **metadata: object) -> None:
 
     @staticmethod
     def instant(
-        event_name: str, metadata: dict[str, Any], time_ns: Optional[int] = None
+        event_name: str, metadata: dict[str, Any], time_ns: int | None = None
     ) -> None:
         """
         Log an instant event to chromium logs with name <event_name> at time <time_ns>. The `args` field in
@@ -714,14 +720,14 @@ def compile_time_record_function(name: str) -> Generator[Any, None, None]:
 def dynamo_timed(
     key: str,
     # TODO(masneral): Deprecate this param.
-    phase_name: Optional[str] = None,
+    phase_name: str | None = None,
     log_pt2_compile_event: bool = False,
-    metadata: Optional[dict[str, object]] = None,
-    dynamo_compile_column_us: Optional[str] = None,
-    compile_id: Optional[CompileId] = None,
-    is_backward: Optional[bool] = None,
+    metadata: dict[str, object] | None = None,
+    dynamo_compile_column_us: str | None = None,
+    compile_id: CompileId | None = None,
+    is_backward: bool | None = None,
     log_waitcounter: bool = False,
-    waitcounter_name_override: Optional[str] = None,
+    waitcounter_name_override: str | None = None,
 ) -> Generator[Any, None, None]:
     """
     dynamo_timed is a context manager
@@ -868,7 +874,7 @@ def compile_times(
 
 def compile_times(  # type: ignore[misc]
     repr: str = "str", aggregate: bool = False
-) -> Union[str, None, tuple[list[str], list[str]]]:
+) -> str | None | tuple[list[str], list[str]]:
     """
     Get metrics about torchdynamo frontend/backend compilation times.
 
@@ -932,7 +938,7 @@ def __init__(self, maxsize: int = 4096) -> None:
     def reset(self) -> None:
         self.set: OrderedDict[Any, Any] = OrderedDict()
 
-    def add(self, key: Union[str, tuple[object, object]]) -> bool:
+    def add(self, key: str | tuple[object, object]) -> bool:
         if key in self.set:
             self.set.move_to_end(key, last=True)
             if not config.verbose:
@@ -1096,8 +1102,95 @@ def istype(
 ) -> TypeIs[T]: ...
 
 
+# This can be simplified once TypeVarTuple objects can be expanded into TypeIs.
 @overload
-def istype(obj: object, allowed_types: Iterable[type]) -> bool: ...
+def istype(
+    obj: object, allowed_types: tuple[type[T1], type[T2]]
+) -> TypeIs[T1 | T2]: ...
+
+
+@overload
+def istype(
+    obj: object, allowed_types: tuple[type[T1], type[T2], type[T3]]
+) -> TypeIs[T1 | T2 | T3]: ...
+
+
+@overload
+def istype(
+    obj: object, allowed_types: tuple[type[T1], type[T2], type[T3], type[T4]]
+) -> TypeIs[T1 | T2 | T3 | T4]: ...
+
+
+@overload
+def istype(
+    obj: object, allowed_types: tuple[type[T1], type[T2], type[T3], type[T4], type[T5]]
+) -> TypeIs[T1 | T2 | T3 | T4 | T5]: ...
+
+
+@overload
+def istype(
+    obj: object,
+    allowed_types: tuple[type[T1], type[T2], type[T3], type[T4], type[T5], type[T6]],
+) -> TypeIs[T1 | T2 | T3 | T4 | T5 | T6]: ...
+
+
+@overload
+def istype(
+    obj: object,
+    allowed_types: tuple[
+        type[T1], type[T2], type[T3], type[T4], type[T5], type[T6], type[T7]
+    ],
+) -> TypeIs[T1 | T2 | T3 | T4 | T5 | T6 | T7]: ...
+
+
+@overload
+def istype(
+    obj: object,
+    allowed_types: tuple[
+        type[T1], type[T2], type[T3], type[T4], type[T5], type[T6], type[T7], type[T8]
+    ],
+) -> TypeIs[T1 | T2 | T3 | T4 | T5 | T6 | T7 | T8]: ...
+
+
+@overload
+def istype(
+    obj: object,
+    allowed_types: tuple[
+        type[T1],
+        type[T2],
+        type[T3],
+        type[T4],
+        type[T5],
+        type[T6],
+        type[T7],
+        type[T8],
+        type[T9],
+    ],
+) -> TypeIs[T1 | T2 | T3 | T4 | T5 | T6 | T7 | T8 | T9]: ...
+
+
+@overload
+def istype(
+    obj: object,
+    allowed_types: tuple[
+        type[T1],
+        type[T2],
+        type[T3],
+        type[T4],
+        type[T5],
+        type[T6],
+        type[T7],
+        type[T8],
+        type[T9],
+        type[T10],
+    ],
+) -> TypeIs[T1 | T2 | T3 | T4 | T5 | T6 | T7 | T8 | T9 | T10]: ...
+
+
+@overload
+def istype(
+    obj: object, allowed_types: tuple[type, ...] | list[type] | set[type]
+) -> bool: ...
 
 
 def istype(obj: object, allowed_types: Any) -> bool:
@@ -1126,11 +1219,11 @@ def istype(obj: object, allowed_types: Any) -> bool:
 def get_inputs_devices(
     inputs: collections.abc.Sequence[object],
     model: torch.fx.GraphModule,
-) -> list[Optional[torch.device]]:
+) -> list[torch.device | None]:
     all_inputs = pytree.tree_flatten(inputs)[0] + [
         node.meta["val"] for node in list(model.graph.nodes) if "val" in node.meta
     ]
-    devices: list[Optional[torch.device]] = list(
+    devices: list[torch.device | None] = list(
         OrderedSet([i.device for i in all_inputs if hasattr(i, "device")])
     )
     return [
@@ -1213,17 +1306,17 @@ def is_lru_cache_wrapped_function(
     )
 
 
-_FuncTypes: TypeAlias = Union[
-    types.FunctionType,
-    types.BuiltinFunctionType,
-    types.MethodDescriptorType,
-    types.WrapperDescriptorType,
-]
+_FuncTypes: TypeAlias = (
+    types.FunctionType
+    | types.BuiltinFunctionType
+    | types.MethodDescriptorType
+    | types.WrapperDescriptorType
+)
 
 
 def is_function_or_wrapper(
     value: Any,
-) -> TypeIs[Union[_FuncTypes, torch._ops.OpOverloadPacket, torch._ops.OpOverload]]:
+) -> TypeIs[_FuncTypes | torch._ops.OpOverloadPacket | torch._ops.OpOverload]:
     return is_function(value) or isinstance(
         value, (torch._ops.OpOverloadPacket, torch._ops.OpOverload)
     )
@@ -1266,13 +1359,11 @@ def is_function(
 def is_wrapper_or_member_descriptor(
     value: Any,
 ) -> TypeIs[
-    Union[
-        types.GetSetDescriptorType,
-        types.MethodDescriptorType,
-        types.WrapperDescriptorType,
-        types.MemberDescriptorType,
-        types.MethodWrapperType,
-    ]
+    types.GetSetDescriptorType
+    | types.MethodDescriptorType
+    | types.WrapperDescriptorType
+    | types.MemberDescriptorType
+    | types.MethodWrapperType
 ]:
     return isinstance(
         value,
@@ -1295,7 +1386,7 @@ def unwrap_if_wrapper(fn: Any) -> Any:
     return unwrap_with_attr_name_if_wrapper(fn)[0]
 
 
-def unwrap_with_attr_name_if_wrapper(fn: Any) -> tuple[Any, Optional[str]]:
+def unwrap_with_attr_name_if_wrapper(fn: Any) -> tuple[Any, str | None]:
     # TODO(anijain2305) - Investigate if we can get rid of this function
     # unpack @torch._dynamo.optimize()(fn) wrapped function
     if is_function(fn) and inspect.getattr_static(fn, "_torchdynamo_inline", False):
@@ -1362,13 +1453,13 @@ def proxy_args_kwargs(args: Any, kwargs: Any) -> tuple[tuple[Any, ...], dict[str
         )
 
 
-def to_int_ms(v: Optional[float]) -> Optional[int]:
+def to_int_ms(v: float | None) -> int | None:
     return None if v is None else int(v * 1000)
 
 
 # float64 timestamp has a quarter microsecond precision in 2024, so while
 # this is suboptimal we shouldn't meaningfully lose precision
-def to_int_us(v: Optional[float]) -> Optional[int]:
+def to_int_us(v: float | None) -> int | None:
     return None if v is None else int(v * 1_000_000)
 
 
@@ -1379,113 +1470,113 @@ def to_int_us(v: Optional[float]) -> Optional[int]:
 
 @dataclasses.dataclass
 class CompilationMetrics:
-    compile_id: Optional[str] = None
-    frame_key: Optional[str] = None
-    co_name: Optional[str] = None
-    co_filename: Optional[str] = None
-    co_firstlineno: Optional[int] = None
-    cache_size: Optional[int] = None
-    accumulated_cache_size: Optional[int] = None
-    guard_count: Optional[int] = None
-    shape_env_guard_count: Optional[int] = None
-    graph_op_count: Optional[int] = None
-    graph_node_count: Optional[int] = None
-    graph_input_count: Optional[int] = None
-    start_time: Optional[float] = None
-    entire_frame_compile_time_s: Optional[float] = None
-    backend_compile_time_s: Optional[float] = None
-    inductor_compile_time_s: Optional[float] = None
-    code_gen_time_s: Optional[float] = None
-    fail_type: Optional[str] = None
-    fail_reason: Optional[str] = None
-    fail_user_frame_filename: Optional[str] = None
-    fail_user_frame_lineno: Optional[int] = None
-    non_compliant_ops: Optional[set[str]] = None
-    compliant_custom_ops: Optional[set[str]] = None
-    restart_reasons: Optional[set[str]] = None
-    dynamo_time_before_restart_s: Optional[float] = None
-    stack_trace: Optional[list[str]] = None
-    exception_stack_trace: Optional[list[str]] = None
-    graph_node_shapes: Optional[str] = None
+    compile_id: str | None = None
+    frame_key: str | None = None
+    co_name: str | None = None
+    co_filename: str | None = None
+    co_firstlineno: int | None = None
+    cache_size: int | None = None
+    accumulated_cache_size: int | None = None
+    guard_count: int | None = None
+    shape_env_guard_count: int | None = None
+    graph_op_count: int | None = None
+    graph_node_count: int | None = None
+    graph_input_count: int | None = None
+    start_time: float | None = None
+    entire_frame_compile_time_s: float | None = None
+    backend_compile_time_s: float | None = None
+    inductor_compile_time_s: float | None = None
+    code_gen_time_s: float | None = None
+    fail_type: str | None = None
+    fail_reason: str | None = None
+    fail_user_frame_filename: str | None = None
+    fail_user_frame_lineno: int | None = None
+    non_compliant_ops: set[str] | None = None
+    compliant_custom_ops: set[str] | None = None
+    restart_reasons: set[str] | None = None
+    dynamo_time_before_restart_s: float | None = None
+    stack_trace: list[str] | None = None
+    exception_stack_trace: list[str] | None = None
+    graph_node_shapes: str | None = None
     # Sometimes, we will finish analyzing a frame but conclude we don't want
     # to install any guarded code.  True means we actually decided to install
     # a compiled frame
-    has_guarded_code: Optional[bool] = None
-    remote_cache_time_saved_s: Optional[float] = None
-    structured_logging_overhead_s: Optional[float] = None
-    config_suppress_errors: Optional[bool] = None
-    config_inline_inbuilt_nn_modules: Optional[bool] = None
-    specialize_float: Optional[bool] = None
-    dynamo_config: Optional[str] = None
-    compiler_config: Optional[str] = None
-    is_forward: Optional[bool] = None
-    num_triton_bundles: Optional[int] = None
-    remote_fx_graph_cache_get_time_ms: Optional[int] = None
-    remote_fx_graph_cache_put_time_ms: Optional[int] = None
-    start_time_us: Optional[int] = None
-    duration_us: Optional[int] = None
-    dynamo_cumulative_compile_time_us: Optional[int] = None
-    aot_autograd_cumulative_compile_time_us: Optional[int] = None
-    inductor_cumulative_compile_time_us: Optional[int] = None
-    inductor_code_gen_cumulative_compile_time_us: Optional[int] = None
-    triton_compile_time_us: Optional[int] = None
-    runtime_cudagraphify_time_us: Optional[int] = None
-    runtime_triton_autotune_time_us: Optional[int] = None
-    dynamo_compile_time_before_restart_us: Optional[int] = None
-    distributed_ephemeral_timeout_us: Optional[int] = None
-    structured_logging_overhead_us: Optional[int] = None
-    remote_fx_graph_cache_get_time_us: Optional[int] = None
-    remote_fx_graph_cache_put_time_us: Optional[int] = None
-    backward_cumulative_compile_time_us: Optional[int] = None
-    end_time_us: Optional[int] = None
-    pre_grad_pass_time_us: Optional[int] = None
-    post_grad_pass_time_us: Optional[int] = None
-    joint_graph_pass_time_us: Optional[int] = None
+    has_guarded_code: bool | None = None
+    remote_cache_time_saved_s: float | None = None
+    structured_logging_overhead_s: float | None = None
+    config_suppress_errors: bool | None = None
+    config_inline_inbuilt_nn_modules: bool | None = None
+    specialize_float: bool | None = None
+    dynamo_config: str | None = None
+    compiler_config: str | None = None
+    is_forward: bool | None = None
+    num_triton_bundles: int | None = None
+    remote_fx_graph_cache_get_time_ms: int | None = None
+    remote_fx_graph_cache_put_time_ms: int | None = None
+    start_time_us: int | None = None
+    duration_us: int | None = None
+    dynamo_cumulative_compile_time_us: int | None = None
+    aot_autograd_cumulative_compile_time_us: int | None = None
+    inductor_cumulative_compile_time_us: int | None = None
+    inductor_code_gen_cumulative_compile_time_us: int | None = None
+    triton_compile_time_us: int | None = None
+    runtime_cudagraphify_time_us: int | None = None
+    runtime_triton_autotune_time_us: int | None = None
+    dynamo_compile_time_before_restart_us: int | None = None
+    distributed_ephemeral_timeout_us: int | None = None
+    structured_logging_overhead_us: int | None = None
+    remote_fx_graph_cache_get_time_us: int | None = None
+    remote_fx_graph_cache_put_time_us: int | None = None
+    backward_cumulative_compile_time_us: int | None = None
+    end_time_us: int | None = None
+    pre_grad_pass_time_us: int | None = None
+    post_grad_pass_time_us: int | None = None
+    joint_graph_pass_time_us: int | None = None
     log_format_version: int = LOG_FORMAT_VERSION
-    inductor_config: Optional[str] = None
-    remote_cache_version: Optional[int] = None
-    inductor_fx_remote_cache_hit_count: Optional[int] = 0
-    inductor_fx_remote_cache_miss_count: Optional[int] = 0
-    inductor_fx_remote_cache_backend_type: Optional[str] = None
-    inductor_fx_remote_cache_hit_keys: Optional[str] = None
-    inductor_fx_remote_cache_miss_keys: Optional[str] = None
-    inductor_fx_local_cache_hit_count: Optional[int] = 0
-    inductor_fx_local_cache_miss_count: Optional[int] = 0
-    aotautograd_remote_cache_hit_count: Optional[int] = 0
-    aotautograd_remote_cache_miss_count: Optional[int] = 0
-    aotautograd_local_cache_hit_count: Optional[int] = 0
-    aotautograd_local_cache_miss_count: Optional[int] = 0
-    cuda_version: Optional[str] = None
-    triton_version: Optional[str] = None
-    feature_usage: Optional[dict[str, bool]] = None
-    compile_time_autotune_time_us: Optional[int] = None
-    is_runtime: Optional[bool] = False
-    gc_time_us: Optional[int] = None
-    tensorify_float_attempt: Optional[bool] = None
-    tensorify_float_success: Optional[bool] = None
-    tensorify_float_failure: Optional[set[str]] = None
-    guard_latency_us: Optional[float] = None
-    recompile_reason: Optional[str] = None
-    num_graph_breaks: Optional[int] = None
-    triton_kernel_compile_times_us: Optional[str] = None
-    ir_count: Optional[int] = None
-    cudagraph_skip_reason: Optional[str] = None
-    python_version: Optional[str] = None
-    pgo_put_remote_code_state_time_us: Optional[int] = None
-    pgo_get_remote_code_state_time_us: Optional[int] = None
+    inductor_config: str | None = None
+    remote_cache_version: int | None = None
+    inductor_fx_remote_cache_hit_count: int | None = 0
+    inductor_fx_remote_cache_miss_count: int | None = 0
+    inductor_fx_remote_cache_backend_type: str | None = None
+    inductor_fx_remote_cache_hit_keys: str | None = None
+    inductor_fx_remote_cache_miss_keys: str | None = None
+    inductor_fx_local_cache_hit_count: int | None = 0
+    inductor_fx_local_cache_miss_count: int | None = 0
+    aotautograd_remote_cache_hit_count: int | None = 0
+    aotautograd_remote_cache_miss_count: int | None = 0
+    aotautograd_local_cache_hit_count: int | None = 0
+    aotautograd_local_cache_miss_count: int | None = 0
+    cuda_version: str | None = None
+    triton_version: str | None = None
+    feature_usage: dict[str, bool] | None = None
+    compile_time_autotune_time_us: int | None = None
+    is_runtime: bool | None = False
+    gc_time_us: int | None = None
+    tensorify_float_attempt: bool | None = None
+    tensorify_float_success: bool | None = None
+    tensorify_float_failure: set[str] | None = None
+    guard_latency_us: float | None = None
+    recompile_reason: str | None = None
+    num_graph_breaks: int | None = None
+    triton_kernel_compile_times_us: str | None = None
+    ir_count: int | None = None
+    cudagraph_skip_reason: str | None = None
+    python_version: str | None = None
+    pgo_put_remote_code_state_time_us: int | None = None
+    pgo_get_remote_code_state_time_us: int | None = None
     # The number of elements within parameters. This is classically what people
     # think of when they think of parameters in a ML model.
-    param_numel: Optional[int] = None
+    param_numel: int | None = None
     # The number of elements counted by bytes - i.e. a float32 is 4 bytes
     # per element.
-    param_bytes: Optional[int] = None
+    param_bytes: int | None = None
     # The number of parameters counted by fields. This is mostly a proxy for
     # the number of distinct type of params.
-    param_count: Optional[int] = None
-    recompile_user_contexts: Optional[set[str]] = None
-    inline_inbuilt_nn_modules_candidate: Optional[bool] = False
-    pytorch_version: Optional[str] = None
-    inductor_provenance: Optional[set[str]] = None
+    param_count: int | None = None
+    recompile_user_contexts: set[str] | None = None
+    inline_inbuilt_nn_modules_candidate: bool | None = False
+    pytorch_version: str | None = None
+    inductor_provenance: set[str] | None = None
 
     @classmethod
     def create(cls, metrics: dict[str, Any]) -> CompilationMetrics:
@@ -1495,13 +1586,13 @@ def create(cls, metrics: dict[str, Any]) -> CompilationMetrics:
         we transform some fields to comma-separated strings for scuba logging.
         """
 
-        def us_to_s(metric: Optional[int]) -> Optional[float]:
+        def us_to_s(metric: int | None) -> float | None:
             return metric / 1e6 if metric is not None else None
 
-        def us_to_ms(metric: Optional[int]) -> Optional[int]:
+        def us_to_ms(metric: int | None) -> int | None:
             return metric // 1000 if metric is not None else None
 
-        def collection_to_str(metric: Optional[Any]) -> Optional[str]:
+        def collection_to_str(metric: Any | None) -> str | None:
             def safe_str(item: Any) -> str:
                 try:
                     return str(item)
@@ -1516,7 +1607,7 @@ def safe_str(item: Any) -> str:
 
             return ",".join(safe_str(item) for item in sorted(metric))
 
-        def collection_to_json_str(metric: Optional[Any]) -> Optional[str]:
+        def collection_to_json_str(metric: Any | None) -> str | None:
             if metric is None:
                 return None
             try:
@@ -1629,7 +1720,7 @@ def add_compilation_metrics_to_chromium(c: CompilationMetrics) -> None:
     )
 
 
-def _get_dynamo_config_for_logging() -> Optional[str]:
+def _get_dynamo_config_for_logging() -> str | None:
     def clean_for_json(d: dict[str, Any]) -> dict[str, Any]:
         blocklist = {
             "TYPE_CHECKING",
@@ -1649,6 +1740,7 @@ def clean_for_json(d: dict[str, Any]) -> dict[str, Any]:
             "_autograd_backward_strict_mode_banned_ops",
             "reorderable_logging_functions",
             "ignore_logger_methods",
+            "ignore_logging_functions",
             "traceable_tensor_subclasses",
             "nontraceable_tensor_subclasses",
             "_custom_ops_profile",
@@ -1664,7 +1756,7 @@ def clean_for_json(d: dict[str, Any]) -> dict[str, Any]:
     return json.dumps(config_dict, sort_keys=True)
 
 
-def _compiler_config_for_logging() -> Optional[str]:
+def _compiler_config_for_logging() -> str | None:
     def clean_for_json(d: dict[str, Any]) -> dict[str, Any]:
         blocklist = {
             "TYPE_CHECKING",
@@ -1688,7 +1780,7 @@ def clean_for_json(d: dict[str, Any]) -> dict[str, Any]:
     return json.dumps(config_dict, sort_keys=True)
 
 
-def _scrubbed_inductor_config_for_logging() -> Optional[str]:
+def _scrubbed_inductor_config_for_logging() -> str | None:
     """
     Method to parse and scrub uninteresting configs from inductor config
     """
@@ -1740,8 +1832,8 @@ def record_compilation_metrics(
     start_time_ns: int,
     end_time_ns: int,
     metrics: dict[str, Any],
-    exc_type: Optional[type[BaseException]],
-    exc_value: Optional[BaseException],
+    exc_type: type[BaseException] | None,
+    exc_value: BaseException | None,
 ) -> None:
     if torch._inductor.utils.should_use_remote_fx_graph_cache():
         try:
@@ -1773,7 +1865,7 @@ def record_compilation_metrics(
         ),
         "dynamo_config": _get_dynamo_config_for_logging(),
         "config_suppress_errors": config.suppress_errors,
-        "config_inline_inbuilt_nn_modules": config.inline_inbuilt_nn_modules,
+        "config_inline_inbuilt_nn_modules": True,
         "inductor_config": _scrubbed_inductor_config_for_logging(),
         "compiler_config": _compiler_config_for_logging(),
         "cuda_version": torch.version.cuda,
@@ -1855,7 +1947,7 @@ def get_stack(self) -> list[str]:
             self.tls.stack = []
             return self.tls.stack
 
-    def get_outermost_event(self) -> Optional[str]:
+    def get_outermost_event(self) -> str | None:
         """
         Get the outermost event name (i.e. the longest running event)
         or None if the stack is empty.
@@ -1973,7 +2065,7 @@ def log_event_start(
         time_ns: int,
         metadata: dict[str, Any],
         log_pt2_compile_event: bool = False,
-        compile_id: Optional[CompileId] = None,
+        compile_id: CompileId | None = None,
     ) -> None:
         """
         Logs the start of a single event.
@@ -2030,7 +2122,7 @@ def log_event_end(
         metadata: dict[str, Any],
         start_time_ns: int,
         log_pt2_compile_event: bool,
-        compile_id: Optional[CompileId] = None,
+        compile_id: CompileId | None = None,
     ) -> None:
         """
         Logs the end of a single event. This function should only be
@@ -2108,7 +2200,7 @@ def _log_timed_event(
         event_name: str,
         time_ns: int,
         phase: str,
-        metadata: Optional[dict[str, Any]] = None,
+        metadata: dict[str, Any] | None = None,
     ) -> dict[str, Any]:
         """
         Logs a timed event in chromium format. See log_event_start, log_event_end, etc.
@@ -2136,7 +2228,7 @@ def log_instant_event(
         self,
         event_name: str,
         time_ns: int,
-        metadata: Optional[dict[str, Any]] = None,
+        metadata: dict[str, Any] | None = None,
         # By default, an instant event isn't logged internally, only to structured logging.
         log_pt2_compile_event: bool = False,
     ) -> None:
@@ -2175,7 +2267,7 @@ def log_instant_event(
             )
 
 
-CHROMIUM_EVENT_LOG: Optional[ChromiumEventLogger] = None
+CHROMIUM_EVENT_LOG: ChromiumEventLogger | None = None
 
 
 def get_chromium_event_logger() -> ChromiumEventLogger:
@@ -2288,9 +2380,7 @@ def copy_dynamo_tensor_attributes(src: torch.Tensor, dst: torch.Tensor) -> None:
     _copy_dynamo_attr(src, dst, "_dynamo_weak_dynamic_indices")
 
 
-def clone_input(
-    x: torch.Tensor, *, dtype: Optional[torch.dtype] = None
-) -> torch.Tensor:
+def clone_input(x: torch.Tensor, *, dtype: torch.dtype | None = None) -> torch.Tensor:
     """copy while preserving strides"""
     # TODO: this is questionable
     if is_fake(x):
@@ -2368,7 +2458,7 @@ def torch_clone(x: torch.Tensor) -> torch.Tensor:
 
 @overload
 def clone_inputs(
-    example_inputs: dict[str, Union[T, tuple[T, ...]]],
+    example_inputs: dict[str, T | tuple[T, ...]],
 ) -> dict[str, list[T]]: ...
 
 
@@ -2377,7 +2467,7 @@ def clone_inputs(example_inputs: Sequence[T]) -> list[T]: ...
 
 
 def clone_inputs(example_inputs: Any) -> Any:
-    res: Union[dict[str, Any], list[Any]]
+    res: dict[str, Any] | list[Any]
     if type(example_inputs) is dict:
         res = dict(example_inputs)
         for key, value in res.items():
@@ -2439,13 +2529,10 @@ def preserve_rng_state() -> Generator[None, None, None]:
 def is_jit_model(
     model0: Any,
 ) -> TypeIs[
-    Union[
-        torch.jit._trace.TopLevelTracedModule,
-        torch.jit._script.RecursiveScriptModule,
-        # pyrefly: ignore [invalid-param-spec]
-        torch.jit.ScriptFunction[Any, Any],
-        torch.jit.ScriptModule,
-    ]
+    torch.jit._trace.TopLevelTracedModule
+    | torch.jit._script.RecursiveScriptModule
+    | torch.jit.ScriptFunction[Any, Any]
+    | torch.jit.ScriptModule
 ]:
     return isinstance(
         model0,
@@ -2476,7 +2563,7 @@ def torchscript(model: Any, example_inputs: Any, verbose: bool = False) -> Any:
     return None
 
 
-def getfile(obj: Any) -> Optional[str]:
+def getfile(obj: Any) -> str | None:
     try:
         return inspect.getfile(obj)
     except (TypeError, OSError):
@@ -2620,6 +2707,10 @@ def rot_n_helper(n: int) -> Callable[..., Any]:
     torch.iinfo,
     torch.nn.attention.SDPBackend,
     torch.cuda._CudaDeviceProperties,
+    # Pytree key types (frozen dataclasses used in tree_map_with_path)
+    torch.utils._pytree.SequenceKey,
+    torch.utils._pytree.MappingKey,
+    torch.utils._pytree.GetAttrKey,
 }
 
 if has_triton_package():
@@ -2663,7 +2754,7 @@ def common_constants() -> set[int]:
     }
 
 
-def is_torch_sym(value: Any) -> TypeGuard[Union[torch.SymBool, torch.SymInt]]:
+def is_torch_sym(value: Any) -> TypeGuard[torch.SymBool | torch.SymInt]:
     return isinstance(value, (torch.SymBool, torch.SymInt)) and not isinstance(
         value.node, torch.nested._internal.nested_int.NestedIntNode
     )
@@ -2818,22 +2909,37 @@ def builtin_dict_keys(d: dict[K, V]) -> KeysView[K]:
     return dict.keys(d)
 
 
-def get_items_from_dict(obj: dict[K, V]) -> Iterable[tuple[K, Union[V, Any]]]:
+def get_items_from_dict(obj: dict[K, V]) -> Iterable[tuple[K, V | Any]]:
     # Get items without calling the user defined __getitem__ or keys method.
     assert isinstance(obj, dict)
     if istype(obj, (dict, OrderedDict)):
         return obj.items()
     elif isinstance(obj, OrderedDict):
-        # pyrefly: ignore [bad-argument-type]
         return [(k, OrderedDict.__getitem__(obj, k)) for k in OrderedDict.keys(obj)]
     else:
-        # pyrefly: ignore [bad-argument-type]
         return [(k, dict.__getitem__(obj, k)) for k in dict.keys(obj)]
 
 
+def enumerate_items_with_dict_position(
+    obj: dict[K, V],
+) -> Iterable[tuple[int, K, V | Any]]:
+    """Enumerate dict items yielding (dict_keys_position, key, value).
+
+    For OrderedDicts where move_to_end/prepend has been used, the OrderedDict
+    iteration order can differ from dict.keys() order.  We iterate in
+    OrderedDict order (correct execution semantics) but return each key's
+    dict.keys() position so that ConstDictKeySource indices stay consistent
+    with PyDict_Next / C++ DictGuardManager.
+    """
+    items = get_items_from_dict(obj)
+    if isinstance(obj, OrderedDict):
+        key_to_pos = {k: i for i, k in enumerate(dict.keys(obj))}
+        return ((key_to_pos[k], k, v) for k, v in items)
+    return ((i, k, v) for i, (k, v) in enumerate(items))
+
+
 def nn_module_new(cls: Any) -> Any:
     obj = object_new(cls)
-    # pyrefly: ignore [bad-argument-type]
     torch.nn.Module.__init__(obj)
     return obj
 
@@ -2854,6 +2960,29 @@ def dataclass_fields(cls: Any) -> Any:
 iter_next = next
 
 
+def normalize_count_iter(count_iter: Iterator[Any]) -> tuple[Any, Any]:
+    try:
+        _, args = count_iter.__reduce__()
+    except TypeError:
+        # Python 3.14 no longer pickles itertools.count, so fall back to the
+        # repr and only recover literal arguments. Non-literal arguments still
+        # fall back to user-defined handling via the NotImplemented sentinel.
+        import ast
+
+        count_repr = repr(count_iter)
+        if not count_repr.startswith("count(") or not count_repr.endswith(")"):
+            return (NotImplemented, NotImplemented)
+        try:
+            args = ast.literal_eval(f"({count_repr[6:-1]},)")
+        except (SyntaxError, ValueError):
+            return (NotImplemented, NotImplemented)
+        if not isinstance(args, tuple) or not 1 <= len(args) <= 2:
+            return (NotImplemented, NotImplemented)
+    if len(args) == 1:
+        return (args[0], 1)
+    return (args[0], args[1])
+
+
 def normalize_range_iter(range_iter: Any) -> tuple[int, int, int]:
     _, (range_obj,), maybe_idx = range_iter.__reduce__()
     # In 3.12+, `maybe_idx` could be None, and `range_obj.start` would've been
@@ -2877,12 +3006,10 @@ def to_subclass(t: Any, cls: type) -> Any:
 
 @torch.fx.wrap
 def dict_keys_getitem(d: dict[Any, Any], n: int) -> Any:
-    # Call dict(d) to prevent calling overridden __iter__/keys
-    dict_class = dict
-    if isinstance(d, OrderedDict):
-        dict_class = OrderedDict
+    # Use dict.keys() to match the iteration order of PyDict_Next used by
+    # the C++ DictGuardManager and by ConstDictKeySource._name_template.
     # pyrefly: ignore [bad-argument-type]
-    return next(itertools.islice(dict_class.keys(d), n, n + 1))
+    return next(itertools.islice(dict.keys(d), n, n + 1))
 
 
 def set_getitem(s: set[T], n: int) -> T:
@@ -2947,7 +3074,6 @@ def raise_args_mismatch(
     actual: str = "",
 ) -> None:
     from torch._dynamo.exc import raise_observed_exception
-    from torch._dynamo.variables import ConstantVariable
 
     msg_str = (
         f"wrong number of arguments or keyword arguments for {name}() call.\n"
@@ -2958,7 +3084,7 @@ def raise_args_mismatch(
     raise_observed_exception(
         TypeError,
         tx,
-        args=[ConstantVariable(msg_str)],
+        args=[msg_str],
     )
 
 
@@ -2984,7 +3110,7 @@ def iter_contains(
         # Match of Tensor means match of FakeTensor
         search = _get_fake_tensor(search)
 
-    found: Optional[VariableTracker] = None
+    found: VariableTracker | None = None
     for x in items:
         if must_check_tensor_id:
             if x.is_tensor():
@@ -3009,7 +3135,7 @@ def iter_contains(
 
 def key_is_id(
     k: Any,
-) -> TypeIs[Union[torch.Tensor, torch.nn.Module, MethodWrapperType]]:
+) -> TypeIs[torch.Tensor | torch.nn.Module | MethodWrapperType]:
     """Returns whether it indexes dictionaries using its id"""
     return isinstance(k, (torch.Tensor, torch.nn.Module, MethodWrapperType))
 
@@ -3571,9 +3697,7 @@ def visit(n: torch.fx.Node) -> Any:
     return torch.fx.node.map_arg(nodes, visit)
 
 
-def get_concrete_sizes_from_symints(
-    msg: str, fake_mode: Optional[FakeTensorMode]
-) -> str:
+def get_concrete_sizes_from_symints(msg: str, fake_mode: FakeTensorMode | None) -> str:
     """
     Replace symbolic size expressions (like 's0', 's94') in error messages
     with their concrete runtime values for better readability.
@@ -3799,11 +3923,11 @@ def _get_fake_value_impl(
         elif isinstance(
             cause, torch.fx.experimental.symbolic_shapes.GuardOnDataDependentSymNode
         ):
-            raise UserError(  # noqa: B904
+            raise UserError(
                 UserErrorType.CONSTRAINT_VIOLATION,
                 str(cause),
                 case_name="constrain_as_size_example",
-            )
+            ) from cause
         elif isinstance(cause, ValueRangeError):
             raise UserError(UserErrorType.CONSTRAINT_VIOLATION, e.args[0]) from e
         elif isinstance(cause, TypeError) and "argument" in str(cause):
@@ -3847,7 +3971,7 @@ def _get_fake_value_impl(
 _current_node = threading.local()
 
 
-def get_current_node() -> Optional[torch.fx.Node]:
+def get_current_node() -> torch.fx.Node | None:
     return getattr(_current_node, "value", None)
 
 
@@ -4062,7 +4186,7 @@ def class_has_getattribute(cls: type) -> bool:
 
 def get_custom_getattr(
     value: Any, ignore_nn_module_getattr: bool = False
-) -> Optional[Any]:
+) -> Any | None:
     try:
         getattr_fn = inspect.getattr_static(type(value), "__getattr__")
     except AttributeError:
@@ -4090,10 +4214,10 @@ def tensor_static_reason_to_message(reason: TensorStaticReason) -> str:
 
 
 def tensor_always_has_static_shape(
-    tensor: Union[torch.Tensor, Any],
+    tensor: torch.Tensor | Any,
     is_tensor: bool,
     tensor_source: Source,
-) -> tuple[bool, Optional[TensorStaticReason]]:
+) -> tuple[bool, TensorStaticReason | None]:
     """
     Given a tensor, source, and is_tensor flag, determine if a shape should be static.
 
@@ -4321,18 +4445,16 @@ def defake(x: Any) -> Any:
     size: torch._prims_common.ShapeType
     stride: torch._prims_common.StrideType
     if x._has_symbolic_sizes_strides:
-        size = []
-        for s in x.size():
-            if isinstance(s, torch.SymInt):
-                size.append(s.node.shape_env.size_hint(s.node.expr))
-            else:
-                size.append(s)
-        stride = []
-        for s in x.stride():
-            if isinstance(s, torch.SymInt):
-                stride.append(s.node.shape_env.size_hint(s.node.expr))
-            else:
-                stride.append(s)
+        # optimization_hint is appropriate here because defake only needs a
+        # plausible concrete shape to allocate a real tensor; it does not need
+        # to install guards. For unbacked symbols the heuristic fallback is fine.
+        size = [
+            torch.fx.experimental.symbolic_shapes.optimization_hint(s) for s in x.size()
+        ]
+        stride = [
+            torch.fx.experimental.symbolic_shapes.optimization_hint(s)
+            for s in x.stride()
+        ]
     else:
         size = x.size()
         stride = x.stride()
@@ -4434,7 +4556,7 @@ class _Anchors:
     right_start_offset: int
 
 
-def _extract_anchors_from_expr(segment: str) -> Optional[_Anchors]:
+def _extract_anchors_from_expr(segment: str) -> _Anchors | None:
     """
     Given source code `segment` corresponding to a bytecode
     instruction, determine:
@@ -4645,7 +4767,7 @@ def get_instruction_source_311(code: types.CodeType, inst: Instruction) -> str:
         num_spaces = len(last_line) - len(last_line.lstrip())
         markers.append(" " * num_spaces + "~" * (end_offset - num_spaces))
 
-    anchors: Optional[_Anchors] = None
+    anchors: _Anchors | None = None
     try:
         anchors = _extract_anchors_from_expr(segment)
     except AssertionError:
@@ -4722,6 +4844,7 @@ def is_tensor_base_attr_getter(value: Any) -> bool:
     return (
         isinstance(value, types.MethodWrapperType)
         and value.__name__ == "__get__"
+        and hasattr(value.__self__, "__objclass__")
         and value.__self__.__objclass__ is torch._C._TensorBase  # type: ignore[attr-defined]
     )
 
@@ -4859,7 +4982,7 @@ def __init__(
         self.unflatten_fn = unflatten_fn
 
     def forward(self, *args: Any) -> Any:
-        # pyrefly: ignore [annotation-mismatch, redefinition]
+        # pyrefly: ignore [redefinition]
         args: list[Any] = list(args)
         return self.gm(*self.unflatten_fn(args))
 
@@ -5182,7 +5305,7 @@ def maybe_disable_inference_mode_for_fake_prop() -> Generator[None, None, None]:
         yield
 
 
-def is_node_meta_valid(node: Optional[torch.fx.Node]) -> bool:
+def is_node_meta_valid(node: torch.fx.Node | None) -> bool:
     return node is None or "example_value" in node.meta or "val" in node.meta
 
 
@@ -5217,7 +5340,7 @@ def record_pregraph_bytecode_exit(cm: AbstractContextManager[None]) -> None:
 
 # Returns a set of code objects present traced in the current TracingContext, or None
 # if there is no current TracingContext.
-def get_traced_code() -> Optional[list[CodeType]]:
+def get_traced_code() -> list[CodeType] | None:
     from torch._guards import TracingContext
 
     return TracingContext.get_traced_code()
@@ -5239,3 +5362,20 @@ def raise_on_overridden_hash(obj: Any, vt: VariableTracker) -> None:
                 *graph_break_hints.SUPPORTABLE,
             ],
         )
+
+
+def _make_inlined(
+    tx: InstructionTranslator, f: Callable[..., Any]
+) -> Callable[..., VariableTracker]:
+    assert callable(f), "Expect f to be a python callable."
+
+    def inline_call(
+        *args: VariableTracker, **kwargs: VariableTracker
+    ) -> VariableTracker:
+        from torch._dynamo.trace_rules import _force_inline
+        from torch._dynamo.variables.functions import UserFunctionVariable
+
+        with _force_inline():
+            return UserFunctionVariable(f).call_function(tx, args, kwargs)  # type: ignore[arg-type]
+
+    return inline_call
diff --git a/torch/_dynamo/variables/__init__.py b/torch/_dynamo/variables/__init__.py
index b73570410442b..c8529febd62a9 100644
--- a/torch/_dynamo/variables/__init__.py
+++ b/torch/_dynamo/variables/__init__.py
@@ -17,12 +17,19 @@
 """
 
 from .base import VariableTracker
-from .builtin import BuiltinVariable
-from .constant import CONSTANT_VARIABLE_NONE, ConstantVariable, EnumVariable
+from .builtin import (
+    BaseBuiltinVariable,
+    BuiltinVariable,
+    DictBuiltinVariable,
+    IterBuiltinVariable,
+    ListBuiltinVariable,
+)
+from .constant import ConstantVariable
 from .ctx_manager import (
     CatchWarningsCtxManagerVariable,
     ContextWrappingVariable,
     CUDADeviceVariable,
+    CudagraphOverrideVariable,
     DisabledSavedTensorsHooksVariable,
     DualLevelContextManager,
     DynamoConfigPatchVariable,
@@ -46,23 +53,19 @@
     ConstDictVariable,
     DefaultDictVariable,
     DictItemsVariable,
-    DictKeySetVariable,
-    FrozensetVariable,
+    DunderDictVariable,
     MappingProxyVariable,
     NNModuleHooksDictVariable,
-    OrderedSetClassVariable,
-    OrderedSetVariable,
-    SetVariable,
 )
-from .distributed import BackwardHookVariable, DistributedVariable, PlacementVariable
+from .distributed import BackwardHookVariable, DistributedVariable
 from .functions import (
+    BaseUserFunctionVariable,
     BuiltinMethodVariable,
     CollectionsNamedTupleFunction,
     CreateTMADescriptorExperimentalVariable,
     CreateTMADescriptorStableVariable,
     FunctionDecoratedByContextlibContextManagerVariable,
     FunctoolsPartialVariable,
-    FunctoolsWrapsVariable,
     InspectSignatureVariable,
     LocalGeneratorFunctionVariable,
     LocalGeneratorObjectVariable,
@@ -74,7 +77,7 @@
     SparseTensorCreationSkipVariable,
     TMADescriptorExperimentalVariable,
     TMADescriptorStableVariable,
-    TritonSetAllocatorSkipVariable,
+    TritonSetAllocatorVariable,
     UserFunctionVariable,
     UserMethodVariable,
     WrapperUserFunctionVariable,
@@ -101,7 +104,6 @@
     BaseListVariable,
     ListIteratorVariable,
     ListVariable,
-    NamedTupleVariable,
     RangeVariable,
     SliceVariable,
     TupleIteratorVariable,
@@ -138,7 +140,19 @@
 )
 from .optimizer import OptimizerVariable
 from .sdpa import SDPAParamsVariable
-from .streams import EventVariable, StreamContextVariable, StreamVariable
+from .sets import (
+    DictKeySetVariable,
+    FrozensetVariable,
+    OrderedSetClassVariable,
+    OrderedSetVariable,
+    SetVariable,
+)
+from .streams import (
+    CudaStreamVariable,
+    EventVariable,
+    StreamContextVariable,
+    StreamVariable,
+)
 from .tensor import (
     DataPtrVariable,
     FakeItemVariable,
@@ -153,8 +167,11 @@
     FrozenDataClassVariable,
     InspectVariable,
     MutableMappingVariable,
+    NamedTupleVariable,
     RemovableHandleVariable,
+    StructSequenceVariable,
     UserDefinedClassVariable,
+    UserDefinedConstantVariable,
     UserDefinedDictVariable,
     UserDefinedExceptionClassVariable,
     UserDefinedExceptionObjectVariable,
@@ -162,6 +179,7 @@
     UserDefinedObjectVariable,
     UserDefinedSetVariable,
     UserDefinedTupleVariable,
+    UserDefinedVariable,
 )
 
 
@@ -169,36 +187,38 @@
     "AutogradFunctionContextVariable",
     "AutogradFunctionVariable",
     "BackwardHookVariable",
+    "BaseBuiltinVariable",
     "BaseListVariable",
     "BuiltinVariable",
     "CatchWarningsCtxManagerVariable",
-    "CONSTANT_VARIABLE_NONE",
     "ConstantVariable",
     "ConstDictVariable",
+    "DictBuiltinVariable",
     "ContextWrappingVariable",
     "CountIteratorVariable",
     "CreateTMADescriptorExperimentalVariable",
     "CreateTMADescriptorStableVariable",
     "CUDADeviceVariable",
+    "CudagraphOverrideVariable",
     "DataPtrVariable",
     "DefaultDictVariable",
     "DeletedVariable",
     "DictKeySetVariable",
     "DynamoConfigPatchVariable",
-    "EnumVariable",
     "FakeItemVariable",
     "GetAttrVariable",
     "GradModeVariable",
     "InspectSignatureVariable",
     "InspectVariable",
+    "IterBuiltinVariable",
     "IteratorVariable",
     "ItertoolsVariable",
     "LambdaVariable",
     "LazyConstantVariable",
     "LazyVariableTracker",
+    "ListBuiltinVariable",
     "ListIteratorVariable",
     "ListVariable",
-    "NamedTupleVariable",
     "NestedUserFunctionVariable",
     "CellVariable",
     "NewGlobalVariable",
@@ -206,7 +226,6 @@
     "NumpyNdarrayVariable",
     "NumpyVariable",
     "OptimizerVariable",
-    "PlacementVariable",
     "PolyfilledFunctionVariable",
     "PythonModuleVariable",
     "RangeVariable",
@@ -232,6 +251,8 @@
     "UntypedStorageVariable",
     "UserDefinedClassVariable",
     "UserDefinedTupleVariable",
+    "NamedTupleVariable",
+    "StructSequenceVariable",
     "UserDefinedObjectVariable",
     "UserFunctionVariable",
     "UserMethodVariable",
diff --git a/torch/_dynamo/variables/base.py b/torch/_dynamo/variables/base.py
index 5f6dd6630bdf7..d6101f8c8b5af 100644
--- a/torch/_dynamo/variables/base.py
+++ b/torch/_dynamo/variables/base.py
@@ -13,16 +13,15 @@
 computations.
 """
 
+from __future__ import annotations
+
 import collections
 import functools
 import logging
 from collections.abc import Callable, ItemsView, KeysView, Sequence, ValuesView
 from contextvars import ContextVar
 from enum import Enum
-from typing import Any, NoReturn, Optional, TYPE_CHECKING
-
-from torch._guards import Guard
-from torch.fx.proxy import Node
+from typing import Any, NoReturn, TYPE_CHECKING
 
 from .. import graph_break_hints, variables
 from ..current_scope_id import current_scope_id
@@ -33,7 +32,11 @@
 
 
 if TYPE_CHECKING:
+    from torch._guards import Guard
+    from torch.fx.proxy import Node
+
     from ..codegen import PyCodegen
+    from ..side_effects import SideEffects
     from ..symbolic_convert import InstructionTranslator
     from .constant import ConstantVariable
     from .functions import UserFunctionVariable
@@ -43,8 +46,11 @@
 
 # Tracks active method calls on VariableTracker instances to detect self-referential
 # calls (e.g., as_python_constant on a list that contains itself). Maps
-# (id(instance), method_name) tuples to track which calls are in progress.
-_vt_active_calls: ContextVar[set[tuple[int, str]] | None] = ContextVar(
+# (id(instance), id(original_method)) tuples to track which calls are in progress.
+# We use id(original_method) rather than the method name string so that super()
+# delegation within a class hierarchy (e.g. TorchScriptObjectVariable.as_python_constant
+# calling UserDefinedObjectVariable.as_python_constant) is not a false positive.
+_vt_active_calls: ContextVar[set[tuple[int, int]] | None] = ContextVar(
     "_vt_active_calls", default=None
 )
 
@@ -195,7 +201,7 @@ class AttributeMutationNew(AttributeMutation):
     the Python world.
     """
 
-    def __init__(self, cls_source: Optional[Source] = None) -> None:
+    def __init__(self, cls_source: Source | None = None) -> None:
         super().__init__(SourceType.New)
         self.cls_source = cls_source
 
@@ -216,13 +222,17 @@ def is_side_effect_safe(m: MutationType) -> bool:
     return m.scope == scope_id
 
 
+class NO_SUCH_SUBOBJ:
+    """Sentinel indicating no concrete Python object is available."""
+
+
 # This helps users of `as_python_constant` to catch unimplemented error with
 # more information; it inherits `NotImplementedError` for backward
 # compatibility reasons.
 class AsPythonConstantNotImplementedError(NotImplementedError):
-    vt: "VariableTracker"
+    vt: VariableTracker
 
-    def __init__(self, vt: "VariableTracker", msg: str | None = None) -> None:
+    def __init__(self, vt: VariableTracker, msg: str | None = None) -> None:
         msg = f"{vt} is not a constant" if msg is None else msg
         super().__init__(msg)
         self.vt = vt
@@ -275,6 +285,11 @@ class VariableTracker(metaclass=VariableTrackerMeta):
     Prefer the factory function VariableTracker.build() over VariableTracker.__init__().
     """
 
+    # The CPython type(s) this VT is designed to represent.
+    # Single type or tuple of types. None means no static CPython type mapping
+    # (e.g., dynamic types like UserDefinedObjectVariable, or Dynamo-internal VTs).
+    _cpython_type: type | tuple[type, ...] | None = None
+
     # fields to leave unmodified in apply()
     _nonvar_fields = {
         "value",
@@ -285,7 +300,7 @@ class VariableTracker(metaclass=VariableTrackerMeta):
         "user_code_variable_name",
     }
 
-    def clone(self, **kwargs: Any) -> "VariableTracker":
+    def clone(self, **kwargs: Any) -> VariableTracker:
         """Shallow copy with some (optional) changes"""
         args = dict(self.__dict__)
         args.update(kwargs)
@@ -294,12 +309,17 @@ def clone(self, **kwargs: Any) -> "VariableTracker":
     @classmethod
     def visit(
         cls,
-        fn: Callable[["VariableTracker"], None],
+        fn: Callable[[VariableTracker], None],
         value: Any,
-        cache: Optional[dict[int, Any]] = None,
+        cache: dict[int, Any] | None = None,
+        side_effects: SideEffects | None = None,
     ) -> None:
         """
-        Walk value and call fn on all the VariableTracker instances
+        Walk value and call fn on all the VariableTracker instances.
+
+        When side_effects is provided, also walks attributes stored in
+        store_attr_mutations (e.g. dataclass fields set during tracing
+        that aren't in the VT's __dict__).
         """
         if cache is None:
             cache = {}
@@ -317,13 +337,16 @@ def visit(
             nonvars = value._nonvar_fields
             for key, subvalue in value.__dict__.items():
                 if key not in nonvars:
-                    cls.visit(fn, subvalue, cache)
+                    cls.visit(fn, subvalue, cache, side_effects)
+            if side_effects is not None and value in side_effects.store_attr_mutations:
+                for attr_vt in side_effects.store_attr_mutations[value].values():
+                    cls.visit(fn, attr_vt, cache, side_effects)
         elif istype(value, (list, tuple)):
             for subvalue in value:
-                cls.visit(fn, subvalue, cache)
+                cls.visit(fn, subvalue, cache, side_effects)
         elif istype(value, (dict, collections.OrderedDict)):
             for subvalue in value.values():
-                cls.visit(fn, subvalue, cache)
+                cls.visit(fn, subvalue, cache, side_effects)
 
     def __repr__(self) -> str:
         return f"{self.__class__.__name__}()"
@@ -393,6 +416,14 @@ def is_python_constant(self) -> bool:
         except NotImplementedError:
             return False
 
+    def bool_impl(self, tx: InstructionTranslator) -> VariableTracker | None:
+        # Mirrors CPython's tp_as_number->nb_bool slot.
+        # https://github.com/python/cpython/blob/c09ccd9c429/Objects/object.c#L2135-L2158
+        #
+        # Returns None when the type has no nb_bool, causing generic_bool to
+        # fall through to length check, then truthy default.
+        return None
+
     def is_constant_match(self, *values: Any) -> bool:
         """
         Check if this variable is a python constant matching one of the given values.
@@ -415,7 +446,7 @@ def make_guard(self, fn: Callable[..., Any]) -> Guard:
 
     # TODO[@lucaskabela] - change this type to `InstructionTranslatorBase`
     # and cascade that (large blast radius)
-    def const_getattr(self, tx: "InstructionTranslator", name: str) -> Any:
+    def const_getattr(self, tx: InstructionTranslator, name: str) -> Any:
         """getattr(self, name) returning a python constant"""
         raise NotImplementedError
 
@@ -427,7 +458,7 @@ def is_tensor(self) -> bool:
         """Return True for TensorVariable instances"""
         return False
 
-    def var_getattr(self, tx: "InstructionTranslator", name: str) -> "VariableTracker":
+    def var_getattr(self, tx: InstructionTranslator, name: str) -> VariableTracker:
         """getattr(self, name) returning a new variable"""
         value = self.const_getattr(tx, name)
         if not variables.ConstantVariable.is_literal(value):
@@ -449,7 +480,7 @@ def is_proxy(self) -> bool:
     def as_proxy(self) -> Any:
         raise NotImplementedError(str(self))
 
-    def maybe_fx_node(self) -> Optional[Node]:
+    def maybe_fx_node(self) -> Node | None:
         try:
             proxy = self.as_proxy()
             import torch.fx
@@ -464,7 +495,7 @@ def _contains_self_reference(self) -> bool:
         """Check if this variable references itself (directly or indirectly)."""
         found_self = False
 
-        def check(vt: "VariableTracker") -> None:
+        def check(vt: VariableTracker) -> None:
             nonlocal found_self
             if vt is self:
                 found_self = True
@@ -476,13 +507,13 @@ def check(vt: "VariableTracker") -> None:
 
         return found_self
 
-    def reconstruct(self, codegen: "PyCodegen") -> None:
+    def reconstruct(self, codegen: PyCodegen) -> None:
         raise NotImplementedError
 
-    def unpack_var_sequence(self, tx: Any) -> list["VariableTracker"]:
+    def unpack_var_sequence(self, tx: Any) -> list[VariableTracker]:
         raise NotImplementedError
 
-    def force_unpack_var_sequence(self, tx: Any) -> list["VariableTracker"]:
+    def force_unpack_var_sequence(self, tx: Any) -> list[VariableTracker]:
         # like unpack_var_sequence, but should only be used when it is
         # safe to eagerly (vs. lazily) unpack this variable.
         # e.g. map(f, x) is normally evaluated lazily but sometimes
@@ -506,15 +537,15 @@ def has_force_unpack_var_sequence(self, tx: Any) -> bool:
     # Only use when it is safe to eagerly unpack this variable (like force_unpack_var_sequence).
     # INVARIANT: variable must satisfy has_force_unpack_var_sequence() == True!
     def force_apply_to_var_sequence(
-        self, tx: Any, fn: Callable[["VariableTracker"], Any]
+        self, tx: Any, fn: Callable[[VariableTracker], Any]
     ) -> None:
         assert self.has_force_unpack_var_sequence(tx)
         for v in self.unpack_var_sequence(tx):
             fn(v)
 
     def call_obj_hasattr(
-        self, tx: "InstructionTranslator", name: str
-    ) -> "ConstantVariable":
+        self, tx: InstructionTranslator, name: str
+    ) -> ConstantVariable:
         unimplemented(
             gb_type="Unsupported hasattr call",
             context=f"call_obj_hasattr {self} {name}",
@@ -528,9 +559,9 @@ def call_obj_hasattr(
     def call_function(
         self,
         tx: Any,
-        args: Sequence["VariableTracker"],
-        kwargs: dict[str, "VariableTracker"],
-    ) -> "VariableTracker":
+        args: Sequence[VariableTracker],
+        kwargs: dict[str, VariableTracker],
+    ) -> VariableTracker:
         unimplemented(
             gb_type="Unsupported function call",
             context=f"call_function {self} {args} {kwargs}",
@@ -541,16 +572,59 @@ def call_function(
             ],
         )
 
+    def sq_length(self, tx: Any) -> VariableTracker:
+        """Called when sq_length is not implemented."""
+        raise_observed_exception(
+            TypeError,
+            tx,
+            args=[f"object of type '{self.python_type_name()}' has no len()"],
+        )
+
+    def mp_length(self, tx: Any) -> VariableTracker:
+        """Called when mp_length is not implemented."""
+        raise_observed_exception(
+            TypeError,
+            tx,
+            args=[f"object of type '{self.python_type_name()}' has no len()"],
+        )
+
+    def mp_subscript_impl(
+        self,
+        tx: InstructionTranslator,
+        key: VariableTracker,
+    ) -> VariableTracker:
+        # PyObject_GetItem: https://github.com/python/cpython/blob/62a6e898e01/Objects/abstract.c#L155-L206
+        # TODO: raise TypeError for non-subscriptable objects (blocked on
+        # branch 3 __class_getitem__ support for type objects).
+        unimplemented(
+            gb_type="missing_mp_subscript",
+            context=f"mp_subscript_impl not defined for {type(self).__name__}",
+            explanation=f"Dynamo does not yet support subscripting '{self.python_type_name()}'.",
+            hints=[*graph_break_hints.SUPPORTABLE],
+        )
+
     def call_method(
         self,
         tx: Any,
         name: str,
-        args: list["VariableTracker"],
-        kwargs: dict[str, "VariableTracker"],
-    ) -> "VariableTracker":
-        if name == "__len__" and self.has_unpack_var_sequence(tx):
-            assert not (args or kwargs)
-            return variables.ConstantVariable.create(len(self.unpack_var_sequence(tx)))
+        args: list[VariableTracker],
+        kwargs: dict[str, VariableTracker],
+    ) -> VariableTracker:
+        if name == "__getitem__":
+            if len(args) == 1 and not kwargs:
+                return self.mp_subscript_impl(tx, args[0])
+            from ..utils import raise_args_mismatch
+
+            raise_args_mismatch(
+                tx,
+                name,
+                "1 args and 0 kwargs",
+                f"{len(args)} args and {len(kwargs)} kwargs",
+            )
+        elif name == "__len__" and not (args or kwargs):
+            from .object_protocol import generic_len
+
+            return generic_len(tx, self)
         elif (
             name == "__getattr__"
             and len(args) == 1
@@ -558,6 +632,12 @@ def call_method(
             and not kwargs
         ):
             return self.var_getattr(tx, args[0].as_python_constant())
+        elif name == "__index__" and not args and not kwargs:
+            return self.nb_index_impl(tx)
+        elif name == "__int__" and not args and not kwargs:
+            return self.nb_int_impl(tx)
+        elif name == "__float__" and not args and not kwargs:
+            return self.nb_float_impl(tx)
         elif name in cmp_name_to_op_mapping and len(args) == 1 and not kwargs:
             other = args[0]
             if not isinstance(self, type(other)) and not (
@@ -595,8 +675,22 @@ def call_method(
                 raise_observed_exception(
                     type(e),
                     tx,
-                    args=[list(map(variables.ConstantVariable.create, e.args))],
+                    args=list(e.args),
                 )
+        # __reduce_ex__ is a C builtin (object.__reduce_ex__) that Dynamo
+        # cannot trace into.  Constant-fold it for VTs backed by a real
+        # Python object so that copy.deepcopy can trace through.
+        if (
+            name == "__reduce_ex__"
+            and len(args) == 1
+            and not kwargs
+            and self.is_python_constant()
+        ):
+            protocol = args[0].as_python_constant()
+            return VariableTracker.build(
+                tx, self.as_python_constant().__reduce_ex__(protocol)
+            )
+
         hints = [
             f"Avoid calling `{self.python_type_name()}.{name}` in your code.",
             "Please report an issue to PyTorch.",
@@ -633,11 +727,11 @@ def call_method(
     def call_tree_map(
         self,
         tx: Any,
-        tree_map_fn: "UserFunctionVariable",
-        map_fn: "VariableTracker",
-        rest: Sequence["VariableTracker"],
-        tree_map_kwargs: dict[str, "VariableTracker"],
-    ) -> "VariableTracker":
+        tree_map_fn: UserFunctionVariable,
+        map_fn: VariableTracker,
+        rest: Sequence[VariableTracker],
+        tree_map_kwargs: dict[str, VariableTracker],
+    ) -> VariableTracker:
         """Performance optimization to implement optree.tree_map faster than tracing it"""
         is_leaf_var = tree_map_kwargs.get("is_leaf")
         if is_leaf_var is not None and not is_leaf_var.is_constant_none():
@@ -666,11 +760,11 @@ def call_tree_map(
     def call_tree_map_branch(
         self,
         tx: Any,
-        tree_map_fn: "UserFunctionVariable",
-        map_fn: "VariableTracker",
-        rest: Sequence["VariableTracker"],
-        tree_map_kwargs: dict[str, "VariableTracker"],
-    ) -> "VariableTracker":
+        tree_map_fn: UserFunctionVariable,
+        map_fn: VariableTracker,
+        rest: Sequence[VariableTracker],
+        tree_map_kwargs: dict[str, VariableTracker],
+    ) -> VariableTracker:
         """Emulate optree.tree_map without is_leaf/none_is_leaf checks (handled above)"""
         return self._tree_map_fallback(
             tx,
@@ -683,11 +777,11 @@ def call_tree_map_branch(
     def _tree_map_fallback(
         self,
         tx: Any,
-        tree_map_fn: "UserFunctionVariable",
-        map_fn: "VariableTracker",
-        rest: Sequence["VariableTracker"],
-        tree_map_kwargs: dict[str, "VariableTracker"],
-    ) -> "VariableTracker":
+        tree_map_fn: UserFunctionVariable,
+        map_fn: VariableTracker,
+        rest: Sequence[VariableTracker],
+        tree_map_kwargs: dict[str, VariableTracker],
+    ) -> VariableTracker:
         tree_map_fn_copy = tree_map_fn.clone()
         tree_map_fn_copy._maybe_call_tree_map_fastpath = lambda *args, **kwargs: None  # type: ignore[missing-attribute]
         log.debug(
@@ -702,14 +796,92 @@ def _tree_map_fallback(
             tree_map_kwargs,
         )
 
+    def call_tree_map_with_path(
+        self,
+        tx: Any,
+        tree_map_fn: UserFunctionVariable,
+        map_fn: VariableTracker,
+        rest: Sequence[VariableTracker],
+        tree_map_kwargs: dict[str, VariableTracker],
+        keypath: tuple[Any, ...],
+    ) -> VariableTracker:
+        """Performance optimization to implement tree_map_with_path faster than tracing it"""
+        is_leaf_var = tree_map_kwargs.get("is_leaf")
+        if is_leaf_var is not None and not is_leaf_var.is_constant_none():
+            pred_result = is_leaf_var.call_function(tx, [self], {})
+            try:
+                leaf_decision = pred_result.as_python_constant()
+            except NotImplementedError:
+                return self._tree_map_with_path_fallback(
+                    tx,
+                    tree_map_fn,
+                    map_fn,
+                    rest,
+                    tree_map_kwargs,
+                    keypath,
+                )
+            if leaf_decision:
+                keypath_var = variables.TupleVariable(
+                    [VariableTracker.build(tx, k) for k in keypath]
+                )
+                return map_fn.call_function(tx, [keypath_var, self, *rest], {})
+
+        return self.call_tree_map_with_path_branch(
+            tx,
+            tree_map_fn,
+            map_fn,
+            rest,
+            tree_map_kwargs,
+            keypath,
+        )
+
+    def call_tree_map_with_path_branch(
+        self,
+        tx: Any,
+        tree_map_fn: UserFunctionVariable,
+        map_fn: VariableTracker,
+        rest: Sequence[VariableTracker],
+        tree_map_kwargs: dict[str, VariableTracker],
+        keypath: tuple[Any, ...],
+    ) -> VariableTracker:
+        """Handle tree_map_with_path for leaf nodes (default behavior)"""
+        keypath_var = variables.TupleVariable(
+            [VariableTracker.build(tx, k) for k in keypath]
+        )
+        return map_fn.call_function(tx, [keypath_var, self, *rest], {})
+
+    def _tree_map_with_path_fallback(
+        self,
+        tx: Any,
+        tree_map_fn: UserFunctionVariable,
+        map_fn: VariableTracker,
+        rest: Sequence[VariableTracker],
+        tree_map_kwargs: dict[str, VariableTracker],
+        keypath: tuple[Any, ...],
+    ) -> VariableTracker:
+        tree_map_fn_copy = tree_map_fn.clone()
+        tree_map_fn_copy._maybe_call_tree_map_fastpath = lambda *args, **kwargs: None  # type: ignore[missing-attribute]
+        log.debug(
+            "tree_map_with_path fastpath fallback triggered for %s (rest=%s, kwargs=%s, keypath=%s)",
+            self,
+            rest,
+            tree_map_kwargs,
+            keypath,
+        )
+        return tree_map_fn_copy.call_function(
+            tx,
+            [map_fn, self, *rest],
+            tree_map_kwargs,
+        )
+
     def set_name_hint(self, name: str) -> None:
         pass
 
-    def realize(self) -> "VariableTracker":
+    def realize(self) -> VariableTracker:
         """Used by LazyVariableTracker to build the real VariableTracker"""
         return self
 
-    def unwrap(self) -> "VariableTracker":
+    def unwrap(self) -> VariableTracker:
         """Used by LazyVariableTracker to return the real VariableTracker if it already exists"""
         return self
 
@@ -717,7 +889,7 @@ def is_realized(self) -> bool:
         """Used by LazyVariableTracker to indicate an unrealized node"""
         return True
 
-    def next_variable(self, tx: Any) -> "VariableTracker":
+    def next_variable(self, tx: Any) -> VariableTracker:
         unimplemented(
             gb_type="Unsupported next() call",
             context=f"next({self})",
@@ -740,7 +912,7 @@ def is_immutable(self) -> bool:
     def build(
         tx: Any,
         value: Any,
-        source: Optional[Source] = None,
+        source: Source | None = None,
         realize: bool = False,
     ) -> Any:
         """Create a new VariableTracker from a value and optional Source"""
@@ -795,6 +967,13 @@ def get_python_hash(self) -> int:
             ],
         )
 
+    def get_real_python_backed_value(self) -> object:
+        """Return the Python object this VT wraps, for `is` comparison.
+
+        Returns NO_SUCH_SUBOBJ if no concrete Python object is available.
+        """
+        return NO_SUCH_SUBOBJ
+
     def is_python_equal(self, other: object) -> bool:
         """
         NB - Deliberately not overriding the __eq__ method because that can
@@ -812,11 +991,64 @@ def is_python_equal(self, other: object) -> bool:
             ],
         )
 
+    def nb_index_impl(
+        self,
+        tx: Any,
+    ) -> VariableTracker:
+        """Mirrors CPython's PyNumber_Index / nb_index slot.
+
+        https://github.com/python/cpython/blob/c09ccd9c429/Objects/abstract.c#L1411-L1450
+
+        The base implementation raises TypeError, matching CPython's behavior
+        when tp_as_number->nb_index is NULL (_PyIndex_Check fails).
+        """
+        raise_observed_exception(
+            TypeError,
+            tx,
+            args=[
+                f"'{self.python_type_name()}' object cannot be interpreted as an integer"
+            ],
+        )
+
+    def nb_int_impl(
+        self,
+        tx: Any,
+    ) -> VariableTracker:
+        """Mirrors CPython's tp_as_number->nb_int slot.
+
+        Called when type_implements_nb_int returns True for this type.
+        Subclasses override to provide the actual conversion.
+        """
+        unimplemented(
+            gb_type="nb_int_impl not implemented",
+            context=f"{type(self).__name__} has nb_int slot but no nb_int_impl override",
+            explanation=f"The type {self.python_type_name()} has an nb_int C slot but "
+            "the corresponding VariableTracker doesn't implement nb_int_impl.",
+            hints=[*graph_break_hints.SUPPORTABLE],
+        )
+
+    def nb_float_impl(
+        self,
+        tx: Any,
+    ) -> VariableTracker:
+        """Mirrors CPython's tp_as_number->nb_float slot.
+
+        Called when type_implements_nb_float returns True for this type.
+        Subclasses override to provide the actual conversion.
+        """
+        unimplemented(
+            gb_type="nb_float_impl not implemented",
+            context=f"{type(self).__name__} has nb_float slot but no nb_float_impl override",
+            explanation=f"The type {self.python_type_name()} has an nb_float C slot but "
+            "the corresponding VariableTracker doesn't implement nb_float_impl.",
+            hints=[*graph_break_hints.SUPPORTABLE],
+        )
+
     def __init__(
         self,
         *,
-        source: Optional[Source] = None,
-        mutation_type: Optional[MutationType] = None,
+        source: Source | None = None,
+        mutation_type: MutationType | None = None,
     ) -> None:
         super().__init__()
         self.source = source
@@ -877,9 +1109,9 @@ def reconstruct_failure(self) -> NoReturn:
 
     @staticmethod
     def _add_call_once_guard(
-        cls: type["VariableTracker"],
+        cls: type[VariableTracker],
         method: str,
-        callback: Callable[["VariableTracker"], Any],
+        callback: Callable[[VariableTracker], Any],
     ) -> None:
         original_method = getattr(cls, method)
 
@@ -895,7 +1127,7 @@ def guarded_method(self, *args: Any, **kwargs: Any) -> VariableTracker:
                 active = set()
                 _vt_active_calls.set(active)
 
-            key = (id(self), method)
+            key = (id(self), id(original_method))
             if key in active:
                 callback(self)
 
@@ -909,11 +1141,6 @@ def guarded_method(self, *args: Any, **kwargs: Any) -> VariableTracker:
         setattr(cls, method, guarded_method)
 
 
-def raise_type_error_exc(tx: Any, msg_str: str) -> NoReturn:
-    msg = variables.ConstantVariable.create(msg_str)
-    raise_observed_exception(TypeError, tx, args=[msg])
-
-
 def typestr(*objs: object) -> str:
     if len(objs) == 1:
         (obj,) = objs
@@ -926,5 +1153,52 @@ def typestr(*objs: object) -> str:
 
 
 instancecheck = type.__instancecheck__
+
+
+_CPYTHON_BASE_URL = "https://github.com/python/cpython/blob/v3.13.0/"
+
+
+def print_cpython_to_vt_mapping() -> None:
+    """Print the mapping from CPython types to Dynamo VariableTracker classes.
+
+    Reads _cpython_type from each VT subclass (own attribute only, not inherited).
+    """
+    # Ensure all VT modules are imported so all_subclasses is complete
+    from . import (  # noqa: F401
+        builtin,
+        constant,
+        ctx_manager,
+        dicts,
+        distributed,
+        functions,
+        higher_order_ops,
+        iter,
+        lazy,
+        lists,
+        misc,
+        nn_module,
+        streams,
+        tensor,
+        torch,
+        torch_function,
+        user_defined,
+    )
+
+    mapping: dict[type, list[type]] = {}
+    for vt_cls in VariableTrackerMeta.all_subclasses:
+        cpython_type = vt_cls.__dict__.get("_cpython_type")
+        if cpython_type is None:
+            continue
+        types = cpython_type if isinstance(cpython_type, tuple) else (cpython_type,)
+        for t in types:
+            mapping.setdefault(t, []).append(vt_cls)
+
+    for py_type, vt_classes in sorted(mapping.items(), key=lambda x: x[0].__qualname__):
+        for vt_cls in vt_classes:
+            print(
+                f"{py_type.__module__}.{py_type.__qualname__:30s} -> {vt_cls.__name__}"
+            )
+
+
 from . import builder
 from .lazy import LazyVariableTracker
diff --git a/torch/_dynamo/variables/builder.py b/torch/_dynamo/variables/builder.py
index 9e5b6ed8e2cd0..482af7ae13480 100644
--- a/torch/_dynamo/variables/builder.py
+++ b/torch/_dynamo/variables/builder.py
@@ -37,7 +37,7 @@
 import weakref
 from collections.abc import Callable, MutableMapping
 from types import ModuleType
-from typing import Any, NamedTuple, NoReturn, Optional, overload, TYPE_CHECKING, Union
+from typing import Any, NamedTuple, NoReturn, overload, TYPE_CHECKING, Union
 
 import sympy
 
@@ -45,6 +45,7 @@
 from torch import SymInt
 from torch._dispatch.python import enable_python_dispatcher
 from torch._dynamo.graph_bytecode_inputs import (
+    CURRENT_STREAM_INDEX,
     get_external_object_by_index,
     register_user_object,
 )
@@ -52,6 +53,7 @@
     get_metrics_context,
     is_int_specialization_case,
     is_torch_sym,
+    normalize_count_iter,
     set_feature_use,
 )
 from torch._guards import TracingContext
@@ -63,6 +65,7 @@
     is_opaque_value_type,
     should_hoist,
 )
+from torch._opaque_base import OpaqueBase
 from torch._ops import HigherOrderOperator, OpOverload, OpOverloadPacket
 from torch._subclasses.fake_tensor import (
     FakeTensor,
@@ -115,6 +118,7 @@
     ChainedSource,
     ConstDictKeySource,
     ConvertIntSource,
+    CurrentStreamSource,
     DictGetItemSource,
     DictSubclassGetItemSource,
     DynamicScalarSource,
@@ -147,8 +151,8 @@
     clone_input,
     common_constant_types,
     dict_keys,
+    enumerate_items_with_dict_position,
     get_fake_value,
-    get_items_from_dict,
     get_locals_to_steal,
     get_static_address_type,
     is_frozen_dataclass,
@@ -175,6 +179,7 @@
     wrap_fake_exception,
 )
 from .base import (
+    AttributeMutationExisting,
     AttributeMutationNew,
     typestr,
     ValueMutationExisting,
@@ -183,9 +188,10 @@
     VariableTrackerMeta,
 )
 from .builtin import BuiltinVariable
-from .constant import ConstantVariable, EnumVariable
+from .constant import ConstantVariable
 from .ctx_manager import (
     AutocastModeVariable,
+    CudagraphOverrideVariable,
     DynamoConfigPatchVariable,
     ErrorOnGraphBreakVariable,
     NullContextVariable,
@@ -194,20 +200,10 @@
 from .dicts import (
     ConstDictVariable,
     DefaultDictVariable,
-    DictKeySetVariable,
-    FrozensetVariable,
     MappingProxyVariable,
-    OrderedSetClassVariable,
-    OrderedSetVariable,
     SetVariable,
 )
-from .distributed import (
-    DeviceMeshVariable,
-    PlacementClassVariable,
-    PlacementVariable,
-    ProcessGroupVariable,
-    WorldMetaClassVariable,
-)
+from .distributed import WorldMetaClassVariable
 from .functions import (
     BuiltinMethodVariable,
     CollectionsNamedTupleFunction,
@@ -215,25 +211,22 @@
     CreateTMADescriptorExperimentalVariable,
     CreateTMADescriptorStableVariable,
     FunctoolsPartialVariable,
-    FunctoolsWrapsVariable,
     SysFunctionVariable,
     TritonKernelVariable,
-    TritonSetAllocatorSkipVariable,
+    TritonSetAllocatorVariable,
     UserFunctionVariable,
-    UserMethodVariable,
     WrapperUserFunctionVariable,
 )
 from .higher_order_ops import (
     LocalMapWrappedHigherOrderVariable,
     TorchHigherOrderOperatorVariable,
 )
-from .iter import ItertoolsVariable
+from .iter import CountIteratorVariable, ItertoolsVariable
 from .lazy import LazyConstantVariable, LazyVariableTracker
 from .lists import (
     BaseListVariable,
     ListIteratorVariable,
     ListVariable,
-    NamedTupleVariable,
     RangeVariable,
     SizeVariable,
     SliceVariable,
@@ -261,6 +254,7 @@
     RandomClassVariable,
     RandomVariable,
     SavedTensorBox,
+    StringFormatVariable,
     TorchVersionVariable,
     TypingVariable,
     WeakRefVariable,
@@ -273,6 +267,12 @@
 from .optimizer import OptimizerVariable
 from .script_object import OpaqueObjectClassVariable, TorchScriptObjectVariable
 from .sdpa import SDPAParamsVariable
+from .sets import (
+    DictKeySetVariable,
+    FrozensetVariable,
+    OrderedSetClassVariable,
+    OrderedSetVariable,
+)
 from .streams import EventVariable, StreamContextVariable, StreamVariable
 from .tensor import (
     NumpyNdarrayVariable,
@@ -301,8 +301,8 @@
     MutableMappingVariable,
     SourcelessGraphModuleVariable,
     UserDefinedClassVariable,
+    UserDefinedConstantVariable,
     UserDefinedDictVariable,
-    UserDefinedEnumClassVariable,
     UserDefinedExceptionClassVariable,
     UserDefinedListVariable,
     UserDefinedObjectVariable,
@@ -493,6 +493,7 @@ def __call__(self, value: object) -> VariableTracker:
             )
 
     def _call_impl(self, value: object) -> VariableTracker:
+        self.tx.output.current_tracer.traced_sources.add(self.source)
         if value in self.tx.output.side_effects:
             side_effect_result = self.tx.output.side_effects[value]
             dup_guard = make_dupe_guard(self.source, side_effect_result.source)
@@ -598,6 +599,7 @@ def _type_dispatch_impl(cls, trace_numpy: bool) -> dict[object, Callable[..., An
                 (tuple, list, odict_values, collections.deque, torch.Size),
                 cls.wrap_listlike,
             ),
+            (itertools.count, cls.wrap_itertools_count),
             (tuple_iterator, cls.wrap_tuple_iterator),
             (range_iterator, cls.wrap_range_iterator),
             ((slice, range), cls.wrap_slice_range),
@@ -647,9 +649,20 @@ def wrap_removable_handle(
         )
 
     def wrap_jit_function(self, value: Any) -> WrapperUserFunctionVariable:
+        if not hasattr(value, "_torchdynamo_inline"):
+            unimplemented(
+                gb_type="wrap_jit_function: missing _torchdynamo_inline",
+                context=f"type: {type(value).__name__}",
+                explanation="Dynamo expected a JIT function with a _torchdynamo_inline attribute, "
+                "but the object does not have one.",
+                hints=[*graph_break_hints.SUPPORTABLE],
+            )
         self.install_guards(GuardBuilder.TYPE_MATCH)
         return WrapperUserFunctionVariable(
-            value, "_torchdynamo_inline", source=self.source
+            value,
+            "_torchdynamo_inline",
+            source=self.source,
+            mutation_type=AttributeMutationExisting(),
         )
 
     def wrap_mapping_proxy(self, value: Any) -> VariableTracker:
@@ -681,7 +694,6 @@ def build_key_value(k: Any, v: Any) -> tuple[VariableTracker, VariableTracker]:
         items = dict(build_key_value(k, v) for k, v in value.items())
 
         # Create a dict_vt to be used in the mapping proxy variable
-        # pyrefly: ignore[bad-argument-type]
         dict_vt = ConstDictVariable(items, source=None)
         result = MappingProxyVariable(dict_vt, source=self.source)
         return self.tx.output.side_effects.track_mutable(value, result)
@@ -724,6 +736,7 @@ def _wrap(self, value: Any) -> VariableTracker:
         )
 
         from ..decorators import (
+            CudagraphOverrideContextManager,
             DynamoConfigPatchProxy,
             ErrorOnGraphBreakDecoratorContextManager,
         )
@@ -813,9 +826,8 @@ def set_allocator() -> None:
                 source=self.source,
                 mutation_type=ValueMutationExisting(),
             )
-            result = NamedTupleVariable(
-                output,
-                tuple_cls=type(value),
+            result = UserDefinedTupleVariable.get_vt_cls(type(value))(
+                value,
                 source=self.source,
                 tuple_vt=tuple_vt,
             )
@@ -848,7 +860,7 @@ def set_allocator() -> None:
                 self.tx.output.guard_on_key_order.add(self.source)
 
             # We need all the keys to be hashable. We do this within the
-            # _HashableTracker class in dicts.py
+            # HashableTracker class in hashable.py
             def build_key_value(
                 i: Any, k: Any, v: Any
             ) -> tuple[VariableTracker, VariableTracker]:
@@ -864,13 +876,9 @@ def build_key_value(
 
                 return key, res_value
 
-            # Ensure that we call dict.keys and not value.keys (which can call
-            # overridden keys method). In the C++ guards, we relied on
-            # PyDict_Next to traverse the dictionary, which uses the internal
-            # data structure and does not call the overridden keys method.
             result = dict(
                 build_key_value(i, k, v)
-                for i, (k, v) in enumerate(get_items_from_dict(value))
+                for i, k, v in enumerate_items_with_dict_position(value)
             )
 
             if istype(value, collections.defaultdict):
@@ -962,10 +970,11 @@ def build_key_value(
             self.install_guards(GuardBuilder.EQUALS_MATCH)
             return FrozensetVariable(items, source=self.source)
         elif isinstance(
-            value, (enum.Enum, torch.DispatchKey, torch._C._functorch.TransformType)
+            value,
+            (enum.Enum, torch.DispatchKey, torch._C._functorch.TransformType),
         ):
             self.install_guards(GuardBuilder.ID_MATCH)
-            return EnumVariable(value=value, source=self.source)
+            return UserDefinedObjectVariable(value, source=self.source)
         elif DebuggingVariable.is_reorderable_logging_function(value):
             # Put this above builtin_callable so that print() can be handled
             # along with other builtin debugging functions
@@ -1126,6 +1135,7 @@ def build_key_value(
                     source=AttrSource(self.source, member="__self__"),
                 ),
                 "apply",
+                py_type=type(value),
             )
         elif isinstance(value, torch._C._ImperativeEngine):
             self.install_guards(GuardBuilder.ID_MATCH)
@@ -1148,6 +1158,8 @@ def build_key_value(
             return DynamoConfigPatchVariable(value.changes)
         elif isinstance(value, ErrorOnGraphBreakDecoratorContextManager):
             return ErrorOnGraphBreakVariable(value.error_on_graph_break)
+        elif isinstance(value, CudagraphOverrideContextManager):
+            return CudagraphOverrideVariable(value.fwd, value.bwd)
         elif callable(value) and trace_rules.lookup_callable(value) is not None:
             if trace_rules.is_callable_allowed(value):
                 self.tx.output.has_user_defined_allowed_in_graph = True
@@ -1174,9 +1186,12 @@ def build_key_value(
             # type: ignore[arg-type]
             return StreamContextVariable.create(self.tx, stream_var)
         elif isinstance(value, torch.Stream):
-            # This refers to the device-agnostic torch.Stream
             self.install_guards(GuardBuilder.TYPE_MATCH)
-            index = register_user_object(value, self.source)
+            if isinstance(self.source, CurrentStreamSource):
+                # Reuse the index pre-allocated in SymbolicStreamState.__init__
+                index = CURRENT_STREAM_INDEX
+            else:
+                index = register_user_object(value, self.source)
             stream_proxy = self.tx.output.create_proxy(
                 "call_function", get_external_object_by_index, (index,), {}
             )
@@ -1221,30 +1236,16 @@ def build_key_value(
         elif isinstance(value, torch.optim.Optimizer):
             self.install_guards(GuardBuilder.ID_MATCH)
             self.source = OptimizerSource(self.source)
-            return OptimizerVariable(value, source=self.source)
+            return OptimizerVariable(
+                value,
+                source=self.source,
+                mutation_type=AttributeMutationExisting(),
+            )
         elif isinstance(value, torch.DispatchKeySet):
             self.install_guards(GuardBuilder.DISPATCH_KEY_SET_MATCH)
             return DispatchKeySetVariable(value)
         elif WorldMetaClassVariable.is_group_member_type(value):
             return WorldMetaClassVariable(value, source=self.source)
-        elif ProcessGroupVariable.is_process_group(value):
-            self.install_guards(GuardBuilder.ID_MATCH)
-            return ProcessGroupVariable(value, source=self.source)
-        elif DeviceMeshVariable.is_device_mesh(value):
-            # TODO: see if we need to add custom guard instead of a simple ID_MATCH
-            self.install_guards(GuardBuilder.EQUALS_MATCH)
-            return DeviceMeshVariable(value, source=self.source)
-        elif PlacementClassVariable.is_placement_type(value):
-            # TODO: see if we need to add custom guard instead of a simple ID_MATCH
-            self.install_guards(GuardBuilder.ID_MATCH)
-            return PlacementClassVariable(value, source=self.source)
-        elif PlacementVariable.is_placement(value):
-            # TODO: see if we need to add custom guard instead of a simple ID_MATCH
-            self.install_guards(GuardBuilder.EQUALS_MATCH)
-            return PlacementVariable(
-                value,
-                source=self.source,
-            )
         elif value is OrderedSet:
             self.install_guards(GuardBuilder.ID_MATCH)
             return OrderedSetClassVariable()
@@ -1272,6 +1273,8 @@ def build_key_value(
                     hint=value.real,  # type: ignore[attr-defined]
                     source=source,
                 )
+                if not isinstance(node, SymInt):
+                    raise AssertionError(f"Expected SymInt, got {type(node)}")
 
             # Bind to graph input
             sym_node_proxy = self.tx.output.root_tracer.create_graph_input(
@@ -1334,6 +1337,8 @@ def build_key_value(
                         hints=[*graph_break_hints.SUPPORTABLE],
                     )
             assert new_symint is not None
+            if not isinstance(new_symint, SymInt):
+                raise AssertionError(f"Expected SymInt, got {type(new_symint)}")
             sym_node_proxy = self.tx.output.root_tracer.create_graph_input(
                 re.sub(r"[^a-zA-Z0-9]+", "_", self.name),
                 type(new_symint),
@@ -1376,18 +1381,21 @@ def build_key_value(
         elif value is TensorDescriptor.from_tensor:
             return CreateTMADescriptorStableVariable()
         elif value is set_allocator:
-            return TritonSetAllocatorSkipVariable(value)
+            return TritonSetAllocatorVariable(value)
         elif isinstance(value, torch.amp.autocast_mode.autocast):
-            self.install_guards(GuardBuilder.ID_MATCH)
-            return AutocastModeVariable(
-                target_values=[
-                    value.device,
-                    value.fast_dtype,
-                    value._enabled,
-                    value._cache_enabled,
-                ],
-                source=self.source,
-            )
+            if isinstance(value, torch.amp.autocast_mode._UnmanagedAutocast):
+                return self.wrap_user_defined(value)
+            else:
+                self.install_guards(GuardBuilder.ID_MATCH)
+                return AutocastModeVariable(
+                    target_values=[
+                        value.device,
+                        value.fast_dtype,
+                        value._enabled,
+                        value._cache_enabled,
+                    ],
+                    source=self.source,
+                )
         elif TorchCtxManagerClassVariable.is_matching_cls(value):
             if inspect.isclass(value):
                 self.install_guards(GuardBuilder.CLASS_MATCH)
@@ -1397,11 +1405,19 @@ def build_key_value(
         elif inspect.getattr_static(value, "__script_if_tracing_wrapper", False):
             self.install_guards(GuardBuilder.TYPE_MATCH)
             return WrapperUserFunctionVariable(
-                value, "__original_fn", source=self.source
+                value,
+                "__original_fn",
+                source=self.source,
+                mutation_type=AttributeMutationExisting(),
             )
         elif is_lru_cache_wrapped_function(value):
             self.install_guards(GuardBuilder.TYPE_MATCH)
-            return WrapperUserFunctionVariable(value, "__wrapped__", source=self.source)
+            return WrapperUserFunctionVariable(
+                value,
+                "__wrapped__",
+                source=self.source,
+                mutation_type=AttributeMutationExisting(),
+            )
         elif value is sys.exc_info or (
             sys.version_info >= (3, 11) and value is sys.exception
         ):
@@ -1411,11 +1427,11 @@ def build_key_value(
         ):
             self.install_guards(GuardBuilder.TYPE_MATCH)
             return WrapperUserFunctionVariable(
-                value, "_torchdynamo_inline", source=self.source
+                value,
+                "_torchdynamo_inline",
+                source=self.source,
+                mutation_type=AttributeMutationExisting(),
             )
-        elif value is functools.wraps:
-            self.install_guards(GuardBuilder.ID_MATCH)
-            return FunctoolsWrapsVariable(value, source=self.source)
         elif value is collections.namedtuple:
             self.install_guards(GuardBuilder.ID_MATCH)
             return CollectionsNamedTupleFunction(value, source=self.source)
@@ -1429,6 +1445,7 @@ def build_key_value(
             return GetAttrVariable(
                 BuiltinVariable(float, source=self.source),
                 value.__name__,
+                py_type=type(value),
             )
         elif is_function_or_wrapper(value):
             value, attr_name = unwrap_with_attr_name_if_wrapper(value)
@@ -1462,31 +1479,6 @@ def build_key_value(
             )
             self.tx.output.side_effects.track_object_existing(value, result)
             return result
-        elif isinstance(value, types.MethodType) and isinstance(
-            value.__self__, (torch.nn.Module, torch.utils._pytree.TreeSpec)
-        ):
-            # don't let MethodTypes fall through to UserDefinedObject,
-            # which doesn't support 'CALL_FUNCTION'
-
-            # TODO(whc): Why do we limit this to methods on NNModules?
-            # I don't have a good reason for this, but it preserves the existing behavior
-            # for MBartForConditionalGeneration, which generates many graph breaks and OOMs otherwise.
-            # I suspect we probably want to relax this check and dig deeper there.
-
-            # In order to construct a MethodVariable in Dynamo, we start with an actual method obj from python,
-            # but need to separately wrap its underlying `__func__` and its `self` argument.  We wrap `self` here
-            # and then `__func__` gets wrapped inside UserMethodVariable.
-            self_obj = VariableBuilder(
-                self.tx, source=AttrSource(self.source, "__self__")
-            )(value.__self__)
-            assert self_obj and isinstance(self_obj, VariableTracker), (
-                "Failed to produce a valid self obj"
-            )
-            return UserMethodVariable(
-                value.__func__,
-                self_obj,
-                source=self.source,
-            )
         elif isinstance(value, types.GetSetDescriptorType):
             # GetSet descriptors are C functions attached to an attribute lookup
             # using PyGetSetDef. Python, on attribute lookup, can decide to
@@ -1543,14 +1535,14 @@ def build_key_value(
                 # ID_MATCH even if its a global variable.
                 self.install_guards(GuardBuilder.CLASS_MATCH)
 
-            if is_opaque_type(value):
-                return OpaqueObjectClassVariable(
+            if isinstance(value, type) and issubclass(value, enum.Enum):
+                return UserDefinedClassVariable(
                     value,
                     source=self.source,
                 )
 
-            if isinstance(value, type) and issubclass(value, enum.Enum):
-                return UserDefinedEnumClassVariable(
+            if is_opaque_type(value):
+                return OpaqueObjectClassVariable(
                     value,
                     source=self.source,
                 )
@@ -1565,6 +1557,12 @@ def build_key_value(
                 ScriptObjectQualifiedNameSource,
             )
 
+            # Unwrap FakeScriptObject to the underlying real object so the
+            # rest of this branch (guards, graph inputs, etc.) operates on
+            # the real opaque object type.
+            if isinstance(value, torch._library.fake_class_registry.FakeScriptObject):
+                value = value.real_obj
+
             # type: ignore[arg-type]
             if torch._library.fake_class_registry.tracing_with_real(value):
                 proxy = self.tx.output.root_tracer.create_graph_input(
@@ -1626,7 +1624,18 @@ def build_key_value(
                 self.tx.output.fake_mode, value
             )
             if is_opaque_value_type(type(value)) and not should_hoist(type(value)):
+                fake_script_obj = value
                 proxy = value
+
+            elif config.install_free_tensors and (
+                is_from_global_source(self.source)
+                or is_from_nonlocal_source(self.source)
+                or is_from_unspecialized_nn_module_source(self.source)
+            ):
+                return self.tx.output.register_attr_or_module(
+                    value, self.name, source=self.source
+                )
+
             else:
                 proxy = self.tx.output.root_tracer.create_graph_input(
                     re.sub(r"[^a-zA-Z0-9]+", "_", self.name),
@@ -1663,7 +1672,7 @@ def build_key_value(
             self.tx.output.guard_on_key_order.add(self.source)
 
             # We need all the keys to be hashable. We do this within the
-            # _HashableTracker class in dicts.py
+            # HashableTracker class in hashable.py
             def build_key_value(
                 i: Any, k: Any, v: Any
             ) -> tuple[VariableTracker, VariableTracker]:
@@ -1676,17 +1685,12 @@ def build_key_value(
 
                 return key, res_value
 
-            # Ensure that we call dict.keys and not value.keys (which can call
-            # overridden keys method). In the C++ guards, we relied on
-            # PyDict_Next to traverse the dictionary, which uses the internal
-            # data structure and does not call the overridden keys method.
             result = dict(
                 build_key_value(i, k, v)
-                for i, (k, v) in enumerate(get_items_from_dict(value))
+                for i, k, v in enumerate_items_with_dict_position(value)
             )
 
             dict_vt = ConstDictVariable(
-                # pyrefly: ignore[bad-argument-type]
                 result,
                 user_cls=(
                     collections.OrderedDict
@@ -1774,7 +1778,7 @@ def build_key_value(
             return self.tx.output.side_effects.track_object_existing(value, result)
         elif is_frozen_dataclass(value):
             self.install_guards(GuardBuilder.TYPE_MATCH)
-            result = FrozenDataClassVariable.create(self.tx, value, source=self.source)
+            result = FrozenDataClassVariable(value, source=self.source)
             return self.tx.output.side_effects.track_object_existing(value, result)
         elif isinstance(value, dict_keys):
             if all(ConstantVariable.is_literal(k) for k in value):
@@ -1832,6 +1836,8 @@ def build_key_value(
             return self.wrap_user_defined(value)
 
     def wrap_user_defined(self, value: Any) -> VariableTracker:
+        from .user_defined import _CONSTANT_BASE_TYPES
+
         self.install_guards(GuardBuilder.TYPE_MATCH)
         if InspectVariable.is_matching_object(value):
             # Skip guards on inspect related variable trackers because they are
@@ -1839,6 +1845,12 @@ def wrap_user_defined(self, value: Any) -> VariableTracker:
             # cause recompiles) and can cause a large number of OBJECT_ALIASING
             # guards.
             result = InspectVariable(value, source=SkipGuardSource(self.source))
+        elif (
+            isinstance(value, _CONSTANT_BASE_TYPES)
+            and type(value) not in common_constant_types
+        ):
+            self.install_guards(GuardBuilder.CONSTANT_SUBCLASS_MATCH)
+            result = UserDefinedConstantVariable(value, source=self.source)
         else:
             result = UserDefinedObjectVariable(value, source=self.source)
         if not SideEffects.cls_supports_mutation_side_effects(type(value)):
@@ -1975,10 +1987,28 @@ def wrap_range_iterator(self, value: range_iterator) -> VariableTracker:
         self.install_guards(GuardBuilder.RANGE_ITERATOR_MATCH)
         # Get all the values from the range iterator; no need to install guards
         # on items since `RANGE_ITERATOR_MATCH` guarantees the same items.
-        items = [ConstantVariable.create(v) for v in copy.deepcopy(value)]
+        items: list[VariableTracker] = [
+            ConstantVariable.create(v) for v in copy.deepcopy(value)
+        ]
         result = ListIteratorVariable(items, source=self.source)
         return self.tx.output.side_effects.track_mutable(value, result)
 
+    def wrap_itertools_count(self, value: Any) -> VariableTracker:
+        current_item, step = normalize_count_iter(value)
+        if not (
+            ConstantVariable.is_literal(current_item)
+            and ConstantVariable.is_literal(step)
+        ):
+            return self.wrap_user_defined(value)
+
+        self.install_guards(GuardBuilder.COUNT_ITERATOR_MATCH)
+        result = CountIteratorVariable(
+            ConstantVariable.create(current_item),
+            ConstantVariable.create(step),
+            source=self.source,
+        )
+        return self.tx.output.side_effects.track_mutable(value, result)
+
     def wrap_slice_range(self, value: slice | range) -> SliceVariable | RangeVariable:
         items = [
             VariableBuilder(self.tx, AttrSource(self.get_source(), k))(
@@ -2042,6 +2072,19 @@ def wrap_module(self, value: torch.nn.Module) -> VariableTracker:
             self.source = AttrSource(self.source, "_orig_mod")
             return self.wrap_module(value._orig_mod)
 
+        if type(value) is torch.jit._script.RecursiveScriptModule:
+            unimplemented(
+                gb_type="torch.jit.script/freeze modules unsupported",
+                context=str(value),
+                explanation="Dynamo does not support tracing into torch.jit.script or "
+                "torch.jit.freeze modules because they execute in the TorchScript "
+                "runtime, not Python. Replace the ScriptModule submodule with the "
+                "original eager nn.Module.",
+                hints=[
+                    *graph_break_hints.FUNDAMENTAL,
+                ],
+            )
+
         if (
             isinstance(value, (torch.nn.RNN, torch.nn.GRU, torch.nn.LSTM))
             and not config.allow_rnn
@@ -2096,62 +2139,54 @@ def wrap_module(self, value: torch.nn.Module) -> VariableTracker:
                 value = value.get_base()
                 self.source = AttrProxySource(self.source)
 
-            if torch._dynamo.config.inline_inbuilt_nn_modules:
-                freezing = is_parameter_freezing()
+            freezing = is_parameter_freezing()
 
-                # Guard against the case where user may overwrite named parameters
-                # / named buffers
-                # NOTE: This is not likely to happen but worth guarding to avoid
-                # exception
-                if (
-                    callable(value.named_parameters)
+            # Guard against the case where user may overwrite named parameters
+            # / named buffers
+            # NOTE: This is not likely to happen but worth guarding to avoid
+            # exception
+            if (
+                callable(value.named_parameters)
+                # type: ignore[attr-defined]
+                and value.named_parameters.__func__ is og_module_named_parameters_fn_ptr
+            ):
+                try:  # catch TypeErrors in named_parameters() from unserializable nn modules
                     # type: ignore[attr-defined]
-                    and value.named_parameters.__func__
-                    is og_module_named_parameters_fn_ptr
-                ):
-                    try:  # catch TypeErrors in named_parameters() from unserializable nn modules
-                        # type: ignore[attr-defined]
-                        for _, p in value.named_parameters():
-                            self.mark_static_input(p, guard=freezing)
-                    except TypeError as e:
-                        raise_observed_exception(type(e), self.tx, args=list(e.args))
-
-                if (
-                    callable(value.named_buffers)
+                    for _, p in value.named_parameters():
+                        self.mark_static_input(p, guard=freezing)
+                except TypeError as e:
+                    raise_observed_exception(type(e), self.tx, args=list(e.args))
+
+            if (
+                callable(value.named_buffers)
+                # type: ignore[attr-defined]
+                and value.named_buffers.__func__ is og_module_named_buffers_fn_ptr
+            ):
+                try:  # catch TypeErrors in named_parameters() from unserializable nn modules
                     # type: ignore[attr-defined]
-                    and value.named_buffers.__func__ is og_module_named_buffers_fn_ptr
-                ):
-                    try:  # catch TypeErrors in named_parameters() from unserializable nn modules
-                        # type: ignore[attr-defined]
-                        for _, b in value.named_buffers():
-                            self.mark_static_input(b, guard=freezing)
-                    except TypeError as e:
-                        raise_observed_exception(type(e), self.tx, args=list(e.args))
-
-                if freezing:
-                    # we need to add the module to tracing context
-                    # in order to allow its params to get invalidated
-                    # this will get cleaned up once compile ends
-                    self.tx.output.nn_modules[self.name] = value
+                    for _, b in value.named_buffers():
+                        self.mark_static_input(b, guard=freezing)
+                except TypeError as e:
+                    raise_observed_exception(type(e), self.tx, args=list(e.args))
+
+            if freezing:
+                # we need to add the module to tracing context
+                # in order to allow its params to get invalidated
+                # this will get cleaned up once compile ends
+                self.tx.output.nn_modules[self.name] = value
 
             if (
                 value.__module__.startswith(("torch.nn.modules", "torch.ao."))
                 and not value.__module__.startswith("torch.nn.modules.container")
             ) or getattr(value.__class__, "_dynamo_marked_static", False):
                 new_source = self.source
-                if config.inline_inbuilt_nn_modules and (
-                    not self.tx.output.export or config.install_free_tensors
-                ):
-                    # Export corner case - look at test_repros.py test_inlining_cornercase
+                if not self.tx.output.export or config.install_free_tensors:
                     new_source = UnspecializedBuiltinNNModuleSource(self.source)
                 result = UnspecializedBuiltinNNModuleVariable(value, source=new_source)
                 install_guard(new_source.make_guard(GuardBuilder.TYPE_MATCH))
             else:
                 new_source = self.source
-                if config.inline_inbuilt_nn_modules and (
-                    not self.tx.output.export or config.install_free_tensors
-                ):
-                    # Export corner case - look at test_repros.py test_inlining_cornercase
+                if not self.tx.output.export or config.install_free_tensors:
                     new_source = UnspecializedNNModuleSource(self.source)
                 result = UnspecializedNNModuleVariable(value, source=new_source)
                 install_guard(new_source.make_guard(GuardBuilder.TYPE_MATCH))
@@ -2263,14 +2298,14 @@ def _wrap_lazy_constant(
     @overload
     def _wrap_lazy_constant(
         self,
-        value: Union[int, float, bool, str],
+        value: int | float | bool | str,
         wrap_fn: None = None,
     ) -> VariableTracker: ...
 
     def _wrap_lazy_constant(
         self,
-        value: Union[int, float, bool, str],
-        wrap_fn: Optional[Callable[[Any], VariableTracker]] = None,
+        value: int | float | bool | str,
+        wrap_fn: Callable[[Any], VariableTracker] | None = None,
     ) -> VariableTracker:
         """Wrap a primitive constant, deferring guard installation if allowed."""
         if not self.allow_lazy_constant:
@@ -2296,15 +2331,10 @@ def wrap_tensor(self, value: torch.Tensor) -> VariableTracker:
 
         is_static_input = get_static_address_type(value) is not None
 
-        if (
-            config.inline_inbuilt_nn_modules
-            and not is_static_input
-            and (
-                isinstance(value, torch.nn.Parameter)
-                # mark tensor attributes of nn modules static. This is done to keep inline_inbuilt_nn_modules behavior
-                # compatible with previous behavior.
-                or (source and source.guard_source.is_unspecialized_nn_module())
-            )
+        if not is_static_input and (
+            isinstance(value, torch.nn.Parameter)
+            # mark tensor attributes of nn modules static
+            or (source and source.guard_source.is_unspecialized_nn_module())
         ):
             self.mark_static_input(value, guard=is_parameter_freezing())
             is_static_input = True
@@ -2320,9 +2350,7 @@ def wrap_tensor(self, value: torch.Tensor) -> VariableTracker:
         )
 
         make_graph_attribute = is_static_input and (
-            not config.inline_inbuilt_nn_modules
-            or is_parameter_freezing()
-            or torch._dynamo.config.prepare_freezing
+            is_parameter_freezing() or torch._dynamo.config.prepare_freezing
         )
 
         if should_install_free_tensor or (
@@ -2452,6 +2480,12 @@ def wrap_tensor(self, value: torch.Tensor) -> VariableTracker:
             **options,
         )
 
+        # Track input tensors for attribute mutation, matching how
+        # handle_traced_output tracks intermediate tensors with AttributeMutationNew.
+        # This enables setattr on input tensors (e.g. tensor.custom_attr = val)
+        # without graph breaking.
+        self.tx.output.side_effects.track_object_existing(value, tensor_variable)
+
         if value._is_view():
             # If value is a view, add its base tensor to the tracked fakes list.
             # This is so we are able to access the correct source for its symbolic
@@ -2494,25 +2528,22 @@ def wrap_tensor(self, value: torch.Tensor) -> VariableTracker:
             if is_dtensor:
                 self.install_guards(GuardBuilder.TYPE_MATCH)
 
-                # The inner tensor name is always _local_tensor. If its not, we
-                # raise assertion to update the check accordingly.
-                inner_tensor_name = value.__tensor_flatten__()[0][0]
-                if inner_tensor_name != "_local_tensor":
+                inner_attrs = value.__tensor_flatten__()[0]
+                if inner_attrs != ["_local_tensor", "device_mesh"]:
                     raise RuntimeError(
-                        "Expecting Dtensor inner tensor name to be _local_tensor"
+                        "Expecting DTensor inner attrs to be ['_local_tensor', 'device_mesh']"
                     )
 
-                # Now selectively guard on the flattening context
                 flattening_ctx = value.__tensor_flatten__()[1]
-                # This is supposed to be (self._spec, self.requires_grad)
                 if not (
-                    len(flattening_ctx) == 2
-                    and flattening_ctx[0] == value._spec
-                    and flattening_ctx[1] == value.requires_grad
+                    len(flattening_ctx) == 4
+                    and flattening_ctx[0] == value._spec.placements
+                    and flattening_ctx[1] == value._spec.tensor_meta
+                    and flattening_ctx[2] == value._spec.shard_order
+                    and flattening_ctx[3] == value.requires_grad
                 ):
-                    # If not, raise an assertion to update to the new guards
                     raise RuntimeError(
-                        "Expecting Dtensor flattening ctx to be _spec, requires_grad"
+                        "Expecting DTensor flattening ctx to be (placements, tensor_meta, shard_order, requires_grad)"
                     )
                 # Guard on the dtensor spec
                 install_guard(
@@ -2536,6 +2567,23 @@ def wrap_tensor(self, value: torch.Tensor) -> VariableTracker:
             attrs, _ = value.__tensor_flatten__()
             for attr in attrs:
                 inner_value = getattr(value, attr)
+                # FakeScriptObject wraps the real opaque object during
+                # fake-mode tracing; unwrap before the type check.
+                inner_type = type(inner_value)
+                if isinstance(
+                    inner_value,
+                    torch._library.fake_class_registry.FakeScriptObject,
+                ):
+                    inner_type = type(inner_value.real_obj)
+                if not isinstance(
+                    inner_value, torch.Tensor
+                ) and not is_opaque_reference_type(inner_type):
+                    raise RuntimeError(
+                        f"{type(inner_value).__name__!r} found in tensor attrs of "
+                        f"{type(value).__name__}.__tensor_flatten__(). "
+                        "Only tensors and reference-type opaques are allowed "
+                        "in tensor attrs."
+                    )
                 inner_source = AttrSource(self.source, attr)
                 LazyVariableTracker.realize_all(
                     VariableBuilder(self.tx, inner_source)(inner_value)
@@ -2652,6 +2700,7 @@ def wrap_symint(
             return self.tx.output.unspec_variable_map[self.name]
 
         shape_env = self.tx.output.shape_env
+        frame_state_entry: FrameStateSizeEntry | None = None
         if TracingContext.get().force_unspec_int_unbacked_size_like:
             wrapped_value = shape_env.create_unbacked_symint()
             _constrain_range_for_size(wrapped_value)
@@ -2715,11 +2764,20 @@ def wrap_symint(
                 self.install_guards(GuardBuilder.CONSTANT_MATCH)
                 return ConstantVariable.create(value=value)
 
+            excluded_scalar = (
+                frame_state_entry.excluded_scalar
+                if config.automatic_dynamic_exclusion_guard
+                and frame_state_entry is not None
+                else None
+            )
             wrapped_value = shape_env.create_unspecified_symint_and_symbol(
                 value,
                 source=self.source,
                 dynamic_dim=dynamic_dim,
+                excluded_value=excluded_scalar,
             )
+            if not isinstance(wrapped_value, SymInt):
+                raise AssertionError(f"Expected SymInt, got {type(wrapped_value)}")
 
             self.tx.output.tracked_fakes.append(
                 TrackedFake(wrapped_value, self.source, context)
@@ -3130,7 +3188,7 @@ def wrap_fx_proxy_cls(
         )
         and proxy.node.op != "placeholder"
     ):
-        tx.output.current_tracer.record_tensor_or_symint_vt(out)
+        tx.output.current_tracer.record_proxyable_vt(out)
     return out
 
 
@@ -3190,7 +3248,7 @@ def _wrap_fx_preexisting_tensor(
             }
             assert "source" in options and options["source"] is not None
             kwargs["source"] = options["source"]
-            # pyrefly: ignore[missing-argument, bad-argument-type]
+            # pyrefly: ignore [missing-argument]
             tensor = wrap_to_fake_tensor_and_record(tensor, tx=tx, **kwargs)
 
         if tensor.device.type != "meta" and (
@@ -3339,6 +3397,7 @@ def handle_traced_output(
 
                 # WARNING: this assumes the same target_cls as this tuple/list call
                 unpacked.append(
+                    # pyrefly: ignore [bad-argument-type]
                     wrap_fx_proxy_cls(
                         # pyrefly: ignore[bad-argument-type]
                         target_cls=target_cls,
@@ -3357,13 +3416,18 @@ def handle_traced_output(
         elif istype(example_value, (list, immutable_list)):
             return ListVariable(unpacked, **options)
         else:
-            assert (
-                example_value.__class__.__module__ == "torch.return_types"
-                or hasattr(example_value, "_fields")
-            ), (
-                f"expected {example_value.__class__.__module__} == torch.return_types or named tuple but got {type(example_value)}"
+            assert is_namedtuple(example_value), (
+                f"expected namedtuple or structseq but got {type(example_value)}"
+            )
+            tuple_vt = TupleVariable(
+                unpacked,
+                mutation_type=options.get("mutation_type", ValueMutationNew()),
+            )
+            return UserDefinedTupleVariable.get_vt_cls(type(example_value))(
+                example_value,
+                tuple_vt=tuple_vt,
+                **options,  # type: ignore[arg-type]
             )
-            return NamedTupleVariable(unpacked, example_value.__class__, **options)  # type: ignore[arg-type]
     elif example_value is None or proxy.node.target is torch.manual_seed:
         return ConstantVariable.create(None, **options)
     elif isinstance(example_value, (torch.SymInt, torch.SymFloat, torch.SymBool)):
@@ -3420,6 +3484,8 @@ def handle_traced_output(
         proxy.node.target
         in [
             torch.sym_int,
+            torch.sym_max,
+            torch.sym_min,
             getattr,
             operator.getitem,
             torch._utils._element_size,
@@ -3460,6 +3526,7 @@ def handle_traced_output(
             torch._C._get_mem_efficient_sdp_enabled,
             torch._C._get_math_sdp_enabled,
             torch._C._get_overrideable_sdp_enabled,
+            torch._C._is_autocast_available,
             "is_integer",
         ]
         + list(supported_const_comparison_op_values.keys())
@@ -3476,10 +3543,21 @@ def handle_traced_output(
     elif isinstance(example_value, float) or proxy.node.target in ["hex", "__round__"]:
         set_example_value(proxy.node, example_value)
         return ConstantVariable.create(example_value, **options)
+    elif isinstance(example_value, torch._library.fake_class_registry.FakeScriptObject):
+        # example_value is already a FakeScriptObject (e.g. returned by getitem
+        # on a container whose fake kernel returns a FakeScriptObject).  No need
+        # to convert it — just wrap the proxy directly.
+        return TorchScriptObjectVariable.create(
+            proxy,
+            example_value,
+        )
     elif is_opaque_type(type(example_value)):
         # This is for handling opaque objects in custom ops
         if is_opaque_value_type(type(example_value)):
-            proxy = example_value  # pyrefly: ignore[bad-assignment]
+            return TorchScriptObjectVariable.create(
+                example_value,  # pyrefly: ignore[bad-argument-type]
+                example_value,
+            )
         fake_script_obj = torch._library.fake_class_registry.maybe_to_fake_obj(
             tx.output.fake_mode, example_value
         )
@@ -3574,6 +3652,7 @@ def construct_tensor_variable(
     if proxy.node.op != "placeholder":
         tx.output.current_tracer.track_produced_symints(example_value, proxy)
     options.update(get_specialized_props(target_cls, tx, example_value, subclass_type))
+    # pyrefly: ignore [bad-argument-count]
     return target_cls(proxy, **options)
 
 
@@ -3737,11 +3816,18 @@ def _automatic_dynamic(
         inner_contexts = {}  # mapping from attr -> symbolic context
         attrs, _ = type(e).__tensor_flatten__(e)
         for attr in attrs:
-            inner_tensor = getattr(e, attr)
-            inner_source = AttrSource(source, attr)
-            inner_contexts[attr] = _automatic_dynamic(
-                inner_tensor, tx, inner_source, static_shapes
-            )
+            match getattr(e, attr):
+                case torch.Tensor() as inner_value:
+                    inner_source = AttrSource(source, attr)
+                    inner_contexts[attr] = _automatic_dynamic(
+                        inner_value, tx, inner_source, static_shapes
+                    )
+                case OpaqueBase():
+                    pass
+                case unexpected:
+                    raise AssertionError(
+                        f"expected Tensor or OpaqueBase, got {type(unexpected)}"
+                    )
 
         return SubclassSymbolicContext(
             dynamic_sizes=outer_context.dynamic_sizes,
@@ -3973,6 +4059,8 @@ def update_dim2constraint(
         tensor_source=source,
         shape_env_to_source_to_symbol_cache=shape_env_to_source_to_symbol_cache,
         shape_ids=getattr(e, "_dynamo_shape_ids", None),
+        unbacked_bounds=getattr(e, "_dynamo_unbacked_bounds", None),
+        excluded_sizes=frame_state_entry.excluded_sizes,
     )
 
 
@@ -4057,8 +4145,16 @@ def _wrap_to_fake_tensor_and_record_impl(
             and isinstance(fake_e, FakeTensor)
             and (sym_val := fake_e.item_memo) is not None
         ):
+            # Match the peephole in FakeTensorConverter.from_real_tensor that
+            # strips FloatTensorSource before calling create_symbol.  Without
+            # this, the tracked fake source name won't match source_to_var and
+            # produce_guards_verbose will report "(unknown source)".
+            if isinstance(source, FloatTensorSource):
+                item_source = source.base
+            else:
+                item_source = CallMethodItemSource(source)
             tx.output.tracked_fakes.append(
-                TrackedFake(sym_val, CallMethodItemSource(source), symbolic_context)
+                TrackedFake(sym_val, item_source, symbolic_context)
             )
 
         if is_traceable_wrapper_subclass(fake_e):
@@ -4166,13 +4262,15 @@ def create(tx: "InstructionTranslatorBase", value: Any) -> VariableTracker:
         if isinstance(value, VariableTracker):
             # This is always valid to call, and useful for recursive calls.
             return value
-        elif is_opaque_value_type(type(value)):
+        elif is_opaque_value_type(type(value)) and not isinstance(value, enum.Enum):
+            return TorchScriptObjectVariable.create(value, value)
+        elif is_opaque_reference_type(type(value)):
             # This is for handling opaque objects in custom ops
             fake_script_obj = torch._library.fake_class_registry.maybe_to_fake_obj(
                 tx.output.fake_mode, value
             )
             return TorchScriptObjectVariable.create(
-                value,
+                value,  # pyrefly: ignore[bad-argument-type]
                 fake_script_obj,
             )
         # type: ignore[attr-defined]
@@ -4191,48 +4289,58 @@ def create(tx: "InstructionTranslatorBase", value: Any) -> VariableTracker:
             # NamedTuple._make uses an alias of tuple.__new__
             # pyrefly: ignore[not-callable, bad-argument-count, missing-attribute]
             obj = trace_rules.lookup_callable(value.__self__)(value.__self__)
-            return GetAttrVariable(obj, "__new__")
+            return GetAttrVariable(obj, "__new__", py_type=type(value))
         elif is_function_or_wrapper(value):
             # pyrefly: ignore[not-callable, bad-argument-count]
             return trace_rules.lookup(value)(value)
         elif isinstance(
-            value, (enum.Enum, torch.DispatchKey, torch._C._functorch.TransformType)
+            value,
+            (enum.Enum, torch.DispatchKey, torch._C._functorch.TransformType),
         ):
-            return EnumVariable(value)
+            return UserDefinedObjectVariable(value)
         elif isinstance(value, (type, abc.ABCMeta)):
-            if isinstance(value, type) and issubclass(value, enum.Enum):
-                return UserDefinedEnumClassVariable(value)
+            if issubclass(type(value), type) and issubclass(value, BaseException):
+                return UserDefinedExceptionClassVariable(value)
             return UserDefinedClassVariable(value)
         elif isinstance(value, types.MethodWrapperType):
             return MethodWrapperVariable(value)
-        elif (
-            isinstance(value, types.MethodType)
-            # We only want to support sourceless class objects here
-            # An instance variable is not allowed and it should have source
-            and isinstance(value.__self__, (type, abc.ABCMeta))
-        ):
-            # value is a classmethod
-            assert getattr(value.__self__, value.__func__.__name__) == value
-            cls_obj_vt = SourcelessBuilder.create(tx, value.__self__)
-            try:
-                # pyrefly: ignore[bad-argument-type]
-                return cls_obj_vt.var_getattr(tx, value.__func__.__name__)
-            except NotImplementedError:
-                pass  # failthrough to unimplemented branch
+        elif isinstance(value, types.MethodType):
+            if isinstance(value.__self__, (type, abc.ABCMeta)):
+                # value is a classmethod
+                assert getattr(value.__self__, value.__func__.__name__) == value
+                cls_obj_vt = SourcelessBuilder.create(tx, value.__self__)
+                try:
+                    # pyrefly: ignore[bad-argument-type]
+                    return cls_obj_vt.var_getattr(tx, value.__func__.__name__)
+                except NotImplementedError:
+                    pass  # failthrough to unimplemented branch
+            else:
+                # Instance method — look up the VT for __self__ via side effects
+                obj_vt = tx.output.side_effects.id_to_variable.get(id(value.__self__))
+                if obj_vt is not None:
+                    return torch._dynamo.variables.UserMethodVariable(
+                        value.__func__, obj_vt
+                    )
         elif isinstance(value, torch.fx.graph_module.GraphModule):
             return SourcelessGraphModuleVariable(value)
         elif isinstance(value, torch.utils._pytree.TreeSpec):
             return UserDefinedObjectVariable(value)
-        elif PlacementVariable.is_placement(value):
-            return PlacementVariable(value)
-        elif DeviceMeshVariable.is_device_mesh(value):
-            return DeviceMeshVariable(value)
-        elif value is functools.wraps:
-            return FunctoolsWrapsVariable(value)
         elif isinstance(value, re.Pattern):
             return ConstantLikeVariable(value)
         elif isinstance(value, torch._dynamo.variables.lazy.LazySymNodeFormatString):
-            return ConstantVariable.create(str(value))
+            try:
+                return ConstantVariable.create(str(value))
+            # If we cannot create due to error in str() call, we should
+            # try explicitly for string format variable
+            except (
+                torch._dynamo.exc.UserError,
+                torch.fx.experimental.symbolic_shapes.GuardOnDataDependentSymNode,
+            ):
+                return StringFormatVariable.create(
+                    value.fmt_var.as_python_constant(),
+                    [value.sym_node_var],
+                    {},
+                )
         elif isinstance(value, type(torch._higher_order_ops.flex_attention_backward)):
             return torch._dynamo.variables.higher_order_ops.FlexAttentionBackwardHighOrderVariable(
                 value
@@ -4244,13 +4352,22 @@ def create(tx: "InstructionTranslatorBase", value: Any) -> VariableTracker:
                 SourcelessBuilder.create(tx, getattr(value, name))
                 for name in namedtuple_fields(type(value))
             ]
-            return NamedTupleVariable(output, tuple_cls=type(value))
+            tuple_vt = TupleVariable(output, mutation_type=ValueMutationNew())
+            return UserDefinedTupleVariable.get_vt_cls(type(value))(
+                value, tuple_vt=tuple_vt
+            )
         elif (
             isinstance(value, torch.SymInt)
             and value.node.expr in tx.output.bound_symbols
         ):
             proxy = tx.output.bound_symbols[value.node.expr]
             return SymNodeVariable.create(tx, proxy)
+        elif isinstance(value, slice):
+            items = [
+                SourcelessBuilder.create(tx, getattr(value, k))
+                for k in ("start", "stop", "step")
+            ]
+            return SliceVariable(items, tx)  # pyrefly: ignore[bad-argument-type]
         elif istype(value, object):
             return ObjectVariable(value)
         unimplemented(
diff --git a/torch/_dynamo/variables/builtin.py b/torch/_dynamo/variables/builtin.py
index f7f2050a81229..37526143ad46c 100644
--- a/torch/_dynamo/variables/builtin.py
+++ b/torch/_dynamo/variables/builtin.py
@@ -32,10 +32,9 @@
 import unittest
 from collections import defaultdict, OrderedDict
 from collections.abc import Callable, Iterable, KeysView, Sequence
-from typing import Any, cast, Literal, TYPE_CHECKING, Union
+from typing import Any, cast, TYPE_CHECKING
 
 import torch
-from torch import sym_float, sym_int
 from torch._subclasses.meta_utils import is_sparse_any
 from torch.overrides import BaseTorchFunctionMode
 from torch.utils._python_dispatch import is_traceable_wrapper_subclass
@@ -78,21 +77,25 @@
     proxy_args_kwargs,
     raise_args_mismatch,
     set_methods,
+    specialize_symnode,
     str_methods,
     tensortype_to_dtype,
 )
-from .base import AsPythonConstantNotImplementedError, ValueMutationNew, VariableTracker
-from .constant import CONSTANT_VARIABLE_NONE, ConstantVariable, EnumVariable
+from .base import (
+    AsPythonConstantNotImplementedError,
+    NO_SUCH_SUBOBJ,
+    ValueMutationNew,
+    VariableTracker,
+)
+from .constant import ConstantVariable, FakeIdVariable
 from .dicts import (
     ConstDictVariable,
     DefaultDictVariable,
+    DictItemsVariable,
     DictKeysVariable,
     DictViewVariable,
-    FrozensetVariable,
-    is_hashable,
-    OrderedSetClassVariable,
-    SetVariable,
 )
+from .hashable import is_hashable
 from .lists import (
     BaseListVariable,
     ListIteratorVariable,
@@ -102,7 +105,8 @@
     TupleIteratorVariable,
     TupleVariable,
 )
-from .streams import EventVariable, StreamVariable
+from .misc import NullVariable
+from .sets import FrozensetVariable, OrderedSetClassVariable, SetVariable
 from .tensor import (
     FakeItemVariable,
     supported_comparison_ops,
@@ -142,11 +146,19 @@
     operator.ixor: operator.xor,
 }
 
+_BUILTIN_CONSTANT_FOLDABLE_METHODS: dict[type, frozenset[str]] = {
+    int: frozenset({"__new__", "from_bytes"}),
+    bool: frozenset({"__new__", "from_bytes"}),
+    float: frozenset({"fromhex", "hex"}),
+}
+if sys.version_info >= (3, 14):
+    _BUILTIN_CONSTANT_FOLDABLE_METHODS[complex] = frozenset({"from_number"})
+
 
 _HandlerCallback = Callable[
     ["InstructionTranslator", typing.Any, typing.Any], VariableTracker | None
 ]
-_TrackersType = Union[type[VariableTracker], tuple[type[VariableTracker], ...]]
+_TrackersType = type[VariableTracker] | tuple[type[VariableTracker], ...]
 polyfill_fn_mapping = {
     operator.eq: polyfills.cmp_eq,
     operator.ne: polyfills.cmp_ne,
@@ -208,6 +220,13 @@
     operator.length_hint,
 )
 
+_SET_LIKE_OP_SUPPORT: tuple[type[VariableTracker], ...] = (
+    DictItemsVariable,
+    DictKeysVariable,
+    SetVariable,
+    UserDefinedObjectVariable,
+)
+
 BUILTIN_TO_TENSOR_FN_MAP: dict[Callable[..., Any], Callable[..., Any]] = {}
 
 # These functions represent the r* versions of the above ops
@@ -283,7 +302,59 @@ def __torch_function__(
                     BUILTIN_TO_TENSOR_RFN_MAP[op] = most_recent_func
 
 
-class BuiltinVariable(VariableTracker):
+class BaseBuiltinVariable(VariableTracker):
+    """
+    Common base class for all builtin variable trackers (BuiltinVariable,
+    DictBuiltinVariable, IterBuiltinVariable, and future specialized builtins).
+
+    Provides shared implementations for guard installation, hasattr tracing,
+    and Python-level hashability/equality.
+
+    Specialized subclasses (e.g. DictBuiltinVariable) set `_fn` as a class
+    attribute. BuiltinVariable stores the callable on the instance as `self.fn`
+    and overrides as_python_constant / reconstruct / var_getattr accordingly.
+    """
+
+    _fn: Any = None
+
+    @classmethod
+    def create_with_source(cls, value: Any, source: Source) -> "BaseBuiltinVariable":
+        install_guard(source.make_guard(GuardBuilder.BUILTIN_MATCH))
+        return cls(source=source)
+
+    def as_python_constant(self) -> Any:
+        return self._fn
+
+    def reconstruct(self, codegen: "PyCodegen") -> None:
+        name = self.as_python_constant().__name__
+        assert name not in codegen.tx.f_globals, "shadowed global"
+        codegen.append_output(codegen.create_load_global(name, add=True))
+
+    def var_getattr(self, tx: "InstructionTranslator", name: str) -> VariableTracker:
+        source = self.source and AttrSource(self.source, name)
+        attr = getattr(self._fn, name, None)
+        return variables.GetAttrVariable(
+            self, name, py_type=type(attr) if attr is not None else None, source=source
+        )
+
+    def call_obj_hasattr(
+        self, tx: "InstructionTranslator", name: str
+    ) -> ConstantVariable:
+        return VariableTracker.build(tx, hasattr(self.as_python_constant(), name))  # type: ignore[return-value]
+
+    def is_python_hashable(self) -> bool:
+        return True
+
+    def get_python_hash(self) -> int:
+        return hash(self.as_python_constant())
+
+    def is_python_equal(self, other: object) -> bool:
+        return isinstance(other, BaseBuiltinVariable) and (
+            self.as_python_constant() is other.as_python_constant()  # type: ignore[union-attr]
+        )
+
+
+class BuiltinVariable(BaseBuiltinVariable):
     """
     A VariableTracker that represents a built-in value (functions and operators).
     A lot of the code here assumes it will be a function object.
@@ -479,15 +550,11 @@ def _binop_handlers() -> dict[
         # combinations. Handlers are attempted in order, and will be used if the type checks
         # match. They are expected to have the signature:
         # fn(tx, arg0: VariableTracker, arg1: VariableTracker) -> VariableTracker
-        from .functions import BaseUserFunctionVariable, UserFunctionVariable
+        from .functions import BaseUserFunctionVariable
         from .nn_module import NNModuleVariable
         from .tensor import supported_const_comparison_ops
         from .torch import BaseTorchVariable
-        from .user_defined import (
-            UserDefinedClassVariable,
-            UserDefinedObjectVariable,
-            UserDefinedVariable,
-        )
+        from .user_defined import UserDefinedVariable
 
         # Override table contains: op_fn -> [list of handlers]
         op_handlers: dict[Any, list[Any]] = {}
@@ -703,7 +770,7 @@ def expand_list_like(
                 raise_observed_exception(
                     type(exc),
                     tx,
-                    args=list(map(ConstantVariable.create, exc.args)),
+                    args=list(exc.args),
                 )
 
         list_like_expansion_handlers: list[
@@ -726,12 +793,12 @@ def compare_by_value(
                 tx: "InstructionTranslator", a: VariableTracker, b: VariableTracker
             ) -> VariableTracker:
                 try:
-                    return ConstantVariable(op(a.value, b.value))  # type: ignore[attr-defined]
+                    return VariableTracker.build(tx, op(a.value, b.value))  # type: ignore[attr-defined]
                 except TypeError as exc:
                     raise_observed_exception(
                         type(exc),
                         tx,
-                        args=list(map(ConstantVariable.create, exc.args)),
+                        args=list(exc.args),
                     )
 
             result: list[
@@ -750,7 +817,6 @@ def compare_by_value(
                 # time benchmark - add_loop_eager.
                 result = [
                     ((ConstantVariable, ConstantVariable), compare_by_value),
-                    ((EnumVariable, EnumVariable), compare_by_value),
                 ]
 
                 op_var = BuiltinVariable(op)
@@ -789,7 +855,7 @@ def handler(
                 def never(
                     tx: "InstructionTranslator", a: VariableTracker, b: VariableTracker
                 ) -> VariableTracker:
-                    return ConstantVariable(none_result)
+                    return VariableTracker.build(tx, none_result)
 
                 obj_op_none = never
                 none_op_obj = never
@@ -820,40 +886,6 @@ def never(
                 op_var = BuiltinVariable(op)
                 result.extend(
                     [
-                        (
-                            (
-                                (UserFunctionVariable, BuiltinVariable),
-                                (UserFunctionVariable, BuiltinVariable),
-                            ),
-                            lambda tx, a, b: ConstantVariable(op(a.fn, b.fn)),
-                        ),
-                        (
-                            (
-                                NNModuleVariable,
-                                NNModuleVariable,
-                            ),
-                            lambda tx, a, b: ConstantVariable(
-                                op(
-                                    tx.output.get_submodule(a.module_key),
-                                    tx.output.get_submodule(b.module_key),
-                                )
-                            ),
-                        ),
-                        (
-                            (UserDefinedObjectVariable, UserDefinedObjectVariable),
-                            compare_by_value,
-                        ),
-                        (
-                            (UserDefinedClassVariable, UserDefinedClassVariable),
-                            compare_by_value,
-                        ),
-                        (
-                            (
-                                (StreamVariable, EventVariable, ConstantVariable),
-                                (StreamVariable, EventVariable, ConstantVariable),
-                            ),
-                            compare_by_value,
-                        ),
                         (
                             (TensorVariable, VariableTracker),
                             op_var._comparison_with_tensor,
@@ -878,22 +910,15 @@ def handle_is(
                     left: VariableTracker,
                     right: VariableTracker,
                 ) -> VariableTracker | None:
-                    # If the two objects are of different type, we can safely return False
-                    # and True for `is` and `is not`, respectively
-                    if type(left) is not type(right):
-                        return ConstantVariable.create(op.__name__ != "is_")
-                    if left is right:
-                        return ConstantVariable.create(op(left, right))
-                    if istype(left, variables.ObjectVariable) and istype(
-                        right, variables.ObjectVariable
-                    ):
-                        return ConstantVariable.create(op(left.value, right.value))
-                    if (
-                        istype(left, variables.ExceptionVariable)
-                        and istype(right, variables.ExceptionVariable)
-                        and left.exc_type is not right.exc_type
-                    ):
-                        return ConstantVariable.create(op(left, right))
+                    from .object_protocol import vt_identity_compare
+
+                    result = vt_identity_compare(left, right)
+                    if result is None:
+                        return None
+                    is_same = result.as_python_constant()
+                    return VariableTracker.build(
+                        tx, is_same if op.__name__ == "is_" else not is_same
+                    )
 
                 result.append(((VariableTracker, VariableTracker), handle_is))  # type: ignore[arg-type]
 
@@ -923,7 +948,15 @@ def _find_binop_handler(
     def can_insert_in_graph(self) -> bool:
         return self.fn in self._fx_graph_functions()
 
+    # Builtins that have been promoted to their own VT classes. Creating a
+    # BuiltinVariable for these is a bug; use the specialized class instead.
+    MUST_USE_SPECIALIZED: frozenset[Any] = frozenset({dict, iter, list})
+
     def __init__(self, fn: Any, **kwargs: Any) -> None:
+        assert fn not in self.MUST_USE_SPECIALIZED, (
+            f"Use the specialized VT class for {fn!r}, not BuiltinVariable. "
+            f"E.g. DictBuiltinVariable for dict."
+        )
         super().__init__(**kwargs)
         self.fn = fn
 
@@ -938,6 +971,9 @@ def __repr__(self) -> str:
     def as_python_constant(self) -> Any:
         return self.fn
 
+    def get_real_python_backed_value(self) -> Any:
+        return self.fn
+
     def as_proxy(self) -> Any:
         DTYPE = {
             bool: torch.bool,
@@ -1009,7 +1045,7 @@ def _make_handler(
     ) -> Callable[
         [
             "InstructionTranslator",
-            tuple[VariableTracker, ...],
+            list[VariableTracker],
             dict[str, VariableTracker],
         ],
         VariableTracker | None,
@@ -1034,7 +1070,7 @@ def _make_handler(
 
             def create_exception_class_object(
                 tx: "InstructionTranslator",
-                args: tuple[VariableTracker, ...],
+                args: list[VariableTracker],
                 kwargs: dict[str, VariableTracker],
             ) -> VariableTracker:
                 if fn is AssertionError and not all(
@@ -1149,7 +1185,7 @@ def constant_fold_handler(
                         raise_observed_exception(
                             type(exc),
                             tx,
-                            args=list(map(ConstantVariable.create, exc.args)),
+                            args=list(exc.args),
                         )
                     except AsPythonConstantNotImplementedError as exc:
                         unimplemented(
@@ -1159,7 +1195,6 @@ def constant_fold_handler(
                             hints=[*graph_break_hints.DYNAMO_BUG],
                             from_exc=exc,
                         )
-                    # pyrefly: ignore [unbound-name]
                     return VariableTracker.build(tx, res)
 
             else:
@@ -1190,9 +1225,8 @@ def constant_fold_handler(
                             raise_observed_exception(
                                 type(exc),
                                 tx,
-                                args=list(map(ConstantVariable.create, exc.args)),
+                                args=list(exc.args),
                             )
-                        # pyrefly: ignore [unbound-name]
                         return VariableTracker.build(tx, res)
                     return None
 
@@ -1249,12 +1283,7 @@ def builtin_dispatch(
 
     def call_vars(self, tx: "InstructionTranslator", *args: Any) -> VariableTracker:
         if len(args) == 0:
-            unimplemented(
-                gb_type="unimplemented builtin op vars() with no arguments",
-                context=f"vars: {self} {args}",
-                explanation=f"Dynamo does not know how to trace builtin operator {self.fn} with no arguments",
-                hints=[*graph_break_hints.SUPPORTABLE],
-            )
+            return self._call_frame_locals_snapshot(tx)
         assert len(args) == 1
         # vars(obj) is obj.__dict__ if __dict__ is present else TypeError
         try:
@@ -1262,6 +1291,34 @@ def call_vars(self, tx: "InstructionTranslator", *args: Any) -> VariableTracker:
         except ObservedAttributeError:
             raise_observed_exception(TypeError, tx)
 
+    def call_locals(
+        self, tx: "InstructionTranslator", *args: VariableTracker
+    ) -> VariableTracker:
+        if len(args) != 0:
+            raise_observed_exception(TypeError, tx)
+        return self._call_frame_locals_snapshot(tx)
+
+    @staticmethod
+    def _call_frame_locals_snapshot(tx: "InstructionTranslator") -> VariableTracker:
+        frame_local_names = set(tx.f_code.co_varnames) | set(tx.cell_and_freevars())
+        cell_and_freevars = set(tx.cell_and_freevars())
+        frame_locals = {}
+        for name, value in tx.symbolic_locals.items():
+            if name not in frame_local_names:
+                continue
+            if name in cell_and_freevars:
+                value = tx.output.side_effects.load_cell(value)
+            if type.__instancecheck__(NullVariable, value) or isinstance(
+                value, variables.DeletedVariable
+            ):
+                continue
+            frame_locals[ConstantVariable.create(name)] = value
+        return ConstDictVariable(
+            frame_locals,
+            dict,
+            mutation_type=ValueMutationNew(),
+        )
+
     def _handle_insert_op_in_graph(
         self,
         tx: "InstructionTranslator",
@@ -1334,7 +1391,7 @@ def _handle_insert_op_in_graph(
                 fn = torch.select
                 args = [
                     args[0],
-                    variables.ConstantVariable.create(0),
+                    variables.VariableTracker.build(tx, 0),
                     args[1],
                 ]  # type: ignore[assignment]
 
@@ -1466,19 +1523,6 @@ def call_method(
                     self, args[0], args[1:]
                 )
 
-            if self.fn is dict and len(args) == 1 and not kwargs:
-                dict_vt = ConstDictVariable({}, dict, mutation_type=ValueMutationNew())
-                if isinstance(args[0], BuiltinVariable) and args[0].fn is dict:
-                    return dict_vt
-                # We don't have to set the underlying dict_vt in
-                # UserDefinedDictVariable because it will be set to empty
-                # ConstDictVariableTracker in the constructor.
-                return tx.output.side_effects.track_new_user_defined_object(
-                    self,
-                    args[0],
-                    args[1:],
-                )
-
             if (
                 self.fn is tuple
                 and len(args) == 2
@@ -1497,56 +1541,34 @@ def call_method(
                     args[1:],
                 )
 
-            if self.fn is list:
-                list_vt = ListVariable([], mutation_type=ValueMutationNew())
-                if isinstance(args[0], BuiltinVariable) and args[0].fn is list:
-                    return list_vt
-                return tx.output.side_effects.track_new_user_defined_object(
-                    self,
-                    args[0],
-                    args[1:],
-                )
-
-        if (
-            self.fn in (float, complex)
-            and len(args) == 1
-            and (
-                (self.fn is float and name in ("fromhex", "hex"))
-                or (name == "from_number" and sys.version_info >= (3, 14))
-            )
-        ):
-            if args[0].is_python_constant():
+        if name in _BUILTIN_CONSTANT_FOLDABLE_METHODS.get(self.fn, ()):
+            if all(a.is_python_constant() for a in args) and all(
+                v.is_python_constant() for v in kwargs.values()
+            ):
                 try:
                     fn = getattr(self.fn, name)
-                    res = fn(args[0].as_python_constant())
-                    return variables.ConstantVariable.create(res)
-                except (OverflowError, ValueError) as e:
+                    res = fn(
+                        *(a.as_python_constant() for a in args),
+                        **{k: v.as_python_constant() for k, v in kwargs.items()},
+                    )
+                    return VariableTracker.build(tx, res)
+                except Exception as e:
                     raise_observed_exception(
                         type(e),
                         tx,
-                        args=list(map(ConstantVariable.create, e.args)),
+                        args=list(e.args),
                     )
 
         if self.fn is object and name == "__init__":
             # object.__init__ is a no-op
-            return variables.CONSTANT_VARIABLE_NONE
-
-        if self.fn is dict and name == "fromkeys":
-            return BuiltinVariable.call_custom_dict_fromkeys(tx, dict, *args, **kwargs)
-
-        if self.fn is dict:
-            resolved_fn = getattr(self.fn, name)
-            if resolved_fn in dict_methods:
-                if isinstance(args[0], variables.UserDefinedDictVariable):
-                    return args[0]._dict_vt.call_method(tx, name, args[1:], kwargs)
-                elif isinstance(args[0], variables.ConstDictVariable):
-                    return args[0].call_method(tx, name, args[1:], kwargs)
+            return variables.ConstantVariable.create(None)
 
         if self.fn is set:
             resolved_fn = getattr(self.fn, name)
             if resolved_fn in set_methods:
                 if isinstance(args[0], variables.UserDefinedSetVariable):
-                    return args[0]._set_vt.call_method(tx, name, args[1:], kwargs)
+                    assert args[0]._base_vt is not None
+                    return args[0]._base_vt.call_method(tx, name, args[1:], kwargs)
                 elif isinstance(args[0], variables.SetVariable):
                     return args[0].call_method(tx, name, args[1:], kwargs)
 
@@ -1566,73 +1588,60 @@ def call_method(
         if self.fn is float and len(args) >= 1:
             # Only delegate to ConstantVariable, not other types that happen to be constants
             if isinstance(args[0], ConstantVariable):
-                return ConstantVariable.create(
-                    getattr(float, name)(args[0].as_python_constant())
+                return VariableTracker.build(
+                    tx, getattr(float, name)(args[0].as_python_constant())
                 )
 
+        if name == "__len__" and len(args) == 1 and not kwargs:
+            # type.__len__(instance) → len(instance)
+            # e.g. list.__len__(my_list) → len(my_list)
+            from .object_protocol import generic_len
+
+            return generic_len(tx, args[0])
+
         return super().call_method(tx, name, args, kwargs)
 
-    def _call_int_float(
+    def call_int(
         self, tx: "InstructionTranslator", arg: VariableTracker
     ) -> VariableTracker | None:
-        # Handle cases like int(torch.seed())
-        # Also handle sym_float to sym_int cases
-        if arg.is_tensor() or isinstance(arg, SymNodeVariable):
-            if arg.is_tensor():
-                item = arg.call_method(tx, "item", [], {})
-            else:
-                item = arg
-            fn_ = sym_int if self.fn is int else sym_float
-            from torch._dynamo.variables.builder import wrap_fx_proxy
+        from .object_protocol import generic_int
 
-            return wrap_fx_proxy(
-                tx=tx,
-                proxy=tx.output.create_proxy(
-                    "call_function",
-                    fn_,
-                    (item.as_proxy(),),
-                    {},
-                ),
-            )
-        return None
+        return generic_int(tx, arg)
+
+    def call_float(
+        self, tx: "InstructionTranslator", arg: VariableTracker
+    ) -> VariableTracker | None:
+        from .object_protocol import generic_float
 
-    call_int = _call_int_float
-    call_float = _call_int_float
+        return generic_float(tx, arg)
 
     def call_bool(
         self, tx: "InstructionTranslator", arg: VariableTracker
     ) -> VariableTracker | None:
-        if arg.is_tensor():
-            item = arg.call_method(tx, "item", [], {})
-            if isinstance(item, SymNodeVariable) and isinstance(
-                item.sym_num, torch.SymBool
-            ):
-                return item
-            if isinstance(item, variables.ConstantVariable):
-                return variables.ConstantVariable.create(bool(item.value))
-            return SymNodeVariable.create(tx, item.as_proxy() != 0)
-        # Emulate `PyBool_Type.tp_vectorcall` which boils down to `PyObject_IsTrue`.
-        # https://github.com/python/cpython/blob/3.12/Objects/object.c#L1674-L1697
-        if isinstance(arg, SymNodeVariable):
-            # Note that we delay specializing on symbolic values to avoid
-            # unnecessary guards. Specialization will happen later if, e.g., the
-            # resulting boolean is used for branching.
-            if isinstance(arg.sym_num, torch.SymBool):
-                return arg
-
-            # Emulate `nb_bool` of int/float objects
-            # - https://github.com/python/cpython/blob/3.12/Objects/longobject.c#L4940-L4944
-            # - https://github.com/python/cpython/blob/3.12/Objects/floatobject.c#L878-L882
-            assert istype(arg.sym_num, (torch.SymInt, torch.SymFloat))
-            return SymNodeVariable.create(tx, arg.as_proxy() != 0)
-
-        # TODO handle more cases and merge this with this with `generic_jump`.
-        return None
+        # Emulate PyBool_Type.tp_vectorcall which boils down to PyObject_IsTrue.
+        from .object_protocol import generic_bool
+
+        return generic_bool(tx, arg)
 
     def call_repr(
         self, tx: "InstructionTranslator", arg: VariableTracker
     ) -> VariableTracker | None:
         """Handle repr() on user defined objects."""
+        if isinstance(
+            arg,
+            (variables.ExceptionVariable, variables.UserDefinedExceptionObjectVariable),
+        ):
+            try:
+                const_args = tuple(a.as_python_constant() for a in arg.args)
+            except NotImplementedError:
+                return None
+            if len(const_args) == 0:
+                value = f"{arg.exc_type.__name__}()"
+            elif len(const_args) == 1:
+                value = f"{arg.exc_type.__name__}({const_args[0]!r})"
+            else:
+                value = f"{arg.exc_type.__name__}{const_args!r}"
+            return VariableTracker.build(tx, value)
         if isinstance(arg, variables.UserDefinedObjectVariable):
             repr_method = arg.value.__repr__
 
@@ -1653,7 +1662,7 @@ def call_repr(
                 return fn_vt.call_function(tx, [arg], {})
         if isinstance(arg, variables.UserDefinedClassVariable):
             if type(arg.value).__repr__ is type.__repr__:
-                return variables.ConstantVariable.create(repr(arg.value))
+                return VariableTracker.build(tx, repr(arg.value))
         if isinstance(
             arg,
             (
@@ -1664,15 +1673,27 @@ def call_repr(
                 DictViewVariable,
             ),
         ):
-            return variables.ConstantVariable.create(arg.debug_repr())
+            return VariableTracker.build(tx, arg.debug_repr())
         return None
 
     def call_str(
         self, tx: "InstructionTranslator", arg: VariableTracker
     ) -> VariableTracker | None:
+        if isinstance(
+            arg,
+            (variables.ExceptionVariable, variables.UserDefinedExceptionObjectVariable),
+        ):
+            if len(arg.args) == 0:
+                return VariableTracker.build(tx, "")
+            elif len(arg.args) == 1:
+                return BuiltinVariable(str).call_function(tx, [arg.args[0]], {})
+            else:
+                tuple_var = variables.TupleVariable(list(arg.args))
+                return BuiltinVariable(str).call_function(tx, [tuple_var], {})
+
         # Handle `str` on a user defined function or object
         if isinstance(arg, (variables.UserFunctionVariable)):
-            return variables.ConstantVariable.create(value=str(arg.fn))
+            return VariableTracker.build(tx, str(arg.fn))
         elif isinstance(arg, (variables.UserDefinedObjectVariable)):
             # Check if object has __str__ method
             if hasattr(arg.value, "__str__"):
@@ -1692,11 +1713,10 @@ def call_str(
                 # Rely on the object str method
                 try:
                     # pyrefly: ignore [unbound-name]
-                    return variables.ConstantVariable.create(value=str_method())
+                    return VariableTracker.build(tx, str_method())
                 except AttributeError:
                     # Graph break
                     return None
-            # pyrefly: ignore [unbound-name]
             elif is_wrapper_or_member_descriptor(str_method):
                 unimplemented(
                     gb_type="Attempted to a str() method implemented in C/C++",
@@ -1719,12 +1739,6 @@ def call_str(
 
                 # Inline the user function
                 return user_func_variable.call_function(tx, [arg], {})
-        elif isinstance(arg, (variables.ExceptionVariable,)):
-            if len(arg.args) == 0:
-                value = f"{arg.exc_type}"
-            else:
-                value = ", ".join(a.as_python_constant() for a in arg.args)
-            return variables.ConstantVariable.create(value=value)
         return None
 
     def _call_min_max(
@@ -1853,7 +1867,7 @@ def _call_min_max_binary(
                 a.as_python_constant(),
                 b.as_python_constant(),
             )
-            return ConstantVariable.create(value)
+            return VariableTracker.build(tx, value)
         return None
 
     call_min = _call_min_max
@@ -1866,7 +1880,7 @@ def call_abs(
 
         # Call arg.__abs__()
         abs_method = SourcelessBuilder.create(tx, getattr).call_function(
-            tx, [arg, ConstantVariable.create("__abs__")], {}
+            tx, [arg, VariableTracker.build(tx, "__abs__")], {}
         )
         return abs_method.call_function(tx, [], {})
 
@@ -1877,24 +1891,17 @@ def call_pos(
 
         # Call arg.__pos__()
         pos_method = SourcelessBuilder.create(tx, getattr).call_function(
-            tx, [arg, ConstantVariable.create("__pos__")], {}
+            tx, [arg, VariableTracker.build(tx, "__pos__")], {}
         )
         return pos_method.call_function(tx, [], {})
 
     def call_index(
         self, tx: "InstructionTranslator", arg: VariableTracker
     ) -> VariableTracker:
-        if arg.is_tensor():
-            unimplemented(
-                gb_type="unsupported index(Tensor)",
-                context="",
-                explanation="Dynamo does not support tracing builtin index() on a Tensor",
-                hints=[],
-            )
-
-        arg = guard_if_dyn(arg)
-        constant_value = operator.index(arg)
-        return variables.ConstantVariable.create(constant_value)
+        # Specialize SymNodeVariable to a constant first, matching CPython's
+        # PyNumber_Index which forces a concrete int.
+        arg = specialize_symnode(arg)
+        return arg.nb_index_impl(tx)
 
     def call_round(
         self,
@@ -1907,7 +1914,7 @@ def call_round(
 
         # Call arg.__round__()
         round_method = SourcelessBuilder.create(tx, getattr).call_function(
-            tx, [arg, ConstantVariable.create("__round__")], {}
+            tx, [arg, VariableTracker.build(tx, "__round__")], {}
         )
         return round_method.call_function(tx, args, kwargs)
 
@@ -1917,9 +1924,7 @@ def call_range(
         if check_unspec_or_constant_args(args, {}):
             return variables.RangeVariable(args)
         elif self._dynamic_args(*args):
-            args = tuple(
-                variables.ConstantVariable.create(guard_if_dyn(arg)) for arg in args
-            )
+            args = tuple(VariableTracker.build(tx, guard_if_dyn(arg)) for arg in args)
             return variables.RangeVariable(args)
         # None no-ops this handler and lets the driving function proceed
         return None
@@ -1972,10 +1977,13 @@ def _call_iter_tuple_list(
                         obj.source.make_guard(GuardBuilder.TUPLE_ITERATOR_LEN)
                     )
                 else:
-                    if (
-                        getattr(obj, "source", False)
-                        and isinstance(obj, ConstDictVariable)
-                        and not istype(obj, (SetVariable, FrozensetVariable))
+                    if getattr(obj, "source", False) and isinstance(
+                        obj,
+                        (
+                            ConstDictVariable,
+                            variables.OrderedSetVariable,
+                            variables.DictKeySetVariable,
+                        ),
                     ):
                         tx.output.guard_on_key_order.add(obj.source)
 
@@ -2032,48 +2040,7 @@ def _call_tuple_list(
         else:
             return self._call_iter_tuple_list(tx, obj, *args, **kwargs)
 
-    def call_iter(
-        self,
-        tx: "InstructionTranslator",
-        obj: VariableTracker,
-        *args: VariableTracker,
-        **kwargs: VariableTracker,
-    ) -> VariableTracker:
-        # avoid the overhead of tracing the polyfill if we already know the class implemented __iter__
-        if isinstance(
-            obj,
-            (
-                variables.ListVariable,
-                variables.RangeVariable,
-                variables.IteratorVariable,
-                variables.ConstDictVariable,
-                variables.NNModuleVariable,
-                variables.TensorVariable,
-                variables.TupleVariable,
-                DictViewVariable,
-            ),
-        ):
-            return obj.call_method(tx, "__iter__", [], {})
-        else:
-            # If the object doesn't implement a __iter__ method, it will be an error in eager mode when calling iter on it anyway.
-            # If the object implements a __iter__ method, inlining effectively forwards the call to another iter call
-            # (e.g. when __iter__ just returns iter(self.list)) or return a user-defined iterator.
-            # If the object implements a __getitem__ method, iter(...) will call obj.__getitem__()
-            # with an integer argument starting at 0, until __getitem__ raises IndexError
-            ret = variables.UserFunctionVariable(
-                polyfills.builtins.iter_  # type: ignore[arg-type]
-            ).call_function(tx, [obj, *args], {})
-
-            if args:
-                # iter(obj, sentinel) returns an object that implements
-                # __iter__ and __next__ methods (UserDefinedObjectVariable)
-                # Wrap the return value in a IteratorVariable subclass (LazyObjectIteratorVariable)
-                # that forwards the next_variable call to the object.
-                ret = variables.ObjectIteratorVariable(ret)
-            return ret
-
     call_tuple = _call_tuple_list
-    call_list = _call_tuple_list
 
     def call_callable(
         self, tx: "InstructionTranslator", arg: VariableTracker
@@ -2092,7 +2059,7 @@ def call_callable(
         ):
             return variables.ConstantVariable.create(True)
         elif isinstance(arg, UserDefinedVariable):
-            return variables.ConstantVariable.create(callable(arg.value))
+            return VariableTracker.build(tx, callable(arg.value))
         elif isinstance(
             arg,
             (
@@ -2128,214 +2095,93 @@ def call_dir(
             return VariableTracker.build(tx, dir(arg.value))
         if isinstance(arg, BuiltinVariable):
             return VariableTracker.build(tx, dir(arg.fn))
+        # Enable specialized VTs for constants to work with dir()
+        if arg.is_python_constant():
+            return VariableTracker.build(tx, dir(arg.as_python_constant()))
         return None
 
-    def call_dict(
+    def call_set(
         self,
         tx: "InstructionTranslator",
-        /,
         *args: VariableTracker,
         **kwargs: VariableTracker,
     ) -> VariableTracker:
-        return BuiltinVariable.call_custom_dict(tx, dict, *args, **kwargs)
+        from .builder import SourcelessBuilder
 
-    @staticmethod
-    def call_custom_dict(
+        assert not kwargs
+        if not args:
+            return SetVariable([], mutation_type=ValueMutationNew())
+        if len(args) != 1:
+            raise_observed_exception(
+                TypeError,
+                tx,
+                args=[f"set() takes 1 positional argument but {len(args)} were given"],
+            )
+        arg = args[0]
+        if istype(arg, variables.SetVariable):
+            return arg.clone(mutation_type=ValueMutationNew())
+        elif arg.has_force_unpack_var_sequence(tx):
+            items = arg.force_unpack_var_sequence(tx)
+            return SetVariable(items, mutation_type=ValueMutationNew())
+        elif isinstance(arg, variables.UserDefinedObjectVariable) and isinstance(
+            arg.value, KeysView
+        ):
+            iter_fn = arg.var_getattr(tx, "__iter__")
+            if isinstance(iter_fn, variables.UserMethodVariable):
+                out = tx.inline_user_function_return(iter_fn, args, kwargs)
+                if isinstance(out, SetVariable):
+                    return out
+                return SourcelessBuilder.create(tx, set).call_set(tx, out)
+        raise_observed_exception(
+            TypeError,
+            tx,
+            args=["failed to construct builtin set()"],
+        )
+
+    def call_frozenset(
+        self,
         tx: "InstructionTranslator",
-        user_cls: type,
-        /,
         *args: VariableTracker,
         **kwargs: VariableTracker,
     ) -> VariableTracker:
-        args_list = list(args)
-        if (
-            len(args_list) == 1
-            and isinstance(args_list[0], variables.GetAttrVariable)
-            and isinstance(args_list[0].obj, variables.UserDefinedClassVariable)
-            and not tx.output.side_effects.has_pending_mutation(args_list[0].obj)
-        ):
-            # Forward the GetAttrVariable(foo, "__dict__") to a realized vt of
-            # VT(foo.__dict__). This simplifies the construction of the new
-            # dict.
-            args_list[0] = args_list[0].get_forwarded_dict(tx)
-        return tx.inline_user_function_return(
-            VariableTracker.build(tx, polyfills.construct_dict),
-            [VariableTracker.build(tx, user_cls), *args_list],
-            kwargs,
+        assert not kwargs
+        if not args:
+            return FrozensetVariable([])
+        if len(args) != 1:
+            raise_observed_exception(
+                TypeError,
+                tx,
+                args=[
+                    f"frozenset() takes 1 positional argument but {len(args)} were given"
+                ],
+            )
+        arg = args[0]
+        if istype(arg, variables.FrozensetVariable):
+            return FrozensetVariable([x.vt for x in arg.set_items])
+        elif arg.has_force_unpack_var_sequence(tx):
+            items = arg.force_unpack_var_sequence(tx)
+            return FrozensetVariable(items)
+        raise_observed_exception(
+            TypeError,
+            tx,
+            args=["failed to construct builtin frozenset()"],
         )
 
-    @staticmethod
-    def call_custom_dict_fromkeys(
+    def call_zip(
+        self,
         tx: "InstructionTranslator",
-        user_cls: type,
-        /,
         *args: VariableTracker,
         **kwargs: VariableTracker,
     ) -> VariableTracker:
-        if user_cls not in {dict, OrderedDict, defaultdict}:
-            unimplemented(
-                gb_type="Unsupported dict type for fromkeys()",
-                context=f"{user_cls.__name__}.fromkeys(): {args} {kwargs}",
-                explanation=f"Failed to call {user_cls.__name__}.fromkeys() because "
-                f"{user_cls.__name__} is not any type of dict, OrderedDict, or defaultdict",
-                hints=[
-                    f"Ensure {user_cls.__name__} is a type of dict, OrderedDict, or defaultdict.",
-                ],
-            )
+        from .builder import SourcelessBuilder
+
         if kwargs:
-            # Only `OrderedDict.fromkeys` accepts `value` passed by keyword
-            if (
-                user_cls is not OrderedDict
-                or len(args) != 1
-                or len(kwargs) != 1
-                or "value" not in kwargs
-            ):
+            if not (len(kwargs) == 1 and "strict" in kwargs):
                 raise_args_mismatch(
                     tx,
-                    f"{user_cls.__name__}.fromkeys",
-                    "1 args and 1 kwargs (`value`)",
-                    f"{len(args)} args and {len(kwargs)} kwargs",
-                )
-            args = (*args, kwargs.pop("value"))
-        if len(args) == 0:
-            raise_args_mismatch(
-                tx,
-                f"{user_cls.__name__}.fromkeys",
-                "at least 1 args",
-                f"{len(args)} args",
-            )
-        if len(args) == 1:
-            args = (*args, CONSTANT_VARIABLE_NONE)
-        if len(args) != 2:
-            raise_args_mismatch(
-                tx,
-                f"{user_cls.__name__}.fromkeys",
-                "2 args",
-                f"{len(args)} args",
-            )
-
-        arg, value = args
-        DictVariableType = (
-            ConstDictVariable if user_cls is not defaultdict else DefaultDictVariable
-        )
-
-        if isinstance(arg, dict):
-            arg_list = [ConstantVariable.create(k) for k in arg]
-            return DictVariableType(
-                dict.fromkeys(arg_list, value),
-                user_cls,
-                mutation_type=ValueMutationNew(),
-            )
-        elif arg.has_force_unpack_var_sequence(tx):
-            keys = arg.force_unpack_var_sequence(tx)
-            if all(is_hashable(v) for v in keys):
-                return DictVariableType(
-                    dict.fromkeys(keys, value),
-                    user_cls,
-                    mutation_type=ValueMutationNew(),
-                )
-
-        unimplemented(
-            gb_type="failed to call dict.fromkeys()",
-            context=f"{user_cls.__name__}.fromkeys(): {args} {kwargs}",
-            explanation=f"Failed to call {user_cls.__name__}.fromkeys() because "
-            "arguments could not be automatically converted to a list, "
-            "or some dict key is not hashable.",
-            hints=[
-                "Manually convert the argument to a list.",
-                "Ensure all keys are hashable.",
-            ],
-        )
-
-    def call_set(
-        self,
-        tx: "InstructionTranslator",
-        *args: VariableTracker,
-        **kwargs: VariableTracker,
-    ) -> VariableTracker:
-        from .builder import SourcelessBuilder
-
-        # Can we merge this implementation and call_dict's one?
-        assert not kwargs
-        if not args:
-            return SetVariable([], mutation_type=ValueMutationNew())
-        if len(args) != 1:
-            raise_observed_exception(
-                TypeError,
-                tx,
-                args=[
-                    ConstantVariable.create(
-                        f"set() takes 1 positional argument but {len(args)} were given"
-                    )
-                ],
-            )
-        arg = args[0]
-        if istype(arg, variables.SetVariable):
-            return arg.clone(mutation_type=ValueMutationNew())
-        elif arg.has_force_unpack_var_sequence(tx):
-            items = arg.force_unpack_var_sequence(tx)
-            return SetVariable(items, mutation_type=ValueMutationNew())
-        elif isinstance(arg, variables.UserDefinedObjectVariable) and isinstance(
-            arg.value, KeysView
-        ):
-            iter_fn = arg.var_getattr(tx, "__iter__")
-            if isinstance(iter_fn, variables.UserMethodVariable):
-                out = tx.inline_user_function_return(iter_fn, args, kwargs)
-                if isinstance(out, SetVariable):
-                    return out
-                return SourcelessBuilder.create(tx, set).call_set(tx, out)
-        raise_observed_exception(
-            TypeError,
-            tx,
-            args=[ConstantVariable.create("failed to construct builtin set()")],
-        )
-
-    def call_frozenset(
-        self,
-        tx: "InstructionTranslator",
-        *args: VariableTracker,
-        **kwargs: VariableTracker,
-    ) -> VariableTracker:
-        assert not kwargs
-        if not args:
-            return FrozensetVariable([])
-        if len(args) != 1:
-            raise_observed_exception(
-                TypeError,
-                tx,
-                args=[
-                    ConstantVariable.create(
-                        f"frozenset() takes 1 positional argument but {len(args)} were given"
-                    )
-                ],
-            )
-        arg = args[0]
-        if istype(arg, variables.FrozensetVariable):
-            return FrozensetVariable([x.vt for x in arg.set_items])
-        elif arg.has_force_unpack_var_sequence(tx):
-            items = arg.force_unpack_var_sequence(tx)
-            return FrozensetVariable(items)
-        raise_observed_exception(
-            TypeError,
-            tx,
-            args=[ConstantVariable.create("failed to construct builtin frozenset()")],
-        )
-
-    def call_zip(
-        self,
-        tx: "InstructionTranslator",
-        *args: VariableTracker,
-        **kwargs: VariableTracker,
-    ) -> VariableTracker:
-        from .builder import SourcelessBuilder
-
-        if kwargs:
-            if not (len(kwargs) == 1 and "strict" in kwargs):
-                raise_args_mismatch(
-                    tx,
-                    "zip",
-                    "1 kwargs (`strict`)",
-                    f"{len(kwargs)} kwargs",
+                    "zip",
+                    "1 kwargs (`strict`)",
+                    f"{len(kwargs)} kwargs",
                 )
         strict = kwargs.pop("strict", ConstantVariable.create(False))
         iter_args = [
@@ -2354,10 +2200,9 @@ def call_len(
         *args: VariableTracker,
         **kwargs: VariableTracker,
     ) -> VariableTracker:
-        try:
-            return args[0].call_method(tx, "__len__", list(args[1:]), kwargs)
-        except AttributeError as e:
-            raise_observed_exception(type(e), tx, args=list(e.args))
+        from .object_protocol import generic_len
+
+        return generic_len(tx, args[0])
 
     def call_getitem(
         self,
@@ -2365,7 +2210,9 @@ def call_getitem(
         *args: VariableTracker,
         **kwargs: VariableTracker,
     ) -> VariableTracker:
-        return args[0].call_method(tx, "__getitem__", list(args[1:]), kwargs)
+        from .object_protocol import vt_getitem
+
+        return vt_getitem(tx, args[0], args[1])
 
     def call_isinstance(
         self,
@@ -2411,9 +2258,7 @@ def check_type(ty: Any) -> bool:
                 else:
                     return check_type(tensor_type)
 
-            return variables.ConstantVariable.create(
-                _tensor_isinstance(arg, isinstance_type)
-            )
+            return VariableTracker.build(tx, _tensor_isinstance(arg, isinstance_type))
         # UserDefinedObject with C extensions can have torch.Tensor attributes,
         # so break graph.
         if isinstance(arg, variables.UserDefinedObjectVariable) and isinstance(
@@ -2431,13 +2276,14 @@ def check_type(ty: Any) -> bool:
             isinstance(arg, variables.UserDefinedObjectVariable)
             and "__instancecheck__" in isinstance_type.__class__.__dict__
         ):
-            return variables.ConstantVariable.create(
-                isinstance_type.__class__.__instancecheck__(isinstance_type, arg.value)
+            return VariableTracker.build(
+                tx,
+                isinstance_type.__class__.__instancecheck__(isinstance_type, arg.value),
             )
 
         if isinstance(arg, variables.UserDefinedExceptionClassVariable):
             # pyrefly: ignore [unbound-name]
-            return ConstantVariable.create(isinstance(arg_type, isinstance_type))
+            return VariableTracker.build(tx, isinstance(arg_type, isinstance_type))
 
         isinstance_type_tuple: tuple[type, ...]
         if isinstance(isinstance_type, type) or callable(
@@ -2469,12 +2315,10 @@ def check_type(ty: Any) -> bool:
             # through it. This is a limitation of the current implementation.
             # Usually `__subclasscheck__` and `__instancecheck__` can be constant fold through, it
             # might not be a big issue and we trade off it for performance.
-            # pyrefly: ignore [unbound-name]
             val = issubclass(arg_type, isinstance_type_tuple)
         except TypeError:
-            # pyrefly: ignore [unbound-name]
             val = arg_type in isinstance_type_tuple
-        return variables.ConstantVariable.create(val)
+        return VariableTracker.build(tx, val)
 
     def call_issubclass(
         self,
@@ -2500,7 +2344,7 @@ def call_issubclass(
         # WARNING: This might run arbitrary user code `__subclasscheck__`.
         # See the comment in call_isinstance above.
         # pyrefly: ignore [unbound-name]
-        return variables.ConstantVariable(issubclass(left_ty_py, right_ty_py))
+        return VariableTracker.build(tx, issubclass(left_ty_py, right_ty_py))
 
     def call_super(
         self, tx: "InstructionTranslator", a: VariableTracker, b: VariableTracker
@@ -2528,8 +2372,6 @@ def call_hasattr(
     ) -> VariableTracker | None:
         if attr.is_python_constant():
             name = attr.as_python_constant()
-            if isinstance(obj, variables.BuiltinVariable):
-                return variables.ConstantVariable(hasattr(obj.fn, name))
             return obj.call_obj_hasattr(tx, name)
         return None
 
@@ -2584,17 +2426,20 @@ def call_filter(
 
     def var_getattr(self, tx: "InstructionTranslator", name: str) -> VariableTracker:
         source = self.source and AttrSource(self.source, name)
+        if name == "__name__":
+            return VariableTracker.build(tx, self.fn.__name__, source)
         if self.fn is object:
             # for object, we can just directly read the attribute
             try:
                 value = getattr(self.fn, name)
             except AttributeError:
                 raise_observed_exception(AttributeError, tx)
-            # pyrefly: ignore [unbound-name]
             if not callable(value):
-                # pyrefly: ignore [unbound-name]
                 return VariableTracker.build(tx, value, source)
-        return variables.GetAttrVariable(self, name, source=source)
+        attr = getattr(self.fn, name, None)
+        return variables.GetAttrVariable(
+            self, name, py_type=type(attr) if attr is not None else None, source=source
+        )
 
     def call_getattr(
         self,
@@ -2668,7 +2513,7 @@ def call_getattr(
                     if name == "__base__":
                         return VariableTracker.build(tx, value.__base__, source)
                     if name == "__flags__":
-                        return ConstantVariable.create(value.__flags__)
+                        return VariableTracker.build(tx, value.__flags__)
             except NotImplementedError:
                 pass
 
@@ -2678,7 +2523,6 @@ def call_getattr(
             obj,
             (
                 variables.TensorVariable,
-                variables.NamedTupleVariable,
                 variables.ConstantVariable,
                 variables.DefaultDictVariable,
                 variables.DistributedVariable,
@@ -2732,13 +2576,20 @@ def call_getattr(
                 return variables.GetAttrVariable(obj, name, source=source)
         elif isinstance(obj, variables.TorchInGraphFunctionVariable):
             # Get OpOverload from an OpOverloadPacket, e.g., torch.ops.aten.add.default.
-            member = getattr(obj.value, name)
+            try:
+                member = getattr(obj.value, name)
+            except AttributeError:
+                raise_observed_exception(AttributeError, tx)
+                raise
+
             if isinstance(
                 member, (torch._ops.OpOverloadPacket, torch._ops.OpOverload)
             ) and torch._dynamo.trace_rules.is_aten_op_or_tensor_method(member):
                 return variables.TorchInGraphFunctionVariable(member, source=source)
             elif name in cmp_name_to_op_mapping:
-                return variables.GetAttrVariable(obj, name, source=source)
+                return variables.GetAttrVariable(
+                    obj, name, py_type=type(member), source=source
+                )
             else:
                 return None
         elif isinstance(obj, DummyModule):
@@ -2751,12 +2602,6 @@ def call_getattr(
             if config.replay_record_enabled:
                 tx.exec_recorder.record_module_access(obj.value, name, member)  # type: ignore[arg-type, union-attr]
             return VariableTracker.build(tx, member, source)
-
-        elif istype(obj, variables.UserFunctionVariable) and name in (
-            "__name__",
-            "__module__",
-        ):
-            return ConstantVariable.create(getattr(obj.fn, name))
         else:
             try:
                 return obj.var_getattr(tx, name)
@@ -2774,8 +2619,6 @@ def call_setattr(
             obj,
             (
                 variables.DefaultDictVariable,
-                variables.PlacementVariable,
-                variables.NamedTupleVariable,
                 variables.UserDefinedObjectVariable,
                 variables.NestedUserFunctionVariable,
                 variables.ExceptionVariable,
@@ -3027,27 +2870,30 @@ def call_id(
         if len(args) > 0 and isinstance(args[0], variables.NNModuleVariable):
             nn_mod_variable = args[0]
             mod = tx.output.get_submodule(nn_mod_variable.module_key)
-            return variables.ConstantVariable.create(id(mod))
-        elif len(args) == 1 and isinstance(
-            args[0],
-            (variables.UserDefinedClassVariable, variables.UserDefinedObjectVariable),
-        ):
-            if args[0].source:
-                if isinstance(args[0], variables.UserDefinedClassVariable):
-                    install_guard(args[0].source.make_guard(GuardBuilder.CLASS_MATCH))
-                else:
-                    install_guard(args[0].source.make_guard(GuardBuilder.ID_MATCH))
-            constant_result = id(args[0].value)
-            return variables.ConstantVariable.create(constant_result)
+            return VariableTracker.build(tx, id(mod))
         elif len(args) == 1 and args[0].is_tensor():
             tensor_variable = cast(TensorVariable, args[0])
             return tensor_variable.call_id(tx)
-        elif istype(args[0], variables.UserFunctionVariable):
-            return variables.ConstantVariable.create(id(args[0].fn))
-        elif istype(args[0], variables.SkipFunctionVariable):
-            return variables.ConstantVariable.create(id(args[0].value))
         elif istype(args[0], variables.FunctoolsPartialVariable):
-            return variables.ConstantVariable.create(id(args[0].fake_value))
+            return VariableTracker.build(tx, id(args[0].fake_value))
+        elif len(args) == 1:
+            arg = args[0]
+            if isinstance(
+                arg,
+                (
+                    variables.UserDefinedClassVariable,
+                    variables.UserDefinedObjectVariable,
+                ),
+            ):
+                if arg.source:
+                    if isinstance(arg, variables.UserDefinedClassVariable):
+                        install_guard(arg.source.make_guard(GuardBuilder.CLASS_MATCH))
+                    else:
+                        install_guard(arg.source.make_guard(GuardBuilder.ID_MATCH))
+            real_val = arg.get_real_python_backed_value()
+            if real_val is not NO_SUCH_SUBOBJ:
+                return VariableTracker.build(tx, id(real_val))
+            return FakeIdVariable(id(arg))
         else:
             unimplemented(
                 gb_type="id() with unsupported args",
@@ -3089,9 +2935,9 @@ def _comparison_with_tensor(
                 == id(extract_fake_example_value(right.as_proxy().node))
             )
             if op is operator.is_:
-                return ConstantVariable.create(is_result)
+                return VariableTracker.build(tx, is_result)
             else:
-                return ConstantVariable.create(not is_result)
+                return VariableTracker.build(tx, not is_result)
 
         if op not in supported_tensor_comparison_op_values:
             unimplemented(
@@ -3144,9 +2990,16 @@ def _comparison_with_symnode(
                 hints=[*graph_break_hints.SUPPORTABLE],
             )
 
-        # This is seen in inspect signature where we check if the value is a default value
-        if isinstance(right, variables.UserDefinedClassVariable):
-            return variables.ConstantVariable(op(object(), None))
+        # SymNodes are numeric (int/float/bool). The non-SymNode operand
+        # must be a type that can participate in a traced numeric comparison.
+        # Anything else (classes, DataPtrVariable, etc.) is a different type
+        # entirely — the comparison result is known at compile time.
+        non_symnode = right if isinstance(left, SymNodeVariable) else left
+        if not isinstance(
+            non_symnode, (SymNodeVariable, ConstantVariable, TensorVariable)
+        ):
+            # pyrefly: ignore [bad-argument-type]
+            return VariableTracker.build(tx, op(object(), None))
 
         proxy = tx.output.create_proxy(
             "call_function", op, (left.as_proxy(), right.as_proxy()), {}
@@ -3172,31 +3025,28 @@ def call_xor(
                 sym_num=None,
             )
 
-        if isinstance(
-            a,
-            (DictKeysVariable, SetVariable, UserDefinedObjectVariable),
-        ):
+        if isinstance(a, _SET_LIKE_OP_SUPPORT):
             return a.call_method(tx, "__xor__", [b], {})
         return None
 
     def call_ixor(
         self, tx: "InstructionTranslator", a: VariableTracker, b: VariableTracker
     ) -> VariableTracker | None:
-        if isinstance(a, (DictKeysVariable, SetVariable, UserDefinedObjectVariable)):
+        if isinstance(a, _SET_LIKE_OP_SUPPORT):
             return a.call_method(tx, "__ixor__", [b], {})
         return None
 
     def call_sub(
         self, tx: "InstructionTranslator", a: VariableTracker, b: VariableTracker
     ) -> VariableTracker | None:
-        if isinstance(a, (DictKeysVariable, SetVariable, UserDefinedObjectVariable)):
+        if isinstance(a, _SET_LIKE_OP_SUPPORT):
             return a.call_method(tx, "__sub__", [b], {})
         return None
 
     def call_isub(
         self, tx: "InstructionTranslator", a: VariableTracker, b: VariableTracker
     ) -> VariableTracker | None:
-        if isinstance(a, (DictKeysVariable, SetVariable, UserDefinedObjectVariable)):
+        if isinstance(a, _SET_LIKE_OP_SUPPORT):
             return a.call_method(tx, "__isub__", [b], {})
         return None
 
@@ -3214,7 +3064,7 @@ def call_and_(
                 ),
                 sym_num=None,
             )
-        if isinstance(a, (DictKeysVariable, SetVariable, UserDefinedObjectVariable)):
+        if isinstance(a, _SET_LIKE_OP_SUPPORT):
             return a.call_method(tx, "__and__", [b], {})
         # None no-ops this handler and lets the driving function proceed
         return None
@@ -3233,7 +3083,7 @@ def call_iand(
                 ),
                 sym_num=None,
             )
-        if isinstance(a, (DictKeysVariable, SetVariable, UserDefinedObjectVariable)):
+        if isinstance(a, _SET_LIKE_OP_SUPPORT):
             return a.call_method(tx, "__iand__", [b], {})
         return None
 
@@ -3243,6 +3093,19 @@ def call_or_(
         # Rely on constant_handler
         if isinstance(a, ConstantVariable) and isinstance(b, ConstantVariable):
             return None
+
+        # Constant fold or_ for class/type variables (e.g. Shard | _StridedShard
+        # producing a types.UnionType for isinstance checks). This handles cases
+        # like OpaqueObjectClassVariable where is_python_constant() returns False
+        # but as_python_constant() works.
+        try:
+            a_const = a.as_python_constant()
+            b_const = b.as_python_constant()
+            if isinstance(a_const, type) and isinstance(b_const, type):
+                return VariableTracker.build(tx, a_const | b_const)
+        except NotImplementedError:
+            pass
+
         if a.is_symnode_like() and b.is_symnode_like():
             return SymNodeVariable.create(
                 tx,
@@ -3256,12 +3119,10 @@ def call_or_(
         if isinstance(
             a,
             (
+                *_SET_LIKE_OP_SUPPORT,
                 ConstDictVariable,
-                DictKeysVariable,
                 MutableMappingVariable,
-                SetVariable,
                 UserDefinedDictVariable,
-                UserDefinedObjectVariable,
             ),
         ):
             # TODO(guilhermeleobas): forward the call to b.__ror__(a) if
@@ -3290,11 +3151,9 @@ def call_ior(
         if isinstance(
             a,
             (
+                *_SET_LIKE_OP_SUPPORT,
                 ConstDictVariable,
-                DictKeysVariable,
                 MutableMappingVariable,
-                SetVariable,
-                UserDefinedObjectVariable,
             ),
         ):
             return a.call_method(tx, "__ior__", [b], {})
@@ -3317,8 +3176,11 @@ def call_not_(
         # Unwrap the underlying ConstDictVariable
         if isinstance(a, DictViewVariable):
             a = a.dv_dict
-        if isinstance(a, (ListVariable, ConstDictVariable)):
-            return ConstantVariable.create(len(a.items) == 0)
+        if isinstance(a, (ListVariable, ConstDictVariable, SetVariable)):
+            return VariableTracker.build(tx, len(a.items) == 0)
+        if isinstance(a, UserDefinedObjectVariable):
+            bool_result = self.call_bool(tx, a)
+            return VariableTracker.build(tx, not bool_result.value)  # type: ignore[missing-attribute]
 
         return None
 
@@ -3327,16 +3189,309 @@ def call_contains(
     ) -> VariableTracker:
         return a.call_method(tx, "__contains__", [b], {})
 
-    def is_python_hashable(self) -> Literal[True]:
-        return True
-
-    def get_python_hash(self) -> int:
-        return hash(self.fn)
-
     def is_python_equal(self, other: object) -> bool:
         return isinstance(other, variables.BuiltinVariable) and self.fn is other.fn
 
 
+class DictBuiltinVariable(BaseBuiltinVariable):
+    """Variable tracker for the `dict` builtin constructor."""
+
+    _fn = dict
+
+    def __init__(self, value: type = dict, **kwargs: Any) -> None:
+        assert value is dict
+        super().__init__(**kwargs)
+
+    def __repr__(self) -> str:
+        return "DictBuiltinVariable()"
+
+    def call_function(
+        self,
+        tx: "InstructionTranslator",
+        args: Sequence[VariableTracker],
+        kwargs: dict[str, VariableTracker],
+    ) -> VariableTracker:
+        return DictBuiltinVariable.call_custom_dict(tx, dict, *args, **kwargs)
+
+    def call_method(
+        self,
+        tx: "InstructionTranslator",
+        name: str,
+        args: list[VariableTracker],
+        kwargs: dict[str, VariableTracker],
+    ) -> VariableTracker:
+        if name == "__new__":
+            if len(args) == 1 and not kwargs:
+                dict_vt = ConstDictVariable({}, dict, mutation_type=ValueMutationNew())
+                if isinstance(args[0], DictBuiltinVariable):
+                    return dict_vt
+                return tx.output.side_effects.track_new_user_defined_object(
+                    self, args[0], args[1:]
+                )
+
+        if name == "fromkeys":
+            return DictBuiltinVariable.call_custom_dict_fromkeys(
+                tx, dict, *args, **kwargs
+            )
+
+        resolved_fn = getattr(dict, name, None)
+        if resolved_fn is not None and resolved_fn in dict_methods:
+            if isinstance(args[0], variables.UserDefinedDictVariable):
+                assert args[0]._base_vt is not None
+                return args[0]._base_vt.call_method(tx, name, args[1:], kwargs)
+            elif isinstance(args[0], ConstDictVariable):
+                return args[0].call_method(tx, name, args[1:], kwargs)
+
+        return super().call_method(tx, name, args, kwargs)
+
+    @staticmethod
+    def call_custom_dict(
+        tx: "InstructionTranslator",
+        user_cls: type,
+        /,
+        *args: VariableTracker,
+        **kwargs: VariableTracker,
+    ) -> VariableTracker:
+        args_list = list(args)
+        return tx.inline_user_function_return(
+            VariableTracker.build(tx, polyfills.construct_dict),
+            [VariableTracker.build(tx, user_cls), *args_list],
+            kwargs,
+        )
+
+    @staticmethod
+    def call_custom_dict_fromkeys(
+        tx: "InstructionTranslator",
+        user_cls: type,
+        /,
+        *args: VariableTracker,
+        **kwargs: VariableTracker,
+    ) -> VariableTracker:
+        if user_cls not in {dict, OrderedDict, defaultdict}:
+            unimplemented(
+                gb_type="Unsupported dict type for fromkeys()",
+                context=f"{user_cls.__name__}.fromkeys(): {args} {kwargs}",
+                explanation=f"Failed to call {user_cls.__name__}.fromkeys() because "
+                f"{user_cls.__name__} is not any type of dict, OrderedDict, or defaultdict",
+                hints=[
+                    f"Ensure {user_cls.__name__} is a type of dict, OrderedDict, or defaultdict.",
+                ],
+            )
+        if kwargs:
+            # Only `OrderedDict.fromkeys` accepts `value` passed by keyword
+            if (
+                user_cls is not OrderedDict
+                or len(args) != 1
+                or len(kwargs) != 1
+                or "value" not in kwargs
+            ):
+                raise_args_mismatch(
+                    tx,
+                    f"{user_cls.__name__}.fromkeys",
+                    "1 args and 1 kwargs (`value`)",
+                    f"{len(args)} args and {len(kwargs)} kwargs",
+                )
+            args = (*args, kwargs.pop("value"))
+        if len(args) == 0:
+            raise_args_mismatch(
+                tx,
+                f"{user_cls.__name__}.fromkeys",
+                "at least 1 args",
+                f"{len(args)} args",
+            )
+        if len(args) == 1:
+            args = (*args, ConstantVariable.create(None))
+        if len(args) != 2:
+            raise_args_mismatch(
+                tx,
+                f"{user_cls.__name__}.fromkeys",
+                "2 args",
+                f"{len(args)} args",
+            )
+
+        arg, value = args
+        DictVariableType = (
+            ConstDictVariable if user_cls is not defaultdict else DefaultDictVariable
+        )
+
+        if isinstance(arg, dict):
+            arg_list = [VariableTracker.build(tx, k) for k in arg]
+            return DictVariableType(
+                dict.fromkeys(arg_list, value),
+                user_cls,
+                mutation_type=ValueMutationNew(),
+            )
+        elif arg.has_force_unpack_var_sequence(tx):
+            keys = arg.force_unpack_var_sequence(tx)
+            if all(is_hashable(v) for v in keys):
+                return DictVariableType(
+                    dict.fromkeys(keys, value),
+                    user_cls,
+                    mutation_type=ValueMutationNew(),
+                )
+
+        unimplemented(
+            gb_type="failed to call dict.fromkeys()",
+            context=f"{user_cls.__name__}.fromkeys(): {args} {kwargs}",
+            explanation=f"Failed to call {user_cls.__name__}.fromkeys() because "
+            "arguments could not be automatically converted to a list, "
+            "or some dict key is not hashable.",
+            hints=[
+                "Manually convert the argument to a list.",
+                "Ensure all keys are hashable.",
+            ],
+        )
+
+
+class IterBuiltinVariable(BaseBuiltinVariable):
+    """Variable tracker for the `iter` builtin."""
+
+    _fn = iter
+
+    def __init__(self, value: Any = iter, **kwargs: Any) -> None:
+        assert value is iter
+        super().__init__(**kwargs)
+
+    def __repr__(self) -> str:
+        return "IterBuiltinVariable()"
+
+    def call_function(
+        self,
+        tx: "InstructionTranslator",
+        args: Sequence[VariableTracker],
+        kwargs: dict[str, VariableTracker],
+    ) -> VariableTracker:
+        if not args:
+            unimplemented(
+                gb_type="iter() with no arguments",
+                context="iter()",
+                explanation="iter() requires at least one argument",
+                hints=[*graph_break_hints.USER_ERROR],
+            )
+        obj, *rest = args
+
+        # Fast path: for known iterable VT types, call __iter__ directly
+        # instead of going through the polyfill, saving tracing overhead.
+        if (
+            not rest
+            and not kwargs
+            and isinstance(
+                obj,
+                (
+                    variables.ListVariable,
+                    variables.RangeVariable,
+                    variables.IteratorVariable,
+                    variables.ConstDictVariable,
+                    variables.NNModuleVariable,
+                    variables.TensorVariable,
+                    variables.TupleVariable,
+                    variables.UserDefinedClassVariable,
+                    DictViewVariable,
+                ),
+            )
+        ):
+            return obj.call_method(tx, "__iter__", [], {})
+
+        # General case: inline the polyfill which handles __iter__ and __getitem__
+        ret = variables.UserFunctionVariable(
+            polyfills.builtins.iter_  # type: ignore[arg-type]
+        ).call_function(tx, [obj, *rest], {})
+
+        if rest:
+            # iter(obj, sentinel) returns a callable iterator; wrap it so
+            # Dynamo knows to forward __next__ calls to the returned object.
+            ret = variables.ObjectIteratorVariable(ret)
+        return ret
+
+
+class ListBuiltinVariable(BaseBuiltinVariable):
+    """Variable tracker for the `list` builtin constructor."""
+
+    _fn = list
+
+    def __init__(self, value: type = list, **kwargs: Any) -> None:
+        assert value is list
+        super().__init__(**kwargs)
+
+    def __repr__(self) -> str:
+        return "ListBuiltinVariable()"
+
+    def call_function(
+        self,
+        tx: "InstructionTranslator",
+        args: Sequence[VariableTracker],
+        kwargs: dict[str, VariableTracker],
+    ) -> VariableTracker:
+        from .user_defined import UserDefinedObjectVariable
+
+        obj = args[0] if args else None
+
+        if isinstance(
+            obj, (variables.IteratorVariable, variables.LocalGeneratorObjectVariable)
+        ) or (
+            isinstance(obj, UserDefinedObjectVariable)
+            and obj.has_force_unpack_var_sequence(tx)
+        ):
+            return ListVariable(
+                list(obj.force_unpack_var_sequence(tx)),
+                mutation_type=ValueMutationNew(),
+            )
+
+        if obj is None:
+            return ListVariable([], mutation_type=ValueMutationNew())
+
+        if obj.has_unpack_var_sequence(tx):
+            if obj.source and not is_constant_source(obj.source):
+                if isinstance(obj, TupleIteratorVariable):
+                    install_guard(
+                        obj.source.make_guard(GuardBuilder.TUPLE_ITERATOR_LEN)
+                    )
+                else:
+                    if isinstance(obj, ConstDictVariable) and not istype(
+                        obj, (SetVariable, FrozensetVariable)
+                    ):
+                        tx.output.guard_on_key_order.add(obj.source)
+                    if isinstance(obj, variables.MappingProxyVariable):
+                        install_guard(
+                            obj.source.make_guard(GuardBuilder.MAPPING_KEYS_CHECK)
+                        )
+                    elif not isinstance(obj, variables.UnspecializedNNModuleVariable):
+                        install_guard(
+                            obj.source.make_guard(GuardBuilder.SEQUENCE_LENGTH)
+                        )
+            return ListVariable(
+                list(obj.unpack_var_sequence(tx)),
+                mutation_type=ValueMutationNew(),
+            )
+
+        arg_types = [type(a).__name__ for a in args]
+        unimplemented(
+            gb_type="Failed to trace list()",
+            context=f"list({arg_types})",
+            explanation=f"Dynamo does not know how to construct a list from argument types {arg_types}",
+            hints=[*graph_break_hints.SUPPORTABLE],
+        )
+
+    def call_method(
+        self,
+        tx: "InstructionTranslator",
+        name: str,
+        args: list[VariableTracker],
+        kwargs: dict[str, VariableTracker],
+    ) -> VariableTracker:
+        if name == "__new__":
+            if len(args) == 1 and not kwargs:
+                list_vt = ListVariable([], mutation_type=ValueMutationNew())
+                if isinstance(args[0], ListBuiltinVariable):
+                    return list_vt
+                return tx.output.side_effects.track_new_user_defined_object(
+                    self, args[0], args[1:]
+                )
+
+        return super().call_method(tx, name, args, kwargs)
+
+
+# pyrefly: ignore [deprecated]
 @contextlib.contextmanager
 def dynamo_disable_grad(tx: "InstructionTranslator") -> typing.Iterator[None]:
     from . import GradModeVariable
diff --git a/torch/_dynamo/variables/constant.py b/torch/_dynamo/variables/constant.py
index 68492e49f57c0..c8061f4e7336a 100644
--- a/torch/_dynamo/variables/constant.py
+++ b/torch/_dynamo/variables/constant.py
@@ -1,34 +1,29 @@
 """
-Constant and enum variable tracking in Dynamo.
+Constant variable tracking in Dynamo.
 
 This module is fundamental to Dynamo's ability to track and propagate constant
 values during compilation, ensuring proper handling of Python literals and
 maintaining type safety through the compilation process.
 """
 
-import enum
+from __future__ import annotations
+
 import operator
-from collections.abc import Sequence
-from typing import Any, Literal, Optional, overload, TYPE_CHECKING, Union
-from typing_extensions import Never, override
+from typing import Any, Literal, overload, TYPE_CHECKING
+from typing_extensions import override
 
 import torch
-from torch._dynamo.source import AttrSource, GetItemSource
+from torch._dynamo.source import GetItemSource
 
-from .. import graph_break_hints, variables
+from .. import variables
 from ..exc import raise_observed_exception, unimplemented
-from ..utils import (
-    cmp_name_to_op_mapping,
-    common_constant_types,
-    istype,
-    np,
-    raise_args_mismatch,
-    raise_on_overridden_hash,
-)
+from ..utils import common_constant_types, istype, np, raise_args_mismatch
 from .base import ValueMutationNew, VariableTracker
 
 
 if TYPE_CHECKING:
+    from collections.abc import Sequence
+
     from torch._dynamo.symbolic_convert import InstructionTranslator
 
     from .functions import UserFunctionVariable
@@ -43,15 +38,23 @@ class ConstantVariable(VariableTracker):
     nested collections.
     """
 
+    # PyLong_Type: https://github.com/python/cpython/blob/v3.13.0/Objects/longobject.c#L6585
+    # PyFloat_Type: https://github.com/python/cpython/blob/v3.13.0/Objects/floatobject.c#L1880
+    # PyBool_Type: https://github.com/python/cpython/blob/v3.13.0/Objects/boolobject.c#L171
+    # PyUnicode_Type: https://github.com/python/cpython/blob/v3.13.0/Objects/unicodeobject.c#L14931
+    # PyBytes_Type: https://github.com/python/cpython/blob/v3.13.0/Objects/bytesobject.c#L3017
+    # PyComplex_Type: https://github.com/python/cpython/blob/v3.13.0/Objects/complexobject.c#L1099
+    # _PyNone_Type: https://github.com/python/cpython/blob/v3.13.0/Objects/object.c#L2022
+    _cpython_type = (int, float, str, bytes, bool, type(None), complex, type(...))
+
     @overload
     @staticmethod
-    def create(value: None) -> Never: ...
+    def create(value: None) -> ConstantVariable: ...
 
     @overload
     @staticmethod
-    def create(value: bool) -> "ConstantVariable": ...
+    def create(value: bool) -> ConstantVariable: ...
 
-    # TODO: Refactor to make these return ConstantVariable
     @overload
     @staticmethod
     def create(value: Any, **kwargs: Any) -> VariableTracker: ...
@@ -66,6 +69,17 @@ def create(value: Any, **kwargs: Any) -> VariableTracker:
         NOTE: the caller must install the proper guards if needed; most often
         the guard will be `CONSTANT_MATCH`.
         """
+        # Return pre-allocated sentinels for None/True/False when there are
+        # no extra kwargs (source, etc.) that would differentiate the instance.
+        if not kwargs:
+            match value:
+                case None:
+                    return CONSTANT_VARIABLE_NONE
+                case True:
+                    return CONSTANT_VARIABLE_TRUE
+                case False:
+                    return CONSTANT_VARIABLE_FALSE
+
         source = kwargs.get("source")
 
         # Routing for supported collection literals.
@@ -139,7 +153,7 @@ def items(self) -> list[VariableTracker]:
         return self.unpack_var_sequence(tx=None)
 
     def getitem_const(
-        self, tx: "InstructionTranslator", arg: VariableTracker
+        self, tx: InstructionTranslator, arg: VariableTracker
     ) -> VariableTracker:
         return ConstantVariable.create(
             self.value[arg.as_python_constant()],
@@ -162,14 +176,29 @@ def is_literal(obj: object, cache: dict[int, object] | None = None) -> bool:
         return ConstantVariable.is_base_literal(obj)
 
     def unpack_var_sequence(
-        self, tx: Optional["InstructionTranslator"]
+        self, tx: InstructionTranslator | None
     ) -> list[VariableTracker]:
         try:
             return [ConstantVariable.create(x) for x in self.as_python_constant()]
         except TypeError as e:
             raise NotImplementedError from e
 
-    def const_getattr(self, tx: "InstructionTranslator", name: str) -> VariableTracker:
+    def len_impl(self, tx: InstructionTranslator) -> VariableTracker:
+        """Generic len for any constant value (sequence or mapping)."""
+        try:
+            return ConstantVariable.create(len(self.value))
+        except TypeError as e:
+            raise_observed_exception(type(e), tx, args=list(e.args))
+
+    def sq_length(self, tx: InstructionTranslator) -> VariableTracker:
+        """Sequence length - delegates to len_impl for constants."""
+        return self.len_impl(tx)
+
+    def mp_length(self, tx: InstructionTranslator) -> VariableTracker:
+        """Mapping length - delegates to len_impl for constants."""
+        return self.len_impl(tx)
+
+    def const_getattr(self, tx: InstructionTranslator, name: str) -> VariableTracker:
         if not hasattr(self.value, name):
             raise_observed_exception(AttributeError, tx, args=[name])
         member = getattr(self.value, name)
@@ -179,7 +208,7 @@ def const_getattr(self, tx: "InstructionTranslator", name: str) -> VariableTrack
 
     def call_method(
         self,
-        tx: "InstructionTranslator",
+        tx: InstructionTranslator,
         name: str,
         args: list[VariableTracker],
         kwargs: dict[str, VariableTracker],
@@ -239,7 +268,7 @@ def call_method(
                     raise_observed_exception(
                         type(exc),
                         tx,
-                        args=list(map(ConstantVariable.create, exc.args)),
+                        args=list(exc.args),
                     )
             if (
                 hasattr(operator, name)
@@ -260,9 +289,7 @@ def call_method(
                     try:
                         return ConstantVariable.create(op(self.value, add_target))
                     except Exception as e:
-                        raise_observed_exception(
-                            type(e), tx, args=list(map(ConstantVariable.create, e.args))
-                        )
+                        raise_observed_exception(type(e), tx, args=list(e.args))
         elif isinstance(self.value, bytes) and name == "decode":
             method = getattr(self.value, name)
             return ConstantVariable.create(method(*const_args, **const_kwargs))
@@ -273,39 +300,27 @@ def call_method(
             except Exception as e:
                 raise_observed_exception(type(e), tx)
 
-        if name == "__len__" and not (args or kwargs):
-            try:
-                # pyrefly: ignore [bad-argument-type]
-                return ConstantVariable.create(len(self.value))
-            except TypeError as e:
-                raise_observed_exception(type(e), tx, args=list(e.args))
-        elif name == "__round__" and len(args) == 1 and args[0].is_python_constant():
+        if name == "__round__" and len(args) == 1 and args[0].is_python_constant():
             try:
                 return ConstantVariable.create(
-                    # pyrefly: ignore [no-matching-overload]
                     round(self.value, args[0].as_python_constant())
                 )
             except Exception as e:
-                raise_observed_exception(
-                    type(e), tx, args=list(map(ConstantVariable.create, e.args))
-                )
+                raise_observed_exception(type(e), tx, args=list(e.args))
         elif name == "__contains__" and len(args) == 1 and args[0].is_python_constant():
             assert not kwargs
             search = args[0].as_python_constant()
             try:
-                # pyrefly: ignore [not-iterable, unsupported-operation]
                 result = search in self.value
                 return ConstantVariable.create(result)
             except TypeError as e:
-                raise_observed_exception(
-                    type(e), tx, args=list(map(ConstantVariable.create, e.args))
-                )
+                raise_observed_exception(type(e), tx, args=list(e.args))
         return super().call_method(tx, name, args, kwargs)
 
     def call_tree_map(
         self,
-        tx: "InstructionTranslator",
-        tree_map_fn: "UserFunctionVariable",
+        tx: InstructionTranslator,
+        tree_map_fn: UserFunctionVariable,
         map_fn: VariableTracker,
         rest: Sequence[VariableTracker],
         tree_map_kwargs: dict[str, VariableTracker],
@@ -356,8 +371,8 @@ def call_tree_map(
 
     @override
     def call_obj_hasattr(
-        self, tx: "InstructionTranslator", name: str
-    ) -> "ConstantVariable":
+        self, tx: InstructionTranslator, name: str
+    ) -> ConstantVariable:
         result = hasattr(self.value, name)
         return variables.ConstantVariable.create(result)
 
@@ -368,7 +383,6 @@ def get_python_hash(self) -> int:
         return hash(self.value)
 
     def is_python_equal(self, other: object) -> bool:
-        # Could be an EnumVariable as well
         from .tensor import SymNodeVariable
 
         if isinstance(other, SymNodeVariable):
@@ -378,67 +392,91 @@ def is_python_equal(self, other: object) -> bool:
             and self.as_python_constant() == other.as_python_constant()
         )
 
+    def get_real_python_backed_value(self) -> object:
+        return self.value
+
+    def nb_index_impl(
+        self,
+        tx: Any,
+    ) -> VariableTracker:
+        # CPython: int and bool define nb_index (returns self for int,
+        # int(self) for bool). All other constant types do not.
+        if isinstance(self.value, (int, bool)):
+            return ConstantVariable.create(operator.index(self.value))
+        return super().nb_index_impl(tx)
+
+    def nb_int_impl(
+        self,
+        tx: Any,
+    ) -> VariableTracker:
+        # CPython: int defines nb_int (long_long, returns copy).
+        # bool inherits nb_int from int via slot inheritance.
+        # float defines nb_int (truncates toward zero via PyLong_FromDouble).
+        return ConstantVariable.create(int(self.value))
+
+    def nb_float_impl(
+        self,
+        tx: Any,
+    ) -> VariableTracker:
+        # CPython: float defines nb_float (float_float, returns copy).
+        # int defines nb_float (long_float, converts to float).
+        # bool inherits nb_float from int via slot inheritance.
+        return ConstantVariable.create(float(self.value))
+
 
 CONSTANT_VARIABLE_NONE = ConstantVariable(None)
+CONSTANT_VARIABLE_TRUE = ConstantVariable(True)
+CONSTANT_VARIABLE_FALSE = ConstantVariable(False)
 
 
-class EnumVariable(VariableTracker):
-    """VariableTracker for enum.Enum and enum.IntEnum instances
+class FakeIdVariable(VariableTracker):
+    """A compile-time-only id value that can be used as a dict key but cannot
+    be reconstructed across graph breaks.
 
-    Provides specialized handling for Python enum types, supporting
-    both standard Enum and IntEnum with proper value tracking and comparison.
+    When dynamo evaluates ``id(x)`` on a variable tracker that has no
+    corresponding runtime object (e.g. a ``ConstDictVariable`` created during
+    tracing), we mint a fake integer id.  This variable holds that id and
+    supports the minimal interface needed to participate as a dict key
+    (hashing and equality).  It intentionally blocks reconstruction so that a
+    graph break does not silently bake a stale id into the resumed bytecode.
     """
 
-    def __init__(self, value: Union[enum.Enum, enum.IntEnum], **kwargs: Any) -> None:
+    # PyLong_Type: https://github.com/python/cpython/blob/v3.13.0/Objects/longobject.c#L6585
+    _cpython_type = int
+
+    def __init__(self, value: int, **kwargs: Any) -> None:
         super().__init__(**kwargs)
         self.value = value
 
-    @classmethod
-    def create(
-        cls, cls_type: Any, value_vt: VariableTracker, options: Any
-    ) -> "EnumVariable":
-        if value_vt.is_python_constant():
-            for member in list(cls_type):
-                if member.value == value_vt.as_python_constant():
-                    return cls(member, **options)
-        unimplemented(
-            gb_type="Failed to construct Enum variable",
-            context=f"value: {value_vt}, allowed enum values: {list(cls_type)}",
-            explanation="Attempted to construct an Enum value that is non-constant (e.g. int, string) "
-            "or is not an acceptable value for the Enum. "
-            f"Acceptable values for Enum `{cls_type}`: {list(cls_type)}.",
-            hints=[*graph_break_hints.USER_ERROR, *graph_break_hints.SUPPORTABLE],
-        )
-
-    def as_proxy(self) -> Union[enum.Enum, int]:
-        if isinstance(self.value, int):
-            return int(self.value)  # convert IntEnum to a normal int
+    def as_python_constant(self) -> int:
         return self.value
 
-    def __repr__(self) -> str:
-        return f"EnumVariable({type(self.value)})"
-
-    def as_python_constant(self) -> Union[enum.Enum, enum.IntEnum]:
-        return self.value
+    def is_python_constant(self) -> bool:
+        return False
 
-    def var_getattr(self, tx: "InstructionTranslator", name: str) -> VariableTracker:
-        if not hasattr(self.value, name):
-            raise NotImplementedError
-        if name in cmp_name_to_op_mapping:
-            return variables.GetAttrVariable(self, name)
-        member = getattr(self.value, name)
-        source = self.source and AttrSource(self.source, name)
-        return VariableTracker.build(tx, member, source=source)
+    def python_type(self) -> type:
+        return int
 
-    def is_python_hashable(self) -> Literal[True]:
-        raise_on_overridden_hash(self.value, self)
+    def is_python_hashable(self) -> bool:
         return True
 
     def get_python_hash(self) -> int:
-        return hash(self.as_python_constant())
+        return hash(self.value)
 
     def is_python_equal(self, other: object) -> bool:
-        return (
-            isinstance(other, VariableTracker)
-            and self.as_python_constant() == other.as_python_constant()
+        if isinstance(other, (FakeIdVariable, ConstantVariable)):
+            return self.value == other.as_python_constant()
+        return False
+
+    def reconstruct(self, codegen: Any) -> None:
+        unimplemented(
+            gb_type="Reconstruction of FakeIdVariable",
+            context=str(self.value),
+            explanation=(
+                "A fake id produced by id() on a compile-time container "
+                "cannot be reconstructed across a graph break."
+            ),
+            hints=[
+                "Avoid using id() on containers in code that may graph-break.",
+            ],
         )
diff --git a/torch/_dynamo/variables/ctx_manager.py b/torch/_dynamo/variables/ctx_manager.py
index 7ced6bd665b73..30ebc0ca8735b 100644
--- a/torch/_dynamo/variables/ctx_manager.py
+++ b/torch/_dynamo/variables/ctx_manager.py
@@ -18,12 +18,14 @@
 restoring state changes.
 """
 
+import contextlib
 import inspect
 import logging
+import types
 import warnings
 from collections.abc import Callable, Sequence, Sized
 from contextlib import AbstractContextManager, ExitStack
-from typing import Any, Optional, TYPE_CHECKING, Union
+from typing import Any, TYPE_CHECKING
 
 import torch._C
 from torch._dynamo import config
@@ -67,7 +69,7 @@ class ContextWrappingVariable(VariableTracker):
     }
 
     def __init__(
-        self, target_values: Any, initial_values: Optional[Any] = None, **kwargs: Any
+        self, target_values: Any, initial_values: Any | None = None, **kwargs: Any
     ) -> None:
         super().__init__(**kwargs)
         self.target_values = target_values
@@ -77,10 +79,10 @@ def enter(self, tx: "InstructionTranslator") -> VariableTracker:
         if hasattr(self, "_call_func"):
             self._call_func(tx, self.target_values)
         self.set_cleanup_hook(tx)
-        return variables.CONSTANT_VARIABLE_NONE
+        return variables.ConstantVariable.create(None)
 
     def set_cleanup_hook(
-        self, tx: "InstructionTranslator", fn: Optional[Callable[..., Any]] = None
+        self, tx: "InstructionTranslator", fn: Callable[..., Any] | None = None
     ) -> None:
         if fn is None:
 
@@ -88,14 +90,14 @@ def fn() -> None:
                 if hasattr(self, "_call_func"):
                     self._call_func(tx, self.initial_values)
 
-        self.cleanup_fn: Optional[Callable[..., Any]] = fn
+        self.cleanup_fn: Callable[..., Any] | None = fn
         tx.output.add_cleanup_hook(self.cleanup)
 
     def exit(
         self, tx: "InstructionTranslator", *args: VariableTracker
     ) -> VariableTracker:
         self.cleanup_assert()
-        return variables.CONSTANT_VARIABLE_NONE
+        return variables.ConstantVariable.create(None)
 
     def reconstruct_type(self, codegen: "PyCodegen") -> None:
         codegen(
@@ -271,7 +273,7 @@ def enter(self, tx: "InstructionTranslator") -> VariableTracker:
             (enabled,),
             {},
         )
-        return variables.CONSTANT_VARIABLE_NONE
+        return variables.ConstantVariable.create(None)
 
     def exit(
         self, tx: "InstructionTranslator", *args: VariableTracker
@@ -283,7 +285,7 @@ def exit(
             (self.prev_state,),
             {},
         )
-        return variables.CONSTANT_VARIABLE_NONE
+        return variables.ConstantVariable.create(None)
 
 
 class TemporarilyPopInterpreterStackCtxManagerVariable(ContextWrappingVariable):
@@ -311,7 +313,7 @@ def enter(self, tx: "InstructionTranslator") -> VariableTracker:
             (),
             {},
         )
-        return variables.CONSTANT_VARIABLE_NONE
+        return variables.ConstantVariable.create(None)
 
     def exit(
         self, tx: "InstructionTranslator", *args: VariableTracker
@@ -323,7 +325,7 @@ def exit(
             (self.proxy,),
             {},
         )
-        return variables.CONSTANT_VARIABLE_NONE
+        return variables.ConstantVariable.create(None)
 
 
 class JvpIncrementNestingCtxManagerVariable(ContextWrappingVariable):
@@ -359,7 +361,7 @@ def enter(self, tx: "InstructionTranslator") -> VariableTracker:
             (),
             {},
         )
-        return variables.ConstantVariable.create(jvp_level)
+        return variables.VariableTracker.build(tx, jvp_level)
 
     def exit(
         self, tx: "InstructionTranslator", *args: VariableTracker
@@ -368,7 +370,7 @@ def exit(
         tx.output.create_node(
             "call_function", torch._C._functorch._jvp_decrement_nesting, (), {}
         )
-        return variables.CONSTANT_VARIABLE_NONE
+        return variables.ConstantVariable.create(None)
 
 
 class SetFwdGradEnabledContextManager(ContextWrappingVariable):
@@ -398,7 +400,7 @@ def enter(self, tx: "InstructionTranslator") -> VariableTracker:
             (mode,),
             {},
         )
-        return variables.CONSTANT_VARIABLE_NONE
+        return variables.ConstantVariable.create(None)
 
     def exit(
         self, tx: "InstructionTranslator", *args: VariableTracker
@@ -410,7 +412,10 @@ def exit(
             (self.prev_state,),
             {},
         )
-        return variables.CONSTANT_VARIABLE_NONE
+        return variables.ConstantVariable.create(None)
+
+    def python_type(self) -> type:
+        return torch.autograd.forward_ad._set_fwd_grad_enabled
 
 
 class DualLevelContextManager(ContextWrappingVariable):
@@ -438,7 +443,7 @@ def enter(self, tx: "InstructionTranslator") -> VariableTracker:
             (),
             {},
         )
-        return variables.ConstantVariable.create(self.new_level)
+        return variables.VariableTracker.build(tx, self.new_level)
 
     def exit(
         self, tx: "InstructionTranslator", *args: VariableTracker
@@ -450,7 +455,10 @@ def exit(
             (self.new_level,),
             {},
         )
-        return variables.CONSTANT_VARIABLE_NONE
+        return variables.ConstantVariable.create(None)
+
+    def python_type(self) -> type:
+        return torch.autograd.forward_ad.dual_level
 
 
 class GradIncrementNestingCtxManagerVariable(ContextWrappingVariable):
@@ -484,7 +492,7 @@ def enter(self, tx: "InstructionTranslator") -> VariableTracker:
             (),
             {},
         )
-        return variables.ConstantVariable.create(grad_level)
+        return variables.VariableTracker.build(tx, grad_level)
 
     def exit(
         self, tx: "InstructionTranslator", *args: VariableTracker
@@ -493,7 +501,7 @@ def exit(
         tx.output.create_node(
             "call_function", torch._C._functorch._grad_decrement_nesting, (), {}
         )
-        return variables.CONSTANT_VARIABLE_NONE
+        return variables.ConstantVariable.create(None)
 
 
 class CatchWarningsCtxManagerVariable(ContextWrappingVariable):
@@ -512,8 +520,8 @@ def create(
     def __init__(
         self,
         catch_warnings_args: dict[str, VariableTracker],
-        target_values: Optional[Any] = None,
-        initial_values: Optional[Any] = None,
+        target_values: Any | None = None,
+        initial_values: Any | None = None,
         **kwargs: Any,
     ) -> None:
         assert isinstance(catch_warnings_args, dict), catch_warnings_args
@@ -528,7 +536,7 @@ def enter(self, tx: "InstructionTranslator") -> VariableTracker:
         }
         ctx_val = warnings.catch_warnings(**kwargs)
         self.set_cleanup_hook(tx, lambda: ctx_val.__exit__(None, None, None))
-        return variables.ConstantVariable.create(ctx_val.__enter__())
+        return variables.VariableTracker.build(tx, ctx_val.__enter__())
 
     def reconstruct(self, codegen: "PyCodegen") -> None:
         codegen.add_push_null(
@@ -538,6 +546,9 @@ def reconstruct(self, codegen: "PyCodegen") -> None:
         keys = tuple(self.catch_warnings_args.keys())
         codegen.extend_output(codegen.create_call_function_kw(len(keys), keys, False))
 
+    def python_type(self) -> type:
+        return warnings.catch_warnings
+
 
 class VmapIncrementNestingCtxManagerVariable(ContextWrappingVariable):
     """represents torch VMap increment/decrement nesting"""
@@ -580,7 +591,7 @@ def enter(self, tx: "InstructionTranslator") -> VariableTracker:
             (batch_size.as_proxy(), randomness),
             {},
         )
-        return variables.ConstantVariable.create(vmap_level)
+        return variables.VariableTracker.build(tx, vmap_level)
 
     def exit(
         self, tx: "InstructionTranslator", *args: VariableTracker
@@ -592,7 +603,7 @@ def exit(
             (),
             {},
         )
-        return variables.CONSTANT_VARIABLE_NONE
+        return variables.ConstantVariable.create(None)
 
 
 class GradModeVariable(ContextWrappingVariable):
@@ -619,7 +630,7 @@ def create(
     def __init__(
         self,
         target_values: Any,
-        initial_values: Optional[Sequence[bool]] = None,
+        initial_values: Sequence[bool] | None = None,
         initialized: bool = True,
         **kwargs: Any,
     ) -> None:
@@ -630,13 +641,13 @@ def __init__(
 
     def enter(self, tx: "InstructionTranslator") -> VariableTracker:
         self._call_func(tx, self.target_values)
-        return variables.CONSTANT_VARIABLE_NONE
+        return variables.ConstantVariable.create(None)
 
     def exit(
         self, tx: "InstructionTranslator", *args: VariableTracker
     ) -> VariableTracker:
         self._call_func(tx, self.initial_values)
-        return variables.CONSTANT_VARIABLE_NONE
+        return variables.ConstantVariable.create(None)
 
     def call_function(
         self,
@@ -663,6 +674,9 @@ def module_name(self) -> str:
     def fn_name(self) -> str:
         return "set_grad_enabled"
 
+    def python_type(self) -> type:
+        return torch.set_grad_enabled
+
 
 class InferenceModeVariable(ContextWrappingVariable):
     @staticmethod
@@ -677,7 +691,7 @@ def create(
     def __init__(
         self,
         target_values: Any,
-        initial_values: Optional[bool] = None,
+        initial_values: bool | None = None,
         **kwargs: Any,
     ) -> None:
         if initial_values is None:
@@ -697,7 +711,7 @@ def exit(
             (self.proxy,),
             {},
         )
-        return variables.CONSTANT_VARIABLE_NONE
+        return variables.ConstantVariable.create(None)
 
     def enter(self, tx: "InstructionTranslator") -> VariableTracker:
         disabled_inference_mode_forcibly = False
@@ -727,7 +741,7 @@ def cleanup_hook() -> None:
             (*self.target_values,),
             {},
         )
-        return variables.CONSTANT_VARIABLE_NONE
+        return variables.ConstantVariable.create(None)
 
     def module_name(self) -> str:
         return "torch"
@@ -735,6 +749,9 @@ def module_name(self) -> str:
     def fn_name(self) -> str:
         return "inference_mode"
 
+    def python_type(self) -> type:
+        return torch.inference_mode
+
 
 class CUDADeviceVariable(ContextWrappingVariable):
     """represents torch.cuda.device"""
@@ -753,7 +770,7 @@ def create(
     def __init__(
         self,
         target_values: Any,
-        initial_values: Optional[Any] = None,
+        initial_values: Any | None = None,
         **kwargs: Any,
     ) -> None:
         super().__init__(
@@ -781,7 +798,7 @@ def enter(self, tx: "InstructionTranslator") -> VariableTracker:
             (*self.target_values,),
             {},
         )
-        return variables.CONSTANT_VARIABLE_NONE
+        return variables.ConstantVariable.create(None)
 
     def module_name(self) -> str:
         return "torch.cuda"
@@ -789,6 +806,9 @@ def module_name(self) -> str:
     def fn_name(self) -> str:
         return "device"
 
+    def python_type(self) -> type:
+        return torch.cuda.device
+
 
 class TorchFunctionDisableVariable(ContextWrappingVariable):
     """represents whether torch function overrides are enabled or not"""
@@ -811,7 +831,7 @@ def __init__(
         self,
         tx: "InstructionTranslator",
         target_values: Sized,
-        initial_values: Optional[Sized] = None,
+        initial_values: Sized | None = None,
         only_subclass: bool = True,
         **kwargs: Any,
     ) -> None:
@@ -843,7 +863,7 @@ def fn() -> None:
                 )
                 if not self.only_subclass:
                     tx.symbolic_torch_function_state.torch_function_mode_enabled = (
-                        self.initial_torch_function_subclass_enabled
+                        self.initial_torch_function_mode_enabled
                     )
 
         self.cleanup_fn = fn
@@ -863,13 +883,18 @@ def fn_name(self) -> str:
             return "DisableTorchFunctionSubclass"
         return "DisableTorchFunction"
 
+    def python_type(self) -> type:
+        if self.only_subclass:
+            return torch._C.DisableTorchFunctionSubclass  # pyrefly: ignore[bad-return]
+        return torch._C.DisableTorchFunction  # pyrefly: ignore[bad-return]
+
 
 class DisabledSavedTensorsHooksVariable(ContextWrappingVariable):
     """represents torch.autograd.graph.disable_saved_tensors_hook."""
 
     @staticmethod
     def create(
-        tx: "InstructionTranslator", target_value: Optional[str], **kwargs: Any
+        tx: "InstructionTranslator", target_value: str | None, **kwargs: Any
     ) -> "DisabledSavedTensorsHooksVariable":
         var = DisabledSavedTensorsHooksVariable(
             target_values=[target_value],
@@ -884,8 +909,8 @@ def create(
 
     def __init__(
         self,
-        target_values: Sequence[Optional[str]],
-        initial_values: Optional[Sequence[Optional[str]]] = None,
+        target_values: Sequence[str | None],
+        initial_values: Sequence[str | None] | None = None,
         **kwargs: Any,
     ) -> None:
         super().__init__(
@@ -893,10 +918,10 @@ def __init__(
         )
 
     def enter(self, tx: "InstructionTranslator") -> VariableTracker:
-        return variables.CONSTANT_VARIABLE_NONE
+        return variables.ConstantVariable.create(None)
 
     def _call_func(
-        self, tx: "InstructionTranslator", values: Sequence[Optional[str]]
+        self, tx: "InstructionTranslator", values: Sequence[str | None]
     ) -> None:
         assert len(values) == 1
         value = values[0]
@@ -924,6 +949,9 @@ def module_name(self) -> str:
     def fn_name(self) -> str:
         return "disable_saved_tensors_hooks"
 
+    def python_type(self) -> type:
+        return contextlib._GeneratorContextManager
+
 
 class AutocastModeVariable(ContextWrappingVariable):
     @staticmethod
@@ -966,7 +994,7 @@ def create(
     def __init__(
         self,
         target_values: Sequence[Any],
-        initial_values: Optional[Any] = None,
+        initial_values: Any | None = None,
         **kwargs: Any,
     ) -> None:
         super().__init__(
@@ -980,7 +1008,7 @@ def exit(
         tx.output.create_node(
             "call_function", torch.amp._exit_autocast, (self.proxy,), {}
         )
-        return variables.CONSTANT_VARIABLE_NONE
+        return variables.ConstantVariable.create(None)
 
     def enter(self, tx: "InstructionTranslator") -> VariableTracker:
         ctx = torch.amp._enter_autocast(*self.target_values)
@@ -988,7 +1016,7 @@ def enter(self, tx: "InstructionTranslator") -> VariableTracker:
         self.proxy = tx.output.create_node(
             "call_function", torch.amp._enter_autocast, (*self.target_values,), {}
         )
-        return variables.CONSTANT_VARIABLE_NONE
+        return variables.ConstantVariable.create(None)
 
     def module_name(self) -> str:
         return "torch.amp.autocast_mode"
@@ -996,23 +1024,26 @@ def module_name(self) -> str:
     def fn_name(self) -> str:
         return "autocast"
 
+    def python_type(self) -> type:
+        return torch.amp.autocast_mode.autocast
+
 
 class NullContextVariable(ContextWrappingVariable):
     """
     This class represents Python contextlib.nullcontext.
     """
 
-    def __init__(self, target_values: Optional[Any] = None, **kwargs: Any) -> None:
+    def __init__(self, target_values: Any | None = None, **kwargs: Any) -> None:
         super().__init__(target_values=target_values, **kwargs)
 
     def enter(self, tx: "InstructionTranslator") -> VariableTracker:
-        none = variables.CONSTANT_VARIABLE_NONE
+        none = variables.ConstantVariable.create(None)
         return self.target_values if self.target_values else none
 
     def exit(
         self, tx: "InstructionTranslator", *args: VariableTracker
     ) -> VariableTracker:
-        return variables.CONSTANT_VARIABLE_NONE
+        return variables.ConstantVariable.create(None)
 
     def module_name(self) -> str:
         return "contextlib"
@@ -1020,6 +1051,9 @@ def module_name(self) -> str:
     def fn_name(self) -> str:
         return "nullcontext"
 
+    def python_type(self) -> type:
+        return contextlib.nullcontext
+
 
 class ProfilerContextVariable(ContextWrappingVariable):
     """
@@ -1033,13 +1067,16 @@ class ProfilerContextVariable(ContextWrappingVariable):
     def __init__(self, **kwargs: Any) -> None:
         super().__init__(target_values=None, **kwargs)
 
+    def python_type(self) -> type:
+        return torch.profiler.profile
+
     def enter(self, tx: "InstructionTranslator") -> VariableTracker:
         return self
 
     def exit(
         self, tx: "InstructionTranslator", *args: VariableTracker
     ) -> VariableTracker:
-        return variables.CONSTANT_VARIABLE_NONE
+        return variables.ConstantVariable.create(None)
 
     def module_name(self) -> str:
         return "contextlib"
@@ -1089,7 +1126,7 @@ def create(
                 record_args[0].as_python_constant()
                 if record_args
                 else kwargs.get(
-                    "name", variables.ConstantVariable("unknown")
+                    "name", variables.ConstantVariable.create("unknown")
                 ).as_python_constant()
             )
             record_args_const = None
@@ -1109,13 +1146,16 @@ def create(
     def __init__(
         self,
         target_values: Any = None,
-        initial_values: Optional[Any] = None,
+        initial_values: Any | None = None,
         **kwargs: Any,
     ) -> None:
         super().__init__(
             target_values=target_values, initial_values=initial_values, **kwargs
         )
 
+    def python_type(self) -> type:
+        return torch.autograd.profiler.record_function
+
     def enter(self, tx: "InstructionTranslator") -> VariableTracker:
         if config.capture_profiler_record_function:
             name, args = self.target_values
@@ -1139,7 +1179,7 @@ def exit(
                 (self.proxy,),
                 {},
             )
-        return variables.CONSTANT_VARIABLE_NONE
+        return variables.ConstantVariable.create(None)
 
     def module_name(self) -> str:
         return (
@@ -1220,8 +1260,11 @@ def __init__(
         if self.prev_versions.is_symnode_like():
             self.prev_versions = variables.TupleVariable([self.prev_versions])
 
+    def python_type(self) -> type:
+        return torch.autograd.grad_mode._unsafe_preserve_version_counter
+
     def enter(self, tx: "InstructionTranslator") -> VariableTracker:
-        return variables.CONSTANT_VARIABLE_NONE
+        return variables.ConstantVariable.create(None)
 
     def exit(
         self, tx: "InstructionTranslator", *args: VariableTracker
@@ -1268,7 +1311,7 @@ def __init__(
         self,
         param_group_var: Any,
         target_values: Sequence[Any],
-        initial_values: Optional[Sequence[Any]] = None,
+        initial_values: Sequence[Any] | None = None,
         **kwargs: Any,
     ) -> None:
         super().__init__(
@@ -1279,13 +1322,13 @@ def __init__(
 
     def enter(self, tx: "InstructionTranslator") -> VariableTracker:
         self._call_func(tx, self.target_values)
-        return variables.CONSTANT_VARIABLE_NONE
+        return variables.ConstantVariable.create(None)
 
     def exit(
         self, tx: "InstructionTranslator", *args: VariableTracker
     ) -> VariableTracker:
         self._call_func(tx, self.initial_values)  # type: ignore[arg-type]
-        return variables.CONSTANT_VARIABLE_NONE
+        return variables.ConstantVariable.create(None)
 
     def call_function(
         self,
@@ -1305,8 +1348,8 @@ def _call_func(self, tx: "InstructionTranslator", values: Sequence[Any]) -> None
                 tx,
                 "__setattr__",
                 (
-                    variables.ConstantVariable.create("_training_state"),
-                    variables.EnumVariable(value),
+                    variables.VariableTracker.build(tx, "_training_state"),
+                    variables.VariableTracker.build(tx, value),
                 ),
                 {},
             )
@@ -1388,7 +1431,7 @@ def enter(self, tx: "InstructionTranslator") -> VariableTracker:
             (arg, bool(self.set_priority)),
             {},
         )
-        return variables.CONSTANT_VARIABLE_NONE
+        return variables.ConstantVariable.create(None)
 
     def exit(
         self, tx: "InstructionTranslator", *args: VariableTracker
@@ -1401,7 +1444,7 @@ def exit(
             (arg, bool(self.set_priority)),
             {},
         )
-        return variables.CONSTANT_VARIABLE_NONE
+        return variables.ConstantVariable.create(None)
 
     def module_name(self) -> str:
         return "torch.nn.attention"
@@ -1411,6 +1454,9 @@ def module_name(self) -> str:
     def fn_name(self) -> str:
         return "_sdpa_kernel_variadic"
 
+    def python_type(self) -> type:
+        return contextlib._GeneratorContextManager
+
 
 class FxTracebackAnnotateVariable(ContextWrappingVariable):
     """
@@ -1439,7 +1485,7 @@ def enter(
         stack.enter_context(torch.fx.traceback.annotate(self.target_values))
         stack.enter_context(torch.fx.traceback.preserve_node_meta())
         self.set_cleanup_hook(tx, lambda: stack.close())
-        return variables.CONSTANT_VARIABLE_NONE
+        return variables.ConstantVariable.create(None)
 
     def module_name(self) -> str:
         return "torch.fx.traceback"
@@ -1447,6 +1493,9 @@ def module_name(self) -> str:
     def fn_name(self) -> str:
         return "annotate"
 
+    def python_type(self) -> type:
+        return contextlib._GeneratorContextManager
+
     def reconstruct_type(self, codegen: "PyCodegen") -> None:
         unimplemented(
             gb_type="torch.fx.traceback.annotate escaped from compiled region",
@@ -1491,6 +1540,11 @@ def module_name(self) -> str:
     def fn_name(self) -> str:
         return "patch_dynamo_config"
 
+    def python_type(self) -> type:
+        from torch._dynamo.decorators import DynamoConfigPatchProxy
+
+        return DynamoConfigPatchProxy
+
 
 class ErrorOnGraphBreakVariable(ContextWrappingVariable):
     """represents torch._dynamo.error_on_graph_break"""
@@ -1512,16 +1566,82 @@ def module_name(self) -> str:
     def fn_name(self) -> str:
         return "error_on_graph_break"
 
+    def python_type(self) -> type:
+        from torch._dynamo.decorators import ErrorOnGraphBreakDecoratorContextManager
+
+        return ErrorOnGraphBreakDecoratorContextManager
+
+
+class CudagraphOverrideVariable(ContextWrappingVariable):
+    """represents torch._dynamo.override_cudagraphs"""
+
+    def __init__(self, fwd: bool | None, bwd: bool | None, **kwargs: Any) -> None:
+        super().__init__(
+            target_values=(fwd, bwd),
+            initial_values=None,  # Captured in enter()
+            **kwargs,
+        )
+
+    def enter(self, tx: "InstructionTranslator") -> "VariableTracker":
+        # Capture current annotation before overwriting
+        prev = tx.output.cudagraph_annotation
+        if prev is not None:
+            self.initial_values = (prev.fwd, prev.bwd)
+        else:
+            self.initial_values = (None,)
+        return super().enter(tx)
+
+    def _call_func(self, tx: "InstructionTranslator", values: tuple[Any, ...]) -> None:
+        from torch._inductor import _CudagraphAnnotation
+
+        if len(values) == 1 and values[0] is None:
+            # Restoring to no annotation
+            tx.output.cudagraph_annotation = None
+            tx.output.tracing_context.cudagraph_annotation = None
+        else:
+            assert len(values) == 2
+            fwd, bwd = values
+            annotation = _CudagraphAnnotation(fwd=fwd, bwd=bwd)
+            tx.output.cudagraph_annotation = annotation
+            tx.output.tracing_context.cudagraph_annotation = annotation
+
+    def module_name(self) -> str:
+        return "torch._dynamo"
+
+    def fn_name(self) -> str:
+        return "override_cudagraphs"
+
+    def python_type(self) -> type:
+        from torch._dynamo.decorators import CudagraphOverrideContextManager
+
+        return CudagraphOverrideContextManager
+
+    def exit_on_graph_break(self) -> bool:
+        # Annotation persists until graph is compiled; each resume function
+        # will reconstruct the context manager and call enter() again
+        return False
+
+    def exit(
+        self, tx: "InstructionTranslator", *args: VariableTracker
+    ) -> "VariableTracker":
+        # Override to NOT call cleanup here. The cleanup will happen in
+        # call_cleanup_hooks() during compile_subgraph, which is AFTER
+        # the annotation is copied to gm.meta.
+        return variables.ConstantVariable.create(None)
+
 
 class WithEnterFunctionVariable(VariableTracker):
     def __init__(
         self,
-        ctx: Union[ContextWrappingVariable, GenericContextWrappingVariable],
+        ctx: ContextWrappingVariable | GenericContextWrappingVariable,
         **kwargs: Any,
     ) -> None:
         super().__init__(**kwargs)
         self.ctx = ctx
 
+    def python_type(self) -> type:
+        return types.MethodType
+
     def call_function(
         self,
         tx: "InstructionTranslator",
@@ -1562,7 +1682,7 @@ class WithExitFunctionVariable(VariableTracker):
 
     def __init__(
         self,
-        ctx: Union[ContextWrappingVariable, GenericContextWrappingVariable],
+        ctx: ContextWrappingVariable | GenericContextWrappingVariable,
         target: Any,
         **kwargs: Any,
     ) -> None:
@@ -1573,6 +1693,9 @@ def __init__(
         self.ctx = ctx
         self.target = target
 
+    def python_type(self) -> type:
+        return types.MethodType
+
     def call_function(
         self,
         tx: "InstructionTranslator",
diff --git a/torch/_dynamo/variables/dicts.py b/torch/_dynamo/variables/dicts.py
index f934369aca817..27f674b2e159b 100644
--- a/torch/_dynamo/variables/dicts.py
+++ b/torch/_dynamo/variables/dicts.py
@@ -6,26 +6,24 @@
 - Ordered dictionaries (collections.OrderedDict)
 - Default dictionaries (collections.defaultdict)
 - Dictionary views (keys and values)
-- Sets and frozensets (implemented internally using dictionaries)
 
 These classes are responsible for tracking dictionary operations during graph compilation,
 maintaining proper guards for dictionary mutations and key existence checks. They handle
 dictionary creation, modification, key/value access, and view operations while ensuring
 correct behavior in the compiled code through appropriate guard installation.
 
-The implementation uses a special _HashableTracker wrapper to handle dictionary keys
-while preserving proper aliasing semantics. Sets are implemented as dictionaries with
-None values for efficiency and code reuse.
+The implementation uses a special HashableTracker wrapper to handle
+dictionary keys while preserving proper aliasing semantics. Set-related classes live
+in sets.py.
 """
 
 import collections
 import functools
-import operator
 import types
-from collections.abc import Iterable, Sequence
-from typing import Any, Literal, Optional, TYPE_CHECKING, Union
+from collections.abc import Callable, Iterator, Sequence
+from typing import Any, Literal, TYPE_CHECKING, Union
 
-from torch.utils._ordered_set import OrderedSet
+from torch.utils._pytree import MappingKey
 
 from .. import graph_break_hints, polyfills, variables
 from ..bytecode_transformation import (
@@ -36,7 +34,12 @@
 )
 from ..exc import raise_observed_exception, unimplemented
 from ..guards import GuardBuilder, install_guard
-from ..source import AttrSource, is_constant_source, is_from_local_source
+from ..source import (
+    AttrSource,
+    DictGetItemSource,
+    is_constant_source,
+    is_from_local_source,
+)
 from ..utils import (
     cmp_name_to_op_mapping,
     dict_items,
@@ -44,16 +47,22 @@
     dict_values,
     istype,
     raise_args_mismatch,
-    specialize_symnode,
 )
-from .base import ValueMutationNew, VariableTracker
-from .constant import CONSTANT_VARIABLE_NONE, ConstantVariable
+from .base import (
+    AttributeMutationExisting,
+    AttributeMutationNew,
+    NO_SUCH_SUBOBJ,
+    ValueMutationNew,
+    VariableTracker,
+)
+from .constant import ConstantVariable
+from .hashable import HashableTracker, is_hashable, raise_unhashable
+from .sets import SetVariable
 
 
 if TYPE_CHECKING:
     from torch._dynamo.codegen import PyCodegen
     from torch._dynamo.symbolic_convert import InstructionTranslator
-    from torch._dynamo.variables.builtin import BuiltinVariable
 
     from .functions import UserFunctionVariable
 
@@ -63,112 +72,18 @@
 # - Implement get_python_hash() and is_python_equal() methods for hashable types
 
 
-def was_instancecheck_override(obj: Any) -> bool:
-    return type(obj).__dict__.get("__instancecheck__", False)
-
-
-def raise_unhashable(
-    arg: VariableTracker, tx: Optional["InstructionTranslator"] = None
-) -> None:
-    from .builder import SourcelessBuilder
-
-    if tx is None:
-        from torch._dynamo.symbolic_convert import InstructionTranslator
-
-        tx = InstructionTranslator.current_tx()
-    try:
-        arg_type = arg.python_type()
-    except Exception:
-        arg_type = type(arg)
-
-    raise_observed_exception(
-        TypeError,
-        tx,
-        args=[
-            SourcelessBuilder.create(
-                tx,
-                f"unhashable type: {arg_type!r} and variable tracker = {type(arg.realize())}",
-            )
-        ],
-    )
-
-
-def is_hashable(x: VariableTracker) -> bool:
-    # NB - performing isinstance check on a LazVT realizes the VT, accidentally
-    # inserting the guard. To avoid this, lazyVT `is_hashable` methods looks at
-    # the underlying value without realizing the VT. Consider updating the
-    # lazyVT `is_hashable` method if you see unnecessary guarding for a key VT.
-    if (
-        isinstance(x, variables.LazyVariableTracker)
-        and not x.is_realized()
-        and x.is_hashable()
-    ):
-        return True
-    return x.is_python_hashable()
-
-
 class ConstDictVariable(VariableTracker):
+    # PyDict_Type: https://github.com/python/cpython/blob/v3.13.0/Objects/dictobject.c#L4825
+    _cpython_type = dict
+
     CONTAINS_GUARD = GuardBuilder.DICT_CONTAINS
+    NOT_CONTAINS_GUARD = GuardBuilder.DICT_NOT_CONTAINS
 
     _nonvar_fields = {
         "user_cls",
         *VariableTracker._nonvar_fields,
     }
 
-    class _HashableTracker:
-        """
-        Auxiliary opaque internal class that wraps a VariableTracker and makes it hashable
-        This should not be seen or touched by anything outside of ConstDictVariable and its children
-        Note that it's also fine to put VTs into dictionaries and sets, but doing so does not take into account aliasing
-        """
-
-        def __init__(self, vt: VariableTracker) -> None:
-            # We specialize SymNodes
-            vt = specialize_symnode(vt)
-
-            # If Dynamo does not know the hashability of the vt, it will raise unsupported here
-            if not is_hashable(vt):
-                raise_unhashable(vt)
-            self.vt = vt
-
-        def __hash__(self) -> int:
-            """
-            Computes the hash value for the wrapped VariableTracker.
-
-            For unrealized LazyVariableTrackers, uses the hash of the original value
-            to avoid realizing the tracker and inserting unnecessary guards.
-            For all other cases, delegates to the VariableTracker's get_python_hash method.
-
-            Returns:
-                The hash value of the underlying variable tracker
-            """
-            if (
-                isinstance(self.vt, variables.LazyVariableTracker)
-                and not self.vt.is_realized()
-                and self.vt.is_hashable()
-            ):
-                return hash(self.vt.original_value())
-            return self.vt.get_python_hash()
-
-        def __eq__(self, other: object) -> bool:
-            """
-            Checks equality between two _HashableTracker instances.
-
-            Delegates to the VariableTracker's is_python_equal method to compare
-            the underlying variable trackers for Python-level equality.
-
-            Args:
-                other: Another _HashableTracker instance to compare with
-
-            Returns:
-                True if the underlying variable trackers are Python-equal, False otherwise
-            """
-            if not isinstance(other, ConstDictVariable._HashableTracker):
-                return False
-            if self.vt is other.vt:
-                return True
-            return self.vt.is_python_equal(other.vt)
-
     def __init__(
         self,
         items: dict[VariableTracker, VariableTracker],
@@ -184,7 +99,7 @@ def __init__(
 
         super().__init__(**kwargs)
 
-        Hashable = ConstDictVariable._HashableTracker
+        Hashable = HashableTracker
 
         # Keys will just be HashableTrackers when cloning, in any other case they'll be VariableTrackers
         assert all(
@@ -194,8 +109,8 @@ def __init__(
         )
 
         def make_hashable(
-            key: Union[VariableTracker, "ConstDictVariable._HashableTracker"],
-        ) -> "ConstDictVariable._HashableTracker":
+            key: Union[VariableTracker, "HashableTracker"],
+        ) -> "HashableTracker":
             return key if isinstance(key, Hashable) else Hashable(key)
 
         dict_cls = self._get_dict_cls_from_user_cls(user_cls)
@@ -253,11 +168,12 @@ def python_type(self) -> type:
 
     def __contains__(self, vt: VariableTracker) -> bool:
         assert isinstance(vt, VariableTracker)
-        Hashable = ConstDictVariable._HashableTracker
-        return (
-            vt.is_python_hashable()
-            and Hashable(vt) in self.items
-            and not isinstance(self.items[Hashable(vt)], variables.DeletedVariable)
+        Hashable = HashableTracker
+        if not is_hashable(vt):
+            return False
+        key = Hashable(vt)
+        return key in self.items and not isinstance(
+            self.items[key], variables.DeletedVariable
         )
 
     def call_tree_map_branch(
@@ -310,6 +226,60 @@ def call_tree_map_branch(
             mutation_type=ValueMutationNew(),
         )
 
+    def call_tree_map_with_path_branch(
+        self,
+        tx: "InstructionTranslator",
+        tree_map_fn: "UserFunctionVariable",
+        map_fn: VariableTracker,
+        rest: Sequence[VariableTracker],
+        tree_map_kwargs: dict[str, VariableTracker],
+        keypath: tuple[Any, ...],
+    ) -> VariableTracker:
+        other_dicts: list[ConstDictVariable] = []
+        for candidate in rest:
+            candidate = candidate.realize()
+            if not isinstance(candidate, ConstDictVariable) or len(
+                candidate.items
+            ) != len(self.items):
+                return self._tree_map_with_path_fallback(
+                    tx, tree_map_fn, map_fn, rest, tree_map_kwargs, keypath
+                )
+            other_dicts.append(candidate)
+
+        new_items_hashed = type(self.items)()
+        for key_tracker, value in self.items.items():
+            sibling_leaves: list[VariableTracker] = []
+            for candidate in other_dicts:
+                try:
+                    sibling_leaves.append(candidate.items[key_tracker])
+                except KeyError:
+                    return self._tree_map_with_path_fallback(
+                        tx, tree_map_fn, map_fn, rest, tree_map_kwargs, keypath
+                    )
+            key_const = key_tracker.vt.as_python_constant()
+            child_keypath = keypath + (MappingKey(key_const),)
+            new_items_hashed[key_tracker] = value.call_tree_map_with_path(
+                tx,
+                tree_map_fn,
+                map_fn,
+                sibling_leaves,
+                tree_map_kwargs,
+                child_keypath,
+            )
+
+        updated_original_items = {
+            key_tracker.vt: new_items_hashed[key_tracker]
+            for key_tracker in new_items_hashed
+        }
+
+        return self.clone(
+            items=new_items_hashed,
+            original_items=updated_original_items,
+            should_reconstruct_all=True,
+            source=None,
+            mutation_type=ValueMutationNew(),
+        )
+
     def len(self) -> int:
         return sum(
             not isinstance(x, variables.DeletedVariable) for x in self.items.values()
@@ -322,7 +292,7 @@ def has_new_items(self) -> bool:
         )
 
     def is_new_item(
-        self, value: Optional[VariableTracker], other: VariableTracker
+        self, value: VariableTracker | None, other: VariableTracker
     ) -> bool:
         # compare the id of the realized values if both values are not lazy VTs
         if value and value.is_realized() and other.is_realized():
@@ -391,24 +361,15 @@ def reconstruct(self, codegen: "PyCodegen") -> None:
     def getitem_const_raise_exception_if_absent(
         self, tx: "InstructionTranslator", arg: VariableTracker
     ) -> VariableTracker:
-        key = ConstDictVariable._HashableTracker(arg)
+        key = HashableTracker(arg)
         if key not in self.items:
-            try:
-                error_message = (
-                    f"Dict key lookup failed for {str(arg)}. "
-                    f"Debug representation of the key is {arg.debug_repr()!r}"
-                )
-            except Exception:
-                error_message = ConstantVariable.create(
-                    f"Dict key lookup failed for {str(arg)}"
-                )
-            raise_observed_exception(KeyError, tx, args=[error_message])
+            raise_observed_exception(KeyError, tx, args=[arg])
         return self.items[key]
 
     def getitem_const(
         self, tx: "InstructionTranslator", arg: VariableTracker
     ) -> VariableTracker:
-        key = ConstDictVariable._HashableTracker(arg)
+        key = HashableTracker(arg)
         if key not in self.items:
             msg = f"Dictionary key {arg.value} not found during tracing"  # type: ignore[attr-defined]
             unimplemented(
@@ -422,8 +383,8 @@ def getitem_const(
             )
         return self.items[key]
 
-    def maybe_getitem_const(self, arg: VariableTracker) -> Optional[VariableTracker]:
-        key = ConstDictVariable._HashableTracker(arg)
+    def maybe_getitem_const(self, arg: VariableTracker) -> VariableTracker | None:
+        key = HashableTracker(arg)
         if key not in self.items:
             return None
         return self.items[key]
@@ -431,7 +392,7 @@ def maybe_getitem_const(self, arg: VariableTracker) -> Optional[VariableTracker]
     def realize_key_vt(self, arg: VariableTracker) -> None:
         # Realize the LazyVT on a particular index
         assert arg in self
-        key = ConstDictVariable._HashableTracker(arg)
+        key = HashableTracker(arg)
         index = tuple(self.items.keys()).index(key)
         original_key_vt = tuple(self.original_items.keys())[index]
         if isinstance(original_key_vt, variables.LazyVariableTracker):
@@ -466,12 +427,14 @@ def install_dict_contains_guard(
 
         contains = args[0] in self
         if args[0].source is None and args[0].is_python_constant():
+            guard_fn = (
+                type(self).CONTAINS_GUARD if contains else type(self).NOT_CONTAINS_GUARD
+            )
             install_guard(
                 self.make_guard(
                     functools.partial(
-                        type(self).CONTAINS_GUARD,
+                        guard_fn,
                         key=args[0].as_python_constant(),
-                        invert=not contains,
                     )
                 )
             )
@@ -481,6 +444,15 @@ def install_dict_contains_guard(
             else:
                 self.install_dict_keys_match_guard()
 
+    def mp_subscript_impl(
+        self,
+        tx: "InstructionTranslator",
+        key: VariableTracker,
+    ) -> VariableTracker:
+        # dict_subscript: https://github.com/python/cpython/blob/62a6e898e01/Objects/dictobject.c#L3673-L3706
+        # Unhashable key check happens inside _HashableTracker (raise_unhashable → TypeError).
+        return self.getitem_const_raise_exception_if_absent(tx, key)
+
     def call_method(
         self,
         tx: "InstructionTranslator",
@@ -495,23 +467,18 @@ def call_method(
         # corresponding value VT. For __contains__, we add a DICT_CONTAINS
         # guard. But for all the other methods, we insert the DICT_KEYS_MATCH
         # guard to be conservative.
-        from . import BuiltinVariable, ConstantVariable
+        from . import DictBuiltinVariable
         from .builder import SourcelessBuilder
 
-        Hashable = ConstDictVariable._HashableTracker
+        Hashable = HashableTracker
 
         if name == "__init__":
-            temp_dict_vt = SourcelessBuilder.create(tx, dict).call_dict(
-                tx, *args, **kwargs
+            temp_dict_vt = DictBuiltinVariable.call_custom_dict(
+                tx, dict, *args, **kwargs
             )
             tx.output.side_effects.mutation(self)
             self.items.update(temp_dict_vt.items)  # type: ignore[attr-defined]
-            return CONSTANT_VARIABLE_NONE
-        elif name == "__getitem__":
-            # Key guarding - Nothing to do. LazyVT for value will take care.
-            if len(args) != 1:
-                raise_args_mismatch(tx, name, "1 args", f"{len(args)} args")
-            return self.getitem_const_raise_exception_if_absent(tx, args[0])
+            return ConstantVariable.create(None)
         elif name == "items":
             if args or kwargs:
                 raise_args_mismatch(
@@ -557,16 +524,6 @@ def call_method(
             return self.clone(
                 items=self.items.copy(), mutation_type=ValueMutationNew(), source=None
             )
-        elif name == "__len__":
-            if args or kwargs:
-                raise_args_mismatch(
-                    tx,
-                    name,
-                    "0 args and 0 kwargs",
-                    f"{len(args)} args and {len(kwargs)} kwargs",
-                )
-            self.install_dict_keys_match_guard()
-            return ConstantVariable.create(len(self.items))
         elif name == "__setitem__" and self.is_mutable():
             arg_hashable = args and is_hashable(args[0])
             if not arg_hashable:
@@ -582,7 +539,7 @@ def call_method(
                 )
             tx.output.side_effects.mutation(self)
             self.items[Hashable(args[0])] = args[1]
-            return CONSTANT_VARIABLE_NONE
+            return ConstantVariable.create(None)
         elif name == "__delitem__" and self.is_mutable():
             arg_hashable = args and is_hashable(args[0])
             if arg_hashable:
@@ -590,7 +547,7 @@ def call_method(
                 self.should_reconstruct_all = True
                 tx.output.side_effects.mutation(self)
                 self.items.__delitem__(Hashable(args[0]))
-                return CONSTANT_VARIABLE_NONE
+                return ConstantVariable.create(None)
             else:
                 return super().call_method(tx, name, args, kwargs)
         elif name == "get":
@@ -605,7 +562,7 @@ def call_method(
                 self.install_dict_contains_guard(tx, args)
                 if len(args) == 1:
                     # if default is not given, return None
-                    return CONSTANT_VARIABLE_NONE
+                    return ConstantVariable.create(None)
                 return args[1]
             # Key guarding - Nothing to do.
             return self.getitem_const(tx, args[0])
@@ -637,8 +594,13 @@ def call_method(
                 raise_args_mismatch(tx, name)
 
             if not self.items:
-                msg = ConstantVariable.create("popitem(): dictionary is empty")
-                raise_observed_exception(KeyError, tx, args=[msg])
+                raise_observed_exception(
+                    KeyError,
+                    tx,
+                    args=[
+                        "popitem(): dictionary is empty",
+                    ],
+                )
 
             if self.user_cls is collections.OrderedDict and (
                 len(args) == 1 or "last" in kwargs
@@ -668,7 +630,7 @@ def call_method(
             self.should_reconstruct_all = True
             tx.output.side_effects.mutation(self)
             self.items.clear()
-            return CONSTANT_VARIABLE_NONE
+            return ConstantVariable.create(None)
         elif name == "update" and self.is_mutable():
             # In general, this call looks like `a.update(b, x=1, y=2, ...)`.
             # Either `b` or the kwargs is omittable, but not both.
@@ -678,22 +640,25 @@ def call_method(
             if has_arg or has_kwargs:
                 tx.output.side_effects.mutation(self)
                 if has_arg:
+                    dict_vt: VariableTracker
                     if isinstance(args[0], ConstDictVariable):
                         # NB - Guard on all the keys of the other dict to ensure
                         # correctness.
                         args[0].install_dict_keys_match_guard()
-                        dict_vt: ConstDictVariable = args[0]
+                        dict_vt = args[0]
                     else:
-                        dict_vt = BuiltinVariable.call_custom_dict(tx, dict, args[0])  # type: ignore[assignment]
+                        dict_vt = DictBuiltinVariable.call_custom_dict(
+                            tx, dict, args[0]
+                        )
                     self.items.update(dict_vt.items)  # type: ignore[attr-defined]
                 if has_kwargs:
                     # Handle kwargs
                     kwargs_hashable = {
-                        Hashable(ConstantVariable.create(k)): v
+                        Hashable(VariableTracker.build(tx, k)): v
                         for k, v in kwargs.items()
                     }
                     self.items.update(kwargs_hashable)
-                return CONSTANT_VARIABLE_NONE
+                return ConstantVariable.create(None)
             else:
                 return super().call_method(tx, name, args, kwargs)
         elif name == "__contains__":
@@ -711,7 +676,7 @@ def call_method(
 
             self.install_dict_contains_guard(tx, args)
             contains = args[0] in self
-            return ConstantVariable.create(contains)
+            return VariableTracker.build(tx, contains)
         elif name == "setdefault" and self.is_mutable():
             if len(args) not in (1, 2):
                 raise_args_mismatch(
@@ -738,7 +703,7 @@ def call_method(
                 return value
             else:
                 if len(args) == 1:
-                    x = CONSTANT_VARIABLE_NONE
+                    x = ConstantVariable.create(None)
                 else:
                     x = args[1]
                 tx.output.side_effects.mutation(self)
@@ -759,7 +724,7 @@ def call_method(
 
             key = Hashable(args[0])
             self.items.move_to_end(key, last=last)
-            return CONSTANT_VARIABLE_NONE
+            return ConstantVariable.create(None)
         elif name == "__eq__" and istype(
             self, ConstDictVariable
         ):  # don't let Set use this function
@@ -770,8 +735,9 @@ def call_method(
                 tx, [self, args[0]], {}
             )
         elif name == "__ne__":
-            return ConstantVariable.create(
-                not self.call_method(tx, "__eq__", args, kwargs).value  # type: ignore[attr-defined]
+            return VariableTracker.build(
+                tx,
+                not self.call_method(tx, "__eq__", args, kwargs).value,  # type: ignore[attr-defined]
             )
         elif name == "__or__":
             if len(args) != 1:
@@ -803,6 +769,12 @@ def call_method(
                     variables.DefaultDictVariable,
                 ),
             ):
+                # Unwrap UserDefinedDictVariable to its underlying ConstDictVariable
+                if isinstance(other, variables.UserDefinedDictVariable):
+                    assert other._base_vt is not None
+                    assert isinstance(other._base_vt, ConstDictVariable)
+                    other = other._base_vt
+
                 # Always return the specialized dictionary, and in the case
                 # both are specialized, take the first to be the type of the
                 # new dictionary
@@ -810,7 +782,6 @@ def call_method(
                     user_cls = self.user_cls
                     to_cpy = self
                 else:
-                    assert isinstance(other, ConstDictVariable)
                     user_cls = other.user_cls
                     to_cpy = other
 
@@ -828,11 +799,14 @@ def call_method(
                 new_dict_vt.items.update(args[0].items)  # type: ignore[attr-defined]
                 return new_dict_vt
             else:
-                err_msg = (
-                    f"unsupported operand type(s) for |: '{self.python_type().__name__}'"
-                    f"and '{other.python_type().__name__}'"
+                raise_observed_exception(
+                    TypeError,
+                    tx,
+                    args=[
+                        f"unsupported operand type(s) for |: '{self.python_type().__name__}'"
+                        f"and '{other.python_type().__name__}'"
+                    ],
                 )
-                raise_observed_exception(TypeError, tx, args=[err_msg])
         elif name == "__ior__":
             self.call_method(tx, "update", args, kwargs)
             return self
@@ -851,6 +825,11 @@ def unpack_var_sequence(self, tx: "InstructionTranslator") -> list[VariableTrack
         self.install_dict_keys_match_guard()
         return [x.vt for x in self.items]
 
+    def mp_length(self, tx: "InstructionTranslator") -> VariableTracker:
+        """Mapping length for dict objects."""
+        self.install_dict_keys_match_guard()
+        return VariableTracker.build(tx, len(self.items))
+
     def call_obj_hasattr(
         self, tx: "InstructionTranslator", name: str
     ) -> ConstantVariable:
@@ -886,8 +865,16 @@ def is_python_hashable(self) -> bool:
         """
         return False
 
+    def var_getattr(self, tx: "InstructionTranslator", name: str):
+        if name == "__class__":
+            return VariableTracker.build(tx, self.python_type())
+        return super().var_getattr(tx, name)
+
 
 class MappingProxyVariable(VariableTracker):
+    # PyDictProxy_Type: https://github.com/python/cpython/blob/v3.13.0/Objects/descrobject.c#L1995
+    _cpython_type = types.MappingProxyType
+
     # proxies to the original dict_vt
     def __init__(self, dv_dict: ConstDictVariable, **kwargs: Any) -> None:
         super().__init__(**kwargs)
@@ -927,13 +914,7 @@ def reconstruct(self, codegen: "PyCodegen") -> None:
         codegen(self.dv_dict)
         codegen.extend_output(create_call_function(1, False))
 
-    def call_method(
-        self,
-        tx: "InstructionTranslator",
-        name: str,
-        args: list[VariableTracker],
-        kwargs: dict[str, VariableTracker],
-    ) -> VariableTracker:
+    def _check_mutation_guard(self, tx: "InstructionTranslator") -> None:
         if self.source and tx.output.side_effects.has_existing_dict_mutation():
             msg = (
                 "A dict has been modified while we have an existing mappingproxy object. "
@@ -954,13 +935,35 @@ def call_method(
                     "Or avoid using the mapping proxy objects after modifying its underlying dictionary",
                 ],
             )
+
+    def mp_subscript_impl(
+        self,
+        tx: "InstructionTranslator",
+        key: VariableTracker,
+    ) -> VariableTracker:
+        # mappingproxy_getitem: https://github.com/python/cpython/blob/62a6e898e01/Objects/descrobject.c#L1052-L1056
+        # TODO(follow-up): add tests for invalid key type, missing key
+        self._check_mutation_guard(tx)
+        return self.dv_dict.mp_subscript_impl(tx, key)
+
+    def call_method(
+        self,
+        tx: "InstructionTranslator",
+        name: str,
+        args: list[VariableTracker],
+        kwargs: dict[str, VariableTracker],
+    ) -> VariableTracker:
+        self._check_mutation_guard(tx)
         return self.dv_dict.call_method(tx, name, args, kwargs)
 
+    def mp_length(self, tx: "InstructionTranslator") -> VariableTracker:
+        return self.dv_dict.mp_length(tx)
+
     def call_obj_hasattr(
         self, tx: "InstructionTranslator", name: str
     ) -> ConstantVariable:
         if self.python_type() is types.MappingProxyType:
-            return ConstantVariable.create(name in types.MappingProxyType.__dict__)
+            return VariableTracker.build(tx, name in types.MappingProxyType.__dict__)
         return super().call_obj_hasattr(tx, name)
 
 
@@ -976,17 +979,19 @@ def install_dict_contains_guard(
 
 
 class DefaultDictVariable(ConstDictVariable):
+    _cpython_type = collections.defaultdict
+
     def __init__(
         self,
         items: dict[VariableTracker, VariableTracker],
         user_cls: type,
-        default_factory: Optional[VariableTracker] = None,
+        default_factory: VariableTracker | None = None,
         **kwargs: Any,
     ) -> None:
         super().__init__(items, user_cls, **kwargs)
         assert user_cls is collections.defaultdict
         if default_factory is None:
-            default_factory = CONSTANT_VARIABLE_NONE
+            default_factory = ConstantVariable.create(None)
         self.default_factory = default_factory
 
     def is_python_constant(self) -> bool:
@@ -1007,12 +1012,30 @@ def is_supported_arg(arg: VariableTracker) -> bool:
         return isinstance(
             arg,
             (
-                variables.BuiltinVariable,
+                variables.BaseBuiltinVariable,
                 variables.functions.BaseUserFunctionVariable,
                 variables.functions.PolyfilledFunctionVariable,
             ),
         ) or (isinstance(arg, variables.ConstantVariable) and arg.value is None)
 
+    def mp_subscript_impl(
+        self,
+        tx: "InstructionTranslator",
+        key: VariableTracker,
+    ) -> VariableTracker:
+        # Mirrors CPython's defaultdict.__getitem__ (dict_subscript → __missing__).
+        # defaultdict.__missing__: https://github.com/python/cpython/blob/62a6e898e01/Modules/_collectionsmodule.c#L2233-L2254
+        # Key present → normal dict lookup (same as ConstDictVariable.mp_subscript_impl).
+        if key in self:
+            return self.getitem_const(tx, key)
+
+        if self.default_factory.is_constant_none():
+            raise_observed_exception(KeyError, tx, args=[key])
+        else:
+            default_var = self.default_factory.call_function(tx, [], {})
+            super().call_method(tx, "__setitem__", [key, default_var], {})
+            return default_var
+
     def call_method(
         self,
         tx: "InstructionTranslator",
@@ -1020,25 +1043,7 @@ def call_method(
         args: list[VariableTracker],
         kwargs: dict[str, VariableTracker],
     ) -> VariableTracker:
-        if name == "__getitem__":
-            if len(args) != 1:
-                raise_args_mismatch(tx, name, "1 args", f"{len(args)} args")
-
-            if args[0] in self:
-                return self.getitem_const(tx, args[0])
-            else:
-                if (
-                    istype(self.default_factory, ConstantVariable)
-                    and self.default_factory.value is None
-                ):
-                    raise_observed_exception(KeyError, tx, args=[args[0]])
-                else:
-                    default_var = self.default_factory.call_function(tx, [], {})
-                    super().call_method(
-                        tx, "__setitem__", [args[0], default_var], kwargs
-                    )
-                    return default_var
-        elif name == "__setattr__" and self.is_mutable:
+        if name == "__setattr__" and self.is_mutable:
             if len(args) != 2:
                 raise_args_mismatch(tx, name, "2 args", f"{len(args)} args")
             # Setting a default factory must be a callable or None type
@@ -1047,13 +1052,13 @@ def call_method(
             ) and self.is_supported_arg(args[1]):
                 tx.output.side_effects.mutation(self)
                 self.default_factory = args[1]
-                return CONSTANT_VARIABLE_NONE
+                return ConstantVariable.create(None)
             return super().call_method(tx, name, args, kwargs)
         elif name == "__eq__":
             if len(args) != 1:
                 raise_args_mismatch(tx, name, "1 args", f"{len(args)} args")
 
-            return variables.UserFunctionVariable(polyfills.dict___eq__).call_function(
+            return VariableTracker.build(tx, polyfills.dict___eq__).call_function(
                 tx, [self, args[0]], {}
             )
         else:
@@ -1098,558 +1103,6 @@ def reconstruct(self, codegen: "PyCodegen") -> None:
         )
 
 
-# TODO: Implementing this via inheritance rather than composition is a
-# footgun, because self method calls in dict will route back to the set
-# implementation, which is almost assuredly wrong
-class SetVariable(ConstDictVariable):
-    """We model a sets as dictionary with None values"""
-
-    CONTAINS_GUARD = GuardBuilder.SET_CONTAINS
-
-    def __init__(
-        self,
-        items: Iterable[VariableTracker],
-        **kwargs: Any,
-    ) -> None:
-        # Items can be either VariableTrackers or _HashableTrackers (from set ops).
-        # For VariableTrackers, realize them to ensure aliasing guards are installed
-        # when the same object appears multiple times.
-        realized_items = []
-        for item in items:
-            if isinstance(item, ConstDictVariable._HashableTracker):
-                # Already a _HashableTracker from a set operation
-                realized_items.append(item)
-            else:
-                # VariableTracker - realize to install guards
-                # pyrefly: ignore [bad-argument-type]
-                realized_items.append(item.realize())
-        # pyrefly: ignore[bad-assignment]
-        items = dict.fromkeys(realized_items, SetVariable._default_value())
-        # pyrefly: ignore[bad-argument-type]
-        super().__init__(items, **kwargs)
-
-    def debug_repr(self) -> str:
-        if not self.items:
-            return "set()"
-        else:
-            items: list[str] = []
-            for v in self.items:
-                vt = v.vt if isinstance(v, ConstDictVariable._HashableTracker) else v
-                val_str = repr(vt.value) if hasattr(vt, "value") else vt.debug_repr()
-                items.append(val_str)
-            return "{" + ",".join(items) + "}"
-
-    @property
-    def set_items(self) -> set["ConstDictVariable._HashableTracker"]:
-        return set(self.items.keys())
-
-    @staticmethod
-    def _default_value() -> VariableTracker:
-        # Variable to fill in he keys of the dictionary
-        return CONSTANT_VARIABLE_NONE
-
-    def as_proxy(self) -> Any:
-        return {k.vt.as_proxy() for k in self.set_items}
-
-    def python_type(self) -> type:
-        return set
-
-    def as_python_constant(self) -> Any:
-        return {k.vt.as_python_constant() for k in self.set_items}
-
-    def reconstruct(self, codegen: "PyCodegen") -> None:
-        codegen.foreach([x.vt for x in self.set_items])
-        codegen.append_output(create_instruction("BUILD_SET", arg=len(self.set_items)))
-
-    def _fast_set_method(
-        self,
-        tx: "InstructionTranslator",
-        fn: Any,
-        args: list[VariableTracker],
-        kwargs: dict[str, VariableTracker],
-    ) -> VariableTracker:
-        try:
-            res = fn(
-                *[x.as_python_constant() for x in [self, *args]],
-                **{k: v.as_python_constant() for k, v in kwargs.items()},
-            )
-        except Exception as exc:
-            raise_observed_exception(
-                type(exc), tx, args=list(map(ConstantVariable.create, exc.args))
-            )
-        # pyrefly: ignore[unbound-name]
-        return VariableTracker.build(tx, res)
-
-    def call_method(
-        self,
-        tx: "InstructionTranslator",
-        name: str,
-        args: list[VariableTracker],
-        kwargs: dict[str, VariableTracker],
-    ) -> VariableTracker:
-        # We forward the calls to the dictionary model
-        from ..utils import check_constant_args
-        from .builder import SourcelessBuilder
-
-        if (
-            name
-            in (
-                "isdisjoint",
-                "union",
-                "intersection",
-                "difference",
-                "symmetric_difference",
-            )
-            and check_constant_args(args, kwargs)
-            and self.python_type() is set
-        ):
-            py_type = self.python_type()
-            return self._fast_set_method(tx, getattr(py_type, name), args, kwargs)
-
-        if name == "__init__":
-            temp_set_vt = SourcelessBuilder.create(tx, set).call_set(
-                tx, *args, **kwargs
-            )
-            tx.output.side_effects.mutation(self)
-            self.items.clear()
-            self.items.update(temp_set_vt.items)  # type: ignore[attr-defined]
-            return CONSTANT_VARIABLE_NONE
-        elif name == "add":
-            if kwargs or len(args) != 1:
-                raise_args_mismatch(
-                    tx,
-                    name,
-                    "1 args and 0 kwargs",
-                    f"{len(args)} args and {len(kwargs)} kwargs",
-                )
-            name = "__setitem__"
-            args = [args[0], SetVariable._default_value()]
-        elif name == "pop":
-            if kwargs or args:
-                raise_args_mismatch(
-                    tx,
-                    name,
-                    "0 args and 0 kwargs",
-                    f"{len(args)} args and {len(kwargs)} kwargs",
-                )
-            # Choose an item at random and pop it via the Dict.pop method
-            try:
-                result: VariableTracker = self.set_items.pop().vt  # type: ignore[assignment]
-            except KeyError as e:
-                raise_observed_exception(
-                    KeyError, tx, args=list(map(ConstantVariable.create, e.args))
-                )
-            # pyrefly: ignore[unbound-name]
-            super().call_method(tx, name, [result], kwargs)
-            # pyrefly: ignore[unbound-name]
-            return result
-        elif name == "isdisjoint":
-            if kwargs or len(args) != 1:
-                raise_args_mismatch(
-                    tx,
-                    name,
-                    "1 args and 0 kwargs",
-                    f"{len(args)} args and {len(kwargs)} kwargs",
-                )
-            return SourcelessBuilder.create(tx, polyfills.set_isdisjoint).call_function(
-                tx, [self, args[0]], {}
-            )
-        elif name == "intersection":
-            if kwargs:
-                raise_args_mismatch(tx, name, "0 kwargs", f"{len(kwargs)} kwargs")
-            return SourcelessBuilder.create(
-                tx, polyfills.set_intersection
-            ).call_function(
-                tx,
-                [self, *args],
-                {"cls": self.python_type_var()},
-            )
-        elif name == "intersection_update":
-            if kwargs:
-                raise_args_mismatch(tx, name, "0 kwargs", f"{len(kwargs)} kwargs")
-            return SourcelessBuilder.create(
-                tx, polyfills.set_intersection_update
-            ).call_function(tx, [self, *args], {})
-        elif name == "union":
-            if kwargs:
-                raise_args_mismatch(tx, name, "0 kwargs", f"{len(kwargs)} kwargs")
-            return SourcelessBuilder.create(tx, polyfills.set_union).call_function(
-                tx,
-                [self, *args],
-                {"cls": self.python_type_var()},
-            )
-        elif name == "difference":
-            if kwargs:
-                raise_args_mismatch(
-                    tx, name, f"Expect: 0 kwargs, Actual: {len(kwargs)} kwargs"
-                )
-            return SourcelessBuilder.create(tx, polyfills.set_difference).call_function(
-                tx,
-                [self, *args],
-                {"cls": self.python_type_var()},
-            )
-        elif name == "difference_update":
-            if kwargs:
-                raise_args_mismatch(tx, name, "0 kwargs", f"{len(kwargs)} kwargs")
-            return SourcelessBuilder.create(
-                tx, polyfills.set_difference_update
-            ).call_function(tx, [self, *args], {})
-        elif name == "symmetric_difference":
-            if kwargs or len(args) != 1:
-                raise_args_mismatch(
-                    tx,
-                    name,
-                    "1 args and 0 kwargs",
-                    f"{len(args)} args and {len(kwargs)} kwargs",
-                )
-            return SourcelessBuilder.create(
-                tx, polyfills.set_symmetric_difference
-            ).call_function(
-                tx,
-                [self, *args],
-                {"cls": self.python_type_var()},
-            )
-        elif name == "symmetric_difference_update":
-            if kwargs or len(args) != 1:
-                raise_args_mismatch(
-                    tx,
-                    name,
-                    "1 args and 0 kwargs",
-                    f"{len(args)} args and {len(kwargs)} kwargs",
-                )
-            return SourcelessBuilder.create(
-                tx, polyfills.set_symmetric_difference_update
-            ).call_function(tx, [self, *args], {})
-        elif name == "update" and self.is_mutable():
-            if kwargs:
-                raise_args_mismatch(tx, name, "0 kwargs", f"{len(kwargs)} kwargs")
-            return SourcelessBuilder.create(tx, polyfills.set_update).call_function(
-                tx, [self, *args], {}
-            )
-        elif name == "remove":
-            if kwargs or len(args) != 1:
-                raise_args_mismatch(
-                    tx,
-                    name,
-                    "1 args and 0 kwargs",
-                    f"{len(args)} args and {len(kwargs)} kwargs",
-                )
-            if args[0] not in self:
-                raise_observed_exception(KeyError, tx, args=args)
-            return super().call_method(tx, "pop", args, kwargs)
-        elif name == "discard":
-            if kwargs or len(args) != 1:
-                raise_args_mismatch(
-                    tx,
-                    name,
-                    "1 args and 0 kwargs",
-                    f"{len(args)} args and {len(kwargs)} kwargs",
-                )
-            if args[0] in self:
-                return super().call_method(tx, "pop", args, kwargs)
-            else:
-                return CONSTANT_VARIABLE_NONE
-        elif name in ("issubset", "issuperset"):
-            if len(args) != 1:
-                raise_args_mismatch(tx, name, "1 args", f"{len(args)} args")
-
-            op = {
-                "issubset": operator.le,
-                "issuperset": operator.ge,
-            }
-            other = args[0].realize()
-            if not istype(other, SetVariable):
-                other = SourcelessBuilder.create(tx, set).call_function(tx, [other], {})
-            return SourcelessBuilder.create(tx, op.get(name)).call_function(
-                tx, [self, other], {}
-            )
-        elif name in ("__and__", "__or__", "__xor__", "__sub__"):
-            m = {
-                "__and__": "intersection",
-                "__or__": "union",
-                "__xor__": "symmetric_difference",
-                "__sub__": "difference",
-            }.get(name)
-            if not isinstance(args[0], (SetVariable, variables.UserDefinedSetVariable)):
-                msg = ConstantVariable.create(
-                    f"unsupported operand type(s) for {name}: '{self.python_type_name()}' and '{args[0].python_type_name()}'"
-                )
-                raise_observed_exception(TypeError, tx, args=[msg])
-            assert m is not None
-            return self.call_method(tx, m, args, kwargs)
-        elif name in ("__iand__", "__ior__", "__ixor__", "__isub__"):
-            if not isinstance(args[0], (SetVariable, variables.UserDefinedSetVariable)):
-                msg = ConstantVariable.create(
-                    f"unsupported operand type(s) for {name}: '{self.python_type_name()}' and '{args[0].python_type_name()}'"
-                )
-                raise_observed_exception(TypeError, tx, args=[msg])
-            m = {
-                "__iand__": "intersection_update",
-                "__ior__": "update",
-                "__ixor__": "symmetric_difference_update",
-                "__isub__": "difference_update",
-            }.get(name)
-            assert m is not None
-            self.call_method(tx, m, args, kwargs)
-            return self
-        elif name == "__eq__":
-            if not isinstance(args[0], (SetVariable, variables.UserDefinedSetVariable)):
-                return ConstantVariable.create(False)
-            r = self.call_method(tx, "symmetric_difference", args, kwargs)
-            return ConstantVariable.create(len(r.set_items) == 0)  # type: ignore[attr-defined]
-        elif name in cmp_name_to_op_mapping:
-            if not isinstance(args[0], (SetVariable, variables.UserDefinedSetVariable)):
-                return ConstantVariable.create(NotImplemented)
-            return ConstantVariable.create(
-                cmp_name_to_op_mapping[name](self.set_items, args[0].set_items)  # type: ignore[attr-defined]
-            )
-        return super().call_method(tx, name, args, kwargs)
-
-    def python_type_var(self) -> "BuiltinVariable":
-        return variables.BuiltinVariable(set)
-
-    def getitem_const(
-        self, tx: "InstructionTranslator", arg: VariableTracker
-    ) -> VariableTracker:
-        raise RuntimeError("Illegal to getitem on a set")
-
-    def install_dict_keys_match_guard(self) -> None:
-        # Already EQUALS_MATCH guarded
-        pass
-
-
-class OrderedSetClassVariable(VariableTracker):
-    def __init__(self, **kwargs: Any) -> None:
-        super().__init__(**kwargs)
-
-    def as_python_constant(self) -> type[OrderedSet[Any]]:
-        return OrderedSet
-
-    def var_getattr(self, tx: "InstructionTranslator", name: str) -> VariableTracker:
-        if name == "__new__":
-            from .misc import GetAttrVariable
-
-            if self.source:
-                attr_source = AttrSource(self.source, name)
-            else:
-                attr_source = None
-            return GetAttrVariable(self, name, source=attr_source)
-        else:
-            return super().var_getattr(tx, name)
-
-    def call_method(
-        self,
-        tx: "InstructionTranslator",
-        name: str,
-        args: list[VariableTracker],
-        kwargs: dict[str, VariableTracker],
-    ) -> VariableTracker:
-        from .builtin import set_methods
-
-        if name == "__new__":
-            if len(args) != 2 or kwargs:
-                raise_args_mismatch(
-                    tx,
-                    name,
-                    "OrderedSet.__new__ only accepts one arg"
-                    f"{len(args)} args and {len(kwargs)} kwargs",
-                )
-
-            return variables.OrderedSetVariable([], mutation_type=ValueMutationNew())
-
-        resolved_fn = getattr(set, name)
-        if resolved_fn in set_methods and isinstance(args[0], variables.SetVariable):
-            return args[0].call_method(tx, name, args[1:], kwargs)
-
-        return super().call_method(tx, name, args, kwargs)
-
-    def call_function(
-        self,
-        tx: "InstructionTranslator",
-        args: Sequence[VariableTracker],
-        kwargs: dict[str, VariableTracker],
-    ) -> "OrderedSetVariable":
-        if len(args) > 1 or kwargs:
-            raise_args_mismatch(
-                tx,
-                "OrderedSet",
-                "OrderedSet only accepts one arg"
-                f"{len(args)} args and {len(kwargs)} kwargs",
-            )
-
-        if len(args) == 0:
-            # pyrefly: ignore [implicit-any]
-            items = []
-        else:
-            items = args[0].force_unpack_var_sequence(tx)
-        return variables.OrderedSetVariable(items, mutation_type=ValueMutationNew())
-
-
-class OrderedSetVariable(SetVariable):
-    def debug_repr(self) -> str:
-        if not self.items:
-            return "OrderedSet([])"
-        else:
-            items: list[str] = []
-            for k, v in self.items:
-                key_str = (
-                    repr(k.vt.value) if hasattr(k.vt, "value") else k.vt.debug_repr()
-                )
-                items.append(key_str)
-            return "OrderedSet([" + ",".join(items) + "])"
-
-    def as_python_constant(self) -> OrderedSet[Any]:
-        return OrderedSet([k.vt.as_python_constant() for k in self.set_items])
-
-    def python_type(self) -> type[OrderedSet[Any]]:
-        return OrderedSet
-
-    # pyrefly: ignore[bad-override]
-    def python_type_var(self) -> OrderedSetClassVariable:
-        return OrderedSetClassVariable()
-
-    def reconstruct(self, codegen: "PyCodegen") -> None:
-        codegen.add_push_null(
-            lambda: codegen.load_import_from("torch.utils._ordered_set", "OrderedSet")
-        )
-        codegen.foreach([x.vt for x in self.set_items])
-        codegen.append_output(create_instruction("BUILD_LIST", arg=len(self.set_items)))
-        codegen.extend_output(create_call_function(1, False))
-
-
-class FrozensetVariable(SetVariable):
-    def debug_repr(self) -> str:
-        if not self.items:
-            return "frozenset()"
-        else:
-            items: list[str] = []
-            for k in self.items:
-                key_str = (
-                    repr(k.vt.value) if hasattr(k.vt, "value") else k.vt.debug_repr()
-                )
-                items.append(key_str)
-            return "{" + ",".join(items) + "}"
-
-    @property
-    def set_items(self) -> set["ConstDictVariable._HashableTracker"]:
-        return self.items.keys()
-
-    def python_type(self) -> type:
-        return frozenset
-
-    def python_type_var(self) -> "BuiltinVariable":
-        return variables.BuiltinVariable(frozenset)
-
-    def as_python_constant(self) -> Any:
-        return frozenset({k.vt.as_python_constant() for k in self.set_items})
-
-    def reconstruct(self, codegen: "PyCodegen") -> None:
-        codegen.add_push_null(
-            lambda: codegen.extend_output(
-                [
-                    codegen.create_load_global("frozenset"),
-                ]
-            )
-        )
-        codegen.foreach([x.vt for x in self.set_items])
-        codegen.extend_output(
-            [
-                create_instruction("BUILD_LIST", arg=len(self.set_items)),
-                *create_call_function(1, False),
-            ]
-        )
-
-    def call_method(
-        self,
-        tx: "InstructionTranslator",
-        name: str,
-        args: list[VariableTracker],
-        kwargs: dict[str, VariableTracker],
-    ) -> VariableTracker:
-        if name in ["add", "pop", "update", "remove", "discard", "clear"]:
-            raise RuntimeError(f"Illegal call_method {name} on a frozenset")
-        elif name == "__init__":
-            # frozenset is immutable. Calling __init__ again shouldn't have any effect
-            # In[1]: s = frozenset([1, 2])
-            #
-            # In[2]: s.__init__([3, 4])
-            #
-            # In[3]: s
-            # frozenset({1, 2})
-            return CONSTANT_VARIABLE_NONE
-        elif name in (
-            "copy",
-            "difference",
-            "intersection",
-            "symmetric_difference",
-        ):
-            r = super().call_method(tx, name, args, kwargs)
-            return FrozensetVariable(r.items)  # type: ignore[attr-defined]
-        return super().call_method(tx, name, args, kwargs)
-
-    def is_python_hashable(self) -> Literal[True]:
-        """
-        Frozensets are immutable and hashable in Python.
-        """
-        return True
-
-    def get_python_hash(self) -> int:
-        return hash(self.as_python_constant())
-
-    def is_python_equal(self, other: object) -> bool:
-        return (
-            isinstance(other, VariableTracker)
-            and self.as_python_constant() == other.as_python_constant()
-        )
-
-
-class DictKeySetVariable(SetVariable):
-    def debug_repr(self) -> str:
-        if not self.items:
-            return "dict_keys([])"
-        else:
-            items: list[str] = []
-            for k in self.items:
-                key_str = (
-                    repr(k.vt.value) if hasattr(k.vt, "value") else k.vt.debug_repr()
-                )
-                items.append(key_str)
-            return "dict_keys([" + ",".join(items) + "])"
-
-    def install_dict_keys_match_guard(self) -> None:
-        # Already EQUALS_MATCH guarded
-        pass
-
-    def install_dict_contains_guard(
-        self, tx: "InstructionTranslator", args: list[VariableTracker]
-    ) -> None:
-        # Already EQUALS_MATCH guarded
-        pass
-
-    @property
-    def set_items(self) -> Any:
-        return self.items
-
-    def python_type(self) -> type:
-        return dict_keys
-
-    def as_python_constant(self) -> Any:
-        return dict.fromkeys(
-            {k.vt.as_python_constant() for k in self.set_items}, None
-        ).keys()
-
-    def call_method(
-        self,
-        tx: "InstructionTranslator",
-        name: str,
-        args: list[VariableTracker],
-        kwargs: dict[str, VariableTracker],
-    ) -> VariableTracker:
-        if name in ["add", "pop", "update", "remove", "discard", "clear"]:
-            raise RuntimeError(f"Illegal call_method {name} on a dict_keys")
-        return super().call_method(tx, name, args, kwargs)
-
-
 class DictViewVariable(VariableTracker):
     """
     Models _PyDictViewObject
@@ -1657,7 +1110,7 @@ class DictViewVariable(VariableTracker):
     This is an "abstract" class. Subclasses will override kv and the items method
     """
 
-    kv: Optional[str] = None
+    kv: str | None = None
 
     def __init__(self, dv_dict: ConstDictVariable, **kwargs: Any) -> None:
         super().__init__(**kwargs)
@@ -1700,20 +1153,25 @@ def call_method(
         args: list[VariableTracker],
         kwargs: dict[str, VariableTracker],
     ) -> VariableTracker:
-        if name == "__len__":
-            return self.dv_dict.call_method(tx, name, args, kwargs)
-        elif name == "__iter__":
+        if name == "__iter__":
             from .lists import ListIteratorVariable
 
             return ListIteratorVariable(
                 self.view_items_vt, mutation_type=ValueMutationNew()
             )
         elif name == "__repr__":
-            return ConstantVariable.create(self.debug_repr())
+            return VariableTracker.build(tx, self.debug_repr())
         return super().call_method(tx, name, args, kwargs)
 
+    def sq_length(self, tx: "InstructionTranslator") -> VariableTracker:
+        """Sequence length for dict view objects."""
+        return VariableTracker.build(tx, len(self.view_items))
+
 
 class DictKeysVariable(DictViewVariable):
+    # PyDictKeys_Type: https://github.com/python/cpython/blob/v3.13.0/Objects/dictobject.c#L6365
+    _cpython_type = dict_keys
+
     kv = "keys"
 
     @property
@@ -1763,16 +1221,28 @@ def call_method(
             m = getattr(self.set_items, name)
             r = m(args[0].set_items)  # type: ignore[attr-defined]
             return SetVariable(r)
-        if name in cmp_name_to_op_mapping:
-            if not isinstance(args[0], (SetVariable, DictKeysVariable)):
-                return ConstantVariable.create(NotImplemented)
-            return ConstantVariable.create(
-                cmp_name_to_op_mapping[name](self.set_items, args[0].set_items)  # type: ignore[attr-defined]
+        elif name in cmp_name_to_op_mapping:
+            if not isinstance(
+                args[0],
+                (
+                    SetVariable,
+                    variables.UserDefinedSetVariable,
+                    DictItemsVariable,
+                    DictKeysVariable,
+                ),
+            ):
+                return VariableTracker.build(tx, NotImplemented)
+            return VariableTracker.build(
+                tx,
+                cmp_name_to_op_mapping[name](self.set_items, args[0].set_items),  # type: ignore[attr-defined]
             )
         return super().call_method(tx, name, args, kwargs)
 
 
 class DictValuesVariable(DictViewVariable):
+    # PyDictValues_Type: https://github.com/python/cpython/blob/v3.13.0/Objects/dictobject.c#L6567
+    _cpython_type = dict_values
+
     # DictValuesVariable is an iterable but cannot be compared.
     kv = "values"
 
@@ -1795,8 +1265,18 @@ def debug_repr(self) -> str:
 
 
 class DictItemsVariable(DictViewVariable):
+    # PyDictItems_Type: https://github.com/python/cpython/blob/v3.13.0/Objects/dictobject.c#L6477
+    _cpython_type = dict_items
+
     kv = "items"
 
+    @property
+    def set_items(self) -> set["HashableTracker"]:
+        return {
+            HashableTracker(variables.TupleVariable([k.vt, v]))
+            for k, v in self.view_items
+        }
+
     @property
     def view_items_vt(self) -> list[VariableTracker]:
         # Returns an iterable of the unpacked items
@@ -1832,6 +1312,19 @@ def call_method(
                 raise_args_mismatch(tx, name, "1 args", f"{len(args)} args")
             if isinstance(args[0], DictItemsVariable):
                 return self.dv_dict.call_method(tx, "__eq__", [args[0].dv_dict], {})
+            elif isinstance(
+                args[0],
+                (
+                    SetVariable,
+                    variables.UserDefinedSetVariable,
+                    DictItemsVariable,
+                    DictKeysVariable,
+                ),
+            ):
+                return VariableTracker.build(
+                    tx,
+                    len(self.set_items ^ args[0].set_items) == 0,
+                )
             return ConstantVariable.create(False)
         elif name == "__iter__":
             from .lists import ListIteratorVariable
@@ -1839,6 +1332,35 @@ def call_method(
             return ListIteratorVariable(
                 self.view_items_vt, mutation_type=ValueMutationNew()
             )
+        elif name in (
+            "__and__",
+            "__iand__",
+            "__or__",
+            "__ior__",
+            "__sub__",
+            "__isub__",
+            "__xor__",
+            "__ixor__",
+        ):
+            # These methods always returns a set
+            fn_hdl = getattr(self.set_items, name)
+            ret_val = fn_hdl(args[0].set_items)  # type: ignore[attr-defined]
+            return SetVariable(ret_val)
+        elif name in cmp_name_to_op_mapping:
+            if not isinstance(
+                args[0],
+                (
+                    SetVariable,
+                    variables.UserDefinedSetVariable,
+                    DictItemsVariable,
+                    DictKeysVariable,
+                ),
+            ):
+                return VariableTracker.build(tx, NotImplemented)
+            return VariableTracker.build(
+                tx,
+                cmp_name_to_op_mapping[name](self.set_items, args[0].set_items),  # type: ignore[attr-defined]
+            )
         return super().call_method(tx, name, args, kwargs)
 
     def is_python_hashable(self) -> Literal[False]:
@@ -1846,3 +1368,168 @@ def is_python_hashable(self) -> Literal[False]:
         Dictionary item views are not hashable in Python.
         """
         return False
+
+
+kV = HashableTracker | str
+
+
+class SideEffectsProxyDict(collections.abc.MutableMapping[kV, VariableTracker]):
+    """
+    A proxy dict that allows us to track mutations to the dict using side
+    effects table as storage.
+    """
+
+    @staticmethod
+    def get_example_value_dict(vt: VariableTracker) -> dict[str, object]:
+        if istype(vt, variables.NestedUserFunctionVariable):
+            # NestedUserFunctionVariable is created with MAKE_FUNCTION and its
+            # __dict__ starts empty. Any mutation will actually be recorded in
+            # the side effects table.
+            return {}
+        elif isinstance(vt, variables.LocalGeneratorFunctionVariable):
+            return SideEffectsProxyDict.get_example_value_dict(vt.vt)
+        else:
+            value = vt.get_real_python_backed_value()
+            if value is not NO_SUCH_SUBOBJ:
+                if isinstance(vt, variables.UserDefinedObjectVariable):
+                    return vt._getattr_static("__dict__")  # type: ignore[bad-return]
+                else:
+                    return object.__getattribute__(value, "__dict__")
+            else:
+                unimplemented(
+                    gb_type="unsupported variable type for __dict__ access",
+                    context=f"VariableTracker type: {type(vt)}",
+                    explanation=f"Dynamo does not know how to get __dict__ from {type(vt)}",
+                    hints=[
+                        *graph_break_hints.DYNAMO_BUG,
+                    ],
+                )
+
+    @staticmethod
+    def get_value___dict__(
+        tx: "InstructionTranslator", vt: VariableTracker
+    ) -> dict[str, VariableTracker]:
+        example_value_dict = SideEffectsProxyDict.get_example_value_dict(vt)
+
+        return {
+            key: VariableTracker.build(
+                tx,
+                value,
+                source=vt.source
+                and DictGetItemSource(AttrSource(vt.source, "__dict__"), key),
+            )
+            for key, value in example_value_dict.items()
+        }
+
+    def __init__(self, item: VariableTracker, tx: "InstructionTranslator") -> None:
+        self.item = item
+        self.side_effects = tx.output.side_effects
+        self.item_dict = self.get_value___dict__(tx, item)
+
+    def _maybe_unwrap_key(self, key: kV) -> str:
+        Hasher = HashableTracker
+        return key.vt.as_python_constant() if istype(key, Hasher) else key
+
+    def side_effects_table(self) -> dict[str, VariableTracker]:
+        return self.side_effects.store_attr_mutations.get(self.item, {})
+
+    def __getitem__(self, key: kV) -> VariableTracker:
+        name = self._maybe_unwrap_key(key)
+        if self.side_effects.has_pending_mutation_of_attr(self.item, name):
+            return self.side_effects.load_attr(self.item, name, deleted_ok=True)
+        return self.item_dict[name]
+
+    def __setitem__(self, key: kV, value: VariableTracker) -> None:
+        # Find a way to not hash the key using HashableTracker
+        name = self._maybe_unwrap_key(key)
+        assert istype(name, str)
+        self.side_effects.store_attr(self.item, name, value)
+
+    def __delitem__(self, key: kV) -> None:
+        name = self._maybe_unwrap_key(key)
+        self.side_effects.store_attr(self.item, name, variables.DeletedVariable())
+
+    def __contains__(self, key: kV) -> bool:  # type: ignore[bad-override]
+        name = self._maybe_unwrap_key(key)
+        table = self.side_effects_table()
+        # if name in side effects, then it is only contained if it's not a DeletedVariable
+        # even if the original dict contains it
+        if name in table:
+            return not isinstance(table[name], variables.DeletedVariable)
+        else:
+            return name in self.item_dict
+
+    def __len__(self) -> int:
+        return sum(1 for _ in self)
+
+    def __iter__(self) -> Iterator[HashableTracker]:
+        Hasher = HashableTracker
+        d = self.side_effects_table()
+        for k, v in d.items():
+            if isinstance(v, variables.DeletedVariable):
+                continue
+            yield Hasher(ConstantVariable.create(k))
+
+        for k, v in self.item_dict.items():
+            if k not in d:
+                yield Hasher(ConstantVariable.create(k))
+
+
+class DunderDictVariable(ConstDictVariable):
+    """represents object.__dict__"""
+
+    @classmethod
+    def create(
+        cls,
+        tx: "InstructionTranslator",
+        vt: VariableTracker,
+    ) -> "DunderDictVariable":
+        mutation = AttributeMutationExisting() if vt.source else AttributeMutationNew()
+        source = vt.source and AttrSource(vt.source, "__dict__")
+
+        return cls(
+            vt,
+            tx=tx,
+            mutation_type=mutation,
+            source=source,
+        )
+
+    def __init__(
+        self,
+        vt: VariableTracker,
+        tx: "InstructionTranslator",
+        **kwargs: Any,
+    ) -> None:
+        super().__init__({}, **kwargs)
+        self.items = SideEffectsProxyDict(vt, tx)
+
+    def setitem(self, name: str, value: VariableTracker) -> None:
+        self.items[name] = value
+
+    def getitem(self, name: str) -> VariableTracker:
+        return self.items[name]
+
+    def contains(self, name: str) -> bool:
+        return name in self.items
+
+    def getitem_or_default(
+        self,
+        name: str,
+        default: Callable[[], VariableTracker],
+    ) -> VariableTracker:
+        if self.contains(name):
+            return self.getitem(name)
+        else:
+            value = default()
+            self.items[name] = value
+            return value
+
+    # Mutations to __dict__ are tracked through side effects (SideEffectsProxyDict),
+    # so we don't need to install guards. Guard installation is overridden to no-op.
+    def install_dict_keys_match_guard(self) -> None:
+        pass
+
+    def install_dict_contains_guard(
+        self, tx: "InstructionTranslator", args: list[VariableTracker]
+    ) -> None:
+        pass
diff --git a/torch/_dynamo/variables/distributed.py b/torch/_dynamo/variables/distributed.py
index 09caddeb7f5e8..9714a9960ed0d 100644
--- a/torch/_dynamo/variables/distributed.py
+++ b/torch/_dynamo/variables/distributed.py
@@ -20,26 +20,21 @@
 
 import functools
 import inspect
-from collections.abc import Sequence
 from typing import Any, Literal, TYPE_CHECKING
 
 import torch
 from torch.fx.experimental._backward_state import BackwardState
 
-from .. import compiled_autograd, variables
+from .. import compiled_autograd
 from .._trace_wrapped_higher_order_op import trace_wrapped
-from ..bytecode_transformation import create_call_function
 from ..exc import unimplemented
 from ..external_utils import call_module_hooks_from_backward_state
 from ..guards import GuardBuilder, install_guard
 from ..source import AttrSource
-from ..utils import istype
 from .base import VariableTracker
-from .constant import CONSTANT_VARIABLE_NONE, ConstantVariable, EnumVariable
 
 
 if TYPE_CHECKING:
-    from torch._dynamo.codegen import PyCodegen
     from torch._dynamo.symbolic_convert import InstructionTranslator
 
 
@@ -134,279 +129,24 @@ def is_group_member_type(cls, value: object) -> bool:
 
         return type(value) is _WorldMeta
 
+    def python_type(self) -> type:
+        return type(self.value)
+
     def var_getattr(self, tx: "InstructionTranslator", name: str) -> VariableTracker:
         if name == "WORLD":
+            from .builder import SourcelessBuilder
+
             assert self.source
             source = AttrSource(base=self.source, member="WORLD")
             install_guard(source.make_guard(GuardBuilder.ID_MATCH))
-            return ProcessGroupVariable(self.value.WORLD)
+            return SourcelessBuilder.create(tx, self.value.WORLD)
         elif name == "NON_GROUP_MEMBER":
             assert self.source
             source = AttrSource(base=self.source, member="NON_GROUP_MEMBER")
             install_guard(source.make_guard(GuardBuilder.ID_MATCH))
-            return EnumVariable(self.value.NON_GROUP_MEMBER)
-        return super().var_getattr(tx, name)
-
-
-class PlacementClassVariable(DistributedVariable):
-    @staticmethod
-    def is_placement_type(value: object) -> bool:
-        # we can't rely on importing/accessing torch distributed, it is not always built.
-        if not DistributedVariable.is_available():
-            return False
-
-        from torch.distributed.tensor.placement_types import Placement
-
-        return isinstance(value, type) and issubclass(value, Placement)
-
-    def as_python_constant(self) -> Any:
-        return self.value
-
-    def call_function(
-        self,
-        tx: "InstructionTranslator",
-        args: Sequence[VariableTracker],
-        kwargs: dict[str, VariableTracker],
-    ) -> VariableTracker:
-        if self.source:
-            # NOTE: we don't need to track mutations to the placement class as they
-            # are supposed to be immutable.
-            new_obj = self.value.__new__(self.value)
-            var = PlacementVariable(new_obj)
-            if inspect.getattr_static(self.value, "__init__", None):
-                var.call_method(tx, "__init__", args, kwargs)
-                return var
-
-        return super().call_function(tx, args, kwargs)
-
-
-class PlacementVariable(DistributedVariable):
-    @staticmethod
-    def is_placement(value: object) -> bool:
-        # we can't rely on importing/accessing torch distributed, it is not always built.
-        if not DistributedVariable.is_available():
-            return False
-        from torch.distributed.tensor.placement_types import Placement
-
-        return isinstance(value, Placement)
-
-    def as_python_constant(self) -> Any:
-        return self.value
-
-    def var_getattr(self, tx: "InstructionTranslator", name: str) -> VariableTracker:
-        if name == "dim":
-            return ConstantVariable.create(self.value.dim)
-        return super().var_getattr(tx, name)
-
-    def call_method(
-        self,
-        tx: "InstructionTranslator",
-        name: str,
-        args: Sequence[VariableTracker],
-        kwargs: dict[str, VariableTracker],
-    ) -> VariableTracker:
-        from . import ConstantVariable
-
-        # Placement types dynamo tracking only allows following methods
-        # and __setattr__  is for case like `Shard(dim)` and methods.
-        # Methods in the list must satisfy:
-        #    1. Input arguments are constants and do not need to be guarded on;
-        #    2. Output is constant with respect to their inputs
-        constant_fold_functions = [
-            "__init__",
-            "__setattr__",
-            "is_shard",
-            "is_partial",
-            "is_replicate",
-        ]
-
-        if name in constant_fold_functions:
-            try:
-                value_type = type(self.value)
-                if inspect.getattr_static(value_type, "__getattr__", None) is not None:
-                    unimplemented(
-                        gb_type="Placement with custom __getattr__ not supported",
-                        context=f"{value_type.__name__} with custom __getattr__",
-                        explanation="Dynamo does not support Placement types with custom __getattr__ methods",
-                        hints=[
-                            "Use Placement types without custom __getattr__ methods",
-                            "Move the Placement usage outside the compiled region",
-                        ],
-                    )
-                method = inspect.getattr_static(value_type, name)
-            except AttributeError:
-                method = None
-            if method is object.__init__:
-                return CONSTANT_VARIABLE_NONE
-
-            args = [x.as_python_constant() for x in args]
-            kwargs = {k: v.as_python_constant() for k, v in kwargs.items()}
-            assert method is not None
-            if name == "__setattr__":
-                method(self.value, *args, **kwargs)
-                return self
-            constant_val = method(self.value, *args, **kwargs)
-            return ConstantVariable.create(constant_val)
-
-        return super().call_method(tx, name, args, kwargs)  # type: ignore[arg-type]
-
-    def reconstruct(self, codegen: "PyCodegen") -> None:
-        # Reconstruct the Placement object by calling its constructor
-        # e.g., Shard(0), Replicate(), Partial()
-        from torch.distributed.tensor.placement_types import Partial, Replicate, Shard
-
-        placement_type = type(self.value)
-
-        # Load the placement class
-        codegen.add_push_null(
-            lambda: codegen.load_import_from(
-                "torch.distributed.tensor.placement_types", placement_type.__name__
-            )
-        )
-
-        # For Shard, we need to pass the dim argument
-        if isinstance(self.value, Shard):
-            codegen(ConstantVariable.create(self.value.dim))
-            codegen.extend_output(create_call_function(1, False))
-        # Replicate and Partial have no required args
-        elif istype(self.value, (Replicate, Partial)):
-            codegen.extend_output(create_call_function(0, False))
-        else:
-            super().reconstruct(codegen)
-
-
-class DeviceMeshVariable(DistributedVariable):
-    @staticmethod
-    def is_device_mesh(value: object) -> bool:
-        # we can't rely on importing/accessing torch distributed, it is not always built.
-        if not DistributedVariable.is_available():
-            return False
-
-        from torch.distributed.device_mesh import DeviceMesh
-
-        return istype(value, DeviceMesh)
-
-    def as_python_constant(self) -> Any:
-        return self.value
-
-    def var_getattr(self, tx: "InstructionTranslator", name: str) -> VariableTracker:
-        if name == "ndim":
-            return ConstantVariable.create(self.value.ndim)
-        if name == "device_type":
-            return ConstantVariable.create(self.value.device_type)
-        if name == "mesh_dim_names":
-            source = self.source
-            if source:
-                source = AttrSource(base=source, member="mesh_dim_names")
-            return VariableTracker.build(tx, self.value.mesh_dim_names, source)
-        return super().var_getattr(tx, name)
-
-    def call_method(
-        self,
-        tx: "InstructionTranslator",
-        name: str,
-        args: list[VariableTracker],
-        kwargs: dict[str, VariableTracker],
-    ) -> VariableTracker:
-        if name == "size":
-            const_args = [x.as_python_constant() for x in args]
-            const_kwargs = {k: v.as_python_constant() for k, v in kwargs.items()}
-            return ConstantVariable.create(self.value.size(*const_args, **const_kwargs))
-        if name == "get_coordinate":
-            return ConstantVariable.create(self.value.get_coordinate())
-        if name == "get_rank":
-            return ConstantVariable.create(self.value.get_rank())
-        if name == "get_local_rank":
-            const_args = [x.as_python_constant() for x in args]
-            const_kwargs = {k: v.as_python_constant() for k, v in kwargs.items()}
-            return ConstantVariable.create(
-                self.value.get_local_rank(*const_args, **const_kwargs)
-            )
-        if name == "get_group":
-            const_args = [x.as_python_constant() for x in args]
-            const_kwargs = {k: v.as_python_constant() for k, v in kwargs.items()}
-            return ProcessGroupVariable(
-                self.value.get_group(*const_args, **const_kwargs)
-            )
-        if name == "_is_current_rank_part_of_mesh":
-            return ConstantVariable.create(self.value._is_current_rank_part_of_mesh())
-        if name == "_get_or_create_default_group":
-            return ProcessGroupVariable(self.value._get_or_create_default_group())
-        if name == "_flatten":
-            from .builder import SourcelessBuilder
-
-            const_args = [x.as_python_constant() for x in args]
-            const_kwargs = {k: v.as_python_constant() for k, v in kwargs.items()}
-            return SourcelessBuilder.create(
-                tx, self.value._flatten(*const_args, **const_kwargs)
-            )
-        if name == "_sym_get_coordinate":
-            const_args = [x.as_python_constant() for x in args]
-            const_kwargs = {k: v.as_python_constant() for k, v in kwargs.items()}
-            return ConstantVariable.create(
-                self.value._sym_get_coordinate(*const_args, **const_kwargs)
-            )
-        return super().call_method(tx, name, args, kwargs)
-
-
-class ProcessGroupVariable(DistributedVariable):
-    """
-    We don't want a ProcessGroup object to end up in our output graph.
-
-    But it's common for dynamo to intercept a PG that is then used to get info like
-    rank() or world_size(), as well as passed to utility functions in distributed_c10d
-    which desugar it into plain types like a ranklist and tag.
-
-    For convenience and proper guarding, we construct a variable type.
-
-    TODO: make it possible to use ProcessGroupVariable as input to simple functions
-          like _expand_group without dynamo complaining about making a proxy for it.
-          It is not a tensor-like type, and we don't want a proxy- but dynamo assumes
-          torch library functions are dealing with tensor-like types and would have proxies
-          for their args.
-    TODO: should we make this inherit VT instead of UDOV? Do we want any of the default behaviors
-          or just graph-break whenever one of our special cases is not hit?
-    """
-
-    def as_python_constant(self) -> Any:
-        return self.value
-
-    def call_method(
-        self,
-        tx: "InstructionTranslator",
-        name: str,
-        args: list[VariableTracker],
-        kwargs: dict[str, VariableTracker],
-    ) -> VariableTracker:
-        if name == "rank":
-            return variables.ConstantVariable.create(self.value.rank())
-        if name == "size":
-            return variables.ConstantVariable.create(self.value.size())
-        if name == "_get_backend_name":
-            return variables.ConstantVariable.create(self.value._get_backend_name())
-
-        return super().call_method(tx, name, args, kwargs)
-
-    def var_getattr(self, tx: "InstructionTranslator", name: str) -> VariableTracker:
-        if name == "group_name":
-            return variables.ConstantVariable.create(self.value.group_name)
-        if name in ["rank", "size"]:
-            return variables.LambdaVariable(
-                lambda *args, **kwargs: self.call_method(tx, name, args, kwargs)
-            )
-        # TODO should this just raise unimplemented?
+            return VariableTracker.build(tx, self.value.NON_GROUP_MEMBER)
         return super().var_getattr(tx, name)
 
-    @staticmethod
-    def is_process_group(value: object) -> bool:
-        # we can't rely on importing/accessing torch distributed, it is not always built.
-        if not DistributedVariable.is_available():
-            return False
-        from torch._C._distributed_c10d import ProcessGroup
-        from torch.testing._internal.distributed.fake_pg import FakeProcessGroup
-
-        return istype(value, (ProcessGroup, FakeProcessGroup))
-
 
 class BackwardHookVariable(VariableTracker):
     """
diff --git a/torch/_dynamo/variables/functions.py b/torch/_dynamo/variables/functions.py
index 3e6cea61a4280..9f17481accb5e 100644
--- a/torch/_dynamo/variables/functions.py
+++ b/torch/_dynamo/variables/functions.py
@@ -53,6 +53,7 @@
     ObservedGeneratorExit,
     ObservedUserStopIteration,
     raise_observed_exception,
+    raise_type_error,
     StepUnsupported,
     unimplemented,
     Unsupported,
@@ -82,11 +83,11 @@
 from .base import (
     AsPythonConstantNotImplementedError,
     AttributeMutationNew,
-    raise_type_error_exc,
     ValueMutationNew,
     VariableTracker,
 )
-from .constant import CONSTANT_VARIABLE_NONE, ConstantVariable
+from .constant import ConstantVariable
+from .user_defined import UserDefinedObjectVariable
 
 
 try:
@@ -109,6 +110,7 @@
         TritonKernelType,
     )
 
+    from .dicts import DunderDictVariable
     from .lists import BaseListVariable, ListVariable
     from .tensor import TensorVariable
 
@@ -335,12 +337,14 @@ def _create_nested_fn(
 
 fn_known_dunder_attrs = {
     "__annotations__",
-    "__defaults__",
-    "__kwdefaults__",
-    "__code__",
-    "__globals__",
     "__closure__",
+    "__code__",
+    "__defaults__",
     "__doc__",
+    "__globals__",
+    "__kwdefaults__",
+    "__name__",
+    "__module__",
 }
 
 
@@ -364,6 +368,8 @@ def fn_var_getattr(
         raise_observed_exception(AttributeError, tx)
 
     # Special handling for known dunder attributes
+    # TODO(guilhermeleobas): this check should go through fn.__dict__ first as
+    # functools.partial can override it
     if name in fn_known_dunder_attrs:
         subobj = getattr(fn, name)
     if source:
@@ -372,12 +378,51 @@ def fn_var_getattr(
 
 
 class BaseUserFunctionVariable(VariableTracker):
+    def __init__(
+        self, dict_vt: "DunderDictVariable | None" = None, **kwargs: Any
+    ) -> None:
+        super().__init__(**kwargs)
+        self.dict_vt: DunderDictVariable | None = dict_vt
+
+    def get_source(self) -> Source | None:
+        return self.source
+
+    def get_dict_vt(self, tx: "InstructionTranslator") -> "DunderDictVariable":
+        if self.dict_vt is None:
+            self.dict_vt = variables.DunderDictVariable.create(tx, self)
+        return self.dict_vt
+
+    def call_method(
+        self,
+        tx: "InstructionTranslator",
+        name: str,
+        args: Sequence[VariableTracker],
+        kwargs: dict[str, VariableTracker],
+    ) -> VariableTracker:
+        if name == "__setattr__":
+            return self.get_dict_vt(tx).call_method(
+                tx, "__setitem__", list(args), kwargs
+            )
+        elif name == "__delattr__":
+            return self.get_dict_vt(tx).call_method(tx, "__delitem__", list(args), {})
+        return super().call_method(tx, name, list(args), kwargs)
+
     def get_filename(self) -> str:
         return self.get_code().co_filename
 
     def get_name(self) -> str:
         return self.get_code().co_name
 
+    def get_qualname(self) -> str:
+        if sys.version_info >= (3, 11):
+            return self.get_code().co_qualname
+        else:
+            return self.get_name()
+
+    def get_doc(self) -> str | None:
+        # stored in code.co_consts[0]
+        return self.get_code().co_consts[0]
+
     def get_globals(self) -> dict[str, Any]:
         raise NotImplementedError
 
@@ -387,6 +432,47 @@ def get_code(self) -> types.CodeType:
     def has_self(self) -> bool:
         raise NotImplementedError
 
+    def get_function(self) -> types.FunctionType:
+        raise NotImplementedError
+
+    def get_module(self) -> str:
+        return self.get_globals()["__name__"]
+
+    def var_getattr(self, tx: "InstructionTranslator", name: str):
+        fn_dict = self.get_dict_vt(tx)
+
+        # missing: __globals__, __closure__, __kwdefautls__, __defaults__
+        if name in ("__name__", "__qualname__", "__doc__", "__module__", "__code__"):
+            val = getattr(self, f"get_{name[2:-2]}")()
+            if fn_dict.contains(name):
+                return fn_dict.getitem(name)
+            return ConstantVariable.create(
+                val, source=self.source and AttrSource(self.source, name)
+            )
+        elif name == "__dict__":
+            return fn_dict
+        elif name == "__annotations__":
+            return fn_dict.getitem_or_default(
+                name,
+                lambda: variables.ConstDictVariable(
+                    {},
+                    mutation_type=ValueMutationNew(),
+                ),
+            )
+        elif name == "__type_params__":
+            return fn_dict.getitem_or_default(
+                name,
+                lambda: variables.TupleVariable(
+                    [],
+                    mutation_type=ValueMutationNew(),
+                ),
+            )
+        else:
+            if fn_dict.contains(name):
+                return fn_dict.getitem(name)
+            else:
+                raise_observed_exception(AttributeError, tx)
+
     def call_function(
         self,
         tx: "InstructionTranslator",
@@ -399,7 +485,7 @@ def call_function(
             self.get_name() == "patch_track_step_called"
             and self.get_filename().endswith("torch/optim/lr_scheduler.py")
         ):
-            return CONSTANT_VARIABLE_NONE
+            return ConstantVariable.create(None)
         return tx.inline_user_function_return(self, [*self.self_args(), *args], kwargs)  # type: ignore[attr-defined]
 
     def call_obj_hasattr(
@@ -407,12 +493,14 @@ def call_obj_hasattr(
     ) -> ConstantVariable:
         result = False
 
-        try:
-            result = hasattr(self.get_function(), name)  # type: ignore[attr-defined]
-        except NotImplementedError:
-            if name == "__name__" and isinstance(self, NestedUserFunctionVariable):
-                result = True
-        return variables.ConstantVariable.create(result)
+        if name in fn_known_dunder_attrs or name == "__dict__":
+            result = True
+        else:
+            try:
+                result = hasattr(self.get_function(), name)  # type: ignore[attr-defined]
+            except NotImplementedError:
+                result = False
+        return VariableTracker.build(tx, result)
 
     def closure_vars(self, tx: "InstructionTranslator") -> dict[str, VariableTracker]:
         return {}
@@ -427,6 +515,9 @@ def should_allow_nested_graph_breaks(self) -> bool:
 class UserFunctionVariable(BaseUserFunctionVariable):
     """Some unsupported user-defined global function"""
 
+    # PyFunction_Type: https://github.com/python/cpython/blob/v3.13.0/Objects/funcobject.c#L1046
+    _cpython_type = types.FunctionType
+
     _nonvar_fields = {
         "fn",
         "is_constant",
@@ -484,6 +575,11 @@ def as_python_constant(self) -> Any:
         # subclasses (such as methods) usually aren't a constant
         return super().as_python_constant()
 
+    def get_real_python_backed_value(self) -> Any:
+        if istype(self, UserFunctionVariable):
+            return self.fn
+        return super().get_real_python_backed_value()
+
     def self_args(self) -> list[VariableTracker]:
         return []
 
@@ -572,8 +668,12 @@ def bind_args(
         return result
 
     def var_getattr(self, tx: "InstructionTranslator", name: str) -> VariableTracker:
-        if name in cmp_name_to_op_mapping:
-            return variables.GetAttrVariable(self, name)
+        if name == "__dict__":
+            return super().var_getattr(tx, name)
+        elif name in cmp_name_to_op_mapping:
+            return variables.GetAttrVariable(
+                self, name, py_type=type(getattr(self.fn, name))
+            )
         source = self.get_source()
         return fn_var_getattr(tx, self.fn, source, name)
 
@@ -581,7 +681,7 @@ def call_obj_hasattr(
         self, tx: "InstructionTranslator", name: str
     ) -> ConstantVariable:
         result = hasattr(self.fn, name)
-        return variables.ConstantVariable.create(result)
+        return VariableTracker.build(tx, result)
 
     def call_function(
         self,
@@ -619,6 +719,25 @@ def call_function(
                     "Improper error_on_graph_break() call. Please fix your call to error_on_graph_break(). "
                     f"args: {args}, kwargs: {kwargs}"
                 ) from e
+        elif self.fn is torch._dynamo.override_cudagraphs:
+            try:
+                bound = inspect.signature(self.fn).bind(*args, **kwargs)
+                bound.apply_defaults()
+                fwd = bound.arguments["fwd"]
+                bwd = bound.arguments["bwd"]
+                if isinstance(fwd, VariableTracker):
+                    fwd = fwd.as_python_constant()
+                if isinstance(bwd, VariableTracker):
+                    bwd = bwd.as_python_constant()
+                return variables.CudagraphOverrideVariable(fwd, bwd)
+            except Exception as e:
+                raise RuntimeError(
+                    "Improper override_cudagraphs() call. Please fix your call to override_cudagraphs(). "
+                    f"args: {args}, kwargs: {kwargs}"
+                ) from e
+        elif self.fn is torch._dynamo.bytecode_debugger.breakpoint:
+            tx.output._emit_debugger_breakpoint = True
+            return variables.ConstantVariable.create(None)
         # Handle a `nonstrict_trace(fn)` call
         elif self.fn is torch._dynamo.nonstrict_trace:
             bound = inspect.signature(self.fn).bind(*args, **kwargs)
@@ -704,24 +823,43 @@ def _maybe_call_tree_map_fastpath(
             tree_map_args = args
             tree_map_kwargs = kwargs
 
-        if not (
+        is_tree_map = (
             isinstance(tree_map_fn, UserFunctionVariable)
             and tree_map_fn._is_tree_map_function()
-            and not ({*tree_map_kwargs} - _SUPPORTED_TREE_MAP_KWARGS)
-            and len(tree_map_args) >= 2
-        ):
+        )
+        is_tree_map_with_path = (
+            isinstance(tree_map_fn, UserFunctionVariable)
+            and tree_map_fn._is_tree_map_with_path_function()
+        )
+
+        if not (is_tree_map or is_tree_map_with_path):
+            return None
+        if {*tree_map_kwargs} - _SUPPORTED_TREE_MAP_KWARGS:
+            return None
+        if len(tree_map_args) < 2:
             return None
 
         map_fn = tree_map_args[0]
         first_tree = tree_map_args[1]
         rest = tree_map_args[2:]
-        return first_tree.call_tree_map(
-            tx,
-            tree_map_fn,
-            map_fn,
-            rest,
-            tree_map_kwargs,
-        )
+
+        if is_tree_map_with_path:
+            return first_tree.call_tree_map_with_path(
+                tx,
+                tree_map_fn,
+                map_fn,
+                rest,
+                tree_map_kwargs,
+                keypath=(),
+            )
+        else:
+            return first_tree.call_tree_map(
+                tx,
+                tree_map_fn,
+                map_fn,
+                rest,
+                tree_map_kwargs,
+            )
 
     def _is_tree_map_function(self) -> bool:
         return (
@@ -729,6 +867,12 @@ def _is_tree_map_function(self) -> bool:
             and getattr(self.fn, "__module__", None) in self._TREE_MAP_MODULES
         )
 
+    def _is_tree_map_with_path_function(self) -> bool:
+        return (
+            getattr(self.fn, "__name__", None) == "tree_map_with_path"
+            and getattr(self.fn, "__module__", None) in self._TREE_MAP_MODULES
+        )
+
     def _is_tree_map_only_function(self) -> bool:
         return (
             getattr(self.fn, "__name__", None) == "tree_map_only"
@@ -770,7 +914,7 @@ def _rewrite_tree_map_only_call(
             map_fn,
             source=getattr(map_fn, "source", None),
         )
-        tree_map_variable = variables.UserFunctionVariable(tree_map_callable)
+        tree_map_variable = VariableTracker.build(tx, tree_map_callable)
         return tree_map_variable, [wrapped_map_fn, tree_arg], dict(kwargs)
 
     def _lookup_tree_map_function(self) -> types.FunctionType | None:
@@ -922,6 +1066,9 @@ def __init__(
         assert isinstance(fn, types.BuiltinMethodType)
         self.fn = fn
 
+    def python_type(self) -> type:
+        return types.BuiltinMethodType
+
     @staticmethod
     def is_supported_builtin_method(obj: Any) -> bool:
         method_self = obj.__self__
@@ -947,6 +1094,9 @@ def call_function(
 
 
 class LocalGeneratorObjectVariable(VariableTracker):
+    # PyGen_Type: https://github.com/python/cpython/blob/v3.13.0/Objects/genobject.c#L814
+    _cpython_type = types.GeneratorType
+
     def __init__(
         self,
         code: types.CodeType,
@@ -1131,13 +1281,13 @@ def call_method(
             tracer = self.inline_tracer
             if self._is_generator_just_started() or self._is_generator_exhausted():
                 tracer.generator_exhausted = True
-                return variables.CONSTANT_VARIABLE_NONE
+                return variables.ConstantVariable.create(None)
 
             # Raise GeneratorExit to see if user code catches it. Any other exception
             # is propagated to the parent frame.
             try:
                 self._setup_exception(
-                    tx, variables.ExceptionVariable(GeneratorExit, ())
+                    tx, variables.ExceptionVariable(GeneratorExit, [])
                 )
                 # There's an extra block on Python 3.12+ to handle StopIteration
                 # see: https://github.com/python/cpython/blob/8f93dd8a8f237b277abad20d566df90c5cbd7f1e/Objects/genobject.c#L394-L397
@@ -1160,11 +1310,11 @@ def call_method(
                     and tracer.next_instruction.opname == "CALL_INTRINSIC_1"
                 ):
                     tracer.generator_exhausted = True
-                    return variables.CONSTANT_VARIABLE_NONE
+                    return variables.ConstantVariable.create(None)
             except ObservedGeneratorExit:
                 # If it doesn't catch, we just return None, as per the text above
                 tracer.generator_exhausted = True
-                return variables.CONSTANT_VARIABLE_NONE
+                return variables.ConstantVariable.create(None)
 
             try:
                 # Raise RuntimeError if the generator yields any other value
@@ -1172,7 +1322,7 @@ def call_method(
                     raise_observed_exception(RuntimeError, tx)
             except ObservedGeneratorExit:
                 tracer.generator_exhausted = True
-                return variables.CONSTANT_VARIABLE_NONE
+                return variables.ConstantVariable.create(None)
             except ObservedUserStopIteration:
                 # In Python 3.13+, one can capture GeneratorExit and return a value
                 # See test_generator.py::test_close_capture_GeneratorExit_return
@@ -1264,7 +1414,7 @@ def call_method(
             exc_type = type("__InternalThrowException", (Exception,), {})
 
             try:
-                self._setup_exception(tx, variables.ExceptionVariable(exc_type, ()))
+                self._setup_exception(tx, variables.ExceptionVariable(exc_type, []))
                 self.next_variable(tx)
             except get_dynamo_observed_exception(exc_type):
                 # We should get back the exception raised before.
@@ -1297,6 +1447,9 @@ class LocalGeneratorFunctionVariable(BaseUserFunctionVariable):
         This is a wrapper around (Nested)UserFunctionVariable
     """
 
+    def python_type(self) -> type:
+        return types.FunctionType
+
     def __init__(
         self,
         vt: BaseUserFunctionVariable,
@@ -1368,6 +1521,9 @@ def call_function(
             source=self.source,
         )
 
+    def get_real_python_backed_value(self) -> object:
+        return self.vt.get_real_python_backed_value()
+
 
 class FunctionDecoratedByContextlibContextManagerVariable(
     LocalGeneratorFunctionVariable
@@ -1407,6 +1563,9 @@ def _build_inline_tracer(
 class UserMethodVariable(UserFunctionVariable):
     """Some unsupported user-defined method"""
 
+    # PyMethod_Type: https://github.com/python/cpython/blob/v3.13.0/Objects/classobject.c#L332
+    _cpython_type = types.MethodType
+
     def __init__(
         self,
         fn: Callable[..., Any],
@@ -1519,6 +1678,9 @@ def var_getattr(self, tx: "InstructionTranslator", name: str) -> VariableTracker
             return VariableTracker.build(tx, self.fn, self.source_fn)  # type: ignore[arg-type]
         return super().var_getattr(tx, name)
 
+    def get_real_python_backed_value(self) -> Any:
+        return self.fn
+
 
 class WrappedUserMethodVariable(UserMethodVariable):
     def __init__(
@@ -1539,6 +1701,11 @@ def call_function(
         args: Sequence[VariableTracker],
         kwargs: dict[str, VariableTracker],
     ) -> VariableTracker:
+        if config.nested_graph_breaks:
+            wrapper_fn = UserFunctionVariable(polyfills._fn_with_ctx)
+            return wrapper_fn.call_function(
+                tx, [self.context, self.wrapped] + list(args), kwargs
+            )
         self.context.enter(tx)
         result = super().call_function(tx, args, kwargs)
         self.context.exit(tx)
@@ -1568,6 +1735,11 @@ def call_function(
         args: Sequence[VariableTracker],
         kwargs: dict[str, VariableTracker],
     ) -> VariableTracker:
+        if config.nested_graph_breaks:
+            wrapper_fn = UserFunctionVariable(polyfills._fn_with_ctx)
+            return wrapper_fn.call_function(
+                tx, [self.context, self.wrapped] + list(args), kwargs
+            )
         self.context.enter(tx)
         result = super().call_function(tx, args, kwargs)
         self.context.exit(tx)
@@ -1589,7 +1761,24 @@ def invoke_and_store_as_constant(
     def convert(x: VariableTracker) -> Any:
         if x.is_tensor():
             return cast("TensorVariable", x).get_real_value()
-        return x.as_python_constant()
+        if isinstance(x, UserDefinedObjectVariable):
+            if x.source is not None:
+                install_guard(x.make_guard(GuardBuilder.ID_MATCH))
+            return x.value
+        try:
+            return x.as_python_constant()
+        except AsPythonConstantNotImplementedError:
+            unimplemented(
+                gb_type="assume_constant_result argument conversion failed",
+                context=f"function {name}, variable type {type(x).__name__}",
+                explanation=f"Cannot convert argument of type {type(x).__name__} to a Python constant "
+                f"for function {name} marked with torch._dynamo.assume_constant_result. "
+                f"The variable tracker does not support constant conversion.",
+                hints=[
+                    "Remove torch._dynamo.assume_constant_result from this function",
+                    "Ensure all arguments passed to the function can be converted to constants",
+                ],
+            )
 
     args = [convert(x) for x in args]
     kwargs = {k: convert(v) for k, v in kwargs.items()}
@@ -1614,7 +1803,6 @@ def __init__(
         f_globals: dict[str, Any],
         defaults: VariableTracker | None,
         kwdefaults: VariableTracker | None,
-        annotations: VariableTracker | None,
         closure: VariableTracker | None,
         # This is present when this function is created by
         # `functools.wrap(wrapped_fn)(this_fn)`.
@@ -1632,7 +1820,6 @@ def __init__(
         self.f_globals = f_globals
         self.defaults = defaults
         self.kwdefaults = kwdefaults
-        self.annotations = annotations
         self.closure = closure
         self.wrapped_fn: VariableTracker | None = wrapped_fn
 
@@ -1721,8 +1908,8 @@ def _get_function_impl(self, _converting: set[int]) -> types.FunctionType:
             func.__defaults__ = self.defaults.as_python_constant()
         if self.kwdefaults:
             func.__kwdefaults__ = self.kwdefaults.as_python_constant()
-        if self.annotations:
-            annotations = self.annotations.as_python_constant()
+        if self.dict_vt and self.dict_vt.contains("__annotations__"):
+            annotations = self.dict_vt.getitem("__annotations__").as_python_constant()
             if isinstance(annotations, tuple):
                 from itertools import pairwise
 
@@ -1733,25 +1920,27 @@ def _get_function_impl(self, _converting: set[int]) -> types.FunctionType:
             func.__annotations__ = annotations
         return func
 
-    def call_setattr(
-        self,
-        tx: "InstructionTranslator",
-        name_var: VariableTracker,
-        val: VariableTracker,
-    ) -> VariableTracker:
-        tx.output.side_effects.store_attr(self, name_var.value, val)  # type: ignore[attr-defined]
-        return CONSTANT_VARIABLE_NONE
-
-    def call_method(
-        self,
-        tx: "InstructionTranslator",
-        name: str,
-        args: Sequence[VariableTracker],
-        kwargs: dict[str, VariableTracker],
-    ) -> VariableTracker:
-        if name == "__setattr__":
-            return self.call_setattr(tx, *args)
-        return super().call_method(tx, name, list(args), kwargs)
+    def var_getattr(self, tx: "InstructionTranslator", name: str) -> VariableTracker:
+        if name in (
+            "__annotations__",
+            "__dict__",
+            "__doc__",
+            "__code__",
+            "__module__",
+            "__name__",
+            "__qualname__",
+            "__type_params__",
+        ):
+            return super().var_getattr(tx, name)
+        if name == "__defaults__":
+            d = getattr(self, "defaults", None)
+            return d.as_python_constant() if d else ConstantVariable.create(None)
+        elif name in cmp_name_to_op_mapping:
+            return variables.GetAttrVariable(
+                self, name, py_type=type(getattr(types.FunctionType, name))
+            )
+        else:
+            return super().var_getattr(tx, name)
 
     def has_closure(self) -> bool:
         return self.closure is not None
@@ -1770,9 +1959,12 @@ def call_obj_hasattr(
         self, tx: "InstructionTranslator", name: str
     ) -> ConstantVariable:
         if name == "__code__":
-            return variables.ConstantVariable.create(hasattr(self, "code"))
+            return VariableTracker.build(tx, hasattr(self, "code"))
         if name == "__defaults__":
-            return variables.ConstantVariable.create(hasattr(self, "defaults"))
+            return VariableTracker.build(tx, hasattr(self, "defaults"))
+        vt = ConstantVariable.create(name)
+        if vt in self.get_dict_vt(tx):
+            return ConstantVariable.create(True)
         return super().call_obj_hasattr(tx, name)
 
     def has_self(self) -> bool:
@@ -1833,14 +2025,15 @@ def reconstruct(self, codegen: "PyCodegen") -> None:
         else:
             codegen.extend_output([codegen.create_load_const(None)])
 
-        if self.annotations:
+        if self.dict_vt and self.dict_vt.contains("__annotations__"):
+            annotations = self.dict_vt.getitem("__annotations__")
             try:
-                annotations = self.annotations.as_python_constant()
+                annotations = annotations.as_python_constant()
                 codegen.extend_output(
                     [codegen.create_load_const_unchecked(annotations)]
                 )
             except NotImplementedError:
-                codegen(self.annotations)
+                codegen(annotations)
         else:
             codegen.extend_output([codegen.create_load_const(None)])
 
@@ -1888,7 +2081,6 @@ def __init__(
             wrapped.f_globals,
             wrapped.defaults,
             wrapped.kwdefaults,
-            wrapped.annotations,
             wrapped.closure,
             wrapped.wrapped_fn,
         )
@@ -1901,6 +2093,11 @@ def call_function(
         args: Sequence[VariableTracker],
         kwargs: dict[str, VariableTracker],
     ) -> VariableTracker:
+        if config.nested_graph_breaks:
+            wrapper_fn = UserFunctionVariable(polyfills._fn_with_ctx)
+            return wrapper_fn.call_function(
+                tx, [self.context, self.wrapped] + list(args), kwargs
+            )
         self.context.enter(tx)
         result = super().call_function(tx, args, kwargs)
         self.context.exit(tx)
@@ -1927,6 +2124,9 @@ def __init__(self, value: Any, reason: str | None = None, **kwargs: Any) -> None
     def as_python_constant(self) -> Any:
         return self.value
 
+    def get_real_python_backed_value(self) -> Any:
+        return self.value
+
     @classmethod
     def create_with_source(cls, value: Any, source: Source) -> "SkipFunctionVariable":
         # Use closure match guard (i.e. guard on __code__ object instead of
@@ -1946,7 +2146,14 @@ def create_with_source(cls, value: Any, source: Source) -> "SkipFunctionVariable
 
             guard_on_source.make_guard(GuardBuilder.CLOSURE_MATCH)
         elif inspect.isbuiltin(value):
-            install_guard(source.make_guard(GuardBuilder.BUILTIN_MATCH))
+            # Bound builtin methods (e.g. obj.__reduce_ex__) are created fresh
+            # on every attribute access, so their id() is unstable.  Skip the
+            # id-based BUILTIN_MATCH guard for them — the type guard on
+            # the owner object is sufficient.
+            if not hasattr(value, "__self__") or isinstance(
+                value.__self__, types.ModuleType
+            ):
+                install_guard(source.make_guard(GuardBuilder.BUILTIN_MATCH))
         elif not is_wrapper_or_member_descriptor(value):
             # These descriptors are not guaranteed to return the same object on
             # attribute lookup. They are unlikely to be changed, so we can skip
@@ -2011,6 +2218,28 @@ def call_function(
                 )
             except Unsupported as e:
                 raise StepUnsupported(e.msg) from None
+        elif self.value is types.FunctionType.__get__:
+            # function.__get__(func, obj[, cls]) produces a bound method.
+            # This is called by inspect._descriptor_get when resolving
+            # descriptors during inspect.signature().
+            # Note that function.__get__ does not use the 3rd argument. The
+            # reason it still has the 3rd argument is because descriptors follow
+            # a function signature that takes 3 arguments, and other descriptors
+            # (not function.__get__) can use the 3rd argument.
+            if len(args) in (2, 3) and not kwargs:
+                func_var = args[0]
+                obj_var = args[1]
+                if isinstance(func_var, UserFunctionVariable):
+                    return UserMethodVariable(
+                        func_var.fn, obj_var, source_fn=func_var.source
+                    )
+            unimplemented(
+                gb_type="unsupported function.__get__ call",
+                context=f"call_function {self}, args: {args}, kwargs: {kwargs}",
+                explanation="Dynamo only supports function.__get__(func, obj[, cls]) "
+                "where func is a user-defined function.",
+                hints=[*graph_break_hints.SUPPORTABLE],
+            )
         else:
             if config.dont_skip_tracing:
                 from .builder import SourcelessBuilder
@@ -2082,13 +2311,20 @@ def call_function(
                     torch._dynamo.utils.warn_once(explanation + "\n" + "\n".join(hints))
             if qualname == "allow_in_graph":
                 explanation = (
-                    "Found an allow_in_graph decorator to a function which "
-                    "is created inside the parent function that is getting "
-                    "compiled. This is not supported for now."
+                    "torch.compiler.allow_in_graph (or torch._dynamo.allow_in_graph) "
+                    "was called inside a compiled region. Dynamically annotating functions "
+                    "inside a compiled region is not supported."
                 )
-                # pyrefly: ignore [implicit-any]
-                hints = []
-            reason = self.reason if self.reason else "<missing reason>"
+                hints = [
+                    "Apply @torch.compiler.allow_in_graph as a decorator before compilation, "
+                    "not inside the compiled function.",
+                ]
+            if self.reason:
+                reason = self.reason
+            else:
+                from ..trace_rules import get_skip_reason
+
+                reason = get_skip_reason(self.value)
             unimplemented(
                 gb_type="Attempted to call function marked as skipped",
                 context=f"module: {module_name}, qualname: {qualname}, skip reason: {reason}",
@@ -2099,11 +2335,13 @@ def call_function(
     def call_obj_hasattr(
         self, tx: "InstructionTranslator", name: str
     ) -> ConstantVariable:
-        return variables.ConstantVariable.create(hasattr(self.value, name))
+        return VariableTracker.build(tx, hasattr(self.value, name))
 
     def var_getattr(self, tx: "InstructionTranslator", name: str) -> VariableTracker:
         if name in cmp_name_to_op_mapping:
-            return variables.GetAttrVariable(self, name)
+            return variables.GetAttrVariable(
+                self, name, py_type=type(getattr(self.value, name))
+            )
 
         return fn_var_getattr(tx, self.value, self.source, name)
 
@@ -2139,6 +2377,11 @@ def call_function(
         args: Sequence[VariableTracker],
         kwargs: dict[str, VariableTracker],
     ) -> VariableTracker:
+        if config.nested_graph_breaks:
+            wrapper_fn = UserFunctionVariable(polyfills._fn_with_ctx)
+            return wrapper_fn.call_function(
+                tx, [self.context, self.wrapped] + list(args), kwargs
+            )
         self.context.enter(tx)
         result = super().call_function(tx, args, kwargs)
         self.context.exit(tx)
@@ -2150,7 +2393,7 @@ def reconstruct(self, codegen: "PyCodegen") -> None:
         codegen.extend_output(create_call_function(1, False))
 
 
-class WrapperUserFunctionVariable(VariableTracker):
+class WrapperUserFunctionVariable(BaseUserFunctionVariable):
     """
     Used to represent a wrapper object that contains the actual callable as an
     attribute. For example, torch.jit.script/trace have the original function at
@@ -2158,19 +2401,33 @@ class WrapperUserFunctionVariable(VariableTracker):
     __script_if_tracing_wrapper have the original attr at "__original_fn".
     """
 
+    def python_type(self) -> type:
+        return types.FunctionType
+
     def __init__(self, wrapper_obj: Any, attr_to_trace: str, **kwargs: Any) -> None:
         super().__init__(**kwargs)
         self.wrapper_obj = wrapper_obj
         self.attr_to_trace = attr_to_trace
 
+    def get_module(self) -> str:
+        return self.wrapper_obj.__module__
+
+    def get_name(self) -> str:
+        return self.wrapper_obj.__name__
+
+    def get_code(self) -> types.CodeType:
+        return self.get_function().__code__
+
     def var_getattr(self, tx: "InstructionTranslator", name: str) -> VariableTracker:
         if name == self.attr_to_trace:
             val = getattr(self.wrapper_obj, self.attr_to_trace)
             source = self.source and AttrSource(self.source, name)
             return VariableTracker.build(tx, val, source)
-
         return super().var_getattr(tx, name)
 
+    def get_function(self):
+        return getattr(self.wrapper_obj, self.attr_to_trace)
+
     def self_args(self) -> list[VariableTracker]:
         return []
 
@@ -2212,14 +2469,18 @@ def call_function(
                     dynamo_logger.debug(user_stack_trace)
 
         all_args = self.self_args() + list(args)
-        return variables.UserFunctionVariable(
-            polyfills.getattr_and_trace  # type: ignore[arg-type]
+        return VariableTracker.build(
+            tx,
+            polyfills.getattr_and_trace,  # type: ignore[arg-type]
         ).call_function(
             tx,
-            [self, variables.ConstantVariable(self.attr_to_trace), *all_args],
+            [self, VariableTracker.build(tx, self.attr_to_trace), *all_args],
             kwargs,
         )
 
+    def get_real_python_backed_value(self) -> object:
+        return getattr(self.wrapper_obj, self.attr_to_trace)
+
 
 class WrapperUserMethodVariable(WrapperUserFunctionVariable):
     """
@@ -2228,6 +2489,9 @@ class WrapperUserMethodVariable(WrapperUserFunctionVariable):
     WrapperUserFunctionVariable in `call_function` method.
     """
 
+    def python_type(self) -> type:
+        return types.MethodType
+
     def __init__(
         self,
         wrapper_obj: Any,
@@ -2342,6 +2606,81 @@ def call_function(
                 ],
             )
 
+        if self.fn == dist.batch_isend_irecv:
+            if not config.enable_p2p_compilation:
+                unimplemented(
+                    gb_type="P2P compilation disabled for batch_isend_irecv",
+                    context=f"{self.fn}",
+                    explanation="P2P compilation is disabled.",
+                    hints=[
+                        "Set TORCHDYNAMO_ENABLE_P2P_COMPILATION=1 to enable.",
+                    ],
+                )
+
+            p2p_ops = kwargs["p2p_op_list"]
+            if not isinstance(p2p_ops, variables.ListVariable):
+                raise torch._dynamo.exc.InternalTorchDynamoError(
+                    "`P2POp` used incorrectly"
+                )
+
+            ops: list[VariableTracker] = list()
+            peers = list()
+            tags = list()
+            tensors = list()
+            group_var: VariableTracker | None = None
+
+            for item in p2p_ops.items:
+                if item.python_type() is not dist.P2POp:
+                    raise torch._dynamo.exc.InternalTorchDynamoError(
+                        "`P2POp` used incorrectly"
+                    )
+
+                op_var = item.var_getattr(tx, "op")
+                if op_var.is_python_constant():
+                    op = op_var.as_python_constant()
+                    if op not in (dist.isend, dist.irecv):
+                        raise torch._dynamo.exc.InternalTorchDynamoError(
+                            f"unexpected P2POp op {op}"
+                        )
+                    op_var = variables.ConstantVariable.create(op.__name__)
+                elif hasattr(op_var, "get_name"):
+                    op_var = variables.ConstantVariable.create(op_var.get_name())
+                else:
+                    raise torch._dynamo.exc.InternalTorchDynamoError(
+                        f"unexpected P2POp op variable {op_var}"
+                    )
+
+                ops.append(op_var)
+                tensors.append(item.var_getattr(tx, "tensor"))
+                peers.append(item.var_getattr(tx, "peer"))
+                tags.append(item.var_getattr(tx, "tag"))
+                if group_var is None:
+                    group_var = item.var_getattr(tx, "group")
+
+            assert group_var is not None
+            new_args: tuple[VariableTracker, ...] = ()
+            new_kwargs: dict[str, VariableTracker] = {
+                "op_list": variables.ListVariable(ops),
+                "peer_list": variables.ListVariable(peers),
+                "tag_list": variables.ListVariable(tags),
+                "tensors": variables.ListVariable(tensors),
+                "group_name": group_var,
+            }
+            return self.replacement_var.call_function(tx, new_args, new_kwargs)
+
+        if self.fn in (dist.isend, dist.irecv):
+            if not config.enable_p2p_compilation:
+                unimplemented(
+                    gb_type="P2P compilation disabled for isend/irecv",
+                    context=f"{self.fn}",
+                    explanation="P2P compilation is disabled.",
+                    hints=[
+                        "Set TORCHDYNAMO_ENABLE_P2P_COMPILATION=1 to enable.",
+                    ],
+                )
+
+            return self.replacement_var.call_function(tx, args, kwargs)
+
         if self.fn in (
             dist.all_reduce,
             dist.reduce_scatter_tensor,
@@ -2356,38 +2695,10 @@ def call_function(
             )
             if reduce_op not in REDUCE_OP_TO_STR:
                 raise ValueError(f"Unsupported all_reduce op: {reduce_op}")
-            kwargs["op"] = variables.ConstantVariable.create(
-                REDUCE_OP_TO_STR[reduce_op]
-            )
+            kwargs["op"] = VariableTracker.build(tx, REDUCE_OP_TO_STR[reduce_op])
         return self.replacement_var.call_function(tx, args, kwargs)
 
 
-class FunctoolsWrapsVariable(UserFunctionVariable):
-    def call_function(
-        self,
-        tx: "InstructionTranslator",
-        args: Sequence[VariableTracker],
-        kwargs: dict[str, VariableTracker],
-    ) -> VariableTracker:
-        if not kwargs and len(args) == 1:
-
-            def wraps(fn: Any) -> VariableTracker:
-                if isinstance(fn, variables.NestedUserFunctionVariable):
-                    return fn.clone(wrapped_fn=args[0])
-                unimplemented(
-                    gb_type="functools.wraps",
-                    context=f"{fn}",
-                    explanation="`torch.compile` can't trace `functools.wraps` on functions defined outside the compile region",
-                    hints=[
-                        *graph_break_hints.SUPPORTABLE,
-                    ],
-                )
-
-            return variables.LambdaVariable(wraps)
-
-        return super().call_function(tx, args, kwargs)
-
-
 class CollectionsNamedTupleFunction(UserFunctionVariable):
     def as_python_constant(self) -> Any:
         return self.fn
@@ -2409,10 +2720,9 @@ def call_function(
                 raise_observed_exception(
                     type(exc),
                     tx,
-                    args=list(map(ConstantVariable.create, exc.args)),
+                    args=list(exc.args),
                 )
             return variables.UserDefinedClassVariable(
-                # pyrefly: ignore[unbound-name]
                 value,
                 mutation_type=ValueMutationNew(),
             )
@@ -2427,6 +2737,9 @@ def call_function(
 
 
 class FunctoolsPartialVariable(VariableTracker):
+    # partial_type_spec: https://github.com/python/cpython/blob/v3.13.0/Modules/_functoolsmodule.c#L538
+    _cpython_type = functools.partial
+
     _nonvar_fields = {
         "original_cache_hash",
         *VariableTracker._nonvar_fields,
@@ -2487,9 +2800,7 @@ def call_obj_hasattr(
         self, tx: "InstructionTranslator", name: str
     ) -> ConstantVariable:
         # functools.partial uses slots, so attributes are constant
-        return variables.ConstantVariable.create(
-            hasattr(functools.partial(identity), name)
-        )
+        return VariableTracker.build(tx, hasattr(functools.partial(identity), name))
 
     def var_getattr(self, tx: "InstructionTranslator", name: str) -> VariableTracker:
         source = self.source and AttrSource(self.source, name)
@@ -2499,10 +2810,12 @@ def var_getattr(self, tx: "InstructionTranslator", name: str) -> VariableTracker
         if name == "args":
             return variables.ListVariable(self.args, source=source)
         if name == "keywords":
-            items = {ConstantVariable.create(k): v for k, v in self.keywords.items()}
+            items = {VariableTracker.build(tx, k): v for k, v in self.keywords.items()}
             return variables.ConstDictVariable(items, source=source)
         if name in cmp_name_to_op_mapping:
-            return variables.GetAttrVariable(self, name)
+            return variables.GetAttrVariable(
+                self, name, py_type=type(getattr(functools.partial, name))
+            )
         raise_observed_exception(AttributeError, tx)
 
     def as_python_constant(self) -> Any:
@@ -2677,11 +2990,10 @@ def call_method(
 
         method = getattr(self.fn, name, None)
         if not (method or is_function(method)):
-            raise_type_error_exc(tx, f"Cannot find callable {name} in {self.fn}")
+            raise_type_error(tx, f"Cannot find callable {name} in {self.fn}")
         options = {}
         if self.source:
             options["source"] = AttrSource(self.source, name)
-        # pyrefly: ignore[bad-specialization]
         polyfilled_method_variable = PolyfilledFunctionVariable(method, **options)
         return polyfilled_method_variable.call_function(tx, args, kwargs)
 
@@ -2694,6 +3006,9 @@ def __init__(self, value: Any, **kwargs: Any) -> None:
         super().__init__(**kwargs)
         self.value = value
 
+    def python_type(self) -> type:
+        return types.BuiltinFunctionType
+
     def exc_info(self, tx: "InstructionTranslator") -> "variables.TupleVariable":
         if len(tx.exn_vt_stack):
             exn = tx.exn_vt_stack[-1]
@@ -2702,9 +3017,9 @@ def exc_info(self, tx: "InstructionTranslator") -> "variables.TupleVariable":
             items = [VariableTracker.build(tx, typ), exn, tb]
         else:
             items = [
-                variables.CONSTANT_VARIABLE_NONE,
-                variables.CONSTANT_VARIABLE_NONE,
-                variables.CONSTANT_VARIABLE_NONE,
+                ConstantVariable.create(None),
+                ConstantVariable.create(None),
+                ConstantVariable.create(None),
             ]
         return variables.TupleVariable(items)  # type: ignore[arg-type]
 
@@ -2767,7 +3082,7 @@ def check_grid(self, grid: "BaseListVariable") -> tuple[torch.fx.proxy.Proxy, ..
     def call_grid(
         self, grid: Any, meta: dict[str, Any], tx: "InstructionTranslator"
     ) -> Any:
-        meta_var = {variables.ConstantVariable.create(k): v for k, v in meta.items()}
+        meta_var = {VariableTracker.build(tx, k): v for k, v in meta.items()}
         grid = grid.call_function(tx, [meta_var], {})
         return grid
 
@@ -2796,9 +3111,12 @@ def wrap_user_defined_obj(
         from .builder import VariableBuilder
 
         assert tx is not None
+        # Route through VariableBuilder.__call__ so already-tracked mutable
+        # objects (for example autotuner config lists) are reused instead of
+        # being registered for mutation twice in the same trace.
         wrapped_user_obj = VariableBuilder(
             tx, AttrSource(variable.kernel_source, f"{name}")
-        )._wrap(user_obj)
+        )(user_obj)
         return wrapped_user_obj
 
     def maybe_unpack_configs(
@@ -2837,7 +3155,7 @@ def call_getitem(  # type: ignore[override]
             kernel=variable.kernel,
             kernel_idx=variable.kernel_idx,
             grid=args[0],
-            kernel_source=variable.source,
+            kernel_source=variable.kernel_source,
         )
 
     def call_HOP(
@@ -2865,7 +3183,7 @@ def call_HOP(
                 combined_args[k] = v.get_tensor()
 
         combined_args_vt = {
-            variables.ConstantVariable.create(k): v for k, v in combined_args.items()
+            VariableTracker.build(tx, k): v for k, v in combined_args.items()
         }
 
         from torch._higher_order_ops.triton_kernel_wrap import (
@@ -2909,7 +3227,8 @@ def call_HOP(
             },
         )
 
-        return variables.ConstantVariable(
+        return VariableTracker.build(
+            tx,
             None,
         )
 
@@ -2921,12 +3240,12 @@ class TritonKernelVariable(VariableTracker):
     grid: "TritonGridType"
     kernel: "TritonKernelType"
     kernel_idx: int | None
-    kernel_source: "AttrSource"
+    kernel_source: Source | None
 
     def __init__(
         self, kernel: Any, kernel_idx: int | None, grid: Any, **kwargs: Any
     ) -> None:
-        self.kernel_source = kwargs.pop("kernel_source", None)
+        self.kernel_source = kwargs.pop("kernel_source", kwargs.get("source"))
         super().__init__(**kwargs)
         dynamo_triton_hopifier_singleton.init_variable(self, kernel, kernel_idx, grid)
 
@@ -2940,6 +3259,15 @@ def call_function(
             self, args, kwargs, tx
         )
 
+    def mp_subscript_impl(
+        self,
+        tx: "InstructionTranslator",
+        key: VariableTracker,
+    ) -> VariableTracker:
+        # Triton kernel[grid] — triton-specific, not a CPython slot.
+        # TODO(follow-up): add test for invalid key type
+        return dynamo_triton_hopifier_singleton.call_getitem(self, [key])
+
     def call_method(
         self,
         tx: "InstructionTranslator",
@@ -2947,9 +3275,7 @@ def call_method(
         args: list[VariableTracker],
         kwargs: dict[str, VariableTracker],
     ) -> VariableTracker:
-        if name == "__getitem__":
-            return dynamo_triton_hopifier_singleton.call_getitem(self, args)
-        elif name == "run":
+        if name == "run":
             return dynamo_triton_hopifier_singleton.call_run(self, args, kwargs, tx)  # type: ignore[return-value]
 
         # Bail out to parent's implementation
@@ -3047,6 +3373,9 @@ def __init__(
         super().__init__(**kwargs)
         self.rank = rank
 
+    def python_type(self) -> type:
+        return types.FunctionType
+
     def call_function(
         self,
         tx: "InstructionTranslator",
@@ -3071,7 +3400,7 @@ def call_function(
 
         if self.rank == 1:
             if len(args) + len(kwargs) != 4:
-                raise_type_error_exc(
+                raise_type_error(
                     tx,
                     f"TMA metadata rank=1 requires exactly 4 arguments, got {len(args) + len(kwargs)}",
                 )
@@ -3083,7 +3412,7 @@ def call_function(
             ]
         else:
             if len(args) + len(kwargs) != 6:
-                raise_type_error_exc(
+                raise_type_error(
                     tx,
                     f"TMA metadata rank=2 requires exactly 6 arguments, got {len(args) + len(kwargs)}",
                 )
@@ -3109,6 +3438,9 @@ def call_function(
 
 
 class CreateTMADescriptorStableVariable(VariableTracker):
+    def python_type(self) -> type:
+        return types.FunctionType
+
     def call_function(
         self,
         tx: "InstructionTranslator",
@@ -3146,9 +3478,8 @@ def call_function(
         kwargs: dict[str, VariableTracker],
     ) -> VariableTracker:
         if len(args) != 1:
-            raise_type_error_exc(
-                tx,
-                f"pytree_get_node_type requires exactly 1 argument, got {len(args)}",
+            raise_type_error(
+                tx, f"pytree_get_node_type requires exactly 1 argument, got {len(args)}"
             )
         type_source = None
         if args[0].source:
@@ -3185,13 +3516,12 @@ def call_function(
     ) -> VariableTracker:
         # tree_is_leaf(tree, is_leaf=None)
         if len(args) < 1 or len(args) > 2:
-            raise_type_error_exc(
-                tx,
-                f"tree_is_leaf requires 1 or 2 arguments, got {len(args)}",
+            raise_type_error(
+                tx, f"tree_is_leaf requires 1 or 2 arguments, got {len(args)}"
             )
 
         # Check if is_leaf parameter is provided
-        is_leaf = kwargs.get("is_leaf", CONSTANT_VARIABLE_NONE)
+        is_leaf = kwargs.get("is_leaf", ConstantVariable.create(None))
         if len(args) == 2:
             is_leaf = args[1]
 
@@ -3213,7 +3543,7 @@ def call_function(
             source=get_pytree_SUPPORTED_NODES_source(),
         )
         out = supported_nodes_var.call_method(tx, "__contains__", [node_type_var], {})
-        return ConstantVariable.create(not out.value)
+        return VariableTracker.build(tx, not out.value)
 
 
 class SparseTensorCreationSkipVariable(SkipFunctionVariable):
@@ -3245,14 +3575,67 @@ def call_function(
         )
 
 
-class TritonSetAllocatorSkipVariable(SkipFunctionVariable):
-    """
-    Skip variable for triton.set_allocator with a clear message to move it outside the compiled region.
+def emit_noargs_leaf_function_to_graph(
+    tx: "InstructionTranslator",
+    real_impl: Callable[[], None],
+    name: str,
+) -> None:
+    """Emit an invoke_leaf_function node for a side-effectful function with no
+    tensor inputs or outputs.
+
+    The function is captured as a closure inside _LeafCallable objects and
+    registered as a static attribute on the graph module.  Because
+    invoke_leaf_function is registered as EffectType.ORDERED, effect tokens
+    prevent DCE and maintain execution ordering relative to other ops.
+
+    Use this when Dynamo needs to preserve a pure-side-effect call (like
+    setting global runtime state) in the compiled graph so that it replays
+    at the correct position at runtime.
     """
+    import torch.utils._pytree as pytree
+    from torch._higher_order_ops.invoke_leaf_function import (
+        _LeafCallable,
+        invoke_leaf_function,
+        make_leaf_function_wrappers,
+    )
+
+    def fake_impl():
+        return None
+
+    captured_out_spec: list[pytree.TreeSpec | None] = [None]
+    wrapped_real, wrapped_fake = make_leaf_function_wrappers(
+        real_impl, fake_impl, captured_out_spec
+    )
+
+    real_callable = _LeafCallable(wrapped_real)
+    fake_callable = _LeafCallable(wrapped_fake)
+    input_spec = pytree.tree_flatten(((), {}))[1]
+
+    def make_proxy(attr_name: str, val: Any) -> Any:
+        proxy = tx.output.register_static_attr_and_return_proxy(attr_name, val)
+        proxy.node.type = type(val)
+        return proxy
+
+    invoke_args = (
+        make_proxy(f"{name}_real_fn", real_callable),
+        make_proxy(f"{name}_fake_fn", fake_callable),
+        make_proxy(f"{name}_input_spec", input_spec),
+        "",  # mutated_flat_indices
+    )
+    tx.output.create_proxy("call_function", invoke_leaf_function, invoke_args, {})
+
+
+class TritonSetAllocatorVariable(VariableTracker):
+    """Trace triton.set_allocator as an invoke_leaf_function node in the
+    graph so that it executes at the right point at runtime, ordered by
+    effect tokens."""
 
     def __init__(self, value: Any, **kwargs: Any) -> None:
-        reason = "triton.set_allocator is not supported inside torch.compile"
-        super().__init__(value, reason=reason, **kwargs)
+        super().__init__(**kwargs)
+        self.value = value
+
+    def python_type(self) -> type:
+        return type(self.value)
 
     def call_function(
         self,
@@ -3260,15 +3643,16 @@ def call_function(
         args: Sequence[VariableTracker],
         kwargs: dict[str, VariableTracker],
     ) -> VariableTracker:
-        unimplemented(
-            gb_type="triton.set_allocator not supported",
-            context="triton.set_allocator called inside compiled region",
-            explanation=(
-                "triton.set_allocator is not supported inside torch.compile. "
-                "It modifies global Triton allocator state and cannot be traced."
-            ),
-            hints=[
-                "Move triton.set_allocator() outside of the torch.compile region "
-                "(call it before the compiled function)."
-            ],
-        )
+        assert len(args) == 1 and not kwargs
+        alloc_fn = args[0].as_python_constant()
+
+        # Emit an invoke_leaf_function node so it runs at runtime.
+        set_allocator = self.value
+
+        def real_impl():
+            set_allocator(alloc_fn)
+            return None
+
+        emit_noargs_leaf_function_to_graph(tx, real_impl, "set_alloc")
+
+        return ConstantVariable.create(None)
diff --git a/torch/_dynamo/variables/hashable.py b/torch/_dynamo/variables/hashable.py
new file mode 100644
index 0000000000000..5896479eaa536
--- /dev/null
+++ b/torch/_dynamo/variables/hashable.py
@@ -0,0 +1,163 @@
+"""
+Hashability utilities for PyTorch Dynamo variable tracking.
+
+This module provides the HashableTracker wrapper class and associated utilities
+for making VariableTracker instances usable as dictionary keys and set elements
+during symbolic execution. Used by both ConstDictVariable and SetVariable.
+"""
+
+from typing import TYPE_CHECKING
+
+import torch
+
+from .. import variables
+from ..exc import raise_observed_exception
+from ..utils import specialize_symnode
+from .base import VariableTracker
+
+
+if TYPE_CHECKING:
+    from torch._dynamo.symbolic_convert import InstructionTranslator
+
+
+def raise_unhashable(
+    arg: VariableTracker, tx: "InstructionTranslator | None" = None
+) -> None:
+    if tx is None:
+        from torch._dynamo.symbolic_convert import InstructionTranslator
+
+        tx = InstructionTranslator.current_tx()
+    try:
+        arg_type = arg.python_type()
+    except Exception:
+        arg_type = type(arg)
+
+    raise_observed_exception(
+        TypeError,
+        tx,
+        args=[
+            f"unhashable type: {arg_type!r} and variable tracker = {type(arg.realize())}",
+        ],
+    )
+
+
+def is_hashable(x: VariableTracker) -> bool:
+    # NB - performing isinstance check on a LazVT realizes the VT, accidentally
+    # inserting the guard. To avoid this, lazyVT `is_hashable` methods looks at
+    # the underlying value without realizing the VT. Consider updating the
+    # lazyVT `is_hashable` method if you see unnecessary guarding for a key VT.
+    if (
+        isinstance(x, variables.LazyVariableTracker)
+        and not x.is_realized()
+        and x.is_hashable()
+    ):
+        return True
+    return x.is_python_hashable()
+
+
+class HashableTracker:
+    """
+    Class that wraps a VariableTracker and makes it hashable.
+    Note that it's fine to put VTs into dictionaries and sets, but doing so
+    does not take into account aliasing.
+    """
+
+    _MISSING = object()
+
+    def __init__(self, vt: VariableTracker) -> None:
+        # We specialize SymNodes
+        vt = specialize_symnode(vt)
+
+        # If Dynamo does not know the hashability of the vt, it will raise unsupported here
+        # TODO(follow-up): check tp_hash via C-level slot detection — unhashable keys
+        # (e.g. list) should raise TypeError, not graph break via is_python_hashable/unimplemented.
+        if not is_hashable(vt):
+            raise_unhashable(vt)
+        self.vt = vt
+
+    @classmethod
+    def _maybe_constant_torch_size(cls, vt: VariableTracker) -> object:
+        from .lists import SizeVariable
+        from .tensor import TensorVariable
+
+        if (
+            isinstance(vt, variables.LazyVariableTracker)
+            and not vt.is_realized()
+            and isinstance(vt.original_value(), torch.Size)
+        ):
+            return vt.original_value()
+
+        if not isinstance(vt, SizeVariable):
+            return cls._MISSING
+
+        items = []
+        for item in vt.items:
+            if item.is_python_constant():
+                items.append(item.as_python_constant())
+                continue
+
+            if isinstance(item, TensorVariable):
+                proxy = getattr(item, "proxy", None)
+                node = getattr(proxy, "node", None)
+                meta = getattr(node, "meta", None) if node is not None else None
+                example_value = (
+                    meta.get("example_value") if isinstance(meta, dict) else None
+                )
+                constant = getattr(example_value, "constant", None)
+
+                if isinstance(constant, torch.Tensor) and constant.numel() == 1:
+                    items.append(constant.item())
+                    continue
+
+            return cls._MISSING
+
+        return torch.Size(items)
+
+    def __hash__(self) -> int:
+        """
+        Computes the hash value for the wrapped VariableTracker.
+
+        For unrealized LazyVariableTrackers, uses the hash of the original value
+        to avoid realizing the tracker and inserting unnecessary guards.
+        For all other cases, delegates to the VariableTracker's get_python_hash method.
+
+        Returns:
+            The hash value of the underlying variable tracker
+        """
+        if (
+            isinstance(self.vt, variables.LazyVariableTracker)
+            and not self.vt.is_realized()
+            and self.vt.is_hashable()
+        ):
+            return hash(self.vt.original_value())
+
+        maybe_constant = self._maybe_constant_torch_size(self.vt)
+        if maybe_constant is not self._MISSING:
+            return hash(maybe_constant)
+
+        return self.vt.get_python_hash()
+
+    def __eq__(self, other: object) -> bool:
+        """
+        Checks equality between two HashableTracker instances.
+
+        Delegates to the VariableTracker's is_python_equal method to compare
+        the underlying variable trackers for Python-level equality.
+
+        Args:
+            other: Another HashableTracker instance to compare with
+
+        Returns:
+            True if the underlying variable trackers are Python-equal, False otherwise
+        """
+        if not isinstance(other, HashableTracker):
+            return False
+        if self.vt is other.vt:
+            return True
+
+        self_constant = self._maybe_constant_torch_size(self.vt)
+        other_constant = self._maybe_constant_torch_size(other.vt)
+        if self_constant is not self._MISSING and other_constant is not self._MISSING:
+            return self_constant == other_constant
+
+        return self.vt.is_python_equal(other.vt)
diff --git a/torch/_dynamo/variables/higher_order_ops.py b/torch/_dynamo/variables/higher_order_ops.py
index f01360ce3748e..18f99046ce767 100644
--- a/torch/_dynamo/variables/higher_order_ops.py
+++ b/torch/_dynamo/variables/higher_order_ops.py
@@ -23,21 +23,23 @@
 import inspect
 import itertools
 import logging
+import traceback
 import types
 import warnings
 from collections.abc import Sequence
-from dataclasses import dataclass
-from typing import Any, Literal, Optional, TYPE_CHECKING, Union
+from dataclasses import dataclass, field
+from typing import Any, cast, Literal, Optional, TYPE_CHECKING, Union
 
 import torch._C
 import torch.fx
 import torch.nn
 from torch._dispatch.python import enable_python_dispatcher
 from torch._dynamo.utils import get_fake_value
-from torch._dynamo.variables.constant import CONSTANT_VARIABLE_NONE, ConstantVariable
+from torch._dynamo.variables.constant import ConstantVariable
 from torch._dynamo.variables.ctx_manager import RepararametrizeModuleContextVariable
 from torch._dynamo.variables.functions import UserFunctionVariable
 from torch._dynamo.variables.nn_module import UnspecializedNNModuleVariable
+from torch._dynamo.variables.script_object import TorchScriptObjectVariable
 from torch._dynamo.variables.tensor import SymNodeVariable, TensorVariable
 from torch._guards import Source
 from torch._higher_order_ops.invoke_subgraph import NestedCompileRegionOptions
@@ -62,6 +64,7 @@
 from .dicts import ConstDictVariable
 from .lazy import LazyVariableTracker
 from .lists import ListVariable, TupleVariable
+from .sets import SetVariable
 
 
 if TYPE_CHECKING:
@@ -70,7 +73,7 @@
     from torch._subclasses.fake_tensor import FakeTensor, FakeTensorMode
     from . import AutogradFunctionContextVariable
 
-from collections.abc import Callable, Generator, Iterable
+from collections.abc import Generator, Iterable
 from typing import ParamSpec, TypeVar
 
 
@@ -111,6 +114,21 @@ def __post_init__(self) -> None:
             assert len(self.masks_to_filter_const_values) == len(self.const_values)
 
 
+@dataclass(frozen=True)
+class SubgraphTracingInfo:
+    """Properties observed during subgraph tracing.
+
+    Expandable container for metadata collected while tracing a HOP subgraph.
+    Future additions may include aliasing info, mutation info, etc.
+    """
+
+    # User code stack at the point where the first externally-visible side
+    # effect was detected.  None means no side effect was observed.
+    side_effect_stack: traceback.StackSummary | None = None
+    # All sources accessed via VariableBuilder during the subgraph trace.
+    traced_sources: OrderedSet[Any] = field(default_factory=OrderedSet)
+
+
 # This function is a syntax sugar for creating a dummy new subtracer so that
 # newly added nodes are added to a separate subgraph in this subtracer instead of affecting
 # the main graph. This is useful for creating sample inputs for tracing the subgraph.
@@ -216,6 +234,9 @@ def find_mismatched_vars(
     elif isinstance(var, ConstDictVariable):
         for value in var.items.values():
             mismatched_vars.update(find_mismatched_vars(value, types, allow_none))
+    elif isinstance(var, SetVariable):
+        for key in var.items:
+            mismatched_vars.update(find_mismatched_vars(key.vt, types, allow_none))
     else:
         if not isinstance(var, types) and not (allow_none and var.is_constant_none()):
             mismatched_vars.add(var)
@@ -232,22 +253,7 @@ def only_consist_of(
 # A more read-able syntax sugar for creating a UserFunctionVariable for f
 # and run call_function on it. Make it return a function to preserve the calling
 # convention of the original f.
-def _make_inlined(
-    tx: "InstructionTranslator", f: Callable[..., Any]
-) -> Callable[..., VariableTracker]:
-    assert callable(f), "Expect f to be a python callable."
-
-    def inline_call(
-        *args: VariableTracker, **kwargs: VariableTracker
-    ) -> VariableTracker:
-        from torch._dynamo.trace_rules import _force_inline
-
-        with _force_inline():
-            # We cannot use SourcelessBuilder here because it ignores _force_inline.
-            # since the rules are static.
-            return UserFunctionVariable(f).call_function(tx, args, kwargs)  # type: ignore[arg-type]
-
-    return inline_call
+from torch._dynamo.utils import _make_inlined
 
 
 def add_call_function(
@@ -256,7 +262,7 @@ def add_call_function(
     args: tuple[Any, ...],
     kwargs: dict[str, Any],
     flat_example_value: Any,
-    config: Optional[NestedCompileRegionOptions] = None,
+    config: NestedCompileRegionOptions | None = None,
 ) -> VariableTracker:
     from .builder import wrap_fx_proxy
 
@@ -321,8 +327,17 @@ def overwrite_tensor_vt_proxy(
     # while still allowing `body_r` to contain arbitrary Python objects.
     # pyrefly: ignore[missing-attribute]
     for orig_vt, subgraph_vt in zip(graph_output_vts, flat_variable.items):
-        if isinstance(orig_vt, (variables.SymNodeVariable, variables.TensorVariable)):
-            assert subgraph_vt.is_tensor() or isinstance(subgraph_vt, SymNodeVariable)
+        if isinstance(
+            orig_vt,
+            (
+                variables.SymNodeVariable,
+                variables.TensorVariable,
+                TorchScriptObjectVariable,
+            ),
+        ):
+            assert subgraph_vt.is_tensor() or isinstance(
+                subgraph_vt, (SymNodeVariable, TorchScriptObjectVariable)
+            )
             orig_vt.proxy = subgraph_vt.proxy
 
 
@@ -573,7 +588,7 @@ def collect_intermediate_outputs(
         tracker.collect_from_inputs(tx)
         tracker.collect_from_outputs(graph_output_vts)
 
-    for out in subtracer.tracked_tensor_or_symint_vt:
+    for out in subtracer.tracked_proxyable_vt:
         proxy = out.as_proxy()
 
         # Skip if already in output
@@ -897,6 +912,7 @@ def are_same_graph_modules(
     from torch._subclasses.fake_tensor import extract_tensor_metadata
 
     # Maps the equivalent nodes from a to b
+    # pyrefly: ignore [implicit-any]
     node_map = {}
 
     def check_all_args(a_nodes: Iterable[Any], b_nodes: Iterable[Any]) -> bool:
@@ -1469,8 +1485,8 @@ def speculate_subgraph_with_auto_output_flattening(
     *,
     # source_target is the .value of HigherOrderOpVariable and is the
     # target of the proxy that we created for the higherOrderOperator.
-    source_target: Optional[HigherOrderOperator] = None,
-    enable_grad: Optional[bool] = None,
+    source_target: HigherOrderOperator | None = None,
+    enable_grad: bool | None = None,
     # automatic: relies on Dynamo to find the used tensors and lift them as
     # inputs.
     #
@@ -1544,6 +1560,7 @@ def speculate_subgraph_with_auto_output_flattening(
     | tuple[
         VariableTracker, ...
     ],  # graph_output_vts: Tensor/symint VTs that are actual FX graph outputs
+    SubgraphTracingInfo,  # tracing_info: Properties observed during subgraph tracing
 ]:
     """
     Speculate subgraph for Higher-Order Operators (HOPs) with automatic output flattening.
@@ -1632,14 +1649,15 @@ def gn(x):
         )
 
     try:
-        # ensure guards on args get installed in parent subgraph
-        f, sub_args, sub_kwargs = LazyVariableTracker.realize_all(
-            (f, sub_args, sub_kwargs),
-        )
-
         with tx.output.subtracer(source_target, tracer, description) as subtracer:
             args = get_hop_args(
-                tx, f, subtracer, sub_args, sub_kwargs, set_subgraph_inputs, description
+                tx,
+                f,
+                subtracer,
+                list(sub_args),
+                sub_kwargs,
+                set_subgraph_inputs,
+                description,
             )
 
             # Special case - if users uses
@@ -1690,14 +1708,15 @@ def gn(x):
             # be actual FX graph outputs.
             # Collect only tensor and symint VTs that should be graph outputs.
             # We walk the output structure and extract proxyable VTs.
-            # pyrefly: ignore [implicit-any]
             graph_output_vt_list = []
 
             def visit(vt: VariableTracker) -> None:
-                if vt.is_tensor() or isinstance(vt, SymNodeVariable):
+                if vt.is_tensor() or isinstance(
+                    vt, (SymNodeVariable, TorchScriptObjectVariable)
+                ):
                     graph_output_vt_list.append(vt)
 
-            VariableTracker.visit(visit, output)
+            VariableTracker.visit(visit, output, side_effects=tx.output.side_effects)
             graph_output_vts = tuple(graph_output_vt_list)
 
             # NOTE - [Return subgraph intermediates as subgraph outputs]
@@ -1738,11 +1757,13 @@ def visit(vt: VariableTracker) -> None:
             # nested_compile_region and autograd.Function. Today, its safe
             # because we error out on seeing a side-effect.
 
-            allow_side_effects = (
-                allow_side_effects
-                or tx.output.current_tracer.traced_with_externally_visible_side_effects
+            traced_externally = (
+                tx.output.current_tracer.traced_with_externally_visible_side_effects
+            )
+            has_side_effects = (
+                subtracer.side_effect_stack is not None or traced_externally
             )
-            if allow_side_effects:
+            if (allow_side_effects or traced_externally) and has_side_effects:
                 extra_outputs = collect_intermediate_outputs(
                     tx, subtracer, graph_output_vts, filter_aliased_intermediates
                 )
@@ -1797,11 +1818,18 @@ def visit(vt: VariableTracker) -> None:
             # - `lifted_freevars`: Free variables lifted as inputs to the subgraph
             # - `graph_output_vts`: Only the tensor/symint VTs that are actual
             #   FX graph outputs (basically the vts associated with graph outputs)
+            # - `tracing_info`: Properties observed during subgraph tracing
+            tracing_info = SubgraphTracingInfo(
+                side_effect_stack=subtracer.side_effect_stack,
+                traced_sources=subtracer.traced_sources,
+            )
+
             return (
                 output,
                 graph,
                 lifted_freevars,
                 graph_output_vts,
+                tracing_info,
             )
     except Unsupported as ex:
         f_name = f"{type(f).__name__}"
@@ -2084,7 +2112,7 @@ def wrapped_call_function(self, *args: Any, **kwargs: Any) -> VariableTracker:
 
 class TorchHigherOrderOperatorVariable(VariableTracker):
     # Subclasses should set _HOP_NAME to enable automatic HOP context in error messages
-    _HOP_NAME: Optional[str] = None
+    _HOP_NAME: str | None = None
     # Set to False for HOPs that hard error on graph break (e.g., cond, map, scan); otherwise
     # HOPs will fall back to eager.
     _ALLOW_FALLBACK_TO_EAGER: bool = True
@@ -2202,6 +2230,9 @@ def _call_function(
     ) -> VariableTracker:
         from . import ListVariable
 
+        self.supports_input_mutation = not torch.is_grad_enabled()
+        self.supports_aliasing = not torch.is_grad_enabled()
+
         args, kwargs = LazyVariableTracker.realize_all((args, kwargs))
 
         for i, k in enumerate(["pred", "true_fn", "false_fn", "operands"]):
@@ -2474,10 +2505,12 @@ def validate_subgraph_output_types(
     ):
         for out in non_tensor_output:
             if (
-                isinstance(out, SymNodeVariable) and out.python_type() in (int, bool)
-            ) or (
-                out.is_python_constant()
-                and isinstance(out.as_python_constant(), (int, bool))
+                (isinstance(out, SymNodeVariable) and out.python_type() in (int, bool))
+                or (
+                    out.is_python_constant()
+                    and isinstance(out.as_python_constant(), (int, bool))
+                )
+                or isinstance(out, TorchScriptObjectVariable)
             ):
                 continue
             unimplemented(
@@ -3242,15 +3275,6 @@ def call_function(
         args: Sequence[VariableTracker],
         kwargs: dict[str, VariableTracker],
     ) -> VariableTracker:
-        if not torch._dynamo.config.inline_inbuilt_nn_modules:
-            unimplemented(
-                gb_type="torch.func.functional_call capture is disabled",
-                context="",
-                explanation="torch.func.functional_call capture is disabled",
-                hints=[
-                    "Set `torch._dynamo.config.inline_inbuilt_nn_modules=True` to enable.",
-                ],
-            )
         return super().call_function(tx, args, kwargs)
 
 
@@ -3305,6 +3329,7 @@ def create_wrapped_node(
         GraphModule,
         str,
         VariableTracker | tuple[VariableTracker, ...],
+        SubgraphTracingInfo,
     ]:
         # See NOTE [HigherOrderOperator tracing design] for more details
         (
@@ -3312,6 +3337,7 @@ def create_wrapped_node(
             body_graph,
             body_lifted_freevars,
             body_graph_output_vts,
+            tracing_info,
         ) = speculate_subgraph_with_auto_output_flattening(
             tx,
             fn_vt,
@@ -3358,6 +3384,7 @@ def create_wrapped_node(
             body_gmod,
             body_name,
             body_graph_output_vts,
+            tracing_info,
         )
 
     def _call_function(
@@ -3375,6 +3402,7 @@ def _call_function(
             _,
             _,
             body_graph_output_vts,
+            _,
         ) = self.create_wrapped_node(tx, args[0], args[1:], kwargs, "wrap")
 
         if len(p_kwargs) > 0:
@@ -3644,6 +3672,7 @@ def _call_function(
             body_gmod,
             _,
             body_graph_output_vts,
+            _,
         ) = self.create_wrapped_node(
             tx,
             args[0],  # function
@@ -3732,11 +3761,13 @@ def _call_function(
         unpacked_sequence = args[1].unpack_var_sequence(tx)
         # TODO (tmanlaibaatar) support pytree here
         for arg in unpacked_sequence:
-            if isinstance(arg, (ListVariable, TupleVariable, ConstDictVariable)):
+            if isinstance(
+                arg, (ListVariable, TupleVariable, ConstDictVariable, SetVariable)
+            ):
                 unimplemented(
                     gb_type="strict_mode: improper args",
                     context=f"args: {args}, kwargs: {kwargs}",
-                    explanation="strict_mode higher order op expects flat inputs (list/tuple/dict)",
+                    explanation="strict_mode higher order op expects flat inputs (list/tuple/dict/set)",
                     hints=[
                         *graph_break_hints.USER_ERROR,
                     ],
@@ -3841,6 +3872,7 @@ def _call_function(
             checkpointed_gmod,
             _,
             body_graph_output_vts,
+            _,
         ) = self.create_wrapped_node(
             tx,
             args[0],
@@ -3896,6 +3928,7 @@ def _call_function(
             gmod,
             _,
             body_graph_output_vts,
+            _,
         ) = self.create_wrapped_node(
             tx,
             args[1],
@@ -3970,6 +4003,31 @@ def _call_function(
         )
 
 
+class InlineAsmElementwiseHigherOrderVariable(TorchHigherOrderOperatorVariable):
+    _HOP_NAME = "inline_asm_elementwise"
+
+    def _call_function(
+        self,
+        tx: "InstructionTranslator",
+        args: Sequence[VariableTracker],
+        kwargs: dict[str, VariableTracker],
+    ) -> VariableTracker:
+        from .builder import wrap_fx_proxy
+
+        p_args = tuple(arg.as_proxy() for arg in args)
+        p_kwargs = {key: arg.as_proxy() for key, arg in kwargs.items()}
+        return wrap_fx_proxy(
+            tx=tx,
+            proxy=tx.output.create_proxy(
+                "call_function",
+                self.value,
+                args=p_args,
+                kwargs=p_kwargs,
+            ),
+            example_value=None,
+        )
+
+
 class AutoFunctionalizeHigherOrderVariable(TorchHigherOrderOperatorVariable):
     _HOP_NAME = "torch.ops.higher_order.auto_functionalized"
 
@@ -3998,6 +4056,14 @@ def _call_function(
 class FlexAttentionBackwardHighOrderVariable(TorchHigherOrderOperatorVariable):
     _HOP_NAME = "torch.ops.higher_order.flex_attention_backward"
 
+    @staticmethod
+    def _uses_pretraced_graphs(
+        fw_graph: VariableTracker, joint_graph: VariableTracker
+    ) -> bool:
+        return not joint_graph.is_constant_none() or isinstance(
+            fw_graph, UnspecializedNNModuleVariable
+        )
+
     def proxy_submod(
         self, tx: "InstructionTranslator", arg: UnspecializedNNModuleVariable
     ) -> Proxy:
@@ -4017,6 +4083,95 @@ def to_proxy(self, tx: "InstructionTranslator", arg: VariableTracker) -> Any:
         else:
             return arg.as_proxy()
 
+    def create_wrapped_node(
+        self,
+        tx: "InstructionTranslator",
+        query: VariableTracker,
+        fn: VariableTracker,
+        fn_name: str,
+        other_buffers: Sequence[VariableTracker],
+    ) -> tuple[Proxy, tuple[Proxy, ...], torch.fx.GraphModule]:
+        from .._trace_wrapped_higher_order_op import TransformGetItemToIndex
+
+        def create_scalar() -> VariableTracker:
+            return query.call_method(
+                tx,
+                "new_empty",
+                [VariableTracker.build(tx, [])],
+                {"dtype": VariableTracker.build(tx, torch.int32)},
+            )
+
+        with discard_graph_changes(tx):
+            bhmn = [create_scalar() for _ in range(4)]
+            if fn_name == "score_mod":
+                scores_require_grad: bool = query.requires_grad  # type: ignore[attr-defined]
+                score = query.call_method(
+                    tx,
+                    "new_empty",
+                    [VariableTracker.build(tx, [])],
+                    {"requires_grad": VariableTracker.build(tx, scores_require_grad)},
+                )
+                new_args = [score, *bhmn, *other_buffers]
+            else:
+                assert fn_name == "mask_fn", "Illegal function name: " + fn_name
+                new_args = [*bhmn, *other_buffers]
+
+        with TransformGetItemToIndex():
+            (
+                (_body_output, _body_spec),
+                body_graph,
+                body_lifted_freevars,
+            ) = speculate_subgraph(
+                tx,
+                fn,
+                new_args,
+                {},
+                description=f"{self._HOP_NAME}: {fn_name}",
+                source_target=self.value,
+                set_subgraph_inputs="flatten_manual",
+            )
+
+        gm = torch.fx.GraphModule(tx.output.nn_modules, body_graph)
+        body_name = tx.output.install_subgraph(fn_name, gm)
+        return make_attr(tx, body_name), tuple(body_lifted_freevars), gm
+
+    @staticmethod
+    def _buffer_example_value(buffer: VariableTracker) -> Any:
+        proxy = buffer.as_proxy()
+        if isinstance(proxy, Proxy):
+            return proxy.node.meta["example_value"]
+        return proxy
+
+    def _derive_joint_graph(
+        self,
+        tx: "InstructionTranslator",
+        query: VariableTracker,
+        fw_graph_gm: torch.fx.GraphModule,
+        score_mod_other_buffers: TupleVariable,
+        fw_graph_lifted_args: tuple[Proxy, ...],
+    ) -> Proxy:
+        from torch._higher_order_ops.flex_attention import create_fw_bw_graph
+
+        query_example = query.as_proxy().node.meta["example_value"]
+        example_vals = (
+            query_example.new_zeros((), requires_grad=True),
+            query_example.new_zeros((), dtype=torch.int),
+            query_example.new_zeros((), dtype=torch.int),
+            query_example.new_zeros((), dtype=torch.int),
+            query_example.new_zeros((), dtype=torch.int),
+        )
+        all_buffer_examples = tuple(
+            self._buffer_example_value(buf) for buf in score_mod_other_buffers.items
+        ) + tuple(buf.node.meta["example_value"] for buf in fw_graph_lifted_args)
+
+        _, joint_gm = create_fw_bw_graph(fw_graph_gm, example_vals, all_buffer_examples)
+        joint_gm = cast(GraphModule, joint_gm)
+
+        submod_name = tx.output.install_subgraph("joint_graph", joint_gm)
+        p_submod = make_attr(tx, submod_name)
+        set_example_value(p_submod.node, joint_gm)
+        return p_submod
+
     def _call_function(
         self,
         tx: "InstructionTranslator",
@@ -4025,6 +4180,130 @@ def _call_function(
     ) -> VariableTracker:
         from .builder import wrap_fx_proxy
 
+        if len(args) != 14:
+            return self._call_function_fallback(tx, args, kwargs)
+
+        (
+            query,
+            key,
+            value,
+            out,
+            logsumexp,
+            grad_out,
+            grad_logsumexp,
+            fw_graph,
+            joint_graph,
+            block_mask,
+            scale,
+            kernel_options,
+            score_mod_other_buffers,
+            mask_mod_other_buffers,
+        ) = args
+
+        if (
+            not isinstance(block_mask, TupleVariable)
+            or not isinstance(score_mod_other_buffers, TupleVariable)
+            or not isinstance(mask_mod_other_buffers, TupleVariable)
+            or len(block_mask.items) < 1
+        ):
+            return self._call_function_fallback(tx, args, kwargs)
+
+        if self._uses_pretraced_graphs(fw_graph, joint_graph):
+            return self._call_function_fallback(tx, args, kwargs)
+
+        fw_graph_node, fw_graph_lifted_args, fw_graph_gm = self.create_wrapped_node(
+            tx, query, fw_graph, "score_mod", score_mod_other_buffers.items
+        )
+
+        joint_graph_node = self._derive_joint_graph(
+            tx,
+            query,
+            fw_graph_gm,
+            score_mod_other_buffers,
+            fw_graph_lifted_args,
+        )
+
+        mask_fn = block_mask.items[-1]
+        if mask_fn.is_python_constant() and mask_fn.as_python_constant() is None:
+            mask_fn = VariableTracker.build(
+                tx,
+                torch.nn.attention.flex_attention.noop_mask,
+                source=mask_fn.source,
+            )
+        mask_fn_node, mask_fn_lifted_args, _ = self.create_wrapped_node(
+            tx, query, mask_fn, "mask_fn", mask_mod_other_buffers.items
+        )
+
+        proxied_args = [
+            query,
+            key,
+            value,
+            out,
+            logsumexp,
+            grad_out,
+            grad_logsumexp,
+            TupleVariable(block_mask.items[:-1], source=block_mask.source),
+            scale,
+            kernel_options,
+        ]
+        inp_args, _ = proxy_args_kwargs(proxied_args, {})
+        proxied_score_mod_other_buffers = tuple(
+            self.to_proxy(tx, arg) for arg in score_mod_other_buffers.items
+        )
+        proxied_mask_mod_other_buffers = tuple(
+            self.to_proxy(tx, arg) for arg in mask_mod_other_buffers.items
+        )
+
+        (
+            inp_q,
+            inp_k,
+            inp_v,
+            inp_out,
+            inp_lse,
+            inp_grad_out,
+            inp_grad_lse,
+            inp_block_mask,
+            inp_scale,
+            inp_kernel_options,
+        ) = inp_args
+
+        block_mask_proxy = tuple(inp_block_mask + (mask_fn_node,))
+
+        with torch.fx.experimental.proxy_tensor.set_original_aten_op(self.value):
+            return wrap_fx_proxy(
+                tx=tx,
+                proxy=tx.output.create_proxy(
+                    "call_function",
+                    self.value,
+                    args=(
+                        inp_q,
+                        inp_k,
+                        inp_v,
+                        inp_out,
+                        inp_lse,
+                        inp_grad_out,
+                        inp_grad_lse,
+                        fw_graph_node,
+                        joint_graph_node,
+                        block_mask_proxy,
+                        inp_scale,
+                        inp_kernel_options,
+                        proxied_score_mod_other_buffers + fw_graph_lifted_args,
+                        proxied_mask_mod_other_buffers + mask_fn_lifted_args,
+                    ),
+                    kwargs={},
+                ),
+                example_value=None,
+            )
+
+    def _call_function_fallback(
+        self,
+        tx: "InstructionTranslator",
+        args: Sequence[VariableTracker],
+        kwargs: dict[str, VariableTracker],
+    ) -> VariableTracker:
+        from .builder import wrap_fx_proxy
+
         p_args, p_kwargs = None, None
         try:
             p_args = tuple(self.to_proxy(tx, arg) for arg in args)
@@ -4245,6 +4524,9 @@ def __init__(
         self.bwd_fn = bwd_fn
         self.parent_source = parent_source
 
+    def python_type(self) -> type:
+        return types.BuiltinMethodType
+
     def call_function(
         self,
         tx: "InstructionTranslator",
@@ -4473,7 +4755,7 @@ def trace_forward_graph(
         )
 
         # Speculate subgraph on the fwd
-        fwd_out, fwd_graph, fwd_freevars, fwd_graph_output_vts = (
+        fwd_out, fwd_graph, fwd_freevars, fwd_graph_output_vts, _ = (
             speculate_subgraph_with_auto_output_flattening(
                 tx,
                 fwd_fn,
@@ -4483,6 +4765,7 @@ def trace_forward_graph(
                 enable_grad=None,
                 set_subgraph_inputs="automatic",
                 allow_side_effects=True,
+                filter_aliased_intermediates=True,
                 tracer=fwd_tracer,
             )
         )
@@ -4528,7 +4811,7 @@ def trace_backward_graph(
         """
         Traces the backward method of the autograd.Function object.
         """
-        from . import UserDefinedClassVariable, UserMethodVariable
+        from . import UserMethodVariable
 
         # Note that for the forward, we do not restore side effects, because we
         # want the later tracing to see the side-effects. But for backward, we
@@ -4554,7 +4837,7 @@ def trace_backward_graph(
                 if i.is_tensor():
                     bwd_args.append(i)
                 else:
-                    bwd_args.append(CONSTANT_VARIABLE_NONE)
+                    bwd_args.append(ConstantVariable.create(None))
 
         bwd_fn, bwd_args = self.prepare_fn_vt(tx, ctx, "backward", bwd_args)
 
@@ -4573,7 +4856,7 @@ def is_strict_for(v: VariableTracker) -> bool:
             tx.strict_translation_mode(is_strict_for),
         ):
             try:
-                bwd_out, bwd_graph, bwd_freevars, bwd_graph_output_vts = (
+                bwd_out, bwd_graph, bwd_freevars, bwd_graph_output_vts, _ = (
                     speculate_subgraph_with_auto_output_flattening(
                         tx,
                         bwd_fn,
@@ -4610,7 +4893,7 @@ def is_strict_for(v: VariableTracker) -> bool:
                 elif isinstance(self.bwd_fn, types.MethodType):
                     bwd_fn = UserMethodVariable(
                         autograd_function_backward_rewritten(self.bwd_fn.__func__),
-                        UserDefinedClassVariable(self.bwd_fn.__class__),
+                        VariableTracker.build(tx, self.bwd_fn.__class__),
                     )
                 else:
                     unimplemented(
@@ -4625,7 +4908,7 @@ def is_strict_for(v: VariableTracker) -> bool:
                     "torch._dynamo.config._autograd_backward_strict_mode_conditional_banned_ops",
                     [],
                 ):
-                    bwd_out, bwd_graph, bwd_freevars, bwd_graph_output_vts = (
+                    bwd_out, bwd_graph, bwd_freevars, bwd_graph_output_vts, _ = (
                         speculate_subgraph_with_auto_output_flattening(
                             tx,
                             bwd_fn,
@@ -5000,7 +5283,7 @@ def prepare_fn_vt(
         method_name: str,
         args: Sequence[VariableTracker],
     ) -> tuple[VariableTracker, Sequence[VariableTracker]]:
-        from . import UserDefinedClassVariable, UserMethodVariable
+        from . import UserMethodVariable
 
         source = None
         if self.parent_source:
@@ -5016,7 +5299,7 @@ def prepare_fn_vt(
             fn_vt = VariableTracker.build(tx, fn, source=source)
             fn_args = [ctx, *args]
         elif isinstance(fn, types.MethodType):
-            cls_vt = UserDefinedClassVariable(fn.__class__)
+            cls_vt = VariableTracker.build(tx, fn.__class__)
             fn_vt = UserMethodVariable(
                 fn.__func__,
                 cls_vt,
@@ -5098,6 +5381,7 @@ def _call_function(
             _,
             _,
             body_graph_output_vts,
+            _,
         ) = self.create_wrapped_node(
             tx, args[0], args[1:], {}, self.value._name, subgraph_name="subgraph"
         )
@@ -5115,150 +5399,6 @@ def _call_function(
         )
 
 
-class InvokeSubgraphHigherOrderVariable(WrapHigherOrderVariable):
-    _HOP_NAME = "torch.ops.higher_order.invoke_subgraph"
-    _ALLOW_FALLBACK_TO_EAGER = False
-    supports_input_mutation = True
-    supports_aliasing = False
-    allow_side_effects = True
-    # invoke_subgraph is NOT desugared in AOTAutograd, so the HOP input/output
-    # shouldn't alias. For checkpoint HOP, we inline it so we don't need
-    # alias analysis as functionalization would just work on the flat graph.
-    filter_aliased_intermediates = True
-
-    # pyrefly: ignore[bad-override]
-    def install_subgraph_in_output_graph(
-        self,
-        tx: "InstructionTranslator",
-        fn_vt: VariableTracker,
-        fn_args_vt: Sequence[VariableTracker],
-        kwargs: dict[str, VariableTracker],
-        body_gmod: GraphModule,
-        attr_name: str,
-    ) -> str:
-        # Check if the subgraph from speculate_subgraph (body_gmod) and the fake
-        # inputs have already been seen before. If yes, the subgraph is already
-        # installed in the output graph and we can just access the subgraph
-        # using the saved attr name.
-
-        if not isinstance(fn_vt, (UnspecializedNNModuleVariable, UserFunctionVariable)):
-            unimplemented(
-                gb_type="Encountered non user function variable during invoke_subgraph HOP tracing",
-                context=str(fn_vt),
-                explanation="invoke_subgraph does not support non user function variable",
-                hints=[*graph_break_hints.SUPPORTABLE],
-            )
-
-        invoke_subgraph_cache = (
-            tx.output.tracing_context.hop_dispatch_set_cache.get_cache(
-                torch._higher_order_ops.invoke_subgraph
-            )
-        )
-
-        if isinstance(fn_vt, UserFunctionVariable):
-            fn_id = id(fn_vt.get_function())
-            fn_name = fn_vt.get_function().__name__
-        else:
-            assert isinstance(fn_vt, UnspecializedNNModuleVariable)
-            fn_id = id(fn_vt.value.forward.__func__)  # type: ignore[attr-defined]
-            fn_name = fn_vt.value.forward.__name__  # type: ignore[attr-defined]
-        # pyrefly: ignore [implicit-any]
-        previously_installed_submodules = []
-        if invoke_subgraph_cache:
-            previously_installed_submodules = (
-                invoke_subgraph_cache.get_dynamo_installed_submodules(fn_id)
-            )
-            current_mod = body_gmod
-            # NB - reverse is more likely to cause a hit sooner because first
-            # graph can have requires_grad=False for a few inputs
-            for submodule_name in reversed(previously_installed_submodules):
-                assert submodule_name in tx.output.nn_modules
-                previous_mod = tx.output.nn_modules[submodule_name]
-                assert tx.fake_mode
-                if are_same_graph_modules(
-                    fn_name, previous_mod, current_mod, tx.fake_mode
-                ):
-                    return submodule_name
-
-        body_name = super().install_subgraph_in_output_graph(
-            tx, fn_vt, fn_args_vt, kwargs, body_gmod, "subgraph"
-        )
-        hc_log.debug(
-            "%s: Installing subgraph with identifier '%s', bringing total count for '%s' function to %s",
-            fn_name,
-            body_name,
-            fn_name,
-            len(previously_installed_submodules) + 1,
-        )
-        if invoke_subgraph_cache:
-            invoke_subgraph_cache.add_dynamo_installed_submodule(fn_id, body_name)
-
-        return body_name
-
-    def _call_function(
-        self,
-        tx: "InstructionTranslator",
-        args: Sequence[VariableTracker],
-        kwargs: dict[str, VariableTracker],
-    ) -> VariableTracker:
-        # This flattens the kwargs into lifted args
-        assert self._HOP_NAME is not None
-        (
-            p_args,
-            p_kwargs,
-            example_value,
-            body_r,
-            body_gmod,
-            body_name,
-            body_graph_output_vts,
-        ) = self.create_wrapped_node(tx, args[0], args[1:], kwargs, self._HOP_NAME)
-
-        if len(p_kwargs) > 0:
-            unimplemented(
-                gb_type="invoke_subgraph: kwargs unexpected",
-                context=f"args: {args}, kwargs: {kwargs}",
-                explanation="kwargs should have been flattened into lifted args.",
-                hints=[
-                    *graph_break_hints.DYNAMO_BUG,
-                ],
-            )
-
-        # Extract nested compile config and store in node meta
-        # This will be used in regional_inductor_invoke_subgraph
-        config = None
-        fn_var = args[0]
-        if hasattr(fn_var, "get_function"):
-            try:
-                fn = fn_var.get_function()
-
-                if hasattr(fn, "__marked_compile_region_config__"):
-                    config = fn.__marked_compile_region_config__
-                    if config is not None:
-                        body_gmod.meta["nested_region_config"] = config
-            except Exception:
-                log.warning(
-                    "Failed to extract nested_compile_region() config from InvokeSubgraphHigherOrderVariable. ",
-                    exc_info=True,
-                )
-                raise
-
-        p_args = (
-            p_args[0],
-            body_name,
-            *p_args[1:],
-        )
-        return _call_function_with_auto_output_flattening(  # type: ignore[return-value]
-            tx,
-            torch._higher_order_ops.invoke_subgraph,
-            tuple(p_args),
-            p_kwargs,
-            example_value,
-            body_r,
-            body_graph_output_vts,
-            config=config,
-        )
-
-
 class LocalMapWrappedHigherOrderVariable(WrapHigherOrderVariable):
     _HOP_NAME = "torch.ops.higher_order.local_map_hop"
     supports_input_mutation = False
@@ -5404,6 +5544,7 @@ def check_none_last(placements: Sequence[Any | None]) -> int:
 
             priors[vt] = global_tensor
             vt.as_proxy().node.meta["example_value"] = local_tensor
+            # pyrefly: ignore [missing-attribute]
             vt.synchronize_attributes(tx)
 
         # Step 3: Trace local_map subgraph with local tensors
@@ -5415,6 +5556,7 @@ def check_none_last(placements: Sequence[Any | None]) -> int:
             body_gmod,
             body_name,
             body_graph_output_vts,
+            _,
         ) = self.create_wrapped_node(
             tx, user_func, user_args, kwargs, self.value._name, subgraph_name="subgraph"
         )
@@ -5487,6 +5629,7 @@ def make_error_msg(*args: Any) -> str:
         # Step 6: Restore inputs and outputs to global shapes
         for vt, global_tensor in priors.items():
             vt.as_proxy().node.meta["example_value"] = global_tensor
+            # pyrefly: ignore [missing-attribute]
             vt.synchronize_attributes(tx)
 
         outs = out.items if isinstance(out, TupleVariable) else [out]
@@ -5526,6 +5669,9 @@ def make_error_msg(*args: Any) -> str:
         return out
 
 
+from .invoke_subgraph import InvokeSubgraphHigherOrderVariable  # noqa: E402
+
+
 # Map operator names to their corresponding variable for fast TorchHigherOrderOperatorVariable.make()
 _hop_name_to_variable_class = {
     "cond": CondHigherOrderVariable,
@@ -5553,6 +5699,7 @@ def make_error_msg(*args: Any) -> str:
     "dynamo_bypassing_wrapper": DynamoBypassingWrapperHigherOrderVariable,
     "auto_functionalized": AutoFunctionalizeHigherOrderVariable,
     "auto_functionalized_v2": AutoFunctionalizeHigherOrderVariable,
+    "inline_asm_elementwise": InlineAsmElementwiseHigherOrderVariable,
     "invoke_subgraph": InvokeSubgraphHigherOrderVariable,
     "custom_function_call": CustomFunctionHigherOrderOperatorVariable,
     "local_map_hop": LocalMapWrappedHigherOrderVariable,
diff --git a/torch/_dynamo/variables/invoke_subgraph.py b/torch/_dynamo/variables/invoke_subgraph.py
new file mode 100644
index 0000000000000..0c135b3060d29
--- /dev/null
+++ b/torch/_dynamo/variables/invoke_subgraph.py
@@ -0,0 +1,1344 @@
+"""
+This module contains the InvokeSubgraphHigherOrderVariable class and its
+supporting helpers for subgraph reuse (auto-cache) in Dynamo's invoke_subgraph
+higher-order operator.
+"""
+
+import enum
+import logging
+import traceback
+import types
+from dataclasses import dataclass
+from typing import Any, cast, NamedTuple, TYPE_CHECKING
+
+import torch
+import torch._higher_order_ops
+from torch._dynamo import graph_break_hints
+from torch._dynamo.exc import unimplemented
+from torch._dynamo.guards import (
+    extract_tensor_metadata,
+    GUARD_VALUE_DISPATCH,
+    GuardCheckSpec,
+    SKIP_GUARD,
+    UnsupportedGuardCheckSpec,
+)
+from torch._dynamo.source import SyntheticLocalSource
+from torch._dynamo.variables.base import VariableTracker
+from torch._dynamo.variables.constant import ConstantVariable
+from torch._dynamo.variables.functions import UserFunctionVariable
+from torch._dynamo.variables.higher_order_ops import WrapHigherOrderVariable
+from torch._dynamo.variables.lists import ListVariable, TupleVariable
+from torch._dynamo.variables.nn_module import UnspecializedNNModuleVariable
+from torch._dynamo.variables.tensor import SymNodeVariable, TensorVariable
+from torch._guards import (
+    Guard,
+    InvokeSubgraphReuseCondition,
+    InvokeSubgraphReuseEntry,
+    Source,
+)
+from torch._higher_order_ops.invoke_subgraph import NestedCompileRegionOptions
+from torch.fx.graph_module import GraphModule
+from torch.fx.proxy import Proxy
+from torch.utils import _pytree as pytree
+from torch.utils._ordered_set import OrderedSet
+
+
+if TYPE_CHECKING:
+    from collections.abc import Sequence
+
+    from torch._dynamo.symbolic_convert import InstructionTranslator
+    from torch._dynamo.variables.higher_order_ops import SubgraphTracingInfo
+
+log = logging.getLogger(__name__)
+hc_log = torch._logging.getArtifactLogger(__name__, "hierarchical_compile")
+
+# Note: [invoke_subgraph subgraph reuse]
+#
+# When mark_compile_region wraps a function called N times (e.g. 80 identical
+# transformer layers), Dynamo traces the subgraph once and stamps out cached
+# copies for subsequent calls. It does safety checks to ensure that a subgraph
+# is reusable, if not (e.g. side-effect), it will fallback to tracing the
+# next invocation.
+#
+# HIGH-LEVEL FLOW
+# ===============
+#   User code: model.layers[0](x), model.layers[1](x), ..., model.layers[79](x)
+#                     |                     |                        |
+#                     v                     v                        v
+#              +--------------+     +--------------+        +--------------+
+#              |  First Call  |     |  Second Call  |  ...   |  80th Call   |
+#              +------+-------+     +------+-------+        +------+-------+
+#                     |                    |                        |
+#                     v                    v                        v
+#              +--------------+     +--------------+        +--------------+
+#              | Full subgraph|     | Cache lookup  |        | Cache lookup  |
+#              |   trace      |     | (is_reusable) |        | (is_reusable) |
+#              +------+-------+     +------+-------+        +------+-------+
+#                     |                    |                        |
+#                     v                    v                        v
+#              +--------------+     +--------------+        +--------------+
+#              | save_reuse_  |     | stamp_out_   |        | stamp_out_   |
+#              | entry()      |     | subgraph()   |        | subgraph()   |
+#              +--------------+     +--------------+        +--------------+
+#
+# WHAT GETS CACHED
+# ================
+# After the first trace, save_reuse_entry stores an InvokeSubgraphReuseEntry
+# (in _guards.py) containing:
+#   - body_name/body_gmod: the traced subgraph
+#   - arg_sources: sources of the original call's arguments
+#   - subgraph_input_mapping: how each lifted arg maps back to user inputs or captures
+#   - output_metadata: shape/stride/dtype/device of outputs
+#
+# Paired with an InvokeSubgraphReuseCondition containing:
+#   - input_checks: (tag, tensor_metadata) per input
+#   - guards: (source, handler, expected, guard) tuples
+#   - treespec: pytree structure of the args
+#   - traced_sources: sources accessed during the trace
+#
+# CACHE LOOKUP (is_reusable)
+# ==========================
+# On subsequent calls:
+#   1. Input structure match -- same treespec, tags, tensor metadata.
+#   2. Source replacement -- clone each guard's source with a replacement map
+#      (old: L['self'].layers[0].weight -> new: L['self'].layers[1].weight),
+#      then evaluate against the new source's runtime value.
+#   3. Mutation check -- reject if the subgraph mutated any captured var.
+#
+# A shared resolve_cache memoizes intermediate source resolution (e.g.
+# L['self'].layers is evaluated once and reused across all guards).
+#
+# STAMP OUT (stamp_out_subgraph)
+# ==============================
+# On cache hit, reconstruct the argument list using the freevar mapping
+# (list[LiftedArgOrigin]):
+#
+#   LiftedUserArg(index)
+#       User arg (activation / explicit input).
+#       Looked up from new call's flat proxies.
+#
+#   LiftedCapturedSource(source)
+#       Sourceful captured var (weight, param, etc).
+#       Source is cloned with replacement map, resolved via
+#       VariableBuilder. Deduplicates via input_source_to_var.
+#
+#   LiftedSyntheticObject(ctor_fn, ctor_args, ctor_arg_sources)
+#       Synthetic object (opaque type with SyntheticLocalSource).
+#       Reconstructed via synthetic_graph_input with cached constructor info.
+#
+# SAFETY
+# ======
+# In normal Dynamo compilation, safety is enforced at runtime: guards are
+# installed during tracing and re-evaluated on every subsequent call against
+# real Python objects.  Subgraph reuse operates differently — we are in the
+# middle of tracing, there are no real Python objects, only VariableTrackers.
+# We must answer: what could cause the second invocation of a nested compile
+# region to produce a different trace than the first?
+#
+# VariableTrackers fall into two categories:
+#
+# 1. Intermediates — values produced during tracing with no originating source
+#    (e.g. the result of a prior FX op). These can reach a nested compile region
+#    only via (a) the region's explicit function arguments, or (b) closure
+#    capture. We do not support nested-function regions that close over tensors,
+#    so only (a) applies. For explicit arguments, the set of types that can
+#    appear is small and well-defined: TensorVariable, SymNodeVariable,
+#    ConstantVariable, and NNModuleVariable. Each has a cheap structural
+#    comparison (tensor metadata, symnode identity, constant value equality).
+#    We also snapshot the pytree treespec of the argument list and verify it
+#    matches on lookup, ensuring the flattened structure is identical.
+#
+# 2. Sourceful variables — values with a known originating source (e.g. a
+#    module attribute or a local variable visible in the outer frame). For these
+#    we collect the guard delta from the first trace, parameterize the guard
+#    sources by replacing the original arg sources with the new arg sources, and
+#    re-evaluate the guards by resolving each source against the live f_locals /
+#    f_globals. The one extra hazard here is mutation: if the outer trace
+#    mutates a sourceful object between the first and second invocations, the
+#    cached guards would evaluate against stale values. We therefore also check
+#    that none of the sources read by the cached subgraph have been mutated in
+#    the outer SideEffects tracker before accepting a reuse.
+#
+# - max_reuse_entries (default 8, configurable via nested_compile_region arg)
+#   caps cache entries per function. Exceeding it raises RuntimeError.
+# - Guard failures logged with guard type + user stack trace.
+#   Enable: TORCH_LOGS='+hierarchical_compile'
+# ---------------------------------------------------------------------------
+# Auto-cache helpers for invoke_subgraph
+# ---------------------------------------------------------------------------
+
+
+class InputTag(enum.Enum):
+    TENSOR = "tensor"
+    SYMNODE = "symnode"
+    CONSTANT = "constant"
+    MODULE = "module"
+
+
+class InputFingerprint(NamedTuple):
+    # (InputTag, VariableTracker) pairs for each leaf input.
+    flat_vts: list[tuple[InputTag, VariableTracker]]
+    # 1-1 mapping to flat_vts: source for each leaf, or None if the VT has no source.
+    arg_sources: list[Source | None]
+    # True if any leaf VT had an unsupported type for reuse.
+    has_unknown: bool = False
+    # TreeSpec from pytree.tree_flatten of the (args, kwargs) structure.
+    treespec: pytree.TreeSpec | None = None
+
+
+def classify_vt(vt: Any) -> InputTag | None:
+    """Return the tag for a leaf VT, or None if unsupported."""
+    if isinstance(vt, TensorVariable):
+        return InputTag.TENSOR
+    elif isinstance(vt, SymNodeVariable):
+        return InputTag.SYMNODE
+    elif isinstance(vt, ConstantVariable):
+        return InputTag.CONSTANT
+    elif isinstance(vt, UnspecializedNNModuleVariable):
+        return InputTag.MODULE
+    return None
+
+
+def build_input_fingerprint(
+    tx: "InstructionTranslator",
+    fn_args_vt: Any,
+    kwargs: dict[str, Any],
+) -> InputFingerprint:
+    """Build an InputFingerprint by flattening (args, kwargs) via pytree.
+
+    Uses _make_inlined(tx, pytree.tree_flatten) to recursively flatten
+    the argument structure into leaf VTs, classifying each leaf as
+    tensor/symnode/constant/module. Also records the TreeSpec so that
+    cache lookups can verify structural equivalence.
+
+    Fast path: when kwargs is empty and all args are already leaf VTs
+    (tensor/symnode/constant/module), skip the expensive pytree flatten.
+    """
+    # Fast path: flat args, no kwargs — skip pytree machinery.
+    if not kwargs:
+        all_leaf = True
+        for vt in fn_args_vt:
+            if classify_vt(vt) is None:
+                all_leaf = False
+                break
+        if all_leaf:
+            return build_fingerprint_fast(fn_args_vt)
+
+    return build_fingerprint_with_pytree(tx, fn_args_vt, kwargs)
+
+
+def build_fingerprint_fast(fn_args_vt: Any) -> InputFingerprint:
+    """Build fingerprint for the common case of flat leaf args, no kwargs."""
+    flat_vts: list[tuple[InputTag, VariableTracker]] = []
+    arg_sources: list[Source | None] = []
+    for vt in fn_args_vt:
+        tag = classify_vt(vt)
+        assert tag is not None
+        flat_vts.append((tag, vt))
+        # Always append (even None) to keep positional alignment with flat_vts
+        # so that source_replacement zip pairing is correct across calls.
+        arg_sources.append(getattr(vt, "source", None))
+    return InputFingerprint(flat_vts, arg_sources)
+
+
+def build_fingerprint_with_pytree(
+    tx: "InstructionTranslator",
+    fn_args_vt: Any,
+    kwargs: dict[str, Any],
+) -> InputFingerprint:
+    """Build fingerprint via pytree flatten for nested/kwargs cases."""
+    from torch._dynamo.variables.builder import SourcelessBuilder
+    from torch._dynamo.variables.higher_order_ops import _make_inlined
+
+    container_vt = SourcelessBuilder.create(tx, (list(fn_args_vt), kwargs))
+    flat_list_vt, treespec_vt = _make_inlined(tx, pytree.tree_flatten)(
+        container_vt
+    ).unpack_var_sequence(tx)
+    treespec = treespec_vt.as_python_constant()
+
+    flat_vts: list[tuple[InputTag, VariableTracker]] = []
+    arg_sources: list[Source | None] = []
+    has_unknown = False
+
+    for vt in flat_list_vt.unpack_var_sequence(tx):
+        tag = classify_vt(vt)
+        if tag is not None:
+            flat_vts.append((tag, vt))
+        else:
+            has_unknown = True
+            continue
+
+        # Always append (even None) to keep positional alignment with flat_vts.
+        arg_sources.append(getattr(vt, "source", None))
+
+    return InputFingerprint(flat_vts, arg_sources, has_unknown, treespec)
+
+
+def get_flat_proxies(fingerprint: InputFingerprint) -> list[Proxy]:
+    """Collect deduplicated proxies from tensor/symnode leaves."""
+    seen: set[torch.fx.Node] = set()
+    flat_proxies: list[Proxy] = []
+    for tag, vt in fingerprint.flat_vts:
+        if tag in (InputTag.TENSOR, InputTag.SYMNODE):
+            proxy = vt.as_proxy()
+            if proxy.node not in seen:
+                seen.add(proxy.node)
+                flat_proxies.append(proxy)
+    return flat_proxies
+
+
+@dataclass
+class LiftedUserArg:
+    """Lifted arg that came from a user argument (intermediate activation or explicit input)."""
+
+    index: int
+
+
+@dataclass
+class LiftedCapturedSource:
+    """Lifted arg that is a captured variable (e.g. a weight or parameter) with a Source."""
+
+    source: Any  # Source
+
+
+@dataclass
+class LiftedSyntheticObject:
+    """Lifted arg that is a TorchScriptObject with a SyntheticLocalSource."""
+
+    ctor_fn: Any  # Callable
+    ctor_args: tuple[Any, ...]
+    ctor_arg_sources: tuple[Any, ...] | None
+
+
+@dataclass
+class LiftedBoundSymbol:
+    """Lifted arg that is a SymInt already bound as a graph input.
+
+    SymInt graph inputs are created during tensor wrapping (not through
+    VariableBuilder.wrap_symint), so they aren't registered in
+    unspec_variable_map or variable_tracker_cache. Using LiftedCapturedSource
+    for these would resolve the source to a concrete Python int via
+    source.get_value() instead of reusing the existing symbolic proxy.
+    """
+
+    expr: Any  # sympy.Expr
+
+
+LiftedArgOrigin = (
+    LiftedUserArg | LiftedCapturedSource | LiftedSyntheticObject | LiftedBoundSymbol
+)
+
+
+def get_fn_code(fn_var: Any) -> types.CodeType | None:
+    if isinstance(fn_var, UserFunctionVariable):
+        return fn_var.get_function().__code__
+    elif isinstance(fn_var, UnspecializedNNModuleVariable):
+        return (
+            fn_var.value.forward.__func__.__code__  # pyrefly: ignore[missing-attribute]
+        )
+    return None
+
+
+def has_mutated_vars(
+    tx: "InstructionTranslator",
+    traced_sources: OrderedSet[Source],
+) -> bool:
+    """Check if any source accessed by the subgraph has been mutated.
+
+    SideEffects.mutated_sources records the exact AttrSource for every
+    store_attr call. A simple set intersection with traced_sources tells
+    us whether any source the subgraph read was later written to.
+    """
+    overlap = tx.output.side_effects.mutated_sources & traced_sources
+    if overlap:
+        hc_log.debug(
+            "subgraph_reuse: mutated sources detected -- %s",
+            overlap,
+        )
+        return True
+    return False
+
+
+def is_reuse_eligible(
+    tx: "InstructionTranslator",
+    body_r: Any,
+    fingerprint: InputFingerprint,
+    tracing_info: "SubgraphTracingInfo",
+    traced_sources: OrderedSet[Source] | None = None,
+    has_reuse_hash_fn: bool = False,
+) -> bool:
+    """Best-effort check for whether a traced subgraph result can be reused.
+
+    It is possible that a subgraph is morally reusable but does not fall
+    into the limited support that Dynamo has today. Current limitations:
+      - The subgraph must not have side effects.
+      - No sourceful variable accessed by the subgraph may have been
+        mutated, because guards are snapshotted on source values at trace
+        time — if the underlying object changed since then, the cached
+        guards would silently evaluate against stale values.
+      - Output must be a single tensor, or a tuple/list of plain tensors.
+      - All flattened inputs must be one of: tensor, symnode, constant,
+        unspecialized NN module — for sourceless or other input types we
+        rely on the treespec and tags for structural matching, so only
+        types with well-defined comparison semantics are supported.
+
+    When ``has_reuse_hash_fn`` is True, side-effect and mutation checks are
+    skipped because the hash key replaces guards — there are no guards to
+    go stale from mutations.
+    """
+    if not has_reuse_hash_fn:
+        if tracing_info.side_effect_stack is not None:
+            stack_msg = "\n" + "".join(
+                traceback.format_list(tracing_info.side_effect_stack)
+            )
+            hc_log.debug(
+                "subgraph_reuse: not eligible -- subgraph has side effects%s",
+                stack_msg,
+            )
+            return False
+
+        if traced_sources and has_mutated_vars(tx, traced_sources):
+            return False
+
+    if isinstance(body_r, TensorVariable):
+        pass
+    elif isinstance(body_r, (TupleVariable, ListVariable)):
+        non_tensor = [
+            type(item).__name__
+            for item in body_r.items
+            if not isinstance(item, TensorVariable)
+        ]
+        if non_tensor:
+            hc_log.debug(
+                "subgraph_reuse: not eligible -- output contains non-tensor types: %s",
+                non_tensor,
+            )
+            return False
+    else:
+        hc_log.debug(
+            "subgraph_reuse: not eligible -- output type %s is not tensor or tuple/list",
+            type(body_r).__name__,
+        )
+        return False
+
+    if fingerprint.has_unknown:
+        hc_log.debug(
+            "subgraph_reuse: not eligible -- unsupported input VT types",
+        )
+        return False
+
+    return True
+
+
+def build_reuse_condition(
+    tx: "InstructionTranslator",
+    fingerprint: InputFingerprint,
+    traced_sources: OrderedSet[Source],
+) -> InvokeSubgraphReuseCondition | None:
+    """Build an InvokeSubgraphReuseCondition from a traced subgraph.
+
+    A reuse condition is a mix of two kinds of checks:
+
+    1. **Input tag checks** (from flat_vts): For each flattened leaf VT,
+       we record its tag (_VtTag.TENSOR/SYMNODE/CONSTANT/MODULE) and
+       metadata (e.g. tensor shape/stride/dtype/device/requires_grad).
+       At lookup time, the treespec ensures structural equivalence, and
+       then we compare tags and metadata leaf-by-leaf.
+
+    2. **Guard checks** (from traced_sources): During the subgraph trace,
+       every source accessed via VariableBuilder is recorded. We look up
+       all guards installed on those sources (and on the arg_sources) to
+       build the set of guards that must be re-evaluated on cache hit.
+       This is more robust than guard diffing because it catches guards
+       that were already installed before the subgraph trace began.
+
+    Raise if any guard type is unsupported, as a feedback for compiler
+    developers to support that guard type.
+    """
+    from torch._guards import InvokeSubgraphReuseCondition
+
+    input_checks: list[tuple[InputTag, object]] = []
+    for tag, vt in fingerprint.flat_vts:
+        if tag == InputTag.TENSOR:
+            assert isinstance(vt, TensorVariable)
+            example = vt.proxy.node.meta.get("example_value", None)
+            if example is None:
+                hc_log.debug(
+                    "subgraph_reuse: cannot build condition -- tensor input has no example_value"
+                )
+                return None
+            input_checks.append((InputTag.TENSOR, extract_tensor_metadata(example)))
+        elif tag == InputTag.SYMNODE:
+            assert isinstance(vt, SymNodeVariable)
+            # Store the SymInt/SymFloat/SymBool object itself. Two accesses to
+            # the same symbolic dimension (e.g. x.shape[0] twice) produce the
+            # same Python object, so identity comparison in is_reusable is
+            # correct and avoids false matches between distinct symbols.
+            input_checks.append((InputTag.SYMNODE, vt.sym_num))
+        elif tag == InputTag.CONSTANT:
+            assert isinstance(vt, ConstantVariable)
+            input_checks.append((InputTag.CONSTANT, vt.value))
+        elif tag == InputTag.MODULE:
+            input_checks.append((InputTag.MODULE, None))
+        else:
+            raise AssertionError(
+                f"Unexpected input tag '{tag}' for {type(vt).__name__} -- "
+                f"is_reuse_eligible should have rejected this"
+            )
+
+    # Collect all guards for sources accessed during the subgraph trace
+    # and for the flattened arg sources.
+    all_sources = set(traced_sources)
+    all_sources.update(s for s in fingerprint.arg_sources if s is not None)
+    all_relevant_guards: set[Guard] = set()
+    for source in all_sources:
+        all_relevant_guards.update(tx.output.guards.get_guards_for_source(source))
+
+    guard_tuples: list[tuple[Source, GuardCheckSpec, object, Guard]] = []
+    for guard in all_relevant_guards:
+        source = guard.originating_source
+        type_str = guard.create_fn_name()
+        handler = GUARD_VALUE_DISPATCH.get(type_str)
+
+        if handler is SKIP_GUARD:
+            continue
+
+        if handler is None or isinstance(handler, UnsupportedGuardCheckSpec):
+            raise RuntimeError(
+                f"subgraph_reuse: unsupported guard type '{type_str}' on source '{source.name}'"
+            )
+
+        try:
+            value = tx.output.resolve_source_value(source)
+        except Exception:
+            raise RuntimeError(
+                f"subgraph_reuse: failed to resolve source '{source.name}' for {type_str} guard"
+            ) from None
+
+        # TODO(anijain2305): vLLM workaround -- skip CONSTANT_MATCH on
+        # strings. Re-evaluate once vLLM migrates off this pattern.
+        # if type_str == "CONSTANT_MATCH" and isinstance(value, str):
+        #     continue
+
+        handler = cast(GuardCheckSpec, handler)
+        expected = handler.get_metadata_fn(guard, value)
+        guard_tuples.append((source, handler, expected, guard))
+
+    hc_log.debug("Number of guards %s", len(guard_tuples))
+
+    return InvokeSubgraphReuseCondition(
+        input_checks=input_checks,
+        guards=guard_tuples,
+        treespec=fingerprint.treespec,
+        traced_sources=traced_sources,
+    )
+
+
+def build_source_replacement(
+    old_arg_sources: list[Source | None],
+    new_arg_sources: list[Source | None],
+) -> dict[Source, Source]:
+    """Map old arg sources to new arg sources for remapping captured variable sources."""
+    return {
+        old: new
+        for old, new in zip(old_arg_sources, new_arg_sources)
+        if old is not None and new is not None and old != new
+    }
+
+
+def is_reusable(
+    tx: "InstructionTranslator",
+    condition: "InvokeSubgraphReuseCondition",
+    fingerprint: InputFingerprint,
+    cached_entry: InvokeSubgraphReuseEntry,
+) -> bool:
+    """Check if a cached subgraph can be reused for the current call.
+
+    Three-phase check:
+    (1) Verify that intermediates (tensor metadata, symnode types, constant
+        values) match the cached input_checks — these are lightweight
+        structural comparisons that don't require source resolution.
+    (2) Check for mutations on the remapped traced_sources — if any source
+        the subgraph read has been mutated since the original trace, the
+        cached guards would evaluate against stale values.
+    (3) Build a source replacement mapping (old sources → new sources) and
+        re-evaluate the snapshotted guards under the new sources.
+    """
+    # Structural check: treespec must match first.
+    if condition.treespec is not None and fingerprint.treespec != condition.treespec:
+        hc_log.debug(
+            "subgraph_reuse: reuse failed -- treespec mismatch",
+        )
+        return False
+
+    # Input count, tags, and metadata must match.
+    # Tensor metadata (shape, stride, dtype, device, requires_grad) is checked
+    # here because TENSOR_MATCH guards for subgraph inputs typically already
+    # exist in the outer graph before tracing and thus won't appear in the
+    # guard delta.
+    if len(condition.input_checks) != len(fingerprint.flat_vts):
+        hc_log.debug(
+            "subgraph_reuse: reuse failed -- input count mismatch: cached %d vs current %d",
+            len(condition.input_checks),
+            len(fingerprint.flat_vts),
+        )
+        return False
+
+    for i, ((cached_tag, cached_val), (cur_tag, cur_vt)) in enumerate(
+        zip(condition.input_checks, fingerprint.flat_vts)
+    ):
+        if cached_tag != cur_tag:
+            hc_log.debug(
+                "subgraph_reuse: reuse failed -- input %d tag mismatch: cached '%s' vs current '%s'",
+                i,
+                cached_tag,
+                cur_tag,
+            )
+            return False
+        if cached_tag == InputTag.TENSOR:
+            assert isinstance(cur_vt, TensorVariable)
+            example = cur_vt.proxy.node.meta.get("example_value", None)
+            if example is None:
+                hc_log.debug(
+                    "subgraph_reuse: reuse failed -- input %d tensor has no example_value",
+                    i,
+                )
+                return False
+            cur_meta = extract_tensor_metadata(example)
+            if cur_meta != cached_val:
+                hc_log.debug(
+                    "subgraph_reuse: reuse failed -- input %d tensor metadata mismatch",
+                    i,
+                )
+                return False
+        elif cached_tag == InputTag.SYMNODE:
+            assert isinstance(cur_vt, SymNodeVariable)
+            if cur_vt.sym_num is not cached_val:
+                return False
+        elif cached_tag == InputTag.CONSTANT:
+            assert isinstance(cur_vt, ConstantVariable)
+            if cur_vt.value != cached_val:
+                # If both the cached and current arg have sources, source
+                # replacement in stamp_out will resolve the correct value.
+                cached_src = (
+                    cached_entry.arg_sources[i]
+                    if i < len(cached_entry.arg_sources)
+                    else None
+                )
+                new_src = (
+                    fingerprint.arg_sources[i]
+                    if i < len(fingerprint.arg_sources)
+                    else None
+                )
+                if cached_src is None or new_src is None:
+                    return False
+
+    source_replacement = build_source_replacement(
+        cached_entry.arg_sources, fingerprint.arg_sources
+    )
+
+    # Parameterized source - this function gives you new sources parameterized
+    # on the arg_sources. For example, if the input to the nested compile region
+    # is a nn Module layer with source `layers[0]`, then old source
+    # `layers[0].weight` gets remapped to `layers[1].weight`. This
+    # parameterization is central in getting the new sources and then running
+    # guards on them.
+    def replacement_fn(s: Source) -> Source:
+        return source_replacement.get(s, s)
+
+    # Check for mutations on remapped traced_sources.
+    if source_replacement:
+        remapped = OrderedSet(s.clone(replacement_fn) for s in condition.traced_sources)
+    else:
+        remapped = condition.traced_sources
+    if has_mutated_vars(tx, remapped):
+        return False
+
+    # If no sources changed, all guards were already checked during the
+    # original trace and will trivially pass again.
+    if not source_replacement:
+        return True
+
+    # Shared resolution context so source.get_value memoizes intermediate
+    # results (e.g. common base sources) across all guards in this check.
+    resolve_globals: dict[str, Any] = {
+        "G": tx.output.root_tx.f_globals,
+        "L": tx.output.root_tx.f_locals,
+    }
+    resolve_locals: dict[str, Any] = {}
+    resolve_cache: dict[Source, Any] = {}
+
+    for source, handler, expected, guard in condition.guards:
+        new_source = source.clone(replacement_fn)
+        # Source unchanged after replacement — guard already passed during
+        # the original trace, skip re-evaluation.
+        if new_source == source:
+            continue
+
+        try:
+            value = new_source.get_value(resolve_globals, resolve_locals, resolve_cache)
+        except Exception:
+            hc_log.debug(
+                "subgraph_reuse: reuse failed -- cannot resolve source\n"
+                "  guard type: %s\n"
+                "  guard source: %s\n"
+                "  guard source name: %s\n"
+                "  user stack:\n%s",
+                guard.create_fn_name(),
+                new_source,
+                new_source.name,
+                "".join(guard.user_stack.format())
+                if guard.user_stack
+                else "<no stack>",
+            )
+            return False
+
+        if not handler.eval_fn(value, expected):
+            hc_log.debug(
+                "subgraph_reuse: reuse failed --\n"
+                "  guard type: %s\n"
+                "  guard source: %s\n"
+                "  guard source name: %s\n"
+                "  expected: %s\n"
+                "  got: %s\n"
+                "  user stack:\n%s",
+                guard.create_fn_name(),
+                new_source,
+                new_source.name,
+                expected,
+                value,
+                "".join(guard.user_stack.format())
+                if guard.user_stack
+                else "<no stack>",
+            )
+            return False
+
+    return True
+
+
+def has_reuse_entries(
+    tx: "InstructionTranslator",
+    fn_var: Any,
+) -> bool:
+    """Cheap check: does the cache have any entries for this function?"""
+    from torch._guards import InvokeSubgraphCache
+
+    invoke_subgraph_cache = tx.output.tracing_context.hop_dispatch_set_cache.get_cache(
+        torch._higher_order_ops.invoke_subgraph
+    )
+    if not isinstance(invoke_subgraph_cache, InvokeSubgraphCache):
+        return False
+    fn_code = get_fn_code(fn_var)
+    return fn_code is not None and fn_code in invoke_subgraph_cache.subgraph_reuse_cache
+
+
+def find_reuse_match(
+    tx: "InstructionTranslator",
+    fn_var: Any,
+    fingerprint: InputFingerprint,
+) -> InvokeSubgraphReuseEntry | None:
+    from torch._guards import InvokeSubgraphCache
+
+    invoke_subgraph_cache = tx.output.tracing_context.hop_dispatch_set_cache.get_cache(
+        torch._higher_order_ops.invoke_subgraph
+    )
+    if not isinstance(invoke_subgraph_cache, InvokeSubgraphCache):
+        return None
+    fn_code = get_fn_code(fn_var)
+    if fn_code is None:
+        return None
+
+    # this evaluator function is called one by one for all the invoke subgraph
+    # reuse entries - the one that evaluates to True is stamped out in the
+    # graph.
+    def evaluator(
+        cond: "InvokeSubgraphReuseCondition", entry: InvokeSubgraphReuseEntry
+    ) -> bool:
+        return is_reusable(tx, cond, fingerprint, entry)
+
+    return invoke_subgraph_cache.find_reuse_entry(fn_code, evaluator)
+
+
+def save_reuse_entry(
+    tx: "InstructionTranslator",
+    fn_var: Any,
+    fingerprint: InputFingerprint,
+    body_name: str,
+    body_gmod: torch.fx.GraphModule,
+    config: NestedCompileRegionOptions | None,
+    p_args: tuple[Any, ...],
+    body_r: VariableTracker,
+    example_value: Any,
+    max_reuse_entries: int = 8,
+    condition: "InvokeSubgraphReuseCondition | None" = None,
+    hash_key: int | None = None,
+) -> None:
+    """Save a traced subgraph into the reuse cache for future cache hits.
+
+    Builds an InvokeSubgraphReuseEntry with the freevar mapping (how each
+    lifted arg maps back to user inputs or captured variables), output
+    metadata, and arg sources. On a future cache hit, stamp_out_subgraph
+    uses this entry to emit a new invoke_subgraph call without re-tracing.
+
+    Exactly one of ``condition`` or ``hash_key`` must be provided.
+    ``condition`` stores the entry in the guard-based cache (linear scan);
+    ``hash_key`` stores it in the hash-key cache (O(1) lookup).
+    """
+    from torch._guards import InvokeSubgraphCache
+
+    assert (condition is None) != (hash_key is None), (
+        "Exactly one of condition or hash_key must be provided"
+    )
+
+    invoke_subgraph_cache = tx.output.tracing_context.hop_dispatch_set_cache.get_cache(
+        torch._higher_order_ops.invoke_subgraph
+    )
+    if not isinstance(invoke_subgraph_cache, InvokeSubgraphCache):
+        return
+
+    fn_code = get_fn_code(fn_var)
+    if fn_code is None:
+        return
+
+    subgraph_input_mapping = build_subgraph_input_mapping(
+        tx, p_args, fingerprint.flat_vts
+    )
+    single_tensor_output = isinstance(body_r, TensorVariable)
+
+    # Count user-visible outputs from body_r. The graph may have additional
+    # outputs from side-effect intermediates that stamp_out_subgraph must
+    # not include when reconstructing the user-visible return value.
+    user_output_vts: list[VariableTracker] = []
+    VariableTracker.visit(
+        lambda vt: user_output_vts.append(vt)
+        if vt.is_tensor() or isinstance(vt, SymNodeVariable)
+        else None,
+        body_r,
+    )
+    num_user_outputs = len(user_output_vts)
+
+    # Cache output tensor metadata so we can construct fresh FakeTensors on
+    # cache hit without re-running the subgraph. This is safe because
+    # invoke_subgraph does not support aliasing between inputs and outputs
+    # (speculate_subgraph will fail if that happens).
+    # example_value may contain SymInts (e.g. shape values for backward);
+    # only record metadata for actual tensors.
+    output_metadata = [
+        (t.shape, t.stride(), t.dtype, t.device, t.requires_grad)
+        for t in example_value
+        if isinstance(t, torch.Tensor)
+    ]
+
+    entry = InvokeSubgraphReuseEntry(
+        body_name=body_name,
+        body_gmod=body_gmod,
+        config=config,
+        subgraph_input_mapping=subgraph_input_mapping,
+        single_tensor_output=single_tensor_output,
+        output_metadata=output_metadata,
+        # Record arg sources so that on cache hit we can build a
+        # source replacement mapping (old sources → new sources) to
+        # rewrite captured variable sources for the current invocation.
+        arg_sources=fingerprint.arg_sources,
+        num_user_outputs=num_user_outputs,
+    )
+    if condition is not None:
+        invoke_subgraph_cache.add_reuse_entry(
+            fn_code, condition, entry, max_reuse_entries
+        )
+    else:
+        assert hash_key is not None
+        invoke_subgraph_cache.add_reuse_entry_by_key(
+            fn_code, hash_key, entry, max_reuse_entries
+        )
+
+
+def trace_reuse_hash_fn(
+    tx: "InstructionTranslator",
+    reuse_hash_fn: Any,
+    fn_args_vt: "Sequence[VariableTracker]",
+    kwargs: dict[str, VariableTracker],
+) -> int:
+    """Trace the user's reuse_hash_fn to get a constant integer hash key.
+
+    Guards installed during the hash function tracing are skipped — the hash
+    key itself is the reuse condition, not the guards.
+    """
+    from torch._dynamo.exc import Unsupported
+    from torch._dynamo.utils import _make_inlined
+
+    with tx.output.tracing_context.guards_context.skip_guard_install():
+        try:
+            result = _make_inlined(tx, reuse_hash_fn)(*fn_args_vt, **kwargs)
+        except Unsupported as e:
+            raise RuntimeError(
+                f"reuse_hash_fn must be fully traceable without graph breaks. Got: {e}"
+            ) from e
+
+    if not isinstance(result, ConstantVariable) or not isinstance(result.value, int):
+        raise RuntimeError(
+            f"reuse_hash_fn must return a constant integer, got {result}"
+        )
+
+    return result.value
+
+
+def find_reuse_entry_by_key(
+    tx: "InstructionTranslator",
+    fn_var: Any,
+    hash_key: int,
+) -> InvokeSubgraphReuseEntry | None:
+    from torch._guards import InvokeSubgraphCache
+
+    invoke_subgraph_cache = tx.output.tracing_context.hop_dispatch_set_cache.get_cache(
+        torch._higher_order_ops.invoke_subgraph
+    )
+    if not isinstance(invoke_subgraph_cache, InvokeSubgraphCache):
+        return None
+    fn_code = get_fn_code(fn_var)
+    if fn_code is None:
+        return None
+    return invoke_subgraph_cache.find_reuse_entry_by_key(fn_code, hash_key)
+
+
+def stamp_out_subgraph(
+    tx: "InstructionTranslator",
+    fingerprint: InputFingerprint,
+    cached: InvokeSubgraphReuseEntry,
+) -> VariableTracker:
+    """Emit a new invoke_subgraph call by stamping out a cached subgraph.
+
+    Sources in the cached entry are parameterized: they refer to the original
+    call's sources and must be rewritten to the current call's sources via
+    source replacement before we can look up or create the corresponding
+    graph placeholders.
+    """
+    from torch._dynamo.variables.builder import VariableBuilder
+    from torch._dynamo.variables.higher_order_ops import add_call_function, make_attr
+
+    flat_proxies = get_flat_proxies(fingerprint)
+    new_arg_sources = fingerprint.arg_sources
+
+    source_replacement = build_source_replacement(cached.arg_sources, new_arg_sources)
+
+    new_lifted_args = []
+    # Shared resolution context so get_value memoizes intermediate results
+    # (e.g. L['self'].layers) across all freevars in this stamp-out.
+    resolve_globals: dict[str, Any] = {
+        "G": tx.output.root_tx.f_globals,
+        "L": tx.output.root_tx.f_locals,
+    }
+    resolve_locals: dict[str, Any] = {}
+    resolve_cache: dict[Source, Any] = {}
+
+    # Find the args for the about-to-be-inserted invoke_subgraph call.
+    for subgraph_input in cached.subgraph_input_mapping:
+        if isinstance(subgraph_input, LiftedUserArg):
+            new_lifted_args.append(flat_proxies[subgraph_input.index])
+        elif isinstance(subgraph_input, LiftedBoundSymbol):
+            from torch._dynamo.output_graph import LazyProxy
+
+            proxy = tx.output.current_tracer.bound_symbols[subgraph_input.expr]
+            if isinstance(proxy, LazyProxy):
+                proxy = proxy()
+                tx.output.current_tracer.bound_symbols[subgraph_input.expr] = proxy
+            new_lifted_args.append(proxy)
+        elif isinstance(subgraph_input, LiftedSyntheticObject):
+            ctor_args = subgraph_input.ctor_args
+            ctor_arg_sources = subgraph_input.ctor_arg_sources
+            if ctor_arg_sources and source_replacement:
+                new_ctor_args = []
+                new_ctor_arg_sources = []
+                for val, arg_src in zip(ctor_args, ctor_arg_sources):
+                    if arg_src is not None:
+                        new_src = arg_src.clone(lambda s: source_replacement.get(s, s))
+                        val = new_src.get_value(
+                            resolve_globals, resolve_locals, resolve_cache
+                        )
+                        arg_src = new_src
+                    new_ctor_args.append(val)
+                    new_ctor_arg_sources.append(arg_src)
+                ctor_args = tuple(new_ctor_args)
+                ctor_arg_sources = tuple(new_ctor_arg_sources)
+            vt = tx.output.synthetic_graph_input(
+                subgraph_input.ctor_fn, ctor_args, ctor_arg_sources
+            )
+            new_lifted_args.append(vt.as_proxy())
+        elif isinstance(subgraph_input, LiftedCapturedSource):
+            new_source = subgraph_input.source
+            if source_replacement:
+                new_source = new_source.clone(lambda s: source_replacement.get(s, s))
+            # VariableBuilder deduplicates via input_source_to_var,
+            # so this reuses existing graph placeholders automatically.
+            value = new_source.get_value(resolve_globals, resolve_locals, resolve_cache)
+            vt = VariableBuilder(tx, new_source)(value)
+            new_lifted_args.append(vt.as_proxy())
+
+    # Generate fake tensor outputs
+    assert tx.fake_mode is not None
+    with tx.fake_mode:
+        example_value = tuple(
+            torch.empty_strided(
+                shape,
+                stride,
+                dtype=dtype,
+                device=device,
+                requires_grad=req_grad,
+            )
+            for shape, stride, dtype, device, req_grad in cached.output_metadata
+        )
+
+    # Install the invoke_subgraph call
+    body_node = make_attr(tx, cached.body_name)
+    p_args = (body_node, cached.body_name, *new_lifted_args)
+    flat_variable = add_call_function(
+        tx,
+        torch._higher_order_ops.invoke_subgraph,
+        tuple(p_args),
+        {},
+        example_value,
+        cached.config,
+    )
+
+    # Return only the user-visible outputs. The graph may have extra
+    # intermediate outputs from side effects (allow_side_effects=True)
+    # that should not be part of the user-facing return value.
+    if cached.single_tensor_output:
+        items = flat_variable.items  # pyrefly: ignore[missing-attribute]
+        assert isinstance(items[0], TensorVariable), (
+            f"Expected tensor output but got {type(items[0]).__name__}"
+        )
+        return items[0]
+
+    items = flat_variable.items  # pyrefly: ignore[missing-attribute]
+    n = cached.num_user_outputs
+    if n > 0 and n < len(items):
+        from .builder import SourcelessBuilder
+
+        return SourcelessBuilder.create(tx, tuple(items[:n]))
+    return flat_variable
+
+
+def build_subgraph_input_mapping(
+    tx: "InstructionTranslator",
+    p_args: tuple[Any, ...],
+    flat_vts: list[tuple[InputTag, VariableTracker]],
+) -> list[LiftedArgOrigin]:
+    """Build a mapping that records the origin of each lifted arg for a subgraph.
+
+    On a cache hit, we stamp out a new invoke_subgraph call and need to
+    reconstruct its argument list in the correct order. Each lifted arg
+    (p_args[2:], skipping body_node and body_name) comes from one of:
+
+    - LiftedUserArg: a user argument (intermediate activation or explicit input)
+    - LiftedCapturedSource: a captured variable (e.g. a weight or parameter)
+    - LiftedSyntheticObject: a TorchScriptObject with a SyntheticLocalSource
+    - LiftedBoundSymbol: a SymInt already bound as a graph input
+    """
+    proxy_node_to_idx: dict[torch.fx.Node, int] = {}
+    idx = 0
+    for tag, vt in flat_vts:
+        if tag in (InputTag.TENSOR, InputTag.SYMNODE):
+            node = vt.as_proxy().node
+            if node not in proxy_node_to_idx:
+                proxy_node_to_idx[node] = idx
+                idx += 1
+
+    subgraph_input_mapping: list[LiftedArgOrigin] = []
+    for outer_proxy in p_args[2:]:
+        matched_idx = proxy_node_to_idx.get(outer_proxy.node, -1)
+        if matched_idx >= 0:
+            subgraph_input_mapping.append(LiftedUserArg(matched_idx))
+        else:
+            grapharg = outer_proxy.node.meta.get("grapharg", None)
+            source = grapharg.source if grapharg is not None else None
+            # SymInt freevars must reuse the existing symbolic proxy rather
+            # than resolving via source.get_value() (which returns the
+            # concrete int). They appear as either:
+            # - placeholder nodes with grapharg.example being a SymInt
+            # - call_function nodes (e.g. sym_size_int) with no grapharg
+            # In both cases, store the sympy expression and look it up in
+            # bound_symbols during stamp-out.
+            example = (
+                grapharg.example
+                if grapharg is not None
+                else outer_proxy.node.meta.get("example_value", None)
+            )
+            if isinstance(example, torch.SymInt):
+                subgraph_input_mapping.append(LiftedBoundSymbol(example.node.expr))
+                continue
+            assert source is not None, (
+                f"Freevar has no source: node.op={outer_proxy.node.op} "
+                f"node.name={outer_proxy.node.name} -- this likely means a "
+                f"function argument was not included in the proxy matching"
+            )
+            if isinstance(source, SyntheticLocalSource):
+                ctor_info = tx.output.synthetic_source_ctor_info.get(source)
+                if ctor_info is not None:
+                    ctor_fn, ctor_args, ctor_arg_sources = ctor_info
+                    subgraph_input_mapping.append(
+                        LiftedSyntheticObject(ctor_fn, ctor_args, ctor_arg_sources)
+                    )
+                    continue
+            subgraph_input_mapping.append(LiftedCapturedSource(source))
+    return subgraph_input_mapping
+
+
+class InvokeSubgraphHigherOrderVariable(WrapHigherOrderVariable):
+    _HOP_NAME = "torch.ops.higher_order.invoke_subgraph"
+    _ALLOW_FALLBACK_TO_EAGER = False
+    supports_input_mutation = True
+    supports_aliasing = False
+    allow_side_effects = True
+    # invoke_subgraph is NOT desugared in AOTAutograd, so the HOP input/output
+    # shouldn't alias. For checkpoint HOP, we inline it so we don't need
+    # alias analysis as functionalization would just work on the flat graph.
+    filter_aliased_intermediates = True
+
+    # pyrefly: ignore[bad-override]
+    def install_subgraph_in_output_graph(
+        self,
+        tx: "InstructionTranslator",
+        fn_vt: VariableTracker,
+        fn_args_vt: "Sequence[VariableTracker]",
+        kwargs: dict[str, VariableTracker],
+        body_gmod: GraphModule,
+        attr_name: str,
+    ) -> str:
+        # Check if the subgraph from speculate_subgraph (body_gmod) and the fake
+        # inputs have already been seen before. If yes, the subgraph is already
+        # installed in the output graph and we can just access the subgraph
+        # using the saved attr name.
+
+        if not isinstance(fn_vt, (UnspecializedNNModuleVariable, UserFunctionVariable)):
+            unimplemented(
+                gb_type="Encountered non user function variable during invoke_subgraph HOP tracing",
+                context=str(fn_vt),
+                explanation="invoke_subgraph does not support non user function variable",
+                hints=[*graph_break_hints.SUPPORTABLE],
+            )
+
+        invoke_subgraph_cache = (
+            tx.output.tracing_context.hop_dispatch_set_cache.get_cache(
+                torch._higher_order_ops.invoke_subgraph
+            )
+        )
+
+        if isinstance(fn_vt, UserFunctionVariable):
+            fn_code = fn_vt.get_function().__code__
+            fn_name = fn_vt.get_function().__name__
+        else:
+            assert isinstance(fn_vt, UnspecializedNNModuleVariable)
+            fn_code = fn_vt.value.forward.__func__.__code__  # type: ignore[attr-defined]
+            fn_name = fn_vt.value.forward.__name__  # type: ignore[attr-defined]
+        # pyrefly: ignore [implicit-any]
+        previously_installed_submodules = []
+        if invoke_subgraph_cache:
+            previously_installed_submodules = (
+                invoke_subgraph_cache.get_dynamo_installed_submodules(fn_code)
+            )
+            current_mod = body_gmod
+            # NB - reverse is more likely to cause a hit sooner because first
+            # graph can have requires_grad=False for a few inputs
+            for submodule_name in reversed(previously_installed_submodules):
+                assert submodule_name in tx.output.nn_modules
+                previous_mod = tx.output.nn_modules[submodule_name]
+                assert tx.fake_mode
+                from torch._dynamo.variables.higher_order_ops import (
+                    are_same_graph_modules,
+                )
+
+                if are_same_graph_modules(
+                    fn_name, previous_mod, current_mod, tx.fake_mode
+                ):
+                    return submodule_name
+
+        body_name = super().install_subgraph_in_output_graph(
+            tx, fn_vt, fn_args_vt, kwargs, body_gmod, "subgraph"
+        )
+        hc_log.debug(
+            "%s: Installing subgraph with identifier '%s', bringing total count for '%s' function to %s",
+            fn_name,
+            body_name,
+            fn_name,
+            len(previously_installed_submodules) + 1,
+        )
+        if invoke_subgraph_cache:
+            invoke_subgraph_cache.add_dynamo_installed_submodule(fn_code, body_name)
+
+        return body_name
+
+    def _call_function(
+        self,
+        tx: "InstructionTranslator",
+        args: "Sequence[VariableTracker]",
+        kwargs: dict[str, VariableTracker],
+    ) -> VariableTracker:
+        from torch._dynamo.utils import dynamo_timed
+        from torch._dynamo.variables.higher_order_ops import (
+            _call_function_with_auto_output_flattening,
+        )
+
+        fn_var = args[0]
+        fn_args_vt = args[1:]
+
+        config = None
+        max_reuse_entries = 8
+        reuse_hash_fn = None
+        if hasattr(fn_var, "get_function"):
+            try:
+                fn = fn_var.get_function()
+                config = getattr(fn, "__marked_compile_region_config__", None)
+                max_reuse_entries = getattr(
+                    fn, "__marked_compile_region_max_reuse_entries__", 8
+                )
+                reuse_hash_fn = getattr(
+                    fn, "__marked_compile_region_reuse_hash_fn__", None
+                )
+            except Exception:
+                log.warning(
+                    "Failed to extract nested_compile_region() config from InvokeSubgraphHigherOrderVariable. ",
+                    exc_info=True,
+                )
+                raise
+
+        # TODO (anijain2305) - Collect issues why this does not work for export,
+        # and enable if request arises.
+        reuse = not tx.output.export
+
+        # User-provided reuse_hash_fn path: hash key determines cache lookup.
+        if reuse and reuse_hash_fn is not None:
+            with dynamo_timed("invoke_subgraph_reuse_hash_fn"):
+                hash_key = trace_reuse_hash_fn(tx, reuse_hash_fn, fn_args_vt, kwargs)
+
+            cached = find_reuse_entry_by_key(tx, fn_var, hash_key)
+            if cached is not None:
+                hc_log.debug(
+                    "subgraph_reuse: hash key %d hit for '%s', reusing subgraph '%s'",
+                    hash_key,
+                    fn_var,
+                    cached.body_name,
+                )
+                fingerprint = build_input_fingerprint(tx, fn_args_vt, kwargs)
+                with dynamo_timed("invoke_subgraph_reuse_stamp_out"):
+                    return stamp_out_subgraph(tx, fingerprint, cached)
+
+        # Automatic reuse lookup (guard-based): check fn_code first (cheap) to
+        # avoid the expensive pytree flatten in build_input_fingerprint on
+        # the first call when there's nothing in the cache yet.
+        elif reuse and has_reuse_entries(tx, fn_var):
+            with dynamo_timed("invoke_subgraph_reuse_lookup"):
+                fingerprint = build_input_fingerprint(tx, fn_args_vt, kwargs)
+                match = find_reuse_match(
+                    tx,
+                    fn_var,
+                    fingerprint,
+                )
+            if match is not None:
+                hc_log.debug(
+                    "subgraph_reuse: cache hit for '%s', reusing subgraph '%s'",
+                    fn_var,
+                    match.body_name,
+                )
+                with dynamo_timed("invoke_subgraph_reuse_stamp_out"):
+                    return stamp_out_subgraph(tx, fingerprint, match)
+
+        assert self._HOP_NAME is not None
+        with dynamo_timed("invoke_subgraph_trace"):
+            (
+                p_args,
+                p_kwargs,
+                example_value,
+                body_r,
+                body_gmod,
+                body_name,
+                body_graph_output_vts,
+                tracing_info,
+            ) = self.create_wrapped_node(tx, fn_var, fn_args_vt, kwargs, self._HOP_NAME)
+
+        if len(p_kwargs) > 0:
+            unimplemented(
+                gb_type="invoke_subgraph: kwargs unexpected",
+                context=f"args: {args}, kwargs: {kwargs}",
+                explanation="kwargs should have been flattened into lifted args.",
+                hints=[
+                    *graph_break_hints.DYNAMO_BUG,
+                ],
+            )
+
+        # Store config in the body graph module meta
+        if isinstance(config, NestedCompileRegionOptions):
+            body_gmod.meta["nested_region_config"] = config
+
+        p_args = (
+            p_args[0],
+            body_name,
+            *p_args[1:],
+        )
+
+        # Subgraph reuse: save entry for future cache hits
+        if reuse:
+            fingerprint = build_input_fingerprint(tx, fn_args_vt, kwargs)
+            if reuse_hash_fn is not None:
+                traced_sources = tracing_info.traced_sources
+                if not is_reuse_eligible(
+                    tx,
+                    body_r,
+                    fingerprint,
+                    tracing_info,
+                    traced_sources,
+                    has_reuse_hash_fn=True,
+                ):
+                    raise RuntimeError(
+                        "reuse_hash_fn was provided but the subgraph is not "
+                        "eligible for reuse. Check the logs with "
+                        "TORCH_LOGS='+hierarchical_compile' for details."
+                    )
+                save_reuse_entry(
+                    tx,
+                    fn_var,
+                    fingerprint,
+                    body_name,
+                    body_gmod,
+                    config,
+                    p_args,
+                    body_r,
+                    example_value,
+                    max_reuse_entries,
+                    hash_key=hash_key,  # type: ignore[possibly-undefined]
+                )
+            else:
+                traced_sources = tracing_info.traced_sources
+                if is_reuse_eligible(
+                    tx, body_r, fingerprint, tracing_info, traced_sources
+                ):
+                    condition = build_reuse_condition(
+                        tx,
+                        fingerprint,
+                        traced_sources,
+                    )
+                    if condition is not None:
+                        save_reuse_entry(
+                            tx,
+                            fn_var,
+                            fingerprint,
+                            body_name,
+                            body_gmod,
+                            config,
+                            p_args,
+                            body_r,
+                            example_value,
+                            max_reuse_entries,
+                            condition=condition,
+                        )
+
+        return _call_function_with_auto_output_flattening(  # type: ignore[return-value]
+            tx,
+            torch._higher_order_ops.invoke_subgraph,
+            tuple(p_args),
+            p_kwargs,
+            example_value,
+            body_r,
+            body_graph_output_vts,
+            config=config,
+        )
diff --git a/torch/_dynamo/variables/iter.py b/torch/_dynamo/variables/iter.py
index 4a3c0247add1b..f4c3ebba8ddf0 100644
--- a/torch/_dynamo/variables/iter.py
+++ b/torch/_dynamo/variables/iter.py
@@ -16,7 +16,7 @@
 import itertools
 import sys
 from collections.abc import Callable, Sequence
-from typing import Any, TYPE_CHECKING, Union
+from typing import Any, TYPE_CHECKING
 
 from .. import graph_break_hints, polyfills, variables
 from ..bytecode_transformation import (
@@ -55,6 +55,9 @@ def __repr__(self) -> str:
     def as_python_constant(self) -> Any:
         return self.value
 
+    def get_real_python_backed_value(self) -> Any:
+        return self.value
+
     def call_function(
         self,
         tx: "InstructionTranslator",
@@ -158,7 +161,6 @@ def keyfunc(x: VariableTracker) -> Any:
 
             result = []
             try:
-                # pyrefly: ignore [unbound-name]
                 for k, v in itertools.groupby(seq, key=keyfunc):
                     result.append(
                         variables.TupleVariable(
@@ -196,10 +198,20 @@ def keyfunc(x: VariableTracker) -> Any:
             return tx.inline_user_function_return(
                 VariableTracker.build(tx, polyfills.repeat), args, kwargs
             )
-        elif self.value is itertools.count:
-            return variables.CountIteratorVariable(
-                *args, mutation_type=ValueMutationNew()
-            )
+        elif self.value is itertools.count and not kwargs:
+            if len(args) == 0:
+                return variables.CountIteratorVariable(mutation_type=ValueMutationNew())
+            if len(args) == 1:
+                return variables.CountIteratorVariable(
+                    item=args[0], mutation_type=ValueMutationNew()
+                )
+            if len(args) == 2:
+                return variables.CountIteratorVariable(
+                    item=args[0],
+                    step=args[1],
+                    mutation_type=ValueMutationNew(),
+                )
+            return super().call_function(tx, args, kwargs)
         elif (
             self.value is itertools.permutations
             and (len(args) == 1 or (len(args) == 2 and args[1].is_python_constant()))
@@ -319,6 +331,9 @@ def __init__(self, item: VariableTracker, **kwargs: Any) -> None:
         super().__init__(**kwargs)
         self.item = item
 
+    def python_type(self) -> type:
+        return itertools.repeat
+
     # Repeat needs no mutation, clone self
     def next_variable(self, tx: "InstructionTranslator") -> VariableTracker:
         return self.item
@@ -337,10 +352,21 @@ def reconstruct(self, codegen: "PyCodegen") -> None:
 
 
 class CountIteratorVariable(IteratorVariable):
+    # advance_count tracks how many next() calls were made during tracing,
+    # used by side_effects.py to replay them on the real iterator post-execution.
+    _nonvar_fields = {
+        "advance_count",
+        *IteratorVariable._nonvar_fields,
+    }
+
+    def python_type(self) -> type:
+        return itertools.count
+
     def __init__(
         self,
-        item: Union[int, VariableTracker] = 0,
-        step: Union[int, VariableTracker] = 1,
+        item: int | VariableTracker = 0,
+        step: int | VariableTracker = 1,
+        advance_count: int = 0,
         **kwargs: Any,
     ) -> None:
         super().__init__(**kwargs)
@@ -350,12 +376,14 @@ def __init__(
             step = ConstantVariable.create(step)
         self.item = item
         self.step = step
+        self.advance_count = advance_count
 
     def next_variable(self, tx: "InstructionTranslator") -> VariableTracker:
         assert self.is_mutable()
         old_item = self.item
         tx.output.side_effects.mutation(self)
         self.item = self.item.call_method(tx, "__add__", [self.step], {})
+        self.advance_count += 1
         return old_item
 
     def reconstruct(self, codegen: "PyCodegen") -> None:
@@ -377,6 +405,9 @@ class ZipVariable(IteratorVariable):
     Represents zip(*iterables)
     """
 
+    # PyZip_Type: https://github.com/python/cpython/blob/v3.13.0/Python/bltinmodule.c#L3011
+    _cpython_type = zip
+
     _nonvar_fields = {
         "index",
         "strict",
@@ -429,7 +460,7 @@ def next_variable(self, tx: "InstructionTranslator") -> VariableTracker:
         args = []
 
         def get_item(
-            it: Union[list[VariableTracker], VariableTracker],
+            it: list[VariableTracker] | VariableTracker,
         ) -> VariableTracker:
             if isinstance(it, list):
                 if old_index >= len(it):
@@ -498,6 +529,9 @@ class MapVariable(ZipVariable):
     Represents map(fn, *iterables)
     """
 
+    # PyMap_Type: https://github.com/python/cpython/blob/v3.13.0/Python/bltinmodule.c#L1484
+    _cpython_type = map
+
     def __init__(
         self,
         fn: VariableTracker,
@@ -545,6 +579,9 @@ class FilterVariable(IteratorVariable):
     Represents filter(fn, iterable)
     """
 
+    # PyFilter_Type: https://github.com/python/cpython/blob/v3.13.0/Python/bltinmodule.c#L630
+    _cpython_type = filter
+
     _nonvar_fields = {
         "index",
         *IteratorVariable._nonvar_fields,
diff --git a/torch/_dynamo/variables/lazy.py b/torch/_dynamo/variables/lazy.py
index 3ec97c3a5cd33..def844adda73a 100644
--- a/torch/_dynamo/variables/lazy.py
+++ b/torch/_dynamo/variables/lazy.py
@@ -257,7 +257,7 @@ class LazyConstantVariable(LazyVariableTracker):
     supported_types = (int, float, bool, str)
 
     @staticmethod
-    def create(  # pyrefly: ignore[bad-override]
+    def create(
         value: Any,
         source: Any,
         **options: Any,
diff --git a/torch/_dynamo/variables/lists.py b/torch/_dynamo/variables/lists.py
index d0e569a873aca..4149d56f45076 100644
--- a/torch/_dynamo/variables/lists.py
+++ b/torch/_dynamo/variables/lists.py
@@ -22,6 +22,7 @@ class that handles its unique behaviors while integrating with Dynamo's
 
 import torch
 import torch.fx
+from torch.utils._pytree import SequenceKey
 
 from .. import graph_break_hints, polyfills, variables
 from ..bytecode_transformation import (
@@ -30,27 +31,24 @@ class that handles its unique behaviors while integrating with Dynamo's
     create_call_method,
     create_dup_top,
     create_instruction,
-    create_rot_n,
 )
-from ..exc import raise_observed_exception, unimplemented
-from ..source import AttrSource, NamedTupleFieldsSource
+from ..exc import raise_observed_exception, raise_type_error, unimplemented
+from ..source import AttrSource
 from ..utils import (
     cmp_name_to_op_mapping,
     cmp_name_to_op_str_mapping,
     get_fake_value,
     guard_if_dyn,
     iter_contains,
-    namedtuple_fields,
     odict_values,
     raise_args_mismatch,
     range_iterator,
     set_example_value,
 )
 from .base import AsPythonConstantNotImplementedError, ValueMutationNew, VariableTracker
-from .constant import CONSTANT_VARIABLE_NONE, ConstantVariable
+from .constant import ConstantVariable
 from .functions import UserFunctionVariable
 from .iter import IteratorVariable
-from .user_defined import UserDefinedTupleVariable
 
 
 if TYPE_CHECKING:
@@ -112,6 +110,9 @@ def as_proxy(self) -> Any:
     def getitem_const(
         self, tx: "InstructionTranslator", arg: VariableTracker
     ) -> VariableTracker:
+        # TODO(follow-up): this assumes the caller (mp_subscript_impl) has already
+        # run _PyIndex_Check → nb_index_impl. Direct callers bypassing
+        # mp_subscript_impl will skip that validation.
         from .tensor import SymNodeVariable
 
         if isinstance(arg, SymNodeVariable):
@@ -121,8 +122,9 @@ def getitem_const(
 
         if isinstance(index, slice):
             if index.step == 0:
-                msg = ConstantVariable.create("slice step cannot be zero")
-                raise_observed_exception(ValueError, tx, args=[msg])
+                raise_observed_exception(
+                    ValueError, tx, args=["slice step cannot be zero"]
+                )
             # Set source to None because slicing a list gives a new local
             return self.clone(
                 items=self.items[index],
@@ -141,6 +143,10 @@ def getitem_const(
     def unpack_var_sequence(self, tx: "InstructionTranslator") -> list[VariableTracker]:
         return list(self.items)
 
+    def sq_length(self, tx: "InstructionTranslator") -> VariableTracker:
+        """Sequence length for lists, tuples, and range objects."""
+        return VariableTracker.build(tx, len(self.items))
+
     def call_tree_map_branch(
         self,
         tx: "InstructionTranslator",
@@ -185,48 +191,90 @@ def call_tree_map_branch(
             mutation_type=ValueMutationNew(),
         )
 
-    def call_method(
+    def call_tree_map_with_path_branch(
         self,
         tx: "InstructionTranslator",
-        name: str,
-        args: list[VariableTracker],
-        kwargs: dict[str, VariableTracker],
+        tree_map_fn: UserFunctionVariable,
+        map_fn: VariableTracker,
+        rest: Sequence[VariableTracker],
+        tree_map_kwargs: dict[str, VariableTracker],
+        keypath: tuple[Any, ...],
     ) -> VariableTracker:
-        from .builder import SourcelessBuilder
+        if not isinstance(self, (ListVariable, TupleVariable)):
+            return self._tree_map_with_path_fallback(
+                tx, tree_map_fn, map_fn, rest, tree_map_kwargs, keypath
+            )
 
-        if name == "__getitem__":
-            if kwargs or len(args) != 1:
-                raise_args_mismatch(
+        other_lists: list[BaseListVariable] = []
+        for candidate in rest:
+            if (
+                not isinstance(candidate, BaseListVariable)
+                or len(candidate.items) != len(self.items)
+                or self.python_type() != candidate.python_type()
+            ):
+                return self._tree_map_with_path_fallback(
+                    tx, tree_map_fn, map_fn, rest, tree_map_kwargs, keypath
+                )
+            other_lists.append(candidate)
+
+        new_items: list[VariableTracker] = []
+        for idx, item in enumerate(self.items):
+            sibling_leaves = [candidate.items[idx] for candidate in other_lists]
+            child_keypath = keypath + (SequenceKey(idx),)
+            new_items.append(
+                item.call_tree_map_with_path(
                     tx,
-                    name,
-                    "1 args and 0 kwargs",
-                    f"{len(args)} args and {len(kwargs)} kwargs",
+                    tree_map_fn,
+                    map_fn,
+                    sibling_leaves,
+                    tree_map_kwargs,
+                    child_keypath,
                 )
+            )
 
-            if args[0].is_tensor():
-                value = get_fake_value(args[0].as_proxy().node, tx)
-                if value.constant is not None and value.constant.numel() == 1:
-                    value = variables.ConstantVariable.create(value.constant.item())
-                else:
-                    unimplemented(
-                        gb_type="Indexing list with non-scalar tensor",
-                        context=f"call_method {self} {name} {args} {kwargs}",
-                        explanation=(
-                            "Attempted to index list-like object with tensor with > 1 element."
-                        ),
-                        hints=[*graph_break_hints.USER_ERROR],
-                    )
-            else:
-                value = args[0]
+        return self.clone(
+            items=new_items,
+            source=None,
+            mutation_type=ValueMutationNew(),
+        )
 
-            if value.python_type() not in (int, slice):
-                msg = f"indices must be integers or slices, not {value.python_type()}"
+    def mp_subscript_impl(
+        self,
+        tx: "InstructionTranslator",
+        key: VariableTracker,
+    ) -> VariableTracker:
+        # list_subscript: https://github.com/python/cpython/blob/62a6e898e01/Objects/listobject.c#L3689-L3710
+        # _PyIndex_Check: https://github.com/python/cpython/blob/62a6e898e01/Include/internal/pycore_abstract.h#L13-L17
+        # TODO(follow-up): replace hasattr(key_type, "__index__") with
+        # has_slot(num_slots, PyNumberSlots.NB_INDEX) for C extension types.
+        try:
+            key_type = key.python_type()
+        except NotImplementedError:
+            key_type = None
+        if key_type not in (int, bool, slice):
+            if key_type is not None and not hasattr(key_type, "__index__"):
+                container_name = self.python_type_name()
                 raise_observed_exception(
-                    TypeError, tx, args=[SourcelessBuilder.create(tx, msg)]
+                    TypeError,
+                    tx,
+                    args=[
+                        f"{container_name} indices must be integers or slices, not {key.python_type_name()}"
+                    ],
                 )
+            key = key.nb_index_impl(tx)
+
+        return self.getitem_const(tx, key)
+
+    def call_method(
+        self,
+        tx: "InstructionTranslator",
+        name: str,
+        args: list[VariableTracker],
+        kwargs: dict[str, VariableTracker],
+    ) -> VariableTracker:
+        from .builder import SourcelessBuilder
 
-            return self.getitem_const(tx, value)
-        elif name == "__contains__":
+        if name == "__contains__":
             if kwargs or len(args) != 1:
                 raise_args_mismatch(
                     tx,
@@ -250,12 +298,14 @@ def call_method(
                 const_args = [arg.as_python_constant() for arg in args]
                 const_kwargs = {k: v.as_python_constant() for k, v in kwargs.items()}
                 try:
-                    return ConstantVariable.create(
-                        items.index(*const_args, **const_kwargs)
+                    return VariableTracker.build(
+                        tx, items.index(*const_args, **const_kwargs)
                     )
                 except ValueError:
                     raise_observed_exception(
-                        ValueError, tx, args=[ConstantVariable.create("tuple.index()")]
+                        ValueError,
+                        tx,
+                        args=["tuple.index()"],
                     )
             except AsPythonConstantNotImplementedError:
                 return tx.inline_user_function_return(
@@ -288,10 +338,13 @@ def call_method(
             if type(self) is not type(args[0]):
                 tp_name = self.python_type_name()
                 other = args[0].python_type_name()
-                msg_vt = ConstantVariable.create(
-                    f'can only concatenate {tp_name} (not "{other}") to {tp_name}'
+                raise_observed_exception(
+                    TypeError,
+                    tx,
+                    args=[
+                        f'can only concatenate {tp_name} (not "{other}") to {tp_name}'
+                    ],
                 )
-                raise_observed_exception(TypeError, tx, args=[msg_vt])
 
             if name == "__add__":
                 return type(self)(self.items + args[0].items, source=self.source)  # type: ignore[attr-defined]
@@ -308,10 +361,13 @@ def call_method(
                 )
 
             if not (args[0].is_python_constant() and args[0].python_type() is int):
-                msg_vt = ConstantVariable.create(
-                    f"can't multiply sequence by non-int type of '{args[0].python_type_name()}'"
+                raise_observed_exception(
+                    TypeError,
+                    tx,
+                    args=[
+                        f"can't multiply sequence by non-int type of '{args[0].python_type_name()}'"
+                    ],
                 )
-                raise_observed_exception(TypeError, tx, args=[msg_vt])
 
             val = args[0].as_python_constant()
 
@@ -350,8 +406,13 @@ def call_method(
                     op_str = cmp_name_to_op_str_mapping[name]
                     left_ty = left.python_type_name()
                     right_ty = right.python_type_name()
-                    msg = f"{op_str} not supported between instances of '{left_ty}' and '{right_ty}'"
-                    raise_observed_exception(TypeError, tx, args=[msg])
+                    raise_observed_exception(
+                        TypeError,
+                        tx,
+                        args=[
+                            f"{op_str} not supported between instances of '{left_ty}' and '{right_ty}'"
+                        ],
+                    )
 
             return SourcelessBuilder.create(tx, polyfills.list_cmp).call_function(
                 tx,
@@ -369,6 +430,9 @@ def call_method(
 
 
 class RangeVariable(BaseListVariable):
+    # PyRange_Type: https://github.com/python/cpython/blob/v3.13.0/Objects/rangeobject.c#L767
+    _cpython_type = range
+
     def __init__(self, items: Sequence[VariableTracker], **kwargs: Any) -> None:
         items_to_map = items
         start = variables.ConstantVariable.create(0)
@@ -481,8 +545,6 @@ def _get_slice_indices(self, length: int, slice: slice) -> list[int]:
         return [start, stop, step]
 
     def apply_index(self, tx: "InstructionTranslator", index: int) -> VariableTracker:
-        from .builder import SourcelessBuilder
-
         length = self.range_length()
         if index < 0:
             index = length + index
@@ -491,10 +553,10 @@ def apply_index(self, tx: "InstructionTranslator", index: int) -> VariableTracke
             raise_observed_exception(
                 IndexError,
                 tx,
-                args=[SourcelessBuilder.create(tx, "range object index out of range")],
+                args=["range object index out of range"],
             )
 
-        return variables.ConstantVariable.create(self.start() + (index * self.step()))
+        return VariableTracker.build(tx, self.start() + (index * self.step()))
 
     def apply_slice(self, slice: slice) -> "RangeVariable":
         (slice_start, slice_stop, slice_step) = self._get_slice_indices(
@@ -523,9 +585,10 @@ def as_python_constant(self) -> range:
     def getitem_const(
         self, tx: "InstructionTranslator", arg: VariableTracker
     ) -> VariableTracker:
-        from .builder import SourcelessBuilder
-
-        # implementations mimics https://github.com/python/cpython/blob/main/Objects/rangeobject.c
+        # range_subscript: https://github.com/python/cpython/blob/main/Objects/rangeobject.c
+        # TODO(follow-up): this assumes the caller (mp_subscript_impl) has already
+        # run _PyIndex_Check → nb_index_impl. Direct callers bypassing
+        # mp_subscript_impl will skip that validation.
         index = arg.as_python_constant()
 
         if isinstance(index, slice):
@@ -533,10 +596,9 @@ def getitem_const(
         elif isinstance(index, int):
             return self.apply_index(tx, index)
         else:
-            msg = SourcelessBuilder.create(
-                tx, "range indices must be integers or slices"
+            raise_observed_exception(
+                TypeError, tx, args=["range indices must be integers or slices"]
             )
-            raise_observed_exception(TypeError, tx, args=[msg])
 
     def as_proxy(self) -> range:
         return self.python_type()(*self._as_proxy())
@@ -546,6 +608,13 @@ def unpack_var_sequence(
     ) -> list[VariableTracker]:
         return [variables.ConstantVariable.create(x) for x in self.as_python_constant()]
 
+    def sq_length(self, tx: "InstructionTranslator") -> VariableTracker:
+        """Sequence length for range objects."""
+        length = self.range_length()
+        if length > sys.maxsize:
+            raise_observed_exception(OverflowError, tx)
+        return VariableTracker.build(tx, length)
+
     def reconstruct(self, codegen: "PyCodegen") -> None:
         assert "range" not in codegen.tx.f_globals
         codegen.add_push_null(
@@ -558,7 +627,7 @@ def call_obj_hasattr(
         self, tx: "InstructionTranslator", name: str
     ) -> ConstantVariable:
         if self.python_type() is range:
-            return variables.ConstantVariable.create(name in range.__dict__)
+            return VariableTracker.build(tx, name in range.__dict__)
         return super().call_obj_hasattr(tx, name)
 
     def range_equals(self, other: "RangeVariable") -> bool:
@@ -594,6 +663,29 @@ def range_count(self, x: VariableTracker) -> int:
             return int(re)
         return 0
 
+    def mp_subscript_impl(
+        self,
+        tx: "InstructionTranslator",
+        key: VariableTracker,
+    ) -> VariableTracker:
+        # range_subscript: https://github.com/python/cpython/blob/62a6e898e01/Objects/rangeobject.c#L729-L748
+        # CPython: range_subscript checks _PyIndex_Check → PyNumber_Index for non-slice keys
+        try:
+            key_type = key.python_type()
+        except NotImplementedError:
+            key_type = None
+        if key_type not in (int, bool, slice):
+            if key_type is not None and not hasattr(key_type, "__index__"):
+                raise_observed_exception(
+                    TypeError,
+                    tx,
+                    args=[
+                        f"range indices must be integers or slices, not {key.python_type_name()}"
+                    ],
+                )
+            key = key.nb_index_impl(tx)
+        return self.getitem_const(tx, key)
+
     def call_method(
         self,
         tx: "InstructionTranslator",
@@ -612,11 +704,6 @@ def call_method(
             return RangeIteratorVariable(
                 self.start(), self.stop(), self.step(), self.range_length()
             )
-        elif name == "__len__":
-            length = self.range_length()
-            if length > sys.maxsize:
-                raise_observed_exception(OverflowError, tx)
-            return ConstantVariable.create(self.range_length())
         elif name in ("count", "__contains__"):
             return SourcelessBuilder.create(tx, self.range_count(*args))
         elif name == "index":
@@ -624,28 +711,25 @@ def call_method(
             start, stop, step = self.start(), self.stop(), self.step()
             in_range = (start <= x < stop) if step > 0 else (stop < x <= start)
             if in_range and ((x - start) % step) == 0:
-                return ConstantVariable.create((x - start) // step)
+                return VariableTracker.build(tx, (x - start) // step)
             raise_observed_exception(
                 ValueError,
                 tx,
-                args=[ConstantVariable.create(f"{x} is not in range")],
+                args=[f"{x} is not in range"],
             )
-        elif name == "__getitem__":
-            return self.getitem_const(tx, *args)
         elif name in cmp_name_to_op_mapping:
             other = args[0]
             pt = other.python_type()
             if name not in ("__eq__", "__ne__"):
-                # ranges are only comparable to other ranges
                 msg = f"{name} not supported between instances of 'range' and '{pt}'"
                 raise_observed_exception(
                     TypeError,
                     tx,
-                    args=[ConstantVariable.create(msg)],
+                    args=[msg],
                 )
 
             if pt is not range:
-                return ConstantVariable.create(NotImplemented)
+                return VariableTracker.build(tx, NotImplemented)
 
             if isinstance(other, RangeVariable):
                 cmp = self.range_equals(other)
@@ -710,7 +794,7 @@ def call_method(
             (arg,) = args
             tx.output.side_effects.mutation(self)
             self.items.append(arg)
-            return CONSTANT_VARIABLE_NONE
+            return ConstantVariable.create(None)
         elif name == "extend" and self.is_mutable():
             if kwargs or len(args) != 1:
                 raise_args_mismatch(
@@ -721,14 +805,15 @@ def call_method(
                 )
 
             if not args[0].has_force_unpack_var_sequence(tx):
-                msg = ConstantVariable.create(f"{type(args[0])} object is not iterable")
-                raise_observed_exception(TypeError, tx, args=[msg])
+                raise_observed_exception(
+                    TypeError, tx, args=[f"{type(args[0])} object is not iterable"]
+                )
 
             (arg,) = args
             arg.force_apply_to_var_sequence(
                 tx, lambda item: self.call_method(tx, "append", [item], {})
             )
-            return CONSTANT_VARIABLE_NONE
+            return ConstantVariable.create(None)
         elif name == "insert" and self.is_mutable():
             if kwargs or len(args) != 2:
                 raise_args_mismatch(
@@ -745,7 +830,7 @@ def call_method(
             tx.output.side_effects.mutation(self)
             # type: ignore[arg-type]
             self.items.insert(const_idx, value)
-            return CONSTANT_VARIABLE_NONE
+            return ConstantVariable.create(None)
         elif name == "pop" and self.is_mutable():
             if kwargs or len(args) > 1:
                 raise_args_mismatch(
@@ -756,14 +841,14 @@ def call_method(
                 )
 
             if len(self.items) == 0:
-                msg = ConstantVariable.create("pop from empty list")
-                raise_observed_exception(IndexError, tx, args=[msg])
+                raise_observed_exception(IndexError, tx, args=["pop from empty list"])
 
             if len(args):
                 idx = args[0].as_python_constant()
-                if idx > len(self.items):
-                    msg = ConstantVariable.create("pop index out of range")
-                    raise_observed_exception(IndexError, tx, args=[msg])
+                if idx >= len(self.items):
+                    raise_observed_exception(
+                        IndexError, tx, args=["pop index out of range"]
+                    )
             tx.output.side_effects.mutation(self)
             return self.items.pop(*[a.as_python_constant() for a in args])
         elif name == "clear" and self.is_mutable():
@@ -776,7 +861,7 @@ def call_method(
                 )
             tx.output.side_effects.mutation(self)
             self.items.clear()
-            return CONSTANT_VARIABLE_NONE
+            return ConstantVariable.create(None)
         elif name == "__setitem__" and self.is_mutable() and args:
             # Realize args[0] to get the concrete type for proper type checking
             key = args[0].realize()
@@ -816,7 +901,7 @@ def call_method(
                     self.items[items_slice] = list(value.items)  # type: ignore[attr-defined]
             else:
                 self.items[key.as_python_constant()] = value
-            return CONSTANT_VARIABLE_NONE
+            return ConstantVariable.create(None)
         elif name == "__delitem__" and self.is_mutable():
             if kwargs or len(args) != 1:
                 raise_args_mismatch(
@@ -842,14 +927,12 @@ def call_method(
                     raise_observed_exception(
                         type(exc),
                         tx,
-                        args=list(map(ConstantVariable.create, exc.args)),
+                        args=list(exc.args),
                     )
             else:
-                msg = ConstantVariable.create(
-                    f"list indices must be integers or slices, not {args[0].python_type_name()}"
-                )
-                raise_observed_exception(TypeError, tx, args=[msg])
-            return CONSTANT_VARIABLE_NONE
+                msg = f"list indices must be integers or slices, not {args[0].python_type_name()}"
+                raise_type_error(tx, msg)
+            return ConstantVariable.create(None)
         elif name == "copy":
             # List copy() doesn't have args and kwargs
             if args or kwargs:
@@ -871,7 +954,7 @@ def call_method(
                 )
             self.items.reverse()
             tx.output.side_effects.mutation(self)
-            return CONSTANT_VARIABLE_NONE
+            return ConstantVariable.create(None)
         elif name == "remove" and self.is_mutable():
             if kwargs or len(args) != 1:
                 raise_args_mismatch(
@@ -883,12 +966,16 @@ def call_method(
 
             idx = self.call_method(tx, "index", args, kwargs)
             self.call_method(tx, "pop", [idx], {})
-            return CONSTANT_VARIABLE_NONE
+            return ConstantVariable.create(None)
         else:
             return super().call_method(tx, name, args, kwargs)
 
 
 class ListVariable(CommonListMethodsVariable):
+    # PyList_Type: https://github.com/python/cpython/blob/v3.13.0/Objects/listobject.c#L3776
+    _cpython_type = list
+    _has_instance_dict = False
+
     def python_type(self) -> type:
         return list
 
@@ -947,13 +1034,15 @@ def call_method(
             tx.output.side_effects.mutation(self)
             if isinstance(key, SliceVariable):
                 if not value.has_force_unpack_var_sequence(tx):
-                    msg = ConstantVariable.create("can only assign an iterable")
-                    raise_observed_exception(TypeError, tx, args=[msg])
+                    raise_observed_exception(
+                        TypeError, tx, args=["can only assign an iterable"]
+                    )
 
                 key_as_const = key.as_python_constant()
                 if key_as_const.step == 0:
-                    msg = ConstantVariable.create("slice step cannot be zero")
-                    raise_observed_exception(ValueError, tx, args=[msg])
+                    raise_observed_exception(
+                        ValueError, tx, args=["slice step cannot be zero"]
+                    )
 
                 value_unpack = value.force_unpack_var_sequence(tx)
                 try:
@@ -962,7 +1051,7 @@ def call_method(
                     raise_observed_exception(
                         type(exc),
                         tx,
-                        args=list(map(ConstantVariable.create, exc.args)),
+                        args=list(exc.args),
                     )
             else:
                 if isinstance(key, SymNodeVariable):
@@ -975,14 +1064,14 @@ def call_method(
                     self.items[key] = value
                 except (IndexError, TypeError) as e:
                     raise_observed_exception(
-                        type(e), tx, args=list(map(ConstantVariable.create, e.args))
+                        type(e), tx, args=[VariableTracker.build(tx, a) for a in e.args]
                     )
-            return CONSTANT_VARIABLE_NONE
+            return ConstantVariable.create(None)
 
         if name == "sort" and self.is_mutable():
             if len(args) != 0:
                 raise_args_mismatch(tx, name, "0 args", f"{len(args)} args")
-            key_fn_var = kwargs.pop("key", CONSTANT_VARIABLE_NONE)
+            key_fn_var = kwargs.pop("key", ConstantVariable.create(None))
             reverse = kwargs.pop(
                 "reverse", ConstantVariable.create(False)
             ).as_python_constant()
@@ -1035,18 +1124,18 @@ def call_method(
                 self.items[:] = [x for x, *_ in sorted_items_with_keys]
             except Exception as e:
                 raise_observed_exception(type(e), tx, args=list(e.args))
-            return CONSTANT_VARIABLE_NONE
+            return ConstantVariable.create(None)
 
         if name == "__init__" and self.is_mutable():
             if kwargs:
                 raise_args_mismatch(tx, name, "0 kwargs", f"{len(kwargs)} kwargs")
             if len(args) == 0:
-                return CONSTANT_VARIABLE_NONE
+                return ConstantVariable.create(None)
             elif len(args) == 1 and args[0].has_force_unpack_var_sequence(tx):
                 (arg,) = args
                 tx.output.side_effects.mutation(self)
                 self.items[:] = arg.force_unpack_var_sequence(tx)
-                return CONSTANT_VARIABLE_NONE
+                return ConstantVariable.create(None)
 
         return super().call_method(tx, name, args, kwargs)
 
@@ -1057,7 +1146,7 @@ def var_getattr(self, tx: "InstructionTranslator", name: str) -> VariableTracker
             if class_type is list:
                 return VariableTracker.build(tx, class_type, source=source)
             else:
-                return variables.UserDefinedClassVariable(class_type, source=source)
+                return VariableTracker.build(tx, class_type, source)
         return super().var_getattr(tx, name)
 
     def call_obj_hasattr(
@@ -1065,21 +1154,29 @@ def call_obj_hasattr(
     ) -> ConstantVariable:
         if self.python_type() is not list:
             return super().call_obj_hasattr(tx, name)
-        return variables.ConstantVariable.create(hasattr([], name))
+        return VariableTracker.build(tx, hasattr([], name))
 
     def is_python_hashable(self) -> bool:
         return False
 
 
+# TODO(follow-up): DequeVariable inherits BaseListVariable.mp_subscript_impl which
+# accepts slices. CPython's deque only has sq_item (Modules/_collectionsmodule.c:1888),
+# not mp_subscript — deque[slice] should raise TypeError. Override mp_subscript_impl
+# to reject slices and only accept integer-like keys via _PyIndex_Check → nb_index_impl.
+# Also add tests for: negative index, __index__ object key, invalid type key.
 class DequeVariable(CommonListMethodsVariable):
+    # deque_spec: https://github.com/python/cpython/blob/v3.13.0/Modules/_collectionsmodule.c#L1866
+    _cpython_type = collections.deque
+
     def __init__(
         self,
         items: list[VariableTracker],
-        maxlen: Optional[VariableTracker] = None,
+        maxlen: VariableTracker | None = None,
         **kwargs: Any,
     ) -> None:
         if maxlen is None:
-            maxlen = CONSTANT_VARIABLE_NONE
+            maxlen = ConstantVariable.create(None)
         assert maxlen.is_python_constant(), (
             f"maxlen must be a constant, got: {maxlen.debug_repr()}"
         )
@@ -1164,7 +1261,7 @@ def call_method(
             assert isinstance(key.as_python_constant(), int)
             tx.output.side_effects.mutation(self)
             self.items[key.as_python_constant()] = value
-            return CONSTANT_VARIABLE_NONE
+            return ConstantVariable.create(None)
 
         maxlen = self.maxlen.as_python_constant()
         if maxlen is not None:
@@ -1191,7 +1288,7 @@ def call_method(
                 tx, lambda item: self.call_method(tx, "appendleft", [item], {})
             )
             slice_within_maxlen = slice(None, maxlen)
-            result = CONSTANT_VARIABLE_NONE
+            result = ConstantVariable.create(None)
         elif name == "popleft" and self.is_mutable():
             if kwargs or len(args) > 0:
                 raise_args_mismatch(
@@ -1213,7 +1310,7 @@ def call_method(
             tx.output.side_effects.mutation(self)
             self.items[:] = [args[0], *self.items]
             slice_within_maxlen = slice(None, maxlen)
-            result = CONSTANT_VARIABLE_NONE
+            result = ConstantVariable.create(None)
         elif name == "insert" and len(args) > 0 and self.is_mutable():
             if kwargs or len(args) != 2:
                 raise_args_mismatch(
@@ -1242,11 +1339,14 @@ def call_obj_hasattr(
         self, tx: "InstructionTranslator", name: str
     ) -> ConstantVariable:
         if self.python_type() is collections.deque:
-            return variables.ConstantVariable.create(name in collections.deque.__dict__)
+            return VariableTracker.build(tx, name in collections.deque.__dict__)
         return super().call_obj_hasattr(tx, name)
 
 
 class TupleVariable(BaseListVariable):
+    # PyTuple_Type: https://github.com/python/cpython/blob/v3.13.0/Objects/tupleobject.c#L846
+    _cpython_type = tuple
+
     def python_type(self) -> type[tuple]:  # type: ignore[type-arg]
         return tuple
 
@@ -1264,10 +1364,7 @@ def var_getattr(self, tx: "InstructionTranslator", name: str) -> VariableTracker
         if name == "__class__":
             source = AttrSource(self.source, name) if self.source else None
             class_type = self.python_type()
-            if class_type is tuple:
-                return VariableTracker.build(tx, class_type, source=source)
-            else:
-                return variables.UserDefinedClassVariable(class_type, source=source)
+            return VariableTracker.build(tx, class_type, source=source)
         return super().var_getattr(tx, name)
 
     def call_obj_hasattr(
@@ -1275,7 +1372,7 @@ def call_obj_hasattr(
     ) -> ConstantVariable:
         if self.python_type() is not tuple:
             return super().call_obj_hasattr(tx, name)
-        return variables.ConstantVariable.create(hasattr((), name))
+        return VariableTracker.build(tx, hasattr((), name))
 
     def is_python_hashable(self) -> bool:
         return all(item.is_python_hashable() for item in self.items)
@@ -1293,6 +1390,8 @@ def is_python_equal(self, other: object) -> bool:
 class SizeVariable(TupleVariable):
     """torch.Size(...)"""
 
+    _cpython_type = torch.Size
+
     _nonvar_fields = {
         "proxy",
         *TupleVariable._nonvar_fields,
@@ -1301,7 +1400,7 @@ class SizeVariable(TupleVariable):
     def __init__(
         self,
         items: list[VariableTracker],
-        proxy: Optional[torch.fx.Proxy] = None,
+        proxy: torch.fx.Proxy | None = None,
         **kwargs: Any,
     ) -> None:
         self.proxy = proxy
@@ -1387,7 +1486,7 @@ def numel(self, tx: "InstructionTranslator") -> VariableTracker:
                 # Delay proxy calls  until we know it will be necessary
                 sym_sizes.append(v)
 
-        result = ConstantVariable.create(const_result)
+        result = VariableTracker.build(tx, const_result)
         if sym_sizes and const_result == 1:
             # Skip multiplying by 1
             result, *sym_sizes = sym_sizes
@@ -1400,6 +1499,29 @@ def numel(self, tx: "InstructionTranslator") -> VariableTracker:
             result = mul.call_function(tx, [result, v], {})
         return result
 
+    def mp_subscript_impl(
+        self,
+        tx: "InstructionTranslator",
+        key: VariableTracker,
+    ) -> VariableTracker:
+        # tuple_subscript: https://github.com/python/cpython/blob/62a6e898e01/Objects/tupleobject.c#L877-L930
+        # CPython: tuplesubscript checks _PyIndex_Check → PyNumber_AsSsize_t for non-slice keys
+        try:
+            key_type = key.python_type()
+        except NotImplementedError:
+            key_type = None
+        if key_type not in (int, bool, slice):
+            if key_type is not None and not hasattr(key_type, "__index__"):
+                raise_observed_exception(
+                    TypeError,
+                    tx,
+                    args=[
+                        f"tuple indices must be integers or slices, not {key.python_type_name()}"
+                    ],
+                )
+            key = key.nb_index_impl(tx)
+        return self.get_item_dyn(tx, key)
+
     def call_method(
         self,
         tx: "InstructionTranslator",
@@ -1407,17 +1529,7 @@ def call_method(
         args: list[VariableTracker],
         kwargs: dict[str, VariableTracker],
     ) -> VariableTracker:
-        if name == "__getitem__":
-            if kwargs or len(args) != 1:
-                raise_args_mismatch(
-                    tx,
-                    name,
-                    "1 args and 0 kwargs",
-                    f"{len(args)} args and {len(kwargs)} kwargs",
-                )
-            out = self.get_item_dyn(tx, args[0])
-            return out
-        elif name == "numel":
+        if name == "numel":
             if args or kwargs:
                 raise_args_mismatch(
                     tx,
@@ -1460,251 +1572,13 @@ def get_item_dyn(
     def call_obj_hasattr(
         self, tx: "InstructionTranslator", name: str
     ) -> ConstantVariable:
-        return variables.ConstantVariable.create(hasattr(torch.Size, name))
-
-
-class NamedTupleVariable(UserDefinedTupleVariable):
-    _nonvar_fields = {
-        "tuple_cls",
-        "dynamic_attributes",
-        *UserDefinedTupleVariable._nonvar_fields,
-    }
-
-    def __init__(
-        self,
-        items: list[VariableTracker],
-        # pyrefly: ignore [implicit-any]
-        tuple_cls: type[tuple],
-        dynamic_attributes: Optional[dict[str, VariableTracker]] = None,
-        tuple_vt: Optional[TupleVariable] = None,
-        **kwargs: Any,
-    ) -> None:
-        if tuple_vt is None:
-            assert getattr(kwargs, "source", None) is None
-            tuple_vt = variables.TupleVariable(
-                items, mutation_type=kwargs.get("mutation_type", ValueMutationNew())
-            )
-
-        if tuple_cls.__module__ == "torch.return_types":
-            # Structseq: single iterable argument
-            dummy_value = tuple_cls(items)
-        else:
-            # Namedtuple: positional arguments
-            dummy_value = tuple_cls(*items)  # type: ignore[arg-type]
-
-        super().__init__(
-            value=dummy_value,
-            tuple_vt=tuple_vt,
-            init_args=items,
-            **kwargs,
-        )
-        self.tuple_cls = tuple_cls
-        if len(self.tuple_cls.__mro__) < 3:
-            raise ValueError("NamedTuple should inherit from Tuple and Object.")
-        self.dynamic_attributes = dynamic_attributes if dynamic_attributes else {}
-
-    @property
-    def items(self) -> list[VariableTracker]:
-        return self._tuple_vt.items
-
-    def is_namedtuple(self) -> bool:
-        return isinstance(getattr(self.tuple_cls, "_fields", None), tuple) and callable(
-            getattr(self.tuple_cls, "_make", None)
-        )
-
-    def is_structseq(self) -> bool:
-        return not self.is_namedtuple()
-
-    def fields(self) -> tuple[str, ...]:
-        return namedtuple_fields(self.tuple_cls)
-
-    def as_python_constant(self) -> Any:
-        if self.is_structseq():
-            # StructSequenceType(iterable)
-            result = self.python_type()([x.as_python_constant() for x in self.items])
-        else:
-            # NamedTupleType(*iterable)
-            result = self.python_type()(*[x.as_python_constant() for x in self.items])
-
-        # Apply dynamic attributes if any were set
-        if self.dynamic_attributes:
-            for attr_name, attr_value in self.dynamic_attributes.items():
-                # Convert VariableTracker to Python constant if needed
-                if hasattr(attr_value, "as_python_constant"):
-                    python_value = attr_value.as_python_constant()
-                else:
-                    raise NotImplementedError(
-                        "Can not convert dynamic attribute without python constant value to python constant."
-                    )
-                setattr(result, attr_name, python_value)
-
-        return result
-
-    def as_proxy(self) -> Any:
-        if self.is_structseq():
-            return self.python_type()([x.as_proxy() for x in self._tuple_vt.items])
-        return self.python_type()(*[x.as_proxy() for x in self._tuple_vt.items])
-
-    def call_tree_map(
-        self,
-        tx: "InstructionTranslator",
-        tree_map_fn: UserFunctionVariable,
-        map_fn: VariableTracker,
-        rest: Sequence[VariableTracker],
-        tree_map_kwargs: dict[str, VariableTracker],
-    ) -> VariableTracker:
-        is_leaf_var = tree_map_kwargs.get("is_leaf")
-        if is_leaf_var is not None and not is_leaf_var.is_constant_none():
-            pred_result = is_leaf_var.call_function(tx, [self], {})
-            try:
-                leaf_decision = pred_result.as_python_constant()
-            except NotImplementedError:
-                return self._tree_map_fallback(
-                    tx, tree_map_fn, map_fn, rest, tree_map_kwargs
-                )
-            if leaf_decision:
-                return map_fn.call_function(tx, [self, *rest], {})
-
-        return self.call_tree_map_branch(
-            tx,
-            tree_map_fn,
-            map_fn,
-            rest,
-            tree_map_kwargs,
-        )
-
-    def call_tree_map_branch(
-        self,
-        tx: "InstructionTranslator",
-        tree_map_fn: UserFunctionVariable,
-        map_fn: VariableTracker,
-        rest: Sequence[VariableTracker],
-        tree_map_kwargs: dict[str, VariableTracker],
-    ) -> VariableTracker:
-        other_tuples: list[NamedTupleVariable] = []
-        for candidate in rest:
-            if (
-                not isinstance(candidate, NamedTupleVariable)
-                or len(candidate.items) != len(self.items)
-                or candidate.tuple_cls is not self.tuple_cls
-            ):
-                return self._tree_map_fallback(
-                    tx, tree_map_fn, map_fn, rest, tree_map_kwargs
-                )
-            other_tuples.append(candidate)
-
-        new_items: list[VariableTracker] = []
-        for idx, item in enumerate(self.items):
-            sibling_leaves = [candidate.items[idx] for candidate in other_tuples]
-            new_items.append(
-                item.call_tree_map(
-                    tx,
-                    tree_map_fn,
-                    map_fn,
-                    sibling_leaves,
-                    tree_map_kwargs,
-                )
-            )
-
-        return NamedTupleVariable(
-            new_items,
-            self.tuple_cls,
-            mutation_type=ValueMutationNew(),
-        )
-
-    def reconstruct(self, codegen: "PyCodegen") -> None:
-        if self.is_structseq():
-            create_fn = self.tuple_cls
-        else:
-            create_fn = self.tuple_cls._make  # type: ignore[attr-defined]
-
-        codegen.add_push_null(
-            lambda: codegen.append_output(
-                codegen.create_load_const_unchecked(create_fn)
-            )
-        )
-        codegen.foreach(self._tuple_vt.items)
-        codegen.extend_output(
-            [
-                create_build_tuple(len(self._tuple_vt.items)),
-            ]
-            + create_call_function(1, False)
-        )
-
-        # Apply initial dynamic attributes after construction (if any)
-        # Runtime dynamic attributes are tracked via side effects system
-        for name, value in self.dynamic_attributes.items():
-            codegen.dup_top()
-            codegen(value)
-            codegen.extend_output(create_rot_n(2))
-            codegen.store_attr(name)
-
-    def _is_method_overridden(self, method_name: str) -> bool:
-        if len(self.tuple_cls.__mro__) < 3:
-            raise ValueError("NamedTuple should inherit from Tuple and Object.")
-        if getattr(self.tuple_cls, method_name, None) == getattr(
-            self.tuple_cls.__mro__[-3], method_name, None
-        ):
-            return False
-        return True
-
-    def call_method(
-        self,
-        tx: "InstructionTranslator",
-        name: str,
-        args: list[VariableTracker],
-        kwargs: dict[str, VariableTracker],
-    ) -> VariableTracker:
-        if self._is_method_overridden(name):
-            # Fall back to UserDefinedTupleVariable
-            return super().call_method(tx, name, args, kwargs)
-        elif name == "__setattr__":
-            if kwargs or len(args) != 2:
-                raise_args_mismatch(
-                    tx,
-                    name,
-                    "2 args and 0 kwargs",
-                    f"{len(args)} args and {len(kwargs)} kwargs",
-                )
-            attr_var, value = args
-            attr = attr_var.as_python_constant()
-
-            if (
-                # structseq is immutable
-                self.is_structseq()
-                # namedtuple directly created by `collections.namedtuple` is immutable
-                or self.tuple_cls.__bases__ == (tuple,)
-                or attr in self.fields()
-            ):
-                raise_observed_exception(AttributeError, tx)
-
-            result = self.method_setattr_standard(tx, attr_var, value)
-            # Also update self.dynamic_attributes
-            self.dynamic_attributes[attr] = value
-            return result
-
-        return super().call_method(tx, name, args, kwargs)
-
-    def python_type(self) -> type:
-        return self.tuple_cls
-
-    def var_getattr(self, tx: "InstructionTranslator", name: str) -> "VariableTracker":
-        if name == "_fields":
-            source = NamedTupleFieldsSource(self.source) if self.source else None
-            return VariableTracker.build(tx, self.fields(), source=source)
-
-        if name in self.dynamic_attributes:
-            return self.dynamic_attributes[name]
-
-        fields = self.fields()
-        if name in fields:
-            field_index = fields.index(name)
-            return self._tuple_vt.items[field_index]
-
-        return super().var_getattr(tx, name)
+        return VariableTracker.build(tx, hasattr(torch.Size, name))
 
 
 class SliceVariable(VariableTracker):
+    # PySlice_Type: https://github.com/python/cpython/blob/v3.13.0/Objects/sliceobject.c#L689
+    _cpython_type = slice
+
     def __init__(
         self,
         items: Sequence[VariableTracker],
@@ -1712,7 +1586,7 @@ def __init__(
         **kwargs: Any,
     ) -> None:
         items_to_map = items
-        start, stop, step = [variables.CONSTANT_VARIABLE_NONE] * 3
+        start, stop, step = [variables.ConstantVariable.create(None)] * 3
 
         if len(items_to_map) == 1:
             (stop,) = items_to_map
@@ -1763,7 +1637,9 @@ def reconstruct(self, codegen: "PyCodegen") -> None:
 
     def var_getattr(self, tx: "InstructionTranslator", name: str) -> VariableTracker:
         if name in cmp_name_to_op_mapping:
-            return variables.GetAttrVariable(self, name)
+            return variables.GetAttrVariable(
+                self, name, py_type=type(getattr(slice, name))
+            )
         fields = ["start", "stop", "step"]
         if name not in fields:
             unimplemented(
@@ -1777,6 +1653,9 @@ def var_getattr(self, tx: "InstructionTranslator", name: str) -> VariableTracker
 
 
 class ListIteratorVariable(IteratorVariable):
+    # PyListIter_Type: https://github.com/python/cpython/blob/v3.13.0/Objects/listobject.c#L3842
+    _cpython_type = type(iter([]))
+
     _nonvar_fields = {
         "index",
         *IteratorVariable._nonvar_fields,
@@ -1812,7 +1691,7 @@ def next_variable(self, tx: "InstructionTranslator") -> VariableTracker:
     def call_obj_hasattr(
         self, tx: "InstructionTranslator", name: str
     ) -> ConstantVariable:
-        return variables.ConstantVariable.create(hasattr(iter([]), name))
+        return VariableTracker.build(tx, hasattr(iter([]), name))
 
     def python_type(self) -> type:
         return type(iter([]))
@@ -1852,10 +1731,14 @@ def reconstruct(self, codegen: "PyCodegen") -> None:
 
 
 class TupleIteratorVariable(ListIteratorVariable):
-    pass
+    # PyTupleIter_Type: https://github.com/python/cpython/blob/v3.13.0/Objects/tupleobject.c#L1067
+    _cpython_type = type(iter(()))
 
 
 class RangeIteratorVariable(IteratorVariable):
+    # PyRangeIter_Type: https://github.com/python/cpython/blob/v3.13.0/Objects/rangeobject.c#L896
+    _cpython_type = type(iter(range(0)))
+
     # only needed for isinstance(..., range_iterator) to work
     _nonvar_fields = {
         "iter_obj",
@@ -1886,11 +1769,9 @@ def call_method(
     def call_obj_hasattr(
         self, tx: "InstructionTranslator", name: str
     ) -> ConstantVariable:
-        from .builder import SourcelessBuilder
-
         if self.python_type() is range_iterator:
             ri = iter(range(0))
-            return SourcelessBuilder.create(tx, hasattr(ri, name))
+            return VariableTracker.build(tx, hasattr(ri, name))
         return super().call_obj_hasattr(tx, name)
 
     def next_variable(self, tx: "InstructionTranslator") -> VariableTracker:
@@ -1900,7 +1781,7 @@ def next_variable(self, tx: "InstructionTranslator") -> VariableTracker:
         self.len -= 1
         current = self.start
         self.start += self.step
-        return ConstantVariable.create(current)
+        return VariableTracker.build(tx, current)
 
     def python_type(self) -> type:
         return range_iterator
diff --git a/torch/_dynamo/variables/misc.py b/torch/_dynamo/variables/misc.py
index 574a8cac16500..c297b0a7136ef 100644
--- a/torch/_dynamo/variables/misc.py
+++ b/torch/_dynamo/variables/misc.py
@@ -15,6 +15,7 @@
 - DebuggingVariable: Handles print and logging
 """
 
+import builtins
 import dataclasses
 import enum
 import functools
@@ -30,7 +31,7 @@
 from collections.abc import Callable, Sequence
 from random import Random
 from types import BuiltinFunctionType
-from typing import Any, Literal, NoReturn, TYPE_CHECKING, TypeGuard, Union
+from typing import Any, Literal, TYPE_CHECKING, TypeGuard, Union
 
 import torch._C
 import torch._numpy as tnp
@@ -46,7 +47,7 @@
     create_instruction,
 )
 from ..create_parameter_op import do_not_convert_to_tracable_parameter
-from ..exc import raise_observed_exception, unimplemented
+from ..exc import raise_observed_exception, raise_type_error, unimplemented
 from ..guards import GuardBuilder, install_guard
 from ..mutation_guard import unpatched_nn_module_init
 from ..source import (
@@ -63,17 +64,11 @@
     identity,
     is_tensor_base_attr_getter,
     istype,
-    list_methods,
     proxy_args_kwargs,
     raise_args_mismatch,
-    tuple_methods,
 )
-from .base import (
-    AsPythonConstantNotImplementedError,
-    raise_type_error_exc,
-    VariableTracker,
-)
-from .constant import CONSTANT_VARIABLE_NONE, ConstantVariable
+from .base import AsPythonConstantNotImplementedError, NO_SUCH_SUBOBJ, VariableTracker
+from .constant import ConstantVariable
 from .functions import NestedUserFunctionVariable, UserFunctionVariable
 from .user_defined import call_random_fn, is_standard_setattr, UserDefinedObjectVariable
 
@@ -83,11 +78,10 @@
     from torch._dynamo.symbolic_convert import InstructionTranslator
 
 
-class NO_SUCH_SUBOBJ:
-    pass
-
-
 class SuperVariable(VariableTracker):
+    # PySuper_Type: https://github.com/python/cpython/blob/v3.13.0/Objects/typeobject.c#L11511
+    _cpython_type = super
+
     _nonvar_fields = {
         *VariableTracker._nonvar_fields,
     }
@@ -109,6 +103,9 @@ def __init__(
         # cls for a classmethod)
         self.objvar = objvar
 
+    def python_type(self) -> type:
+        return builtins.super
+
     def reconstruct(self, codegen: "PyCodegen") -> None:
         codegen.add_push_null(lambda: codegen(variables.BuiltinVariable(super)))
         codegen(self.typevar)
@@ -197,7 +194,7 @@ def var_getattr(self, tx: "InstructionTranslator", name: str) -> VariableTracker
         # not just AttrSource).
         value, source = self._resolved_getattr_and_source(tx, name)
         if not variables.ConstantVariable.is_literal(value):
-            return GetAttrVariable(self, name)
+            return GetAttrVariable(self, name, py_type=type(value))
         if source:
             install_guard(source.make_guard(GuardBuilder.CONSTANT_MATCH))
         return variables.ConstantVariable.create(value, source=source)
@@ -252,7 +249,7 @@ def call_method(
         ):
             user_cls = inner_fn.__self__
             if hasattr(user_cls, "__module__") and user_cls.__module__ == "builtins":
-                user_cls_vt: VariableTracker = variables.BuiltinVariable(user_cls)
+                user_cls_vt: VariableTracker = VariableTracker.build(tx, user_cls)
             else:
                 assert source is not None
                 user_cls_source = source.member
@@ -342,27 +339,14 @@ def call_method(
             tx.output.side_effects.store_attr(
                 self.objvar, attr, variables.DeletedVariable()
             )
-            return variables.CONSTANT_VARIABLE_NONE
-        elif (
-            isinstance(self.objvar, variables.UserDefinedDictVariable)
-            and inner_fn in self.objvar._dict_methods
-        ):
-            return self.objvar._dict_vt.call_method(tx, name, args, kwargs)
+            return variables.ConstantVariable.create(None)
         elif (
-            isinstance(self.objvar, variables.UserDefinedSetVariable)
-            and inner_fn in self.objvar._set_methods
+            isinstance(self.objvar, variables.UserDefinedObjectVariable)
+            and self.objvar._base_vt is not None
+            and self.objvar._base_methods is not None
+            and inner_fn in self.objvar._base_methods
         ):
-            return self.objvar._set_vt.call_method(tx, name, args, kwargs)
-        elif (
-            isinstance(self.objvar, variables.UserDefinedTupleVariable)
-            and inner_fn in tuple_methods
-        ):
-            return self.objvar._tuple_vt.call_method(tx, name, args, kwargs)
-        elif (
-            isinstance(self.objvar, variables.UserDefinedListVariable)
-            and inner_fn in list_methods
-        ):
-            return self.objvar._list_vt.call_method(tx, name, args, kwargs)
+            return self.objvar._base_vt.call_method(tx, name, args, kwargs)
         elif inner_fn is object.__getattribute__:
             # object.__getattribute__ has no side-effects. We can directly call
             # __getattribute__ to access the attribute.
@@ -444,13 +428,13 @@ def python_type(self) -> type:
 
     def var_getattr(self, tx: "InstructionTranslator", name: str) -> VariableTracker:
         if name == "lineno":
-            return variables.ConstantVariable.create(self.frame_summary.lineno)
+            return VariableTracker.build(tx, self.frame_summary.lineno)
         elif name == "filename":
-            return variables.ConstantVariable.create(self.frame_summary.filename)
+            return VariableTracker.build(tx, self.frame_summary.filename)
         elif name == "name":
-            return variables.ConstantVariable.create(self.frame_summary.name)
+            return VariableTracker.build(tx, self.frame_summary.name)
         elif name == "line":
-            return variables.ConstantVariable.create(self.frame_summary.line)
+            return VariableTracker.build(tx, self.frame_summary.line)
         return super().var_getattr(tx, name)
 
 
@@ -483,9 +467,7 @@ def from_frame_summary(
 
     @staticmethod
     def is_valid_traceback(obj: VariableTracker) -> bool:
-        return istype(obj, TracebackVariable) or (
-            istype(obj, ConstantVariable) and obj.is_constant_none()
-        )
+        return istype(obj, TracebackVariable) or obj.is_constant_none()
 
     def extract_tb(self) -> list[traceback.FrameSummary | FrameSummaryVariable]:
         if istype(self.tb_next, ConstantVariable):
@@ -520,7 +502,7 @@ def call_setattr(
             ):
                 raise_observed_exception(ValueError, tx)
             self.tb_next = val
-        return variables.CONSTANT_VARIABLE_NONE
+        return variables.ConstantVariable.create(None)
 
     def var_getattr(self, tx: "InstructionTranslator", name: str) -> VariableTracker:
         if name == "tb_next":
@@ -547,18 +529,21 @@ def call_method(
     ) -> VariableTracker:
         if name == "__eq__":
             # Two traceback variables are only equal if they are the same object
-            return variables.ConstantVariable.create(self is args[0])
+            return VariableTracker.build(tx, self is args[0])
         elif name == "__setattr__":
             return self.call_setattr(tx, *args)
         return super().call_method(tx, name, args, kwargs)
 
 
 class ExceptionVariable(VariableTracker):
+    # _PyExc_BaseException: https://github.com/python/cpython/blob/v3.13.0/Objects/exceptions.c
+    _cpython_type = BaseException
+
     # The ExceptionVariable corresponds to the BaseException class in Python
     def __init__(
         self,
         exc_type: Any,
-        args: tuple[VariableTracker, ...],
+        args: list[VariableTracker],
         init_kwargs: dict[str, VariableTracker] | None = None,
         source: Source | None = None,
         mutation_type: MutationType | None = None,
@@ -576,14 +561,14 @@ def __init__(
         # When raising a new exception while another exception is already being
         # handled, the new exception's __context__ attribute is automatically
         # set to the handled exception.
-        self.__context__: VariableTracker = CONSTANT_VARIABLE_NONE
+        self.__context__: VariableTracker = ConstantVariable.create(None)
         # Set when user raised an exception from another:
         # raise ... from ...
-        self.__cause__: VariableTracker = CONSTANT_VARIABLE_NONE
+        self.__cause__: VariableTracker = ConstantVariable.create(None)
         # Boolean flag that controls whether the __context__ attribute is set
-        self.__suppress_context__: VariableTracker = ConstantVariable(False)
+        self.__suppress_context__: VariableTracker = ConstantVariable.create(False)
         # Contains the call stack where the exception was raised.
-        self.__traceback__: VariableTracker = CONSTANT_VARIABLE_NONE
+        self.__traceback__: VariableTracker = ConstantVariable.create(None)
         # The user stack at the time this exception was first raised.
         # Used to preserve the original exception location when re-raising.
         self.python_stack: traceback.StackSummary | None = None
@@ -621,13 +606,17 @@ def call_setattr(
         name_var: VariableTracker,
         val: VariableTracker,
     ) -> VariableTracker:
-        def raise_error(msg: str) -> NoReturn:
-            raise_observed_exception(TypeError, tx, args=[ConstantVariable(msg)])
-
         name = name_var.as_python_constant()
         if name == "__context__":
             # Constant can be either an Exceptior or None
-            assert isinstance(val, (ExceptionVariable, ConstantVariable))
+            assert val.is_constant_none() or isinstance(
+                val,
+                (
+                    variables.ExceptionVariable,
+                    variables.UserDefinedExceptionClassVariable,
+                    variables.UserDefinedExceptionObjectVariable,
+                ),
+            ), f"{val} is not a valid exception context"
             self.set_context(val)
         elif name == "__cause__":
             if val.is_constant_none() or isinstance(
@@ -640,25 +629,21 @@ def raise_error(msg: str) -> NoReturn:
                 ),
             ):
                 self.__cause__ = val
-                self.__suppress_context__ = variables.ConstantVariable(True)
+                self.__suppress_context__ = variables.ConstantVariable.create(True)
             else:
-                raise_error("exception cause must be None or derive from BaseException")
+                raise_type_error(
+                    tx, "exception cause must be None or derive from BaseException"
+                )
         elif name == "__suppress_context__":
             if val.is_constant_match(True, False):
                 self.__suppress_context__ = val
             else:
-                raise_error("exception cause must be None or derive from BaseException")
+                raise_type_error(
+                    tx, "exception cause must be None or derive from BaseException"
+                )
         elif name == "__traceback__":
             if not TracebackVariable.is_valid_traceback(val):
-                raise_observed_exception(
-                    TypeError,
-                    tx,
-                    args=[
-                        ConstantVariable.create(
-                            "__traceback__ must be a traceback object or None"
-                        )
-                    ],
-                )
+                raise_type_error(tx, "__traceback__ must be a traceback object or None")
             self.__traceback__ = val
         else:
             unimplemented(
@@ -669,7 +654,7 @@ def raise_error(msg: str) -> NoReturn:
                 "`__cause__`, `__suppress_context__`, and `__traceback__` are supported.",
                 hints=[*graph_break_hints.SUPPORTABLE],
             )
-        return variables.CONSTANT_VARIABLE_NONE
+        return variables.ConstantVariable.create(None)
 
     def call_method(
         self,
@@ -682,13 +667,15 @@ def call_method(
             return self.call_setattr(tx, *args)
         elif name == "with_traceback":
             [tb] = args
-            self.call_setattr(tx, ConstantVariable("__traceback__"), tb)
+            self.call_setattr(tx, VariableTracker.build(tx, "__traceback__"), tb)
             return self
         else:
             return super().call_method(tx, name, args, kwargs)
 
     def var_getattr(self, tx: "InstructionTranslator", name: str) -> VariableTracker:
-        if name == "__context__":
+        if name == "__class__":
+            return VariableTracker.build(tx, self.exc_type)
+        elif name == "__context__":
             return self.__context__
         elif name == "__cause__":
             return self.__cause__
@@ -697,7 +684,11 @@ def var_getattr(self, tx: "InstructionTranslator", name: str) -> VariableTracker
         elif name == "__traceback__":
             return self.__traceback__
         elif name == "args":
-            return variables.ListVariable(list(self.args), source=self.source)
+            return VariableTracker.build(
+                tx,
+                tuple(self.args),
+                source=self.source and AttrSource(self.source, "args"),
+            )
         return super().var_getattr(tx, name)
 
     def __str__(self) -> str:
@@ -705,6 +696,17 @@ def __str__(self) -> str:
 
     __repr__ = __str__
 
+    @staticmethod
+    def _debug_format_arg(arg: VariableTracker) -> str:
+        try:
+            return repr(arg.as_python_constant())
+        except Exception:
+            return arg.debug_repr()
+
+    def debug_repr(self) -> str:
+        args = ", ".join(self._debug_format_arg(arg) for arg in self.args)
+        return f"{self.python_type_name()}({args})"
+
 
 class UnknownVariable(VariableTracker):
     """
@@ -779,7 +781,7 @@ def call_function(
             # We have to manually bind the freevars ourselves
             code = fn.get_code()
             if fn.closure:
-                raise_type_error_exc(
+                raise_type_error(
                     tx,
                     f"comptime function must not have free variables, but these variables were free: {code.co_freevars}",
                 )
@@ -800,10 +802,13 @@ def call_function(
         else:
             raise RuntimeError(f"unsupported argument to comptime: {type(fn)}")
 
-        return variables.CONSTANT_VARIABLE_NONE
+        return variables.ConstantVariable.create(None)
 
 
 class CellVariable(VariableTracker):
+    # PyCell_Type: https://github.com/python/cpython/blob/v3.13.0/Objects/cellobject.c#L151
+    _cpython_type = types.CellType
+
     # If the cell existed before Dynamo tracing started, this will be the
     # VariableTracker that represents the cell content.
     #
@@ -822,6 +827,9 @@ def __init__(
         super().__init__(**kwargs)
         self.pre_existing_contents = pre_existing_contents
 
+    def python_type(self) -> type:
+        return types.CellType
+
 
 class NewGlobalVariable(VariableTracker):
     def __init__(self, **kwargs: Any) -> None:
@@ -849,6 +857,9 @@ def __init__(self, fn_cls: Any, **kwargs: Any) -> None:
         super().__init__(**kwargs)
         self.fn_cls = fn_cls
 
+    def python_type(self) -> type:
+        return type
+
     def call_apply(
         self,
         tx: "InstructionTranslator",
@@ -1146,7 +1157,7 @@ def call_method(
             if kwargs:
                 raise_args_mismatch(tx, name, "0 kwargs", f"{len(kwargs)} kwargs")
             self.non_differentiable = proxy_args_kwargs(args, {})[0]
-            return variables.CONSTANT_VARIABLE_NONE
+            return variables.ConstantVariable.create(None)
 
         if name != "save_for_backward":
             unimplemented(
@@ -1173,7 +1184,7 @@ def call_method(
         assert self.saved_tensors is not None
         if not self.inference:
             if kwargs or not self.source:
-                raise_type_error_exc(
+                raise_type_error(
                     tx, "save_for_backward() requires a source and no keyword arguments"
                 )
             tx.output.side_effects.track_save_for_backward(self, args)
@@ -1184,7 +1195,7 @@ def call_method(
             self.saved_tensors.tensors = []
         for arg in args:
             self.saved_tensors.tensors.append(arg)
-        return variables.CONSTANT_VARIABLE_NONE
+        return variables.ConstantVariable.create(None)
 
     def var_getattr(self, tx: "InstructionTranslator", name: str) -> VariableTracker:
         if name in ["save_for_backward", "mark_non_differentiable"]:
@@ -1263,6 +1274,9 @@ def __init__(self, fn: Callable[..., VariableTracker], **kwargs: Any) -> None:
         super().__init__(**kwargs)
         self.fn = fn
 
+    def python_type(self) -> type:
+        return types.FunctionType
+
     def call_function(
         self,
         tx: "InstructionTranslator",
@@ -1339,89 +1353,15 @@ def call_function(
     ) -> VariableTracker:
         return self.obj.call_method(tx, self.name, list(args), kwargs)
 
-    def call_method(
-        self,
-        tx: "InstructionTranslator",
-        name: str,
-        args: list[VariableTracker],
-        kwargs: dict[str, VariableTracker],
-    ) -> VariableTracker:
-        if (
-            name in ("__getitem__", "get")
-            and self.name == "__dict__"
-            and not kwargs
-            and args[0].is_python_constant()
-            and isinstance(
-                self.obj,
-                (
-                    variables.UserDefinedObjectVariable,
-                    variables.NNModuleVariable,
-                    variables.UserDefinedClassVariable,
-                ),
-            )
-        ):
-            obj = self.obj
-            key = args[0].as_python_constant()
-            if obj.has_key_in_generic_dict(tx, key):
-                # redirect to var_getattr on the original obj
-                return obj.var_getattr(tx, key)
-
-            # Return the default value for get
-            if name == "get":
-                if len(args) == 2:
-                    return args[1]
-                else:
-                    return variables.CONSTANT_VARIABLE_NONE
-
-        elif (
-            name == "__contains__"
-            and self.name == "__dict__"
-            and len(args) == 1
-            and args[0].is_python_constant()
-            and not kwargs
-            and isinstance(
-                self.obj,
-                (
-                    variables.UserDefinedObjectVariable,
-                    variables.NNModuleVariable,
-                    variables.UserDefinedClassVariable,
-                ),
-            )
-        ):
-            obj = self.obj
-            key = args[0].as_python_constant()
-            if obj.has_key_in_generic_dict(tx, key):
-                return variables.ConstantVariable(True)
-            else:
-                return variables.ConstantVariable(False)
-
-        elif name == "__setitem__" and self.name == "__dict__" and not kwargs:
-            if isinstance(self.obj, variables.UserDefinedObjectVariable):
-                # Bypass any custom setattr as we are updating the `__dict__` itself
-                return self.obj.method_setattr_standard(
-                    tx, args[0], args[1], directly_update_dict=True
-                )
-            if isinstance(self.obj, variables.NNModuleVariable):
-                # This matches how `setattr` is handled for NNModuleVariable
-                self.obj.convert_to_unspecialized(tx)
-
-        return super().call_method(tx, name, args, kwargs)
-
-    def get_forwarded_dict(self, tx: "InstructionTranslator") -> VariableTracker:
-        assert (
-            self.name == "__dict__"
-            and isinstance(self.obj, variables.UserDefinedClassVariable)
-            and not tx.output.side_effects.has_pending_mutation(self.obj)
-        )
-        self.obj.ban_mutation = True
-        return VariableTracker.build(tx, self.obj.value.__dict__, self.source)
-
 
 class MethodWrapperVariable(VariableTracker):
     def __init__(self, method_wrapper: types.MethodWrapperType, **kwargs: Any) -> None:
         super().__init__(**kwargs)
         self.method_wrapper = method_wrapper
 
+    def get_real_python_backed_value(self) -> types.MethodWrapperType:
+        return self.method_wrapper
+
     def call_function(
         self,
         tx: "InstructionTranslator",
@@ -1432,7 +1372,7 @@ def call_function(
             args[0], variables.TensorVariable
         ):
             if not (len(args) == 1 and len(kwargs) == 0):
-                raise_type_error_exc(
+                raise_type_error(
                     tx, "tensor attribute getter takes exactly one argument"
                 )
             # type: ignore[arg-type, attr-defined]
@@ -1450,7 +1390,7 @@ def call_function(
         if wrapper_name == "__init__":
             fn_obj = type(self_obj).__init__
             if fn_obj is object.__init__:
-                return variables.BuiltinVariable(object).call_method(
+                return VariableTracker.build(tx, object).call_method(
                     tx,
                     wrapper_name,
                     # type: ignore[arg-type, list-item]
@@ -1492,19 +1432,33 @@ def call_function(
         elif (self_obj is type.__dict__["__mro__"] and wrapper_name == "__get__") or (
             self_obj is type.__dict__["__dict__"] and wrapper_name == "__get__"
         ):
-            from .builder import SourcelessBuilder
+            attr_name = (
+                "__mro__" if self_obj is type.__dict__["__mro__"] else "__dict__"
+            )
 
             if len(args) == 1 and not kwargs:
                 try:
-                    return SourcelessBuilder.create(
-                        tx, self.method_wrapper(args[0].as_python_constant())
-                    )
+                    value = self.method_wrapper(args[0].as_python_constant())
                 except AsPythonConstantNotImplementedError:
                     pass
+                else:
+                    # Use a sourced variable when the descriptor is the
+                    # standard one from type (not overridden by a metaclass).
+                    source = args[0].source
+                    if source is not None:
+                        cls_val = args[0].as_python_constant()
+                        static_desc = inspect.getattr_static(type(cls_val), attr_name)
+                        if static_desc is self_obj:
+                            if attr_name == "__mro__":
+                                source = TypeMROSource(source)
+                            else:
+                                source = AttrSource(source, attr_name)
+                            return VariableTracker.build(tx, value, source)
+
+                    from .builder import SourcelessBuilder
+
+                    return SourcelessBuilder.create(tx, value)
 
-            attr_name = (
-                "__mro__" if self_obj is type.__dict__["__mro__"] else "__dict__"
-            )
             unimplemented(
                 gb_type=f"unsupported type.__dict__['{attr_name}'].__get__ call",
                 context=f"call_function {self}, args: {args}, kwargs: {kwargs}",
@@ -1545,10 +1499,16 @@ def __init__(self, desc: types.GetSetDescriptorType, **kwargs: Any) -> None:
         super().__init__(**kwargs)
         self.desc = desc
 
+    def get_real_python_backed_value(self) -> types.GetSetDescriptorType:
+        return self.desc
+
     def var_getattr(self, tx: "InstructionTranslator", name: str) -> VariableTracker:
         if name == "__get__" and self.source:
             source = AttrSource(self.source, "__get__")
             return VariableTracker.build(tx, self.desc.__get__, source)
+        elif name in ("__objclass__", "__name__"):
+            source = self.source and AttrSource(self.source, name)
+            return VariableTracker.build(tx, getattr(self.desc, name), source)
         else:
             return super().var_getattr(tx, name)
 
@@ -1560,6 +1520,9 @@ def as_python_constant(self) -> types.GetSetDescriptorType:
 
 
 class PythonModuleVariable(VariableTracker):
+    # PyModule_Type: https://github.com/python/cpython/blob/v3.13.0/Objects/moduleobject.c#L1203
+    _cpython_type = types.ModuleType
+
     _nonvar_fields = {
         "value",
         "is_torch",
@@ -1577,14 +1540,17 @@ def python_type(self) -> type[types.ModuleType]:
     def as_python_constant(self) -> types.ModuleType:
         return self.value
 
+    def get_real_python_backed_value(self) -> types.ModuleType:
+        return self.value
+
     def __repr__(self) -> str:
         return f"PythonModuleVariable({self.value})"
 
     def call_obj_hasattr(
         self, tx: "InstructionTranslator", name: str
-    ) -> ConstantVariable:
+    ) -> "ConstantVariable":
         result = hasattr(self.value, name)
-        return variables.ConstantVariable.create(result)
+        return VariableTracker.build(tx, result)
 
     def var_getattr(self, tx: "InstructionTranslator", name: str) -> VariableTracker:
         if tx.output.side_effects.has_pending_mutation_of_attr(self, name):
@@ -1608,6 +1574,16 @@ def __init__(self, value: Any, **kwargs: Any) -> None:
         super().__init__(**kwargs)
         self.value = value
 
+    def mp_subscript_impl(
+        self,
+        tx: "InstructionTranslator",
+        key: VariableTracker,
+    ) -> VariableTracker:
+        # e.g., List[int] → typing.List[int]
+        # TODO(follow-up): add test for invalid subscript type
+        new_typing = self.value[key.as_python_constant()]
+        return TypingVariable(new_typing)
+
     def call_method(
         self,
         tx: "InstructionTranslator",
@@ -1615,11 +1591,7 @@ def call_method(
         args: list[VariableTracker],
         kwargs: dict[str, VariableTracker],
     ) -> VariableTracker:
-        # Create a new typing variable, e.g., `List[int]`
-        if name == "__getitem__" and len(args) == 1:
-            new_typing = self.value[args[0].as_python_constant()]
-            return TypingVariable(new_typing)
-        elif name == "__eq__":
+        if name == "__eq__":
             if len(args) == 1 and not kwargs:
                 result = istype(args[0], TypingVariable) and self.value == args[0].value
                 return variables.ConstantVariable.create(result)
@@ -1637,7 +1609,9 @@ def var_getattr(self, tx: "InstructionTranslator", name: str) -> VariableTracker
         from .builder import SourcelessBuilder, VariableBuilder
 
         if name in cmp_name_to_op_mapping:
-            return variables.GetAttrVariable(self, name)
+            return variables.GetAttrVariable(
+                self, name, py_type=type(getattr(self.value, name))
+            )
 
         if tx.output.side_effects.has_pending_mutation_of_attr(self, name):
             return tx.output.side_effects.load_attr(self, name)
@@ -1652,6 +1626,9 @@ def var_getattr(self, tx: "InstructionTranslator", name: str) -> VariableTracker
     def as_python_constant(self) -> Any:
         return self.value
 
+    def get_real_python_backed_value(self) -> Any:
+        return self.value
+
     def reconstruct(self, codegen: "PyCodegen") -> None:
         if not isinstance(self.value, types.GenericAlias):
             return super().reconstruct(codegen)
@@ -1735,6 +1712,9 @@ def __init__(self, value: Any, **kwargs: Any) -> None:
         super().__init__(**kwargs)
         self.value = value
 
+    def get_real_python_backed_value(self) -> Any:
+        return self.value
+
     @classmethod
     def can_constant_fold_through(cls, fn: types.FunctionType) -> bool:
         mod = fn.__module__.split(".")
@@ -1823,7 +1803,8 @@ def call_function(
                 check_unspec_or_constant_args(args, kwargs)
             ):
                 # constant fold
-                return variables.ConstantVariable.create(
+                return VariableTracker.build(
+                    tx,
                     self.as_python_constant()(
                         *[x.as_python_constant() for x in args],
                         **{k: v.as_python_constant() for k, v in kwargs.items()},
@@ -1917,6 +1898,9 @@ class StringFormatVariable(VariableTracker):
 
     _nonvar_fields = {"format_string", *VariableTracker._nonvar_fields}
 
+    def python_type(self) -> type:
+        return str
+
     @classmethod
     def create(
         cls,
@@ -1952,6 +1936,26 @@ def __init__(
     def __repr__(self) -> str:
         return f"{self.__class__.__name__}({self.format_string!r}, {self.sym_args!r}, {self.sym_kwargs!r})"
 
+    @staticmethod
+    def _debug_format_arg(arg: VariableTracker) -> object:
+        try:
+            return arg.as_python_constant()
+        except Exception:
+            return arg.debug_repr()
+
+    def debug_repr(self) -> str:
+        try:
+            rendered = self.format_string.format(
+                *[self._debug_format_arg(arg) for arg in self.sym_args],
+                **{
+                    key: self._debug_format_arg(value)
+                    for key, value in self.sym_kwargs.items()
+                },
+            )
+        except Exception:
+            return repr(self)
+        return repr(rendered)
+
     def reconstruct(self, codegen: "PyCodegen") -> None:
         codegen.add_push_null(
             lambda: codegen.extend_output(
@@ -1971,11 +1975,17 @@ def reconstruct(self, codegen: "PyCodegen") -> None:
 
 
 class ObjectVariable(VariableTracker):
+    # PyBaseObject_Type: https://github.com/python/cpython/blob/v3.13.0/Objects/typeobject.c#L7243
+    _cpython_type = object
+
     # placeholder for unknown / opaque values
     def __init__(self, value: object, **kwargs: Any) -> None:
         super().__init__(**kwargs)
         self.value = value
 
+    def get_real_python_backed_value(self) -> object:
+        return self.value
+
     def python_type(self) -> type[object]:
         return object
 
@@ -1990,6 +2000,9 @@ def __init__(self, value: Any, **kwargs: Any) -> None:
         super().__init__(**kwargs)
         self.value = value
 
+    def python_type(self) -> type:
+        return type(self.value)
+
     @staticmethod
     def is_reorderable_logging_function(
         obj: Any,
@@ -2000,16 +2013,15 @@ def is_reorderable_logging_function(
             and obj in torch._dynamo.config.reorderable_logging_functions
         )
 
-    # type: ignore[override]
     def call_function(
         self,
         tx: "InstructionTranslator",
         args: Sequence[VariableTracker],
         kwargs: dict[str, VariableTracker],
-    ) -> None:
+    ) -> VariableTracker:
         if tx.export:
             # For export cases, we can just make debugging functions no-ops
-            return
+            return ConstantVariable.create(None)
 
         if not self.can_reorder_logs(self.value, args, kwargs):
             unimplemented(
@@ -2023,6 +2035,7 @@ def call_function(
             )
 
         tx.debug_locals.append((self, list(args)))
+        return ConstantVariable.create(None)
 
     def reconstruct(self, codegen: "PyCodegen") -> None:
         assert self.source is not None
@@ -2058,13 +2071,19 @@ def __init__(self, value: Any, **kwargs: Any) -> None:
         super().__init__(**kwargs)
         self.value = value
 
+    def python_type(self) -> type:
+        return type(self.value)
+
+    def get_real_python_backed_value(self) -> Any:
+        return self.value
+
     def call_function(
         self,
         tx: "InstructionTranslator",
         args: Sequence[VariableTracker],
         kwargs: dict[str, VariableTracker],
     ) -> VariableTracker:
-        return variables.CONSTANT_VARIABLE_NONE
+        return variables.ConstantVariable.create(None)
 
 
 class LoggingLoggerVariable(VariableTracker):
@@ -2076,6 +2095,12 @@ def __init__(self, value: logging.Logger, **kwargs: Any) -> None:
         super().__init__(**kwargs)
         self.value = value
 
+    def python_type(self) -> type:
+        return type(self.value)
+
+    def get_real_python_backed_value(self) -> logging.Logger:
+        return self.value
+
     def call_method(
         self,
         tx: "InstructionTranslator",
@@ -2085,7 +2110,7 @@ def call_method(
     ) -> VariableTracker:
         if tx.export:
             # For export cases, we can just make logging functions no-ops.
-            return variables.CONSTANT_VARIABLE_NONE
+            return variables.ConstantVariable.create(None)
 
         method = getattr(self.value, name, None)
         function = getattr(method, "__func__", None)
@@ -2094,7 +2119,7 @@ def call_method(
         ignore_set = torch._dynamo.config.ignore_logging_functions
 
         if method in ignore_set or function in ignore_set:
-            return variables.CONSTANT_VARIABLE_NONE
+            return variables.ConstantVariable.create(None)
 
         unimplemented(
             gb_type="logging.Logger method not supported for non-export cases",
@@ -2165,7 +2190,7 @@ def call_method(
         result = getattr(self.value, name)(*cargs, **ckwargs)
 
         if variables.ConstantVariable.is_literal(result):
-            return variables.ConstantVariable.create(result)
+            return VariableTracker.build(tx, result)
         if isinstance(result, re.Match):
             return ConstantLikeVariable(result)
 
@@ -2188,8 +2213,8 @@ def var_getattr(self, tx: "InstructionTranslator", name: str) -> VariableTracker
             # things like x.dtype.type
             return NumpyVariable(result)
         if variables.ConstantVariable.is_literal(result):
-            return variables.ConstantVariable.create(result)
-        return GetAttrVariable(self, name)
+            return VariableTracker.build(tx, result)
+        return GetAttrVariable(self, name, py_type=type(result))
 
 
 class TorchVersionVariable(ConstantLikeVariable):
@@ -2224,6 +2249,9 @@ class RandomClassVariable(VariableTracker):
     def __init__(self, **kwargs: Any) -> None:
         super().__init__(**kwargs)
 
+    def python_type(self) -> type:
+        return type
+
     def call_function(
         self,
         tx: "InstructionTranslator",
@@ -2239,7 +2267,7 @@ def call_function(
                     *graph_break_hints.USER_ERROR,
                 ],
             )
-        seed = variables.CONSTANT_VARIABLE_NONE if len(args) == 0 else args[0]
+        seed = variables.ConstantVariable.create(None) if len(args) == 0 else args[0]
         return RandomVariable(
             seed=seed, mutation_type=variables.base.ValueMutationNew()
         )
@@ -2253,6 +2281,8 @@ class RandomVariable(VariableTracker):
     Assumes that random objects behave the same given a set seed or state.
     """
 
+    _cpython_type = random.Random
+
     _nonvar_fields = {
         "random",
         *VariableTracker._nonvar_fields,
@@ -2347,13 +2377,13 @@ def call_method(
                 *[x.as_python_constant() for x in args],
                 **{key: val.as_python_constant() for key, val in kwargs.items()},
             )
-            return variables.CONSTANT_VARIABLE_NONE
+            return variables.ConstantVariable.create(None)
         elif name == "getstate":
             return self.wrap_state(self.random.getstate())
         elif name == "setstate":
             tx.output.side_effects.mutation(self)
             self.random.setstate(self.unwrap_state(args[0]))
-            return variables.CONSTANT_VARIABLE_NONE
+            return variables.ConstantVariable.create(None)
         elif name in self._supported_fn_names:
             tx.output.side_effects.mutation(self)
             state = self.random.getstate()
@@ -2393,8 +2423,11 @@ def reconstruct(self, codegen: "PyCodegen") -> None:
 
 
 class WeakRefVariable(VariableTracker):
+    def python_type(self) -> type:
+        return weakref.ref
+
     @staticmethod
-    # pyrefly: ignore[bad-param-name-override]
+    # pyrefly: ignore [bad-override, bad-param-name-override]
     def build(
         tx: "InstructionTranslator",
         weakref_value: weakref.ReferenceType[Any],
diff --git a/torch/_dynamo/variables/nn_module.py b/torch/_dynamo/variables/nn_module.py
index 82b24288d70df..ece303e249f60 100644
--- a/torch/_dynamo/variables/nn_module.py
+++ b/torch/_dynamo/variables/nn_module.py
@@ -28,13 +28,21 @@
 import types
 from collections.abc import Iterable, Sequence
 from contextlib import contextmanager, nullcontext
-from typing import Any, Optional, TYPE_CHECKING
+from typing import Any, TYPE_CHECKING
 
 import torch.nn
 from torch._guards import Source
 
 from .. import graph_break_hints, trace_rules, variables
-from ..exc import raise_observed_exception, unimplemented, UnspecializeRestartAnalysis
+from ..exc import (
+    handle_observed_exception,
+    ObservedAttributeError,
+    raise_observed_exception,
+    raise_type_error,
+    unimplemented,
+    UnspecializeRestartAnalysis,
+    Unsupported,
+)
 from ..guards import GuardBuilder, install_guard
 from ..mutation_guard import GenerationTracker
 from ..source import (
@@ -47,6 +55,7 @@
     UnspecializedNNModuleSource,
 )
 from ..utils import (
+    enumerate_items_with_dict_position,
     get_custom_getattr,
     get_fake_value,
     is_lazy_module,
@@ -62,10 +71,9 @@
     unpatched_nn_module_call,
     unpatched_nn_module_call_impl,
 )
-from .base import raise_type_error_exc, typestr, ValueMutationNew, VariableTracker
+from .base import typestr, ValueMutationNew, VariableTracker
 from .functions import invoke_and_store_as_constant
 from .lazy import LazyVariableTracker
-from .lists import SliceVariable
 from .user_defined import UserDefinedObjectVariable
 
 
@@ -73,6 +81,7 @@
     from torch._dynamo.symbolic_convert import InstructionTranslator
 
     from .constant import ConstantVariable
+    from .dicts import DunderDictVariable
 
 
 def initialize_lazy_module(
@@ -107,16 +116,12 @@ def convert_to_fake(x: Any) -> Any:
         fake_kwargs = {k: convert_to_fake(v) for k, v in proxy_kwargs.items()}
         try:
             mod._infer_parameters(mod, fake_args, fake_kwargs)  # type: ignore[operator]
-        except AttributeError as e:
+        except AttributeError:
             # Re-raise with the original error message from the AttributeError
             raise_observed_exception(
                 AttributeError,
                 tx,
-                args=[
-                    str(e)
-                    if str(e)
-                    else "AttributeError during lazy module initialization"
-                ],
+                args=["AttributeError during lazy module initialization"],
             )
 
 
@@ -142,7 +147,7 @@ def record_nn_module_stack(
 
 
 def guard_to_detect_forward_monkeypatching(
-    source: Optional[Source], mod: torch.nn.Module
+    source: Source | None, mod: torch.nn.Module
 ) -> None:
     # Users sometimes patch the forward method of a nn module instance to
     # perform optimizations like quantization. Though this is not a good
@@ -198,6 +203,11 @@ def __init__(
         self.source: Source = self.source
         self.nn_module_stack_source = self.source
 
+    def get_dict_vt(self, tx: "InstructionTranslator") -> "DunderDictVariable":
+        if not hasattr(self, "dict_vt"):
+            self.dict_vt = variables.DunderDictVariable.create(tx, self)
+        return self.dict_vt
+
     def get_nn_module_stack_source(self) -> Source:
         res = self.nn_module_stack_source or self.source
         assert res
@@ -209,6 +219,22 @@ def set_nn_module_stack_source(self, source: Source) -> None:
     def python_type(self) -> type:
         return self.module_type
 
+    def get_real_python_backed_value(self) -> object:
+        return self.value
+
+    def bool_impl(self, tx: "InstructionTranslator") -> VariableTracker:
+        """nb_bool for nn.Module.
+
+        nn.Module itself has no __bool__ or __len__, so bare modules are always
+        truthy.  Subclasses like ModuleList/ModuleDict define __len__, so
+        bool(module) calls PyObject_IsTrue which falls through nb_bool (NULL)
+        to sq_length/mp_length.  We evaluate on the real module to capture this.
+        """
+        from .constant import ConstantVariable
+
+        mod = tx.output.get_submodule(self.module_key)
+        return ConstantVariable.create(bool(mod))
+
     def _wrap_submodule(
         self,
         tx: "InstructionTranslator",
@@ -225,7 +251,7 @@ def unpack_var_sequence(self, tx: "InstructionTranslator") -> list[VariableTrack
         result: list[VariableTracker] = []
         if isinstance(base, torch.nn.ModuleDict):
             for name, submod in base.items():
-                name_var = variables.ConstantVariable.create(name)
+                name_var = VariableTracker.build(tx, name)
                 tx.output.register_attr_or_module(
                     submod,
                     self.module_key,
@@ -255,11 +281,9 @@ def call_obj_hasattr(
         mod = tx.output.get_submodule(self.module_key)
         result = hasattr(mod, name)
         install_guard(
-            NNModuleSource(AttrSource(self.source, name)).make_guard(
-                GuardBuilder.HASATTR
-            )
+            self.source.make_guard(functools.partial(GuardBuilder.HASATTR, attr=name))
         )
-        return variables.ConstantVariable.create(result)
+        return VariableTracker.build(tx, result)
 
     def is_training(self, tx: "InstructionTranslator") -> bool:
         mod = tx.output.get_submodule(self.module_key)
@@ -275,49 +299,39 @@ def convert_to_unspecialized(self, tx: "InstructionTranslator") -> None:
             GenerationTracker.mark_class_dynamic(type(mod))
         raise UnspecializeRestartAnalysis
 
-    def has_key_in_generic_dict(self, tx: "InstructionTranslator", key: str) -> bool:
-        base = tx.output.get_submodule(self.module_key)
-
-        if object_has_getattribute(base):
-            unimplemented(
-                gb_type="Custom __getattribute__ in nn.Module dict key check",
-                context=f"has_key_in_generic_dict {self} {key}",
-                explanation="Dynamo does not support checking key existence "
-                "on `nn.Module` instances that have a custom "
-                "`__getattribute__` method defined.",
-                hints=[
-                    "Avoid defining `__getattribute__` in your module.",
-                    *graph_break_hints.SUPPORTABLE,
-                ],
-            )
-
-        if tx.output.side_effects.has_pending_mutation_of_attr(self, key):
-            mutated_attr = tx.output.side_effects.load_attr(self, key, deleted_ok=True)
-            return not isinstance(mutated_attr, variables.DeletedVariable)
-
-        base_dict = object.__getattribute__(base, "__dict__")
-        return key in base_dict
-
     def _custom_getattr_fallback(
         self,
         base: torch.nn.Module,
         tx: "InstructionTranslator",
         name: str,
         obj_source: Source,
-    ) -> Optional[VariableTracker]:
+    ) -> VariableTracker | None:
         """Check for a __getattr__ and handle it specially if it is implemented"""
         if object_has_getattribute(base):
-            unimplemented(
-                gb_type="Custom __getattribute__ in nn.Module attribute access",
-                context=f"var_getattr {self} {name}",
-                explanation="Dynamo does not support checking key existence "
-                "on `nn.Module` instances that have a custom "
-                "`__getattribute__` method defined.",
-                hints=[
-                    "Avoid defining `__getattribute__` in your module.",
-                    *graph_break_hints.SUPPORTABLE,
-                ],
+            getattribute_fn = inspect.getattr_static(type(base), "__getattribute__")
+            new_source = (
+                AttrSource(obj_source, "__getattribute__") if obj_source else None
             )
+            try:
+                return variables.UserMethodVariable(
+                    getattribute_fn,
+                    self,
+                    source=new_source,
+                ).call_function(tx, [variables.ConstantVariable.create(name)], {})
+            except ObservedAttributeError:
+                handle_observed_exception(tx)
+            except Unsupported:
+                unimplemented(
+                    gb_type="Custom __getattribute__ in nn.Module attribute access",
+                    context=f"var_getattr {self} {name}",
+                    explanation="Dynamo could not trace through the custom "
+                    "`__getattribute__` method on this `nn.Module`.",
+                    hints=[
+                        "Simplify your `__getattribute__` implementation, "
+                        "or replace it with a targeted `@property`.",
+                        *graph_break_hints.SUPPORTABLE,
+                    ],
+                )
 
         getattr_fn = get_custom_getattr(base, ignore_nn_module_getattr=True)
         if getattr_fn is None:
@@ -345,13 +359,16 @@ def _custom_getattr_fallback(
         options = {"source": AttrSource(obj_source, "__getattr__")}
 
         return variables.UserMethodVariable(getattr_fn, self, **options).call_function(
-            tx, [variables.ConstantVariable.create(name)], {}
+            tx, [VariableTracker.build(tx, name)], {}
         )
 
     def var_getattr(self, tx: "InstructionTranslator", name: str) -> VariableTracker:
         source = self.source and AttrSource(self.source, name)
 
         base = tx.output.get_submodule(self.module_key)
+        # NB: We look up attributes in __dict__ directly, bypassing any custom
+        # __getattribute__. Custom __getattribute__ is only traced through as a
+        # fallback (via _custom_getattr_fallback) for attributes not found here.
         base_dict = object.__getattribute__(base, "__dict__")
         object_member = True
         all_class_attribute_names = set()
@@ -369,7 +386,7 @@ def var_getattr(self, tx: "InstructionTranslator", name: str) -> VariableTracker
             )
 
         if name == "__dict__":
-            return variables.GetAttrVariable(self, name, source=source)
+            return self.get_dict_vt(tx)
 
         subobj = None
         if name in base_dict:
@@ -406,7 +423,7 @@ def var_getattr(self, tx: "InstructionTranslator", name: str) -> VariableTracker
             guard_to_detect_forward_monkeypatching(self.source, base)
 
         if name == "__class__" and not object_member:
-            return variables.UserDefinedClassVariable(base.__class__, source=source)
+            return VariableTracker.build(tx, base.__class__, source=source)
 
         if object_member:
             out = VariableTracker.build(tx, subobj, NNModuleSource(source))  # type: ignore[arg-type]
@@ -460,7 +477,7 @@ def var_getattr(self, tx: "InstructionTranslator", name: str) -> VariableTracker
                     ],
                 )
 
-        return variables.GetAttrVariable(self, name, source=source)
+        return super().var_getattr(tx, name)
 
     def call_function(
         self,
@@ -571,6 +588,104 @@ def call_function(
                     kwargs,
                 )
 
+    def mp_subscript_impl(
+        self,
+        tx: "InstructionTranslator",
+        key: "VariableTracker",
+    ) -> "VariableTracker":
+        # nn.Module containers (ModuleList/Dict/Sequential/ParameterDict/ParameterList)
+        # These are Python-level __getitem__, not CPython C slots.
+        # TODO(follow-up): add tests for ModuleList negative index, ModuleList/Sequential
+        # slice, ModuleDict missing key, invalid key types
+        from .lists import SliceVariable
+        from .tensor import SymNodeVariable
+
+        module = tx.output.get_submodule(self.module_key)
+
+        builtin_supported = (
+            torch.nn.ModuleDict.__getitem__,
+            torch.nn.ModuleList.__getitem__,
+            torch.nn.ParameterDict.__getitem__,
+            torch.nn.ParameterList.__getitem__,
+            torch.nn.Sequential.__getitem__,
+        )
+        # pyrefly: ignore[missing-attribute]
+        if type(module).__getitem__ not in builtin_supported:
+            if not (
+                key.is_python_constant()
+                and isinstance(key.as_python_constant(), (str, int))
+            ):
+                unimplemented(
+                    gb_type="Invalid or non-const argument in nn.Module __getitem__",
+                    context=f"mp_subscript_impl: {self} {key}",
+                    explanation="Dynamo does not support calling "
+                    f"method `__getitem__` of ``nn.Module`` {module} with a non-constant or non-(str, int) key.",
+                    hints=["Use constant arguments of type str or int for __getitem__"],
+                )
+            fn = module.__getitem__.__func__  # pyrefly: ignore[missing-attribute]
+
+            assert isinstance(fn, types.FunctionType)
+
+            src = AttrSource(AttrSource(self.source, "__getitem__"), "__func__")  # type: ignore[arg-type]
+            return tx.inline_user_function_return(
+                variables.UserFunctionVariable(fn, source=src),
+                [self, key],
+                {},
+            )
+
+        if isinstance(key, SliceVariable):
+            # TODO(anijain2305,export-team) - Remove this if condition when inlining of inbuilt nn modules is
+            # enabled for export.
+            if tx.output.export:
+                result = []
+                keys = list(range(len(module)))[key.as_python_constant()]  # type: ignore[arg-type]
+                for idx, submod in enumerate(module[key.as_python_constant()]):  # type: ignore[arg-type]
+                    k = keys[idx]
+                    src = NNModuleSource(GetItemSource(self.source, k))
+                    result.append(
+                        tx.output.register_attr_or_module(
+                            submod,
+                            k,
+                            source=src,
+                        )
+                    )
+
+                new_module = module[key.as_python_constant()]  # type: ignore[index]
+                new_module_variable = tx.output.register_attr_or_module(
+                    new_module,
+                    f"{self}.__getitem__(slice)",
+                    source=NNModuleSource(
+                        GetItemSource(self.source, key.as_python_constant())
+                    ),
+                )
+                return new_module_variable
+            else:
+                # slice on nn module results in a creation of new module instance, so we need to make it sourceless.
+                # Convert to unspecialized so that UnspecializedNNModule variable can take care of it.
+                self.convert_to_unspecialized(tx)
+
+        key_value = 0
+        if isinstance(key, SymNodeVariable):
+            key_value = key.evaluate_expr(tx.output)
+        elif key.is_python_constant():
+            key_value = key.as_python_constant()
+        else:
+            unimplemented(
+                gb_type="Unsupported key type for nn.Module.__getitem__",
+                context=f"mp_subscript_impl: {self} {key}",
+                explanation="Dynamo does not support getitem on "
+                "`nn.Module` with non-constant key.",
+                hints=[],
+            )
+
+        submod = module[key_value]  # type: ignore[index]
+        return tx.output.register_attr_or_module(
+            submod,
+            self.module_key,
+            key_value,
+            source=NNModuleSource(GetItemSource(self.source, key_value)),
+        )
+
     def call_method(
         self,
         tx: "InstructionTranslator",
@@ -579,7 +694,8 @@ def call_method(
         kwargs: dict[str, VariableTracker],
         constant: bool = False,
     ) -> VariableTracker:
-        from . import ConstantVariable, ListIteratorVariable, TupleVariable
+        from . import ListIteratorVariable, TupleVariable
+        from .constant import ConstantVariable
 
         key = self.module_key
         module = tx.output.get_submodule(key)
@@ -630,12 +746,12 @@ def generic_call_method_helper(name: str) -> VariableTracker:
 
         if name == "_get_item_by_idx":
             if not args[1].is_python_constant():
-                raise_type_error_exc(
+                raise_type_error(
                     tx,
                     f"``nn.Module`` {module}'s call method {name} requires a constant index argument",
                 )
             if not isinstance(args[0], TupleVariable):
-                raise_type_error_exc(
+                raise_type_error(
                     tx,
                     f"``nn.Module`` {module}'s call method {name} requires a tuple as first argument",
                 )
@@ -699,7 +815,7 @@ def wrap_values(
         def named_embed(name: str, obj: Any) -> "variables.TupleVariable":
             return TupleVariable(
                 [
-                    ConstantVariable.create(name),
+                    VariableTracker.build(tx, name),
                     tx.output.register_attr_or_module(
                         obj,
                         key,
@@ -791,7 +907,7 @@ def gen_source(source: Source, name: str) -> Source:
             result = []
             # pyrefly: ignore[not-iterable]
             for tmp in module:
-                result.append(ConstantVariable.create(tmp))
+                result.append(VariableTracker.build(tx, tmp))
             return ListIteratorVariable(result, mutation_type=ValueMutationNew())
         elif name == "values":
             if args or kwargs:
@@ -814,15 +930,6 @@ def gen_source(source: Source, name: str) -> Source:
             for name, submod in module.items():  # type: ignore[operator]
                 items_result.append(named_embed(name, submod))
             return ListIteratorVariable(items_result, mutation_type=ValueMutationNew())
-        elif name == "__len__":
-            if args or kwargs:
-                raise_args_mismatch(
-                    tx,
-                    name,
-                    "0 args and 0 kwargs",
-                    f"{len(args)} args and {len(kwargs)} kwargs",
-                )
-            return ConstantVariable.create(len(module))  # type: ignore[arg-type]
         elif name == "__iter__":
             return ListIteratorVariable(
                 self.unpack_var_sequence(tx), mutation_type=ValueMutationNew()
@@ -833,106 +940,8 @@ def gen_source(source: Source, name: str) -> Source:
             and args
             and args[0].is_python_constant()
         ):
-            return ConstantVariable.create(
-                args[0].as_python_constant() in module._modules
-            )
-        elif name == "__getitem__":
-            if kwargs or len(args) != 1:
-                raise_args_mismatch(
-                    tx,
-                    name,
-                    "1 args and 0 kwargs",
-                    f"{len(args)} args and {len(kwargs)} kwargs",
-                )
-            builtin_supported = (
-                torch.nn.ModuleDict.__getitem__,
-                torch.nn.ModuleList.__getitem__,
-                torch.nn.ParameterDict.__getitem__,
-                torch.nn.ParameterList.__getitem__,
-                torch.nn.Sequential.__getitem__,
-            )
-            # pyrefly: ignore[missing-attribute]
-            if type(module).__getitem__ not in builtin_supported:
-                if not (
-                    args[0].is_python_constant()
-                    and isinstance(args[0].as_python_constant(), (str, int))
-                ):
-                    unimplemented(
-                        gb_type="Invalid or non-const argument in nn.Module __getitem__",
-                        context=f"call_method: {self} {name} {args} {kwargs}",
-                        explanation="Dynamo does not support calling "
-                        f"method `{name}` of ``nn.Module`` {module} with a non-constant or non-(str, int) key.",
-                        hints=[
-                            "Use constant arguments of type str or int for __getitem__"
-                        ],
-                    )
-                fn = getattr(module, name).__func__
-
-                assert isinstance(fn, types.FunctionType)
-
-                src = AttrSource(AttrSource(self.source, name), "__func__")  # type: ignore[arg-type]
-                return tx.inline_user_function_return(
-                    variables.UserFunctionVariable(fn, source=src),
-                    [self] + list(args),
-                    kwargs,
-                )
-
-            if isinstance(args[0], SliceVariable):
-                # TODO(anijain2305,export-team) - Remove this if condition when inlining of inbuilt nn modules is
-                # enabled for export.
-                if tx.output.export:
-                    # Build a TupleVariable of NNModules
-                    result = []
-
-                    # Turn the slice into the list of integers
-                    keys = list(range(len(module)))[args[0].as_python_constant()]  # type: ignore[arg-type]
-                    for idx, submod in enumerate(module[args[0].as_python_constant()]):  # type: ignore[arg-type]
-                        key = keys[idx]
-                        src = NNModuleSource(GetItemSource(self.source, key))
-                        result.append(
-                            tx.output.register_attr_or_module(
-                                submod,
-                                key,
-                                source=src,
-                            )
-                        )
-
-                    new_module = module[args[0].as_python_constant()]  # type: ignore[index]
-                    new_module_variable = tx.output.register_attr_or_module(
-                        new_module,
-                        f"{self}.__getitem__(slice)",
-                        source=NNModuleSource(
-                            GetItemSource(self.source, args[0].as_python_constant())
-                        ),
-                    )
-                    return new_module_variable
-                else:
-                    # slice on nn module results in a creation of new module instance, so we need to make it sourceless.
-                    # Convert to unspecialized so that UnspecializedNNModule variable can take care of it.
-                    self.convert_to_unspecialized(tx)
-
-            from .tensor import SymNodeVariable
-
-            key_value = 0
-            if isinstance(args[0], SymNodeVariable):
-                key_value = args[0].evaluate_expr(tx.output)
-            elif args[0].is_python_constant():
-                key_value = args[0].as_python_constant()
-            else:
-                unimplemented(
-                    gb_type="Unsupported key type for nn.Module.__getitem__",
-                    context=f"call_method: {self} {name} {args} {kwargs}",
-                    explanation="Dynamo does not support getitem on "
-                    "`nn.Module` with non-constant key.",
-                    hints=[],
-                )
-
-            submod = module[key_value]  # type: ignore[index]
-            return tx.output.register_attr_or_module(
-                submod,
-                self.module_key,
-                key_value,
-                source=NNModuleSource(GetItemSource(self.source, key_value)),
+            return VariableTracker.build(
+                tx, args[0].as_python_constant() in module._modules
             )
         elif (
             name == "_get_abs_string_index"
@@ -964,6 +973,11 @@ def gen_source(source: Source, name: str) -> Source:
         else:
             return super().call_method(tx, name, list(args), kwargs)
 
+    def sq_length(self, tx: "InstructionTranslator") -> "VariableTracker":
+        """Sequence length for container modules (e.g., nn.Sequential)."""
+        module = tx.output.get_submodule(self.module_key)
+        return VariableTracker.build(tx, len(module))  # type: ignore[arg-type]
+
 
 class UnspecializedNNModuleVariable(UserDefinedObjectVariable):
     _nonvar_fields = {
@@ -982,16 +996,6 @@ class UnspecializedNNModuleVariable(UserDefinedObjectVariable):
     """
 
     def __init__(self, value: torch.nn.Module, **kwargs: Any) -> None:
-        if type(value) is torch.jit._script.RecursiveScriptModule:
-            unimplemented(
-                gb_type="UnspecializedNNModuleVariable wrapped around ScriptModules unsupported",
-                context=str(value),
-                explanation="ScriptModules aren't supported in UnspecializedNNModuleVariable"
-                " because their .forward function isn't a static member of their type.",
-                hints=[
-                    *graph_break_hints.DIFFICULT,
-                ],
-            )
         if "value_type" in kwargs:
             lazy_value_to_become = getattr(kwargs["value_type"], "cls_to_become", None)
             if type(value) is lazy_value_to_become:
@@ -1089,7 +1093,9 @@ def call_function(
             and istype(mod._call_impl, types.MethodType)  # type: ignore[attr-defined]
             and mod.__call__.__func__ is unpatched_nn_module_call  # type: ignore[operator]
             and mod._call_impl.__func__ is unpatched_nn_module_call_impl  # type: ignore[attr-defined]
-            and "forward" not in mod.__dict__
+            # Consult pending STORE_ATTR side effects too. During tracing the
+            # patched forward may not be visible in mod.__dict__ yet.
+            and not self.has_key_in_generic_dict(tx, "forward")
         ):
             forward_method = inspect.getattr_static(mod, "forward")
             if isinstance(forward_method, types.FunctionType):
@@ -1118,7 +1124,7 @@ def call_function(
                     fn = self.value_type.forward  # type: ignore[attr-defined]
 
         if self.source:
-            source = self.get_source_by_walking_mro(name)
+            source = self.get_source_by_walking_mro(tx, name)
         else:
             source = None
 
@@ -1157,21 +1163,23 @@ def call_method(
         if name in ["_call_impl", "_wrapped_call_impl"]:
             fn = getattr(self.value_type, name)
             if self.source:
-                source = self.get_source_by_walking_mro(name)
+                source = self.get_source_by_walking_mro(tx, name)
             else:
                 source = None
 
             fn_vt = VariableTracker.build(tx, fn, source=source, realize=True)
             return fn_vt.call_function(tx, [self] + list(args), kwargs)
 
-        if name not in getattr(self.value, "__dict__", {}):
+        if not self.has_key_in_generic_dict(tx, name):
             try:
                 method = inspect.getattr_static(type(self.value), name)
             except AttributeError:
                 method = None
 
             if isinstance(method, staticmethod):
-                source = AttrSource(self.get_source_by_walking_mro(name), "__func__")
+                source = AttrSource(
+                    self.get_source_by_walking_mro(tx, name), "__func__"
+                )
                 fn_vt = VariableTracker.build(
                     tx, method.__func__, source=source, realize=True
                 )
@@ -1241,7 +1249,7 @@ def call_method(
 
     def getattr_helper(
         self, tx: "InstructionTranslator", field: str, name_vt: VariableTracker
-    ) -> Optional[VariableTracker]:
+    ) -> VariableTracker | None:
         dict_vt = self.var_getattr(tx, field)
         if isinstance(dict_vt, variables.ConstDictVariable):
             return dict_vt.maybe_getitem_const(name_vt)
@@ -1296,7 +1304,7 @@ def build_key_value(
                 i: int, k: Any, v: Any
             ) -> tuple[VariableTracker, VariableTracker]:
                 # Make key sourceless to avoid any guard on it
-                key = variables.ConstantVariable.create(k)
+                key = VariableTracker.build(tx, k)
 
                 # Instead of using dict[key] to access the value, use a dict[dict.keys()[index]] to access the
                 # value. This removes the reliance on the actual key value.
@@ -1306,7 +1314,8 @@ def build_key_value(
                 return key, value
 
             result = dict(
-                build_key_value(i, k, v) for i, (k, v) in enumerate(hooks_dict.items())
+                build_key_value(i, k, v)
+                for i, k, v in enumerate_items_with_dict_position(hooks_dict)
             )
 
             return variables.NNModuleHooksDictVariable(
diff --git a/torch/_dynamo/variables/object_protocol.py b/torch/_dynamo/variables/object_protocol.py
new file mode 100644
index 0000000000000..217a7a796f582
--- /dev/null
+++ b/torch/_dynamo/variables/object_protocol.py
@@ -0,0 +1,319 @@
+"""
+Dynamo implementations of CPython's PyObject_* default slot algorithms.
+
+Analogous to CPython's Objects/object.c, this module holds the general
+dispatch machinery that is independent of any specific type.
+Per-type hook implementations (bool_impl, richcompare_impl, etc.)
+live in their respective VT files.
+"""
+
+from functools import lru_cache
+from typing import TYPE_CHECKING
+
+from torch._C._dynamo import (
+    get_type_slots,
+    has_slot,
+    PyMappingSlots,
+    PyNumberSlots,
+    PySequenceSlots,
+)
+
+from .. import graph_break_hints
+from ..exc import (
+    handle_observed_exception,
+    ObservedTypeError,
+    raise_observed_exception,
+    raise_type_error,
+    unimplemented,
+)
+from ..utils import istype
+from .base import NO_SUCH_SUBOBJ, VariableTracker
+from .constant import ConstantVariable
+
+
+if TYPE_CHECKING:
+    from ..symbolic_convert import InstructionTranslator
+
+
+def vt_identity_compare(
+    left: VariableTracker,
+    right: VariableTracker,
+) -> "VariableTracker | None":
+    """Try to determine Python identity (left is right) at trace time.
+
+    Returns ConstantVariable(True/False) if determinable, else None.
+    Mirrors the logic in BuiltinVariable's handle_is handler.
+    """
+    if left is right:
+        return ConstantVariable.create(True)
+
+    left_val = left.get_real_python_backed_value()
+    right_val = right.get_real_python_backed_value()
+    left_known = left_val is not NO_SUCH_SUBOBJ
+    right_known = right_val is not NO_SUCH_SUBOBJ
+
+    if left_known and right_known:
+        return (
+            ConstantVariable.create(True)
+            if left_val is right_val
+            else ConstantVariable.create(False)
+        )
+
+    # One side has a concrete backing object, the other doesn't — they can't
+    # be the same object.
+    if left_known != right_known:
+        return ConstantVariable.create(False)
+
+    # Mutable containers created during tracing: VT identity = Python identity.
+    from .dicts import ConstDictVariable
+    from .lists import ListVariable
+    from .sets import SetVariable
+
+    if isinstance(left, (ConstDictVariable, ListVariable, SetVariable)):
+        return ConstantVariable.create(False)
+
+    # Different Python types can never be the same object.
+    try:
+        if left.python_type() is not right.python_type():
+            return ConstantVariable.create(False)
+    except NotImplementedError:
+        pass
+
+    # Different exception types are never identical.
+    from .. import variables
+
+    if (
+        istype(left, variables.ExceptionVariable)
+        and istype(right, variables.ExceptionVariable)
+        and left.exc_type is not right.exc_type  # type: ignore[attr-defined]
+    ):
+        return ConstantVariable.create(False)
+
+    return None
+
+
+@lru_cache(maxsize=256)
+def _get_cached_slots(obj_type: type) -> tuple[int, int, int, int]:
+    """Get all type slots for a type (cached)."""
+    return get_type_slots(obj_type)
+
+
+def type_implements_sq_length(obj_type: type) -> bool:
+    """Check whether obj_type implements __len__ as sequence protocol"""
+    seq_slots, _, _, _ = _get_cached_slots(obj_type)
+    return has_slot(seq_slots, PySequenceSlots.SQ_LENGTH)
+
+
+def type_implements_mp_length(obj_type: type) -> bool:
+    """Check whether obj_type implements __len__ as mapping protocol"""
+    _, map_slots, _, _ = _get_cached_slots(obj_type)
+    return has_slot(map_slots, PyMappingSlots.MP_LENGTH)
+
+
+def type_implements_nb_bool(obj_type: type) -> bool:
+    """Check whether obj_type implements the nb_bool slot (i.e. has __bool__ or __len__)."""
+    _, _, number_slots, _ = _get_cached_slots(obj_type)
+    return has_slot(number_slots, PyNumberSlots.NB_BOOL)
+
+
+def type_implements_nb_int(obj_type: type) -> bool:
+    """Check whether obj_type implements the nb_int slot."""
+    _, _, number_slots, _ = _get_cached_slots(obj_type)
+    return has_slot(number_slots, PyNumberSlots.NB_INT)
+
+
+def type_implements_nb_index(obj_type: type) -> bool:
+    """Check whether obj_type implements the nb_index slot."""
+    _, _, number_slots, _ = _get_cached_slots(obj_type)
+    return has_slot(number_slots, PyNumberSlots.NB_INDEX)
+
+
+def type_implements_nb_float(obj_type: type) -> bool:
+    """Check whether obj_type implements the nb_float slot."""
+    _, _, number_slots, _ = _get_cached_slots(obj_type)
+    return has_slot(number_slots, PyNumberSlots.NB_FLOAT)
+
+
+def maybe_get_python_type(obj: VariableTracker) -> type:
+    try:
+        return obj.python_type()
+    except NotImplementedError:
+        unimplemented(
+            gb_type="Unsupported python_type() call",
+            context=f"{obj} does not implement python_type()",
+            explanation="This VariableTracker does not implement python_type(), "
+            "which is required for object protocol operations.",
+            hints=[
+                *graph_break_hints.DYNAMO_BUG,
+            ],
+        )
+
+
+def vt_mapping_size(
+    tx: "InstructionTranslator", obj: "VariableTracker"
+) -> "VariableTracker":
+    # ref: https://github.com/python/cpython/blob/v3.13.3/Objects/abstract.c#L2308-L2330
+    T = maybe_get_python_type(obj)
+    if type_implements_mp_length(T):
+        return obj.mp_length(tx)
+
+    if type_implements_sq_length(T):
+        raise_type_error(tx, f"{obj.python_type_name()} is not a mapping")
+
+    raise_type_error(tx, f"object of type {obj.python_type_name()} has no len()")
+
+
+def generic_len(
+    tx: "InstructionTranslator", obj: "VariableTracker"
+) -> "VariableTracker":
+    # ref: https://github.com/python/cpython/blob/v3.13.3/Objects/abstract.c#L53-L69
+    """
+    Implements PyObject_Size/PyObject_Length semantics for VariableTracker objects.
+    Dispatches to sq_length (sequences) or mp_length (mappings) depending on the VT type.
+    """
+
+    T = maybe_get_python_type(obj)
+    if type_implements_sq_length(T):
+        return obj.sq_length(tx)
+    return vt_mapping_size(tx, obj)
+
+
+def generic_bool(tx: "InstructionTranslator", obj: VariableTracker) -> VariableTracker:
+    """Mirrors PyObject_IsTrue.
+
+    https://github.com/python/cpython/blob/c09ccd9c429/Objects/object.c#L2135-L2158
+
+    Resolution order: constants → nb_bool → mp_length/sq_length → truthy.
+    """
+    from .constant import ConstantVariable
+
+    if obj.is_python_constant():
+        return ConstantVariable.create(bool(obj.as_python_constant()))
+
+    obj_type = maybe_get_python_type(obj)
+
+    if type_implements_nb_bool(obj_type):
+        result = obj.bool_impl(tx)
+        if result is not None:
+            return result
+
+    try:
+        length = generic_len(tx, obj)
+        from .tensor import SymNodeVariable
+
+        if isinstance(length, SymNodeVariable):
+            return SymNodeVariable.create(tx, length.as_proxy() > 0)
+        return ConstantVariable.create(length.as_python_constant() > 0)
+    except ObservedTypeError:
+        handle_observed_exception(tx)
+
+    return ConstantVariable.create(True)
+
+
+def vt_getitem(
+    tx: "InstructionTranslator",
+    obj: VariableTracker,
+    key: VariableTracker,
+) -> VariableTracker:
+    """CPython's PyObject_GetItem — dispatch to the type's mp_subscript/sq_item.
+
+    PyObject_GetItem: https://github.com/python/cpython/blob/62a6e898e01/Objects/abstract.c#L155-L206
+
+    CPython checks three branches in order:
+      1. tp_as_mapping->mp_subscript  (L161-166)
+      2. tp_as_sequence->sq_item      (L168-181) — only if key passes _PyIndex_Check
+      3. PyType_Check(o)              (L183-203) — type[int] → GenericAlias/__class_getitem__
+
+    Branch 1 is the common path (list, tuple, dict, range all have mp_subscript).
+    TODO(follow-up): use has_slot(map_slots, PyMappingSlots.MP_SUBSCRIPT) to gate
+    Branch 1 and has_slot(seq_slots, PySequenceSlots.SQ_ITEM) to gate Branch 2,
+    matching CPython's dispatch order.
+    TODO(follow-up): Branch 2 (sq_item) for C extension types that only have
+    tp_as_sequence (e.g. deque — Modules/_collectionsmodule.c:1888).
+    Branch 3 is handled by TypingVariable.mp_subscript_impl for typing module types
+    and by BuiltinVariable for builtin types like list[int].
+
+    Types that work via constant fold fallback (no dedicated mp_subscript_impl):
+    TODO(follow-up): str (unicode_subscript, Objects/unicodeobject.c:13809)
+    TODO(follow-up): bytes (bytes_subscript, Objects/bytesobject.c)
+    """
+    return obj.mp_subscript_impl(tx, key)
+
+
+def generic_int(tx: "InstructionTranslator", obj: VariableTracker) -> VariableTracker:
+    """Mirrors PyNumber_Long (int(x) dispatch).
+
+    https://github.com/python/cpython/blob/v3.13.0/Objects/abstract.c#L1520-L1632
+
+    Resolution: nb_int → nb_index → str/bytes/bytearray parsing → TypeError.
+    """
+    from .constant import ConstantVariable
+
+    # Fast path for int (sub)class instances — mirrors PyLong_Check at the
+    # top of PyNumber_Long (abstract.c:1531). Avoids infinite recursion for
+    # int subclasses like IntEnum whose __int__ calls int() again.
+    if obj.is_python_constant() and isinstance(obj.as_python_constant(), int):
+        return ConstantVariable.create(int(obj.as_python_constant()))
+
+    obj_type = maybe_get_python_type(obj)
+
+    if type_implements_nb_int(obj_type):
+        return obj.nb_int_impl(tx)
+
+    if type_implements_nb_index(obj_type):
+        return obj.nb_index_impl(tx)
+
+    # String/bytes/bytearray parsing fallback.
+    # https://github.com/python/cpython/blob/v3.13.0/Objects/abstract.c#L1598-L1612
+    if obj.is_python_constant() and isinstance(
+        obj.as_python_constant(), (str, bytes, bytearray)
+    ):
+        try:
+            return ConstantVariable.create(int(obj.as_python_constant()))
+        except ValueError as e:
+            raise_observed_exception(ValueError, tx, args=[str(e)])
+
+    raise_type_error(
+        tx,
+        f"int() argument must be a string, a bytes-like object "
+        f"or a real number, not '{obj.python_type_name()}'",
+    )
+
+
+def generic_float(tx: "InstructionTranslator", obj: VariableTracker) -> VariableTracker:
+    """Mirrors PyNumber_Float (float(x) dispatch).
+
+    https://github.com/python/cpython/blob/v3.13.0/Objects/abstract.c#L1635-L1692
+
+    Resolution: nb_float → nb_index → str parsing → TypeError.
+    """
+    from .constant import ConstantVariable
+
+    # Fast path: if the value is already a float constant, return it directly.
+    # Mirrors PyFloat_CheckExact fast path at the top of PyNumber_Float
+    # (abstract.c:1641-1643).
+    if obj.is_python_constant() and isinstance(obj.as_python_constant(), float):
+        return ConstantVariable.create(float(obj.as_python_constant()))
+
+    obj_type = maybe_get_python_type(obj)
+
+    if type_implements_nb_float(obj_type):
+        return obj.nb_float_impl(tx)
+
+    # https://github.com/python/cpython/blob/v3.13.0/Objects/abstract.c#L1674-L1685
+    if type_implements_nb_index(obj_type):
+        return obj.nb_index_impl(tx)
+
+    # PyFloat_FromString fallback — handles str and bytes.
+    # https://github.com/python/cpython/blob/v3.13.0/Objects/abstract.c#L1691
+    if obj.is_python_constant() and isinstance(obj.as_python_constant(), (str, bytes)):
+        try:
+            return ConstantVariable.create(float(obj.as_python_constant()))
+        except ValueError as e:
+            raise_observed_exception(ValueError, tx, args=[str(e)])
+
+    raise_type_error(
+        tx,
+        f"float() argument must be a string or a real number, "
+        f"not '{obj.python_type_name()}'",
+    )
diff --git a/torch/_dynamo/variables/optimizer.py b/torch/_dynamo/variables/optimizer.py
index a6a05236fc27e..137f6f158f580 100644
--- a/torch/_dynamo/variables/optimizer.py
+++ b/torch/_dynamo/variables/optimizer.py
@@ -23,7 +23,7 @@
 import logging
 import weakref
 from collections.abc import Iterable
-from typing import Any, Optional, TYPE_CHECKING
+from typing import Any, TYPE_CHECKING
 
 import torch
 from torch._dynamo.variables.tensor import TensorVariable
@@ -44,6 +44,7 @@
 from .base import VariableTracker
 from .constant import ConstantVariable
 from .dicts import ConstDictVariable
+from .hashable import HashableTracker
 from .lists import ListVariable
 from .misc import GetAttrVariable
 from .user_defined import UserDefinedObjectVariable
@@ -94,9 +95,9 @@ class OptimizerVariable(UserDefinedObjectVariable):
     def __init__(
         self,
         value: torch.optim.Optimizer,
-        grad_to_source: Optional[dict[Any, GradSource]] = None,
-        static_tensor_names: Optional[set[str]] = None,
-        tensor_to_source: Optional[dict[torch.Tensor, Source]] = None,
+        grad_to_source: dict[Any, GradSource] | None = None,
+        static_tensor_names: set[str] | None = None,
+        tensor_to_source: dict[torch.Tensor, Source] | None = None,
         **kwargs: Any,
     ) -> None:
         super().__init__(value, **kwargs)
@@ -148,7 +149,12 @@ def var_getattr(self, tx: "InstructionTranslator", name: str) -> VariableTracker
         # which will directly inline
         if name in ("_init_group"):
             assert self.source
-            return GetAttrVariable(self, name, source=AttrSource(self.source, name))
+            return GetAttrVariable(
+                self,
+                name,
+                py_type=type(getattr(self.value, name)),
+                source=AttrSource(self.source, name),
+            )
 
         if name == "param_groups":
             from ..decorators import mark_static_address
@@ -207,9 +213,7 @@ def safe_to_set_capturable(group: dict[str, Any]) -> bool:
             VariableTracker.build(tx, self.value.param_groups, source)
         )
         for param_group_vt in param_groups_vt.items:
-            key = ConstDictVariable._HashableTracker(
-                ConstantVariable.create("capturable")
-            )
+            key = HashableTracker(ConstantVariable.create("capturable"))
             param_group_vt.items[key] = ConstantVariable.create(True)
 
     def get_python_args(
diff --git a/torch/_dynamo/variables/script_object.py b/torch/_dynamo/variables/script_object.py
index 5109847497b43..e0d4f6cb45d11 100644
--- a/torch/_dynamo/variables/script_object.py
+++ b/torch/_dynamo/variables/script_object.py
@@ -18,15 +18,18 @@
 by limiting operations to known-safe patterns and failing fast for unsafe usage.
 """
 
+import enum
 import functools
 import inspect
+import types
 from collections.abc import Callable, Iterable, Sequence
-from typing import Any, Optional, TYPE_CHECKING, TypeVar
+from typing import Any, TYPE_CHECKING, TypeVar
 from typing_extensions import ParamSpec
 
 import torch
 import torch.utils._pytree as pytree
 from torch._guards import Source
+from torch._library.fake_class_registry import FakeScriptObject
 from torch._library.opaque_object import (
     get_member_type,
     is_opaque_reference_type,
@@ -39,7 +42,12 @@
 
 from .. import graph_break_hints
 from ..eval_frame import skip_code
-from ..exc import unimplemented, UnsafeScriptObjectError, Unsupported
+from ..exc import (
+    raise_observed_exception,
+    unimplemented,
+    UnsafeScriptObjectError,
+    Unsupported,
+)
 from ..source import AttrSource
 from ..utils import proxy_args_kwargs
 from .base import VariableTracker
@@ -82,6 +90,10 @@ class OpaqueObjectClassVariable(UserDefinedVariable):
     """
 
     def __init__(self, value: Any, **kwargs: Any) -> None:
+        assert not (isinstance(value, type) and issubclass(value, enum.Enum)), (
+            f"Enum class {value} should use UserDefinedClassVariable, "
+            "not OpaqueObjectClassVariable"
+        )
         super().__init__(**kwargs)
         self.value = value
 
@@ -95,7 +107,10 @@ def is_python_constant(self) -> bool:
         return False
 
     def is_python_hashable(self) -> bool:
-        return is_opaque_value_type(self.value)
+        return is_opaque_value_type(self.value)  # pyrefly: ignore[bad-argument-type]
+
+    def get_python_hash(self) -> int:
+        return hash(self.value)
 
     def as_proxy(self) -> Any:
         return self.value
@@ -122,12 +137,11 @@ def var_getattr(self, tx: "InstructionTranslator", name: str) -> VariableTracker
         elif isinstance(obj, property):
             obj = obj.__get__(None, self.value)  # pyrefly: ignore[no-matching-overload]
         elif hasattr(obj, "__get__"):
-            # Check for pybind11 static properties (common in PyTorch C++ bindings)
-            # Reference: https://github.com/python/mypy/blob/131f9d92da58294bb2f273425e8778bd7d5b861f/mypy/stubgenc.py#L590
-            type_name = type(obj).__name__
-            if type_name == "pybind11_static_property":
-                obj = obj.__get__(None, self.value)
+            if not isinstance(type(obj).__dict__.get("__get__"), types.FunctionType):
+                # C-level descriptors are safe to resolve dynamically.
+                obj = getattr(self.value, name)
             else:
+                type_name = type(obj).__name__
                 unimplemented(
                     gb_type="Unsupported descriptor on opaque class",
                     context=f"class={self.value}, attr={name}, descriptor={type_name}",
@@ -139,7 +153,7 @@ def var_getattr(self, tx: "InstructionTranslator", name: str) -> VariableTracker
                 )
 
         if ConstantVariable.is_literal(obj):
-            return ConstantVariable.create(obj)
+            return VariableTracker.build(tx, obj)
 
         source = AttrSource(self.source, name) if self.source else None
         return VariableTracker.build(tx, obj, source)
@@ -171,19 +185,35 @@ def call_function(
 
         var_args = TupleVariable(list(args))
         var_kwargs = ConstDictVariable(
-            {ConstantVariable(k): v for k, v in kwargs.items()}
+            {VariableTracker.build(tx, k): v for k, v in kwargs.items()}
         )
-        constant_args = var_args.as_python_constant()
-        constant_kwargs = var_kwargs.as_python_constant()
+        if should_hoist(self.value):
+            with tx.output.tracing_context.guards_context.skip_guard_install():
+                constant_args = var_args.as_python_constant()
+                constant_kwargs = var_kwargs.as_python_constant()
+        else:
+            constant_args = var_args.as_python_constant()
+            constant_kwargs = var_kwargs.as_python_constant()
         opaque_obj = self.value(  # pyrefly: ignore[not-callable]
             *constant_args, **constant_kwargs
         )
-        fake_script_obj = torch._library.fake_class_registry.maybe_to_fake_obj(
-            tx.output.fake_mode, opaque_obj
-        )
+
+        # Capture sources from the VT args so subgraph reuse can apply
+        # source replacement to resolve new ctor arg values on stamp-out.
+        ctor_arg_sources = tuple(getattr(a, "source", None) for a in args)
+
+        if is_opaque_value_type(type(opaque_obj)):
+            fake_script_obj = opaque_obj
+        else:
+            fake_script_obj = torch._library.fake_class_registry.maybe_to_fake_obj(
+                tx.output.fake_mode, opaque_obj
+            )
 
         return TorchScriptObjectVariable.create(
-            opaque_obj, fake_script_obj, (constant_args, constant_kwargs)
+            opaque_obj,
+            fake_script_obj,
+            (constant_args, constant_kwargs),
+            ctor_arg_sources=ctor_arg_sources,
         )
 
 
@@ -192,20 +222,40 @@ class TorchScriptObjectVariable(UserDefinedObjectVariable):
 
     @classmethod
     def is_matching_cls(cls, user_cls: type) -> bool:
-        return issubclass(user_cls, torch.ScriptObject) or is_opaque_type(user_cls)
+        return (
+            issubclass(user_cls, torch.ScriptObject)
+            or is_opaque_type(user_cls)
+            or issubclass(user_cls, FakeScriptObject)
+        )
 
     @staticmethod
     def create(
-        proxy: Proxy, value: Any, ctor_args_kwargs: Any = None, **options: Any
+        proxy: Proxy,
+        value: Any,
+        ctor_args_kwargs: Any = None,
+        ctor_arg_sources: tuple[Source | None, ...] | None = None,
+        **options: Any,
     ) -> "TorchScriptObjectVariable":
-        return TorchScriptObjectVariable(proxy, value, ctor_args_kwargs, **options)
+        assert not isinstance(value, enum.Enum), (
+            f"Enum {type(value)} should use UserDefinedObjectVariable, not TorchScriptObjectVariable"
+        )
+        out = TorchScriptObjectVariable(
+            proxy, value, ctor_args_kwargs, ctor_arg_sources=ctor_arg_sources, **options
+        )
+        if isinstance(proxy, torch.fx.Proxy) and proxy.node.op != "placeholder":
+            from torch._dynamo.symbolic_convert import InstructionTranslator
+
+            tx = InstructionTranslator.current_tx()
+            tx.output.current_tracer.record_proxyable_vt(out)
+        return out
 
     def __init__(
         self,
         proxy: Proxy,
         value: Any,
         ctor_args_kwargs: Any = None,
-        source: Optional[Source] = None,
+        source: Source | None = None,
+        ctor_arg_sources: tuple[Source | None, ...] | None = None,
         **kwargs: Any,
     ) -> None:
         super().__init__(value, **kwargs)
@@ -216,6 +266,9 @@ def __init__(
         # If the OpaqueObject is sourceless, then this is
         # the constant (args, kwargs) that Dynamo used to construct it.
         self.ctor_args_kwargs = ctor_args_kwargs
+        # Sources of the constructor args, used by subgraph reuse to
+        # resolve new values via source replacement on stamp-out.
+        self.ctor_arg_sources = ctor_arg_sources
 
     def as_proxy(self) -> Proxy:
         if not isinstance(self.proxy, torch.fx.Proxy):
@@ -234,12 +287,24 @@ def as_proxy(self) -> Proxy:
                         "NYI: hoisted opaque objects that accept kwargs, please pass as args"
                     )
                 hoisted_vt = tx.output.synthetic_graph_input(
-                    type(self.proxy), self.ctor_args_kwargs[0]
+                    type(self.proxy),
+                    self.ctor_args_kwargs[0],
+                    ctor_arg_sources=self.ctor_arg_sources,
                 )
                 self.proxy = hoisted_vt.as_proxy()
 
         return self.proxy
 
+    def __str__(self) -> str:
+        value = (
+            self.value.real_obj
+            if isinstance(self.value, FakeScriptObject)
+            else self.value
+        )
+        return f"{self.__class__.__name__}({value})"
+
+    __repr__ = __str__
+
     @_raise_hard_error_if_graph_break(
         "Dynamo cannot safely trace script object due to graph break."
     )
@@ -248,28 +313,16 @@ def var_getattr(self, tx: "InstructionTranslator", name: str) -> VariableTracker
 
         from .higher_order_ops import TorchHigherOrderOperatorVariable
 
-        if hasattr(self.value, "script_class_name") and is_opaque_type(
-            self.value.script_class_name
-        ):
-            real_obj = self.value.real_obj  # pyrefly: ignore[missing-attribute]
-
-            member_type = get_member_type(
-                type(real_obj),
-                name,
-            )
-            if member_type is None:
-                unimplemented(
-                    gb_type="Attempted to access unregistered member on an OpaqueObject",
-                    context=f"value={real_obj}, attr={name}",
-                    explanation=f"Member '{name}' is not registered for this opaque object type.",
-                    hints=[
-                        f"Register '{name}' with a MemberType in register_opaque_type(members=...).",
-                    ],
-                )
+        real_obj = self.as_python_constant()
+        real_obj_type = type(real_obj)
+        if is_opaque_type(real_obj_type):
+            member_type = get_member_type(real_obj_type, name)
 
             if member_type == MemberType.USE_REAL:
                 value = getattr(real_obj, name)
-                if inspect.ismethod(value):
+                if inspect.ismethod(value) or isinstance(
+                    value, types.MethodWrapperType
+                ):
                     return LambdaVariable(
                         lambda *args, **kwargs: self.call_method(tx, name, args, kwargs)
                     )
@@ -278,7 +331,10 @@ def var_getattr(self, tx: "InstructionTranslator", name: str) -> VariableTracker
 
             elif member_type == MemberType.INLINED:
                 value = getattr(real_obj, name)
-                if inspect.ismethod(value) and self.source is None:
+                if (
+                    inspect.ismethod(value)
+                    or isinstance(value, types.MethodWrapperType)
+                ) and self.source is None:
                     # When we don't have a source, fall back to call_method
                     # which creates a proxy node.
                     return LambdaVariable(
@@ -286,6 +342,26 @@ def var_getattr(self, tx: "InstructionTranslator", name: str) -> VariableTracker
                     )
                 return super().var_getattr(tx, name)
 
+            elif is_opaque_value_type(real_obj_type):
+                return super().var_getattr(tx, name)
+
+            elif name in ("__bool__", "__len__") and not hasattr(real_obj, name):
+                # Special case: __bool__ and __len__ are used for truthiness checks.
+                # If they're not registered and the real object doesn't have them,
+                # raise ObservedAttributeError so the caller can fall back to
+                # treating the object as truthy (Python default behavior
+                raise_observed_exception(AttributeError, tx)
+
+            else:
+                unimplemented(
+                    gb_type="Attempted to access unregistered member on an OpaqueObject",
+                    context=f"value={real_obj}, attr={name}",
+                    explanation=f"Member '{name}' is not registered for this opaque object type.",
+                    hints=[
+                        f"Register '{name}' with a MemberType in register_opaque_type(members=...).",
+                    ],
+                )
+
         method = getattr(self.value, name, None)
         if method is None:
             unimplemented(
@@ -316,9 +392,18 @@ def var_getattr(self, tx: "InstructionTranslator", name: str) -> VariableTracker
             method_name=name,
         )
 
+    def mp_subscript_impl(
+        self,
+        tx: "InstructionTranslator",
+        key: "VariableTracker",
+    ) -> "VariableTracker":
+        # Call call_method directly on this class to avoid the __getitem__ →
+        # mp_subscript_impl loop in VariableTracker.call_method.
+        return TorchScriptObjectVariable.call_method(self, tx, "__getitem__", [key], {})
+
     # We only support method calls on script objects. Interpreting the bytecodes
     # should go through var_getattr then call_function instead of call_method.
-    #
+
     # However, it's possible for call_method to be used directly e.g. for __setattr__.
     @_raise_hard_error_if_graph_break(
         "Dynamo cannot safely trace script object due to graph break."
@@ -332,43 +417,19 @@ def call_method(
     ) -> VariableTracker:
         from .builder import wrap_fx_proxy
 
-        if hasattr(self.value, "script_class_name") and is_opaque_type(
-            self.value.script_class_name
-        ):
-            real_obj = self.value.real_obj  # pyrefly: ignore[missing-attribute]
-            value_type = type(real_obj)
-
-            member_type = get_member_type(
-                value_type,
-                name,
-            )
-            if member_type is None:
-                unimplemented(
-                    gb_type="Attempted to access unregistered member on an OpaqueObject",
-                    context=f"value={real_obj}, attr={name}",
-                    explanation=f"Member '{name}' is not registered for this opaque object type.",
-                    hints=[
-                        f"Register '{name}' with a MemberType in register_opaque_type(members=...).",
-                    ],
-                )
-
-            if member_type == MemberType.INLINED:
-                proxy_args, proxy_kwargs = proxy_args_kwargs(args, kwargs)
-
-                proxy = tx.output.create_proxy(
-                    "call_method",
-                    name,
-                    args=(self.proxy, *proxy_args),
-                    kwargs=proxy_kwargs,
-                )
-
-                return wrap_fx_proxy(tx=tx, proxy=proxy)
+        real_obj = self.as_python_constant()
+        real_obj_type = type(real_obj)
+        if is_opaque_type(real_obj_type):
+            member_type = get_member_type(real_obj_type, name)
 
-            elif member_type == MemberType.USE_REAL:
-                if inspect.getattr_static(value_type, "__getattr__", None) is not None:
+            if member_type == MemberType.USE_REAL:
+                if (
+                    inspect.getattr_static(real_obj_type, "__getattr__", None)
+                    is not None
+                ):
                     unimplemented(
                         gb_type="Opaque object with custom __getattr__ not supported",
-                        context=f"{value_type.__name__} with custom __getattr__",
+                        context=f"{real_obj_type.__name__} with custom __getattr__",
                         explanation="Dynamo does not support opaque objects types with custom __getattr__ methods",
                         hints=[],
                     )
@@ -390,7 +451,7 @@ def call_method(
                 ):
                     unimplemented(
                         gb_type="Opaque object member with method-type USE_REAL returned a reference-type opaque object.",
-                        context=f"Opaque object type: {value_type}. Method name: '{name}'",
+                        context=f"Opaque object type: {real_obj_type}. Method name: '{name}'",
                         explanation=(
                             "To properly guard reference-type opaque objects, "
                             "we must lift them as inputs to the graph. In order "
@@ -399,18 +460,34 @@ def call_method(
                         ),
                         hints=[
                             f"Register member '{name}' with MemberType.INLINED in "
-                            "register_opaque_type({value_type}, members=...).",
+                            f"register_opaque_type({real_obj_type}, members=...).",
                         ],
                     )
 
                 return VariableTracker.build(tx, constant_val)
 
+            elif member_type == MemberType.INLINED or is_opaque_value_type(
+                real_obj_type
+            ):
+                proxy_args, proxy_kwargs = proxy_args_kwargs(args, kwargs)
+
+                proxy = tx.output.create_proxy(
+                    "call_method",
+                    name,
+                    args=(self.proxy, *proxy_args),
+                    kwargs=proxy_kwargs,
+                )
+
+                return wrap_fx_proxy(tx=tx, proxy=proxy)
+
             else:
                 unimplemented(
-                    gb_type="Unsupported member type on OpaqueObject",
-                    context=f"value={real_obj}, attr={name}, member_type={member_type}",
-                    explanation=f"Member type '{member_type}' is not supported for this operation.",
-                    hints=[],
+                    gb_type="Attempted to access unregistered member on an OpaqueObject",
+                    context=f"value={real_obj}, attr={name}",
+                    explanation=f"Member '{name}' is not registered for this opaque object type.",
+                    hints=[
+                        f"Register '{name}' with a MemberType in register_opaque_type(members=...).",
+                    ],
                 )
 
         unimplemented(
@@ -426,13 +503,30 @@ def call_method(
         )
 
     def as_python_constant(self) -> Any:
-        if is_opaque_value_type(
-            type(self.value.real_obj)  # pyrefly: ignore[missing-attribute]
-        ):
-            return self.value.real_obj  # pyrefly: ignore[missing-attribute]
+        if isinstance(self.value, FakeScriptObject):
+            return self.value.real_obj
+        elif is_opaque_value_type(type(self.value)):
+            return self.value
+        elif isinstance(self.value, torch.ScriptObject):
+            return self.value
         return super().as_python_constant()
 
     def is_python_hashable(self) -> bool:
-        return is_opaque_value_type(
-            type(self.value.real_obj)  # pyrefly: ignore[missing-attribute]
-        )
+        try:
+            self.get_python_hash()
+            return True
+        except TypeError:
+            return False
+
+    def get_python_hash(self) -> int:
+        real_obj = self.as_python_constant()
+        return hash(real_obj)
+
+    def is_python_equal(self, other: object) -> bool:
+        assert isinstance(other, VariableTracker)
+        real_self = self.as_python_constant()
+        real_other = other.as_python_constant()
+        return real_self == real_other
+
+    def get_real_value(self) -> Any:
+        return self.as_python_constant()
diff --git a/torch/_dynamo/variables/sdpa.py b/torch/_dynamo/variables/sdpa.py
index 1a7006f5d56ab..48de0f9ab4e3c 100644
--- a/torch/_dynamo/variables/sdpa.py
+++ b/torch/_dynamo/variables/sdpa.py
@@ -50,6 +50,9 @@ def __init__(
         self.param_vars = param_vars
         super().__init__(**kwargs)
 
+    def python_type(self) -> type:
+        return SDPAParams
+
     def reconstruct(self, codegen: "PyCodegen") -> None:
         assert self.source is None
         assert self.param_vars is not None
diff --git a/torch/_dynamo/variables/sets.py b/torch/_dynamo/variables/sets.py
new file mode 100644
index 0000000000000..acbcd62333e32
--- /dev/null
+++ b/torch/_dynamo/variables/sets.py
@@ -0,0 +1,797 @@
+"""
+Set-related variable tracking classes for PyTorch Dynamo.
+
+This module implements variable tracking for different types of set-like objects:
+- Regular Python sets (set)
+- Frozen sets (frozenset)
+- Ordered sets (torch.utils._ordered_set.OrderedSet)
+- Dictionary key sets (dict_keys views used as sets)
+
+These classes are responsible for tracking set operations during graph compilation,
+maintaining proper guards for set mutations and element existence checks.
+
+The implementation uses a special HashableTracker wrapper to handle set elements
+while preserving proper aliasing semantics. Sets are modeled internally as
+dictionaries with None values.
+"""
+
+import functools
+import operator
+from collections.abc import Iterable, Sequence
+from typing import Any, Literal, TYPE_CHECKING
+
+from torch.utils._ordered_set import OrderedSet
+
+from .. import polyfills, variables
+from ..bytecode_transformation import create_call_function, create_instruction
+from ..exc import raise_observed_exception
+from ..guards import GuardBuilder, install_guard
+from ..source import AttrSource, is_constant_source, is_from_local_source
+from ..utils import cmp_name_to_op_mapping, istype, raise_args_mismatch
+from .base import ValueMutationNew, VariableTracker
+from .constant import ConstantVariable
+from .hashable import HashableTracker, is_hashable, raise_unhashable
+
+
+if TYPE_CHECKING:
+    from torch._dynamo.codegen import PyCodegen
+    from torch._dynamo.symbolic_convert import InstructionTranslator
+    from torch._dynamo.variables.builtin import BuiltinVariable
+
+
+# [Adding a new supported class within the keys of SetVariable]
+# see steps outlined for ConstDictVariable
+
+
+class SetVariable(VariableTracker):
+    """Represents a Python set during symbolic execution."""
+
+    # PySet_Type: https://github.com/python/cpython/blob/v3.13.0/Objects/setobject.c#L2436
+    _cpython_type = set
+
+    CONTAINS_GUARD = GuardBuilder.SET_CONTAINS
+    NOT_CONTAINS_GUARD = GuardBuilder.SET_NOT_CONTAINS
+
+    def __init__(
+        self,
+        items: Iterable[VariableTracker | HashableTracker],
+        **kwargs: Any,
+    ) -> None:
+        # .clone() passes these arguments in kwargs but they're recreated below
+        if "original_items" in kwargs:
+            kwargs.pop("original_items")
+        if "should_reconstruct_all" in kwargs:
+            kwargs.pop("should_reconstruct_all")
+
+        super().__init__(**kwargs)
+
+        # Items can be either VariableTrackers or HashableTrackers (from set ops).
+        # For VariableTrackers, realize them to ensure aliasing guards are installed
+        # when the same object appears multiple times.
+        hashable_items = []
+        for item in items:
+            if isinstance(item, HashableTracker):
+                # Already a HashableTracker from a set operation
+                hashable_items.append(item)
+            else:
+                # VariableTracker - realize to install guards, then wrap
+                # pyrefly: ignore [bad-argument-type]
+                hashable_items.append(HashableTracker(item.realize()))
+        self.items = dict.fromkeys(hashable_items, SetVariable._default_value())
+        self.should_reconstruct_all = (
+            not is_from_local_source(self.source) if self.source else True
+        )
+        self.original_items = dict.fromkeys(
+            hashable_items, SetVariable._default_value()
+        )
+
+    def debug_repr(self) -> str:
+        if not self.items:
+            return "set()"
+        else:
+            items: list[str] = []
+            for v in self.items:
+                vt = v.vt if isinstance(v, HashableTracker) else v
+                val_str = repr(vt.value) if hasattr(vt, "value") else vt.debug_repr()
+                items.append(val_str)
+            return "{" + ",".join(items) + "}"
+
+    @property
+    def set_items(self) -> set["HashableTracker"]:
+        return set(self.items.keys())
+
+    @staticmethod
+    def _default_value() -> VariableTracker:
+        # Variable to fill in the keys of the dictionary
+        return ConstantVariable.create(None)
+
+    def as_proxy(self) -> Any:
+        return {k.vt.as_proxy() for k in self.set_items}
+
+    def python_type(self) -> type:
+        return set
+
+    def as_python_constant(self) -> Any:
+        return {k.vt.as_python_constant() for k in self.set_items}
+
+    def reconstruct(self, codegen: "PyCodegen") -> None:
+        codegen.foreach([x.vt for x in self.set_items])
+        codegen.append_output(create_instruction("BUILD_SET", arg=len(self.set_items)))
+
+    def __contains__(self, vt: VariableTracker) -> bool:
+        assert isinstance(vt, VariableTracker)
+        if not is_hashable(vt):
+            return False
+        key = HashableTracker(vt)
+        return key in self.items and not isinstance(
+            self.items[key], variables.DeletedVariable
+        )
+
+    def len(self) -> int:
+        return sum(
+            not isinstance(x, variables.DeletedVariable) for x in self.items.values()
+        )
+
+    def has_new_items(self) -> bool:
+        return self.should_reconstruct_all or any(
+            self.is_new_item(self.original_items.get(key.vt), value)
+            for key, value in self.items.items()
+        )
+
+    def is_new_item(
+        self, value: VariableTracker | None, other: VariableTracker
+    ) -> bool:
+        if value and value.is_realized() and other.is_realized():
+            return id(value.realize()) != id(other.realize())
+        return id(value) != id(other)
+
+    def unpack_var_sequence(self, tx: "InstructionTranslator") -> list[VariableTracker]:
+        return [x.vt for x in self.items]
+
+    def clone(self, **kwargs: Any) -> VariableTracker:
+        return super().clone(**kwargs)
+
+    def is_python_hashable(self) -> bool:
+        return False
+
+    def var_getattr(self, tx: "InstructionTranslator", name: str):
+        if name == "__class__":
+            return VariableTracker.build(tx, self.python_type())
+        return super().var_getattr(tx, name)
+
+    def call_obj_hasattr(
+        self, tx: "InstructionTranslator", name: str
+    ) -> ConstantVariable:
+        return VariableTracker.build(tx, hasattr(set, name))
+
+    def install_set_contains_guard(
+        self, tx: "InstructionTranslator", args: list[VariableTracker]
+    ) -> None:
+        if not self.source:
+            return
+
+        if tx.output.side_effects.is_modified(self):
+            return
+
+        contains = args[0] in self
+        if args[0].source is None and args[0].is_python_constant():
+            guard_fn = (
+                type(self).CONTAINS_GUARD if contains else type(self).NOT_CONTAINS_GUARD
+            )
+            install_guard(
+                self.make_guard(
+                    functools.partial(
+                        guard_fn,
+                        key=args[0].as_python_constant(),
+                    )
+                )
+            )
+
+    def _fast_set_method(
+        self,
+        tx: "InstructionTranslator",
+        fn: Any,
+        args: list[VariableTracker],
+        kwargs: dict[str, VariableTracker],
+    ) -> VariableTracker:
+        try:
+            res = fn(
+                *[x.as_python_constant() for x in [self, *args]],
+                **{k: v.as_python_constant() for k, v in kwargs.items()},
+            )
+        except Exception as exc:
+            raise_observed_exception(type(exc), tx, args=list(exc.args))
+        return VariableTracker.build(tx, res)
+
+    def call_method(
+        self,
+        tx: "InstructionTranslator",
+        name: str,
+        args: list[VariableTracker],
+        kwargs: dict[str, VariableTracker],
+    ) -> VariableTracker:
+        from ..utils import check_constant_args
+        from .builder import SourcelessBuilder
+
+        if (
+            name
+            in (
+                "isdisjoint",
+                "union",
+                "intersection",
+                "difference",
+                "symmetric_difference",
+            )
+            and check_constant_args(args, kwargs)
+            and self.python_type() is set
+        ):
+            py_type = self.python_type()
+            return self._fast_set_method(tx, getattr(py_type, name), args, kwargs)
+
+        # Lazy imports to avoid circular dependencies
+        from .dicts import DictItemsVariable, DictKeysVariable
+
+        if name == "__init__":
+            temp_set_vt = SourcelessBuilder.create(tx, set).call_set(
+                tx, *args, **kwargs
+            )
+            tx.output.side_effects.mutation(self)
+            self.items.clear()
+            self.items.update(temp_set_vt.items)  # type: ignore[attr-defined]
+            return ConstantVariable.create(None)
+        elif name == "add":
+            if kwargs or len(args) != 1:
+                raise_args_mismatch(
+                    tx,
+                    name,
+                    "1 args and 0 kwargs",
+                    f"{len(args)} args and {len(kwargs)} kwargs",
+                )
+            # Convert add to __setitem__ with None value
+            if not is_hashable(args[0]):
+                raise_unhashable(args[0], tx)
+            tx.output.side_effects.mutation(self)
+            self.items[HashableTracker(args[0])] = SetVariable._default_value()
+            return ConstantVariable.create(None)
+        elif name == "pop":
+            if kwargs or args:
+                raise_args_mismatch(
+                    tx,
+                    name,
+                    "0 args and 0 kwargs",
+                    f"{len(args)} args and {len(kwargs)} kwargs",
+                )
+            # Choose an item at random and pop it
+            try:
+                result: VariableTracker = self.set_items.pop().vt  # type: ignore[assignment]
+            except KeyError as e:
+                raise_observed_exception(KeyError, tx, args=list(e.args))
+            self.should_reconstruct_all = True
+            tx.output.side_effects.mutation(self)
+            self.items.pop(HashableTracker(result))
+            return result
+        elif name == "isdisjoint":
+            if kwargs or len(args) != 1:
+                raise_args_mismatch(
+                    tx,
+                    name,
+                    "1 args and 0 kwargs",
+                    f"{len(args)} args and {len(kwargs)} kwargs",
+                )
+            return SourcelessBuilder.create(tx, polyfills.set_isdisjoint).call_function(
+                tx, [self, args[0]], {}
+            )
+        elif name == "intersection":
+            if kwargs:
+                raise_args_mismatch(tx, name, "0 kwargs", f"{len(kwargs)} kwargs")
+            return SourcelessBuilder.create(
+                tx, polyfills.set_intersection
+            ).call_function(
+                tx,
+                [self, *args],
+                {"cls": self.python_type_var()},
+            )
+        elif name == "intersection_update":
+            if kwargs:
+                raise_args_mismatch(tx, name, "0 kwargs", f"{len(kwargs)} kwargs")
+            return SourcelessBuilder.create(
+                tx, polyfills.set_intersection_update
+            ).call_function(tx, [self, *args], {})
+        elif name == "union":
+            if kwargs:
+                raise_args_mismatch(tx, name, "0 kwargs", f"{len(kwargs)} kwargs")
+            return SourcelessBuilder.create(tx, polyfills.set_union).call_function(
+                tx,
+                [self, *args],
+                {"cls": self.python_type_var()},
+            )
+        elif name == "difference":
+            if kwargs:
+                raise_args_mismatch(
+                    tx, name, f"Expect: 0 kwargs, Actual: {len(kwargs)} kwargs"
+                )
+            return SourcelessBuilder.create(tx, polyfills.set_difference).call_function(
+                tx,
+                [self, *args],
+                {"cls": self.python_type_var()},
+            )
+        elif name == "difference_update":
+            if kwargs:
+                raise_args_mismatch(tx, name, "0 kwargs", f"{len(kwargs)} kwargs")
+            return SourcelessBuilder.create(
+                tx, polyfills.set_difference_update
+            ).call_function(tx, [self, *args], {})
+        elif name == "symmetric_difference":
+            if kwargs or len(args) != 1:
+                raise_args_mismatch(
+                    tx,
+                    name,
+                    "1 args and 0 kwargs",
+                    f"{len(args)} args and {len(kwargs)} kwargs",
+                )
+            return SourcelessBuilder.create(
+                tx, polyfills.set_symmetric_difference
+            ).call_function(
+                tx,
+                [self, *args],
+                {"cls": self.python_type_var()},
+            )
+        elif name == "symmetric_difference_update":
+            if kwargs or len(args) != 1:
+                raise_args_mismatch(
+                    tx,
+                    name,
+                    "1 args and 0 kwargs",
+                    f"{len(args)} args and {len(kwargs)} kwargs",
+                )
+            return SourcelessBuilder.create(
+                tx, polyfills.set_symmetric_difference_update
+            ).call_function(tx, [self, *args], {})
+        elif name == "update" and self.is_mutable():
+            if kwargs:
+                raise_args_mismatch(tx, name, "0 kwargs", f"{len(kwargs)} kwargs")
+            return SourcelessBuilder.create(tx, polyfills.set_update).call_function(
+                tx, [self, *args], {}
+            )
+        elif name == "remove":
+            if kwargs or len(args) != 1:
+                raise_args_mismatch(
+                    tx,
+                    name,
+                    "1 args and 0 kwargs",
+                    f"{len(args)} args and {len(kwargs)} kwargs",
+                )
+            if args[0] not in self:
+                raise_observed_exception(KeyError, tx, args=args)
+            self.should_reconstruct_all = True
+            tx.output.side_effects.mutation(self)
+            self.items.pop(HashableTracker(args[0]))
+            return ConstantVariable.create(None)
+        elif name == "discard":
+            if kwargs or len(args) != 1:
+                raise_args_mismatch(
+                    tx,
+                    name,
+                    "1 args and 0 kwargs",
+                    f"{len(args)} args and {len(kwargs)} kwargs",
+                )
+            if args[0] in self:
+                self.should_reconstruct_all = True
+                tx.output.side_effects.mutation(self)
+                self.items.pop(HashableTracker(args[0]))
+            return ConstantVariable.create(None)
+        elif name in ("issubset", "issuperset"):
+            if len(args) != 1:
+                raise_args_mismatch(tx, name, "1 args", f"{len(args)} args")
+
+            op = {
+                "issubset": operator.le,
+                "issuperset": operator.ge,
+            }
+            other = args[0].realize()
+            if not istype(other, SetVariable):
+                other = SourcelessBuilder.create(tx, set).call_function(tx, [other], {})
+            return SourcelessBuilder.create(tx, op.get(name)).call_function(
+                tx, [self, other], {}
+            )
+        elif name in ("__and__", "__or__", "__xor__", "__sub__"):
+            m = {
+                "__and__": "intersection",
+                "__or__": "union",
+                "__xor__": "symmetric_difference",
+                "__sub__": "difference",
+            }.get(name)
+            if not isinstance(
+                args[0],
+                (
+                    SetVariable,
+                    variables.UserDefinedSetVariable,
+                    DictItemsVariable,
+                    DictKeysVariable,
+                ),
+            ):
+                raise_observed_exception(
+                    TypeError,
+                    tx,
+                    args=[
+                        f"unsupported operand type(s) for {name}: '{self.python_type_name()}' and '{args[0].python_type_name()}'"
+                    ],
+                )
+            assert m is not None
+            return self.call_method(tx, m, args, kwargs)
+        elif name in ("__rand__", "__ror__", "__rxor__", "__rsub__"):
+            m = {
+                "__rand__": "__and__",
+                "__ror__": "__or__",
+                "__rxor__": "__xor__",
+                "__rsub__": "__sub__",
+            }.get(name)
+            if not isinstance(
+                args[0],
+                (
+                    SetVariable,
+                    variables.UserDefinedSetVariable,
+                    DictItemsVariable,
+                    DictKeysVariable,
+                ),
+            ):
+                raise_observed_exception(
+                    TypeError,
+                    tx,
+                    args=[
+                        f"unsupported operand type(s) for {name}: '{args[0].python_type_name()}' and '{self.python_type_name()}'"
+                    ],
+                )
+            assert m is not None
+            return args[0].call_method(tx, m, [self], kwargs)
+        elif name in ("__iand__", "__ior__", "__ixor__", "__isub__"):
+            if not isinstance(
+                args[0],
+                (
+                    SetVariable,
+                    variables.UserDefinedSetVariable,
+                    DictItemsVariable,
+                    DictKeysVariable,
+                ),
+            ):
+                raise_observed_exception(
+                    TypeError,
+                    tx,
+                    args=[
+                        f"unsupported operand type(s) for {name}: '{self.python_type_name()}' and '{args[0].python_type_name()}'"
+                    ],
+                )
+            m = {
+                "__iand__": "intersection_update",
+                "__ior__": "update",
+                "__ixor__": "symmetric_difference_update",
+                "__isub__": "difference_update",
+            }.get(name)
+            assert m is not None
+            self.call_method(tx, m, args, kwargs)
+            return self
+        elif name == "__eq__":
+            if not isinstance(
+                args[0],
+                (
+                    SetVariable,
+                    variables.UserDefinedSetVariable,
+                    DictItemsVariable,
+                    DictKeysVariable,
+                ),
+            ):
+                return ConstantVariable.create(False)
+            r = self.call_method(tx, "symmetric_difference", args, kwargs)
+            return VariableTracker.build(tx, len(r.set_items) == 0)  # type: ignore[attr-defined]
+        elif name == "__ne__":
+            eq_result = self.call_method(tx, "__eq__", args, kwargs)
+            return VariableTracker.build(tx, not eq_result.value)  # type: ignore[attr-defined]
+        elif name in cmp_name_to_op_mapping:
+            if not isinstance(
+                args[0],
+                (
+                    SetVariable,
+                    variables.UserDefinedSetVariable,
+                    DictItemsVariable,
+                    DictKeysVariable,
+                ),
+            ):
+                return VariableTracker.build(tx, NotImplemented)
+            return VariableTracker.build(
+                tx,
+                cmp_name_to_op_mapping[name](self.set_items, args[0].set_items),  # type: ignore[attr-defined]
+            )
+        elif name == "__contains__":
+            if not len(args):
+                raise_args_mismatch(
+                    tx,
+                    name,
+                    "more than 1 args and 0 kwargs",
+                    f"{len(args)} args and {len(kwargs)} kwargs",
+                )
+            if not (args and is_hashable(args[0])):
+                raise_unhashable(args[0], tx)
+            self.install_set_contains_guard(tx, args)
+            contains = args[0] in self
+            return VariableTracker.build(tx, contains)
+        elif name == "__len__":
+            if args or kwargs:
+                raise_args_mismatch(
+                    tx,
+                    name,
+                    "0 args and 0 kwargs",
+                    f"{len(args)} args and {len(kwargs)} kwargs",
+                )
+            return VariableTracker.build(tx, len(self.items))
+        elif name == "copy":
+            if args or kwargs:
+                raise_args_mismatch(
+                    tx,
+                    name,
+                    "0 args and 0 kwargs",
+                    f"{len(args)} args and {len(kwargs)} kwargs",
+                )
+            return self.clone(
+                items=self.items.copy(), mutation_type=ValueMutationNew(), source=None
+            )
+        elif name == "clear":
+            if args or kwargs:
+                raise_args_mismatch(
+                    tx,
+                    name,
+                    "0 args and 0 kwargs",
+                    f"{len(args)} args and {len(kwargs)} kwargs",
+                )
+            self.should_reconstruct_all = True
+            tx.output.side_effects.mutation(self)
+            self.items.clear()
+            return ConstantVariable.create(None)
+        elif name == "__iter__":
+            from .lists import ListIteratorVariable
+
+            if self.source and not is_constant_source(self.source):
+                tx.output.guard_on_key_order.add(self.source)
+            return ListIteratorVariable(
+                self.unpack_var_sequence(tx), mutation_type=ValueMutationNew()
+            )
+        return super().call_method(tx, name, args, kwargs)
+
+    def python_type_var(self) -> "BuiltinVariable":
+        return variables.BuiltinVariable(set)
+
+    def getitem_const(
+        self, tx: "InstructionTranslator", arg: VariableTracker
+    ) -> VariableTracker:
+        raise RuntimeError("Illegal to getitem on a set")
+
+    def sq_length(self, tx: "InstructionTranslator") -> VariableTracker:
+        return VariableTracker.build(tx, len(self.set_items))
+
+
+class OrderedSetClassVariable(VariableTracker):
+    def __init__(self, **kwargs: Any) -> None:
+        super().__init__(**kwargs)
+
+    def as_python_constant(self) -> type[OrderedSet[Any]]:
+        return OrderedSet
+
+    def var_getattr(self, tx: "InstructionTranslator", name: str) -> VariableTracker:
+        if name == "__new__":
+            from .misc import GetAttrVariable
+
+            if self.source:
+                attr_source = AttrSource(self.source, name)
+            else:
+                attr_source = None
+            return GetAttrVariable(
+                self, name, py_type=type(getattr(OrderedSet, name)), source=attr_source
+            )
+        else:
+            return super().var_getattr(tx, name)
+
+    def call_method(
+        self,
+        tx: "InstructionTranslator",
+        name: str,
+        args: list[VariableTracker],
+        kwargs: dict[str, VariableTracker],
+    ) -> VariableTracker:
+        from .builtin import set_methods
+
+        if name == "__new__":
+            if len(args) != 2 or kwargs:
+                raise_args_mismatch(
+                    tx,
+                    name,
+                    "OrderedSet.__new__ only accepts one arg"
+                    f"{len(args)} args and {len(kwargs)} kwargs",
+                )
+
+            return variables.OrderedSetVariable([], mutation_type=ValueMutationNew())
+
+        resolved_fn = getattr(set, name)
+        if resolved_fn in set_methods and isinstance(args[0], variables.SetVariable):
+            return args[0].call_method(tx, name, args[1:], kwargs)
+
+        return super().call_method(tx, name, args, kwargs)
+
+    def call_function(
+        self,
+        tx: "InstructionTranslator",
+        args: Sequence[VariableTracker],
+        kwargs: dict[str, VariableTracker],
+    ) -> "OrderedSetVariable":
+        if len(args) > 1 or kwargs:
+            raise_args_mismatch(
+                tx,
+                "OrderedSet",
+                "OrderedSet only accepts one arg"
+                f"{len(args)} args and {len(kwargs)} kwargs",
+            )
+
+        if len(args) == 0:
+            # pyrefly: ignore [implicit-any]
+            items = []
+        else:
+            items = args[0].force_unpack_var_sequence(tx)
+        return variables.OrderedSetVariable(items, mutation_type=ValueMutationNew())
+
+
+class OrderedSetVariable(SetVariable):
+    def debug_repr(self) -> str:
+        if not self.items:
+            return "OrderedSet([])"
+        else:
+            items: list[str] = []
+            for k in self.items:
+                key_str = (
+                    repr(k.vt.value) if hasattr(k.vt, "value") else k.vt.debug_repr()
+                )
+                items.append(key_str)
+            return "OrderedSet([" + ",".join(items) + "])"
+
+    def as_python_constant(self) -> OrderedSet[Any]:
+        return OrderedSet([k.vt.as_python_constant() for k in self.set_items])
+
+    def python_type(self) -> type[OrderedSet[Any]]:
+        return OrderedSet
+
+    # pyrefly: ignore[bad-override]
+    def python_type_var(self) -> OrderedSetClassVariable:
+        return OrderedSetClassVariable()
+
+    def reconstruct(self, codegen: "PyCodegen") -> None:
+        codegen.add_push_null(
+            lambda: codegen.load_import_from("torch.utils._ordered_set", "OrderedSet")
+        )
+        codegen.foreach([x.vt for x in self.set_items])
+        codegen.append_output(create_instruction("BUILD_LIST", arg=len(self.set_items)))
+        codegen.extend_output(create_call_function(1, False))
+
+
+class FrozensetVariable(SetVariable):
+    # PyFrozenSet_Type: https://github.com/python/cpython/blob/v3.13.0/Objects/setobject.c#L2526
+    _cpython_type = frozenset
+
+    def debug_repr(self) -> str:
+        if not self.items:
+            return "frozenset()"
+        else:
+            items: list[str] = []
+            for k in self.items:
+                key_str = (
+                    repr(k.vt.value) if hasattr(k.vt, "value") else k.vt.debug_repr()
+                )
+                items.append(key_str)
+            return "{" + ",".join(items) + "}"
+
+    @property
+    def set_items(self) -> set["HashableTracker"]:
+        return set(self.items.keys())
+
+    def python_type(self) -> type:
+        return frozenset
+
+    def python_type_var(self) -> "BuiltinVariable":
+        return variables.BuiltinVariable(frozenset)
+
+    def as_python_constant(self) -> Any:
+        return frozenset({k.vt.as_python_constant() for k in self.set_items})
+
+    def reconstruct(self, codegen: "PyCodegen") -> None:
+        codegen.add_push_null(
+            lambda: codegen.extend_output(
+                [
+                    codegen.create_load_global("frozenset"),
+                ]
+            )
+        )
+        codegen.foreach([x.vt for x in self.set_items])
+        codegen.extend_output(
+            [
+                create_instruction("BUILD_LIST", arg=len(self.set_items)),
+                *create_call_function(1, False),
+            ]
+        )
+
+    def call_method(
+        self,
+        tx: "InstructionTranslator",
+        name: str,
+        args: list[VariableTracker],
+        kwargs: dict[str, VariableTracker],
+    ) -> VariableTracker:
+        if name in ["add", "pop", "update", "remove", "discard", "clear"]:
+            raise RuntimeError(f"Illegal call_method {name} on a frozenset")
+        elif name == "__init__":
+            # frozenset is immutable. Calling __init__ again shouldn't have any effect
+            return ConstantVariable.create(None)
+        elif name in (
+            "copy",
+            "difference",
+            "intersection",
+            "symmetric_difference",
+        ):
+            r = super().call_method(tx, name, args, kwargs)
+            return FrozensetVariable(r.items)  # type: ignore[attr-defined]
+        return super().call_method(tx, name, args, kwargs)
+
+    def is_python_hashable(self) -> Literal[True]:
+        """
+        Frozensets are immutable and hashable in Python.
+        """
+        return True
+
+    def get_python_hash(self) -> int:
+        return hash(self.as_python_constant())
+
+    def is_python_equal(self, other: object) -> bool:
+        return (
+            isinstance(other, VariableTracker)
+            and self.as_python_constant() == other.as_python_constant()
+        )
+
+
+class DictKeySetVariable(SetVariable):
+    def debug_repr(self) -> str:
+        if not self.items:
+            return "dict_keys([])"
+        else:
+            items: list[str] = []
+            for k in self.items:
+                key_str = (
+                    repr(k.vt.value) if hasattr(k.vt, "value") else k.vt.debug_repr()
+                )
+                items.append(key_str)
+            return "dict_keys([" + ",".join(items) + "])"
+
+    def install_set_contains_guard(
+        self, tx: "InstructionTranslator", args: list[VariableTracker]
+    ) -> None:
+        # Already EQUALS_MATCH guarded
+        pass
+
+    @property
+    def set_items(self) -> Any:
+        return self.items
+
+    def python_type(self) -> type:
+        from ..utils import dict_keys
+
+        return dict_keys
+
+    def as_python_constant(self) -> Any:
+        return dict.fromkeys(
+            {k.vt.as_python_constant() for k in self.set_items}, None
+        ).keys()
+
+    def call_method(
+        self,
+        tx: "InstructionTranslator",
+        name: str,
+        args: list[VariableTracker],
+        kwargs: dict[str, VariableTracker],
+    ) -> VariableTracker:
+        if name in ["add", "pop", "update", "remove", "discard", "clear"]:
+            raise RuntimeError(f"Illegal call_method {name} on a dict_keys")
+        return super().call_method(tx, name, args, kwargs)
diff --git a/torch/_dynamo/variables/streams.py b/torch/_dynamo/variables/streams.py
index 4fe62978ffa6d..c42e6a84d1aaa 100644
--- a/torch/_dynamo/variables/streams.py
+++ b/torch/_dynamo/variables/streams.py
@@ -11,12 +11,15 @@
 from ..bytecode_transformation import create_call_function
 from ..exc import TYPE_CHECKING, unimplemented
 from ..graph_bytecode_inputs import (
+    CURRENT_STREAM_INDEX,
     get_external_object_by_index,
     register_graph_created_object,
+    register_user_object,
+    reset_user_object_tracking,
 )
 from ..source import CurrentStreamSource
 from .base import VariableTracker
-from .constant import CONSTANT_VARIABLE_NONE, ConstantVariable
+from .constant import ConstantVariable
 from .ctx_manager import FxTracebackAnnotateVariable
 from .lazy import LazyVariableTracker
 
@@ -125,7 +128,7 @@ def _(
 def record_event(event_index: int, stream_index: int) -> None:
     event = _get_event_by_index(event_index)
     stream = _get_stream_by_index(stream_index)
-    stream.record_event(event)
+    event.record(stream)
 
 
 @record_event.register_fake
@@ -143,7 +146,7 @@ def _(
 def wait_event(event_index: int, stream_index: int) -> None:
     event = _get_event_by_index(event_index)
     stream = _get_stream_by_index(stream_index)
-    stream.wait_event(event)
+    event.wait(stream)
 
 
 @wait_event.register_fake
@@ -157,6 +160,47 @@ def _(
 has_side_effect(torch.ops.streams.wait_event.default)
 
 
+@custom_op("streams::synchronize_event", mutates_args=())
+def synchronize_event(event_index: int) -> None:
+    event = _get_event_by_index(event_index)
+    event.synchronize()
+
+
+@synchronize_event.register_fake
+def _(event_index: int) -> None:
+    pass
+
+
+has_side_effect(torch.ops.streams.synchronize_event.default)
+
+
+@custom_op("streams::synchronize_device", mutates_args=())
+def synchronize_device(device_type: str, device_index: int) -> None:
+    torch.accelerator.synchronize(torch.device(device_type, device_index))
+
+
+@synchronize_device.register_fake
+def _(device_type: str, device_index: int) -> None:
+    pass
+
+
+has_side_effect(torch.ops.streams.synchronize_device.default)
+
+
+@custom_op("streams::synchronize_stream", mutates_args=())
+def synchronize_stream(stream_index: int) -> None:
+    stream = _get_stream_by_index(stream_index)
+    stream.synchronize()
+
+
+@synchronize_stream.register_fake
+def _(stream_index: int) -> None:
+    pass
+
+
+has_side_effect(torch.ops.streams.synchronize_stream.default)
+
+
 @custom_op("streams::wait_stream", mutates_args=())
 def wait_stream(waiting_stream_index: int, waited_on_stream_index: int) -> None:
     waiting = _get_stream_by_index(waiting_stream_index)
@@ -188,6 +232,15 @@ def sync_dealloc(
     torch.ops.streams.wait_event.default(wait_event_index, src_stream_index)
 
 
+@sync_dealloc.register_fake
+def _(
+    wait_event_index: int,
+    src_stream_index: int,
+    to_dealloc: torch.Tensor,
+) -> None:
+    pass
+
+
 has_side_effect(torch.ops.streams.sync_dealloc.default)
 
 
@@ -198,13 +251,15 @@ def record_stream(tensor: torch.Tensor, stream_index: int) -> None:
 
 @record_stream.register_fake
 def _(
-    src_stream_index: int,
-    wait_event_index: int,
-    to_dealloc: torch.Tensor,
+    tensor: torch.Tensor,
+    stream_index: int,
 ) -> None:
     pass
 
 
+has_side_effect(torch.ops.streams.record_stream.default)
+
+
 class SymbolicStreamState:
     """Track the currently entered stream if any"""
 
@@ -213,10 +268,23 @@ def __init__(self) -> None:
 
         cur_stack: list[StreamVariable] = []
         if torch.accelerator.is_available():
-            stream_var = LazyVariableTracker.create(
-                torch.accelerator.current_stream(),
-                source=CurrentStreamSource(torch.accelerator.current_stream().device),
+            # Reset the registry so the current stream is guaranteed index 0.
+            reset_user_object_tracking()
+            stream = torch.accelerator.current_stream()
+            source = CurrentStreamSource(stream.device)
+            # Register the current stream so it gets index 0 (registry is
+            # fresh at tracing start).  The inductor wrapper updates this
+            # entry at runtime so cudagraph capture uses the capture stream
+            # instead of this stale trace-time stream.
+            index = register_user_object(stream, source)
+            assert index == CURRENT_STREAM_INDEX, (
+                f"Current stream must be registered at index {CURRENT_STREAM_INDEX}, "
+                f"got {index}"
             )
+            stream_var = LazyVariableTracker.create(stream, source=source)
+            # Set user_object_index as an instance attribute so accessing it
+            # does NOT trigger LazyVariableTracker realization.
+            stream_var.user_object_index = index  # type: ignore[union-attr]
             cur_stack = [stream_var]  # type: ignore[list-item]
 
         self.cur_stream_stack: collections.deque[StreamVariable] = collections.deque(
@@ -229,7 +297,7 @@ def enter_stream(self, stream: "StreamVariable") -> None:
     def exit_stream(self) -> None:
         self.cur_stream_stack.pop()
 
-    def cur_stream(self, device: Optional[torch.device] = None) -> "StreamVariable":
+    def cur_stream(self, device: torch.device | None = None) -> "StreamVariable":
         if device is not None:
             for stream in reversed(self.cur_stream_stack):
                 if stream.device == device:
@@ -240,6 +308,13 @@ def cur_stream(self, device: Optional[torch.device] = None) -> "StreamVariable":
     def in_stream_context(self) -> bool:
         return len(self.cur_stream_stack) > 0
 
+    def cur_stream_id(self) -> int:
+        """Get a Python object id for the current stream without realizing lazy variables."""
+        stream = self.cur_stream_stack[-1]
+        if isinstance(stream, LazyVariableTracker) and not stream.is_realized():
+            return id(stream.peek_value())
+        return id(stream.value)
+
 
 class StreamContextVariable(FxTracebackAnnotateVariable):
     """This represents torch.cuda.StreamContext"""
@@ -279,6 +354,9 @@ def exit(
         tx.symbolic_stream_state.exit_stream()
         return super().exit(tx, *args)
 
+    def python_type(self) -> type:
+        return torch.cuda.StreamContext
+
     def supports_graph_breaks(self) -> bool:
         return True
 
@@ -290,11 +368,13 @@ def get_stream(self) -> "StreamVariable":
 class StreamVariable(StreamContextVariable):
     """Represents the device-agnostic torch.Stream class"""
 
+    _cpython_type = torch.Stream
+
     def __init__(
         self,
         proxy: Proxy,
         value: torch.Stream,
-        user_object_index: Optional[int] = None,
+        user_object_index: int | None = None,
         **kwargs: Any,
     ) -> None:
         # Index into the user object table
@@ -304,7 +384,6 @@ def __init__(
 
         self.proxy = proxy
         self.value = value
-        # pyrefly: ignore [read-only]
         self.device = value.device
 
         self.user_object_index = user_object_index
@@ -313,6 +392,9 @@ def __init__(
     def python_type(self) -> type:
         return torch.Stream
 
+    def get_real_python_backed_value(self) -> object:
+        return self.value
+
     def call_method(
         self,
         tx: "InstructionTranslator",
@@ -325,11 +407,34 @@ def call_method(
         from ..utils import cmp_name_to_op_mapping, proxy_args_kwargs
         from .builder import wrap_fx_proxy_cls
 
-        if name in ("wait_stream", "synchronize", "wait_event"):
+        if name == "wait_event":
+            event_arg = args[0]
+            assert isinstance(event_arg, EventVariable)
+            tx.output.create_proxy(
+                "call_function",
+                torch.ops.streams.wait_event,
+                (event_arg.user_object_index, self.user_object_index),
+                {},
+            )
+            return ConstantVariable.create(None)
+        elif name == "wait_stream":
+            other_stream = args[0]
+            assert isinstance(other_stream, StreamVariable)
+            tx.output.create_proxy(
+                "call_function",
+                torch.ops.streams.wait_stream,
+                (self.user_object_index, other_stream.user_object_index),
+                {},
+            )
+            return ConstantVariable.create(None)
+        elif name == "synchronize":
             tx.output.create_proxy(
-                "call_method", name, *proxy_args_kwargs([self] + args, kwargs)
+                "call_function",
+                torch.ops.streams.synchronize_stream,
+                (self.user_object_index,),
+                {},
             )
-            return CONSTANT_VARIABLE_NONE
+            return ConstantVariable.create(None)
         elif name == "query":
             return wrap_fx_proxy_cls(
                 target_cls=ConstantVariable,
@@ -339,11 +444,34 @@ def call_method(
                 ),
             )
         elif name == "record_event":
-            return wrap_fx_proxy_cls(
-                target_cls=EventVariable,
+            from .builder import wrap_fx_proxy
+
+            tx.output.check_event_record_after_input_mutation(id(self.value))
+            if args and isinstance(args[0], EventVariable):
+                event_var = args[0]
+                event = event_var.value
+                event_index = event_var.user_object_index
+            else:
+                event = self.value.record_event()
+                event_index = register_graph_created_object(
+                    event,
+                    EventVariable.make_construct_in_graph_event_fn(
+                        TupleVariable([]), ConstDictVariable({})
+                    ),
+                )
+            tx.output.create_proxy(
+                "call_function",
+                torch.ops.streams.record_event,
+                (event_index, self.user_object_index),
+                {},
+            )
+            return wrap_fx_proxy(
                 tx=tx,
                 proxy=tx.output.create_proxy(
-                    "call_method", name, *proxy_args_kwargs([self] + args, kwargs)
+                    "call_function",
+                    get_external_object_by_index,
+                    (event_index,),
+                    {},
                 ),
             )
         elif name in cmp_name_to_op_mapping and len(args) == 1 and not kwargs:
@@ -356,13 +484,14 @@ def call_method(
             # constant values
             other = args[0]
             if not isinstance(other, StreamVariable):
-                return ConstantVariable.create(NotImplemented)
+                return VariableTracker.build(tx, NotImplemented)
 
             if other.source:
                 assert self.source is not None
                 install_guard(self.source.make_guard(GuardBuilder.EQUALS_MATCH))
-            return ConstantVariable.create(
-                cmp_name_to_op_mapping[name](self.value, other.value)  # type: ignore[arg-type]
+            return VariableTracker.build(
+                tx,
+                cmp_name_to_op_mapping[name](self.value, other.value),  # type: ignore[arg-type]
             )
 
         return super().call_method(tx, name, args, kwargs)
@@ -422,12 +551,38 @@ def fn(index: int, codegen: "PyCodegen") -> None:
         return fn
 
 
+class CudaStreamVariable(StreamVariable):
+    """Represents torch.cuda.Stream, preserving device-specific type and attributes."""
+
+    _cpython_type = torch.cuda.Stream
+
+    def python_type(self) -> type:
+        return torch.cuda.Stream
+
+    def var_getattr(self, tx: "InstructionTranslator", name: str) -> "VariableTracker":
+        from . import ConstantVariable
+
+        if name == "cuda_stream":
+            from ..guards import GuardBuilder, install_guard
+
+            if self.source:
+                install_guard(self.source.make_guard(GuardBuilder.EQUALS_MATCH))
+
+            if hasattr(self.value, "cuda_stream"):
+                return ConstantVariable.create(self.value.cuda_stream)
+
+            if hasattr(self.value, "native_handle"):
+                return ConstantVariable.create(self.value.native_handle)
+
+        return super().var_getattr(tx, name)
+
+
 class EventVariable(VariableTracker):
     def __init__(
         self,
         proxy: Proxy,
         value: torch.Event,
-        user_object_index: Optional[int],
+        user_object_index: int | None,
         **kwargs: Any,
     ) -> None:
         if proxy is not None and "example_value" in proxy.node.meta:
@@ -437,6 +592,12 @@ def __init__(
         self.value = value
         self.user_object_index = user_object_index
 
+    def python_type(self) -> type:
+        return torch.Event
+
+    def get_real_python_backed_value(self) -> object:
+        return self.value
+
     def call_method(
         self,
         tx: "InstructionTranslator",
@@ -448,32 +609,38 @@ def call_method(
         from .builder import wrap_fx_proxy_cls
 
         if name == "wait":
+            _, stream_index = EventVariable._get_stream_arg(tx, args, kwargs)
             tx.output.create_proxy(
                 "call_function",
                 torch.ops.streams.wait_event,
                 (
                     self.user_object_index,
-                    EventVariable._get_stream_arg(tx, args, kwargs).user_object_index,
+                    stream_index,
                 ),
                 {},
             )
-            return CONSTANT_VARIABLE_NONE
+            return ConstantVariable.create(None)
         elif name == "record":
+            stream_arg, stream_index = EventVariable._get_stream_arg(tx, args, kwargs)
+            tx.output.check_event_record_after_input_mutation(id(stream_arg.value))
             tx.output.create_proxy(
                 "call_function",
                 torch.ops.streams.record_event,
                 (
                     self.user_object_index,
-                    EventVariable._get_stream_arg(tx, args, kwargs).user_object_index,
+                    stream_index,
                 ),
                 {},
             )
-            return CONSTANT_VARIABLE_NONE
+            return ConstantVariable.create(None)
         elif name == "synchronize":
             tx.output.create_proxy(
-                "call_method", name, *proxy_args_kwargs([self] + args, kwargs)
+                "call_function",
+                torch.ops.streams.synchronize_event,
+                (self.user_object_index,),
+                {},
             )
-            return CONSTANT_VARIABLE_NONE
+            return ConstantVariable.create(None)
         elif name == "query":
             return wrap_fx_proxy_cls(
                 target_cls=ConstantVariable,
@@ -504,7 +671,14 @@ def _get_stream_arg(
         tx: "InstructionTranslator",
         args: list[VariableTracker],
         kwargs: dict[str, VariableTracker],
-    ) -> "StreamVariable":
+    ) -> tuple["StreamVariable", int]:
+        """Returns (stream_variable, stream_index_for_op).
+
+        The ambient current stream is registered at index 0 in the external
+        object registry.  The inductor wrapper updates index 0 at runtime so
+        that cudagraph capture sees the capture stream, not the stale
+        trace-time default stream.
+        """
         stream_arg = None
         if args:
             stream_arg = args[0]
@@ -512,9 +686,10 @@ def _get_stream_arg(
             stream_arg = kwargs.get("stream")
 
         if not stream_arg:
-            stream_arg = tx.symbolic_stream_state.cur_stream()
+            stream_var = tx.symbolic_stream_state.cur_stream()
+            return stream_var, stream_var.user_object_index  # type: ignore[return-value]
 
-        return stream_arg  # type: ignore[return-value]
+        return stream_arg, stream_arg.user_object_index  # type: ignore[return-value]
 
     @staticmethod
     def make_construct_in_graph_event_fn(
diff --git a/torch/_dynamo/variables/tensor.py b/torch/_dynamo/variables/tensor.py
index 703dcd5a4c519..45aff0fae9fba 100644
--- a/torch/_dynamo/variables/tensor.py
+++ b/torch/_dynamo/variables/tensor.py
@@ -32,8 +32,10 @@
 import torch._numpy as tnp
 import torch.fx
 import torch.random
+from torch import sym_float, sym_int
 from torch._dynamo import compiled_autograd
 from torch._library.opaque_object import is_opaque_reference_type
+from torch._opaque_base import OpaqueBase
 from torch._subclasses.meta_utils import is_sparse_any
 from torch.fx.experimental.symbolic_shapes import (
     guard_scalar,
@@ -47,6 +49,8 @@
 from .. import config, graph_break_hints, variables
 from .._trace_wrapped_higher_order_op import trace_wrapped
 from ..exc import (
+    ObservedAttributeError,
+    raise_observed_exception,
     TorchRuntimeError,
     unimplemented,
     UnknownPropertiesDuringBackwardTrace,
@@ -70,7 +74,7 @@
     tensortype_to_dtype,
 )
 from .base import AttributeMutationNew, ValueMutationNew, VariableTracker
-from .constant import CONSTANT_VARIABLE_NONE, ConstantVariable
+from .constant import ConstantVariable
 from .lists import ListIteratorVariable, SizeVariable
 from .script_object import TorchScriptObjectVariable
 from .user_defined import UserDefinedClassVariable
@@ -197,7 +201,6 @@ def __init__(
         super().__init__(**kwargs)
         self.proxy = proxy
         self.dtype = dtype
-        # pyrefly: ignore[read-only]
         self.device = device
         self.layout = layout
         self.ndim = ndim
@@ -251,9 +254,10 @@ def _sync_if_inplace_mutation(
             version_before is not None
             and version_after is not None
             and version_after > version_before
-            and has_tensor_arg
         ):
-            self.synchronize_attributes(tx)
+            if has_tensor_arg:
+                self.synchronize_attributes(tx)
+            tx.output.check_input_mutation_on_current_stream(tx)
 
     def debug_repr(self) -> str:
         # TODO: strip off fake tensor from repr here
@@ -268,6 +272,19 @@ def python_type(self) -> type:
     def is_tensor(self) -> bool:
         return True
 
+    def bool_impl(self, tx: "InstructionTranslator") -> VariableTracker:
+        # THPVariable_bool calls at::Tensor::is_nonzero(), i.e. .item() != 0.
+        from .constant import ConstantVariable
+
+        item = self.call_method(tx, "item", [], {})
+        if isinstance(item, SymNodeVariable) and isinstance(
+            item.sym_num, torch.SymBool
+        ):
+            return item
+        if isinstance(item, ConstantVariable):
+            return VariableTracker.build(tx, bool(item.value))
+        return SymNodeVariable.create(tx, item.as_proxy() != 0)
+
     @staticmethod
     def specialize(value: torch.Tensor) -> dict[str, Any]:
         props: dict[str, Any] = {
@@ -332,8 +349,8 @@ def dynamic_getattr(
             proxy = getattr(self.as_proxy(), name)
             example_value = getattr(fake_val, name)
             if name in attrs:
-                # attrs returned from tensor_flatten are always tensors
-                assert isinstance(example_value, torch.Tensor)
+                # attrs returned from tensor_flatten are always tensors or opaques
+                assert isinstance(example_value, (torch.Tensor, OpaqueBase))
                 from .builder import wrap_fx_proxy
 
                 return wrap_fx_proxy(tx=tx, proxy=proxy, example_value=example_value)
@@ -342,6 +359,11 @@ def dynamic_getattr(
                     tx.output.fake_mode, example_value
                 )
                 return TorchScriptObjectVariable.create(proxy, fake_script_obj)
+            elif isinstance(
+                example_value,
+                torch._library.fake_class_registry.FakeScriptObject,
+            ):
+                return TorchScriptObjectVariable.create(proxy, example_value)
             # any other attributes on the subclass (that are not methods)
             # are assumed to be constant metadata.
             elif not callable(example_value):
@@ -375,7 +397,16 @@ def dynamic_getattr(
         if get_custom_getattr(_input_associated_real_value):
             raise NotImplementedError
 
-        real_value = getattr(_input_associated_real_value, name)
+        try:
+            real_value = getattr(_input_associated_real_value, name)
+        except AttributeError:
+            raise_observed_exception(
+                AttributeError,
+                tx,
+                args=[
+                    f"'{type(_input_associated_real_value).__name__}' object has no attribute '{name}'"
+                ],
+            )
 
         attr_source = AttrSource(self.source, name)
 
@@ -389,41 +420,43 @@ def dynamic_getattr(
                 self, name, source=attr_source, py_type=type(real_value)
             )
 
-        install_guard(attr_source.make_guard(GuardBuilder.HASATTR))
+        install_guard(
+            self.source.make_guard(functools.partial(GuardBuilder.HASATTR, attr=name))
+        )
         return VariableTracker.build(tx, real_value, attr_source)
 
     def method_attr_ndim(self, tx: "InstructionTranslator") -> VariableTracker:
         if self.ndim is not None:
-            return ConstantVariable.create(self.ndim)
+            return VariableTracker.build(tx, self.ndim)
         else:
             return self.call_method(tx, "dim", [], {})
 
     def method_attr_dtype(self, tx: "InstructionTranslator") -> VariableTracker | None:
         if self.dtype is not None:
-            return ConstantVariable.create(self.dtype)
+            return VariableTracker.build(tx, self.dtype)
         return None
 
     def method_attr_device(self, tx: "InstructionTranslator") -> VariableTracker | None:
         if self.device is not None:
-            return ConstantVariable.create(self.device)
+            return VariableTracker.build(tx, self.device)
         return None
 
     def method_attr_layout(self, tx: "InstructionTranslator") -> VariableTracker | None:
         if self.layout is not None:
-            return ConstantVariable.create(self.layout)
+            return VariableTracker.build(tx, self.layout)
         return None
 
     def method_attr_is_cuda(
         self, tx: "InstructionTranslator"
     ) -> ConstantVariable | None:
         if self.device is not None:
-            return ConstantVariable.create(self.device.type == "cuda")
+            return VariableTracker.build(tx, self.device.type == "cuda")
         return None
 
     def method_attr_shape(self, tx: "InstructionTranslator") -> VariableTracker:
         if self.valid_size():
             sizes: list[VariableTracker] = [
-                variables.ConstantVariable.create(x) for x in self.size
+                VariableTracker.build(tx, x) for x in self.size
             ]
             return SizeVariable(sizes)
         else:
@@ -433,28 +466,28 @@ def method_attr_requires_grad(
         self, tx: "InstructionTranslator"
     ) -> ConstantVariable | None:
         if self.requires_grad is not None:
-            return ConstantVariable.create(self.requires_grad)
+            return VariableTracker.build(tx, self.requires_grad)
         return None
 
     def method_attr_is_quantized(
         self, tx: "InstructionTranslator"
     ) -> ConstantVariable | None:
         if self.is_quantized is not None:
-            return ConstantVariable.create(self.is_quantized)
+            return VariableTracker.build(tx, self.is_quantized)
         return None
 
     def method_attr_is_sparse(
         self, tx: "InstructionTranslator"
     ) -> ConstantVariable | None:
         if self.is_sparse is not None:
-            return ConstantVariable.create(self.is_sparse)
+            return VariableTracker.build(tx, self.is_sparse)
         return None
 
     def method_attr_is_nested(
         self, tx: "InstructionTranslator"
     ) -> ConstantVariable | None:
         if self.is_nested is not None:
-            return ConstantVariable.create(self.is_nested)
+            return VariableTracker.build(tx, self.is_nested)
         return None
 
     def method_attr_retain_grad(self, tx: "InstructionTranslator") -> NoReturn:
@@ -465,6 +498,12 @@ def method_attr_retain_grad(self, tx: "InstructionTranslator") -> NoReturn:
             hints=[],
         )
 
+    def method_attr_grad(self, tx: "InstructionTranslator") -> VariableTracker | None:
+        if tx.output.side_effects.has_pending_mutation_of_attr(self, "grad"):
+            return tx.output.side_effects.load_attr(self, "grad")
+        # None tells var_getattr to use default .grad handling
+        return None
+
     def method_attr_data(self, tx: "InstructionTranslator") -> VariableTracker:
         return variables.TorchInGraphFunctionVariable(
             torch._C._autograd._get_data_attr  # type: ignore[attr-defined]
@@ -481,7 +520,7 @@ def method_attr_grad_fn(
                 hints=[],
             )
         else:
-            return variables.CONSTANT_VARIABLE_NONE
+            return variables.ConstantVariable.create(None)
 
     def method_attr__version(self, tx: "InstructionTranslator") -> VariableTracker:
         from ..tensor_version_op import _tensor_version
@@ -494,31 +533,32 @@ def call_obj_hasattr(
         self, tx: "InstructionTranslator", name: str
     ) -> ConstantVariable:
         from . import GetAttrVariable
-        from .builtin import BuiltinVariable
 
         # TODO - This is not a good solution but solves an accuracy issue.
         # Today, var_getattr returns GetAttrVariable for both non-existent
         # attributes and existing attributes. This is a bug and requires more
         # deep dive.
         if name in all_tensor_attrs:
-            return ConstantVariable(True)
+            return ConstantVariable.create(True)
 
         try:
-            var = BuiltinVariable(getattr).call_function(
-                tx, [self, ConstantVariable(name)], {}
+            var = VariableTracker.build(tx, getattr).call_function(
+                tx, [self, VariableTracker.build(tx, name)], {}
             )
             # in the event that TensorVariable returns NotImplemented
             # BuiltinVariable.call_getattr returns GetAttrVariable
             ret_val = not isinstance(var, GetAttrVariable)
-        except AttributeError:
+        except (AttributeError, ObservedAttributeError):
             ret_val = False
 
         if self.source:
             install_guard(
-                AttrSource(self.source, name).make_guard(GuardBuilder.HASATTR)
+                self.source.make_guard(
+                    functools.partial(GuardBuilder.HASATTR, attr=name)
+                )
             )
 
-        return ConstantVariable(ret_val)
+        return VariableTracker.build(tx, ret_val)
 
     def var_getattr(self, tx: "InstructionTranslator", name: str) -> VariableTracker:
         if self.is_strict_mode(tx):
@@ -538,7 +578,7 @@ def var_getattr(self, tx: "InstructionTranslator", name: str) -> VariableTracker
                 )
 
         if name == "__class__":
-            return UserDefinedClassVariable(self.python_type())
+            return VariableTracker.build(tx, self.python_type())
 
         handler = getattr(self, f"method_attr_{name}", None)
         result = handler(tx) if handler is not None else None
@@ -646,7 +686,7 @@ def call_id(self, tx: "InstructionTranslator") -> VariableTracker:
 
         install_guard(self.source.make_guard(GuardBuilder.ID_MATCH))
         id_value = id(_input_associated_real_value)
-        return ConstantVariable.create(id_value)
+        return VariableTracker.build(tx, id_value)
 
     def has_unpack_var_sequence(self, tx: "InstructionTranslator") -> bool:
         return self.ndim > 0
@@ -669,7 +709,9 @@ def unpack_var_sequence(
         if self.valid_size():
             length = self.size[0]
         else:
-            dyn_length = self.call_method(tx, "size", [ConstantVariable.create(0)], {})
+            dyn_length = self.call_method(
+                tx, "size", [VariableTracker.build(tx, 0)], {}
+            )
             # SymNodeVariable for symbolic sizes, ConstantVariable for constants OR values produced through
             # symbolic_shapes, but that end up as int/sympy.Integer
             assert (
@@ -734,6 +776,41 @@ def _strict_mode_conditional_banned_ops(self) -> list[str]:
             torch._dynamo.config._autograd_backward_strict_mode_conditional_banned_ops
         )
 
+    def mp_subscript_impl(
+        self,
+        tx: "InstructionTranslator",
+        key: VariableTracker,
+    ) -> VariableTracker:
+        # Tensor.__getitem__ is a custom C slot, not CPython's mp_subscript.
+        # TODO(follow-up): add tests for negative index, bool index, invalid key type
+        from .builder import SourcelessBuilder, VariableBuilder
+        from .torch_function import can_dispatch_torch_function, dispatch_torch_function
+
+        if self.is_strict_mode(tx) and "__getitem__" in self._strict_mode_banned_ops():
+            unimplemented(
+                gb_type="Illegal __getitem__ invocation in strict mode",
+                context=f"mp_subscript_impl {self} {key}",
+                explanation="Dynamo currently does not support __getitem__ "
+                "invocation in strict mode.",
+                hints=[],
+            )
+
+        empty_kwargs: dict[str, VariableTracker] = {}
+        static_attr = all_tensor_attrs.get("__getitem__", None)
+        if static_attr is not None and can_dispatch_torch_function(
+            tx, (self, key), empty_kwargs
+        ):
+            if self.source:
+                func_var = VariableBuilder(
+                    tx,
+                    AttrSource(AttrSource(self.source, "__class__"), "__getitem__"),
+                )(static_attr)
+            else:
+                func_var = SourcelessBuilder.create(tx, torch.Tensor.__getitem__)
+            return dispatch_torch_function(tx, func_var, (self, key), empty_kwargs)
+
+        return self.method___getitem__(tx, key)
+
     def call_method(
         self,
         tx: "InstructionTranslator",
@@ -753,6 +830,16 @@ def call_method(
                 hints=[],
             )
 
+        if name == "__deepcopy__":
+            unimplemented(
+                gb_type="Attempted to copy.deepcopy a tensor",
+                context=f"copy.deepcopy({self})",
+                explanation="Dynamo does not support copy.deepcopy() on tensors.",
+                hints=[
+                    "Avoid calling copy.deepcopy() on tensors inside compiled regions.",
+                ],
+            )
+
         # Only override builtin tensor methods
         # The user can manually add override handling
         # with a decorator for other methods (e.g. a dispatch subclass with other methods)
@@ -782,7 +869,23 @@ def call_method(
 
         # This is seen in inspect signature where we check if the value is a default value
         if name == "__eq__" and isinstance(args[0], UserDefinedClassVariable):
-            return variables.ConstantVariable(False)
+            return variables.ConstantVariable.create(False)
+
+        if name == "wait":
+            if args or kwargs:
+                raise torch._dynamo.exc.InternalTorchDynamoError(
+                    "`wait` and `wait_tensor` do not take any arguments"
+                )
+            from torch.distributed._functional_collectives import wait_tensor
+
+            from .builder import wrap_fx_proxy
+
+            return wrap_fx_proxy(
+                tx,
+                tx.output.create_proxy(
+                    "call_function", wait_tensor, (self.as_proxy(),), {}
+                ),
+            )
 
         # For historical reasons, these ops decompose down to syntactically
         # invalid aten ops because they contain the python keyword `from`, see
@@ -830,6 +933,23 @@ def call_method(
                     from_exc=e,
                 )
 
+        # Guard against unknown methods reaching the generic proxy path.
+        # For traceable wrapper subclasses (DTensor, NestedTensor), class_type
+        # is torch.Tensor, so check the example_value's actual type instead.
+        example_value = self.proxy.node.meta.get("example_value")
+        check_type = (
+            type(example_value) if example_value is not None else self.class_type
+        )
+        if not hasattr(check_type, name):
+            unimplemented(
+                gb_type="Unhandled tensor method",
+                context=f"call_method {self} {name} {args} {kwargs}",
+                explanation=f"Tensor method `{name}` is not defined on "
+                f"{check_type.__name__} and does not have an explicit "
+                "handler in TensorVariable.",
+                hints=[*graph_break_hints.SUPPORTABLE],
+            )
+
         from .builder import wrap_fx_proxy
 
         proxy = tx.output.create_proxy(
@@ -910,20 +1030,20 @@ def make_const_size_variable(x: Sequence[Any], **options: Any) -> SizeVariable:
 
     def method_numel(self, tx: "InstructionTranslator") -> VariableTracker | None:
         if self.valid_size():
-            return ConstantVariable.create(product(self.size))
+            return VariableTracker.build(tx, product(self.size))
 
         # It might still be constant!  Consult the fake tensor and see
         if (fake := self.proxy.node.meta.get("example_value")) is not None:
             fake_r = fake.numel()
             if not has_free_symbols(fake_r):
-                return ConstantVariable.create(int(fake_r))
+                return VariableTracker.build(tx, int(fake_r))
         return None
 
     method_nelement = method_numel
 
     def method_dim(self, tx: "InstructionTranslator") -> VariableTracker | None:
         if self.ndim is not None:
-            return ConstantVariable.create(self.ndim)
+            return VariableTracker.build(tx, self.ndim)
         return None
 
     method_ndimension = method_dim
@@ -932,7 +1052,7 @@ def method_is_floating_point(
         self, tx: "InstructionTranslator"
     ) -> ConstantVariable | None:
         if self.dtype is not None:
-            return ConstantVariable.create(self.dtype.is_floating_point)
+            return VariableTracker.build(tx, self.dtype.is_floating_point)
         return None
 
     def method_is_inference(
@@ -949,12 +1069,12 @@ def method_is_inference(
                 ],
             )
         if (fake := self.proxy.node.meta.get("example_value")) is not None:
-            return ConstantVariable.create(fake.is_inference())
+            return VariableTracker.build(tx, fake.is_inference())
         return None
 
     def method_is_complex(self, tx: "InstructionTranslator") -> ConstantVariable | None:
         if self.dtype is not None:
-            return ConstantVariable.create(self.dtype.is_complex)
+            return VariableTracker.build(tx, self.dtype.is_complex)
         return None
 
     def method_is_contiguous(
@@ -967,10 +1087,10 @@ def method_is_contiguous(
         )
         if self.is_contiguous is not None:
             # pyrefly: ignore[not-iterable]
-            return ConstantVariable.create(memory_format_const in self.is_contiguous)
+            return VariableTracker.build(tx, memory_format_const in self.is_contiguous)
         elif (fake := self.proxy.node.meta.get("example_value")) is not None:
-            return ConstantVariable.create(
-                fake.is_contiguous(memory_format=memory_format_const)
+            return VariableTracker.build(
+                tx, fake.is_contiguous(memory_format=memory_format_const)
             )
         return None
 
@@ -990,10 +1110,10 @@ def method_type(
                 k for k, v in tensortype_to_dtype.items() if self.dtype in v
             )
             if self.device.type == "cpu":
-                return ConstantVariable.create(f"torch.{tensortype.__name__}")
+                return VariableTracker.build(tx, f"torch.{tensortype.__name__}")
             else:
-                return ConstantVariable.create(
-                    f"torch.{self.device.type}.{tensortype.__name__}"
+                return VariableTracker.build(
+                    tx, f"torch.{self.device.type}.{tensortype.__name__}"
                 )
         elif (
             dtype is not None
@@ -1003,7 +1123,7 @@ def method_type(
             # torch.fx's tracer fails on these types, because it doesn't support arguments of torch.tensortype type.
             # So, we pass it in as a string (which is also supported, see above implementation for .type() with 0 args)
             tensor_type = dtype.as_python_constant()
-            tensor_type_const = ConstantVariable.create(fqn(tensor_type))
+            tensor_type_const = VariableTracker.build(tx, fqn(tensor_type))
 
             from .builder import wrap_fx_proxy
 
@@ -1048,11 +1168,11 @@ def method_as_subclass(
     def method_get_device(self, tx: "InstructionTranslator") -> VariableTracker | None:
         if isinstance(self.device, torch.device):
             index = self.device.index if self.device.type != "cpu" else -1
-            return ConstantVariable.create(index)
+            return VariableTracker.build(tx, index)
         return None
 
     def method_element_size(self, tx: "InstructionTranslator") -> VariableTracker:
-        return ConstantVariable.create(self.dtype.itemsize)
+        return VariableTracker.build(tx, self.dtype.itemsize)
 
     def method_numpy(
         self, tx: "InstructionTranslator", *, force: VariableTracker | bool = False
@@ -1140,7 +1260,7 @@ def wrap(i: Any, sub_proxy: torch.fx.Proxy) -> VariableTracker:
 
     def _collect_backward_inputs(
         self, vars_iter: Iterable[VariableTracker], error_on_non_leaf: bool = False
-    ) -> Optional[list[VariableTracker]]:
+    ) -> list[VariableTracker] | None:
         """
         Collect unique leaf tensors from vars_iter for backward.
 
@@ -1164,7 +1284,11 @@ def _collect_backward_inputs(
                 #    (Dynamo creates GetAttrVariable instead of TensorVariable)
                 #
                 # In-graph created tensors without proper source also can't be handled
-                # because subguards_allowed() returns False for SyntheticLocalSource.
+                # when user explicitly passes them as inputs, because
+                # subguards_allowed() returns False for SyntheticLocalSource.
+                # However, in auto-detect mode (error_on_non_leaf=False), source-less
+                # leaves are valid backward targets — they gained requires_grad via
+                # requires_grad_() and accumulate_grad_ writes to .grad directly.
                 if var.has_grad_fn:
                     if error_on_non_leaf:
                         unimplemented(
@@ -1175,17 +1299,18 @@ def _collect_backward_inputs(
                                 "Only pass leaf tensors (parameters, graph inputs) to backward(inputs=...)",
                             ],
                         )
-                elif not var.source or isinstance(var.source, SyntheticLocalSource):
-                    if error_on_non_leaf:
-                        unimplemented(
-                            gb_type="backward() with in-graph created tensor",
-                            context=f"backward(inputs=[...]) with in-graph created tensor: {var}",
-                            explanation="backward(inputs=[...]) with tensors created inside the "
-                            "compiled function is not yet supported.",
-                            hints=[
-                                "Only pass tensors that are inputs to the compiled function or captured from outside",
-                            ],
-                        )
+                elif error_on_non_leaf and (
+                    not var.source or isinstance(var.source, SyntheticLocalSource)
+                ):
+                    unimplemented(
+                        gb_type="backward() with in-graph created tensor",
+                        context=f"backward(inputs=[...]) with in-graph created tensor: {var}",
+                        explanation="backward(inputs=[...]) with tensors created inside the "
+                        "compiled function is not yet supported.",
+                        hints=[
+                            "Only pass tensors that are inputs to the compiled function or captured from outside",
+                        ],
+                    )
                 else:
                     node = var.proxy.node
                     if node not in seen_nodes:
@@ -1197,11 +1322,11 @@ def _collect_backward_inputs(
     def method_backward(
         self,
         tx: "InstructionTranslator",
-        gradient: Optional[VariableTracker] = None,
-        retain_graph: Optional[VariableTracker] = None,
-        create_graph: Optional[VariableTracker] = None,
-        inputs: Optional[VariableTracker] = None,
-    ) -> Optional[VariableTracker]:
+        gradient: VariableTracker | None = None,
+        retain_graph: VariableTracker | None = None,
+        create_graph: VariableTracker | None = None,
+        inputs: VariableTracker | None = None,
+    ) -> VariableTracker | None:
         """
         Trace tensor.backward() by rewriting as autograd.grad() + accumulate_grad.
 
@@ -1253,7 +1378,7 @@ def method_backward(
                 # No leaf tensors found - nothing to accumulate gradients into.
                 # This matches eager behavior where backward() is a no-op if there
                 # are no leaves requiring grad.
-                return CONSTANT_VARIABLE_NONE
+                return ConstantVariable.create(None)
         else:
             provided_vars = (
                 inputs.items
@@ -1309,7 +1434,7 @@ def method_backward(
 
         grad_mode_var.exit(tx)
 
-        return VariableTracker.build(tx, None)
+        return ConstantVariable.create(None)
 
     def method_data_ptr(
         self,
@@ -1319,6 +1444,25 @@ def method_data_ptr(
     ) -> "DataPtrVariable":
         return DataPtrVariable(self)
 
+    def method_record_stream(
+        self,
+        tx: "InstructionTranslator",
+        stream: VariableTracker,
+    ) -> VariableTracker:
+        from .streams import StreamVariable
+
+        if not isinstance(stream, StreamVariable):
+            raise RuntimeError(
+                f"record_stream() expects a Stream argument, got {stream.python_type().__name__}"
+            )
+        tx.output.create_proxy(
+            "call_function",
+            torch.ops.streams.record_stream,
+            (self.as_proxy(), stream.user_object_index),
+            {},
+        )
+        return ConstantVariable.create(None)
+
     def method_item(
         self,
         tx: "InstructionTranslator",
@@ -1341,6 +1485,95 @@ def method_item(
             )
         return None
 
+    def nb_index_impl(
+        self,
+        tx: "InstructionTranslator",
+    ) -> VariableTracker:
+        # CPython: only integer tensors of a single element can be converted
+        # to an index. Mirrors THPVariable_index_scalar in
+        # https://github.com/pytorch/pytorch/blob/7cfd054075b/tools/autograd/templates/python_variable_methods.cpp#L372-L385
+        if self.dtype is not None and (
+            not self.dtype.is_floating_point and not self.dtype.is_complex
+        ):
+            item = self.call_method(tx, "item", [], {})
+            from .builder import wrap_fx_proxy
+
+            return wrap_fx_proxy(
+                tx=tx,
+                proxy=tx.output.create_proxy(
+                    "call_function",
+                    sym_int,
+                    (item.as_proxy(),),
+                    {},
+                ),
+            )
+        raise_observed_exception(
+            TypeError,
+            tx,
+            args=[
+                "only integer tensors of a single element can be converted to an index"
+            ],
+        )
+
+    def nb_int_impl(
+        self,
+        tx: "InstructionTranslator",
+    ) -> VariableTracker:
+        # CPython: THPVariable_integral_scalar handles both integer and float
+        # tensors (floats are truncated to int). Complex tensors raise
+        # RuntimeError at runtime.
+        if self.dtype is not None and self.dtype.is_complex:
+            raise_observed_exception(
+                RuntimeError,
+                tx,
+                args=["value cannot be converted to type int64_t without overflow"],
+            )
+        # For known non-complex dtypes and unknown dtype (None), emit the
+        # proxy and let it fail at runtime if the dtype is unsupported.
+        item = self.call_method(tx, "item", [], {})
+        from .builder import wrap_fx_proxy
+
+        return wrap_fx_proxy(
+            tx=tx,
+            proxy=tx.output.create_proxy(
+                "call_function",
+                sym_int,
+                (item.as_proxy(),),
+                {},
+            ),
+        )
+
+    def method___int__(self, tx: "InstructionTranslator") -> VariableTracker:
+        return self.nb_int_impl(tx)
+
+    def nb_float_impl(
+        self,
+        tx: "InstructionTranslator",
+    ) -> VariableTracker:
+        # CPython: THPVariable_float_scalar dispatches to double.
+        # Complex tensors raise RuntimeError at runtime.
+        if self.dtype is not None and self.dtype.is_complex:
+            raise_observed_exception(
+                RuntimeError,
+                tx,
+                args=["value cannot be converted to type double without overflow"],
+            )
+        item = self.call_method(tx, "item", [], {})
+        from .builder import wrap_fx_proxy
+
+        return wrap_fx_proxy(
+            tx=tx,
+            proxy=tx.output.create_proxy(
+                "call_function",
+                sym_float,
+                (item.as_proxy(),),
+                {},
+            ),
+        )
+
+    def method___float__(self, tx: "InstructionTranslator") -> VariableTracker:
+        return self.nb_float_impl(tx)
+
     def method___getitem__(
         self,
         tx: "InstructionTranslator",
@@ -1356,7 +1589,7 @@ def method___getitem__(
             fn, args = (  # type: ignore[assignment]
                 torch.select,
                 [
-                    variables.ConstantVariable.create(0),
+                    VariableTracker.build(tx, 0),
                     args[0],
                 ],
             )
@@ -1393,7 +1626,11 @@ def _warn_capture_scalar_outputs() -> None:
         )
 
     def method___len__(self, tx: "InstructionTranslator") -> VariableTracker:
-        return self.call_method(tx, "size", [ConstantVariable.create(0)], {})
+        return self.sq_length(tx)
+
+    def sq_length(self, tx: "InstructionTranslator") -> VariableTracker:
+        """Sequence length for tensors (size along first dimension)."""
+        return self.call_method(tx, "size", [VariableTracker.build(tx, 0)], {})
 
     def method___iter__(self, tx: "InstructionTranslator") -> ListIteratorVariable:
         return ListIteratorVariable(
@@ -1409,13 +1646,13 @@ def method_addcmul_(
         value: Any | None = None,
     ) -> Any | None:
         if value is not None and config.enable_dynamo_decompositions:
-            from .. import polyfills
+            from torch._inductor import inductor_prims
 
-            return tx.inline_user_function_return(
-                VariableTracker.build(tx, polyfills.addcmul_inplace),
-                [self, tensor1, tensor2, value],
-                {},
-            )
+            mul_var = variables.TorchInGraphFunctionVariable(torch.mul)
+            fma_var = variables.TorchInGraphFunctionVariable(inductor_prims.fma)
+            product = mul_var.call_function(tx, [tensor1, tensor2], {})
+            result = fma_var.call_function(tx, [product, value, self], {})
+            return self.call_method(tx, "copy_", [result], {})
         return None
 
     def method___setitem__(
@@ -1449,7 +1686,7 @@ def method___setitem__(
         if config.use_graph_deduplication or config.track_nodes_for_deduplication:
             tx.output.region_tracker.add_node_mutation(proxy.node, 0)
 
-        return CONSTANT_VARIABLE_NONE
+        return ConstantVariable.create(None)
 
     def method_resize_(
         self,
@@ -1532,11 +1769,18 @@ def method_add_(
         *,
         alpha: VariableTracker | None = None,
     ) -> VariableTracker | None:
-        if alpha is not None and config.enable_dynamo_decompositions:
-            result = variables.TorchInGraphFunctionVariable(torch.mul).call_function(
-                tx, [other, alpha], {}
-            )
-            return self.call_method(tx, "add_", [result], {})
+        # Tensor alpha: ATen would call .item() and fma internally.
+        # Use fma directly to avoid the item() graph break.
+        if (
+            alpha is not None
+            and isinstance(alpha, TensorVariable)
+            and config.enable_dynamo_decompositions
+        ):
+            from torch._inductor import inductor_prims
+
+            fma_var = variables.TorchInGraphFunctionVariable(inductor_prims.fma)
+            result = fma_var.call_function(tx, [other, alpha, self], {})
+            return self.call_method(tx, "copy_", [result], {})
         return None
 
     def method_addcdiv_(
@@ -1551,10 +1795,11 @@ def method_addcdiv_(
             result = variables.TorchInGraphFunctionVariable(torch.div).call_function(
                 tx, [tensor1, tensor2], {}
             )
-            result = variables.TorchInGraphFunctionVariable(torch.mul).call_function(
-                tx, [result, value], {}
-            )
-            return self.call_method(tx, "add_", [result], {})
+            from torch._inductor import inductor_prims
+
+            fma_var = variables.TorchInGraphFunctionVariable(inductor_prims.fma)
+            fma_result = fma_var.call_function(tx, [result, value, self], {})
+            return self.call_method(tx, "copy_", [fma_result], {})
         return None
 
     def method___contains__(
@@ -1572,81 +1817,6 @@ def method___contains__(
         )
         return result.call_method(tx, "item", [], {})
 
-    def method_redistribute(
-        self,
-        tx: "InstructionTranslator",
-        *args: VariableTracker,
-        **kwargs: VariableTracker,
-    ) -> VariableTracker:
-        # rewrite non-primitive args/kwargs to be included in the on-the-fly prim function
-        # and rewrite args to have only proxyable args, then insert call_function
-        args_as_value = [x.as_python_constant() for x in args]
-        kwargs_as_value = {k: v.as_python_constant() for k, v in kwargs.items()}
-
-        def redistribute_fn_with_prim_types(x: Any) -> Any:
-            return x.redistribute(*args_as_value, **kwargs_as_value)
-
-        # attach the same function name for better debugging
-        redistribute_fn_with_prim_types.__name__ = "prim_redistribute"
-
-        from .builder import wrap_fx_proxy
-
-        return wrap_fx_proxy(
-            tx=tx,
-            proxy=tx.output.create_proxy(
-                "call_function",
-                redistribute_fn_with_prim_types,
-                *proxy_args_kwargs([self], {}),
-            ),
-        )
-
-    def method_to_local(
-        self,
-        tx: "InstructionTranslator",
-        *args: VariableTracker,
-        **kwargs: VariableTracker,
-    ) -> VariableTracker:
-        # rewrite non-primitive args/kwargs to be included in the on-the-fly prim function
-        # and rewrite args to have only proxyable args, then insert call_function
-
-        # custom grad_placements do not work with  as_python_constant,
-        # to support them we need to handle UserDefinedObject
-        def extract_python_value(vt: VariableTracker) -> Any:
-            if isinstance(vt, variables.UserDefinedObjectVariable):
-                return vt.value
-
-            return vt.as_python_constant()
-
-        grad_placements_vt = kwargs.get("grad_placements", CONSTANT_VARIABLE_NONE)
-        if isinstance(grad_placements_vt, variables.UserDefinedObjectVariable):
-            # grad_placement is a sequence-like structure, iterate over the value
-            grad_placements_vt = variables.BuiltinVariable(tuple).call_function(
-                tx, [grad_placements_vt], {}
-            )
-
-        if kwargs.get("grad_placements") is not None:
-            kwargs["grad_placements"] = grad_placements_vt
-
-        args_as_value = [extract_python_value(x) for x in args]
-        kwargs_as_value = {k: extract_python_value(v) for k, v in kwargs.items()}
-
-        def to_local_fn_with_prim_types(x: Any) -> Any:
-            return x.to_local(*args_as_value, **kwargs_as_value)
-
-        # attach the same function name for better debugging
-        to_local_fn_with_prim_types.__name__ = "prim_to_local"
-
-        from .builder import wrap_fx_proxy
-
-        return wrap_fx_proxy(
-            tx=tx,
-            proxy=tx.output.create_proxy(
-                "call_function",
-                to_local_fn_with_prim_types,
-                *proxy_args_kwargs([self], {}),
-            ),
-        )
-
     def method_register_hook(
         self,
         tx: "InstructionTranslator",
@@ -1809,16 +1979,100 @@ def method_requires_grad_(
         if requires_grad is not True:
             requires_grad = requires_grad.as_python_constant()  # type: ignore[attr-defined]
 
-        if self.as_proxy().node.meta["example_value"].requires_grad != requires_grad:
-            unimplemented(
-                gb_type="Unsupported Tensor.requires_grad_() call",
-                context=f"call_method {self} requires_grad_",
-                explanation="Dynamo does not support changes to a Tensor's "
-                "`requires_grad` through calling `requires_grad_()`.",
-                hints=[],
-            )
-        else:
-            return self
+        node = self.as_proxy().node
+        example_value = node.meta["example_value"]
+        if example_value.requires_grad != requires_grad:
+            # For graph inputs (tensors with source), requires_grad_() is a
+            # metadata mutation that we can't trace — graph break as before.
+            if self.source:
+                unimplemented(
+                    gb_type="Unsupported Tensor.requires_grad_() call",
+                    context=f"call_method {self} requires_grad_",
+                    explanation="Dynamo does not support changes to a Tensor's "
+                    "`requires_grad` through calling `requires_grad_()`.",
+                    hints=[],
+                )
+            # On a previous attempt, we traced through requires_grad_() but
+            # discovered at compile time that the tainted intermediate leaked
+            # as a graph output. Graph break here to preserve partial
+            # acceleration for code before requires_grad_().
+            if tx.speculation_log.graph_break_on_requires_grad_:
+                unimplemented(
+                    gb_type="requires_grad_() intermediate leaked as output",
+                    context=f"call_method {self} requires_grad_",
+                    explanation="An intermediate tensor with requires_grad_() called "
+                    "on it (or a tensor derived from it) is returned from the "
+                    "compiled region. Graph breaking here to preserve partial "
+                    "acceleration.",
+                    hints=[
+                        "Call .detach() before returning if you only need values.",
+                        "Consume the gradient inside the compiled function "
+                        "(call backward() and use .grad), "
+                        "or move requires_grad_() outside torch.compile.",
+                    ],
+                )
+            # AOTAutograd re-traces the FX graph under functorch transforms
+            # (functionalization). Functorch's checkSupportsInplaceRequiresGrad()
+            # rejects requires_grad_() when the dynamic layer stack is non-empty.
+            # We wrap the call with set_inplace_requires_grad_allowed(True) to
+            # bypass this check, matching GradInplaceRequiresGradCtxManagerVariable
+            # in ctx_manager.py (which handles the explicit context manager case).
+            #
+            # Lines below do two things in parallel:
+            # 1. Mutate trace-time state so example_value.requires_grad_() works
+            # 2. Emit FX nodes so the same toggle happens during AOTAutograd re-trace
+            prev_state = torch._C._functorch.get_inplace_requires_grad_allowed()
+            torch._C._functorch.set_inplace_requires_grad_allowed(True)
+            try:
+                tx.output.create_node(
+                    "call_function",
+                    torch._C._functorch.set_inplace_requires_grad_allowed,
+                    (True,),
+                    {},
+                )
+                tx.output.create_proxy(
+                    "call_method",
+                    "requires_grad_",
+                    (self.as_proxy(),),
+                    {},
+                )
+                tx.output.create_node(
+                    "call_function",
+                    torch._C._functorch.set_inplace_requires_grad_allowed,
+                    (prev_state,),
+                    {},
+                )
+            finally:
+                torch._C._functorch.set_inplace_requires_grad_allowed(prev_state)
+            example_value.requires_grad_(requires_grad)
+            self.requires_grad = requires_grad
+            if requires_grad:
+                tx.output.leaf_var_creation_order.append(self)
+                # For source-less intermediates, initialize .grad = None in
+                # side effects so the accumulate_grad polyfill can read/write
+                # .grad naturally. Graph inputs don't need this — they handle
+                # .grad through their source.
+                if not self.source and tx.output.side_effects.is_attribute_mutation(
+                    self
+                ):
+                    tx.output.side_effects.store_attr(
+                        self, "grad", variables.ConstantVariable.create(None)
+                    )
+        return self
+
+    def method_detach_(self, tx: "InstructionTranslator") -> "TensorVariable":
+        from .builder import wrap_fx_proxy
+
+        proxy = tx.output.create_proxy(
+            "call_method",
+            "detach_",
+            (self.as_proxy(),),
+            {},
+        )
+        # Run the fake op so the proxy metadata reflects the detached tensor state.
+        wrap_fx_proxy(tx, proxy)
+        self.synchronize_attributes(tx)
+        return self
 
     def method_share_memory_(self) -> NoReturn:
         unimplemented(
@@ -1849,6 +2103,31 @@ def method_new(
             return self.call_method(tx, "new_empty", args, kwargs)
         return None
 
+    def method_new_tensor(
+        self: "TensorVariable",
+        tx: "InstructionTranslator",
+        *args: VariableTracker,
+        **kwargs: VariableTracker,
+    ) -> VariableTracker | None:
+        if len(args) != 1:
+            return None
+        data_arg = args[0]
+        layout = kwargs.get("layout")
+        if layout is not None:
+            if not layout.is_python_constant():
+                return None
+            if layout.as_python_constant() != torch.strided:
+                return None
+        fwd_kwargs = dict(kwargs)
+        fwd_kwargs.pop("layout", None)
+        fwd_kwargs.setdefault("dtype", self.var_getattr(tx, "dtype"))
+        fwd_kwargs.setdefault("device", self.var_getattr(tx, "device"))
+        return variables.TorchInGraphFunctionVariable(torch.tensor).call_function(
+            tx,
+            [data_arg],
+            fwd_kwargs,
+        )
+
     def method_untyped_storage(
         self, tx: "InstructionTranslator"
     ) -> "UntypedStorageVariable":
@@ -1911,11 +2190,11 @@ def create(
 
         if isinstance(sym_num, (sympy.Integer, int, bool)):
             sym_num = int(sym_num) if isinstance(sym_num, sympy.Integer) else sym_num
-            return ConstantVariable.create(sym_num)
+            return VariableTracker.build(tx, sym_num)
 
         out = SymNodeVariable(proxy, sym_num, **options)
         if proxy.node.op != "placeholder":
-            tx.output.current_tracer.record_tensor_or_symint_vt(out)
+            tx.output.current_tracer.record_proxyable_vt(out)
         return out
 
     def __init__(self, proxy: Any, sym_num: Any, **kwargs: Any) -> None:
@@ -1937,6 +2216,18 @@ def is_symnode_like(self) -> bool:
     def as_proxy(self) -> Any:
         return self.proxy
 
+    def bool_impl(
+        self,
+        tx: "InstructionTranslatorBase",
+    ) -> VariableTracker:
+        # long_bool / float_bool: non-zero check. SymBool is already boolean.
+        # https://github.com/python/cpython/blob/c09ccd9c429/Objects/longobject.c#L5200
+        # https://github.com/python/cpython/blob/c09ccd9c429/Objects/floatobject.c#L853
+        if isinstance(self.sym_num, torch.SymBool):
+            return self
+        assert isinstance(self.sym_num, (torch.SymInt, torch.SymFloat))
+        return SymNodeVariable.create(tx, self.as_proxy() != 0)
+
     def as_tensor(self, tx: "InstructionTranslatorBase", dtype: Any) -> TensorVariable:
         if self._tensor_var is None:
             self._tensor_var = VariableTracker.build(
@@ -1977,6 +2268,52 @@ def call_method(
             ),
         )
 
+    def nb_int_impl(
+        self,
+        tx: "InstructionTranslator",
+    ) -> VariableTracker:
+        # SymInt.__int__: https://github.com/pytorch/pytorch/blob/ee336ca5440939b8ad65e916d47421f849e56178/torch/__init__.py#L462
+        # SymFloat.__int__: https://github.com/pytorch/pytorch/blob/ee336ca5440939b8ad65e916d47421f849e56178/torch/__init__.py#L682
+        # SymBool.__int__: https://github.com/pytorch/pytorch/blob/ee336ca5440939b8ad65e916d47421f849e56178/torch/__init__.py#L784
+        from .builder import wrap_fx_proxy
+
+        return wrap_fx_proxy(
+            tx=tx,
+            proxy=tx.output.create_proxy(
+                "call_function",
+                sym_int,
+                (self.as_proxy(),),
+                {},
+            ),
+        )
+
+    def method___int__(
+        self, tx: "InstructionTranslator", *args: Any, **kwargs: Any
+    ) -> VariableTracker:
+        return self.nb_int_impl(tx)
+
+    def nb_float_impl(
+        self,
+        tx: "InstructionTranslator",
+    ) -> VariableTracker:
+        # SymFloat.__float__: https://github.com/pytorch/pytorch/blob/ee336ca5440939b8ad65e916d47421f849e56178/torch/__init__.py#L679
+        from .builder import wrap_fx_proxy
+
+        return wrap_fx_proxy(
+            tx=tx,
+            proxy=tx.output.create_proxy(
+                "call_function",
+                sym_float,
+                (self.as_proxy(),),
+                {},
+            ),
+        )
+
+    def method___float__(
+        self, tx: "InstructionTranslator", *args: Any, **kwargs: Any
+    ) -> VariableTracker:
+        return self.nb_float_impl(tx)
+
     def is_python_hashable(self) -> bool:
         return True
 
@@ -2035,7 +2372,7 @@ def insert_into_graph() -> VariableTracker:
                 ),
             )
 
-        if name in ["T", "real", "imag"]:
+        if name in ["T", "real", "imag", "flat"]:
             proxy = tx.output.create_proxy(
                 "call_function",
                 numpy_attr_wrapper,
@@ -2058,14 +2395,14 @@ def insert_into_graph() -> VariableTracker:
         # NB: only ALWAYS specialized attributes can go here; notably,
         # size/shape not allowed!
         elif name in ("ndim", "itemsize"):
-            return ConstantVariable.create(getattr(example_ndarray, name))
+            return VariableTracker.build(tx, getattr(example_ndarray, name))
         elif name in ("shape", "stride"):
             if not has_free_symbols(r := getattr(example_ndarray, name)):
-                return ConstantVariable.create(tuple(int(r) for r in r))
+                return VariableTracker.build(tx, tuple(int(r) for r in r))
             return insert_into_graph()
         elif name == "size":
             if not has_free_symbols(r := example_ndarray.size):
-                return ConstantVariable.create(int(r))
+                return VariableTracker.build(tx, int(r))
             return insert_into_graph()
         elif name in ["base", "flags", "dtype"]:
             unimplemented(
@@ -2283,6 +2620,9 @@ def __init__(
         # Example_value will always have device="meta"
         self.example_value = example_value
 
+    def python_type(self) -> type:
+        return torch.UntypedStorage
+
     def call_method(
         self,
         tx: "InstructionTranslator",
@@ -2301,7 +2641,7 @@ def call_method(
             result = self.example_value.size()
             if not has_free_symbols(result):
                 # avoid creating a node in the graph
-                return ConstantVariable.create(int(result))
+                return VariableTracker.build(tx, int(result))
             else:
                 from ..external_utils import untyped_storage_size
                 from .builder import wrap_fx_proxy
@@ -2343,6 +2683,9 @@ def __init__(
         super().__init__(**kwargs)
         self.from_tensor = from_tensor
 
+    def python_type(self) -> type:
+        return int
+
     def reconstruct(self, codegen: "PyCodegen") -> None:
         codegen(self.from_tensor)
         codegen.load_method("data_ptr")
diff --git a/torch/_dynamo/variables/torch.py b/torch/_dynamo/variables/torch.py
index 694edb4c69198..a69112187f60d 100644
--- a/torch/_dynamo/variables/torch.py
+++ b/torch/_dynamo/variables/torch.py
@@ -33,7 +33,7 @@
 import re
 from collections.abc import Callable, Iterable, Sequence
 from contextlib import nullcontext
-from typing import Any, NoReturn, Optional, TYPE_CHECKING, TypeVar, Union
+from typing import Any, NoReturn, TYPE_CHECKING, TypeVar, Union
 from typing_extensions import TypeIs
 
 import torch._C
@@ -42,7 +42,7 @@
 import torch.nn
 import torch.utils._pytree as _pytree
 from torch._C import DispatchKeySet
-from torch._dynamo.variables.constant import CONSTANT_VARIABLE_NONE, ConstantVariable
+from torch._dynamo.variables.constant import ConstantVariable
 from torch._dynamo.variables.streams import StreamVariable
 from torch._dynamo.variables.torch_function import TorchFunctionModeVariable
 from torch._guards import Guard, Source, TracingContext
@@ -58,7 +58,13 @@
     tracable_create_parameter,
 )
 from ..device_interface import get_registered_device_interfaces
-from ..exc import raise_observed_exception, unimplemented
+from ..exc import (
+    raise_observed_exception,
+    raise_type_error,
+    unimplemented,
+    UserError,
+    UserErrorType,
+)
 from ..guards import GuardBuilder, install_guard
 from ..source import (
     AttrSource,
@@ -77,23 +83,24 @@
     proxy_args_kwargs,
     unwrap_if_wrapper,
 )
-from .base import raise_type_error_exc, typestr, VariableTracker
+from .base import typestr, VariableTracker
 from .ctx_manager import (
     AutocastModeVariable,
     ProfilerContextVariable,
     ProfilerRecordFunctionContextVariable,
     TorchFunctionDisableVariable,
 )
-from .dicts import ConstDictVariable
-from .distributed import DistributedVariable, ProcessGroupVariable
+from .distributed import DistributedVariable
 from .functions import bind_args_cached, NestedUserFunctionVariable
-from .lists import ListVariable, NamedTupleVariable, TupleVariable
+from .lists import ListVariable, TupleVariable
+from .script_object import TorchScriptObjectVariable
 from .torch_function import (
     can_dispatch_torch_function,
     dispatch_torch_function,
     TensorWithTFOverrideVariable,
     TorchFunctionModeStackVariable,
 )
+from .user_defined import UserDefinedTupleVariable
 
 
 try:
@@ -109,7 +116,7 @@
 
 if TYPE_CHECKING:
     from torch._dynamo.symbolic_convert import InstructionTranslator
-    from torch._library.opaque_object import OpaqueType
+    from torch._opaque_base import OpaqueBase
     from torch.utils._pytree import TreeSpec
 
 
@@ -166,6 +173,7 @@
     torch.cuda.is_initialized,
     torch.xpu.current_device,
     torch.xpu.is_initialized,
+    torch.__future__.get_overwrite_module_params_on_conversion,
 ]
 
 constant_fold_functions = [
@@ -206,7 +214,7 @@
 
 
 @functools.cache
-def tracing_state_functions() -> dict[Callable[[], Any], Optional[bool]]:
+def tracing_state_functions() -> dict[Callable[[], Any], bool | None]:
     # Defined as a function to avoid circular import like torch.onnx
     return {
         torch.jit.is_scripting: False,
@@ -246,7 +254,7 @@ def _check_for_gradient_edge(var: VariableTracker, arg_name: str) -> None:
     """
     from .lists import BaseListVariable
 
-    if isinstance(var, NamedTupleVariable) and var.tuple_cls is GradientEdge:
+    if isinstance(var, UserDefinedTupleVariable) and type(var.value) is GradientEdge:
         # Try to get source info for context
         source_info = var.source.name if var.source else None
         context = f"GradientEdge in {arg_name}"
@@ -279,7 +287,7 @@ def _collect_all_grad_fns(tensor: torch.Tensor) -> set[torch.autograd.graph.Node
 
     grad_fns: set[torch.autograd.graph.Node] = set()
 
-    plain_tensors: list[torch.SymInt | torch.Tensor | int | OpaqueType] = []
+    plain_tensors: list[torch.SymInt | torch.Tensor | int | OpaqueBase] = []
     # Get all plain tensors (handles nested subclasses)
     if is_traceable_wrapper_subclass(tensor):
         get_plain_tensors(tensor, out=plain_tensors)
@@ -302,20 +310,44 @@ def _collect_all_grad_fns(tensor: torch.Tensor) -> set[torch.autograd.graph.Node
 
 def _collect_tensors_with_sources(
     var: VariableTracker,
-) -> list[tuple[torch.Tensor, Optional[str]]]:
+) -> list[tuple[torch.Tensor, str | None]]:
     """Extract (fake_tensor, source_name) pairs from a VariableTracker.
 
     Used by handle_autograd_grad to collect tensors from the outputs and inputs
     arguments for grad_fn reachability analysis.
     """
+    from torch.utils._python_dispatch import is_traceable_wrapper_subclass
+
     from .lazy import LazyVariableTracker
     from .lists import BaseListVariable
     from .tensor import TensorVariable
 
-    results: list[tuple[torch.Tensor, Optional[str]]] = []
+    results: list[tuple[torch.Tensor, str | None]] = []
     if isinstance(var, TensorVariable):
         fake_tensor = var.as_proxy().node.meta.get("example_value")
-        assert isinstance(fake_tensor, torch._subclasses.fake_tensor.FakeTensor)
+        assert isinstance(fake_tensor, torch.Tensor)
+        if isinstance(fake_tensor, torch._subclasses.fake_tensor.FakeTensor):
+            pass
+        elif is_traceable_wrapper_subclass(fake_tensor):
+            # For tensor subclasses (e.g. DTensor), verify the inner tensors
+            # are FakeTensors but keep the original subclass for grad_fn
+            # reachability analysis.
+            plain: list[object] = []
+            torch._subclasses.fake_tensor.get_plain_tensors(
+                fake_tensor,  # pyrefly: ignore[bad-argument-type]
+                out=plain,  # pyrefly: ignore[bad-argument-type]
+            )
+            assert all(
+                isinstance(t, torch._subclasses.fake_tensor.FakeTensor)
+                for t in plain
+                if isinstance(t, torch.Tensor)
+            ), (
+                f"Expected all plain tensors to be FakeTensors, got {[type(t) for t in plain]}"
+            )
+        else:
+            raise AssertionError(
+                f"Expected FakeTensor or subclass, got {type(fake_tensor)}"
+            )
         source_name = var.source.name if var.source else None
         results.append((fake_tensor, source_name))
     elif isinstance(var, LazyVariableTracker):
@@ -339,20 +371,57 @@ def _collect_tensors_with_sources(
     return results
 
 
+def _collect_placeholder_nodes(var: "VariableTracker") -> list[torch.fx.Node]:
+    """Recursively collect FX placeholder nodes from a VariableTracker.
+
+    The returned placeholder nodes carry grapharg.example (real tensor) and
+    example_value (FakeTensor) metadata — comparing these reveals lost
+    autograd linkage (e.g., grad_fn dropped during tracing).
+    See NOTE [Detecting lost autograd linkage in closure-captured tensors].
+    """
+    from .lazy import LazyVariableTracker
+    from .lists import BaseListVariable
+    from .tensor import TensorVariable
+
+    result: list[torch.fx.Node] = []
+    if isinstance(var, TensorVariable):
+        node = var.as_proxy().node
+        if node.op == "placeholder":
+            result.append(node)
+    elif isinstance(var, LazyVariableTracker):
+        result.extend(_collect_placeholder_nodes(var.realize()))
+    elif isinstance(var, BaseListVariable):
+        for item in var.items:
+            result.extend(_collect_placeholder_nodes(item))
+    else:
+        unimplemented(
+            gb_type="_autograd_grad with unsupported argument type",
+            context=f"got {type(var).__name__}",
+            explanation=(
+                f"_autograd_grad() received an argument of type {type(var).__name__} "
+                "which is not supported. Expected tensor or sequence of tensors."
+            ),
+            hints=[
+                "Ensure outputs and inputs arguments are tensors or sequences of tensors.",
+            ],
+        )
+    return result
+
+
 @functools.cache
 def get_overridable_functions() -> set[Callable[..., Any]]:
     from itertools import chain
 
     from torch.overrides import get_overridable_functions as get_overridable_functions_
+    from torch.utils._device import _device_constructors
 
     funcs = set(chain.from_iterable(get_overridable_functions_().values()))
+    funcs.update(_device_constructors())
     more: set[Callable[..., Any]] = {
-        torch.ones,
         torch.ones_like,
-        torch.zeros,
         torch.zeros_like,
-        torch.empty,
-        torch.full,
+        torch.empty_like,
+        torch.full_like,
     }
     funcs.update(more)
     return funcs
@@ -402,11 +471,14 @@ def as_proxy(self) -> Any:
     def as_python_constant(self) -> Any:
         return self.value
 
+    def get_real_python_backed_value(self) -> Any:
+        return self.value
+
     def call_obj_hasattr(
         self, tx: "InstructionTranslator", name: str
     ) -> ConstantVariable:
         result = hasattr(self.value, name)
-        return variables.ConstantVariable.create(result)
+        return VariableTracker.build(tx, result)
 
     def can_constant_fold_through(self) -> bool:
         if self.value in constant_fold_functions:
@@ -594,7 +666,6 @@ def call_function(
             )
         elif self.value is torch.nn.attention.sdpa_kernel.__wrapped__:  # type: ignore[attr-defined]
             name_to_arg_map = bind_args_cached(
-                # pyrefly: ignore[bad-argument-type]
                 self.value,
                 tx,
                 self.source,
@@ -647,6 +718,7 @@ def get_function(self) -> Callable[..., Any]:
     def _get_handlers() -> dict[Callable[..., Any], Callable[..., Any]]:
         """Build a dict from function -> method to handle it so that we are O(1)
         in terms of the number of function with special handling."""
+        # pyrefly: ignore [implicit-any]
         handlers = {}
 
         def register(
@@ -693,7 +765,11 @@ def handle_tracing_state_functions(
                 torch._dynamo.eval_frame._is_in_optimized_module,
             ):
                 tx.mark_inconsistent_side_effects()
-            return ConstantVariable.create(tracing_state_functions()[self.value])
+            # Read dynamically: the cached dict always has True, but
+            # is_exporting() should only be True during torch.export.
+            if self.value is torch.compiler.is_exporting:
+                return VariableTracker.build(tx, torch.compiler._is_exporting_flag)
+            return VariableTracker.build(tx, tracing_state_functions()[self.value])
 
         @register(*dispatch_key_set_functions)
         def handle_dispatch_key_set_functions(
@@ -823,9 +899,9 @@ def handle_is_floating_point(
             input_arg = input
             if input_arg.is_tensor() and input_arg.dtype is not None:
                 if self.value is torch.is_floating_point:
-                    return ConstantVariable.create(input_arg.dtype.is_floating_point)
+                    return VariableTracker.build(tx, input_arg.dtype.is_floating_point)
                 elif self.value is torch.is_complex:
-                    return ConstantVariable.create(input_arg.dtype.is_complex)
+                    return VariableTracker.build(tx, input_arg.dtype.is_complex)
                 else:
                     raise AssertionError(f"calling {self.value}")
             return None
@@ -835,7 +911,7 @@ def handle_numel(
             self, tx: "InstructionTranslator", input: Any
         ) -> VariableTracker | None:
             if input.is_tensor() and input.valid_size():
-                return ConstantVariable.create(product(input.size))
+                return VariableTracker.build(tx, product(input.size))
             elif input.is_tensor():
                 # Workaround dynamic shapes issue
                 return input.call_method(tx, "numel", [], {})
@@ -910,7 +986,7 @@ def handle_is_grad_enabled(
             self, tx: "InstructionTranslator"
         ) -> ConstantVariable:
             install_guard(GradModeVariable._guards_singleton)
-            return ConstantVariable.create(torch.is_grad_enabled())
+            return VariableTracker.build(tx, torch.is_grad_enabled())
 
         @register(torch.use_deterministic_algorithms)
         def handle_use_deterministic_algorithms(
@@ -933,7 +1009,101 @@ def handle_use_deterministic_algorithms(
                 "call_function", torch._C._set_deterministic_algorithms, (value,), {}
             )
             torch._C._set_deterministic_algorithms(value)
-            return CONSTANT_VARIABLE_NONE
+            return ConstantVariable.create(None)
+
+        @register(torch.autocast_increment_nesting)
+        def handle_autocast_increment_nesting(
+            self, tx: "InstructionTranslator"
+        ) -> VariableTracker:
+            tx.output.create_node(
+                "call_function", torch.autocast_increment_nesting, (), {}
+            )
+            prev = torch.autocast_increment_nesting()
+            tx.output.add_cleanup_hook(lambda: torch.autocast_decrement_nesting())
+            return VariableTracker.build(tx, prev)
+
+        @register(torch.autocast_decrement_nesting)
+        def handle_autocast_decrement_nesting(
+            self, tx: "InstructionTranslator"
+        ) -> VariableTracker:
+            tx.output.create_node(
+                "call_function", torch.autocast_decrement_nesting, (), {}
+            )
+            prev = torch.autocast_decrement_nesting()
+            tx.output.add_cleanup_hook(lambda: torch.autocast_increment_nesting())
+            return VariableTracker.build(tx, prev)
+
+        @register(torch.set_autocast_enabled)
+        def handle_set_autocast_enabled(
+            self,
+            tx: "InstructionTranslator",
+            device_type: VariableTracker,
+            enabled: VariableTracker,
+        ) -> VariableTracker:
+            tx.output.create_node(
+                "call_function",
+                torch.set_autocast_enabled,
+                (device_type.as_proxy(), enabled.as_proxy()),
+            )
+            dev_py_const = device_type.as_python_constant()
+            prev = torch.is_autocast_enabled(dev_py_const)
+            torch.set_autocast_enabled(dev_py_const, enabled.as_python_constant())
+            tx.output.add_cleanup_hook(
+                lambda: torch.set_autocast_enabled(dev_py_const, prev)
+            )
+            return ConstantVariable.create(None)
+
+        @register(torch.set_autocast_cache_enabled)
+        def handle_set_autocast_cache_enabled(
+            self, tx: "InstructionTranslator", enabled: VariableTracker
+        ) -> VariableTracker:
+            tx.output.create_node(
+                "call_function", torch.set_autocast_cache_enabled, (enabled.as_proxy(),)
+            )
+            prev = torch.is_autocast_cache_enabled()
+            torch.set_autocast_cache_enabled(enabled.as_python_constant())
+            tx.output.add_cleanup_hook(lambda: torch.set_autocast_cache_enabled(prev))
+            return ConstantVariable.create(None)
+
+        @register(torch._C._functorch._grad_increment_nesting)
+        def handle_grad_increment_nesting(
+            self, tx: "InstructionTranslator"
+        ) -> VariableTracker:
+            tx.output.create_node(
+                "call_function", torch._C._functorch._grad_increment_nesting
+            )
+            level = torch._C._functorch._grad_increment_nesting()
+            tx.output.add_cleanup_hook(torch._C._functorch._grad_decrement_nesting)
+            return VariableTracker.build(tx, level)
+
+        @register(torch._C._functorch._grad_decrement_nesting)
+        def handle_grad_decrement_nesting(
+            self, tx: "InstructionTranslator"
+        ) -> VariableTracker:
+            tx.output.create_node(
+                "call_function", torch._C._functorch._grad_decrement_nesting
+            )
+            level = torch._C._functorch._grad_decrement_nesting()
+            tx.output.add_cleanup_hook(torch._C._functorch._grad_increment_nesting)
+            return VariableTracker.build(tx, level)
+
+        @register(torch._C._functorch.set_inplace_requires_grad_allowed)
+        def handle_set_inplace_requires_grad_allowed(
+            self, tx: "InstructionTranslator", allowed: VariableTracker
+        ) -> VariableTracker:
+            tx.output.create_node(
+                "call_function",
+                torch._C._functorch.set_inplace_requires_grad_allowed,
+                (allowed.as_proxy(),),
+            )
+            prev = torch._C._functorch.get_inplace_requires_grad_allowed()
+            torch._C._functorch.set_inplace_requires_grad_allowed(
+                allowed.as_python_constant()
+            )
+            tx.output.add_cleanup_hook(
+                lambda: torch._C._functorch.set_inplace_requires_grad_allowed(prev)
+            )
+            return ConstantVariable.create(None)
 
         @register(torch.are_deterministic_algorithms_enabled)
         def handle_are_deterministic_algorithms_enabled(
@@ -944,7 +1114,9 @@ def handle_are_deterministic_algorithms_enabled(
                 GuardBuilder.DETERMINISTIC_ALGORITHMS,  # type: ignore[arg-type]
             )
             install_guard(guard)
-            return ConstantVariable.create(torch.are_deterministic_algorithms_enabled())
+            return VariableTracker.build(
+                tx, torch.are_deterministic_algorithms_enabled()
+            )
 
         @register(torch._C._is_torch_function_enabled)
         def handle_is_torch_function_enabled(
@@ -953,8 +1125,8 @@ def handle_is_torch_function_enabled(
             install_guard(TorchFunctionDisableVariable._guards_singleton)
             # see comment on SymbolicTorchFunctionState class as to why
             # this is not a bug
-            return ConstantVariable.create(
-                tx.symbolic_torch_function_state.torch_function_subclass_enabled
+            return VariableTracker.build(
+                tx, tx.symbolic_torch_function_state.torch_function_subclass_enabled
             )
 
         @register(torch._C._is_torch_function_all_disabled)
@@ -962,8 +1134,8 @@ def handle_is_torch_function_all_disabled(
             self, tx: "InstructionTranslator"
         ) -> ConstantVariable:
             install_guard(TorchFunctionDisableVariable._guards_singleton)
-            return ConstantVariable.create(
-                not tx.symbolic_torch_function_state.torch_function_mode_enabled
+            return VariableTracker.build(
+                tx, not tx.symbolic_torch_function_state.torch_function_mode_enabled
             )
 
         @register(torch._C._is_torch_function_mode_enabled)
@@ -993,9 +1165,7 @@ def handle_has_torch_function(
                 if len(args) == 1 and isinstance(args[0], TupleVariable)
                 else args
             )
-            return ConstantVariable.create(
-                any(has_torch_function(x) for x in elems),
-            )
+            return VariableTracker.build(tx, any(has_torch_function(x) for x in elems))
 
         @register(
             *dict.fromkeys(  # remove duplicates
@@ -1065,8 +1235,8 @@ def handle_cudnn_is_acceptable(
                 "Expect input to cudnn.is_acceptable to be a tensor"
             )
             tensor_inp = torch.tensor(0, dtype=tensor.dtype, device=tensor.device)
-            return ConstantVariable.create(
-                torch.backends.cudnn.is_acceptable(tensor_inp)
+            return VariableTracker.build(
+                tx, torch.backends.cudnn.is_acceptable(tensor_inp)
             )
 
         @register(torch.utils.hooks.BackwardHook)
@@ -1150,9 +1320,9 @@ def handle_inplace_foreach_lerp_scalar(
             *args: VariableTracker,
             **kwargs: VariableTracker,
         ) -> VariableTracker | None:
-            if not config.enable_dynamo_decompositions:
-                return None
-
+            # Decompose via addcmul_ so tensor weights (e.g. 0-dim tensor
+            # from tensor betas in Adam) stay in tensor arguments instead of
+            # hitting float() in the native lerp_scalar lowering.
             if len(args) == 3 and not isinstance(args[2], ListVariable) and not kwargs:
                 return tx.inline_user_function_return(
                     VariableTracker.build(tx, polyfills.foreach_lerp_inplace),
@@ -1205,7 +1375,7 @@ def handle_assert(
                 isinstance(condition, variables.SymNodeVariable)
                 and condition.evaluate_expr()
             ):
-                return CONSTANT_VARIABLE_NONE
+                return ConstantVariable.create(None)
             return None
 
         @register(SDPAParams)
@@ -1232,8 +1402,9 @@ def handle_sdpa_params(
                 _rank_not_in_group,
                 _resolve_group_name_by_ranks_and_tag,
                 get_process_group_ranks,
+                get_rank,
+                get_world_size,
             )
-            from torch.distributed.tensor import DTensor
 
             @register(
                 _get_group_size_by_name,
@@ -1241,102 +1412,60 @@ def handle_sdpa_params(
                 _rank_not_in_group,
                 get_process_group_ranks,
                 _resolve_group_name_by_ranks_and_tag,
+                get_rank,
+                get_world_size,
             )
             def handle_constant_processgroup_functions(
-                self, tx: "InstructionTranslator", *args: VariableTracker
+                self,
+                tx: "InstructionTranslator",
+                *args: VariableTracker,
+                **kwargs: VariableTracker,
             ) -> VariableTracker:
-                # because the input is a "ProcessGroupVariable", we'll be guarding on its
-                # ID_MATCH based on how it was constructed.
-
                 # We desugar it at trace-time into ranks by directly calling util
                 # bake the result into the trace
-                if len(args) == 1:
+                if len(args) == 0 and len(kwargs) == 0:
+                    # get_rank() or get_world_size() with no args (uses default group)
+                    pass
+                elif len(args) == 1 and len(kwargs) == 0:
                     # group or group name
-                    assert (
-                        isinstance(args[0], ProcessGroupVariable)
-                        or args[0].is_python_constant()
+                    assert args[0].is_python_constant() or (
+                        isinstance(args[0], TorchScriptObjectVariable)
+                        and args[  # pyrefly: ignore[missing-attribute]
+                            0
+                        ].value.script_class_name  # pyrefly: ignore[missing-attribute]
+                        == "torch.distributed.distributed_c10d.ProcessGroup"
                     )
-                elif len(args) == 2:
+                elif len(args) == 2 and len(kwargs) == 0:
                     # ranks + tag
                     assert (
                         isinstance(args[0], ListVariable)
                         and args[1].is_python_constant()
                     )
+                elif len(args) == 0 and len(kwargs) > 0:
+                    # All keyword arguments (e.g., get_world_size(group=...))
+                    pass
                 else:
                     raise AssertionError(
-                        f"Invalid group value ({args}) for constant pg "
+                        f"Invalid group value ({args}, {kwargs}) for constant pg "
                         f"function {self.value}"
                     )
-                args_as_value = [arg.as_python_constant() for arg in args]
-                invocation_result = self.value(*args_as_value)
+
+                def get_arg_value(arg: VariableTracker) -> Any:
+                    # TorchScriptObjectVariable for ProcessGroup doesn't support
+                    # as_python_constant(), so extract real_obj directly
+                    if isinstance(arg, TorchScriptObjectVariable):
+                        return arg.value.real_obj  # pyrefly: ignore[missing-attribute]
+                    return arg.as_python_constant()
+
+                args_as_value = [get_arg_value(arg) for arg in args]
+                kwargs_as_value = {k: get_arg_value(v) for k, v in kwargs.items()}
+                invocation_result = self.value(*args_as_value, **kwargs_as_value)
 
                 # Note - while we *could* cook up sources around invocations, like a FunctionSource
                 # the space of invoking functions in the middle of the guard chain is very iffy. As such,
                 # guard propagation via options is the best we can do.
                 return VariableTracker.build(tx, invocation_result)
 
-            @register(DTensor.from_local)
-            def handle_from_local(
-                self,
-                tx: "InstructionTranslator",
-                *args: VariableTracker,
-                **kwargs: VariableTracker,
-            ) -> VariableTracker:
-                from .builder import SourcelessBuilder
-
-                # rewrite non-primitive args/kwargs to be included in the on-the-fly prim function
-                # and rewrite args to have only proxyable args, then insert call_function
-                placements_vt = kwargs.get("placements")
-
-                if placements_vt is None and len(args) >= 3:
-                    placements_vt = args[2]
-
-                if placements_vt is None:
-                    placements_vt = CONSTANT_VARIABLE_NONE
-                elif isinstance(placements_vt, variables.UserDefinedObjectVariable):
-                    placements_vt = SourcelessBuilder.create(tx, tuple).call_function(
-                        tx, [placements_vt], {}
-                    )
-
-                new_args = list(args)
-                if len(new_args) >= 3:
-                    new_args[2] = placements_vt
-                elif kwargs.get("placements") is not None:
-                    kwargs["placements"] = placements_vt
-
-                args_as_value = [x.as_python_constant() for x in new_args[1:]]
-                kwargs_as_value = {
-                    k: v.as_python_constant()
-                    for k, v in kwargs.items()
-                    if k not in ["shape", "stride"]
-                }
-
-                kwargs_to_be_proxied = {
-                    k: kwargs[k] for k in ["shape", "stride"] if k in kwargs
-                }
-
-                def fn_with_prim_types(
-                    x: Any, shape: Any | None = None, stride: Any | None = None
-                ) -> Any:
-                    return self.value(
-                        x, *args_as_value, **kwargs_as_value, shape=shape, stride=stride
-                    )
-
-                # attach the same function name for better debugging
-                fn_with_prim_types.__name__ = "prim " + self.value.__name__
-
-                return wrap_fx_proxy(
-                    tx=tx,
-                    proxy=tx.output.create_proxy(
-                        "call_function",
-                        fn_with_prim_types,
-                        *proxy_args_kwargs(
-                            [args[0]],
-                            kwargs_to_be_proxied,
-                        ),
-                    ),
-                )
-
         @register(torch.nested.nested_tensor)
         def handle_nested_tensor(
             self,
@@ -1391,19 +1520,38 @@ def handle_one_hot(
                 )
             return None
 
-        @register(torch.fx.experimental.symbolic_shapes.size_hint)
-        def handle_size_hint(
+        @register(torch.fx.experimental.symbolic_shapes.guarding_hint_or_throw)
+        def handle_guarding_hint_or_throw(
             self,
             tx: "InstructionTranslator",
             expr: VariableTracker,
-            fallback: Optional[VariableTracker] = None,
+        ) -> VariableTracker | None:
+            if isinstance(expr, SymNodeVariable):
+                return VariableTracker.build(
+                    tx,
+                    torch.fx.experimental.symbolic_shapes.guarding_hint_or_throw(
+                        expr.sym_num
+                    ),
+                )
+            elif expr.is_python_constant():
+                return expr
+            else:
+                return None
+
+        @register(torch.fx.experimental.symbolic_shapes.optimization_hint)
+        def handle_optimization_hint(
+            self,
+            tx: "InstructionTranslator",
+            expr: VariableTracker,
+            fallback: VariableTracker | None = None,
         ) -> VariableTracker | None:
             fallback_int = fallback.as_python_constant() if fallback else None
             if isinstance(expr, SymNodeVariable):
-                return variables.ConstantVariable.create(
-                    torch.fx.experimental.symbolic_shapes.size_hint(
+                return VariableTracker.build(
+                    tx,
+                    torch.fx.experimental.symbolic_shapes.optimization_hint(
                         expr.sym_num, fallback_int
-                    )
+                    ),
                 )
             elif expr.is_python_constant():
                 return expr
@@ -1417,10 +1565,11 @@ def handle_guard_size_oblivious(
             if isinstance(expr, SymNodeVariable):
                 # TODO: this probably should be folded somewhere else but I'm not sure where
                 # TODO: some of the other symbolic_shapes special tools can also get this treatment too
-                return variables.ConstantVariable.create(
+                return VariableTracker.build(
+                    tx,
                     torch.fx.experimental.symbolic_shapes.guard_size_oblivious(  # type: ignore[deprecated]
                         expr.sym_num
-                    )
+                    ),
                 )
             elif expr.is_python_constant():
                 return expr
@@ -1434,8 +1583,9 @@ def handle_guard_or_true(
             if isinstance(expr, SymNodeVariable):
                 # TODO: this probably should be folded somewhere else but I'm not sure where
                 # TODO: some of the other symbolic_shapes special tools can also get this treatment too
-                return variables.ConstantVariable.create(
-                    torch.fx.experimental.symbolic_shapes.guard_or_true(expr.sym_num)
+                return VariableTracker.build(
+                    tx,
+                    torch.fx.experimental.symbolic_shapes.guard_or_true(expr.sym_num),
                 )
             elif expr.is_python_constant():
                 return expr
@@ -1449,8 +1599,9 @@ def handle_guard_or_false(
             if isinstance(expr, SymNodeVariable):
                 # TODO: this probably should be folded somewhere else but I'm not sure where
                 # TODO: some of the other symbolic_shapes special tools can also get this treatment too
-                return variables.ConstantVariable.create(
-                    torch.fx.experimental.symbolic_shapes.guard_or_false(expr.sym_num)
+                return VariableTracker.build(
+                    tx,
+                    torch.fx.experimental.symbolic_shapes.guard_or_false(expr.sym_num),
                 )
             elif expr.is_python_constant():
                 return expr
@@ -1462,16 +1613,33 @@ def handle_statically_known_false(
             self, tx: "InstructionTranslator", expr: VariableTracker
         ) -> VariableTracker | None:
             if isinstance(expr, SymNodeVariable):
-                return variables.ConstantVariable.create(
+                return VariableTracker.build(
+                    tx,
                     torch.fx.experimental.symbolic_shapes.statically_known_false(
                         expr.sym_num
-                    )
+                    ),
                 )
             elif expr.is_python_constant():
                 return expr
             else:
                 return None
 
+        @register(torch.fx.experimental.symbolic_shapes.has_free_unbacked_symbols)
+        def handle_has_free_unbacked_symbols(
+            self, tx: "InstructionTranslator", x: VariableTracker
+        ) -> VariableTracker | None:
+            from .tensor import TensorVariable
+
+            if isinstance(x, TensorVariable):
+                example_value = x.as_proxy().node.meta.get("example_value")
+                if example_value is not None:
+                    return variables.ConstantVariable.create(
+                        torch.fx.experimental.symbolic_shapes.has_free_unbacked_symbols(
+                            example_value
+                        )
+                    )
+            return None
+
         @register(torch.fx.experimental.symbolic_shapes.guard_scalar)
         def guard_scalar(
             self, tx: "InstructionTranslator", expr: VariableTracker
@@ -1487,9 +1655,10 @@ def guard_scalar(
                     explanation="Expected `expr` to be a symbolic variable or constant.",
                     hints=[],
                 )
-            return variables.ConstantVariable.create(
+            return VariableTracker.build(
+                tx,
                 # pyrefly: ignore [bad-argument-type, unbound-name]
-                torch.fx.experimental.symbolic_shapes.guard_scalar(val)
+                torch.fx.experimental.symbolic_shapes.guard_scalar(val),
             )
 
         @register(torch.fx.experimental.symbolic_shapes.statically_known_true)
@@ -1497,10 +1666,11 @@ def handle_statically_known_true(
             self, tx: "InstructionTranslator", expr: VariableTracker
         ) -> VariableTracker | None:
             if isinstance(expr, SymNodeVariable):
-                return variables.ConstantVariable.create(
+                return VariableTracker.build(
+                    tx,
                     torch.fx.experimental.symbolic_shapes.statically_known_true(
                         expr.sym_num
-                    )
+                    ),
                 )
             elif expr.is_python_constant():
                 return expr
@@ -1546,8 +1716,8 @@ def handle_has_static_value(
             else:
                 return None
 
-            return variables.ConstantVariable.create(
-                torch.fx.experimental.symbolic_shapes.has_static_value(val)
+            return VariableTracker.build(
+                tx, torch.fx.experimental.symbolic_shapes.has_static_value(val)
             )
 
         @register(torch._C._autograd._unsafe_set_version_counter)
@@ -1656,14 +1826,14 @@ def handle_push_torch_function(
             **kwargs: VariableTracker,
         ) -> VariableTracker:
             if len(args) != 1 or kwargs:
-                raise_type_error_exc(
+                raise_type_error(
                     tx,
                     f"push_torch_function takes exactly one argument ({len(args)} given)",
                 )
             TorchFunctionModeStackVariable.register_mutation(tx)
             # type: ignore[arg-type]
             tx.symbolic_torch_function_state.push_torch_function_mode(args[0])
-            return CONSTANT_VARIABLE_NONE
+            return ConstantVariable.create(None)
 
         @register(torch._C._len_torch_function_stack)
         def handle_len_torch_function(
@@ -1673,9 +1843,9 @@ def handle_len_torch_function(
             **kwargs: VariableTracker,
         ) -> VariableTracker:
             if args or kwargs:
-                raise_type_error_exc(tx, "len_torch_function_stack takes no arguments")
-            return ConstantVariable.create(
-                len(tx.symbolic_torch_function_state.mode_stack)
+                raise_type_error(tx, "len_torch_function_stack takes no arguments")
+            return VariableTracker.build(
+                tx, len(tx.symbolic_torch_function_state.mode_stack)
             )
 
         @register(torch._C._get_function_stack_at)
@@ -1686,7 +1856,7 @@ def handle_get_stack_at(
             **kwargs: VariableTracker,
         ) -> TorchFunctionModeVariable:
             if len(args) != 1 or kwargs:
-                raise_type_error_exc(
+                raise_type_error(
                     tx,
                     f"get_function_stack_at takes exactly one argument ({len(args)} given)",
                 )
@@ -1728,17 +1898,14 @@ def handle_get_device_module(
                 )
 
             # need to guard only on no-arg get_device_module
-            # pyrefly: ignore [unbound-name]
             if device is None:
                 source = CallFunctionNoArgsSource(self.source)
                 install_guard(source.make_guard(GuardBuilder.ID_MATCH))
             # assumes `module` is in the form `torch.xyz`
             new_source = AttrSource(
                 ImportSource("torch"),
-                # pyrefly: ignore [unbound-name]
                 module.__name__.rsplit(".", maxsplit=1)[-1],
             )
-            # pyrefly: ignore [unbound-name]
             return VariableTracker.build(tx, module, new_source)
 
         @register(torch.accelerator.current_stream, torch.cuda.current_stream)
@@ -1748,6 +1915,8 @@ def handle_current_stream(
             *args: VariableTracker,
             **kwargs: VariableTracker,
         ) -> StreamVariable:
+            from .streams import CudaStreamVariable
+
             if len(args) + len(kwargs) > 1 or (kwargs and "device" not in kwargs):
                 unimplemented(
                     gb_type="unsupported arguments to torch.accelerator.current_stream",
@@ -1765,7 +1934,17 @@ def handle_current_stream(
                 else:
                     device = None
 
-                return tx.symbolic_stream_state.cur_stream(device)
+                stream_var = tx.symbolic_stream_state.cur_stream(device)
+                if self.value is torch.cuda.current_stream and not isinstance(
+                    stream_var, CudaStreamVariable
+                ):
+                    stream_var = CudaStreamVariable(
+                        stream_var.proxy,
+                        stream_var.value,
+                        stream_var.user_object_index,
+                        source=stream_var.source,
+                    )
+                return stream_var
             except Exception as e:
                 unimplemented(
                     gb_type="bad device argument to torch.accelerator.current_stream",
@@ -1775,6 +1954,53 @@ def handle_current_stream(
                     from_exc=e,
                 )
 
+        _synchronize_fn_to_device_type = {
+            torch.cuda.synchronize: "cuda",
+            torch.xpu.synchronize: "xpu",
+            torch.mps.synchronize: "mps",
+            torch.cpu.synchronize: "cpu",
+        }
+
+        @register(
+            torch.accelerator.synchronize,
+            torch.cuda.synchronize,
+            torch.xpu.synchronize,
+            torch.mps.synchronize,
+            torch.cpu.synchronize,
+        )
+        def handle_synchronize(
+            self,
+            tx: "InstructionTranslator",
+            *args: VariableTracker,
+            **kwargs: VariableTracker,
+        ) -> VariableTracker:
+            device = None
+            if kwargs and "device" in kwargs:
+                device = torch.device(kwargs["device"].as_python_constant())
+            elif args:
+                device = torch.device(args[0].as_python_constant())
+
+            if device is None:
+                device_type = _synchronize_fn_to_device_type.get(self.value)
+                if device_type is None:
+                    # torch.accelerator.synchronize with no args
+                    accelerator = torch.accelerator.current_accelerator()
+                    assert accelerator is not None
+                    device_type = accelerator.type
+                device = torch.device(device_type)
+
+            # CPU synchronize is a no-op, skip emitting the op
+            if device.type == "cpu":
+                return ConstantVariable.create(None)
+
+            tx.output.create_proxy(
+                "call_function",
+                torch.ops.streams.synchronize_device,
+                (device.type, device.index or 0),
+                {},
+            )
+            return ConstantVariable.create(None)
+
         @register(torch.set_default_device)
         def handle_set_default_device(
             self,
@@ -1795,7 +2021,40 @@ def handle_set_default_device(
             else:
                 TorchFunctionModeStackVariable.register_device_context_insertion(tx)
 
-            return CONSTANT_VARIABLE_NONE
+            return ConstantVariable.create(None)
+
+        from torch._prims_common import elementwise_dtypes
+
+        @register(elementwise_dtypes)
+        def handle_elementwise_dtypes(
+            self,
+            tx: "InstructionTranslator",
+            *args: VariableTracker,
+            **kwargs: VariableTracker,
+        ) -> VariableTracker:
+            from .builder import SourcelessBuilder
+
+            type_promotion_kind = kwargs["type_promotion_kind"].as_python_constant()
+            real_args = []
+            for arg in args:
+                if isinstance(arg, TensorVariable):
+                    real_args.append(arg.as_proxy().node.meta["example_value"])
+                elif arg.is_python_constant():
+                    real_args.append(arg.as_python_constant())
+                else:
+                    unimplemented(
+                        gb_type="elementwise_dtypes unsupported arg type",
+                        context=str(arg),
+                        explanation=(
+                            "elementwise_dtypes requires tensor or constant arguments, "
+                            f"got {type(arg).__name__}"
+                        ),
+                        hints=[*graph_break_hints.SUPPORTABLE],
+                    )
+            result = elementwise_dtypes(
+                *real_args, type_promotion_kind=type_promotion_kind
+            )
+            return SourcelessBuilder.create(tx, result)
 
         @register(torch._check)
         def handle_check(
@@ -1854,7 +2113,6 @@ def handle_check(
                             *graph_break_hints.SUPPORTABLE,
                         ],
                     )
-                # pyrefly: ignore[missing-attribute]
                 message_eager = message_vt.get_function()
 
                 message_graph_proxy = tx.output.register_static_attr_and_return_proxy(
@@ -1863,7 +2121,7 @@ def handle_check(
 
             if predicate_vt.is_python_constant():
                 self.value(predicate_vt.as_python_constant(), message_eager)
-                return CONSTANT_VARIABLE_NONE
+                return ConstantVariable.create(None)
 
             predicate_proxy = predicate_vt.as_proxy()
 
@@ -1883,8 +2141,53 @@ def handle_check(
                 ),
             )
 
+        def exchange_device_helper(
+            tx: "InstructionTranslator",
+            args: Sequence[VariableTracker],
+            kwargs: dict[str, VariableTracker],
+            fn: Callable[[int], int | None],
+        ) -> VariableTracker:
+            if len(args) != 1 or kwargs:
+                raise_type_error(
+                    tx,
+                    f"{fn.__name__} takes exactly one argument ({len(args)} given)",
+                )
+            current_device_source = CallFunctionNoArgsSource(
+                AttrSource(AttrSource(ImportSource("torch"), "cuda"), "current_device")
+            )
+            install_guard(current_device_source.make_guard(GuardBuilder.EQUALS_MATCH))
+            arg = args[0].as_python_constant()
+            prev = fn(arg)
+            tx.output.create_node(
+                "call_function",
+                fn,
+                (arg,),
+                {},
+            )
+            tx.output.add_cleanup_hook(lambda: torch.cuda.set_device(prev))
+            return VariableTracker.build(tx, prev)
+
+        @register(torch.cuda._exchange_device)
+        def handle_exchange_device(
+            self,
+            tx: "InstructionTranslator",
+            *args: VariableTracker,
+            **kwargs: VariableTracker,
+        ) -> VariableTracker:
+            return exchange_device_helper(tx, args, kwargs, torch.cuda._exchange_device)
+
+        @register(torch.cuda._maybe_exchange_device)
+        def handle_maybe_exchange_device(
+            self,
+            tx: "InstructionTranslator",
+            *args: VariableTracker,
+            **kwargs: VariableTracker,
+        ) -> VariableTracker:
+            return exchange_device_helper(
+                tx, args, kwargs, torch.cuda._maybe_exchange_device
+            )
+
         @register(torch.autograd.grad)
-        # pyrefly: ignore[implicit-any]
         def handle_autograd_grad(self, tx: "InstructionTranslator", *args, **kwargs):
             """
             Handle torch.autograd.grad() calls within compiled regions.
@@ -1962,17 +2265,21 @@ def fn(x):
             # consumed grad_fns of returned tensors. This gives better compile
             # coverage than failing the entire compile.
             if tx.speculation_log.graph_break_on_autograd_grad:
+                leaked = tx.speculation_log.autograd_grad_leaked_tensors
+                leaked_str = ", ".join(leaked) if leaked else "unknown"
                 unimplemented(
                     gb_type="autograd.grad consumed returned tensor's grad_fn",
-                    context="",
+                    context=f"Leaked output tensors: {leaked_str}",
                     explanation=(
                         "torch.autograd.grad() consumes grad_fns that are needed by tensors "
                         "returned from this compiled function. This would cause 'backward "
-                        "through graph a second time' errors."
+                        "through graph a second time' errors.\n"
+                        f"  The following returned tensors have consumed grad_fns: {leaked_str}"
                     ),
                     hints=[
-                        "If you don't need to backward through the returned tensor, "
-                        "call .detach() before returning: `return loss.detach()`",
+                        f"Detach the problematic tensor(s) before returning: e.g. `{leaked[0]}.detach()`"
+                        if leaked
+                        else "Call .detach() on the tensor before returning.",
                         "If you need to backward through the returned tensor, use retain_graph=True in autograd.grad().",
                     ],
                 )
@@ -2126,14 +2433,78 @@ def fn(x):
                     )
                 tx.output.autograd_grad_consumed_grad_fns.update(non_leaf_consumed)
 
-            return wrap_fx_proxy(
-                tx=tx,
-                proxy=tx.output.create_proxy(
+            with (
+                torch.fx.traceback.preserve_node_meta(),
+                torch.fx.traceback._set_autograd_backward(),
+            ):
+                proxy = tx.output.create_proxy(
                     "call_function",
                     torch.autograd.grad,
                     *proxy_args_kwargs(args, kwargs),
-                ),
-            )
+                )
+            return wrap_fx_proxy(tx=tx, proxy=proxy)
+
+        @register(torch._functorch.eager_transforms._autograd_grad)
+        def handle_functorch_autograd_grad(
+            self,
+            tx: "InstructionTranslator",
+            *args: VariableTracker,
+            **kwargs: VariableTracker,
+        ) -> VariableTracker | None:
+            """Graph-break when closure-captured tensors lose their grad_fn.
+
+            NOTE [Detecting lost autograd linkage in closure-captured tensors]
+
+            Functorch transforms (vjp, grad, jacrev) return closures that capture
+            tensors with grad_fn. When such a closure is compiled separately, those
+            tensors become graph placeholders whose FakeTensors lose grad_fn,
+            causing _autograd_grad to silently return zeros.
+
+            _collect_placeholder_nodes gathers placeholder nodes from the
+            outputs/inputs args. For each, we compare grapharg.example (the real
+            tensor, retains grad_fn) against example_value (FakeTensor, grad_fn
+            is None). A mismatch means autograd linkage was lost, so we
+            graph-break.
+
+            This is a pre-check only: kwargs (retain_graph, create_graph,
+            grad_outputs) don't affect linkage detection and are handled by
+            the default proxy path when this returns None.
+            """
+            outputs_var = args[0] if len(args) >= 1 else None
+            inputs_var = args[1] if len(args) >= 2 else None
+
+            if outputs_var is None or inputs_var is None:
+                return None
+
+            output_placeholder_nodes = _collect_placeholder_nodes(outputs_var)
+            input_placeholder_nodes = _collect_placeholder_nodes(inputs_var)
+
+            if output_placeholder_nodes and input_placeholder_nodes:
+                for node in output_placeholder_nodes:
+                    fake = node.meta.get("example_value")
+                    grapharg = node.meta.get("grapharg")
+                    if (
+                        grapharg is not None
+                        and isinstance(fake, torch.Tensor)
+                        and fake.grad_fn is None
+                    ):
+                        real = grapharg.example
+                        if isinstance(real, torch.Tensor) and real.grad_fn is not None:
+                            unimplemented(
+                                gb_type="_autograd_grad with lost grad_fn linkage",
+                                context="outputs lost autograd linkage during tracing",
+                                explanation=(
+                                    "_autograd_grad() received tensors whose grad_fn "
+                                    "was lost during tracing - this silently produces "
+                                    "zero gradients."
+                                ),
+                                hints=[
+                                    "Compile the full transform instead of the returned "
+                                    "closure: torch.compile(lambda x: torch.func.vjp(f, x))",
+                                    *graph_break_hints.SUPPORTABLE,
+                                ],
+                            )
+            return None
 
         return handlers
 
@@ -2143,7 +2514,7 @@ def call_function(
         args: Sequence[VariableTracker],
         kwargs: "dict[str, VariableTracker]",
     ) -> "VariableTracker":
-        from . import ConstantVariable, SymNodeVariable
+        from . import SymNodeVariable
         from .builder import wrap_fx_proxy
 
         if self.kind == AllowInGraphKind.NONSTRICT_TRACE:
@@ -2165,7 +2536,8 @@ def call_function(
                 install_guard(source.make_guard(GuardBuilder.EQUALS_MATCH))
             # constant fold
             try:
-                return ConstantVariable.create(
+                return VariableTracker.build(
+                    tx,
                     self.as_python_constant()(
                         *[x.as_python_constant() for x in args],
                         **{k: v.as_python_constant() for k, v in kwargs.items()},
@@ -2175,7 +2547,7 @@ def call_function(
                 raise_observed_exception(
                     type(exc),
                     tx,
-                    args=list(map(ConstantVariable.create, exc.args)),
+                    args=list(exc.args),
                 )
 
         if self.is_tensor_method():
@@ -2434,13 +2806,14 @@ def _call_nonstrict_traceable_function(
         args: Sequence[VariableTracker],
         kwargs: "dict[str, VariableTracker]",
     ) -> VariableTracker:
+        from torch._dynamo.utils import _make_inlined
         from torch._higher_order_ops.flat_apply import (
             flat_apply,
-            func_to_graphable,
             is_graphable_type,
             is_valid_output,
             to_graphable,
         )
+        from torch._higher_order_ops.invoke_leaf_function import _LeafCallable
         from torch._subclasses.fake_tensor import fake_tensor_tls
         from torch.utils._pytree import tree_flatten
 
@@ -2453,14 +2826,12 @@ def _call_nonstrict_traceable_function(
         # then tree_flatten them, we just let Dynamo symbolically interpret
         # `tree_flatten((args, kwargs))`. This saves us from having to
         # worry about the reconstruction logic, side effects, and guards.
-        packed_input_vt = TupleVariable.build(
-            tx, (TupleVariable.build(tx, args), ConstDictVariable.build(tx, kwargs))
-        )
-        out_vt = SourcelessBuilder.create(tx, tree_flatten).call_function(  # type: ignore[arg-type]
-            tx, [packed_input_vt], {}
+        args_with_states, kwargs_with_states = self._extract_nn_module_states(
+            tx, args, kwargs
         )
-        assert isinstance(out_vt, TupleVariable) and len(out_vt.items) == 2
-        flat_args_vts, input_spec_vt = out_vt.items
+        flat_args_vts, input_spec_vt = _make_inlined(tx, tree_flatten)(
+            VariableTracker.build(tx, (args_with_states, kwargs_with_states))
+        ).unpack_var_sequence(tx)
         assert isinstance(flat_args_vts, ListVariable)
 
         # Handle the case when the input contains a non-graphable type.
@@ -2551,23 +2922,20 @@ def patched_fn(
                 fake_tensor_tls.allow_non_fake_inputs_override = old_val
             return res
 
-        # `flat_apply` wants a TreeSpec for the function input.
-        _, f_spec = func_to_graphable(patched_fn)
+        f_callable = _LeafCallable(patched_fn)
 
-        # TreeSpec isn't graphable, so we register the function and input
-        # specs as attributes on the graph module.
-        f_spec_proxy = tx.output.register_static_attr_and_return_proxy(
-            f"{fn.__name__}_spec", f_spec
+        f_callable_proxy = tx.output.register_static_attr_and_return_proxy(
+            f"{fn.__name__}_callable", f_callable
         )
         input_spec_proxy = tx.output.register_static_attr_and_return_proxy(
             fn.__name__ + "_input_spec",
             # pyrefly: ignore [unbound-name]
             input_spec,
         )
-        f_spec_proxy.node.type = type(f_spec)
+        f_callable_proxy.node.type = type(f_callable)
         # pyrefly: ignore [unbound-name]
         input_spec_proxy.node.type = type(input_spec)
-        all_args = (f_spec_proxy, input_spec_proxy, *proxified_flat_args)
+        all_args = (f_callable_proxy, input_spec_proxy, *proxified_flat_args)
 
         # 2. Create a proxy call to `flat_apply`, then fake-tensor propagate
         # the call and wrap output into a VariableTracker.
@@ -2653,9 +3021,11 @@ def _extract_nn_module_states(
         """
         import torch.utils._pytree as pytree
         from torch._dynamo.graph_bytecode_inputs import register_user_object
-        from torch._higher_order_ops.invoke_leaf_function import LeafModuleState
+        from torch._dynamo.utils import _make_inlined
+        from torch._higher_order_ops.invoke_leaf_function import (
+            convert_modules_to_states,
+        )
 
-        from .higher_order_ops import _make_inlined
         from .nn_module import NNModuleVariable, UnspecializedNNModuleVariable
 
         def is_module_variable(
@@ -2695,20 +3065,6 @@ def is_module_variable(
             kwargs_var = VariableTracker.build(tx, kwargs)
             return args_var, kwargs_var
 
-        def convert_modules_to_states(
-            values: Any, module_to_index: dict[int, int]
-        ) -> Any:
-            def module_to_state(val: Any) -> Any:
-                if isinstance(val, torch.nn.Module):
-                    return LeafModuleState(
-                        nn_module_index=module_to_index[id(val)],
-                        named_parameters=dict(val.named_parameters()),
-                        named_buffers=dict(val.named_buffers()),
-                    )
-                return val
-
-            return pytree.tree_map(module_to_state, values)
-
         module_to_index_var = VariableTracker.build(tx, module_to_index)
 
         result_var = _make_inlined(tx, convert_modules_to_states)(
@@ -2723,18 +3079,19 @@ def _call_leaf_function(
         kwargs: dict[str, VariableTracker],
     ) -> VariableTracker:
         import torch.utils._pytree as pytree
-        from torch._higher_order_ops.flat_apply import func_to_graphable
+        from torch._dynamo.utils import _make_inlined
         from torch._higher_order_ops.invoke_leaf_function import (
+            _LeafCallable,
             invoke_leaf_function,
-            reconstruct_original_args,
+            make_leaf_function_wrappers,
         )
 
         from .builder import wrap_fx_proxy
-        from .higher_order_ops import _make_inlined
 
         decorated_fn = self.value
         real_impl = decorated_fn._torchdynamo_leaf_real_fn
         fake_impl = decorated_fn._torchdynamo_leaf_fake_fn
+        mutates_args = decorated_fn._torchdynamo_leaf_mutates_args
 
         if fake_impl is None:
             raise ValueError(
@@ -2746,56 +3103,66 @@ def _call_leaf_function(
         args_with_states, kwargs_with_states = self._extract_nn_module_states(
             tx, args, kwargs
         )
+
         flat_args_var, input_spec_var = _make_inlined(tx, pytree.tree_flatten)(
             VariableTracker.build(tx, (args_with_states, kwargs_with_states))
         ).unpack_var_sequence(tx)
         flat_arg_proxies = [
             arg.as_proxy() for arg in flat_args_var.unpack_var_sequence(tx)
         ]
-        input_spec = (
-            input_spec_var.as_python_constant()
-        )  # pyrefly: ignore [unbound-name]
-
-        def wrap_impl_for_leaf_module_state(
-            impl: Callable[..., Any],
-        ) -> Callable[..., Any]:
-            def wrapped_impl(*flat_args: Any) -> Any:
-                # NB: The flat_args contain flattened LeafModuleState objects.
-                # We need to unflatten them with input_spec then convert back to nn.Module
-                # before calling the original impl.
-                #
-                # input_spec is captured from the outer scope
-                with reconstruct_original_args(input_spec, flat_args) as (
-                    args_with_modules,
-                    kwargs_with_modules,
-                ):
-                    return impl(*args_with_modules, **kwargs_with_modules)
+        input_spec = input_spec_var.as_python_constant()
 
-            return wrapped_impl
+        mutated_flat_indices = ""
+        if mutates_args:
+            from torch._higher_order_ops.invoke_leaf_function import (
+                _resolve_mutated_flat_indices,
+            )
 
-        wrapped_real_impl = wrap_impl_for_leaf_module_state(real_impl)
-        wrapped_fake_impl = wrap_impl_for_leaf_module_state(fake_impl)
+            try:
+                mutated_flat_indices = _resolve_mutated_flat_indices(
+                    real_impl, mutates_args, len(flat_arg_proxies), input_spec
+                )
+            except ValueError as e:
+                raise UserError(UserErrorType.INVALID_INPUT, str(e)) from e
+
+        # Single-element mutable list so the wrappers can write back the output
+        # TreeSpec. Read captured_out_spec[0] after the wrappers have been called.
+        captured_out_spec: list[pytree.TreeSpec | None] = [None]
+        wrapped_real_impl, wrapped_fake_impl = make_leaf_function_wrappers(
+            real_impl, fake_impl, captured_out_spec
+        )
 
-        _, real_impl_spec = func_to_graphable(wrapped_real_impl)
-        _, fake_impl_spec = func_to_graphable(wrapped_fake_impl)
+        real_impl_callable = _LeafCallable(wrapped_real_impl)
+        fake_impl_callable = _LeafCallable(wrapped_fake_impl)
 
-        def make_spec_proxy(name: str, spec: Any) -> Any:
+        def make_callable_proxy(name: str, spec: Any) -> Any:
             proxy = tx.output.register_static_attr_and_return_proxy(name, spec)
             proxy.node.type = type(spec)
             return proxy
 
-        real_impl_proxy = make_spec_proxy("real_fn", real_impl_spec)
-        fake_impl_proxy = make_spec_proxy("fake_fn", fake_impl_spec)
+        real_impl_proxy = make_callable_proxy("real_fn", real_impl_callable)
+        fake_impl_proxy = make_callable_proxy("fake_fn", fake_impl_callable)
+        input_spec_proxy = make_callable_proxy("input_spec", input_spec)
 
         invoke_args = (
             real_impl_proxy,
             fake_impl_proxy,
+            input_spec_proxy,
+            mutated_flat_indices,
             *flat_arg_proxies,
         )
         result_proxy = tx.output.create_proxy(
             "call_function", invoke_leaf_function, invoke_args, {}
         )
-        return wrap_fx_proxy(tx, result_proxy)
+
+        flat_output_vt = wrap_fx_proxy(tx, result_proxy)
+
+        assert captured_out_spec[0] is not None, (
+            "Output spec was not captured during fake tensor propagation. "
+            "This should not happen - please report a bug."
+        )
+        out_spec_vt = VariableTracker.build(tx, captured_out_spec[0])
+        return _make_inlined(tx, _pytree.tree_unflatten)(flat_output_vt, out_spec_vt)
 
     def _call_ntuple(
         self,
@@ -2818,7 +3185,8 @@ def handle_ntuple(value: VariableTracker) -> VariableTracker:
                 )
             elif value.is_python_constant():
                 # constant prop through it
-                return variables.ConstantVariable.create(
+                return VariableTracker.build(
+                    tx,
                     torch.nn.modules.utils._ntuple(count)(value.as_python_constant()),
                 )
             else:
@@ -2881,7 +3249,6 @@ def call_nn_parameter(
             )
 
         # this results in cleaner graphs, but only works for inputs
-        # pyrefly: ignore [missing-attribute]
         if data.source:
             return cls._nn_param_via_prefix_insert(tx, data, requires_grad)
 
@@ -2902,7 +3269,6 @@ def call_nn_parameter(
         if isinstance(
             data,
             TensorWithTFOverrideVariable,
-            # pyrefly: ignore [missing-attribute]
         ) or is_traceable_wrapper_subclass_type(data.class_type):
             unimplemented(
                 gb_type="Attempted to use torch.nn.Parameter constructor with tensor subclass",
@@ -2925,11 +3291,8 @@ def call_nn_parameter(
             )
 
         try:
-            # pyrefly: ignore [missing-attribute]
             shape = tuple(data.var_getattr(tx, "shape").as_python_constant())
-            # pyrefly: ignore [missing-attribute]
             dtype = data.var_getattr(tx, "dtype").as_python_constant()
-            # pyrefly: ignore [missing-attribute]
             device = data.var_getattr(tx, "device").as_python_constant()
         except NotImplementedError as e:
             unimplemented(
@@ -2945,12 +3308,9 @@ def call_nn_parameter(
 
         placeholder = tx.output.synthetic_graph_input(
             new_parameter_placeholder,
-            # pyrefly: ignore [unbound-name]
             (shape, dtype, device, requires_grad),
         )
-        # pyrefly: ignore [missing-attribute]
         if data.requires_grad:
-            # pyrefly: ignore [missing-attribute]
             data = data.call_method(tx, "detach", [], {})
 
         from .builder import wrap_fx_proxy
@@ -2960,7 +3320,6 @@ def call_nn_parameter(
             tx.output.create_proxy(
                 "call_function",
                 tracable_create_parameter,
-                # pyrefly: ignore [missing-attribute]
                 (data.as_proxy(), placeholder.as_proxy()),
                 {},
             ),
@@ -2998,7 +3357,7 @@ def _nn_param_via_prefix_insert(
         cg = PyCodegen(tx.output.root_tx)
         cg.add_push_null(lambda: cg.load_import_from("torch.nn", "Parameter"))
         cg(data.source)
-        cg(variables.ConstantVariable(requires_grad))
+        cg(VariableTracker.build(tx, requires_grad))
         cg.call_function(2, False)
         cg.store(varname)
         tx.output.pregraph_bytecode.extend(cg.get_instructions())
@@ -3099,14 +3458,15 @@ def call_method(
             args, kwargs
         ):
             method = getattr(self.value, name)
-            return variables.ConstantVariable.create(
+            return VariableTracker.build(
+                tx,
                 method(
                     *[x.as_python_constant() for x in args],
                     **{k: v.as_python_constant() for k, v in kwargs.items()},
                 ),
             )
         elif name == "highestPriorityTypeId":
-            return variables.EnumVariable(self.value.highestPriorityTypeId())
+            return VariableTracker.build(tx, self.value.highestPriorityTypeId())
         return super().call_method(tx, name, args, kwargs)
 
 
@@ -3128,7 +3488,7 @@ def call_method(
         kwargs: dict[str, VariableTracker],
     ) -> "VariableTracker":
         if name == "key":
-            return variables.EnumVariable(self.value.key())
+            return VariableTracker.build(tx, self.value.key())
         elif name == "process":
             return tx.inline_user_function_return(
                 VariableTracker.build(tx, self.value.process.__func__),
@@ -3136,7 +3496,7 @@ def call_method(
                 kwargs,
             )
         elif name in ["level", "batch_size", "randomness"]:
-            return variables.ConstantVariable.create(getattr(self.value, name)())
+            return VariableTracker.build(tx, getattr(self.value, name)())
         elif name == "lower":
             assert not args and not kwargs
             return variables.TemporarilyPopInterpreterStackCtxManagerVariable.create(
diff --git a/torch/_dynamo/variables/torch_function.py b/torch/_dynamo/variables/torch_function.py
index 86df07571c410..c379158a1fa0f 100644
--- a/torch/_dynamo/variables/torch_function.py
+++ b/torch/_dynamo/variables/torch_function.py
@@ -31,7 +31,7 @@
 import operator
 from collections.abc import Generator, Iterable, Sequence
 from types import TracebackType
-from typing import Any, Optional, TYPE_CHECKING
+from typing import Any, TYPE_CHECKING
 
 import torch._C
 import torch.utils._pytree as pytree
@@ -57,7 +57,7 @@
     set_torch_function_mode_stack,
 )
 from .base import VariableTracker
-from .constant import CONSTANT_VARIABLE_NONE, ConstantVariable
+from .constant import ConstantVariable
 from .ctx_manager import GenericContextWrappingVariable
 from .functions import UserMethodVariable
 from .lazy import LazyVariableTracker
@@ -157,8 +157,8 @@ def is_supported_torch_function_mode(ty: type[TorchFunctionMode]) -> bool:
 
     def __init__(
         self,
-        value: Optional[TorchFunctionMode],
-        source: Optional[Source] = None,
+        value: TorchFunctionMode | None,
+        source: Source | None = None,
         **kwargs: Any,
     ) -> None:
         if value is not None:
@@ -203,12 +203,12 @@ def enter(self, tx: "InstructionTranslator") -> VariableTracker:
         from .torch import TorchInGraphFunctionVariable
 
         if isinstance(self.value, NoEnterTorchFunctionMode):
-            return CONSTANT_VARIABLE_NONE
+            return ConstantVariable.create(None)
 
         TorchInGraphFunctionVariable(
             torch._C._push_on_torch_function_stack
         ).call_function(tx, [self], {})
-        return CONSTANT_VARIABLE_NONE
+        return ConstantVariable.create(None)
 
     def exit(self, tx: "InstructionTranslator", *args: Any) -> VariableTracker:
         from .torch import TorchInGraphFunctionVariable
@@ -216,7 +216,7 @@ def exit(self, tx: "InstructionTranslator", *args: Any) -> VariableTracker:
         TorchInGraphFunctionVariable(torch._C._pop_torch_function_stack).call_function(
             tx, [], {}
         )
-        return CONSTANT_VARIABLE_NONE
+        return ConstantVariable.create(None)
 
     def reconstruct_type(self, codegen: "PyCodegen") -> None:
         ty = NoEnterTorchFunctionMode
@@ -245,9 +245,9 @@ def __enter__(self) -> None:
 
     def __exit__(
         self,
-        exc_type: Optional[type[BaseException]],
-        exc_val: Optional[BaseException],
-        exc_tb: Optional[TracebackType],
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: TracebackType | None,
     ) -> None:
         set_torch_function_mode_stack(self.stack)
         self.stack = []
@@ -481,6 +481,8 @@ def call_torch_function(
     types: TupleVariable,
     args: Iterable[Any],
     kwargs: dict[str, Any],
+    *,
+    is_subclass_dispatch: bool = False,
 ) -> Any:
     # This emulates calling __torch_function__, which has a signature
     #   def __torch_function__(cls, func, types, args=(), kwargs=None):
@@ -495,7 +497,20 @@ def call_torch_function(
         VariableTracker.build(tx, tuple(args)),
         VariableTracker.build(tx, kwargs),
     ]
-    return torch_function_var.call_function(tx, tf_args, {})
+    # Mirror the C++ THPModule_disable_torch_function behavior: disable
+    # __torch_function__ subclass dispatch during the call to prevent
+    # re-entrant dispatch on operations inside __torch_function__.
+    # Only do this for subclass dispatch, not mode dispatch. Modes need
+    # subclass dispatch to remain enabled because the mode's
+    # __torch_function__ may re-dispatch to the subclass.
+    tf_state = tx.symbolic_torch_function_state
+    old_subclass_enabled = tf_state.torch_function_subclass_enabled
+    if is_subclass_dispatch and old_subclass_enabled:
+        tf_state.torch_function_subclass_enabled = False
+    try:
+        return torch_function_var.call_function(tx, tf_args, {})
+    finally:
+        tf_state.torch_function_subclass_enabled = old_subclass_enabled
 
 
 def get_torch_function_fn(
@@ -504,10 +519,9 @@ def get_torch_function_fn(
     # The underlying function could be a classmethod, staticmethod, regular
     # function or a function with C-implementation. It doesn't matter as long as
     # they satisfy the calling convention in `call_torch_function`.
-    from .builtin import BuiltinVariable
 
-    args = [vt, ConstantVariable("__torch_function__")]
-    func_vt = BuiltinVariable(getattr).call_function(tx, args, {})
+    args = [vt, VariableTracker.build(tx, "__torch_function__")]
+    func_vt = VariableTracker.build(tx, getattr).call_function(tx, args, {})
     return func_vt
 
 
@@ -556,6 +570,7 @@ def dispatch_torch_function(
         )
 
         if not res.is_constant_match(NotImplemented):
+            tx.output.torch_function_subclass_inlined = True
             return res
 
     unimplemented(
@@ -722,6 +737,7 @@ def call_torch_function(
             types,
             args,
             kwargs,
+            is_subclass_dispatch=True,
         )
 
     def call_method(
diff --git a/torch/_dynamo/variables/user_defined.py b/torch/_dynamo/variables/user_defined.py
index 1ab15902784a9..c8b8a3b1023e7 100644
--- a/torch/_dynamo/variables/user_defined.py
+++ b/torch/_dynamo/variables/user_defined.py
@@ -29,7 +29,6 @@
 import enum
 import functools
 import inspect
-import itertools
 import random
 import sys
 import threading
@@ -45,8 +44,9 @@
 import torch.nn
 from torch._guards import Source, TracingContext
 from torch.utils._python_dispatch import is_traceable_wrapper_subclass_type
+from torch.utils._pytree import GetAttrKey, is_structseq_class
 
-from .. import graph_break_hints, polyfills, variables
+from .. import config, graph_break_hints, polyfills, variables
 from ..bytecode_transformation import create_call_function
 from ..create_parameter_op import do_not_convert_to_tracable_parameter
 from ..exc import (
@@ -56,6 +56,7 @@
     ObservedTypeError,
     ObservedUserStopIteration,
     raise_observed_exception,
+    raise_type_error,
     unimplemented,
 )
 from ..graph_bytecode_inputs import get_external_object_by_index
@@ -63,7 +64,6 @@
 from ..source import (
     AttrSource,
     CallFunctionNoArgsSource,
-    DataclassFieldsSource,
     DictGetItemSource,
     GetItemSource,
     RandomValueSource,
@@ -76,11 +76,9 @@
     check_constant_args,
     cmp_name_to_op_mapping,
     dict_methods,
-    enum_type_methods,
     frozenset_methods,
     get_custom_getattr,
     has_torch_function,
-    is_frozen_dataclass,
     is_lru_cache_wrapped_function,
     is_namedtuple_cls,
     is_wrapper_or_member_descriptor,
@@ -96,8 +94,10 @@
     tuple_methods,
     unpatched_nn_module_getattr,
 )
-from .base import MutationType, raise_type_error_exc, ValueMutationNew, VariableTracker
-from .dicts import ConstDictVariable, DefaultDictVariable, SetVariable
+from .base import MutationType, NO_SUCH_SUBOBJ, ValueMutationNew, VariableTracker
+from .dicts import ConstDictVariable, DefaultDictVariable
+from .hashable import HashableTracker
+from .sets import SetVariable
 
 
 try:
@@ -117,6 +117,7 @@
     from torch._dynamo.symbolic_convert import InstructionTranslator
     from torch._dynamo.variables.constant import ConstantVariable
 
+    from .dicts import DunderDictVariable
     from .lists import ListVariable, TupleVariable
 
 
@@ -154,6 +155,38 @@ def is_cython_function(obj: object) -> bool:
     )
 
 
+def is_pydantic_dataclass_cls(value: object) -> bool:
+    return (
+        inspect.isclass(value)
+        and dataclasses.is_dataclass(value)
+        and "__is_pydantic_dataclass__" in getattr(value, "__dict__", {})
+    )
+
+
+# Types whose instances are data descriptors (have __get__ + (__set__ or __delete__)).
+# CPython invokes data descriptors found on the type MRO *before* checking
+# the instance __dict__.  This set is used by is_data_descriptor for a fast
+# O(1) check before falling back to the generic hasattr probe.
+KNOWN_DATA_DESCRIPTOR_TYPES: frozenset[type] = frozenset(
+    {
+        property,
+        _collections._tuplegetter,
+        types.MemberDescriptorType,
+        types.GetSetDescriptorType,
+    }
+)
+
+
+def is_data_descriptor(obj: object) -> bool:
+    """Return True if *obj* is a data descriptor (has __get__ and (__set__ or __delete__))."""
+    tp = type(obj)
+    if tp in KNOWN_DATA_DESCRIPTOR_TYPES:
+        return True
+    return hasattr(tp, "__get__") and (
+        hasattr(tp, "__set__") or hasattr(tp, "__delete__")
+    )
+
+
 class UserDefinedVariable(VariableTracker):
     value: object
 
@@ -245,108 +278,295 @@ def supported_c_new_functions() -> set[Any]:
             frozenset.__new__,
             tuple.__new__,
             list.__new__,
+            int.__new__,
+            float.__new__,
+            str.__new__,
         }
         return c_new_fns.union(exceptions)
 
     @staticmethod
     def is_supported_new_method(value: object) -> bool:
-        # TODO(anijain2305) - Extend this to support objects with default tp_new
-        # functions.
-        return value in UserDefinedClassVariable.supported_c_new_functions()
+        if value in UserDefinedClassVariable.supported_c_new_functions():
+            return True
+        # Structseq types each define their own C tp_new.
+        owner = getattr(value, "__self__", None)
+        return isinstance(owner, type) and is_structseq_class(owner)
 
     def can_constant_fold_through(self) -> bool:
-        return self.value in self._constant_fold_classes()
-
-    def has_key_in_generic_dict(self, tx: "InstructionTranslator", key: str) -> bool:
-        if tx.output.side_effects.has_pending_mutation_of_attr(self, key):
-            mutated_attr = tx.output.side_effects.load_attr(self, key, deleted_ok=True)
-            return not isinstance(mutated_attr, variables.DeletedVariable)
+        if self.value in self._constant_fold_classes():
+            return True
+        # Enum class calls (e.g., Color(1)) are value lookups that return
+        # existing singleton members, so they can always be constant-folded.
+        return isinstance(self.value, type) and issubclass(self.value, enum.Enum)
+
+    def lookup_cls_mro_attr(self, name: str) -> object:
+        """Walk cls.__mro__ only (not the metaclass chain) to find *name*."""
+        for base in self.value.__mro__:
+            if name in base.__dict__:
+                return base.__dict__[name]
+        return NO_SUCH_SUBOBJ
+
+    def lookup_metaclass_attr(self, name: str) -> object:
+        """Walk type(cls).__mro__ (the metaclass chain) to find *name*."""
+        for base in type(self.value).__mro__:
+            if name in base.__dict__:
+                return base.__dict__[name]
+        return NO_SUCH_SUBOBJ
+
+    def bool_impl(
+        self,
+        tx: "InstructionTranslator",
+    ) -> "VariableTracker":
+        from .constant import ConstantVariable
 
-        return key in self.value.__dict__
+        # bool() on a class consults the metaclass __bool__.
+        # If the metaclass is the default `type`, all classes are truthy.
+        metaclass = type(self.value)
+        if hasattr(metaclass, "__bool__") and metaclass is not type:
+            return self.call_method(tx, "__bool__", [], {})
+        return ConstantVariable.create(True)
 
     def var_getattr(self, tx: "InstructionTranslator", name: str) -> VariableTracker:
-        from . import ConstantVariable, EnumVariable
-
         source = AttrSource(self.source, name) if self.source is not None else None
 
-        if name == "__name__":
-            return ConstantVariable.create(self.value.__name__)
-        elif name == "__qualname__":
-            return ConstantVariable.create(self.value.__qualname__)
-        elif name == "__dict__":
-            options = {"source": source}
-            return variables.GetAttrVariable(self, name, None, **options)
-        elif name == "__mro__":
-            attr_source = self.source and TypeMROSource(self.source)
-            return VariableTracker.build(tx, self.value.__mro__, attr_source)
+        # --- Dynamo-specific pre-checks ---
 
-        # Special handling of collections.OrderedDict.fromkeys()
-        # Wrap it as GetAttrVariable(collections.OrderedDict, "fromkeys") to make it consistent with
-        # collections.defaultdict, and both will be handled at UserDefinedClassVariable.call_method().
-        # Otherwise, it would be wrapped as UserDefinedObjectVariable(collections.OrderedDict.fromkeys),
-        # and we need duplicate code to handle both cases.
+        # Wrap OrderedDict/defaultdict.fromkeys as GetAttrVariable so it's
+        # handled uniformly in call_method().
         if (
             self.value in {collections.OrderedDict, collections.defaultdict}
             and name == "fromkeys"
         ):
             return super().var_getattr(tx, name)
 
-        obj = None
-        try:
-            obj = inspect.getattr_static(self.value, name)
-        except AttributeError:
-            if type(self.value) is type:
-                raise_observed_exception(
-                    AttributeError,
-                    tx,
-                    args=[
-                        f"type object '{self.value.__name__}' has no attribute '{name}'"
-                    ],
+        # Custom metaclasses that override __getattribute__ replace the entire
+        # lookup algorithm; bail out for those. Standard metaclasses (ABCMeta,
+        # EnumType, etc.) that don't override __getattribute__ use
+        # type.__getattribute__ which is the algorithm we implement below.
+        metacls = type(self.value)
+        if metacls is not type and "__getattribute__" in metacls.__dict__:
+            unimplemented(
+                gb_type="Custom metaclass with __getattribute__",
+                context=f"type({self.value}) = {metacls}",
+                explanation="Dynamo does not trace attribute access on classes whose "
+                "metaclass overrides __getattribute__",
+                hints=graph_break_hints.SUPPORTABLE,
+            )
+
+        # ---- CPython type_getattro algorithm ----
+        # https://github.com/python/cpython/blob/3.13/Objects/typeobject.c#L5417-L5505
+        # 1. meta_attr = lookup name in type(cls).__mro__  (metaclass chain)
+        # 2. if meta_attr is a DATA descriptor → invoke
+        # 3. cls_attr = lookup name in cls.__mro__  (class chain)
+        # 4. if cls_attr has __get__ → invoke cls_attr.__get__(None, cls)
+        # 5. if cls_attr exists (plain) → return as-is
+        # 6. if meta_attr is a non-data descriptor or plain → return
+        # 7. raise AttributeError
+
+        # Step 1-2: Metaclass data descriptors.
+        # For type(cls) is type, these are C-level getset/member descriptors
+        # for __dict__, __mro__, __name__, __qualname__, __doc__, etc.
+        meta_attr = self.lookup_metaclass_attr(name)
+        if meta_attr is not NO_SUCH_SUBOBJ and is_data_descriptor(meta_attr):
+            return self.resolve_meta_data_descriptor(tx, name, meta_attr, source)
+
+        # Step 3-5: Class MRO lookup.
+        cls_attr = self.lookup_cls_mro_attr(name)
+        if cls_attr is not NO_SUCH_SUBOBJ:
+            if hasattr(type(cls_attr), "__get__"):
+                # Step 4: Descriptor — invoke __get__(None, cls).
+                return self.resolve_cls_descriptor(tx, name, cls_attr, source)
+            # Step 5: Plain attribute.
+            return self.resolve_cls_plain_attr(tx, name, cls_attr, source)
+
+        # Step 6: Metaclass non-data descriptor or plain attr.
+        # These are non-data descriptors on the metaclass (e.g. type.__call__,
+        # type.__subclasses__, type.mro).  We use GetAttrVariable to defer to
+        # runtime rather than VariableTracker.build, because build would create
+        # a variable for the raw C-level descriptor which then fails when
+        # called (e.g. type.__subclasses__ is a method_descriptor that dynamo
+        # can't trace).  GetAttrVariable defers the access and lets
+        # call_method handle it.
+        if meta_attr is not NO_SUCH_SUBOBJ:
+            return variables.GetAttrVariable(self, name, type(meta_attr), source=source)
+
+        # __getattr__ on metaclass (not part of type_getattro proper —
+        # CPython handles this via slot_tp_getattr_hook).
+        metacls = type(self.value)
+        if metacls is not type:
+            meta_getattr = self.lookup_metaclass_attr("__getattr__")
+            if meta_getattr is not NO_SUCH_SUBOBJ and isinstance(
+                meta_getattr, types.FunctionType
+            ):
+                return variables.UserMethodVariable(meta_getattr, self).call_function(
+                    tx, [variables.ConstantVariable.create(name)], {}
                 )
 
-        if name == "__new__" and UserDefinedClassVariable.is_supported_new_method(obj):
-            return super().var_getattr(tx, name)
+        # Step 7: AttributeError.
+        raise_observed_exception(
+            AttributeError,
+            tx,
+            args=[f"type object '{self.value.__name__}' has no attribute '{name}'"],
+        )
+
+    def resolve_meta_data_descriptor(
+        self,
+        tx: "InstructionTranslator",
+        name: str,
+        meta_attr: object,
+        source: Source | None,
+    ) -> VariableTracker:
+        """Handle data descriptors from the metaclass MRO (type.__dict__ slots)."""
+        if name == "__dict__":
+            return VariableTracker.build(
+                tx,
+                self.value.__dict__,
+                source=self.source and AttrSource(self.source, "__dict__"),
+            )
+        if name == "__mro__":
+            attr_source = self.source and TypeMROSource(self.source)
+            return VariableTracker.build(tx, self.value.__mro__, attr_source)
+        # __name__, __qualname__, __doc__, __module__, __bases__,
+        # __abstractmethods__, etc. — all C-level getset descriptors on type.
+        resolved = type.__getattribute__(self.value, name)
+        if source:
+            return VariableTracker.build(tx, resolved, source)
+        from . import ConstantVariable
 
-        if name in cmp_name_to_op_mapping and not isinstance(obj, types.FunctionType):
-            return variables.GetAttrVariable(self, name, None, source=source)
+        if ConstantVariable.is_literal(resolved):
+            return VariableTracker.build(tx, resolved)
+        return variables.GetAttrVariable(self, name, type(resolved), source=source)
+
+    def resolve_cls_descriptor(
+        self,
+        tx: "InstructionTranslator",
+        name: str,
+        cls_attr: object,
+        source: Source | None,
+    ) -> VariableTracker:
+        """Handle descriptors found in cls.__mro__."""
+        if isinstance(cls_attr, staticmethod):
+            return VariableTracker.build(tx, cls_attr.__get__(self.value), source)
 
-        if isinstance(obj, staticmethod):
-            return VariableTracker.build(tx, obj.__get__(self.value), source)
-        elif isinstance(obj, classmethod):
-            if isinstance(obj.__func__, property):
-                fget_vt = VariableTracker.build(tx, obj.__func__.fget)
+        if isinstance(cls_attr, classmethod):
+            if isinstance(cls_attr.__func__, property):
+                fget_vt = VariableTracker.build(tx, cls_attr.__func__.fget)
                 return fget_vt.call_function(tx, [self], {})
-            return variables.UserMethodVariable(obj.__func__, self, source=source)
-        elif isinstance(obj, types.ClassMethodDescriptorType):
-            # e.g.: inspect.getattr_static(dict, "fromkeys")
-            #       inspect.getattr_static(itertools.chain, "from_iterable")
-            func = obj.__get__(None, self.value)
+            return variables.UserMethodVariable(cls_attr.__func__, self, source=source)
+
+        if isinstance(cls_attr, types.ClassMethodDescriptorType):
+            func = cls_attr.__get__(None, self.value)
             return VariableTracker.build(tx, func, source)
-        elif source:
-            if inspect.ismemberdescriptor(obj):
-                return VariableTracker.build(tx, obj.__get__(self.value), source)
-
-        if ConstantVariable.is_literal(obj):
-            return ConstantVariable.create(obj)
-        elif isinstance(obj, enum.Enum):
-            return EnumVariable(obj)
-        elif self.value is collections.OrderedDict:
-            return variables.GetAttrVariable(self, name)
-        elif name in getattr(self.value, "__dict__", {}) or (
-            self.value.__module__.startswith("torch.")
-            or self.value.__module__ == "torch"
-        ):
+
+        # property and _tuplegetter accessed on the class return the
+        # descriptor itself (descriptor.__get__(None, cls) is descriptor).
+        # Build directly — no need to invoke __get__.
+        if isinstance(cls_attr, (property, _collections._tuplegetter)):
             if source:
-                return VariableTracker.build(tx, obj, source)
+                return VariableTracker.build(tx, cls_attr, source)
+            return UserDefinedObjectVariable(cls_attr)
 
-        if (
-            source
-            and not inspect.ismethoddescriptor(obj)
-            and not is_wrapper_or_member_descriptor(obj)
+        # Comparison dunders inherited from object — defer to runtime.
+        if name in cmp_name_to_op_mapping and not isinstance(
+            cls_attr, types.FunctionType
         ):
-            return VariableTracker.build(tx, obj, source)
+            return variables.GetAttrVariable(
+                self, name, py_type=type(cls_attr), source=source
+            )
 
-        return super().var_getattr(tx, name)
+        # User-defined descriptor with Python __get__.
+        # For torch-internal classes or attributes in the class's own __dict__,
+        # defer descriptor invocation to runtime via VariableTracker.build to
+        # avoid compile-time side effects (e.g. deprecation warnings from
+        # _ClassPropertyDescriptor on torch.FloatStorage.dtype).
+        get_fn = inspect.getattr_static(type(cls_attr), "__get__", None)
+        if isinstance(get_fn, types.FunctionType):
+            if source and (
+                name in getattr(self.value, "__dict__", {})
+                or self.value.__module__.startswith("torch.")
+                or self.value.__module__ == "torch"
+            ):
+                return VariableTracker.build(tx, cls_attr, source)
+            return self.invoke_cls_descriptor_get(tx, name, cls_attr, source)
+
+        # C-level descriptors (WrapperDescriptor, MethodDescriptor, etc.)
+        # Build directly when the attribute lives in the class's own __dict__
+        # or the class belongs to torch (needed for e.g. torch.Tensor.dim).
+        # OrderedDict's C-level methods are handled at runtime.
+        if inspect.ismethoddescriptor(cls_attr) or is_wrapper_or_member_descriptor(
+            cls_attr
+        ):
+            if (
+                source
+                and self.value is not collections.OrderedDict
+                and (
+                    name in getattr(self.value, "__dict__", {})
+                    or self.value.__module__.startswith("torch.")
+                    or self.value.__module__ == "torch"
+                )
+            ):
+                return VariableTracker.build(tx, cls_attr, source)
+            return variables.GetAttrVariable(self, name, type(cls_attr), source=source)
+
+        # Everything else: FunctionType, etc.
+        return VariableTracker.build(tx, cls_attr, source)
+
+    def resolve_cls_plain_attr(
+        self,
+        tx: "InstructionTranslator",
+        name: str,
+        cls_attr: object,
+        source: Source | None,
+    ) -> VariableTracker:
+        """Handle non-descriptor attributes from cls.__mro__."""
+        if name == "__new__" and UserDefinedClassVariable.is_supported_new_method(
+            cls_attr
+        ):
+            return super().var_getattr(tx, name)
+        if self.value is collections.OrderedDict:
+            return variables.GetAttrVariable(self, name, py_type=type(cls_attr))
+        return VariableTracker.build(tx, cls_attr, source)
+
+    def invoke_cls_descriptor_get(
+        self,
+        tx: "InstructionTranslator",
+        name: str,
+        descriptor: object,
+        source: Source | None,
+    ) -> VariableTracker:
+        """Trace a class-MRO descriptor's __get__(None, cls) call."""
+        from .constant import ConstantVariable
+
+        descriptor_source = None
+        descriptor_get_source = None
+        if self.source:
+            descriptor_source = AttrSource(self.source, name)
+            descriptor_get_source = AttrSource(TypeSource(descriptor_source), "__get__")
+            descriptor_var = VariableTracker.build(tx, descriptor, descriptor_source)
+        else:
+            descriptor_var = UserDefinedObjectVariable(descriptor)
+
+        none_var = ConstantVariable.create(None)
+        return variables.UserMethodVariable(
+            descriptor.__get__.__func__,  # type: ignore[union-attr]
+            descriptor_var,
+            source=descriptor_get_source,
+        ).call_function(tx, [none_var, self], {})
+
+    def len_impl(self, tx: "InstructionTranslator") -> VariableTracker:
+        m = self._maybe_get_baseclass_method("__len__")
+        if m:
+            source = self.source and AttrSource(self.source, "__len__")
+            return variables.UserMethodVariable(
+                m, self, source_fn=source
+            ).call_function(tx, [], {})
+        raise_type_error(tx, f"object of type {self.python_type_name()} has no length")
+
+    def sq_length(self, tx: "InstructionTranslator") -> VariableTracker:
+        return self.len_impl(tx)
+
+    def mp_length(self, tx: "InstructionTranslator") -> VariableTracker:
+        return self.len_impl(tx)
 
     def _call_cross_entropy_loss(
         self,
@@ -363,13 +583,13 @@ def _call_cross_entropy_loss(
 
         non functional loss call: input, target, optional_output
         """
-        from . import CONSTANT_VARIABLE_NONE, ConstantVariable
+        from . import ConstantVariable
 
         def normalize_args(
-            weight: VariableTracker = CONSTANT_VARIABLE_NONE,
-            size_average: VariableTracker = CONSTANT_VARIABLE_NONE,
+            weight: VariableTracker = ConstantVariable.create(None),
+            size_average: VariableTracker = ConstantVariable.create(None),
             ignore_index: VariableTracker = ConstantVariable.create(-100),
-            reduce: VariableTracker = CONSTANT_VARIABLE_NONE,
+            reduce: VariableTracker = ConstantVariable.create(None),
             reduction: VariableTracker = ConstantVariable.create("mean"),
             label_smoothing: VariableTracker = ConstantVariable.create(0.0),
         ) -> tuple[VariableTracker, ...]:
@@ -443,15 +663,19 @@ def call_method(
             self.value in {collections.OrderedDict, collections.defaultdict}
             and name == "fromkeys"
         ):
-            return variables.BuiltinVariable.call_custom_dict_fromkeys(
+            return variables.DictBuiltinVariable.call_custom_dict_fromkeys(
                 tx, self.value, *args, **kwargs
             )
         elif self.value is collections.OrderedDict and name == "move_to_end":
             return args[0].call_method(tx, name, [*args[1:]], kwargs)
+        elif name == "__len__" and len(args) == 1 and not kwargs:
+            from .object_protocol import generic_len
+
+            return generic_len(tx, args[0])
         elif name == "__eq__" and len(args) == 1 and hasattr(args[0], "value"):
-            return variables.ConstantVariable(self.value == args[0].value)
+            return VariableTracker.build(tx, self.value == args[0].value)
         elif name == "__ne__" and len(args) == 1 and hasattr(args[0], "value"):
-            return variables.ConstantVariable(self.value != args[0].value)
+            return VariableTracker.build(tx, self.value != args[0].value)
         elif issubclass(self.value, dict) and name != "__new__":
             # __new__ is handled below
             return SourcelessBuilder.create(tx, dict).call_method(
@@ -497,8 +721,31 @@ def call_method(
                 explanation="Dyanmo does not support tracing mutations on a class when its __dict__ is materialized",
                 hints=graph_break_hints.SUPPORTABLE,
             )
+
+        # Dispatch dunder methods defined on the metaclass (e.g., EnumType.__contains__).
+        # In Python, `x in Color` calls `type(Color).__contains__(Color, x)`.
+        metaclass = type(self.value)
+        if metaclass is not type:
+            # Look up the method on the metaclass MRO, not the class MRO
+            for klass in metaclass.__mro__:
+                if name in klass.__dict__:
+                    method = klass.__dict__[name]
+                    if isinstance(method, types.FunctionType):
+                        source = self.source and AttrSource(self.source, name)
+                        return variables.UserMethodVariable(
+                            method, self, source=source
+                        ).call_function(tx, args, kwargs)
+                    break
+
         return super().call_method(tx, name, args, kwargs)
 
+    def unpack_var_sequence(
+        self, tx: "InstructionTranslator"
+    ) -> list["VariableTracker"]:
+        if isinstance(self.value, type) and issubclass(self.value, enum.Enum):
+            return [VariableTracker.build(tx, item) for item in self.value]
+        raise NotImplementedError
+
     def call_function(
         self,
         tx: "InstructionTranslator",
@@ -511,9 +758,28 @@ def call_function(
 
         constant_args = check_constant_args(args, kwargs)
 
+        if torch.distributed.is_available() and self.value is torch.distributed.P2POp:
+            if not config.enable_p2p_compilation:
+                unimplemented(
+                    gb_type="P2P compilation disabled for P2POp construction",
+                    context="torch.distributed.P2POp",
+                    explanation="P2P compilation is disabled.",
+                    hints=[
+                        "Set TORCHDYNAMO_ENABLE_P2P_COMPILATION=1 to enable.",
+                    ],
+                )
+            var = tx.output.side_effects.track_new_user_defined_object(
+                SourcelessBuilder.create(tx, object),
+                self,
+                [],
+            )
+            var.call_method(tx, "__init__", list(args), kwargs)  # type: ignore[arg-type]
+            return var
+
         if self.can_constant_fold_through() and constant_args:
             # constant fold
-            return variables.ConstantVariable.create(
+            return VariableTracker.build(
+                tx,
                 self.as_python_constant()(  # type: ignore[operator]
                     *[x.as_python_constant() for x in args],
                     **{k: v.as_python_constant() for k, v in kwargs.items()},
@@ -534,7 +800,7 @@ def call_function(
             )
         elif self.value is collections.defaultdict:
             if len(args) == 0:
-                default_factory = variables.CONSTANT_VARIABLE_NONE
+                default_factory = variables.ConstantVariable.create(None)
             elif len(args) == 1:
                 # In the case the argument is a builtin, then we can take the callable as the factory method.
                 # Otherwise, it must be a ConstantVariable holding None.
@@ -544,7 +810,7 @@ def call_function(
                 args = []
             else:
                 default_factory, *args = args
-            dict_vt = variables.BuiltinVariable.call_custom_dict(
+            dict_vt = variables.DictBuiltinVariable.call_custom_dict(
                 tx, dict, *args, **kwargs
             )
             return DefaultDictVariable(
@@ -564,9 +830,11 @@ def call_function(
                         *graph_break_hints.SUPPORTABLE,
                     ],
                 )
-            return SourcelessBuilder.create(tx, dict).call_dict(tx, *args, **kwargs)
+            return variables.DictBuiltinVariable.call_custom_dict(
+                tx, dict, *args, **kwargs
+            )
         elif self.value is collections.deque:
-            maxlen = variables.CONSTANT_VARIABLE_NONE
+            maxlen = variables.ConstantVariable.create(None)
 
             def deque_signature(
                 iterable: Iterable[Any] | None = None, maxlen: int | None = None
@@ -629,7 +897,7 @@ def deque_signature(
             if len(args) > 1:
                 callback = args[1]
             else:
-                callback = variables.CONSTANT_VARIABLE_NONE
+                callback = variables.ConstantVariable.create(None)
             return variables.WeakRefVariable(args[0], callback)
         elif self.value is functools.partial:
             if not args:
@@ -654,9 +922,7 @@ def deque_signature(
             return variables.CatchWarningsCtxManagerVariable.create(tx, kwargs)
         elif self.value is torch.cuda.device and not kwargs and len(args) == 1:
             if not args[0].is_python_constant():
-                raise_type_error_exc(
-                    tx, "torch.cuda.device() requires a constant argument"
-                )
+                raise_type_error(tx, "torch.cuda.device() requires a constant argument")
             return variables.CUDADeviceVariable.create(tx, args[0].as_python_constant())
         elif (
             issubclass(type(self.value), type)
@@ -746,17 +1012,15 @@ def deque_signature(
                     )
                 ] + args[1:]
 
-            cm_obj = tx.output.side_effects.track_new_user_defined_object(
-                SourcelessBuilder.create(tx, object),
-                self,
-                arg_new,  # type: ignore[arg-type]
+            return tx.inline_user_function_return(
+                VariableTracker.build(
+                    tx, polyfills.instantiate_user_defined_class_object
+                ),
+                [self, *arg_new],
+                kwargs,
             )
-            cm_obj.call_method(tx, "__init__", arg_new, kwargs)  # type: ignore[arg-type]
-            return cm_obj
         elif is_namedtuple_cls(self.value):
-            fields = namedtuple_fields(self.value)  # type: ignore[arg-type]
-            # check if this a quasi-namedtuple or a real one
-            if self.value.__module__ == "torch.return_types":
+            if is_structseq_class(self.value):
                 if kwargs or len(args) != 1:
                     raise_args_mismatch(
                         tx,
@@ -764,91 +1028,42 @@ def deque_signature(
                         "1 args and 0 kwargs",
                         f"{len(args)} args and {len(kwargs)} kwargs",
                     )
-                items = args[0].force_unpack_var_sequence(tx)
+                # Structseq tp_new is a C function, so we can't trace into
+                # it like namedtuples. Use track_new_user_defined_object
+                # directly with self as both base_cls_vt and cls_vt.
+                return tx.output.side_effects.track_new_user_defined_object(
+                    self,
+                    self,
+                    list(args),
+                )
             else:
-                field_defaults = self.value._field_defaults  # type: ignore[attr-defined]
-
-                items = list(args)
-                # pyrefly: ignore[bad-argument-type]
-                items.extend([None] * (len(fields) - len(items)))
-
-                var_tracker_kwargs: dict[str, VariableTracker] = {}
-                for field_name, var_tracker in zip(fields, items):
-                    if var_tracker is None:
-                        if field_name in kwargs:
-                            field_var = kwargs[field_name]
-                        else:
-                            assert field_name in field_defaults
-                            field_var = VariableTracker.build(
-                                tx, field_defaults[field_name]
-                            )
-                        var_tracker_kwargs[field_name] = field_var
-
-                for name, value in var_tracker_kwargs.items():
-                    assert name in fields
-                    items[fields.index(name)] = value  # type: ignore[call-overload]
-
-                assert all(x is not None for x in items)
-
-            # Modify mutability of namedtuple for sourcelesss instantiations.
-            from .base import AttributeMutationNew
-            from .lists import NamedTupleVariable
-
-            return NamedTupleVariable(
-                items,
-                self.value,  # type: ignore[arg-type]
-                mutation_type=AttributeMutationNew(),
-            )
+                # Namedtuple __new__ is a Python function that calls
+                # tuple.__new__(cls, (field_values,)). Let Dynamo trace
+                # into it so default values and kwargs are handled by
+                # the generated __new__ itself.
+                return tx.inline_user_function_return(
+                    VariableTracker.build(
+                        tx, polyfills.instantiate_user_defined_class_object
+                    ),
+                    [self, *args],
+                    kwargs,
+                )
         elif self.value is torch.Size:
             # This simulates `THPSize_pynew`, the C impl for `Size.__new__`.
             from .lists import SizeVariable
 
             tup = SourcelessBuilder.create(tx, tuple).call_function(tx, args, kwargs)
             return SizeVariable(tup.items)  # type: ignore[missing-attribute]
-        elif is_frozen_dataclass(self.value) and self.is_standard_new():
-            fields = dataclasses.fields(self.value)  # type: ignore[arg-type]
-            assert self.source is not None
-            fields_source = DataclassFieldsSource(self.source)
-            items = list(args)
-            items.extend([None] * (len(fields) - len(items)))  # type: ignore[arg-type]
-
-            default_kwargs = {}
-            for ind, field, var_tracker in zip(itertools.count(), fields, items):
-                if var_tracker is None:
-                    if field.name in kwargs:
-                        var_tracker = kwargs[field.name]
-                    else:
-                        if not field.init:
-                            continue
-
-                        if field.default is not dataclasses.MISSING:
-                            var_tracker = VariableTracker.build(
-                                tx,
-                                field.default,
-                                source=AttrSource(
-                                    GetItemSource(fields_source, ind), "default"
-                                ),
-                            )
-                        elif field.default_factory is not dataclasses.MISSING:
-                            factory_fn = VariableTracker.build(
-                                tx, field.default_factory
-                            )
-                            var_tracker = factory_fn.call_function(tx, [], {})
-                        else:
-                            # if we are subclass, the constructor could possibly
-                            # be missing args
-                            continue
-
-                    default_kwargs[field.name] = var_tracker
-            kwargs.update(default_kwargs)
-
-            var = tx.output.side_effects.track_new_user_defined_object(
-                SourcelessBuilder.create(tx, object),
-                self,
-                args,  # type: ignore[arg-type]
+        elif is_pydantic_dataclass_cls(self.value):
+            # Pydantic populates dataclass fields through an external validator,
+            # so tracing through the constructor misses the instance mutations.
+            unimplemented(
+                gb_type="Pydantic dataclass constructor",
+                context=f"{self.value}",
+                explanation="Dynamo graph breaks on pydantic dataclass constructors "
+                "because validation mutates the instance outside traced bytecode.",
+                hints=graph_break_hints.SUPPORTABLE,
             )
-            var.call_method(tx, "__init__", args, kwargs)  # type: ignore[arg-type]
-            return var
         elif (
             self.value in self._in_graph_classes()
             or is_traceable_wrapper_subclass_type(self.value)
@@ -877,11 +1092,10 @@ def deque_signature(
                 args = [stacked]
 
             if issubclass(self.value, torch.Stream):
-                from .constant import ConstantVariable
                 from .lists import TupleVariable
 
                 var_kwargs = ConstDictVariable(
-                    {ConstantVariable(k): v for k, v in kwargs.items()}
+                    {VariableTracker.build(tx, k): v for k, v in kwargs.items()}
                 )
                 var_args = TupleVariable(list(args))
                 stream = self.value(
@@ -904,12 +1118,11 @@ def deque_signature(
                     ),
                 )
             elif issubclass(self.value, torch.Event):
-                from .constant import ConstantVariable
                 from .lists import TupleVariable
 
                 # Register newly created event for reconstruction
                 var_kwargs = ConstDictVariable(
-                    {ConstantVariable(k): v for k, v in kwargs.items()}
+                    {VariableTracker.build(tx, k): v for k, v in kwargs.items()}
                 )
                 var_args = TupleVariable(list(args))
                 event = self.value(
@@ -966,6 +1179,7 @@ def deque_signature(
                     [self, *args],
                     kwargs,
                 )
+
         return super().call_function(tx, args, kwargs)
 
     def is_standard_new(self) -> bool:
@@ -979,10 +1193,12 @@ def call_obj_hasattr(
         self, tx: "InstructionTranslator", name: str
     ) -> "ConstantVariable":
         if self.source:
-            source = AttrSource(self.source, name)
-            install_guard(source.make_guard(GuardBuilder.HASATTR))
-            return variables.ConstantVariable(hasattr(self.value, name))
-        return super().call_obj_hasattr(tx, name)
+            install_guard(
+                self.source.make_guard(
+                    functools.partial(GuardBuilder.HASATTR, attr=name)
+                )
+            )
+        return VariableTracker.build(tx, hasattr(self.value, name))
 
     def const_getattr(self, tx: "InstructionTranslator", name: str) -> Any:
         if name == "__name__":
@@ -1001,54 +1217,34 @@ def is_python_equal(self, other: object) -> bool:
             and self.value is other.value
         )
 
+    def get_real_python_backed_value(self) -> object:
+        return self.value
+
 
 class UserDefinedExceptionClassVariable(UserDefinedClassVariable):
     @property
     def fn(self) -> type[object]:
         return self.value
 
-
-class UserDefinedEnumClassVariable(UserDefinedClassVariable):
-    """
-    Represents Enum class objects (the class itself, not instances).
-
-    Handles Enum metaclass methods like __contains__ by checking if the method
-    is from the standard EnumType metaclass and executing it directly.
-
-    Not yet supported:
-    - __iter__: iteration over enum members (e.g., `for x in SomeEnum`)
-    - __reversed__: reversed iteration
-    - Flag enum membership checks (e.g., `Flag.A in combined_flags`)
-    """
-
-    # pyrefly: ignore[bad-override]
-    value: type[enum.Enum]
-
-    def call_method(
+    def call_function(
         self,
         tx: "InstructionTranslator",
-        name: str,
-        args: list[VariableTracker],
+        args: Sequence[VariableTracker],
         kwargs: dict[str, VariableTracker],
     ) -> VariableTracker:
-        method = self._maybe_get_baseclass_method(name)
-        if method in enum_type_methods:
-            if name == "__contains__" and len(args) == 1 and not kwargs:
-                arg = args[0]
-                if isinstance(arg, variables.EnumVariable):
-                    # Check if the enum value is a member of this enum class
-                    return variables.ConstantVariable.create(arg.value in self.value)
-                elif arg.is_python_constant():
-                    # Check if a constant value is in the enum
-                    return variables.ConstantVariable.create(
-                        arg.as_python_constant() in self.value
-                    )
-
-        return super().call_method(tx, name, args, kwargs)
-
+        from .builder import SourcelessBuilder
 
-class NO_SUCH_SUBOBJ:
-    pass
+        if self.source is None:
+            # NB: If source is added via side effects, create the exception
+            # object through side_effects as well. See FrozenDataClass creation
+            var = tx.output.side_effects.track_new_user_defined_object(
+                SourcelessBuilder.create(tx, BaseException),
+                self,
+                list(args),
+            )
+            var.call_method(tx, "__init__", list(args), dict(kwargs))
+            return var
+        return super().call_function(tx, args, kwargs)
 
 
 class RemovableHandleClass:
@@ -1069,6 +1265,8 @@ def call_random_fn(
     args = [x.as_python_constant() for x in args]
     kwargs = {k: v.as_python_constant() for k, v in kwargs.items()}
     random_call_index = len(tx.output.random_calls)
+    # NB: it is probably not important for the example_value to be exactly correct,
+    # we just need the right type
     example_value = fn(*args, **kwargs)
     source = RandomValueSource(random_call_index)
     tx.output.random_calls.append((fn, args, kwargs))  # type: ignore[arg-type]
@@ -1083,10 +1281,18 @@ class UserDefinedObjectVariable(UserDefinedVariable):
     Mostly objects of defined type.  Catch-all for something where we only know the type.
     """
 
+    # VT representing the base built-in type's data for subclassed built-in types
+    # (e.g., ConstDictVariable for dict subclasses, ListVariable for list subclasses).
+    # None for plain user-defined objects that don't subclass a built-in container.
+    _base_vt: VariableTracker | None = None
+
+    # Set of base class methods that can be delegated to _base_vt.
+    # Used to check whether a method is overridden before delegating.
+    _base_methods: set[Any] | None = None
+
     _nonvar_fields = {
         "value",
         "value_type",
-        "attrs_directly_modifed_on_dict",
         *UserDefinedVariable._nonvar_fields,
     }
 
@@ -1116,12 +1322,9 @@ def __init__(
         self.base_cls_vt = base_cls_vt
         self.init_args = init_args
 
-        # This records names of the attributes that were modified via instance
+        # This records the attributes that were modified via instance
         # `__dict__` directly, rather than the normal setattr path.
-        #
-        # TODO consider emulating `obj.__dict__` as a `ConstDictVariable` to get
-        # rid of these workarounds here and in `GetAttrVariable`.
-        self.attrs_directly_modifed_on_dict: set[str] = set()
+        self.dict_vt: DunderDictVariable | None = None
 
         # Cache inspect.getattr_static outputs for the same name. This is fine
         # because if there is a mutation for the name, we use side-effects infra
@@ -1153,13 +1356,29 @@ def __str__(self) -> str:
     def __repr__(self) -> str:
         return f"{self.__class__.__name__}({self.value_type.__name__})"
 
-    def is_underlying_vt_modified(self, side_effects: "SideEffects") -> bool:
+    def get_dict_vt(self, tx: "InstructionTranslator") -> "DunderDictVariable":
+        if self.dict_vt is None:
+            self.dict_vt = variables.DunderDictVariable.create(tx, self)
+        return self.dict_vt
+
+    def is_base_vt_modified(self, side_effects: "SideEffects") -> bool:
+        if self._base_vt is not None:
+            return side_effects.is_modified(self._base_vt)
         return False
 
     def python_type(self) -> type:
         return self.value_type  # type: ignore[return-value]
 
+    def get_real_python_backed_value(self) -> object:
+        return self.value
+
     def as_python_constant(self) -> object:
+        if isinstance(
+            self.value,
+            (enum.Enum, torch.DispatchKey, torch._C._functorch.TransformType),
+        ):
+            return self.value
+
         if self.is_pytree_constant_class and self.source:
             # NOTE pytree constants created in the torch.compile region will
             # NOT be guarded (even though they have a source set)
@@ -1186,12 +1405,114 @@ def as_python_constant(self) -> object:
 
         return super().as_python_constant()
 
+    def as_proxy(self) -> object:
+        if isinstance(self.value, enum.Enum):
+            if isinstance(self.value, int):
+                return int(self.value)
+            return self.value
+        return super().as_proxy()
+
     def guard_as_python_constant(self) -> object:
         if self.source:
             install_guard(self.source.make_guard(GuardBuilder.ID_MATCH))
             return self.value
         return super().guard_as_python_constant()
 
+    def bool_impl(
+        self,
+        tx: "InstructionTranslator",
+    ) -> "VariableTracker | None":
+        # Mirrors slot_nb_bool:
+        # https://github.com/python/cpython/blob/c09ccd9c429/Objects/typeobject.c#L9408-L9458
+        if self._maybe_get_baseclass_method("__bool__"):
+            result = self.call_method(tx, "__bool__", [], {})
+            if result.is_python_constant():
+                result_value = result.as_python_constant()
+                if not isinstance(result_value, bool):
+                    raise_observed_exception(
+                        TypeError,
+                        tx,
+                        args=[
+                            f"__bool__ should return bool, returned {type(result_value).__name__}"
+                        ],
+                    )
+            return result
+        return None
+
+    def nb_index_impl(
+        self,
+        tx: "InstructionTranslator",
+    ) -> VariableTracker:
+        # CPython: PyNumber_Index checks tp_as_number->nb_index.
+        # For user-defined types, __index__ in tp_dict means nb_index is set.
+        type_attr = inspect.getattr_static(type(self.value), "__index__", None)
+        if type_attr is None:
+            return super().nb_index_impl(tx)
+        source = self.source and self.get_source_by_walking_mro(tx, "__index__")
+        method_var = self.resolve_type_attr(tx, "__index__", type_attr, source)
+        result = method_var.call_function(tx, [], {})
+        # CPython validates that __index__ returns an int.
+        # https://github.com/python/cpython/blob/c09ccd9c429/Objects/abstract.c#L1433-L1438
+        if result.is_python_constant() and not isinstance(
+            result.as_python_constant(), int
+        ):
+            raise_observed_exception(
+                TypeError,
+                tx,
+                args=[
+                    f"__index__ returned non-int (type {type(result.as_python_constant()).__name__})"
+                ],
+            )
+        return result
+
+    def nb_int_impl(
+        self,
+        tx: "InstructionTranslator",
+    ) -> VariableTracker:
+        # CPython: slot_nb_int calls __int__(), PyNumber_Long validates the return type.
+        # https://github.com/python/cpython/blob/v3.13.0/Objects/abstract.c#L1538-L1550
+        source = self.source and self.get_source_by_walking_mro(tx, "__int__")
+        method_var = self.resolve_type_attr(
+            tx,
+            "__int__",
+            inspect.getattr_static(type(self.value), "__int__"),
+            source,
+        )
+        result = method_var.call_function(tx, [], {})
+        if not issubclass(result.python_type(), int):
+            raise_observed_exception(
+                TypeError,
+                tx,
+                args=[
+                    f"__int__ returned non-int (type {result.python_type().__name__})"
+                ],
+            )
+        return result
+
+    def nb_float_impl(
+        self,
+        tx: "InstructionTranslator",
+    ) -> VariableTracker:
+        # CPython: slot_nb_float calls __float__(), PyNumber_Float validates the return type.
+        # https://github.com/python/cpython/blob/v3.13.0/Objects/abstract.c#L1647-L1658
+        source = self.source and self.get_source_by_walking_mro(tx, "__float__")
+        method_var = self.resolve_type_attr(
+            tx,
+            "__float__",
+            inspect.getattr_static(type(self.value), "__float__"),
+            source,
+        )
+        result = method_var.call_function(tx, [], {})
+        if not issubclass(result.python_type(), float):
+            raise_observed_exception(
+                TypeError,
+                tx,
+                args=[
+                    f"__float__ returned non-float (type {result.python_type().__name__})"
+                ],
+            )
+        return result
+
     def torch_function_check(self) -> None:
         assert has_torch_function(self), (
             f"calling torch function on object without __torch_function__ {self}"
@@ -1235,6 +1556,28 @@ def _supported_random_functions() -> set[Any]:
         }
         return fns
 
+    def mp_subscript_impl(
+        self,
+        tx: "InstructionTranslator",
+        key: VariableTracker,
+    ) -> VariableTracker:
+        # PyObject_GetItem: https://github.com/python/cpython/blob/62a6e898e01/Objects/abstract.c#L155-L206
+        method = self._maybe_get_baseclass_method("__getitem__")
+        if (
+            self._base_vt is not None
+            and self._base_methods is not None
+            and method in self._base_methods
+        ):
+            return self._base_vt.mp_subscript_impl(tx, key)
+        if isinstance(method, types.FunctionType):
+            source_fn = self.source and self.get_source_by_walking_mro(
+                tx, "__getitem__"
+            )
+            return variables.UserMethodVariable(
+                method, self, source_fn=source_fn, source=self.source
+            ).call_function(tx, [key], {})
+        return super().mp_subscript_impl(tx, key)
+
     def call_method(
         self,
         tx: "InstructionTranslator",
@@ -1242,12 +1585,14 @@ def call_method(
         args: list[Any],
         kwargs: dict[str, Any],
     ) -> VariableTracker:
-        from . import CONSTANT_VARIABLE_NONE, ConstantVariable, UserMethodVariable
+        from .. import trace_rules
+        from . import UserMethodVariable
+        from .constant import ConstantVariable
 
         method = self._maybe_get_baseclass_method(name)
         if method is not None:
             if method is object.__init__:
-                return CONSTANT_VARIABLE_NONE
+                return ConstantVariable.create(None)
 
             if is_standard_setattr(method) or isinstance(self.value, threading.local):
                 return self.method_setattr_standard(tx, *args, **kwargs)
@@ -1260,11 +1605,11 @@ def call_method(
             if method is object.__eq__ and len(args) == 1 and not kwargs:
                 other = args[0]
                 if not isinstance(other, UserDefinedObjectVariable):
-                    return variables.ConstantVariable.create(NotImplemented)
+                    return VariableTracker.build(tx, NotImplemented)
 
                 # TODO(anijain2305) - Identity checking should already be a part
                 # of the cmp_eq  polyfill function.
-                return ConstantVariable.create(self.value is other.value)
+                return VariableTracker.build(tx, self.value is other.value)
 
             if torch._dynamo.config.enable_faithful_generator_behavior and isinstance(
                 self.value, types.GeneratorType
@@ -1281,12 +1626,34 @@ def call_method(
                     ],
                 )
 
+            # torch.Generator methods like manual_seed(), get_state(), etc.
+            # are stateful RNG operations that cannot be soundly traced.
+            if (
+                isinstance(self.value, torch._C.Generator)
+                and name in trace_rules._GENERATOR_METHODS_THAT_GRAPH_BREAK
+            ):
+                unimplemented(
+                    gb_type="torch.Generator method",
+                    context=f"torch.Generator.{name}",
+                    explanation=f"torch.Generator.{name}() is a stateful RNG "
+                    "operation that cannot be soundly traced in the FX graph.",
+                    hints=[*graph_break_hints.FUNDAMENTAL],
+                )
+
+            # Delegate to _base_vt for non-overridden base-class methods
+            if (
+                self._base_vt is not None
+                and self._base_methods is not None
+                and method in self._base_methods
+            ):
+                return self._base_vt.call_method(tx, name, args, kwargs)
+
             # check for methods implemented in C++
             if isinstance(method, types.FunctionType):
                 source = self.source
                 source_fn = None
                 if source:
-                    source_fn = self.get_source_by_walking_mro(name)
+                    source_fn = self.get_source_by_walking_mro(tx, name)
                 # TODO(jansel): add a guard to check for monkey patching?
                 from ..mutation_guard import unpatched_nn_module_init
 
@@ -1298,10 +1665,50 @@ def call_method(
 
             if method is list.__len__ and self.source and not (args or kwargs):
                 install_guard(self.source.make_guard(GuardBuilder.SEQUENCE_LENGTH))
-                return ConstantVariable(len(self.value))  # type: ignore[arg-type]
+                return VariableTracker.build(tx, len(self.value))  # type: ignore[arg-type]
 
         return super().call_method(tx, name, args, kwargs)
 
+    def len_impl(self, tx: "InstructionTranslator") -> VariableTracker:
+        method = self._maybe_get_baseclass_method("__len__")
+        if method is not None:
+            type_attr = self.lookup_class_mro_attr("__len__")
+            source = self.source and self.get_source_by_walking_mro(tx, "__len__")
+            method_var = self.resolve_type_attr(tx, "__len__", type_attr, source)
+            if not isinstance(method_var, variables.GetAttrVariable):
+                return method_var.call_function(tx, [], {})
+
+        unimplemented(
+            gb_type="Cannot trace user-defined __len__",
+            context=f"{self.python_type_name()}.__len__()",
+            explanation=(
+                f"Dynamo cannot trace len() on {self.python_type_name()} because the __len__ "
+                "method is either not traceable (e.g., defined in C or built-in) or returns a "
+                "non-constant value."
+            ),
+            hints=[
+                *graph_break_hints.SUPPORTABLE,
+            ],
+        )
+
+    def sq_length(self, tx: "InstructionTranslator") -> VariableTracker:
+        if (
+            self._base_vt is not None
+            and self._base_methods is not None
+            and self._maybe_get_baseclass_method("__len__") in self._base_methods
+        ):
+            return self._base_vt.sq_length(tx)
+        return self.len_impl(tx)
+
+    def mp_length(self, tx: "InstructionTranslator") -> VariableTracker:
+        if (
+            self._base_vt is not None
+            and self._base_methods is not None
+            and self._maybe_get_baseclass_method("__len__") in self._base_methods
+        ):
+            return self._base_vt.mp_length(tx)
+        return self.len_impl(tx)
+
     def method_setattr_standard(
         self,
         tx: "InstructionTranslator",
@@ -1324,10 +1731,27 @@ def method_setattr_standard(
             "an AttributeMutation mutation_type"
         )
 
-        if name_str == "__class__":
+        if (
+            torch.distributed.is_available()
+            and type(self.value) is torch.distributed.P2POp
+            and (
+                tx.output.side_effects.has_pending_mutation_of_attr(self, name_str)
+                or name_str in self.value.__dict__
+            )
+        ):
             unimplemented(
-                gb_type="__class__ assignment on user-defined object",
-                context=f"object={self}, value={value}",
+                gb_type="P2POp mutation",
+                context=f"object={self}, name={name}, value={value}",
+                explanation="Dynamo does not support mutating torch.distributed.P2POp instances.",
+                hints=[
+                    "Construct a new torch.distributed.P2POp instead of mutating an existing one inside torch.compile.",
+                ],
+            )
+
+        if name_str == "__class__":
+            unimplemented(
+                gb_type="__class__ assignment on user-defined object",
+                context=f"object={self}, value={value}",
                 explanation="Dynamo does not support reassigning __class__ on user-defined objects.",
                 hints=[
                     "Move the __class__ assignment outside of the torch.compile region.",
@@ -1335,7 +1759,7 @@ def method_setattr_standard(
             )
 
         if directly_update_dict:
-            self.attrs_directly_modifed_on_dict.add(name_str)
+            self.get_dict_vt(tx).setitem(name_str, value)
         else:
             tmp = self.try_get_descritor_and_setter_py_func(name_str)
             if tmp:
@@ -1345,19 +1769,51 @@ def method_setattr_standard(
                 desc_source = None
                 func_source = None
                 if self.cls_source:
-                    desc_source = self.get_source_by_walking_mro(name_str)
+                    desc_source = self.get_source_by_walking_mro(tx, name_str)
                     # use `type(...)` to ignore instance attrs.
                     func_source = AttrSource(TypeSource(desc_source), "__set__")
                 desc_var = VariableTracker.build(tx, descriptor, desc_source)
                 func_var = VariableTracker.build(tx, setter, func_source, realize=True)
-                args = [desc_var, self, value]
+                if isinstance(descriptor, property):
+                    args = [self, value]  # property.fset(self, value)
+                else:
+                    args = [desc_var, self, value]  # __set__(desc, self, value)
                 return func_var.call_function(tx, args, {})
+
+            # Handle Python property descriptors whose __set__ is a C slot
+            # wrapper (not a Python function), which the above check misses.
+            # Mirrors the property getter handling in var_getattr.
+            descriptor = inspect.getattr_static(type(self.value), name_str, None)
+            if isinstance(descriptor, property) and descriptor.fset is not None:
+                fset_source = None
+                if self.cls_source:
+                    fset_source = AttrSource(
+                        self.get_source_by_walking_mro(tx, name_str), "fset"
+                    )
+                fset_var = VariableTracker.build(
+                    tx, descriptor.fset, source=fset_source
+                )
+                return fset_var.call_function(tx, [self, value], {})
+
             # NOTE: else we assume the descriptor (if any) has a
             # side-effect-free `__set__` as far as Dynamo tracing is concerned.
 
-        # Emulate the standard setattr on instance dict.
+        # If the code reaches here, the attribute is either:
+        #  1) a slot descriptor
+        #  2) a plain attribute with no descriptor
+        # If the object has no __dict__, only slot descriptors (member_descriptor)
+        # allow mutation. Any other attribute assignment raises AttributeError.
+        if not hasattr(self.value, "__dict__"):
+            descriptor = self.lookup_class_mro_attr(name_str)
+            if not inspect.ismemberdescriptor(descriptor):
+                error_msg = VariableTracker.build(
+                    tx,
+                    f"'{type(self.value).__name__}' object has no attribute '{name_str}'",
+                )
+                raise_observed_exception(AttributeError, tx, args=[error_msg])
+
         tx.output.side_effects.store_attr(self, name_str, value)
-        return variables.CONSTANT_VARIABLE_NONE
+        return variables.ConstantVariable.create(None)
 
     def needs_slow_setattr(self) -> bool:
         return not is_standard_setattr(
@@ -1365,6 +1821,10 @@ def needs_slow_setattr(self) -> bool:
         ) and not isinstance(self.value, threading.local)
 
     def unpack_var_sequence(self, tx: "InstructionTranslator") -> list[VariableTracker]:
+        if self._base_vt is not None and self._base_methods is not None:
+            iter_method = self._maybe_get_baseclass_method("__iter__")
+            if iter_method is not None and iter_method in self._base_methods:
+                return self._base_vt.unpack_var_sequence(tx)
         if (
             self.source
             and self._maybe_get_baseclass_method("__iter__") is list.__iter__
@@ -1449,7 +1909,7 @@ def call_function(
                 and obj.__class__ is torch.autograd.grad_mode.inference_mode
             ):
                 # simulate the inference_mode.clone implementation
-                var = variables.ConstantVariable(obj.mode)  # type: ignore[attr-defined]
+                var = VariableTracker.build(tx, obj.mode)  # type: ignore[attr-defined]
                 return variables.TorchCtxManagerClassVariable(
                     obj.__class__
                 ).call_function(tx, [var], kwargs)
@@ -1523,6 +1983,26 @@ def _getattr_static(self, name: str) -> object:
         self._looked_up_attrs[name] = subobj
         return subobj
 
+    def lookup_class_mro_attr(self, name: str) -> object:
+        """Walk type(obj).__mro__ to find *name* in the class hierarchy.
+
+        This only searches the class chain (type(obj).__mro__), NOT the
+        metaclass chain (type(type(obj)).__mro__).  The distinction matters
+        because inspect.getattr_static conflates both chains — it can return
+        metaclass descriptors (e.g. type.__dict__['__annotations__'], a
+        getset_descriptor) when the attribute doesn't exist on the class MRO.
+        Walking cls.__mro__ directly avoids that leak.
+        """
+        if name in self._subobj_from_class:
+            return self._subobj_from_class[name]
+        result = NO_SUCH_SUBOBJ
+        for base in self.value.__class__.__mro__:
+            if name in base.__dict__:
+                result = base.__dict__[name]
+                break
+        self._subobj_from_class[name] = result
+        return result
+
     def should_skip_descriptor_setter(self, attr_name: str) -> bool:
         # Check if `attr_name` corresponds to a descriptor.
         descriptor = inspect.getattr_static(type(self.value), attr_name, None)
@@ -1533,7 +2013,7 @@ def should_skip_descriptor_setter(self, attr_name: str) -> bool:
                 return True
             # For untraceable `__set__` we should still skip if the attribute
             # was mutated via instance `__dict__`.
-            elif attr_name in self.attrs_directly_modifed_on_dict:
+            elif self.dict_vt and self.dict_vt.contains(attr_name):
                 return True
         return False
 
@@ -1541,6 +2021,9 @@ def try_get_descritor_and_setter_py_func(
         self, attr_name: str
     ) -> tuple[object, object] | None:
         descriptor = inspect.getattr_static(type(self.value), attr_name, None)
+        # Handle property descriptors with setters - call fset directly
+        if isinstance(descriptor, property) and descriptor.fset is not None:
+            return (descriptor, descriptor.fset)
         setter = inspect.getattr_static(type(descriptor), "__set__", None)
         if inspect.isfunction(setter):
             return (descriptor, setter)
@@ -1551,38 +2034,49 @@ def has_key_in_generic_dict(self, tx: "InstructionTranslator", key: str) -> bool
             mutated_attr = tx.output.side_effects.load_attr(self, key, deleted_ok=True)
             return not isinstance(mutated_attr, variables.DeletedVariable)
 
+        # TODO(guilhermeleobas): This can trigger a side effect
         return key in self.value.__dict__
 
-    def get_source_by_walking_mro(self, name: str) -> DictGetItemSource:
+    def get_source_by_walking_mro(
+        self, tx: "InstructionTranslator", name: str
+    ) -> DictGetItemSource:
         assert self.cls_source is not None
 
         for idx, klass in enumerate(type(self.value).__mro__):
             if name in klass.__dict__:
-                if idx != 0:
-                    mro_source = TypeMROSource(self.cls_source)
-                    klass_source: Source = GetItemSource(mro_source, idx)
-                else:
-                    klass_source = self.cls_source
-                dict_source = TypeDictSource(klass_source)
-                out_source = DictGetItemSource(dict_source, name)
+                descriptor = klass.__dict__[name]
 
+                # Guard that intermediate MRO classes don't shadow this
+                # attribute, deduplicating by (id(klass), name) across
+                # subclasses that share the same intermediate MRO class.
+                # Safe because TYPE_MATCH guards fix the MRO, so the same
+                # id(klass) always refers to the same class object.
                 for absent_idx in range(1, idx):
-                    # Insert a guard that the name is not present in the mro hierarchy
+                    absent_klass = type(self.value).__mro__[absent_idx]
+                    cache_key = (id(absent_klass), name)
+                    if cache_key in tx.output.guarded_mro_absent_keys:
+                        continue
+                    tx.output.guarded_mro_absent_keys.add(cache_key)
                     mro_source = TypeMROSource(self.cls_source)
-                    klass_source = GetItemSource(mro_source, absent_idx)
+                    klass_source: Source = GetItemSource(mro_source, absent_idx)
                     dict_source = TypeDictSource(klass_source)
                     install_guard(
                         dict_source.make_guard(
-                            functools.partial(
-                                GuardBuilder.DICT_CONTAINS, key=name, invert=True
-                            )
+                            functools.partial(GuardBuilder.DICT_NOT_CONTAINS, key=name)
                         )
                     )
-                # Insert a guard that the name is not present in the object __dict__
+
+                # Guard that the instance __dict__ does not shadow the
+                # class attribute.  Skipped for data descriptors (those
+                # with __set__, e.g. property) because Python gives data
+                # descriptors priority over instance __dict__ in attribute
+                # lookup — the instance dict can only be populated by
+                # directly writing to obj.__dict__, not via setattr.
                 if (
                     self.source
                     and hasattr(self.value, "__dict__")
                     and name not in self.value.__dict__
+                    and not hasattr(descriptor, "__set__")
                 ):
                     install_guard(
                         self.source.make_guard(
@@ -1591,6 +2085,26 @@ def get_source_by_walking_mro(self, name: str) -> DictGetItemSource:
                             )
                         )
                     )
+
+                # Reuse the source if we've already resolved the same
+                # descriptor object for the same attribute name (e.g. same
+                # property reached via different subclasses) to avoid
+                # redundant ID_MATCH guards.  We include name in the key
+                # because distinct attributes can point to the same object
+                # (e.g. a = b = some_obj, or interned small integers).
+                cache_key = (id(descriptor), name)
+                cache = tx.output.mro_source_cache
+                if cache_key in cache:
+                    return cache[cache_key]
+
+                if idx != 0:
+                    mro_source = TypeMROSource(self.cls_source)
+                    klass_source = GetItemSource(mro_source, idx)
+                else:
+                    klass_source = self.cls_source
+                dict_source = TypeDictSource(klass_source)
+                out_source = DictGetItemSource(dict_source, name)
+                cache[cache_key] = out_source
                 return out_source
 
         unimplemented(
@@ -1604,8 +2118,6 @@ def get_source_by_walking_mro(self, name: str) -> DictGetItemSource:
         )
 
     def var_getattr(self, tx: "InstructionTranslator", name: str) -> VariableTracker:
-        from . import ConstantVariable
-
         source: Source | None = AttrSource(self.source, name) if self.source else None
 
         if self._object_has_getattribute:
@@ -1620,9 +2132,8 @@ def var_getattr(self, tx: "InstructionTranslator", name: str) -> VariableTracker
                 return variables.UserMethodVariable(
                     getattribute_fn,
                     self,
-                    # pyrefly: ignore[unbound-name]
                     source=new_source,
-                ).call_function(tx, [ConstantVariable.create(name)], {})
+                ).call_function(tx, [VariableTracker.build(tx, name)], {})
             except ObservedAttributeError:
                 # Pass through to __getattr__ if __getattribute__ fails
                 handle_observed_exception(tx)
@@ -1634,14 +2145,15 @@ def var_getattr(self, tx: "InstructionTranslator", name: str) -> VariableTracker
                     AttributeError,
                     tx,
                     args=[
-                        f"'{type(self.value).__name__}' object has no attribute '{name}'"
+                        f"'{type(self.value).__name__}' object has no attribute '{name}'",
                     ],
                 )
             return result
 
         if name == "__dict__":
-            options_dict = {"source": source}
-            return variables.GetAttrVariable(self, name, None, **options_dict)
+            if not hasattr(self.value, "__dict__"):
+                raise_observed_exception(AttributeError, tx)
+            return self.get_dict_vt(tx)
 
         # TODO(anijain2305) - Investigate if we need specialization for more
         # dunder attrs. inspect.getattr_static does not return correct value for
@@ -1652,291 +2164,320 @@ def var_getattr(self, tx: "InstructionTranslator", name: str) -> VariableTracker
                 cls_source = self.cls_source
             else:
                 cls_source = source
-            options = {"source": cls_source}
-            return UserDefinedClassVariable(type(self.value), **options)
-
-        try:
-            subobj = self._getattr_static(name)
-        except AttributeError:
-            subobj = NO_SUCH_SUBOBJ
-            getattr_fn = self._check_for_getattr()
-            if isinstance(getattr_fn, types.FunctionType):
-                # Dynamo is going to trace the __getattr__ function with
-                # args=name. Set the source accordingly.
-                if (
-                    getattr_fn is unpatched_nn_module_getattr
-                    and isinstance(self, variables.UnspecializedNNModuleVariable)
-                    # prevent against overwriting of params/buffers/submodules
-                    and istype(self.value._parameters, dict)  # type: ignore[attr-defined]
-                    and istype(self.value._buffers, dict)  # type: ignore[attr-defined]
-                    and istype(self.value._modules, dict)  # type: ignore[attr-defined]
-                ):
-                    # Manually trace out the nn module __getattr__ to avoid large compilation latency.
-                    out = self.manually_trace_nn_module_getattr(tx, name)
-                else:
-                    new_source = None
-                    if self.source:
-                        new_source = AttrSource(self.source, "__getattr__")
-                    out = variables.UserMethodVariable(
-                        getattr_fn, self, source=new_source
-                    ).call_function(tx, [ConstantVariable.create(name)], {})
-
-                if self.source and getattr_fn is torch.nn.Module.__getattr__:
-                    if isinstance(
-                        out,
-                        (
-                            variables.UnspecializedNNModuleVariable,
-                            variables.NNModuleVariable,
-                        ),
-                    ):
-                        # nn_module_stack source is BC surface area. Ensure that
-                        # mod._modules["linear"] is reflected as mod.linear for
-                        # nn_module_stack.
-                        out.set_nn_module_stack_source(  # type: ignore[attr-defined]
-                            AttrSource(self.get_nn_module_stack_source(), name)  # type: ignore[attr-defined]
-                        )
-                return out
-
-            elif getattr_fn is not None:
-                unimplemented(
-                    gb_type="User-defined object with non-function __getattr__",
-                    context=f"object={self.value}, name={name}, getattr_fn={getattr_fn}",
-                    explanation=f"Found a non-function __getattr__ {getattr_fn} from a user-defined object {self.value} "
-                    f" when attempting to getattr `{name}`",
-                    hints=[
-                        "Ensure the object's __getattr__ is a function type.",
-                    ],
-                )
+            return VariableTracker.build(tx, type(self.value), cls_source)
 
         from ..mutation_guard import unpatched_nn_module_init
 
-        if subobj is torch.nn.Module.__init__:
-            subobj = unpatched_nn_module_init
+        # ---- CPython attribute lookup algorithm ----
+        # Mirror object.__getattribute__ (PyObject_GenericGetAttr):
+        #   1. type_attr = lookup name in type(obj).__mro__
+        #   2. if type_attr is a DATA descriptor → invoke it
+        #   3. if name in obj.__dict__ → return as-is (no descriptor invocation)
+        #   4. if type_attr is a non-data descriptor → invoke it
+        #   5. if type_attr is a plain class variable → return it
+        #   6. __getattr__ fallback
+        #   7. raise AttributeError
+        #
+        # Between steps 5 and 6, we also handle objects with custom storage
+        # that aren't visible via the MRO walk or instance __dict__ (step 5b).
+        #
+        # Step 1: Single MRO walk on the type (cached).
+        type_attr = self.lookup_class_mro_attr(name)
+
+        # Dynamo patches nn.Module.__init__ at import time to inject tracing
+        # hooks.  Undo that here so the unpatched original is traced instead.
+        if type_attr is torch.nn.Module.__init__:
+            type_attr = unpatched_nn_module_init
+
+        # Step 2: Data descriptors on the type take priority over instance dict.
+        if type_attr is not NO_SUCH_SUBOBJ and is_data_descriptor(type_attr):
+            return self.resolve_data_descriptor(tx, name, type_attr, source)
+
+        # Step 3: Instance __dict__ — return as-is, no descriptor invocation.
+        # TODO(guilhermeleobas): step 3 should look into dict_vt and not self.value.__dict__
+        # as the object could have mutated an attribute via setattr
+        if hasattr(self.value, "__dict__") and name in self.value.__dict__:
+            subobj = self.value.__dict__[name]
+            source = self.maybe_wrap_nn_module_source_for_instance(tx, name, source)
+            return VariableTracker.build(tx, subobj, source)
+
+        # Step 4-5: Non-data descriptor or plain class attribute.
+        if type_attr is not NO_SUCH_SUBOBJ:
+            return self.resolve_type_attr(tx, name, type_attr, source)
+
+        # Step 5b: Dynamic fallback for attributes that exist on the live
+        # object but aren't visible to the static MRO walk or instance
+        # __dict__ check above.  This covers objects with custom storage
+        # backends (e.g. threading.local uses a per-thread dict not
+        # accessible via obj.__dict__) and C extensions that store data
+        # outside the normal Python object layout.
+        #
+        # This is NOT the same as the C-level data descriptor fallback in
+        # resolve_data_descriptor (step 2): that handles descriptors found
+        # on the type MRO (like member_descriptor for __slots__), while this
+        # handles attributes that aren't on the type MRO at all.
+        #
+        # Only safe when the class doesn't override __getattribute__,
+        # otherwise we'd run arbitrary user code.
+        if not self._object_has_getattribute:
+            try:
+                resolved = type(self.value).__getattribute__(self.value, name)
+                source = self.maybe_wrap_nn_module_source_for_instance(tx, name, source)
+                return VariableTracker.build(tx, resolved, source)
+            except AttributeError:
+                pass
 
-        # Check if its already saved, avoids inspect getattr_static call
-        if name in self._subobj_from_class:
-            subobj_from_class = self._subobj_from_class[name]
-        else:
-            subobj_from_class = inspect.getattr_static(
-                self.value.__class__, name, NO_SUCH_SUBOBJ
+        # Step 6: __getattr__ fallback.
+        getattr_fn = self._check_for_getattr()
+        if isinstance(getattr_fn, types.FunctionType):
+            if (
+                getattr_fn is unpatched_nn_module_getattr
+                and isinstance(self, variables.UnspecializedNNModuleVariable)
+                and istype(self.value._parameters, dict)  # type: ignore[attr-defined]
+                and istype(self.value._buffers, dict)  # type: ignore[attr-defined]
+                and istype(self.value._modules, dict)  # type: ignore[attr-defined]
+            ):
+                out = self.manually_trace_nn_module_getattr(tx, name)
+            else:
+                new_source = None
+                if self.source:
+                    new_source = AttrSource(self.source, "__getattr__")
+                out = variables.UserMethodVariable(
+                    getattr_fn, self, source=new_source
+                ).call_function(tx, [variables.ConstantVariable.create(name)], {})
+
+            if self.source and getattr_fn is torch.nn.Module.__getattr__:
+                if isinstance(
+                    out,
+                    (
+                        variables.UnspecializedNNModuleVariable,
+                        variables.NNModuleVariable,
+                    ),
+                ):
+                    out.set_nn_module_stack_source(  # type: ignore[attr-defined]
+                        AttrSource(self.get_nn_module_stack_source(), name)  # type: ignore[attr-defined]
+                    )
+            return out
+
+        elif getattr_fn is not None:
+            unimplemented(
+                gb_type="User-defined object with non-function __getattr__",
+                context=f"object={self.value}, name={name}, getattr_fn={getattr_fn}",
+                explanation=f"Found a non-function __getattr__ {getattr_fn} from a user-defined object {self.value} "
+                f" when attempting to getattr `{name}`",
+                hints=[
+                    "Ensure the object's __getattr__ is a function type.",
+                ],
             )
-            self._subobj_from_class[name] = subobj_from_class
 
-        is_accessible_from_type_mro = (
-            subobj_from_class is subobj
-            and self.cls_source is not None
-            and self.source is not None
-            and hasattr(self.value, "__dict__")
-            and name not in self.value.__dict__
+        # Step 7: AttributeError.
+        raise_observed_exception(
+            AttributeError,
+            tx,
+            args=[f"'{type(self.value).__name__}' object has no attribute '{name}'"],
         )
 
-        if isinstance(subobj, property):
+    def resolve_data_descriptor(
+        self,
+        tx: "InstructionTranslator",
+        name: str,
+        type_attr: object,
+        source: Source | None,
+    ) -> VariableTracker:
+        """Handle data descriptors found on the type MRO (property, _tuplegetter, etc.)."""
+        if isinstance(type_attr, property) and not self._is_c_defined_property(
+            type_attr
+        ):
+            # Python property — trace fget directly.
             if self.source:
-                # Read the class attribute to reach the property
-                source = AttrSource(self.get_source_by_walking_mro(name), "fget")
+                source = AttrSource(self.get_source_by_walking_mro(tx, name), "fget")
             fget_vt = VariableTracker.build(
-                tx, subobj.fget, source=source, realize=True
+                tx, type_attr.fget, source=source, realize=True
             )
             return fget_vt.call_function(tx, [self], {})
-        elif isinstance(subobj, _collections._tuplegetter):
-            # namedtuple fields are represented by _tuplegetter, and here we
-            # emulate its `__get__`, which is implemented in C.
-            _, (idx, _) = subobj.__reduce__()
-            # Don't go through the `__getitem__` method anymore, see
-            # https://github.com/python/cpython/blob/470941782f74288823b445120f6383914b659f23/Modules/_collectionsmodule.c#L2690
-            assert isinstance(self, UserDefinedTupleVariable)
-            return self._tuple_vt.items[idx]  # type: ignore[union-attr]
-        elif isinstance(subobj, staticmethod):
-            # Safe because `staticmethod.__get__` basically won't trigger user
-            # code and just returns the underlying `__func__`:
-            # https://github.com/python/cpython/blob/3.11/Objects/funcobject.c#L1088-L1100
-            if is_accessible_from_type_mro:
-                # Accessing from __dict__ does not resolve the descriptor, it
-                # returns a staticmethod object, so access the __func__
-                # attribute to get to the actual function.
-                source = AttrSource(self.get_source_by_walking_mro(name), "__func__")
-            func = subobj.__get__(self.value)
+
+        get_fn = inspect.getattr_static(type(type_attr), "__get__", None)
+        if isinstance(get_fn, types.FunctionType):
+            # User-defined data descriptor with a Python __get__.
+            return self.invoke_descriptor_get(tx, name, type_attr, source)
+
+        # C-level data descriptor (property with C fget, member/getset
+        # descriptors, Cython attrs, etc.) — resolve via
+        # object.__getattribute__ which is side-effect free.
+        # Uninitialized slots raise AttributeError which must be surfaced
+        # as ObservedAttributeError so dynamo's try/except tracing works.
+        try:
+            resolved = type(self.value).__getattribute__(self.value, name)
+        except AttributeError:
+            raise_observed_exception(
+                AttributeError,
+                tx,
+                args=[
+                    f"'{type(self.value).__name__}' object has no attribute '{name}'"
+                ],
+            )
+        return VariableTracker.build(tx, resolved, source)
+
+    def resolve_type_attr(
+        self,
+        tx: "InstructionTranslator",
+        name: str,
+        type_attr: object,
+        source: Source | None,
+    ) -> VariableTracker:
+        """Handle non-data descriptors and plain class attributes from the type MRO."""
+        from ..mutation_guard import unpatched_nn_module_init
+
+        if (
+            type_attr is unpatched_nn_module_init
+            or type_attr is torch.nn.Module.__init__
+        ):
+            type_attr = unpatched_nn_module_init
+
+        can_use_mro_source = self.cls_source is not None and self.source is not None
+
+        if isinstance(type_attr, staticmethod):
+            # type_attr is the raw staticmethod wrapper from cls.__dict__
+            # (not the unwrapped function).  We call __get__ to unwrap it,
+            # but the *source* must go through __func__ on the descriptor
+            # (not the resolved function) because the guard needs to watch
+            # the descriptor object in the class dict, not the result.
+            if can_use_mro_source:
+                source = AttrSource(
+                    self.get_source_by_walking_mro(tx, name), "__func__"
+                )
+            func = type_attr.__get__(self.value)
             return VariableTracker.build(tx, func, source)
-        elif isinstance(subobj, classmethod):
+        elif isinstance(type_attr, classmethod):
             source_fn = None
-            if is_accessible_from_type_mro:
-                # Accessing from __dict__ does not resolve the descriptor, it
-                # returns a classmethod object, so access the __func__
-                # attribute to get to the actual function.
-                source_fn = AttrSource(self.get_source_by_walking_mro(name), "__func__")  # type: ignore[assignment]
+            if can_use_mro_source:
+                source_fn = AttrSource(
+                    self.get_source_by_walking_mro(tx, name), "__func__"
+                )  # type: ignore[assignment]
             return variables.UserMethodVariable(
-                subobj.__func__,
+                type_attr.__func__,
                 self.var_getattr(tx, "__class__"),
                 source_fn=source_fn,
                 source=source,
             )
-        elif isinstance(subobj, types.ClassMethodDescriptorType):
-            # e.g.: inspect.getattr_static({}, "fromkeys")
-            func = subobj.__get__(self.value, None)
+        elif isinstance(type_attr, types.ClassMethodDescriptorType):
+            func = type_attr.__get__(self.value, None)
             return VariableTracker.build(tx, func, source)
-        elif is_lru_cache_wrapped_function(subobj):
-            # getattr_static returns the lru_wrapped function, and we cannot
-            # extract the underlying method from the wrapped function. To handle
-            # it, manually create a wrapped user method vt.
+        elif is_lru_cache_wrapped_function(type_attr):
             return variables.WrapperUserMethodVariable(
-                subobj, "__wrapped__", self, source=source
+                type_attr, "__wrapped__", self, source=source
             )
-        elif inspect.getattr_static(
-            type(subobj), "__get__", NO_SUCH_SUBOBJ
-        ) is not NO_SUCH_SUBOBJ and not is_wrapper_or_member_descriptor(
-            type(subobj).__get__  # type: ignore[attr-defined]
-        ):
-            # Emulate https://github.com/python/cpython/blob/3.11/Objects/object.c#L1271-L1285
-            #
-            # Attribute has a __get__ method. Create a user defined object vt
-            # for the subobj, and then trace the __get__ method.
-            descriptor_source = None
-            descriptor_get_source = None
-            if self.cls_source:
-                # To access the method descriptor from the udf object w/o using
-                # inspect.getattr_static, we can look into the class mro
-                descriptor_source = self.get_source_by_walking_mro(name)
-                descriptor_get_source = AttrSource(
-                    TypeSource(descriptor_source), "__get__"
-                )
-                descriptor_var = VariableTracker.build(tx, subobj, descriptor_source)
-            else:
-                # Sourceless Builder does not support user defined objects
-                descriptor_var = UserDefinedObjectVariable(subobj)
-
-            # The arguments of the __get__ function are (self, instance, owner)
-            # self - descriptor_var
-            # instance - instance of the class, represented by self here
-            # owner - class object
-            owner_var = UserDefinedClassVariable(type(self.value))
+        elif isinstance(type_attr, types.FunctionType):
+            while hasattr(type_attr, "_torchdynamo_inline"):
+                type_attr = type_attr._torchdynamo_inline  # type: ignore[union-attr]
+                source = AttrSource(source, "_torchdynamo_inline") if source else None
+            # Function on the type MRO + not in instance dict → bound method.
+            var_source = None
+            if can_use_mro_source:
+                var_source = self.get_source_by_walking_mro(tx, name)
             return variables.UserMethodVariable(
-                subobj.__get__.__func__,  # type: ignore[attr-defined]
-                descriptor_var,
-                source=descriptor_get_source,
-            ).call_function(tx, [self, owner_var], {})
-        elif isinstance(subobj, types.FunctionType) or (
-            isinstance(subobj, types.MethodType)
-            and isinstance(self.value, torch.nn.Module)
+                type_attr, self, source_fn=var_source, source=source
+            )
+        # Check for a Python-level __get__ (non-data descriptor with traceable __get__).
+        get_fn = inspect.getattr_static(type(type_attr), "__get__", None)
+        if isinstance(get_fn, types.FunctionType):
+            return self.invoke_descriptor_get(tx, name, type_attr, source)
+
+        # C-level non-data descriptors / opaque callables — defer to runtime.
+        # MethodDescriptorType: e.g. list.append (PyMethodDef)
+        # WrapperDescriptorType: e.g. list.__add__ (slot wrappers)
+        # MethodWrapperType: e.g. [].__add__ (bound slot wrappers)
+        #
+        # Exception: if the descriptor has a registered polyfill, return the
+        # polyfill as a bound method so Dynamo can trace through it.
+        if (
+            isinstance(
+                type_attr,
+                (
+                    types.MethodDescriptorType,
+                    types.WrapperDescriptorType,
+                    types.MethodWrapperType,
+                ),
+            )
+            or torch._C._dynamo.utils.is_instancemethod(type_attr)  # type: ignore[attr-defined]
+            or is_cython_function(type_attr)
         ):
-            # Since we get subobj via self._getattr_static, which may not trigger dynamic lookup.
-            # Static lookup can't tell us it's a method or function correctly,
-            # so we trigger dynamic lookup here to get the correct type.
-            dynamic_subobj = getattr(self.value, name)
-
-            while dynamic_subobj is subobj and hasattr(subobj, "_torchdynamo_inline"):
-                subobj = subobj._torchdynamo_inline
-                dynamic_subobj = subobj
-                source = AttrSource(source, "_torchdynamo_inline") if source else None
-
-            if isinstance(subobj, types.MethodType):
-                if dynamic_subobj.__self__ is not self.value:
-                    if not isinstance(dynamic_subobj.__func__, types.FunctionType):
-                        unimplemented(
-                            gb_type="User-defined object method with non-function __func__",
-                            context=f"object={self.value}, name={name}, method={dynamic_subobj}, "
-                            f"method.__self__={dynamic_subobj.__self__}, method.__func__={dynamic_subobj.__func__}",
-                            explanation=f"Method {dynamic_subobj} (name={name}) of user-defined object {self.value} has a "
-                            f"__func__ ({dynamic_subobj.__func__}) that is not a function type.",
-                            hints=[
-                                "Ensure that the method's __func__ is a function type.",
-                            ],
-                        )
-
-                    # Use the __self__ attribute of the method to find the
-                    # source of the new self object.
-                    self_source = None
-                    if source is not None:
-                        self_source = AttrSource(source, "__self__")
-                    object_vt = VariableTracker.build(
-                        tx, dynamic_subobj.__self__, self_source
-                    )
-
-                    return variables.UserMethodVariable(
-                        dynamic_subobj.__func__,
-                        object_vt,
-                    )
-                func = subobj.__func__
-            else:
-                assert isinstance(subobj, types.FunctionType)
-                func = subobj
+            from .. import trace_rules
+
+            if trace_rules.is_polyfilled_callable(type_attr):  # type: ignore[arg-type]
+                from .functions import PolyfilledFunctionVariable
+
+                polyfill_handlers = PolyfilledFunctionVariable._get_polyfill_handlers()
+                wrapped: Any = polyfill_handlers.get(type_attr)  # type: ignore[arg-type]
+                if wrapped is not None:
+                    traceable_fn = wrapped.__torch_dynamo_polyfill__
+                    return variables.UserMethodVariable(traceable_fn, self)
+            return variables.GetAttrVariable(self, name, type(type_attr), source=source)
+
+        # Plain class variable (or MethodType, C-level non-data descriptor
+        # without __get__, etc.).
+        if can_use_mro_source:
+            source = self.get_source_by_walking_mro(tx, name)
+        elif not source and self.cls_source is not None:
+            source = AttrSource(self.cls_source, name)
+        return VariableTracker.build(tx, type_attr, source)
+
+    def invoke_descriptor_get(
+        self,
+        tx: "InstructionTranslator",
+        name: str,
+        descriptor: object,
+        source: Source | None,
+    ) -> VariableTracker:
+        """Trace a descriptor's __get__(instance, owner) call."""
+        descriptor_source = None
+        descriptor_get_source = None
+        if self.cls_source:
+            descriptor_source = self.get_source_by_walking_mro(tx, name)
+            descriptor_get_source = AttrSource(TypeSource(descriptor_source), "__get__")
+            descriptor_var = VariableTracker.build(tx, descriptor, descriptor_source)
+        else:
+            descriptor_var = UserDefinedObjectVariable(descriptor)
 
-            if inspect.ismethod(dynamic_subobj):
-                var_source = None
-                if is_accessible_from_type_mro:
-                    var_source = self.get_source_by_walking_mro(name)
-                return variables.UserMethodVariable(
-                    func, self, source_fn=var_source, source=source
-                )
-            elif inspect.isfunction(dynamic_subobj):
-                return VariableTracker.build(tx, func, source)
+        owner_var = UserDefinedClassVariable(type(self.value))
+        return variables.UserMethodVariable(
+            descriptor.__get__.__func__,  # type: ignore[union-attr]
+            descriptor_var,
+            source=descriptor_get_source,
+        ).call_function(tx, [self, owner_var], {})
 
+    def maybe_wrap_nn_module_source_for_instance(
+        self,
+        tx: "InstructionTranslator",
+        name: str,
+        source: Source | None,
+    ) -> Source | None:
+        """Wrap source for nn.Module instance dict attribute access if needed."""
         if (
-            # wrap the source only if inline_inbuilt_nn_modules is set or fsdp modules. This is a temporary solution to
-            # keep Dynamo behavior compatible with no inlining, as there will be some delay to turn on the flag in
-            # fbcode.
-            (
-                torch._dynamo.config.inline_inbuilt_nn_modules
-                or isinstance(self, variables.FSDPManagedNNModuleVariable)
-            )
-            and source
+            source
             and isinstance(self, variables.UnspecializedNNModuleVariable)
-            # export has some awkwardness around specialized and unspecialized modules. Skip wrapping source for export
-            # usecase for now.
             and (not tx.output.export or torch._dynamo.config.install_free_tensors)
         ):
-            # Recalculate source for params/buffers
             if name in ("_buffers", "_parameters"):
                 assert self.source is not None
                 source = UnspecializedParamBufferSource(self.source, name)
             source = self._wrap_source(source)
-
-        if subobj is not NO_SUCH_SUBOBJ:
-            if (
-                is_wrapper_or_member_descriptor(subobj)
-                or torch._C._dynamo.utils.is_instancemethod(subobj)  # type: ignore[attr-defined]
-                or is_cython_function(subobj)
-            ):
-                options = {"source": source}
-                return variables.GetAttrVariable(self, name, None, **options)
-            if source:
-                if is_accessible_from_type_mro:
-                    source = self.get_source_by_walking_mro(name)
-
-                return variables.LazyVariableTracker.create(subobj, source)
-            else:
-                # Check if the subobj is accessible from the class itself. If the class source is known, we can create a
-                # sourceful variable tracker.
-                if self.cls_source is not None:
-                    subobj_from_class = inspect.getattr_static(
-                        self.value.__class__, name, NO_SUCH_SUBOBJ
-                    )
-                    if subobj_from_class is subobj:
-                        src_from_class = AttrSource(self.cls_source, name)
-                        return variables.LazyVariableTracker.create(
-                            subobj_from_class, src_from_class
-                        )
-
-                return VariableTracker.build(tx, subobj)
-
-        # Earlier we were returning GetAttrVariable but its incorrect. In absence of attr, Python raises AttributeError.
-        raise_observed_exception(
-            AttributeError,
-            tx,
-            args=[f"'{type(self.value).__name__}' object has no attribute '{name}'"],
-        )
+        return source
 
     def call_obj_hasattr(
         self, tx: "InstructionTranslator", name: str
     ) -> "ConstantVariable":
         if self.source:
             install_guard(
-                AttrSource(self.source, name).make_guard(GuardBuilder.HASATTR)
+                self.source.make_guard(
+                    functools.partial(GuardBuilder.HASATTR, attr=name)
+                )
             )
 
         try:
             var_vt = self.var_getattr(tx, name)
-            return variables.ConstantVariable.create(
-                not isinstance(var_vt, variables.DeletedVariable)
+            return VariableTracker.build(
+                tx, not isinstance(var_vt, variables.DeletedVariable)
             )
         except ObservedAttributeError:
             handle_observed_exception(tx)
@@ -1944,13 +2485,22 @@ def call_obj_hasattr(
 
     def is_python_hashable(self) -> bool:
         raise_on_overridden_hash(self.value, self)
+        if self._base_vt is not None:
+            return self._base_vt.is_python_hashable()
         return True
 
     def get_python_hash(self) -> int:
-        # default hash
+        if self._base_vt is not None:
+            return self._base_vt.get_python_hash()
         return hash(self.value)
 
     def is_python_equal(self, other: object) -> bool:
+        if (
+            isinstance(other, VariableTracker)
+            and self.is_python_constant()
+            and other.is_python_constant()
+        ):
+            return self.as_python_constant() == other.as_python_constant()
         # id check
         if not isinstance(other, UserDefinedVariable):
             return False
@@ -2053,38 +2603,100 @@ def call_tree_map_branch(
             tree_map_kwargs,
         )
 
+    def call_tree_map_with_path_branch(
+        self,
+        tx: "InstructionTranslator",
+        tree_map_fn: "variables.functions.UserFunctionVariable",
+        map_fn: "VariableTracker",
+        rest: "collections.abc.Sequence[VariableTracker]",
+        tree_map_kwargs: "dict[str, VariableTracker]",
+        keypath: "tuple[Any, ...]",
+    ) -> "VariableTracker":
+        """Emulate tree_map_with_path behavior for user-defined objects.
 
-class FrozenDataClassVariable(UserDefinedObjectVariable):
-    @staticmethod
-    def create(
-        tx: "InstructionTranslator", value: object, source: Source
-    ) -> "FrozenDataClassVariable":
-        from dataclasses import fields
+        Same logic as call_tree_map_branch but passes keypath to the map function.
+        """
+        tree_map_module = tree_map_fn.get_module()
+        is_optree = tree_map_module.startswith("optree")
 
-        assert is_frozen_dataclass(value)
+        if is_optree:
+            try:
+                import optree
+                from optree.registry import _NODETYPE_REGISTRY
+
+                is_registered = (
+                    self.value_type in _NODETYPE_REGISTRY
+                    or optree.is_namedtuple_class(self.value_type)
+                    or optree.is_structseq_class(self.value_type)
+                )
 
-        field_map = {}
-        for field in fields(value):  # type: ignore[arg-type]
-            if hasattr(value, field.name):
-                field_map[field.name] = VariableTracker.build(
+                if not is_registered:
+                    namespace_var = tree_map_kwargs.get("namespace")
+                    if namespace_var is not None:
+                        try:
+                            namespace = namespace_var.as_python_constant()
+                            is_registered = (
+                                namespace,
+                                self.value_type,
+                            ) in _NODETYPE_REGISTRY
+                        except NotImplementedError:
+                            return self._tree_map_with_path_fallback(
+                                tx,
+                                tree_map_fn,
+                                map_fn,
+                                rest,
+                                tree_map_kwargs,
+                                keypath,
+                            )
+            except ImportError:
+                return self._tree_map_with_path_fallback(
                     tx,
-                    getattr(value, field.name),
-                    source and AttrSource(source, field.name),
+                    tree_map_fn,
+                    map_fn,
+                    rest,
+                    tree_map_kwargs,
+                    keypath,
                 )
+        else:
+            import torch.utils._pytree as pytree
+
+            is_registered = (
+                self.value_type in pytree.SUPPORTED_NODES
+                or pytree.is_namedtuple_class(self.value_type)
+                or pytree.is_structseq_class(self.value_type)
+            )
 
-        return FrozenDataClassVariable(value, fields=field_map, source=source)
+        if not is_registered:
+            keypath_var = variables.TupleVariable(
+                [VariableTracker.build(tx, k) for k in keypath]
+            )
+            return map_fn.call_function(tx, [keypath_var, self, *rest], {})
 
-    def __init__(
-        self, value: object, fields: dict[str, Any] | None = None, **kwargs: Any
-    ) -> None:
-        super().__init__(value, **kwargs)
-        if fields is None:
-            fields = {}
-        self.fields = fields
+        return self._tree_map_with_path_fallback(
+            tx,
+            tree_map_fn,
+            map_fn,
+            rest,
+            tree_map_kwargs,
+            keypath,
+        )
+
+
+class FrozenDataClassVariable(UserDefinedObjectVariable):
+    """Frozen dataclass variable for as_proxy/as_python_constant/hashability.
+
+    Construction is handled by the generic polyfill path (tracing through
+    the auto-generated __init__). Field values are retrieved dynamically
+    via var_getattr using InstructionTranslator.current_tx().
+    """
+
+    def _get_field_vt(self, field_name: str) -> VariableTracker:
+        from torch._dynamo.symbolic_convert import InstructionTranslator
+
+        tx = InstructionTranslator.current_tx()
+        return self.var_getattr(tx, field_name)
 
     def as_python_constant(self) -> object:
-        # NOTE: this is an intentionally limited version of
-        # `as_python_constant` for `nonstrict_trace` implementation.
         from dataclasses import fields
 
         import torch.utils._pytree as pytree
@@ -2092,154 +2704,74 @@ def as_python_constant(self) -> object:
         if not istype(
             self.value, (pytree.TreeSpec, pytree.LeafSpec, pytree.ConstantNode)
         ):
-            # TODO loosen this restriction and fix `as_proxy`.
             raise NotImplementedError(
                 "currently can't reconstruct arbitrary frozen dataclass instances"
             )
 
-        # LeafSpec is deprecated, use treespec_leaf() instead
         if istype(self.value, pytree.LeafSpec):
             return pytree.treespec_leaf()
 
-        args = []
-        kwargs = {}
+        args: list[object] = []
+        kwargs: dict[str, object] = {}
         for field in fields(self.value):  # type: ignore[arg-type]
             if field.init:
-                data = self.fields[field.name].as_python_constant()
+                data = self._get_field_vt(field.name).as_python_constant()
                 if getattr(field, "kw_only", False):
                     kwargs[field.name] = data
                 else:
                     args.append(data)
 
-        # This is safe because we know the TreeSpec classes constructors don't
-        # have external side effects.
-        ctor = self.python_type()
-        return ctor(*args, **kwargs)
+        return self.python_type()(*args, **kwargs)
 
     def as_proxy(self) -> object:
         from dataclasses import fields
 
-        args = []
-        kwargs = {}
+        args: list[object] = []
+        kwargs: dict[str, object] = {}
         for field in fields(self.value):  # type: ignore[arg-type]
-            proxy = self.fields[field.name].as_proxy()
+            proxy = self._get_field_vt(field.name).as_proxy()
             if hasattr(field, "kw_only") and field.kw_only:
                 kwargs[field.name] = proxy
             else:
                 args.append(proxy)
 
-        # TODO this isn't really safe, because
-        # 1. it could invoke a user defined `__post_init__`.
-        # 2. it could invoke a user defined `__init__` if the class _subclasses_
-        #    a frozen dataclass.
-        # Either of the above could end up mutating external state.
-        ctor = self.python_type()
-        return ctor(*args, **kwargs)
+        return self.python_type()(*args, **kwargs)
 
     def reconstruct(self, codegen: "PyCodegen") -> None:
-        from dataclasses import fields
-
-        # Handle specific pytree classes
-        import torch.utils._pytree as pytree
-
-        if isinstance(self.value, pytree.TreeSpec) and self.value.is_leaf():
-            # Create a new LeafSpec instance by calling the constructor
-            codegen.add_push_null(
-                lambda: codegen.load_import_from("torch.utils._pytree", "LeafSpec")
-            )
-            codegen.extend_output(create_call_function(0, False))
+        if self.source is not None:
+            codegen(self.source)
             return
-
-        # For general frozen dataclasses, reconstruct by calling the constructor
-        # with the field values as arguments
-        dataclass_cls = self.python_type()
-
-        if hasattr(dataclass_cls, "__post_init__"):
-            unimplemented(
-                gb_type="Frozen dataclass with __post_init__",
-                context=f"dataclass={dataclass_cls.__name__}",
-                explanation="Cannot reconstruct frozen dataclass with __post_init__ method, "
-                "as it may have side effects that would be incorrectly replayed.",
-                hints=[
-                    "Remove the __post_init__ method from the frozen dataclass.",
-                    *graph_break_hints.SUPPORTABLE,
-                ],
-            )
-
-        # Collect positional and keyword-only arguments
-        pos_args = []
-        # pyrefly: ignore [implicit-any]
-        kw_args = []
-        for field in fields(dataclass_cls):
-            if not field.init:
-                continue
-            field_vt = self.fields.get(field.name)
-            if field_vt is None:
-                unimplemented(
-                    gb_type="Frozen dataclass with missing field",
-                    context=f"dataclass={dataclass_cls.__name__}, field={field.name}",
-                    explanation=f"Cannot reconstruct frozen dataclass: field '{field.name}' "
-                    "was not tracked during tracing.",
-                    hints=[*graph_break_hints.SUPPORTABLE],
-                )
-            if getattr(field, "kw_only", False):
-                kw_args.append((field.name, field_vt))
-            else:
-                pos_args.append(field_vt)
-
-        # Load the dataclass constructor
-        codegen.add_push_null(
-            lambda: codegen.append_output(
-                codegen.create_load_const_unchecked(dataclass_cls)
-            )
+        codegen.append_output(
+            codegen.create_load_const_unchecked(self.as_python_constant())
         )
-        # Reconstruct all arguments
-        for arg_vt in pos_args:
-            codegen(arg_vt)
-        for _, arg_vt in kw_args:
-            codegen(arg_vt)
-        # Call the constructor
-        total_args = len(pos_args) + len(kw_args)
-        if kw_args:
-            kw_names = tuple(name for name, _ in kw_args)
-            codegen.extend_output(
-                codegen.create_call_function_kw(total_args, kw_names, push_null=False)
-            )
-        else:
-            codegen.extend_output(create_call_function(total_args, False))
-
-    # NB: This is called during __init__ for a frozen dataclass
-    # use this to accumulate the most up-to-date field values
-    def method_setattr_standard(
-        self,
-        tx: "InstructionTranslator",
-        name: VariableTracker,
-        value: VariableTracker,
-        directly_update_dict: bool = False,
-    ) -> VariableTracker:
-        self.fields[name.as_python_constant()] = value
-        return super().method_setattr_standard(tx, name, value, directly_update_dict)
 
     def __repr__(self) -> str:
         return f"{self.__class__.__name__}({self.value_type.__name__})"
 
     def is_python_hashable(self) -> Literal[True]:
-        # TODO - Check corner cases like eq=False, hash=False etc
         return True
 
     def get_python_hash(self) -> int:
-        return hash(tuple(arg.get_python_hash() for arg in self.fields.values()))
+        from dataclasses import fields as dc_fields
+
+        return hash(
+            tuple(
+                self._get_field_vt(f.name).get_python_hash()
+                for f in dc_fields(self.value)  # type: ignore[arg-type]
+            )
+        )
 
     def is_python_equal(self, other: object) -> bool:
         if not isinstance(other, FrozenDataClassVariable):
             return False
-        is_class_same = self.python_type() is other.python_type()
-        is_field_name_same = self.fields.keys() == other.fields.keys()
-        is_field_value_same = all(
-            value_a.is_python_equal(value_b)
-            for value_a, value_b in zip(self.fields.values(), other.fields.values())
+        if self.python_type() is not other.python_type():
+            return False
+        from dataclasses import fields as dc_fields
+
+        return all(
+            self._get_field_vt(f.name).is_python_equal(other._get_field_vt(f.name))
+            for f in dc_fields(self.value)  # type: ignore[arg-type]
         )
-        return is_class_same and is_field_name_same and is_field_value_same
 
 
 class SourcelessGraphModuleVariable(UserDefinedObjectVariable):
@@ -2269,7 +2801,8 @@ def call_method(
 class UserDefinedExceptionObjectVariable(UserDefinedObjectVariable):
     def __init__(self, value: object, **kwargs: Any) -> None:
         super().__init__(value, **kwargs)
-        self.exc_vt = variables.ExceptionVariable(self.value_type, ())
+        init_args = kwargs.get("init_args", [])
+        self.exc_vt = variables.ExceptionVariable(self.value_type, init_args)
 
     @property
     def fn(self) -> Callable[..., object]:
@@ -2288,10 +2821,7 @@ def call_method(
             and inspect.ismethoddescriptor(method)
             and len(kwargs) == 0
         ):
-            self.exc_vt.args = tuple(args)
-            # pyrefly: ignore[missing-attribute]
-            self.value.args = args
-            return variables.CONSTANT_VARIABLE_NONE
+            return variables.ConstantVariable.create(None)
         elif (
             name == "__setattr__"
             and len(args) == 2
@@ -2304,13 +2834,24 @@ def call_method(
             return self.exc_vt.call_method(tx, name, args, kwargs)
         return super().call_method(tx, name, args, kwargs)
 
+    def var_getattr(self, tx: "InstructionTranslator", name: str):
+        if name in (
+            "args",
+            "__cause__",
+            "__context__",
+            "__suppress_context__",
+            "__traceback__",
+        ):
+            return self.exc_vt.var_getattr(tx, name)
+        return super().var_getattr(tx, name)
+
     @property
     def __context__(self) -> "ConstantVariable":
         # type: ignore[return-value]
         return self.exc_vt.__context__
 
     @property
-    def args(self) -> tuple[VariableTracker, ...]:
+    def args(self) -> list[VariableTracker]:
         return self.exc_vt.args
 
     def set_context(self, context: "variables.ExceptionVariable") -> None:
@@ -2324,6 +2865,9 @@ def exc_type(self) -> type[BaseException]:
     def python_stack(self) -> traceback.StackSummary | None:
         return self.exc_vt.python_stack
 
+    def debug_repr(self) -> str:
+        return self.exc_vt.debug_repr()
+
     @python_stack.setter
     def python_stack(self, value: traceback.StackSummary) -> None:
         self.exc_vt.python_stack = value
@@ -2381,6 +2925,40 @@ def var_getattr(self, tx: "InstructionTranslator", name: str) -> VariableTracker
         return super().var_getattr(tx, name)
 
 
+_CONSTANT_BASE_TYPES = (int, float, str)
+
+_constant_base_methods: dict[type, set[Any]] = {
+    t: {m for m in t.__dict__.values() if callable(m)} for t in _CONSTANT_BASE_TYPES
+}
+
+
+class UserDefinedConstantVariable(UserDefinedObjectVariable):
+    """
+    Represents user-defined objects that subclass immutable constant types
+    (int, float, str).
+
+    Uses a ConstantVariable as _base_vt for the underlying constant value.
+    """
+
+    def __init__(self, value: Any, **kwargs: Any) -> None:
+        from .constant import ConstantVariable
+
+        super().__init__(value, **kwargs)
+        for base in type(value).__mro__:
+            if base in _CONSTANT_BASE_TYPES:
+                self._base_vt = ConstantVariable.create(base(value))
+                self._base_methods = _constant_base_methods[base]
+                break
+        assert self._base_vt is not None
+
+    def as_python_constant(self) -> Any:
+        return self.value
+
+    def as_proxy(self) -> object:
+        assert self._base_vt is not None
+        return self._base_vt.as_proxy()
+
+
 class IntWrapperVariable(UserDefinedObjectVariable):
     # Dummy class to check if the object is an IntWrapper, and turn it into a
     # symint
@@ -2416,7 +2994,7 @@ def call_method(
                 assert self.idx is not None
                 tx.output.side_effects.remove_hook(self.idx)
                 self.idx = self.REMOVED
-            return variables.CONSTANT_VARIABLE_NONE
+            return variables.ConstantVariable.create(None)
         return super().call_method(tx, name, args, kwargs)
 
     def reconstruct(self, codegen: "PyCodegen") -> None:
@@ -2456,12 +3034,45 @@ def __init__(
             assert self.source is None, (
                 "dict_vt must be constructed by builder.py when source is present"
             )
-            self._dict_vt = ConstDictVariable(
-                {}, type(value), mutation_type=ValueMutationNew()
+            self._base_vt = ConstDictVariable(
+                {},
+                user_cls=(
+                    collections.OrderedDict
+                    if isinstance(value, collections.OrderedDict)
+                    else dict
+                ),
+                mutation_type=ValueMutationNew(),
             )
         else:
-            self._dict_vt = dict_vt
-        self._dict_methods = dict_methods
+            self._base_vt = dict_vt
+        self._base_methods = dict_methods
+        assert self._base_vt is not None
+
+    def sq_length(self, tx: "InstructionTranslator") -> VariableTracker:
+        # Dict implements __len__ via mp_length (mapping protocol), not
+        # sq_length (sequence protocol). Redirect so generic_len works.
+        return self.mp_length(tx)
+
+    def mp_subscript_impl(
+        self,
+        tx: "InstructionTranslator",
+        key: VariableTracker,
+    ) -> VariableTracker:
+        # dict_subscript: https://github.com/python/cpython/blob/62a6e898e01/Objects/dictobject.c#L3673-L3706
+        # TODO(follow-up): add test for unhashable/invalid key type, Counter missing key
+        method = self._maybe_get_baseclass_method("__getitem__")
+        if method in self._base_methods:
+            assert self._base_vt is not None
+            try:
+                return self._base_vt.mp_subscript_impl(tx, key)
+            except ObservedKeyError:
+                if issubclass(
+                    self.python_type(), dict
+                ) and self._maybe_get_baseclass_method("__missing__"):
+                    return self.call_method(tx, "__missing__", [key], {})
+                else:
+                    raise
+        return super().mp_subscript_impl(tx, key)
 
     def call_method(
         self,
@@ -2470,55 +3081,22 @@ def call_method(
         args: list[VariableTracker],
         kwargs: dict[str, VariableTracker],
     ) -> VariableTracker:
-        method = self._maybe_get_baseclass_method(name)
-        if method in self._dict_methods:
-            # Dict subclasses can override __missing__ to provide fallback
-            # behavior instead of raising a KeyError. This is used, for example,
-            # by collections.Counter.
+        # Dict subclasses can override __missing__ to provide fallback
+        # behavior instead of raising a KeyError. This is used, for example,
+        # by collections.Counter.
+        if (
+            name == "__getitem__"
+            and self._maybe_get_baseclass_method("__getitem__") in self._base_methods
+            and self._maybe_get_baseclass_method("__missing__")
+        ):
+            assert self._base_vt is not None
             try:
-                return self._dict_vt.call_method(tx, name, args, kwargs)
+                return self._base_vt.call_method(tx, name, args, kwargs)
             except ObservedKeyError:
-                if (
-                    name == "__getitem__"
-                    and issubclass(self.python_type(), dict)
-                    and self._maybe_get_baseclass_method("__missing__")
-                ):
-                    return self.call_method(tx, "__missing__", args, kwargs)
-                else:
-                    raise
+                handle_observed_exception(tx)
+                return self.call_method(tx, "__missing__", args, kwargs)
         return super().call_method(tx, name, args, kwargs)
 
-    def unpack_var_sequence(self, tx: "InstructionTranslator") -> list[VariableTracker]:
-        if type(self.value).__iter__ in (  # type: ignore[attr-defined]
-            dict.__iter__,
-            collections.OrderedDict.__iter__,
-        ):
-            return self._dict_vt.unpack_var_sequence(tx)
-        raise NotImplementedError
-
-    def is_underlying_vt_modified(self, side_effects: "SideEffects") -> bool:
-        return side_effects.is_modified(self._dict_vt)
-
-    @property
-    def user_cls(self) -> type[object]:
-        return self._dict_vt.user_cls
-
-    @property
-    def items(self) -> dict[VariableTracker, VariableTracker]:
-        return self._dict_vt.items
-
-    def install_dict_keys_match_guard(self) -> None:
-        return self._dict_vt.install_dict_keys_match_guard()
-
-    def install_dict_contains_guard(
-        self, tx: "InstructionTranslator", args: list[VariableTracker]
-    ) -> None:
-        return self._dict_vt.install_dict_contains_guard(tx, args)
-
-    def is_python_hashable(self) -> Literal[False]:
-        raise_on_overridden_hash(self.value, self)
-        return False
-
 
 class UserDefinedSetVariable(UserDefinedObjectVariable):
     """
@@ -2538,7 +3116,7 @@ def __init__(
         super().__init__(value, **kwargs)
 
         python_type = set if isinstance(value, set) else frozenset
-        self._set_methods = set_methods if python_type is set else frozenset_methods
+        self._base_methods = set_methods if python_type is set else frozenset_methods
 
         if set_vt is None:
             assert self.source is None, (
@@ -2546,7 +3124,7 @@ def __init__(
             )
             if python_type is set:
                 # set is initialized later
-                self._set_vt = variables.SetVariable(
+                self._base_vt = variables.SetVariable(
                     set(),
                     mutation_type=ValueMutationNew(),
                 )
@@ -2554,65 +3132,32 @@ def __init__(
                 init_args = kwargs.get("init_args", {})
                 if tx is None:
                     tx = torch._dynamo.symbolic_convert.InstructionTranslator.current_tx()
-                self._set_vt = SourcelessBuilder.create(tx, python_type).call_function(  # type: ignore[assignment]
+                self._base_vt = SourcelessBuilder.create(tx, python_type).call_function(  # type: ignore[assignment]
                     tx, init_args, {}
                 )
         else:
-            self._set_vt = set_vt
-
-    def call_method(
-        self,
-        tx: "InstructionTranslator",
-        name: str,
-        args: list[VariableTracker],
-        kwargs: dict[str, VariableTracker],
-    ) -> VariableTracker:
-        method = self._maybe_get_baseclass_method(name)
-        if method in self._set_methods:
-            return self._set_vt.call_method(tx, name, args, kwargs)
-        return super().call_method(tx, name, args, kwargs)
+            self._base_vt = set_vt
+        assert self._base_vt is not None
 
     def as_python_constant(self) -> object:
-        return self._set_vt.as_python_constant()
-
-    def unpack_var_sequence(self, tx: "InstructionTranslator") -> list[VariableTracker]:
-        if inspect.getattr_static(self.value, "__iter__") in (
-            set.__iter__,
-            frozenset.__iter__,
-        ):
-            return self._set_vt.unpack_var_sequence(tx)
-        raise NotImplementedError
+        assert self._base_vt is not None
+        return self._base_vt.as_python_constant()
 
     @property
     def set_items(self) -> set[Any]:
-        return self._set_vt.set_items
+        assert self._base_vt is not None
+        return self._base_vt.set_items  # pyrefly: ignore[missing-attribute]
 
     @property
-    def items(self) -> list[VariableTracker]:
-        return self._set_vt.items
-
-    def is_underlying_vt_modified(self, side_effects: "SideEffects") -> bool:
-        return side_effects.is_modified(self._set_vt)
-
-    def install_dict_keys_match_guard(self) -> None:
-        return self._set_vt.install_dict_keys_match_guard()
-
-    def install_dict_contains_guard(
-        self, tx: "InstructionTranslator", args: list[VariableTracker]
-    ) -> None:
-        return self._set_vt.install_dict_contains_guard(tx, args)
-
-    def is_python_hashable(self) -> bool:
-        raise_on_overridden_hash(self.value, self)
-        return self._set_vt.is_python_hashable()
-
-    def get_python_hash(self) -> int:
-        return self._set_vt.get_python_hash()
+    def items(self) -> dict[HashableTracker, VariableTracker]:
+        assert self._base_vt is not None
+        return self._base_vt.items  # pyrefly: ignore[missing-attribute]
 
     def is_python_equal(self, other: object) -> bool:
+        assert self._base_vt is not None
         return isinstance(
             other, UserDefinedSetVariable
-        ) and self._set_vt.is_python_equal(other._set_vt)
+        ) and self._base_vt.is_python_equal(other._base_vt)
 
 
 class UserDefinedListVariable(UserDefinedObjectVariable):
@@ -2634,35 +3179,11 @@ def __init__(
             assert self.source is None, (
                 "list_vt must be constructed by builder.py when source is present"
             )
-            self._list_vt = ListVariable([], mutation_type=ValueMutationNew())
+            self._base_vt = ListVariable([], mutation_type=ValueMutationNew())
         else:
-            self._list_vt = list_vt
-
-    def call_method(
-        self,
-        tx: "InstructionTranslator",
-        name: str,
-        args: list[VariableTracker],
-        kwargs: dict[str, VariableTracker],
-    ) -> VariableTracker:
-        assert self._list_vt is not None
-        method = self._maybe_get_baseclass_method(name)
-        if method in list_methods:
-            return self._list_vt.call_method(tx, name, args, kwargs)
-        return super().call_method(tx, name, args, kwargs)
-
-    def unpack_var_sequence(self, tx: "InstructionTranslator") -> list[VariableTracker]:
-        assert self._list_vt is not None
-        if type(self.value).__iter__ is list.__iter__:  # type: ignore[attr-defined]
-            return self._list_vt.unpack_var_sequence(tx)
-        raise NotImplementedError
-
-    def is_underlying_vt_modified(self, side_effects: "SideEffects") -> bool:
-        return side_effects.is_modified(self._list_vt)
-
-    def is_python_hashable(self) -> Literal[False]:
-        raise_on_overridden_hash(self.value, self)
-        return False
+            self._base_vt = list_vt
+        self._base_methods = list_methods
+        assert self._base_vt is not None
 
 
 class UserDefinedTupleVariable(UserDefinedObjectVariable):
@@ -2672,9 +3193,21 @@ class UserDefinedTupleVariable(UserDefinedObjectVariable):
     Internally, it uses a TupleVariable to represent the tuple part of the
     variable tracker. For everything else, it falls back to
     UserDefinedObjectVariable.
+
+    NamedTupleVariable and StructSequenceVariable are subclasses that handle
+    namedtuples and structseqs (torch.return_types.*) respectively.
     """
 
-    _nonvar_fields = UserDefinedObjectVariable._nonvar_fields
+    _nonvar_fields = {
+        "tuple_cls",
+        *UserDefinedObjectVariable._nonvar_fields,
+    }
+
+    @staticmethod
+    def get_vt_cls(cls: type) -> type["UserDefinedTupleVariable"]:
+        if is_structseq_class(cls):
+            return StructSequenceVariable
+        return NamedTupleVariable
 
     def __init__(self, value, tuple_vt=None, init_args=None, **kwargs):  # type: ignore[all]
         from .lists import TupleVariable
@@ -2695,9 +3228,17 @@ def __init__(self, value, tuple_vt=None, init_args=None, **kwargs):  # type: ign
 
                 tx = InstructionTranslator.current_tx()
             elems = init_args[0].force_unpack_var_sequence(tx)
-            self._tuple_vt = TupleVariable(elems, mutation_type=ValueMutationNew())
+            self._base_vt = TupleVariable(elems, mutation_type=ValueMutationNew())
         else:
-            self._tuple_vt = tuple_vt
+            self._base_vt = tuple_vt
+        self.tuple_cls = type(value)
+        self._base_methods = tuple_methods
+        assert self._base_vt is not None
+
+    @property
+    def items(self) -> list[VariableTracker]:
+        assert self._base_vt is not None
+        return self._base_vt.items  # type: ignore[return-value]
 
     def call_method(
         self,
@@ -2706,43 +3247,214 @@ def call_method(
         args: list[VariableTracker],
         kwargs: dict[str, VariableTracker],
     ) -> VariableTracker:
-        assert self._tuple_vt is not None
         if name == "__eq__":
             if len(args) != 1 or kwargs:
                 raise ValueError("Improper arguments for method.")
-            return variables.ConstantVariable(self.is_python_equal(args[0]))
+            return VariableTracker.build(tx, self.is_python_equal(args[0]))
         elif name == "__ne__":
             if len(args) != 1 or kwargs:
                 raise ValueError("Improper arguments for method.")
-            return variables.ConstantVariable(not self.is_python_equal(args[0]))
-        method = self._maybe_get_baseclass_method(name)
-        if method in tuple_methods:
-            return self._tuple_vt.call_method(tx, name, args, kwargs)
+            return VariableTracker.build(tx, not self.is_python_equal(args[0]))
         return super().call_method(tx, name, args, kwargs)
 
-    def unpack_var_sequence(self, tx: "InstructionTranslator") -> list[VariableTracker]:
-        assert self._tuple_vt is not None
-        if type(self.value).__iter__ is tuple.__iter__:  # type: ignore[attr-defined]
-            return self._tuple_vt.unpack_var_sequence(tx)
+    def reconstruct(self, codegen: "PyCodegen") -> None:
+        # Sourceless namedtuples/structseqs (e.g. tensor subclass metadata from
+        # SourcelessBuilder) aren't in id_to_variable so codegen_save_tempvars
+        # never processes them. When they appear in return values, codegen falls
+        # through to call_reconstruct. This is the same pattern as other
+        # sourceless containers (ConstDictVariable, TupleVariable, etc.).
+        # UserDefinedDictVariable doesn't need this because it's never created
+        # sourceless — it only comes from VariableBuilder which always has a
+        # source.
+        assert self.source is None
+        create_fn = self.get_construct_fn()
+        codegen.add_push_null(
+            lambda: codegen.append_output(
+                codegen.create_load_const_unchecked(create_fn)
+            )
+        )
+        codegen(self._base_vt)
+        codegen.extend_output(create_call_function(1, False))
+
+    def get_construct_fn(self) -> Callable[..., Any]:
         raise NotImplementedError
 
-    def is_python_hashable(self) -> bool:
-        raise_on_overridden_hash(self.value, self)
-        return self._tuple_vt.is_python_hashable()
+    def _validate_rest_for_tree_map(
+        self, rest: "collections.abc.Sequence[VariableTracker]"
+    ) -> list["UserDefinedTupleVariable"] | None:
+        """Validate that rest args are compatible for tree_map fast-path."""
+        others: list[UserDefinedTupleVariable] = []
+        n = len(self.items)
+        for candidate in rest:
+            if (
+                not isinstance(candidate, UserDefinedTupleVariable)
+                or len(candidate.items) != n
+                or candidate.tuple_cls is not self.tuple_cls
+            ):
+                return None
+            others.append(candidate)
+        return others
 
-    def get_python_hash(self) -> int:
-        return self._tuple_vt.get_python_hash()
+    def _make_tree_map_result(
+        self, new_items: list[VariableTracker]
+    ) -> "UserDefinedTupleVariable":
+        from .lists import TupleVariable
+
+        tuple_vt = TupleVariable(new_items, mutation_type=ValueMutationNew())
+        return type(self)(
+            self.value,
+            tuple_vt=tuple_vt,
+            mutation_type=ValueMutationNew(),
+        )
+
+    def _is_pytree_node(self) -> bool:
+        from torch.utils._pytree import is_namedtuple_class
+
+        return is_namedtuple_class(self.tuple_cls) or is_structseq_class(self.tuple_cls)
+
+    def call_tree_map_branch(
+        self,
+        tx: "InstructionTranslator",
+        tree_map_fn: "variables.functions.UserFunctionVariable",
+        map_fn: "VariableTracker",
+        rest: "collections.abc.Sequence[VariableTracker]",
+        tree_map_kwargs: "dict[str, VariableTracker]",
+    ) -> "VariableTracker":
+        if not self._is_pytree_node():
+            return super().call_tree_map_branch(
+                tx, tree_map_fn, map_fn, rest, tree_map_kwargs
+            )
+        others = self._validate_rest_for_tree_map(rest)
+        if others is None:
+            return self._tree_map_fallback(
+                tx, tree_map_fn, map_fn, rest, tree_map_kwargs
+            )
+
+        new_items: list[VariableTracker] = []
+        for idx, item in enumerate(self.items):
+            sibling_leaves = [o.items[idx] for o in others]
+            new_items.append(
+                item.call_tree_map(
+                    tx, tree_map_fn, map_fn, sibling_leaves, tree_map_kwargs
+                )
+            )
+
+        return self._make_tree_map_result(new_items)
+
+    def call_tree_map_with_path_branch(
+        self,
+        tx: "InstructionTranslator",
+        tree_map_fn: "variables.functions.UserFunctionVariable",
+        map_fn: "VariableTracker",
+        rest: "collections.abc.Sequence[VariableTracker]",
+        tree_map_kwargs: "dict[str, VariableTracker]",
+        keypath: "tuple[Any, ...]",
+    ) -> "VariableTracker":
+        if not self._is_pytree_node():
+            return super().call_tree_map_with_path_branch(
+                tx, tree_map_fn, map_fn, rest, tree_map_kwargs, keypath
+            )
+        others = self._validate_rest_for_tree_map(rest)
+        if others is None:
+            return self._tree_map_with_path_fallback(
+                tx, tree_map_fn, map_fn, rest, tree_map_kwargs, keypath
+            )
+
+        fields = namedtuple_fields(self.tuple_cls)
+        new_items: list[VariableTracker] = []
+        for idx, item in enumerate(self.items):
+            sibling_leaves = [o.items[idx] for o in others]
+            child_keypath = keypath + (GetAttrKey(fields[idx]),)
+            new_items.append(
+                item.call_tree_map_with_path(
+                    tx,
+                    tree_map_fn,
+                    map_fn,
+                    sibling_leaves,
+                    tree_map_kwargs,
+                    child_keypath,
+                )
+            )
+
+        return self._make_tree_map_result(new_items)
 
     def is_python_equal(self, other: object) -> bool:
-        return isinstance(
-            other, UserDefinedTupleVariable
-        ) and self._tuple_vt.is_python_equal(other._tuple_vt)
+        assert self._base_vt is not None
+        other = other._base_vt if isinstance(other, UserDefinedTupleVariable) else other
+        return self._base_vt.is_python_equal(other)
+
+
+class NamedTupleVariable(UserDefinedTupleVariable):
+    """Represents Python namedtuples (created via collections.namedtuple).
+
+    Namedtuples use _tuplegetter descriptors for field access and
+    Type(*args) / Type._make(iterable) for construction.
+    """
+
+    def resolve_data_descriptor(
+        self,
+        tx: "InstructionTranslator",
+        name: str,
+        type_attr: object,
+        source: Source | None,
+    ) -> VariableTracker:
+        if isinstance(type_attr, _collections._tuplegetter):
+            # namedtuple fields are _tuplegetter descriptors implemented in C.
+            # We emulate _tuplegetter.__get__ by indexing into the tracked
+            # tuple items, because self.value may not hold actual runtime values.
+            _, (idx, _) = type_attr.__reduce__()
+            return self.items[idx]  # type: ignore[union-attr]
+        return super().resolve_data_descriptor(tx, name, type_attr, source)
+
+    def get_construct_fn(self) -> Callable[..., Any]:
+        return self.tuple_cls._make  # type: ignore[attr-defined]
+
+    def as_python_constant(self) -> Any:
+        items = [x.as_python_constant() for x in self.items]
+        return self.tuple_cls(*items)  # type: ignore[arg-type]
+
+    def as_proxy(self) -> Any:
+        items = [x.as_proxy() for x in self.items]
+        return self.tuple_cls(*items)  # type: ignore[arg-type]
+
+
+class StructSequenceVariable(UserDefinedTupleVariable):
+    """Represents C-implemented PyStructSequence types (torch.return_types.*).
+
+    Structseqs use Type(iterable) calling convention and reject tuple.__new__.
+    """
+
+    def resolve_data_descriptor(
+        self,
+        tx: "InstructionTranslator",
+        name: str,
+        type_attr: object,
+        source: Source | None,
+    ) -> VariableTracker:
+        if isinstance(type_attr, types.MemberDescriptorType):
+            # Structseq fields are member_descriptor objects. We emulate
+            # field access by looking up the field name in _fields and
+            # indexing into the tracked tuple items.
+            fields = namedtuple_fields(self.tuple_cls)
+            if name in fields:
+                return self.items[fields.index(name)]
+        return super().resolve_data_descriptor(tx, name, type_attr, source)
+
+    def get_construct_fn(self) -> Callable[..., Any]:
+        return self.tuple_cls
+
+    def as_python_constant(self) -> Any:
+        items = [x.as_python_constant() for x in self.items]
+        return self.tuple_cls(items)
+
+    def as_proxy(self) -> Any:
+        items = [x.as_proxy() for x in self.items]
+        return self.tuple_cls(items)
 
 
 class MutableMappingVariable(UserDefinedObjectVariable):
     def __init__(self, value: object, **kwargs: Any) -> None:
         super().__init__(value, **kwargs)
-        self.generic_dict_vt = ConstDictVariable({})
 
     def method_setattr_standard(
         self,
@@ -2775,7 +3487,7 @@ def method_setattr_standard(
             if isinstance(descriptor, property) and descriptor.fset is not None:
                 fset_source = None
                 if self.cls_source:
-                    desc_source = self.get_source_by_walking_mro(name_str)
+                    desc_source = self.get_source_by_walking_mro(tx, name_str)
                     fset_source = AttrSource(desc_source, "fset")
                 fset_vt = VariableTracker.build(tx, descriptor.fset, fset_source)
                 return fset_vt.call_function(tx, [self, value], {})
@@ -2784,29 +3496,22 @@ def method_setattr_standard(
 
     def var_getattr(self, tx: "InstructionTranslator", name: str) -> VariableTracker:
         # A common pattern in the init code of MutableMapping objects is to
-        # update the __dict__ attribute. To prevent graph break, we directly
-        # return a ConstDictVariable for the __dict__attr.
-        #
-        # However, users can try to add a new attribute to the class using the
-        # __dict__ attribute. To catch this, we save the ConstDictVariable for
-        # the __dict__ and then lookup into this vt for each attr lookup.
+        # update the __dict__ attribute. The parent class
+        # (UserDefinedObjectVariable) implements __dict__ lookups using a VT
+        # (self.dict_vt) that uses the side effects table as source of truth.
         if name == "get" and type(self.value).get in (  # type: ignore[attr-defined]
             collections.abc.Mapping.get,
             dict.get,
         ):
             return variables.UserMethodVariable(polyfills.mapping_get, self)
-        elif name == "__dict__" and self.source:
-            self.generic_dict_vt = variables.LazyVariableTracker.create(  # type: ignore[assignment]
-                self.value.__dict__, AttrSource(self.source, "__dict__")
-            )
-            return self.generic_dict_vt
-        elif out := self.generic_dict_vt.maybe_getitem_const(
-            variables.ConstantVariable(name)
-        ):
-            return out
         else:
             return super().var_getattr(tx, name)
 
+    def mp_length(self, tx: "InstructionTranslator") -> VariableTracker:
+        if self._maybe_get_baseclass_method("__len__") in dict_methods:
+            return VariableTracker.build(tx, len(self.value))  # type: ignore[bad-argument-type]
+        return super().mp_length(tx)
+
 
 class RandomVariable(UserDefinedObjectVariable):
     pass
diff --git a/torch/_export/__init__.py b/torch/_export/__init__.py
index 3f4d286b94179..7a46a50c04b3f 100644
--- a/torch/_export/__init__.py
+++ b/torch/_export/__init__.py
@@ -80,14 +80,14 @@ def aot_compile_warning():
 def aot_compile(
     f: Callable,
     args: tuple[Any, ...],
-    kwargs: Optional[dict[str, Any]] = None,
+    kwargs: dict[str, Any] | None = None,
     *,
-    dynamic_shapes: Optional[dict[str, Any]] = None,
-    options: Optional[dict[str, Any]] = None,
+    dynamic_shapes: dict[str, Any] | None = None,
+    options: dict[str, Any] | None = None,
     remove_runtime_assertions: bool = False,
     disable_constraint_solver: bool = False,
     same_signature: bool = True,
-) -> Union[list[Any], str]:
+) -> list[Any] | str:
     """
     Note: this function is not stable yet
 
diff --git a/torch/_export/converter.py b/torch/_export/converter.py
index 54e46bd37f794..e69ca58d3ae7e 100644
--- a/torch/_export/converter.py
+++ b/torch/_export/converter.py
@@ -6,7 +6,7 @@
 import warnings
 from collections.abc import Callable, Sequence
 from contextlib import contextmanager
-from typing import Any, Optional, Union
+from typing import Any
 
 import torch
 import torch.export._trace
@@ -72,8 +72,8 @@ def _trace_and_get_graph_from_model(model, args):
 
 
 def _create_jit_graph(
-    model: Union[torch.nn.Module, torch.jit.ScriptFunction], args: Sequence[Any]
-) -> tuple[torch.Graph, list["_C.IValue"], Any, Optional[torch.ScriptModule]]:
+    model: torch.nn.Module | torch.jit.ScriptFunction, args: Sequence[Any]
+) -> tuple[torch.Graph, list["_C.IValue"], Any, torch.ScriptModule | None]:
     if isinstance(model, (torch.jit.ScriptFunction, torch.jit.ScriptModule)):
         flattened_args = tuple(torch.jit._flatten(tuple(args))[0])
         torch_out = None
@@ -393,7 +393,7 @@ def get_op_overload(node: torch._C.Node):
 class TS2FXGraphConverter:
     def __init__(
         self,
-        ts_graph: Union[torch._C.Graph, torch._C.Block],
+        ts_graph: torch._C.Graph | torch._C.Block,
         name_to_param: dict[str, torch.Tensor],
         name_to_buffer: dict[str, torch.Tensor],
         blocks_to_lifted_attrs: dict[torch._C.Block, set[str]],
@@ -413,7 +413,7 @@ def __init__(
 
         # Mapping of TS node name to converted FX node
         self.name_to_node: dict[
-            str, Union[torch.fx.Node, list[torch.fx.Node], dict[Any, torch.fx.Node]]
+            str, torch.fx.Node | list[torch.fx.Node] | dict[Any, torch.fx.Node]
         ] = {}
         # Mapping of TS node name to constant value (int, str, TorchBind obj,
         # tensor constants ...)
@@ -1338,7 +1338,7 @@ def __contains__(self, key):
 
     def __init__(
         self,
-        ts_graph: Union[torch._C.Graph, torch._C.Block],
+        ts_graph: torch._C.Graph | torch._C.Block,
         name_to_param: dict[str, torch.Tensor],
         name_to_buffer: dict[str, torch.Tensor],
         blocks_to_lifted_attrs: dict[torch._C.Block, set[str]],
@@ -1400,9 +1400,9 @@ class TS2EPConverter:
     # TorchScript model to ExportedProgram converter
     def __init__(
         self,
-        ts_model: Union[torch.jit.ScriptModule, torch.jit.ScriptFunction],
+        ts_model: torch.jit.ScriptModule | torch.jit.ScriptFunction,
         sample_args: tuple[Any, ...],
-        sample_kwargs: Optional[dict[str, Any]] = None,
+        sample_kwargs: dict[str, Any] | None = None,
     ):
         self.ts_model = ts_model
         self.ts_graph, self.params, _, _ = _create_jit_graph(ts_model, sample_args)
diff --git a/torch/_export/db/case.py b/torch/_export/db/case.py
index a4c22ad534f59..476fe9793a55c 100644
--- a/torch/_export/db/case.py
+++ b/torch/_export/db/case.py
@@ -4,7 +4,7 @@
 import string
 from dataclasses import dataclass, field
 from enum import Enum
-from typing import Any, Optional
+from typing import Any
 from types import ModuleType
 
 import torch
@@ -79,11 +79,11 @@ class ExportCase:
     model: torch.nn.Module
     name: str
     example_kwargs: dict[str, Any] = field(default_factory=dict)
-    extra_args: Optional[ArgsType] = None  # For testing graph generalization.
+    extra_args: ArgsType | None = None  # For testing graph generalization.
     # Tags associated with the use case. (e.g dynamic-shape, escape-hatch)
     tags: set[str] = field(default_factory=set)
     support_level: SupportLevel = SupportLevel.SUPPORTED
-    dynamic_shapes: Optional[dict[str, Any]] = None
+    dynamic_shapes: dict[str, Any] | None = None
 
     def __post_init__(self):
         check_inputs_type(self.example_args, self.example_kwargs)
diff --git a/torch/_export/db/examples/unsupported_operator.py b/torch/_export/db/examples/unsupported_operator.py
index f5a52d80b895b..ba5475c5f15fe 100644
--- a/torch/_export/db/examples/unsupported_operator.py
+++ b/torch/_export/db/examples/unsupported_operator.py
@@ -5,7 +5,7 @@
 
 class TorchSymMin(torch.nn.Module):
     """
-    torch.sym_min operator is not supported in export.
+    torch.sym_min operator is supported in export.
     """
 
     def forward(self, x):
@@ -14,5 +14,5 @@ def forward(self, x):
 
 example_args = (torch.randn(3, 2),)
 tags = {"torch.operator"}
-support_level = SupportLevel.NOT_SUPPORTED_YET
+support_level = SupportLevel.SUPPORTED
 model = TorchSymMin()
diff --git a/torch/_export/db/logging.py b/torch/_export/db/logging.py
index 652972b62e3f0..1416b1e99c2d1 100644
--- a/torch/_export/db/logging.py
+++ b/torch/_export/db/logging.py
@@ -1,4 +1,3 @@
-from typing import Optional
 
 def exportdb_error_message(case_name: str) -> str:
     from .examples import all_examples
@@ -18,7 +17,7 @@ def exportdb_error_message(case_name: str) -> str:
         return f"{case_name} is unsupported."
 
 
-def get_class_if_classified_error(e: Exception) -> Optional[str]:
+def get_class_if_classified_error(e: Exception) -> str | None:
     """
     Returns a string case name if the export error e is classified.
     Returns None otherwise.
diff --git a/torch/_export/non_strict_utils.py b/torch/_export/non_strict_utils.py
index 055b3b9a24184..2f1aae39d0a1a 100644
--- a/torch/_export/non_strict_utils.py
+++ b/torch/_export/non_strict_utils.py
@@ -9,7 +9,7 @@
 from collections import defaultdict
 from collections.abc import Callable, Sequence
 from contextlib import contextmanager
-from typing import Any, Optional, TYPE_CHECKING, Union
+from typing import Any, TYPE_CHECKING, TypeGuard
 
 import torch
 import torch.utils._pytree as pytree
@@ -25,7 +25,8 @@
 from torch._export.utils import _fakify_params_buffers
 from torch._guards import Source
 from torch._library.fake_class_registry import FakeScriptObject
-from torch._library.opaque_object import is_opaque_type
+from torch._library.opaque_object import is_opaque_value
+from torch._opaque_base import OpaqueBase
 from torch._subclasses.fake_tensor import FakeTensorMode
 from torch.export import Constraint
 from torch.export.dynamic_shapes import (
@@ -138,7 +139,7 @@ def make_sourced_prefixes(nn_module, args, kwargs) -> _KeyPathTrie:
 
 
 def key_path_to_source(
-    kp: KeyPath, sourced_prefixes: Optional[_KeyPathTrie] = None
+    kp: KeyPath, sourced_prefixes: _KeyPathTrie | None = None
 ) -> Source:
     """
     Given a key path, return the source for the key path.
@@ -171,13 +172,13 @@ def fakify(
     t: Any,
     t_constraints: dict[int, dict[int, Constraint]],
     sources: dict[tuple[int, int], list[Source]],
-    sourced_prefixes: Optional[_KeyPathTrie] = None,
+    sourced_prefixes: _KeyPathTrie | None = None,
 ):
     source = key_path_to_source(kp, sourced_prefixes=sourced_prefixes)
     if (
         _is_constant_argument(t)
         or isinstance(t, (torch.ScriptObject, torch.nn.Module))
-        or is_opaque_type(type(t))
+        or is_opaque_value(t)
     ):
         return t
 
@@ -257,11 +258,18 @@ def _create_symbolic_context_for_tensor(t, source, t_constraints, sources, mode)
 
         # Propagate outer tensor constraints to inner tensors if not already present
         for attr in attrs:
-            inner_tensor = getattr(t, attr)
-            inner_source = AttrSource(source, attr)
-            inner_contexts[attr] = _create_symbolic_context_for_tensor(
-                inner_tensor, inner_source, t_constraints, sources, mode
-            )
+            match getattr(t, attr):
+                case torch.Tensor() as inner_value:
+                    inner_source = AttrSource(source, attr)
+                    inner_contexts[attr] = _create_symbolic_context_for_tensor(
+                        inner_value, inner_source, t_constraints, sources, mode
+                    )
+                case OpaqueBase():
+                    pass
+                case unexpected:
+                    raise AssertionError(
+                        f"expected Tensor or OpaqueBase, got {type(unexpected)}"
+                    )
 
         symbolic_context = SubclassSymbolicContext(
             dynamic_sizes=dynamic_sizes,
@@ -485,7 +493,7 @@ def make_fake_inputs(
 
         names: dict[str, tuple[int, int]] = {}
         source_pairs: list[tuple[Source, Source]] = []
-        derived_equalities: list[tuple[Source, Union[Source, Symbol], Callable]] = []
+        derived_equalities: list[tuple[Source, Source | Symbol, Callable]] = []
         phantom_symbols: dict[str, Symbol] = {}
         relaxed_sources: set[Source] = set()
         for constraint in constraints:
@@ -519,7 +527,7 @@ def make_fake_inputs(
 
 def _flatten_dynamic_shapes(
     combined_args: dict[str, Any],
-    dynamic_shapes: Union[dict[str, Any], tuple[Any], list[Any]],
+    dynamic_shapes: dict[str, Any] | tuple[Any] | list[Any],
 ) -> list[Any]:
     flat_shapes = []
 
@@ -546,7 +554,7 @@ def _clean_dynamic_markers(tensor: torch.Tensor) -> None:
 def produce_guards_and_solve_constraints(
     fake_mode: FakeTensorMode,
     gm: torch.fx.GraphModule,
-    dynamic_shapes: Union[dict[str, Any], tuple[Any], list[Any], None],
+    dynamic_shapes: dict[str, Any] | tuple[Any] | list[Any] | None,
     equalities_inputs: EqualityConstraint,
     original_signature: inspect.Signature,
 ):
@@ -615,7 +623,7 @@ def produce_guards_and_solve_constraints(
         raise constraint_violation_error
 
 
-def is_int(x: object) -> bool:
+def is_int(x: object) -> TypeGuard[int | torch.SymInt]:
     return isinstance(x, int) or (isinstance(x, torch.SymInt) and x.node.expr.is_number)
 
 
@@ -626,8 +634,8 @@ def _constrain_user_specified_dimhint_range(
     range_constraints,
     shape_env,
     keypath: KeyPath,
-    i: Optional[int] = None,
-) -> Optional[str]:
+    i: int | None = None,
+) -> str | None:
     trace_vr = (
         range_constraints[symint.node.expr]
         if not is_int(symint)
@@ -696,7 +704,7 @@ def make_constraints(
     fake_mode: FakeTensorMode,
     gm: torch.fx.GraphModule,
     combined_args: dict[str, Any],
-    dynamic_shapes: Union[dict[str, Any], tuple[Any], list[Any], None],
+    dynamic_shapes: dict[str, Any] | tuple[Any] | list[Any] | None,
     num_lifted_inputs: int,
 ):
     """
@@ -865,7 +873,7 @@ def inner(m: torch.nn.Module, prefix_atoms: list[str], constants):
 
 
 def _get_graph_inputs_of_type_nn_module(
-    args: Optional[tuple[tuple[Any], dict[Any, Any]]],
+    args: tuple[tuple[Any], dict[Any, Any]] | None,
 ) -> set[type[torch.nn.Module]]:
     if args is None:
         return set()
@@ -892,7 +900,7 @@ def _exit_enable_graph_inputs_of_type_nn_module(
 
 @contextlib.contextmanager
 def _enable_graph_inputs_of_type_nn_module(
-    args: Optional[tuple[tuple[Any], dict[Any, Any]]],
+    args: tuple[tuple[Any], dict[Any, Any]] | None,
 ):
     if args is None:
         yield
@@ -941,7 +949,7 @@ def _fakify_script_objects(
     mod: torch.nn.Module,
     args: Sequence[Any],
     kwargs: dict[Any, Any],
-    fake_mode: Optional[torch._subclasses.fake_tensor.FakeTensorMode],
+    fake_mode: torch._subclasses.fake_tensor.FakeTensorMode | None,
 ):
     # This context manager is used to fakify script objects into FakeScriptObject.
     # Inputs:
@@ -984,7 +992,7 @@ def _leaf_mod_and_attr(
         for obj, fqns in constant_attrs.items():
             if torch._library.fake_class_registry._is_script_object(
                 obj
-            ) or is_opaque_type(obj):
+            ) or is_opaque_value(obj):
                 fake_script_obj = _maybe_fakify_obj(obj)
                 for fqn in fqns:
                     cur_mod, attr = _leaf_mod_and_attr(mod, fqn)
diff --git a/torch/_export/pass_base.py b/torch/_export/pass_base.py
index 27a3f01aac47a..6200264118fc7 100644
--- a/torch/_export/pass_base.py
+++ b/torch/_export/pass_base.py
@@ -1,10 +1,9 @@
 # mypy: allow-untyped-defs
 import operator
 import traceback
-import typing
 from collections.abc import Callable
 from contextlib import nullcontext
-from typing import Any, Optional, Union
+from typing import Any
 
 import torch
 from torch import fx
@@ -32,7 +31,7 @@
 Argument = Any
 Value = Any
 Fn = Callable[..., Any]
-PassType = Callable[[torch.fx.GraphModule], Optional[PassResult]]
+PassType = Callable[[torch.fx.GraphModule], PassResult | None]
 
 
 _TORCH_SYM_OPS: set[Callable] = {
@@ -70,7 +69,7 @@ def __init__(
             self.graph = torch.fx.Graph()
             self.graph.set_codegen(codegen)
             self.tensor_attrs: dict[str, torch.Tensor] = {}  # type: ignore[assignment]
-            self.fake_tensor_mode: Optional[FakeTensorMode] = None
+            self.fake_tensor_mode: FakeTensorMode | None = None
             self.submodules: dict[torch.nn.Module, str] = {}
 
         def trace(self) -> None:  # type: ignore[override]
@@ -104,17 +103,17 @@ def set_metadata(
             # propagate the fake tensor or sym nodes
             def make_val(
                 x: Argument,
-            ) -> Union[
-                FakeTensor,
-                torch.SymInt,
-                torch.SymFloat,
-                torch.SymBool,
-                int,
-                float,
-                bool,
-                str,
-                None,
-            ]:
+            ) -> (
+                FakeTensor
+                | torch.SymInt
+                | torch.SymFloat
+                | torch.SymBool
+                | int
+                | float
+                | bool
+                | str
+                | None
+            ):
                 if isinstance(x, FakeTensor):
                     return x
                 elif isinstance(x, torch.Tensor):
@@ -161,7 +160,7 @@ def make_val(
             node.meta["val"] = pytree.tree_map(make_val, value)
 
             # Set the tensor_metadata for values that do not have a corresponding FakeTensor
-            def make_tensor_meta(x: Argument) -> Optional[TensorMetadata]:
+            def make_tensor_meta(x: Argument) -> TensorMetadata | None:
                 if not isinstance(x, FakeTensor) and isinstance(x, torch.Tensor):
                     if x.is_quantized:
                         # TODO (tmanlaibaatar) properly support Quantized FakeTensor
@@ -292,9 +291,9 @@ def __init__(self) -> None:
             torch.fx.GraphModule(torch.nn.Module(), torch.fx.Graph())
         )
         self.tracer = self.ExportTracer(self, CodeGen())
-        self.fake_tensor_mode: Optional[FakeTensorMode] = None
+        self.fake_tensor_mode: FakeTensorMode | None = None
         self._initialized = True
-        self.node_debug_str: typing.Optional[str] = None
+        self.node_debug_str: str | None = None
 
     def _fx(
         self,
@@ -332,7 +331,7 @@ def inputs(self, graph_module: torch.fx.GraphModule) -> list[Argument]:
         if (args := graph_module.meta.get("args", None)) is not None:
             return list(args)
 
-        def extract_input(node: torch.fx.Node) -> Optional[FakeTensor]:
+        def extract_input(node: torch.fx.Node) -> FakeTensor | None:
             if "val" in node.meta:
                 fake = node.meta["val"]
                 if hasattr(fake, "constant") and fake.constant is not None:
diff --git a/torch/_export/pass_infra/proxy_value.py b/torch/_export/pass_infra/proxy_value.py
index 3f31a1cf46463..d4a3e1cc4a63c 100644
--- a/torch/_export/pass_infra/proxy_value.py
+++ b/torch/_export/pass_infra/proxy_value.py
@@ -1,6 +1,6 @@
 # pyre-strict
 from collections.abc import Iterable, Iterator
-from typing import Generic, TypeVar, Union
+from typing import Generic, TypeVar
 
 import torch
 
@@ -10,7 +10,7 @@
 
 class ProxyValue(Generic[_T]):
     # pyre-ignore
-    def __init__(self, data: Iterable[_T], proxy: Union[torch.fx.Proxy, torch.fx.Node]):
+    def __init__(self, data: Iterable[_T], proxy: torch.fx.Proxy | torch.fx.Node):
         # pyre-ignore
         self.data = data
         self.proxy_or_node = proxy
diff --git a/torch/_export/passes/_node_metadata_hook.py b/torch/_export/passes/_node_metadata_hook.py
index 9fea0b5d7d9e1..414bc04d1bad9 100644
--- a/torch/_export/passes/_node_metadata_hook.py
+++ b/torch/_export/passes/_node_metadata_hook.py
@@ -1,6 +1,6 @@
 # mypy: allow-untyped-defs
 import contextlib
-from typing import Any, Optional
+from typing import Any
 
 import torch
 import torch.utils._pytree as pytree
@@ -14,8 +14,8 @@
 
 def _node_metadata_hook(
     node: torch.fx.Node,
-    metadata: Optional[dict[str, Any]] = None,
-    fake_mode: Optional[FakeTensorMode] = None,
+    metadata: dict[str, Any] | None = None,
+    fake_mode: FakeTensorMode | None = None,
 ) -> None:
     """
     Hook for adding the appropriate metadata to nodes that are created during a
@@ -87,6 +87,8 @@ def _node_metadata_hook(
         ),
     )
 
+    node.meta["custom"] = node.meta.get("custom", arg_meta.get("custom", {}))
+
 
 @contextlib.contextmanager
 def _set_node_metadata_hook(gm: torch.fx.GraphModule, f):
diff --git a/torch/_export/passes/constant_folding.py b/torch/_export/passes/constant_folding.py
index 58534856422c7..ab0f8c33f6f26 100644
--- a/torch/_export/passes/constant_folding.py
+++ b/torch/_export/passes/constant_folding.py
@@ -2,7 +2,7 @@
 import collections
 from collections import defaultdict
 from collections.abc import Callable
-from typing import Any, Optional
+from typing import Any
 
 import torch
 import torch.utils._pytree as pytree
@@ -206,7 +206,7 @@ def run(self):  # type: ignore[override]
 
 def constant_fold(
     gm: torch.fx.GraphModule,
-    constraint_fn: Optional[Callable[[torch.fx.Node], bool]] = None,
+    constraint_fn: Callable[[torch.fx.Node], bool] | None = None,
 ):
     with torch.utils._python_dispatch._disable_current_modes():
         cf = ConstantFolder(gm, skip_constructors=True)
diff --git a/torch/_export/passes/functionalize_side_effectful_ops_pass.py b/torch/_export/passes/functionalize_side_effectful_ops_pass.py
index afe5298dfb0c0..c612aaf348887 100644
--- a/torch/_export/passes/functionalize_side_effectful_ops_pass.py
+++ b/torch/_export/passes/functionalize_side_effectful_ops_pass.py
@@ -1,5 +1,4 @@
 import copy
-from typing import Optional
 
 import torch
 from torch._export.pass_base import (
@@ -45,8 +44,8 @@ def f(x):
 
     def __init__(self) -> None:
         super().__init__()
-        self._dep_token: Optional[ProxyValue] = None
-        self._next_dep_token_index: Optional[int] = None
+        self._dep_token: ProxyValue | None = None
+        self._next_dep_token_index: int | None = None
 
     def call(self, graph_module: torch.fx.GraphModule) -> PassResult:
         # Early return if no non-functional assertions.
diff --git a/torch/_export/passes/lift_constants_pass.py b/torch/_export/passes/lift_constants_pass.py
index 32b8de4a5d916..7c5690d7c6b41 100644
--- a/torch/_export/passes/lift_constants_pass.py
+++ b/torch/_export/passes/lift_constants_pass.py
@@ -1,13 +1,17 @@
 # mypy: allow-untyped-defs
 import collections
 import logging
-from typing import Any, Optional, Union
+from typing import Any
 
 import torch
 from torch._export.verifier import SpecViolationError
 from torch._guards import detect_fake_mode
-from torch._library.fake_class_registry import FakeScriptObject
-from torch._library.opaque_object import is_opaque_reference_type
+from torch._library.fake_class_registry import FakeScriptObject, maybe_to_fake_obj
+from torch._library.opaque_object import (
+    get_opaque_type_name,
+    is_opaque_reference_type,
+    is_opaque_type,
+)
 from torch._subclasses.fake_tensor import unset_fake_temporarily
 from torch.export.exported_program import (
     ArgumentSpec,
@@ -26,16 +30,17 @@
 
 class ConstantAttrMap(collections.abc.MutableMapping):
     """A mapping class that understands how to use module constants (tensors,
-    ScriptObjects, FakeScriptObjects) as keys. We store tensors and FakeScriptObjects normally,
-    but ScriptObjects are stored by hash, because different torch.ScriptObjects can point to
-    the same underlying value (but we guarantee that they will `hash()` to the same value
+    ScriptObjects, FakeScriptObjects, opaque objects) as keys. We store tensors,
+    FakeScriptObjects, and opaque objects normally, but ScriptObjects are stored
+    by hash, because different torch.ScriptObjects can point to the same
+    underlying value (but we guarantee that they will `hash()` to the same value
     if that's the case).
     """
 
     def __init__(self) -> None:
         # Underlying dict that we use to implement this mapping.
         self._constant_attrs: dict[
-            Union[int, torch.Tensor, FakeScriptObject, torch.utils._pytree.TreeSpec],
+            int | torch.Tensor | FakeScriptObject | torch.utils._pytree.TreeSpec,
             list[Any],
         ] = {}
         # Map from the hash(ScriptObject) to the ScriptObject itself. Used for
@@ -45,9 +50,11 @@ def __init__(self) -> None:
 
     def __getitem__(self, key: _ConstantAttributeType) -> Any:
         real_key = hash(key) if isinstance(key, torch.ScriptObject) else key
-        if not isinstance(real_key, (int, torch.Tensor, FakeScriptObject)):
+        if not isinstance(
+            real_key, (int, torch.Tensor, FakeScriptObject)
+        ) and not is_opaque_type(type(real_key)):
             raise AssertionError(
-                f"expected int, Tensor, or FakeScriptObject key, got {type(real_key)}"
+                f"expected int, Tensor, FakeScriptObject, or opaque type key, got {type(real_key)}"
             )
         return self._constant_attrs[real_key]
 
@@ -64,7 +71,9 @@ def add(self, key: _ConstantAttributeType, value: Any) -> None:
                 self._constant_attrs[hash(key)] = []
             self._constant_attrs[hash(key)].append(value)
             self._script_object_map[hash(key)] = key
-        elif isinstance(key, (torch.Tensor, FakeScriptObject)):
+        elif isinstance(key, (torch.Tensor, FakeScriptObject)) or is_opaque_type(
+            type(key)
+        ):
             if key not in self._constant_attrs:
                 self._constant_attrs[key] = []
             self._constant_attrs[key].append(value)
@@ -114,7 +123,7 @@ def _get_first_fqn(
     return fqns[0] if fqns else None
 
 
-def _unused_constant(node: torch.fx.Node) -> Optional[list[torch.fx.Node]]:
+def _unused_constant(node: torch.fx.Node) -> list[torch.fx.Node] | None:
     """
     If there is a tensor constant created while tracing, here is how the graph
     looks like:
@@ -346,6 +355,23 @@ def lift_constants_pass(
                         class_fqn=class_fqn,
                         fake_val=constant_val,
                     )
+                elif is_opaque_type(type(constant_val)):
+                    class_fqn = get_opaque_type_name(type(constant_val))
+                    fake_val = (
+                        maybe_to_fake_obj(fake_mode, constant_val)
+                        if fake_mode
+                        else None
+                    )
+                    const_placeholder_node.meta["val"] = CustomObjArgument(
+                        constant_fqn,
+                        class_fqn,
+                        fake_val,  # pyrefly: ignore[bad-argument-type]
+                    )
+                    input_spec_arg = CustomObjArgument(
+                        name=const_placeholder_node.name,
+                        class_fqn=class_fqn,
+                        fake_val=fake_val,  # pyrefly: ignore[bad-argument-type]
+                    )
                 else:
                     raise SpecViolationError(
                         f"tried to lift unsupported type {type(constant_val)} from node {node.format_node()}"
diff --git a/torch/_export/passes/replace_quantized_ops_with_standard_ops_pass.py b/torch/_export/passes/replace_quantized_ops_with_standard_ops_pass.py
index 2ab441e8c2e0c..9be342f0d188d 100644
--- a/torch/_export/passes/replace_quantized_ops_with_standard_ops_pass.py
+++ b/torch/_export/passes/replace_quantized_ops_with_standard_ops_pass.py
@@ -1,7 +1,6 @@
 # mypy: allow-untyped-defs
 import logging
 import operator
-from typing import Optional, Union
 
 import torch
 import torch.export._trace
@@ -18,9 +17,9 @@
 log = logging.getLogger(__name__)
 
 # Those values will need to be carried over multiple operators.
-_INPUT_Q_DTYPE: Optional[Union[torch.dtype, torch.fx.Node]] = None
-_SCALE: Optional[Union[float, torch.fx.Node]] = None
-_ZERO_POINT: Optional[Union[float, torch.fx.Node]] = None
+_INPUT_Q_DTYPE: torch.dtype | torch.fx.Node | None = None
+_SCALE: float | torch.fx.Node | None = None
+_ZERO_POINT: float | torch.fx.Node | None = None
 
 
 def int_to_valid_dtype(val: int) -> torch.dtype:
@@ -43,12 +42,12 @@ def fx_enum_to_dtype(gm: torch.fx.GraphModule, val: int) -> torch.fx.Node:
 def insert_quantized_node(
     gm: torch.fx.GraphModule,
     val_node: torch.fx.Node,
-    scale_node: Union[float, torch.fx.Node],
-    zero_point_node: Union[float, torch.fx.Node],
-    qmin_node: Union[float, int, torch.fx.Node],
-    qmax_node: Union[float, int, torch.fx.Node],
-    dtype_node: Union[torch.dtype, torch.fx.Node],
-    qscheme: Optional[torch.qscheme],
+    scale_node: float | torch.fx.Node,
+    zero_point_node: float | torch.fx.Node,
+    qmin_node: float | int | torch.fx.Node,
+    qmax_node: float | int | torch.fx.Node,
+    dtype_node: torch.dtype | torch.fx.Node,
+    qscheme: torch.qscheme | None,
 ) -> torch.fx.Node:
     return gm.graph.call_function(
         quantize_per_tensor,
@@ -65,13 +64,13 @@ def insert_quantized_node(
 
 def get_dequantized(
     val: torch.Tensor,
-    scale: Union[float, torch.Tensor],
-    zero_point: Union[float, torch.Tensor],
-    qmin: Union[float, int],
-    qmax: Union[float, int],
+    scale: float | torch.Tensor,
+    zero_point: float | torch.Tensor,
+    qmin: float | int,
+    qmax: float | int,
     dtype: torch.dtype,
-    axis: Optional[int],
-    qscheme: Optional[torch.qscheme],
+    axis: int | None,
+    qscheme: torch.qscheme | None,
 ) -> torch.Tensor:
     if qscheme is torch.per_tensor_affine:
         return dequantize_per_tensor(
@@ -99,13 +98,13 @@ def get_dequantized(
 def insert_dequantized_node(
     gm: torch.fx.GraphModule,
     val_node: torch.fx.Node,
-    scale_node: Union[float, torch.fx.Node],
-    zero_point_node: Union[float, torch.fx.Node],
-    qmin_node: Union[float, int, torch.fx.Node],
-    qmax_node: Union[float, int, torch.fx.Node],
-    dtype_node: Union[torch.dtype, torch.fx.Node],
-    axis_node: Optional[Union[int, torch.fx.Node]],
-    qscheme: Optional[torch.qscheme],
+    scale_node: float | torch.fx.Node,
+    zero_point_node: float | torch.fx.Node,
+    qmin_node: float | int | torch.fx.Node,
+    qmax_node: float | int | torch.fx.Node,
+    dtype_node: torch.dtype | torch.fx.Node,
+    axis_node: int | torch.fx.Node | None,
+    qscheme: torch.qscheme | None,
 ) -> torch.fx.Node:
     if qscheme is torch.per_tensor_affine:
         return gm.graph.call_function(
@@ -136,12 +135,12 @@ def insert_dequantized_node(
         raise RuntimeError(f"Unsupported dequantization scheme: {qscheme}")
 
 
-def get_qmin_qmax(dtype: torch.dtype) -> tuple[Union[int, float], Union[int, float]]:
+def get_qmin_qmax(dtype: torch.dtype) -> tuple[int | float, int | float]:
     return calculate_qmin_qmax(None, None, False, dtype, False)  # type: ignore[arg-type]
 
 
 def insert_qmin_qmax_node(
-    gm: torch.fx.GraphModule, dtype_node: Union[torch.dtype, torch.fx.Node]
+    gm: torch.fx.GraphModule, dtype_node: torch.dtype | torch.fx.Node
 ) -> tuple[torch.fx.Node, torch.fx.Node]:
     q_min_max_node = gm.graph.call_function(
         calculate_qmin_qmax, (None, None, False, dtype_node, False)
@@ -173,7 +172,7 @@ def get_script_object(
 def insert_weight_and_bias_get_attr_node_from_get_attr_to_scriptobject(
     gm: torch.fx.GraphModule,
     param_node: torch.fx.Node,
-) -> tuple[torch.fx.Node, Optional[torch.fx.Node]]:
+) -> tuple[torch.fx.Node, torch.fx.Node | None]:
     """Directly inline tensor from a get_attr fx node."""
     mod = get_script_object(gm, param_node)
     w_qtensor, b_qtensor = mod.unpack()  # type: ignore[attr-defined]
@@ -189,8 +188,8 @@ def insert_weight_and_bias_get_attr_node_from_get_attr_to_scriptobject(
 def insert_weight_and_bias_get_attr_node_from_get_attr_to_qtensor(
     gm: torch.fx.GraphModule,
     get_attr_to_weight_node: torch.fx.Node,
-    get_attr_to_bias_node: Optional[torch.fx.Node],
-) -> tuple[torch.fx.Node, Optional[torch.fx.Node]]:
+    get_attr_to_bias_node: torch.fx.Node | None,
+) -> tuple[torch.fx.Node, torch.fx.Node | None]:
     if not isinstance(get_attr_to_weight_node.target, str):
         raise AssertionError(
             f"expected str target, got {type(get_attr_to_weight_node.target).__name__}"
@@ -216,10 +215,10 @@ def insert_weight_and_bias_get_attr_node_from_get_attr_to_qtensor(
 def insert_weight_and_bias_get_attr_node(
     gm: torch.fx.GraphModule,
     w_qtensor: torch.Tensor,
-    b_qtensor: Optional[torch.Tensor],
+    b_qtensor: torch.Tensor | None,
     w_attr_name: str,
     b_attr_name: str,
-) -> tuple[torch.fx.Node, Optional[torch.fx.Node]]:
+) -> tuple[torch.fx.Node, torch.fx.Node | None]:
     w_tensor = get_tensor_from_qtensor(w_qtensor)
     _assign_attr(w_tensor, gm, w_attr_name)
     w_tensor_attr = gm.graph.get_attr(w_attr_name)
@@ -278,7 +277,7 @@ def insert_fused_activation_node(
 def _conv1d_op_with_squeeze(
     inp: torch.Tensor,
     weight: torch.Tensor,
-    bias: Optional[torch.Tensor],
+    bias: torch.Tensor | None,
     stride: list[int],
     padding: list[int],
     dilation: list[int],
diff --git a/torch/_export/passes/replace_view_ops_with_view_copy_ops_pass.py b/torch/_export/passes/replace_view_ops_with_view_copy_ops_pass.py
index 489bc19ed1d50..98d93e16559b7 100644
--- a/torch/_export/passes/replace_view_ops_with_view_copy_ops_pass.py
+++ b/torch/_export/passes/replace_view_ops_with_view_copy_ops_pass.py
@@ -1,5 +1,4 @@
 # mypy: allow-untyped-defs
-from typing import Optional
 
 import torch
 from torch._export.error import InternalError
@@ -22,7 +21,7 @@ def is_view_op(schema: torch._C.FunctionSchema) -> bool:
     return (alias_info is not None) and (not alias_info.is_write)
 
 
-def get_view_copy_of_view_op(schema: torch._C.FunctionSchema) -> Optional[OpOverload]:
+def get_view_copy_of_view_op(schema: torch._C.FunctionSchema) -> OpOverload | None:
     if is_view_op(schema) and schema.name.startswith("aten::"):
         view_op_name = schema.name.split("::")[1]
         view_op_overload = (
diff --git a/torch/_export/serde/dynamic_shapes.py b/torch/_export/serde/dynamic_shapes.py
index 36e726b94ffb0..32fe6014b0444 100644
--- a/torch/_export/serde/dynamic_shapes.py
+++ b/torch/_export/serde/dynamic_shapes.py
@@ -1,5 +1,5 @@
 import dataclasses
-from typing import Any, Optional, Union
+from typing import Any
 
 import torch
 from torch._dynamo.exc import UserError, UserErrorType
@@ -22,7 +22,7 @@ class RootDim:
     """
 
     min: int
-    max: Union[int, None]
+    max: int | None
     derived: list[str]
 
 
@@ -32,15 +32,15 @@ class DynamicShapesSpec:
     This stores a dynamic_shapes spec for de/serialization.
     """
 
-    dynamic_shapes: Union[dict[str, Any], tuple[Any], list[Any], None]
+    dynamic_shapes: dict[str, Any] | tuple[Any] | list[Any] | None
     dims: dict[str, RootDim]
 
 
 def _postprocess_serialized_shapes(
-    dynamic_shapes: Union[dict[str, Any], tuple[Any], list[Any], None],
-    dims: dict[str, dict[str, Union[int, list[str], None]]],
-    to_dict: Optional[bool] = False,
-) -> Union[DynamicShapesSpec, dict[str, Any]]:
+    dynamic_shapes: dict[str, Any] | tuple[Any] | list[Any] | None,
+    dims: dict[str, dict[str, int | list[str] | None]],
+    to_dict: bool | None = False,
+) -> DynamicShapesSpec | dict[str, Any]:
     """
     Sorts dims and dumps to dictionary format.
     """
@@ -64,11 +64,11 @@ def _postprocess_serialized_shapes(
 
 
 def _dump_dynamic_shapes(
-    dynamic_shapes: Union[dict[str, Any], tuple[Any], list[Any], None],
+    dynamic_shapes: dict[str, Any] | tuple[Any] | list[Any] | None,
     args: tuple[Any],
-    kwargs: Optional[dict[str, Any]] = None,
-    to_dict: Optional[bool] = False,
-) -> Union[DynamicShapesSpec, dict[str, Any]]:
+    kwargs: dict[str, Any] | None = None,
+    to_dict: bool | None = False,
+) -> DynamicShapesSpec | dict[str, Any]:
     """
     Utility function for dynamic shapes serialization, serializing a dynamic_shapes spec.
     Returns a DynamicShapesSpec dataclass containing 2 fields, "dynamic_shapes" and "dims".
@@ -152,8 +152,8 @@ def _standardize_shapes(path, tensor, shape):  # type: ignore[no-untyped-def]
         return out
 
     def _track_dim_from_dims(
-        val: Union[None, int, _DimHint, Dim],
-    ) -> Union[None, int, str]:
+        val: None | int | _DimHint | Dim,
+    ) -> None | int | str:
         """
         Tracks dims, ranges, derived dims from the standardized dynamic_shapes spec.
         """
@@ -202,9 +202,9 @@ def _track_dim_from_dims(
 
 
 def _load_dynamic_shapes(
-    spec: Union[DynamicShapesSpec, dict[str, Any]],
-    from_dict: Optional[bool] = False,
-) -> Union[dict[str, Any], tuple[Any], list[Any], None]:
+    spec: DynamicShapesSpec | dict[str, Any],
+    from_dict: bool | None = False,
+) -> dict[str, Any] | tuple[Any] | list[Any] | None:
     """
     Utility function for dynamic shapes serialization.
     Deserializes a DynamicShapesSpec or corresponding dictionary into a dynamic_shapes input to export().
@@ -300,8 +300,8 @@ def _load_dynamic_shapes(
             dim_cache[_expr] = ddim  # cache derived dims
 
     def deserialize_shape(
-        val: Union[None, int, str],
-    ) -> Union[None, int, Dim, _DimHint]:
+        val: None | int | str,
+    ) -> None | int | Dim | _DimHint:
         if val is None or isinstance(val, int):
             return val
         elif val == "_DimHint.AUTO":
diff --git a/torch/_export/serde/export_schema.thrift b/torch/_export/serde/export_schema.thrift
index b8fb0c507e20c..5268efc883809 100644
--- a/torch/_export/serde/export_schema.thrift
+++ b/torch/_export/serde/export_schema.thrift
@@ -1,5 +1,5 @@
 // @generated by update_schema.py
-// checksum<<971a23161cd2f3ca95f10592dd3e8eabe8fc2074acfead937ed0b765aae6611a>>
+// checksum<<ffbfd7b406d10b13faddd5faafc7296f90cf65af9409b166c135f6169946a08c>>
 
 namespace py3 torch._export
 namespace cpp2 torch._export.schema
@@ -52,6 +52,9 @@ enum ScalarType {
   FLOAT8E5M2 = 30,
   FLOAT8E4M3FNUZ = 31,
   FLOAT8E5M2FNUZ = 32,
+  FLOAT8E8M0FNU = 33,
+  UINT32 = 34,
+  UINT64 = 35,
 }
 
 
@@ -170,6 +173,7 @@ union Argument {
   270: list<list<TensorArgument>> as_nested_tensors;
   280: list<list<i64>> as_int_lists;
   290: map<string, Argument> as_string_to_argument;
+  300: list<list<double>> as_float_lists;
 }
 
 struct NamedArgument {
diff --git a/torch/_export/serde/schema.py b/torch/_export/serde/schema.py
index 0e468cd593415..c7ee7b43495bd 100644
--- a/torch/_export/serde/schema.py
+++ b/torch/_export/serde/schema.py
@@ -3,18 +3,23 @@
 
 from dataclasses import dataclass, field
 from enum import IntEnum
-from typing import Annotated, Optional
+from typing import Annotated
 
 from torch._export.serde.union import _Union, _union_dataclass
 
 
 # NOTE: Please update this value if any modifications are made to the schema
-SCHEMA_VERSION = (8, 17)
+SCHEMA_VERSION = (8, 20)
 TREESPEC_VERSION = 1
 
 
 # NOTE: If you updated the schema, please run `scripts/export/update_schema.py`
 # to update the auto generated files.
+#
+# There are also mappings from serialized enum values to c10 enum member names.
+# These are used by scripts/export/update_schema.py to generate C++ conversion
+# functions for oss_proxy_executor.cpp.
+# When adding new enum values, update both the enum AND the mapping.
 class ScalarType(IntEnum):
     UNKNOWN = 0
     BYTE = 1
@@ -35,6 +40,9 @@ class ScalarType(IntEnum):
     FLOAT8E5M2 = 30
     FLOAT8E4M3FNUZ = 31
     FLOAT8E5M2FNUZ = 32
+    FLOAT8E8M0FNU = 33
+    UINT32 = 34
+    UINT64 = 35
 
 
 class Layout(IntEnum):
@@ -56,10 +64,52 @@ class MemoryFormat(IntEnum):
     PreserveFormat = 4
 
 
+SCALAR_TYPE_TO_C10: dict[int, str] = {
+    ScalarType.BYTE: "Byte",
+    ScalarType.CHAR: "Char",
+    ScalarType.SHORT: "Short",
+    ScalarType.INT: "Int",
+    ScalarType.LONG: "Long",
+    ScalarType.HALF: "Half",
+    ScalarType.FLOAT: "Float",
+    ScalarType.DOUBLE: "Double",
+    ScalarType.COMPLEXHALF: "ComplexHalf",
+    ScalarType.COMPLEXFLOAT: "ComplexFloat",
+    ScalarType.COMPLEXDOUBLE: "ComplexDouble",
+    ScalarType.BOOL: "Bool",
+    ScalarType.BFLOAT16: "BFloat16",
+    ScalarType.UINT16: "UInt16",
+    ScalarType.FLOAT8E4M3FN: "Float8_e4m3fn",
+    ScalarType.FLOAT8E5M2: "Float8_e5m2",
+    ScalarType.FLOAT8E4M3FNUZ: "Float8_e4m3fnuz",
+    ScalarType.FLOAT8E5M2FNUZ: "Float8_e5m2fnuz",
+    ScalarType.FLOAT8E8M0FNU: "Float8_e8m0fnu",
+    ScalarType.UINT32: "UInt32",
+    ScalarType.UINT64: "UInt64",
+}
+
+LAYOUT_TO_C10: dict[int, str] = {
+    Layout.SparseCoo: "Sparse",
+    Layout.SparseCsr: "SparseCsr",
+    Layout.SparseCsc: "SparseCsc",
+    Layout.SparseBsr: "SparseBsr",
+    Layout.SparseBsc: "SparseBsc",
+    Layout._mkldnn: "Mkldnn",
+    Layout.Strided: "Strided",
+}
+
+MEMORY_FORMAT_TO_C10: dict[int, str] = {
+    MemoryFormat.ContiguousFormat: "Contiguous",
+    MemoryFormat.ChannelsLast: "ChannelsLast",
+    MemoryFormat.ChannelsLast3d: "ChannelsLast3d",
+    MemoryFormat.PreserveFormat: "Preserve",
+}
+
+
 @dataclass
 class Device:
     type: Annotated[str, 10]
-    index: Annotated[Optional[int], 20] = None
+    index: Annotated[int | None, 20] = None
 
 
 @_union_dataclass
@@ -76,7 +126,7 @@ class SymExprHint(_Union):
 @dataclass
 class SymExpr:
     expr_str: Annotated[str, 10]
-    hint: Annotated[Optional[SymExprHint], 20] = None
+    hint: Annotated[SymExprHint | None, 20] = None
 
 
 @_union_dataclass
@@ -215,6 +265,7 @@ class Argument(_Union):
     as_nested_tensors: Annotated[list[list[TensorArgument]], 270]
     as_int_lists: Annotated[list[list[int]], 280]
     as_string_to_argument: Annotated[dict[str, "Argument"], 290]
+    as_float_lists: Annotated[list[list[float]], 300]
 
 
 class ArgumentKind(IntEnum):
@@ -228,7 +279,7 @@ class NamedArgument:
     # Argument name from the operator schema
     name: Annotated[str, 10]
     arg: Annotated[Argument, 20]
-    kind: Annotated[Optional[ArgumentKind], 30] = None
+    kind: Annotated[ArgumentKind | None, 30] = None
 
 
 @dataclass
@@ -237,9 +288,9 @@ class Node:
     inputs: Annotated[list[NamedArgument], 20]
     outputs: Annotated[list[Argument], 30]
     metadata: Annotated[dict[str, str], 40]
-    is_hop_single_tensor_return: Annotated[Optional[bool], 50] = None
+    is_hop_single_tensor_return: Annotated[bool | None, 50] = None
     # For BC, default is None so older serialized models without 'name' can be loaded.
-    name: Annotated[Optional[str], 60] = None
+    name: Annotated[str | None, 60] = None
 
 
 @dataclass
@@ -388,8 +439,8 @@ class GraphSignature:
 
 @dataclass
 class RangeConstraint:
-    min_val: Annotated[Optional[int], 10]
-    max_val: Annotated[Optional[int], 20]
+    min_val: Annotated[int | None, 10]
+    max_val: Annotated[int | None, 20]
 
 
 @dataclass
@@ -404,13 +455,13 @@ class ModuleCallSignature:
 
     # This field is used to prettify the graph placeholders
     # after we Ser/Der and retrace
-    forward_arg_names: Annotated[Optional[list[str]], 50] = None
+    forward_arg_names: Annotated[list[str] | None, 50] = None
 
 
 @dataclass
 class ModuleCallEntry:
     fqn: Annotated[str, 10]
-    signature: Annotated[Optional[ModuleCallSignature], 30] = None
+    signature: Annotated[ModuleCallSignature | None, 30] = None
 
 
 @dataclass
@@ -474,7 +525,7 @@ class PayloadMeta:
     # are serialized using pickle.
     use_pickle: Annotated[bool, 30]
     # Custom Objects don't have tensor_meta and will be serialized using pickle
-    tensor_meta: Annotated[Optional[TensorMeta], 40]
+    tensor_meta: Annotated[TensorMeta | None, 40]
 
 
 # The mapping from payload FQN to its metadata.
@@ -504,11 +555,11 @@ class AOTInductorModelPickleData:
     # These fields tell whether floating point inputs/outputs should be converted to
     # a certain type. If None, the dtypes that the AOTInductor engine inferred from the sample
     # inputs are used.
-    floating_point_input_dtype: Annotated[Optional[int], 4] = None
-    floating_point_output_dtype: Annotated[Optional[int], 5] = None
+    floating_point_input_dtype: Annotated[int | None, 4] = None
+    floating_point_output_dtype: Annotated[int | None, 5] = None
 
     # Whether AOTInductor runtime is for CPU.
-    aot_inductor_model_is_cpu: Annotated[Optional[bool], 6] = None
+    aot_inductor_model_is_cpu: Annotated[bool | None, 6] = None
 
 
 @dataclass
diff --git a/torch/_export/serde/schema.yaml b/torch/_export/serde/schema.yaml
index c1b8b1f5b2e77..3b9c4d3051a87 100644
--- a/torch/_export/serde/schema.yaml
+++ b/torch/_export/serde/schema.yaml
@@ -1,5 +1,5 @@
 # @generated by update_schema.py
-# checksum<<5b9bcd54d4ce78ce8bf3ec2e2cad62f19889537e14223d9e9adcfba1f2edc82c>>
+# checksum<<eb9645ae5f539caace53ada70fb262097ce780e6edf685149dafc5a5abe8d31b>>
 AOTInductorModelPickleData:
   kind: struct
   fields:
@@ -81,6 +81,8 @@ Argument:
       type: List[List[int]]
     as_string_to_argument:
       type: Dict[str, Argument]
+    as_float_lists:
+      type: List[List[float]]
 ArgumentKind:
   kind: enum
   fields:
@@ -448,6 +450,9 @@ ScalarType:
     FLOAT8E5M2: 30
     FLOAT8E4M3FNUZ: 31
     FLOAT8E5M2FNUZ: 32
+    FLOAT8E8M0FNU: 33
+    UINT32: 34
+    UINT64: 35
 SchemaVersion:
   kind: struct
   fields:
@@ -560,5 +565,5 @@ UserOutputSpec:
       type: Argument
 SCHEMA_VERSION:
 - 8
-- 17
+- 20
 TREESPEC_VERSION: 1
diff --git a/torch/_export/serde/schema_check.py b/torch/_export/serde/schema_check.py
index a535e0792c570..548ceb5435df8 100644
--- a/torch/_export/serde/schema_check.py
+++ b/torch/_export/serde/schema_check.py
@@ -3,9 +3,10 @@
 import hashlib
 import inspect
 import re
+import types
 import typing
 from enum import IntEnum
-from typing import Annotated, Any, ForwardRef, Optional, Union
+from typing import Annotated, Any, ForwardRef, Union
 
 from torch._export.serde import schema
 from torch._export.serde.union import _Union
@@ -81,7 +82,7 @@ def dump_type(t, level: int) -> tuple[str, str, str]:
                         "map<",
                         ">",
                     )
-                elif o == Union:
+                elif o is Union or o is types.UnionType:
                     if level != 0:
                         raise AssertionError(
                             f"Optional is only supported at the top level, got level={level}"
@@ -134,10 +135,10 @@ def dump_cpp_value(v) -> str:
                     f"Default value {v} is not supported yet in export schema."
                 )
 
-        def dump_field(f) -> tuple[dict[str, Any], str, Optional[str], str, int]:
+        def dump_field(f) -> tuple[dict[str, Any], str, str | None, str, int]:
             t, cpp_type, thrift_type = dump_type(f.type, 0)
             ret = {"type": t}
-            cpp_default: Optional[str] = None
+            cpp_default: str | None = None
             if typing.get_origin(f.type) is not Annotated:
                 raise AssertionError(
                     f"Field {f.name} must be annotated with an integer id."
@@ -394,6 +395,9 @@ class {name} {{
                 raise AssertionError(
                     f"expected SCHEMA_VERSION or TREESPEC_VERSION, got {name}"
                 )
+        elif isinstance(value, dict):
+            # Skip mapping dictionaries used for codegen
+            pass
         else:
             raise AssertionError(f"Unknown variable {name}: {value}")
 
@@ -627,6 +631,120 @@ def _hash_content(s: str):
     return hashlib.sha256(s.strip().encode("utf-8")).hexdigest()
 
 
+def _generate_enum_converters() -> str:
+    """Generate C++ converter functions from serialized enum values to c10 enums."""
+
+    def validate_mapping(
+        enum_class: type[IntEnum],
+        mapping: dict[int, str],
+        enum_name: str,
+        skip_values: set[int],
+    ) -> None:
+        """Validate that all enum values have corresponding c10 mappings."""
+        for member in enum_class:
+            if member.value in skip_values:
+                continue
+            if member.value not in mapping:
+                raise SchemaUpdateError(
+                    f"{enum_name}.{member.name} (value={member.value}) is missing "
+                    f"from {enum_name.upper()}_TO_C10 mapping in schema.py. "
+                    f"Please add the mapping to the c10 enum name."
+                )
+
+    # Validate that all enum values have mappings (except UNKNOWN values)
+    validate_mapping(
+        schema.ScalarType,
+        schema.SCALAR_TYPE_TO_C10,
+        "ScalarType",
+        {schema.ScalarType.UNKNOWN},
+    )
+    validate_mapping(
+        schema.Layout,
+        schema.LAYOUT_TO_C10,
+        "Layout",
+        {schema.Layout.Unknown},
+    )
+    validate_mapping(
+        schema.MemoryFormat,
+        schema.MEMORY_FORMAT_TO_C10,
+        "MemoryFormat",
+        {schema.MemoryFormat.Unknown},
+    )
+
+    def generate_converter(
+        name: str,
+        c10_type: str,
+        mapping: dict[int, str],
+        max_value: int,
+    ) -> str:
+        lines: list[str] = []
+        for i in range(max_value + 1):
+            if i in mapping:
+                lines.append(
+                    f"      static_cast<int>(c10::{c10_type}::{mapping[i]}), // {i}"
+                )
+            else:
+                lines.append(f"      kInvalid, // {i}")
+
+        return f"""
+inline c10::{c10_type} convertSerialized{name}(int serialized_value) {{
+  constexpr int kInvalid = -1;
+  constexpr int k{name}Map[] = {{
+{chr(10).join(lines)}
+  }};
+  constexpr int kMapSize = sizeof(k{name}Map) / sizeof(k{name}Map[0]);
+
+  TORCH_CHECK(
+      serialized_value >= 0 && serialized_value < kMapSize,
+      "Serialized {name} value out of range: ",
+      serialized_value);
+  int result = k{name}Map[serialized_value];
+  TORCH_CHECK(
+      result != kInvalid,
+      "Invalid serialized {name} value: ",
+      serialized_value);
+  return static_cast<c10::{c10_type}>(result);
+}}
+"""
+
+    scalar_type_converter = generate_converter(
+        "ScalarType",
+        "ScalarType",
+        schema.SCALAR_TYPE_TO_C10,
+        max(schema.SCALAR_TYPE_TO_C10.keys()),
+    )
+    layout_converter = generate_converter(
+        "Layout",
+        "Layout",
+        schema.LAYOUT_TO_C10,
+        max(schema.LAYOUT_TO_C10.keys()),
+    )
+    memory_format_converter = generate_converter(
+        "MemoryFormat",
+        "MemoryFormat",
+        schema.MEMORY_FORMAT_TO_C10,
+        max(schema.MEMORY_FORMAT_TO_C10.keys()),
+    )
+
+    return f"""
+#pragma once
+
+#include <c10/core/Layout.h>
+#include <c10/core/MemoryFormat.h>
+#include <c10/core/ScalarType.h>
+#include <c10/util/Exception.h>
+
+// Converter functions from serialized enum values (torch._export.serde.schema)
+// to c10 enums. The serialized format has different enum values than c10.
+
+namespace torch::aot_inductor {{
+{scalar_type_converter}
+{layout_converter}
+{memory_format_converter}
+}} // namespace torch::aot_inductor
+"""
+
+
 @dataclasses.dataclass
 class _Commit:
     result: dict[str, Any]
@@ -635,11 +753,13 @@ class _Commit:
     additions: dict[str, Any]
     subtractions: dict[str, Any]
     base: dict[str, Any]
-    checksum_head: Optional[str]
+    checksum_head: str | None
     cpp_header: str
     cpp_header_path: str
-    thrift_checksum_head: Optional[str]
-    thrift_checksum_real: Optional[str]
+    enum_converter_header: str
+    enum_converter_header_path: str
+    thrift_checksum_head: str | None
+    thrift_checksum_real: str | None
     thrift_checksum_next: str
     thrift_schema: str
     thrift_schema_path: str
@@ -691,6 +811,7 @@ def update_schema():
         dst = {"SCHEMA_VERSION": None, "TREESPEC_VERSION": None}
 
     src, cpp_header, thrift_schema = _staged_schema()
+    enum_converter_header = _generate_enum_converters()
     additions, subtractions = _diff_schema(dst, src)
     # pyrefly: ignore [missing-attribute]
     yaml_path = __package__.replace(".", "/") + "/schema.yaml"
@@ -716,6 +837,9 @@ def update_schema():
         checksum_head=checksum_head,
         cpp_header=cpp_header,
         cpp_header_path=torch_prefix + "csrc/utils/generated_serialization_types.h",
+        enum_converter_header=enum_converter_header,
+        enum_converter_header_path=torch_prefix
+        + "csrc/inductor/aoti_torch/generated_enum_converters.h",
         thrift_checksum_head=thrift_checksum_head,
         thrift_checksum_real=thrift_checksum_real,
         thrift_checksum_next=_hash_content(thrift_schema),
diff --git a/torch/_export/serde/serialize.py b/torch/_export/serde/serialize.py
index 152302d06e97e..12d373c11c7a7 100644
--- a/torch/_export/serde/serialize.py
+++ b/torch/_export/serde/serialize.py
@@ -13,13 +13,14 @@
 import operator
 import re
 import traceback
+import types
 import typing
 from collections import namedtuple, OrderedDict
 from collections.abc import Callable, Iterable, Iterator, Sequence
 from contextlib import contextmanager
 from dataclasses import dataclass, field
 from enum import Enum
-from typing import Annotated, Any, cast, final, Optional, Union
+from typing import Annotated, Any, cast, final, Optional
 
 import sympy
 
@@ -117,16 +118,16 @@ def _reverse_map(d: dict[Any, Enum]):
     return {v.value: k for k, v in d.items()}
 
 
-MetaType = Union[
-    FakeTensor,
-    int,
-    torch.SymInt,
-    float,
-    torch.SymFloat,
-    bool,
-    torch.SymBool,
-    ep.CustomObjArgument,
-]
+MetaType = (
+    FakeTensor
+    | int
+    | torch.SymInt
+    | float
+    | torch.SymFloat
+    | bool
+    | torch.SymBool
+    | ep.CustomObjArgument
+)
 
 DEFAULT_PICKLE_PROTOCOL = 2
 
@@ -151,6 +152,9 @@ def _reverse_map(d: dict[Any, Enum]):
     torch.float8_e5m2: ScalarType.FLOAT8E5M2,
     torch.float8_e4m3fnuz: ScalarType.FLOAT8E4M3FNUZ,
     torch.float8_e5m2fnuz: ScalarType.FLOAT8E5M2FNUZ,
+    torch.float8_e8m0fnu: ScalarType.FLOAT8E8M0FNU,
+    torch.uint32: ScalarType.UINT32,
+    torch.uint64: ScalarType.UINT64,
 }
 
 
@@ -287,13 +291,13 @@ def deserialize_storage_offset(offset: SymInt) -> int:
     return offset.as_int
 
 
-def _print_sympy(s: Union[torch.SymInt, torch.SymBool, torch.SymFloat, sympy.Expr]):
+def _print_sympy(s: torch.SymInt | torch.SymBool | torch.SymFloat | sympy.Expr):
     if isinstance(s, (torch.SymInt, torch.SymBool, torch.SymFloat)):
         s = s.node.expr
     return sympy.printing.repr.srepr(s)
 
 
-def serialize_sym_int(s: Union[int, torch.SymInt]) -> SymInt:
+def serialize_sym_int(s: int | torch.SymInt) -> SymInt:
     if isinstance(s, (torch.SymInt, sympy.Symbol, int)):
         if symbolic_shapes.is_concrete_int(s):
             return SymInt.create(as_int=int(s))
@@ -317,7 +321,7 @@ def serialize_sym_int(s: Union[int, torch.SymInt]) -> SymInt:
         )
 
 
-def serialize_sym_float(s: Union[float, torch.SymFloat]) -> SymFloat:
+def serialize_sym_float(s: float | torch.SymFloat) -> SymFloat:
     if isinstance(s, (torch.SymFloat, sympy.Symbol, float)):
         if symbolic_shapes.is_concrete_float(s):
             return SymFloat.create(as_float=float(s))
@@ -341,7 +345,7 @@ def serialize_sym_float(s: Union[float, torch.SymFloat]) -> SymFloat:
         )
 
 
-def serialize_sym_bool(s: Union[bool, torch.SymBool]) -> SymBool:
+def serialize_sym_bool(s: bool | torch.SymBool) -> SymBool:
     if isinstance(s, (torch.SymBool, bool)):
         if symbolic_shapes.is_concrete_bool(s):
             return SymBool.create(as_bool=bool(s))
@@ -397,7 +401,7 @@ def _reconstruct_fake_tensor(
 
 
 def serialize_torch_artifact(
-    artifact: Optional[Any], pickle_protocol: int = DEFAULT_PICKLE_PROTOCOL
+    artifact: Any | None, pickle_protocol: int = DEFAULT_PICKLE_PROTOCOL
 ) -> bytes:
     if artifact is None:
         return b""
@@ -420,7 +424,7 @@ def serialize_torch_artifact(
 
 
 def deserialize_torch_artifact(
-    serialized: Union[dict[str, Any], tuple[Any, ...], bytes],
+    serialized: dict[str, Any] | tuple[Any, ...] | bytes,
 ):
     if isinstance(serialized, (dict, tuple)):
         return serialized
@@ -445,7 +449,7 @@ def deserialize_torch_artifact(
     return artifact
 
 
-def _sympy_int_to_int(val: sympy.Expr, adjust: str) -> Optional[int]:
+def _sympy_int_to_int(val: sympy.Expr, adjust: str) -> int | None:
     # Convert simple sympy Integers into concrete int
     if val in (sympy.oo, int_oo):
         return None
@@ -472,7 +476,7 @@ def _sympy_int_to_int(val: sympy.Expr, adjust: str) -> Optional[int]:
         raise RuntimeError(f"Got invalid adjustment {adjust}")
 
 
-def _int_to_sympy_int(val: Optional[int], default) -> sympy.Expr:
+def _int_to_sympy_int(val: int | None, default) -> sympy.Expr:
     # Convert concrete int into simple sympy Integers
     if val is None:
         return default
@@ -1025,7 +1029,7 @@ def serialize_tensor_list_output(node):
     def handle_get_attr(self, node):
         log.debug("[handle_get_attr] %s", node.name)
 
-    def _output_node_at_index(self, node, index) -> Optional[torch.fx.Node]:
+    def _output_node_at_index(self, node, index) -> torch.fx.Node | None:
         user_node = None
         for user in node.users:
             if user.target is not operator.getitem:
@@ -1105,8 +1109,8 @@ def export_nn_module_stack(val):
         return ret
 
     def _serialize_from_node(
-        self, from_node: Optional[list[NodeSource]]
-    ) -> Optional[list[dict[str, Any]]]:
+        self, from_node: list[NodeSource] | None
+    ) -> list[dict[str, Any]] | None:
         """
         Serialize from_node metadata from a list of NodeSource objects to a list of dictionaries.
         """
@@ -1245,7 +1249,7 @@ def is_sym_bool_arg(self, arg) -> bool:
         )
 
     # should be torch._C.JitType but that annotation is busted
-    def serialize_input(self, arg, arg_type: Optional[Any] = None) -> Argument:
+    def serialize_input(self, arg, arg_type: Any | None = None) -> Argument:
         import torch._inductor.ir as inductor_ir
 
         inductor_tensor_buffers = (
@@ -1518,6 +1522,12 @@ def serialize_optional_tensor_args(a):
             ):
                 # list of int tuples
                 return Argument.create(as_int_lists=[list(t) for t in arg])
+            elif all(
+                isinstance(a, (list, tuple)) and all(isinstance(x, float) for x in a)
+                for a in arg
+            ):
+                # list of float lists (List[List[float]])
+                return Argument.create(as_float_lists=[list(t) for t in arg])
             else:
                 raise SerializeError(
                     f"Unsupported list/tuple argument type: {[type(a) for a in arg]}"
@@ -1974,17 +1984,15 @@ def _is_single_tensor_list_return(target: Any) -> bool:
                         f"expected ListType with TensorType element, got {type(return_schema.real_type).__name__}"
                     )
                 user_node = self._output_node_at_index(node, idx)
-                if user_node is None:
-                    raise AssertionError(
-                        f"user_node should not be None for list output at index {idx}"
-                    )
-
                 args = []
                 for i, m in enumerate(meta):
                     if m is None:
                         continue
-                    sub_user_node_name = self._output_node_name_at_index(user_node, i)
-                    args.append(self.serialize_tensor_output(sub_user_node_name, m))
+                    if user_node is None:
+                        name = f"{node.name}_unused_{idx}_{i}"
+                    else:
+                        name = self._output_node_name_at_index(user_node, i)
+                    args.append(self.serialize_tensor_output(name, m))
                 output_arguments.append(Argument.create(as_tensors=args))
             elif isinstance(meta, (int, SymInt, float, SymFloat)):
                 user_node_name = self._output_node_name_at_index(node, idx)
@@ -2008,11 +2016,6 @@ def serialize_hoo_outputs(self, node: torch.fx.Node) -> list[Argument]:
                 user_node = self._output_node_at_index(node, i)
                 if isinstance(element_meta_val, list):
                     # e.g "-> Tensor[]"
-                    if user_node is None:
-                        raise AssertionError(
-                            f"user_node should not be None for list output at index {i}"
-                        )
-
                     tensors = []
                     for j, m in enumerate(element_meta_val):
                         if not isinstance(m, torch.Tensor):
@@ -2020,7 +2023,10 @@ def serialize_hoo_outputs(self, node: torch.fx.Node) -> list[Argument]:
                                 f"Serialize list output with type {type(m)} nyi"
                             )
 
-                        name = self._output_node_name_at_index(user_node, j)
+                        if user_node is None:
+                            name = f"{node.name}_unused_{i}_{j}"
+                        else:
+                            name = self._output_node_name_at_index(user_node, j)
                         tensors.append(self.serialize_tensor_output(name, m))
                     outputs.append(Argument.create(as_tensors=tensors))
 
@@ -2166,7 +2172,7 @@ def serialize(self, graph_module: torch.fx.GraphModule) -> GraphModule:
 class ExportedProgramSerializer(metaclass=Final):
     def __init__(
         self,
-        opset_version: Optional[dict[str, int]] = None,
+        opset_version: dict[str, int] | None = None,
         pickle_protocol: int = DEFAULT_PICKLE_PROTOCOL,
     ):
         self.opset_version: dict[str, int] = {}
@@ -2177,7 +2183,14 @@ def __init__(
 
         self.pickle_protocol = pickle_protocol
 
-    def serialize(self, exported_program: ep.ExportedProgram) -> _SerializedProgram:
+    def serialize(
+        self,
+        exported_program: ep.ExportedProgram,
+        *,
+        serialize_state_dict: bool = True,
+        serialize_constants: bool = True,
+        serialize_example_inputs: bool = True,
+    ) -> _SerializedProgram:
         """
         Args:
             exported_program: Exported Program to serialize
@@ -2221,13 +2234,29 @@ def serialize(self, exported_program: ep.ExportedProgram) -> _SerializedProgram:
         new_state_dict = remove_proxy_from_state_dict(
             exported_program.state_dict, in_place=False
         )
+        serialized_state_dict = b""
+        if serialize_state_dict:
+            serialized_state_dict = serialize_torch_artifact(
+                new_state_dict, self.pickle_protocol
+            )
+
+        serialized_constants = b""
+        if serialize_constants:
+            serialized_constants = serialize_torch_artifact(
+                constants, self.pickle_protocol
+            )
+
+        serialized_example_inputs = b""
+        if serialize_example_inputs:
+            serialized_example_inputs = serialize_torch_artifact(
+                exported_program.example_inputs, self.pickle_protocol
+            )
+
         return _SerializedProgram(
             serialized_ep,
-            serialize_torch_artifact(new_state_dict, self.pickle_protocol),
-            serialize_torch_artifact(constants, self.pickle_protocol),
-            serialize_torch_artifact(
-                exported_program.example_inputs, self.pickle_protocol
-            ),
+            serialized_state_dict,
+            serialized_constants,
+            serialized_example_inputs,
         )
 
 
@@ -2239,9 +2268,9 @@ class Result:
         signature: ep.ExportGraphSignature
         module_call_graph: list[ep.ModuleCallEntry]
         names_to_symbols: dict[str, sympy.Symbol]
-        state_dict: dict[str, Union[torch.Tensor, torch.nn.Parameter]]
+        state_dict: dict[str, torch.Tensor | torch.nn.Parameter]
         constants: dict[str, _ConstantAttributeType]
-        example_inputs: Optional[tuple[tuple[torch.Tensor, ...], dict[str, Any]]]
+        example_inputs: tuple[tuple[torch.Tensor, ...], dict[str, Any]] | None
 
     def __init__(self) -> None:
         self.serialized_name_to_node: dict[str, torch.fx.Node] = {}
@@ -2306,7 +2335,7 @@ def deserialize_operator(self, serialized_target: str):
         return target
 
     def _parse_sym_expr(
-        self, expr_str: str, hint: Optional[Union[int, bool, float]] = None
+        self, expr_str: str, hint: int | bool | float | None = None
     ) -> sympy.Expr:
         """
         Parses and does bottom-up processing of sympy.Expr nodes,
@@ -2314,7 +2343,7 @@ def _parse_sym_expr(
         """
 
         def _process_sym_expr(
-            sym: sympy.Expr, hint: Optional[Union[int, bool, float]] = None
+            sym: sympy.Expr, hint: int | bool | float | None = None
         ) -> sympy.Expr:
             if sym.is_Integer or sym.is_Float or sym.is_Boolean:  # base case
                 return sym
@@ -2354,7 +2383,7 @@ def _process_sym_expr(
         )
         return _process_sym_expr(expr, hint)
 
-    def deserialize_sym_int(self, s: SymInt) -> Union[int, torch.SymInt]:
+    def deserialize_sym_int(self, s: SymInt) -> int | torch.SymInt:
         val = s.value
         if s.type == "as_expr":
             if val.hint is None:
@@ -2377,7 +2406,7 @@ def deserialize_sym_int(self, s: SymInt) -> Union[int, torch.SymInt]:
                 f"SymInt has invalid field type {s.type} with value {s.value}"
             )
 
-    def deserialize_sym_float(self, s: SymFloat) -> Union[float, torch.SymFloat]:
+    def deserialize_sym_float(self, s: SymFloat) -> float | torch.SymFloat:
         val = s.value
         if s.type == "as_expr":
             hint = val.hint.as_float if val.hint else None
@@ -2392,7 +2421,7 @@ def deserialize_sym_float(self, s: SymFloat) -> Union[float, torch.SymFloat]:
                 f"SymFloat has invalid field type {s.type} with value {s.value}"
             )
 
-    def deserialize_sym_bool(self, s: SymBool) -> Union[bool, torch.SymBool]:
+    def deserialize_sym_bool(self, s: SymBool) -> bool | torch.SymBool:
         val = s.value
         if s.type == "as_expr":
             expr = self._parse_sym_expr(val.expr_str)
@@ -2430,7 +2459,7 @@ def deserialize_script_obj_meta(
             class_fqn=script_obj_meta.class_fqn,
         )
 
-    def deserialize_graph_output(self, output) -> Optional[Union[torch.fx.Node, int]]:
+    def deserialize_graph_output(self, output) -> torch.fx.Node | int | None:
         if output.type == "as_tensor":
             return self.serialized_name_to_node[output.as_tensor.name]
         elif output.type == "as_sym_int":
@@ -2823,12 +2852,12 @@ def deserialize_signature(self, sig: GraphSignature) -> ep.ExportGraphSignature:
     def deserialize(
         self,
         serialized_graph_module: GraphModule,
-        serialized_state_dict: Union[dict[str, torch.Tensor], bytes],
-        constants: Union[dict[str, Any], bytes],
-        example_inputs: Optional[
-            Union[tuple[tuple[torch.Tensor, ...], dict[str, Any]], bytes]
-        ] = None,
-        symbol_name_to_range: Optional[dict[str, symbolic_shapes.ValueRanges]] = None,
+        serialized_state_dict: dict[str, torch.Tensor] | bytes,
+        constants: dict[str, Any] | bytes,
+        example_inputs: tuple[tuple[torch.Tensor, ...], dict[str, Any]]
+        | bytes
+        | None = None,
+        symbol_name_to_range: dict[str, symbolic_shapes.ValueRanges] | None = None,
     ) -> Result:
         global _CURRENT_DESERIALIZER
         if _CURRENT_DESERIALIZER is not None:
@@ -3073,6 +3102,8 @@ def deserialize_input(self, inp: Argument) -> Any:
             elif typ_ == "as_int_lists":
                 # Convert list of lists back to list of tuples for Triton grids
                 return [tuple(dims) for dims in value]
+            elif typ_ == "as_float_lists":
+                return [list(floats) for floats in value]
             elif typ_ == "as_nested_tensors":
                 # nested list of tensors (List[List[Tensor]])
                 return [
@@ -3201,7 +3232,7 @@ def generate_getitem(
         self,
         meta_val,
         fx_node: torch.fx.Node,
-        arg: Union[TensorArgument, SymIntArgument, SymFloatArgument],
+        arg: TensorArgument | SymIntArgument | SymFloatArgument,
         idx: int,
         deserialized_metadata: dict[str, Any],
     ):
@@ -3392,8 +3423,8 @@ def metadata_split(metadata):
         return ret
 
     def _deserialize_from_node(
-        self, from_node_data: Optional[list[dict[str, Any]]]
-    ) -> Optional[list[NodeSource]]:
+        self, from_node_data: list[dict[str, Any]] | None
+    ) -> list[NodeSource] | None:
         """
         Deserialize from_node metadata from JSON data.
         """
@@ -3460,7 +3491,7 @@ def deserialize_module_call_graph(
 
 @final
 class ExportedProgramDeserializer(metaclass=Final):
-    def __init__(self, expected_opset_version: Optional[dict[str, int]] = None):
+    def __init__(self, expected_opset_version: dict[str, int] | None = None):
         self.expected_opset_version: dict[str, int] = {}
         if expected_opset_version:
             self.expected_opset_version.update(expected_opset_version)
@@ -3487,11 +3518,11 @@ def deserialize_range_constraints(
     def deserialize(
         self,
         exported_program: ExportedProgram,
-        state_dict: Union[dict[str, torch.Tensor], bytes],
-        constants: Union[dict[str, torch.Tensor], bytes],
-        example_inputs: Optional[
-            Union[tuple[tuple[torch.Tensor, ...], dict[str, Any]], bytes]
-        ] = None,
+        state_dict: dict[str, torch.Tensor] | bytes,
+        constants: dict[str, torch.Tensor] | bytes,
+        example_inputs: tuple[tuple[torch.Tensor, ...], dict[str, Any]]
+        | bytes
+        | None = None,
         *,
         _unsafe_skip_version_check=False,
     ) -> ep.ExportedProgram:
@@ -3590,13 +3621,22 @@ def _to_json_bytes(obj: Any) -> bytes:
 
 def serialize(
     exported_program: ep.ExportedProgram,
-    opset_version: Optional[dict[str, int]] = None,
+    opset_version: dict[str, int] | None = None,
     pickle_protocol: int = DEFAULT_PICKLE_PROTOCOL,
+    *,
+    serialize_state_dict: bool = True,
+    serialize_constants: bool = True,
+    serialize_example_inputs: bool = True,
 ) -> SerializedArtifact:
     with _enable_graph_inputs_of_type_nn_module(exported_program.example_inputs):
         serialized_program = ExportedProgramSerializer(
             opset_version, pickle_protocol
-        ).serialize(exported_program)
+        ).serialize(
+            exported_program,
+            serialize_state_dict=serialize_state_dict,
+            serialize_constants=serialize_constants,
+            serialize_example_inputs=serialize_example_inputs,
+        )
     if not isinstance(serialized_program.exported_program, ExportedProgram):
         raise AssertionError(
             f"expected ExportedProgram, got {type(serialized_program.exported_program).__name__}"
@@ -3628,7 +3668,9 @@ def _dict_to_dataclass(cls, data):
         raise AssertionError(f"Unresolved class type: '{cls}'.")
     if typing.get_origin(cls) is Annotated:
         return _dict_to_dataclass(cls.__origin__, data)
-    if typing.get_origin(cls) == typing.Union and type(None) in typing.get_args(cls):
+    if typing.get_origin(cls) in (typing.Union, types.UnionType) and type(
+        None
+    ) in typing.get_args(cls):
         if data is None:
             return None
         ty_args = typing.get_args(cls)
@@ -3682,7 +3724,7 @@ def _bytes_to_dataclass(cls: Any, artifact_bytes: bytes) -> Any:
 
 def deserialize(
     artifact: SerializedArtifact,
-    expected_opset_version: Optional[dict[str, int]] = None,
+    expected_opset_version: dict[str, int] | None = None,
     *,
     _unsafe_skip_version_check=False,
 ) -> ep.ExportedProgram:
@@ -3760,6 +3802,8 @@ def _get_argument(a: Argument):
             return None
         elif a.type == "as_int_lists":
             return None
+        elif a.type == "as_float_lists":
+            return None
         elif a.type == "as_string_to_argument":
             return None
         elif a.type == "as_nested_tensors":
@@ -3786,7 +3830,7 @@ class Edges:
         rank: dict[str, int] = {}
         ret: list[Node] = []
 
-        def get_name(a) -> Optional[str]:
+        def get_name(a) -> str | None:
             if a is None:
                 return None
             if isinstance(a, TensorArgument):
@@ -4026,7 +4070,7 @@ def replace_use(a):
 
 
 def canonicalize(
-    ep: ExportedProgram, constants: Optional[set[str]] = None
+    ep: ExportedProgram, constants: set[str] | None = None
 ) -> ExportedProgram:
     """
     Normalize a serialized ExportedProgram, so that different eager program which
@@ -4069,7 +4113,7 @@ def canonicalize(
             f"graph.outputs length {len(graph.outputs)} != signature.output_specs length {len(signature.output_specs)}"
         )
 
-    def rank_input(inp) -> tuple[int, Optional[str], int]:
+    def rank_input(inp) -> tuple[int, str | None, int]:
         idx, (_arg, spec) = inp
         if not isinstance(spec, InputSpec):
             raise AssertionError(f"expected InputSpec, got {type(spec).__name__}")
@@ -4090,7 +4134,7 @@ def rank_input(inp) -> tuple[int, Optional[str], int]:
         else:
             raise AssertionError(f"Unknown input type: {spec}")
 
-    def rank_output(out) -> tuple[int, Optional[str], int]:
+    def rank_output(out) -> tuple[int, str | None, int]:
         idx, (_arg, spec) = out
         if not isinstance(spec, OutputSpec):
             raise AssertionError(f"expected OutputSpec, got {type(spec).__name__}")
diff --git a/torch/_export/tools.py b/torch/_export/tools.py
index b254fd62e3b2d..9e846de96c465 100644
--- a/torch/_export/tools.py
+++ b/torch/_export/tools.py
@@ -2,7 +2,7 @@
 import logging
 import warnings
 from collections.abc import Iterable
-from typing import Any, Optional
+from typing import Any
 
 import torch
 import torch.export
@@ -19,7 +19,7 @@ def _generate_inputs_for_submodules(
     model: torch.nn.Module,
     target_submodules: Iterable[str],
     args: tuple[Any, ...],
-    kwargs: Optional[dict[str, Any]] = None,
+    kwargs: dict[str, Any] | None = None,
 ) -> dict[str, tuple[Any, Any]]:
     """
     Generate inputs for targeting submdoules in the given model. Note that if two submodules refer to the same obj, this
@@ -63,11 +63,11 @@ def pre_forward(module, module_args, module_kwargs):
 def report_exportability(
     mod: torch.nn.Module,
     args: tuple[Any, ...],
-    kwargs: Optional[dict[str, Any]] = None,
+    kwargs: dict[str, Any] | None = None,
     *,
     strict: bool = True,
     pre_dispatch: bool = False,
-) -> dict[str, Optional[Exception]]:
+) -> dict[str, Exception | None]:
     """
     Report exportability issues for a module in one-shot.
 
@@ -94,7 +94,7 @@ def report_exportability(
     submod_inputs = _generate_inputs_for_submodules(mod, all_submod_names, args, kwargs)
 
     tried_module_types = set()
-    report: dict[str, Optional[Exception]] = {}
+    report: dict[str, Exception | None] = {}
 
     def try_export(module, module_name, args, kwargs):
         nonlocal submod_inputs, report, strict, pre_dispatch, tried_module_types
diff --git a/torch/_export/utils.py b/torch/_export/utils.py
index 022ec0b990fd3..c0ebd5d5938fb 100644
--- a/torch/_export/utils.py
+++ b/torch/_export/utils.py
@@ -12,7 +12,7 @@
 from collections.abc import Callable, Iterable
 from contextlib import contextmanager
 from inspect import ismethod, Parameter
-from typing import Any, Optional, TYPE_CHECKING, Union
+from typing import Any, TYPE_CHECKING
 
 import torch
 from torch._guards import detect_fake_mode
@@ -40,7 +40,7 @@
     _deregister_pytree_node,
     _register_pytree_node,
     Context,
-    FlattenFunc,
+    FlattenFn,
     FromDumpableContextFn,
     GetAttrKey,
     KeyPath,
@@ -49,7 +49,7 @@
     SequenceKey,
     ToDumpableContextFn,
     tree_flatten_with_path,
-    UnflattenFunc,
+    UnflattenFn,
 )
 
 
@@ -347,12 +347,12 @@ def get_keystr(key_path: KeyPath) -> str:
 
 
 def _check_symint(
-    symint: Union[int, torch.SymInt],
+    symint: int | torch.SymInt,
     arg: int,
     range_constraints,
     unification_map,
     keypath: KeyPath,
-    i: Optional[int] = None,
+    i: int | None = None,
 ) -> None:
     from torch.export.dynamic_shapes import _IntWrapper
 
@@ -489,12 +489,12 @@ def _check_input_constraints_for_graph(
 
 def register_dataclass_as_pytree_node(
     cls: type[Any],
-    flatten_fn: Optional[FlattenFunc] = None,
-    unflatten_fn: Optional[UnflattenFunc] = None,
+    flatten_fn: FlattenFn | None = None,
+    unflatten_fn: UnflattenFn | None = None,
     *,
-    serialized_type_name: Optional[str] = None,
-    to_dumpable_context: Optional[ToDumpableContextFn] = None,
-    from_dumpable_context: Optional[FromDumpableContextFn] = None,
+    serialized_type_name: str | None = None,
+    to_dumpable_context: ToDumpableContextFn | None = None,
+    from_dumpable_context: FromDumpableContextFn | None = None,
     return_none_fields: bool = False,
 ) -> None:
     if not dataclasses.is_dataclass(cls):
@@ -557,7 +557,7 @@ def is_param(program: "ExportedProgram", node: torch.fx.Node) -> bool:
 def get_param(
     program: "ExportedProgram",
     node: torch.fx.Node,
-) -> Optional[torch.nn.Parameter]:
+) -> torch.nn.Parameter | None:
     """
     Returns the parameter associated with the given node in the exported program.
     Returns None if the node is not a parameter within the exported program
@@ -581,7 +581,7 @@ def is_buffer(program: "ExportedProgram", node: torch.fx.Node) -> bool:
 def get_buffer(
     program: "ExportedProgram",
     node: torch.fx.Node,
-) -> Optional[torch.Tensor]:
+) -> torch.Tensor | None:
     """
     Returns the buffer associated with the given node in the exported program.
     Returns None if the node is not a buffer within the exported program
@@ -611,7 +611,7 @@ def is_lifted_tensor_constant(
 def get_lifted_tensor_constant(
     program: "ExportedProgram",
     node: torch.fx.Node,
-) -> Optional[torch.Tensor]:
+) -> torch.Tensor | None:
     """
     Returns the lifted tensor constant associated with the given node in the exported program.
     Returns None if the node is not a lifted tensor constant within the exported program
@@ -628,7 +628,7 @@ def get_lifted_tensor_constant(
 
 def sequential_split(
     gm: torch.fx.GraphModule,
-    node_call_back: Callable[[torch.fx.Node], Union[torch.fx.Node, bool]],
+    node_call_back: Callable[[torch.fx.Node], torch.fx.Node | bool],
 ) -> torch.fx.GraphModule:
     """
     sequential_split creates a new graph module that splits the input graph module into multiple submodules
@@ -759,7 +759,7 @@ def apply_runtime_assertion_pass(gm: torch.fx.GraphModule, graph_signature):
 
 def nodes_first(
     nodes: list[torch.fx.Node], node_call_back=None
-) -> Optional[torch.fx.Node]:
+) -> torch.fx.Node | None:
     """
     Returns the first node that matches the node_call_back. If no node matches, returns None.
     When node_call_back is None, returns the first node in the node list.
@@ -803,7 +803,7 @@ def _update_gm_meta_if_possible(gm: torch.fx.GraphModule, mod: torch.nn.Module)
         gm.meta.update({"custom": mod.meta["custom"]})
 
 
-def node_inline_(call_mod_node: torch.fx.Node) -> Optional[torch.fx.GraphModule]:
+def node_inline_(call_mod_node: torch.fx.Node) -> torch.fx.GraphModule | None:
     """
     Inline the submodule of the given node into the parent module.
     Note: we only support the case where submodule takes tensors inputs.
@@ -1211,7 +1211,7 @@ def remove_proxy_from_state_dict(state_dict: dict, in_place: bool) -> dict:
 
 def _detect_fake_mode_from_gm(
     gm: torch.fx.GraphModule,
-) -> Optional[torch._subclasses.fake_tensor.FakeTensorMode]:
+) -> torch._subclasses.fake_tensor.FakeTensorMode | None:
     """
     For a given graph module, we look at the "val" of placeholder nodes to find the fake inputs.
     Additionally, if gm doesn't have placeholders, we further look at the "example_value" or "val" of other nodes.
@@ -1434,7 +1434,7 @@ def _compiling_state_context():
 def _fakify_params_buffers(
     fake_mode: FakeTensorMode,
     mod: torch.nn.Module,
-) -> dict[str, Union[torch.Tensor, torch.nn.Parameter]]:
+) -> dict[str, torch.Tensor | torch.nn.Parameter]:
     params_buffers = {
         **dict(mod.named_parameters(remove_duplicate=False)),
         **dict(mod.named_buffers(remove_duplicate=False)),
diff --git a/torch/_functorch/_activation_checkpointing/graph_info_provider.py b/torch/_functorch/_activation_checkpointing/graph_info_provider.py
index 2aa0ea6291ed6..0294e32f292e5 100644
--- a/torch/_functorch/_activation_checkpointing/graph_info_provider.py
+++ b/torch/_functorch/_activation_checkpointing/graph_info_provider.py
@@ -1,4 +1,4 @@
-from typing import Any, Optional
+from typing import Any
 
 import networkx as nx
 
@@ -23,11 +23,11 @@ def __init__(
         graph_nodes_in_order: list[str],
         graph_edges: list[tuple[str, str]],
         all_recomputable_banned_nodes: list[str],
-        all_node_runtimes: Optional[dict[str, float]] = None,
-        all_node_memories: Optional[dict[str, float]] = None,
-        recorded_knapsack_input_memories: Optional[list[float]] = None,
-        recorded_knapsack_input_runtimes: Optional[list[float]] = None,
-        joint_graph: Optional[Graph] = None,
+        all_node_runtimes: dict[str, float] | None = None,
+        all_node_memories: dict[str, float] | None = None,
+        recorded_knapsack_input_memories: list[float] | None = None,
+        recorded_knapsack_input_runtimes: list[float] | None = None,
+        joint_graph: Graph | None = None,
     ) -> None:
         self.graph_nodes_in_order = graph_nodes_in_order
         self.graph_edges = graph_edges
diff --git a/torch/_functorch/_activation_checkpointing/knapsack_evaluator.py b/torch/_functorch/_activation_checkpointing/knapsack_evaluator.py
index d3621c8207525..86ad04ed2f5c7 100644
--- a/torch/_functorch/_activation_checkpointing/knapsack_evaluator.py
+++ b/torch/_functorch/_activation_checkpointing/knapsack_evaluator.py
@@ -207,7 +207,7 @@ def evaluate_distribution_of_results_for_knapsack_algo(
             knapsack_algo (Callable): The knapsack algorithm to use for evaluation.
             memory_budget_values (List[float]): A list of memory budgets to evaluate.
         """
-        results = list()
+        results = []
         for memory_budget in memory_budget_values:
             _, saved_nodes, recomputed_nodes = knapsack_algo(
                 self._graph_info_provider.get_knapsack_memory_input(),
diff --git a/torch/_functorch/_activation_checkpointing/remat_using_tags_for_fwd_loss_bwd_graph_pass.py b/torch/_functorch/_activation_checkpointing/remat_using_tags_for_fwd_loss_bwd_graph_pass.py
index 346f217d8fda3..82ef44a23ff0f 100644
--- a/torch/_functorch/_activation_checkpointing/remat_using_tags_for_fwd_loss_bwd_graph_pass.py
+++ b/torch/_functorch/_activation_checkpointing/remat_using_tags_for_fwd_loss_bwd_graph_pass.py
@@ -1,17 +1,15 @@
-"""
-AC rematerialize pass: Duplicates checkpointed nodes for backward, then DCE removes unused forward versions.
-"""
+"""AC rematerialize pass: Duplicates recompute nodes for backward, then DCE removes unused forward versions."""
 
+import itertools
+import logging
 from typing import Any, overload
 
 import torch
 import torch.fx as fx
-from torch._functorch import config
 from torch._functorch.compile_utils import raise_getitems
 from torch._functorch.partitioners import (
     cleanup_recompute_tags,
     force_save_bw_mutation_src,
-    force_save_collectives,
     has_recomputable_ops,
     has_recomputable_rng_ops,
     is_not_collective,
@@ -19,6 +17,10 @@
 )
 
 
+log = logging.getLogger(__name__)
+_EMPTY_CUSTOM_META: dict[str, object] = {}
+
+
 def is_impure_node_for_dce(node: fx.Node) -> bool:
     # Check for special collectives that should be treated as pure
     if not is_not_collective(node):
@@ -34,44 +36,100 @@ def is_impure_node_for_dce(node: fx.Node) -> bool:
     return node.is_impure(impure_random)
 
 
-def _is_backward_node(node: fx.Node) -> bool:
-    """Check if node is in backward region via annotation"""
-    return node.meta.get("custom", {}).get("remat_pass_tag", None) == "is_backward"
+def _is_backward_node(node: fx.Node, use_phase: bool = False) -> bool:
+    """Check if node is in backward region.
+
+    If use_phase is True, only checks custom["phase"] == "backward"
+    (user annotation). Otherwise falls back to node.meta["autograd_backward"],
+    which Dynamo adds when tracing torch.autograd.grad.
+    """
+    custom = node.meta.get("custom", _EMPTY_CUSTOM_META)
+    if use_phase:
+        return custom.get("phase") == "backward"
+    return node.meta.get("autograd_backward", False)
+
+
+def _has_user_phase_annotation(gm: fx.GraphModule) -> bool:
+    """Check if any node has the user-level phase: backward annotation."""
+    return any(
+        node.meta.get("custom", _EMPTY_CUSTOM_META).get("phase") == "backward"
+        for node in gm.graph.nodes
+    )
+
+
+def _collect_backward_region_data(
+    gm: fx.GraphModule,
+) -> tuple[bool, int | None, int | None, int]:
+    use_phase = _has_user_phase_annotation(gm)
+    bwd_start: int | None = None
+    bwd_end: int | None = None
+    num_regions = 0
+    in_backward = False
+
+    for idx, node in enumerate(gm.graph.nodes):
+        is_bwd = _is_backward_node(node, use_phase=use_phase)
+        if is_bwd:
+            if bwd_start is None:
+                bwd_start = idx
+            bwd_end = idx + 1
+            if not in_backward:
+                num_regions += 1
+        in_backward = is_bwd
+
+    return use_phase, bwd_start, bwd_end, num_regions
 
 
 def remat_using_tags_for_fwd_loss_bwd_graph(gm: fx.GraphModule) -> fx.GraphModule:
     """
-    Duplicate checkpointed nodes for backward use. DCE removes unused forward versions. We assume that
-    you already annotated your backward region with fx.traceback.annotate({"remat_pass_tag": "is_backward"})
-    which helps us identify the backward region.
+    Duplicate recompute nodes for backward use. DCE removes unused forward versions.
+
+    Backward regions are identified by custom["phase"] == "backward" (user
+    annotation) or node.meta["autograd_backward"] == True (set automatically when
+    Dynamo traces torch.autograd.grad). When the user provides phase
+    annotations, only those annotated regions are used.
+
+    Only a single contiguous backward region is supported. If multiple disjoint
+    backward regions are detected, an error is raised. Consecutive backward
+    operations without non-backward nodes between them are treated as a single
+    backward region.
     """
     if not has_recomputable_ops(gm):
         return gm
 
-    # Find backward boundary and build ordering
-    bwd_start: int | None = None
-    order = {}
-    for idx, node in enumerate(gm.graph.nodes):
-        order[node] = idx
-        if _is_backward_node(node) and bwd_start is None:
-            bwd_start = idx
+    use_phase, bwd_start, bwd_end, num_regions = _collect_backward_region_data(gm)
+    if num_regions > 1:
+        if use_phase:
+            raise RuntimeError(
+                f"Detected {num_regions} disjoint backward regions annotated with "
+                'phase: "backward" but remat only supports a single backward region. '
+                "Please ensure only one contiguous region is annotated."
+            )
+        raise RuntimeError(
+            f"Detected {num_regions} disjoint backward regions in the graph but remat only supports "
+            "a single backward region. This can happen when non-backward computation appears "
+            "between backward sections. Please annotate the real backward with "
+            'torch.fx.traceback.annotate({"phase": "backward"}).'
+        )
 
     if bwd_start is None:
         return gm
+    if bwd_end is None:
+        raise AssertionError(
+            "backward region should end somewhere when there was explicit backward region start."
+        )
+
+    order = {node: idx for idx, node in enumerate(gm.graph.nodes)}
 
     if has_recomputable_rng_ops(gm):
         raise RuntimeError(
-            "Activation checkpoint rematerializing in `forward-loss-backward` graph does not support RNG ops "
-            "in checkpointed regions. Please move RNG operations outside "
-            "of checkpoint regions, or use joint graph mode (where partitioner handles RNG)."
+            "Activation checkpoint rematerialization in `forward-loss-backward` graph does not support RNG ops "
+            "in recompute regions. Please move RNG operations outside "
+            "of recompute regions, or use joint graph mode (where partitioner handles RNG)."
         )
 
     # Use partitioner pass to normalize AC node tags.
     gm = cleanup_recompute_tags(gm, is_default_partition=True)
 
-    if not config.unsafe_allow_optimization_of_collectives:
-        force_save_collectives(gm)
-
     force_save_bw_mutation_src(gm)
 
     new_graph = fx.Graph()
@@ -79,7 +137,7 @@ def remat_using_tags_for_fwd_loss_bwd_graph(gm: fx.GraphModule) -> fx.GraphModul
     recomputed_nodes: dict[fx.Node, fx.Node] = {}
 
     # Insert forward nodes
-    for node in list(gm.graph.nodes)[:bwd_start]:
+    for node in itertools.islice(gm.graph.nodes, 0, bwd_start):
         env[node] = new_graph.node_copy(node, lambda x: env[x])
 
     @overload
@@ -93,37 +151,47 @@ def remat_input(x: object) -> object:
             return x
         return recomputed_nodes.get(x, env[x])
 
-    def gather_checkpointed_deps(node: fx.Node, visited: set[fx.Node]) -> None:
-        if node in visited or node in recomputed_nodes:
-            return
-        visited.add(node)
+    def gather_recompute_deps(node: fx.Node) -> set[fx.Node]:
+        deps: set[fx.Node] = set()
+
+        def _gather(n: fx.Node) -> None:
+            if n in deps or n in recomputed_nodes or not must_recompute(n):
+                return
+            deps.add(n)
+            for inp in n.all_input_nodes:
+                _gather(inp)
+
+        # Can't call _gather(node) directly: node itself may not be must_recompute
+        # (e.g. backward nodes), so _gather would return early without visiting inputs.
         for inp in node.all_input_nodes:
-            if must_recompute(inp):
-                gather_checkpointed_deps(inp, visited)
+            _gather(inp)
+        return deps
 
     # Insert backward nodes
-    for node in list(gm.graph.nodes)[bwd_start:]:
-        # Gather all checkpointed deps needed by this node
-        deps = set()
-        for inp in node.all_input_nodes:
-            if must_recompute(inp):
-                gather_checkpointed_deps(inp, deps)
+    for node in itertools.islice(gm.graph.nodes, bwd_start, bwd_end):
+        # Gather all deps that need to be recomputed for this node
+        deps = gather_recompute_deps(node)
 
         # Insert deps in forward order (guaranteed disjoint from already-inserted)
         # This is not as inefficient as it looks, because we only add fresh dependencies
         # when they are not yet processed as recomputed nodes.
-        for dep in sorted(deps, key=lambda n: order[n]):
-            if dep in recomputed_nodes:
-                raise AssertionError(
-                    f"We shouldn't have recomputed {dep} before, "
-                    f"but found it in recomputed_nodes"
-                )
+        new_deps = sorted(deps, key=lambda n: order[n])
+        if new_deps:
+            log.debug(
+                "To compute backward node %s, recomputing [%s]",
+                node.name,
+                ", ".join(dep.name for dep in new_deps),
+            )
+        for dep in new_deps:
             dup = new_graph.node_copy(dep, remat_input)
             dup.name = dep.name + "_recomputed"
             recomputed_nodes[dep] = dup
 
         env[node] = new_graph.node_copy(node, remat_input)
 
+    for node in itertools.islice(gm.graph.nodes, bwd_end, None):
+        env[node] = new_graph.node_copy(node, lambda x: env[x])
+
     new_gm = torch.fx.GraphModule(gm, new_graph)
 
     # DCE with custom is_impure_node (like default_partition)
diff --git a/torch/_functorch/_activation_offloading/activation_offloading.py b/torch/_functorch/_activation_offloading/activation_offloading.py
index 6296c9f4525a3..8f9cd2203d12f 100644
--- a/torch/_functorch/_activation_offloading/activation_offloading.py
+++ b/torch/_functorch/_activation_offloading/activation_offloading.py
@@ -1,14 +1,8 @@
-"""
-Activation offloading for memory optimization in (more like post) partitioners.
-
-This module provides functionality to offload activations to CPU during forward pass
-and reload them during backward pass, reducing GPU memory usage.
+"""Activation offloading for memory optimization during compilation.
 
-Additional TODO:
-* given the fact that PT2 stream support is in active development, testings should
-  be done once that is more finalized. A issue currently known is that with streams,
-  each iteration will have its own offload streams, but the streams should be shared
-  across the iterations.
+This module provides functionality to offload activations to CPU during the forward
+pass and reload them during the backward pass, reducing GPU memory usage. It can be
+applied to graphs produced by both AOT Autograd partitioners and make_fx-based tracing.
 """
 
 import logging
@@ -17,8 +11,11 @@
 
 import torch
 import torch.fx as fx
-from torch._dynamo.variables.streams import get_current_stream, new_event, new_stream
-from torch._inductor import config as inductor_config
+from torch._functorch._activation_offloading.offload_ops import (  # noqa: F401 -- registers ao::offload, ao::reload, ao::wait_tensor ops
+    offload,
+    reload,
+    wait_tensor,
+)
 from torch._inductor.fx_passes.overlap_scheduling import benchmark_node, is_compute_node
 from torch._subclasses.fake_tensor import extract_tensor_metadata
 from torch.utils._ordered_set import OrderedSet
@@ -36,20 +33,31 @@
 GPU_RELOAD_PREFIX = "gpu_reload_"
 
 
+def _find_all_effective_users(node: fx.Node, op_types: OpTypes) -> OrderedSet[fx.Node]:
+    """Find all effective users of a node, where view ops extend the lifetime
+    of the original node. If a user is a view op, recursively find users of
+    the view."""
+    effective_users: OrderedSet[fx.Node] = OrderedSet()
+    for user in node.users:
+        if user.op == "output":
+            continue
+        effective_users.add(user)
+        if op_types.is_view(user):
+            effective_users.update(_find_all_effective_users(user, op_types))
+    return effective_users
+
+
 @dataclass
 class ReloadNodeInfo:
     """
     Information about backward reload related nodes for each reload operation.
 
-    Pattern: fork → wait_stream → device_put → record_event → join → wait_event
+    Pattern: ao.reload → ao.wait_tensor
 
-    This pattern is divided into two logical groups for optimization purposes:
-    - Reload group (fork → wait_stream → device_put → record_event → join):
-      Performs the actual asynchronous data transfer on a separate stream.
-      These nodes can be moved earlier in the graph to overlap with computation.
-    - Wait group (wait_event):
-      Synchronization point that blocks until the data transfer completes.
-      This must remain at the point where the reloaded data is first needed.
+    - Reload group (ao.reload): Performs the actual asynchronous data transfer.
+      Can be moved earlier in the graph to overlap with computation.
+    - Wait node (ao.wait_tensor): Synchronization point that blocks until the data
+      transfer completes. Must remain at the point where the data is first needed.
     """
 
     reload_group_nodes: list[fx.Node]
@@ -92,21 +100,6 @@ def offload_activation_fw(graph: fx.Graph) -> None:
 
     op_types: OpTypes = get_default_op_list()
 
-    def find_all_effective_users(node: fx.Node) -> OrderedSet[fx.Node]:
-        """
-        Find all effective users of a node, where view ops extend the lifetime of the
-        original node. If a user is a view op, recursively find users of the view.
-        """
-        effective_users: OrderedSet[fx.Node] = OrderedSet()
-        for user in node.users:
-            if user.op == "output":
-                continue
-            effective_users.add(user)
-            if op_types.is_view(user):
-                effective_users.update(find_all_effective_users(user))
-
-        return effective_users
-
     output_node: fx.Node = graph.find_nodes(op="output")[0]
     # pyrefly: ignore [bad-assignment]
     fwd_outputs: tuple[fx.Node, ...] = output_node.args[
@@ -122,8 +115,7 @@ def find_all_effective_users(node: fx.Node) -> OrderedSet[fx.Node]:
             continue
 
         # Find insertion point, which is the last use
-        all_effective_users: OrderedSet[fx.Node] = find_all_effective_users(node)
-        if all_effective_users := find_all_effective_users(node):
+        if all_effective_users := _find_all_effective_users(node, op_types):
             last_user = max(all_effective_users, key=lambda n: node_to_index[n])
         else:
             last_user: fx.Node = node
@@ -193,6 +185,116 @@ def reload_activation_bw(graph: fx.Graph) -> None:
                 user.replace_input_with(node, gpu_node)
 
 
+def offload_activation_fw_async(graph: fx.Graph) -> None:
+    """Insert async CPU offload operations in the forward pass graph.
+
+    Uses ao.offload + ao.wait_tensor ops which encapsulate stream management
+    internally, producing a clean 2-node IR per offloaded tensor.
+    """
+
+    op_types: OpTypes = get_default_op_list()
+
+    output_node: fx.Node = graph.find_nodes(op="output")[0]
+    # pyrefly: ignore [bad-assignment]
+    fwd_outputs: tuple[fx.Node, ...] = output_node.args[
+        0
+    ]  # pyrefly: ignore [bad-assignment]
+    node_to_offload: dict[fx.Node, fx.Node] = dict()
+    node_to_index: dict[fx.Node, int] = {
+        node: idx for idx, node in enumerate(graph.nodes)
+    }
+
+    if not any(n.meta.get("saved_for_offloading", False) for n in fwd_outputs):
+        return
+
+    for node in fwd_outputs:
+        if node.meta.get("saved_for_offloading", False) is False:
+            continue
+
+        if all_effective_users := _find_all_effective_users(node, op_types):
+            last_user = max(all_effective_users, key=lambda n: node_to_index[n])
+        else:
+            last_user: fx.Node = node
+
+        with graph.inserting_after(last_user):
+            offload_node: fx.Node = graph.call_function(
+                torch.ops.ao.offload.default,
+                args=(node,),
+                name=f"async_{CPU_OFFLOAD_PREFIX}{node.name}",
+            )
+            offload_node.meta["val"] = node.meta["val"].to(torch.device("cpu"))
+            offload_node.meta["tensor_meta"] = extract_tensor_metadata(
+                offload_node.meta["val"]
+            )
+        # The keepalive=node arg extends the GPU tensor's lifetime in the
+        # graph so the allocator doesn't reclaim it before the async D2H
+        # copy completes.
+        with graph.inserting_after(offload_node):
+            wait_node: fx.Node = graph.call_function(
+                torch.ops.ao.wait_tensor.default,
+                args=(offload_node, node),
+                name=CPU_OFFLOAD_PREFIX + str(node.name),
+            )
+            wait_node.meta["val"] = offload_node.meta["val"]
+            wait_node.meta["tensor_meta"] = offload_node.meta["tensor_meta"]
+
+        node_to_offload[node] = wait_node
+
+    output_node.update_arg(
+        0, tuple(node_to_offload.get(node, node) for node in fwd_outputs)
+    )
+
+
+def reload_activation_bw_async(graph: fx.Graph) -> None:
+    """Insert async GPU reload operations in the backward pass graph.
+
+    Uses ao.reload + ao.wait_tensor ops which encapsulate stream management internally,
+    producing a clean 2-node IR per reloaded tensor.
+    """
+
+    node_to_index: dict[fx.Node, int] = {
+        node: idx for idx, node in enumerate(graph.nodes)
+    }
+
+    nodes_to_reload = [
+        n
+        for n in graph.find_nodes(op="placeholder")
+        if n.meta.get("saved_for_offloading", False)
+    ]
+    if not nodes_to_reload:
+        return
+
+    for node in nodes_to_reload:
+        if not node.users:
+            raise RuntimeError(
+                f"Offloaded tensor {node.name} has no users in the backward graph"
+            )
+        insert_point: fx.Node = min(node.users.keys(), key=lambda n: node_to_index[n])
+
+        original_device: torch.device = node.meta["original_device"]
+        with graph.inserting_before(insert_point):
+            reload_node: fx.Node = graph.call_function(
+                torch.ops.ao.reload.default,
+                args=(node, original_device),
+                name=f"async_{str(node.name).replace(CPU_OFFLOAD_PREFIX, GPU_RELOAD_PREFIX)}",
+            )
+            reload_node.meta["val"] = node.meta["val"].to(original_device)
+            reload_node.meta["tensor_meta"] = extract_tensor_metadata(
+                reload_node.meta["val"]
+            )
+            wait_node: fx.Node = graph.call_function(
+                torch.ops.ao.wait_tensor.default,
+                args=(reload_node,),
+                name=str(node.name).replace(CPU_OFFLOAD_PREFIX, GPU_RELOAD_PREFIX),
+            )
+            wait_node.meta["val"] = reload_node.meta["val"]
+            wait_node.meta["tensor_meta"] = reload_node.meta["tensor_meta"]
+
+        for user in list(node.users.keys()):
+            if user != reload_node:
+                user.replace_input_with(node, wait_node)
+
+
 def can_offload(
     node: fx.Node,
     fwd_outputs: OrderedSet[fx.Node],
@@ -341,245 +443,129 @@ def offload_chosen_sets(
     reload_activation_bw(bwd_module.graph)
 
 
-def add_forward_offload_stream_ops(graph: fx.Graph) -> None:
+def offload_chosen_sets_async(
+    fwd_module: fx.GraphModule,
+    bwd_module: fx.GraphModule,
+) -> None:
     """
-    Add stream operations for forward pass CPU offloading.
-
-    Pattern: record_event → fork → wait_event → record_stream → device_put → record_event_2 → join → wait_event_2
-
-    This ensures that:
-    1. Offloading waits for the last use to complete (record_event on default stream)
-    2. Offloading happens on a separate stream (fork → wait_event → device_put)
-    3. The tensor is marked as used in the offload stream (record_stream)
-    4. Execution returns to the default stream after offloading and
-       waits for offload to complete (record_event_2 → join → wait_event_2)
+    Add async offload and reload nodes using ao ops.
 
-    NOTE: For stream optimization and overlapping compute with communication,
-          the "wait_event_2" ops can be sinked to the end of the graph.
-
-    Args:
-        graph: The forward graph to modify
+    Uses ao.offload/ao.reload + ao.wait_tensor which encapsulate stream management,
+    instead of device_put + explicit stream operations. Can be applied to
+    partitioned forward/backward graphs or to a joint graph produced by make_fx.
     """
 
-    # Find all CPU offload nodes
-    offload_nodes: list[fx.Node] = [
-        node
-        for node in graph.nodes
-        if CPU_OFFLOAD_PREFIX in node.name and node.op == "call_function"
-    ]
-    if not offload_nodes:
-        return
-
-    # Get default stream id and offload stream id
-    current_stream_id: int = get_current_stream(
-        offload_nodes[0].args[0].meta["val"].device  # type: ignore[assignment]
-    )
-    offload_stream_id: int = new_stream()
+    offload_activation_fw_async(fwd_module.graph)
 
-    for offload_node in offload_nodes:
-        offload_ready_event_id: int = new_event()
-        offload_completion_event_id: int = new_event()
-
-        # Get the tensor being offloaded
-        tensor_node: fx.Node = offload_node.args[0]  # type: ignore[assignment]
-
-        with graph.inserting_before(offload_node):
-            # Record event on default stream to ensure last use completes
-            graph.call_function(
-                torch.ops.streams.record_event.default,
-                args=(offload_ready_event_id, current_stream_id),
-            )
-            # Fork to offload stream
-            graph.call_function(
-                torch.ops.streams.fork.default,
-                args=(current_stream_id, offload_stream_id),
-                name=f"stream_in_{offload_node.name}",
-            )
-            # Wait for the event on offload stream
-            graph.call_function(
-                torch.ops.streams.wait_event.default,
-                args=(offload_ready_event_id, offload_stream_id),
-            )
-            # Inform the CUDA Caching Allocator that this tensor will be accessed in the
-            # offload stream. Without this, the program may prematurely free its memory
-            # even though the async offload operation is still in progress, and this can
-            # lead to memory corruption, especially with reordering for compute and
-            # communication overlaps.
-            graph.call_function(
-                torch.ops.streams.record_stream.default,
-                args=(tensor_node, offload_stream_id),
-                name=f"record_stream_{tensor_node.name}",
-            )
-        with graph.inserting_after(offload_node):
-            # Record event on offload stream after device_put completes
-            record_event_node = graph.call_function(
-                torch.ops.streams.record_event.default,
-                args=(offload_completion_event_id, offload_stream_id),
-            )
-        with graph.inserting_after(record_event_node):
-            # Join back to default stream
-            join_node = graph.call_function(
-                torch.ops.streams.join.default,
-                args=(offload_stream_id, current_stream_id),
-                name=f"stream_out_{offload_node.name}",
-            )
-        with graph.inserting_after(join_node):
-            # Wait for the offload to complete on default stream
-            graph.call_function(
-                torch.ops.streams.wait_event.default,
-                args=(offload_completion_event_id, current_stream_id),
-            )
+    # Replace backward graph placeholders with their offloaded (CPU) counterparts.
+    # For each offloaded forward output, find the matching backward input and swap
+    # it with a new placeholder carrying the CPU tensor's metadata, then mark it
+    # for reloading.
+    bwd_inputs: dict[str, fx.Node] = {
+        node.name: node for node in bwd_module.graph.find_nodes(op="placeholder")
+    }
+    for fwd_node in fwd_module.graph.find_nodes(op="output")[0].args[0]:
+        if CPU_OFFLOAD_PREFIX not in fwd_node.name:
+            continue
 
+        bwd_node: fx.Node = bwd_inputs[fwd_node.name.replace(CPU_OFFLOAD_PREFIX, "")]
+        with bwd_module.graph.inserting_after(bwd_node):
+            bwd_offload_node: fx.Node = bwd_module.graph.placeholder(name=fwd_node.name)
 
-def add_backward_reload_stream_ops(graph: fx.Graph) -> None:
-    """
-    Add stream operations for backward pass GPU reloading.
+        bwd_offload_node.meta.update(fwd_node.meta)
+        bwd_offload_node.meta["saved_for_offloading"] = True
+        bwd_offload_node.meta["original_device"] = bwd_node.meta["val"].device
+        bwd_node.replace_all_uses_with(bwd_offload_node)
+        bwd_module.graph.erase_node(bwd_node)
 
-    Pattern: fork → wait_stream → device_put → record_event → join → wait_event
+    reload_activation_bw_async(bwd_module.graph)
 
-    This ensures that:
-    1. Reloading doesn't start prematurely (fork → wait_stream)
-    2. Reloading happens on a separate stream (device_put)
-    3. First use waits for reload completion (record_event → join → wait_event)
 
-    NOTE: The pattern consists of two logical groups:
-          - First group (fork → wait_stream → device_put → record_event → join):
-            Performs asynchronous data transfer on a separate stream
-          - Second group (wait_event):
-            Data transfer completion check when the data is actually needed
+def activation_offload_sink_wait_async(fwd_module: fx.GraphModule) -> None:
+    """Sink ao.wait_tensor operations for offload completion to the end of the graph.
 
-          For prefetch optimization, the first group can be moved earlier in the graph
-          to overlap computation with data transfer, while the wait_event must remain
-          at its current position to prevent blocking computation unnecessarily.
+    This allows computation to overlap with offload operations.
 
-    Args:
-        graph: The backward graph to modify
+    NOTE: Sinking waits to the end delays GPU memory release of the source
+    tensor (kept alive via the wait's keepalive arg) until the end of the
+    compiled graph. For per-layer compile this is fine (one layer's worth of
+    memory), but for full-model compile this means offloaded GPU tensors are
+    not freed until the entire forward pass completes.
     """
+    graph: fx.Graph = fwd_module.graph
+    output_node: fx.Node = graph.find_nodes(op="output")[0]
 
-    # Find all GPU reload nodes
-    reload_nodes: list[fx.Node] = [
+    wait_nodes_to_sink: list[fx.Node] = [
         node
         for node in graph.nodes
-        if GPU_RELOAD_PREFIX in node.name and node.op == "call_function"
+        if node.op == "call_function"
+        and node.target == torch.ops.ao.wait_tensor.default
+        and isinstance(node.args[0], fx.Node)
+        and node.args[0].op == "call_function"
+        and node.args[0].target == torch.ops.ao.offload.default
     ]
-    if not reload_nodes:
-        return
-
-    # Get default stream id and offload stream id
-    current_stream_id: int = get_current_stream(
-        reload_nodes[0].args[0].meta["original_device"]  # type: ignore[assignment]
-    )
-    reload_stream_id: int = new_stream()
-
-    for reload_node in reload_nodes:
-        event_id: int = new_event()
-
-        with graph.inserting_before(reload_node):
-            # Fork to reload stream
-            graph.call_function(
-                torch.ops.streams.fork.default,
-                args=(current_stream_id, reload_stream_id),
-                name=f"stream_in_{reload_node.name}",
-            )
-            # Wait for default stream to prevent premature reloading
-            graph.call_function(
-                torch.ops.streams.wait_stream.default,
-                args=(reload_stream_id, current_stream_id),
-            )
-        with graph.inserting_after(reload_node):
-            # Record event on reload stream after device_put
-            record_event_node = graph.call_function(
-                torch.ops.streams.record_event.default,
-                args=(event_id, reload_stream_id),
-            )
-        with graph.inserting_after(record_event_node):
-            # Join back to default stream
-            join_node = graph.call_function(
-                torch.ops.streams.join.default,
-                args=(reload_stream_id, current_stream_id),
-                name=f"stream_out_{reload_node.name}",
-            )
-        with graph.inserting_after(join_node):
-            # Wait for the event on default stream
-            graph.call_function(
-                torch.ops.streams.wait_event.default,
-                args=(event_id, current_stream_id),
-            )
 
-
-def put_offload_nodes_on_separate_stream(
-    fwd_module: fx.GraphModule,
-    bwd_module: fx.GraphModule,
-) -> None:
-    """
-    Add stream and event related operations around offload nodes.
-
-    Args:
-        fwd_module: Forward module graph
-        bwd_module: Backward module graph
-    """
-
-    add_forward_offload_stream_ops(fwd_module.graph)
-    add_backward_reload_stream_ops(bwd_module.graph)
+    # prepend moves the node from its current position (no manual removal needed)
+    for wait_node in wait_nodes_to_sink:
+        output_node.prepend(wait_node)
 
 
-def _validate_pattern_nodes(
-    fork_node: fx.Node,
-    wait_stream_node: fx.Node,
-    record_event_node: fx.Node,
-    join_node: fx.Node,
-    wait_event_node: fx.Node,
-) -> None:
+def activation_reload_prefetch_async(bwd_module: fx.GraphModule) -> None:
     """
-    Validate that the pattern nodes match the expected structure.
-
-    Raises ValueError if any node doesn't match expectations.
+    Prefetch backward reload operations by moving ao.reload nodes earlier
+    in the graph to overlap data transfer with computation, while keeping
+    ao.wait_tensor at its original position.
     """
+    graph: fx.Graph = bwd_module.graph
+    nodes_list: list[fx.Node] = list(graph.nodes)
 
-    if not (
-        fork_node.op == "call_function"
-        and fork_node.target == torch.ops.streams.fork.default
-    ):
-        raise ValueError("Expected fork node two nodes before device_put node")
-
-    if not (
-        wait_stream_node.op == "call_function"
-        and wait_stream_node.target == torch.ops.streams.wait_stream.default
-    ):
-        raise ValueError("Expected wait_stream node one node before device_put node")
-
-    if not (
-        record_event_node.op == "call_function"
-        and record_event_node.target == torch.ops.streams.record_event.default
-    ):
-        raise ValueError("Expected record_event node one node after device_put node")
-
-    if not (
-        join_node.op == "call_function"
-        and join_node.target == torch.ops.streams.join.default
-    ):
-        raise ValueError("Expected join node two nodes after device_put node")
+    # Identify reload + wait pairs
+    reload_patterns: dict[fx.Node, ReloadNodeInfo] = {}
+    for node in graph.nodes:
+        if not (
+            node.op == "call_function" and node.target == torch.ops.ao.reload.default
+        ):
+            continue
+        wait_node = next(
+            (u for u in node.users if u.target == torch.ops.ao.wait_tensor.default),
+            None,
+        )
+        if wait_node is None:
+            continue
+        transfer_size_bytes: int = _calculate_transfer_size(node)
+        transfer_time_ms: float = _estimate_transfer_time_in_ms(transfer_size_bytes)
+        reload_patterns[node] = ReloadNodeInfo(
+            reload_group_nodes=[node],
+            wait_event_node=wait_node,
+            transfer_size_bytes=transfer_size_bytes,
+            transfer_time_ms=transfer_time_ms,
+        )
 
-    if not (
-        wait_event_node.op == "call_function"
-        and wait_event_node.target == torch.ops.streams.wait_event.default
-    ):
-        raise ValueError("Expected wait_event node three nodes after device_put node")
+    reorder_for_prefetch(nodes_list, reload_patterns)
 
 
 def _calculate_transfer_size(device_put_node: fx.Node) -> int:
     """Calculate the size in bytes of data being transferred."""
 
-    return _size_of(device_put_node.args[0])  # pyrefly: ignore [bad-argument-type]
+    # ao.offload(tensor) -> tensor at args[0]
+    # ao.reload(tensor, device) -> tensor at args[0]
+    if device_put_node.target in (
+        torch.ops.ao.offload.default,
+        torch.ops.ao.reload.default,
+    ):
+        return _size_of(device_put_node.args[0])  # pyrefly: ignore [bad-argument-type]
+    raise ValueError(f"Unexpected transfer op: {device_put_node.target}")
 
 
 def _estimate_transfer_time_in_ms(transfer_size_bytes: int) -> float:
-    """
-    Estimate transfer time in milliseconds based on size and bandwidth.
-    NOTE: potentially could be standardized in node estimator class
-    """
+    """Estimate transfer time in milliseconds based on size and bandwidth.
 
-    return transfer_size_bytes / (1024**3) * 1_000 / inductor_config.cpu_gpu_bw
+    Uses config.activation_offload_cpu_gpu_bw (GB/s) which should be set by
+    the user to match their hardware.
+    """
+    return (
+        transfer_size_bytes / (1024**3) * 1_000 / config.activation_offload_cpu_gpu_bw
+    )
 
 
 def identify_reload_patterns(
@@ -628,7 +614,8 @@ def identify_reload_patterns(
         wait_event_node: fx.Node = nodes_list[reload_node_idx + 3]
 
         # Validate the nodes are what we expect
-        _validate_pattern_nodes(
+        # Removed in follow-up commit
+        _validate_pattern_nodes(  # noqa: F821  # pyrefly: ignore [unknown-name]
             fork_node,
             wait_stream_node,
             record_event_node,
@@ -811,15 +798,16 @@ def enable_activation_offloading(
         return
 
     # Step 2: Add offload and reload nodes to the graphs
-    offload_chosen_sets(fwd_module, bwd_module)
-
-    # Step 3: Put offload nodes on separate stream if configured
     if config.activation_offload_separate_stream:
-        put_offload_nodes_on_separate_stream(fwd_module, bwd_module)
+        # Use async ao ops (2 nodes each: offload/reload + wait_tensor)
+        offload_chosen_sets_async(fwd_module, bwd_module)
         if config.activation_offload_sink_wait:
-            activation_offload_sink_wait(fwd_module)
+            activation_offload_sink_wait_async(fwd_module)
         if config.activation_reload_prefetch:
-            activation_reload_prefetch(bwd_module)
+            activation_reload_prefetch_async(bwd_module)
+    else:
+        # Use synchronous device_put (1 node each)
+        offload_chosen_sets(fwd_module, bwd_module)
 
     fwd_module.graph.lint()
     bwd_module.graph.lint()
diff --git a/torch/_functorch/_activation_offloading/offload_ops.py b/torch/_functorch/_activation_offloading/offload_ops.py
new file mode 100644
index 0000000000000..5259b97adca3d
--- /dev/null
+++ b/torch/_functorch/_activation_offloading/offload_ops.py
@@ -0,0 +1,180 @@
+"""Custom ops for async activation offloading between GPU and CPU.
+
+These ops encapsulate stream management internally, producing a clean 2-node
+IR pattern (offload/reload + wait_tensor) similar to c10d functional collectives.
+
+A single dedicated transfer stream handles all D2H/H2D copies.
+Completion events are keyed by output tensor data_ptr() and stored in a
+module-level registry, so ``ao.wait_tensor`` takes only the tensor itself
+(plus an optional keepalive).
+
+Offload pattern:
+    cpu_tensor = ao.offload(gpu_tensor)
+    cpu_tensor = ao.wait_tensor(cpu_tensor, gpu_tensor)
+                               (keepalive arg extends gpu_tensor lifetime past the async D2H copy)
+
+Reload pattern:
+    gpu_tensor = ao.reload(cpu_tensor, device)
+    gpu_tensor = ao.wait_tensor(gpu_tensor)
+"""
+
+import torch
+from torch._library.custom_ops import custom_op
+from torch.fx import has_side_effect
+
+
+# --- Global transfer stream (one per device, lazily created) ---
+_transfer_streams: dict[torch.device, torch.Stream] = {}
+
+
+def _get_or_create_transfer_stream(device: torch.device) -> torch.Stream:
+    if device not in _transfer_streams:
+        _transfer_streams[device] = torch.Stream(device=device)
+    return _transfer_streams[device]
+
+
+# --- Wait registry: maps data_ptr() -> (completion_event, device) ---
+# Created by ao.offload / ao.reload, consumed (popped) by ao.wait_tensor.
+# Not thread-safe — graph execution is single-threaded Python.
+_wait_registry: dict[int, tuple[torch.Event, torch.device]] = {}
+
+
+def _register_wait(tensor: torch.Tensor, device: torch.device) -> torch.Event:
+    """Create an event for an async transfer and register it for wait_tensor."""
+    event = torch.Event()
+    _wait_registry[tensor.data_ptr()] = (event, device)
+    return event
+
+
+def _pop_wait(tensor: torch.Tensor) -> tuple[torch.Event, torch.device]:
+    key = tensor.data_ptr()
+    try:
+        return _wait_registry.pop(key)
+    except KeyError:
+        raise RuntimeError(
+            f"ao.wait_tensor: no pending transfer for tensor with data_ptr={key}. "
+            "Every ao.wait_tensor must be paired with a preceding ao.offload or ao.reload."
+        ) from None
+
+
+def _clear_wait_registry() -> None:
+    _wait_registry.clear()
+
+
+@custom_op("ao::offload", mutates_args=())
+def offload(tensor: torch.Tensor) -> torch.Tensor:
+    """Async offload a GPU tensor to CPU on the dedicated transfer stream.
+
+    Callers MUST pair this with an ``ao.wait_tensor`` that passes the source GPU
+    tensor as ``keepalive`` to extend its lifetime past the async D2H copy.
+    Do NOT use ``record_stream`` — it causes memory fragmentation and
+    unbounded memory growth.
+
+    Uses pinned-memory allocation + copy_ so the transfer is compatible
+    with CUDA graph capture.
+    """
+    device = tensor.device
+    transfer_stream = _get_or_create_transfer_stream(device)
+    current_stream = torch.accelerator.current_stream(device)
+
+    transfer_stream.wait_stream(current_stream)
+
+    torch.accelerator.set_stream(transfer_stream)
+    result = torch.empty_like(tensor, device="cpu", pin_memory=True)
+    completion_event = _register_wait(result, device)
+    result.copy_(tensor, non_blocking=True)
+    transfer_stream.record_event(completion_event)
+    torch.accelerator.set_stream(current_stream)
+
+    return result
+
+
+@offload.register_fake
+def _(tensor: torch.Tensor) -> torch.Tensor:
+    return torch.empty_like(tensor, device="cpu")
+
+
+@custom_op("ao::reload", mutates_args=())
+def reload(
+    tensor: torch.Tensor,
+    device: torch.device,
+) -> torch.Tensor:
+    """Async reload a CPU tensor to GPU on the dedicated transfer stream.
+
+    The GPU tensor is allocated on the compute stream to avoid cross-stream
+    allocator ownership issues. The H2D copy runs on the transfer stream.
+    The completion event is keyed by the output tensor's data_ptr.
+    """
+    transfer_stream = _get_or_create_transfer_stream(device)
+    current_stream = torch.accelerator.current_stream(device)
+
+    # Allocate on compute stream so the allocator tracks ownership correctly
+    result = torch.empty_like(tensor, device=device)
+    completion_event = _register_wait(result, device)
+
+    transfer_stream.wait_stream(current_stream)
+
+    torch.accelerator.set_stream(transfer_stream)
+    result.copy_(tensor, non_blocking=True)
+    transfer_stream.record_event(completion_event)
+    torch.accelerator.set_stream(current_stream)
+
+    return result
+
+
+@reload.register_fake
+def _(
+    tensor: torch.Tensor,
+    device: torch.device,
+) -> torch.Tensor:
+    return torch.empty_like(tensor, device=device)
+
+
+# ao::wait_tensor is defined via torch.library with an aliasing schema so the
+# output can alias the input (custom_op forbids this).
+#
+# Uses CompositeExplicitAutograd (single impl for all devices) because the
+# offload case has mixed-device args: ``tensor`` is CPU (the offload result)
+# while ``keepalive`` is CUDA (the source GPU tensor). A single impl avoids
+# relying on device-priority dispatch ordering.
+#
+# Synchronization details (completion event, device) are looked up from
+# ``_wait_registry`` keyed on ``tensor.data_ptr()``.
+#
+# ``keepalive`` is not read by the op — its sole purpose is to create a graph
+# dependency that extends the tensor's lifetime in the FX graph. For offload,
+# this keeps the source GPU tensor alive until the compute stream has waited
+# on the D2H completion event, preventing the allocator from reclaiming it
+# while the async copy is still in flight.
+_lib = torch.library.Library("ao", "DEF")
+_lib.define("wait_tensor(Tensor(a) tensor, Tensor? keepalive=None) -> Tensor(a)")
+
+
+@torch.library.impl("ao::wait_tensor", "CompositeExplicitAutograd")
+def _ao_wait_tensor(
+    tensor: torch.Tensor,
+    keepalive: torch.Tensor | None = None,
+) -> torch.Tensor:
+    completion_event, device = _pop_wait(tensor)
+    current_stream = torch.accelerator.current_stream(device)
+    current_stream.wait_event(completion_event)
+    return tensor
+
+
+@torch.library.register_fake("ao::wait_tensor")
+def _ao_wait_tensor_fake(
+    tensor: torch.Tensor,
+    keepalive: torch.Tensor | None = None,
+) -> torch.Tensor:
+    return tensor
+
+
+has_side_effect(torch.ops.ao.wait_tensor.default)
+
+
+def wait_tensor(
+    tensor: torch.Tensor,
+    keepalive: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Callable wrapper so ``wait_tensor`` can be imported by name for op registration."""
+    return torch.ops.ao.wait_tensor.default(tensor, keepalive)
diff --git a/torch/_functorch/_aot_autograd/aot_autograd_result.py b/torch/_functorch/_aot_autograd/aot_autograd_result.py
index adc63662593f4..83ab2240b01c9 100644
--- a/torch/_functorch/_aot_autograd/aot_autograd_result.py
+++ b/torch/_functorch/_aot_autograd/aot_autograd_result.py
@@ -24,7 +24,7 @@
 from collections.abc import Callable, Sequence
 from copy import copy
 from dataclasses import dataclass
-from typing import Any, Generic, Optional, TYPE_CHECKING, TypeVar
+from typing import Any, Generic, TYPE_CHECKING, TypeVar
 
 import torch
 from torch._dynamo.precompile_context import BackendCacheArtifact
@@ -35,9 +35,11 @@
     OutputCode,
 )
 from torch._inductor.utils import should_use_remote_fx_graph_cache
+from torch._logging import getArtifactLogger
 
 from .runtime_wrappers import (
     AOTDispatchAutograd,
+    AOTDispatchAutogradCompileSpec,
     AOTDispatchSubclassWrapper,
     CachedAutogradLazyBackwardCompileInfo,
     CompilerWrapper,
@@ -57,6 +59,7 @@
     from .schemas import AOTConfig, ViewAndMutationMeta
 
 log = logging.getLogger(__name__)
+aot_graphs_log = getArtifactLogger(__name__, "aot_graphs")
 
 
 TOut = TypeVar("TOut", bound=OutputCode)
@@ -122,9 +125,22 @@ def post_compile(
                 payload_fn=lambda: json.dumps(cache_info),
             )
             result = graph  # type: ignore[assignment]
+            result.compile_region_name = (  # pyrefly: ignore[missing-attribute]
+                fx_config.get("compile_region_name")
+            )
 
         # Run normal post compile
         result.post_compile(self.example_inputs, constants, fx_config)
+
+        # Let the CUDAGraph policy do outer-level wrapping (e.g. wrapping
+        # an entire RegionalOutputCode as a single CUDA graph instead of
+        # per-inner-region).
+        import torch._inductor.config as _inductor_config
+
+        policy = _inductor_config.cudagraph_policy
+        if policy is not None:
+            result = policy.wrap_output(result)
+
         return result
 
 
@@ -137,7 +153,7 @@ def post_compile(
 @dataclass
 class FxGraphCacheLoadable(InductorOutput[CompiledFxGraph]):
     fx_graph_cache_info: tuple[str, list[str]]
-    fx_graph_guard_expr: Optional[str]
+    fx_graph_guard_expr: str | None
 
     def pre_save(self) -> None:
         return
@@ -216,7 +232,17 @@ def post_compile(
         """
         Called after FXGraphCacheLoadable.load, mutates fx_config
         """
+        result.compile_region_name = fx_config.get(  # pyrefly: ignore[bad-assignment]
+            "compile_region_name"
+        )
         result.post_compile(self.example_inputs, self.constants, fx_config)
+
+        import torch._inductor.config as _inductor_config
+
+        policy = _inductor_config.cudagraph_policy
+        if policy is not None:
+            result = policy.wrap_output(result)
+
         return result
 
 
@@ -236,27 +262,28 @@ class GenericCompiledBackward(InductorOutput[TOut]):
     backward_state_indices: list[int]
     num_symints_saved_for_bw_: int
 
+    def post_compile(self, result: TOut, fx_config: _CompileFxKwargs) -> TOut:
+        # The concrete post_compile comes from the loadable mixin in each subclass MRO.
+        compiled_bw = super().post_compile(  # pyrefly: ignore[missing-attribute]
+            result, fx_config
+        )
+        # See note [Wrapping bw_compiler in disable]
+        # This is done by _wrapped_bw_compiler in torch/_dynamo/backends/common.py
+        # But since on cache hit we do not call the bw_compiler, we need to reapply the disable
+        return torch._dynamo.disable(  # type: ignore[return-value]
+            compiled_bw, reason="do not trace generated backwards pass"
+        )
+
 
 @dataclass
 class CompiledBackward(GenericCompiledBackward[CompiledFxGraph], FxGraphCacheLoadable):
     """
-    Cacheable entry for a forward function
+    Cacheable entry for a backward function
     """
 
     def _is_backward(self) -> bool:
         return True
 
-    def post_compile(
-        self, result: CompiledFxGraph, fx_config: _CompileFxKwargs
-    ) -> CompiledFxGraph:
-        compiled_bw = super().post_compile(result, fx_config)
-        # See note [Wrapping bw_compiler in disable]
-        # This is done by _wrapped_bw_compiler in torch/_dynamo/backends/common.py
-        # But since on cache hit we do not call the bw_compiler, we need to reapply the disable
-        return torch._dynamo.disable(  # type: ignore[return-value]
-            compiled_bw, reason="do not trace generated backwards pass"
-        )
-
 
 # Generic bundled forward/backward classes that work with any OutputCode type
 @dataclass
@@ -280,17 +307,6 @@ class BundledCompiledBackward(
     Works with any OutputCode type (CompiledFxGraph, RegionalOutputCode, etc.)
     """
 
-    def post_compile(
-        self, result: TOutputCode, fx_config: _CompileFxKwargs
-    ) -> TOutputCode:
-        compiled_bw = super().post_compile(result, fx_config)
-        # See note [Wrapping bw_compiler in disable]
-        # This is done by _wrapped_bw_compiler in torch/_dynamo/backends/common.py
-        # But since on cache hit we do not call the bw_compiler, we need to reapply the disable
-        return torch._dynamo.disable(  # type: ignore[return-value]
-            compiled_bw, reason="do not trace generated backwards pass"
-        )
-
 
 @dataclass
 class SerializedGraphModule:
@@ -337,13 +353,13 @@ class GenericAOTAutogradResult(Generic[TForward, TBackward]):
 
     # Forward and Backward info
     compiled_fw: TForward
-    compiled_bw: Optional[TBackward]
+    compiled_bw: TBackward | None
 
     # Code of the joint graph using print_readable()
     # Used for logging purposes
-    aot_joint_graph_str: Optional[str]
-    aot_forward_graph_str: Optional[str]
-    aot_backward_graph_str: Optional[str]
+    aot_joint_graph_str: str | None
+    aot_forward_graph_str: str | None
+    aot_backward_graph_str: str | None
 
     # Runtime_metadata saved right before compilation
     runtime_metadata: ViewAndMutationMeta
@@ -352,8 +368,8 @@ class GenericAOTAutogradResult(Generic[TForward, TBackward]):
     dispatch_wrappers: list[CompilerWrapper]
 
     # Used by AOTSubclassWrapper
-    maybe_subclass_meta: Optional[SubclassMeta]
-    num_fw_outs_saved_for_bw: Optional[int]
+    maybe_subclass_meta: SubclassMeta | None
+    num_fw_outs_saved_for_bw: int | None
 
     # Used by RuntimeWrapper
     indices_of_inps_to_detach: list[int]
@@ -367,10 +383,10 @@ class GenericAOTAutogradResult(Generic[TForward, TBackward]):
     # Used by standalone_compile
     sanitized_aot_config: AOTConfig
 
-    guards_expr: Optional[str]
+    guards_expr: str | None
 
     # Used by Compiled Autograd
-    serialized_bw_module: Optional[SerializedGraphModule]
+    serialized_bw_module: SerializedGraphModule | None
 
     def pre_save(self) -> None:
         """
@@ -380,118 +396,115 @@ def pre_save(self) -> None:
         if self.compiled_bw is not None:
             self.compiled_bw.pre_save()
 
-    # Turn result into the original callable
-    def wrap_post_compile(
-        self,
-        args: list[torch.Tensor],
-        aot_config: AOTConfig,
-        fx_config: _CompileFxKwargs,
-        # pyrefly: ignore [implicit-any]
-    ) -> Callable:
-        """
-        This function takes a result and carefully reconstructs the original callable
-        that AOTAutograd returned the first time it was run. It does this by running the various
-        post compile steps that AOTAutograd runs on its compiled artifact after running the fw/bw compilers.
-
-        In the inference path, this consists of the Subclass, FunctionalzedRngRuntime, and RuntimeWrappers.
-        In the autograd path, this consists of AOTAutogradDispatch.post_compile.
-
-        The steps here should match exactly the steps that are run in aot_dispatch_base and aot_dispatch_autograd.
+    def _log_cached_graphs(self, aot_config: AOTConfig) -> None:
+        if not aot_config.enable_log:
+            return
 
-        Notably absent from the cached path are:
-        - DebugAssertWrapper
-        - FakifiedOutWrapper
-
-        Which we'll handle separately later on, if necessary.
-        """
-        from torch._dynamo.utils import CompileEventLogger, dynamo_timed
-
-        # Log the output of AOTAutogradCache
-        if aot_config.enable_log:
-            # TODO: maybe also log to aot_graphs_log
-            # Unfortunately aot_graphs_log uses
-            # slightly different formatting though
-            if self.aot_joint_graph_str is not None:
-                torch._logging.trace_structured(
-                    "aot_joint_graph", payload_fn=lambda: self.aot_joint_graph_str
-                )
+        if self.aot_joint_graph_str is not None:
+            torch._logging.trace_structured(
+                "aot_joint_graph", payload_fn=lambda: self.aot_joint_graph_str
+            )
+            aot_graphs_log.info(
+                "Joint graph (from cache)\n\n%s", self.aot_joint_graph_str
+            )
 
-            if self.aot_forward_graph_str is not None:
-                from torchgen.utils import dataclass_repr
+        if self.aot_forward_graph_str is not None:
+            from torchgen.utils import dataclass_repr
 
+            torch._logging.trace_structured(
+                "artifact",
+                metadata_fn=lambda: {
+                    "name": "aot_forward_graph_fw_metadata",
+                    "encoding": "string",
+                },
+                payload_fn=lambda: dataclass_repr(self.runtime_metadata),
+            )
+            if self.maybe_subclass_meta is not None:
                 torch._logging.trace_structured(
                     "artifact",
                     metadata_fn=lambda: {
-                        "name": "aot_forward_graph_fw_metadata",
+                        "name": "aot_forward_graph_fw_subclass_metadata",
                         "encoding": "string",
                     },
-                    payload_fn=lambda: dataclass_repr(self.runtime_metadata),
-                )
-                if self.maybe_subclass_meta is not None:
-                    torch._logging.trace_structured(
-                        "artifact",
-                        metadata_fn=lambda: {
-                            "name": "aot_forward_graph_fw_subclass_metadata",
-                            "encoding": "string",
-                        },
-                        payload_fn=lambda: dataclass_repr(self.maybe_subclass_meta),
-                    )
-
-                # It's called an inference graph if not running with autograd
-                name = (
-                    "aot_forward_graph"
-                    if self.aot_backward_graph_str is not None
-                    else "aot_inference_graph"
-                )
-                torch._logging.trace_structured(
-                    name, payload_fn=lambda: self.aot_forward_graph_str
+                    payload_fn=lambda: dataclass_repr(self.maybe_subclass_meta),
                 )
 
-            if self.aot_backward_graph_str is not None:
-                torch._logging.trace_structured(
-                    "aot_backward_graph", payload_fn=lambda: self.aot_backward_graph_str
-                )
-        with dynamo_timed("AOTAutogradCache.inductor_load"):
-            compiled_fw_func = self.compiled_fw.load(args)
-            compiled_bw_func = None
-            if self.compiled_bw is not None:
-                compiled_bw_func = self.compiled_bw.load(args)
-                needs_autograd = True
-                CompileEventLogger.try_add_pt2_compile(
-                    "backend_compile", dispatch_mode="autograd"
-                )
-                # Now that we've loaded forward and backward, call post compile on both
-                # This avoids setting things like BoxedBools in fx_config until
-                # after both forward and backward cache hit
-                fw_fx_config: _CompileFxKwargs = {
-                    **fx_config,
-                    "is_backward": False,
-                }
-                bw_fx_config: _CompileFxKwargs = {
-                    **fx_config,
-                    "is_backward": True,
-                }
-                compiled_fw_func = self.compiled_fw.post_compile(
-                    compiled_fw_func, fw_fx_config
-                )
-                compiled_bw_func = self.compiled_bw.post_compile(
-                    compiled_bw_func, bw_fx_config
-                )
-            else:
-                inference_fx_config: _CompileFxKwargs = {
-                    **fx_config,
-                    "is_backward": False,
-                }
-
-                needs_autograd = False
-                CompileEventLogger.try_add_pt2_compile(
-                    "backend_compile", dispatch_mode="inference"
-                )
-                compiled_fw_func = self.compiled_fw.post_compile(
-                    compiled_fw_func, inference_fx_config
-                )
+            # It's called an inference graph if not running with autograd
+            has_backward = self.aot_backward_graph_str is not None
+            torch._logging.trace_structured(
+                "aot_forward_graph" if has_backward else "aot_inference_graph",
+                payload_fn=lambda: self.aot_forward_graph_str,
+            )
+            aot_graphs_log.info(
+                "Forward graph (from cache)\n\n%s",
+                self.aot_forward_graph_str,
+            )
+
+        if self.aot_backward_graph_str is not None:
+            torch._logging.trace_structured(
+                "aot_backward_graph", payload_fn=lambda: self.aot_backward_graph_str
+            )
+            aot_graphs_log.info(
+                "Backward graph (from cache)\n\n%s",
+                self.aot_backward_graph_str,
+            )
+
+    def _load_and_post_compile(
+        self,
+        args: list[torch.Tensor],
+        fx_config: _CompileFxKwargs,
+    ) -> tuple[Callable[..., Any], Callable[..., Any] | None, bool]:
+        from torch._dynamo.utils import CompileEventLogger
+
+        compiled_fw_func = self.compiled_fw.load(args)
+        if self.compiled_bw is not None:
+            compiled_bw_func = self.compiled_bw.load(args)
+            needs_autograd = True
+            CompileEventLogger.try_add_pt2_compile(
+                "backend_compile", dispatch_mode="autograd"
+            )
+            # Now that we've loaded forward and backward, call post compile on both
+            # This avoids setting things like BoxedBools in fx_config until
+            # after both forward and backward cache hit
+            fw_fx_config: _CompileFxKwargs = {
+                **fx_config,
+                "is_backward": False,
+            }
+            bw_fx_config: _CompileFxKwargs = {
+                **fx_config,
+                "is_backward": True,
+            }
+            compiled_fw_func = self.compiled_fw.post_compile(
+                compiled_fw_func, fw_fx_config
+            )
+            compiled_bw_func = self.compiled_bw.post_compile(
+                compiled_bw_func, bw_fx_config
+            )
+            return compiled_fw_func, compiled_bw_func, needs_autograd
+
+        inference_fx_config: _CompileFxKwargs = {
+            **fx_config,
+            "is_backward": False,
+        }
+
+        needs_autograd = False
+        CompileEventLogger.try_add_pt2_compile(
+            "backend_compile", dispatch_mode="inference"
+        )
+        compiled_fw_func = self.compiled_fw.post_compile(
+            compiled_fw_func, inference_fx_config
+        )
+        return compiled_fw_func, None, needs_autograd
+
+    def _apply_runtime_wrappers(
+        self,
+        compiled_fw_func: Callable[..., Any],
+        compiled_bw_func: Callable[..., Any] | None,
+        needs_autograd: bool,
+        aot_config: AOTConfig,
+    ) -> Callable[..., Any]:
+        from torch._dynamo.utils import CompileEventLogger
 
-        # Wrap the forward function in post compile wrappers
         compiled_fw_func = AOTDispatchSubclassWrapper(
             trace_joint=needs_autograd,
             fw_only=None,
@@ -531,19 +544,20 @@ def wrap_post_compile(
             # 1. the bw is already compiled
             # 2. we don't need to save to the cache again
             # so those corresponding arguments are set to None.
-            compiled_function = AOTDispatchAutograd.post_compile(
-                compiled_fw_func,
-                compiled_bw_func,
-                self.maybe_subclass_meta,
-                self.compiled_bw.num_symints_saved_for_bw_,
-                self.compiled_bw.backward_state_indices,
-                disable_amp,
-                self.indices_of_inps_to_detach,
-                cached_lazy_backward,
-                aot_config,
+            compile_spec = AOTDispatchAutogradCompileSpec(
+                compiled_fw_func=compiled_fw_func,
+                compiled_bw_func=compiled_bw_func,
+                maybe_subclass_meta=self.maybe_subclass_meta,
+                num_symints_saved_for_bw=self.compiled_bw.num_symints_saved_for_bw_,
+                backward_state_indices=self.compiled_bw.backward_state_indices,
+                disable_amp=disable_amp,
+                indices_of_inps_to_detach=self.indices_of_inps_to_detach,
+                lazy_backward_info=cached_lazy_backward,
+                aot_config=aot_config,
                 fw_metadata=self.runtime_metadata,
                 try_save_cache_entry=None,
             )
+            compiled_function = AOTDispatchAutograd.post_compile(compile_spec)
 
         else:
             compiled_function = RuntimeWrapper(
@@ -561,9 +575,9 @@ def wrap_post_compile(
             aot_config,
             runtime_metadata=self.runtime_metadata,
         )
+        return compiled_function
 
-        # Now that we're pretty sure it's a successful load, add guards
-        # to the existing shape environment from the cache
+    def _check_guards(self, args: list[torch.Tensor]) -> None:
         if self.guards_expr:
             from .autograd_cache import AOTAutogradCache
 
@@ -572,6 +586,43 @@ def wrap_post_compile(
             if check is not True:
                 raise AssertionError(f"guards check failed: {check}")
 
+    # Turn result into the original callable
+    def wrap_post_compile(
+        self,
+        args: list[torch.Tensor],
+        aot_config: AOTConfig,
+        fx_config: _CompileFxKwargs,
+    ) -> Callable[..., Any]:
+        """
+        This function takes a result and carefully reconstructs the original callable
+        that AOTAutograd returned the first time it was run. It does this by running the various
+        post compile steps that AOTAutograd runs on its compiled artifact after running the fw/bw compilers.
+
+        In the inference path, this consists of the Subclass, FunctionalzedRngRuntime, and RuntimeWrappers.
+        In the autograd path, this consists of AOTAutogradDispatch.post_compile.
+
+        The steps here should match exactly the steps that are run in aot_dispatch_base and aot_dispatch_autograd.
+
+        Notably absent from the cached path are:
+        - DebugAssertWrapper
+        - FakifiedOutWrapper
+
+        Which we'll handle separately later on, if necessary.
+        """
+        from torch._dynamo.utils import dynamo_timed
+
+        self._log_cached_graphs(aot_config)
+        with dynamo_timed("AOTAutogradCache.inductor_load"):
+            compiled_fw_func, compiled_bw_func, needs_autograd = (
+                self._load_and_post_compile(args, fx_config)
+            )
+
+        compiled_function = self._apply_runtime_wrappers(
+            compiled_fw_func, compiled_bw_func, needs_autograd, aot_config
+        )
+        # Now that we're pretty sure it's a successful load, add guards
+        # to the existing shape environment from the cache.
+        self._check_guards(args)
         return compiled_function
 
 
diff --git a/torch/_functorch/_aot_autograd/autograd_cache.py b/torch/_functorch/_aot_autograd/autograd_cache.py
index 9d50d694b6b2a..08843076fe01e 100644
--- a/torch/_functorch/_aot_autograd/autograd_cache.py
+++ b/torch/_functorch/_aot_autograd/autograd_cache.py
@@ -7,6 +7,7 @@
 import base64
 import contextlib
 import functools
+import hashlib
 import json
 import logging
 import os
@@ -16,13 +17,18 @@
 import time
 import traceback
 from copy import copy
-from typing import Any, Optional, TYPE_CHECKING
+from typing import Any, TYPE_CHECKING
 from typing_extensions import override
 
 import torch
 from torch._dynamo.precompile_context import PrecompileContext
 from torch._dynamo.trace_rules import torch_non_c_binding_in_graph_functions
-from torch._dynamo.utils import chromium_event_log_active, CompileEventLogger, counters
+from torch._dynamo.utils import (
+    chromium_event_log_active,
+    CompileEventLogger,
+    counters,
+    warn_once,
+)
 from torch._functorch import config
 from torch._inductor.codecache import (
     _ident,
@@ -52,7 +58,7 @@
     CacheArtifactFactory,
     CacheArtifactManager,
 )
-from torch.fx.experimental.symbolic_shapes import size_hint
+from torch.fx.experimental.symbolic_shapes import guarding_hint_or_throw
 from torch.fx.node import Node
 from torch.utils._triton import has_triton_package
 
@@ -78,8 +84,7 @@
 if TYPE_CHECKING:
     from collections.abc import Callable, Generator, Sequence
 
-    from torch._inductor.compile_fx import _CompileFxKwargs
-    from torch._inductor.cudagraph_utils import BoxedDeviceIndex
+    from torch._inductor.compile_fx import _CompileFxKwargs, CompilerConfigExtra
     from torch._inductor.remote_cache import JsonDataTy, RemoteCache
 
 
@@ -154,6 +159,7 @@ def check_node_safe(node: Node) -> None:
         "torch.sym_float",
         "torch.sym_sum",
         "torch.autograd.grad",
+        "torch.distributed.tensor._api.from_local",
     )
     SAFE_NON_TORCH_FUNCTIONS = (
         "einops.einops.rearrange",
@@ -515,6 +521,107 @@ def __init__(self, gm: torch.fx.GraphModule) -> None:
             }
         )
 
+    # pyrefly: ignore [bad-override]
+    def reducer_override(self, obj: Any) -> Any:
+        """
+        Override to handle tensor subclasses (like DTensor) that aren't caught
+        by the dispatch_table's exact type matching.
+
+        The dispatch_table only matches exact types, so subclasses like DTensor
+        fall through to the default __reduce_ex__ which includes non-deterministic
+        storage addresses. This method catches those cases using isinstance checks.
+        """
+        from torch.utils._python_dispatch import is_traceable_wrapper_subclass
+
+        # Handle tensor subclasses that aren't exactly torch.Tensor
+        # dispatch_table already handles torch.Tensor exactly
+        if isinstance(obj, torch.Tensor) and type(obj) is not torch.Tensor:
+            if hasattr(obj, "_stable_hash_for_caching"):
+                return (_ident, (obj._stable_hash_for_caching(),))
+            if is_traceable_wrapper_subclass(obj):
+                warn_once(
+                    f"{type(obj).__name__} does not implement _stable_hash_for_caching. "
+                    "For PT2-compatible tensor subclasses, it is recommended to implement "
+                    "_stable_hash_for_caching(self) -> str for stable AOT autograd caching."
+                )
+                return (_ident, (self._default_stable_hash_for_caching(obj),))
+            return self._reduce_tensor(obj)
+        # Return NotImplemented to fall back to default behavior
+        return NotImplemented
+
+    # [NOTE] Tensor subclass stable hashing for AOT autograd cache
+    # Python's hash() varies with PYTHONHASHSEED, making cache keys unstable
+    # across processes. We use blake2b for cross-process determinism.
+    #
+    # EXTENSION POINT: Traceable wrapper subclasses can override cache key
+    # generation by implementing _stable_hash_for_caching(self) -> str.
+    # This method should return a deterministic string that uniquely identifies
+    # the tensor's metadata for caching purposes. See DTensor for an example.
+    #
+    # We can't define a default method on subclasses because there is no abstract
+    # base subclass, and we don't want to pollute torch.Tensor. Instead, we provide
+    # a default implementation here that uses __tensor_flatten__ to recursively
+    # hash inner tensors and metadata.
+
+    def _get_stable_hash(self, obj: Any) -> str:
+        """
+        Get stable hash for a tensor or opaque object, dispatching to custom or default implementation.
+        """
+        from torch._opaque_base import OpaqueBase
+        from torch.utils._python_dispatch import is_traceable_wrapper_subclass
+
+        if hasattr(obj, "_stable_hash_for_caching"):
+            return obj._stable_hash_for_caching()
+        elif isinstance(obj, torch.Tensor) and is_traceable_wrapper_subclass(obj):
+            return self._default_stable_hash_for_caching(obj)
+        elif isinstance(obj, OpaqueBase):
+            # Opaque objects are runtime pass-throughs; only the type matters
+            # for cache key purposes, not the instance identity or value.
+            type_name = type(obj).__qualname__
+            return hashlib.blake2b(type_name.encode(), digest_size=16).hexdigest()
+        elif isinstance(obj, torch.Tensor):
+            metadata = extract_tensor_metadata_for_cache_key(obj)
+            return hashlib.blake2b(pickle.dumps(metadata), digest_size=16).hexdigest()
+        else:
+            return hashlib.blake2b(pickle.dumps(obj), digest_size=16).hexdigest()
+
+    def _default_stable_hash_for_caching(self, tensor: torch.Tensor) -> str:
+        """
+        Default stable hash implementation for traceable wrapper subclasses.
+        """
+        from torch._opaque_base import OpaqueBase
+
+        inner_tensor_names, subclass_metadata = tensor.__tensor_flatten__()  # type: ignore[attr-defined]
+
+        # Recursively get hashes of inner tensors/opaque objects
+        inner_hashes: dict[str, str] = {}
+        for name in inner_tensor_names:
+            inner = getattr(tensor, name)
+            inner_hashes[name] = self._get_stable_hash(inner)
+
+        # Stabilize metadata: replace OpaqueBase instances with their type name
+        # since their repr includes memory addresses
+        def _stabilize(obj: Any) -> Any:
+            if isinstance(obj, OpaqueBase):
+                return type(obj).__qualname__
+            if isinstance(obj, tuple):
+                return tuple(_stabilize(x) for x in obj)
+            if isinstance(obj, list):
+                return [_stabilize(x) for x in obj]
+            if isinstance(obj, dict):
+                return {k: _stabilize(v) for k, v in obj.items()}
+            return obj
+
+        cache_data = pickle.dumps(
+            (
+                tensor.shape,
+                tensor.requires_grad,
+                _stabilize(subclass_metadata),
+                inner_hashes,
+            )
+        )
+        return hashlib.blake2b(cache_data, digest_size=16).hexdigest()
+
     def _reduce_aot_config(
         self, aot_config: AOTConfig
     ) -> tuple[Callable[..., Any], tuple[Any, ...]]:
@@ -599,55 +706,83 @@ def normalize_placeholder_names(
         gm.recompile()
 
 
+def create_fx_config(
+    compiler_config_extra: CompilerConfigExtra | None = None,
+    compile_region_name: str | None = None,
+) -> _CompileFxKwargs:
+    if compiler_config_extra is None:
+        cudagraphs = BoxedBool(torch._inductor.config.triton.cudagraphs)
+        boxed_forward_device_index = None
+    else:
+        cudagraphs = compiler_config_extra.cudagraphs
+        boxed_forward_device_index = compiler_config_extra.forward_device
+    return {
+        "cudagraphs": cudagraphs,
+        "boxed_forward_device_index": boxed_forward_device_index,
+        "compile_region_name": compile_region_name,  # pyrefly: ignore[bad-typed-dict-key]
+    }
+
+
 def autograd_cache_key(
-    gm: torch.fx.GraphModule,
+    mod: torch.fx.GraphModule | torch._dynamo.utils.GmWrapper,
     example_inputs: Sequence[Any],
     config: AOTConfig,
-    fx_config: _CompileFxKwargs,
+    compiler_config_extra: CompilerConfigExtra | None = None,
     # TODO: add args and parameters
 ) -> tuple[str, list[str]]:
     """
     Generate a unique hash of the FX graph for caching.
     """
 
-    try:
-        check_cacheable(gm)
-        if has_triton_package():
-            # Due to https://github.com/triton-lang/triton/issues/3729,
-            # if triton is < 3.2.0, AOTAutogradCache may cause us to
-            # attempt to load a cache entry without initializing
-            # the CUDA context on the autograd thread.
-
-            # Without caching, we naturally do this initialization when
-            # tracing through the graph with the autograd engine.
-            import triton
-
-            if triton.__version__ < "3.2.0":
-                raise BypassAOTAutogradCache("AOTAutogradCache requires triton 3.2.0")
-        details = AOTAutogradCacheDetails(gm, example_inputs, config, fx_config)
-        pickler = AOTAutogradCachePickler(gm)
-        # The prefix distinguishes among the other kinds of objects we cache
-        key = "a" + pickler.get_hash(details)
-        debug_lines = pickler.debug_lines(details)
-        log.debug(
-            "Autograd graph cache hash details for key %s:\n%s",
-            key,
-            LazyString(lambda: "\n".join(debug_lines)),
-        )
-        return key, debug_lines
-    except Exception:
-        # If enable_aot_compile is set, we're in AOT precompile mode where we always
-        # want to use fallback nonce keys. Unlike caching, it's fine if we can't generate
-        # a proper key because we are guaranteed in an AOT precompile world users are in
-        # complete control of distributing and loading artifacts.
-        if torch._functorch.config.bypass_autograd_cache_key:
-            log.info(
-                "Failed to generate AOTAutograd cache key; falling back to nonce due to enable_aot_compile",
-                exc_info=True,
+    gm = mod.gm if isinstance(mod, torch._dynamo.utils.GmWrapper) else mod
+    with sanitize_gm_for_cache(gm):
+        try:
+            check_cacheable(gm)
+            if has_triton_package():
+                # Due to https://github.com/triton-lang/triton/issues/3729,
+                # if triton is < 3.2.0, AOTAutogradCache may cause us to
+                # attempt to load a cache entry without initializing
+                # the CUDA context on the autograd thread.
+
+                # Without caching, we naturally do this initialization when
+                # tracing through the graph with the autograd engine.
+                import triton
+
+                if triton.__version__ < "3.2.0":
+                    raise BypassAOTAutogradCache(
+                        "AOTAutogradCache requires triton 3.2.0"
+                    )
+            details = AOTAutogradCacheDetails(
+                gm, example_inputs, config, create_fx_config(compiler_config_extra)
             )
-            return str(random.random()), []
-        else:
-            raise
+            pickler = AOTAutogradCachePickler(gm)
+            # The prefix distinguishes among the other kinds of objects we cache
+            key = "a" + pickler.get_hash(details)
+            # debug_lines re-hashes every attribute individually and is
+            # expensive. Only compute when debug logging is enabled.
+            if log.isEnabledFor(logging.DEBUG):
+                debug_lines = pickler.debug_lines(details)
+                log.debug(
+                    "Autograd graph cache hash details for key %s:\n%s",
+                    key,
+                    LazyString(lambda: "\n".join(debug_lines)),
+                )
+            else:
+                debug_lines: list[str] = []
+            return key, debug_lines
+        except Exception:
+            # If enable_aot_compile is set, we're in AOT precompile mode where we always
+            # want to use fallback nonce keys. Unlike caching, it's fine if we can't generate
+            # a proper key because we are guaranteed in an AOT precompile world users are in
+            # complete control of distributing and loading artifacts.
+            if torch._functorch.config.bypass_autograd_cache_key:
+                log.info(
+                    "Failed to generate AOTAutograd cache key; falling back to nonce due to enable_aot_compile",
+                    exc_info=True,
+                )
+                return str(random.random()), []
+            else:
+                raise
 
 
 @contextlib.contextmanager
@@ -666,6 +801,7 @@ def sanitize_gm_for_cache(
     """
     # Mapping from each field to a default value
     IGNORED_FIELDS: dict[str, Any] = {
+        # pyrefly: ignore [implicit-any]
         "meta": {},  # metadata used by export
         "compile_subgraph_reason": None,  # Used by dynamo only for logging, no change in inductor/autograd behavior
         "_param_name_to_source": None,  # Encapsulated by aot_config.aot_autograd_arg_pos_to_source
@@ -745,165 +881,152 @@ def try_load(
         mod: torch.fx.GraphModule | torch._dynamo.utils.GmWrapper,
         args: list[Any],
         aot_config: AOTConfig,
-        cudagraphs: BoxedBool,
-        boxed_forward_device_index: BoxedDeviceIndex | None,
+        compiler_config_extra: CompilerConfigExtra | None,
         local: bool,
         remote: bool,
+        compile_region_name: str | None = None,
     ) -> Callable[..., Any] | None:
         """
         Load a result from the cache, and reconstruct a runtime wrapper around the object
         """
-        gm = mod.gm if isinstance(mod, torch._dynamo.utils.GmWrapper) else mod
-        with sanitize_gm_for_cache(gm):
-            compiled_fn = None
-            cache_info: dict[str, Any] = {}
-            cache_key = None
-            debug_lines: list[str] = []
-            cache_event_time = time.time_ns()
-            cache_state = None
-            fx_config: _CompileFxKwargs = {
-                "cudagraphs": cudagraphs,
-                "boxed_forward_device_index": boxed_forward_device_index,
-            }
-            try:
-                cache_key, debug_lines = autograd_cache_key(
-                    gm, args, aot_config, fx_config
+        compiled_fn = None
+        cache_info: dict[str, Any] = {}
+        cache_key = None
+        debug_lines: list[str] = []
+        cache_event_time = time.time_ns()
+        cache_state = None
+        try:
+            cache_key, debug_lines = autograd_cache_key(
+                mod, args, aot_config, compiler_config_extra
+            )
+            result: tuple[GenericAOTAutogradResult[Any, Any], bytes] | None = (
+                AOTAutogradCache._lookup(
+                    cache_key, local, remote, args, cache_info, aot_config
                 )
-                result: Optional[tuple[GenericAOTAutogradResult[Any, Any], bytes]] = (
-                    AOTAutogradCache._lookup(
-                        cache_key, local, remote, args, cache_info, aot_config
-                    )
+            )
+            if result is not None:
+                (entry, pickled_content) = result
+                fx_config = create_fx_config(compiler_config_extra, compile_region_name)
+                compiled_fn = entry.wrap_post_compile(args, aot_config, fx_config)
+                # Make the compiled_fn serializable, where the serialize function just
+                # makes a copy of the original entry before post compile via the pickled content
+                compiled_fn = SerializableCompiledFunction(
+                    compiled_fn, lambda: pickle.loads(pickled_content)
                 )
-                if result is not None:
-                    (entry, pickled_content) = result
-                    compiled_fn = entry.wrap_post_compile(args, aot_config, fx_config)
-                    # Make the compiled_fn serializable, where the serialize function just
-                    # makes a copy of the original entry before post compile via the pickled content
-                    compiled_fn = SerializableCompiledFunction(
-                        compiled_fn, lambda: pickle.loads(pickled_content)
-                    )
-                    log.info("AOTAutograd cache hit for key %s", cache_key)
-
-                    counters["aot_autograd"]["autograd_cache_hit"] += 1
-                    cache_state = "hit"
-                    cache_event_time = time.time_ns()
-                    forward_time_saved = entry.forward_time_taken_ns // 1e6
-                    backward_time_saved = entry.backward_time_taken_ns // 1e6
-                    cache_info.update(
-                        {
-                            "forward_time_saved_ms": forward_time_saved,
-                            "backward_time_saved_ms": backward_time_saved,
-                            "time_saved_ms": forward_time_saved + backward_time_saved,
-                        }
-                    )
-                    time_saved_ns = (
-                        entry.forward_time_taken_ns + entry.backward_time_taken_ns
-                    )
-                    # TODO: should we use the same field for remote cache time saved for both
-                    # FXGraphCache and AOTAutogradCache?
-                    # get_metrics_context().increment(...)
-                    if (
-                        ephemeral_increase
-                        := add_ephemeral_timeout_increase_for_distributed(time_saved_ns)
-                    ) != 0:
-                        cache_info["ephemeral_timeout_increase"] = ephemeral_increase
-
-                if compiled_fn is None:
-                    log.info("AOTAutograd cache miss for key %s", cache_key)
-                    counters["aot_autograd"]["autograd_cache_miss"] += 1
-                    cache_state = "miss"
-                    cache_event_time = time.time_ns()
-            # Count missing the FXGraphCache as a miss not a bypass
-            except FXGraphCacheMiss as e:
-                counters["aot_autograd"]["autograd_cache_miss"] += 1
-                cache_state = "miss"
-                if (
-                    config.strict_autograd_cache
-                    or torch._dynamo.config.strict_precompile
-                ):
-                    raise e
-            # Most often this is BypassAOTAutogradCache, but
-            # if there's ever different reason we can't cache,
-            # we still never want to hard throw an exception, since
-            # we can always fallback to a cache bypass.
-            # As an example, if the user calls autograd via
-            # standalone inductor, we will sometimes get a GraphModule
-            # that doesn't actually have a `.graph` on it. Instead
-            # of checking every single case, we safely catch the exception
-            # in those cases.
-            except Exception as e:
-                cache_key = None
-                counters["aot_autograd"]["autograd_cache_bypass"] += 1
-                log.info("Bypassing autograd cache due to: %s", e)  # noqa: G200
-                cache_state = "bypass"
+                log.info("AOTAutograd cache hit for key %s", cache_key)
+
+                counters["aot_autograd"]["autograd_cache_hit"] += 1
+                cache_state = "hit"
                 cache_event_time = time.time_ns()
-                cache_info["cache_bypass_reason"] = str(e)
-                cache_info["cache_bypass_exception_type"] = type(e).__name__
-                cache_info["cache_bypass_traceback"] = traceback.format_exc().split(
-                    "\n"
+                forward_time_saved = entry.forward_time_taken_ns // 1e6
+                backward_time_saved = entry.backward_time_taken_ns // 1e6
+                cache_info.update(
+                    {
+                        "forward_time_saved_ms": forward_time_saved,
+                        "backward_time_saved_ms": backward_time_saved,
+                        "time_saved_ms": forward_time_saved + backward_time_saved,
+                    }
                 )
-                # TODO: this gets logged implicitly by cache_bypass_reason,
-                # and here we explicitly log it into tlparse.
-                # We may want to log this as an extra column in Scuba, though.
-                cache_info["cache_bypass_hard_exception"] = not isinstance(
-                    e, BypassAOTAutogradCache
+                time_saved_ns = (
+                    entry.forward_time_taken_ns + entry.backward_time_taken_ns
                 )
-                if remote:
-                    log_cache_bypass("bypass_aot_autograd", str(e))
+                # TODO: should we use the same field for remote cache time saved for both
+                # FXGraphCache and AOTAutogradCache?
+                # get_metrics_context().increment(...)
                 if (
-                    config.strict_autograd_cache
-                    or torch._dynamo.config.strict_precompile
-                ):
-                    raise e
-            if compiled_fn is None:
-                # Set the cache key so we can save a cache result later
-                symints = AOTAutogradCache._filter_backed_symints(args)
-                if cache_key is not None:
-                    aot_config.cache_info = AOTAutogradCacheInfo(
-                        cache_key,
-                        time.time_ns(),
-                        forward_symints=symints,
-                    )
+                    ephemeral_increase
+                    := add_ephemeral_timeout_increase_for_distributed(time_saved_ns)
+                ) != 0:
+                    cache_info["ephemeral_timeout_increase"] = ephemeral_increase
 
-            cache_info.update(
-                {
-                    "key": cache_key,
-                    "cache_state": cache_state,
-                    "components": debug_lines,
-                }
+            if compiled_fn is None:
+                log.info("AOTAutograd cache miss for key %s", cache_key)
+                counters["aot_autograd"]["autograd_cache_miss"] += 1
+                cache_state = "miss"
+                cache_event_time = time.time_ns()
+        # Count missing the FXGraphCache as a miss not a bypass
+        except FXGraphCacheMiss as e:
+            counters["aot_autograd"]["autograd_cache_miss"] += 1
+            cache_state = "miss"
+            if config.strict_autograd_cache or torch._dynamo.config.strict_precompile:
+                raise e
+        # Most often this is BypassAOTAutogradCache, but
+        # if there's ever different reason we can't cache,
+        # we still never want to hard throw an exception, since
+        # we can always fallback to a cache bypass.
+        # As an example, if the user calls autograd via
+        # standalone inductor, we will sometimes get a GraphModule
+        # that doesn't actually have a `.graph` on it. Instead
+        # of checking every single case, we safely catch the exception
+        # in those cases.
+        except Exception as e:
+            cache_key = None
+            counters["aot_autograd"]["autograd_cache_bypass"] += 1
+            log.info("Bypassing autograd cache due to: %s", e)  # noqa: G200
+            cache_state = "bypass"
+            cache_event_time = time.time_ns()
+            cache_info["cache_bypass_reason"] = str(e)
+            cache_info["cache_bypass_exception_type"] = type(e).__name__
+            cache_info["cache_bypass_traceback"] = traceback.format_exc().split("\n")
+            # TODO: this gets logged implicitly by cache_bypass_reason,
+            # and here we explicitly log it into tlparse.
+            # We may want to log this as an extra column in Scuba, though.
+            cache_info["cache_bypass_hard_exception"] = not isinstance(
+                e, BypassAOTAutogradCache
             )
-            if chromium_event_log_active():
-                CompileEventLogger.instant(
-                    f"autograd_cache_{cache_state}",
-                    metadata=cache_info,
-                    time_ns=cache_event_time,
-                )
-                CompileEventLogger.try_add_pt2_compile(
-                    "backend_compile",
-                    cache_state=cache_state,
-                    cache_event_time=cache_event_time,
-                    key=cache_info.get("key"),
-                    components=cache_info.get("components"),
-                    cache_bypass_reason=cache_info.get("cache_bypass_reason"),
-                    remote_cache_enabled=remote,
-                    local_cache_enabled=local,
+            if remote:
+                log_cache_bypass("bypass_aot_autograd", str(e))
+            if config.strict_autograd_cache or torch._dynamo.config.strict_precompile:
+                raise e
+        if compiled_fn is None:
+            # Set the cache key so we can save a cache result later
+            symints = AOTAutogradCache._filter_backed_symints(args)
+            if cache_key is not None:
+                aot_config.cache_info = AOTAutogradCacheInfo(
+                    cache_key,
+                    time.time_ns(),
+                    forward_symints=symints,
                 )
 
-            torch._logging.trace_structured(
-                "artifact",
-                metadata_fn=lambda: {
-                    "name": f"aotautograd_cache_{cache_state}",
-                    "encoding": "json",
-                },
-                payload_fn=lambda: json.dumps(cache_info),
+        cache_info.update(
+            {
+                "key": cache_key,
+                "cache_state": cache_state,
+                "components": debug_lines,
+            }
+        )
+        if chromium_event_log_active():
+            CompileEventLogger.instant(
+                f"autograd_cache_{cache_state}",
+                metadata=cache_info,
+                time_ns=cache_event_time,
+            )
+            CompileEventLogger.try_add_pt2_compile(
+                "backend_compile",
+                cache_state=cache_state,
+                cache_event_time=cache_event_time,
+                key=cache_info.get("key"),
+                components=cache_info.get("components"),
+                cache_bypass_reason=cache_info.get("cache_bypass_reason"),
+                remote_cache_enabled=remote,
+                local_cache_enabled=local,
             )
 
-            return compiled_fn
+        torch._logging.trace_structured(
+            "artifact",
+            metadata_fn=lambda: {
+                "name": f"aotautograd_cache_{cache_state}",
+                "encoding": "json",
+            },
+            payload_fn=lambda: json.dumps(cache_info),
+        )
+
+        return compiled_fn
 
     @classmethod
     def generate_guards_expression(
         cls: type[AOTAutogradCache], cache_info: AOTAutogradCacheInfo
-    ) -> Optional[str]:
+    ) -> str | None:
         shape_env = cls._get_shape_env()
 
         if shape_env is None:
@@ -980,12 +1103,12 @@ def _lookup(
         aot_config: AOTConfig | None,
     ) -> tuple[GenericAOTAutogradResult[Any, Any], bytes] | None:
         """Given a key generated by AOTAutogradCachePickler, look up its location in the cache."""
-        remote_cache: Optional[RemoteCache[JsonDataTy]] = None
+        remote_cache: RemoteCache[JsonDataTy] | None = None
         if remote:
             remote_cache = AOTAutogradCache.get_remote_cache()
 
         symints = AOTAutogradCache._filter_backed_symints(args)
-        hints = [size_hint(s) for s in symints]
+        hints = [guarding_hint_or_throw(s) for s in symints]
         entry = None
         pickled_content = None
         try:
@@ -1047,7 +1170,7 @@ def _write_to_local_cache(key: str, content: bytes) -> None:
     @staticmethod
     def _find_unpicklable_field(
         entry: GenericAOTAutogradResult[Any, Any],
-    ) -> Optional[str]:
+    ) -> str | None:
         """Find which field of entry is causing pickle to fail."""
         # pyrefly: ignore [implicit-any]
         fields = []
@@ -1066,10 +1189,11 @@ def _find_unpicklable_field(
     @staticmethod
     def _pickle_entry(
         entry: GenericAOTAutogradResult[Any, Any], remote: bool
-    ) -> Optional[bytes]:
+    ) -> bytes | None:
         """Pickle entry, returning None on failure."""
         try:
-            return pickle.dumps(entry)
+            with torch.utils._python_dispatch._disable_current_modes():
+                return pickle.dumps(entry)
         except (pickle.PicklingError, TypeError, AttributeError) as e:
             bad_field = AOTAutogradCache._find_unpicklable_field(entry)
             error_str = str(e)
@@ -1108,7 +1232,7 @@ def _handle_save_error(e: Exception, remote: bool, is_bypass: bool) -> None:
     @staticmethod
     def save(key: str, entry: GenericAOTAutogradResult[Any, Any], remote: bool) -> None:
         """Save a single entry into the cache."""
-        content: Optional[bytes] = None
+        content: bytes | None = None
         try:
             entry.pre_save()
             content = AOTAutogradCache._pickle_entry(entry, remote)
@@ -1148,7 +1272,7 @@ def save(key: str, entry: GenericAOTAutogradResult[Any, Any], remote: bool) -> N
 
     @staticmethod
     @functools.cache
-    def get_remote_cache() -> Optional[RemoteCache[JsonDataTy]]:
+    def get_remote_cache() -> RemoteCache[JsonDataTy] | None:
         """
         Attempts to load the remote cache, returns None on error.
         """
@@ -1163,22 +1287,22 @@ def get_remote_cache() -> Optional[RemoteCache[JsonDataTy]]:
     @staticmethod
     def make_entry(
         compiled_fw_func: OutputCode,
-        compiled_bw_func: Optional[OutputCode],
-        aot_joint_graph_str: Optional[str],
-        aot_forward_graph_str: Optional[str],
-        aot_backward_graph_str: Optional[str],
+        compiled_bw_func: OutputCode | None,
+        aot_joint_graph_str: str | None,
+        aot_forward_graph_str: str | None,
+        aot_backward_graph_str: str | None,
         runtime_metadata: ViewAndMutationMeta,
         dispatch_wrappers: list[CompilerWrapper],
-        maybe_subclass_meta: Optional[SubclassMeta],
-        num_fw_outs_saved_for_bw: Optional[int],
+        maybe_subclass_meta: SubclassMeta | None,
+        num_fw_outs_saved_for_bw: int | None,
         indices_of_inps_to_detach: list[int],
         forward_time_taken_ns: int,
         backward_time_taken_ns: int,
         sanitized_aot_config: AOTConfig,
-        guards_expr: Optional[str],
-        backward_state_indices: Optional[list[int]],
-        num_symints_saved_for_bw: Optional[int],
-        serialized_bw_module: Optional[SerializedGraphModule],
+        guards_expr: str | None,
+        backward_state_indices: list[int] | None,
+        num_symints_saved_for_bw: int | None,
+        serialized_bw_module: SerializedGraphModule | None,
     ) -> GenericAOTAutogradResult[Any, Any]:
         if should_bundle_autograd_cache():
             # Helper function to unwrap all the wrappers we added during aotdispatch
diff --git a/torch/_functorch/_aot_autograd/collect_metadata_analysis.py b/torch/_functorch/_aot_autograd/collect_metadata_analysis.py
index fe7cdb83478c1..d4d56919424eb 100644
--- a/torch/_functorch/_aot_autograd/collect_metadata_analysis.py
+++ b/torch/_functorch/_aot_autograd/collect_metadata_analysis.py
@@ -13,7 +13,7 @@
 import collections
 import contextlib
 import logging
-from typing import Any, Optional, TYPE_CHECKING
+from typing import Any, TYPE_CHECKING
 
 import torch
 import torch.utils._pytree as pytree
@@ -168,10 +168,8 @@ def run_functionalized_fw_and_collect_metadata(
     *,
     flat_args_descs: list[AOTInput],
     keep_input_mutations: bool,
-    # TODO: refactor to kill this flag
-    is_train: bool = False,
     # Note: this is guaranteed to be set when running under dynamo
-    static_input_indices: Optional[list[int]] = None,
+    static_input_indices: list[int] | None = None,
     pre_dispatch: bool = False,
 ) -> Callable[..., ViewAndMutationMeta]:
     memo: dict[Tensor, Tensor] = {}
@@ -325,7 +323,7 @@ def inner(*flat_args: Any) -> ViewAndMutationMeta:
         ] = collections.defaultdict(int)
 
         out_storage_to_metadata_key_to_tensors: collections.defaultdict[
-            Optional[StorageWeakRef],
+            StorageWeakRef | None,
             collections.defaultdict[MetadataKey, set[torch.Tensor]],
         ] = collections.defaultdict(lambda: collections.defaultdict(set))
 
@@ -690,12 +688,18 @@ def inner(*flat_args: Any) -> ViewAndMutationMeta:
                 if isinstance(o, FunctionalTensor):
                     view_meta_sequence = ViewMetaSequence(o)
 
+            requires_grad = isinstance(o, torch.Tensor) and o.requires_grad
             out_info = OutputAliasInfo(
                 output_type=output_type,
                 raw_type=type(o),
                 base_idx=base_idx,
                 dynamic_dims=dynamic_dims,
-                requires_grad=isinstance(o, torch.Tensor) and o.requires_grad,
+                requires_grad=requires_grad,
+                # A view created under no_grad() inherits requires_grad from
+                # its base but has no grad_fn and does not participate in
+                # differentiation.
+                requires_grad_for_backward=requires_grad
+                and (o._base is None or grad_fn is not None),
                 view_meta_sequence=view_meta_sequence,
             )
             output_info.append(out_info)
@@ -784,7 +788,7 @@ def _is_subclass_mutated_input_tangent_always_subclass(inp: object) -> bool:
                 OutputType.custom_function_view,
             ]
             and issubclass(info.raw_type, torch.Tensor)
-            and info.requires_grad
+            and info.requires_grad_for_backward
         ]
         f_output_tangents, f_output_tangents_descs = (
             [x[0] for x in f_output_tangents_pairs],
@@ -829,25 +833,15 @@ def _is_subclass_mutated_input_tangent_always_subclass(inp: object) -> bool:
             for inp, info in zip(flat_f_args, input_info)
             if info.mutation_type == MutationType.MUTATED_OUT_GRAPH
         ]
-        f_metadata_mutated_inputs = [
-            inp for inp, info in zip(flat_f_args, input_info) if info.mutates_metadata
-        ]
-        # This logic (annoyingly) re-figures out exactly what the outputs to the compiled fw graph will be.
-        # When handling subclasses, we need info about **all** outputs of compiled forward graph,
-        # so we know precisely which graph outputs to wrap back into tensor subclasses
-        # Ideally we would refactor this so not have an is_train flag, and have the separate
-        # inference and training paths decide which inputs/output to ask for subclass info on.
-        # However, we currently stash indexing information on each SubclassMeta about its order
-        # in the graph outputs list.
-        f_fw_graph_outs = list(flat_f_outs)
-        if is_train or not keep_input_mutations:
-            f_fw_graph_outs = f_mutated_inputs + f_fw_graph_outs
-        else:
-            # even when "keep_input_mutations" is True,
-            # we never keep metadata-only mutations in the fw graph
-            f_fw_graph_outs = f_metadata_mutated_inputs + f_fw_graph_outs
-        if is_train:
-            f_fw_graph_outs = f_fw_graph_outs + intermediate_bases
+        # Build the full list of forward graph outputs so the subclass wrapping
+        # code knows exactly which graph outputs to wrap back into subclasses.
+        # Including intermediate_bases unconditionally is safe: they are only
+        # populated when outputs require grad (line ~539), so they are naturally
+        # empty during pure inference.  In the "downgrade from training to
+        # inference" path, num_intermediate_bases > 0 is already gated behind
+        # `assert not req_subclass_dispatch` (aot_autograd.py), so the subclass
+        # wrapping code that consumes subclass_fw_graph_out_meta never sees them.
+        f_fw_graph_outs = [*f_mutated_inputs, *flat_f_outs, *intermediate_bases]
         fw_graph_outs = pytree.tree_map(from_fun, f_fw_graph_outs)
 
         grad_enabled_mutation = None
@@ -864,6 +858,12 @@ def _is_subclass_mutated_input_tangent_always_subclass(inp: object) -> bool:
                 grad_enabled_mutation,
             )
 
+        subclass_inp_meta = create_subclass_meta(flat_args)
+        subclass_fw_graph_out_meta = create_subclass_meta(fw_graph_outs)
+        subclass_tangent_meta = create_subclass_meta(
+            traced_tangents, count_symints=False, with_memory_format=True
+        )
+
         metadata = ViewAndMutationMeta(
             input_info=input_info,
             output_info=output_info,
@@ -871,12 +871,9 @@ def _is_subclass_mutated_input_tangent_always_subclass(inp: object) -> bool:
             keep_input_mutations=keep_input_mutations,
             traced_tangents=traced_tangents,
             traced_tangents_descs=traced_tangents_descs,
-            subclass_inp_meta=create_subclass_meta(flat_args),
-            subclass_fw_graph_out_meta=create_subclass_meta(fw_graph_outs),
-            subclass_tangent_meta=create_subclass_meta(
-                traced_tangents, count_symints=False, with_memory_format=True
-            ),
-            is_train=is_train,
+            subclass_inp_meta=subclass_inp_meta,
+            subclass_fw_graph_out_meta=subclass_fw_graph_out_meta,
+            subclass_tangent_meta=subclass_tangent_meta,
             grad_enabled_mutation=grad_enabled_mutation,
             static_input_indices=static_input_indices,
             tokens=mode._tokens,
diff --git a/torch/_functorch/_aot_autograd/frontend_utils.py b/torch/_functorch/_aot_autograd/frontend_utils.py
index 5547aabd5bc5a..5691a2641a301 100644
--- a/torch/_functorch/_aot_autograd/frontend_utils.py
+++ b/torch/_functorch/_aot_autograd/frontend_utils.py
@@ -2,12 +2,13 @@
 
 import warnings
 from contextlib import contextmanager
-from typing import Any, cast, Optional, TYPE_CHECKING
+from typing import Any, cast, TYPE_CHECKING
 
 import torch
 import torch.utils._pytree as pytree
 from torch._guards import detect_fake_mode
 from torch._library.opaque_object import is_opaque_type
+from torch._opaque_base import OpaqueBase
 from torch._subclasses import FakeTensor, FakeTensorMode
 from torch.fx.experimental.proxy_tensor import _pytree_subclasses_that_lose_info
 from torch.fx.experimental.symbolic_shapes import ShapeEnv
@@ -31,12 +32,54 @@ def process_inputs(
     flat_args: list[Any],
     aot_config: AOTConfig,
     fake_mode: FakeTensorMode,
-    shape_env: Optional[ShapeEnv],
+    shape_env: ShapeEnv | None,
     ignore_shape_env: bool = False,
-) -> FakifiedFlatArgs:
+) -> tuple[FakifiedFlatArgs, list[int]]:
+    """Convert real tensor inputs into fake tensors for AOT autograd tracing.
+
+    Called at compile time (not runtime) to produce the fake inputs that AOT
+    autograd traces through. Each real tensor is converted to a FakeTensor
+    via ``fake_mode.from_tensor``, preserving shape, dtype, device, and
+    symbolic shape information from the ShapeEnv. Non-tensor inputs (ints,
+    SymInts, ScriptObjects) are converted or passed through as appropriate.
+
+    Tensor subclass inputs (DTensor, etc.) are fakified recursively by
+    walking their ``__tensor_flatten__`` attrs. AsyncCollectiveTensors are
+    resolved via ``trigger_wait()`` before fakification so they don't appear
+    in the traced metadata (see below).
+
+    Called from ``aot_function``, ``aot_module_simplified``, and
+    ``aot_export_module`` — anywhere AOT autograd needs fake inputs before
+    graph capture.
+
+    Returns:
+        A tuple of (fakified_args, act_input_indices) where act_input_indices
+        records which positions held AsyncCollectiveTensors. These indices are
+        stored on ViewAndMutationMeta so that the runtime wrapper can emit
+        direct trigger_wait() calls on those positions.
+    """
+    # Resolve AsyncCollectiveTensors before tracing. ACTs are transient
+    # eager-mode wrappers for async collective overlap; if they leak into the
+    # traced graph as input types, AOT autograd records them in
+    # SubclassCreationMeta for output tangent metadata. At runtime, autograd
+    # produces plain tensor tangents, causing a type mismatch. Unwrapping
+    # here prevents ACT from appearing in the traced metadata.
+    try:
+        from torch.distributed._functional_collectives import AsyncCollectiveTensor
+    except ImportError:
+        AsyncCollectiveTensor = None
+
+    act_input_indices: list[int] = []
+    if AsyncCollectiveTensor is not None:
+        for i, a in enumerate(flat_args):
+            if isinstance(a, AsyncCollectiveTensor):
+                act_input_indices.append(i)
+                flat_args[i] = a.trigger_wait()
+
     with fake_mode:
 
         def convert(idx: int, x: Any) -> Any:
+            nonlocal ignore_shape_env
             if shape_env is not None and not ignore_shape_env:
                 from torch._dynamo.source import ConstantSource
 
@@ -66,11 +109,28 @@ def convert(idx: int, x: Any) -> Any:
                 return x
             if is_traceable_wrapper_subclass(x):
                 attrs, _ = x.__tensor_flatten__()
-                if all(isinstance(getattr(x, attr), FakeTensor) for attr in attrs):
-                    if all(getattr(x, attr).fake_mode is fake_mode for attr in attrs):
-                        return x
-                    # FakeTensor subclass from a different mode.
-                    # Fall through to refakify.
+                # See if all inner tensors are FakeTensors from this mode
+                all_this_fake = True
+                for a in attrs:
+                    match getattr(x, a):
+                        case FakeTensor() as v:
+                            if v.fake_mode is not fake_mode:
+                                # FakeTensor subclass from a different mode.
+                                # Fall through to refakify.
+                                all_this_fake = False
+                                break
+                        case torch.Tensor():
+                            all_this_fake = False
+                            break
+                        case OpaqueBase():
+                            pass
+                        case unexpected:
+                            raise AssertionError(
+                                f"expected Tensor or OpaqueBase, got {type(unexpected)}"
+                            )
+
+                if all_this_fake:
+                    return x
 
             # see note [Tensor Fakification and Symbol Caching]
             symbolic_context = None
@@ -101,12 +161,14 @@ def convert(idx: int, x: Any) -> Any:
             )
             return result
 
-        return FakifiedFlatArgs([convert(idx, x) for idx, x in enumerate(flat_args)])
+        return FakifiedFlatArgs(
+            [convert(idx, x) for idx, x in enumerate(flat_args)]
+        ), act_input_indices
 
 
 def construct_fake_mode(
     flat_args: list[Any], aot_config: AOTConfig
-) -> tuple[FakeTensorMode, Optional[ShapeEnv]]:
+) -> tuple[FakeTensorMode, ShapeEnv | None]:
     fake_mode = detect_fake_mode(flat_args)
     if fake_mode is None:
         shape_env = ShapeEnv() if aot_config.dynamic_shapes else None
diff --git a/torch/_functorch/_aot_autograd/functional_utils.py b/torch/_functorch/_aot_autograd/functional_utils.py
index ba2d01bf7cb00..9770fc3956686 100644
--- a/torch/_functorch/_aot_autograd/functional_utils.py
+++ b/torch/_functorch/_aot_autograd/functional_utils.py
@@ -15,6 +15,7 @@
 from torch import Tensor
 from torch._C import _functionalization
 from torch._logging import getArtifactLogger
+from torch._opaque_base import OpaqueBase
 from torch._subclasses.fake_tensor import FakeTensor
 from torch._subclasses.functional_tensor import FunctionalTensor
 from torch._subclasses.meta_utils import is_sparse_any
@@ -49,7 +50,15 @@ def sync_functional_tensor(t: torch.Tensor) -> None:
     if is_traceable_wrapper_subclass(t):
         attrs, _ctx = t.__tensor_flatten__()  # type: ignore[attr-defined]
         for attr in attrs:
-            sync_functional_tensor(getattr(t, attr))
+            match getattr(t, attr):
+                case Tensor() as inner:
+                    sync_functional_tensor(inner)
+                case OpaqueBase():
+                    pass
+                case unexpected:
+                    raise AssertionError(
+                        f"expected Tensor or OpaqueBase, got {type(unexpected)}"
+                    )
     else:
         torch._sync(t)
 
@@ -83,14 +92,24 @@ def is_fun(t: object) -> TypeGuard[FunctionalTensor | Tensor]:
         # goes at the bottom.
         # recurse here, so we can support nested wrapper subclasses
         t_attrs, _ = t.__tensor_flatten__()  # type: ignore[attr-defined]
-        t_inners = [getattr(t, attr) for attr in t_attrs]
-        any_fun = any(is_fun(x) for x in t_inners)
-        all_fun = all(is_fun(x) for x in t_inners)
-        if any_fun != all_fun:
-            raise AssertionError(
-                f"mixed functional/non-functional inner tensors: any_fun={any_fun}, all_fun={all_fun}"
-            )
-        return any_fun
+        got_fun: bool | None = None
+        for attr in t_attrs:
+            match getattr(t, attr):
+                case Tensor() as v:
+                    fun = is_fun(v)
+                    if got_fun is None:
+                        got_fun = fun
+                    elif got_fun != fun:
+                        raise AssertionError(
+                            "mixed functional/non-functional inner tensors"
+                        )
+                case OpaqueBase():
+                    pass
+                case unexpected:
+                    raise AssertionError(
+                        f"expected Tensor or OpaqueBase, got {type(unexpected)}"
+                    )
+        return got_fun or False
 
     return isinstance(t, FunctionalTensor)
 
@@ -103,7 +122,18 @@ def has_data_mutation(t: object) -> bool:
     if is_traceable_wrapper_subclass(t):
         attrs, _ = t.__tensor_flatten__()
         # A tensor subclass was updated if any of its inner elements were updated
-        return any(has_data_mutation(getattr(t, attr)) for attr in attrs)
+        for attr in attrs:
+            match getattr(t, attr):
+                case Tensor() as v:
+                    if has_data_mutation(v):
+                        return True
+                case OpaqueBase():
+                    pass
+                case unexpected:
+                    raise AssertionError(
+                        f"expected Tensor or OpaqueBase, got {type(unexpected)}"
+                    )
+        return False
     else:
         if isinstance(t, torch.Tensor):
             if not isinstance(t, FunctionalTensor):
@@ -116,9 +146,18 @@ def are_all_mutations_hidden_from_autograd(t: object) -> bool:
     if is_traceable_wrapper_subclass(t):
         attrs, _ = t.__tensor_flatten__()
         # If all inner elements are mutations hidden from autograd, then it is a mutation hidden from autograd.
-        return all(
-            are_all_mutations_hidden_from_autograd(getattr(t, attr)) for attr in attrs
-        )
+        for attr in attrs:
+            match getattr(t, attr):
+                case Tensor() as v:
+                    if not are_all_mutations_hidden_from_autograd(v):
+                        return False
+                case OpaqueBase():
+                    pass
+                case unexpected:
+                    raise AssertionError(
+                        f"expected Tensor or OpaqueBase, got {type(unexpected)}"
+                    )
+        return True
     elif isinstance(t, torch.Tensor):
         if not isinstance(t, FunctionalTensor):
             raise AssertionError(f"expected FunctionalTensor, got {type(t)}")
@@ -130,10 +169,18 @@ def are_all_mutations_hidden_from_autograd(t: object) -> bool:
 def are_all_mutations_under_no_grad_or_inference_mode(t: torch.Tensor) -> bool:
     if is_traceable_wrapper_subclass(t):
         attrs, _ = t.__tensor_flatten__()
-        return all(
-            are_all_mutations_under_no_grad_or_inference_mode(getattr(t, attr))
-            for attr in attrs
-        )
+        for attr in attrs:
+            match getattr(t, attr):
+                case Tensor() as v:
+                    if not are_all_mutations_under_no_grad_or_inference_mode(v):
+                        return False
+                case OpaqueBase():
+                    pass
+                case unexpected:
+                    raise AssertionError(
+                        f"expected Tensor or OpaqueBase, got {type(unexpected)}"
+                    )
+        return True
     else:
         if not isinstance(t, FunctionalTensor):
             raise AssertionError(f"expected FunctionalTensor, got {type(t)}")
@@ -145,10 +192,19 @@ def are_all_mutations_under_no_grad_or_inference_mode(t: torch.Tensor) -> bool:
 def was_inductor_storage_resized(t: object) -> bool:
     if is_traceable_wrapper_subclass(t):
         attrs, _ = t.__tensor_flatten__()
-        if any(was_inductor_storage_resized(getattr(t, attr)) for attr in attrs):
-            raise RuntimeError(
-                f"storage resizing is not supported on tensor subclass: {type(t)}"
-            )
+        for attr in attrs:
+            match getattr(t, attr):
+                case Tensor() as v:
+                    if was_inductor_storage_resized(v):
+                        raise RuntimeError(
+                            f"storage resizing is not supported on tensor subclass: {type(t)}"
+                        )
+                case OpaqueBase():
+                    pass
+                case unexpected:
+                    raise AssertionError(
+                        f"expected Tensor or OpaqueBase, got {type(unexpected)}"
+                    )
         return False
     elif not isinstance(t, torch.Tensor):
         return False
@@ -172,16 +228,24 @@ def has_metadata_mutation(
     if is_traceable_wrapper_subclass(f_arg):
         attrs, _ = f_arg.__tensor_flatten__()
         # A tensor subclass was updated if any of its inner elements were updated
-        f_inner_ts = [getattr(f_arg, attr) for attr in attrs]
-        inner_ts = [getattr(arg, attr) for attr in attrs]
-        return any(
-            has_metadata_mutation(
-                f_inner_t,
-                inner_t,
-                check_only_storage_mutation=check_only_storage_mutation,
-            )
-            for f_inner_t, inner_t in zip(f_inner_ts, inner_ts)
-        )
+        for attr in attrs:
+            match getattr(f_arg, attr):
+                case Tensor():
+                    f_inner_t = getattr(f_arg, attr)
+                    inner_t = getattr(arg, attr)
+                    if has_metadata_mutation(
+                        f_inner_t,
+                        inner_t,
+                        check_only_storage_mutation=check_only_storage_mutation,
+                    ):
+                        return True
+                case OpaqueBase():
+                    pass
+                case unexpected:
+                    raise AssertionError(
+                        f"expected Tensor or OpaqueBase, got {type(unexpected)}"
+                    )
+        return False
     else:
         if not isinstance(f_arg, torch.Tensor):
             if isinstance(arg, torch.Tensor):
@@ -440,10 +504,18 @@ def was_tensor_updated(arg: torch.Tensor, new_arg: torch.Tensor) -> bool:
         if attrs != new_attrs:
             raise AssertionError(f"attrs mismatch: {attrs} != {new_attrs}")
         # A tensor subclass was updated if any of its inner elements were updated
-        return any(
-            was_tensor_updated(getattr(arg, attr), getattr(new_arg, attr))
-            for attr in attrs
-        )
+        for attr in attrs:
+            match getattr(arg, attr):
+                case Tensor() as v:
+                    if was_tensor_updated(v, getattr(new_arg, attr)):
+                        return True
+                case OpaqueBase():
+                    pass
+                case unexpected:
+                    raise AssertionError(
+                        f"expected Tensor or OpaqueBase, got {type(unexpected)}"
+                    )
+        return False
     else:
         return arg is not new_arg
 
@@ -466,10 +538,18 @@ def was_tensor_metadata_updated(arg: Any, new_arg: Any) -> bool:
         if attrs != new_attrs:
             raise AssertionError(f"attrs mismatch: {attrs} != {new_attrs}")
         # A tensor subclass was updated if any of its inner elements were updated
-        return any(
-            was_tensor_metadata_updated(getattr(arg, attr), getattr(new_arg, attr))
-            for attr in attrs
-        )
+        for attr in attrs:
+            match getattr(arg, attr):
+                case Tensor() as v:
+                    if was_tensor_metadata_updated(v, getattr(new_arg, attr)):
+                        return True
+                case OpaqueBase():
+                    pass
+                case unexpected:
+                    raise AssertionError(
+                        f"expected Tensor or OpaqueBase, got {type(unexpected)}"
+                    )
+        return False
     else:
         return arg is not new_arg and StorageWeakRef(
             arg.untyped_storage()
diff --git a/torch/_functorch/_aot_autograd/fx_utils.py b/torch/_functorch/_aot_autograd/fx_utils.py
index a8945f1c503cf..fc97619de7d3f 100644
--- a/torch/_functorch/_aot_autograd/fx_utils.py
+++ b/torch/_functorch/_aot_autograd/fx_utils.py
@@ -5,7 +5,7 @@
 recommend reading :mod:torch._functorch._aot_autograd.descriptors`.
 """
 
-from typing import NoReturn, Optional, Union
+from typing import NoReturn
 
 import torch.fx as fx
 
@@ -26,7 +26,7 @@
 
 
 def _raise_autograd_subclass_not_implemented(
-    n: fx.Node, desc: Union[AOTInput, AOTOutput]
+    n: fx.Node, desc: AOTInput | AOTOutput
 ) -> NoReturn:
     raise RuntimeError(
         "Subclasses are currently not supported by this function, but a desugared subclass input "
@@ -43,7 +43,7 @@ def _raise_autograd_subclass_not_implemented(
 
 def get_all_input_and_grad_nodes(
     g: fx.Graph,
-) -> dict[DifferentiableAOTInput, tuple[fx.Node, Optional[fx.Node]]]:
+) -> dict[DifferentiableAOTInput, tuple[fx.Node, fx.Node | None]]:
     """
     Given a joint graph with descriptors (meta['desc'] on placeholders and
     output), returns the node for every input and its corresponding grad
@@ -69,7 +69,7 @@ def get_all_input_and_grad_nodes(
         is not supported by API as there is not necessarily a 1-1 correspondence
         between inputs and grads when subclasses are involved.
     """
-    input_index: dict[DifferentiableAOTInput, tuple[fx.Node, Optional[fx.Node]]] = {}
+    input_index: dict[DifferentiableAOTInput, tuple[fx.Node, fx.Node | None]] = {}
     for n in g.nodes:
         if n.op == "placeholder":
             desc = n.meta["desc"]
@@ -99,7 +99,7 @@ def get_all_input_and_grad_nodes(
 
 def get_all_output_and_tangent_nodes(
     g: fx.Graph,
-) -> dict[DifferentiableAOTOutput, tuple[fx.Node, Optional[fx.Node]]]:
+) -> dict[DifferentiableAOTOutput, tuple[fx.Node, fx.Node | None]]:
     """Get all output nodes and their corresponding tangent nodes from a joint graph.
 
     Similar to get_all_input_and_grad_nodes, but returns output nodes paired with
@@ -124,7 +124,7 @@ def get_all_output_and_tangent_nodes(
         is not supported by API as there is not necessarily a 1-1 correspondence
         between outputs and tangents when subclasses are involved.
     """
-    output_index: dict[DifferentiableAOTOutput, tuple[fx.Node, Optional[fx.Node]]] = {}
+    output_index: dict[DifferentiableAOTOutput, tuple[fx.Node, fx.Node | None]] = {}
     for n in g.nodes:
         if n.op == "output":
             desc = n.meta["desc"]
@@ -153,7 +153,7 @@ def get_all_output_and_tangent_nodes(
 
 def get_param_and_grad_nodes(
     graph: fx.Graph,
-) -> dict[ParamAOTInput, tuple[fx.Node, Optional[fx.Node]]]:
+) -> dict[ParamAOTInput, tuple[fx.Node, fx.Node | None]]:
     """Get parameter nodes and their corresponding gradient nodes from a joint graph.
 
     Args:
@@ -173,7 +173,7 @@ def get_param_and_grad_nodes(
 
 def get_plain_input_and_grad_nodes(
     graph: fx.Graph,
-) -> dict[PlainAOTInput, tuple[fx.Node, Optional[fx.Node]]]:
+) -> dict[PlainAOTInput, tuple[fx.Node, fx.Node | None]]:
     """Get plain input nodes and their corresponding gradient nodes from a joint graph.
 
     Args:
@@ -193,7 +193,7 @@ def get_plain_input_and_grad_nodes(
 
 def get_plain_output_and_tangent_nodes(
     graph: fx.Graph,
-) -> dict[PlainAOTOutput, tuple[fx.Node, Optional[fx.Node]]]:
+) -> dict[PlainAOTOutput, tuple[fx.Node, fx.Node | None]]:
     """Get plain output nodes and their corresponding tangent nodes from a joint graph.
 
     Args:
@@ -212,7 +212,7 @@ def get_plain_output_and_tangent_nodes(
 
 
 def _raise_fqn_subclass_not_implemented(
-    n: fx.Node, desc: Union[AOTInput, AOTOutput]
+    n: fx.Node, desc: AOTInput | AOTOutput
 ) -> NoReturn:
     raise RuntimeError(
         "Subclasses are currently not supported by this function, but a desugared subclass input "
diff --git a/torch/_functorch/_aot_autograd/graph_capture.py b/torch/_functorch/_aot_autograd/graph_capture.py
index 9c4ab11ff8bd6..f1752bd4d13e7 100644
--- a/torch/_functorch/_aot_autograd/graph_capture.py
+++ b/torch/_functorch/_aot_autograd/graph_capture.py
@@ -6,7 +6,7 @@
 import contextlib
 import dataclasses
 from collections.abc import Callable
-from typing import Any, Optional
+from typing import Any
 
 import torch
 import torch.utils._pytree as pytree
@@ -39,6 +39,7 @@
     insert_backward_syncs,
     populate_fw_metadata_with_stream_indices,
     sync_deallocations,
+    wrap_all_sync_nodes_with_control_deps,
 )
 from .utils import (
     call_and_expect_output_descs,
@@ -54,12 +55,45 @@
 aot_graphs_log = getArtifactLogger(__name__, "aot_graphs")
 
 
+def _extract_tangent_source_stack_traces(
+    fx_g: torch.fx.GraphModule,
+    fw_metadata: ViewAndMutationMeta,
+) -> None:
+    from .descriptors import PlainAOTOutput, TangentAOTInput
+
+    if not fw_metadata.traced_tangents_descs:
+        return
+
+    output_node = list(fx_g.graph.nodes)[-1]
+    all_outputs = output_node.args[0]
+
+    stack_traces: list[str | None] = []
+    got_one = False
+
+    for desc in fw_metadata.traced_tangents_descs:
+        stack_trace = None
+
+        if isinstance(desc, TangentAOTInput):
+            output_desc = desc.output
+            if isinstance(output_desc, PlainAOTOutput) and output_desc.idx < len(
+                all_outputs
+            ):
+                output_arg = all_outputs[output_desc.idx]
+                if isinstance(output_arg, torch.fx.Node):
+                    stack_trace = output_arg.meta.get("stack_trace", None)
+                    got_one = True
+
+        stack_traces.append(stack_trace)
+
+    if got_one:
+        fw_metadata.tangent_source_stack_traces = stack_traces
+
+
 def _create_graph(
     f: Callable[..., Any],
     args: list[torch.Tensor],
-    args_descs: Optional[
-        list[AOTInput]
-    ] = None,  # keep compat with old clients; maybe we should split into two impls
+    args_descs: list[AOTInput]
+    | None = None,  # keep compat with old clients; maybe we should split into two impls
     *,
     aot_config: AOTConfig,
 ) -> torch.fx.GraphModule:
@@ -157,6 +191,96 @@ def _detach_and_copy_item_memo(t: torch.Tensor) -> torch.Tensor:
     return detached_t
 
 
+@dataclasses.dataclass
+class _GraphCaptureTracingResult:
+    fn_to_trace: Callable[..., Any]
+    flat_args: Any
+    flat_args_descs: Any
+    maybe_subclass_meta: SubclassMeta | None
+
+
+def _detach_traced_inputs(flat_args: Any) -> Any:
+    if detect_fake_mode():
+        detach_tensor = _detach_and_copy_item_memo
+    else:
+
+        def detach_tensor(t: torch.Tensor) -> torch.Tensor:
+            return t.detach()
+
+    return pytree.tree_map_only(torch.Tensor, detach_tensor, flat_args)
+
+
+def _prepare_graph_capture_tracing(
+    fn_to_trace: Callable[..., Any],
+    flat_args: Any,
+    flat_args_descs: Any,
+    flat_fn: TraceFn,
+    *,
+    fw_metadata: ViewAndMutationMeta,
+    aot_config: AOTConfig,
+    trace_joint: bool,
+    joint_fn_handle: Any | None = None,
+) -> _GraphCaptureTracingResult:
+    if aot_config.disable_functionalization:
+        updated_flat_args, updated_flat_args_descs = flat_args, flat_args_descs
+    else:
+        fn_to_trace, updated_flat_args, updated_flat_args_descs = (
+            create_functionalized_fn(
+                fn_to_trace,
+                flat_args,
+                flat_args_descs,
+                meta=fw_metadata,
+                aot_config=aot_config,
+                trace_joint=trace_joint,
+                joint_fn_handle=joint_fn_handle,
+            )
+        )
+
+    subclass_tracing_info = aot_dispatch_subclass(
+        fn_to_trace,
+        updated_flat_args,
+        updated_flat_args_descs,
+        is_joint_structure=trace_joint,
+        meta=fw_metadata,
+        fw_only=flat_fn,
+    )
+    fn_to_trace = subclass_tracing_info.plain_tensor_trace_fn
+    updated_flat_args = subclass_tracing_info.plain_tensor_args
+    updated_flat_args_descs = subclass_tracing_info.plain_tensor_args_descs
+
+    if not aot_config.disable_functionalization:
+        fn_to_trace, updated_flat_args, updated_flat_args_descs = (
+            handle_effect_tokens_fn(
+                fn_to_trace,
+                updated_flat_args,
+                updated_flat_args_descs,
+                meta=fw_metadata,
+                trace_joint=trace_joint,
+            )
+        )
+
+    return _GraphCaptureTracingResult(
+        fn_to_trace=fn_to_trace,
+        flat_args=updated_flat_args,
+        flat_args_descs=updated_flat_args_descs,
+        maybe_subclass_meta=subclass_tracing_info.maybe_subclass_meta,
+    )
+
+
+def _create_graph_and_save_traced_inputs(
+    fn_to_trace: Callable[..., Any],
+    flat_args: Any,
+    flat_args_descs: Any,
+    *,
+    aot_config: AOTConfig,
+) -> tuple[torch.fx.GraphModule, Any]:
+    saved_flat_args = _detach_traced_inputs(flat_args)
+    return (
+        _create_graph(fn_to_trace, flat_args, flat_args_descs, aot_config=aot_config),
+        saved_flat_args,
+    )
+
+
 def aot_dispatch_base_graph(
     flat_fn: TraceFn,
     flat_args: list[FxValue],
@@ -164,7 +288,7 @@ def aot_dispatch_base_graph(
     aot_config: AOTConfig,
     *,
     fw_metadata: ViewAndMutationMeta,
-) -> tuple[torch.fx.GraphModule, list[FxValue], list[AOTInput], Optional[SubclassMeta]]:
+) -> tuple[torch.fx.GraphModule, list[FxValue], list[AOTInput], SubclassMeta | None]:
     # aot_dispatch_base requires functionalization, but doesn't need to handle as many cases as the autograd case.
     # The cases that aot_dispatch_base doesn't need to handle include:
     # - outputs that are aliases of graph intermediates
@@ -178,59 +302,28 @@ def aot_dispatch_base_graph(
         fw_metadata,
         keep_data_input_mutations=aot_config.keep_inference_input_mutations,
     )
-
-    if aot_config.disable_functionalization:
-        updated_flat_args, updated_flat_args_descs = (
-            flat_args,
-            flat_args_descs,
-        )
-    else:
-        fn_to_trace, updated_flat_args, updated_flat_args_descs = (
-            create_functionalized_fn(
-                fn_to_trace,
-                flat_args,
-                flat_args_descs,
-                meta=fw_metadata,
-                aot_config=aot_config,
-                trace_joint=False,
-            )
-        )
-
     # TODO: replace with AOTDispatchSubclassWrapper once we refactor
     # fn_input_mutations_to_outputs and create_functionalized_fn
     # into CompilerWrappers.
-    (
+    tracing_state = _prepare_graph_capture_tracing(
         fn_to_trace,
-        updated_flat_args_subclasses_desugared,
-        updated_flat_args_subclasses_desugared_descs,
-        maybe_subclass_meta,
-    ) = aot_dispatch_subclass(
-        fn_to_trace,
-        updated_flat_args,
-        updated_flat_args_descs,
-        is_joint_structure=False,
-        meta=fw_metadata,
-        fw_only=flat_fn,
+        flat_args,
+        flat_args_descs,
+        flat_fn,
+        fw_metadata=fw_metadata,
+        aot_config=aot_config,
+        trace_joint=False,
     )
-
-    if not aot_config.disable_functionalization:
-        (
-            fn_to_trace,
-            updated_flat_args_subclasses_desugared,
-            updated_flat_args_subclasses_desugared_descs,
-        ) = handle_effect_tokens_fn(
-            fn_to_trace,
-            updated_flat_args_subclasses_desugared,
-            updated_flat_args_subclasses_desugared_descs,
-            meta=fw_metadata,
-            trace_joint=False,
-        )
+    fn_to_trace = tracing_state.fn_to_trace
+    updated_flat_args_subclasses_desugared = tracing_state.flat_args
+    updated_flat_args_subclasses_desugared_descs = tracing_state.flat_args_descs
+    maybe_subclass_meta = tracing_state.maybe_subclass_meta
 
     aot_graphs_log.debug(
         "aot_config id: %s, fw_metadata=%s,subclass_metadata=%s",
-        str(aot_config.aot_id),
-        str(fw_metadata),
-        str(maybe_subclass_meta),
+        aot_config.aot_id,
+        fw_metadata,
+        maybe_subclass_meta,
     )
 
     # We track buffer assignments when exporting in non-strict mode.
@@ -244,27 +337,18 @@ def aot_dispatch_base_graph(
             mod_when_exporting_non_strict, assigned_buffers
         )
 
-    fake_mode = detect_fake_mode()
-    if fake_mode:
-        saved_updated_flat_args_subclasses_desugared = pytree.tree_map_only(
-            torch.Tensor,
-            _detach_and_copy_item_memo,
-            updated_flat_args_subclasses_desugared,
-        )
-    else:
-        saved_updated_flat_args_subclasses_desugared = pytree.tree_map_only(
-            torch.Tensor, lambda t: t.detach(), updated_flat_args_subclasses_desugared
-        )
-    saved_updated_flat_args_subclasses_desugared_descs = (
-        updated_flat_args_subclasses_desugared_descs
-    )
-
-    fw_module = _create_graph(
+    (
+        fw_module,
+        saved_updated_flat_args_subclasses_desugared,
+    ) = _create_graph_and_save_traced_inputs(
         fn_to_trace,
         updated_flat_args_subclasses_desugared,
         updated_flat_args_subclasses_desugared_descs,
         aot_config=aot_config,
     )
+    saved_updated_flat_args_subclasses_desugared_descs = (
+        updated_flat_args_subclasses_desugared_descs
+    )
 
     if aot_config.is_export and mod_when_exporting_non_strict is not None:
         # We update metadata to consider any assigned buffers as buffer mutations.
@@ -294,6 +378,8 @@ def aot_dispatch_base_graph(
     if not aot_config.disable_functionalization:
         copy_count = assert_functional_graph(fw_module.graph)
         assign_epilogue_copy_streams(fw_module)
+        # Wrap sync nodes with control_deps to prevent reordering
+        wrap_all_sync_nodes_with_control_deps(fw_module)
         # Populate fw_metadata with stream indices from the compiled graph
         populate_fw_metadata_with_stream_indices(fw_module, fw_metadata)
         fw_module.graph.eliminate_dead_code()
@@ -363,11 +449,10 @@ def aot_dispatch_base_graph(
         )
 
     # TODO: should factor this into a separate function for export that always only returns just the graph.
-    if aot_config.is_export:
-        if maybe_subclass_meta is not None:
-            raise AssertionError(
-                "aot_export_module does not support tensor subclass inputs for now."
-            )
+    if aot_config.is_export and maybe_subclass_meta is not None:
+        raise AssertionError(
+            "aot_export_module does not support tensor subclass inputs for now."
+        )
     return (
         fw_module,
         saved_updated_flat_args_subclasses_desugared,
@@ -391,7 +476,7 @@ def aot_dispatch_autograd_graph(
     torch.fx.GraphModule,
     tuple[list[Any], list[Any]],
     tuple[list[AOTInput], list[AOTInput]],
-    Optional[SubclassMeta],
+    SubclassMeta | None,
 ]:
     # NB: flat_fn here is the original user function (as far as
     # aot_module_simplified is concerned)
@@ -413,51 +498,23 @@ def aot_dispatch_autograd_graph(
     )
     # pyrefly: ignore[missing-attribute]
     joint_fn_handle = joint_fn_to_trace.handle
-
-    if aot_config.disable_functionalization:
-        updated_joint_inputs, updated_joint_inputs_descs = (
-            joint_inputs,
-            joint_inputs_descs,
-        )
-    else:
-        joint_fn_to_trace, updated_joint_inputs, updated_joint_inputs_descs = (
-            create_functionalized_fn(
-                joint_fn_to_trace,
-                joint_inputs,
-                joint_inputs_descs,
-                meta=fw_metadata,
-                aot_config=aot_config,
-                trace_joint=True,
-                joint_fn_handle=joint_fn_handle,
-            )
-        )
-
     # TODO: replace with AOTDispatchSubclassWrapper once we refactor
     # fn_input_mutations_to_outputs and create_functionalized_fn
     # into CompilerWrappers.
-    subclass_tracing_info = aot_dispatch_subclass(
+    tracing_state = _prepare_graph_capture_tracing(
         joint_fn_to_trace,
-        updated_joint_inputs,
-        updated_joint_inputs_descs,
-        is_joint_structure=True,
-        meta=fw_metadata,
-        fw_only=flat_fn,
+        joint_inputs,
+        joint_inputs_descs,
+        flat_fn,
+        fw_metadata=fw_metadata,
+        aot_config=aot_config,
+        trace_joint=True,
+        joint_fn_handle=joint_fn_handle,
     )
-
-    joint_fn_to_trace = subclass_tracing_info.plain_tensor_trace_fn
-    updated_joint_inputs = subclass_tracing_info.plain_tensor_args
-    updated_joint_inputs_descs = subclass_tracing_info.plain_tensor_args_descs
-
-    if not aot_config.disable_functionalization:
-        (joint_fn_to_trace, updated_joint_inputs, updated_joint_inputs_descs) = (
-            handle_effect_tokens_fn(
-                joint_fn_to_trace,
-                updated_joint_inputs,
-                updated_joint_inputs_descs,
-                meta=fw_metadata,
-                trace_joint=True,
-            )
-        )
+    joint_fn_to_trace = tracing_state.fn_to_trace
+    updated_joint_inputs = tracing_state.flat_args
+    updated_joint_inputs_descs = tracing_state.flat_args_descs
+    maybe_subclass_meta = tracing_state.maybe_subclass_meta
 
     # When we call _create_graph, this may mutate the metadata of joint
     # inputs.  But callers are expecting to get the original joint inputs.  So
@@ -467,19 +524,7 @@ def aot_dispatch_autograd_graph(
     # This destroys requires_grad/grad_fn information.  However, backends
     # beneath AOTAutograd are indifferent to this information, so it doesn't
     # matter.
-
-    fake_mode = detect_fake_mode()
-    if fake_mode:
-        saved_updated_joint_inputs = pytree.tree_map_only(
-            torch.Tensor, _detach_and_copy_item_memo, updated_joint_inputs
-        )
-    else:
-        saved_updated_joint_inputs = pytree.tree_map_only(
-            torch.Tensor, lambda t: t.detach(), updated_joint_inputs
-        )
-    maybe_subclass_meta = subclass_tracing_info.maybe_subclass_meta
-
-    fx_g = _create_graph(
+    fx_g, saved_updated_joint_inputs = _create_graph_and_save_traced_inputs(
         joint_fn_to_trace,
         updated_joint_inputs,
         updated_joint_inputs_descs,
@@ -507,10 +552,17 @@ def aot_dispatch_autograd_graph(
     # is distinct from their allocation stream
     sync_deallocations(fx_g)
 
+    # Wrap sync nodes with control_deps to prevent reordering
+    # (must be after sync_deallocations which inserts additional sync nodes)
+    wrap_all_sync_nodes_with_control_deps(fx_g)
+
     # Populate fw_metadata with stream indices from the compiled graph
     # NB: This needs to be done after the above stream assignments
     populate_fw_metadata_with_stream_indices(fx_g, fw_metadata)
 
+    # this helps users identify which forward output to call .detach() on.
+    _extract_tangent_source_stack_traces(fx_g, fw_metadata)
+
     fx_g.graph.eliminate_dead_code()
     if not aot_config.disable_functionalization:
         # There should be *NO* mutating ops in the graph at this point.
@@ -521,11 +573,10 @@ def aot_dispatch_autograd_graph(
     # TODO: in AOTAutograd, we create metadata like _indices_of_inps_to_detach to detect
     # when we need to manually detach() some inputs in the forward.
     # Higher order ops might eventually need to do the same.
-    if aot_config.is_export:
-        if maybe_subclass_meta is not None:
-            raise AssertionError(
-                "aot_export_module does not support tensor subclass inputs for now."
-            )
+    if aot_config.is_export and maybe_subclass_meta is not None:
+        raise AssertionError(
+            "aot_export_module does not support tensor subclass inputs for now."
+        )
     return (
         fx_g,
         saved_updated_joint_inputs,
diff --git a/torch/_functorch/_aot_autograd/graph_capture_wrappers.py b/torch/_functorch/_aot_autograd/graph_capture_wrappers.py
index 69b766fda0904..bb26cca92171c 100644
--- a/torch/_functorch/_aot_autograd/graph_capture_wrappers.py
+++ b/torch/_functorch/_aot_autograd/graph_capture_wrappers.py
@@ -10,11 +10,12 @@
 4. dispatching subclasses
 """
 
+import typing
 import warnings
 from collections.abc import Callable, Generator
 from contextlib import AbstractContextManager, contextmanager, ExitStack, nullcontext
 from dataclasses import dataclass
-from typing import Any, Optional, TypeVar, Union
+from typing import Any, TypeVar
 from unittest.mock import patch
 
 import torch
@@ -23,6 +24,7 @@
 from torch import Tensor
 from torch._decomp.decompositions_for_rng import PhiloxStateTracker
 from torch._guards import detect_fake_mode
+from torch._opaque_base import OpaqueBase
 from torch._prims_common import CUDARngStateHelper
 from torch.fx.experimental.proxy_tensor import (
     _proxy_tensor_disable_update_tensor_tracker,
@@ -240,7 +242,7 @@ def inner_fn(
             # Also, only tensor outputs should participate in the backward
             # (in particular, Symint outputs in the forward graph shouldn't get tangents)
             and issubclass(meta.output_info[i].raw_type, Tensor)
-            and meta.output_info[i].requires_grad
+            and meta.output_info[i].requires_grad_for_backward
             for (i, x) in enumerate(outs)
         ]
 
@@ -291,7 +293,7 @@ class JointFnHandle:
 #     (the way this is handled is that we ensure any inputs that normally get mutated are cloned first)
 def create_joint(
     fn: Callable[..., Any],
-    primals_descs: Optional[list[AOTInput]] = None,
+    primals_descs: list[AOTInput] | None = None,
     *,
     aot_config: AOTConfig,
 ) -> Callable[..., Any]:
@@ -303,8 +305,8 @@ def create_joint(
     def inner_fn(
         primals: list[FxValue], tangents: list[FxValue]
     ) -> tuple[
-        tuple[list[FxValue], list[Optional[Tensor]]],
-        tuple[list[AOTOutput], list[Optional[AOTOutput]]],
+        tuple[list[FxValue], list[Tensor | None]],
+        tuple[list[AOTOutput], list[AOTOutput | None]],
     ]:
         outs_descs = None
         if primals_descs is None:
@@ -478,8 +480,8 @@ def inner_fn(
     def inner_fn_with_anomaly(
         primals: list[FxValue], tangents: list[FxValue]
     ) -> tuple[
-        tuple[list[FxValue], list[Optional[Tensor]]],
-        tuple[list[AOTOutput], list[Optional[AOTOutput]]],
+        tuple[list[FxValue], list[Tensor | None]],
+        tuple[list[AOTOutput], list[AOTOutput | None]],
     ]:
         with fx_traceback.preserve_node_meta(), warnings.catch_warnings():
             warnings.filterwarnings("ignore", "Anomaly Detection has been enabled.")
@@ -489,8 +491,8 @@ def inner_fn_with_anomaly(
     def joint_helper(
         primals: list[FxValue], tangents: list[FxValue]
     ) -> tuple[
-        tuple[list[FxValue], list[Optional[Tensor]]],
-        tuple[list[AOTOutput], list[Optional[AOTOutput]]],
+        tuple[list[FxValue], list[Tensor | None]],
+        tuple[list[AOTOutput], list[AOTOutput | None]],
     ]:
         return inner_fn_with_anomaly(primals, tangents)
 
@@ -516,13 +518,13 @@ def create_functionalized_rng_ops_wrapper(
         fake_mode = fake_mode_det
 
     def override_get_rng_state(
-        device: Union[int, str, torch.device] = "cuda",
+        device: int | str | torch.device = "cuda",
     ) -> Tensor:
         out = PhiloxStateTracker.get_state_as_tensor()
         return out
 
     def override_set_rng_state(
-        x: Tensor, device: Union[int, str, torch.device] = "cuda"
+        x: Tensor, device: int | str | torch.device = "cuda"
     ) -> None:
         PhiloxStateTracker.set_state_from_tensor(x)
 
@@ -671,7 +673,15 @@ def visit(e: Any) -> None:
             return
 
         for a in e.__tensor_flatten__()[0]:
-            visit(getattr(e, a))
+            match getattr(e, a):
+                case torch.Tensor() as inner:
+                    visit(inner)
+                case OpaqueBase():
+                    pass
+                case unexpected:
+                    raise AssertionError(
+                        f"expected Tensor or OpaqueBase, got {type(unexpected)}"
+                    )
 
     visit(t)
     return accum
@@ -718,8 +728,8 @@ def apply_in_graph_mutations(
     inpt_new: Tensor,
     f_inpt: Tensor,
     input_idx: int,
-    mcs: Optional[MutationCounters] = None,
-    applied_mcs: Optional[MutationCounters] = None,
+    mcs: MutationCounters | None = None,
+    applied_mcs: MutationCounters | None = None,
 ) -> None:
     if input_info.mutation_type != MutationType.MUTATED_IN_GRAPH:
         raise AssertionError(
@@ -742,7 +752,7 @@ def apply_in_graph_mutations(
     if input_info.mutates_storage_metadata:
         if mcs is None or mcs.mc_storage > applied_mcs.mc_storage:  # type: ignore[union-attr]
             with torch.no_grad():
-                # pyrefly: ignore[no-matching-overload]
+                # pyrefly: ignore [bad-argument-type, no-matching-overload]
                 inpt_old.set_(inpt_new)
 
     # Note [Ordering of resize_() and set_()]
@@ -843,11 +853,11 @@ def create_functionalized_fn(
     meta: ViewAndMutationMeta,
     aot_config: AOTConfig,
     trace_joint: bool,
-    joint_fn_handle: Optional[JointFnHandle] = None,
+    joint_fn_handle: JointFnHandle | None = None,
 ) -> Any:
     primals_after_forward = None
     f_args_after_forward = None
-    f_args_mutation_counters_after_forward: Optional[list[MutationCounters]] = None
+    f_args_mutation_counters_after_forward: list[MutationCounters] | None = None
     inputs_mutated_in_graph = [
         info.mutation_type == MutationType.MUTATED_IN_GRAPH for info in meta.input_info
     ]
@@ -856,7 +866,7 @@ def create_functionalized_fn(
     @simple_wraps(fn)
     def _functionalized_f_helper(
         *args: list[FxValue],
-    ) -> tuple[tuple[list[FxValue], list[Tensor]], list[Optional[AOTOutput]]]:
+    ) -> tuple[tuple[list[FxValue], list[Tensor]], list[AOTOutput | None]]:
         with maybe_enable_thunkify():
             # See Note [Disabling Functionalize TLS Above Python Functionalization]
             disable_above = torch._C._ExcludeDispatchKeyGuard(
@@ -960,7 +970,11 @@ def _post_forward(primals: Any) -> None:
                                 raise AssertionError(
                                     f"expected both before and after to be Tensors, got {type(before)} and {type(after)}"
                                 )
-                            before.copy_(after)
+                            # no_grad prevents the FakeTensor's requires_grad from
+                            # triggering check_inplace during tracing.  The
+                            # requires_grad case is checked at runtime instead
+                            with torch.no_grad():
+                                before.copy_(after)
                         meta.indices_of_inputs_that_requires_grad_with_mutations_in_bw.append(
                             idx
                         )
@@ -1076,6 +1090,7 @@ def _post_forward(primals: Any) -> None:
                         ):
                             apply_in_graph_mutations(
                                 inpt_info,
+                                # pyrefly: ignore [bad-argument-type]
                                 before,
                                 after,
                                 f_inpt,
@@ -1100,7 +1115,7 @@ def _post_forward(primals: Any) -> None:
                         != MutationType.MUTATED_IN_GRAPH
                     ):
                         continue
-                    mcs: Optional[MutationCounters] = None
+                    mcs: MutationCounters | None = None
                     if f_args_mutation_counters_after_forward is not None:
                         # This could happen for subclasses tracing
                         # Subclasses support for mutations in fw and bw is TBD.
@@ -1271,6 +1286,11 @@ def inner_fn(*args: Any) -> Any:
     else:
         args = [*additional_fwd_token_inputs, *args]
         args_descs = [*additional_fwd_token_inputs_descs, *args_descs]
+
+        if num_tokens > 0:
+            meta.static_input_indices = [
+                idx + num_tokens for idx in meta.static_input_indices
+            ]
     return inner_fn, args, args_descs
 
 
@@ -1285,9 +1305,9 @@ def inner_fn(*args: Any) -> Any:
 #   Why do we need this? We need to collect updated ViewAndMutationMeta on our new dense -> dense functions.
 #   In particular, we need this to tell the partitioner how many dense forward outputs there are.
 def aot_dispatch_subclass(
-    flat_fn_maybe_joint: Union[JointTraceFn, TraceFn],
-    args: Union[list[FxValue], tuple[list[FxValue], list[FxValue]]],
-    args_descs: Union[list[AOTInput], tuple[list[AOTInput], list[AOTInput]]],
+    flat_fn_maybe_joint: JointTraceFn | TraceFn,
+    args: list[FxValue] | tuple[list[FxValue], list[FxValue]],
+    args_descs: list[AOTInput] | tuple[list[AOTInput], list[AOTInput]],
     *,
     is_joint_structure: bool,
     meta: ViewAndMutationMeta,
@@ -1338,11 +1358,15 @@ def inner_fn(fn: Callable[..., Any], args: Any, *, use_trace_joint: bool) -> Any
             # Add extra symints as outputs to the forward/backward graphs
             # ignore nested ints here
             forward_outs, forward_outs_descs = unwrap_tensor_subclasses(
-                wrapped_outs[0], wrapped_outs_descs[0], append_symints=True
+                wrapped_outs[0],
+                wrapped_outs_descs[0],
+                append_symints=True,
             )
             # ignore nested ints here
             backward_outs, backward_outs_descs = unwrap_tensor_subclasses(
-                wrapped_outs[1], wrapped_outs_descs[1], append_symints=True
+                wrapped_outs[1],
+                wrapped_outs_descs[1],
+                append_symints=True,
             )
             return (
                 (forward_outs, backward_outs),
@@ -1376,42 +1400,53 @@ def inner_fw_only(*args: Any) -> Any:
         return inner_fn(inner_fw_only, primals, use_trace_joint=False)
 
     if is_joint_structure:
+        primals_wrapped: list[FxValue] = typing.cast(list[FxValue], args[0])
+        primals_wrapped_descs: list[AOTInput] = typing.cast(
+            list[AOTInput], args_descs[0]
+        )
+        tangents_wrapped: list[FxValue] = typing.cast(list[FxValue], args[1])
+        tangents_wrapped_descs: list[AOTInput] = typing.cast(
+            list[AOTInput], args_descs[1]
+        )
+
         # Add extra symints (size/strides) as input to the forward graph
         primals_unwrapped_pair = unwrap_tensor_subclasses(
-            args[0],  # type: ignore[arg-type]
-            args_descs[0],  # type: ignore[arg-type]
+            primals_wrapped,
+            primals_wrapped_descs,
             append_symints=True,
         )
         # We pass append_symints=False here because the partitioner will
-        # capture and add any extra argument
+        # capture and add any extra argument.
         tangents_unwrapped_pair = unwrap_tensor_subclasses(
-            args[1],  # type: ignore[arg-type]
-            args_descs[1],  # type: ignore[arg-type]
+            tangents_wrapped,
+            tangents_wrapped_descs,
             append_symints=False,
         )
 
         args_unwrapped = (primals_unwrapped_pair[0], tangents_unwrapped_pair[0])
         args_descs_unwrapped = (primals_unwrapped_pair[1], tangents_unwrapped_pair[1])
         remapped_static_indices = remap_unwrapped_subclass_arg_indices(
-            args[0],  # type: ignore[arg-type]
+            primals_wrapped,
             meta.static_input_indices,  # type: ignore[arg-type]
         )
+
+        primals_unwrapped = args_unwrapped[0]  # type: ignore[assignment]
+        primals_unwrapped_descs = args_descs_unwrapped[0]  # type: ignore[assignment]
+        fn_to_trace = joint_fn  # type: ignore[assignment]
     else:
+        primals_wrapped: list[FxValue] = typing.cast(list[FxValue], args)
+        primals_wrapped_descs: list[AOTInput] = typing.cast(list[AOTInput], args_descs)
+
         args_unwrapped, args_descs_unwrapped = unwrap_tensor_subclasses(  # type: ignore[assignment]
-            args,  # type: ignore[arg-type]
-            args_descs,  # type: ignore[arg-type]
+            primals_wrapped,
+            primals_wrapped_descs,
             append_symints=True,
         )
         remapped_static_indices = remap_unwrapped_subclass_arg_indices(
-            args,  # type: ignore[arg-type]
+            primals_wrapped,
             meta.static_input_indices,  # type: ignore[arg-type]
         )
 
-    if is_joint_structure:
-        primals_unwrapped = args_unwrapped[0]  # type: ignore[assignment]
-        primals_unwrapped_descs = args_descs_unwrapped[0]  # type: ignore[assignment]
-        fn_to_trace = joint_fn  # type: ignore[assignment]
-    else:
         primals_unwrapped = args_unwrapped  # type: ignore[assignment]
         primals_unwrapped_descs = args_descs_unwrapped  # type: ignore[assignment]
         fn_to_trace = fw_fn  # type: ignore[assignment]
@@ -1439,7 +1474,6 @@ def inner_fw_only(*args: Any) -> Any:
         flat_args_descs=primals_unwrapped_descs,
         static_input_indices=remapped_static_indices,
         keep_input_mutations=meta.keep_input_mutations,
-        is_train=meta.is_train,
         # pyrefly: ignore [not-iterable]
     )(*primals_unwrapped)
 
diff --git a/torch/_functorch/_aot_autograd/graph_compile.py b/torch/_functorch/_aot_autograd/graph_compile.py
index e55caaddf603d..10a6a42918448 100644
--- a/torch/_functorch/_aot_autograd/graph_compile.py
+++ b/torch/_functorch/_aot_autograd/graph_compile.py
@@ -13,21 +13,13 @@
 import itertools
 import logging
 import operator
+import threading
 import time
 import traceback
 from collections import defaultdict
 from collections.abc import Callable, Generator
-from contextlib import nullcontext
-from typing import Any, Optional, TYPE_CHECKING, Union
-
-from torch._library.fake_class_registry import FakeScriptObject
-
-
-if TYPE_CHECKING:
-    from collections.abc import Sequence
-
-import threading
-from contextlib import contextmanager
+from contextlib import contextmanager, nullcontext
+from typing import Any
 
 import torch
 import torch.utils._pytree as pytree
@@ -40,7 +32,10 @@
     lazy_format_graph_code,
 )
 from torch._guards import CompileContext, TracingContext
+from torch._library.fake_class_registry import FakeScriptObject
+from torch._library.opaque_object import is_opaque_value
 from torch._logging import getArtifactLogger, trace_structured
+from torch._opaque_base import OpaqueBase
 from torch._subclasses import FakeTensor
 from torch._subclasses.meta_utils import is_sparse_any
 from torch.fx.experimental._backward_state import BackwardState
@@ -66,6 +61,7 @@
 from .runtime_wrappers import (
     AOTDedupeWrapper,
     AOTDispatchAutograd,
+    AOTDispatchAutogradCompileSpec,
     AOTDispatchSubclassWrapper,
     AOTSyntheticBaseWrapper,
     AutogradLazyBackwardCompileInfo,
@@ -101,6 +97,22 @@
 )
 
 
+def is_opaque_node(node: Any) -> bool:
+    """Check if a node contains an opaque or non-tensor value (e.g., ProcessGroup)."""
+    from torch._library.fake_class_registry import FakeScriptObject
+
+    if not isinstance(node, torch.fx.Node):
+        return False
+    if "val" not in getattr(node, "meta", {}):
+        return False
+    val = node.meta["val"]
+    if is_opaque_value(val):
+        return True
+    if isinstance(val, (torch.ScriptObject, FakeScriptObject)):
+        return True
+    return False
+
+
 _thread_local = threading.local()
 
 
@@ -207,7 +219,7 @@ def orig_flat_fn2(*args: FxValue) -> tuple[list[FxValue], list[AOTOutput]]:
     # NB: This is currently only used for backwards, where fwd/bwd
     # deterministic TLS can be different
     aot_state.fw_metadata.deterministic = torch.are_deterministic_algorithms_enabled()
-    updated_flat_args: Union[list[Any], tuple[list[Any], list[Any]]]
+    updated_flat_args: list[Any] | tuple[list[Any], list[Any]]
 
     with maybe_skip_decompose(aot_config):
         # if config.selective_decompose, skip decomposition and apply selective_decompose
@@ -319,7 +331,7 @@ def sanitize_aot_config(input: AOTConfig) -> AOTConfig:
 
 
 def _get_inner_meta(
-    maybe_subclass_meta: Optional[SubclassMeta],
+    maybe_subclass_meta: SubclassMeta | None,
     fw_metadata: ViewAndMutationMeta,
 ) -> ViewAndMutationMeta:
     """
@@ -348,9 +360,9 @@ def aot_stage2_compile(
     # pyrefly: ignore [implicit-any]
     fw_compiler: Callable,
     # pyrefly: ignore [implicit-any]
-    bw_compiler: Optional[Callable] = None,
+    bw_compiler: Callable | None = None,
     # pyrefly: ignore [implicit-any]
-    inference_compiler: Optional[Callable] = None,
+    inference_compiler: Callable | None = None,
 ) -> DispatchReturn:
     if bw_compiler is None:
         bw_compiler = fw_compiler
@@ -371,7 +383,7 @@ def aot_stage2_compile(
 def _log_inference_graph(
     fw_module: torch.fx.GraphModule,
     aot_config: AOTConfig,
-) -> Optional[str]:
+) -> str | None:
     """
     Log the inference graph to the structured logger.
     Return a str representation of the graph.
@@ -404,7 +416,7 @@ def _log_inference_graph(
 def _aot_stage2b_inference_compile(
     fw_module: torch.fx.GraphModule,
     updated_flat_args: list[Any],
-    maybe_subclass_meta: Optional[SubclassMeta],
+    maybe_subclass_meta: SubclassMeta | None,
     fw_metadata: ViewAndMutationMeta,
     aot_config: AOTConfig,
     # pyrefly: ignore [implicit-any]
@@ -472,11 +484,11 @@ def aot_stage2_inference(
 def _cache_inference_info(
     aot_config: AOTConfig,
     fw_metadata: ViewAndMutationMeta,
-    maybe_subclass_meta: Optional[SubclassMeta],
+    maybe_subclass_meta: SubclassMeta | None,
     compiled_fw: Callable[..., Any],
-    aot_forward_graph_str: Optional[str],
+    aot_forward_graph_str: str | None,
     wrappers: list[CompilerWrapper],
-) -> Optional[GenericAOTAutogradResult[Any, Any]]:
+) -> GenericAOTAutogradResult[Any, Any] | None:
     make_runtime_safe(fw_metadata, maybe_subclass_meta)
 
     cache_info = aot_config.cache_info
@@ -487,7 +499,7 @@ def should_save_cache() -> bool:
         else:
             return hasattr(compiled_fw, "_fx_graph_cache_key")
 
-    entry: Optional[GenericAOTAutogradResult[Any, Any]] = None
+    entry: GenericAOTAutogradResult[Any, Any] | None = None
     if cache_info is not None and should_save_cache():
         time_taken_ns = time.time_ns() - cache_info.start_time_ns
         guards_expr = AOTAutogradCache.generate_guards_expression(cache_info)
@@ -524,7 +536,7 @@ def _aot_stage2c_make_inference_function(
     fw_metadata: ViewAndMutationMeta,
     compiled_fw: Callable[..., Any],
     wrappers: list[CompilerWrapper],
-    entry: Optional[GenericAOTAutogradResult[Any, Any]],
+    entry: GenericAOTAutogradResult[Any, Any] | None,
 ) -> DispatchReturn:
     if entry is not None:
         compiled_fw = SerializableCompiledFunction(compiled_fw, lambda: entry)
@@ -547,9 +559,9 @@ def _aot_stage2c_make_inference_function(
 
 
 def collect_fw_donated_buffer_idxs(
-    fw_ins: list[Optional[FakeTensor]],
-    user_fw_outs: list[Optional[FakeTensor]],
-    bw_outs: list[Optional[FakeTensor]],
+    fw_ins: list[FakeTensor | None],
+    user_fw_outs: list[FakeTensor | None],
+    bw_outs: list[FakeTensor | None],
     saved_tensors: list[FakeTensor | None],
 ) -> list[int]:
     """
@@ -570,6 +582,7 @@ def collect_fw_donated_buffer_idxs(
         t = saved_tensors[i]
         if (
             t is not None
+            and isinstance(t, FakeTensor)
             and not is_sparse_any(t)
             and StorageWeakRef(t.untyped_storage()) not in storage_refs
         ):
@@ -646,13 +659,13 @@ class InvokeSubgraphHopGraphs:
 
     # To avoid re-partitioning subgraphs
     partitioning_done: bool = False
-    old_num_fw_outputs: Optional[int] = None
-    old_num_fw_inputs: Optional[int] = None
+    old_num_fw_outputs: int | None = None
+    old_num_fw_inputs: int | None = None
 
-    new_fw_hop_gm: Optional[torch.fx.GraphModule] = None
-    new_bw_hop_gm: Optional[torch.fx.GraphModule] = None
-    new_num_sym_nodes: Optional[int] = None
-    new_num_saved_nodes: Optional[int] = None
+    new_fw_hop_gm: torch.fx.GraphModule | None = None
+    new_bw_hop_gm: torch.fx.GraphModule | None = None
+    new_num_sym_nodes: int | None = None
+    new_num_saved_nodes: int | None = None
 
 
 def prepare_for_partitioner(
@@ -1098,7 +1111,7 @@ def maybe_log_graph(
     graph_name: str,
     aot_config: AOTConfig,
     structured_log_prefix_fn: Callable[[], str],
-    out_structured_logs: Optional[list[str]] = None,
+    out_structured_logs: list[str] | None = None,
 ) -> None:
     if not aot_config.enable_log:
         return
@@ -1419,7 +1432,7 @@ def _get_extra_info() -> dict[str, Any]:
 
         def find_saved_in_bw_inputs(
             bw_inputs: list[torch.fx.Node],
-        ) -> Optional[torch.fx.Node]:
+        ) -> torch.fx.Node | None:
             for n in bw_inputs:
                 if n.name == saved.name:  # type: ignore[union-attr]
                     return n
@@ -1592,7 +1605,7 @@ def _log_structured_logs() -> None:
 def _log_joint_graph(
     fx_g: torch.fx.GraphModule,
     aot_config: AOTConfig,
-) -> Optional[str]:
+) -> str | None:
     """
     Log the joint graph to the structured logger.
     Return a str representation of the graph.
@@ -1626,10 +1639,10 @@ def _log_joint_graph(
 def _log_fw_bw_graphs(
     fw_module: torch.fx.GraphModule,
     bw_module: torch.fx.GraphModule,
-    maybe_subclass_meta: Optional[SubclassMeta],
+    maybe_subclass_meta: SubclassMeta | None,
     fw_metadata: ViewAndMutationMeta,
     aot_config: AOTConfig,
-) -> tuple[Optional[str], Optional[str]]:
+) -> tuple[str | None, str | None]:
     """
     Log the fw and bw graphs to the structured logger.
     Return str representations of the graphs.
@@ -1647,9 +1660,9 @@ def _log_fw_bw_graphs(
         )
         aot_graphs_log.info(
             "aot_config id: %s, fw_metadata=%s, inner_meta=%s",
-            str(aot_config.aot_id),
-            str(fw_metadata),
-            str(_get_inner_meta(maybe_subclass_meta, fw_metadata)),
+            aot_config.aot_id,
+            fw_metadata,
+            _get_inner_meta(maybe_subclass_meta, fw_metadata),
         )
 
         aot_graphs_log.info(
@@ -1716,10 +1729,269 @@ def _log_fw_bw_graphs(
     return fw_module_str, bw_module_str
 
 
+def _partition_joint_graph_into_fw_bw(
+    fx_g: torch.fx.GraphModule,
+    joint_inputs: list[Any] | tuple[list[Any], list[Any]],
+    inner_meta: ViewAndMutationMeta,
+    fw_metadata: ViewAndMutationMeta,
+    aot_config: AOTConfig,
+) -> tuple[torch.fx.GraphModule, torch.fx.GraphModule, int]:
+    # See Note: [Partitioner handling for Subclasses, Part 1]
+    # See Note: [Recomputing subclass mutation handling]
+    mutated_inp_runtime_indices = compute_inner_mutated_inp_indices_from_subclass_meta(
+        fw_metadata, inner_meta
+    )
+    num_tokens = len(fw_metadata.tokens)
+    num_inner_fwd_outputs = (
+        len(mutated_inp_runtime_indices)
+        + inner_meta.num_outputs
+        + inner_meta.num_intermediate_bases
+        + inner_meta.num_outputs_rng_offset
+        + num_tokens  # See Note [Side-Effectful Tokens in AOTAutograd]
+    )
+
+    fx_g = run_joint_graph_passes_on_hops(fx_g, joint_inputs, aot_config)
+
+    # apply joint_gm callback here
+    if callable(torch._functorch.config.joint_custom_pass):
+        # pyrefly: ignore [bad-assignment]
+        fx_g = torch._functorch.config.joint_custom_pass(fx_g, joint_inputs)
+
+    if aot_config.partition_fn is None:
+        raise AssertionError("aot_config.partition_fn must not be None")
+    fw_module, bw_module = aot_config.partition_fn(
+        fx_g,
+        joint_inputs,
+        num_fwd_outputs=num_inner_fwd_outputs,
+        static_lifetime_input_indices=fw_metadata.static_input_indices,
+    )
+
+    rng_states = [
+        n
+        for n in fw_module.graph.find_nodes(op="placeholder")
+        if "fwd_rng_state" in n.name
+    ]
+    fw_metadata.num_graphsafe_rng_states = len(rng_states)
+    if rng_states:
+        fw_metadata.graphsafe_rng_state_index = rng_states[0].meta["val"].device.index
+
+    return fw_module, bw_module, num_inner_fwd_outputs
+
+
+def _joint_inputs_for_forward(
+    joint_inputs: list[Any] | tuple[list[Any], list[Any]],
+) -> list[Any]:
+    return joint_inputs[0] if isinstance(joint_inputs, tuple) else joint_inputs
+
+
+def _maybe_unlift_partitioned_effect_tokens(
+    fw_module: torch.fx.GraphModule,
+    bw_module: torch.fx.GraphModule,
+    joint_inputs: list[Any] | tuple[list[Any], list[Any]],
+    fw_metadata: ViewAndMutationMeta,
+    aot_config: AOTConfig,
+    num_inner_fwd_outputs: int,
+) -> tuple[int, list[Any] | tuple[list[Any], list[Any]]]:
+    num_tokens = len(fw_metadata.tokens)
+
+    # See Note [Side-Effectful Tokens in AOTAutograd]
+    if config.unlift_effect_tokens and (
+        num_tokens > 0 or fw_metadata.num_backward_tokens > 0
+    ):
+        unlift_tokens(fw_module, fw_metadata, aot_config, bw_module)
+        num_inner_fwd_outputs -= num_tokens
+        if isinstance(joint_inputs, tuple):
+            joint_inputs = (
+                _joint_inputs_for_forward(joint_inputs)[num_tokens:],
+                joint_inputs[1],
+            )
+        else:
+            joint_inputs = joint_inputs[num_tokens:]
+
+    return num_inner_fwd_outputs, joint_inputs
+
+
+def _categorize_saved_tensors_for_backward(
+    fw_module: torch.fx.GraphModule,
+    bw_module: torch.fx.GraphModule,
+    inner_meta: ViewAndMutationMeta,
+    fw_metadata: ViewAndMutationMeta,
+    num_inner_fwd_outputs: int,
+) -> tuple[int, int]:
+    fw_outs = next(iter(fw_module.graph.find_nodes(op="output"))).args[0]
+    # we only need to bookkeep the symints that are saved for bw, not any symints
+    # the user forward might have returned in its own output
+    fw_outs_saved_for_bw = fw_outs[num_inner_fwd_outputs:]
+    num_fw_outs_saved_for_bw = len(fw_outs_saved_for_bw)
+
+    num_symints_saved_for_bw = 0
+    num_opaque_objects_saved_for_bw = 0
+    for idx, node in enumerate(fw_outs_saved_for_bw):
+        if is_sym_node(node):
+            num_symints_saved_for_bw += 1
+        elif is_opaque_node(node):
+            num_opaque_objects_saved_for_bw += 1
+        elif isinstance(node, torch.fx.Node) and "val" in getattr(node, "meta", {}):
+            if isinstance(node.meta["val"], FakeTensor):
+                # record dynamic tensor activations
+                dynamic_dims: set[int] = {
+                    dim
+                    for dim, size in enumerate(node.meta["val"].shape)
+                    if not isinstance(size, int)
+                }
+                if dynamic_dims:
+                    fw_metadata.dynamic_saved_tensors_idxs[idx] = dynamic_dims
+            elif isinstance(node.meta["val"], (FakeScriptObject, OpaqueBase)):
+                num_opaque_objects_saved_for_bw += 1
+
+    fw_metadata.num_symints_saved_for_bw = num_symints_saved_for_bw
+    fw_metadata.num_opaque_objects_saved_for_bw = num_opaque_objects_saved_for_bw
+    inner_meta.num_symints_saved_for_bw = num_symints_saved_for_bw
+    inner_meta.num_opaque_objects_saved_for_bw = num_opaque_objects_saved_for_bw
+
+    # See Note [Activations with no version counter checks in eager]
+    # Count tensors saved with no version counter check.
+    # These are tensors that were stashed on ctx (e.g., ctx.x = x) rather than
+    # via save_for_backward in an autograd.Function.
+    # The partitioner sorts these to be at the end of saved_values.
+    num_tensors_saved_with_no_vc_check = sum(
+        1
+        for node in fw_outs_saved_for_bw
+        if isinstance(node, torch.fx.Node)
+        and node.meta.get("saved_tensor_with_no_vc_check", False)
+    )
+    fw_metadata.num_tensors_saved_with_no_vc_check = num_tensors_saved_with_no_vc_check
+    inner_meta.num_tensors_saved_with_no_vc_check = num_tensors_saved_with_no_vc_check
+
+    if torch._functorch.config.donated_buffer:
+        fw_metadata.bw_donated_idxs = collect_bw_donated_buffer_idxs(
+            fw_module,
+            bw_module,
+            inner_meta,
+        )
+        inner_meta.bw_donated_idxs = fw_metadata.bw_donated_idxs
+
+    return num_fw_outs_saved_for_bw, num_symints_saved_for_bw
+
+
+# Note [Detaching inputs that never need gradients]
+# See https://github.com/pytorch/pytorch/issues/97745
+# Suppose we have a function like this that we want to compile:
+#
+# def f(x, y):
+#     return torch.mul(x, y.detach())
+#
+# What gradients should we compute for x and y?
+# By default, AOTAutograd will compute a gradient for **every** input that requires gradients,
+# and so we'll compute:
+#    x_grad_input = y
+#    y_grad_input = None
+# Does this preserve the semantics of eager mode?
+# Unfortunately, no.
+# Doing the above will cause autograd to **continue** to backprop the autograd tape
+# that was generated from constructing y.
+#
+# This is **different** from what would have happened in eager mode.
+# In eager mode, if we backprop through the output of this function, autograd will only traverse
+# the bit of the autograd tape corresponding to "x".
+# In particular, if a user had previously backpropped through y's autograd tape,
+# And then they try to backprop through the output of the above function,
+# then we'll hit the dreaded "Trying to backward through the graph a second time" error.
+#
+# You might think: If autograd sees that a gradient is None, shouldn't it stop early,
+# instead of continuing the backprop through the ancestors of that node in the graph?
+#
+# Autograd has two passes:
+# (1) a first pass that traverses the autograd graph and figures out which nodes need to be executed
+# (2) a second pass that actually goes ahead and executes each node when it becomes ready,
+#     propagating gradients
+# By the time we're executing a node and we see that it produces a None, the set of nodes to execute
+# is already locked-in.
+#
+# The fix: instead, we can recognize statically that the graph we're compiling will never contribute
+# gradients to y, and prevent autograd from trying to traverse y's autograd tape at all.
+# We can do this by manually detach'ing y before sending it through the `CompiledFunction`.
+#
+# Note that this solution is not bulletproof.
+# It's possible to construct a case where eager may or may not have have tried to autograd through y,
+# depending on the actual grad_outputs that were passed in during the backward.
+# There is no easy fix for this: the simplest fix would be to run with `retain_graph=True`,
+# allowing autograd to reuse the graph.
+#
+# An example of this case is:
+# def f(x):
+#     return x.detach() * 2, x * 3
+# If we were to only backprop through outs[0], in eager, we would stop
+# If we backward only on the first output, we shouldn't send a grad through x.
+# But the custom autograd function doesn't know that: it will materialize zero grads for x * 3
+# and we will end up with a zero grad at x.
+# If we later backprop through the second output, this will also require backprop'ing through x.
+# Meaning we'll need to use `retain_graph=True` to be able to backprop through x the second time.
+def _compute_indices_of_inps_to_detach(
+    bw_module: torch.fx.GraphModule,
+    maybe_subclass_meta: SubclassMeta | None,
+    inner_meta: ViewAndMutationMeta,
+    fw_metadata: ViewAndMutationMeta,
+) -> list[int]:
+    # TODO: we should apply the below "detach inputs if their gradients are statically known to be None"
+    # optimization even if we have subclass inputs/outputs (we do not handle this today).
+    # Computing which our our inputs get None gradients is a bit more complicated,
+    # if any of our inputs are subclasses. Why?
+    # (a) we need to make sure that we call .detach() on the input subclasses, since autograd sees subclasses.
+    # (b) The grad_outputs that we AOT computed in our backward graph are the desugared tensor tensors,
+    #     so we need to figure out which subclass fw inputs they map to.
+    if maybe_subclass_meta is not None:
+        return []
+
+    indices_of_inps_to_detach: list[int] = []
+
+    # reversed() since we expect output at end of graph
+    bw_output = next(reversed(bw_module.graph.find_nodes(op="output")))
+    bw_outs = bw_output.args[0]
+
+    num_backward_tokens = inner_meta.num_backward_tokens
+    expected_bw_outs = (
+        len(fw_metadata.input_info)
+        + inner_meta.num_outputs_rng_offset
+        + num_backward_tokens
+    )
+    if len(bw_outs) != expected_bw_outs:
+        raise AssertionError(
+            f"expected len(bw_outs) == {expected_bw_outs}, got {len(bw_outs)}"
+        )
+
+    bw_outs_no_rng_no_tokens = bw_outs
+    if (inner_meta.num_outputs_rng_offset + num_backward_tokens) > 0:
+        bw_outs_no_rng_no_tokens = bw_outs[
+            : -(inner_meta.num_outputs_rng_offset + num_backward_tokens)
+        ]
+    if len(bw_outs_no_rng_no_tokens) != len(fw_metadata.input_info):
+        raise AssertionError(
+            f"expected len(bw_outs_no_rng_no_tokens) == {len(fw_metadata.input_info)}, "
+            f"got {len(bw_outs_no_rng_no_tokens)}"
+        )
+
+    for i, bw_out in enumerate(bw_outs_no_rng_no_tokens):
+        # If our input experiences a metadata mutation inside the graph (e.g. set_()),
+        # we *must* not detach, otherwise it will be the detach'd input that gets the metadata mutation
+        metadata_mutation_in_graph = (
+            fw_metadata.input_info[i].mutation_type == MutationType.MUTATED_IN_GRAPH
+            and fw_metadata.input_info[i].mutates_storage_metadata
+        )
+        is_non_leaf = (
+            fw_metadata.input_info[i].requires_grad
+            and not fw_metadata.input_info[i].is_leaf
+        )
+        if bw_out is None and not metadata_mutation_in_graph and is_non_leaf:
+            indices_of_inps_to_detach.append(i)
+
+    return indices_of_inps_to_detach
+
+
 def _aot_stage2a_partition(
     fx_g: torch.fx.GraphModule,
-    joint_inputs: Union[list[Any], tuple[list[Any], list[Any]]],
-    maybe_subclass_meta: Optional[SubclassMeta],
+    joint_inputs: list[Any] | tuple[list[Any], list[Any]],
+    maybe_subclass_meta: SubclassMeta | None,
     fw_metadata: ViewAndMutationMeta,
     aot_config: AOTConfig,
 ) -> tuple[torch.fx.GraphModule, torch.fx.GraphModule, int, int, list[int], list[Any]]:
@@ -1736,60 +2008,25 @@ def _aot_stage2a_partition(
     with torch.no_grad():
         context = torch._C._DisableAutocast if disable_amp else nullcontext
         with context(), track_graph_compiling(aot_config, "joint"):
-            # See Note: [Partitioner handling for Subclasses, Part 1]
-            # See Note: [Recomputing subclass mutation handling]
-            mutated_inp_runtime_indices = (
-                compute_inner_mutated_inp_indices_from_subclass_meta(
-                    fw_metadata, inner_meta
+            fw_module, bw_module, num_inner_fwd_outputs = (
+                _partition_joint_graph_into_fw_bw(
+                    fx_g,
+                    joint_inputs,
+                    inner_meta,
+                    fw_metadata,
+                    aot_config,
                 )
             )
-            num_tokens = len(fw_metadata.tokens)
-            num_mutated_inp_runtime_indices = len(mutated_inp_runtime_indices)
-            num_inner_fwd_outputs = (
-                num_mutated_inp_runtime_indices
-                + inner_meta.num_outputs
-                + inner_meta.num_intermediate_bases
-                + inner_meta.num_outputs_rng_offset
-                + num_tokens  # See Note [Side-Effectful Tokens in AOTAutograd]
-            )
-            fx_g = run_joint_graph_passes_on_hops(fx_g, joint_inputs, aot_config)
-
-            # apply joint_gm callback here
-            if callable(torch._functorch.config.joint_custom_pass):
-                # pyrefly: ignore [bad-assignment]
-                fx_g = torch._functorch.config.joint_custom_pass(fx_g, joint_inputs)
-
-            static_lifetime_input_indices = fw_metadata.static_input_indices
-            if aot_config.partition_fn is None:
-                raise AssertionError("aot_config.partition_fn must not be None")
-            fw_module, bw_module = aot_config.partition_fn(
-                fx_g,
-                joint_inputs,
-                num_fwd_outputs=num_inner_fwd_outputs,
-                static_lifetime_input_indices=static_lifetime_input_indices,
-            )
-            rng_states = [
-                n
-                for n in fw_module.graph.find_nodes(op="placeholder")
-                if "fwd_rng_state" in n.name
-            ]
-            fw_metadata.num_graphsafe_rng_states = len(rng_states)
-            if rng_states:
-                fw_metadata.graphsafe_rng_state_index = (
-                    rng_states[0].meta["val"].device.index
-                )
-
-            # See Note [Side-Effectful Tokens in AOTAutograd]
-            if config.unlift_effect_tokens and (
-                num_tokens > 0 or fw_metadata.num_backward_tokens > 0
-            ):
-                unlift_tokens(fw_module, fw_metadata, aot_config, bw_module)
-
-                num_inner_fwd_outputs -= num_tokens
-                joint_inputs = (
-                    joint_inputs[0][num_tokens:],
-                    joint_inputs[1],
+            num_inner_fwd_outputs, joint_inputs = (
+                _maybe_unlift_partitioned_effect_tokens(
+                    fw_module,
+                    bw_module,
+                    joint_inputs,
+                    fw_metadata,
+                    aot_config,
+                    num_inner_fwd_outputs,
                 )
+            )
 
             maybe_inline_graph_saved_tensors_hooks(
                 fw_module,
@@ -1799,170 +2036,22 @@ def _aot_stage2a_partition(
                 aot_config,
                 fw_metadata.static_input_indices,
             )
-            static_lifetime_input_indices = fw_metadata.static_input_indices
-
-            fw_outs = next(iter(fw_module.graph.find_nodes(op="output"))).args[0]
-            # we only need to bookkeep the symints that are saved for bw, not any symints
-            # the user forward might have returned in its own output
-            fw_outs_saved_for_bw = fw_outs[num_inner_fwd_outputs:]
-            num_fw_outs_saved_for_bw = len(fw_outs_saved_for_bw)
-            symint_outs_saved_for_bw = []
-            opaque_outs_saved_for_bw = []
-            for idx, node in enumerate(fw_outs_saved_for_bw):
-                if is_sym_node(node):
-                    symint_outs_saved_for_bw.append(node)
-                elif isinstance(node, torch.fx.Node) and "val" in getattr(
-                    node, "meta", {}
-                ):
-                    if isinstance(node.meta["val"], FakeTensor):
-                        # record dynamic tensor activations
-                        dynamic_dims: set[int] = {
-                            dim
-                            for dim, size in enumerate(node.meta["val"].shape)
-                            if not isinstance(size, int)
-                        }
-                        if dynamic_dims:
-                            fw_metadata.dynamic_saved_tensors_idxs[idx] = dynamic_dims
-                    elif isinstance(node.meta["val"], FakeScriptObject):
-                        opaque_outs_saved_for_bw.append(node)
-
-            num_symints_saved_for_bw = len(symint_outs_saved_for_bw)
-            num_opaque_objects_saved_for_bw = len(opaque_outs_saved_for_bw)
-            fw_metadata.num_symints_saved_for_bw = num_symints_saved_for_bw
-            fw_metadata.num_opaque_objects_saved_for_bw = (
-                num_opaque_objects_saved_for_bw
-            )
-            inner_meta.num_symints_saved_for_bw = num_symints_saved_for_bw
-            inner_meta.num_opaque_objects_saved_for_bw = num_opaque_objects_saved_for_bw
-
-            # See Note [Activations with no version counter checks in eager]
-            # Count tensors saved with no version counter check.
-            # These are tensors that were stashed on ctx (e.g., ctx.x = x) rather than
-            # via save_for_backward in an autograd.Function.
-            # The partitioner sorts these to be at the end of saved_values.
-            num_tensors_saved_with_no_vc_check = 0
-            for node in fw_outs_saved_for_bw:
-                if isinstance(node, torch.fx.Node) and node.meta.get(
-                    "saved_tensor_with_no_vc_check", False
-                ):
-                    num_tensors_saved_with_no_vc_check += 1
-            fw_metadata.num_tensors_saved_with_no_vc_check = (
-                num_tensors_saved_with_no_vc_check
-            )
-            inner_meta.num_tensors_saved_with_no_vc_check = (
-                num_tensors_saved_with_no_vc_check
-            )
-
-            if torch._functorch.config.donated_buffer:
-                fw_metadata.bw_donated_idxs = collect_bw_donated_buffer_idxs(
+            num_fw_outs_saved_for_bw, num_symints_saved_for_bw = (
+                _categorize_saved_tensors_for_backward(
                     fw_module,
                     bw_module,
                     inner_meta,
+                    fw_metadata,
+                    num_inner_fwd_outputs,
                 )
-                inner_meta.bw_donated_idxs = fw_metadata.bw_donated_idxs
-
-        # Note [Detaching inputs that never need gradients]
-        # See https://github.com/pytorch/pytorch/issues/97745
-        # Suppose we have a function like this that we want to compile:
-        #
-        # def f(x, y):
-        #     return torch.mul(x, y.detach())
-        #
-        # What gradients should we compute for x and y?
-        # By default, AOTAutograd will compute a gradient for **every** input that requires gradients,
-        # and so we'll compute:
-        #    x_grad_input = y
-        #    y_grad_input = None
-        # Does this preserve the semantics of eager mode?
-        # Unfortunately, no.
-        # Doing the above will cause autograd to **continue** to backprop the autograd tape
-        # that was generated from constructing y.
-        #
-        # This is **different** from what would have happened in eager mode.
-        # In eager mode, if we backprop through the output of this function, autograd will only traverse
-        # the bit of the autograd tape corresponding to "x".
-        # In particular, if a user had previously backpropped through y's autograd tape,
-        # And then they try to backprop through the output of the above function,
-        # then we'll hit the dreaded "Trying to backward through the graph a second time" error.
-        #
-        # You might think: If autograd sees that a gradient is None, shouldn't it stop early,
-        # instead of continuing the backprop through the ancestors of that node in the graph?
-        #
-        # Autograd has two passes:
-        # (1) a first pass that traverses the autograd graph and figures out which nodes need to be executed
-        # (2) a second pass that actually goes ahead and executes each node when it becomes ready,
-        #     propagating gradients
-        # By the time we're executing a node and we see that it produces a None, the set of nodes to execute
-        # is already locked-in.
-        #
-        # The fix: instead, we can recognize statically that the graph we're compiling will never contribute
-        # gradients to y, and prevent autograd from trying to traverse y's autograd tape at all.
-        # We can do this by manually detach'ing y before sending it through the `CompiledFunction`.
-        #
-        # Note that this solution is not bulletproof.
-        # It's possible to construct a case where eager may or may not have have tried to autograd through y,
-        # depending on the actual grad_outputs that were passed in during the backward.
-        # There is no easy fix for this: the simplest fix would be to run with `retain_graph=True`,
-        # allowing autograd to reuse the graph.
-        #
-        # An example of this case is:
-        # def f(x):
-        #     return x.detach() * 2, x * 3
-        # If we were to only backprop through outs[0], in eager, we would stop
-        # If we backward only on the first output, we shouldn't send a grad through x.
-        # But the custom autograd function doesn't know that: it will materialize zero grads for x * 3
-        # and we will end up with a zero grad at x.
-        # If we later backprop through the second output, this will also require backprop'ing through x.
-        # Meaning we'll need to use `retain_graph=True` to be able to backprop through x the second time.
-        _indices_of_inps_to_detach: list[int] = []
-
-        # reversed() since we expect output at end of graph
-        bw_output = next(reversed(bw_module.graph.find_nodes(op="output")))
-        bw_outs: Sequence[torch.fx.Node] = bw_output.args[0]  # type: ignore[assignment]
-
-        # TODO: we should apply the below "detach inputs if their gradients are statically known to be None"
-        # optimization even if we have subclass inputs/outputs (we do not handle this today).
-        # Computing which our our inputs get None gradients is a bit more complicated,
-        # if any of our inputs are subclasses. Why?
-        # (a) we need to make sure that we call .detach() on the input subclasses, since autograd sees subclasses.
-        # (b) The grad_outputs that we AOT computed in our backward graph are the desugared tensor tensors,
-        #     so we need to figure out which subclass fw inputs they map to.
-        if maybe_subclass_meta is None:
-            num_backward_tokens: int = inner_meta.num_backward_tokens
-            expected_bw_outs = (
-                len(fw_metadata.input_info)
-                + inner_meta.num_outputs_rng_offset
-                + num_backward_tokens
             )
-            if len(bw_outs) != expected_bw_outs:
-                raise AssertionError(
-                    f"expected len(bw_outs) == {expected_bw_outs}, got {len(bw_outs)}"
-                )
-            bw_outs_no_rng_no_tokens = bw_outs
-            if (inner_meta.num_outputs_rng_offset + num_backward_tokens) > 0:
-                bw_outs_no_rng_no_tokens = bw_outs[
-                    : -(inner_meta.num_outputs_rng_offset + num_backward_tokens)
-                ]
-            if len(bw_outs_no_rng_no_tokens) != len(fw_metadata.input_info):
-                raise AssertionError(
-                    f"expected len(bw_outs_no_rng_no_tokens) == {len(fw_metadata.input_info)}, "
-                    f"got {len(bw_outs_no_rng_no_tokens)}"
-                )
 
-            for i, (bw_out) in enumerate(bw_outs_no_rng_no_tokens):
-                # If our input experiences a metadata mutation inside the graph (e.g. set_()),
-                # we *must* not detach, otherwise it will be the detach'd input that gets the metadata mutation
-                metadata_mutation_in_graph = (
-                    fw_metadata.input_info[i].mutation_type
-                    == MutationType.MUTATED_IN_GRAPH
-                    and fw_metadata.input_info[i].mutates_storage_metadata
-                )
-                is_non_leaf = (
-                    fw_metadata.input_info[i].requires_grad
-                    and not fw_metadata.input_info[i].is_leaf
-                )
-                if bw_out is None and not metadata_mutation_in_graph and is_non_leaf:
-                    _indices_of_inps_to_detach.append(i)
+        _indices_of_inps_to_detach = _compute_indices_of_inps_to_detach(
+            bw_module,
+            maybe_subclass_meta,
+            inner_meta,
+            fw_metadata,
+        )
 
     return (
         fw_module,
@@ -1970,19 +2059,19 @@ def _aot_stage2a_partition(
         num_fw_outs_saved_for_bw,
         num_symints_saved_for_bw,
         _indices_of_inps_to_detach,
-        joint_inputs[0],
+        _joint_inputs_for_forward(joint_inputs),
     )
 
 
 def _aot_stage2b_fw_compile(
     fw_module: torch.fx.GraphModule,
     adjusted_flat_args: list[Any],
-    maybe_subclass_meta: Optional[SubclassMeta],
+    maybe_subclass_meta: SubclassMeta | None,
     fw_metadata: ViewAndMutationMeta,
     num_fw_outs_saved_for_bw: int,
     aot_config: AOTConfig,
     # pyrefly: ignore [implicit-any]
-) -> tuple[Optional[list[Optional[tuple[int, ...]]]], Callable]:
+) -> tuple[list[tuple[int, ...] | None] | None, Callable]:
     return _aot_stage2b_compile_forward_or_inference(
         fw_module,
         adjusted_flat_args,
@@ -1996,13 +2085,13 @@ def _aot_stage2b_fw_compile(
 
 def _aot_stage2b_bw_compile(
     bw_module: torch.fx.GraphModule,
-    maybe_subclass_meta: Optional[SubclassMeta],
+    maybe_subclass_meta: SubclassMeta | None,
     fw_metadata: ViewAndMutationMeta,
-    fwd_output_strides: Optional[list[Optional[tuple[int, ...]]]],
+    fwd_output_strides: list[tuple[int, ...] | None] | None,
     num_symints_saved_for_bw: int,
     aot_config: AOTConfig,
     # pyrefly: ignore [implicit-any]
-) -> tuple[AutogradLazyBackwardCompileInfo, Optional[Callable]]:
+) -> tuple[AutogradLazyBackwardCompileInfo, Callable | None]:
     """
     Compile the backward graph. Returns:
     - the placeholder list for the backward graph
@@ -2249,13 +2338,13 @@ def _aot_stage2c_make_autograd_function(
     aot_config: AOTConfig,
     flat_args: list[Any],
     fw_metadata: ViewAndMutationMeta,
-    maybe_subclass_meta: Optional[SubclassMeta],
+    maybe_subclass_meta: SubclassMeta | None,
     wrappers: list[CompilerWrapper],
     compiled_fw_func: Callable[..., Any],
-    compiled_bw_func: Optional[Callable[..., Any]],
-    lazy_backward_info: Optional[AutogradLazyBackwardCompileInfo],
+    compiled_bw_func: Callable[..., Any] | None,
+    lazy_backward_info: AutogradLazyBackwardCompileInfo | None,
     try_save_cache_entry: Callable[..., Any],
-    entry: Optional[GenericAOTAutogradResult[Any, Any]],
+    entry: GenericAOTAutogradResult[Any, Any] | None,
     _indices_of_inps_to_detach: list[int],
     num_symints_saved_for_bw: int,
 ) -> DispatchReturn:
@@ -2268,25 +2357,26 @@ def _aot_stage2c_make_autograd_function(
         )
 
     disable_amp = torch._C._is_any_autocast_enabled()
-    compiled_fn = AOTDispatchAutograd.post_compile(
-        compiled_fw_func,
-        compiled_bw_func,
-        maybe_subclass_meta,
-        num_symints_saved_for_bw,
-        backward_state_indices,
-        disable_amp,
-        _indices_of_inps_to_detach,
-        lazy_backward_info,
-        aot_config,
+    compile_spec = AOTDispatchAutogradCompileSpec(
+        compiled_fw_func=compiled_fw_func,
+        compiled_bw_func=compiled_bw_func,
+        maybe_subclass_meta=maybe_subclass_meta,
+        num_symints_saved_for_bw=num_symints_saved_for_bw,
+        backward_state_indices=backward_state_indices,
+        disable_amp=disable_amp,
+        indices_of_inps_to_detach=_indices_of_inps_to_detach,
+        lazy_backward_info=lazy_backward_info,
+        aot_config=aot_config,
         fw_metadata=fw_metadata,
         try_save_cache_entry=try_save_cache_entry,
     )
+    compiled_fn = AOTDispatchAutograd.post_compile(compile_spec)
 
     if entry is not None:
         compiled_fn = SerializableCompiledFunction(compiled_fn, lambda: entry)
 
     if config.debug_assert:
-        flat_requires_grad: list[Optional[bool]] = [
+        flat_requires_grad: list[bool | None] = [
             a.requires_grad if isinstance(a, Tensor) else None for a in flat_args
         ]
         compiled_fn = DebugAssertWrapper(
@@ -2306,19 +2396,19 @@ def _cache_autograd_info(
     aot_config: AOTConfig,
     flat_args: list[Any],
     compiled_fw_func: Callable[..., Any],
-    compiled_bw_func: Optional[Callable[..., Any]],
-    fw_module_str: Optional[str],
-    bw_module_str: Optional[str],
-    joint_graph_str: Optional[str],
+    compiled_bw_func: Callable[..., Any] | None,
+    fw_module_str: str | None,
+    bw_module_str: str | None,
+    joint_graph_str: str | None,
     wrappers: list[CompilerWrapper],
-    maybe_subclass_meta: Optional[SubclassMeta],
+    maybe_subclass_meta: SubclassMeta | None,
     fw_metadata: ViewAndMutationMeta,
     num_fw_outs_saved_for_bw: int,
     _indices_of_inps_to_detach: list[int],
     num_symints_saved_for_bw: int,
-    bw_module: Optional[torch.fx.GraphModule],
+    bw_module: torch.fx.GraphModule | None,
 ) -> tuple[
-    Optional[GenericAOTAutogradResult[Any, Any]],
+    GenericAOTAutogradResult[Any, Any] | None,
     Callable[..., Any],
 ]:
     backward_state_indices = [
@@ -2331,8 +2421,8 @@ def _cache_autograd_info(
 
     make_runtime_safe(fw_metadata, maybe_subclass_meta)
 
-    try_save_cache_entry: Optional[Callable[..., Any]] = None
-    entry: Optional[GenericAOTAutogradResult[Any, Any]] = None
+    try_save_cache_entry: Callable[..., Any] | None = None
+    entry: GenericAOTAutogradResult[Any, Any] | None = None
 
     if aot_config.cache_info is not None:
         forward_time_taken_ns = time.time_ns() - aot_config.cache_info.start_time_ns
@@ -2345,7 +2435,7 @@ def try_save_cache_entry(  # noqa: F811
             bw_module: torch.fx.GraphModule,
             _fw_metadata: ViewAndMutationMeta,
             aot_config: AOTConfig,
-        ) -> Optional[GenericAOTAutogradResult[Any, Any]]:
+        ) -> GenericAOTAutogradResult[Any, Any] | None:
             cache_info = aot_config.cache_info
 
             def should_save_cache() -> bool:
@@ -2367,9 +2457,9 @@ def should_save_cache() -> bool:
                 # update backward_time_taken_ns to be more inclusive
                 backward_time_taken_ns = getattr(compiled_bw_func, "_time_taken_ns", 0)
 
-                aot_forward_graph_str: Optional[str] = fw_module_str
-                aot_backward_graph_str: Optional[str] = bw_module_str
-                aot_joint_graph_str: Optional[str] = joint_graph_str
+                aot_forward_graph_str: str | None = fw_module_str
+                aot_backward_graph_str: str | None = bw_module_str
+                aot_joint_graph_str: str | None = joint_graph_str
                 guards_expr = AOTAutogradCache.generate_guards_expression(cache_info)
 
                 entry = AOTAutogradCache.make_entry(
@@ -2419,14 +2509,14 @@ def should_save_cache() -> bool:
 def _aot_stage2b_compile_forward_or_inference(
     fw_module: torch.fx.GraphModule,
     adjusted_flat_args: list[Any],
-    maybe_subclass_meta: Optional[SubclassMeta],
+    maybe_subclass_meta: SubclassMeta | None,
     fw_metadata: ViewAndMutationMeta,
     aot_config: AOTConfig,
     *,
     is_inference: bool,
-    num_fw_outs_saved_for_bw: Optional[int] = None,
+    num_fw_outs_saved_for_bw: int | None = None,
     # pyrefly: ignore [implicit-any]
-) -> tuple[Optional[list[Optional[tuple[int, ...]]]], Callable]:
+) -> tuple[list[tuple[int, ...] | None] | None, Callable]:
     """
     Compile the forward or inference graph. Returns:
     - the output strides of the forward graph
diff --git a/torch/_functorch/_aot_autograd/indexed_dict.py b/torch/_functorch/_aot_autograd/indexed_dict.py
index 39a06996c6e08..748818ff4b16c 100644
--- a/torch/_functorch/_aot_autograd/indexed_dict.py
+++ b/torch/_functorch/_aot_autograd/indexed_dict.py
@@ -1,5 +1,5 @@
 from collections.abc import Iterator, MutableMapping
-from typing import Generic, Optional, TypeVar
+from typing import Generic, TypeVar
 
 
 K = TypeVar("K")
@@ -39,14 +39,14 @@ def __iter__(self) -> Iterator[K]:
     def __contains__(self, key: object) -> bool:
         return key in self._dict
 
-    def next_key(self, key: K) -> Optional[K]:
+    def next_key(self, key: K) -> K | None:
         """Get the next key in insertion order. O(1)."""
         idx = self._key_to_index.get(key)
         if idx is not None and idx + 1 < len(self._keys):
             return self._keys[idx + 1]
         return None
 
-    def prev_key(self, key: K) -> Optional[K]:
+    def prev_key(self, key: K) -> K | None:
         """Get the previous key in insertion order. O(1)."""
         idx = self._key_to_index.get(key)
         if idx is not None and idx > 0:
diff --git a/torch/_functorch/_aot_autograd/input_output_analysis.py b/torch/_functorch/_aot_autograd/input_output_analysis.py
index 50e0870a78b7e..0c0721fcb73f5 100644
--- a/torch/_functorch/_aot_autograd/input_output_analysis.py
+++ b/torch/_functorch/_aot_autograd/input_output_analysis.py
@@ -96,6 +96,7 @@ def remove_dupe_metadata(
                 dynamic_dims=o.dynamic_dims,
                 base_idx=None if o.base_idx is None else add_dupe_map[o.base_idx],
                 requires_grad=o.requires_grad,
+                requires_grad_for_backward=o.requires_grad_for_backward,
                 view_meta_sequence=o.view_meta_sequence,
             )
             for o in m.output_info
@@ -108,7 +109,6 @@ def remove_dupe_metadata(
         subclass_inp_meta=[],
         subclass_fw_graph_out_meta=[],
         subclass_tangent_meta=subclass_tangent_meta,
-        is_train=m.is_train,
     )
 
 
@@ -224,7 +224,8 @@ def create_synthetic_base_metadata(
                 if not is_concrete_int(s)
             },
             base_idx=synthetic_base_info[outer_idx][0],  # type: ignore[index]
-            requires_grad=outer_args[outer_idx].requires_grad,
+            requires_grad=(requires_grad := outer_args[outer_idx].requires_grad),
+            requires_grad_for_backward=requires_grad,
         )
         for outer_idx in outer_aliased_arg_idx_with_metadata_mutations
     ]
@@ -253,6 +254,7 @@ def create_synthetic_base_metadata(
                 # Map the input idx pre-synthetic-bases to the new idx post-synthetic-bases
                 base_idx=new_base_idx,  # type: ignore[arg-type]
                 requires_grad=o.requires_grad,
+                requires_grad_for_backward=o.requires_grad_for_backward,
                 view_meta_sequence=o.view_meta_sequence,
             )
         )
@@ -305,7 +307,6 @@ def create_synthetic_base_metadata(
             subclass_inp_meta=[],
             subclass_fw_graph_out_meta=[],
             subclass_tangent_meta=subclass_tangent_meta,
-            is_train=m.is_train,
         ),
         outer_aliased_arg_idx_with_metadata_mutations,
     )
diff --git a/torch/_functorch/_aot_autograd/logging_utils.py b/torch/_functorch/_aot_autograd/logging_utils.py
index 550f99e6dd213..a5dda6c67dfa7 100644
--- a/torch/_functorch/_aot_autograd/logging_utils.py
+++ b/torch/_functorch/_aot_autograd/logging_utils.py
@@ -117,6 +117,7 @@ def prehook(grad_output: Any) -> None:
 
             fx_traceback.set_stack_trace(stack_)
             fx_traceback.set_grad_fn_seq_nr(seq_nr)
+            fx_traceback._mark_autograd_backward()
 
         return prehook
 
@@ -126,6 +127,7 @@ def get_posthook(
         def posthook(grad_input: Any, grad_output: Any) -> None:
             fx_traceback.set_stack_trace(special_stack_)
             fx_traceback.reset_grad_fn_seq_nr()
+            fx_traceback._reset_autograd_backward()
 
         return posthook
 
@@ -139,6 +141,16 @@ def posthook(grad_input: Any, grad_output: Any) -> None:
         node.register_hook(get_posthook(special_stack, node._sequence_nr()))
 
 
+def setup_stacktrace_preservation_hooks_from_tensors(outputs: Any) -> None:
+    roots = [
+        t.grad_fn
+        for t in (outputs if isinstance(outputs, (list, tuple)) else (outputs,))
+        if isinstance(t, torch.Tensor) and t.grad_fn is not None
+    ]
+    if roots:
+        setup_stacktrace_preservation_hooks(roots)
+
+
 def describe_input(i: int, aot_config: AOTConfig) -> str:
     if i < aot_config.num_params_buffers:
         return f"parameter/buffer {i}"
diff --git a/torch/_functorch/_aot_autograd/runtime_wrappers.py b/torch/_functorch/_aot_autograd/runtime_wrappers.py
index 1b3abce2f5b56..069ebf615db68 100644
--- a/torch/_functorch/_aot_autograd/runtime_wrappers.py
+++ b/torch/_functorch/_aot_autograd/runtime_wrappers.py
@@ -12,7 +12,6 @@
 import copy
 import functools
 import itertools
-import logging
 import pprint
 import typing
 import warnings
@@ -20,7 +19,7 @@
 from contextlib import AbstractContextManager, nullcontext
 from dataclasses import dataclass, field
 from functools import wraps
-from typing import Any, Optional, Union
+from typing import Any
 
 import torch
 import torch.fx as fx
@@ -37,9 +36,11 @@
     tracing,
     TracingContext,
 )
+from torch._library.fake_class_registry import FakeScriptObject
 from torch._library.opaque_object import is_opaque_type
 from torch._library.utils import is_builtin
 from torch._logging import getArtifactLogger
+from torch._opaque_base import OpaqueBase
 from torch._ops import OpOverload
 from torch._prims_common import CUDARngStateHelper
 from torch._subclasses import FakeTensor
@@ -79,6 +80,7 @@
     InputAliasInfo,
     MemoryFormatMeta,
     MutationType,
+    OpaqueMeta,
     OutputType,
     PlainTensorMeta,
     SubclassCreationMeta,
@@ -103,21 +105,31 @@
 )
 
 
+def _unwrap_tensor_subclasses_no_symints(
+    args: list[Any],
+) -> list[Any]:
+    return runtime_unwrap_tensor_subclasses(args, append_symints=False)  # type: ignore[arg-type]
+
+
 zip = strict_zip
 
 aot_graphs_log = getArtifactLogger(__name__, "aot_graphs")
 
 
+def _unwrap_no_symints(args: list[Any]) -> list[Any]:
+    return runtime_unwrap_tensor_subclasses(args, append_symints=False)
+
+
 def _describe_arg_for_logging(arg: object) -> str:
     from torch._library import opaque_object
 
     try:
-        is_dtensor = isinstance(arg, torch.distributed._tensor.DTensor)
+        is_dtensor = isinstance(arg, torch.distributed.tensor.DTensor)
     except AttributeError:
         is_dtensor = False
 
     if is_dtensor:
-        arg = typing.cast(torch.distributed._tensor.DTensor, arg)
+        arg = typing.cast(torch.distributed.tensor.DTensor, arg)
         mesh = arg.device_mesh
         return (
             f"DTensor(shape={arg.shape}, dtype={arg.dtype}, "
@@ -460,98 +472,26 @@ def __call__(self) -> AbstractContextManager[Any]:
         return nullcontext()
 
 
-def _create_runtime_wrapper(
-    compiled_fn: Callable[..., Any],
-    *,
-    runtime_metadata: ViewAndMutationMeta,
-    indices_of_inps_to_detach: list[int],
-    trace_joint: bool,
-    keep_input_mutations: bool,
-    disable_amp: bool,
-) -> Callable[..., Any]:
-    if not getattr(compiled_fn, "_boxed_call", False):
-        compiled_fn = make_boxed_func(compiled_fn)
-
-    # We only want to run debugmode on custom ops at the first invocation of
-    # runtime wrapper. For all subsequent uses, we should no-op for performance
-    # See: https://github.com/pytorch/pytorch/issues/165349
-    first_invocation_ctx = _FirstInvocationContext()
-
-    # Note [Inputs needed in runtime epilogue after list clearing]
-    # In Python functions, you can't free the input arguments of a function within the scope of that function. A workaround is to
-    # wrap the input arguments in a list, and clear the list from within the function.
-    # Here, this is implemented as `call_func_at_runtime_with_args(..., steal_args=True)`.
-    #
-    # This is needed for Compiled Autograd since some of the inputs (activations) should be freed early.
-    # However, we cannot blindly clear the entire list, because AOTAutograd may need access to some of the graph inputs
-    # **after** the compiled function has finished running. There are two main cases:
-    #   (1) Input mutations: If there are an input mutations that we must run outside of the graph, we need access to the input.
-    #   (2) Output aliasing: Outputs that aliases graph inputs generally must be regenerated outside of the `autograd.Function`,
-    #       and doing so requires us accessing the corresponding input after the compiled artifact has run.
-    epilogue_args_idx = []
-    epilogue_args_idx.extend(runtime_metadata.mutated_inp_runtime_indices)
-    for info in runtime_metadata.output_info:
-        if (
-            info.output_type == OutputType.alias_of_input
-            or info.output_type == OutputType.is_input
-        ):
-            if not isinstance(info.base_idx, int):
-                raise AssertionError(
-                    f"expected info.base_idx to be int, got {type(info.base_idx)}"
-                )
-            epilogue_args_idx.append(info.base_idx)
-
-    if config.unlift_effect_tokens:
-        if len(runtime_metadata.tokens) != 0:
-            raise AssertionError(
-                f"expected no tokens when unlift_effect_tokens is True, got {len(runtime_metadata.tokens)}"
-            )
-
-    if runtime_metadata.num_outputs_aliased > 0:
-        output_handlers = tuple(
-            make_output_handler(info, runtime_metadata, trace_joint)
-            for info in runtime_metadata.output_info
-        )
-
-    def record_runtime_wrapper_prologue_enter() -> AbstractContextManager[None] | None:
-        if (
-            torch.autograd.profiler._is_profiler_enabled
-            and dynamo_config.record_runtime_overhead
-        ):
-            cm = torch._C._profiler._RecordFunctionFast(
-                "AOTDispatcher Runtime Wrapper Prologue"
-            )
-            cm.__enter__()
-            return cm
-        return None
-
-    def record_runtime_wrapper_prologue_exit(
-        cm: AbstractContextManager[None] | None,
-    ) -> None:
-        if cm is not None:
-            cm.__exit__(None, None, None)
-
-    @simple_wraps(compiled_fn)
-    def runtime_wrapper(args: list[Any]) -> Any:
-        # Create context manager for profiler
-        cm = record_runtime_wrapper_prologue_enter()
-
-        # stash a ref to each input tensor we plan to use after the compiled function
-        orig_inputs = {i: args[i] for i in epilogue_args_idx}
+@dataclass
+class _RuntimeCompiledFnInvoker:
+    compiled_fn: Callable[..., Any]
+    indices_of_inps_to_detach: list[int]
+    trace_joint: bool
+    disable_amp: bool
+    first_invocation_ctx: _FirstInvocationContext = field(
+        default_factory=_FirstInvocationContext
+    )
 
-        if keep_input_mutations:
-            mutated_args = (
-                args[i]
-                for i in runtime_metadata.mutated_graph_handled_indices_seen_by_autograd
-            )
-            torch.autograd.graph.increment_version(mutated_args)
+    def __post_init__(self) -> None:
+        if not getattr(self.compiled_fn, "_boxed_call", False):
+            self.compiled_fn = make_boxed_func(self.compiled_fn)
 
-        # Enable _AnalyzeCustomOpInputOutputMode on first invocation to check aliasing constraints for custom ops
-        with first_invocation_ctx():
-            if trace_joint:
+    def run(self, args: list[Any], *, on_before_call: Callable[[], None]) -> list[Any]:
+        with self.first_invocation_ctx():
+            if self.trace_joint:
                 args_ = list(args)
                 # See Note [Detaching inputs that never need gradients]
-                for idx in indices_of_inps_to_detach:
+                for idx in self.indices_of_inps_to_detach:
                     if isinstance(args_[idx], torch.Tensor):
                         args_[idx] = args_[idx].detach()
 
@@ -562,157 +502,284 @@ def runtime_wrapper(args: list[Any]) -> Any:
                     torch.autograd._force_original_view_tracking(True),
                     torch.enable_grad(),
                 ):
-                    record_runtime_wrapper_prologue_exit(cm)
-                    all_outs = call_func_at_runtime_with_args(
-                        compiled_fn, args_, disable_amp=disable_amp, steal_args=True
+                    on_before_call()
+                    return call_func_at_runtime_with_args(
+                        self.compiled_fn,
+                        args_,
+                        disable_amp=self.disable_amp,
+                        steal_args=True,
                     )
-            else:
-                # When we have an inference graph, we run with grad disabled.
-                # It's possible to get an inference graph with inputs that require grad,
-                # in which case we want to make sure autograd is disabled
-                # (since e.g., inductor will generate aten.addmm.out calls which autograd will complain on)
-                # NOTE: We use _set_grad_enabled directly to reduce runtime overhead
-                grad_enabled = torch.is_grad_enabled()
-                try:
-                    if grad_enabled:
-                        torch._C._set_grad_enabled(False)
-                    record_runtime_wrapper_prologue_exit(cm)
-                    all_outs = call_func_at_runtime_with_args(
-                        compiled_fn, args, disable_amp=disable_amp, steal_args=True
+
+            # When we have an inference graph, we run with grad disabled.
+            # It's possible to get an inference graph with inputs that require grad,
+            # in which case we want to make sure autograd is disabled
+            # (since e.g., inductor will generate aten.addmm.out calls which autograd will complain on)
+            # NOTE: We use _set_grad_enabled directly to reduce runtime overhead
+            grad_enabled = torch.is_grad_enabled()
+            try:
+                if grad_enabled:
+                    torch._C._set_grad_enabled(False)
+                on_before_call()
+                return call_func_at_runtime_with_args(
+                    self.compiled_fn,
+                    args,
+                    disable_amp=self.disable_amp,
+                    steal_args=True,
+                )
+            finally:
+                if grad_enabled:
+                    torch._C._set_grad_enabled(True)
+
+
+@dataclass
+class _RuntimeForwardEpilogue:
+    runtime_metadata: ViewAndMutationMeta
+    trace_joint: bool
+    keep_input_mutations: bool
+    epilogue_args_idx: tuple[int, ...] = field(init=False)
+    output_handlers: tuple[Any, ...] = field(init=False)
+
+    def __post_init__(self) -> None:
+        epilogue_args_idx = list(self.runtime_metadata.mutated_inp_runtime_indices)
+        for info in self.runtime_metadata.output_info:
+            if (
+                info.output_type == OutputType.alias_of_input
+                or info.output_type == OutputType.is_input
+            ):
+                if not isinstance(info.base_idx, int):
+                    raise AssertionError(
+                        f"expected info.base_idx to be int, got {type(info.base_idx)}"
                     )
-                finally:
-                    if grad_enabled:
-                        torch._C._set_grad_enabled(True)
+                epilogue_args_idx.append(info.base_idx)
+        self.epilogue_args_idx = tuple(epilogue_args_idx)
 
-        del args
+        if config.unlift_effect_tokens:
+            if len(self.runtime_metadata.tokens) != 0:
+                raise AssertionError(
+                    "expected no tokens when unlift_effect_tokens is True, "
+                    f"got {len(self.runtime_metadata.tokens)}"
+                )
+
+        if self.runtime_metadata.num_outputs_aliased > 0:
+            self.output_handlers = tuple(
+                make_output_handler(info, self.runtime_metadata, self.trace_joint)
+                for info in self.runtime_metadata.output_info
+            )
+        else:
+            self.output_handlers = ()
+
+    def capture_orig_inputs(self, args: list[Any]) -> dict[int, Any]:
+        return {i: args[i] for i in self.epilogue_args_idx}
+
+    def increment_mutation_versions(self, args: list[Any]) -> None:
+        if self.keep_input_mutations:
+            mutated_args = (
+                args[i]
+                for i in self.runtime_metadata.mutated_graph_handled_indices_seen_by_autograd
+            )
+            torch.autograd.graph.increment_version(mutated_args)
 
-        num_mutated_runtime_inps = runtime_metadata.num_mutated_inp_runtime_indices
-        num_intermediate_bases = runtime_metadata.num_intermediate_bases
+    def finalize(self, orig_inputs: dict[int, Any], all_outs: list[Any]) -> Any:
+        self._validate_compiled_output_arity(all_outs)
+        updated_inputs, fw_outs = self._split_mutated_inputs(all_outs)
+        if updated_inputs is not None:
+            self._apply_input_mutations(orig_inputs, updated_inputs)
 
+        ret_outs = self._replay_output_aliases(orig_inputs, fw_outs)
+        if self.runtime_metadata.dynamic_outputs:
+            for t, o in zip(ret_outs, self.runtime_metadata.output_info):
+                if o.dynamic_dims is None:
+                    continue
+                maybe_mark_dynamic_helper(t, o.dynamic_dims)
+        if self.runtime_metadata.grad_enabled_mutation is not None:
+            torch._C._set_grad_enabled(self.runtime_metadata.grad_enabled_mutation)
+        return ret_outs
+
+    def _validate_compiled_output_arity(self, all_outs: list[Any]) -> None:
         expected_outs = (
-            num_mutated_runtime_inps
-            + runtime_metadata.num_outputs
-            + num_intermediate_bases
+            self.runtime_metadata.num_mutated_inp_runtime_indices
+            + self.runtime_metadata.num_outputs
+            + self.runtime_metadata.num_intermediate_bases
         )
         if len(all_outs) != expected_outs:
             raise AssertionError(
                 f"expected {expected_outs} outputs, got {len(all_outs)}"
             )
 
-        # Step 3: After running the compiled fw, apply updates to mutated inputs
-        if num_mutated_runtime_inps > 0:
-            updated_inputs = all_outs[:num_mutated_runtime_inps]
-            fw_outs = all_outs[num_mutated_runtime_inps:]
+    def _split_mutated_inputs(
+        self, all_outs: list[Any]
+    ) -> tuple[list[Any] | None, list[Any]]:
+        num_mutated_runtime_inps = self.runtime_metadata.num_mutated_inp_runtime_indices
+        if num_mutated_runtime_inps == 0:
+            return None, all_outs
+        return (
+            all_outs[:num_mutated_runtime_inps],
+            all_outs[num_mutated_runtime_inps:],
+        )
 
-            for i, inpt_idx in enumerate(runtime_metadata.mutated_inp_runtime_indices):
-                meta = runtime_metadata.input_info[inpt_idx]
-                if not meta.mutates_data and not meta.mutates_metadata:
-                    continue
-                original_inpt = orig_inputs[inpt_idx]
-                updated_inpt = updated_inputs[i]
-                if meta.mutates_storage_metadata:
-                    # See Note [set_() Input Mutations in AOTAutograd]
-                    # mutates_storage_metadata means our input saw a x.set_(y) call.
-                    # What if x **also** saw a data and/or a metadata mutation?
-                    # (1) If the [meta]data mutation occurred after the set_(),
-                    #     then there is no need to copy_() the data.
-                    #     When we perform x.set_(x_updated), we are guaranteed that
-                    #     x_updated already has the final version of the data/metadata
-                    # (2) If a data mutation occurred before the set_().
-                    #     This case seems very difficult to support.
-                    #     TODO: discuss on the PR and decide if we want to tr to
-                    #     either support it, or detect and ban it.
-                    if trace_joint:
-                        if not isinstance(updated_inpt, TensorAlias):
-                            raise AssertionError(
-                                f"expected TensorAlias for updated_inpt, got {type(updated_inpt)}"
-                            )
-                        updated_inpt = updated_inpt.alias
-                    with torch.no_grad():
-                        original_inpt.set_(updated_inpt)
-                    continue
-                if meta.mutates_metadata and not meta.mutates_data:
-                    if trace_joint:
-                        if not isinstance(updated_inpt, TensorAlias):
-                            raise AssertionError(
-                                f"expected TensorAlias for updated_inpt, got {type(updated_inpt)}"
-                            )
-                        updated_inpt = updated_inpt.alias
-                    # We need to grab the size/stride/storage_offset from the compiled forward,
-                    # and use that to mutate the metadata of the input
+    def _apply_input_mutations(
+        self, orig_inputs: dict[int, Any], updated_inputs: list[Any]
+    ) -> None:
+        for i, inpt_idx in enumerate(self.runtime_metadata.mutated_inp_runtime_indices):
+            meta = self.runtime_metadata.input_info[inpt_idx]
+            if not meta.mutates_data and not meta.mutates_metadata:
+                continue
+            original_inpt = orig_inputs[inpt_idx]
+            updated_inpt = updated_inputs[i]
+            if meta.mutates_storage_metadata:
+                # See Note [set_() Input Mutations in AOTAutograd]
+                # mutates_storage_metadata means our input saw a x.set_(y) call.
+                # What if x **also** saw a data and/or a metadata mutation?
+                # (1) If the [meta]data mutation occurred after the set_(),
+                #     then there is no need to copy_() the data.
+                #     When we perform x.set_(x_updated), we are guaranteed that
+                #     x_updated already has the final version of the data/metadata
+                # (2) If a data mutation occurred before the set_().
+                #     This case seems very difficult to support.
+                #     TODO: discuss on the PR and decide if we want to tr to
+                #     either support it, or detect and ban it.
+                if self.trace_joint:
+                    if not isinstance(updated_inpt, TensorAlias):
+                        raise AssertionError(
+                            f"expected TensorAlias for updated_inpt, got {type(updated_inpt)}"
+                        )
+                    updated_inpt = updated_inpt.alias
+                with torch.no_grad():
+                    original_inpt.set_(updated_inpt)
+                continue
+            if meta.mutates_metadata and not meta.mutates_data:
+                if self.trace_joint:
+                    if not isinstance(updated_inpt, TensorAlias):
+                        raise AssertionError(
+                            f"expected TensorAlias for updated_inpt, got {type(updated_inpt)}"
+                        )
+                    updated_inpt = updated_inpt.alias
+                # We need to grab the size/stride/storage_offset from the compiled forward,
+                # and use that to mutate the metadata of the input
+                original_inpt.as_strided_(
+                    updated_inpt.size(),
+                    updated_inpt.stride(),
+                    updated_inpt.storage_offset(),
+                )
+            else:
+                if meta.mutates_data and meta.mutates_metadata:
                     original_inpt.as_strided_(
                         updated_inpt.size(),
                         updated_inpt.stride(),
                         updated_inpt.storage_offset(),
                     )
                 else:
-                    if meta.mutates_data and meta.mutates_metadata:
-                        original_inpt.as_strided_(
-                            updated_inpt.size(),
-                            updated_inpt.stride(),
-                            updated_inpt.storage_offset(),
+                    if not meta.mutates_data:
+                        raise AssertionError("expected meta.mutates_data to be True")
+                if meta.is_leaf and original_inpt.requires_grad:
+                    # We can hit this situation in this case:
+                    #   def f(x):
+                    #       x.detach().mul_(2)
+                    #       return x + 1
+                    # AOTAutograd will see a mutation in the above case, and try to
+                    # apply a copy_() here, in the epilogue.
+                    # But if x required gradients, and is a leaf, then autograd
+                    # will yell at us for trying to mutate it.
+                    # However, it's only possible to end up in this scenario (like the above)
+                    # if all of the mutations to the leaf input were non-autograd-tracking mutations
+                    # (aka mutations under no_grad(), or on detached views).
+                    # In that case, we fully want to hide the mutation from autograd, so detaching is ok.
+                    original_inpt.detach().copy_(updated_inpt)
+                else:
+                    # Check if we have stream index information for this mutated input
+                    if (
+                        self.runtime_metadata.mutated_inp_stream_indices is not None
+                        and i < len(self.runtime_metadata.mutated_inp_stream_indices)
+                        and self.runtime_metadata.mutated_inp_stream_indices[i]
+                        is not None
+                    ):
+                        raise RuntimeError(
+                            "Mutations on inputs with user-specified streams are not yet supported. "
+                            "See: https://github.com/pytorch/pytorch/issues/172522"
                         )
-                    else:
-                        if not meta.mutates_data:
-                            raise AssertionError(
-                                "expected meta.mutates_data to be True"
-                            )
-                    if meta.is_leaf and original_inpt.requires_grad:
-                        # We can hit this situation in this case:
-                        #   def f(x):
-                        #       x.detach().mul_(2)
-                        #       return x + 1
-                        # AOTAutograd will see a mutation in the above case, and try to
-                        # apply a copy_() here, in the epilogue.
-                        # But if x required gradients, and is a leaf, then autograd
-                        # will yell at us for trying to mutate it.
-                        # However, it's only possible to end up in this scenario (like the above)
-                        # if all of the mutations to the leaf input were non-autograd-tracking mutations
-                        # (aka mutations under no_grad(), or on detached views).
-                        # In that case, we fully want to hide the mutation from autograd, so detaching is ok.
-                        original_inpt.detach().copy_(updated_inpt)
-                    else:
-                        # Check if we have stream index information for this mutated input
-                        if (
-                            runtime_metadata.mutated_inp_stream_indices is not None
-                            and i < len(runtime_metadata.mutated_inp_stream_indices)
-                            and runtime_metadata.mutated_inp_stream_indices[i]
-                            is not None
-                        ):
-                            raise RuntimeError(
-                                "Mutations on inputs with user-specified streams are not yet supported. "
-                                "See: https://github.com/pytorch/pytorch/issues/172522"
-                            )
-                        original_inpt.copy_(updated_inpt)
-        else:
-            fw_outs = all_outs
-
-        # Step 4: Manually regenerate any outputs that are aliased to inputs, instead of
-        # compiling them.
-        if runtime_metadata.num_outputs_aliased > 0:
-            # The compiled forward also returned intermediate bases. We don't want to return them to the user.
-            expect_num_outputs = (
-                len(output_handlers) + runtime_metadata.num_intermediate_bases
+                    original_inpt.copy_(updated_inpt)
+
+    def _replay_output_aliases(
+        self, orig_inputs: dict[int, Any], fw_outs: list[Any]
+    ) -> Any:
+        if self.runtime_metadata.num_outputs_aliased == 0:
+            return fw_outs
+
+        # The compiled forward also returned intermediate bases. We don't want to return them to the user.
+        expect_num_outputs = (
+            len(self.output_handlers) + self.runtime_metadata.num_intermediate_bases
+        )
+        if len(fw_outs) != expect_num_outputs:
+            raise AssertionError(
+                f"expected {expect_num_outputs} fw_outs, got {len(fw_outs)}"
             )
-            if len(fw_outs) != expect_num_outputs:
-                raise AssertionError(
-                    f"expected {expect_num_outputs} fw_outs, got {len(fw_outs)}"
-                )
-            ret_outs = [
-                handler(orig_inputs, fw_outs, out)
-                for out, handler in builtins.zip(fw_outs, output_handlers)
-            ]
-        else:
-            ret_outs = fw_outs
+        return [
+            handler(orig_inputs, fw_outs, out)
+            for out, handler in builtins.zip(fw_outs, self.output_handlers)
+        ]
 
-        if runtime_metadata.dynamic_outputs:
-            for t, o in zip(ret_outs, runtime_metadata.output_info):
-                if o.dynamic_dims is None:
-                    continue
-                maybe_mark_dynamic_helper(t, o.dynamic_dims)
-        if runtime_metadata.grad_enabled_mutation is not None:
-            torch._C._set_grad_enabled(runtime_metadata.grad_enabled_mutation)
-        return ret_outs
+
+def _create_runtime_wrapper(
+    compiled_fn: Callable[..., Any],
+    *,
+    runtime_metadata: ViewAndMutationMeta,
+    indices_of_inps_to_detach: list[int],
+    trace_joint: bool,
+    keep_input_mutations: bool,
+    disable_amp: bool,
+) -> Callable[..., Any]:
+    compiled_invoker = _RuntimeCompiledFnInvoker(
+        compiled_fn=compiled_fn,
+        indices_of_inps_to_detach=indices_of_inps_to_detach,
+        trace_joint=trace_joint,
+        disable_amp=disable_amp,
+    )
+    runtime_epilogue = _RuntimeForwardEpilogue(
+        runtime_metadata=runtime_metadata,
+        trace_joint=trace_joint,
+        keep_input_mutations=keep_input_mutations,
+    )
+
+    def record_runtime_wrapper_prologue_enter() -> AbstractContextManager[None] | None:
+        if (
+            torch.autograd.profiler._is_profiler_enabled
+            and dynamo_config.record_runtime_overhead
+        ):
+            cm = torch._C._profiler._RecordFunctionFast(
+                "AOTDispatcher Runtime Wrapper Prologue"
+            )
+            cm.__enter__()
+            return cm
+        return None
+
+    def record_runtime_wrapper_prologue_exit(
+        cm: AbstractContextManager[None] | None,
+    ) -> None:
+        if cm is not None:
+            cm.__exit__(None, None, None)
+
+    @simple_wraps(compiled_invoker.compiled_fn)
+    def runtime_wrapper(args: list[Any]) -> Any:
+        # Create context manager for profiler
+        cm = record_runtime_wrapper_prologue_enter()
+        prologue_exited = False
+
+        def exit_prologue() -> None:
+            nonlocal prologue_exited
+            if not prologue_exited:
+                record_runtime_wrapper_prologue_exit(cm)
+                prologue_exited = True
+
+        try:
+            # stash a ref to each input tensor we plan to use after the compiled function
+            orig_inputs = runtime_epilogue.capture_orig_inputs(args)
+            runtime_epilogue.increment_mutation_versions(args)
+            all_outs = compiled_invoker.run(args, on_before_call=exit_prologue)
+        finally:
+            exit_prologue()
+
+        del args
+        return runtime_epilogue.finalize(orig_inputs, all_outs)
 
     if not (trace_joint and _should_disable_saved_tensors_hooks()):
         return runtime_wrapper
@@ -931,6 +998,17 @@ def pre_compile(
         self.maybe_subclass_meta = subclass_meta
         return new_flat_fn, new_flat_args, new_flat_args_descs, fw_metadata
 
+    @staticmethod
+    def _get_frozen_inp_indices() -> frozenset[int]:
+        # fw_compiler_freezing (compile_fx.py) bakes frozen params into the
+        # graph and sets their TracingContext.params_flat entries to None
+        # before post_compile runs.  We pass these indices to codegen so it
+        # can emit straight-line code instead of a runtime None check.
+        tc = TracingContext.try_get()
+        if tc is None or tc.params_flat is None:
+            return frozenset()
+        return frozenset(i for i, p in enumerate(tc.params_flat) if p is None)
+
     def post_compile(
         self,
         compiled_fn: Callable[..., Any],
@@ -938,53 +1016,19 @@ def post_compile(
         *,
         runtime_metadata: ViewAndMutationMeta,
     ) -> Callable[..., Any]:
-        if self.maybe_subclass_meta is None:
+        if self.maybe_subclass_meta is None and not runtime_metadata.act_input_indices:
             return compiled_fn
 
-        subclass_metas = runtime_metadata.subclass_fw_graph_out_meta
+        from .subclass_codegen import codegen_subclass_wrapper
 
-        @wraps(compiled_fn)
-        def inner_fn(args: list[Any]) -> Any:
-            if aot_graphs_log.isEnabledFor(logging.DEBUG):
-                aot_graphs_log.debug(
-                    "=== AOTDispatchSubclassWrapper.inner_fn START ==="
-                )
-                _log_input_metadata(runtime_metadata)
-                _log_args_list(args, "Incoming args")
-
-            unwrapped_args = runtime_unwrap_tensor_subclasses(
-                args,
-                subclass_metas=runtime_metadata.subclass_inp_meta,
-                append_symints=True,
-            )
-
-            if aot_graphs_log.isEnabledFor(logging.DEBUG):
-                _log_args_list(unwrapped_args, "After unwrapping, unwrapped_args")
-
-            args.clear()
-            # expectation: runtime_fn is a boxed fn
-            unwrapped_outs = compiled_fn(unwrapped_args)
-
-            if aot_graphs_log.isEnabledFor(logging.DEBUG):
-                _log_args_maybe_list(
-                    unwrapped_outs, "After compiled_fn, unwrapped_outs"
-                )
-
-            wrapped_outs = wrap_tensor_subclasses(
-                unwrapped_outs,
-                subclass_metas=subclass_metas,
-                num_fw_outs_saved_for_bw=self.num_fw_outs_saved_for_bw,
-                is_runtime=True,
-                included_subclass_symints=True,
-            )
-
-            if aot_graphs_log.isEnabledFor(logging.DEBUG):
-                _log_args_maybe_list(wrapped_outs, "After wrapping, wrapped_outs")
-                aot_graphs_log.debug("=== AOTDispatchSubclassWrapper.inner_fn END ===")
-
-            return wrapped_outs
-
-        # box it
+        inner_fn = codegen_subclass_wrapper(
+            compiled_fn=compiled_fn,
+            inp_metas=runtime_metadata.subclass_inp_meta,
+            out_metas=runtime_metadata.subclass_fw_graph_out_meta,
+            num_fw_outs_saved_for_bw=self.num_fw_outs_saved_for_bw,
+            frozen_inp_indices=self._get_frozen_inp_indices(),
+            act_input_indices=runtime_metadata.act_input_indices,
+        )
         inner_fn._boxed_call = True  # type: ignore[attr-defined]
         return inner_fn
 
@@ -1272,7 +1316,6 @@ def wrapped_flat_fn(
                 flat_args_descs=deduped_flat_args_descs,
                 static_input_indices=aot_config.static_input_indices,
                 keep_input_mutations=fw_metadata.keep_input_mutations,
-                is_train=fw_metadata.is_train,
             )(*deduped_flat_args)
             if ref_fw_metadata != updated_fw_metadata:
                 raise AssertionError(
@@ -1296,11 +1339,23 @@ def post_compile(
         if not self.needs_post_compile:
             return compiled_fn
 
-        @wraps(compiled_fn)
-        def wrapped_compiled_fn(args: list[Any]) -> Any:
-            deduped_args = self.remove_dupe_args(args)
-            args.clear()
-            return compiled_fn(deduped_args)
+        keep_indices = [i for i, keep in enumerate(self.keep_arg_mask) if keep]
+        idx_list = ", ".join(f"args[{i}]" for i in keep_indices)
+        source = (
+            f"def inner_fn(args):\n"
+            f"    deduped_args = [{idx_list}]\n"
+            f"    args.clear()\n"
+            f"    return compiled_fn(deduped_args)\n"
+        )
+        from .subclass_codegen import _compile_and_exec_source
+
+        wrapped_compiled_fn: Callable[..., Any] = _compile_and_exec_source(  # type: ignore[assignment]
+            source,
+            {"compiled_fn": compiled_fn},
+            "inner_fn",
+            "dedup_wrapper",
+            wrapped_fn=compiled_fn,
+        )
 
         wrapped_compiled_fn._boxed_call = True  # type: ignore[attr-defined]
 
@@ -1495,7 +1550,6 @@ def wrapped_flat_fn(*args: Any) -> Any:
                 flat_args_descs=flat_args_descs_with_synthetic_bases,
                 static_input_indices=aot_config.static_input_indices,
                 keep_input_mutations=fw_metadata.keep_input_mutations,
-                is_train=fw_metadata.is_train,
             )(*flat_args_with_synthetic_bases)
             if ref_fw_metadata != fw_metadata_updated:
                 raise AssertionError(
@@ -1690,7 +1744,7 @@ def _same_dtype_views(view1: torch.Tensor, view2: torch.Tensor) -> bool:
     # - another int (corresponding to the index in the argument list of the element from the outer calling convention)
     # - idx, view_tensor, where we can generate the new output with view_tensor._view_func(old_args[idx])
     #   idx corresponds to which synthetic base from the outer calling context to view
-    inner_calling_convention_meta: dict[int, Union[int, tuple[int, torch.Tensor]]] = {}
+    inner_calling_convention_meta: dict[int, int | tuple[int, torch.Tensor]] = {}
     for aliased_input_indices in storage_ref_to_idx.values():
         if len(aliased_input_indices) <= 1 or not any(
             # We only care about mutations that affect all aliases,
@@ -1855,9 +1909,9 @@ def make_hashable(arg: Any) -> Any:
             inner_calling_convention_meta[old_idx] = new_idx
 
         # post process into a list
-        post_processed_calling_convention_meta: list[
-            Union[int, tuple[int, torch.Tensor]]
-        ] = [-1 for _ in range(len(inner_calling_convention_meta))]
+        post_processed_calling_convention_meta: list[int | tuple[int, torch.Tensor]] = [
+            -1 for _ in range(len(inner_calling_convention_meta))
+        ]
         for k, v in inner_calling_convention_meta.items():
             post_processed_calling_convention_meta[k] = v
         # Quick assert: every argument in the inner calling convention should be accounted for.
@@ -1866,6 +1920,7 @@ def make_hashable(arg: Any) -> Any:
                 raise AssertionError(
                     "every argument in the inner calling convention should be accounted for"
                 )
+        # pyrefly: ignore [bad-return]
         return (
             args_to_functionalization,
             args_to_functionalization_descs,
@@ -1917,6 +1972,7 @@ def _backward_prologue_functional(
     metadata: ViewAndMutationMeta,
     maybe_subclass_metadata: SubclassMeta | None,
     *flat_args: Any,
+    codegen_unwrap_fn: Callable[..., Any] | None = None,
 ) -> list[Any]:
     # Calling convention: we expect a grad_out passed to the backward:
     # - for every output of the fw that does *not* alias an input or graph intermediate
@@ -1983,7 +2039,7 @@ def _backward_prologue_functional(
             OutputType.custom_function_view,
         ]
         and issubclass(info.raw_type, torch.Tensor)
-        and info.requires_grad
+        and info.requires_grad_for_backward
     ]
     # intermediate bases always require gradients, and always participate in the backward graph.
     flat_bw_args_with_grads = [
@@ -2011,7 +2067,7 @@ def _backward_prologue_functional(
 
     bw_tokens = [None] * metadata.num_backward_tokens
 
-    # - note: donated buffer logic requires (*ctx.symints, *ctx.saved_tensors) showing up first
+    # - note: donated buffer logic requires (*ctx.symints, *ctx.saved_tensors, *ctx.opaques) showing up first
     #   in the bw output order.
 
     # Every dereference of ctx.saved_tensors incurs saved_tensors_hooks calls
@@ -2073,44 +2129,53 @@ def _backward_prologue_functional(
                 "The grad inputs should be same number as forward output tangents"
             )
 
+        stack_traces = metadata.tangent_source_stack_traces or ()
+
         flat_processed_tangents = list(
             itertools.chain.from_iterable(
                 (
                     AOTDispatchAutograd.process_runtime_tangent(
                         t,
                         m,
+                        tangent_idx=idx,
+                        tangent_desc=desc,
+                        compile_id_str=metadata.compile_id_str,
+                        tangent_stack_trace=stack_traces[idx] if stack_traces else None,
                     )[1]
                 )
-                for t, m in zip(
-                    tangents,
-                    metadata.subclass_tangent_meta,
+                for idx, (t, m, desc) in enumerate(
+                    zip(
+                        tangents,
+                        metadata.subclass_tangent_meta,
+                        metadata.traced_tangents_descs,
+                    )
                 )
             )
         )
 
+        if codegen_unwrap_fn is not None:
+            unwrap = codegen_unwrap_fn
+        else:
+            unwrap = _unwrap_no_symints
         all_args = (
-            runtime_unwrap_tensor_subclasses(
-                all_args[:tangents_start_idx],  # type: ignore[arg-type]
-                # SymInts that are inputs to the backward graph are
-                # already included in the "all_args" list.
-                # Any symints coming from tensor subclasses should always
-                # come from primals, and so they will show up as extra
-                # arguments to the forward graph, and they will be saved
-                # as activation in the backward graph.
-                append_symints=False,
-            )
+            unwrap(all_args[:tangents_start_idx])
             + flat_processed_tangents
-            + runtime_unwrap_tensor_subclasses(
-                all_args[tangents_end_idx:],  # type: ignore[arg-type]
-                append_symints=False,
-            )
+            + unwrap(all_args[tangents_end_idx:])
         )
     else:
+        stack_traces = metadata.tangent_source_stack_traces or ()
+
         all_args = [
             (
                 AOTDispatchAutograd.process_runtime_tangent(
                     t,
                     metadata.subclass_tangent_meta[i - tangents_start_idx],
+                    tangent_idx=i - tangents_start_idx,
+                    tangent_desc=metadata.traced_tangents_descs[i - tangents_start_idx],
+                    compile_id_str=metadata.compile_id_str,
+                    tangent_stack_trace=(
+                        stack_traces[i - tangents_start_idx] if stack_traces else None
+                    ),
                 )[0]
                 if (tangents_start_idx <= i < tangents_end_idx)
                 else t
@@ -2175,7 +2240,9 @@ def _backward_epilogue_functional(
     maybe_subclass_metadata: SubclassMeta | None,
     out: Any,
     *,
+    ctx_opaque_objects: Sequence[Any] = (),
     make_subclass_override: Callable[..., Any] | None = None,
+    codegen_wrap_fn: Callable[..., Any] | None = None,
 ) -> tuple[Any, ...]:
     # Toss out the backward output tokens
     num_bw_tokens = metadata.num_backward_tokens
@@ -2188,10 +2255,29 @@ def _backward_epilogue_functional(
     )
     out = tuple(out)
 
+    # Replace compile-time opaque constants in the backward output with the
+    # real runtime opaques saved from the forward pass. During joint graph
+    # tracing, backward output opaques come from tangent constants (baked at
+    # compile time). At runtime we need the actual opaque objects that were
+    # saved for backward from the forward pass.
+    if ctx_opaque_objects:
+        opaque_iter = iter(ctx_opaque_objects)
+        out = tuple(
+            next(opaque_iter) if isinstance(v, FakeScriptObject) else v for v in out
+        )
+        remaining = list(opaque_iter)
+        if remaining:
+            raise AssertionError(
+                f"ctx_opaque_objects had {len(remaining)} leftover entries "
+                "(expected all to be consumed by FakeScriptObject slots in backward output)"
+            )
+
     # TODO: figure out how to refactor the backward properly so I can use aot_dispatch_subclass_wrapper() here.
     if maybe_subclass_metadata is not None:
         if maybe_subclass_metadata.grad_input_metas is None:
             raise AssertionError("grad_input_metas must not be None")
+        if codegen_wrap_fn is not None and make_subclass_override is None:
+            return codegen_wrap_fn(out)
         outs_wrapped = wrap_tensor_subclasses(
             out,
             subclass_metas=maybe_subclass_metadata.grad_input_metas,
@@ -2298,161 +2384,437 @@ def __call__(self, *args: Any, **kwargs: Any) -> Any:
         return self.compiled_fn(*args, **kwargs)
 
 
-# This is wrapped in a class just for namespacing purposes
-# No need to make it into an actual CompilerWrapper because it doesn't fit the abstract as cleanly
-class AOTDispatchAutograd:
-    @staticmethod
-    def process_runtime_tangent(
-        x: Any, meta: PlainTensorMeta | SubclassCreationMeta
-    ) -> tuple[Any, list[Any]]:
-        if not isinstance(x, torch.Tensor):
-            return x, [x]
+@dataclass
+class AOTDispatchAutogradCompileSpec:
+    compiled_fw_func: Callable[..., Any]
+    compiled_bw_func: Callable[..., Any] | None
+    maybe_subclass_meta: SubclassMeta | None
+    num_symints_saved_for_bw: int
+    backward_state_indices: list[int]
+    disable_amp: bool
+    indices_of_inps_to_detach: list[int]
+    lazy_backward_info: (
+        AutogradLazyBackwardCompileInfo | CachedAutogradLazyBackwardCompileInfo | None
+    )
+    aot_config: AOTConfig
+    fw_metadata: ViewAndMutationMeta
+    try_save_cache_entry: Callable[..., Any] | None
 
-        if isinstance(x, FakeTensor):
-            if not meta.memory_format:
-                raise AssertionError(
-                    "meta.memory_format must not be None for FakeTensor"
-                )
-            x = coerce_to_expected_memory_format(x, meta.memory_format)
-            return x, [x]
 
-        expected_type: Optional[type] = torch.Tensor
-        expected_meta = None
-        if isinstance(meta, SubclassCreationMeta):
-            expected_type = meta.original_subclass_type
-            expected_meta = meta.meta
+@dataclass
+class _AutogradSavedState:
+    metadata: ViewAndMutationMeta
 
-        runtime_type = type(x)
-        # When we're inside compiled autograd's AOTDispatcher step,
-        # regular Tensors look like FunctionalTensors.
-        # Tensor subclasses still look like Tensor subclasses though.
-        if isinstance(x, torch._subclasses.functional_tensor.FunctionalTensor):
-            runtime_type = torch.Tensor
+    def save_from_forward(self, ctx: Any, fw_outs: Sequence[Any]) -> None:
+        tensors_saved_with_vc_check = fw_outs[
+            self.metadata.tensors_saved_for_backwards_with_vc_check_slice
+        ]
+        tensors_saved_no_vc_check = fw_outs[
+            self.metadata.tensors_saved_for_backwards_no_vc_check_slice
+        ]
+        if not all(isinstance(x, torch.Tensor) for x in tensors_saved_with_vc_check):
+            raise AssertionError(
+                "expected all tensors_saved_with_vc_check to be Tensors, "
+                f"got types: {[type(x) for x in tensors_saved_with_vc_check]}"
+            )
+        if not all(isinstance(x, torch.Tensor) for x in tensors_saved_no_vc_check):
+            raise AssertionError(
+                "expected all tensors_saved_no_vc_check to be Tensors, "
+                f"got types: {[type(x) for x in tensors_saved_no_vc_check]}"
+            )
 
-        runtime_meta = None
-        runtime_subclass_keys: Sequence[str] = []
+        # See Note [Detaching saved tensors in AOTAutograd]
+        num_vc_check = len(tensors_saved_with_vc_check)
+        tensors_to_save = [
+            x.detach() if x._is_view() else x for x in tensors_saved_with_vc_check
+        ]
+        tensors_no_vc_check = [
+            x.detach() if x._is_view() else x for x in tensors_saved_no_vc_check
+        ]
 
-        if is_traceable_wrapper_subclass(x):
-            runtime_subclass_keys, runtime_meta = x.__tensor_flatten__()
+        # dynamic_saved_tensors_idxs has indices relative to all saved tensors
+        # (vc_check + no_vc_check combined). Mark dynamics on the detached tensors.
+        for idx, dims in self.metadata.dynamic_saved_tensors_idxs.items():
+            if idx < num_vc_check:
+                maybe_mark_dynamic_helper(tensors_to_save[idx], dims)
+            else:
+                maybe_mark_dynamic_helper(tensors_no_vc_check[idx - num_vc_check], dims)
 
-        def maybe_coerce(x: torch.Tensor) -> torch.Tensor | None:
-            same_type: bool = expected_type == runtime_type
-            same_meta: bool = expected_meta == runtime_meta
+        ctx.save_for_backward(*tensors_to_save)
+        ctx._tensors_no_vc_check = tensors_no_vc_check
 
-            if same_type and same_meta:
-                return x
+        symint_outs = fw_outs[self.metadata.symints_saved_for_backwards_slice]
+        if not all(
+            isinstance(x, (int, float, torch.SymInt, torch.SymFloat))
+            for x in symint_outs
+        ):
+            raise AssertionError(
+                "expected all symint_outs to be int/float/SymInt/SymFloat, "
+                f"got types: {[type(x) for x in symint_outs]}"
+            )
+        ctx.symints = symint_outs
 
-            if not hasattr(x, "__coerce_same_metadata_as_tangent__"):
-                return None
+        opaque_object_outs = fw_outs[
+            self.metadata.opaque_objects_saved_for_backwards_slice
+        ]
+        if not all(
+            is_opaque_type(type(obj)) or isinstance(obj, OpaqueBase)
+            for obj in opaque_object_outs
+        ):
+            raise AssertionError(
+                "expected all opaque_object_outs to be opaque types, "
+                f"got types: {[type(obj) for obj in opaque_object_outs]}"
+            )
+        ctx.opaque_objects = opaque_object_outs
 
-            if same_type:
-                # Backward Compatibility, as some Subclass impls can have original 1-arg function.
-                return x.__coerce_same_metadata_as_tangent__(expected_meta)
+    def load_tensors(self, ctx: Any) -> Sequence[torch.Tensor]:
+        if len(ctx._tensors_no_vc_check) > 0:
+            return list(ctx.saved_tensors) + ctx._tensors_no_vc_check
+        return ctx.saved_tensors
 
-            return x.__coerce_same_metadata_as_tangent__(expected_meta, expected_type)
 
-        # Coerce to expected type and metadata
-        orig_x = x
-        x = maybe_coerce(x)
-        if x is None:
-            raise RuntimeError(
-                f"""
-During the backward, we encountered a tensor subclass where we guessed its
-metadata incorrectly.
+@dataclass
+class _AutogradForwardEpilogue:
+    metadata: ViewAndMutationMeta
 
-Expected metadata: {str(expected_meta)}, expected type: {str(expected_type)}
+    def finalize(self, ctx: Any, fw_outs: Sequence[Any]) -> tuple[Any, ...]:
+        num_outputs = self.metadata.num_outputs
+        num_outputs_aliased = self.metadata.num_outputs_aliased
+        num_mutated_runtime_inps = self.metadata.num_mutated_inp_runtime_indices
+        num_forward_returns = self.metadata.num_forward_returns
 
-Runtime metadata: {str(runtime_meta)}, runtime type: {str(runtime_type)}
+        raw_returns = list(fw_outs[:num_forward_returns])
 
-shape: {str(orig_x.shape)}
-To fix this, your tensor subclass must implement the dunder method __force_to_same_metadata__.
-"""
+        # Wrap all autograd.Function.forward() outputs that are aliases
+        # so that autograd.Function doesn't treat them as tensors
+        if num_mutated_runtime_inps > 0:
+            for i, idx in enumerate(self.metadata.mutated_inp_runtime_indices):
+                # We could make this faster by only looping over inputs with metadata-only mutations
+                # (instead of looping over inputs with either data or metadata mutations), but there shouldn't be many.
+                info = self.metadata.input_info[idx]
+                if info.mutates_metadata and not info.mutates_data:
+                    raw_returns[i] = TensorAlias(raw_returns[i])
+
+            if config.debug_assert:
+                user_mutated_inputs_raw = raw_returns[0:num_mutated_runtime_inps]
+                mut_inp_infos = [
+                    x
+                    for x in self.metadata.input_info
+                    if x.mutates_data or x.mutates_metadata
+                ]
+                if len(user_mutated_inputs_raw) != len(mut_inp_infos):
+                    raise AssertionError(
+                        "expected len(user_mutated_inputs_raw) == len(mut_inp_infos), "
+                        f"got {len(user_mutated_inputs_raw)} != {len(mut_inp_infos)}"
+                    )
+
+        if self.metadata.num_unsafe_view_outputs > 0:
+            for idx in self.metadata.unsafe_view_out_indices:
+                raw_return_idx = num_mutated_runtime_inps + idx
+                o = raw_returns[raw_return_idx]
+                raw_returns[raw_return_idx] = torch.ops.aten._unsafe_view(o, o.shape)
+
+        if num_outputs_aliased > 0:
+            for idx in self.metadata.aliased_out_indices:
+                raw_return_idx = num_mutated_runtime_inps + idx
+                raw_returns[raw_return_idx] = TensorAlias(raw_returns[raw_return_idx])
+
+            if config.debug_assert:
+                intermediates_raw = raw_returns[
+                    num_mutated_runtime_inps + num_outputs :
+                ]
+                if any(isinstance(x, TensorAlias) for x in intermediates_raw):
+                    raise AssertionError("expected no TensorAlias in intermediates_raw")
+
+        # invariant: intermediate bases always require gradients, so we don't have to
+        # consider marking them as non-differentiable.
+        raw_returns_not_including_intermediate_bases = raw_returns[
+            : num_mutated_runtime_inps + num_outputs
+        ]
+        raw_returns_meta = [
+            x
+            for x in self.metadata.input_info
+            if x.mutation_type == MutationType.MUTATED_OUT_GRAPH
+        ] + self.metadata.output_info
+
+        fw_outs_not_requiring_grad = [
+            x
+            for (i, x) in enumerate(raw_returns_not_including_intermediate_bases)
+            if isinstance(x, torch.Tensor) and not raw_returns_meta[i].requires_grad
+        ]
+        ctx.mark_non_differentiable(*fw_outs_not_requiring_grad)
+        ctx._materialize_non_diff_grads = False
+        return tuple(raw_returns)
+
+
+@dataclass
+class _AutogradRngStateTracker:
+    num_rng: int
+    graphsafe_idx: int | None
+    fwd_rng_states: list[torch.Generator] = field(default_factory=list)
+    bwd_rng_states: list[torch.Generator] = field(default_factory=list)
+    curr_fwd_iter: Any = field(default_factory=lambda: itertools.count(0))
+    backward_state_position: int = 0
+    pending_forwards: set[int] = field(default_factory=set)
+    saved_backward_tensor_states: dict[int, list[torch.Tensor]] = field(
+        default_factory=dict
+    )
+
+    def add_forward_args(self, ctx: Any, args: tuple[Any, ...]) -> tuple[Any, ...]:
+        if self.num_rng == 0:
+            return args
+
+        if len(self.fwd_rng_states) == 0:
+            if self.graphsafe_idx is None:
+                raise AssertionError("graphsafe_idx must not be None when num_rng > 0")
+            initialize_rng_states(
+                self.num_rng,
+                self.graphsafe_idx,
+                self.fwd_rng_states,
+                self.bwd_rng_states,
             )
 
-        # Coerce to expected memory format
-        if not meta.memory_format:
-            raise AssertionError("meta.memory_format must not be None")
-        x = coerce_to_expected_memory_format(x, meta.memory_format)
+        curr_iter = next(self.curr_fwd_iter)
+        ctx._curr_iter = curr_iter
 
-        if not is_traceable_wrapper_subclass(x):
-            return x, [x]
+        # if this state is not contained in the backward,
+        # we need to save it for when its backward pass happens
+        if curr_iter != self.backward_state_position:
+            self.saved_backward_tensor_states[curr_iter] = [
+                rng_state.get_state() for rng_state in self.fwd_rng_states
+            ]
 
-        if not isinstance(meta, SubclassCreationMeta):
-            raise AssertionError(f"expected SubclassCreationMeta, got {type(meta)}")
-        if orig_x is not x:
-            runtime_subclass_keys = x.__tensor_flatten__()[0]
+        self.pending_forwards.add(curr_iter)
+        return (*args, *self.fwd_rng_states)
 
-        if len(meta.attrs) != len(runtime_subclass_keys):
+    def add_backward_args(self, ctx: Any, all_args: list[Any]) -> None:
+        if self.num_rng == 0:
+            return
+
+        curr_backward_iter = ctx._curr_iter
+        retain_graph = torch._C._autograd._get_current_graph_task_keep_graph()
+
+        # Save current state if we have a pending forward that needs this state
+        # or this state may be needed again because of retain graph
+        if (
+            self.backward_state_position in self.pending_forwards
+            and self.backward_state_position not in self.saved_backward_tensor_states
+            and (self.backward_state_position != curr_backward_iter or retain_graph)
+        ):
+            self.saved_backward_tensor_states[self.backward_state_position] = [
+                rng_state.get_state() for rng_state in self.bwd_rng_states
+            ]
+
+        # Restore saved states if needed
+        if curr_backward_iter in self.saved_backward_tensor_states:
+            if self.backward_state_position != curr_backward_iter:
+                for bwd_state, saved_state in zip(
+                    self.bwd_rng_states,
+                    self.saved_backward_tensor_states[curr_backward_iter],
+                ):
+                    bwd_state.set_state(saved_state)
+            if not retain_graph:
+                del self.saved_backward_tensor_states[curr_backward_iter]
+        else:
+            if self.backward_state_position != curr_backward_iter:
+                raise AssertionError(
+                    "expected backward_state_position == curr_backward_iter, "
+                    f"got {self.backward_state_position} != {curr_backward_iter}"
+                )
+
+        self.backward_state_position = curr_backward_iter + 1
+        if not retain_graph:
+            self.pending_forwards.remove(curr_backward_iter)
+        all_args.extend(self.bwd_rng_states)
+
+
+@dataclass
+class _AutogradBackwardCompiler:
+    compiled_bw: Callable[..., Any] | None
+    lazy_backward_info: (
+        AutogradLazyBackwardCompileInfo | CachedAutogradLazyBackwardCompileInfo | None
+    )
+    disable_amp: bool
+    aot_config: AOTConfig
+    fw_metadata: ViewAndMutationMeta
+    try_save_cache_entry: Callable[..., Any] | None
+
+    def get_or_compile(self, *, saved_tensors_use_once: bool) -> Callable[..., Any]:
+        if self.compiled_bw is not None:
+            return self.compiled_bw
+
+        if self.lazy_backward_info is None:
+            raise AssertionError("lazy_backward_info must not be None")
+        if not isinstance(self.lazy_backward_info, AutogradLazyBackwardCompileInfo):
             raise AssertionError(
-                f"expected len(meta.attrs) == len(runtime_subclass_keys), "
-                f"got {len(meta.attrs)} != {len(runtime_subclass_keys)}"
+                "expected AutogradLazyBackwardCompileInfo, "
+                f"got {type(self.lazy_backward_info)}"
             )
-        leaves = []
-        for attr, attr_meta in meta.attrs.items():
-            elem = getattr(x, attr)
-            new_elem, elem_leaves = AOTDispatchAutograd.process_runtime_tangent(
-                elem, attr_meta
+
+        self._prepare_lazy_backward_context(saved_tensors_use_once)
+
+        bw_module = self.lazy_backward_info.bw_module
+        placeholder_list = self.lazy_backward_info.placeholder_list
+        saved_context = self.lazy_backward_info.saved_context
+        saved_compile_context = self.lazy_backward_info.saved_compile_context
+
+        context = torch._C._DisableAutocast if self.disable_amp else nullcontext
+        metrics_context = get_metrics_context()
+        with (
+            tracing(saved_context),
+            compile_context(saved_compile_context),
+            context(),
+            track_graph_compiling(self.aot_config, "backward"),
+            metrics_context,
+            dynamo_timed(
+                "backward._backward_impl",
+                phase_name="entire_backward_compile",
+                log_pt2_compile_event=True,
+                dynamo_compile_column_us="backward_cumulative_compile_time_us",
+                log_waitcounter=True,
+                waitcounter_name_override="entire_backward_compile",
+            ),
+            callback_handler.install_callbacks(
+                CallbackTrigger.LAZY_BACKWARD,
+                str(CompileContext.current_compile_id()),
+            ),
+        ):
+            CompileEventLogger.compilation_metric(is_forward=False)
+            # See Note: [Backward graph lazy lowering]
+            if self.aot_config.bw_compiler is None:
+                raise AssertionError("aot_config.bw_compiler must not be None")
+            self.compiled_bw = self.aot_config.bw_compiler(
+                copy.deepcopy(bw_module), placeholder_list
             )
-            if new_elem is not elem:
-                setattr(x, attr, new_elem)
-            leaves.extend(elem_leaves)
+            # Maybe save cache entry
+            if self.try_save_cache_entry is not None:
+                self.try_save_cache_entry(
+                    self.compiled_bw,
+                    bw_module,
+                    self.fw_metadata,
+                    self.aot_config,
+                )
 
-        return x, leaves
+        return self.compiled_bw
 
-    @staticmethod
-    def post_compile(
-        compiled_fw_func: Callable[..., Any],  # fw_module after compilation + wrappers
-        compiled_bw_func: Callable[..., Any]
-        | None,  # bw_module after compilation + wrappers
-        maybe_subclass_meta: SubclassMeta | None,
-        num_symints_saved_for_bw_: int,
-        backward_state_indices: list[int],
-        disable_amp: bool,
-        indices_of_inps_to_detach: list[int],
-        lazy_backward_info: AutogradLazyBackwardCompileInfo
-        | CachedAutogradLazyBackwardCompileInfo
-        | None,
-        aot_config: AOTConfig,
-        *,
-        fw_metadata: ViewAndMutationMeta,  # runtime metadata
-        try_save_cache_entry: Callable[..., Any] | None,  # Serialization function
-    ) -> Callable[..., Any]:
-        # For additional context see Note [CUDA Graph Safe RNG Functionalization]
-        # Each pair forward, backward rng states must be equal prior to its invocation on any
-        # iteration of forward, backward. Because they are initialized equal, and are computing the same rng op,
-        # running forward then backward advances them the same amount and keeps them equal.
-        # However, a user may invoke multiple forwards, then backwards, such that they are not in sync.
-        # Initially we have:
-        # fwd_state0 == bwd_state0.
-        # Lets say we run:
-        # fwd0: fwd_state0 -> fwd_state1
-        # fwd1: fwd_state1 -> fwd_state2
-        # fwd2: fwd_state2 -> fwd_state3
-        # If we now invoke bwd2,
-        # we need to update bwd_state equal to the rng that was observed in fwd2.
-        # we save the rng_state fwd_state2 in forward because we detect that it is not the
-        # current backward state and therefore would not be accessible if we do not save it.
-        # Similarly, if we are going to update the backward state to a new value, and there is a pending
-        # forwards which needs its current state, we will save it.
-        # Within the autograd context, we keep track of the curr iteration so that on backward
-        # we know what the generator state must be before the backward is run.
-        num_rng = fw_metadata.num_graphsafe_rng_states
-        graphsafe_idx = fw_metadata.graphsafe_rng_state_index
-        fwd_rng_states: list[torch.Generator] = []
-        bwd_rng_states: list[torch.Generator] = []
-        curr_fwd_iter = itertools.count(0)
-        backward_state_position = 0
-        pending_forwards: set[int] = set()
-        saved_backward_tensor_states: dict[int, list[torch.Tensor]] = {}
+    def _prepare_lazy_backward_context(self, saved_tensors_use_once: bool) -> None:
+        if self.lazy_backward_info is None:
+            raise AssertionError("lazy_backward_info must not be None")
+        if not isinstance(self.lazy_backward_info, AutogradLazyBackwardCompileInfo):
+            raise AssertionError(
+                "expected AutogradLazyBackwardCompileInfo, "
+                f"got {type(self.lazy_backward_info)}"
+            )
+
+        if (
+            hasattr(self.lazy_backward_info, "saved_context")
+            and self.lazy_backward_info.saved_context is not None
+        ):
+            if not isinstance(self.lazy_backward_info.saved_context, TracingContext):
+                raise AssertionError(
+                    f"expected TracingContext, got {type(self.lazy_backward_info.saved_context)}"
+                )
+            ddp_ctx = self.lazy_backward_info.saved_context.ddp_optimizer_ctx
+            if ddp_ctx is not None:
+                if ddp_ctx.curr_bucket < 0:
+                    raise AssertionError(
+                        "expected same # of fw and bw compiles, "
+                        f"but found bucket {ddp_ctx.curr_bucket}"
+                    )
+                curr_fw_meta = ddp_ctx.metadata_per_bucket[ddp_ctx.curr_bucket]
+                # Note [DDPOptimizer and fw_metadata]
+                # When using the DDPOptimizer, we have a single dynamo graph (and TracingContext),
+                # but multiple AOTDispatcher graph.
+                #
+                # One consequence is that there will be **multiple** fw_metadata objects, one per AOT graph,
+                # which we stash the fw_metadata on the TracingContext.
+                #
+                # Normally what happens is that as we compile AOT graphs 1...N, we clobber the fw_metadata
+                # for graph i-1 when we start running AOT for graph i.
+                # Ordinarily this is fine, because inductor no longer needs the metadata from graph i-1.
+                #
+                # However, this is a problem for lazy compilation of the backward. During backward compilation,
+                # we compile the backward lazily at backward runtime, meaning that we will first compile
+                # backward graph N, N-1, ..., 1.
+                # We need to ensure that at the time inductor compiles bw graph N-1, it can access
+                # the corresponding fw_metadta for graph N-1.
+                #
+                # We do this by stashing a DDPOptimizerContext, which tracks:
+                # - the metadata of all N graphs
+                # - the graph we are currently compiling in our DDPOptimizer region.
+                ddp_ctx.curr_bucket -= 1
+                self.lazy_backward_info.saved_context.fw_metadata = curr_fw_meta
+
+        if not saved_tensors_use_once:
+            self.fw_metadata.bw_donated_idxs = []
+            # Update bw_donated_idxs if using lazy_backward_info from `aot_dispatch_autograd`
+            if (
+                hasattr(self.lazy_backward_info, "saved_context")
+                and hasattr(self.lazy_backward_info.saved_context, "fw_metadata")
+                and hasattr(
+                    self.lazy_backward_info.saved_context.fw_metadata,  # type: ignore[union-attr]
+                    "bw_donated_idxs",
+                )
+            ):
+                self.lazy_backward_info.saved_context.fw_metadata.bw_donated_idxs = (  # type: ignore[union-attr]
+                    # pyrefly: ignore [implicit-any]
+                    []
+                )
+
+
+@dataclass
+class _AOTDispatchAutogradFunctionFactory:
+    spec: AOTDispatchAutogradCompileSpec
+
+    def build(self) -> type[torch.autograd.Function]:
+        compile_id = CompileContext.current_compile_id()
+        compile_id_str = str(compile_id) if compile_id is not None else None
+        self.spec.fw_metadata.compile_id_str = compile_id_str
+
+        saved_state = _AutogradSavedState(self.spec.fw_metadata)
+        forward_epilogue = _AutogradForwardEpilogue(self.spec.fw_metadata)
+        rng_state = _AutogradRngStateTracker(
+            num_rng=self.spec.fw_metadata.num_graphsafe_rng_states,
+            graphsafe_idx=self.spec.fw_metadata.graphsafe_rng_state_index,
+        )
+        backward_compiler = _AutogradBackwardCompiler(
+            compiled_bw=self.spec.compiled_bw_func,
+            lazy_backward_info=self.spec.lazy_backward_info,
+            disable_amp=self.spec.disable_amp,
+            aot_config=self.spec.aot_config,
+            fw_metadata=self.spec.fw_metadata,
+            try_save_cache_entry=self.spec.try_save_cache_entry,
+        )
+
+        compiled_fw_func = self.spec.compiled_fw_func
+        compiled_bw_func = self.spec.compiled_bw_func
+        maybe_subclass_meta = self.spec.maybe_subclass_meta
+        num_symints_saved_for_bw_ = self.spec.num_symints_saved_for_bw
+        backward_state_indices = self.spec.backward_state_indices
+        disable_amp = self.spec.disable_amp
+        lazy_backward_info = self.spec.lazy_backward_info
+        aot_config = self.spec.aot_config
+        fw_metadata = self.spec.fw_metadata
+
+        _codegen_bw_unwrap_fn = None
+        _codegen_bw_wrap_fn = None
+        if maybe_subclass_meta is not None:
+            from .subclass_codegen import codegen_backward_subclass_fns
+
+            _codegen_bw_unwrap_fn, _codegen_bw_wrap_fn = codegen_backward_subclass_fns(
+                grad_input_metas=maybe_subclass_meta.grad_input_metas,
+            )
 
         class CompiledFunction(torch.autograd.Function):
             compiled_fw = compiled_fw_func
             compiled_bw = compiled_bw_func
             metadata: ViewAndMutationMeta = fw_metadata  # type: ignore[assignment]
-            maybe_subclass_metadata: Optional[SubclassMeta] = maybe_subclass_meta
+            maybe_subclass_metadata: SubclassMeta | None = maybe_subclass_meta
             num_symints_saved_for_bw = num_symints_saved_for_bw_
             _aot_id = aot_config.aot_id
             _lazy_backward_info = lazy_backward_info
+            _bw_epilogue_wrap_fn = _codegen_bw_wrap_fn
+            _bw_prologue_unwrap_fn = _codegen_bw_unwrap_fn
 
             @staticmethod
             def _compiled_autograd_key(ctx: Any) -> tuple[Any, ...]:
@@ -2470,28 +2832,7 @@ def forward(ctx: Any, *deduped_flat_tensor_args: Any) -> Any:
                         )
                     ctx._compiled_autograd_backward_state = bw_state
 
-                if num_rng:
-                    if len(fwd_rng_states) == 0:
-                        if graphsafe_idx is None:
-                            raise AssertionError(
-                                "graphsafe_idx must not be None when num_rng > 0"
-                            )
-                        initialize_rng_states(
-                            num_rng, graphsafe_idx, fwd_rng_states, bwd_rng_states
-                        )
-
-                    _curr_iter = next(curr_fwd_iter)
-                    ctx._curr_iter = _curr_iter
-
-                    # if this state is not contained in the backward,
-                    # we need to save it for when its backward pass happens
-                    if _curr_iter != backward_state_position:
-                        saved_backward_tensor_states[_curr_iter] = [
-                            rng_state.get_state() for rng_state in fwd_rng_states
-                        ]
-
-                    pending_forwards.add(_curr_iter)
-                    args = (*args, *fwd_rng_states)
+                args = rng_state.add_forward_args(ctx, args)
 
                 # There is a pretty complicated calling convention around what the compiled fw returns.
                 # The full list of outputs and their relative order is:
@@ -2507,228 +2848,21 @@ def forward(ctx: Any, *deduped_flat_tensor_args: Any) -> Any:
                     disable_amp=disable_amp,
                 )
 
-                num_outputs = CompiledFunction.metadata.num_outputs
-                num_outputs_aliased = CompiledFunction.metadata.num_outputs_aliased
-                num_mutated_runtime_inps = (
-                    CompiledFunction.metadata.num_mutated_inp_runtime_indices
-                )
-                num_forward_returns = CompiledFunction.metadata.num_forward_returns
-
-                # See Note [Activations with no version counter checks in eager]
-                # Partitioners must put symint arguments at the end separate from tensor arguments
-                # Split tensors into those that need VC checks (via save_for_backward)
-                # and those that don't (stashed directly on ctx).
-                # The partitioner sorts tensors so that no-VC-check tensors are at the end.
-                tensors_saved_with_vc_check = fw_outs[
-                    CompiledFunction.metadata.tensors_saved_for_backwards_with_vc_check_slice
-                ]
-                tensors_saved_no_vc_check = fw_outs[
-                    CompiledFunction.metadata.tensors_saved_for_backwards_no_vc_check_slice
-                ]
-                if not all(
-                    isinstance(x, torch.Tensor) for x in tensors_saved_with_vc_check
-                ):
-                    raise AssertionError(
-                        f"expected all tensors_saved_with_vc_check to be Tensors, "
-                        f"got types: {[type(x) for x in tensors_saved_with_vc_check]}"
-                    )
-                if not all(
-                    isinstance(x, torch.Tensor) for x in tensors_saved_no_vc_check
-                ):
-                    raise AssertionError(
-                        f"expected all tensors_saved_no_vc_check to be Tensors, "
-                        f"got types: {[type(x) for x in tensors_saved_no_vc_check]}"
-                    )
-
-                # See Note [Detaching saved tensors in AOTAutograd]
-                num_vc_check = len(tensors_saved_with_vc_check)
-                tensors_to_save = [
-                    x.detach() if x._is_view() else x
-                    for x in tensors_saved_with_vc_check
-                ]
-                tensors_no_vc = [
-                    x.detach() if x._is_view() else x for x in tensors_saved_no_vc_check
-                ]
-
-                # dynamic_saved_tensors_idxs has indices relative to all saved tensors
-                # (vc_check + no_vc_check combined). Mark dynamics on the detached tensors.
-                for (
-                    idx,
-                    dims,
-                ) in CompiledFunction.metadata.dynamic_saved_tensors_idxs.items():
-                    if idx < num_vc_check:
-                        maybe_mark_dynamic_helper(tensors_to_save[idx], dims)
-                    else:
-                        maybe_mark_dynamic_helper(
-                            tensors_no_vc[idx - num_vc_check], dims
-                        )
-
-                # Only save tensors that need VC checks via save_for_backward
-                ctx.save_for_backward(*tensors_to_save)
-                ctx._tensors_no_vc_check = tensors_no_vc
-
-                symint_outs = fw_outs[
-                    CompiledFunction.metadata.symints_saved_for_backwards_slice
-                ]
-                if not all(
-                    isinstance(x, (int, float, torch.SymInt, torch.SymFloat))
-                    for x in symint_outs
-                ):
-                    raise AssertionError(
-                        f"expected all symint_outs to be int/float/SymInt/SymFloat, "
-                        f"got types: {[type(x) for x in symint_outs]}"
-                    )
-                ctx.symints = symint_outs
-
-                opaque_object_outs = fw_outs[
-                    CompiledFunction.metadata.opaque_objects_saved_for_backwards_slice
-                ]
-                if not all(is_opaque_type(type(obj)) for obj in opaque_object_outs):
-                    raise AssertionError(
-                        f"expected all opaque_object_outs to be opaque types, "
-                        f"got types: {[type(obj) for obj in opaque_object_outs]}"
-                    )
-                ctx.opaque_objects = opaque_object_outs
-
-                raw_returns = fw_outs[0:num_forward_returns]
-
-                # Wrap all autograd.Function.forward() outputs that are aliases
-                # so that autograd.Function doesn't treat them as tensors
-                if num_mutated_runtime_inps > 0:
-                    for i, idx in enumerate(
-                        CompiledFunction.metadata.mutated_inp_runtime_indices
-                    ):
-                        # We could make this faster by only looping over inputs with metadata-only mutations
-                        # (instead of looping over inputs with either data or metadata mutations), but there shouldn't be many.
-                        info = CompiledFunction.metadata.input_info[idx]
-                        if info.mutates_metadata and not info.mutates_data:
-                            raw_return_idx = i
-                            raw_returns[raw_return_idx] = TensorAlias(
-                                raw_returns[raw_return_idx]
-                            )
-
-                    if config.debug_assert:
-                        user_mutated_inputs_raw = raw_returns[
-                            0:num_mutated_runtime_inps
-                        ]
-                        mut_inp_infos = [
-                            x
-                            for x in CompiledFunction.metadata.input_info
-                            if x.mutates_data or x.mutates_metadata
-                        ]
-                        if len(user_mutated_inputs_raw) != len(mut_inp_infos):
-                            raise AssertionError(
-                                f"expected len(user_mutated_inputs_raw) == len(mut_inp_infos), "
-                                f"got {len(user_mutated_inputs_raw)} != {len(mut_inp_infos)}"
-                            )
-
-                if CompiledFunction.metadata.num_unsafe_view_outputs > 0:
-                    for idx in CompiledFunction.metadata.unsafe_view_out_indices:
-                        raw_return_idx = num_mutated_runtime_inps + idx
-                        o = raw_returns[raw_return_idx]
-                        raw_returns[raw_return_idx] = torch.ops.aten._unsafe_view(
-                            o, o.shape
-                        )
-
-                if num_outputs_aliased > 0:
-                    for idx in CompiledFunction.metadata.aliased_out_indices:
-                        raw_return_idx = num_mutated_runtime_inps + idx
-                        raw_returns[raw_return_idx] = TensorAlias(
-                            raw_returns[raw_return_idx]
-                        )
-
-                    if config.debug_assert:
-                        intermediates_raw = raw_returns[
-                            num_mutated_runtime_inps + num_outputs :
-                        ]
-                        if any(isinstance(x, TensorAlias) for x in intermediates_raw):
-                            raise AssertionError(
-                                "expected no TensorAlias in intermediates_raw"
-                            )
-
-                # invariant: intermediate bases always require gradients, so we don't have to
-                # consider marking them as non-differentiable.
-                raw_returns_not_including_intermediate_bases = raw_returns[
-                    : num_mutated_runtime_inps + num_outputs
-                ]
-                raw_returns_meta = [
-                    x
-                    for x in CompiledFunction.metadata.input_info
-                    if x.mutation_type == MutationType.MUTATED_OUT_GRAPH
-                ] + CompiledFunction.metadata.output_info
-
-                fw_outs_not_requiring_grad = [
-                    x
-                    for (i, x) in enumerate(
-                        raw_returns_not_including_intermediate_bases
-                    )
-                    if isinstance(x, torch.Tensor)
-                    and not raw_returns_meta[i].requires_grad
-                ]
-                ctx.mark_non_differentiable(*fw_outs_not_requiring_grad)
-                ctx._materialize_non_diff_grads = False
-                return tuple(raw_returns)
+                saved_state.save_from_forward(ctx, fw_outs)
+                return forward_epilogue.finalize(ctx, fw_outs)
 
             @staticmethod
             def backward(ctx: Any, *flat_args: Any) -> tuple[Any, ...]:
-                # Combine tensors from both sources:
-                # 1. ctx.saved_tensors - tensors that went through save_for_backward (with VC check)
-                # 2. ctx._tensors_no_vc_check - tensors stashed directly on ctx (no VC check)
                 all_args = _backward_prologue_functional(
-                    (
-                        list(ctx.saved_tensors) + ctx._tensors_no_vc_check
-                        if len(ctx._tensors_no_vc_check) > 0
-                        else ctx.saved_tensors
-                    ),
+                    saved_state.load_tensors(ctx),
                     ctx.symints,
                     ctx.opaque_objects,
                     CompiledFunction.metadata,
                     CompiledFunction.maybe_subclass_metadata,
                     *flat_args,
+                    codegen_unwrap_fn=CompiledFunction._bw_prologue_unwrap_fn,
                 )
-
-                if num_rng:
-                    nonlocal backward_state_position, bwd_rng_states
-                    curr_backward_iter = ctx._curr_iter
-                    retain_graph = (
-                        torch._C._autograd._get_current_graph_task_keep_graph()
-                    )
-
-                    # Save current state if we have a pending forward that needs this state
-                    # or this state may be needed again because of retain graph
-                    if (
-                        backward_state_position in pending_forwards
-                        and backward_state_position not in saved_backward_tensor_states
-                        and (
-                            backward_state_position != curr_backward_iter
-                            or retain_graph
-                        )
-                    ):
-                        saved_backward_tensor_states[backward_state_position] = [
-                            rng_state.get_state() for rng_state in bwd_rng_states
-                        ]
-
-                    # Restore saved states if needed
-                    if curr_backward_iter in saved_backward_tensor_states:
-                        if backward_state_position != curr_backward_iter:
-                            for bwd_state, saved_state in zip(
-                                bwd_rng_states,
-                                saved_backward_tensor_states[curr_backward_iter],
-                            ):
-                                bwd_state.set_state(saved_state)
-                        if not retain_graph:
-                            del saved_backward_tensor_states[curr_backward_iter]
-                    else:
-                        if backward_state_position != curr_backward_iter:
-                            raise AssertionError(
-                                f"expected backward_state_position == curr_backward_iter, "
-                                f"got {backward_state_position} != {curr_backward_iter}"
-                            )
-
-                    backward_state_position = curr_backward_iter + 1
-                    if not retain_graph:
-                        pending_forwards.remove(curr_backward_iter)
-                    all_args.extend(bwd_rng_states)
+                rng_state.add_backward_args(ctx, all_args)
 
                 def impl_fn(double_ctx: Any = None) -> Any:
                     out = CompiledFunction._backward_impl(ctx, all_args)
@@ -2736,16 +2870,22 @@ def impl_fn(double_ctx: Any = None) -> Any:
                         CompiledFunction.metadata,
                         CompiledFunction.maybe_subclass_metadata,
                         out,
+                        codegen_wrap_fn=CompiledFunction._bw_epilogue_wrap_fn,
                     )
 
+                if (
+                    torch._C._is_key_in_tls("context")
+                    and (config_ctx := torch._C._get_obj_in_tls("context")) is not None
+                ):
+                    impl_fn = functools.partial(config_ctx.run, impl_fn)
+
                 needs_grad = torch.is_grad_enabled() and any(
                     t.requires_grad for t in all_args if isinstance(t, torch.Tensor)
                 )
                 if needs_grad:
                     # double backward
                     return CompiledFunction._double_backward(ctx, impl_fn, all_args)
-                else:
-                    return impl_fn()
+                return impl_fn()
 
             @staticmethod
             def _double_backward(
@@ -2785,118 +2925,10 @@ def _backward_impl(ctx: Any, all_args: list[Any]) -> Any:
                 saved_tensors_use_once = (
                     not torch._C._autograd._get_current_graph_task_keep_graph()
                 )
-
-                if CompiledFunction.compiled_bw is None:
-                    if lazy_backward_info is None:
-                        raise AssertionError("lazy_backward_info must not be None")
-                    if not isinstance(
-                        lazy_backward_info, AutogradLazyBackwardCompileInfo
-                    ):
-                        raise AssertionError(
-                            f"expected AutogradLazyBackwardCompileInfo, got {type(lazy_backward_info)}"
-                        )
-
-                    if (
-                        hasattr(lazy_backward_info, "saved_context")
-                        and lazy_backward_info.saved_context is not None
-                    ):
-                        if not isinstance(
-                            lazy_backward_info.saved_context, TracingContext
-                        ):
-                            raise AssertionError(
-                                f"expected TracingContext, got {type(lazy_backward_info.saved_context)}"
-                            )
-                        ddp_ctx = lazy_backward_info.saved_context.ddp_optimizer_ctx
-                        if ddp_ctx is not None:
-                            if ddp_ctx.curr_bucket < 0:
-                                raise AssertionError(
-                                    f"expected same # of fw and bw compiles, but found bucket {ddp_ctx.curr_bucket}"
-                                )
-                            curr_fw_meta = ddp_ctx.metadata_per_bucket[
-                                ddp_ctx.curr_bucket
-                            ]
-                            # Note [DDPOptimizer and fw_metadata]
-                            # When using the DDPOptimizer, we have a single dynamo graph (and TracingContext),
-                            # but multiple AOTDispatcher graph.
-                            #
-                            # One consequence is that there will be **multiple** fw_metadata objects, one per AOT graph,
-                            # which we stash the fw_metadata on the TracingContext.
-                            #
-                            # Normally what happens is that as we compile AOT graphs 1...N, we clobber the fw_metadata
-                            # for graph i-1 when we start running AOT for graph i.
-                            # Ordinarily this is fine, because inductor no longer needs the metadata from graph i-1.
-                            #
-                            # However, this is a problem for lazy compilation of the backward. During backward compilation,
-                            # we compile the backward lazily at backward runtime, meaning that we will first compile
-                            # backward graph N, N-1, ..., 1.
-                            # We need to ensure that at the time inductor compiles bw graph N-1, it can access
-                            # the corresponding fw_metadta for graph N-1.
-                            #
-                            # We do this by stashing a DDPOptimizerContext, which tracks:
-                            # - the metadata of all N graphs
-                            # - the graph we are currently compiling in our DDPOptimizer region.
-                            ddp_ctx.curr_bucket -= 1
-                            lazy_backward_info.saved_context.fw_metadata = curr_fw_meta
-
-                    if not saved_tensors_use_once:
-                        fw_metadata.bw_donated_idxs = []
-                        # Update bw_donated_idxs if using lazy_backward_info from `aot_dispatch_autograd`
-                        if (
-                            hasattr(lazy_backward_info, "saved_context")
-                            and hasattr(lazy_backward_info.saved_context, "fw_metadata")
-                            and hasattr(
-                                lazy_backward_info.saved_context.fw_metadata,  # type: ignore[union-attr]
-                                "bw_donated_idxs",
-                            )
-                        ):
-                            lazy_backward_info.saved_context.fw_metadata.bw_donated_idxs = (  # type: ignore[union-attr]
-                                # pyrefly: ignore [implicit-any]
-                                []
-                            )
-
-                    bw_module = lazy_backward_info.bw_module
-                    placeholder_list = lazy_backward_info.placeholder_list
-                    saved_context = lazy_backward_info.saved_context
-                    saved_compile_context = lazy_backward_info.saved_compile_context
-
-                    context = torch._C._DisableAutocast if disable_amp else nullcontext
-                    metrics_context = get_metrics_context()
-                    with (
-                        tracing(saved_context),
-                        compile_context(saved_compile_context),
-                        context(),
-                        track_graph_compiling(aot_config, "backward"),
-                        metrics_context,
-                        dynamo_timed(
-                            "backward._backward_impl",
-                            phase_name="entire_backward_compile",
-                            log_pt2_compile_event=True,
-                            dynamo_compile_column_us="backward_cumulative_compile_time_us",
-                            log_waitcounter=True,
-                            waitcounter_name_override="entire_backward_compile",
-                        ),
-                        callback_handler.install_callbacks(
-                            CallbackTrigger.LAZY_BACKWARD,
-                            str(CompileContext.current_compile_id()),
-                        ),
-                    ):
-                        CompileEventLogger.compilation_metric(is_forward=False)
-                        # See Note: [Backward graph lazy lowering]
-                        if aot_config.bw_compiler is None:
-                            raise AssertionError(
-                                "aot_config.bw_compiler must not be None"
-                            )
-                        CompiledFunction.compiled_bw = aot_config.bw_compiler(
-                            copy.deepcopy(bw_module), placeholder_list
-                        )
-                        # Maybe save cache entry
-                        if try_save_cache_entry is not None:
-                            try_save_cache_entry(
-                                CompiledFunction.compiled_bw,
-                                bw_module,
-                                fw_metadata,
-                                aot_config,
-                            )
+                compiled_bw = backward_compiler.get_or_compile(
+                    saved_tensors_use_once=saved_tensors_use_once
+                )
+                CompiledFunction.compiled_bw = compiled_bw
 
                 if (
                     torch._functorch.config.donated_buffer
@@ -2915,26 +2947,212 @@ def _backward_impl(ctx: Any, all_args: list[Any]) -> Any:
                         ),
                     )
 
-                out = call_func_at_runtime_with_args(
-                    CompiledFunction.compiled_bw,
+                return call_func_at_runtime_with_args(
+                    compiled_bw,
                     all_args,
                     steal_args=True,
                     disable_amp=disable_amp,
                 )
-                return out
 
-        compiled_function = RuntimeWrapper(
-            indices_of_inps_to_detach=indices_of_inps_to_detach,
+        return CompiledFunction
+
+
+# This is wrapped in a class just for namespacing purposes
+# No need to make it into an actual CompilerWrapper because it doesn't fit the abstract as cleanly
+class AOTDispatchAutograd:
+    @staticmethod
+    def _raise_tangent_metadata_error(
+        expected_type: type | None,
+        expected_meta: Any,
+        runtime_type: type,
+        runtime_meta: Any,
+        orig_x: torch.Tensor,
+        tangent_idx: int | None,
+        tangent_desc: Any | None,
+        compile_id_str: str | None,
+        tangent_stack_trace: str | None,
+    ) -> RuntimeError:
+        expected_subclass_got_plain_tensor = (
+            expected_type is not None
+            and expected_type is not torch.Tensor
+            and runtime_type is torch.Tensor
+        )
+        if expected_subclass_got_plain_tensor:
+            tangent_msg = ""
+            if tangent_idx is not None:
+                tangent_msg = f" (tangent index: {tangent_idx})"
+
+            output_hint = ""
+            if tangent_desc is not None:
+                from .descriptors import PlainAOTOutput, TangentAOTInput
+
+                if isinstance(tangent_desc, TangentAOTInput) and isinstance(
+                    tangent_desc.output, PlainAOTOutput
+                ):
+                    idx = tangent_desc.output.idx
+                    output_hint = f"\n\nThe problematic output is: forward output at index {idx} (0-indexed)"
+                else:
+                    output_hint = (
+                        f"\n\nThe problematic output is: {tangent_desc.expr()}"
+                    )
+
+            graph_hint = ""
+            if compile_id_str is not None:
+                graph_hint = (
+                    f"\n\nThis error occurred in compiled graph [{compile_id_str}]."
+                )
+
+            stack_trace_hint = ""
+            if tangent_stack_trace is not None:
+                stack_trace_hint = (
+                    f"\n\nThe forward output was created here:\n{tangent_stack_trace}"
+                )
+
+            return RuntimeError(
+                f"""
+During the backward, we encountered a tensor subclass where we guessed its
+metadata incorrectly.
+Expected a {expected_type.__name__} tangent but got a plain Tensor{tangent_msg}.
+This happens when a compiled function returns multiple outputs that
+require gradients, but .backward() is only called on some of them.
+To fix: call .detach() on forward outputs you don't need gradients for.{output_hint}{graph_hint}{stack_trace_hint}
+
+This error is also more likely to occur if your compiled model is suffering
+from a large number of graph breaks. For more advice on finding and fixing
+graph breaks, see:
+https://docs.pytorch.org/docs/stable/user_guide/torch_compiler/compile/programming_model.graph_breaks_index.html
+
+For more info about this error, see:
+https://github.com/pytorch/pytorch/issues/172556"""
+            )
+        else:
+            return RuntimeError(
+                f"""
+During the backward, we encountered a tensor subclass where we guessed its
+metadata incorrectly.
+Expected: {expected_meta} (type {expected_type}),
+got: {runtime_meta} (type {runtime_type}), shape: {orig_x.shape}.
+Your tensor subclass must implement __coerce_same_metadata_as_tangent__."""
+            )
+
+    @staticmethod
+    def process_runtime_tangent(
+        x: Any,
+        meta: PlainTensorMeta | SubclassCreationMeta,
+        tangent_idx: int | None = None,
+        tangent_desc: Any | None = None,
+        compile_id_str: str | None = None,
+        tangent_stack_trace: str | None = None,
+    ) -> tuple[Any, list[Any]]:
+        if not isinstance(x, torch.Tensor):
+            return x, [x]
+
+        if isinstance(x, FakeTensor):
+            if not meta.memory_format:
+                raise AssertionError(
+                    "meta.memory_format must not be None for FakeTensor"
+                )
+            x = coerce_to_expected_memory_format(x, meta.memory_format)
+            return x, [x]
+
+        expected_type: type | None = torch.Tensor
+        expected_meta = None
+        if isinstance(meta, SubclassCreationMeta):
+            expected_type = meta.original_subclass_type
+            expected_meta = meta.meta
+
+        runtime_type = type(x)
+        # When we're inside compiled autograd's AOTDispatcher step,
+        # regular Tensors look like FunctionalTensors.
+        # Tensor subclasses still look like Tensor subclasses though.
+        if isinstance(x, torch._subclasses.functional_tensor.FunctionalTensor):
+            runtime_type = torch.Tensor
+
+        runtime_meta = None
+        runtime_subclass_keys: Sequence[str] = []
+
+        if is_traceable_wrapper_subclass(x):
+            runtime_subclass_keys, runtime_meta = x.__tensor_flatten__()
+
+        def maybe_coerce(x: torch.Tensor) -> torch.Tensor | None:
+            same_type: bool = expected_type == runtime_type
+            same_meta: bool = expected_meta == runtime_meta
+
+            if same_type and same_meta:
+                return x
+
+            if not hasattr(x, "__coerce_same_metadata_as_tangent__"):
+                return None
+
+            if same_type:
+                # Backward Compatibility, as some Subclass impls can have original 1-arg function.
+                return x.__coerce_same_metadata_as_tangent__(expected_meta)
+
+            return x.__coerce_same_metadata_as_tangent__(expected_meta, expected_type)
+
+        # Coerce to expected type and metadata
+        orig_x = x
+        x = maybe_coerce(x)
+        if x is None:
+            raise AOTDispatchAutograd._raise_tangent_metadata_error(
+                expected_type,
+                expected_meta,
+                runtime_type,
+                runtime_meta,
+                orig_x,
+                tangent_idx,
+                tangent_desc,
+                compile_id_str,
+                tangent_stack_trace,
+            )
+
+        # Coerce to expected memory format
+        if not meta.memory_format:
+            raise AssertionError("meta.memory_format must not be None")
+        x = coerce_to_expected_memory_format(x, meta.memory_format)
+
+        if not is_traceable_wrapper_subclass(x):
+            return x, [x]
+
+        if not isinstance(meta, SubclassCreationMeta):
+            raise AssertionError(f"expected SubclassCreationMeta, got {type(meta)}")
+        if orig_x is not x:
+            runtime_subclass_keys = x.__tensor_flatten__()[0]
+
+        if len(meta.attrs) != len(runtime_subclass_keys):
+            raise AssertionError(
+                f"expected len(meta.attrs) == len(runtime_subclass_keys), "
+                f"got {len(meta.attrs)} != {len(runtime_subclass_keys)}"
+            )
+        leaves = []
+        for attr, attr_meta in meta.attrs.items():
+            if isinstance(attr_meta, OpaqueMeta):
+                # Opaques aren't differentiable but occupy a flat arg slot.
+                leaves.append(getattr(x, attr))
+                continue
+            elem = getattr(x, attr)
+            new_elem, elem_leaves = AOTDispatchAutograd.process_runtime_tangent(
+                elem, attr_meta
+            )
+            if new_elem is not elem:
+                setattr(x, attr, new_elem)
+            leaves.extend(elem_leaves)
+
+        return x, leaves
+
+    @staticmethod
+    def post_compile(spec: AOTDispatchAutogradCompileSpec) -> Callable[..., Any]:
+        compiled_function_cls = _AOTDispatchAutogradFunctionFactory(spec).build()
+        return RuntimeWrapper(
+            indices_of_inps_to_detach=spec.indices_of_inps_to_detach,
             trace_joint=True,
-            disable_amp=disable_amp,
+            disable_amp=spec.disable_amp,
         ).post_compile(
-            CompiledFunction.apply,
-            aot_config,
-            runtime_metadata=fw_metadata,
+            compiled_function_cls.apply,
+            spec.aot_config,
+            runtime_metadata=spec.fw_metadata,
         )
 
-        return compiled_function
-
 
 @dataclass
 class DebugAssertWrapper(CompilerWrapper):
@@ -2947,36 +3165,38 @@ def post_compile(
         *,
         runtime_metadata: ViewAndMutationMeta,
     ) -> Callable[..., Any]:
-        @wraps(compiled_fn)
-        def debug_compiled_function(args: list[Any]) -> Any:
-            # TODO: T253242027 Check aliasing relationships
-            # TODO: Check strides for metadata mutation
-            # (NB: ideally, this logic is factored out of this function and
-            # you move these debug checks there)
-
-            # Check requires grad.  Bad case is when we compiled with
-            # requires_grad = False, but input requires_grad = True
-            # (vice versa is OK; we compute a gradient and then throw
-            # it away when it hits the input.)
-            for i, a in enumerate(args):
-                can_require_grad = self.flat_requires_grad[i]
-                if can_require_grad is None:
-                    if isinstance(a, Tensor):
-                        raise AssertionError(
-                            f"expected non-Tensor for arg {i}, got Tensor"
-                        )
-                elif not can_require_grad:
-                    if a.requires_grad:
-                        raise AssertionError(
-                            format_guard_bug_msg(
-                                aot_config,
-                                f"{describe_input(i, aot_config)} would not require grad",
-                            )
-                        )
+        lines = ["def inner_fn(args):"]
+        globals_dict: dict[str, object] = {"compiled_fn": compiled_fn}
+        for i, can_require_grad in enumerate(self.flat_requires_grad):
+            if can_require_grad is None:
+                lines.append(
+                    f"    if isinstance(args[{i}], Tensor):"
+                    f" raise AssertionError("
+                    f"'expected non-Tensor for arg {i}, got Tensor')"
+                )
+            elif not can_require_grad:
+                msg_name = f"_msg_{i}"
+                globals_dict[msg_name] = format_guard_bug_msg(
+                    aot_config,
+                    f"{describe_input(i, aot_config)} would not require grad",
+                )
+                lines.append(
+                    f"    if args[{i}].requires_grad: raise AssertionError({msg_name})"
+                )
+        lines.append("    return compiled_fn(args)")
 
-            return compiled_fn(args)
+        source = "\n".join(lines)
+        globals_dict["Tensor"] = Tensor
 
-        return debug_compiled_function
+        from .subclass_codegen import _compile_and_exec_source
+
+        return _compile_and_exec_source(
+            source,
+            globals_dict,
+            "inner_fn",
+            "debug_assert_wrapper",
+            wrapped_fn=compiled_fn,
+        )
 
 
 def pre_compile(
diff --git a/torch/_functorch/_aot_autograd/schemas.py b/torch/_functorch/_aot_autograd/schemas.py
index c630fa520b726..809f389b3ed56 100644
--- a/torch/_functorch/_aot_autograd/schemas.py
+++ b/torch/_functorch/_aot_autograd/schemas.py
@@ -9,13 +9,13 @@
 import functools
 from dataclasses import dataclass, field, replace
 from enum import Enum
-from typing import Any, NewType, Optional, Protocol, TYPE_CHECKING, TypeVar, Union
+from typing import Any, NewType, Protocol, TYPE_CHECKING, TypeVar
 from typing_extensions import ParamSpec
 
 import torch
 import torch.utils._pytree as pytree
 from torch import SymInt, Tensor
-from torch._library.opaque_object import OpaqueType
+from torch._opaque_base import OpaqueBase
 from torch._subclasses import FakeTensor, FakeTensorMode
 from torch._subclasses.fake_tensor import is_fake
 from torch.fx.experimental._backward_state import BackwardState
@@ -101,11 +101,17 @@ class OutputAliasInfo:
     # If (6) above, then:
     # - Tells us that the base of this alias is output_user_fwds[base_idx]
     #   here, this refers to the index of the *direct* traced
-    base_idx: Optional[int]
+    base_idx: int | None
     # If it is a Tensor, what the dynamic dims are (otherwise is None)
-    dynamic_dims: Optional[set[int]]
-    # requires_grad
+    dynamic_dims: set[int] | None
+    # Whether the output tensor has requires_grad=True. Use for
+    # preserving output properties (mark_non_differentiable, alias reconstruction).
     requires_grad: bool
+    # Whether this output participates in backward as
+    # a tangent. Normally the same as requires_grad, but False for views created
+    # under no_grad() that inherit requires_grad from their base without having a
+    # grad_fn. Use when constructing tangent lists for torch.autograd.grad().
+    requires_grad_for_backward: bool
     # Sequence of ViewMeta objects.
     #
     # Provides us the means to re-run view functions on other tensors.
@@ -113,7 +119,7 @@ class OutputAliasInfo:
     # We need to wrap the actual list of ViewMeta with this class so that
     # we compare the ViewMeta elements appropriately, i.e. their type and
     # the elements returned by the `as_tuple()` call.
-    view_meta_sequence: Optional[ViewMetaSequence] = None
+    view_meta_sequence: ViewMetaSequence | None = None
 
 
 class MutationType(Enum):
@@ -173,21 +179,25 @@ def mutation_type(self) -> MutationType:
 @dataclass
 class MemoryFormatMeta:
     # For static shapes we assume tangents have the same strideness as outputs
-    size: Optional[Sequence[int]] = None
-    stride: Optional[Sequence[int]] = None
+    size: Sequence[int] | None = None
+    stride: Sequence[int] | None = None
 
     # For dynamic shapes we assume the same memory format: contiguous, channels_last etc.
-    memory_format: Optional[torch.memory_format] = None
+    memory_format: torch.memory_format | None = None
 
     @staticmethod
-    def from_tensor(t: torch.Tensor) -> Optional[MemoryFormatMeta]:
+    def from_tensor(
+        t: torch.Tensor, force_use_memory_format: bool = False
+    ) -> MemoryFormatMeta | None:
         # We only memorize expected memory format for
         # 1. Traceable wrapper subclasses
         # We can not create restrided subclass tensor, as torch.empty_strided works only with dense tensors.
         # 2. Dynamic shape tensors
         # Support for symbolic shapes is not implemented yet.
+        # 3. force_use_memory_format=True (e.g., local_map where shapes change)
         use_memory_format: bool = (
-            not torch._functorch.config.guess_tangent_strides_as_outputs
+            force_use_memory_format
+            or not torch._functorch.config.guess_tangent_strides_as_outputs
             or is_traceable_wrapper_subclass(t)
         )
         if not use_memory_format:
@@ -208,7 +218,12 @@ def from_tensor(t: torch.Tensor) -> Optional[MemoryFormatMeta]:
 @dataclass
 class PlainTensorMeta:
     unwrapped_idx: int
-    memory_format: Optional[MemoryFormatMeta] = None
+    memory_format: MemoryFormatMeta | None = None
+
+
+@dataclass
+class OpaqueMeta:
+    pass
 
 
 @dataclass
@@ -241,7 +256,7 @@ class SubclassCreationMeta:
     # meta and attrs are produced by the subclass's __tensor_flatten__.
     # We need to keep them around along with outer_size / outer_stride to plumb them
     # into __tensor_unflatten__
-    attrs: dict[str, SubclassCreationMeta | PlainTensorMeta]
+    attrs: dict[str, SubclassCreationMeta | PlainTensorMeta | OpaqueMeta]
     outer_size: Iterable[IntLikeType | None]
     outer_stride: Iterable[IntLikeType | None]
     meta: Any
@@ -250,15 +265,15 @@ class SubclassCreationMeta:
     # (this is guaranteed to be a wrapper subclass that holds a fake tensor,
     #  so holding onto this at runtime shouldn't leak memory)
     # This field is nulled out after calling make_runtime_safe()
-    original_subclass: Optional[torch.Tensor]
+    original_subclass: torch.Tensor | None
 
     # Used at runtime to determine the subclass type, so we don't need to save the original subclass
-    original_subclass_type: Optional[type] = None
-    memory_format: Optional[MemoryFormatMeta] = None
+    original_subclass_type: type | None = None
+    memory_format: MemoryFormatMeta | None = None
 
     def compute_outer_size_and_stride(
         self,
-        all_args: Sequence[torch.Tensor | IntLikeType],
+        all_args: Sequence[torch.Tensor | IntLikeType | OpaqueBase],
         *,
         curr_start_idx: int,
     ) -> tuple[
@@ -274,13 +289,12 @@ def compute(
             has_symbolic = any(placeholders)
 
             if has_symbolic:
-                start = curr_start_idx
                 end = start_idx + sum(placeholders)
-                it_args = iter(all_args[start:end])
+                it_args = iter(all_args[start_idx:end])
                 it_placeholders = iter(placeholders)
                 return pytree.tree_map_only(
                     lambda _: next(it_placeholders), lambda _: next(it_args), outer
-                ), start + len(placeholders)
+                ), end
             else:
                 return outer, start_idx
 
@@ -290,16 +304,25 @@ def compute(
 
     def creation_fn(
         self,
-        all_args: Sequence[torch.Tensor | IntLikeType],
+        all_args: Sequence[torch.Tensor | IntLikeType | OpaqueBase],
         *,
         is_runtime: bool,
     ) -> torch.Tensor:
-        inner_tensors = {}
+        inner_tensors: dict[str, torch.Tensor | OpaqueBase] = {}
 
         curr_start_idx = self.flat_tensor_start_idx
         for attr, creation_meta in self.attrs.items():
+            if isinstance(creation_meta, OpaqueMeta):
+                opaque = all_args[curr_start_idx]
+                if not isinstance(opaque, OpaqueBase):
+                    raise AssertionError(f"OpaqueBase expected, got {type(opaque)}")
+                inner_tensors[attr] = opaque
+                curr_start_idx += 1
+                continue
             if isinstance(creation_meta, PlainTensorMeta):
                 subclass = all_args[curr_start_idx]
+                if not isinstance(subclass, Tensor):
+                    raise AssertionError("Tensor expected")
                 curr_start_idx += 1
             else:
                 subclass = creation_meta.creation_fn(
@@ -340,7 +363,7 @@ def creation_fn(
         return rebuilt
 
     def make_runtime_safe(self) -> None:
-        def _make_size_runtime_safe(x: Union[None, int, torch.SymInt]) -> Optional[int]:
+        def _make_size_runtime_safe(x: None | int | torch.SymInt) -> int | None:
             dummy = -1
             if isinstance(x, torch.SymInt):
                 # Replace nested ints by a dummy value (-1) as NJT ignores
@@ -425,7 +448,7 @@ class ViewAndMutationMeta:
     #      inputs[3] and inputs[4] of the plain-tensor graph".
 
     # length = # user inputs
-    subclass_inp_meta: list[Union[PlainTensorMeta, SubclassCreationMeta]]
+    subclass_inp_meta: list[PlainTensorMeta | SubclassCreationMeta]
     # So, the full set of outputs to the forward graph looks something like:
     # (*mutated_inps, *user_outs, *intermediate_bases, *saved_for_bw_tensors)
     # where the first 3 of those 4 can be subclasses
@@ -433,47 +456,49 @@ class ViewAndMutationMeta:
     # and not user visible, so there's no point in wrapping/unwrapping them at runtime).
     # This list contains subclass information on all of the fw graph outputs
     # except for saved_for_bw_tensors.
-    subclass_fw_graph_out_meta: list[Union[PlainTensorMeta, SubclassCreationMeta]]
+    subclass_fw_graph_out_meta: list[PlainTensorMeta | SubclassCreationMeta]
     # length = # backward graph inputs
-    subclass_tangent_meta: list[Union[PlainTensorMeta, SubclassCreationMeta]]
-    # TODO: we should kill this
-    # (need to default it to not break internal)
-    is_train: bool = False
+    subclass_tangent_meta: list[PlainTensorMeta | SubclassCreationMeta]
 
     # length = (# inputs w data mutations) + (# user outputs that are non_aliasing tensors)
     #        + (# intermediate bases)
     # At runtime, we don't keep the traced_tangents around since they're not serializable.
     # Instead, we keep any necessary subclass metadata necessary about each traced_tangent.
     # This list is generated after calling make_runtime_safe().
-    traced_tangent_metas: Optional[list[Any]] = None
+    traced_tangent_metas: list[Any] | None = None
 
-    num_symints_saved_for_bw: Optional[int] = None
+    num_symints_saved_for_bw: int | None = None
 
     # See Note [Activations with no version counter checks in eager]
     # Number of tensors saved for backward that were stashed on ctx (e.g., ctx.x = x)
     # rather than via save_for_backward in an autograd.Function.
     # These tensors are placed at the end of the saved tensors and should skip
     # version counter checks at runtime.
-    num_tensors_saved_with_no_vc_check: Optional[int] = None
+    num_tensors_saved_with_no_vc_check: int | None = None
 
     # Number of opaque objects saved for backward
-    num_opaque_objects_saved_for_bw: Optional[int] = None
+    num_opaque_objects_saved_for_bw: int | None = None
     # The grad_enabled mutation that will be emitted in the runtime_wrapper epilogue
     # NOTE: AOTAutograd will assume that the ambient `is_grad_enabled` is the grad mode
     # that is intended to be in effect prior to running the graph, in keeping with
     # equivalence to eager mode. It is the responsibility of upstream graph acquisition
     # to reset the grad mode to its pre-graph value prior to calling aot_autograd.
-    grad_enabled_mutation: Optional[bool] = None
+    grad_enabled_mutation: bool | None = None
 
     # Keeps track of whether `torch.use_deterministic_algorithms` was turned on
     # when the forward was run. If deterministic mode was turned off during the
     # forward, but is turned on during the backward call, then an error is
     # raised
-    deterministic: Optional[bool] = None
+    deterministic: bool | None = None
 
     # Keeps track of which input indices store parameters (which we will treat as static)
     static_input_indices: list[int] = field(default_factory=list)
 
+    # Input indices that held AsyncCollectiveTensors at compile time.
+    # Used to emit direct trigger_wait() calls at runtime instead of
+    # scanning every arg on every graph invocation.
+    act_input_indices: list[int] = field(default_factory=list)
+
     # Map of effect type (ex. _EffectType.ORDERED) to token.  If there are
     # side-effectful operators, FunctionalTensorMode will populate this
     # dictionary telling us how many tokens we will need during tracing.
@@ -492,7 +517,7 @@ class ViewAndMutationMeta:
     # Indexes of saved tensors which are donated buffer.
     # Donated buffer means the tensor is not alias of any forward user input, forward user output,
     # and backward output.
-    bw_donated_idxs: Optional[list[int]] = None
+    bw_donated_idxs: list[int] | None = None
 
     # Number of tokens used in backward, appended at the end of backward outputs.
     # Filled after tracing joint function.
@@ -502,14 +527,20 @@ class ViewAndMutationMeta:
     # cudagraph compatible run_and_save_rng
     num_graphsafe_rng_states: int = 0
 
-    graphsafe_rng_state_index: Optional[int] = None
+    graphsafe_rng_state_index: int | None = None
 
     # Stream indices for mutated inputs in the epilogue
     # Maps from index in mutated_inp_runtime_indices to the stream index that last touched
     # the storage of the tensor that will be copied back into the original input
     # None means use the current/default stream
     # This is populated during graph compilation when stream assignments are made
-    mutated_inp_stream_indices: Optional[list[Optional[int]]] = None
+    mutated_inp_stream_indices: list[int | None] | None = None
+
+    # compile ID string (e.g., "1/0") for error messages
+    compile_id_str: str | None = None
+
+    # help users identify where to add .detach() in their code
+    tangent_source_stack_traces: list[str | None] | None = None
 
     def __post_init__(self) -> None:
         # pre-compute the indices of the inputs that are mutated.
@@ -837,9 +868,7 @@ class SubclassMeta:
     # in case we made incorrect assumptions about the subclass-ness of our grad_outputs
     #
     # Optional field because we don't compute for inference graphs
-    grad_input_metas: Optional[list[Union[PlainTensorMeta, SubclassCreationMeta]]] = (
-        None
-    )
+    grad_input_metas: list[PlainTensorMeta | SubclassCreationMeta] | None = None
 
     def __init__(self) -> None:
         # The fields in this class get set after its construction.
@@ -920,7 +949,7 @@ class GraphSignature:
     in_spec: pytree.TreeSpec
     out_spec: pytree.TreeSpec
 
-    backward_signature: Optional[BackwardSignature]
+    backward_signature: BackwardSignature | None
 
     input_tokens: list[GraphInputName]
     output_tokens: list[GraphOutputName]
@@ -939,8 +968,8 @@ def from_tracing_metadata(
         num_user_inputs: int,
         num_user_outputs: int,
         trace_joint: bool,
-        loss_index: Optional[int],
-        backward_signature: Optional[BackwardSignature],
+        loss_index: int | None,
+        backward_signature: BackwardSignature | None,
     ) -> GraphSignature:
         graph_inputs = graph_input_names
         graph_outputs = graph_output_names
@@ -1078,20 +1107,20 @@ class AOTConfig:
     is_export: bool = False
     no_tangents: bool = False
     dynamic_shapes: bool = False
-    aot_autograd_arg_pos_to_source: Optional[list[Source]] = None
-    static_input_indices: Optional[list[int]] = None
+    aot_autograd_arg_pos_to_source: list[Source] | None = None
+    static_input_indices: list[int] | None = None
     inference_compiler: Callable[..., Any] | None = None
     enable_log: bool = True
     # this is always false outside of export.
     pre_dispatch: bool = False
     # Key to use for AOTAutogradCache
-    cache_info: Optional[AOTAutogradCacheInfo] = None
+    cache_info: AOTAutogradCacheInfo | None = None
     # If we should ignore the shape_env in the ambient tracing_context.
     # The net effect is that if dynamic shapes are on, we end up
     # specializing on example_inputs.
     # Used only by standalone_compile.
     ignore_shape_env: bool = False
-    precompile_backend_id: Optional[str] = None
+    precompile_backend_id: str | None = None
     force_non_lazy_backward_lowering: bool = False
     # This config makes sure to check certain things like
     # mutating input with req_grad in export joint tracing.
@@ -1184,7 +1213,7 @@ class AOTState:
     fake_mode: FakeTensorMode
 
 
-FxValue = Tensor | int | SymInt | BackwardState | OpaqueType
+FxValue = Tensor | int | SymInt | BackwardState | OpaqueBase
 
 
 class CompilerWrapper:
@@ -1339,11 +1368,9 @@ class AOTGraphCapture:  # Produced by aot_stage1_graph_capture
     # larger than the original flat_args as all tangents get inputs.  The
     # tuple organizes into primals and tangents.  When not autograd it's just
     # a plain list.
-    updated_flat_args: Union[list[Any], tuple[list[Any], list[Any]]]
+    updated_flat_args: list[Any] | tuple[list[Any], list[Any]]
 
-    updated_flat_args_descs: Union[
-        list[AOTInput], tuple[list[AOTInput], list[AOTInput]]
-    ]
+    updated_flat_args_descs: list[AOTInput] | tuple[list[AOTInput], list[AOTInput]]
 
     # Metadata about subclass inputs/outputs in the graph trace.
     maybe_subclass_meta: Any
@@ -1416,8 +1443,8 @@ class JointTraceFn(Protocol):
     def __call__(
         self, primals: list[FxValue], tangents: list[FxValue]
     ) -> tuple[
-        tuple[list[FxValue], list[Optional[Tensor]]],
-        tuple[list[AOTOutput], list[Optional[AOTOutput]]],
+        tuple[list[FxValue], list[Tensor | None]],
+        tuple[list[AOTOutput], list[AOTOutput | None]],
     ]: ...
 
 
diff --git a/torch/_functorch/_aot_autograd/streams.py b/torch/_functorch/_aot_autograd/streams.py
index 823efde534781..8e94bef9370de 100644
--- a/torch/_functorch/_aot_autograd/streams.py
+++ b/torch/_functorch/_aot_autograd/streams.py
@@ -1,10 +1,12 @@
-from typing import Any, Optional, TYPE_CHECKING, TypeAlias
+import operator
+from typing import Any, TYPE_CHECKING, TypeAlias
 
 import torch.fx
 import torch.fx.traceback
 import torch.utils._pytree as pytree
 from torch._dynamo.graph_utils import _get_flat_args
 from torch._dynamo.variables.streams import get_current_stream, new_event
+from torch.fx.node import map_arg
 from torch.utils._runtime_estimation import (
     _FLOAT_TYPES,
     _IGNORE_OPS,
@@ -24,6 +26,14 @@
 Node: TypeAlias = torch.fx.Node
 Graph: TypeAlias = torch.fx.Graph
 
+_SYNC_OPS = (
+    torch.ops.streams.record_event.default,
+    torch.ops.streams.wait_event.default,
+    torch.ops.streams.synchronize_event.default,
+    torch.ops.streams.synchronize_device.default,
+    torch.ops.streams.synchronize_stream.default,
+)
+
 
 def get_roofline_estimate(node: Node) -> float:
     if node.op != "call_function":
@@ -69,7 +79,7 @@ def get_device(node: Node) -> torch.device:
     return node.meta["val"].device
 
 
-def get_stream(node: Node) -> Optional[int]:
+def get_stream(node: Node) -> int | None:
     maybe_annotation = node.meta.get("custom", None)
     if maybe_annotation is not None:
         return node.meta["custom"].get("stream", None)
@@ -120,9 +130,9 @@ def insert_wait_event_before_node(graph: Graph, node: Node, event_ind: int) -> N
 
 
 def populate_stream_timeline(
-    stream_to_timeline: dict[Optional[int], IndexedDict[Node, float]],
+    stream_to_timeline: dict[int | None, IndexedDict[Node, float]],
     graph: Graph,
-    stream_index: Optional[int],
+    stream_index: int | None,
 ) -> IndexedDict[Node, float]:
     if stream_index not in stream_to_timeline:
         stream_to_timeline[stream_index] = IndexedDict()
@@ -149,7 +159,7 @@ def populate_stream_timeline(
 # we attempt to find the point which to deallocate based on the estimated timestamps.
 def handle_synced_deallocation(
     graph: Graph,
-    stream_to_exec_trace: dict[Optional[int], IndexedDict[Node, float]],
+    stream_to_exec_trace: dict[int | None, IndexedDict[Node, float]],
     node: Node,
     last_usage: Node,
 ) -> None:
@@ -278,7 +288,7 @@ def sync_deallocations(gm: torch.fx.GraphModule) -> None:
     # I think this is fine because you should have large tensors if you're using streams
     # although perhaps I could add a constant 10us per op ahead of the first stream op?
     # a trace of all the nodes running in a given stream
-    stream_to_exec_trace: dict[Optional[int], IndexedDict[Node, float]] = {}
+    stream_to_exec_trace: dict[int | None, IndexedDict[Node, float]] = {}
     for node in gm.graph.nodes:
         if node.op == "call_function" and is_bwd_node(node):
             allocating_stream = get_stream(node)
@@ -350,3 +360,237 @@ def populate_fw_metadata_with_stream_indices(
             stream_indices.append(None)
 
     fw_metadata.mutated_inp_stream_indices = stream_indices
+
+
+def _wrap_sync_node(
+    gm: torch.fx.GraphModule,
+    sync_node: Node,
+    deps_before_sync: list[Node],
+    visited: set[Node],
+) -> tuple[Node, list[Node]]:
+    """
+    Core logic: wrap a single sync node in control_deps.
+
+    Returns (control_deps_node, passthrough_getitems) where passthrough_getitems
+    are the getitem nodes that thread dependencies through the control_deps node.
+    ``visited`` is the set of nodes at or before the sync node in graph order,
+    used to distinguish pre-sync vs post-sync users.
+    """
+    from torch._inductor.fx_passes.control_dependencies import (
+        _create_subgraph_for_node,
+        control_deps,
+        get_subgraph_name,
+    )
+
+    graph = gm.graph
+
+    # Use dep.users to find deps with uses after the sync — avoids a forward walk.
+    deps_with_uses_after_sync = [
+        dep
+        for dep in deps_before_sync
+        if any(user not in visited for user in dep.users)
+    ]
+
+    # Create subgraph that executes sync and passes through only used dependencies
+    subgraph_module = _create_subgraph_for_node(
+        graph, sync_node, deps_with_uses_after_sync
+    )
+    subgraph_attr_name = get_subgraph_name(gm, sync_node.name)
+    setattr(gm, subgraph_attr_name, subgraph_module)
+
+    # Create control_deps call
+    # Note: sync nodes (record_event/wait_event) only take int args, no Node args.
+    with graph.inserting_before(sync_node):
+        get_subgraph = graph.get_attr(subgraph_attr_name)
+        control_deps_node = graph.call_function(
+            control_deps,
+            args=(
+                tuple(deps_before_sync),  # additional_deps (all deps for ordering)
+                get_subgraph,  # subgraph
+                *deps_with_uses_after_sync,  # only pass through deps that are used
+            ),
+            kwargs={},
+        )
+
+    # Mark newly created nodes as visited so subsequent syncs don't
+    # misclassify them as "after the sync" during replacement.
+    visited.add(get_subgraph)
+    visited.add(control_deps_node)
+
+    # The output is (sync_result, *deps_with_uses_after_sync)
+    # Create getitem nodes only for dependencies that have uses after sync
+    replacements: dict[Node, Node] = {}
+    with graph.inserting_after(control_deps_node):
+        for i, dep in enumerate(deps_with_uses_after_sync):
+            getitem_node = graph.call_function(
+                operator.getitem,
+                args=(control_deps_node, i + 1),  # +1 because index 0 is sync result
+            )
+            getitem_node.meta.update(dep.meta)
+            replacements[dep] = getitem_node
+            visited.add(getitem_node)
+
+    # Replace uses of dependencies that come after sync_node.
+    # Use map_arg to handle nested structures (e.g. output node's list args).
+    for dep, getitem_node in replacements.items():
+        for user in list(dep.users.keys()):
+            if user is control_deps_node:
+                continue
+            if user in visited:
+                continue
+            # Don't replace forward outputs in the output node — they belong
+            # to the forward partition and must not reference backward nodes.
+            if user.op == "output" and not is_bwd_node(dep):
+                continue
+
+            def _replace(n: Node) -> Node:
+                return getitem_node if n is dep else n
+
+            user.args = map_arg(user.args, _replace)
+            user.kwargs = map_arg(user.kwargs, _replace)
+
+    # Remove original sync node
+    sync_node.replace_all_uses_with(control_deps_node)
+    graph.erase_node(sync_node)
+    return control_deps_node, list(replacements.values())
+
+
+def wrap_all_sync_nodes_with_control_deps(gm: torch.fx.GraphModule) -> None:
+    """
+    Single-pass wrap of all sync nodes in control_deps.
+
+    Iterates through the graph once, accumulating per-stream node lists.
+    When a sync node is encountered, it is wrapped using the accumulated deps
+    for that stream, then the deps are reset to the control_deps node
+    (maintaining the ordering chain for subsequent syncs on the same stream).
+    """
+    graph = gm.graph
+    if len(graph.nodes) == 0:
+        raise RuntimeError("Expected a non-empty graph")
+    stream_to_nodes: dict[int | None, list[Node]] = {}
+    # Maps event_index -> control_deps node that wrapped its record_event,
+    # so the corresponding wait_event/synchronize_event can depend on the record.
+    event_to_ctrl: dict[int, Node] = {}
+    # Maps event_index -> getitem nodes threaded through record_event's control_deps,
+    # so synchronize_event can thread them through to subsequent ops.
+    event_to_passthrough: dict[int, list[Node]] = {}
+    # Maps event_index -> stream that the event was recorded on,
+    # so synchronize_event can infer its stream.
+    event_to_stream: dict[int, int | None] = {}
+    visited: set[Node] = set()
+    found_sync = False
+
+    # Walk the node linked-list manually so we can mutate the graph
+    # (wrapping sync nodes inserts/erases nodes) without losing our place.
+    node = next(iter(graph.nodes))
+    while node.op != "root":
+        next_node = node.next
+        visited.add(node)
+
+        if node.op == "call_function":
+            if node.target in _SYNC_OPS:
+                # synchronize_device and synchronize_stream block the CPU,
+                # so all subsequent kernel launches are host-ordered after
+                # them. Treat both as full barriers across all streams.
+                if node.target in (
+                    torch.ops.streams.synchronize_device.default,
+                    torch.ops.streams.synchronize_stream.default,
+                ):
+                    all_stream_deps: list[Node] = [
+                        n for nodes in stream_to_nodes.values() for n in nodes
+                    ]
+                    if all_stream_deps:
+                        found_sync = True
+                        _wrap_sync_node(gm, node, all_stream_deps, visited)
+                    stream_to_nodes.clear()
+                    node = next_node
+                    continue
+
+                event_index: int = node.args[0]  # type: ignore[assignment]
+
+                # synchronize_event blocks the CPU thread, so it acts
+                # as a barrier across all streams. Collect deps from every
+                # stream and reset them all afterward. If the event was
+                # recorded externally, thread the graph inputs through so
+                # that any post-sync uses depend on the synchronize.
+                if node.target is torch.ops.streams.synchronize_event.default:
+                    sync_stream: int | None = event_to_stream.get(event_index)
+                    all_stream_deps: list[Node] = [
+                        n for nodes in stream_to_nodes.values() for n in nodes
+                    ]
+                    if event_index not in event_to_stream:
+                        placeholders = [n for n in graph.nodes if n.op == "placeholder"]
+                        deps_before_sync = [*placeholders, *all_stream_deps]
+                    else:
+                        deps_before_sync = all_stream_deps
+                else:
+                    sync_stream = node.args[1]  # type: ignore[assignment]
+                    deps_before_sync = list(stream_to_nodes.get(sync_stream, ()))
+                    # Nodes without explicit stream annotation (custom.stream=None)
+                    # run on the current/default stream. Include them when the sync
+                    # op references a stream, since the unannotated nodes are
+                    # implicitly on that stream.
+                    if None in stream_to_nodes and sync_stream is not None:
+                        deps_before_sync.extend(stream_to_nodes[None])
+
+                # For wait_event and synchronize_event, add a cross-event
+                # dependency on the matching record_event's control_deps node
+                # so they cannot be reordered before the record.
+                if (
+                    node.target
+                    in (
+                        torch.ops.streams.wait_event.default,
+                        torch.ops.streams.synchronize_event.default,
+                    )
+                    and event_index in event_to_ctrl
+                ):
+                    deps_before_sync = [
+                        event_to_ctrl[event_index],
+                        *deps_before_sync,
+                    ]
+
+                # For synchronize_event, also include the getitem nodes
+                # threaded through record_event's control_deps. This ensures
+                # subsequent ops that depend on recorded values get rewired
+                # through synchronize_event.
+                if (
+                    node.target is torch.ops.streams.synchronize_event.default
+                    and event_index in event_to_passthrough
+                ):
+                    deps_before_sync = [
+                        *deps_before_sync,
+                        *event_to_passthrough[event_index],
+                    ]
+
+                if deps_before_sync:
+                    found_sync = True
+                    ctrl_node, passthrough = _wrap_sync_node(
+                        gm, node, deps_before_sync, visited
+                    )
+                else:
+                    ctrl_node = None
+                    passthrough: list[torch.fx.Node] = []
+
+                if node.target is torch.ops.streams.record_event.default:
+                    event_to_stream[event_index] = sync_stream
+                    if ctrl_node is not None:
+                        event_to_ctrl[event_index] = ctrl_node
+                    event_to_passthrough[event_index] = passthrough
+
+                # Reset: ops between this sync and the next will accumulate
+                # fresh. Ordering with prior ops is already enforced because
+                # their uses were rewired through getitems from control_deps.
+                if node.target is torch.ops.streams.synchronize_event.default:
+                    stream_to_nodes.clear()
+                else:
+                    stream_to_nodes[sync_stream] = []
+                    if None in stream_to_nodes:
+                        stream_to_nodes[None] = []
+            elif "val" in node.meta:
+                stream = get_stream(node)
+                stream_to_nodes.setdefault(stream, []).append(node)
+
+        node = next_node
+
+    if found_sync:
+        gm.recompile()
diff --git a/torch/_functorch/_aot_autograd/subclass_codegen.py b/torch/_functorch/_aot_autograd/subclass_codegen.py
new file mode 100644
index 0000000000000..f6a47790569ad
--- /dev/null
+++ b/torch/_functorch/_aot_autograd/subclass_codegen.py
@@ -0,0 +1,397 @@
+"""
+Codegen for AOTDispatchSubclassWrapper.
+
+Generates a Python function that replaces the data-driven
+runtime_unwrap_tensor_subclasses / wrap_tensor_subclasses loop with
+a straight-line function where all metadata (indices, attr names,
+subclass types, symint positions) is baked in at compile time.
+"""
+
+import functools
+import keyword
+import logging
+from collections.abc import Callable, Iterable
+
+import torch
+from torch import SymInt
+
+from .schemas import OpaqueMeta, PlainTensorMeta, SubclassCreationMeta
+
+
+log = logging.getLogger(__name__)
+
+
+def _is_symint_placeholder(x: None | int | SymInt) -> bool:
+    """Check whether a size/stride entry is symbolic and needs a runtime value.
+
+    Works both before make_runtime_safe() (entries are SymInt) and after
+    (symbolic entries replaced with None, nested ints with -1).
+    """
+    if x is None:
+        return True
+    if isinstance(x, SymInt) and not x.node.is_nested_int():
+        return True
+    return False
+
+
+def _compute_placeholders(outer: Iterable[None | int | SymInt]) -> list[bool]:
+    return [_is_symint_placeholder(s) for s in outer]
+
+
+def _safe_attr_access(var: str, attr: str) -> str:
+    if attr.isidentifier() and not keyword.iskeyword(attr):
+        return f"{var}.{attr}"
+    return f"getattr({var}, {attr!r})"
+
+
+class _CodegenState:
+    """Accumulates lines of generated source and global bindings."""
+
+    def __init__(self) -> None:
+        self.lines: list[str] = []
+        self.globals: dict[str, object] = {}
+        self._name_counter: int = 0
+
+    def emit(self, line: str, indent: int = 1) -> None:
+        self.lines.append("    " * indent + line)
+
+    def fresh_name(self, prefix: str) -> str:
+        name = f"{prefix}_{self._name_counter}"
+        self._name_counter += 1
+        return name
+
+    def add_global(self, name: str, value: object) -> str:
+        self.globals[name] = value
+        return name
+
+
+def _codegen_unwrap_subclass(
+    state: _CodegenState,
+    meta: SubclassCreationMeta,
+    var: str,
+    indent: int = 1,
+    include_symints: bool = True,
+) -> None:
+    """Emit code to recursively unwrap a single subclass input."""
+    for attr, attr_meta in meta.attrs.items():
+        match attr_meta:
+            case PlainTensorMeta() | OpaqueMeta():
+                state.emit(
+                    f"unwrapped_args.append({_safe_attr_access(var, attr)})",
+                    indent=indent,
+                )
+            case SubclassCreationMeta():
+                inner_var = state.fresh_name("_inner")
+                state.emit(
+                    f"{inner_var} = {_safe_attr_access(var, attr)}", indent=indent
+                )
+                _codegen_unwrap_subclass(
+                    state,
+                    attr_meta,
+                    inner_var,
+                    indent=indent,
+                    include_symints=include_symints,
+                )
+
+    # Emit symint extraction
+    if include_symints:
+        size_placeholders = _compute_placeholders(meta.outer_size)
+        stride_placeholders = _compute_placeholders(meta.outer_stride)
+        has_size_symints = any(size_placeholders)
+        has_stride_symints = any(stride_placeholders)
+
+        if has_size_symints or has_stride_symints:
+            size_var = state.fresh_name("_size")
+            state.emit(f"{size_var} = {var}.size()", indent=indent)
+            for i, is_sym in enumerate(size_placeholders):
+                if is_sym:
+                    state.emit(f"unwrapped_args.append({size_var}[{i}])", indent=indent)
+
+            stride_var = state.fresh_name("_stride")
+            state.emit(f"{stride_var} = {var}.stride()", indent=indent)
+            for i, is_sym in enumerate(stride_placeholders):
+                if is_sym:
+                    state.emit(
+                        f"unwrapped_args.append({stride_var}[{i}])", indent=indent
+                    )
+
+
+def _concrete_value(val: None | int | SymInt) -> int:
+    """Get the concrete int value for a non-symbolic size/stride entry.
+
+    Used for entries that are NOT symbolic placeholders, meaning they are
+    concrete ints or nested ints (represented as -1 after make_runtime_safe).
+    """
+    if isinstance(val, int):
+        return val
+    # Before make_runtime_safe: nested ints are SymInts; use -1 as dummy.
+    # After make_runtime_safe: they're already -1.
+    if isinstance(val, SymInt) and val.node.is_nested_int():
+        return -1
+    raise AssertionError(f"Expected concrete int, got {type(val)}: {val}")
+
+
+def _codegen_wrap_subclass(
+    state: _CodegenState,
+    meta: SubclassCreationMeta,
+    out_idx_ref: list[int],
+) -> str:
+    """Emit code to reconstruct one subclass output. Returns the variable name."""
+    inner_dict_var = state.fresh_name("_out_inner")
+    entries: list[str] = []
+
+    for attr, attr_meta in meta.attrs.items():
+        match attr_meta:
+            case PlainTensorMeta() | OpaqueMeta():
+                idx = out_idx_ref[0]
+                out_idx_ref[0] += 1
+                entries.append(f"{attr!r}: unwrapped_outs[{idx}]")
+            case SubclassCreationMeta():
+                nested_var = _codegen_wrap_subclass(state, attr_meta, out_idx_ref)
+                entries.append(f"{attr!r}: {nested_var}")
+
+    state.emit(f"{inner_dict_var} = {{{', '.join(entries)}}}")
+
+    # Reconstruct outer_size and outer_stride
+    size_placeholders = _compute_placeholders(meta.outer_size)
+    stride_placeholders = _compute_placeholders(meta.outer_stride)
+
+    def _build_tuple(
+        outer: Iterable[None | int | SymInt], placeholders: list[bool]
+    ) -> str:
+        parts: list[str] = []
+        for val, is_sym in zip(outer, placeholders):
+            if is_sym:
+                idx = out_idx_ref[0]
+                out_idx_ref[0] += 1
+                parts.append(f"unwrapped_outs[{idx}]")
+            else:
+                parts.append(repr(_concrete_value(val)))
+        if len(parts) == 1:
+            return f"({parts[0]},)"
+        return f"({', '.join(parts)})"
+
+    size_expr = _build_tuple(meta.outer_size, size_placeholders)
+    stride_expr = _build_tuple(meta.outer_stride, stride_placeholders)
+
+    type_name = state.add_global(
+        state.fresh_name("_subclass_type"),
+        meta.original_subclass_type or type(meta.original_subclass),
+    )
+    meta_name = state.add_global(state.fresh_name("_meta"), meta.meta)
+
+    result_var = state.fresh_name("_out")
+    state.emit(
+        f"{result_var} = {type_name}.__tensor_unflatten__("
+        f"{inner_dict_var}, {meta_name}, {size_expr}, {stride_expr})"
+    )
+    return result_var
+
+
+def _emit_output_wrapping(
+    state: _CodegenState,
+    out_metas: list[PlainTensorMeta | SubclassCreationMeta],
+) -> tuple[list[str], int]:
+    """Emit wrapping code for output metas.
+
+    Returns (result_exprs, num_args_tallied) where result_exprs are Python
+    expression strings referencing each wrapped output.
+    """
+    out_idx_ref = [0]
+    result_exprs: list[str] = []
+    num_args_tallied = 0
+
+    for meta in out_metas:
+        if isinstance(meta, PlainTensorMeta):
+            result_exprs.append(f"unwrapped_outs[{meta.unwrapped_idx}]")
+            num_args_tallied += 1
+            out_idx_ref[0] = max(out_idx_ref[0], meta.unwrapped_idx + 1)
+        else:
+            result_var = _codegen_wrap_subclass(state, meta, out_idx_ref)
+            result_exprs.append(result_var)
+            num_args_tallied += meta.arg_count
+
+    return result_exprs, num_args_tallied
+
+
+def _emit_input_unwrapping(
+    state: _CodegenState,
+    inp_metas: list[PlainTensorMeta | SubclassCreationMeta],
+    frozen_inp_indices: frozenset[int] = frozenset(),
+    include_symints: bool = True,
+) -> None:
+    """Emit unwrapping code for input metas into unwrapped_args.
+
+    Caller must have already emitted ``unwrapped_args = []``.
+    """
+    for i, meta in enumerate(inp_metas):
+        if isinstance(meta, PlainTensorMeta):
+            state.emit(f"unwrapped_args.append(args[{i}])")
+        elif i in frozen_inp_indices:
+            # Frozen by inductor freezing: constant already baked into graph.
+            state.emit("unwrapped_args.append(None)")
+        else:
+            inp_var = state.fresh_name("_inp")
+            type_name = state.add_global(
+                state.fresh_name("_expected_type"),
+                meta.original_subclass_type or type(meta.original_subclass),
+            )
+            state.emit(f"{inp_var} = args[{i}]")
+            state.emit(
+                f"assert type({inp_var}) is {type_name}, "
+                f"f'expected {{{type_name}}}, got {{type({inp_var})}}'",
+            )
+            _codegen_unwrap_subclass(
+                state, meta, inp_var, indent=1, include_symints=include_symints
+            )
+
+
+def _codegen_subclass_wrapper_source(
+    inp_metas: list[PlainTensorMeta | SubclassCreationMeta],
+    out_metas: list[PlainTensorMeta | SubclassCreationMeta],
+    num_fw_outs_saved_for_bw: int | None,
+    frozen_inp_indices: frozenset[int] = frozenset(),
+    act_input_indices: list[int] | None = None,
+) -> tuple[str, dict[str, object]]:
+    """Generate source and globals for a subclass wrapper.
+
+    Returns (source, globals_dict).  The globals_dict will NOT contain
+    ``compiled_fn`` — the caller is responsible for adding it before exec.
+    """
+    state = _CodegenState()
+
+    state.emit("def inner_fn(args):", indent=0)
+
+    # --- Resolve AsyncCollectiveTensors ---
+    # ACTs are transient eager-mode wrappers for async collective overlap.
+    # Inductor triton kernels bypass __torch_dispatch__, so we must call
+    # trigger_wait() before the compiled graph uses the data.
+    if act_input_indices:
+        for i in act_input_indices:
+            state.emit(f"args[{i}] = args[{i}].trigger_wait()")
+
+    # --- Input unwrapping ---
+    state.emit("unwrapped_args = []")
+    _emit_input_unwrapping(state, inp_metas, frozen_inp_indices=frozen_inp_indices)
+
+    # Pass through any trailing args not covered by inp_metas
+    # (e.g. rng seed/offset added by FunctionalizedRngRuntimeWrapper).
+    num_inp_metas = len(inp_metas)
+    state.emit(f"unwrapped_args.extend(args[{num_inp_metas}:])")
+    state.emit("args.clear()")
+
+    # --- Call compiled function ---
+    state.emit("unwrapped_outs = compiled_fn(unwrapped_args)")
+
+    # --- Output wrapping ---
+    result_exprs, num_args_tallied = _emit_output_wrapping(state, out_metas)
+    result_tuple = f"({', '.join(result_exprs)},)" if result_exprs else "()"
+    if num_fw_outs_saved_for_bw is not None:
+        state.emit(
+            f"return {result_tuple} + tuple(unwrapped_outs[{num_args_tallied}:])"
+        )
+    else:
+        state.emit(f"return {result_tuple}")
+
+    source = "\n".join(state.lines)
+    return source, state.globals
+
+
+def _codegen_subclass_wrap_source(
+    out_metas: list[PlainTensorMeta | SubclassCreationMeta],
+) -> tuple[str, dict[str, object]]:
+    """Generate source for wrapping flat outputs into subclasses.
+
+    Used for the backward epilogue. Shares output-wrapping logic with
+    _codegen_subclass_wrapper_source via _emit_output_wrapping.
+    """
+    state = _CodegenState()
+    state.emit("def wrap_fn(unwrapped_outs):", indent=0)
+    result_exprs, _ = _emit_output_wrapping(state, out_metas)
+    result_tuple = f"({', '.join(result_exprs)},)" if result_exprs else "()"
+    state.emit(f"return {result_tuple}")
+    source = "\n".join(state.lines)
+    return source, state.globals
+
+
+def _compile_and_exec_source(
+    source: str,
+    globals_dict: dict[str, object],
+    fn_name: str,
+    artifact_name: str,
+    wrapped_fn: Callable[..., object] | None = None,
+) -> Callable[..., object]:
+    """Compile generated source, exec it, and return the named function.
+
+    If wrapped_fn is provided, applies functools.update_wrapper so that
+    __wrapped__ and __dict__ (e.g. _fx_graph_cache_key) propagate to the
+    generated function.
+    """
+    if log.isEnabledFor(logging.DEBUG):
+        log.debug("Generated %s:\n%s", artifact_name, source)
+
+    torch._logging.trace_structured(
+        "artifact",
+        metadata_fn=lambda: {
+            "name": artifact_name,
+            "encoding": "string",
+        },
+        payload_fn=lambda: source,
+    )
+
+    code = compile(source, f"<{artifact_name}>", "exec")
+    local_dict: dict[str, object] = {}
+    exec(code, globals_dict, local_dict)  # noqa: S102
+    fn = local_dict[fn_name]
+    if wrapped_fn is not None:
+        functools.update_wrapper(fn, wrapped_fn)  # type: ignore[arg-type]
+    return fn  # type: ignore[return-value]
+
+
+def codegen_backward_subclass_fns(
+    grad_input_metas: list[PlainTensorMeta | SubclassCreationMeta] | None = None,
+) -> tuple[Callable[..., object], Callable[..., object] | None]:
+    """Generate codegen'd unwrap and wrap functions for the backward pass.
+
+    Returns (unwrap_fn, wrap_fn). unwrap_fn is used by the backward prologue
+    to unwrap non-tangent subclass inputs (always an identity in AOT dispatch
+    since the compiled forward operates on unwrapped inner tensors). wrap_fn
+    is used by the backward epilogue to wrap flat grad inputs back into
+    subclasses; it is None when grad_input_metas is None.
+    """
+    source = "def unwrap_fn(args):\n    return list(args)"
+    globals_dict: dict[str, object] = {}
+    unwrap_fn = _compile_and_exec_source(
+        source, globals_dict, "unwrap_fn", "backward_subclass_unwrap"
+    )
+
+    wrap_fn = None
+    if grad_input_metas is not None:
+        wrap_source, wrap_globals = _codegen_subclass_wrap_source(grad_input_metas)
+        wrap_fn = _compile_and_exec_source(
+            wrap_source, wrap_globals, "wrap_fn", "backward_subclass_wrapper"
+        )
+
+    return unwrap_fn, wrap_fn
+
+
+def codegen_subclass_wrapper(
+    compiled_fn: Callable[..., object],
+    inp_metas: list[PlainTensorMeta | SubclassCreationMeta],
+    out_metas: list[PlainTensorMeta | SubclassCreationMeta],
+    num_fw_outs_saved_for_bw: int | None,
+    frozen_inp_indices: frozenset[int] = frozenset(),
+    act_input_indices: list[int] | None = None,
+) -> Callable[..., object]:
+    """Generate a specialized wrapper function for subclass unwrap/wrap."""
+    source, globals_dict = _codegen_subclass_wrapper_source(
+        inp_metas,
+        out_metas,
+        num_fw_outs_saved_for_bw,
+        frozen_inp_indices,
+        act_input_indices=act_input_indices,
+    )
+    globals_dict["compiled_fn"] = compiled_fn
+    return _compile_and_exec_source(
+        source, globals_dict, "inner_fn", "subclass_wrapper", wrapped_fn=compiled_fn
+    )
diff --git a/torch/_functorch/_aot_autograd/subclass_parametrization.py b/torch/_functorch/_aot_autograd/subclass_parametrization.py
index d870061734da7..0911526b91242 100644
--- a/torch/_functorch/_aot_autograd/subclass_parametrization.py
+++ b/torch/_functorch/_aot_autograd/subclass_parametrization.py
@@ -1,11 +1,20 @@
+from __future__ import annotations
+
 import dataclasses
 import itertools
-from collections.abc import Iterable
-from typing import Any, Union
+from typing import Any, TYPE_CHECKING
 
 import torch
+from torch._library.opaque_object import is_opaque_reference_type
+from torch._opaque_base import OpaqueBase
 from torch.utils._python_dispatch import is_traceable_wrapper_subclass
 
+from .schemas import OpaqueMeta
+
+
+if TYPE_CHECKING:
+    from collections.abc import Iterable
+
 
 # This is technically very similar to SubclassCreatingMeta
 # in aot_autograd, but we don't need all the stuff in there
@@ -15,23 +24,30 @@ class SubclassCreationMeta:
     start_idx: int
     num_tensors: int
     class_type: Any
-    attrs: dict[str, "SubclassCreationMeta"]
+    # None means the attr is a plain tensor (base case of recursion)
+    attrs: dict[str, SubclassCreationMeta | OpaqueMeta | None]
     metadata: Any
-    outer_size: Iterable[Union[None, int, torch.SymInt]]
-    outer_stride: Iterable[Union[None, int, torch.SymInt]]
+    outer_size: Iterable[None | int | torch.SymInt]
+    outer_stride: Iterable[None | int | torch.SymInt]
 
 
 class UnwrapTensorSubclass(torch.nn.Module):
     def forward(self, *tensors) -> torch.Tensor:  # type: ignore[no-untyped-def]
-        todo: list[torch.Tensor] = list(tensors)
+        todo: list[torch.Tensor | OpaqueBase] = list(tensors)
 
         def _unwrap_tensor_subclasses(subclass_meta, tensors, offset):  # type: ignore[no-untyped-def]
             if subclass_meta is None:
                 return tensors[offset], offset + 1
             inner_tensors = {}
             for attr, meta in subclass_meta.attrs.items():
-                built_tensor, offset = _unwrap_tensor_subclasses(meta, tensors, offset)
-                inner_tensors[attr] = built_tensor
+                if isinstance(meta, OpaqueMeta):
+                    inner_tensors[attr] = tensors[offset]
+                    offset += 1
+                else:
+                    built_tensor, offset = _unwrap_tensor_subclasses(
+                        meta, tensors, offset
+                    )
+                    inner_tensors[attr] = built_tensor
             rebuilt = subclass_meta.class_type.__tensor_unflatten__(
                 inner_tensors,
                 subclass_meta.metadata,
@@ -42,10 +58,10 @@ def _unwrap_tensor_subclasses(subclass_meta, tensors, offset):  # type: ignore[n
 
         return _unwrap_tensor_subclasses(self.subclass_meta, todo, 0)[0]
 
-    def right_inverse(self, tensor: torch.Tensor) -> list[torch.Tensor]:
+    def right_inverse(self, tensor: torch.Tensor) -> list[torch.Tensor | OpaqueBase]:
         if type(tensor) is torch.Tensor:
             raise AssertionError("tensor must be a subclass, not torch.Tensor")
-        plain_tensors: list[torch.Tensor] = []
+        plain_tensors: list[torch.Tensor | OpaqueBase] = []
 
         def _create_subclass_meta(tensor, idx, plain_tensor_container):  # type: ignore[no-untyped-def]
             if type(tensor) is torch.Tensor:
@@ -53,13 +69,30 @@ def _create_subclass_meta(tensor, idx, plain_tensor_container):  # type: ignore[
                 return None, idx + 1
             inner_tensors_attrnames, metadata = tensor.__tensor_flatten__()  # type: ignore[attr-defined]
             new_idx = idx
-            attr_to_meta = {}
+            attr_to_meta: dict[str, SubclassCreationMeta | OpaqueMeta | None] = {}
             for attr in inner_tensors_attrnames:
                 val = getattr(tensor, attr)
-                subclass_meta, new_idx = _create_subclass_meta(
-                    val, new_idx, plain_tensor_container
-                )
-                attr_to_meta[attr] = subclass_meta
+                match val:
+                    case OpaqueBase():
+                        if not is_opaque_reference_type(type(val)):
+                            raise ValueError(
+                                f"{type(val).__name__!r} found in tensor attrs of "
+                                f"{type(tensor).__name__}.__tensor_flatten__(). "
+                                "Only tensors and reference-type opaques are allowed "
+                                "in tensor attrs."
+                            )
+                        attr_to_meta[attr] = OpaqueMeta()
+                        plain_tensor_container.append(val)
+                        new_idx += 1
+                    case torch.Tensor():
+                        subclass_meta, new_idx = _create_subclass_meta(
+                            val, new_idx, plain_tensor_container
+                        )
+                        attr_to_meta[attr] = subclass_meta
+                    case _:
+                        raise AssertionError(
+                            f"expected Tensor or OpaqueBase, got {type(val)}"
+                        )
             return (
                 SubclassCreationMeta(
                     start_idx=idx,
@@ -91,7 +124,7 @@ def unwrap_tensor_subclass_parameters(module: torch.nn.Module) -> torch.nn.Modul
     """
     for name, tensor in itertools.chain(
         list(module.named_parameters(recurse=False)),
-        # pyrefly: ignore [no-matching-overload]
+        # pyrefly: ignore [bad-argument-type, no-matching-overload]
         list(module.named_buffers(recurse=False)),
     ):
         if is_traceable_wrapper_subclass(tensor):
diff --git a/torch/_functorch/_aot_autograd/subclass_utils.py b/torch/_functorch/_aot_autograd/subclass_utils.py
index 6defde0ca3853..d019d9f5649e6 100644
--- a/torch/_functorch/_aot_autograd/subclass_utils.py
+++ b/torch/_functorch/_aot_autograd/subclass_utils.py
@@ -7,11 +7,14 @@
 import collections
 import typing
 from collections.abc import Callable, Iterable, Sequence
-from typing import Any, Optional, TYPE_CHECKING, TypeGuard, TypeVar, Union
+from typing import Any, TypeGuard, TypeVar
 
 import torch
 import torch.utils._pytree as pytree
 from torch import SymInt, Tensor
+from torch._library.fake_class_registry import maybe_unwrap_fake_script_object
+from torch._library.opaque_object import is_opaque_reference_type
+from torch._opaque_base import OpaqueBase
 from torch._subclasses.fake_tensor import get_plain_tensors
 from torch.types import IntLikeType
 from torch.utils._python_dispatch import is_traceable_wrapper_subclass
@@ -31,6 +34,7 @@
     FakifiedFlatArgs,
     FxValue,
     MutationType,
+    OpaqueMeta,
     PlainTensorMeta,
     SubclassCreationMeta,
     ViewAndMutationMeta,
@@ -38,10 +42,6 @@
 from .utils import strict_zip
 
 
-if TYPE_CHECKING:
-    from torch._library.opaque_object import OpaqueType
-
-
 zip = strict_zip
 
 T = TypeVar("T", bound=torch.Tensor)
@@ -70,7 +70,7 @@ def requires_subclass_dispatch(
 
 def maybe_suggest_memory_format(
     t: Tensor, with_memory_format: bool
-) -> Optional[MemoryFormatMeta]:
+) -> MemoryFormatMeta | None:
     if not with_memory_format:
         return None
 
@@ -91,8 +91,15 @@ def _get_types_for_subclass(tensor_subclass: torch.Tensor) -> None:
         tracker[type(tensor_subclass)].append(tensor_subclass)
         inner_keys, _ = tensor_subclass.__tensor_flatten__()
         for key in inner_keys:
-            inner_tensor = getattr(tensor_subclass, key)
-            _get_types_for_subclass(inner_tensor)
+            match getattr(tensor_subclass, key):
+                case torch.Tensor() as inner_value:
+                    _get_types_for_subclass(inner_value)
+                case OpaqueBase():
+                    pass
+                case unexpected:
+                    raise AssertionError(
+                        f"expected Tensor or OpaqueBase, got {type(unexpected)}"
+                    )
 
     tracker: dict[Any, list[Any]] = collections.defaultdict(list)
     _get_types_for_subclass(tensor_subclass)
@@ -100,7 +107,10 @@ def _get_types_for_subclass(tensor_subclass: torch.Tensor) -> None:
 
 
 def create_subclass_metadata(
-    a: Any, start_idx: int, count_symints: bool, with_memory_format: bool = False
+    a: Any,
+    start_idx: int,
+    count_symints: bool,
+    with_memory_format: bool = False,
 ) -> tuple[Any, int]:
     if not is_traceable_wrapper_subclass(a):
         idx = start_idx + 1
@@ -114,16 +124,36 @@ def create_subclass_metadata(
 
     inner_keys, metadata = a.__tensor_flatten__()
     new_start_idx = start_idx
-    attrs = {}
+    attrs: dict[str, SubclassCreationMeta | PlainTensorMeta | OpaqueMeta] = {}
 
     for key in inner_keys:
-        new_subclass_meta, new_start_idx = create_subclass_metadata(
-            getattr(a, key),
-            new_start_idx,
-            count_symints=count_symints,
-            with_memory_format=with_memory_format,
-        )
-        attrs[key] = new_subclass_meta
+        inner_value = getattr(a, key)
+        match inner_value:
+            case OpaqueBase():
+                # During tracing, opaques are wrapped in FakeScriptObject;
+                # unwrap to check the real type.
+                real_type = type(maybe_unwrap_fake_script_object(inner_value))
+                if not is_opaque_reference_type(real_type):
+                    raise RuntimeError(
+                        f"{real_type.__name__!r} found in tensor attrs of "
+                        f"{type(a).__name__}.__tensor_flatten__(). "
+                        "Only tensors and reference-type opaques are allowed "
+                        "in tensor attrs."
+                    )
+                attrs[key] = OpaqueMeta()
+                new_start_idx += 1
+            case Tensor():
+                new_subclass_meta, new_start_idx = create_subclass_metadata(
+                    inner_value,
+                    new_start_idx,
+                    count_symints=count_symints,
+                    with_memory_format=with_memory_format,
+                )
+                attrs[key] = new_subclass_meta
+            case _:
+                raise AssertionError(
+                    f"expected Tensor or OpaqueBase, got {type(inner_value)}"
+                )
 
     # It *must* be because is_traceable_wrapper_subclass() - but mypy is not smart.
     if not isinstance(a, Tensor):
@@ -155,13 +185,13 @@ def create_subclass_metadata(
 # computes metadata about "how to reconstruct the current list of subclasses,
 # if we were given their flattened dense tensors instead"
 def create_subclass_meta(
-    curr_args: Union[list[Any], tuple[Any, ...]],
+    curr_args: list[Any] | tuple[Any, ...],
     *,
     count_symints: bool = True,
     with_memory_format: bool = False,
-) -> list[Union[PlainTensorMeta, SubclassCreationMeta]]:
+) -> list[PlainTensorMeta | SubclassCreationMeta]:
     idx = 0
-    infos: list[Union[PlainTensorMeta, SubclassCreationMeta]] = []
+    infos: list[PlainTensorMeta | SubclassCreationMeta] = []
     for a in curr_args:
         if is_traceable_wrapper_subclass(a):
             if not isinstance(a, Tensor):
@@ -197,7 +227,7 @@ def symint_check(s: IntLikeType) -> TypeGuard[SymInt]:
     return [(i, s) for i, s in enumerate(lst) if symint_check(s)]
 
 
-def compute_symint_placeholders(lst: Iterable[Union[None, int, SymInt]]) -> list[bool]:
+def compute_symint_placeholders(lst: Iterable[None | int | SymInt]) -> list[bool]:
     # Non-nested symints are replaced with None in `make_runtime_safe()`
     return [s is None for s in lst]
 
@@ -223,10 +253,23 @@ def compute_symint_placeholders(lst: Iterable[Union[None, int, SymInt]]) -> list
 # this function below.
 def unwrap_tensor_subclasses(
     wrapped_args: list[FxValue],
-    wrapped_args_descs: list[AOTDescriptor],
+    wrapped_args_descs: Sequence[AOTDescriptor],
     *,
     append_symints: bool,
 ) -> tuple[list[FxValue], list[AOTDescriptor]]:
+    def _maybe_fakeify_opaque(v: Any) -> Any:
+        # Registered opaque types need to be wrapped as FakeScriptObject for
+        # compile-time FX tracing (proxy slot tracking, hashability, etc.).
+        if isinstance(v, OpaqueBase):
+            from torch._guards import detect_fake_mode
+            from torch._library.fake_class_registry import maybe_to_fake_obj
+            from torch._library.opaque_object import is_opaque_type
+
+            fake_mode = detect_fake_mode()
+            if fake_mode is not None and is_opaque_type(type(v)):
+                return maybe_to_fake_obj(fake_mode, v)
+        return v
+
     def flatten_subclass(
         t: FxValue,
         desc: AOTDescriptor,
@@ -236,40 +279,42 @@ def flatten_subclass(
         # unwrap a subclass into plain tensors and their size/stride if "append_symint"
         # is True
         if not is_traceable_wrapper_subclass(t):
-            out[0].append(t)
+            out[0].append(_maybe_fakeify_opaque(t))
             out[1].append(desc)
             return
 
         attrs, _ = t.__tensor_flatten__()
 
+        SubclassGetAttr: Callable[[AOTInput | AOTOutput, str], AOTDescriptor]
+        SubclassSize: Callable[[AOTInput | AOTOutput, int], AOTDescriptor]
+        SubclassStride: Callable[[AOTInput | AOTOutput, int], AOTDescriptor]
+        if isinstance(desc, AOTInput):
+            SubclassGetAttr = SubclassGetAttrAOTInput  # type: ignore[bad-assignment]
+            SubclassSize = SubclassSizeAOTInput  # type: ignore[bad-assignment]
+            SubclassStride = SubclassStrideAOTInput  # type: ignore[bad-assignment]
+        else:
+            SubclassGetAttr = SubclassGetAttrAOTOutput  # type: ignore[bad-assignment]
+            SubclassSize = SubclassSizeAOTOutput  # type: ignore[bad-assignment]
+            SubclassStride = SubclassStrideAOTOutput  # type: ignore[bad-assignment]
+
         for attr in attrs:
-            inner_tensor = getattr(t, attr)
-            n_desc: Any = (
-                SubclassGetAttrAOTInput(desc, attr)
-                if isinstance(desc, AOTInput)
-                # pyrefly: ignore [bad-argument-type]
-                else SubclassGetAttrAOTOutput(desc, attr)
-            )
-            flatten_subclass(inner_tensor, n_desc, out=out)
+            inner_value = getattr(t, attr)
+            n_desc: Any = SubclassGetAttr(desc, attr)
+            flatten_subclass(inner_value, n_desc, out=out)
 
         if append_symints:
             sizes = enumerate_filter_symints(t.size())
             strides = enumerate_filter_symints(t.stride())
             out[0].extend(s for _, s in sizes)
             out[0].extend(s for _, s in strides)
-            if isinstance(desc, AOTInput):
-                out[1].extend(SubclassSizeAOTInput(desc, i) for i, _ in sizes)  # type: ignore[misc]
-                out[1].extend(SubclassStrideAOTInput(desc, i) for i, _ in strides)  # type: ignore[misc]
-            else:
-                out[1].extend(SubclassSizeAOTOutput(desc, i) for i, _ in sizes)  # type: ignore[misc]
-                out[1].extend(SubclassStrideAOTOutput(desc, i) for i, _ in strides)  # type: ignore[misc]
+            out[1].extend(SubclassSize(desc, i) for i, _ in sizes)
+            out[1].extend(SubclassStride(desc, i) for i, _ in strides)
 
     xs_inner: list[FxValue] = []
     descs_inner: list[AOTDescriptor] = []
 
     for x, desc in zip(wrapped_args, wrapped_args_descs):
-        # pyrefly: ignore [bad-argument-type]
-        flatten_subclass(typing.cast(Tensor, x), desc, out=(xs_inner, descs_inner))
+        flatten_subclass(x, desc, out=(xs_inner, descs_inner))
 
     return xs_inner, descs_inner
 
@@ -281,31 +326,41 @@ def runtime_unwrap_tensor_subclasses(
     *,
     append_symints: bool,
     subclass_metas: list[PlainTensorMeta | SubclassCreationMeta] | None = None,
-) -> list[Any]:
+) -> list[int | Tensor | SymInt | OpaqueBase]:
     def flatten_subclass(
-        x: Tensor, meta: SubclassCreationMeta | None, *, out: list[Any]
-    ) -> list[Any]:
+        x: Tensor,
+        subclass_meta: PlainTensorMeta | SubclassCreationMeta | OpaqueMeta | None,
+        *,
+        out: list[OpaqueBase | SymInt | Tensor | int],
+    ) -> list[OpaqueBase | SymInt | Tensor | int]:
         if not is_traceable_wrapper_subclass(x):
             out.append(x)
             return out
 
         if not isinstance(x, Tensor):
             raise AssertionError(f"expected Tensor, got {type(x)}")
+        if not isinstance(subclass_meta, SubclassCreationMeta):
+            raise AssertionError("subclass_meta should be a SubclassCreationMeta")
 
         attrs, _ = x.__tensor_flatten__()
 
         for attr in attrs:
-            inner_tensor = getattr(x, attr)
-            # pyrefly: ignore [missing-attribute]
-            inner_meta = meta.attrs.get(attr)
-            flatten_subclass(inner_tensor, inner_meta, out=out)
+            inner_value = getattr(x, attr)
+            match inner_value:
+                case OpaqueBase():
+                    out.append(inner_value)
+                case Tensor():
+                    inner_meta = subclass_meta.attrs.get(attr)
+                    flatten_subclass(inner_value, inner_meta, out=out)
+                case _:
+                    raise AssertionError(
+                        f"expected Tensor or OpaqueBase, got {type(inner_value)}"
+                    )
 
         if append_symints:
-            if not isinstance(meta, SubclassCreationMeta):
-                raise AssertionError(f"expected SubclassCreationMeta, got {type(meta)}")
             # outer_size
             size = x.size()
-            symint_placeholders = compute_symint_placeholders(meta.outer_size)
+            symint_placeholders = compute_symint_placeholders(subclass_meta.outer_size)
             if len(size) != len(symint_placeholders):
                 raise AssertionError(
                     f"size length mismatch: {len(size)} != {len(symint_placeholders)}"
@@ -316,7 +371,9 @@ def flatten_subclass(
 
             # outer_stride
             stride = x.stride()
-            symint_placeholders = compute_symint_placeholders(meta.outer_stride)
+            symint_placeholders = compute_symint_placeholders(
+                subclass_meta.outer_stride
+            )
             if len(stride) != len(symint_placeholders):
                 raise AssertionError(
                     f"stride length mismatch: {len(stride)} != {len(symint_placeholders)}"
@@ -326,7 +383,7 @@ def flatten_subclass(
             )
         return out
 
-    xs_inner: list[int | Tensor | SymInt | OpaqueType] = []
+    xs_inner: list[int | Tensor | SymInt | OpaqueBase] = []
 
     if append_symints:
         if subclass_metas is None:
@@ -342,10 +399,12 @@ def flatten_subclass(
         if subclass_metas is None:
             get_plain_tensors(typing.cast(Tensor, x), out=xs_inner)
         else:
-            meta = subclass_metas[idx]
-            if not isinstance(meta, SubclassCreationMeta):
-                raise AssertionError(f"expected SubclassCreationMeta, got {type(meta)}")
-            flatten_subclass(typing.cast(Tensor, x), meta, out=xs_inner)
+            subclass_meta = subclass_metas[idx]
+            if not isinstance(subclass_meta, SubclassCreationMeta):
+                raise AssertionError(
+                    f"expected SubclassCreationMeta, got {type(subclass_meta)}"
+                )
+            flatten_subclass(typing.cast(Tensor, x), subclass_meta, out=xs_inner)
 
     return xs_inner
 
@@ -424,7 +483,10 @@ def wrap_tensor_subclasses(
                 )
             else:
                 wrapped_args.append(
-                    subclass_meta.creation_fn(unwrapped_args, is_runtime=is_runtime)
+                    subclass_meta.creation_fn(
+                        unwrapped_args,
+                        is_runtime=is_runtime,
+                    )
                 )
             num_args_tallied += subclass_meta.arg_count
 
diff --git a/torch/_functorch/_aot_autograd/utils.py b/torch/_functorch/_aot_autograd/utils.py
index cc1eb87c4952e..a06c82307ffdf 100644
--- a/torch/_functorch/_aot_autograd/utils.py
+++ b/torch/_functorch/_aot_autograd/utils.py
@@ -16,7 +16,7 @@
 import torch
 import torch.utils._pytree as pytree
 from torch._library.fake_class_registry import FakeScriptObject
-from torch._library.opaque_object import is_opaque_type
+from torch._library.opaque_object import is_opaque_value
 from torch._logging import getArtifactLogger
 from torch._subclasses.fake_tensor import FakeTensor
 from torch._subclasses.functional_tensor import FunctionalTensor
@@ -48,18 +48,6 @@
 strict_zip = partial(zip, strict=True)
 
 
-def _get_symint_hints(exprs: Any) -> Any:
-    """
-    Get the hints of a list/tuple of int/SymInt.
-    """
-    if isinstance(exprs, (list, tuple)):
-        return type(exprs)(_get_symint_hints(e) for e in exprs)
-    elif isinstance(exprs, torch.SymInt):
-        return exprs.node.shape_env.size_hint(exprs.node.expr)
-    else:
-        return exprs
-
-
 def partial_flatten_asdict(obj: object) -> Any:
     if dataclasses.is_dataclass(obj):
         return {
@@ -204,7 +192,7 @@ def flat_fn(*flat_args: Any) -> list[Any]:
         tree_out = fn(*args, **kwargs)
         flat_out, spec = pytree.tree_flatten(tree_out)
         for i in flat_out:
-            is_known_type = isinstance(i, tuple(KNOWN_TYPES)) or is_opaque_type(i)
+            is_known_type = isinstance(i, tuple(KNOWN_TYPES)) or is_opaque_value(i)
             if not is_known_type:
                 raise RuntimeError(
                     f"Found {type(i)} in output, which is not a known type. "
diff --git a/torch/_functorch/aot_autograd.py b/torch/_functorch/aot_autograd.py
index 52fb04cdfde0f..278bbf812cbb0 100644
--- a/torch/_functorch/aot_autograd.py
+++ b/torch/_functorch/aot_autograd.py
@@ -5,7 +5,7 @@
 import time
 from contextlib import nullcontext
 from functools import wraps
-from typing import Any, TYPE_CHECKING
+from typing import Any, Literal, TYPE_CHECKING
 from typing_extensions import ParamSpec, TypeVar
 from unittest.mock import patch
 
@@ -25,19 +25,17 @@
     set_feature_use,
 )
 from torch._guards import detect_fake_mode
-from torch._inductor.utils import BoxedBool
+from torch._inductor.codecache import resolve_pre_grad_pass_timing
+
+# Runtime annotation consumers still resolve BoxedBool from module globals.
 from torch._subclasses import FakeTensor, FakeTensorMode
 from torch.export._tree_utils import reorder_kwargs
 from torch.fx.experimental.proxy_tensor import make_fx
 
-
-static_inputs_log = torch._logging.getArtifactLogger(
-    __name__, "cudagraph_static_inputs"
-)
 from . import config
+from ._aot_autograd import autograd_cache
 from ._aot_autograd.autograd_cache import (  # noqa: F401
     AOTAutogradCache,
-    autograd_cache_key,
     should_use_local_autograd_cache,
     should_use_remote_autograd_cache,
 )
@@ -123,6 +121,7 @@
     InputAliasInfo,
     JointWithDescriptors,
     MutationType,
+    OpaqueMeta,
     OutputAliasInfo,
     OutputType,
     SerializableAOTDispatchCompiler,
@@ -140,7 +139,6 @@
 )
 from ._aot_autograd.utils import (  # noqa: F401
     _get_autocast_states,
-    _get_symint_hints,
     call_func_at_runtime_with_args,
     create_tree_flattened_fn,
     KNOWN_TYPES,
@@ -160,7 +158,7 @@
 if TYPE_CHECKING:
     from collections.abc import Callable, Iterable, Sequence
 
-    from torch._inductor.cudagraph_utils import BoxedDeviceIndex
+    from torch._inductor.compile_fx import CompilerConfigExtra
     from torch._inductor.output_code import OutputCode
     from torch._inductor.utils import InputType
     from torch._ops import OpOverload
@@ -184,6 +182,7 @@
 # may involve compiling multiple subgraphs; e.g., for forwards/backwards)
 AOT_COUNTER = itertools.count()
 
+
 # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 #
@@ -585,7 +584,6 @@ def _dup_fake_script_obj(fake_flat_args: FakifiedFlatArgs) -> list[Any]:
                     flat_args_descs=flat_args_descs,
                     static_input_indices=aot_config.static_input_indices,
                     keep_input_mutations=aot_config.keep_inference_input_mutations,
-                    is_train=needs_autograd,
                     pre_dispatch=aot_config.pre_dispatch,
                 )(*_dup_fake_script_obj(fake_flat_args))
 
@@ -621,35 +619,6 @@ def _dup_fake_script_obj(fake_flat_args: FakifiedFlatArgs) -> list[Any]:
                 # and none of the inputs that require grad are mutated.
                 # so we actually have an inference graph.
                 needs_autograd = False
-                # A bit silly: right now in the subclass codepath, our ViewAndMutationMeta
-                # changes depending on whether we pass in is_train / keep_input_mutations,
-                # so we're forced to recompute the metadata.
-                # TODO: refactor the subclass path of run_functionalized_fw_and_collect_metadata
-                # so that this is unnecessary.
-                if req_subclass_dispatch:
-                    fw_metadata = run_functionalized_fw_and_collect_metadata(
-                        flat_fn,
-                        flat_args_descs=flat_args_descs,
-                        keep_input_mutations=aot_config.keep_inference_input_mutations,
-                        is_train=False,
-                        pre_dispatch=aot_config.pre_dispatch,
-                        static_input_indices=aot_config.static_input_indices,
-                    )(*fake_flat_args)
-                else:
-                    fw_metadata = ViewAndMutationMeta(
-                        input_info=fw_metadata.input_info,
-                        output_info=fw_metadata.output_info,
-                        num_intermediate_bases=fw_metadata.num_intermediate_bases,
-                        keep_input_mutations=aot_config.keep_inference_input_mutations,
-                        traced_tangents=fw_metadata.traced_tangents,
-                        traced_tangents_descs=fw_metadata.traced_tangents_descs,
-                        subclass_inp_meta=fw_metadata.subclass_inp_meta,
-                        subclass_fw_graph_out_meta=fw_metadata.subclass_fw_graph_out_meta,
-                        subclass_tangent_meta=fw_metadata.subclass_tangent_meta,
-                        is_train=False,
-                        tokens=fw_metadata.tokens,
-                        static_input_indices=fw_metadata.static_input_indices,
-                    )
 
     if fw_metadata.num_intermediate_bases > 0:
         if req_subclass_dispatch:
@@ -823,7 +792,8 @@ def returned_function(*args: _P.args, **kwargs: _P.kwargs) -> Any:
         if cached_res is None:
             flat_fn, out_spec = create_tree_flattened_fn(fn, args, kwargs)
             (fake_mode, shape_env) = construct_fake_mode(flat_args, aot_config)
-            fake_flat_args: FakifiedFlatArgs = process_inputs(
+            fake_flat_args: FakifiedFlatArgs
+            fake_flat_args, act_input_indices = process_inputs(
                 flat_args, aot_config, fake_mode, shape_env
             )
             # TODO: We actually could use the pytree path to make better descs.
@@ -841,6 +811,7 @@ def returned_function(*args: _P.args, **kwargs: _P.kwargs) -> Any:
                     fake_mode,
                     shape_env,
                 )
+                aot_state.fw_metadata.act_input_indices = act_input_indices
                 aot_graph_capture = aot_stage1_graph_capture(aot_state, flat_fn)
                 compiled_fn, _ = aot_stage2_compile(
                     aot_state,
@@ -920,39 +891,60 @@ def forward(self, *args: Any, **kwargs: Any) -> Any:
     return AOTModule()
 
 
-def prepare_aot_module_simplified(
+def autograd_cache_key(
+    graph,
+    example_inputs,
+    ignore_shape_env: bool,
+    decompositions,
+    compiler_config_extra: CompilerConfigExtra,
+    keep_inference_input_mutations: bool = False,
+    disable_functionalization: bool = False,
+):
+    (
+        _params_buffers_flat,
+        _params_spec,
+        _buffers_spec,
+        full_args,
+        _full_args_descs,
+        aot_config,
+    ) = prepare_aot_config(
+        graph,
+        example_inputs,
+        decompositions,
+        keep_inference_input_mutations,
+        ignore_shape_env,
+        force_non_lazy_backward_lowering=config.force_non_lazy_backward_lowering,
+        disable_functionalization=disable_functionalization,
+    )
+
+    fake_mode, shape_env = construct_fake_mode(full_args, aot_config)
+    fake_flat_args, _act_input_indices = process_inputs(
+        full_args, aot_config, fake_mode, shape_env, ignore_shape_env
+    )
+
+    return autograd_cache.autograd_cache_key(
+        graph, fake_flat_args, aot_config, compiler_config_extra
+    )
+
+
+def prepare_aot_config(
     mod: nn.Module,
     args: Iterable[Any],
-    kwargs: dict[str, Any] | None,
     decompositions: dict[OpOverload, Callable[..., Any]] | None,
     keep_inference_input_mutations: bool,
-    boxed_forward_device_index: BoxedDeviceIndex | None,
     ignore_shape_env: bool,
-    flatten: bool,
     *,
     force_non_lazy_backward_lowering: bool = False,
     disable_functionalization: bool = False,
-    _record_nn_module_stack: bool = False,
     _disable_torch_fn_metadata_mode: bool = False,
 ) -> tuple[
-    Any,
     list[torch.nn.Parameter | Tensor],
     list[str],
     list[str],
-    FakifiedFlatArgs,
+    list[Any],
     list[Any],
     AOTConfig,
-    FakeTensorMode,
-    ShapeEnv | None,
-    pytree.TreeSpec | None,
-    PytreeThunk | None,
 ]:
-    if not flatten:
-        if kwargs is not None:
-            raise AssertionError("kwargs must be None when flatten=False")
-    elif kwargs is None:
-        kwargs = {}
-
     # TODO: There's something a bit suspicious here; typically simplified
     # module shouldn't actually have any parameters...
     params = dict(mod.named_parameters(remove_duplicate=False))
@@ -966,41 +958,14 @@ def prepare_aot_module_simplified(
 
     params_buffers = {**params, **buffers}
     params_buffers_flat = params_flat + buffers_flat
-    params_buffers_spec = params_spec + buffers_spec
-
-    # Take a break to figure what we're doing with the module
-
-    # NB: This doesn't change the in/out convention, except adding the
-    # parameters as explicit arguments
-    functional_call = create_functional_call(
-        mod,
-        params_buffers_spec,
-        params_len + buffers_len,
-        strict_out_tuple=not flatten,
-        # We need this for export to run ModuleStackTracer
-        # instead of PythonKeyTracer
-        store_orig_mod=_record_nn_module_stack,
-    )
 
     full_args = [*params_flat, *buffers_flat, *args]
-    in_spec, out_spec = None, None
-    if flatten:
-        functional_call, out_spec = create_tree_flattened_fn(
-            functional_call, full_args, kwargs
-        )
-        full_args, in_spec = pytree.tree_flatten((full_args, kwargs))
-
-    del kwargs
 
     # OK, set up the descs
 
     full_args_descs: list[DifferentiableAOTInput] = []
     full_args_descs.extend(ParamAOTInput(fqn) for fqn in params_spec)
     full_args_descs.extend(BufferAOTInput(fqn) for fqn in buffers_spec)
-    # TODO: it would be better to put pytree information in here
-    full_args_descs.extend(
-        PlainAOTInput(i) for i in range(len(full_args) - len(full_args_descs))
-    )
 
     # TODO: These tracing_context fields should become unnecessary once we
     # always maintain sources on all arguments
@@ -1049,9 +1014,97 @@ def prepare_aot_module_simplified(
         disable_functionalization=disable_functionalization,
         _disable_torch_fn_metadata_mode=_disable_torch_fn_metadata_mode,
     )
+
+    return (
+        params_buffers_flat,
+        params_spec,
+        buffers_spec,
+        full_args,
+        full_args_descs,
+        aot_config,
+    )
+
+
+def prepare_aot_module_simplified(
+    mod: nn.Module,
+    args: Iterable[Any],
+    kwargs: dict[str, Any] | None,
+    decompositions: dict[OpOverload, Callable[..., Any]] | None,
+    keep_inference_input_mutations: bool,
+    ignore_shape_env: bool,
+    flatten: bool,
+    *,
+    force_non_lazy_backward_lowering: bool = False,
+    disable_functionalization: bool = False,
+    _record_nn_module_stack: bool = False,
+    _disable_torch_fn_metadata_mode: bool = False,
+) -> tuple[
+    Any,
+    list[torch.nn.Parameter | Tensor],
+    list[str],
+    list[str],
+    FakifiedFlatArgs,
+    list[Any],
+    AOTConfig,
+    FakeTensorMode,
+    ShapeEnv | None,
+    pytree.TreeSpec | None,
+    PytreeThunk | None,
+    list[int],
+]:
+    if not flatten:
+        if kwargs is not None:
+            raise AssertionError("kwargs must be None when flatten=False")
+    elif kwargs is None:
+        kwargs = {}
+
+    (
+        params_buffers_flat,
+        params_spec,
+        buffers_spec,
+        full_args,
+        full_args_descs,
+        aot_config,
+    ) = prepare_aot_config(
+        mod,
+        args,
+        decompositions,
+        keep_inference_input_mutations,
+        ignore_shape_env,
+        force_non_lazy_backward_lowering=force_non_lazy_backward_lowering,
+        disable_functionalization=disable_functionalization,
+        _disable_torch_fn_metadata_mode=_disable_torch_fn_metadata_mode,
+    )
+
+    params_buffers_spec = params_spec + buffers_spec
+
+    # NB: This doesn't change the in/out convention, except adding the
+    # parameters as explicit arguments
+    functional_call = create_functional_call(
+        mod,
+        params_buffers_spec,
+        aot_config.num_params_buffers,
+        strict_out_tuple=not flatten,
+        # We need this for export to run ModuleStackTracer
+        # instead of PythonKeyTracer
+        store_orig_mod=_record_nn_module_stack,
+    )
+
+    in_spec, out_spec = None, None
+    if flatten:
+        functional_call, out_spec = create_tree_flattened_fn(
+            functional_call, full_args, kwargs
+        )
+        full_args, in_spec = pytree.tree_flatten((full_args, kwargs))
+
+    # TODO: it would be better to put pytree information in here
+    full_args_descs.extend(
+        PlainAOTInput(i) for i in range(len(full_args) - len(full_args_descs))
+    )
+
     fake_mode, shape_env = construct_fake_mode(full_args, aot_config)
     # NB: full_args_descs not needed here, fake_flat_args is 1:1 with full_args
-    fake_flat_args = process_inputs(
+    fake_flat_args, act_input_indices = process_inputs(
         full_args, aot_config, fake_mode, shape_env, ignore_shape_env
     )
 
@@ -1067,12 +1120,13 @@ def prepare_aot_module_simplified(
         shape_env,
         in_spec,
         out_spec,
+        act_input_indices,
     )
 
 
 def aot_module_simplified(
     mod: torch.fx.GraphModule | torch._dynamo.utils.GmWrapper,
-    args: Iterable[Any],
+    args: Sequence[Any],
     fw_compiler: AOTDispatchCompiler,
     bw_compiler: AOTDispatchCompiler | None = None,
     partition_fn: Callable[..., Any] = default_partition,
@@ -1081,10 +1135,15 @@ def aot_module_simplified(
     inference_compiler: AOTDispatchCompiler | None = None,
     # TODO: This doesn't seem to be used in any nontrivial way, check if it's
     # actually needed
-    cudagraphs: BoxedBool | None = None,
-    boxed_forward_device_index: BoxedDeviceIndex | None = None,
+    compiler_config_extra: CompilerConfigExtra | None = None,
     ignore_shape_env: bool = False,
     disable_functionalization: bool = False,
+    # Optional callback to run passes on the module at the start of AOT autograd.
+    pre_grad_passes: Callable[
+        [torch.fx.GraphModule, Sequence[InputType]], torch.fx.GraphModule
+    ]
+    | None = None,
+    compile_region_name: str | None = None,
 ) -> Callable[..., Any]:
     """
     This is the simplified or low overhead version of aot_module. For frontends
@@ -1097,8 +1156,14 @@ def aot_module_simplified(
     :func:`aot_module_simplified` removes these overheads.
     """
 
-    if cudagraphs is None:
-        cudagraphs = BoxedBool(torch._inductor.config.triton.cudagraphs)
+    pre_grad_pass_timing: Literal["early", "late"] = resolve_pre_grad_pass_timing()
+
+    if (
+        pre_grad_pass_timing == "early"
+        and pre_grad_passes
+        and isinstance(mod, torch.fx.GraphModule)
+    ):
+        mod = pre_grad_passes(mod, args)
 
     with contextlib.ExitStack() as stack:
         (
@@ -1113,13 +1178,13 @@ def aot_module_simplified(
             shape_env,
             _in_spec,
             _out_spec,
+            act_input_indices,
         ) = prepare_aot_module_simplified(
             mod,
             args,
             None,
             decompositions,
             keep_inference_input_mutations,
-            boxed_forward_device_index,
             ignore_shape_env,
             flatten=False,
             force_non_lazy_backward_lowering=config.force_non_lazy_backward_lowering,
@@ -1140,13 +1205,20 @@ def aot_module_simplified(
                     mod,
                     fake_flat_args,
                     aot_config,
-                    cudagraphs,
-                    boxed_forward_device_index,
+                    compiler_config_extra,
                     local,
                     remote,
+                    compile_region_name=compile_region_name,
                 )
 
         if compiled_fn is None:
+            if (
+                pre_grad_pass_timing == "late"
+                and pre_grad_passes
+                and isinstance(mod, torch.fx.GraphModule)
+            ):
+                mod = pre_grad_passes(mod, args)
+
             stack.enter_context(compiled_autograd._disable())
             aot_state = create_aot_state(
                 stack,
@@ -1157,6 +1229,7 @@ def aot_module_simplified(
                 fake_mode,
                 shape_env,
             )
+            aot_state.fw_metadata.act_input_indices = act_input_indices
             aot_graph_capture = aot_stage1_graph_capture(aot_state, functional_call)
             compiled_fn, _ = aot_stage2_compile(
                 aot_state,
@@ -1220,7 +1293,7 @@ def grab_serialize_fn(fn: Any) -> Callable[..., Any] | None:
 def boxed_nop_preserve_node_meta(
     gm: torch.fx.GraphModule, example_inputs: Sequence[InputType]
 ) -> Any:
-    def run(args: Sequence[Any]) -> OutputCode:
+    def run(args: list[Any]) -> OutputCode:
         with torch.fx.traceback.preserve_node_meta():
             return torch.fx.Interpreter(gm).boxed_run(args)
 
@@ -1313,6 +1386,7 @@ def aot_export_joint_with_descriptors(
         shape_env,
         in_spec,
         out_spec,
+        act_input_indices,
     ) = prepare_aot_module_simplified(
         mod,
         args,
@@ -1320,7 +1394,6 @@ def aot_export_joint_with_descriptors(
         # In contrast, decompositions are needed at this stage.
         decompositions,
         keep_inference_input_mutations,
-        None,
         ignore_shape_env,
         flatten=True,
         # Without this, we will attempt to "compile" the backward lazily
@@ -1347,6 +1420,7 @@ def aot_export_joint_with_descriptors(
         fake_mode,
         shape_env,
     )
+    aot_state.fw_metadata.act_input_indices = act_input_indices
 
     # NB: no cache lookup!
     aot_graph_capture = aot_stage1_graph_capture(aot_state, functional_call)
@@ -1837,9 +1911,6 @@ def _aot_export_function(
         decompositions=decompositions,
         num_params_buffers=num_params_buffers,
         aot_id=next(AOT_COUNTER),
-        # For now there's no use case involving keeping input mutations in the graph
-        # (which we can only do in the inference case anyway).
-        # We can add this later if we need to.
         keep_inference_input_mutations=keep_input_mutations,
         dynamic_shapes=dynamic_shapes,
         aot_autograd_arg_pos_to_source=None,
@@ -1852,7 +1923,9 @@ def _aot_export_function(
         fake_mode, shape_env = construct_fake_mode(flat_args, aot_config)
     else:
         shape_env = fake_mode.shape_env
-    fake_flat_args = process_inputs(flat_args, aot_config, fake_mode, shape_env)
+    fake_flat_args, act_input_indices = process_inputs(
+        flat_args, aot_config, fake_mode, shape_env
+    )
     # TODO: Improve the descs here with pytree information
     fake_flat_args_descs: list[AOTInput] = [
         PlainAOTInput(i) for i in range(len(fake_flat_args))
@@ -1868,6 +1941,7 @@ def _aot_export_function(
             fake_mode,
             shape_env,
         )
+        aot_state.fw_metadata.act_input_indices = act_input_indices
         aot_graph_capture = aot_stage1_graph_capture(aot_state, flat_fn)
         fx_g, meta = aot_stage2_export(aot_state, aot_graph_capture)
 
diff --git a/torch/_functorch/compilers.py b/torch/_functorch/compilers.py
index 5cbf0c5269e43..687aabbb9c9cf 100644
--- a/torch/_functorch/compilers.py
+++ b/torch/_functorch/compilers.py
@@ -334,7 +334,8 @@ def debug_compile(
   minifier(fx.symbolic_trace(mod), inps, check_nvfuser_subprocess)
 """
     )
-    from foo import FxModule  # pyrefly: ignore[missing-import]
+    # pyrefly: ignore[missing-import, missing-module-attribute]
+    from foo import FxModule
 
     FxModule().cuda()(*inps)
 
diff --git a/torch/_functorch/config.py b/torch/_functorch/config.py
index 6f30a0933f80e..7f834556485d6 100644
--- a/torch/_functorch/config.py
+++ b/torch/_functorch/config.py
@@ -13,7 +13,7 @@
 
 import os
 import sys
-from typing import Literal, Optional, TYPE_CHECKING
+from typing import Literal, TYPE_CHECKING
 
 from torch.utils._config_module import Config, install_config_module
 
@@ -99,7 +99,7 @@
 error_on_custom_op_aliasing = bool(os.getenv("CI"))
 
 
-def remote_autograd_cache_default() -> Optional[bool]:
+def remote_autograd_cache_default() -> bool | None:
     if os.environ.get("TORCHINDUCTOR_AUTOGRAD_REMOTE_CACHE") == "1":
         return True
     if os.environ.get("TORCHINDUCTOR_AUTOGRAD_REMOTE_CACHE") == "0":
@@ -230,6 +230,10 @@ def remote_autograd_cache_default() -> Optional[bool]:
 # activation reloading with prefetching when using separate streams (bwd graph)
 activation_reload_prefetch = False
 
+# CPU ↔ GPU bandwidth in GB/s, used to estimate transfer times for prefetch
+# scheduling. This is hardware-specific and should be set by the user.
+activation_offload_cpu_gpu_bw: float = 50.0
+
 # If FakeTensor.data_ptr() should error.
 # This option is independent of AOTAutograd and torch.compile, but our policy
 # is to turn it off during torch.compile.
@@ -347,7 +351,7 @@ def remote_autograd_cache_default() -> Optional[bool]:
 # compiler to proceed with compilation by choosing the preferred device type
 # for consistency. For example, set to "mtia" to prefer MTIA devices over
 # CPU, or "cuda" to prefer CUDA devices over CPU.
-fake_tensor_prefer_device_type: Optional[str] = None
+fake_tensor_prefer_device_type: str | None = None
 
 # CUDAGraph safe run_with_rng functionalization.
 # TODO: turn on by default
@@ -399,7 +403,8 @@ def remote_autograd_cache_default() -> Optional[bool]:
 # At runtime non contiguous tangents will be coerced to be contiguous.
 # This config changes this guess for tangents strides to be the same as outputs.
 # TODO(ivankobzarev): Remove this config once extra memory usage is investigated.
-guess_tangent_strides_as_outputs = False
+guess_tangent_strides_as_outputs = not is_fbcode()
+
 
 # This is a temporary config to ensure all ranks take the same decision in the partitioner
 # it will ultimately be removed once we share size_hints across ranks through compiler collectives
diff --git a/torch/_functorch/eager_transforms.py b/torch/_functorch/eager_transforms.py
index 1c7f81bf94a4d..b3bf0f078f8e6 100644
--- a/torch/_functorch/eager_transforms.py
+++ b/torch/_functorch/eager_transforms.py
@@ -8,7 +8,7 @@
 
 import contextlib
 from functools import partial, wraps
-from typing import Any, Optional, overload, TYPE_CHECKING, Union
+from typing import Any, overload, TYPE_CHECKING
 from typing_extensions import ParamSpec, TypeVar
 
 import torch
@@ -335,6 +335,23 @@ def vjp(
     return _vjp_with_argnums(func, *primals, has_aux=has_aux)
 
 
+@contextlib.contextmanager
+def _disable_inference_mode() -> Generator[None, None, None]:
+    # Disable inference_mode without clobbering grad_mode / fw_grad_mode.
+    # torch.inference_mode(False) unconditionally sets grad_mode=True and
+    # fw_grad_mode=True; we save and restore those to avoid that.
+    # No-op when inference_mode is already off.
+    if not torch.is_inference_mode_enabled():
+        yield
+        return
+    prev_grad = torch.is_grad_enabled()
+    prev_fw_grad = torch._C._is_fwd_grad_enabled()
+    with torch.inference_mode(False):
+        torch._C._set_grad_enabled(prev_grad)
+        torch._C._set_fwd_grad_enabled(prev_fw_grad)
+        yield
+
+
 @contextlib.contextmanager
 def grad_increment_nesting() -> Generator[int, None, None]:
     try:
@@ -369,12 +386,12 @@ def jvp_increment_nesting() -> Generator[int, None, None]:
 def _vjp_with_argnums(
     func: Callable[..., Any],
     *primals: Any,
-    argnums: Optional[argnums_t] = None,
+    argnums: argnums_t | None = None,
     has_aux: bool = False,
 ) -> tuple[Any, Callable[..., Any]] | tuple[Any, Callable[..., Any], Any]:
     # This is the same function as vjp but also accepts an argnums argument
     # All args are the same as vjp except for the added argument
-    # argnums (Optional[int or tuple[int,...]]): Optional, specifies the argument(s) to compute gradients with respect to.
+    # argnums (int or tuple[int,...] | None): Optional, specifies the argument(s) to compute gradients with respect to.
     #         If None, computes the gradients with respect to all inputs (used for vjp). Default: None
     #
     # WARN: Users should NOT call this function directly and should just be calling vjp.
@@ -438,13 +455,23 @@ def wrapper(
                     f"cotangents: {treespec_pprint(cotangents_spec)}, "
                     f"primal output: {treespec_pprint(primals_out_spec)}"
                 )
-            result = _autograd_grad(
-                flat_primals_out,
-                flat_diff_primals,
-                flat_cotangents,
-                retain_graph=retain_graph,
-                create_graph=create_graph,
+            # This closure runs after grad_increment_nesting exits, so
+            # inference_mode may have been restored. Disable it for autograd.
+            # Skip under Dynamo — tracing through the generator CM emits
+            # spurious _enter_inference_mode nodes.
+            ctx = (
+                contextlib.nullcontext()
+                if torch.compiler.is_compiling()
+                else _disable_inference_mode()
             )
+            with ctx:
+                result = _autograd_grad(
+                    flat_primals_out,
+                    flat_diff_primals,
+                    flat_cotangents,
+                    retain_graph=retain_graph,
+                    create_graph=create_graph,
+                )
             return tree_unflatten(result, primals_spec)
 
     if has_aux:
@@ -476,10 +503,10 @@ def error_if_complex(func_name: str, args: Any, is_input: bool) -> None:
 @exposed_in("torch.func")
 def jacrev(
     func: Callable[..., Any],
-    argnums: Union[int, tuple[int, ...]] = 0,
+    argnums: int | tuple[int, ...] = 0,
     *,
     has_aux: bool = False,
-    chunk_size: Optional[int] = None,
+    chunk_size: int | None = None,
     _preallocate_and_copy: bool = False,
 ) -> Callable[..., Any]:
     """
@@ -1128,7 +1155,7 @@ def _jvp_with_argnums(
     func: Callable[..., Any],
     primals: Any,
     tangents: Any,
-    argnums: Optional[argnums_t],
+    argnums: argnums_t | None,
     *,
     strict: bool = False,
     has_aux: bool,
diff --git a/torch/_functorch/partitioners.py b/torch/_functorch/partitioners.py
index 6327dd405c8f0..6042049bebe12 100644
--- a/torch/_functorch/partitioners.py
+++ b/torch/_functorch/partitioners.py
@@ -44,7 +44,7 @@
     find_symbol_binding_fx_nodes,
     free_symbols,
     is_symbol_binding_fx_node,
-    size_hint,
+    optimization_hint,
     statically_known_false,
     statically_known_true,
 )
@@ -67,6 +67,7 @@
     SavedForBackwardsNoVcCheckAOTOutput,
 )
 from ._aot_autograd.functional_utils import _is_functional_graph
+from ._aot_autograd.graph_compile import is_opaque_node
 from ._aot_autograd.logging_utils import get_aot_graph_name
 from ._aot_autograd.utils import (
     _is_bwd_seed_offset,
@@ -382,6 +383,12 @@ def _extract_graph_with_inputs_outputs(
             output_values.append(x)
     out = new_graph.output(tuple(output_values))
     out.meta["desc"] = outputs_descs
+    # Snapshot stack traces on the output node before passes run,
+    # as later passes may strip stack_trace from individual nodes.
+    out.meta["output_stack_traces"] = [
+        v.meta.get("stack_trace") if isinstance(v, fx.Node) else None
+        for v in output_values
+    ]
 
     new_graph.eliminate_dead_code()
     new_graph.lint()
@@ -697,7 +704,7 @@ def calculate_range(dtype: torch.dtype) -> tuple[float, float]:
     return info.min, info.max
 
 
-def quantize_activation_fw(graph: torch.fx.Graph) -> None:
+def quantize_activation_fw(graph: torch.fx.Graph, num_fwd_outputs: int = 0) -> None:
     output = graph.find_nodes(op="output")[0]
     fwd_outputs = output.args[0]
     quant_type = get_quant_type()
@@ -706,6 +713,14 @@ def quantize_activation_fw(graph: torch.fx.Graph) -> None:
     tensor_scale_nodes: list[fx.Node] = []
     sym_scale_nodes: list[fx.Node] = []
     for position, node in enumerate(fwd_outputs):
+        # Don't quantize user-visible forward outputs. A tensor may appear as
+        # both a user output and a saved-for-backward activation (same FX node
+        # at two positions). Quantizing the user output position would:
+        # 1. Return fp8 to the user instead of the original precision
+        # 2. Create duplicate fp8_quant/fp8_scale backward placeholders that
+        #    shift the stride mapping in _aot_stage2b_bw_compile (T264303372)
+        if position < num_fwd_outputs:
+            continue
         # check if the activation node is the node saved for quantization
         if node.meta.get("saved_for_quantization", False):
             # case: use scaling
@@ -846,6 +861,7 @@ def perform_fp8_activation_quantization(
     fwd_module: fx.GraphModule,
     bwd_module: fx.GraphModule,
     bwd_module_inputs: dict[str, fx.Node],
+    num_fwd_outputs: int = 0,
 ) -> None:
     trace_structured(
         "artifact",
@@ -858,7 +874,7 @@ def perform_fp8_activation_quantization(
         ),
     )
 
-    quantize_activation_fw(fwd_module.graph)
+    quantize_activation_fw(fwd_module.graph, num_fwd_outputs)
 
     trace_structured(
         "artifact",
@@ -937,15 +953,8 @@ def enable_activation_quantization(
     fwd_module: fx.GraphModule,
     bwd_module: fx.GraphModule,
     static_lifetime_input_nodes: OrderedSet[fx.Node] | None = None,
+    num_fwd_outputs: int = 0,
 ) -> None:
-    if (
-        inductor_config.post_grad_fusion_options.get(
-            "activation_quantization_aten_pass", None
-        )
-        is None
-    ):
-        return
-
     static_input_names: list[str] = (
         [node.name for node in static_lifetime_input_nodes]
         if static_lifetime_input_nodes
@@ -976,33 +985,62 @@ def enable_activation_quantization(
             should_perform_fp8_quant = True
 
     if should_perform_fp8_quant:
-        perform_fp8_activation_quantization(fwd_module, bwd_module, bwd_module_inputs)
+        perform_fp8_activation_quantization(
+            fwd_module, bwd_module, bwd_module_inputs, num_fwd_outputs
+        )
 
 
 def _extract_fwd_bwd_modules(
     joint_module: fx.GraphModule,
     saved_values: list[fx.Node],
     saved_sym_nodes: list[fx.Node],
+    saved_opaque_nodes: list[fx.Node] | None = None,
     *,
     num_fwd_outputs: int,
     static_lifetime_input_nodes: OrderedSet[fx.Node] | None = None,
+    ignore_must_be_in_fw_bw: bool = False,
+    omit_aot_autograd_runtime: bool = False,
 ) -> tuple[fx.GraphModule, fx.GraphModule]:
+    """Extract forward and backward graph modules from a joint graph.
+
+    Args:
+        ignore_must_be_in_fw_bw: When True, disables forward/backward placement
+            enforcement in _extract_graph_with_inputs_outputs. Needed when the
+            joint_module is not an original fwd+bwd joint graph (e.g. a backward
+            graph being re-partitioned for dI/dW splitting).
+        omit_aot_autograd_runtime: When True, skips postprocessing that is
+            only needed when the resulting modules will be wrapped in a custom
+            autograd.Function (the AOTAutograd path). This includes: tangent input
+            handling, version-counter check sorting of saved tensors, opaque object
+            (FakeScriptObject) separation, and fp8 activation quantization. Set this
+            to True when the fwd/bwd modules will be executed directly without autograd.
+    """
     fwd_outputs, bwd_outputs, fwd_outputs_descs, bwd_outputs_descs = (
         _extract_fwd_bwd_outputs(joint_module, num_fwd_outputs=num_fwd_outputs)
     )
     placeholders = joint_module.graph.find_nodes(op="placeholder")
     primal_inputs = [*filter(_is_primal, placeholders)]
-    tangent_inputs = [*filter(_is_tangent, placeholders)]
+    tangent_inputs = (
+        [] if omit_aot_autograd_runtime else [*filter(_is_tangent, placeholders)]
+    )
     fwd_seed_offset_inputs = [*filter(_is_fwd_seed_offset, placeholders)]
     bwd_seed_offset_inputs = [*filter(_is_bwd_seed_offset, placeholders)]
     backward_state_inputs = [*filter(_is_backward_state, placeholders)]
 
+    if saved_opaque_nodes is None:
+        saved_opaque_nodes = []
+
     bwd_graph = _extract_graph_with_inputs_outputs(
         joint_module.graph,
-        saved_sym_nodes + saved_values + tangent_inputs + bwd_seed_offset_inputs,
+        saved_sym_nodes
+        + saved_opaque_nodes
+        + saved_values
+        + tangent_inputs
+        + bwd_seed_offset_inputs,
         bwd_outputs,
         bwd_outputs_descs,
         "backward",
+        ignore_must_be_in_fw_bw=ignore_must_be_in_fw_bw,
     )
 
     distributed_enabled = torch.distributed.is_available()
@@ -1012,6 +1050,7 @@ def _extract_fwd_bwd_modules(
         if not node.users:
             _remove_by_name(saved_values, node.name)
             _remove_by_name(saved_sym_nodes, node.name)
+            _remove_by_name(saved_opaque_nodes, node.name)
         # wait_tensor is a bit special: if we have a "dead activation" that is not used in the bw,
         # but this dead activation is actually a collective,
         # then the collective will generally by followed by a wait_tensor() call.
@@ -1023,6 +1062,7 @@ def _extract_fwd_bwd_modules(
         ):
             _remove_by_name(saved_values, node.name)
             _remove_by_name(saved_sym_nodes, node.name)
+            _remove_by_name(saved_opaque_nodes, node.name)
         elif _is_backward_state(node):
             # BackwardState is saved directly
             _remove_by_name(saved_values, node.name)
@@ -1070,75 +1110,121 @@ def _extract_fwd_bwd_modules(
     saved_sym_nodes.clear()
     saved_sym_nodes.extend(saved_sym_nodes_binding + saved_sym_nodes_derived)
 
-    # See Note [Activations with no version counter checks in eager]
-    # Sort saved_values so that tensors with saved_tensor_with_no_vc_check=True
-    # are at the end. This allows us to have two consecutive slices:
-    # 1. tensors_saved_with_vc_check_slice - tensors saved via save_for_backward
-    # 2. tensors_saved_with_no_vc_check_slice - tensors stashed on ctx without save_for_backward
-    # The sort is stable, so the relative order within each group is preserved.
-    #
-    # Additionally, separate out opaque objects (FakeScriptObject) from tensors.
-    # Opaque objects should be placed after tensors in the forward outputs.
-    saved_values_with_vc_check = []
-    saved_values_no_vc_check = []
-    saved_opaque_objects = []
-    for node in saved_values:
-        # Check if this is an opaque object
-        if isinstance(node.meta.get("val"), FakeScriptObject):
-            saved_opaque_objects.append(node)
-        elif node.meta.get("saved_tensor_with_no_vc_check", False):
-            saved_values_no_vc_check.append(node)
-        else:
-            saved_values_with_vc_check.append(node)
-    saved_values.clear()
-    saved_values.extend(saved_values_with_vc_check + saved_values_no_vc_check)
-    no_vc_check_start_idx = len(saved_values_with_vc_check)
-
-    # debug assert: given saved_values where the last k of them are expected to not
-    # require VC checks, they should all have node metadata indicating so.
-    for i, node in enumerate(saved_values):
-        if i >= no_vc_check_start_idx:
-            if not node.meta.get("saved_tensor_with_no_vc_check", False):
-                raise AssertionError(
-                    f"i={i}, no_vc_check_start_idx={no_vc_check_start_idx}, len(saved_values)={len(saved_values)}"
-                )
+    if not omit_aot_autograd_runtime:
+        # See Note [Activations with no version counter checks in eager]
+        # Sort saved_values so that tensors with saved_tensor_with_no_vc_check=True
+        # are at the end. This allows us to have two consecutive slices:
+        # 1. tensors_saved_with_vc_check_slice - tensors saved via save_for_backward
+        # 2. tensors_saved_with_no_vc_check_slice - tensors stashed on ctx without save_for_backward
+        # The sort is stable, so the relative order within each group is preserved.
+        #
+        # Additionally, separate out opaque objects (FakeScriptObject) from tensors.
+        # Opaque objects should be placed after tensors in the forward outputs.
+        saved_values_with_vc_check = []
+        saved_values_no_vc_check = []
+        saved_opaque_objects = []
+        for node in saved_values:
+            # Check if this is an opaque object
+            if isinstance(node.meta.get("val"), FakeScriptObject):
+                saved_opaque_objects.append(node)
+            elif node.meta.get("saved_tensor_with_no_vc_check", False):
+                saved_values_no_vc_check.append(node)
+            else:
+                saved_values_with_vc_check.append(node)
+        saved_values.clear()
+        saved_values.extend(saved_values_with_vc_check + saved_values_no_vc_check)
+        no_vc_check_start_idx = len(saved_values_with_vc_check)
+
+        # debug assert: given saved_values where the last k of them are expected to not
+        # require VC checks, they should all have node metadata indicating so.
+        for i, node in enumerate(saved_values):
+            if i >= no_vc_check_start_idx:
+                if not node.meta.get("saved_tensor_with_no_vc_check", False):
+                    raise AssertionError(
+                        f"i={i}, no_vc_check_start_idx={no_vc_check_start_idx}, len(saved_values)={len(saved_values)}"
+                    )
 
-    # Now, we re-generate the fwd/bwd graphs.
-    # NB: This might increase compilation time, but I doubt it matters
-    # Convention for saved acts is (tensors_with_vc_check, tensors_no_vc_check, opaque_objects, symints)
-    fwd_graph = _extract_graph_with_inputs_outputs(
-        joint_module.graph,
-        primal_inputs + fwd_seed_offset_inputs,
-        fwd_outputs + saved_values + saved_opaque_objects + saved_sym_nodes,
-        fwd_outputs_descs
-        + [
-            SavedForBackwardsNoVcCheckAOTOutput(i)
-            if i >= no_vc_check_start_idx and i < len(saved_values)
-            else SavedForBackwardsAOTOutput(i)
-            for i in range(
-                len(saved_values) + len(saved_opaque_objects) + len(saved_sym_nodes)
-            )
-        ],
-        "forward",
-    )
-    bwd_graph = _extract_graph_with_inputs_outputs(
-        joint_module.graph,
-        saved_sym_nodes
-        + saved_values
-        + saved_opaque_objects
-        + tangent_inputs
-        + bwd_seed_offset_inputs
-        + backward_state_inputs,
-        bwd_outputs,
-        bwd_outputs_descs,
-        "backward",
-    )
+        # Now, we re-generate the fwd/bwd graphs.
+        # NB: This might increase compilation time, but I doubt it matters
+        # Convention for saved acts is (tensors_with_vc_check, tensors_no_vc_check, opaque_objects, symints, opaque_nodes)
+        fwd_graph = _extract_graph_with_inputs_outputs(
+            joint_module.graph,
+            primal_inputs + fwd_seed_offset_inputs,
+            fwd_outputs
+            + saved_values
+            + saved_opaque_objects
+            + saved_opaque_nodes
+            + saved_sym_nodes,
+            fwd_outputs_descs
+            + [
+                SavedForBackwardsNoVcCheckAOTOutput(i)
+                if i >= no_vc_check_start_idx and i < len(saved_values)
+                else SavedForBackwardsAOTOutput(i)
+                for i in range(
+                    len(saved_values)
+                    + len(saved_opaque_objects)
+                    + len(saved_opaque_nodes)
+                    + len(saved_sym_nodes)
+                )
+            ],
+            "forward",
+            ignore_must_be_in_fw_bw=ignore_must_be_in_fw_bw,
+        )
+        bwd_graph = _extract_graph_with_inputs_outputs(
+            joint_module.graph,
+            saved_sym_nodes
+            + saved_values
+            + saved_opaque_objects
+            + saved_opaque_nodes
+            + tangent_inputs
+            + bwd_seed_offset_inputs
+            + backward_state_inputs,
+            bwd_outputs,
+            bwd_outputs_descs,
+            "backward",
+            ignore_must_be_in_fw_bw=ignore_must_be_in_fw_bw,
+        )
+    else:
+        # Raw fwd/bwd split for direct execution without autograd
+        fwd_graph = _extract_graph_with_inputs_outputs(
+            joint_module.graph,
+            primal_inputs + fwd_seed_offset_inputs,
+            fwd_outputs + saved_values + saved_sym_nodes,
+            fwd_outputs_descs
+            + [
+                SavedForBackwardsAOTOutput(i)
+                for i in range(len(saved_values) + len(saved_sym_nodes))
+            ],
+            "forward",
+            ignore_must_be_in_fw_bw=ignore_must_be_in_fw_bw,
+        )
+        bwd_graph = _extract_graph_with_inputs_outputs(
+            joint_module.graph,
+            saved_values
+            + saved_sym_nodes
+            + bwd_seed_offset_inputs
+            + backward_state_inputs,
+            bwd_outputs,
+            bwd_outputs_descs,
+            "backward",
+            ignore_must_be_in_fw_bw=ignore_must_be_in_fw_bw,
+        )
 
     fwd_module = fx._lazy_graph_module._make_graph_module(joint_module, fwd_graph)
     bwd_module = fx._lazy_graph_module._make_graph_module(joint_module, bwd_graph)
-    enable_activation_quantization(
-        saved_values, fwd_module, bwd_module, static_lifetime_input_nodes
-    )
+    if (
+        inductor_config.post_grad_fusion_options.get(
+            "activation_quantization_aten_pass", None
+        )
+        is not None
+    ):
+        enable_activation_quantization(
+            saved_values,
+            fwd_module,
+            bwd_module,
+            static_lifetime_input_nodes,
+            num_fwd_outputs,
+        )
     return fwd_module, bwd_module
 
 
@@ -1222,6 +1308,7 @@ def default_partition(
 
     saved_values = []
     saved_sym_nodes = []
+    saved_opaque_nodes = []
 
     distributed_enabled = torch.distributed.is_available()
 
@@ -1281,14 +1368,23 @@ def is_impure(node: fx.Node) -> bool:
             # Must be ordered before MUST_SAVE tags to avoid saving tuples marked MUST_SAVE.
             continue
         if node.meta.get("recompute") == CheckpointPolicy.MUST_SAVE:
-            saved_values.append(node)
+            if is_opaque_node(node):
+                saved_opaque_nodes.append(node)
+            else:
+                saved_values.append(node)
             continue
         if is_impure(node):
             if graph_has_recomputable_ops:
                 raise AssertionError(
                     f"Trying to apply AC on a graph with impure op: {node}, {node.target}"
                 )
-            saved_values.append(node)
+            if is_opaque_node(node):
+                saved_opaque_nodes.append(node)
+            else:
+                saved_values.append(node)
+            continue
+        if is_opaque_node(node):
+            saved_opaque_nodes.append(node)
             continue
         if not is_tensor(node) and node.op == "call_function":
             raise AssertionError(f"Expected {node} to be a tensor")
@@ -1309,6 +1405,7 @@ def is_impure(node: fx.Node) -> bool:
 
     saved_values = list(dict.fromkeys(saved_values).keys())
     saved_sym_nodes = list(dict.fromkeys(saved_sym_nodes).keys())
+    saved_opaque_nodes = list(dict.fromkeys(saved_opaque_nodes).keys())
 
     if config._sync_decision_cross_ranks:
         saved_values = _sync_decision_cross_ranks(joint_module.graph, saved_values)
@@ -1319,6 +1416,7 @@ def is_impure(node: fx.Node) -> bool:
         joint_module,
         saved_values,
         saved_sym_nodes=saved_sym_nodes,
+        saved_opaque_nodes=saved_opaque_nodes,
         num_fwd_outputs=num_fwd_outputs,
         static_lifetime_input_nodes=static_lifetime_input_nodes,
     )
@@ -1370,7 +1468,7 @@ def _size_of(node: fx.Node) -> int:
     def object_nbytes(x: object) -> int:
         if not isinstance(x, torch.Tensor):
             return 0
-        return _tensor_nbytes(size_hint(x.numel(), fallback=4096), x.dtype)
+        return _tensor_nbytes(optimization_hint(x.numel(), fallback=4096), x.dtype)
 
     if "val" in node.meta:
         val = node.meta["val"]
@@ -2255,9 +2353,12 @@ def ban_recomputation_if_allowed(node: fx.Node, reason: str = "") -> bool:
             weight = float(sym_node_size(node))
             cannot_save_reason = None
         elif is_non_tensor_node:
-            # FakeScriptObjects (opaque objects) should have weight 0.0 so they can be
-            # properly partitioned between forward and backward, like BackwardState.
-            if isinstance(node.meta.get("val"), (BackwardState, FakeScriptObject)):
+            # FakeScriptObjects and opaque objects should have weight 0.0
+            # so they can be properly partitioned between forward and
+            # backward, like BackwardState.
+            if isinstance(
+                node.meta.get("val"), (BackwardState, FakeScriptObject)
+            ) or is_opaque_node(node):
                 weight = 0.0
                 cannot_save_reason = None
             else:
@@ -2913,7 +3014,7 @@ def _remove_symbols_without_guarding(x: torch.Tensor, fallback: int) -> torch.Te
     shape = list(x.shape)
 
     def realize_symbol(d: torch.SymInt | int) -> int:
-        return size_hint(d, fallback=fallback)
+        return optimization_hint(d, fallback=fallback)
 
     shape = [realize_symbol(s) for s in shape]
     stride = [realize_symbol(s) for s in x.stride()]
@@ -2927,7 +3028,7 @@ def materialize_arg(x: Any) -> Any:
         if isinstance(x, fx.Node) and isinstance(x.meta["val"], torch.Tensor):
             return _remove_symbols_without_guarding(x.meta["val"], fallback=4096)
         elif isinstance(x, fx.Node) and isinstance(x.meta["val"], torch.SymInt):
-            return size_hint(x.meta["val"], fallback=4096)
+            return optimization_hint(x.meta["val"], fallback=4096)
         elif isinstance(x, fx.Node) and isinstance(x.meta["val"], torch.SymFloat):
             return 1.0
         elif isinstance(x, fx.Node) and isinstance(x.meta["val"], torch.SymBool):
@@ -3542,9 +3643,27 @@ def min_cut_rematerialization_partition(
     # pyrefly: ignore [unbound-name]
     if config._sync_decision_cross_ranks:
         saved_values = _sync_decision_cross_ranks(joint_graph, saved_values)
+
     # save_for_backward on tensors and stashes symints in autograd .ctx
-    saved_sym_nodes = list(filter(is_sym_node, saved_values))
-    saved_values = list(filter(lambda n: not is_sym_node(n), saved_values))
+    # Skip SymBool nodes whose only consumers are _assert_scalar calls.
+    # These are runtime assertion intermediates and are not needed in backward
+    # for any real computation.
+    def _is_assert_only_symbool(n: fx.Node) -> bool:
+        return (
+            isinstance(n.meta.get("val"), torch.SymBool)
+            and len(n.users) > 0
+            and all(u.target is torch.ops.aten._assert_scalar.default for u in n.users)
+        )
+
+    saved_sym_nodes = list(
+        filter(
+            lambda n: is_sym_node(n) and not _is_assert_only_symbool(n), saved_values
+        )
+    )
+    saved_opaque_nodes = list(filter(is_opaque_node, saved_values))
+    saved_values = list(
+        filter(lambda n: not is_sym_node(n) and not is_opaque_node(n), saved_values)
+    )
 
     # NB: saved_sym_nodes will be mutated to reflect the actual saved symbols
     fw_module, bw_module = _extract_fwd_bwd_modules(
@@ -3552,6 +3671,7 @@ def min_cut_rematerialization_partition(
         saved_values,
         # pyrefly: ignore [bad-argument-type]
         saved_sym_nodes=saved_sym_nodes,
+        saved_opaque_nodes=saved_opaque_nodes,
         num_fwd_outputs=num_fwd_outputs,
         static_lifetime_input_nodes=node_info.static_lifetime_input_nodes,
     )
diff --git a/torch/_functorch/utils.py b/torch/_functorch/utils.py
index a2790a0fdd743..f7f00384c6013 100644
--- a/torch/_functorch/utils.py
+++ b/torch/_functorch/utils.py
@@ -1,6 +1,6 @@
 import contextlib
 from collections.abc import Generator
-from typing import Any, Union
+from typing import Any
 
 import torch
 from torch._C._functorch import (
@@ -37,4 +37,4 @@ def unwrap_dead_wrappers(args: tuple[Any, ...]) -> tuple[Any, ...]:
     return result
 
 
-argnums_t = Union[int, tuple[int, ...]]
+argnums_t = int | tuple[int, ...]
diff --git a/torch/_guards.py b/torch/_guards.py
index 4210f1ff89895..1ae422d999c6e 100644
--- a/torch/_guards.py
+++ b/torch/_guards.py
@@ -14,7 +14,7 @@
 from collections import defaultdict
 from contextlib import contextmanager
 from dataclasses import dataclass
-from typing import Any, Generic, NamedTuple, Optional, overload, TYPE_CHECKING, TypeVar
+from typing import Any, Generic, NamedTuple, overload, TYPE_CHECKING, TypeVar
 from typing_extensions import dataclass_transform
 
 import torch
@@ -36,7 +36,9 @@
 
     from torch._dynamo.backends.distributed import DDPOptimizerContext
     from torch._dynamo.codegen import PyCodegen
+    from torch._dynamo.guards import GuardCheckSpec
     from torch._functorch._aot_autograd.schemas import ViewAndMutationMeta
+    from torch._higher_order_ops.invoke_subgraph import NestedCompileRegionOptions
     from torch._subclasses.fake_tensor import FakeTensorMode
 
 
@@ -512,7 +514,7 @@ class GuardsCheckpointState:
     def __init__(self, dynamo_guards: OrderedSet[Guard]) -> None:
         self.dynamo_guards = dynamo_guards
 
-    def diff(self, other: GuardsCheckpointState) -> Optional[OrderedSet[Guard]]:
+    def diff(self, other: GuardsCheckpointState) -> OrderedSet[Guard] | None:
         """
         Produces a delta against another GuardsCheckpointState.
 
@@ -635,11 +637,20 @@ def restore_graphstate(self, state: GlobalContextCheckpointState) -> None:
 # Like a Set[Guard] but will record the user stack on all guards at the
 # time they were installed at their destination
 class GuardsSet:
-    def __init__(self, inner: Optional[OrderedSet[Guard]] = None) -> None:
+    def __init__(self, inner: OrderedSet[Guard] | None = None) -> None:
         if inner is None:
             self.inner: OrderedSet[Guard] = OrderedSet()
         else:
             self.inner = inner
+        # Map from source to list of guards with that source. Incrementally
+        # built in `add` method.
+        self.source_to_guards: defaultdict[Source, list[Guard]] = defaultdict(list)
+        for guard in self.inner:
+            self.track_guard_by_source(guard)
+
+    def track_guard_by_source(self, guard: Guard) -> None:
+        source = guard.originating_source
+        self.source_to_guards[source].append(guard)
 
     def __iter__(self) -> Iterator[Guard]:
         return iter(self.inner)
@@ -655,6 +666,10 @@ def __sub__(self, other: GuardsSet) -> GuardsSet:
     def __bool__(self) -> bool:
         return bool(self.inner)
 
+    def clear(self) -> None:
+        self.inner = OrderedSet()
+        self.source_to_guards = defaultdict(list)
+
     def add(
         self, guard: Guard, *, collect_debug_stack: bool = True, skip: int = 0
     ) -> None:
@@ -666,12 +681,17 @@ def add(
         if guard.user_stack is None:
             guard.user_stack = TracingContext.extract_stack()
         self.inner.add(guard)
+        self.track_guard_by_source(guard)
 
     def update(self, *others: set[Guard]) -> None:
         for o in others:
             for g in o:
                 self.add(g, skip=1)
 
+    def get_guards_for_source(self, source: Source) -> list[Guard]:
+        """Return all guards with the given originating_source."""
+        return list(self.source_to_guards[source])
+
     def remove_guards_with_source(self, source: Source) -> None:
         """Delete all guards that contains a given source"""
         from ._dynamo.source import is_from_source
@@ -679,6 +699,11 @@ def remove_guards_with_source(self, source: Source) -> None:
         self.inner = OrderedSet(
             g for g in self.inner if not is_from_source(g.originating_source, source)
         )
+        # Rebuild the index since is_from_source walks the chain, so
+        # multiple source keys may need removal.
+        self.source_to_guards = defaultdict(list)
+        for guard in self.inner:
+            self.track_guard_by_source(guard)
 
 
 """
@@ -693,6 +718,16 @@ class GuardsContext(Checkpointable[GuardsCheckpointState]):
     def __init__(self) -> None:
         self.dynamo_guards: GuardsSet = GuardsSet()
         self.aotautograd_guards: list[GuardEnvExpr] = []
+        self.skip_install: bool = False
+
+    @contextlib.contextmanager
+    def skip_guard_install(self) -> Generator[None, None, None]:
+        old = self.skip_install
+        self.skip_install = True
+        try:
+            yield
+        finally:
+            self.skip_install = old
 
     def copy_graphstate(self) -> GuardsCheckpointState:
         return GuardsCheckpointState(OrderedSet(self.dynamo_guards.inner))
@@ -706,10 +741,12 @@ def restore_graphstate(self, state: GuardsCheckpointState) -> None:
 
 class HopSubgraphCache:
     @abstractmethod
-    def add_dynamo_installed_submodule(self, fn_id: int, identifier: str) -> None: ...
+    def add_dynamo_installed_submodule(
+        self, fn_code: CodeType, identifier: str
+    ) -> None: ...
 
     @abstractmethod
-    def get_dynamo_installed_submodules(self, fn_id: int) -> list[str]: ...
+    def get_dynamo_installed_submodules(self, fn_code: CodeType) -> list[str]: ...
 
     @abstractmethod
     def add_autograd_key_entry(self, identifier: str, key: Callable) -> None: ...
@@ -737,23 +774,89 @@ def get_lazy_bwd_entry(
     ) -> tuple[torch.fx.GraphModule | None, int | None]: ...
 
 
+@dataclass
+class InvokeSubgraphReuseEntry:
+    body_name: str
+    body_gmod: torch.fx.GraphModule
+    config: NestedCompileRegionOptions | None
+    subgraph_input_mapping: list[
+        Any
+    ]  # list[LiftedArgOrigin] - defined in invoke_subgraph.py
+    single_tensor_output: bool
+    # Per-output tensor metadata (shape, stride, dtype, device, requires_grad)
+    # cached from the first trace so we can construct fresh FakeTensors on
+    # cache hit without re-running body_gmod.
+    output_metadata: list[
+        tuple[torch.Size, tuple[int, ...], torch.dtype, torch.device, bool]
+    ]
+    # 1-1 mapping to flat_vts: source for each flattened arg/kwarg, or None if
+    # the VT has no source. On cache hit, we build a source replacement mapping
+    # (old arg sources → new arg sources) to rewrite captured variable sources
+    # for the current invocation.
+    arg_sources: list[Source | None]
+    # Number of user-visible outputs (from the function return value).
+    # The graph may have additional outputs from side-effect intermediates;
+    # stamp_out_subgraph uses this to return only the user-visible slice.
+    num_user_outputs: int = 0
+
+
+@dataclass
+class InvokeSubgraphReuseCondition:
+    # Per flattened input VT: (InputTag, metadata).
+    #   (InputTag.TENSOR, TensorMetadata)
+    #   (InputTag.SYMNODE, sym_num — same object implies same symbol)
+    #   (InputTag.CONSTANT, value)
+    #   (InputTag.MODULE, None)
+    # Tensor metadata is checked here because TENSOR_MATCH guards for
+    # subgraph inputs may already exist before tracing and thus won't
+    # appear in the guard delta.
+    input_checks: list[tuple[Any, object]]  # list[tuple[InputTag, object]]
+
+    # Guards captured during the trace (delta + source-mapped).
+    # Each entry: (source, handler, expected_value, guard)
+    # handler is a pre-resolved GuardCheckSpec from GUARD_VALUE_DISPATCH.
+    guards: list[tuple[Source, GuardCheckSpec, object, Guard]]
+
+    # TreeSpec from pytree.tree_flatten of the (args, kwargs) structure.
+    # On cache hit, we verify the new call has the same treespec.
+    treespec: pytree.TreeSpec | None = None
+
+    # All sources accessed via VariableBuilder during the subgraph trace.
+    # On cache hit, we check if any modified VT's source is a base of one
+    # of these to detect mutations on captured variables.
+    traced_sources: OrderedSet[Source] = dataclasses.field(default_factory=OrderedSet)
+
+
 class InvokeSubgraphCache(HopSubgraphCache):
     def __init__(self) -> None:
         self.autograd_cache: dict[str, Callable] = {}
         self.proxy_dispatch_cache: dict[str, Callable] = {}
-        self.dynamo_installed_submodules: dict[int, list[str]] = defaultdict(list)
+        self.dynamo_installed_submodules: dict[CodeType, list[str]] = defaultdict(list)
         self.lazy_bwd_cache: dict[
             str, dict[tuple[object], tuple[torch.fx.GraphModule, int]]
         ] = defaultdict(dict)
         self.effects_cache: dict[
             str, set
         ] = {}  # Maps identifier -> set of effect types
+        # fn.__code__ → list of (condition, cache_entry) pairs. Walked linearly
+        # on lookup; first matching condition wins.
+        self.subgraph_reuse_cache: dict[
+            CodeType,
+            list[tuple[InvokeSubgraphReuseCondition, InvokeSubgraphReuseEntry]],
+        ] = defaultdict(list)
+        # fn_code → {hash_key → cache_entry}. Used by user-provided
+        # reuse_hash_fn for O(1) subgraph reuse lookup.
+        self.subgraph_reuse_key_cache: dict[
+            CodeType, dict[int, InvokeSubgraphReuseEntry]
+        ] = defaultdict(dict)
 
-    def add_dynamo_installed_submodule(self, fn_id: int, identifier: str) -> None:
-        self.dynamo_installed_submodules[fn_id].append(identifier)
+    def add_dynamo_installed_submodule(
+        self, fn_code: CodeType, identifier: str
+    ) -> None:
+        self.dynamo_installed_submodules[fn_code].append(identifier)
 
-    def get_dynamo_installed_submodules(self, fn_id: int) -> list[str]:
-        return self.dynamo_installed_submodules.get(fn_id, [])
+    def get_dynamo_installed_submodules(self, fn_code: CodeType) -> list[str]:
+        return self.dynamo_installed_submodules.get(fn_code, [])
 
     def add_autograd_key_entry(self, identifier: str, key: Callable) -> None:
         self.autograd_cache[identifier] = key
@@ -802,6 +905,65 @@ def get_effects(self, identifier: str) -> set | None:
         """Retrieve the effect types for a given invoke_subgraph identifier."""
         return self.effects_cache.get(identifier, None)
 
+    def add_reuse_entry(
+        self,
+        fn_code: CodeType,
+        condition: InvokeSubgraphReuseCondition,
+        entry: InvokeSubgraphReuseEntry,
+        max_reuse_entries: int = 8,
+    ) -> None:
+        entries = self.subgraph_reuse_cache[fn_code]
+        if len(entries) >= max_reuse_entries:
+            raise RuntimeError(
+                f"invoke_subgraph: exceeded maximum reuse entries "
+                f"({max_reuse_entries}) for function code {fn_code}. "
+                f"This most likely means a guard keeps failing on every "
+                f"invocation, preventing subgraph reuse. "
+                f"Set TORCH_LOGS='+hierarchical_compile' to identify which "
+                f"guard is failing. If reuse is genuinely not possible and "
+                f"you need more cache entries, increase the limit via the "
+                f"max_reuse_entries argument to nested_compile_region()."
+            )
+        entries.append((condition, entry))
+
+    def find_reuse_entry(
+        self,
+        fn_code: CodeType,
+        evaluator: Callable[
+            [InvokeSubgraphReuseCondition, InvokeSubgraphReuseEntry], bool
+        ],
+    ) -> InvokeSubgraphReuseEntry | None:
+        entries = self.subgraph_reuse_cache.get(fn_code, [])
+        for i, (condition, entry) in enumerate(entries):
+            if evaluator(condition, entry):
+                # MRU: move the hit entry to the front for faster future lookups
+                if i > 0:
+                    entries.insert(0, entries.pop(i))
+                return entry
+        return None
+
+    def find_reuse_entry_by_key(
+        self, fn_code: CodeType, hash_key: int
+    ) -> InvokeSubgraphReuseEntry | None:
+        return self.subgraph_reuse_key_cache.get(fn_code, {}).get(hash_key)
+
+    def add_reuse_entry_by_key(
+        self,
+        fn_code: CodeType,
+        hash_key: int,
+        entry: InvokeSubgraphReuseEntry,
+        max_reuse_entries: int = 8,
+    ) -> None:
+        key_cache = self.subgraph_reuse_key_cache[fn_code]
+        if len(key_cache) >= max_reuse_entries and hash_key not in key_cache:
+            raise RuntimeError(
+                f"invoke_subgraph: exceeded maximum reuse entries "
+                f"({max_reuse_entries}) for function code {fn_code} (hash-key path). "
+                f"Increase the limit via the max_reuse_entries argument to "
+                f"nested_compile_region()."
+            )
+        key_cache[hash_key] = entry
+
 
 class HopDispatchSetCache:
     def __init__(self) -> None:
@@ -956,6 +1118,7 @@ def __init__(self, fake_mode: FakeTensorMode | None) -> None:
         self.hop_dispatch_set_cache = HopDispatchSetCache()
         # list of code objects for inlined functions
         self.traced_code: list[CodeType] = []
+        self.cudagraph_annotation: Any = None
 
     def clear(self) -> None:
         # Look at the note in output_graph.py in function `save_global_state`
@@ -1154,7 +1317,7 @@ def __reduce__(self):
             # The _hash is a cached value that can be nondeterministically computed
             # (e.g., based on id() of objects), so it should not affect pickling.
             fields = dataclasses.fields(self)
-            field_values = tuple(getattr(self, f.name) for f in fields)
+            field_values = tuple(getattr(self, f.name) for f in fields if f.init)
             return (self.__class__, field_values)
 
         new_cls.__hash__ = __hash__
@@ -1206,7 +1369,7 @@ def get_value(
         self,
         globals: dict[str, Any],
         locals: dict[str, Any],
-        cache: weakref.WeakKeyDictionary[Source, Any],
+        cache: dict[Source, Any],
     ) -> Any:
         if self in cache:
             return cache[self]
@@ -1226,6 +1389,13 @@ def subguards_allowed(self) -> bool:
         """True if you can guard on attributes of this"""
         return self.guard_source != GuardSource.SYNTHETIC_LOCAL
 
+    def clone(self, transform_fn: Callable[[Source], Source] | None = None) -> Source:
+        # Frozen dataclass, so returning self is effectively a clone.
+        # Subclasses with mutable state should override.
+        if transform_fn is not None:
+            return transform_fn(self)
+        return self
+
 
 # Subclasses can be found in torch/_dynamo/source.py
 @dataclass_with_cached_hash(frozen=True)
@@ -1257,7 +1427,7 @@ def get_value(
         self,
         globals: dict[str, Any],
         locals: dict[str, Any],
-        cache: weakref.WeakKeyDictionary[Source, Any],
+        cache: dict[Source, Any],
     ) -> Any:
         if self in cache:
             return cache[self]
@@ -1272,6 +1442,19 @@ def get_value(
         cache[self] = value
         return value
 
+    def clone(self, transform_fn: Callable[[Source], Source] | None = None) -> Source:
+        cloned_fields: dict[str, Any] = {"base": self.base.clone(transform_fn)}
+        for f in dataclasses.fields(self):
+            if f.name == "base":
+                continue
+            val = getattr(self, f.name)
+            if isinstance(val, Source):
+                cloned_fields[f.name] = val.clone(transform_fn)
+        result = dataclasses.replace(self, **cloned_fields)
+        if transform_fn is not None:
+            result = transform_fn(result)
+        return result
+
 
 def detect_fake_mode(inputs: Any = None) -> FakeTensorMode | None:
     """
diff --git a/torch/_higher_order_ops/__init__.py b/torch/_higher_order_ops/__init__.py
index 29b2fc1ba4b33..09bd8d188112e 100644
--- a/torch/_higher_order_ops/__init__.py
+++ b/torch/_higher_order_ops/__init__.py
@@ -20,6 +20,7 @@
 )
 from torch._higher_order_ops.foreach_map import _foreach_map, foreach_map
 from torch._higher_order_ops.hints_wrap import hints_wrapper
+from torch._higher_order_ops.inline_asm_elementwise import inline_asm_elementwise
 from torch._higher_order_ops.invoke_leaf_function import invoke_leaf_function
 from torch._higher_order_ops.invoke_subgraph import invoke_subgraph
 from torch._higher_order_ops.local_map import local_map_hop
@@ -81,4 +82,5 @@
     "local_map_hop",
     "print",
     "inductor_compiled_code",
+    "inline_asm_elementwise",
 ]
diff --git a/torch/_higher_order_ops/_invoke_quant.py b/torch/_higher_order_ops/_invoke_quant.py
index b7a9fb94b93e2..4e125e68fa4f4 100644
--- a/torch/_higher_order_ops/_invoke_quant.py
+++ b/torch/_higher_order_ops/_invoke_quant.py
@@ -2,7 +2,6 @@
 # need to fix prim_hop_base type annotations first
 
 import dataclasses
-from typing import Optional
 
 import torch
 from torch._higher_order_ops.base_hop import BaseHOP, FunctionWithNoFreeVars
@@ -50,7 +49,7 @@ class InvokeQuant:
     def __call__(
         self,
         *args,
-        scheme: Optional[str] = None,
+        scheme: str | None = None,
         **kwargs,
     ):
         if not torch.compiler.is_compiling():
diff --git a/torch/_higher_order_ops/auto_functionalize.py b/torch/_higher_order_ops/auto_functionalize.py
index d58acb9fa89a0..7c2d8d9828f87 100644
--- a/torch/_higher_order_ops/auto_functionalize.py
+++ b/torch/_higher_order_ops/auto_functionalize.py
@@ -3,7 +3,7 @@
 from abc import ABC, abstractmethod
 from collections.abc import Callable, Sequence
 from dataclasses import dataclass
-from typing import Any, get_args, Optional, Union
+from typing import Any, get_args
 
 import torch
 import torch._library.utils as library_utils
@@ -71,8 +71,8 @@ def regenerate_view(self, bases_list: list[Tensor]):
 
 @dataclass
 class AsStridedViewInfo(ViewInfo):
-    size: Sequence[Union[int, torch.SymInt]]
-    stride: Sequence[Union[int, torch.SymInt]]
+    size: Sequence[int | torch.SymInt]
+    stride: Sequence[int | torch.SymInt]
     storage_offset: int
 
     def __init__(self, base_index, size, stride, storage_offset):
@@ -92,9 +92,9 @@ def regenerate_view(self, bases_list: list[Tensor]):
 
 @dataclass
 class SliceViewInfo(ViewInfo):
-    dim: Union[int, torch.SymInt]
-    start: Union[int, torch.SymInt]
-    end: Union[int, torch.SymInt]
+    dim: int | torch.SymInt
+    start: int | torch.SymInt
+    end: int | torch.SymInt
 
     def __init__(self, base_index, dim, start, end):
         super().__init__(base_index)
@@ -377,7 +377,7 @@ def __call__(
 auto_functionalized.fallthrough(DispatchKey.AutogradCUDA)
 
 
-_MutableOpType = Union[OpOverload, HigherOrderOperator]
+_MutableOpType = OpOverload | HigherOrderOperator
 
 
 class AutoFunctionalizedV2(HigherOrderOperator):
@@ -397,7 +397,7 @@ def __call__(
         _mutable_op: _MutableOpType,
         **kwargs: Any,
     ) -> tuple[Any, tuple[Tensor, ...]]:
-        _op_to_check: Optional[Union[OpOverload, HopInstance]] = None
+        _op_to_check: OpOverload | HopInstance | None = None
         if isinstance(_mutable_op, HigherOrderOperator):
             _op_to_check = HopInstance(
                 _mutable_op,
@@ -424,7 +424,7 @@ def __call__(
 
 
 def can_auto_functionalize(
-    op: Union[OperatorBase, HopInstance],
+    op: OperatorBase | HopInstance,
 ) -> bool:
     if isinstance(op, HopInstance):
         # HOPs that implement gen_schema and schema is not functional are auto_functionalizable.
@@ -545,9 +545,7 @@ def do_auto_functionalize(
     # List of the name of args that get mutated (according to the schema)
     mutable_args_names, _ = get_mutable_args(op)
 
-    unwrapped_actual_out: Union[Any, tuple[Any]] = unwrapped_outs[
-        : -len(mutable_args_names)
-    ]
+    unwrapped_actual_out: Any | tuple[Any] = unwrapped_outs[: -len(mutable_args_names)]
     unwrapped_mutable_out = unwrapped_outs[-len(mutable_args_names) :]
 
     if len(op._schema.returns) == 0:
@@ -608,13 +606,19 @@ def sync_update(o, orig_arg):
 class FunctionalCallableWithEpilogue:
     def __init__(self, orig_callable: Callable):
         self.orig_callable = orig_callable
+        # Propagate so callers pass inputs as a list, enabling input deallocation.
+        self._boxed_call = getattr(orig_callable, "_boxed_call", False)
 
     def __call__(self, *args, **kwargs):
-        # We call torch.func.functionalize. This allows us to inline the epilogue graph.
-        # Inlining has the benefit of allowing easiser fusion inside subgraph.
-        # Though the epilogue graph contains copy_, it is OK because inductor can handle it
-        # and this is also how we have been supporting top-level graph input mutation.
-        return tuple(torch.func.functionalize(self.orig_callable)(*args, **kwargs))
+        # Functionalize to inline the epilogue graph (copy_ ops) for better fusion.
+        functionalized = torch.func.functionalize(self.orig_callable)
+        if self._boxed_call:
+            # Not all callers respect _boxed_call (e.g. reenter_make_fx
+            # always calls f(*args)). Detect which convention was used.
+            if len(args) == 1 and isinstance(args[0], list):
+                return tuple(functionalized(args[0]))
+            return tuple(functionalized(list(args)))
+        return tuple(functionalized(*args, **kwargs))
 
     def __hash__(self):
         return id(self.orig_callable)
@@ -622,7 +626,7 @@ def __hash__(self):
 
 def do_auto_functionalize_v2(
     mode: "torch._subclasses.functional_tensor.FunctionalTensorMode",
-    op: Union[OpOverload, HopInstance],
+    op: OpOverload | HopInstance,
     args: tuple[Any, ...],
     kwargs: dict[str, Any],
 ) -> Any:
@@ -640,12 +644,24 @@ def do_auto_functionalize_v2(
     if not isinstance(op, get_args(_MutableOpType)):
         raise AssertionError(f"Expected _MutableOpType, got {type(op)}")
 
-    def _functionalize_callable(arg: Any):
-        if callable(arg):
+    subgraph_arg_names = {
+        arg_info.name
+        for arg_info in schema.arguments
+        if isinstance(arg_info.type, torch._C.AnyType)
+    }
+
+    def _maybe_functionalize(name: str, arg: Any) -> Any:
+        if name in subgraph_arg_names and callable(arg):
             return FunctionalCallableWithEpilogue(arg)
         return arg
 
-    args, kwargs = pytree.tree_map(_functionalize_callable, (args, kwargs))
+    args = tuple(
+        _maybe_functionalize(schema.arguments[i].name, a)
+        if i < len(schema.arguments)
+        else a
+        for i, a in enumerate(args)
+    )
+    kwargs = {k: _maybe_functionalize(k, v) for k, v in kwargs.items()}
 
     for idx, arg in enumerate(schema.arguments):
         # NB: torch_dispatch kwargs are the args defined as kwarg-only in the schema
@@ -740,7 +756,7 @@ def set_result(base_index):
             **auto_func_kwargs,  # type: ignore[arg-type]
         )
 
-    unwrapped_actual_out: Union[Any, tuple[Any]] = (
+    unwrapped_actual_out: Any | tuple[Any] = (
         unwrapped_outs if len(all_bases) == 0 else unwrapped_outs[: -len(all_bases)]
     )
 
@@ -810,7 +826,7 @@ def sync_update(o, orig_arg):
 @auto_functionalized.py_impl(DispatchKey.CompositeExplicitAutograd)
 def auto_functionalized_dense(
     _mutable_op: OpOverload,
-    _only_clone_these_tensors: Optional[tuple[str, ...]] = None,
+    _only_clone_these_tensors: tuple[str, ...] | None = None,
     **kwargs: Any,
 ) -> tuple[Any, tuple[Tensor, ...]]:
     new_kwargs = dict(**kwargs)
@@ -887,7 +903,7 @@ def auto_functionalized_func(ctx, _mutable_op, **kwargs):
 @auto_functionalized_v2.py_impl(DispatchKey.CompositeExplicitAutograd)
 def auto_functionalized_v2_dense(
     _mutable_op: _MutableOpType,
-    _only_clone_these_bases: Optional[tuple[int, ...]] = None,
+    _only_clone_these_bases: tuple[int, ...] | None = None,
     **kwargs: Any,
 ) -> tuple[Any, tuple[Tensor, ...]]:
     _all_bases: list[Tensor] = kwargs.pop("_all_bases", [])
@@ -900,7 +916,7 @@ def auto_functionalized_v2_dense(
         schema = pytree.tree_unflatten([], kwargs.pop("_op_schema")).schema
 
     if isinstance(_mutable_op, OpOverload):
-        _callable_op: Union[HopInstance, OpOverload] = _mutable_op
+        _callable_op: HopInstance | OpOverload = _mutable_op
     else:
         if not isinstance(schema, HopSchema):
             raise AssertionError(f"Expected HopSchema, got {type(schema)}")
@@ -1021,10 +1037,15 @@ def auto_functionalized_v2_proxy(
             HopInstance(_mutable_op, schema), tuple(), new_kwargs
         )
 
-        # Only replace the callabes in kwargs with the materialized subgraphs.
+        # Only replace the callables in kwargs with the materialized subgraphs.
         # The rest of the kwargs are kept unchanged.
+        subgraph_arg_names = {
+            arg_info.name
+            for arg_info in schema.arguments
+            if isinstance(arg_info.type, torch._C.AnyType)
+        }
         for k, v in kwargs.items():
-            if callable(v):
+            if k in subgraph_arg_names and callable(v):
                 if k not in materialized_kwargs or not isinstance(
                     materialized_kwargs[k], torch.fx.GraphModule
                 ):
diff --git a/torch/_higher_order_ops/cond.py b/torch/_higher_order_ops/cond.py
index f1e4a5b520bd3..4f05618d38806 100644
--- a/torch/_higher_order_ops/cond.py
+++ b/torch/_higher_order_ops/cond.py
@@ -5,7 +5,7 @@
 import logging
 import warnings
 from collections.abc import Callable
-from typing import Any, Optional, Union
+from typing import Any
 
 import torch
 import torch.utils._pytree as pytree
@@ -94,10 +94,10 @@ def gen_schema(self, pred, true_fn, false_fn, operands):
 
 @exposed_in("torch")
 def cond(
-    pred: Union[bool, int, float, torch.Tensor],
+    pred: bool | int | float | torch.Tensor,
     true_fn: Callable,
     false_fn: Callable,
-    operands: Union[tuple, list] = (),
+    operands: tuple | list = (),
 ) -> Any:
     r"""
     Conditionally applies `true_fn` or `false_fn`.
@@ -163,10 +163,34 @@ def false_fn(x: torch.Tensor):
           - The function must return a tensor with the same metadata, e.g. shape,
             dtype, etc.
 
-          - The function cannot have in-place mutations on inputs or global variables.
+          - The function cannot have in-place mutations on global variables.
             (Note: in-place tensor operations such as `add_` for intermediate results
             are allowed in a branch)
 
+          - The function can perform in-place mutations on its input tensors during inference (i.e.,
+            when `torch.is_grad_enabled()` is False).
+            Note: When using `torch.compile()` with a non-constant predicate, the outputs will always
+            be new tensors that do not share object identity with the original inputs.
+
+            Example::
+
+                def true_fn(x):
+                    return x.sin_()
+
+
+                def false_fn(x):
+                    return x + 1
+
+
+                def f(x):
+                    return cond(x.sum() > 0, true_fn, false_fn, (x,))
+
+
+                x = torch.ones(4)
+                with torch.no_grad():
+                    result = torch.compile(f)(x)
+                assert result is not x  # result is a new tensor, not the original x
+
     """
     if torch.compiler.is_dynamo_compiling():
         return cond_op(pred, true_fn, false_fn, operands)
@@ -217,12 +241,9 @@ def _validate_input(pred, true_fn, false_fn, operands):
     def _cond_op_wrapper(*args, **kwargs):
         return cond_op(*args, **kwargs)
 
-    from torch._higher_order_ops.utils import setup_compilation_env
+    from torch._higher_order_ops.utils import _hop_compile_and_call
 
-    with setup_compilation_env() as backend:
-        return torch.compile(_cond_op_wrapper, backend=backend, fullgraph=True)(
-            pred, true_fn, false_fn, operands
-        )
+    return _hop_compile_and_call(_cond_op_wrapper, (pred, true_fn, false_fn, operands))
 
 
 def trace_cond(proxy_mode, func_overload, pred, true_fn, false_fn, operands):
@@ -428,8 +449,8 @@ def _get_attr_maybe_call(t: torch.Tensor, attr_name: str) -> Any:
 
 
 def _merge_output(
-    a: Optional[Union[torch.Tensor, int]],
-    b: Optional[Union[torch.Tensor, int]],
+    a: torch.Tensor | int | None,
+    b: torch.Tensor | int | None,
     mode: FakeTensorMode,
 ):
     from torch.fx.experimental.symbolic_shapes import (
@@ -507,9 +528,9 @@ def _bound(s0, lower_bound: bool):
         u2 has range [5, 8]
         u3 has range [5, 7]
     """
-    merged_size: list[Union[int, torch.SymInt]] = []
+    merged_size: list[int | torch.SymInt] = []
 
-    def _has_unbacked_symbols(s: Union[int, torch.SymInt]) -> bool:
+    def _has_unbacked_symbols(s: int | torch.SymInt) -> bool:
         if isinstance(s, int):
             return False
         else:
@@ -585,19 +606,19 @@ def _bound_stride(
         b_ex_size: torch.Size,
         a_ex_stride: tuple[int, ...],
         b_ex_stride: tuple[int, ...],
-        merged_size: list[Union[int, torch.SymInt]],
-    ) -> list[Union[int, torch.SymInt]]:
+        merged_size: list[int | torch.SymInt],
+    ) -> list[int | torch.SymInt]:
         from torch._inductor.ir import get_stride_order
 
         a_sorted_stride_idx = get_stride_order(a_ex_stride, mode.shape_env)
         b_sorted_stride_idx = get_stride_order(b_ex_stride, mode.shape_env)
 
-        a_stride_li: list[Optional[tuple[Union[int, torch.SymInt], int]]] = [
-            None
-        ] * len(a_ex_stride)
-        b_stride_li: list[Optional[tuple[Union[int, torch.SymInt], int]]] = [
-            None
-        ] * len(b_ex_stride)
+        a_stride_li: list[tuple[int | torch.SymInt, int] | None] = [None] * len(
+            a_ex_stride
+        )
+        b_stride_li: list[tuple[int | torch.SymInt, int] | None] = [None] * len(
+            b_ex_stride
+        )
         for i, idx in enumerate(a_sorted_stride_idx):
             a_stride_li[idx] = (a_ex_stride[i], -i)
         for i, idx in enumerate(b_sorted_stride_idx):
@@ -619,14 +640,14 @@ def _bound_stride(
                     f"Consider using contiguous() to make the two branches have the same contiguousness."
                 )
 
-        def _maybe_expr(s: Union[int, torch.SymInt]):
+        def _maybe_expr(s: int | torch.SymInt):
             if isinstance(s, int):
                 return s
             return s.node.expr
 
-        a_stride_expr: dict[Any, Union[int, torch.SymInt]] = {}
-        b_stride_expr: dict[Any, Union[int, torch.SymInt]] = {}
-        merged_strides: list[Union[int, torch.SymInt]] = [None] * len(a_ex_stride)  # type: ignore[list-item]
+        a_stride_expr: dict[Any, int | torch.SymInt] = {}
+        b_stride_expr: dict[Any, int | torch.SymInt] = {}
+        merged_strides: list[int | torch.SymInt] = [None] * len(a_ex_stride)  # type: ignore[list-item]
         for a_pair, b_pair in zip(a_stride_li, b_stride_li):
             if a_pair is None or b_pair is None:
                 raise AssertionError(
@@ -676,7 +697,7 @@ def _maybe_expr(s: Union[int, torch.SymInt]):
             b_stride_expr[_maybe_expr(b_val * b_ex_size[i])] = nxt_merged_stride_expr
         return merged_strides
 
-    merged_stride: list[Union[int, torch.SymInt]] = _bound_stride(
+    merged_stride: list[int | torch.SymInt] = _bound_stride(
         a.size(), b.size(), a.stride(), b.stride(), merged_size
     )
 
@@ -688,7 +709,22 @@ def _maybe_expr(s: Union[int, torch.SymInt]):
 
 @cond_op.py_functionalize_impl
 def cond_func(ctx, pred, true_fn, false_fn, inputs):
-    from torch._higher_order_ops.utils import _check_alias_and_mutation
+    from torch._higher_order_ops.auto_functionalize import (
+        can_auto_functionalize,
+        do_auto_functionalize_v2,
+    )
+    from torch._higher_order_ops.utils import _check_alias_and_mutation, HopInstance
+
+    hop_instance = HopInstance.create(cond_op, pred, true_fn, false_fn, inputs)
+    # For now, we only support auto-functionalization for cond when using python
+    # functionalization mode
+    if can_auto_functionalize(hop_instance) and hasattr(ctx, "mode"):
+        return do_auto_functionalize_v2(
+            ctx.mode,
+            hop_instance,
+            tuple(pytree.tree_flatten((pred, true_fn, false_fn, inputs))[0]),
+            {},
+        )
 
     unwrapped_inputs = ctx.unwrap_tensors(inputs)
     unwrapped_pred = ctx.unwrap_tensors(pred)
diff --git a/torch/_higher_order_ops/cudagraph_conditional_nodes.py b/torch/_higher_order_ops/cudagraph_conditional_nodes.py
new file mode 100644
index 0000000000000..393119c740162
--- /dev/null
+++ b/torch/_higher_order_ops/cudagraph_conditional_nodes.py
@@ -0,0 +1,130 @@
+# mypy: allow-untyped-defs
+from collections.abc import Generator
+from contextlib import contextmanager
+
+import torch
+import torch.utils._pytree as pytree
+from torch.utils._python_dispatch import TorchDispatchMode
+
+
+class CUDAGraphCaptureControlFlowOpDispatchMode(TorchDispatchMode):
+    @classmethod
+    def ignore_compile_internals(cls) -> bool:
+        return True
+
+    def __init__(
+        self,
+    ) -> None:
+        self.supports_higher_order_operators = True
+        super().__init__()
+
+    def __torch_dispatch__(
+        self,
+        func,
+        types,
+        args=(),
+        kwargs=None,
+    ):
+        if func is torch.ops.higher_order.cond:
+            # Re-enter the mode to support nested conditionals
+            with self:
+                return if_else_node(*args)
+        kwargs = {} if kwargs is None else kwargs
+        return func(*args, **kwargs)
+
+
+class ControlFlowOpWarmupDispatchMode(TorchDispatchMode):
+    """The purpose of this TodchDispatchMode is to "warm up" both sides of a torch.cond() statement.
+
+    For data-dependent control flow code, only one side will be
+    executed. Therefore, it is not safe to stream capture a
+    torch.cond() statement naively, since we don't have a guarantee
+    that all ops will have been "warmed up". The clever workaround is
+    to use a "relaxed" stream capture whose final cuda graph we throw
+    away. This works because stream capture does not actually execute
+    any GPU code, and because true_fn and false_fn are both fxgraphs,
+    which do not have any CPU side effects.
+    """
+
+    @classmethod
+    def ignore_compile_internals(cls) -> bool:
+        return True
+
+    def __init__(
+        self,
+    ) -> None:
+        super().__init__()
+        self.supports_higher_order_operators = True
+        self.capture_stream = torch.cuda.Stream()
+
+    def __torch_dispatch__(
+        self,
+        func,
+        types,
+        args=(),
+        kwargs=None,
+    ):
+        kwargs = {} if kwargs is None else kwargs
+
+        # Warm up both sides of this torch.cond()
+        if func is torch.ops.higher_order.cond:
+            if torch.cuda.is_current_stream_capturing():
+                # This is a call to torch.cond() nested within another
+                # torch.cond() function.
+                with self:
+                    # We re-enter the mode in case of nested calls to torch.cond()
+                    return if_else_node(*args)
+            else:
+                with (
+                    torch.cuda.graph(
+                        torch.cuda.CUDAGraph(),
+                        pool=None,
+                        stream=self.capture_stream,
+                        capture_error_mode="relaxed",
+                    ),
+                    self,
+                ):
+                    if_else_node(*args)
+
+                return func(*args, **kwargs)
+        else:
+            return func(*args, **kwargs)
+
+
+@contextmanager
+def _if_body(pred: torch.Tensor) -> Generator[None, None, None]:
+    current_cuda_graph = torch.cuda.CUDAGraph.get_currently_capturing_graph()
+    current_cuda_graph.begin_capture_to_if_node(pred)
+    try:
+        yield
+    finally:
+        current_cuda_graph.end_capture_to_conditional_node()
+
+
+def if_else_node(pred: torch.Tensor, true_fn, false_fn, operands):
+    if not pred.is_cuda:
+        raise ValueError(
+            "Conditions must be on a cuda device to use conditional node in cuda graphs"
+        )
+    # if-else is not supported until CUDA 12.8. Therefore, we use two
+    # if conditions, where one evaluates !pred
+    outs = []
+
+    for lazy_pred, fn in [
+        (lambda: pred, true_fn),
+        (lambda: torch.logical_not(pred), false_fn),
+    ]:
+        with _if_body(lazy_pred()):
+            outs.append(fn(*operands))
+
+            # The output of the else branch gets copied into the
+            # output of the if branch. This is done because the rest
+            # of the cudagraph after the conditional node has fixed
+            # inputs, so we need to merge the two outputs into a
+            # single output.
+            if len(outs) == 2:
+                for if_out, else_out in zip(
+                    pytree.tree_iter(outs[0]), pytree.tree_iter(outs[1])
+                ):
+                    if_out.copy_(else_out)
+    return outs[0]
diff --git a/torch/_higher_order_ops/effects.py b/torch/_higher_order_ops/effects.py
index 783605ad9f13f..7f2d203df50c7 100644
--- a/torch/_higher_order_ops/effects.py
+++ b/torch/_higher_order_ops/effects.py
@@ -1,5 +1,5 @@
 # mypy: allow-untyped-defs
-from typing import Any, Optional, Union
+from typing import Any, Union
 
 import torch
 import torch.utils._pytree as pytree
@@ -46,7 +46,7 @@ def _get_op_qualname(op: _op_identifier) -> str:
 
 
 def _register_effectful_op(
-    op: _op_identifier, effect: Optional[EffectType]
+    op: _op_identifier, effect: EffectType | None
 ) -> RegistrationHandle:
     qualname = _get_op_qualname(op)
     entry = torch._library.simple_registry.singleton.find(qualname)
@@ -54,7 +54,7 @@ def _register_effectful_op(
     return handle
 
 
-def _get_effect(op: _op_identifier) -> Optional[_EffectType]:
+def _get_effect(op: _op_identifier) -> _EffectType | None:
     qualname = _get_op_qualname(op)
     entry = torch._library.simple_registry.singleton.find(qualname)
     return entry.effect.effect
@@ -220,7 +220,7 @@ def with_effects_functional(
 _EFFECTFUL_HOPS_WITH_SCHEMA = {hop_print, invoke_leaf_function}
 
 
-def _get_schema(op, args, kwargs: Optional[dict] = None) -> torch.FunctionSchema:
+def _get_schema(op, args, kwargs: dict | None = None) -> torch.FunctionSchema:
     if isinstance(op, torch._ops.OpOverload):
         return op._schema
     elif op == call_torchbind:
diff --git a/torch/_higher_order_ops/flat_apply.py b/torch/_higher_order_ops/flat_apply.py
index 10e9f711aadfa..5f93ec46a19da 100644
--- a/torch/_higher_order_ops/flat_apply.py
+++ b/torch/_higher_order_ops/flat_apply.py
@@ -123,9 +123,13 @@ def __call__(
         registered with pytree.register_constant. The constant type goes directly
         into the spec.
         """
-        if not (isinstance(func, _op_types) or pytree._is_constant_holder(func)):
+        if not (
+            isinstance(func, _op_types)
+            or is_opaque_type(type(func))
+            or pytree._is_constant_holder(func)
+        ):
             raise AssertionError(
-                f"func must be an op type or constant holder, got {type(func)}"
+                f"func must be an op type, constant holder, or opaque callable, got {type(func)}"
             )
         if len(_unused) != 0:
             raise AssertionError(f"unexpected keyword arguments: {_unused}")
@@ -161,9 +165,10 @@ def impl(
                 f"expected retrieved constant to be _ConstantFunction, got {type(func)}"
             )
 
-    # pyrefly: ignore[bad-argument-type]  # pyrefly bug?
-    args, kwargs = from_graphable(flat_args, in_spec)
-    out = func(*args, **kwargs)
+    from torch._higher_order_ops.invoke_leaf_function import unflatten_args_with_modules
+
+    with unflatten_args_with_modules(flat_args, in_spec) as (args, kwargs):
+        out = func(*args, **kwargs)
 
     if checked_output:
         # For "normal" usage all outputs must either be graphable or
diff --git a/torch/_higher_order_ops/flex_attention.py b/torch/_higher_order_ops/flex_attention.py
index 41266cc66acaf..9949aba8b69b2 100644
--- a/torch/_higher_order_ops/flex_attention.py
+++ b/torch/_higher_order_ops/flex_attention.py
@@ -1,6 +1,6 @@
 import math
 from collections.abc import Callable, Sequence
-from typing import Any, Optional, Union
+from typing import Any
 
 import torch
 import torch.utils._pytree as pytree
@@ -86,7 +86,7 @@ def _permute_strides(out: torch.Tensor, query_strides: tuple[int, ...]) -> torch
         fill_order = [last_dim] + fill_order
         out_strides = _construct_strides(out.shape, fill_order)
 
-    new_out = out.new_empty(out.shape).as_strided(out.shape, out_strides)
+    new_out = out.new_empty_strided(out.shape, out_strides)
     new_out.copy_(out)
     return new_out
 
@@ -138,7 +138,7 @@ def __call__(
         logsumexp: torch.Tensor,
         grad_out: torch.Tensor | None,
         grad_logsumexp: torch.Tensor | None,
-        fw_graph: Union[Callable, GraphModule],
+        fw_graph: Callable | GraphModule,
         joint_graph: GraphModule,
         block_mask: tuple,
         scale: float,
@@ -146,7 +146,7 @@ def __call__(
         score_mod_other_buffers: tuple = (),
         mask_mod_other_buffers: tuple = (),
     ) -> tuple[
-        torch.Tensor, torch.Tensor, torch.Tensor, tuple[Optional[torch.Tensor], ...]
+        torch.Tensor, torch.Tensor, torch.Tensor, tuple[torch.Tensor | None, ...]
     ]:
         validate_subgraph_args_types(score_mod_other_buffers + mask_mod_other_buffers)
 
@@ -669,8 +669,8 @@ def create_fw_bw_graph(
         with disable_proxy_modes_tracing():
 
             def _from_fun(
-                t: Union[Tensor, torch.SymInt, int],
-            ) -> Union[Tensor, torch.SymInt, int]:
+                t: Tensor | torch.SymInt | int,
+            ) -> Tensor | torch.SymInt | int:
                 if isinstance(t, torch.Tensor):
                     return torch.empty_strided(
                         t.size(),
@@ -812,7 +812,7 @@ def backward(  # type: ignore[override]
         grad_out: Tensor,
         grad_logsumexp: Tensor,
         grad_max_scores: Tensor,
-    ) -> tuple[Optional[Tensor], ...]:
+    ) -> tuple[Tensor | None, ...]:
         fw_args = saved_values(ctx)
         (
             query,
@@ -959,15 +959,24 @@ def sdpa_dense_backward(
     kernel_options: dict[str, Any],
     score_mod_other_buffers: tuple,
     mask_mod_other_buffers: tuple,
-) -> tuple[
-    torch.Tensor, torch.Tensor, torch.Tensor, tuple[Optional[torch.Tensor], ...]
-]:
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, tuple[torch.Tensor | None, ...]]:
     if query.dtype != key.dtype or query.dtype != value.dtype:
         raise ValueError(
             f"Backward pass with mixed query, key, and value dtype is not supported, "
             f"got query.dtype={query.dtype}, key.dtype={key.dtype}, "
             f"and value.dtype={value.dtype}"
         )
+    if joint_graph is None:
+        example_vals = (
+            query.new_zeros((), requires_grad=True),
+            query.new_zeros((), dtype=torch.int),
+            query.new_zeros((), dtype=torch.int),
+            query.new_zeros((), dtype=torch.int),
+            query.new_zeros((), dtype=torch.int),
+        )
+        _, joint_graph = create_fw_bw_graph(
+            fw_graph, example_vals, score_mod_other_buffers
+        )
     from torch._dynamo._trace_wrapped_higher_order_op import TransformGetItemToIndex
 
     Bq, Hq, seq_len_q, qk_head_dim = query.shape
@@ -984,8 +993,8 @@ def sdpa_dense_backward(
     actual_grad_value = _permute_strides(actual_grad_value, value.stride())
 
     def _maybe_new_buffer(
-        buffer: Union[torch.Tensor, torch.SymInt, int],
-    ) -> Optional[Union[torch.Tensor, torch.SymInt, int]]:
+        buffer: torch.Tensor | torch.SymInt | int,
+    ) -> torch.Tensor | torch.SymInt | int | None:
         if isinstance(buffer, torch.Tensor):
             return (
                 torch.empty_like(buffer, memory_format=torch.contiguous_format)
@@ -1014,7 +1023,9 @@ def _maybe_new_buffer(
     if grad_logsumexp is None:
         grad_logsumexp = torch.zeros_like(logsumexp)
 
-    # We're undoing the log -> log2 change of base in the forwards
+    # logsumexp is expected in log2 scale (as returned by the forward HOP).
+    # The public flex_attention API converts lse to natural log before returning,
+    # so callers using the public API must not pass that value here directly.
     logsumexp = logsumexp * math.log(2)
     # The backwards formula for the log -> log2 change of base in the forwards
     grad_logsumexp = grad_logsumexp / math.log(2)
@@ -1135,16 +1146,14 @@ def trace_flex_attention_backward(
     logsumexp: torch.Tensor,
     grad_out: torch.Tensor,
     grad_logsumexp: torch.Tensor,
-    fw_graph: Union[Callable, GraphModule],
+    fw_graph: Callable | GraphModule,
     joint_graph: GraphModule,
     block_mask: tuple,
     scale: float,
     kernel_options: dict[str, Any],
     score_mod_other_buffers: tuple = (),
     mask_mod_other_buffers: tuple = (),
-) -> tuple[
-    torch.Tensor, torch.Tensor, torch.Tensor, tuple[Optional[torch.Tensor], ...]
-]:
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, tuple[torch.Tensor | None, ...]]:
     """We already have the forward graph and joint graph from the forward pass, so we create a proxy attach both graphs"""
     from torch._dynamo._trace_wrapped_higher_order_op import TransformGetItemToIndex
 
@@ -1239,16 +1248,14 @@ def flex_attention_backward_proxy_torch_dispatch_mode(
     logsumexp: torch.Tensor,
     grad_out: torch.Tensor,
     grad_logsumexp: torch.Tensor,
-    fw_graph: Union[Callable, GraphModule],
+    fw_graph: Callable | GraphModule,
     joint_graph: GraphModule,
     block_mask: tuple,
     scale: float,
     kernel_options: dict[str, Any],
     score_mod_other_buffers: tuple = (),
     mask_mod_other_buffers: tuple = (),
-) -> tuple[
-    torch.Tensor, torch.Tensor, torch.Tensor, tuple[Optional[torch.Tensor], ...]
-]:
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, tuple[torch.Tensor | None, ...]]:
     if mode is None:
         raise AssertionError("Mode should always be enabled for python fallback key")
     with torch.fx.experimental.proxy_tensor.set_original_aten_op(
@@ -1283,16 +1290,14 @@ def flex_attention_backward_functionalize(
     logsumexp: torch.Tensor,
     grad_out: torch.Tensor,
     grad_logsumexp: torch.Tensor,
-    fw_graph: Union[Callable, GraphModule],
+    fw_graph: Callable | GraphModule,
     joint_graph: GraphModule,
     block_mask: tuple,
     scale: float,
     kernel_options: dict[str, Any],
     score_mod_other_buffers: tuple = (),
     mask_mod_other_buffers: tuple = (),
-) -> tuple[
-    torch.Tensor, torch.Tensor, torch.Tensor, tuple[Optional[torch.Tensor], ...]
-]:
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, tuple[torch.Tensor | None, ...]]:
     """Defines the functionalization rules for the flex_attention operator.
 
     Write now we are unwrapping each tensor and then redispatching to the next,
@@ -1397,16 +1402,14 @@ def flex_attention_backward_fake_tensor_mode(
     logsumexp: torch.Tensor,
     grad_out: torch.Tensor,
     grad_logsumexp: torch.Tensor,
-    fw_graph: Union[Callable, GraphModule],
+    fw_graph: Callable | GraphModule,
     joint_graph: GraphModule,
     block_mask: tuple,
     scale: float,
     kernel_options: dict[str, Any],
     score_mod_other_buffers: tuple = (),
     mask_mod_other_buffers: tuple = (),
-) -> tuple[
-    torch.Tensor, torch.Tensor, torch.Tensor, tuple[Optional[torch.Tensor], ...]
-]:
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, tuple[torch.Tensor | None, ...]]:
     if has_user_subclass(
         (
             query,
diff --git a/torch/_higher_order_ops/inline_asm_elementwise.py b/torch/_higher_order_ops/inline_asm_elementwise.py
new file mode 100644
index 0000000000000..f23f8dc626503
--- /dev/null
+++ b/torch/_higher_order_ops/inline_asm_elementwise.py
@@ -0,0 +1,299 @@
+# mypy: allow-untyped-defs
+import functools
+import re
+
+import torch
+import torch.utils._pytree as pytree
+from torch._C import DispatchKey
+from torch._higher_order_ops.utils import autograd_not_implemented
+from torch._ops import HigherOrderOperator
+from torch._subclasses.fake_tensor import FakeTensorMode
+from torch.fx.experimental.proxy_tensor import ProxyTorchDispatchMode, track_tensor_tree
+
+
+__all__ = ["inline_asm_elementwise"]
+
+
+class InlineAsmElementwiseOp(HigherOrderOperator):
+    """Execute inline PTX assembly elementwise over tensors.
+
+    This is an elementwise map where the function body is inline assembly.
+    Input tensors are implicitly broadcast to the same shape.
+
+    Each invocation of the inline asm processes ``pack`` elements at a time.
+    Exactly which set of inputs a given invocation receives is unspecified.
+
+    Output strides follow PyTorch's standard pointwise striding propagation
+    rules.
+
+    In eager mode, the assembly is executed via the CUDA Jiterator.  Under
+    ``torch.compile`` the assembly is lowered to Triton's
+    ``tl.inline_asm_elementwise`` via Inductor, which allows fusion with
+    surrounding operators.
+
+    Args:
+        *inputs: Input tensors whose values are passed to the asm block.
+        asm_str: PTX assembly string. Operands use ``$N`` syntax
+            (e.g. ``$0`` for the first output, ``$1`` for the first input).
+        constraints: Inline-asm constraints in LLVM format. Output constraints
+            are prefixed with ``=`` (e.g. ``"=f,f,f"`` for one float output
+            and two float inputs).
+        dtype: Element type of the returned tensor.
+        is_pure: Must be ``True``. If true, the compiler may assume the asm
+            block has no side-effects.
+        pack: Number of elements processed per asm invocation.  When
+            ``pack > 1``, the constraint string must list ``pack`` outputs
+            and ``pack`` copies of each input.  Requires ``torch.compile``.
+
+    Returns:
+        A tensor with the broadcast shape of the inputs and the given dtype.
+
+    Example::
+
+        >>> # xdoctest: +SKIP(requires CUDA)
+        >>> # Float32 fused multiply-add via PTX
+        >>> result = inline_asm_elementwise(
+        ...     a, b, c,
+        ...     asm_str="fma.rn.f32 $0, $1, $2, $3;",
+        ...     constraints="=f,f,f,f",
+        ...     dtype=torch.float32,
+        ... )
+
+        >>> # xdoctest: +SKIP(requires CUDA)
+        >>> # pack=2: each asm invocation processes two elements
+        >>> result = inline_asm_elementwise(
+        ...     x,
+        ...     asm_str="mov.b32 $0, $2; mov.b32 $1, $3;",
+        ...     constraints="=r,=r,r,r",
+        ...     dtype=torch.float32,
+        ...     pack=2,
+        ... )
+    """
+
+    def __init__(self):
+        super().__init__("inline_asm_elementwise")
+
+    def __call__(
+        self,
+        *inputs: torch.Tensor,
+        asm_str: str,
+        constraints: str,
+        dtype: torch.dtype,
+        is_pure: bool = True,
+        pack: int = 1,
+    ) -> torch.Tensor:
+        if not is_pure:
+            raise ValueError("inline_asm_elementwise only supports is_pure=True")
+        # pyrefly: ignore [missing-attribute]
+        return super().__call__(
+            *inputs,
+            asm_str=asm_str,
+            constraints=constraints,
+            dtype=dtype,
+            is_pure=True,
+            pack=pack,
+        )
+
+
+inline_asm_elementwise = InlineAsmElementwiseOp()
+
+
+def _parse_constraints(constraints: str) -> tuple[int, int]:
+    parts = [p.strip() for p in constraints.split(",")]
+    n_outputs = sum(1 for p in parts if p.startswith("="))
+    n_inputs = len(parts) - n_outputs
+    return n_outputs, n_inputs
+
+
+_DTYPE_TO_CUDA_TYPE = {
+    torch.float32: "float",
+    torch.float64: "double",
+    torch.float16: "__half",
+    torch.bfloat16: "__nv_bfloat16",
+    torch.int32: "int",
+    torch.int64: "long long",
+    torch.int16: "short",
+    torch.int8: "signed char",
+    torch.uint8: "unsigned char",
+    torch.uint16: "unsigned short",
+    torch.uint32: "unsigned int",
+    torch.bool: "bool",
+}
+
+
+_TRITON_ARG_RE = re.compile(r"\$(\d+)")
+
+
+def _triton_asm_to_cuda_asm(asm_str: str) -> str:
+    return _TRITON_ARG_RE.sub(r"%\1", asm_str)
+
+
+@functools.lru_cache
+def _get_jiterator_fn(
+    asm_str: str,
+    constraints: str,
+    n_inputs: int,
+    input_dtype: torch.dtype,
+    output_dtype: torch.dtype,
+):
+    from torch.cuda.jiterator import _create_jit_fn
+
+    cuda_asm = _triton_asm_to_cuda_asm(asm_str)
+
+    constraint_parts = [p.strip() for p in constraints.split(",")]
+    output_constraints = [p.lstrip("=") for p in constraint_parts if p.startswith("=")]
+    input_constraints = [p for p in constraint_parts if not p.startswith("=")]
+
+    if input_dtype not in _DTYPE_TO_CUDA_TYPE:
+        raise ValueError(f"Unsupported input dtype for inline asm: {input_dtype}")
+    if output_dtype not in _DTYPE_TO_CUDA_TYPE:
+        raise ValueError(f"Unsupported output dtype for inline asm: {output_dtype}")
+
+    input_type = _DTYPE_TO_CUDA_TYPE[input_dtype]
+    output_type = _DTYPE_TO_CUDA_TYPE[output_dtype]
+
+    input_params = ", ".join(f"{input_type} in{i}" for i in range(n_inputs))
+    out_constraints_str = ", ".join(f'"={c}"(result)' for c in output_constraints)
+    in_constraints_str = ", ".join(
+        f'"{c}"(in{i})' for i, c in enumerate(input_constraints)
+    )
+    escaped_asm = (
+        cuda_asm.replace("\\", "\\\\").replace('"', '\\"').replace("\n", "\\n")
+    )
+
+    code = f"""
+template <typename T>
+{output_type} inline_asm_kernel({input_params}) {{
+    {output_type} result;
+    asm volatile (
+        "{escaped_asm}"
+        : {out_constraints_str}
+        : {in_constraints_str}
+    );
+    return result;
+}}
+"""
+
+    return _create_jit_fn(code)
+
+
+def _inline_asm_dense(*inputs, asm_str, constraints, dtype, is_pure, pack):
+    if not inputs:
+        raise ValueError("inline_asm_elementwise requires at least one input tensor")
+
+    inputs = torch.broadcast_tensors(*inputs)
+
+    if not inputs[0].is_cuda:
+        raise RuntimeError("inline_asm_elementwise only supports CUDA tensors")
+
+    if pack > 1:
+        raise RuntimeError(
+            "inline_asm_elementwise with pack > 1 requires torch.compile"
+        )
+
+    n_outputs, n_inputs = _parse_constraints(constraints)
+
+    if n_outputs != 1:
+        raise ValueError(f"Expected 1 output constraint, got {n_outputs}")
+
+    if n_inputs != len(inputs):
+        raise ValueError(
+            f"Constraint string specifies {n_inputs} inputs but got "
+            f"{len(inputs)} tensor(s)"
+        )
+
+    # Jiterator generates a single input type for all inputs — mixed dtypes
+    # would produce incorrect CUDA code.
+    input_dtypes = {inp.dtype for inp in inputs}
+    if len(input_dtypes) > 1:
+        raise ValueError(
+            f"All inputs must have the same dtype for eager execution, "
+            f"got {sorted(str(d) for d in input_dtypes)}"
+        )
+
+    jit_fn = _get_jiterator_fn(
+        asm_str=asm_str,
+        constraints=constraints,
+        n_inputs=len(inputs),
+        input_dtype=inputs[0].dtype,
+        output_dtype=dtype,
+    )
+
+    return jit_fn(*inputs)
+
+
+@inline_asm_elementwise.py_impl(DispatchKey.CompositeExplicitAutograd)
+def _(*inputs, asm_str, constraints, dtype, is_pure=True, pack=1):
+    return _inline_asm_dense(
+        *inputs,
+        asm_str=asm_str,
+        constraints=constraints,
+        dtype=dtype,
+        is_pure=is_pure,
+        pack=pack,
+    )
+
+
+inline_asm_elementwise.py_autograd_impl(
+    autograd_not_implemented(inline_asm_elementwise, deferred_error=True)
+)
+
+
+def _elementwise_output_like(*inputs, dtype):
+    from torch._prims_common import compute_elementwise_output_logical_to_physical_perm
+
+    broadcasted = torch.broadcast_tensors(*inputs)
+    l2p_perm, _ = compute_elementwise_output_logical_to_physical_perm(*broadcasted)
+    return torch.empty_permuted(
+        broadcasted[0].shape, l2p_perm, dtype=dtype, device=broadcasted[0].device
+    )
+
+
+@inline_asm_elementwise.py_impl(FakeTensorMode)
+def _(mode, *inputs, asm_str, constraints, dtype, is_pure=True, pack=1):
+    with mode:
+        return _elementwise_output_like(*inputs, dtype=dtype)
+
+
+@inline_asm_elementwise.py_impl(ProxyTorchDispatchMode)
+def _(mode, *inputs, asm_str, constraints, dtype, is_pure=True, pack=1):
+    proxy_args = pytree.tree_map(mode.tracer.unwrap_proxy, inputs)
+
+    out_proxy = mode.tracer.create_proxy(
+        "call_function",
+        inline_asm_elementwise,
+        proxy_args,
+        {
+            "asm_str": asm_str,
+            "constraints": constraints,
+            "dtype": dtype,
+            "is_pure": is_pure,
+            "pack": pack,
+        },
+        name="inline_asm_elementwise",
+    )
+
+    out = inline_asm_elementwise(
+        *inputs,
+        asm_str=asm_str,
+        constraints=constraints,
+        dtype=dtype,
+        is_pure=is_pure,
+        pack=pack,
+    )
+    return track_tensor_tree(out, out_proxy, constant=None, tracer=mode.tracer)
+
+
+@inline_asm_elementwise.py_functionalize_impl
+def _(ctx, *inputs, asm_str, constraints, dtype, is_pure=True, pack=1):
+    unwrapped_inputs = ctx.unwrap_tensors(inputs)
+
+    with ctx.redispatch_to_next():
+        res = inline_asm_elementwise(
+            *unwrapped_inputs,
+            asm_str=asm_str,
+            constraints=constraints,
+            dtype=dtype,
+            pack=pack,
+        )
+    return ctx.wrap_tensors(res)
diff --git a/torch/_higher_order_ops/invoke_leaf_function.py b/torch/_higher_order_ops/invoke_leaf_function.py
index 9c9d2fe93ccbf..9746e4696cfb8 100644
--- a/torch/_higher_order_ops/invoke_leaf_function.py
+++ b/torch/_higher_order_ops/invoke_leaf_function.py
@@ -8,16 +8,56 @@
 import torch.utils._pytree as pytree
 from torch._C import DispatchKey, DispatchKeySet
 from torch._higher_order_ops.utils import register_fake
+from torch._library.opaque_object import OpaqueBase, register_opaque_type
 from torch._ops import HigherOrderOperator
 from torch.autograd.graph import get_gradient_edge
 from torch.fx.experimental.proxy_tensor import ProxyTorchDispatchMode, track_tensor_tree
 from torch.nn.utils.stateless import _reparametrize_module
 
-from .flat_apply import func_to_graphable
-
 
 _leaf_function_module_retriever: Callable[[int], Any] | None = None
 
+# Separate storage for the make_fx / _invoke_leaf_function_python path.
+# We can't use Dynamo's register_user_object because it requires a Source
+# (for bytecode generation) and adds entries to index_to_bytecode_constructor.
+# In the make_fx path we have neither a Source nor bytecode generation, so we
+# maintain our own dict keyed by negative indices to avoid collisions with
+# Dynamo's non-negative indices.
+_makefx_module_storage: dict[int, torch.nn.Module] = {}
+_makefx_next_index = 0
+
+
+def store_makefx_modules(modules: list[torch.nn.Module]) -> tuple[int, ...]:
+    """Store modules for the make_fx path and return their assigned indices.
+
+    Uses negative indices to avoid collisions with Dynamo's register_user_object
+    which uses non-negative indices.
+    """
+    global _makefx_next_index
+    indices = []
+    for mod in modules:
+        _makefx_next_index -= 1
+        _makefx_module_storage[_makefx_next_index] = mod
+        indices.append(_makefx_next_index)
+    return tuple(indices)
+
+
+def reset_makefx_module_storage() -> None:
+    global _makefx_next_index
+    _makefx_next_index = 0
+    _makefx_module_storage.clear()
+
+
+class _LeafCallable(OpaqueBase):
+    def __init__(self, fn: Callable) -> None:
+        self._fn = fn
+
+    def __call__(self, *args: Any, **kwargs: Any) -> Any:
+        return self._fn(*args, **kwargs)
+
+
+register_opaque_type(_LeafCallable, typ="reference")
+
 
 def set_leaf_function_module_retriever(retriever: Callable[[int], Any]) -> None:
     global _leaf_function_module_retriever
@@ -38,6 +78,92 @@ class LeafModuleState(NamedTuple):
     named_buffers: dict[str, torch.Tensor]
 
 
+def convert_modules_to_states(values: Any, module_to_index: dict[int, int]) -> Any:
+    """Replace nn.Module instances in a pytree with LeafModuleState objects.
+
+    Args:
+        values: A pytree of values that may contain nn.Module instances.
+        module_to_index: Mapping from id(module) to its integer index.
+    """
+
+    def module_to_state(val: Any) -> Any:
+        if isinstance(val, torch.nn.Module):
+            return LeafModuleState(
+                nn_module_index=module_to_index[id(val)],
+                named_parameters=dict(val.named_parameters()),
+                named_buffers=dict(val.named_buffers()),
+            )
+        return val
+
+    return pytree.tree_map(module_to_state, values)
+
+
+def _resolve_mutated_flat_indices(
+    fn: Callable,
+    mutates_args: frozenset[str],
+    num_flat_args: int,
+    input_spec: pytree.TreeSpec,
+) -> str:
+    """Resolve mutates_args expressions to a comma-separated string of flat-arg indices.
+
+    Each expression in mutates_args (e.g. "x", "model.running_mean") is evaluated
+    against sentinel values to determine which flat-arg positions are mutated.
+
+    Example: for ``def fn(x, model)`` where model is an nn.Module with parameters
+    ``weight`` and ``bias``, the flat args are ``[x, nn_module_index, weight, bias]``.
+    Given ``mutates_args={"model.weight"}``, this assigns sentinels ``[0, 1, 2, 3]``
+    to the flat args, evaluates ``model.weight`` to ``2``, and returns ``"2"``.
+    """
+    import inspect
+
+    class _AttrDict:
+        pass
+
+    def _set_nested_attr(obj: _AttrDict, fqn: str, value: Any) -> None:
+        parts = fqn.split(".")
+        for part in parts[:-1]:
+            if not hasattr(obj, part):
+                setattr(obj, part, _AttrDict())
+            obj = getattr(obj, part)
+        setattr(obj, parts[-1], value)
+
+    def _lms_to_attr_dict(val: Any) -> Any:
+        if isinstance(val, LeafModuleState):
+            target = _AttrDict()
+            for fqn, sentinel in val.named_parameters.items():
+                _set_nested_attr(target, fqn, sentinel)
+            for fqn, sentinel in val.named_buffers.items():
+                _set_nested_attr(target, fqn, sentinel)
+            return target
+        return val
+
+    sig = inspect.signature(fn)
+    sentinels = list(range(num_flat_args))
+    args_struct, kwargs_struct = pytree.tree_unflatten(sentinels, input_spec)
+    args_eval, kwargs_eval = pytree.tree_map(
+        _lms_to_attr_dict,
+        (args_struct, kwargs_struct),
+        is_leaf=lambda x: isinstance(x, LeafModuleState),
+    )
+    namespace = dict(sig.bind(*args_eval, **kwargs_eval).arguments)
+
+    indices: list[int] = []
+    for expr in mutates_args:
+        # Empty __builtins__ prevents access to builtins like __import__, open, exec.
+        result = eval(expr, {"__builtins__": {}}, namespace)  # noqa: S307
+        leaves = pytree.tree_leaves(result)
+        for sentinel in leaves:
+            if not isinstance(sentinel, int):
+                raise ValueError(
+                    f"mutates_args expression '{expr}' resolved to a non-leaf value "
+                    f"of type {type(sentinel).__name__}. Expressions must resolve to "
+                    f"individual tensor positions, e.g. 'model.weight' not 'model'."
+                )
+            indices.append(sentinel)
+    indices.sort()
+    return ",".join(str(i) for i in indices)
+
+
 @dataclass
 class GradientInfo:
     """
@@ -55,11 +181,16 @@ class GradientInfo:
     device: torch.device
 
 
-def unwrap_fn_spec(fn_spec: pytree.TreeSpec) -> Callable:
-    return pytree.tree_unflatten((), fn_spec)
-
-
 def _retrieve_module_by_index(nn_module_index: int) -> torch.nn.Module:
+    # Check make_fx storage first (used by _invoke_leaf_function_python).
+    # Fall back to the Dynamo retriever (used by the compiled path).
+    if nn_module_index in _makefx_module_storage:
+        if nn_module_index >= 0:
+            raise RuntimeError(
+                f"Expected negative nn_module_index for non-strict trace over leaf_function, but got {nn_module_index}."
+            )
+        return _makefx_module_storage[nn_module_index]
+
     if _leaf_function_module_retriever is None:
         raise RuntimeError("Leaf function module retriever not set.")
 
@@ -151,13 +282,9 @@ def check_escaped_gradients(
 
 
 @contextlib.contextmanager
-def reconstruct_original_args(
-    input_spec: pytree.TreeSpec | None, flat_args: tuple[Any, ...]
+def unflatten_args_with_modules(
+    flat_args: tuple[Any, ...], input_spec: pytree.TreeSpec
 ) -> Generator[tuple[list[Any] | tuple[Any, ...], dict[str, Any]], None, None]:
-    if input_spec is None:
-        yield flat_args, {}
-        return
-
     args, kwargs = pytree.tree_unflatten(flat_args, input_spec)
 
     with contextlib.ExitStack() as stack:
@@ -181,6 +308,65 @@ def process_module_state(state: LeafModuleState) -> torch.nn.Module:
         yield new_args, new_kwargs
 
 
+def flatten_args_with_modules(
+    args_kwargs: tuple[Any, ...],
+) -> list[Any]:
+    def expand_module(x: Any) -> Any:
+        if isinstance(x, torch.nn.Module):
+            return LeafModuleState(
+                nn_module_index=-1,
+                named_parameters=dict(x.named_parameters()),
+                named_buffers=dict(x.named_buffers()),
+            )
+        return x
+
+    expanded = pytree.tree_map(expand_module, args_kwargs)
+    return pytree.tree_leaves(expanded)
+
+
+def make_leaf_function_wrappers(
+    real_fn: Callable[..., Any],
+    fake_fn: Callable[..., Any],
+    captured_out_spec: list[pytree.TreeSpec | None],
+) -> tuple[Callable[..., tuple[Any, ...]], Callable[..., tuple[Any, ...]]]:
+    """Wrap real_fn and fake_fn to flatten outputs and capture the output TreeSpec.
+
+    Both wrappers share the same captured output spec: the first call (typically
+    fake_fn during tracing) records it, and subsequent calls verify consistency.
+    The caller passes in a single-element list and reads captured_out_spec[0]
+    after the wrappers have been called.
+
+    Used by both the Dynamo path (_call_leaf_function in torch.py) and the
+    make_fx path (_invoke_leaf_function_python in decorators.py).
+    """
+
+    def _wrap(fn: Callable[..., Any]) -> Callable[..., tuple[Any, ...]]:
+        if len(captured_out_spec) != 1:
+            raise RuntimeError(
+                f"captured_out_spec must be a single-element list, got length {len(captured_out_spec)}"
+            )
+
+        def wrapper(*args: Any, **kwargs: Any) -> tuple[Any, ...]:
+            out = fn(*args, **kwargs)
+
+            flat_out, out_spec = pytree.tree_flatten(out)
+            if captured_out_spec[0] is None:
+                captured_out_spec[0] = out_spec
+            elif captured_out_spec[0] != out_spec:
+                raise AssertionError(
+                    f"leaf_function output structure mismatch: "
+                    f"expected {captured_out_spec[0]}, got {out_spec}. "
+                    f"This can happen if the real function and fake function return "
+                    f"different pytree structures (e.g., dict vs tuple, different number "
+                    f"of elements). Ensure both functions return the same structure."
+                )
+            return tuple(flat_out)
+
+        return wrapper
+
+    return _wrap(real_fn), _wrap(fake_fn)
+
+
 def autograd_grad_with_gradient_info(
     output_infos: Sequence[GradientInfo | None],
     input_infos: Sequence[GradientInfo | None],
@@ -259,42 +445,29 @@ def autograd_grad_with_gradient_info(
 
 def _make_forward(
     fn: Callable,
-    requires_grad_indices: set[int],
     include_keys: DispatchKeySet,
     exclude_keys: DispatchKeySet,
 ) -> tuple[Callable, dict[str, Any]]:
     state: dict[str, Any] = {"inputs": None, "outputs": None}
 
     @functools.wraps(fn)
-    def forward(*flat_args):
-        # Detach all tensor inputs to isolate from external autograd context
-        flat_args = [
-            arg.detach() if isinstance(arg, torch.Tensor) else arg for arg in flat_args
-        ]
-
-        # Restore requires_grad state that was captured at tracing time.
-        # We need this because at runtime with aot_eager, the forward graph
-        # runs inside autograd.Function.forward(), which is a no-grad context.
-        # This means intermediate tensors can lose their requires_grad status
-        # by the time they reach invoke_leaf_function, even though they had it
-        # during tracing.
-        inputs = [
-            arg.requires_grad_(True)
-            if isinstance(arg, torch.Tensor) and idx in requires_grad_indices
-            else arg
-            for idx, arg in enumerate(flat_args)
-        ]
-
-        # NB: we capture dispatch keys at creation time (during tracing), where PythonDispatcher
-        # is active, but at runtime it's not so we remove it from the effective keys.
+    def forward(*args, **kwargs):
         effective_keys = include_keys
         if include_keys.has(DispatchKey.PythonDispatcher):
-            effective_keys = include_keys.remove(DispatchKey.PythonDispatcher)
+            effective_keys = effective_keys.remove(DispatchKey.PythonDispatcher)
+        if effective_keys.has(DispatchKey.Python):
+            effective_keys = effective_keys.remove(DispatchKey.Python)
         with torch._C._ForceDispatchKeyGuard(effective_keys, exclude_keys):
             with torch.enable_grad():
-                outputs = fn(*inputs)
+                outputs = fn(*args, **kwargs)
 
-                check_escaped_gradients(outputs, inputs, requires_grad_indices)
+                flat_inputs = flatten_args_with_modules((args, kwargs))
+                requires_grad_indices = {
+                    i
+                    for i, inp in enumerate(flat_inputs)
+                    if isinstance(inp, torch.Tensor) and inp.requires_grad
+                }
+                check_escaped_gradients(outputs, flat_inputs, requires_grad_indices)
 
                 state["inputs"] = tuple(
                     GradientInfo(
@@ -306,7 +479,7 @@ def forward(*flat_args):
                     )
                     if isinstance(inp, torch.Tensor) and inp.requires_grad
                     else None
-                    for inp in inputs
+                    for inp in flat_inputs
                 )
 
                 if outputs is None:
@@ -338,28 +511,85 @@ def forward(*flat_args):
 
 class InvokeLeafFunction(HigherOrderOperator):
     def __init__(self):
-        super().__init__("invoke_leaf_function")
-
-    def __call__(self, real_fn_spec, fake_fn_spec, *flat_args):
+        super().__init__("invoke_leaf_function", supports_training_input_mutation=True)
+
+    def __call__(
+        self,
+        real_fn_callable,
+        fake_fn_callable,
+        input_spec,
+        mutated_arg_indices,
+        *flat_args,
+        requires_grad_indices="",
+    ):
         """
-        real_fn_spec: pytree.TreeSpec for the real function that's wrapped in dynamo
-        fake_fn_spec: pytree.TreeSpec for the fake function that's wrapped in dynamo
+        real_fn_callable: _LeafCallable wrapping the real function
+        fake_fn_callable: _LeafCallable wrapping the fake function
+        input_spec: pytree.TreeSpec for unflattening flat_args back to (args, kwargs)
+        mutated_arg_indices: comma-separated string of flat-arg indices that are
+            declared as mutated (e.g. "1,2"), or "" for no mutations. Encoded as a
+            string so it is a pytree leaf for the HOP schema infrastructure.
+        requires_grad_indices: comma-separated string of flat-arg indices that
+            require grad (e.g. "0,1"), or "" for none.
         """
-        return super().__call__(real_fn_spec, fake_fn_spec, *flat_args)  # type: ignore[attr-defined]
+        return super().__call__(  # type: ignore[attr-defined]
+            real_fn_callable,
+            fake_fn_callable,
+            input_spec,
+            mutated_arg_indices,
+            *flat_args,
+            requires_grad_indices=requires_grad_indices,
+        )
 
     # pyrefly: ignore [bad-override]
-    def gen_schema(self, real_fn_spec, fake_fn_spec, *flat_args):
+    def gen_schema(
+        self,
+        real_fn_callable,
+        fake_fn_callable,
+        input_spec,
+        mutated_arg_indices,
+        *flat_args,
+        requires_grad_indices="",
+    ):
         from torch._higher_order_ops.schema import HopSchemaGenerator
         from torch._higher_order_ops.utils import _maybe_fake_prop_ignore_unbacked
+        from torch.fx.experimental.proxy_tensor import disable_proxy_modes_tracing
 
-        fake_fn = unwrap_fn_spec(fake_fn_spec)
-        fake_outputs = _maybe_fake_prop_ignore_unbacked(fake_fn, flat_args)
+        mutated_set = _parse_mutated_arg_indices(mutated_arg_indices)
+
+        with disable_proxy_modes_tracing():
+            if mutated_set:
+                schema_flat_args = tuple(
+                    arg.detach().clone()
+                    if isinstance(arg, torch.Tensor) and i in mutated_set
+                    else arg
+                    for i, arg in enumerate(flat_args)
+                )
+            else:
+                schema_flat_args = flat_args
+
+            def run_fake(*unfunc_flat_args):
+                with unflatten_args_with_modules(unfunc_flat_args, input_spec) as (
+                    args,
+                    kwargs,
+                ):
+                    return fake_fn_callable(*args, **kwargs)
+
+            fake_outputs = _maybe_fake_prop_ignore_unbacked(run_fake, schema_flat_args)
 
         gen = HopSchemaGenerator(self)
-        gen.add_arg("real_fn_spec", real_fn_spec)
-        gen.add_arg("fake_fn_spec", fake_fn_spec)
+        gen.add_arg("real_fn_callable", real_fn_callable)
+        gen.add_arg("fake_fn_callable", fake_fn_callable)
+        gen.add_arg("input_spec", input_spec)
+        gen.add_arg("mutated_arg_indices", mutated_arg_indices)
         for i, arg in enumerate(flat_args):
-            gen.add_arg(f"arg{i}", arg)
+            gen.add_arg(f"arg{i}", arg, is_mutated=i in mutated_set)
+        gen.add_arg(
+            "requires_grad_indices",
+            requires_grad_indices,
+            default_value="",
+            kw_only=True,
+        )
 
         if isinstance(fake_outputs, tuple):
             for out in fake_outputs:
@@ -379,12 +609,12 @@ def gen_schema(self, real_fn_spec, fake_fn_spec, *flat_args):
 
 # NOTE: [Autograd support for invoke_leaf_function]
 #
-# The oveerall idea is that when the real forward executes, we are going to build an autograd graph
+# The overall idea is that when the real forward executes, we are going to build an autograd graph
 # and save it.  When the real backward executes, we are going to invoke the autograd graph.
 # We need to build these "real_forward" and "real_backward" functions from the "real_fn".
 #
 # Inputs:
-# real_fn_spec/fake_fn_spec are pytree.TreeSpecs that contain real_fn and fake_fn.
+# real_fn_callable/fake_fn_callable are _LeafCallable objects that wrap real_fn and fake_fn.
 # These functions were created in dynamo by wrapping the user's original leaf function and fake function:
 #   - They accept *flat_args (flattened LeafModuleState objects + other args)
 #   - They unflatten flat_args and convert LeafModuleState back to nn.Modules
@@ -407,20 +637,25 @@ def gen_schema(self, real_fn_spec, fake_fn_spec, *flat_args):
 class InvokeLeafFunctionAutogradOp(torch.autograd.Function):
     @staticmethod
     # pyrefly: ignore [bad-override]
-    def forward(ctx, real_fn_spec, fake_fn_spec, *flat_args):
-        real_fn = unwrap_fn_spec(real_fn_spec)
-
+    def forward(
+        ctx,
+        real_fn_callable,
+        fake_fn_callable,
+        input_spec,
+        mutated_arg_indices,
+        *flat_args,
+    ):
         include_keys = torch._C._dispatch_tls_local_include_set()
         exclude_keys = torch._C._dispatch_tls_local_exclude_set()
 
-        requires_grad_indices = {
-            i
+        requires_grad_indices = ",".join(
+            str(i)
             for i, arg in enumerate(flat_args)
             if isinstance(arg, torch.Tensor) and arg.requires_grad
-        }
+        )
 
         real_forward, real_state = _make_forward(
-            real_fn, requires_grad_indices, include_keys, exclude_keys
+            real_fn_callable, include_keys, exclude_keys
         )
 
         def real_backward(*grads):
@@ -458,11 +693,16 @@ def fake_backward(*grads):
                 for info in input_infos_for_fake
             )
 
-        _, new_real_fn_spec = func_to_graphable(real_forward)
+        new_real_fn_callable = _LeafCallable(real_forward)
 
         with torch._C._AutoDispatchBelowAutograd():
             fw_outputs = invoke_leaf_function(
-                new_real_fn_spec, fake_fn_spec, *flat_args
+                new_real_fn_callable,
+                fake_fn_callable,
+                input_spec,
+                mutated_arg_indices,
+                *flat_args,
+                requires_grad_indices=requires_grad_indices,
             )
 
         ctx.real_backward = real_backward
@@ -473,20 +713,43 @@ def fake_backward(*grads):
     @staticmethod
     # pyrefly: ignore [bad-override]
     def backward(ctx, *grads):
-        _, real_bw_spec = func_to_graphable(ctx.real_backward)
-        _, fake_bw_spec = func_to_graphable(ctx.fake_backward)
-        fw_grads = invoke_leaf_function(real_bw_spec, fake_bw_spec, *grads)
-        return None, None, *fw_grads
+        real_bw_callable = _LeafCallable(ctx.real_backward)
+        fake_bw_callable = _LeafCallable(ctx.fake_backward)
+        _, bw_input_spec = pytree.tree_flatten((grads, {}))
+        fw_grads = invoke_leaf_function(
+            real_bw_callable, fake_bw_callable, bw_input_spec, "", *grads
+        )
+        return None, None, None, None, *fw_grads
 
 
 @invoke_leaf_function.py_autograd_impl
-def invoke_leaf_function_autograd(real_fn_spec, fake_fn_spec, *flat_args):
-    return InvokeLeafFunctionAutogradOp.apply(real_fn_spec, fake_fn_spec, *flat_args)
+def invoke_leaf_function_autograd(
+    real_fn_callable,
+    fake_fn_callable,
+    input_spec,
+    mutated_arg_indices,
+    *flat_args,
+    requires_grad_indices="",
+):
+    return InvokeLeafFunctionAutogradOp.apply(
+        real_fn_callable, fake_fn_callable, input_spec, mutated_arg_indices, *flat_args
+    )
 
 
-# TODO: allow user annotated mutation and aliasing info
+# TODO: aliasing is not allowed
 @invoke_leaf_function.py_functionalize_impl
-def invoke_leaf_function_functionalization(ctx, *all_args):
+def invoke_leaf_function_functionalization(ctx, *all_args, **kwargs):
+    from torch._higher_order_ops.auto_functionalize import (
+        can_auto_functionalize,
+        do_auto_functionalize_v2,
+    )
+    from torch._higher_order_ops.utils import HopInstance
+
+    unwrapped_args = ctx.unwrap_tensors(all_args)
+    hop_instance = HopInstance.create(invoke_leaf_function, *unwrapped_args, **kwargs)
+    if can_auto_functionalize(hop_instance):
+        return do_auto_functionalize_v2(ctx.mode, hop_instance, all_args, kwargs)
+
     from torch._higher_order_ops.effects import handle_effects
 
     return handle_effects(
@@ -494,16 +757,16 @@ def invoke_leaf_function_functionalization(ctx, *all_args):
         ctx.mode._tokens,
         invoke_leaf_function,
         all_args,
-        {},
+        kwargs,
     )
 
 
 @invoke_leaf_function.py_impl(ProxyTorchDispatchMode)
-def invoke_leaf_function_proxy_mode(proxy_mode, *all_args):
-    out = invoke_leaf_function(*all_args)
+def invoke_leaf_function_proxy_mode(proxy_mode, *all_args, **kwargs):
+    out = invoke_leaf_function(*all_args, **kwargs)
     proxies = pytree.tree_map(proxy_mode.tracer.unwrap_proxy, all_args)
     proxy = proxy_mode.tracer.create_proxy(
-        "call_function", invoke_leaf_function, proxies, {}
+        "call_function", invoke_leaf_function, proxies, kwargs
     )
     return track_tensor_tree(out, proxy, constant=None, tracer=proxy_mode.tracer)
 
@@ -566,43 +829,69 @@ def _validate_outputs_match(
                 )
 
 
+def _parse_mutated_arg_indices(s: str) -> set[int]:
+    return {int(x) for x in s.split(",") if x}
+
+
 def _check_no_input_mutation(
     flat_args: tuple[Any, ...],
     version_before: list[int],
+    mutated_arg_indices: str = "",
 ) -> None:
+    mutated_set = _parse_mutated_arg_indices(mutated_arg_indices)
     for i, arg in enumerate(flat_args):
-        if isinstance(arg, torch.Tensor):
-            if arg._version != version_before[i]:
+        if isinstance(arg, torch.Tensor) and arg._version != version_before[i]:
+            if i not in mutated_set:
                 raise RuntimeError(
-                    f"In-place mutation detected on input tensor at position {i} "
-                    f"(in the pytree-flattened inputs with nn.Module states expanded) in "
-                    f"@leaf_function. In-place mutations on inputs are not supported yet."
-                    f"Consider cloning the input before mutating it."
+                    f"Undeclared in-place mutation on input tensor at position {i}. "
+                    f"Declare it in @leaf_function(mutates_args=...) or avoid mutating inputs."
                 )
 
 
 @register_fake(invoke_leaf_function)
-def invoke_leaf_function_fake(real_fn_spec, fake_fn_spec, *flat_args):
-    fake_fn = unwrap_fn_spec(fake_fn_spec)
-    return fake_fn(*flat_args)
+def invoke_leaf_function_fake(
+    real_fn_callable,
+    fake_fn_callable,
+    input_spec,
+    mutated_arg_indices,
+    *flat_args,
+    requires_grad_indices="",
+):
+    with unflatten_args_with_modules(flat_args, input_spec) as (args, kwargs):
+        return fake_fn_callable(*args, **kwargs)
 
 
 @invoke_leaf_function.py_impl(DispatchKey.CompositeExplicitAutograd)
-def invoke_leaf_function_dense(real_fn_spec, fake_fn_spec, *flat_args):
+def invoke_leaf_function_dense(
+    real_fn_callable,
+    fake_fn_callable,
+    input_spec,
+    mutated_arg_indices,
+    *flat_args,
+    requires_grad_indices="",
+):
     from torch._dynamo import config as dynamo_config
 
     version_before = [
         arg._version if isinstance(arg, torch.Tensor) else 0 for arg in flat_args
     ]
 
-    real_fn = unwrap_fn_spec(real_fn_spec)
-    real_output = real_fn(*flat_args)
+    flat_args = tuple(
+        arg.detach() if isinstance(arg, torch.Tensor) else arg for arg in flat_args
+    )
+    requires_grad_indices_set = _parse_mutated_arg_indices(requires_grad_indices)
+    flat_args = tuple(
+        arg.requires_grad_(True) if idx in requires_grad_indices_set else arg
+        for idx, arg in enumerate(flat_args)
+    )
+
+    with unflatten_args_with_modules(flat_args, input_spec) as (args, kwargs):
+        real_output = real_fn_callable(*args, **kwargs)
 
-    _check_no_input_mutation(flat_args, version_before)
+        _check_no_input_mutation(flat_args, version_before, mutated_arg_indices)
 
-    if dynamo_config.leaf_function_validate_outputs:
-        fake_fn = unwrap_fn_spec(fake_fn_spec)
-        fake_output = fake_fn(*flat_args)
-        _validate_outputs_match(fake_output, real_output)
+        if dynamo_config.leaf_function_validate_outputs:
+            fake_output = fake_fn_callable(*args, **kwargs)
+            _validate_outputs_match(fake_output, real_output)
 
     return real_output
diff --git a/torch/_higher_order_ops/invoke_subgraph.py b/torch/_higher_order_ops/invoke_subgraph.py
index 53c4171c87d9c..6a79caeb6315a 100644
--- a/torch/_higher_order_ops/invoke_subgraph.py
+++ b/torch/_higher_order_ops/invoke_subgraph.py
@@ -7,7 +7,7 @@
 from collections.abc import Callable
 from contextlib import nullcontext
 from dataclasses import dataclass, field
-from typing import Any, Optional, Union
+from typing import Any
 
 import torch
 import torch.utils._pytree as pytree
@@ -50,7 +50,7 @@
 # InvokeSubgraphAutogradOp.
 @dataclass
 class OutputMetadata:
-    num_fw_outs: Optional[int] = None
+    num_fw_outs: int | None = None
     indexes_with_symint: set[int] = field(default_factory=set)
     indexes_with_no_grad: set[int] = field(default_factory=set)
 
@@ -61,8 +61,8 @@ class OutputMetadata:
 class NestedCompileRegionOptions:
     # A Callable that takes (gm, example_inputs, decompositions=None, **kwargs) as inputs.
     # Returns AOTCompiledArtifact
-    fw_compiler: Optional[Callable] = None
-    bw_compiler: Optional[Callable] = None
+    fw_compiler: Callable | None = None
+    bw_compiler: Callable | None = None
 
     # Note: [InvokeSubgraphHOP Partitioner]
     # If not None, add "partitioner" to HOP node meta.
@@ -70,11 +70,11 @@ class NestedCompileRegionOptions:
     # If str, the options are "default_partition" and "min_cut_rematerialization_partition".
     # The HOP joint graph will be partitioned using the corresponding functions in
     # torch/_functorch/partitioners.py
-    partitioner: Optional[Callable | str] = None
+    partitioner: Callable | str | None = None
 
     # If it's None, we'll inherit the parent call's decompositions.
     # Otherwise, the nested region will use this decompositions.
-    decompositions: Optional[dict[str, Any]] = None
+    decompositions: dict[str, Any] | None = None
 
 
 def _extract_nested_region_config(fn):
@@ -115,8 +115,8 @@ def __init__(self) -> None:
     # identifying two invoke_subgraph calls have same subgraph.
     def __call__(
         self,
-        subgraph: Union[GraphModule, FunctionalizeCtxWrapper],
-        identifier: Optional[str],
+        subgraph: GraphModule | FunctionalizeCtxWrapper,
+        identifier: str | None,
         *operands,
     ):
         if identifier is not None and not isinstance(identifier, str):
@@ -174,7 +174,7 @@ def gen_schema(self, subgraph, identifier, *operands):
 
 
 def invoke_subgraph_infer(
-    subgraph: Union[GraphModule, FunctionalizeCtxWrapper],
+    subgraph: GraphModule | FunctionalizeCtxWrapper,
     *operands,
 ):
     """Inference-only entrypoint for invoke_subgraph that auto-generates identifier.
@@ -252,19 +252,19 @@ def invoke_subgraph_placeholder(func, *args, **kwargs):
         def _invoke_subgraph_placeholder_wrapper(func, args):
             return invoke_subgraph_placeholder(func, *args)
 
-        from torch._higher_order_ops.utils import setup_compilation_env
+        from torch._higher_order_ops.utils import _hop_compile_and_call
 
-        with setup_compilation_env() as backend:
-            return torch.compile(
-                _invoke_subgraph_placeholder_wrapper,
-                backend=backend,
-                fullgraph=True,
-            )(func, args)
+        return _hop_compile_and_call(_invoke_subgraph_placeholder_wrapper, (func, args))
 
     return func(*args, **kwargs)
 
 
-def mark_compile_region(fn=None, options: Optional[NestedCompileRegionOptions] = None):
+def mark_compile_region(
+    fn=None,
+    options: NestedCompileRegionOptions | None = None,
+    max_reuse_entries: int = 8,
+    reuse_hash_fn=None,
+):
     """
     This wrapper instructs torch.compile to compile the wrapped region once and
     reuse the compiled artifact, instead of the usual way of aggressively
@@ -278,6 +278,9 @@ def mark_compile_region(fn=None, options: Optional[NestedCompileRegionOptions] =
         options: Optional config to use for compiling the subgraph.
             Warning: this is an experimental feature under development and
             not ready for use yet.
+        max_reuse_entries: Maximum number of reuse cache entries per function
+            before raising an error. If this limit is hit, guards keep failing
+            across invocations and hierarchical compilation is not effective.
     """
 
     def wrap(func):
@@ -290,6 +293,8 @@ def inner(*args, **kwargs):
 
         inner.__marked_compile_region_fn__ = func  # type: ignore[attr-defined]
         func.__marked_compile_region_config__ = options  # type: ignore[attr-defined]
+        func.__marked_compile_region_max_reuse_entries__ = max_reuse_entries  # type: ignore[attr-defined]
+        func.__marked_compile_region_reuse_hash_fn__ = reuse_hash_fn  # type: ignore[attr-defined]
 
         return inner
 
@@ -897,6 +902,8 @@ def _(subgraph, identifier, *operands):
     from torch._dynamo.utils import dynamo_timed
 
     with dynamo_timed("invoke_subgraph_fake_tensor", log_pt2_compile_event=True):
+        if getattr(subgraph, "_boxed_call", False):
+            return subgraph(list(operands))
         return subgraph(*operands)
 
 
diff --git a/torch/_higher_order_ops/local_map.py b/torch/_higher_order_ops/local_map.py
index 205a049021159..e691eda36b057 100644
--- a/torch/_higher_order_ops/local_map.py
+++ b/torch/_higher_order_ops/local_map.py
@@ -9,7 +9,7 @@
 import functools
 from collections.abc import Callable, Generator, Sequence
 from contextlib import contextmanager
-from typing import Any, Optional, TypeAlias
+from typing import Any, TypeAlias
 
 import torch
 import torch.utils._pytree as pytree
@@ -49,8 +49,8 @@ def defer_inlining() -> Generator[None, None, None]:
 # Used to unwrap tensors classes like FunctionalTensor and Parameter
 def _new_tensor(
     t: Any,
-    new_shape: Optional[Sequence[int]] = None,
-    new_stride: Optional[Sequence[int]] = None,
+    new_shape: Sequence[int] | None = None,
+    new_stride: Sequence[int] | None = None,
 ) -> Any:
     if isinstance(t, torch.Tensor):
         if type(t) not in (FunctionalTensor, FakeTensor, torch.Tensor):
@@ -113,7 +113,7 @@ def _redistribute(
 
 
 def redistribute_fw_inputs(
-    global_args: Any, all_placements: Any, mesh: Any, _: Optional[int] = None
+    global_args: Any, all_placements: Any, mesh: Any, _: int | None = None
 ) -> GraphArg:
     if len(global_args) != len(all_placements):
         raise AssertionError(
@@ -174,7 +174,7 @@ def redistribute_bw_inputs(
 
 
 def redistribute_bw_outputs(
-    local_outs: Any, all_placements: Any, mesh: Any, _: Optional[int] = None
+    local_outs: Any, all_placements: Any, mesh: Any, _: int | None = None
 ) -> GraphArg:
     if len(local_outs) != len(all_placements):
         raise AssertionError(
@@ -463,7 +463,7 @@ def forward(
         filtered_grads_idx: set[int],
         *args: Any,
         **kwargs: Any,
-    ) -> tuple[Optional[torch.Tensor], ...]:
+    ) -> tuple[torch.Tensor | None, ...]:
         from torch._functorch._aot_autograd.schemas import MemoryFormatMeta
 
         ctx.bw_gm = bw_gm
@@ -477,15 +477,19 @@ def forward(
         saved_activations = fw_outs_with_saved_activations[num_fw_outs:]
         save_values_for_backward(ctx, saved_activations)
 
+        # Force memory_format path (not exact size/stride) because local_map forward
+        # operates on local shapes but backward receives global-shaped tangents.
+        # TODO(ivankobzarev): Support exact size/stride by converting between local/global shapes.
         ctx.expected_tangent_metadata = {
-            i: MemoryFormatMeta.from_tensor(fw_outs[i]) for i in filtered_grads_idx
+            i: MemoryFormatMeta.from_tensor(fw_outs[i], force_use_memory_format=True)
+            for i in filtered_grads_idx
         }
         return fw_outs
 
     @staticmethod
     def backward(
         ctx: Any, *_grads: tuple[torch.Tensor]
-    ) -> tuple[Optional[torch.Tensor], ...]:
+    ) -> tuple[torch.Tensor | None, ...]:
         from torch._functorch._aot_autograd.runtime_wrappers import (
             coerce_to_expected_memory_format,
         )
diff --git a/torch/_higher_order_ops/map.py b/torch/_higher_order_ops/map.py
index a63a135ff38cc..cb54eddfb7a13 100644
--- a/torch/_higher_order_ops/map.py
+++ b/torch/_higher_order_ops/map.py
@@ -1,7 +1,6 @@
 # mypy: allow-untyped-defs
 import functools
 from collections.abc import Callable
-from typing import Union
 from typing_extensions import TypeVarTuple
 
 import torch
@@ -47,7 +46,7 @@ def __call__(self, *args, **kwargs):
 
 def map(
     f: Callable[[pytree.PyTree, tuple[pytree.PyTree, ...]], pytree.PyTree],
-    xs: Union[pytree.PyTree, torch.Tensor],
+    xs: pytree.PyTree | torch.Tensor,
     *args: TypeVarTuple,
 ):
     r"""
diff --git a/torch/_higher_order_ops/partitioner.py b/torch/_higher_order_ops/partitioner.py
index 575142f16dda8..3517ca2b35127 100644
--- a/torch/_higher_order_ops/partitioner.py
+++ b/torch/_higher_order_ops/partitioner.py
@@ -1,6 +1,6 @@
 import logging
 from collections.abc import Callable
-from typing import Any, Union
+from typing import Any
 
 import torch
 from torch._higher_order_ops.utils import create_bw_fn, materialize_as_graph
@@ -83,8 +83,8 @@ def _check_partition_boundary(self) -> None:
         bw_grads = bw_phs[-self.n_fw_outputs :]
 
         def _match_size_or_expr(
-            val1: Union[torch.SymInt, torch.Tensor],
-            val2: Union[torch.SymInt, torch.Tensor],
+            val1: torch.SymInt | torch.Tensor,
+            val2: torch.SymInt | torch.Tensor,
         ) -> bool:
             if type(val1) is not type(val2):
                 return False
@@ -109,8 +109,7 @@ def _match_size_or_expr(
         if len(invalid_reasons) > 0:
             newline = "\n"
             raise RuntimeError(
-                "Invalid HopPartitionedGraph. Reasons:\n",
-                f"{newline.join(invalid_reasons)}",
+                f"Invalid HopPartitionedGraph. Reasons:\n{newline.join(invalid_reasons)}"
             )
 
     def _reorder_fw_output(self) -> None:
@@ -312,7 +311,7 @@ def partition(
 
 def create_hop_joint_graph(
     fw_fn: Callable,
-    fw_args: tuple[Union[torch.Tensor, torch.SymInt], ...],
+    fw_args: tuple[torch.Tensor | torch.SymInt, ...],
     functionalize: bool,
 ) -> HopJointGraph:
     fw_gm = materialize_as_graph(fw_fn, fw_args, force_enable_grad=True)
@@ -357,7 +356,7 @@ class HopGraphMinCutPartitioner:
     @staticmethod
     def create_partitioned_graph(
         fw_fn: Callable,
-        fw_args: tuple[Union[torch.Tensor, torch.SymInt], ...],
+        fw_args: tuple[torch.Tensor | torch.SymInt, ...],
         *,
         always_recompute_complex_exprs: bool = False,
     ) -> HopPartitionedGraph:
diff --git a/test/dynamo_expected_failures/CPython313-test_exception_variations-ExceptStarTestCases.test_try_except_else b/torch/_higher_order_ops/passes/__init__.py
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_exception_variations-ExceptStarTestCases.test_try_except_else
rename to torch/_higher_order_ops/passes/__init__.py
diff --git a/torch/_higher_order_ops/passes/inline_invoke_subgraph.py b/torch/_higher_order_ops/passes/inline_invoke_subgraph.py
new file mode 100644
index 0000000000000..54b2c4ab81b31
--- /dev/null
+++ b/torch/_higher_order_ops/passes/inline_invoke_subgraph.py
@@ -0,0 +1,69 @@
+import operator
+from typing import Any, TYPE_CHECKING
+
+import torch
+from torch.fx.graph_module import GraphModule
+
+
+if TYPE_CHECKING:
+    from torch.fx.node import Node
+
+
+def inline_invoke_subgraph(gm: GraphModule) -> GraphModule:
+    """Inline all invoke_subgraph HOPs, producing a flat FX graph.
+
+    This is useful when downstream compilers (like vllm-compile) don't support
+    HOPs or prefer a flat graph, but we still want the Dynamo tracing-time
+    benefits of auto-caching (trace once, stamp out cached calls).
+    """
+    invoke_nodes = list(
+        gm.graph.find_nodes(
+            op="call_function", target=torch.ops.higher_order.invoke_subgraph
+        )
+    )
+    if not invoke_nodes:
+        return gm
+
+    # Recursively inline any nested invoke_subgraph calls inside
+    # subgraph modules themselves.
+    for name, mod in gm.named_modules():
+        if name and isinstance(mod, GraphModule):
+            inline_invoke_subgraph(mod)
+
+    for node in invoke_nodes:
+        get_attr_node = node.args[0]
+        # args[1] is the identifier string, args[2:] are operands
+        operands = node.args[2:]
+
+        subgraph: GraphModule = getattr(gm, get_attr_node.target)
+
+        # Build mapping from subgraph placeholder nodes -> parent operands
+        env: dict[Node, Any] = dict(
+            zip(subgraph.graph.find_nodes(op="placeholder"), operands)
+        )
+
+        # Copy subgraph nodes into parent graph, inserting before the
+        # invoke_subgraph node.
+        with gm.graph.inserting_before(node):
+            for sub_node in subgraph.graph.nodes:
+                if sub_node.op in ("placeholder", "output"):
+                    continue
+                env[sub_node] = gm.graph.node_copy(sub_node, lambda n: env[n])
+
+        output_values = subgraph.graph.output_node().args[0]
+
+        # Replace getitem users of the invoke_subgraph result with the
+        # corresponding inlined output.
+        for user in list(node.users):
+            if user.op == "call_function" and user.target is operator.getitem:
+                idx = user.args[1]
+                user.replace_all_uses_with(env[output_values[idx]])  # pyrefly: ignore
+                gm.graph.erase_node(user)
+
+        gm.graph.erase_node(node)
+
+        # Remove the get_attr node if it has no other users.
+        if not get_attr_node.users:
+            gm.graph.erase_node(get_attr_node)
+
+    return gm
diff --git a/torch/_higher_order_ops/print.py b/torch/_higher_order_ops/print.py
index 86c692b9037b0..0399400acb4bb 100644
--- a/torch/_higher_order_ops/print.py
+++ b/torch/_higher_order_ops/print.py
@@ -26,6 +26,16 @@ class Print(HigherOrderOperator):
        torch._higher_order_ops.print("moo {} {y}", 1, y=2)
        Output: "moo 1 2"
 
+    4. DTensor support:
+       DTensor args are unwrapped to local tensors via to_local() (no collective).
+       Each rank prints its own local view, prefixed with [rank N].
+       For the global view of a sharded tensor, call full_tensor() before
+       passing to print.
+
+       dt = DTensor.from_local(local_shard, mesh, [Shard(0)])
+       torch._higher_order_ops.print("activations: {}", dt)
+       # Output: [rank 0] activations: tensor([0., 1.])
+
     This HOP enables printing without causing graph break.
     """
 
@@ -103,6 +113,33 @@ def print_impl(format_str: str, *args: object, **kwargs: object) -> None:
 print.fallthrough(torch._C.DispatchKey.AutogradCUDA)
 
 
+def _register_dtensor_impl() -> None:
+    from torch.distributed.tensor import DTensor
+
+    @print.py_impl(DTensor)  # pyrefly: ignore [missing-attribute]
+    # pyre-ignore
+    def print_dtensor(format_str: str, *args: object, **kwargs: object) -> None:
+        # Unwrap DTensors to local tensors via to_local() — no collective is
+        # introduced so there is no OOM or performance risk.  Every rank prints
+        # its own local view (including Replicate, where to_local() already
+        # holds the full tensor).
+        #
+        # The output is prefixed with [rank N] so users can identify which rank
+        # produced each line.
+        #
+        # If the user needs the global view of a sharded tensor, they can call
+        # full_tensor() explicitly before passing it to print.
+        import torch.distributed as dist
+
+        local_args = pytree.tree_map_only(DTensor, DTensor.to_local, args)
+        local_kwargs = pytree.tree_map_only(DTensor, DTensor.to_local, kwargs)
+        if dist.is_initialized() and dist.get_world_size() > 1:
+            format_str = f"[rank {dist.get_rank()}] {format_str}"
+        print(  # pyrefly: ignore [no-matching-overload]
+            format_str, *local_args, **local_kwargs
+        )
+
+
 @print.py_functionalize_impl
 def print_func(ctx, format_str: str, *args: object, **kwargs: object):
     from torch._higher_order_ops.effects import handle_effects
diff --git a/torch/_higher_order_ops/scan.py b/torch/_higher_order_ops/scan.py
index 40fbf14d64b33..0d7015f266316 100644
--- a/torch/_higher_order_ops/scan.py
+++ b/torch/_higher_order_ops/scan.py
@@ -298,22 +298,26 @@ def _scan(init, xs):
             return carry, []
 
         num_elems = xs[0].shape[dim]
-        ind = 0
-
-        # Compute dummy shapes for the pre-allocation
         num_init_leaves = len(init)
-        dummy_carry, dummy_out = _extract_carry_and_out(
+
+        # Process element 0 to infer output shapes for pre-allocation
+        # AND produce the first real result in a single call.  The previous
+        # approach used first_slice_copy() for shape inference and then
+        # re-processed element 0 in the main loop, calling the operator
+        # num_elems+1 times.  That extra invocation is incorrect for
+        # operators with side effects.
+        carry, out_0 = _extract_carry_and_out(
             call_operator(
                 operator,
                 *carry,
-                *[first_slice_copy(elem, dim) for elem in xs],
+                *[elem.select(dim, 0) for elem in xs],
                 *additional_inputs,
             ),
             num_init_leaves,
         )
 
-        out_tensor_mask = get_tensor_mask(dummy_out)
-        dummy_out_masked = mask_list(out_tensor_mask, dummy_out)
+        out_tensor_mask = get_tensor_mask(out_0)
+        out_0_masked = mask_list(out_tensor_mask, out_0)
 
         # Pre-allocate
         # outs -> Output matrix
@@ -321,16 +325,15 @@ def _scan(init, xs):
         # out: (num_elems, M, N, ...)
         # idx: (1, M, N)
         outs = [
-            torch.zeros(
+            torch.empty(
                 [num_elems] + list(e.size()),
                 dtype=e.dtype,
                 device=e.device,
             )
-            for i, e in enumerate(dummy_out_masked)
+            for e in out_0_masked
         ]
         idxs = [
-            torch.ones_like(e, dtype=torch.int64).unsqueeze(0)
-            for i, e in enumerate(dummy_out_masked)
+            torch.ones_like(e, dtype=torch.int64).unsqueeze(0) for e in out_0_masked
         ]
 
         def store_out_in_outs(out, ind):
@@ -342,20 +345,21 @@ def store_out_in_outs(out, ind):
                 # essentially: o[ind][n][k] = x[0][n][k]
                 o.scatter_(0, ind * idx, x.unsqueeze(0))
 
-        for i in range(num_elems):
-            ind = i
+        # Store element 0's result, then continue from element 1.
+        store_out_in_outs(out_0_masked, 0)
+
+        for i in range(1, num_elems):
             carry, out = _extract_carry_and_out(
                 call_operator(
                     operator,
                     *carry,
-                    *[elem.select(dim, ind) for elem in xs],
+                    *[elem.select(dim, i) for elem in xs],
                     *additional_inputs,
                 ),
                 num_init_leaves,
             )
 
-            # Store the inits in the outs matrix.
-            store_out_in_outs(mask_list(out_tensor_mask, out), ind)
+            store_out_in_outs(mask_list(out_tensor_mask, out), i)
 
         # Expand outs with None depending on the tensor mask of the output
         outs_expanded = [outs.pop(0) if out_m else None for out_m in out_tensor_mask]
diff --git a/torch/_higher_order_ops/schema.py b/torch/_higher_order_ops/schema.py
index d49b400635388..550a64b9e8249 100644
--- a/torch/_higher_order_ops/schema.py
+++ b/torch/_higher_order_ops/schema.py
@@ -1,6 +1,6 @@
 import copy
 from dataclasses import dataclass
-from typing import Any, Optional
+from typing import Any
 
 import torch
 import torch.utils._pytree as pytree
@@ -32,7 +32,7 @@ def from_example(
         example_value: Any,
         *,
         name: str = "",
-        default_value: Optional[Any] = None,
+        default_value: Any | None = None,
         is_mutated: bool = False,
         kw_only: bool = False,
     ) -> HopArgumentInfo:
@@ -104,18 +104,18 @@ class HopSchemaGenerator:
     def __init__(self, hop: torch._ops.HigherOrderOperator):
         self.arg_infos: list[HopArgumentInfo] = []
         self.example_outputs: list[Any] = []
-        self.schema_tree_spec: Optional[pytree.TreeSpec] = None
+        self.schema_tree_spec: pytree.TreeSpec | None = None
         self.hop = hop
 
     def add_arg(
         self,
         name: str,
         example_value: Any,
-        default_value: Optional[Any] = None,
+        default_value: Any | None = None,
         is_mutated: bool = False,
         kw_only: bool = False,
     ) -> None:
-        if callable(example_value):
+        if callable(example_value) and not is_opaque_type(type(example_value)):
             if not isinstance(
                 example_value, (torch.fx.GraphModule, torch._ops.OperatorBase)
             ):
@@ -212,7 +212,7 @@ def from_hop_argument_info(
         op_name: str,
         inp_argument_info: list[HopArgumentInfo],
         out_argument_info: HopArgumentInfo,
-        schema_tree_spec: Optional[pytree.TreeSpec],
+        schema_tree_spec: pytree.TreeSpec | None,
     ) -> Any:
         args = []
         for i, arg_info in enumerate(inp_argument_info):
@@ -265,7 +265,7 @@ def __init__(
         returns: list[torch._C.Argument],
         is_vararg: bool,
         is_varret: bool,
-        schema_tree_spec: Optional[pytree.TreeSpec],
+        schema_tree_spec: pytree.TreeSpec | None,
     ):
         self.tree_spec = schema_tree_spec
         self.is_vararg = is_vararg
diff --git a/torch/_higher_order_ops/strict_mode.py b/torch/_higher_order_ops/strict_mode.py
index f9ae59c37bc17..ac0f2952f2c23 100644
--- a/torch/_higher_order_ops/strict_mode.py
+++ b/torch/_higher_order_ops/strict_mode.py
@@ -22,12 +22,9 @@ def strict_mode(callable, operands):
     if torch.compiler.is_dynamo_compiling():
         return strict_mode_op(callable, operands)
 
-    from torch._higher_order_ops.utils import setup_compilation_env
+    from torch._higher_order_ops.utils import _hop_compile_and_call
 
-    with setup_compilation_env() as backend:
-        return torch.compile(strict_mode_op, backend=backend, fullgraph=True)(
-            callable, operands
-        )
+    return _hop_compile_and_call(strict_mode_op, (callable, operands))
 
 
 class StrictMode(HigherOrderOperator):
diff --git a/torch/_higher_order_ops/triton_kernel_wrap.py b/torch/_higher_order_ops/triton_kernel_wrap.py
index ae1e13bd1faa2..bb884b3445b8d 100644
--- a/torch/_higher_order_ops/triton_kernel_wrap.py
+++ b/torch/_higher_order_ops/triton_kernel_wrap.py
@@ -1,3 +1,4 @@
+import ast
 import collections
 import copy
 import dataclasses
@@ -7,6 +8,7 @@
 import logging
 import operator
 import threading
+import typing
 from collections import defaultdict
 from collections.abc import Callable, Sequence
 from typing import Any, Optional, TYPE_CHECKING, Union
@@ -29,6 +31,7 @@
 )
 from torch.fx.experimental.symbolic_shapes import guard_scalar
 from torch.types import IntLikeType
+from torch.utils._ordered_set import OrderedSet
 from torch.utils.checkpoint import _CachedTorchDispatchMode, _CachingTorchDispatchMode
 
 
@@ -41,14 +44,16 @@
     from torch._dynamo.symbolic_convert import InstructionTranslator
     from torch._dynamo.variables.constant import ConstantVariable
     from torch._dynamo.variables.functions import TritonKernelVariable
+    from torch._guards import Source
+    from torch._inductor.dependencies import ReadWrites
     from torch._subclasses.functional_tensor import BaseFunctionalizeAPI
     from torch.fx.proxy import Proxy
     from torch.utils._triton import has_triton
 
     TritonMetaParamsType = dict[str, int]
-    TritonGridTupleType = tuple[Union[int, sympy.Expr, SymInt], ...]
+    TritonGridTupleType = tuple[int | sympy.Expr | SymInt, ...]
     TritonGridCallableType = Callable[[TritonMetaParamsType], tuple[int, ...]]
-    TritonGridType = Union[TritonGridTupleType, TritonGridCallableType]
+    TritonGridType = TritonGridTupleType | TritonGridCallableType
 
     if has_triton():
         from triton.runtime.autotuner import Autotuner, Config as TritonConfig
@@ -61,9 +66,9 @@ class Autotuner:  # type: ignore[no-redef]
         class JITFunction:  # type: ignore[no-redef]
             pass
 
-    TritonKernelType = Union[Autotuner, JITFunction]
+    TritonKernelType = Autotuner | JITFunction
     # mypy specifically complains that TritonAutotunerType is not a valid type if Autotuner is not inside of a Union.
-    TritonAutotunerType = Union[Autotuner]
+    TritonAutotunerType = Union[Autotuner]  # noqa: UP007
 
 log = logging.getLogger("torch._dynamo")
 
@@ -95,8 +100,8 @@ def create_tma_experimental_metadata(
 
 
 def maybe_unpack_tma_experimental_metadata(
-    tma_meta: Union[TMAExperimentalMetadata, TMAStableMetadata],
-) -> Optional[tuple[list[IntLikeType], list[IntLikeType], IntLikeType]]:
+    tma_meta: TMAExperimentalMetadata | TMAStableMetadata,
+) -> tuple[list[IntLikeType], list[IntLikeType], IntLikeType] | None:
     if not tma_meta or len(tma_meta) != 2:
         return None
     if tma_meta[0] == "experimental":
@@ -111,8 +116,8 @@ def create_tma_stable_metadata(
 
 
 def maybe_unpack_tma_stable_metadata(
-    tma_meta: Union[TMAExperimentalMetadata, TMAStableMetadata],
-) -> Optional[tuple[list[IntLikeType]]]:
+    tma_meta: TMAExperimentalMetadata | TMAStableMetadata,
+) -> tuple[list[IntLikeType]] | None:
     if not tma_meta or len(tma_meta) != 2:
         return None
     if tma_meta[0] == "stable":
@@ -132,7 +137,7 @@ def maybe_unpack_tma_stable_metadata(
 # These are stored as raw tuples (instead of classes) for ease of serialization.
 TMADescriptorMetadata = dict[
     str,  # kernel parameter name
-    Union[TMAExperimentalMetadata, TMAStableMetadata],
+    TMAExperimentalMetadata | TMAStableMetadata,
 ]
 
 
@@ -216,11 +221,11 @@ def fake(self) -> bool:
 @dataclasses.dataclass(frozen=True, slots=True)
 class Op:
     name: str
-    fn_call_name: Optional[str]
-    args: list[Union[Param, Intermediate]]
+    fn_call_name: str | None
+    args: list[Param | Intermediate]
     ret: Intermediate = dataclasses.field(repr=False)
     # used for scf.yield: see [Note: scf.yield fix-up]
-    sub_idx: Optional[int] = None
+    sub_idx: int | None = None
     # used for tt.elementwise_inline_asm
     # `is_pure = True` assumes the asm block has no side-effects
     is_pure: bool = False
@@ -330,10 +335,9 @@ def is_stable_tensor_descriptor_arg(arg: Any) -> bool:
                 return True
         return False
 
-    def is_tensor_like_arg(arg: Any) -> bool:
-        if isinstance(arg, Tensor) or is_stable_tensor_descriptor_arg(arg):
-            return True
-        return False
+    def _is_constexpr_or_none(name: str, arg: Any) -> bool:
+        param_idx = kernel.arg_names.index(name)
+        return kernel.params[param_idx].is_constexpr or arg is None
 
     # Note: one would expect that each input to the triton kernel maps to
     # one input parameter in the TTIR. This is _not_ true for TMA descriptors:
@@ -343,9 +347,15 @@ def is_tensor_like_arg(arg: Any) -> bool:
     #   * N sizes, for a rank-N tensor
     # To account for this, we inject some fake arg names as placeholders for
     # the stride and size parameters.
-    def get_tensor_names(name: str, arg: Any) -> list[str]:
-        if isinstance(arg, Tensor):
-            return [name]
+    #
+    # Additionally, tensors and scalars are both included as TTIR parameters,
+    # whereas `constexpr` are inlined, and None are excluded. We both preserve
+    # scalars and tensors as this matters for "odd" ordering,
+    # eg. [tensor, scalar, tensor].
+    def get_arg_names(name: str, arg: Any) -> list[str]:
+        if _is_constexpr_or_none(name, arg):
+            return []
+
         if is_stable_tensor_descriptor_arg(arg):
             stable_meta = maybe_unpack_tma_stable_metadata(
                 tma_descriptor_metadata[name]
@@ -358,11 +368,12 @@ def get_tensor_names(name: str, arg: Any) -> list[str]:
             names.extend(name + f" STRIDE PLACEHOLDER {i}" for i in range(tensor_rank))
             names.extend(name + f" SIZE PLACEHOLDER {i}" for i in range(tensor_rank))
             return names
-        return []
 
-    ordered_tensor_names = list(
+        return [name]
+
+    ordered_arg_names = list(
         itertools.chain.from_iterable(
-            get_tensor_names(name, arg) for name, arg in ordered_args.items()
+            get_arg_names(name, arg) for name, arg in ordered_args.items()
         )
     )
 
@@ -452,8 +463,12 @@ def _native_specialize_impl(
             return attrs
 
     specialization = _get_specialization(ordered_args.values())
+    # Triton explicitly interprets ASTSource.constants entries as constexpr
+    # Thus, only None and arguments marked `is_constexpr` should be treated as such.
     constants = {
-        name: arg for name, arg in ordered_args.items() if not is_tensor_like_arg(arg)
+        name: arg
+        for name, arg in ordered_args.items()
+        if _is_constexpr_or_none(name, arg)
     }
 
     if (mangle_type := getattr(triton.runtime.jit, "mangle_type", None)) is not None:
@@ -462,7 +477,14 @@ def get_signature_value(idx: int, arg: Any) -> str:
             if kernel.params[idx].is_constexpr:
                 return "constexpr"
             # pyrefly: ignore [not-callable]
-            return mangle_type(arg)
+            result = mangle_type(arg)
+            # Workaround for Triton i1/u1 AOTI bug: PyTorch stores bool
+            # tensors as uint8 (1 byte per element), but *i1/*u1 causes
+            # the compiled kernel to generate bit-packed loads. Use *u8
+            # so loads correctly read 1 byte per element.
+            if result in ("*i1", "*u1"):
+                result = "*u8"
+            return result
 
     else:
 
@@ -518,7 +540,7 @@ def get_signature_value(idx: int, arg: Any) -> str:
     if not ttir_module.verify():
         raise RuntimeError("Verification for TTIR module has failed")
 
-    return ttir_module, ordered_tensor_names
+    return ttir_module, ordered_arg_names
 
 
 def ttir_to_functions(
@@ -537,7 +559,7 @@ def ttir_to_functions(
     )
     region_id_to_block_ids: dict[int, list[int]] = defaultdict(list)
     block_id_to_block_arg_ids: dict[int, list[int]] = {}
-    replacements: dict[int, Union[Intermediate, Param]] = {}
+    replacements: dict[int, Intermediate | Param] = {}
     reindex_map: dict[int, int] = {}
     next_fake_intermediate = 0
 
@@ -757,7 +779,7 @@ def mlir_to_functions(op: "TritonIROperation") -> None:
             callee = None
             if name == "tt.call":
                 callee = op.get_flat_symbol_ref_attr("callee")
-            args: list[Union[Param, Intermediate]] = [
+            args: list[Param | Intermediate] = [
                 Intermediate(operand) for operand in operand_ids
             ]
             block_ops = op_stack[parent_block_id]
@@ -796,7 +818,7 @@ def __call__(
         functions: dict[str, dict[Intermediate, list[Op]]],
         fn_name: str,
         *args: Any,
-    ) -> list[bool]:
+    ) -> Any:
         key: tuple[Any, ...] = (fn_name, *args)
         if key not in self.cache:
             self.cache[key] = None
@@ -812,7 +834,7 @@ def reset(self) -> None:
 @MemoizeWithCycleCheck
 def get_tma_stores(
     functions: dict[str, dict[Intermediate, list[Op]]], fn_name: str
-) -> set[Union[Intermediate, Param]]:
+) -> set[Intermediate | Param]:
     """
     Identifies all intermediates and parameters that are written to by a
     `tt.experimental_descriptor_store`. It tracks only the specific values
@@ -837,7 +859,7 @@ def get_tma_stores(
     function will also be marked.
     """
 
-    result: set[Union[Intermediate, Param]] = set()
+    result: set[Intermediate | Param] = set()
 
     ops = functions[fn_name]
     for op_list in ops.values():
@@ -881,21 +903,38 @@ def get_tma_stores(
     return result
 
 
+@dataclasses.dataclass
+class TensorAccesses:
+    read_writes: "ReadWrites"
+    can_fuse_epilogue: bool
+
+
 @MemoizeWithCycleCheck
-def analyze_kernel_mutations(
-    functions: dict[str, dict[Intermediate, list[Op]]], fn_name: str, num_args: int
-) -> list[bool]:
+def analyze_kernel_access(
+    functions: dict[str, dict[Intermediate, list[Op]]],
+    fn_name: str,
+    num_args: int,
+    tensor_names: tuple[str, ...],
+    tensor_arg_indices: frozenset[int] | None,
+) -> TensorAccesses:
     """
-    Analyzes the graph to detect all sinks from a predefined list of sinks
-    by using triton's MemWrite trait list. NOTE: What if triton exposed this?
-    From each sink, it traverses the CFG backwards to identify all the input
-    pointers that are mutated.
+    Analyzes the graph to detect which arguments are written to and which are read.
+
+    For writes: traverses from write sinks (tt.store, tt.atomic_cas, etc.) backwards
+    to identify input pointers that are written to.
+
+    For reads: traverses from read operations (tt.load) backwards to identify
+    input pointers that are read from.
+
+    Returns ReadWrites with StarDep objects for each accessed tensor.
     """
+    from torch._inductor.dependencies import Dep, ReadWrites, StarDep
+
     # Name of mutation op to mutated parameter indices
     # List from Triton Github include/triton/Dialect/Triton/IR/TritonOps.td
     # All the OPs that have MemWrite trait.
     # What if Triton exposed this?
-    MUTATION_OPS = {
+    WRITE_OPS = {
         "tt.store": [0],
         "tt.atomic_cas": [0],
         "tt.atomic_rmw": [0],
@@ -903,11 +942,16 @@ def analyze_kernel_mutations(
         "tt.experimental_tensormap_create": [0],
         "tt.descriptor_store": [0],
     }
-    # Ops that we want to bail out on
+    READ_OPS = {
+        "tt.load": [0],
+        "tt.load_tensor_descriptor": [0],
+        "tt.descriptor_load": [0],
+    }
     UNKNOWN_OPS = {"tt.elementwise_inline_asm"}
 
-    stack: list[Union[Param, Intermediate]] = []
-    visited = set()
+    write_stack: list[Param | Intermediate] = []
+    read_stack: list[Param | Intermediate] = []
+
     ops = functions[fn_name]
     tma_stores = get_tma_stores(functions, fn_name)
 
@@ -937,61 +981,135 @@ def analyze_kernel_mutations(
                         f"got {len(op.args)}"
                     )
                 if op.args[0] in tma_stores:
-                    stack.append(op.args[1])
+                    write_stack.append(op.args[1])
 
             if op.name == "tt.call":
                 if op.fn_call_name not in functions:
                     raise AssertionError(
                         f"Function {op.fn_call_name} not found in functions dict"
                     )
-                mutations = analyze_kernel_mutations(
+                # Create placeholder names for nested function arguments
+                nested_names = tuple(f"_arg{i}" for i in range(len(op.args)))
+
+                # Do not pass tensor_arg_indices, most outer call of
+                # analyze_kernel_access will filter Param nodes.
+                accesses = analyze_kernel_access(
                     functions,
                     # pyrefly: ignore [bad-argument-type]
                     op.fn_call_name,
                     len(op.args),
+                    nested_names,
+                    None,
                 )
-                stack.extend(arg for arg, mutated in zip(op.args, mutations) if mutated)
+                # Map back from StarDep names to args
+                written_set = {dep.name for dep in accesses.read_writes.writes}
+                read_set = {dep.name for dep in accesses.read_writes.reads}
+                for arg, name in zip(op.args, nested_names):
+                    if name in written_set:
+                        write_stack.append(arg)
+                    if name in read_set:
+                        read_stack.append(arg)
             else:
-                stack.extend(op.args[idx] for idx in MUTATION_OPS.get(op.name, []))
+                write_stack.extend(op.args[idx] for idx in WRITE_OPS.get(op.name, []))
+                read_stack.extend(op.args[idx] for idx in READ_OPS.get(op.name, []))
+
+    # For these ops, only the first argument (base pointer) refers to actual
+    # memory. The remaining arguments are shape/stride/offset metadata and
+    # should not be traced during mutation analysis.
+    POINTER_ONLY_OPS = {
+        "tt.make_tensor_ptr",
+        "tt.advance",
+        "tt.make_tensor_descriptor",
+    }
 
-    # The following is an iterative DFS algorithm
-    mutated = [False] * num_args
-    while stack:
-        arg = stack.pop()
-        if arg in visited:
-            continue
+    def _find_arg_access_count(
+        initial_stack: list[Param | Intermediate],
+        skip_loads: bool,
+    ) -> dict[int, int]:
+        """DFS traversal to find argument indices that are accessed (and how many times they are accessed)."""
+        access_count = dict()
+        stack = initial_stack[:]
 
-        visited.add(arg)
+        while stack:
+            arg = stack.pop()
 
-        if isinstance(arg, Param):
-            if arg.idx >= num_args:
-                # This is an argument defined in the kernel, not passed in
-                continue
-            mutated[arg.idx] = True
-        elif isinstance(arg, Intermediate) and not arg.fake():
-            for op in ops[arg]:
-                # Skip arguments to load
-                if op.name != "tt.load":
-                    stack.extend(op.args)
-    return mutated
+            if isinstance(arg, Param):
+                if arg.idx >= num_args:
+                    continue
+                if tensor_arg_indices is not None and arg.idx not in tensor_arg_indices:
+                    continue
+                if arg.idx not in access_count:
+                    access_count[arg.idx] = 1
+                else:
+                    access_count[arg.idx] += 1
+            elif isinstance(arg, Intermediate) and not arg.fake():
+                for op in ops[arg]:
+                    if skip_loads and op.name == "tt.load":
+                        continue
+                    if op.name in POINTER_ONLY_OPS:
+                        stack.append(op.args[0])
+                    else:
+                        stack.extend(op.args)
+
+        return access_count
 
+    write_count = _find_arg_access_count(write_stack, skip_loads=True)
+    read_count = _find_arg_access_count(read_stack, skip_loads=False)
 
-def identify_mutated_tensors(
+    writes: OrderedSet[Dep] = OrderedSet(
+        StarDep(tensor_names[i]) for i in sorted(write_count.keys())
+    )
+    reads: OrderedSet[Dep] = OrderedSet(
+        StarDep(tensor_names[i]) for i in sorted(read_count.keys())
+    )
+
+    read_writes = ReadWrites(
+        reads=reads,
+        writes=writes,
+        index_exprs=OrderedSet(),
+    )
+
+    def _decide_can_fuse_epilogue():
+        # only do epilogue fusion if the kernel has a single output tensor
+        if len(write_count) != 1:
+            return False
+
+        written_arg_index = next(iter(write_count))
+        # only do epilogue fusion if the written tensor is written exactly once
+        if write_count[written_arg_index] != 1:
+            return False
+
+        written_arg_name = next(iter(writes)).name
+        #  cannot fuse if the kernel also reads from the output buffer
+        if any(read_dep.name == written_arg_name for read_dep in reads):
+            return False
+
+        return True
+
+    can_fuse_epilogue = _decide_can_fuse_epilogue()
+
+    return TensorAccesses(read_writes=read_writes, can_fuse_epilogue=can_fuse_epilogue)
+
+
+def identify_accessed_tensors(
     kernel: "TritonKernelType",
     kwargs: dict[str, Any],
     tma_descriptor_metadata: TMADescriptorMetadata,
-) -> list[str]:
+) -> TensorAccesses:
     """
     Given a triton kernel and the arguments for this kernel, this function
     1) Retrieves the TTIR converted version of the kernel from Triton's API.
     2) Parses the TTIR and creates a control flow graph
-    3) Analyzes the graph to detect all input tensor mutations
+    3) Analyzes the graph to detect which input tensors are read and/or written
     """
 
+    from torch._inductor.dependencies import Dep, ReadWrites, StarDep
+    from torch._inductor.ir import TensorBox
+
     ttir_module = None
     functions = None
     try:
-        ttir_module, ordered_tensor_names = generate_ttir(
+        ttir_module, ordered_arg_names = generate_ttir(
             kernel, kwargs, tma_descriptor_metadata
         )
 
@@ -1009,37 +1127,119 @@ def identify_mutated_tensors(
                 f"Kernel name {kernel_fn_name} not found in TTIR kernel name {kernel_name}"
             )
         # Reset the cache between top level invocations
-        # The cache for analyze kernel mutations is mainly used for cycle
+        # The cache for analyze kernel access is mainly used for cycle
         # detection, so each top level invocation needs a clean cache
-        analyze_kernel_mutations.reset()
+        analyze_kernel_access.reset()
         get_tma_stores.reset()
-        mutations = analyze_kernel_mutations(
-            functions, kernel_name, len(ordered_tensor_names)
+
+        # Build frozenset of indices corresponding to tensor args only.
+        # Used to filter out scalars which are transitively captured as mutated
+        # during traversal.
+        tensor_arg_indices = frozenset(
+            i
+            for i, name in enumerate(ordered_arg_names)
+            if isinstance(kwargs.get(name), (Tensor, TensorBox))
         )
 
-        return [
-            ordered_tensor_names[i] for i, mutated in enumerate(mutations) if mutated
-        ]
+        return analyze_kernel_access(
+            functions,
+            kernel_name,
+            len(ordered_arg_names),
+            tuple(ordered_arg_names),
+            tensor_arg_indices,
+        )
     except Exception:
-        import torch._inductor.ir
-
         log.warning(
-            "Encountered an exception in identify_mutated_tensors, assuming every input is mutated",
+            "Encountered an exception in identify_accessed_tensors, assuming every input is mutated",
             exc_info=True,
         )
         if ttir_module is not None:
-            log.debug("TTIR:\n%s", str(ttir_module))
+            log.debug("TTIR:\n%s", ttir_module)
         if functions is not None:
             log.debug("functions:")
             for name, fn in functions.items():
                 log.debug("===\t%s\t===", name)
                 for ret, ops in fn.items():
                     log.debug("%s\t=>\t%s", ret, ops)
-        return [
+
+        all_tensor_names = [
             key
             for key, value in kwargs.items()
-            if isinstance(value, (Tensor, torch._inductor.ir.TensorBox))
+            if isinstance(value, (Tensor, TensorBox))
         ]
+        all_deps = OrderedSet(StarDep(name) for name in all_tensor_names)
+        all_deps = typing.cast(OrderedSet[Dep], all_deps)
+        return TensorAccesses(
+            ReadWrites(
+                reads=all_deps,
+                writes=all_deps,
+                index_exprs=OrderedSet(),
+            ),
+            can_fuse_epilogue=False,
+        )
+
+
+@dataclasses.dataclass
+class TritonStore:
+    store_node: ast.Call
+    store_pointer_node: ast.Expr
+    store_value_node: ast.Expr
+
+
+@dataclasses.dataclass
+class TritonStores:
+    stores: list[TritonStore]
+
+
+@functools.cache
+def identify_triton_stores(source_code: str) -> TritonStores:
+    """
+    Parse Python source code of triton kernel and find all tl.store calls.
+    Returns a TritonStores object containing information about pointer, value, and mask.
+
+    tl.store signature: store(pointer, value, mask=None, boundary_check=(), ...)
+    """
+    return identify_triton_stores_from_ast(ast.parse(source_code))
+
+
+def identify_triton_stores_from_ast(tree: ast.Module) -> TritonStores:
+    stores = []
+
+    def _extract_arg(node, arg_name, positional_index):
+        """
+        Extract an argument from a Call node, checking both positional and keyword args.
+        Returns the AST node for the argument, or None if not found.
+        """
+        # Check positional args first
+        if len(node.args) > positional_index:
+            return node.args[positional_index]
+
+        # Check keyword args
+        for keyword in node.keywords:
+            if keyword.arg == arg_name:
+                return keyword.value
+
+        return None
+
+    for node in ast.walk(tree):
+        if isinstance(node, ast.Call):
+            # Check if this is a tl.store call
+            if (
+                isinstance(node.func, ast.Attribute)
+                and isinstance(node.func.value, ast.Name)
+                and node.func.value.id == "tl"
+                and node.func.attr == "store"
+            ):
+                # Extract required arguments
+                pointer_node = _extract_arg(node, "pointer", 0)
+                value_node = _extract_arg(node, "value", 1)
+
+                if pointer_node is None or value_node is None:
+                    continue
+
+                stores.append(TritonStore(node, pointer_node, value_node))
+
+    return TritonStores(stores=stores)
 
 
 ###############################################################################
@@ -1222,7 +1422,7 @@ def trace_triton_kernel_wrapper(
     proxy_mode: ProxyTorchDispatchMode,
     func_overload: Callable[..., Any],
     node_args: dict[str, Any],
-) -> Optional[dict[str, Any]]:
+) -> dict[str, Any] | None:
     with disable_proxy_modes_tracing():
         out = func_overload(**node_args)
 
@@ -1275,9 +1475,16 @@ def get_mutated_tensors(
 ) -> list[str]:
     kernel = kernel_side_table.get_kernel(kernel_idx)
     constant_args = kernel_side_table.get_constant_args(constant_args_idx)
-    return identify_mutated_tensors(
+    tensor_accesses = identify_accessed_tensors(
         kernel, {**kwargs, **constant_args}, tma_descriptor_metadata
     )
+    # Filter to only tensor kwargs: with Triton 3.7+, ordered_arg_names
+    # includes scalars, so writes may reference non-tensor args like SymInts.
+    return [
+        dep.name
+        for dep in tensor_accesses.read_writes.writes
+        if isinstance(kwargs.get(dep.name), Tensor)
+    ]
 
 
 @triton_kernel_wrapper_mutation.py_functionalize_impl
@@ -1492,16 +1699,14 @@ def call_grid(  # type: ignore[no-untyped-def]
         grid,
         meta,
         tx,
-    ) -> Union[tuple[Union[int, sympy.Expr, SymInt], ...], tuple["Proxy", ...]]:
+    ) -> tuple[int | sympy.Expr | SymInt, ...] | tuple["Proxy", ...]:
         raise NotImplementedError("abstract method")
 
     def wrap_user_defined_obj(
         self,
         user_obj: Any,
         tx: Optional["InstructionTranslator"],
-        variable: Optional[
-            Union["TritonKernelVariable", "TraceableTritonKernelWrapper"]
-        ],
+        variable: Union["TritonKernelVariable", "TraceableTritonKernelWrapper"] | None,
         name: str,
     ) -> Any:
         raise NotImplementedError("abstract method")
@@ -1512,9 +1717,7 @@ def call_user_defined_fn(
         args: list,
         kwargs: dict,
         tx: Optional["InstructionTranslator"],
-        variable: Optional[
-            Union["TritonKernelVariable", "TraceableTritonKernelWrapper"]
-        ],
+        variable: Union["TritonKernelVariable", "TraceableTritonKernelWrapper"] | None,
     ) -> Any:
         raise NotImplementedError("abstract method")
 
@@ -1529,8 +1732,8 @@ def maybe_unpack_heuristic_result(self, result: Any) -> Any:
     @staticmethod
     def do_prune_configs(  # type: ignore[no-untyped-def]
         autotuner: "TritonAutotunerType",
-        early_config_prune: Optional[Callable],
-        perf_model: Optional[Callable],
+        early_config_prune: Callable | None,
+        perf_model: Callable | None,
         top_k: float,
         configs: list,
         named_args: dict,
@@ -1582,14 +1785,14 @@ def call_HOP(  # type: ignore[no-untyped-def]
 
     def check_grid(  # type: ignore[no-untyped-def]
         self, grid
-    ) -> Union[tuple[Union[int, sympy.Expr, SymInt], ...], tuple["Proxy", ...]]:
+    ) -> tuple[int | sympy.Expr | SymInt, ...] | tuple["Proxy", ...]:
         raise NotImplementedError("abstract method")
 
     def init_variable(
         self,
         variable: Union["TraceableTritonKernelWrapper", "TritonKernelVariable"],
         kernel: "TritonKernelType",
-        kernel_idx: Optional[int],
+        kernel_idx: int | None,
         grid: Optional["TritonGridType"],
     ) -> None:
         from triton.runtime.autotuner import Autotuner
@@ -1664,6 +1867,30 @@ def defaults_ok(
                     "pre_hook and post_hook are not supported in triton.Autotune or triton.Config"
                 )
 
+    @staticmethod
+    def get_kernel_source(
+        variable: "TritonKernelVariable | TraceableTritonKernelWrapper",
+    ) -> "Source | None":
+        kernel_source = getattr(variable, "kernel_source", None)
+        if kernel_source is None:
+            kernel_source = getattr(variable, "source", None)
+        return kernel_source
+
+    def recreate_variable(
+        self,
+        variable: "TritonKernelVariable | TraceableTritonKernelWrapper",
+        *,
+        kernel: "TritonKernelType",
+        kernel_idx: int | None,
+        grid: "TritonGridType | None",
+    ) -> "TritonKernelVariable | TraceableTritonKernelWrapper":
+        return type(variable)(
+            kernel=kernel,
+            kernel_idx=kernel_idx,
+            grid=grid,
+            kernel_source=self.get_kernel_source(variable),
+        )
+
     def call_getitem(
         self,
         variable: Union["TritonKernelVariable", "TraceableTritonKernelWrapper"],
@@ -1676,7 +1903,8 @@ def call_getitem(
                 "Triton kernels should be called with only a single grid"
             )
 
-        return type(variable)(
+        return self.recreate_variable(
+            variable,
             kernel=variable.kernel,
             kernel_idx=variable.kernel_idx,
             grid=args[0],
@@ -1695,8 +1923,11 @@ def call_run(
         kwargs.pop("warmup", None)
         # rewrite kernel.run(*args, grid=grid) to kernel[grid](*args)
         return self.call_triton_kernel(
-            type(variable)(
-                kernel=variable.kernel, kernel_idx=variable.kernel_idx, grid=grid
+            self.recreate_variable(
+                variable,
+                kernel=variable.kernel,
+                kernel_idx=variable.kernel_idx,
+                grid=grid,
             ),
             args,
             kwargs,
@@ -1801,7 +2032,12 @@ def call_triton_kernel(
             )(iter_kernel)
             # create a new variable to contain the new (wrapped) kernel;
             # skip kernel_idx to get a new record in the kernel side table
-            new_var = type(variable)(new_kernel, None, variable.grid)
+            new_var = self.recreate_variable(
+                variable,
+                kernel=new_kernel,
+                kernel_idx=None,
+                grid=variable.grid,
+            )
             return self.call_triton_kernel(new_var, args, kwargs, tx)
 
         SPECIAL_CONFIG_NAMES = {
@@ -1846,7 +2082,12 @@ def call_triton_kernel(
 
             # create a new variable to contain the new (wrapped) kernel;
             # skip kernel_idx to get a new record in the kernel side table
-            new_var = type(variable)(new_kernel, None, variable.grid)
+            new_var = self.recreate_variable(
+                variable,
+                kernel=new_kernel,
+                kernel_idx=None,
+                grid=variable.grid,
+            )
             return self.call_triton_kernel(new_var, args, kwargs, tx)
 
         if isinstance(variable.kernel, Autotuner):
@@ -1886,7 +2127,12 @@ def call_triton_kernel(
                     new_kernel = autotune(
                         configs=new_configs, prune_configs_by=prune_configs_by, key=[]
                     )(variable.kernel.fn)
-                    new_var = type(variable)(new_kernel, None, variable.grid)
+                    new_var = self.recreate_variable(
+                        variable,
+                        kernel=new_kernel,
+                        kernel_idx=None,
+                        grid=variable.grid,
+                    )
                     return self.call_triton_kernel(new_var, args, kwargs, tx)
 
         # These are the default values in upstream Triton
@@ -1946,7 +2192,12 @@ def call_triton_kernel(
             new_kernel = autotune(configs=pruned_configs, key=[])(variable.kernel.fn)
             # create a new variable to contain the new (wrapped) kernel;
             # skip kernel_idx to get a new record in the kernel side table
-            new_var = type(variable)(new_kernel, None, variable.grid)
+            new_var = self.recreate_variable(
+                variable,
+                kernel=new_kernel,
+                kernel_idx=None,
+                grid=variable.grid,
+            )
             return self.call_triton_kernel(new_var, args, kwargs, tx)
 
         # Both for grid's meta as well as for the kernel, we need combined
@@ -2038,7 +2289,7 @@ def call_grid(
         grid: "TritonGridCallableType",
         meta: "TritonMetaParamsType",
         tx: None,
-    ) -> tuple[Union[int, sympy.Expr, SymInt], ...]:
+    ) -> tuple[int | sympy.Expr | SymInt, ...]:
         if tx is not None:
             raise AssertionError("tx must be None for TracingTritonHOPifier")
         if not isinstance(meta, dict):
@@ -2051,9 +2302,7 @@ def wrap_user_defined_obj(
         self,
         user_obj: Any,
         tx: Optional["InstructionTranslator"],
-        variable: Optional[
-            Union["TritonKernelVariable", "TraceableTritonKernelWrapper"]
-        ],
+        variable: Union["TritonKernelVariable", "TraceableTritonKernelWrapper"] | None,
         name: str,
     ) -> Any:
         if tx is not None:
@@ -2066,9 +2315,7 @@ def call_user_defined_fn(
         args: list,
         kwargs: dict,
         tx: Optional["InstructionTranslator"],
-        variable: Optional[
-            Union["TritonKernelVariable", "TraceableTritonKernelWrapper"]
-        ],
+        variable: Union["TritonKernelVariable", "TraceableTritonKernelWrapper"] | None,
     ) -> Any:
         if not isinstance(args, list):
             raise AssertionError(f"args must be a list, got {type(args)}")
@@ -2091,7 +2338,7 @@ def maybe_unpack_heuristic_result(self, result: Any) -> Any:
     def check_grid(
         self,
         grid: "TritonGridType",
-    ) -> tuple[Union[int, sympy.Expr, SymInt], ...]:
+    ) -> tuple[int | sympy.Expr | SymInt, ...]:
         if not isinstance(grid, collections.abc.Sequence):
             raise RuntimeError(
                 "wrap_triton can only handle grids that resolve to Sequence[int]."
@@ -2156,17 +2403,20 @@ def call_HOP(
 
 class TraceableTritonKernelWrapper:
     kernel: "TritonKernelType"
-    kernel_idx: Optional[int]
+    kernel_idx: int | None
     grid: Optional["TritonGridType"]
+    kernel_source: "Source | None"
 
     def __init__(
         self,
         kernel: "TritonKernelType",
-        kernel_idx: Optional[int],
+        kernel_idx: int | None,
         grid: Optional["TritonGridType"],
+        kernel_source: "Source | None" = None,
     ) -> None:
         self.kernel = None
         self.grid = None
+        self.kernel_source = kernel_source
         tracing_triton_hopifier_singleton.init_variable(self, kernel, kernel_idx, grid)
         if self.kernel is None:
             raise AssertionError("kernel was not initialized properly")
diff --git a/torch/_higher_order_ops/utils.py b/torch/_higher_order_ops/utils.py
index 54b582f1e1afa..a4baf0be6f2bb 100644
--- a/torch/_higher_order_ops/utils.py
+++ b/torch/_higher_order_ops/utils.py
@@ -4,7 +4,7 @@
 from collections.abc import Callable, Iterable, Mapping, Sequence
 from contextlib import AbstractContextManager, contextmanager, ExitStack, nullcontext
 from dataclasses import dataclass
-from typing import Any, Optional, overload, TypeVar, Union
+from typing import Any, overload, TypeVar
 
 import torch
 import torch.fx.traceback as fx_traceback
@@ -96,10 +96,29 @@ def graph_with_interpreter(*args):
     return maybe_interpreted_fn
 
 
+def _hop_compile_and_call(fn, args, kwargs=None):
+    """Compile and call fn with fullgraph=True for HOP eager execution.
+
+    Pre-activates the fullgraph counter so that compile_wrapper treats this as
+    a nested compile.  This avoids erroring when a non-infra dispatch mode
+    causes the frame to be skipped — the function still executes eagerly within
+    compile_wrapper and returns normally.
+    """
+    from torch._dynamo.eval_frame import set_fullgraph_compiled_frame_count
+
+    with setup_compilation_env() as backend:
+        old_count = set_fullgraph_compiled_frame_count(0)
+        try:
+            return torch.compile(fn, backend=backend, fullgraph=True)(
+                *args, **(kwargs or {})
+            )
+        finally:
+            set_fullgraph_compiled_frame_count(old_count)
+
+
 def _maybe_compile_and_run_fn(fn, *args):
     if not torch.compiler.is_dynamo_compiling():
-        with setup_compilation_env() as backend:  # type: ignore[attr-defined]
-            return torch.compile(fn, backend=backend, fullgraph=True)(*args)
+        return _hop_compile_and_call(fn, args)
     else:
         return fn(*args)
 
@@ -157,19 +176,19 @@ def wrapped(*args):
 
 
 def check_meta_consistency(
-    lhs_list: list[Union[torch.Tensor, torch.SymInt, int]],
-    rhs_list: list[Union[torch.Tensor, torch.SymInt, int]],
+    lhs_list: list[torch.Tensor | torch.SymInt | int],
+    rhs_list: list[torch.Tensor | torch.SymInt | int],
     lhs_name: str,
     rhs_name: str,
     include_contiguity: bool = True,
 ) -> None:
     def diff_meta_pairs(
-        lhs_list: list[Union[torch.Tensor, torch.SymInt, int]],
-        rhs_list: list[Union[torch.Tensor, torch.SymInt, int]],
+        lhs_list: list[torch.Tensor | torch.SymInt | int],
+        rhs_list: list[torch.Tensor | torch.SymInt | int],
     ) -> list[str]:
         def diff_meta(
-            lhs: Union[torch.Tensor, torch.SymInt, int],
-            rhs: Union[torch.Tensor, torch.SymInt, int],
+            lhs: torch.Tensor | torch.SymInt | int,
+            rhs: torch.Tensor | torch.SymInt | int,
         ) -> str:
             if isinstance(lhs, torch.Tensor) and isinstance(rhs, torch.Tensor):
                 return ", ".join(
@@ -202,8 +221,8 @@ def _both_tensor(lhs, rhs):
 
         # Manually check the device of lhs and rhs as this field is currently not part of TensorMetadata
         def diff_device(
-            lhs: Union[torch.Tensor, torch.SymInt, int],
-            rhs: Union[torch.Tensor, torch.SymInt, int],
+            lhs: torch.Tensor | torch.SymInt | int,
+            rhs: torch.Tensor | torch.SymInt | int,
         ) -> str:
             if isinstance(lhs, torch.Tensor) and isinstance(rhs, torch.Tensor):
                 if (
@@ -312,7 +331,12 @@ def _set_compilation_env():
 
 # The invariant here is that we always trace the branch with fake tensor
 def _maybe_fake_tracing(fn, inputs: list[Any], pre_dispatch):
-    fake_mode_det = detect_fake_mode(inputs)
+    fake_mode_det = None
+    for inp in pytree.tree_leaves(inputs):
+        if isinstance(inp, FakeTensor):
+            fake_mode_det = inp.fake_mode
+            break
+
     fake_mode: AbstractContextManager = nullcontext()
     tracing_mode = "fake"
     if fake_mode_det is not None:
@@ -414,7 +438,7 @@ def _collect_fake_inputs(inputs):
     from torch._subclasses.fake_tensor import FakeTensor
 
     # Get the example values of the inputs.
-    inputs_fake: list[Union[FakeTensor, torch.Tensor, int]] = []
+    inputs_fake: list[FakeTensor | torch.Tensor | int] = []
     for inp in inputs:
         if isinstance(inp, (torch.fx.proxy.Proxy, torch.fx.node.Node)):
             inp = inp.node if isinstance(inp, torch.fx.proxy.Proxy) else inp
@@ -599,14 +623,20 @@ def unmask_none_gradients(grads, operands):
 
 
 def _maybe_fake_prop_ignore_unbacked(fn, args):
-    with ExitStack() as ctx_stack:
-        if (fake_mode := detect_fake_mode(args)) is not None:
-            ctx_stack.enter_context(fake_mode)
-            if fake_mode.shape_env is not None:
-                ctx_stack.enter_context(
-                    fake_mode.shape_env.ignore_fresh_unbacked_symbols()
-                )
-        return fn(*args)
+    with suspend_functionalization(), disable_functional_mode():
+        with disable_proxy_modes_tracing():
+            unfunc_args = [_from_fun(arg) for arg in args]
+        with ExitStack() as ctx_stack:
+            ctx_stack.enter_context(
+                torch.utils._python_dispatch._disable_current_modes()
+            )
+            if (fake_mode := detect_fake_mode(unfunc_args)) is not None:
+                ctx_stack.enter_context(fake_mode)
+                if fake_mode.shape_env is not None:
+                    ctx_stack.enter_context(
+                        fake_mode.shape_env.ignore_fresh_unbacked_symbols()
+                    )
+            return fn(*unfunc_args)
 
 
 def redirect_to_mode(hop: OperatorBase, mode):
@@ -917,7 +947,7 @@ def get_tensor_mask(tensor_list: Iterable[Any]) -> list[bool]:
 
 
 def mask_list(
-    mask: list[bool], inp: list[Any], other: Optional[list[Any]] = None
+    mask: list[bool], inp: list[Any], other: list[Any] | None = None
 ) -> list[Any]:
     # Masks elements on an `inp` list.
     # If other is None, then the elements of the `inp` list where the mask is False are removed
@@ -977,7 +1007,7 @@ def diff_tensor_meta(
 #      to support int arguments. In the eager run case, we re-trace the subgraph in AutogradKey, so inner
 #      hops may receive int inputs from the shape of outer tensor inputs.
 #      However, CompositeExplicitAutograd won't receive SymInt inputs because it only accepts real tensor inputs.
-def validate_subgraph_args_types(lifted_args: Union[tuple[Any, ...], list[Any]]):
+def validate_subgraph_args_types(lifted_args: tuple[Any, ...] | list[Any]):
     allowed_types = (torch.Tensor, int, torch.SymInt)
     if not all(
         isinstance(arg, (torch.Tensor, int, torch.SymInt)) for arg in lifted_args
@@ -1014,13 +1044,16 @@ def check_input_alias_and_mutation_return_outputs(
     dict[int, int],
     dict[int, int],
     list[int],
-    Union[tuple[Any, ...], list[Any]],
+    tuple[Any, ...] | list[Any],
 ]:
     def _get_example_value(n):
         if not isinstance(n, torch.fx.Node):
             return n
-        else:
-            return n.meta["val"] if "val" in n.meta else n.meta["example_value"]
+        if "val" in n.meta:
+            return n.meta["val"]
+        if "example_value" in n.meta:
+            return n.meta["example_value"]
+        return None
 
     fake_args = [
         _get_example_value(n)
@@ -1137,6 +1170,8 @@ class and then use the subgraph as the hash.
     def __init__(self, ctx, subgraph):
         self.ctx = ctx
         self.subgraph = subgraph
+        # Propagate so callers pass inputs as a list, enabling input deallocation.
+        self._boxed_call = getattr(subgraph, "_boxed_call", False)
 
     def __hash__(self):
         return id(self.subgraph)
@@ -1146,12 +1181,23 @@ def __repr__(self):
 
     def __call__(self, *args, **kwargs):
         if isinstance(self.subgraph, torch.fx.GraphModule):
-            # Running graph with interpreter is needed for propagating the stack_trace
-            with fx_traceback.preserve_node_meta():
-                return self.ctx.functionalize(torch.fx.Interpreter(self.subgraph).run)(
-                    *args, **kwargs
-                )
-        return self.ctx.functionalize(self.subgraph)(*args, **kwargs)
+            if self._boxed_call:
+                # Not all callers respect _boxed_call (e.g. reenter_make_fx).
+                if len(args) == 1 and isinstance(args[0], list):
+                    return self.ctx.functionalize(self.subgraph)(args[0])
+                return self.ctx.functionalize(self.subgraph)(list(args))
+            else:
+                # Running graph with interpreter is needed for propagating the stack_trace
+                with fx_traceback.preserve_node_meta():
+                    return self.ctx.functionalize(
+                        torch.fx.Interpreter(self.subgraph).run
+                    )(*args, **kwargs)
+        functionalized = self.ctx.functionalize(self.subgraph)
+        if self._boxed_call:
+            if len(args) == 1 and isinstance(args[0], list):
+                return functionalized(args[0])
+            return functionalized(list(args))
+        return functionalized(*args, **kwargs)
 
 
 # A wrapper over HigherOrderOperator that also carries its schema
@@ -1174,7 +1220,7 @@ def create(hop: HigherOrderOperator, *args, **kwargs):
 # This call_op can be used to call a HopInstance with
 # flat args and kwargs. We need to make use of the hop's schema's tree_spec
 # to unflatten the args and kwargs before calling the hop.
-def call_op(op: Union[OpOverload, HopInstance], args, kwargs):
+def call_op(op: OpOverload | HopInstance, args, kwargs):
     if isinstance(op, OpOverload):
         return op(*args, **kwargs)
 
@@ -1212,10 +1258,10 @@ def call_op(op: Union[OpOverload, HopInstance], args, kwargs):
 def materialize_as_graph(
     fn: Callable,
     args: tuple[Any, ...],
-    include_key_set: Optional[torch._C.DispatchKeySet] = None,
-    exclude_key_set: Optional[torch._C.DispatchKeySet] = None,
+    include_key_set: torch._C.DispatchKeySet | None = None,
+    exclude_key_set: torch._C.DispatchKeySet | None = None,
     force_enable_grad=False,
-    subgraph_decomp_table: Optional[Mapping[OpOverload, Callable]] = None,
+    subgraph_decomp_table: Mapping[OpOverload, Callable] | None = None,
 ) -> torch.fx.GraphModule:
     if include_key_set is None:
         include_key_set = torch._C._dispatch_tls_local_include_set()
@@ -1224,25 +1270,31 @@ def materialize_as_graph(
 
     @torch._dynamo.disable(recursive=True, reason=None)
     def _materialize_as_graph_inner():
+        from torch._guards import active_fake_mode
+        from torch.fx.experimental.proxy_tensor import _CURRENT_MAKE_FX_TRACER
+
         with suspend_functionalization(), disable_functional_mode():
-            with disable_proxy_modes_tracing():
-                unfunc_t = [_from_fun(arg) for arg in args]
+            fake_mode = None
+            if _CURRENT_MAKE_FX_TRACER is not None:
+                fake_mode = _CURRENT_MAKE_FX_TRACER.fake_tensor_mode
+            if fake_mode is None:
+                fake_mode = active_fake_mode()
 
             with contextlib.ExitStack() as stack:
                 stack.enter_context(
                     torch.utils._python_dispatch._disable_current_modes()
                 )
+                if fake_mode is not None:
+                    stack.enter_context(fake_mode)
+
+                with disable_proxy_modes_tracing():
+                    unfunc_t = [_from_fun(arg) for arg in args]
+
                 stack.enter_context(
                     torch._C._ForceDispatchKeyGuard(include_key_set, exclude_key_set),
                 )
                 if force_enable_grad:
                     stack.enter_context(torch.enable_grad())
-                # fake_mode is needed because parent tracer's fake_mode might
-                # be None but the associated inputs have fake mode or there
-                # is a global tracing context with fake mode. We nneed to
-                # make sure the fake mode when tracing subgraph is consistent.
-                if fake_mode := detect_fake_mode(unfunc_t):
-                    stack.enter_context(fake_mode)
                 return _maybe_reenter_make_fx(
                     fn, subgraph_decomp_table=subgraph_decomp_table
                 )(*unfunc_t)
@@ -1326,7 +1378,7 @@ def _has_gen_schema(op: HigherOrderOperator):
     )
 
 
-def filter_with_masks(data: list[Optional[torch.Tensor]], masks: list[bool]):
+def filter_with_masks(data: list[torch.Tensor | None], masks: list[bool]):
     if len(data) != len(masks):
         raise AssertionError(
             f"data length ({len(data)}) != masks length ({len(masks)})"
@@ -1334,6 +1386,6 @@ def filter_with_masks(data: list[Optional[torch.Tensor]], masks: list[bool]):
     return [item for item, keep in zip(data, masks) if keep]
 
 
-def fill_none_with_masks(data: list[Optional[torch.Tensor]], masks: list[bool]):
+def fill_none_with_masks(data: list[torch.Tensor | None], masks: list[bool]):
     data_iter = iter(data)
     return [next(data_iter) if kept else None for kept in masks]
diff --git a/torch/_higher_order_ops/while_loop.py b/torch/_higher_order_ops/while_loop.py
index c23d7ad7cc6b2..5484c7ab6e486 100644
--- a/torch/_higher_order_ops/while_loop.py
+++ b/torch/_higher_order_ops/while_loop.py
@@ -2,7 +2,6 @@
 import contextlib
 import functools
 from collections.abc import Callable
-from typing import Union
 
 import torch
 import torch.utils._pytree as pytree
@@ -35,8 +34,8 @@ def __call__(
         self,
         cond_fn: Callable,
         body_fn: Callable,
-        carried_inputs: tuple[Union[torch.Tensor, int, float, bool]],
-        additional_inputs: tuple[Union[torch.Tensor, torch.SymInt, int], ...],
+        carried_inputs: tuple[torch.Tensor | int | float | bool],
+        additional_inputs: tuple[torch.Tensor | torch.SymInt | int, ...],
         /,
     ):
         if not isinstance(carried_inputs, (tuple, list)):
@@ -249,12 +248,12 @@ def _validate_input(cond_fn, body_fn, carried_inputs):
     def _while_loop_op_wrapper(*args, **kwargs):
         return while_loop_op(*args, **kwargs)
 
-    from torch._higher_order_ops.utils import setup_compilation_env
+    from torch._higher_order_ops.utils import _hop_compile_and_call
 
-    with setup_compilation_env() as backend:
-        return torch.compile(_while_loop_op_wrapper, backend=backend, fullgraph=True)(
-            flat_cond_fn, flat_body_fn, tuple(flat_inputs), tuple()
-        )
+    return _hop_compile_and_call(
+        _while_loop_op_wrapper,
+        (flat_cond_fn, flat_body_fn, tuple(flat_inputs), tuple()),
+    )
 
 
 @while_loop_op.py_impl(DispatchKey.CompositeExplicitAutograd)
@@ -627,8 +626,8 @@ def __call__(
         self,
         cond_fn: Callable,
         body_fn: Callable,
-        carried_inputs: tuple[Union[torch.Tensor, int, float, bool]],
-        additional_inputs: tuple[Union[torch.Tensor, torch.SymInt, int], ...],
+        carried_inputs: tuple[torch.Tensor | int | float | bool],
+        additional_inputs: tuple[torch.Tensor | torch.SymInt | int, ...],
         /,
     ):
         if not isinstance(carried_inputs, (tuple, list)):
diff --git a/torch/_higher_order_ops/wrap.py b/torch/_higher_order_ops/wrap.py
index 09bee9eaea886..37a7cb6cb993f 100644
--- a/torch/_higher_order_ops/wrap.py
+++ b/torch/_higher_order_ops/wrap.py
@@ -2,7 +2,10 @@
 import inspect
 import itertools
 import logging
-from typing import Any, Optional
+import weakref
+from collections.abc import Callable
+from typing import Any
+from typing_extensions import ParamSpec, TypeVar
 
 import torch
 import torch.utils._pytree as pytree
@@ -21,8 +24,13 @@
 from torch.utils.checkpoint import _CachedTorchDispatchMode, _CachingTorchDispatchMode
 
 
+_P = ParamSpec("_P")
+_R = TypeVar("_R")
+
+
 log = logging.getLogger(__name__)
 
+
 uid = itertools.count(1)
 
 
@@ -31,7 +39,9 @@ class Wrap(HigherOrderOperator):
     def __init__(self) -> None:
         super().__init__("wrap")
 
-    def __call__(self, func, *args, **kwargs):
+    def __call__(
+        self, func: Callable[_P, _R], *args: _P.args, **kwargs: _P.kwargs
+    ) -> _R:
         # Dynamo already traces the body of HigherOrderOp beforehand when it
         # so no need to trace into it.
         import torch._dynamo  # noqa: F401
@@ -58,9 +68,9 @@ class InductorCompiledCode(HigherOrderOperator):
     def __init__(self) -> None:
         super().__init__("inductor_compiled_code")
 
-    def __call__(self, func, *args, **kwargs):
+    def __call__(self, func, inputs, *, name: str | None = None):
         # pyrefly: ignore [missing-attribute]
-        return super().__call__(func, *args, **kwargs)
+        return super().__call__(func, inputs, name=name)
 
 
 inductor_compiled_code = InductorCompiledCode()
@@ -68,9 +78,97 @@ def __call__(self, func, *args, **kwargs):
 inductor_compiled_code.fallthrough(DispatchKey.AutogradCUDA)
 
 
+_inductor_compiled_callable_id = itertools.count()
+
+
+class InductorCompiledCallable:
+    """
+    A wrapper class that holds both the Inductor-compiled callable and the
+    original FX graph for fake tensor propagation.
+    Each instance gets a globally unique idx at creation (via atomic itertools.count).
+    """
+
+    def __init__(
+        self,
+        compiled_callable,
+        original_gm=None,
+        compile_region_name: str | None = None,
+    ):
+        self.idx = next(_inductor_compiled_callable_id)
+        self.compiled_callable = compiled_callable
+        self.original_gm = original_gm
+        self.compile_region_name = compile_region_name
+        # AOT autograd needs this to know inputs are passed as a list
+        self._boxed_call = True
+
+    def __call__(self, inputs):
+        return self.compiled_callable(inputs)
+
+
+class InductorCodeSideTable:
+    """
+    Side table for storing InductorCompiledCallable objects.
+
+    We cannot put InductorCompiledCallable objects directly into the FX graph
+    as graph nodes do not support arbitrary objects. We use this side table
+    and pass callable.idx instead.
+
+    Uses WeakValueDictionary so entries are automatically removed when
+    the InductorCompiledCallable is no longer referenced elsewhere
+    (e.g. after dynamo.reset() drops the compiled code cache).
+    """
+
+    def __init__(self):
+        self.id_to_callable: weakref.WeakValueDictionary[
+            int, InductorCompiledCallable
+        ] = weakref.WeakValueDictionary()
+
+    def add_callable(self, callable_obj: InductorCompiledCallable) -> int:
+        """Register a callable and return its idx."""
+        self.id_to_callable[callable_obj.idx] = callable_obj
+        return callable_obj.idx
+
+    def get_callable(self, idx: int) -> InductorCompiledCallable:
+        """Get the callable at the given index."""
+        assert idx in self.id_to_callable, f"Invalid inductor code index: {idx}"  # noqa: S101
+        return self.id_to_callable[idx]
+
+    def __getstate__(self):
+        # Convert WeakValueDictionary to regular dict for pickling
+        return {"id_to_callable": dict(self.id_to_callable)}
+
+    def __setstate__(self, state):
+        self.id_to_callable = weakref.WeakValueDictionary(state["id_to_callable"])
+
+    def reset_table(self) -> None:
+        """Reset the table."""
+        self.id_to_callable = weakref.WeakValueDictionary()
+
+
+inductor_code_side_table = InductorCodeSideTable()
+
+
+def _resolve_inductor_callable(
+    func: int | InductorCompiledCallable,
+) -> InductorCompiledCallable:
+    """
+    Resolve func to an InductorCompiledCallable.
+
+    func is either an InductorCompiledCallable directly (from post_compile)
+    or an int index into the side table (from a traced FX graph node).
+    """
+    if isinstance(func, int):
+        return inductor_code_side_table.get_callable(func)
+    assert isinstance(func, InductorCompiledCallable), (  # noqa: S101
+        f"Unexpected func type: {type(func)}"
+    )
+    return func
+
+
 @inductor_compiled_code.py_impl(DispatchKey.CompositeExplicitAutograd)
-def inductor_compiled_code_impl(func, inputs):
-    return func(inputs)
+def inductor_compiled_code_impl(func, inputs, *, name=None):
+    resolved = _resolve_inductor_callable(func)
+    return resolved.compiled_callable(inputs)
 
 
 redirect_to_mode(inductor_compiled_code, DebugMode)
@@ -79,19 +177,64 @@ def inductor_compiled_code_impl(func, inputs):
 
 
 @register_fake(inductor_compiled_code)
-def inductor_compiled_code_fake(func, inputs):
-    raise RuntimeError(
-        "Inductor compiled code cannot be run with FakeTensor inputs. "
-        "This can happen when torch.compile is called inside a FakeTensorMode. "
-        "Consider using backend='eager' or backend='aot_eager' instead."
+def inductor_compiled_code_fake(func, inputs, *, name=None):
+    resolved = _resolve_inductor_callable(func)
+    if resolved.original_gm is None:
+        raise RuntimeError(
+            "inductor_compiled_code original_gm is None — the compiled graph may "
+            "have been serialized without it. Recompile to restore."
+        )
+    # Run the original FX graph under FakeTensorMode to re-derive output
+    # shapes, dtypes, and aliasing from the input fake tensors.
+    return tuple(resolved.original_gm(*inputs))
+
+
+@inductor_compiled_code.py_functionalize_impl
+def inductor_compiled_code_functionalize(ctx, func, inputs, *, name=None):
+    # Unwrap the functional tensors to get the underlying tensors
+    unwrapped_inputs = ctx.unwrap_tensors(inputs)
+
+    # Redispatch to the next handler in the dispatch chain
+    with ctx.redispatch_to_next():
+        kwargs = {"name": name} if name is not None else {}
+        result = inductor_compiled_code(func, unwrapped_inputs, **kwargs)
+        return ctx.wrap_tensors(result)
+
+
+@inductor_compiled_code.py_impl(ProxyTorchDispatchMode)
+def inductor_compiled_code_proxy(mode, func, inputs, *, name=None):
+    resolved = _resolve_inductor_callable(func)
+
+    # Run the fake impl to get example outputs for tracing
+    kwargs = {"name": name} if name is not None else {}
+    example_out = inductor_compiled_code(func, inputs, **kwargs)
+
+    # Register in side table so the FX node stores a serializable int
+    callable_idx = inductor_code_side_table.add_callable(resolved)
+
+    proxy_inputs = pytree.tree_map(mode.tracer.unwrap_proxy, inputs)
+
+    out_proxy = mode.tracer.create_proxy(
+        "call_function",
+        inductor_compiled_code,
+        (callable_idx, proxy_inputs),
+        kwargs,
     )
 
+    return track_tensor_tree(example_out, out_proxy, constant=None, tracer=mode.tracer)
+
 
 class WrapWithSetGradEnabled(HigherOrderOperator):
     def __init__(self) -> None:
         super().__init__("wrap_with_set_grad_enabled")
 
-    def __call__(self, enable_grad, wrapped_func, *args, **kwargs):
+    def __call__(
+        self,
+        enable_grad: bool,
+        wrapped_func: Callable[_P, _R],
+        *args: _P.args,
+        **kwargs: _P.kwargs,
+    ) -> _R:
         # Dynamo already traces the body of HigherOrderOp beforehand when it
         # so no need to trace into it.
         import torch._dynamo  # noqa: F401
@@ -118,13 +261,13 @@ def __init__(self):
     def __call__(
         self,
         device_type: str,
-        dtype: Optional[_dtype],
+        dtype: _dtype | None,
         enabled: bool,
-        cache_enabled: Optional[bool],
-        wrapped_func,
-        *args,
-        **kwargs,
-    ):
+        cache_enabled: bool | None,
+        wrapped_func: Callable[_P, _R],
+        *args: _P.args,
+        **kwargs: _P.kwargs,
+    ) -> _R:
         # Dynamo already traces the body of HigherOrderOp beforehand when it
         # so no need to trace into it.
         import torch._dynamo  # noqa: F401
@@ -204,7 +347,7 @@ class WrapActivationCheckpoint(HigherOrderOperator):
     def __init__(self) -> None:
         super().__init__("wrap_activation_checkpoint", cacheable=False)
 
-    def __call__(self, function, *args, **kwargs):
+    def __call__(self, function: GraphModule, *args: Any, **kwargs: Any) -> Any:
         # use_reentrant is set to False because this op is going to be traced.
         # And we ensure that AOT Autograd traces through the non reentrant
         # version of checkpointing.
@@ -284,22 +427,6 @@ def divide_kwargs(kwargs):
         }
         return checkpoint_kwargs, gmod_kwargs
 
-    @staticmethod
-    def tag_nodes(gmod, is_sac):
-        from torch.utils.checkpoint import CheckpointPolicy
-
-        unique_graph_id = next(uid)
-        for node in gmod.graph.nodes:
-            if node.op in ("call_function", "call_method", "call_module"):
-                node.meta["ac_graph_id"] = unique_graph_id
-                if is_sac:
-                    # For selective checkpointing, we will populate this tag later in _CachingTorchDispatchMode.
-                    node.meta["recompute"] = None
-                else:
-                    # Under vanilla activation checkpointing, all nodes should be recomputed.
-                    node.meta["recompute"] = CheckpointPolicy.PREFER_RECOMPUTE
-        return gmod
-
     def __call__(self, gmod, *args, **kwargs):
         dispatch_key_set = torch._ops._compute_keyset(
             args, kwargs, self.non_fallthrough_keys
@@ -315,11 +442,22 @@ def __call__(self, gmod, *args, **kwargs):
 tag_activation_checkpoint = TagActivationCheckpoint()
 
 
+def _always_prefer_recompute(ctx, op, *args, **kwargs):
+    from torch.utils.checkpoint import CheckpointPolicy
+
+    return CheckpointPolicy.PREFER_RECOMPUTE
+
+
 def tag_activation_checkpoint_impl(gmod, *args, **kwargs):
+    import functools
+
     import torch.fx.traceback as fx_traceback
     from torch.fx import Interpreter
+    from torch.utils.checkpoint import create_selective_checkpoint_contexts
 
+    unique_graph_id = next(uid)
     if "_checkpoint_context_fn" in gmod.meta:
+        context_fn = gmod.meta["_checkpoint_context_fn"]
         warning_once(
             log,
             """
@@ -327,33 +465,47 @@ def tag_activation_checkpoint_impl(gmod, *args, **kwargs):
 Please make sure the checkpointed region does not contain in-place ops (e.g. torch.relu_).
 """,
         )
-        # use_reentrant is set to False because this op is going to be traced.
-        # And we ensure that AOT Autograd traces through the non reentrant
-        # version of checkpointing.
-        kwargs["use_reentrant"] = False
-        # preserve_rng_state is set to False because we want to prevent AOTAutograd from tracing through
-        # `torch.random.fork_rng` op (which is not supported yet under CUDA).
-        # This doesn't mean that we don't preserve RNG state. Instead, we will always preserve RNG state
-        # regardless of this flag (by doing RNG functionalization via `replace_random_passes` in Inductor
-        # instead of in AOTAutograd).
-        kwargs["preserve_rng_state"] = False
-        kwargs["context_fn"] = gmod.meta["_checkpoint_context_fn"]
-        # We first tag all nodes as "recompute" in this graph, and then we undo the "recompute" tag
-        # for specific nodes in _CachingTorchDispatchMode in torch/utils/checkpoint.py.
-        gmod = TagActivationCheckpoint.tag_nodes(gmod, is_sac=True)
-        # Using interpreter allows preservation of metadata through torch.compile stack.
-        with fx_traceback.preserve_node_meta():
-            from torch.utils.checkpoint import checkpoint
-
-            return checkpoint(Interpreter(gmod).run, *args, **kwargs)
     else:
-        gmod = TagActivationCheckpoint.tag_nodes(gmod, is_sac=False)
-        # Using interpreter allows preservation of metadata through torch.compile stack.
-        # TODO: We want to use the same `checkpoint(Interpreter(gmod).run, *args, **kwargs)` here
-        # as the `context_fn != None` case, but that depends on in-place op support in TorchDispatchMode + torch.compile.
-        # (for details on in-place op issue, run `test_compile_selective_checkpoint_inplace_op` unit test)
-        with fx_traceback.preserve_node_meta():
-            return Interpreter(gmod).run(*args)
+        # Vanilla AC: use the SAC path with a policy that always recomputes.
+        context_fn = functools.partial(
+            create_selective_checkpoint_contexts, _always_prefer_recompute
+        )
+
+    def context_fn_with_graph_id():
+        fwd_ctx, recomp_ctx = context_fn()
+        # Plumb ac_graph_id so _CachingTorchDispatchMode tags all nodes
+        # (including ops from desugared HOPs like custom autograd.Function).
+        fwd_ctx.ac_graph_id = unique_graph_id
+        return fwd_ctx, recomp_ctx
+
+    # use_reentrant is set to False because this op is going to be traced.
+    # And we ensure that AOT Autograd traces through the non reentrant
+    # version of checkpointing.
+    kwargs["use_reentrant"] = False
+    # preserve_rng_state is set to False because we want to prevent AOTAutograd from tracing through
+    # `torch.random.fork_rng` op (which is not supported yet under CUDA).
+    # This doesn't mean that we don't preserve RNG state. Instead, we will always preserve RNG state
+    # regardless of this flag (by doing RNG functionalization via `replace_random_passes` in Inductor
+    # instead of in AOTAutograd).
+    kwargs["preserve_rng_state"] = False
+    kwargs["context_fn"] = context_fn_with_graph_id
+    # Disable early stop to prevent _StopRecomputationError from interrupting
+    # recomputation between _vmap_increment_nesting and _vmap_decrement_nesting,
+    # which would leak a functorch dynamic layer.
+    kwargs["early_stop"] = False
+    # Using interpreter allows preservation of metadata through torch.compile stack.
+    # We use a wrapper instead of passing Interpreter(gmod).run directly because
+    # checkpoint's recompute_fn captures the function in a closure. A bound method
+    # reference would keep the Interpreter alive, whose env dict retains the output
+    # tensors and prevents the autograd graph from being freed.
+
+    def run_with_interpreter(*args):
+        return Interpreter(gmod).run(*args)
+
+    with fx_traceback.preserve_node_meta():
+        from torch.utils.checkpoint import checkpoint
+
+        return checkpoint(run_with_interpreter, *args, **kwargs)
 
 
 @tag_activation_checkpoint.py_impl(ProxyTorchDispatchMode)
diff --git a/torch/_inductor/__init__.py b/torch/_inductor/__init__.py
index 165bf91caae7c..09a0de086adbc 100644
--- a/torch/_inductor/__init__.py
+++ b/torch/_inductor/__init__.py
@@ -1,6 +1,7 @@
 # mypy: allow-untyped-defs
 from __future__ import annotations
 
+import dataclasses
 import io
 import logging
 import os
@@ -33,7 +34,7 @@
 def compile(
     gm: torch.fx.GraphModule,
     example_inputs: list[InputType],
-    options: Optional[dict[str, Any]] = None,
+    options: dict[str, Any] | None = None,
 ):
     """
     Compile a given FX graph with TorchInductor.  This allows compiling
@@ -57,8 +58,8 @@ def aoti_compile_and_package(
     _deprecated_unused_args=None,
     _deprecated_unused_kwargs=None,
     *,
-    package_path: Optional[FileLike] = None,
-    inductor_configs: Optional[dict[str, Any]] = None,
+    package_path: FileLike | None = None,
+    inductor_configs: dict[str, Any] | None = None,
 ) -> str:
     """
     Compiles the exported program with AOTInductor, and packages it into a .pt2
@@ -160,12 +161,12 @@ def _aoti_compile_and_package_inner(
     gm: torch.nn.Module,
     # flat_example_inputs: List[Any],
     args: tuple[Any],
-    kwargs: Optional[dict[str, Any]] = None,
+    kwargs: dict[str, Any] | None = None,
     *,
     load_and_run: bool = False,
-    check_accuracy: Optional[str] = None,
-    package_path: Optional[Union[str, io.BytesIO]] = None,
-    inductor_configs: Optional[dict[str, Any]] = None,
+    check_accuracy: str | None = None,
+    package_path: str | io.BytesIO | None = None,
+    inductor_configs: dict[str, Any] | None = None,
 ):
     """
     See docstring for aoti_compile_and_package.
@@ -272,10 +273,10 @@ def aoti_load_package(
 def aot_compile(
     gm: torch.fx.GraphModule,
     args: tuple[Any, ...],
-    kwargs: Optional[dict[str, Any]] = None,
+    kwargs: dict[str, Any] | None = None,
     *,
-    options: Optional[dict[str, Any]] = None,
-) -> Union[str, list[Union[str, Weights]], torch.fx.GraphModule]:
+    options: dict[str, Any] | None = None,
+) -> str | list[str | Weights] | torch.fx.GraphModule:
     """
     Ahead-of-time compile a given FX graph with TorchInductor into a shared library.
 
@@ -334,7 +335,7 @@ def aot_compile(
 
 
 def list_mode_options(
-    mode: Optional[str] = None, dynamic: Optional[bool] = None
+    mode: str | None = None, dynamic: bool | None = None
 ) -> dict[str, Any]:
     r"""Returns a dictionary describing the optimizations that each of the available
     modes passed to `torch.compile()` performs.
@@ -409,8 +410,9 @@ def standalone_compile(
     dynamic_shapes: Literal[
         "from_example_inputs", "from_tracing_context", "from_graph"
     ] = "from_graph",
-    options: Optional[dict[str, Any]] = None,
+    options: dict[str, Any] | None = None,
     aot: bool = False,  # AOT mode, which uses BundledAOTAutogradCache
+    donate_graph_module: bool = False,
 ) -> CompiledArtifact:
     """
     Precompilation API for inductor.
@@ -434,6 +436,9 @@ def standalone_compile(
             If "from_example_inputs", we will specialize the graph on the
             example_inputs.
         options: Inductor compilation options
+        donate_graph_module: If True, standalone_compile takes ownership of
+            the graph module and may mutate it, avoiding an internal deepcopy.
+            Defaults to False for backwards compatibility.
 
     Returns:
         CompiledArtifact that can be saved to disk or invoked directly.
@@ -442,5 +447,16 @@ def standalone_compile(
 
     options = options if options else {}
     return standalone_compile(
-        gm, example_inputs, dynamic_shapes=dynamic_shapes, options=options, aot=aot
+        gm,
+        example_inputs,
+        dynamic_shapes=dynamic_shapes,
+        options=options,
+        aot=aot,
+        donate_graph_module=donate_graph_module,
     )
+
+
+@dataclasses.dataclass
+class _CudagraphAnnotation:
+    fwd: bool | None
+    bwd: bool | None
diff --git a/torch/_inductor/analysis/device_info.py b/torch/_inductor/analysis/device_info.py
index 8d5edf1e7fd26..6b8dcfc2e385c 100644
--- a/torch/_inductor/analysis/device_info.py
+++ b/torch/_inductor/analysis/device_info.py
@@ -1,6 +1,5 @@
 import logging
 from dataclasses import dataclass
-from typing import Optional, Union
 
 import torch
 
@@ -19,7 +18,7 @@ class DeviceInfo:
     For example,
     """
 
-    tops: dict[Union[torch.dtype, str], float]
+    tops: dict[torch.dtype | str, float]
     dram_bw_gbs: float
     dram_gb: float
 
@@ -178,7 +177,7 @@ class DeviceInfo:
 _device_mapping["AMD INSTINCT MI210X"] = _device_mapping["AMD MI210X"]
 
 
-def lookup_device_info(name: str) -> Optional[DeviceInfo]:
+def lookup_device_info(name: str) -> DeviceInfo | None:
     """
     Problem: when diffing profiles between amd and nvidia, we don't have access to the device information
     of the other one. Also, since the analysis is static, we should be able to do it on another device unrelated
@@ -189,12 +188,12 @@ def lookup_device_info(name: str) -> Optional[DeviceInfo]:
     return _device_mapping.get(name)
 
 
-def datasheet_tops(dtype: torch.dtype, is_tf32: bool = False) -> Optional[float]:
+def datasheet_tops(dtype: torch.dtype, is_tf32: bool = False) -> float | None:
     """
     Get the theoretical TFLOPS of the device for a given dtype. This can throw an exception if the device
     is not in the datasheet list above.
     """
-    name: Optional[str] = torch.cuda.get_device_name()
+    name: str | None = torch.cuda.get_device_name()
     if name is None:
         log.info("No device found, returning None")
         return None
diff --git a/torch/_inductor/analysis/profile_analysis.py b/torch/_inductor/analysis/profile_analysis.py
index 6a6ec39003bdb..3dd20fba88a92 100644
--- a/torch/_inductor/analysis/profile_analysis.py
+++ b/torch/_inductor/analysis/profile_analysis.py
@@ -4,7 +4,7 @@
 from collections import defaultdict
 from collections.abc import Callable
 from dataclasses import dataclass
-from typing import Any, Optional, Union
+from typing import Any, Optional
 
 import torch
 from torch._inductor.analysis.device_info import DeviceInfo, lookup_device_info
@@ -44,7 +44,7 @@ def parse_list(lst: str) -> list[int]:
 
 
 def register_adapter(
-    aten: Union[str, list[str]],
+    aten: str | list[str],
 ) -> Callable[
     [AdapterType],
     AdapterType,
@@ -144,7 +144,7 @@ def mm_adapter(
     return shapes, {}
 
 
-def _parse_kernel_name(name: str) -> Optional[str]:
+def _parse_kernel_name(name: str) -> str | None:
     """
     parse the name of the kernel from the event name.
     """
@@ -385,7 +385,7 @@ class KernelStats:
 class Device:
     name: str
     index: int
-    info: Optional[DeviceInfo]
+    info: DeviceInfo | None
     stats: KernelNameMap
 
     def __repr__(self) -> str:
@@ -402,8 +402,8 @@ class JsonProfile:
     def __init__(
         self,
         path: str,
-        benchmark_name: Optional[str] = None,
-        dtype: Optional[Union[torch.dtype, str]] = None,
+        benchmark_name: str | None = None,
+        dtype: torch.dtype | str | None = None,
     ):
         """
         Convenience class for running common operations on chrome/perfetto json traces.
@@ -423,7 +423,7 @@ def __init__(
             self.dtype = _dtype_map.get(dtype)
         self._create_devices()
 
-    def convert_dtype(self, event: dict[str, Any]) -> Optional[torch.dtype]:
+    def convert_dtype(self, event: dict[str, Any]) -> torch.dtype | None:
         """
         Each op has a list of dtypes for each input arg. We need to convert these into a single dtype for flop estimation.
         Issues:
diff --git a/torch/_inductor/analyze_preserves_zero_mask.py b/torch/_inductor/analyze_preserves_zero_mask.py
index 0674d1566c33b..b1721a4f8290f 100644
--- a/torch/_inductor/analyze_preserves_zero_mask.py
+++ b/torch/_inductor/analyze_preserves_zero_mask.py
@@ -1,6 +1,6 @@
 import dataclasses
 import itertools
-from typing import Any, Optional, TYPE_CHECKING
+from typing import Any, TYPE_CHECKING
 
 import sympy
 
@@ -29,7 +29,7 @@ class PreservesZeros(SymPyOps, DefaultHandler):
 
     def __init__(self) -> None:
         self.count = itertools.count(0)
-        self.store_preserves_zeros: Optional[bool] = None
+        self.store_preserves_zeros: bool | None = None
         self.dtype_prop = DtypePropagationOpsHandler()
 
     def load(self, name: str, index: sympy.Expr) -> TypedExpr:
diff --git a/torch/_inductor/aoti_eager.py b/torch/_inductor/aoti_eager.py
index 991f1caaecbb9..69f6a45afa576 100644
--- a/torch/_inductor/aoti_eager.py
+++ b/torch/_inductor/aoti_eager.py
@@ -2,8 +2,9 @@
 import logging
 import os
 from collections.abc import Callable
+from dataclasses import dataclass
 from pathlib import Path
-from typing import Any, Optional
+from typing import Any
 from unittest import mock
 
 import torch
@@ -16,6 +17,26 @@
 log = logging.getLogger(__name__)
 
 
+@dataclass
+class AOTICompileBackend:
+    compile_fn: Callable[..., str]
+    load_fn: Callable[[str, str, str], list[dict[str, Any] | None]]
+
+
+_aoti_compile_backends: dict[str, AOTICompileBackend] = {}
+
+
+def register_aoti_compile_backend(
+    device_type: str,
+    compile_fn: Callable[..., str],
+    load_fn: Callable[[str, str, str], list[dict[str, Any] | None]],
+) -> None:
+    _aoti_compile_backends[device_type] = AOTICompileBackend(
+        compile_fn=compile_fn,
+        load_fn=load_fn,
+    )
+
+
 def aoti_eager_cache_dir(namespace: str, device: str) -> Path:
     return Path(cache_dir()) / "aoti_eager" / namespace / device
 
@@ -32,7 +53,11 @@ def aoti_eager_op_conf_lock(op_func_name_with_overload: str) -> Any:
 
 def load_aoti_eager_cache(
     ns: str, op_func_name_with_overload: str, device_type: str
-) -> list[Optional[dict[str, Any]]]:
+) -> list[dict[str, Any] | None]:
+    backend = _aoti_compile_backends.get(device_type)
+    if backend:
+        return backend.load_fn(ns, op_func_name_with_overload, device_type)
+
     device_kernel_cache = aoti_eager_cache_dir(ns, device_type)
     op_conf = device_kernel_cache / f"{op_func_name_with_overload}.json"
     if not op_conf.exists():
@@ -174,14 +199,30 @@ def aoti_compile_with_persistent_cache(
     args: tuple[Any],
     kwargs: dict[str, Any],
     *,
-    dynamic_shapes: Optional[dict[str, Any]] = None,
-    options: Optional[dict[str, Any]] = None,
+    dynamic_shapes: dict[str, Any] | None = None,
+    options: dict[str, Any] | None = None,
     remove_runtime_assertions: bool = False,
     disable_constraint_solver: bool = False,
 ) -> str:
     """
     Compile the given function with persistent cache for AOTI eager mode.
     """
+    backend = _aoti_compile_backends.get(device_type)
+    if backend:
+        return backend.compile_fn(
+            ns,
+            op_func_name_with_overload,
+            device_type,
+            dynamic,
+            f,
+            args,
+            kwargs,
+            dynamic_shapes=dynamic_shapes,
+            options=options,
+            remove_runtime_assertions=remove_runtime_assertions,
+            disable_constraint_solver=disable_constraint_solver,
+        )
+
     assert not dynamic, "Only support static shape for now"
     flattened_inputs = list(args) + list(kwargs.values())
     if not all(
diff --git a/torch/_inductor/async_compile.py b/torch/_inductor/async_compile.py
index 7539a0c7b6c05..1e6dcbece6b63 100644
--- a/torch/_inductor/async_compile.py
+++ b/torch/_inductor/async_compile.py
@@ -13,7 +13,7 @@
 from concurrent.futures.process import BrokenProcessPool
 from functools import partial
 from time import time, time_ns
-from typing import Any, Optional, TYPE_CHECKING
+from typing import Any, TYPE_CHECKING
 
 import torch
 from torch._dynamo.device_interface import get_registered_device_interfaces
@@ -36,6 +36,7 @@
     ROCmCodeCache,
     StaticAutotunerFuture,
     torch_key,
+    XPUCodeCache,
 )
 from torch._inductor.compile_worker.subproc_pool import (
     AnyPool,
@@ -47,6 +48,7 @@
 )
 from torch._inductor.compile_worker.utils import _async_compile_initializer
 from torch._inductor.runtime.compile_tasks import (
+    _set_triton_libdevice_path,
     _set_triton_ptxas_path,
     _worker_compile_triton,
 )
@@ -66,13 +68,13 @@
 
 # timing metrics for time spent in the compilation
 _cumulative_compile_time = 0.0
-_t0: Optional[float] = None
+_t0: float | None = None
 
 kernel_code_log = torch._logging.getArtifactLogger(__name__, "kernel_code")
 
 log = logging.getLogger(__name__)
 
-_triton_kernel_metrics: Optional[dict[str, dict[str, Any]]] = None
+_triton_kernel_metrics: dict[str, dict[str, Any]] | None = None
 
 size_hints_regex = re.compile(
     r"size_hints=(\{.*?\})",
@@ -212,7 +214,7 @@ def save(kernel_src: str, future: CodeCacheFuture):
         CompiledTritonKernels._cache[key] = future
 
     @staticmethod
-    def get(kernel_src: str) -> Optional[CodeCacheFuture]:
+    def get(kernel_src: str) -> CodeCacheFuture | None:
         key = CompiledTritonKernels.key(kernel_src)
         return CompiledTritonKernels._cache.get(key, None)
 
@@ -234,7 +236,8 @@ class AsyncCompile:
     Utilities to compile in thread pools or subprocess pools (in the case of Triton).
     """
 
-    _ready_future: Optional[Future[Any]] = None
+    _ready_future: Future[Any] | None = None
+    _metal_sources: list[tuple[str, str, list[str]]] | None = None
 
     def __init__(self) -> None:
         pass
@@ -401,9 +404,15 @@ def reload_kernel_in_parent():
 
         # Cache miss
         if is_parallel:
+            # Ensure libdevice path is set in os.environ before passing to workers
+            _set_triton_libdevice_path()
             # We want to support changing these env vars after (and while) the
             # process pool is running, so pass them to the subprocess to reset.
-            env_vars = ["TORCHINDUCTOR_CACHE_DIR", "TRITON_CACHE_DIR"]
+            env_vars = [
+                "TORCHINDUCTOR_CACHE_DIR",
+                "TRITON_CACHE_DIR",
+                "TRITON_LIBDEVICE_PATH",
+            ]
             extra_env = {v: os.environ[v] for v in env_vars if v in os.environ}
             extra_config = {
                 "use_static_triton_launcher": torch._inductor.config.use_static_triton_launcher
@@ -478,6 +487,7 @@ def get_result() -> CachingAutotuner:
                 try:
                     start_ns = time_ns()
                     _set_triton_ptxas_path()
+                    _set_triton_libdevice_path()
                     kernel = load_kernel()
                     kernel.set_compile_info(compile_id, is_backward)
                     kernel.precompile(
@@ -527,19 +537,25 @@ def cpp_pybinding(self, argtypes: list[str], source_code: str):
             )
             return LambdaFuture(get_result)
 
-    def cuda(self, source_code, dst_file_ext, aot_compile=False):
-        kernel_code_log.info("CUDA Kernel:\n%s", source_code)
-
+    def cutlass(self, cache_cls, source_code, dst_file_ext, aot_compile=False):
         def task():
             if aot_compile:
                 # We rely on JITInductor to compile the CUDA code,
                 # so that we can load it into AOTInductor.
-                output_path, *_ = CUDACodeCache.compile(source_code, "o")
-                CUDACodeCache.aot_kernels_o.append(output_path)
-            return CUDACodeCache.load(source_code, dst_file_ext)[0]
+                output_path, *_ = cache_cls.compile(source_code, "o")
+                cache_cls.aot_kernels_o.append(output_path)
+            return cache_cls.load(source_code, dst_file_ext)[0]
 
         return self.submit(task)
 
+    def cuda(self, source_code, dst_file_ext, aot_compile=False):
+        kernel_code_log.info("CUDA Kernel:\n%s", source_code)
+        return self.cutlass(CUDACodeCache, source_code, dst_file_ext, aot_compile)
+
+    def xpu(self, source_code, dst_file_ext, aot_compile=False):
+        kernel_code_log.info("XPU Kernel:\n%s", source_code)
+        return self.cutlass(XPUCodeCache, source_code, dst_file_ext, aot_compile)
+
     def rocm(
         self,
         source_code,
@@ -687,6 +703,12 @@ def task():
             future = self.submit(task)
             return LambdaFuture(lambda: future.result())
 
+    def metal(self, kernel_name: str, source: str, headers: list[str]) -> None:
+        """Register a Metal kernel body; wait() compiles all registered kernels into one library."""
+        if self._metal_sources is None:
+            self._metal_sources = []
+        self._metal_sources.append((kernel_name, source, headers))
+
     def wait(self, scope: dict[str, Any]) -> None:
         if get_compile_threads() > 1:
             with dynamo_timed(
@@ -698,6 +720,12 @@ def wait(self, scope: dict[str, Any]) -> None:
             ):
                 self._wait_futures(scope)
 
+        if self._metal_sources:
+            from torch._inductor.runtime.runtime_utils import compile_mps_shaders
+
+            scope.update(compile_mps_shaders(self._metal_sources))
+            self._metal_sources.clear()
+
         _compile_end()
 
     def _wait_futures(self, scope: dict[str, Any]) -> None:
diff --git a/torch/_inductor/augmented_graph_helper.py b/torch/_inductor/augmented_graph_helper.py
index 227b8ad664d6e..e6ff7cbe01e41 100644
--- a/torch/_inductor/augmented_graph_helper.py
+++ b/torch/_inductor/augmented_graph_helper.py
@@ -1,11 +1,15 @@
+import logging
 from collections import defaultdict
-from typing import Optional
 
 import torch
 import torch.fx as fx
+from torch._logging import trace_structured
 from torch.utils._ordered_set import OrderedSet
 
 
+log = logging.getLogger(__name__)
+
+
 class AugmentedGraphHelper:
     """
     Graph helper that augments the original graph with additional
@@ -18,7 +22,7 @@ class AugmentedGraphHelper:
     def __init__(
         self,
         graph: fx.Graph,
-        node_ancestors: Optional[dict[fx.Node, OrderedSet[fx.Node]]] = None,
+        node_ancestors: dict[fx.Node, OrderedSet[fx.Node]] | None = None,
     ):
         # Each node starts in its own singleton set
         self.graph = graph
@@ -90,8 +94,22 @@ def get_merged_deps(self, node: fx.Node) -> OrderedSet[fx.Node]:
         return deps
 
     def has_cycle(self) -> bool:
-        merged_deps = {n: self.get_merged_deps(n) for n in self.graph.nodes}
-        return torch._dynamo.graph_deduplication._has_cycle(self.graph, merged_deps)
+        return torch._dynamo.graph_deduplication._has_cycle(
+            self.graph, self.get_all_extra_deps()
+        )
+
+    def _get_all_ancestors(self, node: fx.Node) -> OrderedSet[fx.Node]:
+        """Transitive ancestors through both data deps and extra deps."""
+        ancestors: OrderedSet[fx.Node] = OrderedSet()
+        stack: list[fx.Node] = list(node.all_input_nodes)
+        stack.extend(self.extra_deps.get(node, ()))
+        while stack:
+            n = stack.pop()
+            if n not in ancestors:
+                ancestors.add(n)
+                stack.extend(n.all_input_nodes)
+                stack.extend(self.extra_deps.get(n, ()))
+        return ancestors
 
     def has_path(self, source: fx.Node, target: fx.Node) -> bool:
         """Check if there's a path from source to target."""
@@ -138,6 +156,10 @@ def transfer_erased_node_deps(
         """
         Transfer all extra dependencies from erased nodes to their replacements, handling
         cross-dependencies between erased nodes correctly.
+
+        Skips deps where both endpoints resolve to replacement nodes from the
+        same erasure batch — these are intra-bucket deps that would create
+        cycles (e.g. new_start <-> new_wait within the same bucket).
         """
         erased_merge_sets: dict[fx.Node, fx.Node | None] = {}
 
@@ -162,6 +184,11 @@ def transfer_erased_node_deps(
                 for extra_dep in self.extra_deps[old_node]:
                     updated_dep = erased_merge_sets.get(extra_dep, extra_dep)
                     if updated_dep is not None and updated_dep != new_node:
+                        # Skip if reverse dep already exists (extra or data)
+                        if new_node in self.extra_deps.get(
+                            updated_dep, ()
+                        ) or new_node in OrderedSet(updated_dep.all_input_nodes):
+                            continue
                         self.extra_deps[new_node].add(updated_dep)
                         self.extra_uses[updated_dep].discard(old_node)
                         self.extra_uses[updated_dep].add(new_node)
@@ -170,6 +197,11 @@ def transfer_erased_node_deps(
                 for extra_use in self.extra_uses[old_node]:
                     updated_use = erased_merge_sets.get(extra_use, extra_use)
                     if updated_use is not None and updated_use != new_node:
+                        # Skip if reverse dep already exists (extra or data)
+                        if updated_use in self.extra_deps.get(
+                            new_node, ()
+                        ) or updated_use in OrderedSet(new_node.all_input_nodes):
+                            continue
                         self.extra_deps[updated_use].discard(old_node)
                         self.extra_deps[updated_use].add(new_node)
                         self.extra_uses[new_node].add(updated_use)
@@ -180,6 +212,58 @@ def transfer_erased_node_deps(
             self.extra_uses[old_node].clear()
             del self.merge_sets[old_node]
 
+    def remove_erased_extra_deps(self) -> None:
+        """Remove extra deps referencing erased nodes."""
+        for node in list(self.extra_deps):
+            if node._erased:
+                for dep in list(self.extra_deps[node]):
+                    self.remove_extra_dep(n=node, dep=dep)
+                continue
+            for dep in list(self.extra_deps[node]):
+                if dep._erased:
+                    self.remove_extra_dep(n=node, dep=dep)
+
+    def check_and_maybe_autofix_cyclic_extra_deps(
+        self, *, autofix: bool = False
+    ) -> None:
+        """Check for and optionally remove extra deps that create cycles.
+
+        Args:
+            autofix: If True, silently remove cyclic deps.  If False (default),
+                raise an error so the root cause gets investigated.
+        """
+        if not self.has_cycle():
+            return
+        removed = []
+        for node in list(self.extra_deps):
+            for dep in list(self.extra_deps[node]):
+                ancestors = self._get_all_ancestors(dep)
+                if node in ancestors:
+                    removed.append((node.name, dep.name))
+                    self.remove_extra_dep(n=node, dep=dep)
+        if not removed:
+            return
+        msg = (
+            f"Overlap scheduling: detected {len(removed)} cyclic extra "
+            f"dep(s): {removed}. Please report this to the overlap "
+            f"scheduling developers."
+        )
+        log.warning(msg)
+        trace_structured(
+            "artifact",
+            metadata_fn=lambda: {
+                "name": "inductor_overlap_cyclic_extra_deps",
+                "encoding": "string",
+            },
+            payload_fn=lambda: msg,
+        )
+        if not autofix:
+            raise RuntimeError(
+                f"{msg}\nTo unblock, set "
+                f"torch._inductor.config.aten_distributed_optimizations"
+                f".overlap_scheduling_autofix_cycles = True"
+            )
+
     def get_all_extra_deps(self) -> dict[fx.Node, OrderedSet[fx.Node]]:
         """
         Get all extra dependencies in a format suitable for topological sort.
diff --git a/torch/_inductor/autoheuristic/artifacts/_MMRankingA100.py b/torch/_inductor/autoheuristic/artifacts/_MMRankingA100.py
index 7ebf134c83d7c..c87561d6eb3e4 100644
--- a/torch/_inductor/autoheuristic/artifacts/_MMRankingA100.py
+++ b/torch/_inductor/autoheuristic/artifacts/_MMRankingA100.py
@@ -30,7 +30,7 @@ def check_precondition(self, metadata: AHMetadata, context: AHContext,) -> bool:
     def get_confidence_threshold(self) -> float:
         return 0.0
 
-    def get_choice(self, idx: int) -> Optional[str]:
+    def get_choice(self, idx: int) -> str | None:
         if idx < len(self.choices):
             return self.choices[idx]
         return None
@@ -238,7 +238,7 @@ def fill_choices(self) -> None:
     def get_name(self) -> str:
         return 'mm'
 
-    def get_best_choices(self, context: AHContext) -> Optional[list[tuple[float, int]]]:
+    def get_best_choices(self, context: AHContext) -> list[tuple[float, int]] | None:
         if context.get_value('arith_intensity') <= 52.6245059967041:
             if context.get_value('n') <= 34.0:
                 if context.get_value('n') <= 18.0:
diff --git a/torch/_inductor/autoheuristic/artifacts/_MMRankingH100.py b/torch/_inductor/autoheuristic/artifacts/_MMRankingH100.py
index 6201acc4213aa..74812d9f0b404 100644
--- a/torch/_inductor/autoheuristic/artifacts/_MMRankingH100.py
+++ b/torch/_inductor/autoheuristic/artifacts/_MMRankingH100.py
@@ -30,7 +30,7 @@ def check_precondition(self, metadata: AHMetadata, context: AHContext,) -> bool:
     def get_confidence_threshold(self) -> float:
         return 0.0
 
-    def get_choice(self, idx: int) -> Optional[str]:
+    def get_choice(self, idx: int) -> str | None:
         if idx < len(self.choices):
             return self.choices[idx]
         return None
@@ -242,7 +242,7 @@ def fill_choices(self) -> None:
     def get_name(self) -> str:
         return 'mm'
 
-    def get_best_choices(self, context: AHContext) -> Optional[list[tuple[float, int]]]:
+    def get_best_choices(self, context: AHContext) -> list[tuple[float, int]] | None:
         if context.get_value('arith_intensity') <= 29.89772129058838:
             if context.get_value('n') <= 34.0:
                 if context.get_value('n') <= 18.0:
diff --git a/torch/_inductor/autoheuristic/artifacts/_MixedMMA100.py b/torch/_inductor/autoheuristic/artifacts/_MixedMMA100.py
index 1ba7cbaf90275..a781767a3ce52 100644
--- a/torch/_inductor/autoheuristic/artifacts/_MixedMMA100.py
+++ b/torch/_inductor/autoheuristic/artifacts/_MixedMMA100.py
@@ -30,7 +30,7 @@ def check_precondition(self, metadata: AHMetadata, context: AHContext,) -> bool:
     def get_confidence_threshold(self) -> float:
         return 0.0
 
-    def get_choice(self, idx: int) -> Optional[str]:
+    def get_choice(self, idx: int) -> str | None:
         if idx < len(self.choices):
             return self.choices[idx]
         return None
@@ -62,7 +62,7 @@ def fill_choices(self) -> None:
     def get_name(self) -> str:
         return 'mixed_mm'
 
-    def get_best_choices(self, context: AHContext) -> Optional[list[tuple[float, int]]]:
+    def get_best_choices(self, context: AHContext) -> list[tuple[float, int]] | None:
         if str(context.get_value('1LEQmLEQ16')) != 'True':
             if context.get_value('m') <= 32.5:
                 if context.get_value('n') <= 6976.0:
diff --git a/torch/_inductor/autoheuristic/artifacts/_MixedMMH100.py b/torch/_inductor/autoheuristic/artifacts/_MixedMMH100.py
index 8fe46cf75d8c6..a149cc39d7f6c 100644
--- a/torch/_inductor/autoheuristic/artifacts/_MixedMMH100.py
+++ b/torch/_inductor/autoheuristic/artifacts/_MixedMMH100.py
@@ -30,7 +30,7 @@ def check_precondition(self, metadata: AHMetadata, context: AHContext,) -> bool:
     def get_confidence_threshold(self) -> float:
         return 0.0
 
-    def get_choice(self, idx: int) -> Optional[str]:
+    def get_choice(self, idx: int) -> str | None:
         if idx < len(self.choices):
             return self.choices[idx]
         return None
@@ -61,7 +61,7 @@ def fill_choices(self) -> None:
     def get_name(self) -> str:
         return 'mixed_mm'
 
-    def get_best_choices(self, context: AHContext) -> Optional[list[tuple[float, int]]]:
+    def get_best_choices(self, context: AHContext) -> list[tuple[float, int]] | None:
         if context.get_value('arith_intensity') <= 15.988086223602295:
             if context.get_value('n') <= 25280.0:
                 if context.get_value('n') <= 1344.0:
diff --git a/torch/_inductor/autoheuristic/artifacts/_PadMMA100.py b/torch/_inductor/autoheuristic/artifacts/_PadMMA100.py
index b61f8a9dd1e99..c19817f98484a 100644
--- a/torch/_inductor/autoheuristic/artifacts/_PadMMA100.py
+++ b/torch/_inductor/autoheuristic/artifacts/_PadMMA100.py
@@ -2,16 +2,23 @@
 # fmt: off
 # This file was generated by AutoHeuristic. Do not modify it manually!
 # To regenerate this file, take a look at the steps in the README.md file inside torchgen/_autoheuristic/pad_mm/
-from torch._inductor.autoheuristic.autoheuristic_utils import AHContext, AHMetadata, Choice, CHOICE_COL
+from typing import Optional
+
+from torch._inductor.autoheuristic.autoheuristic_utils import (
+    AHContext,
+    AHMetadata,
+    Choice,
+)
 from torch._inductor.autoheuristic.learnedheuristic_interface import (
-    LearnedHeuristicRegression,
+    LearnedHeuristicDecision,
 )
 
 
-class PadMMA100(LearnedHeuristicRegression):
+class PadMMA100(LearnedHeuristicDecision):
 
     def __init__(self) -> None:
-        pass
+        self.choices: list[Choice] = []
+        self.fill_choices()
 
     def check_precondition(self, metadata: AHMetadata, context: AHContext,) -> bool:
         return (
@@ -20,90 +27,239 @@ def check_precondition(self, metadata: AHMetadata, context: AHContext,) -> bool:
             and str(metadata.device_capa) == "(8, 0)"
         )
 
-    def get_feedback(self, context: AHContext, choice: Choice) -> float:
-        context.context_dict[CHOICE_COL] = choice
-        return self.predict(context)
-
     def get_confidence_threshold(self) -> float:
-        return 1.7025303314066
+        return 0.9294871794871795
+
+    def get_choice(self, idx: int) -> Optional[str]:
+        if idx < len(self.choices):
+            return self.choices[idx]
+        return None
+
+    def fill_choices(self) -> None:
+        self.choices.append('orig')
+        self.choices.append('pad')
 
     def get_name(self) -> str:
         return 'pad_mm'
 
-    def predict(self, context: AHContext) -> float:
-        if str(context.get_value('choice')) != 'pad':
-            if str(context.get_value('using_tf32')) != 'False':
-                if context.get_value('m*n') <= 4171264.0:
-                    if context.get_value('m*k') <= 3999308.0:
-                        return 1.8751469764071178
+    def get_best_choices(self, context: AHContext) -> Optional[list[tuple[float, int]]]:
+        if str(context.get_value('mat1_innermost_needs_padding')) != 'False':
+            if context.get_value('arith_intensity') <= 880.0238037109375:
+                if str(context.get_value('m_multiple_2')) != 'True':
+                    if context.get_value('n') <= 652.0:
+                        if context.get_value('m') <= 2022.0:
+                            if str(context.get_value('using_tf32')) != 'False':
+                                return [(0.579, 1), (0.421, 0)]
+                            else:
+                                return [(1.000, 0)]
+                        else:
+                            return [(1.000, 0)]
                     else:
-                        if str(context.get_value('n_multiple_32')) != 'True':
-                            return 0.9117231355626345
+                        if context.get_value('m*k') <= 107278336.0:
+                            if str(context.get_value('using_tf32')) != 'False':
+                                if context.get_value('m') <= 23691.0:
+                                    return [(0.993, 1), (0.007, 0)]
+                                else:
+                                    return [(0.840, 1), (0.160, 0)]
+                            else:
+                                if context.get_value('arith_intensity') <= 793.3185424804688:
+                                    return [(0.958, 0), (0.042, 1)]
+                                else:
+                                    return [(0.792, 1), (0.208, 0)]
                         else:
-                            return 1.1607689608873861
+                            if context.get_value('arith_intensity') <= 795.6242370605469:
+                                if context.get_value('mat2_stride_1') <= 2048.5:
+                                    return [(0.929, 0), (0.071, 1)]
+                                else:
+                                    return [(1.000, 1)]
+                            else:
+                                if context.get_value('arith_intensity') <= 796.3460388183594:
+                                    return [(0.957, 1), (0.043, 0)]
+                                else:
+                                    return [(0.778, 0), (0.222, 1)]
                 else:
-                    if str(context.get_value('n_multiple_2')) != 'True':
-                        if str(context.get_value('using_tf32')) != 'True':
-                            return 0.7430382200435992
+                    if context.get_value('mat2_stride_0') <= 2432.0:
+                        if str(context.get_value('k_multiple_2')) != 'False':
+                            if context.get_value('n') <= 1024.5:
+                                if str(context.get_value('prepadded_mat1')) != 'False':
+                                    return [(0.580, 0), (0.420, 1)]
+                                else:
+                                    return [(0.986, 0), (0.014, 1)]
+                            else:
+                                if context.get_value('mat1_stride_0') <= 5125.0:
+                                    return [(0.551, 0), (0.449, 1)]
+                                else:
+                                    return [(0.916, 0), (0.084, 1)]
                         else:
-                            return 0.8531269794448678
+                            if context.get_value('mat2_align_size') <= 6.0:
+                                if str(context.get_value('using_tf32')) != 'True':
+                                    return [(1.000, 0)]
+                                else:
+                                    return [(0.800, 0), (0.200, 1)]
+                            else:
+                                if context.get_value('mat2_stride_1') <= 3820.0:
+                                    return [(0.986, 1), (0.014, 0)]
+                                else:
+                                    return [(0.532, 0), (0.468, 1)]
                     else:
-                        if str(context.get_value('k_multiple_2')) != 'True':
-                            return 0.7577181972719917
+                        if str(context.get_value('using_tf32')) != 'False':
+                            if context.get_value('m*n') <= 5244928.0:
+                                if str(context.get_value('k_multiple_2')) != 'True':
+                                    return [(0.971, 1), (0.029, 0)]
+                                else:
+                                    return [(0.646, 1), (0.354, 0)]
+                            else:
+                                if context.get_value('k/(m*n)') <= 9.468618827668251e-06:
+                                    return [(0.800, 1), (0.200, 0)]
+                                else:
+                                    return [(1.000, 1)]
                         else:
-                            return 0.8977349440424219
+                            if context.get_value('mat1_stride_1') <= 1288.0:
+                                if context.get_value('k') <= 5717.0:
+                                    return [(0.983, 0), (0.017, 1)]
+                                else:
+                                    return [(0.800, 0), (0.200, 1)]
+                            else:
+                                return [(0.588, 0), (0.412, 1)]
             else:
-                if context.get_value('m*n') <= 1299712.0:
-                    return 1.1669723418995592
-                else:
-                    if context.get_value('mat2_stride_1') <= 45217.5:
-                        if context.get_value('m*n') <= 55884158.0:
-                            return 1.0262769936909601
+                if str(context.get_value('using_tf32')) != 'False':
+                    if context.get_value('n') <= 2640.0:
+                        if context.get_value('m_padded_length') <= 1.5:
+                            if context.get_value('mat2_stride_1') <= 6021.0:
+                                return [(1.000, 1)]
+                            else:
+                                if context.get_value('mat2_stride_1') <= 6069.0:
+                                    return [(0.900, 1), (0.100, 0)]
+                                else:
+                                    return [(1.000, 1)]
                         else:
-                            return 1.0022677428470845
+                            if str(context.get_value('m_multiple_2')) != 'False':
+                                if context.get_value('m*k') <= 24444928.0:
+                                    return [(0.593, 0), (0.407, 1)]
+                                else:
+                                    return [(0.923, 0), (0.077, 1)]
+                            else:
+                                return [(1.000, 1)]
                     else:
-                        if context.get_value('m') <= 18478.0:
-                            return 1.1127066261894312
+                        if context.get_value('m*k') <= 404182016.0:
+                            if str(context.get_value('mat2_innermost_needs_padding')) != 'False':
+                                if context.get_value('m*k') <= 12328960.0:
+                                    return [(0.732, 1), (0.268, 0)]
+                                else:
+                                    return [(0.989, 1), (0.011, 0)]
+                            else:
+                                if context.get_value('m*k') <= 389028864.0:
+                                    return [(0.998, 1), (0.002, 0)]
+                                else:
+                                    return [(0.922, 1), (0.078, 0)]
                         else:
-                            return 1.0337740659894263
+                            if context.get_value('m*n') <= 137631744.0:
+                                if context.get_value('m*k') <= 405715968.0:
+                                    return [(0.611, 1), (0.389, 0)]
+                                else:
+                                    return [(0.946, 1), (0.054, 0)]
+                            else:
+                                return [(0.714, 0), (0.286, 1)]
+                else:
+                    if context.get_value('mat2_stride_0') <= 3902.5:
+                        return [(0.941, 0), (0.059, 1)]
+                    else:
+                        return [(0.583, 1), (0.417, 0)]
         else:
-            if str(context.get_value('mat1_dtype')) != 'torch.float32':
-                if str(context.get_value('n_multiple_2')) != 'False':
+            if context.get_value('n_padded_length') <= 0.5:
+                if str(context.get_value('mat2_innermost_needs_padding')) != 'False':
                     if str(context.get_value('k_multiple_2')) != 'True':
-                        if context.get_value('mat1_stride_0') <= 561.0:
-                            return 1.2900382135142956
+                        if context.get_value('arith_intensity') <= 884.5185852050781:
+                            if context.get_value('arith_intensity') <= 743.931884765625:
+                                return [(1.000, 0)]
+                            else:
+                                return [(0.583, 0), (0.417, 1)]
                         else:
-                            return 1.5761737616057887
+                            return [(1.000, 1)]
                     else:
-                        if context.get_value('num_dims_needs_padding') <= 1.5:
-                            return 1.0472263310239422
+                        if context.get_value('k/(m*n)') <= 0.00023481580865336582:
+                            return [(0.900, 0), (0.100, 1)]
                         else:
-                            return 1.1727673465762514
+                            return [(1.000, 0)]
                 else:
-                    if context.get_value('k') <= 28238.5:
-                        if context.get_value('k/(m*n)') <= 0.00026227018679492176:
-                            return 1.6770542505397175
+                    if str(context.get_value('using_tf32')) != 'False':
+                        if context.get_value('m*k') <= 93734912.0:
+                            if context.get_value('mat1_stride_0') <= 1344.0:
+                                if context.get_value('n') <= 7168.0:
+                                    return [(1.000, 0)]
+                                else:
+                                    return [(0.970, 0), (0.030, 1)]
+                            else:
+                                if context.get_value('m') <= 22883.5:
+                                    return [(0.977, 0), (0.023, 1)]
+                                else:
+                                    return [(0.800, 0), (0.200, 1)]
                         else:
-                            return 1.3974785435105923
+                            if context.get_value('arith_intensity') <= 1914.3681030273438:
+                                if str(context.get_value('prepadded_mat1')) != 'False':
+                                    return [(0.981, 0), (0.019, 1)]
+                                else:
+                                    return [(0.995, 0), (0.005, 1)]
+                            else:
+                                return [(1.000, 0)]
                     else:
-                        if str(context.get_value('mat1_dtype')) != 'torch.bfloat16':
-                            return 1.3952699800111992
+                        if str(context.get_value('prepadded_mat1')) != 'False':
+                            if context.get_value('mat2_stride_1') <= 256.5:
+                                if context.get_value('m') <= 5880.5:
+                                    return [(1.000, 0)]
+                                else:
+                                    return [(0.800, 0), (0.200, 1)]
+                            else:
+                                if context.get_value('m*k') <= 6318080.0:
+                                    return [(0.618, 1), (0.382, 0)]
+                                else:
+                                    return [(0.880, 0), (0.120, 1)]
                         else:
-                            return 1.5759286511628336
+                            if context.get_value('k/(m*n)') <= 0.0009747986623551697:
+                                if context.get_value('k/(m*n)') <= 0.0006397514371201396:
+                                    return [(0.951, 0), (0.049, 1)]
+                                else:
+                                    return [(0.857, 0), (0.143, 1)]
+                            else:
+                                return [(1.000, 0)]
             else:
                 if str(context.get_value('using_tf32')) != 'False':
-                    if context.get_value('m*n') <= 14119424.0:
-                        return 0.8875772670422478
+                    if str(context.get_value('n_multiple_2')) != 'False':
+                        if context.get_value('m') <= 2024.0:
+                            if context.get_value('mat2_stride_0') <= 1629.0:
+                                if context.get_value('k*n') <= 1288704.0:
+                                    return [(0.600, 0), (0.400, 1)]
+                                else:
+                                    return [(0.982, 0), (0.018, 1)]
+                            else:
+                                if context.get_value('mat1_stride_0') <= 768.0:
+                                    return [(0.619, 0), (0.381, 1)]
+                                else:
+                                    return [(0.812, 1), (0.188, 0)]
+                        else:
+                            if context.get_value('m*n') <= 5803008.0:
+                                return [(0.500, 0), (0.500, 1)]
+                            else:
+                                if context.get_value('mat2_stride_1') <= 896.0:
+                                    return [(1.000, 1)]
+                                else:
+                                    return [(0.818, 1), (0.182, 0)]
                     else:
-                        if str(context.get_value('mat2_innermost_needs_padding')) != 'True':
-                            return 1.1467728924377265
+                        if context.get_value('mat2_stride_1') <= 2560.0:
+                            if context.get_value('num_dims_needs_padding') <= 1.5:
+                                if context.get_value('mat2_stride_1') <= 896.0:
+                                    return [(1.000, 1)]
+                                else:
+                                    return [(0.857, 1), (0.143, 0)]
+                            else:
+                                return [(0.727, 1), (0.273, 0)]
                         else:
-                            return 1.215842963532998
+                            return [(0.667, 1), (0.333, 0)]
                 else:
-                    if context.get_value('arith_intensity') <= 396.8774871826172:
-                        return 0.89940161869551
+                    if context.get_value('k/(m*n)') <= 0.00015462777810171247:
+                        return [(0.857, 0), (0.143, 1)]
                     else:
-                        if context.get_value('mat2_stride_1') <= 45217.5:
-                            return 0.9964328169353532
+                        if context.get_value('k/(m*n)') <= 0.0019917909521609545:
+                            return [(1.000, 0)]
                         else:
-                            return 0.9493479238294826
+                            return [(0.900, 0), (0.100, 1)]
diff --git a/torch/_inductor/autoheuristic/artifacts/_PadMMH200.py b/torch/_inductor/autoheuristic/artifacts/_PadMMH200.py
new file mode 100644
index 0000000000000..1527b02319c10
--- /dev/null
+++ b/torch/_inductor/autoheuristic/artifacts/_PadMMH200.py
@@ -0,0 +1,175 @@
+# flake8: noqa: B950
+# fmt: off
+# This file was generated by AutoHeuristic. Do not modify it manually!
+# To regenerate this file, take a look at the steps in the README.md file inside torchgen/_autoheuristic/pad_mm/
+from typing import Optional
+
+from torch._inductor.autoheuristic.autoheuristic_utils import (
+    AHContext,
+    AHMetadata,
+    Choice,
+)
+from torch._inductor.autoheuristic.learnedheuristic_interface import (
+    LearnedHeuristicDecision,
+)
+
+
+class PadMMH200(LearnedHeuristicDecision):
+
+    def __init__(self) -> None:
+        self.choices: list[Choice] = []
+        self.fill_choices()
+
+    def check_precondition(self, metadata: AHMetadata, context: AHContext,) -> bool:
+        return (
+            metadata.name == self.get_name()
+            and metadata.shared_memory == 232448
+            and str(metadata.device_capa) == "(9, 0)"
+        )
+
+    def get_confidence_threshold(self) -> float:
+        return 0.7710651828298887
+
+    def get_choice(self, idx: int) -> Optional[str]:
+        if idx < len(self.choices):
+            return self.choices[idx]
+        return None
+
+    def fill_choices(self) -> None:
+        self.choices.append('orig')
+        self.choices.append('pad')
+
+    def get_name(self) -> str:
+        return 'pad_mm'
+
+    def get_best_choices(self, context: AHContext) -> Optional[list[tuple[float, int]]]:
+        if str(context.get_value('mat1_innermost_needs_padding')) != 'True':
+            if str(context.get_value('mat2_innermost_needs_padding')) != 'True':
+                if context.get_value('n_padded_length') <= 0.5:
+                    if str(context.get_value('prepadded_mat1')) != 'True':
+                        if str(context.get_value('using_tf32')) != 'False':
+                            return [(1.000, 0)]
+                        else:
+                            if context.get_value('mat1_stride_0') <= 3584.0:
+                                return [(1.000, 0)]
+                            else:
+                                if context.get_value('mat2_stride_0') <= 3584.0:
+                                    return [(1.000, 0)]
+                                else:
+                                    return [(0.528, 0), (0.472, 1)]
+                    else:
+                        if context.get_value('n') <= 2304.0:
+                            if context.get_value('m*k') <= 25198592.0:
+                                if context.get_value('arith_intensity') <= 1103.9319458007812:
+                                    return [(1.000, 0)]
+                                else:
+                                    return [(0.885, 0), (0.115, 1)]
+                            else:
+                                if context.get_value('m*k') <= 25688064.0:
+                                    return [(1.000, 1)]
+                                else:
+                                    return [(0.771, 0), (0.229, 1)]
+                        else:
+                            if str(context.get_value('using_tf32')) != 'False':
+                                if context.get_value('m') <= 27825.0:
+                                    return [(0.948, 0), (0.052, 1)]
+                                else:
+                                    return [(0.855, 0), (0.145, 1)]
+                            else:
+                                if context.get_value('mat2_stride_0') <= 3584.0:
+                                    return [(1.000, 0)]
+                                else:
+                                    return [(0.917, 1), (0.083, 0)]
+                else:
+                    if context.get_value('m') <= 1823.5:
+                        if str(context.get_value('n_multiple_2')) != 'False':
+                            if context.get_value('k*n') <= 7859200.0:
+                                return [(0.600, 0), (0.400, 1)]
+                            else:
+                                return [(1.000, 0)]
+                        else:
+                            if context.get_value('k/(m*n)') <= 0.00040277576772496104:
+                                return [(1.000, 1)]
+                            else:
+                                return [(0.800, 1), (0.200, 0)]
+                    else:
+                        if context.get_value('n') <= 3602.0:
+                            return [(0.800, 1), (0.200, 0)]
+                        else:
+                            return [(1.000, 1)]
+            else:
+                if str(context.get_value('using_tf32')) != 'False':
+                    if str(context.get_value('n_multiple_16')) != 'False':
+                        if str(context.get_value('k_multiple_2')) != 'True':
+                            if context.get_value('arith_intensity') <= 744.8332214355469:
+                                return [(0.600, 0), (0.400, 1)]
+                            else:
+                                return [(1.000, 1)]
+                        else:
+                            if context.get_value('m*n') <= 8912896.0:
+                                if context.get_value('m*k') <= 5934080.0:
+                                    return [(0.800, 0), (0.200, 1)]
+                                else:
+                                    return [(1.000, 0)]
+                            else:
+                                return [(1.000, 1)]
+                    else:
+                        return [(1.000, 1)]
+                else:
+                    return [(1.000, 0)]
+        else:
+            if context.get_value('arith_intensity') <= 895.8767395019531:
+                if str(context.get_value('m_multiple_2')) != 'False':
+                    if context.get_value('mat1_stride_1') <= 3421.0:
+                        if str(context.get_value('using_tf32')) != 'False':
+                            if context.get_value('mat2_stride_1') <= 10706.5:
+                                if context.get_value('mat2_stride_0') <= 1024.5:
+                                    return [(0.816, 1), (0.184, 0)]
+                                else:
+                                    return [(1.000, 1)]
+                            else:
+                                if str(context.get_value('k_multiple_2')) != 'True':
+                                    return [(0.905, 1), (0.095, 0)]
+                                else:
+                                    return [(1.000, 0)]
+                        else:
+                            if str(context.get_value('prepadded_mat2')) != 'True':
+                                if str(context.get_value('mat2_innermost_needs_padding')) != 'False':
+                                    return [(1.000, 0)]
+                                else:
+                                    return [(0.932, 0), (0.068, 1)]
+                            else:
+                                if context.get_value('arith_intensity') <= 742.1241760253906:
+                                    return [(0.889, 0), (0.111, 1)]
+                                else:
+                                    return [(0.765, 1), (0.235, 0)]
+                    else:
+                        if context.get_value('n') <= 1216.0:
+                            if str(context.get_value('using_tf32')) != 'True':
+                                if context.get_value('mat1_stride_1') <= 5567.0:
+                                    return [(0.896, 0), (0.104, 1)]
+                                else:
+                                    return [(0.999, 0), (0.001, 1)]
+                            else:
+                                return [(1.000, 1)]
+                        else:
+                            if str(context.get_value('using_tf32')) != 'False':
+                                return [(1.000, 1)]
+                            else:
+                                return [(1.000, 0)]
+                else:
+                    if str(context.get_value('using_tf32')) != 'False':
+                        return [(1.000, 1)]
+                    else:
+                        if context.get_value('mat2_stride_1') <= 2688.0:
+                            return [(1.000, 0)]
+                        else:
+                            return [(0.500, 0), (0.500, 1)]
+            else:
+                if str(context.get_value('using_tf32')) != 'False':
+                    return [(1.000, 1)]
+                else:
+                    if str(context.get_value('mat2_innermost_needs_padding')) != 'True':
+                        return [(1.000, 0)]
+                    else:
+                        return [(0.800, 0), (0.200, 1)]
diff --git a/torch/_inductor/autoheuristic/autoheuristic.py b/torch/_inductor/autoheuristic/autoheuristic.py
index 898915acaf991..d0072dd64310a 100644
--- a/torch/_inductor/autoheuristic/autoheuristic.py
+++ b/torch/_inductor/autoheuristic/autoheuristic.py
@@ -2,7 +2,7 @@
 import os
 from collections.abc import Callable
 from functools import partial
-from typing import Any, Optional
+from typing import Any
 
 import torch
 from torch._inductor.autoheuristic.autoheuristic_utils import (
@@ -57,11 +57,11 @@ def __init__(
         self,
         fallback: Callable[[], Choice],
         choices: list[Choice],
-        feedback: Optional[LocalFeedback],
+        feedback: LocalFeedback | None,
         context: AHContext,
         name: str,
-        augment_context: Optional[list[AHOperation]] = None,
-        precondition: Optional[Callable[[AHMetadata, AHContext], bool]] = None,
+        augment_context: list[AHOperation] | None = None,
+        precondition: Callable[[AHMetadata, AHContext], bool] | None = None,
     ) -> None:
         """
         Initializes an instance of the AutoHeuristic class.
@@ -136,8 +136,8 @@ def get_choice(self) -> Choice:
         return self.fallback()
 
     def get_top_k_choices(
-        self, top_k: int, always_included: Optional[list[str]] = None
-    ) -> Optional[list[Choice]]:
+        self, top_k: int, always_included: list[str] | None = None
+    ) -> list[Choice] | None:
         if not self.satisfies_precondition():
             return None
         if torch._inductor.config.use_autoheuristic(self.name):
@@ -223,13 +223,13 @@ class AutoHeuristicSelectAlgorithm(AutoHeuristic):
 
     def __init__(
         self,
-        fallback: Callable[[], Optional[ChoiceCaller]],
+        fallback: Callable[[], ChoiceCaller | None],
         choices: list[ChoiceCaller],
         input_nodes: list[Any],
         context: AHContext,
         name: str,
-        augment_context: Optional[list[AHOperation]] = None,
-        precondition: Optional[Callable[[AHMetadata, AHContext], bool]] = None,
+        augment_context: list[AHOperation] | None = None,
+        precondition: Callable[[AHMetadata, AHContext], bool] | None = None,
     ) -> None:
         """
         The arguments choices, input_nodes and name have to match the ones used in the call to
@@ -305,13 +305,13 @@ def store_global_feedback(
         feedback_saver = partial(store_global_feedback, inputs_key, precompile_key)
         add_feedback_saver(feedback_saver)
 
-    def get_choice_caller(self) -> Optional[ChoiceCaller]:
+    def get_choice_caller(self) -> ChoiceCaller | None:
         choice = self.get_choice()
         return self.choicestr2choice.get(choice, None)
 
     def get_top_k_choices_caller(
-        self, top_k: int, always_included: Optional[list[str]] = None
-    ) -> Optional[list[ChoiceCaller]]:
+        self, top_k: int, always_included: list[str] | None = None
+    ) -> list[ChoiceCaller] | None:
         choices = self.get_top_k_choices(top_k, always_included)
         if choices is None:
             return None
diff --git a/torch/_inductor/autoheuristic/learned_heuristic_controller.py b/torch/_inductor/autoheuristic/learned_heuristic_controller.py
index 50c11eb9a712a..26f03e3838c33 100644
--- a/torch/_inductor/autoheuristic/learned_heuristic_controller.py
+++ b/torch/_inductor/autoheuristic/learned_heuristic_controller.py
@@ -2,7 +2,7 @@
 import inspect
 import pkgutil
 from collections import defaultdict
-from typing import Any, Optional
+from typing import Any
 
 from torch._inductor.autoheuristic.autoheuristic_utils import (
     AHContext,
@@ -93,7 +93,7 @@ def get_heuristics(self, name: str) -> list[LearnedHeuristic]:
 
         return LearnedHeuristicController.existing_heuristics[name]
 
-    def get_decision(self) -> Optional[Choice]:
+    def get_decision(self) -> Choice | None:
         """
         Returns the decision made by the learned heuristic or None if no heuristic was found or the heuristic is unsure
         which choice to make.
@@ -105,7 +105,7 @@ def get_decision(self) -> Optional[Choice]:
                 return heuristic.get_decision(self.context, self.metadata.choices)
         return None
 
-    def get_decisions_ranked(self, top_k: int) -> Optional[list[Choice]]:
+    def get_decisions_ranked(self, top_k: int) -> list[Choice] | None:
         heuristics = self.get_heuristics(self.metadata.name)
         for heuristic in heuristics:
             if heuristic.check_precondition(self.metadata, self.context):
diff --git a/torch/_inductor/autoheuristic/learnedheuristic_interface.py b/torch/_inductor/autoheuristic/learnedheuristic_interface.py
index 84a941b076c31..85558cf8e488c 100644
--- a/torch/_inductor/autoheuristic/learnedheuristic_interface.py
+++ b/torch/_inductor/autoheuristic/learnedheuristic_interface.py
@@ -1,5 +1,4 @@
 import operator
-from typing import Optional
 
 from torch._inductor.autoheuristic.autoheuristic_utils import (
     AHContext,
@@ -23,9 +22,7 @@ def check_precondition(
     ) -> bool:
         return True
 
-    def get_decision(
-        self, context: AHContext, choices: list[Choice]
-    ) -> Optional[Choice]:
+    def get_decision(self, context: AHContext, choices: list[Choice]) -> Choice | None:
         return None
 
     def get_confidence_threshold(self) -> float:
@@ -34,7 +31,7 @@ def get_confidence_threshold(self) -> float:
     def get_name(self) -> str:
         return ""
 
-    def get_decisions_ranked(self, context: AHContext) -> Optional[list[str]]:
+    def get_decisions_ranked(self, context: AHContext) -> list[str] | None:
         return None
 
 
@@ -42,9 +39,7 @@ class LearnedHeuristicRegression(LearnedHeuristic):
     def get_feedback(self, context: AHContext, choice: Choice) -> float:
         return 1.0
 
-    def get_decision(
-        self, context: AHContext, choices: list[Choice]
-    ) -> Optional[Choice]:
+    def get_decision(self, context: AHContext, choices: list[Choice]) -> Choice | None:
         choice2feedback = {}
         for choice in choices:
             predicted_feedback = self.get_feedback(context, choice)
@@ -61,12 +56,10 @@ def get_decision(
 
 
 class LearnedHeuristicDecision(LearnedHeuristic):
-    def get_choice(self, idx: int) -> Optional[str]:
+    def get_choice(self, idx: int) -> str | None:
         return None
 
-    def get_decision(
-        self, context: AHContext, choices: list[Choice]
-    ) -> Optional[Choice]:
+    def get_decision(self, context: AHContext, choices: list[Choice]) -> Choice | None:
         best_choices = self.get_best_choices(context)
         if not best_choices:
             return None
@@ -75,7 +68,7 @@ def get_decision(
             return None
         return self.get_choice(best_choice_idx)
 
-    def get_decisions_ranked(self, context: AHContext) -> Optional[list[str]]:
+    def get_decisions_ranked(self, context: AHContext) -> list[str] | None:
         feedback_idx_list = self.get_best_choices(context)
         if feedback_idx_list is None:
             return None
@@ -85,5 +78,5 @@ def get_decisions_ranked(self, context: AHContext) -> Optional[list[str]]:
         choices = [choice for choice in choices if choice is not None]
         return choices
 
-    def get_best_choices(self, context: AHContext) -> Optional[list[tuple[float, int]]]:
+    def get_best_choices(self, context: AHContext) -> list[tuple[float, int]] | None:
         return []
diff --git a/torch/_inductor/autotune_process.py b/torch/_inductor/autotune_process.py
index 232849d640fe6..781c872a7d114 100644
--- a/torch/_inductor/autotune_process.py
+++ b/torch/_inductor/autotune_process.py
@@ -2,6 +2,7 @@
 from __future__ import annotations
 
 import atexit
+import contextvars
 import ctypes
 import dataclasses
 import functools
@@ -19,7 +20,7 @@
 from collections.abc import Callable, Iterable, Sequence
 from concurrent.futures import Future, ProcessPoolExecutor, ThreadPoolExecutor
 from ctypes import byref, c_size_t, c_void_p, CDLL
-from typing import Any, IO, Optional, TYPE_CHECKING, Union
+from typing import Any, IO, TYPE_CHECKING
 
 import torch
 import torch._inductor.async_compile  # noqa: F401 required to warm up AsyncCompile pools
@@ -32,6 +33,7 @@
     DLLWrapper,
     get_hash,
     PyCodeCache,
+    XPUCodeCache,
 )
 from torch._inductor.compile_worker.timer import Timer
 from torch._inductor.utils import (
@@ -121,7 +123,7 @@ def send(
     def recv(read_pipe: IO[bytes]) -> Any:
         return pickle.load(read_pipe)
 
-    def __init__(self, device: Optional[int]):
+    def __init__(self, device: int | None):
         self.device = device
         self.start()
 
@@ -284,7 +286,7 @@ def __init__(self) -> None:
         self.executor = ThreadPoolExecutor(max_workers=len(devices))
 
     @staticmethod
-    def get_device_list() -> Sequence[Optional[int]]:
+    def get_device_list() -> Sequence[int | None]:
         """
         Gather the list of devices to be used in the pool.
         """
@@ -343,10 +345,13 @@ def target(self, choice: TritonTemplateCaller) -> float:
                 f"Failed to benchmark choice '{choice}'. It will be ignored. "
                 "Please debug the root cause in case the choice can bring perf gains."
             )
-            # An unspecified launch failure (cudaErrorLaunchFailure) corrupts the
-            # CUDA context, making it unrecoverable. All subsequent CUDA calls will
-            # fail as well. The process must be restarted to restore CUDA functionality.
-            if "cudaErrorLaunchFailure" in str(process_exception):
+            # Sticky CUDA errors corrupt the context, making it unrecoverable.
+            # The process must be restarted to restore CUDA functionality.
+            error_msg = str(process_exception)
+            if (
+                "cudaErrorLaunchFailure" in error_msg
+                or "cudaErrorIllegalAddress" in error_msg
+            ):
                 process.restart()
             # Set to INF so this choice will be ignored
             return float("inf")
@@ -368,7 +373,7 @@ def benchmark(
         return results
 
 
-LayoutOrBuffer = Union[ir.Layout, ir.Buffer]
+LayoutOrBuffer = ir.Layout | ir.Buffer
 
 
 @dataclasses.dataclass
@@ -378,12 +383,14 @@ class TensorMeta:
     sizes: torch._prims_common.ShapeType
     strides: torch._prims_common.StrideType
     offset: int
-    name: Optional[str] = None
+    name: str | None = None
 
     @classmethod
     def from_irnodes(
-        cls, irnodes: Union[LayoutOrBuffer, Sequence[LayoutOrBuffer]]
-    ) -> Union[TensorMeta, list[TensorMeta]]:
+        cls, irnodes: LayoutOrBuffer | Sequence[LayoutOrBuffer]
+    ) -> TensorMeta | list[TensorMeta]:
+        from torch._inductor.select_algorithm import get_strides_with_layout_constraints
+
         if isinstance(irnodes, Sequence):
             result: list[Any] = [cls.from_irnodes(x) for x in irnodes]
             assert all(isinstance(x, TensorMeta) for x in result)
@@ -402,7 +409,9 @@ def from_irnodes(
             device=device,
             dtype=dtype,
             sizes=V.graph.sizevars.optimization_hints(node.get_size()),
-            strides=V.graph.sizevars.optimization_hints(node.get_stride()),
+            strides=V.graph.sizevars.optimization_hints(
+                get_strides_with_layout_constraints(node)
+            ),
             offset=V.graph.sizevars.optimization_hint(node.get_layout().offset),
             name=node.get_name(),
         )
@@ -430,8 +439,8 @@ class BenchmarkRequest:
     def __init__(
         self,
         kernel_name: str,
-        input_tensor_meta: Union[TensorMeta, list[TensorMeta]],
-        output_tensor_meta: Union[TensorMeta, list[TensorMeta]],
+        input_tensor_meta: TensorMeta | list[TensorMeta],
+        output_tensor_meta: TensorMeta | list[TensorMeta],
         extra_args: Iterable[Any],
     ) -> None:
         # the kernel name defined in the module
@@ -456,6 +465,7 @@ def __init__(
             self.output_tensor_meta: TensorMeta = output_tensor_meta
 
         self.extra_args = extra_args
+        self.benchmark_with_cudagraphs = False
 
     def make_run_fn(
         self, *input_tensors: torch.Tensor, out: torch.Tensor
@@ -469,14 +479,14 @@ def do_bench(
         self,
         fn,
         *input_tensors: torch.Tensor,
-        out: Optional[torch.Tensor] = None,
+        out: torch.Tensor | None = None,
     ) -> float:
         raise NotImplementedError
 
     def benchmark(
         self,
         *input_tensors: torch.Tensor,
-        out: Optional[torch.Tensor] = None,
+        out: torch.Tensor | None = None,
     ) -> float:
         debug = autotuning_log.isEnabledFor(logging.DEBUG)
         if debug:
@@ -495,29 +505,34 @@ def benchmark(
             create_tensor_elapse = time.time() - start_ts  # type: ignore[possibly-undefined]
             start_ts = time.time()
         try:
-            fn = self.make_run_fn(*input_tensors, out=out)
-        except NonzeroWorkspaceNotSupportedError:
-            # Skipping all ops with nonzero workspace requirements
-            autotuning_log.info("Skipping op due to nonzero workspace requirement")
-            return float("inf")
-
-        if debug:
-            load_elapse = time.time() - start_ts  # type: ignore[possibly-undefined]
-            start_ts = time.time()
-
-        res = self.do_bench(fn, *input_tensors, out)
-
-        if debug:
-            bench_elapse = time.time() - start_ts  # type: ignore[possibly-undefined]
-            autotuning_log.debug(
-                "InChildProcess %s: load %f, create tensor %f, bench %f",
-                str(self),
-                load_elapse,  # type: ignore[possibly-undefined]
-                create_tensor_elapse,  # type: ignore[possibly-undefined]
-                bench_elapse,
-            )
-        self.cleanup_run_fn()
-        return res
+            try:
+                fn = self.make_run_fn(*input_tensors, out=out)
+            except NonzeroWorkspaceNotSupportedError:
+                # Skipping all ops with nonzero workspace requirements
+                autotuning_log.info("Skipping op due to nonzero workspace requirement")
+                return float("inf")
+
+            if debug:
+                load_elapse = time.time() - start_ts  # type: ignore[possibly-undefined]
+                start_ts = time.time()
+
+            if self.benchmark_with_cudagraphs:
+                res = benchmarker.benchmark_gpu_with_cuda_graph(fn)
+            else:
+                res = self.do_bench(fn, *input_tensors, out)
+
+            if debug:
+                bench_elapse = time.time() - start_ts  # type: ignore[possibly-undefined]
+                autotuning_log.debug(
+                    "InChildProcess %s: load %f, create tensor %f, bench %f",
+                    self,
+                    load_elapse,  # type: ignore[possibly-undefined]
+                    create_tensor_elapse,  # type: ignore[possibly-undefined]
+                    bench_elapse,
+                )
+            return res
+        finally:
+            self.cleanup_run_fn()
 
 
 class _TestBenchmarkRequest(BenchmarkRequest):
@@ -529,9 +544,9 @@ class _TestBenchmarkRequest(BenchmarkRequest):
     def __init__(
         self,
         result: float = 0.0,
-        device: Optional[int] = None,
-        sleep: Optional[float] = None,
-        exc: Optional[Exception] = None,
+        device: int | None = None,
+        sleep: float | None = None,
+        exc: Exception | None = None,
         crash: bool = False,
     ):
         self.result = result
@@ -541,7 +556,7 @@ def __init__(
         self.crash = crash
 
     def benchmark(
-        self, *input_tensors: torch.Tensor, out: Optional[torch.Tensor] = None
+        self, *input_tensors: torch.Tensor, out: torch.Tensor | None = None
     ) -> float:
         if self.device is not None:
             assert os.environ.get(CUDA_VISIBLE_DEVICES, None) == str(self.device)
@@ -559,7 +574,7 @@ def do_bench(
         self,
         fn,
         *input_tensors: torch.Tensor,
-        out: Optional[torch.Tensor] = None,
+        out: torch.Tensor | None = None,
     ) -> float:
         device_idx_set = OrderedSet(
             tensor.device.index
@@ -594,7 +609,7 @@ def do_bench(
         self,
         fn,
         *input_tensors: torch.Tensor,
-        out: Optional[torch.Tensor] = None,
+        out: torch.Tensor | None = None,
     ) -> float:
         return benchmarker.benchmark_cpu(fn)
 
@@ -610,8 +625,8 @@ class TritonBenchmarkRequest(BenchmarkRequest):
     def __init__(
         self,
         kernel_name: str,
-        input_tensor_meta: Union[TensorMeta, list[TensorMeta]],
-        output_tensor_meta: Union[TensorMeta, list[TensorMeta]],
+        input_tensor_meta: TensorMeta | list[TensorMeta],
+        output_tensor_meta: TensorMeta | list[TensorMeta],
         extra_args: Iterable[Any],
         module_path: str,  # the path of the module defining the triton kernel
         module_cache_key: str,
@@ -622,7 +637,7 @@ def __init__(
         matrix_instr_nonkdim: int = 0,  # only used for hip to choose the shape of mfma instruction.
         waves_per_eu: int = 0,  # only used for hip to schedule waves per execution unit
         kpack: int = 0,  # ROCm specific gemm parameter
-        workspace_size: Optional[int] = None,  # size of workspace buffer in bytes
+        workspace_size: int | None = None,  # size of workspace buffer in bytes
         workspace_zero_fill: bool = False,  # whether to zero-fill workspace
     ) -> None:
         super().__init__(kernel_name, input_tensor_meta, output_tensor_meta, extra_args)
@@ -637,11 +652,17 @@ def __init__(
         self.kpack = kpack
         self.workspace_size = workspace_size
         self.workspace_zero_fill = workspace_zero_fill
+        self._benchmark_module: Any | None = None
 
     def make_run_fn(
         self, *input_tensors: torch.Tensor, out: torch.Tensor
     ) -> Callable[[], None]:
-        mod = PyCodeCache.load_by_key_path(self.module_cache_key, self.module_path)
+        mod = PyCodeCache.load_by_key_path(
+            self.module_cache_key,
+            self.module_path,
+            set_sys_modules=False,
+        )
+        self._benchmark_module = mod
         autotuning_log.debug(
             "benchmark module key: %s, path: %s",
             self.module_cache_key,
@@ -709,9 +730,64 @@ def make_run_fn(
                 benchmark_run=True,
             )
 
+    def cleanup_run_fn(self) -> None:
+        # Authoritative cleanup for one benchmark module load. Higher-level
+        # finally blocks may call this again as safety nets, so keep it idempotent.
+        mod = self._benchmark_module
+        self._benchmark_module = None
+
+        cached_mod = PyCodeCache.modules_no_attr.pop(self.module_path, None)
+        if mod is None:
+            mod = cached_mod
+
+        if mod is not None:
+            kernel = getattr(mod, self.kernel_name, None)
+            release_benchmark_artifacts = getattr(
+                kernel, "release_benchmark_artifacts", None
+            )
+            if release_benchmark_artifacts is not None:
+                release_benchmark_artifacts()
+            release_static_launchers = getattr(
+                kernel, "_release_static_launchers_except", None
+            )
+            if (
+                release_benchmark_artifacts is None
+                and release_static_launchers is not None
+            ):
+                release_static_launchers(None)
+            elif release_benchmark_artifacts is None:
+                for launcher in getattr(kernel, "launchers", ()) or ():
+                    close = getattr(getattr(launcher, "__self__", None), "close", None)
+                    if close is not None:
+                        close()
+                for compile_result in getattr(kernel, "compile_results", ()) or ():
+                    close = getattr(
+                        getattr(compile_result, "kernel", None), "close", None
+                    )
+                    if close is not None:
+                        close()
+
+        PyCodeCache.modules[:] = [
+            module
+            for module in PyCodeCache.modules
+            if getattr(module, "__file__", None) != self.module_path
+        ]
+        PyCodeCache.linemaps.pop(self.module_path, None)
+
     def precompile(self):
-        mod = PyCodeCache.load_by_key_path(self.module_cache_key, self.module_path)
-        getattr(mod, self.kernel_name).precompile()
+        try:
+            mod = PyCodeCache.load_by_key_path(
+                self.module_cache_key,
+                self.module_path,
+                set_sys_modules=False,
+            )
+            self._benchmark_module = mod
+            kernel = getattr(mod, self.kernel_name)
+            kernel.precompile()
+
+            self.n_regs = kernel.launchers[0].n_regs
+        finally:
+            self.cleanup_run_fn()
 
     def __str__(self) -> str:
         return f"{self.kernel_name=}, {self.module_path=}, {self.module_cache_key=}"
@@ -737,11 +813,11 @@ class ExternKernelBenchmarkRequest(BenchmarkRequest):
     def __init__(
         self,
         kernel_name: str,
-        input_tensor_meta: Union[TensorMeta, list[TensorMeta]],
-        output_tensor_meta: Union[TensorMeta, list[TensorMeta]],
+        input_tensor_meta: TensorMeta | list[TensorMeta],
+        output_tensor_meta: TensorMeta | list[TensorMeta],
         extra_args: Iterable[Any],
         callable_path: str,  # Module path to the callable (e.g., "extern_kernels.mm")
-        kwargs: Optional[dict[str, Any]] = None,
+        kwargs: dict[str, Any] | None = None,
         has_out_variant: bool = True,
     ) -> None:
         super().__init__(kernel_name, input_tensor_meta, output_tensor_meta, extra_args)
@@ -760,9 +836,7 @@ def make_run_fn(
             # For non-out variant, just call with inputs
             return functools.partial(fn, *input_tensors)
 
-    def benchmark(
-        self, *input_tensors: torch.Tensor, out: Optional[torch.Tensor] = None
-    ):
+    def benchmark(self, *input_tensors: torch.Tensor, out: torch.Tensor | None = None):
         if out is not None and out.numel() == 0:
             # no need to run the kernel of do benchmarking
             return 0.0
@@ -776,6 +850,10 @@ def benchmark(
                     out_new, tuple(out.size()), tuple(out.stride())
                 )
                 out.copy_(out_new)  # for correctness checking
+            if self.benchmark_with_cudagraphs:
+                return benchmarker.benchmark_gpu_with_cuda_graph(
+                    lambda: algo(*input_tensors)
+                )
             if config.profile_bandwidth_with_do_bench_using_profiling:
                 return do_bench_using_profiling(lambda: algo(*input_tensors))
             return benchmarker.benchmark(algo, input_tensors, {})
@@ -812,6 +890,57 @@ class ExternKernelCPUBenchmarkRequest(
     pass
 
 
+class SubgraphBenchmarkRequest(BenchmarkRequest):
+    """
+    Benchmark request for subgraph choices.
+
+    Pre-compiles the subgraph in the main process and stores
+    the module path/cache key for loading in subprocess.
+    """
+
+    def __init__(
+        self,
+        kernel_name: str,
+        input_tensor_meta: TensorMeta | list[TensorMeta],
+        output_tensor_meta: TensorMeta | list[TensorMeta],
+        extra_args: Iterable[Any],
+        module_path: str,
+        module_cache_key: str,
+        sym_input_values: list[int],
+    ) -> None:
+        super().__init__(kernel_name, input_tensor_meta, output_tensor_meta, extra_args)
+        self.module_path = module_path
+        self.module_cache_key = module_cache_key
+        self.sym_input_values = sym_input_values
+
+    def make_run_fn(
+        self, *input_tensors: torch.Tensor, out: torch.Tensor
+    ) -> Callable[[], None]:
+        mod = PyCodeCache.load_by_key_path(
+            self.module_cache_key,
+            self.module_path,
+            set_sys_modules=False,
+        )
+        sym_input_values = self.sym_input_values
+        # Create a new list each call since mod.call does args.clear()
+        return lambda: mod.call([*sym_input_values, *input_tensors])
+
+    def precompile(self) -> None:
+        # Module is already compiled in main process, no precompilation needed
+        pass
+
+    def __str__(self) -> str:
+        return f"SubgraphBenchmarkRequest({self.kernel_name}, {self.module_path})"
+
+
+class SubgraphGPUBenchmarkRequest(GPUDeviceBenchmarkMixin, SubgraphBenchmarkRequest):
+    pass
+
+
+class SubgraphCPUBenchmarkRequest(CPUDeviceBenchmarkMixin, SubgraphBenchmarkRequest):
+    pass
+
+
 class CUTLASSBenchmarkRequest(GPUDeviceBenchmarkMixin, BenchmarkRequest):
     """
     A class to handle CUDA (CUTLASS) benchmark requests. This class is for
@@ -825,8 +954,8 @@ class CUTLASSBenchmarkRequest(GPUDeviceBenchmarkMixin, BenchmarkRequest):
     def __init__(
         self,
         kernel_name: str,
-        input_tensor_meta: Union[TensorMeta, list[TensorMeta]],
-        output_tensor_meta: Union[TensorMeta, list[TensorMeta]],
+        input_tensor_meta: TensorMeta | list[TensorMeta],
+        output_tensor_meta: TensorMeta | list[TensorMeta],
         extra_args: Iterable[Any],
         source_code: str,
         device_type: str = "cuda",
@@ -834,13 +963,13 @@ def __init__(
         super().__init__(kernel_name, input_tensor_meta, output_tensor_meta, extra_args)
         self.source_code = source_code
         self.workspace_size: int = 0
-        self.workspace: Optional[torch.Tensor] = None
-        self.DLL: Optional[DLLWrapper] = None
+        self.workspace: torch.Tensor | None = None
+        self.DLL: DLLWrapper | None = None
         self._workspace_size_updated = False
         self.hash_key: str = ""
         self.source_file: str = ""
         self.device_type = device_type
-        self.codecache_cls = CUDACodeCache
+        self.codecache_cls = XPUCodeCache if device_type == "xpu" else CUDACodeCache
         self.device_interface = get_interface_for_device(device_type)
         self.hash_key, self.source_file = self.codecache_cls.write(
             self.source_code, "so"
@@ -852,7 +981,7 @@ def precompile(self):
         This may happen in a separate thread pool.
         """
         autotuning_log.debug("Precompiling %s", self)
-        CUDACodeCache.compile(self.source_code, "so")
+        self.codecache_cls.compile(self.source_code, "so")
         autotuning_log.debug("Done precompiling %s", self)
 
     def make_run_fn(
@@ -874,8 +1003,9 @@ def make_run_fn(
             args,
             self.extra_args,
         )
-        current_stream = self.device_interface.current_stream()
-        stream_ptr = c_void_p(current_stream.cuda_stream)  # type: ignore[attr-defined]
+        stream_ptr = c_void_p(
+            self.device_interface.get_raw_stream(self.device_interface.current_device())
+        )
         run_method = getattr(self.DLL, self.kernel_name)
         workspace_ptr = c_void_p(0)
         if self.workspace_size > 0:
@@ -918,8 +1048,9 @@ def update_workspace_size(self) -> None:
             dict.fromkeys(meta.name for meta in self.input_tensor_meta)
         )
         args = [c_void_p(None) for _ in range(unique_input_count + 1)]
-        current_stream = self.device_interface.current_stream()
-        stream_ptr = c_void_p(current_stream.cuda_stream)  # type: ignore[attr-defined]
+        stream_ptr = c_void_p(
+            self.device_interface.get_raw_stream(self.device_interface.current_device())
+        )
 
         run_method = getattr(self.DLL, self.kernel_name)
         # Retrieve workspace_size and initialize workspace.
@@ -952,6 +1083,7 @@ def ensure_dll_loaded(self):
             self.DLL, self.hash_key, self.source_file = self.codecache_cls.load(
                 self.source_code, "so"
             )
+        self.DLL.open()
 
     def cleanup_run_fn(self) -> None:
         if self.DLL is not None:
@@ -962,6 +1094,16 @@ def cleanup_run_fn(self) -> None:
     def __str__(self) -> str:
         return f"{self.kernel_name=}, {self.source_file=}, {self.hash_key=}"
 
+    def __getstate__(self) -> dict[str, Any]:
+        state = self.__dict__.copy()
+        state["DLL"] = None
+        state["workspace"] = None
+        state["_workspace_size_updated"] = False
+        return state
+
+    def __setstate__(self, state: dict[str, Any]) -> None:
+        self.__dict__.update(state)
+
 
 class CppBenchmarkRequest(CPUDeviceBenchmarkMixin, BenchmarkRequest):
     # Important: Instances of this class have to be serializable
@@ -970,15 +1112,15 @@ class CppBenchmarkRequest(CPUDeviceBenchmarkMixin, BenchmarkRequest):
     def __init__(
         self,
         kernel_name: str,
-        input_tensor_meta: Union[TensorMeta, list[TensorMeta]],
-        output_tensor_meta: Union[TensorMeta, list[TensorMeta]],
+        input_tensor_meta: TensorMeta | list[TensorMeta],
+        output_tensor_meta: TensorMeta | list[TensorMeta],
         extra_args: Iterable[Any],
         source_code: str,
     ) -> None:
         super().__init__(kernel_name, input_tensor_meta, output_tensor_meta, extra_args)
         self.source_code = source_code
         self.hash_key = get_hash(source_code)
-        self.DLL: Optional[Union[CDLL, ModuleType]] = None
+        self.DLL: CDLL | ModuleType | None = None
 
     def precompile(self):
         # Prepopulate CppCodeCache
@@ -1024,8 +1166,8 @@ class CuteDSLBenchmarkRequest(GPUDeviceBenchmarkMixin, BenchmarkRequest):
     def __init__(
         self,
         kernel_name: str,
-        input_tensor_meta: Union[TensorMeta, list[TensorMeta]],
-        output_tensor_meta: Union[TensorMeta, list[TensorMeta]],
+        input_tensor_meta: TensorMeta | list[TensorMeta],
+        output_tensor_meta: TensorMeta | list[TensorMeta],
         extra_args: tuple[Any, ...],
         source_code: PartialRender,
     ) -> None:
@@ -1041,7 +1183,11 @@ def make_run_fn(
         Create a function to run the CuteDSL kernel with the given input and output tensors.
         Similar to TritonBenchmarkRequest.make_run_fn but for CuteDSL kernels.
         """
-        mod = PyCodeCache.load_by_key_path(self.module_cache_key, self.module_path)
+        mod = PyCodeCache.load_by_key_path(
+            self.module_cache_key,
+            self.module_path,
+            set_sys_modules=False,
+        )
 
         # Logic replicated async_compile
         from .codegen.cutedsl.cutedsl_kernel import MAIN_SUFFIX
@@ -1086,7 +1232,7 @@ class AutotuneProcessPool:
     in a separate process.
     """
 
-    _instance: Optional[AutotuneProcessPool] = None
+    _instance: AutotuneProcessPool | None = None
     _lock: threading.Lock = threading.Lock()
     _shutdown_for_inactivity: bool = False
 
@@ -1176,7 +1322,7 @@ def warm_up(self) -> Future[Any]:
                     self._warmup_start_time = time.perf_counter()
                     self._warmup_future = self.pool.submit(
                         _init_autotune_subprocess,
-                        allow_tf32=torch.backends.cuda.matmul.allow_tf32,
+                        fp32_precision=torch.backends.cuda.matmul.fp32_precision,
                     )
                     self._warmup_future.add_done_callback(self._on_warmup_complete)
                     autotuning_log.info("Warmup job submitted")
@@ -1196,6 +1342,7 @@ def _on_warmup_complete(self, future: Future[Any]) -> None:
                 warmup_elapsed_time,
                 result,
             )
+            self._record_activity()
         except Exception as e:
             autotuning_log.error(
                 "AutotuneProcessPool warmup failed after %.4f seconds",
@@ -1236,7 +1383,7 @@ def use_pipelined_autotuning() -> bool:
     )
 
 
-def _init_autotune_subprocess(allow_tf32: bool) -> bool:
+def _init_autotune_subprocess(fp32_precision: str) -> bool:
     """
     Warmup function run in the autotune subprocess.
     """
@@ -1246,7 +1393,7 @@ def _init_autotune_subprocess(allow_tf32: bool) -> bool:
     if torch.cuda.is_available():
         torch.zeros(1, device="cuda")
 
-    torch.backends.cuda.matmul.allow_tf32 = allow_tf32
+    torch.backends.cuda.matmul.fp32_precision = fp32_precision
 
     return True
 
@@ -1274,9 +1421,10 @@ def run_autotune_in_subprocess(
         return timing
 
     except Exception:
-        autotuning_log.error(
+        autotuning_log.warning(
             "Failed to benchmark choice %s",
             benchmark_request,
+            exc_info=True,
         )
         # Use infinity for failed benchmarks so they're not selected
         return float("inf")
@@ -1290,7 +1438,7 @@ class PrecompileThreadPool:
     precompilation happens in background threads.
     """
 
-    _instance: Optional[PrecompileThreadPool] = None
+    _instance: PrecompileThreadPool | None = None
     _lock = threading.Lock()
 
     def __init__(self, max_workers: int = 4):
@@ -1307,6 +1455,9 @@ def get_instance(cls) -> PrecompileThreadPool:
         return cls._instance
 
     def submit(self, fn, *args, **kwargs):
+        ctx = contextvars.copy_context()
+        # Need to copy context so workers have access to the correct config settings
+        fn = functools.partial(ctx.run, fn)
         return self._executor.submit(fn, *args, **kwargs)
 
     def _shutdown(self, wait: bool = False):
diff --git a/torch/_inductor/await_utils.py b/torch/_inductor/await_utils.py
index 2468b0039a18d..fcb4549066edd 100644
--- a/torch/_inductor/await_utils.py
+++ b/torch/_inductor/await_utils.py
@@ -5,7 +5,7 @@
 from collections.abc import Awaitable, Callable, Coroutine, Generator, Iterator
 from contextlib import contextmanager, ExitStack
 from contextvars import Context
-from typing import Any, Optional, Protocol, TypeVar
+from typing import Any, Protocol, TypeVar
 
 from torch.utils._ordered_set import OrderedSet
 
@@ -81,7 +81,7 @@ def _restore_running_loop() -> Iterator[None]:
 
 @contextmanager
 def _new_loop(
-    task_factory: Optional[TaskFactoryType] = None,
+    task_factory: TaskFactoryType | None = None,
 ) -> Iterator[asyncio.AbstractEventLoop]:
     loop = asyncio.new_event_loop()
     tasks = _patch_loop(loop)
@@ -133,12 +133,12 @@ def _cancel_all_tasks(
 def _patch_loop(loop: AbstractEventLoop) -> OrderedSet[Future]:  # type: ignore[type-arg]
     tasks: weakref.WeakSet[Future] = weakref.WeakSet()  # type: ignore[type-arg]
 
-    task_factories: list[Optional[TaskFactoryType]] = [None]
+    task_factories: list[TaskFactoryType | None] = [None]
 
-    def _set_task_factory(factory: Optional[TaskFactoryType]) -> None:
+    def _set_task_factory(factory: TaskFactoryType | None) -> None:
         task_factories[0] = factory
 
-    def _get_task_factory() -> Optional[TaskFactoryType]:
+    def _get_task_factory() -> TaskFactoryType | None:
         return task_factories[0]
 
     def _safe_task_factory(
diff --git a/torch/_inductor/bounds.py b/torch/_inductor/bounds.py
index bc8dba5119252..50b0fbaf96c12 100644
--- a/torch/_inductor/bounds.py
+++ b/torch/_inductor/bounds.py
@@ -2,7 +2,7 @@
 import operator
 from collections.abc import Callable
 from functools import partial
-from typing import Any, Optional, Union
+from typing import Any
 
 import sympy
 from sympy import Expr
@@ -36,7 +36,7 @@ class BoundVars:
     """
 
     def __init__(self, loop_body: LoopBody) -> None:
-        def upper_bound(v: Union[Expr, int]) -> int:
+        def upper_bound(v: Expr | int) -> int:
             return bound_sympy(v).upper if isinstance(v, Expr) else v
 
         self.loop_body = loop_body
@@ -195,10 +195,11 @@ def index_expr(cls, index: Any, dtype: torch.dtype) -> ValueRanges[Any]:
         return cls.to_dtype(index, dtype)
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def to_dtype(
         x: Any,
         dtype: torch.dtype,
-        src_dtype: Optional[torch.dtype] = None,
+        src_dtype: torch.dtype | None = None,
         use_compute_types: bool = True,
     ) -> ValueRanges[Any]:
         x = ValueRanges.wrap(x)
@@ -237,10 +238,12 @@ def cast(x: Any, dtype: torch.dtype) -> sympy.Expr:
             return ValueRanges(cast(x.lower, dtype), cast(x.upper, dtype))
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def square(x: Any) -> ValueRanges[Any]:
         return ValueRanges.convex_min_zero_map(x, lambda y: PowByNatural(y, 2))
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def neg(x: Any) -> ValueRanges[Any]:
         return ValueRanges.decreasing_map(x, operator.neg)
 
diff --git a/torch/_inductor/choices.py b/torch/_inductor/choices.py
index bcbb38bcbf0cb..2265ad4befe97 100644
--- a/torch/_inductor/choices.py
+++ b/torch/_inductor/choices.py
@@ -2,7 +2,7 @@
 
 import dataclasses
 import typing
-from typing import Any, Optional, TYPE_CHECKING, Union
+from typing import Any, TYPE_CHECKING
 
 import sympy
 
@@ -58,12 +58,17 @@ class FusionScore:
     template_score: int
     node_type_score: bool
     memory_score: int
+    buffer_overlap_score: int
     proximity_score: int
 
     def __lt__(self, other):
         """
         node_type_score has higher priority than memory_score unless
-        the memory_score differs too much
+        the memory_score differs too much.
+
+        buffer_overlap_score is prioritized below memory_score so that
+        strict global memory savings (exact dep matches) are preferred
+        over buffer overlap scoring (same buffer, different indexing).
         """
         threshold = 16
         if self.template_score != other.template_score:
@@ -75,16 +80,22 @@ def __lt__(self, other):
         ):
             return self.memory_score < other.memory_score
 
-        return (self.node_type_score, self.memory_score, self.proximity_score) < (
+        return (
+            self.node_type_score,
+            self.memory_score,
+            self.buffer_overlap_score,
+            self.proximity_score,
+        ) < (
             other.node_type_score,
             other.memory_score,
+            other.buffer_overlap_score,
             other.proximity_score,
         )
 
 
 class InductorChoices:
     """
-    This class contains a collection of default heuristics that effect performance of our generated
+    This class contains a collection of default heuristics that affect performance of our generated
     code.  We try to not put correctness requirements in this file.
 
     You can override the choices made here by doing:
@@ -93,10 +104,13 @@ class MyHeuristics(InductorChoices):
                 ...
 
             torch._inductor.virtualized.V.set_choices_handler(MyHeuristics())
+
+    Subclasses used with inductor_choices_class must implement uuid() for
+    cache key computation.
     """
 
     def get_config_heuristics(
-        self, device_type: Optional[str] = "cuda"
+        self, device_type: str | None = "cuda"
     ) -> BaseConfigHeuristic:
         if device_type == "cuda":
             if torch.version.hip is None:
@@ -114,27 +128,31 @@ def get_config_heuristics(
 
     # Conv configs
     def get_conv_configs(
-        self, device_type: Optional[str] = "cuda"
+        self, device_type: str | None = "cuda"
     ) -> partial[Generator[TritonConfig, None, None]]:
         conv_heuristics = self.get_config_heuristics(device_type)
         return conv_heuristics.get_conv_configs()
 
+    def get_depthwise_conv_configs(self, device_type: str | None = "cuda") -> list[Any]:
+        heuristics = self.get_config_heuristics(device_type)
+        return heuristics.get_depthwise_conv_configs()
+
     # Flex attention configs
     # TODO(coconutruben): break out flexattention/decode configs into the new retrieval mechanism
     def get_flex_attention_fwd_configs(
-        self, head_dim: int, dtype: torch.dtype, device_type: Optional[str] = "cuda"
+        self, head_dim: int, dtype: torch.dtype, device_type: str | None = "cuda"
     ) -> list[Any]:
         flex_heuristics = self.get_config_heuristics(device_type)
         return flex_heuristics.get_flex_attn_fwd_configs(head_dim, dtype)
 
     def get_flex_attention_bwd_configs(
-        self, head_dim: int, dtype: torch.dtype, device_type: Optional[str] = "cuda"
+        self, head_dim: int, dtype: torch.dtype, device_type: str | None = "cuda"
     ) -> list[Any]:
         flex_heuristics = self.get_config_heuristics(device_type)
         return flex_heuristics.get_flex_attn_bwd_configs(head_dim, dtype)
 
     def get_flex_decode_configs(
-        self, head_dim: int, dtype: torch.dtype, device_type: Optional[str] = "cuda"
+        self, head_dim: int, dtype: torch.dtype, device_type: str | None = "cuda"
     ) -> list[Any]:
         flex_heuristics = self.get_config_heuristics(device_type)
         return flex_heuristics.get_flex_decode_configs(head_dim, dtype)
@@ -143,9 +161,9 @@ def _finalize_template_configs(
         self,
         template_choices: dict[str, Generator[KernelTemplateChoice, None, None]],
         kernel_inputs: KernelInputs,
-        templates: list[Union[KernelTemplate, ExternKernelChoice]],
+        templates: list[KernelTemplate | ExternKernelChoice],
         op_name: str,
-        kwarg_overrides: Optional[dict[str, dict[str, Any]]] = None,
+        kwarg_overrides: dict[str, dict[str, Any]] | None = None,
     ) -> list[KernelTemplateChoice]:
         """
         This method can be subclassed to perform any override/modification of the choices.
@@ -174,9 +192,9 @@ def _finalize_template_configs(
     def get_ktc(
         self,
         kernel_inputs: KernelInputs,
-        template: Union[KernelTemplate, ExternKernelChoice],
+        template: KernelTemplate | ExternKernelChoice,
         op_name: str,
-        kwarg_overrides: Optional[dict[str, Any]] = None,
+        kwarg_overrides: dict[str, Any] | None = None,
     ) -> Generator[KernelTemplateChoice, None, None]:
         """
         Utility to get the KernelTemplateChoice generator for a specific input.
@@ -269,9 +287,9 @@ def _need_to_fix_layout(
     def get_template_configs(
         self,
         kernel_inputs: KernelInputs,
-        templates: list[Union[KernelTemplate, ExternKernelChoice]],
+        templates: list[KernelTemplate | ExternKernelChoice],
         op_name: str,
-        kwarg_overrides: Optional[dict[str, dict[str, Any]]] = None,
+        kwarg_overrides: dict[str, dict[str, Any]] | None = None,
     ) -> list[ChoiceCaller]:
         """
         Get list of ChoiceCallers for MM templates using template-specific heuristics.
@@ -332,6 +350,18 @@ def triton_kernel_kwargs(
         """Hook to change the kwargs passed to TritonKernel, used to apply fixed configurations"""
         return kernel_kwargs
 
+    def override_best_choice(
+        self,
+        best_choice: ChoiceCaller,
+        timings: dict[ChoiceCaller, float],
+    ) -> ChoiceCaller:
+        """Hook to override the autotuning best choice after benchmarking."""
+        return best_choice
+
+    def customize_fused_kernel_name(self, fused_name: str, src_code: str) -> str:
+        """Hook to transform fused kernel names during codegen"""
+        return fused_name
+
     @staticmethod
     def should_use_cooperative_reduction(features: SIMDKernelFeatures) -> bool:
         """Heuristic to decide if a cooperative reduction should be used."""
@@ -429,7 +459,7 @@ def reduction_split_factor(
         so we will do the reduction in two phases."""
         props = DeviceProperties.create(device)
         num_sm = props.multi_processor_count
-        warp_size = props.warp_size if props.warp_size is not None else 32
+        warp_size = props.warp_size_or_default
         max_threads_per_sm = (
             props.max_threads_per_multi_processor
             if props.max_threads_per_multi_processor is not None
@@ -628,8 +658,8 @@ def score_fusion(
         - Fusions closer together in original graph order
         """
 
-        memory_score, is_mix_order_reduction = typing.cast(
-            tuple[int, bool],
+        memory_score, buffer_overlap_score, is_mix_order_reduction = typing.cast(
+            tuple[int, int, bool],
             scheduler.score_fusion_memory(
                 node1, node2, return_is_mix_order_reduction=True
             ),
@@ -654,5 +684,6 @@ def score_fusion(
             template_score,
             type_score,
             memory_score,
+            buffer_overlap_score,
             proximity_score,
         )
diff --git a/torch/_inductor/codecache.py b/torch/_inductor/codecache.py
index 772536ffb34fd..28cae0be378be 100644
--- a/torch/_inductor/codecache.py
+++ b/torch/_inductor/codecache.py
@@ -34,7 +34,7 @@
 from tempfile import _TemporaryFileWrapper
 from time import time, time_ns
 from types import ModuleType
-from typing import Any, cast, Generic, NoReturn, Optional, TYPE_CHECKING, TypeVar, Union
+from typing import Any, cast, Generic, Literal, NoReturn, TYPE_CHECKING, TypeVar
 from typing_extensions import override, Self
 
 import torch
@@ -65,11 +65,13 @@
     _LINKER_SCRIPT,
     _set_gpu_runtime_env,
     _TORCH_PATH,
+    batch_convert_cubins_to_obj,
     convert_cubin_to_obj,
     CppBuilder,
     CppOptions,
     CppTorchDeviceOptions,
     get_compiler_version_info,
+    get_cpp_compiler,
     get_ld_and_objcopy,
     get_name_and_dir_from_output_file_path,
     normalize_path_separator,
@@ -95,6 +97,7 @@
     XPU_KERNEL_FORMAT,
 )
 from torch._library.fake_class_registry import FakeScriptObject
+from torch._library.opaque_object import is_opaque_reference_type
 from torch._logging import trace_structured
 from torch._subclasses.fake_tensor import (
     extract_tensor_metadata,
@@ -110,7 +113,11 @@
 )
 from torch.export.pt2_archive._package_weights import TensorProperties, Weights
 from torch.export.pt2_archive.constants import CUSTOM_OBJ_FILENAME_PREFIX
-from torch.fx.experimental.symbolic_shapes import has_hint, ShapeEnv, size_hint
+from torch.fx.experimental.symbolic_shapes import (
+    guarding_hint_or_throw,
+    has_guarding_hint,
+    ShapeEnv,
+)
 from torch.utils._ordered_set import OrderedSet
 
 from .output_code import CompiledFxGraph
@@ -494,14 +501,19 @@ def __init__(
         self,
         gm: torch.fx.GraphModule,
         has_user_defined_triton_kernels: bool = False,
+        device_id_agnostic: bool = False,
     ) -> None:
         """
         Create an FX graph pickler. If include_non_inlined=True, then pickling will
         include the _values_ for all Tensors. (Note that any tensors are constants
         attached as attributes to the GraphModule). Otherwise, pickling will include
         only the metadata for these tensors.
+
+        If device_id_agnostic=True, device indices in TensorMetadata are normalized
+        to 0, so that the same graph on different GPUs produces identical bytes.
         """
         self._stream = io.BytesIO()
+        self._device_id_agnostic = device_id_agnostic
         super().__init__(self._stream)
 
         self.dispatch_table = copyreg.dispatch_table.copy()
@@ -511,6 +523,7 @@ def __init__(
                 torch.Tensor: functools.partial(self._reduce_tensor),
                 torch.nn.parameter.Parameter: functools.partial(self._reduce_tensor),
                 torch.SymInt: functools.partial(self._reduce_symint),
+                torch.SymBool: functools.partial(self._reduce_symbool),
                 torch.fx.experimental._backward_state.BackwardState: functools.partial(
                     self._reduce_unsupported
                 ),
@@ -534,6 +547,10 @@ def _reduce_fake_tensor(
         Custom reducer to pickle FakeTensors.
         """
         metadata = extract_tensor_metadata_for_cache_key(t)
+        if self._device_id_agnostic:
+            metadata = dataclasses.replace(
+                metadata, device=torch.device(metadata.device.type, 0)
+            )
         return (_ident, (metadata,))
 
     def _reduce_tensor(
@@ -552,6 +569,10 @@ def _reduce_tensor(
             raise BypassFxGraphCache("mkldnn tensors unpickleable")
 
         metadata = extract_tensor_metadata_for_cache_key(t)
+        if self._device_id_agnostic:
+            metadata = dataclasses.replace(
+                metadata, device=torch.device(metadata.device.type, 0)
+            )
 
         # If this is a non-inlined frozen parameter, we consider the metadata only.
         if is_frozen_param(t) and not GraphLowering.can_inline_constant(t):
@@ -579,6 +600,14 @@ def _reduce_symint(self, s: SymInt) -> tuple[Callable[[T], T], tuple[str]]:
         # entity with SymInt args is safe to reuse.
         return (_ident, (str(s),))
 
+    def _reduce_symbool(self, s: torch.SymBool) -> tuple[Callable[[T], T], tuple[str]]:
+        """
+        Custom reducer to pickle SymBools.
+        """
+        # Same approach as _reduce_symint: use the string representation for
+        # hashing.  Guards ensure correctness on cache reload.
+        return (_ident, (str(s),))
+
     def _reduce_unsupported(self, s: Any) -> NoReturn:
         """
         Custom reducer to handle any objects that we don't support and therefore
@@ -781,14 +810,71 @@ class BypassFxGraphCache(Exception):
     """
 
 
+_warned_pre_grad_pass_missing_uuid: OrderedSet[str] = OrderedSet()
+
+
+def resolve_pre_grad_pass_timing() -> Literal["early", "late"]:
+    """Resolve the effective pre-grad pass timing from the config.
+
+    "default" is resolved based on whether the custom pass provides a UUID:
+    passes with a UUID (or no custom pass) run "late" (after cache lookup),
+    passes without a UUID run "early" (before cache lookup).
+
+    Raises RuntimeError if a custom pass without a UUID is explicitly set to
+    run "late", since the cache key cannot account for it.
+    """
+    timing: Literal["early", "late", "default"] = config.pre_grad_pass_timing
+    custom_pass = config.pre_grad_custom_pass
+    has_uuid = (
+        custom_pass
+        and isinstance(custom_pass, CustomGraphPass)
+        and custom_pass.uuid() is not None
+    )
+
+    if timing == "default":
+        supports_late = custom_pass is None or has_uuid
+        timing = "late" if supports_late else "early"
+        if timing == "early" and custom_pass:
+            pass_name = type(custom_pass).__qualname__
+            if pass_name not in _warned_pre_grad_pass_missing_uuid:
+                _warned_pre_grad_pass_missing_uuid.add(pass_name)
+                log.warning(
+                    "pre_grad_custom_pass %s does not implement uuid(); "
+                    "falling back to early timing (pre-grad pass cache will be bypassed). "
+                    "Implement uuid() on your CustomGraphPass to enable caching.",
+                    pass_name,
+                )
+                CompileEventLogger.try_add_pt2_compile(
+                    "backend_compile",
+                    pre_grad_pass_missing_uuid=True,
+                    pre_grad_pass_name=pass_name,
+                )
+
+    if timing == "late" and custom_pass and not has_uuid:
+        raise RuntimeError(
+            "pre_grad_custom_pass must implement uuid() to run late "
+            "(after cache lookup). Either implement uuid() or set "
+            "pre_grad_pass_timing to 'early'."
+        )
+
+    return timing
+
+
+@dataclasses.dataclass
+class HashableOpaqueValue:
+    ordinal: int
+
+
 class FxGraphHashDetails:
     """
     Object to capture all the details for a compiled FX graph relevant to computing
     a safe and stable cache key.
     """
 
-    # Excluded kwargs param that are not stable between runs
-    EXCLUDED_KWARGS = ["graph_id"]
+    # Excluded kwargs param that are not stable between runs or that
+    # don't affect compiled output (like compile_region_name which is
+    # just a debug label).
+    EXCLUDED_KWARGS = ["graph_id", "compile_region_name"]
 
     def __init__(
         self,
@@ -798,7 +884,19 @@ def __init__(
         inputs_to_check: Sequence[int],
     ) -> None:
         self.gm = gm
-        self.example_inputs = example_inputs
+        # Replace opaque references with hashable ordinals. What's important
+        # is that if the same reference appears twice then it's the same hash
+        # value for each.
+        processed_inputs: list[InputType | HashableOpaqueValue] = []
+        seen_opaques: dict[int, HashableOpaqueValue] = {}
+        for inp in example_inputs:
+            if is_opaque_reference_type(type(inp)):
+                if id(inp) not in seen_opaques:
+                    seen_opaques[id(inp)] = HashableOpaqueValue(len(seen_opaques))
+                processed_inputs.append(seen_opaques[id(inp)])
+            else:
+                processed_inputs.append(inp)
+        self.example_inputs = processed_inputs
         self.cache_key_tag = cconfig.cache_key_tag
 
         # Order kwargs so hashing is stable to changes in kwarg order. Although
@@ -881,6 +979,12 @@ def __init__(
             torch.utils.deterministic.fill_uninitialized_memory,  # type: ignore[attr-defined]
         )
 
+        # Provenance tracking level affects whether provenance data is stored
+        # in the CompiledFxGraph, so it must be part of the cache key.
+        # Note: the "trace" prefix is excluded from _cache_config_ignore_prefix,
+        # so we add this explicitly.
+        self.provenance_tracking_level = config.trace.provenance_tracking_level
+
         # Global settings affecting matmul codegen.
         self.cuda_matmul_settings = (
             torch.backends.cuda.matmul.fp32_precision,
@@ -888,11 +992,31 @@ def __init__(
             torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction,
         )
 
+        # Include cudagraph annotation in cache key only when it changes
+        # behavior. When both fwd and bwd are overridden to the same value,
+        # normalize to a simple boolean (equivalent to flipping the config).
+        # When fwd and bwd differ, include the full annotation.
+        if gm is not None:
+            annotation = gm.meta.get("cudagraph_annotation")
+            if annotation is not None:
+                default = config.triton.cudagraphs
+                if annotation.fwd == annotation.bwd and annotation.fwd is not None:
+                    if annotation.fwd != default:
+                        self.cudagraph_override = annotation.fwd
+                elif (annotation.fwd is not None and annotation.fwd != default) or (
+                    annotation.bwd is not None and annotation.bwd != default
+                ):
+                    self.cudagraph_annotation = annotation
+
         # Also hash on various system info (including the triton compiler version).
         self.torch_version = torch_key()
         self.system_info = CacheBase.get_system()
         self.inductor_config = config.save_config_portable(ignore_private_configs=False)
-        # Custom post grad passes should provide an ID to hash.
+        # Custom passes should provide an ID to hash when they run late (after cache lookup).
+        if resolve_pre_grad_pass_timing() != "early":
+            self.pre_grad_custom_pass = self._get_custom_pass_detail(
+                config.pre_grad_custom_pass
+            )
         self.post_grad_custom_pre_pass = self._get_custom_pass_detail(
             config.post_grad_custom_pre_pass
         )
@@ -1178,7 +1302,9 @@ def _filter_backed_symints(
         Get the backed SymInt objects from the input list. Note that we can never
         have guards that depend on unbacked symint.
         """
-        return [s for s in inputs if isinstance(s, torch.SymInt) and has_hint(s)]
+        return [
+            s for s in inputs if isinstance(s, torch.SymInt) and has_guarding_hint(s)
+        ]
 
     @classmethod
     def _get_shape_env(cls: type[GuardedCache[T]]) -> ShapeEnv | None:
@@ -1415,7 +1541,7 @@ def _lookup_graph(
         assert shape_env is not None
 
         symints = FxGraphCache._filter_backed_symints(example_inputs)
-        hints = [size_hint(s) for s in symints]
+        hints = [guarding_hint_or_throw(s) for s in symints]
 
         # If this config is turned on, everything is a guard hit and we check nothing
         if config.unsafe_skip_cache_dynamic_shape_guards:
@@ -1436,6 +1562,17 @@ def _lookup_graph(
         if graph is None:
             return None, cache_info
 
+        # Validate extern_libs (e.g. libdevice) match the current env.
+        if graph.extern_libs_key is not None:
+            try:
+                backend = torch.utils._triton.triton_backend()
+                current = torch.utils._triton._extern_libs_key(backend)
+            except Exception:
+                current = None
+            if current != graph.extern_libs_key:
+                cache_info["cache_status_detailed"] = "guard_miss"
+                return None, cache_info
+
         if pickled_content is not None:
             CacheArtifactManager.record_artifact(
                 InductorCacheArtifact.type(), key, pickled_content
@@ -1492,6 +1629,13 @@ def _save_graph(
         compiled_graph.guards_expr = shape_env.produce_guards_expression(
             placeholders=symints, guards=guards
         )
+        try:
+            backend = torch.utils._triton.triton_backend()
+            compiled_graph.extern_libs_key = torch.utils._triton._extern_libs_key(
+                backend
+            )
+        except Exception:
+            pass
         disk_compiled_graph = copy(compiled_graph)
         disk_compiled_graph.prepare_for_serialization()
 
@@ -1535,6 +1679,11 @@ def _check_for_hop(gm: torch.fx.GraphModule) -> None:
                     raise BypassFxGraphCache(
                         f"Can't cache HigherOrderOperator: {node.target.name()}"
                     )
+                # TODO: this check is broken in two ways:
+                # 1. FX uses "get_attr" (with underscore), not "getattr"
+                # 2. It only checks for ScriptObject, not FakeScriptObject
+                # Fixing it would also bypass AOTAutogradCache (which calls
+                # _check_can_cache), so we'd need to decouple the two first.
                 if node.op == "getattr" and isinstance(
                     getattr(gm, node.target), torch._C.ScriptObject
                 ):
@@ -1546,8 +1695,15 @@ def _check_can_cache(gm: torch.fx.GraphModule) -> None:
         Check some conditions that would preclude caching and raise BypassFxGraphCache
         to bypass in case caching is not possible.
         """
-        # Post grad custom passes must implement the CustomGraphPass or we don't
+        # Custom passes must implement the CustomGraphPass or we don't
         # know how to include them in the cache key calculation.
+        # When timing is EARLY, pre-grad passes already ran before the cache
+        # lookup so there's nothing to validate here.
+        if resolve_pre_grad_pass_timing() != "early":
+            assert not config.pre_grad_custom_pass or (
+                isinstance(config.pre_grad_custom_pass, CustomGraphPass)
+                and config.pre_grad_custom_pass.uuid()
+            ), "Unsupported pre grad custom pass"
         for p in (config.post_grad_custom_pre_pass, config.post_grad_custom_post_pass):
             if p and (not isinstance(p, CustomGraphPass) or not p.uuid()):
                 raise BypassFxGraphCache("Unsupported post grad custom pass")
@@ -1824,7 +1980,7 @@ def compile(
         *,
         device_type: str,
         additional_files: list[str],
-    ) -> list[Union[str, Weights]] | str:
+    ) -> list[str | Weights] | str:
         """
         Returns the .so path, or returns a list of files that were generated if
         config.aot_inductor.package=True.
@@ -2443,9 +2599,13 @@ def _pad_to_alignment(raw_bytes: bytes) -> bytes:
                     f.write(json.dumps(qual_name_to_id))
                 generated_files.append(constants_config_json)
 
-            gpu_codecache: ROCmCodeCache | CUDACodeCache = (
-                ROCmCodeCache() if torch.version.hip else CUDACodeCache()
-            )
+            cache_cls = {
+                "rocm": ROCmCodeCache,
+                "cuda": CUDACodeCache,
+                "xpu": XPUCodeCache,
+            }.get("rocm" if torch.version.hip else device_type, CUDACodeCache)
+
+            gpu_codecache = cache_cls()
             gpu_kernels_o = gpu_codecache.aot_kernels_o.copy()
             # clear the list of aot kernels after each linking
             gpu_codecache.aot_kernels_o.clear()
@@ -2457,7 +2617,9 @@ def _pad_to_alignment(raw_bytes: bytes) -> bytes:
 
             cubins_o = []
             asm_files = []
+            fatbin_cmds: list[tuple[str, str]] = []
             if not _IS_WINDOWS:
+                cubins_to_embed: list[tuple[str, str]] = []
                 ld, objcopy = get_ld_and_objcopy(use_relative_path)
                 kernels = getattr(V.graph.wrapper_code, "_kernel_name_to_body", {})
                 for kernel_name, value in CudaKernelParamCache.cache.items():
@@ -2475,30 +2637,7 @@ def _pad_to_alignment(raw_bytes: bytes) -> bytes:
                         and device_type == "cuda"
                     ):
                         if torch.version.hip is None:
-                            current_arch = (
-                                cuda_compile_utils._nvcc_arch_as_compile_option()
-                            )
-                            cmd = (
-                                # pyrefly: ignore [unbound-name]
-                                f"{cuda_compile_utils._cuda_compiler()} -fatbin {asm_file} -o {cubin_file} "
-                                # Triton only allows generating PTX version as same as the current arch
-                                f"-gencode arch=compute_{current_arch},code=compute_{current_arch} "
-                                # Include SASS for the current specific arch
-                                f"-gencode arch=compute_{current_arch},code=sm_{current_arch} "
-                            )
-                            try:
-                                subprocess.run(
-                                    cmd.split(),
-                                    capture_output=True,
-                                    text=True,
-                                    check=True,
-                                )
-                            except subprocess.CalledProcessError as e:
-                                print(
-                                    f"{cmd} failed with:\nstdout:\n{e.stdout}\nstderr:\n{e.stderr}",
-                                    file=sys.stderr,
-                                )
-                                raise
+                            fatbin_cmds.append((asm_file, cubin_file))
 
                         else:
                             # ROCm multi-arch: compile LLVM IR to multi-arch bundle
@@ -2532,10 +2671,59 @@ def _pad_to_alignment(raw_bytes: bytes) -> bytes:
                             log.info("Created multi-arch bundle: %s", cubin_file)
 
                     if config.aot_inductor.embed_kernel_binary:
-                        # Embed cubin files into model.so using objcopy
-                        cubins_o.append(
-                            convert_cubin_to_obj(cubin_file, kernel_name, ld, objcopy)
+                        cubins_to_embed.append((cubin_file, kernel_name))
+
+                # Compile PTX → fatbin in parallel (each nvcc call is independent).
+                # Must complete before cubin embedding below.
+                if fatbin_cmds:
+                    from concurrent.futures import ThreadPoolExecutor
+
+                    current_arch = cuda_compile_utils._nvcc_arch_as_compile_option()
+                    nvcc = cuda_compile_utils._cuda_compiler()
+
+                    def _compile_fatbin(asm_and_cubin: tuple[str, str]) -> None:
+                        asm_f, cubin_f = asm_and_cubin
+                        cmd = (
+                            f"{nvcc} -fatbin {asm_f} -o {cubin_f} "
+                            f"-gencode arch=compute_{current_arch},code=compute_{current_arch} "
+                            f"-gencode arch=compute_{current_arch},code=sm_{current_arch} "
                         )
+                        try:
+                            subprocess.run(
+                                cmd.split(), capture_output=True, text=True, check=True
+                            )
+                        except subprocess.CalledProcessError as e:
+                            print(
+                                f"{cmd} failed with:\nstdout:\n{e.stdout}\nstderr:\n{e.stderr}",
+                                file=sys.stderr,
+                            )
+                            raise
+
+                    with ThreadPoolExecutor() as pool:
+                        list(pool.map(_compile_fatbin, fatbin_cmds))
+
+                if cubins_to_embed:
+                    # Batch all cubins into a single .o using .incbin assembly.
+                    # This replaces N * 3 subprocess calls (ld + 2x objcopy per
+                    # cubin) with a single compiler invocation.
+                    try:
+                        combined_obj = batch_convert_cubins_to_obj(
+                            cubins_to_embed,
+                            os.path.dirname(output_so),
+                            cpp_compiler=get_cpp_compiler(),
+                        )
+                        cubins_o.append(combined_obj)
+                    except subprocess.CalledProcessError:
+                        log.warning(
+                            "Batched cubin embedding failed, "
+                            "falling back to per-cubin objcopy"
+                        )
+                        for cubin_file, kernel_name in cubins_to_embed:
+                            cubins_o.append(
+                                convert_cubin_to_obj(
+                                    cubin_file, kernel_name, ld, objcopy
+                                )
+                            )
 
             output_name, output_dir = get_name_and_dir_from_output_file_path(output_so)
             so_build_options = CppTorchDeviceOptions(
@@ -2545,6 +2733,26 @@ def _pad_to_alignment(raw_bytes: bytes) -> bytes:
                 use_relative_path=use_relative_path,
             )
 
+            if gpu_kernels_o and device_type == "xpu":
+                so_build_options = CppTorchDeviceOptions(
+                    compiler="icpx",
+                    vec_isa=picked_vec_isa,
+                    device_type=device_type,
+                    aot_mode=graph.aot_mode,
+                    use_relative_path=use_relative_path,
+                    extra_flags=[
+                        "-fsycl",
+                        "-fsycl-targets=intel_gpu_pvc",
+                        "-Xspirv-translator",
+                        (
+                            "-spirv-ext="
+                            "+SPV_INTEL_split_barrier,"
+                            "+SPV_INTEL_2d_block_io,"
+                            "+SPV_INTEL_subgroup_matrix_multiply_accumulate"
+                        ),
+                    ],
+                )
+
             obj_srcs = [wrapper_o, kernel_o, consts_o, *gpu_kernels_o, *cubins_o]
             so_builder = CppBuilder(
                 name=output_name,
@@ -2904,9 +3112,12 @@ def load_async(
         # the optimized_code argument is present at all, since that's how the user of
         # this function opts in, but we do compilation and linking in one step if the
         # optimized_code argument is empty (as a micro-optimization).
+        # On GPU the C++ wrapper is just glue — the real kernels are compiled
+        # separately by Triton/CUDA.  Always use -O1 to cut compile time.
+        min_optimize = optimized_code is not None or device_type != "cpu"
         main_build_option = CppTorchDeviceOptions(
             compile_only=bool(optimized_code),
-            min_optimize=optimized_code is not None,
+            min_optimize=min_optimize,
             # pyrefly: ignore [bad-argument-type]
             **compile_command,
         )
@@ -2954,7 +3165,7 @@ def get_hashable_command_line(build_option: BuildOptionsBase) -> str:
                     main_build_option.precompiled_header = _precompile_header(
                         header,
                         main_cmd_line,
-                        min_optimize=optimized_code is not None,
+                        min_optimize=min_optimize,
                         **compile_command,
                     )
 
@@ -3054,7 +3265,15 @@ def _worker_compile_cpp(
 @clear_on_fresh_cache
 class CppPythonBindingsCodeCache(CppCodeCache):
     cache: dict[str, Callable[[], CDLL | ModuleType]] = {}
-    cache_clear = staticmethod(cache.clear)
+    _loaded_module_names: OrderedSet[str] = OrderedSet()
+
+    @staticmethod
+    def cache_clear() -> None:
+        CppPythonBindingsCodeCache.cache.clear()
+        for name in CppPythonBindingsCodeCache._loaded_module_names:
+            sys.modules.pop(name, None)
+        CppPythonBindingsCodeCache._loaded_module_names.clear()
+
     cpp_compile_command_flags = {
         # kernels have no dependency on libtorch
         "include_pytorch": False,
@@ -3072,18 +3291,6 @@ class CppPythonBindingsCodeCache(CppCodeCache):
         #include <cstdlib>
         #include <cerrno>
 
-        #ifndef _MSC_VER
-        #if __cplusplus < 202002L
-        // C++20 (earlier) code
-        // https://en.cppreference.com/w/cpp/language/attributes/likely
-        #define likely(x)       __builtin_expect(!!(x), 1)
-        #define unlikely(x)     __builtin_expect(!!(x), 0)
-        #endif
-        #else
-        #define likely(x) (x)
-        #define unlikely(x) (x)
-        #endif
-
         // This is defined in guards.cpp so we don't need to import PyTorch headers that are slooow.
         // We manually link it below to workaround issues with fbcode build.
         static void* (*_torchinductor_pyobject_tensor_data_ptr)(PyObject* obj);
@@ -3094,19 +3301,19 @@ class CppPythonBindingsCodeCache(CppCodeCache):
         }}
         template <> inline int64_t parse_arg<int64_t>(PyObject* args, size_t n) {{
             auto result = PyLong_AsSsize_t(PyTuple_GET_ITEM(args, n));
-            if(unlikely(result == -1 && PyErr_Occurred()))
+            if(result == -1 && PyErr_Occurred()) [[unlikely]]
                 throw std::runtime_error("expected int arg");
             return result;
         }}
         template <> inline uintptr_t parse_arg<uintptr_t>(PyObject* args, size_t n) {{
             auto result = PyLong_AsVoidPtr(PyTuple_GET_ITEM(args, n));
-            if(unlikely(result == reinterpret_cast<void*>(-1) && PyErr_Occurred()))
+            if(result == reinterpret_cast<void*>(-1) && PyErr_Occurred()) [[unlikely]]
                 throw std::runtime_error("expected int arg");
             return reinterpret_cast<uintptr_t>(result);
         }}
         template <> inline float parse_arg<float>(PyObject* args, size_t n) {{
             auto result = PyFloat_AsDouble(PyTuple_GET_ITEM(args, n));
-            if(unlikely(result == -1.0 && PyErr_Occurred()))
+            if(result == -1.0 && PyErr_Occurred()) [[unlikely]]
                 throw std::runtime_error("expected float arg");
             return static_cast<float>(result);
         }}
@@ -3115,9 +3322,9 @@ class CppPythonBindingsCodeCache(CppCodeCache):
 
         static PyObject* {entry_func}_py(PyObject* self, PyObject* args) {{
             try {{
-                if(unlikely(!PyTuple_CheckExact(args)))
+                if(!PyTuple_CheckExact(args)) [[unlikely]]
                     throw std::runtime_error("tuple args required");
-                if(unlikely(PyTuple_GET_SIZE(args) != {arg_len}))
+                if(PyTuple_GET_SIZE(args) != {arg_len}) [[unlikely]]
                     throw std::runtime_error("requires {arg_len} args");
                 {call_entry_func}
             }} catch(std::exception const& e) {{
@@ -3179,6 +3386,7 @@ def _load_library_inner(cls, path: str, key: str) -> ModuleType:
         assert spec is not None
         module = importlib.util.module_from_spec(spec)
         sys.modules[module_name] = module
+        CppPythonBindingsCodeCache._loaded_module_names.add(module_name)
         assert spec.loader is not None
         spec.loader.exec_module(module)
         return module
@@ -3248,7 +3456,11 @@ def load_pybinding(cls, *args: Any, **kwargs: Any) -> Any:
 @clear_on_fresh_cache
 class CppWrapperCodeCache(CppPythonBindingsCodeCache):
     cache: dict[str, Callable[[], CDLL | ModuleType]] = {}
-    cache_clear = staticmethod(cache.clear)
+
+    @staticmethod
+    def cache_clear() -> None:
+        CppWrapperCodeCache.cache.clear()
+
     cpp_compile_command_flags = {
         "include_pytorch": True,
         "shared": True,
@@ -3724,6 +3936,8 @@ def touch(filename: str) -> None:
 
 @clear_on_fresh_cache
 class PyCodeCache:
+    """Caches generated Python modules and their source mappings."""
+
     # Track the loaded modules so we can remove the on-disk artifacts when
     # clearing the cache. Note also that we may load the same path more
     # than once, but attach different attributes, i.e., due to different
@@ -3741,9 +3955,15 @@ def write(cls, source_code: str, extra: str = "") -> tuple[str, str]:
         return write(source_code, "py", extra=extra)
 
     @classmethod
-    def load(cls, source_code: str, extra: str = "") -> ModuleType:
+    def load(
+        cls,
+        source_code: str,
+        extra: str = "",
+        *,
+        set_sys_modules: bool | None = None,
+    ) -> ModuleType:
         key, path = write(source_code, "py", extra=extra)
-        return cls.load_by_key_path(key, path)
+        return cls.load_by_key_path(key, path, set_sys_modules=set_sys_modules)
 
     @classmethod
     def load_by_key_path(
@@ -3752,19 +3972,26 @@ def load_by_key_path(
         path: str,
         linemap: list[tuple[int, str]] | None = None,
         attrs: dict[str, Any] | None = None,
+        *,
+        set_sys_modules: bool | None = None,
     ) -> ModuleType:
         if linemap is None:
             linemap = []
 
+        in_toplevel = in_toplevel_process()
+        set_sys_modules = in_toplevel if set_sys_modules is None else set_sys_modules
+
         # we only cache when attrs is None
         if attrs is None and path in cls.modules_no_attr:
-            return cls.modules_no_attr[path]
+            mod = cls.modules_no_attr[path]
+            if set_sys_modules:
+                sys.modules.setdefault(mod.__name__, mod)
+            return mod
 
-        in_toplevel = in_toplevel_process()
-        mod = _reload_python_module(key, path, set_sys_modules=in_toplevel)
+        mod = _reload_python_module(key, path, set_sys_modules=set_sys_modules)
 
         # unzip into separate lines/nodes lists
-        if in_toplevel:
+        if set_sys_modules:
             cls.linemaps[path] = list(zip(*linemap))
 
         if attrs is not None:
@@ -3862,8 +4089,12 @@ def __init__(
     ) -> None:
         self.lib_path = lib_path
         self.is_open = False
-        self.DLL = cdll.LoadLibrary(lib_path)
-        self.is_open = True
+        self.open()
+
+    def open(self) -> None:
+        if not self.is_open:
+            self.DLL = cdll.LoadLibrary(self.lib_path)
+            self.is_open = True
 
     def close(self) -> None:
         if self.is_open:
@@ -3872,7 +4103,16 @@ def close(self) -> None:
 
     def _dlclose(self) -> None:
         f_dlclose = None
+        # During Python interpreter shutdown, importing modules or calling
+        # dlclose is unsafe. Silently skip cleanup in that case.
+        try:
+            import sys
 
+            if sys.is_finalizing():
+                return
+        except Exception:
+            # import machinery may already be torn down
+            return
         if is_linux():
             syms = CDLL(None)
             if not hasattr(syms, "dlclose"):
@@ -3960,6 +4200,7 @@ class CacheEntry:
     def cache_clear(cls) -> None:
         cls.cache.clear()
         cls.aot_kernels_o.clear()
+        cls.write.cache_clear()
 
     @staticmethod
     @lru_cache(maxsize=4)
@@ -4004,7 +4245,7 @@ def _compile_command(
         src_files: list[str],
         dst_file: str,
         dst_file_ext: str,
-        extra_args: Optional[list[str]] = None,
+        extra_args: list[str] | None = None,
     ) -> str:
         raise NotImplementedError
 
@@ -4221,7 +4462,7 @@ def _compile_command(
         src_files: list[str],
         dst_file: str,
         dst_file_ext: str,
-        extra_args: Optional[list[str]] = None,
+        extra_args: list[str] | None = None,
     ) -> str:
         return cuda_compile_utils.cuda_compile_command(
             src_files, dst_file, dst_file_ext, extra_args=extra_args
@@ -4245,6 +4486,63 @@ def _source_code_extra(cls) -> str:
         return extra
 
 
+from torch._inductor.codegen.xpu import compile_utils as xpu_compile_utils
+
+
+@clear_on_fresh_cache
+class XPUCodeCache(CUTLASSCodeCache):
+    _SOURCE_CODE_SUFFIX = "cpp"
+    _BACKEND = "XPU"
+    dll_cache: dict[str, DLLWrapper] = {}
+
+    @classmethod
+    def _use_re_build(cls) -> bool:
+        return False
+
+    @classmethod
+    def _compile_command(
+        cls,
+        src_files: list[str],
+        dst_file: str,
+        dst_file_ext: str,
+        extra_args: list[str] | None = None,
+    ) -> str:
+        return xpu_compile_utils.xpu_compile_command(
+            src_files, dst_file, dst_file_ext, extra_args=extra_args
+        )
+
+    @classmethod
+    def _source_code_extra(cls) -> str:
+        extra = repr(
+            [
+                xpu_compile_utils._sycl_compiler(),
+                xpu_compile_utils._sycl_compiler_options(),
+                cutlass_key(),
+            ]
+        )
+        return extra
+
+    @classmethod
+    def load(cls, source_code: str, dst_file_ext: str) -> tuple[DLLWrapper, str, str]:
+        """
+        Compiles source code and loads the generated .so file.
+        Returns a tuple of DLLWrapper, hash_key, source_code_path
+        """
+
+        if dst_file_ext != "so":
+            raise RuntimeError(
+                f"Only support loading a .so file for now. "
+                f"Requested file extension: {dst_file_ext}. Source code: {source_code}"
+            )
+        dst_file_path, hash_key, source_code_path = cls.compile(
+            source_code, dst_file_ext
+        )
+        if dst_file_path not in cls.dll_cache:
+            cls.dll_cache[dst_file_path] = DLLWrapper(dst_file_path)
+
+        return (cls.dll_cache[dst_file_path], hash_key, source_code_path)
+
+
 @clear_on_fresh_cache
 class ROCmCodeCache:
     @dataclasses.dataclass
diff --git a/torch/_inductor/codegen/block_analysis.py b/torch/_inductor/codegen/block_analysis.py
index 5b1e774958823..a62c219b2c5a0 100644
--- a/torch/_inductor/codegen/block_analysis.py
+++ b/torch/_inductor/codegen/block_analysis.py
@@ -1,7 +1,6 @@
 import collections
 import functools
 import textwrap
-from typing import Optional
 
 import sympy
 from sympy import Expr, Symbol
@@ -91,7 +90,7 @@ def match_mod_div_block_expr(
         index_var: Symbol,
         numel: Expr,
         num_dims: int,
-    ) -> Optional[tuple[list[Expr], list[Expr], list[Expr]]]:
+    ) -> tuple[list[Expr], list[Expr], list[Expr]] | None:
         """
         Matches modular indexing expressions, converting them to implied block dimensions and strides.
         See triton.py for more information.
@@ -205,7 +204,7 @@ def match_affine_block_expr(
         cls,
         index: Expr,
         index_var: Symbol,
-    ) -> Optional[Expr]:
+    ) -> Expr | None:
         """
         Matches simple expressions of the form stride * index, returning the
         stride.
diff --git a/torch/_inductor/codegen/common.py b/torch/_inductor/codegen/common.py
index c482d966eb20e..16b7b03ecc6aa 100644
--- a/torch/_inductor/codegen/common.py
+++ b/torch/_inductor/codegen/common.py
@@ -15,16 +15,7 @@
 from abc import ABC, abstractmethod
 from enum import auto, Enum
 from itertools import chain
-from typing import (
-    Any,
-    cast,
-    ClassVar,
-    Generic,
-    NamedTuple,
-    Optional,
-    TYPE_CHECKING,
-    Union,
-)
+from typing import Any, cast, ClassVar, Generic, NamedTuple, TYPE_CHECKING
 from typing_extensions import Self, TypeVar
 
 import sympy
@@ -82,9 +73,9 @@
     from .wrapper import PythonWrapperCodegen
 
     _T = TypeVar("_T")
-    SchedulingConstructor = Callable[[Optional[Scheduler]], BaseScheduling]
+    SchedulingConstructor = Callable[[Scheduler | None], BaseScheduling]
     WrapperConstructor = type[PythonWrapperCodegen]
-    SymbolLike = Union[str, sympy.Symbol]
+    SymbolLike = str | sympy.Symbol
 
     # OpVarT should really be Union[CSEVariable, str], however this
     # causes typing errors in subclasses (defined in other files).
@@ -161,7 +152,7 @@ def get_name(self) -> str:
         pass
 
     @abstractmethod
-    def get_example(self) -> Union[torch.Tensor, sympy.Symbol]:
+    def get_example(self) -> torch.Tensor | sympy.Symbol:
         pass
 
 
@@ -229,7 +220,7 @@ def get_device(self) -> torch.device:
     def get_dtype(self) -> torch.dtype:
         return self.dtype
 
-    def get_example(self) -> Union[torch.Tensor, sympy.Symbol]:
+    def get_example(self) -> torch.Tensor | sympy.Symbol:
         return self.get_layout().get_example()
 
     def get_layout(self) -> FixedLayout:
@@ -284,7 +275,7 @@ class TensorArg:
     buffer: str
     dtype: torch.dtype
     offset: sympy.Expr = sympy.S.Zero  # c++ only
-    alias_of: Optional[str] = None  # halide only
+    alias_of: str | None = None  # halide only
 
 
 @dataclasses.dataclass
@@ -293,7 +284,7 @@ class SizeArg:
     expr: sympy.Expr
 
     @property
-    def alias_of(self) -> Optional[str]:
+    def alias_of(self) -> str | None:
         return None
 
 
@@ -306,19 +297,19 @@ class ConstexprArg:
 class TMADescriptorArg:
     name: str
     api_type: str  # "experimental" or "stable"
-    block_shape: Optional[list[sympy.Expr]]  # only needed for "stable"
-    dtype: Optional[torch.dtype]  # only needed for "stable"
+    block_shape: list[sympy.Expr] | None  # only needed for "stable"
+    dtype: torch.dtype | None  # only needed for "stable"
 
 
 @dataclasses.dataclass
 class DeviceCodegen:
     scheduling: SchedulingConstructor
     wrapper_codegen: WrapperConstructor
-    cpp_wrapper_codegen: Optional[WrapperConstructor] = None
-    fx_wrapper_codegen: Optional[WrapperConstructor] = None
+    cpp_wrapper_codegen: WrapperConstructor | None = None
+    fx_wrapper_codegen: WrapperConstructor | None = None
 
 
-KernelArgType = Union[WorkspaceArg, TensorArg, SizeArg, TMADescriptorArg, ConstexprArg]
+KernelArgType = WorkspaceArg | TensorArg | SizeArg | TMADescriptorArg | ConstexprArg
 
 device_codegens: dict[str, DeviceCodegen] = {}
 
@@ -373,15 +364,16 @@ def tma_descriptor_helpers(self) -> str:
         raise NotImplementedError
 
     def cpp_scratch(
-        self, idx: int, workspace: TritonScratchWorkspace, prefix: Optional[str] = None
-    ) -> Optional[tuple[list[str], str]]:
+        self, idx: int, workspace: TritonScratchWorkspace, prefix: str | None = None
+    ) -> tuple[list[str], str] | None:
         # optionally return (scratch definition, arg name)
         raise NotImplementedError
 
 
 device_op_overrides_dict: dict[str, DeviceOpOverrides] = {}
-custom_backend_passes: dict[str, Optional[CustomGraphModulePass]] = {}
-custom_backend_codegen_configs: dict[str, Optional[ConfigModule]] = {}
+_device_op_overrides_initialized = False
+custom_backend_passes: dict[str, CustomGraphModulePass | None] = {}
+custom_backend_codegen_configs: dict[str, ConfigModule | None] = {}
 
 
 # The code generated by Inductor consists of two main parts: kernel code and wrapper code.
@@ -409,10 +401,10 @@ def register_backend_for_device(
     device: str,
     device_scheduling: SchedulingConstructor,
     device_wrapper_codegen: WrapperConstructor,
-    device_cpp_wrapper_codegen: Optional[WrapperConstructor] = None,
-    device_fx_wrapper_codegen: Optional[WrapperConstructor] = None,
-    device_custom_pass: Optional[CustomGraphModulePass] = None,
-    device_custom_config: Optional[ConfigModule] = None,
+    device_cpp_wrapper_codegen: WrapperConstructor | None = None,
+    device_fx_wrapper_codegen: WrapperConstructor | None = None,
+    device_custom_pass: CustomGraphModulePass | None = None,
+    device_custom_config: ConfigModule | None = None,
 ) -> None:
     device_codegens[device] = DeviceCodegen(
         device_scheduling,
@@ -445,7 +437,7 @@ class BackendFeature(Enum):
 
 
 def get_backend_features(
-    device: Union[torch.device, str, None],
+    device: torch.device | str | None,
 ) -> OrderedSet[BackendFeature]:
     if device is None:
         return OrderedSet()
@@ -463,20 +455,20 @@ def get_backend_features(
 
 
 def has_backend_feature(
-    device: Union[torch.device, str, None], feature: BackendFeature
+    device: torch.device | str | None, feature: BackendFeature
 ) -> bool:
     """See also V.graph.has_feature"""
     assert isinstance(feature, BackendFeature)
     return feature in get_backend_features(device)
 
 
-def get_scheduling_for_device(device: str) -> Optional[SchedulingConstructor]:
+def get_scheduling_for_device(device: str) -> SchedulingConstructor | None:
     return device_codegens[device].scheduling if device in device_codegens else None
 
 
 def get_wrapper_codegen_for_device(
     device: str, cpp_wrapper: bool = False, fx_wrapper: bool = False
-) -> Optional[WrapperConstructor]:
+) -> WrapperConstructor | None:
     if device in device_codegens:
         wrapper_codegen_obj: DeviceCodegen = device_codegens[device]
         if fx_wrapper:
@@ -488,11 +480,11 @@ def get_wrapper_codegen_for_device(
     return None
 
 
-def get_custom_backend_pass_for_device(device: str) -> Optional[CustomGraphModulePass]:
+def get_custom_backend_pass_for_device(device: str) -> CustomGraphModulePass | None:
     return custom_backend_passes.get(device)
 
 
-def get_custom_backend_config_for_device(device: str) -> Optional[ConfigModule]:
+def get_custom_backend_config_for_device(device: str) -> ConfigModule | None:
     return custom_backend_codegen_configs.get(device)
 
 
@@ -515,6 +507,7 @@ def init_backend_registration() -> None:
     from .triton import TritonScheduling
     from .wrapper import PythonWrapperCodegen
     from .wrapper_fxir import WrapperFxCodegen
+    from .xpu.xpu_combined_scheduling import XPUCombinedScheduling
 
     if get_scheduling_for_device("cpu") is None:
         cpu_backends = {
@@ -548,10 +541,19 @@ def init_backend_registration() -> None:
             WrapperFxCodegen,
         )
 
+    if get_scheduling_for_device("tpu") is None:
+        register_backend_for_device(
+            "tpu",
+            PallasScheduling,
+            PythonWrapperCodegen,
+            # CppWrapperGpu,
+            # WrapperFxCodegen,
+        )
+
     if get_scheduling_for_device("xpu") is None:
         register_backend_for_device(
             "xpu",
-            TritonScheduling,
+            XPUCombinedScheduling,
             PythonWrapperCodegen,
             CppWrapperGpu,
             WrapperFxCodegen,
@@ -626,21 +628,28 @@ def register_device_op_overrides(
     device_op_overrides_dict[device] = device_op_overrides
 
 
-def get_device_op_overrides(device: str) -> DeviceOpOverrides:
-    assert isinstance(device, str), type(device)
+def _initialize_device_op_overrides():
+    # Use a flag rather than checking device_op_overrides_dict, since external/test
+    # code may partially populate it before we are called.
+    global _device_op_overrides_initialized
+    if _device_op_overrides_initialized:
+        return
+
+    from . import mps_device_op_overrides  # noqa: F401
+    from .cpu_device_op_overrides import CpuDeviceOpOverrides
+    from .cuda import device_op_overrides  # noqa: F401
+    from .mtia import device_op_overrides as mtia_op_overrides  # noqa: F401
+    from .xpu import device_op_overrides as xpu_op_overrides  # noqa: F401
 
-    if not device_op_overrides_dict:
-        from . import cpu_device_op_overrides, mps_device_op_overrides  # noqa: F401
-        from .cuda import device_op_overrides  # noqa: F401
-        from .mtia import device_op_overrides as mtia_op_overrides  # noqa: F401
-        from .xpu import device_op_overrides as xpu_op_overrides  # noqa: F401
+    # TPU uses Pallas for codegen and only needs no-op overrides
+    register_device_op_overrides("tpu", CpuDeviceOpOverrides())
 
-    if device not in device_op_overrides_dict:
-        # For backends like TPU that only need no-op overrides (Pallas handles codegen)
-        from .cpu_device_op_overrides import CpuDeviceOpOverrides
+    _device_op_overrides_initialized = True
 
-        register_device_op_overrides(device, CpuDeviceOpOverrides())
 
+def get_device_op_overrides(device: str) -> DeviceOpOverrides:
+    assert isinstance(device, str), type(device)
+    _initialize_device_op_overrides()
     return device_op_overrides_dict[device]
 
 
@@ -670,7 +679,7 @@ def deduce_output_dtype_by_name(
     op_name: str,
     *args: Any,
     **kwargs: Any,
-) -> Optional[torch.dtype]:
+) -> torch.dtype | None:
     """
     Given op name and a list of input dtypes, deduce the output dtype
     """
@@ -757,13 +766,13 @@ def check_nan(buffer: IndentedBuffer, var: CSEVariableType) -> None:
 class DataTypePropagation:
     def __init__(self, body: LoopBody) -> None:
         self.body = body
-        self.graphs: dict[Union[Callable[..., Any], str], Any] = {
+        self.graphs: dict[Callable[..., Any] | str, Any] = {
             "root": body.root_block.graph
         }
         for k, v in body.subblocks.items():
             self.graphs[k] = v.graph
 
-    def deduce_node_dtype_by_inputs(self, node: torch.fx.Node) -> Optional[torch.dtype]:
+    def deduce_node_dtype_by_inputs(self, node: torch.fx.Node) -> torch.dtype | None:
         inputs = node.all_input_nodes
         input_nodes = [
             n for n in inputs if isinstance(n, torch.fx.Node) and n.op != "placeholder"
@@ -790,7 +799,7 @@ def deduce_node_dtype_by_subgraph(self, node: torch.fx.Node) -> torch.dtype:
         assert dtype
         return dtype
 
-    def deduce_node_dtype(self, node: torch.fx.Node) -> Optional[torch.dtype]:
+    def deduce_node_dtype(self, node: torch.fx.Node) -> torch.dtype | None:
         if node.op == "placeholder":
             return None
 
@@ -819,9 +828,9 @@ def deduce_node_dtype(self, node: torch.fx.Node) -> Optional[torch.dtype]:
 
         return self.deduce_node_dtype_by_inputs(node)
 
-    def propagate_graph(self, graph: torch.fx.Graph) -> Optional[torch.dtype]:
+    def propagate_graph(self, graph: torch.fx.Graph) -> torch.dtype | None:
         assert graph.nodes
-        graph_dtype: Optional[torch.dtype] = None
+        graph_dtype: torch.dtype | None = None
         # For masked_subblock, we use output's dtype to represent
         # the dtype of this subgraph. For other cases, graph_dtype
         # might be None
@@ -837,15 +846,15 @@ def propagate_graph(self, graph: torch.fx.Graph) -> Optional[torch.dtype]:
                 graph_dtype = opt_ctx.dtype
         return graph_dtype
 
-    def propagate(self) -> Optional[torch.dtype]:
+    def propagate(self) -> torch.dtype | None:
         return self.propagate_graph(self.graphs["root"])
 
     @classmethod
-    def propagate_loopbody(cls, body: LoopBody) -> Optional[torch.dtype]:
+    def propagate_loopbody(cls, body: LoopBody) -> torch.dtype | None:
         return cls(body).propagate()
 
     @classmethod
-    def propagate_scheduler_node(cls, node: SchedulerNode) -> Optional[torch.dtype]:
+    def propagate_scheduler_node(cls, node: SchedulerNode) -> torch.dtype | None:
         from ..loop_body import LoopBody
         from ..scheduler import SchedulerNode
 
@@ -885,7 +894,9 @@ def identity(value: OpVarT) -> OpVarT:
 
     @staticmethod
     def reciprocal(x: OpVarT) -> OpVarT:
-        return ops.truediv(ops.constant(1, torch.int32), x)
+        # Use float32 constant so that div_rn can be applied when
+        # eager_numerics.division_rounding is enabled
+        return ops.truediv(ops.constant(1.0, torch.float32), x)
 
     @staticmethod
     def square(x: OpVarT) -> OpVarT:
@@ -938,6 +949,11 @@ def mul_rn(x: OpVarT, y: OpVarT) -> OpVarT:
         # for backends that don't override this, just use regular mul
         return ops.mul(x, y)
 
+    @staticmethod
+    def div_rn(x: OpVarT, y: OpVarT) -> OpVarT:
+        # for backends that don't override this, just use regular div
+        return ops.truediv(x, y)
+
     @staticmethod
     def floor_to_int(a: OpVarT, dtype: torch.dtype) -> OpVarT:
         return ops.to_dtype(ops.floor(a), dtype)
@@ -997,38 +1013,46 @@ def paren(string: OpVarT) -> OpVarT:
         return f"({string})"
 
     @staticmethod
-    def constant(value: Union[bool, float, int], dtype: torch.dtype) -> OpVarT:
+    def constant(value: bool | float | int, dtype: torch.dtype) -> OpVarT:
         return repr(value)
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def bitwise_not(x: OpVarT) -> OpVarT:
         return f"~{OpOverrides.paren(x)}"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def logical_not(a: OpVarT) -> OpVarT:
         return f"{OpOverrides.paren(a)} == 0"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def bitwise_and(x: OpVarT, y: OpVarT) -> OpVarT:
         return f"{OpOverrides.paren(x)} & {OpOverrides.paren(y)}"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def bitwise_or(x: OpVarT, y: OpVarT) -> OpVarT:
         return f"{OpOverrides.paren(x)} | {OpOverrides.paren(y)}"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def bitwise_xor(x: OpVarT, y: OpVarT) -> OpVarT:
         return f"{OpOverrides.paren(x)} ^ {OpOverrides.paren(y)}"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def bitwise_left_shift(x: OpVarT, y: OpVarT) -> OpVarT:
         return f"{OpOverrides.paren(x)} << {OpOverrides.paren(y)}"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def bitwise_right_shift(x: OpVarT, y: OpVarT) -> OpVarT:
         return f"{OpOverrides.paren(x)} >> {OpOverrides.paren(y)}"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def int_truediv(a: OpVarT, b: OpVarT) -> OpVarT:
         # TODO: this is wrong
         # TODO: an easy bandaid is to generate runtime asserts that it's
@@ -1042,7 +1066,7 @@ def load_seed(name: str, offset: OpVarT) -> OpVarT:
     def indirect_indexing(
         self,
         var: OpVarT,
-        size: Union[sympy.Expr, int],
+        size: sympy.Expr | int,
         check: bool = True,
         wrap_neg: bool = True,
     ) -> sympy.Symbol:
@@ -1082,8 +1106,8 @@ def reduction(
         dtype: torch.dtype,
         src_dtype: torch.dtype,
         reduction_type: ReductionType,
-        value: Union[OpVarT, tuple[OpVarT, ...]],
-    ) -> Union[OpVarT, tuple[OpVarT, ...]]:
+        value: OpVarT | tuple[OpVarT, ...],
+    ) -> OpVarT | tuple[OpVarT, ...]:
         raise NotImplementedError(
             f"{type(self).__name__}: reduction should be handled by CSEProxy"
         )
@@ -1119,8 +1143,8 @@ def bucketize(
         boundary_indices: OpVarT,
         indexing_dtype: torch.dtype,
         right: bool,
-        sorter: Optional[tuple[str, sympy.Expr]] = None,
-        sorter_indices: Optional[OpVarT] = None,
+        sorter: tuple[str, sympy.Expr] | None = None,
+        sorter_indices: OpVarT | None = None,
     ) -> OpVarT:
         raise NotImplementedError(
             f"{type(self).__name__}: bucketize should be handled by CSEProxy"
@@ -1140,10 +1164,11 @@ def inline_asm_elementwise(
         self,
         *inputs: OpVarT,
         asm: str,
-        constraints: Optional[str] = None,
+        constraints: str | None = None,
         dtype: torch.dtype = torch.float32,
         is_pure: bool = True,
         pack: int = 1,
+        input_dtypes: tuple[torch.dtype, ...] | None = None,
     ) -> OpVarT:
         raise NotImplementedError(
             f"{type(self).__name__}: inline_asm_elementwise only implemented for Triton backend"
@@ -1198,14 +1223,14 @@ class OverridesData:
     name: str
     cpp: Callable[..., str]
     # None when not impl in libdevice/triton
-    triton: Optional[Callable[..., str]] = None
+    triton: Callable[..., str] | None = None
     # None when not impl in aten/.../vec
-    cppvec: Optional[Callable[..., str]] = None
+    cppvec: Callable[..., str] | None = None
     type_promotion_kind: ELEMENTWISE_TYPE_PROMOTION_KIND = (
         ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT
     )
-    halide: Optional[Callable[..., str]] = None
-    mps: Optional[Callable[..., str]] = None
+    halide: Callable[..., str] | None = None
+    mps: Callable[..., str] | None = None
 
 
 # NB: if you add a new special function, don't forget to update
@@ -1475,7 +1500,7 @@ def __init__(self, name: str, line: str):
         self.name = name
         assert not isinstance(line, DeferredLineBase)
 
-    def __call__(self) -> Optional[str]:
+    def __call__(self) -> str | None:
         if not is_buffer_removed(self.name):
             return self.line
         return None
@@ -1532,10 +1557,10 @@ class KernelArgs:
     @staticmethod
     def _lookup(
         prefix: str,
-        odict: Union[dict[_T, Union[str, RemovedArg]], dict[_T, str]],
+        odict: dict[_T, str | RemovedArg] | dict[_T, str],
         name: _T,
     ) -> str:
-        result: Union[str, RemovedArg] = odict.get(name, REMOVED)
+        result: str | RemovedArg = odict.get(name, REMOVED)
         if isinstance(result, RemovedArg):
             odict[name] = new_result = f"{prefix}{len(odict)}"
             return new_result
@@ -1543,8 +1568,8 @@ def _lookup(
 
     def __init__(self) -> None:
         self.input_buffers: dict[str, str] = {}
-        self.output_buffers: dict[str, Union[str, RemovedArg]] = {}
-        self.inplace_buffers: dict[str, Union[InplacedBuffer, RemovedArg]] = {}
+        self.output_buffers: dict[str, str | RemovedArg] = {}
+        self.inplace_buffers: dict[str, InplacedBuffer | RemovedArg] = {}
         self.sizevars: dict[sympy.Expr, str] = {}
         self.workspace_args: list[WorkspaceArg] = []
 
@@ -1716,7 +1741,7 @@ def call_names(self) -> Iterator[str]:
             self.input_buffers.keys(), self.output_buffers.keys(), self.sizevars.keys()
         )
 
-    def arg_name(self, name: str) -> Optional[str]:
+    def arg_name(self, name: str) -> str | None:
         """
         Returns inner name of a given outer name.
         """
@@ -1735,7 +1760,7 @@ def wrap_size_arg(self, size: SymbolLike) -> str:
         return str(size)
 
     def cpp_argdefs(
-        self, dtype_to_cpp_type: Optional[dict[torch.dtype, str]] = None
+        self, dtype_to_cpp_type: dict[torch.dtype, str] | None = None
     ) -> tuple[list[str], list[str], list[str]]:
         from .cpp_utils import INDEX_TYPE
 
@@ -1886,7 +1911,7 @@ def __init__(
         self,
         name: str,
         bounds: ValueRanges[Any],
-        dtype: Optional[torch.dtype] = None,
+        dtype: torch.dtype | None = None,
         shape: BlockShapeType = None,
     ):
         super().__init__()
@@ -1920,7 +1945,7 @@ def __repr__(self) -> str:
     ReductionCacheKey = tuple[
         torch.dtype,
         ReductionType,
-        Union[CSEVariable, tuple[CSEVariable, ...]],
+        CSEVariable | tuple[CSEVariable, ...],
     ]
 
 
@@ -1932,12 +1957,11 @@ def __init__(
         prefix: str = "",
         suffix: str = "",
         name_prefix: str = "tmp",
-        iter_buffers: Optional[itertools.count[int]] = None,
-        store_cache: Optional[MutableMapping[str, CSEVariableType]] = None,
-        reduction_cache: Optional[
-            MutableMapping[ReductionCacheKey, CSEVariableType]
-        ] = None,
-        varname_map: Optional[dict[str, CSEVariableType]] = None,
+        iter_buffers: itertools.count[int] | None = None,
+        store_cache: MutableMapping[str, CSEVariableType] | None = None,
+        reduction_cache: MutableMapping[ReductionCacheKey, CSEVariableType]
+        | None = None,
+        varname_map: dict[str, CSEVariableType] | None = None,
     ):
         self.prefix = prefix
         self.suffix = suffix
@@ -1990,7 +2014,7 @@ def put(self, cache_key: str, val: CSEVariableType) -> None:
     def contains(self, cache_key: str) -> bool:
         return self.augment_key(cache_key) in self._cache
 
-    def try_get(self, cache_key: str) -> Optional[CSEVariableType]:
+    def try_get(self, cache_key: str) -> CSEVariableType | None:
         return self._cache.get(self.augment_key(cache_key), None)
 
     def get(self, cache_key: str) -> CSEVariableType:
@@ -1999,12 +2023,12 @@ def get(self, cache_key: str) -> CSEVariableType:
     def generate(
         self,
         buffer: IndentedBuffer,
-        expr: Union[str, CSEVariable, OpsValue, IndentedBuffer, DeferredLineBase],
+        expr: str | CSEVariable | OpsValue | IndentedBuffer | DeferredLineBase,
         *,
         bounds: ValueRanges[Any] = ValueRanges.unknown(),
         write: bool = True,
         assignment: bool = True,
-        dtype: Optional[torch.dtype] = None,
+        dtype: torch.dtype | None = None,
         shape: BlockShapeType = None,
     ) -> CSEVariableType:
         if isinstance(expr, OpsValue):
@@ -2076,7 +2100,7 @@ def generate(
     def newvar(
         self,
         bounds: ValueRanges[Any] = ValueRanges.unknown(),
-        dtype: Optional[torch.dtype] = None,
+        dtype: torch.dtype | None = None,
         shape: BlockShapeType = None,
     ) -> CSEVariableType:
         var_name = f"{self.name_prefix}{next(self.iter_buffer_ids)}"
@@ -2088,7 +2112,7 @@ def namedvar(
         self,
         name: str,
         bounds: ValueRanges[Any] = ValueRanges.unknown(),
-        dtype: Optional[torch.dtype] = None,
+        dtype: torch.dtype | None = None,
         shape: BlockShapeType = None,
     ) -> CSEVariableType:
         torch._check_value(
@@ -2115,10 +2139,10 @@ def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
 class Kernel(CodeGen, Generic[CSEVariableType]):
     newvar_prefix: str = ""
     suffix: str = ""
-    overrides: Optional[Callable[[], OpsHandler[Any]]] = None
+    overrides: Callable[[], OpsHandler[Any]] | None = None
 
     def __init__(
-        self, args: Optional[KernelArgs] = None, increase_kernel_count: bool = True
+        self, args: KernelArgs | None = None, increase_kernel_count: bool = True
     ) -> None:
         super().__init__()
         if increase_kernel_count:
@@ -2137,11 +2161,11 @@ def __init__(
         self.cse: CSE[CSEVariableType, Any] = CSE(self.newvar_prefix, self.suffix)
         self.must_keep_buffers: OrderedSet[str] = OrderedSet()
         self.store_buffer_names: OrderedSet[str] = OrderedSet()
-        self._load_mask: Optional[str] = None
-        self._load_other: Union[None, int, float] = None
+        self._load_mask: str | None = None
+        self._load_other: None | int | float = None
         # OrderedSet in set_current_node
-        self.current_node: Optional[SchedulerNode] = None
-        self.node_to_bounds: Optional[dict[torch.fx.Node, ValueRanges[Any]]] = None
+        self.current_node: SchedulerNode | None = None
+        self.node_to_bounds: dict[torch.fx.Node, ValueRanges[Any]] | None = None
 
         self.removed_buffers: OrderedSet[str] = OrderedSet()
         self.inplaced_to_remove: OrderedSet[str] = OrderedSet()
@@ -2152,7 +2176,7 @@ def __init__(
         self.inplace_update_buffers: dict[str, str] = {}
         # Set minimum number of elements processed per thread.
         self.min_elem_per_thread = 1
-        self.kernel_name: Optional[str] = None
+        self.kernel_name: str | None = None
 
     @contextlib.contextmanager
     def set_current_node(self, node: SchedulerNode) -> Iterator[None]:
@@ -2168,8 +2192,8 @@ def set_current_node(self, node: SchedulerNode) -> Iterator[None]:
     def swap_buffers(
         self,
         lb: IndentedBuffer,
-        cb: Optional[IndentedBuffer] = None,
-        sb: Optional[IndentedBuffer] = None,
+        cb: IndentedBuffer | None = None,
+        sb: IndentedBuffer | None = None,
     ) -> Iterator[None]:
         if cb is None:
             cb = lb
@@ -2241,8 +2265,8 @@ def reduction(
         dtype: torch.dtype,
         src_dtype: torch.dtype,
         reduction_type: ReductionType,
-        value: Union[CSEVariable, tuple[CSEVariable, ...]],
-    ) -> Union[CSEVariable, tuple[CSEVariable, ...]]:
+        value: CSEVariable | tuple[CSEVariable, ...],
+    ) -> CSEVariable | tuple[CSEVariable, ...]:
         raise NotImplementedError
 
     def partial_accumulate(
@@ -2283,8 +2307,8 @@ def bucketize(
         boundary_indices: CSEVariable,
         indexing_dtype: torch.dtype,
         right: bool,
-        sorter: Optional[tuple[str, sympy.Expr]] = None,
-        sorter_indices: Optional[CSEVariable] = None,
+        sorter: tuple[str, sympy.Expr] | None = None,
+        sorter_indices: CSEVariable | None = None,
     ) -> CSEVariable:
         """
         See [Note: Inductor bucketize op]
@@ -2297,10 +2321,10 @@ def assert_function(self) -> str:
 
     def indirect_assert(
         self,
-        var: Union[CSEVariable, str],
-        lower: Optional[str],
-        upper: Optional[str],
-        mask: Optional[Union[CSEVariable, str]] = None,
+        var: CSEVariable | str,
+        lower: str | None,
+        upper: str | None,
+        mask: CSEVariable | str | None = None,
     ) -> str:
         if isinstance(var, CSEVariable):
             var = str(var)
@@ -2400,7 +2424,7 @@ def remove_inplace_buffer(self, name: str) -> None:
         self.removed_buffers.add(name)
 
     def rename_indexing(
-        self, index: Union[list[sympy.Expr], tuple[sympy.Expr, ...], sympy.Expr]
+        self, index: list[sympy.Expr] | tuple[sympy.Expr, ...] | sympy.Expr
     ) -> sympy.Expr:
         # adds the necessary kernel args for index expressions
         # and renames variables in index expressions to kernel arg names
@@ -2426,7 +2450,7 @@ def rename_indexing(
     def create_cse_var(self, *args: Any, **kwargs: Any) -> CSEVariable:
         return CSEVariable(*args, **kwargs)
 
-    def arg_name(self, node: IRNode) -> Optional[str]:
+    def arg_name(self, node: IRNode) -> str | None:
         """
         Returns arg name of a given input or output node.
         """
@@ -2439,7 +2463,7 @@ def arg_name(self, node: IRNode) -> Optional[str]:
 class OptimizationContext:
     key: ClassVar[str] = "opt_ctx"
 
-    dtype: Optional[torch.dtype] = None
+    dtype: torch.dtype | None = None
     ops_name: str = ""
 
 
@@ -2520,7 +2544,7 @@ def __str__(self) -> str:
 
     @staticmethod
     def _fake_get_dtype(
-        fake_outs: Union[list[Buffer], Buffer],
+        fake_outs: list[Buffer] | Buffer,
     ) -> Callable[[str], torch.dtype]:
         _get_dtype_real = V.graph.get_dtype
         if isinstance(fake_outs, (list, tuple)):
@@ -2536,7 +2560,7 @@ def get_dtype(name: str) -> torch.dtype:
 
         return get_dtype
 
-    def __init__(self, name: str, hash: Optional[str] = None) -> None:
+    def __init__(self, name: str, hash: str | None = None) -> None:
         self.name = name
         self._hash = hash
 
@@ -2552,7 +2576,7 @@ def uid(self) -> str:
         return self.name
 
     @property
-    def src_hash(self) -> Union[str, None]:
+    def src_hash(self) -> str | None:
         """
         source hash for a Template.
 
@@ -2562,7 +2586,7 @@ def src_hash(self) -> Union[str, None]:
         """
         return self._hash
 
-    def choice_or_none(self, **kwargs: Any) -> Optional[ChoiceCaller]:
+    def choice_or_none(self, **kwargs: Any) -> ChoiceCaller | None:
         """
         Maybe generates a new ChoiceCaller and returns it, or None if generation fails.
 
@@ -2576,7 +2600,7 @@ def choice_or_none(self, **kwargs: Any) -> Optional[ChoiceCaller]:
 
     def maybe_append_choice(
         self, choices: list[Any], **kwargs: Any
-    ) -> Optional[NotImplementedError]:
+    ) -> NotImplementedError | None:
         """
         Maybe generates a new ChoiceCaller and appends it into existing choices.
         Returns None if success, otherwise returns the error.
@@ -2653,10 +2677,10 @@ def _default(self, name: str, args: tuple[Any, ...], kwargs: dict[str, Any]) ->
 
         output_idx = 0
 
-        def do_cse(v: Union[str, CSEVariable]) -> CSEVariable:
+        def do_cse(v: str | CSEVariable) -> CSEVariable:
             # we tree_map over the output, so we need to fetch corresponding dtype
             nonlocal output_idx
-            var_dtype: Optional[torch.dtype] = (
+            var_dtype: torch.dtype | None = (
                 output_dtype[output_idx]
                 if isinstance(output_dtype, (list, tuple))
                 else output_dtype
@@ -2756,7 +2780,7 @@ def arg_to_bound(x: Any) -> Any:
     def indirect_indexing(
         self,
         var: CSEVariable,
-        size: Union[sympy.Expr, int],
+        size: sympy.Expr | int,
         check: bool = True,
         wrap_neg: bool = True,
     ) -> sympy.Symbol:
@@ -2841,7 +2865,8 @@ def store(
         self, name: str, index: sympy.Expr, value: CSEVariable, mode: StoreMode = None
     ) -> None:
         self.kernel.store_buffer_names.add(name)
-        if mode is None:
+        # Update store cache when mode is None or "tma"
+        if mode != "atomic_add":
             self._update_store_cache(name, value)
         if name not in V.graph.removed_buffers:
             self.kernel.store(name, index, value, mode=mode)
@@ -2867,8 +2892,8 @@ def reduction(
         dtype: torch.dtype,
         src_dtype: torch.dtype,
         reduction_type: ReductionType,
-        value: Union[CSEVariable, tuple[CSEVariable, ...]],
-    ) -> Union[CSEVariable, tuple[CSEVariable, ...]]:
+        value: CSEVariable | tuple[CSEVariable, ...],
+    ) -> CSEVariable | tuple[CSEVariable, ...]:
         self.kernel.num_reduction += 1
         return self.kernel.reduction(dtype, src_dtype, reduction_type, value)
 
@@ -2899,8 +2924,8 @@ def bucketize(
         boundary_indices: CSEVariable,
         indexing_dtype: torch.dtype,
         right: bool,
-        sorter: Optional[tuple[str, sympy.Expr]] = None,
-        sorter_indices: Optional[CSEVariable] = None,
+        sorter: tuple[str, sympy.Expr] | None = None,
+        sorter_indices: CSEVariable | None = None,
     ) -> CSEVariable:
         """
         [Note: Inductor bucketize op]
diff --git a/torch/_inductor/codegen/cpp.py b/torch/_inductor/codegen/cpp.py
index e81f2c53bf717..a392c0c184576 100644
--- a/torch/_inductor/codegen/cpp.py
+++ b/torch/_inductor/codegen/cpp.py
@@ -5,12 +5,13 @@
 import itertools
 import math
 import operator
+import os
 import re
 import sys
 import warnings
 from collections.abc import Callable, Sequence
 from enum import Enum
-from typing import Any, cast, Optional, Union
+from typing import Any, cast, Optional
 
 import sympy
 
@@ -209,7 +210,7 @@ def reduction_combine(
     var,
     next_value,
     helper_val=None,
-    index: Optional[sympy.Symbol] = None,
+    index: sympy.Symbol | None = None,
     src_dtype=None,
 ):
     is_bool = src_dtype == torch.bool
@@ -314,11 +315,11 @@ def move_code_under_inner_loop(
 
 
 def reduction_prefix_array(
-    acc_var: Union[str, CSEVariable],
+    acc_var: str | CSEVariable,
     acc_type: str,
     reduction_type: str,
     dtype: torch.dtype,
-    len: Union[str, int],
+    len: str | int,
     init_fn,
 ):
     """
@@ -463,7 +464,7 @@ def visit_modular_indexing(divisor, modulus):
 
 @functools.lru_cache
 def stride_at_vec_range(
-    index: sympy.Expr, var: sympy.Symbol, vec_length: Optional[int] = None
+    index: sympy.Expr, var: sympy.Symbol, vec_length: int | None = None
 ):
     if vec_length:
         index = simplify_index_in_vec_range(index, var, vec_length)
@@ -522,10 +523,10 @@ def fuse(  # type: ignore[override]
     def __init__(
         self,
         scheduler: "Scheduler",
-        outer_fused_nodes: list[Union[FusedSchedulerNode, SchedulerNode]],
+        outer_fused_nodes: list[FusedSchedulerNode | SchedulerNode],
         outer_loop_fusion_depth,
     ):
-        self.outer_fused_nodes: list[Union[FusedSchedulerNode, SchedulerNode]] = (
+        self.outer_fused_nodes: list[FusedSchedulerNode | SchedulerNode] = (
             outer_fused_nodes
         )
         self.outer_loop_fusion_depth = outer_loop_fusion_depth
@@ -645,8 +646,8 @@ def merge_outer_fusion_kernels(
 class RecordOptimizationContext:
     def __init__(self, func_name: str = ""):
         self.func_name = func_name
-        self.current_node: Optional[torch.fx.Node] = None
-        self.opt_ctx: Optional[OptimizationContext] = None
+        self.current_node: torch.fx.Node | None = None
+        self.opt_ctx: OptimizationContext | None = None
 
     def __enter__(self):
         assert V.interpreter
@@ -744,20 +745,36 @@ def to_dtype(x, dtype, src_dtype=None, use_compute_types=True):
             V.kernel.cache_dtype_convert(x, src_dtype, csevar, dtype)
         return csevar
 
+    @staticmethod
+    # pyrefly: ignore [bad-override]
+    def round_to_int(x, dtype, src_dtype=None, use_compute_types=True):
+        assert isinstance(x, CppCSEVariable)
+        if src_dtype is None:
+            src_dtype = x.dtype
+        expr = V.kernel.get_to_dtype_expr(x, dtype, src_dtype, rounding=True)
+        csevar = V.kernel.cse.generate(V.kernel.compute, expr)
+        csevar.update_on_args("round_to_int", (x, dtype), {"src_dtype": src_dtype})
+        if dtype in DTYPE_LOWP_FP and src_dtype == torch.float:
+            V.kernel.cache_dtype_convert(x, src_dtype, csevar, dtype)
+        return csevar
+
     @staticmethod
     def to_dtype_bitcast(x, dtype, src_dtype):
         assert dtype in DTYPE_TO_CPP, f"{dtype} missing from {__name__}.DTYPE_TO_CPP"
         return f"c10::bit_cast<{DTYPE_TO_CPP[dtype]}>({x})"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def abs(x):
         return f"std::abs({x})"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def sin(x):
         return f"std::sin({x})"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def cos(x):
         return f"std::cos({x})"
 
@@ -766,6 +783,7 @@ def neg(x):
         return f"decltype({x})(-{x})"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def exp(x):
         # return f"Sleef_expf_u10({x})"
         return f"std::exp({x})"
@@ -779,6 +797,7 @@ def expm1(x):
         return f"std::expm1({x})"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def erf(x):
         return f"std::erf({x})"
 
@@ -787,14 +806,17 @@ def erfc(x):
         return f"std::erfc({x})"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def erfinv(x):
         return f"calc_erfinv({x})"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def sqrt(x):
         return f"std::sqrt({x})"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def rsqrt(x):
         return f"1 / std::sqrt({x})"
 
@@ -811,14 +833,17 @@ def log1p(x):
             )
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def tan(x):
         return f"std::tan({x})"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def tanh(x):
         return f"std::tanh({x})"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def signbit(x):
         """
         On windows std::signbit only support float type.
@@ -835,14 +860,17 @@ def pow(a, b):
         return f"std::pow({a}, {b})"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def log(x):
         return f"std::log({x})"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def round(x):
         return f"std::nearbyint({x})"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def floor(x):
         return f"std::floor({x})"
 
@@ -854,75 +882,93 @@ def floordiv(a, b):
         return f"(({a} < 0) != ({b} < 0) ? ({rem} != 0 ? {quot} - 1 : {quot}) : {quot})"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def ceil(x):
         return f"std::ceil({x})"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def trunc(x):
         return f"std::trunc({x})"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def truncdiv(a, b):
         # a and b are integer type
         return f"{a} / {b}"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def fmod(a, b):
         return f"std::fmod({a}, {b})"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def isinf(x):
         return f"std::isinf({x})"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def isnan(x):
         return f"std::isnan({x})"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def lgamma(x):
         return f"std::lgamma({x})"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def acos(x):
         return f"std::acos({x})"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def acosh(x):
         return f"std::acosh({x})"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def cosh(x):
         return f"std::cosh({x})"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def sinh(x):
         return f"std::sinh({x})"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def asin(x):
         return f"std::asin({x})"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def asinh(x):
         return f"std::asinh({x})"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def atan2(x, y):
         return f"std::atan2({x}, {y})"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def atan(x):
         return f"std::atan({x})"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def atanh(x):
         return f"std::atanh({x})"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def copysign(x, y):
         return f"std::copysign({x}, {y})"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def frexp(x):
         cache_keys = f"frexp({x})[0]", f"frexp({x})[1]"
         if all(V.kernel.cse.try_get(cache_key) is not None for cache_key in cache_keys):
@@ -940,6 +986,7 @@ def frexp(x):
         return mantissa, exponent
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def hypot(x, y):
         return f"std::hypot({x}, {y})"
 
@@ -952,10 +999,12 @@ def log2(x):
         return f"std::log2({x})"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def ldexp(x, n):
         return f"std::ldexp({x}, {n})"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def nextafter(x, y):
         return f"std::nextafter({x}, {y})"
 
@@ -976,14 +1025,17 @@ def relu(x):
             )
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def minimum(a, b):
         return f"min_propagate_nan({a}, {b})"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def maximum(a, b):
         return f"max_propagate_nan({a}, {b})"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def where(a, b, c):
         return f"{a} ? {b} : {c}"
 
@@ -1021,6 +1073,7 @@ def masked(mask, body, other):
         return f"{mask} ? {body_var}() : {other_code}"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def logical_and(a, b):
         return f"{a} && {b}"
 
@@ -1029,10 +1082,12 @@ def logical_not(a):
         return f"!{a}"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def logical_or(a, b):
         return f"{a} || {b}"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def logical_xor(a, b):
         return f"{a} != {b}"
 
@@ -1094,6 +1149,19 @@ def bitwise_right_shift(a, b):
     def rand(seed: sympy.Expr, offset: sympy.Expr):
         return f"normalized_rand_cpu({seed}, {offset})"
 
+    @staticmethod
+    def rand_eager(
+        seed: sympy.Expr,
+        base_offset: sympy.Expr,
+        threads_per_round: sympy.Expr,
+        tid: sympy.Expr,
+        vec: sympy.Expr,
+    ):
+        # NOTE: This is a codegen fallback used by the C++ backend for eager random.
+        # It is not intended to provide CPU parity; parity target is CUDA eager vs compiled.
+        # Keep this hook to satisfy codegen paths and CI.
+        return f"normalized_rand_cpu({seed}, {base_offset})"
+
     @staticmethod
     def randn(seed: sympy.Expr, offset: sympy.Expr):
         return f"randn_cpu({seed}, {offset})"
@@ -1107,6 +1175,7 @@ def sigmoid(x):
         return f"decltype({x})(1) / (decltype({x})(1) + std::exp(-{x}))"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def sign(x):
         code = BracesBuffer()
         scalar_zero = f"decltype({x})(0)"
@@ -1659,6 +1728,22 @@ def to_dtype(x, dtype, src_dtype=None, use_compute_types=True):
             V.kernel.cache_dtype_convert(x, src_dtype, csevar, dtype)
         return csevar
 
+    @staticmethod
+    def round_to_int(x, dtype, src_dtype=None, use_compute_types=True):
+        assert dtype in [
+            torch.uint8,
+            torch.int8,
+            torch.int32,
+        ], f"{__name__} does not support {dtype}"
+        assert isinstance(x, CppCSEVariable)
+        src_dtype = x.dtype
+        expr = V.kernel.get_to_dtype_expr(x, dtype, src_dtype, rounding=True)
+        csevar = V.kernel.cse.generate(V.kernel.compute, expr)
+        csevar.update_on_args("round_to_int", (x, dtype), {"src_dtype": src_dtype})
+        if dtype in DTYPE_LOWP_FP and src_dtype == torch.float:
+            V.kernel.cache_dtype_convert(x, src_dtype, csevar, dtype)
+        return csevar
+
     @staticmethod
     def log1p(x):
         bug = config.cpp.inject_log1p_bug_TESTING_ONLY
@@ -1724,7 +1809,7 @@ def maskify_or_vecify(code):
                     assert isinstance(other_vec_var, CppCSEVariable), other_vec_var
                     body_vec_var.dtype = dtype
                     other_vec_var.dtype = dtype
-                    overrides: type[Union[CppOverrides, CppVecOverrides]] = (
+                    overrides: type[CppOverrides | CppVecOverrides] = (
                         # pyrefly: ignore [bad-assignment]
                         V.kernel.overrides
                     )  # type: ignore[has-type]
@@ -1736,7 +1821,8 @@ def maskify_or_vecify(code):
                 V.kernel.compute,
                 code,
             )
-            result.is_vec = True
+            assert isinstance(csevar, CppCSEVariable)
+            csevar.is_vec = True
         elif result.is_vec:
             csevar = V.kernel.cse.generate(
                 V.kernel.compute, f"{mask} ? {body_code_vec} : {other_code_vec}"
@@ -1934,7 +2020,7 @@ def __init__(self, args, num_threads):
         # Indicate this kernel will be moved under the inner for-loop
         # See move_code_under_inner_loop
         self.inner_itervars: list[sympy.Symbol] = []
-        self.call_ranges: Optional[tuple[sympy.Expr, ...]] = None
+        self.call_ranges: tuple[sympy.Expr, ...] | None = None
         self.ranges: list[sympy.Expr] = []
         self.itervars: list[sympy.Symbol] = []
         self.reduction_depth = None
@@ -2012,7 +2098,7 @@ def update_stores_with_parallel_reduction(self):
         for var_name in self.reduction_var_names:
             replace_acc_name(self.stores, var_name, f"{var_name}_local")
 
-    def gen_body(self, code: Optional[BracesBuffer] = None):
+    def gen_body(self, code: BracesBuffer | None = None):
         assert code is None
         code = BracesBuffer()
         with contextlib.ExitStack() as stack:
@@ -2150,7 +2236,7 @@ def device_assert_async(self, cond, msg):
 
     def _gen_reduction_prefix(
         self,
-        acc: Union[CSEVariable, str],
+        acc: CSEVariable | str,
         acc_type: str,
         rtype: str,
         dtype: torch.dtype,
@@ -2162,7 +2248,7 @@ def _gen_reduction_prefix(
         # Otherwise, we will define and initialize a reduction array
         # => float tmp_acc0_arr[size];
         # => for (int i = 0; i < size; i++) tmp_acc0_arr[i] = 0;
-        def inner(size: Optional[int] = None):
+        def inner(size: int | None = None):
             if size is None:
                 return f"{acc_type} {acc} = {init_fn(rtype, dtype)};"
             else:
@@ -2177,7 +2263,7 @@ def inner(size: Optional[int] = None):
 
         return inner
 
-    def finalize_reduction_prefix(self, size: Optional[int] = None):
+    def finalize_reduction_prefix(self, size: int | None = None):
         for gen_fn in self.reduction_prefix_generators:
             self.reduction_prefix.splice(gen_fn(size))
 
@@ -2200,18 +2286,7 @@ def need_use_acc_helper(self, reduction_type, dtype, use_scalar):
             # chunk size to balance accuracy and performance
             chunk_size = 4096
 
-            # use acc helper If cannot get size_hint
-            try:
-                reduction_size_hint = V.graph.sizevars.size_hint(reduction_size)
-            except Exception:
-                return True
-
-            if reduction_size_hint > chunk_size:
-                # use helper if the reduction size is too large
-                V.graph.sizevars.check_lt(chunk_size, reduction_size)
-                return True
-            else:
-                V.graph.sizevars.check_leq(reduction_size, chunk_size)
+            return V.graph.sizevars.guard_or_false(sympy.Gt(reduction_size, chunk_size))
         return False
 
     def _acc_helper_init(
@@ -2379,9 +2454,8 @@ def set_ranges(self, lengths, reduction_lengths):
 
     def size_hint(self):
         assert self.call_ranges is not None
-        return V.graph.sizevars.size_hint(
-            sympy_product(self.call_ranges), fallback=8192
-        )
+        expr = sympy_product(self.call_ranges)
+        return V.graph.sizevars.optimization_hint(expr)
 
     def codegen_loops_impl(self, loop_nest, code, worksharing):
         assert isinstance(self, CppKernelProxy)
@@ -2584,7 +2658,13 @@ def write_to_suffix(self):
     def create_cse_var(self, *args, **kwargs):
         return CppCSEVariable(*args, **kwargs)
 
-    def get_to_dtype_expr(self, src, dtype, src_dtype):
+    def get_to_dtype_expr(self, src, dtype, src_dtype, rounding=False):
+        if (
+            rounding
+            and dtype in [torch.int8, torch.uint8]
+            and src_dtype in [torch.float, torch.double]
+        ):
+            return f"c10::convert<{DTYPE_TO_CPP[dtype]}>(std::round({src}))"
         return f"c10::convert<{DTYPE_TO_CPP[dtype]}>({src})"
 
     def cache_dtype_convert(self, dst, dst_dtype, src, src_dtype):
@@ -2594,8 +2674,8 @@ def cache_dtype_convert(self, dst, dst_dtype, src, src_dtype):
     def codegen_conditions(
         self,
         code: BracesBuffer,
-        prefix: Optional[str] = None,
-        var: Optional[sympy.Symbol] = None,
+        prefix: str | None = None,
+        var: sympy.Symbol | None = None,
     ):
         if prefix is None:
             prefix = ""
@@ -2711,7 +2791,7 @@ def _get_vec_load_line(
         var: str,
         index: sympy.Expr,
         dtype: torch.dtype,
-        load_mask: Optional[CppCSEVariable] = None,
+        load_mask: CppCSEVariable | None = None,
     ):
         """
         Get a load line str that loads a vector from `var` at `index` of type `dtype`.
@@ -2745,13 +2825,13 @@ def _get_vec_load_line(
 
     def _load_or_store_non_contiguous(
         self,
-        var: Optional[str],
+        var: str | None,
         index: sympy.Expr,
         dtype: torch.dtype,
-        buffer: Optional[IndentedBuffer] = None,
-        store_value: Optional[Union[str, CppCSEVariable]] = None,
+        buffer: IndentedBuffer | None = None,
+        store_value: str | CppCSEVariable | None = None,
         accu_store: bool = False,
-    ) -> Optional[CppCSEVariable]:
+    ) -> CppCSEVariable | None:
         """
         Load or store a vector in a non-contiguous way. The vector is initialized from an array that is
         filled in an inner loop over the tiling factor.
@@ -2908,7 +2988,7 @@ def load(self, name: str, index: sympy.Expr):
 
     def _get_store_line(
         self,
-        value: Union[str, CppCSEVariable],
+        value: str | CppCSEVariable,
         var: str,
         index: sympy.Expr,
         dtype: torch.dtype,
@@ -3165,7 +3245,7 @@ def reduction(self, dtype, src_dtype, reduction_type, value):
                 reduction_combine_fn=self.reduction_combine_vec,
                 reduction_init_fn=self.reduction_init_vec,
             )
-        tmpvar: Union[str, CSEVariable]
+        tmpvar: str | CSEVariable
         is_bool = dtype == torch.bool
         if horizontal_reduction:
             # Horizontal reduction
@@ -3307,18 +3387,20 @@ def reduction_init_vec(self, reduction_type, dtype):
             return f"Welford<{vec_type}>()"
 
         if reduction_type in ("argmin", "argmax"):
-            cdtype = DTYPE_TO_CPP[scalar_type]
+            # For bool argmin/argmax, we use float for computations
+            compute_dtype = torch.float if dtype == torch.bool else scalar_type
+            cdtype = DTYPE_TO_CPP[compute_dtype]
             acc_type = self.reduction_acc_type_vec(reduction_type, dtype)
             if reduction_type == "argmin":
                 val = (
                     f"std::numeric_limits<{cdtype}>::infinity()"
-                    if is_float_dtype(dtype)
+                    if is_float_dtype(dtype) or dtype == torch.bool
                     else f"std::numeric_limits<{cdtype}>::max()"
                 )
             else:
                 val = (
                     f"-std::numeric_limits<{cdtype}>::infinity()"
-                    if is_float_dtype(dtype)
+                    if is_float_dtype(dtype) or dtype == torch.bool
                     else f"std::numeric_limits<{cdtype}>::min()"
                 )
             return f"{acc_type}({val})"
@@ -3339,10 +3421,13 @@ def reduction_acc_type_vec(self, reduction_type, dtype):
         if is_welford_reduction(reduction_type):
             return f"Welford<{vec_type}>"
         if reduction_type in ("argmin", "argmax"):
-            n_src = self._get_num_vectors(scalar_type)
             n_idx = self._get_num_vectors(torch.int64)
             if dtype == torch.bool:
+                # For bool argmin/argmax, we use float for computations
+                # so n_src must be computed from float, not bool
+                n_src = self._get_num_vectors(torch.float)
                 return f"IndexValueVec<{DTYPE_TO_CPP[torch.float]}, {n_src}, {n_idx}>"
+            n_src = self._get_num_vectors(scalar_type)
             return f"IndexValueVec<{DTYPE_TO_CPP[scalar_type]}, {n_src}, {n_idx}>"
         if dtype == torch.bool:
             assert reduction_type in ("min", "max", "any", "sum")
@@ -3355,9 +3440,9 @@ def reduction_combine_vec(
         var,
         next_value,
         helper_val=None,
-        index: Optional[sympy.Symbol] = None,
-        horizontal_reduction: Optional[bool] = None,
-        src_dtype: Optional[torch.dtype] = torch.float32,
+        index: sympy.Symbol | None = None,
+        horizontal_reduction: bool | None = None,
+        src_dtype: torch.dtype | None = torch.float32,
     ):
         is_bool = src_dtype == torch.bool
         if reduction_type == "max":
@@ -3425,9 +3510,15 @@ def reduction_combine_vec(
         elif reduction_type in ("argmin", "argmax"):
             assert src_dtype is not None
             cdtype = DTYPE_TO_CPP[src_dtype]
+            compute_dtype = src_dtype
             if src_dtype == torch.bool:
+                # For bool argmin/argmax, we use float for computations
                 cdtype = DTYPE_TO_CPP[torch.float]
-            n_src = self._get_num_vectors(src_dtype)
+                compute_dtype = torch.float
+                # Convert bool VecMask to float vector for argmax_combine_vec
+                if isinstance(next_value, CppCSEVariable) and next_value.is_vec:
+                    (next_value,) = unify_mask_base_type(self.compute, (next_value,))
+            n_src = self._get_num_vectors(compute_dtype)
             n_idx = self._get_num_vectors(torch.int64)
             t_extra = ""
             arg_extra = ""
@@ -3479,7 +3570,7 @@ def indirect_assert(self, var, lower, upper, mask=None):
         cond = f"{self._get_mask_type(var.dtype)}({cond})"
         if mask:
             if not mask.is_vec:
-                mask = f"{self._get_mask_type(var.dtype)}({mask})"
+                mask = f"{self._get_mask_type(var.dtype)}::from({mask})"
             # We need not check when the mask is False
             cond = f"({cond}) | ~({mask})"
         if self.tail_size:
@@ -3490,10 +3581,10 @@ def indirect_assert(self, var, lower, upper, mask=None):
         cond = f"({cond}).all_masked()"
         return f'{self.assert_function}({cond}, "index out of bounds: {cond_print}")'
 
-    def get_to_dtype_expr(self, src, dtype, src_dtype):
+    def get_to_dtype_expr(self, src, dtype, src_dtype, rounding=False):
         assert isinstance(src, CppCSEVariable)
         if not src.is_vec:
-            return super().get_to_dtype_expr(src, dtype, src_dtype)
+            return super().get_to_dtype_expr(src, dtype, src_dtype, rounding)
         src_cpp_type = DTYPE_TO_CPP[src_dtype]
         src_num_vectors = self._get_num_vectors(src_dtype)
         dst_cpp_type = DTYPE_TO_CPP[dtype]
@@ -3504,10 +3595,22 @@ def get_to_dtype_expr(self, src, dtype, src_dtype):
         elif src_dtype == torch.bool and dtype != torch.bool:
             expr = f"{src}.to<{dst_cpp_type},{dst_num_vectors}>()"
         elif src_dtype != dtype:
+            expr = ""
+            if (
+                rounding
+                and src_dtype in [torch.float, torch.double]
+                and dtype in [torch.int8, torch.uint8]
+            ):
+                expr = "at::vec::round_convert"
+            else:
+                expr = "at::vec::convert"
             if src_num_vectors == dst_num_vectors == 1:
-                expr = f"at::vec::convert<{dst_cpp_type}>({src})"
+                expr = expr + f"<{dst_cpp_type}>({src})"
             else:
-                expr = f"at::vec::convert<{dst_cpp_type},{dst_num_vectors},{src_cpp_type},{src_num_vectors}>({src})"
+                expr = (
+                    expr
+                    + f"<{dst_cpp_type},{dst_num_vectors},{src_cpp_type},{src_num_vectors}>({src})"
+                )
         return expr
 
 
@@ -3733,7 +3836,7 @@ def transform_indexing(self, index: sympy.Expr) -> sympy.Expr:
         )
 
 
-def get_loop_body_lowp_fp(_body: LoopBody) -> tuple[Optional[torch.dtype], bool]:
+def get_loop_body_lowp_fp(_body: LoopBody) -> tuple[torch.dtype | None, bool]:
     """
     Returns the low precision data type (torch.float16/torch.bfloat16) contained in the nodes
     and if all the nodes can codegen with this data type without converting to float.
@@ -3741,7 +3844,7 @@ def get_loop_body_lowp_fp(_body: LoopBody) -> tuple[Optional[torch.dtype], bool]
     """
     sub_blocks = [_body.root_block] + list(_body.subblocks.values())
 
-    _lowp_fp_type: Optional[torch.dtype] = None
+    _lowp_fp_type: torch.dtype | None = None
     _use_fp32 = False
     for sub_block in sub_blocks:
         for _node in sub_block.graph.nodes:
@@ -3937,7 +4040,7 @@ def _is_valid_indices(
                     if tiling_indice < 0 or tiling_indice >= len(call_ranges):
                         continue
                     if has_free_symbols(call_ranges):
-                        call_range = V.graph.sizevars.size_hint(
+                        call_range = V.graph.sizevars.optimization_hint(
                             call_ranges[tiling_indice], fallback=0
                         )
                         if call_range < factor_lowp:
@@ -4040,7 +4143,7 @@ def is_lowp_fp_scheduler(self, scheduler_node: SchedulerNode):
 
     def legalize_lowp_fp_dtype_loopbody(self, loop_body: LoopBody):
         def add_to_dtype(sub_graph: torch.fx.Graph):
-            def get_input_dtype(node: torch.fx.Node) -> Optional[torch.dtype]:
+            def get_input_dtype(node: torch.fx.Node) -> torch.dtype | None:
                 """Get input dtype for nodes that may consumes lowp fp dt"""
                 if node.target == "store":
                     return V.graph.get_dtype(node.args[1])  # type: ignore[arg-type]
@@ -4054,7 +4157,7 @@ def get_input_dtype(node: torch.fx.Node) -> Optional[torch.dtype]:
                 else:
                     return None
 
-            def get_output_dtype(node: torch.fx.Node) -> Optional[torch.dtype]:
+            def get_output_dtype(node: torch.fx.Node) -> torch.dtype | None:
                 """Get output dtype for nodes that may produce lowp fp dt"""
                 if node.target == "load":
                     assert len(node.args) == 3
@@ -4523,7 +4626,7 @@ def update_stores_with_parallel_reduction(self):
         for kernel in self.kernels:
             kernel.update_stores_with_parallel_reduction()
 
-    def gen_body(self, code: Optional[BracesBuffer] = None):
+    def gen_body(self, code: BracesBuffer | None = None):
         assert code is not None
         if_prefix = "C10_LIKELY"
         for kernel in self.kernels:
@@ -4779,7 +4882,7 @@ def get_indexing_ranges_exprs(node):
             else:
                 return FusedSchedulerNode.fuse(node1, node2)
 
-    def _why_fuse_nodes(self, node1, node2) -> Optional[ReasonFusedNodes]:
+    def _why_fuse_nodes(self, node1, node2) -> ReasonFusedNodes | None:
         _, (vars1, reduce1) = node1.group
         _, (vars2, reduce2) = node2.group
 
@@ -4849,10 +4952,13 @@ def _can_fuse_nodes_with_compatible_ranges(self, node1, node2):
         return True
 
     def _can_fuse_horizontal_impl(self, node1, node2):
-        assert isinstance(node1, (FusedSchedulerNode, SchedulerNode))
+        assert isinstance(
+            node1, (FusedSchedulerNode, SchedulerNode, ExternKernelSchedulerNode)
+        )
         assert isinstance(node2, (FusedSchedulerNode, SchedulerNode))
         if any(
-            isinstance(node, OuterLoopFusedSchedulerNode) for node in (node1, node2)
+            isinstance(node, (OuterLoopFusedSchedulerNode, ExternKernelSchedulerNode))
+            for node in (node1, node2)
         ):
             return False
         return self._why_fuse_nodes(node1, node2) is not None
@@ -5304,7 +5410,7 @@ def try_share_local_buffer(local_buffer_layout, local_buffers):
 
     def codegen_node(
         self,
-        node: Union[OuterLoopFusedSchedulerNode, FusedSchedulerNode, SchedulerNode],
+        node: OuterLoopFusedSchedulerNode | FusedSchedulerNode | SchedulerNode,
     ):
         """
         Turn an set of pre-fused nodes into a C++ kernel.
@@ -5357,9 +5463,7 @@ def codegen_template(
         _, (_, rnumel) = template_node.group
         assert rnumel == ()
         ctb: ir.CppTemplateBuffer = cast(ir.CppTemplateBuffer, template_node.node)
-        epilogue_ir_nodes: list[Optional[ir.Operation]] = [
-            n.node for n in epilogue_nodes
-        ]
+        epilogue_ir_nodes: list[ir.Operation | None] = [n.node for n in epilogue_nodes]
         assert all(isinstance(n, ir.ComputedBuffer) for n in epilogue_ir_nodes), (
             "Epilogue nodes must all be instances of ir.ComputedBuffer"
         )
@@ -5597,7 +5701,19 @@ def parallel(self, threads):
         if not self.in_parallel:
             self.num_threads = threads
             self.in_parallel = True
-            if config.cpp.dynamic_threads:
+            # Decide whether to use dynamic threading
+            use_dynamic = False
+            if config.cpp.threads >= 1:
+                # User explicitly set config.cpp.threads (hardcode it)
+                use_dynamic = False
+            elif threads == os.cpu_count():
+                # Thread count matches system CPU count (most likely default, use dynamic)
+                use_dynamic = True
+            else:
+                # Thread count differs from system (user probably set it so hardcode)
+                use_dynamic = False
+
+            if use_dynamic or config.cpp.dynamic_threads:
                 self.code.writeline("#pragma omp parallel")
             else:
                 self.code.writeline(f"#pragma omp parallel num_threads({threads})")
@@ -5625,8 +5741,8 @@ def __exit__(self, exc_type, exc_val, exc_tb):
 
 @dataclasses.dataclass
 class LoopLevel:
-    var: Optional[sympy.Expr] = None
-    size: Optional[sympy.Expr] = None
+    var: sympy.Expr | None = None
+    size: sympy.Expr | None = None
     offset: sympy.Expr = sympy.S.Zero
     # Note [tiled_size]
     # We may do loop-tiling at this loop level.
@@ -5720,8 +5836,8 @@ class LoopNest:
     2D tiling at both the innermost and outer levels.
     """
 
-    loops: Optional[list[LoopLevel]] = None
-    kernel: Optional[CppKernel] = None
+    loops: list[LoopLevel] | None = None
+    kernel: CppKernel | None = None
 
     @staticmethod
     def build(kernel: CppKernel):
@@ -5731,7 +5847,7 @@ def build(kernel: CppKernel):
         reduction_depth = kernel.reduction_depth
         assert reduction_depth is not None
 
-        loops: Optional[list[LoopLevel]] = None
+        loops: list[LoopLevel] | None = None
         for loop_idx, (var, size) in enumerate(zip(itervars, ranges)):
             loop = LoopLevel(var, size)
             if not loops:
diff --git a/torch/_inductor/codegen/cpp_bmm_template.py b/torch/_inductor/codegen/cpp_bmm_template.py
index f4a7c2ef16406..e516a3ac31146 100644
--- a/torch/_inductor/codegen/cpp_bmm_template.py
+++ b/torch/_inductor/codegen/cpp_bmm_template.py
@@ -2,7 +2,7 @@
 import contextlib
 import itertools
 from collections.abc import Callable
-from typing import Any, Optional
+from typing import Any
 from unittest.mock import patch
 
 import sympy
@@ -49,7 +49,12 @@
     constexpr int64_t num_threads = {{num_threads}};
     int64_t B_single_thread_block = (B / num_threads) * num_threads;
 
+    {%- set use_dynamic_threads = ((config.cpp.threads < 1) and (num_threads == cpu_count)) or config.cpp.dynamic_threads %}
+    {%- if use_dynamic_threads %}
+    #pragma omp parallel for
+    {%- else %}
     #pragma omp parallel for num_threads({{num_threads}})
+    {%- endif %}
     {%- else %}
     int64_t B_single_thread_block = B;
     {%- endif %}
@@ -83,7 +88,7 @@ def __init__(
         beta=1,
         alpha=1,
         has_bias=False,
-        epilogue_creator: Optional[Callable[[ir.Buffer], ir.Pointwise]] = None,
+        epilogue_creator: Callable[[ir.Buffer], ir.Pointwise] | None = None,
         should_block_weights: bool = False,
         name="bmm",
     ):
@@ -176,9 +181,9 @@ def reindexer(args):
     def get_options(
         self,
         kernel: CppTemplateKernel,
-        template_buffer_node: Optional[ir.CppTemplateBuffer] = None,
-        flag_template_buffer_has_other_users: Optional[bool] = None,
-        epilogue_nodes: Optional[list[ir.IRNode]] = None,
+        template_buffer_node: ir.CppTemplateBuffer | None = None,
+        flag_template_buffer_has_other_users: bool | None = None,
+        epilogue_nodes: list[ir.IRNode] | None = None,
         **kwargs,
     ) -> dict[str, Any]:
         options = super().get_options(
@@ -210,9 +215,9 @@ def get_options(
     def render(  # type: ignore[override, return]
         self,
         kernel: CppTemplateKernel,
-        template_buffer_node: Optional[ir.CppTemplateBuffer] = None,
-        flag_template_buffer_has_other_users: Optional[bool] = None,
-        epilogue_nodes: Optional[list[ir.IRNode]] = None,
+        template_buffer_node: ir.CppTemplateBuffer | None = None,
+        flag_template_buffer_has_other_users: bool | None = None,
+        epilogue_nodes: list[ir.IRNode] | None = None,
         **kwargs,
     ) -> str:
         options = self.get_options(
diff --git a/torch/_inductor/codegen/cpp_flex_attention_template.py b/torch/_inductor/codegen/cpp_flex_attention_template.py
index 2b417097b3b0e..30104954bb188 100644
--- a/torch/_inductor/codegen/cpp_flex_attention_template.py
+++ b/torch/_inductor/codegen/cpp_flex_attention_template.py
@@ -2,7 +2,6 @@
 import contextlib
 import logging
 import re
-from typing import Optional
 from unittest.mock import patch
 
 import sympy
@@ -963,8 +962,8 @@ def apply_score_mod(self, score, b, h, q_idx, kv_idx):
     def render(  # type: ignore[override,return]
         self,
         kernel,
-        template_buffer_node: Optional[ir.CppTemplateBuffer] = None,
-        epilogue_nodes: Optional[list[ir.IRNode]] = None,
+        template_buffer_node: ir.CppTemplateBuffer | None = None,
+        epilogue_nodes: list[ir.IRNode] | None = None,
         **kwargs,
     ) -> str:
         if epilogue_nodes is not None and epilogue_nodes != []:
diff --git a/torch/_inductor/codegen/cpp_gemm_template.py b/torch/_inductor/codegen/cpp_gemm_template.py
index 20b9c19923062..911d566a1f987 100644
--- a/torch/_inductor/codegen/cpp_gemm_template.py
+++ b/torch/_inductor/codegen/cpp_gemm_template.py
@@ -2,9 +2,10 @@
 import contextlib
 import logging
 import math
+import os
 from collections.abc import Callable
 from functools import lru_cache
-from typing import Any, cast, Optional, TypeVar, Union
+from typing import Any, cast, TypeVar
 from unittest.mock import patch
 
 import torch
@@ -207,7 +208,12 @@
 {%- endif %}
 
 {%- if num_threads > 1 %}
+    {%- set use_dynamic_threads = ((config.cpp.threads < 1) and (num_threads == cpu_count)) or config.cpp.dynamic_threads %}
+    {%- if use_dynamic_threads %}
+    #pragma omp parallel
+    {%- else %}
     #pragma omp parallel num_threads({{num_threads}})
+    {%- endif %}
     {
         {{ template.codegen_multi_threads_params()|indent(8, false) }}
 {%- else %}
@@ -409,7 +415,7 @@ def transpose_w(W: _T, trans_w: bool) -> _T:
     return W
 
 
-def expand_bias(B: Optional[_T], X: _T) -> Optional[_T]:
+def expand_bias(B: _T | None, X: _T) -> _T | None:
     """
     Expand Bias to the same size of X.
     """
@@ -504,11 +510,11 @@ def gen_2d_view_of_epilogue_buf(
     Y: ir.Buffer,
     template_buffer: ir.Buffer,
     epilogue_nodes: list[ir.IRNode],
-    reindexers: list[Optional[Callable[[list[Any]], list[Any]]]],
-    default_reindexers: list[Optional[Callable[[list[Any]], list[Any]]]],
+    reindexers: list[Callable[[list[Any]], list[Any]] | None],
+    default_reindexers: list[Callable[[list[Any]], list[Any]] | None],
 ) -> tuple[
-    Union[ir.Buffer, ir.ReinterpretView],
-    list[Optional[Callable[[list[Any]], list[Any]]]],
+    ir.Buffer | ir.ReinterpretView,
+    list[Callable[[list[Any]], list[Any]] | None],
 ]:
     """
     The dimension and the indexing could be different between the GEMM output, i.e. `template_buffer`, which is
@@ -518,7 +524,7 @@ def gen_2d_view_of_epilogue_buf(
     In this function, we return a 2D buffer (`Y_2d`) according to GEMM output (reinterpreted from `Y` if needed) and
     build a reindexer that converts the indexing of `Y` into `Y_2d`.
     """
-    Y_2d: Union[ir.Buffer, ir.ReinterpretView] = Y
+    Y_2d: ir.Buffer | ir.ReinterpretView = Y
     if (
         Y.get_size() == template_buffer.get_size()
         and Y.get_stride() == template_buffer.get_stride()
@@ -535,7 +541,7 @@ def get_reindexer(epilogue_node, default_reindexer=None):
             #       size (1, 18, 18, 512), stride (165888, 9216, 512, 1)
             stride_order = list(
                 ir.get_stride_order(
-                    V.graph.sizevars.size_hints(epilogue_node.get_stride())
+                    V.graph.sizevars.guarding_hints_or_throw(epilogue_node.get_stride())
                 )
             )
             fill_order = ir.stride_order2fill_order(stride_order)
@@ -598,7 +604,7 @@ def __init__(
         beta=1,
         alpha=1,
         has_bias=False,
-        epilogue_creator: Optional[Callable[[ir.Buffer], ir.Pointwise]] = None,
+        epilogue_creator: Callable[[ir.Buffer], ir.Pointwise] | None = None,
         should_block_weights: bool = True,
         name="packed_gemm",
     ) -> None:
@@ -927,8 +933,8 @@ def add_choices(
         has_bias=False,
         trans_w=False,
         input_indices=None,
-        epilogue_creator: Optional[Callable[[ir.Buffer], ir.Pointwise]] = None,
-        act_mapping: Optional[dict[int, ir.IRNode]] = None,
+        epilogue_creator: Callable[[ir.Buffer], ir.Pointwise] | None = None,
+        act_mapping: dict[int, ir.IRNode] | None = None,
     ):
         """
         Add choices for the GEMM template.
@@ -994,8 +1000,10 @@ def normalize_shapes(inputs, layout_or_out):
                 if has_free_symbols(view_size):
                     # If batch size B is dynamic, we need to set the batch size and possibly stride
                     assert not has_free_symbols(view_size[1:])
-                    view_size[:] = V.graph.sizevars.size_hints(view_size)
-                    view_stride[:] = V.graph.sizevars.size_hints(view_stride)
+                    view_size[:] = V.graph.sizevars.guarding_hints_or_throw(view_size)
+                    view_stride[:] = V.graph.sizevars.guarding_hints_or_throw(
+                        view_stride
+                    )
                 # With the assumptation that W is the storage of unwrap view
                 # thus view it back here
                 new_inputs[1] = new_inputs[1].as_strided(
@@ -1351,9 +1359,9 @@ def get_default_reindexers(self, epilogue_nodes):
     def get_options(
         self,
         kernel: CppTemplateKernel,
-        template_buffer_node: Optional[ir.CppTemplateBuffer] = None,
-        flag_template_buffer_has_other_users: Optional[bool] = None,
-        epilogue_nodes: Optional[list[ir.IRNode]] = None,
+        template_buffer_node: ir.CppTemplateBuffer | None = None,
+        flag_template_buffer_has_other_users: bool | None = None,
+        epilogue_nodes: list[ir.IRNode] | None = None,
     ) -> dict[str, Any]:
         assert len(self.input_nodes) >= 2
 
@@ -1398,7 +1406,7 @@ def get_options(
         gemm_output_buffer = template_buffer
 
         epilogues: list[ir.IRNode] = []
-        reindexers: list[Optional[Callable[[list[Any]], list[Any]]]] = []
+        reindexers: list[Callable[[list[Any]], list[Any]] | None] = []
         epilogue_creators: list[Callable[[ir.Buffer], ir.Pointwise]] = []
         fake_buffers: list[ir.Buffer] = []
         Y_aliases: OrderedSet[str] = OrderedSet()
@@ -1514,7 +1522,7 @@ def copy_inner(index):
                     )
 
         assert isinstance(Y, (ir.Buffer, ir.ReinterpretView))
-        Y_2d: Union[ir.Buffer, ir.ReinterpretView] = Y
+        Y_2d: ir.Buffer | ir.ReinterpretView = Y
 
         if epilogue_nodes:
             if not template_buffer_has_other_users:
@@ -1604,6 +1612,7 @@ def copy_inner(index):
             is_woq_int4=self.is_woq_int4(),
             q_group_size=q_group_size_node,
             qscale_and_zeros=qscale_and_zeros,
+            cpu_count=os.cpu_count(),
         )
         return options
 
@@ -1628,9 +1637,9 @@ def is_int8_woq_gemm_small_m_dim(
     def render(  # type: ignore[override, return]
         self,
         kernel: CppTemplateKernel,
-        template_buffer_node: Optional[ir.CppTemplateBuffer] = None,
-        flag_template_buffer_has_other_users: Optional[bool] = None,
-        epilogue_nodes: Optional[list[ir.IRNode]] = None,
+        template_buffer_node: ir.CppTemplateBuffer | None = None,
+        flag_template_buffer_has_other_users: bool | None = None,
+        epilogue_nodes: list[ir.IRNode] | None = None,
         **kwargs,
     ) -> str:
         options = self.get_options(
diff --git a/torch/_inductor/codegen/cpp_grouped_gemm_template.py b/torch/_inductor/codegen/cpp_grouped_gemm_template.py
index 9dc4255b10a94..d5fca03d48929 100644
--- a/torch/_inductor/codegen/cpp_grouped_gemm_template.py
+++ b/torch/_inductor/codegen/cpp_grouped_gemm_template.py
@@ -1,7 +1,8 @@
 import contextlib
 import logging
+import os
 from collections.abc import Callable
-from typing import Any, cast, Optional, TypeVar
+from typing import Any, cast, TypeVar
 from unittest.mock import patch
 
 import torch
@@ -46,7 +47,12 @@
         num_threads, N, K, micro_gemm, is_dynamic_M, kernel, GemmOuts[0], config, L1_cache_size, L2_cache_size, X_list[0], W_list[0]
     ) }}
 {%- if num_threads > 1 %}
+    {%- set use_dynamic_threads = ((config.cpp.threads < 1) and (num_threads == cpu_count)) or config.cpp.dynamic_threads %}
+    {%- if use_dynamic_threads %}
+    #pragma omp parallel
+    {%- else %}
     #pragma omp parallel num_threads({{num_threads}})
+    {%- endif %}
     {
         {{ template.codegen_multi_threads_params()|indent(8, false) }}
 {%- else %}
@@ -160,8 +166,8 @@ def __init__(
         beta: int = 1,
         alpha: int = 1,
         has_bias: bool = False,
-        epilogue_creator: Optional[Callable[[ir.Buffer], ir.Pointwise]] = None,
-        act_mapping: Optional[dict[int, ir.IRNode]] = None,
+        epilogue_creator: Callable[[ir.Buffer], ir.Pointwise] | None = None,
+        act_mapping: dict[int, ir.IRNode] | None = None,
         gemm_grouped_num: int = 1,
     ) -> None:
         """
@@ -201,9 +207,9 @@ def add_choices(
         alpha: int = 1,
         has_bias: tuple[bool, ...] = (False, False),
         trans_w: bool = False,
-        input_indices: Optional[list[int]] = None,
-        epilogue_creator: Optional[Callable[[ir.Buffer], ir.Pointwise]] = None,
-        act_mapping: Optional[dict[int, ir.IRNode]] = None,  # gemm idx to its act buf
+        input_indices: list[int] | None = None,
+        epilogue_creator: Callable[[ir.Buffer], ir.Pointwise] | None = None,
+        act_mapping: dict[int, ir.IRNode] | None = None,  # gemm idx to its act buf
     ) -> DataProcessorTemplateWrapper:
         # Input nodes order: x, optional[x1], ... w0, w1, ... optional[b0], optional[b1], ...
         gemm_grouped_num = len(has_bias)
@@ -356,9 +362,9 @@ def postprocessor(output: _T) -> _T:
     def render(  # type: ignore[override,return,no-untyped-def]
         self,
         kernel: CppTemplateKernel,
-        template_buffer_node: Optional[ir.CppTemplateBuffer] = None,
-        flag_template_buffer_has_other_users: Optional[bool] = None,
-        epilogue_nodes: Optional[list[ir.IRNode]] = None,
+        template_buffer_node: ir.CppTemplateBuffer | None = None,
+        flag_template_buffer_has_other_users: bool | None = None,
+        epilogue_nodes: list[ir.IRNode] | None = None,
         **kwargs,
     ) -> str:
         assert self.act_mapping
@@ -423,7 +429,7 @@ def render(  # type: ignore[override,return,no-untyped-def]
         assert L2_cache_size > 0, f"Expect L2_cache_size > 0 but got {L2_cache_size}"
 
         epilogues: list[ir.IRNode] = []
-        reindexers: list[Optional[Callable[[list[Any]], list[Any]]]] = []
+        reindexers: list[Callable[[list[Any]], list[Any]] | None] = []
         gemm_output_buffers: list[ir.Buffer] = []
         for out_buf_idx in range(self.gemm_grouped_num):
             gemm_output_name = f"{template_buffer.get_name()}_GemmOut" + str(
@@ -437,7 +443,7 @@ def render(  # type: ignore[override,return,no-untyped-def]
             "epilogue_creator is not supported yet in Grouped GEMM Template"
         )
 
-        kernel_args: dict[str, Optional[ir.IRNode]] = {}
+        kernel_args: dict[str, ir.IRNode | None] = {}
         for x_idx in range(wgt_start_idx):
             kernel_args["X" + str(x_idx)] = act_deduplicated[x_idx]
         for w_idx in range(self.gemm_grouped_num):
@@ -506,6 +512,7 @@ def _bias_add_epilogue(buf: ir.IRNode, inp: ir.IRNode) -> ir.Pointwise:
             Y_list={"Y" + str(idx): Y for idx, Y in enumerate(Y_list)},
             Y_2d_list=Y_2d_list,
             multi_output_buffers=multi_output_buffers,
+            cpu_count=os.cpu_count(),
         )
         with contextlib.ExitStack() as stack:
             stack.enter_context(
diff --git a/torch/_inductor/codegen/cpp_micro_gemm.py b/torch/_inductor/codegen/cpp_micro_gemm.py
index 8289d1afb2646..a6ba0d93f22f6 100644
--- a/torch/_inductor/codegen/cpp_micro_gemm.py
+++ b/torch/_inductor/codegen/cpp_micro_gemm.py
@@ -4,7 +4,6 @@
 import sys
 from collections.abc import Callable
 from enum import Enum
-from typing import Optional
 
 import torch
 
@@ -233,7 +232,7 @@ class CppMicroGemmConfig:
     compute_dtype: torch.dtype
     vec_isa_cls: type[VecISA]
     register_blocking: GemmBlocking
-    extra_check: Optional[Callable[..., bool]] = None
+    extra_check: Callable[..., bool] | None = None
 
 
 micro_gemm_configs: dict[type[CppMicroGemm], list[CppMicroGemmConfig]] = {}
@@ -2106,7 +2105,7 @@ def create_micro_gemm(
     num_threads=-1,
     use_ref=True,
     q_group_size=None,
-) -> Optional[CppMicroGemm]:
+) -> CppMicroGemm | None:
     """
     Based on the provided info, try to find the config of the micro-kernel that would
     deliver the best performance in terms of lower latency for this case.
diff --git a/torch/_inductor/codegen/cpp_template.py b/torch/_inductor/codegen/cpp_template.py
index 736c3b87dee09..5b71d24a81572 100644
--- a/torch/_inductor/codegen/cpp_template.py
+++ b/torch/_inductor/codegen/cpp_template.py
@@ -5,7 +5,6 @@
 import logging
 import sys
 from collections.abc import Callable, Iterable
-from typing import Optional, Union
 from unittest.mock import patch
 
 import sympy
@@ -30,12 +29,12 @@ def __init__(
         input_nodes,
         layout: ir.Layout,
         num_threads: int,
-        epilogue_creator: Optional[Callable[[ir.Buffer], ir.Pointwise]] = None,
+        epilogue_creator: Callable[[ir.Buffer], ir.Pointwise] | None = None,
     ) -> None:
         super().__init__(name)
         self.input_nodes = input_nodes
         self.index = next(self.index_counter)
-        self.output_node: Union[ir.Buffer, list[ir.Buffer]] = ir.Buffer(
+        self.output_node: ir.Buffer | list[ir.Buffer] = ir.Buffer(
             name=f"buf_out{self.index}", layout=layout
         )
         self.layout = layout
@@ -72,7 +71,8 @@ def generate(self, **kwargs):
             call_args,
             expected_args,
         )
-        extra_args = V.graph.sizevars.size_hints(
+        # extra_args are only used for benchmarking, not compiled kernel correctness
+        extra_args = V.graph.sizevars.optimization_hints(
             map(sympy.expand, call_args[len(expected_args) :])
         )
         # Cast the size hint from int to ctypes.c_ulonglong explicitly
@@ -93,7 +93,7 @@ def generate(self, **kwargs):
         def make_kernel_render(
             template_node: ir.CppTemplateBuffer,
             flag_template_buffer_has_other_users: bool,
-            epilogue_nodes: Optional[list[ir.IRNode]] = None,
+            epilogue_nodes: list[ir.IRNode] | None = None,
         ):
             kernel = CppTemplateKernel(
                 kernel_name=str(Placeholder.KERNEL_NAME), num_threads=self.num_threads
diff --git a/torch/_inductor/codegen/cpp_template_kernel.py b/torch/_inductor/codegen/cpp_template_kernel.py
index ae2536e2efbb6..836d50b5984a6 100644
--- a/torch/_inductor/codegen/cpp_template_kernel.py
+++ b/torch/_inductor/codegen/cpp_template_kernel.py
@@ -1,7 +1,7 @@
 # mypy: allow-untyped-defs
 import itertools
 from collections.abc import Callable, Iterable
-from typing import Any, Optional, Union
+from typing import Any
 from unittest.mock import patch
 
 import sympy
@@ -56,9 +56,9 @@ def def_kernel(
         self,
         inputs: dict[str, ir.Buffer],
         outputs: dict[str, ir.Buffer],
-        aliases: Optional[dict[str, str]] = None,
+        aliases: dict[str, str] | None = None,
         function_name: str = "",
-        extra_sizevars: Optional[list[sympy.Expr]] = None,
+        extra_sizevars: list[sympy.Expr] | None = None,
         placeholder: str = "<DEF_KERNEL>",
     ) -> str:
         if len(function_name) == 0:
@@ -245,8 +245,8 @@ def store_pointwise_nodes(
         self,
         dst: ir.Buffer,
         nodes: list[ir.IRNode],
-        offsets: Optional[list[sympy.Expr]] = None,
-        reindexers: Optional[list[Optional[Callable[[list[Any]], list[Any]]]]] = None,
+        offsets: list[sympy.Expr] | None = None,
+        reindexers: list[Callable[[list[Any]], list[Any]] | None] | None = None,
     ) -> str:
         var_sizes = (tuple(dst.get_size()), ())
         var_ranges = {
@@ -309,7 +309,7 @@ def store_grouped_gemm_pointwise_nodes(
         dst: tuple[ir.Buffer],
         nodes: list[ir.IRNode],
         offsets: list[sympy.Expr],
-        reindexers: list[Optional[Callable[[list[Any]], list[Any]]]],
+        reindexers: list[Callable[[list[Any]], list[Any]] | None],
         output_names: list[str],
     ) -> str:
         ref_dst = dst[0]
@@ -370,10 +370,10 @@ def store_output(
         self,
         dst: ir.Buffer,
         src: ir.Buffer,
-        orig_src: Optional[ir.Buffer] = None,
-        epilogue_nodes: Optional[list[ir.IRNode]] = None,
-        offsets: Optional[list[Any]] = None,
-        reindexers: Optional[list[Optional[Callable[[list[Any]], list[Any]]]]] = None,
+        orig_src: ir.Buffer | None = None,
+        epilogue_nodes: list[ir.IRNode] | None = None,
+        offsets: list[Any] | None = None,
+        reindexers: list[Callable[[list[Any]], list[Any]] | None] | None = None,
     ):
         """
         Store the `src` buffer to the `dst` buffer. The size of `src` and `dst` should match.
@@ -432,11 +432,11 @@ def store_outputs(
         self,
         dst: tuple[ir.Buffer],
         src: tuple[ir.IRNode],
-        orig_src: Optional[tuple[ir.IRNode]] = None,
-        epilogue_nodes: Optional[list[ir.IRNode]] = None,
-        offsets: Optional[list[Any]] = None,
-        reindexers: Optional[list[Optional[Callable[[list[Any]], list[Any]]]]] = None,
-        multi_output_buffers: Optional[tuple[ir.MultiOutput, ...]] = None,
+        orig_src: tuple[ir.IRNode] | None = None,
+        epilogue_nodes: list[ir.IRNode] | None = None,
+        offsets: list[Any] | None = None,
+        reindexers: list[Callable[[list[Any]], list[Any]] | None] | None = None,
+        multi_output_buffers: tuple[ir.MultiOutput, ...] | None = None,
     ):
         assert isinstance(dst, Iterable)
         assert all(_dst.get_size() == _src.get_size() for _src, _dst in zip(src, dst))
@@ -567,15 +567,14 @@ def __init__(
             [
                 ir.CppTemplateBuffer,
                 bool,
-                Optional[list[ir.IRNode]],
+                list[ir.IRNode] | None,
             ],
             str,
         ],
         bmreq: CppBenchmarkRequest,
         template: "CppTemplate",  # type: ignore[name-defined]  # noqa: F821
-        info_kwargs: Optional[
-            dict[str, Union[ir.PrimitiveInfoType, list[ir.PrimitiveInfoType]]]
-        ] = None,
+        info_kwargs: dict[str, ir.PrimitiveInfoType | list[ir.PrimitiveInfoType]]
+        | None = None,
     ):
         super().__init__(name, input_nodes, layout, description="")
         self.category = category
@@ -605,7 +604,7 @@ def hash_key(self) -> str:
 
     def info_dict(
         self,
-    ) -> dict[str, Union[ir.PrimitiveInfoType, list[ir.PrimitiveInfoType]]]:
+    ) -> dict[str, ir.PrimitiveInfoType | list[ir.PrimitiveInfoType]]:
         return {"backend": "CPP", "op_type": "unknown"}
 
     def output_node(self) -> ir.TensorBox:
diff --git a/torch/_inductor/codegen/cpp_utils.py b/torch/_inductor/codegen/cpp_utils.py
index dfa22c69ba37c..e5299897ead7e 100644
--- a/torch/_inductor/codegen/cpp_utils.py
+++ b/torch/_inductor/codegen/cpp_utils.py
@@ -6,7 +6,7 @@
 import sys
 from collections import namedtuple
 from collections.abc import Callable, Sequence
-from typing import Any, Optional
+from typing import Any
 from unittest.mock import patch
 
 import sympy
@@ -104,6 +104,7 @@ def get_promote_dtype(args):
         # pyrefly: ignore [no-matching-overload]
         functools.reduce(
             torch.promote_types,  # type: ignore[arg-type]
+            # pyrefly: ignore [bad-argument-type]
             [n.dtype for n in args if isinstance(n, CppCSEVariable)],
         )
         if all(n.dtype is not None for n in args if isinstance(n, CppCSEVariable))
@@ -146,7 +147,7 @@ def __init__(
         self,
         name,
         bounds: ValueRanges[Any],
-        dtype: Optional[torch.dtype] = None,
+        dtype: torch.dtype | None = None,
         shape: BlockShapeType = None,
     ) -> None:
         super().__init__(name, bounds, dtype, shape=shape)
@@ -378,7 +379,7 @@ def __exit__(self, exc_type, exc_val, exc_tb):
         self.exit_stack.__exit__(exc_type, exc_val, exc_tb)
 
     def add_local_buffer(
-        self, local_buffer: ir.Buffer, global_buffers: Optional[list[ir.Buffer]] = None
+        self, local_buffer: ir.Buffer, global_buffers: list[ir.Buffer] | None = None
     ):
         assert local_buffer.get_name() not in self.local_buffers
         self.local_buffers[local_buffer.get_name()] = local_buffer
diff --git a/torch/_inductor/codegen/cpp_wrapper_cpu.py b/torch/_inductor/codegen/cpp_wrapper_cpu.py
index b04baf00eb8c6..dc3e3c94d5809 100644
--- a/torch/_inductor/codegen/cpp_wrapper_cpu.py
+++ b/torch/_inductor/codegen/cpp_wrapper_cpu.py
@@ -8,7 +8,7 @@
 import sys
 import textwrap
 from itertools import chain, count
-from typing import Any, Optional, Protocol, TYPE_CHECKING, Union
+from typing import Any, Protocol, TYPE_CHECKING
 
 import sympy
 
@@ -19,6 +19,7 @@
 from torch._inductor.runtime.runtime_utils import dynamo_timed
 from torch.fx.experimental.symbolic_shapes import ConvertIntKey, DivideByKey, SymTypes
 from torch.utils._ordered_set import OrderedSet
+from torch.utils._sympy.functions import CleanDiv, FloorDiv, Mod, ModularIndexing
 from torch.utils._sympy.symbol import symbol_is_type, SymT
 
 from .. import config, cpp_builder, ir
@@ -43,13 +44,13 @@
     from ..graph import GraphLowering
 
     # At most, the list nesting can go one layer deep.
-    _OUTPUT_ARGS_TYPE = list[Union[Optional[str], list[Optional[str]]]]
+    _OUTPUT_ARGS_TYPE = list[str | None | list[str | None]]
 
     from ..scheduler import BaseSchedulerNode
 
 
 class HasWriteLine(Protocol):
-    def writeline(self, line: Union[LineContext, DeferredLineBase, str]) -> None: ...
+    def writeline(self, line: LineContext | DeferredLineBase | str) -> None: ...
 
 
 class CppWrapperCpu(PythonWrapperCodegen):
@@ -102,9 +103,9 @@ def __init__(self):
     @staticmethod
     def create(
         is_subgraph: bool,
-        subgraph_name: Optional[str],
-        parent_wrapper: Optional[PythonWrapperCodegen],
-        partition_signatures: Optional[ir.GraphPartitionSignature] = None,
+        subgraph_name: str | None,
+        parent_wrapper: PythonWrapperCodegen | None,
+        partition_signatures: ir.GraphPartitionSignature | None = None,
     ):
         # TODO - support subgraph codegen by lifting functions. Check the
         # comment at CppWrapperCpu `codegen_subgraph` function.
@@ -146,6 +147,7 @@ def _generate_kernel_call_helper(
         inductor_meta=None,
         graph_name="",
         original_fxnode_name=None,
+        current_stream_idx=None,
     ):
         """
         Generates kernel call code.
@@ -216,13 +218,13 @@ def write_header(self):
 
         if not V.graph.aot_mode:
             self.header.splice(
-                """
+                '''
                 import torch
                 from torch._inductor.codecache import CppWrapperCodeCache
 
                 cpp_wrapper_src = (
-                r'''
-                """
+                r"""
+                '''
             )
 
         for device in V.graph.device_types:
@@ -321,7 +323,7 @@ def strideof(name):
             return f"{name}_stride"
 
         def codegen_symbol(
-            sym_or_exp: Union[sympy.Symbol, sympy.Expr],
+            sym_or_exp: sympy.Symbol | sympy.Expr,
             base_name: str,
             name_fn: Callable[[str], str],
             dim: int,
@@ -994,7 +996,7 @@ def codegen_const_run_driver(self):
 
         with self.prefix.indent():
             # This is a mapping to the index of constant folding graph's output
-            const_index_mapping: list[Optional[tuple[int, str]]] = [None] * len(
+            const_index_mapping: list[tuple[int, str] | None] = [None] * len(
                 V.graph.const_output_index
             )
             for idx, (name, _) in enumerate(V.graph.constants.items()):
@@ -1069,9 +1071,9 @@ def _define_kernel_helper(
         self,
         kernel_name: str,
         kernel_body: str,
-        metadata: Optional[str] = None,
+        metadata: str | None = None,
         gpu: bool = False,
-        cpp_definition: Optional[str] = None,
+        cpp_definition: str | None = None,
     ):
         if cpp_definition is not None:
             self.header.splice(cpp_definition)
@@ -1174,11 +1176,11 @@ def generate_end(self, result):
 
         if config.cpp_wrapper_build_separate:
             # Close the wrapper code block, then write any kernel definitions.
-            result.splice("'''\n)")
+            result.splice('"""\n)')
             if self.kernel_declarations:
-                result.splice("\nkernel_src = (\nr'''")
+                result.splice('\nkernel_src = (\nr"""')
                 result.splice(self.kernel_declarations.getvalue())
-                result.splice("'''\n)")
+                result.splice('"""\n)')
             else:
                 result.splice(
                     """
@@ -1190,7 +1192,7 @@ def generate_end(self, result):
             result.splice(self.kernel_declarations.getvalue())
             self.kernel_declarations.clear()
             # Close the wrapper code block
-            result.splice("'''\n)")
+            result.splice('"""\n)')
 
         kernel_code = "kernel_src" if config.cpp_wrapper_build_separate else "None"
         # Cpp entry function for JIT with cpp wrapper
@@ -1283,8 +1285,8 @@ def generate_c_shim_extern_kernel_call(
         args: list[str],
         device: str,
         *,
-        debug_args: Optional[list[str]] = None,
-        stack_traces: Optional[OrderedSet[str]] = None,
+        debug_args: list[str] | None = None,
+        stack_traces: OrderedSet[str] | None = None,
     ) -> None:
         """debug_args kwarg allows CppWrapperCpuArrayRef to pass in wrapped arguments in
         place of args while preserving debug printer output."""
@@ -1410,10 +1412,10 @@ def _generate_extern_kernel_out_helper(
         self,
         kernel: str,
         out: str,
-        out_view: Optional[str],
+        out_view: str | None,
         args: list[str],
         device: str,
-        stack_traces: Optional[OrderedSet[str]] = None,
+        stack_traces: OrderedSet[str] | None = None,
     ) -> None:
         if out_view:
             out_name = f"{out}_as_strided"
@@ -1491,8 +1493,68 @@ def add_benchmark_harness(self, output):
             return
         super().add_benchmark_harness(output)
 
+    def _extract_divisors_from_expr(
+        self, expr: sympy.Expr
+    ) -> list[tuple[sympy.Expr, str]]:
+        """
+        Walk the sympy expression and extract all divisors/modulos from
+        FloorDiv, CleanDiv, Mod, and ModularIndexing nodes.
+
+        Returns a list of (divisor_expr, operation_name) tuples for divisors
+        that are not constant positive integers (which cannot be zero).
+        """
+        divisors: list[tuple[sympy.Expr, str]] = []
+        seen: OrderedSet[str] = OrderedSet()
+
+        def is_constant_nonzero(e: sympy.Expr) -> bool:
+            """Check if expression is a constant that cannot be zero."""
+            if isinstance(e, sympy.Integer):
+                return int(e) != 0
+            if isinstance(e, sympy.Number):
+                return float(e) != 0
+            return False
+
+        def maybe_add_divisor(divisor: sympy.Expr, op_name: str) -> None:
+            """Add divisor to list if it could potentially be zero."""
+            if not is_constant_nonzero(divisor):
+                key = str(divisor)
+                if key not in seen:
+                    seen.add(key)
+                    divisors.append((divisor, op_name))
+
+        def walk(e: sympy.Expr) -> None:
+            if isinstance(e, (FloorDiv, CleanDiv)):
+                _, div = e.args
+                maybe_add_divisor(div, "floor division")
+            elif isinstance(e, Mod):
+                _, mod = e.args
+                maybe_add_divisor(mod, "modulo")
+            elif isinstance(e, ModularIndexing):
+                _, div, mod = e.args
+                maybe_add_divisor(div, "modular indexing division")
+                maybe_add_divisor(mod, "modular indexing modulo")
+
+            # Recurse into arguments
+            if hasattr(e, "args"):
+                for arg in e.args:
+                    if isinstance(arg, sympy.Expr):
+                        walk(arg)
+
+        walk(expr)
+        return divisors
+
     def codegen_cpp_sizevar(self, x: sympy.Expr, *, simplify: bool = True) -> str:
-        return cexpr(V.graph.sizevars.simplify(x) if simplify else x)
+        maybe_simplified_x = V.graph.sizevars.simplify(x) if simplify else x
+        # In AOT mode, emit runtime checks for potential division/modulo by zero
+        # to prevent SIGFPE crashes when symbolic tensor shapes can be 0
+        if V.graph.aot_mode:
+            divisors = self._extract_divisors_from_expr(maybe_simplified_x)
+            for divisor, op_name in divisors:
+                divisor_str = cexpr(divisor)
+                self.writeline(
+                    f'AOTI_TORCH_CHECK({divisor_str} != 0, "Integer {op_name} by zero");'
+                )
+        return cexpr(maybe_simplified_x)
 
     def codegen_sizevar(self, x: sympy.Expr) -> str:
         return self.codegen_cpp_sizevar(x)
@@ -1520,8 +1582,14 @@ def ensure_size_computed(self, sym: sympy.Symbol):
     def _generate_symbolic_call_arg_helper(
         self, arg: SymbolicCallArg, graph: GraphLowering
     ) -> None:
-        if (arg.inner, graph) not in self.kernel_numel_expr:
-            # declare expr once in each graph (scope)
+        enable_kernel_profile = config.cpp.enable_kernel_profile and sys.platform in [
+            "linux",
+            "win32",
+        ]
+        if enable_kernel_profile or (arg.inner, graph) not in self.kernel_numel_expr:
+            # When enable_kernel_profile is on, each kernel call is wrapped in
+            # its own {} scope block, so we must redeclare the variable each
+            # time since prior declarations are no longer visible.
             self.kernel_numel_expr.add((arg.inner, graph))
             self.writeline(f"int64_t {arg.inner} = {cexpr(arg.inner_expr)};")
         else:
@@ -1952,7 +2020,7 @@ def create_new_tensor_handle() -> tuple[str, list[str]]:
         # ```
         return final_tensor_str
 
-    def codegen_device_copy(self, src, dst, non_blocking: Union[bool, str]):
+    def codegen_device_copy(self, src, dst, non_blocking: bool | str):
         """This function is overridden by cpp_wrapper_cpu_array_ref, so we don't need to
         handle cases where dst is not an AtenTensorHandle."""
         self.writeline(
@@ -2121,7 +2189,7 @@ def codegen_while_loop(self, while_loop, stack_output=False):
 
     def generate_extern_kernel_args_decl_if_needed(
         self,
-        op_overload: Union[torch._ops.OpOverload, torch._ops.HigherOrderOperator],
+        op_overload: torch._ops.OpOverload | torch._ops.HigherOrderOperator,
         raw_args: Sequence[Any],
         output_args: _OUTPUT_ARGS_TYPE,
         raw_outputs: Sequence[ir.Buffer],
@@ -2323,7 +2391,7 @@ def generate_fallback_kernel_with_runtime_lookup(
         buf_name: str,
         python_kernel_name: str,
         get_args: Callable[[], Sequence[str]],
-        op_overload: Union[torch._ops.OpOverload, torch._ops.HigherOrderOperator],
+        op_overload: torch._ops.OpOverload | torch._ops.HigherOrderOperator,
         raw_args: Sequence[Any],
         outputs: Sequence[ir.Buffer],
     ) -> None:
@@ -2331,8 +2399,8 @@ def generate_fallback_kernel_with_runtime_lookup(
         different code paths for AOT Inductor vs cpp_wrapper Inductor mode."""
 
         def extract_output_name(
-            out: Optional[Union[ir.Buffer, Sequence[ir.Buffer]]],
-        ) -> Union[Optional[str], _OUTPUT_ARGS_TYPE]:
+            out: ir.Buffer | Sequence[ir.Buffer] | None,
+        ) -> str | None | _OUTPUT_ARGS_TYPE:
             if out is None:
                 return None
             if isinstance(out, (ir.MultiOutput, ir._CollectiveKernel)):
@@ -2460,6 +2528,9 @@ def generate_float_value(self, val):
             return f"{val}"
 
     def generate_py_arg(self, py_args_var, idx, raw_arg, arg_type):
+        """Generate C++ code that converts a single operator argument into a Python
+        object and inserts it into a PyTuple at the given index."""
+
         def generate_py_arg_inner(lines, raw_arg, arg_type):
             def handle_scalar(scalar):
                 if isinstance(scalar, int):
@@ -2537,19 +2608,31 @@ def handle_scalar(scalar):
                     f"arg type {arg_type} is not yet supported by custom_op_wrapper"
                 )
 
-        lines = []
-        if isinstance(arg_type, torch.ListType):
-            assert isinstance(raw_arg, (list, tuple)), str(raw_arg) + " is not a list"
-            lines.append(
+        def handle_sequence_arg(raw_arg_, arg_type_, lines_):
+            assert isinstance(raw_arg_, (list, tuple)), str(raw_arg) + " is not a list"
+            lines_.append(
                 f"PyObject* {py_args_var}_{idx} = PyList_New({len(raw_arg)});\n"
             )
-            for i, elem in enumerate(raw_arg):
-                lines.append(
-                    f"PyList_SetItem({py_args_var}_{idx}, {i}, {generate_py_arg_inner(lines, elem, arg_type.getElementType())});\n"
+            for i, elem in enumerate(raw_arg_):
+                lines_.append(
+                    f"PyList_SetItem({py_args_var}_{idx}, {i}, {generate_py_arg_inner(lines, elem, arg_type_.getElementType())});\n"
                 )
-            lines.append(
+            lines_.append(
                 f"PyTuple_SetItem({py_args_var}, {idx}, {py_args_var}_{idx});\n"
             )
+
+        lines = []
+        if isinstance(arg_type, torch.ListType):
+            handle_sequence_arg(raw_arg, arg_type, lines)
+        elif isinstance(arg_type, torch.OptionalType) and isinstance(
+            arg_type.getElementType(), torch.ListType
+        ):
+            if raw_arg is None:
+                lines.append(
+                    f"PyTuple_SetItem({py_args_var}, {idx}, {generate_py_arg_inner(lines, raw_arg, arg_type)});\n"
+                )
+            else:
+                handle_sequence_arg(raw_arg, arg_type.getElementType(), lines)
         else:
             lines.append(
                 f"PyTuple_SetItem({py_args_var}, {idx}, {generate_py_arg_inner(lines, raw_arg, arg_type)});\n"
@@ -2560,7 +2643,7 @@ def generate_fallback_kernel_with_runtime_lookup_nopython(
         self,
         get_args: Callable[[], Sequence[str]],
         op_overload: torch._ops.OpOverload,
-        output_args: Sequence[Optional[str]],
+        output_args: Sequence[str | None],
         raw_outputs: Sequence[ir.Buffer],
     ) -> None:
         """Generate fallback kernel calls with runtime (non-AOT) dispatch.  This can
@@ -2672,7 +2755,7 @@ def generate_fallback_kernel_with_runtime_lookup_python(
         python_kernel_name: str,
         op_overload: torch._ops.OpOverload,
         raw_args: Sequence[Any],
-        output_args: Sequence[Optional[str]],
+        output_args: Sequence[str | None],
         raw_outputs: Sequence[ir.Buffer],
     ) -> None:
         """Generate fallback kernel calls with runtime (non-AOT) dispatch.  This can
@@ -2746,7 +2829,7 @@ def generate_fallback_kernel_with_runtime_lookup_python(
 
     def generate_fallback_kernel_with_runtime_lookup_aot(
         self,
-        op_overload: Union[torch._ops.OpOverload, torch._ops.HigherOrderOperator],
+        op_overload: torch._ops.OpOverload | torch._ops.HigherOrderOperator,
         raw_args: Sequence[Any],
         output_args: _OUTPUT_ARGS_TYPE,
         raw_outputs: Sequence[ir.Buffer],
@@ -2940,7 +3023,7 @@ def val_to_arg_str(self, val, type_=None) -> str:
         return self.val_to_arg_str_for_prim_type(val, type_)
 
     def create_tmp_raii_handle_var_if_needed(
-        self, handle: str, writer: Optional[Union[HasWriteLine, list[str]]] = None
+        self, handle: str, writer: HasWriteLine | list[str] | None = None
     ) -> str:
         """If the input handle is an rvalue RAII tensor, creates an lvalue variable for
         it in writer.  Returns a variable name that can be used to access handle."""
@@ -2985,10 +3068,10 @@ def write_kernel_context_guard_end(
     def write_kernel_context_guard(
         self,
         kernel_name: str,
-        node_schedule: Union[Sequence[BaseSchedulerNode], ExternKernel],
+        node_schedule: Sequence[BaseSchedulerNode] | ExternKernel,
     ):
         def aggregate_stack_traces(
-            node_schedule: Union[Sequence[BaseSchedulerNode], ExternKernel],
+            node_schedule: Sequence[BaseSchedulerNode] | ExternKernel,
         ) -> OrderedSet[str]:
             if isinstance(node_schedule, list):
                 return functools.reduce(
diff --git a/torch/_inductor/codegen/cpp_wrapper_cpu_array_ref.py b/torch/_inductor/codegen/cpp_wrapper_cpu_array_ref.py
index 0cf5f95e76b62..86827fb701290 100644
--- a/torch/_inductor/codegen/cpp_wrapper_cpu_array_ref.py
+++ b/torch/_inductor/codegen/cpp_wrapper_cpu_array_ref.py
@@ -1,6 +1,6 @@
 # mypy: allow-untyped-defs
 from collections.abc import Callable, Sequence
-from typing import Any, Optional, Union
+from typing import Any
 
 import sympy
 
@@ -50,9 +50,9 @@ def __init__(self):
     @staticmethod
     def create(
         is_subgraph: bool,
-        subgraph_name: Optional[str],
-        parent_wrapper: Optional[PythonWrapperCodegen],
-        partition_signatures: Optional[ir.GraphPartitionSignature] = None,
+        subgraph_name: str | None,
+        parent_wrapper: PythonWrapperCodegen | None,
+        partition_signatures: ir.GraphPartitionSignature | None = None,
     ):
         # TODO - support subgraph codegen by lifting functions. Check the
         # comment at CppWrapperCpu `codegen_subgraph` function.
@@ -117,6 +117,7 @@ def _generate_kernel_call_helper(
         inductor_meta=None,
         graph_name="",
         original_fxnode_name=None,
+        current_stream_idx=None,
     ):
         """
         Generates kernel call code.
@@ -771,7 +772,7 @@ def generate_fallback_kernel_with_runtime_lookup(
         buf_name: str,
         python_kernel_name: str,
         get_args: Callable[[], Sequence[str]],
-        op_overload: Union[torch._ops.OpOverload, torch._ops.HigherOrderOperator],
+        op_overload: torch._ops.OpOverload | torch._ops.HigherOrderOperator,
         raw_args: Sequence[Any],
         outputs: Sequence[ir.Buffer],
     ) -> None:
@@ -781,7 +782,7 @@ def generate_fallback_kernel_with_runtime_lookup(
             buf_name, python_kernel_name, get_args, op_overload, raw_args, outputs
         )
 
-    def codegen_device_copy(self, src, dst, non_blocking: Union[bool, str]):
+    def codegen_device_copy(self, src, dst, non_blocking: bool | str):
         # aoti_torch_tensor_copy_ takes AtenTensorHandle as input,
         # while stack-allocation results in ArrayRefTensor
         # so disable stack allocation here
diff --git a/torch/_inductor/codegen/cpp_wrapper_gpu.py b/torch/_inductor/codegen/cpp_wrapper_gpu.py
index 459b777be75c6..94a6d2ca2d4c9 100644
--- a/torch/_inductor/codegen/cpp_wrapper_gpu.py
+++ b/torch/_inductor/codegen/cpp_wrapper_gpu.py
@@ -5,7 +5,7 @@
 import re
 import sys
 from itertools import count, zip_longest
-from typing import Any, Optional, Union
+from typing import Any
 from typing_extensions import Self
 
 import sympy
@@ -14,6 +14,7 @@
 from torch import dtype as torch_dtype
 from torch._inductor.codecache import get_cpp_wrapper_cubin_path_name
 from torch._inductor.runtime.runtime_utils import dynamo_timed
+from torch.utils._ordered_set import OrderedSet
 
 from .. import config
 from ..codecache import CudaKernelParamCache
@@ -65,7 +66,7 @@ def cpp_string_literal(s: str) -> str:
 }
 
 
-def signature_is_tma_desc(sig: Optional[str]) -> bool:
+def signature_is_tma_desc(sig: str | None) -> bool:
     """Check if a Triton signature represents a TMA descriptor."""
     if not sig:
         return False
@@ -76,6 +77,24 @@ def signature_is_tma_desc(sig: Optional[str]) -> bool:
     return False
 
 
+def _unpack_tma_descriptor_args(var_name: str, sig_type: str) -> list[str]:
+    """Unpack a StableTMADescriptor into kernel launch args.
+
+    Given a variable name holding a StableTMADescriptor and its tensordesc<...>
+    signature, returns the list of pointer args: &var.m, &var.block_shape[i]...,
+    &var.strides[i]...
+    """
+    match = re.match(r"tensordesc<[^[]*\[([^\]]*)\]", sig_type)
+    assert match is not None, f"Cannot parse tensordesc signature: {sig_type}"
+    ndim = match.group(1).count(",") + 1
+    result = [f"&{var_name}.m"]
+    for i in range(ndim):
+        result.append(f"&{var_name}.block_shape[{i}]")
+    for i in range(ndim):
+        result.append(f"&{var_name}.strides[{i}]")
+    return result
+
+
 @dataclasses.dataclass
 class DeferredTritonCallWrapper:
     """
@@ -88,6 +107,94 @@ class DeferredTritonCallWrapper:
     kernel_name: str
     kernel_name_to_body: dict[str, str]
     arg_types: list[Any]
+    triton_meta: dict[str, Any] | None = None
+    inductor_meta: dict[str, Any] | None = None
+    tma_tensor_args: dict[str, str] = dataclasses.field(default_factory=dict)
+
+    @cache_on_self
+    def _get_tma_args(self) -> dict[str, str]:
+        """Get mapping of TMA descriptor arg names to their signature types."""
+        triton_meta = self.triton_meta or {}
+        signature = triton_meta.get("signature", {})
+        for name, sig_type in signature.items():
+            if sig_type == "nvTmaDesc":
+                raise RuntimeError(
+                    f"nvTmaDesc (experimental TMA API) is not supported in lazy compile "
+                    f"for arg '{name}'. Use the stable tensordesc API instead."
+                )
+        return {
+            name: sig_type
+            for name, sig_type in signature.items()
+            if isinstance(sig_type, str) and sig_type.startswith("tensordesc<")
+        }
+
+    def _get_cpp_param_type(
+        self, name: str, arg_type: Any, signature: dict[str, str] | None = None
+    ) -> str:
+        """Get the C++ parameter declaration for a given arg type."""
+        if isinstance(arg_type, (torch_dtype, UnwrapUnspecArg)):
+            # TMA descriptors need non-const references since their fields
+            # are passed as void* pointers to kernel launch args
+            if signature and signature_is_tma_desc(signature.get(name)):
+                return f"{name}_type_& {name}"
+            return f"const {name}_type_& {name}"
+        elif issubclass(arg_type, (SymbolicCallArg, sympy.Expr, int)):
+            return f"int64_t {name}"
+        elif arg_type is float:
+            return f"float {name}"
+        elif arg_type is bool:
+            return f"bool {name}"
+        else:
+            raise ValueError(f"Unexpected arg type {arg_type}")
+
+    def _write_wrapper_signature(
+        self,
+        prefix: IndentedBuffer,
+        wrapper: CppWrapperGpu,
+        arg_names: list[str],
+        arg_types: list[Any] | None = None,
+        signature: dict[str, str] | None = None,
+    ) -> None:
+        """Write the wrapper function signature including template and parameters."""
+        if arg_types is None:
+            arg_types = self.arg_types
+
+        # Generate template types for tensor arguments
+        template_types = [
+            f"typename {name}_type_"
+            for name, arg_type in zip(arg_names, arg_types)
+            if isinstance(arg_type, (torch_dtype, UnwrapUnspecArg))
+        ]
+        if V.graph.aot_mode:
+            template_types.append("typename kernels_type_")
+
+        if template_types:
+            prefix.writeline(f"template <{', '.join(template_types)}>")
+
+        # Build parameter list
+        param_lines = [
+            self._get_cpp_param_type(name, arg_type, signature)
+            for name, arg_type in zip(arg_names, arg_types)
+        ]
+        param_lines.append("int32_t device_idx_")
+        param_lines.append(
+            maybe_hipify_code_wrapper(
+                f"{wrapper.device_codegen.cpp_stream_type()} stream_"
+            )
+        )
+        if V.graph.aot_mode:
+            param_lines.append("kernels_type_& kernels_")
+        param_lines.append(
+            "const std::optional<std::string>& cubin_dir_ = std::nullopt"
+        )
+
+        # Write function signature
+        prefix.writeline(f"static __attribute__((noinline)) void {self.wrapper_name}(")
+        with prefix.indent():
+            for i, param in enumerate(param_lines):
+                comma = "," if i < len(param_lines) - 1 else ""
+                prefix.writeline(f"{param}{comma}")
+        prefix.writeline("){")
 
     def generate(self, wrapper: CppWrapperGpu):
         """
@@ -97,6 +204,11 @@ def generate(self, wrapper: CppWrapperGpu):
         if self.kernel_name.startswith("multi_kernel_"):
             # MultiKernel will select one kernel after running the autotune block
             self.kernel_name = MultiKernelCall.lookup_choice(self.kernel_name)
+
+        # Defer compilation to runtime if autotune_at_compile_time is False (JIT only)
+        if not V.graph.aot_mode and config.triton.autotune_at_compile_time is False:
+            return self.generate_lazy(wrapper)
+
         params = CudaKernelParamCache.get(self.kernel_name)
         assert params, f"CudaKernelParamCache not populated for {self.kernel_name}"
         def_args = params["def_args"]
@@ -122,42 +234,9 @@ def generate(self, wrapper: CppWrapperGpu):
         else:
             kernel_var_name = f"kernels_.{self.kernel_name}"
 
-        # tensors can be RAIIAtenTensorHandle or ConstantHandle, so make them template types
-        template_types = [
-            f"typename {name}_type_"
-            for name, arg_type in zip(def_args, arg_types)
-            if isinstance(arg_type, (torch_dtype, UnwrapUnspecArg))
-        ]
-        if V.graph.aot_mode:
-            template_types.append("typename kernels_type_")
-        if template_types:
-            prefix.writeline(f"template <{', '.join(template_types)}>")
-        prefix.writeline(f"static inline void {self.wrapper_name}(")
-        with prefix.indent():
-            assert len(def_args) == len(arg_types), (def_args, arg_types)
-            for name, arg_type in zip(def_args, arg_types):
-                if isinstance(arg_type, (torch_dtype, UnwrapUnspecArg)):
-                    prefix.writeline(f"const {name}_type_& {name},")
-                elif issubclass(arg_type, (SymbolicCallArg, sympy.Expr, int)):
-                    prefix.writeline(f"int64_t {name},")
-                elif arg_type is float:
-                    prefix.writeline(f"float {name},")
-                elif arg_type is bool:
-                    prefix.writeline(f"bool {name},")
-                else:
-                    raise ValueError(f"Unexpected arg type {arg_type}")
-            prefix.writeline("int32_t device_idx_,")
-            prefix.writeline(
-                maybe_hipify_code_wrapper(
-                    f"{wrapper.device_codegen.cpp_stream_type()} stream_,"
-                )
-            )
-            if V.graph.aot_mode:
-                prefix.writeline("kernels_type_& kernels_,")
-            prefix.writeline(
-                "const std::optional<std::string>& cubin_dir_ = std::nullopt"
-            )
-        prefix.writeline("){")
+        # Write wrapper function signature
+        self._write_wrapper_signature(prefix, wrapper, def_args, arg_types)
+
         with prefix.indent():
             if V.graph.aot_mode:
                 # Emit the original Triton kernel for debugging purposes
@@ -175,6 +254,314 @@ def generate(self, wrapper: CppWrapperGpu):
                 params[get_cpp_wrapper_cubin_path_name()]
             )
 
+    def _resolve_lazy_arg_names(self) -> tuple[list[str], list[str]]:
+        """Compute wrapper and kernel arg names from triton_meta signature.
+
+        Returns (wrapper_arg_names, kernel_arg_names) where:
+        - wrapper_arg_names: params accepted by the C++ wrapper function
+        - kernel_arg_names: params passed to the GPU kernel launch (non-constexpr only)
+        """
+        assert self.triton_meta is not None, (
+            f"triton_meta is required for lazy compile of {self.kernel_name}"
+        )
+        signature = self.triton_meta.get("signature", {})
+        inductor_meta = self.inductor_meta or {}
+        extra_launcher_args_count = len(inductor_meta.get("extra_launcher_args", []))
+        tma_tensor_args = self.tma_tensor_args
+        num_tma_tensor_args = len(tma_tensor_args)
+
+        internal_config_suffixes = ("BLOCK", "RSPLIT", "RSPLIT_SIZE")
+        # Declared constexpr params (tl.constexpr in kernel signature) are excluded
+        # from arg_types for user-defined kernels, while value-based constexpr params
+        # (e.g. numel=1, arg=None) are still in arg_types.
+        declared_constexpr_names = OrderedSet(
+            inductor_meta.get("declared_constexpr_names", [])
+        )
+        wrapper_arg_names = []
+        kernel_arg_names = []
+        for name, sig_type in signature.items():
+            if name.endswith(internal_config_suffixes):
+                continue
+            if sig_type != "constexpr":
+                kernel_arg_names.append(name)
+            if name not in declared_constexpr_names:
+                wrapper_arg_names.append(name)
+
+        num_wrapper_args = (
+            len(self.arg_types) - extra_launcher_args_count - num_tma_tensor_args
+        )
+        if num_wrapper_args != len(wrapper_arg_names):
+            raise AssertionError(
+                f"Mismatch between ({num_wrapper_args}) arg_types and "
+                f"{len(wrapper_arg_names)} wrapper_arg_names for {self.kernel_name}."
+            )
+
+        # Append grid args: passed to wrapper. Kernel args will handle grids separately.
+        for i in range(extra_launcher_args_count):
+            wrapper_arg_names.append(f"_grid_{i}")
+
+        # Add TMA tensor args after grid args
+        if tma_tensor_args:
+            sig_tma_keys = list(self._get_tma_args().keys())
+            assert list(tma_tensor_args.keys()) == sig_tma_keys, (
+                f"TMA tensor args order mismatch for {self.kernel_name}: "
+                f"{list(tma_tensor_args.keys())} vs signature order {sig_tma_keys}"
+            )
+        for desc_name in tma_tensor_args:
+            wrapper_arg_names.append(f"_tma_tensor_{desc_name}")
+
+        return wrapper_arg_names, kernel_arg_names
+
+    def _generate_lazy_grid(self, prefix: IndentedBuffer) -> None:
+        """Generate grid computation code for lazy-compiled kernels."""
+        kernel_name = self.kernel_name
+        grid_type = self.inductor_meta.get("grid_type") if self.inductor_meta else None
+
+        # For PrecomputedGrid, generate switch statement on config_index
+        if grid_type == "PrecomputedGrid":
+            assert self.inductor_meta is not None
+            precomputed_grids = self.inductor_meta.get("precomputed_grids", [])
+            extra_launcher_args = self.inductor_meta.get("extra_launcher_args", [])
+
+            switch_cases = []
+            for idx, entry in enumerate(precomputed_grids):
+                cpp_grids = list(entry.get("cpp", ["1L", "1L", "1L"]))
+                # Replace internal arg names with C++ parameter names
+                # e.g., _launcher_s0 -> _grid_0
+                for i, arg_name in enumerate(extra_launcher_args):
+                    cpp_grids = [g.replace(arg_name, f"_grid_{i}") for g in cpp_grids]
+                g0 = cpp_grids[0]
+                g1 = cpp_grids[1] if len(cpp_grids) > 1 else "1"
+                g2 = cpp_grids[2] if len(cpp_grids) > 2 else "1"
+                switch_cases.append(
+                    f"case {idx}: grid_0 = {g0}; grid_1 = {g1}; grid_2 = {g2}; break;"
+                )
+            switch_cases.append("default: grid_0 = 1; grid_1 = 1; grid_2 = 1; break;")
+            switch_body = "\n                        ".join(switch_cases)
+
+            prefix.splice(
+                f"""\
+                uint32_t grid_0, grid_1, grid_2;
+                switch ({kernel_name}_result.config_index) {{
+                    {switch_body}
+                }}
+                if (grid_0 == 0) return;
+                """
+            )
+        else:
+            from ..runtime.triton_heuristics import GridExpr
+
+            grid = GridExpr.from_meta_lazy(self.inductor_meta, kernel_name)
+            for line in grid.prefix:
+                prefix.writeline(line)
+
+            prefix.splice(
+                f"""\
+                uint32_t grid_0 = {grid.x_grid};
+                uint32_t grid_1 = {grid.y_grid};
+                uint32_t grid_2 = {grid.z_grid};
+                if (grid_0 == 0) return;
+                """
+            )
+
+    def _generate_lazy_tma_args(
+        self,
+        prefix: IndentedBuffer,
+        call_args_str: str,
+        kernel_arg_names: list[str],
+        tma_arg_names: OrderedSet[str],
+        signature: dict[str, str],
+    ) -> str:
+        """Unpack TMA descriptor args into kernel launch args."""
+        for arg_name in kernel_arg_names:
+            if arg_name in tma_arg_names:
+                tma_parts = _unpack_tma_descriptor_args(arg_name, signature[arg_name])
+                tma_str = ", ".join(tma_parts)
+                call_args_str = (
+                    f"{call_args_str}, {tma_str}" if call_args_str else tma_str
+                )
+        return call_args_str
+
+    def _generate_lazy_scratch(
+        self,
+        prefix: IndentedBuffer,
+        wrapper: CppWrapperGpu,
+        call_args_str: str,
+    ) -> str:
+        """Generate scratch space allocations with runtime-known sizes."""
+        kernel_name = self.kernel_name
+        dtype_str = wrapper.codegen_dtype(torch.uint8)
+        device_type, _ = wrapper.codegen_device(torch.device(get_gpu_type())).split(
+            ", "
+        )
+        device_ptr_type = wrapper.device_codegen.cpp_device_ptr()
+        for scratch_name in ("global_scratch", "profile_scratch"):
+            size_expr = f"{kernel_name}_result.{scratch_name}"
+            var = f"{scratch_name}_ptr"
+            prefix.splice(
+                maybe_hipify_code_wrapper(
+                    f"""\
+                {device_ptr_type} {var} = 0;
+                RAIIAtenTensorHandle {var}_tensor;
+                if ({size_expr} > 0) {{
+                    int64_t {var}_size[] = {{{size_expr}}};
+                    int64_t {var}_stride[] = {{1}};
+                    AtenTensorHandle {var}_handle;
+                    AOTI_TORCH_ERROR_CODE_CHECK(aoti_torch_empty_strided(
+                        1, {var}_size, {var}_stride, {dtype_str},
+                        {device_type}, device_idx_, &{var}_handle));
+                    {var}_tensor = RAIIAtenTensorHandle({var}_handle);
+                    {var} = reinterpret_cast<{device_ptr_type}>({var}_tensor.data_ptr());
+                }}
+            """
+                )
+            )
+            call_args_str += f", &{var}"
+        return call_args_str
+
+    def _generate_lazy_launch(
+        self,
+        prefix: IndentedBuffer,
+        wrapper: CppWrapperGpu,
+        wrapper_arg_names: list[str],
+        kernel_arg_names: list[str],
+    ) -> None:
+        """Generate kernel launch code for lazy-compiled kernels."""
+        kernel_name = self.kernel_name
+        signature = (self.triton_meta or {}).get("signature", {})
+        tma_tensor_args = self.tma_tensor_args
+        num_tma_tensor_args = len(tma_tensor_args)
+
+        # wrapper_arg_names may include grid and TMA tensor args at the end;
+        # only the leading portion maps 1:1 to kernel signature params.
+        num_signature_args = len(wrapper_arg_names) - num_tma_tensor_args
+        inductor_meta = self.inductor_meta or {}
+        num_signature_args -= len(inductor_meta.get("extra_launcher_args", []))
+
+        arg_type_lookup = dict(
+            zip(wrapper_arg_names, self.arg_types[:num_signature_args])
+        )
+
+        # Identify TMA args — they are already passed as StableTMADescriptor params,
+        # so we just unpack them directly (no need to reconstruct from tensors).
+        tma_arg_names = OrderedSet(self._get_tma_args().keys())
+
+        # Non-TMA args go through generate_args_decl
+        non_tma_arg_names = [n for n in kernel_arg_names if n not in tma_arg_names]
+        non_tma_arg_types = [
+            arg_type_lookup[n] for n in non_tma_arg_names if n in arg_type_lookup
+        ]
+        non_tma_arg_sigs = [signature.get(n) for n in non_tma_arg_names]
+
+        call_args_str = wrapper.generate_args_decl(
+            prefix,
+            non_tma_arg_names,
+            non_tma_arg_types,
+            non_tma_arg_sigs,
+        )
+
+        call_args_str = self._generate_lazy_tma_args(
+            prefix, call_args_str, kernel_arg_names, tma_arg_names, signature
+        )
+        call_args_str = self._generate_lazy_scratch(prefix, wrapper, call_args_str)
+
+        launch_args = (
+            f"{kernel_name}, grid_0, grid_1, grid_2,"
+            f" {kernel_name}_result.num_warps,"
+            f" {kernel_name}_result.shared_mem,"
+            f" kernel_args_, stream_"
+        )
+
+        prefix.splice(
+            f"""\
+            void* kernel_args_[] = {{{call_args_str}}};
+            launchKernel({launch_args});
+            """
+        )
+
+    def generate_lazy(self, wrapper: CppWrapperGpu):
+        """
+        Generate C++ code that embeds Triton source and compiles it at runtime.
+        """
+        prefix = wrapper.prefix
+        kernel_name = self.kernel_name
+        # Track kernel names for parallel initialization
+        wrapper._lazy_kernel_names.append(kernel_name)
+
+        # Include TMA helpers if any args use TMA descriptors
+        tma_signature_types = self._get_tma_args()
+        if tma_signature_types:
+            wrapper.write_tma_descriptor_helpers_once()
+
+        kernel_var_decl = maybe_hipify_code_wrapper(
+            f"static {wrapper.device_codegen.cpp_kernel_type()} {kernel_name} = nullptr;"
+        )
+        prefix.writeline(kernel_var_decl)
+        # Use delimited raw string to handle )" in kernel source
+        kernel_source_str = self.kernel_name_to_body.get(kernel_name, "")
+        kernel_body = f'R"TRITON(\n{kernel_source_str}\n)TRITON"'
+        prefix.writeline(f"static const char* {kernel_name}_source = {kernel_body};")
+        prefix.writeline(f"static LazyKernelCompileResult {kernel_name}_result;")
+
+        wrapper_arg_names, kernel_arg_names = self._resolve_lazy_arg_names()
+        signature = (self.triton_meta or {}).get("signature", {})
+        self._write_wrapper_signature(
+            prefix, wrapper, wrapper_arg_names, self.arg_types, signature
+        )
+
+        # Build autotune args - for TMA, pass tensors instead of descriptors.
+        # Only iterate over signature params and grid args, not the trailing
+        # TMA tensor params (those are only in the C++ wrapper signature).
+        tma_tensor_args = self.tma_tensor_args
+        num_autotune_args = len(wrapper_arg_names) - len(tma_tensor_args)
+        autotune_arg_list = []
+        # Track which args need scalar extraction for the autotune call.
+        # UnwrapUnspecArg args are 0-dim tensors in C++ that Triton expects
+        # as Python scalars; we use codegen_tensor_item to extract them.
+        scalar_extractions: list[tuple[str, str, torch_dtype]] = []
+        for idx, name in enumerate(wrapper_arg_names[:num_autotune_args]):
+            if name in tma_signature_types:
+                autotune_arg_list.append(f"_tma_tensor_{name}")
+            elif isinstance(self.arg_types[idx], UnwrapUnspecArg):
+                scalar_var = f"_autotune_scalar_{name}"
+                scalar_extractions.append((name, scalar_var, self.arg_types[idx].dtype))
+                autotune_arg_list.append(scalar_var)
+            else:
+                autotune_arg_list.append(name)
+        autotune_args = ", ".join(autotune_arg_list)
+        # Lazy compile with autotuning on first invocation
+        with prefix.indent():
+            prefix.writeline(f"if ({kernel_name} == nullptr) {{")
+            with prefix.indent():
+                for tensor_name, scalar_var, dtype in scalar_extractions:
+                    wrapper.codegen_tensor_item(
+                        dtype, tensor_name, scalar_var, indented_buffer=prefix
+                    )
+                prefix.splice(
+                    f"""\
+                    {kernel_name}_result = runTritonKernelWithAutotune(
+                        _module_pending_kernels, "{kernel_name}", stream_, {autotune_args});
+
+                    {kernel_name} = loadKernel(
+                        {kernel_name}_result.cubin_path,
+                        {kernel_name}_result.mangled_name,
+                        {kernel_name}_result.shared_mem);
+
+                    // First invocation already ran the kernel, so return early
+                    return;
+                    """
+                )
+            prefix.writeline("}")
+
+            self._generate_lazy_grid(prefix)
+            self._generate_lazy_launch(
+                prefix,
+                wrapper,
+                wrapper_arg_names,
+                kernel_arg_names,
+            )
+        prefix.writeline("}")
+
     def generate_grid(
         self,
         prefix: IndentedBuffer,
@@ -247,8 +634,9 @@ def generate_launch_kernel(self, prefix, wrapper, kernel_var_name, params):
         ]
         arg_types = [arg_type_lookup[name] for name in call_args]
         arg_signatures = [triton_meta["signature"][name] for name in call_args]
+        num_ctas = params.get("config", {}).get("num_ctas", 1)
         scratch_spaces = {
-            name: params[name]
+            name: params[name] * num_ctas
             for name in ["global_scratch", "profile_scratch"]
             if params.get(name, None) is not None
         }
@@ -270,8 +658,6 @@ def generate_launch_kernel(self, prefix, wrapper, kernel_var_name, params):
             "kernel_args_",
             "stream_",
         ]
-        if wrapper.device == "xpu":
-            launch_kernel_args.append(str(params["threads_per_warp"]))
 
         enable_kernel_profile = config.cpp.enable_kernel_profile and sys.platform in [
             "linux",
@@ -426,13 +812,14 @@ def __init__(self) -> None:
         self._kernel_name_to_body: dict[str, str] = {}
         self._triton_call_wrappers: dict[str, DeferredTritonCallWrapper] = {}
         self.autotune_input_prefix = "_REAL_AUTOTUNE_INPUT"
+        self._lazy_kernel_names: list[str] = []
 
     @staticmethod
     def create(
         is_subgraph: bool,
-        subgraph_name: Optional[str],
-        parent_wrapper: Optional[PythonWrapperCodegen],
-        partition_signatures: Optional[GraphPartitionSignature] = None,
+        subgraph_name: str | None,
+        parent_wrapper: PythonWrapperCodegen | None,
+        partition_signatures: GraphPartitionSignature | None = None,
     ):
         # TODO - support subgraph codegen by lifting functions. Check the
         # comment at CppWrapperCpu `codegen_subgraph` function.
@@ -510,9 +897,9 @@ def _define_kernel_helper(
         self,
         kernel_name: str,
         kernel_body: str,
-        metadata: Optional[str] = None,
+        metadata: str | None = None,
         gpu: bool = True,
-        cpp_definition: Optional[str] = None,
+        cpp_definition: str | None = None,
     ):
         if gpu:
             self._kernel_name_to_body[kernel_name] = kernel_body
@@ -541,6 +928,32 @@ def finalize_prefix(self):
         for kernel in self._triton_call_wrappers.values():
             self.prefix.writeline("\n")
             kernel.generate(self)
+
+        # Generate parallel kernel compilation initialization function
+        if self._lazy_kernel_names:
+            start_compile_calls = "\n    ".join(
+                f'startKernelCompile(_module_pending_kernels, "{name}", {name}_source);'
+                for name in self._lazy_kernel_names
+            )
+            self.prefix.splice(
+                f"""\
+// Start parallel compilation of all Triton kernels
+static inline void start_all_triton_kernel_compiles() {{
+    loadLazyCompileFuncs();
+    _module_pending_kernels = PyDict_New();
+    AOTI_TORCH_CHECK(_module_pending_kernels, "Failed to create pending kernels dict");
+    {start_compile_calls}
+}}
+
+// Static initializer to start kernel compilation on module load
+static struct TritonKernelCompileInit {{
+    TritonKernelCompileInit() {{
+        start_all_triton_kernel_compiles();
+    }}
+}} __triton_kernel_compile_init;
+"""
+            )
+
         triton_prefix = self.prefix
 
         self.prefix = IndentedBuffer()
@@ -630,12 +1043,12 @@ def fill_array(name, values):
 
     def generate_args_decl(
         self,
-        code: Union[IndentedBuffer, Self],
+        code: IndentedBuffer | Self,
         call_args,
         arg_types,
         arg_signatures,
         is_triton_kernel=True,
-        scratch_spaces: Optional[dict[str, int]] = None,
+        scratch_spaces: dict[str, int] | None = None,
     ):
         """
         Generates any declarations of args to pass into a kernel call, and then returns the arg names.
@@ -664,23 +1077,7 @@ def process_tma_stable_arg(arg, arg_type, arg_signature, var_name):
             # (StableTMADescriptor) containing the necessary information; and then
             # when we call the function (i.e. here), we unpack the struct members.
             code.writeline(f"auto {var_name} = {cexpr(arg)};")
-
-            result = []
-            result.append(f"&{var_name}.m")
-
-            # from https://github.com/triton-lang/triton/blob/16961b79bdac1b774b42d44e52fd55a266ec2866/third_party/nvidia/backend/driver.py#L111  # noqa: B950
-            match = re.match("tensordesc<([^[>]*)\\[([^]]*)\\]", arg_signature)
-            assert match is not None
-            shape = match.group(2)
-            ndim = shape.count(",") + 1
-
-            for i in range(ndim):
-                result.append(f"&{var_name}.block_shape[{i}]")
-
-            for i in range(ndim):
-                result.append(f"&{var_name}.strides[{i}]")
-
-            return result
+            return _unpack_tma_descriptor_args(var_name, arg_signature)
 
         def process_args(arg, arg_type, arg_signature=None):
             var_name = f"var_{next(self.arg_var_id)}"
@@ -778,6 +1175,7 @@ def _generate_kernel_call_helper(
         inductor_meta=None,
         graph_name="",
         original_fxnode_name=None,
+        current_stream_idx=None,
     ):
         """
         Override the default value of argument 'gpu' to True here.
@@ -832,6 +1230,25 @@ def _generate_kernel_call_helper(
                 # pyrefly: ignore [bad-argument-type]
                 arg_types,
             )
+
+            # For lazy compile mode with TMA, extract underlying tensor names
+            tma_tensor_args: dict[str, str] = {}
+            is_lazy_compile = (
+                not V.graph.aot_mode and config.triton.autotune_at_compile_time is False
+            )
+            if is_lazy_compile and raw_args and triton_meta:
+                signature = triton_meta.get("signature", {})
+                raw_keys_list = raw_keys or []
+                for key, raw_arg in zip(raw_keys_list, raw_args):
+                    sig_type = signature.get(key, "")
+                    if isinstance(sig_type, str) and signature_is_tma_desc(sig_type):
+                        if isinstance(raw_arg, TMADescriptorStable):
+                            # Get the underlying tensor name
+                            tensor_name = raw_arg.get_tensor().get_name()
+                            tma_tensor_args[key] = tensor_name
+                        else:
+                            raise AssertionError("Unsupported TMA descriptor type")
+
             wrapper_name = f"call_{kernel_name}"
             if wrapper_name not in self._triton_call_wrappers:
                 self._triton_call_wrappers[wrapper_name] = DeferredTritonCallWrapper(
@@ -839,7 +1256,19 @@ def _generate_kernel_call_helper(
                     kernel_name,
                     self._kernel_name_to_body,
                     arg_types,
+                    triton_meta=triton_meta,
+                    inductor_meta=inductor_meta,
+                    tma_tensor_args=tma_tensor_args,
                 )
+
+            # For TMA in lazy compile mode, add tensor args to the call
+            if is_lazy_compile and tma_tensor_args:
+                for tensor_name in tma_tensor_args.values():
+                    call_args.append(tensor_name)
+                    arg_types.append(
+                        torch.float32
+                    )  # dtype doesn't matter, just need tensor type
+
             device_idx = "this->device_idx_" if V.graph.aot_mode else str(device.index)
             call_args.append(device_idx)
             call_args.append(stream)
@@ -854,7 +1283,7 @@ def _generate_kernel_call_helper(
                 self.writeline(f"{wrapper_name}({', '.join(call_args)});")
         else:
             casted = []
-            # pyrefly: ignore [no-matching-overload]
+            # pyrefly: ignore [bad-argument-type, no-matching-overload]
             for arg_type, arg in zip(arg_types, call_args):
                 new_arg = arg
                 if arg_type.endswith("*") and arg != "nullptr":
diff --git a/torch/_inductor/codegen/cpp_wrapper_mps.py b/torch/_inductor/codegen/cpp_wrapper_mps.py
index 3cc131d041b88..86b6cf92870cf 100644
--- a/torch/_inductor/codegen/cpp_wrapper_mps.py
+++ b/torch/_inductor/codegen/cpp_wrapper_mps.py
@@ -1,4 +1,4 @@
-from typing import Any, Optional
+from typing import Any
 
 import sympy
 
@@ -25,9 +25,9 @@ def __init__(self) -> None:
     @staticmethod
     def create(
         is_subgraph: bool,
-        subgraph_name: Optional[str],
-        parent_wrapper: Optional[PythonWrapperCodegen],
-        partition_signatures: Optional[GraphPartitionSignature] = None,
+        subgraph_name: str | None,
+        parent_wrapper: PythonWrapperCodegen | None,
+        partition_signatures: GraphPartitionSignature | None = None,
     ) -> "CppWrapperMps":
         return CppWrapperMps()
 
@@ -36,15 +36,16 @@ def _generate_kernel_call_helper(
         kernel_name: str,
         call_args: list[str],
         *,
-        device: Optional[torch.device] = None,
+        device: torch.device | None = None,
         triton: bool = True,
-        arg_types: Optional[tuple[Any, ...]] = None,
-        raw_keys: Optional[tuple[Any, ...]] = None,
-        raw_args: Optional[tuple[Any, ...]] = None,
-        triton_meta: Optional[dict[str, Any]] = None,
-        inductor_meta: Optional[dict[str, Any]] = None,
+        arg_types: tuple[Any, ...] | None = None,
+        raw_keys: tuple[Any, ...] | None = None,
+        raw_args: tuple[Any, ...] | None = None,
+        triton_meta: dict[str, Any] | None = None,
+        inductor_meta: dict[str, Any] | None = None,
         graph_name: str = "",
-        original_fxnode_name: Optional[str] = None,
+        original_fxnode_name: str | None = None,
+        current_stream_idx: int | None = None,
     ) -> None:
         """
         Generates MPS kernel call code. It should look something like:
diff --git a/torch/_inductor/codegen/cuda/compile_utils.py b/torch/_inductor/codegen/cuda/compile_utils.py
index 9d32e6cae9448..cee25029c86f3 100644
--- a/torch/_inductor/codegen/cuda/compile_utils.py
+++ b/torch/_inductor/codegen/cuda/compile_utils.py
@@ -3,7 +3,6 @@
 import os
 import shutil
 from pathlib import Path
-from typing import Optional
 
 import torch
 from torch._inductor import config
@@ -68,7 +67,7 @@ def _cutlass_include_paths() -> list[str]:
     ]
 
 
-def _cuda_compiler() -> Optional[str]:
+def _cuda_compiler() -> str | None:
     if cuda_env.nvcc_exist(config.cuda.cuda_cxx):
         return config.cuda.cuda_cxx
     if config.is_fbcode():
@@ -129,12 +128,18 @@ def _nvcc_arch_as_compile_option() -> str:
     if arch == "90":
         # Required by cutlass compilation.
         return "90a"
-    if arch == "103":
-        return "100f"
     if arch == "100":
         return "100a"
+    if arch == "101":
+        return "101a"
+    if arch == "103":
+        return "103a"
+    if arch == "110":
+        return "110a"
     if arch == "120":
         return "120a"
+    if arch == "121":
+        return "121a"
     return arch
 
 
@@ -151,7 +156,7 @@ def _nvcc_compiler_options() -> list[str]:
         "-w",
         f"-gencode=arch=compute_{arch},code=[{','.join(code)}]",
         config.cutlass.compile_opt_level,
-        "-std=c++17",
+        "-std=c++20",
         "--expt-relaxed-constexpr",
         "-DNDEBUG",
     ]
@@ -240,9 +245,7 @@ def __enter__(self, *args, **kwargs):
 
         _compile_method_orig = torch._inductor.codecache.CUDACodeCache.compile
 
-        def my_compile(
-            source_code, dst_file_ext, extra_args: Optional[list[str]] = None
-        ):
+        def my_compile(source_code, dst_file_ext, extra_args: list[str] | None = None):
             self.sources.append(source_code)
             return _compile_method_orig(source_code, dst_file_ext)
 
diff --git a/torch/_inductor/codegen/cuda/cuda_env.py b/torch/_inductor/codegen/cuda/cuda_env.py
index 9ca3afbd9ca57..e12bd94d2a677 100644
--- a/torch/_inductor/codegen/cuda/cuda_env.py
+++ b/torch/_inductor/codegen/cuda/cuda_env.py
@@ -1,7 +1,6 @@
 import functools
 import logging
 import shutil
-from typing import Optional
 
 import torch
 from torch._inductor.utils import clear_on_fresh_cache
@@ -14,7 +13,7 @@
 
 @clear_on_fresh_cache
 @functools.lru_cache(1)
-def get_cuda_arch() -> Optional[str]:
+def get_cuda_arch() -> str | None:
     try:
         cuda_arch = config.cuda.arch
         if cuda_arch is None:
@@ -39,7 +38,7 @@ def is_datacenter_blackwell_arch() -> bool:
 
 @clear_on_fresh_cache
 @functools.lru_cache(1)
-def get_cuda_version() -> Optional[str]:
+def get_cuda_version() -> str | None:
     try:
         cuda_version = config.cuda.version
         if cuda_version is None:
@@ -51,5 +50,5 @@ def get_cuda_version() -> Optional[str]:
 
 
 @functools.cache
-def nvcc_exist(nvcc_path: Optional[str] = "nvcc") -> bool:
+def nvcc_exist(nvcc_path: str | None = "nvcc") -> bool:
     return nvcc_path is not None and shutil.which(nvcc_path) is not None
diff --git a/torch/_inductor/codegen/cuda/device_op_overrides.py b/torch/_inductor/codegen/cuda/device_op_overrides.py
index 147515e0decfe..dde9b0924cc8c 100644
--- a/torch/_inductor/codegen/cuda/device_op_overrides.py
+++ b/torch/_inductor/codegen/cuda/device_op_overrides.py
@@ -1,9 +1,8 @@
 from __future__ import annotations
 
-from typing import Optional
-
 import torch
 
+from ...runtime.hints import get_warp_size
 from ..common import (
     DeviceOpOverrides,
     register_device_op_overrides,
@@ -52,6 +51,10 @@ def kernel_header(self) -> str:
         return source_codes
 
     def kernel_driver(self) -> str:
+        """Return C++ host-side helpers (loadKernel, launchKernel, CUDA_DRIVER_CHECK)
+        embedded in AOTI-generated wrapper code."""
+        # NVIDIA devices have a warp size of 32, while AMD devices can have a
+        # warp size of 32 or 64 depending on the architecture.
         source_codes = """
             #define CUDA_DRIVER_CHECK(EXPR)                    \\
             do {                                               \\
@@ -120,17 +123,16 @@ def kernel_driver(self) -> str:
                     void* args[],
                     cudaStream_t stream) {
                 CUDA_DRIVER_CHECK(cuLaunchKernel(
-                    func, gridX, gridY, gridZ, 32*numWarps, 1, 1, sharedMemBytes, stream, args, nullptr
+                    func, gridX, gridY, gridZ, __WARP_SIZE__*numWarps, 1, 1, sharedMemBytes, stream, args, nullptr
                 ));
             }
         """
-        if torch.version.hip is not None:
-            # Adjusting the warp size to GPU supported wavefront size on AMD GPU
-            prop = torch.cuda.get_device_properties(torch.cuda.current_device())
-            source_codes = source_codes.replace(
-                "32*numWarps", str(prop.warp_size) + "*numWarps"
-            )
-        return source_codes
+        if torch.version.hip is not None and torch.cuda.is_available():
+            device = torch.device("cuda", torch.cuda.current_device())
+            warp_size = get_warp_size(device)
+        else:
+            warp_size = 32
+        return source_codes.replace("__WARP_SIZE__", str(warp_size))
 
     def tma_descriptor_helpers(self) -> str:
         """
@@ -333,12 +335,15 @@ def cpp_device_ptr(self) -> str:
         return "CUdeviceptr"
 
     def cpp_scratch(
-        self, idx: int, workspace: TritonScratchWorkspace, prefix: Optional[str] = None
-    ) -> Optional[tuple[list[str], str]]:
+        self, idx: int, workspace: TritonScratchWorkspace, prefix: str | None = None
+    ) -> tuple[list[str], str] | None:
         prefix = f"{prefix}_" if prefix else ""
         var_name = f"{prefix}scratch_{idx}"
         if workspace.size > 0:
-            size_array = f"int64_t {var_name}_size[] = {{{workspace.size}}};"
+            size_expr = (
+                f"static_cast<int64_t>({workspace.size}) * grid_0 * grid_1 * grid_2"
+            )
+            size_array = f"int64_t {var_name}_size[] = {{{size_expr}}};"
             stride_array = f"int64_t {var_name}_stride[] = {{1}};"
             device_type = "cached_torch_device_type_cuda"
             device_idx = "device_idx_"
diff --git a/torch/_inductor/codegen/cuda_combined_scheduling.py b/torch/_inductor/codegen/cuda_combined_scheduling.py
index 293ec4597180b..b69225b78a2c3 100644
--- a/torch/_inductor/codegen/cuda_combined_scheduling.py
+++ b/torch/_inductor/codegen/cuda_combined_scheduling.py
@@ -1,7 +1,7 @@
 # mypy: allow-untyped-defs
 from __future__ import annotations
 
-from typing import Any, Optional, TYPE_CHECKING, Union
+from typing import Any, TYPE_CHECKING
 
 from ..scheduler import (
     BaseSchedulerNode,
@@ -28,7 +28,7 @@
 
     from .common import BackendFeature
 
-    _IntLike: TypeAlias = Union[int, Expr]
+    _IntLike: TypeAlias = int | Expr
 
 
 class CUDACombinedScheduling(BaseScheduling):
@@ -41,7 +41,7 @@ class CUDACombinedScheduling(BaseScheduling):
     this would also be the place to do it.
     """
 
-    def __init__(self, scheduler: Optional[Scheduler]) -> None:
+    def __init__(self, scheduler: Scheduler | None) -> None:
         super().__init__(scheduler)
         self._triton_scheduling = TritonScheduling(scheduler)
         self._cutlass_scheduling = CUTLASSScheduling(scheduler)
@@ -112,7 +112,7 @@ def codegen_template(
         template_node: BaseSchedulerNode,
         epilogue_nodes: Sequence[BaseSchedulerNode],
         prologue_nodes: Sequence[BaseSchedulerNode],
-    ) -> Optional[str]:
+    ) -> str | None:
         if self._cutlass_scheduling.is_cutlass_template(template_node):
             assert not prologue_nodes
             return self._cutlass_scheduling.codegen_template(
@@ -148,7 +148,7 @@ def codegen_template(
     def codegen_mix_order_reduction(self, node):
         return self._triton_scheduling.codegen_mix_order_reduction(node)
 
-    def codegen_node(self, node: Union[FusedSchedulerNode, SchedulerNode]) -> None:
+    def codegen_node(self, node: FusedSchedulerNode | SchedulerNode) -> None:
         return self._triton_scheduling.codegen_node(node)
 
     def codegen_sync(self) -> None:
@@ -172,7 +172,7 @@ def generate_kernel_code_from_nodes(
         self,
         nodes: Sequence[Any],
         benchmark_kernel: bool = False,
-        hint_override: Optional[int] = None,
+        hint_override: int | None = None,
     ) -> str:
         return self._triton_scheduling.generate_kernel_code_from_nodes(
             nodes, benchmark_kernel, hint_override=hint_override
@@ -180,7 +180,7 @@ def generate_kernel_code_from_nodes(
 
     def benchmark_combo_kernel(
         self, node_list: Sequence[BaseSchedulerNode], node_benchmark_results
-    ) -> tuple[float, float, list[Optional[str]]]:
+    ) -> tuple[float, float, list[str | None]]:
         return self._triton_scheduling.benchmark_combo_kernel(
             node_list, node_benchmark_results
         )
diff --git a/torch/_inductor/codegen/cutedsl/cutedsl_kernel.py b/torch/_inductor/codegen/cutedsl/cutedsl_kernel.py
index 460b78340d461..175860efbaff0 100644
--- a/torch/_inductor/codegen/cutedsl/cutedsl_kernel.py
+++ b/torch/_inductor/codegen/cutedsl/cutedsl_kernel.py
@@ -4,7 +4,7 @@
 import logging
 import textwrap
 from collections.abc import Callable
-from typing import Any, Optional
+from typing import Any
 
 import sympy
 
@@ -45,9 +45,7 @@
 class CuteDSLKernelWrapper:
     """Wrapper to provide .run() interface for CuteDSL kernels"""
 
-    def __init__(
-        self, kernel_fn: Callable[..., Any], kernel_path: Optional[str] = None
-    ):
+    def __init__(self, kernel_fn: Callable[..., Any], kernel_path: str | None = None):
         self.kernel_fn = kernel_fn
         self.kernel_path = kernel_path
         kernel_code_log.info("CuteDSL kernel path: %s", kernel_path)
@@ -72,9 +70,9 @@ class CuteDSLSubgraphInfo:
     """Minimal subgraph info for CuteDSL kernels."""
 
     body: IndentedBuffer
-    template_mask: Optional[str] = None
-    template_out: Optional[str] = None
-    cse: Optional[CSE[Any]] = None
+    template_mask: str | None = None
+    template_out: str | None = None
+    cse: CSE[Any] | None = None
 
     def __post_init__(self):
         self.only_copy_if_non_none_fields = ("cse",)
@@ -97,7 +95,7 @@ def __init__(
         kernel_name: str,
         input_nodes: list[Buffer],
         output_node: Buffer,
-        subgraphs: Optional[list[Buffer]] = None,
+        subgraphs: list[Buffer] | None = None,
     ) -> None:
         # Call parent Kernel constructor
         super().__init__()
@@ -109,9 +107,9 @@ def __init__(
 
         # Template attributes
         self.body: IndentedBuffer = IndentedBuffer()
-        self.template_mask: Optional[str] = None
-        self.template_out: Optional[str] = None
-        self.template_indices: Optional[list[Any]] = None
+        self.template_mask: str | None = None
+        self.template_out: str | None = None
+        self.template_indices: list[Any] | None = None
         self.render_hooks: dict[str, Any] = {}
 
         # TODO Additional attributes needed by template system
@@ -404,8 +402,8 @@ def _get_subgraph(self, subgraph_number: int):
     def modification(
         self,
         subgraph_number: int,
-        output_name: Optional[str],
-        mask: Optional[str] = None,
+        output_name: str | None,
+        mask: str | None = None,
         **fixed_inputs,
     ) -> str:
         """Generate CuteDSL code for a subgraph modification."""
@@ -469,7 +467,7 @@ def __init__(
         kernel,
         subgraph_number: int,
         fixed_inputs: dict[str, Any],
-        mask: Optional[str],
+        mask: str | None,
     ):
         cutedsl_ops = CuteDSLOpOverrides()
         super().__init__(cutedsl_ops)
@@ -546,11 +544,19 @@ def _emit_scalar_fragment(
         self, expr_str: str, cute_dtype: str, torch_dtype: torch.dtype
     ) -> str:
         """
-        Convert SSA expression to indexable scalar for tensor loads.
+        Convert expression to indexable scalar for tensor loads.
 
         Workaround for lack of gather support: SSA values cannot be used directly
-        as indices. This generates code to convert SSA → indexable scalar.
+        as indices in tensor loads. This generates code to convert SSA → indexable
+        scalar. Compile-time integer constants are already indexable and are
+        returned directly without the SSA round-trip.
         """
+        # Constant integer expressions (e.g. sympy-folded offsets like "0")
+        # are already valid indices — skip the ssa_to_indexable round-trip
+        # which only accepts TensorSSA, not bare Python ints.
+        if expr_str.lstrip("-").isdigit():
+            return expr_str
+
         result = self.kernel.cse.newvar(dtype=torch_dtype)
         self.kernel.body.writeline(
             f"{result} = ssa_to_indexable({expr_str}, {cute_dtype})"
diff --git a/torch/_inductor/codegen/cutedsl/cutedsl_op_overrides.py b/torch/_inductor/codegen/cutedsl/cutedsl_op_overrides.py
index c855767ec174b..ebfb2e0c08507 100644
--- a/torch/_inductor/codegen/cutedsl/cutedsl_op_overrides.py
+++ b/torch/_inductor/codegen/cutedsl/cutedsl_op_overrides.py
@@ -7,7 +7,6 @@
 """
 
 import math
-from typing import Optional, Union
 
 import sympy
 
@@ -17,7 +16,7 @@
 from torch.utils._sympy.value_ranges import ValueRanges
 
 
-CuteDSLArg = Union[CSEVariable, str]
+CuteDSLArg = CSEVariable | str | bool | float | int
 
 
 def upcast_compute_type(dtype: torch.dtype) -> torch.dtype:
@@ -44,6 +43,7 @@ class CuteDSLOpOverrides(OpOverrides):
         torch.int16: "cutlass.Int16",
         torch.int32: "cutlass.Int32",
         torch.int64: "cutlass.Int64",
+        torch.uint8: "cutlass.Uint8",
         torch.bool: "cutlass.Boolean",
         torch.float8_e4m3fn: "cutlass.Float8E4M3FN",
         torch.float8_e5m2: "cutlass.Float8E5M2",
@@ -53,7 +53,7 @@ class CuteDSLOpOverrides(OpOverrides):
     LOG2_E = 1.4426950408889634  # 1/ln(2) for converting natural exp to base-2 exp
 
     @staticmethod
-    def _get_cse_var(arg: CuteDSLArg) -> Optional[CSEVariable]:
+    def _get_cse_var(arg: CuteDSLArg) -> CSEVariable | None:
         """Extract CSEVariable from arg if it's a tensor (either direct or wrapped in OpsValue)."""
         if isinstance(arg, CSEVariable):
             return arg
@@ -62,7 +62,30 @@ def _get_cse_var(arg: CuteDSLArg) -> Optional[CSEVariable]:
         return None
 
     @staticmethod
-    def _ensure_tensor_ssa(arg: CuteDSLArg, template_tensor: CuteDSLArg) -> str:
+    def _as_expr(arg: CuteDSLArg) -> str:
+        cse_var = CuteDSLOpOverrides._get_cse_var(arg)
+        if cse_var is not None:
+            return str(cse_var)
+        return str(arg)
+
+    @staticmethod
+    def _node_tensor_flags() -> tuple[bool, bool] | None:
+        node = V.current_node
+        if not isinstance(node, torch.fx.Node) or len(node.args) < 2:
+            return None
+
+        def _is_tensor(raw_arg: object) -> bool:
+            if hasattr(raw_arg, "meta"):
+                val = raw_arg.meta.get("val")
+                return isinstance(val, torch.Tensor) and val.ndim > 0
+            return False
+
+        return _is_tensor(node.args[0]), _is_tensor(node.args[1])
+
+    @staticmethod
+    def _ensure_tensor_ssa(
+        arg: CuteDSLArg, template_tensor: CuteDSLArg, *, is_tensor: bool
+    ) -> str:
         """
         Convert scalar arguments to TensorSSA using cute.full_like if needed.
 
@@ -73,21 +96,16 @@ def _ensure_tensor_ssa(arg: CuteDSLArg, template_tensor: CuteDSLArg) -> str:
         Returns:
             String representation suitable for CuteDSL operations
         """
-        if isinstance(arg, CSEVariable):
-            return str(arg)
-
-        if isinstance(arg, OpsValue) and isinstance(arg.value, CSEVariable):
-            return str(arg.value)
-
-        if isinstance(template_tensor, CSEVariable):
-            return f"cute.full_like({template_tensor}, {arg})"
-
-        return str(arg)
+        arg_expr = CuteDSLOpOverrides._as_expr(arg)
+        if is_tensor:
+            return arg_expr
+        template_expr = CuteDSLOpOverrides._as_expr(template_tensor)
+        return f"cute.full_like({template_expr}, {arg_expr})"
 
     @staticmethod
     def _extract_dtype_and_bounds(
         *args: CuteDSLArg,
-    ) -> tuple[Optional[torch.dtype], ValueRanges[sympy.Expr]]:
+    ) -> tuple[torch.dtype | None, ValueRanges[sympy.Expr]]:
         """Extract dtype and bounds from CSEVariable arguments (including OpsValue wrappers)."""
         for arg in args:
             cse_var = CuteDSLOpOverrides._get_cse_var(arg)
@@ -112,24 +130,69 @@ def _apply_binary_op(a: CuteDSLArg, b: CuteDSLArg, op_format: str) -> CuteDSLArg
         Returns:
             CSEVariable if at least one operand is a tensor, otherwise string
         """
-        # Check for CSEVariable directly or wrapped in OpsValue
-        tensor_arg = CuteDSLOpOverrides._get_cse_var(
-            a
-        ) or CuteDSLOpOverrides._get_cse_var(b)
+        a_cse = CuteDSLOpOverrides._get_cse_var(a)
+        b_cse = CuteDSLOpOverrides._get_cse_var(b)
+
+        node_flags = CuteDSLOpOverrides._node_tensor_flags()
+        if node_flags is not None:
+            a_is_tensor, b_is_tensor = node_flags
+        else:
+            a_is_tensor = a_cse is not None
+            b_is_tensor = b_cse is not None
+
+        tensor_arg = a if a_is_tensor else (b if b_is_tensor else None)
+        if tensor_arg is None:
+            tensor_arg = a_cse or b_cse
+
         if tensor_arg is not None:
-            a_ssa = CuteDSLOpOverrides._ensure_tensor_ssa(a, tensor_arg)
-            b_ssa = CuteDSLOpOverrides._ensure_tensor_ssa(b, tensor_arg)
+            if a_cse is None and b_cse is None:
+                return op_format.format(
+                    a=CuteDSLOpOverrides._as_expr(a),
+                    b=CuteDSLOpOverrides._as_expr(b),
+                )
+
+            a_ssa = CuteDSLOpOverrides._ensure_tensor_ssa(
+                a, tensor_arg, is_tensor=a_is_tensor
+            )
+            b_ssa = CuteDSLOpOverrides._ensure_tensor_ssa(
+                b, tensor_arg, is_tensor=b_is_tensor
+            )
             result_expr = op_format.format(a=a_ssa, b=b_ssa)
 
             dtype, bounds = CuteDSLOpOverrides._extract_dtype_and_bounds(a, b)
+            expected = CuteDSLOpOverrides._expected_tensor_val()
+            if dtype is None:
+                dtype = expected.dtype if expected is not None else torch.int32
+            if a_cse is not None:
+                shape = a_cse.shape
+            elif b_cse is not None:
+                shape = b_cse.shape
+            else:
+                shape = tuple(expected.size()) if expected is not None else None
 
             # Create and return CSEVariable using CSE generation for caching
             return V.kernel.cse.generate(
-                V.kernel.body, result_expr, bounds=bounds, dtype=dtype
+                V.kernel.body, result_expr, bounds=bounds, dtype=dtype, shape=shape
             )
 
         return op_format.format(a=a, b=b)
 
+    @staticmethod
+    def _expected_tensor_val() -> torch.Tensor | None:
+        """Return the fake-tensor value from the current FX node's metadata, if any."""
+        node = V.current_node
+        if not isinstance(node, torch.fx.Node):
+            return None
+        val = node.meta.get("val")
+        return val if isinstance(val, torch.Tensor) else None
+
+    @staticmethod
+    def _cast_expr(expr: str, dtype: torch.dtype) -> str:
+        cute_type = CuteDSLOpOverrides.TORCH_TO_CUTE_DTYPE.get(dtype)
+        if cute_type is None:
+            return expr
+        return f"{cute_type}({expr})"
+
     @staticmethod
     def _apply_unary_op(x: CuteDSLArg, op_format: str) -> CuteDSLArg:
         """
@@ -152,7 +215,7 @@ def _apply_unary_op(x: CuteDSLArg, op_format: str) -> CuteDSLArg:
         return op_format.format(x=x)
 
     @staticmethod
-    def constant(value: Union[bool, float, int], dtype: torch.dtype) -> str:
+    def constant(value: bool | float | int, dtype: torch.dtype) -> str:
         """Generate CuteDSL constant representation."""
         if value == float("-inf"):
             return "float('-inf')"
@@ -191,64 +254,128 @@ def remainder(a, b):
         return CuteDSLOpOverrides._apply_binary_op(a, b, "({a} % {b})")
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def exp(x: CuteDSLArg) -> CuteDSLArg:
-        """Exponential using CuteDSL cute.math.exp function."""
+        """Exponential using CuteDSL cute.math.exp2 with log2(e) scaling."""
+        if CuteDSLOpOverrides._get_cse_var(x) is None:
+            x = CuteDSLOpOverrides._cast_expr(str(x), torch.float32)
         return CuteDSLOpOverrides._apply_unary_op(
             x, f"cute.math.exp2({{x}} * {CuteDSLOpOverrides.LOG2_E})"
         )
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def sqrt(x: CuteDSLArg) -> CuteDSLArg:
         """Square root using CuteDSL cute.math.sqrt function."""
         return CuteDSLOpOverrides._apply_unary_op(x, "cute.math.sqrt({x})")
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def log(x: CuteDSLArg) -> CuteDSLArg:
         """Natural logarithm using CuteDSL cute.math.log function."""
         return CuteDSLOpOverrides._apply_unary_op(x, "cute.math.log({x})")
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def cos(x: CuteDSLArg) -> CuteDSLArg:
         """Cosine using CuteDSL cute.math.cos function."""
         return CuteDSLOpOverrides._apply_unary_op(x, "cute.math.cos({x})")
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def sin(x: CuteDSLArg) -> CuteDSLArg:
         """Sine using CuteDSL cute.math.sin function."""
         return CuteDSLOpOverrides._apply_unary_op(x, "cute.math.sin({x})")
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def erf(x: CuteDSLArg) -> CuteDSLArg:
         """Error function using CuteDSL cute.math.erf function."""
         return CuteDSLOpOverrides._apply_unary_op(x, "cute.math.erf({x})")
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
+    def sigmoid(x: CuteDSLArg) -> CuteDSLArg:
+        """Sigmoid with fp32 compute and cast-back to expected output dtype."""
+        x_cse = CuteDSLOpOverrides._get_cse_var(x)
+        if x_cse is not None:
+            x_fp32: CuteDSLArg = CuteDSLOpOverrides.to_dtype(
+                x_cse,
+                torch.float32,
+                use_compute_types=False,
+            )
+        else:
+            x_fp32 = CuteDSLOpOverrides._cast_expr(str(x), torch.float32)
+
+        result = CuteDSLOpOverrides._apply_unary_op(
+            x_fp32,
+            f"(1.0 / (1.0 + cute.math.exp2(-{{x}} * {CuteDSLOpOverrides.LOG2_E})))",
+        )
+
+        expected = CuteDSLOpOverrides._expected_tensor_val()
+        expected_dtype = expected.dtype if expected is not None else None
+        if expected_dtype is not None and expected_dtype != torch.float32:
+            result_cse = CuteDSLOpOverrides._get_cse_var(result)
+            if result_cse is not None:
+                return CuteDSLOpOverrides.to_dtype(
+                    result_cse,
+                    expected_dtype,
+                    use_compute_types=False,
+                )
+            return CuteDSLOpOverrides._cast_expr(str(result), expected_dtype)
+        return result
+
+    @staticmethod
+    def _minmax(a: CuteDSLArg, b: CuteDSLArg, *, op: str) -> CuteDSLArg:
+        tensor_arg = CuteDSLOpOverrides._get_cse_var(
+            a
+        ) or CuteDSLOpOverrides._get_cse_var(b)
+        if tensor_arg is not None:
+            return CuteDSLOpOverrides._apply_binary_op(
+                a, b, f"cute.where(({{a}}) {op} ({{b}}), {{a}}, {{b}})"
+            )
+
+        lhs = str(a)
+        rhs = str(b)
+        expected = CuteDSLOpOverrides._expected_tensor_val()
+        if expected is not None:
+            lhs = CuteDSLOpOverrides._cast_expr(lhs, expected.dtype)
+            rhs = CuteDSLOpOverrides._cast_expr(rhs, expected.dtype)
+        return f"({lhs} if {lhs} {op} {rhs} else {rhs})"
+
+    @staticmethod
+    # pyrefly: ignore [bad-override]
     def maximum(a: CuteDSLArg, b: CuteDSLArg) -> CuteDSLArg:
-        raise NotImplementedError("TODO: maximum is not supported yet for TensorSSA")
+        return CuteDSLOpOverrides._minmax(a, b, op=">")
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def minimum(a: CuteDSLArg, b: CuteDSLArg) -> CuteDSLArg:
-        raise NotImplementedError("TODO: minimum is not supported yet for TensorSSA")
+        return CuteDSLOpOverrides._minmax(a, b, op="<")
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def where(
         condition: CuteDSLArg,
         a: CuteDSLArg,
         b: CuteDSLArg,
     ) -> CuteDSLArg:
         """Conditional selection - handles CSEVariable, OpsValue, and string inputs."""
-        # Find a tensor argument to use as template for full_like
-        # Priority: use 'a' if it's a tensor, else use 'b', else condition
-        tensor_arg = (
-            CuteDSLOpOverrides._get_cse_var(a)
-            or CuteDSLOpOverrides._get_cse_var(b)
-            or CuteDSLOpOverrides._get_cse_var(condition)
-        )
+        a_cse = CuteDSLOpOverrides._get_cse_var(a)
+        b_cse = CuteDSLOpOverrides._get_cse_var(b)
+        cond_cse = CuteDSLOpOverrides._get_cse_var(condition)
+        tensor_arg = a_cse or b_cse or cond_cse
 
         if tensor_arg is not None:
-            a_ssa = CuteDSLOpOverrides._ensure_tensor_ssa(a, tensor_arg)
-            b_ssa = CuteDSLOpOverrides._ensure_tensor_ssa(b, tensor_arg)
-            cond_ssa = CuteDSLOpOverrides._ensure_tensor_ssa(condition, tensor_arg)
+            a_ssa = CuteDSLOpOverrides._ensure_tensor_ssa(
+                a, tensor_arg, is_tensor=a_cse is not None
+            )
+            b_ssa = CuteDSLOpOverrides._ensure_tensor_ssa(
+                b, tensor_arg, is_tensor=b_cse is not None
+            )
+            cond_ssa = CuteDSLOpOverrides._ensure_tensor_ssa(
+                condition, tensor_arg, is_tensor=cond_cse is not None
+            )
             result_expr = f"cute.where({cond_ssa}, {a_ssa}, {b_ssa})"
 
             dtype, bounds = CuteDSLOpOverrides._extract_dtype_and_bounds(
@@ -266,6 +393,7 @@ def pow(a: CuteDSLArg, b: CuteDSLArg):
         return CuteDSLOpOverrides._apply_binary_op(a, b, "({a} ** {b})")
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def abs(x: CuteDSLArg) -> CuteDSLArg:
         """Absolute value using CuteDSL cute.math.abs function."""
         if isinstance(x, CSEVariable):
@@ -287,11 +415,14 @@ def abs(x: CuteDSLArg) -> CuteDSLArg:
 
     @staticmethod
     def neg(x: CuteDSLArg) -> CuteDSLArg:
-        """Negation using CuteDSL TensorSSA __neg__ operator."""
-        # TODO: See https://github.com/NVIDIA/cutlass/issues/2584
-        return CuteDSLOpOverrides._apply_unary_op(
-            x, "cute.TensorSSA(-{x}, {x}.shape, {x}.dtype)"
-        )
+        """Negation for both TensorSSA and scalar-like expressions."""
+        # TensorSSA path: avoid relying on __neg__ directly due upstream issue.
+        if CuteDSLOpOverrides._get_cse_var(x) is not None:
+            return CuteDSLOpOverrides._apply_unary_op(
+                x, "cute.TensorSSA(-{x}, {x}.shape, {x}.dtype)"
+            )
+        # Scalar path: shape/dtype attributes are unavailable.
+        return CuteDSLOpOverrides._apply_unary_op(x, "(-{x})")
 
     @staticmethod
     def to_dtype(
@@ -304,8 +435,8 @@ def to_dtype(
 
         Raises NotImplementedError for unsigned integer and unsupported dtypes.
         """
-        # Always convert up from bf16 and fp16 TODO on configuring
-        dtype = upcast_compute_type(dtype)
+        if use_compute_types:
+            dtype = upcast_compute_type(dtype)
 
         cute_type = CuteDSLOpOverrides.TORCH_TO_CUTE_DTYPE.get(dtype)
         if cute_type is None:
diff --git a/torch/_inductor/codegen/cutedsl/cutedsl_template.py b/torch/_inductor/codegen/cutedsl/cutedsl_template.py
index f7f4f89b8ba8f..17f681a302730 100644
--- a/torch/_inductor/codegen/cutedsl/cutedsl_template.py
+++ b/torch/_inductor/codegen/cutedsl/cutedsl_template.py
@@ -2,7 +2,7 @@
 import functools
 import itertools
 from collections.abc import Iterable
-from typing import Any, Optional
+from typing import Any
 from unittest.mock import patch
 
 from torch._inductor.utils import Placeholder
@@ -29,8 +29,8 @@ def __init__(
         self,
         name: str,
         source: str,
-        subgraph_fn: Optional[Any] = None,
-        mask_fn: Optional[Any] = None,
+        subgraph_fn: Any | None = None,
+        mask_fn: Any | None = None,
     ) -> None:
         super().__init__(name)
         self.source = source
@@ -48,7 +48,7 @@ def _template_from_string(source: str) -> Any:
 
     def maybe_append_choice(
         self, choices: list[Any], **kwargs: Any
-    ) -> Optional[NotImplementedError]:
+    ) -> NotImplementedError | None:
         """
         Maybe generates a new ChoiceCaller and appends it into existing choices.
         Returns None if success, otherwise returns the error.
@@ -69,6 +69,7 @@ def generate(self, **kwargs: Any) -> ChoiceCaller:
         layout = kwargs.pop("layout")
         mutated_inputs = kwargs.pop("mutated_inputs", None)
         subgraphs = kwargs.pop("subgraphs", None)
+        template_kwargs = dict(kwargs)
 
         kernel_name = f"cutedsl_{self.name}_{next(self.index_counter)}"
 
@@ -98,7 +99,7 @@ def generate(self, **kwargs: Any) -> ChoiceCaller:
                 source_code=code,
             )
 
-            def make_kernel_render(out_node, hint_override: Optional[int] = None):
+            def make_kernel_render(out_node, hint_override: int | None = None):
                 """
                 Factory function that creates a kernel renderer for the final output.
 
@@ -126,6 +127,7 @@ def render():
                 bmreq=bmreq,
                 template=self,
                 mutated_inputs=mutated_inputs,
+                template_kwargs=template_kwargs,
             )
 
 
@@ -140,19 +142,29 @@ def __init__(
         make_kernel_render: Any,
         bmreq: CuteDSLBenchmarkRequest,
         template: "CuteDSLTemplate",
-        mutated_inputs: Optional[Iterable[IRNode]] = None,
+        mutated_inputs: Iterable[IRNode] | None = None,
+        template_kwargs: dict[str, Any] | None = None,
     ):
+        description = self._build_description(name, template_kwargs)
         super().__init__(
             name=name,
             input_nodes=input_nodes,
             layout=layout,
-            description=f"CuteDSL template {name}",
+            description=description,
         )
         self.make_kernel_render = make_kernel_render
         self.bmreq = bmreq
         self.template = template
         self.mutated_inputs = mutated_inputs
 
+    def _build_description(
+        self, name: str, template_kwargs: dict[str, Any] | None
+    ) -> str:
+        if not template_kwargs:
+            return f"CuteDSL template {name}"
+        kwargs_desc = ", ".join(f"{k}={v}" for k, v in template_kwargs.items())
+        return f"CuteDSL template {name} ({kwargs_desc})"
+
     def __str__(self) -> str:
         return f"CuteDSLTemplateCaller({self.name})"
 
diff --git a/torch/_inductor/codegen/cutlass/cache.py b/torch/_inductor/codegen/cutlass/cache.py
index 9de1a6257c2d1..b990f052e1d16 100644
--- a/torch/_inductor/codegen/cutlass/cache.py
+++ b/torch/_inductor/codegen/cutlass/cache.py
@@ -6,11 +6,10 @@
 import logging
 import os
 import time
-from typing import Any, Optional
+from typing import Any
 
 import torch._inductor.config as config
 from torch._inductor.codecache import cutlass_key
-from torch._inductor.codegen.cuda.cuda_env import get_cuda_arch, get_cuda_version
 from torch._inductor.codegen.cutlass import serialization, utils
 from torch._inductor.codegen.cutlass.serialization import (
     get_cutlass_operation_serializer,
@@ -27,11 +26,11 @@
 
 def get_config_request_key(
     arch: str,
-    cuda_version: str,
+    toolkit_version: str,
     instantiation_level: str,
 ) -> str:
     """
-    Return a key for the full ops, based on cutlass key, arch, cuda version, instantiation level, and serialization.py file hash.
+    Return a key for the full ops, based on cutlass key, arch, toolkit version, instantiation level, and serialization.py file hash.
     """
 
     # Get hash of serialization.py and cutlass_utils.py files using their module file paths
@@ -47,7 +46,7 @@ def get_file_hash(file_module):
         [
             cutlass_key().hex(),
             arch,
-            cuda_version,
+            toolkit_version,
             instantiation_level,
             serialization_hash,
             cutlass_utils_hash,
@@ -65,7 +64,7 @@ def _generate_config_filename(request_key: str) -> str:
 
 @clear_on_fresh_cache
 @functools.cache
-def maybe_fetch_ops() -> Optional[list[Any]]:
+def maybe_fetch_ops(device_type: str) -> list[Any] | None:
     """
     Fetch ops from databases.
     """
@@ -73,10 +72,12 @@ def maybe_fetch_ops() -> Optional[list[Any]]:
         return None
 
     # setup
-    arch: str = get_cuda_arch()
-    # get_cuda_version might return "12.4.0" or "12.4"
-    # but we want to use "12.4"
-    version: str = ".".join(get_cuda_version().split(".")[:2])
+    arch: str = utils.cutlass_arch(device_type)
+    version: str = utils.toolkit_version(device_type)
+    if device_type == "cuda":
+        # get_cuda_version might return "12.4.0" or "12.4"
+        # but we want to use "12.4"
+        version = ".".join(version.split(".")[:2])
     instantiation_level: str = config.cutlass.cutlass_instantiation_level
 
     # filename and filepath
@@ -85,7 +86,7 @@ def maybe_fetch_ops() -> Optional[list[Any]]:
     filepath: str = os.path.join(cache_dir(), filename)
 
     # try fetch
-    serialized_ops: Optional[list[str]] = None
+    serialized_ops: list[str] | None = None
     start_time = time.time()
     if os.path.isfile(filepath):
         # locally
diff --git a/torch/_inductor/codegen/cutlass/gemm_template.py b/torch/_inductor/codegen/cutlass/gemm_template.py
index 136144ce32d06..365ce8449320a 100644
--- a/torch/_inductor/codegen/cutlass/gemm_template.py
+++ b/torch/_inductor/codegen/cutlass/gemm_template.py
@@ -6,7 +6,7 @@
 import re
 import time
 from abc import ABC, abstractmethod
-from typing import Any, Optional, Union
+from typing import Any
 
 import torch
 import torch.utils._pytree as pytree
@@ -32,6 +32,7 @@
 from ...utils import is_dynamic, Placeholder
 from ...virtualized import V
 from ..common import IndentedBuffer
+from ..cuda import cuda_env
 from . import utils as cutlass_utils
 from .kernel import CUTLASSTemplateKernel
 from .python_evt import CutlassEVTCodegen, scaled_mm_evt
@@ -430,8 +431,8 @@ def __init__(
         layout: Layout,
         alpha: float,
         beta: float,
-        input_reorder: Optional[list[int]] = None,
-        use_fast_accum: Optional[bool] = None,
+        input_reorder: list[int] | None = None,
+        use_fast_accum: bool | None = None,
     ) -> None:
         """
         Args:
@@ -461,10 +462,10 @@ def add_cutlass_gemm_choices(
         choices: list[ChoiceCaller],
         layout: ir.Layout,
         input_nodes: list[Buffer],
-        alpha: Union[float, int] = 1,
-        beta: Union[float, int] = 0,
-        input_reorder: Optional[list[int]] = None,
-        use_fast_accum: Optional[bool] = None,
+        alpha: float | int = 1,
+        beta: float | int = 0,
+        input_reorder: list[int] | None = None,
+        use_fast_accum: bool | None = None,
         **extra_kwargs,
     ) -> None:
         raise NotImplementedError
@@ -487,7 +488,7 @@ def _get_template(self) -> str:
     def _get_template_args(
         self,
         op: "cutlass_library.gemm_op.GemmOperation",  # type: ignore[name-defined]  # noqa: F821
-    ) -> tuple[str, Optional[str]]:
+    ) -> tuple[str, str | None]:
         raise NotImplementedError
 
     @abstractmethod
@@ -519,7 +520,7 @@ def _set_bias_layout_and_alignment(
     def _define_gemm_instance(
         self,
         op: GemmOperation,
-        evt_name: Optional[str] = None,
+        evt_name: str | None = None,
     ) -> tuple[str, str]:
         raise NotImplementedError
 
@@ -527,7 +528,7 @@ def _define_gemm_instance(
     def _get_extra_inputs_and_names(
         self,
         op: "cutlass_gemm_op.GemmOperation" = None,  # type: ignore[name-defined]  # noqa: F821
-    ) -> tuple[Optional[Buffer], list[Optional[Buffer]], list[str]]:
+    ) -> tuple[Buffer | None, list[Buffer | None], list[str]]:
         raise NotImplementedError
 
     @abstractmethod
@@ -543,9 +544,9 @@ def _add_cutlass_gemm_choices(
         choices: list[ChoiceCaller],
         layout: ir.Layout,
         input_nodes: list[Buffer],
-        alpha: Union[float, int] = 1,
-        beta: Union[float, int] = 0,
-        input_reorder: Optional[list[int]] = None,
+        alpha: float | int = 1,
+        beta: float | int = 0,
+        input_reorder: list[int] | None = None,
         **extra_kwargs,
     ) -> None:
         """
@@ -568,11 +569,20 @@ def _add_cutlass_gemm_choices(
 
         # pre-computation
         layout_repr: str = str(layout)
-        input_tensor_meta: Union[TensorMeta, list[TensorMeta]] = (
-            TensorMeta.from_irnodes(self.input_nodes)
+        input_tensor_meta: TensorMeta | list[TensorMeta] = TensorMeta.from_irnodes(
+            self.input_nodes
         )
-        output_tensor_meta: Union[TensorMeta, list[TensorMeta]] = (
-            TensorMeta.from_irnodes(self.output_node)
+        # When input_reorder is set (e.g. [2, 0, 1] for addmm), the kernel
+        # function signature is reordered (e.g. from [X, W, Bias] to
+        # [Bias, X, W]).  input_tensor_meta must follow the same order
+        # because subprocess benchmarking creates tensors from this metadata
+        # and passes them positionally to the compiled kernel.  Without this
+        # reorder the kernel receives mismatched pointers/strides, causing
+        # out-of-bounds GPU memory access for large shapes.
+        if self.input_reorder is not None and isinstance(input_tensor_meta, list):
+            input_tensor_meta = [input_tensor_meta[idx] for idx in self.input_reorder]
+        output_tensor_meta: TensorMeta | list[TensorMeta] = TensorMeta.from_irnodes(
+            self.output_node
         )
 
         with dynamo_timed("CUTLASSGemmTemplate.maybe_append_choice"):
@@ -622,7 +632,6 @@ def header(self) -> IndentedBuffer:
                 #include "cutlass/gemm/device/gemm_universal.h"
                 #include "cutlass/gemm/device/gemm_universal_adapter.h"
                 #include "cutlass/gemm/kernel/gemm_universal.hpp"
-                #include "cutlass/gemm/device/gemm_sparse.h"
                 #include "cutlass/gemm/collective/collective_builder.hpp"
                 #include "cutlass/epilogue/collective/collective_builder.hpp"
                 #include "cutlass/epilogue/collective/default_epilogue.hpp"
@@ -636,6 +645,13 @@ def header(self) -> IndentedBuffer:
                 #include "cutlass/util/tensor_view_io.h"
             """
         )
+        if self.device_type != "xpu":
+            # XPU SYCL-TLA does not support sparse gemm yet
+            res.splice(
+                """
+                #include "cutlass/gemm/device/gemm_sparse.h"
+                """
+            )
         if inductor_cutlass_config.generate_test_runner and not is_dynamic(
             *self.input_nodes, self.output_node
         ):
@@ -643,7 +659,7 @@ def header(self) -> IndentedBuffer:
         return res
 
     @staticmethod
-    def cutlass_layout(torch_layout: ir.Layout) -> "Optional[cutlass_lib.LayoutType]":  # type: ignore[name-defined]  # noqa: F821
+    def cutlass_layout(torch_layout: ir.Layout) -> "cutlass_lib.LayoutType | None":  # type: ignore[name-defined]  # noqa: F821
         """
         Converts an ir.Layout instance into the corresponding cutlass_library.LayoutType enum value
         (RowMajor, ColumnMajor, or None if no matching value is found ).
@@ -713,12 +729,14 @@ def set_alignment(torch_layout, op_element) -> bool:
             bool: True if the alignment was successfully updated, False otherwise.
         """
         alignment = cutlass_utils.get_max_alignment(torch_layout)
-        cuda_arch = cutlass_utils.get_cuda_arch()
-        if cuda_arch and int(cuda_arch) >= 90 and alignment < op_element.alignment:
-            return False
-        else:
-            op_element.alignment = alignment
-            return True
+        if torch.cuda.is_available():
+            cuda_arch = cuda_env.get_cuda_arch()
+            cuda_arch = cutlass_utils._normalize_cuda_arch(cuda_arch)
+            if cuda_arch and int(cuda_arch) >= 90 and alignment < op_element.alignment:
+                return False
+
+        op_element.alignment = alignment
+        return True
 
     @staticmethod
     def should_swap_XW(
@@ -767,8 +785,8 @@ def fix_op_layout(
         op: "cutlass_library.gemm_op.GemmOperation",  # type: ignore[name-defined] # noqa: F821
         X: Buffer,
         W: Buffer,
-        Bias: Optional[Buffer],
-        Y: Union[Buffer, ReinterpretView],
+        Bias: Buffer | None,
+        Y: Buffer | ReinterpretView,
     ) -> "cutlass_library.gemm_op.GemmOperation":  # type: ignore[name-defined]  # noqa: F821
         # This is a workaround to deal with cases where the input layouts have changed
         # between autotuning and rendering. This happens if the inputs layout
@@ -989,6 +1007,12 @@ def filter_op(
             ):
                 return None
 
+        # `_procedural_name` is decorated with @functools.cached_property in cutlass, and its value is
+        # cached based on the key `self`. After we modify some attributes of
+        # `self` (e.g., layout or alignment), the `self` itself doesn’t change, so the
+        # cached value remains stale. We therefore need to clear the cached value so that
+        # `_procedural_name` can be recomputed with the updated attributes.
+        del op._procedural_name
         return op
 
     def gen_ops(self) -> "list[tuple[str, cutlass_gemm_op.GemmOperation]]":  # type: ignore[name-defined]  # noqa: F821
@@ -1010,10 +1034,10 @@ def gen_ops(self) -> "list[tuple[str, cutlass_gemm_op.GemmOperation]]":  # type:
             return self.filtered_ops_cache[self.cache_key]
 
         with dynamo_timed("CUTLASSGemmTemplate.maybe_fetch_ops"):
-            maybe_ops = maybe_fetch_ops()
+            maybe_ops = maybe_fetch_ops(self.device_type)
         if maybe_ops is None:
             log.debug("Cannot fetch ops from cache, generating ops from scratch")
-            full_ops = cutlass_utils.gen_ops()
+            full_ops = cutlass_utils.gen_ops(self.device_type)
             ops = pytree.tree_flatten(full_ops)[0]
         else:
             log.debug("Using cached ops from cache")
@@ -1095,13 +1119,13 @@ def render(  # type: ignore[override]
         self,
         kernel: CUTLASSTemplateKernel,
         op: "cutlass_gemm_op.GemmOperation" = None,  # type: ignore[name-defined]  # noqa: F821
-        template_buffer_node: Optional[CUTLASSTemplateBuffer] = None,
-        epilogue_nodes: Optional[list[BaseSchedulerNode]] = None,
+        template_buffer_node: CUTLASSTemplateBuffer | None = None,
+        epilogue_nodes: list[BaseSchedulerNode] | None = None,
         **kwargs,
     ) -> str:
         """
         The primary entry point for the code rendering process used in this template.
-        Renders the Cutlass based CUDA C++ code for the GEMM Kernel that this template is designed to implement,
+        Renders the Cutlass based CUDA/XPU C++ code for the GEMM Kernel that this template is designed to implement,
         including potentially fused epilogues.
 
         Args:
@@ -1111,7 +1135,7 @@ def render(  # type: ignore[override]
             **kwargs: Additional keyword arguments. Currently unused.
 
         Returns:
-            str: Cutlass based CUDA C++ code fragment as a string, to be used by the current
+            str: Cutlass based CUDA/XPU C++ code fragment as a string, to be used by the current
             CUTLASSTemplateKernel or autotuning code.
 
         Note:
@@ -1332,7 +1356,7 @@ def test_call_statement(
         names_str: str = "",
     ) -> str:
         """
-        Helper method to render the Cutlass CUDA C++ code required for calling the GEMM operation in the standalone
+        Helper method to render the Cutlass CUDA/XPU C++ code required for calling the GEMM operation in the standalone
         test runner that might also be generated along with the rest of the code, if the corresponding config is
         enabled.
 
@@ -1372,10 +1396,10 @@ def add_cutlass_gemm_choices(
         choices: list[ChoiceCaller],
         layout: ir.Layout,
         input_nodes: list[Buffer],
-        alpha: Union[float, int] = 1,
-        beta: Union[float, int] = 0,
-        input_reorder: Optional[list[int]] = None,
-        use_fast_accum: Optional[bool] = None,
+        alpha: float | int = 1,
+        beta: float | int = 0,
+        input_reorder: list[int] | None = None,
+        use_fast_accum: bool | None = None,
         **extra_kwargs,
     ) -> None:
         template = CUTLASS3xGemmTemplate(
@@ -1403,7 +1427,7 @@ def _get_template(self) -> str:
     def _get_template_args(
         self,
         op: "cutlass_library.gemm_op.GemmOperation",  # type: ignore[name-defined]  # noqa: F821
-    ) -> tuple[str, Optional[str]]:
+    ) -> tuple[str, str | None]:
         return (GEMM_ARGS_CUTLASS_3X, GEMM_ARGS_CUTLASS_3X_EPILOGUE)
 
     @staticmethod
@@ -1448,8 +1472,8 @@ def _are_inputs_layout_compatible(self, layouts: list[Layout]) -> bool:
             return False
         if len(B_layout.size) < 1:
             return False
-        A_size = list(V.graph.sizevars.size_hints(A_layout.size))
-        B_size = list(V.graph.sizevars.size_hints(B_layout.size))
+        A_size = list(V.graph.sizevars.guarding_hints_or_throw(A_layout.size))
+        B_size = list(V.graph.sizevars.guarding_hints_or_throw(B_layout.size))
         if len(A_size) < 2:
             A_size.insert(0, 1)
         if len(B_size) < 2:
@@ -1472,7 +1496,7 @@ def _are_inputs_layout_compatible(self, layouts: list[Layout]) -> bool:
                 return False
         if len(layouts) == 3:
             C_layout = layouts[2]
-            C_size = [V.graph.sizevars.size_hint(i) for i in C_layout.size]
+            C_size = list(V.graph.sizevars.guarding_hints_or_throw(C_layout.size))
             while len(C_size) < len(A_size):
                 C_size.insert(0, 1)
             # check batch dims
@@ -1512,10 +1536,18 @@ def _render_evt(
         acc_dtype = torch_dtype_to_cutlass_type(accumulator_dtype)
         output_dtype = torch_dtype_to_cutlass_type(output_dtype)
 
+        # TODO: size_hint_fn is passed to both create_example_tensors (just for
+        # tracing examples) and trace -> _render_argument_type -> _get_arg_from_node
+        # where stride values are baked as C++ literals into the generated CUTLASS
+        # argument struct. For the latter, accessing hint is wrong: even for
+        # backed symbols we access the hint without installing a guard, so the
+        # generated strides could be incorrect for different dynamic shape inputs.
+        # guarding_hint_or_throw would be wrong here — it does not install
+        # guards either. This should use guard_int to properly guard on the values.
         examples = create_example_tensors(
             var_name_to_buffer_name,
             name_to_buffer,  # type: ignore[arg-type]
-            V.graph.sizevars.size_hint,
+            V.graph.sizevars.optimization_hint,
         )
         evt_name, evt_args, evt_code, arg_renames = trace(
             evt_py_code,
@@ -1525,7 +1557,8 @@ def _render_evt(
             op.tile_description,  # type: ignore[attr-defined]
             op.epilogue_schedule,  # type: ignore[attr-defined]
             {k: name_to_buffer[v] for k, v in var_name_to_buffer_name.items()},  # type: ignore[arg-type,misc]
-            V.graph.sizevars.size_hint,
+            V.graph.sizevars.guarding_hint_or_throw,
+            kernel_schedule=op.kernel_schedule,
         )
 
         return (
@@ -1576,9 +1609,9 @@ def _set_bias_layout_and_alignment(
     def _define_gemm_instance(
         self,
         op: GemmOperation,
-        evt_name: Optional[str] = None,
+        evt_name: str | None = None,
     ) -> tuple[str, str]:
-        """Defines and renders the Cutlass / CUDA C++ code for a given GEMM operation instance.
+        """Defines and renders the Cutlass / CUDA/XPU C++ code for a given GEMM operation instance.
 
         This function uses the Cutlass library to generate key parts of the codegen process. General Matrix Multiply
         forms a core part of a number of scientific applications, so this efficient and adaptable implementation is
@@ -1596,7 +1629,9 @@ def _define_gemm_instance(
 
         from .lib_extensions import gemm_operation_extensions as gemm_extensions
 
-        emitter = gemm_extensions.EmitGemmUniversal3xInstanceWithEVT(evt_name=evt_name)  # type: ignore[call-arg]
+        emitter = gemm_extensions.EmitGemmUniversal3xInstanceWithEVT(
+            evt_name=evt_name, device_type=self.device_type
+        )  # type: ignore[call-arg]
 
         if not hasattr(op, "epilogue_functor") or not isinstance(
             op.epilogue_functor, enum.Enum
@@ -1621,9 +1656,9 @@ def _define_gemm_instance(
     def _get_extra_inputs_and_names(
         self,
         op: "cutlass_gemm_op.GemmOperation" = None,  # type: ignore[name-defined]  # noqa: F821
-    ) -> tuple[Optional[Buffer], list[Optional[Buffer]], list[str]]:
+    ) -> tuple[Buffer | None, list[Buffer | None], list[str]]:
         Bias = self.input_nodes[2] if len(self.input_nodes) == 3 else None
-        inputs: list[Optional[Buffer]] = []
+        inputs: list[Buffer | None] = []
         names: list[str] = []
         return (Bias, inputs, names)
 
@@ -1657,7 +1692,7 @@ def render_gemm_arguments(
         epilogue_args,
     ) -> str:
         """
-        Render the Cutlass CUDA C++ code required for passing arguments to the GEMM operation.
+        Render the Cutlass CUDA/XPU C++ code required for passing arguments to the GEMM operation.
 
         Args:
             argument_template (str): Template for the GEMM operation arguments.
@@ -1670,11 +1705,11 @@ def render_gemm_arguments(
             Y (IRNode): The output tensor.
             alpha (float): Scaling factor for the product of the inputs.
             beta (float): Scaling factor for the output tensor.
-            kernel (CUTLASSTemplateKernel): CUDA Template kernel for the operation.
+            kernel (CUTLASSTemplateKernel): CUDA/XPU Template kernel for the operation.
             epilogue_args (any): Additional arguments for the epilogue state.
 
         Returns:
-            str: A block of CUDA C++ code as a string, ready to be used as arguments for the GEMM operation.
+            str: A block of CUDA/XPU C++ code as a string, ready to be used as arguments for the GEMM operation.
 
         Note: If `should_swap_xw` is True, a transpose operation will be applied to the X, W, Bias, and Y
         tensors. This operation also implies the M and N dimensions of Bias and GEMM output to be swapped
@@ -1742,7 +1777,7 @@ def __init__(
         layout: Layout,
         alpha: float,
         beta: float,
-        input_reorder: Optional[list[int]] = None,
+        input_reorder: list[int] | None = None,
     ):
         super().__init__(input_nodes, layout, alpha, beta, input_reorder)
 
@@ -1751,10 +1786,10 @@ def add_cutlass_gemm_choices(
         choices: list[ChoiceCaller],
         layout: ir.Layout,
         input_nodes: list[Buffer],
-        alpha: Union[float, int] = 1,
-        beta: Union[float, int] = 0,
-        input_reorder: Optional[list[int]] = None,
-        use_fast_accum: Optional[bool] = False,
+        alpha: float | int = 1,
+        beta: float | int = 0,
+        input_reorder: list[int] | None = None,
+        use_fast_accum: bool | None = False,
         **extra_kwargs,
     ) -> None:
         template = CUTLASS2xGemmTemplate(
@@ -1780,7 +1815,7 @@ def _get_template(self) -> str:
     def _get_template_args(
         self,
         op: "cutlass_library.gemm_op.GemmOperation",  # type: ignore[name-defined]  # noqa: F821
-    ) -> tuple[str, Optional[str]]:
+    ) -> tuple[str, str | None]:
         import cutlass_library.library as cutlass_lib
 
         if op.gemm_kind == cutlass_lib.GemmKind.Sparse:
@@ -1866,7 +1901,7 @@ def _set_bias_layout_and_alignment(
     def _define_gemm_instance(
         self,
         op: GemmOperation,
-        evt_name: Optional[str] = None,
+        evt_name: str | None = None,
     ) -> tuple[str, str]:
         """Defines and renders the Cutlass / CUDA C++ code for a given GEMM operation instance.
 
@@ -1907,7 +1942,7 @@ def _define_gemm_instance(
     def _get_extra_inputs_and_names(
         self,
         op: "cutlass_gemm_op.GemmOperation" = None,  # type: ignore[name-defined]  # noqa: F821
-    ) -> tuple[Optional[Buffer], list[Optional[Buffer]], list[str]]:
+    ) -> tuple[Buffer | None, list[Buffer | None], list[str]]:
         import cutlass_library.library as cutlass_lib
 
         if op.gemm_kind == cutlass_lib.GemmKind.Sparse:
diff --git a/torch/_inductor/codegen/cutlass/kernel.py b/torch/_inductor/codegen/cutlass/kernel.py
index 615d1c4b4c76e..5498fc7129083 100644
--- a/torch/_inductor/codegen/cutlass/kernel.py
+++ b/torch/_inductor/codegen/cutlass/kernel.py
@@ -5,12 +5,13 @@
 from collections import defaultdict
 from collections.abc import Callable
 from dataclasses import dataclass
-from typing import Any, Literal, Optional, TYPE_CHECKING, Union
+from typing import Any, Literal, TYPE_CHECKING
 
 from sympy import Expr, symbols
 
 import torch._inductor.config as config
 from torch import dtype as torch_dtype
+from torch._inductor.codegen.common import get_device_op_overrides
 from torch._inductor.codegen.cpp_wrapper_cpu import CppWrapperCpu
 from torch._inductor.scheduler import BaseSchedulerNode
 from torch._inductor.utils import do_bench_using_profiling, OrderedSet, Placeholder
@@ -82,19 +83,17 @@ class CUTLASSKernel(Kernel):
     def __init__(self, *args, **kwargs) -> None:
         super().__init__(*args, **kwargs)
         self.layout_args: dict[str, list[LayoutArg]] = defaultdict(list)
-        self.size_args: list[Union[Expr, int]] = []
+        self.size_args: list[Expr | int] = []
         # Mapping from arg name to IRNode.
         self.named_nodes: dict[str, IRNode] = {}
 
-    def find_symbol(
-        self, node: IRNode, attr: ValidLayoutAttrs, dim: int
-    ) -> Optional[str]:
+    def find_symbol(self, node: IRNode, attr: ValidLayoutAttrs, dim: int) -> str | None:
         arg = self.find_layout_arg(node, attr, dim)
         return arg.symbol if arg else None
 
     def find_layout_arg(
         self, node: IRNode, attr: ValidLayoutAttrs, dim: int
-    ) -> Optional[LayoutArg]:
+    ) -> LayoutArg | None:
         matches = [
             arg
             for arg in itertools.chain.from_iterable(self.layout_args.values())
@@ -151,7 +150,7 @@ def init_layout_args(self) -> None:
             self.add_layout_arg("ldc", Bias, "stride", ldc_dim)
         self.add_layout_arg("ldd", Y, "stride", ldd_dim)
 
-    def get_layout_args(self) -> tuple[Union[Expr, int], ...]:
+    def get_layout_args(self) -> tuple[Expr | int, ...]:
         X = self.named_nodes["X"]
         W = self.named_nodes["W"]
         Y = self.named_nodes["Y"]
@@ -160,7 +159,7 @@ def get_layout_args(self) -> tuple[Union[Expr, int], ...]:
         ndim = _normalize_idx(-1, len(W.get_size()))
         kdim = _normalize_idx(-1, len(X.get_size()))
 
-        def get_ld(node) -> Union[Expr, int]:
+        def get_ld(node) -> Expr | int:
             dim = self.find_ld_idx(node)
             return node.get_stride()[dim]
 
@@ -174,7 +173,7 @@ def get_ld(node) -> Union[Expr, int]:
         LDD = get_ld(Y)
         return (M, N, K, B, LDA, LDB, LDC, LDD)
 
-    def get_dynamic_shape_args(self) -> list[Union[Expr, int]]:
+    def get_dynamic_shape_args(self) -> list[Expr | int]:
         return [*self.get_layout_args(), *self.size_args]
 
     def get_offset_args(self) -> list[Expr]:
@@ -196,13 +195,12 @@ class CUTLASSTemplateKernel(CUTLASSKernel):
     Template kernels defined by Cutlass in C++.
     """
 
-    _EXTRA_CPP_ARGS = "size_t* workspace_size, uint8_t* workspace, cudaStream_t stream"
-
     def __init__(
         self,
         kernel_name: str,
         runtime_arg_info: list["ArgInfo"],
         runtime_arg_values: list[Any],
+        device_type: str = "cuda",  # type: ignore[assignment]
     ) -> None:
         """
         Initializes a new instance of the CUTLASSTemplateKernel class.
@@ -214,6 +212,9 @@ def __init__(
         self.kernel_name = kernel_name
         self.runtime_arg_info = runtime_arg_info
         self.runtime_arg_values = runtime_arg_values
+        self.device_type = device_type
+        self.device_codegen = get_device_op_overrides(self.device_type)
+        self._EXTRA_CPP_ARGS = f"size_t* workspace_size, uint8_t* workspace, {self.device_codegen.cpp_stream_type()} stream"
 
     def check_not_null(self, node: IRNode) -> str:
         """
@@ -251,7 +252,7 @@ def def_kernel(
         inputs: list[IRNode],
         outputs: list[IRNode],
         names_str: str = "",
-        input_reorder: Optional[list[int]] = None,
+        input_reorder: list[int] | None = None,
     ) -> str:
         """
         Hook called from template code to generate function definition and
@@ -410,7 +411,7 @@ def call_kernel(
         if ws:
             wrapper.generate_workspace_deallocation(ws)
 
-    def dtype(self, node: IRNode) -> Optional[str]:
+    def dtype(self, node: IRNode) -> str | None:
         """
         Generates code which represents dtype of a given node.
         """
@@ -419,7 +420,7 @@ def dtype(self, node: IRNode) -> Optional[str]:
             return "void"
         return DTYPE_TO_CPP.get(node.get_layout().dtype)
 
-    def cutlass_dtype(self, node: IRNode, default_dtype="void") -> Optional[str]:
+    def cutlass_dtype(self, node: IRNode, default_dtype="void") -> str | None:
         # Helper method, called into from CUTLASSGemmTemplate
         if node is None:
             return default_dtype
@@ -452,7 +453,7 @@ def size(
         self,
         node: IRNode,
         start_index: int,
-        end_index: Optional[int] = None,
+        end_index: int | None = None,
         default_value: int = 0,
     ) -> str:
         """
@@ -581,15 +582,13 @@ def __init__(
         input_nodes: list[Buffer],
         layout: Layout,
         make_kernel_render: Callable[
-            [CUTLASSTemplateBuffer, Optional[list[BaseSchedulerNode]]],
+            [CUTLASSTemplateBuffer, list[BaseSchedulerNode] | None],
             tuple[CUTLASSTemplateKernel, functools.partial[str]],
         ],
         bmreq: CUTLASSBenchmarkRequest,
         supports_epilogue_fusion: bool,
         template: "CUTLASSTemplate",  # type: ignore[name-defined]
-        info_kwargs: Optional[
-            dict[str, Union[PrimitiveInfoType, list[PrimitiveInfoType]]]
-        ],  # type: ignore[type-arg]
+        info_kwargs: dict[str, PrimitiveInfoType | list[PrimitiveInfoType]] | None,  # type: ignore[type-arg]
         description: str,
     ) -> None:
         super().__init__(name, input_nodes, layout, description)
@@ -645,7 +644,7 @@ def hash_key(self) -> str:
             ]
         )
 
-    def info_dict(self) -> dict[str, Union[PrimitiveInfoType, list[PrimitiveInfoType]]]:
+    def info_dict(self) -> dict[str, PrimitiveInfoType | list[PrimitiveInfoType]]:
         """
         Information returned here is logged to the autotune log file when that is enabled.
 
diff --git a/torch/_inductor/codegen/cutlass/lib_extensions/evt_extensions.py b/torch/_inductor/codegen/cutlass/lib_extensions/evt_extensions.py
index b454c6e8fff37..5357dc0bd98a5 100644
--- a/torch/_inductor/codegen/cutlass/lib_extensions/evt_extensions.py
+++ b/torch/_inductor/codegen/cutlass/lib_extensions/evt_extensions.py
@@ -1,5 +1,5 @@
 from collections.abc import Callable
-from typing import Any, Union
+from typing import Any
 
 from sympy import Expr
 
@@ -14,7 +14,7 @@
 
 
 EpilogueFunctor = Any  # EpilogueFunctor local class defined in _trace
-Buffer = Union[ComputedBuffer, InputBuffer]
+Buffer = ComputedBuffer | InputBuffer
 CutlassTupleType = Any  # cutlass.backend.c_types.tuple_factory_.<locals>.TupleType
 CutlassVisitorType = Any  # cutlass.backend.c_types.visitor_factory.<locals>.VisitorType
 CutlassArgType = (
@@ -26,7 +26,6 @@
     import ast
     import ctypes
     import textwrap
-    from typing import Union
 
     from cutlass_cppgen.backend.c_types import (  # type: ignore[import-not-found]
         EmptyByte,
@@ -40,6 +39,9 @@
     from cutlass_cppgen.backend.evt.backend.emitter_base import (  # type: ignore[import-not-found]
         FusionCallbacks,
     )
+    from cutlass_cppgen.backend.evt.backend.sm100_emitter import (  # type: ignore[import-not-found]
+        Sm100CollectiveEpilogue,
+    )
     from cutlass_cppgen.backend.evt.backend.sm90_emitter import (  # type: ignore[import-not-found]
         CollectiveEpilogue,
     )
@@ -81,7 +83,7 @@ def get(self, name: str) -> str:
     def create_example_tensors(
         var_name_to_buffer_name: dict[str, str],
         name_to_buffer: dict[str, Buffer],
-        size_hint_fn: Callable[[Union[Expr, int]], int],
+        size_hint_fn: Callable[[Expr | int], int],
     ) -> dict[str, CutlassTensor]:
         def cutlass_tensor_from_buffer(
             buffer: Buffer,
@@ -121,7 +123,8 @@ def trace(
         tile_description: TileDescription,
         epilogue_schedule: EpilogueScheduleType,
         name_to_buffer: dict[str, Buffer],
-        size_hint_fn: Callable[[Union[Expr, int]], int],
+        size_hint_fn: Callable[[Expr | int], int],
+        kernel_schedule: Any | None = None,
         **kwargs: dict[str, Any],
     ) -> tuple[str, str, str, EVTArgRenames]:
         cuda_arch = int(cuda_env.get_cuda_arch())  # type: ignore[arg-type]
@@ -129,13 +132,23 @@ def trace(
         epilogue_functor = _trace(fn_src, example_tensors, cuda_arch, **kwargs)
         visitor = EpilogueFunctorVisitor(cuda_arch, epilogue_functor)
         fusion_callbacks = FusionCallbacks(visitor.graph, cuda_arch, emit_CD=False)
-        collective_epilogue = CollectiveEpilogue(
-            tile_description,
-            epilogue_schedule,
-            accum_type,
-            output_type,
-            fusion_callbacks,
-        )
+        if cuda_arch < 100:
+            collective_epilogue = CollectiveEpilogue(
+                tile_description,
+                epilogue_schedule,
+                accum_type,
+                output_type,
+                fusion_callbacks,
+            )
+        else:
+            collective_epilogue = Sm100CollectiveEpilogue(
+                tile_description=tile_description,
+                kernel_schedule=kernel_schedule,
+                epilogue_schedule=epilogue_schedule,
+                element_accumulator=accum_type,
+                element_d=output_type,
+                fusion_callbacks=fusion_callbacks,
+            )
         evt_name, evt_code = collective_epilogue.emit()
         evt_args, arg_renames = _render_argument_type(
             epilogue_functor, name_to_buffer, size_hint_fn
@@ -174,7 +187,7 @@ def parse(
     def _render_argument_type(
         epilogue_functor: EpilogueFunctor,
         name_to_buffer: dict[str, Buffer],
-        size_hint_fn: Callable[[Union[Expr, int]], int],
+        size_hint_fn: Callable[[Expr | int], int],
     ) -> tuple[str, EVTArgRenames]:
         epilogue_thread_type = epilogue_functor.epilogue_thread_type
         arg_renames = EVTArgRenames()
@@ -234,7 +247,7 @@ def render_thread_type(name: str, t: CutlassArgType) -> None:
     def _get_arg_from_node(
         arg_ty: type,
         node: Buffer,
-        size_hint_fn: Callable[[Union[Expr, int]], int],
+        size_hint_fn: Callable[[Expr | int], int],
         arg_renames: EVTArgRenames,
     ) -> str:
         from ..template import CUTLASSTemplate
diff --git a/torch/_inductor/codegen/cutlass/lib_extensions/gemm_operation_extensions.py b/torch/_inductor/codegen/cutlass/lib_extensions/gemm_operation_extensions.py
index d10669d40bea0..f0f229136b295 100644
--- a/torch/_inductor/codegen/cutlass/lib_extensions/gemm_operation_extensions.py
+++ b/torch/_inductor/codegen/cutlass/lib_extensions/gemm_operation_extensions.py
@@ -16,7 +16,7 @@
     class EmitGemmUniversal3xInstanceWithEVT:
         """Responsible for emitting a CUTLASS 3.x template definition"""
 
-        def __init__(self, operation_suffix="", evt_name=None):
+        def __init__(self, operation_suffix="", evt_name=None, device_type="cuda"):
             self.operation_suffix = operation_suffix
             self.includes = [
                 "cutlass/cutlass.h",
@@ -32,8 +32,8 @@ def __init__(self, operation_suffix="", evt_name=None):
             ${element_c},
             ${element_epilogue}
             >"""
-
             self.evt_name = evt_name
+            self.device_type = device_type
             self.gemm_template = """
 using ${operation_name}_epilogue =
 typename cutlass::epilogue::collective::CollectiveBuilder<
@@ -175,6 +175,8 @@ def emit(self, operation):
                     f"cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(\
 sizeof(typename {str(operation.procedural_name())}_epilogue::SharedStorage))>"
                 )
+            if self.device_type == "xpu":
+                stage_count_string = "cutlass::gemm::collective::StageCountAuto"
 
             epi_tile_mn = "cutlass::epilogue::collective::EpilogueTileAuto"
 
@@ -350,6 +352,11 @@ def emit(self, operation):
             if self.evt_name:
                 epilogue_functor = self.evt_name
 
+            if self.device_type == "xpu":
+                arch = f"cutlass::arch::Xe{operation.arch}"
+            else:
+                arch = f"cutlass::arch::Sm{operation.arch}"
+
             values = {
                 "operation_name": operation_name_str,
                 "operation_suffix": self.operation_suffix,
@@ -369,7 +376,7 @@ def emit(self, operation):
                 "element_accumulator": DataTypeTag[operation.accumulator_type()],
                 "opcode_class_main": OpcodeClassTag[opcode_class_main],
                 "opcode_class_epi": OpcodeClassTag[opcode_class_epi],
-                "arch": f"cutlass::arch::Sm{operation.arch}",
+                "arch": arch,
                 "tile_shape_m": str(tile_shape_m),
                 "tile_shape_n": str(tile_shape_n),
                 "tile_shape_k": str(tile_shape_k),
diff --git a/torch/_inductor/codegen/cutlass/python_evt.py b/torch/_inductor/codegen/cutlass/python_evt.py
index bbc759cc929ef..26e4a9b9cc062 100644
--- a/torch/_inductor/codegen/cutlass/python_evt.py
+++ b/torch/_inductor/codegen/cutlass/python_evt.py
@@ -2,7 +2,7 @@
 from collections.abc import Generator, Iterable, Iterator, Sequence
 from contextlib import contextmanager
 from os import linesep
-from typing import Any, Optional
+from typing import Any
 
 import sympy
 
@@ -21,7 +21,7 @@
 
 
 def scaled_mm_evt(
-    scale_A_name: str, scale_B_name: str, bias_name: Optional[str], output_name: str
+    scale_A_name: str, scale_B_name: str, bias_name: str | None, output_name: str
 ) -> tuple[list[str], dict[str, Any], str]:
     evt_read_names = [scale_A_name, scale_B_name]
     var_name_to_buffer_name = {n: n for n in [scale_A_name, scale_B_name]}
@@ -57,7 +57,7 @@ def _prefix_un_op(op: str, a: str) -> str:
     def to_dtype(
         x: str,
         dtype: Any,
-        src_dtype: Optional[torch.dtype] = None,
+        src_dtype: torch.dtype | None = None,
         use_compute_types: bool = False,
     ) -> str:
         return x
@@ -166,9 +166,10 @@ def __init__(self, accumulator_node_name: str, removed_buffers: OrderedSet[str])
             _ACCUMULATOR_ARG_NAME: accumulator_node_name
         }
         self.removed_buffers: OrderedSet[str] = removed_buffers
-        self.cur_node: Optional[ComputedBuffer] = None
+        self.cur_node: ComputedBuffer | None = None
         self.name_to_buffer = V.graph.name_to_buffer | V.graph.graph_inputs
         for name in V.graph.constants:
+            # pyrefly: ignore [unsupported-operation]
             self.name_to_buffer[name] = V.graph.add_tensor_constant(
                 V.graph.constants[name], name
             )
diff --git a/torch/_inductor/codegen/cutlass/scheduling.py b/torch/_inductor/codegen/cutlass/scheduling.py
index 74f3cd9504482..933a6ed33593b 100644
--- a/torch/_inductor/codegen/cutlass/scheduling.py
+++ b/torch/_inductor/codegen/cutlass/scheduling.py
@@ -2,7 +2,7 @@
 import hashlib
 import logging
 from collections.abc import Sequence
-from typing import cast
+from typing import cast, TypeGuard
 
 from torch._inductor.codegen.cutlass.python_evt import (
     CutlassEVTCodegen,
@@ -40,7 +40,7 @@ class CUTLASSScheduling(BaseScheduling):
     """
     Partial Scheduling implementation for cutlass C++ Kernels.
     This class is intended to be used in combination with TritonScheduling,
-    and delegated to by CUDACombinedScheduling.
+    and delegated to by CUDACombinedScheduling/XPUCombinedScheduling.
 
     It handles fusion decisions and cutlass C++ specific template code generation.
     """
@@ -53,7 +53,7 @@ def group_fn(self, sizes):
         return tuple(V.graph.sizevars.simplify(sympy_product(s)) for s in sizes)
 
     @staticmethod
-    def is_cutlass_template(node: BaseSchedulerNode) -> bool:
+    def is_cutlass_template(node: BaseSchedulerNode) -> TypeGuard[SchedulerNode]:
         return isinstance(node, SchedulerNode) and isinstance(
             node.node, CUTLASSTemplateBuffer
         )
@@ -109,7 +109,7 @@ def define_kernel(self, src_code: str, node_schedule) -> str:
             _, _, kernel_path = get_path(code_hash(src_code), "py")
 
             compile_wrapper = IndentedBuffer()
-            compile_wrapper.writeline("async_compile.cuda(r'''")
+            compile_wrapper.writeline(f"async_compile.{V.graph.device_type}(r'''")
             compile_wrapper.splice(src_code, strip=True)
             compile_wrapper.writeline(
                 f"''', 'so', aot_compile={str(V.graph.aot_mode)})"
@@ -136,7 +136,6 @@ def codegen_template(
         assert self.is_cutlass_template(template_node), (
             "Template node passed to CUTLASSScheduling.codegen_template must be a SchedulerNode that wraps a CUTLASSTemplateBuffer"
         )
-        template_node = cast(SchedulerNode, template_node)
         _, (_numel, rnumel) = template_node.group
         assert rnumel == 1
         ctb: CUTLASSTemplateBuffer = cast(CUTLASSTemplateBuffer, template_node.node)
diff --git a/torch/_inductor/codegen/cutlass/serialization.py b/torch/_inductor/codegen/cutlass/serialization.py
index 39184e4e6e2c6..c1649b1b2fffa 100644
--- a/torch/_inductor/codegen/cutlass/serialization.py
+++ b/torch/_inductor/codegen/cutlass/serialization.py
@@ -256,7 +256,7 @@ def _tile_description_to_json(cls, tile_desc: "TileDescription") -> str:  # type
     @classmethod
     @functools.lru_cache(None)
     def _json_to_tile_description(
-        cls, json_dict: Optional[str]
+        cls, json_dict: str | None
     ) -> Optional["TileDescription"]:  # type: ignore[name-defined]  # noqa: F821
         """
         Convert JSON dict to TileDescription object.
@@ -313,7 +313,7 @@ def _json_to_tile_description(
     def _math_instruction_to_json(
         cls,
         math_instruction: Optional["MathInstruction"],  # type: ignore[name-defined]  # noqa: F821
-    ) -> Optional[str]:
+    ) -> str | None:
         """Convert MathInstruction to JSON string.
 
         Args:
@@ -344,7 +344,7 @@ def _math_instruction_to_json(
     @classmethod
     @functools.lru_cache(None)
     def _json_to_math_instruction(
-        cls, json_dict: Optional[str]
+        cls, json_dict: str | None
     ) -> Optional["MathInstruction"]:  # type: ignore[name-defined]  # noqa: F821
         """Convert JSON string to MathInstruction object.
 
@@ -399,7 +399,7 @@ def _json_to_math_instruction(
     def _tensor_description_to_json(
         cls,
         tensor_desc: Optional["TensorDescription"],  # type: ignore[name-defined]  # noqa: F821
-    ) -> Optional[str]:
+    ) -> str | None:
         """Convert TensorDescription to JSON string.
 
         Args:
@@ -424,8 +424,8 @@ def _tensor_description_to_json(
     @functools.lru_cache(None)
     def _json_to_tensor_description(
         cls,
-        json_dict: Optional[str],
-        tensor_name: Optional[str] = None,
+        json_dict: str | None,
+        tensor_name: str | None = None,
     ) -> Optional["TensorDescription"]:  # type: ignore[name-defined]  # noqa: F821
         """Convert JSON string to TensorDescription object.
 
@@ -459,7 +459,7 @@ def _json_to_tensor_description(
 
     @classmethod
     @functools.lru_cache(None)
-    def _enum_to_json(cls, enum_value: Optional[Enum]) -> Optional[str]:
+    def _enum_to_json(cls, enum_value: Enum | None) -> str | None:
         """Convert enum value to JSON string.
 
         Args:
@@ -480,7 +480,7 @@ def _enum_to_json(cls, enum_value: Optional[Enum]) -> Optional[str]:
 
     @classmethod
     @functools.lru_cache(None)
-    def _json_to_enum(cls, json_dict: Optional[str], enum_class: Any) -> Optional[Enum]:
+    def _json_to_enum(cls, json_dict: str | None, enum_class: Any) -> Enum | None:
         """Convert JSON string to enum value.
 
         Format: {name: "EnumName", value: 1}
@@ -501,7 +501,7 @@ def _json_to_enum(cls, json_dict: Optional[str], enum_class: Any) -> Optional[En
 
 
 @functools.lru_cache(1)
-def get_cutlass_operation_serializer() -> Optional[CUTLASSOperationSerializer]:
+def get_cutlass_operation_serializer() -> CUTLASSOperationSerializer | None:
     if not try_import_cutlass():
         return None
     return CUTLASSOperationSerializer()
diff --git a/torch/_inductor/codegen/cutlass/template.py b/torch/_inductor/codegen/cutlass/template.py
index 5511dd4b0d153..189338cb8f2d7 100644
--- a/torch/_inductor/codegen/cutlass/template.py
+++ b/torch/_inductor/codegen/cutlass/template.py
@@ -3,7 +3,7 @@
 import hashlib
 import itertools
 from dataclasses import dataclass
-from typing import Any, Optional, TYPE_CHECKING, Union
+from typing import Any, TYPE_CHECKING
 from unittest.mock import patch
 
 import sympy
@@ -55,8 +55,7 @@ def __init__(
         name: str,
         input_nodes: list[Buffer],
         layout: Layout,
-        input_reorder: Optional[list[int]] = None,
-        device_type: str = "cuda",
+        input_reorder: list[int] | None = None,
     ) -> None:
         """
         Baseclass for CUTLASS C++ Templates, derived from KernelTemplate.
@@ -74,7 +73,7 @@ def __init__(
         self.output_node: Buffer = Buffer(name="buf_out", layout=layout)
         self.input_reorder = input_reorder
         self.layout = layout
-        self.device_type = device_type
+        self.device_type = layout.device.type
 
     @classmethod
     @functools.lru_cache(None)
@@ -115,7 +114,7 @@ def generate_code_and_args(
         Generate code and args with caching. We cache the code even if runtime
         args are different.
         """
-        key: Optional[str] = None
+        key: str | None = None
         if config.cutlass.enable_caching_codegen:
             key = self.make_key(name=name, input_key=input_key, layout_repr=layout_repr)
 
@@ -133,6 +132,7 @@ def generate_code_and_args(
             kernel_name=kernel_name,
             runtime_arg_info=self.get_runtime_arg_info(),
             runtime_arg_values=self.get_runtime_arg_values(**kwargs),
+            device_type=self.device_type,
         )
         with patch.object(V.graph, "get_dtype", self._fake_get_dtype(self.output_node)):
             code = self.render(kernel=kernel, **kwargs)
@@ -157,9 +157,12 @@ def generate_code_and_args(
             call_args,
             expected_args,
         )
-        V.graph.sizevars.size_hints(map(sympy.expand, call_args[len(expected_args) :]))
-        size_args = V.graph.sizevars.size_hints(kernel.get_dynamic_shape_args())
-        offset_args = V.graph.sizevars.size_hints(kernel.get_offset_args())
+        # Resolve symbolic sizes to concrete ints for benchmarking only.
+        V.graph.sizevars.optimization_hints(
+            map(sympy.expand, call_args[len(expected_args) :])
+        )
+        size_args = V.graph.sizevars.optimization_hints(kernel.get_dynamic_shape_args())
+        offset_args = V.graph.sizevars.optimization_hints(kernel.get_offset_args())
 
         if key is not None:
             self.code_cache[key] = code, size_args, offset_args
@@ -177,8 +180,8 @@ def generate(  # type: ignore[override]
         description: str,
         input_key: str,
         layout_repr: str,
-        input_tensor_meta: Union[TensorMeta, list[TensorMeta]],
-        output_tensor_meta: Union[TensorMeta, list[TensorMeta]],
+        input_tensor_meta: TensorMeta | list[TensorMeta],
+        output_tensor_meta: TensorMeta | list[TensorMeta],
         **kwargs,
     ) -> CUTLASSTemplateCaller:
         """
@@ -225,7 +228,7 @@ def generate(  # type: ignore[override]
 
         def make_kernel_render(
             template_node: CUTLASSTemplateBuffer,
-            epilogue_nodes: Optional[list[BaseSchedulerNode]] = None,
+            epilogue_nodes: list[BaseSchedulerNode] | None = None,
         ) -> tuple[CUTLASSTemplateKernel, functools.partial[str]]:
             assert supports_epilogue_fusion or not epilogue_nodes, (
                 "epilogue fusion is not supported for this kernel"
@@ -234,6 +237,7 @@ def make_kernel_render(
                 kernel_name=str(Placeholder.KERNEL_NAME),
                 runtime_arg_info=self.get_runtime_arg_info(),
                 runtime_arg_values=self.get_runtime_arg_values(**kwargs),
+                device_type=self.device_type,
             )
             render = functools.partial(
                 self.render,
diff --git a/torch/_inductor/codegen/cutlass/utils.py b/torch/_inductor/codegen/cutlass/utils.py
index da4b7d2842d36..5fac9f392abb9 100644
--- a/torch/_inductor/codegen/cutlass/utils.py
+++ b/torch/_inductor/codegen/cutlass/utils.py
@@ -7,7 +7,7 @@
 import sys
 import time
 from dataclasses import dataclass
-from typing import Any, Optional
+from typing import Any
 from typing_extensions import TypeIs
 
 import sympy
@@ -22,6 +22,7 @@
 from ...runtime.runtime_utils import cache_dir
 from ...virtualized import V
 from ..cuda.cuda_env import get_cuda_arch, get_cuda_version
+from ..xpu.xpu_env import get_xpu_arch, get_xpu_version
 
 
 log = logging.getLogger(__name__)
@@ -82,7 +83,7 @@ def try_import_cutlass() -> bool:
         except ImportError as e:
             log.warning(  # noqa: G200
                 "Failed to import CUTLASS packages in fbcode: %s, ignoring the CUTLASS backend.",
-                str(e),
+                e,
             )
             return False
 
@@ -99,7 +100,12 @@ def path_join(path0, path1):
 
     # contains both cutlass and cutlass_library
     # we need cutlass for eVT
-    cutlass_python_path = path_join(config.cutlass.cutlass_dir, "python")
+    cutlass_dir = (
+        config.xpu.cutlass_dir
+        if torch.xpu._is_compiled()
+        else config.cutlass.cutlass_dir
+    )
+    cutlass_python_path = path_join(cutlass_dir, "python")
     torch_root = os.path.abspath(os.path.dirname(torch.__file__))
     mock_src_path = os.path.join(
         torch_root,
@@ -167,7 +173,7 @@ def link_and_append(dst_link, src_path, parent_dir):
         except ImportError as e:
             log.debug(  # noqa: G200
                 "Failed to import CUTLASS packages: %s, ignoring the CUTLASS backend.",
-                str(e),
+                e,
             )
     else:
         log.debug(
@@ -177,7 +183,15 @@ def link_and_append(dst_link, src_path, parent_dir):
     return False
 
 
-@functools.lru_cache(8)
+def _normalize_xpu_arch(arch: str) -> str:
+    if arch.startswith("Xe"):
+        return arch[2:]
+    if 12 <= int(arch) and int(arch) <= 50:
+        return arch
+    else:
+        raise NotImplementedError(f"Unsupported xpu arch: {arch}")
+
+
 def _normalize_cuda_arch(arch: str) -> str:
     arch_num = arch
     if isinstance(arch, str):
@@ -206,16 +220,34 @@ def _normalize_cuda_arch(arch: str) -> str:
     raise NotImplementedError(f"Unsupported cuda arch: {arch}")
 
 
+@functools.lru_cache(8)
+def cutlass_arch(device_type: str) -> str:
+    if device_type == "xpu":
+        arch = get_xpu_arch()
+        return _normalize_xpu_arch(arch)
+    else:
+        arch = get_cuda_arch()
+        return _normalize_cuda_arch(arch)
+
+
+@functools.lru_cache(1)
+def toolkit_version(device_type: str) -> str:
+    if device_type == "xpu":
+        return get_xpu_version()
+    else:
+        return get_cuda_version()
+
+
 @dataclass
 class CUTLASSArgs:
     """
     CUTLASS args used to initialize a CUTLASS Manifest.
     """
 
-    architectures: Optional[str] = None
-    cuda_version: Optional[str] = None
-    instantiation_level: Optional[str] = None
-    operations: Optional[str] = None
+    architectures: str | None = None
+    toolkit_version: str | None = None
+    instantiation_level: str | None = None
+    operations: str | None = None
 
     build_dir = ""
     curr_build_dir = ""
@@ -229,18 +261,18 @@ class CUTLASSArgs:
     interface_dir: None = None
     filter_by_cc = True
     disable_full_archs_compilation = False
+    device_type: str = "cuda"
 
     def __post_init__(self):
-        if self.architectures is None or self.cuda_version is None:
+        if self.architectures is None or self.toolkit_version is None:
             raise RuntimeError(
-                f"{self.architectures=} or {self.cuda_version=} is None!"
+                f"{self.architectures=} or {self.toolkit_version=} is None!"
             )
-        self.architectures = _normalize_cuda_arch(self.architectures)
 
 
 @clear_on_fresh_cache
 @functools.cache
-def _gen_ops_cached(arch, version) -> dict[Any, Any]:
+def _gen_ops_cached(arch: str, version: str, device_type: str) -> dict[Any, Any]:
     # Note: Cache needs to be specific for cuda architecture and version
 
     # Import cutlass python scripts.
@@ -250,35 +282,46 @@ def _gen_ops_cached(arch, version) -> dict[Any, Any]:
 
     if arch is None or version is None:
         log.error(
-            "Cannot detect cuda arch %s or cuda version %s. "
+            "Cannot detect cuda arch %s or version %s. "
             "Will discard all cutlass ops. "
             "Please consider setting _inductor.cuda.arch and _inductor.cuda.version configs.",
             arch,
             version,
         )
         return {}
-    arch = _normalize_cuda_arch(arch)
+
     gen_arch = (
         "100" if arch == "103" else arch
     )  # CUTLASS SM103 generator only covers NVFB4; fallback to SM100 set
     instantiation_level: str = config.cutlass.cutlass_instantiation_level
     args = CUTLASSArgs(
         architectures=gen_arch,
-        cuda_version=version,
+        toolkit_version=version,
         instantiation_level=instantiation_level,
         operations=CUTLASS_OPERATION_KIND,
+        device_type=device_type,
     )
     manifest = cutlass_manifest.Manifest(args)
 
     start_time = time.time()
-    if gen_arch == "100":
+    if device_type == "xpu":
+        if hasattr(cutlass_generator, "GenerateIntelXe"):
+            cutlass_generator.GenerateIntelXe(
+                manifest, args.toolkit_version, arch=int(arch)
+            )
+        else:
+            raise NotImplementedError(
+                "Arch " + arch + " is not supported by current cutlass lib."
+            )
+
+    elif arch == "100":
         if hasattr(cutlass_generator, "GenerateSM100"):
-            cutlass_generator.GenerateSM100(manifest, args.cuda_version)
-        cutlass_generator.GenerateSM90(manifest, args.cuda_version)
+            cutlass_generator.GenerateSM100(manifest, args.toolkit_version)
+        cutlass_generator.GenerateSM90(manifest, args.toolkit_version)
     else:
         try:
             func = getattr(cutlass_generator, "GenerateSM" + gen_arch)
-            func(manifest, args.cuda_version)
+            func(manifest, args.toolkit_version)
         except AttributeError as e:
             raise NotImplementedError(
                 "Arch " + gen_arch + " is not supported by current cutlass lib."
@@ -292,26 +335,35 @@ def _gen_ops_cached(arch, version) -> dict[Any, Any]:
     return manifest.operations
 
 
-def gen_ops() -> dict[Any, Any]:
+def gen_ops(device_type: str) -> dict[Any, Any]:
     """
     Generates all supported CUTLASS operations.
     """
     with dynamo_timed("cutlass_utils.gen_ops"):
-        arch = get_cuda_arch()
-        version = get_cuda_version()
-        return _gen_ops_cached(arch, version)
+        arch = cutlass_arch(device_type)
+        version = toolkit_version(device_type)
+        return _gen_ops_cached(arch, version, device_type)
 
 
 from ..cpp_utils import DTYPE_TO_CPP
 
 
-DTYPE_TO_CUTLASS_TYPE = {
-    **DTYPE_TO_CPP,
-    torch.float16: "__half",
-    torch.bfloat16: "__nv_bfloat16",
-    torch.float8_e4m3fn: "__nv_fp8_e4m3",
-    torch.float8_e5m2: "__nv_fp8_e5m2",
-}
+if torch.xpu._is_compiled():
+    DTYPE_TO_CUTLASS_TYPE = {
+        **DTYPE_TO_CPP,
+        torch.float16: "uint16_t",
+        torch.bfloat16: "uint16_t",
+        torch.float8_e4m3fn: "uint8_t",
+        torch.float8_e5m2: "uint8_t",
+    }
+else:
+    DTYPE_TO_CUTLASS_TYPE = {
+        **DTYPE_TO_CPP,
+        torch.float16: "__half",
+        torch.bfloat16: "__nv_bfloat16",
+        torch.float8_e4m3fn: "__nv_fp8_e4m3",
+        torch.float8_e5m2: "__nv_fp8_e5m2",
+    }
 
 
 @functools.lru_cache(32)
@@ -328,13 +380,17 @@ def torch_dtype_to_cutlass_type(
         return cutlass_library.library.DataType.f16
     elif torch_dtype == torch.bfloat16:
         return cutlass_library.library.DataType.bf16
+    elif torch_dtype == torch.float8_e4m3fn:
+        return cutlass_library.library.DataType.e4m3
+    elif torch_dtype == torch.float8_e5m2:
+        return cutlass_library.library.DataType.e5m2
     else:
         raise NotImplementedError(f"Unsupported data type: {torch_dtype=}")
 
 
 @functools.lru_cache(32)
 def dtype_match(
-    torch_dtype: Optional[torch.dtype],
+    torch_dtype: torch.dtype | None,
     cutlass_dtype: "cutlass_library.library.DataType",  # type: ignore[name-defined]  # noqa: F821
 ) -> bool:
     # Import cutlass python scripts.
@@ -366,7 +422,7 @@ def dtype_match(
 
 def get_accumulator_dtype(
     input_torch_dtypes: list[torch.dtype],
-) -> Optional[torch.dtype]:
+) -> torch.dtype | None:
     """
     Given a pair of input torch dtypes, returns the inferred accumulator torch dtype.
     """
diff --git a/torch/_inductor/codegen/debug_utils.py b/torch/_inductor/codegen/debug_utils.py
index 9b465e3d1ffab..24f2c88add6b8 100644
--- a/torch/_inductor/codegen/debug_utils.py
+++ b/torch/_inductor/codegen/debug_utils.py
@@ -5,7 +5,7 @@
 import logging
 import os
 from enum import Enum
-from typing import Optional, TYPE_CHECKING
+from typing import TYPE_CHECKING
 
 import torch
 from torch import dtype as torch_dtype
@@ -61,11 +61,11 @@ def __init__(
         self,
         debug_printer_level,
         use_array_ref: bool,
-        writeline: Optional[Callable[..., None]] = None,
-        args_to_print_or_save: Optional[list[str]] = None,
+        writeline: Callable[..., None] | None = None,
+        args_to_print_or_save: list[str] | None = None,
         kernel_name: str = "",
         kernel=None,
-        arg_signatures: Optional[list[type]] = None,
+        arg_signatures: list[type] | None = None,
         kernel_type=None,
     ):
         self.debug_printer_level = IntermediateValueDebuggingLevel(debug_printer_level)
@@ -74,7 +74,7 @@ def __init__(
             args_to_print_or_save = []
         self.args_to_print_or_save = args_to_print_or_save
         self.kernel_name = kernel_name
-        self.arg_signatures: Optional[list[type]] = None
+        self.arg_signatures: list[type] | None = None
         self.kernel = kernel
         self.filtered_kernel_names_to_print = self._get_debug_filtered_kernel_names()
         self.kernel_type = None
@@ -100,7 +100,7 @@ def _perform_debug_print_or_save_helper(
         args_to_print_or_save,
         kernel_name,
         before_launch,
-        arg_signatures: Optional[list[type]] = None,
+        arg_signatures: list[type] | None = None,
     ):
         if self.debug_printer_level == IntermediateValueDebuggingLevel.OFF:
             return
@@ -144,7 +144,7 @@ def set_printer_args(
         self,
         args_to_print_or_save: list[str],
         kernel_name: str,
-        arg_signatures: Optional[list[type]],
+        arg_signatures: list[type] | None,
         kernel,
         kernel_type=None,
     ):
@@ -196,7 +196,7 @@ def codegen_intermediate_tensor_value_save(
         args_to_save,
         kernel_name,
         before_launch=True,
-        arg_signatures: Optional[list[type]] = None,
+        arg_signatures: list[type] | None = None,
     ) -> None:
         for i, arg in enumerate(args_to_save):
             if arg_signatures is not None and not isinstance(
@@ -233,7 +233,7 @@ def codegen_intermediate_tensor_value_print(
         args_to_print,
         kernel_name,
         before_launch=True,
-        arg_signatures: Optional[list[type]] = None,
+        arg_signatures: list[type] | None = None,
     ) -> None:
         launch_prefix = "before_launch" if before_launch else "after_launch"
 
diff --git a/torch/_inductor/codegen/halide.py b/torch/_inductor/codegen/halide.py
index 06e57dab76d31..f782dea55fd4e 100644
--- a/torch/_inductor/codegen/halide.py
+++ b/torch/_inductor/codegen/halide.py
@@ -6,9 +6,9 @@
 import itertools
 import logging
 import re
+import sys
 from collections import defaultdict
-from math import inf
-from typing import Any, cast, Optional, TYPE_CHECKING, Union
+from typing import Any, cast, TYPE_CHECKING
 
 import sympy
 
@@ -268,7 +268,7 @@ class HalideOverrides(OpOverrides):
     def to_dtype(
         x,
         dtype: torch.dtype,
-        src_dtype: Optional[torch.dtype] = None,
+        src_dtype: torch.dtype | None = None,
         use_compute_types=True,
     ):
         if dtype == torch.bool:
@@ -289,20 +289,24 @@ def constant(cls, value, dtype):
         return cls.to_dtype(halide_constant(value), dtype)
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def abs(x):
         return f"hl.abs({x})"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def exp(x):
         if not hasattr(x, "name"):
             return f"hl.exp({x})"
         return f"hl.fast_exp(hl.cast(hl.Float(32), {x})) if {x.name}.type().bits() <= 32 else hl.exp({x})"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def sqrt(x):
         return f"hl.sqrt({x})"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def minimum(a, b):
         # return f"hl.min({a}, {b})"  <== handles nan wrong
         if not hasattr(a, "name"):
@@ -311,6 +315,7 @@ def minimum(a, b):
         return f"hl.select(({a}<{b})|hl.is_nan({a}), {a}, {b}) if {a.name}.type().is_float() else hl.min({a}, {b})"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def maximum(a, b):
         # return f"hl.max({a}, {b})"  <== handles nan wrong
         if not hasattr(a, "name"):
@@ -319,80 +324,99 @@ def maximum(a, b):
         return f"hl.select(({a}>{b})|hl.is_nan({a}), {a}, {b}) if {a.name}.type().is_float() else hl.max({a}, {b})"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def where(a, b, c):
         if hasattr(b, "name"):
             c = f"hl.cast({b.name}.type(), {c})"
         return f"hl.select({a}, {b}, {c})"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def cos(x):
         return f"hl.cos({x})"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def sin(x):
         return f"hl.sin({x})"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def lgamma(x):
         raise Unsupported("lgamma")
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def erf(x):
         return f"hl.erf({x})"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def cosh(x):
         return f"hl.cosh({x})"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def sinh(x):
         return f"hl.sinh({x})"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def acos(x):
         return f"hl.acos({x})"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def acosh(x):
         return f"hl.acosh({x})"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def asin(x):
         return f"hl.asin({x})"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def asinh(x):
         return f"hl.asinh({x})"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def atan2(x, y):
         return f"hl.atan2({x}, {y})"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def atan(x):
         return f"hl.atan({x})"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def atanh(x):
         return f"hl.atanh({x})"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def copysign(x, y):
         raise Unsupported("copysign")
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def erfinv(x):
         raise Unsupported("erfinv")
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def hypot(x, y):
         return f"hl.hypot({x}, {y})"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def nextafter(x, y):
         raise Unsupported("nextafter")
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def logical_and(a, b):
         return f"{a} & {b}"
 
@@ -401,10 +425,12 @@ def logical_not(a):
         return f"{a} == 0"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def logical_or(a, b):
         return f"{a} | {b}"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def logical_xor(a, b):
         return f"({a} ^ {b})"
 
@@ -440,6 +466,10 @@ def rand(seed, offset):
     def randn(seed, offset):
         return f"halide_helpers.randn({seed}, {offset})"
 
+    @staticmethod
+    def rand_eager(seed, base_offset, threads_per_round, tid, vec):
+        return f"halide_helpers.rand_eager_kernel({seed}, {base_offset}, {threads_per_round}, {tid}, {vec})"
+
     @staticmethod
     def randint64(seed, offset, low, high):
         return f"halide_helpers.randint64({seed}, {offset}, {low}, {high})"
@@ -449,23 +479,28 @@ def load_seed(name, offset):
         return f"{ops.load(name, 0)} + {V.kernel.args.seed_offset('load_seed_offset', offset)}"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def rsqrt(x):
         # return f"hl.fast_inverse_sqrt({x})"  <== accuracy issues
         return f"1./hl.sqrt({x})"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def tan(x):
         return f"hl.tan({x})"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def tanh(x):
         return f"hl.tanh({x})"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def signbit(x):
         return f"(hl.reinterpret(hl.UInt(32), hl.cast(hl.Float(32), {x})) >> 31) != 0"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def fmod(a, b):
         # TODO(jansel): find a better way to do this, builtin % has wrong sign
         return f"{a} - hl.trunc({a}/{b})*{b}"
@@ -475,10 +510,12 @@ def pow(a, b):
         return f"hl.pow({a}, {b})"  # hl.fast_pow fails accuracy
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def ldexp(x, n):
         raise Unsupported("ldexp")
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def log(x):
         return f"hl.log({x})"  # hl.fast_log fails accuracy
 
@@ -487,20 +524,24 @@ def log2(x):
         raise NotImplementedError("log2")
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def isinf(x):
         # workaround https://github.com/halide/Halide/issues/8309
         return f"hl.is_inf(hl.cast(hl.Float(32), {x}))"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def isnan(x):
         # workaround https://github.com/halide/Halide/issues/8309
         return f"hl.is_nan(hl.cast(hl.Float(32), {x}))"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def round(x):
         return f"hl.round({x})"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def floor(x):
         return f"hl.floor({x})"
 
@@ -523,10 +564,12 @@ def sign(cls, x):
         return f"hl.cast({x.name}.type(), {sub})"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def trunc(x):
         return f"hl.trunc({x})"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def truncdiv(a, b):
         # this causes crashes with floating point exception, see test_div_zero_dim_cpu
         # return f"hl.div_round_to_zero({a}, {b})"
@@ -535,6 +578,7 @@ def truncdiv(a, b):
         )
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def ceil(x):
         return f"hl.ceil({x})"
 
@@ -590,6 +634,7 @@ def masked(mask, body, other):
         return ops.where(new_mask, result, other)
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def frexp(x):
         raise NotImplementedError("frexp")
 
@@ -618,11 +663,11 @@ def __init__(
         self,
         name,
         bounds: ValueRanges[Any],
-        dtype: Optional[torch.dtype] = None,
+        dtype: torch.dtype | None = None,
         shape: BlockShapeType = None,
     ) -> None:
         super().__init__(name, bounds, dtype, shape=shape)
-        self.used_dims: Optional[list[sympy.Symbol]] = None
+        self.used_dims: list[sympy.Symbol] | None = None
 
     def update_on_args(self, name, args, kwargs):
         used = OrderedSet(self.used_dims or ())
@@ -653,7 +698,7 @@ def subs_str(self, replacements):
 
 @dataclasses.dataclass
 class DimensionInfo:
-    expr: Optional[sympy.Expr]
+    expr: sympy.Expr | None
     size: sympy.Expr
     stride: sympy.Expr
 
@@ -685,32 +730,27 @@ def index_str(self, replacements=None, zero_vars=False):
 
 
 def eq(left, right):
-    if V.graph.sizevars.statically_known_equals(left, right):
-        return True
-    try:
-        a = V.graph.sizevars.size_hint_or_throw(left)
-        b = V.graph.sizevars.size_hint_or_throw(right)
-    except TypeError:  # unbacked symints
-        return False
-    if a == b:
-        V.graph.sizevars.check_equals(left, right)
-    return a == b
+    return V.graph.sizevars.guard_or_false(sympy.Eq(left, right))
 
 
 def lt(left, right):
-    if V.graph.sizevars.statically_known_lt(left, right):
+    """Compare sizes: only use on inputs known to be >= 0."""
+    if V.graph.sizevars.guard_or_false(sympy.Lt(left, right)):
         return True
-    try:
-        a = V.graph.sizevars.size_hint_or_throw(left)
-        b = V.graph.sizevars.size_hint_or_throw(right)
-    except TypeError:  # unbacked symints
-        gcd = sympy.gcd(left, right)
-        if gcd == left:
-            return left != right
-        return False
-    if a < b:
-        V.graph.sizevars.check_lt(left, right)
-    return a < b
+
+    # GCD fallback: if gcd(left, right) == left then left divides right,
+    # so left <= right.  Combined with left != right this gives left < right.
+    #
+    # TODO: This is NOT always sound for unbacked symints.
+    # e.g. lt(u0, 10*u0): gcd=u0, gcd==left, u0 != 10*u0 → returns True,
+    # but if u0=0 then 0 < 0 is False.  The >= 0 checks mitigate the
+    # negative case but the zero case remains.
+    # TODO shall we add a runtime assertion at least.
+    gcd = sympy.gcd(left, right)
+    if gcd == left:
+        return left != right
+
+    return False
 
 
 class HalideKernel(SIMDKernel):
@@ -767,7 +807,9 @@ def finalize_indexing(self, indices: Sequence[sympy.Expr]):
         assert not (
             self.index_replacements or self.halide_vars or self.reduction_renames
         )
-        size_hint = functools.partial(V.graph.sizevars.size_hint, fallback=inf)  # type: ignore[arg-type]
+        size_hint = functools.partial(
+            V.graph.sizevars.optimization_hint, fallback=sys.maxsize
+        )
         # pyrefly: ignore [bad-assignment]
         indices = dict.fromkeys(map(super().prepare_indexing, indices))
         all_used_symbols = OrderedSet[Any]()
@@ -790,7 +832,8 @@ def visit_modular_indexing(base, divisor, modulus):
                     node.root.lookup(
                         node.divisor * divisor,
                         V.graph.sizevars.evaluate_min(
-                            modulus, FloorDiv(node.length, divisor)
+                            modulus,
+                            FloorDiv(node.length, divisor),
                         ),
                     ).symbol()
                 )
@@ -847,7 +890,8 @@ def visit_floor_div(base, divisor):
                 handled_count += len(sizes_to_add)
                 assert sizes_to_add, nodes
                 end = divisor * functools.reduce(
-                    V.graph.sizevars.evaluate_max, sizes_to_add
+                    lambda a, b: V.graph.sizevars.evaluate_max(a, b),
+                    sizes_to_add,
                 )
                 sizes_to_add.extend(
                     [
@@ -1040,7 +1084,11 @@ def expr_to_dimension(expr, syms):
             dims.append(expr_to_dimension(expr, syms))
         for sym, expr in split_expr.items():
             dims.append(expr_to_dimension(expr, [sym]))
-        dims.sort(key=lambda d: V.graph.sizevars.size_hint(d.stride, fallback=inf))  # type: ignore[arg-type]
+        dims.sort(
+            key=lambda d: V.graph.sizevars.optimization_hint(
+                d.stride, fallback=sys.maxsize
+            )
+        )  # type: ignore[arg-type]
 
         if not dims:  # scalar load/store
             if self.has_indirect_indexing:
@@ -1227,8 +1275,8 @@ def reduction(
         dtype: torch.dtype,
         src_dtype: torch.dtype,
         reduction_type: ReductionType,
-        value: Union[CSEVariable, tuple[CSEVariable, ...]],
-    ) -> Union[CSEVariable, tuple[CSEVariable, ...]]:
+        value: CSEVariable | tuple[CSEVariable, ...],
+    ) -> CSEVariable | tuple[CSEVariable, ...]:
         """Codegen a reduction operation"""
         assert self.inside_reduction
         assert not self._load_mask
@@ -1451,7 +1499,7 @@ def arg_order(arg_tuple):
                 assert "in_ptr" in arg.name
                 return 0
 
-        result: list[tuple[Optional[str], KernelArgType]] = []
+        result: list[tuple[str | None, KernelArgType]] = []
         _, a, b, _ = self.args.python_argdefs()
         for call_str, arg in sorted(zip(a, b), key=arg_order):
             result.append((call_str, arg))
@@ -1612,14 +1660,14 @@ def update_index(m):
             # This causes crashes if our estimate is greater than the vector length
             # https://github.com/halide/Halide/issues/3103
             if isinstance(arg, SizeArg):
-                hint = V.graph.sizevars.size_hint(arg.expr, fallback=1)
+                hint = V.graph.sizevars.optimization_hint(arg.expr, fallback=1)
                 code.writeline(f"{arg.name}.set_estimate({hint})")
             else:
                 dims = self.buffer_dimensions[arg.name]
                 range_hints = []
                 for i, dim in enumerate(dims):
                     hint = self._autoscheduler_workarounds(
-                        V.graph.sizevars.size_hint(dim.size, fallback=1), dims
+                        V.graph.sizevars.optimization_hint(dim.size, fallback=1), dims
                     )
                     # pyrefly: ignore [bad-argument-type]
                     range_hints.append(f"hl.Range(0, {hint})")
diff --git a/torch/_inductor/codegen/memory_planning.py b/torch/_inductor/codegen/memory_planning.py
index dc606f20babfe..4f2f362b5b2b5 100644
--- a/torch/_inductor/codegen/memory_planning.py
+++ b/torch/_inductor/codegen/memory_planning.py
@@ -5,7 +5,7 @@
 import dataclasses
 import itertools
 import pprint
-from typing import Any, Optional, Protocol, TYPE_CHECKING
+from typing import Any, Protocol, TYPE_CHECKING
 
 import sympy
 
@@ -141,9 +141,9 @@ class Allocation(AllocationTreeNode):
     size_hint: int
     symbolic_size: sympy.Expr
     allocated: bool = False
-    pool: Optional[AllocationPool] = None
-    offset: Optional[sympy.Expr] = None
-    earliest_available: Optional[float] = None
+    pool: AllocationPool | None = None
+    offset: sympy.Expr | None = None
+    earliest_available: float | None = None
 
     def __post_init__(self) -> None:
         has_unbacked_sym = False
@@ -387,8 +387,8 @@ class AllocationPool:
     device: torch.device
     root: TemporalSplit
     can_expand: bool = True
-    restrict_live_range: Optional[LiveRange] = None
-    name: Optional[str] = None
+    restrict_live_range: LiveRange | None = None
+    name: str | None = None
     names_to_del: list[str] = dataclasses.field(default_factory=list)
     creation_cache: dict[str, str] = dataclasses.field(default_factory=dict)
 
@@ -553,7 +553,7 @@ def __init__(self, node: BufferLike):
         self.node = node
         self.names = [node.get_name()]
         self.is_output = False
-        self.allocation: Optional[Allocation] = None
+        self.allocation: Allocation | None = None
         self.live_range = LiveRange(float("inf"), -float("inf"))
 
     def update_usage(self, timestep: int):
@@ -592,7 +592,7 @@ class PoolMemoryPlanningLine(MemoryPlanningLine):
     """Abstract base class for {Alloc,Dealloc}FromPoolLine"""
 
     group: BufferGroup
-    timestep: Optional[int] = None
+    timestep: int | None = None
 
     @property
     def node(self):
@@ -653,7 +653,7 @@ class MemoryPlanner:
 
     wrapper: Any
     pools: AllocationPools = dataclasses.field(default_factory=AllocationPools)
-    buffer_groups: Optional[list[BufferGroup]] = None
+    buffer_groups: list[BufferGroup] | None = None
 
     def plan(self, lines: list[Any]) -> list[Any]:
         """Call all the memory planning passes in sequence"""
diff --git a/torch/_inductor/codegen/mps.py b/torch/_inductor/codegen/mps.py
index 05d0e84c681ad..a6661976387b4 100644
--- a/torch/_inductor/codegen/mps.py
+++ b/torch/_inductor/codegen/mps.py
@@ -7,7 +7,7 @@
 import logging
 import math
 from pathlib import Path
-from typing import Any, Optional, TYPE_CHECKING
+from typing import Any, TYPE_CHECKING
 
 import sympy
 from sympy.printing.precedence import PRECEDENCE
@@ -32,8 +32,6 @@
 
 
 if TYPE_CHECKING:
-    from typing import Union
-
     from ..ops_handler import ReductionType, StoreMode
     from ..scheduler import Scheduler, SchedulerNode
     from .common import OpVarT
@@ -53,7 +51,7 @@
 }
 
 
-def value_to_metal(val: Union[float, int, bool, str, CSEVariable]) -> str:
+def value_to_metal(val: float | int | bool | str | CSEVariable) -> str:
     if isinstance(val, float):
         if val == torch.inf:
             return "HUGE_VALF"
@@ -80,6 +78,9 @@ def _print_FloorDiv(self, expr: sympy.Expr) -> str:
 
     def _print_ModularIndexing(self, expr: sympy.Expr) -> str:
         x, div, mod = expr.args
+        # Workaround for Metal compiler bug with fused (x / A) % B, see PR 175481
+        use_safe_mod = div == 65536 and (mod & (mod - 1)) != 0
+
         x = self.doprint(x)
         if div != 1:
             div = self.doprint(div)
@@ -88,6 +89,8 @@ def _print_ModularIndexing(self, expr: sympy.Expr) -> str:
             else:
                 x = f"metal::floor({x}) / ({div})"
         mod = self.doprint(mod)
+        if use_safe_mod:
+            return f"c10::metal::safe_mod({x}, {mod})"
         return f"({x}) % ({mod})"
 
     def _print_Min(self, expr: sympy.Expr) -> str:
@@ -186,7 +189,7 @@ class MetalOverrides(OpOverrides):
     def to_dtype(
         x: CSEVariable,
         dtype: torch.dtype,
-        src_dtype: Optional[torch.dtype] = None,
+        src_dtype: torch.dtype | None = None,
         use_compute_types: bool = True,
     ) -> str:
         if dtype == torch.double:
@@ -203,7 +206,7 @@ def to_dtype_bitcast(
         return f"as_type<{DTYPE_TO_METAL[dtype]}>(static_cast<{DTYPE_TO_METAL[src_dtype]}>({x}))"
 
     @staticmethod
-    def constant(val: Union[bool, float, int], dtype: torch.dtype) -> str:
+    def constant(val: bool | float | int, dtype: torch.dtype) -> str:
         return value_to_metal(val)
 
     @staticmethod
@@ -226,7 +229,8 @@ def masked(mask: CSEVariable, body: sympy.Expr, other: CSEVariable) -> str:
             # generates identical variable names. Without this reset, repeated calls to
             # body() would keep incrementing the counter, resulting in different cache key.
             V.kernel.cse.iter_buffer_ids = itertools.count()
-            V.kernel.cse.name_prefix = "tmp_scoped_"
+            # Append "_scoped" to the current prefix so each nesting level gets unique vars
+            V.kernel.cse.name_prefix += "_scoped"
             rc = body()
 
         # Compute cache key manually as variable name is needed to actually generate the code
@@ -240,63 +244,79 @@ def masked(mask: CSEVariable, body: sympy.Expr, other: CSEVariable) -> str:
             )
             with V.kernel.compute.indent():
                 V.kernel.compute.splice(scoped_body)
-                V.kernel.compute.writeline(f"{var} = {rc};")
-            V.kernel.compute.writeline(f"}} else {var} = {other_str};")
+                V.kernel.compute.writeline(
+                    f"{var} = static_cast<decltype({var})>({rc});"
+                )
+            V.kernel.compute.writeline(
+                f"}} else {var} = static_cast<decltype({var})>({other_str});"
+            )
         return var
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def where(a: OpVarT, b: OpVarT, c: OpVarT) -> str:
-        return f"{a} ? {b} : {value_to_metal(c)}"
+        return f"{a} ? {b} : static_cast<decltype({b})>({value_to_metal(c)})"
 
     @staticmethod
     def remainder(a: OpVarT, b: OpVarT) -> str:
         return f"c10::metal::remainder({a}, {b})"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def maximum(a: CSEVariable, b: CSEVariable) -> str:
         typecast_a = f"static_cast<decltype({a}+{b})>({a})"
         typecast_b = f"static_cast<decltype({a}+{b})>({b})"
         return f"c10::metal::max({typecast_a}, {typecast_b})"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def minimum(a: CSEVariable, b: CSEVariable) -> str:
         typecast_a = f"static_cast<decltype({a}+{b})>({a})"
         typecast_b = f"static_cast<decltype({a}+{b})>({b})"
         return f"c10::metal::min({typecast_a}, {typecast_b})"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def logical_or(a: CSEVariable, b: CSEVariable) -> str:
         return f"{a} || {b}"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def logical_and(a: CSEVariable, b: CSEVariable) -> str:
         return f"{a} && {b}"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def isnan(x: CSEVariable) -> str:
         return f"metal::isnan({x})"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def isinf(x: CSEVariable) -> str:
         return f"metal::isinf({x})"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def log(x: CSEVariable) -> str:
         return f"metal::precise::log({x})"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def exp(x: CSEVariable) -> str:
         return f"metal::precise::exp({x})"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def abs(x: CSEVariable) -> str:
         return f"metal::abs({x})"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def signbit(x: CSEVariable) -> str:
         return f"metal::signbit({x})"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def sin(x: CSEVariable) -> str:
         return f"metal::precise::sin({x})"
 
@@ -305,30 +325,37 @@ def sinc(x: CSEVariable) -> str:
         return f"c10::metal::sinc({x})"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def cos(x: CSEVariable) -> str:
         return f"metal::precise::cos({x})"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def tan(x: CSEVariable) -> str:
         return f"metal::precise::tan({x})"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def asin(x: CSEVariable) -> str:
         return f"metal::precise::asin({x})"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def acos(x: CSEVariable) -> str:
         return f"metal::precise::acos({x})"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def atan(x: CSEVariable) -> str:
         return f"metal::precise::atan({x})"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def atan2(x: CSEVariable, y: CSEVariable) -> str:
         return f"::metal::precise::atan2({x}, {y})"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def sqrt(x: CSEVariable) -> str:
         return f"metal::precise::sqrt({x})"
 
@@ -339,14 +366,17 @@ def neg(x: CSEVariable) -> str:
         return f"static_cast<decltype({x})>(-{x})"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def rsqrt(x: CSEVariable) -> str:
         return f"metal::precise::rsqrt({x})"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def tanh(x: CSEVariable) -> str:
         return f"metal::precise::tanh({x})"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def atanh(x: CSEVariable) -> str:
         return f"metal::precise::atanh({x})"
 
@@ -356,24 +386,29 @@ def floordiv(a: CSEVariable, b: CSEVariable) -> str:
         return f"c10::metal::floor_divide({a}, {b})"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def floor(x: CSEVariable) -> str:
         return f"metal::floor({x})"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def sign(x: CSEVariable) -> str:
         return f"metal::sign({x})"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def fmod(a: CSEVariable, b: CSEVariable) -> str:
         typecast_a = f"static_cast<decltype({a}+{b})>({a})"
         typecast_b = f"static_cast<decltype({a}+{b})>({b})"
         return f"metal::fmod({typecast_a}, {typecast_b})"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def trunc(x: CSEVariable) -> str:
         return f"metal::trunc({x})"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def truncdiv(a: CSEVariable, b: CSEVariable) -> str:
         quot = f"{a} / {b}"
         if (a.dtype is not None and a.dtype.is_floating_point) or (
@@ -383,6 +418,7 @@ def truncdiv(a: CSEVariable, b: CSEVariable) -> str:
         return quot
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def ceil(x: CSEVariable) -> str:
         return f"metal::ceil({x})"
 
@@ -404,6 +440,7 @@ def randint64(
         return f"c10::metal::randint64({seed}, {offset}, {low}, {high})"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def round(x: CSEVariable) -> str:
         return f"metal::rint({x})"
 
@@ -564,9 +601,9 @@ def store_reduction(self, name: str, index: sympy.Expr, value: CSEVariable) -> N
 
     def _new_idxvar(
         self,
-        dtype: Union[str | torch.dtype],
-        elem_count: Optional[int] = None,
-        default_value: Optional[Any] = None,
+        dtype: str | torch.dtype,
+        elem_count: int | None = None,
+        default_value: Any | None = None,
         is_threadgroup: bool = True,
         bounds: ValueRanges[Any] = ValueRanges.unknown(),
     ) -> CSEVariable:
@@ -589,8 +626,8 @@ def reduction(
         dtype: torch.dtype,
         src_dtype: torch.dtype,
         reduction_type: ReductionType,
-        value: Union[CSEVariable, tuple[CSEVariable, ...]],
-    ) -> Union[CSEVariable, tuple[CSEVariable, ...]]:
+        value: CSEVariable | tuple[CSEVariable, ...],
+    ) -> CSEVariable | tuple[CSEVariable, ...]:
         "Caching wrapper around _reduction_nocache"
         cache_key = (src_dtype, reduction_type, value)
         # Return cached reduction
@@ -605,8 +642,8 @@ def _reduction_nocache(
         dtype: torch.dtype,
         src_dtype: torch.dtype,
         reduction_type: ReductionType,
-        value: Union[CSEVariable, tuple[CSEVariable, ...]],
-    ) -> Union[CSEVariable, tuple[CSEVariable, ...]]:
+        value: CSEVariable | tuple[CSEVariable, ...],
+    ) -> CSEVariable | tuple[CSEVariable, ...]:
         """Codegen a reduction operation.
         Only sum and prod operations are somewhat reasonable optimized"""
         assert self.inside_reduction
@@ -892,10 +929,11 @@ def codegen_body(self) -> None:
         self.compute.clear()
         self.stores.clear()
 
-    def codegen_kernel(self, name: Optional[str] = None) -> str:
+    def codegen_kernel(self, name: str = "generated_kernel") -> str:
         """Called at the end to generate a final kernel string"""
         self.codegen_body()
         code = IndentedBuffer()
+        fn_name = name
 
         if V.graph.cpp_wrapper:
             code.writeline('(R"MTL(')
@@ -932,7 +970,7 @@ def codegen_kernel(self, name: Optional[str] = None) -> str:
                 code.writeline(
                     f"[[max_total_threads_per_threadgroup({threadgroup_size})]]"
                 )
-            code.writeline("kernel void generated_kernel(")
+            code.writeline(f"kernel void {fn_name}(")
             with code.indent():
                 for outer, inner in self.args.output_buffers.items():
                     if outer in self.removed_buffers:
@@ -1090,6 +1128,20 @@ def format_threads(threads: list[str], kwarg: str) -> str:
             arg_types=arg_types,
         )
 
+    def device_assert_async(self, cond: CSEVariable, msg: str) -> None:
+        if V.graph.cpp_wrapper:
+            self.cse.generate(self.compute, f"if (!{cond}) return", assignment=False)
+        else:
+            self.headers.add("error")
+            self.compute.writelines(
+                [
+                    f"if (!{cond}) {{",
+                    f"    TORCH_REPORT_ERROR(error_buf, {repr(msg)});",
+                    "    return;",
+                    "}",
+                ]
+            )
+
     def check_bounds(
         self, expr: sympy.Expr, size: sympy.Expr, lower: bool, upper: bool
     ) -> None:
@@ -1129,36 +1181,55 @@ def check_bounds(
 
 class MetalScheduling(SIMDScheduling):
     kernel_type = MetalKernel  # type: ignore[assignment]
+    _kernel_fn_counter: int = 0
 
-    def __init__(self, scheduler: Optional[Scheduler]) -> None:
+    def __init__(self, scheduler: Scheduler | None) -> None:
         super().__init__(scheduler)
-        wrapper = V.graph.wrapper_code
-        if wrapper is not None:
-            if not V.graph.cpp_wrapper:
-                wrapper.header.splice(
-                    "from torch._inductor.runtime.runtime_utils import compile_mps_shader"
-                )
 
     def define_kernel(
         self, src_code: str, node_schedule: list[SchedulerNode], kernel: MetalKernel
     ) -> str:
         wrapper = V.graph.wrapper_code
         if src_code in wrapper.src_to_kernel:
-            kernel_name = wrapper.src_to_kernel[src_code]
-        else:
-            # TODO: Merge multiple kernels into a single library
-            # Either using MultiKernel concept or overriding SIMDScheduling.codegen_node_scheduling
-            mps_lib_name = f"mps_lib_{wrapper.next_kernel_suffix()}"
+            return wrapper.src_to_kernel[src_code]
 
-            kernel_name = f"{mps_lib_name}"
+        if V.graph.cpp_wrapper:
+            # C++ path: one library per kernel (each has a single "generated_kernel" function)
+            mps_lib_name = f"mps_lib_{wrapper.next_kernel_suffix()}"
+            kernel_name = mps_lib_name
             wrapper.src_to_kernel[src_code] = kernel_name
-
-            if V.graph.cpp_wrapper:
-                # For shimified version, generate source constant instead of direct instantiation
-                src_code = f"const char* {mps_lib_name}_source = " + src_code
-
+            src_code = f"const char* {mps_lib_name}_source = " + src_code
             origins, detailed_origins = get_kernel_metadata(node_schedule, wrapper)
-            metadata_comment = f"{origins}\n{detailed_origins}"
-            wrapper.define_kernel(mps_lib_name, src_code, metadata_comment, gpu=False)
+            wrapper.define_kernel(
+                mps_lib_name, src_code, f"{origins}\n{detailed_origins}", gpu=False
+            )
+            return kernel_name
+
+        # Python path: register kernel with async_compile; wait() will compile all
+        # accumulated Metal kernels into a single library and replace each placeholder.
+        fn_name = f"generated_kernel_{self._kernel_fn_counter}"
+        self._kernel_fn_counter += 1
+        wrapper.src_to_kernel[src_code] = fn_name
+
+        # Extract Metal source from compile_mps_shader('''...''') call
+        metal_src_start = "compile_mps_shader('''"
+        start = src_code.index(metal_src_start) + len(metal_src_start)
+        end = src_code.rindex("''')")
+        metal_src = src_code[start:end]
+
+        # Strip #include lines and rename the kernel function
+        body_lines = []
+        for line in metal_src.split("\n"):
+            if line.strip().startswith("#include"):
+                continue
+            body_lines.append(
+                line.replace("kernel void generated_kernel(", f"kernel void {fn_name}(")
+            )
+
+        headers_repr = repr(sorted(kernel.headers))
+        wrapper.header.writeline(f"{fn_name} = async_compile.metal({fn_name!r}, '''")
+        for line in body_lines:
+            wrapper.header.writeline(line)
+        wrapper.header.writeline(f"''', {headers_repr})")
 
-        return kernel_name
+        return fn_name
diff --git a/torch/_inductor/codegen/multi_kernel.py b/torch/_inductor/codegen/multi_kernel.py
index 094164a1f08ca..f4ba4ad542d55 100644
--- a/torch/_inductor/codegen/multi_kernel.py
+++ b/torch/_inductor/codegen/multi_kernel.py
@@ -4,7 +4,7 @@
 import math
 import os
 import pathlib
-from typing import Any, Optional, Union
+from typing import Any
 
 from torch._inductor.ir import MultiTemplateBuffer
 from torch._inductor.metrics import get_metric_table, is_metric_table_enabled
@@ -36,9 +36,7 @@ def __init__(self):
     def define_kernel(
         self,
         kernels: list[Any],
-        kernel_shape_keys: Optional[
-            list[Union[None, tuple[tuple[int, ...], ...]]]
-        ] = None,
+        kernel_shape_keys: list[None | tuple[tuple[int, ...], ...]] | None = None,
     ) -> str:
         """
         Previously we name the multi kernel as "multi_kernel_{kernel_names[0]}".
diff --git a/torch/_inductor/codegen/nv_universal_gemm/kernel_cache.py b/torch/_inductor/codegen/nv_universal_gemm/kernel_cache.py
index f4fe8e38f6225..f0a8cf924b6bd 100644
--- a/torch/_inductor/codegen/nv_universal_gemm/kernel_cache.py
+++ b/torch/_inductor/codegen/nv_universal_gemm/kernel_cache.py
@@ -12,13 +12,13 @@
 
 import logging
 from collections.abc import Callable
-from typing import Any, Optional
+from typing import Any
 
 
 log = logging.getLogger(__name__)
 
 # Global cache: kernel_name -> kernel object
-_kernel_by_name_cache: Optional[dict[str, Any]] = None
+_kernel_by_name_cache: dict[str, Any] | None = None
 
 
 def _build_kernel_cache() -> dict[str, Any]:
@@ -26,6 +26,14 @@ def _build_kernel_cache() -> dict[str, Any]:
     import cutlass_api
 
     log.debug("Building NVGEMM kernel cache (this may take a few seconds)...")
+
+    try:
+        from torch._inductor.kernel.vendored_templates.cutedsl import (  # noqa: F401
+            wrappers,
+        )
+    except ImportError:
+        log.debug("Vendored kernel wrappers not available")
+
     all_kernels = cutlass_api.get_kernels()
     cache = {k.metadata.kernel_name: k for k in all_kernels}
     log.debug("NVGEMM kernel cache built: %d kernels", len(cache))
@@ -35,7 +43,7 @@ def _build_kernel_cache() -> dict[str, Any]:
 def get_compatible_kernels(
     args: Any,
     cc: int,
-    metadata_filter: Optional[Callable[[Any], bool]] = None,
+    metadata_filter: Callable[[Any], bool] | None = None,
 ) -> list[Any]:
     """Get kernels compatible with the given arguments from the cache."""
     global _kernel_by_name_cache
diff --git a/torch/_inductor/codegen/nv_universal_gemm/nv_universal_gemm.py b/torch/_inductor/codegen/nv_universal_gemm/nv_universal_gemm.py
index f9d1bf6da24c5..bf934fab6bc95 100644
--- a/torch/_inductor/codegen/nv_universal_gemm/nv_universal_gemm.py
+++ b/torch/_inductor/codegen/nv_universal_gemm/nv_universal_gemm.py
@@ -8,7 +8,7 @@
 
 import itertools
 from enum import auto, Enum
-from typing import Any, Optional, Union
+from typing import Any
 
 import torch
 from torch._inductor import config
@@ -65,22 +65,22 @@ class NVUniversalGemmBenchmarkRequest(GPUDeviceBenchmarkMixin, BenchmarkRequest)
     def __init__(
         self,
         kernel_name: str,
-        input_tensor_meta: Union[TensorMeta, list[TensorMeta]],
-        output_tensor_meta: Union[TensorMeta, list[TensorMeta]],
+        input_tensor_meta: TensorMeta | list[TensorMeta],
+        output_tensor_meta: TensorMeta | list[TensorMeta],
         kernel,  # cutlass_api.Kernel object
         accumulator_type: torch.dtype,
         variant: GemmVariant,
         workspace_size: int = 0,
-        scale_type_a: Optional[Any] = None,
-        scale_type_b: Optional[Any] = None,
-        swizzle_type_a: Optional[Any] = None,
-        swizzle_type_b: Optional[Any] = None,
+        scale_type_a: Any | None = None,
+        scale_type_b: Any | None = None,
+        swizzle_type_a: Any | None = None,
+        swizzle_type_b: Any | None = None,
     ) -> None:
         super().__init__(kernel_name, input_tensor_meta, output_tensor_meta, ())
         self.kernel = kernel
         self.accumulator_type = accumulator_type
         self._compiled_artifact = None
-        self._workspace: Optional[torch.Tensor] = None
+        self._workspace: torch.Tensor | None = None
         self.workspace_size = workspace_size
         self.variant = variant
         self.scale_type_a = scale_type_a
@@ -91,7 +91,7 @@ def __init__(
     def benchmark(
         self,
         *input_tensors: torch.Tensor,
-        out: Optional[torch.Tensor] = None,
+        out: torch.Tensor | None = None,
     ) -> float:
         """Benchmark the NVIDIA Universal GEMM kernel.
 
@@ -205,10 +205,10 @@ def __init__(
         accumulator_type: torch.dtype,
         variant: GemmVariant,
         workspace_size: int = 0,
-        scale_type_a: Optional[Any] = None,
-        scale_type_b: Optional[Any] = None,
-        swizzle_type_a: Optional[Any] = None,
-        swizzle_type_b: Optional[Any] = None,
+        scale_type_a: Any | None = None,
+        scale_type_b: Any | None = None,
+        swizzle_type_a: Any | None = None,
+        swizzle_type_b: Any | None = None,
     ) -> None:
         super().__init__(
             name=name,
@@ -284,7 +284,7 @@ def info_dict(self) -> dict[str, Any]:
         }
 
 
-def _create_dummy_tensor_from_layout(layout: Layout) -> Optional[torch.Tensor]:
+def _create_dummy_tensor_from_layout(layout: Layout) -> torch.Tensor | None:
     """
     Create a FakeTensor from a Layout for kernel filtering.
 
@@ -316,11 +316,11 @@ def _add_nv_gemm_choices_impl(
     input_nodes: list[Buffer],
     variant: GemmVariant,
     accumulator_type: torch.dtype,
-    mm_inputs: Optional[MMKernelInputs] = None,
-    scale_type_a: Optional[Any] = None,
-    scale_type_b: Optional[Any] = None,
-    swizzle_type_a: Optional[Any] = None,
-    swizzle_type_b: Optional[Any] = None,
+    mm_inputs: MMKernelInputs | None = None,
+    scale_type_a: Any | None = None,
+    scale_type_b: Any | None = None,
+    swizzle_type_a: Any | None = None,
+    swizzle_type_b: Any | None = None,
 ) -> None:
     """
     Unified implementation for adding NVIDIA Universal GEMM choices.
@@ -456,7 +456,7 @@ def add_nv_universal_gemm_choices(
     choices: list[ChoiceCaller],
     layout: Layout,
     inputs: MMKernelInputs,
-    accumulator_type: Optional[torch.dtype] = None,
+    accumulator_type: torch.dtype | None = None,
 ) -> None:
     """
     Add NVIDIA Universal GEMM kernels to the autotune choices.
@@ -481,7 +481,7 @@ def add_nv_universal_grouped_gemm_choices(
     choices: list[ChoiceCaller],
     layout: Layout,
     input_nodes: list[Buffer],
-    accumulator_type: Optional[torch.dtype] = None,
+    accumulator_type: torch.dtype | None = None,
 ) -> None:
     """
     Add NVIDIA Universal Grouped GEMM kernels to the autotune choices.
@@ -513,8 +513,8 @@ def add_nv_universal_scaled_gemm_choices(
     choices: list[ChoiceCaller],
     layout: Layout,
     input_nodes: list[Buffer],
-    accumulator_type: Optional[torch.dtype] = None,
-    kernel_inputs: Optional[MMKernelInputs] = None,
+    accumulator_type: torch.dtype | None = None,
+    kernel_inputs: MMKernelInputs | None = None,
 ) -> None:
     """
     Add NVIDIA Universal Scaled GEMM (FP8) kernels to the autotune choices.
diff --git a/torch/_inductor/codegen/nv_universal_gemm/nv_universal_gemm_kernel.py b/torch/_inductor/codegen/nv_universal_gemm/nv_universal_gemm_kernel.py
index ed2a471112187..aae76dfd3246c 100644
--- a/torch/_inductor/codegen/nv_universal_gemm/nv_universal_gemm_kernel.py
+++ b/torch/_inductor/codegen/nv_universal_gemm/nv_universal_gemm_kernel.py
@@ -8,7 +8,7 @@
 from __future__ import annotations
 
 import logging
-from typing import Any, Optional, TYPE_CHECKING, Union
+from typing import Any, TYPE_CHECKING
 
 from torch._inductor.codegen.common import (
     IndentedBuffer,
@@ -41,7 +41,7 @@
 class NVUniversalGemmKernelWrapper:
     """Wrapper to provide .run() interface for NVIDIA Universal GEMM kernels."""
 
-    def __init__(self, kernel_fn, kernel_path: Optional[str] = None):
+    def __init__(self, kernel_fn, kernel_path: str | None = None):
         self.kernel_fn = kernel_fn
         self.kernel_path = kernel_path
 
@@ -68,10 +68,10 @@ def __init__(
         accumulator_type: Any,
         variant: GemmVariant,
         workspace_size: int = 0,
-        scale_type_a: Optional[Any] = None,
-        scale_type_b: Optional[Any] = None,
-        swizzle_type_a: Optional[Any] = None,
-        swizzle_type_b: Optional[Any] = None,
+        scale_type_a: Any | None = None,
+        scale_type_b: Any | None = None,
+        swizzle_type_a: Any | None = None,
+        swizzle_type_b: Any | None = None,
     ) -> None:
         super().__init__()
         self.kernel_name = kernel_name
@@ -225,7 +225,7 @@ def {self.kernel_name}_main({params_str}):
 
         return code.getvalue()
 
-    def _get_reinterpret_view(self, node) -> Optional[ReinterpretView]:
+    def _get_reinterpret_view(self, node) -> ReinterpretView | None:
         """Extract or convert to ReinterpretView from a node, handling all views."""
         while isinstance(node, MutableBox):
             node = node.data
@@ -243,7 +243,7 @@ def call_kernel(self, name: str, node=None):
 
         call_args: list[str] = []
         arg_types: list[Any] = []
-        raw_args: list[Union[Buffer, ReinterpretView, None]] = []
+        raw_args: list[Buffer | ReinterpretView | None] = []
 
         for _, input_node in self._template_input_args:
             reinterpret_view = self._get_reinterpret_view(input_node)
@@ -263,7 +263,7 @@ def call_kernel(self, name: str, node=None):
         raw_args.append(None)  # Output buffer is findable by name
 
         # Allocate workspace if needed
-        ws: Optional[WorkspaceArg] = None
+        ws: WorkspaceArg | None = None
         if self.workspace_size > 0:
             ws = WorkspaceArg(
                 count=self.workspace_size,
diff --git a/torch/_inductor/codegen/nv_universal_gemm/nv_universal_gemm_utils.py b/torch/_inductor/codegen/nv_universal_gemm/nv_universal_gemm_utils.py
index 65368f779e1a3..077e03dfb3a7c 100644
--- a/torch/_inductor/codegen/nv_universal_gemm/nv_universal_gemm_utils.py
+++ b/torch/_inductor/codegen/nv_universal_gemm/nv_universal_gemm_utils.py
@@ -3,14 +3,14 @@
 Utility functions for NVIDIA Universal GEMM.
 """
 
-from typing import Any, Optional
+from typing import Any
 
 from torch.nn.functional import ScalingType, SwizzleType
 
 
 def to_cutlass_scale_mode(
     scale_type: Any, swizzle_type: Any
-) -> tuple[Optional[Any], Optional[Any]]:
+) -> tuple[Any | None, Any | None]:
     """
     Map PyTorch ScalingType/SwizzleType to cutlass_api ScaleMode/ScaleSwizzleMode.
 
diff --git a/torch/_inductor/codegen/pallas.py b/torch/_inductor/codegen/pallas.py
index 9df33d439f0e0..2bbf781e334d0 100644
--- a/torch/_inductor/codegen/pallas.py
+++ b/torch/_inductor/codegen/pallas.py
@@ -2,10 +2,10 @@
 
 import dataclasses
 import hashlib
-import itertools
 import math
+import re
 import typing_extensions
-from typing import Any, Optional, TYPE_CHECKING, Union
+from typing import Any, cast, TYPE_CHECKING
 
 import sympy  # noqa: TC002
 
@@ -14,7 +14,6 @@
 from torch.utils._sympy.functions import ModularIndexing
 
 from .. import config
-from ..ir import ComputedBuffer
 from ..runtime.runtime_utils import torch_dtype_to_jax
 from ..utils import get_fused_kernel_name, get_kernel_metadata
 from ..virtualized import V
@@ -26,7 +25,7 @@
     OpOverrides,
     PythonPrinter,
 )
-from .simd import SIMDKernel, SIMDScheduling
+from .simd import IterationRangesEntry, SIMDKernel, SIMDScheduling
 
 
 class PallasPrinter(PythonPrinter):
@@ -73,17 +72,6 @@ def _print_Max(self, expr: sympy.Expr) -> str:
 # Main function suffix used in generated Pallas code
 MAIN_SUFFIX = "main"
 
-# Mosaic GPU warpgroup size: 4 warps × 32 threads = 128 threads per warpgroup.
-# This is a hardware constant for Hopper and Blackwell GPUs.
-# See: jax/_src/pallas/mosaic_gpu/lowering.py
-WARPGROUP_SIZE = 128
-
-
-def _align_to_warpgroup(size: int) -> int:
-    """Align size to WARPGROUP_SIZE (128) for Mosaic GPU compatibility."""
-    return ((size + WARPGROUP_SIZE - 1) // WARPGROUP_SIZE) * WARPGROUP_SIZE
-
-
 # Logger for Pallas kernel code
 kernel_code_log = torch._logging.getArtifactLogger(__name__, "kernel_code")
 
@@ -91,9 +79,7 @@ def _align_to_warpgroup(size: int) -> int:
 class PallasKernelWrapper:
     """Wrapper to provide .run() interface for Pallas kernels"""
 
-    def __init__(
-        self, kernel_fn: Callable[..., Any], kernel_path: Optional[str] = None
-    ):
+    def __init__(self, kernel_fn: Callable[..., Any], kernel_path: str | None = None):
         self.kernel_fn = kernel_fn
         self.kernel_path = kernel_path
         kernel_code_log.info("Pallas kernel path: %s", kernel_path)
@@ -126,42 +112,52 @@ class PallasKernelOverrides(OpOverrides):
     """
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def sin(x: str) -> str:
         return f"jnp.sin({x})"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def cos(x: str) -> str:
         return f"jnp.cos({x})"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def tan(x: str) -> str:
         return f"jnp.tan({x})"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def sinh(x: str) -> str:
         return f"jnp.sinh({x})"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def cosh(x: str) -> str:
         return f"jnp.cosh({x})"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def tanh(x: str) -> str:
         return f"jnp.tanh({x})"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def asin(x: str) -> str:
         return f"jnp.arcsin({x})"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def acos(x: str) -> str:
         return f"jnp.arccos({x})"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def atan(x: str) -> str:
         return f"jnp.arctan({x})"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def exp(x: str) -> str:
         return f"jnp.exp({x})"
 
@@ -174,6 +170,7 @@ def expm1(x: str) -> str:
         return f"jnp.expm1({x})"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def log(x: str) -> str:
         return f"jnp.log({x})"
 
@@ -190,14 +187,17 @@ def log1p(x: str) -> str:
         return f"jnp.log1p({x})"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def sqrt(x: str) -> str:
         return f"jnp.sqrt({x})"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def rsqrt(x: str) -> str:
         return f"jax.lax.rsqrt({x})"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def abs(x: str) -> str:
         return f"jnp.abs({x})"
 
@@ -206,18 +206,22 @@ def neg(x: str) -> str:
         return f"(-{x})"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def floor(x: str) -> str:
         return f"jnp.floor({x})"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def ceil(x: str) -> str:
         return f"jnp.ceil({x})"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def trunc(x: str) -> str:
         return f"jnp.trunc({x})"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def round(x: str) -> str:
         return f"jnp.round({x})"
 
@@ -234,14 +238,17 @@ def pow(a: str, b: str) -> str:
         return f"jnp.power({a}, {b})"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def maximum(a: str, b: str) -> str:
         return f"jnp.maximum({a}, {b})"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def minimum(a: str, b: str) -> str:
         return f"jnp.minimum({a}, {b})"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def where(cond: str, a: str, b: str) -> str:
         return f"jnp.where({cond}, {a}, {b})"
 
@@ -269,9 +276,12 @@ def masked(mask: str, body: Callable[[], str], other: float) -> str:
     def to_dtype(
         x: str,
         dtype: torch.dtype,
-        src_dtype: Optional[torch.dtype] = None,
+        src_dtype: torch.dtype | None = None,
         use_compute_types: bool = True,
     ) -> str:
+        # TPU doesn't support 64-bit types
+        if dtype == torch.int64 and V.graph.get_current_device_or_throw().type == "tpu":
+            dtype = torch.int32
         jax_dtype = torch_dtype_to_jax(dtype)
         # Wrap in jnp.asarray to handle scalars from integer indexing
         return f"jnp.asarray({x}).astype({jax_dtype})"
@@ -363,10 +373,12 @@ def gt(a: str, b: str) -> str:
         return f"({a} > {b})"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def isnan(x: str) -> str:
         return f"jnp.isnan({x})"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def isinf(x: str) -> str:
         return f"jnp.isinf({x})"
 
@@ -380,10 +392,12 @@ def ge(a: str, b: str) -> str:
 
     # Logical operations
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def logical_and(a: str, b: str) -> str:
         return f"jnp.logical_and({a}, {b})"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def logical_or(a: str, b: str) -> str:
         return f"jnp.logical_or({a}, {b})"
 
@@ -392,19 +406,23 @@ def logical_not(x: str) -> str:
         return f"jnp.logical_not({x})"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def logical_xor(a: str, b: str) -> str:
         return f"jnp.logical_xor({a}, {b})"
 
     # Math operations
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def atan2(a: str, b: str) -> str:
         return f"jnp.arctan2({a}, {b})"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def hypot(a: str, b: str) -> str:
         return f"jnp.hypot({a}, {b})"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def fmod(a: str, b: str) -> str:
         return f"jnp.fmod({a}, {b})"
 
@@ -413,6 +431,7 @@ def remainder(a: str, b: str) -> str:
         return f"jnp.remainder({a}, {b})"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def truncdiv(a: str, b: str) -> str:
         # Truncated division (rounds toward zero)
         # For integers: sign(a)*sign(b) * (abs(a) // abs(b))
@@ -430,15 +449,19 @@ def clamp(x: str, min_val: str, max_val: str) -> str:
 
     # Sign operations
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def sign(x: str) -> str:
-        return f"jnp.sign({x})"
+        # PyTorch returns 0 for NaN, JAX returns NaN
+        return f"jnp.where(jnp.isnan({x}), 0.0, jnp.sign({x}))"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def signbit(x: str) -> str:
         return f"jnp.signbit({x})"
 
     # Special math functions
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def erf(x: str) -> str:
         return f"jax.scipy.special.erf({x})"
 
@@ -447,10 +470,12 @@ def erfc(x: str) -> str:
         return f"jax.scipy.special.erfc({x})"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def erfinv(x: str) -> str:
         return f"jax.scipy.special.erfinv({x})"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def lgamma(x: str) -> str:
         return f"jax.scipy.special.gammaln({x})"
 
@@ -556,6 +581,7 @@ def xlog1py(x: str, y: str) -> str:
         return f"jax.scipy.special.xlog1py({x}, {y})"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def chebyshev_polynomial_t(x: str, n: str) -> str:
         # Chebyshev polynomial of the first kind T_n(x)
         # For |x| <= 1: T_n(x) = cos(n * arccos(x))
@@ -570,6 +596,7 @@ def chebyshev_polynomial_t(x: str, n: str) -> str:
         )
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def chebyshev_polynomial_u(x: str, n: str) -> str:
         # Chebyshev polynomial of the second kind U_n(x)
         # For |x| < 1: U_n(x) = sin((n+1) * arccos(x)) / sqrt(1 - x^2)
@@ -591,6 +618,7 @@ def chebyshev_polynomial_u(x: str, n: str) -> str:
         )
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def chebyshev_polynomial_v(x: str, n: str) -> str:
         # Chebyshev polynomial of the third kind V_n(x)
         # V_n(x) = (T_n(x) - T_{n+1}(x)) / (1 - x) for x != 1
@@ -607,6 +635,7 @@ def chebyshev_polynomial_v(x: str, n: str) -> str:
         )
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def chebyshev_polynomial_w(x: str, n: str) -> str:
         # Chebyshev polynomial of the fourth kind W_n(x)
         # W_n(x) = (T_n(x) + T_{n+1}(x)) / (1 + x) for x != -1
@@ -623,22 +652,27 @@ def chebyshev_polynomial_w(x: str, n: str) -> str:
         )
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def shifted_chebyshev_polynomial_t(x: str, n: str) -> str:
         return PallasKernelOverrides.chebyshev_polynomial_t(f"(2 * {x} - 1)", n)
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def shifted_chebyshev_polynomial_u(x: str, n: str) -> str:
         return PallasKernelOverrides.chebyshev_polynomial_u(f"(2 * {x} - 1)", n)
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def shifted_chebyshev_polynomial_v(x: str, n: str) -> str:
         return PallasKernelOverrides.chebyshev_polynomial_v(f"(2 * {x} - 1)", n)
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def shifted_chebyshev_polynomial_w(x: str, n: str) -> str:
         return PallasKernelOverrides.chebyshev_polynomial_w(f"(2 * {x} - 1)", n)
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def hermite_polynomial_h(x: str, n: str) -> str:
         # Physicist's Hermite polynomial H_n(x)
         # H_n(x) = 2^n * x^n - n*(n-1)/2 * 2^(n-2) * x^(n-2) + ...
@@ -657,6 +691,7 @@ def hermite_polynomial_h(x: str, n: str) -> str:
         )
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def hermite_polynomial_he(x: str, n: str) -> str:
         # Probabilist's Hermite polynomial He_n(x)
         # He_0 = 1, He_1 = x, He_2 = x^2 - 1, He_3 = x^3 - 3x
@@ -671,6 +706,7 @@ def hermite_polynomial_he(x: str, n: str) -> str:
         )
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def laguerre_polynomial_l(x: str, n: str) -> str:
         # Laguerre polynomial L_n(x)
         # L_0 = 1, L_1 = 1 - x, L_2 = (x^2 - 4x + 2)/2, L_3 = (-x^3 + 9x^2 - 18x + 6)/6
@@ -685,6 +721,7 @@ def laguerre_polynomial_l(x: str, n: str) -> str:
         )
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def legendre_polynomial_p(x: str, n: str) -> str:
         # Legendre polynomial P_n(x)
         # P_0 = 1, P_1 = x, P_2 = (3x^2 - 1)/2, P_3 = (5x^3 - 3x)/2
@@ -718,18 +755,22 @@ def fma(a: str, b: str, c: str) -> str:
         return f"(({a}) * ({b}) + ({c}))"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def copysign(a: str, b: str) -> str:
         return f"jnp.copysign({a}, {b})"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def nextafter(a: str, b: str) -> str:
         return f"jnp.nextafter({a}, {b})"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def ldexp(a: str, b: str) -> str:
         return f"jnp.ldexp({a}, {b})"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def frexp(x: str) -> str:
         return f"jnp.frexp({x})"
 
@@ -812,6 +853,37 @@ def randint64(seed: str, offset: str, low: str, high: str) -> str:
         )
 
 
+@dataclasses.dataclass
+class _IndirectAccessInfo:
+    """Describes a detected indirect (data-dependent) buffer access."""
+
+    table_param: str
+    table_buf_name: str
+    table_shape: tuple
+    indirect_dim: int
+    indirect_var: str
+    indices_param: str
+
+
+@dataclasses.dataclass
+class _BufferIndexing:
+    """Encapsulates index string and flattening requirements for buffer access."""
+
+    index_str: str
+    needs_flatten: bool
+
+
+@dataclasses.dataclass
+class _BroadcastedIterVar:
+    """Encapsulates information needed to codegen a broadcasted iteration var"""
+
+    # index of this var in `self.range_tree_nodes.items()``
+    idx: int
+    var_sym: sympy.Symbol
+    entry: IterationRangesEntry
+    length_val: int | None
+
+
 @dataclasses.dataclass
 class _CodegenContext:
     """Bundles local state shared across codegen_kernel helper methods."""
@@ -833,22 +905,14 @@ class _CodegenContext:
     full_kernel_params: list[str]
     non_alias_out_set: OrderedSet[str]
     copy_output_indices: list[int]
+    alias_pairs: list[tuple[int, int]]
 
 
 class PallasKernel(SIMDKernel):
-    """
-    Pallas kernel for elementwise operations with support for strided/scatter access.
-
-    Strategy:
-    - Convert index expressions to JAX-compatible array slicing
-    - Load/store using indexed access: "in_ptrX[slice]" or full-array "in_ptrX[...]"
-    - Compute expression with Python operators (compatible with jax.numpy broadcasting)
-    - Generate Python code that defines a Pallas kernel and a host entrypoint.
-    - Use async_compile.pallas path to compile and load Python code.
-
-    For GPU (Mosaic backend):
-    - Use TMA (Tensor Memory Accelerator) for automatic OOB masking
-    - Falls back to legacy padding approach for reductions, broadcasting, non-contiguous tensors
+    """Pallas kernel codegen for TPU and GPU (Mosaic backend).
+
+    Generates Python code that defines a Pallas kernel and a host entrypoint,
+    compiled and loaded via async_compile.pallas.
     """
 
     overrides = PallasKernelOverrides  # type: ignore[assignment]
@@ -859,23 +923,48 @@ def __init__(self, *args, **kwargs):
         # Determine device type once at initialization
         device = V.graph.get_current_device_or_throw()
         self.is_gpu = device.type == "cuda"
+        self.is_tpu = device.type == "tpu"
         # Use TMA (Tensor Memory Accelerator) for GPU to handle non-aligned tensor sizes
         # TMA automatically masks OOB accesses, eliminating the need for explicit
         # padding to multiples of 128. Uses lax.fori_loop with direct TMA primitives.
         self.use_emit_pipeline = self.is_gpu  # Enable TMA approach for GPU
-        # Legacy: warpgroup padding (enabled when TMA approach is disabled)
-        self.use_warpgroup_padding = self.is_gpu and not self.use_emit_pipeline
         # Track which output param each store uses: list of (out_ptr_name, store_line)
         self.store_with_output: list[tuple[str, str]] = []
-        # Track load index expressions for argmax/argmin axis detection
+        # Track load index expressions for reduction axis detection
         self.load_index_exprs: dict[str, sympy.Expr] = {}
         # Track outputs that need to be readable (for scatter operations)
         self.outputs_need_read: OrderedSet[str] = OrderedSet()
-        # Track if any load in this kernel used transpose
-        # Used to avoid double transpose (load + store)
-        self.has_transposed_load = False
+        # Map input buffer names to their detected permutation tuples.
+        self.permuted_input_buffers: dict[str, tuple[int, ...]] = {}
+        self.collapsed_reshape_inputs: dict[str, tuple[int, ...]] = {}
+        self.collapsed_output_shape: tuple[int, ...] | None = None
+        self._cpu_max_grid_product: int | None = None
+        # Precompute output buffer names from scheduler nodes so that the
+        # load path can check output shapes before stores are processed.
+        self._output_buffer_names: list[str] = []
+        for snode in self.features.scheduler_nodes():
+            for dep in snode.read_writes.writes:
+                self._output_buffer_names.append(dep.name)
         # Track which iteration variables are actually used in the kernel
         self.used_iter_vars: OrderedSet[sympy.Symbol] = OrderedSet()
+        # Iteration vars that have been emitted in tile-relative form
+        # (safe for tiling even when they appear in the compute body)
+        self.tile_relative_iter_vars: OrderedSet[sympy.Symbol] = OrderedSet()
+        # Track if any load/store uses flatten-based indexing (buf[...].flatten()[idx])
+        self.has_flatten_indexing = False
+        # Strided input buffers: map graph buffer name -> per-dim
+        # (stride, offset, skip) triples.  Used to reshape inputs outside
+        # the kernel and generate static indexing inside
+        # (e.g. in_ref[:, :, offset] instead of in_ref[...].flatten()[idx]).
+        self.strided_input_buffers: dict[str, list[tuple[int, int, int]]] = {}
+        # Buffers that already use flatten+gather indexing; strided
+        # decomposition must not reshape these (it would break flat offsets).
+        self.flatten_indexed_buffers: OrderedSet[str] = OrderedSet()
+        # Indirect (data-dependent) access info for scalar prefetch
+        self.indirect_access: _IndirectAccessInfo | None = None
+        self._cse_to_param: dict[str, str] = {}
+        self._param_to_graph_name: dict[str, str] = {}
+        self.is_tpu = device.type == "tpu"
 
     def check_bounds(
         self, expr: sympy.Expr, size: sympy.Expr, lower: bool, upper: bool
@@ -1069,130 +1158,506 @@ def _has_indirect_vars(self, index: sympy.Expr) -> bool:
         """Check if index expression contains indirect variables."""
         return len(self._get_indirect_vars(index)) > 0
 
-    def _get_expected_output_shape(self) -> list:
-        """Get the expected output shape from iteration variables.
-
-        Iteration variables are shaped for broadcasting. For 2D outputs:
-        - First var (e.g., y0) gets shape (1, N) - innermost dimension
-        - Second var (e.g., x1) gets shape (M, 1) - outermost dimension
-        The broadcast result is (M, N).
-        """
-        # Collect variable lengths
-        var_items = list(self.range_tree_nodes.items())
-        broadcast_vars = []
-        for var_sym, entry in var_items:
-            length = self._safe_int(entry.length)
-            if length is not None:
-                broadcast_vars.append(length)
-
-        if len(broadcast_vars) <= 1:
-            return broadcast_vars
-
-        # For 2D case: variables are reshaped in reverse order
-        # First var is innermost (last dim), second var is outermost (first dim)
-        # So output shape is [second_var_length, first_var_length, ...]
-        return list(reversed(broadcast_vars))
-
-    def _is_transposed_access(self, name: str, index: sympy.Expr) -> bool:
-        """Check if buffer access needs transpose.
-
-        Transpose on load is needed when:
-        1. Non-square buffers: dimensions are swapped relative to iteration vars
-        2. Square buffers: index coefficient pattern indicates transposed access
-           (first iteration var has larger coefficient than second)
+    def _decompose_strided_access(
+        self, index: sympy.Expr, name: str
+    ) -> list[tuple[int, int, int]] | None:
+        """Decompose a flat index into per-dimension (stride, offset, skip) triples.
+
+        Given flat index like ``64*x0 + 2*x1 + 5`` and buffer shape ``(32, 64)``
+        with C-contiguous strides ``[64, 1]``:
+          - x0 coefficient 64 / buffer_stride[0]=64 -> dim 0: stride=1
+          - x1 coefficient 2 / buffer_stride[1]=1  -> dim 1: stride=2
+          - constant 5: dim 0 gets 5//64=0, dim 1 gets 5//1=5
+          - dim 1 offset 5 with stride 2: skip=5//2=2, offset=5%2=1
+
+        Returns per-dim ``[(stride, offset, skip), ...]`` where:
+          - stride: access stride on this dim (1 = contiguous)
+          - offset: static index into the stride dim (0 <= offset < stride)
+          - skip: number of stride-blocks to skip at the start of this dim
+        Returns None if decomposition fails.
         """
+        if self._has_indirect_vars(index) or index.has(ModularIndexing):
+            return None
+
+        # Don't reshape a buffer that already has flatten+gather loads;
+        # the reshape would change the flat layout and break those loads.
+        if name in self.flatten_indexed_buffers:
+            return None
+
         info = self._get_buffer_info(name)
         if info is None:
-            return False
+            return None
+        _, buf_size, _, _, _ = info
 
-        _, buf_size, _, actual_strides, _ = info
+        buf_shape_or_none = [self._safe_int(s) for s in buf_size]
+        if any(s is None or s <= 0 for s in buf_shape_or_none):
+            return None
+        buf_shape: list[int] = cast(list[int], buf_shape_or_none)
+        ndim = len(buf_shape)
+        if ndim == 0:
+            return None
 
-        # Only handle 2D buffers
-        if len(buf_size) != 2 or len(actual_strides) != 2:
-            return False
+        c_strides = self._c_contiguous_strides(buf_shape)
 
-        size0 = self._safe_int(buf_size[0])
-        size1 = self._safe_int(buf_size[1])
-        if size0 is None or size1 is None or size0 <= 1 or size1 <= 1:
-            return False
+        # Extract per-variable coefficients
+        used_vars = self._get_used_iter_vars(index)
+        if not used_vars:
+            return None
 
-        s0 = actual_strides[0]
-        s1 = actual_strides[1]
-        if s0 is None or s1 is None:
-            return False
+        # [stride, raw_offset] per dim — raw_offset may be >= stride
+        result: list[list[int]] = [[1, 0] for _ in range(ndim)]
+        # Track which variable maps to which dimension
+        var_to_dim: dict[sympy.Symbol, int] = {}
 
-        # Get iteration variable info
-        var_items = list(self.range_tree_nodes.items())
-        if len(var_items) < 2:
-            return False
+        remaining = V.graph.sizevars.simplify(index)
+        for var in used_vars:
+            var_expr = BlockPatternMatcher.get_subexpr_involving_symbol(remaining, var)
+            coeff = BlockPatternMatcher.match_affine_block_expr(var_expr, var)
+            if coeff is None:
+                return None
+            coeff_int = self._safe_int(coeff)
+            if coeff_int is None or coeff_int <= 0:
+                return None
 
-        # Skip for reduction variables
+            # Find which buffer dim this variable maps to
+            dim = None
+            for d in range(ndim):
+                if c_strides[d] == 0:
+                    continue
+                if coeff_int % c_strides[d] == 0:
+                    per_dim_stride = coeff_int // c_strides[d]
+                    if per_dim_stride >= 1:
+                        dim = d
+                        break
+            if dim is None:
+                return None
 
-        if any(entry.is_reduction for _, entry in var_items):
-            return False
+            per_dim_stride = coeff_int // c_strides[dim]
+            if per_dim_stride < 1 or buf_shape[dim] % per_dim_stride != 0:
+                return None
 
-        # Extract coefficients from index expression
-        inner_var = var_items[0][0]
-        outer_var = var_items[1][0]
-        index = V.graph.sizevars.simplify(index)
+            result[dim][0] = per_dim_stride
+            var_to_dim[var] = dim
+            remaining = V.graph.sizevars.simplify(remaining - var_expr)
 
-        inner_coeff = self._get_index_coefficient(index, inner_var)
-        outer_coeff = self._get_index_coefficient(index, outer_var)
+        # Remaining is the constant offset; distribute across dims using
+        # divmod with C-contiguous strides (largest stride first).
+        offset_val = self._safe_int(remaining)
+        if offset_val is None:
+            return None
+        if offset_val < 0:
+            return None
+        for d in range(ndim):
+            if c_strides[d] > 0:
+                result[d][1] = offset_val // c_strides[d]
+                offset_val = offset_val % c_strides[d]
+        if offset_val != 0:
+            return None
 
-        if inner_coeff != 0 and outer_coeff != 0:
-            # Only transpose for standard row-major buffers (stride[0] = size[1], stride[1] = 1)
-            is_standard_row_major = s0 == size1 and s1 == 1
-            if not is_standard_row_major:
-                return False
+        # Only return if there's at least one dim with stride > 1.
+        # Contiguous accesses with just an offset (stride=1 everywhere)
+        # are handled by the normal tiling path.
+        if all(s == 1 for s, _ in result):
+            return None
 
-            # Only transpose if output is column-major (indicates actual transpose op)
-            output_is_column_major = self._has_column_major_output()
-            if not output_is_column_major:
-                return False
+        # Decompose each dim's raw offset into (skip, offset) where
+        # offset < stride: raw = skip * stride + offset.
+        # Then validate the output numel matches.
+        decomposed: list[tuple[int, int, int]] = []
+        output_numel_expected = 1
+        for d in range(ndim):
+            stride, raw_offset = result[d]
+            offset = raw_offset % stride
+            skip = raw_offset // stride
+            n_blocks = buf_shape[d] // stride
+            if skip >= n_blocks:
+                return None
+            output_numel_expected *= n_blocks - skip
+            decomposed.append((stride, offset, skip))
 
-            # Check if coefficients indicate transposed access
-            inner_matches_s0 = abs(inner_coeff - s0) < abs(inner_coeff - s1)
-            outer_matches_s1 = abs(outer_coeff - s1) < abs(outer_coeff - s0)
-            return inner_matches_s0 and outer_matches_s1
+        output_numel, _ = self._compute_output_numel_from_index(index)
+        if output_numel != output_numel_expected:
+            return None
 
-        return False
+        # Verify each variable's range matches its assigned dimension's
+        # effective size.  When the kernel collapses multiple buffer dims
+        # into one iteration variable (e.g. batch*channels), the variable
+        # range won't match any single buffer dimension and we must bail
+        # out to avoid shape mismatches in the generated code.
+        for var, dim in var_to_dim.items():
+            if var not in self.range_tree_nodes:
+                return None
+            var_range = self._safe_int(self.range_tree_nodes[var].length)
+            if var_range is None:
+                return None
+            stride_d, _offset_d, skip_d = decomposed[dim]
+            effective_size = buf_shape[dim] // stride_d - skip_d
+            if var_range != effective_size:
+                return None
+
+        return decomposed
+
+    @staticmethod
+    def _strided_load_expr(buf: str, decomp: list[tuple[int, int, int]]) -> str:
+        """Build ``buf[:, :, offset]`` for strided dims, ``:`` for others."""
+        parts: list[str] = []
+        for stride, offset, _skip in decomp:
+            if stride == 1:
+                parts.append(":")
+            else:
+                parts.append(":")  # the halved dim
+                parts.append(str(offset))  # static index into stride dim
+        return f"{buf}[{', '.join(parts)}]"
+
+    def _codegen_strided_reshapes(
+        self, code: IndentedBuffer, params: list[str]
+    ) -> None:
+        """Emit reshape + optional slice for strided input parameters.
+
+        For each strided param, reshapes ``(M, N)`` to ``(M, N/stride, stride)``
+        and, when ``skip > 0``, slices off leading blocks so the remaining
+        elements align with the output.
+        """
+        for param in params:
+            buf_name = self._param_to_buf_name(param)
+            if buf_name is None or buf_name not in self.strided_input_buffers:
+                continue
+            strides = self.strided_input_buffers[buf_name]
+            info = self._get_buffer_info(buf_name)
+            if info is None:
+                continue
+            _, buf_size, _, _, _ = info
+            new_shape_parts: list[str] = []
+            for d, (stride, _offset, _skip) in enumerate(strides):
+                dim = self._safe_int(buf_size[d])
+                if dim is None:
+                    break
+                if stride > 1:
+                    new_shape_parts.append(str(dim // stride))
+                    new_shape_parts.append(str(stride))
+                else:
+                    new_shape_parts.append(str(dim))
+            else:
+                code.writeline(
+                    f"{param} = {param}.reshape({', '.join(new_shape_parts)})"
+                )
+                if any(skip > 0 for _, _, skip in strides):
+                    slice_parts: list[str] = []
+                    for stride, _offset, skip in strides:
+                        if stride == 1:
+                            slice_parts.append(":")
+                        else:
+                            slice_parts.append(f"{skip}:" if skip > 0 else ":")
+                            slice_parts.append(":")
+                    code.writeline(f"{param} = {param}[{', '.join(slice_parts)}]")
+
+    @staticmethod
+    def _get_actual_out_strides(out_buf, n: int) -> list[int] | None:
+        """Extract actual output buffer strides from its layout."""
+        layout = getattr(out_buf, "get_layout", lambda: None)()
+        if layout is None:
+            return None
+        stride_raw = getattr(layout, "stride", None)
+        if stride_raw is None or len(stride_raw) != n:
+            return None
+        strides: list[int] = []
+        for s in stride_raw:
+            v = int(s) if isinstance(s, (int, sympy.Integer)) else None
+            if v is None:
+                return None
+            strides.append(v)
+        return strides
+
+    def _compute_store_coeffs(self, ordered: list) -> dict | None:
+        """Compute store-side linearization coefficients from range tree nesting.
+
+        The tree structure encodes the output iteration order: later
+        trees (prefix ``x``) are innermost, earlier trees (``y``, ``z``)
+        are outer.  Within a tree, dict order goes inner-to-outer.
+        The innermost variable gets coefficient 1; each successive
+        variable (moving outward) multiplies by the previous range.
+
+        Returns ``{sympy.Symbol: int}`` mapping each RT var to its store
+        coefficient, or ``None`` on failure.
+        """
+        prefix_groups: dict[str, list] = {}
+        prefix_order: list[str] = []
+        for v in ordered:
+            node = self.range_tree_nodes[v]
+            p = node.prefix
+            if p not in prefix_groups:
+                prefix_groups[p] = []
+                prefix_order.append(p)
+            prefix_groups[p].append(v)
+        inner_to_outer: list = []
+        for p in reversed(prefix_order):
+            inner_to_outer.extend(prefix_groups[p])
+        coeffs: dict = {}
+        coeff = 1
+        for v in inner_to_outer:
+            sz = self._safe_int(self.range_tree_nodes[v].length)
+            if sz is None:
+                return None
+            coeffs[v] = coeff
+            coeff *= sz
+        return coeffs
+
+    def _get_full_load_permutation(
+        self, name: str, index: sympy.Expr
+    ) -> tuple[int, ...] | None:
+        """Return permutation for a full-array load, or None.
+
+        Computes the permutation by mapping each range-tree variable to
+        both an output dimension (via store coefficients + actual output
+        strides) and an input dimension (via load coefficients + input
+        C-contiguous strides).  The permutation is then:
+
+            perm[out_dim] = in_dim   for each RT variable
+
+        Using actual output strides (not C-contiguous) is critical: the
+        scheduler may choose a non-standard output layout (e.g. column-
+        major) to optimise for transposed inputs.
+
+        When all dimensions collapse to a single flat RT variable (e.g.
+        (2,2,2,2,2) with all dims size 2), infers the permutation
+        directly from output strides vs input C-contiguous strides.
+        """
+        info = self._get_buffer_info(name)
+        if not info:
+            return None
+        _, buf_size, _, _, is_contiguous = info
+        in_shape_raw = [self._safe_int(s) for s in buf_size]
+        if len(in_shape_raw) < 2 or None in in_shape_raw:
+            return None
+        in_shape: list[int] = cast(list[int], in_shape_raw)
+        if not is_contiguous:
+            return None  # .contiguous() at JAX boundary handles this
+
+        # Extract index coefficients for each non-reduction RT variable.
+        iter_used = self._get_used_iter_vars(index)
+        ordered = [
+            s
+            for s, e in self.range_tree_nodes.items()
+            if s in iter_used and not e.is_reduction
+        ]
+        if len(ordered) != len(in_shape):
+            # All dims may have collapsed to a single flat RT variable
+            # (e.g. (2,2,2,2,2) → single x0 of length 32).  In this
+            # case, infer the permutation directly from output strides
+            # vs input C-contiguous strides.
+            n = len(in_shape)
+            if len(ordered) == 1 and self._safe_int(
+                self.range_tree_nodes[ordered[0]].length
+            ) == math.prod(in_shape):
+                in_strides = self._c_contiguous_strides(in_shape)
+                for out_name in self._output_buffer_names:
+                    out_buf = V.graph.get_buffer(out_name)
+                    if out_buf is None:
+                        continue
+                    out_shape = [self._safe_int(s) for s in out_buf.get_size()]
+                    if any(s is None for s in out_shape) or len(out_shape) != n:
+                        continue
+                    actual = self._get_actual_out_strides(out_buf, n)
+                    if actual is None:
+                        break
+                    # Map each output dim to the input dim with the
+                    # same stride.
+                    perm = self._map_coeffs_to_dims(actual, in_strides)
+                    if perm is None:
+                        break
+                    if list(perm) == list(range(n)):
+                        return None
+                    return tuple(perm)
+            return None
+        coeffs_raw = [
+            self._get_index_coefficient(V.graph.sizevars.simplify(index), v)
+            for v in ordered
+        ]
+        if not all(isinstance(c, int) and c > 0 for c in coeffs_raw):
+            return None
+        coeffs: list[int] = cast(list[int], coeffs_raw)
+
+        n = len(ordered)
+        in_strides = self._c_contiguous_strides(in_shape)
+        store_coeffs = self._compute_store_coeffs(ordered)
+
+        # --- Primary path: dimension-mapping with actual output strides ---
+        if store_coeffs is not None:
+            for out_name in self._output_buffer_names:
+                out_buf = V.graph.get_buffer(out_name)
+                if out_buf is None:
+                    continue
+                out_shape = [self._safe_int(s) for s in out_buf.get_size()]
+                if any(s is None for s in out_shape) or len(out_shape) != n:
+                    continue
+
+                actual = self._get_actual_out_strides(out_buf, n)
+                if actual is not None:
+                    rt_to_out = self._map_coeffs_to_dims(
+                        [store_coeffs[v] for v in ordered], actual
+                    )
+                    rt_to_in = self._map_coeffs_to_dims(list(coeffs), in_strides)
+                    if rt_to_out is not None and rt_to_in is not None:
+                        perm = [0] * n
+                        for k in range(n):
+                            perm[rt_to_out[k]] = rt_to_in[k]
+                        if list(perm) == list(range(n)):
+                            return None
+                        return tuple(perm)
+                break
+
+        return None
+
+    def _get_collapsed_load_permutation(
+        self, name: str, index: sympy.Expr
+    ) -> tuple[tuple[int, ...], tuple[int, ...]] | None:
+        """Handle permutation when range tree has collapsed dimensions.
+
+        When simplify_and_reorder merges contiguous dims, the range tree
+        has fewer variables than the buffer's rank.  This method detects
+        the permutation in the collapsed space and returns
+        (collapsed_input_shape, perm) so the caller can generate:
+            jnp.permute_dims(load.reshape(collapsed_shape), perm)
+
+        Uses index coefficients on both sides: load-index coefficients
+        map vars to collapsed input dims, and store-side coefficients
+        (derived from the range tree nesting) map vars to collapsed
+        output dims.  Both sets of strides are always unique, so
+        matching is unambiguous even with duplicate group sizes.
+        """
+        info = self._get_buffer_info(name)
+        if not info:
+            return None
+        _, buf_size, _, _, is_contiguous = info
+        in_shape_raw = [self._safe_int(s) for s in buf_size]
+        if len(in_shape_raw) < 2 or None in in_shape_raw:
+            return None
+        in_shape: list[int] = cast(list[int], in_shape_raw)
+        if not is_contiguous:
+            return None
+
+        iter_used = self._get_used_iter_vars(index)
+        ordered = [
+            s
+            for s, e in self.range_tree_nodes.items()
+            if s in iter_used and not e.is_reduction
+        ]
+        n = len(ordered)
+        if n < 2 or n >= len(in_shape):
+            return None
+        ranges_raw = [self._safe_int(self.range_tree_nodes[v].length) for v in ordered]
+        if None in ranges_raw:
+            return None
+        ranges: list[int] = cast(list[int], ranges_raw)
+        if math.prod(ranges) != math.prod(in_shape):
+            return None
+
+        # Group consecutive input dims (right-to-left) to match ranges
+        in_groups = self._group_dims_to_ranges(in_shape, ranges)
+        if in_groups is None:
+            return None
+
+        # Compute collapsed input strides (row-major) and use load-index
+        # coefficients to map each range tree var to a collapsed input dim.
+        # Strides are always unique, so this is unambiguous even when
+        # group sizes are duplicated.
+        collapsed_in_strides = [0] * n
+        stride = 1
+        for i in range(n - 1, -1, -1):
+            collapsed_in_strides[i] = stride
+            stride *= in_groups[i]
+
+        simplified = V.graph.sizevars.simplify(index)
+        in_coeffs_raw = [self._get_index_coefficient(simplified, v) for v in ordered]
+        if not all(isinstance(c, int) and c > 0 for c in in_coeffs_raw):
+            return None
+        in_coeffs: list[int] = cast(list[int], in_coeffs_raw)
+
+        in_stride_to_dim = {s: i for i, s in enumerate(collapsed_in_strides)}
+        var_to_in_dim = []
+        for coeff in in_coeffs:
+            dim = in_stride_to_dim.get(coeff)
+            if dim is None:
+                return None
+            var_to_in_dim.append(dim)
 
-    def _has_column_major_output(self) -> bool:
-        """Check if any output buffer has column-major stride layout."""
-        output_buffers = getattr(self.args, "output_buffers", {})
-        # Check both explicit output buffers and graph buffers (which may not
-        # be populated during load).
-        buf_names = itertools.chain(output_buffers, V.graph.name_to_buffer)
-        for buf_name in buf_names:
-            out_buf = V.graph.get_buffer(buf_name)
+        store_coeffs = self._compute_store_coeffs(ordered)
+        if store_coeffs is None:
+            return None
+
+        # Find the output-side mapping using store coefficients.
+        for out_name in self._output_buffer_names:
+            out_buf = V.graph.get_buffer(out_name)
             if out_buf is None:
                 continue
-            if buf_name not in output_buffers and not isinstance(
-                out_buf, ComputedBuffer
-            ):
+            out_shape_raw = [self._safe_int(s) for s in out_buf.get_size()]
+            if any(s is None for s in out_shape_raw) or len(out_shape_raw) < 2:
                 continue
-            layout = getattr(out_buf, "get_layout", lambda: None)()
-            if layout is None:
+            out_shape: list[int] = cast(list[int], out_shape_raw)
+            if math.prod(out_shape) != math.prod(in_shape):
                 continue
-            out_stride = getattr(layout, "stride", None)
-            if out_stride is None or len(out_stride) < 2:
+            out_groups = self._group_dims_to_ranges(out_shape, list(in_groups))
+            if out_groups is None:
                 continue
-            out_s0 = self._safe_int(out_stride[0])
-            out_s1 = self._safe_int(out_stride[1])
-            if out_s0 is not None and out_s1 is not None and out_s0 < out_s1:
-                return True
 
-        return False
+            # Compute collapsed output strides and match store coefficients.
+            collapsed_out_strides = [0] * n
+            stride = 1
+            for i in range(n - 1, -1, -1):
+                collapsed_out_strides[i] = stride
+                stride *= out_groups[i]
+
+            out_stride_to_dim = {s: j for j, s in enumerate(collapsed_out_strides)}
+            var_to_out_dim = []
+            for v in ordered:
+                j = out_stride_to_dim.get(store_coeffs[v])
+                if j is None:
+                    return None
+                var_to_out_dim.append(j)
+
+            # Build perm: perm[out_dim] = in_dim
+            perm = [0] * n
+            for k in range(n):
+                perm[var_to_out_dim[k]] = var_to_in_dim[k]
+            if perm == list(range(n)):
+                return None
+            return (tuple(in_groups), tuple(perm))
+        return None
+
+    @staticmethod
+    def _group_dims_to_ranges(dims: list[int], ranges: list[int]) -> list[int] | None:
+        """Group consecutive dims (right-to-left) to match range values.
+
+        Returns collapsed shape (left-to-right) or None if no valid grouping.
+        """
+        available = list(ranges)
+        groups: list[int] = []
+        product = 1
+        for i in range(len(dims) - 1, -1, -1):
+            product *= dims[i]
+            try:
+                idx = available.index(product)
+            except ValueError:
+                continue
+            groups.append(product)
+            available.pop(idx)
+            product = 1
+        if product != 1 or available:
+            return None
+        groups.reverse()
+        return groups
 
-    def _get_index_expr(self, index: sympy.Expr) -> tuple[str, bool]:
+    def _get_index_expr(self, index: sympy.Expr) -> _BufferIndexing:
         """Get the index expression string and whether it needs flattening."""
         has_indirect = self._has_indirect_vars(index)
         has_iter_vars = self._has_iteration_vars(index)
 
         if has_indirect and has_iter_vars:
-            return self._handle_mixed_indexing(index), True
+            return _BufferIndexing(
+                index_str=self._handle_mixed_indexing(index), needs_flatten=True
+            )
         elif has_indirect:
-            return self.kexpr(index), False
+            return _BufferIndexing(index_str=self.kexpr(index), needs_flatten=False)
         else:
             index_str = self._get_index_str(index)
             # Check if index contains ModularIndexing - this requires flattened access
@@ -1204,17 +1669,166 @@ def _get_index_expr(self, index: sympy.Expr) -> tuple[str, bool]:
                 # Check if it's a simple slice pattern (::N or M::N)
                 if not ("::" in index_str or index_str.lstrip("-").isdigit()):
                     needs_flatten = True
-            return index_str, needs_flatten
+            return _BufferIndexing(index_str=index_str, needs_flatten=needs_flatten)
 
     @staticmethod
-    def _safe_int(val: Any) -> Optional[int]:
+    def _safe_int(val: Any) -> int | None:
         """Convert value to int, returning None on failure."""
         try:
             return int(val)
         except (TypeError, ValueError):
             return None
 
-    def _compute_prefix_numel(self, prefixes: OrderedSet) -> Optional[int]:
+    @staticmethod
+    def _c_contiguous_strides(shape: list[int]) -> list[int]:
+        """Return C-contiguous strides for the given shape."""
+        n = len(shape)
+        strides = [1] * n
+        for i in range(n - 2, -1, -1):
+            strides[i] = strides[i + 1] * shape[i + 1]
+        return strides
+
+    @staticmethod
+    def _map_coeffs_to_dims(coeffs: list[int], strides: list[int]) -> list[int] | None:
+        """Map coefficient values to dimension indices via stride matching.
+
+        Returns a list where entry k is the dimension whose stride equals
+        coeffs[k], or None if the mapping is ambiguous or incomplete.
+        """
+        stride_to_dim: dict[int, int] = {}
+        for d, s in enumerate(strides):
+            if s in stride_to_dim:
+                return None  # duplicate strides
+            stride_to_dim[s] = d
+        mapping: list[int] = []
+        for c in coeffs:
+            d = stride_to_dim.get(c)
+            if d is None:
+                return None
+            mapping.append(d)
+        if len(OrderedSet(mapping)) != len(coeffs):
+            return None
+        return mapping
+
+    def _zero_dim_output_flags(self, ctx: _CodegenContext) -> tuple[bool, bool]:
+        """Return whether an output has a zero or unknown dimension."""
+        has_unknown_dim = False
+        for buf_name in ctx.output_buffer_lookup.values():
+            buf = V.graph.try_get_buffer(buf_name)
+            if buf is None:
+                has_unknown_dim = True
+                continue
+            for dim in buf.get_size():
+                dim_int = self._safe_int(dim)
+                if dim_int == 0:
+                    return True, has_unknown_dim
+                if dim_int is None:
+                    has_unknown_dim = True
+        return False, has_unknown_dim
+
+    def _get_reduction_axes(self) -> tuple[int, ...]:
+        """Determine which axes of the loaded array are reduction axes.
+
+        Finds the innermost reduction stride from the load index
+        expression, then walks outward through the buffer's dims
+        using stride ratios until the accumulated product reaches
+        red_numel.  Falls back to stride-direction analysis for
+        gather/flatten loads.
+        """
+        if not self.load_index_exprs:
+            return (-1,)
+
+        r_vars = [v for v, e in self.range_tree_nodes.items() if e.is_reduction]
+        pw_vars = [v for v, e in self.range_tree_nodes.items() if not e.is_reduction]
+        if not r_vars or not pw_vars:
+            return (-1,)
+
+        red_numel = self._compute_reduction_numel()
+        if not red_numel or red_numel <= 1:
+            return (-1,)
+
+        for buf_name, load_index in self.load_index_exprs.items():
+            info = self._get_buffer_info(buf_name)
+            if info is None:
+                continue
+            _, buf_size, _, actual_strides, _ = info
+            nd = len(buf_size)
+            if nd < 2:
+                continue
+            strides_or_none = [self._safe_int(s) for s in actual_strides]
+            if any(s is None for s in strides_or_none):
+                continue
+            strides: list[int] = cast(list[int], strides_or_none)
+
+            # Get reduction stride coefficients by zeroing pw_vars.
+            r_only = load_index
+            for pv in pw_vars:
+                r_only = r_only.subs(pv, 0)
+            r_coeffs: OrderedSet[int] = OrderedSet()
+            for term in sympy.Add.make_args(r_only):
+                if term.is_number:
+                    continue
+                coeff, _ = term.as_coeff_Mul()
+                c = self._safe_int(coeff)
+                if c is not None and c > 0:
+                    r_coeffs.add(c)
+            if not r_coeffs:
+                continue
+
+            # Match all coefficients against buffer strides
+            matched = sorted(
+                (i for i in range(nd) if strides[i] in r_coeffs),
+            )
+            if not matched:
+                continue
+
+            # Multiple r_vars each map to a distinct dim — return directly.
+            # Single r_var with multiple coefficients (transposed access)
+            # → skip to fallback.
+            if len(r_coeffs) > 1:
+                if len(r_coeffs) == len(matched) and len(r_vars) > 1:
+                    return tuple(i - nd for i in matched)
+                continue
+
+            # Single coefficient: walk outward from the matched dim
+            # using span to find flattened contiguous dims.
+            r_stride = next(iter(r_coeffs))
+            span = (red_numel - 1) * r_stride
+            is_contiguous = all(strides[i] > strides[i + 1] for i in range(nd - 1))
+            if is_contiguous:
+                # Walk by dim index (strides are in descending order)
+                inner = matched[-1]
+                start = inner
+                while start > 0 and span > strides[start - 1]:
+                    start -= 1
+                axes = list(range(start, inner + 1))
+            else:
+                # Non-contiguous layout: collect dims whose strides
+                # fall within the r_var's traversal range
+                axes = sorted(
+                    i
+                    for i in range(nd)
+                    if r_stride <= strides[i] and strides[i] < span + r_stride
+                )
+                if not axes:
+                    axes = list(matched)
+            return tuple(i - nd for i in axes)
+
+        # Fallback: stride-direction for gather/flatten loads
+        load_index = next(iter(self.load_index_exprs.values()))
+        r_coeff = load_index.coeff(r_vars[0])
+        r_stride = self._safe_int(r_coeff) if r_coeff != 0 else 1
+        if r_stride is None:
+            r_stride = 1
+        pw_coeff = load_index.coeff(pw_vars[0])
+        pw_stride = self._safe_int(pw_coeff) if pw_coeff != 0 else 1
+        if pw_stride is None:
+            pw_stride = 1
+        if r_stride > pw_stride:
+            return (0,)
+        return (-1,)
+
+    def _compute_prefix_numel(self, prefixes: OrderedSet) -> int | None:
         """Compute total numel for given prefixes (e.g., pointwise prefixes)."""
         result = 1
         for p in prefixes:
@@ -1225,7 +1839,7 @@ def _compute_prefix_numel(self, prefixes: OrderedSet) -> Optional[int]:
                 result *= numel
         return result
 
-    def _compute_reduction_numel(self) -> Optional[int]:
+    def _compute_reduction_numel(self) -> int | None:
         """Compute total reduction numel."""
         result = 1
         for tree in self.range_trees:
@@ -1245,6 +1859,11 @@ def _can_use_tma_approach(self) -> bool:
         - Broadcasting (inputs have different shapes or output differs)
         - Non-contiguous tensors (strided, transposed)
         """
+        # TMA flattens to 1D tiles, incompatible with permutation detection
+        # which emits jnp.permute_dims expecting N-D input.
+        if self.permuted_input_buffers:
+            return False
+
         # Check for reductions
         reduction_numel = self._compute_reduction_numel()
         if reduction_numel is not None and reduction_numel > 1:
@@ -1300,7 +1919,7 @@ def _can_use_tma_approach(self) -> bool:
 
         return True
 
-    def _get_buffer_info(self, name: str) -> Optional[tuple[Any, Any, Any, list, bool]]:
+    def _get_buffer_info(self, name: str) -> tuple[Any, Any, Any, list, bool] | None:
         """Get buffer metadata (buf_obj, buf_size, buf_numel, actual_strides, is_contiguous).
 
         Returns None if the buffer doesn't exist.
@@ -1419,24 +2038,16 @@ def _needs_strided_indexing(
         self,
         name: str,
         index: sympy.Expr,
-        index_str: str,
-        needs_flatten: bool,
-    ) -> tuple[str, bool]:
-        """
-        Check if buffer access needs strided indexing due to size mismatch or gather patterns.
-
-        This handles cases like:
-        - Pooling operations where input/output have different sizes
-        - im2col-like gather patterns
-        - Transposed or strided buffer access
-        """
+        indexing: _BufferIndexing,
+    ) -> _BufferIndexing:
+        """Check if buffer access needs strided indexing due to size mismatch or gather patterns."""
         # Only applies when full array access is indicated
-        if index_str != "..." or needs_flatten:
-            return index_str, needs_flatten
+        if indexing.index_str != "..." or indexing.needs_flatten:
+            return indexing
 
         info = self._get_buffer_info(name)
         if info is None:
-            return index_str, needs_flatten
+            return indexing
 
         buf_obj, buf_size, buf_numel, actual_strides, is_contiguous = info
         output_numel, used_vars = self._compute_output_numel_from_index(index)
@@ -1475,77 +2086,85 @@ def _needs_strided_indexing(
             and not skip_for_non_contiguous
             and not has_symbolic_coef
         ):
-            return self._generate_strided_index(index), True
+            return _BufferIndexing(
+                index_str=self._generate_strided_index(index), needs_flatten=True
+            )
 
-        return index_str, needs_flatten
+        return indexing
 
     def _adjust_index_for_buffer_shape(
         self,
         name: str,
         index: sympy.Expr,
-        index_str: str,
-        needs_flatten: bool,
-    ) -> tuple[str, bool]:
+        indexing: _BufferIndexing,
+    ) -> _BufferIndexing:
         """
         Adjust index expression based on buffer shape (0-dim scalar, multi-dim, etc.).
         """
-        if needs_flatten or index_str == "...":
-            return index_str, needs_flatten
+        if indexing.needs_flatten or indexing.index_str == "...":
+            return indexing
 
         buf_obj = V.graph.get_buffer(name)
         if buf_obj is None:
-            return index_str, needs_flatten
+            return indexing
 
         buf_size = buf_obj.get_size()
 
         # 0-dimensional (scalar) buffer - use [...] to access it
         if len(buf_size) == 0:
-            return "...", needs_flatten
+            return _BufferIndexing(
+                index_str="...", needs_flatten=indexing.needs_flatten
+            )
 
         # Multi-dimensional buffer with constant/scalar index
         if len(buf_size) > 1:
             has_iter_vars = self._has_iteration_vars(index)
             if not has_iter_vars:
-                return index_str, True  # Use flattened access
-            elif "::" in index_str:
+                return _BufferIndexing(
+                    index_str=indexing.index_str, needs_flatten=True
+                )  # Use flattened access
+            elif "::" in indexing.index_str:
                 # Strided slice patterns need flattened indexing for multi-dim
-                return self._generate_strided_index(index), True
+                return _BufferIndexing(
+                    index_str=self._generate_strided_index(index), needs_flatten=True
+                )
 
         # GPU doesn't support gather from slice patterns on 1D buffers
-        if self.is_gpu and "::" in index_str:
-            return self._generate_strided_index(index), True
+        if self.is_gpu and "::" in indexing.index_str:
+            return _BufferIndexing(
+                index_str=self._generate_strided_index(index), needs_flatten=True
+            )
 
-        return index_str, needs_flatten
+        return indexing
 
     def _try_multidim_slice(
         self,
         name: str,
         index: sympy.Expr,
-        index_str: str,
-        needs_flatten: bool,
-    ) -> tuple[str, bool]:
+        indexing: _BufferIndexing,
+    ) -> _BufferIndexing:
         """
         Try to emit multi-dim slice notation instead of flatten + gather.
 
         For a buffer with shape (d0, ..., dk) and index `stride * var + offset`,
         emit `buf[:, ..., :, offset::stride]` when stride divides dk.
         """
-        if not needs_flatten:
-            return index_str, needs_flatten
+        if not indexing.needs_flatten:
+            return indexing
 
         buf_obj = V.graph.get_buffer(name)
         if buf_obj is None:
-            return index_str, needs_flatten
+            return indexing
 
         buf_size = buf_obj.get_size()
         ndim = len(buf_size)
         if ndim < 2:
-            return index_str, needs_flatten
+            return indexing
 
         # Need a single iteration variable with an affine index
         used_vars = self._get_used_iter_vars(index)
         if len(used_vars) != 1:
-            return index_str, needs_flatten
+            return indexing
 
         var = next(iter(used_vars))
         var_expr = BlockPatternMatcher.get_subexpr_involving_symbol(index, var)
@@ -1553,20 +2172,20 @@ def _try_multidim_slice(
             BlockPatternMatcher.match_affine_block_expr(var_expr, var)
         )
         if stride is None or stride <= 1:
-            return index_str, needs_flatten
+            return indexing
 
         offset = V.graph.sizevars.simplify(index - var_expr)
         try:
             offset_val = int(offset)
         except (TypeError, ValueError):
-            return index_str, needs_flatten
+            return indexing
 
         if offset_val < 0 or offset_val >= stride:
-            return index_str, needs_flatten
+            return indexing
 
         last_dim = self._safe_int(buf_size[-1])
         if last_dim is None or last_dim % stride != 0:
-            return index_str, needs_flatten
+            return indexing
 
         # Verify the iteration variable covers all buffer elements at the
         # given stride: var_length * stride == buf_numel. This ensures
@@ -1574,48 +2193,224 @@ def _try_multidim_slice(
         # to buf[:, ..., :, offset::stride].
         entry = self.range_tree_nodes.get(var)
         if entry is None:
-            return index_str, needs_flatten
+            return indexing
         var_length = self._safe_int(entry.length)
         buf_numel = 1
         for s in buf_size:
             d = self._safe_int(s)
             if d is None:
-                return index_str, needs_flatten
+                return indexing
             buf_numel *= d
         if var_length is None or var_length * stride != buf_numel:
-            return index_str, needs_flatten
+            return indexing
 
         prefix = ":, " * (ndim - 1)
         if offset_val == 0:
             slice_str = f"{prefix}::{stride}"
         else:
             slice_str = f"{prefix}{offset_val}::{stride}"
-        return slice_str, False
+        return _BufferIndexing(index_str=slice_str, needs_flatten=False)
+
+    @staticmethod
+    def _gather_permute_expr(load_expr: str, perm: tuple[int, ...]) -> str:
+        """Generate gather-based permutation instead of jnp.permute_dims.
+
+        Avoids a Mosaic compiler bug where jnp.permute_dims produces
+        corrupted output tensors on TPU for 3D+ arrays.  Uses
+        pallas_permute which flattens to 1D and does a 1D gather.
+        """
+        return f"pallas_permute({load_expr}, {perm})"
+
+    def _trace_to_load_source(self, var_name: str) -> str | None:
+        """Trace a tmp variable back to its source buffer's kernel param.
+
+        Follows CSE assignments backward through bounds-checking (where/clamp)
+        until it finds a variable that was directly loaded from a buffer.
+        """
+        if var_name in self._cse_to_param:
+            return self._cse_to_param[var_name]
+        for line in self.compute._lines:
+            line_str = str(line).lstrip()
+            if not line_str.startswith(f"{var_name} = "):
+                continue
+            for ref in re.findall(r"\btmp\d+\b", line_str.split(" = ", 1)[1]):
+                result = self._trace_to_load_source(ref)
+                if result is not None:
+                    return result
+        return None
+
+    def _detect_indirect_access(
+        self, buf: str, name: str, index: sympy.Expr
+    ) -> _IndirectAccessInfo | None:
+        """Detect a load with data-dependent indexing suitable for scalar prefetch.
+
+        Matches exactly one indirect variable whose coefficient corresponds to
+        a C-contiguous stride dimension.  Rejects 1-to-1 gather patterns where
+        the indices buffer covers the full iteration space.
+        """
+        buf_info = self._get_buffer_info(name)
+        if buf_info is None:
+            return None
+        _, buf_size, _, _, _ = buf_info
+        buf_size_raw = [self._safe_int(s) for s in buf_size]
+        if len(buf_size_raw) < 2 or any(s is None for s in buf_size_raw):
+            return None
+        buf_size_ints: list[int] = cast(list[int], buf_size_raw)
+
+        indirect_vars = self._get_indirect_vars(index)
+        if len(indirect_vars) != 1:
+            return None
+        indirect_var = indirect_vars[0]
+
+        coeff = self._get_index_coefficient(index, indirect_var)
+        if coeff == 0 or not isinstance(coeff, int):
+            return None
+
+        # Use existing stride mapping to find which dimension is indirected
+        strides = self._c_contiguous_strides(buf_size_ints)
+        mapping = self._map_coeffs_to_dims([coeff], strides)
+        if mapping is None:
+            return None
+        indirect_dim = mapping[0]
+
+        ndim = len(buf_size_ints)
+        if indirect_dim >= max(1, ndim - 2):
+            return None
+
+        indirect_var_name = str(indirect_var)
+        indices_param = self._trace_to_load_source(indirect_var_name)
+        if indices_param is None:
+            return None
+
+        # Reject gather patterns: only 1-D static index tensors supported
+        indices_graph_name = self._param_to_graph_name.get(indices_param)
+        if indices_graph_name is not None:
+            indices_info = self._get_buffer_info(indices_graph_name)
+            if indices_info is not None:
+                _, indices_size, _, _, _ = indices_info
+                if len(indices_size) != 1:
+                    return None
+                if self._safe_int(indices_size[0]) is None:
+                    return None
+                indices_numel = math.prod(
+                    v for s in indices_size if (v := self._safe_int(s)) is not None
+                )
+                iter_product = math.prod(
+                    length
+                    for var in self._get_used_iter_vars(index)
+                    if var in self.range_tree_nodes
+                    if (length := self._safe_int(self.range_tree_nodes[var].length))
+                    is not None
+                )
+                if indices_numel >= iter_product:
+                    return None
+
+        return _IndirectAccessInfo(
+            table_param=buf,
+            table_buf_name=name,
+            table_shape=tuple(buf_size_ints),
+            indirect_dim=indirect_dim,
+            indirect_var=indirect_var_name,
+            indices_param=indices_param,
+        )
+
+    def _eliminate_dead_indirect_code(self) -> None:
+        """Remove dead compute lines after scalar prefetch replaces indirect load.
+
+        When the table load is simplified to buf[0] (scalar prefetch handles
+        indexing), the indices load and all derived bounds-checking code become
+        dead.  This performs backward liveness analysis from the store variables
+        to identify and remove dead lines.
+        """
+        # Collect variables used by stores (live roots)
+        live_vars: OrderedSet[str] = OrderedSet()
+        for _, store_line in self.store_with_output:
+            for m in re.finditer(r"\btmp\d+\b", store_line):
+                live_vars.add(m.group())
+
+        # Parse assignments from compute lines
+        assignments: list[tuple[str | None, str, Any]] = []
+        for line in self.compute._lines:
+            line_str = str(line).lstrip()
+            m = re.match(r"^(tmp\d+)\s*=\s*(.*)", line_str, re.DOTALL)
+            if m:
+                assignments.append((m.group(1), m.group(2), line))
+            else:
+                assignments.append((None, line_str, line))
+
+        # Propagate liveness backward
+        changed = True
+        while changed:
+            changed = False
+            for var_name, rhs, _ in reversed(assignments):
+                if var_name and var_name in live_vars:
+                    for m in re.finditer(r"\btmp\d+\b", rhs):
+                        if m.group() not in live_vars:
+                            live_vars.add(m.group())
+                            changed = True
+
+        # Keep only live assignments (and non-assignment lines)
+        self.compute._lines = [
+            line
+            for var_name, _, line in assignments
+            if var_name is None or var_name in live_vars
+        ]
 
     def _build_load_expr(
         self,
         buf: str,
         name: str,
         index: sympy.Expr,
-        index_str: str,
-        needs_flatten: bool,
+        indexing: _BufferIndexing,
     ) -> str:
         """
         Build the load expression based on indexing mode.
         """
-        if needs_flatten:
+        if indexing.needs_flatten:
+            # Detect indirect (data-dependent) access for scalar prefetch
+            indirect = self._detect_indirect_access(buf, name, index)
+            if indirect is not None:
+                if self.indirect_access is not None:
+                    # Fused nodes may re-visit the same indirect load (e.g.
+                    # a reduction + pointwise over the same embedding).
+                    # Allow that, but reject truly different indirect accesses.
+                    assert indirect == self.indirect_access, (
+                        "only one indirect access per kernel supported"
+                    )
+                self.indirect_access = indirect
+                return f"{buf}[0]"
+
+            self.has_flatten_indexing = True
+            self.flatten_indexed_buffers.add(name)
             # Flatten then index for non-contiguous access (gather operation)
             has_minmax = index.has(sympy.Min) or index.has(sympy.Max)
-            idx = f"({index_str}).astype(jnp.int64)" if has_minmax else index_str
+            idx_dtype = "jnp.int32" if self.is_tpu else "jnp.int64"
+            idx = (
+                f"({indexing.index_str}).astype({idx_dtype})"
+                if has_minmax
+                else indexing.index_str
+            )
             return f"{buf}[...].flatten()[{idx}]"
         else:
             # Direct indexing for contiguous access
-            load_expr = f"{buf}[{index_str}]"
+            load_expr = f"{buf}[{indexing.index_str}]"
 
-            # Check for transposed access
-            if index_str == "..." and self._is_transposed_access(name, index):
-                load_expr = f"jnp.transpose({load_expr})"
-                self.has_transposed_load = True
+            if indexing.index_str == "..." and not self.is_gpu:
+                perm = self._get_full_load_permutation(name, index)
+                if perm is not None:
+                    load_expr = self._gather_permute_expr(load_expr, perm)
+                    self.permuted_input_buffers[name] = perm
+                else:
+                    collapsed = self._get_collapsed_load_permutation(name, index)
+                    if collapsed is not None:
+                        collapsed_shape, cperm = collapsed
+                        load_expr = f"jnp.permute_dims({load_expr}, {cperm})"
+                        # Don't store cperm in permuted_input_buffers as it's for the reshaped tensor
+                        # not the original shape, which causes issues later when used for tiling
+                        self.collapsed_reshape_inputs[name] = collapsed_shape
+                        self.collapsed_output_shape = tuple(
+                            collapsed_shape[p] for p in cperm
+                        )
 
             return load_expr
 
@@ -1648,18 +2443,28 @@ def _maybe_squeeze_intermediate_buffer(self, name: str, load_expr: str) -> str:
     def _maybe_broadcast_1d_buffer(
         self, name: str, index: sympy.Expr, load_expr: str
     ) -> str:
-        """Reshape 1D buffers (e.g., batch norm mean) for higher-dim broadcasting."""
+        """Reshape 1D buffers for higher-dim broadcasting in reduction kernels.
+
+        When a 1D buffer (e.g. a reduction result from a prior kernel, or a
+        batch-norm parameter) is loaded into a kernel with 2+ iteration dims,
+        JAX right-aligns it for broadcasting: (N,) becomes (1, N).  This is
+        wrong when the buffer corresponds to a non-trailing axis; we reshape
+        to (N, 1, ...) so broadcasting matches the correct axis.
+        """
         buf_obj = V.graph.get_buffer(name)
         if buf_obj is None or len(buf_obj.get_size()) != 1:
             return load_expr
 
+        # Only graph inputs, not intermediate buffers — intermediates are
+        # already shaped by the IR and their dim order may not match the
+        # reference buffer used below for axis inference.
+        if name.startswith("buf"):
+            return load_expr
+
         buf_length = self._safe_int(buf_obj.get_size()[0])
         if buf_length is None:
             return load_expr
 
-        # Only graph inputs, not intermediate buffers or index tensors
-        if name.startswith("buf"):
-            return load_expr
         dtype = V.graph.get_dtype(name)
         if dtype is not None and not dtype.is_floating_point:
             return load_expr
@@ -1689,7 +2494,9 @@ def _maybe_broadcast_1d_buffer(
         if self._safe_int(entry.length) != buf_length:
             return load_expr
 
-        # Buffer length must uniquely match one iteration variable
+        # Buffer length must uniquely match one non-reduction iteration variable.
+        # If multiple pointwise vars share the same length (e.g. 2D pointwise
+        # kernel with both dims equal), the axis is ambiguous and we bail out.
         matching_vars = [
             v
             for v, e in self.range_tree_nodes.items()
@@ -1698,12 +2505,25 @@ def _maybe_broadcast_1d_buffer(
         if len(matching_vars) != 1:
             return load_expr
 
-        # Buffer length must uniquely match one ref buffer dimension
+        # Determine axis position from the iteration variable's position
+        # in the range tree (pointwise vars first, then reduction vars).
+        axis_pos = None
         matching_dims = [i for i, s in enumerate(ref_buf_size) if s == buf_length]
-        if len(matching_dims) != 1:
-            return load_expr
+        if len(matching_dims) == 1:
+            axis_pos = matching_dims[0]
+        else:
+            # Ambiguous by size (e.g. square tensor with reduction).
+            # Use the variable's position in the range tree.
+            pw_idx = 0
+            for sym, e in self.range_tree_nodes.items():
+                if sym == used_var:
+                    axis_pos = pw_idx
+                    break
+                if not e.is_reduction:
+                    pw_idx += 1
 
-        axis_pos = matching_dims[0]
+        if axis_pos is None:
+            return load_expr
         if axis_pos == len(ref_buf_size) - 1:
             return load_expr  # Last dim uses default broadcasting
 
@@ -1712,16 +2532,16 @@ def _maybe_broadcast_1d_buffer(
         return f"{load_expr}.reshape({', '.join(map(str, reshape_dims))})"
 
     def _check_im2col_pattern(
-        self, index: sympy.Expr, index_str: str, needs_flatten: bool
-    ) -> tuple[str, bool]:
+        self, index: sympy.Expr, indexing: _BufferIndexing
+    ) -> _BufferIndexing:
         """
         Check for im2col-like patterns where store uses block variables but load doesn't.
 
         For cat/expand patterns, both load and store prepared indices share block vars.
         For im2col patterns, store compresses to block vars but load doesn't.
         """
-        if index_str != "..." or needs_flatten:
-            return index_str, needs_flatten
+        if indexing.index_str != "..." or indexing.needs_flatten:
+            return indexing
 
         prepared_index = self.prepare_indexing(index)
         iter_vars = self._get_iter_vars()
@@ -1735,7 +2555,7 @@ def _check_im2col_pattern(
 
         # Only trigger if store introduces new block vars
         if not new_vars or len(store_orig_vars) <= 1:
-            return index_str, needs_flatten
+            return indexing
 
         # Check if loads are compatible with broadcast or cat pattern
         has_im2col_pattern = False
@@ -1768,9 +2588,12 @@ def _check_im2col_pattern(
                 break
 
         if has_im2col_pattern:
-            return self._generate_strided_index(prepared_index), True
+            return _BufferIndexing(
+                index_str=self._generate_strided_index(prepared_index),
+                needs_flatten=True,
+            )
 
-        return index_str, needs_flatten
+        return indexing
 
     def _check_load_is_strided_input(
         self, buf_name: str, load_index: sympy.Expr, load_orig_vars: OrderedSet
@@ -1821,7 +2644,7 @@ def _check_store_needs_transpose(self, name: str) -> bool:
         - But input(s) have row-major stride
         - And we haven't already transposed on load
         """
-        if self.has_transposed_load:
+        if self.permuted_input_buffers:
             return False
 
         info = self._get_buffer_info(name)
@@ -1895,8 +2718,7 @@ def _build_store_expr(
         name: str,
         index: sympy.Expr,
         value: CSEVariable,
-        index_str: str,
-        needs_flatten: bool,
+        indexing: _BufferIndexing,
         mode: Any = None,
     ) -> list[str]:
         """
@@ -1904,16 +2726,17 @@ def _build_store_expr(
         mode can be None (set) or "atomic_add" (accumulate).
         Returns a list of lines to emit.
         """
-        if index_str == "...":
+        if indexing.index_str == "...":
             # Full array store with shape matching
             needs_transpose = self._check_store_needs_transpose(name)
             return self._build_full_array_store_expr(out, value, needs_transpose)
 
-        if needs_flatten:
+        if indexing.needs_flatten:
+            self.has_flatten_indexing = True
             # Block variable indexing (e.g., im2col) - use flattened scatter
             scatter_op = "add" if mode == "atomic_add" else "set"
             return [
-                f"{out}[...] = {out}[...].flatten().at[({index_str}).flatten()].{scatter_op}("
+                f"{out}[...] = {out}[...].flatten().at[({indexing.index_str}).flatten()].{scatter_op}("
                 f"jnp.asarray({value}).flatten()).reshape({out}.shape)"
             ]
 
@@ -1931,22 +2754,20 @@ def _build_store_expr(
             # Indirect indexed store (scatter): use .add() for atomic_add, .set() otherwise
             scatter_op = "add" if mode == "atomic_add" else "set"
             lines = [f"_val = jnp.asarray({value})"]
-            value_expr = (
-                f"(jnp.full({index_str}.shape, _val) if _val.ndim == 0 else {value})"
-            )
+            value_expr = f"(jnp.full({indexing.index_str}.shape, _val) if _val.ndim == 0 else {value})"
             if mode == "atomic_add":
                 # For atomic_add, mark output as needing to be readable (for aliasing)
                 self.outputs_need_read.add(out)
                 alias_param = f"{out}_alias"
                 lines.append(
-                    f"{out}[...] = {alias_param}[...].flatten().at[({index_str}).flatten()].{scatter_op}("
+                    f"{out}[...] = {alias_param}[...].flatten().at[({indexing.index_str}).flatten()].{scatter_op}("
                     f"{value_expr}.flatten()).reshape({out}.shape)"
                 )
             else:
-                lines.append(f"{out}[{index_str}] = {value_expr}")
+                lines.append(f"{out}[{indexing.index_str}] = {value_expr}")
             return lines
 
-        return [f"{out}[{index_str}] = {value}"]
+        return [f"{out}[{indexing.index_str}] = {value}"]
 
     def _build_scatter_store_expr(
         self,
@@ -2037,37 +2858,44 @@ def load(self, name: str, index: sympy.Expr) -> CSEVariable:
         self.load_index_exprs[name] = index
 
         # Get base index expression
-        index_str, needs_flatten = self._get_index_expr(index)
+        indexing = self._get_index_expr(index)
 
         # Check for buffer size mismatch requiring strided indexing
-        index_str, needs_flatten = self._needs_strided_indexing(
-            name, index, index_str, needs_flatten
-        )
-
-        # Adjust index for buffer shape (scalar, multi-dim, etc.)
-        index_str, needs_flatten = self._adjust_index_for_buffer_shape(
-            name, index, index_str, needs_flatten
-        )
+        indexing = self._needs_strided_indexing(name, index, indexing)
+
+        # Try strided decomposition before multidim slice or flatten.
+        # This generates reshape + static indexing which works on both
+        # CPU and TPU (unlike slice notation which fails on Mosaic).
+        decomp = self._decompose_strided_access(index, name)
+        if decomp is not None:
+            self.strided_input_buffers[name] = decomp
+            load_expr = self._strided_load_expr(buf, decomp)
+        else:
+            # Adjust index for buffer shape (scalar, multi-dim, etc.)
+            indexing = self._adjust_index_for_buffer_shape(name, index, indexing)
 
-        # Try to emit multi-dim slice instead of flatten + gather
-        index_str, needs_flatten = self._try_multidim_slice(
-            name, index, index_str, needs_flatten
-        )
+            # Try to emit multi-dim slice instead of flatten + gather
+            indexing = self._try_multidim_slice(name, index, indexing)
 
-        # Build the load expression
-        load_expr = self._build_load_expr(buf, name, index, index_str, needs_flatten)
+            # Build the load expression
+            load_expr = self._build_load_expr(buf, name, index, indexing)
 
         # Handle intermediate buffer squeezing for correct broadcasting
-        if not needs_flatten and index_str == "...":
+        if not indexing.needs_flatten and indexing.index_str == "...":
             load_expr = self._maybe_squeeze_intermediate_buffer(name, load_expr)
             # Handle 1D buffer broadcasting for higher-dimensional kernels
             load_expr = self._maybe_broadcast_1d_buffer(name, index, load_expr)
 
-        return self.cse.generate(
+        cse_var = self.cse.generate(
             self.compute,
             load_expr,
             dtype=dtype,
         )
+        # Track CSE var -> param -> graph name for indirect access detection
+        buf_param = self.args.input(name)
+        self._cse_to_param[str(cse_var)] = buf_param
+        self._param_to_graph_name[buf_param] = name
+        return cse_var
 
     def _handle_mixed_indexing(self, index: sympy.Expr) -> str:
         """
@@ -2283,28 +3111,36 @@ def store(
                 f"{out}[...] = jnp.full({out}.shape, _val) if _val.ndim == 0 else _val.reshape({out}.shape)",
             ]
         else:
-            # Check for scatter pattern (indirect indexing for stores)
-            scatter_info = self._detect_scatter_pattern(index, name)
-
-            if scatter_info is not None:
-                # Track iteration variables used in scatter index
-                self.used_iter_vars.update(self._get_used_iter_vars(index))
-                store_lines = [
-                    self._build_scatter_store_expr(out, value, scatter_info, name, mode)
-                ]
+            # When collapsed_output_shape is set, the load-side permutation
+            # already produces data in the correct layout for the collapsed
+            # output.  Force a full-array store ("...") so the scatter index
+            # (which was computed for the original output layout) does not
+            # rearrange the permuted data.
+            if self.collapsed_output_shape is not None:
+                store_lines = self._build_full_array_store_expr(out, value, False)
             else:
-                # Get base index expression
-                index_str, needs_flatten = self._get_index_expr(index)
+                # Check for scatter pattern (indirect indexing for stores)
+                scatter_info = self._detect_scatter_pattern(index, name)
+
+                if scatter_info is not None:
+                    # Track iteration variables used in scatter index
+                    self.used_iter_vars.update(self._get_used_iter_vars(index))
+                    store_lines = [
+                        self._build_scatter_store_expr(
+                            out, value, scatter_info, name, mode
+                        )
+                    ]
+                else:
+                    # Get base index expression
+                    indexing = self._get_index_expr(index)
 
-                # Check for im2col-like patterns
-                index_str, needs_flatten = self._check_im2col_pattern(
-                    index, index_str, needs_flatten
-                )
+                    # Check for im2col-like patterns
+                    indexing = self._check_im2col_pattern(index, indexing)
 
-                # Build the store expression
-                store_lines = self._build_store_expr(
-                    out, name, index, value, index_str, needs_flatten, mode
-                )
+                    # Build the store expression
+                    store_lines = self._build_store_expr(
+                        out, name, index, value, indexing, mode
+                    )
 
         for line in store_lines:
             self.stores.writeline(line)
@@ -2326,7 +3162,7 @@ def _get_index_coefficient(
 
     def _detect_scatter_pattern(
         self, index: sympy.Expr, output_name: str = ""
-    ) -> Optional[dict[str, Any]]:
+    ) -> dict[str, Any] | None:
         """Detect scatter operation pattern. Returns scatter info dict or None."""
         indirect_syms = self._get_indirect_vars(index)
         if len(indirect_syms) != 1:
@@ -2347,7 +3183,7 @@ def _detect_scatter_pattern(
 
     def _detect_point_scatter(
         self, output_name: str, indirect_var: str, indirect_coeff: int
-    ) -> Optional[dict[str, Any]]:
+    ) -> dict[str, Any] | None:
         """Detect single-element scatter pattern."""
         if not output_name:
             return None
@@ -2380,7 +3216,7 @@ def _detect_point_scatter(
 
     def _detect_iter_scatter(
         self, index: sympy.Expr, indirect_var: str, indirect_coeff: int
-    ) -> Optional[dict[str, Any]]:
+    ) -> dict[str, Any] | None:
         """Detect scatter pattern with iteration variables."""
         used_iter_vars = self._get_used_iter_vars(index)
 
@@ -2428,8 +3264,8 @@ def reduction(
         dtype: torch.dtype,
         src_dtype: torch.dtype,
         reduction_type: ReductionType,
-        value: Union[CSEVariable, tuple[CSEVariable, ...]],
-    ) -> Union[CSEVariable, tuple[CSEVariable, ...]]:  # type: ignore[override]
+        value: CSEVariable | tuple[CSEVariable, ...],
+    ) -> CSEVariable | tuple[CSEVariable, ...]:  # type: ignore[override]
         """
         Generate code for reduction operations in JAX/Pallas.
 
@@ -2451,7 +3287,7 @@ def reduction(
                 "Tuple reductions (e.g., welford_combine) not supported in Pallas backend"
             )
 
-        # Check if this reduction is already cached
+        # Check if this reduction is already cached.
         cache_key = (src_dtype, reduction_type, value)
         if cache_key in self.cse.reduction_cache:
             return self.cse.reduction_cache[cache_key]
@@ -2471,102 +3307,52 @@ def reduction(
         # or a full reduction to scalar
         pointwise_prefixes = OrderedSet(["x", "y", "z"])
         has_pointwise = any(p in self.numels for p in pointwise_prefixes)
-
-        # Get the pointwise and reduction numels
-        pointwise_numel: Optional[int] = self._compute_prefix_numel(pointwise_prefixes)
-        reduction_numel: Optional[int] = self._compute_reduction_numel()
-
-        # Count the number of pointwise and reduction dimensions
+        pointwise_numel: int | None = self._compute_prefix_numel(pointwise_prefixes)
+        reduction_numel: int | None = self._compute_reduction_numel()
         n_reduction_dims = sum(
             1 for var, entry in self.range_tree_nodes.items() if entry.is_reduction
         )
 
+        is_partial_reduction = (
+            has_pointwise
+            and pointwise_numel is not None
+            and pointwise_numel > 1
+            and reduction_numel
+            and n_reduction_dims > 0
+        )
+        is_symbolic_partial = (
+            has_pointwise and n_reduction_dims > 0 and pointwise_numel is None
+        )
+
         if reduction_type == "xor_sum":
-            if has_pointwise and pointwise_numel and reduction_numel:
-                reduction_expr = f"jnp.bitwise_xor.reduce({value}.reshape({pointwise_numel}, -1), axis=-1)"
+            if is_partial_reduction:
+                axes = self._get_reduction_axes()
+                axis_expr = axes[0] if len(axes) == 1 else axes
+                reduction_expr = f"jnp.bitwise_xor.reduce({value}, axis={axis_expr})"
             else:
                 reduction_expr = f"jnp.bitwise_xor.reduce({value})"
         elif reduction_type in ("argmax", "argmin"):
-            # For argmax/argmin, the result is indices into the reduction dimension.
-            # Unlike sum/max/min, we can't just reshape because the indices depend
-            # on which axis we reduce over. We need to determine the correct axis.
             reduction_op = reduction_ops[reduction_type]
-            # Check if this is a true partial reduction (pointwise numel > 1)
-            # When pointwise_numel == 1, it's effectively a full reduction to scalar
-            is_partial_reduction = (
-                has_pointwise
-                and pointwise_numel
-                and pointwise_numel > 1
-                and reduction_numel
-            )
-            if is_partial_reduction and n_reduction_dims > 0:
-                # Partial reduction: determine the reduction axis from load index
-                # The reduction variable's coefficient in the index expression tells us its stride
-                # Higher stride = outer axis (lower axis number in row-major order)
-                reduction_axis = -1  # Default to last axis
-                if self.load_index_exprs:
-                    # Get the first load index expression
-                    load_index = next(iter(self.load_index_exprs.values()))
-                    # Find the reduction variable (starts with 'r')
-                    reduction_vars = [
-                        var
-                        for var, entry in self.range_tree_nodes.items()
-                        if entry.is_reduction
-                    ]
-                    if reduction_vars:
-                        r_var = reduction_vars[0]
-                        # Get the coefficient (stride) of the reduction variable
-                        r_coeff = load_index.coeff(r_var)
-                        r_stride = self._safe_int(r_coeff) if r_coeff != 0 else 1
-                        if r_stride is None:
-                            r_stride = 1
-                        # Get pointwise variable
-                        pw_vars = [
-                            var
-                            for var, entry in self.range_tree_nodes.items()
-                            if not entry.is_reduction
-                        ]
-                        if pw_vars:
-                            pw_var = pw_vars[0]
-                            pw_coeff = load_index.coeff(pw_var)
-                            pw_stride = self._safe_int(pw_coeff) if pw_coeff != 0 else 1
-                            if pw_stride is None:
-                                pw_stride = 1
-                            # Higher stride = earlier (outer) axis
-                            # For 2D: axis 0 has stride = dim1_size, axis 1 has stride = 1
-                            reduction_axis = 0 if r_stride > pw_stride else -1
-                reduction_expr = f"{reduction_op}({value}, axis={reduction_axis})"
+            if is_partial_reduction:
+                # argmax/argmin only accept a single axis
+                axes = self._get_reduction_axes()
+                reduction_expr = f"{reduction_op}({value}, axis={axes[-1]})"
             else:
-                # Full reduction to scalar
                 reduction_expr = f"{reduction_op}({value})"
         elif reduction_type in reduction_ops:
-            # Check for true partial reduction (pointwise_numel > 1 means we have
-            # actual pointwise dimensions, not just a scalar placeholder)
-            is_partial_reduction = (
-                has_pointwise
-                and pointwise_numel is not None
-                and pointwise_numel > 1
-                and reduction_numel
-            )
-            # Also check for symbolic partial reduction (has both pw and reduction vars)
-            is_symbolic_partial = (
-                has_pointwise and n_reduction_dims > 0 and pointwise_numel is None
-            )
+            reduction_op = reduction_ops[reduction_type]
             if is_partial_reduction:
-                # For partial reductions, we need to:
-                # 1. Find which axes are reduction axes (contiguous axes whose product = reduction_numel)
-                # 2. Move pointwise axes to front, reduction axes to back
-                # 3. Reshape to (pointwise_numel, reduction_numel) and reduce over last axis
-                # 4. Reshape output with 1s in reduced dims for proper broadcasting
-                reduction_op = reduction_ops[reduction_type]
-                # Use a helper to find reduction axes by product matching
-                reduction_expr = f"pallas_partial_reduce({reduction_op}, {value}, {pointwise_numel}, {reduction_numel})"
+                axes = self._get_reduction_axes()
+                axis_expr = axes[0] if len(axes) == 1 else axes
+                reduction_expr = (
+                    f"{reduction_op}({value}, axis={axis_expr}, keepdims=True)"
+                )
             elif is_symbolic_partial:
-                # Symbolic sizes: use axis-based reduction (axis=0 for outer reduction)
-                reduction_expr = f"{reduction_ops[reduction_type]}({value}, axis=0)"
+                # With symbolic shapes, strided loads produce a degenerate
+                # batch dim at axis=0 that just needs squeezing.
+                reduction_expr = f"{reduction_op}({value}, axis=0)"
             else:
-                # Full reduction to scalar
-                reduction_expr = f"{reduction_ops[reduction_type]}({value})"
+                reduction_expr = f"{reduction_op}({value})"
         else:
             raise Unsupported(
                 f"Reduction type '{reduction_type}' not yet supported in Pallas backend. "
@@ -2590,7 +3376,205 @@ def _buffer_is_contiguous(buffer_name: str) -> bool:
         layout = buf.get_layout()
         return layout.is_contiguous()
 
-    def codegen_kernel(self, name: Optional[str] = None) -> str:  # type: ignore[override]
+    def _can_tile_cpu_tpu(self) -> bool:
+        """Check if this kernel can use tiling on CPU/TPU.
+
+        Tiling is compatible with reductions, transpositions, and multi-range-tree
+        kernels as long as no flatten-based indexing is used (buf[...].flatten()[idx]).
+        Flatten indexing requires global flat indices which don't decompose into
+        per-tile local indices.
+
+        Reject:
+        - GPU (has its own TMA / padding path)
+        - Flatten-based indexing
+        - Scatter outputs (indirect indexing complicates tile boundaries)
+        """
+        if self.is_gpu:
+            return False
+        if self.has_flatten_indexing:
+            return False
+        if self.outputs_need_read:
+            return False
+
+        # If iteration variables appear in the compute body (not just in
+        # load/store index resolution that collapses to [...]), tiling is
+        # unsafe because the arange-based vars have full-tensor shapes.
+        # Exception: vars emitted in tile-relative form are safe.
+        if self.used_iter_vars:
+            compute_text = "\n".join(str(line) for line in self.compute._lines)
+            for var_sym in self.used_iter_vars:
+                if var_sym in self.tile_relative_iter_vars:
+                    continue
+                if str(var_sym) in compute_text:
+                    return False
+
+        # Determine the reference output shape (highest-ndim output).
+        out_bufs = list(self.args.output_buffers.keys())
+
+        # Only check the current kernel's actual output buffers for transpose,
+        # not _has_column_major_output() which scans all graph buffers and can
+        # be triggered by unrelated intermediates (e.g., (N,1) reductions with
+        # degenerate column-major strides).
+        has_col_major_out = False
+        for buf_name in out_bufs:
+            info = self._get_buffer_info(buf_name)
+            if info is None:
+                continue
+            _, buf_size, _, actual_strides, _ = info
+            if len(actual_strides) >= 2 and len(buf_size) >= 2:
+                s0 = actual_strides[0]
+                s1 = actual_strides[1]
+                d0 = self._safe_int(buf_size[0])
+                d1 = self._safe_int(buf_size[1])
+                if (
+                    s0 is not None
+                    and s1 is not None
+                    and s0 < s1
+                    and d0 is not None
+                    and d1 is not None
+                    and d0 > 1
+                    and d1 > 1
+                ):
+                    has_col_major_out = True
+                    break
+        self.tile_has_transpose = bool(self.permuted_input_buffers) or has_col_major_out
+
+        # Count trailing reduction dimensions in the output shape that must
+        # not be tiled (the kernel body needs the full reduction range).
+        # Only count when the kernel actually performs reduction (numel > 1).
+        reduction_numel = self._compute_reduction_numel()
+        has_reduction = reduction_numel is not None and reduction_numel > 1
+        self.tile_skip_last_n = (
+            sum(1 for tree in self.range_trees if tree.is_reduction)
+            if has_reduction
+            else 0
+        )
+
+        ref_shape: list[int] = []
+        for buf_name in out_bufs:
+            info = self._get_buffer_info(buf_name)
+            if info is None:
+                return False
+            _, buf_size, _, _, _ = info
+            int_size = [self._safe_int(s) for s in buf_size]
+            if any(s is None for s in int_size):
+                return False
+            if len(int_size) > len(ref_shape):
+                ref_shape = int_size  # type: ignore[assignment]
+
+        if not ref_shape:
+            return False
+
+        # For collapsed permutation kernels, override ref_shape with the
+        # collapsed output shape so all compatibility checks operate in
+        # collapsed-shape space.
+        if self.collapsed_output_shape is not None:
+            ref_shape = list(self.collapsed_output_shape)
+
+        ref_nd = len(ref_shape)
+
+        all_bufs = list(self.args.input_buffers) + out_bufs
+        has_tileable = False
+        for buf_name in all_bufs:
+            info = self._get_buffer_info(buf_name)
+            if info is None:
+                return False
+            _, buf_size, _, _, _ = info
+            if len(buf_size) == 0:
+                continue  # scalar
+
+            # Use collapsed shapes when available so dimension checks
+            # operate in the same space as the kernel.
+            if buf_name in self.collapsed_reshape_inputs:
+                int_size = list(self.collapsed_reshape_inputs[buf_name])
+            elif self.collapsed_output_shape is not None and buf_name in out_bufs:
+                int_size = list(self.collapsed_output_shape)
+            else:
+                int_size = [self._safe_int(s) for s in buf_size]
+                if any(s is None for s in int_size):
+                    return False
+            buf_nd = len(int_size)
+
+            if buf_nd == ref_nd:
+                # Same ndim: check dimensions match or are broadcast (1).
+                # Allow strided buffers (dims may differ after reshape).
+                is_strided = buf_name in self.strided_input_buffers
+                mismatch = False
+                for i in range(ref_nd):
+                    if (
+                        int_size[i] == ref_shape[i]
+                        or int_size[i] == 1
+                        or ref_shape[i] == 1
+                        or is_strided
+                    ):
+                        continue
+                    mismatch = True
+                    break
+
+                if mismatch and buf_name in self.permuted_input_buffers:
+                    perm = self.permuted_input_buffers[buf_name]
+                    if not (
+                        len(perm) == ref_nd
+                        and all(
+                            int_size[perm[i]] == ref_shape[i]
+                            or int_size[perm[i]] == 1
+                            or ref_shape[i] == 1
+                            for i in range(ref_nd)
+                        )
+                    ):
+                        return False
+                elif mismatch:
+                    return False
+
+                # At least one buffer with a tileable dim
+                if is_strided or any(
+                    int_size[i] == ref_shape[i] and ref_shape[i] > 1
+                    for i in range(ref_nd)
+                ):
+                    has_tileable = True
+
+            elif buf_nd > ref_nd:
+                # Reduction input with extra dims. Find an alignment offset k
+                # such that buf_shape[k+i] == ref_shape[i] for all i (skipping
+                # broadcast dims where ref_shape[i] == 1).
+                found = False
+                for k in range(buf_nd - ref_nd + 1):
+                    ok = True
+                    for i in range(ref_nd):
+                        if ref_shape[i] == 1:
+                            continue
+                        if int_size[k + i] != ref_shape[i]:
+                            ok = False
+                            break
+                    if ok:
+                        found = True
+                        break
+                if not found:
+                    return False
+                has_tileable = True
+
+            else:
+                # Fewer dims: verify numpy-style broadcastability
+                for a, b in zip(reversed(int_size), reversed(ref_shape)):
+                    if a != b and a != 1 and b != 1:
+                        return False
+
+        if not has_tileable:
+            return False
+
+        # On CPU (interpret mode) each tile iteration has significant
+        # Python/JAX overhead, so cap the grid size.  Store the cap
+        # so _codegen_tiled_specs can pass it to pallas_compute_tiling,
+        # which will scale up tiles to stay within the limit.
+        is_tpu = V.graph.get_current_device_or_throw().type == "tpu"
+        if not is_tpu:
+            self._cpu_max_grid_product = 64
+        else:
+            self._cpu_max_grid_product = None
+
+        return True
+
+    def codegen_kernel(self, name: str | None = None) -> str:  # type: ignore[override]
         """
         Generate the complete Pallas kernel code as a Python string.
 
@@ -2627,7 +3611,6 @@ def codegen_kernel(self, name: Optional[str] = None) -> str:  # type: ignore[ove
         }
 
         kernel_name = name or "<KERNEL_NAME>"
-        is_tpu = V.graph.get_current_device_or_throw().type == "tpu"
         interpret_is_cpu = V.graph.get_current_device_or_throw().type == "cpu"
         interpret_literal = "True" if interpret_is_cpu else "False"
 
@@ -2645,12 +3628,15 @@ def codegen_kernel(self, name: Optional[str] = None) -> str:  # type: ignore[ove
         non_alias_out_set = OrderedSet(
             [name for name, flag in aliasable_flags.items() if not flag]
         )
-        # On CPU (interpret=True) and TPU, pallas_call returns new arrays so
-        # we must copy back every output.  On CUDA, aliased outputs are
-        # mutated in-place by the donated-buffer mechanism so only
-        # non-aliased outputs need an explicit copy.
-        if interpret_is_cpu or is_tpu:
+        # On CPU (interpret=True), pallas_call returns new arrays so we must
+        # copy back every output.  On TPU, call_custom_kernel with
+        # input_output_aliases handles donation (zero-copy), so no copy is
+        # needed.  On CUDA, aliased outputs are mutated in-place by the
+        # donated-buffer mechanism so only non-aliased outputs need a copy.
+        if interpret_is_cpu:
             copy_output_indices = list(range(len(output_params)))
+        elif self.is_tpu:
+            copy_output_indices = []
         else:
             copy_output_indices = [
                 idx
@@ -2661,7 +3647,7 @@ def codegen_kernel(self, name: Optional[str] = None) -> str:  # type: ignore[ove
         ctx = _CodegenContext(
             code=code,
             kernel_name=kernel_name,
-            is_tpu=is_tpu,
+            is_tpu=self.is_tpu,
             interpret_is_cpu=interpret_is_cpu,
             interpret_literal=interpret_literal,
             kernel_params=kernel_params,
@@ -2676,14 +3662,12 @@ def codegen_kernel(self, name: Optional[str] = None) -> str:  # type: ignore[ove
             full_kernel_params=full_kernel_params,
             non_alias_out_set=non_alias_out_set,
             copy_output_indices=copy_output_indices,
+            alias_pairs=[],
         )
         self.aliasable_out_ptrs = aliasable_flags
 
         self._codegen_imports(ctx)
 
-        # Generate kernel body into a separate buffer first.
-        # This allows us to discover all size variables (registered via rename_indexing)
-        # before generating the kernel signature.
         kernel_body = IndentedBuffer()
         with kernel_body.indent():
             self._codegen_iteration_vars(kernel_body, ctx)
@@ -2705,23 +3689,37 @@ def codegen_kernel(self, name: Optional[str] = None) -> str:  # type: ignore[ove
         ctx.kernel_input_params = alias_params + ctx.pointer_tail
         ctx.full_kernel_params = alias_params + kernel_params
 
+        # Decide whether to use tiling for CPU/TPU after kernel body is fully
+        # generated (used_iter_vars is populated during load/store codegen).
+        self.tile_cpu_tpu = self._can_tile_cpu_tpu()
+
+        extra_kernel_params = ""
+        if self.tile_relative_iter_vars:
+            extra_kernel_params = ", _pallas_tile=None, _pallas_ax2g=None"
+
+        ctx.alias_pairs = self._compute_alias_pairs(ctx, aliasable_flags)
+
+        use_scalar_prefetch = bool(self.indirect_access)
+
+        if use_scalar_prefetch:
+            self._eliminate_dead_indirect_code()
+            kernel_body_sp = IndentedBuffer()
+            with kernel_body_sp.indent():
+                for line in self.compute._lines:
+                    kernel_body_sp.writeline(str(line))
+            self._codegen_scalar_prefetch_wrapper(
+                ctx,
+                kernel_name,
+                kernel_body_sp,
+            )
+            return code.getvalue()
+
         # Emit the kernel function with the correct signature
-        kernel_signature = (
-            f"def {kernel_name}_kernel({', '.join(ctx.full_kernel_params)}):"
-        )
+        kernel_signature = f"def {kernel_name}_kernel({', '.join(ctx.full_kernel_params)}{extra_kernel_params}):"
         code.writeline(kernel_signature)
 
         with code.indent():
-            for line in kernel_body._lines:
-                if isinstance(line, str):
-                    code.writeline(line.lstrip())
-                else:
-                    code._lines.append(line)
-
-            # Filter stores to only emit those for outputs that are in kernel params.
-            for out_ptr, store_line in self.store_with_output:
-                if out_ptr in ctx.full_kernel_params:
-                    code.writeline(store_line)
+            self._emit_kernel_body(code, kernel_body, ctx)
 
         code.writeline("")
         jit_wrapper_name = f"{kernel_name}_jit_wrapper"
@@ -2745,57 +3743,317 @@ def codegen_kernel(self, name: Optional[str] = None) -> str:  # type: ignore[ove
             ["out_shapes", "out_dtypes"] + ctx.size_var_params + ctx.kernel_input_params
         )
         code.writeline(f"def {jit_wrapper_name}({', '.join(wrapper_params)}):")
+
+        alias_map_literal = ", ".join(f"{i}: {o}" for (i, o) in ctx.alias_pairs)
+
+        has_zero_dim, has_unknown_dim = self._zero_dim_output_flags(ctx)
+
+        zero_dim_return = (
+            "results = tuple(jnp.empty(s, dtype=dt) "
+            "for s, dt in zip(out_shapes, out_dtypes))",
+            "return results if len(results) > 1 else results[0]",
+        )
+
         with code.indent():
-            code.writeline("out_shapes_pallas = tuple(")
-            code.writeline("    jax.ShapeDtypeStruct(shape, dtype)")
-            code.writeline("    for shape, dtype in zip(out_shapes, out_dtypes)")
-            code.writeline(")")
-            code.writeline("indexer = lambda n: lambda i: [i] * n")
-            code.writeline("out_specs_pallas = tuple(")
-            code.writeline("    pl.BlockSpec(shape, indexer(len(shape)))")
-            code.writeline("    for shape, dtype in zip(out_shapes, out_dtypes)")
-            code.writeline(")")
-            code.writeline("in_specs_pallas = tuple(")
-            code.writeline("    pl.BlockSpec(i.shape, indexer(len(i.shape)))")
-            code.writeline("    for i in [" + ", ".join(ctx.kernel_input_params) + "]")
-            code.writeline(")")
+            if has_zero_dim:
+                code.writelines(zero_dim_return)
+            else:
+                if has_unknown_dim:
+                    code.writeline("if any(0 in shape for shape in out_shapes):")
+                    with code.indent():
+                        code.writelines(zero_dim_return)
+                # Pallas requires >= 1-d tensors; promote 0-d to (1,)
+                code.writeline(
+                    "_pallas_out_shapes = tuple("
+                    "s if len(s) > 0 else (1,) for s in out_shapes)"
+                )
+                if self.collapsed_output_shape is not None:
+                    code.writeline(
+                        f"_pallas_out_shapes = ({self.collapsed_output_shape},)"
+                    )
+                # Reshape aliased inputs to match promoted output shapes
+                for input_idx, out_idx in ctx.alias_pairs:
+                    param = ctx.kernel_input_params[input_idx]
+                    code.writeline(
+                        f"{param} = {param}.reshape(_pallas_out_shapes[{out_idx}])"
+                    )
+                code.writeline("out_shapes_pallas = tuple(")
+                code.writeline("    jax.ShapeDtypeStruct(shape, dtype)")
+                code.writeline(
+                    "    for shape, dtype in zip(_pallas_out_shapes, out_dtypes)"
+                )
+                code.writeline(")")
+                if self.tile_cpu_tpu:
+                    self._codegen_tiled_specs(ctx)
+                else:
+                    self._codegen_strided_reshapes(code, ctx.kernel_input_params)
+                    for param in ctx.kernel_input_params:
+                        buf_name = self._param_to_buf_name(param)
+                        cshape = (
+                            self.collapsed_reshape_inputs.get(buf_name)
+                            if buf_name
+                            else None
+                        )
+                        if cshape is not None:
+                            code.writeline(f"{param} = {param}.reshape({cshape})")
 
-            alias_pairs: list[tuple[int, int]] = []
-            for out_idx, name in enumerate(ctx.output_params):
-                if name.startswith("out_ptr"):
-                    if aliasable_flags.get(name, False):
-                        alias_name = f"{name}_alias"
-                        input_idx = ctx.kernel_input_params.index(alias_name)
-                        alias_pairs.append((input_idx, out_idx))
+                    code.writeline("out_specs_pallas = tuple(")
+                    code.writeline("    pallas_make_block_spec_non_tiled(shape)")
+                    code.writeline(
+                        "    for shape, dtype in zip(_pallas_out_shapes, out_dtypes)"
+                    )
+                    code.writeline(")")
+                    code.writeline("in_specs_pallas = tuple(")
+                    code.writeline("    pallas_make_block_spec_non_tiled(i.shape)")
+                    code.writeline(
+                        "    for i in [" + ", ".join(ctx.kernel_input_params) + "]"
+                    )
+                    code.writeline(")")
+
+                if self.tile_relative_iter_vars:
+                    if self.tile_cpu_tpu:
+                        code.writeline("_pallas_tile = _tile")
+                        code.writeline("_pallas_ax2g = _ax2g")
+                    else:
+                        code.writeline("_pallas_tile = _pallas_out_shapes[0]")
+                        code.writeline("_pallas_ax2g = {}")
+
+                # Wrap kernel with functools.partial to pass scalar arguments (size variables)
+                partial_args = []
+                for sv_param in ctx.size_var_params:
+                    partial_args.append(f"{sv_param}={sv_param}")
+
+                if self.tile_relative_iter_vars:
+                    partial_args.append("_pallas_tile=_pallas_tile")
+                    partial_args.append("_pallas_ax2g=_pallas_ax2g")
+
+                if partial_args:
+                    kernel_arg = f"functools.partial({kernel_name}_kernel, {', '.join(partial_args)}),"
+                else:
+                    kernel_arg = f"{kernel_name}_kernel,"
+
+                use_tma = (
+                    self.is_gpu
+                    and self.use_emit_pipeline
+                    and self._can_use_tma_approach()
+                )
+                if use_tma:
+                    self._codegen_jit_wrapper_tma(ctx, kernel_arg)
+                elif self.is_gpu:
+                    self._codegen_jit_wrapper_legacy_gpu(ctx, kernel_arg)
                 else:
-                    input_idx = ctx.kernel_input_params.index(name)
+                    self._codegen_jit_wrapper_cpu_tpu(
+                        ctx, kernel_arg, ctx.alias_pairs, alias_map_literal
+                    )
+
+        self._codegen_main_entry(ctx, jit_wrapper_name)
+        return code.getvalue()
+
+    @staticmethod
+    def _compute_alias_pairs(
+        ctx: _CodegenContext, aliasable_flags: dict[str, bool]
+    ) -> list[tuple[int, int]]:
+        alias_pairs: list[tuple[int, int]] = []
+        for out_idx, name in enumerate(ctx.output_params):
+            if name.startswith("out_ptr"):
+                if aliasable_flags.get(name, False):
+                    alias_name = f"{name}_alias"
+                    input_idx = ctx.kernel_input_params.index(alias_name)
                     alias_pairs.append((input_idx, out_idx))
-            alias_map_literal = ", ".join(f"{i}: {o}" for (i, o) in alias_pairs)
+            else:
+                input_idx = ctx.kernel_input_params.index(name)
+                alias_pairs.append((input_idx, out_idx))
+        return alias_pairs
+
+    def _emit_kernel_body(
+        self,
+        code: IndentedBuffer,
+        kernel_body: IndentedBuffer,
+        ctx: _CodegenContext,
+    ) -> None:
+        """Emit the kernel body lines and store operations into code."""
+        for line in kernel_body._lines:
+            if isinstance(line, str):
+                code.writeline(line.lstrip())
+            else:
+                code._lines.append(line)
+        for out_ptr, store_line in self.store_with_output:
+            if out_ptr in ctx.full_kernel_params:
+                code.writeline(store_line)
+
+    def _codegen_scalar_prefetch_wrapper(
+        self,
+        ctx: _CodegenContext,
+        kernel_name: str,
+        kernel_body: IndentedBuffer,
+    ) -> None:
+        """Emit kernel, JIT wrapper, and main entry for scalar prefetch."""
+        assert self.indirect_access is not None
+        indirect = self.indirect_access
+        code = ctx.code
+
+        alias_set = OrderedSet(ctx.alias_params)
+        other_input_params = [
+            p
+            for p in ctx.kernel_input_params
+            if p != indirect.indices_param
+            and p != indirect.table_param
+            and p not in alias_set
+        ]
+
+        # Emit kernel function with params reordered for PrefetchScalarGridSpec:
+        # [scalar_prefetch] + [in_specs refs] + [out_specs refs]
+        prefetch_kernel_params = (
+            [indirect.indices_param]
+            + [indirect.table_param]
+            + other_input_params
+            + list(ctx.alias_params)
+            + ctx.output_params
+        )
+        code.writeline(
+            f"def {kernel_name}_kernel({', '.join(prefetch_kernel_params)}):"
+        )
+        with code.indent():
+            self._emit_kernel_body(code, kernel_body, ctx)
+
+        # Emit JIT wrapper
+        code.writeline("")
+        jit_wrapper_name = f"{kernel_name}_jit_wrapper"
+        wrapper_params = (
+            ["out_shapes", "out_dtypes"] + ctx.size_var_params + ctx.kernel_input_params
+        )
+        static_argnums = list(range(2 + len(ctx.size_var_params)))
+        static_argnums_literal = "(" + ", ".join(str(x) for x in static_argnums) + ",)"
+        code.writeline(
+            f"@functools.partial(jax.jit, static_argnums={static_argnums_literal})"
+        )
+        code.writeline(f"def {jit_wrapper_name}({', '.join(wrapper_params)}):")
 
-            # Wrap kernel with functools.partial to pass scalar arguments (size variables)
-            partial_args = []
-            for sv_param in ctx.size_var_params:
-                partial_args.append(f"{sv_param}={sv_param}")
+        with code.indent():
+            table = indirect.table_param
+            indices = indirect.indices_param
+
+            ind_dim = indirect.indirect_dim
+            ndim = len(indirect.table_shape)
+            code.writeline("_D = 1")
+            for i in range(ndim):
+                if i != ind_dim:
+                    code.writeline(f"_D = _D * {table}.shape[{i}]")
+            code.writeline(f"_seq = {indices}.shape[0]")
+
+            if ind_dim == 0:
+                code.writeline(f"_table_3d = {table}.reshape({table}.shape[0], 1, _D)")
+            else:
+                perm = (ind_dim, *[d for d in range(ndim) if d != ind_dim])
+                code.writeline(
+                    f"_table_3d = {table}.transpose{perm}.reshape("
+                    f"{table}.shape[{ind_dim}], 1, _D)"
+                )
 
+            # Reshape other (non-table, non-indices) inputs to 3D to match the
+            # table's (seq, 1, D) layout.  Currently handles:
+            #   - 2D with leading dim == seq: row-aligned, reshape to (seq, 1, D)
+            #   - 1D: broadcast scalar/vector, reshape to (1, 1, numel)
+            #   - else: flatten to (1, 1, -1) — assumes broadcastable with
+            #     (seq, 1, D).  This may not work correctly for 3D+ inputs.
+            pallas_call_other_args = []
+            for p in other_input_params:
+                p3d = f"_{p}_3d"
+                code.writeline(f"if {p}.ndim == 2 and {p}.shape[0] == _seq:")
+                code.writeline(f"    {p3d} = {p}.reshape(_seq, 1, _D)")
+                code.writeline(f"elif {p}.ndim == 1:")
+                code.writeline(f"    {p3d} = {p}.reshape(1, 1, {p}.shape[0])")
+                code.writeline("else:")
+                code.writeline(f"    {p3d} = {p}.reshape(1, 1, -1)")
+                pallas_call_other_args.append(p3d)
+
+            pallas_call_alias_args = []
+            for p in ctx.alias_params:
+                p3d = f"_{p}_3d"
+                code.writeline(f"{p3d} = {p}.reshape(_seq, 1, _D)")
+                pallas_call_alias_args.append(p3d)
+
+            partial_args = [f"{sv}={sv}" for sv in ctx.size_var_params]
             if partial_args:
-                kernel_arg = f"functools.partial({kernel_name}_kernel, {', '.join(partial_args)}),"
+                kernel_ref = (
+                    f"functools.partial({kernel_name}_kernel,"
+                    f" {', '.join(partial_args)})"
+                )
             else:
-                kernel_arg = f"{kernel_name}_kernel,"
+                kernel_ref = f"{kernel_name}_kernel"
 
-            use_tma = (
-                self.is_gpu and self.use_emit_pipeline and self._can_use_tma_approach()
+            # Reusable row-tiled BlockSpec (all i32 index_map for Mosaic compat)
+            code.writeline(
+                "_ROW_SPEC = pl.BlockSpec((1, 1, _D),"
+                " lambda i, _: (i, jnp.int32(0), jnp.int32(0)))"
             )
-            if use_tma:
-                self._codegen_jit_wrapper_tma(ctx, kernel_arg)
-            elif self.is_gpu:
-                self._codegen_jit_wrapper_legacy_gpu(ctx, kernel_arg)
-            else:
-                self._codegen_jit_wrapper_cpu_tpu(
-                    ctx, kernel_arg, alias_pairs, alias_map_literal
+
+            num_non_alias_in_specs = 1 + len(pallas_call_other_args)
+            code.writeline("_in_specs = [")
+            with code.indent():
+                code.writeline(
+                    "pl.BlockSpec((1, 1, _D),"
+                    " lambda gi, idx: (idx[gi], jnp.int32(0), jnp.int32(0))),"
                 )
+                for p3d in pallas_call_other_args:
+                    code.writeline(
+                        f"_ROW_SPEC"
+                        f" if {p3d}.shape[0] == _seq else"
+                        f" pl.BlockSpec({p3d}.shape,"
+                        f" lambda i, _: (jnp.int32(0), jnp.int32(0), jnp.int32(0))),"
+                    )
+                for _ in ctx.alias_params:
+                    code.writeline("_ROW_SPEC,")
+            code.writeline("]")
+
+            num_outputs = len(ctx.output_params)
+            code.writeline(
+                "_out_specs = [" + ", ".join(["_ROW_SPEC"] * num_outputs) + "]"
+            )
+
+            # input_output_aliases: pallas_call arg index -> output index
+            # (offset by 1 for scalar prefetch arg)
+            alias_map_parts = []
+            for out_idx, _ in enumerate(ctx.alias_params):
+                arg_idx = 1 + num_non_alias_in_specs + out_idx
+                alias_map_parts.append(f"{arg_idx}: {out_idx}")
+            alias_map_literal = ", ".join(alias_map_parts)
+
+            out_shape_parts = [
+                f"jax.ShapeDtypeStruct((_seq, 1, _D), out_dtypes[{i}])"
+                for i in range(num_outputs)
+            ]
+            out_shape_expr = "[" + ", ".join(out_shape_parts) + "]"
+
+            code.writeline("_result = pl.pallas_call(")
+            with code.indent():
+                code.writeline(f"{kernel_ref},")
+                code.writeline(f"out_shape={out_shape_expr},")
+                code.writeline("grid_spec=pltpu.PrefetchScalarGridSpec(")
+                with code.indent():
+                    code.writeline("num_scalar_prefetch=1,")
+                    code.writeline("grid=(_seq,),")
+                    code.writeline("in_specs=_in_specs,")
+                    code.writeline("out_specs=_out_specs,")
+                code.writeline("),")
+                if alias_map_parts:
+                    code.writeline(f"input_output_aliases={{ {alias_map_literal} }},")
+                if not self.is_tpu:
+                    code.writeline(f"interpret={ctx.interpret_literal},")
+
+            all_pallas_args = (
+                [indices]
+                + ["_table_3d"]
+                + pallas_call_other_args
+                + pallas_call_alias_args
+            )
+            code.writeline(f")({', '.join(all_pallas_args)})")
+
+            code.writeline(
+                "return tuple(r.reshape(s) for r, s in zip(_result, out_shapes))"
+            )
 
         self._codegen_main_entry(ctx, jit_wrapper_name)
-        return code.getvalue()
 
     def _codegen_imports(self, ctx: _CodegenContext) -> None:
         imports = """
@@ -2805,29 +4063,51 @@ def _codegen_imports(self, ctx: _CodegenContext) -> None:
 import jax
 import jax.numpy as jnp
 from jax.experimental import pallas as pl
+from torch.utils._ordered_set import OrderedSet
 from torch._inductor.runtime.runtime_utils import (
+    pallas_compute_tiling, pallas_make_block_spec, pallas_permute,
     pallas_gpu_align_output_specs, pallas_gpu_pad_inputs,
-    pallas_gpu_unpad_results, pallas_partial_reduce,
+    pallas_gpu_unpad_results,
+    pallas_ensure_nonzero_rank,
+    pallas_make_block_spec_non_tiled,
     torch_dtype_to_jax_runtime,
 )
 """
         if ctx.is_tpu:
             imports += "\nimport jax.export"
+            imports += "\nfrom jax.experimental.pallas import tpu as pltpu"
             imports += "\nfrom torch_tpu._internal.pallas import tpu_torch_pallas"
         elif not ctx.interpret_is_cpu:
             imports += "\nfrom jax.experimental.pallas import mosaic_gpu as plgpu"
+        if self.indirect_access and not ctx.is_tpu:
+            imports += (
+                "\nimport os as _os; _os.environ.setdefault('JAX_PLATFORMS', 'cpu')"
+            )
+            imports += "\nfrom jax.experimental.pallas import tpu as pltpu"
         ctx.code.splice(imports, strip=True)
 
-    def _codegen_iteration_vars(
-        self, kernel_body: IndentedBuffer, ctx: _CodegenContext
-    ) -> None:
-        # Generate iteration variables as jnp.arange arrays
-        # Skip on GPU - jnp.arange is not supported by Pallas Mosaic backend
-        if not (self.range_tree_nodes and not self.is_gpu and self.used_iter_vars):
-            return
+    def _get_iter_var_axis(self, var_sym: sympy.Symbol) -> int | None:
+        """Map an iteration variable to its output tensor axis index.
 
-        kernel_body.writeline("# Define iteration variables as JAX arrays")
+        Non-reduction variables map to axes 0, 1, 2, ... in order.
+        Reduction variables map to axes after all pointwise axes.
+        Returns None if the mapping cannot be determined.
+        """
+        pw_idx = 0
+        r_idx = 0
+        n_pw = sum(1 for _, e in self.range_tree_nodes.items() if not e.is_reduction)
+        for sym, entry in self.range_tree_nodes.items():
+            if sym == var_sym:
+                return pw_idx if not entry.is_reduction else n_pw + r_idx
+            if entry.is_reduction:
+                r_idx += 1
+            else:
+                pw_idx += 1
+        return None
 
+    def _get_reshape_target_shape_and_numel(
+        self,
+    ) -> tuple[tuple[int, ...] | None, int | None]:
         # Find reshape target: N-D shape whose numel matches an iteration
         # var. Try output first (repeat/upsample), then inputs (reductions).
         iter_lengths = OrderedSet(
@@ -2849,11 +4129,7 @@ def _get_nd_shape_if_matches(buf_name):
             numel = math.prod(shape)
             return (shape, numel) if numel in iter_lengths else (None, None)
 
-        candidate_buf_names = []
-        if ctx.output_params:
-            buf_name = ctx.output_buffer_lookup.get(ctx.output_params[0])
-            if buf_name:
-                candidate_buf_names.append(buf_name)
+        candidate_buf_names = self._output_buffer_names.copy()
         candidate_buf_names.extend(self.args.input_buffers)
 
         reshape_target_shape, reshape_target_numel = None, None
@@ -2863,6 +4139,41 @@ def _get_nd_shape_if_matches(buf_name):
                 reshape_target_shape, reshape_target_numel = result
                 break
 
+        return reshape_target_shape, reshape_target_numel
+
+    def _make_broadcasted_iteration_var_expr(
+        self, broadcast_vars: list[_BroadcastedIterVar], broadcast_idx: int
+    ) -> str:
+        bv = broadcast_vars[broadcast_idx]
+        length = bv.entry.length
+        renamed_length = self.rename_indexing(length)
+        length_str = self.kexpr(renamed_length)
+
+        num_broadcast_dims = len(broadcast_vars)
+        axis_idx = self._broadcast_axis_idx(
+            broadcast_vars, broadcast_idx, num_broadcast_dims
+        )
+        shape_parts = ["1"] * num_broadcast_dims
+        shape_parts[axis_idx] = length_str
+        shape_str = ", ".join(shape_parts)
+        arange = f"jnp.arange({length_str})"
+        reshaped = f"{arange}.reshape({shape_str})"
+        return reshaped
+
+    def _codegen_iteration_vars(
+        self, kernel_body: IndentedBuffer, ctx: _CodegenContext
+    ) -> None:
+        # Generate iteration variables as jnp.arange arrays
+        # Skip on GPU - jnp.arange is not supported by Pallas Mosaic backend
+        if not (self.range_tree_nodes and not self.is_gpu and self.used_iter_vars):
+            return
+
+        kernel_body.writeline("# Define iteration variables as JAX arrays")
+
+        reshape_target_shape, reshape_target_numel = (
+            self._get_reshape_target_shape_and_numel()
+        )
+
         var_items = list(self.range_tree_nodes.items())
 
         broadcast_vars = []
@@ -2872,7 +4183,9 @@ def _get_nd_shape_if_matches(buf_name):
             if length_val is not None and length_val == reshape_target_numel:
                 total_var_idx = idx
             else:
-                broadcast_vars.append((idx, var_sym, entry, length_val))
+                broadcast_vars.append(
+                    _BroadcastedIterVar(idx, var_sym, entry, length_val)
+                )
 
         num_broadcast_dims = len(broadcast_vars)
 
@@ -2892,24 +4205,14 @@ def _get_nd_shape_if_matches(buf_name):
                     and idx != total_var_idx
                 ):
                     broadcast_idx = next(
-                        (
-                            i
-                            for i, (vidx, _, _, _) in enumerate(broadcast_vars)
-                            if vidx == idx
-                        ),
+                        (i for i, v in enumerate(broadcast_vars) if v.idx == idx),
                         None,
                     )
                     if broadcast_idx is not None:
-                        axis_idx = self._broadcast_axis_idx(
-                            broadcast_vars, broadcast_idx, num_broadcast_dims
-                        )
-                        shape_parts = ["1"] * num_broadcast_dims
-                        shape_parts[axis_idx] = length_str
-                        shape_str = ", ".join(shape_parts)
-                        arange = f"jnp.arange({length_str})"
-                        kernel_body.writeline(
-                            f"{var_name} = {arange}.reshape({shape_str})"
+                        expr = self._make_broadcasted_iteration_var_expr(
+                            broadcast_vars, broadcast_idx
                         )
+                        kernel_body.writeline(f"{var_name} = {expr}")
                         continue
                 kernel_body.writeline(f"{var_name} = jnp.arange({length_str})")
                 continue
@@ -2924,22 +4227,45 @@ def _get_nd_shape_if_matches(buf_name):
                 kernel_body.writeline(f"{var_name} = {arange}.reshape({shape_str})")
             elif num_broadcast_dims > 1 and idx != total_var_idx:
                 broadcast_idx = next(
-                    i for i, (vidx, _, _, _) in enumerate(broadcast_vars) if vidx == idx
+                    i for i, v in enumerate(broadcast_vars) if v.idx == idx
                 )
-                axis_idx = self._broadcast_axis_idx(
-                    broadcast_vars, broadcast_idx, num_broadcast_dims
+                expr = self._make_broadcasted_iteration_var_expr(
+                    broadcast_vars, broadcast_idx
                 )
-                shape_parts = ["1"] * num_broadcast_dims
-                shape_parts[axis_idx] = length_str
-                shape_str = ", ".join(shape_parts)
-                arange = f"jnp.arange({length_str})"
-                kernel_body.writeline(f"{var_name} = {arange}.reshape({shape_str})")
+                kernel_body.writeline(f"{var_name} = {expr}")
             else:
-                kernel_body.writeline(f"{var_name} = jnp.arange({length_str})")
+                # Simple 1D arange — emit tile-relative form so tiling is safe.
+                # When grid=(1,), _pallas_tile[ax] == full length and
+                # pl.program_id(0) == 0, so this degenerates to jnp.arange(N).
+                # Only do this when the var actually appears in compute body
+                # (otherwise tiling is not blocked and the full arange is fine).
+                # Skip for scatter/index kernels where the iter var is used
+                # as a global index, not a data value.
+                compute_text = "\n".join(str(line) for line in self.compute._lines)
+                var_in_compute = var_name in compute_text
+                can_tile_relative = (
+                    var_in_compute
+                    and not self.is_gpu
+                    and not self.outputs_need_read
+                    and not self.has_flatten_indexing
+                    and not entry.is_reduction
+                )
+                axis_idx = (
+                    self._get_iter_var_axis(var_sym) if can_tile_relative else None
+                )
+                if axis_idx is not None:
+                    kernel_body.writeline(
+                        f"{var_name} = jnp.arange(_pallas_tile[{axis_idx}])"
+                        f" + pl.program_id(_pallas_ax2g.get({axis_idx}, 0))"
+                        f" * _pallas_tile[{axis_idx}]"
+                    )
+                    self.tile_relative_iter_vars.add(var_sym)
+                else:
+                    kernel_body.writeline(f"{var_name} = jnp.arange({length_str})")
 
     @staticmethod
     def _broadcast_axis_idx(
-        broadcast_vars: list[tuple[int, Any, Any, Any]],
+        broadcast_vars: list[_BroadcastedIterVar],
         broadcast_idx: int,
         num_broadcast_dims: int,
     ) -> int:
@@ -2947,10 +4273,10 @@ def _broadcast_axis_idx(
         # - Mixed: pointwise first, reduction last for output reshape
         # - Same-type: reverse order, first var innermost
         has_reduction_vars = any(
-            str(v).startswith("r") for _, v, _, _ in broadcast_vars
+            str(bv.var_sym).startswith("r") for bv in broadcast_vars
         )
         has_pointwise_vars = any(
-            not str(v).startswith("r") for _, v, _, _ in broadcast_vars
+            not str(bv.var_sym).startswith("r") for bv in broadcast_vars
         )
         is_mixed = has_reduction_vars and has_pointwise_vars
         if is_mixed:
@@ -3093,7 +4419,7 @@ def _codegen_jit_wrapper_legacy_gpu(
             code.writeline(f"_all_sizes.append({param}.size)")
         code.writeline("for shape in out_shapes:")
         code.writeline("    _all_sizes.append(math.prod(shape))")
-        code.writeline("_unique_sizes = set(_all_sizes)")
+        code.writeline("_unique_sizes = OrderedSet(_all_sizes)")
         code.writeline(
             "_can_pad = len(_unique_sizes) == 1 and all(s > 1 for s in _unique_sizes)"
         )
@@ -3155,6 +4481,86 @@ def _codegen_jit_wrapper_legacy_gpu(
             "        return pallas_gpu_unpad_results(_result, out_shapes, _is_scalar)"
         )
 
+    def _param_to_buf_name(self, param: str) -> str | None:
+        """Map a kernel parameter name back to its graph buffer name."""
+        for graph_name, inner_name in self.args.input_buffers.items():
+            if inner_name == param:
+                return graph_name
+        return None
+
+    def _codegen_tiled_specs(self, ctx: _CodegenContext) -> None:
+        """Generate tiled BlockSpec and grid variables for CPU/TPU.
+
+        Tiles the last 1–2 dimensions of each tensor, respecting TPU
+        alignment constraints (last dim multiple of 128, second-to-last
+        multiple of 8).  Lower-ndim inputs are right-aligned with the
+        reference output shape per numpy broadcast rules.
+        """
+        code = ctx.code
+        skip_n = self.tile_skip_last_n
+        has_transpose = "True" if self.tile_has_transpose else "False"
+        is_tpu_literal = "True" if ctx.is_tpu else "False"
+
+        # Collect per-input permutations for tiling alignment.
+        all_perms: list[tuple[int, ...] | None] = []
+        for p in ctx.kernel_input_params:
+            buf_name = self._param_to_buf_name(p)
+            all_perms.append(
+                self.permuted_input_buffers.get(buf_name) if buf_name else None
+            )
+
+        mgp = self._cpu_max_grid_product
+        mgp_arg = f", max_grid_product={mgp}" if mgp else ""
+        perms_arg = (
+            f", permutations={repr(all_perms)}"
+            if any(p is not None for p in all_perms)
+            else ""
+        )
+        code.writeline(
+            f"_tile, _grid, _ax2g = pallas_compute_tiling("
+            f"_pallas_out_shapes[0], "
+            f"transpose={has_transpose}, "
+            f"skip_last_n={skip_n}, exact_only=len(_pallas_out_shapes[0]) < 2, "
+            f"is_tpu={is_tpu_literal}"
+            f"{perms_arg}{mgp_arg})"
+        )
+        code.writeline("_ng = len(_grid)")
+        code.writeline("_ref = _pallas_out_shapes[0]")
+
+        code.writeline("out_specs_pallas = tuple(")
+        code.writeline(
+            "    pallas_make_block_spec(s, _ref, _tile, _ax2g, _ng, is_output=True)"
+        )
+        code.writeline("    for s in _pallas_out_shapes")
+        code.writeline(")")
+
+        self._codegen_strided_reshapes(code, ctx.kernel_input_params)
+
+        # Reshape collapsed inputs before building specs.
+        for param in ctx.kernel_input_params:
+            buf_name = self._param_to_buf_name(param)
+            cshape = self.collapsed_reshape_inputs.get(buf_name) if buf_name else None
+            if cshape is not None:
+                code.writeline(f"{param} = {param}.reshape({cshape})")
+
+        # Build input BlockSpecs (with per-input permutation when needed).
+        input_list = ", ".join(ctx.kernel_input_params)
+        if any(p is not None for p in all_perms):
+            perm_list = ", ".join(repr(p) for p in all_perms)
+            code.writeline(f"_perm_flags = [{perm_list}]")
+            code.writeline("in_specs_pallas = tuple(")
+            code.writeline(
+                f"    pallas_make_block_spec(i.shape, _ref, _tile, _ax2g, _ng, permutation=p)"
+                f" for i, p in zip([{input_list}], _perm_flags)"
+            )
+            code.writeline(")")
+        else:
+            code.writeline("in_specs_pallas = tuple(")
+            code.writeline(
+                f"    pallas_make_block_spec(i.shape, _ref, _tile, _ax2g, _ng) for i in [{input_list}]"
+            )
+            code.writeline(")")
+
     def _codegen_jit_wrapper_cpu_tpu(
         self,
         ctx: _CodegenContext,
@@ -3163,13 +4569,14 @@ def _codegen_jit_wrapper_cpu_tpu(
         alias_map_literal: str,
     ) -> None:
         code = ctx.code
-        code.writeline("return pl.pallas_call(")
+        grid_expr = "_grid" if self.tile_cpu_tpu else "(1,)"
+        code.writeline("_result = pl.pallas_call(")
         code.writeline("    " + kernel_arg)
         code.writeline("    out_shape=out_shapes_pallas,")
         code.writeline("    out_specs=out_specs_pallas,")
         code.writeline("    in_specs=in_specs_pallas,")
         code.writeline(f"    interpret={ctx.interpret_literal},")
-        code.writeline("    grid=(1,),")
+        code.writeline(f"    grid={grid_expr},")
         code.writeline(
             f"    input_output_aliases={{ {alias_map_literal} }},"
             if alias_pairs
@@ -3177,8 +4584,19 @@ def _codegen_jit_wrapper_cpu_tpu(
         )
         code.writeline(")(")
         if ctx.kernel_input_params:
-            code.writeline(f"    {', '.join(ctx.kernel_input_params)},")
+            kernel_input_params_nonzero_rank = [
+                f"pallas_ensure_nonzero_rank({p})" for p in ctx.kernel_input_params
+            ]
+            code.writeline(f"    {', '.join(kernel_input_params_nonzero_rank)},")
         code.writeline(")")
+        # Reshape results back to original shapes (restores 0-d from promoted (1,))
+        code.writeline("if isinstance(_result, tuple):")
+        code.writeline(
+            "    _result = tuple(r.reshape(s) for r, s in zip(_result, out_shapes))"
+        )
+        code.writeline("else:")
+        code.writeline("    _result = _result.reshape(out_shapes[0])")
+        code.writeline("return _result")
 
     def _codegen_main_entry(self, ctx: _CodegenContext, jit_wrapper_name: str) -> None:
         if ctx.is_tpu:
@@ -3197,9 +4615,21 @@ def _codegen_main_entry_tpu(
             f"def {main_name}({', '.join(ctx.full_kernel_params)}, stream=None):"
         )
         with code.indent():
-            code.writeline("jax.config.update('jax_enable_x64', True)")
+            # `jax_enable_x64` is per-process. The CPU path sets it to True,
+            # so running both CPU and TPU tests in one process can cause
+            # x64-related TPU crashes if we do not explicitly set it to
+            # False here.
+            code.writeline("jax.config.update('jax_enable_x64', False)")
             code.writeline("jax.clear_caches()")
 
+            # Convert int64 inputs to int32 (TPU doesn't support int64)
+            all_input_params = list(ctx.alias_params) + list(ctx.pointer_tail)
+            for param_name in all_input_params:
+                code.writeline(
+                    f"{param_name} = {param_name}.to(torch.int32) "
+                    f"if {param_name}.dtype == torch.int64 else {param_name}"
+                )
+
             # Build JAX placeholders for all inputs
             code.writeline("# Build JAX placeholders for export tracing")
             all_jax_input_names = []
@@ -3244,18 +4674,45 @@ def _codegen_main_entry_tpu(
             )
 
             # Register and call via tpu_torch_pallas
+            # Include all output and input shapes in the key to avoid stale
+            # cache hits when the same kernel name is compiled with different
+            # input/output ranks (e.g. broadcasting vs non-broadcasting calls).
+            shape_key_parts = []
+            for p in ctx.output_params:
+                shape_key_parts.append(f"'_'.join(str(s) for s in {p}.shape)")
+            output_key_expr = (
+                " + 'x' + ".join(shape_key_parts) if shape_key_parts else "''"
+            )
+            input_key_parts = []
+            for p in ctx.kernel_input_params:
+                input_key_parts.append(f"'_'.join(str(s) for s in {p}.shape)")
+            input_key_expr = (
+                " + 'x' + ".join(input_key_parts) if input_key_parts else "''"
+            )
             code.writeline(
-                f"kernel_key = '{kernel_name_str}_' + "
-                f"'_'.join(str(s) for s in {ctx.output_params[0]}.shape)"
+                f"kernel_key = '{kernel_name_str}_out_' + "
+                f"{output_key_expr}"
+                f" + '_in_' + {input_key_expr}"
             )
+
             code.writeline(
                 f"if not tpu_torch_pallas.lookup_custom_kernel('{kernel_name_str}', kernel_key):"
             )
             with code.indent():
-                code.writeline(
-                    f"tpu_torch_pallas.register_custom_kernel("
-                    f"'{kernel_name_str}', kernel_key, exported.mlir_module_serialized)"
-                )
+                code.writeline("try:")
+                with code.indent():
+                    code.writeline(
+                        f"tpu_torch_pallas.register_custom_kernel("
+                        f"'{kernel_name_str}', kernel_key, "
+                        f"serialized_mlir_module=exported.mlir_module_serialized)"
+                    )
+                code.writeline("except TypeError:")
+                with code.indent():
+                    code.writeline(
+                        f"tpu_torch_pallas.register_custom_kernel("
+                        f"'{kernel_name_str}', kernel_key, "
+                        f"exported.mlir_module_serialized)"
+                    )
 
             # Build input tensor list (all non-size-var inputs)
             input_tensor_names = list(ctx.alias_params) + list(ctx.pointer_tail)
@@ -3278,16 +4735,30 @@ def _codegen_main_entry_tpu(
                     )
             code.writeline("]")
 
-            code.writeline(
-                f"results = tpu_torch_pallas.call_custom_kernel("
-                f"input_tensors, output_shape_tensors, "
-                f"'{kernel_name_str}', kernel_key)"
-            )
+            # Build input_output_aliases for zero-copy donation
+            if ctx.alias_pairs:
+                alias_map_str = ", ".join(f"{i}: {o}" for (i, o) in ctx.alias_pairs)
+                code.writeline(f"_input_output_aliases = {{ {alias_map_str} }}")
+            else:
+                code.writeline("_input_output_aliases = {}")
 
-            # Copy results to output buffers
-            for idx in ctx.copy_output_indices:
-                out_name = ctx.output_params[idx]
-                code.writeline(f"{out_name}.copy_(results[{idx}])")
+            code.writeline("try:")
+            with code.indent():
+                code.writeline(
+                    f"tpu_torch_pallas.call_custom_kernel("
+                    f"'{kernel_name_str}', kernel_key, "
+                    f"inputs=input_tensors, "
+                    f"output_shapes=output_shape_tensors, "
+                    f"input_output_aliases=_input_output_aliases)"
+                )
+            code.writeline("except TypeError:")
+            with code.indent():
+                code.writeline(
+                    f"tpu_torch_pallas.call_custom_kernel("
+                    f"input_tensors, output_shape_tensors, "
+                    f"'{kernel_name_str}', kernel_key, "
+                    f"_input_output_aliases)"
+                )
 
     def _codegen_main_entry_default(
         self, ctx: _CodegenContext, jit_wrapper_name: str
@@ -3299,8 +4770,11 @@ def _codegen_main_entry_default(
             f"def {main_name}({', '.join(ctx.full_kernel_params)}, stream=None):"
         )
         with code.indent():
-            code.writeline("# Enable JAX x64 mode for float64/int64 support")
             code.writeline("jax.config.update('jax_enable_x64', True)")
+            if ctx.interpret_is_cpu:
+                code.writeline(
+                    "jax.config.update('jax_default_device', jax.devices('cpu')[0])"
+                )
             code.writeline("jax.clear_caches()")
             if ctx.alias_params:
                 code.writeline("# Convert Torch -> JAX for donated outputs")
@@ -3373,7 +4847,7 @@ def _emit_torch_to_jax(
         suffix = ".detach().contiguous()" if contiguous else ".detach()"
         code.writeline(f"{var_name}_jax = jax.dlpack.from_dlpack({var_name}{suffix})")
 
-    def call_kernel(self, name: str, node: Optional[IRNode] = None) -> None:  # type: ignore[override]
+    def call_kernel(self, name: str, node: IRNode | None = None) -> None:  # type: ignore[override]
         """Generate the Python code that calls this Pallas kernel."""
         wrapper = V.graph.wrapper_code
         arg_defs, call_args, _, _ = self.args.python_argdefs()
@@ -3403,6 +4877,30 @@ def get_backend_features(cls, device: torch.device) -> OrderedSet[BackendFeature
         # without requiring split reductions
         return OrderedSet([BackendFeature.REDUCE_TO_SINGLE_ELEMENT])
 
+    def can_fuse(self, node1, node2):  # type: ignore[override]
+        if not super().can_fuse(node1, node2):
+            return False
+        # Pallas partial reductions use keepdims, so fusing two reductions
+        # that read the same buffer with different index patterns produces
+        # intermediates with incompatible shapes (e.g. (1,8) + (8,1) = (8,8)
+        # instead of (8,)).  Prevent this by rejecting fusion when the read
+        # indices differ.
+        if node1.is_reduction() and node2.is_reduction():
+            from torch._inductor.dependencies import MemoryDep
+
+            reads1 = {}
+            for dep in node1.read_writes.reads:
+                if isinstance(dep, MemoryDep):
+                    reads1[dep.name] = dep.index
+            for dep in node2.read_writes.reads:
+                if isinstance(dep, MemoryDep) and dep.name in reads1:
+                    if reads1[dep.name] != dep.index:
+                        return False
+        return True
+
+    can_fuse_vertical = can_fuse  # type: ignore[assignment]
+    can_fuse_horizontal = can_fuse  # type: ignore[assignment]
+
     def define_kernel(
         self,
         src_code: str,
diff --git a/torch/_inductor/codegen/python_wrapper_mtia.py b/torch/_inductor/codegen/python_wrapper_mtia.py
index 00833e1de702c..4a0309f2088ae 100644
--- a/torch/_inductor/codegen/python_wrapper_mtia.py
+++ b/torch/_inductor/codegen/python_wrapper_mtia.py
@@ -1,4 +1,3 @@
-from typing import Optional
 from typing_extensions import override
 
 from torch._inductor import ir
@@ -22,9 +21,9 @@ def write_header(self) -> None:
     @staticmethod
     def create(
         is_subgraph: bool,
-        subgraph_name: Optional[str],
-        parent_wrapper: Optional[PythonWrapperCodegen],
-        partition_signatures: Optional[ir.GraphPartitionSignature] = None,
+        subgraph_name: str | None,
+        parent_wrapper: PythonWrapperCodegen | None,
+        partition_signatures: ir.GraphPartitionSignature | None = None,
     ) -> PythonWrapperCodegen:
         if is_subgraph:
             # Delegate to the parent class to handle the case of subgraph
diff --git a/torch/_inductor/codegen/rocm/ck_universal_gemm_template.py b/torch/_inductor/codegen/rocm/ck_universal_gemm_template.py
index e9f8ff54f9f45..57e5053acd607 100644
--- a/torch/_inductor/codegen/rocm/ck_universal_gemm_template.py
+++ b/torch/_inductor/codegen/rocm/ck_universal_gemm_template.py
@@ -4,7 +4,6 @@
 import math
 import random
 from collections import namedtuple
-from typing import Optional
 
 import sympy
 
@@ -317,7 +316,7 @@ def __init__(
         layout: Layout,
         alpha: float,
         beta: float,
-        input_reorder: Optional[list[int]] = None,
+        input_reorder: list[int] | None = None,
     ) -> None:
         is_batched = len(layout.size) == 3
         name = "ck_batched_gemm_template" if is_batched else "ck_gemm_template"
diff --git a/torch/_inductor/codegen/rocm/compile_command.py b/torch/_inductor/codegen/rocm/compile_command.py
index aa935b14af23c..2aa4494e0aa32 100644
--- a/torch/_inductor/codegen/rocm/compile_command.py
+++ b/torch/_inductor/codegen/rocm/compile_command.py
@@ -1,7 +1,6 @@
 # mypy: allow-untyped-defs
 import logging
 import os
-from typing import Optional
 
 from torch._inductor import config
 from torch._inductor.utils import is_linux, try_import_ck_lib
@@ -81,7 +80,7 @@ def _rocm_compiler_options() -> list[str]:
         config.rocm.compile_opt_level,
         "-x",
         "hip",
-        "-std=c++17",
+        "-std=c++20",
         *gpu_arch_flags,
         "-fno-gpu-rdc",
         "-fPIC",
@@ -106,7 +105,7 @@ def _rocm_compiler_options() -> list[str]:
     return opts
 
 
-def rocm_compiler() -> Optional[str]:
+def rocm_compiler() -> str | None:
     if is_linux():
         if config.rocm.rocm_home:
             return os.path.realpath(
@@ -128,7 +127,7 @@ def rocm_compile_command(
     src_files: list[str],
     dst_file: str,
     dst_file_ext: str,
-    extra_args: Optional[list[str]] = None,
+    extra_args: list[str] | None = None,
 ) -> str:
     include_paths = _rocm_include_paths(dst_file_ext)
     lib_options = _rocm_lib_options(dst_file_ext)
diff --git a/torch/_inductor/codegen/rocm/rocm_benchmark_request.py b/torch/_inductor/codegen/rocm/rocm_benchmark_request.py
index c5a87bef820df..b07f33654ca27 100644
--- a/torch/_inductor/codegen/rocm/rocm_benchmark_request.py
+++ b/torch/_inductor/codegen/rocm/rocm_benchmark_request.py
@@ -4,7 +4,7 @@
 import functools
 import logging
 from ctypes import byref, c_int, c_size_t, c_void_p
-from typing import Any, Optional, TYPE_CHECKING, Union
+from typing import Any, TYPE_CHECKING
 
 import torch
 from torch._inductor import config
@@ -30,16 +30,16 @@ class ROCmBenchmarkRequest(GPUDeviceBenchmarkMixin, BenchmarkRequest):
     def __init__(
         self,
         kernel_name: str,
-        input_tensor_meta: Union[TensorMeta, list[TensorMeta]],
-        output_tensor_meta: Union[TensorMeta, list[TensorMeta]],
+        input_tensor_meta: TensorMeta | list[TensorMeta],
+        output_tensor_meta: TensorMeta | list[TensorMeta],
         extra_args: Iterable[Any],
         source_code: str,
     ) -> None:
         super().__init__(kernel_name, input_tensor_meta, output_tensor_meta, extra_args)
         self.source_code = source_code
         self.workspace_size: int = 0
-        self.workspace: Optional[torch.Tensor] = None
-        self.DLL: Optional[DLLWrapper] = None
+        self.workspace: torch.Tensor | None = None
+        self.DLL: DLLWrapper | None = None
         self._workspace_size_updated = False
         self.hash_key: str = ""
         self.source_file: str = ""
diff --git a/torch/_inductor/codegen/rocm/rocm_cpp_scheduling.py b/torch/_inductor/codegen/rocm/rocm_cpp_scheduling.py
index ec58e458df6b1..eb8769de8db79 100644
--- a/torch/_inductor/codegen/rocm/rocm_cpp_scheduling.py
+++ b/torch/_inductor/codegen/rocm/rocm_cpp_scheduling.py
@@ -1,7 +1,7 @@
 # mypy: allow-untyped-defs
 import logging
 from collections.abc import Sequence
-from typing import cast
+from typing import cast, TypeGuard
 
 from ... import config
 from ...codecache import code_hash, get_path
@@ -28,7 +28,7 @@ def group_fn(self, sizes):
         return tuple(V.graph.sizevars.simplify(sympy_product(s)) for s in sizes)
 
     @staticmethod
-    def is_rocm_cpp_template(node: BaseSchedulerNode) -> bool:
+    def is_rocm_cpp_template(node: BaseSchedulerNode) -> TypeGuard[SchedulerNode]:
         return isinstance(node, SchedulerNode) and isinstance(
             node.node, ROCmTemplateBuffer
         )
@@ -82,7 +82,6 @@ def codegen_template(
         assert self.is_rocm_cpp_template(template_node), (
             "Template node passed to ROCmScheduler.codegen_template must be a SchedulerNode that wraps a ROCmTemplateBuffer"
         )
-        template_node = cast(SchedulerNode, template_node)
         _, (_numel, rnumel) = template_node.group
         assert rnumel == 1
         ctb: ROCmTemplateBuffer = cast(ROCmTemplateBuffer, template_node.node)
diff --git a/torch/_inductor/codegen/rocm/rocm_kernel.py b/torch/_inductor/codegen/rocm/rocm_kernel.py
index 598a1c189ee6d..99373589c227a 100644
--- a/torch/_inductor/codegen/rocm/rocm_kernel.py
+++ b/torch/_inductor/codegen/rocm/rocm_kernel.py
@@ -1,7 +1,7 @@
 # mypy: allow-untyped-defs
 import logging
 from collections.abc import Callable, Sequence
-from typing import Any, Optional, TYPE_CHECKING, Union
+from typing import Any, TYPE_CHECKING
 
 import torch._inductor.config as config
 from torch._inductor.codegen.cpp_wrapper_cpu import CppWrapperCpu
@@ -71,7 +71,7 @@ def def_kernel(
         outputs: list[IRNode],
         size_args: list[str],
         names_str: str = "",
-        input_reorder: Optional[list[int]] = None,
+        input_reorder: list[int] | None = None,
     ) -> str:
         """
         Hook called from template code to generate function definition and
@@ -228,13 +228,11 @@ def __init__(
         input_nodes: list[Buffer],
         layout: Layout,
         make_kernel_render: Callable[
-            [ROCmTemplateBuffer, Optional[Sequence[IRNode]]], str
+            [ROCmTemplateBuffer, Sequence[IRNode] | None], str
         ],
         bmreq: ROCmBenchmarkRequest,
         template: "ROCmTemplate",  # type: ignore[name-defined]
-        info_kwargs: Optional[
-            dict[str, Union[PrimitiveInfoType, list[PrimitiveInfoType]]]
-        ],  # type: ignore[type-arg]
+        info_kwargs: dict[str, PrimitiveInfoType | list[PrimitiveInfoType]] | None,  # type: ignore[type-arg]
     ) -> None:
         super().__init__(name, input_nodes, layout, description="")
         self.category = category
@@ -268,7 +266,7 @@ def hash_key(self) -> str:
             ]
         )
 
-    def info_dict(self) -> dict[str, Union[PrimitiveInfoType, list[PrimitiveInfoType]]]:
+    def info_dict(self) -> dict[str, PrimitiveInfoType | list[PrimitiveInfoType]]:
         """Information returned here is logged to the autotune log file when that is enabled."""
         return {
             "backend": "ROCm",
diff --git a/torch/_inductor/codegen/rocm/rocm_template.py b/torch/_inductor/codegen/rocm/rocm_template.py
index bfeb03eabc72d..ac5f514fa6460 100644
--- a/torch/_inductor/codegen/rocm/rocm_template.py
+++ b/torch/_inductor/codegen/rocm/rocm_template.py
@@ -4,7 +4,7 @@
 import logging
 from collections.abc import Sequence
 from dataclasses import dataclass
-from typing import Any, Optional
+from typing import Any
 from unittest.mock import patch
 
 from ...autotune_process import TensorMeta
@@ -37,7 +37,7 @@ def __init__(
         name: str,
         input_nodes: list[Buffer],
         layout: Layout,
-        input_reorder: Optional[list[int]] = None,
+        input_reorder: list[int] | None = None,
     ) -> None:
         """
 
@@ -106,9 +106,8 @@ def generate(  # type: ignore[override]
         size_args = (
             self.size_args() if hasattr(self, "size_args") else ()
         )  # subclass should define def size_args()
-        size_args_ints = [
-            V.graph.sizevars.size_hint(arg) for arg in size_args
-        ]  # resolve to ints for benchmarking
+        # Resolve symbolic sizes to concrete ints for benchmarking only.
+        size_args_ints = list(V.graph.sizevars.optimization_hints(size_args))
         # The runtime args come right after the size args
         runtime_args = self.get_runtime_arg_values(**kwargs)
         extra_args = size_args_ints + runtime_args
@@ -122,7 +121,7 @@ def generate(  # type: ignore[override]
 
         def make_kernel_render(
             template_node: ROCmTemplateBuffer,
-            epilogue_nodes: Optional[Sequence[IRNode]] = None,
+            epilogue_nodes: Sequence[IRNode] | None = None,
         ):
             kernel = ROCmTemplateKernel(
                 kernel_name="KERNEL_NAME",
diff --git a/torch/_inductor/codegen/segmented_tree.py b/torch/_inductor/codegen/segmented_tree.py
index d6e5c86d18109..40d132d66ea23 100644
--- a/torch/_inductor/codegen/segmented_tree.py
+++ b/torch/_inductor/codegen/segmented_tree.py
@@ -1,11 +1,11 @@
 from collections.abc import Callable
-from typing import Generic, Optional, TypeVar
+from typing import Generic, TypeVar
 
 
 T = TypeVar("T")
 
 
-def _value_or(opt: Optional[T], default: T) -> T:
+def _value_or(opt: T | None, default: T) -> T:
     return opt if opt is not None else default
 
 
@@ -55,7 +55,7 @@ def __init__(
         # we receive an interval query that neither fully
         # contains the node nor fully doesn't contain the
         # node
-        self.lazy: list[Optional[T]] = [None] * self.size
+        self.lazy: list[T | None] = [None] * self.size
 
         # Build the tree
         self._build(values, 1, 0, self.n - 1)
diff --git a/torch/_inductor/codegen/simd.py b/torch/_inductor/codegen/simd.py
index 6c25a38aef418..406d59ae8ec2c 100644
--- a/torch/_inductor/codegen/simd.py
+++ b/torch/_inductor/codegen/simd.py
@@ -11,7 +11,7 @@
 import operator
 import textwrap
 from collections import Counter
-from typing import Any, Generic, NamedTuple, Optional, TYPE_CHECKING, Union
+from typing import Any, Generic, NamedTuple, TYPE_CHECKING
 from typing_extensions import TypeVar
 
 import sympy
@@ -165,11 +165,11 @@ def __init__(
         prefix: str,
         index: int,
         kernel: SIMDKernel,
-        pid_cache: Optional[dict[str, str]] = None,
+        pid_cache: dict[str, str] | None = None,
         *,
         is_loop: bool,
-        tensor_dim: Optional[int],
-        grid_dim: Optional[int],
+        tensor_dim: int | None,
+        grid_dim: int | None,
         has_zdim: bool,
     ) -> None:
         if pid_cache is None:
@@ -352,7 +352,7 @@ def __eq__(self, other: object) -> bool:
         return self.name == other.name
 
 
-def constant_repr(value: Union[int, float]) -> str:
+def constant_repr(value: int | float) -> str:
     if value == float("inf"):
         return 'float("inf")'
     elif value == float("-inf"):
@@ -400,10 +400,10 @@ def __init__(
         self,
         tiling: dict[str, sympy.Expr],
         features: SIMDKernelFeatures,
-        pid_cache: Optional[dict[str, str]] = None,
-        override_persistent_reduction: Optional[bool] = None,
-        override_cooperative_reduction: Optional[bool] = None,
-        tiling_scores: Optional[dict[str, sympy.Expr]] = None,
+        pid_cache: dict[str, str] | None = None,
+        override_persistent_reduction: bool | None = None,
+        override_cooperative_reduction: bool | None = None,
+        tiling_scores: dict[str, sympy.Expr] | None = None,
         mix_order_reduction: bool = False,
     ) -> None:
         if pid_cache is None:
@@ -425,7 +425,7 @@ def __init__(
             if override_cooperative_reduction is not None
             else self.should_use_cooperative_reduction()
         )
-        self.tiling_scores: Optional[dict[str, sympy.Expr]] = tiling_scores
+        self.tiling_scores: dict[str, sympy.Expr] | None = tiling_scores
         self.tiling: dict[str, sympy.Expr] = tiling
         self.persistent_reduction: bool = (
             override_persistent_reduction
@@ -434,7 +434,7 @@ def __init__(
         )
         self.mix_order_reduction: bool = mix_order_reduction
         self.no_x_dim = self.want_no_x_dim()
-        self.code_hash: Optional[str] = None
+        self.code_hash: str | None = None
         # Info to enable multiple store_output calls for epilogue subtiling
         self.store_output_ctr = itertools.count()
         self.is_native_matmul = False
@@ -463,7 +463,7 @@ def simplify_indexing(index: sympy.Expr):
         self.rsplit_size = 0
         self.saved_partial_accumulate: list[PartialAccumulate] = []
 
-    def codegen_template_override(
+    def codegen_template_body(
         self,
         scheduling,
         template_node,
@@ -472,14 +472,26 @@ def codegen_template_override(
         buf_name_to_prologue_group,
         prologue_preserves_zero_mask_fn,
         render,
-        only_gen_src_code: bool,
-    ) -> str | None:
-        """Override template codegen. Return None to use default flow.
+    ) -> str:
+        """Generate template source code with fused prologues and epilogues.
 
-        External template handlers (e.g. Helion) can override this method
-        to implement custom code generation.
+        Subclasses override this to implement custom code generation.
+        The default implementation raises NotImplementedError — the actual
+        standard path lives in ``TritonTemplateKernel.codegen_template_body``.
         """
-        return None
+        raise NotImplementedError
+
+    def get_unfused_epilogues(self) -> list[Any]:
+        """Return epilogue nodes that were not fused into the kernel.
+
+        These nodes need separate codegen (via ``call_kernel``) and must
+        be excluded from ``mark_run`` in ``_codegen_single_template``.
+
+        The standard path fuses all epilogues, so this returns ``[]``.
+        ``ExternalTritonTemplateKernel`` overrides this for epilogues that
+        don't read exactly one template output and cannot be fused.
+        """
+        return []
 
     def _get_store_output_subgraph_name(self, i: int) -> str:
         return f"<STORE_OUTPUT_{i}>"
@@ -509,7 +521,7 @@ def want_no_x_dim(self) -> bool:
 
     def construct_range_trees(
         self,
-        pid_cache: Optional[dict[str, str]],
+        pid_cache: dict[str, str] | None,
         inside_reduction: bool,
         is_reduction: bool,
         numels: dict[str, sympy.Expr],
@@ -790,7 +802,10 @@ def getter(flat_vars: list[sympy.Expr]) -> sympy.Expr:
                     if not sv.statically_known_multiple_of(
                         size, remaining[current_group] * remaining[current_group + 1]
                     ):
-                        raise CantSplit
+                        raise CantSplit(
+                            size,
+                            remaining[current_group] * remaining[current_group + 1],
+                        )
 
                     size1 = remaining[current_group]
                     size2 = remaining[current_group + 1]
@@ -844,16 +859,17 @@ def getter(flat_vars: list[sympy.Expr]) -> sympy.Expr:
                         )
                     )
                 else:
-                    if current_group < len(remaining):
-                        return_getters.append(
-                            # pyrefly: ignore [bad-argument-type]
-                            operator.itemgetter(add_range(current_group, size))
-                        )
+                    if current_group >= len(remaining):
+                        raise CantSplit(size, 0)
+                    return_getters.append(
+                        # pyrefly: ignore [bad-argument-type]
+                        operator.itemgetter(add_range(current_group, size))
+                    )
             return_getters_groups.append(return_getters)
 
-        assert all(V.graph.sizevars.size_hint(s) == 1 for s in remaining), (
-            f"failed to set ranges {remaining} {lengths}"
-        )
+        assert all(
+            V.graph.sizevars.guarding_hint_or_throw(s) == 1 for s in remaining
+        ), f"failed to set ranges {remaining} {lengths}"
         # pyrefly: ignore [bad-return]
         return new_ranges, return_getters_groups
 
@@ -1063,14 +1079,12 @@ def deallocate_workspaces(self):
             wrapper.generate_workspace_deallocation(ws)
 
     def call_kernel(
-        self, name: str, node: Optional[IRNode] = None, deallocate_ws: bool = True
+        self, name: str, node: IRNode | None = None, deallocate_ws: bool = True
     ) -> None:
         raise NotImplementedError("NYI: call_kernel")
 
     @contextlib.contextmanager
-    def mask_loads(
-        self, mask: Union[str, OpsWrapper], value: Union[int, float]
-    ) -> Iterator[str]:
+    def mask_loads(self, mask: str | OpsWrapper, value: int | float) -> Iterator[str]:
         """Context manager to add an additional mask to tl.load/store"""
         prior = self._load_mask
         prior_val = self._load_other
@@ -1117,7 +1131,7 @@ def _map_tuple_or_scalar(fn, value):
             return tuple(map(fn, value))
         return fn(value)
 
-    def estimate_flops(self) -> Optional[int]:
+    def estimate_flops(self) -> int | None:
         flops = [
             node.estimate_flops()
             for node in NodeScheduleMarker.only_nodes(self.features.node_schedule)
@@ -1474,7 +1488,7 @@ def generate_node_schedule(self, nodes, numel, rnumel):
         # reduction loop has ended
         not_ready_yet_nodes: OrderedSet[str] = OrderedSet()
         current_loop_buffer_usage: OrderedSet[str] = OrderedSet()
-        maybe_split_index: Optional[int] = None
+        maybe_split_index: int | None = None
 
         def fits_in_main_body(n):
             _, (node_numel, node_rnumel) = n.group
@@ -1637,8 +1651,9 @@ def _generate_kernel_code_for_mix_order_reduction(
             src_code = src_code.replace(str(Placeholder.KERNEL_NAME), "triton_")
         return kernel, ws_name, src_code
 
+    # pyrefly: ignore [bad-override]
     def benchmark_codegened_module(
-        self, mod, n_spills_threshold=8, node_names: Optional[OrderedSet[str]] = None
+        self, mod, n_spills_threshold=8, node_names: OrderedSet[str] | None = None
     ) -> tuple[float, str]:
         raise NotImplementedError
 
@@ -1655,8 +1670,10 @@ def _pick_split_size():
             num_sm = device_prop.multi_processor_count
             estimated_num_splits = num_sm * 8
 
-            # split_size is decided based on hint
-            numel_hint = V.graph.sizevars.size_hint(numel)
+            # split_size is decided based on hint.
+            # optimization_hint is fine here: the result is clamped to [16, 128],
+            # so any fallback value still produces a valid split size.
+            numel_hint = V.graph.sizevars.optimization_hint(numel)
             split_size = max(last_power_of_2(numel_hint // estimated_num_splits), 16)
             split_size = min(split_size, 128)
             return split_size
@@ -1767,7 +1784,7 @@ def _bench(candidate_split_size):
         for idx, partial_accum in enumerate(kernel.saved_partial_accumulate):
             buffer_name = partial_accum.buffer_name
 
-            stride_str = f"{nsplit} * {rnumel}"
+            stride_str = f"({nsplit}) * ({rnumel})"
             start = f"{idx} * {stride_str}"
             end = f"({idx} + 1) * {stride_str}"
             reduction_type2op = {
@@ -1778,7 +1795,13 @@ def _bench(candidate_split_size):
                 partial_accum.reduction_type, partial_accum.reduction_type
             )
 
-            final_reduce = f"{buffer_name} = {ws_name}[{start} : {end}].view({nsplit}, {rnumel}).{opname}(dim=0)"
+            # Check if the original reduction used keepdim=True by comparing dimensions.
+            # Without keepdim, reduction produces [rnumel]; with keepdim, [1, rnumel].
+            buffer = V.graph.get_buffer(buffer_name)
+            keepdim = buffer is not None and len(buffer.get_layout().size) > 1
+
+            final_reduce = f"{buffer_name} = {ws_name}[{start} : {end}].view({nsplit}, {rnumel}).{opname}(dim=0, keepdim={keepdim})"
+
             # The workspace tensor is in torch.float, need a cast if the buffer is
             # not.
             if (buffer_dtype := V.graph.get_dtype(buffer_name)) != torch.float:
@@ -1798,7 +1821,7 @@ def _bench(candidate_split_size):
     def _codegen_nodes(
         self,
         nodes: Sequence[scheduler.SchedulerNode],
-        coalesce_analysis: Optional[CoalesceVarAnalysis] = None,
+        coalesce_analysis: CoalesceVarAnalysis | None = None,
     ):
         assert self.scheduler
         nodes = [
@@ -1816,7 +1839,7 @@ def _codegen_nodes(
         )
 
     def codegen_node(
-        self, node: Union[scheduler.FusedSchedulerNode, scheduler.SchedulerNode]
+        self, node: scheduler.FusedSchedulerNode | scheduler.SchedulerNode
     ):
         """
         Given a set of pre-fused nodes, generate a Triton kernel.
@@ -1843,9 +1866,7 @@ def codegen_node(
     @staticmethod
     def can_use_32bit_indexing(
         numel: sympy.Expr,
-        buffers: Iterable[
-            Union[ir.Buffer, ir.TensorBox, ir.TorchBindObject, ir.IRNode]
-        ],
+        buffers: Iterable[ir.Buffer | ir.TensorBox | ir.TorchBindObject | ir.IRNode],
     ) -> bool:
         int_max = torch.iinfo(torch.int32).max
 
@@ -1896,6 +1917,20 @@ def process_kernel(
         V.graph.removed_buffers |= kernel.removed_buffers
         V.graph.inplaced_to_remove |= kernel.inplaced_to_remove
 
+    def _collect_config_patches(self, node_schedule: list[Any]) -> dict[str, Any]:
+        """Collect and merge config_patches from all operations in the node schedule.
+
+        This enables scoped config (e.g., coordinate_descent_tuning) for kernels
+        that contain decomposition operations.
+        """
+        merged_patches: dict[str, Any] = {}
+        for node in node_schedule:
+            if isinstance(node, BaseSchedulerNode) and node.node is not None:
+                patches = node.node.get_config_patches()
+                if patches:
+                    merged_patches.update(patches)
+        return merged_patches
+
     def codegen_node_schedule(self, kernel_features: SIMDKernelFeatures):
         """
         Generate code for nodes in kernel_features
@@ -1916,8 +1951,13 @@ def codegen_node_schedule(self, kernel_features: SIMDKernelFeatures):
         for kernel in kernels:
             self.codegen_node_schedule_with_kernel(node_schedule, kernel)
         MultiKernel.merge_workspaces_inplace(kernels)
+
+        # Collect config_patches from operations (e.g., decomposition ops with
+        # coordinate_descent_tuning) and apply during kernel codegen
+        config_patches = self._collect_config_patches(node_schedule)
+
         for kernel in kernels:
-            with V.set_kernel_handler(kernel):
+            with V.set_kernel_handler(kernel), config.patch(**config_patches):
                 src_code = kernel.codegen_kernel()
             kernel_name = self.define_kernel(src_code, node_schedule, kernel)
             log.debug("Generating kernel code with kernel_name: %s", kernel_name)
@@ -1925,7 +1965,7 @@ def codegen_node_schedule(self, kernel_features: SIMDKernelFeatures):
             kernel.code_hash = code_hash(src_code)
         del kernel
 
-        final_kernel: Union[SIMDKernel, MultiKernel]
+        final_kernel: SIMDKernel | MultiKernel
         if len(kernels) > 1:
             final_kernel = MultiKernel(kernels)
         else:
@@ -2052,8 +2092,15 @@ def _codegen_single_template(
         # all prologue groups should have finalized with use in template
         assert len(prologue_group) == 0
 
-        # External template handlers (e.g. Helion) can override codegen
-        result = kernel.codegen_template_override(
+        # Remove prologue-fused inputs from input_buffers so that
+        # remove_kernel_local_buffers can remove them.
+        for buf_name in kernel.prologue_fused_inputs:
+            kernel.args.input_buffers.pop(buf_name, None)
+
+        # Dispatch to the kernel for source generation.  TritonTemplateKernel
+        # handles the standard Triton path; ExternalTritonTemplateKernel
+        # handles external backends (e.g. Helion).
+        src_code = kernel.codegen_template_body(
             self,
             template_node,
             epilogue_nodes,
@@ -2061,110 +2108,35 @@ def _codegen_single_template(
             buf_name_to_prologue_group,
             prologue_preserves_zero_mask,
             render,
-            only_gen_src_code,
         )
-        if result is not None:
-            return result
 
-        with kernel:
-            if not only_gen_src_code:
-                # prologue nodes can only be fused if their only use is in the template,
-                # so they are necessarily not allocated
-                for node in [template_node, *epilogue_nodes]:
-                    node.mark_run()
-
-            partial_code = render()
+        if config.benchmark_kernel:
+            num_gb = kernel.estimate_kernel_num_bytes() / 1e9
+            src_code = (
+                f"{kernel.imports_for_benchmark_kernel()}\n"
+                f"{src_code}\n"
+                f"{kernel.codegen_kernel_benchmark(num_gb).getvalue()}"
+            )
 
-            num_store_subgraphs = kernel.get_store_output_count()
-            for i in range(num_store_subgraphs):
-                subgraph_name = kernel._get_store_output_subgraph_name(i)
-                with kernel.set_subgraph_body(subgraph_name):
-                    for node in epilogue_nodes:
-                        node.codegen(kernel.split_and_set_ranges(node.get_ranges()))
-                    kernel.cse.invalidate(OrderedSet())
+        node_schedule = [*prologue_nodes, template_node, *epilogue_nodes]
 
-            for input_name, buffer in kernel.named_input_nodes.items():
-                subgraph_name = f"<LOAD_INPUT_{input_name}>"
-                if prologue_group := buf_name_to_prologue_group.get(
-                    buffer.get_name(), []
-                ):
-                    can_codegen_without_upcast = all(
-                        p_n.can_codegen_without_upcasts() for p_n in prologue_group
-                    )
+        if only_gen_src_code:
+            return src_code
 
-                    # TODO - this doesn't work with libdevice calls, potentially other bugs
-                    # upcasting to fp32 and downcasting gives large slowdown
-                    with config.patch(
-                        "triton.codegen_upcast_to_fp32", not can_codegen_without_upcast
-                    ):
-                        with kernel.set_subgraph_body(subgraph_name):
-                            for prologue_node in prologue_group:
-                                if (
-                                    len(prologue_node.get_buffer_names()) == 1
-                                    and len(prologue_group) == 1
-                                ):
-                                    if prologue_preserves_zero_mask(prologue_node):
-                                        kernel.prologue_fused_inputs_preserve_zero |= (
-                                            prologue_node.get_buffer_names()
-                                        )
-
-                                prologue_node.codegen(
-                                    kernel.split_and_set_ranges(
-                                        prologue_node.get_ranges()
-                                    )
-                                )
-                            kernel.cse.invalidate(OrderedSet())
-
-        # Template hooks must be finalised after kernel.remove_kernel_local_buffers
-        # is called (this is called when the kernel context is exited above), and when
-        # the kernel handler is set (as below). This is because the hooks may add
-        # DeferredLine type lines, which preclude lines involving buffers that have
-        # been removed
-
-        # finalize must be called after adding epilogue above
+        # Unfused epilogues are codegen'd separately in call_kernel;
+        # exclude them from mark_run.
+        unfused_set = OrderedSet([id(n) for n in kernel.get_unfused_epilogues()])
         with V.set_kernel_handler(kernel):
-            if not isinstance(partial_code, str):
-                # This is used to calculate flops in TritonTemplateKernels
-                with ir.IRNode.current_origins(template_node.node.origins):
-                    partial_code.finalize_hook("<DEF_KERNEL>")
-                partial_code.finalize_hook("<ARGDEFS>", strict=False)
-
-            # TODO: Maybe unify CUTLASSTemplateKernel to also use PartialRender for flexible epilogue fusion.
-
-            for input_name in kernel.named_input_nodes:
-                subgraph_name = f"<LOAD_INPUT_{input_name}>"
-
-                partial_code.finalize_hook(subgraph_name, strict=False)
-
-            num_store_subgraphs = kernel.get_store_output_count()
-            for i in range(num_store_subgraphs):
-                subgraph_name = kernel._get_store_output_subgraph_name(i)
-
-                partial_code.finalize_hook(subgraph_name)
-
-            if isinstance(partial_code, str):
-                src_code = partial_code
-            else:
-                # Ensure all hooks are finalized before the kernel is defined.
-                # Note: some of these hooks may have been registered by a kernel subclass
-                src_code = partial_code.finalize_remaining()
-
-            node_schedule = [*prologue_nodes, template_node, *epilogue_nodes]
-
-            if config.benchmark_kernel:
-                num_gb = kernel.estimate_kernel_num_bytes() / 1e9
-                src_code = (
-                    f"{kernel.imports_for_benchmark_kernel()}\n"
-                    f"{src_code}\n"
-                    f"{kernel.codegen_kernel_benchmark(num_gb).getvalue()}"
-                )
-
-            if only_gen_src_code:
-                return src_code
+            template_node.mark_run()
+            for node in epilogue_nodes:
+                if id(node) not in unfused_set:
+                    node.mark_run()
+            for node in prologue_nodes:
+                node.mark_run()
 
-            kernel.kernel_name = self.define_kernel(src_code, node_schedule, kernel)
+        kernel.kernel_name = self.define_kernel(src_code, node_schedule, kernel)
 
-            return kernel
+        return kernel
 
     def _get_multikernel_shapes(
         self, node: MultiTemplateBuffer
@@ -2222,8 +2194,8 @@ def codegen_template(
         prologue_nodes,
         *,
         only_gen_src_code=False,
-        hint_override: Optional[int] = None,
-    ) -> Optional[str]:
+        hint_override: int | None = None,
+    ) -> str | None:
         """
         Codegen a triton template with multi-kernel dispatch support
 
@@ -2337,7 +2309,7 @@ def generate_combo_kernel_code(
         enable_autotune: bool,
         mixed_sizes: bool,
         only_gen_src_code: bool = False,
-    ) -> list[tuple[Optional[str], Any, Any]]:
+    ) -> list[tuple[str | None, Any, Any]]:
         """
         Generate kernel code for combo kernel partitions.
 
@@ -2521,7 +2493,7 @@ def collapse_ranges(ranges: Sequence[sympy.Expr]) -> sympy.Expr:
                 )
 
                 # score by number of elements
-                score = V.graph.sizevars.size_hint(
+                score = V.graph.sizevars.optimization_hint(
                     sympy_product(
                         size for size, stride in zip(ranges, strides) if stride != 0
                     )
@@ -2535,7 +2507,7 @@ def collapse_ranges(ranges: Sequence[sympy.Expr]) -> sympy.Expr:
                     score *= 2
 
                 if (
-                    V.graph.sizevars.size_hint(
+                    V.graph.sizevars.optimization_hint(
                         score - sympy_product(itertools.chain(ranges, reduction_ranges))
                     )
                     >= 0
@@ -2593,10 +2565,16 @@ def create_tiling(
         """
         Create a tiling dict from pointwise and reduction splits.
         """
-        pw_prefixes = ["z", "y", "x"][-len(pw_tiling) :]
-        reduction_prefixes = ["r0_", "r1_"][: len(reduction_tiling)]
+        pw_prefixes = ("z", "y", "x")
+        reduction_prefixes = ("r0_", "r1_")
+        assert len(pw_tiling) <= len(pw_prefixes)
+        assert len(reduction_tiling) <= len(reduction_prefixes)
+
         return immutable_dict(
-            [*zip(pw_prefixes, pw_tiling), *zip(reduction_prefixes, reduction_tiling)]
+            [
+                *zip(pw_prefixes[-len(pw_tiling) :], pw_tiling, strict=False),
+                *zip(reduction_prefixes, reduction_tiling, strict=False),
+            ]
         )
 
     @classmethod
@@ -2644,6 +2622,59 @@ def get_nd_tilings(
 
         Returns a list of tilings ranked by dimensionality.
         """
+
+        def collapse_dims(
+            dims: Sequence[sympy.Expr], fallback_numel: sympy.Expr
+        ) -> tuple[sympy.Expr, ...]:
+            """
+            Collapse dimensions to the maximum allowed number of tiles.
+            """
+            if not dims:
+                return (fallback_numel,)
+            max_tiles = get_max_tiles(2)
+            if V.graph.sizevars.statically_known_equals(
+                pointwise_numel, 1
+            ) and V.graph.sizevars.statically_known_gt(reduction_numel, 1):
+                # We only have at most two dimensions to tile over when emitting a
+                # reduction-only kernel.
+                max_tiles = min(max_tiles, 2)
+            num_leading_dims = max(0, len(dims) - max_tiles)
+            first_trailing_dim = num_leading_dims + 1
+            collapsed_leading_dim = sympy_product(dims[:first_trailing_dim])
+            return (collapsed_leading_dim,) + tuple(dims[first_trailing_dim:])
+
+        def tile_var_ranges(
+            var_ranges, ranges_to_tile, total_numel
+        ) -> tuple[sympy.Expr, ...]:
+            # Pattern match the subexpression pertaining to each index variable.
+            tiling = []
+            for var, numel in var_ranges:
+                index = BlockPatternMatcher.get_subexpr_involving_symbol(dep.index, var)
+
+                # Heuristic to bound the maximum dimensionality of the block.
+                num_dims = max(
+                    2,
+                    index.count(FloorDiv) + index.count(ModularIndexing),
+                    len(ranges_to_tile),
+                )
+
+                # Attempt to pattern match the index expr.
+                # Failed matches default to the full range.
+                match_result = BlockPatternMatcher.match_mod_div_block_expr(
+                    index, var, numel, num_dims
+                )
+                dims = match_result[0] if match_result is not None else [numel]
+                tiling.extend(dims)
+
+            # Prune dimensions of size 1.
+            tiling = [
+                dim
+                for dim in tiling
+                if not V.graph.sizevars.statically_known_equals(dim, sympy.S.One)
+            ]
+
+            return collapse_dims(tiling, total_numel)
+
         is_pointwise = reduction_numel == 1
         tilings = OrderedSet[immutable_dict[str, sympy.Expr]]()
         for node in EnableReduction.filter(node_schedule):
@@ -2657,8 +2688,9 @@ def get_nd_tilings(
                 continue
 
             # Use the node ranges as the default tiling candidate.
-            ranges_to_tile = node_ranges[0 if is_pointwise else 1]
-            node_tilings = [ranges_to_tile]
+            default_pointwise_tiling = collapse_dims(node_ranges[0], pointwise_numel)
+            default_reduction_tiling = collapse_dims(node_ranges[1], reduction_numel)
+            node_tilings = [(default_pointwise_tiling, default_reduction_tiling)]
 
             # Search the indexing expressions for more candidates.
             # If we see modular indexing, try to subdivide ranges into their implied
@@ -2689,62 +2721,34 @@ def get_nd_tilings(
                 ):
                     continue
 
-                # Partition var ranges into pointwise and reduction splits.
+                # Partition var ranges into pointwise and reduction splits, tiling them
+                # separately.
                 reduction_start_idx = pointwise_end_idx + 1
-                var_ranges = (
-                    all_var_ranges[:reduction_start_idx]
-                    if is_pointwise
-                    else all_var_ranges[reduction_start_idx:]
+                pointwise_var_ranges = all_var_ranges[:reduction_start_idx]
+                reduction_var_ranges = (
+                    None if is_pointwise else all_var_ranges[reduction_start_idx:]
                 )
-
-                # Pattern match the subexpression pertaining to each index variable.
-                index_tiling = []
-                for var, numel in var_ranges:
-                    index = BlockPatternMatcher.get_subexpr_involving_symbol(
-                        dep.index, var
-                    )
-
-                    # Heuristic to bound the maximum dimensionality of the block.
-                    num_dims = max(
-                        2,
-                        index.count(FloorDiv) + index.count(ModularIndexing),
-                        len(ranges_to_tile),
-                    )
-
-                    # Attempt to pattern match the index expr.
-                    # Failed matches default to the full range.
-                    match_result = BlockPatternMatcher.match_mod_div_block_expr(
-                        index, var, numel, num_dims
-                    )
-                    dims = match_result[0] if match_result is not None else [numel]
-                    index_tiling.extend(dims)
-
-                # Prune dimensions of size 1.
-                index_tiling = [
-                    dim
-                    for dim in index_tiling
-                    if not V.graph.sizevars.statically_known_equals(dim, sympy.S.One)
-                ]
-
-                if len(index_tiling) > 0:
-                    node_tilings.append(index_tiling)
-
-            # Flatten leading dimensions, assigning labels to each dim.
-            for node_tiling in node_tilings:
-                num_leading_dims = max(0, len(node_tiling) - get_max_tiles(2))
-                first_trailing_dim = num_leading_dims + 1
-                collapsed_leading_dim = sympy_product(node_tiling[:first_trailing_dim])
-                collapsed_splits = (collapsed_leading_dim,) + tuple(
-                    node_tiling[first_trailing_dim:]
+                pointwise_tiling = tile_var_ranges(
+                    pointwise_var_ranges, node_ranges[0], pointwise_numel
                 )
-                tilings.add(
-                    cls.complete_partial_tiling(
-                        cls.create_partial_tiling(collapsed_splits, is_pointwise),
-                        pointwise_numel,
-                        reduction_numel,
+                reduction_tiling = (
+                    (sympy.S.One,)
+                    if is_pointwise
+                    else tile_var_ranges(
+                        reduction_var_ranges, node_ranges[1], reduction_numel
                     )
                 )
 
+                if len(pointwise_tiling) and len(reduction_tiling) > 0:
+                    node_tilings.append((pointwise_tiling, reduction_tiling))
+
+            # Each memory dependency contributes one pointwise and reduction tiling. We
+            # take the Cartesian product of these, yielding all possible joint tilings.
+            for pointwise_tiling, reduction_tiling in itertools.product(
+                *zip(*node_tilings)
+            ):
+                tilings.add(cls.create_tiling(pointwise_tiling, reduction_tiling))
+
         # Rank tilings by the number of dimensions. E.g., prefer 2D to 1D.
         # Since this is a stable sort, ties are broken by schedule order.
         ranked_tilings = sorted(
@@ -2762,11 +2766,11 @@ def compute_tiling_strategy(
         pointwise_numel: sympy.Expr,
         reduction_numel: sympy.Expr,
         coalesce_analysis: CoalesceVarAnalysis,
-    ) -> tuple[dict[str, sympy.Expr], Optional[dict[str, sympy.Expr]]]:
+    ) -> tuple[dict[str, sympy.Expr], dict[str, sympy.Expr] | None]:
         """
         Generates a tiling, and a score of each tile according to each tile's coalesced memory accesses.
         """
-        tiling_var: Optional[sympy.Expr] = (
+        tiling_var: sympy.Expr | None = (
             None
             if not coalesce_analysis.suggested_split
             else coalesce_analysis.suggested_split.var
@@ -3012,7 +3016,7 @@ def select_tiling(
         node_schedule,
         numel,
         reduction_numel=sympy.S.One,
-        coalesce_analysis: Optional[CoalesceVarAnalysis] = None,
+        coalesce_analysis: CoalesceVarAnalysis | None = None,
     ) -> dict[str, sympy.Expr]:
         return cls.get_tiling_and_scores(
             node_schedule, numel, reduction_numel, coalesce_analysis
@@ -3024,8 +3028,8 @@ def get_tiling_and_scores(
         node_schedule,
         numel,
         reduction_numel=sympy.S.One,
-        coalesce_analysis: Optional[CoalesceVarAnalysis] = None,
-    ) -> tuple[dict[str, sympy.Expr], Optional[dict[str, sympy.Expr]]]:
+        coalesce_analysis: CoalesceVarAnalysis | None = None,
+    ) -> tuple[dict[str, sympy.Expr], dict[str, sympy.Expr] | None]:
         """
         Heuristics to decide how to tile kernels.
         Currently, we tile based on stride-1 dimensions.
@@ -3112,20 +3116,24 @@ def get_tiling_and_scores(
 
             def convert_tiling_to_3d(
                 tiling0: dict[str, sympy.Expr], tiling1: dict[str, sympy.Expr]
-            ) -> Optional[dict[str, sympy.Expr]]:
+            ) -> dict[str, sympy.Expr] | None:
                 a0, a1 = tiling0["x"], tiling0.get("y", 1)
                 b0, b1 = tiling1["x"], tiling1.get("y", 1)
 
-                if (
-                    free_unbacked_symbols([a1, b1])
-                    or V.graph.sizevars.size_hint(a1 - b1) == 0
-                ):
+                # TODO: These tiling decisions (equality, ordering, divisibility)
+                # are structural — they determine the kernel's iteration space
+                # decomposition — but are NOT guarded here. If the relationship
+                # between a1 and b1 changes across dynamic shape inputs, the
+                # compiled kernel could be wrong. Seems scary, unless there is a reason
+                # this is safe in that case probably we need better comment here !
+                hint = V.graph.sizevars.guarding_hint_or_throw
+                if hint(a1 - b1) == 0:
                     return None
-                if V.graph.sizevars.size_hint(a1 - b1) < 0:
+                if hint(a1 - b1) < 0:
                     # swap so a0 is bigger
                     (a0, a1), (b0, b1) = (b0, b1), (a0, a1)
 
-                assert V.graph.sizevars.size_hint(a1 - b1) > 0
+                assert hint(a1 - b1) > 0
                 if not V.graph.sizevars.statically_known_multiple_of(a1, b1):
                     return None
 
@@ -3171,7 +3179,7 @@ def ready_to_flush(self) -> bool:
         return False
 
     def generate_kernel_code_from_nodes(
-        self, nodes, benchmark_kernel=False, hint_override: Optional[int] = None
+        self, nodes, benchmark_kernel=False, hint_override: int | None = None
     ):
         if not any(n.is_template() for n in nodes):
             _, (numel, rnumel) = max(nodes, key=lambda x: int(x.is_reduction())).group
@@ -3182,8 +3190,11 @@ def generate_kernel_code_from_nodes(
                 features=SIMDKernelFeatures(node_schedule, numel, rnumel),
             )
             self.codegen_node_schedule_with_kernel(node_schedule, kernel)
+            # Collect config_patches from operations
+            config_patches = self._collect_config_patches(node_schedule)
+            config_patches["benchmark_kernel"] = benchmark_kernel
             with (
-                config.patch("benchmark_kernel", benchmark_kernel),
+                config.patch(**config_patches),
                 V.set_kernel_handler(kernel),
             ):
                 src_code = kernel.codegen_kernel()
@@ -3212,7 +3223,7 @@ def define_kernel(self, src_code, node_schedule, kernel):
 class CandidateTiling:
     tiling: dict[str, sympy.Expr]
     score: int  # higher is better
-    name: Optional[str] = None
+    name: str | None = None
 
     @staticmethod
     def is_good_size(s):
diff --git a/torch/_inductor/codegen/simd_kernel_features.py b/torch/_inductor/codegen/simd_kernel_features.py
index 6f6a9e1a77279..5b89828ef199e 100644
--- a/torch/_inductor/codegen/simd_kernel_features.py
+++ b/torch/_inductor/codegen/simd_kernel_features.py
@@ -5,7 +5,7 @@
 import functools
 import itertools
 import typing
-from typing import Any, Optional, Union
+from typing import Any
 
 import sympy
 
@@ -39,7 +39,7 @@ def is_reduction() -> bool:
         return False
 
 
-NodeScheduleEntry = Union[SchedulerNode, type[NodeScheduleMarker]]
+NodeScheduleEntry = SchedulerNode | type[NodeScheduleMarker]
 
 
 class DisableReduction(NodeScheduleMarker):
@@ -82,7 +82,7 @@ def __init__(
         node_schedule: list[NodeScheduleEntry],
         numel: sympy.Expr,
         reduction_numel: sympy.Expr = sympy.S.One,
-        coalesce_analysis: Optional[CoalesceVarAnalysis] = None,
+        coalesce_analysis: CoalesceVarAnalysis | None = None,
     ):
         self.node_schedule = node_schedule
         # numel excludes reduction_numel
@@ -150,7 +150,7 @@ def select_index_dtype(self) -> torch.dtype:
         return torch.int64
 
     def get_reduction_hint(
-        self, tiling_scores: Optional[dict[str, int]] = None
+        self, tiling_scores: dict[str, int] | None = None
     ) -> ReductionHint:
         reductions = self.reduction_nodes()
         if len(reductions) > 0:
@@ -230,7 +230,7 @@ def reduction_hint(node: Any) -> ReductionHint:
             return node.node.data.reduction_hint
 
     def memory_stats(
-        self, groups_dict: Optional[dict[str, sympy.Expr]] = None
+        self, groups_dict: dict[str, sympy.Expr] | None = None
     ) -> MemoryStats:
         """Analysis to generate features that can be used in heuristics"""
         if groups_dict is None:
diff --git a/torch/_inductor/codegen/subgraph.py b/torch/_inductor/codegen/subgraph.py
index 689321599e241..167891ba04ded 100644
--- a/torch/_inductor/codegen/subgraph.py
+++ b/torch/_inductor/codegen/subgraph.py
@@ -1,3 +1,5 @@
+import contextlib
+import functools
 import itertools
 import logging
 from collections.abc import Callable
@@ -7,7 +9,13 @@
 
 import torch
 import torch._inductor.config as config
+from torch._dynamo.utils import counters
 from torch._inductor import ir
+from torch._inductor.autotune_process import (
+    SubgraphCPUBenchmarkRequest,
+    SubgraphGPUBenchmarkRequest,
+    TensorMeta,
+)
 from torch._inductor.codegen.common import KernelTemplate
 from torch._inductor.ir import (
     Buffer,
@@ -66,7 +74,14 @@ def __init__(
     ) -> None:
         super().__init__(name, input_nodes, layout, description)
 
-        self.example_inputs = []
+        # Create inputs for tracing and benchmarking:
+        # - trace_inputs: Always use symbolic inputs for tracing so the graph works
+        #   for any size at runtime. This correctly captures shape-dependent operations
+        #   like x * x.shape[0].
+        # - benchmark_inputs: Use input_gen_fns if provided (for range-specific benchmarks)
+        #   or concrete inputs from ir_node_to_tensor with guard_shape=False.
+        trace_inputs = []
+        self.benchmark_inputs = []
         with V.fake_mode:
             for i, inp in enumerate(self.input_nodes):
                 # Here there will be no unbacked symbols, as SubgraphBuffer does not support them
@@ -75,15 +90,22 @@ def __init__(
 
                 inp.data.freeze_layout()  # type: ignore[attr-defined]
 
-                # Use input_gen_fn if provided for this index (e.g., for range-based autotuning)
-                # This ensures the graph is traced with correct target shapes
+                # Always use symbolic inputs for tracing
+                trace_inputs.append(ir_node_to_tensor(inp))
+
+                # Use input_gen_fn for benchmarking if provided, otherwise concrete sizes
                 if input_gen_fns is not None and i in input_gen_fns:
-                    self.example_inputs.append(input_gen_fns[i](inp))
+                    self.benchmark_inputs.append(input_gen_fns[i](inp))
                 else:
-                    self.example_inputs.append(ir_node_to_tensor(inp))
+                    self.benchmark_inputs.append(
+                        ir_node_to_tensor(inp, replace_symbols_with_hints=True)
+                    )
 
-        self.gm = make_fx_graph(*self.example_inputs)
+        # Trace with symbolic inputs for guard detection
+        self.gm = make_fx_graph(*trace_inputs)
         gm_original_output_strides(self.gm)
+        # Store symbolic inputs for sym_input computation
+        self.example_inputs = trace_inputs
 
         self.sym_inputs = get_symbolic_inputs(self.input_nodes)
         self.sym_input_values = self._compute_sym_input_values()
@@ -91,14 +113,27 @@ def __init__(
         # Cached decomposition info for range-based dispatch (set via cache_decomposition)
         self.decomposition: Callable[..., Any] | None = None
         self.decomposition_kwargs: dict[str, Any] = {}
+        # Config patches to apply during kernel codegen (e.g., coordinate_descent_tuning)
+        self.config_patches: dict[str, Any] = {}
         # Cache compiled module to avoid recompiling on every benchmark call
         self._compiled_module: Any = None
+        # Cache benchmark request for async autotuning
+        self._bmreq: (
+            SubgraphGPUBenchmarkRequest | SubgraphCPUBenchmarkRequest | None
+        ) = None
+
+        # Pre-compile only if using async pipelined autotuning
+        # Must happen in __init__ because compilation requires virtualized context (V.graph, V.debug)
+        if config.pipeline_max_autotune_gemm:
+            with V.fake_mode:
+                self._compiled_module = self._compile_for_benchmarking()
+                self._bmreq = self._create_benchmark_request()
 
     def _compute_sym_input_values(self) -> list[int]:
-        """Extract concrete dimension values for sym_inputs from example_inputs.
+        """Extract concrete dimension values for sym_inputs from benchmark_inputs.
 
         The compiled module expects symbolic dimension values as runtime arguments.
-        This maps each symbolic variable to its concrete value from the example tensors.
+        This maps each symbolic variable to its concrete value from the benchmark tensors.
         Used for range based autotuning.
         """
         sym_input_names = OrderedSet(
@@ -107,9 +142,11 @@ def _compute_sym_input_values(self) -> list[int]:
 
         # Build mapping: symbolic dimension name -> concrete value
         sym_name_to_value: dict[str, int] = {}
-        for inp_node, example_inp in zip(self.input_nodes, self.example_inputs):
-            if isinstance(example_inp, torch.Tensor):
-                for sym_dim, actual_dim in zip(inp_node.get_size(), example_inp.shape):
+        for inp_node, benchmark_inp in zip(self.input_nodes, self.benchmark_inputs):
+            if isinstance(benchmark_inp, torch.Tensor):
+                for sym_dim, actual_dim in zip(
+                    inp_node.get_size(), benchmark_inp.shape
+                ):
                     if isinstance(sym_dim, sympy.Symbol):
                         sym_name_to_value[sym_dim.name] = int(actual_dim)
                     elif str(sym_dim) in sym_input_names:
@@ -120,8 +157,8 @@ def _compute_sym_input_values(self) -> list[int]:
             if isinstance(sym_var, sympy.Symbol) and sym_var.name in sym_name_to_value:
                 result.append(sym_name_to_value[sym_var.name])
             else:
-                hint = V.graph.sizevars.shape_env.size_hint(sym_var)
-                result.append(int(hint) if hint is not None else 1)
+                hint = V.graph.sizevars.shape_env.optimization_hint(sym_var, fallback=1)
+                result.append(int(hint))
         return result
 
     def cache_decomposition(
@@ -140,9 +177,18 @@ def _compile_for_benchmarking(self) -> Any:
 
         safe_name = self.name.replace("::", "_").replace(".", "_")
 
+        # Symbolic inputs produce code with symbolic sizes resolved at runtime
+        # via sym_input_values. Without symbols (static shapes), compile with
+        # benchmark_inputs whose sizes match the actual benchmark tensors.
+        compile_inputs = (
+            self.example_inputs if self.sym_inputs else self.benchmark_inputs
+        )
+        log.debug("Benchmark compile %s: sym_inputs=%s", self.name, self.sym_inputs)
+
+        assert self.gm is not None
         bm_graph_lowering = GraphLowering(
             gm=self.gm,
-            example_inputs=self.example_inputs,
+            example_inputs=compile_inputs,
             shape_env=V.graph._shape_env,
             cpp_wrapper=V.graph.cpp_wrapper,
             aot_mode=V.graph.aot_mode,
@@ -157,38 +203,86 @@ def _compile_for_benchmarking(self) -> Any:
             bm_graph_lowering.graph_input_names.append(sym_inp.name)
 
         with V.set_graph_handler(bm_graph_lowering):
-            # Don't bother autotuning on Triton here
-            with config.patch(
-                max_autotune=False,
-                max_autotune_gemm=False,
-                max_autotune_gemm_backends="ATEN",
-            ):
-                bm_graph_lowering.run(*self.example_inputs)
+            # Apply config_patches during benchmarking (e.g., coordinate_descent_tuning)
+            # Also disable max_autotune to avoid nested autotuning
+            benchmark_config: dict[str, Any] = {
+                "max_autotune": False,
+                "max_autotune_gemm": False,
+                "max_autotune_gemm_backends": "ATEN",
+                "benchmark_fusion": False,
+                "pipeline_max_autotune_gemm": False,
+                **self.config_patches,
+            }
+            with config.patch(benchmark_config):
+                bm_graph_lowering.run(*compile_inputs)
                 return bm_graph_lowering.compile_to_module()
 
-    def benchmark(self, *args: list[Any], out: torch.Tensor) -> float:
-        """Regular benchmarking: compile and use benchmarker with warmup/rep."""
+    def _create_benchmark_request(
+        self,
+    ) -> SubgraphGPUBenchmarkRequest | SubgraphCPUBenchmarkRequest:
+        """Create a benchmark request for async autotuning."""
+        assert self._compiled_module is not None, (
+            "Module must be compiled before creating benchmark request"
+        )
+        input_tensor_meta = TensorMeta.from_irnodes(self.input_nodes)
+        output_tensor_meta = TensorMeta.from_irnodes(self.layout)
+
+        if self.layout.device.type == "cpu":
+            bmreq_cls = SubgraphCPUBenchmarkRequest
+        else:
+            bmreq_cls = SubgraphGPUBenchmarkRequest
+
+        return bmreq_cls(
+            kernel_name=self.name,
+            input_tensor_meta=input_tensor_meta,
+            output_tensor_meta=output_tensor_meta,
+            extra_args=tuple(),
+            module_path=self._compiled_module.__file__,
+            module_cache_key=self._compiled_module.key,
+            sym_input_values=self.sym_input_values,
+        )
+
+    @property
+    def bmreq(
+        self,
+    ) -> SubgraphGPUBenchmarkRequest | SubgraphCPUBenchmarkRequest:
+        """Benchmark request for async autotuning. Pre-compiled when pipeline_max_autotune_gemm is enabled."""
+        assert self._bmreq is not None, (
+            "bmreq accessed but pipeline_max_autotune_gemm was not enabled during __init__"
+        )
+        return self._bmreq
+
+    def _ensure_compiled(self) -> None:
+        """Ensure the module is compiled. Used for lazy compilation in non-async path."""
         if self._compiled_module is None:
             self._compiled_module = self._compile_for_benchmarking()
 
+    def benchmark(self, *args: list[Any], out: torch.Tensor) -> float:
+        """Regular benchmarking: compile if needed, then use benchmarker."""
+        self._ensure_compiled()
         bm_func = self._compiled_module.call
         sym_inputs = self.sym_input_values
+
+        def fn() -> Any:
+            return bm_func([*sym_inputs, *args])
+
+        if self._benchmark_with_cudagraphs:
+            return benchmarker.benchmark_gpu_with_cuda_graph(fn)
+
         if config.profile_bandwidth_with_do_bench_using_profiling:
-            return do_bench_using_profiling(lambda: bm_func([*sym_inputs, *args]))
+            return do_bench_using_profiling(fn)
         return benchmarker.benchmark(
-            # Shallow clone args since bm_func may clear args
-            lambda: bm_func([*sym_inputs, *args]),
+            fn,
             device=benchmarker.infer_device(*sym_inputs, *args),
         )
 
     def benchmark_collective(self, *args: list[Any], out: torch.Tensor) -> None:
         """Run once for collective benchmarking (barrier sync handled by caller)."""
-        if self._compiled_module is None:
-            self._compiled_module = self._compile_for_benchmarking()
-
+        self._ensure_compiled()
         self._compiled_module.call([*self.sym_input_values, *args])
 
     def hash_key(self) -> str:
+        assert self.gm is not None
         return "-".join(
             [
                 self.name.rsplit("_", 1)[0],
@@ -199,6 +293,7 @@ def hash_key(self) -> str:
         )
 
     def output_node(self) -> ir.TensorBox:
+        assert self.gm is not None
         return ir.TensorBox.create(
             ir.SubgraphBuffer(
                 layout=self.layout,
@@ -206,6 +301,7 @@ def output_node(self) -> ir.TensorBox:
                 gm=self.gm,
                 example_inputs=self.example_inputs,
                 subgraph_name=self.name,
+                config_patches=self.config_patches if self.config_patches else None,
             )
         )
 
@@ -287,6 +383,7 @@ def generate_custom_op_choices(
         non_tensor_args: list[dict[str, Any]],
         default_impl: Callable[..., Any] | None = None,
         input_gen_fns: dict[int, Callable[[Any], torch.Tensor]] | None = None,
+        config_patches_list: list[dict[str, Any]] | None = None,
     ) -> list[SubgraphChoiceCaller]:
         """
         Generate multiple SubgraphChoiceCaller instances for custom op autotuning.
@@ -301,6 +398,7 @@ def generate_custom_op_choices(
             non_tensor_args: List of non-tensor kwargs only, one dict per corresponding decomposition.
             default_impl: Default implementation for layout inference
             input_gen_fns: Optional dict mapping input indices to tensor generators
+            config_patches_list: Optional list of config patches per decomposition
 
         Returns:
             List of SubgraphChoiceCaller instances for autotuning
@@ -313,6 +411,10 @@ def generate_custom_op_choices(
             f"got {len(decompositions)} decompositions and {len(non_tensor_args)} kwargs"
         )
 
+        # Default to empty config_patches if not provided
+        if config_patches_list is None:
+            config_patches_list = [{} for _ in decompositions]
+
         # Infer layouts and ensure layout consistency for fair autotuning comparison
         layouts = [
             self._infer_custom_op_layout(
@@ -326,9 +428,12 @@ def generate_custom_op_choices(
         layout = layouts[0]  # All layouts are now validated to be equivalent
 
         choices: list[SubgraphChoiceCaller] = []
-        for decomp, decomp_kwargs in zip(decompositions, non_tensor_args):
+        for decomp, decomp_kwargs, config_patches in zip(
+            decompositions, non_tensor_args, config_patches_list
+        ):
             # Create make_fx_graph function for this decomposition
-            import functools
+            # Uses error_on_new_guards to detect impls that add guards
+            from torch.fx.experimental.symbolic_shapes import _ShapeEnvGuardError
 
             def make_fx_graph(
                 *args: Any,
@@ -336,30 +441,52 @@ def make_fx_graph(
                 decomp_kwargs: dict[str, Any] = decomp_kwargs,
             ) -> Any:
                 # decomp_kwargs contains all merged parameters: CustomOpConfig params + runtime kwargs
+
                 from torch.fx.experimental.proxy_tensor import make_fx
 
                 from ..decomposition import select_decomp_table
 
                 decomposition_table = select_decomp_table()
+                shape_env = V.fake_mode.shape_env
 
-                return make_fx(
-                    functools.partial(decomp, **decomp_kwargs),
-                    decomposition_table=decomposition_table,
-                )(*args)
+                # Use error_on_new_guards to detect impls that add guards during tracing
+                guard_ctx = (
+                    shape_env.error_on_new_guards()
+                    if shape_env is not None
+                    else contextlib.nullcontext()
+                )
+                with guard_ctx:
+                    return make_fx(
+                        functools.partial(decomp, **decomp_kwargs),
+                        decomposition_table=decomposition_table,
+                        tracing_mode="symbolic",
+                    )(*args)
 
             # Generate descriptive name for this variant
             variant_name = self._generate_variant_name(decomp, decomp_kwargs)
 
-            choice = self.generate(
-                name=f"{name}_{variant_name}",
-                input_nodes=input_nodes,
-                layout=layout,
-                make_fx_graph=make_fx_graph,
-                description=f"CustomOp {decomp.__name__}",
-                input_gen_fns=input_gen_fns,
-            )
+            # Try to create choice; skip if it adds guards
+            try:
+                choice = self.generate(
+                    name=f"{name}_{variant_name}",
+                    input_nodes=input_nodes,
+                    layout=layout,
+                    make_fx_graph=make_fx_graph,
+                    description=f"CustomOp {decomp.__name__}",
+                    input_gen_fns=input_gen_fns,
+                )
+            except _ShapeEnvGuardError:
+                log.info(
+                    "Skipping decomposition %s: adds guards during tracing",
+                    decomp.__name__,
+                )
+                counters["inductor"]["custom_op_decomp_guard_skips"] += 1
+                continue
+
             # Cache decomposition info for range-based dispatch
             choice.cache_decomposition(decomp, decomp_kwargs)
+            # Store config_patches for this choice
+            choice.config_patches = config_patches
             choices.append(choice)
 
         return choices
@@ -368,10 +495,25 @@ def _generate_variant_name(
         self, decomp: Callable[..., Any], kwargs: dict[str, Any]
     ) -> str:
         """Generate a descriptive name for a decomposition variant with its parameters."""
+        import re
+
         base_name = decomp.__name__
         if not kwargs:
             return base_name
-        param_suffix = "_".join(f"{k}_{v}" for k, v in sorted(kwargs.items()))
+
+        def sanitize_value(v: Any) -> str:
+            """Convert a value to a valid Python identifier component."""
+            s = str(v)
+            # Replace invalid characters with underscores
+            s = re.sub(r"[^a-zA-Z0-9_]", "_", s)
+            # Ensure it doesn't start with a digit
+            if s and s[0].isdigit():
+                s = "_" + s
+            return s
+
+        param_suffix = "_".join(
+            f"{k}_{sanitize_value(v)}" for k, v in sorted(kwargs.items())
+        )
         return f"{base_name}_{param_suffix}"
 
     def _validate_non_tensor_kwargs(self, kwargs: dict[str, Any]) -> None:
diff --git a/torch/_inductor/codegen/triton.py b/torch/_inductor/codegen/triton.py
index bca5da15e3e22..f052d585d3b45 100644
--- a/torch/_inductor/codegen/triton.py
+++ b/torch/_inductor/codegen/triton.py
@@ -1,8 +1,10 @@
 # mypy: allow-untyped-defs
 from __future__ import annotations
 
+import ast
 import collections
 import contextlib
+import copy
 import dataclasses
 import functools
 import itertools
@@ -14,7 +16,7 @@
 from abc import abstractmethod
 from collections.abc import Callable, Iterable, Sequence
 from functools import lru_cache
-from typing import Any, cast, Optional, TYPE_CHECKING, TypeVar, Union
+from typing import Any, cast, TYPE_CHECKING, TypeVar
 
 import sympy
 from sympy.printing.precedence import PRECEDENCE
@@ -24,9 +26,15 @@
 import torch.utils._pytree as pytree
 from torch._dynamo.device_interface import get_interface_for_device
 from torch._dynamo.utils import identity, preserve_rng_state
-from torch._prims_common import is_integer_dtype
+from torch._prims_common import is_integer_dtype, type_to_dtype
 from torch.utils._ordered_set import OrderedSet
-from torch.utils._sympy.functions import CeilDiv, FloorDiv, ModularIndexing
+from torch.utils._sympy.functions import (
+    CeilDiv,
+    FloorDiv,
+    ModularIndexing,
+    TruncToFloat,
+    TruncToInt,
+)
 from torch.utils._triton import (
     get_triton_version,
     has_triton_package,
@@ -45,13 +53,21 @@
 from ..runtime.hints import (
     AutotuneHint,
     DeviceProperties,
+    get_warp_size,
     ReductionHint,
     TRITON_MAX_BLOCK,
     TRITON_MAX_RSPLIT,
 )
 from ..runtime.runtime_utils import get_max_y_grid, next_power_of_2
-from ..scheduler import BaseSchedulerNode, FusedSchedulerNode, Scheduler, SchedulerNode
+from ..scheduler import (
+    BaseSchedulerNode,
+    FusedExternTritonKernelSchedulerNode,
+    FusedSchedulerNode,
+    Scheduler,
+    SchedulerNode,
+)
 from ..shape_propagation import get_broadcasted_shape
+from ..stream_utils import get_raw_stream_name
 from ..utils import (
     cache_on_self,
     DelayReplaceLine,
@@ -152,6 +168,47 @@ def is_sympy_integer_like(expr: object):
     )
 
 
+def _materialize_trunc_to_float_expr(
+    expr: sympy.Expr, dtype: torch.dtype
+) -> sympy.Expr:
+    if not dtype.is_floating_point or not expr.has(TruncToInt):
+        return expr
+
+    # Preserve float truncation semantics when materializing symbolic scalars
+    # into floating tensors. Casting to the kernel index dtype first can
+    # overflow before the requested floating-point conversion happens. Only
+    # rewrite truncations that are already participating in floating-point
+    # computation; integer subexpressions and predicates must keep exact
+    # integer semantics until the final materialization cast.
+    if expr.func is TruncToInt:
+        return TruncToFloat(*expr.args)
+
+    def is_predicate_expr(node: sympy.Basic) -> bool:
+        return bool(
+            getattr(node, "is_Boolean", False) or getattr(node, "is_Relational", False)
+        )
+
+    def rewrite_float_subexpr(node: sympy.Expr) -> sympy.Expr:
+        if not node.has(TruncToInt):
+            return node
+        if node.func is TruncToInt:
+            return TruncToFloat(*node.args)
+        if is_predicate_expr(node) or node.is_integer:
+            return node
+
+        new_args = tuple(
+            rewrite_float_subexpr(arg)
+            if isinstance(arg, sympy.Expr) and not is_predicate_expr(arg)
+            else arg
+            for arg in node.args
+        )
+        if new_args == node.args:
+            return node
+        return node.func(*new_args)
+
+    return rewrite_float_subexpr(expr)
+
+
 class OpDtypeSupport:
     """
     Some Triton ops such as libdevice and tl.math only support float32 and float64.
@@ -275,10 +332,10 @@ def get_block_offset(cls, tree: IterationRanges) -> sympy.Symbol:
 class IndexingOptions:
     index_str: str
     mask_vars: OrderedSet[str]
-    expand_str: Optional[str]
+    expand_str: str | None
     _has_rindex: bool
     index: sympy.Expr
-    expand_shape: Optional[Sequence[Union[int, str]]]
+    expand_shape: Sequence[int | str] | None
 
     def has_mask(self) -> bool:
         return bool(self.mask_vars)
@@ -324,7 +381,7 @@ class BlockDescriptorOptions:
     # transpose load / store blocks at runtime using the information in
     # stride_sorter.
     stride_sorter: BlockParameters.StrideSorter
-    _boundary_check: Optional[list[int]] = None
+    _boundary_check: list[int] | None = None
     # Can we safely lift the constructor
     # to the top of the kernel?
     can_lift: bool = False
@@ -760,6 +817,15 @@ def _print_TruncToInt(self, expr: sympy.Expr) -> str:
             f"libdevice.trunc({self._print(expr.args[0])}).to({V.kernel.index_dtype})"
         )
 
+    def _print_TruncToFloat(self, expr: sympy.Expr) -> str:
+        assert len(expr.args) == 1
+        # pyrefly: ignore [missing-attribute]
+        value = self._print(expr.args[0])
+        # Adding +0.0 preserves large floating results while canonicalizing
+        # libdevice.trunc(-0.0) back to Python's +0.0 materialization behavior.
+        # pyrefly: ignore [missing-attribute]
+        return f"(libdevice.trunc({value}) + tl.zeros_like({value}))"
+
     def _print_Float(self, expr: sympy.Expr) -> str:
         if expr.is_integer:
             # sympy considers 0.0 to be integer, but triton doesn't.
@@ -1002,7 +1068,7 @@ def low_precision_fp(dtype: torch.dtype) -> bool:
     return dtype.itemsize <= 2 and dtype.is_floating_point
 
 
-def low_precision_fp_var(var: Union[CSEVariable, Any]) -> bool:
+def low_precision_fp_var(var: CSEVariable | Any) -> bool:
     if not isinstance(var, CSEVariable):
         return False
 
@@ -1010,6 +1076,22 @@ def low_precision_fp_var(var: Union[CSEVariable, Any]) -> bool:
     return low_precision_fp(dtype) if isinstance(dtype, torch.dtype) else False
 
 
+def triton_arg_dtype(arg: Any) -> torch.dtype | None:
+    if isinstance(arg, CSEVariable):
+        return arg.dtype
+    if isinstance(arg, torch._prims_common.Number):
+        return type_to_dtype(type(arg))
+    return None
+
+
+def needs_upcast_to_float32(arg: Any) -> bool:
+    return (
+        not config.triton.codegen_upcast_to_fp32
+        and isinstance(arg, CSEVariable)
+        and arg.dtype in (torch.float16, torch.bfloat16)
+    )
+
+
 class TritonCSEVariable(CSEVariable):
     def __init__(
         self,
@@ -1050,15 +1132,8 @@ def maybe_upcast_float32(convert_output: bool = True) -> Callable[[_T], _T]:
     This decorates tl.math/libdevice codegen functions.
     """
 
-    def needs_upcast(var) -> bool:
-        return (
-            not config.triton.codegen_upcast_to_fp32
-            and isinstance(var, CSEVariable)
-            and var.dtype in (torch.float16, torch.bfloat16)
-        )
-
     def maybe_upcast_arg(var) -> str:
-        upcast_string = ".to(tl.float32)" if needs_upcast(var) else ""
+        upcast_string = ".to(tl.float32)" if needs_upcast_to_float32(var) else ""
         return f"{var}{upcast_string}"
 
     def decorator(func: Callable[..., Any]) -> Callable[..., Any]:
@@ -1073,7 +1148,8 @@ def wrapped(*args, **kwargs) -> str:
             # Call the decorated function, optionally downcasting the result.
             result = func(*upcast_args, **upcast_kwargs)
             any_needs_upcast = convert_output and any(
-                needs_upcast(var) for var in itertools.chain(args, kwargs.values())
+                needs_upcast_to_float32(var)
+                for var in itertools.chain(args, kwargs.values())
             )
             result_dtype = (
                 None
@@ -1102,7 +1178,7 @@ class TritonOverrides(OpOverrides):
     def to_dtype(
         x,
         dtype: torch.dtype,
-        src_dtype: Optional[torch.dtype] = None,
+        src_dtype: torch.dtype | None = None,
         use_compute_types=True,
     ):
         def _get_min_elements_per_thread(
@@ -1182,6 +1258,14 @@ def _shaped_constant(value, dtype, shape):
         triton_val = constant_repr(type_(value))
         triton_type = triton_compute_type(dtype)
 
+        # Triton's scalar_constant() treats -0.0 as 0 (since -0.0 == 0 in Python),
+        # Work around by encoding -0.0 as its IEEE 754 hex in uint and bitcasting.
+        if value == 0 and math.copysign(1.0, value) < 0:
+            if triton_type == "tl.float32":
+                return f"tl.full({shape}, 0x80000000, tl.uint32).to({triton_type}, bitcast=True)"
+            elif triton_type == "tl.float64":
+                return f"tl.full({shape}, 0x8000000000000000, tl.uint64).to({triton_type}, bitcast=True)"
+
         # NOTE: We use tl.full here to get the expected type.
         # Otherwise, subnormal float32 values are treated as fp64
         # causing fp32 * fp64 promotion and different numerical results.
@@ -1195,8 +1279,17 @@ def _shaped_constant(value, dtype, shape):
     def constant(cls, value, dtype):
         return cls._shaped_constant(value, dtype, shape=[])
 
+    @classmethod
+    def _cast_libdevice_arg(cls, arg, dtype: torch.dtype) -> str:
+        if isinstance(arg, torch._prims_common.Number):
+            return cls.constant(arg, dtype)
+        if triton_arg_dtype(arg) == dtype:
+            return f"{arg}"
+        return f"({arg}).to({triton_type(dtype)})"
+
     @staticmethod
     @maybe_upcast_float32()
+    # pyrefly: ignore [bad-override]
     def abs(x):
         return f"tl_math.abs({x})"
 
@@ -1226,6 +1319,28 @@ def truediv(x, y):
 
         return out
 
+    @staticmethod
+    def div_rn(x, y):
+        """
+        Division with round-to-nearest rounding mode.
+        Always uses triton.language.div_rn for float32 inputs to match eager CUDA behavior.
+        """
+        x_dtype = getattr(x, "dtype", None)
+        y_dtype = getattr(y, "dtype", None)
+
+        if x_dtype == torch.float32 and y_dtype == torch.float32:
+            out = f"triton.language.div_rn({x}, {y})"
+        else:
+            # Fall back to regular division for non-float32 types
+            out = f"({x} / {y})"
+
+        if low_precision_fp_var(x) or low_precision_fp_var(y):
+            out_dtype = get_dtype_handler().truediv(x, y)
+            if out_dtype in (torch.float16, torch.float32):
+                out = f"{out}.to({triton_type(out_dtype)})"
+
+        return out
+
     @staticmethod
     def mod(x, y):
         out = f"({x} % {y})"
@@ -1237,6 +1352,7 @@ def mod(x, y):
 
     @staticmethod
     @maybe_upcast_float32()
+    # pyrefly: ignore [bad-override]
     def exp(x):
         """
         When use_fast_math, use the ftz (flushing to zero) variant
@@ -1262,6 +1378,7 @@ def expm1(x):
 
     @staticmethod
     @maybe_upcast_float32()
+    # pyrefly: ignore [bad-override]
     def sqrt(x):
         return f"tl.sqrt_rn({x})"
 
@@ -1284,6 +1401,7 @@ def relu(x):
             )
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def minimum(a, b):
         if torch.version.hip:
             return f"tl.minimum({a}, {b}, tl.PropagateNan.ALL)"
@@ -1291,6 +1409,7 @@ def minimum(a, b):
             return f"triton_helpers.minimum({a}, {b})"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def maximum(a, b):
         if torch.version.hip:
             return f"tl.maximum({a}, {b}, tl.PropagateNan.ALL)"
@@ -1298,10 +1417,12 @@ def maximum(a, b):
             return f"triton_helpers.maximum({a}, {b})"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def where(a, b, c):
         return f"tl.where({a}, {b}, {c})"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def dot(a, b):
         """
         Triton code generation for lowering ops.dot to tl.dot.
@@ -1490,21 +1611,83 @@ def reshape_transpose_broadcast_for_dot(
 
     @staticmethod
     def inline_asm_elementwise(
-        *inputs, asm, constraints=None, dtype=torch.float32, is_pure=True, pack=1
+        *inputs,
+        asm,
+        constraints=None,
+        dtype=torch.float32,
+        is_pure=True,
+        pack=1,
+        input_dtypes=None,
     ):
-        triton_type = triton_compute_type(dtype)
-        input_refs = ", ".join([str(i) for i in inputs])
+        # Use the actual dtype, not the compute type — the asm operates on
+        # specific register types and Triton needs to know the real output type.
+        asm_triton_type = triton_type(dtype)
         if constraints is None:
             constraints = ", ".join(["=r"] + ["r" for _ in inputs])
-        return f"tl.inline_asm_elementwise('{asm}', '{constraints}', [{input_refs}], dtype={triton_type}, is_pure={is_pure}, pack={pack})"  # noqa: B950
+
+        # Inductor computes bf16/fp16 in fp32. For "h" (16-bit register)
+        # constraints, cast back to the original dtype so the asm sees the
+        # right register type.
+        constraint_parts = [p.strip() for p in constraints.split(",")]
+        input_constraints = [p for p in constraint_parts if not p.startswith("=")]
+        cast_inputs = []
+        for i, (inp, c) in enumerate(zip(inputs, input_constraints[: len(inputs)])):
+            if (
+                c == "h"
+                and input_dtypes is not None
+                and isinstance(inp, CSEVariable)
+                and inp.dtype != input_dtypes[i]
+            ):
+                cast_inputs.append(f"{inp}.to({triton_type(input_dtypes[i])})")
+            else:
+                cast_inputs.append(str(inp))
+
+        if torch.version.hip:
+            # AMDGCN asm strings may contain real newlines (instructions are
+            # newline-separated, unlike PTX which uses semicolons).  The
+            # generated code is nested inside two Python string layers:
+            #   Layer 1 : the cached wrapper .py file
+            #   Layer 2 : the Triton kernel source (a triple-quoted string
+            #             inside that wrapper, exec'd / JIT-compiled)
+            # repr() escapes \n -> \\n, then we double the backslashes so
+            # they survive both layers: \\\\n -> (L1 parse) \\n -> (L2 parse) \n.
+            asm_literal = repr(asm).replace("\\", "\\\\")
+            constraints_literal = repr(constraints).replace("\\", "\\\\")
+        else:
+            asm_literal = f"'{asm}'"
+            constraints_literal = f"'{constraints}'"
+
+        def asm_call(args):
+            return (
+                f"tl.inline_asm_elementwise({asm_literal}, {constraints_literal}, "
+                f"[{args}], dtype={asm_triton_type}, is_pure={is_pure}, pack={pack})"
+            )
+
+        if pack <= 1:
+            return asm_call(", ".join(cast_inputs))
+
+        first_input = inputs[0]
+        compute = V.kernel.compute
+        cse = V.kernel.cse
+        result = cse.newvar(dtype=dtype, shape=first_input.shape)
+        packed_args = ", ".join(
+            f"triton_helpers.inline_asm_pack({inp}, {pack})" for inp in cast_inputs
+        )
+        compute.writeline(f"{result} = {asm_call(packed_args)}")
+        compute.writeline(
+            f"{result} = triton_helpers.inline_asm_unpack({result}, {first_input}, {pack})"
+        )
+        return result
 
     @staticmethod
     @maybe_upcast_float32()
+    # pyrefly: ignore [bad-override]
     def cos(x):
         return f"tl_math.cos({x})"
 
     @staticmethod
     @maybe_upcast_float32()
+    # pyrefly: ignore [bad-override]
     def sin(x):
         return f"tl_math.sin({x})"
 
@@ -1518,61 +1701,73 @@ def masked(mask, body, other):
 
     @staticmethod
     @maybe_upcast_float32()
+    # pyrefly: ignore [bad-override]
     def lgamma(x):
         return f"libdevice.lgamma({x})"
 
     @staticmethod
     @maybe_upcast_float32()
+    # pyrefly: ignore [bad-override]
     def erf(x):
         return f"libdevice.erf({x})"
 
     @staticmethod
     @maybe_upcast_float32()
+    # pyrefly: ignore [bad-override]
     def cosh(x):
         return f"libdevice.cosh({x})"
 
     @staticmethod
     @maybe_upcast_float32()
+    # pyrefly: ignore [bad-override]
     def sinh(x):
         return f"libdevice.sinh({x})"
 
     @staticmethod
     @maybe_upcast_float32()
+    # pyrefly: ignore [bad-override]
     def acos(x):
         return f"libdevice.acos({x})"
 
     @staticmethod
     @maybe_upcast_float32()
+    # pyrefly: ignore [bad-override]
     def acosh(x):
         return f"libdevice.acosh({x})"
 
     @staticmethod
     @maybe_upcast_float32()
+    # pyrefly: ignore [bad-override]
     def asin(x):
         return f"libdevice.asin({x})"
 
     @staticmethod
     @maybe_upcast_float32()
+    # pyrefly: ignore [bad-override]
     def asinh(x):
         return f"libdevice.asinh({x})"
 
     @staticmethod
     @maybe_upcast_float32()
+    # pyrefly: ignore [bad-override]
     def atan2(x, y):
         return f"libdevice.atan2({x}, {y})"
 
     @staticmethod
     @maybe_upcast_float32()
+    # pyrefly: ignore [bad-override]
     def atan(x):
         return f"libdevice.atan({x})"
 
     @staticmethod
     @maybe_upcast_float32()
+    # pyrefly: ignore [bad-override]
     def atanh(x):
         return f"libdevice.atanh({x})"
 
     @staticmethod
     @maybe_upcast_float32()
+    # pyrefly: ignore [bad-override]
     def copysign(x, y):
         return f"libdevice.copysign({x}, {y})"
 
@@ -1583,11 +1778,13 @@ def erfc(x):
 
     @staticmethod
     @maybe_upcast_float32()
+    # pyrefly: ignore [bad-override]
     def erfinv(x):
         return f"libdevice.erfinv({x})"
 
     @staticmethod
     @maybe_upcast_float32()
+    # pyrefly: ignore [bad-override]
     def hypot(x, y):
         return f"libdevice.hypot({x}, {y})"
 
@@ -1602,15 +1799,18 @@ def log2(x):
         return f"libdevice.log2({x})"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def ldexp(x, n):
         return f"libdevice.ldexp({x}, {n}.to(tl.int32))"
 
     @staticmethod
     @maybe_upcast_float32()
+    # pyrefly: ignore [bad-override]
     def nextafter(x, y):
         return f"libdevice.nextafter({x}, {y})"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def logical_and(a, b):
         return f"{a} & {b}"
 
@@ -1619,10 +1819,12 @@ def logical_not(a):
         return f"{a} == 0"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def logical_or(a, b):
         return f"{a} | {b}"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def logical_xor(a, b):
         return f"({a} ^ {b})"
 
@@ -1655,6 +1857,16 @@ def rand(seed, offset):
         offset = f"({offset}).to(tl.uint32)"
         return f"tl.rand({seed}, {offset})"
 
+    @staticmethod
+    def rand_eager(seed, base_offset, threads_per_round, tid, vec):
+        # vec: 4 for fp32, 8 for fp16/bf16
+        tid_u32 = f"({tid}).to(tl.uint32)"
+        denom = f"(({vec})*({threads_per_round}))"
+        r = f"(({tid_u32})//({denom})*({vec}//4))"
+        tid_trunc = f"(({tid_u32})%({denom}))"
+
+        return f"triton_helpers.rand_eager_kernel({seed}, {base_offset}+{r}, {tid_trunc}, VEC={vec})"
+
     @staticmethod
     def randn(seed, offset):
         offset = f"({offset}).to(tl.uint32)"
@@ -1671,6 +1883,7 @@ def load_seed(name, offset):
 
     @staticmethod
     @maybe_upcast_float32()
+    # pyrefly: ignore [bad-override]
     def rsqrt(x):
         if torch.version.hip:
             return f"tl.rsqrt({x})"
@@ -1684,11 +1897,13 @@ def log1p(x):
 
     @staticmethod
     @maybe_upcast_float32()
+    # pyrefly: ignore [bad-override]
     def tan(x):
         return f"libdevice.tan({x})"
 
     @staticmethod
     @maybe_upcast_float32()
+    # pyrefly: ignore [bad-override]
     def tanh(x):
         cse_var = V.kernel.cse.varname_map.get(x)
         if cse_var and hasattr(cse_var, "dtype"):
@@ -1714,6 +1929,7 @@ def sigmoid(x):
         return f"tl.sigmoid({x})"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def signbit(x):
         # XX: This is wrong for the value -0.0 in floating point
         return (
@@ -1722,49 +1938,105 @@ def signbit(x):
 
     @staticmethod
     @maybe_upcast_float32()
+    # pyrefly: ignore [bad-override]
     def fmod(a, b):
         return f"libdevice.fmod({a}, {b})"
 
-    @staticmethod
-    @maybe_upcast_float32()
-    def pow(a, b):
-        return f"libdevice.pow({a}, {b})"
+    @classmethod
+    def pow(cls, a, b):
+        result_dtype = get_dtype_handler().pow(a, b)
+        if result_dtype is not None and is_integer_dtype(result_dtype):
+            base = cls._cast_libdevice_arg(a, result_dtype)
+            exponent = (
+                cls.constant(b, torch.int64)
+                if isinstance(b, torch._prims_common.Number)
+                else f"{b}"
+            )
+            return f"triton_helpers.pow_integer({base}, {exponent})"
+
+        any_needs_upcast = needs_upcast_to_float32(a) or needs_upcast_to_float32(b)
+        pow_dtype = result_dtype
+        if pow_dtype not in (torch.float32, torch.float64):
+            # libdevice.pow only accepts fp32/fp64. Keep low-precision floating
+            # cases on the existing fp32 path, and otherwise fall back to fp64
+            # for symbolic integer scalar pow expressions like 2 ** ks0.
+            pow_dtype = (
+                torch.float32
+                if low_precision_fp(result_dtype) or any_needs_upcast
+                else torch.float64
+            )
+
+        cast_a = cls._cast_libdevice_arg(a, pow_dtype)
+        cast_b = cls._cast_libdevice_arg(b, pow_dtype)
+        result = f"libdevice.pow({cast_a}, {cast_b})"
+        if result_dtype is not None and result_dtype != pow_dtype:
+            if low_precision_fp(result_dtype):
+                if any_needs_upcast:
+                    result = f"{result}.to({triton_type(result_dtype)})"
+            else:
+                result = f"{result}.to({triton_type(result_dtype)})"
+        return result
 
     @staticmethod
     @maybe_upcast_float32()
+    # pyrefly: ignore [bad-override]
     def log(x):
         return f"tl_math.log({x})"
 
     @staticmethod
     @maybe_upcast_float32(convert_output=False)
+    # pyrefly: ignore [bad-override]
     def isinf(x):
         return f"libdevice.isinf({x}).to(tl.int1)"
 
     @staticmethod
     @maybe_upcast_float32(convert_output=False)
+    # pyrefly: ignore [bad-override]
     def isnan(x):
         return f"libdevice.isnan({x}).to(tl.int1)"
 
     @staticmethod
     @maybe_upcast_float32()
+    # pyrefly: ignore [bad-override]
     def round(x):
         return f"libdevice.nearbyint({x})"
 
     @staticmethod
     @maybe_upcast_float32()
+    # pyrefly: ignore [bad-override]
     def floor(x):
         return f"libdevice.floor({x})"
 
     @staticmethod
     def floordiv(a, b):
         # See the comment in lowering.div_mode. a and b are integer type.
-        # Similar to div_floor_kernel_cuda in pytorch core.
-        # Notice that // in triton behaves as truncdiv instead of floordiv
-        quot = f"{a} // {b}"
-        rem = f"{a} % {b}"
-        return f"tl.where(({a} < 0) != ({b} < 0), tl.where({rem} != 0, {quot} - 1, {quot}), {quot})"
+        # Notice that // in triton behaves as truncdiv instead of floordiv.
+        #
+        # We avoid feeding negative values into Triton's // operator because
+        # Triton's AxisInfo analysis incorrectly deduplicates signed division
+        # results for contiguous inputs (triton-lang/triton#XXXX). Instead we
+        # use bitwise complement (~) to make the dividend non-negative:
+        #   floor_div(a, b) = ~(~a // b) when a < 0, a // b when a >= 0
+        # For negative b we negate both operands first.
+        zero = ops.constant(0, torch.int32)
+        one = ops.constant(1, torch.int32)
+        # Guard against integer division by zero before the division to
+        # avoid undefined behavior (LLVM may optimize away a post-division
+        # check assuming UB doesn't happen). Replace b with 1 when b is 0
+        # so the division is safe, then select 0 as the final result.
+        b_zero = ops.eq(b, zero)
+        b = ops.where(b_zero, one, b)
+        b_neg = ops.lt(b, zero)
+        a = ops.where(b_neg, ops.sub(zero, a), a)
+        b = ops.where(b_neg, ops.sub(zero, b), b)
+        a_neg = ops.lt(a, zero)
+        a = ops.where(a_neg, ops.bitwise_not(a), a)
+        quot = ops.truncdiv(a, b)
+        quot = ops.where(a_neg, ops.bitwise_not(quot), quot)
+        return ops.where(b_zero, zero, quot)
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def sign(x):
         z = ops.constant(0, torch.int32)
         left = ops.to_dtype((ops.lt(z, x)), torch.int8)
@@ -1774,10 +2046,12 @@ def sign(x):
 
     @staticmethod
     @maybe_upcast_float32()
+    # pyrefly: ignore [bad-override]
     def trunc(x):
         return f"libdevice.trunc({x})"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def truncdiv(a, b):
         # See the comment in lowering.div_mode. a and b are integer type.
         # Notice that // in triton behaves as truncdiv instead of floordiv
@@ -1785,10 +2059,14 @@ def truncdiv(a, b):
 
     @staticmethod
     @maybe_upcast_float32()
+    # pyrefly: ignore [bad-override]
     def ceil(x):
         return f"libdevice.ceil({x})"
 
 
+# Register the custom pow override after class creation so type checkers see
+# a plain callable instead of the class-body staticmethod descriptor.
+OpDtypeSupport.register_upcast(TritonOverrides.pow, True)
 TritonOverrides._initialize_pointwise_overrides("triton")
 
 
@@ -1856,6 +2134,7 @@ def constant(cls, value, dtype):
 
     @classmethod
     def index_expr(cls, expr, dtype):
+        expr = _materialize_trunc_to_float_expr(expr, dtype)
         indexing = V.kernel.indexing(
             expr, block_ptr=False, tma_compatibility_checker=None
         )
@@ -1971,6 +2250,7 @@ def load_seed(name, offset):
         )
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def frexp(x):
         cache_key = f"frexp({x})"
         if cse_val := V.kernel.cse.try_get(cache_key):
@@ -2072,7 +2352,7 @@ def is_identity(self):
         @classmethod
         @abstractmethod
         def create(
-            cls, original_strides: list[Union[int, sympy.Expr]], shape_env: ShapeEnv
+            cls, original_strides: list[int | sympy.Expr], shape_env: ShapeEnv
         ) -> BlockParameters.StrideSorter:
             """Create a `StrideSorter` that can be used to sort block parameters."""
 
@@ -2093,7 +2373,7 @@ def __post_init__(self):
 
         @classmethod
         def create(
-            cls, original_strides: list[Union[int, sympy.Expr]], shape_env: ShapeEnv
+            cls, original_strides: list[int | sympy.Expr], shape_env: ShapeEnv
         ) -> BlockParameters.StrideSorter:
             return cls(
                 original_strides=original_strides,
@@ -2111,7 +2391,7 @@ def __post_init__(self):
 
         @classmethod
         def create(
-            cls, original_strides: list[Union[int, sympy.Expr]], shape_env: ShapeEnv
+            cls, original_strides: list[int | sympy.Expr], shape_env: ShapeEnv
         ) -> BlockParameters.StrideSorter:
             """
             If the strides are not all known constants or if the strides are already
@@ -2227,13 +2507,13 @@ def __contains__(self, item):
         return item in self.config
 
 
-class TritonCSE(CSE[TritonCSEVariable, Union[str, tuple[str, str]]]):
+class TritonCSE(CSE[TritonCSEVariable, str | tuple[str, str]]):
     """
     Subclasses CSE to apply the current load mask to the cache key to avoid CSEing
     variables across separate masked blocks.
     """
 
-    def augment_key(self, cache_key: str) -> Union[str, tuple[str, str]]:
+    def augment_key(self, cache_key: str) -> str | tuple[str, str]:
         if mask := V.kernel._load_mask:
             return (cache_key, mask.name)
         else:
@@ -2306,7 +2586,8 @@ def are_block_parameters_compatible(
         """
         if self.force:
             strides = [
-                V.graph.sizevars.symbolic_hint(st) for st in block_params.strides
+                V.graph.sizevars.replace_backed_symbols_with_hints(st)
+                for st in block_params.strides
             ]
         else:
             strides = block_params.strides
@@ -2422,7 +2703,7 @@ def are_block_parameters_compatible(
                 def indexing_div_rep(
                     x: sympy.Expr,
                     y: sympy.Expr,
-                    z: Optional[sympy.Expr] = None,
+                    z: sympy.Expr | None = None,
                 ) -> sympy.Expr:
                     div = x / y
                     if z:
@@ -2514,15 +2795,15 @@ class TritonKernel(SIMDKernel[TritonCSEVariable]):
     tensor_descriptor_options_cls: type[TensorDescriptorOptions] = (
         TensorDescriptorOptions
     )
-    transpose_discontiguous_tensor_descriptors_override: Optional[bool] = None
+    transpose_discontiguous_tensor_descriptors_override: bool | None = None
 
     def __init__(
         self,
         tiling: dict[str, sympy.Expr],
         min_elem_per_thread=0,
         optimize_mask=True,
-        fixed_config: Optional[FixedTritonConfig] = None,
-        hint_override: Optional[int] = None,
+        fixed_config: FixedTritonConfig | None = None,
+        hint_override: int | None = None,
         is_combo_kernel: bool = False,
         **kwargs,
     ) -> None:
@@ -2552,7 +2833,7 @@ def __init__(
 
         # A set of autotuning hints to pass as part of triton_meta
         self.autotune_hints = OrderedSet[AutotuneHint]()
-        self.triton_meta: Optional[dict[str, Any]] = None
+        self.triton_meta: dict[str, Any] | None = None
 
         if self.inside_reduction:
             self.codegen_reduction_numels(self.body)
@@ -2724,11 +3005,11 @@ def indexing(
         self,
         index: sympy.Expr,
         *,
-        copy_shape: Optional[Union[str, tuple[str]]] = None,
+        copy_shape: str | tuple[str] | None = None,
         dense_indexing=False,
         override_mask=None,
         block_ptr=False,
-        tma_compatibility_checker: Optional[TMACompatibilityChecker] = None,
+        tma_compatibility_checker: TMACompatibilityChecker | None = None,
     ):
         """
         Compute the index and mask to pass to tl.load() or tl.store()
@@ -2809,7 +3090,7 @@ def indexing(
 
             def match_affine_block(
                 index: sympy.Expr, range_tree: IterationRangesRoot
-            ) -> Optional[BlockParameters]:
+            ) -> BlockParameters | None:
                 """
                 Matches expressions of the form:
                     idx = s * xindex
@@ -2831,7 +3112,7 @@ def match_affine_block(
 
             def match_mod_div_block(
                 index: sympy.Expr, range_tree: IterationRangesRoot
-            ) -> Optional[BlockParameters]:
+            ) -> BlockParameters | None:
                 """
                 Matches higher-dimensional blocks coming from FloorDiv and ModularIndexing.
 
@@ -2954,7 +3235,7 @@ def match_mod_div_block(
             def match_block_subexpr(
                 expr: sympy.Expr,
                 range_tree: IterationRangesRoot,
-            ) -> Optional[BlockParameters]:
+            ) -> BlockParameters | None:
                 """
                 Match a block indexing subexpression involving a single range tree.
                 """
@@ -2971,7 +3252,7 @@ def match_block_subexpr(
 
                 return None
 
-            def match_block_expr() -> Optional[BlockDescriptorOptions]:
+            def match_block_expr() -> BlockDescriptorOptions | None:
                 index_relative_to_xyr_index = sympy_subs(
                     index, {v: t.expr for v, t in self.range_tree_nodes.items()}
                 )
@@ -3208,7 +3489,7 @@ def codegen_block_ptr(
         self,
         name: str,
         var: str,
-        indexing: Union[BlockPtrOptions, TensorDescriptorOptions],
+        indexing: BlockPtrOptions | TensorDescriptorOptions,
         other="",
     ) -> tuple[str, str]:
         """Generate a block pointer or tensor descriptor for Triton kernel operations.
@@ -3332,7 +3613,15 @@ def codegen_block_ptr_store_line(self, name, indexing, block_ptr, value, other="
         )
 
         # workaround https://github.com/triton-lang/triton/issues/2814
-        value = f"{value}.to({triton_store_type(V.graph.get_dtype(name))})"
+        # For inplace-mutated buffers, the block pointer element type comes from
+        # the actual tensor (input buffer), which may differ from the graph dtype
+        # (e.g., _to_copy produces fp32 values stored into a bf16 gradient buffer).
+        store_dtype = V.graph.get_dtype(name)
+        if name in self.args.inplace_buffers:
+            buf = self.args.inplace_buffers[name]
+            if not isinstance(buf, RemovedArg):
+                store_dtype = V.graph.get_dtype(buf.other_names[0])
+        value = f"{value}.to({triton_store_type(store_dtype)})"
         if isinstance(indexing, BlockPtrOptions):
             return f"tl.store({block_ptr}, {value}{other})"
         return f"{block_ptr}.store({V.kernel.index_to_str(indexing.offsets)}, {value})"
@@ -3472,7 +3761,7 @@ def load(self, name: str, index: sympy.Expr):
         var = self.args.input(name)
         load_counts = self._load_counts
         load_counts[name] += 1
-        make_line: Callable[[str], Union[str, DelayReplaceLine]] = identity
+        make_line: Callable[[str], str | DelayReplaceLine] = identity
         indirect_indexing = self.is_indirect_indexing(index)
         original_index = index
         dtype = V.graph.get_dtype(name)
@@ -3542,7 +3831,10 @@ def decide_later():
         has_read_deps = True
         if config.triton.skip_l1_cache:
             buffer_read_counts = self.features.buffer_read_counts()
-            has_read_deps = buffer_read_counts[name] > 1
+            # Graph inputs, primals_*, arg*_* would not be tracked by `buffer_read_counts`
+            # and it'd be fair to expect them to be reused.
+            if name in buffer_read_counts:
+                has_read_deps = buffer_read_counts[name] > 1
         """Skip L1 cache if we're (pretty?) sure the data is used only once
         """
         skip_l1_cache = (
@@ -3757,7 +4049,7 @@ def guard_cooperative_store(self, name, buffer):
         buffer.writeline(DeferredLine(name, f"if rsplit_id == ({idx} % RSPLIT):"))
         return buffer.indent()
 
-    def _combine_masks(self, *variables: Optional[CSEVariable]):
+    def _combine_masks(self, *variables: CSEVariable | None):
         masks = None
         for elem in variables:
             if elem is None:
@@ -3776,8 +4068,8 @@ def bucketize(
         boundary_indices: CSEVariable,
         indexing_dtype: torch.dtype,
         right: bool,
-        sorter: Optional[tuple[str, sympy.Expr]] = None,
-        sorter_indices: Optional[CSEVariable] = None,
+        sorter: tuple[str, sympy.Expr] | None = None,
+        sorter_indices: CSEVariable | None = None,
     ) -> CSEVariable:
         """
         See [Note: Inductor bucketize op]
@@ -3874,32 +4166,33 @@ def reduction(
         dtype: torch.dtype,
         src_dtype: torch.dtype,
         reduction_type: ReductionType,
-        value: Union[CSEVariable, tuple[CSEVariable, ...]],
-    ) -> Union[CSEVariable, tuple[CSEVariable, ...]]:
+        value: CSEVariable | tuple[CSEVariable, ...],
+    ) -> CSEVariable | tuple[CSEVariable, ...]:
         """
         codegen reduction of value to Triton according the reduction_type
         """
 
+        def should_upcast(d: torch.dtype | None) -> bool:
+            return d is not None and d.is_floating_point and d.itemsize < 4
+
         def maybe_upcast(value: CSEVariable) -> CSEVariable:
-            # Math reductions in FP16/BF16 are less accurate because the Triton compiler does not
-            # automatically promote to FP32 for accumulation. Additionally, max/min reductions
-            # do not support FP16/BF16. We manually promote to FP32 here.
+            # Math reductions in small floats are less accurate because the Triton
+            # compiler does not automatically promote to FP32 for accumulation.
+            # Additionally, max/min reductions do not support FP16/BF16.  Manually
+            # promote to FP32 here.
             return (
                 ops.to_dtype(value, torch.float32)
-                if value.dtype
-                in [
-                    torch.float16,
-                    torch.bfloat16,
-                ]
+                if should_upcast(value.dtype)
                 else value
             )
 
-        original_dtypes = [val.dtype for val in pytree.tree_leaves(value)]
-        value = pytree.tree_map(maybe_upcast, value)
-        if any(x in [torch.float16, torch.bfloat16] for x in original_dtypes):
+        do_upcast = pytree.tree_any(lambda v: should_upcast(v.dtype), value)
+        original_dtype = dtype
+        if do_upcast:
             # Only promote FB16/BF16; do not promote other integer/boolean dtypes
-            src_dtype = torch.promote_types(src_dtype, torch.float32)
-            dtype = torch.promote_types(dtype, torch.float32)
+            value = pytree.tree_map(maybe_upcast, value)
+            src_dtype = torch.float32 if should_upcast(src_dtype) else src_dtype
+            dtype = torch.float32 if should_upcast(dtype) else dtype
 
         assert self.inside_reduction
         masks = OrderedSet(f"{tree.prefix}mask" for tree in self.range_trees)
@@ -3954,8 +4247,8 @@ def maybe_upcast(value: CSEVariable) -> CSEVariable:
         def final_reduction(
             buffer,
             value: CSEVariable,
-            result_type: Optional[torch.dtype],
-        ) -> tuple[str, Optional[torch.dtype], BlockShapeType]:
+            result_type: torch.dtype | None,
+        ) -> tuple[str, torch.dtype | None, BlockShapeType]:
             """
             Helper to generate a reduction call, e.g. tl.sum.
             """
@@ -3988,7 +4281,7 @@ def final_reduction_define(
             buffer,
             result_var: CSEVariable,
             value: CSEVariable,
-            result_type: Optional[torch.dtype],
+            result_type: torch.dtype | None,
         ) -> None:
             """
             Generate a reduction and assign it to an existing variable.
@@ -4060,7 +4353,7 @@ def _mask_value(value, default) -> CSEVariable:
                     shape=value.shape,
                 )
 
-            masked_value: Union[CSEVariable, Sequence[CSEVariable], None]
+            masked_value: CSEVariable | Sequence[CSEVariable] | None
             if reduction_type == "online_softmax_reduce":
                 # Don't generate mask value for online_softmax since we
                 # will fallback below
@@ -4331,32 +4624,18 @@ def _mask_value(value, default) -> CSEVariable:
 
         self.cse.reduction_cache[cache_key] = result_var
 
-        if isinstance(result_var, tuple):
-            assert all(isinstance(x, TritonCSEVariable) for x in result_var)
-            self.outside_loop_vars.update(result_var)
-
-            # Match output dtype with input dtype
-            if reduction_type in ("welford_reduce", "online_softmax_reduce"):
-                assert len(original_dtypes) == 1
-                original_dtypes = len(result_var) * original_dtypes
+        result_tuple = result_var if isinstance(result_var, tuple) else (result_var,)
+        self.outside_loop_vars.update(result_tuple)
+        assert all(isinstance(x, TritonCSEVariable) for x in result_tuple)
 
-            assert len(result_var) == len(original_dtypes)
-            for var, orig_dtype in zip(result_var, original_dtypes):
-                assert orig_dtype is not None
-                if var.dtype != orig_dtype:
+        # If BF16/F16 upcasting was done, ensure the output is downcast to the
+        # expected dtype.
+        if do_upcast:
+            for result in result_tuple:
+                if result.dtype != original_dtype:
                     self.post_loop_combine.writeline(
-                        f"{var} = {var}.to({triton_compute_type(orig_dtype)})"
+                        f"{result} = {result}.to({triton_compute_type(original_dtype)})"
                     )
-        else:
-            assert isinstance(result_var, TritonCSEVariable)
-            self.outside_loop_vars.add(result_var)
-
-            # Match output dtype with input dtype
-            if result_var.dtype != original_dtypes[0]:
-                assert original_dtypes[0] is not None
-                self.post_loop_combine.writeline(
-                    f"{result_var} = {result_var}.to({triton_compute_type(original_dtypes[0])})"
-                )
 
         return result_var
 
@@ -5018,15 +5297,19 @@ def codegen_body(self):
                         # Subtract any advancements made in the previous loop level.
                         if level < len(loop_trees) - 1:
                             prev_tree = loop_trees[level + 1]
-                            prev_advancement = self.pointer_advancements[
+                            prev_advancements = self.pointer_advancements[
                                 prev_tree.symt
-                            ][block_ptr]
-                            prev_block = TritonSymbols.get_block_size(prev_tree)
-                            prev_num_iter = CeilDiv(prev_tree.numel, prev_block)
-                            advancement = [
-                                cur - prev * prev_num_iter
-                                for cur, prev in zip(advancement, prev_advancement)
                             ]
+                            # block_ptr may not exist in the inner loop's advancements
+                            # if its advancement was identity (zero) and was skipped
+                            if block_ptr in prev_advancements:
+                                prev_advancement = prev_advancements[block_ptr]
+                                prev_block = TritonSymbols.get_block_size(prev_tree)
+                                prev_num_iter = CeilDiv(prev_tree.numel, prev_block)
+                                advancement = [
+                                    cur - prev * prev_num_iter
+                                    for cur, prev in zip(advancement, prev_advancement)
+                                ]
 
                         self.body.writeline(
                             DeferredLine(
@@ -5089,7 +5372,7 @@ def kernel_benchmark_extra_args(self) -> list[str]:
                     raise ValueError(f"Unsupported numel argument type: {type(arg)}")
         return args
 
-    def codegen_kernel_benchmark(self, num_gb: Optional[float]) -> IndentedBuffer:
+    def codegen_kernel_benchmark(self, num_gb: float | None) -> IndentedBuffer:
         """
         Generates Python code for benchmarking this Triton kernel.
         - Creates example inputs (random tensors, constants, sizes).
@@ -5173,7 +5456,7 @@ def codegen_kernel_benchmark(self, num_gb: Optional[float]) -> IndentedBuffer:
                 result.writeline(
                     V.graph.device_ops.set_device(index)
                 )  # no-op to ensure context
-                stream_name = f"stream{index}"
+                stream_name = get_raw_stream_name(index)
                 result.writeline(f"{stream_name} = get_raw_stream({index})")
                 result.writeline(
                     f"{str(Placeholder.KERNEL_NAME)}.run(*args, stream={stream_name})"
@@ -5297,6 +5580,8 @@ def inductor_meta_common(cls):
             "store_cubin": config.triton.store_cubin,
             "deterministic": config.deterministic,
             "force_filter_reduction_configs": config.test_configs.force_filter_reduction_configs,
+            "mix_order_reduction_allow_multi_stages": config.triton.mix_order_reduction_allow_multi_stages,
+            "dynamic_disable_pipelining": config.triton.dynamic_disable_pipelining,
         }
 
         if config.write_are_deterministic_algorithms_enabled:
@@ -5340,22 +5625,8 @@ def codegen_kernel(self, name=None) -> str:
             if prefix_is_reduction(prefix) and not self.inside_reduction:
                 continue
 
-            numel_hint = V.graph.sizevars.symbolic_hint(numel)
-            if not isinstance(numel_hint, (int, sympy.Integer)):
-                # This default heuristic hint was picked carefully: it is
-                # large, to ensure that we don't shrink the block size (since
-                # if you don't have many elements, it'd be wasteful to pick a
-                # large block size).  Since we don't know how many elements we
-                # might have, we should be OK with some inefficiency to make
-                # sure we handle the large case well.  8192 is the largest
-                # block size we support, so we pick that.
-                #
-                # If we have a better hint for unbacked SymInts (e.g., because
-                # a user told us, or we are tracking upper bounds) we could
-                # use that here.
-                size_hint = 8192
-            else:
-                size_hint = next_power_of_2(int(numel_hint))
+            numel_hint = V.graph.sizevars.optimization_hint(numel)
+            size_hint = next_power_of_2(int(numel_hint))
             size_hints[prefix] = size_hint
 
         if name is None:
@@ -5568,8 +5839,6 @@ def add_constexpr_arg(arg_name):
             if flops is not None:
                 inductor_meta["kernel_flop"] = flops
 
-        triton_meta["configs"] = [config_of(signature)]
-
         # Triton compiler includes equal_to_1 args into constants even
         # when they are not constexpr. otherwise there may be a segfault
         # during launching the Inductor-compiled Triton kernel.
@@ -5585,6 +5854,14 @@ def add_constexpr_arg(arg_name):
         self.codegen_body()
         self._filter_pdl(self.body)
 
+        # Compute configs after codegen_body() so we know if the kernel
+        # uses atomic ops. On HIP, buffer ops don't support atomics, so
+        # we must not tag any args with pointer_range_32 in that case.
+        if torch.version.hip is not None and self.atomic_add_found:
+            triton_meta["configs"] = [config_of(signature, pointer_range_override=())]
+        else:
+            triton_meta["configs"] = [config_of(signature)]
+
         for helper in self.helper_functions:
             code.writeline("")
             code.splice(helper)
@@ -5750,7 +6027,7 @@ def add_numel_to_call_args(self, name, call_args, arg_types):
                 arg_types.append(type(expr))
 
     def call_kernel(
-        self, name: str, node: Optional[IRNode] = None, deallocate_ws: bool = True
+        self, name: str, node: IRNode | None = None, deallocate_ws: bool = True
     ):
         wrapper = V.graph.wrapper_code
         wrapper.write_triton_header_once()
@@ -5885,12 +6162,22 @@ def _has_constant_mask(self, tree: IterationRangesRoot) -> bool:
                 return True
         elif not self.is_combo_kernel:
             if V.graph.sizevars.statically_known_equals(tree.numel, 1):
-                return True
+                if not (tree.is_reduction and self.persistent_reduction):
+                    return True
 
         # Masks are superfluous if numel is a multiple of BLOCK
         # (We use the fact that BLOCK is required by triton to be a power of 2)
         if tree.is_reduction and self.persistent_reduction:
             max_block = self._get_persistent_RBLOCK(tree.numel)
+            # Triton's auto-tuner can map a full hardware warp along the
+            # reduction axis.  When RBLOCK < warp_size the excess lanes
+            # would execute out-of-bounds global loads.  This results in
+            # faults on AMD hardware.  Keep the dynamic mask so that all
+            # hardware stays correct.
+            device = V.graph.get_current_device_or_throw()
+            warp_size = get_warp_size(device)
+            if isinstance(max_block, int) and max_block < warp_size:
+                return False
         elif tree.prefix == "x" and self.no_x_dim:
             max_block = 1
         else:
@@ -6032,6 +6319,139 @@ def iteration_ranges_codegen_header(
             code.writeline(f"{x}mask = {entry.name} < {x}numel")
 
 
+class FusedUserDefinedTritonKernel(TritonKernel):
+    """
+    When fusing a user-defined triton kernel with epilogues, we use this class to generate the modified triton kernel source
+    """
+
+    def __init__(
+        self,
+        tiling: dict[str, sympy.Expr],
+        features: SIMDKernelFeatures,
+        scheduler_node: FusedExternTritonKernelSchedulerNode,
+    ) -> None:
+        super().__init__(
+            tiling,
+            features=features,
+            min_elem_per_thread=0,
+            optimize_mask=True,
+            fixed_config=None,
+            hint_override=None,
+            is_combo_kernel=False,
+        )
+        self.scheduler_node = scheduler_node
+        assert isinstance(
+            self.scheduler_node.kernel_node.node, ir.UserDefinedTritonKernel
+        )
+        self.ir_node: ir.UserDefinedTritonKernel = self.scheduler_node.kernel_node.node
+        assert self.ir_node.can_fuse_epilogue()
+
+        # must be true because `self.ir_node.can_fuse_epilogue()`
+        assert len(self.ir_node.kernel_stores.stores) == 1
+        self.original_stored_expr = ast.unparse(
+            self.ir_node.kernel_stores.stores[0].store_value_node
+        )
+        self.original_stored_expr = self.original_stored_expr.replace("\n", "")
+
+    def load(self, name: str, index: sympy.Expr) -> TritonCSEVariable:
+        # must be true because `self.ir_node.can_fuse_epilogue()`
+        assert len(self.ir_node.mutable_args) == 1
+        if name == self.ir_node.mutable_args[0].get_name():
+            # when the fused epilogue nodes tries to load the mutated buffer
+            # we replace the load with the expression stored by the user kernel
+            loaded_expr = self.original_stored_expr
+            indexing = self.indexing(index)
+            dtype = V.graph.get_dtype(name)
+            if indexing.expand_shape:
+                shape = indexing.expand_shape
+            else:
+                shape = TritonSymbols.get_block_shape(indexing.index)
+            result_var = self.cse.generate(
+                self.loads, loaded_expr, dtype=dtype, shape=shape
+            )
+            return result_var
+        else:
+            # The scheduler should prevent this.
+            raise AssertionError(
+                f"Epilogue attempted to load from '{name}'. "
+                "Inductor indexing variables are not defined in user kernel scope. "
+            )
+
+    def store(
+        self, name: str, index: sympy.Expr, value: CSEVariable, mode: StoreMode = None
+    ) -> None:
+        assert isinstance(self.scheduler_node.fused_epilogue.node, ir.ComputedBuffer)
+        if name == self.scheduler_node.fused_epilogue.node.get_name():
+            # when the fused epilogue nodes tries to store to its destination buffer
+            # we remember this expr and then later replace it into the `tl.store` call of the original kernel
+            self.new_store_cse_var = value
+        else:
+            super().store(name, index, value, mode)
+
+    # returns a str which is the src code of a modified version of the user kernel that includes the epilogues
+    def codegen(self) -> str:
+        with self:
+            index_vars = self.split_and_set_ranges(
+                self.scheduler_node.fused_epilogue.get_ranges()
+            )
+            self.scheduler_node.fused_epilogue.codegen(index_vars)
+
+        # Generate a new AST where the store value expr is replaced with the new value
+        new_ast = copy.deepcopy(self.ir_node.kernel_ast)
+        from torch._higher_order_ops.triton_kernel_wrap import (
+            identify_triton_stores,
+            identify_triton_stores_from_ast,
+        )
+
+        # avoid redundant cache entry of new_ast
+        kernel_stores = identify_triton_stores_from_ast(new_ast)
+        assert len(kernel_stores.stores) == 1
+
+        new_store_value_node = ast.Name(self.new_store_cse_var.name)
+
+        def _replace_arg(
+            call_node: ast.Call,
+            arg_name: str,
+            positional_index: int,
+            new_arg,
+        ):
+            if len(call_node.args) > positional_index:
+                call_node.args[positional_index] = new_arg
+
+            # Check keyword args
+            for keyword in call_node.keywords:
+                if keyword.arg == arg_name:
+                    keyword.value = new_arg
+
+        _replace_arg(
+            kernel_stores.stores[0].store_node, "value", 1, new_store_value_node
+        )
+
+        src_with_store_replaced = ast.unparse(new_ast)
+
+        # then, we need to inject the additional `load` and `compute` lines generated when we `codegen` the epilogue nodes
+
+        src_lines = src_with_store_replaced.splitlines()
+
+        # identify the store again, because the previous parse-modify-unparse could've change its location
+        kernel_stores = identify_triton_stores(src_with_store_replaced)
+        # python ast lineno is 1-indexed
+        store_line_index = kernel_stores.stores[0].store_node.lineno - 1
+
+        indentations = " " * kernel_stores.stores[0].store_node.col_offset
+
+        load_lines = [indentations + l for l in self.loads.get_lines_ref()]
+        compute_lines = [indentations + l for l in self.compute.get_lines_ref()]
+
+        new_src_lines = (
+            src_lines[:store_line_index]
+            + load_lines
+            + compute_lines
+            + src_lines[store_line_index:]
+        )
+        return "\n".join(new_src_lines)
+
+
 class TritonScheduling(SIMDScheduling):
     """Scheduling backend for Triton kernel code generation."""
 
@@ -6049,7 +6469,7 @@ class TritonScheduling(SIMDScheduling):
         ]
     )
 
-    def __init__(self, scheduler: Optional[Scheduler]) -> None:
+    def __init__(self, scheduler: Scheduler | None) -> None:
         super().__init__(scheduler)
         if scheduler is None or not hasattr(scheduler, "nodes"):
             return
@@ -6153,6 +6573,8 @@ def define_kernel(self, src_code, node_schedule, kernel):
                 if config.triton.descriptive_names
                 else ""
             )
+            if fused_name:
+                fused_name = V.choices.customize_fused_kernel_name(fused_name, src_code)
             kernel_category = get_kernel_category_by_source_code(src_code)[:3]
             kernel_name = "_".join(
                 ["triton", kernel_category, fused_name, wrapper.next_kernel_suffix()]
@@ -6209,7 +6631,7 @@ def benchmark_fused_nodes(self, nodes, n_spills_threshold=8) -> tuple[float, str
         )
 
     def benchmark_codegened_module(
-        self, mod, n_spills_threshold=8, node_names: Optional[OrderedSet[str]] = None
+        self, mod, n_spills_threshold=8, node_names: OrderedSet[str] | None = None
     ) -> tuple[float, str]:
         """Benchmark an already compiled module"""
         device_interface = get_interface_for_device(V.graph.device_type)
@@ -6512,11 +6934,16 @@ def debug_triton_code(node: BaseSchedulerNode) -> list[str]:
         from torch._inductor.codegen.cuda_combined_scheduling import (
             CUDACombinedScheduling,
         )
+        from torch._inductor.codegen.xpu.xpu_combined_scheduling import (
+            XPUCombinedScheduling,
+        )
 
         device = node.get_device()
         assert device is not None
         backend = node.scheduler.get_backend(device)
-        assert isinstance(backend, (SIMDScheduling, CUDACombinedScheduling)), (
+        assert isinstance(
+            backend, (SIMDScheduling, CUDACombinedScheduling, XPUCombinedScheduling)
+        ), (
             f"Scheduling backend should be SIMD or CUDACombined when generating debug Triton strings, got: {type(backend)}"
         )
 
diff --git a/torch/_inductor/codegen/triton_combo_kernel.py b/torch/_inductor/codegen/triton_combo_kernel.py
index 3c00e74498760..1c21a1594b017 100644
--- a/torch/_inductor/codegen/triton_combo_kernel.py
+++ b/torch/_inductor/codegen/triton_combo_kernel.py
@@ -4,7 +4,7 @@
 from collections import defaultdict
 from collections.abc import Callable
 from dataclasses import dataclass
-from typing import Any, cast, Optional, Union
+from typing import Any, cast
 
 import sympy
 from sympy import Integer, Symbol
@@ -20,6 +20,7 @@
     SequentialFlattenComboKernelGrid,
 )
 from ..scheduler import BaseSchedulerNode
+from ..stream_utils import get_raw_stream_name
 from ..utils import Placeholder, triton_version_uses_attrs_dict
 from ..virtualized import V
 from .common import (
@@ -40,6 +41,11 @@
 from .triton_utils import config_of, equal_1_arg_indices, signature_to_meta
 
 
+# Default block sizes used when combo kernel autotuning is disabled.
+DEFAULT_COMBO_BLOCK_SIZE_1D = 1024
+DEFAULT_COMBO_BLOCK_SIZE_2D = 32
+
+
 log = logging.getLogger(__name__)
 pexpr = PythonPrinter().doprint
 LARGE_NUMELS = 512e5
@@ -84,13 +90,12 @@ def _default_custom_combo_kernel_horizontal_partition(
         not_reduction = [n for n in group_per_dim if n not in reduction]
         # rnumel > 2048 usually has long execution time
         # BaseSchedulerNode.group[-1][-1] is rnumel for reduction nodes
+        # Scheduling heuristic: separate long reductions (rnumel > 2048).
+        # Uses optimization_hint with fallback=1 so unbacked defaults to short reduction.
         long_reduction = [
             n
             for n in reduction
-            if (
-                V.graph.sizevars.shape_env.has_hint(n.group[-1][-1])
-                and V.graph.sizevars.size_hint(n.group[-1][-1]) > 2048  # type: ignore[arg-type]
-            )
+            if V.graph.sizevars.optimization_hint(n.group[-1][-1], fallback=1) > 2048  # type: ignore[arg-type]
         ]
         short_reduction = [n for n in reduction if n not in long_reduction]
         if long_reduction:
@@ -103,8 +108,10 @@ def _default_custom_combo_kernel_horizontal_partition(
             for n in not_reduction
             if not node_info_map[n].features.is_reduction()
             and len(node_info_map[n].tiling) == 2
-            and V.graph.sizevars.shape_env.has_hint(node_info_map[n].tiling["x"])
-            and V.graph.sizevars.size_hint(node_info_map[n].tiling["x"]) > LARGE_NUMELS
+            and V.graph.sizevars.optimization_hint(
+                node_info_map[n].tiling["x"], fallback=1
+            )
+            > LARGE_NUMELS  # type: ignore[arg-type]
         ]
         if large_pointwise:
             # TODO benchmark the performance when large pointwise nodes combining with others
@@ -439,24 +446,23 @@ def __init__(
         self.sub_kernels: list[TritonKernel] = []
         self.iter_vars_count = itertools.count()
         self.grids: list[list[int]] = []
-        self.min_x_blocks_list: list[Union[int, str]] = []
-        self.x_numels_list: list[Union[int, str]] = []
+        self.min_x_blocks_list: list[int | str] = []
+        self.x_numels_list: list[int | str] = []
         self.y_tree_list: list = []
         self.enable_autotune = enable_autotune
         self.mixed_sizes = mixed_sizes
-        self.dispatch_class: Optional[
+        self.dispatch_class: (
             type[
-                Union[
-                    ComboKernel.SequentialDispatch,
-                    ComboKernel.SequentialFlattenGridDispatch,
-                    ComboKernel.RoundRobinDispatch,
-                ]
+                ComboKernel.SequentialDispatch
+                | ComboKernel.SequentialFlattenGridDispatch
+                | ComboKernel.RoundRobinDispatch
             ]
-        ] = None
+            | None
+        ) = None
         self.block_args: list[str] = []
-        # there following are used when autotuning is disabled
-        self.block_size_1d = 1024  # Try tuning this value
-        self.block_size_2d = 32
+        # the following are used when autotuning is disabled
+        self.block_size_1d = DEFAULT_COMBO_BLOCK_SIZE_1D
+        self.block_size_2d = DEFAULT_COMBO_BLOCK_SIZE_2D
         self.num_warps = 8
         self.block_size_reduce = 256
         self.dynamic_shape_args: list[str] = []
@@ -568,8 +574,8 @@ def min_x_blocks_sub_kernel(self, sub_kernel: TritonKernel, num: int) -> None:
         Kernels with no_x_dim being true has no tunable XBLOCK. They have a fixed number of X blocks.
         Grid calculation needs to make sure that they are assigned with enough number of blocks.
         """
-        min_x_blocks: Union[int, str] = 0
-        x_numels: Union[int, str] = 0
+        min_x_blocks: int | str = 0
+        x_numels: int | str = 0
         for tree in sub_kernel.range_trees:
             simplified_tree_numel = V.graph.sizevars.simplify(tree.numel)
             if tree.prefix == "x":
@@ -741,6 +747,22 @@ def jit_line(
         dispatch = self.dispatch_class
         assert dispatch is not None
 
+        # Compute the max persistent R0_BLOCK across sub-kernels.
+        # This is used by _reduction_configs() to avoid generating configs
+        # where XBLOCK * max_persistent_rblock creates pathologically large
+        # tiles that cause extreme ROCm compilation times.
+        # The max_persistent_rblock mirrors how R0_BLOCK is computed in
+        # codegen_static_numels_sub_kernel() for persistent reductions.
+        max_persistent_rblock = 0
+        for sub in self.sub_kernels:
+            if sub.persistent_reduction:
+                for tree in sub.range_trees:
+                    if tree.is_reduction:
+                        simplified_numel = V.graph.sizevars.simplify(tree.numel)
+                        if isinstance(simplified_numel, (Integer, int)):
+                            val = next_power_of_2(int(simplified_numel))
+                            max_persistent_rblock = max(max_persistent_rblock, val)
+
         inductor_meta = {
             "grid_type": dispatch.grid_expr.__name__,
             "combo_grid_meta": self.combo_grid_meta(size_hints_list),
@@ -748,6 +770,8 @@ def jit_line(
             "mutated_arg_names": mutated_args,
             **self.triton_kernel_cls.inductor_meta_common(),
         }
+        if max_persistent_rblock > 0:
+            inductor_meta["max_persistent_rblock"] = max_persistent_rblock
 
         sub_kernel = selected_kernel
         if heuristics == "foreach":
@@ -879,7 +903,7 @@ def kernel_benchmark_extra_args(self) -> list[str]:
                     )
         return extra_args
 
-    def codegen_kernel(self, name: Optional[str] = None) -> str:
+    def codegen_kernel(self, name: str | None = None) -> str:
         """Generate the triton code for a combo kernel that fuses multiple sub-kernels."""
         # TODO: is it correct to use the first sub kernel's heuristics?
         heuristics_list, size_hints_list = [], []
@@ -1000,7 +1024,7 @@ def codegen_kernel_benchmark(self, num_gb: float) -> IndentedBuffer:
                         f"{var_name} = rand_strided({size}, {stride}, device='{const_tensor.device}', dtype={const_tensor.dtype})"  # type: ignore[arg-type]  # noqa: B950 line too long
                     )
                 elif isinstance(arg_sig, SizeArg):
-                    symval_hint = V.graph.sizevars.size_hint(arg_sig.expr)
+                    symval_hint = V.graph.sizevars.optimization_hint(arg_sig.expr)
 
                     # Force the seed_offset to be 0 so calls to the same kernel
                     # using different seed offset will have the same benchmark harness.
@@ -1010,7 +1034,7 @@ def codegen_kernel_benchmark(self, num_gb: float) -> IndentedBuffer:
                     result.writeline(f"{var_name} = {symval_hint}")
                 elif isinstance(arg_sig, WorkspaceArg):
                     device = V.graph.get_current_device_or_throw()
-                    count = V.graph.sizevars.size_hint(arg_sig.count)
+                    count = V.graph.sizevars.optimization_hint(arg_sig.count)
                     # for benchmark harness, we ignore arg_sig.zero_mode and always zero it
                     result.writeline(
                         f"{var_name} = torch.zeros({count}, device='{device}', dtype={arg_sig.dtype})"
@@ -1033,7 +1057,7 @@ def codegen_kernel_benchmark(self, num_gb: float) -> IndentedBuffer:
                 result.writeline(
                     V.graph.device_ops.set_device(index)
                 )  # no-op to ensure context
-                stream_name = f"stream{index}"
+                stream_name = get_raw_stream_name(index)
                 result.writeline(f"{stream_name} = get_raw_stream({index})")
                 result.writeline(
                     f"{str(Placeholder.KERNEL_NAME)}.run(*args, stream={stream_name})"
@@ -1183,6 +1207,10 @@ def combo_grid_meta(self, size_hints_list: list[dict[str, int]]) -> dict[str, An
                         meta[f"tile_hint_{num}"] = "TileHint.SQUARE"
                     else:
                         meta[f"tile_hint_{num}"] = "TileHint.DEFAULT"
+                else:
+                    meta[f"reduction_hint_{num}"] = (
+                        sub_kernel.features.get_reduction_hint().name
+                    )
 
             for tree in sub_kernel.range_trees:
                 if not tree.is_reduction:
diff --git a/torch/_inductor/codegen/triton_split_scan.py b/torch/_inductor/codegen/triton_split_scan.py
index 0abee54393939..e109eab996324 100644
--- a/torch/_inductor/codegen/triton_split_scan.py
+++ b/torch/_inductor/codegen/triton_split_scan.py
@@ -1,6 +1,5 @@
 # mypy: allow-untyped-defs
 import functools
-from typing import Union
 
 import sympy
 
@@ -123,7 +122,7 @@ def scan(self, dtypes, combine_fn, values):
         )
         max_blocks = pointwise_numel * CeilDiv(reduction_numel, min_rblock)
         nbytes = scratch_nbytes_per_block * max_blocks
-        scratch_base: Union[str, TritonCSEVariable]
+        scratch_base: str | TritonCSEVariable
         scratch_base, _, offset = self.args.workspace(nelem=nbytes, zero_fill=True)
         if offset != 0:
             scratch_base = cse_load(
diff --git a/torch/_inductor/codegen/triton_utils.py b/torch/_inductor/codegen/triton_utils.py
index d38798e793512..f89c72ebfb1c4 100644
--- a/torch/_inductor/codegen/triton_utils.py
+++ b/torch/_inductor/codegen/triton_utils.py
@@ -1,5 +1,5 @@
 # mypy: allow-untyped-defs
-from typing import Any, Optional
+from typing import Any
 
 import sympy
 
@@ -32,20 +32,9 @@ def should_unwrap_unspec_arg(name: str):
     return False
 
 
-def signature_of(arg: KernelArgType, *, size_dtype: Optional[str]) -> str:
+def signature_of(arg: KernelArgType, *, size_dtype: str | None) -> str:
     if isinstance(arg, TensorArg):
-        # TODO: Remove fp8 special handling when Triton supports PyTorch fp8 dtypes.
-        # Related PR: https://github.com/triton-lang/triton/pull/2279/
-        if arg.dtype == torch.float8_e4m3fn:
-            typ = "*fp8e4nv"
-        elif arg.dtype == torch.float8_e5m2:
-            typ = "*fp8e5"
-        elif arg.dtype == torch.float8_e4m3fnuz:
-            typ = "*fp8e4b8"
-        elif arg.dtype == torch.float8_e5m2fnuz:
-            typ = "*fp8e5b16"
-        else:
-            typ = _type_of(arg.dtype)
+        typ = _type_of(arg.dtype)
         if should_unwrap_unspec_arg(arg.buffer):
             # had unwrapped 0d tensor as scalar
             new_typ = typ.lstrip("*")
@@ -125,9 +114,9 @@ def non_constexpr_signature(signature):
 def signature_to_meta(
     signature: list[KernelArgType],
     *,
-    size_dtype: Optional[str],
+    size_dtype: str | None,
     argdefs: list[ArgName],
-    indices: Optional[list[int]] = None,
+    indices: list[int] | None = None,
     is_template: bool = False,
 ) -> dict[str, str]:
     if indices is None:
@@ -160,6 +149,21 @@ def _decide_tl_dtype(arg):
     }
 
 
+def _get_buffer_layout(buf_name: str) -> "torch._inductor.ir.Layout":
+    """Get the layout for a buffer, handling both scheduler buffers and graph inputs."""
+    if V.graph.scheduler:
+        layout = V.graph.scheduler.get_buffer_layout(buf_name)
+    else:
+        buffer = V.graph.try_get_buffer(buf_name)
+        # output arg
+        if not buffer:
+            assert buf_name == V.kernel.output_node.name
+            layout = V.kernel.output_node.layout
+        else:
+            layout = buffer.get_layout()
+    return layout
+
+
 def is_unaligned_buffer(arg: TensorArg):
     buf_name = arg.buffer
     if buf_name in V.graph.unaligned_buffers:
@@ -175,17 +179,7 @@ def is_unaligned_buffer(arg: TensorArg):
         # all constants are assumed to be aligned
         return False
 
-    if V.graph.scheduler:
-        layout = V.graph.scheduler.get_buffer_layout(buf_name)
-    else:
-        buffer = V.graph.try_get_buffer(buf_name)
-        # output arg
-        if not buffer:
-            assert buf_name == V.kernel.output_node.name
-            layout = V.kernel.output_node.layout
-        else:
-            layout = buffer.get_layout()
-
+    layout = _get_buffer_layout(buf_name)
     if isinstance(layout, torch._inductor.ir.NonOwningLayout):
         return not layout.maybe_guard_aligned()
     else:
@@ -203,7 +197,7 @@ def _arg_equals_1(arg: KernelArgType) -> bool:
 def equal_1_arg_indices(
     args: list[KernelArgType],
     *,
-    indices: Optional[list[int]] = None,
+    indices: list[int] | None = None,
 ) -> tuple[int, ...]:
     if indices is None:
         indices = list(range(len(args)))
@@ -213,10 +207,36 @@ def equal_1_arg_indices(
     return equal_to_1
 
 
+def _is_tensor_within_2gb(arg: TensorArg) -> bool:
+    """Check if a tensor argument's storage is provably within 2GB.
+
+    Mirrors HIPBackend.is_within_2gb() but uses compile-time symbolic analysis
+    instead of runtime tensor inspection. This enables canonicalize_pointers to
+    decompose pointer arithmetic into (splat(base), offset) form for buffer ops.
+    """
+    MAX_BYTES = 2**31 - 1
+    try:
+        # Graph inputs aren't tracked by the scheduler; get their layout
+        # from the graph_inputs dict to avoid KeyError in get_buffer_layout.
+        if arg.buffer in V.graph.graph_inputs:
+            inp = V.graph.graph_inputs[arg.buffer]
+            if hasattr(inp, "get_layout"):
+                layout = inp.get_layout()
+            else:
+                return False
+        else:
+            layout = _get_buffer_layout(arg.buffer)
+        storage_bytes = layout.storage_size() * arg.dtype.itemsize
+        return V.graph.sizevars.statically_known_true(storage_bytes <= MAX_BYTES)
+    except Exception:
+        return False
+
+
 def config_of(
     args: list[KernelArgType],
     *,
-    indices: Optional[list[int]] = None,
+    indices: list[int] | None = None,
+    pointer_range_override: tuple[int, ...] | None = None,
 ) -> Any:
     if indices is None:
         indices = list(range(len(args)))
@@ -242,7 +262,7 @@ def is_aligned(x: KernelArgType, alignment: int, include_tensor: bool) -> bool:
                 return False
             if x.expr is None:
                 return False
-            if isinstance(x.expr, float):
+            if isinstance(x.expr, (float, bool)):
                 return False
             return V.graph.sizevars.statically_known_multiple_of(x.expr, alignment)  # type: ignore[arg-type]
         if isinstance(x, WorkspaceArg):
@@ -263,5 +283,18 @@ def is_aligned(x: KernelArgType, alignment: int, include_tensor: bool) -> bool:
 
     equal_to_1 = equal_1_arg_indices(args, indices=indices)
 
-    # pyrefly: ignore [bad-argument-type]
-    return AttrsDescriptorWrapper(divisible_by_16, equal_to_1)
+    # On AMD/HIP, tag tensor args whose storage fits in 2GB so Triton
+    # can use 32-bit pointer offsets and emit buffer load/store ops.
+    if pointer_range_override is not None:
+        pointer_range_32 = pointer_range_override
+    elif torch.version.hip is not None:
+        pointer_range_32 = tuple(
+            i
+            for i, arg in zip(indices, args)
+            if isinstance(arg, TensorArg) and _is_tensor_within_2gb(arg)
+        )
+    else:
+        pointer_range_32 = ()
+
+    # pyrefly: ignore [bad-argument-count, bad-argument-type]
+    return AttrsDescriptorWrapper(divisible_by_16, equal_to_1, pointer_range_32)
diff --git a/torch/_inductor/codegen/wrapper.py b/torch/_inductor/codegen/wrapper.py
index 5ba306bda47b6..5ec3c9589d7a0 100644
--- a/torch/_inductor/codegen/wrapper.py
+++ b/torch/_inductor/codegen/wrapper.py
@@ -15,7 +15,7 @@
 import tempfile
 from collections.abc import Callable
 from itertools import chain, count
-from typing import Any, Optional, TYPE_CHECKING, Union
+from typing import Any, TYPE_CHECKING
 
 import sympy
 from sympy import Expr
@@ -47,6 +47,8 @@
 from ..ir import IRNode, ReinterpretView
 from ..runtime import triton_heuristics
 from ..runtime.hints import DeviceProperties
+from ..stream_constants import DEFAULT_STREAM, DEFAULT_STREAM_IDX, STREAM_NAME_TEMPLATE
+from ..stream_utils import get_raw_stream_name, get_stream_name
 from ..utils import (
     cache_on_self,
     DelayReplaceLine,
@@ -76,7 +78,7 @@
 
 
 if TYPE_CHECKING:
-    from collections.abc import Iterator, Sequence
+    from collections.abc import Iterable, Iterator, Sequence
 
     import triton
 
@@ -91,15 +93,16 @@
 pexpr = PythonPrinter().doprint
 
 
-ReuseKey = tuple[torch.device, torch.dtype, str, bool]
+ReuseKey = tuple[torch.device, torch.dtype, str, bool, int]
 CommBufferReuseKey = tuple[torch.device, torch.dtype, str, "ir.CommBufferType", str]
-BufferLike = Union[ir.Buffer, WorkspaceArg]
+BufferLike = ir.Buffer | WorkspaceArg
 FxConversionFunc = Callable[["WrapperLine"], None]
 
 
 def buffer_reuse_key(node: BufferLike) -> ReuseKey:
     storage_size = V.graph.get_allocation_storage_size(node)
     alignment = node.get_name() not in V.graph.unaligned_buffers
+    stream = V.graph.scheduler.get_buf_stream(node.get_name())
     return (
         node.get_device_or_error(),
         node.get_dtype(),
@@ -108,6 +111,7 @@ def buffer_reuse_key(node: BufferLike) -> ReuseKey:
         # size hint
         sympy_str(V.graph.sizevars.simplify(storage_size)),
         alignment,
+        stream,
     )
 
 
@@ -192,26 +196,26 @@ def codegen_reinterpret_view_helper(data):
 
 # TODO: Move to a well known place
 TritonMetaParams = dict[str, int]
-TritonGrid = Union[
-    tuple[Union[int, sympy.Expr], ...], Callable[[TritonMetaParams], tuple[int, ...]]
-]
+TritonGrid = (
+    tuple[int | sympy.Expr, ...] | Callable[[TritonMetaParams], tuple[int, ...]]
+)
 
 
 def user_defined_kernel_grid_fn_code(
     name: str,
     configs: list[triton.Config],  # type: ignore[name-defined]
     grids: list[TritonGrid],
-    wrapper: Optional[PythonWrapperCodegen] = None,
-    original_fxnode_name: Optional[str] = None,
+    wrapper: PythonWrapperCodegen | None = None,
+    original_fxnode_name: str | None = None,
 ) -> tuple[str, str]:
     output = IndentedBuffer()
 
-    def _convert_to_sympy_expr(item: Union[int, sympy.Expr]) -> sympy.Expr:
+    def _convert_to_sympy_expr(item: int | sympy.Expr) -> sympy.Expr:
         return item if isinstance(item, sympy.Expr) else sympy.Integer(item)
 
     def determine_grid(
         grid: TritonGrid,
-        example_grid: Optional[TritonGrid] = None,
+        example_grid: TritonGrid | None = None,
     ):
         """
         This function return a tuple of two values: the first one is for the real grid
@@ -240,7 +244,7 @@ def determine_grid(
             ),
         )
 
-    def writeline(line: str, example_grid: Optional[str] = None):
+    def writeline(line: str, example_grid: str | None = None):
         output.writeline(line)
         if (
             wrapper
@@ -307,18 +311,26 @@ def writeline(line: str, example_grid: Optional[str] = None):
     return fn_name, output.getvalue()
 
 
-def user_defined_triton_kernel_transitive_closure_source_code(kernel) -> str:
+def user_defined_triton_kernel_transitive_closure_source_code(
+    kernel, epilogue_fusion: tuple[ir.ComputedBuffer, str] | None = None
+) -> str:
     """
     Given a triton kernel function pointer collect the transitive closure of
     its dependencies
+
+    epilogue_fusion: Optional[(fused epilogue node, modified kerel src code)]
     """
     compile_wrapper = IndentedBuffer()
-    compile_wrapper.splice(kernel.src, strip=True)
+    kernel_src = kernel.src
+    if epilogue_fusion:
+        kernel_src = epilogue_fusion[1]
+    compile_wrapper.splice(kernel_src, strip=True)
 
     # Also include any possible kernel being called indirectly
     import triton
     from triton import JITFunction  # type: ignore[name-defined, attr-defined]
     from triton.language import constexpr  # type: ignore[name-defined]
+    from triton.language.core import dtype as triton_dtype
 
     # global constexpr vars handled above
     symbols_included = OrderedSet([kernel.__name__])
@@ -342,7 +354,6 @@ def traverse(cur_kernel):
                 if isinstance(symbol, JITFunction):
                     compile_wrapper.newline()
                     compile_wrapper.writeline("@triton.jit")
-
                     compile_wrapper.splice(symbol.src, strip=True)
                     symbols_included.add(symbol_name)
                     traverse(symbol)
@@ -350,9 +361,25 @@ def traverse(cur_kernel):
                     symbol,
                     triton.runtime.jit.ConstexprFunction,
                 ):
+                    # Import dtype class if used in type annotations
+                    if "dtype" in symbol.src and "dtype" not in symbols_included:
+                        dtype_symbol = symbol.fn.__globals__.get("dtype")
+                        if (
+                            dtype_symbol
+                            and hasattr(dtype_symbol, "__module__")
+                            and dtype_symbol.__module__.startswith("triton")
+                        ):
+                            compile_wrapper.writeline(
+                                f"from {dtype_symbol.__module__} import dtype as dtype"
+                            )
+                            symbols_included.add("dtype")
                     compile_wrapper.newline()
                     compile_wrapper.writeline("@triton.constexpr_function")
                     compile_wrapper.splice(symbol.src, strip=True)
+                    if symbol_name != symbol.fn.__name__:
+                        compile_wrapper.writeline(
+                            f"{symbol_name} = {symbol.fn.__name__}"
+                        )
                     symbols_included.add(symbol_name)
                     traverse(symbol)
                 elif isinstance(symbol, (int, str, bool, constexpr)):
@@ -386,9 +413,14 @@ def traverse(cur_kernel):
                     # a global symbol imported from triton is referenced
                     # without module qualification (i.e., `store` instead
                     # of `tl.store`): need to codegen an import
-                    compile_wrapper.writeline(
-                        f"from {symbol.__module__} import {symbol.__name__} as {symbol_name}"
-                    )
+
+                    # Triton dtype instances have .name instead of .__name__
+                    if isinstance(symbol, triton_dtype):
+                        compile_wrapper.writeline(f"{symbol_name} = tl.{symbol.name}")
+                    elif hasattr(symbol, "__name__"):
+                        compile_wrapper.writeline(
+                            f"from {symbol.__module__} import {symbol.__name__} as {symbol_name}"
+                        )
                     symbols_included.add(symbol_name)
 
     traverse(kernel)
@@ -522,7 +554,7 @@ def codegen_fx(self, converter: FxConverter) -> FxConversionFunc:
 @dataclasses.dataclass
 class EnterDeviceContextManagerLine(WrapperLine):
     device_idx: int
-    last_seen_device_guard_index: Optional[int]
+    last_seen_device_guard_index: int | None
 
     def codegen(self, code: IndentedBuffer) -> None:
         if V.graph.cpp_wrapper:
@@ -611,10 +643,38 @@ def codegen_fx(self, converter: FxConverter) -> FxConversionFunc:
         return converter._generate_extern_kernel_out
 
 
+@dataclasses.dataclass
+class ExternKernelMultiOutLine(WrapperLine):
+    """Codegen line for multi-output .out() variant calls.
+
+    Generates a kernel call with pre-allocated output buffers passed as
+    keyword arguments. E.g. kernel(x, out0=buf0, out1=buf1).
+    """
+
+    wrapper: PythonWrapperCodegen
+    node: ir.ExternKernelMultiOut
+
+    def codegen(self, code: IndentedBuffer) -> None:
+        node = self.node
+        kernel_name = node.get_kernel_name()
+
+        args = [*node.codegen_args(), *node.codegen_kwargs()]
+        for out_name, out_node in zip(
+            node.out_arg_names, node.out_variant_output_nodes
+        ):
+            args.append(f"{out_name}={out_node.get_name()}")
+
+        code.writeline(f"{node.get_name()} = {kernel_name}({', '.join(args)})")
+
+        for out_node in node.out_variant_output_nodes:
+            if isinstance(out_node.layout, ir.Layout):
+                out_node.codegen_size_asserts(self.wrapper)
+
+
 @dataclasses.dataclass
 class FreeLine(WrapperLine):
     wrapper: PythonWrapperCodegen
-    node: Union[BufferLike, ir.TorchBindObject]
+    node: BufferLike | ir.TorchBindObject
 
     def codegen(self, code: IndentedBuffer) -> None:
         assert self.node.get_name() not in V.graph.removed_buffers
@@ -634,10 +694,11 @@ class KernelCallLine(WrapperLine):
     arg_types: list[str]
     triton: bool
     triton_meta: dict[str, Any]
-    inductor_meta: Optional[dict[str, Any]]
+    inductor_meta: dict[str, Any] | None
     device: torch.device
     graph_name: str
     original_fxnode_name: str
+    current_stream_idx: int | None = None
 
     def codegen(self, code: IndentedBuffer) -> None:
         self.wrapper._generate_kernel_call_helper(
@@ -652,6 +713,7 @@ def codegen(self, code: IndentedBuffer) -> None:
             device=self.device,
             graph_name=self.graph_name,
             original_fxnode_name=self.original_fxnode_name,
+            current_stream_idx=self.current_stream_idx,
         )
 
     def codegen_fx(self, converter: FxConverter) -> FxConversionFunc:
@@ -663,9 +725,9 @@ class KernelDefinitionLine(WrapperLine):
     wrapper: PythonWrapperCodegen
     kernel_name: str
     kernel_body: str
-    metadata: Optional[str] = None
+    metadata: str | None = None
     gpu: bool = True
-    cpp_definition: Optional[str] = None
+    cpp_definition: str | None = None
 
     def codegen(self, code: IndentedBuffer) -> None:
         self.wrapper._define_kernel_helper(
@@ -706,6 +768,75 @@ def __str__(self) -> str:
         return f"{type(self).__name__}({', '.join(args)})"
 
 
+@dataclasses.dataclass
+class EnterDeviceContextManagerWithStreamInfoLine(EnterDeviceContextManagerLine):
+    """Enter a CUDA device context and retrieve user stream objects.
+
+    Attributes:
+        num_streams: Number of streams (determined by user annotations on nodes).
+        stream_idx_to_user_obj_idx: Maps stream_idx → user_object_index for
+            retrieving user stream objects via get_external_object_by_index.
+    """
+
+    num_streams: int = 1
+    stream_idx_to_user_obj_idx: dict[int, int] = dataclasses.field(default_factory=dict)
+
+    def codegen(self, code: IndentedBuffer) -> None:
+        """Generate context switching and stream retrieval code."""
+        if V.graph.cpp_wrapper:
+            super().codegen(code)
+        else:
+            super().codegen(code)
+            code.writeline(f"{DEFAULT_STREAM} = torch.cuda.current_stream()")
+
+            if self.num_streams > 1:
+                for i in range(1, self.num_streams):
+                    user_obj_idx = self.stream_idx_to_user_obj_idx[i]
+                    code.writeline(
+                        f"{STREAM_NAME_TEMPLATE.format(stream_idx=i)} "
+                        f"= get_external_object_by_index({user_obj_idx})",
+                    )
+
+
+@dataclasses.dataclass
+class ExitDeviceContextManagerWithStreamInfoLine(ExitDeviceContextManagerLine):
+    """Exit a CUDA device context.
+
+    Attributes:
+        num_streams: Number of streams that were allocated (must match Enter).
+    """
+
+    num_streams: int = 1
+
+    def codegen(self, code: IndentedBuffer) -> None:
+        """Generate context exit code."""
+        if not V.graph.cpp_wrapper:
+            code.do_unindent()
+
+
+@dataclasses.dataclass
+class EnterCudaStreamContextLine(WrapperLine):
+    """Enter a context executed by respective CUDA Stream.
+
+    Attributes:
+        stream_idx: The index number corresponds to the entering CUDA Stream context.
+    """
+
+    stream_idx: int
+
+    def codegen(self, code: IndentedBuffer) -> None:
+        code.writeline(f"with torch.cuda.stream({get_stream_name(self.stream_idx)}):")
+        code.do_indent()
+
+
+@dataclasses.dataclass
+class ExitCudaStreamContextLine(WrapperLine):
+    """Generate code to exit the current stream context."""
+
+    def codegen(self, code: IndentedBuffer) -> None:
+        code.do_unindent()
+
+
 class EfficientPeakEstimate:
     def __init__(self):
         from ..memory import estimate_peak_memory, get_freeable_input_buf
@@ -727,7 +858,7 @@ def __init__(self):
         )
 
     def _get_size(self, node: BufferLike) -> int:
-        return V.graph.sizevars.size_hint(
+        return V.graph.sizevars.optimization_hint(
             V.graph.get_allocation_storage_size(node), fallback=0
         ) * get_dtype_size(node.get_dtype())
 
@@ -785,10 +916,11 @@ def plan(self, state: MemoryPlanningState) -> MemoryPlanningLine:
             return self
 
         # Regular buffer reuse
+        # Stream is part of the key, so cross-stream reuse is naturally prevented.
         key = buffer_reuse_key(self.node)
         if config.allow_buffer_reuse and key in state:
             free_line = state.pop(key)
-            size = V.graph.sizevars.size_hint(
+            size = V.graph.sizevars.optimization_hint(
                 V.graph.get_allocation_storage_size(self.node), fallback=0
             ) * get_dtype_size(self.node.get_dtype())
             if self.should_reuse_buffer(free_line, size):
@@ -993,7 +1125,7 @@ def codegen_fx(self, converter: FxConverter) -> FxConversionFunc:
 class IndexPutFallbackLine(WrapperLine):
     wrapper: PythonWrapperCodegen
     node: ir.IndexPutFallback
-    indices: list[Optional[ir.IRNode]]
+    indices: list[ir.IRNode | None]
 
     def codegen(self, code: IndentedBuffer) -> None:
         node = self.node
@@ -1059,7 +1191,7 @@ class UnbackedSymbolDefsLine(WrapperLine):
     wrapper: PythonWrapperCodegen
     output_name: str
     outputs: Any
-    unbacked_bindings: Optional[dict[sympy.Symbol, pytree.KeyPath]]
+    unbacked_bindings: dict[sympy.Symbol, pytree.KeyPath] | None
 
     def codegen(self, code: IndentedBuffer) -> None:
         self.wrapper._codegen_unbacked_symbol_defs_for_outputs(
@@ -1070,8 +1202,22 @@ def codegen_fx(self, converter: FxConverter) -> FxConversionFunc:
         return converter._generate_unbacked_symbol_defs
 
 
+@dataclasses.dataclass
+class AssertSizeStrideLine(WrapperLine):
+    name: str
+    size: str
+    stride: str
+
+    def codegen(self, code: IndentedBuffer) -> None:
+        code.writeline(f"assert_size_stride({self.name}, {self.size}, {self.stride})")
+
+    @staticmethod
+    def codegen_fx(converter: FxConverter) -> FxConversionFunc:
+        return converter._generate_assert_size_stride
+
+
 BufferName = str
-Line = Union[MemoryPlanningLine, LineContext]
+Line = MemoryPlanningLine | LineContext
 
 
 class PythonWrapperCodegen(CodeGen):
@@ -1083,9 +1229,11 @@ class PythonWrapperCodegen(CodeGen):
 
     def __init__(self):
         super().__init__()
+        self._pending_input_asserts: dict[str, tuple[str, str]] = {}
+        self._pending_alignment_copies: OrderedSet[str] = OrderedSet()
         self._names_iter: Iterator[int] = count()
         self.args_to_buffers: dict[
-            str, Union[None, ir.TensorBox, ir.Buffer, ir.TorchBindObject]
+            str, None | ir.TensorBox | ir.Buffer | ir.TorchBindObject
         ] = {}
         self.imports = IndentedBuffer()
         self.header = IndentedBuffer()
@@ -1113,7 +1261,7 @@ def __init__(self):
         self.none_str = "None"
         self.move_begin = "std::move(" if V.graph.cpp_wrapper else ""
         self.move_end = ")" if V.graph.cpp_wrapper else ""
-        self.last_seen_device_guard_index: Optional[int] = None
+        self.last_seen_device_guard_index: int | None = None
         self.supports_intermediate_hooks = True
         self.user_defined_kernel_cache: dict[
             tuple[Any, ...], tuple[str, Any, dict[str, Any]]
@@ -1181,9 +1329,9 @@ def add_import_once(line: str) -> None:
     @staticmethod
     def create(
         is_subgraph: bool,
-        subgraph_name: Optional[str],
-        parent_wrapper: Optional[PythonWrapperCodegen],
-        partition_signatures: Optional[ir.GraphPartitionSignature] = None,
+        subgraph_name: str | None,
+        parent_wrapper: PythonWrapperCodegen | None,
+        partition_signatures: ir.GraphPartitionSignature | None = None,
     ):
         if is_subgraph:
             assert subgraph_name is not None
@@ -1402,7 +1550,7 @@ def mark_output_type(self) -> None:
 
     def get_graph_inputs(
         self,
-    ) -> dict[str, Union[ir.TensorBox, ir.TorchBindObject, sympy.Expr]]:
+    ) -> dict[str, ir.TensorBox | ir.TorchBindObject | sympy.Expr]:
         return V.graph.graph_inputs
 
     def get_graph_outputs(self) -> list[IRNode]:
@@ -1410,13 +1558,19 @@ def get_graph_outputs(self) -> list[IRNode]:
 
     def codegen_input_size_asserts(self) -> None:
         for name, buf in self.get_graph_inputs().items():
-            if isinstance(buf, (sympy.Expr, ir.TorchBindObject)):
+            if isinstance(
+                buf,
+                (
+                    sympy.Basic,
+                    ir.TorchBindObject,
+                    ir.GeneratorState,
+                    ir.OpaqueObjectState,
+                ),
+            ):
                 continue
 
             # a graph partition may take an IRNode output from a previous partition
-            if name not in V.graph.graph_input_names or isinstance(
-                buf, ir.GeneratorState
-            ):
+            if name not in V.graph.graph_input_names:
                 continue
 
             # comparing strides for 0 size tensor is tricky. Ignore them for now.
@@ -1424,14 +1578,13 @@ def codegen_input_size_asserts(self) -> None:
                 continue
             size = self.codegen_python_shape_tuple(buf.get_size())
             stride = self.codegen_python_shape_tuple(buf.get_stride())
-            self.prefix.writeline(f"assert_size_stride({name}, {size}, {stride})")
+            self._pending_input_asserts[name] = (size, stride)
 
     def codegen_input_nan_asserts(self) -> None:
         self.prefix.writeline("# make sure graph inputs are not nan/inf")
         for name, buf in self.get_graph_inputs().items():
-            if isinstance(buf, (sympy.Expr, ir.TorchBindObject)):
+            if isinstance(buf, (sympy.Basic, ir.TorchBindObject)):
                 continue
-
             line = f"assert not {name}.isnan().any().item()"
             self.prefix.writeline(line)
             line = f"assert not {name}.isinf().any().item()"
@@ -1517,12 +1670,53 @@ def codegen_input_size_and_nan_asserts(self) -> None:
         if config.nan_asserts:
             self.codegen_input_nan_asserts()
 
+    # Input size/stride assertions are deferred from the top of call() to just
+    # before the first kernel that uses each input. This avoids a block of N
+    # sequential assert calls (~1 us each) on the critical path before the first
+    # GPU kernel launch. Called from the scheduler codegen loop.
+    def codegen_deferred_input_asserts(self, input_names: Iterable[str]) -> None:
+        for name in input_names:
+            if name in self._pending_input_asserts:
+                size, stride = self._pending_input_asserts.pop(name)
+                self.writeline(AssertSizeStrideLine(name, size, stride))
+
+    def register_alignment_check_inputs(self) -> None:
+        """Populate pending alignment copies for non-mutated inputs.
+        Called from the scheduler after mutated_input_idxs is computed."""
+        if V.graph.cpp_wrapper:
+            return
+        inputs_to_check = V.graph.inputs_to_check
+        if not inputs_to_check:
+            return
+        # Mutated inputs are handled separately by the runtime wrapper,
+        # which needs to copy back the mutation after the call.
+        mutated_idxs = OrderedSet(V.graph.mutated_input_idxs)
+        for idx in inputs_to_check:
+            if idx not in mutated_idxs:
+                name = V.graph.graph_input_names[idx]
+                self._pending_alignment_copies.add(name)
+        if self._pending_alignment_copies:
+            V.graph._defers_input_alignment = True
+            self.imports.writeline(
+                "from torch._C._dynamo.guards import copy_misaligned"
+            )
+
+    def codegen_deferred_alignment_copies(self, input_names: Iterable[str]) -> None:
+        """Emit alignment check + clone just before the first kernel
+        that reads each input, hiding the cost behind GPU execution."""
+        if V.graph.cpp_wrapper:
+            return
+        for name in input_names:
+            if name in self._pending_alignment_copies:
+                self._pending_alignment_copies.discard(name)
+                self.writeline(f"{name} = copy_misaligned({name})")
+
     # this function (and below) takes the graph name as input so
     # that stream caching happens per graph instance. this
     # is important for nested subgraph codegening.
     def write_get_raw_stream(self, device_idx: int, graph_name: str) -> str:
         self.write_get_raw_stream_header()
-        name = f"stream{device_idx}"
+        name = get_raw_stream_name(device_idx)
         if config.triton.autotune_at_compile_time:
             self.kernel_autotune_calls.writeline(
                 f"{name} = get_raw_stream({device_idx})"
@@ -1553,10 +1747,34 @@ def pop_computed_sizes(self):
     def next_kernel_suffix(self) -> str:
         return f"{next(self._names_iter)}"
 
-    def codegen_device_guard_enter(self, device_idx: int) -> None:
-        self.writeline(
-            EnterDeviceContextManagerLine(device_idx, self.last_seen_device_guard_index)
-        )
+    def codegen_device_guard_enter(
+        self,
+        device_idx: int,
+        num_streams: int = 1,
+        stream_idx_to_user_obj_idx: dict[int, int] | None = None,
+    ) -> None:
+        if num_streams > 1:
+            assert stream_idx_to_user_obj_idx is not None
+            import_line = (
+                "from torch._dynamo.graph_bytecode_inputs import "
+                "get_external_object_by_index"
+            )
+            if not self.imports.contains(import_line):
+                self.imports.writeline(import_line)
+            self.writeline(
+                EnterDeviceContextManagerWithStreamInfoLine(
+                    device_idx,
+                    self.last_seen_device_guard_index,
+                    num_streams,
+                    stream_idx_to_user_obj_idx,
+                ),
+            )
+        else:
+            self.writeline(
+                EnterDeviceContextManagerLine(
+                    device_idx, self.last_seen_device_guard_index
+                )
+            )
         if config.triton.autotune_at_compile_time:
             # mimic logic of EnterDeviceContextManagerLine.codegen for the autotune code block
             self.write_triton_header_once()
@@ -1568,15 +1786,47 @@ def codegen_device_guard_enter(self, device_idx: int) -> None:
                 # Need get_raw_stream for subgraph
                 self.write_get_raw_stream_header()
             self.kernel_autotune_calls.writeline(
-                f"stream{device_idx} = get_raw_stream({device_idx})"
+                f"{get_raw_stream_name(device_idx)} = get_raw_stream({device_idx})"
             )
         self.last_seen_device_guard_index = device_idx
+        self._num_streams: int = num_streams
 
     def codegen_device_guard_exit(self) -> None:
-        self.writeline(ExitDeviceContextManagerLine())
+        if hasattr(self, "_num_streams") and self._num_streams > 1:
+            self.writeline(
+                ExitDeviceContextManagerWithStreamInfoLine(self._num_streams)
+            )
+        else:
+            self.writeline(ExitDeviceContextManagerLine())
         if config.triton.autotune_at_compile_time:
             self.kernel_autotune_calls.do_unindent()
 
+    def codegen_cuda_stream_enter(
+        self,
+        stream_idx: int,
+    ) -> EnterCudaStreamContextLine:
+        """Generate data structure for entering a CUDA Stream context.
+
+        Args:
+            stream_idx: The index number of the entering CUDA Stream context.
+
+        Raises:
+            ValueError: If this function is called while the previous stream context isn't exited.
+        """
+        # pyre-fixme[16]: scheduler.current_stream_name added in scheduler commit
+        if (current_stream_name := V.graph.scheduler.current_stream_name) is not None:
+            raise ValueError(
+                f"Nested stream context switching: {current_stream_name} -> "
+                f"{get_stream_name(stream_idx)}",
+            )
+        ctx_entrance = EnterCudaStreamContextLine(stream_idx=stream_idx)
+        self.writeline(ctx_entrance)
+        return ctx_entrance
+
+    def codegen_cuda_stream_exit(self) -> None:
+        """Generate data structure for exiting a CUDA Stream context."""
+        self.writeline(ExitCudaStreamContextLine())
+
     def generate_return(self, output_refs: list[str]) -> None:
         if output_refs:
             if config.nan_asserts:
@@ -1625,6 +1875,12 @@ def generate_fallback_kernel(self, node: ir.FallbackKernel) -> None:
                 return
         self.writeline(ExternKernelAllocLine(self, node))
 
+    def generate_extern_kernel_multi_out(self, node: ir.ExternKernelMultiOut) -> None:
+        """Generate .out() call with pre-allocated output buffers."""
+        for out_node in node.out_variant_output_nodes:
+            self.codegen_allocation(out_node)
+        self.writeline(ExternKernelMultiOutLine(self, node))
+
     def generate_extern_kernel_alloc(self, node: ir.ExternKernelAlloc):
         node.codegen_comment(self)
         self.writeline(ExternKernelAllocLine(self, node))
@@ -1671,10 +1927,10 @@ def _generate_extern_kernel_out_helper(
         self,
         kernel: str,
         out: str,
-        out_view: Optional[str],
+        out_view: str | None,
         args: list[str],
         device: str,
-        stack_traces: Optional[OrderedSet[str]] = None,
+        stack_traces: OrderedSet[str] | None = None,
     ) -> None:
         # add debug printer code for triton kernel calls at (jit) inductor level
         debug_printer_manager = V.graph.wrapper_code.debug_printer
@@ -1753,7 +2009,7 @@ def _generate_scatter_fallback(
 
     def generate_index_put_fallback(self, node: ir.IndexPutFallback) -> None:
         # Collect index tensors into a list.
-        indices: list[Optional[ir.IRNode]] = []
+        indices: list[ir.IRNode | None] = []
         valid_indices = node.inputs[2:]
         iter_valid_indices = iter(valid_indices)
         for i, _ in enumerate(node.indices):
@@ -1776,7 +2032,7 @@ def generate_fallback_kernel_with_runtime_lookup(
         buf_name: str,
         python_kernel_name: str,
         get_args: Callable[[], Sequence[str]],
-        op_overload: Union[torch._ops.OpOverload, torch._ops.HigherOrderOperator],
+        op_overload: torch._ops.OpOverload | torch._ops.HigherOrderOperator,
         raw_args: Sequence[Any],
         outputs: Sequence[ir.Buffer],
     ) -> None:
@@ -2021,9 +2277,9 @@ def strideof(name):
                 if isinstance(stride, sympy.Symbol) and stride not in bound_vars:
                     code.writeline(f"{stride} = {strideof(name)}[{dim}]")
                     bound_vars.add(stride)
-        elif isinstance(value, ir.TorchBindObject):
-            return
-        elif isinstance(value, ir.GeneratorState):
+        elif isinstance(
+            value, (ir.TorchBindObject, ir.GeneratorState, ir.OpaqueObjectState)
+        ):
             return
         else:
             if torch._inductor.config.graph_partition:
@@ -2184,7 +2440,7 @@ def apply_reinterpret(
 
         return apply_reinterpret(name, size, stride, offset, dtype, base_dtype)
 
-    def codegen_device_copy(self, src, dst, non_blocking: Union[bool, str]):
+    def codegen_device_copy(self, src, dst, non_blocking: bool | str):
         self.writeline(f"{dst}.copy_({src}, {non_blocking})")
 
     def codegen_multi_output(self, node: ir.MultiOutput):
@@ -2314,11 +2570,7 @@ def add_torchbind_input(name, value):
                     # the subclass.
                     continue
                 if isinstance(value, ir.TorchBindObject):
-                    if len(V.graph.torchbind_constants) == 0:
-                        # otherwise we have already imported the pickle package
-                        output.writeline("import pickle")
-                    output.writeline(f"global {name}")
-                    add_torchbind_input(name, value.get_real_obj())
+                    output.writeline(f"{name} = None")
                 elif isinstance(value, sympy.Expr):  # Don't need to add symbolic
                     # TODO: this fallback and those below actually will generate possibly
                     # invalid benchmark code, because it's not guaranteed 42
@@ -2327,20 +2579,26 @@ def add_torchbind_input(name, value):
                     add_expr_input(
                         name, V.graph.sizevars.optimization_hint(value, fallback=42)
                     )
+                elif isinstance(value, sympy.Basic):
+                    # sympy.Boolean (e.g. StrictLessThan from torch.cond predicates)
+                    # is not a sympy.Expr so optimization_hint cannot handle it.
+                    # Use False as a fallback for benchmark harness purposes.
+                    add_expr_input(name, False)
                 elif isinstance(value, ir.GeneratorState):
                     add_expr_input(
                         name,
                         f"torch.cuda.default_generators[{value.device.index}].graphsafe_get_state()",
                     )
+                elif isinstance(value, ir.OpaqueObjectState):
+                    output.writeline(f"{name} = None")
                 else:
-                    shape = [
-                        V.graph.sizevars.optimization_hint(x, fallback=42)
-                        for x in value.get_size()
-                    ]
-                    stride = [
-                        V.graph.sizevars.size_hint(x, fallback=42)
-                        for x in value.get_stride()
-                    ]
+                    shape = V.graph.sizevars.optimization_hints(
+                        value.get_size(), fallback=42
+                    )
+                    stride = V.graph.sizevars.optimization_hints(
+                        value.get_stride(), fallback=42
+                    )
+
                     add_fake_input(
                         name,
                         shape,
@@ -2391,9 +2649,9 @@ def define_kernel(
         self,
         kernel_name: str,
         kernel_body: str,
-        metadata: Optional[str] = None,
+        metadata: str | None = None,
         gpu: bool = True,
-        cpp_definition: Optional[str] = None,
+        cpp_definition: str | None = None,
     ):
         self.writeline(
             KernelDefinitionLine(
@@ -2408,7 +2666,7 @@ def define_kernel(
 
     @staticmethod
     def _format_kernel_definition(
-        kernel_name: str, kernel_body: str, metadata: Optional[str] = None
+        kernel_name: str, kernel_body: str, metadata: str | None = None
     ):
         if config.triton.autotune_at_compile_time and metadata:
             # Generating autotune block
@@ -2422,9 +2680,9 @@ def _define_kernel_helper(
         self,
         kernel_name: str,
         kernel_body: str,
-        metadata: Optional[str] = None,
+        metadata: str | None = None,
         gpu: bool = True,
-        cpp_definition: Optional[str] = None,
+        cpp_definition: str | None = None,
     ):
         if config.triton.autotune_at_compile_time and gpu:
             body = self._format_kernel_definition(
@@ -2460,7 +2718,8 @@ def define_user_defined_triton_kernel(
         kwargs,
         restore_value_args,
         reset_to_zero_args,
-        grids: list[list[Union[int, sympy.Expr]]],
+        grids: list[list[int | sympy.Expr]],
+        epilogue_fusion: tuple[ir.ComputedBuffer, str] | None,
     ):
         from ..runtime.triton_heuristics import (
             config_to_dict,
@@ -2622,7 +2881,7 @@ def add_arg(idx, arg, is_constexpr=False, equals_1=False, equals_none=False):
             extra_launcher_call_args = [*map(sympy.sympify, grids[0])]
         else:
 
-            def rename_sizes_for_launcher(expr: Union[int, sympy.Expr]) -> sympy.Expr:
+            def rename_sizes_for_launcher(expr: int | sympy.Expr) -> sympy.Expr:
                 if isinstance(expr, sympy.Expr):
                     symbols = [*expr.free_symbols]
                     if not symbols:
@@ -2661,6 +2920,11 @@ def rename_sizes_for_launcher(expr: Union[int, sympy.Expr]) -> sympy.Expr:
             }
             extra_launcher_call_args = [*extra_launcher_args.keys()]
 
+        if constexprs:
+            inductor_meta["declared_constexpr_names"] = [
+                arg_names[i] for i in constexprs
+            ]
+
         # Distinguish between different functions using function id
         cache_key: Any = [id(kernel.fn)]
         if len(configs) > 0:
@@ -2710,10 +2974,17 @@ def rename_sizes_for_launcher(expr: Union[int, sympy.Expr]) -> sympy.Expr:
             @triton.jit
             """
         )
-        kernel_src = user_defined_triton_kernel_transitive_closure_source_code(kernel)
+        kernel_src = user_defined_triton_kernel_transitive_closure_source_code(
+            kernel, epilogue_fusion
+        )
         if config.triton.unique_user_kernel_names:
             # We replace the original_name with the unique name.
             kernel_src = kernel_src.replace(f"def {original_name}(", f"def {name}(")
+        if config.cpp_wrapper:
+            # With cpp_wrapper + autotune_at_compile_time=False, the source is
+            # further embedded in a C++ raw string inside a Python r"""...""" wrapper.
+            # So we need to add backslash here.
+            kernel_src = kernel_src.replace('"""', '\\"\\"\\"')
         kernel_src = kernel_src.replace("'''", "\\'\\'\\'")
         compile_wrapper.splice(kernel_src)
 
@@ -2731,7 +3002,7 @@ def rename_sizes_for_launcher(expr: Union[int, sympy.Expr]) -> sympy.Expr:
         self.user_defined_kernel_cache[cache_key] = (name, triton_meta, inductor_meta)
         return name, triton_meta, inductor_meta, extra_launcher_call_args
 
-    def generate_numel_expr(self, kernel_name: str, tree, suffix: Optional[str] = None):
+    def generate_numel_expr(self, kernel_name: str, tree, suffix: str | None = None):
         sym_name = f"{kernel_name}_{tree.prefix}numel"
         if suffix is not None:
             sym_name += f"_{suffix}"
@@ -2787,7 +3058,7 @@ def generate_workspace_allocation(self, ws: WorkspaceArg):
                     name,
                     ws.device,
                     ws.dtype,
-                    shape=(V.graph.sizevars.size_hint(ws.count),),
+                    shape=(V.graph.sizevars.optimization_hint(ws.count),),
                     stride=(1,),
                 )
             )
@@ -2972,6 +3243,7 @@ def generate_kernel_call(
         )
 
         device = device or V.graph.get_current_device_or_throw()
+        current_stream_idx = V.graph.scheduler.current_stream_idx
         self.writeline(
             KernelCallLine(
                 self,
@@ -2991,6 +3263,7 @@ def generate_kernel_call(
                 graph_name=V.graph.name,
                 # pyrefly: ignore [bad-argument-type]
                 original_fxnode_name=original_fxnode_name,
+                current_stream_idx=current_stream_idx,
             )
         )
 
@@ -3008,25 +3281,32 @@ def _generate_kernel_call_helper(
         inductor_meta=None,
         graph_name="",
         original_fxnode_name=None,
+        current_stream_idx=None,
     ):
         device = device or V.graph.get_current_device_or_throw()
-        if not triton and device.type != "cuda":
+        if not triton and device.type not in ("cuda", "xpu"):
             if device.type == "cpu":
                 self.writeline(self.wrap_kernel_call(kernel_name, call_args))
             elif device.type == "mps":
                 # TODO: Fix me, MPS does not expose streams now
-                self.writeline(
-                    self.wrap_kernel_call(f"{kernel_name}.generated_kernel", call_args)
-                )
+                self.writeline(self.wrap_kernel_call(kernel_name, call_args))
             else:
                 raise RuntimeError(f"device {device.type} nyi")
             return
 
         call_args_str = self.prepare_triton_kernel_call(call_args)
         call_args_str = ", ".join(call_args_str)
-        stream_name = PythonWrapperCodegen.write_get_raw_stream(
-            self, device.index, graph_name
-        )
+        if current_stream_idx is not None and current_stream_idx != DEFAULT_STREAM_IDX:
+            # Inside a user stream context: emit a fresh get_raw_stream call so
+            # it picks up the active stream at runtime, rather than reusing the
+            # LRU-cached stream0 variable which captured the default stream.
+            self.write_get_raw_stream_header()
+            stream_name = "raw_stream"
+            self.writeline(f"{stream_name} = get_raw_stream({device.index})")
+        else:
+            stream_name = PythonWrapperCodegen.write_get_raw_stream(
+                self, device.index, graph_name
+            )
         if not triton:
             stream_ptr = f"c_void_p({stream_name})"
             self.writeline(
@@ -3111,7 +3391,7 @@ def infer_arg_by_inputs(raw_keys, raw_args, idx, reused_args):
 
             reused_args = {}
             for i, (arg, arg_type, raw_key, raw_arg) in enumerate(
-                # pyrefly: ignore [no-matching-overload]
+                # pyrefly: ignore [bad-argument-type, no-matching-overload]
                 zip(call_args, arg_types, raw_keys, raw_args)
             ):
                 key = None
@@ -3119,7 +3399,7 @@ def infer_arg_by_inputs(raw_keys, raw_args, idx, reused_args):
                     # arg may be passed in a kwarg style, and then we need to extract its value
                     key, arg = arg.split("=")
 
-                triton_input: Optional[str] = None
+                triton_input: str | None = None
                 if autotune_args and raw_key in autotune_args:
                     triton_input = self.get_autotuning_input_name(  # type: ignore[attr-defined]
                         autotune_args[raw_key]
@@ -3152,6 +3432,7 @@ def infer_arg_by_inputs(raw_keys, raw_args, idx, reused_args):
                     self.kernel_autotune_example_args[arg] = (arg_str, kernel_name)
                 else:
                     arg_str = self.generate_example_arg_value(arg, arg_type, raw_arg)
+
                 if isinstance(arg, str) and should_unwrap_unspec_arg(arg):
                     arg_str += ".item()"
                 all_args.append(arg_str if key is None else f"{key}={arg_str}")
@@ -3220,7 +3501,7 @@ def __repr__(self):
             return s.codegen_reference()
         elif has_triton_package() and isinstance(s, triton.language.dtype):  # type: ignore[possibly-undefined]
             return repr(s)
-        elif isinstance(s, ir.GeneratorState):
+        elif isinstance(s, (ir.GeneratorState, ir.OpaqueObjectState)):
             return s.codegen_reference()
         elif is_opaque_value_type(type(s)):
             obj_repr, opaque_types = get_opaque_obj_repr(s)
@@ -3304,7 +3585,7 @@ def make_comment(self, line):
     def make_tensor_alias(self, new_name, old_name, comment=""):
         return f"{self.declare}{new_name} = {old_name}{self.ending}  {self.comment} {comment}"
 
-    def make_buffer_free(self, buffer: Union[BufferLike, ir.TorchBindObject]):
+    def make_buffer_free(self, buffer: BufferLike | ir.TorchBindObject):
         return f"del {buffer.get_name()}"
 
     def make_free_by_names(self, names_to_del: list[str]):
@@ -3316,7 +3597,7 @@ def codegen_exact_buffer_reuse(self, old_name: str, new_name: str, del_line: str
     def write_provenance_debug_handle(
         self,
         kernel_name,
-        debug_handle: Optional[int] = None,
+        debug_handle: int | None = None,
     ):
         if debug_handle is not None:
             self.writeline(
@@ -3467,7 +3748,7 @@ def codegen_unbacked_symbol_defs_for_outputs(
         self,
         output_name: str,
         outputs: Any,
-        unbacked_bindings: Optional[dict[sympy.Symbol, pytree.KeyPath]],
+        unbacked_bindings: dict[sympy.Symbol, pytree.KeyPath] | None,
     ) -> None:
         unbacked_bindings = resolve_unbacked_bindings(
             V.graph.sizevars.shape_env, unbacked_bindings
@@ -3480,7 +3761,7 @@ def _codegen_unbacked_symbol_defs_for_outputs(
         self,
         output_name: str,
         outputs: Any,
-        unbacked_bindings: Optional[dict[sympy.Symbol, pytree.KeyPath]],
+        unbacked_bindings: dict[sympy.Symbol, pytree.KeyPath] | None,
     ) -> None:
         if not unbacked_bindings:
             return
@@ -3892,7 +4173,7 @@ def can_prove_buffer_has_static_shape(buffer):
     def write_kernel_context_guard(
         self,
         kernel_name: str,
-        node_schedule: Union[Sequence[BaseSchedulerNode], ExternKernel],
+        node_schedule: Sequence[BaseSchedulerNode] | ExternKernel,
     ):
         return
 
@@ -3925,7 +4206,7 @@ def __init__(
         self,
         subgraph_name: str,
         parent_wrapper: PythonWrapperCodegen,
-        partition_signatures: Optional[ir.GraphPartitionSignature] = None,
+        partition_signatures: ir.GraphPartitionSignature | None = None,
     ):
         # It is necessary to set the subgraph_name before calling super __init__
         # because __init__ calls set_launcher_fn_name
@@ -3983,7 +4264,7 @@ def get_wrapper_call_indent(self) -> int:
 
     def get_graph_inputs(
         self,
-    ) -> dict[str, Union[ir.TensorBox, ir.TorchBindObject, sympy.Expr, None]]:
+    ) -> dict[str, ir.TensorBox | ir.TorchBindObject | sympy.Expr | None]:
         if signature := self.partition_signatures:
             inputs = signature.input_nodes | {
                 str(s): s for s in signature.symbol_inputs
@@ -4018,6 +4299,27 @@ def codegen_allocation(self, buffer: ir.Buffer):
 
         super().codegen_allocation(buffer)
 
+    def _write_get_raw_stream(
+        self, device_idx: int, graph: GraphLowering | None = None
+    ) -> str:
+        self.write_triton_header_once()
+        # pyre-fixme[16]: scheduler.current_stream_name added in scheduler commit
+        if (current_stream_name := V.graph.scheduler.current_stream_name) is not None:
+            name = f"{current_stream_name}_raw"
+            self.writeline(f"{name} = {current_stream_name}.cuda_stream")
+        else:
+            name = get_raw_stream_name(device_idx)
+            self.writeline(f"{name} = get_raw_stream({device_idx})")
+        return name
+
+    def codegen_graph_nvtx_range_push(self, post_grad_graph_id: int) -> None:
+        """Generate NVTX range push for graph."""
+        self.writeline(f"torch.cuda.nvtx.range_push('graph {post_grad_graph_id}')")
+
+    def codegen_graph_nvtx_range_pop(self) -> None:
+        """Generate NVTX range pop for graph."""
+        self.writeline("torch.cuda.nvtx.range_pop()")
+
     @cache_on_self
     def write_triton_header_once(self) -> None:
         # TODO: Uncomment in future. This will be needed to support subgraph
diff --git a/torch/_inductor/codegen/wrapper_fxir.py b/torch/_inductor/codegen/wrapper_fxir.py
index b395c96ce0893..4a1cce8bdb964 100644
--- a/torch/_inductor/codegen/wrapper_fxir.py
+++ b/torch/_inductor/codegen/wrapper_fxir.py
@@ -4,8 +4,8 @@
 import operator
 import textwrap
 from collections import Counter
-from collections.abc import Callable, Sequence
-from typing import Any, Optional, Union
+from collections.abc import Callable, Iterable, Sequence
+from typing import Any
 
 import sympy
 
@@ -30,9 +30,9 @@
     CallMethodKey,
     ConvertIntKey,
     DivideByKey,
-    free_unbacked_symbols,
 )
 from torch.utils import _pytree as pytree
+from torch.utils._ordered_set import OrderedSet
 from torch.utils._sympy.functions import FloorDiv
 from torch.utils._sympy.interp import _run_sympy_handler, sympy_interp
 from torch.utils._sympy.reference import OptimizedPythonReferenceAnalysis
@@ -94,13 +94,13 @@ class SymbolBuffer(CodegenSymbol):
     def get_name(self) -> str:
         return str(self.symbol)
 
-    def get_example(self) -> Union[torch.Tensor, torch.SymInt]:
+    def get_example(self) -> torch.Tensor | torch.SymInt:
         sym_int = convert_to_symint(self.symbol)
         assert isinstance(sym_int, torch.SymInt)
         return sym_int
 
 
-CodegenBuffer = Union[BufferLike, SymbolBuffer]
+CodegenBuffer = BufferLike | SymbolBuffer
 
 
 @dataclasses.dataclass
@@ -166,7 +166,7 @@ def codegen_conditional(self, conditional: ir.Conditional) -> None:
             self.codegen_subgraph_common(subgraph)
 
     def define_subgraph_launcher_fn(
-        self, name: str, subgraph_code: Union[ValueWithLineMap, FileBackedGraphModule]
+        self, name: str, subgraph_code: ValueWithLineMap | FileBackedGraphModule
     ) -> None:
         """
         Record subgms as they're generated.
@@ -181,7 +181,7 @@ def is_subgraph(self) -> bool:
 
     def get_fx_graph_inputs(
         self,
-    ) -> dict[str, Union[ir.TensorBox, ir.TorchBindObject, sympy.Expr, None]]:
+    ) -> dict[str, ir.TensorBox | ir.TorchBindObject | sympy.Expr | None]:
         """
         Get the input nodes corresponding to FX graph placeholders.
         """
@@ -232,13 +232,20 @@ def write_header(self) -> None:
         """
         PythonWrapperCodegen.write_header(self)
 
+    def register_alignment_check_inputs(self) -> None:
+        """FXIR does not emit deferred alignment copies.
+        Alignment is handled by the runtime wrapper."""
+
+    def codegen_deferred_alignment_copies(self, input_names: Iterable[str]) -> None:
+        """FXIR does not emit deferred alignment copies."""
+
     @classmethod
     def create(
         cls: type["WrapperFxCodegen"],
         is_subgraph: bool,
-        subgraph_name: Optional[str],
-        parent_wrapper: Optional[PythonWrapperCodegen],
-        partition_signatures: Optional[ir.GraphPartitionSignature] = None,
+        subgraph_name: str | None,
+        parent_wrapper: PythonWrapperCodegen | None,
+        partition_signatures: ir.GraphPartitionSignature | None = None,
     ) -> "WrapperFxCodegen":
         if is_subgraph:
             assert subgraph_name is not None
@@ -274,7 +281,7 @@ class FxConverter:
 
     lines: list[Line]
     prologue: str
-    graph_inputs: dict[str, Union[ir.TensorBox, ir.TorchBindObject, sympy.Expr, None]]
+    graph_inputs: dict[str, ir.TensorBox | ir.TorchBindObject | sympy.Expr | None]
     graph_outputs: list[ir.IRNode]
     subgms: dict[str, torch.fx.GraphModule]
     is_subgraph: bool
@@ -283,7 +290,7 @@ def __post_init__(self) -> None:
         graph = torch.fx.Graph()
         self.gm = GraphModule({}, graph)  # Wrapper FX IR.
         self.buffer_to_node: dict[
-            Optional[str], torch.fx.Node
+            str | None, torch.fx.Node
         ] = {}  # Symbol table for codegen.
         self.kernels: dict[str, TritonKernel] = {}  # Table to store Triton kernels.
         self._unique_symbol_ids: Counter[str] = Counter()
@@ -316,7 +323,7 @@ def _create_as_strided(
         input_node: torch.fx.Node,
         size: tuple[Any, ...],
         stride: tuple[Any, ...],
-        offset: Union[int, sympy.Expr],
+        offset: int | sympy.Expr,
     ) -> torch.fx.Node:
         return self.gm.graph.call_function(
             torch.as_strided,
@@ -336,7 +343,7 @@ def _record_allocation(self, buffer: CodegenBuffer, node: torch.fx.Node) -> None
         assert node not in self.buffer_to_node
         self.buffer_to_node[buffer.get_name()] = node
 
-    def _free(self, buffer: Union[CodegenBuffer, ir.TorchBindObject]) -> None:
+    def _free(self, buffer: CodegenBuffer | ir.TorchBindObject) -> None:
         """
         Removes the buffer from the symbol table.
         """
@@ -411,7 +418,7 @@ def _generate_graph_input_shapes(self) -> None:
         """
 
         def _codegen_symbol(
-            sym_or_exp: Union[sympy.Symbol, sympy.Expr],
+            sym_or_exp: sympy.Symbol | sympy.Expr,
             base_node: torch.fx.Node,
             target: torch._ops.OpOverload,
             dim: int,
@@ -492,7 +499,7 @@ def _generate_graph_constants(self) -> None:
             setattr(self.gm, name, value)
             self.buffer_to_node[name] = node
 
-    def _generate_buffer(self, node: ir.IRNode) -> Optional[torch.fx.Node]:
+    def _generate_buffer(self, node: ir.IRNode) -> torch.fx.Node | None:
         """
         Generates FX IR for transformations on a buffer, such as ReinterpretView.
         Does nothing if no such transformations are present.
@@ -502,7 +509,7 @@ def _generate_buffer(self, node: ir.IRNode) -> Optional[torch.fx.Node]:
             # Generate FX nodes to compute the shape expression.
             return self._sympy_interp(node.expr).node
 
-        def generate_to_buffer(node: ir.IRNode) -> Optional[BufferLike]:
+        def generate_to_buffer(node: ir.IRNode) -> BufferLike | None:
             if isinstance(node, (ir.Buffer, WorkspaceArg)):
                 return node
             elif isinstance(node, ir.NoneAsConstantBuffer):
@@ -539,7 +546,7 @@ def generate_to_buffer(node: ir.IRNode) -> Optional[BufferLike]:
 
     def _generate_outputs(
         self,
-    ) -> Union[Optional[torch.fx.Node], list[Optional[torch.fx.Node]]]:
+    ) -> torch.fx.Node | None | list[torch.fx.Node | None]:
         """
         Generate FX IR for graph outputs.
         """
@@ -644,9 +651,7 @@ def _sympy_interp(self, expr: sympy.Expr) -> torch.fx.Proxy:
         )
         return self.expr_to_proxy[expr]
 
-    def _generate_sym_node(
-        self, s: Union[int, sympy.Expr]
-    ) -> Union[int, torch.fx.Node]:
+    def _generate_sym_node(self, s: int | sympy.Expr) -> int | torch.fx.Node:
         if isinstance(s, (int, sympy.Integer)):
             return int(s)
         elif isinstance(s, sympy.Symbol):
@@ -665,7 +670,7 @@ def _generate_sym_node(
 
     def _generate_sym_nodes(
         self, shape: Sequence[sympy.Expr]
-    ) -> list[Union[int, torch.fx.Node]]:
+    ) -> list[int | torch.fx.Node]:
         return [self._generate_sym_node(s) for s in shape]
 
     def _generate_allocate(self, line: WrapperLine) -> None:
@@ -692,7 +697,7 @@ def _generate_allocate(self, line: WrapperLine) -> None:
     def _generate_conditional(self, line: WrapperLine) -> None:
         assert isinstance(line, ConditionalLine)
 
-        def get_subgm_attr(subgraph: Optional[ir.Subgraph]) -> torch.fx.Node:
+        def get_subgm_attr(subgraph: ir.Subgraph | None) -> torch.fx.Node:
             assert subgraph is not None
             return self._get_subgm_attr(subgraph)
 
@@ -703,7 +708,7 @@ def get_subgm_attr(subgraph: Optional[ir.Subgraph]) -> torch.fx.Node:
             for subgraph in (ir_node.true_subgraph, ir_node.false_subgraph)
         ]
 
-        def generate_buffer(node: Optional[ir.IRNode]) -> Optional[torch.fx.Node]:
+        def generate_buffer(node: ir.IRNode | None) -> torch.fx.Node | None:
             assert node is not None
             return self._generate_buffer(node)
 
@@ -716,6 +721,9 @@ def generate_buffer(node: Optional[ir.IRNode]) -> Optional[torch.fx.Node]:
         )
         self._record_allocation(ir_node, fx_node)
 
+    def _generate_assert_size_stride(self, line: WrapperLine) -> None:
+        pass
+
     def _generate_comment(self, line: WrapperLine) -> None:
         assert isinstance(line, CommentLine)
         # We ignore comments in FX IR.
@@ -730,7 +738,7 @@ def _generate_dynamic_scalar(self, line: WrapperLine) -> None:
         keypath = ir_node.keypath
         graph = self.gm.graph
 
-        def generate_item(x: Optional[torch.fx.Node]) -> torch.fx.Node:
+        def generate_item(x: torch.fx.Node | None) -> torch.fx.Node:
             assert x is not None
             return graph.call_function(
                 aten.item.default,
@@ -873,8 +881,8 @@ def _generate_multi_output(self, line: WrapperLine) -> None:
     def _generate_fallback_call(
         self,
         ir_node: ir.ExternKernel,
-        args: Optional[tuple[Any, ...]] = None,
-        kwargs: Optional[dict[str, Any]] = None,
+        args: tuple[Any, ...] | None = None,
+        kwargs: dict[str, Any] | None = None,
     ) -> None:
         fx_node = self.gm.graph.call_function(
             ir_node.op_overload,  # type: ignore[arg-type]
@@ -883,14 +891,19 @@ def _generate_fallback_call(
         )
         result_buffer = ir_node.codegen_reference()
         self.buffer_to_node[result_buffer] = fx_node
+        # For in-place mutation ops (e.g., scatter_reduce_, index_put_),
+        # update the buffer mapping for mutated inputs so downstream
+        # references to the mutated buffer see the post-mutation node.
+        for mutated_name in ir_node.get_mutation_names():
+            self.buffer_to_node[mutated_name] = fx_node
 
     def _generate_index_put_fallback(self, line: WrapperLine) -> None:
         assert isinstance(line, IndexPutFallbackLine)
         ir_node = line.node
 
         def generate_buffer_or_none(
-            x: Union[ir.IRNode, Sequence[ir.IRNode], None],
-        ) -> Optional[torch.fx.Node]:
+            x: ir.IRNode | Sequence[ir.IRNode] | None,
+        ) -> torch.fx.Node | None:
             """
             Handles None before calling _generate_buffer.
             """
@@ -917,6 +930,15 @@ def _generate_scatter_fallback(self, line: WrapperLine) -> None:
         kwargs = {}
         if reduce := ir_node.kwargs.get("reduce"):
             kwargs["reduce"] = reduce
+        # Only pass kwargs that the op's schema actually accepts, since
+        # ScatterFallback stores both reduce and include_self for all
+        # scatter variants, but not all ops support them (e.g.,
+        # scatter_.value has no kwargs, scatter_reduce_.two has both).
+        assert isinstance(ir_node.op_overload, torch._ops.OpOverload)
+        schema_arg_names = OrderedSet(
+            [a.name for a in ir_node.op_overload._schema.arguments]
+        )
+        kwargs = {k: v for k, v in ir_node.kwargs.items() if k in schema_arg_names}
 
         self._generate_fallback_call(ir_node, args, kwargs)
 
@@ -940,9 +962,6 @@ def _generate_triton_call(self, line: WrapperLine) -> None:
         kernel = self.kernels[line.kernel_name]
         tuner = kernel.tuner
 
-        class UnbackedSymintsError(Exception):
-            pass
-
         def tune_kernel(tuner: CachingAutotuner, call_args: Sequence[Any]) -> None:
             from triton.runtime import driver
 
@@ -958,13 +977,10 @@ def node_to_tuning_arg(arg: Any) -> Any:
                 for dynamic shapes.
                 """
 
-                def to_size_hint_sympy_int(arg: Union[sympy.Expr, int]) -> int:
-                    if len(free_unbacked_symbols(arg)) > 0:
-                        # NYI: tuning args require backed symints.
-                        raise UnbackedSymintsError
-                    return V.graph.sizevars.size_hint(arg)
+                def to_size_hint_sympy_int(arg: sympy.Expr | int) -> int:
+                    return V.graph.sizevars.optimization_hint(arg)
 
-                def to_size_hint_list(arg: list[Union[torch.SymInt, int]]) -> list[int]:
+                def to_size_hint_list(arg: list[torch.SymInt | int]) -> list[int]:
                     args_sympy = [
                         x.node.expr if isinstance(x, torch.SymInt) else x for x in arg
                     ]
@@ -989,11 +1005,25 @@ def to_size_hint_list(arg: list[Union[torch.SymInt, int]]) -> list[int]:
         # The FX backend currently only supports compile-time tuning.
         kernel_name = tuner.fn.__name__
         if config.triton.autotune_at_compile_time:
-            try:
+            # Skip compile-time autotuning if any unbacked symbol lacks a user-provided
+            # optimization hint — autotuning with the generic fallback would
+            # produce meaningless results.
+            hinted = V.graph.sizevars.all_unbacked_explicitly_hinted
+            can_tune = True
+            for arg in call_args:
+                if isinstance(arg, torch.fx.Node):
+                    fake = arg.meta["val"]
+                    if not hinted(list(fake.shape) + list(fake.stride())):
+                        can_tune = False
+                        break
+                elif not hinted(arg):
+                    can_tune = False
+                    break
+            if can_tune:
                 tune_kernel(tuner, call_args)
-            except UnbackedSymintsError:
+            else:
                 log.info(
-                    "Detected unbacked symints. Skipping autotuning for kernel %s.",
+                    "Detected unhinted unbacked symints. Skipping compile-time autotuning for kernel %s.",
                     kernel_name,
                 )
         else:
@@ -1114,7 +1144,7 @@ def _generate_extern_kernel_common(
             for k, v in kernel.kwargs.items()
         }
 
-        result_buffer: Optional[str] = None
+        result_buffer: str | None = None
         if isinstance(kernel, ir.ExternKernelOut):
             kwargs["out"] = self.buffer_to_node[out_ir_node.codegen_reference()]
         elif isinstance(kernel.layout, (ir.Layout, ir.MultiOutputLayout)):
diff --git a/torch/_inductor/codegen/xpu/compile_utils.py b/torch/_inductor/codegen/xpu/compile_utils.py
new file mode 100644
index 0000000000000..e979ad570c909
--- /dev/null
+++ b/torch/_inductor/codegen/xpu/compile_utils.py
@@ -0,0 +1,130 @@
+# mypy: allow-untyped-defs
+import logging
+import os
+import shutil
+
+from torch._inductor import config
+from torch._inductor.codegen.xpu.xpu_env import get_xpu_arch
+from torch._inductor.utils import is_linux
+
+from ..cuda.compile_utils import _cutlass_include_paths
+
+
+log = logging.getLogger(__name__)
+
+
+def _sycl_compiler() -> str:
+    # Search order:
+    # 0) which icpx
+    # 1) config.xpu.oneapi_root
+    # 2) ONEAPI_ROOT environment variable
+    # 3) default system search PATH.
+    if shutil.which("icpx"):
+        return "icpx"
+
+    if os.path.exists(config.xpu.oneapi_root or ""):
+        oneapi_root = config.xpu.oneapi_root
+    elif os.path.exists(os.getenv("ONEAPI_ROOT") or ""):
+        oneapi_root = os.getenv("ONEAPI_ROOT")
+    else:
+        oneapi_root = None
+
+    if oneapi_root:
+        oneapi_inclue = os.path.join(oneapi_root, "include")
+        if "CPLUS_INCLUDE_PATH" in os.environ:
+            os.environ["CPLUS_INCLUDE_PATH"] += ":" + oneapi_inclue
+        else:
+            os.environ["CPLUS_INCLUDE_PATH"] = oneapi_inclue
+        return os.path.realpath(os.path.join(oneapi_root, "bin/icpx"))
+    else:
+        raise RuntimeError("Can not find Intel compiler.")
+
+
+def _sycl_lib_options() -> list[str]:
+    """
+    Util function for CUTLASS backend to find the correct XPU libraries.
+    """
+    # _set_gpu_runtime_env()  # cpp_extension consults the env
+    from torch.utils import cpp_extension
+
+    lpaths = cpp_extension.library_paths(device_type="xpu")
+    extra_ldflags: list[str] = []
+    if is_linux():
+        for path in lpaths:
+            if "torch/lib" in path:
+                # don't want to depend on pytorch
+                continue
+            # -rpath ensures the DLL can find its dependencies when loaded, even
+            # if the library path is non-standard.
+            extra_ldflags.extend([f"-L{path}", "-Xlinker", f"-rpath={path}"])
+
+        extra_ldflags.append("-lsycl")
+    else:
+        raise NotImplementedError(
+            "Unsupported env, failed to find xpu libs! Currently only Linux is supported."
+        )
+    return extra_ldflags
+
+
+def _sycl_arch_as_compile_option() -> str:
+    arc_option_map = {"Xe12": "intel_gpu_pvc", "Xe20": "intel_gpu_bmg_g21"}
+    arch = get_xpu_arch()
+    return arc_option_map.get(arch, "intel_gpu_pvc")
+
+
+def _sycl_compiler_options() -> list[str]:
+    options = [
+        "-DCUTLASS_ENABLE_SYCL",
+        "-DSYCL_INTEL_TARGET",
+        "-DCUTLASS_VERSIONS_GENERATED",
+        "-O3",
+        "-DNDEBUG",
+        "-std=c++20",
+        "-fPIC",
+        "-fsycl",
+        f"-fsycl-targets={_sycl_arch_as_compile_option()}",
+        "-Xspirv-translator",
+        "-spirv-ext=+SPV_INTEL_split_barrier,+SPV_INTEL_2d_block_io,+SPV_INTEL_subgroup_matrix_multiply_accumulate",
+        "-fno-sycl-instrument-device-code",
+        "-DMKL_ILP64",
+        "-MD",
+        "-MT",
+    ]
+    if config.cutlass.enable_debug_info:
+        options.extend(["-lineinfo", "-g", "-DCUTLASS_DEBUG_TRACE_LEVEL=1"])
+    return options
+
+
+def xpu_compile_command(
+    src_files: list[str],
+    dst_file: str,
+    dst_file_ext: str,
+    extra_args: list[str] | None = None,
+) -> str:
+    if extra_args is None:
+        extra_args = []
+    include_paths = _cutlass_include_paths()
+    sycl_lib_options = _sycl_lib_options()
+    sycl_compiler_options = _sycl_compiler_options()
+
+    options = (
+        extra_args
+        + ["-I" + path for path in include_paths]
+        + ["-isystem", "/include"]
+        + sycl_compiler_options
+        + sycl_lib_options
+    )
+    src_file = " ".join(src_files)
+    res = ""
+    if dst_file_ext == "o":
+        res = f"{_sycl_compiler()} {' '.join(options)} -c -o {dst_file} {src_file}"
+    elif dst_file_ext == "so":
+        options.append("-shared")
+        res = f"{_sycl_compiler()} {' '.join(options)} -o {dst_file} {src_file}"
+    elif dst_file_ext == "exe":
+        res = f"{_sycl_compiler()} {' '.join(options)} -o {dst_file} {src_file}"
+    else:
+        raise NotImplementedError(f"Unsupported output file suffix {dst_file_ext}!")
+
+    log.debug("XPU command: %s", res)
+    return res
diff --git a/torch/_inductor/codegen/xpu/device_op_overrides.py b/torch/_inductor/codegen/xpu/device_op_overrides.py
index 5d538ec20ca21..f6f2ff3eb6e5c 100644
--- a/torch/_inductor/codegen/xpu/device_op_overrides.py
+++ b/torch/_inductor/codegen/xpu/device_op_overrides.py
@@ -1,7 +1,5 @@
 from __future__ import annotations
 
-from typing import Optional
-
 from ..common import (
     DeviceOpOverrides,
     register_device_op_overrides,
@@ -59,8 +57,8 @@ def cpp_device_ptr(self) -> str:
         return "void *"
 
     def cpp_scratch(
-        self, idx: int, workspace: TritonScratchWorkspace, prefix: Optional[str] = None
-    ) -> Optional[tuple[list[str], str]]:
+        self, idx: int, workspace: TritonScratchWorkspace, prefix: str | None = None
+    ) -> tuple[list[str], str] | None:
         return [f"void *global_scratch_{idx} = 0;"], f"global_scratch_{idx}"
 
 
diff --git a/torch/_inductor/codegen/xpu/xpu_combined_scheduling.py b/torch/_inductor/codegen/xpu/xpu_combined_scheduling.py
new file mode 100644
index 0000000000000..b7c1b5eb3933e
--- /dev/null
+++ b/torch/_inductor/codegen/xpu/xpu_combined_scheduling.py
@@ -0,0 +1,135 @@
+# mypy: allow-untyped-defs
+from __future__ import annotations
+
+from typing import Any, TYPE_CHECKING
+
+from torch._inductor.scheduler import (
+    BaseSchedulerNode,
+    BaseScheduling,
+    FusedSchedulerNode,
+    Scheduler,
+    SchedulerNode,
+)
+
+from ..cutlass.scheduling import CUTLASSScheduling
+from ..triton import TritonScheduling
+
+
+if TYPE_CHECKING:
+    from collections.abc import Sequence
+    from typing import TypeAlias
+
+    from sympy import Expr
+
+    import torch
+    from torch.utils._ordered_set import OrderedSet
+
+    from ..common import BackendFeature
+
+    _IntLike: TypeAlias = int | Expr
+
+
+class XPUCombinedScheduling(BaseScheduling):
+    """
+    Scheduler for XPU Kernels, which delegates calls as appropriate
+    to the SYCL-C++ and Triton Schedulers, which both work for XPU devices
+    and use a unified-wrapper for codegen.
+
+    If Scheduling code needs to be specialized for the case of mixed Triton / SYCL C++ code,
+    this would also be the place to do it.
+    """
+
+    def __init__(self, scheduler: Scheduler | None) -> None:
+        super().__init__(scheduler)
+        self._triton_scheduling = TritonScheduling(scheduler)
+        self._cutlass_scheduling = CUTLASSScheduling(scheduler)
+
+    def get_backend_features(self, device: torch.device) -> OrderedSet[BackendFeature]:
+        return self._triton_scheduling.get_backend_features(device)
+
+    def choose_node_backend(self, node: BaseSchedulerNode) -> BaseScheduling:
+        if self._cutlass_scheduling.is_cutlass_template(node):
+            return self._cutlass_scheduling
+        return self._triton_scheduling
+
+    def can_fuse_vertical(
+        self, node1: BaseSchedulerNode, node2: BaseSchedulerNode
+    ) -> bool:
+        if self._cutlass_scheduling.can_fuse_vertical(node1, node2):
+            return True
+        elif self._cutlass_scheduling.is_cutlass_template(
+            node1
+        ) or self._cutlass_scheduling.is_cutlass_template(node2):
+            return False
+        return self._triton_scheduling.can_fuse_vertical(node1, node2)
+
+    def can_fuse_horizontal(
+        self, node1: BaseSchedulerNode, node2: BaseSchedulerNode
+    ) -> bool:
+        for node in (node1, node2):
+            if self._cutlass_scheduling.is_cutlass_template(node):
+                return self._cutlass_scheduling.can_fuse_horizontal(
+                    node1, node2
+                )  # always False at the moment
+        return self._triton_scheduling.can_fuse_horizontal(node1, node2)
+
+    def group_fn(
+        self, sizes: Sequence[Sequence[_IntLike]]
+    ) -> tuple[tuple[_IntLike, ...], ...]:
+        return self._triton_scheduling.group_fn(sizes)
+
+    def codegen_template(
+        self,
+        template_node: BaseSchedulerNode,
+        epilogue_nodes: Sequence[BaseSchedulerNode],
+        prologue_nodes: Sequence[BaseSchedulerNode],
+    ) -> str | None:
+        if self._cutlass_scheduling.is_cutlass_template(template_node):
+            assert not prologue_nodes
+            return self._cutlass_scheduling.codegen_template(
+                template_node, epilogue_nodes, prologue_nodes
+            )
+        else:
+            return self._triton_scheduling.codegen_template(
+                template_node, epilogue_nodes, prologue_nodes
+            )
+
+    def codegen_mix_order_reduction(self, node):
+        return self._triton_scheduling.codegen_mix_order_reduction(node)
+
+    def codegen_node(self, node: FusedSchedulerNode | SchedulerNode) -> None:
+        return self._triton_scheduling.codegen_node(node)
+
+    def codegen_sync(self) -> None:
+        return self._triton_scheduling.codegen_sync()
+
+    def flush(self) -> None:
+        return self._triton_scheduling.flush()
+
+    def codegen_combo_kernel(self, *args: Any, **kwargs: Any) -> None:
+        return self._triton_scheduling.codegen_combo_kernel(*args, **kwargs)
+
+    def benchmark_fused_nodes(
+        self, nodes: Sequence[BaseSchedulerNode]
+    ) -> tuple[float, str]:
+        return self._triton_scheduling.benchmark_fused_nodes(nodes)
+
+    def benchmark_codegened_module(self, module):
+        return self._triton_scheduling.benchmark_codegened_module(module)
+
+    def generate_kernel_code_from_nodes(
+        self,
+        nodes: Sequence[Any],
+        benchmark_kernel: bool = False,
+        hint_override: int | None = None,
+    ) -> str:
+        return self._triton_scheduling.generate_kernel_code_from_nodes(
+            nodes, benchmark_kernel, hint_override=hint_override
+        )
+
+    def benchmark_combo_kernel(
+        self, node_list: Sequence[BaseSchedulerNode], node_benchmark_results
+    ) -> tuple[float, float, list[str | None]]:
+        return self._triton_scheduling.benchmark_combo_kernel(
+            node_list, node_benchmark_results
+        )
diff --git a/torch/_inductor/codegen/xpu/xpu_env.py b/torch/_inductor/codegen/xpu/xpu_env.py
new file mode 100644
index 0000000000000..a74b4cf2a24c7
--- /dev/null
+++ b/torch/_inductor/codegen/xpu/xpu_env.py
@@ -0,0 +1,38 @@
+import functools
+import logging
+
+import torch
+from torch._inductor.utils import clear_on_fresh_cache
+
+
+log = logging.getLogger(__name__)
+
+
+@clear_on_fresh_cache
+@functools.lru_cache(1)
+def get_xpu_arch() -> str | None:
+    from torch.testing._internal.common_xpu import get_xpu_codename, XPUCodename
+
+    name2arch = {
+        XPUCodename.PVC: "Xe12",
+        XPUCodename.BMG: "Xe20",
+    }
+
+    codename = get_xpu_codename()
+    if not codename or codename not in name2arch:
+        log.warning("Unknown XPU codename, cannot determine architecture")
+        return None
+
+    return name2arch[codename]
+
+
+@clear_on_fresh_cache
+@functools.lru_cache(1)
+def get_xpu_version() -> str | None:
+    # string of version, like 20250101
+    try:
+        xpu_version = torch.version.xpu or ""
+        return xpu_version
+    except Exception:
+        log.exception("Error getting xpu version")
+        return None
diff --git a/torch/_inductor/comm_analysis.py b/torch/_inductor/comm_analysis.py
index ffcfd5d43a2a3..6f277daf47d55 100644
--- a/torch/_inductor/comm_analysis.py
+++ b/torch/_inductor/comm_analysis.py
@@ -3,13 +3,13 @@
 import math
 import operator
 from enum import IntEnum
-from typing import Any, Optional
+from typing import Any
 
 import sympy
 
 import torch
 import torch.utils._pytree as pytree
-from torch.fx.experimental.symbolic_shapes import size_hint
+from torch.fx.experimental.symbolic_shapes import optimization_hint
 from torch.fx.operator_schemas import normalize_function
 
 from . import ir
@@ -26,6 +26,7 @@ class NCCL_COLL(IntEnum):
     REDUCE_SCATTER = 2
     ALL_TO_ALL = 3
     UNSUPPORTED = 4
+    P2P = 5
 
 
 class NVIDIA_GPU_TYPE(IntEnum):
@@ -61,6 +62,8 @@ def get_collective_type_from_kernel_name(kernel_name: str) -> NCCL_COLL:
         return NCCL_COLL.REDUCE_SCATTER
     elif any(comm in kernel_name for comm in ("all_to_all", "alltoall")):
         return NCCL_COLL.ALL_TO_ALL
+    elif any(comm in kernel_name for comm in ("isend", "irecv", "batch_p2p")):
+        return NCCL_COLL.P2P
     else:
         return NCCL_COLL.UNSUPPORTED
 
@@ -83,7 +86,7 @@ def get_ir_node_size_numel(size: torch.Size, fallback: int = 4096 * 4096) -> int
 
 def get_fx_node_size_numel(size: torch.Size, fallback: int = 4096 * 4096) -> int:
     numel = functools.reduce(operator.mul, size, 1)
-    result = size_hint(numel, fallback=fallback)
+    result = optimization_hint(numel, fallback=fallback)
     return result
 
 
@@ -193,7 +196,7 @@ class NCCL_PROTO(IntEnum):
 ]
 
 
-def estimate_nccl_collective_runtime_nccl_estimator(snode) -> Optional[float]:  # type: ignore[no-untyped-def]
+def estimate_nccl_collective_runtime_nccl_estimator(snode) -> float | None:  # type: ignore[no-untyped-def]
     kernel = snode.node
     assert kernel is not None
     py_kernel_name = getattr(kernel, "python_kernel_name", "")
@@ -253,6 +256,9 @@ def estimate_nccl_collective_runtime_impl(
     if nRanks <= 1:
         return 0
 
+    if coll == NCCL_COLL.UNSUPPORTED:
+        return 0
+
     # Assumes ring algorithm
     nccl_algo = NCCL_ALGO.RING
     nccl_proto = NCCL_PROTO.LL
@@ -291,6 +297,9 @@ def estimate_nccl_collective_runtime_impl(
         nsteps = 2 * (nRanks - 1)
     elif coll in (NCCL_COLL.REDUCE_SCATTER, NCCL_COLL.ALL_GATHER):
         nsteps = nRanks - 1
+    elif coll == NCCL_COLL.P2P:
+        # assume 1 hop per pair
+        nsteps = 1
 
     # Convert bus BW to algorithm BW (tensor bytes / algoBW = actual execution time)
     ratio = (1.0 * nRanks) / nsteps  # type: ignore[possibly-undefined]
@@ -308,6 +317,9 @@ def estimate_nccl_collective_runtime_impl(
             nInterSteps = 0
     elif coll in (NCCL_COLL.REDUCE_SCATTER, NCCL_COLL.ALL_GATHER, NCCL_COLL.ALL_TO_ALL):
         nInterSteps = nNodes - 1
+    elif coll == NCCL_COLL.P2P:
+        # p2p may cross node bdry
+        nInterSteps = 1 if nNodes > 1 else 0
 
     # First compute latency in us; then at the end, convert it to ns
     latency = baseLat[nccl_algo][nccl_proto]
@@ -387,10 +399,18 @@ def add_inp_bytes(inp: torch.fx.Node):
 
     output_val = fx_node.meta.get("val", None)
 
-    if input_bytes is None or not isinstance(output_val, torch.Tensor):
+    if input_bytes is None or output_val is None:
         return 0
 
-    output_bytes = tensor_bytes(output_val)
+    # Coalesced collectives return a list of tensors
+    if isinstance(output_val, (list, tuple)):
+        output_bytes = sum(
+            tensor_bytes(t) for t in output_val if isinstance(t, torch.Tensor)
+        )
+    elif isinstance(output_val, torch.Tensor):
+        output_bytes = tensor_bytes(output_val)
+    else:
+        return 0
 
     return input_bytes + output_bytes
 
@@ -411,7 +431,7 @@ def estimate_fx_collective_memory_footprint(fx_node: torch.fx.Node) -> int:
 
 def estimate_nccl_collective_runtime_from_fx_node(
     fx_node: torch.fx.Node,
-    override_size: Optional[int] = None,
+    override_size: int | None = None,
     use_nccl_estimator: bool = True,
 ) -> float:
     """
@@ -453,21 +473,20 @@ def estimate_nccl_collective_runtime_from_fx_node(
     assert isinstance(fx_node.target, torch._ops.OpOverload)
     coll = get_collective_type_from_kernel_name(fx_node.target.name())
 
-    def _nccl_estimate() -> Optional[float]:
+    def _nccl_estimate() -> float | None:
         # TODO: Refactor with estimate_nccl_collective_runtime_nccl_estimator
-        from torch.distributed.distributed_c10d import (
-            _get_pg_default_device,
-            _resolve_process_group,
-            Backend,
-        )
+        from torch.distributed.distributed_c10d import _resolve_process_group, Backend
 
         pg = _resolve_process_group(group_name)
         if torch.distributed.distributed_c10d.get_backend(pg) == Backend.FAKE:
             # nccl estimator requires real process group
             return None
 
-        device = _get_pg_default_device(pg)
-        backend = pg._get_backend(device)
+        device = torch.device("cuda")
+        try:
+            backend = pg._get_backend(device)
+        except RuntimeError:
+            return None
         if not backend.supports_time_estimate:
             return None
 
@@ -480,9 +499,6 @@ def _tensor(size, dtype, device) -> torch.Tensor:  # type: ignore[no-untyped-def
                 device=device,
             )
 
-        def try_size_hint(s: sympy.Expr) -> int:
-            return V.graph.sizevars.optimization_hint(s, fallback=0)
-
         def to_real_tensor(e: Any) -> Any:
             if isinstance(e, torch.fx.Node):
                 return to_real_tensor(e.meta["val"])
@@ -495,9 +511,16 @@ def to_real_tensor(e: Any) -> Any:
 
         fn = fx_node.target
         assert isinstance(fn, torch._ops.OpOverload)
-        with torch.distributed._time_estimator(group=pg) as time_estimator:
+        with torch.distributed._time_estimator(
+            group=pg, device=device
+        ) as time_estimator:
             w = fn(*real_args, **real_kwargs)
-            torch.ops._c10d_functional.wait_tensor.default(w)
+            # Coalesced collectives return a list of tensors
+            if isinstance(w, (list, tuple)):
+                for t in w:
+                    torch.ops._c10d_functional.wait_tensor.default(t)
+            else:
+                torch.ops._c10d_functional.wait_tensor.default(w)
         est_time_us = time_estimator.estimated_time
         # -1000 constant is NCCL return in case of error during estimations.
         # Observed it for all_to_all estimations.
diff --git a/torch/_inductor/comm_lowering.py b/torch/_inductor/comm_lowering.py
index 0b23d92743ba0..e5dfd1fd18127 100644
--- a/torch/_inductor/comm_lowering.py
+++ b/torch/_inductor/comm_lowering.py
@@ -402,6 +402,71 @@ def _wait_tensor(inp):
         ir._WaitKernel.create_wait(c10d.wait_tensor.default, inp)
         return inp
 
+    @register_comm_lowering(c10d.isend)  # type: ignore[misc]
+    def _isend(inp, dst, tag, group_name):
+        inp = ir.ExternKernel.require_contiguous(inp)
+        return _create_out_of_place(c10d.isend.default, inp, dst, tag, group_name)
+
+    @register_comm_lowering(c10d.irecv)  # type: ignore[misc]
+    def _irecv(inp, src, tag, group_name):
+        inp = ir.ExternKernel.require_contiguous(inp)
+        ir._CollectiveKernel.create_inplace(
+            c10d.irecv.default, inp, src, tag, group_name
+        )
+        return inp
+
+    @register_comm_lowering(c10d.batch_p2p_ops)  # type: ignore[misc]
+    def _batch_p2p_ops(op_list, peer_list, tag_list, tensors, group_name):
+        tensors = [ir.ExternKernel.require_contiguous(t) for t in tensors]
+        kernel = c10d.batch_p2p_ops.default
+        with V.graph.fake_mode:
+            (
+                example_output,
+                tensor_args,
+                non_tensor_args,
+                unflatten_args,
+                unbacked_bindings,
+            ) = ir._CollectiveKernel.process_kernel(
+                kernel,
+                op_list,
+                peer_list,
+                tag_list,
+                tensors,
+                group_name,
+            )
+        assert not unbacked_bindings, f"{kernel} {unbacked_bindings}"
+        for op, tensor_arg in zip(op_list, tensor_args):
+            tensor_arg.realize()
+            if op == "irecv":
+                V.graph.mark_buffer_mutated(tensor_arg.get_name())
+
+        device = tensor_args[0].get_device()
+        packed = ir._CollectiveKernel(
+            ir.MultiOutputLayout(device=device),
+            kernel,
+            tensor_args,
+            non_tensor_args,
+            unflatten_args,
+        )
+
+        results = []
+        for i, (op, t, ex_out) in enumerate(zip(op_list, tensors, example_output)):
+            if op == "irecv":
+                packed.mutation_outputs.append(
+                    ir.MutationOutput(ir.NoneLayout(device=device), t, packed)
+                )
+                packed.alias_names.append(t.get_name())
+                results.append(t)
+            else:
+                # isend: 0-element placeholder output connected to the collective
+                placeholder = ir.MultiOutput(
+                    ir._CollectiveKernel.tensor_to_layout(ex_out),
+                    packed,
+                    [(list, i)],
+                )
+                results.append(ir.TensorBox.create(placeholder))
+        return results
+
 
 def register_symm_mem_lowerings():
     """
@@ -418,21 +483,65 @@ def register_symm_mem_lowerings():
         log.info("symm_mem ops not available, skipping symm_mem lowerings")
         return
 
+    from torch._library._out_variant import register_out_variant
+
+    # Register manual out variant mappings for symm_mem ops.
+    register_out_variant(
+        symm_mem.one_shot_all_reduce.default,
+        symm_mem.one_shot_all_reduce_out.default,
+    )
+    register_out_variant(
+        symm_mem.one_shot_all_reduce_copy.default,
+        symm_mem.one_shot_all_reduce_copy_out.default,
+    )
+
     from .lowering import register_lowering
 
+    def _copy_input_to_comm_buffer(
+        inp: ir.TensorBox,
+        comm_buffer_type: ir.CommBufferType,
+        group_name: "torch.distributed.distributed_c10d.GroupName",
+    ) -> ir.TensorBox:
+        """
+        Fallback: insert a Pointwise identity copy allocated in P2P via
+        CommBufferLayout.  Used when we don't control the input's allocation.
+        """
+        inp.realize()
+        copy = ir.Pointwise.create(
+            device=inp.get_device(),
+            dtype=inp.get_dtype(),
+            inner_fn=inp.make_loader(),
+            ranges=inp.get_size(),
+        )
+        realize_as_comm_buffer(copy, comm_buffer_type, group_name)
+        return copy
+
     def _maybe_realize_symm_mem(
         inp: ir.TensorBox,
         group_name: str,  # type: ignore[arg-type]
-    ) -> None:
+    ) -> ir.TensorBox:
         """
-        Helper to realize an input as symmetric memory buffer if possible.
+        Ensure inp is in P2P memory for a symm_mem collective.
+
+        If inductor controls the buffer's allocation (ComputedBuffer,
+        or any buffer with FlexibleLayout/FixedLayout), switch its
+        layout to CommBufferLayout in-place, zero-copy.
+
+        If inductor does not control allocation (e.g. InputBuffer),
+        insert a Pointwise identity copy into a new CommBufferLayout buffer.
+        This adds an extra Triton kernel. Returns the possibly new TensorBox.
+
+        TODO(tianrengao): eliminate the extra kernel for static-shape
+        InputBuffers by pre-allocating P2P memory in the wrapper and DMA .copy_()
         """
         if can_realize_as_comm_buffer(inp, ir.CommBufferType.SYMM_MEM):
             realize_as_comm_buffer(inp, ir.CommBufferType.SYMM_MEM, group_name)  # type: ignore[arg-type]
+            return inp
         else:
-            log.warning(
-                "Failed to realize the input as a symmetric memory buffer for symm_mem operation; "
-                "ensure the input is allocated as a symmetric memory buffer."
+            return _copy_input_to_comm_buffer(
+                inp,
+                ir.CommBufferType.SYMM_MEM,
+                group_name,  # type: ignore[arg-type]
             )
 
     @register_lowering(symm_mem.one_shot_all_reduce)
@@ -441,7 +550,7 @@ def _symm_mem_one_shot_all_reduce(
         reduce_op: str,
         group_name: str,
     ):
-        _maybe_realize_symm_mem(inp, group_name)
+        inp = _maybe_realize_symm_mem(inp, group_name)
         return pytree.tree_map(
             ir.TensorBox.create,
             ir.FallbackKernel.create(
@@ -459,7 +568,7 @@ def _symm_mem_one_shot_all_reduce_out(
         group_name: str,
         out: ir.TensorBox,
     ):
-        _maybe_realize_symm_mem(inp, group_name)
+        inp = _maybe_realize_symm_mem(inp, group_name)
         return pytree.tree_map(
             ir.TensorBox.create,
             ir.FallbackKernel.create(
@@ -478,7 +587,7 @@ def _symm_mem_one_shot_all_reduce_copy(
         reduce_op: str,
         group_name: str,
     ):
-        _maybe_realize_symm_mem(symm_buffer, group_name)
+        symm_buffer = _maybe_realize_symm_mem(symm_buffer, group_name)
         return pytree.tree_map(
             ir.TensorBox.create,
             ir.FallbackKernel.create(
@@ -498,7 +607,7 @@ def _symm_mem_one_shot_all_reduce_copy_out(
         group_name: str,
         out: ir.TensorBox,
     ):
-        _maybe_realize_symm_mem(symm_buffer, group_name)
+        symm_buffer = _maybe_realize_symm_mem(symm_buffer, group_name)
         return pytree.tree_map(
             ir.TensorBox.create,
             ir.FallbackKernel.create(
@@ -517,7 +626,7 @@ def _symm_mem_two_shot_all_reduce_(
         reduce_op: str,
         group_name: str,
     ):
-        _maybe_realize_symm_mem(inp, group_name)
+        inp = _maybe_realize_symm_mem(inp, group_name)
         ir.FallbackKernel.create(
             symm_mem.two_shot_all_reduce_.default,
             inp,
@@ -533,7 +642,7 @@ def _symm_mem_two_shot_all_reduce_out(
         group_name: str,
         output: ir.TensorBox,
     ):
-        _maybe_realize_symm_mem(inp, group_name)
+        inp = _maybe_realize_symm_mem(inp, group_name)
         return pytree.tree_map(
             ir.TensorBox.create,
             ir.FallbackKernel.create(
@@ -551,7 +660,7 @@ def _symm_mem_multimem_all_reduce_(
         reduce_op: str,
         group_name: str,
     ):
-        _maybe_realize_symm_mem(inp, group_name)
+        inp = _maybe_realize_symm_mem(inp, group_name)
         ir.FallbackKernel.create(
             symm_mem.multimem_all_reduce_.default,
             inp,
@@ -566,7 +675,7 @@ def _symm_mem_multimem_one_shot_all_reduce(
         reduce_op: str,
         group_name: str,
     ):
-        _maybe_realize_symm_mem(inp, group_name)
+        inp = _maybe_realize_symm_mem(inp, group_name)
         return pytree.tree_map(
             ir.TensorBox.create,
             ir.FallbackKernel.create(
@@ -584,7 +693,7 @@ def _symm_mem_multimem_one_shot_all_reduce_out(
         group_name: str,
         out: ir.TensorBox,
     ):
-        _maybe_realize_symm_mem(inp, group_name)
+        inp = _maybe_realize_symm_mem(inp, group_name)
         return pytree.tree_map(
             ir.TensorBox.create,
             ir.FallbackKernel.create(
@@ -604,7 +713,7 @@ def _symm_mem_multimem_one_shot_reduce_out(
         group_name: str,
         out: ir.TensorBox,
     ):
-        _maybe_realize_symm_mem(inp, group_name)
+        inp = _maybe_realize_symm_mem(inp, group_name)
         return pytree.tree_map(
             ir.TensorBox.create,
             ir.FallbackKernel.create(
@@ -623,7 +732,7 @@ def _symm_mem_multimem_all_gather_out(
         group_name: str,
         out: ir.TensorBox,
     ):
-        _maybe_realize_symm_mem(inp, group_name)
+        inp = _maybe_realize_symm_mem(inp, group_name)
         return pytree.tree_map(
             ir.TensorBox.create,
             ir.FallbackKernel.create(
@@ -641,7 +750,7 @@ def _symm_mem_reduce_scatter_out(
         split_last_dim: bool,
         output: ir.TensorBox,
     ):
-        _maybe_realize_symm_mem(inp, group_name)
+        inp = _maybe_realize_symm_mem(inp, group_name)
         return pytree.tree_map(
             ir.TensorBox.create,
             ir.FallbackKernel.create(
@@ -661,8 +770,8 @@ def _symm_mem_all_to_all_vdev(
         out_splits_offsets: ir.TensorBox,
         group_name: str,
     ):
-        _maybe_realize_symm_mem(inp, group_name)
-        _maybe_realize_symm_mem(out, group_name)
+        inp = _maybe_realize_symm_mem(inp, group_name)
+        out = _maybe_realize_symm_mem(out, group_name)
         ir.FallbackKernel.create(
             symm_mem.all_to_all_vdev.default,
             inp,
@@ -682,8 +791,8 @@ def _symm_mem_all_to_all_vdev_2d(
         group_name: str,
         major_align=None,
     ):
-        _maybe_realize_symm_mem(inp, group_name)
-        _maybe_realize_symm_mem(out, group_name)
+        inp = _maybe_realize_symm_mem(inp, group_name)
+        out = _maybe_realize_symm_mem(out, group_name)
         ir.FallbackKernel.create(
             symm_mem.all_to_all_vdev_2d.default,
             inp,
@@ -703,8 +812,8 @@ def _symm_mem_all_to_all_vdev_2d_offset(
         out_splits_offsets: ir.TensorBox,
         group_name: str,
     ):
-        _maybe_realize_symm_mem(inp, group_name)
-        _maybe_realize_symm_mem(out, group_name)
+        inp = _maybe_realize_symm_mem(inp, group_name)
+        out = _maybe_realize_symm_mem(out, group_name)
         ir.FallbackKernel.create(
             symm_mem.all_to_all_vdev_2d_offset.default,
             inp,
@@ -723,8 +832,8 @@ def _symm_mem_tile_reduce(
         group_name: str,
         reduce_op: str = "sum",
     ):
-        _maybe_realize_symm_mem(in_tile, group_name)
-        _maybe_realize_symm_mem(out_tile, group_name)
+        in_tile = _maybe_realize_symm_mem(in_tile, group_name)
+        out_tile = _maybe_realize_symm_mem(out_tile, group_name)
         ir.FallbackKernel.create(
             symm_mem.tile_reduce.default,
             in_tile,
@@ -743,9 +852,9 @@ def _symm_mem_multi_root_tile_reduce(
         group_name: str,
         reduce_op: str = "sum",
     ):
-        for in_tile in in_tiles:
-            _maybe_realize_symm_mem(in_tile, group_name)
-        _maybe_realize_symm_mem(out_tile, group_name)
+        for i, in_tile in enumerate(in_tiles):
+            in_tiles[i] = _maybe_realize_symm_mem(in_tile, group_name)
+        out_tile = _maybe_realize_symm_mem(out_tile, group_name)
         ir.FallbackKernel.create(
             symm_mem.multi_root_tile_reduce.default,
             in_tiles,
diff --git a/torch/_inductor/comms.py b/torch/_inductor/comms.py
index 6ea55ec2e87b0..34a1c2a693d8a 100644
--- a/torch/_inductor/comms.py
+++ b/torch/_inductor/comms.py
@@ -10,7 +10,7 @@
 import sys
 from collections import defaultdict
 from dataclasses import dataclass
-from typing import Any, Optional, TYPE_CHECKING, Union
+from typing import Any, TYPE_CHECKING
 
 import torch
 from torch._logging import trace_structured
@@ -183,7 +183,7 @@ def improvement(self):
         return self.initial_exposed - self.final_exposed
 
 
-def is_gemm_like(node: Optional[Union[IRNode, Operation]]) -> bool:
+def is_gemm_like(node: IRNode | Operation | None) -> bool:
     if node is None:
         return False
 
@@ -266,8 +266,8 @@ def _initialize_memory_tracking(snodes, graph_inputs, graph_outputs):
 def _initialize_double_linked_list(
     snodes: list[BaseSchedulerNode],
 ) -> tuple[
-    dict[BaseSchedulerNode, Optional[BaseSchedulerNode]],
-    dict[BaseSchedulerNode, Optional[BaseSchedulerNode]],
+    dict[BaseSchedulerNode, BaseSchedulerNode | None],
+    dict[BaseSchedulerNode, BaseSchedulerNode | None],
     BaseSchedulerNode,
 ]:
     """Create double-linked list structure from snodes"""
@@ -348,9 +348,9 @@ def contains_async_collective(snode):
 
 
 def _group_nodes_from_linked_list(
-    head: Optional[BaseSchedulerNode],
-    tail: Optional[BaseSchedulerNode],
-    next_dict: dict[BaseSchedulerNode, Optional[BaseSchedulerNode]],
+    head: BaseSchedulerNode | None,
+    tail: BaseSchedulerNode | None,
+    next_dict: dict[BaseSchedulerNode, BaseSchedulerNode | None],
 ) -> list[BaseSchedulerNode]:
     """
     Traverse doubly-linked list from head to tail and return nodes as a list.
@@ -391,7 +391,7 @@ def _is_corresponding_collective_wait(
 
 def _coll_exposed_communication_time(
     collective_snode: BaseSchedulerNode,
-    next_dict: dict[BaseSchedulerNode, Optional[BaseSchedulerNode]],
+    next_dict: dict[BaseSchedulerNode, BaseSchedulerNode | None],
     runtimes: dict[BaseSchedulerNode, float],
     node_output_sets: dict[BaseSchedulerNode, frozenset[str]],
     node_dep_sets: dict[BaseSchedulerNode, frozenset[str]],
@@ -452,7 +452,7 @@ def accumulate_time(_snode: BaseSchedulerNode) -> None:
 def _wait_exposed_communication_time(
     wait_snode: BaseSchedulerNode,
     head: BaseSchedulerNode,
-    prev_dict: dict[BaseSchedulerNode, Optional[BaseSchedulerNode]],
+    prev_dict: dict[BaseSchedulerNode, BaseSchedulerNode | None],
     runtimes: dict[BaseSchedulerNode, float],
     node_output_sets: dict[BaseSchedulerNode, frozenset[str]],
     node_dep_sets: dict[BaseSchedulerNode, frozenset[str]],
@@ -513,8 +513,8 @@ def _perform_double_linked_list_swap(
     candidate: BaseSchedulerNode,
     group_head: BaseSchedulerNode,
     group_tail: BaseSchedulerNode,
-    prev_dict: dict[BaseSchedulerNode, Optional[BaseSchedulerNode]],
-    next_dict: dict[BaseSchedulerNode, Optional[BaseSchedulerNode]],
+    prev_dict: dict[BaseSchedulerNode, BaseSchedulerNode | None],
+    next_dict: dict[BaseSchedulerNode, BaseSchedulerNode | None],
     head: BaseSchedulerNode,
 ) -> BaseSchedulerNode:
     """
@@ -707,7 +707,7 @@ def _find_buffers_with_changed_last_use(
     gns: list[BaseSchedulerNode],
     buf_to_snode_last_use: dict,
     candidate_buffer_map: dict[BaseSchedulerNode, OrderedSet],
-) -> dict[BaseSchedulerNode, list[Union[FreeableInputBuffer, Any]]]:
+) -> dict[BaseSchedulerNode, list[FreeableInputBuffer | Any]]:
     """
     Find buffers whose last use will change after swapping candidate with group.
 
@@ -725,7 +725,7 @@ def _find_buffers_with_changed_last_use(
         Dict mapping group nodes to buffers that will change their last-use node
     """
     group_n_to_bufs_after_swap_dealloc_by_candidate: dict[
-        BaseSchedulerNode, list[Union[FreeableInputBuffer, Any]]
+        BaseSchedulerNode, list[FreeableInputBuffer | Any]
     ] = defaultdict(list)
 
     # Optimization: only check buffers where candidate is a successor
@@ -743,7 +743,7 @@ def _find_buffers_with_changed_last_use(
 
 def _is_node_groupable_for_reorder(
     candidate: BaseSchedulerNode,
-) -> tuple[bool, Optional[str]]:
+) -> tuple[bool, str | None]:
     """
     Check if a candidate node can be grouped with collective during reordering.
 
@@ -774,7 +774,7 @@ def _is_node_groupable_for_reorder(
 def _format_and_log_reordering_stats(
     stats: dict[BaseSchedulerNode, ReorderInfo],
     head: BaseSchedulerNode,
-    next_dict: dict[BaseSchedulerNode, Optional[BaseSchedulerNode]],
+    next_dict: dict[BaseSchedulerNode, BaseSchedulerNode | None],
     original_snodes_num: int,
     peak_memory: int,
     name_to_freeable_input_buf: dict,
@@ -925,12 +925,12 @@ def _reorder_communication_preserving_peak_memory_internal(
 
     _prev, _next, _head = _initialize_double_linked_list(snodes)
 
-    debug_num_collectives_to_reorder: Optional[int] = (
+    debug_num_collectives_to_reorder: int | None = (
         config_comms.reorder_iterative_debug_limit_to_reorder
     )
 
     num_processed_collectives: int = 0
-    curr: Optional[BaseSchedulerNode] = _head
+    curr: BaseSchedulerNode | None = _head
     debug_iterative_memory_recompute = (
         config_comms.reorder_iterative_debug_memory_recompute
     )
@@ -1416,7 +1416,7 @@ def improvement(self):
 
 def _is_node_groupable_for_sink_waits(
     candidate: BaseSchedulerNode,
-) -> tuple[bool, Optional[str]]:
+) -> tuple[bool, str | None]:
     """
     Check if a candidate node can be grouped during sink_waits pass.
 
@@ -1584,8 +1584,8 @@ def _perform_double_linked_list_swap_sink_waits(
     candidate: BaseSchedulerNode,
     group_head: BaseSchedulerNode,
     group_tail: BaseSchedulerNode,
-    prev_dict: dict[BaseSchedulerNode, Optional[BaseSchedulerNode]],
-    next_dict: dict[BaseSchedulerNode, Optional[BaseSchedulerNode]],
+    prev_dict: dict[BaseSchedulerNode, BaseSchedulerNode | None],
+    next_dict: dict[BaseSchedulerNode, BaseSchedulerNode | None],
     head: BaseSchedulerNode,
 ) -> BaseSchedulerNode:
     """
@@ -1632,7 +1632,7 @@ def _perform_double_linked_list_swap_sink_waits(
 def _format_and_log_sink_waits_stats(
     stats: dict[BaseSchedulerNode, SinkWaitInfo],
     head: BaseSchedulerNode,
-    next_dict: dict[BaseSchedulerNode, Optional[BaseSchedulerNode]],
+    next_dict: dict[BaseSchedulerNode, BaseSchedulerNode | None],
     original_snodes_num: int,
     peak_memory: int,
     name_to_freeable_input_buf: dict,
@@ -1724,7 +1724,7 @@ def _find_buffers_with_changed_last_use_sink_waits(
     gns: list[BaseSchedulerNode],
     buf_to_snode_last_use: dict,
     candidate_buffer_map: dict[BaseSchedulerNode, OrderedSet],
-) -> dict[BaseSchedulerNode, list[Union[FreeableInputBuffer, Any]]]:
+) -> dict[BaseSchedulerNode, list[FreeableInputBuffer | Any]]:
     """
     Find buffers whose last use will change after swapping in sink_waits pass.
 
@@ -1742,7 +1742,7 @@ def _find_buffers_with_changed_last_use_sink_waits(
         Dict mapping group nodes to buffers that will change their last-use node
     """
     group_n_to_bufs_after_swap_dealloc_instead_of_candidate: dict[
-        BaseSchedulerNode, list[Union[FreeableInputBuffer, Any]]
+        BaseSchedulerNode, list[FreeableInputBuffer | Any]
     ] = defaultdict(list)
 
     # Optimization: only check buffers where candidate is a successor
@@ -1811,13 +1811,13 @@ def _sink_waits_iterative_internal(
         for snode in snodes
     }
 
-    curr: Optional[BaseSchedulerNode] = snodes[-1]
+    curr: BaseSchedulerNode | None = snodes[-1]
 
     processed_waits = OrderedSet()  # type: ignore[var-annotated]
     debug_iterative_memory_recompute = (
         config_comms.reorder_iterative_debug_memory_recompute
     )
-    debug_num_sink_waits_to_reorder: Optional[int] = (
+    debug_num_sink_waits_to_reorder: int | None = (
         config_comms.sink_waits_iterative_debug_limit_to_sink
     )
 
diff --git a/torch/_inductor/comms_debug.py b/torch/_inductor/comms_debug.py
index 20c9779a4ef3f..f40df1a2436fd 100644
--- a/torch/_inductor/comms_debug.py
+++ b/torch/_inductor/comms_debug.py
@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Union
+from typing import TYPE_CHECKING
 
 from torch._logging import trace_structured
 
@@ -26,7 +26,7 @@ def _debug_iterative_memory_recompute(
     snodes_allocfree: dict[BaseSchedulerNode, SNodeMemory],
     tlparse_name: str,
     gn_to_bufs_last_use: dict[
-        BaseSchedulerNode, list[Union[FreeableInputBuffer, SchedulerBuffer]]
+        BaseSchedulerNode, list[FreeableInputBuffer | SchedulerBuffer]
     ],
 ) -> bool:
     iterative_recompute_error = False
diff --git a/torch/_inductor/compile_fx.py b/torch/_inductor/compile_fx.py
index 95b777eab5ad2..bd7e96e0f727a 100644
--- a/torch/_inductor/compile_fx.py
+++ b/torch/_inductor/compile_fx.py
@@ -19,7 +19,7 @@
 from inspect import currentframe
 from itertools import count
 from operator import attrgetter
-from typing import Any, Optional, TYPE_CHECKING, TypeVar, Union
+from typing import Any, TYPE_CHECKING, TypeVar
 from typing_extensions import Never, override, ParamSpec, Protocol, TypedDict, Unpack
 from unittest import mock
 
@@ -35,6 +35,7 @@
     logging as dynamo_logging,
     utils as dynamo_utils,
 )
+from torch._dynamo.backends import common as dynamo_common
 from torch._dynamo.device_interface import get_interface_for_device
 from torch._dynamo.repro.after_aot import wrap_compiler_debug
 from torch._dynamo.utils import (
@@ -46,10 +47,11 @@
     flatten_graph_inputs,
     get_inputs_devices,
     get_metrics_context,
+    GmWrapper,
     lazy_format_graph_code,
     set_feature_use,
 )
-from torch._functorch import config as functorch_config
+from torch._functorch import aot_autograd, config as functorch_config
 from torch._functorch._aot_autograd.subclass_parametrization import (
     unwrap_tensor_subclass_parameters,
 )
@@ -62,6 +64,7 @@
 from torch._inductor.codecache import code_hash, FxGraphCache, output_code_log
 from torch._inductor.cudagraph_utils import (
     BoxedDeviceIndex,
+    cudagraphs_log,
     format_default_skip_message,
     log_cudagraph_skip_and_bump_counter,
     PlaceholderInfo,
@@ -85,6 +88,7 @@
     count_tangents,
     fresh_cache,
     get_all_devices,
+    get_static_bw_input_idxs,
     InputType,
     is_gpu,
     should_assume_input_aligned,
@@ -101,7 +105,6 @@
 from torch.monitor import _WaitCounter
 from torch.utils._ordered_set import OrderedSet
 
-from .._dynamo.backends.common import aot_autograd
 from .._dynamo.exc import ShortenTraceback, SkipFrame
 from ..fx._lazy_graph_module import _use_lazy_graph_module
 from ..fx.graph import _PyTreeCodeGen
@@ -157,20 +160,15 @@ def log_optimus_to_scuba(*args: object, **kwargs: object) -> None:
     from torch._inductor.fb.utils import log_optimus_to_scuba, time_and_log
 
 if TYPE_CHECKING:
-    import types
-
     from torch._functorch._aot_autograd.schemas import (
         FQN,
         GraphInputName,
         GraphSignature,
     )
 
-    CompileFxOutput = Union[
-        Callable[[list[object]], Sequence[torch.Tensor]],
-        str,
-        list[str],
-        Weights,
-    ]
+    CompileFxOutput = (
+        Callable[[list[object]], Sequence[torch.Tensor]] | str | list[str] | Weights
+    )
 
 
 class FxCompileMode(enum.Enum):
@@ -407,7 +405,7 @@ def _unlift_graph(
 
     _resolve_name_collision(mod, gm)
 
-    state_dict: dict[str, Union[torch.nn.parameter.Parameter, torch.Tensor]] = {}
+    state_dict: dict[str, torch.nn.parameter.Parameter | torch.Tensor] = {}
     for name, param in mod.named_parameters(remove_duplicate=False):
         state_dict[name] = param
         _assign_attr(
@@ -426,7 +424,7 @@ def _unlift_graph(
         )
 
     placeholder_nodes = gm.graph.find_nodes(op="placeholder")
-    lifted_inputs: list[Optional[FQN]] = []
+    lifted_inputs: list[FQN | None] = []
 
     # In AOTI, module parameters and buffers are not lifted as graph inputs.
     # As a result, mutation to buffers has side effect which makes their initial
@@ -456,7 +454,7 @@ def _unlift_graph(
     user_input_mutations = graph_signature.user_inputs_to_mutate
     output_tokens = graph_signature.output_tokens
     for idx, out in enumerate(outputs):
-        value: Optional[Union[FQN, GraphInputName]] = None
+        value: FQN | GraphInputName | None = None
 
         if idx < len(buffer_mutations) + len(user_input_mutations) + len(output_tokens):
             name = GraphOutputName(out.name)
@@ -529,7 +527,7 @@ def _recursive_pre_grad_passes(
 def _recursive_joint_graph_passes(
     gm: GraphModule,
     skip_invoke_subgraph: bool = False,
-    input_device: Optional[torch.device] = None,
+    input_device: torch.device | None = None,
 ) -> GraphModule:
     def _run_on_sub_graph_module(subgraph_name: str) -> None:
         subgraph = getattr(gm, subgraph_name)
@@ -586,8 +584,8 @@ def _recursive_post_grad_passes(gm: GraphModule, is_inference: bool = False) ->
 def split_const_gm(
     gm: GraphModule,
     skip_constructor: bool = True,
-    lifted_constant_names: Optional[list[str]] = None,
-    skip_folding_node_fn: Optional[Callable[[torch.fx.Node], bool]] = None,
+    lifted_constant_names: list[str] | None = None,
+    skip_folding_node_fn: Callable[[torch.fx.Node], bool] | None = None,
 ) -> tuple[GraphModule, dict[str, int]]:
     """
     This function takes an GraphModule input "gm".
@@ -743,7 +741,7 @@ def fake_tensor_prop(
 
 # pass config dict back to user
 def get_patched_config_dict(
-    config_patches: Optional[Union[str, dict[str, Any]]] = None,
+    config_patches: str | dict[str, Any] | None = None,
 ) -> dict[str, Any]:
     with config.patch(config_patches):
         return config.get_config_copy()
@@ -762,17 +760,18 @@ def with_fresh_cache_if_config() -> Generator[None, None, None]:
 
 
 class _CompileFxKwargs(TypedDict, total=False):
-    cudagraphs: Optional[BoxedBool]
+    cudagraphs: BoxedBool | None
     static_input_idxs: Sequence[int]
     is_backward: bool
-    graph_id: Optional[int]
+    graph_id: int | None
     cpp_wrapper: bool
     aot_mode: bool
     is_inference: bool
-    layout_opt: Optional[bool]
-    extern_node_serializer: Optional[Callable[[list[ExternKernelNode]], Any]]
-    boxed_forward_device_index: Optional[BoxedDeviceIndex]
+    layout_opt: bool | None
+    extern_node_serializer: Callable[[list[ExternKernelNode]], Any] | None
+    boxed_forward_device_index: BoxedDeviceIndex | None
     fx_wrapper: bool
+    get_decomp_fn: Callable[..., dict[Any, Callable[..., Any]]]
 
 
 class _CompileFxCallable(Protocol):
@@ -780,6 +779,7 @@ def __call__(
         self,
         gm: GraphModule,
         example_inputs: Sequence[InputType],
+        compile_region_name: str | None = None,
         **kwargs: Unpack[_CompileFxKwargs],
     ) -> OutputCode: ...
 
@@ -787,6 +787,7 @@ def __call__(
 def compile_fx_inner(
     gm: GraphModule,
     example_inputs: Sequence[InputType],
+    compile_region_name: str | None = None,
     **kwargs: Unpack[_CompileFxKwargs],
 ) -> OutputCode:
     kwargs.setdefault("cudagraphs", None)
@@ -826,6 +827,7 @@ def compile_fx_inner(
         return wrap_compiler_debug(_compile_fx_inner, compiler_name="inductor")(
             gm,
             example_inputs,
+            compile_region_name=compile_region_name,
             **kwargs,
         )
 
@@ -834,6 +836,7 @@ def compile_fx_inner(
 def _compile_fx_inner(
     gm: GraphModule,
     example_inputs: Sequence[InputType],
+    compile_region_name: str | None = None,
     **graph_kwargs: Unpack[_CompileFxKwargs],
 ) -> OutputCode:
     """
@@ -892,6 +895,7 @@ def _compile_fx_inner(
         save_args_for_compile_fx_inner(
             gm,
             example_inputs,
+            compile_region_name=compile_region_name,
             **graph_kwargs,
         )
 
@@ -945,7 +949,7 @@ def _compile_fx_inner(
             ):
                 input._is_inductor_static = True  # type: ignore[attr-defined]
 
-        mb_compiled_graph: Optional[OutputCode] = None
+        mb_compiled_graph: OutputCode | None = None
         key_info = None
         cache_info = None
         remote_cache = None
@@ -988,7 +992,11 @@ def _compile_fx_inner(
             TritonBundler.begin_compile()
             try:
                 mb_compiled_graph = fx_codegen_and_compile(
-                    gm, example_inputs, inputs_to_check, **graph_kwargs
+                    gm,
+                    example_inputs,
+                    inputs_to_check,
+                    compile_region_name=compile_region_name,
+                    **graph_kwargs,
                 )
                 assert mb_compiled_graph is not None
                 (
@@ -1020,7 +1028,11 @@ def _compile_fx_inner(
             )
             try:
                 mb_compiled_graph = fx_codegen_and_compile(
-                    gm, example_inputs, inputs_to_check, **graph_kwargs
+                    gm,
+                    example_inputs,
+                    inputs_to_check,
+                    compile_region_name=compile_region_name,
+                    **graph_kwargs,
                 )
             except Exception as e:
                 raise InductorError(e, currentframe()).with_traceback(
@@ -1035,7 +1047,11 @@ def _compile_fx_inner(
             TritonBundler.begin_compile()
             try:
                 mb_compiled_graph = fx_codegen_and_compile(
-                    gm, example_inputs, inputs_to_check, **graph_kwargs
+                    gm,
+                    example_inputs,
+                    inputs_to_check,
+                    compile_region_name=compile_region_name,
+                    **graph_kwargs,
                 )
                 assert mb_compiled_graph is not None
                 mb_compiled_graph._time_taken_ns = time.time_ns() - start_time
@@ -1080,6 +1096,8 @@ def _compile_fx_inner(
 
         assert mb_compiled_graph is not None
         compiled_graph = mb_compiled_graph
+        if isinstance(compiled_graph, CompiledFxGraph):
+            compiled_graph.compile_region_name = compile_region_name
 
         # Logging and observability: we log a single chromium event
         # and a tlparse log for every cache action.
@@ -1127,6 +1145,10 @@ def _compile_fx_inner(
             )
         compiled_graph.post_compile(example_inputs, constants, graph_kwargs)
 
+        policy = config.cudagraph_policy
+        if policy is not None:
+            compiled_graph = policy.wrap_output(compiled_graph)
+
     log.debug("FX codegen and compilation took %.3fs", time.time() - start)
 
     # This message is for printing overview information of inductor mm counts, shapes,etc after lowering
@@ -1194,6 +1216,7 @@ class FxCompile(ABC):
 
     # Some stats for logging/debugging
     _compile_stats: dict[type[FxCompile], _FxCompileStat] = defaultdict(_FxCompileStat)
+    compile_region_name: str | None = None
 
     # TODO: We should probably eventually add some kind of async version of this
     # so we can kick off a compile and then go do other things - but we'll need
@@ -1233,15 +1256,17 @@ def codegen_and_compile(
         cudagraphs: BoxedBool = graph_kwargs["cudagraphs"]
         static_input_idxs: Sequence[int] = graph_kwargs.get("static_input_idxs", ())
         is_backward: bool = graph_kwargs.get("is_backward", False)
-        graph_id: Optional[int] = graph_kwargs.get("graph_id", None)
+        graph_id: int | None = graph_kwargs.get("graph_id", None)
         cpp_wrapper: bool = graph_kwargs.get("cpp_wrapper", False)
         fx_wrapper: bool = graph_kwargs.get("fx_wrapper", False)
         aot_mode: bool = V.aot_compilation
         is_inference: bool = graph_kwargs.get("is_inference", False)
-        extern_node_serializer: Optional[Callable[[list[ExternKernelNode]], Any]] = (
+        extern_node_serializer: Callable[[list[ExternKernelNode]], Any] | None = (
             graph_kwargs.get("extern_node_serializer", None)
         )
-
+        get_decomp_fn: Callable[..., dict[Any, Callable[..., Any]]] = graph_kwargs.get(
+            "get_decomp_fn", select_decomp_table
+        )
         with (
             _WaitCounter("pytorch.wait_counter.actual_codegen_and_compile").guard(),
             dynamo_utils.preserve_rng_state(),
@@ -1437,6 +1462,7 @@ def codegen_and_compile(
                         is_backward=is_backward,
                         is_const_graph=True,
                         fx_wrapper=fx_wrapper,
+                        get_decomp_fn=get_decomp_fn,
                     )
                     with (
                         V.set_graph_handler(const_graph),
@@ -1471,6 +1497,7 @@ def codegen_and_compile(
                     const_module=const_graph,
                     inputs_to_check=inputs_to_check,
                     fx_wrapper=fx_wrapper,
+                    get_decomp_fn=get_decomp_fn,
                 )
                 metrics_helper = metrics.CachedMetricsHelper()
 
@@ -1483,7 +1510,7 @@ def codegen_and_compile(
                     distributed_autotune.graph_context(),
                 ):
                     graph.run(*example_inputs)
-                    output_strides: list[Optional[tuple[_StrideExprStr, ...]]] = []
+                    output_strides: list[tuple[_StrideExprStr, ...] | None] = []
                     if graph.graph_outputs is not None:
                         # We'll put the output strides in the compiled graph so we
                         # can later return them to the caller via TracingContext
@@ -1735,6 +1762,7 @@ def codegen_and_compile(
                         cudagraphs,
                         example_inputs,
                         static_input_idxs,
+                        self.compile_region_name,
                         graph_kwargs,
                         inputs_to_check,
                         runnable_graph_str,
@@ -1751,6 +1779,7 @@ def fx_codegen_and_compile(
     # This is derivable from the other inputs to this function, but we pass it
     # in explicitly because it's nontrivial to compute
     inputs_to_check: Sequence[int],
+    compile_region_name: str | None = None,
     **graph_kwargs: Unpack[_CompileFxKwargs],
 ) -> OutputCode:
     scheme: FxCompile
@@ -1776,6 +1805,9 @@ def fx_codegen_and_compile(
         )
         # pyrefly: ignore [unbound-name]
         scheme = _AsyncFxCompile(scheme)
+        scheme._compile.compile_region_name = (
+            compile_region_name  # pyrefly: ignore[attr-defined]
+        )
 
     if fx_compile_progressive:
         from .compile_fx_async import _ProgressiveFxCompile
@@ -1790,9 +1822,15 @@ def fx_codegen_and_compile(
 
         # Use in-process compile for the fast version
         fast_scheme = _InProcessFxCompile()
+        fast_scheme.compile_region_name = compile_region_name
 
         # pyrefly: ignore [unbound-name]
         scheme = _ProgressiveFxCompile(fast_scheme, scheme, progression_configs)
+        scheme._optimized_compile.compile_region_name = (
+            compile_region_name  # pyrefly: ignore[attr-defined]
+        )
+
+    scheme.compile_region_name = compile_region_name  # pyrefly: ignore[unbound-name]
 
     # pyrefly: ignore [unbound-name]
     return scheme.codegen_and_compile(gm, example_inputs, inputs_to_check, graph_kwargs)
@@ -1838,7 +1876,7 @@ def cudagraphify(
     static_input_idxs: Sequence[int] = (),
     *,
     device_index: int,
-    stack_traces: list[Optional[str]],
+    stack_traces: list[str | None],
     is_backward: bool,
     is_inference: bool,
     constants: tuple[torch.Tensor, ...] = (),
@@ -1997,8 +2035,8 @@ def compile_fx_aot(
     model_: GraphModule,
     example_inputs_: list[InputType],
     inner_compile: _CompileFxCallable = compile_fx_inner,
-    config_patches: Optional[dict[str, Any]] = None,
-) -> Union[list[Union[str, Weights]], str, GraphModule]:
+    config_patches: dict[str, Any] | None = None,
+) -> list[str | Weights] | str | GraphModule:
     assert isinstance(model_, GraphModule), model_
 
     # [See NOTE] Unwrapping subclasses AOT
@@ -2069,6 +2107,7 @@ def fw_compiler_freezing(
     dynamo_model: GraphModule,
     num_example_inputs: int,
     inner_compile: Callable[..., Any],
+    # TODO: Take compiler_config_extra instead
     cudagraphs: BoxedBool,
     graph_id: int,
     forward_device: BoxedDeviceIndex,
@@ -2078,7 +2117,8 @@ def fw_compiler_freezing(
     # partition_fn won't be called
     inputs_devices = get_inputs_devices(aot_example_inputs, aot_autograd_model)
     aot_autograd_model = _recursive_joint_graph_passes(
-        aot_autograd_model, input_device=next(iter(inputs_devices))
+        aot_autograd_model,
+        input_device=next(iter(inputs_devices)),
     )
 
     layout_opt = GraphLowering.decide_layout_opt(aot_autograd_model, is_inference=True)
@@ -2169,20 +2209,27 @@ def wrapper(args: list[object]) -> Sequence[torch.Tensor]:
 
 
 def get_cpp_wrapper_config() -> dict[str, object]:
-    if config.triton.cudagraphs:
+    if config.triton.cudagraphs and config.graph_partition:
         log_cudagraph_skip_and_bump_counter(
-            format_default_skip_message("cpp wrapper enabled")
+            format_default_skip_message(
+                "cpp-wrapper does not support graph partition yet"
+            )
         )
 
+    autotune_at_compile_time = (
+        config.triton.autotune_at_compile_time
+        if config.triton.autotune_at_compile_time is not None
+        # Default to True for AOTI. Subject to change in future.
+        else has_triton() and V.aot_compilation
+    )
     return {
-        # Set autotune_at_compile_time to True as default if the option is not explicitly set
-        "triton.autotune_at_compile_time": (
-            config.triton.autotune_at_compile_time
-            if config.triton.autotune_at_compile_time is not None
-            else has_triton()
+        "triton.autotune_at_compile_time": autotune_at_compile_time,
+        "triton.autotune_cublasLt": not autotune_at_compile_time,
+        "triton.cudagraphs": (
+            config.triton.cudagraphs
+            and not V.aot_compilation
+            and not config.graph_partition
         ),
-        "triton.autotune_cublasLt": False,
-        "triton.cudagraphs": False,  # TODO: to be removed
         "triton.store_cubin": True,
     }
 
@@ -2217,10 +2264,12 @@ def partition_fn(
         # in partitioning.
         inputs_devices = get_inputs_devices(joint_inputs, gm)
         gm = _recursive_joint_graph_passes(
-            gm, skip_invoke_subgraph=True, input_device=next(iter(inputs_devices))
+            gm,
+            skip_invoke_subgraph=True,
+            input_device=next(iter(inputs_devices)),
         )
 
-    static_lifetime_input_indices: Optional[list[int]] = kwargs.pop(  # type: ignore[assignment]
+    static_lifetime_input_indices: list[int] | None = kwargs.pop(  # type: ignore[assignment]
         "static_lifetime_input_indices", None
     )
 
@@ -2257,20 +2306,65 @@ def get_num_model_outputs(model: GraphModule) -> int:
     return len(model_outputs)
 
 
+def cudagraph_annotation_context(
+    cudagraphs: BoxedBool,
+) -> contextlib.AbstractContextManager[None]:
+    # When an annotation force-enables cudagraphs but the global config has them
+    # off, patch config.triton.cudagraphs for the duration of compilation,
+    # so existing codepaths that access config.triton.cudagraphs work
+    if cudagraphs.value and not config.triton.cudagraphs:
+        return config.patch({"triton.cudagraphs": True})
+    return contextlib.nullcontext()
+
+
 @dataclass(frozen=True)
 class CompilerConfigExtra:
     cudagraphs: BoxedBool
     graph_id: int
     forward_device: BoxedDeviceIndex
+    forward_is_partitioned: BoxedBool
+    cudagraphs_bwd_override: bool | None = None
+
 
+def create_compiler_config_extra(
+    gm: GraphModule | GmWrapper,
+) -> CompilerConfigExtra:
+    gm_meta = gm.meta if isinstance(gm, GraphModule) else None
 
-def create_compiler_config_extra(config: types.ModuleType) -> CompilerConfigExtra:
     # Although cudagraphs may have been enabled via config, various
     # conditions (which are tested within the bowels of Inductor) may
     # force cudagraphs to be disabled.  This mutable box lets us retrieve
     # the final determination if cudagraphs actually can be used or not.
     cudagraphs = BoxedBool(config.triton.cudagraphs)
 
+    cudagraphs_bwd_override: bool | None = None
+
+    # Override cudagraphs BoxedBool based on override_cudagraphs annotation.
+    # Disabling fwd disables bwd (copying activations isn't profitable),
+    # so cudagraphs_bwd_override is only needed for fwd=True / bwd=False.
+    if (
+        gm_meta is not None
+        and (annotation := gm_meta.get("cudagraph_annotation")) is not None
+    ):
+        if annotation.fwd is not None and annotation.fwd != config.triton.cudagraphs:
+            cudagraphs = BoxedBool(annotation.fwd)
+            if annotation.fwd:
+                cudagraphs_log.info(
+                    "enabling cudagraphs due to override_cudagraphs annotation"
+                )
+            else:
+                log_cudagraph_skip_and_bump_counter(
+                    "disabling cudagraphs due to override_cudagraphs annotation"
+                )
+
+        # bwd override only matters when fwd enables cudagraphs but bwd
+        # explicitly disables them.
+        if cudagraphs.value and annotation.bwd is not None and not annotation.bwd:
+            cudagraphs_bwd_override = annotation.bwd
+            log_cudagraph_skip_and_bump_counter(
+                "disabling cudagraphs for backward due to override_cudagraphs annotation"
+            )
+
     # TODO: The modern style is to use CompileId from TracingContext to
     # identify Inductor compilation.  However, this CompileId cannot
     # uniquely identify multiple Inductor compilations that arise from
@@ -2280,10 +2374,17 @@ def create_compiler_config_extra(config: types.ModuleType) -> CompilerConfigExtr
     # See [Backward Generation Handling]
     forward_device = BoxedDeviceIndex(None)
 
+    # Set by the forward compilation when it is partitioned for CUDA graphs.
+    # The backward reads this to decide whether saved tensors can be assumed
+    # to have fixed addresses.
+    forward_is_partitioned = BoxedBool(False)
+
     return CompilerConfigExtra(
         cudagraphs=cudagraphs,
         graph_id=graph_id,
         forward_device=forward_device,
+        cudagraphs_bwd_override=cudagraphs_bwd_override,
+        forward_is_partitioned=forward_is_partitioned,
     )
 
 
@@ -2322,6 +2423,18 @@ def compile_fx_forward(
             ),
         )
 
+        # Snapshot stack traces on the output node before passes run,
+        # as later passes may strip stack_trace from individual nodes.
+        output = output_node(gm)
+        output.meta["output_stack_traces"] = [
+            (
+                arg.meta.get("stack_trace")
+                if isinstance(arg, torch.fx.node.Node)
+                else None
+            )
+            for arg in output.args[0]  # type: ignore[union-attr]
+        ]
+
         inputs_devices = get_inputs_devices(example_inputs, gm)
         gm = _recursive_joint_graph_passes(gm, input_device=next(iter(inputs_devices)))
 
@@ -2387,15 +2500,26 @@ def compile_fx_forward(
     # original strides
     _recursive_record_user_visible_output_idxs(gm)
 
-    return inner_compile(
-        gm,
-        example_inputs,
-        static_input_idxs=get_static_input_idxs(fixed),
-        cudagraphs=compiler_config_extra.cudagraphs,
-        graph_id=compiler_config_extra.graph_id,
-        is_inference=is_inference,
-        boxed_forward_device_index=compiler_config_extra.forward_device,
-    )
+    with cudagraph_annotation_context(compiler_config_extra.cudagraphs):
+        result = inner_compile(
+            gm,
+            example_inputs,
+            static_input_idxs=get_static_input_idxs(fixed),
+            cudagraphs=compiler_config_extra.cudagraphs,
+            graph_id=compiler_config_extra.graph_id,
+            is_inference=is_inference,
+            boxed_forward_device_index=compiler_config_extra.forward_device,
+        )
+
+        if (
+            not is_inference
+            and isinstance(result, CompiledFxGraph)
+            and result.partition_maps
+            and len(result.partition_maps) > 1
+        ):
+            compiler_config_extra.forward_is_partitioned.value = True
+
+        return result
 
 
 def compile_fx_backward(
@@ -2428,16 +2552,32 @@ def compile_fx_backward(
             model_outputs_node.meta["user_visible_output_idxs"] = []
 
         fixed = count_tangents(gm)
+
+        # Check if cudagraphs should be overridden for backward via annotation
+        cudagraphs = compiler_config_extra.cudagraphs
+        if compiler_config_extra.cudagraphs_bwd_override is not None:
+            cudagraphs = BoxedBool(compiler_config_extra.cudagraphs_bwd_override)
+
+        # When the forward was partitioned, saved activations from inline
+        # code between partitions are NOT at fixed addresses. Only mark
+        # primals (params/buffers) as static.
+        if compiler_config_extra.forward_is_partitioned.value:
+            static_input_idxs: Sequence[int] = get_static_bw_input_idxs(gm)
+        else:
+            static_input_idxs = list(range(fixed))
         with (
-            config.patch(get_cpp_wrapper_config())
-            if config.cpp_wrapper
-            else contextlib.nullcontext()
+            (
+                config.patch(get_cpp_wrapper_config())
+                if config.cpp_wrapper
+                else contextlib.nullcontext()
+            ),
+            cudagraph_annotation_context(cudagraphs),
         ):
             return inner_compile(
                 gm,
                 example_inputs,
-                static_input_idxs=list(range(fixed)),
-                cudagraphs=compiler_config_extra.cudagraphs,
+                static_input_idxs=static_input_idxs,
+                cudagraphs=cudagraphs,
                 is_backward=True,
                 graph_id=compiler_config_extra.graph_id,
                 boxed_forward_device_index=compiler_config_extra.forward_device,
@@ -2498,9 +2638,10 @@ def compile_fx(
     model_: GraphModule,
     example_inputs_: Sequence[InputType],
     inner_compile: Callable[..., OutputCode] = compile_fx_inner,
-    config_patches: Optional[dict[str, Any]] = None,
-    decompositions: Optional[dict[OpOverload, Callable[..., Any]]] = None,
+    config_patches: dict[str, Any] | None = None,
+    decompositions: dict[OpOverload, Callable[..., Any]] | None = None,
     ignore_shape_env: bool = False,
+    compile_region_name: str | None = None,
 ) -> CompileFxOutput:
     """
     Main entry point for compiling given FX graph.  Despite the fact that this
@@ -2513,6 +2654,13 @@ def compile_fx(
     NB: This function TAKES OWNERSHIP of the input ``model_`` and can potentially
     mutate it!  Make a copy if you need to preserve the original GraphModule.
     """
+    if decompositions is not None:
+
+        def get_decomp_fn() -> dict[Any, Callable[..., Any]]:
+            return decompositions  # pyrefly: ignore[bad-return]
+    else:
+        get_decomp_fn = select_decomp_table
+
     # Some arguments trigger a recursive call to compile_fx.  Handle these
     # short circuits first, before anything else
 
@@ -2531,8 +2679,15 @@ def compile_fx(
                 inner_compile=config.patch(config_patches)(inner_compile),
                 decompositions=decompositions,
                 ignore_shape_env=ignore_shape_env,
+                compile_region_name=compile_region_name,
             )
 
+    # Keep region names out of graph_kwargs so they don't perturb FX cache keys.
+    inner_compile = functools.partial(
+        inner_compile,
+        compile_region_name=compile_region_name,
+    )
+
     # Wake up the AsyncCompile subproc pool as early as possible (if there's cuda).
     if any(
         isinstance(e, torch.Tensor) and e.device.type in ("cuda", "xpu")
@@ -2571,16 +2726,18 @@ def compile_fx(
                         cpp_wrapper=cpp_wrapper_config,
                         fx_wrapper=fx_wrapper_config,
                     ),
-                    decompositions=decompositions,
                     ignore_shape_env=ignore_shape_env,
+                    get_decomp_fn=get_decomp_fn,
+                    compile_region_name=compile_region_name,
                 )
 
     return _maybe_wrap_and_compile_fx_main(
         model_,
         example_inputs_,
         inner_compile,
-        decompositions,
         ignore_shape_env,
+        get_decomp_fn=get_decomp_fn,
+        compile_region_name=compile_region_name,
     )
 
 
@@ -2619,8 +2776,10 @@ def _maybe_wrap_and_compile_fx_main(
     model_: GraphModule,
     example_inputs_: Sequence[InputType],
     inner_compile: Callable[..., OutputCode],
-    decompositions: Optional[dict[OpOverload, Callable[..., Any]]],
     ignore_shape_env: bool,
+    *,
+    get_decomp_fn: Callable[..., dict[Any, Callable[..., Any]]] = select_decomp_table,
+    compile_region_name: str | None = None,
 ) -> CompileFxOutput:
     """
     Part of compile_fx, called after patching configs.
@@ -2634,8 +2793,9 @@ def _maybe_wrap_and_compile_fx_main(
     compile_gm = functools.partial(
         _maybe_wrap_and_compile_fx_main,
         inner_compile=inner_compile,
-        decompositions=decompositions,
         ignore_shape_env=ignore_shape_env,
+        get_decomp_fn=get_decomp_fn,
+        compile_region_name=compile_region_name,
     )
     if not graph_returns_tuple(model_):
         return make_graph_return_tuple(model_, example_inputs_, compile_gm)
@@ -2656,8 +2816,9 @@ def _maybe_wrap_and_compile_fx_main(
         model_,
         example_inputs_,
         inner_compile,
-        decompositions,
         ignore_shape_env,
+        get_decomp_fn=get_decomp_fn,
+        compile_region_name=compile_region_name,
     )
 
 
@@ -2665,8 +2826,10 @@ def _compile_fx_main(
     model_: GraphModule,
     example_inputs_: Sequence[InputType],
     inner_compile: Callable[..., OutputCode],
-    decompositions: Optional[dict[OpOverload, Callable[..., Any]]],
     ignore_shape_env: bool,
+    *,
+    get_decomp_fn: Callable[..., dict[Any, Callable[..., Any]]] = select_decomp_table,
+    compile_region_name: str | None = None,
 ) -> CompileFxOutput:
     """
     Main part of compile_fx, called after wrapping is done.
@@ -2688,23 +2851,17 @@ def _compile_fx_main(
         ),
         torch._inductor.debug.reset_provenance_globals(),
     ):
-        # Pre-grad passes cannot be run if we weren't given a GraphModule.
-        # Dynamo will always produce a GraphModule, but this handles cases
-        # where a user directly passes a plain Module with the intention of
-        # having AOTAutograd trace it.
-        # TODO: Get rid of this?
-        if isinstance(model_, GraphModule):
-            model_ = run_pre_grad_passes(model_, example_inputs_)
+        # Note: Pre-grad passes are now run inside aot_module_simplified (via the
+        # pre_grad_passes callback) after the cache lookup.
 
         assert not config._raise_error_for_testing
 
         num_example_inputs = len(example_inputs_)
 
-        compiler_config_extra = create_compiler_config_extra(config)
+        compiler_config_extra = create_compiler_config_extra(model_)
 
-        decompositions = (
-            decompositions if decompositions is not None else select_decomp_table()
-        )
+        decompositions = get_decomp_fn()
+        inner_compile = functools.partial(inner_compile, get_decomp_fn=get_decomp_fn)
 
         def fw_compiler_base(
             gm: GraphModule,
@@ -2776,6 +2933,12 @@ def bw_compiler(
 
             is_valid_aoti_model_name()
 
+            # Pre-grad passes for the aot_autograd path are run inside
+            # aot_module_simplified after the cache lookup.  The
+            # aot_export_module path doesn't use that cache so run them here.
+            if isinstance(model_, GraphModule):
+                model_ = run_pre_grad_passes(model_, example_inputs_)
+
             with functorch_config.patch(
                 unlift_effect_tokens=True,
                 selective_decompose=config.selective_decompose,
@@ -2847,16 +3010,17 @@ def bw_compiler(
             ),
         ):
             try:
-                return aot_autograd(
+                return dynamo_common.aot_autograd(
                     fw_compiler=fw_compiler,
                     bw_compiler=bw_compiler,
                     inference_compiler=inference_compiler,
                     decompositions=decompositions,
                     partition_fn=partition_fn,
                     keep_inference_input_mutations=True,
-                    cudagraphs=compiler_config_extra.cudagraphs,
-                    boxed_forward_device_index=compiler_config_extra.forward_device,
+                    compiler_config_extra=compiler_config_extra,
                     ignore_shape_env=ignore_shape_env,
+                    pre_grad_passes=run_pre_grad_passes,
+                    compile_region_name=compile_region_name,
                 )(model_, example_inputs_)
             except ShortenTraceback as e:
                 # We will also shorten the traceback inside dynamo.
@@ -2931,7 +3095,7 @@ def wrapper(*args: Any) -> Any:
 
 
 def _check_triton_bf16_support(graph: GraphLowering) -> None:
-    def warn_and_skip(device: Optional[torch.device]) -> Never:
+    def warn_and_skip(device: torch.device | None) -> Never:
         from torch._dynamo.exc import SkipFrame
 
         assert device is not None
@@ -2963,10 +3127,10 @@ def warn_and_skip(device: Optional[torch.device]) -> Never:
 
 def _aoti_flatten_inputs(
     gm: torch.fx.GraphModule,
-    args: Union[list[Any], tuple[Any, ...]],
-    kwargs: Optional[dict[str, Any]] = None,
+    args: list[Any] | tuple[Any, ...],
+    kwargs: dict[str, Any] | None = None,
     *,
-    options: Optional[dict[str, Any]] = None,
+    options: dict[str, Any] | None = None,
 ) -> tuple[list[Any], dict[str, Any]]:
     """
     Flatten the inputs to the graph module and return the flat inputs and options.
@@ -3045,3 +3209,73 @@ def _aoti_flatten_inputs(
         }
     )
     return flat_example_inputs, options
+
+
+def autograd_cache_key(
+    graph,
+    example_inputs,
+    ignore_shape_env: bool,
+    decompositions=None,
+):
+    if config.cpp_wrapper or config.fx_wrapper:
+        raise RuntimeError(
+            "autograd_cache_key is not supported with cpp_wrapper or fx_wrapper"
+        )
+
+    decompositions = (
+        decompositions if decompositions is not None else select_decomp_table()
+    )
+    # compile_fx applies these graph transforms before reaching _compile_fx_main.
+    # Neither occurs on the torch.compile/Dynamo path (which always produces
+    # tuple-returning, pre-flattened graphs). Not supported by this API.
+    if isinstance(graph, GraphModule) and not graph_returns_tuple(graph):
+        raise NotImplementedError(
+            "autograd_cache_key does not support graphs that don't return a tuple"
+        )
+    if any(isinstance(x, (list, tuple, dict)) for x in example_inputs):
+        raise NotImplementedError(
+            "autograd_cache_key does not support nested container inputs"
+        )
+
+    compiler_config_extra = create_compiler_config_extra(graph)
+
+    # These context managers replicate the ones that _compile_fx_main sets up
+    # before calling aot_autograd, so that the config snapshot captured by
+    # autograd_cache_key is identical to a real compile_fx run:
+    #   _compile_fx_main outer with-block: _use_lazy_graph_module,
+    #       enable_python_dispatcher, preserve_node_meta,
+    #       reset_provenance_globals
+    #   _compile_fx_main aot_autograd with-block: V.set_fake_mode,
+    #       torch._guards.tracing, compiled_autograd._disable,
+    #       functorch_config.patch
+
+    fake_mode = detect_fake_mode(example_inputs) or torch._subclasses.FakeTensorMode(
+        allow_non_fake_inputs=True
+    )
+    tracing_context = (
+        torch._guards.TracingContext.try_get()
+        or torch._guards.TracingContext(fake_mode)
+    )
+
+    with (
+        functorch_config.patch(
+            unlift_effect_tokens=True, selective_decompose=config.selective_decompose
+        ),
+        _use_lazy_graph_module(dynamo_config.use_lazy_graph_module),
+        enable_python_dispatcher(),
+        torch.fx.traceback.preserve_node_meta(
+            config.trace.provenance_tracking_level == 1
+        ),
+        torch._inductor.debug.reset_provenance_globals(),
+        V.set_fake_mode(fake_mode),
+        torch._guards.tracing(tracing_context),
+        compiled_autograd._disable(),
+    ):
+        return aot_autograd.autograd_cache_key(
+            graph,
+            example_inputs,
+            ignore_shape_env=ignore_shape_env,
+            decompositions=decompositions,
+            compiler_config_extra=compiler_config_extra,
+            keep_inference_input_mutations=True,
+        )
diff --git a/torch/_inductor/compile_fx_async.py b/torch/_inductor/compile_fx_async.py
index 95a0832349b1c..732a2c0a52659 100644
--- a/torch/_inductor/compile_fx_async.py
+++ b/torch/_inductor/compile_fx_async.py
@@ -2,11 +2,15 @@
 
 from collections import deque
 from dataclasses import dataclass
-from typing import Any, Optional, TYPE_CHECKING
+from typing import Any, TYPE_CHECKING
 from typing_extensions import final, override
 
 import torch._inductor.async_compile  # noqa: F401 required to warm up AsyncCompile pools
-from torch._inductor.output_code import CompiledFxGraphConstants, OutputCode
+from torch._inductor.output_code import (
+    CompiledFxGraph,
+    CompiledFxGraphConstants,
+    OutputCode,
+)
 
 from .compile_fx import _CompileFxKwargs, _InProcessFxCompile, FxCompile
 from .output_code import complex_memory_overlap  # noqa: F401
@@ -37,7 +41,7 @@ class _PostCompileData:
 class ProgressiveCompilationState:
     progression_futures: deque[Future[_WireProtocolPickledOutput]]
     callback: Callable[[_WireProtocolPickledOutput], OutputCode]
-    post_compile_data: Optional[_PostCompileData]
+    post_compile_data: _PostCompileData | None
 
     def check_and_get_ready_stage(self) -> int:
         """Check if any progression stage is ready and return its index, or -1 if none are ready."""
@@ -79,11 +83,11 @@ def switch_to_progression_stage(self, stage_index: int) -> tuple[OutputCode, boo
 # out-of-process compile to finish and then switching over to it.
 @final
 class _AsyncOutputCode(OutputCode):
-    _eager_fn: Optional[Callable[..., Any]]
-    _output_code: Optional[OutputCode]
-    _future: Optional[Future[_WireProtocolPickledOutput]]
+    _eager_fn: Callable[..., Any] | None
+    _output_code: OutputCode | None
+    _future: Future[_WireProtocolPickledOutput] | None
     _callback: Callable[[_WireProtocolPickledOutput], OutputCode]
-    _post_compile_data: Optional[_PostCompileData] = None
+    _post_compile_data: _PostCompileData | None = None
     _boxed_call: bool  # Copied from the forward/output_code
 
     def __init__(
@@ -199,7 +203,9 @@ def codegen_and_compile(
         inputs_to_check: Sequence[int],
         graph_kwargs: _CompileFxKwargs,
     ) -> OutputCode:
-        eager_output_code = _InProcessFxCompile().codegen_and_compile(
+        eager_compile = _InProcessFxCompile()
+        eager_compile.compile_region_name = self.compile_region_name
+        eager_output_code = eager_compile.codegen_and_compile(
             gm, example_inputs, inputs_to_check, graph_kwargs
         )
 
@@ -222,6 +228,8 @@ def codegen_and_compile(
         def callback(pickled_output: _WireProtocolPickledOutput) -> OutputCode:
             _AsyncFxCompile._stat_bg_finished += 1
             output = pickled_output.deserialize(constants)
+            if isinstance(output.graph, CompiledFxGraph):
+                output.graph.compile_region_name = self.compile_region_name
             self._compile._postprocess(output)
             return output.graph
 
@@ -232,9 +240,9 @@ def callback(pickled_output: _WireProtocolPickledOutput) -> OutputCode:
 # to a more optimized version when the expensive compile finishes.
 @final
 class _ProgressiveOutputCode(OutputCode):
-    _fast_output_code: Optional[OutputCode]
-    _optimized_output_code: Optional[OutputCode]
-    _compilation_state: Optional[ProgressiveCompilationState]
+    _fast_output_code: OutputCode | None
+    _optimized_output_code: OutputCode | None
+    _compilation_state: ProgressiveCompilationState | None
     # _boxed_call state is effectively cached (we sometimes wrap unboxed w/
     # lambdas to box them) so we can't change it mid-way. Since _boxed_call=True
     # is more common let's default to that and we'll convert if necessary.
@@ -392,6 +400,8 @@ def codegen_and_compile(
         def callback(pickled_output: _WireProtocolPickledOutput) -> OutputCode:
             _ProgressiveFxCompile._stat_bg_finished += 1
             output = pickled_output.deserialize(constants)
+            if isinstance(output.graph, CompiledFxGraph):
+                output.graph.compile_region_name = self.compile_region_name
             self._optimized_compile._postprocess(output)
             return output.graph
 
diff --git a/torch/_inductor/compile_fx_ext.py b/torch/_inductor/compile_fx_ext.py
index d7f5ea79b5fe6..1ed07fdc0672b 100644
--- a/torch/_inductor/compile_fx_ext.py
+++ b/torch/_inductor/compile_fx_ext.py
@@ -11,7 +11,7 @@
 import warnings
 from abc import abstractmethod
 from dataclasses import dataclass
-from typing import Any, Optional, TYPE_CHECKING, TypeGuard, Union
+from typing import Any, TYPE_CHECKING, TypeGuard
 from typing_extensions import final, override, Self
 
 import torch._inductor.async_compile  # noqa: F401 required to warm up AsyncCompile pools
@@ -214,12 +214,12 @@ class _WireProtocolInput:
     example_inputs: Sequence[InputType]
     inputs_to_check: Sequence[int]
     graph_kwargs: _CompileFxKwargs
-    tracing_context: Optional[torch._guards.TracingContext]
+    tracing_context: torch._guards.TracingContext | None
     config: dict[str, object]
     virtualized: _VirtualizedSerializer
-    deterministic_guard_for_testing: Optional[  # type: ignore[name-defined]  # mypy bug
-        torch.testing._internal.common_utils.DeterministicGuard
-    ]
+    deterministic_guard_for_testing: (  # type: ignore[name-defined]  # mypy bug
+        torch.testing._internal.common_utils.DeterministicGuard | None
+    )
     logger_state: _LoggerState
     lowering: _LoweringSerializer
     fake_tensor_mode: _FakeTensorModeSerializer
@@ -271,8 +271,8 @@ class _WireProtocolOutput:
     graph: OutputCode
     metrics: CachedMetricsDeltas
     logs: list[logging.LogRecord]
-    warning_replay: Optional[list[warnings.WarningMessage]]
-    shape_env: Optional[torch.fx.experimental.symbolic_shapes.ShapeEnv]
+    warning_replay: list[warnings.WarningMessage] | None
+    shape_env: torch.fx.experimental.symbolic_shapes.ShapeEnv | None
 
     def serialize(self) -> _WireProtocolPickledOutput:
         """
@@ -314,14 +314,14 @@ class _LoggerState:
     loggers: dict[str, int]
     # The actual log capturing mechanism - this should be None when we're not
     # actively capturing logs.
-    captured_logs: Optional[_CapturedLogs] = None
+    captured_logs: _CapturedLogs | None = None
 
     def __init__(self) -> None:
         # Mapping from logger name to level.
         self.loggers = {}
 
         def filter(
-            logger: Union[logging.Logger, logging.PlaceHolder],
+            logger: logging.Logger | logging.PlaceHolder,
         ) -> TypeGuard[logging.Logger]:
             if not isinstance(logger, logging.Logger):
                 # Assume that Placeholders propagate
@@ -360,9 +360,9 @@ def __enter__(self) -> _CapturedLogs:
 
     def __exit__(
         self,
-        exc_type: Optional[type[BaseException]],
-        exc_value: Optional[BaseException],
-        traceback: Optional[types.TracebackType],
+        exc_type: type[BaseException] | None,
+        exc_value: BaseException | None,
+        traceback: types.TracebackType | None,
     ) -> None:
         assert self.captured_logs is not None
         self.captured_logs.remove()
@@ -376,7 +376,7 @@ class _CapturedLogs:
 
     state: _LoggerState
     queue: queue.Queue[logging.LogRecord]
-    handlers: Optional[dict[str, logging.Handler]]
+    handlers: dict[str, logging.Handler] | None
 
     def __init__(self, state: _LoggerState) -> None:
         self.state = state
@@ -438,12 +438,16 @@ def codegen_and_compile(
             gm, example_inputs, inputs_to_check, graph_kwargs
         )
         if not serialized:
-            return _InProcessFxCompile().codegen_and_compile(
+            eager_compile = _InProcessFxCompile()
+            eager_compile.compile_region_name = self.compile_region_name
+            return eager_compile.codegen_and_compile(
                 gm, example_inputs, inputs_to_check, graph_kwargs
             )
 
         inputs, constants = serialized
         output = self._send_to_child(inputs).deserialize(constants)
+        if isinstance(output.graph, CompiledFxGraph):
+            output.graph.compile_region_name = self.compile_region_name
 
         self._postprocess(output)
         self._compile_stats[type(self)].codegen_and_compile += 1
@@ -459,7 +463,7 @@ def serialize_compile(
         example_inputs: Sequence[InputType],
         inputs_to_check: Sequence[int],
         graph_kwargs: _CompileFxKwargs,
-    ) -> Optional[tuple[_WireProtocolPickledInput, CompiledFxGraphConstantsWithGm]]:
+    ) -> tuple[_WireProtocolPickledInput, CompiledFxGraphConstantsWithGm] | None:
         """
         Prepare a _WireProtocolInput to compile. If None is returned then it
         wasn't possible to serialize and we should fallback to in-process.
@@ -487,9 +491,9 @@ def serialize_compile(
 
         # If we're running tests then grab the DeterministicGuard (don't want to
         # import this if it isn't already imported because it has side-effects)
-        deterministic_guard_for_testing: Optional[  # type: ignore[name-defined]  # mypy bug
-            torch.testing._internal.common_utils.DeterministicGuard
-        ] = None
+        deterministic_guard_for_testing: (  # type: ignore[name-defined]  # mypy bug
+            torch.testing._internal.common_utils.DeterministicGuard | None
+        ) = None
         try:
             deterministic_guard_for_testing = (
                 torch.testing._internal.common_utils.DeterministicGuard._current_state()  # type: ignore[attr-defined]  # mypy bug
@@ -541,7 +545,7 @@ def _postprocess(self, output: _WireProtocolOutput) -> None:
     def _run_in_child(
         cls,
         pickled_input: _WireProtocolPickledInput,
-        extra_env: Optional[Mapping[str, str]] = None,
+        extra_env: Mapping[str, str] | None = None,
     ) -> _WireProtocolPickledOutput:
         metrics = CachedMetricsHelper()
 
diff --git a/torch/_inductor/compile_fx_subproc.py b/torch/_inductor/compile_fx_subproc.py
index 58d5195046fd1..0dbf71cbb175c 100644
--- a/torch/_inductor/compile_fx_subproc.py
+++ b/torch/_inductor/compile_fx_subproc.py
@@ -3,7 +3,7 @@
 import atexit
 import functools
 import os
-from typing import Optional, TYPE_CHECKING
+from typing import TYPE_CHECKING
 from typing_extensions import final, override
 
 import torch._inductor.async_compile  # noqa: F401 required to warm up AsyncCompile pools
@@ -65,7 +65,7 @@ def process_pool() -> AnyPool:
     def _run_in_child_subprocess(
         cls,
         pickled_input: _WireProtocolPickledInput,
-        extra_env: Optional[Mapping[str, str]],
+        extra_env: Mapping[str, str] | None,
     ) -> _WireProtocolPickledOutput:
         # TODO: In subprocess mode we need to clear the inductor caches.
         # The problem:
diff --git a/torch/_inductor/compile_worker/subproc_pool.py b/torch/_inductor/compile_worker/subproc_pool.py
index 07c59b8cbb860..319c2d6918a9b 100644
--- a/torch/_inductor/compile_worker/subproc_pool.py
+++ b/torch/_inductor/compile_worker/subproc_pool.py
@@ -15,7 +15,7 @@
 from concurrent.futures import Future, ProcessPoolExecutor
 from concurrent.futures.process import BrokenProcessPool
 from enum import Enum, IntEnum
-from typing import Any, IO, Optional, TypeVar
+from typing import Any, IO, TypeVar
 from typing_extensions import Never, ParamSpec
 
 # _thread_safe_fork is needed because the subprocesses in the pool can read
@@ -131,7 +131,7 @@ class SubprocPool:
     def __init__(
         self,
         nprocs: int,
-        pickler: Optional[SubprocPickler] = None,
+        pickler: SubprocPickler | None = None,
         kind: SubprocKind = SubprocKind.FORK,
         quiesce: bool = False,
     ) -> None:
@@ -209,17 +209,17 @@ def __init__(
         self.running_waitcounter.__enter__()
 
         # The quiesce waitcounter indicates when the job is in a quiesced state.
-        self.quiesce_waitcounter: Optional[_WaitCounterTracker] = None
+        self.quiesce_waitcounter: _WaitCounterTracker | None = None
 
         # Firstjob is used to capture the time from when the firstjob is queued, to when the first job is done.
         self.firstjob = True
-        self.firstjob_id: Optional[int] = None
+        self.firstjob_id: int | None = None
         self.firstjob_waitcounter = _WaitCounter(
             "pytorch.wait_counter.subproc_pool.first_job"
         ).guard()
 
         if quiesce:
-            self.timer: Optional[Timer] = Timer(
+            self.timer: Timer | None = Timer(
                 config.quiesce_async_compile_time, self.quiesce
             )
         else:
@@ -369,7 +369,7 @@ def __init__(
         self.write_pipe = write_pipe
         self.write_lock = threading.Lock()
         self.nprocs = nprocs
-        self.pool: Optional[ProcessPoolExecutor] = None
+        self.pool: ProcessPoolExecutor | None = None
         self.running = True
 
     def main(self) -> None:
@@ -460,7 +460,7 @@ def do_job(pickler: SubprocPickler, data: bytes) -> bytes:
         return pickler.dumps(result)
 
 
-AnyPool = typing.Union[ProcessPoolExecutor, SubprocPool]
+AnyPool = ProcessPoolExecutor | SubprocPool
 
 
 def _warm_process_pool(pool: ProcessPoolExecutor, n: int) -> None:
diff --git a/torch/_inductor/compile_worker/timer.py b/torch/_inductor/compile_worker/timer.py
index 7c495403b3a55..35ed0b810439c 100644
--- a/torch/_inductor/compile_worker/timer.py
+++ b/torch/_inductor/compile_worker/timer.py
@@ -1,7 +1,6 @@
 from collections.abc import Callable
 from threading import Lock, Thread
 from time import monotonic, sleep
-from typing import Optional, Union
 
 
 class Timer:
@@ -11,12 +10,12 @@ class Timer:
 
     def __init__(
         self,
-        duration: Union[int, float],  # Duration in seconds
+        duration: int | float,  # Duration in seconds
         call: Callable[[], None],  # Function to call when we expire
     ) -> None:
         # We don't start the background thread until we actually get an event.
-        self.background_thread: Optional[Thread] = None
-        self.last_called: Optional[float] = None
+        self.background_thread: Thread | None = None
+        self.last_called: float | None = None
         self.duration = duration
         self.sleep_time = duration / 2
         self.call = call
diff --git a/torch/_inductor/compile_worker/tracked_process_pool.py b/torch/_inductor/compile_worker/tracked_process_pool.py
index 546a5cbc6395a..44f2457cd5d82 100644
--- a/torch/_inductor/compile_worker/tracked_process_pool.py
+++ b/torch/_inductor/compile_worker/tracked_process_pool.py
@@ -8,7 +8,7 @@
 from dataclasses import dataclass
 from multiprocessing.context import BaseContext
 from time import time
-from typing import Any, Optional, TypeVar
+from typing import Any, TypeVar
 from typing_extensions import ParamSpec
 
 # _thread_safe_fork is needed because the subprocesses in the pool can read
@@ -44,9 +44,9 @@ class _QueueStats:
 class TrackedProcessPoolExecutor(ProcessPoolExecutor):
     def __init__(
         self,
-        max_workers: Optional[int] = None,
-        mp_context: Optional[BaseContext] = None,
-        initializer: Optional[Callable[[], object]] = None,
+        max_workers: int | None = None,
+        mp_context: BaseContext | None = None,
+        initializer: Callable[[], object] | None = None,
     ) -> None:
         with _queue_stats_lock:
             _queue_stats.pool_count += 1
diff --git a/torch/_inductor/compile_worker/utils.py b/torch/_inductor/compile_worker/utils.py
index b4b5e21630c27..dc6c3aff0176f 100644
--- a/torch/_inductor/compile_worker/utils.py
+++ b/torch/_inductor/compile_worker/utils.py
@@ -2,7 +2,6 @@
 import signal
 from threading import Thread
 from time import sleep
-from typing import Optional
 
 
 _IN_TOPLEVEL_PROCESS = True
@@ -46,8 +45,8 @@ def run() -> None:
     _IN_TOPLEVEL_PROCESS = False
 
 
-_watchdog_thread: Optional[Thread] = None
-_original_parent: Optional[int] = None
+_watchdog_thread: Thread | None = None
+_original_parent: int | None = None
 
 
 def has_parent_changed() -> bool:
diff --git a/torch/_inductor/compiler_bisector.py b/torch/_inductor/compiler_bisector.py
index 011f746faebab..f3977179583eb 100644
--- a/torch/_inductor/compiler_bisector.py
+++ b/torch/_inductor/compiler_bisector.py
@@ -8,7 +8,6 @@
 import tempfile
 from collections.abc import Callable
 from dataclasses import dataclass, field
-from typing import Optional
 
 from torch._inductor.runtime.cache_dir_utils import cache_dir
 
@@ -83,7 +82,7 @@ def reset_counters() -> None:
 
 
 @functools.cache
-def get_env_val(env_str: str) -> Optional[str]:
+def get_env_val(env_str: str) -> str | None:
     return os.environ.get(env_str, None)
 
 
@@ -97,9 +96,9 @@ class BisectionResult:
     """
 
     backend: str
-    subsystem: Optional[str] = None
-    bisect_number: Optional[int] = None
-    debug_info: Optional[str] = None
+    subsystem: str | None = None
+    bisect_number: int | None = None
+    debug_info: str | None = None
 
 
 class CompilerBisector:
@@ -131,7 +130,7 @@ class CompilerBisector:
 
     bisection_enabled: bool = False
 
-    in_process_cache: Optional[str] = None
+    in_process_cache: str | None = None
 
     @classmethod
     def get_dir(cls) -> str:
@@ -194,7 +193,7 @@ def update_bisect_range(
         cls.write_lines_to_file(file_path, lines)
 
     @classmethod
-    def get_backend(cls) -> Optional[str]:
+    def get_backend(cls) -> str | None:
         """
         Returns the active backend, if any
         """
@@ -209,7 +208,7 @@ def get_backend(cls) -> Optional[str]:
         return None
 
     @classmethod
-    def get_subsystem(cls) -> Optional[str]:
+    def get_subsystem(cls) -> str | None:
         """
         Returns the active subsystem, if any
         """
@@ -230,7 +229,7 @@ def get_subsystem_object(cls, backend_name: str, subsystem_name: str) -> Subsyst
         return next(obj for obj in BACKENDS[backend_name] if obj.name == subsystem_name)
 
     @classmethod
-    def get_run_state(cls, backend_name: str, subsystem_name: str) -> Optional[str]:
+    def get_run_state(cls, backend_name: str, subsystem_name: str) -> str | None:
         """
         Returns the current stage of bisecting, if Any
         """
@@ -283,7 +282,7 @@ def update_config_change(cls, backend: str, subsystem: ConfigChange) -> None:
         cls.write_lines_to_file(file_path, lines)
 
     @classmethod
-    def get_config_change(cls, config_name: str) -> Optional[dict[str, object]]:
+    def get_config_change(cls, config_name: str) -> dict[str, object] | None:
         backend = cls.get_backend()
         subsystem = cls.get_subsystem()
 
@@ -326,7 +325,7 @@ def disable_subsystem(
         cls,
         backend: str,
         subsystem: str,
-        debug_info: Optional[Callable[[], str]] = None,
+        debug_info: Callable[[], str] | None = None,
     ) -> bool:
         if not cls.bisection_enabled:
             return False
@@ -375,7 +374,7 @@ def disable_subsystem(
     @classmethod
     def advance_subsystem(
         cls, curr_backend: str, curr_subsystem: Subsystem
-    ) -> Optional[Subsystem]:
+    ) -> Subsystem | None:
         """
         Tries to move to the next subsystem within the current system.
         """
@@ -403,7 +402,7 @@ def advance_subsystem(
             return None
 
     @classmethod
-    def advance_backend(cls, curr_backend: str) -> Optional[str]:
+    def advance_backend(cls, curr_backend: str) -> str | None:
         """
         Tries Move to the next backend.
         """
@@ -498,7 +497,7 @@ def initialize_system(cls) -> None:
     @classmethod
     def do_bisect(
         cls, fn: Callable[[], bool], cli_interface: bool = False
-    ) -> Optional[BisectionResult]:
+    ) -> BisectionResult | None:
         """
         Run fn repeatedly attempting to bisect torch.compile. fn should return True on success and False on failure.
         """
diff --git a/torch/_inductor/config.py b/torch/_inductor/config.py
index 9bf1f70fa4c4c..1ba123b66e9e5 100644
--- a/torch/_inductor/config.py
+++ b/torch/_inductor/config.py
@@ -1,7 +1,7 @@
 import os
 import sys
 from collections.abc import Callable
-from typing import Any, cast, Literal, Optional, TYPE_CHECKING, Union
+from typing import Any, cast, Literal, TYPE_CHECKING
 
 import torch
 import torch._inductor.custom_graph_pass
@@ -16,16 +16,17 @@
 
 if TYPE_CHECKING:
     from torch._inductor.choices import InductorChoices
+    from torch._inductor.cudagraph_utils import CUDAGraphPolicy
 
 inplace_padding = os.environ.get("TORCHINDUCTOR_INPLACE_PADDING", "1") == "1"
 can_inplace_pad_graph_input = False  # ease testing
 
 
-def fx_graph_remote_cache_default() -> Optional[bool]:
+def fx_graph_remote_cache_default() -> bool | None:
     return get_tristate_env("TORCHINDUCTOR_FX_GRAPH_REMOTE_CACHE")
 
 
-def vec_isa_ok_default() -> Optional[bool]:
+def vec_isa_ok_default() -> bool | None:
     if os.environ.get("TORCHINDUCTOR_VEC_ISA_OK") == "1":
         return True
     if os.environ.get("TORCHINDUCTOR_VEC_ISA_OK") == "0":
@@ -33,21 +34,25 @@ def vec_isa_ok_default() -> Optional[bool]:
     return None
 
 
-def autotune_remote_cache_default() -> Optional[bool]:
+def autotune_remote_cache_default() -> bool | None:
     return get_tristate_env("TORCHINDUCTOR_AUTOTUNE_REMOTE_CACHE")
 
 
-def bundled_autotune_remote_cache_default() -> Optional[bool]:
+def bundled_autotune_remote_cache_default() -> bool | None:
     return get_tristate_env("TORCHINDUCTOR_BUNDLED_AUTOTUNE_REMOTE_CACHE")
 
 
-def bundle_triton_into_fx_graph_cache_default() -> Optional[bool]:
+def bundle_triton_into_fx_graph_cache_default() -> bool | None:
     return get_tristate_env(
         "TORCHINDUCTOR_BUNDLE_TRITON_INTO_FX_GRAPH_CACHE",
         True if not is_fbcode() else None,
     )
 
 
+def autotune_at_compile_time_default() -> bool | None:
+    return get_tristate_env("TORCHINDUCTOR_AUTOTUNE_AT_COMPILE_TIME")
+
+
 def static_cuda_launcher_default() -> bool:
     STATIC_CUDA_LAUNCHER_VERSION = 2
 
@@ -114,10 +119,10 @@ def prologue_fusion_enabled() -> bool:
 # False: Disables the cache
 # True: Enables the cache
 # None: Not set -- Off for OSS, JustKnobs based for internal
-fx_graph_remote_cache: Optional[bool] = fx_graph_remote_cache_default()
+fx_graph_remote_cache: bool | None = fx_graph_remote_cache_default()
 
 # should we bundle triton caching into fx graph cache
-bundle_triton_into_fx_graph_cache: Optional[bool] = (
+bundle_triton_into_fx_graph_cache: bool | None = (
     bundle_triton_into_fx_graph_cache_default()
 )
 
@@ -143,7 +148,7 @@ def prologue_fusion_enabled() -> bool:
 # False: Disables the cache
 # True: Enables the cache
 # None: Not set -- Off for OSS, JustKnobs based for internal
-autotune_remote_cache: Optional[bool] = autotune_remote_cache_default()
+autotune_remote_cache: bool | None = autotune_remote_cache_default()
 
 # Enable bundled autotune cache.
 #
@@ -155,7 +160,7 @@ def prologue_fusion_enabled() -> bool:
 # False: Disables the cache
 # True: Enables the cache (requires autotune_local_cache)
 # None: Not set -- Off for OSS, JustKnobs based for internal
-bundled_autotune_remote_cache: Optional[bool] = bundled_autotune_remote_cache_default()
+bundled_autotune_remote_cache: bool | None = bundled_autotune_remote_cache_default()
 
 # See torch.compiler.config.force_disable_caches
 force_disable_caches: bool = Config(alias="torch.compiler.config.force_disable_caches")
@@ -178,7 +183,7 @@ def prologue_fusion_enabled() -> bool:
 unsafe_marked_cacheable_functions: dict[str, str] = {}
 
 # sleep in inductor for testing
-sleep_sec_TESTING_ONLY: Optional[int] = None
+sleep_sec_TESTING_ONLY: int | None = None
 
 # The default layout constraint for user-defined triton kernels.
 # See "The default layout constraint for custom operators" for options.
@@ -204,7 +209,12 @@ def prologue_fusion_enabled() -> bool:
 # Controls automatic precompiling of common include files for codecache.CppCodeCache
 # (i.e. for cpp_wrapper mode and for cpp kernels on CPU).  AOTI header precompiling is
 # controlled by a separate flag.
-cpp_cache_precompile_headers: bool = not is_fbcode()
+cpp_cache_precompile_headers: bool = (
+    os.environ.get(
+        "TORCHINDUCTOR_CPP_CACHE_PRECOMPILE_HEADERS", "0" if is_fbcode() else "1"
+    )
+    == "1"
+)
 
 online_softmax = os.environ.get("TORCHINDUCTOR_ONLINE_SOFTMAX", "1") == "1"
 
@@ -269,6 +279,9 @@ def prologue_fusion_enabled() -> bool:
 # do epilogue fusions before other fusions
 epilogue_fusion_first = False
 
+# do epilogue fusions for user defined triton kernels
+epilogue_fusion_user_defined_triton_kernel = False
+
 # enable pattern match+replace optimizations
 pattern_matcher = True
 
@@ -293,27 +306,29 @@ def prologue_fusion_enabled() -> bool:
 # Registers a custom pregrad pass. Note that the pre-grad IR is 1.
 # non-functional, 2. non-normalized, and 3. prone to change. Ideally we should
 # use post-grad passes.
-pre_grad_custom_pass: Optional[Callable[[torch.fx.graph.Graph], None]] = None
+pre_grad_custom_pass: torch._inductor.custom_graph_pass.CustomGraphPassType = None
 
 # Registers a custom pass to be run right before fusion in Inductor scheduler.
 # WARNING: Inductor scheduler IR is at prototype stage and subject to change,
 # hence custom IR passes built on top of it might break in the future.
-_pre_fusion_custom_pass: Optional[
+_pre_fusion_custom_pass: (
     Callable[
         [list["torch._inductor.scheduler.BaseSchedulerNode"]],
         list["torch._inductor.scheduler.BaseSchedulerNode"],
     ]
-] = None
+    | None
+) = None
 
 # Registers a custom pass to be run right after fusion in Inductor scheduler.
 # WARNING: Inductor scheduler IR is at prototype stage and subject to change,
 # hence custom IR passes built on top of it might break in the future.
-_post_fusion_custom_pass: Optional[
+_post_fusion_custom_pass: (
     Callable[
         [list["torch._inductor.scheduler.BaseSchedulerNode"]],
         list["torch._inductor.scheduler.BaseSchedulerNode"],
     ]
-] = None
+    | None
+) = None
 
 # Deprecated
 split_cat_fx_passes = True
@@ -364,6 +379,11 @@ def prologue_fusion_enabled() -> bool:
 # but the mul gets fused with other pointwise ops instead.
 force_fuse_int_mm_with_mul = False
 
+# Prevent unfusing addmm into mm+add for bf16/fp16 to avoid precision loss
+# from extra truncation at the mm output. Set to False to allow unfusing
+# (may improve perf at the cost of accuracy for some models).
+keep_addmm_fused_for_half_dtypes = True
+
 # DEPRECATED. This setting is ignored.
 use_mixed_mm = True
 
@@ -399,17 +419,15 @@ def prologue_fusion_enabled() -> bool:
 #     "reorder_communication_preserving_peak_memory",
 # ]
 reorder_for_compute_comm_overlap_passes: list[
-    Union[
-        str,
-        Callable[
-            [list["torch._inductor.scheduler.BaseSchedulerNode"]],
-            list["torch._inductor.scheduler.BaseSchedulerNode"],
-        ],
+    str
+    | Callable[
+        [list["torch._inductor.scheduler.BaseSchedulerNode"]],
+        list["torch._inductor.scheduler.BaseSchedulerNode"],
     ]
 ] = []
 
 # Maximum number of positions to advance a given collective, unlimited by default
-reorder_prefetch_limit: Optional[int] = None
+reorder_prefetch_limit: int | None = None
 
 # enable operator reordering for peak memory optimization
 reorder_for_peak_memory = True
@@ -425,17 +443,24 @@ def prologue_fusion_enabled() -> bool:
 
 bucket_all_gathers_fx: Literal["none", "all", "only_fsdp"] = "none"
 # By default torch._inductor.fx_passes.bucketing.bucket_size_determinator is used
-bucket_all_gathers_fx_bucket_size_determinator: Optional[Callable[[int], int]] = None
+bucket_all_gathers_fx_bucket_size_determinator: Callable[[int], int] | None = None
+
+bucket_all_gathers_bucket_mode: Literal[
+    "default", "custom_ops", "custom_ops_multidtype"
+] = "default"
 
 bucket_reduce_scatters_fx: Literal["none", "all"] = "none"
 # By default torch._inductor.fx_passes.bucketing.bucket_size_determinator is used
-bucket_reduce_scatters_fx_bucket_size_determinator: Optional[Callable[[int], int]] = (
-    None
-)
+bucket_reduce_scatters_fx_bucket_size_determinator: Callable[[int], int] | None = None
+
+bucket_reduce_scatters_bucket_mode: Literal[
+    "default", "custom_ops", "custom_ops_multidtype"
+] = "default"
+
 
 bucket_all_reduces_fx: Literal["none", "all"] = "none"
 # By default torch._inductor.fx_passes.bucketing.bucket_size_determinator is used
-bucket_all_reduces_fx_bucket_size_determinator: Optional[Callable[[int], int]] = None
+bucket_all_reduces_fx_bucket_size_determinator: Callable[[int], int] | None = None
 
 # runtime estimation function for ops
 # for built-in estimation function, pass in "default"; for user-defined estimation function, pass in the function handle
@@ -496,7 +521,7 @@ def prologue_fusion_enabled() -> bool:
 
 
 # Modifies the number of autotuning choices displayed, set to None for all
-def _autotune_num_choices_displayed_default() -> Optional[int]:
+def _autotune_num_choices_displayed_default() -> int | None:
     env_val = os.environ.get("TORCHINDUCTOR_AUTOTUNE_NUM_CHOICES_DISPLAYED")
     if env_val is None:
         return 10
@@ -505,9 +530,7 @@ def _autotune_num_choices_displayed_default() -> Optional[int]:
     return int(env_val)
 
 
-autotune_num_choices_displayed: Optional[int] = (
-    _autotune_num_choices_displayed_default()
-)
+autotune_num_choices_displayed: int | None = _autotune_num_choices_displayed_default()
 
 # Report the autotune choices and their benchmark results. Default is True.
 max_autotune_report_choices_stats = (
@@ -533,6 +556,15 @@ def _autotune_num_choices_displayed_default() -> Optional[int]:
     == "1"
 )
 
+# Pluggable CUDAGraph wrapping policy.  When set to a ``CUDAGraphPolicy``
+# instance, ``post_compile`` delegates cudagraph wrapping to the policy
+# instead of the built-in ``cudagraphify`` pipeline.  This allows custom
+# cudagraph implementations, selective inner-vs-outer wrapping for
+# regional compilation, and shared memory pool management.
+#
+# See ``torch._inductor.cudagraph_utils.CUDAGraphPolicy`` for the base class.
+cudagraph_policy: "CUDAGraphPolicy | None" = None
+
 # register ops upon which inductor should partition the graph. name format should be
 # "namespace::kernel_name" (e.g., aten::mm) for op overload packet, or
 # "namespace::kernel_name.overload" (e.g., aten::mm.default).
@@ -568,11 +600,14 @@ def _autotune_num_choices_displayed_default() -> Optional[int]:
 # as expected before turning it on for everyone.
 multi_kernel_hints: list[int] = []
 
+
 # Specify candidate backends for gemm autotune.
-# Possible choices are combinations of: ATen, Triton, CUTLASS, CK, CKTILE, CPP.
+# Possible choices are combinations of: ATen, Triton, CUTLASS, CUTEDSL, NVGEMM, CK, CKTILE, CPP.
 # ATen: default Pytorch ATen kernels.
 # Triton: Triton templates defined in torch inductor (AMD and NVidia GPUs).
 # CUTLASS: Cutlass templates and kernels (NVidia GPUs only).
+# CUTEDSL: CuteDSL templates for Blackwell GPUs (NVidia SM100-SM109 only).
+# NVGEMM: NVIDIA Universal GEMM via cutlass_api (NVidia GPUs only).
 # CK: Composable Kernel templates and kernels (AMD Instinct GPUs only).
 # CKTILE: Composable Kernel templates and kernels, new API (AMD Instinct GPUs only).
 # CPP: CPP templates and kernels for CPU.
@@ -584,14 +619,14 @@ def _autotune_num_choices_displayed_default() -> Optional[int]:
 # Configures the maximum number of NVIDIA Universal GEMM (NVGEMM) configs to profile
 # in max_autotune. By default it's 5, to keep compile time reasonable.
 # Set to None (or env var "none"/"all") to tune all configs.
-def _nvgemm_max_profiling_configs_default() -> Optional[int]:
+def _nvgemm_max_profiling_configs_default() -> int | None:
     env_val = os.environ.get("TORCHINDUCTOR_NVGEMM_MAX_PROFILING_CONFIGS", "5")
     if env_val.lower() in ("none", "all"):
         return None
     return int(env_val)
 
 
-nvgemm_max_profiling_configs: Optional[int] = _nvgemm_max_profiling_configs_default()
+nvgemm_max_profiling_configs: int | None = _nvgemm_max_profiling_configs_default()
 
 
 # As above, specify candidate backends for conv autotune.
@@ -640,6 +675,15 @@ def _nvgemm_max_profiling_configs_default() -> Optional[int]:
 
 # Use fx graph passes
 use_pre_grad_passes: bool = True
+
+# "early": pre-grad passes run before cache lookup (every compile).
+# "late": pre-grad passes run after cache lookup (only on cache miss);
+#   requires custom passes to implement uuid() for the cache key.
+# "default": resolves to "late" when possible (no custom pass, or custom pass
+#   with uuid), falls back to "early" otherwise.
+pre_grad_pass_timing: Literal["early", "late", "default"] = "default"
+
+
 use_joint_graph_passes: bool = True
 use_post_grad_passes: bool = True
 
@@ -697,10 +741,35 @@ def _nvgemm_max_profiling_configs_default() -> Optional[int]:
 
 # AutoHeuristic is a framework that allows one to collect data from autotuning, use the data to learn a heuristic, and
 # generate the learned heuristic to code which is shipped with the compiler
-# Specify a list of comma separated optimizations to collect data for
-autoheuristic_collect = os.environ.get("TORCHINDUCTOR_AUTOHEURISTIC_COLLECT", "")
-# Specify a list of comma separated optimizations to use learned heuristics for
-autoheuristic_use = os.environ.get("TORCHINDUCTOR_AUTOHEURISTIC_USE", "mixed_mm")
+
+
+def _parse_autoheuristic_collect_env():
+    collect_env = os.environ.get("TORCHINDUCTOR_AUTOHEURISTIC_COLLECT", "").split(",")
+    return collect_env
+
+
+def _parse_autoheuristic_use_env():
+    use_env = os.environ.get("TORCHINDUCTOR_AUTOHEURISTIC_USE", "mixed_mm").split(",")
+    return use_env
+
+
+class autoheuristic_collect:
+    """
+    Config for which autoheuristic optimizations should collect training data.
+    """
+
+    pad_mm = "pad_mm" in _parse_autoheuristic_collect_env()
+    mixed_mm = "mixed_mm" in _parse_autoheuristic_collect_env()
+
+
+class autoheuristic_use:
+    """
+    Config for which autoheuristic optimizations should use learned heuristics.
+    """
+
+    pad_mm = "pad_mm" in _parse_autoheuristic_use_env()
+    mixed_mm = "mixed_mm" in _parse_autoheuristic_collect_env()
+
 
 # If set to 1, will run a JIT post compile hook if one is set.
 run_jit_post_compile_hook = (
@@ -713,11 +782,23 @@ def run_autoheuristic(name: str) -> bool:
 
 
 def collect_autoheuristic(name: str) -> bool:
-    return name in torch._inductor.config.autoheuristic_collect.split(",")
+    if name == "pad_mm":
+        return autoheuristic_collect.pad_mm
+    elif name == "mixed_mm":
+        return autoheuristic_collect.mixed_mm
+    else:
+        # For test compatibility with non-standard ops (e.g. "test", "foo" used in tests)
+        return name in _parse_autoheuristic_collect_env()
 
 
 def use_autoheuristic(name: str) -> bool:
-    return name in torch._inductor.config.autoheuristic_use.split(",")
+    if name == "pad_mm":
+        return autoheuristic_use.pad_mm
+    elif name == "mixed_mm":
+        return autoheuristic_use.mixed_mm
+    else:
+        # For test compatibility with non-standard ops (e.g. "test", "foo" used in tests)
+        return name in _parse_autoheuristic_use_env()
 
 
 # If set to "DEFAULT", this will use the default log path specified in autoheuristic.py.
@@ -734,6 +815,13 @@ def use_autoheuristic(name: str) -> bool:
 
 force_layout_optimization = os.environ.get("TORCHINDUCTOR_FORCE_LAYOUT_OPT", "0") == "1"
 
+# Cache SDPA constraint results keyed by (tensor identity, stride_order) to avoid
+# creating duplicate buffers when the same tensor feeds multiple SDPA positions
+# (e.g., key=value in simplified PMA attention).
+cache_sdpa_constraint = (
+    os.environ.get("TORCHINDUCTOR_CACHE_SDPA_CONSTRAINT", "0") == "1"
+)
+
 
 # Whether to keep the output strides the same as eager after layout optimization.
 keep_output_stride = os.environ.get("TORCHINDUCTOR_KEEP_OUTPUT_STRIDE", "1") == "1"
@@ -752,13 +840,26 @@ def use_autoheuristic(name: str) -> bool:
 
 # Threshold to prevent excessive accumulation of ops in one buffer during lowering
 realize_acc_reads_threshold = 8
-realize_acc_reads_size_threshold: Optional[int] = (
+realize_acc_reads_size_threshold: int | None = (
     None  # TODO(xuanzh): harden this to make it non optional
 )
 
+# Defer early realization of cheap output nodes (0 buffer reads, small opcount)
+# to prevent cascade materialization in fullgraph compilation.
+# Shared constants/indices saved for backward get eagerly materialized because
+# they are graph outputs with multiple users, which inflates downstream read
+# counts and can trigger suboptimal Triton block size heuristics.
+delay_realize_cheap_outputs: bool = Config(
+    env_name_force="TORCHINDUCTOR_DELAY_REALIZE_CHEAP_OUTPUTS",
+    default=True,
+)
+
 # fallback to eager for random/dropout, this is slow but useful for debugging
 fallback_random = False
 
+# align random/dropout as eager mode(aten) behavior, maintaining fused possibility and faster gpu kernel
+align_random_eager = False
+
 # fallback embedding_bag_byte_unpack to eager
 fallback_embedding_bag_byte_unpack = False
 
@@ -768,8 +869,18 @@ def use_autoheuristic(name: str) -> bool:
     os.environ.get("TORCHINDUCTOR_ASSUME_UNALIGNED_FALLBACK_OUTPUT") == "1"
 )
 
-# Custom InductorChoices callable to use (can be a class or functools.partial with kwargs)
-inductor_choices_class: Optional[Callable[[], "InductorChoices"]] = None
+# Factory callable that returns a custom InductorChoices instance.
+# A callable (rather than a class) is used to defer imports and avoid circular
+# dependencies between config and the choices module. Example:
+#
+#     def _custom_choices_factory():
+#         from my_package.choices import MyInductorChoices
+#         return MyInductorChoices()
+#
+#     config.inductor_choices_class = _custom_choices_factory
+#
+# The returned instance must implement uuid() for cache key serialization.
+inductor_choices_class: Callable[[], "InductorChoices"] | None = None
 
 # fuse even in cases without common reads
 aggressive_fusion = False
@@ -819,12 +930,16 @@ def use_autoheuristic(name: str) -> bool:
 # how many nodes to allow into a single fusion
 max_fusion_size = 64
 
+# Minimum overlap ratio to consider fusion beneficial when inputs are shared by no indices overlapped.
+# Valid range: [0, 1]. Default to not fusion.
+min_overlap_ratio = 1.1
+
 # how many nodes to attempt pairwise fusion with in a buffer group
 max_fusion_buffer_group_pairwise_attempts = 64
 
 # maximum number of unique input/output buffers allowed in fused kernels.
 # The check is disabled if set to None.
-max_fusion_unique_io_buffers: Optional[int] = None
+max_fusion_unique_io_buffers: int | None = None
 
 # max number of inputs to generate cat as a pointwise op with masked loads
 max_pointwise_cat_inputs = 8
@@ -871,6 +986,9 @@ def use_autoheuristic(name: str) -> bool:
 # assert that indirect indexing does not read / write out of bounds
 assert_indirect_indexing = True
 
+# skip emitting runtime assertions for unbacked symbols in generated code
+do_not_emit_runtime_assertions = False
+
 # compute CSE bounds on variables that do not appear in the FX graph
 compute_all_bounds = False
 
@@ -917,8 +1035,8 @@ def use_autoheuristic(name: str) -> bool:
 )
 
 # options in caffe2/torch/_inductor/fx_passes/pre_grad.py
-add_pre_grad_passes: Optional[str] = None
-remove_pre_grad_passes: Optional[str] = None
+add_pre_grad_passes: str | None = None
+remove_pre_grad_passes: str | None = None
 
 # Comma-separated list of pass names to disable. Passes disabled via this config
 # will be skipped when they go through GraphTransformObserver.
@@ -981,7 +1099,7 @@ def decide_worker_start_method() -> str:
 # overlapping. At this moment, this pass performs better than
 # reorder_for_compute_comm_overlap_passes but we will add the logic of
 # "schedule_comm_wait" in the future and remove the one here.
-_fuse_ddp_communication_passes: list[Union[Callable[..., None], str]] = [
+_fuse_ddp_communication_passes: list[Callable[..., None] | str] = [
     "fuse_ddp_with_concat_op",
     "schedule_comm_wait",
 ]
@@ -1017,54 +1135,89 @@ class aten_distributed_optimizations:
     enable_overlap_scheduling: bool = False
 
     # Enable overlap-preserving collective bucketing
-    collective_bucketing: Optional[bool] = None
+    collective_bucketing: bool | None = None
 
     # Insert ordering dependencies to preserve overlap relationships. This should only be used if
     # compiling with inductor, or for subsequent passes before removing the ops prior to execution
-    insert_overlap_deps: Optional[bool] = None
+    insert_overlap_deps: bool | None = None
 
     # Maximum compute node prefetch distance for overlap scheduling
-    max_compute_pre_fetch: Optional[int] = None
+    max_compute_pre_fetch: int | None = None
 
-    compute_overlap_multipler: Optional[float] = None
+    compute_overlap_multipler: float | None = None
 
     # Custom runtime estimation function for ops
     # For user-defined estimation function, pass in the function handle
     # None means use default estimations
     # TODO - need estimated and profile based version
-    custom_runtime_estimation: Optional[Callable[[torch.fx.Node], Optional[float]]] = (
-        None
-    )
+    custom_runtime_estimation: Callable[[torch.fx.Node], float | None] | None = None
 
     # Method for estimating collective runtime
     # "analytical": Use bandwidth formulas (default)
     # "benchmark": Use CUDA events with power-of-2 rounding and interpolation
+    # In deterministic mode, this setting is ignored and "analytical" is used.
     collective_estimator: Literal["analytical", "benchmark"] = "analytical"
 
+    # Method for estimating compute (ATen op) runtime
+    # "analytical": Use roofline model estimates (deterministic, no GPU sync)
+    # "benchmark": Use GPU benchmarking (more accurate, requires GPU sync)
+    # In deterministic mode, this setting is ignored and "analytical" is used.
+    compute_estimator: Literal["analytical", "benchmark"] = "benchmark"
+
     # Maximum memory increase above baseline for prefetch operations
     # Uses minimum of absolute cap and ratio of baseline
-    max_memory_increase_gb: Optional[float] = None  # Absolute cap in GB
-    max_memory_increase_ratio: Optional[float] = None  # Ratio of baseline peak memory
+    max_memory_increase_gb: float | None = None  # Absolute cap in GB
+    max_memory_increase_ratio: float | None = None  # Ratio of baseline peak memory
 
     # Maximum GB of concurrent collective data in flight. Too much in flight memory
     # can cause memory fragmentation within the CUDA Caching Allocator.
-    max_in_flight_gb: Optional[float] = None
+    max_in_flight_gb: float | None = None
 
     # Maximum prefetch or bucketing candidates. Mainly intended for compile time.
-    max_coll_distance: Optional[int] = None
+    max_coll_distance: int | None = None
     log_final_collectives_estimations: bool = False
 
-    # Bucket exposed collectives first
-    bucket_exposed_first: bool = True
+    # Bucket exposed collectives first (None means auto)
+    bucket_exposed_first: bool | None = None
+
+    # Experimental setting to bucket only internode communications
+    bucket_only_internode_comms: bool = False
 
     # Enable fusion region detection for overlap scheduling cost estimation.
     # When enabled, groups of fusible ops (pointwise, reduction, etc.) are treated
     # as atomic units with memory-bound runtime estimates.
-    enable_fusion_regions: Optional[bool] = None
+    enable_fusion_regions: bool | None = None
+
+    # Default bucketing mode for auto and manual overlap scheduling
+    # "default": traced bucketing, fully lowered by inductor during compilation
+    # "custom_ops": temporary bucketing using custom ops to hide parts from inductor
+    # "custom_ops_multidtype": same as custom_ops but buckets multiple dtypes
+    #     (e.g. bf16 and fp32) into one bucket
+    # "coalesced": zero-copy batching via reduce_scatter_tensor_coalesced
+    #     (reduce_scatter only; all_gather falls back to default)
+    # None means "auto" — the compiler picks the best mode
+    bucket_mode: (
+        Literal["default", "custom_ops", "custom_ops_multidtype", "coalesced"] | None
+    ) = None
 
     # Prioritize bucketing during overlap scheduling by grouping candidates by bucket key
     prioritize_bucketing_during_scheduling: bool = True
 
+    # Verify FX graphs are identical across ranks before overlap scheduling.
+    # Detects non-SPMD graphs that would cause NCCL collective ordering
+    # mismatches and hangs.
+    spmd_check: bool = True
+
+    # Action on SPMD graph mismatch: "warn" logs a warning, "error" raises
+    # RuntimeError. "error" fails fast instead of risking silent NCCL hang.
+    # TODO(ivankobzarev): change default to "error" after real-world testing.
+    spmd_mismatch: Literal["warn", "error"] = "warn"
+
+    # When True, automatically remove extra deps that create cycles instead of
+    # raising an error.  Set this to True as a workaround if overlap scheduling
+    # fails with a cycle error, and file a bug so the root cause can be fixed.
+    overlap_scheduling_autofix_cycles: bool = False
+
 
 def parallel_compile_enabled_internally() -> bool:
     """
@@ -1104,11 +1257,7 @@ def decide_compile_threads() -> int:
         compile_threads = 1
         log.info("compile_threads set to 1 in fbcode")
     else:
-        cpu_count = (
-            len(os.sched_getaffinity(0))
-            if hasattr(os, "sched_getaffinity")
-            else os.cpu_count()
-        )
+        cpu_count = torch._utils.cpu_count()
         assert cpu_count
         compile_threads = min(32, cpu_count)
         log.info("compile_threads set to %d", compile_threads)
@@ -1117,7 +1266,7 @@ def decide_compile_threads() -> int:
 
 
 # TODO: Set directly after internal rollout.
-compile_threads: Optional[int] = None if is_fbcode() else decide_compile_threads()
+compile_threads: int | None = None if is_fbcode() else decide_compile_threads()
 
 # Whether to quiesce the Triton-compile subprocess pool at the end of each compilation.
 quiesce_async_compile_pool: bool = Config(
@@ -1159,7 +1308,7 @@ def decide_compile_threads() -> int:
 )
 
 # gemm autotuning global cache dir
-global_cache_dir: Optional[str]
+global_cache_dir: str | None
 if is_fbcode():
     try:
         from libfb.py import parutil
@@ -1260,7 +1409,7 @@ def decide_compile_threads() -> int:
 profile_bandwidth_regex = "" if _profile_var == "1" else _profile_var
 # Specify a file where we print out the profiling results.
 # None means we do not dump results to a file.
-profile_bandwidth_output: Optional[str] = os.environ.get(
+profile_bandwidth_output: str | None = os.environ.get(
     "TORCHINDUCTOR_PROFILE_OUTPUT", None
 )
 # Switch to do_bench_using_profiling to exclude the CPU overheads
@@ -1336,7 +1485,7 @@ def decide_compile_threads() -> int:
 enable_autograd_for_aot: bool = False
 
 
-def get_worker_log_path() -> Optional[str]:
+def get_worker_log_path() -> str | None:
     log_loc = None
     if is_fbcode():
         mast_job_name = os.environ.get("MAST_HPC_JOB_NAME", None)
@@ -1392,7 +1541,7 @@ class cpp:
     # performance degradation.
     dynamic_threads = os.environ.get("TORCHINDUCTOR_CPP_DYNAMIC_THREADS", "0") == "1"
 
-    simdlen: Optional[int] = None
+    simdlen: int | None = None
     min_chunk_size = int(os.environ.get("TORCHINDUCTOR_CPP_MIN_CHUNK_SIZE", "512"))
 
     cxx: tuple[None, str] = (
@@ -1411,12 +1560,12 @@ class cpp:
     # Inject a bug into our relu implementation; useful for testing our repro
     # extraction and minification functionality.
     # Valid values: "compile_error", "runtime_error", "accuracy"
-    inject_relu_bug_TESTING_ONLY: Optional[str] = None
-    inject_log1p_bug_TESTING_ONLY: Optional[str] = None
+    inject_relu_bug_TESTING_ONLY: str | None = None
+    inject_log1p_bug_TESTING_ONLY: str | None = None
 
     # If None, autodetect whether or not AVX512/AVX2 can be used.  Otherwise,
     # force usage as specified, without testing. Default None.
-    vec_isa_ok: Optional[bool] = get_tristate_env("TORCHINDUCTOR_VEC_ISA_OK")
+    vec_isa_ok: bool | None = get_tristate_env("TORCHINDUCTOR_VEC_ISA_OK")
 
     # similar to config.triton.descriptive_names
     descriptive_names: Literal["torch", "original_aten", "inductor_node"] = (
@@ -1501,6 +1650,9 @@ class cpp:
         os.environ.get("TORCHINDUCTOR_CPP_USE_CONSTEXPR_FOR_INT_ARRAY", "1") == "1"
     )
 
+    # threshold between two step reduction algorithm and welford reduction algorithm
+    use_two_step_variance_threshold = 1024
+
 
 class triton:
     """
@@ -1519,7 +1671,14 @@ class triton:
 
     # Specify dynamic shapes to capture cudagraphs and skip cudagraph for other shapes.
     # Default to None, which means we capture cudagraphs for all shapes.
-    cudagraph_capture_sizes: Optional[tuple[Union[int, tuple[int, ...]]]] = None
+    cudagraph_capture_sizes: tuple[int | tuple[int, ...]] | None = None
+
+    # Minimum number of nodes (kernels) required for a cudagraph partition.
+    # If a partition has fewer nodes than this threshold, it won't be cudagraphed.
+    # This helps avoid overhead for very small partitions where cudagraph
+    # recording/replay cost outweighs the benefits.
+    # Set to 0 to disable this check.
+    cudagraph_min_partition_size = 0
 
     # assertions not on the fast path, steady state
     slow_path_cudagraph_asserts = True
@@ -1542,7 +1701,7 @@ class triton:
 
     # Warn loudly when the number of cudagraphs due to dynamic shape
     # exceeds this limit
-    cudagraph_dynamic_shape_warn_limit: Optional[int] = 8
+    cudagraph_dynamic_shape_warn_limit: int | None = 8
 
     # synchronize after cudagraph invocation
     force_cudagraph_sync = False
@@ -1595,7 +1754,7 @@ class triton:
 
     #  We use a max of 3 if coalesce_tiling_analysis is True, and 2 otherwise.
     #  Note - coalesce_tiling_analysis does not yet apply to dynamic shapes.
-    max_tiles: Optional[int] = None
+    max_tiles: int | None = None
 
     # Prefer higher dimensional tilings. This simplifies indexing expressions, making
     # it easier to identify block pointers.
@@ -1610,7 +1769,7 @@ class triton:
 
     # Tune the generated Triton kernels at compile time instead of first time they run
     # Setting to None means uninitialized
-    autotune_at_compile_time: Optional[bool] = None
+    autotune_at_compile_time: bool | None = autotune_at_compile_time_default()
 
     # We use random tensors for autotune by default. Setting this as true will let us
     # use inputs from sample inputs to autotune user defined triton kernels.
@@ -1673,6 +1832,15 @@ class triton:
         os.environ.get("TORCHINDUCTOR_PERSISTENT_REDUCTIONS", "1") == "1"
     )
 
+    # Decompose sort-based ops (sort, mode, median) to generate Triton
+    # kernels instead of falling back to ATen eager.  When enabled, sort
+    # removes the default 512-element dimension limit and uses int32
+    # indices (up to 2^31-1 elements), and mode/median decompose into
+    # sort + reduction / pointwise ops that Inductor can lower to Triton.
+    decompose_sort_ops: bool = (
+        os.environ.get("TORCHINDUCTOR_DECOMPOSE_SORT_OPS", "0") == "1"
+    )
+
     # For small output size reductions uses cross thread-block synchronization to gain more parallelism
     cooperative_reductions = (
         os.environ.get("TORCHINDUCTOR_COOPERATIVE_REDUCTIONS", "0") == "1"
@@ -1737,13 +1905,14 @@ class triton:
     # Inject a bug into our relu implementation; useful for testing our repro
     # extraction and minification functionality.
     # Valid values: "compile_error", "runtime_error", "accuracy"
-    inject_relu_bug_TESTING_ONLY: Optional[str] = None
+    inject_relu_bug_TESTING_ONLY: str | None = None
 
     # Whether to upcast float16 / bfloat16 to float32 in triton codegen (Experimental)
     codegen_upcast_to_fp32 = True
 
-    # Whether persistent matmul kernels should be enabled this flag only has effect when on h100
-    # with a version of triton new enough to support TMA
+    # Whether persistent matmul kernels should be enabled. On NVIDIA H100+ with TMA support,
+    # this enables TMA persistent kernels. On AMD GPUs without TMA, this enables
+    # non-TMA persistent kernels as a fallback.
     enable_persistent_tma_matmul = (
         os.environ.get("ENABLE_PERSISTENT_TMA_MATMUL", "0") == "1"
     )
@@ -1784,7 +1953,7 @@ class triton:
     )
     mix_order_reduction_initial_xblock = 1
 
-    mix_order_reduction_split_size: Optional[int] = None
+    mix_order_reduction_split_size: int | None = None
     mix_order_reduction_autotune_split_size = (
         os.environ.get("TORCHINDUCTOR_MIX_ORDER_REDUCTION_AUTOTUNE_SPLIT_SIZE", "0")
         == "1"
@@ -1793,8 +1962,9 @@ class triton:
     # this could be helpful to avoid recompilations in some cases
     mix_order_reduction_non_strict_mode = False
 
-    enable_tlx_templates: bool = (
-        os.environ.get("TORCHINDUCTOR_ENABLE_TLX_TEMPLATES", "0") == "1"
+    # Don't allow multi-stages by default to avoid out of shared memory
+    mix_order_reduction_allow_multi_stages = (
+        os.environ.get("TORCHINDUCTOR_MIX_ORDER_REDUCTION_ALLOW_MULTI_STAGES") == "1"
     )
 
     # Map for storing the amount of kernel runs with dumped input tensors
@@ -1810,7 +1980,7 @@ class triton:
         os.environ.get("TORCHINDUCTOR_TRITON_PROTON_PROFILING", "0") == "1"
     )
     # If not specified, proton traces will be saved to the debug directory
-    proton_output_dir: Optional[str] = os.environ.get(
+    proton_output_dir: str | None = os.environ.get(
         "TORCHINDUCTOR_TRITON_PROTON_OUTPUT_DIR"
     )
     # Group CTAs by SM in proton trace files.
@@ -1827,6 +1997,8 @@ class triton:
         os.environ.get("TORCHINDUCTOR_TRITON_PROTON_PER_CTA_OCCUPANCY", "1") == "1"
     )
 
+    dynamic_disable_pipelining = True
+
 
 class aot_inductor:
     """
@@ -1844,6 +2016,11 @@ class aot_inductor:
     debug_compile = os.environ.get("AOT_INDUCTOR_DEBUG_COMPILE", "0") == "1"
     debug_symbols = os.environ.get("AOT_INDUCTOR_DEBUG_SYMBOLS", "0") == "1"
 
+    # Enable frame pointers for profiling tools (e.g. strobelight)
+    enable_frame_pointer = (
+        os.environ.get("AOT_INDUCTOR_ENABLE_FRAME_POINTER", "0") == "1"
+    )
+
     # Annotate generated main wrapper function, i.e. AOTInductorModel::run_impl,
     # to use which cpp compiler optimization level, default to O1
     compile_wrapper_opt_level = os.environ.get(
@@ -1885,7 +2062,7 @@ class aot_inductor:
     use_consts_asm_build = True
 
     package: bool = False
-    package_cpp_only: Optional[bool] = None
+    package_cpp_only: bool | None = None
 
     # If package_cpp_only is True, whether cpp files will be compiled to a
     # dynamically linked library or static linked library
@@ -1955,20 +2132,20 @@ class aot_inductor:
     # "binary_blob":
     #       Stores all weights in a single binary blob in data/aot_inductor/model folder for each model.
     #       This option and config.aot_inductor.force_mmap_weights cannot both be True
-    package_constants_on_disk_format: Optional[str] = None
+    package_constants_on_disk_format: str | None = None
 
     # Experimental.  Controls automatic precompiling of common AOTI include files.
     precompile_headers: bool = not is_fbcode()
 
     # Embed generated kernel binary files into model.so
-    embed_kernel_binary: Optional[bool] = None
+    embed_kernel_binary: bool | None = None
 
     # Generate kernel files that support multiple archs
     # For CUDA, this means generating fatbin files for kernels, and the fatbin files
     # contains PTX and SASS for the current architecture.
     # For XPU, this means generating SPIR-V files for kernels, and the SPIR-V files
     # will be compiled to target different XPU architectures at runtime.
-    emit_multi_arch_kernel: Optional[bool] = None
+    emit_multi_arch_kernel: bool | None = None
 
     # If not None, the generated files with use this name in file stem.
     # If None, we will use a hash to name files.
@@ -1979,12 +2156,12 @@ class aot_inductor:
     # If compile_standalone, the aoti model class name is f"AOTInductorModel{name}"
     #
     # This name can only contain letters, numbers, and underscores.
-    model_name_for_generated_files: Optional[str] = None
+    model_name_for_generated_files: str | None = None
 
     # Custom ops that have implemented C shim wrappers, defined as an op to C shim declaration dict
     custom_ops_to_c_shims: dict[torch._ops.OpOverload, list[str]] = {}
     # custom op libs that have implemented C shim wrappers
-    custom_op_libs: Optional[list[str]] = None
+    custom_op_libs: list[str] | None = None
 
     # Whether to enable link-time-optimization
     enable_lto = os.environ.get("AOT_INDUCTOR_ENABLE_LTO", "0") == "1"
@@ -1998,12 +2175,12 @@ class aot_inductor:
     # to point to windows CUDA toolkit.
     # Example: WINDOWS_CUDA_HOME=cuda-windows-base/cuda_cudart/cudart/
     # The path should contain lib cuda and lib cudart
-    cross_target_platform: Optional[str] = None
+    cross_target_platform: str | None = None
 
     # If link_libtorch is False and cross_target_platform is windows,
     # a library needs to be provided to provide the shim implementations.
-    aoti_shim_library: Optional[str | list[str]] = None
-    aoti_shim_library_path: Optional[str] = None
+    aoti_shim_library: str | list[str] | None = None
+    aoti_shim_library_path: str | None = None
 
 
 # a convenient class that automatically sets a group of the configs in aot_inductor
@@ -2046,7 +2223,7 @@ class cutlass:
     # Configures the maximum number of CUTLASS configs to profile in max_autotune.
     # By default it's None, so that all CUTLASS configs are tuned.
     # This is mainly used to reduce test time in CI.
-    cutlass_max_profiling_configs: Optional[int] = None
+    cutlass_max_profiling_configs: int | None = None
 
     # The L2 swizzle values to consider when profiling CUTLASS configs in max_autotune.
     cutlass_max_profiling_swizzle_options: list[int] = [1, 2, 4, 8]
@@ -2090,7 +2267,7 @@ class cutlass:
 
     # Keep only Cutlass op configs which contain this regular expression pattern
     # Set this to "warpspecialized_cooperative_epi_tma" to enable only SM90 TMA Cutlass Kernels for large GEMMs
-    cutlass_op_allowlist_regex: Optional[str] = os.environ.get(
+    cutlass_op_allowlist_regex: str | None = os.environ.get(
         "TORCHINDUCTOR_CUTLASS_ALLOWLIST"
     )
 
@@ -2103,7 +2280,7 @@ class cutlass:
     # Set this to "pingpong" to avoid numerical issues
     # caused by the op ordering of the "pingpong" memory access
     # pattern used by some Cutlass Kernels.
-    cutlass_op_denylist_regex: Optional[str] = os.environ.get(
+    cutlass_op_denylist_regex: str | None = os.environ.get(
         "TORCHINDUCTOR_CUTLASS_DENYLIST"
     )
 
@@ -2153,12 +2330,12 @@ class cuda(cutlass):
     # CUDA arch to use for CUDA template kernel compilation.
     # e.g. "70", "75", "80", "90", etc.
     # When arch is None, Inductor uses torch.cuda.get_device_capability(0).
-    arch: Optional[str] = None
+    arch: str | None = None
 
     # CUDA version to use for CUDA template kernel compilation.
     # e.g. "11.4", "12.1", etc.
     # When version is None, Inductor uses torch.version.cuda.
-    version: Optional[str] = None
+    version: str | None = None
 
     # Path to CUDA NVCC.
     # NVCC search order:
@@ -2166,7 +2343,7 @@ class cuda(cutlass):
     # 2) CUDACXX environment variable
     # 3) CUDA_HOME environment variable
     # 4) default system search PATH.
-    cuda_cxx: Optional[str] = None
+    cuda_cxx: str | None = None
 
     # Whether to enable device LTO (link-time-optimization).
     enable_cuda_lto = False
@@ -2176,17 +2353,22 @@ class cuda(cutlass):
 
     # Configures the maximum number of NVIDIA Universal GEMM (NVGEMM) configs to profile in max_autotune.
     # By default it's 5, to keep compile time to a reasonable level.
-    nvgemm_max_profiling_configs: Optional[int] = 5
+    nvgemm_max_profiling_configs: int | None = 5
 
 
 @inherit_fields_from(cutlass)
 class xpu(cutlass):
     # Xe arch to use for SYCL kernel compilation.
     # eg. 12, 20, which corresponding to Xe12(PVC) and Xe20 (BMG)
-    arch: Optional[str] = None
+    arch: str | None = None
     # oneAPI version to use for SYCL kernel compilation.
     # e.g. "20250201".
-    version: Optional[str] = None
+    version: str | None = None
+
+    # Path to Intel OneAPI.
+    oneapi_root: str | None = None
+
+    cutlass_dir = os.path.realpath(os.environ.get("TORCHINDUCTOR_CUTLASS_DIR", ""))
 
 
 class rocm:
@@ -2225,7 +2407,7 @@ class rocm:
 
     # Path to ROCm installation, if None, use env variable ROCM_HOME.
     # In fbcode see triton/fb/TARGETS for how ROCM_HOME gets set.
-    rocm_home: Optional[str] = None
+    rocm_home: str | None = None
 
     # Path to Composable Kernel library.
     # Install with `pip install git+https://github.com/rocm/composable_kernel@develop`.
@@ -2237,15 +2419,15 @@ class rocm:
     )
 
     # Deprecated, use CK and/or CK-tile specific settings
-    n_max_profiling_configs: Optional[int] = None
+    n_max_profiling_configs: int | None = None
 
     # Number of op instance choices to trade off between runtime perf and compilation time
     # For CK Kernels
-    ck_max_profiling_configs: Optional[int] = None
+    ck_max_profiling_configs: int | None = None
 
     # Number of op instance choices to trade off between runtime perf and compilation time
     # For CK-Tile Kernels
-    ck_tile_max_profiling_configs: Optional[int] = None
+    ck_tile_max_profiling_configs: int | None = None
 
     # Flag to use a short list of CK instances which perform well across a variety of shapes.
     # Currently RCR and F16 only
@@ -2253,7 +2435,7 @@ class rocm:
 
     # List to determine kBatch parameters to sweep over. By default, we calculate one in splitK
     # scenarios, and run on kBatch=1 in non-splitK scenarios
-    kBatch_sweep: Optional[list[int]] = None
+    kBatch_sweep: list[int] | None = None
 
     # The threshold at which we trigger a splitK config - K // max(M,N) has to be greater than this
     split_k_threshold: int = 16
@@ -2276,6 +2458,11 @@ class rocm:
 xpu_backend: Literal["triton"] = "triton"
 
 
+class mtia:
+    # Configuration to force inductor to never use welford reductions for MTIA backend
+    disable_welford_reduction = False
+
+
 class halide:
     # Base halide target to use for CPU devices
     cpu_target = "host"
@@ -2313,7 +2500,7 @@ class trace:
 
     # Save debug information to a temporary directory
     # If not specified, a temp directory will be created by system
-    debug_dir: Optional[str] = None
+    debug_dir: str | None = None
 
     # Save python logger call >=logging.DEBUG
     debug_log = False
@@ -2363,7 +2550,7 @@ class trace:
 
     # Upload the .tar.gz file
     # Needs to be overridden based on specific environment needs
-    upload_tar: Optional[Callable[[str], None]] = None
+    upload_tar: Callable[[str], None] | None = None
 
     log_autotuning_results = os.environ.get("LOG_AUTOTUNE_RESULTS", "0") == "1"
 
@@ -2394,6 +2581,9 @@ class trace:
     "post_grad_custom_post_pass",
     "_fuse_ddp_communication_passes",
     "_pre_fusion_custom_pass",
+    # CUDAGraphPolicy objects are not picklable and only affect
+    # post_compile wrapping, not compiled code itself.
+    "cudagraph_policy",
 ]
 
 _cache_config_ignore_prefix: list[str] = [
@@ -2411,10 +2601,15 @@ class trace:
     "post_grad_custom_pre_pass",
     "joint_custom_pre_pass",
     "joint_custom_post_pass",
+    "pre_grad_custom_pass",
     "_fuse_ddp_communication_passes",
     "_pre_fusion_custom_pass",
+    # CUDAGraphPolicy only affects post_compile, not compiled output
+    "cudagraph_policy",
     # tests assume that changes here don't invalidate cache
     "always_complex_memory_overlap_TESTING_ONLY",
+    # timing affects cache structure, not cache content
+    "pre_grad_pass_timing",
     # cache related options are not relevant to cache results
     "fx_graph_cache",
     "fx_graph_remote_cache",
@@ -2422,6 +2617,12 @@ class trace:
     "autotune_remote_cache",
 ]
 
+# Config keys whose values are callable factories. save_config_portable will
+# instantiate the factory and use .uuid() for serialization.
+_cache_config_factory_keys: list[str] = [
+    "inductor_choices_class",
+]
+
 # External callable for matmul tuning candidates
 external_matmul: list[Callable[[torch.Tensor, torch.Tensor, torch.Tensor], None]] = []
 
@@ -2432,7 +2633,7 @@ class trace:
 
 class lookup_table:
     # Lookup table for template config overrides
-    table: Optional[dict[str, list[dict[str, Any]]]] = None
+    table: dict[str, list[dict[str, Any]]] | None = None
 
     # Enable template src_hash checking in lookup table to prevent using stale configs.
     # If True, configs with 'template_hash' field will be compared against the template's
@@ -2444,7 +2645,17 @@ class lookup_table:
 class test_configs:
     force_extern_kernel_in_multi_template: bool = False
 
-    max_mm_configs: Optional[int] = None
+    # Force custom op autotuning choice selection:
+    # - None: normal autotuning (default)
+    # - True: force decomposition to win
+    # - False: force fallback to win
+    force_custom_op_decomposition: bool | None = None
+
+    # Force custom op autotuning to prevent grouping ranges by implementation.
+    # When True, each range is treated as a separate impl group, forcing torch.cond dispatch.
+    force_no_impl_grouping: bool = False
+
+    max_mm_configs: int | None = None
 
     runtime_triton_dtype_assert = False
     runtime_triton_shape_assert = False
@@ -2455,16 +2666,16 @@ class test_configs:
     # Can be set via TORCHINDUCTOR_AUTOTUNE_CHOICE_NAME_REGEX and
     # TORCHINDUCTOR_AUTOTUNE_CHOICE_DESC_REGEX environment variables
 
-    autotune_choice_name_regex: Optional[str] = os.environ.get(
+    autotune_choice_name_regex: str | None = os.environ.get(
         "TORCHINDUCTOR_AUTOTUNE_CHOICE_NAME_REGEX"
     )
-    autotune_choice_desc_regex: Optional[str] = os.environ.get(
+    autotune_choice_desc_regex: str | None = os.environ.get(
         "TORCHINDUCTOR_AUTOTUNE_CHOICE_DESC_REGEX"
     )
 
     graphsafe_rng_func_ignores_fallback_random = False
 
-    track_memory_lifecycle: Optional[Literal["assert", "log"]] = None
+    track_memory_lifecycle: Literal["assert", "log"] | None = None
 
     # If set to True, AOTI-generated CMakelists.txt will still use libtorch
     # for unit testing
@@ -2504,6 +2715,11 @@ class eager_numerics:
 
     disable_ftz: bool = False
 
+    # Use the CUDA toolkit's libdevice instead of Triton's bundled version.
+    # Triton bundles its own libdevice.10.bc which may use different polynomial
+    # coefficients than CUDA's version, causing ~1 ULP differences in pow.
+    use_pytorch_libdevice: bool = False
+
 
 # Mode to emulate PyTorch eager numerics when doing lower precision compute
 # (fp16, bf16).  PyTorch eager computes bf16/fp16 by upcasting inputs to fp32
diff --git a/torch/_inductor/config_comms.py b/torch/_inductor/config_comms.py
index 849a6b6b23cd0..deed06c54990c 100644
--- a/torch/_inductor/config_comms.py
+++ b/torch/_inductor/config_comms.py
@@ -1,6 +1,5 @@
 import os
 import sys
-from typing import Optional
 
 from torch.utils._config_module import install_config_module
 
@@ -14,12 +13,12 @@
 runtime_estimations_align_across_all_distributed_ranks: bool = False
 
 reorder_iterative_debug_memory_recompute: bool = False
-reorder_iterative_debug_limit_to_reorder: Optional[int] = (
+reorder_iterative_debug_limit_to_reorder: int | None = (
     None
     if (env_str := os.getenv("PYTORCH_REORDER_COLLECTIVES_LIMIT")) is None
     else int(env_str)
 )
-sink_waits_iterative_debug_limit_to_sink: Optional[int] = (
+sink_waits_iterative_debug_limit_to_sink: int | None = (
     # pyrefly: ignore[unbound-name]
     None if (env_str := os.getenv("PYTORCH_SINK_WAITS_LIMIT")) is None else int(env_str)
 )
diff --git a/torch/_inductor/constant_folding.py b/torch/_inductor/constant_folding.py
index f76897695240b..6e0ea4ebf46de 100644
--- a/torch/_inductor/constant_folding.py
+++ b/torch/_inductor/constant_folding.py
@@ -1,6 +1,7 @@
 import collections
+import typing_extensions
 from collections.abc import Callable
-from typing import Any, Optional
+from typing import Any
 
 import torch
 import torch.utils._pytree as pytree
@@ -33,8 +34,8 @@ def clear_dont_constant_fold() -> None:
 def replace_node_with_constant(
     gm: torch.fx.GraphModule,
     node: torch.fx.Node,
-    constant: Optional[torch.Tensor] = None,
-    name: Optional[str] = None,
+    constant: torch.Tensor | None = None,
+    name: str | None = None,
 ) -> None:
     g = gm.graph
 
@@ -73,7 +74,7 @@ def replace_node_with_constant(
 
 
 def is_const_source(
-    node: torch.fx.Node, lifted_constant_names: Optional[list[str]]
+    node: torch.fx.Node, lifted_constant_names: list[str] | None
 ) -> bool:
     return node.op == "get_attr" or node.name in (lifted_constant_names or ())
 
@@ -83,8 +84,8 @@ def __init__(
         self,
         gm: torch.fx.GraphModule,
         skip_constructors: bool = False,
-        lifted_constant_names: Optional[list[str]] = None,
-        skip_folding_node_fn: Optional[Callable[[torch.fx.Node], bool]] = None,
+        lifted_constant_names: list[str] | None = None,
+        skip_folding_node_fn: Callable[[torch.fx.Node], bool] | None = None,
     ) -> None:
         super().__init__(gm)
         self.node_replacements: dict[torch.fx.Node, Any] = {}
@@ -115,7 +116,7 @@ def _deduce_value(self, node: torch.fx.Node) -> Any:
             if (
                 isinstance(inp, torch.fx.Node)
                 and inp.name not in (self.lifted_constant_names or ())
-                and self.env[inp] != self.deferred_value
+                and self.env[inp] is not self.deferred_value
             ):
                 return self.unknown_value
         return self.deferred_value
@@ -189,6 +190,7 @@ def add_use(inp: torch.fx.Node) -> None:
 
         return last_non_output_use
 
+    @typing_extensions.override
     def run_node(self, node: torch.fx.Node) -> Any:
         if node.target == "output":
             # because we remove nodes from env on last non output use,
@@ -250,25 +252,36 @@ def set_env(arg: torch.fx.Node) -> None:
 
         out = self._deduce_value(node)
 
-        if isinstance(out, torch._C.ScriptObject):
+        if isinstance(
+            out,
+            (
+                torch._C.ScriptObject,
+                torch._library.fake_class_registry.FakeScriptObject,
+            ),
+        ):
             return out
 
-        if out == self.unknown_value:
+        if out is self.unknown_value:
             return self.unknown_value
 
         if not is_const_source(node, self.lifted_constant_names) and (
-            isinstance(out, torch.Tensor) or out == self.deferred_value
+            isinstance(out, torch.Tensor) or out is self.deferred_value
         ):
-            if out != self.deferred_value and out.device.type == "meta":
+            if (
+                out is not self.deferred_value
+                and out.device.type == "meta"  # pyrefly: ignore[missing-attribute]
+            ):
                 return out
 
-            if not self.insertable_tensor_check(out):
+            if not self.insertable_tensor_check(
+                out  # pyrefly: ignore[bad-argument-type]
+            ):
                 return out
 
             if self.is_impure(node):
                 return self.unknown_value
 
-            self.add_node_replacement(node, out)
+            self.add_node_replacement(node, out)  # pyrefly: ignore[bad-argument-type]
 
             flattened_node_inps = pytree.arg_tree_leaves(*node.args, **node.kwargs)
 
@@ -307,7 +320,7 @@ def insert_placerholder_values(self, env: dict[torch.fx.Node, Any]) -> None:
 
 def constant_fold(
     gm: torch.fx.GraphModule,
-    constraint_fn: Optional[Callable[[torch.fx.Node], bool]] = None,
+    constraint_fn: Callable[[torch.fx.Node], bool] | None = None,
 ) -> None:
     with torch.utils._python_dispatch._disable_current_modes():
         cf = ConstantFolder(gm, skip_constructors=True)
@@ -336,8 +349,8 @@ def constant_fold(
 def constant_graph_tag(
     gm: torch.fx.GraphModule,
     skip_constructors: bool = True,
-    lifted_constant_names: Optional[list[str]] = None,
-    skip_folding_node_fn: Optional[Callable[[torch.fx.Node], bool]] = None,
+    lifted_constant_names: list[str] | None = None,
+    skip_folding_node_fn: Callable[[torch.fx.Node], bool] | None = None,
 ) -> None:
     with torch.utils._python_dispatch._disable_current_modes():
         cf = ConstantFolder(
@@ -365,8 +378,8 @@ def constant_graph_tag(
 def run_and_get_constant_graph(
     gm: torch.fx.GraphModule,
     skip_constructors: bool = True,
-    lifted_constant_names: Optional[list[str]] = None,
-    skip_folding_node_fn: Optional[Callable[[torch.fx.Node], bool]] = None,
+    lifted_constant_names: list[str] | None = None,
+    skip_folding_node_fn: Callable[[torch.fx.Node], bool] | None = None,
 ) -> torch.fx.GraphModule:
     """
     Construct a GraphModule which corresponds to the part which could be
diff --git a/torch/_inductor/cpp_builder.py b/torch/_inductor/cpp_builder.py
index 6dd6e0d2b5c5f..c4a801d9ba674 100644
--- a/torch/_inductor/cpp_builder.py
+++ b/torch/_inductor/cpp_builder.py
@@ -23,7 +23,7 @@
 from ctypes import cdll, wintypes
 from ctypes.util import find_library
 from pathlib import Path
-from typing import Any, Optional, Union
+from typing import Any
 
 import torch
 from torch._dynamo.utils import dynamo_timed
@@ -437,6 +437,56 @@ def convert_cubin_to_obj(
     return obj_file
 
 
+def batch_convert_cubins_to_obj(
+    cubins: list[tuple[str, str]],
+    output_dir: str,
+    cpp_compiler: str = "gcc",
+) -> str:
+    """Convert multiple cubin files to a single .o using batched .incbin assembly.
+
+    Instead of spawning 3 subprocesses per cubin (ld + 2x objcopy), generates
+    a single .S file with .incbin directives for all cubins and compiles it
+    with one compiler invocation. Produces bit-identical rodata and symbols
+    as the per-cubin convert_cubin_to_obj approach.
+
+    Args:
+        cubins: list of (cubin_file_path, kernel_name) tuples.
+        output_dir: directory for the generated .S and .o files.
+        cpp_compiler: C compiler to use for assembling (default: gcc).
+
+    Returns:
+        Path to the combined .o file.
+    """
+    asm_path = os.path.join(output_dir, "cubins_combined.S")
+    obj_path = os.path.join(output_dir, "cubins_combined.o")
+
+    with open(asm_path, "w") as f:
+        f.write(".section .rodata\n")
+        for cubin_file, kernel_name in cubins:
+            # Use absolute path to avoid issues with working directory
+            abs_cubin = os.path.abspath(cubin_file)
+            escaped_path = abs_cubin.replace("\\", "\\\\").replace('"', '\\"')
+            f.write(
+                f".balign 16\n"
+                f".global __{kernel_name}_start\n"
+                f".global __{kernel_name}_end\n"
+                f"__{kernel_name}_start:\n"
+                f'.incbin "{escaped_path}"\n'
+                f"__{kernel_name}_end:\n"
+                f".global __{kernel_name}_size\n"
+                f".set __{kernel_name}_size, "
+                f"__{kernel_name}_end - __{kernel_name}_start\n"
+            )
+
+    subprocess.run(
+        [cpp_compiler, "-c", asm_path, "-o", obj_path],
+        capture_output=True,
+        text=True,
+        check=True,
+    )
+    return obj_path
+
+
 @functools.cache
 def _is_apple_clang(cpp_compiler: str) -> bool:
     version_string = subprocess.check_output([cpp_compiler, "--version"]).decode("utf8")
@@ -654,13 +704,13 @@ class BuildOptionsBase:
     def __init__(
         self,
         compiler: str = "",
-        definitions: Optional[list[str]] = None,
-        include_dirs: Optional[list[str]] = None,
-        cflags: Optional[list[str]] = None,
-        ldflags: Optional[list[str]] = None,
-        libraries_dirs: Optional[list[str]] = None,
-        libraries: Optional[list[str]] = None,
-        passthrough_args: Optional[list[str]] = None,
+        definitions: list[str] | None = None,
+        include_dirs: list[str] | None = None,
+        cflags: list[str] | None = None,
+        ldflags: list[str] | None = None,
+        libraries_dirs: list[str] | None = None,
+        libraries: list[str] | None = None,
+        passthrough_args: list[str] | None = None,
         aot_mode: bool = False,
         use_relative_path: bool = False,
         compile_only: bool = False,
@@ -679,7 +729,7 @@ def __init__(
 
         # Optionally, the path to a precompiled header which should be included on the
         # build command line.
-        self.precompiled_header: Optional[str] = None
+        self.precompiled_header: str | None = None
 
         self._aot_mode: bool = aot_mode
         self._use_relative_path: bool = use_relative_path
@@ -688,9 +738,10 @@ def __init__(
         self._preprocessing: bool = preprocessing
 
     def _process_compile_only_options(self) -> None:
-        if self._compile_only:
+        if self._compile_only or self._precompiling or self._preprocessing:
             self._libraries_dirs = []
             self._libraries = []
+            self._ldflags = []
 
     def _remove_duplicate_options(self) -> None:
         self._definitions = _remove_duplication_in_list(self._definitions)
@@ -770,7 +821,7 @@ def _get_warning_all_cflag(warning_all: bool = True) -> list[str]:
         return []
 
 
-def _get_cpp_std_cflag(std_num: str = "c++17") -> list[str]:
+def _get_cpp_std_cflag(std_num: str = "c++20") -> list[str]:
     if _IS_WINDOWS:
         """
         On Windows, only c++20 can support `std::enable_if_t`.
@@ -912,6 +963,12 @@ def _get_optimization_cflags(
         cflags += debug_cflags
         ldflags += debug_ldflags
 
+    if config.aot_inductor.enable_frame_pointer:
+        if _IS_WINDOWS:
+            cflags.append("Oy-")
+        else:
+            cflags.append("fno-omit-frame-pointer")
+
     cflags += _get_ffast_math_flags()
 
     if _IS_WINDOWS:
@@ -1063,7 +1120,10 @@ def __init__(
 
 
 def _get_torch_cpp_wrapper_definition() -> list[str]:
-    return ["TORCH_INDUCTOR_CPP_WRAPPER", "STANDALONE_TORCH_HEADER"]
+    defs = ["TORCH_INDUCTOR_CPP_WRAPPER", "STANDALONE_TORCH_HEADER"]
+    if config.cpp_cache_precompile_headers:
+        defs.append("TORCH_INDUCTOR_PRECOMPILE_HEADERS")
+    return defs
 
 
 def _use_custom_generated_macros() -> list[str]:
@@ -1078,6 +1138,11 @@ def _use_fb_internal_macros() -> list[str]:
                 "C10_USE_MINIMAL_GLOG",
                 "C10_DISABLE_TENSORIMPL_EXTENSIBILITY",
             ]
+            if platform.machine() == "x86_64":
+                fb_internal_macros += [
+                    "ATEN_MKL_ENABLED_FBCODE=1",
+                    "ATEN_MKLDNN_ENABLED_FBCODE=1",
+                ]
             return fb_internal_macros
         else:
             return []
@@ -1089,12 +1154,13 @@ def _setup_standard_sys_libs(
     cpp_compiler: str,
     aot_mode: bool,
     use_relative_path: bool,
-) -> tuple[list[str], list[str], list[str]]:
+) -> tuple[list[str], list[str], list[str], list[str]]:
     cflags: list[str] = []
     include_dirs: list[str] = []
     passthrough_args: list[str] = []
+    ldflags: list[str] = []
     if _IS_WINDOWS:
-        return cflags, include_dirs, passthrough_args
+        return cflags, include_dirs, passthrough_args, ldflags
 
     if config.is_fbcode():
         # TODO(T203137008) Can we unify these flags with triton_cc_command?
@@ -1119,12 +1185,12 @@ def _setup_standard_sys_libs(
 
         if _is_clang(cpp_compiler):
             passthrough_args.append(" --rtlib=compiler-rt")
-            passthrough_args.append(" -fuse-ld=lld")
-            passthrough_args.append(f" -Wl,--script={linker_script}")
             passthrough_args.append(" -B" + build_paths.glibc_lib)
-            passthrough_args.append(" -L" + build_paths.glibc_lib)
+            ldflags.append("fuse-ld=lld")
+            ldflags.append(f"Wl,--script={linker_script}")
+            ldflags.append("L" + build_paths.glibc_lib)
 
-    return cflags, include_dirs, passthrough_args
+    return cflags, include_dirs, passthrough_args, ldflags
 
 
 def _get_build_args_of_chosen_isa(vec_isa: VecISA) -> tuple[list[str], list[str]]:
@@ -1159,6 +1225,8 @@ def _get_torch_related_args(
         libraries_dirs = [TORCH_LIB_PATH]
         if sys.platform != "darwin" and not config.is_fbcode():
             libraries.extend(["torch", "torch_cpu"])
+            if _IS_WINDOWS:
+                libraries.append("c10")
             if not aot_mode:
                 libraries.append("torch_python")
     else:
@@ -1393,9 +1461,8 @@ def _get_openmp_args(
         if config.is_fbcode():
             include_dir_paths.append(build_paths.openmp_include)
 
-            openmp_lib = build_paths.openmp_lib_so
-            fb_openmp_extra_flags = f"-Wp,-fopenmp {openmp_lib}"
-            passthrough_args.append(fb_openmp_extra_flags)
+            passthrough_args.append("-Wp,-fopenmp")
+            lib_dir_paths.append(os.path.dirname(build_paths.openmp_lib_so))
 
             libs.append("omp")
         else:
@@ -1483,6 +1550,7 @@ def get_cpp_torch_options(
         sys_libs_cflags,
         sys_libs_include_dirs,
         sys_libs_passthrough_args,
+        sys_libs_ldflags,
     ) = _setup_standard_sys_libs(cpp_compiler, aot_mode, use_relative_path)
 
     isa_macros, isa_ps_args_build_flags = _get_build_args_of_chosen_isa(vec_isa)
@@ -1524,7 +1592,7 @@ def get_cpp_torch_options(
         + omp_include_dir_paths
     )
     cflags = sys_libs_cflags + omp_cflags
-    ldflags = omp_ldflags
+    ldflags = sys_libs_ldflags + omp_ldflags
     libraries_dirs = python_libraries_dirs + torch_libraries_dirs + omp_lib_dir_paths
     libraries = torch_libraries + omp_lib
     passthrough_args = (
@@ -1623,7 +1691,7 @@ def _set_gpu_runtime_env() -> None:
 
 
 @functools.lru_cache(8)
-def _find_libcudart_static(path: str) -> Optional[Path]:
+def _find_libcudart_static(path: str) -> Path | None:
     lib_dirs = list(Path(path).rglob("libcudart_static.a"))
     if lib_dirs:
         return lib_dirs[0].resolve().parent
@@ -1632,13 +1700,190 @@ def _find_libcudart_static(path: str) -> Optional[Path]:
     return None
 
 
+def _gen_mingw_import_lib(dll_path: str, def_path: str, import_lib_path: str) -> None:
+    """Generate a MinGW import library (.a) from a DLL using gendef and dlltool."""
+    dll_name = os.path.basename(dll_path)
+    with open(def_path, "w") as def_file:
+        subprocess.run(
+            ["gendef", "-", dll_path],
+            stdout=def_file,
+            stderr=subprocess.PIPE,
+            check=True,
+        )
+
+    subprocess.run(
+        [
+            "x86_64-w64-mingw32-dlltool",
+            "-d",
+            def_path,
+            "-l",
+            import_lib_path,
+            "-D",
+            dll_name,
+        ],
+        stderr=subprocess.PIPE,
+        check=True,
+    )
+    log.info("Generated MinGW import library %s from %s", import_lib_path, dll_name)
+
+
+# MSVC /GS buffer security check stubs for MinGW cross-compilation.
+# CUDA 13.0+ cudart.lib contains MSVC-compiled static objects that reference
+# these symbols (__security_cookie, __security_check_cookie, __GSHandlerCheck).
+# When cross-compiling with MinGW, we provide no-op stubs so the linker can
+# resolve them. At runtime on Windows, the CUDA runtime DLL handles its own
+# security checks internally; the static loader code that references these
+# symbols is a thin shim whose /GS instrumentation is safe to stub out.
+_MSVC_GS_STUBS_SOURCE = """\
+#include <stdint.h>
+uint64_t __security_cookie = 0x00002B992DDFA232ULL;
+void __security_check_cookie(uint64_t cookie) { (void)cookie; }
+void __GSHandlerCheck(void) {}
+"""
+
+
+def _create_msvc_gs_stubs_lib(output_dir: str) -> str | None:
+    """
+    Create a static library with MSVC GS security symbol stubs for MinGW.
+
+    Returns the library name (without lib prefix / .a suffix) if successful,
+    or None on failure.
+    """
+    stubs_lib = os.path.join(output_dir, "libmsvc_gs_stubs.a")
+    if os.path.exists(stubs_lib):
+        return "msvc_gs_stubs"
+
+    src_path = ""
+    obj_path = ""
+    try:
+        src_path = os.path.join(output_dir, "_msvc_gs_stubs.c")
+        obj_path = os.path.join(output_dir, "_msvc_gs_stubs.o")
+        with open(src_path, "w") as f:
+            f.write(_MSVC_GS_STUBS_SOURCE)
+
+        mingw_gcc = MINGW_GXX.replace("g++", "gcc")
+        subprocess.run(
+            [mingw_gcc, "-c", src_path, "-o", obj_path],
+            stderr=subprocess.PIPE,
+            check=True,
+        )
+        mingw_ar = MINGW_GXX.replace("g++", "ar")
+        subprocess.run(
+            [mingw_ar, "rcs", stubs_lib, obj_path],
+            stderr=subprocess.PIPE,
+            check=True,
+        )
+        log.info("Created MSVC GS stubs library: %s", stubs_lib)
+        return "msvc_gs_stubs"
+    except (FileNotFoundError, subprocess.CalledProcessError):
+        log.warning(
+            "Failed to create MSVC GS stubs library.",
+            exc_info=True,
+        )
+        if os.path.exists(stubs_lib):
+            os.remove(stubs_lib)
+        return None
+    finally:
+        for f in [src_path, obj_path]:
+            if f and os.path.exists(f):
+                os.remove(f)
+
+
+def _ensure_mingw_cudart_import_lib(libraries_dirs: list[str]) -> list[str]:
+    """
+    Auto-generate a MinGW-compatible import library (libcudart.a)
+    from the CUDA runtime DLL. This avoids linking against the hybrid cudart.lib
+    which contains MSVC-compiled static objects with /GS security symbols that
+    MinGW cannot resolve.
+
+    Falls back to creating MSVC GS security stubs if the DLL is unavailable,
+    and falls back gracefully to the original cudart.lib if that also fails.
+
+    Returns a list of extra library names to link (e.g. ["msvc_gs_stubs"]).
+    """
+    import glob
+
+    windows_cuda_home = os.environ.get("WINDOWS_CUDA_HOME")
+    if not windows_cuda_home:
+        log.debug(
+            "WINDOWS_CUDA_HOME not set, skipping MinGW cudart import lib generation"
+        )
+        return []
+
+    for lib_dir in libraries_dirs:
+        if os.path.exists(os.path.join(lib_dir, "libcudart.a")):
+            log.debug("libcudart.a already exists in %s, skipping generation", lib_dir)
+            return []
+
+    # Find the CUDA runtime DLL for import lib generation
+    bin_dir = os.path.join(windows_cuda_home, "bin", "x64")
+    if not os.path.isdir(bin_dir):
+        bin_dir = os.path.join(windows_cuda_home, "bin")
+    dll_candidates = glob.glob(os.path.join(bin_dir, "cudart64_*.dll"))
+
+    # Find a writable directory containing cudart.lib for output
+    output_dir = None
+    for lib_dir in libraries_dirs:
+        if os.path.isdir(lib_dir) and os.access(lib_dir, os.W_OK):
+            if os.path.exists(os.path.join(lib_dir, "cudart.lib")):
+                output_dir = lib_dir
+                break
+
+    if not dll_candidates:
+        log.warning(
+            "No cudart64_*.dll found in %s. Cannot generate MinGW import library. "
+            "Will create MSVC GS security stubs as fallback.",
+            bin_dir,
+        )
+        # Fallback: create GS stubs so the hybrid cudart.lib can link
+        if output_dir is not None:
+            stub_lib = _create_msvc_gs_stubs_lib(output_dir)
+            if stub_lib:
+                return [stub_lib]
+        return []
+
+    if output_dir is None:
+        log.warning(
+            "No writable directory containing cudart.lib found. "
+            "Cannot generate MinGW import library. "
+            "If linking fails with undefined references to __security_cookie, "
+            "ensure cudart.lib is present in one of: %s",
+            libraries_dirs,
+        )
+        return []
+
+    dll_path = dll_candidates[0]
+    dll_name = os.path.basename(dll_path)
+
+    def_path = os.path.join(output_dir, dll_name.replace(".dll", ".def"))
+    import_lib_path = os.path.join(output_dir, "libcudart.a")
+
+    try:
+        _gen_mingw_import_lib(dll_path, def_path, import_lib_path)
+        return []
+    except (FileNotFoundError, subprocess.CalledProcessError):
+        log.warning(
+            "Failed to generate MinGW cudart import library. "
+            "Falling back to MSVC GS stubs.",
+            exc_info=True,
+        )
+        for f in [def_path, import_lib_path]:
+            if os.path.exists(f):
+                os.remove(f)
+        # Fallback: create GS stubs
+        stub_lib = _create_msvc_gs_stubs_lib(output_dir)
+        if stub_lib:
+            return [stub_lib]
+        return []
+
+
 def _transform_cuda_paths(lpaths: list[str]) -> None:
     # This handles two cases:
     # 1. Cases where libs are in (e.g.) lib/cuda-12 and lib/cuda-12/stubs
     # 2. Linux machines may have CUDA installed under either lib64/ or lib/
     for i, path in enumerate(lpaths):
         if "CUDA_HOME" in os.environ and path.startswith(os.environ["CUDA_HOME"]):
-            lib_dir: Optional[Path] = _find_libcudart_static(path)
+            lib_dir: Path | None = _find_libcudart_static(path)
             if lib_dir is None:
                 continue
             lpaths[i] = str(lib_dir)
@@ -1700,7 +1945,8 @@ def get_cpp_torch_device_options(
             else:
                 libraries += ["cuda", "torch_cuda"]
             if config.aot_inductor.cross_target_platform == "windows":
-                libraries += ["cudart"]
+                extra_libs = _ensure_mingw_cudart_import_lib(libraries_dirs)
+                libraries += ["cudart"] + extra_libs
             _transform_cuda_paths(libraries_dirs)
 
     if device_type == "xpu":
@@ -1781,6 +2027,7 @@ def __init__(
         min_optimize: bool = False,
         precompiling: bool = False,
         preprocessing: bool = False,
+        compiler: str = "",
     ) -> None:
         super().__init__(
             vec_isa=vec_isa,
@@ -1794,6 +2041,7 @@ def __init__(
             min_optimize=min_optimize,
             precompiling=precompiling,
             preprocessing=preprocessing,
+            compiler=compiler,
         )
 
         device_definitions: list[str] = []
@@ -1906,7 +2154,7 @@ def __get_preprocessor_output_flags() -> tuple[str, str]:
     def __init__(
         self,
         name: str,
-        sources: Union[str, list[str]],
+        sources: str | list[str],
         BuildOption: BuildOptionsBase,
         output_dir: str = "",
     ) -> None:
@@ -1994,6 +2242,11 @@ def __init__(
             assert len(sources) == 1
             # See above; we can currently assume this is not on MSVC.
             self._sources_args = f"-x c++-header {sources[0]}"
+            if self._use_relative_path and _is_clang(BuildOption.get_compiler()):
+                # Store PCH paths relative to -isysroot so the .pch can
+                # be used from a different build directory.  The matching
+                # -isysroot is injected by build_fbcode_re().
+                self._cflags_args += " -relocatable-pch -Xclang -fno-pch-timestamp "
         else:
             self._sources_args = " ".join(sources)
 
@@ -2017,6 +2270,13 @@ def __init__(
                 )
             else:
                 self._include_dirs_args = f"-include {precompiled_header} "
+                if self._use_relative_path and _is_clang(BuildOption.get_compiler()):
+                    # Skip clang's own PCH validation during consumption.
+                    # _precompile_header() already handles cache invalidation
+                    # via content hashing, and -fno-validate-pch allows the
+                    # PCH to be used even when the original source file is at
+                    # a different path (e.g. across Remote Execution workers).
+                    self._cflags_args += " -Xclang -fno-validate-pch "
 
         for inc_dir in BuildOption.get_include_dirs():
             if _IS_WINDOWS:
@@ -2098,7 +2358,7 @@ def build_fbcode_re(
         self,
     ) -> None:
         with dynamo_timed("compile_file"):
-            command = self.get_command_line().split()
+            command = shlex.split(self.get_command_line())
             try:
                 output_path = self._target_file
                 # When we build remotely, we need to make sure to carefully copy any files
@@ -2112,9 +2372,43 @@ def build_fbcode_re(
                         shutil.copy(src, os.path.join(tmp_dir, os.path.basename(src)))
                     dest_include_path = os.path.join(tmp_dir, "include")
                     shutil.copytree(torch_includes_path, dest_include_path)
-                    # Run the build
+
+                    # Copy precompiled header (.h and .gch/.pch) into the
+                    # build directory and rewrite the -include flag so the
+                    # compiler can find it.
+                    pch_header = self._build_option.precompiled_header
+                    if pch_header and os.path.isfile(pch_header):
+                        pch_ext = ".pch" if _IS_WINDOWS or not is_gcc() else ".gch"
+                        pch_compiled = pch_header + pch_ext
+                        pch_basename = os.path.basename(pch_header)
+                        shutil.copy(pch_header, os.path.join(tmp_dir, pch_basename))
+                        if os.path.isfile(pch_compiled):
+                            shutil.copy(
+                                pch_compiled,
+                                os.path.join(tmp_dir, pch_basename + pch_ext),
+                            )
+                        command = [
+                            pch_basename if arg == pch_header else arg
+                            for arg in command
+                        ]
+
+                    # Relocatable PCH stores include paths relative to
+                    # -isysroot.  Set sysroot to the tmp build dir so
+                    # paths resolve correctly in both precompilation and
+                    # later kernel compilations that consume the PCH.
+                    if self._precompiling or (
+                        pch_header and os.path.isfile(pch_header)
+                    ):
+                        command[1:1] = ["-isysroot", "."]
+
+                    # Run the build, raising RuntimeError on failure instead of
+                    # SkipFrame so compilation errors propagate rather than
+                    # silently falling back to eager execution.
                     tmp_output_path = _run_build_command(
-                        command, tmp_dir, os.path.basename(output_path)
+                        command,
+                        tmp_dir,
+                        os.path.basename(output_path),
+                        exception_class=RuntimeError,
                     )
                     # Copy output from the build
                     if os.path.exists(output_path):
@@ -2125,8 +2419,7 @@ def build_fbcode_re(
                     elif output_path.endswith(".so"):
                         os.chmod(output_path, 0o755)
             except subprocess.CalledProcessError as e:
-                output = e.output.decode("utf-8")
-                raise exc.CppCompileError(command, output) from e
+                raise exc.CppCompileError(command, e.output.decode("utf-8")) from e
 
     def build(self) -> None:
         """
@@ -2165,7 +2458,7 @@ def save_compile_cmd_to_cmake(
             f"""
             cmake_minimum_required(VERSION 3.27 FATAL_ERROR)
             project({self._target_name} LANGUAGES CXX)
-            set(CMAKE_CXX_STANDARD 17)
+            set(CMAKE_CXX_STANDARD 20)
 
             # Set a library target
             add_library({self._target_name} {target_library_type})
diff --git a/torch/_inductor/cpu_vec_isa.py b/torch/_inductor/cpu_vec_isa.py
index 9793d165c7b24..e4aee2d6c0258 100644
--- a/torch/_inductor/cpu_vec_isa.py
+++ b/torch/_inductor/cpu_vec_isa.py
@@ -8,7 +8,7 @@
 import sys
 import warnings
 from collections.abc import Callable
-from typing import Any, Union
+from typing import Any
 
 import torch
 from torch._inductor import config
@@ -466,7 +466,7 @@ def _check_and_append_supported_isa(
 
 
 def get_isa_from_cpu_capability(
-    capability: Union[str, None],
+    capability: str | None,
     vec_isa_list: list[VecISA],
     invalid_vec_isa: InvalidVecISA,
 ):
diff --git a/torch/_inductor/cudagraph_trees.py b/torch/_inductor/cudagraph_trees.py
index c239132fcc4f9..d962820cdc9e2 100644
--- a/torch/_inductor/cudagraph_trees.py
+++ b/torch/_inductor/cudagraph_trees.py
@@ -51,13 +51,21 @@
 from collections import defaultdict
 from contextlib import AbstractContextManager
 from enum import auto, Enum
-from typing import Any, cast, Optional, TYPE_CHECKING, TypeVar, Union
+from typing import Any, cast, TYPE_CHECKING, TypeVar
 
 import torch.fx
 from torch import Tensor
 from torch._dynamo.callback import CallbackTrigger
+from torch._dynamo.graph_bytecode_inputs import (
+    CURRENT_STREAM_INDEX,
+    set_external_object_by_index,
+)
 from torch._dynamo.mutation_guard import GenerationTracker
 from torch._dynamo.utils import counters, dynamo_timed, preserve_rng_state
+from torch._higher_order_ops.cudagraph_conditional_nodes import (
+    ControlFlowOpWarmupDispatchMode,
+    CUDAGraphCaptureControlFlowOpDispatchMode,
+)
 from torch._inductor.compile_fx import (
     align_inputs_from_check_idxs,
     copy_misaligned_inputs,
@@ -80,6 +88,8 @@
     PlaceholderInfo,
     WrappedFunction,
 )
+from torch._library.opaque_object import is_opaque_value
+from torch._opaque_base import OpaqueBase
 from torch.multiprocessing.reductions import StorageWeakRef
 from torch.storage import UntypedStorage
 from torch.utils import _pytree as pytree
@@ -234,7 +244,7 @@ def __init__(self, device_index: int) -> None:
         # This class keeps a strong reference to tree_manager,
         # but upon all other strong references to the tree_manager will reset it to None.
         # We need a strong reference so that we can still access its attributes upon cleanup.
-        self.tree_manager: Optional[CUDAGraphTreeManager] = None
+        self.tree_manager: CUDAGraphTreeManager | None = None
 
         # Number of outstanding references to the current tree manager
         self.live_cudagraphify_fns = 0
@@ -245,7 +255,7 @@ def __init__(self, device_index: int) -> None:
         # the cudagraphify_fns. Reference to the Graph is needed to keep the private pool from
         # deallocation.
         self.live_storages_count = 0
-        self.graph: Optional[torch.cuda.CUDAGraph] = None
+        self.graph: torch.cuda.CUDAGraph | None = None
 
         self.lock = threading.Lock()
 
@@ -373,13 +383,13 @@ def get_container(device_index: int) -> TreeManagerContainer:
 
 def get_manager(
     device_index: int, create_if_none_exists: bool = True
-) -> Optional[CUDAGraphTreeManager]:
+) -> CUDAGraphTreeManager | None:
     if create_if_none_exists:
         return get_container(device_index).get_tree_manager()
     return get_container(device_index).tree_manager
 
 
-def is_cudagraph_capture_sizes(int_key: Union[int, tuple[int, ...]]) -> bool:
+def is_cudagraph_capture_sizes(int_key: int | tuple[int, ...]) -> bool:
     """
     Returns true if all dynamic shapes should be captured or the dynamic shape
     int_key should be captured.
@@ -450,8 +460,8 @@ def deferred_cudagraphify(inputs: list[InputType]) -> OutputType:
 @contextlib.contextmanager
 def dynamo_timed_cudagraph(
     name: str,
-    compile_id: Optional[CompileId],
-    mode: Optional[CompilationMode],
+    compile_id: CompileId | None,
+    mode: CompilationMode | None,
 ) -> Generator[Any, None, None]:
     """
     Makes usages of dynamo_timed in this file less verbose. NOTE: This CM sums
@@ -477,11 +487,11 @@ def cudagraphify(
     device_index: int,
     is_backward: bool,
     is_inference: bool,
-    stack_traces: Optional[StackTraces] = None,
+    stack_traces: StackTraces | None = None,
     constants: tuple[torch.Tensor, ...] = (),
     placeholders: tuple[PlaceholderInfo, ...] = (),
     mutated_input_idxs: tuple[int, ...] = (),
-    compile_id: Optional[CompileId] = None,
+    compile_id: CompileId | None = None,
 ) -> tuple[ModelType, OutputType]:
     assert not (is_backward and is_inference)
     mode = (
@@ -513,12 +523,12 @@ class StorageWeakRefWrapper:
 
     __slots__ = ["ref", "_data_ptr", "extra_ref_check"]
 
-    storage_ref: Optional[StorageWeakRef]
+    storage_ref: StorageWeakRef | None
 
     def __init__(
         self,
-        inp: Union[Tensor, UntypedStorage],
-        extra_ref_check: Optional[Callable[[], bool]] = None,
+        inp: Tensor | UntypedStorage,
+        extra_ref_check: Callable[[], bool] | None = None,
     ) -> None:
         """
         extra_ref_check is an additional check we need to run to check if the
@@ -539,7 +549,7 @@ def from_weakref_and_data_ptr(
         cls: type[StorageWeakRefWrapper],
         cdata: Any,
         data_ptr: int,
-        extra_ref_check: Optional[Callable[[], bool]] = None,
+        extra_ref_check: Callable[[], bool] | None = None,
     ) -> StorageWeakRefWrapper:
         instance = cls.__new__(cls)
         instance._data_ptr = data_ptr
@@ -547,7 +557,7 @@ def from_weakref_and_data_ptr(
         instance.extra_ref_check = extra_ref_check
         return instance
 
-    def __call__(self) -> Optional[StorageWeakRefPointer]:
+    def __call__(self) -> StorageWeakRefPointer | None:
         if self.expired():
             return None
 
@@ -584,13 +594,13 @@ def __repr__(self) -> str:
             return f"StorageWeakRefWrapper to {self.data_ptr()}; alive"
 
 
-def is_live(weak_ref: Optional[StorageWeakRefWrapper]) -> bool:
+def is_live(weak_ref: StorageWeakRefWrapper | None) -> bool:
     return maybe_deref(weak_ref) is not None
 
 
 def maybe_deref(
-    weak_ref: Optional[StorageWeakRefWrapper],
-) -> Optional[tuple[StorageWeakRefPointer, int]]:
+    weak_ref: StorageWeakRefWrapper | None,
+) -> tuple[StorageWeakRefPointer, int] | None:
     if weak_ref is None:
         return None
     r = weak_ref()
@@ -627,7 +637,20 @@ def _use_cuda_memory_pool_manager(
     torch.cuda.current_stream().wait_stream(stream)
 
 
-def map_to_ref(t: Optional[Tensor]) -> Optional[StorageWeakRefWrapper]:
+@contextlib.contextmanager
+def _update_current_stream_external_object() -> Generator[None, None, None]:
+    """Update the external object registry so custom ops see the capture stream.
+
+    During cudagraph recording/warmup the current stream differs from the
+    trace-time default stream.  The external object at CURRENT_STREAM_INDEX
+    must reflect the actual current stream so that custom ops (e.g. event
+    record/wait) executed during capture use the right stream.
+    """
+    set_external_object_by_index(CURRENT_STREAM_INDEX, torch.cuda.current_stream())
+    yield
+
+
+def map_to_ref(t: Tensor | None) -> StorageWeakRefWrapper | None:
     if not isinstance(t, torch.Tensor):
         assert t is None
         return None
@@ -641,7 +664,7 @@ def map_to_ref(t: Optional[Tensor]) -> Optional[StorageWeakRefWrapper]:
 # For each node in the path, for each output, is the output alive
 PathLiveness = list[list[bool]]
 
-StackTraces = list[Optional[str]]
+StackTraces = list[str | None]
 
 
 class CUDAWarmupNode:
@@ -666,20 +689,20 @@ class CUDAWarmupNode:
     def __init__(
         self,
         wrapped_function: WrappedFunction,
-        parent: Optional[Union[CUDAGraphNode, CUDAWarmupNode]],
+        parent: CUDAGraphNode | CUDAWarmupNode | None,
         cuda_graphs_pool: tuple[int, int],
-        existing_cuda_graph: Optional[torch.cuda.CUDAGraph],
+        existing_cuda_graph: torch.cuda.CUDAGraph | None,
         device_index: int,
-        stack_traces: Optional[StackTraces],
+        stack_traces: StackTraces | None,
         stream: torch.cuda.Stream,
         already_warm: bool,
         id: GraphID,
     ) -> None:
         self.wrapped_function = wrapped_function
-        self.parent: Optional[Union[CUDAGraphNode, CUDAWarmupNode]] = parent
+        self.parent: CUDAGraphNode | CUDAWarmupNode | None = parent
         self.cuda_graphs_pool = cuda_graphs_pool
-        self.outputs_weakrefs: list[Optional[StorageWeakRefWrapper]] = []
-        self.tensor_weakrefs: list[Optional[TensorWeakRef]] = []
+        self.outputs_weakrefs: list[StorageWeakRefWrapper | None] = []
+        self.tensor_weakrefs: list[TensorWeakRef | None] = []
         self.existing_cuda_graph = existing_cuda_graph
         self.has_run = False
         self.device_index = device_index
@@ -719,6 +742,9 @@ def get_non_cudagraph_inps() -> list[weakref.ReferenceType[UntypedStorage]]:
             _use_cuda_memory_pool_manager(
                 self.device_index, self.cuda_graphs_pool, self.stream
             ),
+            # NB: must go after _use_cuda_memory_pool_manager which switches the stream
+            _update_current_stream_external_object(),
+            ControlFlowOpWarmupDispatchMode(),
             get_history_recording(),
         ):
             out = self.wrapped_function.model(new_inputs)
@@ -762,9 +788,9 @@ def add_ref(o: Any) -> bool:
     @property
     def _path_from_root(
         self,
-    ) -> Generator[Union[CUDAGraphNode, CUDAWarmupNode], None, None]:
+    ) -> Generator[CUDAGraphNode | CUDAWarmupNode, None, None]:
         nodes = []
-        node: Union[CUDAGraphNode, CUDAWarmupNode] = self
+        node: CUDAGraphNode | CUDAWarmupNode = self
         while node:
             nodes.append(node)
             node = node.parent  # type: ignore[assignment]
@@ -854,14 +880,14 @@ def __init__(
         self,
         wrapped_function: WrappedFunction,
         id: GraphID,
-        parent: Optional[CUDAGraphNode],
+        parent: CUDAGraphNode | None,
         inputs: list[InputType],
         cuda_graphs_pool: _POOL_HANDLE,
         device_index: int,
-        stack_traces: Optional[StackTraces],
+        stack_traces: StackTraces | None,
         stream: torch.cuda.Stream,
-        mode: Optional[CompilationMode],
-        compile_id: Optional[CompileId],
+        mode: CompilationMode | None,
+        compile_id: CompileId | None,
     ) -> None:
         assert isinstance(inputs, (list, tuple))
 
@@ -873,10 +899,7 @@ def __init__(
 
         # Enable re-record a cudagraph when static tensor address changed.
         # if not we should error when it changed.
-        self.rerecord_if_static_inputs_change = (
-            torch._dynamo.config.inline_inbuilt_nn_modules
-            or torch._inductor.config.triton.cudagraph_support_input_mutation
-        )
+        self.rerecord_if_static_inputs_change = True
 
         # if this is a root parent will be None. use weakref to prevent reference cycle
         self._parent = weakref.ref(parent) if parent is not None else None
@@ -900,14 +923,14 @@ def __init__(
         # in children to avoid children having to chase parent pointers in the hot path
         # DO NOT reassign output_weakrefs, only call `clear()`
         # Path is a series of nodes from root to the current node
-        self.outputs_weakrefs: OutputList[Optional[StorageWeakRefWrapper]] = []
-        self.path_weakrefs: LevelList[OutputList[Optional[StorageWeakRefWrapper]]] = [
+        self.outputs_weakrefs: OutputList[StorageWeakRefWrapper | None] = []
+        self.path_weakrefs: LevelList[OutputList[StorageWeakRefWrapper | None]] = [
             node.outputs_weakrefs for node in self._path_from_root
         ]
-        self.path_stacktraces: LevelList[Optional[StackTraces]] = [
+        self.path_stacktraces: LevelList[StackTraces | None] = [
             node.stack_traces for node in self._path_from_root
         ]
-        self.tensor_weakrefs: OutputList[Optional[TensorWeakRef]] = []
+        self.tensor_weakrefs: OutputList[TensorWeakRef | None] = []
 
         # tensors which are outputs of previous graphs in the tree
         self.cudagraph_managed_idxs: list[int] = [
@@ -917,7 +940,7 @@ def __init__(
         ]
 
         # (depth, offset) of live tensors which are alias of previous graph outputs
-        self.live_cudagraph_managed_path_refs: InputList[Optional[PathOutputIndex]] = [
+        self.live_cudagraph_managed_path_refs: InputList[PathOutputIndex | None] = [
             (
                 self._is_alias_of_live_recorded_tensor(t)
                 if isinstance(t, torch.Tensor)
@@ -930,9 +953,26 @@ def __init__(
         # and also aliases an output of the current CUDAGraphNode
         self.preserved_aliased_inputs: InputList[bool] = [False] * len(inputs)
 
+        # Opaque values (e.g. DeviceMesh, ProcessGroup) are non-tensor
+        # inputs that cannot be copied like tensors. We include them in
+        # static_input_idxs to keep them out of non_static_input_idx
+        # (which drives the tensor-copy path during replay). "Static"
+        # here just means "don't try to copy this as a tensor" — it
+        # does NOT mean the object is semantically immutable.
+        #
+        # Opaque indices must also be excluded from any list passed to
+        # _tensors_data_ptrs_at_indices_equal (the C++ data-pointer
+        # stability check), because opaque objects have no data_ptr.
+        # That is why tensor_static_input_idxs and
+        # non_managed_static_input_idxs filter them out below.
+        opaque_input_idxs = OrderedSet(
+            i for i, inp in enumerate(inputs) if is_opaque_value(inp)
+        )
+        static_input_idxs = OrderedSet(wrapped_function.static_input_idxs)
+        cudagraph_managed_idxs = OrderedSet(self.cudagraph_managed_idxs)
+
         self.static_input_idxs: list[int] = list(
-            OrderedSet(wrapped_function.static_input_idxs)
-            | OrderedSet(self.cudagraph_managed_idxs)
+            static_input_idxs | cudagraph_managed_idxs | opaque_input_idxs
         )
 
         self.non_static_input_idx: LevelList[int] = [
@@ -943,23 +983,25 @@ def __init__(
             self.non_static_input_idx
         )
 
-        self.non_managed_static_input_idxs: LevelList[int] = [
-            i
-            for i in wrapped_function.static_input_idxs
-            if i not in self.cudagraph_managed_idxs
-        ]
+        self.non_managed_static_input_idxs: LevelList[int] = LevelList(
+            static_input_idxs - cudagraph_managed_idxs - opaque_input_idxs
+        )
+
+        self.tensor_static_input_idxs: list[int] = list(
+            static_input_idxs | cudagraph_managed_idxs
+        )
 
         def maybe_get_static_data_ptr(
             idx: int,
             inputs: list[InputType],
             static_input_idxs: list[int],
-        ) -> Optional[int]:
+        ) -> int | None:
             inp = inputs[idx]
             if isinstance(inp, torch.Tensor) and idx in static_input_idxs:
                 return inp.data_ptr()
             return None
 
-        self.static_input_data_ptrs: InputList[Optional[int]] = [
+        self.static_input_data_ptrs: InputList[int | None] = [
             maybe_get_static_data_ptr(i, inputs, self.static_input_idxs)
             for i in range(len(inputs))
         ]
@@ -1014,7 +1056,7 @@ def maybe_get_static_data_ptr(
         del inputs
 
         # graph used for recording model invocation
-        self.graph: Optional[torch.cuda.CUDAGraph] = torch.cuda.CUDAGraph()
+        self.graph: torch.cuda.CUDAGraph | None = torch.cuda.CUDAGraph()
 
         # TODO: register_generator_state should potentially take explicit device
         with torch.cuda.device(self.device):
@@ -1046,14 +1088,14 @@ def maybe_get_static_data_ptr(
 
         # initialized below in _record
 
-        self.checkpointed_caching_state: Optional[AllocatorState] = None
+        self.checkpointed_caching_state: AllocatorState | None = None
 
         # Output Storage Alias information, can be:
         # - A new, unaliased storage, or the output is None
         # - An alias of an output of a prior graph
         # - An alias of an output already created in the reconstructed outputs
         # This is None if the output in question is an int
-        self.output_storage_alias: OutputList[Optional[OutputAliasInfo]] = []
+        self.output_storage_alias: OutputList[OutputAliasInfo | None] = []
 
         # is the output Storage unaliased in subsequent outputs, of all subsequent paths
         # if it is, we cached the output tensor and adjust storage liveness tracking to also
@@ -1065,19 +1107,19 @@ def maybe_get_static_data_ptr(
         # The cached tensor outputs are added on the first execution, and cleared whenever we need
         # to do subsequent recording
         self.unaliased_in_all_paths: OutputList[bool] = []
-        self.cached_tensor_outputs: OutputList[Optional[Tensor]] = []
+        self.cached_tensor_outputs: OutputList[Tensor | None] = []
 
         # if an output aliases a static, persistent input then the corresponding Tensor will
         # be set here. These are different than cached tensors, because they are tensors that
         # are aliases of parameters that are always live.
-        self.static_output_tensors: OutputList[Optional[Tensor]] = []
+        self.static_output_tensors: OutputList[Tensor | None] = []
 
         # Cleared after recording
         with dynamo_timed_cudagraph("CUDAGraphNode.record", compile_id, mode):
-            self.recording_outputs: Optional[OutputType] = self._record(
+            self.recording_outputs: OutputType | None = self._record(
                 wrapped_function.model, recording_inputs
             )
-        self.outputs_metadata: OutputList[Union[dict[str, Any], int, None]] = []
+        self.outputs_metadata: OutputList[dict[str, Any] | int | None] = []
 
         # As with inputs, we do not want to keep the outputs permanently alive because that would prevent
         # their memory being reclaimed in subsequent cuda graph recordings. We record the tensor metadata
@@ -1223,9 +1265,9 @@ def reconstruct_outputs(self) -> OutputType:
 
     def prepare_alias_info_for_tensor_construction(
         self,
-        out_alias_info: Optional[OutputAliasInfo],
-        metadata: Union[dict[str, Any], int, None],
-    ) -> Union[UntypedStorage, None, int]:
+        out_alias_info: OutputAliasInfo | None,
+        metadata: dict[str, Any] | int | None,
+    ) -> UntypedStorage | None | int:
         if (
             isinstance(metadata, (int, type(None)))
             or out_alias_info is UnaliasedStorage
@@ -1243,7 +1285,7 @@ def prepare_alias_info_for_tensor_construction(
 
     def prepare_storages_for_construction(
         self,
-    ) -> list[Union[UntypedStorage, None, int]]:
+    ) -> list[UntypedStorage | None | int]:
         output_storages = []
         for output_storage_alias, metadata in zip(
             self.output_storage_alias, self.outputs_metadata
@@ -1311,6 +1353,9 @@ def static_input_iter() -> Generator[torch.Tensor, None, None]:
                 pool=self.cuda_graphs_pool,
                 capture_error_mode="thread_local",
             ),
+            # NB: must go after torch.cuda.graph which switches the stream
+            _update_current_stream_external_object(),
+            CUDAGraphCaptureControlFlowOpDispatchMode(),
             get_history_recording(),
         ):
             static_outputs = model(inputs)
@@ -1494,7 +1539,7 @@ def get_output_refcount(self, index: int) -> int:
         return sys.getrefcount(self.cached_tensor_outputs[index])
 
     @property
-    def parent(self) -> Optional[CUDAGraphNode]:
+    def parent(self) -> CUDAGraphNode | None:
         "unwraps the weakref to _parent"
         return self._parent() if self._parent is not None else None
 
@@ -1529,7 +1574,7 @@ def _is_cuda_graph_recorded_tensor(self, t: torch.Tensor) -> bool:
 
     def _is_alias_of_live_recorded_tensor(
         self, t: torch.Tensor
-    ) -> Optional[PathOutputIndex]:
+    ) -> PathOutputIndex | None:
         for depth, output_refs in enumerate(self.path_weakrefs):
             for output_index, storage_ref in enumerate(output_refs):
                 if (storage_and_ptr := maybe_deref(storage_ref)) is not None:
@@ -1542,7 +1587,7 @@ def _is_alias_of_live_recorded_tensor(
     @staticmethod
     def _check_liveness(
         indices: list[PathOutputIndex],
-        output_refs: list[list[Optional[StorageWeakRefWrapper]]],
+        output_refs: list[list[StorageWeakRefWrapper | None]],
     ) -> bool:
         "Check that all of the indices specified are dead references"
         for depth, output_index in indices:
@@ -1573,7 +1618,7 @@ def _get_different_indices(
 
     @staticmethod
     def _get_liveness(
-        weakrefs: list[list[Optional[StorageWeakRefWrapper]]],
+        weakrefs: list[list[StorageWeakRefWrapper | None]],
     ) -> list[list[bool]]:
         "Maps weakrefs to true if the reference is alive and false otherwise"
         if len(weakrefs) == 0:
@@ -1693,7 +1738,7 @@ def _tensor_metadata(
         }
 
     def _reconstruct_from_tensor_metadata(
-        self, metadata: dict[str, Any], storage: Optional[UntypedStorage] = None
+        self, metadata: dict[str, Any], storage: UntypedStorage | None = None
     ) -> Tensor:
         s = self.create_storage(metadata) if storage is None else storage
         return torch._C._construct_CUDA_Tensor_From_Storage_And_Metadata(metadata, s)  # type: ignore[arg-type]
@@ -1726,7 +1771,7 @@ def _allocate_and_copy_recording_inputs(
         ):
             for i, inp in enumerate(inputs):
                 if not isinstance(inp, torch.Tensor):
-                    assert isinstance(inp, (int, torch.Generator))
+                    assert isinstance(inp, (int, torch.Generator, OpaqueBase))
 
                     recording_inputs.append(inp)
                 elif i not in self.static_input_idxs:
@@ -1785,13 +1830,13 @@ def check_invariants(
             and not torch._C._tensors_data_ptrs_at_indices_equal(
                 inputs,  # type: ignore[arg-type]
                 self.static_input_data_ptrs,
-                self.static_input_idxs,
+                self.tensor_static_input_idxs,
             )
         ):
             status = CheckInvariantStatus.StaticInputIdxMismatch
             _logger = functools.partial(
                 _logger,
-                self.static_input_idxs,
+                self.tensor_static_input_idxs,
                 status,
             )
             return status, _logger
@@ -2003,7 +2048,7 @@ def __init__(self, device_index: int) -> None:
         # mapping from function id to wrapped function
         self.ids_to_funcs: dict[FunctionID, WrappedFunction] = {}
 
-        self.ids_to_stack_traces: dict[FunctionID, Optional[StackTraces]] = {}
+        self.ids_to_stack_traces: dict[FunctionID, StackTraces | None] = {}
 
         self.warmed_up_functions: OrderedSet[FunctionID] = OrderedSet()
         # if we fail to increment generation, and are stuck warming up,
@@ -2026,7 +2071,7 @@ def __init__(self, device_index: int) -> None:
             self.stream.wait_stream(torch.cuda.current_stream())
 
             # Keeps Memory Pool Alive
-            self.graph: Optional[torch.cuda.CUDAGraph] = torch.cuda.CUDAGraph()
+            self.graph: torch.cuda.CUDAGraph | None = torch.cuda.CUDAGraph()
             self.cuda_graphs_thread_pool = torch.cuda.graph_pool_handle()
 
             with (
@@ -2046,13 +2091,13 @@ def __init__(self, device_index: int) -> None:
         # mapping from graph_id to (function id to mutation type hint) since we are
         # specializing on a particular combination of Parent Node -> Function ID.
         self.non_cudagraph_managed_mutation_hint: dict[
-            Optional[GraphID], dict[FunctionID, bool]
+            GraphID | None, dict[FunctionID, bool]
         ] = defaultdict(dict)
         self.warmup_node_counter = itertools.count(start=-1, step=-1)
 
         # mapping from graph_id to (function id to re-record count). We fall back to
         # eager function if a function is re-recorded frequently on a node.
-        self.num_rerecord: dict[Optional[GraphID], dict[FunctionID, int]] = defaultdict(
+        self.num_rerecord: dict[GraphID | None, dict[FunctionID, int]] = defaultdict(
             lambda: defaultdict(lambda: 0)
         )
 
@@ -2065,7 +2110,7 @@ def __init__(self, device_index: int) -> None:
         # when there is no output from a previous recording or execution whose memory
         # we need to respect in the cuda caching allocation. If you incremented generation,
         # this will also be none, as ignore those allocations.
-        self.current_node: Optional[Union[CUDAGraphNode, CUDAWarmupNode]] = None
+        self.current_node: CUDAGraphNode | CUDAWarmupNode | None = None
 
         # current generation of cudagraph invocations. when torch.compile is run
         # we increment the current generation. are willing to ignore live outputs
@@ -2079,7 +2124,7 @@ def __init__(self, device_index: int) -> None:
         self.debug_checkpointing_counter = 0
 
         self.id_to_mode: dict[FunctionID, CompilationMode] = {}
-        self.id_to_compile_id: dict[FunctionID, Optional[CompileId]] = {}
+        self.id_to_compile_id: dict[FunctionID, CompileId | None] = {}
 
         # Note: [Backward Generation Handling]
         # We generally perform a sequence of forward executions followed by backward executions.
@@ -2096,7 +2141,7 @@ def __init__(self, device_index: int) -> None:
         # mod2(mod1(x)).sum().backward()
 
         self.running_forwards_with_pending_backwards = False
-        self.mode: Optional[CompilationMode] = None
+        self.mode: CompilationMode | None = None
 
         self.disable_invalidate_aliases = (
             False
@@ -2152,7 +2197,7 @@ def _update_non_cudagraph_managed_mutation(
         else:
             self.non_cudagraph_managed_mutation_hint[node_id][function_id] = False
 
-    def _get_node_id(self) -> Optional[GraphID]:
+    def _get_node_id(self) -> GraphID | None:
         if self.current_node is None:
             return None
         elif isinstance(self.current_node, (CUDAGraphNode, CUDAWarmupNode)):
@@ -2161,15 +2206,9 @@ def _get_node_id(self) -> Optional[GraphID]:
             raise RuntimeError(f"Unknown node type {type(self.current_node)}")
 
     def exceed_rerecord_limit(
-        self, node_id: Optional[GraphID], function_id: FunctionID
+        self, node_id: GraphID | None, function_id: FunctionID
     ) -> bool:
-        if torch._dynamo.config.inline_inbuilt_nn_modules:
-            return False
-
-        return (
-            self.num_rerecord[node_id][function_id]
-            > torch._inductor.config.triton.cudagraph_unexpected_rerecord_limit
-        )
+        return False
 
     def _run(self, new_inputs: list[InputType], function_id: FunctionID) -> OutputType:
         # we will try to end the current execution lazily, since
@@ -2241,9 +2280,10 @@ def _run(self, new_inputs: list[InputType], function_id: FunctionID) -> OutputTy
                     if log.isEnabledFor(logging.DEBUG):
                         unexpected_rerecord_reason = status_logger()
                         log.debug(
-                            "[%s] Re-recording function=%s, reason=%s",
+                            "[%s] Re-recording function=%s, mode=%s, reason=%s",
                             self.compile_id,
                             self.get_func_name(function_id),
+                            self.id_to_mode[function_id].name,
                             unexpected_rerecord_reason,
                         )
                     else:
@@ -2332,9 +2372,10 @@ def record_function(
         ):
             graph_id = self.new_graph_id()
             log.debug(
-                "[%s] Recording function=%s, cuda_graph_id=%d, inputs: %s",
+                "[%s] Recording function=%s, mode=%s, cuda_graph_id=%d, inputs: %s",
                 self.compile_id,
                 self.get_func_name(function_id),
+                self.id_to_mode[function_id].name,
                 graph_id.id,
                 format_inputs_log(new_inputs),
             )
@@ -2378,9 +2419,10 @@ def run_eager(
         func_name = self.get_func_name(function_id)
         if not already_warm:
             log.debug(
-                "[%s] Running warmup function=%s",
+                "[%s] Running warmup function=%s, mode=%s",
                 self.compile_id,
                 func_name,
+                self.id_to_mode[function_id].name,
             )
         else:
             log.debug(
@@ -2419,12 +2461,12 @@ def add_function(
         model: ModelType,
         inputs: list[InputType],
         static_input_idxs: Sequence[int],
-        stack_traces: Optional[StackTraces],
+        stack_traces: StackTraces | None,
         mode: CompilationMode,
         constants: tuple[torch.Tensor, ...],
         placeholders: tuple[PlaceholderInfo, ...],
         mutated_input_idxs: tuple[int, ...],
-        compile_id: Optional[CompileId],
+        compile_id: CompileId | None,
     ) -> tuple[
         ModelType,
         OutputType,
@@ -2460,13 +2502,11 @@ def get_roots(self) -> Iterator[CUDAGraphNode]:
             yield from nodes
 
     @property
-    def current_node(self) -> Optional[Union[CUDAGraphNode, CUDAWarmupNode]]:
+    def current_node(self) -> CUDAGraphNode | CUDAWarmupNode | None:
         return self._current_node
 
     @current_node.setter
-    def current_node(
-        self, value: Optional[Union[CUDAGraphNode, CUDAWarmupNode]]
-    ) -> None:
+    def current_node(self, value: CUDAGraphNode | CUDAWarmupNode | None) -> None:
         self._current_node = value
         if value is None:
             self.path_state = ExecutionState.NONE
@@ -2580,13 +2620,13 @@ def check_warn_on_unable_to_start_executing(self, function_id: FunctionID) -> No
 
         self.warned_functions.add(function_id)
         warnings.warn(
-            "Unable to hit fast path of CUDAGraphs because of pending, uninvoked backwards. "
-            "Consider running with torch.no_grad() or using torch.compiler.cudagraph_mark_step_begin() "
-            "before each model invocation"
+            "Unable to hit fast path of CUDAGraphs because outputs from a previous step "
+            "still require backward. Ensure backward() is invoked or detach outputs. "
+            "You may also call torch.compiler.cudagraph_mark_step_begin() before each model invocation."
         )
 
     @staticmethod
-    def format_dealloc_msg(stack_trace: Optional[str]) -> str:
+    def format_dealloc_msg(stack_trace: str | None) -> str:
         stack_trace = (
             stack_trace.strip() if stack_trace else "[Could not find stack trace]"
         )
@@ -2602,7 +2642,7 @@ def dealloc_current_path_weakrefs(self) -> None:
         # TODO: we could also allow the these weak refs to continue to be allocated,
         # but that adds some complications.
 
-        stor_stack_trace: dict[int, Optional[str]] = {}
+        stor_stack_trace: dict[int, str | None] = {}
         for node in self.current_node._path_from_root:
             assert node.stack_traces is not None
             assert len(node.tensor_weakrefs) == len(node.stack_traces)
diff --git a/torch/_inductor/cudagraph_utils.py b/torch/_inductor/cudagraph_utils.py
index 46ba65bdcb929..d2c443c50cec2 100644
--- a/torch/_inductor/cudagraph_utils.py
+++ b/torch/_inductor/cudagraph_utils.py
@@ -4,7 +4,7 @@
 import dataclasses
 from collections.abc import Callable
 from enum import Enum
-from typing import Any, Optional, TYPE_CHECKING, Union
+from typing import Any, TYPE_CHECKING, TypeVar
 
 import torch
 from torch._dynamo.utils import counters, get_metrics_context
@@ -18,6 +18,10 @@
 if TYPE_CHECKING:
     from collections.abc import Sequence, Set as AbstractSet
 
+    from torch._inductor.output_code import OutputCode
+
+_OC = TypeVar("_OC", bound="OutputCode")
+
 
 cudagraphs_log = torch._logging.getArtifactLogger(__name__, "cudagraphs")
 static_inputs_log = torch._logging.getArtifactLogger(
@@ -25,10 +29,97 @@
 )
 
 
-OutputType = list[Optional[Union[int, torch.Tensor]]]
+OutputType = list[int | torch.Tensor | None]
 ModelType = Callable[[list[InputType]], OutputType]
 
 
+class CUDAGraphPolicy:
+    """Pluggable policy controlling CUDA graph wrapping in Inductor's post_compile.
+
+    Override methods to customize:
+      - HOW compiled functions are cudagraph-wrapped (cudagraphify)
+      - WHETHER inner CompiledFxGraphs should be wrapped (should_wrap)
+      - OUTER wrapping of compound outputs like RegionalOutputCode (wrap_output)
+
+    Set via ``torch._inductor.config.cudagraph_policy``.  When ``None``
+    (the default), the existing built-in behaviour is used unchanged.
+
+    Example usage::
+
+        class MyCUDAGraphPolicy(CUDAGraphPolicy):
+            def cudagraphify(self, model, example_inputs, static_input_idxs, **kwargs):
+                return my_custom_wrapper(model, example_inputs, static_input_idxs)
+
+
+        with torch._inductor.config.patch("cudagraph_policy", MyCUDAGraphPolicy()):
+            compiled_fn = deserialize_artifacts(...)
+    """
+
+    def cudagraphify(
+        self,
+        model: Callable[..., Any],
+        example_inputs: Sequence[InputType],
+        static_input_idxs: Sequence[int],
+        *,
+        device_index: int,
+        is_backward: bool,
+        is_inference: bool,
+        **kwargs: Any,
+    ) -> Callable[..., Any]:
+        """Wrap a single compiled callable with CUDA graph capture/replay.
+
+        Called by ``cudagraph_post_compile`` for each ``CompiledFxGraph``.
+        The default delegates to ``compile_fx.cudagraphify`` (cudagraph_trees).
+
+        ``example_inputs`` are the example inputs at post_compile time.
+        The default implementation does not forward them because
+        ``compile_fx.cudagraphify`` defers graph recording to the first
+        real call via an inner closure.  Subclasses that need the
+        example inputs for warmup or static-input detection may use them.
+
+        When ``config.graph_partition=True``, setting a CUDAGraphPolicy
+        bypasses ``cudagraph_partition_post_compile`` (which wraps each
+        partition individually) and routes through ``cudagraph_post_compile``
+        instead, so this method wraps the *entire* callable, not individual
+        partitions.  Subclasses that need per-partition control should
+        handle partitioning internally.
+        """
+        from torch._inductor.compile_fx import cudagraphify
+
+        return cudagraphify(
+            model,
+            static_input_idxs,
+            device_index=device_index,
+            is_backward=is_backward,
+            is_inference=is_inference,
+            **kwargs,
+        )
+
+    def should_wrap(self, compiled_graph: OutputCode) -> bool:
+        """Whether to apply cudagraph wrapping to this CompiledFxGraph.
+
+        Called for each inner ``CompiledFxGraph`` during ``post_compile``.
+        Return ``False`` to skip wrapping (e.g. when wrapping at the outer
+        level via ``wrap_output`` instead).
+
+        Default: ``True`` (wrap everything, same as current behaviour).
+        """
+        return True
+
+    def wrap_output(self, output_code: _OC) -> _OC:
+        """Optional outer-level wrapping after inner post_compile completes.
+
+        Called by ``_compile_fx_inner``, ``BundledOutputCodeLoadable.post_compile``,
+        and ``FxGraphCacheLoadable.post_compile`` on the ``OutputCode`` returned
+        from ``post_compile``.  Subclasses that only want to wrap specific
+        output types should check ``isinstance`` and return the input
+        unchanged for types they don't handle.
+
+        Default: identity (no outer wrapping).
+        """
+        return output_code
+
+
 @dataclasses.dataclass(frozen=True, slots=True)
 class FunctionID:
     "Unique counter of a function wrapped in cudagraphify_impl"
@@ -45,10 +136,10 @@ class PlaceholderInfo:
     """
 
     name: str
-    stack_trace: Optional[str]
+    stack_trace: str | None
     # This field is recursive, but never cyclic (since a node never uses itself)
     users: list[PlaceholderInfo]
-    mutating_use_stack_trace: Optional[str]
+    mutating_use_stack_trace: str | None
 
 
 @dataclasses.dataclass(frozen=True, slots=True)
@@ -69,7 +160,7 @@ class WrappedFunction:
 
 def get_mutating_use_stack_trace_from_node(
     placeholder_node: torch.fx.Node,
-) -> Optional[str]:
+) -> str | None:
     # reinplaced uses might have a single, non-copy_ use
     if len(placeholder_node.users) == 1:
         return next(iter(placeholder_node.users)).meta.get("stack_trace", None)
@@ -82,7 +173,7 @@ def get_mutating_use_stack_trace_from_node(
     return None
 
 
-def get_mutating_use_stack_trace(placeholder_info: PlaceholderInfo) -> Optional[str]:
+def get_mutating_use_stack_trace(placeholder_info: PlaceholderInfo) -> str | None:
     return placeholder_info.mutating_use_stack_trace
 
 
@@ -113,9 +204,9 @@ def format_default_skip_message(reason: str) -> str:
 
 def get_mutation_stack_trace(
     placeholders: Sequence[PlaceholderInfo],
-    mutation_indices: Union[AbstractSet[int], Sequence[int]],
+    mutation_indices: AbstractSet[int] | Sequence[int],
 ) -> str:
-    stack_trace: Optional[str] = ""
+    stack_trace: str | None = ""
 
     for idx in mutation_indices:
         placeholder = placeholders[idx]
@@ -135,7 +226,7 @@ def check_for_mutation(
     func: WrappedFunction,
     inputs: list[InputType],
     is_cuda_graph_recorded_tensor: Callable[[torch.Tensor], bool],
-) -> Optional[str]:
+) -> str | None:
     # doesn't work for non-trees because the warmup run would apply mutation twice
     if torch._inductor.config.triton.cudagraph_trees:
         # checking if mutation is only on parameters/static inputs
@@ -162,7 +253,7 @@ def check_for_mutation(
     )
 
 
-def _get_use_stack_trace(node: torch.fx.Node) -> Optional[str]:
+def _get_use_stack_trace(node: torch.fx.Node) -> str | None:
     for use in node.users:
         if stack_trace := use.meta.get("stack_trace", None):
             return stack_trace
@@ -171,7 +262,7 @@ def _get_use_stack_trace(node: torch.fx.Node) -> Optional[str]:
 
 def check_multiple_devices_or_any_cpu_nodes(
     device_node_mapping: dict[torch.device, torch.fx.Node],
-) -> Optional[str]:
+) -> str | None:
     # meta tensors are supported since there is no compute
     device_node_mapping.pop(torch.device("meta"), None)
 
@@ -199,7 +290,7 @@ def check_multiple_devices_or_any_cpu_nodes(
 
 def check_lowering_disable_cudagraph(
     device_node_mapping: dict[torch.device, torch.fx.Node],
-) -> Optional[str]:
+) -> str | None:
     return check_multiple_devices_or_any_cpu_nodes(device_node_mapping)
 
 
@@ -217,9 +308,9 @@ def log_cudagraph_skip_and_bump_counter(msg: str) -> None:
 
 @dataclasses.dataclass
 class BoxedDeviceIndex:
-    value: Optional[int]
+    value: int | None
 
-    def set(self, device_idx: Optional[int]) -> None:
+    def set(self, device_idx: int | None) -> None:
         assert device_idx is None or isinstance(device_idx, int)
         self.value = device_idx
 
@@ -229,7 +320,7 @@ def check_for_mutation_ignore_cuda_graph_managed_tensor(
     mutated_inputs: OrderedSet[str],
     mutated_input_idxs: OrderedSet[int],
     static_input_idxs: Sequence[int],
-) -> Optional[str]:
+) -> str | None:
     default_msg = format_default_skip_message("mutated inputs")
 
     # doesn't work for non-trees because the warmup run would apply mutation twice
@@ -248,7 +339,7 @@ def check_for_mutation_ignore_cuda_graph_managed_tensor(
         return None if not has_mutation else default_msg
 
 
-def get_placeholder_stack_trace(placeholder: PlaceholderInfo) -> Optional[str]:
+def get_placeholder_stack_trace(placeholder: PlaceholderInfo) -> str | None:
     """
     Gets the first non-empty stack trace of a placeholder or its users.
     """
@@ -289,7 +380,7 @@ def __str__(self) -> str:
 def log_data_ptr_mismatch(
     placeholders: Sequence[PlaceholderInfo],
     inputs: list[InputType],
-    recorded_data_ptr: Sequence[Optional[int]],
+    recorded_data_ptr: Sequence[int | None],
     target_idxs: Sequence[int],
     mismatch: CheckInvariantStatus,
 ) -> str:
@@ -353,7 +444,7 @@ class CudagraphCachedInfo:
     """
 
     placeholders: Sequence[PlaceholderInfo]
-    stack_traces: list[Optional[str]]
+    stack_traces: list[str | None]
     cudagraph_fail_reasons: list[str]
 
 
@@ -366,7 +457,7 @@ class CudagraphMetadata:
     placeholders: Sequence[PlaceholderInfo]
     static_input_idxs: OrderedSet[int]
     mutated_input_idxs: OrderedSet[int]
-    stack_traces: list[Optional[str]]
+    stack_traces: list[str | None]
     constants: dict[str, torch.Tensor]
 
 
diff --git a/torch/_inductor/custom_graph_pass.py b/torch/_inductor/custom_graph_pass.py
index 5b83de361f7c0..31b842133dc86 100644
--- a/torch/_inductor/custom_graph_pass.py
+++ b/torch/_inductor/custom_graph_pass.py
@@ -2,7 +2,7 @@
 from abc import ABC, abstractmethod
 from collections.abc import Callable, Sequence
 from functools import lru_cache
-from typing import Any, Optional, TYPE_CHECKING, TypeAlias, Union
+from typing import Any, TYPE_CHECKING, TypeAlias
 
 import torch.fx.graph
 
@@ -50,7 +50,7 @@ def __call__(self, graph: torch.fx.graph.Graph) -> None:
         """
 
     @abstractmethod
-    def uuid(self) -> Optional[Any]:
+    def uuid(self) -> Any | None:
         """
         Return an ID to uniquely identify your custom pass implementation. Return None
         to skip inductor code caching entirely.
@@ -82,16 +82,29 @@ def __call__(self, gm: torch.fx.GraphModule) -> None:
         """
 
     @abstractmethod
-    def uuid(self) -> Optional[Any]:
+    def uuid(self) -> Any | None:
         """
         Return an ID to uniquely identify your custom pass implementation. Return None
         to skip inductor code caching entirely.
         """
 
 
-CustomGraphPassType: TypeAlias = Optional[
-    Union[CustomGraphPass, Callable[[torch.fx.graph.Graph], None]]
-]
+class CustomInferenceAwareGraphPass(CustomGraphPass):
+    """
+    Implement this interface for custom inference aware Graph passes.
+
+    """
+
+    @abstractmethod
+    def __call__(self, graph: torch.fx.graph.Graph, is_inference: bool) -> None:
+        """
+        Implementation of the custom pass.
+        """
+
+
+CustomGraphPassType: TypeAlias = (
+    CustomGraphPass | Callable[[torch.fx.graph.Graph], None] | None
+)
 
 
 @lru_cache(1)
@@ -152,14 +165,14 @@ def __call__(
         """
 
     @abstractmethod
-    def uuid(self) -> Optional[Any]:
+    def uuid(self) -> Any | None:
         """
         Return an ID to uniquely identify your custom partitioner implementation.
         Return None to skip inductor code caching entirely.
         """
 
 
-CustomPartitionerFnType: TypeAlias = Optional[CustomPartitionerFn]
+CustomPartitionerFnType: TypeAlias = CustomPartitionerFn | None
 
 
 class CustomRuntimeEstimator(ABC):
@@ -206,7 +219,7 @@ def __call__(self, node: "torch.fx.Node") -> float:
         """
 
     @abstractmethod
-    def uuid(self) -> Optional[Any]:
+    def uuid(self) -> Any | None:
         """
         Return an ID to uniquely identify your custom runtime estimator implementation.
         Return None to skip AOTAutograd caching entirely.
@@ -265,7 +278,7 @@ def __call__(
         """
 
     @abstractmethod
-    def uuid(self) -> Optional[Any]:
+    def uuid(self) -> Any | None:
         """
         Return an ID to uniquely identify your custom knapsack solver implementation.
         Return None to skip AOTAutograd caching entirely.
diff --git a/torch/_inductor/debug.py b/torch/_inductor/debug.py
index bd27e9c3330f1..3187cbe8c8023 100644
--- a/torch/_inductor/debug.py
+++ b/torch/_inductor/debug.py
@@ -15,7 +15,7 @@
 import tempfile
 import traceback
 from collections.abc import Callable, Iterator, Sequence
-from typing import Any, IO, Optional, Union
+from typing import Any, IO
 from unittest.mock import patch
 
 import torch
@@ -49,9 +49,9 @@
 log = logging.getLogger(__name__)
 
 # Graph execution tracking for debugging
-GRAPH_EXECUTION_ORDER: Optional[list[dict[str, object]]] = None
+GRAPH_EXECUTION_ORDER: list[dict[str, object]] | None = None
 RECORD_GRAPH_EXECUTION: bool = False
-GRAPH_COMPILE_IDS: Optional[dict[int, Optional[str]]] = None
+GRAPH_COMPILE_IDS: dict[int, str | None] | None = None
 
 ir_pre_fusion_log = getArtifactLogger(__name__, "ir_pre_fusion")
 ir_post_fusion_log = getArtifactLogger(__name__, "ir_post_fusion")
@@ -68,7 +68,7 @@ def has_dot() -> bool:
 def draw_buffers(
     nodes: list[BaseSchedulerNode],
     print_graph: bool = False,
-    fname: Optional[str] = None,
+    fname: str | None = None,
 ) -> None:
     """
     Draw a graph in fname.svg.
@@ -91,6 +91,16 @@ def draw_buffers(
                 group = (group[1],)
             else:
                 group = group[1]
+        elif isinstance(group, str):
+            # extern / template / nop nodes store a string like "extern"
+            # as the group instead of a shape tuple.  Extract the real
+            # output shape from the scheduler node's first output buffer.
+            try:
+                snode = node.meta["fusion_meta"].snode
+                size = snode.get_outputs()[0].node.maybe_get_size()
+                group = tuple(size) if size else ()
+            except Exception:
+                group = ()
 
         # gather meta data
         dtype = None
@@ -164,7 +174,7 @@ def func1(*args: Any) -> int:
             kwargs = {"device": snode.get_device()}
         fx_node = graph.call_function(node_func, args=(), kwargs=kwargs)  # type: ignore[arg-type]
 
-        def in_output(snode: Union[BaseSchedulerNode, FusedSchedulerNode]) -> bool:
+        def in_output(snode: BaseSchedulerNode | FusedSchedulerNode) -> bool:
             if isinstance(snode, FusedSchedulerNode):
                 return any(in_output(x) for x in snode.snodes)
             return any(
@@ -212,9 +222,9 @@ def in_output(snode: Union[BaseSchedulerNode, FusedSchedulerNode]) -> bool:
 
 
 def update_orig_fx_node_name_to_buf_name(
-    nodes: Optional[SchedulerNodeList],
+    nodes: SchedulerNodeList | None,
     node_name_to_buf_name: dict[str, str],
-    parent_buf_name: Optional[str] = None,
+    parent_buf_name: str | None = None,
     n_origins: int = 0,
 ) -> None:
     if nodes is None:
@@ -328,7 +338,7 @@ def enable_aot_logging() -> Iterator[None]:
 # _inductor_triton_kernel_to_post_grad_node_info's Debug Context
 _inductor_post_to_pre_grad_nodes: dict[str, dict[str, list[str]]] = {}
 _inductor_triton_kernel_to_post_grad_node_info: dict[str, list[str]] = {}
-_pre_grad_graph_id: Optional[int] = None
+_pre_grad_graph_id: int | None = None
 _inductor_pre_grad_node_stack_trace: dict[str, str] = {}
 _inductor_kernel_stack_trace: dict[str, list[str]] = {}
 _inductor_kernel_provenance_debug_handle: int = 0
@@ -394,7 +404,7 @@ class DebugContext:
     _counter = itertools.count()
 
     @staticmethod
-    def create_debug_dir(folder_name: str) -> Optional[str]:
+    def create_debug_dir(folder_name: str) -> str | None:
         debug_dir = config.trace.debug_dir or get_debug_dir()
         for n in DebugContext._counter:
             dirname = os.path.join(
@@ -507,9 +517,9 @@ def _setup_log_capture(
 
     def __exit__(
         self,
-        exc_type: Optional[type[BaseException]],
-        exc_val: Optional[BaseException],
-        exc_tb: Optional[Any],
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: Any | None,
     ) -> None:
         if self._prof:
             self._prof.disable()
@@ -531,7 +541,7 @@ def _save_profile_data(self) -> None:
             stats.sort_stats("tottime")
             stats.print_stats(100)
 
-    def __getattr__(self, name: str) -> Optional[Callable[..., None]]:
+    def __getattr__(self, name: str) -> Callable[..., None] | None:
         if config.trace.enabled and getattr(config.trace, name):
             try:
                 return getattr(DebugFormatter(self), name)
@@ -635,7 +645,7 @@ def log_autotuning_results(
         timings: dict["ChoiceCaller", float],  # type: ignore[name-defined] # noqa: F821
         elapse: float,
         precompile_elapse: float,
-        prescreening_elapse: Optional[float],
+        prescreening_elapse: float | None,
     ) -> None:
         from .ir import FixedLayout
 
@@ -729,7 +739,7 @@ def log_ir_post_fusion(nodes: SchedulerNodeList) -> None:
     V.debug.ir_post_fusion(nodes)
 
 
-def _dump_collective_schedule(schedule: list[Union[str, None]]) -> None:
+def _dump_collective_schedule(schedule: list[str | None]) -> None:
     try:
         trace_structured(
             "artifact",
@@ -764,10 +774,10 @@ def log_runtime_and_tensor_meta(node_runtimes: Sequence[tuple[Any, float]]) -> N
     try:
         to_optimization_hints = V.graph.sizevars.optimization_hints
 
-        def to_list(x: Optional[Sequence[Any]]) -> list[Any]:
+        def to_list(x: Sequence[Any] | None) -> list[Any]:
             return list(to_optimization_hints(x)) if x is not None else []
 
-        def dtype_to_str(dtype: Any) -> Optional[str]:
+        def dtype_to_str(dtype: Any) -> str | None:
             if dtype is None:
                 return None
             s = str(dtype)
@@ -865,7 +875,7 @@ class TensorMetadataHolder:
 
 
 def create_mapping_pre_post_grad_nodes(
-    pre_grad_graph_id: Optional[int],
+    pre_grad_graph_id: int | None,
     post_to_pre_grad_nodes_json: dict[str, Any],
 ) -> dict[str, dict[str, list[str]]]:
     """
@@ -1102,10 +1112,10 @@ def create_kernel_information_json() -> dict[str, dict[str, list[str]]]:
 
 
 def set_kernel_post_grad_provenance_tracing(
-    node_schedule: Union[Sequence[BaseSchedulerNode], ExternKernel],
+    node_schedule: Sequence[BaseSchedulerNode] | ExternKernel,
     kernel_name: str,
     is_extern: bool = False,
-) -> Optional[int]:
+) -> int | None:
     """
     Set the mapping between `kernel_name` and the post_grad nodes in `node_schedule`.
 
@@ -1255,7 +1265,7 @@ def aot_inductor_minifier_wrapper(
     exported_program: torch.export.ExportedProgram,
     *,
     inductor_configs: dict[str, Any],
-    package_path: Optional[FileLike] = None,
+    package_path: FileLike | None = None,
 ) -> str:
     from torch._dynamo.debug_utils import AccuracyError
     from torch._dynamo.repro.aoti import dump_to_minify
diff --git a/torch/_inductor/decomposition.py b/torch/_inductor/decomposition.py
index 8e246c631af77..5c90e4c070238 100644
--- a/torch/_inductor/decomposition.py
+++ b/torch/_inductor/decomposition.py
@@ -4,9 +4,8 @@
 import math
 import operator
 import sys
-import typing
 from collections.abc import Callable
-from typing import Any, Optional, TypeAlias, TypeVar, Union
+from typing import Any, TypeAlias, TypeVar
 from typing_extensions import ParamSpec
 
 import torch
@@ -33,11 +32,11 @@
 from torch._prims_common import (
     elementwise_dtypes,
     ELEMENTWISE_TYPE_PROMOTION_KIND,
+    suggest_memory_format,
     type_to_dtype,
 )
 from torch._refs import native_layer_norm as decomp_native_layer_norm
 from torch.fx.experimental.symbolic_shapes import guard_or_false, statically_known_true
-from torch.utils._ordered_set import OrderedSet
 
 from . import config, inductor_prims
 from .utils import (
@@ -50,9 +49,7 @@
 _T = TypeVar("_T")
 _P = ParamSpec("_P")
 
-_GenericOperator: TypeAlias = Union[
-    torch._ops.OperatorBase, torch._ops.OpOverloadPacket
-]
+_GenericOperator: TypeAlias = torch._ops.OperatorBase | torch._ops.OpOverloadPacket
 
 log = logging.getLogger(__name__)
 aten = torch.ops.aten
@@ -103,6 +100,7 @@
         aten.triu_indices,
         aten.unbind_copy.int,
         aten.upsample_bilinear2d.vec,
+        aten.hann_window,
         quantized.linear_dynamic_fp16_unpacked_weight,
         _quantized.wrapped_quantized_linear,
     ]
@@ -111,7 +109,7 @@
 
 # Remove unwanted decompositions included via the core ATen decompositions from
 # the Inductor decomp table.
-decomps_to_exclude: list[Union[torch._ops.OpOverload, torch._ops.OpOverloadPacket]] = [
+decomps_to_exclude: list[torch._ops.OpOverload | torch._ops.OpOverloadPacket] = [
     aten._unsafe_index,
     aten._unsafe_masked_index,
     aten._unsafe_masked_index_put_accumulate,
@@ -131,13 +129,24 @@
     aten.sum,  # inductor lowers this directly
     aten.unbind,  # inductor lowers this directly
     aten.baddbmm,  # upcasts to fp32, perf issue
+    # FMA ops - we have lowerings that use FMA to match eager CUDA behavior
+    aten.addcmul,
+    aten.addcmul_,
+    aten._foreach_addcmul.Scalar,
+    aten._foreach_addcmul_,
+    aten.addcdiv,
+    aten.addcdiv_,
+    aten._foreach_addcdiv.Scalar,
+    aten._foreach_addcdiv_,
+    aten.lerp,
+    aten.lerp_,
 ]
 
 remove_decompositions(decompositions, decomps_to_exclude)
 
 
 def register_decomposition(
-    ops: Union[_GenericOperator, list[_GenericOperator]],
+    ops: _GenericOperator | list[_GenericOperator],
 ) -> Callable[[Callable[_P, _T]], Callable[_P, _T]]:
     for op in ops if isinstance(ops, list) else [ops]:
         if op in decompositions:
@@ -145,6 +154,38 @@ def register_decomposition(
     return decomp.register_decomposition(ops, decompositions)
 
 
+@register_decomposition([aten.lerp.Scalar])
+def _lerp_scalar(start: torch.Tensor, end: torch.Tensor, weight: float) -> torch.Tensor:
+    # Decompose into sub + add(alpha=weight) so that the add lowering emits FMA,
+    # matching eager CUDA's dual-formula (see aten/src/ATen/native/Lerp.h).
+    # Convert end to start's memory format so the output preserves start's layout,
+    # matching eager TensorIterator behavior.
+    fmt = suggest_memory_format(start)
+    if fmt != torch.contiguous_format:
+        end = end.contiguous(memory_format=fmt)
+    diff = end - start
+    if weight >= 0.5 or weight <= -0.5:
+        return torch.add(end, diff, alpha=-(1.0 - weight))
+    return torch.add(start, diff, alpha=weight)
+
+
+@register_decomposition([aten.lerp.Tensor])
+def _lerp_tensor(
+    start: torch.Tensor, end: torch.Tensor, weight: torch.Tensor
+) -> torch.Tensor:
+    # Same dual-formula as foreach_lerp polyfill: decompose into sub + addcmul.
+    # Convert end to start's memory format so the output preserves start's layout.
+    fmt = suggest_memory_format(start)
+    if fmt != torch.contiguous_format:
+        end = end.contiguous(memory_format=fmt)
+    diff = end - start
+    mask = weight.abs() >= 0.5
+    neg_omw = -(1.0 - weight)
+    w = torch.where(mask, neg_omw, weight)
+    base = torch.where(mask, end, start)
+    return torch.addcmul(base, w, diff, value=1)
+
+
 @register_decomposition([aten.embedding_dense_backward])
 def _embedding_dense_backward(
     grad_output: torch.Tensor,
@@ -167,8 +208,8 @@ def _embedding_dense_backward(
 def _native_layer_norm(
     input: torch.Tensor,
     normalized_shape: utils.ShapeType,
-    weight: Optional[torch.Tensor],
-    bias: Optional[torch.Tensor],
+    weight: torch.Tensor | None,
+    bias: torch.Tensor | None,
     eps: float,
 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
     if input.is_mtia:
@@ -181,8 +222,8 @@ def _native_layer_norm(
 def sym_constrain_range_for_size(
     symbol: torch.SymInt,
     *,
-    min: Optional[torch.types.Number] = None,
-    max: Optional[torch.types.Number] = None,
+    min: torch.types.Number | None = None,
+    max: torch.types.Number | None = None,
 ) -> None:
     return
 
@@ -191,8 +232,8 @@ def sym_constrain_range_for_size(
 @pw_cast_for_opmath_non_tensor_args
 def clamp(
     x: torch.Tensor,
-    min: Optional[torch.types.Number] = None,
-    max: Optional[torch.types.Number] = None,
+    min: torch.types.Number | None = None,
+    max: torch.types.Number | None = None,
 ) -> torch.Tensor:
     if min is not None:
         x = x.clamp_min(min)
@@ -212,7 +253,7 @@ def silu(x: torch.Tensor) -> torch.Tensor:
 
 @register_decomposition([aten.full])
 def full(
-    size: list[Union[int, torch.SymInt]],
+    size: list[int | torch.SymInt],
     fill_value: torch.types.Number,
     **kwargs: Any,
 ) -> torch.Tensor:
@@ -247,7 +288,7 @@ def index_add(
 # cool with strides and everything goes to empty_strided)
 @register_decomposition([aten.empty_permuted.default])
 def empty_permuted(
-    size: list[Union[int, torch.SymInt]],
+    size: list[int | torch.SymInt],
     physical_layout: list[int],
     **kwargs: Any,
 ) -> torch.Tensor:
@@ -268,9 +309,9 @@ def convolution_backward(
     input: torch.Tensor,
     weight: torch.Tensor,
     bias_sizes: list[int],
-    stride: Union[int, list[int]],
-    padding: Union[int, list[int]],
-    dilation: Union[int, list[int]],
+    stride: int | list[int],
+    padding: int | list[int],
+    dilation: int | list[int],
     transposed: bool,
     output_padding: list[int],
     groups: int,
@@ -306,8 +347,15 @@ def round_dec(x: torch.Tensor, decimals: int = 0) -> torch.Tensor:
 def bmm(
     self: torch.Tensor,
     batch2: torch.Tensor,
-    out_dtype: Optional[torch.dtype] = None,
+    out_dtype: torch.dtype | None = None,
 ) -> torch.Tensor:
+    # Outer-product specialization: [B, M, 1] x [B, 1, N] -> [B, M, N].
+    # This avoids introducing a reduction and maps directly to broadcasted mul.
+    if statically_known_true(self.shape[2] == 1) or statically_known_true(
+        batch2.shape[1] == 1
+    ):
+        return (self * batch2).contiguous()
+
     # TODO: Re-enable for mps once our reductions are performant enough
     # (https://github.com/pytorch/pytorch/issues/150121)
     if config.coordinate_descent_tuning and self.device.type not in ["cpu", "mps"]:
@@ -333,10 +381,20 @@ def addmm(
     self: torch.Tensor,
     mat1: torch.Tensor,
     mat2: torch.Tensor,
-    out_dtype: Optional[torch.dtype] = None,
+    out_dtype: torch.dtype | None = None,
     beta: torch.types.Number = 1,
     alpha: torch.types.Number = 1,
 ) -> torch.Tensor:
+    if mat1.device.type not in ["cpu", "mps"]:
+        if (
+            statically_known_true(mat1.size(-1) == 1)
+            and statically_known_true(mat1.size(0) != 1)
+            and statically_known_true(mat2.size(1) != 1)
+        ):
+            counters["inductor"]["decompose_addmm"] += 1
+            out = mat1 * mat2
+            return alpha * out + beta * self
+
     if self.device.type == "cpu":
         if statically_known_true(mat1.size(0) == 1) and statically_known_true(
             mat2.size(-1) == 1
@@ -362,7 +420,7 @@ def addmm(
 def mm(
     self: torch.Tensor,
     input2: torch.Tensor,
-    out_dtype: Optional[torch.dtype] = None,
+    out_dtype: torch.dtype | None = None,
 ) -> torch.Tensor:
     # Our matrix vector multiplies only achieve peak bandwidth with coordinate descent tuning.
     # todo: Look into why and fix it (hopefully)
@@ -374,9 +432,14 @@ def mm(
             input2.shape[1] == 1
         ):
             return (self.unsqueeze(2) * input2.unsqueeze(0)).sum(dim=1)
-    if self.device.type == "cpu":
-        if (
-            statically_known_true(self.size(-1) == 1)
+    # Non-CPU/MPS: always decompose. CPU: only for small tensors.
+    if (
+        statically_known_true(self.size(-1) == 1)
+        and statically_known_true(self.size(0) != 1)
+        and statically_known_true(input2.size(1) != 1)
+    ):
+        if self.device.type not in ["cpu", "mps"] or (
+            self.device.type == "cpu"
             and statically_known_true(self.size(0) > 0)
             and statically_known_true(input2.size(0) == 1)
             and (self.dtype == input2.dtype)
@@ -384,6 +447,7 @@ def mm(
         ):
             counters["inductor"]["decompose_mm"] += 1
             return self * input2
+    if self.device.type == "cpu":
         if statically_known_true(self.size(0) == 1) and statically_known_true(
             input2.size(-1) == 1
         ):
@@ -485,7 +549,7 @@ def add(
     x: torch.Tensor,
     y: torch.Tensor,
     *,
-    alpha: Optional[torch.types.Number] = None,
+    alpha: torch.types.Number | None = None,
 ) -> torch.Tensor:
     # Require both x and y to be complex tensors.
     x_is_complex_tensor = torch.is_tensor(x) and x.is_complex()
@@ -573,7 +637,7 @@ def fmax(self: torch.Tensor, other: torch.Tensor) -> torch.Tensor:
 @register_decomposition(aten.amax)
 def amax(
     self: torch.Tensor,
-    dim: Optional[int] = None,
+    dim: int | None = None,
     keepdim: bool = False,
 ) -> torch.Tensor:
     if self.dtype == torch.bool:
@@ -584,7 +648,7 @@ def amax(
 @register_decomposition(aten.amin)
 def amin(
     self: torch.Tensor,
-    dim: Optional[int] = None,
+    dim: int | None = None,
     keepdim: bool = False,
 ) -> torch.Tensor:
     if self.dtype == torch.bool:
@@ -609,7 +673,7 @@ def narrow_copy(
 @register_decomposition([aten.view_copy.default])
 def view_copy_default(
     self: torch.Tensor,
-    size: list[Union[int, torch.SymInt]],
+    size: list[int | torch.SymInt],
 ) -> torch.Tensor:
     return aten.view(self, size).clone()
 
@@ -638,11 +702,11 @@ def _get_shape_permutation_like(
 @register_decomposition(aten.full_like)
 def full_like(
     self: torch.Tensor,
-    fill_value: Union[int, float],
+    fill_value: int | float,
     *,
-    dtype: Optional[torch.dtype] = None,
-    layout: Optional[torch.layout] = None,
-    device: Optional[torch.device] = None,
+    dtype: torch.dtype | None = None,
+    layout: torch.layout | None = None,
+    device: torch.device | None = None,
     pin_memory: bool = False,
     requires_grad: bool = False,
     memory_format: torch.memory_format = torch.preserve_format,
@@ -684,8 +748,8 @@ def _rand_like(
     rand_fn: Callable[..., torch.Tensor],
     self: torch.Tensor,
     *,
-    dtype: Optional[torch.dtype] = None,
-    device: Optional[torch.device] = None,
+    dtype: torch.dtype | None = None,
+    device: torch.device | None = None,
     memory_format: torch.memory_format = torch.preserve_format,
     **kwargs: Any,
 ) -> torch.Tensor:
@@ -712,22 +776,22 @@ def _rand_like(
     return result.permute(permutation).clone()
 
 
-@register_decomposition(aten.rand_like)
+@decomp.register_decomposition([aten.rand_like], extra_random_decomps)
 def rand_like(self: torch.Tensor, **kwargs: Any) -> torch.Tensor:
     return _rand_like(torch.rand, self, **kwargs)
 
 
-@register_decomposition(aten.randn_like)
+@decomp.register_decomposition([aten.randn_like], extra_random_decomps)
 def randn_like(self: torch.Tensor, **kwargs: Any) -> torch.Tensor:
     return _rand_like(torch.randn, self, **kwargs)
 
 
-@register_decomposition(aten.randint_like.default)
+@decomp.register_decomposition([aten.randint_like.default], extra_random_decomps)
 def randint_like(self: torch.Tensor, high: int, **kwargs: Any) -> torch.Tensor:
     return _rand_like(functools.partial(aten.randint.low, 0, high), self, **kwargs)
 
 
-@register_decomposition(aten.randint_like.low_dtype)
+@decomp.register_decomposition([aten.randint_like.low_dtype], extra_random_decomps)
 def randint_like_low(
     self: torch.Tensor, low: int, high: int, **kwargs: Any
 ) -> torch.Tensor:
@@ -737,7 +801,7 @@ def randint_like_low(
 @register_decomposition(aten.randint.default)
 def randint(
     high: int,
-    size: list[Union[int, torch.SymInt]],
+    size: list[int | torch.SymInt],
     **kwargs: Any,
 ) -> torch.Tensor:
     return aten.randint.low(0, high, size, **kwargs)
@@ -747,7 +811,7 @@ def randint(
 def linear_dynamic_fp16_unpacked_weight(
     input: torch.Tensor,
     weight: torch.Tensor,
-    bias: Optional[torch.Tensor] = None,
+    bias: torch.Tensor | None = None,
 ) -> torch.Tensor:
     packed_weight = torch.ops._quantized.wrapped_fbgemm_pack_gemm_matrix_fp16(weight)
     return torch.ops._quantized.wrapped_fbgemm_linear_fp16_weight(
@@ -826,32 +890,8 @@ def grid_sampler_2d(
     )
     return output
 
-
-# _foreach_addcmul.Scalar decomposition - uses mul+add instead of FMA
-# When emulate_precision_casts is enabled, we skip this decomposition
-# and use the inductor lowering which preserves FMA semantics
-@register_decomposition(aten._foreach_addcmul.Scalar)
-def _foreach_addcmul_scalar(
-    self: list[torch.Tensor],
-    left_tensors: list[torch.Tensor],
-    right_tensors: list[torch.Tensor],
-    scalar: float = 1,
-) -> list[torch.Tensor]:
-    return aten._foreach_add.List(
-        self, aten._foreach_mul.List(left_tensors, right_tensors), alpha=scalar
-    )
-
-
-@register_decomposition(aten._foreach_addcdiv.Scalar)
-def _foreach_addcdiv_scalar(
-    self: list[torch.Tensor],
-    left_tensors: list[torch.Tensor],
-    right_tensors: list[torch.Tensor],
-    scalar: float = 1,
-) -> list[torch.Tensor]:
-    return aten._foreach_add.List(
-        self, aten._foreach_div.List(left_tensors, right_tensors), alpha=scalar
-    )
+    # _foreach_addcmul and _foreach_addcdiv decompositions removed —
+    # the inductor lowering in lowering.py handles these with FMA.
 
 
 @register_decomposition(aten._foreach_lerp.Scalar)
@@ -887,9 +927,9 @@ def _foreach_lerp_scalarlist(
 def miopen_batch_norm(
     input: torch.Tensor,
     weight: torch.Tensor,
-    bias: typing.Optional[torch.Tensor],
-    running_mean: typing.Optional[torch.Tensor],
-    running_var: typing.Optional[torch.Tensor],
+    bias: torch.Tensor | None,
+    running_mean: torch.Tensor | None,
+    running_var: torch.Tensor | None,
     training: bool,
     exponential_average_factor: float,
     epsilon: float,
@@ -930,31 +970,6 @@ def select_decomp_table() -> dict[Any, Callable[..., Any]]:
         decompositions.pop(torch.ops.quantized.embedding_bag_byte_unpack.default, None)
         return decompositions
     result = fast_random_decomps()
-    if config.emulate_precision_casts:
-        # When emulating precision casts, skip decomposition of addcmul ops
-        # so that we use the inductor lowering which preserves FMA semantics.
-        # For _foreach_addcdiv, we use the native CUDA kernel.
-        # The decomposed version uses separate mul+add/div+add ops which don't match
-        # eager's FMA rounding behavior.
-        # Note: We check against OpOverloadPacket to match all overloads (default, out, etc.)
-        ops_to_skip = OrderedSet(
-            [
-                aten.addcmul,
-                aten._foreach_addcmul.Scalar,
-                aten._foreach_addcdiv.Scalar,
-            ]
-        )
-
-        def should_skip(op: Any) -> bool:
-            # Check if op is directly in the skip set
-            if op in ops_to_skip:
-                return True
-            # For OpOverload, also check if its OpOverloadPacket is in the skip set
-            if hasattr(op, "overloadpacket"):
-                return op.overloadpacket in ops_to_skip
-            return False
-
-        result = {k: v for k, v in result.items() if not should_skip(k)}
     return result
 
 
@@ -1097,9 +1112,9 @@ def index_reduce(
 def _max_pool_with_indices(
     x: torch.Tensor,
     kernel_size: list[int],
-    stride: Optional[Union[int, list[int]]],
-    padding: Union[int, list[int]],
-    dilation: Union[int, list[int]],
+    stride: int | list[int] | None,
+    padding: int | list[int],
+    dilation: int | list[int],
     ceil_mode: bool,
     dim: int,
 ) -> tuple[torch.Tensor, torch.Tensor]:
@@ -1154,9 +1169,9 @@ def _max_pool_with_indices(
 def max_pool2d_with_indices(
     x: torch.Tensor,
     kernel_size: list[int],
-    stride: Optional[Union[int, list[int]]] = None,
-    padding: Union[int, list[int]] = 0,
-    dilation: Union[int, list[int]] = 1,
+    stride: int | list[int] | None = None,
+    padding: int | list[int] = 0,
+    dilation: int | list[int] = 1,
     ceil_mode: bool = False,
 ) -> tuple[torch.Tensor, torch.Tensor]:
     return _max_pool_with_indices(
@@ -1168,9 +1183,9 @@ def max_pool2d_with_indices(
 def max_pool3d_with_indices(
     x: torch.Tensor,
     kernel_size: list[int],
-    stride: Optional[Union[int, list[int]]] = None,
-    padding: Union[int, list[int]] = 0,
-    dilation: Union[int, list[int]] = 1,
+    stride: int | list[int] | None = None,
+    padding: int | list[int] = 0,
+    dilation: int | list[int] = 1,
     ceil_mode: bool = False,
 ) -> tuple[torch.Tensor, torch.Tensor]:
     return _max_pool_with_indices(
@@ -1203,8 +1218,8 @@ def searchsorted_scalar(
     *,
     out_int32: bool = False,
     right: bool = False,
-    side: Optional[str] = None,
-    sorter: Optional[torch.Tensor] = None,
+    side: str | None = None,
+    sorter: torch.Tensor | None = None,
 ) -> torch.Tensor:
     return aten.searchsorted(
         sorted_sequence,
@@ -1239,7 +1254,7 @@ def rrelu_with_noise_functional(
     lower: float = 0.125,
     upper: float = 0.3333333333333333,
     training: bool = False,
-    generator: Optional[torch.Generator] = None,
+    generator: torch.Generator | None = None,
 ) -> tuple[torch.Tensor, torch.Tensor]:
     if training:
         not_positive = self <= 0
@@ -1255,7 +1270,7 @@ def rrelu_with_noise_functional(
 @register_decomposition(aten.repeat_interleave.Tensor)
 def repeat_interleave_Tensor(
     repeat: torch.Tensor,
-    output_size: Optional[int] = None,
+    output_size: int | None = None,
 ) -> torch.Tensor:
     if config.triton.autotune_at_compile_time:
         # We can't compile-time auto-tune this because
@@ -1279,7 +1294,7 @@ def repeat_interleave_Tensor(
 def conv1d_to_conv2d(
     input: torch.Tensor,
     weight: torch.Tensor,
-    bias: Optional[torch.Tensor] = None,
+    bias: torch.Tensor | None = None,
     stride: tuple[int] = (1,),
     padding: tuple[int] = (0,),
     dilation: tuple[int] = (1,),
diff --git a/torch/_inductor/dependencies.py b/torch/_inductor/dependencies.py
index 9194f661e1b60..893cf4b6a62e1 100644
--- a/torch/_inductor/dependencies.py
+++ b/torch/_inductor/dependencies.py
@@ -4,7 +4,7 @@
 import logging
 import re
 from collections.abc import Callable, Iterable, Sequence
-from typing import Any, Optional, TypeVar, Union
+from typing import Any, TypeVar
 from typing_extensions import Self
 from unittest.mock import patch
 
@@ -80,7 +80,7 @@ class MemoryDep(Dep):
     index: sympy.Expr
     var_names: tuple[sympy.Symbol, ...]
     size: tuple[sympy.Expr, ...]
-    mode: Optional[str] = None
+    mode: str | None = None
 
     def get_free_symbol_uses(
         self, unbacked_only: bool = False
@@ -101,7 +101,7 @@ def __repr__(self) -> str:
     def num_vars(self) -> int:
         return len(self.var_names)
 
-    def decide_loop_order_to_match(self, other: "MemoryDep") -> Optional[list[int]]:
+    def decide_loop_order_to_match(self, other: "MemoryDep") -> list[int] | None:
         """
         Can return None if not able to decide loop orders.
         """
@@ -320,7 +320,7 @@ def is_indirect(self) -> bool:
 class StarDep(Dep):
     # pyrefly: ignore [bad-override]
     name: str
-    mode: Optional[str] = None
+    mode: str | None = None
 
     # depends on the entire buffer
     @property
@@ -433,8 +433,8 @@ class ReadWrites:
     reads: OrderedSet[Dep]
     writes: OrderedSet[Dep]
     index_exprs: OrderedSet[IndexExprDep]
-    range_vars: Optional[list[sympy.Expr]] = None
-    var_ranges: Optional[VarRanges] = None
+    range_vars: list[sympy.Expr] | None = None
+    var_ranges: VarRanges | None = None
 
     def rename(self, renames: dict[str, str]) -> "ReadWrites":
         return ReadWrites(
@@ -445,7 +445,7 @@ def rename(self, renames: dict[str, str]) -> "ReadWrites":
             self.var_ranges,
         )
 
-    def with_read(self, dep: Union[Dep, OrderedSet[Dep]]) -> "ReadWrites":
+    def with_read(self, dep: Dep | OrderedSet[Dep]) -> "ReadWrites":
         assert isinstance(dep, (WeakDep, StarDep, OrderedSet))
         if not isinstance(dep, OrderedSet):
             dep = OrderedSet([dep])
@@ -517,7 +517,7 @@ def __init__(self, var_ranges: VarRanges, normalize: bool) -> None:
 
     @staticmethod
     def drop_unused_symbols(
-        index: Union[int, sympy.Expr],
+        index: int | sympy.Expr,
         var_names: list[sympy.Expr],
         sizes: list[sympy.Expr],
     ) -> None:
@@ -586,14 +586,14 @@ def load_seed(self, name: str, index: int) -> None:
         self.load(name, sympy.Integer(index))
 
     def store(
-        self, name: str, index: sympy.Expr, value: str, mode: Optional[str] = None
+        self, name: str, index: sympy.Expr, value: str, mode: str | None = None
     ) -> None:
         self._writes.add(MemoryDep(name, *self.canonicalize(index), mode=mode))
 
     def store_reduction(self, name: str, index: sympy.Expr, value: str) -> None:
         self.store(name, index, f"store_reduction({value})")
 
-    def index_expr(self, index: sympy.Expr, dtype: Optional[torch.dtype]) -> None:
+    def index_expr(self, index: sympy.Expr, dtype: torch.dtype | None) -> None:
         self._index_exprs.add(IndexExprDep(*self.canonicalize(index)))
 
     def bucketize(
@@ -603,8 +603,8 @@ def bucketize(
         boundary_indices: T,
         indexing_dtype: torch.dtype,
         right: bool,
-        sorter: Optional[tuple[str, sympy.Expr]] = None,
-        sorter_indices: Optional[T] = None,
+        sorter: tuple[str, sympy.Expr] | None = None,
+        sorter_indices: T | None = None,
     ) -> None:
         """Records the names of the buffers that bucketize will read from."""
         self._reads.add(StarDep(boundaries[0]))
@@ -748,7 +748,7 @@ def extract_loop_body_with_args(
 
 def extract_input_node_reduction_ranges(
     input_node: "torch._inductor.ir.IRNode",
-) -> tuple[Optional[list[sympy.Expr]], Optional[list[sympy.Expr]]]:
+) -> tuple[list[sympy.Expr] | None, list[sympy.Expr] | None]:
     """
     Returns the size and reduction size of all inputs, if the sizes and reduction_sizes (if exist) are all the same.
     It's possible that a node has multiple inputs, some are Reduction nodes and others are Pointwise nodes.
@@ -758,8 +758,8 @@ def extract_input_node_reduction_ranges(
 
     from .ir import ComputedBuffer, ExternKernel, Loops
 
-    size: Optional[list[sympy.Expr]]
-    reduction_size: Optional[list[sympy.Expr]]
+    size: list[sympy.Expr] | None
+    reduction_size: list[sympy.Expr] | None
 
     if isinstance(input_node.get_defining_op(), ComputedBuffer):
         # Input node has already been realized. Return its size and reduction_size.
@@ -778,8 +778,8 @@ def extract_input_node_reduction_ranges(
     # The current method still uses reduction ranges from the dependent realized node, which is not ideal.
     # Is there a way to check whether there are permutations in between?
     reads = input_node.get_reads()
-    reduction_size: Optional[list[sympy.Expr]] = None
-    size: Optional[list[sympy.Expr]] = None
+    reduction_size: list[sympy.Expr] | None = None
+    size: list[sympy.Expr] | None = None
     while reduction_size is None and len(reads) > 0:
         seen: OrderedSet[str] = OrderedSet()
         new_reads: list[Dep] = []
@@ -833,7 +833,7 @@ def _default(self, name: str, args: tuple[Any, ...], kwargs: dict[str, Any]) ->
     def indirect_indexing(
         self,
         index_var: Any,
-        size: Union[int, sympy.Expr],
+        size: int | sympy.Expr,
         check: bool = True,
         wrap_neg: bool = True,
     ) -> sympy.Symbol:
@@ -859,8 +859,8 @@ def reduction(
         dtype: torch.dtype,
         src_dtype: torch.dtype,
         reduction_type: ReductionType,
-        value: Union[None, tuple[None, ...]],
-    ) -> Union[None, tuple[None, ...]]:
+        value: None | tuple[None, ...],
+    ) -> None | tuple[None, ...]:
         num_values = reduction_num_outputs(reduction_type)
         return (None,) * num_values if num_values > 1 else None
 
@@ -873,7 +873,7 @@ def masked(self, mask: Any, body: Callable[..., Any], other: Any) -> None:
 def extract_free_symbols(
     fn: Callable[..., Any],
     index: Sequence[sympy.Expr],
-    rindex: Optional[Sequence[sympy.Expr]] = None,
+    rindex: Sequence[sympy.Expr] | None = None,
     unbacked_only: bool = True,
 ) -> OrderedSet[sympy.Symbol]:
     from .ir import FlexibleLayout
@@ -888,3 +888,16 @@ def extract_free_symbols(
     ):
         fn(*args)
     return handler.symbols
+
+
+class SymbolUsageCollectorOpsHandler(DefaultHandler):
+    usages: OrderedSet[str]
+
+    def __init__(self, symbol: sympy.Symbol) -> None:
+        self.symbol = symbol
+        self.usages = OrderedSet()
+
+    def _default(self, name: str, args: tuple[Any, ...], kwargs: dict[str, Any]) -> Any:
+        used_here = self.symbol in args or self.symbol in kwargs.values()
+        if used_here:
+            self.usages.add(name)
diff --git a/torch/_inductor/distributed_autotune.py b/torch/_inductor/distributed_autotune.py
index 4826f2852ec20..d066244a86344 100644
--- a/torch/_inductor/distributed_autotune.py
+++ b/torch/_inductor/distributed_autotune.py
@@ -2,7 +2,7 @@
 
 import contextlib
 import dataclasses
-from typing import Any, TYPE_CHECKING, Union
+from typing import Any, TYPE_CHECKING
 from unittest.mock import patch
 
 import sympy
@@ -175,7 +175,7 @@ def autotune(self, ser_choice: _SerializedChoice) -> TensorBox:
             kernel_inputs = MMKernelInputs([*self.original_inputs])
             assert isinstance(self.layout, Layout)
             choice = ser_choice.get_choice(self.layout, kernel_inputs)
-            buffer = autotune_select_algorithm(
+            buffer, _ = autotune_select_algorithm(
                 self._kernel_name,
                 [choice],
                 kernel_inputs.nodes(),
@@ -253,7 +253,7 @@ def get_choice(self, layout: Layout, inputs: KernelInputs) -> ChoiceCaller | Non
         return ktc.choice
 
     @staticmethod
-    def _compute_kwargs(description: str) -> dict[str, Union[int, str, bool]]:
+    def _compute_kwargs(description: str) -> dict[str, int | str | bool]:
         """
         Given a template description turn it into input kwargs.
         """
@@ -262,7 +262,7 @@ def _compute_kwargs(description: str) -> dict[str, Union[int, str, bool]]:
 
         # TODO: It seems like it would be better if the template could provide
         # this directly instead of having to parse a string.
-        kwargs: dict[str, Union[int, str, bool]] = {}
+        kwargs: dict[str, int | str | bool] = {}
         for cfg in description.split(","):
             key, val = cfg.split("=", 1)
             key, val = key.strip(), val.strip()
diff --git a/torch/_inductor/dtype_propagation.py b/torch/_inductor/dtype_propagation.py
index 4aae72e08a6c9..312ab76341e3c 100644
--- a/torch/_inductor/dtype_propagation.py
+++ b/torch/_inductor/dtype_propagation.py
@@ -1,12 +1,16 @@
 # mypy: allow-untyped-defs
 import functools
 from collections.abc import Callable, Sequence
-from typing import Any, Optional, Protocol, TYPE_CHECKING, TypeVar, Union
+from typing import Any, Optional, Protocol, TYPE_CHECKING, TypeVar
 
 import sympy
 
 import torch
-from torch._prims_common import ELEMENTWISE_TYPE_PROMOTION_KIND, type_to_dtype
+from torch._prims_common import (
+    ELEMENTWISE_TYPE_PROMOTION_KIND,
+    is_integer_dtype,
+    type_to_dtype,
+)
 from torch.utils._ordered_set import OrderedSet
 
 from .ops_handler import OP_NAMES, OpsHandler
@@ -15,6 +19,15 @@
 
 
 T = TypeVar("T")
+_MISSING_SHAPE = object()
+_UNSIGNED_INT_DTYPES: frozenset[torch.dtype] = frozenset(
+    (
+        torch.uint8,
+        torch.uint16,
+        torch.uint32,
+        torch.uint64,
+    )
+)
 
 
 class DTypeVar(Protocol):
@@ -22,7 +35,7 @@ class DTypeVar(Protocol):
     def dtype(self) -> torch.dtype: ...
 
 
-DTypeArg = Union[DTypeVar, torch.types.Number, str, OpsValue]
+DTypeArg = DTypeVar | torch.types.Number | str | OpsValue
 
 
 # Inputs need to be cacheable (e.g., not a CSEVar) in order for the cache to be effective
@@ -32,7 +45,7 @@ def dtype(self) -> torch.dtype: ...
 @functools.cache
 def get_promoted_dtype(
     *args: Sequence[tuple[torch.dtype, bool]],
-    type_promotion_kind: Optional[ELEMENTWISE_TYPE_PROMOTION_KIND] = None,
+    type_promotion_kind: ELEMENTWISE_TYPE_PROMOTION_KIND | None = None,
 ):
     def construct_input(inp):
         if inp[1]:
@@ -54,7 +67,7 @@ def construct_input(inp):
 
 def promote_types(
     args: Sequence[DTypeArg],
-    type_promotion_kind: Optional[ELEMENTWISE_TYPE_PROMOTION_KIND] = None,
+    type_promotion_kind: ELEMENTWISE_TYPE_PROMOTION_KIND | None = None,
 ):
     dtype_prop_candidates = []
 
@@ -79,6 +92,58 @@ def promote_types(
     return dtype
 
 
+def _unwrap_dtype_arg(arg: DTypeArg) -> DTypeVar | torch.types.Number | str:
+    if isinstance(arg, OpsValue):
+        return arg.value
+    return arg
+
+
+def _is_scalar_dtype_arg(arg: DTypeArg) -> bool:
+    arg = _unwrap_dtype_arg(arg)
+    if isinstance(arg, torch._prims_common.Number):
+        return True
+
+    is_scalar = getattr(arg, "is_scalar", False)
+    if callable(is_scalar):
+        if is_scalar():
+            return True
+    elif is_scalar:
+        return True
+
+    shape = getattr(arg, "shape", _MISSING_SHAPE)
+    if shape is _MISSING_SHAPE:
+        return False
+
+    if shape is None:
+        return True
+    if not isinstance(shape, Sequence):
+        return False
+
+    return len(shape) == 0
+
+
+def _has_known_nonnegative_scalar_int_value(arg: DTypeArg) -> bool:
+    arg = _unwrap_dtype_arg(arg)
+
+    if isinstance(arg, bool):
+        return True
+    if isinstance(arg, int):
+        return arg >= 0
+
+    dtype = getattr(arg, "dtype", None)
+    if dtype is None or not is_integer_dtype(dtype):
+        return False
+    if dtype in _UNSIGNED_INT_DTYPES:
+        return True
+
+    lower = getattr(getattr(arg, "bounds", None), "lower", None)
+    if lower is None:
+        return False
+    if isinstance(lower, sympy.Expr):
+        return lower.is_nonnegative is True
+    return lower >= 0
+
+
 class DtypePropagationOpsHandler:
     """
     Propagate dtype from args to output
@@ -188,7 +253,7 @@ def index_expr(expr: sympy.Expr, dtype: torch.dtype) -> torch.dtype:
     def to_dtype(
         x: DTypeArg,
         dtype: torch.dtype,
-        src_dtype: Optional[torch.dtype] = None,
+        src_dtype: torch.dtype | None = None,
         use_compute_types=True,
     ) -> torch.dtype:
         return upcast_compute_type(dtype) if use_compute_types else dtype
@@ -212,9 +277,24 @@ def truediv(a: DTypeArg, b: DTypeArg) -> torch.dtype:
         return promote_types([a, b])
 
     @staticmethod
-    def pow(a: DTypeArg, b: DTypeArg) -> torch.dtype:
+    def div_rn(a: DTypeArg, b: DTypeArg) -> torch.dtype:
         return promote_types([a, b])
 
+    @staticmethod
+    def pow(a: DTypeArg, b: DTypeArg) -> torch.dtype:
+        dtype = promote_types([a, b])
+        if (
+            is_integer_dtype(dtype)
+            and _is_scalar_dtype_arg(a)
+            and _is_scalar_dtype_arg(b)
+            and not _has_known_nonnegative_scalar_int_value(b)
+        ):
+            # Scalar integer pow follows Python semantics: negative exponents
+            # produce a floating result, even though tensor integer pow stays
+            # integral or errors out in separate lowering paths.
+            return torch.float64
+        return dtype
+
     @staticmethod
     def mod(a: DTypeArg, b: DTypeArg) -> torch.dtype:
         return promote_types([a, b])
@@ -233,6 +313,10 @@ def randn(seed: int, offset: int) -> torch.dtype:
     def rand(seed: int, offset: int) -> torch.dtype:
         return torch.float
 
+    @staticmethod
+    def rand_eager(seed, offset, threads_per_round, tid, vec) -> torch.dtype:
+        return torch.float
+
     @staticmethod
     def store_reduction(name: str, index, value: DTypeArg) -> None:
         return None
@@ -244,7 +328,7 @@ def reduction(
         return dtype
 
     @staticmethod
-    def store(name: str, index, value: DTypeArg, mode: Optional[str] = None) -> None:
+    def store(name: str, index, value: DTypeArg, mode: str | None = None) -> None:
         return None
 
     @staticmethod
@@ -321,8 +405,8 @@ def bucketize(
         boundary_indices: DTypeArg,
         indexing_dtype: torch.dtype,
         right: bool,
-        sorter: Optional[tuple[str, sympy.Expr]] = None,
-        sorter_indices: Optional[T] = None,
+        sorter: tuple[str, sympy.Expr] | None = None,
+        sorter_indices: T | None = None,
     ) -> torch.dtype:
         return indexing_dtype
 
@@ -364,7 +448,13 @@ def dot(x: DTypeArg, y: DTypeArg) -> torch.dtype:
 
     @staticmethod
     def inline_asm_elementwise(
-        *inputs, asm, constraints=None, dtype=torch.float32, is_pure=True, pack=1
+        *inputs,
+        asm,
+        constraints=None,
+        dtype=torch.float32,
+        is_pure=True,
+        pack=1,
+        input_dtypes=None,
     ):
         return dtype
 
diff --git a/torch/_inductor/exc.py b/torch/_inductor/exc.py
index 4d8ad9b7af04e..90fbb70a4740d 100644
--- a/torch/_inductor/exc.py
+++ b/torch/_inductor/exc.py
@@ -4,7 +4,7 @@
 import tempfile
 import textwrap
 from functools import lru_cache
-from typing import Any, Optional, TYPE_CHECKING
+from typing import Any, TYPE_CHECKING
 
 from torch._dynamo.exc import BackendCompilerFailed, ShortenTraceback
 
@@ -67,7 +67,7 @@ def __init__(
         target: Any,
         args: list[Any],
         kwargs: dict[str, Any],
-        stack_trace: Optional[str] = None,
+        stack_trace: str | None = None,
     ) -> None:
         msg = f"{type(exc).__name__}: {exc}\n{self.operator_str(target, args, kwargs)}"
         if stack_trace:
@@ -126,7 +126,7 @@ class CUDACompileError(CppCompileError):
 
 
 class TritonMissing(ShortenTraceback):
-    def __init__(self, first_useful_frame: Optional[types.FrameType]) -> None:
+    def __init__(self, first_useful_frame: types.FrameType | None) -> None:
         super().__init__(
             "Cannot find a working triton installation. "
             "Either the package is not installed or it is too old. "
@@ -140,7 +140,7 @@ def __init__(
         self,
         # pyrefly: ignore [not-a-type]
         device_props: _CudaDeviceProperties,
-        first_useful_frame: Optional[types.FrameType],
+        first_useful_frame: types.FrameType | None,
     ) -> None:
         super().__init__(
             f"Found {device_props.name} which is too old to be supported by the triton GPU compiler, "
@@ -156,7 +156,7 @@ class InductorError(BackendCompilerFailed):
     def __init__(
         self,
         inner_exception: Exception,
-        first_useful_frame: Optional[types.FrameType],
+        first_useful_frame: types.FrameType | None,
     ) -> None:
         self.inner_exception = inner_exception
         ShortenTraceback.__init__(
diff --git a/torch/_inductor/freezing.py b/torch/_inductor/freezing.py
index bffab2470865e..d8545a69bf0c7 100644
--- a/torch/_inductor/freezing.py
+++ b/torch/_inductor/freezing.py
@@ -4,7 +4,7 @@
 import itertools
 import logging
 import weakref
-from typing import Any, Optional
+from typing import Any
 
 import torch
 import torch.utils._pytree as pytree
@@ -145,7 +145,7 @@ class ErasedTensor(torch.Tensor):
     def __new__(cls, elem, name, owning_mod):
         return super().__new__(cls, elem.to(device="meta"))
 
-    def __init__(self, elem, name: Optional[str], mod) -> None:
+    def __init__(self, elem, name: str | None, mod) -> None:
         self.erased_name = name
         self.owning_mod_ref = weakref.ref(mod)
 
diff --git a/torch/_inductor/fuzzer.py b/torch/_inductor/fuzzer.py
index 53985ed3dca22..e9ba9f575e7b1 100644
--- a/torch/_inductor/fuzzer.py
+++ b/torch/_inductor/fuzzer.py
@@ -11,7 +11,7 @@
 from enum import Enum
 from functools import partial, wraps
 from types import FrameType
-from typing import Any, get_args, get_origin, Literal, Optional, TypeVar, Union
+from typing import Any, get_args, get_origin, Literal, TypeVar, Union
 
 import torch
 from functorch.compile import min_cut_rematerialization_partition
@@ -62,7 +62,7 @@ class DummyPass(CustomGraphPass):
     def __call__(self, graph: torch.fx.graph.Graph) -> None:
         return None
 
-    def uuid(self) -> Optional[Any]:
+    def uuid(self) -> Any | None:
         return None
 
 
@@ -76,7 +76,7 @@ def __call__(
     ) -> tuple[torch.fx.GraphModule, torch.fx.GraphModule]:
         return min_cut_rematerialization_partition(gm, joint_inputs, **kwargs)
 
-    def uuid(self) -> Optional[Any]:
+    def uuid(self) -> Any | None:
         return None
 
 
@@ -96,7 +96,7 @@ class TypeExemplars:
     }
 
     @staticmethod
-    def example(t: type[T]) -> Optional[T]:
+    def example(t: type[T]) -> T | None:
         """
         Return an example of a class.
         """
@@ -179,8 +179,18 @@ def failing(self) -> bool:
             "group_linear": {"require_fbgemm": True},
         },
     ],
-    "autoheuristic_collect": ["pad_mm", "mixed_mm"],
-    "autoheuristic_use": ["pad_mm", "mixed_mm"],
+    "autoheuristic_collect": [
+        {"pad_mm": True, "mixed_mm": True},
+        {"pad_mm": True, "mixed_mm": False},
+        {"pad_mm": False, "mixed_mm": True},
+        {"pad_mm": False, "mixed_mm": False},
+    ],
+    "autoheuristic_use": [
+        {"pad_mm": True, "mixed_mm": True},
+        {"pad_mm": True, "mixed_mm": False},
+        {"pad_mm": False, "mixed_mm": True},
+        {"pad_mm": False, "mixed_mm": False},
+    ],
     "traceable_tensor_subclasses": [OrderedSet()],
     "nontraceable_tensor_subclasses": [OrderedSet()],
 }
@@ -435,7 +445,7 @@ def set(self, combo: ComboType, status: Status) -> None:
         combo = tuple(sorted(combo))
         self._vals[combo] = status
 
-    def lookup(self, combo: ComboType) -> Optional[Status]:
+    def lookup(self, combo: ComboType) -> Status | None:
         combo = tuple(sorted(combo))
         return self._vals.get(combo, None)
 
@@ -509,6 +519,7 @@ def keys(self) -> KeysView[ComboType]:
         "pre_grad_custom_pass": DEFAULT,  # Typing
         "custom_partitioner_fn": DEFAULT,  # Typing
         "inductor_choices_class": DEFAULT,  # Typing
+        "cudagraph_policy": DEFAULT,  # Typing
     },
     "torch._dynamo.config": {
         "traceable_tensor_subclasses": DEFAULT,  # Typing
@@ -589,7 +600,7 @@ def __init__(
         config_module: ConfigModule,
         test_model_fn_factory: FactoryType,
         seed: int,
-        default: Optional[ConfigType] = None,
+        default: ConfigType | None = None,
         sm: SamplingMethod = SamplingMethod.TOGGLE,
         test_timeout: int = 3600,
     ):
@@ -719,7 +730,7 @@ def load_state(self, filename: str = "fuzzer_state.pkl") -> None:
             self.results = state["results"]
             self.detailed_results = state.get("detailed_results", {})
 
-    def timeout_handler(self, signum: int, frame: Optional[FrameType]) -> None:
+    def timeout_handler(self, signum: int, frame: FrameType | None) -> None:
         raise TimeoutError("Test execution timed out")
 
     def test_config(self, results: ResultType, config: ConfigType) -> Status:
@@ -749,7 +760,7 @@ def handle_return(
             message: str,
             return_status: Status,
             print_traceback: bool,
-            exc: Optional[Exception],
+            exc: Exception | None,
         ) -> Status:
             signal.signal(signal.SIGALRM, original_handler)
             print(f"{message} with config combination:")
@@ -843,12 +854,12 @@ def bisect(self, num_attempts: int = 100, p: float = 0.5) -> list[ConfigType]:
 
     def _bisect_failing_config(
         self, results: ResultType, failing_config: ConfigType
-    ) -> Optional[ConfigType]:
+    ) -> ConfigType | None:
         return self._bisect_failing_config_helper(results, list(failing_config.items()))
 
     def _bisect_failing_config_helper(
         self, results: ResultType, failing_config: list[tuple[str, Any]]
-    ) -> Optional[ConfigType]:
+    ) -> ConfigType | None:
         """
         Bisect a failing configuration to find minimal set of configs that cause failure.
 
diff --git a/torch/_inductor/fx_passes/auto_chunker/applier.py b/torch/_inductor/fx_passes/auto_chunker/applier.py
index 50e9e8273c677..37582297d6d76 100644
--- a/torch/_inductor/fx_passes/auto_chunker/applier.py
+++ b/torch/_inductor/fx_passes/auto_chunker/applier.py
@@ -1,7 +1,7 @@
 import copy
 import logging
 import operator
-from typing import Any, Optional
+from typing import Any
 
 import torch
 from torch import Tensor
@@ -62,14 +62,14 @@ def __init__(self, parent_gm: GraphModule, num_chunk: int):
         self.num_chunk = num_chunk
 
         # tangent node to the all-one tensor
-        self.overriden_tangent: dict[Node, Optional[Node]] = {}
+        self.overriden_tangent: dict[Node, Node | None] = {}
 
         self.subgraph_input: list[Node] = []
         self.subgraph_body: list[Node] = []
         self.subgraph_output: list[Node] = []
         self._categorize_subgraph_nodes()
 
-        self.chunk_sizes: Optional[list[int]] = None
+        self.chunk_sizes: list[int] | None = None
 
         # First index is node index,
         # Second index is chunk index.
@@ -78,7 +78,7 @@ def __init__(self, parent_gm: GraphModule, num_chunk: int):
         # check self.chunk_subgraph_input for more details
         self.chunked_subgraph_input: list[list[Node]] = []
 
-        self.accumulators: dict[Node, Optional[Node]] = {}
+        self.accumulators: dict[Node, Node | None] = {}
         self.chunks_for_recovering: dict[Node, list[Node]] = {}
         for node in self.subgraph_output:
             meta = get_chunking_meta(node)
@@ -251,6 +251,7 @@ def _create_placeholder_node(input_node: Node) -> Node:
                 and meta.chunk_dim is not None
             ):
                 shape = list(original_node.args[0])  # type: ignore[arg-type]
+                # pyrefly: ignore [unsupported-operation]
                 shape[meta.chunk_dim] = chunk_size
                 env[original_node] = new_graph.call_function(
                     aten.full.default,
@@ -258,15 +259,15 @@ def _create_placeholder_node(input_node: Node) -> Node:
                     original_node.kwargs,
                 )
                 continue
-            # Chunk aten.expand a scalar
+            # Chunk aten.expand: adjust the target shape at the chunk dimension
             if (
                 original_node.target == aten.expand.default
                 and isinstance(original_node.args[0], torch.fx.Node)
-                and original_node.args[0].meta["val"].numel() == 1
                 and (meta := get_chunking_meta(original_node)) is not None
                 and meta.chunk_dim is not None
             ):
                 shape = list(original_node.args[1])  # type: ignore[arg-type]
+                # pyrefly: ignore [unsupported-operation]
                 shape[meta.chunk_dim] = chunk_size
                 env[original_node] = new_graph.call_function(
                     aten.expand.default,
@@ -275,6 +276,23 @@ def _create_placeholder_node(input_node: Node) -> Node:
                 )
                 continue
 
+            # Chunk aten.view: adjust the target shape at the chunk dimension
+            if (
+                original_node.target == aten.view.default
+                and isinstance(original_node.args[0], torch.fx.Node)
+                and (meta := get_chunking_meta(original_node)) is not None
+                and meta.chunk_dim is not None
+            ):
+                shape = list(original_node.args[1])  # type: ignore[arg-type]
+                # pyrefly: ignore [unsupported-operation]
+                shape[meta.chunk_dim] = chunk_size
+                env[original_node] = new_graph.call_function(
+                    aten.view.default,
+                    (env[original_node.args[0]], shape),  # type: ignore[arg-type]
+                    original_node.kwargs,
+                )
+                continue
+
             # create the node with chunked inputs
             env[original_node] = new_graph.node_copy(original_node, lambda x: env[x])
 
diff --git a/torch/_inductor/fx_passes/auto_chunker/common.py b/torch/_inductor/fx_passes/auto_chunker/common.py
index bf810408bcbe2..b2a221c598f51 100644
--- a/torch/_inductor/fx_passes/auto_chunker/common.py
+++ b/torch/_inductor/fx_passes/auto_chunker/common.py
@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 import dataclasses
-from typing import Any, Optional, TYPE_CHECKING
+from typing import Any, TYPE_CHECKING
 
 
 if TYPE_CHECKING:
@@ -18,7 +18,7 @@ class ChunkingMeta:
     # tensor in the `scale_by` field. Need this since we pretend tangent to be 1 first and
     # scale the affected tensor later. Need propagate such information
     # to downstream.
-    scale_by: Optional[Node] = None
+    scale_by: Node | None = None
 
     # The dimension of the current tensor that get chunked.
     # Can be None if the current tensor is not chunked. E.g., when
@@ -26,7 +26,7 @@ class ChunkingMeta:
     #
     # To recover a tensor with non-None chunk_dim, we need concat each
     # chunk at the 'chunk_dim' dimension.
-    chunk_dim: Optional[int] = None
+    chunk_dim: int | None = None
 
     # The original tensor is the sum of each tensor in the chunking subgraph.
     # `chunk_dim` and `need_sum` are exclusive! A node can not have both a non
diff --git a/torch/_inductor/fx_passes/auto_chunker/core.py b/torch/_inductor/fx_passes/auto_chunker/core.py
index 634ebbed9097c..13622253aafd9 100644
--- a/torch/_inductor/fx_passes/auto_chunker/core.py
+++ b/torch/_inductor/fx_passes/auto_chunker/core.py
@@ -1,6 +1,6 @@
 import logging
 from collections.abc import Callable, Sequence
-from typing import Any, Optional
+from typing import Any
 
 import torch
 from torch._inductor import config
@@ -21,7 +21,7 @@
 
 
 def set_chunking_meta(
-    node: Node, meta: Optional[ChunkingMeta] = None, **kwargs: Any
+    node: Node, meta: ChunkingMeta | None = None, **kwargs: Any
 ) -> bool:
     """
     kwargs can override fields in the passed in `meta`
@@ -61,7 +61,7 @@ def update_chunking_meta(node: Node, **kwargs: Any) -> bool:
 def set_chunking_meta_if_none(
     nodes: Sequence[Node],
     meta: ChunkingMeta,
-    filter_for_nop: Optional[Callable[[Node], bool]] = None,
+    filter_for_nop: Callable[[Node], bool] | None = None,
 ) -> bool:
     """
     If filter_fop_nop returns true for a node, we set the chunking
@@ -88,7 +88,7 @@ def copy_chunking_meta(dst_node: Node, src_node: Node | ChunkingMeta) -> bool:
     return set_chunking_meta(dst_node, src_meta)
 
 
-def get_chunking_meta(node: Node) -> Optional[ChunkingMeta]:
+def get_chunking_meta(node: Node) -> ChunkingMeta | None:
     return node.meta.get("chunking")
 
 
@@ -114,7 +114,7 @@ def get_chunking_metas(
 )
 
 
-def find_amplifier_node(graph: Graph) -> Optional[Node]:
+def find_amplifier_node(graph: Graph) -> Node | None:
     r"""
     Find the 'amplifier' node which is a node that generates large
     output with small/medium input.
diff --git a/torch/_inductor/fx_passes/auto_chunker/propagate_scale_by.py b/torch/_inductor/fx_passes/auto_chunker/propagate_scale_by.py
index 0b75c37a5f808..22848193e52a4 100644
--- a/torch/_inductor/fx_passes/auto_chunker/propagate_scale_by.py
+++ b/torch/_inductor/fx_passes/auto_chunker/propagate_scale_by.py
@@ -131,15 +131,50 @@ def propagate_where(where_node: Node) -> bool:
     return False
 
 
+@register_propagate_rule(
+    [
+        aten.exp.default,
+        aten.log.default,
+        aten.tanh.default,
+        aten.eq.Tensor,
+    ]
+)
+def propagate_requires_no_scaling(out_node: Node) -> bool:
+    """
+    For nonlinear ops (exp, log, tanh) scale_by cannot be propagated
+    through since f(S*x) != S*f(x). For boolean-output ops (eq) scale_by
+    is meaningless. These ops only appear in the chunking subgraph when
+    scale_by is None (e.g. the final gradient is 1).
+    """
+    args_node = get_args_of_node_type(out_node)
+    args_meta = get_chunking_metas(args_node)
+    out_meta = get_chunking_meta(out_node)
+
+    scale_by = get_scale_by_from_metas(*args_meta)  # type: ignore[arg-type]
+    assert scale_by is None, (
+        f"Nonlinear op {out_node.target} requires scale_by=None, got {scale_by}"
+    )
+    assert out_meta is not None
+    out_meta.scale_by = None
+    return True
+
+
 @register_propagate_rule(
     [
         aten.mul.Tensor,
         prims.convert_element_type.default,
+        aten.neg.default,
         aten.sum.dim_IntList,
         aten.sum.default,  # sum to scalar
+        aten.amax.default,
         aten.mm.default,
         aten.permute.default,
         aten.expand.default,
+        aten.squeeze.dim,
+        aten.unsqueeze.default,
+        aten.gather.default,
+        aten.scatter_add.default,
+        aten.view.default,
     ]
 )
 def propagate_general_copy(out_node: Node) -> bool:
@@ -158,6 +193,16 @@ def propagate_general_copy(out_node: Node) -> bool:
     return True
 
 
+@register_propagate_rule(aten.scatter.value)
+def propagate_scatter_value(out_node: Node) -> bool:
+    # The backward of scatter.value always has value=0 (gradient of a constant),
+    # so S * scatter(x, idx, 0) = scatter(S*x, idx, 0) holds.
+    value = out_node.args[3]
+    if value != 0:
+        return False
+    return propagate_general_copy(out_node)
+
+
 @register_propagate_rule(
     [
         aten.add.Tensor,
diff --git a/torch/_inductor/fx_passes/auto_chunker/propagator.py b/torch/_inductor/fx_passes/auto_chunker/propagator.py
index 14e055b2004c4..6f5f3915cef51 100644
--- a/torch/_inductor/fx_passes/auto_chunker/propagator.py
+++ b/torch/_inductor/fx_passes/auto_chunker/propagator.py
@@ -1,5 +1,6 @@
 import functools
 import logging
+import math
 from collections.abc import Callable, Sequence
 from enum import Enum
 from queue import Queue
@@ -388,6 +389,7 @@ def bwd() -> PropagateStatus:
         prims.convert_element_type.default,
         aten.exp.default,
         aten.log.default,
+        aten.tanh.default,
         aten.add.Tensor,
         aten.sub.Tensor,
         aten.div.Tensor,
@@ -395,6 +397,7 @@ def bwd() -> PropagateStatus:
         prims.fma.default,
         aten.where.self,
         aten.neg.default,
+        aten.eq.Tensor,
     ]
 )
 def propagate_general_copy_metadata(
@@ -508,6 +511,7 @@ def propagate_bwd() -> PropagateStatus:
         aten.squeeze.dim,
         aten.gather.default,
         aten.scatter.value,
+        aten.scatter_add.default,
     ]
 )
 def propagate_general_copy_metadata_ignore_broadcast(out_node: Node) -> _HandlerRetType:
@@ -635,6 +639,118 @@ def bwd() -> PropagateStatus:
     return fwd(), bwd()
 
 
+@register_propagate_rule(aten.unsqueeze.default)
+def propagate_unsqueeze(unsqueeze_node: Node) -> _HandlerRetType:
+    input_node, unsqueeze_dim = unsqueeze_node.args[:2]
+    assert isinstance(input_node, Node)
+    assert isinstance(unsqueeze_dim, int)
+    input_ndim = get_fake_tensor_from_node_arg(input_node).ndim  # type: ignore[union-attr]
+    # Normalize negative dim: unsqueeze valid range is [-(ndim+1), ndim]
+    normalized_dim = (
+        unsqueeze_dim + input_ndim + 1 if unsqueeze_dim < 0 else unsqueeze_dim
+    )
+
+    def fwd() -> PropagateStatus:
+        assert isinstance(input_node, Node)
+        input_meta = get_chunking_meta(input_node)
+        if input_meta is None:
+            return _bool_to_status(False)
+        if input_meta.chunk_dim is None:
+            return _bool_to_status(copy_chunking_meta(unsqueeze_node, input_meta))
+
+        # pyrefly: ignore[unsupported-operation]
+        new_dim = input_meta.chunk_dim + (
+            1 if input_meta.chunk_dim >= normalized_dim else 0
+        )
+        return _bool_to_status(
+            set_chunking_meta(unsqueeze_node, meta=input_meta, chunk_dim=new_dim)
+        )
+
+    def bwd() -> PropagateStatus:
+        assert isinstance(input_node, Node)
+        output_meta = get_chunking_meta(unsqueeze_node)
+        if output_meta is None:
+            return _bool_to_status(False)
+        if output_meta.chunk_dim is None:
+            return _bool_to_status(copy_chunking_meta(input_node, output_meta))
+        # pyrefly: ignore[unsupported-operation]
+        new_dim = output_meta.chunk_dim - (
+            1 if output_meta.chunk_dim > normalized_dim else 0
+        )
+        return _bool_to_status(
+            set_chunking_meta(input_node, meta=output_meta, chunk_dim=new_dim)
+        )
+
+    return fwd(), bwd()
+
+
+def _find_chunk_dim_after_reshape(
+    old_shape: Sequence[int], new_shape: Sequence[int], chunk_dim: int
+) -> int | None:
+    """
+    Find the equivalent chunk_dim position after a reshape by matching
+    the prefix product (number of elements before the dimension) and
+    the dimension size. Returns None if the chunk dimension is merged
+    or split by the reshape, making it unsafe to propagate.
+
+    Examples:
+      [M, N] -> [M, N, 1], chunk_dim=0: returns 0 (trailing dim added)
+      [M]    -> [M, 1],     chunk_dim=0: returns 0
+      [M, N] -> [M1, M2, N] where M1*M2=M, chunk_dim=0: returns None (split)
+      [M, N] -> [M*N],      chunk_dim=0: returns None (merged)
+    """
+    chunk_size = old_shape[chunk_dim]
+    old_offset = math.prod(old_shape[:chunk_dim])
+    new_offset = 1
+    for new_dim in range(len(new_shape)):
+        if new_offset == old_offset and new_shape[new_dim] == chunk_size:
+            return new_dim
+        new_offset *= new_shape[new_dim]
+    return None
+
+
+@register_propagate_rule(aten.view.default)
+def propagate_view(view_node: Node) -> _HandlerRetType:
+    input_node = view_node.args[0]
+    assert isinstance(input_node, Node)
+    input_shape = list(get_fake_tensor_from_node_arg(input_node).shape)  # type: ignore[union-attr]
+    output_shape = list(get_fake_tensor_from_node_arg(view_node).shape)  # type: ignore[union-attr]
+
+    def fwd() -> PropagateStatus:
+        assert isinstance(input_node, Node)
+        input_meta = get_chunking_meta(input_node)
+        if input_meta is None:
+            return _bool_to_status(False)
+        if input_meta.chunk_dim is None:
+            return _bool_to_status(copy_chunking_meta(view_node, input_meta))
+        new_dim = _find_chunk_dim_after_reshape(
+            input_shape, output_shape, input_meta.chunk_dim
+        )
+        if new_dim is None:
+            return PropagateStatus.FAIL
+        return _bool_to_status(
+            set_chunking_meta(view_node, meta=input_meta, chunk_dim=new_dim)
+        )
+
+    def bwd() -> PropagateStatus:
+        assert isinstance(input_node, Node)
+        output_meta = get_chunking_meta(view_node)
+        if output_meta is None:
+            return _bool_to_status(False)
+        if output_meta.chunk_dim is None:
+            return _bool_to_status(copy_chunking_meta(input_node, output_meta))
+        new_dim = _find_chunk_dim_after_reshape(
+            output_shape, input_shape, output_meta.chunk_dim
+        )
+        if new_dim is None:
+            return PropagateStatus.FAIL
+        return _bool_to_status(
+            set_chunking_meta(input_node, meta=output_meta, chunk_dim=new_dim)
+        )
+
+    return fwd(), bwd()
+
+
 @register_propagate_rule(
     [
         aten.expand.default,
@@ -644,15 +760,50 @@ def propagate_expand(expand_node: Node) -> _HandlerRetType:
     input_node = expand_node.args[0]
     assert isinstance(input_node, Node)
 
-    if input_node.meta["val"].numel() != 1:
-        return PropagateStatus.FAIL
+    input_ft = get_fake_tensor_from_node_arg(input_node)
+    assert input_ft is not None
+    output_ft = get_fake_tensor_from_node_arg(expand_node)
+    assert output_ft is not None
+    input_shape = list(input_ft.shape)
+    output_shape = list(output_ft.shape)
 
-    # Combined fwd/bwd rule
-    output_meta = get_chunking_meta(expand_node)
-    if output_meta is None:
-        return _bool_to_status(False)
+    if input_ft.numel() == 1:
+        # Scalar input: combined fwd/bwd rule
+        output_meta = get_chunking_meta(expand_node)
+        if output_meta is None:
+            return _bool_to_status(False)
+        return _bool_to_status(set_chunking_meta(input_node))
+
+    # How many leading dims are added by expand
+    dim_offset = len(output_shape) - len(input_shape)
+
+    def is_expand_dim(out_dim: int) -> bool:
+        """Check if out_dim is a broadcast dimension (newly added or size 1 in input)."""
+        return out_dim < dim_offset or input_shape[out_dim - dim_offset] == 1
+
+    def fwd() -> PropagateStatus:
+        assert isinstance(input_node, Node)
+        input_meta = get_chunking_meta(input_node)
+        if input_meta is None:
+            return _bool_to_status(False)
+        # Fail if chunk_dim is an expand dimension (input size 1 broadcast to larger size)
+        if input_meta.chunk_dim is not None and is_expand_dim(
+            input_meta.chunk_dim + dim_offset
+        ):
+            return PropagateStatus.FAIL
+        return _bool_to_status(copy_chunking_meta(expand_node, input_meta))
 
-    return _bool_to_status(set_chunking_meta(input_node))
+    def bwd() -> PropagateStatus:
+        assert isinstance(input_node, Node)
+        output_meta = get_chunking_meta(expand_node)
+        if output_meta is None:
+            return _bool_to_status(False)
+        # Fail if chunk_dim is an expand dimension (input size 1 broadcast to larger size)
+        if output_meta.chunk_dim is not None and is_expand_dim(output_meta.chunk_dim):
+            return PropagateStatus.FAIL
+        return _bool_to_status(copy_chunking_meta(input_node, output_meta))
+
+    return fwd(), bwd()
 
 
 @register_propagate_rule(
diff --git a/torch/_inductor/fx_passes/auto_chunker/utils.py b/torch/_inductor/fx_passes/auto_chunker/utils.py
index 4da3b699dd661..d37b32e61e5df 100644
--- a/torch/_inductor/fx_passes/auto_chunker/utils.py
+++ b/torch/_inductor/fx_passes/auto_chunker/utils.py
@@ -1,5 +1,5 @@
 from collections.abc import Sequence
-from typing import Any, Optional
+from typing import Any
 
 import torch
 from torch.fx import Graph, GraphModule, Node
@@ -42,7 +42,7 @@ def compute_tensor_size(*args: Any, count_bytes: bool = True, **kwargs: Any) ->
 
 def get_fake_tensor_from_node_arg(
     node: torch.fx.node.Argument,
-) -> Optional[torch.Tensor]:
+) -> torch.Tensor | None:
     if (
         not hasattr(node, "meta")
         or ("val" not in node.meta)  # type: ignore[union-attr]
@@ -93,7 +93,7 @@ def has_any_chunking_meta(*node_list: Node) -> bool:
     return any(get_chunking_meta(node) for node in node_list)
 
 
-def get_first_chunking_meta(*node_list: Node) -> Optional[ChunkingMeta]:
+def get_first_chunking_meta(*node_list: Node) -> ChunkingMeta | None:
     """
     Get the first non-none chunking metadata if there is any.
     """
@@ -106,7 +106,7 @@ def get_first_chunking_meta(*node_list: Node) -> Optional[ChunkingMeta]:
     return None
 
 
-def get_scale_by_from_metas(*metas: ChunkingMeta) -> Optional[Node]:
+def get_scale_by_from_metas(*metas: ChunkingMeta) -> Node | None:
     """
     If there are multiple ChunkingMeta having the scale_by field,
     raise a CantChunk exception.
@@ -128,7 +128,7 @@ def get_scale_by_from_metas(*metas: ChunkingMeta) -> Optional[Node]:
     return scale_by_list[0] if len(scale_by_list) == 1 else None
 
 
-def get_scale_by_from_node(node: Node) -> Optional[Node]:
+def get_scale_by_from_node(node: Node) -> Node | None:
     from .core import get_chunking_meta
 
     meta = get_chunking_meta(node)
diff --git a/torch/_inductor/fx_passes/b2b_gemm.py b/torch/_inductor/fx_passes/b2b_gemm.py
index 6a7151e59e080..a2e203fcfa0b6 100644
--- a/torch/_inductor/fx_passes/b2b_gemm.py
+++ b/torch/_inductor/fx_passes/b2b_gemm.py
@@ -570,7 +570,8 @@ def tuned_b2b_gemm(
         )
     )
     # autotune
-    return autotune_select_algorithm("b2b_gemm", choices, [A, B, C], layout)
+    node, _ = autotune_select_algorithm("b2b_gemm", choices, [A, B, C], layout)
+    return node
 
 
 # match the inner mm of a potential b2b_gemm
diff --git a/torch/_inductor/fx_passes/bucketing.py b/torch/_inductor/fx_passes/bucketing.py
index d79bef415c4f7..a98793040a194 100644
--- a/torch/_inductor/fx_passes/bucketing.py
+++ b/torch/_inductor/fx_passes/bucketing.py
@@ -26,7 +26,15 @@
 
 overlap_log = torch._logging.getArtifactLogger(__name__, "overlap")
 
-BucketMode: TypeAlias = Literal["default", "custom_ops", "custom_ops_multidtype"]
+BucketMode: TypeAlias = Literal[
+    "default", "custom_ops", "custom_ops_multidtype", "coalesced"
+]
+
+
+def _default_bucket_mode() -> BucketMode:
+    from torch._inductor import config
+
+    return config.aten_distributed_optimizations.bucket_mode or "default"
 
 
 # Helper functions moved to top for better organization
@@ -71,13 +79,13 @@ def _compute_foreach_groups(
     Returns a flat list with -1 as group delimiter, or None if only one group exists.
     For example, groups [[0, 2], [1]] would be encoded as [0, 2, -1, 1].
     """
-    from torch.fx.experimental.symbolic_shapes import size_hint
+    from torch.fx.experimental.symbolic_shapes import guarding_hint_or_throw
 
     groups: defaultdict[tuple[torch.dtype, torch.dtype, tuple[int, ...]], list[int]] = (
         defaultdict(list)
     )
     for i, (ag_in, out_dtype) in enumerate(zip(ag_ins, out_dtypes)):
-        shape = tuple(size_hint(s) for s in ag_in.shape)
+        shape = tuple(guarding_hint_or_throw(s) for s in ag_in.shape)
         key = (ag_in.dtype, out_dtype, shape)
         groups[key].append(i)
 
@@ -94,23 +102,36 @@ def _compute_foreach_groups(
     return result
 
 
-def _schedulable_wait_node(node: torch.fx.Node) -> bool:
-    """
-    Add additional check on if the wait node is schedulable
-    We should not schedule a fx node that is:
-        1. wait on a collective that is not callable
-        2. wait on a non-NCCL communication node
+def _get_collective_node_from_wait(node: torch.fx.Node) -> torch.fx.Node | None:
+    """Given a wait node, return the collective it waits on.
+
+    Handles both standard (wait -> collective) and coalesced
+    (wait -> getitem -> coalesced_collective) patterns.
+    Returns None if the node is not a wait on a recognized NCCL collective.
     """
     if not is_wait_tensor(node):
-        return False
-    assert isinstance(node.args[0], torch.fx.Node)
-    if not isinstance(node.args[0].target, Callable):
-        return False
-    is_callable: bool = node.args[0].op == "call_function"
+        return None
+    arg = node.args[0]
+    assert isinstance(arg, torch.fx.Node)
+    if arg.op != "call_function":
+        return None
+    if arg.target is operator.getitem:
+        assert isinstance(arg.args[0], torch.fx.Node)
+        arg = arg.args[0]
+        if arg.op != "call_function":
+            return None
+    if not isinstance(arg.target, Callable):
+        return None
     # pyrefly: ignore [missing-attribute]
-    coll: NCCL_COLL = get_collective_type_from_kernel_name(node.args[0].target.name())
-    is_collective: bool = coll != NCCL_COLL.UNSUPPORTED
-    return is_callable and is_collective
+    coll: NCCL_COLL = get_collective_type_from_kernel_name(arg.target.name())
+    if coll == NCCL_COLL.UNSUPPORTED:
+        return None
+    return arg
+
+
+def _schedulable_wait_node(node: torch.fx.Node) -> bool:
+    """Check if this wait node is schedulable (waits on a recognized NCCL collective)."""
+    return _get_collective_node_from_wait(node) is not None
 
 
 def _populate_node_meta(
@@ -195,8 +216,9 @@ def bucket_cap_mb_by_bucket_idx_default(bucket_id: int) -> float:
 def bucket_all_gather(
     gm: torch.fx.GraphModule,
     bucket_cap_mb_by_bucket_idx: Callable[[int], float] | None = None,
-    mode: BucketMode = "default",
+    mode: BucketMode | None = None,
 ) -> None:
+    mode = mode or _default_bucket_mode()
     if bucket_cap_mb_by_bucket_idx is None:
         from torch._inductor.fx_passes.bucketing import (
             bucket_cap_mb_by_bucket_idx_default,  # pyrefly: ignore [missing-module-attribute]
@@ -212,8 +234,9 @@ def bucket_all_gather(
 def bucket_reduce_scatter(
     gm: torch.fx.GraphModule,
     bucket_cap_mb_by_bucket_idx: Callable[[int], float] | None = None,
-    mode: BucketMode = "default",
+    mode: BucketMode | None = None,
 ) -> None:
+    mode = mode or _default_bucket_mode()
     if bucket_cap_mb_by_bucket_idx is None:
         from torch._inductor.fx_passes.bucketing import (
             bucket_cap_mb_by_bucket_idx_default,  # pyrefly: ignore [missing-module-attribute]
@@ -275,7 +298,7 @@ def get_collective_type(node: torch.fx.Node) -> str:
 
 
 def get_full_bucket_key(
-    node: torch.fx.Node, bucket_mode: BucketMode
+    node: torch.fx.Node, bucket_mode: BucketMode | None
 ) -> tuple[str, Any]:
     """Get the full bucket key including collective type and bucket key."""
     return (get_collective_type(node), bucket_key(node, mode=bucket_mode))
@@ -453,7 +476,7 @@ def bucket_all_gather_by_mb(
     gm: torch.fx.GraphModule,
     bucket_cap_mb_by_bucket_idx: Callable[[int], float],
     filter_wait_node: Callable[[torch.fx.Node], bool] | None = None,
-    mode: BucketMode = "default",
+    mode: BucketMode | None = None,
 ) -> list[list[torch.fx.Node]]:
     """
     Identifies all all_gather nodes and groups them into buckets,
@@ -472,6 +495,7 @@ def bucket_all_gather_by_mb(
     Returns:
         list[list[torch.fx.Node]]: List of buckets, where each bucket is a list of all_gather nodes.
     """
+    mode = mode or _default_bucket_mode()
 
     group_key_fn = (
         _ag_group_key_multidtype if mode and "multidtype" in mode else _ag_group_key
@@ -490,7 +514,7 @@ def bucket_reduce_scatter_by_mb(
     gm: torch.fx.GraphModule,
     bucket_cap_mb_by_bucket_idx: Callable[[int], float],
     filter_wait_node: Callable[[torch.fx.Node], bool] | None = None,
-    mode: BucketMode = "default",
+    mode: BucketMode | None = None,
 ) -> list[list[torch.fx.Node]]:
     """
     Identifies all reduce_scatter nodes and groups them into buckets,
@@ -507,8 +531,9 @@ def bucket_reduce_scatter_by_mb(
     Returns:
         list[list[torch.fx.Node]]: List of buckets, where each bucket is a list of reduce_scatter nodes.
     """
+    mode = mode or _default_bucket_mode()
 
-    assert "multidtype" not in mode, (
+    assert mode is None or "multidtype" not in mode, (
         "reduce scatter bucketing does not support multidtype"
     )
 
@@ -624,6 +649,29 @@ def reduce_scatter_merge_fn_to_trace(
     return new_outs
 
 
+def reduce_scatter_merge_fn_coalesced(
+    rs_ins: list[torch.Tensor],
+    group_size: int,
+    group_name: str,
+    reduce_op: str,
+    reduce_dtype: torch.dtype,
+    device: torch.device,
+) -> list[torch.Tensor]:
+    """Bucketed RS via NCCL's coalesced API (ncclGroupStart/End).
+
+    Avoids cat-ing inputs into one buffer; instead passes the tensor list
+    directly to reduce_scatter_tensor_coalesced for zero-copy batching.
+    """
+    rs_ins_flat = [x.view(-1) for x in rs_ins]
+    new_out_sizes = [(x.shape[0] // group_size,) + x.shape[1:] for x in rs_ins]
+
+    rs_outs = torch.ops._c10d_functional.reduce_scatter_tensor_coalesced(
+        rs_ins_flat, reduce_op, group_size, group_name
+    )
+    rs_outs = [torch.ops.c10d_functional.wait_tensor(o) for o in rs_outs]
+    return [o.view(s) for o, s in zip(rs_outs, new_out_sizes)]
+
+
 def all_reduce_merge_fn_to_trace(
     ar_ins: list[torch.Tensor],
     group_name: str,
@@ -825,9 +873,11 @@ def all_gather_merge_fn_to_trace(
     device = ag_ins[0].device
     new_ag_out = torch.empty(ag_input_numel * group_size, dtype=dtype, device=device)
     new_ag_in = new_ag_out.narrow(0, ag_input_numel * rank, ag_input_numel)
-    foreach_copy_dsts = torch.split(new_ag_in, ins_split_sizes)
     ag_ins_flattened = [ag_in.reshape(-1) for ag_in in ag_ins]
-    torch._foreach_copy_(foreach_copy_dsts, ag_ins_flattened)
+    # Inductor fuses copy_(cat(...)) into 1 Triton kernel with no allocation for cat.
+    # _foreach_copy_(..., ag_ins_flattened) emits separate kernel per item,
+    # resulting in large number of small triton kernels to launch.
+    new_ag_in.copy_(torch.cat(ag_ins_flattened))
     wait_tensor = torch.ops.c10d_functional.wait_tensor(
         torch.ops._c10d_functional.all_gather_into_tensor_out.default(
             new_ag_in, group_size, group_name, out=new_ag_out
@@ -1066,10 +1116,11 @@ def process_collective_bucket(
 def merge_reduce_scatter_bucket(
     g: torch.fx.Graph,
     rs_nodes: list[torch.fx.Node],
-    mode: BucketMode = "default",
+    mode: BucketMode | None = None,
     insert_before: torch.fx.Node | None = None,
     wait_insertion_point: torch.fx.Node | None = None,
 ) -> tuple[list[torch.fx.Node], dict[torch.fx.Node, torch.fx.Node]]:
+    mode = mode or _default_bucket_mode()
     # Validate bucket consistency
     rs0 = rs_nodes[0]
     rs0_val = rs0.meta["val"]
@@ -1089,7 +1140,9 @@ def merge_reduce_scatter_bucket(
 
     # Choose merge function based on mode
     rs_merge_fn = reduce_scatter_merge_fn_to_trace
-    if mode and "custom_ops" in mode:
+    if mode == "coalesced":
+        rs_merge_fn = reduce_scatter_merge_fn_coalesced
+    elif mode and "custom_ops" in mode:
         rs_merge_fn = reduce_scatter_merge_fn_to_trace_custom_ops
 
     # Process bucket with lazy input collection
@@ -1159,10 +1212,11 @@ def create_trace_args(bucket_ins: list[torch.fx.Node]) -> tuple[Any, ...]:
 def merge_all_gather_bucket(
     g: torch.fx.Graph,
     ag_nodes: list[torch.fx.Node],
-    mode: BucketMode = "default",
+    mode: BucketMode | None = None,
     insert_before: torch.fx.Node | None = None,
     wait_insertion_point: torch.fx.Node | None = None,
 ) -> tuple[list[torch.fx.Node], dict[torch.fx.Node, torch.fx.Node]]:
+    mode = mode or _default_bucket_mode()
     from torch.distributed.distributed_c10d import _resolve_process_group
 
     ag0 = ag_nodes[0]
@@ -1178,7 +1232,9 @@ def merge_all_gather_bucket(
 
     # Choose merge function based on mode
     ag_merge_fn = all_gather_merge_fn_to_trace
-    if mode is not None and "custom_ops" in mode:
+    if mode == "coalesced":
+        logger.info("coalesced bucket_mode not supported for all_gather, using default")
+    elif mode and "custom_ops" in mode:
         ag_merge_fn = all_gather_merge_fn_to_trace_custom_ops  # type: ignore[assignment]
 
     # Process bucket with lazy input collection
@@ -1207,11 +1263,12 @@ def create_trace_args(bucket_ins: list[torch.fx.Node]) -> tuple[Any, ...]:
 def merge_reduce_scatter(
     gm: torch.fx.GraphModule,
     rs_buckets: list[list[torch.fx.Node]],
-    mode: BucketMode = "default",
+    mode: BucketMode | None = None,
 ) -> None:
     """
     Merges specified buckets of reduce_scatter to joint reduce_scatter.
     """
+    mode = mode or _default_bucket_mode()
     with dynamo_timed("fx.bucketing.merge_reduce_scatter", log_pt2_compile_event=True):
         trace_structured(
             "artifact",
@@ -1231,11 +1288,12 @@ def merge_reduce_scatter(
 def merge_all_gather(
     gm: torch.fx.GraphModule,
     ag_buckets: list[list[torch.fx.Node]],
-    mode: BucketMode = "default",
+    mode: BucketMode | None = None,
 ) -> None:
     """
     Merges specified buckets of all_gather to joint all_gather.
     """
+    mode = mode or _default_bucket_mode()
     with dynamo_timed("fx.bucketing.merge_all_gather", log_pt2_compile_event=True):
         trace_structured(
             "artifact",
diff --git a/torch/_inductor/fx_passes/control_dependencies.py b/torch/_inductor/fx_passes/control_dependencies.py
index 132a74fba3d10..d50a034665b2b 100644
--- a/torch/_inductor/fx_passes/control_dependencies.py
+++ b/torch/_inductor/fx_passes/control_dependencies.py
@@ -12,6 +12,7 @@
 
 import torch.fx as fx
 import torch.utils._pytree as pytree
+from torch._C import DispatchKey
 from torch._higher_order_ops.utils import register_fake
 from torch._ops import HigherOrderOperator
 from torch.utils._ordered_set import OrderedSet
@@ -56,6 +57,13 @@ def __call__(self, additional_deps, subgraph, *args, **kwargs):
 
 control_deps = ControlDeps()
 
+# control_deps wraps side-effecting ops (e.g. record_event, wait_event)
+# and must not be eliminated by DCE even when its outputs are unused.
+from torch.fx.node import has_side_effect
+
+
+has_side_effect(control_deps)
+
 
 # Register fake implementation for tracing
 @register_fake(control_deps)
@@ -64,6 +72,20 @@ def _(additional_deps, subgraph, *args, **kwargs):
     return subgraph(*args, **kwargs)
 
 
+# Register eager execution implementation
+@control_deps.py_impl(DispatchKey.CompositeExplicitAutograd)
+def control_deps_eager(additional_deps, subgraph, *args, **kwargs):
+    """Eager implementation - just execute the subgraph."""
+    return subgraph(*args, **kwargs)
+
+
+# Autograd impl needed because additional_deps tensors may have autograd state,
+# causing dispatch through AutogradCUDA even in post-autograd graphs.
+@control_deps.py_impl(DispatchKey.Autograd)
+def control_deps_autograd(additional_deps, subgraph, *args, **kwargs):
+    return subgraph(*args, **kwargs)
+
+
 def get_subgraph_name(gm: fx.GraphModule, name):
     name = f"subgraph_{name}"
 
@@ -185,9 +207,11 @@ def preserve_node_ordering(
         replacements[dependent_node] = ordered_node
 
 
-def _create_subgraph_for_node(graph: fx.Graph, node: fx.Node) -> fx.GraphModule:
+def _create_subgraph_for_node(
+    graph: fx.Graph, node: fx.Node, additional_deps=None
+) -> fx.GraphModule:
     """
-    Create a subgraph that exactly recreates a node's operation.
+    Create a subgraph that exactly recreates a node's operation optionally passing through additional dependencies.
 
     The subgraph takes only the fx.Node arguments and recreates the operation
     with the exact target, args structure, and kwargs.
@@ -195,6 +219,7 @@ def _create_subgraph_for_node(graph: fx.Graph, node: fx.Node) -> fx.GraphModule:
     Args:
         graph: The parent graph
         node: The node to wrap in a subgraph
+        additional_deps: Additional dependencies to pass through the subgraph
 
     Returns:
         A GraphModule containing the subgraph
@@ -222,6 +247,13 @@ def replace_nodes(item: Any) -> Any:
             return node_to_placeholder[item]
         return item
 
+    additional_deps_placeholders = []
+    for idx, dep in enumerate(additional_deps or ()):
+        placeholder = subgraph.placeholder(f"dep_{idx}")
+        if "val" in dep.meta:
+            placeholder.meta.update(dep.meta)
+        additional_deps_placeholders.append(placeholder)
+
     new_flat = [replace_nodes(item) for item in flat_args_kwargs]
     new_args, new_kwargs = pytree.tree_unflatten(new_flat, spec)
 
@@ -236,8 +268,13 @@ def replace_nodes(item: Any) -> Any:
     # Copy metadata from the original node
     result.meta.update(node.meta)
 
-    out = subgraph.output(result)
-    if "val" in result.meta:
-        out.meta["val"] = result.meta["val"]
+    if additional_deps_placeholders:
+        outputs = tuple([result] + additional_deps_placeholders)
+        out = subgraph.output(outputs)
+        out.meta["val"] = tuple(output.meta.get("val") for output in outputs)
+    else:
+        out = subgraph.output(result)
+        if "val" in result.meta:
+            out.meta["val"] = result.meta["val"]
 
     return fx.GraphModule(owning_module, subgraph)
diff --git a/torch/_inductor/fx_passes/freezing_patterns.py b/torch/_inductor/fx_passes/freezing_patterns.py
index 6e22a5896953b..48a940ae625fe 100644
--- a/torch/_inductor/fx_passes/freezing_patterns.py
+++ b/torch/_inductor/fx_passes/freezing_patterns.py
@@ -1,6 +1,5 @@
 # mypy: allow-untyped-defs
 import functools
-from typing import Optional
 
 import torch
 from torch._inductor.compile_fx import fake_tensor_prop
@@ -90,7 +89,7 @@ def freezing_passes(gm: torch.fx.GraphModule, aot_example_inputs):
 
 
 @init_once_fakemode
-def lazy_init(input_device: Optional[torch.device] = None):
+def lazy_init(input_device: torch.device | None = None):
     if torch._C._has_mkldnn and config.cpp.weight_prepack:
         from .mkldnn_fusion import _mkldnn_weight_pack_init
 
@@ -188,7 +187,7 @@ def check_concat_weights(match):
 
             if not all(
                 inp.op == "get_attr"
-                and inp.meta["val"].shape == inps[0].meta["val"].shape
+                and inp.meta["val"].shape[:-1] == inps[0].meta["val"].shape[:-1]
                 for inp in inps
             ):
                 return False
@@ -222,9 +221,10 @@ def matmul_fuse_pattern(inp, w1, w2, w3):
         return (inp @ w1, inp @ w2, inp @ w3)
 
     def matmul_replacement(inp, w1, w2, w3):
-        cat_t = torch.cat((w1, w2, w3), dim=1)
+        weights = (w1, w2, w3)
+        cat_t = torch.cat(weights, dim=1)
         mm = inp @ cat_t
-        return mm.chunk(3, dim=1)
+        return mm.split([w.size(1) for w in weights], dim=1)
 
     register_replacement(
         # pyrefly: ignore [bad-argument-type]
@@ -244,9 +244,10 @@ def matmul_fuse_pattern_two(inp, w1, w2):
         return (inp @ w1, inp @ w2)
 
     def matmul_replacement_two(inp, w1, w2):
-        cat_t = torch.cat((w1, w2), dim=1)
+        weights = (w1, w2)
+        cat_t = torch.cat(weights, dim=1)
         mm = inp @ cat_t
-        return mm.chunk(2, dim=1)
+        return mm.split([w.size(1) for w in weights], dim=1)
 
     register_replacement(
         # pyrefly: ignore [bad-argument-type]
@@ -270,9 +271,10 @@ def addmm_fuse_pattern_second(inp, w1, w2, w3, b1, b2, b3):
         )
 
     def addmm_fuse_replacement_second(inp, w1, w2, w3, b1, b2, b3):
-        cat_w = torch.cat((w1, w2, w3), dim=1)
+        weights = (w1, w2, w3)
+        cat_w = torch.cat(weights, dim=1)
         cat_b = torch.cat((b1, b2, b3))
-        return aten.addmm(cat_b, inp, cat_w).chunk(3, dim=1)
+        return aten.addmm(cat_b, inp, cat_w).split([w.size(1) for w in weights], dim=1)
 
     register_replacement(
         # pyrefly: ignore [bad-argument-type]
diff --git a/torch/_inductor/fx_passes/fsdp.py b/torch/_inductor/fx_passes/fsdp.py
index 1e71c350ed7b6..655d81621f3cb 100644
--- a/torch/_inductor/fx_passes/fsdp.py
+++ b/torch/_inductor/fx_passes/fsdp.py
@@ -6,6 +6,7 @@
     bucket_all_gather_by_mb,
     bucket_reduce_scatter_by_mb,
     BucketMode,
+    is_all_gather_into_tensor as is_all_gather,
     merge_all_gather,
     merge_reduce_scatter,
 )
@@ -19,19 +20,20 @@ def is_graph_input(node: torch.fx.Node) -> bool:
     return node.op == "placeholder"
 
 
+def is_fsdp_all_gather(n):
+    assert is_all_gather(n)
+    while len(n.all_input_nodes) == 1:
+        n = n.all_input_nodes[0]
+        if n.op == "placeholder":
+            return True
+    return False
+
+
 def is_fsdp_all_gather_wait(wait: torch.fx.Node) -> bool:
     # Assume all_gather_into_tensor input is either graph input
     # or dtype conversion of graph input
     ag_node = wait.args[0]  # type: ignore[arg-type, union-attr]
-    return (
-        is_graph_input(ag_node.args[0])  # type: ignore[arg-type, union-attr]
-        or (  # type: ignore[arg-type, union-attr]
-            ag_node.args[0].op == "call_function"  # type: ignore[arg-type, union-attr]
-            and ag_node.args[0].target  # type: ignore[arg-type, union-attr]
-            == torch.ops.prims.convert_element_type.default  # type: ignore[arg-type, union-attr]
-            and is_graph_input(ag_node.args[0].args[0])  # type: ignore[arg-type, union-attr]
-        )
-    )
+    return is_fsdp_all_gather(ag_node)
 
 
 def is_graph_output(node: torch.fx.Node) -> bool:
diff --git a/torch/_inductor/fx_passes/fuse_attention.py b/torch/_inductor/fx_passes/fuse_attention.py
index 621f54edf1db8..4bb2a99455572 100644
--- a/torch/_inductor/fx_passes/fuse_attention.py
+++ b/torch/_inductor/fx_passes/fuse_attention.py
@@ -2,9 +2,7 @@
 import functools
 import inspect
 import logging
-import math
 import warnings
-from typing import Optional
 
 import torch
 
@@ -109,15 +107,15 @@ def _sfdp_replacement_4(query, key, value, scale_factor, dropout_p):
     )
 
 
-def _sfdp_pattern_5(query, key, value, attn_mask):
+def _sfdp_pattern_5(query, key, value, attn_mask, inv_scale):
     attn_weight = torch.softmax(
-        (query @ key.transpose(-2, -1) / math.sqrt(query.size(-1))) + attn_mask, dim=-1
+        (query @ key.transpose(-2, -1) / (inv_scale)) + attn_mask, dim=-1
     )
     # attn_weight = torch.dropout(attn_weight, dropout_p)
     return attn_weight @ value
 
 
-def _sfdp_replacement_5(query, key, value, attn_mask):
+def _sfdp_replacement_5(query, key, value, attn_mask, inv_scale):
     counters["inductor"]["fuse_attention"] += 1
     return _scaled_dot_product_attention(
         query,
@@ -125,19 +123,20 @@ def _sfdp_replacement_5(query, key, value, attn_mask):
         value,
         attn_mask=attn_mask.to(dtype=query.dtype),
         dropout_p=0.0,
+        scale=1.0 / inv_scale,
         is_causal=False,
     )
 
 
-def _sfdp_pattern_6(query, key, value, attn_mask, dropout_p):
+def _sfdp_pattern_6(query, key, value, attn_mask, inv_scale, dropout_p):
     attn_weight = torch.softmax(
-        (query @ key.transpose(-2, -1) / math.sqrt(query.size(-1))) + attn_mask, dim=-1
+        (query @ key.transpose(-2, -1) / inv_scale) + attn_mask, dim=-1
     )
     attn_weight = torch.dropout(attn_weight, dropout_p, True)
     return attn_weight @ value
 
 
-def _sfdp_replacement_6(query, key, value, attn_mask, dropout_p):
+def _sfdp_replacement_6(query, key, value, attn_mask, inv_scale, dropout_p):
     counters["inductor"]["fuse_attention"] += 1
     return _scaled_dot_product_attention(
         query,
@@ -145,26 +144,28 @@ def _sfdp_replacement_6(query, key, value, attn_mask, dropout_p):
         value,
         attn_mask=attn_mask.to(dtype=query.dtype),
         dropout_p=dropout_p,
+        scale=1.0 / inv_scale,
         is_causal=False,
     )
 
 
-def _sfdp_pattern_7(query, key, value, dropout_p):
+def _sfdp_pattern_7(query, key, value, inv_scale, dropout_p):
     # in real workloads inputs to matmul are permuted
     # causing matmul to expand to a series of expand and clone calls
     # we want the same to happen during pattern tracing
     q = query.permute(0, 2, 1, 3)
     k = key.permute(0, 2, 1, 3)
     v = value.permute(0, 2, 1, 3)
-    div = q @ k.transpose(-2, -1) / math.sqrt(q.size(-1))
+    div = q @ k.transpose(-2, -1) / inv_scale
     div = div.to(torch.float32)
     attn_weight = torch.softmax(div, dim=-1)
     attn_weight = torch.dropout(attn_weight, dropout_p, True)
     attn_weight = attn_weight.to(torch.float16)
+    v = v.to(attn_weight.dtype)
     return attn_weight @ v
 
 
-def _sfdp_replacement_7(query, key, value, dropout_p):
+def _sfdp_replacement_7(query, key, value, inv_scale, dropout_p):
     # sdpa prefers inputs in permuted format
     # it makes a copy to put them in this format
     # if they aren't already
@@ -180,23 +181,25 @@ def _sfdp_replacement_7(query, key, value, dropout_p):
         v,
         attn_mask=None,  # attn_mask,
         dropout_p=dropout_p,
+        scale=1.0 / inv_scale,
         is_causal=False,
     )
 
 
-def _sfdp_pattern_8(query, key, value):
+def _sfdp_pattern_8(query, key, value, inv_scale):
     # no dropout version of pattern 7
     q = query.permute(0, 2, 1, 3)
     k = key.permute(0, 2, 1, 3)
     v = value.permute(0, 2, 1, 3)
-    div = q @ k.transpose(-2, -1) / math.sqrt(q.size(-1))
+    div = q @ k.transpose(-2, -1) / inv_scale
     div = div.to(torch.float32)
     attn_weight = torch.softmax(div, dim=-1)
     attn_weight = attn_weight.to(torch.float16)
+    v = v.to(attn_weight.dtype)
     return attn_weight @ v
 
 
-def _sfdp_replacement_8(query, key, value):
+def _sfdp_replacement_8(query, key, value, inv_scale):
     counters["inductor"]["fuse_attention"] += 1
     q = query.permute(0, 2, 1, 3)
     k = key.permute(0, 2, 1, 3)
@@ -207,24 +210,26 @@ def _sfdp_replacement_8(query, key, value):
         v,
         attn_mask=None,  # attn_mask,
         dropout_p=0.0,
+        scale=1.0 / inv_scale,
         is_causal=False,
     )
 
 
-def _sfdp_pattern_9(query, key, value, dropout_p):
+def _sfdp_pattern_9(query, key, value, inv_scale, dropout_p):
     q = query.permute(0, 2, 1, 3)
     k = key.permute(0, 2, 1, 3)
     v = value.permute(0, 2, 1, 3)
-    q = q / math.sqrt(q.size(-1))
+    q = q / inv_scale
     div = q @ k.transpose(-2, -1)
     div = div.to(torch.float32)
     attn_weight = torch.softmax(div, dim=-1)
     attn_weight = torch.dropout(attn_weight, dropout_p, True)
     attn_weight = attn_weight.to(torch.float16)
+    v = v.to(attn_weight.dtype)
     return attn_weight @ v
 
 
-def _sfdp_replacement_9(query, key, value, dropout_p):
+def _sfdp_replacement_9(query, key, value, inv_scale, dropout_p):
     counters["inductor"]["fuse_attention"] += 1
     q = query.permute(0, 2, 1, 3)
     k = key.permute(0, 2, 1, 3)
@@ -235,24 +240,26 @@ def _sfdp_replacement_9(query, key, value, dropout_p):
         v,
         attn_mask=None,  # attn_mask,
         dropout_p=dropout_p,
+        scale=1.0 / inv_scale,
         is_causal=False,
     )
 
 
-def _sfdp_pattern_10(query, key, value):
+def _sfdp_pattern_10(query, key, value, inv_scale):
     # no dropout version of 9
     q = query.permute(0, 2, 1, 3)
     k = key.permute(0, 2, 1, 3)
     v = value.permute(0, 2, 1, 3)
-    q = q / math.sqrt(q.size(-1))
+    q = q / inv_scale
     div = q @ k.transpose(-2, -1)
     div = div.to(torch.float32)
     attn_weight = torch.softmax(div, dim=-1)
     attn_weight = attn_weight.to(torch.float16)
+    v = v.to(attn_weight.dtype)
     return attn_weight @ v
 
 
-def _sfdp_replacement_10(query, key, value):
+def _sfdp_replacement_10(query, key, value, inv_scale):
     counters["inductor"]["fuse_attention"] += 1
     q = query.permute(0, 2, 1, 3)
     k = key.permute(0, 2, 1, 3)
@@ -263,6 +270,7 @@ def _sfdp_replacement_10(query, key, value):
         v,
         attn_mask=None,  # attn_mask,
         dropout_p=0.0,
+        scale=1.0 / inv_scale,
         is_causal=False,
     )
 
@@ -461,7 +469,7 @@ def _sfdp_replacement_17(query, key, value, attn_mask, inv_scale, dropout_p):
     )
 
 
-def _sfdp_pattern_18(query, key, value, causal_mask, dropout_p):
+def _sfdp_pattern_18(query, key, value, causal_mask, inv_scale, dropout_p):
     # for hf_GPT2 with dropout (introduces clone node) for inference
     # it also returns permuted key & value
     query = query.permute([0, 2, 1, 3])
@@ -470,7 +478,7 @@ def _sfdp_pattern_18(query, key, value, causal_mask, dropout_p):
     attn_weights = torch.matmul(query, key.permute(0, 1, 3, 2))
     inv_scale = torch.full(
         [],
-        value.size(-1) ** 0.5,
+        inv_scale,
         dtype=attn_weights.dtype,
         device=attn_weights.device,
     )
@@ -490,7 +498,7 @@ def _sfdp_pattern_18(query, key, value, causal_mask, dropout_p):
     )
 
 
-def _sfdp_replacement_18(query, key, value, causal_mask, dropout_p):
+def _sfdp_replacement_18(query, key, value, causal_mask, inv_scale, dropout_p):
     counters["inductor"]["fuse_attention"] += 1
     permuted_key = key.transpose(1, 2)
     permuted_value = value.transpose(1, 2)
@@ -502,19 +510,19 @@ def _sfdp_replacement_18(query, key, value, causal_mask, dropout_p):
             attn_mask=causal_mask,
             dropout_p=dropout_p,
             is_causal=False,
-            scale=1.0 / math.sqrt(value.size(-1)),
+            scale=1.0 / inv_scale,
         ),
         permuted_key,
         permuted_value,
     )
 
 
-def _sfdp_pattern_19(query, key, value, causal_mask, attn_mask, dropout_p):
+def _sfdp_pattern_19(query, key, value, causal_mask, attn_mask, inv_scale, dropout_p):
     # for token-classification+gpt2 / text-generation+gpt2
     attn_weights = torch.matmul(query, key.permute(0, 1, 3, 2))
     inv_scale = torch.full(
         [],
-        value.size(-1) ** 0.5,
+        inv_scale,
         dtype=attn_weights.dtype,
         device=attn_weights.device,
     )
@@ -528,7 +536,9 @@ def _sfdp_pattern_19(query, key, value, causal_mask, attn_mask, dropout_p):
     return torch.nn.functional.dropout(attn_weights, dropout_p).matmul(value)
 
 
-def _sfdp_replacement_19(query, key, value, causal_mask, attn_mask, dropout_p):
+def _sfdp_replacement_19(
+    query, key, value, causal_mask, attn_mask, inv_scale, dropout_p
+):
     counters["inductor"]["fuse_attention"] += 1
     fill_value = torch.full((), -float("inf"), dtype=query.dtype, device=query.device)
     attn_mask = torch.where(causal_mask, attn_mask, fill_value)
@@ -539,18 +549,18 @@ def _sfdp_replacement_19(query, key, value, causal_mask, attn_mask, dropout_p):
         attn_mask=attn_mask,
         dropout_p=dropout_p,
         is_causal=False,
-        scale=1.0 / math.sqrt(value.size(-1)),
+        scale=1.0 / inv_scale,
     )
 
 
-def _sfdp_pattern_20(query, key, value, attn_mask, dropout_p):
+def _sfdp_pattern_20(query, key, value, attn_mask, inv_scale, dropout_p):
     # for DistilBert with dropout transformers==4.44.2
     q = query.permute([0, 2, 1, 3])
     k = key.permute([0, 2, 1, 3])
     v = value.permute([0, 2, 1, 3])
     bs = q.size(0)
     k_len = k.size(-2)
-    q = q.div(math.sqrt(q.size(-1)))
+    q = q.div(inv_scale)
     scores = q @ k.transpose(-2, -1)
     fill_value = torch.full((), -float("inf"), dtype=query.dtype, device=query.device)
     attn_mask = (attn_mask == 0).view((bs, 1, 1, k_len)).expand_as(scores)
@@ -562,7 +572,7 @@ def _sfdp_pattern_20(query, key, value, attn_mask, dropout_p):
     )
 
 
-def _sfdp_replacement_20(query, key, value, attn_mask, dropout_p):
+def _sfdp_replacement_20(query, key, value, attn_mask, inv_scale, dropout_p):
     counters["inductor"]["fuse_attention"] += 1
     bs = query.size(0)
     n_head = query.size(2)
@@ -579,7 +589,7 @@ def _sfdp_replacement_20(query, key, value, attn_mask, dropout_p):
         attn_mask=attn_mask.to(dtype=torch.bool),
         dropout_p=dropout_p,
         is_causal=False,
-        scale=1.0 / math.sqrt(query.size(-1)),
+        scale=1.0 / inv_scale,
     )
 
 
@@ -812,6 +822,25 @@ def _sfdp_replacement_27(query, key, value, dropout_p):
     )
 
 
+def _sfdp_pattern_28(query, key, value, scale_factor, dropout_p):
+    # Visformer pattern
+    # same as pattern 4 but non-contiguous q/k/v
+    return _sfdp_pattern_4(query, key, value, scale_factor, dropout_p)
+
+
+def _sfdp_replacement_28(query, key, value, scale_factor, dropout_p):
+    counters["inductor"]["fuse_attention"] += 1
+    return _scaled_dot_product_attention(
+        query.contiguous(),
+        key.contiguous(),
+        value.contiguous(),
+        attn_mask=None,
+        dropout_p=dropout_p,
+        is_causal=False,
+        scale=scale_factor,
+    )
+
+
 @functools.lru_cache(None)
 def _warn_tf32_disabled() -> None:
     if (
@@ -917,7 +946,7 @@ def wrapper(*args, **kwargs):
     return wrapper
 
 
-def _get_sfdp_patterns(input_device: Optional[torch.device] = None):
+def _get_sfdp_patterns(input_device: torch.device | None = None):
     from .joint_graph import patterns
 
     if input_device:
@@ -936,6 +965,14 @@ def _get_sfdp_patterns(input_device: Optional[torch.device] = None):
     g_inp = functools.partial(
         torch.empty, (2, 4, 8, 16), device=device, requires_grad=True
     )
+    # non-contiguous input to cover more patterns.
+    gn_inp = functools.partial(
+        torch.empty_strided,
+        (2, 6, 16, 8),
+        (2304, 128, 1, 16),
+        device=device,
+        requires_grad=True,
+    )
     # attn_mask
     b_inp = functools.partial(torch.empty, (1, 1, 8, 8), device=device)
     m_inp = functools.partial(torch.empty, (2, 1, 1, 4), device=device)
@@ -946,6 +983,8 @@ def _get_sfdp_patterns(input_device: Optional[torch.device] = None):
     # workaround https://github.com/pytorch/pytorch/issues/97894
     # 0.113377 is a "magic" value that lets us recover the lost input arg relationship
     d = {"dropout_p": 0.113377}
+    s = {"inv_scale": 0.66666}
+    sd = {"inv_scale": 0.66666, "dropout_p": 0.113377}
 
     # we could also generate all these patterns in 3d.. TODO
     g_3d_inp = functools.partial(
@@ -964,6 +1003,7 @@ def _get_sfdp_patterns(input_device: Optional[torch.device] = None):
     # but will not in float, so we generate a pattern for both
     for dtype in [torch.float, torch.half]:
         g = functools.partial(g_inp, dtype=dtype)
+        gn = functools.partial(gn_inp, dtype=dtype)
         b = functools.partial(b_inp, dtype=dtype)
         b_float = functools.partial(b_inp, dtype=torch.float)
         b_bool = functools.partial(b_inp, dtype=torch.bool)
@@ -1011,42 +1051,42 @@ def _get_sfdp_patterns(input_device: Optional[torch.device] = None):
                 _sfdp_pattern_5,
                 _sfdp_replacement_5,
                 [g(), g(), g(), b()],
-                {},
+                s,
                 _sfdp_params_check,
             ),
             (
                 _sfdp_pattern_6,
                 _sfdp_replacement_6,
                 [g(), g(), g(), b()],
-                d,
+                sd,
                 _sfdp_params_check,
             ),
             (
                 _sfdp_pattern_7,
                 _sfdp_replacement_7,
                 [g(), g(), g()],
-                d,
+                sd,
                 _sfdp_params_check,
             ),
             (
                 _sfdp_pattern_8,
                 _sfdp_replacement_8,
                 [g(), g(), g()],
-                {},
+                s,
                 _sfdp_params_check,
             ),
             (
                 _sfdp_pattern_9,
                 _sfdp_replacement_9,
                 [g(), g(), g()],
-                d,
+                sd,
                 _sfdp_params_check,
             ),
             (
                 _sfdp_pattern_10,
                 _sfdp_replacement_10,
                 [g(), g(), g()],
-                {},
+                s,
                 _sfdp_params_check,
             ),
             (
@@ -1114,29 +1154,29 @@ def _get_sfdp_patterns(input_device: Optional[torch.device] = None):
                 _sfdp_pattern_18,
                 _sfdp_replacement_18,
                 [g(), g(), g(), m_bool()],
-                d,
+                sd,
                 _sfdp_params_check,
             ),
             (
                 _sfdp_pattern_18,
                 _sfdp_replacement_18,
                 [g_bs1(), g_bs1(), g_bs1(), m_bs1_bool()],
-                d,
+                sd,
                 _sfdp_params_check,
             ),
             (
                 _sfdp_pattern_19,
                 _sfdp_replacement_19,
                 [g(), g(), g(), b_bool(), b_float()],
-                d,
+                sd,
                 _sfdp_params_check,
             ),
             (
                 _sfdp_pattern_20,
                 _sfdp_replacement_20,
                 [g(), g(), g(), m_2d()],
-                d,
-                _sfdp_extra_check(aten.div.Tensor),
+                sd,
+                _sfdp_params_check,
             ),
             (
                 _sfdp_pattern_21,
@@ -1229,6 +1269,13 @@ def _get_sfdp_patterns(input_device: Optional[torch.device] = None):
                 d,
                 _sfdp_extra_check(disable_cuda=True),
             ),
+            (
+                _sfdp_pattern_28,
+                _sfdp_replacement_28,
+                [gn(), gn(), gn(), c()],
+                d,
+                _sfdp_extra_check(aten.mul.Tensor),
+            ),
         ]
         mask_fp32_patterns = ["pattern_16"]
         if dtype == torch.half:
@@ -1288,15 +1335,17 @@ def _get_sfdp_patterns(input_device: Optional[torch.device] = None):
                     "skip_duplicates": True,
                 },
             )
-
+            inference_workaround = {}
             if workaround:
-                assert len(workaround) == 1 and "dropout_p" in workaround
-                # functools.partial insufficient because we look at signature downstream
-                pattern = partialize_and_update_signature(pattern, dropout_p=0.0)
-                replacement = partialize_and_update_signature(
-                    replacement, dropout_p=0.0
-                )
-                workaround = {}
+                assert len(workaround) <= 2
+                if "inv_scale" in workaround:
+                    inference_workaround["inv_scale"] = workaround["inv_scale"]
+                if "dropout_p" in workaround:
+                    # functools.partial insufficient because we look at signature downstream
+                    pattern = partialize_and_update_signature(pattern, dropout_p=0.0)
+                    replacement = partialize_and_update_signature(
+                        replacement, dropout_p=0.0
+                    )
 
             inference_name = name + "_inference"
             yield (
@@ -1308,7 +1357,7 @@ def _get_sfdp_patterns(input_device: Optional[torch.device] = None):
                     "trace_fn": fwd_only,
                     "pass_dicts": patterns,
                     "extra_check": extra_check,
-                    "scalar_workaround": workaround,
+                    "scalar_workaround": inference_workaround,
                     # with dropout turned into clone, we end up with a number of
                     # semantically identical graphs
                     "skip_duplicates": True,
@@ -1317,6 +1366,6 @@ def _get_sfdp_patterns(input_device: Optional[torch.device] = None):
 
 
 @functools.cache
-def _sfdp_init(input_device: Optional[torch.device] = None):
+def _sfdp_init(input_device: torch.device | None = None):
     for key, register_replacement_kwargs in _get_sfdp_patterns(input_device):
         gen_register_replacement(key, **register_replacement_kwargs)
diff --git a/torch/_inductor/fx_passes/graph_view.py b/torch/_inductor/fx_passes/graph_view.py
index 69578552dfaf5..58927d7ad4e12 100644
--- a/torch/_inductor/fx_passes/graph_view.py
+++ b/torch/_inductor/fx_passes/graph_view.py
@@ -2,7 +2,7 @@
 
 import itertools
 import re
-from typing import Any, Optional, TYPE_CHECKING, Union
+from typing import Any, TYPE_CHECKING
 
 import torch.fx as fx  # noqa: TC001
 from torch.utils._ordered_set import OrderedSet
@@ -65,9 +65,7 @@ def add(self, data: fx.Node) -> None:
             self.data.append(data)
             self.unique_nodes.add(data)
 
-    def get_child(
-        self, module_stack: str, klass: Optional[type[Any]] = None
-    ) -> GraphView:
+    def get_child(self, module_stack: str, klass: type[Any] | None = None) -> GraphView:
         if module_stack not in self.children:
             new_stack = GraphView(module_stack, klass or self.klass)
             self.children[module_stack] = new_stack
@@ -112,7 +110,7 @@ def _is_root(stack: str) -> bool:
 def make_graph_view(
     graph: fx.Graph,
     module_stack_fn: Callable[[fx.Node], list[tuple[str, type[Any]]]] | None = None,
-) -> Optional[GraphView]:
+) -> GraphView | None:
     """
     Code from: https://github.com/meta-pytorch/autoparallel/pull/158
 
@@ -209,7 +207,7 @@ def nn_module_stack_meta(node: fx.Node) -> list[tuple[str, type[Any]]]:
 
 
 def get_subgraph_by_path(
-    graph_view: GraphView, paths: Union[str, list[str]]
+    graph_view: GraphView, paths: str | list[str]
 ) -> list[fx.Node]:
     """
     Get subgraph by path(s).
diff --git a/torch/_inductor/fx_passes/joint_graph.py b/torch/_inductor/fx_passes/joint_graph.py
index 8151ebf6196c2..16919f540f34b 100644
--- a/torch/_inductor/fx_passes/joint_graph.py
+++ b/torch/_inductor/fx_passes/joint_graph.py
@@ -6,7 +6,7 @@
 import typing
 from collections import Counter
 from collections.abc import Sequence
-from typing import Any, Optional
+from typing import Any
 
 import torch
 import torch._guards
@@ -56,7 +56,7 @@
 
 
 @init_once_fakemode
-def lazy_init(input_device: Optional[torch.device] = None):
+def lazy_init(input_device: torch.device | None = None):
     from .fuse_attention import _sfdp_init
     from .misc_patterns import _misc_patterns_init
     from .pad_mm import _pad_mm_init
@@ -71,8 +71,8 @@ def remove_no_ops(
     zeros: OrderedSet[torch.fx.Node],
     ones: OrderedSet[torch.fx.Node],
 ):
+    """Remove identity arithmetic operations: (+ 0, - 0, * 1, / 1)."""
     with torch.utils._python_dispatch._disable_current_modes():
-        "Removes no-ops: (+ 0, - 0, * 1, / 1)"
         graph = gm.graph
 
         def fake_tensors_eq(t1, t2, fields=("shape", "dtype", "device")):
@@ -83,6 +83,18 @@ def fake_tensors_eq(t1, t2, fields=("shape", "dtype", "device")):
                     return False
             return True
 
+        def is_mutated(n):
+            """Check if a node is mutated by any in-place operation."""
+            for user in n.users:
+                if user.op != "call_function" or not hasattr(user.target, "_schema"):
+                    continue
+                for i, arg in enumerate(user.args):
+                    if arg is n:
+                        schema_arg = user.target._schema.arguments[i]
+                        if schema_arg.alias_info and schema_arg.alias_info.is_write:
+                            return True
+            return False
+
         def replace_no_op(node, replace_input_index):
             replacement = node.args[replace_input_index]
 
@@ -92,6 +104,13 @@ def replace_no_op(node, replace_input_index):
             if not all(isinstance(arg, torch.fx.Node) for arg in node.args):
                 return
 
+            # https://github.com/pytorch/pytorch/issues/174187
+            # Don't replace if the replacement value is mutated in-place.
+            # The original node acts as an implicit copy; removing it would
+            # cause users to observe the post-mutation value instead.
+            if is_mutated(replacement):
+                return
+
             if not fake_tensors_eq(node.meta["val"], replacement.meta["val"]):
                 if fake_tensors_eq(
                     node.meta["val"],
@@ -343,12 +362,10 @@ def _deduce_value(self, node: torch.fx.Node):
         # handle before view ops because this changes value
         if node.target is aten.view.dtype:
             (input_tensor, output_dtype), kwargs = self.fetch_args_kwargs_from_env(node)
-            # view.dtype fails on 0-d tensors when element size changes
-            # (e.g., 0-d complex tensors can't be viewed as float)
-            if (
-                input_tensor.ndim == 0
-                and input_tensor.element_size() != output_dtype.itemsize
-            ):
+            # view.dtype with different element sizes changes element count
+            # (e.g., complex64 [1+0j] viewed as float32 becomes [1.0, 0.0]),
+            # making uniform values non-uniform. Also crashes on 0-d tensors.
+            if input_tensor.element_size() != output_dtype.itemsize:
                 return self.unknown_value
             return super(ConstantFolder, self).run_node(node)
 
@@ -600,7 +617,8 @@ def canonicalize_aten_ir_passes(gm: torch.fx.GraphModule):
 
 
 def joint_graph_passes(
-    graph: torch.fx.GraphModule, input_device: Optional[torch.device] = None
+    graph: torch.fx.GraphModule,
+    input_device: torch.device | None = None,
 ):
     """
     Run FX transformations on the joint forwards+backwards graph.
@@ -610,7 +628,7 @@ def joint_graph_passes(
         subsystem="joint_graph_passes",
     )
 
-    lazy_init(input_device)  # type: ignore[call-arg]
+    lazy_init(input_device)
     count = 0
 
     # must occur before other passes
diff --git a/torch/_inductor/fx_passes/memory_estimator.py b/torch/_inductor/fx_passes/memory_estimator.py
index ca723bb98bdd4..c939621679e09 100644
--- a/torch/_inductor/fx_passes/memory_estimator.py
+++ b/torch/_inductor/fx_passes/memory_estimator.py
@@ -6,7 +6,7 @@
 
 import torch
 import torch.fx as fx
-from torch.fx.experimental.symbolic_shapes import size_hint
+from torch.fx.experimental.symbolic_shapes import optimization_hint
 from torch.utils._ordered_set import OrderedSet
 from torch.utils._pytree import tree_map_only
 
@@ -144,9 +144,7 @@ def get_storages_last_used(
 
 
 def _size_of_default(num_bytes: int | torch.SymInt) -> int:
-    return size_hint(
-        num_bytes, fallback=torch._inductor.config.unbacked_symint_fallback
-    )
+    return optimization_hint(num_bytes)
 
 
 def device_filter(device: torch.device) -> bool:
@@ -387,9 +385,7 @@ def get_current_memory_bytes(self) -> int:
     def _get_storage_size(self, storage_key: StorageKey) -> int:
         """Get the size of a storage in bytes, handling symbolic shapes."""
         size_bytes = storage_key.storage.nbytes()
-        return size_hint(
-            size_bytes, fallback=torch._inductor.config.unbacked_symint_fallback
-        )
+        return optimization_hint(size_bytes)
 
     def _get_storages_freed_by_node(self, node: fx.Node) -> OrderedSet[StorageKey]:
         """Get storages that would be freed if we schedule this node."""
diff --git a/torch/_inductor/fx_passes/misc_patterns.py b/torch/_inductor/fx_passes/misc_patterns.py
index 9b84eabaaf2bf..4179e0441f36d 100644
--- a/torch/_inductor/fx_passes/misc_patterns.py
+++ b/torch/_inductor/fx_passes/misc_patterns.py
@@ -1,6 +1,5 @@
 # mypy: allow-untyped-defs
 import functools
-from typing import Optional
 
 import torch
 from torch._dynamo.utils import counters
@@ -14,7 +13,7 @@
 
 
 @functools.cache
-def _misc_patterns_init(input_device: Optional[torch.device] = None):
+def _misc_patterns_init(input_device: torch.device | None = None):
     from .joint_graph import patterns as joint_graph_patterns
     from .post_grad import pass_patterns as post_grad_patterns_all
 
diff --git a/torch/_inductor/fx_passes/mkldnn_fusion.py b/torch/_inductor/fx_passes/mkldnn_fusion.py
index d61ef27840fa5..c1f0a7284484e 100644
--- a/torch/_inductor/fx_passes/mkldnn_fusion.py
+++ b/torch/_inductor/fx_passes/mkldnn_fusion.py
@@ -6,7 +6,7 @@
 
 import torch
 from torch._dynamo.utils import counters
-from torch.fx.experimental.symbolic_shapes import has_free_symbols
+from torch.fx.experimental.symbolic_shapes import has_free_symbols, optimization_hint
 from torch.utils._ordered_set import OrderedSet
 
 from .. import ir, mkldnn_ir
@@ -104,7 +104,7 @@ def pack_linear_weight(
             # For bfloat16 dynamic shape path, using input size hint to pack weight for a better performance.
             packed_weight_inputs = (
                 transpose_weight_node,
-                batch_size.node.shape_env.size_hint(batch_size.node.expr)
+                optimization_hint(batch_size)
                 if has_free_symbols(batch_size)
                 else batch_size,
             )
diff --git a/torch/_inductor/fx_passes/node_runtime_estimation.py b/torch/_inductor/fx_passes/node_runtime_estimation.py
index 8389c5781b4d4..4d7201d8a8188 100644
--- a/torch/_inductor/fx_passes/node_runtime_estimation.py
+++ b/torch/_inductor/fx_passes/node_runtime_estimation.py
@@ -8,7 +8,7 @@
 import itertools
 import operator
 from functools import lru_cache
-from typing import Any, Optional
+from typing import Any
 
 import torch
 import torch.fx as fx
@@ -41,7 +41,7 @@ def _get_collective_key(coll_node: fx.Node) -> str:
     group_name = kwargs.get("group_name", None)
     group_size = kwargs.get("group_size", None)
 
-    tensor_bytes: Optional[int] = None
+    tensor_bytes: int | None = None
     success, args, kw = fx_utils.get_fake_args_kwargs(coll_node)
     if success:
 
@@ -93,7 +93,7 @@ def _get_collective_cache() -> dict[str, float]:
     return {}
 
 
-def get_cached_runtime(key: str) -> Optional[float]:
+def get_cached_runtime(key: str) -> float | None:
     """Get cached runtime from process-local cache."""
     return _get_collective_cache().get(key)
 
@@ -103,7 +103,7 @@ def set_cached_runtime(key: str, value: float) -> None:
     _get_collective_cache()[key] = value
 
 
-def get_hint(x: int | torch.SymInt) -> Optional[int]:
+def get_hint(x: int | torch.SymInt) -> int | None:
     if isinstance(x, int):
         return x
     assert isinstance(x, torch.SymInt)
@@ -226,7 +226,7 @@ def benchmark_collective_with_cuda_events_impl(
         return None, ""
 
     # Extract actual input size in BYTES (first tensor argument)
-    actual_bytes: Optional[int] = None
+    actual_bytes: int | None = None
 
     def extract_tensor_info(t: torch.Tensor) -> torch.Tensor:
         nonlocal actual_bytes
@@ -242,7 +242,8 @@ def extract_tensor_info(t: torch.Tensor) -> torch.Tensor:
 
             actual_bytes = total_elems * t.dtype.itemsize
         else:
-            raise RuntimeError(f"should only be one input tensor to collective {n}")
+            # out-variants (e.g. all_gather_into_tensor_out) can have multiple tensors
+            pass
         return t
 
     torch.utils._pytree.tree_map_only(torch.Tensor, extract_tensor_info, (args, kwargs))
@@ -353,9 +354,9 @@ def _log_graph_collective_benchmarks(gm: fx.GraphModule, artifact_name: str) ->
 
 def _log_collective_benchmarks(
     collective_nodes: list[fx.Node],
-    collective_keys: Optional[list[str]] = None,
-    benchmarked_medians: Optional[list[float]] = None,
-    world_size: Optional[int] = None,
+    collective_keys: list[str] | None = None,
+    benchmarked_medians: list[float] | None = None,
+    world_size: int | None = None,
     artifact_name: str = "fx_collectives_analytical_estimation",
 ) -> None:
     """Log collective estimations for tlparse. Includes benchmarks if provided."""
diff --git a/torch/_inductor/fx_passes/overlap_manual_scheduling.py b/torch/_inductor/fx_passes/overlap_manual_scheduling.py
index b3c0fa4e8ebcf..f8d4cd6197797 100644
--- a/torch/_inductor/fx_passes/overlap_manual_scheduling.py
+++ b/torch/_inductor/fx_passes/overlap_manual_scheduling.py
@@ -2,13 +2,15 @@
 
 import heapq
 from collections import Counter, defaultdict
-from typing import Any, Optional, TYPE_CHECKING
+from typing import Any, TYPE_CHECKING
 
-import torch
-import torch.fx as fx
+import torch  # noqa: TC001
+import torch.fx as fx  # noqa: TC001
 from torch._dynamo.graph_deduplication import _stable_topological_sort
 from torch._inductor.fx_passes.bucketing import (
+    _get_collective_node_from_wait,
     _schedulable_wait_node,
+    BucketMode,
     is_all_gather_into_tensor as is_all_gather,
     is_fsdp_all_gather,
     is_fsdp_reduce_scatter,
@@ -17,7 +19,7 @@
     merge_reduce_scatter_bucket,
 )
 from torch._inductor.fx_passes.overlap_preserving_bucketer import (
-    bucket_key,
+    get_full_bucket_key,
     OverlapPreservingBucketer,
 )
 from torch._inductor.fx_passes.overlap_scheduling import (
@@ -72,7 +74,7 @@ def _bucket_group(self, coll_nodes: list[fx.Node]) -> None:
                 coll_nodes,
                 wait_insertion_point=first_wait,
                 insert_before=next_node,
-                mode="custom_ops",
+                mode=self.bucket_mode,
             )
         elif is_reduce_scatter(first):
             new_nodes, replacements = merge_reduce_scatter_bucket(
@@ -80,7 +82,7 @@ def _bucket_group(self, coll_nodes: list[fx.Node]) -> None:
                 coll_nodes,
                 wait_insertion_point=first_wait,
                 insert_before=next_node,
-                mode="custom_ops",
+                mode=self.bucket_mode,
             )
         else:
             raise ValueError(
@@ -89,25 +91,31 @@ def _bucket_group(self, coll_nodes: list[fx.Node]) -> None:
 
         logger.debug(f"bucketing nodes: {coll_nodes} into {new_nodes}")  # noqa: G004
 
-        # Identify the new wait and start
-        new_waits = [n for n in new_nodes if _schedulable_wait_node(n)]
-        assert len(new_waits) == 1, f"Expected exactly one new wait, got {new_waits}"
-        new_wait = new_waits[0]
-        new_start = new_wait.args[0]
-        assert isinstance(new_start, fx.Node)
+        # Identify the new wait(s) and their collective start in a single pass
+        wait_to_start = {
+            n: start
+            for n in new_nodes
+            if (start := _get_collective_node_from_wait(n)) is not None
+        }
+        assert len(wait_to_start) >= 1, (
+            f"Expected at least one new wait, got none in {new_nodes}"
+        )
+        new_waits = list(wait_to_start)
+        new_start: fx.Node = wait_to_start[new_waits[0]]
+        # Use last wait as the canonical wait for scheduling (same node when len == 1)
+        new_wait = new_waits[-1]
 
-        # Set manual bucketing-specific metadata
-        # Note: Generic metadata (nn_module_stack, fwd_nn_module_stack, custom, stack_trace)
-        # is now preserved automatically by the bucketing functions in bucketing.py
+        # Set manual bucketing-specific metadata (generic metadata is preserved in bucketing.py)
         node_type = (
             "bucketed_all_gather" if is_all_gather(first) else "bucketed_reduce_scatter"
         )
+        wait_set = OrderedSet(new_waits)
         for n in new_nodes:
-            if n == new_wait:
-                node_type = node_type + "_wait"
-            n.meta["manual_bucket_node_type"] = node_type
-            if "wait" in node_type:
+            if n in wait_set:
+                n.meta["manual_bucket_node_type"] = node_type + "_wait"
                 self.node_to_wait_map[n] = new_wait
+            elif n is new_start:
+                n.meta["manual_bucket_node_type"] = node_type
 
     def manual_bucket_collectives(self, nodes: list[fx.Node]) -> None:
         """
@@ -124,7 +132,7 @@ def manual_bucket_collectives(self, nodes: list[fx.Node]) -> None:
                 or is_fsdp_reduce_scatter(node)
             ):
                 continue
-            key = bucket_key(node)
+            key = get_full_bucket_key(node, "custom_ops")
             if key is not None:
                 grouped_collectives[key].add(node)
 
@@ -143,7 +151,10 @@ def __init__(
         module_bucket_plans: list[list[str] | str],
         insert_overlap_deps: bool,
         module_stack_fn: Callable[[fx.Node], list[tuple[str, type[Any]]]] | None = None,
+        bucket_mode: BucketMode | None = None,
     ):
+        # Manual overlap historically used "custom_ops" mode for bucketing
+        bucket_mode = bucket_mode or "custom_ops"
         super().__init__(
             gm,
             max_in_flight_gb=0.0,
@@ -156,6 +167,7 @@ def __init__(
             collective_estimator="analytical",
             max_memory_increase_gb=None,
             max_memory_increase_ratio=None,
+            bucket_mode=bucket_mode,
         )
         self.module_bucket_plans = module_bucket_plans
         self.nodes_in_subgraph: list[list[fx.Node]] = []
@@ -164,6 +176,7 @@ def __init__(
             graph=self.graph,
             collective_info=self.collective_info,
             scheduled=OrderedSet(self.graph.nodes),
+            bucket_mode=bucket_mode,
         )
         self.insert_overlap_deps = insert_overlap_deps
 
@@ -270,7 +283,7 @@ def _manual_reorder_graph(self) -> None:
 
         self.scheduled = OrderedSet(reversed(list(self.scheduled)))
         picked_ag: list[fx.Node] = []
-        last_compute: Optional[fx.Node] = None
+        last_compute: fx.Node | None = None
 
         for node in self.scheduled:
             node_type = node.meta.get("manual_bucket_node_type", "")
@@ -352,6 +365,7 @@ def manual_overlap_bucketing(
     module_bucket_plans: list[list[str] | str],
     insert_overlap_deps: bool = False,
     module_stack_fn: Callable[[fx.Node], list[tuple[str, type[Any]]]] | None = None,
+    bucket_mode: BucketMode | None = None,
 ) -> torch.fx.GraphModule:
     """Schedule nodes based on user specifications in module_bucket_plans
     The manual overlapping consists of two steps:
@@ -370,10 +384,15 @@ def manual_overlap_bucketing(
 
             See the `module_stack_fn` parameter in `make_graph_view` (graph_view.py) for
             detailed documentation on signature, return format, and usage examples.
+        bucket_mode: Bucket mode for collective bucketing. None uses default.
     """
     # decode abbreviated FQNs to actual FQNs
     overlapped_gm = ManualOverlapScheduler(
-        gm, module_bucket_plans, insert_overlap_deps, module_stack_fn
+        gm,
+        module_bucket_plans,
+        insert_overlap_deps,
+        module_stack_fn,
+        bucket_mode=bucket_mode,
     ).run()
     overlapped_gm.recompile()
     return overlapped_gm
diff --git a/torch/_inductor/fx_passes/overlap_preserving_bucketer.py b/torch/_inductor/fx_passes/overlap_preserving_bucketer.py
index ee9c8446c6b4f..88a6c6d05f8fe 100644
--- a/torch/_inductor/fx_passes/overlap_preserving_bucketer.py
+++ b/torch/_inductor/fx_passes/overlap_preserving_bucketer.py
@@ -9,14 +9,16 @@
 from torch._dynamo.utils import counters
 from torch._inductor.augmented_graph_helper import AugmentedGraphHelper
 from torch._inductor.fx_passes.bucketing import (
+    _default_bucket_mode,
+    _get_collective_node_from_wait,
     _schedulable_wait_node,
-    bucket_key,
     BucketMode,
     get_full_bucket_key,
     has_mergeable_all_gather_convert_dtype,
     is_all_gather_into_tensor as is_all_gather,
     is_reduce_scatter_tensor as is_reduce_scatter,
 )
+from torch._inductor.fx_passes.fsdp import is_fsdp_all_gather
 from torch._inductor.fx_passes.overlap_scheduling import (
     CollBucket,
     CollectiveInfo,
@@ -138,10 +140,11 @@ def __init__(
         max_bucket_memory_gb: float = 1.0,
         max_coll_distance: int = 1000,
         insert_overlap_deps: bool = False,
-        bucket_mode: BucketMode = "custom_ops_multidtype",
-        bucket_exposed_first: bool = True,
         collective_bucketing: bool = True,
+        bucket_mode: BucketMode | None = None,
+        bucket_exposed_first: bool | None = None,
         region_of: dict[fx.Node, Any] | None = None,
+        bucket_only_internode_comms: bool = False,
     ):
         self.graph = graph
         self.collective_info = collective_info
@@ -151,7 +154,8 @@ def __init__(
         self.max_coll_distance = max_coll_distance
         self.insert_overlap_deps = insert_overlap_deps
         self.bucket_exposed_first = bucket_exposed_first
-        self.bucket_mode = bucket_mode
+        self.bucket_only_internode_comms = bucket_only_internode_comms
+        self.bucket_mode = bucket_mode or _default_bucket_mode()
         self.collective_bucketing = collective_bucketing
         self.region_of: dict[fx.Node, Any] = region_of or {}
         self.node_to_event: dict[fx.Node, PGEvent] = {}
@@ -162,7 +166,7 @@ def __init__(
         self.aug_graph = AugmentedGraphHelper(self.graph, self.node_ancestors)
 
         # Build timelines and add constraints to aug_graph
-        self.pg_to_timeline_head: dict[str, Optional[PGEvent]] = self.build_timelines()
+        self.pg_to_timeline_head: dict[str, PGEvent | None] = self.build_timelines()
         self._add_hiding_interval_constraints()
 
     def _compute_node_ancestors(self) -> dict[fx.Node, OrderedSet[fx.Node]]:
@@ -189,20 +193,20 @@ def _compute_node_ancestors(self) -> dict[fx.Node, OrderedSet[fx.Node]]:
 
         return node_ancestors
 
-    def build_timelines(self) -> dict[str, Optional[PGEvent]]:
+    def build_timelines(self) -> dict[str, PGEvent | None]:
         "Construct each process groups ordered series of event"
         all_pgs: OrderedSet[str] = OrderedSet()
         for start in self.collective_info:
             pg = get_group_name(start)
             all_pgs.add(pg)
 
-        pg_timeline: dict[str, Optional[PGEvent]] = {}
+        pg_timeline: dict[str, PGEvent | None] = {}
         for pg in all_pgs:
             pg_timeline[pg] = self.build_timeline(pg)
 
         return pg_timeline
 
-    def build_timeline(self, pg: str) -> Optional[PGEvent]:
+    def build_timeline(self, pg: str) -> PGEvent | None:
         """
         Build a timeline of important events (starts, waits, hiding compute) for this process group
         and constrain this ordering in the augmented graph.
@@ -225,8 +229,8 @@ def build_timeline(self, pg: str) -> Optional[PGEvent]:
                 node_type = "starts"
                 hiding_nodes |= self.collective_info[node].hiding_nodes
             elif _schedulable_wait_node(node):
-                wait_input = node.args[0]
-                if isinstance(wait_input, fx.Node) and get_group_name(wait_input) == pg:
+                wait_coll = _get_collective_node_from_wait(node)
+                if isinstance(wait_coll, fx.Node) and get_group_name(wait_coll) == pg:
                     node_type = "waits"
                 # Wait for a different PG but hiding a collective on this PG
                 elif node in hiding_nodes:
@@ -275,23 +279,42 @@ def _add_hiding_interval_constraints(self) -> None:
 
             self.all_hiding_nodes |= info.hiding_nodes
 
+    def identify_internode_group_names(self) -> OrderedSet[str]:
+        # Identify internode comm groups.
+        # Temporary uses FSDP pattern as heuristic for outer groups.
+        checked_pgs = OrderedSet()
+        internode_pgs = OrderedSet()
+        for start in self.collective_info:
+            pg = get_group_name(start)
+            if is_all_gather(start) and pg not in checked_pgs:
+                if is_fsdp_all_gather(start):
+                    internode_pgs.add(pg)
+                checked_pgs.add(pg)
+        return internode_pgs
+
     def _bucket_collectives_impl(self) -> list[CollBucket]:
         """Find and apply bucket transformations for collectives."""
         pg_collectives: dict[str, OrderedSet[fx.Node]] = defaultdict(OrderedSet)
+        internode_pgs = self.identify_internode_group_names()
+
         for start in self.collective_info:
             pg = get_group_name(start)
             pg_collectives[pg].add(start)
 
         all_buckets: list[CollBucket] = []
         for pg, collectives in pg_collectives.items():
+            if self.bucket_only_internode_comms and pg not in internode_pgs:
+                continue
+
+            # Populate node_to_event for this PG's timeline
             self._populate_node_to_event(pg)
 
             grouped_collectives: dict[object, OrderedSet[fx.Node]] = defaultdict(
                 OrderedSet
             )
             for start in collectives:
-                key = bucket_key(start, self.bucket_mode)
-                if key is not None:
+                key = get_full_bucket_key(start, self.bucket_mode)
+                if key[1] is not None:
                     grouped_collectives[key].add(start)
 
             for key, collective_group in grouped_collectives.items():
@@ -300,7 +323,7 @@ def _bucket_collectives_impl(self) -> list[CollBucket]:
                     key,
                     [n.name for n in collective_group],
                 )
-                buckets = self._find_buckets(collective_group)
+                buckets = self._find_buckets(collective_group, internode_pgs)
                 all_buckets.extend(buckets)
 
         for coll_bucket in all_buckets:
@@ -315,6 +338,10 @@ def _apply_deps_and_effect_tokens(self) -> None:
         """Apply topological sort and effect tokens to preserve overlap."""
         from torch._dynamo.graph_deduplication import _stable_topological_sort
 
+        # Clean up any remaining erased node references and cycles
+        self.aug_graph.remove_erased_extra_deps()
+        autofix = torch._inductor.config.aten_distributed_optimizations.overlap_scheduling_autofix_cycles
+        self.aug_graph.check_and_maybe_autofix_cyclic_extra_deps(autofix=autofix)
         additional_deps = self.aug_graph.get_all_extra_deps()
 
         for n, deps in additional_deps.items():
@@ -439,16 +466,35 @@ def _compute_overlap_ratio(self, node: fx.Node) -> float:
         hidden_time = info.estimated_time_ms - info.exposed_time_ms
         return hidden_time / info.estimated_time_ms
 
+    def _should_bucket_exposed_first(
+        self,
+        collective_group: OrderedSet[fx.Node],
+        current_pg: str,
+        internode_pgs: OrderedSet[str],
+    ) -> bool:
+        """Determine whether to bucket exposed collectives first."""
+        if self.bucket_exposed_first is None:
+            return current_pg in internode_pgs
+        return self.bucket_exposed_first
+
     def _find_buckets(
         self,
         collective_group: OrderedSet[fx.Node],
+        internode_pgs: OrderedSet[str],
     ) -> list[CollBucket]:
         """Find valid buckets within a group of similar collectives."""
         max_bucket_bytes = int(self.max_bucket_memory_gb * 1024 * 1024 * 1024)
         buckets = []
         processed: OrderedSet[fx.Node] = OrderedSet()
+        if len(collective_group) == 0:
+            return []
 
-        if self.bucket_exposed_first:
+        current_pg = get_group_name(next(iter(collective_group)))
+
+        bucket_exposed_first = self._should_bucket_exposed_first(
+            collective_group, current_pg, internode_pgs
+        )
+        if bucket_exposed_first:
             # Sort by overlap ratio (ascending) to bucket least hidden collectives first.
             # Exposed collectives benefit most from bucketing since their latency is on the
             # critical path. Prioritizing them also preserves hiding relationships for
@@ -498,7 +544,7 @@ def _find_buckets(
 
                 candidate_node_idx = self.node_idx[candidate]
                 if (
-                    self.bucket_exposed_first
+                    bucket_exposed_first
                     and abs(candidate_node_idx - start_node_idx)
                     > max_consecutive_failures
                 ):
@@ -528,7 +574,7 @@ def _ancestor_dep(self, n1: fx.Node, n2: fx.Node) -> bool:
 
     def _get_intervals(
         self, event: PGEvent
-    ) -> tuple[Optional[tuple[int, int]], list[tuple[int, int]]]:
+    ) -> tuple[tuple[int, int] | None, list[tuple[int, int]]]:
         """Get (execution_interval, hiding_intervals) for a collective event.
 
         Returns:
@@ -543,10 +589,9 @@ def _get_intervals(
             coll = event.node
         # For wait events, look up the start node from the event's args
         elif event.is_wait:
-            wait_input = event.node.args[0]
-            if not isinstance(wait_input, fx.Node):
+            coll = _get_collective_node_from_wait(event.node)
+            if coll is None:
                 return None, []
-            coll = wait_input
         else:
             return None, []
 
@@ -685,9 +730,7 @@ def enclosed_interval(inner: tuple[int, int], outer: tuple[int, int]) -> bool:
 
         return True
 
-    def remove_from_event(
-        self, node: fx.Node
-    ) -> tuple[Optional[PGEvent], Optional[PGEvent]]:
+    def remove_from_event(self, node: fx.Node) -> tuple[PGEvent | None, PGEvent | None]:
         """Remove node from timeline and return (prev_event, next_event)."""
         event = self.node_to_event[node]
         assert not event.is_compute, "Cannot remove compute events from timeline"
@@ -709,8 +752,8 @@ def remove_from_event(
     def restore_to_event(
         self,
         node: fx.Node,
-        prev_event: Optional[PGEvent],
-        next_event: Optional[PGEvent],
+        prev_event: PGEvent | None,
+        next_event: PGEvent | None,
     ) -> None:
         """Restore node to timeline after failed merge attempt."""
         event = self.node_to_event[node]
@@ -960,13 +1003,13 @@ def _apply_bucket(self, bucket_info: CollBucket) -> None:
                 self.graph,
                 bucket,
                 insert_before=next_node,
-                mode="custom_ops",
+                mode=self.bucket_mode,
             )
         elif is_all_reduce_tensor(bucket[0]):
             new_nodes, replacements = merge_all_reduce_bucket(
                 self.graph,
                 bucket,
-                mode="custom_ops",
+                mode=self.bucket_mode,
                 insert_before=next_node,
             )
         else:
@@ -975,39 +1018,51 @@ def _apply_bucket(self, bucket_info: CollBucket) -> None:
                 self.graph,
                 bucket,
                 insert_before=next_node,
-                mode="custom_ops",
+                mode=self.bucket_mode,
             )
 
-        # Get new nodes
-        new_waits = [n for n in new_nodes if _schedulable_wait_node(n)]
-        assert len(new_waits) == 1
-
-        new_wait = new_waits[0]
-        new_start = new_wait.args[0]
-        assert isinstance(new_start, fx.Node)
+        # Identify the new wait(s) and their collective start in a single pass
+        wait_to_start = {
+            n: start
+            for n in new_nodes
+            if (start := _get_collective_node_from_wait(n)) is not None
+        }
+        new_waits = list(wait_to_start)
 
         # Create mapping of all erased nodes to their replacements
         erased_to_new: dict[fx.Node, fx.Node | None] = {}
-        for old_start in old_starts:
-            erased_to_new[old_start] = new_start
-        for old_wait in old_waits:
-            erased_to_new[old_wait] = new_wait
+        new_start = wait_to_start[new_waits[0]]
+        if len(new_waits) == 1:
+            # Standard bucketing: single start + single wait
+            new_wait = new_waits[0]
+            for old_start in old_starts:
+                erased_to_new[old_start] = new_start
+            for old_wait in old_waits:
+                erased_to_new[old_wait] = new_wait
+        else:
+            # Coalesced bucketing: single start + N waits (one per original tensor)
+            assert len(new_waits) == len(old_waits)
+            for old_start in old_starts:
+                erased_to_new[old_start] = new_start
+            erased_to_new.update(dict(zip(old_waits, new_waits)))
 
         # Handle convert_element_type nodes that were fused and erased
         # The bucketed operation may have a _pre_bucket op that handles dtype conversion
         if fused_convert_dtypes:
-            # all gather bucketing may fuse in dtype conversion into the bucketing
-            # if so, we need to transfer hiding deps from the old dtype conversion
-            # to the new bucketing node
-            new_convert_dtypes_node = new_start.kwargs["out"]
-            assert isinstance(new_convert_dtypes_node, fx.Node)
-            assert (
-                new_convert_dtypes_node.target
+            # In custom_ops mode, the _pre_bucket_all_gather node handles dtype conversion
+            # In default mode, convert nodes are just erased — map them to new_start
+            new_convert_dtypes_node = new_start.kwargs.get("out")
+            if (
+                isinstance(new_convert_dtypes_node, fx.Node)
+                and new_convert_dtypes_node.target
                 == torch.ops.bucketing._pre_bucket_all_gather.default
-            )
+            ):
+                replacement = new_convert_dtypes_node
+            else:
+                replacement = new_start
 
             for n in fused_convert_dtypes:
-                erased_to_new[n] = new_convert_dtypes_node
+                erased_to_new[n] = replacement
 
         # Transfer all dependencies from old nodes to new nodes
         self.aug_graph.transfer_erased_node_deps(erased_to_new)
@@ -1023,7 +1078,9 @@ def finalize_overlap_scheduling(
     max_bucket_memory_gb: float = 2.0,
     max_coll_distance: int = 1000,
     region_of: dict[fx.Node, Any] | None = None,
-    bucket_exposed_first: bool = True,
+    bucket_exposed_first: bool | None = None,
+    bucket_only_internode_comms: bool = False,
+    bucket_mode: BucketMode | None = None,
 ) -> None:
     """
     Finalize overlap scheduling by applying deps, inlining fusions, and optionally bucketing.
@@ -1052,7 +1109,9 @@ def finalize_overlap_scheduling(
         max_coll_distance=max_coll_distance,
         insert_overlap_deps=insert_overlap_deps,
         collective_bucketing=collective_bucketing,
-        region_of=region_of,
         bucket_exposed_first=bucket_exposed_first,
+        bucket_only_internode_comms=bucket_only_internode_comms,
+        region_of=region_of,
+        bucket_mode=bucket_mode,
     )
     bucketer.bucket_collectives()
diff --git a/torch/_inductor/fx_passes/overlap_scheduling.py b/torch/_inductor/fx_passes/overlap_scheduling.py
index c38ed9acb1b3e..08b7853c667ab 100644
--- a/torch/_inductor/fx_passes/overlap_scheduling.py
+++ b/torch/_inductor/fx_passes/overlap_scheduling.py
@@ -11,8 +11,11 @@
 import torch
 import torch.fx as fx
 from torch._dynamo.utils import counters, dynamo_timed
+from torch._inductor import config
 from torch._inductor.comm_analysis import estimate_fx_collective_memory_footprint
 from torch._inductor.fx_passes.bucketing import (
+    _default_bucket_mode,
+    _get_collective_node_from_wait,
     _schedulable_wait_node,
     bucket_key,
     BucketMode,
@@ -28,7 +31,7 @@
 
 log = logging.getLogger(__name__)
 
-from ..pattern_matcher import stable_topological_sort
+from torch._inductor.pattern_matcher import stable_topological_sort
 
 
 def make_all_device_put_sync(gm: torch.fx.GraphModule) -> int:
@@ -84,24 +87,6 @@ def make_all_device_put_sync(gm: torch.fx.GraphModule) -> int:
     return count
 
 
-def estimate_runtime_analytical(n: torch.fx.Node) -> float:
-    """Estimate runtime using analytical roofline model for mm operations."""
-    if n.target != torch.ops.aten.mm.default:
-        return 0.0
-    import torch.utils._pytree as pytree
-    from torch.distributed._tools import RuntimeEstimator
-
-    def _val(node: Any) -> Any:
-        if not isinstance(node, torch.fx.Node):
-            return node
-        return node.meta["val"]
-
-    args = pytree.tree_map(_val, n.args)
-    kwargs = pytree.tree_map(_val, n.kwargs)
-    _, ms = RuntimeEstimator._roofline_estimate(n.target, args, kwargs)
-    return ms
-
-
 @dataclass
 class WhyNoOverlap:
     """Track reasons why a collective cannot overlap with compute."""
@@ -153,6 +138,7 @@ def estimate_collective_time(
     override_size: int | None = None,
     custom_runtime_estimation: Callable[[fx.Node, int | None], float | None]
     | None = None,
+    collective_estimator: Literal["analytical", "benchmark"] = "analytical",
 ) -> float:
     """Estimate the runtime of a collective operation, optionally with an overridden size."""
     if (
@@ -160,7 +146,16 @@ def estimate_collective_time(
     ) is not None:
         return est
 
-    # Use analytical model (benchmarking is handled separately in alignment)
+    if collective_estimator == "benchmark":
+        from torch._inductor.fx_passes.node_runtime_estimation import (
+            benchmark_collective_with_cuda_events,
+        )
+
+        cuda_val, _ = benchmark_collective_with_cuda_events(n, nruns=5)
+        if cuda_val is not None:
+            return cuda_val
+
+    # Analytical model (also fallback when benchmark returns None)
     return torch._inductor.comm_analysis.estimate_nccl_collective_runtime_from_fx_node(
         n, override_size
     )
@@ -177,25 +172,52 @@ def is_compute_node(n: fx.Node) -> bool:
     )
 
 
-def estimate_mem_bound_runtime_ms(node: fx.Node) -> float:
-    """Estimate runtime for a memory-bound node based on input/output bytes.
+def estimate_roofline_runtime_ms(node: fx.Node) -> float:
+    """Estimate runtime using roofline model (max of compute and memory bound).
 
-    Returns 0 for view nodes (no memory cost).
+    Uses FLOPs for compute-bound estimate if op is in flop_registry,
+    and memory bandwidth for memory-bound estimate.
+    Returns 0 for view nodes (no cost).
     """
     from torch._inductor.fx_passes.fusion_regions import is_view_node
-    from torch.utils._pytree import tree_flatten
-    from torch.utils._runtime_estimation import get_transfer_time
+    from torch.utils._pytree import tree_flatten, tree_map
+    from torch.utils._runtime_estimation import get_compute_time, get_transfer_time
+    from torch.utils.flop_counter import flop_registry
 
     if is_view_node(node):
         return 0.0
 
-    input_vals = [inp.meta.get("val") for inp in node.all_input_nodes]
-    output_vals = [node.meta.get("val")]
-    flat_inputs, _ = tree_flatten(input_vals)
-    flat_outputs, _ = tree_flatten(output_vals)
+    def _get_val(n: Any) -> Any:
+        if isinstance(n, fx.Node):
+            return n.meta.get("val")
+        return n
+
+    args = tree_map(_get_val, node.args)
+    kwargs = tree_map(_get_val, node.kwargs)
+    out = _get_val(node)
+
+    if out is None:
+        return 0.0
+
+    flat_args_kwargs, _ = tree_flatten((args, kwargs))
+    flat_outs, _ = tree_flatten(out)
+    out_dtypes = OrderedSet([t.dtype for t in flat_outs if isinstance(t, torch.Tensor)])
 
-    transfer_time_ns = get_transfer_time(flat_inputs, flat_outputs)
-    return transfer_time_ns / 1e6
+    # Compute time (FLOPs-based, only if op is in flop_registry)
+    # May return SymFloat if shapes are symbolic (after flop division)
+    compute_ns: float = 0.0
+    func_packet = getattr(node.target, "overloadpacket", None)
+    if func_packet in flop_registry and len(out_dtypes) == 1:
+        compute_ns = get_compute_time(func_packet, args, kwargs, out, out_dtypes.copy())
+        # Extract hint from symbolic value if needed
+        if isinstance(compute_ns, (torch.SymInt, torch.SymFloat)):
+            compute_ns = compute_ns.node.hint if compute_ns.node.has_hint() else 0.0
+
+    # Transfer time (memory bandwidth-based, uses size_hint internally)
+    transfer_ns = get_transfer_time(flat_args_kwargs, flat_outs)
+
+    # Roofline: max of compute and transfer, convert ns to ms
+    return max(float(compute_ns), float(transfer_ns)) / 1e6
 
 
 def get_hint(x: int | torch.SymInt) -> int | None:
@@ -359,14 +381,18 @@ def __init__(
         insert_overlap_deps: bool,
         compute_overlap_multipler: float,
         max_coll_distance: int,
-        custom_runtime_estimation: Callable[[fx.Node, int | None], float | None] | None,
-        collective_estimator: Literal["analytical", "benchmark"],
+        custom_runtime_estimation: Callable[[fx.Node, int | None], float | None]
+        | None = None,
+        collective_estimator: Literal["analytical", "benchmark"] = "analytical",
+        compute_estimator: Literal["analytical", "benchmark"] = "benchmark",
         max_memory_increase_gb: float | None = 1.0,
         max_memory_increase_ratio: float | None = 0.05,
         log_final_collectives_estimations: bool = False,
-        bucket_exposed_first: bool = True,
+        bucket_exposed_first: bool | None = None,
         enable_fusion_regions: bool = False,
-        bucket_mode: BucketMode = "custom_ops_multidtype",
+        bucket_only_internode_comms: bool = False,
+        bucket_mode: BucketMode | None = None,
+        max_off_bucket_gb: float | None = 0.5,
         prioritize_bucketing_during_scheduling: bool = True,
     ):
         self.gm = gm
@@ -378,10 +404,20 @@ def __init__(
         self.collective_bucketing = collective_bucketing
         self.insert_overlap_deps = insert_overlap_deps
         self.max_compute_pre_fetch = max_compute_pre_fetch
-        self.collective_estimator = collective_estimator
+        # In deterministic mode, force analytical estimation to avoid GPU sync
+        if config.deterministic:
+            self.collective_estimator = "analytical"
+            self.compute_estimator = "analytical"
+        else:
+            self.collective_estimator = collective_estimator
+            self.compute_estimator = compute_estimator
         self.log_final_collectives_estimations = log_final_collectives_estimations
         self.bucket_exposed_first = bucket_exposed_first
-        self.bucket_mode = bucket_mode
+        self.bucket_only_internode_comms = bucket_only_internode_comms
+        self.bucket_mode = bucket_mode or _default_bucket_mode()
+        self.max_off_bucket_bytes: int | None = (
+            gb_to_bytes(max_off_bucket_gb) if max_off_bucket_gb is not None else None
+        )
         self.prioritize_bucketing_during_scheduling = (
             prioritize_bucketing_during_scheduling
         )
@@ -398,20 +434,18 @@ def __init__(
                 num_device_put_converted,
             )
 
-        # Build and collapse fusion regions FIRST so all subsequent operations
-        # work on the collapsed graph where fused ops are atomic units
-        self.region_of: dict[fx.Node, Any] = {}
-        if enable_fusion_regions:
-            from torch._inductor.fx_passes.fusion_regions import (
-                build_fusion_regions,
-                collapse_fusion_regions,
-            )
-
-            self.region_of = build_fusion_regions(self.gm)
-            if self.region_of:
-                self.region_of = collapse_fusion_regions(self.gm, self.region_of)
-                # fuse_by_partitions replaces gm.graph, so we need to update our reference
-                self.graph = gm.graph
+        # Build fusion regions (mutates gm.graph) and compute initial node runtime
+        # estimates. Compute nodes use roofline model here; the alignment step in
+        # run() replaces them with benchmarked + cross-rank-aligned values.
+        self.node_estimations, self.region_of = gather_node_runtime_estimations(
+            gm,
+            custom_runtime_estimation,
+            enable_fusion_regions=enable_fusion_regions,
+            log_estimations=True,
+        )
+        if self.region_of:
+            # fuse_by_partitions replaces gm.graph, so we need to update our reference
+            self.graph = gm.graph
 
         # Build structures
         stable_topological_sort(self.graph)
@@ -487,6 +521,8 @@ def __init__(
         # Two separate queues: on-path (domination-based) and off-path (node_idx-based)
         self.on_path_ready: list[tuple[object, fx.Node]] = []
         self.off_path_ready: list[tuple[object, fx.Node]] = []
+        # Track potential bucket sizes for off-path collectives (for batch scheduling)
+        self.off_path_ready_potential_buckets: dict[object, int] = defaultdict(int)
 
         for node in self.nodes:
             if self.in_degree[node] == 0:
@@ -503,6 +539,11 @@ def _add_to_ready_queue(self, node: fx.Node) -> None:
         if self.off_compute_path(node):
             score = self._compute_off_path_score(node)
             heapq.heappush(self.off_path_ready, (score, node))
+            # Track potential bucket sizes for off-path ready collectives
+            if node in self.collective_info:
+                bucket_key = get_full_bucket_key(node, self.bucket_mode)
+                node_bytes = self.collective_info[node].size_bytes
+                self.off_path_ready_potential_buckets[bucket_key] += node_bytes
         else:
             score = self._compute_on_path_score(node)
             heapq.heappush(self.on_path_ready, (score, node))
@@ -613,10 +654,18 @@ def _identify_collectives(self) -> None:
 
         for node in self.nodes:
             if _schedulable_wait_node(node):
-                start = node.args[0]
-                coll_time_ms = estimate_collective_time(
-                    start, custom_runtime_estimation=self.custom_runtime_estimation
+                start = _get_collective_node_from_wait(node)
+                assert start is not None
+                assert start in self.node_estimations, (
+                    f"Missing estimation for collective {start.name}. "
+                    f"Ensure custom_runtime_estimation returns a value for this node."
                 )
+                self.wait_to_start[node] = start
+                # For coalesced collectives, multiple waits share the same
+                # start node. Only register the first wait as the representative.
+                if start in self.collective_info:
+                    continue
+                coll_time_ms = self.node_estimations[start]
 
                 info = CollectiveInfo(
                     start_node=start,
@@ -626,7 +675,6 @@ def _identify_collectives(self) -> None:
                     exposed_time_ms=coll_time_ms,  # Initially fully exposed
                 )
                 self.collective_info[start] = info
-                self.wait_to_start[node] = start
                 self.unscheduled_collectives.add(start)
                 self.all_pgs.add(get_group_name(start))
 
@@ -673,12 +721,18 @@ def _align_compute_nodes_runtime_estimations_across_all_distributed_ranks(
         runtime_estimations_analytical: list[float] = []
 
         for n in self.compute_nodes:
-            val, key = benchmark_node_with_cache_key(n, self.custom_runtime_estimation)
-
-            # Analytical estimations
-            val_analytical = estimate_runtime_analytical(n)
+            # Compute analytical estimation using roofline model
+            val_analytical = estimate_roofline_runtime_ms(n)
             runtime_estimations_analytical.append(val_analytical)
 
+            if self.compute_estimator == "benchmark":
+                val, key = benchmark_node_with_cache_key(
+                    n, self.custom_runtime_estimation
+                )
+            else:
+                # Use analytical estimation
+                val, key = val_analytical, None
+
             runtime_estimations.append(val)
             runtime_estimations_keys.append(key)
             compute_key_count += 1
@@ -725,24 +779,33 @@ def _align_compute_nodes_runtime_estimations_across_all_distributed_ranks(
                     runtime_estimations_keys.append(cuda_key)
                     benchmarked_collective_nodes.append(n)
 
-        # Single all_gather and compute medians
+        # When both estimators are analytical, estimates are deterministic across ranks
+        # (same shapes = same estimates), so skip the all_gather to avoid sync.
         import torch.distributed as dist
-        from torch._subclasses.fake_tensor import unset_fake_temporarily
-        from torch.distributed.distributed_c10d import _get_default_group
 
         world_size = dist.get_world_size()
-        pg = _get_default_group()
 
-        with unset_fake_temporarily():
-            gathered_runtime_estimations: list[list[float]] = [
-                [] for _ in range(world_size)
-            ]
-            dist.all_gather_object(
-                gathered_runtime_estimations, runtime_estimations, pg
-            )
-            median_runtime_estimations = torch.median(
-                torch.tensor(gathered_runtime_estimations), dim=0
-            ).values.tolist()
+        if (
+            self.compute_estimator == "analytical"
+            and self.collective_estimator == "analytical"
+        ):
+            median_runtime_estimations = runtime_estimations
+        else:
+            # Single all_gather and compute medians
+            from torch._subclasses.fake_tensor import unset_fake_temporarily
+            from torch.distributed.distributed_c10d import _get_default_group
+
+            pg = _get_default_group()
+            with unset_fake_temporarily():
+                gathered_runtime_estimations: list[list[float]] = [
+                    [] for _ in range(world_size)
+                ]
+                dist.all_gather_object(
+                    gathered_runtime_estimations, runtime_estimations, pg
+                )
+                median_runtime_estimations = torch.median(
+                    torch.tensor(gathered_runtime_estimations), dim=0
+                ).values.tolist()
 
         # Cache medians
         collective_keys = []
@@ -754,6 +817,9 @@ def _align_compute_nodes_runtime_estimations_across_all_distributed_ranks(
                 continue
             if idx < compute_key_count:
                 # Compute node
+                self.node_estimations[self.compute_nodes[idx]] = (
+                    median_runtime_estimation
+                )
                 set_cached_node_time(key, median_runtime_estimation)
             else:
                 # Collective CUDA event benchmark
@@ -769,6 +835,7 @@ def _align_compute_nodes_runtime_estimations_across_all_distributed_ranks(
                 info = self.collective_info[coll_node]
                 info.estimated_time_ms = median_runtime_estimation
                 info.exposed_time_ms = median_runtime_estimation
+                self.node_estimations[coll_node] = median_runtime_estimation
 
                 collective_keys.append(key)
                 collective_medians.append(median_runtime_estimation)
@@ -803,8 +870,14 @@ def _align_compute_nodes_runtime_estimations_across_all_distributed_ranks(
 
         log.info("Overlap scheduling: Runtime estimations aligned")
 
-    def _get_next_node(self) -> fx.Node:
-        """Get next node: off-path nodes scheduled near original position, exposed waits deferred."""
+    def _get_next_nodes(self) -> list[fx.Node]:
+        """
+        Get next node(s) to schedule.
+
+        When max_off_bucket_bytes is set, off-path collectives of the same type
+        (same bucket_key) are batched together to enable bucketing them in
+        overlap_preserving_bucketer. Bucket size is limited by max_off_bucket_bytes.
+        """
         if self.off_path_ready:
             _, node = self.off_path_ready[0]
 
@@ -819,15 +892,63 @@ def _get_next_node(self) -> fx.Node:
                     > self.allowed_peak_memory_bytes
                 )
                 should_schedule = not info.is_exposed or over_budget
+            elif self.max_off_bucket_bytes is not None and node in self.collective_info:
+                # Batch off-path collectives: schedule when bucket threshold is reached
+                bucket_key = get_full_bucket_key(node, self.bucket_mode)
+                bucket_size = self.off_path_ready_potential_buckets[bucket_key]
+                should_schedule = bucket_size >= self.max_off_bucket_bytes
             elif self.dominates_reduce_scatter(node):
                 # Only schedule off-path nodes that dominate reduce_scatters after original position
                 should_schedule = self.node_idx[node] <= self.last_on_path_node_idx
 
             if should_schedule:
                 heapq.heappop(self.off_path_ready)
-                return node
 
-        return heapq.heappop(self.on_path_ready)[1]
+                # If batching enabled and this is a collective, gather same-type collectives
+                if (
+                    self.max_off_bucket_bytes is not None
+                    and node in self.collective_info
+                ):
+                    node_key = get_full_bucket_key(node, self.bucket_mode)
+                    if node_key is not None:
+                        same_type_nodes = [node]
+                        total_bytes = self.collective_info[node].size_bytes
+                        indices_to_remove = []
+
+                        # Scan the off_path_ready queue for same-key collectives
+                        for i, (_, candidate) in enumerate(self.off_path_ready):
+                            if candidate in self.scheduled:
+                                continue
+                            if candidate not in self.collective_info:
+                                continue
+                            candidate_key = get_full_bucket_key(
+                                candidate, self.bucket_mode
+                            )
+                            if candidate_key == node_key:
+                                candidate_bytes = self.collective_info[
+                                    candidate
+                                ].size_bytes
+                                # Check bucket size limit before adding
+                                if (
+                                    total_bytes + candidate_bytes
+                                    > self.max_off_bucket_bytes
+                                ):
+                                    continue  # Skip but keep looking for smaller ones
+                                same_type_nodes.append(candidate)
+                                total_bytes += candidate_bytes
+                                indices_to_remove.append(i)
+
+                        # Remove collected nodes from heap (reverse order to preserve indices)
+                        for i in reversed(indices_to_remove):
+                            self.off_path_ready.pop(i)
+                        if indices_to_remove:
+                            heapq.heapify(self.off_path_ready)
+
+                        return same_type_nodes
+
+                return [node]
+
+        return [heapq.heappop(self.on_path_ready)[1]]
 
     def run(self) -> torch.fx.GraphModule:
         """Run the scheduling algorithm."""
@@ -841,26 +962,33 @@ def run(self) -> torch.fx.GraphModule:
                 self._force_oldest_wait()
                 continue
 
-            node = self._get_next_node()
+            nodes = self._get_next_nodes()
 
-            # we don't always remove nodes from the heap when we schedule them
-            if node in self.scheduled:
-                continue
+            for node in nodes:
+                # we don't always remove nodes from the heap when we schedule them
+                if node in self.scheduled:
+                    continue
 
-            if node.op == "placeholder":
-                self._schedule(node)
-            elif node in self.collective_info:
-                self._handle_collective_start(node)
-            elif _schedulable_wait_node(node):
-                self._handle_wait(node)
-            else:
-                self._handle_compute_or_other(node)
-
-            # Track progress for off-path scheduling - only for nodes from main queue
-            if not self.off_compute_path(node):
-                self.last_on_path_node_idx = max(
-                    self.last_on_path_node_idx, self.node_idx[node]
-                )
+                if node.op == "placeholder":
+                    self._schedule(node)
+                elif node in self.collective_info:
+                    self._handle_collective_start(node)
+                elif _schedulable_wait_node(node):
+                    self._handle_wait(node)
+                else:
+                    self._handle_compute_or_other(node)
+
+                # Track progress for off-path scheduling - only for nodes from main queue
+                if not self.off_compute_path(node):
+                    self.last_on_path_node_idx = max(
+                        self.last_on_path_node_idx, self.node_idx[node]
+                    )
+                else:
+                    # Decrement off-path bucket bytes when scheduling
+                    if node in self.collective_info:
+                        bucket_key = get_full_bucket_key(node, self.bucket_mode)
+                        node_bytes = self.collective_info[node].size_bytes
+                        self.off_path_ready_potential_buckets[bucket_key] -= node_bytes
 
         self._reorder_graph()
 
@@ -879,6 +1007,8 @@ def run(self) -> torch.fx.GraphModule:
             max_coll_distance=self.max_node_distance,
             region_of=self.region_of,
             bucket_exposed_first=self.bucket_exposed_first,
+            bucket_only_internode_comms=self.bucket_only_internode_comms,
+            bucket_mode=self.bucket_mode,
         )
 
         if self.log_final_collectives_estimations:
@@ -892,29 +1022,6 @@ def run(self) -> torch.fx.GraphModule:
 
         return self.gm
 
-    def get_non_collective_runtime_estimate(self, node: fx.Node) -> float | None:
-        """Get runtime estimation for a node in ms. Returns None if no estimation is available."""
-        if is_compute_node(node):
-            return benchmark_node(node, self.custom_runtime_estimation)
-
-        # Use precomputed cost for fusion region call_module nodes
-        # This takes priority even over custom estimation since fusion regions
-        # have already computed their cost based on their contents
-        if node in self.region_of:
-            return self.region_of[node].cost_ms
-
-        if self.custom_runtime_estimation is not None:
-            if (est := self.custom_runtime_estimation(node, None)) is not None:
-                return est
-            # Custom estimation provided but returned None - don't fall through to fusible estimation
-            return None
-
-        # assume any node without flop counter is mem bound
-        if node.op == "call_function":
-            return estimate_mem_bound_runtime_ms(node)
-
-        return None
-
     def _reduce_exposed_time_of_in_flight_collectives(
         self,
         node: fx.Node,
@@ -953,7 +1060,7 @@ def _reduce_exposed_time_of_in_flight_collectives(
 
     def _handle_compute_or_other(self, node: fx.Node) -> None:
         """Handle scheduling compute or other nodes and attempt to overlap with collectives."""
-        runtime_estimate = self.get_non_collective_runtime_estimate(node)
+        runtime_estimate = self.node_estimations.get(node)
 
         # TODO: we could consider skipping overlapping for overlapable, unary chains to collectives.
         # using these nodes for overlap prevents bucketing. potentially if chain time < latency
@@ -1081,7 +1188,11 @@ def _handle_wait(self, node: fx.Node) -> None:
         """Handle scheduling a wait."""
         assert node in self.wait_to_start
         coll_start = self.wait_to_start[node]
-        assert coll_start in self.in_flight
+        # For coalesced collectives, multiple waits share the same start node.
+        # The first wait completes the collective; subsequent waits just schedule.
+        if coll_start not in self.in_flight:
+            self._schedule(node)
+            return
 
         # Scheduling a wait of a collective also forces the wait
         # of every node enqueued prior to the collective on the
@@ -1323,12 +1434,12 @@ def should_assume_bucketed(self, node: fx.Node) -> bool:
         if not torch._inductor.config.test_configs.assume_bucketing_reduces_latency:
             return False
 
-        key = bucket_key(node, mode="custom_ops_multidtype")
+        key = bucket_key(node, mode=self.bucket_mode)
         if key is None:
             return False
 
         for in_flight_coll in self.in_flight:
-            if bucket_key(in_flight_coll, mode="custom_ops_multidtype") == key:
+            if bucket_key(in_flight_coll, mode=self.bucket_mode) == key:
                 return True
 
         return False
@@ -1436,8 +1547,9 @@ def _bucket_collectives(self) -> None:
             max_bucket_memory_gb=2.0,  # Could make this configurable
             max_coll_distance=self.max_node_distance,
             insert_overlap_deps=self.insert_overlap_deps,
-            bucket_exposed_first=self.bucket_exposed_first,
             bucket_mode=self.bucket_mode,
+            bucket_exposed_first=self.bucket_exposed_first,
+            bucket_only_internode_comms=self.bucket_only_internode_comms,
         )
         bucketer.bucket_collectives()
 
@@ -1479,6 +1591,149 @@ def compute_potential_hidden_waits(self) -> dict[fx.Node, fx.Node]:
         return self.compute_potential_hidden_nodes(wait_nodes)
 
 
+def gather_node_runtime_estimations(
+    gm: torch.fx.GraphModule,
+    custom_runtime_estimation: Callable[[fx.Node, int | None], float | None]
+    | None = None,
+    collective_estimator: Literal["analytical", "benchmark"] = "analytical",
+    enable_fusion_regions: bool = False,
+    log_estimations: bool = False,
+) -> tuple[dict[fx.Node, float], dict[fx.Node, Any]]:
+    """Gather initial runtime estimations for all nodes without scheduling.
+
+    Uses analytical models (roofline) for compute nodes — the alignment step
+    in OverlapScheduler.run() replaces these with benchmarked + cross-rank-aligned
+    values. Collectives use bandwidth formulas or CUDA events depending on
+    collective_estimator.
+
+    When enable_fusion_regions is True, builds and collapses fusion regions
+    (mutating gm's graph), then includes their costs in the estimations.
+
+    Args:
+        collective_estimator: "analytical" uses bandwidth formulas,
+            "benchmark" uses CUDA events for collectives.
+        log_estimations: When True, log compute and collective estimations
+            via trace_structured for tlparse.
+
+    Returns (estimations, fusion_region_of) where estimations maps fx.Node to
+    runtime in ms, and fusion_region_of maps call_module nodes to FusionRegion
+    objects (empty dict if fusion regions are disabled).
+    """
+    # Build and collapse fusion regions first (mutates gm)
+    fusion_region_of: dict[fx.Node, Any] = {}
+    if enable_fusion_regions:
+        from torch._inductor.fx_passes.fusion_regions import (
+            build_fusion_regions,
+            collapse_fusion_regions,
+        )
+
+        fusion_region_of = build_fusion_regions(gm)
+        if fusion_region_of:
+            fusion_region_of = collapse_fusion_regions(gm, fusion_region_of)
+
+    estimations: dict[fx.Node, float] = {}
+    nodes = list(gm.graph.nodes)
+
+    # Collectives
+    collective_nodes: list[fx.Node] = []
+    for node in nodes:
+        if _schedulable_wait_node(node):
+            start = _get_collective_node_from_wait(node)
+            assert start is not None
+            if start in estimations:
+                continue
+            estimations[start] = estimate_collective_time(
+                start,
+                custom_runtime_estimation=custom_runtime_estimation,
+                collective_estimator=collective_estimator,
+            )
+            collective_nodes.append(start)
+
+    # Compute nodes (matmul, bmm, etc.) — analytical estimates only.
+    # The alignment step in run() replaces these with benchmarked + aligned values.
+    compute_nodes: list[fx.Node] = []
+    compute_analytical: list[float] = []
+
+    for node in nodes:
+        if is_compute_node(node):
+            est = estimate_roofline_runtime_ms(node)
+            if custom_runtime_estimation is not None:
+                custom_est = custom_runtime_estimation(node, None)
+                if custom_est is not None:
+                    est = custom_est
+            estimations[node] = est
+            compute_nodes.append(node)
+            compute_analytical.append(est)
+        elif node.op == "call_function" and node not in estimations:
+            if custom_runtime_estimation is not None:
+                est = custom_runtime_estimation(node, None)
+                if est is not None:
+                    estimations[node] = est
+            else:
+                est = estimate_roofline_runtime_ms(node)
+                if est > 0:
+                    estimations[node] = est
+
+    # Fusion region costs (call_module nodes from collapse_fusion_regions)
+    for node, region in fusion_region_of.items():
+        estimations[node] = region.cost_ms  # pyrefly: ignore[missing-attribute]
+
+    # Logging
+    if log_estimations and compute_nodes:
+        from torch._inductor.fx_passes.node_runtime_estimation import (
+            _log_compute_estimations,
+        )
+
+        _log_compute_estimations(
+            compute_nodes,
+            compute_analytical,
+            compute_analytical,
+        )
+
+    if log_estimations and collective_nodes:
+        from torch._inductor.fx_passes.node_runtime_estimation import (
+            _log_collective_benchmarks,
+        )
+
+        _log_collective_benchmarks(
+            collective_nodes,
+            artifact_name="fx_collectives_analytical_estimation",
+        )
+
+    return estimations, fusion_region_of
+
+
+def align_estimations_across_ranks(
+    estimations: dict[fx.Node, float],
+) -> dict[fx.Node, float]:
+    """Align runtime estimations across distributed ranks using median.
+
+    All ranks must make identical scheduling decisions, so we gather each
+    rank's values and take the median. All nodes in estimations are aligned.
+
+    Returns a new estimations dict with aligned values.
+    """
+    import torch.distributed as dist
+    from torch._subclasses.fake_tensor import unset_fake_temporarily
+    from torch.distributed.distributed_c10d import _get_default_group
+
+    nodes = list(estimations.keys())
+    if not nodes:
+        return {}
+
+    local_values = [estimations[n] for n in nodes]
+
+    world_size = dist.get_world_size()
+    pg = _get_default_group()
+
+    with unset_fake_temporarily():
+        gathered: list[list[float]] = [[] for _ in range(world_size)]
+        dist.all_gather_object(gathered, local_values, pg)
+        medians = torch.median(torch.tensor(gathered), dim=0).values.tolist()
+
+    return dict(zip(nodes, medians))
+
+
 def schedule_overlap_bucketing(
     gm: torch.fx.GraphModule,
     max_in_flight_gb: float = 5,
@@ -1490,12 +1745,16 @@ def schedule_overlap_bucketing(
     custom_runtime_estimation: Callable[[fx.Node, int | None], float | None]
     | None = None,
     collective_estimator: Literal["analytical", "benchmark"] = "analytical",
+    compute_estimator: Literal["analytical", "benchmark"] = "benchmark",
     max_memory_increase_gb: float | None = 1.0,
     max_memory_increase_ratio: float | None = 0.05,
     log_final_collectives_estimations: bool = False,
-    bucket_exposed_first: bool = True,
+    bucket_exposed_first: bool | None = None,
     enable_fusion_regions: bool = False,
+    bucket_only_internode_comms=False,
     prioritize_bucketing_during_scheduling: bool = True,
+    max_off_bucket_gb: float | None = 0.5,
+    bucket_mode: BucketMode | None = None,
 ) -> torch.fx.GraphModule:
     """Schedule nodes to maximize compute-collective overlap.
 
@@ -1511,14 +1770,18 @@ def schedule_overlap_bucketing(
         compute_overlap_multipler: Scale factor for compute time used to hide collectives. This can be used
             to address over or under aggressive overlapping.
         max_coll_distance: Maximum pre fetch or bucketing candidates. Mainly intended for compile time
-        custom_runtime_estimation: Custom runtime estimation function that estimates runtime in ms for an fx node.
-            If None, uses default estimations. This is currently limited to collectives and compute nodes.
+        custom_runtime_estimation: Override runtime estimation for specific nodes. Called as
+            custom_runtime_estimation(node, override_size) -> float | None. To pass pre-computed
+            estimations, wrap a dict: lambda node, _: estimations.get(node).
         collective_estimator: Method for estimating collective runtime. "analytical" uses bandwidth formulas,
             "benchmark" uses CUDA events with power-of-2 rounding and interpolation.
+        compute_estimator: Method for estimating compute (ATen op) runtime. "analytical" uses roofline model
+            estimates (deterministic, no GPU sync), "benchmark" uses GPU benchmarking (more accurate).
         max_memory_increase_gb: Maximum GB increase above baseline memory (absolute cap). If None, no absolute limit.
         max_memory_increase_ratio: Maximum increase as ratio of baseline peak memory. If None, no ratio limit.
             Uses minimum of absolute and ratio limits when both are specified.
         enable_fusion_regions: Enable fusion region detection and cost estimation for fusible ops.
+        bucket_mode: Bucketing mode for grouping collectives.
     """
     if not any(is_wait_tensor(n) for n in gm.graph.nodes):
         return gm
@@ -1541,12 +1804,16 @@ def schedule_overlap_bucketing(
         collective_bucketing=collective_bucketing,
         insert_overlap_deps=insert_overlap_deps,
         collective_estimator=collective_estimator,
+        compute_estimator=compute_estimator,
         max_memory_increase_gb=max_memory_increase_gb,
         max_memory_increase_ratio=max_memory_increase_ratio,
         log_final_collectives_estimations=log_final_collectives_estimations,
         bucket_exposed_first=bucket_exposed_first,
         enable_fusion_regions=enable_fusion_regions,
+        bucket_only_internode_comms=bucket_only_internode_comms,
         prioritize_bucketing_during_scheduling=prioritize_bucketing_during_scheduling,
+        max_off_bucket_gb=max_off_bucket_gb,
+        bucket_mode=bucket_mode,
     ).run()
     trace_structured(
         "artifact",
@@ -1582,6 +1849,7 @@ def schedule_overlap_bucketing_from_inductor_configs(
         "custom_runtime_estimation",
         "insert_overlap_deps",
         "collective_estimator",
+        "compute_estimator",
         "max_memory_increase_gb",
         "max_memory_increase_ratio",
         "compute_overlap_multipler",
@@ -1589,8 +1857,10 @@ def schedule_overlap_bucketing_from_inductor_configs(
         "max_coll_distance",
         "log_final_collectives_estimations",
         "bucket_exposed_first",
+        "bucket_only_internode_comms",
         "enable_fusion_regions",
         "prioritize_bucketing_during_scheduling",
+        "bucket_mode",
     )
     for key in config_keys:
         if (val := getattr(dist_opts, key, None)) is not None:
diff --git a/torch/_inductor/fx_passes/pad_mm.py b/torch/_inductor/fx_passes/pad_mm.py
index 6fbf655e46c87..5a0deb68090f8 100644
--- a/torch/_inductor/fx_passes/pad_mm.py
+++ b/torch/_inductor/fx_passes/pad_mm.py
@@ -3,7 +3,7 @@
 import operator
 import typing
 from collections.abc import Callable, Sequence
-from typing import Any, Optional
+from typing import Any
 
 import torch
 import torch._inductor.runtime.runtime_utils
@@ -84,11 +84,13 @@ def check_dtype(a: Tensor, b: Tensor) -> bool:
     return a.is_floating_point() and b.is_floating_point()
 
 
-def realize_symbols(
-    ds: torch.Size | tuple[torch.SymInt, ...],
+def hint_symbols(
+    ds: Sequence[int | torch.SymInt],
 ) -> list[int]:
     """Helper to convert symbolic dimensions to their concrete hint values."""
-    return [d if isinstance(d, int) else d.node.hint for d in ds]
+    from torch.fx.experimental.symbolic_shapes import optimization_hint
+
+    return [optimization_hint(d) for d in ds]
 
 
 def can_pad(
@@ -102,31 +104,15 @@ def can_pad(
     All logic related to whether it's safe to pad should be here.
     """
 
-    # It's fine we have symbolic shapes or strides as long as they
-    # have hints. Later, we will make sure we only pad non-symbolic dimensions.
-    def valid_shape_and_stride(t: Tensor | None) -> bool:
-        if t is None:
-            return True
-
-        symbolic_cnt = 0
+    # Can't pad if there is no static dims, we pad static dims only.
+    def has_one_static_dim(t: Tensor) -> bool:
+        """Return False if all dimensions are symbolic — nothing concrete to pad."""
         for x in t.size():
             if isinstance(x, int):
-                continue
-            elif utils.is_symbolic(x):
-                # pyrefly: ignore [missing-attribute]
-                if not x.node.has_hint():
-                    return False
-                symbolic_cnt += 1
-            else:
-                return False
-        # filter out cases where all dimensions are symbolic
-        if symbolic_cnt == len(t.size()):
-            return False
-        return all(
-            # pyrefly: ignore [missing-attribute]
-            isinstance(x, int) or (utils.is_symbolic(x) and x.node.has_hint())
-            for x in t.stride()
-        )
+                return True
+            elif not isinstance(x, torch.SymInt):
+                raise RuntimeError("not expected size")
+        return False
 
     # Basic safety checks
     if not torch._inductor.config.shape_padding:
@@ -138,15 +124,16 @@ def valid_shape_and_stride(t: Tensor | None) -> bool:
     if not check_dtype(mat1, mat2):
         return False
 
-    if not all(valid_shape_and_stride(t) for t in (mat1, mat2, input)):
+    # For padding to be vaible each tensor should have at least one static dim.
+    tensors = [t for t in (mat1, mat2, input) if t is not None]
+    if not all(has_one_static_dim(t) for t in tensors):
         return False
 
-    # Check for zero dimensions - not safe to pad
+    # Skip zero-sized dimensions — padding would be wasteful (mm on empty tensors)
+    from torch.fx.experimental.symbolic_shapes import optimization_hint
+
     if any(
-        dim == 0
-        for dim in itertools.chain(
-            realize_symbols(mat1.shape), realize_symbols(mat2.shape)
-        )
+        optimization_hint(dim) == 0 for dim in itertools.chain(mat1.shape, mat2.shape)
     ):
         return False
 
@@ -172,10 +159,11 @@ def valid_shape_and_stride(t: Tensor | None) -> bool:
             return False
 
     # In deterministic mode, we can't safely benchmark - disallow padding
-    # Check this after other basic checks so force_shape_pad can override
+    # Check this after other basic checks so force_shape_pad/autoheuristic can override
     if (
         torch._inductor.config.deterministic
         and not torch._inductor.config.force_shape_pad
+        and not torch._inductor.config.use_autoheuristic("pad_mm")
     ):
         return False
 
@@ -388,11 +376,15 @@ def get_non_view_def(node: torch.fx.Node) -> torch.fx.Node:
 
 
 def should_exclude_padding_time(match: Match, arg_name: str) -> bool:
+    from torch._prims_common import is_contiguous_or_false
+
     node_def = get_non_view_def(match.kwargs[arg_name])
 
     # constant padding converts tensors to contiguous so even if the input tensor
     # can be planned layout transform is not free. TODO - way to pad and preserve layout ?
-    if not fetch_fake_tensors(match, (arg_name,))[0].is_contiguous():
+    # Use is_contiguous_or_false to avoid guarding on data-dependent expressions
+    # with unbacked symints - returns False instead of raising an error.
+    if not is_contiguous_or_false(fetch_fake_tensors(match, (arg_name,))[0]):
         return False
 
     # TODO - see issue https://github.com/pytorch/pytorch/issues/128889
@@ -519,15 +511,19 @@ def _should_pad(
         if torch._inductor.config.force_shape_pad:
             return True
 
+        # Resolve symbolic dims to concrete hints for heuristic checks below.
+        # These are performance decisions, not correctness — optimization_hint is safe.
+        m_concrete, k_concrete, n_concrete = hint_symbols((m, k, n))
+
         # Performance heuristic for bf16 large K scenarios
         if (
             "pad_aten_mm_pass" in torch._inductor.config.post_grad_fusion_options
-            and should_pad_mm_bf16(mat1.dtype, m, n, k)
+            and should_pad_mm_bf16(mat1.dtype, m_concrete, n_concrete, k_concrete)
         ):
             return True
 
         # Check if operation is compute bound (performance check)
-        if not is_mm_compute_bound(m, k, n, mat1.dtype):
+        if not is_mm_compute_bound(m_concrete, k_concrete, n_concrete, mat1.dtype):
             return False
 
         # We don't want to look up the cache for cases that are trivially false
@@ -540,9 +536,9 @@ def _should_pad(
 
         def realize_tensor(t):
             if isinstance(t, FakeTensor):
-                size_hints = realize_symbols(t.size())
+                size_hints = hint_symbols(t.size())
                 # pyrefly: ignore [bad-argument-type]
-                stride_hint = realize_symbols(t.stride())
+                stride_hint = hint_symbols(t.stride())
                 real_size = (
                     sum((d - 1) * s for d, s in zip(size_hints, stride_hint)) + 1
                 )
@@ -677,6 +673,10 @@ def pad_bench_fn():
             if ah_should_pad is not None:
                 return ah_should_pad
 
+        # AH didn't make a decision, so if we're in deterministic mode, we should return false
+        if torch._inductor.config.deterministic:
+            return False
+
         if ori_time is None:
             ori_time = do_bench(orig_bench_fn)
             set_cached_base_mm_benchmark_time(ori_time_key, ori_time)
@@ -917,7 +917,7 @@ def bmm_replace(mat1: Tensor, mat2: Tensor) -> Tensor:
 
 
 @functools.cache
-def _pad_mm_init(input_device: Optional[torch.device] = None) -> None:
+def _pad_mm_init(input_device: torch.device | None = None) -> None:
     from .joint_graph import patterns
 
     if input_device:
diff --git a/torch/_inductor/fx_passes/post_grad.py b/torch/_inductor/fx_passes/post_grad.py
index 427a6918a9cea..02ec541caf7a9 100644
--- a/torch/_inductor/fx_passes/post_grad.py
+++ b/torch/_inductor/fx_passes/post_grad.py
@@ -6,7 +6,7 @@
 import operator
 from collections import Counter, defaultdict
 from collections.abc import Callable
-from typing import Any, Optional, TypeVar
+from typing import Any, TypeVar
 from typing_extensions import ParamSpec
 
 import torch
@@ -16,6 +16,7 @@
 from torch._decomp import register_decomposition
 from torch._dynamo.utils import counters
 from torch._inductor import comms
+from torch._inductor.custom_graph_pass import CustomInferenceAwareGraphPass
 from torch._inductor.virtualized import ops  # noqa: F401
 from torch._logging import trace_structured
 from torch._prims_common import is_boolean_dtype, is_expandable_to, is_integer_dtype
@@ -137,6 +138,10 @@ def post_grad_passes(gm: torch.fx.GraphModule, is_inference: bool):
     fake_tensor_updater = FakeTensorUpdater(gm.graph)
 
     if post_grad_custom_pre_pass := config.post_grad_custom_pre_pass:
+        if isinstance(post_grad_custom_pre_pass, CustomInferenceAwareGraphPass):
+            post_grad_custom_pre_pass = functools.partial(
+                post_grad_custom_pre_pass, is_inference=is_inference
+            )
         GraphTransformObserver(gm, "post_grad_custom_pre_pass").apply_graph_pass(
             post_grad_custom_pre_pass
         )
@@ -217,6 +222,10 @@ def post_grad_passes(gm: torch.fx.GraphModule, is_inference: bool):
         )
 
     if post_grad_custom_post_pass := config.post_grad_custom_post_pass:
+        if isinstance(post_grad_custom_post_pass, CustomInferenceAwareGraphPass):
+            post_grad_custom_post_pass = functools.partial(
+                post_grad_custom_post_pass, is_inference=is_inference
+            )
         GraphTransformObserver(gm, "post_grad_custom_post_pass").apply_graph_pass(
             post_grad_custom_post_pass
         )
@@ -236,6 +245,15 @@ def post_grad_passes(gm: torch.fx.GraphModule, is_inference: bool):
                 pass_name = "custom_backend_passes_" + device
                 GraphTransformObserver(gm, pass_name).apply_gm_pass(custom_backend_pass)
 
+    # SPMD verification — before collective reordering passes.
+    if (
+        config.aten_distributed_optimizations.spmd_check
+        and _needs_spmd_graph_preservation()
+    ):
+        from torch._inductor.fx_passes.spmd_check import spmd_check
+
+        spmd_check(gm)
+
     collectives_bucketing: bool = False
 
     if config.bucket_reduce_scatters_fx != "none":
@@ -251,7 +269,7 @@ def post_grad_passes(gm: torch.fx.GraphModule, is_inference: bool):
             lambda graph: p(
                 graph.owning_module,
                 config.bucket_reduce_scatters_fx_bucket_size_determinator,
-                config.bucket_reduce_scatters_fx,  # type: ignore[arg-type]
+                config.bucket_reduce_scatters_bucket_mode,  # type: ignore[arg-type]
             )
         )
         collectives_bucketing = True
@@ -283,7 +301,7 @@ def post_grad_passes(gm: torch.fx.GraphModule, is_inference: bool):
             lambda graph: p(
                 graph.owning_module,
                 config.bucket_all_gathers_fx_bucket_size_determinator,
-                config.bucket_all_gathers_fx,  # type: ignore[arg-type]
+                config.bucket_all_gathers_bucket_mode,  # type: ignore[arg-type]
             )
         )
         collectives_bucketing = True
@@ -723,12 +741,16 @@ def body_fn(*flat_args):
 
 
 @init_once_fakemode
-def lazy_init(input_device: Optional[torch.device] = None):
+def lazy_init(input_device: torch.device | None = None):
     if torch._C._has_mkldnn:
         from . import decompose_mem_bound_mm  # noqa: F401
         from .mkldnn_fusion import _mkldnn_fusion_init
 
         _mkldnn_fusion_init()
+    else:
+        from .quantization import _register_woq_lowerings
+
+        _register_woq_lowerings()
 
     # Put this patterns in post-grad pass rather than joint-graph
     # pass since otherwise there will be perf/peak-memory regression:
@@ -764,6 +786,13 @@ def check():
         def check():
             return True
 
+    def consumes_rng_state(node: torch.fx.Node) -> bool:
+        return (
+            node.op == "call_function"
+            and isinstance(node.target, torch._ops.OpOverload)
+            and torch.Tag.nondeterministic_seeded in node.target.tags
+        )
+
     def visit(other_node):
         if (
             other_node.op == "call_function"
@@ -773,6 +802,11 @@ def visit(other_node):
             == get_mutation_region_id(graph, other_node)
             and check()
         ):
+            # Ops that consume RNG state are order-sensitive and must not be
+            # reordered during locality optimization.
+            if consumes_rng_state(other_node):
+                return
+
             # move node's producers right before it
             node.prepend(other_node)
 
@@ -1024,8 +1058,20 @@ def register_fun(cond):
     return register_fun
 
 
+def _needs_spmd_graph_preservation() -> bool:
+    """Check if SPMD graph preservation is needed for distributed overlap."""
+    return (
+        config.aten_distributed_optimizations.enable_overlap_scheduling
+        or config.reorder_for_compute_comm_overlap
+    )
+
+
 @register_noop_decomp(aten.slice)
 def slice_noop(self, dim=0, start=None, end=None, step=1):
+    if _needs_spmd_graph_preservation():
+        # Keep no-op slices so all ranks produce identical FX graphs (SPMD)
+        # with matching op counts and runtime estimations.
+        return False
     if start is None or end is None:
         return False
 
@@ -1069,6 +1115,10 @@ def repeat_noop(self, repeats):
 
 @register_noop_decomp(aten.constant_pad_nd)
 def constant_pad_nd(x, padding, fill_value=0):
+    if _needs_spmd_graph_preservation():
+        # Keep no-op pads so all ranks produce identical FX graphs (SPMD)
+        # with matching op counts and runtime estimations.
+        return False
     return all(p == 0 for p in padding)
 
 
@@ -1517,6 +1567,12 @@ def should_prefer_unfused_addmm(match):
     extra_check=should_prefer_unfused_addmm,
 )
 def unfuse_bias_add_to_pointwise(match: Match, mat1, mat2, *, inp, alpha, beta):
+    if config.keep_addmm_fused_for_half_dtypes and inp.meta["val"].dtype in (
+        torch.bfloat16,
+        torch.float16,
+    ):
+        return
+
     def repl(inp, x1, x2, alpha, beta):
         mm_result = x1 @ x2
         if alpha != 1:
@@ -1858,7 +1914,8 @@ def __call__(self, graph: fx.Graph) -> None:
                         lambda x: x
                         not in [cpu_concat, gpu_concat, gpu_split, gpu_node]
                         + unsqueezed_nodes
-                        and x.target != torch.ops.aten.copy_.default,
+                        and x.target != torch.ops.aten.copy_.default
+                        and x.target != "output",
                     )
                     last_node = gpu_node
 
@@ -1934,6 +1991,7 @@ def make_dependencies_equivalent(
                     del cpu_indeg[user]
                 elif (
                     self.allow_inputs
+                    and self.is_on_target_device(user)
                     and self.all_inputs_are_cpu_scalar_or_on_target_device(user)
                 ):
                     # this node takes only cpu scalar tensors or gpu tensors as inputs
diff --git a/torch/_inductor/fx_passes/pre_grad.py b/torch/_inductor/fx_passes/pre_grad.py
index dbe46a9c7753f..efc953f97068f 100644
--- a/torch/_inductor/fx_passes/pre_grad.py
+++ b/torch/_inductor/fx_passes/pre_grad.py
@@ -5,7 +5,6 @@
 import logging
 import types
 from collections.abc import Sequence
-from typing import Optional
 
 import torch
 import torch.nn as nn
@@ -171,7 +170,7 @@ def use_matmul_fuse_lce_replace_first_LCE(graph):
 
 
 @init_once_fakemode
-def lazy_init(input_device: Optional[torch.device] = None):
+def lazy_init(input_device: torch.device | None = None):
     from . import (  # noqa: F401  # noqa: F401
         apply_gumbel_max_trick,
         efficient_conv_bn_eval,
@@ -432,16 +431,20 @@ def remove_identity(gm: torch.fx.GraphModule) -> torch.fx.GraphModule:
     """
     Removes all identity layers from the module.
     """
-
-    class IdentityRemover(torch.fx.Transformer):
-        def call_module(self, target, args, kwargs):
-            if isinstance(self.submodules[target], nn.Identity):
-                assert len(args) == 1
-                return args[0]
-            else:
-                return super().call_module(target, args, kwargs)
-
-    return IdentityRemover(gm).transform()
+    graph = gm.graph
+    work_done = False
+    for module_name, module in gm.named_modules():
+        if type(module) is nn.Identity:
+            for node in list(graph.find_nodes(op="call_module", target=module_name)):
+                assert len(node.args) == 1
+                input_node = node.args[0]
+                node.replace_all_uses_with(input_node)
+                graph.erase_node(node)
+                work_done = True
+    if work_done:
+        graph.lint()
+        gm.recompile()
+    return gm
 
 
 def fuse_conv_bn(gm: torch.fx.GraphModule, inplace=False) -> torch.fx.GraphModule:
diff --git a/torch/_inductor/fx_passes/quantization.py b/torch/_inductor/fx_passes/quantization.py
index 7b4870fff5e02..bde5c94351931 100644
--- a/torch/_inductor/fx_passes/quantization.py
+++ b/torch/_inductor/fx_passes/quantization.py
@@ -1539,7 +1539,7 @@ def _find_first_node_in_dequant_pattern(_node):
         for user_node in user_node_list[1:]:
             _source_node = dequant_pattern_end_node
             _user_node = user_node
-            # pyrefly: ignore [bad-assignment]
+            # pyrefly: ignore [bad-assignment, non-convergent-recursion]
             while _source_node != dequant_pattern_start_node.args[0]:
                 _user_node = clone_to_new_node(graph, _source_node, _user_node)
                 _source_node = _source_node.args[0]  # type: ignore[assignment]
diff --git a/torch/_inductor/fx_passes/reduced_atomic_contention.py b/torch/_inductor/fx_passes/reduced_atomic_contention.py
index 838f1bb2d3ce7..a4961b148190c 100644
--- a/torch/_inductor/fx_passes/reduced_atomic_contention.py
+++ b/torch/_inductor/fx_passes/reduced_atomic_contention.py
@@ -8,7 +8,7 @@
 
 import logging
 import math
-from typing import Any, Optional
+from typing import Any
 
 import torch
 import torch.fx as fx
@@ -355,7 +355,7 @@ def _fit_to_memory_budget(
 
 def _extract_scatter_dim_and_index(
     indices_arg: Any,
-) -> tuple[Optional[int], Optional[fx.Node]]:
+) -> tuple[int | None, fx.Node | None]:
     """Extract scatter dimension and index node from indices argument."""
     # Case 1: Single index → dim=0
     if not isinstance(indices_arg, (list, tuple)):
@@ -377,7 +377,7 @@ def _extract_scatter_dim_and_index(
     return scatter_dim, index_node
 
 
-def _get_tensor_meta(node: fx.Node) -> Optional[dict[str, Any]]:
+def _get_tensor_meta(node: fx.Node) -> dict[str, Any] | None:
     """Extract tensor metadata from FX node."""
     if not hasattr(node, "meta") or "val" not in node.meta:
         return None
diff --git a/torch/_inductor/fx_passes/reinplace.py b/torch/_inductor/fx_passes/reinplace.py
index ceea6fb10d46c..2ca7429aabe51 100644
--- a/torch/_inductor/fx_passes/reinplace.py
+++ b/torch/_inductor/fx_passes/reinplace.py
@@ -761,15 +761,15 @@ def tensor_with_same_storage_already_reinplaced(arg):
                         mutated_tensors_flat.append(arg)
 
                 # Check if all mutated args can be inplaced
-                all_can_inplace = all_can_inplace(node, mutated_tensors_flat)
+                can_inplace_all = all_can_inplace(node, mutated_tensors_flat)
 
                 log.debug(
-                    "reinplace with_effects: mutated_tensors=%s, all_can_inplace=%s",
+                    "reinplace with_effects: mutated_tensors=%s, can_inplace_all=%s",
                     [str(a) for a in mutated_tensors_flat],
-                    all_can_inplace,
+                    can_inplace_all,
                 )
 
-                if all_can_inplace and inplaceable_op.extra_check(node):
+                if can_inplace_all and inplaceable_op.extra_check(node):
                     log.debug(
                         "reinplace with_effects: converting %s -> %s",
                         inner_op,
diff --git a/torch/_inductor/fx_passes/replace_random.py b/torch/_inductor/fx_passes/replace_random.py
index 150ba5cde4a7c..7689b64ed5334 100644
--- a/torch/_inductor/fx_passes/replace_random.py
+++ b/torch/_inductor/fx_passes/replace_random.py
@@ -3,6 +3,7 @@
 import logging
 
 import torch
+from torch.fx.experimental.symbolic_shapes import guard_or_false, statically_known_true
 from torch.fx.passes.graph_transform_observer import GraphTransformObserver
 from torch.fx.passes.shape_prop import _extract_tensor_metadata
 
@@ -21,6 +22,40 @@
 aten = torch.ops.aten
 
 
+def _shape_to_offset(shape, device: torch.device):
+    # Modified from torch/_prims/rng_prims.py:philox_rand_offset
+    nelem = 1
+    for s in shape:
+        nelem *= s
+
+    # Empty tensor: no random numbers are generated/consumed.
+    is_empty = nelem == 0
+    if statically_known_true(is_empty) or guard_or_false(is_empty):
+        return 0
+
+    if device is None:
+        device = torch.device("cpu")
+    elif isinstance(device, str):
+        device = torch.device(device)
+
+    if device.type != "cuda":
+        return 0
+
+    block_size = 256
+    unroll = 4
+    curand4_engine_calls = 4
+
+    device_property = torch.cuda.get_device_properties(device)
+
+    blocks_per_sm = device_property.max_threads_per_multi_processor // block_size
+    max_grid = device_property.multi_processor_count * blocks_per_sm
+    grid_size = (nelem + block_size - 1) // block_size
+    grid_size = -torch.sym_min(-grid_size, -1)
+    grid_size = torch.sym_min(grid_size, max_grid)
+
+    return ((nelem - 1) // (block_size * grid_size * unroll) + 1) * curand4_engine_calls
+
+
 def replace_random_passes(gm: torch.fx.GraphModule):
     """Modify the given FX graph to use backend-native random ops"""
     if config.fallback_random:
@@ -29,10 +64,59 @@ def replace_random_passes(gm: torch.fx.GraphModule):
     count = patterns.apply(gm)
     with GraphTransformObserver(gm, "fuse_seed_creation_pass", "joint_graph_passes"):
         count += fuse_seed_creation_pass(gm.graph)
+    if config.align_random_eager:
+        with GraphTransformObserver(gm, "fuse_offset_creation_pass"):
+            count += fuse_offset_creation_pass(gm.graph)
 
     return count
 
 
+def fuse_offset_creation_pass(graph: torch.fx.Graph) -> int:
+    """
+    Here offset node means seed << 32 + offset, will unpacked in lowering.py:inductor_random()
+    Horizontally fuse all the seed generation on each device
+        a = inductor_prims.rand_eager_offset(offset, dev)
+        b = inductor_prims.rand_eager_offset(offset, dev)
+    Becomes:
+        offsets = inductor_prims.rand_eager_offsets([offset1, offset2...], dev)
+        a = torch.ops.aten.select.int(offsets, 0, 0)
+        b = torch.ops.aten.select.int(offsets, 0, 1)
+    We do this because seed creation is entirely launch overhead bound.
+    """
+    device_offsets = collections.defaultdict(list)
+    for node in graph.nodes:
+        if CallFunctionVarArgs(inductor_prims.rand_eager_offset).match(node):
+            device_offsets[node.args[1]].append(node)
+
+    if not device_offsets:
+        return 0
+
+    for device, offsets in device_offsets.items():
+        with graph.inserting_before(offsets[0]):
+            offs = [n.args[0] for n in offsets]
+            combined = graph.call_function(
+                inductor_prims.rand_eager_offsets, (offs, device)
+            )
+            with V.fake_mode:
+                combined.meta["val"] = torch.empty(
+                    [len(offsets), 2], device=device, dtype=torch.int64
+                )
+                combined.meta["tensor_meta"] = _extract_tensor_metadata(
+                    combined.meta["val"]
+                )
+
+        for idx, offset in enumerate(offsets):
+            with graph.inserting_before(offset):
+                new_state = graph.call_function(
+                    torch.ops.aten.select.int, (combined, 0, idx)
+                )
+            offset.replace_all_uses_with(new_state)
+            new_state.meta.update(offset.meta)
+            graph.erase_node(offset)
+
+    return len(device_offsets)
+
+
 def fuse_seed_creation_pass(graph: torch.fx.Graph):
     """
     Horizontally fuse all the seed generation on each device
@@ -124,8 +208,32 @@ def replacement(size):
         match.output_node().target.overloadpacket  # type: ignore[union-attr]
     ]  # type: ignore[union-attr]
     device = get_device(device)
+    replacement_fn = replacement
+
+    if mode == "rand" and config.align_random_eager and device.type == "cuda":
+        # Only enable when align_random_eager is on.
+        def replacement_align(size):
+            offset = _shape_to_offset(size, device)
+
+            align_dtype = dtype
+            if isinstance(align_dtype, (tuple, list)):
+                align_dtype = align_dtype[0] if len(align_dtype) else None
+
+            result = inductor_prims.random(
+                size,
+                inductor_prims.rand_eager_offset(offset, device),
+                mode,
+                **default_kwargs(device),
+                align_dtype=align_dtype,
+            )
+            if dtype is not None:
+                result = result.to(dtype)
+            return result
+
+        replacement_fn = replacement_align
+
     # pyrefly: ignore [bad-argument-type]
-    match.replace_by_example(replacement, [size])
+    match.replace_by_example(replacement_fn, [size])
 
 
 # pyrefly: ignore [bad-argument-type]
diff --git a/torch/_inductor/fx_passes/serialized_patterns/_sfdp_pattern_10.py b/torch/_inductor/fx_passes/serialized_patterns/_sfdp_pattern_10.py
index 567390838ede7..9e6e4a9ec945c 100644
--- a/torch/_inductor/fx_passes/serialized_patterns/_sfdp_pattern_10.py
+++ b/torch/_inductor/fx_passes/serialized_patterns/_sfdp_pattern_10.py
@@ -32,7 +32,7 @@
    _TargetExprVarArgs,
 )
 permute_default = CallFunction(aten.permute.default, KeywordArg('query'), Ignored())
-div_Tensor = CallFunction(aten.div.Tensor, permute_default, Ignored())
+div_Tensor = CallFunction(aten.div.Tensor, permute_default, KeywordArg('inv_scale'))
 expand_default = CallFunction(aten.expand.default, div_Tensor, Ignored())
 clone_default = CallFunction(aten.clone.default, expand_default, memory_format=torch.contiguous_format)
 view_default = CallFunction(aten.view.default, clone_default, Ignored(), _users=2)
@@ -52,7 +52,8 @@
 expand_default_2 = CallFunction(aten.expand.default, convert_element_type_default, Ignored())
 view_default_3 = CallFunction(aten.view.default, expand_default_2, Ignored(), _users=2)
 permute_default_3 = CallFunction(aten.permute.default, KeywordArg('value'), Ignored())
-expand_default_3 = CallFunction(aten.expand.default, permute_default_3, Ignored())
+convert_element_type_default_1 = CallFunction(prims.convert_element_type.default, permute_default_3, Ignored())
+expand_default_3 = CallFunction(aten.expand.default, convert_element_type_default_1, Ignored())
 clone_default_2 = CallFunction(aten.clone.default, expand_default_3, memory_format=torch.contiguous_format)
 view_default_4 = CallFunction(aten.view.default, clone_default_2, Ignored(), _users=2)
 bmm_default_1 = CallFunction(aten.bmm.default, view_default_3, view_default_4)
@@ -61,8 +62,7 @@
 view_default_6 = CallFunction(aten.view.default, KeywordArg('tangents_1'), Ignored(), _users=2)
 permute_default_4 = CallFunction(aten.permute.default, view_default_4, Ignored())
 bmm_default_2 = CallFunction(aten.bmm.default, view_default_6, permute_default_4)
-convert_element_type_default_1 = CallFunction(prims.convert_element_type.default, bmm_default_2, Ignored())
-view_default_7 = CallFunction(aten.view.default, convert_element_type_default_1, Ignored())
+view_default_7 = CallFunction(aten.view.default, bmm_default_2, Ignored())
 convert_element_type_default_2 = CallFunction(prims.convert_element_type.default, view_default_7, Ignored())
 mul_Tensor = CallFunction(aten.mul.Tensor, convert_element_type_default_2, div_Tensor_1, _users=2)
 sum_dim_IntList_1 = CallFunction(aten.sum.dim_IntList, mul_Tensor, Ignored(), True)
@@ -71,7 +71,7 @@
 permute_default_5 = CallFunction(aten.permute.default, view_default_1, Ignored())
 bmm_default_3 = CallFunction(aten.bmm.default, view_default_8, permute_default_5)
 view_default_9 = CallFunction(aten.view.default, bmm_default_3, Ignored())
-div_Tensor_2 = CallFunction(aten.div.Tensor, view_default_9, Ignored())
+div_Tensor_2 = CallFunction(aten.div.Tensor, view_default_9, KeywordArg('inv_scale'))
 permute_default_6 = CallFunction(aten.permute.default, div_Tensor_2, Ignored())
 permute_default_7 = CallFunction(aten.permute.default, view_default, Ignored())
 bmm_default_4 = CallFunction(aten.bmm.default, permute_default_7, view_default_8)
@@ -81,16 +81,18 @@
 permute_default_10 = CallFunction(aten.permute.default, view_default_3, Ignored())
 bmm_default_5 = CallFunction(aten.bmm.default, permute_default_10, view_default_6)
 view_default_11 = CallFunction(aten.view.default, bmm_default_5, Ignored())
-permute_default_11 = CallFunction(aten.permute.default, view_default_11, Ignored())
+convert_element_type_default_3 = CallFunction(prims.convert_element_type.default, view_default_11, Ignored())
+permute_default_11 = CallFunction(aten.permute.default, convert_element_type_default_3, Ignored())
 _sfdp_pattern_10_training = MultiOutputPattern([view_default_5,
   permute_default_6,
   permute_default_9,
-  permute_default_11
+  permute_default_11,
+  None
 ])
 
 
 permute_default = CallFunction(aten.permute.default, KeywordArg('query'), Ignored())
-div_Tensor = CallFunction(aten.div.Tensor, permute_default, Ignored())
+div_Tensor = CallFunction(aten.div.Tensor, permute_default, KeywordArg('inv_scale'))
 expand_default = CallFunction(aten.expand.default, div_Tensor, Ignored())
 clone_default = CallFunction(aten.clone.default, expand_default, memory_format=torch.contiguous_format)
 view_default = CallFunction(aten.view.default, clone_default, Ignored())
@@ -110,7 +112,8 @@
 expand_default_2 = CallFunction(aten.expand.default, convert_element_type_default, Ignored())
 view_default_3 = CallFunction(aten.view.default, expand_default_2, Ignored())
 permute_default_3 = CallFunction(aten.permute.default, KeywordArg('value'), Ignored())
-expand_default_3 = CallFunction(aten.expand.default, permute_default_3, Ignored())
+convert_element_type_default_1 = CallFunction(prims.convert_element_type.default, permute_default_3, Ignored())
+expand_default_3 = CallFunction(aten.expand.default, convert_element_type_default_1, Ignored())
 clone_default_2 = CallFunction(aten.clone.default, expand_default_3, memory_format=torch.contiguous_format)
 view_default_4 = CallFunction(aten.view.default, clone_default_2, Ignored())
 bmm_default_1 = CallFunction(aten.bmm.default, view_default_3, view_default_4)
@@ -118,7 +121,7 @@
 
 
 permute_default = CallFunction(aten.permute.default, KeywordArg('query'), Ignored())
-div_Tensor = CallFunction(aten.div.Tensor, permute_default, Ignored())
+div_Tensor = CallFunction(aten.div.Tensor, permute_default, KeywordArg('inv_scale'))
 expand_default = CallFunction(aten.expand.default, div_Tensor, Ignored())
 clone_default = CallFunction(aten.clone.default, expand_default, memory_format=torch.contiguous_format)
 view_default = CallFunction(aten.view.default, clone_default, Ignored(), _users=2)
@@ -158,7 +161,7 @@
 permute_default_5 = CallFunction(aten.permute.default, view_default_1, Ignored())
 bmm_default_3 = CallFunction(aten.bmm.default, view_default_8, permute_default_5)
 view_default_9 = CallFunction(aten.view.default, bmm_default_3, Ignored())
-div_Tensor_2 = CallFunction(aten.div.Tensor, view_default_9, Ignored())
+div_Tensor_2 = CallFunction(aten.div.Tensor, view_default_9, KeywordArg('inv_scale'))
 permute_default_6 = CallFunction(aten.permute.default, div_Tensor_2, Ignored())
 permute_default_7 = CallFunction(aten.permute.default, view_default, Ignored())
 bmm_default_4 = CallFunction(aten.bmm.default, permute_default_7, view_default_8)
@@ -172,12 +175,13 @@
 _sfdp_pattern_10_half_training = MultiOutputPattern([view_default_5,
   permute_default_6,
   permute_default_9,
-  permute_default_11
+  permute_default_11,
+  None
 ])
 
 
 permute_default = CallFunction(aten.permute.default, KeywordArg('query'), Ignored())
-div_Tensor = CallFunction(aten.div.Tensor, permute_default, Ignored())
+div_Tensor = CallFunction(aten.div.Tensor, permute_default, KeywordArg('inv_scale'))
 expand_default = CallFunction(aten.expand.default, div_Tensor, Ignored())
 clone_default = CallFunction(aten.clone.default, expand_default, memory_format=torch.contiguous_format)
 view_default = CallFunction(aten.view.default, clone_default, Ignored())
diff --git a/torch/_inductor/fx_passes/serialized_patterns/_sfdp_pattern_18.py b/torch/_inductor/fx_passes/serialized_patterns/_sfdp_pattern_18.py
index 567d898ed2042..252cdb3d147ef 100644
--- a/torch/_inductor/fx_passes/serialized_patterns/_sfdp_pattern_18.py
+++ b/torch/_inductor/fx_passes/serialized_patterns/_sfdp_pattern_18.py
@@ -44,7 +44,7 @@
 view_default_1 = CallFunction(aten.view.default, clone_default_1, Ignored(), _users=2)
 bmm_default = CallFunction(aten.bmm.default, view_default, view_default_1)
 view_default_2 = CallFunction(aten.view.default, bmm_default, Ignored())
-full_default = CallFunction(aten.full.default, [], Ignored(), dtype=Ignored(), device=Ignored(), pin_memory=False, _users=2)
+full_default = CallFunction(aten.full.default, [], KeywordArg('inv_scale'), dtype=Ignored(), device=Ignored(), pin_memory=False, _users=2)
 div_Tensor = CallFunction(aten.div.Tensor, view_default_2, full_default)
 full_default_1 = CallFunction(aten.full.default, [], Ignored(), dtype=Ignored(), device=Ignored(), pin_memory=False)
 where_self = CallFunction(aten.where.self, KeywordArg('causal_mask'), div_Tensor, full_default_1, _users=2)
@@ -98,6 +98,7 @@
   permute_default_9,
   permute_default_11,
   None,
+  None,
   None
 ])
 
@@ -113,7 +114,7 @@
 view_default_1 = CallFunction(aten.view.default, clone_default_1, Ignored())
 bmm_default = CallFunction(aten.bmm.default, view_default, view_default_1)
 view_default_2 = CallFunction(aten.view.default, bmm_default, Ignored())
-full_default = CallFunction(aten.full.default, [], Ignored(), dtype=Ignored(), device=Ignored(), pin_memory=False)
+full_default = CallFunction(aten.full.default, [], KeywordArg('inv_scale'), dtype=Ignored(), device=Ignored(), pin_memory=False)
 div_Tensor = CallFunction(aten.div.Tensor, view_default_2, full_default)
 full_default_1 = CallFunction(aten.full.default, [], Ignored(), dtype=Ignored(), device=Ignored(), pin_memory=False)
 where_self = CallFunction(aten.where.self, KeywordArg('causal_mask'), div_Tensor, full_default_1, _users=2)
@@ -147,7 +148,7 @@
 view_default_1 = CallFunction(aten.view.default, expand_default_1, Ignored(), _users=2)
 bmm_default = CallFunction(aten.bmm.default, view_default, view_default_1)
 view_default_2 = CallFunction(aten.view.default, bmm_default, Ignored())
-full_default = CallFunction(aten.full.default, [], Ignored(), dtype=Ignored(), device=Ignored(), pin_memory=False, _users=2)
+full_default = CallFunction(aten.full.default, [], KeywordArg('inv_scale'), dtype=Ignored(), device=Ignored(), pin_memory=False, _users=2)
 div_Tensor = CallFunction(aten.div.Tensor, view_default_2, full_default)
 full_default_1 = CallFunction(aten.full.default, [], Ignored(), dtype=Ignored(), device=Ignored(), pin_memory=False)
 where_self = CallFunction(aten.where.self, KeywordArg('causal_mask'), div_Tensor, full_default_1, _users=2)
@@ -200,6 +201,7 @@
   permute_default_9,
   permute_default_11,
   None,
+  None,
   None
 ])
 
@@ -213,7 +215,7 @@
 view_default_1 = CallFunction(aten.view.default, expand_default_1, Ignored())
 bmm_default = CallFunction(aten.bmm.default, view_default, view_default_1)
 view_default_2 = CallFunction(aten.view.default, bmm_default, Ignored())
-full_default = CallFunction(aten.full.default, [], Ignored(), dtype=Ignored(), device=Ignored(), pin_memory=False)
+full_default = CallFunction(aten.full.default, [], KeywordArg('inv_scale'), dtype=Ignored(), device=Ignored(), pin_memory=False)
 div_Tensor = CallFunction(aten.div.Tensor, view_default_2, full_default)
 full_default_1 = CallFunction(aten.full.default, [], Ignored(), dtype=Ignored(), device=Ignored(), pin_memory=False)
 where_self = CallFunction(aten.where.self, KeywordArg('causal_mask'), div_Tensor, full_default_1, _users=2)
@@ -248,7 +250,7 @@
 view_default_1 = CallFunction(aten.view.default, clone_default_1, Ignored(), _users=2)
 bmm_default = CallFunction(aten.bmm.default, view_default, view_default_1)
 view_default_2 = CallFunction(aten.view.default, bmm_default, Ignored())
-full_default = CallFunction(aten.full.default, [], Ignored(), dtype=Ignored(), device=Ignored(), pin_memory=False, _users=2)
+full_default = CallFunction(aten.full.default, [], KeywordArg('inv_scale'), dtype=Ignored(), device=Ignored(), pin_memory=False, _users=2)
 div_Tensor = CallFunction(aten.div.Tensor, view_default_2, full_default)
 full_default_1 = CallFunction(aten.full.default, [], Ignored(), dtype=Ignored(), device=Ignored(), pin_memory=False)
 where_self = CallFunction(aten.where.self, KeywordArg('causal_mask'), div_Tensor, full_default_1)
@@ -307,6 +309,7 @@
   permute_default_9,
   permute_default_11,
   None,
+  None,
   None
 ])
 
@@ -322,7 +325,7 @@
 view_default_1 = CallFunction(aten.view.default, clone_default_1, Ignored())
 bmm_default = CallFunction(aten.bmm.default, view_default, view_default_1)
 view_default_2 = CallFunction(aten.view.default, bmm_default, Ignored())
-full_default = CallFunction(aten.full.default, [], Ignored(), dtype=Ignored(), device=Ignored(), pin_memory=False)
+full_default = CallFunction(aten.full.default, [], KeywordArg('inv_scale'), dtype=Ignored(), device=Ignored(), pin_memory=False)
 div_Tensor = CallFunction(aten.div.Tensor, view_default_2, full_default)
 full_default_1 = CallFunction(aten.full.default, [], Ignored(), dtype=Ignored(), device=Ignored(), pin_memory=False)
 where_self = CallFunction(aten.where.self, KeywordArg('causal_mask'), div_Tensor, full_default_1)
@@ -358,7 +361,7 @@
 view_default_1 = CallFunction(aten.view.default, expand_default_1, Ignored(), _users=2)
 bmm_default = CallFunction(aten.bmm.default, view_default, view_default_1)
 view_default_2 = CallFunction(aten.view.default, bmm_default, Ignored())
-full_default = CallFunction(aten.full.default, [], Ignored(), dtype=Ignored(), device=Ignored(), pin_memory=False, _users=2)
+full_default = CallFunction(aten.full.default, [], KeywordArg('inv_scale'), dtype=Ignored(), device=Ignored(), pin_memory=False, _users=2)
 div_Tensor = CallFunction(aten.div.Tensor, view_default_2, full_default)
 full_default_1 = CallFunction(aten.full.default, [], Ignored(), dtype=Ignored(), device=Ignored(), pin_memory=False)
 where_self = CallFunction(aten.where.self, KeywordArg('causal_mask'), div_Tensor, full_default_1)
@@ -416,6 +419,7 @@
   permute_default_9,
   permute_default_11,
   None,
+  None,
   None
 ])
 
@@ -429,7 +433,7 @@
 view_default_1 = CallFunction(aten.view.default, expand_default_1, Ignored())
 bmm_default = CallFunction(aten.bmm.default, view_default, view_default_1)
 view_default_2 = CallFunction(aten.view.default, bmm_default, Ignored())
-full_default = CallFunction(aten.full.default, [], Ignored(), dtype=Ignored(), device=Ignored(), pin_memory=False)
+full_default = CallFunction(aten.full.default, [], KeywordArg('inv_scale'), dtype=Ignored(), device=Ignored(), pin_memory=False)
 div_Tensor = CallFunction(aten.div.Tensor, view_default_2, full_default)
 full_default_1 = CallFunction(aten.full.default, [], Ignored(), dtype=Ignored(), device=Ignored(), pin_memory=False)
 where_self = CallFunction(aten.where.self, KeywordArg('causal_mask'), div_Tensor, full_default_1)
diff --git a/torch/_inductor/fx_passes/serialized_patterns/_sfdp_pattern_19.py b/torch/_inductor/fx_passes/serialized_patterns/_sfdp_pattern_19.py
index 5c6d316351b85..66baf24230d19 100644
--- a/torch/_inductor/fx_passes/serialized_patterns/_sfdp_pattern_19.py
+++ b/torch/_inductor/fx_passes/serialized_patterns/_sfdp_pattern_19.py
@@ -40,7 +40,7 @@
 view_default_1 = CallFunction(aten.view.default, expand_default_1, Ignored(), _users=2)
 bmm_default = CallFunction(aten.bmm.default, view_default, view_default_1)
 view_default_2 = CallFunction(aten.view.default, bmm_default, Ignored())
-full_default = CallFunction(aten.full.default, [], Ignored(), dtype=Ignored(), device=Ignored(), pin_memory=False, _users=2)
+full_default = CallFunction(aten.full.default, [], KeywordArg('inv_scale'), dtype=Ignored(), device=Ignored(), pin_memory=False, _users=2)
 div_Tensor = CallFunction(aten.div.Tensor, view_default_2, full_default)
 full_default_1 = CallFunction(aten.full.default, [], Ignored(), dtype=Ignored(), device=Ignored(), pin_memory=False)
 where_self = CallFunction(aten.where.self, KeywordArg('causal_mask'), div_Tensor, full_default_1)
@@ -89,6 +89,7 @@
   view_default_11,
   None,
   None,
+  None,
   None
 ])
 
@@ -100,7 +101,7 @@
 view_default_1 = CallFunction(aten.view.default, expand_default_1, Ignored())
 bmm_default = CallFunction(aten.bmm.default, view_default, view_default_1)
 view_default_2 = CallFunction(aten.view.default, bmm_default, Ignored())
-full_default = CallFunction(aten.full.default, [], Ignored(), dtype=Ignored(), device=Ignored(), pin_memory=False)
+full_default = CallFunction(aten.full.default, [], KeywordArg('inv_scale'), dtype=Ignored(), device=Ignored(), pin_memory=False)
 div_Tensor = CallFunction(aten.div.Tensor, view_default_2, full_default)
 full_default_1 = CallFunction(aten.full.default, [], Ignored(), dtype=Ignored(), device=Ignored(), pin_memory=False)
 where_self = CallFunction(aten.where.self, KeywordArg('causal_mask'), div_Tensor, full_default_1)
@@ -127,7 +128,7 @@
 view_default_1 = CallFunction(aten.view.default, expand_default_1, Ignored(), _users=2)
 bmm_default = CallFunction(aten.bmm.default, view_default, view_default_1)
 view_default_2 = CallFunction(aten.view.default, bmm_default, Ignored())
-full_default = CallFunction(aten.full.default, [], Ignored(), dtype=Ignored(), device=Ignored(), pin_memory=False, _users=2)
+full_default = CallFunction(aten.full.default, [], KeywordArg('inv_scale'), dtype=Ignored(), device=Ignored(), pin_memory=False, _users=2)
 div_Tensor = CallFunction(aten.div.Tensor, view_default_2, full_default)
 full_default_1 = CallFunction(aten.full.default, [], Ignored(), dtype=Ignored(), device=Ignored(), pin_memory=False)
 where_self = CallFunction(aten.where.self, KeywordArg('causal_mask'), div_Tensor, full_default_1)
@@ -179,6 +180,7 @@
   view_default_11,
   None,
   None,
+  None,
   None
 ])
 
@@ -190,7 +192,7 @@
 view_default_1 = CallFunction(aten.view.default, expand_default_1, Ignored())
 bmm_default = CallFunction(aten.bmm.default, view_default, view_default_1)
 view_default_2 = CallFunction(aten.view.default, bmm_default, Ignored())
-full_default = CallFunction(aten.full.default, [], Ignored(), dtype=Ignored(), device=Ignored(), pin_memory=False)
+full_default = CallFunction(aten.full.default, [], KeywordArg('inv_scale'), dtype=Ignored(), device=Ignored(), pin_memory=False)
 div_Tensor = CallFunction(aten.div.Tensor, view_default_2, full_default)
 full_default_1 = CallFunction(aten.full.default, [], Ignored(), dtype=Ignored(), device=Ignored(), pin_memory=False)
 where_self = CallFunction(aten.where.self, KeywordArg('causal_mask'), div_Tensor, full_default_1)
diff --git a/torch/_inductor/fx_passes/serialized_patterns/_sfdp_pattern_20.py b/torch/_inductor/fx_passes/serialized_patterns/_sfdp_pattern_20.py
index 9185aa3b1e330..207cb76b49b92 100644
--- a/torch/_inductor/fx_passes/serialized_patterns/_sfdp_pattern_20.py
+++ b/torch/_inductor/fx_passes/serialized_patterns/_sfdp_pattern_20.py
@@ -38,7 +38,7 @@
 expand_default = CallFunction(aten.expand.default, view_default, Ignored(), _users=2)
 full_default = CallFunction(aten.full.default, [], Ignored(), dtype=Ignored(), device=Ignored(), pin_memory=False)
 permute_default = CallFunction(aten.permute.default, KeywordArg('query'), Ignored())
-div_Tensor = CallFunction(aten.div.Tensor, permute_default, Ignored())
+div_Tensor = CallFunction(aten.div.Tensor, permute_default, KeywordArg('inv_scale'))
 expand_default_1 = CallFunction(aten.expand.default, div_Tensor, Ignored())
 clone_default = CallFunction(aten.clone.default, expand_default_1, memory_format=torch.contiguous_format)
 view_default_1 = CallFunction(aten.view.default, clone_default, Ignored(), _users=2)
@@ -82,7 +82,7 @@
 permute_default_5 = CallFunction(aten.permute.default, view_default_2, Ignored())
 bmm_default_3 = CallFunction(aten.bmm.default, view_default_9, permute_default_5)
 view_default_10 = CallFunction(aten.view.default, bmm_default_3, Ignored())
-div_Tensor_2 = CallFunction(aten.div.Tensor, view_default_10, Ignored())
+div_Tensor_2 = CallFunction(aten.div.Tensor, view_default_10, KeywordArg('inv_scale'))
 permute_default_6 = CallFunction(aten.permute.default, div_Tensor_2, Ignored())
 permute_default_7 = CallFunction(aten.permute.default, view_default_1, Ignored())
 bmm_default_4 = CallFunction(aten.bmm.default, permute_default_7, view_default_9)
@@ -98,6 +98,7 @@
   permute_default_9,
   permute_default_11,
   None,
+  None,
   None
 ])
 
@@ -107,7 +108,7 @@
 expand_default = CallFunction(aten.expand.default, view_default, Ignored())
 full_default = CallFunction(aten.full.default, [], Ignored(), dtype=Ignored(), device=Ignored(), pin_memory=False)
 permute_default = CallFunction(aten.permute.default, KeywordArg('query'), Ignored())
-div_Tensor = CallFunction(aten.div.Tensor, permute_default, Ignored())
+div_Tensor = CallFunction(aten.div.Tensor, permute_default, KeywordArg('inv_scale'))
 expand_default_1 = CallFunction(aten.expand.default, div_Tensor, Ignored())
 clone_default = CallFunction(aten.clone.default, expand_default_1, memory_format=torch.contiguous_format)
 view_default_1 = CallFunction(aten.view.default, clone_default, Ignored())
@@ -141,7 +142,7 @@
 expand_default = CallFunction(aten.expand.default, view_default, Ignored(), _users=2)
 full_default = CallFunction(aten.full.default, [], Ignored(), dtype=Ignored(), device=Ignored(), pin_memory=False)
 permute_default = CallFunction(aten.permute.default, KeywordArg('query'), Ignored())
-div_Tensor = CallFunction(aten.div.Tensor, permute_default, Ignored())
+div_Tensor = CallFunction(aten.div.Tensor, permute_default, KeywordArg('inv_scale'))
 expand_default_1 = CallFunction(aten.expand.default, div_Tensor, Ignored())
 clone_default = CallFunction(aten.clone.default, expand_default_1, memory_format=torch.contiguous_format)
 view_default_1 = CallFunction(aten.view.default, clone_default, Ignored(), _users=2)
@@ -190,7 +191,7 @@
 permute_default_5 = CallFunction(aten.permute.default, view_default_2, Ignored())
 bmm_default_3 = CallFunction(aten.bmm.default, view_default_9, permute_default_5)
 view_default_10 = CallFunction(aten.view.default, bmm_default_3, Ignored())
-div_Tensor_2 = CallFunction(aten.div.Tensor, view_default_10, Ignored())
+div_Tensor_2 = CallFunction(aten.div.Tensor, view_default_10, KeywordArg('inv_scale'))
 permute_default_6 = CallFunction(aten.permute.default, div_Tensor_2, Ignored())
 permute_default_7 = CallFunction(aten.permute.default, view_default_1, Ignored())
 bmm_default_4 = CallFunction(aten.bmm.default, permute_default_7, view_default_9)
@@ -206,6 +207,7 @@
   permute_default_9,
   permute_default_11,
   None,
+  None,
   None
 ])
 
@@ -215,7 +217,7 @@
 expand_default = CallFunction(aten.expand.default, view_default, Ignored())
 full_default = CallFunction(aten.full.default, [], Ignored(), dtype=Ignored(), device=Ignored(), pin_memory=False)
 permute_default = CallFunction(aten.permute.default, KeywordArg('query'), Ignored())
-div_Tensor = CallFunction(aten.div.Tensor, permute_default, Ignored())
+div_Tensor = CallFunction(aten.div.Tensor, permute_default, KeywordArg('inv_scale'))
 expand_default_1 = CallFunction(aten.expand.default, div_Tensor, Ignored())
 clone_default = CallFunction(aten.clone.default, expand_default_1, memory_format=torch.contiguous_format)
 view_default_1 = CallFunction(aten.view.default, clone_default, Ignored())
diff --git a/torch/_inductor/fx_passes/serialized_patterns/_sfdp_pattern_28.py b/torch/_inductor/fx_passes/serialized_patterns/_sfdp_pattern_28.py
new file mode 100644
index 0000000000000..f414823adcebf
--- /dev/null
+++ b/torch/_inductor/fx_passes/serialized_patterns/_sfdp_pattern_28.py
@@ -0,0 +1,202 @@
+# mypy: ignore-errors
+
+# noqa: F401, E501
+# This is an auto-generated file. Please do not modify it by hand.
+# To re-generate, run:
+# cd ~/pytorch && python torchgen/fuse/gen_patterns.py
+
+import torch
+import torch._inductor
+import operator
+
+aten = torch.ops.aten
+prims = torch.ops.prims
+
+from torch._inductor.pattern_matcher import (
+   Arg,
+   CallFunction,
+   CallFunctionVarArgs,
+   CallMethod,
+   CallMethodVarArgs,
+   CallModule,
+   CallModuleVarArgs,
+   ExclusiveKeywordArg,
+   Ignored,
+   KeywordArg,
+   ListOf,
+   MultiOutputPattern,
+   PatternExpr,
+   RepeatedExpr,
+   _TargetArgsExpr,
+   _TargetExpr,
+   _TargetExprVarArgs,
+)
+rand_default = CallFunction(aten.rand.default, Ignored(), dtype=Ignored(), device=Ignored(), pin_memory=False)
+gt_Scalar = CallFunction(aten.gt.Scalar, rand_default, KeywordArg('dropout_p'), _users=2)
+expand_default = CallFunction(aten.expand.default, KeywordArg('query'), Ignored())
+clone_default = CallFunction(aten.clone.default, expand_default, memory_format=torch.contiguous_format)
+view_default = CallFunction(aten.view.default, clone_default, Ignored(), _users=2)
+permute_default = CallFunction(aten.permute.default, KeywordArg('key'), Ignored())
+expand_default_1 = CallFunction(aten.expand.default, permute_default, Ignored())
+clone_default_1 = CallFunction(aten.clone.default, expand_default_1, memory_format=torch.contiguous_format)
+view_default_1 = CallFunction(aten.view.default, clone_default_1, Ignored(), _users=2)
+bmm_default = CallFunction(aten.bmm.default, view_default, view_default_1)
+view_default_2 = CallFunction(aten.view.default, bmm_default, Ignored())
+mul_Tensor = CallFunction(aten.mul.Tensor, view_default_2, KeywordArg('scale_factor'), _users=2)
+amax_default = CallFunction(aten.amax.default, mul_Tensor, Ignored(), True)
+sub_Tensor = CallFunction(aten.sub.Tensor, mul_Tensor, amax_default)
+exp_default = CallFunction(aten.exp.default, sub_Tensor, _users=2)
+sum_dim_IntList = CallFunction(aten.sum.dim_IntList, exp_default, Ignored(), True)
+div_Tensor = CallFunction(aten.div.Tensor, exp_default, sum_dim_IntList, _users=3)
+mul_Tensor_1 = CallFunction(aten.mul.Tensor, gt_Scalar, div_Tensor)
+mul_Tensor_2 = CallFunction(aten.mul.Tensor, mul_Tensor_1, Ignored())
+expand_default_2 = CallFunction(aten.expand.default, mul_Tensor_2, Ignored())
+view_default_3 = CallFunction(aten.view.default, expand_default_2, Ignored(), _users=2)
+expand_default_3 = CallFunction(aten.expand.default, KeywordArg('value'), Ignored())
+clone_default_2 = CallFunction(aten.clone.default, expand_default_3, memory_format=torch.contiguous_format)
+view_default_4 = CallFunction(aten.view.default, clone_default_2, Ignored(), _users=2)
+bmm_default_1 = CallFunction(aten.bmm.default, view_default_3, view_default_4)
+view_default_5 = CallFunction(aten.view.default, bmm_default_1, Ignored())
+neg_default = CallFunction(aten.neg.default, div_Tensor)
+view_default_6 = CallFunction(aten.view.default, KeywordArg('tangents_1'), Ignored(), _users=2)
+permute_default_1 = CallFunction(aten.permute.default, view_default_4, Ignored())
+bmm_default_2 = CallFunction(aten.bmm.default, view_default_6, permute_default_1)
+view_default_7 = CallFunction(aten.view.default, bmm_default_2, Ignored())
+convert_element_type_default = CallFunction(prims.convert_element_type.default, gt_Scalar, Ignored())
+mul_Tensor_3 = CallFunction(aten.mul.Tensor, convert_element_type_default, Ignored())
+mul_Tensor_4 = CallFunction(aten.mul.Tensor, view_default_7, mul_Tensor_3)
+mul_Tensor_5 = CallFunction(aten.mul.Tensor, mul_Tensor_4, div_Tensor, _users=2)
+sum_dim_IntList_1 = CallFunction(aten.sum.dim_IntList, mul_Tensor_5, Ignored(), True)
+fma_default = CallFunction(prims.fma.default, neg_default, sum_dim_IntList_1, mul_Tensor_5)
+mul_Tensor_6 = CallFunction(aten.mul.Tensor, fma_default, KeywordArg('scale_factor'))
+view_default_8 = CallFunction(aten.view.default, mul_Tensor_6, Ignored(), _users=2)
+permute_default_2 = CallFunction(aten.permute.default, view_default_1, Ignored())
+bmm_default_3 = CallFunction(aten.bmm.default, view_default_8, permute_default_2)
+view_default_9 = CallFunction(aten.view.default, bmm_default_3, Ignored())
+permute_default_3 = CallFunction(aten.permute.default, view_default, Ignored())
+bmm_default_4 = CallFunction(aten.bmm.default, permute_default_3, view_default_8)
+view_default_10 = CallFunction(aten.view.default, bmm_default_4, Ignored())
+permute_default_4 = CallFunction(aten.permute.default, view_default_10, Ignored())
+permute_default_5 = CallFunction(aten.permute.default, view_default_3, Ignored())
+bmm_default_5 = CallFunction(aten.bmm.default, permute_default_5, view_default_6)
+view_default_11 = CallFunction(aten.view.default, bmm_default_5, Ignored())
+_sfdp_pattern_28_training = MultiOutputPattern([view_default_5,
+  view_default_9,
+  permute_default_4,
+  view_default_11,
+  None,
+  None
+])
+
+
+expand_default = CallFunction(aten.expand.default, KeywordArg('query'), Ignored())
+clone_default = CallFunction(aten.clone.default, expand_default, memory_format=torch.contiguous_format)
+view_default = CallFunction(aten.view.default, clone_default, Ignored())
+permute_default = CallFunction(aten.permute.default, KeywordArg('key'), Ignored())
+expand_default_1 = CallFunction(aten.expand.default, permute_default, Ignored())
+clone_default_1 = CallFunction(aten.clone.default, expand_default_1, memory_format=torch.contiguous_format)
+view_default_1 = CallFunction(aten.view.default, clone_default_1, Ignored())
+bmm_default = CallFunction(aten.bmm.default, view_default, view_default_1)
+view_default_2 = CallFunction(aten.view.default, bmm_default, Ignored())
+mul_Tensor = CallFunction(aten.mul.Tensor, view_default_2, KeywordArg('scale_factor'), _users=2)
+amax_default = CallFunction(aten.amax.default, mul_Tensor, Ignored(), True)
+sub_Tensor = CallFunction(aten.sub.Tensor, mul_Tensor, amax_default)
+exp_default = CallFunction(aten.exp.default, sub_Tensor, _users=2)
+sum_dim_IntList = CallFunction(aten.sum.dim_IntList, exp_default, Ignored(), True)
+div_Tensor = CallFunction(aten.div.Tensor, exp_default, sum_dim_IntList)
+expand_default_2 = CallFunction(aten.expand.default, div_Tensor, Ignored())
+view_default_3 = CallFunction(aten.view.default, expand_default_2, Ignored())
+expand_default_3 = CallFunction(aten.expand.default, KeywordArg('value'), Ignored())
+clone_default_2 = CallFunction(aten.clone.default, expand_default_3, memory_format=torch.contiguous_format)
+view_default_4 = CallFunction(aten.view.default, clone_default_2, Ignored())
+bmm_default_1 = CallFunction(aten.bmm.default, view_default_3, view_default_4)
+_sfdp_pattern_28_inference = CallFunction(aten.view.default, bmm_default_1, Ignored(), _users=0)
+
+
+rand_default = CallFunction(aten.rand.default, Ignored(), dtype=Ignored(), device=Ignored(), pin_memory=False)
+gt_Scalar = CallFunction(aten.gt.Scalar, rand_default, KeywordArg('dropout_p'), _users=2)
+expand_default = CallFunction(aten.expand.default, KeywordArg('query'), Ignored())
+clone_default = CallFunction(aten.clone.default, expand_default, memory_format=torch.contiguous_format)
+view_default = CallFunction(aten.view.default, clone_default, Ignored(), _users=2)
+permute_default = CallFunction(aten.permute.default, KeywordArg('key'), Ignored())
+expand_default_1 = CallFunction(aten.expand.default, permute_default, Ignored())
+clone_default_1 = CallFunction(aten.clone.default, expand_default_1, memory_format=torch.contiguous_format)
+view_default_1 = CallFunction(aten.view.default, clone_default_1, Ignored(), _users=2)
+bmm_default = CallFunction(aten.bmm.default, view_default, view_default_1)
+view_default_2 = CallFunction(aten.view.default, bmm_default, Ignored())
+mul_Tensor = CallFunction(aten.mul.Tensor, view_default_2, KeywordArg('scale_factor'))
+convert_element_type_default = CallFunction(prims.convert_element_type.default, mul_Tensor, Ignored(), _users=2)
+amax_default = CallFunction(aten.amax.default, convert_element_type_default, Ignored(), True)
+sub_Tensor = CallFunction(aten.sub.Tensor, convert_element_type_default, amax_default)
+exp_default = CallFunction(aten.exp.default, sub_Tensor, _users=2)
+sum_dim_IntList = CallFunction(aten.sum.dim_IntList, exp_default, Ignored(), True)
+div_Tensor = CallFunction(aten.div.Tensor, exp_default, sum_dim_IntList)
+convert_element_type_default_1 = CallFunction(prims.convert_element_type.default, div_Tensor, Ignored(), _users=2)
+mul_Tensor_1 = CallFunction(aten.mul.Tensor, gt_Scalar, convert_element_type_default_1)
+mul_Tensor_2 = CallFunction(aten.mul.Tensor, mul_Tensor_1, Ignored())
+expand_default_2 = CallFunction(aten.expand.default, mul_Tensor_2, Ignored())
+view_default_3 = CallFunction(aten.view.default, expand_default_2, Ignored(), _users=2)
+expand_default_3 = CallFunction(aten.expand.default, KeywordArg('value'), Ignored())
+clone_default_2 = CallFunction(aten.clone.default, expand_default_3, memory_format=torch.contiguous_format)
+view_default_4 = CallFunction(aten.view.default, clone_default_2, Ignored(), _users=2)
+bmm_default_1 = CallFunction(aten.bmm.default, view_default_3, view_default_4)
+view_default_5 = CallFunction(aten.view.default, bmm_default_1, Ignored())
+convert_element_type_default_2 = CallFunction(prims.convert_element_type.default, convert_element_type_default_1, Ignored(), _users=2)
+neg_default = CallFunction(aten.neg.default, convert_element_type_default_2)
+view_default_6 = CallFunction(aten.view.default, KeywordArg('tangents_1'), Ignored(), _users=2)
+permute_default_1 = CallFunction(aten.permute.default, view_default_4, Ignored())
+bmm_default_2 = CallFunction(aten.bmm.default, view_default_6, permute_default_1)
+view_default_7 = CallFunction(aten.view.default, bmm_default_2, Ignored())
+convert_element_type_default_3 = CallFunction(prims.convert_element_type.default, gt_Scalar, Ignored())
+mul_Tensor_3 = CallFunction(aten.mul.Tensor, convert_element_type_default_3, Ignored())
+mul_Tensor_4 = CallFunction(aten.mul.Tensor, view_default_7, mul_Tensor_3)
+convert_element_type_default_4 = CallFunction(prims.convert_element_type.default, mul_Tensor_4, Ignored())
+mul_Tensor_5 = CallFunction(aten.mul.Tensor, convert_element_type_default_4, convert_element_type_default_2, _users=2)
+sum_dim_IntList_1 = CallFunction(aten.sum.dim_IntList, mul_Tensor_5, Ignored(), True)
+fma_default = CallFunction(prims.fma.default, neg_default, sum_dim_IntList_1, mul_Tensor_5)
+convert_element_type_default_5 = CallFunction(prims.convert_element_type.default, fma_default, Ignored())
+mul_Tensor_6 = CallFunction(aten.mul.Tensor, convert_element_type_default_5, KeywordArg('scale_factor'))
+view_default_8 = CallFunction(aten.view.default, mul_Tensor_6, Ignored(), _users=2)
+permute_default_2 = CallFunction(aten.permute.default, view_default_1, Ignored())
+bmm_default_3 = CallFunction(aten.bmm.default, view_default_8, permute_default_2)
+view_default_9 = CallFunction(aten.view.default, bmm_default_3, Ignored())
+permute_default_3 = CallFunction(aten.permute.default, view_default, Ignored())
+bmm_default_4 = CallFunction(aten.bmm.default, permute_default_3, view_default_8)
+view_default_10 = CallFunction(aten.view.default, bmm_default_4, Ignored())
+permute_default_4 = CallFunction(aten.permute.default, view_default_10, Ignored())
+permute_default_5 = CallFunction(aten.permute.default, view_default_3, Ignored())
+bmm_default_5 = CallFunction(aten.bmm.default, permute_default_5, view_default_6)
+view_default_11 = CallFunction(aten.view.default, bmm_default_5, Ignored())
+_sfdp_pattern_28_half_training = MultiOutputPattern([view_default_5,
+  view_default_9,
+  permute_default_4,
+  view_default_11,
+  None,
+  None
+])
+
+
+expand_default = CallFunction(aten.expand.default, KeywordArg('query'), Ignored())
+clone_default = CallFunction(aten.clone.default, expand_default, memory_format=torch.contiguous_format)
+view_default = CallFunction(aten.view.default, clone_default, Ignored())
+permute_default = CallFunction(aten.permute.default, KeywordArg('key'), Ignored())
+expand_default_1 = CallFunction(aten.expand.default, permute_default, Ignored())
+clone_default_1 = CallFunction(aten.clone.default, expand_default_1, memory_format=torch.contiguous_format)
+view_default_1 = CallFunction(aten.view.default, clone_default_1, Ignored())
+bmm_default = CallFunction(aten.bmm.default, view_default, view_default_1)
+view_default_2 = CallFunction(aten.view.default, bmm_default, Ignored())
+mul_Tensor = CallFunction(aten.mul.Tensor, view_default_2, KeywordArg('scale_factor'))
+convert_element_type_default = CallFunction(prims.convert_element_type.default, mul_Tensor, Ignored(), _users=2)
+amax_default = CallFunction(aten.amax.default, convert_element_type_default, Ignored(), True)
+sub_Tensor = CallFunction(aten.sub.Tensor, convert_element_type_default, amax_default)
+exp_default = CallFunction(aten.exp.default, sub_Tensor, _users=2)
+sum_dim_IntList = CallFunction(aten.sum.dim_IntList, exp_default, Ignored(), True)
+div_Tensor = CallFunction(aten.div.Tensor, exp_default, sum_dim_IntList)
+convert_element_type_default_1 = CallFunction(prims.convert_element_type.default, div_Tensor, Ignored())
+expand_default_2 = CallFunction(aten.expand.default, convert_element_type_default_1, Ignored())
+view_default_3 = CallFunction(aten.view.default, expand_default_2, Ignored())
+expand_default_3 = CallFunction(aten.expand.default, KeywordArg('value'), Ignored())
+clone_default_2 = CallFunction(aten.clone.default, expand_default_3, memory_format=torch.contiguous_format)
+view_default_4 = CallFunction(aten.view.default, clone_default_2, Ignored())
+bmm_default_1 = CallFunction(aten.bmm.default, view_default_3, view_default_4)
+_sfdp_pattern_28_half_inference = CallFunction(aten.view.default, bmm_default_1, Ignored(), _users=0)
diff --git a/torch/_inductor/fx_passes/serialized_patterns/_sfdp_pattern_5.py b/torch/_inductor/fx_passes/serialized_patterns/_sfdp_pattern_5.py
index f211e56b17a0a..a36ff3ae712fb 100644
--- a/torch/_inductor/fx_passes/serialized_patterns/_sfdp_pattern_5.py
+++ b/torch/_inductor/fx_passes/serialized_patterns/_sfdp_pattern_5.py
@@ -38,7 +38,7 @@
 view_default_1 = CallFunction(aten.view.default, expand_default_1, Ignored(), _users=2)
 bmm_default = CallFunction(aten.bmm.default, view_default, view_default_1)
 view_default_2 = CallFunction(aten.view.default, bmm_default, Ignored())
-div_Tensor = CallFunction(aten.div.Tensor, view_default_2, Ignored())
+div_Tensor = CallFunction(aten.div.Tensor, view_default_2, KeywordArg('inv_scale'))
 add_Tensor = CallFunction(aten.add.Tensor, div_Tensor, KeywordArg('attn_mask'), _users=2)
 amax_default = CallFunction(aten.amax.default, add_Tensor, Ignored(), True)
 sub_Tensor = CallFunction(aten.sub.Tensor, add_Tensor, amax_default)
@@ -59,7 +59,7 @@
 mul_Tensor = CallFunction(aten.mul.Tensor, view_default_7, div_Tensor_1, _users=2)
 sum_dim_IntList_1 = CallFunction(aten.sum.dim_IntList, mul_Tensor, Ignored(), True)
 fma_default = CallFunction(prims.fma.default, neg_default, sum_dim_IntList_1, mul_Tensor)
-div_Tensor_2 = CallFunction(aten.div.Tensor, fma_default, Ignored())
+div_Tensor_2 = CallFunction(aten.div.Tensor, fma_default, KeywordArg('inv_scale'))
 view_default_8 = CallFunction(aten.view.default, div_Tensor_2, Ignored(), _users=2)
 permute_default_2 = CallFunction(aten.permute.default, view_default_1, Ignored())
 bmm_default_3 = CallFunction(aten.bmm.default, view_default_8, permute_default_2)
@@ -75,6 +75,7 @@
   view_default_9,
   permute_default_4,
   view_default_11,
+  None,
   None
 ])
 
@@ -86,7 +87,7 @@
 view_default_1 = CallFunction(aten.view.default, expand_default_1, Ignored())
 bmm_default = CallFunction(aten.bmm.default, view_default, view_default_1)
 view_default_2 = CallFunction(aten.view.default, bmm_default, Ignored())
-div_Tensor = CallFunction(aten.div.Tensor, view_default_2, Ignored())
+div_Tensor = CallFunction(aten.div.Tensor, view_default_2, KeywordArg('inv_scale'))
 add_Tensor = CallFunction(aten.add.Tensor, div_Tensor, KeywordArg('attn_mask'), _users=2)
 amax_default = CallFunction(aten.amax.default, add_Tensor, Ignored(), True)
 sub_Tensor = CallFunction(aten.sub.Tensor, add_Tensor, amax_default)
@@ -108,7 +109,7 @@
 view_default_1 = CallFunction(aten.view.default, expand_default_1, Ignored(), _users=2)
 bmm_default = CallFunction(aten.bmm.default, view_default, view_default_1)
 view_default_2 = CallFunction(aten.view.default, bmm_default, Ignored())
-div_Tensor = CallFunction(aten.div.Tensor, view_default_2, Ignored())
+div_Tensor = CallFunction(aten.div.Tensor, view_default_2, KeywordArg('inv_scale'))
 add_Tensor = CallFunction(aten.add.Tensor, div_Tensor, KeywordArg('attn_mask'))
 convert_element_type_default = CallFunction(prims.convert_element_type.default, add_Tensor, Ignored(), _users=2)
 amax_default = CallFunction(aten.amax.default, convert_element_type_default, Ignored(), True)
@@ -134,7 +135,7 @@
 sum_dim_IntList_1 = CallFunction(aten.sum.dim_IntList, mul_Tensor, Ignored(), True)
 fma_default = CallFunction(prims.fma.default, neg_default, sum_dim_IntList_1, mul_Tensor)
 convert_element_type_default_4 = CallFunction(prims.convert_element_type.default, fma_default, Ignored())
-div_Tensor_2 = CallFunction(aten.div.Tensor, convert_element_type_default_4, Ignored())
+div_Tensor_2 = CallFunction(aten.div.Tensor, convert_element_type_default_4, KeywordArg('inv_scale'))
 view_default_8 = CallFunction(aten.view.default, div_Tensor_2, Ignored(), _users=2)
 permute_default_2 = CallFunction(aten.permute.default, view_default_1, Ignored())
 bmm_default_3 = CallFunction(aten.bmm.default, view_default_8, permute_default_2)
@@ -150,6 +151,7 @@
   view_default_9,
   permute_default_4,
   view_default_11,
+  None,
   None
 ])
 
@@ -161,7 +163,7 @@
 view_default_1 = CallFunction(aten.view.default, expand_default_1, Ignored())
 bmm_default = CallFunction(aten.bmm.default, view_default, view_default_1)
 view_default_2 = CallFunction(aten.view.default, bmm_default, Ignored())
-div_Tensor = CallFunction(aten.div.Tensor, view_default_2, Ignored())
+div_Tensor = CallFunction(aten.div.Tensor, view_default_2, KeywordArg('inv_scale'))
 add_Tensor = CallFunction(aten.add.Tensor, div_Tensor, KeywordArg('attn_mask'))
 convert_element_type_default = CallFunction(prims.convert_element_type.default, add_Tensor, Ignored(), _users=2)
 amax_default = CallFunction(aten.amax.default, convert_element_type_default, Ignored(), True)
diff --git a/torch/_inductor/fx_passes/serialized_patterns/_sfdp_pattern_6.py b/torch/_inductor/fx_passes/serialized_patterns/_sfdp_pattern_6.py
index 01304bf415163..48e3b33dfecc8 100644
--- a/torch/_inductor/fx_passes/serialized_patterns/_sfdp_pattern_6.py
+++ b/torch/_inductor/fx_passes/serialized_patterns/_sfdp_pattern_6.py
@@ -40,7 +40,7 @@
 view_default_1 = CallFunction(aten.view.default, expand_default_1, Ignored(), _users=2)
 bmm_default = CallFunction(aten.bmm.default, view_default, view_default_1)
 view_default_2 = CallFunction(aten.view.default, bmm_default, Ignored())
-div_Tensor = CallFunction(aten.div.Tensor, view_default_2, Ignored())
+div_Tensor = CallFunction(aten.div.Tensor, view_default_2, KeywordArg('inv_scale'))
 add_Tensor = CallFunction(aten.add.Tensor, div_Tensor, KeywordArg('attn_mask'), _users=2)
 amax_default = CallFunction(aten.amax.default, add_Tensor, Ignored(), True)
 sub_Tensor = CallFunction(aten.sub.Tensor, add_Tensor, amax_default)
@@ -66,7 +66,7 @@
 mul_Tensor_4 = CallFunction(aten.mul.Tensor, mul_Tensor_3, div_Tensor_1, _users=2)
 sum_dim_IntList_1 = CallFunction(aten.sum.dim_IntList, mul_Tensor_4, Ignored(), True)
 fma_default = CallFunction(prims.fma.default, neg_default, sum_dim_IntList_1, mul_Tensor_4)
-div_Tensor_2 = CallFunction(aten.div.Tensor, fma_default, Ignored())
+div_Tensor_2 = CallFunction(aten.div.Tensor, fma_default, KeywordArg('inv_scale'))
 view_default_8 = CallFunction(aten.view.default, div_Tensor_2, Ignored(), _users=2)
 permute_default_2 = CallFunction(aten.permute.default, view_default_1, Ignored())
 bmm_default_3 = CallFunction(aten.bmm.default, view_default_8, permute_default_2)
@@ -83,6 +83,7 @@
   permute_default_4,
   view_default_11,
   None,
+  None,
   None
 ])
 
@@ -94,7 +95,7 @@
 view_default_1 = CallFunction(aten.view.default, expand_default_1, Ignored())
 bmm_default = CallFunction(aten.bmm.default, view_default, view_default_1)
 view_default_2 = CallFunction(aten.view.default, bmm_default, Ignored())
-div_Tensor = CallFunction(aten.div.Tensor, view_default_2, Ignored())
+div_Tensor = CallFunction(aten.div.Tensor, view_default_2, KeywordArg('inv_scale'))
 add_Tensor = CallFunction(aten.add.Tensor, div_Tensor, KeywordArg('attn_mask'), _users=2)
 amax_default = CallFunction(aten.amax.default, add_Tensor, Ignored(), True)
 sub_Tensor = CallFunction(aten.sub.Tensor, add_Tensor, amax_default)
@@ -118,7 +119,7 @@
 view_default_1 = CallFunction(aten.view.default, expand_default_1, Ignored(), _users=2)
 bmm_default = CallFunction(aten.bmm.default, view_default, view_default_1)
 view_default_2 = CallFunction(aten.view.default, bmm_default, Ignored())
-div_Tensor = CallFunction(aten.div.Tensor, view_default_2, Ignored())
+div_Tensor = CallFunction(aten.div.Tensor, view_default_2, KeywordArg('inv_scale'))
 add_Tensor = CallFunction(aten.add.Tensor, div_Tensor, KeywordArg('attn_mask'))
 convert_element_type_default = CallFunction(prims.convert_element_type.default, add_Tensor, Ignored(), _users=2)
 amax_default = CallFunction(aten.amax.default, convert_element_type_default, Ignored(), True)
@@ -149,7 +150,7 @@
 sum_dim_IntList_1 = CallFunction(aten.sum.dim_IntList, mul_Tensor_4, Ignored(), True)
 fma_default = CallFunction(prims.fma.default, neg_default, sum_dim_IntList_1, mul_Tensor_4)
 convert_element_type_default_5 = CallFunction(prims.convert_element_type.default, fma_default, Ignored())
-div_Tensor_2 = CallFunction(aten.div.Tensor, convert_element_type_default_5, Ignored())
+div_Tensor_2 = CallFunction(aten.div.Tensor, convert_element_type_default_5, KeywordArg('inv_scale'))
 view_default_8 = CallFunction(aten.view.default, div_Tensor_2, Ignored(), _users=2)
 permute_default_2 = CallFunction(aten.permute.default, view_default_1, Ignored())
 bmm_default_3 = CallFunction(aten.bmm.default, view_default_8, permute_default_2)
@@ -166,6 +167,7 @@
   permute_default_4,
   view_default_11,
   None,
+  None,
   None
 ])
 
@@ -177,7 +179,7 @@
 view_default_1 = CallFunction(aten.view.default, expand_default_1, Ignored())
 bmm_default = CallFunction(aten.bmm.default, view_default, view_default_1)
 view_default_2 = CallFunction(aten.view.default, bmm_default, Ignored())
-div_Tensor = CallFunction(aten.div.Tensor, view_default_2, Ignored())
+div_Tensor = CallFunction(aten.div.Tensor, view_default_2, KeywordArg('inv_scale'))
 add_Tensor = CallFunction(aten.add.Tensor, div_Tensor, KeywordArg('attn_mask'))
 convert_element_type_default = CallFunction(prims.convert_element_type.default, add_Tensor, Ignored(), _users=2)
 amax_default = CallFunction(aten.amax.default, convert_element_type_default, Ignored(), True)
diff --git a/torch/_inductor/fx_passes/serialized_patterns/_sfdp_pattern_7.py b/torch/_inductor/fx_passes/serialized_patterns/_sfdp_pattern_7.py
index b463c7e64a613..4f42972a0d383 100644
--- a/torch/_inductor/fx_passes/serialized_patterns/_sfdp_pattern_7.py
+++ b/torch/_inductor/fx_passes/serialized_patterns/_sfdp_pattern_7.py
@@ -44,7 +44,7 @@
 view_default_1 = CallFunction(aten.view.default, clone_default_1, Ignored(), _users=2)
 bmm_default = CallFunction(aten.bmm.default, view_default, view_default_1)
 view_default_2 = CallFunction(aten.view.default, bmm_default, Ignored())
-div_Tensor = CallFunction(aten.div.Tensor, view_default_2, Ignored(), _users=2)
+div_Tensor = CallFunction(aten.div.Tensor, view_default_2, KeywordArg('inv_scale'), _users=2)
 amax_default = CallFunction(aten.amax.default, div_Tensor, Ignored(), True)
 sub_Tensor = CallFunction(aten.sub.Tensor, div_Tensor, amax_default)
 exp_default = CallFunction(aten.exp.default, sub_Tensor, _users=2)
@@ -56,7 +56,8 @@
 expand_default_2 = CallFunction(aten.expand.default, convert_element_type_default, Ignored())
 view_default_3 = CallFunction(aten.view.default, expand_default_2, Ignored(), _users=2)
 permute_default_3 = CallFunction(aten.permute.default, KeywordArg('value'), Ignored())
-expand_default_3 = CallFunction(aten.expand.default, permute_default_3, Ignored())
+convert_element_type_default_1 = CallFunction(prims.convert_element_type.default, permute_default_3, Ignored())
+expand_default_3 = CallFunction(aten.expand.default, convert_element_type_default_1, Ignored())
 clone_default_2 = CallFunction(aten.clone.default, expand_default_3, memory_format=torch.contiguous_format)
 view_default_4 = CallFunction(aten.view.default, clone_default_2, Ignored(), _users=2)
 bmm_default_1 = CallFunction(aten.bmm.default, view_default_3, view_default_4)
@@ -65,8 +66,7 @@
 view_default_6 = CallFunction(aten.view.default, KeywordArg('tangents_1'), Ignored(), _users=2)
 permute_default_4 = CallFunction(aten.permute.default, view_default_4, Ignored())
 bmm_default_2 = CallFunction(aten.bmm.default, view_default_6, permute_default_4)
-convert_element_type_default_1 = CallFunction(prims.convert_element_type.default, bmm_default_2, Ignored())
-view_default_7 = CallFunction(aten.view.default, convert_element_type_default_1, Ignored())
+view_default_7 = CallFunction(aten.view.default, bmm_default_2, Ignored())
 convert_element_type_default_2 = CallFunction(prims.convert_element_type.default, view_default_7, Ignored())
 convert_element_type_default_3 = CallFunction(prims.convert_element_type.default, gt_Scalar, Ignored())
 mul_Tensor_2 = CallFunction(aten.mul.Tensor, convert_element_type_default_3, Ignored())
@@ -74,7 +74,7 @@
 mul_Tensor_4 = CallFunction(aten.mul.Tensor, mul_Tensor_3, div_Tensor_1, _users=2)
 sum_dim_IntList_1 = CallFunction(aten.sum.dim_IntList, mul_Tensor_4, Ignored(), True)
 fma_default = CallFunction(prims.fma.default, neg_default, sum_dim_IntList_1, mul_Tensor_4)
-div_Tensor_2 = CallFunction(aten.div.Tensor, fma_default, Ignored())
+div_Tensor_2 = CallFunction(aten.div.Tensor, fma_default, KeywordArg('inv_scale'))
 view_default_8 = CallFunction(aten.view.default, div_Tensor_2, Ignored(), _users=2)
 permute_default_5 = CallFunction(aten.permute.default, view_default_1, Ignored())
 bmm_default_3 = CallFunction(aten.bmm.default, view_default_8, permute_default_5)
@@ -88,11 +88,13 @@
 permute_default_10 = CallFunction(aten.permute.default, view_default_3, Ignored())
 bmm_default_5 = CallFunction(aten.bmm.default, permute_default_10, view_default_6)
 view_default_11 = CallFunction(aten.view.default, bmm_default_5, Ignored())
-permute_default_11 = CallFunction(aten.permute.default, view_default_11, Ignored())
+convert_element_type_default_4 = CallFunction(prims.convert_element_type.default, view_default_11, Ignored())
+permute_default_11 = CallFunction(aten.permute.default, convert_element_type_default_4, Ignored())
 _sfdp_pattern_7_training = MultiOutputPattern([view_default_5,
   permute_default_6,
   permute_default_9,
   permute_default_11,
+  None,
   None
 ])
 
@@ -108,7 +110,7 @@
 view_default_1 = CallFunction(aten.view.default, clone_default_1, Ignored())
 bmm_default = CallFunction(aten.bmm.default, view_default, view_default_1)
 view_default_2 = CallFunction(aten.view.default, bmm_default, Ignored())
-div_Tensor = CallFunction(aten.div.Tensor, view_default_2, Ignored(), _users=2)
+div_Tensor = CallFunction(aten.div.Tensor, view_default_2, KeywordArg('inv_scale'), _users=2)
 amax_default = CallFunction(aten.amax.default, div_Tensor, Ignored(), True)
 sub_Tensor = CallFunction(aten.sub.Tensor, div_Tensor, amax_default)
 exp_default = CallFunction(aten.exp.default, sub_Tensor, _users=2)
@@ -118,7 +120,8 @@
 expand_default_2 = CallFunction(aten.expand.default, convert_element_type_default, Ignored())
 view_default_3 = CallFunction(aten.view.default, expand_default_2, Ignored())
 permute_default_3 = CallFunction(aten.permute.default, KeywordArg('value'), Ignored())
-expand_default_3 = CallFunction(aten.expand.default, permute_default_3, Ignored())
+convert_element_type_default_1 = CallFunction(prims.convert_element_type.default, permute_default_3, Ignored())
+expand_default_3 = CallFunction(aten.expand.default, convert_element_type_default_1, Ignored())
 clone_default_2 = CallFunction(aten.clone.default, expand_default_3, memory_format=torch.contiguous_format)
 view_default_4 = CallFunction(aten.view.default, clone_default_2, Ignored())
 bmm_default_1 = CallFunction(aten.bmm.default, view_default_3, view_default_4)
@@ -138,7 +141,7 @@
 view_default_1 = CallFunction(aten.view.default, clone_default_1, Ignored(), _users=2)
 bmm_default = CallFunction(aten.bmm.default, view_default, view_default_1)
 view_default_2 = CallFunction(aten.view.default, bmm_default, Ignored())
-div_Tensor = CallFunction(aten.div.Tensor, view_default_2, Ignored())
+div_Tensor = CallFunction(aten.div.Tensor, view_default_2, KeywordArg('inv_scale'))
 convert_element_type_default = CallFunction(prims.convert_element_type.default, div_Tensor, Ignored(), _users=2)
 amax_default = CallFunction(aten.amax.default, convert_element_type_default, Ignored(), True)
 sub_Tensor = CallFunction(aten.sub.Tensor, convert_element_type_default, amax_default)
@@ -169,7 +172,7 @@
 sum_dim_IntList_1 = CallFunction(aten.sum.dim_IntList, mul_Tensor_4, Ignored(), True)
 fma_default = CallFunction(prims.fma.default, neg_default, sum_dim_IntList_1, mul_Tensor_4)
 convert_element_type_default_4 = CallFunction(prims.convert_element_type.default, fma_default, Ignored())
-div_Tensor_2 = CallFunction(aten.div.Tensor, convert_element_type_default_4, Ignored())
+div_Tensor_2 = CallFunction(aten.div.Tensor, convert_element_type_default_4, KeywordArg('inv_scale'))
 view_default_8 = CallFunction(aten.view.default, div_Tensor_2, Ignored(), _users=2)
 permute_default_5 = CallFunction(aten.permute.default, view_default_1, Ignored())
 bmm_default_3 = CallFunction(aten.bmm.default, view_default_8, permute_default_5)
@@ -188,6 +191,7 @@
   permute_default_6,
   permute_default_9,
   permute_default_11,
+  None,
   None
 ])
 
@@ -203,7 +207,7 @@
 view_default_1 = CallFunction(aten.view.default, clone_default_1, Ignored())
 bmm_default = CallFunction(aten.bmm.default, view_default, view_default_1)
 view_default_2 = CallFunction(aten.view.default, bmm_default, Ignored())
-div_Tensor = CallFunction(aten.div.Tensor, view_default_2, Ignored())
+div_Tensor = CallFunction(aten.div.Tensor, view_default_2, KeywordArg('inv_scale'))
 convert_element_type_default = CallFunction(prims.convert_element_type.default, div_Tensor, Ignored(), _users=2)
 amax_default = CallFunction(aten.amax.default, convert_element_type_default, Ignored(), True)
 sub_Tensor = CallFunction(aten.sub.Tensor, convert_element_type_default, amax_default)
diff --git a/torch/_inductor/fx_passes/serialized_patterns/_sfdp_pattern_8.py b/torch/_inductor/fx_passes/serialized_patterns/_sfdp_pattern_8.py
index 3faff67089b17..711f68cc6189e 100644
--- a/torch/_inductor/fx_passes/serialized_patterns/_sfdp_pattern_8.py
+++ b/torch/_inductor/fx_passes/serialized_patterns/_sfdp_pattern_8.py
@@ -42,7 +42,7 @@
 view_default_1 = CallFunction(aten.view.default, clone_default_1, Ignored(), _users=2)
 bmm_default = CallFunction(aten.bmm.default, view_default, view_default_1)
 view_default_2 = CallFunction(aten.view.default, bmm_default, Ignored())
-div_Tensor = CallFunction(aten.div.Tensor, view_default_2, Ignored(), _users=2)
+div_Tensor = CallFunction(aten.div.Tensor, view_default_2, KeywordArg('inv_scale'), _users=2)
 amax_default = CallFunction(aten.amax.default, div_Tensor, Ignored(), True)
 sub_Tensor = CallFunction(aten.sub.Tensor, div_Tensor, amax_default)
 exp_default = CallFunction(aten.exp.default, sub_Tensor, _users=2)
@@ -52,7 +52,8 @@
 expand_default_2 = CallFunction(aten.expand.default, convert_element_type_default, Ignored())
 view_default_3 = CallFunction(aten.view.default, expand_default_2, Ignored(), _users=2)
 permute_default_3 = CallFunction(aten.permute.default, KeywordArg('value'), Ignored())
-expand_default_3 = CallFunction(aten.expand.default, permute_default_3, Ignored())
+convert_element_type_default_1 = CallFunction(prims.convert_element_type.default, permute_default_3, Ignored())
+expand_default_3 = CallFunction(aten.expand.default, convert_element_type_default_1, Ignored())
 clone_default_2 = CallFunction(aten.clone.default, expand_default_3, memory_format=torch.contiguous_format)
 view_default_4 = CallFunction(aten.view.default, clone_default_2, Ignored(), _users=2)
 bmm_default_1 = CallFunction(aten.bmm.default, view_default_3, view_default_4)
@@ -61,13 +62,12 @@
 view_default_6 = CallFunction(aten.view.default, KeywordArg('tangents_1'), Ignored(), _users=2)
 permute_default_4 = CallFunction(aten.permute.default, view_default_4, Ignored())
 bmm_default_2 = CallFunction(aten.bmm.default, view_default_6, permute_default_4)
-convert_element_type_default_1 = CallFunction(prims.convert_element_type.default, bmm_default_2, Ignored())
-view_default_7 = CallFunction(aten.view.default, convert_element_type_default_1, Ignored())
+view_default_7 = CallFunction(aten.view.default, bmm_default_2, Ignored())
 convert_element_type_default_2 = CallFunction(prims.convert_element_type.default, view_default_7, Ignored())
 mul_Tensor = CallFunction(aten.mul.Tensor, convert_element_type_default_2, div_Tensor_1, _users=2)
 sum_dim_IntList_1 = CallFunction(aten.sum.dim_IntList, mul_Tensor, Ignored(), True)
 fma_default = CallFunction(prims.fma.default, neg_default, sum_dim_IntList_1, mul_Tensor)
-div_Tensor_2 = CallFunction(aten.div.Tensor, fma_default, Ignored())
+div_Tensor_2 = CallFunction(aten.div.Tensor, fma_default, KeywordArg('inv_scale'))
 view_default_8 = CallFunction(aten.view.default, div_Tensor_2, Ignored(), _users=2)
 permute_default_5 = CallFunction(aten.permute.default, view_default_1, Ignored())
 bmm_default_3 = CallFunction(aten.bmm.default, view_default_8, permute_default_5)
@@ -81,11 +81,13 @@
 permute_default_10 = CallFunction(aten.permute.default, view_default_3, Ignored())
 bmm_default_5 = CallFunction(aten.bmm.default, permute_default_10, view_default_6)
 view_default_11 = CallFunction(aten.view.default, bmm_default_5, Ignored())
-permute_default_11 = CallFunction(aten.permute.default, view_default_11, Ignored())
+convert_element_type_default_3 = CallFunction(prims.convert_element_type.default, view_default_11, Ignored())
+permute_default_11 = CallFunction(aten.permute.default, convert_element_type_default_3, Ignored())
 _sfdp_pattern_8_training = MultiOutputPattern([view_default_5,
   permute_default_6,
   permute_default_9,
-  permute_default_11
+  permute_default_11,
+  None
 ])
 
 
@@ -100,7 +102,7 @@
 view_default_1 = CallFunction(aten.view.default, clone_default_1, Ignored())
 bmm_default = CallFunction(aten.bmm.default, view_default, view_default_1)
 view_default_2 = CallFunction(aten.view.default, bmm_default, Ignored())
-div_Tensor = CallFunction(aten.div.Tensor, view_default_2, Ignored(), _users=2)
+div_Tensor = CallFunction(aten.div.Tensor, view_default_2, KeywordArg('inv_scale'), _users=2)
 amax_default = CallFunction(aten.amax.default, div_Tensor, Ignored(), True)
 sub_Tensor = CallFunction(aten.sub.Tensor, div_Tensor, amax_default)
 exp_default = CallFunction(aten.exp.default, sub_Tensor, _users=2)
@@ -110,7 +112,8 @@
 expand_default_2 = CallFunction(aten.expand.default, convert_element_type_default, Ignored())
 view_default_3 = CallFunction(aten.view.default, expand_default_2, Ignored())
 permute_default_3 = CallFunction(aten.permute.default, KeywordArg('value'), Ignored())
-expand_default_3 = CallFunction(aten.expand.default, permute_default_3, Ignored())
+convert_element_type_default_1 = CallFunction(prims.convert_element_type.default, permute_default_3, Ignored())
+expand_default_3 = CallFunction(aten.expand.default, convert_element_type_default_1, Ignored())
 clone_default_2 = CallFunction(aten.clone.default, expand_default_3, memory_format=torch.contiguous_format)
 view_default_4 = CallFunction(aten.view.default, clone_default_2, Ignored())
 bmm_default_1 = CallFunction(aten.bmm.default, view_default_3, view_default_4)
@@ -128,7 +131,7 @@
 view_default_1 = CallFunction(aten.view.default, clone_default_1, Ignored(), _users=2)
 bmm_default = CallFunction(aten.bmm.default, view_default, view_default_1)
 view_default_2 = CallFunction(aten.view.default, bmm_default, Ignored())
-div_Tensor = CallFunction(aten.div.Tensor, view_default_2, Ignored())
+div_Tensor = CallFunction(aten.div.Tensor, view_default_2, KeywordArg('inv_scale'))
 convert_element_type_default = CallFunction(prims.convert_element_type.default, div_Tensor, Ignored(), _users=2)
 amax_default = CallFunction(aten.amax.default, convert_element_type_default, Ignored(), True)
 sub_Tensor = CallFunction(aten.sub.Tensor, convert_element_type_default, amax_default)
@@ -154,7 +157,7 @@
 sum_dim_IntList_1 = CallFunction(aten.sum.dim_IntList, mul_Tensor, Ignored(), True)
 fma_default = CallFunction(prims.fma.default, neg_default, sum_dim_IntList_1, mul_Tensor)
 convert_element_type_default_3 = CallFunction(prims.convert_element_type.default, fma_default, Ignored())
-div_Tensor_2 = CallFunction(aten.div.Tensor, convert_element_type_default_3, Ignored())
+div_Tensor_2 = CallFunction(aten.div.Tensor, convert_element_type_default_3, KeywordArg('inv_scale'))
 view_default_8 = CallFunction(aten.view.default, div_Tensor_2, Ignored(), _users=2)
 permute_default_5 = CallFunction(aten.permute.default, view_default_1, Ignored())
 bmm_default_3 = CallFunction(aten.bmm.default, view_default_8, permute_default_5)
@@ -172,7 +175,8 @@
 _sfdp_pattern_8_half_training = MultiOutputPattern([view_default_5,
   permute_default_6,
   permute_default_9,
-  permute_default_11
+  permute_default_11,
+  None
 ])
 
 
@@ -187,7 +191,7 @@
 view_default_1 = CallFunction(aten.view.default, clone_default_1, Ignored())
 bmm_default = CallFunction(aten.bmm.default, view_default, view_default_1)
 view_default_2 = CallFunction(aten.view.default, bmm_default, Ignored())
-div_Tensor = CallFunction(aten.div.Tensor, view_default_2, Ignored())
+div_Tensor = CallFunction(aten.div.Tensor, view_default_2, KeywordArg('inv_scale'))
 convert_element_type_default = CallFunction(prims.convert_element_type.default, div_Tensor, Ignored(), _users=2)
 amax_default = CallFunction(aten.amax.default, convert_element_type_default, Ignored(), True)
 sub_Tensor = CallFunction(aten.sub.Tensor, convert_element_type_default, amax_default)
diff --git a/torch/_inductor/fx_passes/serialized_patterns/_sfdp_pattern_9.py b/torch/_inductor/fx_passes/serialized_patterns/_sfdp_pattern_9.py
index 3bf77120e836a..599851e7af012 100644
--- a/torch/_inductor/fx_passes/serialized_patterns/_sfdp_pattern_9.py
+++ b/torch/_inductor/fx_passes/serialized_patterns/_sfdp_pattern_9.py
@@ -34,7 +34,7 @@
 rand_default = CallFunction(aten.rand.default, Ignored(), dtype=Ignored(), device=Ignored(), pin_memory=False)
 gt_Scalar = CallFunction(aten.gt.Scalar, rand_default, KeywordArg('dropout_p'), _users=2)
 permute_default = CallFunction(aten.permute.default, KeywordArg('query'), Ignored())
-div_Tensor = CallFunction(aten.div.Tensor, permute_default, Ignored())
+div_Tensor = CallFunction(aten.div.Tensor, permute_default, KeywordArg('inv_scale'))
 expand_default = CallFunction(aten.expand.default, div_Tensor, Ignored())
 clone_default = CallFunction(aten.clone.default, expand_default, memory_format=torch.contiguous_format)
 view_default = CallFunction(aten.view.default, clone_default, Ignored(), _users=2)
@@ -56,7 +56,8 @@
 expand_default_2 = CallFunction(aten.expand.default, convert_element_type_default, Ignored())
 view_default_3 = CallFunction(aten.view.default, expand_default_2, Ignored(), _users=2)
 permute_default_3 = CallFunction(aten.permute.default, KeywordArg('value'), Ignored())
-expand_default_3 = CallFunction(aten.expand.default, permute_default_3, Ignored())
+convert_element_type_default_1 = CallFunction(prims.convert_element_type.default, permute_default_3, Ignored())
+expand_default_3 = CallFunction(aten.expand.default, convert_element_type_default_1, Ignored())
 clone_default_2 = CallFunction(aten.clone.default, expand_default_3, memory_format=torch.contiguous_format)
 view_default_4 = CallFunction(aten.view.default, clone_default_2, Ignored(), _users=2)
 bmm_default_1 = CallFunction(aten.bmm.default, view_default_3, view_default_4)
@@ -65,8 +66,7 @@
 view_default_6 = CallFunction(aten.view.default, KeywordArg('tangents_1'), Ignored(), _users=2)
 permute_default_4 = CallFunction(aten.permute.default, view_default_4, Ignored())
 bmm_default_2 = CallFunction(aten.bmm.default, view_default_6, permute_default_4)
-convert_element_type_default_1 = CallFunction(prims.convert_element_type.default, bmm_default_2, Ignored())
-view_default_7 = CallFunction(aten.view.default, convert_element_type_default_1, Ignored())
+view_default_7 = CallFunction(aten.view.default, bmm_default_2, Ignored())
 convert_element_type_default_2 = CallFunction(prims.convert_element_type.default, view_default_7, Ignored())
 convert_element_type_default_3 = CallFunction(prims.convert_element_type.default, gt_Scalar, Ignored())
 mul_Tensor_2 = CallFunction(aten.mul.Tensor, convert_element_type_default_3, Ignored())
@@ -78,7 +78,7 @@
 permute_default_5 = CallFunction(aten.permute.default, view_default_1, Ignored())
 bmm_default_3 = CallFunction(aten.bmm.default, view_default_8, permute_default_5)
 view_default_9 = CallFunction(aten.view.default, bmm_default_3, Ignored())
-div_Tensor_2 = CallFunction(aten.div.Tensor, view_default_9, Ignored())
+div_Tensor_2 = CallFunction(aten.div.Tensor, view_default_9, KeywordArg('inv_scale'))
 permute_default_6 = CallFunction(aten.permute.default, div_Tensor_2, Ignored())
 permute_default_7 = CallFunction(aten.permute.default, view_default, Ignored())
 bmm_default_4 = CallFunction(aten.bmm.default, permute_default_7, view_default_8)
@@ -88,17 +88,19 @@
 permute_default_10 = CallFunction(aten.permute.default, view_default_3, Ignored())
 bmm_default_5 = CallFunction(aten.bmm.default, permute_default_10, view_default_6)
 view_default_11 = CallFunction(aten.view.default, bmm_default_5, Ignored())
-permute_default_11 = CallFunction(aten.permute.default, view_default_11, Ignored())
+convert_element_type_default_4 = CallFunction(prims.convert_element_type.default, view_default_11, Ignored())
+permute_default_11 = CallFunction(aten.permute.default, convert_element_type_default_4, Ignored())
 _sfdp_pattern_9_training = MultiOutputPattern([view_default_5,
   permute_default_6,
   permute_default_9,
   permute_default_11,
+  None,
   None
 ])
 
 
 permute_default = CallFunction(aten.permute.default, KeywordArg('query'), Ignored())
-div_Tensor = CallFunction(aten.div.Tensor, permute_default, Ignored())
+div_Tensor = CallFunction(aten.div.Tensor, permute_default, KeywordArg('inv_scale'))
 expand_default = CallFunction(aten.expand.default, div_Tensor, Ignored())
 clone_default = CallFunction(aten.clone.default, expand_default, memory_format=torch.contiguous_format)
 view_default = CallFunction(aten.view.default, clone_default, Ignored())
@@ -118,7 +120,8 @@
 expand_default_2 = CallFunction(aten.expand.default, convert_element_type_default, Ignored())
 view_default_3 = CallFunction(aten.view.default, expand_default_2, Ignored())
 permute_default_3 = CallFunction(aten.permute.default, KeywordArg('value'), Ignored())
-expand_default_3 = CallFunction(aten.expand.default, permute_default_3, Ignored())
+convert_element_type_default_1 = CallFunction(prims.convert_element_type.default, permute_default_3, Ignored())
+expand_default_3 = CallFunction(aten.expand.default, convert_element_type_default_1, Ignored())
 clone_default_2 = CallFunction(aten.clone.default, expand_default_3, memory_format=torch.contiguous_format)
 view_default_4 = CallFunction(aten.view.default, clone_default_2, Ignored())
 bmm_default_1 = CallFunction(aten.bmm.default, view_default_3, view_default_4)
@@ -128,7 +131,7 @@
 rand_default = CallFunction(aten.rand.default, Ignored(), dtype=Ignored(), device=Ignored(), pin_memory=False)
 gt_Scalar = CallFunction(aten.gt.Scalar, rand_default, KeywordArg('dropout_p'), _users=2)
 permute_default = CallFunction(aten.permute.default, KeywordArg('query'), Ignored())
-div_Tensor = CallFunction(aten.div.Tensor, permute_default, Ignored())
+div_Tensor = CallFunction(aten.div.Tensor, permute_default, KeywordArg('inv_scale'))
 expand_default = CallFunction(aten.expand.default, div_Tensor, Ignored())
 clone_default = CallFunction(aten.clone.default, expand_default, memory_format=torch.contiguous_format)
 view_default = CallFunction(aten.view.default, clone_default, Ignored(), _users=2)
@@ -173,7 +176,7 @@
 permute_default_5 = CallFunction(aten.permute.default, view_default_1, Ignored())
 bmm_default_3 = CallFunction(aten.bmm.default, view_default_8, permute_default_5)
 view_default_9 = CallFunction(aten.view.default, bmm_default_3, Ignored())
-div_Tensor_2 = CallFunction(aten.div.Tensor, view_default_9, Ignored())
+div_Tensor_2 = CallFunction(aten.div.Tensor, view_default_9, KeywordArg('inv_scale'))
 permute_default_6 = CallFunction(aten.permute.default, div_Tensor_2, Ignored())
 permute_default_7 = CallFunction(aten.permute.default, view_default, Ignored())
 bmm_default_4 = CallFunction(aten.bmm.default, permute_default_7, view_default_8)
@@ -188,12 +191,13 @@
   permute_default_6,
   permute_default_9,
   permute_default_11,
+  None,
   None
 ])
 
 
 permute_default = CallFunction(aten.permute.default, KeywordArg('query'), Ignored())
-div_Tensor = CallFunction(aten.div.Tensor, permute_default, Ignored())
+div_Tensor = CallFunction(aten.div.Tensor, permute_default, KeywordArg('inv_scale'))
 expand_default = CallFunction(aten.expand.default, div_Tensor, Ignored())
 clone_default = CallFunction(aten.clone.default, expand_default, memory_format=torch.contiguous_format)
 view_default = CallFunction(aten.view.default, clone_default, Ignored())
diff --git a/torch/_inductor/fx_passes/split_cat.py b/torch/_inductor/fx_passes/split_cat.py
index b4906a2f2ca14..9f548846d1463 100644
--- a/torch/_inductor/fx_passes/split_cat.py
+++ b/torch/_inductor/fx_passes/split_cat.py
@@ -2003,12 +2003,15 @@ def merge_unbind_stack_aten(match: Match, *args, **kwargs):
     cat_dim = get_arg_value(node, 1, "dim")
     # check the unsqueeze nodes come from the select nodes
     if not all(
+        # pyrefly: ignore [bad-argument-type]
         get_arg_value(unsqueeze_node, 0, "input").target is torch.ops.aten.select
         for unsqueeze_node in unsqueeze_nodes
     ):
         return
     select_nodes = [
-        get_arg_value(unsqueeze_node, 0, "input") for unsqueeze_node in unsqueeze_nodes
+        # pyrefly: ignore [bad-argument-type]
+        get_arg_value(unsqueeze_node, 0, "input")
+        for unsqueeze_node in unsqueeze_nodes
     ]
     parent_of_select_node = get_arg_value(select_nodes[0], 0, "input")
     # check the target of select_nodes are the same
@@ -2041,6 +2044,7 @@ def merge_unbind_stack_aten(match: Match, *args, **kwargs):
     node.replace_all_uses_with(parent_of_select_node)
     graph.erase_node(node)
     for unsqueeze_node in unsqueeze_nodes:
+        # pyrefly: ignore [bad-argument-type]
         graph.erase_node(unsqueeze_node)
     for select_node in select_nodes:
         if len(select_node.users) == 0:
diff --git a/torch/_inductor/fx_passes/spmd_check.py b/torch/_inductor/fx_passes/spmd_check.py
new file mode 100644
index 0000000000000..b51004107a107
--- /dev/null
+++ b/torch/_inductor/fx_passes/spmd_check.py
@@ -0,0 +1,228 @@
+"""SPMD graph verification for overlap scheduling.
+
+Verifies all ranks have identical FX graph structure before collective
+reordering passes. Non-SPMD graphs cause NCCL collective ordering
+mismatches and hangs.
+"""
+
+import hashlib
+import logging
+from collections import Counter
+
+import torch
+from torch._inductor import config
+from torch._logging import trace_structured
+
+
+log = logging.getLogger(__name__)
+
+
+def _compute_hash(gm: torch.fx.GraphModule) -> int | None:
+    """Compute a structural hash of the graph including tensor metadata.
+
+    Uses FxGraphCachePickler(device_id_agnostic=True) to serialize
+    (target, val) per call_function node, capturing op targets and
+    FakeTensor metadata (dtype, shape, stride, etc.) with device indices
+    normalized to 0.
+
+    Returns None if the graph contains unpicklable objects.
+    """
+    from torch._inductor.codecache import BypassFxGraphCache, FxGraphCachePickler
+
+    try:
+        pickler = FxGraphCachePickler(gm, device_id_agnostic=True)
+        data = pickler.dumps(
+            tuple(
+                (str(n.target), n.meta.get("val"))
+                for n in gm.graph.nodes
+                if n.op == "call_function"
+            )
+        )
+        digest = hashlib.blake2b(data, digest_size=8).digest()
+        return int.from_bytes(digest, "big", signed=True)
+    except BypassFxGraphCache:
+        # FxGraphCachePickler can't serialize certain objects:
+        # mkldnn tensors, BackwardState, torchbind objects, or general
+        # pickle failures. Skip the SPMD check gracefully.
+        log.warning("SPMD check: skipping, unpicklable graph objects", exc_info=True)
+        return None
+
+
+def _build_diag_fingerprint(
+    gm: torch.fx.GraphModule,
+) -> tuple[tuple[str, str | None], ...]:
+    """Build human-readable fingerprint for mismatch diagnostics.
+
+    Only called on the rare mismatch path.
+    """
+    from torch._inductor.codecache import extract_tensor_metadata_for_cache_key
+
+    entries: list[tuple[str, str | None]] = []
+    for n in gm.graph.nodes:
+        if n.op != "call_function":
+            continue
+        target_str = str(n.target)
+        val = n.meta.get("val")
+        entries.append(
+            (
+                target_str,
+                _format_val_metadata(val, extract_tensor_metadata_for_cache_key),
+            )
+        )
+    return tuple(entries)
+
+
+def _format_val_metadata(val: object, extract_fn: object) -> str | None:
+    """Format node val metadata for human-readable diagnostics."""
+    if val is None:
+        return None
+    if isinstance(val, torch.Tensor):
+        return str(extract_fn(val))  # type: ignore[operator]
+    if isinstance(val, (tuple, list)):
+        parts = []
+        for v in val:
+            if isinstance(v, torch.Tensor):
+                parts.append(str(extract_fn(v)))  # type: ignore[operator]
+            else:
+                parts.append(str(type(v).__name__))
+        return f"({', '.join(parts)})"
+    return str(type(val).__name__)
+
+
+def spmd_check(gm: torch.fx.GraphModule) -> bool:
+    """Verify all ranks have identical FX graph structure (SPMD).
+
+    Computes a structural hash (op targets + tensor metadata including
+    shapes, dtypes, strides) and compares across ranks.
+    On mismatch, emits a diagnostic report to stdout, logging, and
+    trace_structured.
+
+    Returns True if graphs match (SPMD), False on mismatch.
+    """
+    import torch.distributed as dist
+
+    if not dist.is_initialized() or dist.get_world_size() <= 1:
+        return True
+
+    structure_hash = _compute_hash(gm)
+    if structure_hash is None:
+        return True
+
+    from torch._subclasses.fake_tensor import unset_fake_temporarily
+    from torch.distributed.distributed_c10d import _get_default_group
+
+    pg = _get_default_group()
+    world_size = dist.get_world_size()
+    rank = dist.get_rank()
+
+    with unset_fake_temporarily():
+        all_hashes: list[int] = [0] * world_size
+        dist.all_gather_object(all_hashes, structure_hash, pg)
+
+    if all(h == all_hashes[0] for h in all_hashes):
+        return True
+
+    # Mismatch detected — build and gather diagnostic fingerprints
+    fingerprint = _build_diag_fingerprint(gm)
+    with unset_fake_temporarily():
+        all_fingerprints: list[tuple[object, ...]] = [() for _ in range(world_size)]
+        dist.all_gather_object(all_fingerprints, fingerprint, pg)
+
+    report = _build_mismatch_report(all_fingerprints, rank, world_size)
+
+    print(report, flush=True)
+    log.warning("\n%s", report)
+
+    trace_structured(
+        "artifact",
+        metadata_fn=lambda: {
+            "name": "inductor_spmd_graph_mismatch",
+            "encoding": "string",
+        },
+        payload_fn=lambda: report,
+    )
+
+    if config.aten_distributed_optimizations.spmd_mismatch == "error":
+        raise RuntimeError(
+            "SPMD graph verification failed. "
+            'Set aten_distributed_optimizations.spmd_mismatch="warn" '
+            "to warn instead of fail.\n" + report
+        )
+
+    return False
+
+
+def _entry_target(entry: object) -> str:
+    """Extract the target string from a fingerprint entry."""
+    if isinstance(entry, tuple):
+        return str(entry[0])
+    return str(entry)
+
+
+def _entry_metadata(entry: object) -> str:
+    """Format metadata from a fingerprint entry, if present."""
+    if isinstance(entry, tuple) and len(entry) >= 2:
+        meta = entry[1]
+        if meta is not None:
+            return f" meta={meta}"
+    return ""
+
+
+def _build_mismatch_report(
+    all_fingerprints: list[tuple[object, ...]],
+    rank: int,
+    world_size: int,
+) -> str:
+    """Build diagnostic report for SPMD graph mismatch."""
+    lines = [
+        "=" * 80,
+        f"SPMD GRAPH MISMATCH — rank {rank}, world_size={world_size}",
+        "=" * 80,
+    ]
+
+    # Node count per rank
+    counts = [len(t) for t in all_fingerprints]
+    lines.append("NODE COUNTS PER RANK:")
+    for r in range(world_size):
+        marker = " <--" if counts[r] != counts[0] else ""
+        lines.append(f"  rank {r}: {counts[r]} call_function nodes{marker}")
+    lines.append("")
+
+    # Find entries that differ
+    ref = all_fingerprints[0]
+    for r in range(1, world_size):
+        other = all_fingerprints[r]
+        if other == ref:
+            continue
+        lines.append(f"DIFFS rank 0 vs rank {r}:")
+
+        # Show first few positional differences
+        max_diffs = 10
+        shown = 0
+        for i, (a, b) in enumerate(zip(ref, other)):
+            if a != b and shown < max_diffs:
+                lines.append(f"  node {i}:")
+                lines.append(f"    rank 0: {_entry_target(a)}{_entry_metadata(a)}")
+                lines.append(f"    rank {r}: {_entry_target(b)}{_entry_metadata(b)}")
+                shown += 1
+
+        # Also show count-based diffs for op targets
+        ref_targets = [_entry_target(e) for e in ref]
+        other_targets = [_entry_target(e) for e in other]
+
+        ref_counts = Counter(ref_targets)
+        other_counts = Counter(other_targets)
+        only_ref = ref_counts - other_counts
+        only_other = other_counts - ref_counts
+        if only_ref:
+            lines.append("  Only on rank 0:")
+            for op, cnt in only_ref.most_common(10):
+                lines.append(f"    {op} (x{cnt})")
+        if only_other:
+            lines.append(f"  Only on rank {r}:")
+            for op, cnt in only_other.most_common(10):
+                lines.append(f"    {op} (x{cnt})")
+        lines.append("")
+
+    lines.append("=" * 80)
+    return "\n".join(lines)
diff --git a/torch/_inductor/fx_utils.py b/torch/_inductor/fx_utils.py
index 88cb9c7ea08bb..c520deb369a0a 100644
--- a/torch/_inductor/fx_utils.py
+++ b/torch/_inductor/fx_utils.py
@@ -3,7 +3,7 @@
 import operator
 from collections import defaultdict
 from collections.abc import Callable
-from typing import Any, Optional
+from typing import Any
 
 import sympy
 
@@ -91,7 +91,7 @@ def hash_node(self, node: torch.fx.Node):
 
     def incremental_update(self):
         """Update FakeTensors on self.graph. We will try to do the minimum amount of work."""
-        existing_storages: defaultdict[Optional[int], int] = defaultdict(int)
+        existing_storages: defaultdict[int | None, int] = defaultdict(int)
         for node in self.graph.nodes:
             existing_storages[get_node_storage(node)] += 1
 
@@ -253,7 +253,7 @@ def get_storage(t: torch.Tensor) -> int:
     return t.untyped_storage()._cdata
 
 
-def get_node_storage(node: torch.fx.Node) -> Optional[int]:
+def get_node_storage(node: torch.fx.Node) -> int | None:
     if "val" not in node.meta:
         return None
     if not isinstance(node.meta["val"], torch.Tensor):
@@ -315,7 +315,7 @@ def realizes_inputs(node: torch.fx.Node) -> bool:
     return False
 
 
-def count_flops_fx(node: torch.fx.Node) -> Optional[int]:
+def count_flops_fx(node: torch.fx.Node) -> int | None:
     if not countable_fx(node) or isinstance(node.target, str):
         return None
     with FakeTensorMode(allow_non_fake_inputs=True):
diff --git a/torch/_inductor/graph.py b/torch/_inductor/graph.py
index 7c0876faaac85..c774772494545 100644
--- a/torch/_inductor/graph.py
+++ b/torch/_inductor/graph.py
@@ -9,9 +9,10 @@
 import re
 import sys
 import time
+import typing_extensions
 from collections import defaultdict
 from contextlib import contextmanager
-from typing import Any, NoReturn, Optional, TYPE_CHECKING, Union
+from typing import Any, NoReturn, TYPE_CHECKING
 
 import sympy
 from sympy import Expr
@@ -23,7 +24,11 @@
 from torch._decomp import get_decompositions
 from torch._dynamo.utils import defake, dynamo_timed
 from torch._library.fake_class_registry import FakeScriptObject
-from torch._library.opaque_object import is_opaque_type
+from torch._library.opaque_object import (
+    is_opaque_reference_type,
+    is_opaque_type,
+    is_opaque_value_type,
+)
 from torch._library.utils import get_layout_constraint_tag
 from torch._logging import LazyString, trace_structured
 from torch._prims_common import (
@@ -96,6 +101,7 @@
     require_contiguous,
     tag_to_layout_constraint,
     unsupported_output_tensor,
+    user_lowerings,
 )
 from .runtime import autotune_cache
 from .runtime.autotune_cache import AutotuneCacheBundler
@@ -130,7 +136,7 @@
     from .dependencies import Dep
     from .scheduler import BaseSchedulerNode
 
-    CompiledModule = Union[ModuleType, FileBackedGraphModule]
+    CompiledModule = ModuleType | FileBackedGraphModule
 
 from torch._inductor.codecache import output_code_log
 
@@ -143,14 +149,20 @@
 _post_grad_graph_counter = itertools.count()
 
 if config.is_fbcode():
+    from torch._inductor.fb.triton_kernel_metadata import (
+        save_triton_kernel_perf_artifact,
+    )
     from torch._inductor.fb.utils import log_module_code
 else:
 
     def log_module_code(*args: Any, **kwargs: Any) -> None:
         pass
 
+    def save_triton_kernel_perf_artifact(*args: Any, **kwargs: Any) -> None:
+        pass
 
-def may_get_constant_buffer_dtype(constant_buffer: sympy.Expr) -> Optional[torch.dtype]:
+
+def may_get_constant_buffer_dtype(constant_buffer: sympy.Expr) -> torch.dtype | None:
     assert isinstance(
         constant_buffer, (sympy.Symbol, sympy.Expr, sympy.core.numbers.Integer)
     ), (
@@ -177,7 +189,7 @@ def is_magic_method(op: Any) -> bool:
 
 def getattr_recursive(
     obj: GraphModule, target: str
-) -> Union[Tensor, torch._C.ScriptObject, GraphModule]:
+) -> Tensor | torch._C.ScriptObject | GraphModule:
     target_atoms = target.split(".")
     attr_itr = obj
     for i, atom in enumerate(target_atoms):
@@ -275,7 +287,7 @@ def mark_nodes_dislike_padding(
 
     def _get_overload_packet(
         node: torch.fx.Node,
-    ) -> Optional[torch._ops.OpOverloadPacket]:
+    ) -> torch._ops.OpOverloadPacket | None:
         return (
             node.target._overloadpacket
             if node.op == "call_function"
@@ -347,27 +359,28 @@ class GraphLowering(torch.fx.Interpreter):
     def __init__(
         self,
         gm: torch.fx.GraphModule,
-        example_inputs: Optional[Sequence[object]] = None,
-        shape_env: Optional[ShapeEnv] = None,
-        graph_id: Optional[int] = None,
+        example_inputs: Sequence[object] | None = None,
+        shape_env: ShapeEnv | None = None,
+        graph_id: int | None = None,
         cpp_wrapper: bool = False,
         aot_mode: bool = False,
-        layout_opt: Optional[bool] = None,
-        extern_node_serializer: Optional[
-            Callable[[list[ir.ExternKernelNode]], Any]
-        ] = None,
+        layout_opt: bool | None = None,
+        extern_node_serializer: Callable[[list[ir.ExternKernelNode]], Any]
+        | None = None,
         is_inference: bool = False,
         is_backward: bool = False,
         is_const_graph: bool = False,
-        const_output_index: Optional[dict[str, int]] = None,
-        const_wrapper_code: Optional[str] = None,
-        const_kernel_code: Optional[str] = None,
-        const_module: Optional[GraphLowering] = None,
-        name: Optional[str] = None,
-        inputs_to_check: Optional[Sequence[int]] = None,
+        const_output_index: dict[str, int] | None = None,
+        const_wrapper_code: str | None = None,
+        const_kernel_code: str | None = None,
+        const_module: GraphLowering | None = None,
+        name: str | None = None,
+        inputs_to_check: Sequence[int] | None = None,
         fx_wrapper: bool = False,
+        get_decomp_fn: Callable[..., dict[Any, Callable[..., Any]]] | None = None,
     ) -> None:
         super().__init__(gm)
+        self.get_decomp_fn = get_decomp_fn
         self.example_inputs = example_inputs
         self.layout_opt = (
             layout_opt
@@ -382,6 +395,7 @@ def __init__(
         self.const_kernel_code = const_kernel_code
         self.const_module = const_module
         self.inputs_to_check = inputs_to_check
+        self._defers_input_alignment = False
 
         self.extra_traceback = False  # we do our own error wrapping
         if shape_env is None:
@@ -391,16 +405,16 @@ def __init__(
             self.reuse_shape_env = True
         self._shape_env = shape_env
         # We're going to mutate ras_by_symbol as we finish generating them
-        self.ras_by_symbol: dict[Optional[sympy.Symbol], list[RuntimeAssert]] = (
+        self.ras_by_symbol: dict[sympy.Symbol | None, list[RuntimeAssert]] = (
             shape_env.deferred_runtime_asserts.copy()
         )
         self.bound_unbacked_symbols = OrderedSet[sympy.Symbol]()
 
         self.sizevars = SizeVarAllocator(shape_env)
         self.graph_input_names: list[str] = []
-        self.graph_inputs: dict[str, Union[TensorBox, TorchBindObject, sympy.Expr]] = {}
+        self.graph_inputs: dict[str, TensorBox | TorchBindObject | sympy.Expr] = {}
         self.graph_inputs_original: dict[str, InputBuffer] = {}
-        self.partition_maps: Optional[list[GraphPartitionMap]] = None
+        self.partition_maps: list[GraphPartitionMap] | None = None
         self.zero_dim_cpu_tensor_list: OrderedSet[str] = OrderedSet()
         self.device_types: OrderedSet[str] = (
             const_module.device_types if const_module else OrderedSet()
@@ -441,7 +455,7 @@ def __init__(
             const_module.named_parameters if const_module else {}
         )
         self.torchbind_constants: dict[
-            str, Union[torch._C.ScriptObject, FakeScriptObject]
+            str, torch._C.ScriptObject | FakeScriptObject
         ] = {}
         self.opaque_value_type_classes: dict[str, type] = {}
         self.seen_subgraphs: dict[str, ir.Subgraph] = {}
@@ -450,6 +464,7 @@ def __init__(
         self.removed_buffers: OrderedSet[str] = OrderedSet()
         self.removed_inplace_buffers: OrderedSet[str] = OrderedSet()
         self.mutated_buffers: OrderedSet[str] = OrderedSet()
+        self.sdpa_constraint_cache: dict[tuple, ir.IRNode] = {}
         self.never_reuse_buffers: OrderedSet[str] = OrderedSet()
         self.inplaced_to_remove: OrderedSet[str] = OrderedSet()
         self.device_ops: DeviceOpOverrides = None  # type: ignore[assignment]
@@ -488,13 +503,13 @@ def __init__(
 
         # record intermediate results for input of UsedDefinedTritonKernels
         # This will be used if autotuning is done in one pass.
-        self.autotuning_inputs: Optional[list[torch.Tensor]] = None
-        self.autotuning_mapping: Optional[dict[str, dict[str, int]]] = None
-        self.autotuning_grids: Optional[dict[str, Any]] = None
+        self.autotuning_inputs: list[torch.Tensor] | None = None
+        self.autotuning_mapping: dict[str, dict[str, int]] | None = None
+        self.autotuning_grids: dict[str, Any] | None = None
 
         # current_device is set only during codegen of a device-specific kernel
         # a graph can have many devices
-        self.current_device: Optional[torch.device] = None
+        self.current_device: torch.device | None = None
 
         self.nodes_prefer_channels_last = (
             self.find_nodes_prefer_channels_last() if self.layout_opt else OrderedSet()
@@ -508,7 +523,7 @@ def __init__(
             tuple[int, str]
         ] = []  # This is the linemap used by the profiler to mark custom compiled kernels getting run
         # Used if lowering encounters cases where cudagraphs are not supported
-        self.disable_cudagraphs_reason: Optional[str] = None
+        self.disable_cudagraphs_reason: str | None = None
 
         # only keeping one node per device for stack trace purposes
         self.device_node_mapping: dict[torch.device, torch.fx.Node] = {}
@@ -560,7 +575,7 @@ def freeze_runtime_asserts(self) -> None:
 
     def symbolic_sizes_strides(
         self, ex: torch.Tensor
-    ) -> tuple[Sequence[Union[int, Expr]], Sequence[Union[int, Expr]]]:
+    ) -> tuple[Sequence[int | Expr], Sequence[int | Expr]]:
         """
         Support dynamic shapes and dynamic strides by assigning variables
         to each dimension.  We duck-shape tensors, so if two tensors
@@ -607,9 +622,11 @@ def static_sizes_strides(
 
     def get_allocation_size(
         self,
-        node: Union[
-            ir.TensorBox, ir.StorageBox, ir.Buffer, WorkspaceArg, ir.TorchBindObject
-        ],
+        node: ir.TensorBox
+        | ir.StorageBox
+        | ir.Buffer
+        | WorkspaceArg
+        | ir.TorchBindObject,
     ) -> Sequence[Expr]:
         if isinstance(node, ir.TensorBox):
             node = node.data  # type: ignore[assignment]
@@ -625,7 +642,7 @@ def get_allocation_size(
             return node.get_size()
 
     def get_allocation_storage_size(
-        self, node: Union[ir.Buffer, WorkspaceArg, ir.TorchBindObject]
+        self, node: ir.Buffer | WorkspaceArg | ir.TorchBindObject
     ) -> Expr:
         layout = node.get_layout()
         size = self.get_allocation_size(node)  # consider inplace padding
@@ -635,7 +652,7 @@ def get_allocation_storage_size(
 
     def has_feature(
         self,
-        device: Union[torch._inductor.ir.IRNode, device, None],
+        device: torch._inductor.ir.IRNode | device | None,
         feature: BackendFeature,
     ) -> bool:
         assert isinstance(feature, BackendFeature), feature
@@ -647,8 +664,17 @@ def get_dep_size_hint(self, dep: Dep, count_bytes: bool = True) -> int:
         """
         if (dep, count_bytes) not in self.dep_size_hint_cache:
             res = 0
+            # Non-tensor graph inputs (TorchBindObject, OpaqueObjectState)
+            # have no meaningful size — skip the size computation entirely.
+            inp = self.graph_inputs.get(dep.name)
+            if isinstance(inp, ir.NonTensorObj):
+                self.dep_size_hint_cache[(dep, count_bytes)] = 0
+                return 0
             try:
-                if not dep.has_unbacked_symbols():
+                if (
+                    not dep.has_unbacked_symbols()
+                    or self.sizevars.all_unbacked_explicitly_hinted(dep.get_numel())
+                ):
                     if count_bytes:
                         res = dep.numbytes_hint()
                     else:
@@ -959,7 +985,7 @@ def fake_mode(self) -> torch._subclasses.fake_tensor.FakeTensorMode:
 
     def try_get_buffer(
         self, buffer_name: str
-    ) -> Optional[Union[ir.TensorBox, ir.Buffer, ir.TorchBindObject]]:
+    ) -> ir.TensorBox | ir.Buffer | ir.TorchBindObject | None:
         if buffer_name in self.name_to_buffer:
             return self.name_to_buffer[buffer_name]
         if buffer_name in self.graph_inputs:
@@ -980,7 +1006,7 @@ def add_symbol_graph_input(self, symbol: sympy.Expr) -> None:
 
     def get_buffer(
         self, buffer_name: str
-    ) -> Union[ir.TensorBox, ir.Buffer, ir.TorchBindObject]:
+    ) -> ir.TensorBox | ir.Buffer | ir.TorchBindObject:
         buf = self.try_get_buffer(buffer_name)
         if buf is not None:
             return buf
@@ -1008,7 +1034,7 @@ def get_dtype(self, buffer_name: str) -> torch.dtype:
             return self.get_dtype(m.group(1))
         raise KeyError(f"could not find {buffer_name}")
 
-    def get_numel(self, buffer_name: str) -> Union[int, Expr]:
+    def get_numel(self, buffer_name: str) -> int | Expr:
         if buffer_name in self.constants:
             return self.constants[buffer_name].numel()
         if buffer_name in self.name_to_buffer:
@@ -1058,10 +1084,8 @@ def register_operation_list(self, operation_names: list[str]) -> str:
         self.lists[name] = operation_names
         return name
 
-    def register_users_of(
-        self, node_output: Union[Iterable[ir.IRNode], ir.IRNode]
-    ) -> None:
-        def register(value: Union[Iterable[ir.IRNode], ir.IRNode]) -> None:
+    def register_users_of(self, node_output: Iterable[ir.IRNode] | ir.IRNode) -> None:
+        def register(value: Iterable[ir.IRNode] | ir.IRNode) -> None:
             if isinstance(value, (list, tuple)):
                 for x in value:
                     register(x)
@@ -1101,9 +1125,7 @@ def get_original_value_of_constant(self, name: str) -> torch.Tensor:
             else self.constants[name]
         )
 
-    def allocate_non_dup_const_name(
-        self, name: Optional[str], data: Union[Tensor]
-    ) -> str:
+    def allocate_non_dup_const_name(self, name: str | None, data: Tensor) -> str:
         if not config.aot_inductor.use_runtime_constant_folding:
             for constant_name, value in self.constants.items():
                 if is_same_tensor(data, value):
@@ -1132,9 +1154,7 @@ def allocate_non_dup_const_name(
         self.allocated_constant_name[name] = orig_name  # type: ignore[assignment]
         return name
 
-    def add_tensor_constant(
-        self, data: Tensor, name: Optional[str] = None
-    ) -> TensorBox:
+    def add_tensor_constant(self, data: Tensor, name: str | None = None) -> TensorBox:
         new_name = self.allocate_non_dup_const_name(name, data)
         return TensorBox.create(
             ir.ConstantBuffer(
@@ -1145,7 +1165,7 @@ def add_tensor_constant(
             )
         )
 
-    def constant_name(self, name: str, device_override: Optional[torch.device]) -> str:
+    def constant_name(self, name: str, device_override: torch.device | None) -> str:
         """
         We AOT copy constants to the devices they are needed on.
         If device_override doesn't match the constant's device, then
@@ -1191,7 +1211,7 @@ def placeholder(
         target: str,  # type: ignore[override]
         args: tuple[object],  # type: ignore[override]
         kwargs: dict[str, object],
-    ) -> Union[Expr, TensorBox, None]:
+    ) -> Expr | TensorBox | None:
         self.placeholder_idx += 1
         example = super().placeholder(target, args, kwargs)  # type: ignore[arg-type]
         target = self.qualify_name(target)
@@ -1235,6 +1255,11 @@ def placeholder(
             self.graph_inputs[target] = gen  # type: ignore[assignment]
             self.graph_input_names.append(target)
             return gen
+        elif is_opaque_reference_type(type(example)):
+            opaque_obj = ir.OpaqueObjectState(name=target, value=example)
+            self.graph_inputs[target] = opaque_obj  # type: ignore[assignment]
+            self.graph_input_names.append(target)
+            return opaque_obj
 
         assert isinstance(example, torch.Tensor), example
         # todo(chilli): We can remove the last check once we turn buffers into
@@ -1290,7 +1315,8 @@ def placeholder(
                 self.unaligned_buffers.add(target)
         return tensor
 
-    def call_function(self, target: Callable, args: Any, kwargs: dict[str, Any]) -> Any:  # type: ignore[type-arg, override]
+    @typing_extensions.override
+    def call_function(self, target: Callable, args: Any, kwargs: dict[str, Any]) -> Any:  # type: ignore[type-arg]
         if target is operator.getitem and isinstance(args[0], (list, tuple, dict)):
             return super().call_function(target, args, kwargs)
 
@@ -1307,7 +1333,12 @@ def call_function(self, target: Callable, args: Any, kwargs: dict[str, Any]) ->
             )
             base_name = target.name().split(".")[0]
             if base_name in FALLBACK_ALLOW_LIST:
-                make_fallback(target, warn=False, override_decomp=True)
+                make_fallback(
+                    target,
+                    warn=False,
+                    get_decomp_fn=self.get_decomp_fn,
+                    override_decomp=True,
+                )
             elif config.implicit_fallbacks:
                 error = (
                     MissingOperatorWithDecomp
@@ -1319,7 +1350,7 @@ def call_function(self, target: Callable, args: Any, kwargs: dict[str, Any]) ->
                     error.operator_str(target, args, kwargs),
                 )
 
-                tag: Optional[torch._C.Tag] = get_layout_constraint_tag(
+                tag: torch._C.Tag | None = get_layout_constraint_tag(
                     target, with_default=False
                 )
                 if (
@@ -1336,7 +1367,7 @@ def call_function(self, target: Callable, args: Any, kwargs: dict[str, Any]) ->
                     #
                     # TODO: should really switch to "needs_fixed_stride" constraint on these
                     # and identify them one by one.
-                    decided_constraint: Optional[Callable[..., tuple[Any, Any]]] = (
+                    decided_constraint: Callable[..., tuple[Any, Any]] | None = (
                         require_contiguous
                     )
                 else:
@@ -1345,7 +1376,11 @@ def call_function(self, target: Callable, args: Any, kwargs: dict[str, Any]) ->
                     )
                     decided_constraint = tag_to_layout_constraint(default_tag)
 
-                make_fallback(target, layout_constraint=decided_constraint)
+                make_fallback(
+                    target,
+                    layout_constraint=decided_constraint,
+                    get_decomp_fn=self.get_decomp_fn,
+                )
 
             elif get_decompositions([target]):
                 # There isn't a good way to dynamically patch this in
@@ -1396,7 +1431,29 @@ def normalize(args: Any, kwargs: Any) -> tuple[Any, Any]:
                     *args, **kwargs
                 )
             else:
-                out = lowerings[target](*args, **kwargs)  # type: ignore[index]
+                out = None
+
+                if (
+                    target in user_lowerings
+                    and target not in V.active_user_lowering_ops
+                ):
+                    # User-registered lowering takes priority, with recursion guard
+                    V.active_user_lowering_ops.add(target)
+                    try:
+                        # pyrefly: ignore[bad-index]
+                        out = user_lowerings[target](*args, **kwargs)
+                    finally:
+                        V.active_user_lowering_ops.discard(target)
+
+                # If no user_lowering, or it returned None fall back to normal lowering
+                if out is None:
+                    if target in lowerings:
+                        out = lowerings[target](*args, **kwargs)
+                    else:
+                        # Fallback for ops not in lowerings (e.g., custom ops during recursion)
+                        out = fallback_handler(target, add_to_fallback_set=False)(
+                            *args, **kwargs
+                        )
 
             if layout_constraints:
                 # layout_constraints are allowed to make new copies of the inputs.
@@ -1431,9 +1488,7 @@ def get_attr(
         target: str,  # type: ignore[override]
         args: tuple[()],  # type: ignore[override]
         kwargs: dict[str, object],
-    ) -> Union[
-        Constant, TensorBox, ShapeAsConstantBuffer, ir.Subgraph, TorchBindObject
-    ]:
+    ) -> Constant | TensorBox | ShapeAsConstantBuffer | ir.Subgraph | TorchBindObject:
         # this is a constant
         value = getattr_recursive(self.module, target)  # type: ignore[arg-type]
 
@@ -1474,7 +1529,7 @@ def get_attr(
                     value=value.item(), dtype=value.dtype, device=value.device
                 )
             if self.can_inline_constant(value):
-                log.debug("Inlining constant: %s ", str(target))
+                log.debug("Inlining constant: %s ", target)
                 # tensor lowering has constant inlining logic
                 from .lowering import tensor
 
@@ -1488,11 +1543,11 @@ def call_module(self, target: Any, args: Any, kwargs: Any) -> NoReturn:
     def call_method(self, target: Any, args: Any, kwargs: Any) -> NoReturn:
         raise AssertionError
 
-    # pyrefly: ignore [bad-override]
+    @typing_extensions.override
     def output(
         self,
-        target: str,  # type: ignore[override]
-        args: tuple[object],  # type: ignore[override]
+        target: torch.fx.node.Target,
+        args: tuple[torch.fx.node.Argument, ...],
         kwargs: dict[str, object],
     ) -> None:
         result = super().output(target, args, kwargs)  # type: ignore[arg-type]
@@ -1500,6 +1555,10 @@ def output(
             # nested subgraphs can have singleton outputs
             result = (result,)
         assert isinstance(result, (tuple, list)), type(result)
+        result = [
+            ir.OpaqueValueTypeConstant(value=x) if is_opaque_value_type(type(x)) else x
+            for x in result
+        ]
         assert all(
             isinstance(
                 x,
@@ -1514,6 +1573,9 @@ def output(
                     ir.EffectfulKernel,
                     ir.ShapeAsConstantBuffer,
                     TorchBindObject,
+                    ir.OpaqueMultiOutput,
+                    ir.OpaqueValueTypeConstant,
+                    ir.OpaqueObjectState,
                 ),
             )
             for x in result
@@ -1552,13 +1614,19 @@ def output(
         self.graph_outputs = result_correct_strides
         value: ir.IRNode
         for name, value in self.graph_inputs.items():
-            if isinstance(value, TorchBindObject):
-                continue
-            assert isinstance(
-                value, (TensorBox, sympy.Expr, torch._inductor.ir.GeneratorState)
-            ), f"Unsupported inductor graph input type: {type(value)}"
-            if not isinstance(value, TensorBox):
+            if isinstance(
+                value,
+                (
+                    TorchBindObject,
+                    sympy.Basic,
+                    torch._inductor.ir.GeneratorState,
+                    torch._inductor.ir.OpaqueObjectState,
+                ),
+            ):
                 continue
+            assert isinstance(value, TensorBox), (
+                f"Unsupported inductor graph input type: {type(value)}"
+            )
             value.realize()
             assert isinstance(value, TensorBox)
             value = value.data
@@ -1680,6 +1748,29 @@ def maybe_propagate(
             schema_arg = schema_kwargs[key]
             maybe_propagate(schema_arg, old_arg, new_arg)
 
+    @staticmethod
+    def _get_node_stream(n: torch.fx.Node) -> int | None:
+        """Get the user-annotated stream index from FX node metadata."""
+        return n.meta.get("custom", {}).get("stream")
+
+    def _realize_inputs_at_stream_boundaries(self, n: torch.fx.Node) -> None:
+        """Realize IR inputs that are on a different stream.
+
+        Without this, pointwise ops across stream boundaries would be inlined
+        into each other during lowering, making it impossible for the scheduler
+        to split them into separate kernels.
+
+        None means the default stream, so it is compared like any other value.
+        """
+        node_stream = self._get_node_stream(n)
+        for input_node in n.all_input_nodes:
+            input_stream = self._get_node_stream(input_node)
+            if input_stream == node_stream:
+                continue
+            ir_value = self.env.get(input_node)
+            if isinstance(ir_value, ir.TensorBox):
+                ir_value.realize()
+
     def run_node(self, n: torch.fx.Node) -> object:
         """Lower and execute a single FX node into Inductor IR."""
 
@@ -1724,8 +1815,10 @@ def maybe_apply_channels_last_stride_order(
         if is_call_function:
             args, kwargs = self.fetch_args_kwargs_from_env(n)
             origins |= gather_origins(args, kwargs)
+            self._realize_inputs_at_stream_boundaries(n)
         with (
             ir.IRNode.current_origins(origins),
+            ir.IRNode.current_stream_idx(self._get_node_stream(n)),
             self.set_current_node(n),
             V.set_current_node(n),
         ):
@@ -1846,7 +1939,7 @@ def maybe_apply_channels_last_stride_order(
                 result.realize()
 
             if (is_output or is_input_for_as_strided) and isinstance(
-                n.meta["val"], torch.Tensor
+                n.meta.get("val"), torch.Tensor
             ):
                 if is_user_visible:
                     strides = self.user_visible_output_strides.get(n)
@@ -1960,6 +2053,19 @@ def maybe_apply_channels_last_stride_order(
                     if user.op == "output":
                         # pyrefly: ignore [missing-attribute]
                         if isinstance(result.data.data, (Pointwise, Reduction)):
+                            # Cheap-to-recompute nodes (0 buffer reads, e.g.
+                            # index arithmetic or constant fills) can be
+                            # deferred to realize_input at output processing.
+                            # This prevents cascade materialization where
+                            # shared constants inflate downstream read counts.
+                            if (
+                                config.delay_realize_cheap_outputs
+                                # pyrefly: ignore [missing-attribute]
+                                and result.data.num_reads() == 0
+                                # pyrefly: ignore [missing-attribute]
+                                and not result.data.has_large_inner_fn()
+                            ):
+                                continue
                             result.realize()
 
                 _data = result.data  # type: ignore[attr-defined]
@@ -2083,6 +2189,8 @@ def format_new_defs() -> str:
     def create_deferred_runtime_asserts(
         self, n: torch.fx.Node, new_unbacked_defs: OrderedSet[sympy.Symbol]
     ) -> None:
+        if config.do_not_emit_runtime_assertions:
+            return
         # [NOTE] Codegen runtime asserts in Inductor
         #
         # We need to generate runtime asserts directly in Inductor instead
@@ -2175,9 +2283,9 @@ def validate_can_generate_cpp_wrapper(self) -> None:
     def init_wrapper_code(
         self,
         is_subgraph: bool = False,
-        subgraph_name: Optional[str] = None,
-        parent_wrapper_code: Optional[PythonWrapperCodegen] = None,
-        partition_signatures: Optional[GraphPartitionSignature] = None,
+        subgraph_name: str | None = None,
+        parent_wrapper_code: PythonWrapperCodegen | None = None,
+        partition_signatures: GraphPartitionSignature | None = None,
     ) -> None:
         device_types = self.device_types.copy()
         device_types.discard("cpu")
@@ -2210,7 +2318,7 @@ def init_wrapper_code(
             self.wrapper_code._names_iter = self.const_module.wrapper_code._names_iter
 
     def extract_autotune_inputs(
-        self, example_inputs: list[Union[int, float, torch.Tensor]]
+        self, example_inputs: list[int | float | torch.Tensor]
     ) -> None:
         import copy
 
@@ -2312,10 +2420,10 @@ def codegen_with_cpp_wrapper(
         """
         if any(device in self.device_types for device in ["cuda", "xpu"]):
 
-            def extract_real_inputs() -> list[Union[int, float, torch.Tensor]]:
+            def extract_real_inputs() -> list[int | float | torch.Tensor]:
                 def materialize(
-                    x: Union[torch.SymInt, torch.SymFloat, torch.Tensor],
-                ) -> Union[int, float, torch.Tensor]:
+                    x: torch.SymInt | torch.SymFloat | torch.Tensor,
+                ) -> int | float | torch.Tensor:
                     if x is None:
                         # pyrefly: ignore [bad-return]
                         return None
@@ -2394,8 +2502,14 @@ def materialize(
                     if user_defined_kernels:
                         real_inputs = extract_real_inputs()
                         self.extract_autotune_inputs(real_inputs)
+                        save_triton_kernel_perf_artifact(self)
                 return self.codegen()
             else:
+                if not self.aot_mode:
+                    # Lazy kernel compilation does not require two passes
+                    # TODO: need to consolidate the logic between AOT and JIT
+                    return self.codegen()
+
                 # first pass
                 self.cpp_wrapper = False
                 compiled = self.compile_to_module().call
@@ -2487,7 +2601,7 @@ def count_bytes(
         return total_bytes, node_counts, node_runtimes
 
     # No-op to be patched for unit tests
-    save_output_code: Optional[Callable[[str], None]] = None
+    save_output_code: Callable[[str], None] | None = None
 
     def compile_to_module(self) -> CompiledModule:
         with dynamo_timed(
@@ -2597,7 +2711,8 @@ def _compile_to_module_lines(
 
         if config.benchmark_harness and config.profile_bandwidth_output:
             # run the inputs code gen to get the bandwidth info
-            mod.benchmark_compiled_module(times=1, repeat=1)
+            args = mod.get_args()
+            mod.benchmark_compiled_module(args, times=1, repeat=1)
 
         return mod
 
@@ -2641,9 +2756,9 @@ def __init__(self, parent: GraphLowering, *args: Any, **kwargs: Any) -> None:
     def init_wrapper_code(
         self,
         is_subgraph: bool = False,
-        subgraph_name: Optional[str] = None,
-        parent_wrapper_code: Optional[PythonWrapperCodegen] = None,
-        partition_signatures: Optional[GraphPartitionSignature] = None,
+        subgraph_name: str | None = None,
+        parent_wrapper_code: PythonWrapperCodegen | None = None,
+        partition_signatures: GraphPartitionSignature | None = None,
     ) -> None:
         super().init_wrapper_code(
             is_subgraph=True,
diff --git a/torch/_inductor/index_propagation.py b/torch/_inductor/index_propagation.py
index 3711266ae93b0..8ae0435437dbd 100644
--- a/torch/_inductor/index_propagation.py
+++ b/torch/_inductor/index_propagation.py
@@ -24,7 +24,7 @@
 import itertools
 from collections.abc import Sequence
 from dataclasses import dataclass
-from typing import Any, Literal, Optional, overload, TypeAlias, Union
+from typing import Any, Literal, overload, TypeAlias
 
 import sympy
 
@@ -39,7 +39,7 @@
 from .virtualized import V
 
 
-_ExprType = Union[sympy.Expr, float, int, bool]
+_ExprType = sympy.Expr | float | int | bool
 
 
 def _is_constant(val: _ExprType):
@@ -91,18 +91,18 @@ def identity(value: Any) -> Any:
         return value
 
     @staticmethod
-    def constant(value: Union[int, float, bool], dtype: torch.dtype) -> TypedExpr:
+    def constant(value: int | float | bool, dtype: torch.dtype) -> TypedExpr:
         return TypedExpr(value, dtype)
 
     @staticmethod
-    def index_expr(value: Union[sympy.Expr, int], dtype: torch.dtype) -> TypedExpr:
+    def index_expr(value: sympy.Expr | int, dtype: torch.dtype) -> TypedExpr:
         return TypedExpr(value, dtype)
 
     @staticmethod
     def to_dtype(
         value: TypedExpr,
         dtype: torch.dtype,
-        src_dtype: Optional[torch.dtype] = None,
+        src_dtype: torch.dtype | None = None,
         use_compute_types: bool = False,
     ) -> TypedExpr:
         return TypedExpr(value.expr, dtype)
@@ -143,7 +143,7 @@ def floordiv(x: TypedExpr, y: TypedExpr) -> TypedExpr:
         return TypedExpr(FloorDiv(x.expr, y.expr), result_type)
 
     @staticmethod
-    def mod(x: TypedExpr, y: TypedExpr) -> Optional[TypedExpr]:
+    def mod(x: TypedExpr, y: TypedExpr) -> TypedExpr | None:
         result_type = torch.promote_types(x.dtype, y.dtype)
         if not is_integer_dtype(result_type):
             return NotImplemented
@@ -152,7 +152,7 @@ def mod(x: TypedExpr, y: TypedExpr) -> Optional[TypedExpr]:
         return TypedExpr(result_expr, result_type)
 
     @staticmethod
-    def remainder(x: TypedExpr, y: TypedExpr) -> Optional[TypedExpr]:
+    def remainder(x: TypedExpr, y: TypedExpr) -> TypedExpr | None:
         result_type = torch.promote_types(x.dtype, y.dtype)
         if not is_integer_dtype(result_type):
             return NotImplemented
@@ -195,7 +195,7 @@ def __post_init__(self):
         )
 
 
-IndexPropResult: TypeAlias = Union[IndexPropVar, tuple["IndexPropResult", ...]]
+IndexPropResult: TypeAlias = IndexPropVar | tuple["IndexPropResult", ...]
 
 
 class IndexPropagation(DefaultHandler):
@@ -238,7 +238,7 @@ def materialize_expr(self, expr: sympy.Expr, dtype: torch.dtype) -> Any:
             return self._inner.constant(val, dtype)
         return self._inner.index_expr(expr, dtype)
 
-    def unwrap(self, a: Union[Any, IndexPropVar]) -> Any:
+    def unwrap(self, a: Any | IndexPropVar) -> Any:
         if isinstance(a, (list, tuple)):
             return tuple(self.unwrap(v) for v in a)
 
@@ -281,7 +281,7 @@ def propagate_sympy(
         self, name: str, args: Sequence[Any], kwargs: dict[str, Any]
     ) -> IndexPropResult:
         # Build a new SymPy expression from this ops call
-        def unwrap(a: Union[Any, IndexPropVar]) -> Any:
+        def unwrap(a: Any | IndexPropVar) -> Any:
             if not isinstance(a, IndexPropVar):
                 return a
             return a.value
@@ -337,7 +337,7 @@ def statically_true(self, e):
 
     def indirect_indexing(
         self,
-        index: Union[Any, IndexPropVar],
+        index: Any | IndexPropVar,
         size: Any,
         check: bool = True,
         wrap_neg=True,
diff --git a/torch/_inductor/inductor_prims.py b/torch/_inductor/inductor_prims.py
index e797a622fff19..785b575596295 100644
--- a/torch/_inductor/inductor_prims.py
+++ b/torch/_inductor/inductor_prims.py
@@ -4,10 +4,11 @@
 import functools
 import logging
 import operator
-from typing import Optional, TYPE_CHECKING
+from typing import TYPE_CHECKING
 
 import torch
 from torch import _prims, Tensor
+from torch._utils import _get_device_index
 
 
 if TYPE_CHECKING:
@@ -22,7 +23,7 @@ def make_prim(
     impl_aten,
     return_type=_prims.RETURN_TYPE.NEW,
     doc: str = "",
-    tags: Optional[Sequence[torch.Tag]] = None,
+    tags: Sequence[torch.Tag] | None = None,
 ):
     if isinstance(return_type, tuple):
 
@@ -85,8 +86,8 @@ def eager_prepare_softmax(x: Tensor, dim: int) -> tuple[Tensor, Tensor]:
 # the dtype, so it always faithfully produces a float32 tensor during tracing,
 # even if the default dtype is set to something else.
 random = make_prim(
-    "inductor_random(SymInt[] size, Tensor seed, str mode) -> Tensor",
-    lambda size, seed, mode: getattr(torch, mode)(
+    "inductor_random(SymInt[] size, Tensor seed, str mode, *, ScalarType? align_dtype=None) -> Tensor",
+    lambda size, seed, mode, *, align_dtype=None: getattr(torch, mode)(
         size, device=seed.device, dtype=torch.float32
     ),
     doc="torch.rand()/torch.randn() using backend-specific RNG that can be fused",
@@ -96,6 +97,101 @@ def eager_prepare_softmax(x: Tensor, dim: int) -> tuple[Tensor, Tensor]:
     lambda low, high, size, seed: torch.randint(low, high, size, device=seed.device),
     doc="torch.randint() using backend-specific RNG that can be fused",
 )
+
+
+def _reserve_rng_state(device: torch.device, used_offset):
+    """
+    Reserve `used_offset` 32-bit Philox samples on the given CUDA device and
+    return (seed, base), where base is in Philox-4x32 units.
+
+    This mirrors how Inductor accounts for Philox consumption so compiled
+    dropout kernels can reconstruct eager RNG state.
+    """
+    dev = device if isinstance(device, torch.device) else torch.device(device)
+    if dev.type != "cuda":
+        # Only CUDA devices have Philox-based CUDAGenerator. For non-CUDA
+        # devices this prim should be dead code and never actually run.
+        return 0, 0
+
+    dev_index = _get_device_index(dev, optional=True)
+    if dev_index is None:
+        dev_index = torch.cuda.current_device()
+
+    gen = torch.cuda.default_generators[dev_index]
+    seed_t, off_t, intra_t = torch.ops.inductor_prims.inductor_reserve_rng_state(
+        gen, used_offset
+    )
+
+    # NOTE: for correctness in eager, intra_t should be 0.
+    # Keep everything as tensor math to avoid host sync.
+    if intra_t.device.type != off_t.device.type:
+        intra = int(intra_t.item())
+        base = torch.div(off_t + intra, 4, rounding_mode="floor")
+    else:
+        base = torch.div(off_t + intra_t, 4, rounding_mode="floor")
+    return seed_t, base
+
+
+def _rand_eager_offset_impl(offset, device: torch.device) -> Tensor:
+    """
+    Reserve `offset` 32-bit Philox samples and return a 1-element int64 tensor
+    Place-holder: will be replaced by rand_eager_offsets
+    In fx_passes/replace_random.py
+        fuse_offset_creation_pass()
+    """
+    return torch.empty(2, dtype=torch.int64, device=device)
+
+
+def _rand_eager_offsets_impl(offsets, device: torch.device) -> Tensor:
+    """
+    Batched version of _rand_eager_offset_impl. For each entry in `offsets`,
+    reserve that many 32-bit Philox samples and return a 1D int64 tensor
+    containing the packed (seed, base) values for each reservation.
+    """
+    states = [_reserve_rng_state(device, int(off)) for off in offsets]
+    seeds = [s for s, _ in states]
+    bases = [b for _, b in states]
+
+    def _to_i64(x):
+        if isinstance(x, torch.Tensor):
+            return x
+        return torch.as_tensor(x, device=device, dtype=torch.int64)
+
+    seeds_tensor = torch.stack([_to_i64(x) for x in seeds]).view(-1)
+    bases_tensor = torch.stack([_to_i64(x) for x in bases]).view(-1)
+    packed = torch.stack([seeds_tensor, bases_tensor], dim=1)
+    return packed
+
+
+def _rand_eager_offsets_meta(offsets, device: torch.device):
+    return torch.empty((len(offsets), 2), dtype=torch.int64, device=device)
+
+
+rand_eager_offset = make_prim(
+    "inductor_rand_eager_offset(SymInt offset, Device device) -> Tensor",
+    _rand_eager_offset_impl,
+    doc=(
+        "Reserve `offset` 32-bit Philox samples on `device` and return a "
+        "1-element int64 tensor containing packed (seed, base)."
+    ),
+    tags=(torch.Tag.nondeterministic_seeded,),
+)
+
+
+rand_eager_offsets = _prims._make_prim(
+    schema="inductor_rand_eager_offsets(SymInt[] offsets, Device device) -> Tensor",
+    return_type=_prims.RETURN_TYPE.NEW,
+    meta=_rand_eager_offsets_meta,
+    impl_aten=_rand_eager_offsets_impl,
+    doc=(
+        "Batched version of inductor_rand_eager_offset. For each entry in "
+        "`offsets`, reserves that many 64-bit Philox samples and returns "
+        "packed (seed, base) values."
+    ),
+    tags=(torch.Tag.nondeterministic_seeded,),
+)
+
+
 force_stride_order = make_prim(
     "inductor_force_stride_order(Tensor input, SymInt[] stride) -> Tensor",
     eager_force_stride,
@@ -114,6 +210,15 @@ def eager_prepare_softmax(x: Tensor, dim: int) -> tuple[Tensor, Tensor]:
     doc="Fused multiply add: fma(a, b, c) -> (a * b) + c without rounding after the multiplication",
     tags=(torch.Tag.pointwise,),
 )
+
+# Register DTensor sharding strategies for inductor prims ops.
+# This must happen here (not in _pointwise_ops.py) because the ops
+# don't exist until make_prim() is called above.
+if torch.distributed.is_available():
+    from torch.distributed.tensor._ops._pointwise_ops import register_inductor_prims
+
+    register_inductor_prims()
+
 prepare_softmax_online = make_prim(
     "prepare_softmax_online(Tensor a, int dim) -> (Tensor, Tensor)",
     eager_prepare_softmax,
diff --git a/torch/_inductor/invert_expr_analysis.py b/torch/_inductor/invert_expr_analysis.py
index 816482dba020c..904f073f45090 100644
--- a/torch/_inductor/invert_expr_analysis.py
+++ b/torch/_inductor/invert_expr_analysis.py
@@ -1,5 +1,4 @@
 from dataclasses import dataclass
-from typing import Optional
 
 import sympy
 
@@ -16,14 +15,12 @@ def static_eq(a: _IntLike, b: _IntLike) -> bool:
 @dataclass
 class Term:
     coefficient: _IntLike
-    range: Optional[_IntLike]  # None for unbounded
+    range: _IntLike | None  # None for unbounded
     original_expr: sympy.Expr
     reconstruction_multiplier: _IntLike  # The multiplier needed for reconstruction
 
 
-def generate_inverse_formula(
-    expr: sympy.Expr, var: sympy.Symbol
-) -> Optional[sympy.Expr]:
+def generate_inverse_formula(expr: sympy.Expr, var: sympy.Symbol) -> sympy.Expr | None:
     """
      Analyze an expression to see if it matches a specific invertible pattern that we
      know how to reverse.
@@ -89,7 +86,7 @@ def generate_inverse_formula(
     return generate_reconstruction_expr(terms, var)
 
 
-def parse_terms(expr: sympy.Expr, var: sympy.Symbol) -> Optional[list[Term]]:
+def parse_terms(expr: sympy.Expr, var: sympy.Symbol) -> list[Term] | None:
     """Parse expression into terms."""
     if not isinstance(expr, sympy.Add):
         # Single term
@@ -107,7 +104,7 @@ def parse_terms(expr: sympy.Expr, var: sympy.Symbol) -> Optional[list[Term]]:
     return terms
 
 
-def parse_single_term(term: sympy.Expr, var: sympy.Symbol) -> Optional[Term]:
+def parse_single_term(term: sympy.Expr, var: sympy.Symbol) -> Term | None:
     """Parse a single term and extract coefficient, range, and reconstruction multiplier."""
     # Extract coefficient and expression parts
     coefficient, expr_parts = term.as_coeff_mul()
@@ -141,7 +138,7 @@ def parse_single_term(term: sympy.Expr, var: sympy.Symbol) -> Optional[Term]:
 
 def analyze_expression_properties(
     expr: sympy.Expr, var: sympy.Symbol
-) -> tuple[Optional[_IntLike], Optional[_IntLike]]:
+) -> tuple[_IntLike | None, _IntLike | None]:
     """Analyze an expression to determine its range and reconstruction multiplier."""
     # ModularIndexing(var, divisor, modulo) = (var // divisor) % modulo
     if isinstance(expr, ModularIndexing):
diff --git a/torch/_inductor/ir.py b/torch/_inductor/ir.py
index a45641301f088..435e301860d53 100644
--- a/torch/_inductor/ir.py
+++ b/torch/_inductor/ir.py
@@ -18,7 +18,6 @@
     cast,
     ClassVar,
     Literal,
-    Optional,
     overload,
     SupportsFloat,
     SupportsInt,
@@ -43,10 +42,12 @@
 from torch._higher_order_ops.auto_functionalize import can_auto_functionalize
 from torch._inductor import metrics
 from torch._inductor.utils import get_free_symbols
-from torch._library.opaque_object import is_opaque_type
+from torch._library.fake_class_registry import FakeScriptObject
+from torch._library.opaque_object import get_opaque_obj_repr, is_opaque_value
 from torch._prims_common import (
     compute_required_storage_length,
     is_boolean_dtype,
+    is_contiguous_for_memory_format_or_false,
     is_float_dtype,
     make_channels_last_strides_for,
     StrideType,
@@ -56,6 +57,8 @@
     compute_unbacked_bindings,
     free_symbols,
     free_unbacked_symbols,
+    GuardOnDataDependentSymNode,
+    has_free_unbacked_symbols,
     IterateExprs,
     rebind_unbacked,
     resolve_unbacked_bindings,
@@ -81,6 +84,7 @@
     extract_free_symbols,
     extract_input_node_reduction_ranges,
     extract_read_writes,
+    SymbolUsageCollectorOpsHandler,
     var_builder,
 )
 from .loop_body import LoopBody
@@ -115,9 +119,9 @@
 
 
 if TYPE_CHECKING:
-    from torch._library.fake_class_registry import FakeScriptObject
     from torch.fx.experimental.symbolic_shapes import SympyBoolean
     from torch.fx.node import Argument
+    from torch.types import IntLikeType
 
     from .codegen.cutlass.template import CUTLASSTemplate
     from .codegen.wrapper import PythonWrapperCodegen
@@ -143,10 +147,10 @@
 _U = TypeVar("_U")
 _V = TypeVar("_V")
 
-_IntLike: TypeAlias = Union[int, Expr]
-_NumLike: TypeAlias = Union[int, float, Expr]
+_IntLike: TypeAlias = int | Expr
+_NumLike: TypeAlias = int | float | Expr
 
-_OpOverloads: TypeAlias = Union[torch._ops.OpOverload, torch._ops.HigherOrderOperator]
+_OpOverloads: TypeAlias = torch._ops.OpOverload | torch._ops.HigherOrderOperator
 
 log = logging.getLogger(__name__)
 indent = functools.partial(textwrap.indent, prefix="  ")
@@ -201,12 +205,12 @@
     "Symbol",
     "IRNode",
     Sequence[
-        Optional[Union[int, dict[str, "TensorBox"], "TensorBox", "Symbol", "IRNode"]]
+        Union[int, dict[str, "TensorBox"], "TensorBox", "Symbol", "IRNode"] | None
     ],
 ]
 
 
-def _is_static(x: object) -> TypeIs[Union[int, Integer]]:
+def _is_static(x: object) -> TypeIs[int | Integer]:
     return isinstance(x, (int, Integer))
 
 
@@ -217,7 +221,7 @@ class GraphPartitionSignature:
 
     # mapping from partition input name to IRNode or Expr. Need the name str since
     # we cannot get name from Expr.
-    input_nodes: dict[str, Union[IRNode, sympy.Expr, TorchBindObject]]
+    input_nodes: dict[str, IRNode | sympy.Expr | TorchBindObject]
     output_nodes: list[IRNode]
 
     # mapping from partition input name to a boolean for whether deallocating it
@@ -229,8 +233,8 @@ class GraphPartitionSignature:
     constant_names: list[str]
 
 
-def validate_ir(node_or_nodes: Optional[_NodeOrNodes]) -> None:
-    def _check_tensorbox(nodes: Optional[_NodeOrNodes]) -> None:
+def validate_ir(node_or_nodes: _NodeOrNodes | None) -> None:
+    def _check_tensorbox(nodes: _NodeOrNodes | None) -> None:
         # Could expand this to check deeper properties
         # (e.g. TensorBox points to View or StorageBox)
         if nodes is None:
@@ -254,6 +258,7 @@ def _check_tensorbox(nodes: Optional[_NodeOrNodes]) -> None:
                     int,
                     EffectfulKernel,
                     ShapeAsConstantBuffer,
+                    OpaqueMultiOutput,
                 ),
             ), (
                 f"Found {type(nodes)}, which is not a supported top level IR node. See [Note: Inductor IR]"
@@ -305,7 +310,7 @@ def reindex(index: Sequence[_T]) -> Sequence[_V]:
 
 
 def get_fill_order(
-    seq: Sequence[Union[int, torch.SymInt, Expr]], shape_env: Optional[ShapeEnv] = None
+    seq: Sequence[int | torch.SymInt | Expr], shape_env: ShapeEnv | None = None
 ) -> Sequence[int]:
     """
     Convert strides to fill order (argsort)
@@ -318,7 +323,7 @@ def get_fill_order(
     return sorted_idx
 
 
-def stride_order2fill_order(order: Sequence[Union[int, Integer]]) -> Sequence[int]:
+def stride_order2fill_order(order: Sequence[int | Integer]) -> Sequence[int]:
     """
     Convert stride order to fill order
     For channel last format,
@@ -331,7 +336,7 @@ def stride_order2fill_order(order: Sequence[Union[int, Integer]]) -> Sequence[in
 
 
 def get_stride_order(
-    seq: Sequence[Union[int, torch.SymInt, Expr]], shape_env: Optional[ShapeEnv] = None
+    seq: Sequence[int | torch.SymInt | Expr], shape_env: ShapeEnv | None = None
 ) -> Sequence[int]:
     """
     Convert strides to stride order
@@ -344,22 +349,29 @@ def get_stride_order(
 
 
 @overload
-def ir_node_to_tensor(x: None, guard_shape: bool = True) -> None: ...
+def ir_node_to_tensor(x: None, replace_symbols_with_hints: bool = False) -> None: ...
 
 
 @overload
-def ir_node_to_tensor(x: IRNode, guard_shape: bool = True) -> torch.Tensor: ...
+def ir_node_to_tensor(
+    x: IRNode, replace_symbols_with_hints: bool = False
+) -> torch.Tensor: ...
 
 
 def ir_node_to_tensor(
-    x: Optional[IRNode], guard_shape: bool = True
-) -> Optional[torch.Tensor]:
+    x: IRNode | None, replace_symbols_with_hints: bool = False
+) -> torch.Tensor | None:
+    # When replace_symbols_with_hints=False (default), sizes/strides remain as
+    # symbolic expressions, so downstream operations on the resulting tensor (e.g.,
+    # shape comparisons inside a kernel's meta function) may install guards. When
+    # True, symbolic expressions are replaced with concrete integer hints via
+    # size_hint, preventing any downstream guards.
     if x is None:
         return None
 
-    shape_fn: Callable[[Union[int, Expr]], Union[int, Expr]]
-    if not guard_shape:
-        shape_fn = V.graph.sizevars.size_hint
+    shape_fn: Callable[[int | Expr], int | Expr]
+    if replace_symbols_with_hints:
+        shape_fn = V.graph.sizevars.optimization_hint
     else:
         shape_fn = identity
     size = [shape_fn(s) for s in x.get_size()]
@@ -381,8 +393,8 @@ def ir_node_to_tensor(
 
 
 def may_convert_to_optional(
-    value: Optional[Sequence[_T]],
-) -> Optional[Sequence[Optional[_T]]]:
+    value: Sequence[_T] | None,
+) -> Sequence[_T | None] | None:
     if isinstance(value, list) and not value:
         # [None] makes sure the cpp wrapper codegen will generate something like
         # {std::nullopt} instead of {}
@@ -391,8 +403,8 @@ def may_convert_to_optional(
 
 
 def get_device_type(
-    x: Union[IRNode, OutputSpec, torch.device, None, str],
-) -> Optional[str]:
+    x: IRNode | OutputSpec | torch.device | None | str,
+) -> str | None:
     if isinstance(x, str) or x is None:
         return x
     elif isinstance(x, torch.device):
@@ -403,12 +415,12 @@ def get_device_type(
     assert_never(f"get_device_type({x}: {type(x).__name__})")
 
 
-def is_triton(x: Union[IRNode, torch.device, None, str]) -> bool:
+def is_triton(x: IRNode | torch.device | None | str) -> bool:
     device = get_device_type(x)
     # Special case cpu and cuda as using the method below
     # to determine if the scheduler is a triton scheduler subclass
     # requires instantiating a scheduler for them
-    if device in ["cpu", "cuda"]:
+    if device in ["cpu", "cuda", "xpu"]:
         if getattr(config, f"{device}_backend") == "triton":
             return True
         return False
@@ -423,11 +435,11 @@ def is_triton(x: Union[IRNode, torch.device, None, str]) -> bool:
     return issubclass(device_scheduling, TritonScheduling)
 
 
-def is_cpu(x: Union[IRNode, torch.device, None, str]) -> bool:
+def is_cpu(x: IRNode | torch.device | None | str) -> bool:
     return get_device_type(x) == "cpu"
 
 
-def is_aligned_realized_tensor(x: Union[Buffer, TensorBox], alignment: int) -> bool:
+def is_aligned_realized_tensor(x: Buffer | TensorBox, alignment: int) -> bool:
     if (
         not isinstance(x, IRNode)
         or x.maybe_get_stride() is None
@@ -461,17 +473,14 @@ def significant_strides_equal(
         if V.graph.sizevars.statically_known_leq(dim, 1):
             continue
 
-        if not V.graph.sizevars.statically_known_equals(
-            s1, s2
-        ) and V.graph.sizevars.symbolic_hint(s1) != V.graph.sizevars.symbolic_hint(s2):
+        if not V.graph.sizevars.guard_or_false(sympy.Eq(s1, s2)):
             return False
-
     return True
 
 
 def try_match_insignificant_strides(
     tensor: IRNode,
-    strides: Sequence[Union[int, torch.SymInt]],
+    strides: Sequence[int | torch.SymInt],
 ) -> IRNode:
     """
     Tries to match the strides of the tensor to those in the meta_strides. Strides of insignificant
@@ -546,14 +555,17 @@ class IRNode:
     """
 
     _current_origins: ClassVar[OrderedSet[Any]] = OrderedSet()
+    _current_stream_idx: ClassVar[int | None] = None
 
     # NB: These are kinda weird,
     origins: OrderedSet[Any] = dataclasses.field(init=False)
     # traces back to where the IRNode is created in Inductor
-    traceback: Optional[list[str]] = dataclasses.field(init=False)
-    origin_node: Optional[torch.fx.Node] = dataclasses.field(init=False)
+    traceback: list[str] | None = dataclasses.field(init=False)
+    origin_node: torch.fx.Node | None = dataclasses.field(init=False)
     # Annotations dict for storing metadata (e.g., KernelTemplateChoice)
     annotations: dict[str, Any] = dataclasses.field(init=False)
+    # User-annotated stream index from FX node metadata (set during lowering)
+    stream_idx: int | None = dataclasses.field(init=False)
 
     @staticmethod
     @contextlib.contextmanager
@@ -565,6 +577,18 @@ def current_origins(origins: OrderedSet[Node]) -> Generator[None, None, None]:
         finally:
             IRNode._current_origins = old
 
+    @staticmethod
+    @contextlib.contextmanager
+    def current_stream_idx(
+        stream_idx: int | None,
+    ) -> Generator[None, None, None]:
+        old = IRNode._current_stream_idx
+        IRNode._current_stream_idx = stream_idx
+        try:
+            yield
+        finally:
+            IRNode._current_stream_idx = old
+
     @staticmethod
     def is_realized_node(node: IRNode) -> bool:
         return isinstance(
@@ -578,6 +602,9 @@ def is_realized_node(node: IRNode) -> bool:
             ),
         )
 
+    def wrap_for_lowering(self) -> IRNode:
+        return TensorBox.create(self)
+
     def _post_init_setattr(self, attr: str, value: Any) -> None:
         # Intended for use in __post_init__ for enforcing an invariant on a dataclass
         # If you must, can also be used for setting provenance info
@@ -593,19 +620,24 @@ def __post_init__(self) -> None:
         self._post_init_setattr("origin_node", None)
         # Annotations dict for storing metadata (e.g., KernelTemplateChoice)
         self._post_init_setattr("annotations", {})
+        self._post_init_setattr("stream_idx", self._current_stream_idx)
 
     def get_read_names(self) -> OrderedSet[str]:
         return OrderedSet(dep.name for dep in self.get_reads())
 
-    def get_traceback(self) -> Optional[list[str]]:
+    def get_traceback(self) -> list[str] | None:
         return self.traceback
 
-    def get_origin_node(self) -> Optional[torch.fx.Node]:
+    def get_origin_node(self) -> torch.fx.Node | None:
         return self.origin_node
 
-    def get_defining_op(self) -> Optional[Operation]:
+    def get_defining_op(self) -> Operation | None:
         return None
 
+    def get_subgraphs(self) -> list[Subgraph]:
+        """Return subgraphs contained in this node"""
+        return []
+
     def get_stack_traces(self) -> OrderedSet[str]:
         # Return stack traces to user model code
         # A single IRNode could correspond to multiple lines of code
@@ -669,7 +701,7 @@ def str_helper(
     def get_dtype(self) -> torch.dtype:
         return self.dtype
 
-    def maybe_get_dtype(self) -> Optional[torch.dtype]:
+    def maybe_get_dtype(self) -> torch.dtype | None:
         try:
             return self.get_dtype()
         except NotImplementedError:
@@ -678,7 +710,7 @@ def maybe_get_dtype(self) -> Optional[torch.dtype]:
     def get_layout(self) -> Layout:
         raise NotImplementedError(f"get_layout() is not implemented by {type(self)}!")
 
-    def maybe_get_layout(self) -> Optional[Layout]:
+    def maybe_get_layout(self) -> Layout | None:
         try:
             return self.get_layout()
         except NotImplementedError:
@@ -687,7 +719,7 @@ def maybe_get_layout(self) -> Optional[Layout]:
     def get_output_spec(self) -> OutputSpec:
         return self.get_layout()
 
-    def maybe_get_output_spec(self) -> Optional[OutputSpec]:
+    def maybe_get_output_spec(self) -> OutputSpec | None:
         try:
             return self.get_output_spec()
         except NotImplementedError:
@@ -700,14 +732,14 @@ def has_tensor_output(self) -> bool:
     def get_size(self) -> Sequence[Expr]:
         raise NotImplementedError(f"get_size() is not implemented by {type(self)}!")
 
-    def maybe_get_size(self) -> Optional[Sequence[_IntLike]]:
+    def maybe_get_size(self) -> Sequence[_IntLike] | None:
         try:
             return self.get_size()
         except NotImplementedError:
             return None
 
     @property
-    def shape(self) -> Union[_IntLike, sympy.Rel, Sequence[_IntLike]]:
+    def shape(self) -> _IntLike | sympy.Rel | Sequence[_IntLike]:
         return self.get_size()
 
     def get_numel(self) -> Expr:
@@ -716,7 +748,7 @@ def get_numel(self) -> Expr:
     def is_zero_elements(self) -> bool:
         return V.graph.sizevars.statically_known_true(sympy.Eq(self.get_numel(), 0))
 
-    def realize(self) -> Optional[str]:
+    def realize(self) -> str | None:
         """
         If the IRNode refers to data which has not been materialized (e.g.,
         it is a Pointwise/Reduction that could potentially have more
@@ -734,10 +766,10 @@ def realize(self) -> Optional[str]:
         """
         raise NotImplementedError(f"realize NYI on {type(self)}")
 
-    def codegen_reference(self, writer: Optional[IndentedBuffer] = None) -> str:
+    def codegen_reference(self, writer: IndentedBuffer | None = None) -> str:
         raise NotImplementedError(f"codegen_reference NYI on {type(self)}")
 
-    def get_device(self) -> Optional[torch.device]:
+    def get_device(self) -> torch.device | None:
         return None
 
     def get_device_or_error(self) -> torch.device:
@@ -757,7 +789,7 @@ def make_indexer(self) -> Callable[[Sequence[Expr]], Expr]:
     def get_stride(self) -> Sequence[_IntLike]:
         raise NotImplementedError(type(self).__name__)
 
-    def maybe_get_stride(self) -> Optional[Sequence[_IntLike]]:
+    def maybe_get_stride(self) -> Sequence[_IntLike] | None:
         try:
             return self.get_stride()
         except NotImplementedError:
@@ -766,7 +798,7 @@ def maybe_get_stride(self) -> Optional[Sequence[_IntLike]]:
     def get_name(self) -> str:
         raise NotImplementedError(type(self).__name__)
 
-    def maybe_get_name(self) -> Optional[str]:
+    def maybe_get_name(self) -> str | None:
         try:
             return self.get_name()
         except NotImplementedError:
@@ -778,7 +810,7 @@ def is_input_buffer(self) -> bool:
         except NotImplementedError:
             return False
 
-    def has_large_inner_fn(self, threshold: Optional[int] = None) -> bool:
+    def has_large_inner_fn(self, threshold: int | None = None) -> bool:
         return False
 
     def mark_reuse(self, users: int) -> None:
@@ -826,7 +858,7 @@ def get_free_symbol_uses(
     ) -> OrderedSet[sympy.Symbol]:
         raise NotImplementedError(type(self).__name__)
 
-    def get_reduction_type(self) -> Optional[str]:
+    def get_reduction_type(self) -> str | None:
         raise NotImplementedError(type(self).__name__)
 
     def get_reduction_size(self) -> Sequence[Expr]:
@@ -859,12 +891,13 @@ def dtype(self) -> torch.dtype: ...
 @ir_dataclass(frozen=False)
 class Operation:
     def __post_init__(self) -> None:
-        self.operation_name: Optional[str] = None
+        self.operation_name: str | None = None
+        self._config_patches: dict[str, Any] = {}
 
-    def get_device(self) -> Optional[torch.device]:
+    def get_device(self) -> torch.device | None:
         raise NotImplementedError
 
-    def get_origin_node(self) -> Optional[torch.fx.Node]:
+    def get_origin_node(self) -> torch.fx.Node | None:
         assert hasattr(self, "origin_node")
         return self.origin_node
 
@@ -872,10 +905,22 @@ def get_origins(self) -> OrderedSet[Any]:
         assert hasattr(self, "origins")
         return self.origins
 
+    def get_stream_idx(self) -> int | None:
+        assert hasattr(self, "stream_idx")
+        return self.stream_idx
+
     def get_operation_name(self) -> str:
         assert self.operation_name is not None
         return self.operation_name
 
+    def get_config_patches(self) -> dict[str, Any]:
+        """Get config patches for this operation (e.g., coordinate_descent_tuning)."""
+        return self._config_patches
+
+    def set_config_patches(self, patches: dict[str, Any]) -> None:
+        """Set config patches for this operation."""
+        self._config_patches = patches
+
     def is_extern(self) -> bool:
         return False
 
@@ -964,10 +1009,10 @@ def __str__(self) -> str:
 
     __repr__ = __str__
 
-    def get_device(self) -> Optional[torch.device]:
+    def get_device(self) -> torch.device | None:
         return self.device
 
-    def get_origin_node(self) -> Optional[torch.fx.Node]:
+    def get_origin_node(self) -> torch.fx.Node | None:
         return self.origin_node
 
     def get_size(self) -> Sequence[Expr]:
@@ -1014,7 +1059,7 @@ def inner_fn_str(self) -> str:
             self.inner_fn, *self.inner_fn_args()
         )
 
-    def has_large_inner_fn(self, threshold: Optional[int] = None) -> bool:
+    def has_large_inner_fn(self, threshold: int | None = None) -> bool:
         if threshold is None:
             threshold = 0
         threshold = max(threshold, config.realize_opcount_threshold)
@@ -1024,6 +1069,16 @@ def inner_fn_free_symbols(self, unbacked_only: bool = False) -> OrderedSet[Symbo
         index = self._index(self.ranges)
         return extract_free_symbols(self.inner_fn, index, unbacked_only=unbacked_only)
 
+    # returns a list of op names which references the symbol
+    def collect_inner_fn_symbol_usage(self, symbol: Symbol) -> OrderedSet[str]:
+        index = self._index(self.ranges)
+        handler = SymbolUsageCollectorOpsHandler(symbol)
+        with (
+            V.set_ops_handler(handler),
+        ):
+            self.inner_fn(index)
+        return handler.usages
+
     def get_reads(self) -> OrderedSet[Dep]:
         with patch.object(FlexibleLayout, "allow_indexing", True):
             if self.get_reduction_type():
@@ -1049,7 +1104,7 @@ def get_reduction_size(self) -> Sequence[Expr]:
             f"get_reduction_size() is not implemented by {type(self)}!"
         )
 
-    def get_reduction_type(self) -> Optional[str]:
+    def get_reduction_type(self) -> str | None:
         raise NotImplementedError(
             f"get_reduction_type() is not implemented by {type(self)}!"
         )
@@ -1060,7 +1115,7 @@ def constant_to_device(self, device: torch.device) -> IRNode:
         )
 
 
-def nop_loader_fn(idx: Union[Expr, Sequence[Expr]], *, dtype: torch.dtype) -> OpsValue:
+def nop_loader_fn(idx: Expr | Sequence[Expr], *, dtype: torch.dtype) -> OpsValue:
     if dtype.is_floating_point:
         return ops.constant(float("nan"), dtype)
     else:
@@ -1084,12 +1139,12 @@ def __str__(self) -> str:
     def get_reduction_size(self) -> Sequence[sympy.Expr]:
         return []
 
-    def get_reduction_type(self) -> Optional[str]:
+    def get_reduction_type(self) -> str | None:
         return None
 
     def store_output(
         self,
-        output_name: Optional[str],
+        output_name: str | None,
         indexer: Callable[[Sequence[Expr]], Never],
         vars: Sequence[Expr],
     ) -> None:
@@ -1128,7 +1183,7 @@ def constant_to_device(self, device: torch.device) -> IRNode:
 
     def store_output(
         self,
-        output_name: Optional[str],
+        output_name: str | None,
         indexer: Callable[[Sequence[Expr]], Never],
         vars: Sequence[Expr],
     ) -> Any:
@@ -1239,12 +1294,12 @@ def get_free_symbol_uses(self, unbacked_only: bool = False) -> OrderedSet[Symbol
     def get_reduction_size(self) -> Sequence[Expr]:
         return self.reduction_ranges
 
-    def get_reduction_type(self) -> Optional[str]:
+    def get_reduction_type(self) -> str | None:
         return self.reduction_type
 
     def store_reduction(
         self,
-        output_name: Optional[str],
+        output_name: str | None,
         indexer: Callable[[Sequence[Expr]], Never],
         vars: Sequence[Expr],
         reduction_vars: Sequence[Symbol],
@@ -1295,12 +1350,17 @@ def num_splits(
         inner_fn: Callable[_P, OpsValue],
         ranges: Sequence[_IntLike],
         reduction_ranges: Sequence[_IntLike],
-        reduction_type: Union[ReductionType, Literal["scan"]],
+        reduction_type: ReductionType | Literal["scan"],
         reduction_numel: Expr,
-        input_node: Optional[IRNode] = None,
+        input_node: IRNode | None = None,
     ) -> tuple[ReductionHint, _IntLike]:
-        reduction_numel_hint = V.graph.sizevars.symbolic_hint(reduction_numel)
-        numel_hint = V.graph.sizevars.symbolic_hint(sympy_product(ranges))
+        # Use optimization_hint when all unbacked symbols have explicit hints,
+        # otherwise fall back conservatively.
+        exprs = [reduction_numel, sympy_product(ranges)]
+        if not V.graph.sizevars.all_unbacked_explicitly_hinted(exprs):
+            return ReductionHint.DEFAULT, 1
+        reduction_numel_hint = V.graph.sizevars.optimization_hint(reduction_numel)
+        numel_hint = V.graph.sizevars.optimization_hint(sympy_product(ranges))
 
         should_split = reduction_type == "scan" or (
             not V.graph.has_feature(device, BackendFeature.REDUCE_TO_SINGLE_ELEMENT)
@@ -1312,10 +1372,6 @@ def num_splits(
             and config.split_reductions
         )
 
-        if not (_is_static(reduction_numel_hint) and _is_static(numel_hint)):
-            # We don't support unbacked symints
-            return ReductionHint.DEFAULT, 1
-
         if reduction_type == "dot":
             # Don't split when doing native matmul
             return ReductionHint.DEFAULT, 1
@@ -1353,8 +1409,10 @@ def inner_reduction_splits(
                         new_reduction_ranges,
                     ) = extract_input_node_reduction_ranges(input_node)
                 if new_ranges is not None and new_reduction_ranges is not None:
-                    extracted_numel_hint = V.graph.sizevars.symbolic_hint(
-                        sympy_product(new_ranges + new_reduction_ranges)
+                    extracted_numel_hint = (
+                        V.graph.sizevars.replace_backed_symbols_with_hints(
+                            sympy_product(new_ranges + new_reduction_ranges)
+                        )
                     )
                     if reduction_numel_hint == extracted_numel_hint:
                         log.debug(
@@ -1513,7 +1571,7 @@ def create(
         reduction_ranges: Sequence[Expr],
         reduction_type: ReductionType,
         reduction_hint: ReductionHint = ReductionHint.DEFAULT,
-        input_node: Optional[IRNode] = None,
+        input_node: IRNode | None = None,
     ) -> TensorBox:
         """
         Create a reduction node. May split the reduction to multiple layers to expose
@@ -1525,7 +1583,7 @@ def create(
             # N.B. This is a hack to generate the literal of the given type
             # Ideally, we should be fixing `def constant` in triton.py
             # but it breaks due to hardcoded dtypes in other places
-            def py_cnst(val: object) -> Union[bool, float, int]:
+            def py_cnst(val: object) -> bool | float | int:
                 if dst_dtype == torch.bool:
                     return bool(val)
                 elif dst_dtype.is_floating_point:
@@ -1576,8 +1634,7 @@ def fn(index: int) -> OpsValue:
 
         if (
             isinstance(reduction_numel, Integer)
-            and V.graph.sizevars.size_hint_or_throw(reduction_numel)
-            < config.unroll_reductions_threshold
+            and int(reduction_numel) < config.unroll_reductions_threshold
             and (sympy_product(ranges) != 1 or is_gpu(device.type))
             and reduction_type != "dot"
         ):
@@ -1664,7 +1721,7 @@ def _maybe_increase_split(split: int) -> int:
 
                 def _find_split_reduction(
                     cur_node: TensorBox,
-                ) -> Optional[ComputedBuffer]:
+                ) -> ComputedBuffer | None:
                     read_names = cur_node.get_read_names()
                     if len(read_names) != 1:
                         return None
@@ -1715,7 +1772,7 @@ def _find_split_reduction(
     @staticmethod
     def default_accumulator(
         reduction_type: str, dtype: torch.dtype
-    ) -> Union[_NumLike, Sequence[_NumLike]]:
+    ) -> _NumLike | Sequence[_NumLike]:
         if reduction_type in ("max", "argmax"):
             if is_float_dtype(dtype):
                 return float("-inf")
@@ -1747,7 +1804,7 @@ def default_accumulator(
     @staticmethod
     def default_value(
         reduction_type: str, dtype: torch.dtype
-    ) -> Union[_NumLike, Sequence[_NumLike]]:
+    ) -> _NumLike | Sequence[_NumLike]:
         if reduction_type == "welford_reduce":
             return 0
         return Reduction.default_accumulator(reduction_type, dtype)
@@ -1771,8 +1828,8 @@ def _multilayer_second_step_hint(
 
     @classmethod
     def check_for_split_dense_dim_reindexing(
-        cls, reduction_numel: _IntLike, input_node: Optional[IRNode]
-    ) -> Optional[int]:
+        cls, reduction_numel: _IntLike, input_node: IRNode | None
+    ) -> int | None:
         """
         If we are reducing over the full tensor, and it is non-dense in the last dimension,
         reindex so we reduce over the dense dimension. initially just handle complete
@@ -1809,8 +1866,8 @@ def _multilayer_wrap_loader(
         reduction_numel: _IntLike,
         split: _IntLike,
         block_size: _IntLike,
-        default: Union[_NumLike, Sequence[_NumLike]],
-        input_node: Optional[IRNode] = None,
+        default: _NumLike | Sequence[_NumLike],
+        input_node: IRNode | None = None,
     ) -> Callable[..., object]:
         dense_index = cls.check_for_split_dense_dim_reindexing(
             reduction_numel, input_node
@@ -1819,7 +1876,7 @@ def _multilayer_wrap_loader(
             reduction_ranges, [reduction_numel], dense_index
         )
         need_mask = not V.graph.sizevars.statically_known_true(
-            sympy.Eq(reduction_numel % split, 0)
+            sympy.Eq(Mod(reduction_numel, split), 0)
         )
 
         def wrapper_fn(
@@ -1949,7 +2006,7 @@ def create_multilayer(
         reduction_type: ReductionType,
         split: _IntLike,
         reduction_hint: ReductionHint,
-        input_node: Optional[IRNode] = None,
+        input_node: IRNode | None = None,
     ) -> TensorBox:
         """
         Break a large reduction up into multiple smaller reductions
@@ -2025,7 +2082,7 @@ def create_multilayer_existing_ranges(
 
 def _fixed_indexer(
     size: Sequence[int],
-    stride: Optional[Sequence[int]] = None,
+    stride: Sequence[int] | None = None,
     offset: Expr = Integer(0),
 ) -> Callable[[Sequence[Expr]], Expr]:
     """A closure containing math to read a given element"""
@@ -2052,7 +2109,7 @@ def __init__(
         self,
         device: torch.device,
         dst_dtype: torch.dtype,
-        inner_fns: Union[INNER_FN_TY, Sequence[INNER_FN_TY]],
+        inner_fns: INNER_FN_TY | Sequence[INNER_FN_TY],
         ranges: Sequence[Integer],
         reduction_ranges: Sequence[Integer],
         reduction_type: ReductionType,
@@ -2087,7 +2144,7 @@ def loader(
 
     def store_reduction(
         self,
-        output_name: Optional[str],
+        output_name: str | None,
         indexer: Callable[[Sequence[Expr]], Never],
         vars: Sequence[Expr],
         reduction_vars: Sequence[Symbol],
@@ -2115,7 +2172,7 @@ def create(  # type: ignore[override]
         reduction_ranges: Sequence[Expr],
         num_output: int,
         reduction_hint: ReductionHint = ReductionHint.DEFAULT,
-        input_node: Optional[IRNode] = None,
+        input_node: IRNode | None = None,
     ) -> Sequence[TensorBox]:
         """
         Create the reduction disregarding splitting.
@@ -2154,6 +2211,9 @@ def create(  # type: ignore[override]
         reduction_hint: ReductionHint = ReductionHint.DEFAULT,
     ) -> Sequence[TensorBox]:
         assert reduction_type in ("welford_reduce", "welford_combine")
+        assert not config.mtia.disable_welford_reduction, (
+            "welford reduction usage is explicitly disabled, please check you config"
+        )
 
         reduction_numel = V.graph.sizevars.simplify(sympy_product(reduction_ranges))
 
@@ -2201,7 +2261,7 @@ def inner_fn(idx: Sequence[Expr]) -> OpsValue:
         # TODO: Unrolled reduction
         # if (
         #     isinstance(reduction_numel, Integer)
-        #     and V.graph.sizevars.size_hint(reduction_numel)
+        #     and int(reduction_numel)
         #     < config.unroll_reductions_threshold
         #     and sympy_product(ranges) != 1
         # ):
@@ -2266,7 +2326,7 @@ def inner_fn(idx: Sequence[Expr]) -> OpsValue:
     @staticmethod
     def default_value(
         reduction_type: str, dtype: torch.dtype
-    ) -> Union[_NumLike, Sequence[_NumLike]]:
+    ) -> _NumLike | Sequence[_NumLike]:
         return (0, 0, 0)
 
     @classmethod
@@ -2287,7 +2347,7 @@ def create_multilayer(  # type: ignore[override]
         """
         reduction_numel = sympy_product(reduction_ranges)
         need_mask = not V.graph.sizevars.statically_known_true(
-            sympy.Eq(reduction_numel % split, 0)
+            sympy.Eq(Mod(reduction_numel, split), 0)
         )
 
         if need_mask and reduction_type != "welford_combine":
@@ -2344,7 +2404,9 @@ def intermediate_loader_fn(
         ) -> OpsValue:
             return loader([*index, *reduction_index])
 
-        numel_hint = V.graph.sizevars.size_hint(sympy_product(ranges))
+        # numel_hint is only used to choose between ReductionHint.OUTER vs
+        # OUTER_TINY, which is a performance heuristic, not a correctness decision.
+        numel_hint = V.graph.sizevars.optimization_hint(sympy_product(ranges))
         reduction_hint = cls._multilayer_second_step_hint(
             split, numel_hint, reduction_hint
         )
@@ -2398,7 +2460,7 @@ def __post_init__(self) -> None:
 
     def store_reduction(
         self,
-        output_name: Optional[str],
+        output_name: str | None,
         indexer: Callable[[Sequence[_IntLike]], Never],
         vars: Sequence[Expr],
         scan_vars: Sequence[Symbol],
@@ -2410,7 +2472,7 @@ def store_reduction(
             output_name or "unnamed", indexer(idx), result[self.output_index]
         )
 
-    def get_reduction_type(self) -> Optional[str]:
+    def get_reduction_type(self) -> str | None:
         # return self.scan_op
         return "custom"
 
@@ -2452,7 +2514,7 @@ def create(  # type: ignore[override]
         # Whether we have the option to fallback to aten
         can_fallback_to_aten: bool = True,
         **kwargs: Any,
-    ) -> Sequence[Optional[TensorBox]]:
+    ) -> Sequence[TensorBox | None]:
         pointwise_ranges = [*size[:axis], *size[axis + 1 :]]
         scan_ranges = [size[axis]]
 
@@ -2606,7 +2668,7 @@ def __post_init__(self) -> None:
 
     def store_reduction(
         self,
-        output_name: Optional[str],
+        output_name: str | None,
         indexer: Callable[[Sequence[Expr]], Expr],
         vars: Sequence[Expr],
         reduction_vars: Sequence[Expr],
@@ -2618,7 +2680,7 @@ def store_reduction(
             output_name or "unnamed", indexer(idx), result[self.output_index]
         )
 
-    def get_reduction_type(self) -> Optional[str]:
+    def get_reduction_type(self) -> str | None:
         return "sort"
 
     def get_reduction_size(self) -> Sequence[Expr]:
@@ -2657,7 +2719,7 @@ def create(  # type: ignore[override]
         descending: bool,
         reduction_hint: ReductionHint = ReductionHint.DEFAULT,
         **kwargs: Any,
-    ) -> Sequence[Optional[TensorBox]]:
+    ) -> Sequence[TensorBox | None]:
         pointwise_ranges = [*size[:axis], *size[axis + 1 :]]
         sort_ranges = [size[axis]]
 
@@ -2667,13 +2729,18 @@ def create(  # type: ignore[override]
         sizevars = V.graph.sizevars
         sort_numel = sizevars.simplify(sympy_product(sort_ranges))
 
-        # Heuristic, smallest rblock where triton usually outperforms aten.sort
+        # Heuristic, smallest rblock where triton usually outperforms aten.sort.
         # It also isn't bandwidth bound so fusion is unlikely to help.
-        max_rblock = 512
-        is_persistent_kernel = (
-            config.triton.persistent_reductions
-            and sizevars.statically_known_true(sympy.Le(sort_numel, max_rblock))
-        )
+        # When decompose_sort_ops is enabled, skip the size limit to always
+        # attempt Triton sort (index dtype is widened to int32 in lowering).
+        if config.triton.decompose_sort_ops:
+            is_persistent_kernel = config.triton.persistent_reductions
+        else:
+            max_rblock = 512
+            is_persistent_kernel = (
+                config.triton.persistent_reductions
+                and sizevars.statically_known_true(sympy.Le(sort_numel, max_rblock))
+            )
         if not is_persistent_kernel:
             # We only support persistent triton kernels
             return [None] * len(dtypes)
@@ -2749,9 +2816,9 @@ def as_storage_and_layout(
     x: IRNode,
     freeze: bool = True,
     want_contiguous: bool = False,
-    stride_order: Optional[Sequence[Union[int, Integer]]] = None,
+    stride_order: Sequence[int | Integer] | None = None,
     allow_padding: bool = False,
-    exact_strides: Optional[Sequence[Union[int, Integer]]] = None,
+    exact_strides: Sequence[int | Integer] | None = None,
 ) -> tuple[StorageBox, Layout]:
     """
     Try to simplify x into a StorageBox and a Layout.
@@ -2806,7 +2873,7 @@ def as_storage_and_layout(
 
 
 def is_stride_order_storage_and_layout(
-    x: IRNode, stride_order: Sequence[Union[int, Integer]]
+    x: IRNode, stride_order: Sequence[int | Integer]
 ) -> bool:
     try:
         _buffer, layout = as_storage_and_layout(x, freeze=False)
@@ -2869,10 +2936,10 @@ def dtype(self) -> torch.dtype:
     def get_layout(self) -> Layout:
         return self.data.get_layout()
 
-    def get_device(self) -> Optional[torch.device]:
+    def get_device(self) -> torch.device | None:
         return self.data.get_device()
 
-    def get_origin_node(self) -> Optional[torch.fx.Node]:
+    def get_origin_node(self) -> torch.fx.Node | None:
         return None
 
     def get_name(self) -> str:
@@ -2887,7 +2954,7 @@ def mark_reuse(self, users: int) -> None:
     def has_exceeded_max_reads(self) -> bool:
         return self.data.has_exceeded_max_reads()
 
-    def realize(self) -> Optional[str]:
+    def realize(self) -> str | None:
         return self.data.realize()
 
     def realize_hint(self) -> None:
@@ -2951,7 +3018,9 @@ def _normalize_size(x: IRNode, new_size: Sequence[_IntLike]) -> Sequence[_IntLik
                 old_size[i]
             ):
                 pass
-            else:
+            elif not has_free_unbacked_symbols(
+                old_size
+            ) and not has_free_unbacked_symbols(new_size):
                 # Sanity check: Expect broadcast compatibility
                 #
                 # NB: new_size[i] == old_size[i] is expected to already be
@@ -3076,7 +3145,7 @@ def reindex(
 @ir_dataclass
 class SqueezeView(BaseView):
     @classmethod
-    def create(cls, x: IRNode, *, dim: Optional[int] = None) -> IRNode:
+    def create(cls, x: IRNode, *, dim: int | None = None) -> IRNode:
         if is_storage_and_layout(x):
             storage, old_layout = as_storage_and_layout(x)
             new_size = []
@@ -3230,20 +3299,22 @@ def handle_unbacked_or_dynamic_reshape(
         ) -> IRNode:
             """
             Handle the case where view is not possible with current strides.
-            For unbacked symbols, make contiguous; otherwise use dynamic_reshape_indexer.
+            Try dynamic_reshape_indexer first; if it fails with unbacked
+            symbols (guard_or_false can't resolve comparisons), fall back
+            to making the tensor contiguous.
             """
-            nonlocal old_size, new_size, unbacked_symbols_in_sizes
-            if unbacked_symbols_in_sizes:
-                # For unbacked symbols, we must require contiguous
+            nonlocal old_size, new_size
+            try:
+                reindex = cls.dynamic_reshape_indexer(old_size, new_size)
+                return cls(data=x, size=list(new_size), reindex=reindex)
+            except GuardOnDataDependentSymNode:
                 # dynamic_reshape_indexer cannot handle unbacked SymInts
+                # because guard_or_false can't resolve size comparisons.
                 # https://github.com/pytorch/pytorch/issues/145561
                 x = ExternKernel.require_contiguous(x)
                 return create_reinterpret_view(
                     x, new_size, FlexibleLayout.contiguous_strides(new_size)
                 )
-            # For backed symbols, fall back to dynamic_reshape_indexer
-            reindex = cls.dynamic_reshape_indexer(old_size, new_size)
-            return cls(data=x, size=list(new_size), reindex=reindex)
 
         if 0 in new_size:
 
@@ -3326,11 +3397,11 @@ def dynamic_reshape_indexer(
         cls,
         old_size: Sequence[_IntLike],
         new_size: Sequence[_IntLike],
-        dense_dim: Optional[int] = None,
+        dense_dim: int | None = None,
     ) -> Callable[[Sequence[_T]], Sequence[_V]]:
         try:
             reindex = cls._dynamic_reshape_indexer(old_size, new_size, dense_dim)
-        except (AssertionError, IndexError):
+        except (AssertionError, GuardOnDataDependentSymNode, IndexError):
             # optimistic algorithm failed, lets do a fallback
             flat = [sympy_product(old_size)]
             reindex1 = cls._dynamic_reshape_indexer(old_size, flat)
@@ -3342,12 +3413,47 @@ def dynamic_reshape_indexer(
     def _dynamic_reshape_indexer(
         old_size: Sequence[Expr],
         new_size: Sequence[Expr],
-        dense_dim: Optional[int] = None,
+        dense_dim: int | None = None,
     ) -> Callable[[Sequence[Expr]], Sequence[Expr]]:
         """
         Perform a reshape entirely by modifying indexing math
         """
-        size_hint = V.graph.sizevars.size_hint
+        guard_or_false = V.graph.sizevars.guard_or_false
+
+        def compare_sizes(a: Expr, b: Expr) -> int:
+            """
+            Compare two symbolic sizes, returning -1 if a < b, 0 if a == b, 1 if a > b.
+
+            For unbacked symbols, guard_or_false returns False, so we fall back
+            to divisibility checks.
+            """
+            if guard_or_false(sympy.Eq(a, b)):
+                return 0
+            if guard_or_false(sympy.Lt(a, b)):
+                return -1
+            if guard_or_false(sympy.Gt(a, b)):
+                return 1
+
+            # Divisibility fallback for unbacked symbols:
+            # e.g. comparing u0 vs u0*u1, statically_known_multiple_of(u0*u1, u0)
+            # returns True so we take the merge path (return -1).
+            # The merge reindex for old=[u0, u1] -> new=[u0*u1] is:
+            #   for k in range(u0 * u1):
+            #     i = k // u1
+            #     j = k % u1
+            #     z[k] = x[i, j] + 1
+            # Two cases where this could seem wrong but is still safe:
+            #   u1=1: the merge reindex is correct since the extra dim is
+            #         size 1 (a no-op). e.g. k // 1 = k, k % 1 = 0.
+            #   u0=0: loop is range(0), no kernel runs, result doesn't matter.
+
+            if V.graph.sizevars.statically_known_multiple_of(b, a):
+                return -1
+            if V.graph.sizevars.statically_known_multiple_of(a, b):
+                return 1
+
+            raise GuardOnDataDependentSymNode(sympy.Eq(a, b))
+
         # TODO: These symbols may not escape, if they don't assert so and
         # treat them as temporary
         vars = [
@@ -3377,22 +3483,21 @@ def _dynamic_reshape_indexer(
                 stack_new.append((var, size_new))  # re-add
             elif size_new == 1:
                 stack_old.append(size_old)  # re-add
-            elif size_hint(size_new) == size_hint(size_old):
+            elif compare_sizes(size_new, size_old) == 0:
                 view_expr.append(var)
-                V.graph.sizevars.check_equals(size_new, size_old)
-            elif size_hint(size_new) < size_hint(size_old):
-                while size_hint(size_new) < size_hint(size_old):
+            elif compare_sizes(size_new, size_old) < 0:
+                while compare_sizes(size_new, size_old) < 0:
                     var2, size_new2 = stack_new.pop()
                     var = var2 * size_new + var
                     size_new = size_new * size_new2
                 view_expr.append(var)
                 V.graph.sizevars.check_equals(size_new, size_old)
-            elif size_hint(size_new) > size_hint(size_old):
+            elif compare_sizes(size_new, size_old) > 0:
                 divisor = sympy.S.One
                 modulus = size_old
                 view_expr.append(ModularIndexing(var, divisor, modulus))
                 divisor = divisor * modulus
-                while size_hint(size_new) > size_hint(size_old):
+                while compare_sizes(size_new, size_old) > 0:
                     modulus = stack_old.pop()
                     view_expr.append(ModularIndexing(var, divisor, modulus))
                     divisor = divisor * modulus
@@ -3454,10 +3559,10 @@ def __str__(self) -> str:
     def get_name(self) -> str:
         return self.data.get_name()
 
-    def get_device(self) -> Optional[torch.device]:
+    def get_device(self) -> torch.device | None:
         return self.layout.device
 
-    def get_origin_node(self) -> Optional[torch.fx.Node]:
+    def get_origin_node(self) -> torch.fx.Node | None:
         return None
 
     @property
@@ -3500,7 +3605,7 @@ def get_free_symbol_uses(
             | get_free_symbols(self.layout.offset, unbacked_only)
         )
 
-    def codegen_reference(self, writer: Optional[IndentedBuffer] = None) -> str:
+    def codegen_reference(self, writer: IndentedBuffer | None = None) -> str:
         # reinterpret_tensor is similar to as_strided except:
         # - offset is added to the existing offset (rather than replacing it)
         # - view tracking is disabled similar to unsafe_view
@@ -3560,6 +3665,11 @@ def loader(idx: Sequence[Expr]) -> OpsValue:
 
 
 class SliceView(View):
+    """View that represents a slice along a single dimension.
+
+    Corresponds to tensor[..., start:end:step, ...].
+    """
+
     @classmethod
     def normalize_start_end(
         cls, x: IRNode, dim: int, start: int, end: int
@@ -3574,6 +3684,14 @@ def normalize_start_end(
         if any(free_unbacked_symbols(x) for x in (start, end, dim_size)):
             min_func = sympy.Min
             max_func = sympy.Max
+        elif any(
+            # Only needed when backed_size_oblivious is on.
+            x.has(sympy.Min, sympy.Max)
+            for x in (start, end, dim_size)
+            if isinstance(x, Expr)
+        ):
+            min_func = sympy.Min
+            max_func = sympy.Max
         else:
             min_func = sizevars.evaluate_min
             max_func = sizevars.evaluate_max
@@ -3590,8 +3708,8 @@ def clamp(x: Expr, lower: int, upper: int) -> Expr:
             return clamped_full
 
         def clamp_wrap(
-            val: Union[int, None], lower: int, upper: int, default: Union[Expr, int]
-        ) -> Union[Expr, int]:
+            val: int | None, lower: int, upper: int, default: Expr | int
+        ) -> Expr | int:
             if val is None:
                 # TODO(rec): can this really happen?
                 return default
@@ -3665,10 +3783,10 @@ class BaseConstant(IRNode):
     def get_size(self) -> Sequence[Expr]:
         return ()
 
-    def get_device(self) -> Optional[torch.device]:
+    def get_device(self) -> torch.device | None:
         return self.device
 
-    def get_origin_node(self) -> Optional[torch.fx.Node]:
+    def get_origin_node(self) -> torch.fx.Node | None:
         return None
 
     def get_reads(self) -> OrderedSet[Dep]:
@@ -3687,7 +3805,7 @@ def loader(index: Sequence[Expr]) -> OpsValue:
 
         return loader
 
-    def realize(self) -> Optional[str]:
+    def realize(self) -> str | None:
         pass
 
     def constant_to_device(self, device: torch.device) -> IRNode:
@@ -3738,7 +3856,7 @@ class OutputSpec:
     """Abstract base for Layout, MultiOutputLayout, NoneLayout.
     Represents the memory layout of the output of an Operation."""
 
-    def get_device(self) -> Optional[torch.device]:
+    def get_device(self) -> torch.device | None:
         raise NotImplementedError(type(self).__name__)
 
     def storage_size(self) -> int:
@@ -3764,7 +3882,7 @@ def __init__(
         device: torch.device,
         dtype: torch.dtype,
         size: Sequence[Expr],
-        stride: Optional[Sequence[Expr]] = None,
+        stride: Sequence[Expr] | None = None,
         offset: Expr = Integer(0),
         is_pinned: bool = False,
     ) -> None:
@@ -4145,7 +4263,7 @@ def same_ordered(
         the fill order should be [1, 3, 2, 0]
         """
         assert len(sizes) == len(stride)
-        stride = [V.graph.sizevars.size_hint_or_throw(x) for x in stride]
+        stride = V.graph.sizevars.guarding_hints_or_throw(stride)
         fill_order = sorted(range(len(stride)), key=stride.__getitem__)
         return FlexibleLayout.fill_ordered(sizes, fill_order)
 
@@ -4258,7 +4376,7 @@ def __init__(
         device: torch.device,
         dtype: torch.dtype,
         size: Sequence[Expr],
-        stride_order: Optional[Sequence[Union[int, Integer]]] = None,
+        stride_order: Sequence[int | Integer] | None = None,
         is_pinned: bool = False,
     ) -> None:
         if stride_order:
@@ -4275,7 +4393,7 @@ def __init__(
 class NonOwningLayout(Layout):
     """Is a view into the storage of another tensor"""
 
-    def __init__(self, view: Union[BaseView, TensorBox]) -> None:
+    def __init__(self, view: BaseView | TensorBox) -> None:
         layout = view.get_layout()
         super().__init__(
             layout.device,
@@ -4329,7 +4447,7 @@ class CommBufferLayout(FixedLayout):
 
     def __init__(
         self,
-        layout: Union[FlexibleLayout, FixedLayout],
+        layout: FlexibleLayout | FixedLayout,
         comm_buffer_type: CommBufferType,
         group_name: str,
     ):
@@ -4356,7 +4474,7 @@ class NoneLayout(OutputSpec):
     # If you have an ir.Node with NoneLayout, you probably need to setup
     # dependencies manually in scheduler
 
-    device: Optional[torch.device]
+    device: torch.device | None
     size: list[int] = dataclasses.field(default_factory=lambda: [0])
     stride: list[int] = dataclasses.field(default_factory=lambda: [0])
 
@@ -4366,7 +4484,7 @@ def storage_size(self) -> int:
     def as_fixed(self) -> OutputSpec:
         return self
 
-    def get_device(self) -> Optional[torch.device]:
+    def get_device(self) -> torch.device | None:
         return self.device
 
 
@@ -4463,7 +4581,7 @@ def make_indexer(self) -> Callable[[Sequence[Expr]], Expr]:
 class Buffer(IRNode, CodegenSymbol):
     # Name is sometimes None; e.g., ForceInPlace, where there isn't
     # a meaningful name
-    name: Optional[str]
+    name: str | None
     layout: OutputSpec
 
     # Multi-output buffers will define 'outputs: List[Buffer]'. Confusingly,
@@ -4480,15 +4598,15 @@ def get_name(self) -> str:
         assert self.name, self
         return self.name
 
-    def get_example(self) -> Union[torch.Tensor, torch.SymInt]:
+    def get_example(self) -> torch.Tensor | torch.SymInt:
         if isinstance(self.layout, Layout):
             return self.layout.get_example()
         raise NotImplementedError(type(self.layout).__name__)
 
-    def get_device(self) -> Optional[torch.device]:
+    def get_device(self) -> torch.device | None:
         return self.get_output_spec().get_device()
 
-    def get_defining_op(self) -> Optional[Operation]:
+    def get_defining_op(self) -> Operation | None:
         return None
 
     @property
@@ -4560,7 +4678,7 @@ def loader(index: Sequence[Expr]) -> OpsValue:
 
         return loader
 
-    def codegen_reference(self, writer: Optional[IndentedBuffer] = None) -> str:
+    def codegen_reference(self, writer: IndentedBuffer | None = None) -> str:
         return self.get_name()
 
     def decide_layout(self) -> None:
@@ -4588,7 +4706,7 @@ def get_free_symbol_uses(
     def get_unbacked_symbol_defs(self) -> OrderedSet[sympy.Symbol]:
         return OrderedSet()
 
-    def realize(self) -> Optional[str]:
+    def realize(self) -> str | None:
         pass
 
     def should_allocate(self) -> bool:
@@ -4629,7 +4747,7 @@ class DonatedBuffer(InputBuffer):
 
 
 class ConstantBuffer(InputBuffer):
-    override_device: Optional[torch.device] = None
+    override_device: torch.device | None = None
 
     def make_loader(self) -> Callable[[Sequence[Expr]], OpsValue]:
         def loader(index: Sequence[Expr]) -> OpsValue:
@@ -4658,7 +4776,7 @@ def get_free_symbol_uses(
     ) -> OrderedSet[sympy.Symbol]:
         return OrderedSet()
 
-    def codegen_reference(self, writer: Optional[IndentedBuffer] = None) -> str:
+    def codegen_reference(self, writer: IndentedBuffer | None = None) -> str:
         return V.graph.wrapper_code.none_str
 
     def get_output_spec(self) -> OutputSpec:
@@ -4678,7 +4796,7 @@ def get_free_symbol_uses(
     ) -> OrderedSet[sympy.Symbol]:
         return get_free_symbols(self.expr, unbacked_only)
 
-    def codegen_reference(self, writer: Optional[IndentedBuffer] = None) -> str:
+    def codegen_reference(self, writer: IndentedBuffer | None = None) -> str:
         return V.graph.wrapper_code.codegen_sizevar(self.expr)
 
     def has_tensor_output(self) -> bool:
@@ -4695,10 +4813,10 @@ class ComputedBuffer(OperationBuffer):
     _force_realize: ClassVar[bool] = False
 
     # fields for split reduction
-    _split_size: Optional[int] = None
-    _original_inner_fn: Optional[Callable[..., Any]] = None
-    _original_ranges: Optional[Sequence[_IntLike]] = None
-    _original_reduction_ranges: Optional[Sequence[_IntLike]] = None
+    _split_size: int | None = None
+    _original_inner_fn: Callable[..., Any] | None = None
+    _original_ranges: Sequence[_IntLike] | None = None
+    _original_reduction_ranges: Sequence[_IntLike] | None = None
 
     @contextlib.contextmanager
     def with_original_inner_fn(self) -> Iterator[None]:
@@ -4745,7 +4863,7 @@ def force_realize() -> Iterator[None]:
         finally:
             ComputedBuffer._force_realize = old_value
 
-    def get_computed_buffer_name(self) -> Optional[str]:
+    def get_computed_buffer_name(self) -> str | None:
         """
         Returns self.name if it exists, otherwise returns the name of the data node if that exists.
         If neither exist, returns None.
@@ -4834,7 +4952,7 @@ def get_store_function(self) -> Callable[..., None]:
             assert isinstance(self.data, Pointwise), type(self.data)
             return partial(self.data.store_output, self.name, indexer)
 
-    def get_fill_order(self) -> Optional[list[int]]:
+    def get_fill_order(self) -> list[int] | None:
         """
         If our layout is still flexible, try to determine the stride order based on stride orders of reads.
 
@@ -4916,9 +5034,9 @@ def get_default_sizes_body(
 
     def simplify_and_reorder(
         self,
-        extra_indexing_constraints: Optional[tuple[dict[Any, Any], list[Any]]] = None,
-        recompute_sizes_body_func: Optional[Callable[..., Any]] = None,
-    ) -> tuple[tuple[list[Expr], list[Expr]], Optional[LoopBody]]:
+        extra_indexing_constraints: tuple[dict[Any, Any], list[Any]] | None = None,
+        recompute_sizes_body_func: Callable[..., Any] | None = None,
+    ) -> tuple[tuple[list[Expr], list[Expr]], LoopBody | None]:
         """
         This is a main place where we do loop transformations in a
         backend-agnostic way.
@@ -5071,7 +5189,7 @@ def _apply_loop_reordering(
         support_vars: Sequence[sympy.Symbol],
         sizes: Sequence[int],
         memory_addrs: list[sympy.Expr],
-        priority_idx: Optional[list[int]] = None,
+        priority_idx: list[int] | None = None,
     ) -> tuple[
         list[int],
         Callable[[Sequence[int]], Sequence[int]],
@@ -5111,7 +5229,7 @@ def get_pointwise_size(self) -> Sequence[Expr]:
     def get_reduction_size(self) -> Sequence[Expr]:
         return self.data.get_reduction_size()
 
-    def get_reduction_type(self) -> Optional[str]:
+    def get_reduction_type(self) -> str | None:
         return self.data.get_reduction_type()
 
     def is_no_op(self) -> bool:
@@ -5125,17 +5243,31 @@ def constant_to_device(self, device: torch.device) -> IRNode:
         return self.data.constant_to_device(device)
 
 
+@dataclasses.dataclass(frozen=True)
+class FinalizeCodegenResult:
+    """Structured result from TemplateBuffer._finalize_codegen for external backends."""
+
+    source: str
+    imports: list[str]
+    call_preamble: list[str]
+    call_args: list[str]
+
+
 class TemplateBuffer(OperationBuffer):
     """
-    Represents a Triton (in the future other type) of template operator
-    that we can fuse an epilogue onto.
+    Base class for template operators that support epilogue and prologue fusion.
+    Subclasses: TritonTemplateBuffer (built-in Triton templates),
+    HelionTemplateBuffer (Helion kernels), etc.
     """
 
     def __init__(
         self,
         layout: OutputSpec,
         inputs: Sequence[IRNode],
-        make_kernel_render: Optional[Callable[..., Any]],
+        make_kernel_render: Callable[..., Any] | None,
+        mutated_inputs: Iterable[IRNode] | None = None,
+        allowed_prologue_inps: OrderedSet[str] | None = None,
+        named_inputs: dict[str, IRNode] | None = None,
     ) -> None:
         super().__init__(name=None, layout=layout)
         self.inputs = InputsKernel.unwrap_storage(inputs)
@@ -5145,10 +5277,93 @@ def __init__(
         # Annotations dict for storing metadata (e.g., KernelTemplateChoice)
         self.annotations: dict[str, Any] = {}
 
+        # Output buffer names eligible for epilogue fusion.
+        # Maps buffer name → kernel parameter name (e.g. "buf3" → "result").
+        self.epilogue_fusable_outputs: dict[str, str] = {}
+        # For multi-output kernels: maps child buffer name → MultiOutput
+        # node.  Used by call_kernel to emit tuple-unpacking lines.
+        self._multi_output_children: dict[str, MultiOutput] = {}
+        # Maps kernel parameter name → IRNode for each tensor input.
+        # Used by ExternalTritonTemplateKernel to set up prologue fusion and
+        # by HelionTemplateBuffer to resolve call arguments.
+        self._named_inputs: dict[str, IRNode] = (
+            dict(named_inputs) if named_inputs else {}
+        )
+
+        # Inputs that the kernel mutates in-place
+        self.mutated_inputs = mutated_inputs
+        self.mutation_outputs: list[MutationOutput] = []
+        if mutated_inputs is not None:
+            first_input = self.inputs[0]
+            assert isinstance(first_input, IRNode), type(first_input)
+            device = first_input.get_device()
+            self.mutation_outputs = [
+                MutationOutput(NoneLayout(device=device), buf, self)
+                for buf in mutated_inputs
+            ]
+        # Input buffer names eligible for prologue fusion.
+        self.allowed_prologue_inps: OrderedSet[str] = (
+            allowed_prologue_inps or OrderedSet()
+        )
+        # Per-template fusion overrides.  None means fall back to global
+        # config.epilogue_fusion / config.prologue_fusion.
+        self.allow_epilogue_fusion: bool | None = None
+        self.allow_prologue_fusion: bool | None = None
+
+    @property
+    def dtype(self) -> torch.dtype:
+        if isinstance(self.layout, MultiOutputLayout):
+            raise NotImplementedError(
+                "Multi-output templates do not have a single dtype"
+            )
+        return self.get_layout().dtype
+
     def get_read_writes(self) -> dependencies.ReadWrites:
         return self.extract_read_writes(normalize=True)
 
+    def _read_deps_from_inputs(self, normalize: bool) -> OrderedSet[dependencies.Dep]:
+        """Build read dependencies from all inputs."""
+        reads: OrderedSet[dependencies.Dep] = OrderedSet()
+        for inp_raw in self.inputs:
+            assert isinstance(inp_raw, (ReinterpretView, Buffer)), type(inp_raw)
+            inp: ReinterpretView | Buffer = inp_raw
+            assert isinstance(inp.layout, Layout), type(inp.layout)
+            inp_indexer = inp.layout.make_indexer()
+
+            def dummy(index: Sequence[Any], rindex: Sequence[Any]) -> Any:
+                assert len(rindex) == 0
+                return ops.load(inp.get_name(), inp_indexer(index))
+
+            reads |= dependencies.extract_read_writes(
+                dummy, inp.get_size(), (), normalize=normalize
+            ).reads
+        return reads
+
     def extract_read_writes(self, normalize: bool = False) -> dependencies.ReadWrites:
+        """Extract read/write dependencies for this TemplateBuffer.
+
+        When the layout is MultiOutputLayout (multi-output templates), the
+        buffer itself has no data layout, so we cannot build an indexer.
+        Instead, synthesize a trivial write dep and derive read deps from
+        the named tensor inputs (``_named_inputs``).  For single-output
+        templates with a concrete layout, fall through to the standard path.
+        """
+        if isinstance(self.layout, MultiOutputLayout):
+            writes: OrderedSet[dependencies.Dep] = OrderedSet(
+                [
+                    dependencies.MemoryDep(
+                        self.get_name(), sympy.Integer(0), var_names=(), size=()
+                    ),
+                ]
+            )
+            return dependencies.ReadWrites(
+                reads=self._read_deps_from_inputs(normalize),
+                writes=writes,
+                index_exprs=OrderedSet(),
+                range_vars=None,
+                var_ranges=None,
+            )
+
         name = self.get_name()
         indexer = self.get_layout().make_indexer()
 
@@ -5159,28 +5374,13 @@ def dummy(index: Sequence[Any], rindex: Sequence[Any]) -> Any:
         deps = dependencies.extract_read_writes(
             dummy, self.get_size(), (), normalize=normalize
         )
-
-        for inp in self.inputs:
-            assert isinstance(inp, (ReinterpretView, Buffer)), type(inp)
-            assert isinstance(inp.layout, Layout), type(inp.layout)
-
-            indexer = inp.layout.make_indexer()
-
-            def dummy(index: Sequence[Any], rindex: Sequence[Any]) -> Any:
-                assert len(rindex) == 0
-                # pyrefly: ignore [missing-attribute]
-                return ops.load(inp.get_name(), indexer(index))
-
-            deps.reads |= dependencies.extract_read_writes(
-                dummy, inp.get_size(), (), normalize=normalize
-            ).reads
-
+        deps.reads |= self._read_deps_from_inputs(normalize)
         return deps
 
     def get_reduction_size(self) -> Sequence[Expr]:
         return sympy.S.One
 
-    def get_reduction_type(self) -> Optional[str]:
+    def get_reduction_type(self) -> str | None:
         return None
 
     def should_allocate(self) -> bool:
@@ -5188,9 +5388,9 @@ def should_allocate(self) -> bool:
 
     def simplify_and_reorder(
         self,
-        extra_indexing_constraints: Optional[tuple[dict[Any, Any], list[Any]]] = None,
-        recompute_sizes_body_func: Optional[Callable[..., Any]] = None,
-    ) -> tuple[tuple[Sequence[Expr], list[Expr]], Optional[LoopBody]]:
+        extra_indexing_constraints: tuple[dict[Any, Any], list[Any]] | None = None,
+        recompute_sizes_body_func: Callable[..., Any] | None = None,
+    ) -> tuple[tuple[Sequence[Expr], list[Expr]], LoopBody | None]:
         return (
             (
                 self.get_size(),
@@ -5199,15 +5399,92 @@ def simplify_and_reorder(
             None,
         )
 
+    def is_multi_outputs_template(self) -> bool:
+        """Whether this template produces multiple outputs via MultiOutputLayout."""
+        return isinstance(self.layout, MultiOutputLayout)
+
+    def get_allowed_prologue_inps(self) -> OrderedSet[str]:
+        return self.allowed_prologue_inps
+
+    def _finalize_codegen(
+        self, hook_outputs: dict[str, str]
+    ) -> FinalizeCodegenResult | None:
+        """Called after epilogue/prologue subgraph codegen with rendered hook outputs.
+
+        ``hook_outputs`` maps placeholder keys (e.g. ``<STORE_OUTPUT_0>``,
+        ``<LOAD_INPUT_x>``) to the Triton code generated by Inductor for
+        each fused subgraph.
+
+        Return a ``FinalizeCodegenResult`` to provide custom source code and
+        call metadata, or ``None`` to use the default codegen path.
+        """
+        return None
+
+    @classmethod
+    def realize_template_input(cls, tb: TensorBox) -> IRNode:
+        """Realize a TensorBox, preserving MultiOutput layout (unlike ExternKernel.realize_input)."""
+        if isinstance(tb, TensorBox) and isinstance(tb.data, MultiOutput):
+            return tb.data
+        result = ExternKernel.realize_input(tb)
+        if isinstance(result, StorageBox):
+            result = result.data
+        if isinstance(result.layout, FlexibleLayout):  # type: ignore[union-attr]
+            result.freeze_layout()
+        return result
+
+    @classmethod
+    def build_multi_outputs(
+        cls,
+        template_buf: TemplateBuffer,
+        structured: object,
+        *,
+        direct_alias_at_leaf: dict[int, IRNode] | None = None,
+        on_tensor_leaf: Callable[[str, MultiOutput, list[tuple[type, int]], int], None]
+        | None = None,
+        on_non_tensor_leaf: Callable[[int], None] | None = None,
+    ) -> tuple[TensorBox, ...]:
+        """Walk a structured output tree, creating MultiOutput nodes for tensor leaves."""
+        seen_outputs: dict[int, TensorBox] = {}
+        leaf_counter = itertools.count()
+
+        def walk(output: object, indices: list[tuple[type, int]]) -> list[TensorBox]:
+            if isinstance(output, (list, tuple)):
+                results: list[TensorBox] = []
+                for i, item in enumerate(output):
+                    results.extend(walk(item, [*indices, (type(output), i)]))
+                return results
+            leaf_idx = next(leaf_counter)
+            if isinstance(output, torch.Tensor):
+                if direct_alias_at_leaf and leaf_idx in direct_alias_at_leaf:
+                    return [TensorBox.create(direct_alias_at_leaf[leaf_idx])]
+                tid = id(output)
+                if tid in seen_outputs:
+                    return [seen_outputs[tid]]
+                mo = MultiOutput(
+                    FallbackKernel.tensor_to_layout(output), template_buf, indices
+                )
+                template_buf._multi_output_children[mo.get_name()] = mo
+                if on_tensor_leaf is not None:
+                    on_tensor_leaf(mo.get_name(), mo, indices, leaf_idx)
+                tb = TensorBox(mo)
+                seen_outputs[tid] = tb
+                return [tb]
+            # Non-tensor leaf (int, SymInt, None, etc.)
+            if on_non_tensor_leaf is not None:
+                on_non_tensor_leaf(leaf_idx)
+            return []
+
+        return tuple(walk(structured, []))
+
 
 class TritonTemplateBuffer(TemplateBuffer):
     def __init__(
         self,
         layout: Layout,
         inputs: Sequence[IRNode],
-        make_kernel_render: Optional[Callable[_P, _T]],
-        mutated_inputs: Optional[Iterable[IRNode]] = None,
-        allowed_prologue_inps: Optional[OrderedSet[str]] = None,
+        make_kernel_render: Callable[_P, _T] | None,
+        mutated_inputs: Iterable[IRNode] | None = None,
+        allowed_prologue_inps: OrderedSet[str] | None = None,
     ) -> None:
         """
         NOTE:[TritonTemplates with multiple outputs]
@@ -5218,23 +5495,18 @@ def __init__(
         We work around this by creating an extra input buffer during the lowering
         and we mark them as mutated inputs.
         """
-        super().__init__(layout, inputs, make_kernel_render)
-        self.mutated_inputs = mutated_inputs
-        self.outputs: list[Buffer] = [self]
-        if mutated_inputs is not None:
-            assert isinstance(self.inputs[0], IRNode), type(self.inputs[0])
-            device = self.inputs[0].get_device()
-            self.outputs += [
-                MutationOutput(NoneLayout(device=device), buf, self)
-                for buf in mutated_inputs
-            ]
-
-        self.allowed_prologue_inps = (
-            allowed_prologue_inps if allowed_prologue_inps else OrderedSet()
+        super().__init__(
+            layout,
+            inputs,
+            make_kernel_render,
+            mutated_inputs=mutated_inputs,
+            allowed_prologue_inps=allowed_prologue_inps,
         )
+        assert self.name is not None
+        self.epilogue_fusable_outputs = {self.name: self.name}
 
-        self.subgraph_inps: Optional[list[Optional[Union[IRNode, sympy.Expr]]]] = None
-        self.subgraph_outs: Optional[list[Optional[IRNode]]] = None
+        self.subgraph_inps: list[IRNode | sympy.Expr | None] | None = None
+        self.subgraph_outs: list[IRNode | None] | None = None
 
     @cache_on_self_and_args("TritonTemplateBuffer")
     def get_free_symbol_uses(
@@ -5261,17 +5533,14 @@ def get_free_symbol_uses(
         return res
 
     def get_outputs(self) -> list[Buffer]:
-        return self.outputs
-
-    def get_allowed_prologue_inps(self) -> OrderedSet[str]:
-        return self.allowed_prologue_inps
+        return [self, *self.mutation_outputs]
 
     def __str__(self) -> str:
         out = f"TritonTemplateBuffer(layout={self.layout})"
         return out
 
 
-PrimitiveInfoType = Union[int, float, bool, str, list[Union[int, str, float, bool]]]
+PrimitiveInfoType = int | float | bool | str | list[int | str | float | bool]
 
 
 class ChoiceCaller:
@@ -5298,13 +5567,22 @@ def __init__(
         # knowing what autotuning is choosing)
         self.description = description
         self.failed: bool = False
+        # When True, benchmark using CUDA graph capture/replay
+        self._benchmark_with_cudagraphs: bool = False
         # A place to store annotations that can be read post benchmarking
         # Use this to shuttle information between ChoieCaller generation
         # and the end of benchmarking
         self.annotations: dict[Any, Any] = {}
+        # Subclass-overridden attributes for subgraph-based choices
+        self.gm: torch.fx.GraphModule | None = None
+        self.decomposition: Callable[..., Any] | None = None
+        self.decomposition_kwargs: dict[str, Any] = {}
+        self.config_patches: dict[str, Any] = {}
 
     def benchmark(self, *args: Any, out: torch.Tensor) -> float:
         algo = self.to_callable()
+        if self._benchmark_with_cudagraphs:
+            return benchmarker.benchmark_gpu_with_cuda_graph(lambda: algo(*args))
         if config.profile_bandwidth_with_do_bench_using_profiling:
             return do_bench_using_profiling(lambda: algo(*args))  # type: ignore[arg-type]
         return benchmarker.benchmark(algo, args, {"out": out}, device=None)
@@ -5328,7 +5606,7 @@ def hash_key(self) -> str:
     def output_node(self) -> TensorBox:
         raise NotImplementedError
 
-    def info_dict(self) -> dict[str, Union[PrimitiveInfoType, list[PrimitiveInfoType]]]:
+    def info_dict(self) -> dict[str, PrimitiveInfoType | list[PrimitiveInfoType]]:
         """Information returned here is logged to the autotune log file when that is enabled."""
         return {}
 
@@ -5362,7 +5640,7 @@ def __init__(
         self,
         layout: Layout,
         inputs: Sequence[IRNode],
-        choice_timings_fn: Callable[[Optional[int]], dict[ChoiceCaller, float]],
+        choice_timings_fn: Callable[[int | None], dict[ChoiceCaller, float]],
         unfiltered_choices: list[ChoiceCaller],
         allowed_prologue_inps: OrderedSet[str],
     ) -> None:
@@ -5373,7 +5651,7 @@ def __init__(
             allowed_prologue_inps=allowed_prologue_inps,
         )
         self._choice_timings_fn = choice_timings_fn
-        self._choice_timings: dict[Optional[int], dict[ChoiceCaller, float]] = {}
+        self._choice_timings: dict[int | None, dict[ChoiceCaller, float]] = {}
         self._choices: list[ChoiceCaller] = unfiltered_choices
         self.original_inputs = inputs
         self._output_plannable = all(
@@ -5384,7 +5662,7 @@ def __init__(
             )
             for choice in unfiltered_choices
         )
-        self._make_kernel_renders: dict[Optional[int], Any] = {}
+        self._make_kernel_renders: dict[int | None, Any] = {}
 
     @property
     def output_plannable(self) -> bool:
@@ -5398,7 +5676,7 @@ def choices(self) -> list[ChoiceCaller]:
         return self._choices
 
     def choice_timings(
-        self, hint_override: Optional[int] = None
+        self, hint_override: int | None = None
     ) -> dict[ChoiceCaller, float]:
         if hint_override not in self._choice_timings:
             self._choice_timings[hint_override] = self._choice_timings_fn(hint_override)
@@ -5427,14 +5705,14 @@ def finalize_as_triton_caller(self, caller: TritonTemplateCallerBase) -> None:
         self.make_kernel_render = caller.get_make_kernel_render()
 
     def get_min_choice(
-        self, hint_override: Optional[int] = None
+        self, hint_override: int | None = None
     ) -> tuple[ChoiceCaller, float]:
         timings = self.choice_timings(hint_override=hint_override)
         min_choice = min(timings, key=timings.get)  # type: ignore[arg-type]
         return (min_choice, timings[min_choice])
 
     def finalize_as_triton_callers(
-        self, callers: dict[Optional[int], TritonTemplateCallerBase]
+        self, callers: dict[int | None, TritonTemplateCallerBase]
     ) -> None:
         """Finalize with multiple callers for different hint overrides"""
         for hint_override, caller in callers.items():
@@ -5480,7 +5758,7 @@ def __init__(
         super().__init__(layout, inputs, make_kernel_render)
         self.template = template
         self.choice = choice
-        self.outputs: Optional[list[Buffer]] = None
+        self.outputs: list[Buffer] | None = None
 
     def get_layout(self) -> Layout:
         if isinstance(self.layout, MultiOutputLayout):
@@ -5507,7 +5785,7 @@ def __init__(
         inputs: Sequence[IRNode],
         make_kernel_render: Callable[_P, _T],
         template: Any,
-        mutated_inputs: Optional[Iterable[IRNode]] = None,
+        mutated_inputs: Iterable[IRNode] | None = None,
     ) -> None:
         super().__init__(layout, inputs, make_kernel_render)
         self.template = template
@@ -5542,10 +5820,10 @@ def __init__(
         accumulator_type: Any,
         variant: Any,  # GemmVariant, use Any to avoid circular import
         workspace_size: int = 0,
-        scale_type_a: Optional[Any] = None,
-        scale_type_b: Optional[Any] = None,
-        swizzle_type_a: Optional[Any] = None,
-        swizzle_type_b: Optional[Any] = None,
+        scale_type_a: Any | None = None,
+        scale_type_b: Any | None = None,
+        swizzle_type_a: Any | None = None,
+        swizzle_type_b: Any | None = None,
     ) -> None:
         # We pass None initially, then override with our method below
         super().__init__(layout, inputs, make_kernel_render=None)
@@ -5575,7 +5853,7 @@ def get_outputs(self) -> list[Buffer]:
         return self.outputs
 
     def _make_kernel_render(
-        self, out_node: Any, hint_override: Optional[int] = None
+        self, out_node: Any, hint_override: int | None = None
     ) -> tuple[Any, Any]:
         """
         Create a kernel renderer for code generation.
@@ -5620,14 +5898,14 @@ def render():
 
 
 def is_node_sequence(
-    nodes: Sequence[Union[IRNode, Sequence[IRNode]]],
+    nodes: Sequence[IRNode | Sequence[IRNode]],
 ) -> TypeIs[Sequence[IRNode]]:
     return all(isinstance(n, IRNode) for n in nodes)
 
 
 @ir_dataclass(frozen=False)
 class InputsKernel(OperationBuffer):
-    inputs: Sequence[Union[IRNode, Sequence[IRNode]]]
+    inputs: Sequence[IRNode | Sequence[IRNode]]
 
     def input_name(self, i: int) -> str:
         input = self.inputs[i]
@@ -5680,9 +5958,9 @@ def unwrap_storage_for_input(cls, x: IRNode) -> IRNode:
 
     @staticmethod
     def unwrap_storage(
-        inputs: Sequence[Union[IRNode, Sequence[IRNode]]],
-    ) -> list[Union[IRNode, Sequence[IRNode]]]:
-        inputs_new: list[Union[IRNode, Sequence[IRNode]]] = []
+        inputs: Sequence[IRNode | Sequence[IRNode]],
+    ) -> list[IRNode | Sequence[IRNode]]:
+        inputs_new: list[IRNode | Sequence[IRNode]] = []
         for x in inputs:
             if isinstance(x, Sequence):
                 x = [InputsKernel.unwrap_storage_for_input(i) for i in x]
@@ -5771,18 +6049,24 @@ def create(cls, inputs: Sequence[IRNode], dim: int) -> StorageBox:
                     break
         any_input_is_storage_and_layout = any(is_storage_and_layout(x) for x in inputs)
         fx_node_args = V.graph.current_node.args[0]
-        assert isinstance(fx_node_args, list), type(fx_node_args)
         # If any of the inputs has meta tensor and the meta tensor is in CL format, use CL format for the output
-        if any_input_is_storage_and_layout is False and any(
-            # pyrefly: ignore [missing-attribute]
-            "val" in arg.meta
-            and (
-                # pyrefly: ignore [missing-attribute]
-                arg.meta["val"].is_contiguous(memory_format=torch.channels_last)
+        # Skip this check when fx_node_args is not a list (e.g., called from _pad_as_cat).
+        if (
+            any_input_is_storage_and_layout is False
+            and isinstance(fx_node_args, list)
+            and any(
                 # pyrefly: ignore [missing-attribute]
-                or arg.meta["val"].is_contiguous(memory_format=torch.channels_last_3d)
+                "val" in arg.meta
+                and (
+                    # pyrefly: ignore [missing-attribute]
+                    arg.meta["val"].is_contiguous(memory_format=torch.channels_last)
+                    # pyrefly: ignore [missing-attribute]
+                    or arg.meta["val"].is_contiguous(
+                        memory_format=torch.channels_last_3d
+                    )
+                )
+                for arg in fx_node_args
             )
-            for arg in fx_node_args
         ):
             output_stride = make_channels_last_strides_for(new_size)
 
@@ -5841,7 +6125,7 @@ def create(cls, inputs: Sequence[IRNode], dim: int) -> StorageBox:
 
     @classmethod
     def can_realize_into_without_copy(
-        cls, src: IRNode, dst: Optional[IRNode] = None
+        cls, src: IRNode, dst: IRNode | None = None
     ) -> bool:
         if isinstance(src, TensorBox):
             # unwrap a TensorBox
@@ -5928,20 +6212,20 @@ class ExternKernel(InputsKernel):
 
     constant_args: Sequence[Any] = ()
     kwargs: dict[str, Any] = dataclasses.field(default_factory=dict)
-    output_view: Optional[ReinterpretView] = None
-    python_kernel_name: Optional[str] = None
-    cpp_kernel_name: Optional[str] = None
+    output_view: ReinterpretView | None = None
+    python_kernel_name: str | None = None
+    cpp_kernel_name: str | None = None
     # FIXME: in some cases we sill need to explicitly pass in ordered_kwargs_for_cpp_kernel
     # We shouldn't need to do this since the information can be retrieved from op_overload._schema.
     ordered_kwargs_for_cpp_kernel: Iterable[str] = dataclasses.field(
         default_factory=list
     )
-    op_overload: Optional[_OpOverloads] = None
-    arg_properties: Optional[list[dict[str, Any]]] = None
+    op_overload: _OpOverloads | None = None
+    arg_properties: list[dict[str, Any]] | None = None
     allarg_properties: dict[str, dict[str, Any]] = dataclasses.field(
         default_factory=dict
     )
-    kwarg_properties: Optional[dict[str, dict[str, Any]]] = None
+    kwarg_properties: dict[str, dict[str, Any]] | None = None
     unbacked_bindings: dict[sympy.Symbol, pytree.KeyPath] = dataclasses.field(
         default_factory=dict
     )
@@ -5949,16 +6233,16 @@ class ExternKernel(InputsKernel):
 
     def __init__(
         self,
-        name: Optional[str],
+        name: str | None,
         layout: OutputSpec,
-        inputs: Sequence[Union[IRNode, Sequence[IRNode]]],
+        inputs: Sequence[IRNode | Sequence[IRNode]],
         constant_args: Sequence[Any] = (),
         kwargs: dict[str, Any] | None = None,
-        output_view: Optional[ReinterpretView] = None,
-        python_kernel_name: Optional[str] = None,
-        cpp_kernel_name: Optional[str] = None,
+        output_view: ReinterpretView | None = None,
+        python_kernel_name: str | None = None,
+        cpp_kernel_name: str | None = None,
         ordered_kwargs_for_cpp_kernel: Iterable[str] = (),
-        op_overload: Optional[_OpOverloads] = None,
+        op_overload: _OpOverloads | None = None,
     ) -> None:
         super().__init__(
             name=name,
@@ -6028,7 +6312,7 @@ def decide_layout(self) -> None:
             self.freeze_layout()
 
     def codegen_comment(
-        self, wrapper: PythonWrapperCodegen, kernel_name: Optional[str] = None
+        self, wrapper: PythonWrapperCodegen, kernel_name: str | None = None
     ) -> None:
         origin_str, _detailed_origin_str = get_kernel_metadata(self, wrapper)
         if origin_str:
@@ -6047,7 +6331,7 @@ def codegen_comment(
     def codegen(self, wrapper: PythonWrapperCodegen) -> None:
         raise NotImplementedError
 
-    def set_cpp_kernel_name(self, cpp_kernel_name: Optional[str] = None) -> None:
+    def set_cpp_kernel_name(self, cpp_kernel_name: str | None = None) -> None:
         self.cpp_kernel_name = cpp_kernel_name
         if not V.graph.cpp_wrapper or not isinstance(
             self.op_overload, torch._ops.OpOverload
@@ -6067,11 +6351,24 @@ def set_cpp_kernel_name(self, cpp_kernel_name: Optional[str] = None) -> None:
                     if kernel._overloadname == "default"
                     else kernel.__name__.replace(".", "_")
                 )
+                # If the op has a versioned c_shim entry, call the latest _v{N}
+                # variant so new AOTI artifacts match the current op schema. The
+                # unversioned shim is retained solely for BC with existing
+                # already-compiled artifacts.
+                from torchgen.aoti.fallback_ops import inductor_fallback_ops
+
+                version_info = inductor_fallback_ops.get(f"aten.{kernel.__name__}", {})
+                latest_version = max(
+                    (int(v[1:]) for v in version_info if v.startswith("v")),
+                    default=1,
+                )
+                if latest_version > 1:
+                    opname = f"{opname}_v{latest_version}"
                 self.cpp_kernel_name = f"at::_ops::{opname}::call"
             else:
                 self.cpp_kernel_name = kernel._schema.name
 
-    def set_python_kernel_name(self, python_kernel_name: Optional[str]) -> None:
+    def set_python_kernel_name(self, python_kernel_name: str | None) -> None:
         self.python_kernel_name = python_kernel_name
         if python_kernel_name is not None:
             return
@@ -6086,7 +6383,7 @@ def set_python_kernel_name(self, python_kernel_name: Optional[str]) -> None:
                 f"{kernel.__module__.replace('._ops.', '.ops.')}.{kernel.__name__}"
             )
 
-    def try_get_kernel_name(self) -> Optional[str]:
+    def try_get_kernel_name(self) -> str | None:
         from .codegen.cpp_wrapper_cpu import CppWrapperCpu
 
         device = d.type if (d := self.get_device()) else V.graph.device_type
@@ -6130,26 +6427,60 @@ def process_kernel(
         list[Any],
         list[Any],
         Callable[[Any, Any], Any],
-        Optional[dict[sympy.Symbol, pytree.KeyPath]],
+        dict[sympy.Symbol, pytree.KeyPath] | None,
     ]:
+        """Partition kernel args into tensor and non-tensor, realize tensor inputs,
+        re-run fake tensor propagation with the realized strides, and return
+        (example_output, tensor_args, non_tensor_args, unflatten_args, unbacked_bindings).
+
+        unflatten_args(new_tensor_args, new_non_tensor_args) reconstructs the
+        original (args, kwargs) tree from replacement lists.
+        """
         binded_args = {"args": args, "kwargs": kwargs}
 
         args_flat, args_spec = pytree.tree_flatten(binded_args)
 
-        is_arg_tensor = []
+        args_flat_is_tensor: list[bool] = []
         # tensor_args can be either tensor or torchbind objects
-        tensor_args = []
-        non_tensor_args: list[Any] = []
+        tensor_args: list[IRNode] = []
+        non_tensor_args: list[object] = []
+        real_non_tensor_args: list[
+            FakeScriptObject
+            | torch._C.Generator
+            | torch._C.ScriptObject
+            | torch.Tensor
+            | IntLikeType
+        ] = []
         for arg in args_flat:
-            is_arg_tensor.append(
-                isinstance(arg, IRNode) and not isinstance(arg, GeneratorState)
-            )
-            if is_arg_tensor[-1]:
-                tensor_args.append(arg)
-            else:
-                if isinstance(arg, Expr):
-                    arg = V.graph.sizevars.shape_env.create_symintnode(arg, hint=None)
-                non_tensor_args.append(arg)
+            match arg:
+                case Expr():
+                    node = V.graph.sizevars.shape_env.create_symintnode(arg, hint=None)
+                    args_flat_is_tensor.append(False)
+                    non_tensor_args.append(node)
+                    real_non_tensor_args.append(node)
+
+                case GeneratorState():
+                    args_flat_is_tensor.append(False)
+                    non_tensor_args.append(arg)
+                    device_index = arg.device.index
+                    assert arg.device.type == "cuda" and device_index is not None
+                    real_non_tensor_args.append(
+                        torch.cuda.default_generators[device_index].clone_state()
+                    )
+
+                case OpaqueObjectState():
+                    args_flat_is_tensor.append(False)
+                    non_tensor_args.append(arg)
+                    real_non_tensor_args.append(arg.value)
+
+                case IRNode():
+                    args_flat_is_tensor.append(True)
+                    tensor_args.append(arg)
+
+                case _:
+                    args_flat_is_tensor.append(False)
+                    non_tensor_args.append(arg)
+                    real_non_tensor_args.append(arg)
 
         def unflatten_args(
             new_tensor_args: Sequence[_T], new_non_tensor_args: Sequence[_T]
@@ -6157,7 +6488,7 @@ def unflatten_args(
             result = []
             it_tensors = iter(new_tensor_args)
             it_non_tensors = iter(new_non_tensor_args)
-            for is_tensor in is_arg_tensor:
+            for is_tensor in args_flat_is_tensor:
                 if is_tensor:
                     result.append(next(it_tensors))
                 else:
@@ -6177,9 +6508,7 @@ def unflatten_args(
         # strides of inputs and we need to determine accurately what the
         # output stride will be.
         example_args: list[
-            Union[
-                torch.Tensor, torch._C.ScriptObject, FakeScriptObject, torch.Generator
-            ]
+            torch.Tensor | torch._C.ScriptObject | FakeScriptObject | torch.Generator
         ] = []
 
         # We need to retain the constant values of fake tensors that we originally
@@ -6197,6 +6526,8 @@ def unflatten_args(
                 example_args.append(V.graph.torchbind_constants[x.get_name()])
             elif isinstance(x, TorchBindObject):
                 example_args.append(x.get_value())
+            elif isinstance(x, OpaqueMultiOutput):
+                example_args.append(x.opaque_example_value)
             elif isinstance(x, torch._inductor.ir.GeneratorState):
                 device_index = x.device.index
                 assert x.device.type == "cuda" and device_index is not None
@@ -6204,12 +6535,12 @@ def unflatten_args(
                     torch.cuda.default_generators[device_index].clone_state()
                 )
             else:
-                example_args.append(ir_node_to_tensor(x, guard_shape=True))
+                example_args.append(ir_node_to_tensor(x))
 
-        new_args, new_kwargs = unflatten_args(example_args, non_tensor_args)
+        new_args, new_kwargs = unflatten_args(example_args, real_non_tensor_args)
         example_output = kernel(*new_args, **new_kwargs)
 
-        unbacked_bindings: Optional[dict[sympy.Symbol, pytree.KeyPath]] = None
+        unbacked_bindings: dict[sympy.Symbol, pytree.KeyPath] | None = None
         if shape_env := V.fake_mode.shape_env:
             node_meta_val = V.current_node.meta.get("val")
             ctx: AbstractContextManager[None] = nullcontext()
@@ -6273,11 +6604,13 @@ def convert_to_reinterpret_view(cls, x: IRNode) -> ReinterpretView:
             and isinstance(x_unwrap_view, (ReinterpretView, Buffer, MutableBox))
             and isinstance(x_unwrap_view.layout, FlexibleLayout)
             and (
-                x_unwrap_view_fx_node.meta["val"].is_contiguous(
-                    memory_format=torch.channels_last
+                is_contiguous_for_memory_format_or_false(
+                    x_unwrap_view_fx_node.meta["val"],
+                    memory_format=torch.channels_last,
                 )
-                or x_unwrap_view_fx_node.meta["val"].is_contiguous(
-                    memory_format=torch.channels_last_3d
+                or is_contiguous_for_memory_format_or_false(
+                    x_unwrap_view_fx_node.meta["val"],
+                    memory_format=torch.channels_last_3d,
                 )
             )
         ):
@@ -6351,7 +6684,7 @@ def realize_input(cls, x: IRNode) -> IRNode:
             # TODO(jansel): impose layout preference on realized buffer
             x.realize()
             return x
-        if isinstance(x, (NonTensorObj, ShapeAsConstantBuffer)):
+        if isinstance(x, (NonTensorObj, ShapeAsConstantBuffer, OpaqueMultiOutput)):
             return x
         return cls.copy_input(x)
 
@@ -6369,8 +6702,8 @@ def require_stride1(cls, x: IRNode) -> IRNode:
     def require_strides(
         cls,
         x: IRNode,
-        order: Optional[Sequence[int]] = None,
-        exact_strides: Optional[Sequence[_IntLike]] = None,
+        order: Sequence[int] | None = None,
+        exact_strides: Sequence[_IntLike] | None = None,
         allow_padding: bool = False,
     ) -> IRNode:
         assert order is not None or exact_strides is not None
@@ -6400,7 +6733,7 @@ def require_strides(
                         want_contiguous=False,
                         stride_order=(
                             get_stride_order(
-                                V.graph.sizevars.size_hints_or_throw(
+                                V.graph.sizevars.guarding_hints_or_throw(
                                     x.get_layout().stride
                                 )
                             )
@@ -6493,7 +6826,7 @@ def require_strides(
         # across expanded dims.
         # TODO: could also be good to have a codegen fix to recognize overlapping elements
 
-        expanded_dims: Optional[list[int]] = None
+        expanded_dims: list[int] | None = None
         orig_size = x.get_size()
         if exact_strides is not None:
             sizevars = V.graph.sizevars
@@ -6620,7 +6953,7 @@ def fill_non_provided_args(
                 )
         return args
 
-    def codegen_const_args(self, names: Optional[list[str]] = None) -> list[str]:
+    def codegen_const_args(self, names: list[str] | None = None) -> list[str]:
         if V.graph.cpp_wrapper:
             result = []
             # Aten ops follow the convention that tensor args are before non-tensor args,
@@ -6783,7 +7116,9 @@ def canonicalize(self) -> tuple[Expr, Sequence[Expr]]:
         sizevars = V.graph.sizevars
         sizes = self.get_size()
         strides = self.get_stride()
-        strides = [sizevars.size_hint(x) for x in strides]
+        # Stride hints are only used as sort keys to determine dimension
+        # ordering for canonicalization, not for correctness.
+        strides = [sizevars.optimization_hint(x) for x in strides]
         # TODO: I can't tell if the symbols here are temporary
         index_vars = [sympy_index_symbol(f"d{i}") for i in range(len(sizes))]
         # reorder index vars according to stride
@@ -6847,12 +7182,12 @@ def __init__(
         layout: Layout,
         inputs: Sequence[IRNode],
         constant_args: Sequence[Any] = (),
-        kwargs: Optional[dict[str, Any]] = None,
-        output_view: Optional[ReinterpretView] = None,
-        python_kernel_name: Optional[str] = None,
-        cpp_kernel_name: Optional[str] = None,
+        kwargs: dict[str, Any] | None = None,
+        output_view: ReinterpretView | None = None,
+        python_kernel_name: str | None = None,
+        cpp_kernel_name: str | None = None,
         ordered_kwargs_for_cpp_kernel: Sequence[Any] = (),
-        op_overload: Optional[_OpOverloads] = None,
+        op_overload: _OpOverloads | None = None,
     ) -> None:
         unwrapped_inputs = self.unwrap_storage(inputs)
         assert isinstance(unwrapped_inputs, Sequence), type(unwrapped_inputs)
@@ -6904,11 +7239,11 @@ def __init__(
         layout: OutputSpec,
         inputs: Sequence[IRNode],
         constant_args: Sequence[Any] = (),
-        kwargs: Optional[dict[str, Any]] = None,
-        python_kernel_name: Optional[str] = None,
-        cpp_kernel_name: Optional[str] = None,
+        kwargs: dict[str, Any] | None = None,
+        python_kernel_name: str | None = None,
+        cpp_kernel_name: str | None = None,
         ordered_kwargs_for_cpp_kernel: Sequence[Any] = (),
-        op_overload: Optional[_OpOverloads] = None,
+        op_overload: _OpOverloads | None = None,
     ) -> None:
         unwrapped_inputs = self.unwrap_storage(inputs)
         assert all(isinstance(i, IRNode) for i in unwrapped_inputs)
@@ -7046,9 +7381,9 @@ class TMADescriptorExperimental(TMADescriptor):
     def __init__(
         self,
         tensor: IRNode,
-        dims: list[Union[int, torch.SymInt]],
-        block_dims: list[Union[int, torch.SymInt]],
-        element_size: Optional[int] = None,
+        dims: list[int | torch.SymInt],
+        block_dims: list[int | torch.SymInt],
+        element_size: int | None = None,
     ) -> None:
         assert len(dims) in (1, 2)
         assert len(dims) == len(block_dims)
@@ -7083,7 +7418,7 @@ class TMADescriptorStable(TMADescriptor):
     See also TMADescriptorExperimental for the old API.
     """
 
-    def __init__(self, tensor: IRNode, block_shape: list[Union[int, torch.SymInt]]):
+    def __init__(self, tensor: IRNode, block_shape: list[int | torch.SymInt]):
         self.block_shape = block_shape
 
         super().__init__(
@@ -7101,6 +7436,7 @@ def __init__(
         gm: torch.fx.GraphModule,
         example_inputs: list[Any],
         subgraph_name: str,
+        config_patches: dict[str, Any] | None = None,
     ):
         super().__init__(None, layout, input_nodes)
         self.gm = gm
@@ -7122,14 +7458,23 @@ def __init__(
         import torch._inductor.config as inductor_config
 
         with V.set_graph_handler(self.subgraph):
-            # Don't bother autotuning on Triton here
-            with inductor_config.patch(
-                max_autotune=False,
-                max_autotune_gemm=False,
-                max_autotune_gemm_backends="ATEN",
-            ):
+            # Base config: don't autotune Triton, but allow other optimizations
+            base_patches = {
+                "max_autotune": False,
+                "max_autotune_gemm": False,
+                "max_autotune_gemm_backends": "ATEN",
+            }
+            # Merge with user config_patches (e.g., coordinate_descent_tuning)
+            merged_patches: dict[str, Any] = {**base_patches, **(config_patches or {})}
+            with inductor_config.patch(merged_patches):
                 self.subgraph.run(*self.example_inputs)
 
+            # Tag all operations in subgraph with config_patches
+            # These will be applied during kernel codegen via SIMD scheduling
+            if config_patches:
+                for op in self.subgraph.operations:
+                    op.set_config_patches(config_patches.copy())
+
     def codegen(self, wrapper: PythonWrapperCodegen) -> None:
         class CodegenGraph:
             def __init__(self, graph: GraphLowering):
@@ -7138,6 +7483,7 @@ def __init__(self, graph: GraphLowering):
 
         assert is_node_sequence(self.inputs)
         outer_inputs = [t.codegen_reference() for t in self.inputs]
+
         wrapper.codegen_subgraph_with_flattened_outputs(
             CodegenGraph(self.subgraph),
             [*self.sym_inputs, *outer_inputs],
@@ -7146,6 +7492,10 @@ def __init__(self, graph: GraphLowering):
 
 
 class UserDefinedTritonKernel(ExternKernel):
+    """
+    A user-defined triton kernel (e.g. via @triton.jit).
+    """
+
     def get_kernel_and_metadata(self) -> tuple[Kernel, Any, list[str], list[str]]:
         from triton.runtime.autotuner import Autotuner
 
@@ -7178,8 +7528,79 @@ def get_kernel_and_metadata(self) -> tuple[Kernel, Any, list[str], list[str]]:
 
         return kernel, configs, restore_value_args, reset_to_zero_args
 
+    def can_fuse_epilogue(self) -> bool:
+        """
+        For kernels like
+
+        @triton.jit
+        def add_kernel(in_ptr0, in_ptr1, out_ptr, n_elements, BLOCK_SIZE: tl.constexpr):
+            pid = tl.program_id(0)
+            offs = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+            mask = offs < n_elements
+            x = tl.load(in_ptr0 + offs, mask=mask)
+            y = tl.load(in_ptr1 + offs, mask=mask)
+            tl.store(out_ptr + offs, x + y, mask=mask)
+
+        @torch.compile
+        def fn(a, b):
+            out = torch.empty_like(a)
+            grid = (triton.cdiv(a.numel(), 1024),)
+            add_kernel[grid](a, b, out, a.numel(), BLOCK_SIZE=1024)
+            return out.relu()
+
+        We can potentially fuse the relu epilogue into the add_kernel.
+        We do this by pruning the `out` tensor allocation and directly writing the relu-output.
+        """
+
+        if not config.epilogue_fusion_user_defined_triton_kernel:
+            return False
+
+        if not self.arg_accesses.can_fuse_epilogue:
+            return False
+
+        # We achieve fusion by parsing the original src into a python AST,
+        # then identify the expr containing the original value written via tl.store().
+        # We generate an expr for the value after the epilogue and replace that into the tl.store.
+        # So far we only support the simple case where there is a single tl.store in the kernel.
+        if len(self.kernel_stores.stores) != 1:
+            return False
+
+        # Only fuse if the mutated arg is originally an "empty" tensor.
+        # This is because we don't know exactly which element of that tensor is being written to.
+        # If the kernel only writes to a subset of the tensor, then we only apply the epilogue to that subset.
+        # In these edge cases, our fusion is only correct if the original tensor is empty,
+        # where the semantics is that content values are UB, and we can rely on the fact that `epilogue(UB) == UB`.
+        assert len(self.mutable_args) == 1
+        if not isinstance(self.mutable_args[0], TensorBox):
+            return False
+        if not isinstance(self.mutable_args[0].data, StorageBox):
+            return False
+        if not isinstance(self.mutable_args[0].data.data, ComputedBuffer):
+            return False
+        if not isinstance(self.mutable_args[0].data.data.data, Pointwise):
+            return False
+        if not all(r == 0 for r in self.mutable_args[0].data.data.data.ranges):
+            return False
+
+        return True
+
     @override
     def codegen(self, wrapper: PythonWrapperCodegen) -> None:
+        return self._codegen(wrapper, epilogue_fusion=None)
+
+    def codegen_with_epilogue_fusion(
+        self, wrapper: PythonWrapperCodegen, epilogue_fusion: tuple[ComputedBuffer, str]
+    ) -> None:
+        """
+        epilogue_fusion: (fused epilogue node, modified kerel src code)
+        """
+        return self._codegen(wrapper, epilogue_fusion)
+
+    def _codegen(
+        self,
+        wrapper: PythonWrapperCodegen,
+        epilogue_fusion: tuple[ComputedBuffer, str] | None,
+    ) -> None:
         """Overrides the parent member.
         See https://github.com/pytorch/pytorch/issues/151692"""
 
@@ -7205,10 +7626,19 @@ def codegen(self, wrapper: PythonWrapperCodegen) -> None:
             restore_value_args,
             reset_to_zero_args,
             self.grid,
+            epilogue_fusion,
         )
         named_args = {
             k: self.get_kwargs_value(k) for k in self.ordered_kwargs_for_cpp_kernel
         }
+
+        if epilogue_fusion:
+            assert len(self.arg_accesses.read_writes.writes) == 1
+            mutable_arg_name = next(iter(self.arg_accesses.read_writes.writes)).name
+            assert mutable_arg_name in named_args
+            epilogue_computed_buffer, _ = epilogue_fusion
+            named_args[mutable_arg_name] = epilogue_computed_buffer
+
         arg_names = [p.name for p in kernel.params]  # type: ignore[attr-defined]
         constexprs = [p.num for p in kernel.params if p.is_constexpr]  # type: ignore[attr-defined]
         constexpr_names = OrderedSet(arg_names[i] for i in constexprs)
@@ -7327,16 +7757,33 @@ def __init__(
             arg for arg in kernel.arg_names if arg in kernel_args
         ]
 
-        from torch._higher_order_ops.triton_kernel_wrap import identify_mutated_tensors
+        from torch._higher_order_ops.triton_kernel_wrap import (
+            identify_accessed_tensors,
+            identify_triton_stores,
+        )
 
         autotuned_kwargs = configs[0].kwargs if len(configs) > 0 else {}
+
+        import ast
+
+        # pyrefly: ignore [missing-attribute]
+        self.kernel_src = kernel.src
+        self.kernel_ast = ast.parse(self.kernel_src)
+        self.kernel_stores = identify_triton_stores(self.kernel_src)
+        self.kernel_args = kernel_args
+        # names in `arg_accesses.read_writes` are names of formal arguments in the kernel's prototype
+        self.arg_accesses = identify_accessed_tensors(
+            kernel,
+            {**kernel_args, **autotuned_kwargs},
+            tma_descriptor_metadata,
+        )
+
+        # Filter to only tensor args: with Triton 3.7+, ordered_arg_names
+        # includes scalars, so writes may reference non-tensor args like SymInts.
         self.mutable_args = [
-            kernel_args[key]
-            for key in identify_mutated_tensors(
-                kernel,
-                {**kernel_args, **autotuned_kwargs},
-                tma_descriptor_metadata,
-            )
+            kernel_args[key.name]
+            for key in self.arg_accesses.read_writes.writes
+            if isinstance(kernel_args.get(key.name), TensorBox)
         ]
 
         self.mutation_outputs = [
@@ -7345,10 +7792,48 @@ def __init__(
         ]
         V.graph.register_operation(self)
 
+    @override
+    def get_read_writes(self) -> dependencies.ReadWrites:
+        # Limit the new `get_read_writes` to `epilogue_fusion_user_defined_triton_kernel`
+        # to avoid potential regression to existing models.
+        if not config.epilogue_fusion_user_defined_triton_kernel:
+            return super().get_read_writes()
+
+        # maps formal arg name to actual arg name
+        read_renames = {
+            formal_arg_dep.name: self.kernel_args[formal_arg_dep.name].get_name()
+            for formal_arg_dep in self.arg_accesses.read_writes.reads
+        }
+
+        formal_arg_writes = list(self.arg_accesses.read_writes.writes)
+        write_renames = {
+            formal_arg_dep.name: mut_output.get_name()
+            for formal_arg_dep, mut_output in zip(
+                formal_arg_writes, self.mutation_outputs
+            )
+        }
+
+        read_writes = dependencies.ReadWrites(
+            reads=OrderedSet(
+                [
+                    dep.rename(read_renames)
+                    for dep in self.arg_accesses.read_writes.reads
+                ]
+            ),
+            writes=OrderedSet(
+                [
+                    dep.rename(write_renames)
+                    for dep in self.arg_accesses.read_writes.writes
+                ]
+            ),
+            index_exprs=OrderedSet(),
+        )
+        return read_writes
+
     def get_outputs(self) -> list[Buffer]:
         return list(self.mutation_outputs)
 
-    def get_device(self) -> Optional[torch.device]:
+    def get_device(self) -> torch.device | None:
         return self.device
 
 
@@ -7547,7 +8032,7 @@ def __init__(
         index: IRNode,
         src: IRNode,
         *,
-        reduce: Optional[str] = None,
+        reduce: str | None = None,
         include_self: bool = True,
     ) -> None:
         self.src_is_tensor = isinstance(src, TensorBox)
@@ -7698,9 +8183,9 @@ def __init__(
         self,
         unbacked_offset_symbol: sympy.Symbol,
         index: sympy.Symbol,
-        base_offset: Union[sympy.Symbol, int],
-        base_dim_stride: Union[sympy.Symbol, int],
-        size: Union[sympy.Symbol, int],
+        base_offset: sympy.Symbol | int,
+        base_dim_stride: sympy.Symbol | int,
+        size: sympy.Symbol | int,
         clamp: bool,
     ) -> None:
         super().__init__(None, NoneLayout(device=torch.device("cpu")), [])
@@ -7751,10 +8236,10 @@ def should_allocate(self) -> bool:
     def __init__(
         self,
         unbacked_size_symbol: sympy.Symbol,
-        start: Union[sympy.Symbol, int],
-        end: Union[sympy.Symbol, int],
-        step: Union[sympy.Symbol, int],
-        size: Union[sympy.Symbol, int],
+        start: sympy.Symbol | int,
+        end: sympy.Symbol | int,
+        step: sympy.Symbol | int,
+        size: sympy.Symbol | int,
     ):
         super().__init__(None, NoneLayout(device=torch.device("cpu")), [])
         # This node codegen
@@ -7890,9 +8375,9 @@ def __init__(
         tensor_args: Sequence[IRNode],
         nontensor_args: Sequence[Any],
         unflatten_args: Callable[..., Any],
-        kwargs: Optional[dict[str, Any]] = None,
+        kwargs: dict[str, Any] | None = None,
         *,
-        unbacked_bindings: Optional[dict[sympy.Symbol, pytree.KeyPath]] = None,
+        unbacked_bindings: dict[sympy.Symbol, pytree.KeyPath] | None = None,
     ) -> None:
         super().__init__(
             layout,
@@ -8055,7 +8540,7 @@ def __repr__(self) -> str:
 
     @staticmethod
     def find_device(
-        tensor_args: Optional[Sequence[torch.Tensor]], example_output: Sequence[Any]
+        tensor_args: Sequence[torch.Tensor] | None, example_output: Sequence[Any]
     ) -> Any:
         non_torch_bind_tensor_args = (
             [t for t in tensor_args if not isinstance(t, TorchBindObject)]
@@ -8068,6 +8553,10 @@ def find_device(
             return devices[0]
         if isinstance(example_output, torch.Tensor):
             return example_output.device
+        if isinstance(
+            example_output, (torch._C.ScriptObject, FakeScriptObject)
+        ) or is_opaque_value(example_output):
+            return torch.device("cpu")
         if isinstance(example_output, (list, tuple)):
             device_set = OrderedSet(
                 # pyrefly: ignore [bad-argument-type]
@@ -8078,6 +8567,8 @@ def find_device(
             devices = [device for device in device_set if device]
             if len(devices) == 1:
                 return devices[0]
+            if not devices:
+                return None
             for device in devices:
                 assert isinstance(device, torch.device)
                 if is_gpu(device.type):
@@ -8150,8 +8641,8 @@ def export_extern_kernel_node(self):  # type: ignore[no-untyped-def]
 
         # serialize_outputs
         def handle_single_output(
-            return_type: Union[torch.TensorType, torch.ListType, torch.JitType],
-            output: Union[IRNode, Sequence[IRNode]],
+            return_type: torch.TensorType | torch.ListType | torch.JitType,
+            output: IRNode | Sequence[IRNode],
         ) -> export_schema.Argument:
             if isinstance(return_type, (torch.TensorType, torch.NoneType)):
                 # For single Tensor or None
@@ -8358,6 +8849,39 @@ def create(cls, kernel: _OpOverloads, *args: Any, **kwargs: Any) -> FallbackKern
                 unbacked_bindings,
             ) = cls.process_kernel(kernel, *args, **kwargs)
 
+        # Try to lower single output functional custom ops to their out-variant.
+        if isinstance(kernel, torch._ops.OpOverload) and isinstance(
+            example_output, torch.Tensor
+        ):
+            from torch._library._out_variant import (
+                _is_functional,
+                get_out_arg_names,
+                lookup_manual_out_variant,
+                to_out_variant,
+            )
+
+            out_op = None
+            if _is_functional(kernel._schema):
+                out_op = to_out_variant(kernel)
+            if out_op is None:
+                out_op = lookup_manual_out_variant(kernel)
+
+            if out_op is not None and len(get_out_arg_names(out_op)) == 1:
+                layout = FixedLayout(
+                    device=example_output.device,
+                    dtype=example_output.dtype,
+                    size=[*example_output.shape],
+                    stride=[*example_output.stride()],
+                )
+                return ExternKernelOut(  # type: ignore[return-value]
+                    layout=layout,
+                    inputs=list(tensor_args),
+                    constant_args=list(non_tensor_args),
+                    kwargs=kwargs,
+                    python_kernel_name=_make_out_variant_kernel_name(out_op),
+                    op_overload=out_op,
+                )
+
         # We need this extra check for input alignment since the example
         # inputs we created are always aligned.
         has_unaligned_input = any(is_unaligned(arg) for arg in tensor_args)
@@ -8371,6 +8895,26 @@ def create(cls, kernel: _OpOverloads, *args: Any, **kwargs: Any) -> FallbackKern
         ):
             device = torch.device("cpu")
 
+        # Try multi-output .out() lowering for ops with out_variant tag.
+        if (
+            isinstance(kernel, torch._ops.OpOverload)
+            and not V.graph.cpp_wrapper
+            and device
+        ):
+            out_result = ExternKernelMultiOut.try_create(
+                kernel,
+                example_output,
+                device,
+                tensor_args,
+                non_tensor_args,
+                unflatten_args,
+                kwargs,
+                unbacked_bindings=unbacked_bindings,
+                has_unaligned_input=has_unaligned_input,
+            )
+            if out_result is not None:
+                return out_result  # type: ignore[return-value]
+
         if example_output is None:
             packed = cls(
                 NoneLayout(device=device),
@@ -8422,6 +8966,15 @@ def generate_output(output: Any, indices: list[tuple[Any, int]]) -> Any:
                 return output
             elif isinstance(output, torch.SymInt):
                 return output.node.expr
+            elif isinstance(
+                output, (torch._C.ScriptObject, FakeScriptObject)
+            ) or is_opaque_value(output):
+                return OpaqueMultiOutput(
+                    NoneLayout(device=device),
+                    packed,
+                    indices,
+                    output,
+                )
             else:
                 assert output is None, (
                     f"FallbackKernel output type {type(output)} is not supported"
@@ -8459,7 +9012,7 @@ def __init__(
         unflatten_args: Callable[..., Any],
         *,
         kwargs: dict[str, Any] | None = None,
-        unbacked_bindings: Optional[dict[sympy.Symbol, pytree.KeyPath]] = None,
+        unbacked_bindings: dict[sympy.Symbol, pytree.KeyPath] | None = None,
     ) -> None:
         super().__init__(
             layout,
@@ -8501,7 +9054,7 @@ def codegen(self, wrapper: PythonWrapperCodegen) -> None:
 class MultiOutputLayout(OutputSpec):
     device: torch.device
 
-    def get_device(self) -> Optional[torch.device]:
+    def get_device(self) -> torch.device | None:
         return self.device
 
 
@@ -8546,6 +9099,204 @@ def get_inputs_that_alias_output(self) -> Sequence[str]:
             and len(inp.get_inputs_that_alias_output()) > 0
         ]
 
+    def get_read_writes(self) -> dependencies.ReadWrites:
+        # Reads: StarDep on parent (we don't know which elements of the
+        # packed output we index into — conservative is correct).
+        reads: OrderedSet[dependencies.Dep] = OrderedSet()
+        for inp in self.inputs:
+            if isinstance(inp, IRNode):
+                reads.add(dependencies.StarDep(inp.get_name()))
+
+        # Writes: build proper MemoryDep from our FixedLayout so the
+        # scheduler can match our write with downstream epilogue reads.
+        # Normalize using the same policy as SchedulerNode so that the
+        # index expressions are directly comparable during fusion checks.
+        name = self.get_name()
+        indexer = self.get_layout().make_indexer()
+
+        def dummy(index: Sequence[Any], rindex: Sequence[Any]) -> Any:
+            assert len(rindex) == 0
+            return ops.store(name, indexer(index), "fake")
+
+        device = self.get_device()
+        should_normalize = (
+            not config.loop_ordering_after_fusion
+            or device is None
+            or not is_gpu(device.type)
+        )
+        write_rw = dependencies.extract_read_writes(
+            dummy, self.get_size(), (), normalize=should_normalize
+        )
+        return dependencies.ReadWrites(
+            reads=reads,
+            writes=write_rw.writes,
+            index_exprs=OrderedSet(),
+        )
+
+
+class OpaqueMultiOutput(MultiOutput):
+    """MultiOutput for opaque objects."""
+
+    def __init__(
+        self,
+        layout: OutputSpec,
+        input: IRNode,
+        indices: list[tuple[Any, ...]],
+        opaque_value: Any,
+    ) -> None:
+        super().__init__(layout, input, indices, skip_size_stride_alignment_checks=True)
+        self.opaque_example_value = opaque_value
+
+    @property  # type: ignore[override]
+    def dtype(self) -> Never:
+        raise AttributeError("OpaqueMultiOutput has no dtype")
+
+    def wrap_for_lowering(self) -> OpaqueMultiOutput:
+        return self
+
+    def get_read_writes(self) -> dependencies.ReadWrites:
+        reads: OrderedSet[dependencies.Dep] = OrderedSet()
+        for inp in self.inputs:
+            if isinstance(inp, IRNode):
+                reads.add(dependencies.StarDep(inp.get_name()))
+        writes: OrderedSet[dependencies.Dep] = OrderedSet(
+            [dependencies.StarDep(self.get_name())]
+        )
+        return dependencies.ReadWrites(
+            reads=reads,
+            writes=writes,
+            index_exprs=OrderedSet(),
+        )
+
+
+class AllocatingMultiOutput(MultiOutput):
+    """MultiOutput with Inductor-controlled allocation for .out() variant ops.
+
+    Overrides should_allocate()=True so Inductor allocates the output buffer,
+    and skips tuple-indexing codegen since .out() writes directly into these buffers.
+    """
+
+    def should_allocate(self) -> bool:
+        return True
+
+    def codegen(self, wrapper: PythonWrapperCodegen) -> None:
+        if not self.skip_size_stride_alignment_checks:
+            self.codegen_size_asserts(wrapper)
+            self.codegen_alignment_asserts(wrapper)
+
+
+def _make_out_variant_kernel_name(out_op: torch._ops.OpOverload) -> str:
+    """Build fully-qualified kernel name for an out-variant op."""
+    ns = out_op.namespace
+    op_name = out_op._schema.name.split("::")[1]
+    overload = out_op._overloadname
+    return f"torch.ops.{ns}.{op_name}.{overload}"
+
+
+class ExternKernelMultiOut(FallbackKernel):
+    """Multi-output .out() variant lowering.
+
+    Subclass of FallbackKernel that emits .out() calls with pre-allocated
+    output buffers. Uses AllocatingMultiOutput child nodes for each output.
+    """
+
+    out_arg_names: list[str]
+    out_variant_output_nodes: list[AllocatingMultiOutput]
+
+    def __init__(
+        self,
+        *args: Any,
+        out_op: torch._ops.OpOverload,
+        out_arg_names: list[str],
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(*args, **kwargs)
+        self.out_arg_names = out_arg_names
+        self.out_variant_output_nodes = []
+        self.python_kernel_name = _make_out_variant_kernel_name(out_op)
+        self.op_overload = out_op
+
+    def codegen(self, wrapper: PythonWrapperCodegen) -> None:
+        self.codegen_comment(wrapper)
+        wrapper.generate_extern_kernel_multi_out(self)
+
+    @classmethod
+    def try_create(
+        cls,
+        kernel: torch._ops.OpOverload,
+        example_output: Any,
+        device: torch.device,
+        tensor_args: Sequence[IRNode],
+        non_tensor_args: Sequence[Any],
+        unflatten_args: Callable[..., Any],
+        kwargs: dict[str, Any] | None,
+        *,
+        unbacked_bindings: dict[sympy.Symbol, pytree.KeyPath] | None = None,
+        has_unaligned_input: bool = False,
+    ) -> Sequence[AllocatingMultiOutput] | None:
+        """Create an ExternKernelMultiOut if the op has a matching .out() variant."""
+        from torch._library._out_variant import (
+            _is_functional,
+            get_out_arg_names,
+            to_out_variant,
+        )
+
+        if not _is_functional(kernel._schema):
+            return None
+
+        if not isinstance(example_output, (tuple, list)):
+            return None
+
+        out_op = to_out_variant(kernel)
+        if out_op is None:
+            return None
+
+        out_arg_names = get_out_arg_names(out_op)
+        if not all(isinstance(t, torch.Tensor) for t in example_output):
+            return None
+        if len(example_output) != len(out_arg_names):
+            return None
+
+        packed = cls(
+            MultiOutputLayout(device=device),
+            kernel,
+            tensor_args,
+            non_tensor_args,
+            unflatten_args,
+            kwargs=kwargs,
+            unbacked_bindings=unbacked_bindings,
+            out_op=out_op,
+            out_arg_names=out_arg_names,
+        )
+
+        outputs: list[AllocatingMultiOutput] = []
+        for i, tensor_out in enumerate(example_output):
+            layout = FixedLayout(
+                device=tensor_out.device,
+                dtype=tensor_out.dtype,
+                size=[*tensor_out.shape],
+                stride=[*tensor_out.stride()],
+            )
+            multi_out = AllocatingMultiOutput(
+                layout=layout,
+                input=packed,
+                indices=[(type(example_output), i)],
+            )
+            if (
+                config.assume_unaligned_fallback_output
+                or has_unaligned_input
+                or not tensor_is_aligned(tensor_out)
+            ):
+                V.graph.unaligned_buffers.add(multi_out.name)  # type: ignore[arg-type]
+            outputs.append(multi_out)
+
+        packed.out_variant_output_nodes = outputs
+        packed.outputs = outputs
+
+        if isinstance(example_output, tuple):
+            return tuple(outputs)  # type: ignore[return-value]
+        return list(outputs)
+
 
 # We just use a normal dataclass for MutableBox/TensorBox/StorageBox since
 # they're mainly lowering-time constructs that we expect to mutate and such.
@@ -8560,7 +9311,7 @@ class MutableBox(IRNode):
     def has_exceeded_max_reads(self) -> bool:
         return self.data.has_exceeded_max_reads()
 
-    def get_device(self) -> Optional[torch.device]:
+    def get_device(self) -> torch.device | None:
         return self.data.get_device()
 
     def make_loader(self) -> Callable[[Sequence[Expr]], OpsValue]:
@@ -8575,7 +9326,7 @@ def get_stride(self) -> Sequence[_IntLike]:
     def get_name(self) -> str:
         return self.data.get_name()
 
-    def has_large_inner_fn(self, threshold: Optional[int] = None) -> bool:
+    def has_large_inner_fn(self, threshold: int | None = None) -> bool:
         return self.data.has_large_inner_fn(threshold)
 
     def mark_reuse(self, users: int) -> None:
@@ -8621,7 +9372,7 @@ def num_reads(self) -> int:
     def get_storage_numel(self) -> _IntLike:
         return self.data.get_storage_numel()
 
-    def get_reduction_type(self) -> Optional[str]:
+    def get_reduction_type(self) -> str | None:
         return self.data.get_reduction_type()
 
     def get_reduction_size(self) -> Sequence[Expr]:
@@ -8645,7 +9396,7 @@ def get_operation_name(self) -> str:
     def get_inputs_that_alias_output(self) -> Sequence[str]:
         return self.data.get_inputs_that_alias_output()
 
-    def realize(self) -> Optional[str]:
+    def realize(self) -> str | None:
         return self.data.realize()
 
     @cache_on_self_and_args("MutableBox")
@@ -8657,10 +9408,10 @@ def get_free_symbol_uses(
     def get_read_names(self) -> OrderedSet[str]:
         return self.data.get_read_names()
 
-    def get_defining_op(self) -> Optional[Operation]:
+    def get_defining_op(self) -> Operation | None:
         return self.data.get_defining_op()
 
-    def codegen_reference(self, writer: Optional[IndentedBuffer] = None) -> str:
+    def codegen_reference(self, writer: IndentedBuffer | None = None) -> str:
         return self.data.codegen_reference(writer)
 
     @property
@@ -8732,7 +9483,7 @@ def is_module_buffer(self) -> bool:
             and self.data.get_name() in V.graph.constants
         )
 
-    def realize(self) -> Optional[str]:
+    def realize(self) -> str | None:
         if IRNode.is_realized_node(self.data):
             return self.data.get_name()
 
@@ -8759,6 +9510,7 @@ def realize(self) -> Optional[str]:
         self.data.origins = self.origins
         self.data.origin_node = origin_node
         self.data.traceback = traceback
+        self.data.stream_idx = self.data.data.stream_idx
         return self.data.name
 
     def realize_hint(self) -> None:
@@ -8832,7 +9584,7 @@ def num_reads(self) -> int:
 class Subgraph(IRNode):
     name: str
     graph_module: torch.fx.GraphModule
-    graph: Optional[GraphLowering] = None
+    graph: GraphLowering | None = None
 
 
 def _has_aliased_buffers(buffers: Sequence[IRNode]) -> bool:
@@ -8850,9 +9602,9 @@ class InvokeSubgraph(ExternKernel):
     Ir node for the invoke_subgraph HOP.
     """
 
-    subgraph: Optional[Subgraph] = None
-    operands: Optional[Sequence[IRNode]] = None
-    outputs: Optional[Sequence[IRNode]] = None
+    subgraph: Subgraph | None = None
+    operands: Sequence[IRNode] | None = None
+    outputs: Sequence[IRNode] | None = None
 
     def __init__(
         self, subgraph: Subgraph, operands: Sequence[IRNode], layout: MultiOutputLayout
@@ -8866,10 +9618,13 @@ def __init__(
         self.name = V.graph.register_buffer(self)
         V.graph.register_operation(self)
 
+    def get_subgraphs(self) -> list[Subgraph]:
+        return [self.subgraph] if self.subgraph else []
+
     @classmethod
     def create(
         cls, subgraph: Subgraph, *operands: IRNode
-    ) -> list[Union[ShapeAsConstantBuffer, NoneAsConstantBuffer, MultiOutput]]:
+    ) -> list[ShapeAsConstantBuffer | NoneAsConstantBuffer | MultiOutput]:
         """For each operand, get a realized input, force it to have the same
         strides as the subgraph inputs, then use an InvokeSubgraph"""
         from .lowering import constrain_to_fake_tensor
@@ -8906,7 +9661,9 @@ def create(
         new_operands: list[IRNode] = []
 
         for idx, operand in enumerate(operands):
-            if isinstance(operand, (ShapeAsConstantBuffer, GeneratorState)):
+            if isinstance(
+                operand, (ShapeAsConstantBuffer, GeneratorState, OpaqueObjectState)
+            ):
                 new_operands.append(operand)
             else:
                 new_operands.append(
@@ -8944,7 +9701,7 @@ def create(
 
         def create_output(
             output: IRNode, ind: int
-        ) -> Union[ShapeAsConstantBuffer, NoneAsConstantBuffer, MultiOutput]:
+        ) -> ShapeAsConstantBuffer | NoneAsConstantBuffer | MultiOutput:
             if isinstance(output, (ShapeAsConstantBuffer, NoneAsConstantBuffer)):
                 return output
             else:
@@ -8986,11 +9743,11 @@ class Conditional(ExternKernel):
         outputs: MultiOutput nodes representing the conditional's outputs.
     """
 
-    predicate: Optional[IRNode] = None
-    operands: Optional[Sequence[IRNode]] = None
-    true_subgraph: Optional[Subgraph] = None
-    false_subgraph: Optional[Subgraph] = None
-    outputs: Optional[Sequence[MultiOutput]] = None
+    predicate: IRNode | None = None
+    operands: Sequence[IRNode] | None = None
+    true_subgraph: Subgraph | None = None
+    false_subgraph: Subgraph | None = None
+    outputs: Sequence[MultiOutput] | None = None
 
     def __init__(
         self,
@@ -8999,7 +9756,7 @@ def __init__(
         true_subgraph: Subgraph,
         false_subgraph: Subgraph,
         layout: MultiOutputLayout,
-        unbacked_bindings: Optional[dict[sympy.Symbol, pytree.KeyPath]],
+        unbacked_bindings: dict[sympy.Symbol, pytree.KeyPath] | None,
     ) -> None:
         self.predicate = predicate
         self.operands = operands
@@ -9020,8 +9777,16 @@ def __init__(
         self.name = V.graph.register_buffer(self)
         V.graph.register_operation(self)
 
+    def get_subgraphs(self) -> list[Subgraph]:
+        subgraphs = []
+        if self.true_subgraph:
+            subgraphs.append(self.true_subgraph)
+        if self.false_subgraph:
+            subgraphs.append(self.false_subgraph)
+        return subgraphs
+
     @staticmethod
-    def _maybe_expr(s: Union[int, torch.SymInt]) -> Union[int, sympy.Expr]:
+    def _maybe_expr(s: int | torch.SymInt) -> int | sympy.Expr:
         if isinstance(s, int):
             return s
         return s.node.expr
@@ -9042,8 +9807,16 @@ def create(
         fx_operands: Argument = V.graph.current_node.args[-1]
 
         assert isinstance(fx_operands, Sequence), type(fx_operands)
-        assert all(isinstance(n, Node) for n in fx_operands)
-        fake_operands = [cast(Node, x).meta["val"] for x in fx_operands]
+        # Build fake_operands from FX nodes' metadata
+        # For FX Nodes, get the fake tensor from meta["val"]
+        # For non-Nodes (e.g., symbolic integers from sym_size lowering), pass directly
+        fake_operands: list[Any] = []
+        for fx_op in fx_operands:
+            if isinstance(fx_op, Node):
+                fake_operands.append(fx_op.meta["val"])
+            else:
+                # Symbolic integer or constant - pass directly
+                fake_operands.append(fx_op)
         fake_outputs = V.graph.current_node.meta["val"]
 
         def _require_exact_strides(
@@ -9149,6 +9922,28 @@ def _require_exact_strides(
         ]
 
         conditional.outputs = outputs  # type: ignore[assignment]
+
+        from torch._higher_order_ops.utils import (
+            check_input_alias_and_mutation_return_outputs,
+        )
+
+        (_, _, _, true_mutated_inputs, _) = (
+            check_input_alias_and_mutation_return_outputs(true_fn.graph_module)
+        )
+        (_, _, _, false_mutated_inputs, _) = (
+            check_input_alias_and_mutation_return_outputs(false_fn.graph_module)
+        )
+
+        mutated_operand_indices = OrderedSet(true_mutated_inputs) | OrderedSet(
+            false_mutated_inputs
+        )
+
+        # Create MutationOutput for each mutated operand (for scheduler dependencies)
+        conditional.mutation_outputs = [
+            MutationOutput(operands[idx].layout, operands[idx], conditional)  # type: ignore[union-attr]
+            for idx in sorted(mutated_operand_indices)
+        ]
+
         return outputs
 
     def codegen(self, wrapper: PythonWrapperCodegen) -> None:
@@ -9186,11 +9981,11 @@ def _split_by_sym_type(
 class WhileLoop(ExternKernel):
     """The IR node for while_loop and while_loop_stack_output. It supports input mutation."""
 
-    carried_inputs: Optional[Sequence[IRNode]] = None
-    additional_inputs: Optional[Sequence[IRNode]] = None
-    cond_subgraph: Optional[Subgraph] = None
-    body_subgraph: Optional[Subgraph] = None
-    outputs: Optional[Sequence[MultiOutput]] = None
+    carried_inputs: Sequence[IRNode] | None = None
+    additional_inputs: Sequence[IRNode] | None = None
+    cond_subgraph: Subgraph | None = None
+    body_subgraph: Subgraph | None = None
+    outputs: Sequence[MultiOutput] | None = None
 
     def __init__(
         self,
@@ -9199,7 +9994,7 @@ def __init__(
         cond_subgraph: Subgraph,
         body_subgraph: Subgraph,
         layout: MultiOutputLayout,
-        unbacked_bindings: Optional[dict[sympy.Symbol, pytree.KeyPath]],
+        unbacked_bindings: dict[sympy.Symbol, pytree.KeyPath] | None,
         stack_output: bool,
     ) -> None:
         self.carried_inputs = carried_inputs
@@ -9223,6 +10018,14 @@ def __init__(
         self.name = V.graph.register_buffer(self)
         V.graph.register_operation(self)
 
+    def get_subgraphs(self) -> list[Subgraph]:
+        subgraphs = []
+        if self.cond_subgraph:
+            subgraphs.append(self.cond_subgraph)
+        if self.body_subgraph:
+            subgraphs.append(self.body_subgraph)
+        return subgraphs
+
     # Accidental aliasing can be created due to cse, where the empty buffers we
     # allocated for backward to use gets csed into the same buffer in function fx_graph_cse.
     # See test_scan_multiple_layers_gradient for a concrete example.
@@ -9241,7 +10044,7 @@ def _clone_aliased_inputs(carried_inputs: Sequence[IRNode]) -> Sequence[IRNode]:
 
         # Track which buffers we've seen and their indices
         seen_buffers: OrderedSet[int] = OrderedSet()
-        result: list[Union[IRNode, TensorBox]] = []
+        result: list[IRNode | TensorBox] = []
 
         for original_input, unwrapped_buffer in zip(carried_inputs, unwrapped_buffers):
             if id(unwrapped_buffer) in seen_buffers:
@@ -9271,7 +10074,7 @@ def create(
         carried_inputs: Sequence[IRNode],
         additional_inputs: Sequence[IRNode],
         stack_output: bool,
-    ) -> Union[IRNode, Sequence[IRNode]]:
+    ) -> IRNode | Sequence[IRNode]:
         """create the while_loop IR node. stack_output controls whether it stack
         each iterations' output, which is necessary for training.
         """
@@ -9279,7 +10082,7 @@ def create(
 
         def _require_exact_strides(
             tensor_boxes: Sequence[IRNode],
-            fake_tensors: list[Union[int, torch.SymInt, torch.Tensor]],
+            fake_tensors: list[int | torch.SymInt | torch.Tensor],
         ) -> list[IRNode]:
             assert len(tensor_boxes) == len(fake_tensors)
             ret = []
@@ -9380,8 +10183,8 @@ def _require_exact_strides(
         for i, (op, bo) in enumerate(zip(carried_inputs_, body_outputs)):
 
             def _guard_list_equals(
-                lhs_exprs: Sequence[Union[int, sympy.Expr]],
-                rhs_exprs: Sequence[Union[int, sympy.Expr]],
+                lhs_exprs: Sequence[int | sympy.Expr],
+                rhs_exprs: Sequence[int | sympy.Expr],
             ) -> None:
                 assert len(lhs_exprs) == len(rhs_exprs)
                 for lhs, rhs in zip(lhs_exprs, rhs_exprs):
@@ -9506,9 +10309,9 @@ def __init__(
         tensor_args: Sequence[IRNode],
         nontensor_args: Sequence[Any],
         unflatten_args: Callable[..., Any],
-        kwargs: Optional[dict[str, Any]] = None,
+        kwargs: dict[str, Any] | None = None,
         *,
-        unbacked_bindings: Optional[dict[sympy.Symbol, pytree.KeyPath]] = None,
+        unbacked_bindings: dict[sympy.Symbol, pytree.KeyPath] | None = None,
     ) -> None:
         super().__init__(
             layout,
@@ -9553,15 +10356,15 @@ def get_free_symbol_uses(
 @ir_dataclass
 class TorchBindObject(NonTensorObj):
     name: str
-    value: Union[FakeScriptObject, torch.ScriptObject]
+    value: FakeScriptObject | torch.ScriptObject
 
     def get_name(self) -> str:
         return self.name
 
-    def codegen_reference(self, writer: Optional[IndentedBuffer] = None) -> str:
+    def codegen_reference(self, writer: IndentedBuffer | None = None) -> str:
         return self.name
 
-    def get_value(self) -> Union[FakeScriptObject, torch.ScriptObject]:
+    def get_value(self) -> FakeScriptObject | torch.ScriptObject:
         return self.value
 
     def get_real_obj(self) -> torch.ScriptObject:
@@ -9574,7 +10377,7 @@ def get_buf_bytes(self) -> int:
         # Returns the sum of all tensors in the flattened object
         real_script_obj = self.get_real_obj()
 
-        if is_opaque_type(real_script_obj):
+        if is_opaque_value(real_script_obj):
             return 0
 
         assert hasattr(real_script_obj, "__obj_flatten__")
@@ -9588,6 +10391,27 @@ def get_buf_bytes(self) -> int:
         return functools.reduce(operator.add, flat_sizes, 0)
 
 
+@ir_dataclass
+class OpaqueValueTypeConstant(NonTensorObj):
+    """IR node for opaque value type constants that appear directly in graph outputs.
+
+    Unlike TorchBindObject (which references named constants loaded at runtime),
+    this inlines the value's repr into the generated code since value types are
+    reconstructed from their repr.
+    """
+
+    value: Any
+
+    def get_name(self) -> str:
+        return repr(self.value)
+
+    def codegen_reference(self, writer: IndentedBuffer | None = None) -> str:
+        obj_repr, opaque_types = get_opaque_obj_repr(self.value)
+        for n, t in opaque_types.items():
+            V.graph.opaque_value_type_classes[n] = t
+        return obj_repr
+
+
 @ir_dataclass
 class GeneratorState(NonTensorObj):
     name: str
@@ -9596,7 +10420,25 @@ class GeneratorState(NonTensorObj):
     def get_name(self) -> str:
         return self.name
 
-    def codegen_reference(self, writer: Optional[IndentedBuffer] = None) -> str:
+    def codegen_reference(self, writer: IndentedBuffer | None = None) -> str:
+        return self.name
+
+
+@ir_dataclass
+class OpaqueObjectState(NonTensorObj):
+    """
+    Represents an opaque object (e.g., ProcessGroup) that is passed through
+    as a graph input. Similar to GeneratorState, this wraps the object with
+    its placeholder name so codegen can reference it properly.
+    """
+
+    name: str
+    value: Any  # The actual opaque object (for reference, not used in codegen)
+
+    def get_name(self) -> str:
+        return self.name
+
+    def codegen_reference(self, writer: IndentedBuffer | None = None) -> str:
         return self.name
 
 
@@ -9609,7 +10451,7 @@ def has_side_effects(self) -> bool:
 
     # This is identical to FallbackKernel.set_cpp_kernel(), minus the
     # part that checks against input aliasing and mutation.
-    def set_cpp_kernel_name(self, cpp_kernel_name: Optional[str] = None) -> None:
+    def set_cpp_kernel_name(self, cpp_kernel_name: str | None = None) -> None:
         assert type(self.op_overload) is torch._ops.OpOverload, (
             "Setting cpp kernel needs a valid op_overload"
         )
@@ -9633,7 +10475,7 @@ def set_cpp_kernel_name(self, cpp_kernel_name: Optional[str] = None) -> None:
     def create_inplace(
         cls,
         kernel: _OpOverloads,
-        inputs: Union[IRNode, list[IRNode]],
+        inputs: IRNode | list[IRNode],
         *args: Any,
         **kwargs: Any,
     ) -> None:
@@ -9699,10 +10541,10 @@ def create_inplace(
     def create_out_of_place(
         cls,
         kernel: _OpOverloads,
-        inputs: Union[TensorBox, list[TensorBox]],
+        inputs: TensorBox | list[TensorBox],
         *args: Any,
         **kwargs: Any,
-    ) -> Union[list[MultiOutput], _CollectiveKernel]:
+    ) -> list[MultiOutput] | _CollectiveKernel:
         with V.graph.fake_mode:
             (
                 example_output,
@@ -9764,9 +10606,9 @@ def __init__(
         tensor_args: Sequence[IRNode],
         nontensor_args: Sequence[Any],
         unflatten_args: Callable[..., Any],
-        kwargs: Optional[dict[str, Any]] = None,
+        kwargs: dict[str, Any] | None = None,
         *,
-        unbacked_bindings: Optional[dict[sympy.Symbol, pytree.KeyPath]] = None,
+        unbacked_bindings: dict[sympy.Symbol, pytree.KeyPath] | None = None,
     ) -> None:
         super().__init__(
             layout,
@@ -9795,9 +10637,9 @@ def __init__(
         tensor_args: Sequence[IRNode],
         nontensor_args: Sequence[Any],
         unflatten_args: Callable[..., Any],
-        kwargs: Optional[dict[str, Any]] = None,
+        kwargs: dict[str, Any] | None = None,
         *,
-        unbacked_bindings: Optional[dict[sympy.Symbol, pytree.KeyPath]] = None,
+        unbacked_bindings: dict[sympy.Symbol, pytree.KeyPath] | None = None,
     ) -> None:
         super().__init__(
             layout,
@@ -9826,9 +10668,9 @@ def __init__(
         tensor_args: Sequence[IRNode],
         nontensor_args: Sequence[Any],
         unflatten_args: Callable[..., Any],
-        kwargs: Optional[dict[str, Any]] = None,
+        kwargs: dict[str, Any] | None = None,
         *,
-        unbacked_bindings: Optional[dict[sympy.Symbol, pytree.KeyPath]] = None,
+        unbacked_bindings: dict[sympy.Symbol, pytree.KeyPath] | None = None,
     ) -> None:
         super().__init__(
             layout,
diff --git a/torch/_inductor/jagged_lowerings.py b/torch/_inductor/jagged_lowerings.py
index 83023cc562dde..3a8e404133c4e 100644
--- a/torch/_inductor/jagged_lowerings.py
+++ b/torch/_inductor/jagged_lowerings.py
@@ -1,5 +1,4 @@
 # mypy: allow-untyped-defs
-from typing import Optional, Union
 
 import sympy
 
@@ -24,7 +23,7 @@ def dense_idx_to_jagged_idx(batch_idx, seq_idx, offsets_loader, jagged_len):
 
 def get_inverse_offsets(
     offsets: TensorBox,
-    jagged_len: Union[int, sympy.Expr],
+    jagged_len: int | sympy.Expr,
     realize: bool = True,
 ) -> TensorBox:
     """
@@ -92,8 +91,8 @@ def jagged_idx_to_dense_idx(
     jagged_idx,  # pyre-ignore[2]
     inverse_offsets_loader,  # pyre-ignore[2]
     offsets_loader,  # pyre-ignore[2]
-    batch_size: Union[int, sympy.Expr],
-    max_seq_len: Union[int, sympy.Expr],
+    batch_size: int | sympy.Expr,
+    max_seq_len: int | sympy.Expr,
     offsets_dtype: torch.dtype,
 ) -> tuple[sympy.Expr, sympy.Expr]:
     batch_idx = ops.indirect_indexing(
@@ -187,7 +186,7 @@ def _dense_to_jagged_forward_impl(
         fallback_op,  # pyre-ignore[2]
         dense: TensorBox,
         jagged_offsets: list[TensorBox],
-        jagged_len: Optional[int] = None,
+        jagged_len: int | None = None,
     ) -> TensorBox:
         device = dense.get_device_or_error()
         dtype = dense.get_dtype()
@@ -260,7 +259,7 @@ def inner_fn(index):
     def _dense_to_jagged_forward(
         dense: TensorBox,
         jagged_offsets: list[TensorBox],
-        jagged_len: Optional[int] = None,
+        jagged_len: int | None = None,
     ) -> TensorBox:
         return _dense_to_jagged_forward_impl(
             fallback_op=torch.ops.aten._padded_dense_to_jagged_forward.default,
diff --git a/torch/_inductor/kernel/bmm.py b/torch/_inductor/kernel/bmm.py
index 7b6339fe8283e..55e6bcf82810b 100644
--- a/torch/_inductor/kernel/bmm.py
+++ b/torch/_inductor/kernel/bmm.py
@@ -1,6 +1,6 @@
 # mypy: allow-untyped-defs
 import logging
-from typing import TYPE_CHECKING, Union
+from typing import TYPE_CHECKING
 
 import torch
 from torch._dynamo.utils import counters
@@ -10,6 +10,7 @@
 from .. import config as inductor_config, ir, lowering as L
 from ..kernel_inputs import MMKernelInputs
 from ..lowering import lowerings, make_pointwise, make_reduction, transform_args
+from ..runtime.runtime_utils import get_max_y_grid
 from ..select_algorithm import (
     autotune_select_algorithm,
     ExternKernelChoice,
@@ -43,8 +44,14 @@
 
 
 @SymbolicGridFn
-def bmm_grid(b, m, n, meta, *, cdiv):
-    return (cdiv(m, meta["BLOCK_M"]) * cdiv(n, meta["BLOCK_N"]), b, 1)
+def bmm_grid(b, m, n, meta, *, cdiv, max):
+    tiles = cdiv(m, meta["BLOCK_M"]) * cdiv(n, meta["BLOCK_N"])
+    # Split batch across grid_y and grid_z to avoid exceeding CUDA grid_y limit.
+    # When b <= max_y_grid, grid_z = 1 and behavior is identical to the original.
+    max_y_grid = get_max_y_grid()
+    grid_z = max(cdiv(b, max_y_grid), 1)
+    grid_y = cdiv(b, grid_z)
+    return (tiles, grid_y, grid_z)
 
 
 # We define each template kernel in a separate file which is the name of the input to load_kernel_template
@@ -60,7 +67,7 @@ def bmm_grid(b, m, n, meta, *, cdiv):
 aten_bmm = ExternKernelChoice(torch.bmm, "at::bmm_out", op_overload=aten.bmm.out)
 aten_bmm_dtype = ExternKernelChoice(
     torch.bmm,
-    "at::_bmm_out_dtype_xpu" if torch.xpu.is_available() else "at::_bmm_out_dtype_cuda",
+    "at::_bmm_out_dtype_xpu" if torch.xpu._is_compiled() else "at::_bmm_out_dtype_cuda",
     name="bmm_dtype",
     op_overload=aten.bmm.dtype_out,
 )
@@ -172,17 +179,14 @@ def _to_dtype(x):
     choices: list[ChoiceCaller] = []
 
     # Collect all templates for unified call
-    templates_to_use: list[Union[ExternKernelChoice, KernelTemplate]] = []
+    templates_to_use: list[ExternKernelChoice | KernelTemplate] = []
     kwarg_overrides = {}
 
     if use_aten_gemm_kernels():
         templates_to_use.append(aten_handler)
         kwarg_overrides[aten_handler.uid] = aten_extra_kwargs
 
-    if use_triton_template(layout, check_max_autotune=False) and (
-        out_dtype is None or out_dtype == mat1.get_dtype()
-    ):
-        # TODO: add out_dtype support for Triton Template
+    if use_triton_template(layout, check_max_autotune=False):
         templates_to_use.append(bmm_template)
 
     # Single unified call for all templates
@@ -225,7 +229,8 @@ def _to_dtype(x):
 
         add_nv_universal_gemm_choices(choices, layout, kernel_inputs)
 
-    return autotune_select_algorithm(name, choices, kernel_inputs.nodes(), layout)
+    node, _ = autotune_select_algorithm(name, choices, kernel_inputs.nodes(), layout)
+    return node
 
 
 @L.register_lowering(aten.baddbmm)
@@ -273,7 +278,7 @@ def tuned_baddbmm(inp, mat1, mat2, *, alpha=1, beta=1, layout=None):
     choices: list[ChoiceCaller] = []
 
     # Collect all templates for unified call
-    templates_to_use: list[Union[ExternKernelChoice, KernelTemplate]] = []
+    templates_to_use: list[ExternKernelChoice | KernelTemplate] = []
     if use_aten_gemm_kernels():
         templates_to_use.append(aten_baddbmm)
 
@@ -285,4 +290,5 @@ def tuned_baddbmm(inp, mat1, mat2, *, alpha=1, beta=1, layout=None):
         V.choices.get_template_configs(kernel_inputs, templates_to_use, name)
     )
 
-    return autotune_select_algorithm(name, choices, kernel_inputs.nodes(), layout)
+    node, _ = autotune_select_algorithm(name, choices, kernel_inputs.nodes(), layout)
+    return node
diff --git a/torch/_inductor/kernel/conv.py b/torch/_inductor/kernel/conv.py
index 452879d5e302b..1fa9588a7bf63 100644
--- a/torch/_inductor/kernel/conv.py
+++ b/torch/_inductor/kernel/conv.py
@@ -2,7 +2,7 @@
 from __future__ import annotations
 
 import logging
-from typing import Optional, TYPE_CHECKING, TypedDict
+from typing import TYPE_CHECKING, TypedDict
 
 import torch
 from torch._inductor.codegen.rocm.ck_conv_template import CKGroupedConvFwdTemplate
@@ -29,6 +29,7 @@
     use_triton_template,
 )
 from ..virtualized import V
+from .mm_common import load_kernel_template
 
 
 if TYPE_CHECKING:
@@ -60,6 +61,29 @@ def conv3d_grid(n, c, d, h, w, meta, *, cdiv):
     )
 
 
+# =============================================================================
+# Depthwise conv1d (groups == in_channels == out_channels)
+# Uses direct element-wise multiply-accumulate instead of implicit GEMM.
+# Channels-last (NLC) layout with 3D tiling: BLOCK_N x BLOCK_L x BLOCK_C.
+# =============================================================================
+
+
+@SymbolicGridFn
+def depthwise_conv1d_grid(n, c, l, meta, *, cdiv):
+    return (
+        cdiv(n, meta["BLOCK_N"]),
+        cdiv(l, meta["BLOCK_L"]),
+        cdiv(c, meta["BLOCK_C"]),
+    )
+
+
+depthwise_conv1d_template = TritonTemplate(
+    name="depthwise_conv1d",
+    grid=depthwise_conv1d_grid,
+    source=load_kernel_template("triton_depthwise_conv"),
+    cache_codegen_enabled_for_template=True,
+)
+
 LOOP_BODY_2D = """
         idx_x_h = i - PADDING_H + idx_y_h * STRIDE_H
         idx_x_w = j - PADDING_W + idx_y_w * STRIDE_W
@@ -352,7 +376,7 @@ class ConvLayoutParams(TypedDict):
 def conv_layout(
     x: TensorBox,
     weight: TensorBox,
-    bias: Optional[TensorBox],
+    bias: TensorBox | None,
     stride: Sequence[int],
     padding: tuple[int, ...],
     dilation: tuple[int, ...],
@@ -361,16 +385,24 @@ def conv_layout(
     groups: int,
 ) -> ir.Layout:
     """Determine output layout for a convolution"""
+    # We use guard_int_seq rather than size_hints because the output shape
+    # depends on these values — if they ever contained symbols, size_hints
+    # would silently substitute a hint that could be wrong, producing an
+    # incorrect layout. guard_int_seq will install a proper guard instead.
+    # Note: stride and padding are already guarded via guard_int_seq in
+    # convolution() above, but we guard all four here so conv_layout is
+    # self-contained and doesn't rely on callers.
+    guard = V.graph.sizevars.guard_int_seq
     with V.graph.fake_mode:
         output = torch.ops.aten.convolution(
-            ir.ir_node_to_tensor(x, guard_shape=True),
-            ir.ir_node_to_tensor(weight, guard_shape=True),
-            ir.ir_node_to_tensor(bias, guard_shape=True),
-            V.graph.sizevars.size_hints(stride),  # type: ignore[arg-type]
-            V.graph.sizevars.size_hints(padding),  # type: ignore[arg-type]
-            V.graph.sizevars.size_hints(dilation),  # type: ignore[arg-type]
+            ir.ir_node_to_tensor(x),
+            ir.ir_node_to_tensor(weight),
+            ir.ir_node_to_tensor(bias),
+            guard(stride),
+            guard(padding),
+            guard(dilation),
             transposed,
-            V.graph.sizevars.size_hints(output_padding),  # type: ignore[arg-type]
+            guard(output_padding),
             groups,
         )
         sizes = ir.convert_shape_to_inductor(output.size())
@@ -417,7 +449,7 @@ def convert_1x1_conv_to_mm(x, weight, bias):
 def convolution(
     x: TensorBox,
     weight: TensorBox,
-    bias: Optional[TensorBox],
+    bias: TensorBox | None,
     stride: Sequence[int],
     padding: Sequence[int],
     dilation: Sequence[int],
@@ -494,8 +526,10 @@ def channels_last_conv():
             return True
 
         layout = conv_layout(x, weight, None, **kwargs)
+        # TODO: This does not guard on the stride order decision,
+        # shall we use optimization_hint to handle unbacked?
         req_stride_order = ir.get_stride_order(
-            V.graph.sizevars.size_hints(layout.stride)
+            V.graph.sizevars.guarding_hints_or_throw(layout.stride)
         )
         return req_stride_order == ir.NHWC_STRIDE_ORDER
 
@@ -536,8 +570,10 @@ def channels_last_conv():
         layout = conv_layout(x, weight, None, **kwargs)
     else:
         layout = conv_layout(x, weight, None, **kwargs)
+        # TODO: This does not guard on the stride order decision,
+        # shall we use optimization_hint to handle unbacked?
         req_stride_order = ir.get_stride_order(
-            V.graph.sizevars.size_hints(layout.stride)
+            V.graph.sizevars.guarding_hints_or_throw(layout.stride)
         )
         x = ir.ExternKernel.require_stride_order(x, req_stride_order)  # type: ignore[assignment]
         weight = ir.ExternKernel.require_stride_order(weight, req_stride_order)  # type: ignore[assignment]
@@ -589,6 +625,22 @@ def channels_last_conv():
         ):
             choices.append(aten_conv1x1_via_mm.bind(args, layout))
 
+        is_depthwise = groups > 1 and in_chan == 1 and out_chan == groups
+        if is_depthwise and ndim == 1:
+            depthwise_configs = V.choices.get_depthwise_conv_configs(device_type)
+            for cfg in depthwise_configs:
+                depthwise_conv1d_template.maybe_append_choice(
+                    choices,
+                    input_nodes=(x, weight),
+                    layout=layout,
+                    KERNEL_SIZE=kernel_shape[0],
+                    CONV_STRIDE=stride[0],
+                    PADDING=padding[0],
+                    num_stages=cfg.num_stages,
+                    num_warps=cfg.num_warps,
+                    **cfg.kwargs,
+                )
+
         conv_configs = V.choices.get_conv_configs(device_type)
 
         dtype_size = x.get_dtype().itemsize
@@ -652,7 +704,8 @@ def channels_last_conv():
             groups=groups,
             n_spatial_dimensions=ndim,
         )
-    return autotune_select_algorithm("convolution", choices, args, layout)
+    node, _ = autotune_select_algorithm("convolution", choices, args, layout)
+    return node
 
 
 @register_lowering(aten._convolution)
diff --git a/torch/_inductor/kernel/custom_op.py b/torch/_inductor/kernel/custom_op.py
index 91aeba21507a6..dbc10b26af44d 100644
--- a/torch/_inductor/kernel/custom_op.py
+++ b/torch/_inductor/kernel/custom_op.py
@@ -1,22 +1,30 @@
 # Owner(s): ["module: inductor"]
 
+import contextlib
 import functools
 import logging
 from collections.abc import Callable
 from dataclasses import dataclass, field
-from typing import Any, Optional, Union
+from typing import Any
 
 import torch
+from torch._dynamo.utils import counters
 from torch._inductor.codegen.subgraph import SubgraphTemplate
-from torch._inductor.ir import Buffer, FixedLayout, ir_node_to_tensor, TensorBox
-from torch._inductor.lowering import lowerings, validate_ir
+from torch._inductor.ir import (
+    Buffer,
+    ChoiceCaller,
+    FixedLayout,
+    ir_node_to_tensor,
+    StorageBox,
+    TensorBox,
+)
+from torch._inductor.lowering import user_lowerings, validate_ir
 from torch._inductor.select_algorithm import (
     autotune_select_algorithm,
     ExternKernelChoice,
 )
 from torch._inductor.utils import convert_symint_to_expr
 from torch._inductor.virtualized import V
-from torch.utils._ordered_set import OrderedSet
 
 
 log = logging.getLogger(__name__)
@@ -29,7 +37,7 @@ class RangeBounds:
     """Inclusive range [start, end] for dimension-based dispatch."""
 
     start: int
-    end: Union[int, float]  # float('inf') for unbounded
+    end: int | float  # float('inf') for unbounded
 
     def __post_init__(self) -> None:
         if self.start < 1:
@@ -54,6 +62,9 @@ class ImplConfig:
     impl_name: str
     impl_func: Callable[..., Any] = field(compare=False, hash=False, repr=False)
     kwargs: dict[str, Any] = field(default_factory=dict)
+    config_patches: dict[str, Any] = field(
+        default_factory=dict, compare=False, hash=False, repr=False
+    )
 
     def __hash__(self) -> int:
         return hash((self.impl_name, tuple(sorted(self.kwargs.items()))))
@@ -97,6 +108,10 @@ def impl_func(self) -> Callable[..., Any]:
     def impl_kwargs(self) -> dict[str, Any]:
         return self.impl_config.kwargs
 
+    @property
+    def config_patches(self) -> dict[str, Any]:
+        return self.impl_config.config_patches
+
 
 def _detect_collective_ops(choices: list) -> bool:
     """
@@ -128,16 +143,20 @@ class CustomOpConfig:
 
     Args:
         decomposition: Optional functions to autotune. If not provided, default will be used.
+        config_patches: Optional dict of config patches to apply during kernel codegen
+            (e.g., {"coordinate_descent_tuning": True})
         **params: Parameters passed to the function
 
     Examples:
         CustomOpConfig(attention_impl, head_dim=32, method='chunked')
         CustomOpConfig(head_dim=32, method='chunked')
+        CustomOpConfig(decomposition, config_patches={"coordinate_descent_tuning": True})
     """
 
     def __init__(
         self,
-        decomposition: Optional[Callable[..., Any]] = None,
+        decomposition: Callable[..., Any] | None = None,
+        config_patches: dict[str, Any] | None = None,
         **params: Any,
     ):
         if decomposition is not None and not callable(decomposition):
@@ -146,10 +165,11 @@ def __init__(
             )
 
         self.decomposition = decomposition
+        self.config_patches = config_patches or {}
         self.params = params
 
     def get_decomposition(
-        self, default_impl: Optional[Callable[..., Any]] = None
+        self, default_impl: Callable[..., Any] | None = None
     ) -> Callable[..., Any]:
         """Return the decomposition function for this config.
         When decomposition is not specified, return the default implementation.
@@ -206,13 +226,13 @@ def _extract_tensor_inputs(
     ]
 
     for i, arg in enumerate(args):
-        if isinstance(arg, (TensorBox, Buffer)):
+        if isinstance(arg, (TensorBox, Buffer, StorageBox)):
             tensor_inputs.append(arg)
         else:
             non_tensor_kwargs[param_names[i]] = arg
 
     for key, value in kwargs.items():
-        if isinstance(value, (TensorBox, Buffer)):
+        if isinstance(value, (TensorBox, Buffer, StorageBox)):
             tensor_inputs.append(value)
         else:
             non_tensor_kwargs[key] = value
@@ -224,45 +244,30 @@ def _merge_config_and_runtime_kwargs(
     config_params: dict[str, Any],
     runtime_kwargs: dict[str, Any],
 ) -> dict[str, Any]:
-    """Merge config parameters with runtime kwargs. Runtime kwargs take precedence.
-       If there are conflicts, log a warning and use runtime value.
+    """Merge config parameters with runtime kwargs. Config params take precedence,
+    since they represent the values being autotuned.
 
     Args:
-        config_params: Parameters from CustomOpConfig
+        config_params: Parameters from CustomOpConfig (autotuning knobs)
         runtime_kwargs: Runtime non-tensor kwargs from _extract_tensor_inputs
 
     Returns:
-        Merged kwargs dictionary with runtime values taking precedence
+        Merged kwargs dictionary with config values taking precedence
     """
-    merged_kwargs = config_params.copy()
-
-    # Check for conflicts and let runtime kwargs dominate
-    conflicts = OrderedSet(config_params.keys()).intersection(runtime_kwargs.keys())
-
-    for key in conflicts:
-        log.warning(
-            "Parameter '%s' specified both in CustomOpConfig (%s) "
-            "and at runtime (%s). Using runtime value.",
-            key,
-            config_params[key],
-            runtime_kwargs[key],
-        )
-
-    # Runtime kwargs override config params
-    merged_kwargs.update(runtime_kwargs)
-
+    merged_kwargs = runtime_kwargs.copy()
+    merged_kwargs.update(config_params)
     return merged_kwargs
 
 
 def _adapt_user_input_gen_fns(
     inputs: list[Any],
-    arg_names: list[str],
+    op_overload: torch._ops.OpOverload,
     user_input_gen_fns: dict[str, Callable[[torch.Tensor], torch.Tensor]],
 ) -> dict[int, Callable[[Any], torch.Tensor]]:
     """Convert user input generators from name-based to index-based format.
     Inductor autotune's input_gen_fns expects index of arg_names as key.
     """
-
+    arg_names = [arg.name for arg in op_overload._schema.arguments]
     name_to_index = {name: i for i, name in enumerate(arg_names)}
     index_based_fns = {}
 
@@ -270,11 +275,9 @@ def _adapt_user_input_gen_fns(
         if name in name_to_index:
             index_based_fns[name_to_index[name]] = gen_fn
         else:
-            log.warning(
-                "Unknown argument name '%s' in input_gen_fns. "
-                "Available argument names: %s",
-                name,
-                list(name_to_index.keys()),
+            raise ValueError(
+                f"Unknown argument name '{name}' in input_gen_fns. "
+                f"Available argument names: {list(name_to_index.keys())}"
             )
 
     def create_internal_input_gen_fn(
@@ -283,7 +286,7 @@ def create_internal_input_gen_fn(
         """Create internal input generator that converts IR buffer to user's fake tensor."""
 
         def internal_input_gen_fn(ir_buffer: Any) -> torch.Tensor:
-            fake_tensor = ir_node_to_tensor(ir_buffer, guard_shape=False)
+            fake_tensor = ir_node_to_tensor(ir_buffer, replace_symbols_with_hints=True)
             assert fake_tensor is not None, "ir_node_to_tensor returned None"
             return user_function(fake_tensor)
 
@@ -302,9 +305,23 @@ def _group_ranges_by_impl(
     range_to_best_impl: dict[RangeBounds, ImplConfig],
 ) -> list[RangeImplGroup]:
     """Group ranges by implementation using semantic identity (name + kwargs)."""
+    from torch._inductor import config
+
     if not range_to_best_impl:
         return []
 
+    # Test mode: skip grouping to force torch.cond dispatch path
+    if config.test_configs.force_no_impl_grouping:
+        log.info("Test mode: skipping impl grouping, each range is separate group")
+        groups = []
+        for range_bounds, impl_config in sorted(
+            range_to_best_impl.items(), key=lambda x: x[0].start
+        ):
+            group = RangeImplGroup(impl_config)
+            group.add_range(range_bounds)
+            groups.append(group)
+        return groups
+
     # Group ranges by impl_config (uses __hash__ and __eq__ based on semantic identity)
     impl_to_group: dict[ImplConfig, RangeImplGroup] = {}
 
@@ -356,7 +373,7 @@ def _create_range_input_gen_fn(
     base_gen_fn: Callable[[torch.Tensor], torch.Tensor],
     dim_index: int,
     range_start: int,
-    range_end: Union[int, float],
+    range_end: int | float,
     range_upper_bound: int,
 ) -> Callable[[torch.Tensor], torch.Tensor]:
     """Create input generator that modifies target dimension to top of range.
@@ -393,21 +410,27 @@ def _default_input_gen_fn(fake_tensor: torch.Tensor) -> torch.Tensor:
 
 
 def _create_fallback_choice(
-    name: str,
-    default_impl: Callable[..., Any],
-    fake_output: torch.Tensor,
-    kwargs: dict[str, Any],
+    op_overload: torch._ops.OpOverload,
 ) -> ExternKernelChoice:
-    """Create fallback choice for default implementation."""
+    """Create or reuse fallback choice that calls the op eagerly.
 
-    def fallback_wrapper(*args: Any) -> Any:
-        return default_impl(*args, **kwargs)
+    Since kwargs are passed at bind time via maybe_append_choice rather than
+    baked into the kernel, the same ExternKernelChoice is reused across
+    compilations for the same op_overload.
+    """
+    fallback_name = (
+        f"{op_overload.name().replace('::', '_').replace('.', '_')}_fallback"
+    )
+
+    existing = ExternKernelChoice.lookup(fallback_name)
+    if existing is not None:
+        return existing
 
     return ExternKernelChoice(
-        kernel=fallback_wrapper,
-        name=f"{name}_fallback_default",
+        kernel=op_overload,
+        name=fallback_name,
         has_out_variant=False,
-        op_overload=default_impl,
+        op_overload=op_overload,
         use_fallback_kernel=True,
     )
 
@@ -418,11 +441,11 @@ def autotune_custom_op(
     inputs: list[torch.fx.Node],
     non_tensor_args: list[dict[str, Any]],
     op_overload: torch._ops.OpOverload,
-    user_input_gen_fns: Optional[
-        dict[str, Callable[[torch.Tensor], torch.Tensor]]
-    ] = None,
-    return_choice: bool = False,
-) -> Union[TensorBox, Any, tuple[Any, Any]]:
+    user_input_gen_fns: dict[str, Callable[[torch.Tensor], torch.Tensor]] | None = None,
+    config_patches_list: list[dict[str, Any]] | None = None,
+    min_speedup_threshold: float = 1.0,
+    benchmark_with_cudagraphs: bool = False,
+) -> tuple[TensorBox, ChoiceCaller]:
     """Autotune custom operations by comparing multiple decomposition implementations.
 
     Currently supports SINGLE OUTPUT custom ops only.
@@ -443,7 +466,7 @@ def autotune_custom_op(
                            and return real tensors for performance measurement.
 
     Returns:
-        IR node representing the optimized operation result
+        Tuple of (IR node representing the optimized operation result, winning ChoiceCaller)
 
     Raises:
         TypeError: If decompositions is not a list/tuple
@@ -466,56 +489,47 @@ def autotune_custom_op(
     # Convert user input generation functions BEFORE creating choices
     input_gen_fns: dict[int, Callable[[Any], torch.Tensor]] = {}
     if user_input_gen_fns:
-        import inspect
-
-        arg_names = (
-            list(inspect.signature(decompositions[0]).parameters.keys())
-            if decompositions
-            else []
+        input_gen_fns = _adapt_user_input_gen_fns(
+            inputs, op_overload, user_input_gen_fns
         )
-        input_gen_fns = _adapt_user_input_gen_fns(inputs, arg_names, user_input_gen_fns)
 
     template = SubgraphTemplate(name=name)
     choices = template.generate_custom_op_choices(
         name=name,
         decompositions=decompositions,
-        # pyrefly: ignore [no-matching-overload]
+        # pyrefly: ignore [bad-argument-type, no-matching-overload]
         input_nodes=list(inputs),
         non_tensor_args=non_tensor_args,
         input_gen_fns=input_gen_fns if input_gen_fns else None,
+        config_patches_list=config_patches_list,
     )
 
-    # Add default implementation as fallback
-    if op_overload and hasattr(op_overload, "_op"):
-        fallback_name = f"{name}_fallback_default"
-        from torch._inductor.select_algorithm import extern_kernels
-
-        # Skip if extern_kernel already registered to avoid duplicate registration error
-        if not hasattr(extern_kernels, fallback_name):
-            with V.fake_mode:
-                # pyrefly: ignore [no-matching-overload]
-                fake_inputs = [ir_node_to_tensor(inp) for inp in inputs]
-                fallback_kwargs = non_tensor_args[0] if non_tensor_args else {}
-                fake_output = op_overload(*fake_inputs, **fallback_kwargs)
-
-            output_size = tuple(convert_symint_to_expr(s) for s in fake_output.shape)
-            output_stride = tuple(
-                convert_symint_to_expr(s) for s in fake_output.stride()
-            )
+    # Add fallback choice that calls the op eagerly (not through inductor lowering)
+    # This provides a baseline to compare decompositions against
+    from torch._inductor import config
 
-            fallback_choice = _create_fallback_choice(
-                name, op_overload, fake_output, fallback_kwargs
-            )
-            fallback_choice.maybe_append_choice(
-                choices=choices,
-                input_nodes=list(inputs),
-                layout=FixedLayout(
-                    device=fake_output.device,
-                    dtype=fake_output.dtype,
-                    size=output_size,
-                    stride=output_stride,
-                ),
-            )
+    fallback_kwargs = non_tensor_args[0] if non_tensor_args else {}
+
+    with V.fake_mode:
+        # pyrefly: ignore [no-matching-overload]
+        fake_inputs = [ir_node_to_tensor(inp) for inp in inputs]
+        fake_output = op_overload(*fake_inputs, **fallback_kwargs)
+
+    output_size = tuple(convert_symint_to_expr(s) for s in fake_output.shape)
+    output_stride = tuple(convert_symint_to_expr(s) for s in fake_output.stride())
+
+    fallback_choice = _create_fallback_choice(op_overload)
+    fallback_choice.maybe_append_choice(
+        choices=choices,
+        input_nodes=list(inputs),
+        layout=FixedLayout(
+            device=fake_output.device,
+            dtype=fake_output.dtype,
+            size=output_size,
+            stride=output_stride,
+        ),
+        **fallback_kwargs,
+    )
 
     if not choices:
         raise RuntimeError(f"No valid choices generated for {name}")
@@ -529,22 +543,38 @@ def autotune_custom_op(
         input_nodes=list(inputs),
         layout=choices[0].layout,
         input_gen_fns=input_gen_fns,
-        return_choice=True,
         is_collective=is_collective,
+        min_speedup_threshold=min_speedup_threshold,
+        benchmark_with_cudagraphs=benchmark_with_cudagraphs,
     )
 
-    # Apply inlining for fusion if winning_choice has graph; otherwise return result as-is(default fallback impl)
-    if winning_choice.gm is not None:
-        # Skip inlining when return_choice=True since caller only needs choice metadata
-        # (e.g., range-based dispatch builds its own torch.cond from winning choices)
-        if return_choice:
-            log.debug(
-                "Skipping inline for return_choice: %s (name=%s)",
-                getattr(winning_choice, "name", type(winning_choice).__name__),
-                name,
-            )
-            return selected_result, winning_choice
+    # Test mode: force specific choice to win
+    force_choice = config.test_configs.force_custom_op_decomposition
+    if force_choice is True and winning_choice.gm is None:
+        # Force decomposition: pick first choice with a graph
+        for choice in choices:
+            if choice.gm is not None:
+                log.info(
+                    "Test mode: forcing decomposition %s over fallback",
+                    getattr(choice, "name", type(choice).__name__),
+                )
+                winning_choice = choice
+                selected_result = choice.output_node()
+                break
+    elif force_choice is False and winning_choice.gm is not None:
+        # Force fallback: pick first choice without a graph
+        for choice in choices:
+            if choice.gm is None:
+                log.info(
+                    "Test mode: forcing fallback %s over decomposition",
+                    getattr(choice, "name", type(choice).__name__),
+                )
+                winning_choice = choice
+                selected_result = choice.output_node()
+                break
 
+        # Always inline when winning_choice has a graph; callers extract choice metadata separately
+    if winning_choice.gm is not None:
         log.debug(
             "Inlining winning choice: %s (name=%s)",
             getattr(winning_choice, "name", type(winning_choice).__name__),
@@ -552,30 +582,35 @@ def autotune_custom_op(
         )
         from torch._inductor.codegen.subgraph import inline_subgraph_to_ir_nodes
 
+        ops_before = len(V.graph.operations)
         result = inline_subgraph_to_ir_nodes(winning_choice.gm, inputs, name)
-        return result
+
+        # Tag inlined operations with config_patches from the winning choice
+        config_patches = winning_choice.config_patches
+        if config_patches:
+            for op in V.graph.operations[ops_before:]:
+                op.set_config_patches(config_patches.copy())
+
+        return result, winning_choice
 
     log.debug(
         "Winning choice does not support inlining: %s (name=%s)",
         getattr(winning_choice, "name", type(winning_choice).__name__),
         name,
     )
-    if return_choice:
-        return selected_result, winning_choice
-    return selected_result
+    return selected_result, winning_choice
 
 
 def _generate_dynamic_configs(
     tensor_inputs: list[Buffer],
     config_generator: Callable[[dict[str, torch.Tensor]], list[CustomOpConfig]],
-    default_impl: Callable[..., Any],
+    op_overload: torch._ops.OpOverload,
     operation_name: str,
 ) -> list[CustomOpConfig]:
     """Generate configs dynamically based on input tensors at lowering time."""
-    import inspect
-
-    sig = inspect.signature(default_impl)
-    param_names = list(sig.parameters.keys())
+    # Get parameter names from op schema instead of impl signature
+    schema = op_overload._schema
+    param_names = [arg.name for arg in schema.arguments if not arg.kwarg_only]
 
     with V.fake_mode:
         fake_tensors = [ir_node_to_tensor(inp) for inp in tensor_inputs]
@@ -590,21 +625,24 @@ def _generate_dynamic_configs(
             f"got {type(configs)}"
         )
     if not configs:
-        raise ValueError(f"config_generator returned empty list for {operation_name}. ")
+        log.info(
+            "config_generator returned empty list for %s, will use default lowering",
+            operation_name,
+        )
+        return []
 
     return list(configs)
 
 
 def _prepare_configs_and_decompositions(
-    processed_configs: Optional[list[CustomOpConfig]],
-    config_generator: Optional[
-        Callable[[dict[str, torch.Tensor]], list[CustomOpConfig]]
-    ],
+    processed_configs: list[CustomOpConfig] | None,
+    config_generator: Callable[[dict[str, torch.Tensor]], list[CustomOpConfig]] | None,
     tensor_inputs: list[Any],
     default_impl: Callable[..., Any],
+    op_overload: torch._ops.OpOverload,
     runtime_kwargs: dict[str, Any],
     name: str,
-) -> tuple[list[Callable], list[dict[str, Any]]]:
+) -> tuple[list[Callable], list[dict[str, Any]], list[dict[str, Any]]]:
     """Prepare decompositions and merged kwargs from configs.
 
     Handles both static configs and dynamic config generation.
@@ -613,7 +651,7 @@ def _prepare_configs_and_decompositions(
     # Get configs: either generate dynamically or use static configs
     if config_generator is not None:
         configs_to_use = _generate_dynamic_configs(
-            tensor_inputs, config_generator, default_impl, name
+            tensor_inputs, config_generator, op_overload, name
         )
     else:
         assert processed_configs is not None
@@ -622,6 +660,7 @@ def _prepare_configs_and_decompositions(
     # Prepare decompositions and kwargs for autotuning
     decompositions = []
     non_tensor_args = []
+    config_patches_list = []
 
     for cfg in configs_to_use:
         decomp = cfg.get_decomposition(default_impl=default_impl)
@@ -631,7 +670,10 @@ def _prepare_configs_and_decompositions(
         merged_kwargs = _merge_config_and_runtime_kwargs(cfg.params, runtime_kwargs)
         non_tensor_args.append(merged_kwargs)
 
-    return decompositions, non_tensor_args
+        # Collect config_patches for each config
+        config_patches_list.append(cfg.config_patches)
+
+    return decompositions, non_tensor_args, config_patches_list
 
 
 def _standard_lowering_fn(
@@ -639,46 +681,84 @@ def _standard_lowering_fn(
     default_impl: Callable[..., Any],
     name: str,
     op_overload: torch._ops.OpOverload,
-    input_gen_fns: Optional[dict[str, Callable[[torch.Tensor], torch.Tensor]]],
+    input_gen_fns: dict[str, Callable[[torch.Tensor], torch.Tensor]] | None,
     tensor_inputs: list[Any],
     runtime_kwargs: dict[str, Any],
-    config_generator: Optional[
-        Callable[[dict[str, torch.Tensor]], list[CustomOpConfig]]
-    ] = None,
+    config_generator: Callable[[dict[str, torch.Tensor]], list[CustomOpConfig]]
+    | None = None,
+    min_speedup_threshold: float = 1.0,
+    benchmark_with_cudagraphs: bool = False,
 ) -> Any:
-    """Standard autotuning lowering function."""
-    decompositions, non_tensor_args = _prepare_configs_and_decompositions(
-        processed_configs,
-        config_generator,
-        tensor_inputs,
-        default_impl,
-        runtime_kwargs,
-        name,
+    """Standard autotuning lowering function.
+
+    Returns None if no configs/decompositions available, signaling caller to
+    use normal lowering.
+    """
+    decompositions, non_tensor_args, config_patches_list = (
+        _prepare_configs_and_decompositions(
+            processed_configs,
+            config_generator,
+            tensor_inputs,
+            default_impl,
+            op_overload,
+            runtime_kwargs,
+            name,
+        )
     )
 
-    result = autotune_custom_op(
+    # If no decompositions, signal caller to use normal lowering
+    if not decompositions:
+        return None
+
+    result, _ = autotune_custom_op(
         name=name,
         decompositions=decompositions,
         inputs=tensor_inputs,
         non_tensor_args=non_tensor_args,
+        config_patches_list=config_patches_list,
         op_overload=op_overload,
         user_input_gen_fns=input_gen_fns,
+        min_speedup_threshold=min_speedup_threshold,
+        benchmark_with_cudagraphs=benchmark_with_cudagraphs,
     )
 
     validate_ir(result)
     return result
 
 
+def _apply_config_patches_recursive(
+    operations: list,
+    config_patches: dict[str, Any],
+) -> None:
+    """Apply config_patches to operations, including those inside subgraphs."""
+    for op in operations:
+        if hasattr(op, "set_config_patches"):
+            op.set_config_patches(config_patches.copy())
+
+        # Recurse into any subgraphs (Conditional, WhileLoop, InvokeSubgraph, etc.)
+        for subgraph in op.get_subgraphs():
+            if subgraph.graph:
+                _apply_config_patches_recursive(
+                    subgraph.graph.operations, config_patches
+                )
+
+
 def _lower_single_impl(
     impl: Callable[..., Any],
     impl_kwargs: dict[str, Any],
     runtime_kwargs: dict[str, Any],
     tensor_inputs: list[Any],
     name: str,
+    config_patches: dict[str, Any] | None = None,
 ) -> Any:
-    """Lower a single implementation by tracing and inlining it."""
+    """Lower a single implementation by tracing and inlining it.
+
+    Uses error_on_new_guards() during tracing to detect if the impl adds guards.
+    Returns None if the impl adds guards, signaling caller to skip this choice.
+    """
     from torch._inductor.codegen.subgraph import inline_subgraph_to_ir_nodes
     from torch.fx.experimental.proxy_tensor import make_fx
+    from torch.fx.experimental.symbolic_shapes import _ShapeEnvGuardError
 
     from ..decomposition import select_decomp_table
 
@@ -687,17 +767,45 @@ def _lower_single_impl(
     def impl_wrapper(*tensors):
         return impl(*tensors, **merged_kwargs)
 
+    shape_env = V.fake_mode.shape_env
+
     with V.fake_mode:
         fake_inputs = tuple(ir_node_to_tensor(inp) for inp in tensor_inputs)
         decomposition_table = select_decomp_table()
-        impl_gm = make_fx(
-            impl_wrapper,
-            decomposition_table=decomposition_table,
-            tracing_mode="symbolic",
-        )(*fake_inputs)
+
+        context = (
+            shape_env.error_on_new_guards
+            if shape_env is not None
+            else contextlib.nullcontext
+        )
+        try:
+            with context():
+                impl_gm = make_fx(
+                    impl_wrapper,
+                    decomposition_table=decomposition_table,
+                    tracing_mode="symbolic",
+                )(*fake_inputs)
+        except (_ShapeEnvGuardError, AssertionError) as e:
+            is_guard_error = isinstance(e, _ShapeEnvGuardError) or (
+                isinstance(e, AssertionError)
+                and "Guard attempted while ShapeEnv guards are frozen" in str(e)
+            )
+            if not is_guard_error:
+                raise
+            log.info(
+                "Implementation %s adds guards, skipping custom op lowering",
+                impl.__name__,
+            )
+            counters["inductor"]["custom_op_decomp_guard_skips"] += 1
+            return None
 
     log.info("Inlining implementation: %s", impl.__name__)
+    ops_before = len(V.graph.operations)
     result = inline_subgraph_to_ir_nodes(impl_gm, tensor_inputs, name)
+
+    if config_patches:
+        _apply_config_patches_recursive(V.graph.operations[ops_before:], config_patches)
+
     validate_ir(result)
     return result
 
@@ -707,16 +815,17 @@ def _range_based_lowering_fn(
     default_impl: Callable[..., Any],
     name: str,
     op_overload: torch._ops.OpOverload,
-    input_gen_fns: Optional[dict[str, Callable[[torch.Tensor], torch.Tensor]]],
+    input_gen_fns: dict[str, Callable[[torch.Tensor], torch.Tensor]] | None,
     tensor_name: str,
     dim_index: int,
-    ranges: list[tuple[int, Union[int, float]]],
+    ranges: list[tuple[int, int | float]],
     tensor_inputs: list[Any],
     runtime_kwargs: dict[str, Any],
     range_upper_bound: int,
-    config_generator: Optional[
-        Callable[[dict[str, torch.Tensor]], list[CustomOpConfig]]
-    ] = None,
+    config_generator: Callable[[dict[str, torch.Tensor]], list[CustomOpConfig]]
+    | None = None,
+    min_speedup_threshold: float = 1.0,
+    benchmark_with_cudagraphs: bool = False,
 ) -> Any:
     """Range-based autotuning lowering function."""
     from torch._inductor.codegen.subgraph import inline_subgraph_to_ir_nodes
@@ -727,13 +836,16 @@ def _range_based_lowering_fn(
     log.info("=== Range-based Autotuning for %s ===", name)
     log.info("Dispatch on: %s[%d], Ranges: %s", tensor_name, dim_index, ranges)
 
-    decompositions, non_tensor_args = _prepare_configs_and_decompositions(
-        processed_configs,
-        config_generator,
-        tensor_inputs,
-        default_impl,
-        runtime_kwargs,
-        name,
+    decompositions, non_tensor_args, config_patches_list = (
+        _prepare_configs_and_decompositions(
+            processed_configs,
+            config_generator,
+            tensor_inputs,
+            default_impl,
+            op_overload,
+            runtime_kwargs,
+            name,
+        )
     )
 
     range_to_best_impl_map: dict[RangeBounds, ImplConfig] = {}
@@ -760,13 +872,12 @@ def _range_based_lowering_fn(
             non_tensor_args=non_tensor_args,
             op_overload=op_overload,
             user_input_gen_fns=range_input_gen_fns,
-            return_choice=True,
+            min_speedup_threshold=min_speedup_threshold,
+            benchmark_with_cudagraphs=benchmark_with_cudagraphs,
+            config_patches_list=config_patches_list,
         )
 
-        if (
-            hasattr(winning_choice, "decomposition")
-            and winning_choice.decomposition is not None
-        ):
+        if winning_choice.decomposition is not None:
             winning_impl = winning_choice.decomposition
             winning_kwargs = winning_choice.decomposition_kwargs
         else:
@@ -779,12 +890,15 @@ def _range_based_lowering_fn(
                 range_end if range_end != float("inf") else "inf",
             )
 
+        winning_config_patches = winning_choice.config_patches or {}
+
         # Create dataclass instances for cleaner code
         range_bounds = RangeBounds(range_start, range_end)
         impl_config = ImplConfig(
             impl_name=winning_impl.__name__,
             impl_func=winning_impl,
             kwargs=winning_kwargs,
+            config_patches=winning_config_patches,
         )
         range_to_best_impl_map[range_bounds] = impl_config
 
@@ -794,7 +908,9 @@ def _range_based_lowering_fn(
             impl_config.impl_name,
         )
 
-    # Group ranges by implementation (more aggressive than adjacent merging)
+    # Group ranges by implementation
+    from torch.fx.experimental.symbolic_shapes import _ShapeEnvGuardError
+
     impl_groups = _group_ranges_by_impl(range_to_best_impl_map)
 
     log.info("After grouping by implementation: %d impl groups", len(impl_groups))
@@ -806,7 +922,12 @@ def _range_based_lowering_fn(
         group = impl_groups[0]
         log.info("Only one implementation after grouping, directly inlining")
         return _lower_single_impl(
-            group.impl_func, group.impl_kwargs, runtime_kwargs, tensor_inputs, name
+            group.impl_func,
+            group.impl_kwargs,
+            runtime_kwargs,
+            tensor_inputs,
+            name,
+            config_patches=group.config_patches,
         )
 
     def dispatch_fn(*fake_tensors):
@@ -820,7 +941,6 @@ def dispatch_fn(*fake_tensors):
         dim_value = fake_tensors[0].size(dim_index)
 
         def build_range_predicate(ranges_list: list[RangeBounds]) -> torch.Tensor:
-            """Build OR predicate: (dim in range1) | (dim in range2) | ..."""
             predicates = []
             for rb in ranges_list:
                 end = int(rb.end) if rb.end != float("inf") else None
@@ -835,7 +955,6 @@ def build_range_predicate(ranges_list: list[RangeBounds]) -> torch.Tensor:
             return result  # pyrefly: ignore [bad-return]
 
         def build_nested_cond(idx: int):
-            """Recursively build nested torch.cond for impl_groups[idx:]."""
             if idx >= num_impl_groups:
                 raise RuntimeError(f"Invalid impl group index: {idx}")
 
@@ -866,29 +985,55 @@ def cond_wrapper(*ops, _ranges=group.ranges):
 
         return build_nested_cond(0)(*fake_tensors)
 
-    # Trace with make_fx using fake mode
     with V.fake_mode:
         fake_inputs = tuple(ir_node_to_tensor(inp) for inp in tensor_inputs)
         decomposition_table = select_decomp_table()
+        shape_env = V.fake_mode.shape_env
 
         log.info("Tracing torch.cond dispatch with symbolic shapes...")
 
         try:
-            dispatch_gm = make_fx(
-                dispatch_fn,
-                decomposition_table=decomposition_table,
-                tracing_mode="symbolic",
-            )(*fake_inputs)
+            context = (
+                shape_env.error_on_new_guards
+                if shape_env is not None
+                else contextlib.nullcontext
+            )
+            with context():
+                dispatch_gm = make_fx(
+                    dispatch_fn,
+                    decomposition_table=decomposition_table,
+                    tracing_mode="symbolic",
+                )(*fake_inputs)
 
             log.info("Successfully traced torch.cond dispatch")
             log.info("Traced graph:\n%s", dispatch_gm.graph)
 
+        except (_ShapeEnvGuardError, AssertionError) as e:
+            is_guard_error = isinstance(e, _ShapeEnvGuardError) or (
+                isinstance(e, AssertionError)
+                and "Guard attempted while ShapeEnv guards are frozen" in str(e)
+            )
+            if not is_guard_error:
+                raise
+            log.info("Dispatch function adds guards, skipping custom op lowering")
+            counters["inductor"]["custom_op_decomp_guard_skips"] += 1
+            return None
+
         except Exception:
             log.exception("make_fx tracing FAILED")
             raise
 
+    ops_before = len(V.graph.operations)
     result = inline_subgraph_to_ir_nodes(dispatch_gm, tensor_inputs, f"{name}_dispatch")
 
+    # Apply config_patches from all impl groups to inlined operations
+    # TODO - consider conflicting patches
+    merged_patches: dict[str, Any] = {}
+    for group in impl_groups:
+        merged_patches.update(group.config_patches)
+    if merged_patches:
+        _apply_config_patches_recursive(V.graph.operations[ops_before:], merged_patches)
+
     log.info(
         "Successfully created torch.cond dispatch for %d impl groups", len(impl_groups)
     )
@@ -902,14 +1047,15 @@ def _create_autotuning_lowering(
     default_impl: Callable[..., Any],
     name: str,
     op_overload: torch._ops.OpOverload,
-    input_gen_fns: Optional[dict[str, Callable[[torch.Tensor], torch.Tensor]]],
+    input_gen_fns: dict[str, Callable[[torch.Tensor], torch.Tensor]] | None,
     range_upper_bound: int,
     is_range_based: bool = False,
-    config_generator: Optional[
-        Callable[[dict[str, torch.Tensor]], list[CustomOpConfig]]
-    ] = None,
-    dispatch_on: Optional[tuple[str, int]] = None,
-    split_points: Optional[list[int]] = None,
+    config_generator: Callable[[dict[str, torch.Tensor]], list[CustomOpConfig]]
+    | None = None,
+    dispatch_on: tuple[str, int] | None = None,
+    split_points: list[int] | None = None,
+    min_speedup_threshold: float = 1.0,
+    benchmark_with_cudagraphs: bool = False,
 ) -> Callable[..., Any]:
     """Create the lowering function for autotuning."""
     if not is_range_based:
@@ -928,6 +1074,8 @@ def standard_lowering_wrapper(*args: Any, **kwargs: Any) -> Any:
                 tensor_inputs=tensor_inputs,
                 runtime_kwargs=runtime_kwargs,
                 config_generator=config_generator,
+                min_speedup_threshold=min_speedup_threshold,
+                benchmark_with_cudagraphs=benchmark_with_cudagraphs,
             )
 
         return standard_lowering_wrapper
@@ -956,21 +1104,24 @@ def range_based_lowering_wrapper(*args: Any, **kwargs: Any) -> Any:
             runtime_kwargs=runtime_kwargs,
             range_upper_bound=range_upper_bound,
             config_generator=config_generator,
+            min_speedup_threshold=min_speedup_threshold,
+            benchmark_with_cudagraphs=benchmark_with_cudagraphs,
         )
 
     return range_based_lowering_wrapper
 
 
 def register_custom_op_autotuning(
-    custom_op: torch._library.custom_ops.CustomOpDef,
-    configs: Optional[Union[list[CustomOpConfig], list[Callable[..., Any]]]] = None,
-    config_generator: Optional[
-        Callable[[dict[str, torch.Tensor]], list[CustomOpConfig]]
-    ] = None,
-    name: Optional[str] = None,
-    input_gen_fns: Optional[dict[str, Callable[[torch.Tensor], torch.Tensor]]] = None,
-    dispatch_on: Optional[dict[str, Any]] = None,
-    split_points: Optional[list[int]] = None,
+    custom_op: torch._library.custom_ops.CustomOpDef | torch._ops.OpOverload,
+    configs: list[CustomOpConfig] | list[Callable[..., Any]] | None = None,
+    config_generator: Callable[[dict[str, torch.Tensor]], list[CustomOpConfig]]
+    | None = None,
+    name: str | None = None,
+    input_gen_fns: dict[str, Callable[[torch.Tensor], torch.Tensor]] | None = None,
+    dispatch_on: dict[str, Any] | None = None,
+    split_points: list[int] | None = None,
+    min_speedup_threshold: float = 1.0,
+    benchmark_with_cudagraphs: bool = False,
 ) -> None:
     """Register custom op for autotuning with custom_op configs where each config
     specifies a decomposition implementation function with its parameter values.
@@ -978,7 +1129,8 @@ def register_custom_op_autotuning(
     runtime dispatch.
 
     Args:
-        custom_op: Custom operation (decorated function from @torch.library.custom_op)
+        custom_op: Custom operation (CustomOpDef from @torch.library.custom_op) or
+                   OpOverload (e.g., torch.ops.aten.mm.default)
         configs: List of CustomOpConfig objects for static inputs. Mutually exclusive with config_generator.
         config_generator: Dynamic config generator function that takes a dict mapping
                           parameter names to fake tensors, and returns list[CustomOpConfig]
@@ -992,6 +1144,15 @@ def register_custom_op_autotuning(
                 as [2048, inf] -> [2048, unbounded_size]. Set based on your expected workload size.
                 Default is DEFAULT_RANGE_UPPER_BOUND=65536.
         split_points: List of range endpoints in ascending order for range-based autotuning
+        min_speedup_threshold: Only pick a non-fallback choice if it beats the fallback
+            by at least this ratio. Default is 1.0 (any speedup wins). Set to e.g. 1.1
+            to require 10% speedup over fallback.
+        benchmark_with_cudagraphs: If True, benchmark the fallback kernel using CUDA graph
+            capture and replay for fair comparison with compiled kernels. Default is False.
+
+    The default/fallback implementation is automatically derived:
+    - For CustomOpDef: Uses the decorated function
+    - For OpOverload: Traces the op call, which falls through to normal inductor lowering
 
     Examples:
         # Static configs
@@ -1045,12 +1206,21 @@ def generate_k_split_configs(fake_tensors: dict[str, torch.Tensor]) -> list[Cust
     """
     from torch._library.custom_ops import CustomOpDef
 
-    if not isinstance(custom_op, CustomOpDef):
+    # Handle both CustomOpDef and OpOverload - derive impl_fn automatically
+    # Both cases call through op_overload so fake handlers are used during tracing
+    if isinstance(custom_op, CustomOpDef):
+        op_overload = custom_op._opoverload
+    elif isinstance(custom_op, torch._ops.OpOverload):
+        op_overload = custom_op
+    else:
         raise TypeError(
-            f"custom_op must be a CustomOpDef (decorated function from @torch.library.custom_op), "
-            f"got {type(custom_op)}."
+            f"custom_op must be a CustomOpDef or OpOverload, got {type(custom_op)}."
         )
 
+    # impl_fn calls through op_overload so fake handlers are used during tracing
+    def impl_fn(*args, **kwargs):
+        return op_overload(*args, **kwargs)
+
     # Validate configs and config_generator are mutually exclusive
     if configs is not None and config_generator is not None:
         raise ValueError(
@@ -1061,9 +1231,6 @@ def generate_k_split_configs(fake_tensors: dict[str, torch.Tensor]) -> list[Cust
     if configs is None and config_generator is None:
         raise ValueError("Must specify either 'configs' or 'config_generator'")
 
-    op_overload = custom_op._opoverload
-    default_impl = custom_op._init_fn
-
     # Process and validate static configs at registration time
     static_configs = None
     if configs is not None:
@@ -1087,7 +1254,7 @@ def generate_k_split_configs(fake_tensors: dict[str, torch.Tensor]) -> list[Cust
 
     # Validate range-based parameters
     is_range_based = dispatch_on is not None or split_points is not None
-    dispatch_on_tuple: Optional[tuple[str, int]] = None
+    dispatch_on_tuple: tuple[str, int] | None = None
     range_upper_bound = DEFAULT_RANGE_UPPER_BOUND
     if is_range_based:
         if dispatch_on is None or split_points is None:
@@ -1132,7 +1299,7 @@ def generate_k_split_configs(fake_tensors: dict[str, torch.Tensor]) -> list[Cust
     lowering_fn = _create_autotuning_lowering(
         # pyrefly: ignore [bad-argument-type]
         processed_configs=static_configs,
-        default_impl=default_impl,
+        default_impl=impl_fn,
         name=name,
         op_overload=op_overload,
         input_gen_fns=input_gen_fns,
@@ -1141,6 +1308,10 @@ def generate_k_split_configs(fake_tensors: dict[str, torch.Tensor]) -> list[Cust
         dispatch_on=dispatch_on_tuple,
         split_points=split_points,
         range_upper_bound=range_upper_bound,
+        min_speedup_threshold=min_speedup_threshold,
+        benchmark_with_cudagraphs=benchmark_with_cudagraphs,
     )
 
-    lowerings[op_overload] = lowering_fn
+    # Register in user_lowerings which takes priority over built-in lowerings
+    # The dispatch in graph.py checks user_lowerings first with recursion guard
+    user_lowerings[op_overload] = lowering_fn
diff --git a/torch/_inductor/kernel/flex/common.py b/torch/_inductor/kernel/flex/common.py
index bf4006de8399a..1ad6e50239170 100644
--- a/torch/_inductor/kernel/flex/common.py
+++ b/torch/_inductor/kernel/flex/common.py
@@ -5,7 +5,7 @@
 from collections.abc import Sequence
 from functools import partial
 from pathlib import Path
-from typing import Any, Optional, TYPE_CHECKING, Union
+from typing import Any, TYPE_CHECKING
 
 import sympy
 
@@ -18,7 +18,7 @@
 if TYPE_CHECKING:
     from torch._inductor.codegen.cuda_combined_scheduling import _IntLike
 else:
-    _IntLike = Union[int, sympy.Expr]
+    _IntLike = int | sympy.Expr
 
 
 from ...ir import (
@@ -46,7 +46,7 @@
 from ...utils import load_template
 
 
-SubgraphResults = Union[list[Optional[ComputedBuffer]], Optional[ComputedBuffer]]
+SubgraphResults = list[ComputedBuffer | None] | ComputedBuffer | None
 
 
 def zeros_and_scatter_lowering(shape: list[int], indices, values):
@@ -96,7 +96,7 @@ def zeros_and_scatter_lowering(shape: list[int], indices, values):
 
 def get_fwd_subgraph_outputs(
     subgraph_buffer: SubgraphResults, mask_graph_buffer: SubgraphResults
-) -> list[Optional[ComputedBuffer]]:
+) -> list[ComputedBuffer | None]:
     subgraph_buffer = (
         subgraph_buffer if isinstance(subgraph_buffer, Sequence) else [subgraph_buffer]
     )
@@ -134,7 +134,7 @@ def build_subgraph_module_buffer(
     with V.set_graph_handler(pw_subgraph):  # type: ignore[arg-type]
         pw_subgraph.run(*args)
 
-    def convert_output_node_to_buffer(output_buffer) -> Optional[ComputedBuffer]:
+    def convert_output_node_to_buffer(output_buffer) -> ComputedBuffer | None:
         if output_buffer is None:
             return None
         if isinstance(output_buffer, ComputedBuffer):
@@ -168,7 +168,7 @@ def build_subgraph_buffer(args: list[TensorBox], subgraph: Subgraph) -> Subgraph
     return build_subgraph_module_buffer(args, subgraph.graph_module)
 
 
-def maybe_realize(args: list[Optional[IRNode]]):
+def maybe_realize(args: list[IRNode | None]):
     """Accepts a list of optional IRNodes and returns a list of realized IRNodes"""
     return tree_map(
         lambda x: (
@@ -200,7 +200,7 @@ def create_placeholder(
     name: str,
     dtype: torch.dtype,
     device: torch.device,
-    size: Optional[list[int]] = None,
+    size: list[int] | None = None,
 ) -> TensorBox:
     """Creates a placeholder input buffers for producing subgraph_output."""
     input_buffer = InputBuffer(
diff --git a/torch/_inductor/kernel/flex/flex_attention.py b/torch/_inductor/kernel/flex/flex_attention.py
index c6017fadb0387..8c1523591c1ef 100644
--- a/torch/_inductor/kernel/flex/flex_attention.py
+++ b/torch/_inductor/kernel/flex/flex_attention.py
@@ -8,13 +8,14 @@
 import warnings
 from collections.abc import Sequence
 from dataclasses import dataclass
-from typing import Any, cast, Optional, TYPE_CHECKING, Union
+from typing import Any, cast, TYPE_CHECKING
 
 import sympy
 
 import torch
 from torch._inductor.virtualized import V
 from torch.nn.attention.flex_attention import _Backend
+from torch.utils._sympy.functions import FloorDiv, Mod
 
 from ...ir import ComputedBuffer, ExternKernel, FixedLayout, TensorBox
 from ...lowering import empty, empty_strided, lowerings, register_lowering, to_dtype
@@ -101,6 +102,7 @@ def get_float32_precision():
     source=load_flex_template("flex_attention")
     + load_flex_template("utilities")
     + load_flex_template("common"),
+    always_freeze_layout=True,
 )
 
 
@@ -317,10 +319,16 @@ def flex_attention(
 
     B = Bq
 
-    if seq_len_q % 128 != 0 or seq_len_kv % 128 != 0:
-        kernel_options.setdefault("IS_DIVISIBLE", False)
-    else:
+    seq_q_divisible = V.graph.sizevars.statically_known_true(
+        sympy.Eq(Mod(seq_len_q, 128), 0)
+    )
+    seq_kv_divisible = V.graph.sizevars.statically_known_true(
+        sympy.Eq(Mod(seq_len_kv, 128), 0)
+    )
+    if seq_q_divisible and seq_kv_divisible:
         kernel_options.setdefault("IS_DIVISIBLE", True)
+    else:
+        kernel_options.setdefault("IS_DIVISIBLE", False)
 
     # NB it is okay that the v_head_dim is different
     # We are using these to match fill order of the output.
@@ -352,7 +360,7 @@ def flex_attention(
     kernel_options.setdefault("SM_SCALE", scale)
 
     # Determine GQA broadcast factor.
-    gqa_shared_heads = Hq // Hkv
+    gqa_shared_heads = FloorDiv(Hq, Hkv)
     kernel_options.setdefault("GQA_SHARED_HEADS", gqa_shared_heads)
 
     # Inside of Triton kernel, only apply partial masking if partial blocks are computed.
@@ -484,7 +492,7 @@ def flex_attention(
         8: create_indices_fake,
     }
 
-    out = autotune_select_algorithm(
+    out, _ = autotune_select_algorithm(
         "flex_attention",
         choices,
         # Need to filter out symbols since there is an invariant
@@ -530,6 +538,7 @@ def flex_attention_backward_grid(
     name="flex_attention_backward",
     grid=flex_attention_backward_grid,
     source=load_flex_template("flex_backwards") + load_flex_template("utilities"),
+    always_freeze_layout=True,
 )
 
 
@@ -563,7 +572,7 @@ class JointOutputResult:
 
     grad_input: ComputedBuffer
     captured_grads_compute: list[ComputedBuffer]
-    captured_grads: list[Optional[TensorBox]]
+    captured_grads: list[TensorBox | None]
     mutated_grads: list[TensorBox]
 
 
@@ -702,8 +711,16 @@ def flex_attention_backward(*args, **kwargs):
         for k, v in kernel_options.items()
     }
     kernel_options.setdefault("FLOAT32_PRECISION", get_float32_precision())
-    seq_q_divisible = V.graph.sizevars.statically_known_true(seq_len_q % 128 == 0)
-    seq_kv_divisible = V.graph.sizevars.statically_known_true(seq_len_kv % 128 == 0)
+    kernel_options.setdefault("PRESCALE_QK", False)
+    kernel_options.setdefault("ROWS_GUARANTEED_SAFE", False)
+    kernel_options.setdefault("BLOCKS_ARE_CONTIGUOUS", False)
+    kernel_options.setdefault("WRITE_DQ", True)
+    seq_q_divisible = V.graph.sizevars.statically_known_true(
+        sympy.Eq(Mod(seq_len_q, 128), 0)
+    )
+    seq_kv_divisible = V.graph.sizevars.statically_known_true(
+        sympy.Eq(Mod(seq_len_kv, 128), 0)
+    )
     if seq_q_divisible and seq_kv_divisible:
         kernel_options.setdefault("IS_DIVISIBLE", True)
     else:
@@ -863,7 +880,7 @@ def flex_attention_backward(*args, **kwargs):
     kernel_options.setdefault("SM_SCALE", scale)
 
     # Determine GQA factor
-    gqa_shared_heads = Hq // Hkv
+    gqa_shared_heads = FloorDiv(Hq, Hkv)
     kernel_options.setdefault("GQA_SHARED_HEADS", gqa_shared_heads)
 
     # Inside of Triton kernel, only apply partial masking if partial blocks are computed.
@@ -1010,7 +1027,7 @@ def flex_attention_backward(*args, **kwargs):
         15: create_indices_fake,
     }
 
-    broadcasted_grad_key = autotune_select_algorithm(
+    broadcasted_grad_key, _ = autotune_select_algorithm(
         "flex_attention_backward",
         choices,
         [x for x in inputs_for_autotuning if isinstance(x, torch._inductor.ir.IRNode)],
@@ -1054,7 +1071,7 @@ def get_bwd_subgraph_outputs(
     subgraph_buffer: SubgraphResults,
     mask_graph_buffer: SubgraphResults,
     joint_outputs: JointOutputResult,
-) -> list[Optional[Union[ComputedBuffer, TensorBox]]]:
+) -> list[ComputedBuffer | TensorBox | None]:
     subgraph_buffer = (
         subgraph_buffer if isinstance(subgraph_buffer, Sequence) else [subgraph_buffer]
     )
diff --git a/torch/_inductor/kernel/flex/flex_cpu.py b/torch/_inductor/kernel/flex/flex_cpu.py
index 45ea6bc2a6a62..b16975f7113eb 100644
--- a/torch/_inductor/kernel/flex/flex_cpu.py
+++ b/torch/_inductor/kernel/flex/flex_cpu.py
@@ -328,7 +328,7 @@ def convert_mask_graph_module(mask_graph):
         key,
         value,
     ]
-    res = autotune_select_algorithm(
+    res, _ = autotune_select_algorithm(
         "flex_attention",
         _choices,
         inputs_for_autotuning,
diff --git a/torch/_inductor/kernel/flex/flex_decoding.py b/torch/_inductor/kernel/flex/flex_decoding.py
index 4111915c26082..2cd25c5fa1063 100644
--- a/torch/_inductor/kernel/flex/flex_decoding.py
+++ b/torch/_inductor/kernel/flex/flex_decoding.py
@@ -8,6 +8,7 @@
 
 import torch
 from torch._inductor.virtualized import V
+from torch.utils._sympy.functions import FloorDiv, Mod
 
 from ... import ir
 from ...ir import FixedLayout, FlexibleLayout
@@ -69,7 +70,7 @@ def _use_flex_decoding(query, kv_indices, value, kernel_options, enable_gqa) ->
 
     Hq = query.get_size()[1]
     Hkv = value.get_size()[1]
-    ratio = Hq // Hkv
+    ratio = FloorDiv(Hq, Hkv)
 
     pw_of_two = V.graph.sizevars.guard_or_false(
         sympy.And(sympy.Gt(ratio, 0), sympy.Eq(ratio & (ratio - 1), 0))
@@ -116,6 +117,7 @@ def flex_decoding_grid(batch_size, kv_heads, gqa_group_size, seq_len_q, d_model,
     source=load_flex_template("flex_decode")
     + load_flex_template("utilities")
     + load_flex_template("common"),
+    always_freeze_layout=True,
 )
 
 
@@ -180,10 +182,10 @@ def create_flex_decoding_kernel(*args, **kwargs):
     }
 
     seq_q_divisible = V.graph.sizevars.statically_known_true(
-        sympy.Eq(seq_len_q % 128, 0)
+        sympy.Eq(Mod(seq_len_q, 128), 0)
     )
     seq_kv_divisible = V.graph.sizevars.statically_known_true(
-        sympy.Eq(seq_len_kv % 128, 0)
+        sympy.Eq(Mod(seq_len_kv, 128), 0)
     )
     if seq_q_divisible and seq_kv_divisible:
         kernel_options.setdefault("IS_DIVISIBLE", True)
@@ -191,7 +193,7 @@ def create_flex_decoding_kernel(*args, **kwargs):
         kernel_options.setdefault("IS_DIVISIBLE", False)
 
     # Calculate GQA head sharing
-    gqa_shared_heads = Hq // Hkv
+    gqa_shared_heads = FloorDiv(Hq, Hkv)
     if not is_power_of_2(gqa_shared_heads):
         raise ValueError(
             "Number of shared query heads sharing the same KV head must be power of 2. "
@@ -301,7 +303,7 @@ def create_flex_decoding_kernel(*args, **kwargs):
 
     kernel_options.setdefault(
         "SAFE_M_BOUNDARY",
-        ((seq_len_q * gqa_shared_heads) % kernel_options["BLOCK_M"]) == 0,
+        Mod(seq_len_q * gqa_shared_heads, kernel_options["BLOCK_M"]) == 0,
     )
     # TODO: This feels sketchy
     kernel_options.setdefault("SAFE_N_BOUNDARY", True)
@@ -412,7 +414,7 @@ def create_flex_decoding_kernel(*args, **kwargs):
         8: create_indices_fake,
     }
 
-    buf_ACC = autotune_select_algorithm(
+    buf_ACC, _ = autotune_select_algorithm(
         "flex_decoding",
         choices,
         inputs_for_flex_decoding,
diff --git a/torch/_inductor/kernel/flex/flex_flash_attention.py b/torch/_inductor/kernel/flex/flex_flash_attention.py
index b955964f9db51..9401da5155c81 100644
--- a/torch/_inductor/kernel/flex/flex_flash_attention.py
+++ b/torch/_inductor/kernel/flex/flex_flash_attention.py
@@ -1,11 +1,12 @@
 # mypy: allow-untyped-defs
 """Call into flash-attention 4 for flexattention"""
 
+import dataclasses
 import functools
 import importlib
 from collections.abc import Callable, Sequence
 from contextlib import contextmanager
-from typing import Any, cast, Literal, Optional
+from typing import Any, cast, Literal
 
 import sympy
 from sympy import Expr, Integer
@@ -15,8 +16,45 @@
 
 from ...ir import FixedLayout, ShapeAsConstantBuffer, Subgraph, TensorBox
 from ...lowering import empty_strided
+from ...select_algorithm import autotune_select_algorithm
 from ...virtualized import V
-from .common import infer_dense_strides, load_flex_template, SubgraphResults
+from .common import (
+    create_indices_fake,
+    create_num_blocks_fake_generator,
+    infer_dense_strides,
+    load_flex_template,
+    SubgraphResults,
+)
+
+
+@dataclasses.dataclass
+class FlexFlashConfig:
+    """Autotuning configuration for CuteDSL flex flash attention kernels.
+
+    score_mod_vec_size: Number of elements processed per thread in the score_mod
+        application loop. Maps to score_mod.__vec_size__ in CuTe flash attention.
+        None uses the kernel default. Only effective for forward; backward does
+        not currently support vectorized score_mod.
+    """
+
+    score_mod_vec_size: int | None = None
+
+
+def _get_flex_flash_fwd_configs(
+    has_score_mod: bool,
+    has_aux_tensors: bool,
+) -> list[FlexFlashConfig]:
+    if not has_score_mod or not torch._inductor.config.max_autotune:
+        return [FlexFlashConfig()]
+    if has_aux_tensors:
+        return [FlexFlashConfig(score_mod_vec_size=1)]
+    return [
+        FlexFlashConfig(score_mod_vec_size=v) for v in (1, 2, 4, 8, 16, 32, 64, 128)
+    ]
+
+
+def _get_flex_flash_bwd_configs() -> list[FlexFlashConfig]:
+    return [FlexFlashConfig()]
 
 
 aten = torch.ops.aten
@@ -376,29 +414,52 @@ def create_flex_flash_attention_kernel(
         subgraphs.append(subgraph_buffer)
     subgraphs.append(mask_graph_buffer)
 
-    with patch_fixed_layout_indexer_for_cutedsl():
-        error = flash_attention_cutedsl_template.maybe_append_choice(
-            choices,
-            input_nodes=input_nodes,
-            layout=output_layout,
-            mutated_inputs=[lse],
-            subgraphs=subgraphs,
-            SM_SCALE=scale,
-            HAS_SCORE_MOD=has_score_mod,
-            NEEDS_BLOCK_MASK=needs_block_mask,
-            SPARSE_Q_BLOCK_SIZE=sparse_q_block_size,
-            SPARSE_KV_BLOCK_SIZE=sparse_kv_block_size,
-        )
+    configs = _get_flex_flash_fwd_configs(
+        has_score_mod, len(score_mod_other_buffers) > 0
+    )
+
+    error: NotImplementedError | None = None
+    for conf in configs:
+        with patch_fixed_layout_indexer_for_cutedsl():
+            error = flash_attention_cutedsl_template.maybe_append_choice(
+                choices,
+                input_nodes=input_nodes,
+                layout=output_layout,
+                mutated_inputs=[lse],
+                subgraphs=subgraphs,
+                SM_SCALE=scale,
+                HAS_SCORE_MOD=has_score_mod,
+                SCORE_MOD_VEC_SIZE=conf.score_mod_vec_size,
+                NEEDS_BLOCK_MASK=needs_block_mask,
+                SPARSE_Q_BLOCK_SIZE=sparse_q_block_size,
+                SPARSE_KV_BLOCK_SIZE=sparse_kv_block_size,
+            )
+        if error is not None and len(configs) == 1:
+            raise RuntimeError(f"CuteDSL template failed: {error}")
 
     for choice in choices:
         wrap_choice_render_with_cutedsl_indexer(choice)
 
-    if error or not choices:
-        # Fallback to original implementation
+    if not choices:
         raise RuntimeError(f"CuteDSL template failed: {error}")
 
-    # No autotune for now
-    template_output = choices[0].output_node()
+    input_gen_fns: dict[int, Callable] | None = None
+    if has_full_blocks:
+        input_gen_fns = {
+            4: create_num_blocks_fake_generator(kv_indices),
+            5: create_indices_fake,
+            6: create_num_blocks_fake_generator(full_kv_indices),
+            7: create_indices_fake,
+        }
+
+    template_output, _ = autotune_select_algorithm(
+        "flex_flash_attention",
+        choices,
+        input_nodes,
+        output_layout,
+        input_gen_fns=input_gen_fns,
+        return_multi_template=False,
+    )
 
     return (template_output, lse)
 
@@ -406,8 +467,8 @@ def create_flex_flash_attention_kernel(
 def _can_use_flex_flash_attention_backward(
     fw_subgraph: Subgraph,
     mask_graph: Subgraph,
-    joint_outputs: Optional[Any] = None,
-    score_mod_other_buffers: Optional[Sequence[TensorBox]] = None,
+    joint_outputs: Any | None = None,
+    score_mod_other_buffers: Sequence[TensorBox] | None = None,
     num_score_mod_placeholders: int = 5,
 ) -> tuple[bool, str]:
     if not ensure_flash_available():
@@ -440,8 +501,8 @@ def _use_flex_flash_attention_backward(
     fw_subgraph: Subgraph,
     mask_graph: Subgraph,
     backend: Literal["AUTO", "TRITON", "FLASH", "TRITON_DECODE"],
-    joint_outputs: Optional[Any] = None,
-    score_mod_other_buffers: Optional[Sequence[TensorBox]] = None,
+    joint_outputs: Any | None = None,
+    score_mod_other_buffers: Sequence[TensorBox] | None = None,
 ) -> bool:
     """Determine if we should use flex flash attention for the given inputs.
 
@@ -485,14 +546,14 @@ def create_flex_flash_attention_backward_kernel(
     kernel_options: dict[str, Any],
     sparse_q_block_size: int,
     sparse_kv_block_size: int,
-    fw_subgraph_buffer: Optional[SubgraphResults] = None,
-    joint_subgraph_buffer: Optional[Any] = None,
-    score_mod_other_buffers: Optional[list[TensorBox]] = None,
-    mask_graph_buffer: Optional[SubgraphResults] = None,
-    q_num_blocks: Optional[TensorBox] = None,
-    q_indices: Optional[TensorBox] = None,
-    full_q_num_blocks: Optional[TensorBox] = None,
-    full_q_indices: Optional[TensorBox] = None,
+    fw_subgraph_buffer: SubgraphResults | None = None,
+    joint_subgraph_buffer: Any | None = None,
+    score_mod_other_buffers: list[TensorBox] | None = None,
+    mask_graph_buffer: SubgraphResults | None = None,
+    q_num_blocks: TensorBox | None = None,
+    q_indices: TensorBox | None = None,
+    full_q_num_blocks: TensorBox | None = None,
+    full_q_indices: TensorBox | None = None,
 ) -> tuple[TensorBox | ShapeAsConstantBuffer, TensorBox, TensorBox, tuple]:
     """Create a CuteDSL flash attention backward kernel for the default mod path."""
     if not ensure_flash_available():
@@ -580,26 +641,49 @@ def create_flex_flash_attention_backward_kernel(
     if has_block_mask:
         subgraphs.append(mask_graph_buffer)
 
-    with patch_fixed_layout_indexer_for_cutedsl():
-        error = flash_attention_backward_cutedsl_template.maybe_append_choice(
-            choices,
-            input_nodes=input_nodes,
-            layout=output_layout,
-            mutated_inputs=[grad_key, grad_value],
-            subgraphs=subgraphs if subgraphs else None,
-            SM_SCALE=scale,
-            HAS_SCORE_MOD=has_score_mod,
-            HAS_BLOCK_MASK=has_block_mask,
-            SPARSE_Q_BLOCK_SIZE=sparse_q_block_size,
-            SPARSE_KV_BLOCK_SIZE=sparse_kv_block_size,
-        )
+    configs = _get_flex_flash_bwd_configs()
+
+    error: NotImplementedError | None = None
+    for conf in configs:
+        with patch_fixed_layout_indexer_for_cutedsl():
+            error = flash_attention_backward_cutedsl_template.maybe_append_choice(
+                choices,
+                input_nodes=input_nodes,
+                layout=output_layout,
+                mutated_inputs=[grad_key, grad_value],
+                subgraphs=subgraphs or None,
+                SM_SCALE=scale,
+                HAS_SCORE_MOD=has_score_mod,
+                SCORE_MOD_VEC_SIZE=conf.score_mod_vec_size,
+                HAS_BLOCK_MASK=has_block_mask,
+                SPARSE_Q_BLOCK_SIZE=sparse_q_block_size,
+                SPARSE_KV_BLOCK_SIZE=sparse_kv_block_size,
+            )
+        if error is not None and len(configs) == 1:
+            raise RuntimeError(f"CuteDSL template failed: {error}")
 
     for choice in choices:
         wrap_choice_render_with_cutedsl_indexer(choice)
 
-    if error or not choices:
+    if not choices:
         raise RuntimeError(f"CuteDSL template failed: {error}")
 
-    template_output = choices[0].output_node()
+    input_gen_fns: dict[int, Callable] | None = None
+    if has_block_mask:
+        input_gen_fns = {
+            8: create_num_blocks_fake_generator(q_indices),
+            9: create_indices_fake,
+            10: create_num_blocks_fake_generator(full_q_indices),
+            11: create_indices_fake,
+        }
+
+    template_output, _ = autotune_select_algorithm(
+        "flex_flash_attention_backward",
+        choices,
+        input_nodes,
+        output_layout,
+        input_gen_fns=input_gen_fns,
+        return_multi_template=False,
+    )
 
     return (template_output, grad_key, grad_value, tuple())
diff --git a/torch/_inductor/kernel/flex/templates/flash_attention.py.jinja b/torch/_inductor/kernel/flex/templates/flash_attention.py.jinja
index 0ee44ff5c99cd..3061747f12a9c 100644
--- a/torch/_inductor/kernel/flex/templates/flash_attention.py.jinja
+++ b/torch/_inductor/kernel/flex/templates/flash_attention.py.jinja
@@ -27,6 +27,9 @@
         ) | indent_except_first(2) }}
         return tSrS_ssa
     {{ set_cute_hash("score_mod", "score") }}
+    {% if SCORE_MOD_VEC_SIZE %}
+    score_mod.__vec_size__ = {{SCORE_MOD_VEC_SIZE}}
+    {% endif %}
     {% else %}
     score_mod = None
     {% endif %}
diff --git a/torch/_inductor/kernel/flex/templates/flex_attention.py.jinja b/torch/_inductor/kernel/flex/templates/flex_attention.py.jinja
index b92ea6c14a33f..74fa146b13423 100644
--- a/torch/_inductor/kernel/flex/templates/flex_attention.py.jinja
+++ b/torch/_inductor/kernel/flex/templates/flex_attention.py.jinja
@@ -102,6 +102,11 @@
     stride_kv_num_blks_h = {{stride("KV_NUM_BLKS", 1)}}
     stride_kv_idx_h = {{stride("KV_IDX", 1)}}
     stride_kv_idx_m = {{stride("KV_IDX", 2)}}
+    {%- if HAS_FULL_BLOCKS %}
+    stride_full_kv_num_blks_h = {{stride("FULL_KV_NUM_BLKS", 1)}}
+    stride_full_kv_idx_h = {{stride("FULL_KV_IDX", 1)}}
+    stride_full_kv_idx_m = {{stride("FULL_KV_IDX", 2)}}
+    {%- endif %}
 
     # initialize pointer to m and l
     m_i = tl.zeros([BLOCK_M], dtype=tl.float32) - float("inf")
@@ -114,6 +119,10 @@
     sparse_hz_offset = sparse_idx_z * SPARSE_HQ + sparse_idx_hq
     sparse_kv_num_blks_offset = sparse_hz_offset * stride_kv_num_blks_h + q_start // SPARSE_Q_MULTIPLE
     sparse_kv_idx_offset = sparse_hz_offset * stride_kv_idx_h + (q_start // SPARSE_Q_MULTIPLE) * stride_kv_idx_m  # noqa: B950
+    {%- if HAS_FULL_BLOCKS %}
+    sparse_full_kv_num_blks_offset = sparse_hz_offset * stride_full_kv_num_blks_h + q_start // SPARSE_Q_MULTIPLE
+    sparse_full_kv_idx_offset = sparse_hz_offset * stride_full_kv_idx_h + (q_start // SPARSE_Q_MULTIPLE) * stride_full_kv_idx_m  # noqa: B950
+    {%- endif %}
 
     {%- if USE_TMA %}
     q = tl.load_tensor_descriptor(
@@ -159,9 +168,9 @@
     # apply mask_mod to them - only score_mod
     if HAS_FULL_BLOCKS:
         # FULL_KV_IDX and FULL_KV_NUM_BLKS are always contiguous.
-        kv_indices = FULL_KV_IDX + sparse_kv_idx_offset
+        kv_indices = FULL_KV_IDX + sparse_full_kv_idx_offset
         kv_start = tl.load(kv_indices) * SPARSE_KV_BLOCK_SIZE # first kv block we're loading
-        kv_num_blocks = tl.load(FULL_KV_NUM_BLKS + sparse_kv_num_blks_offset)
+        kv_num_blocks = tl.load(FULL_KV_NUM_BLKS + sparse_full_kv_num_blks_offset)
         block_n_end = tl.minimum(kv_num_blocks * SPARSE_KV_MULTIPLE, tl.maximum(tl.cdiv(KV_LEN, BLOCK_N), 1))
         # K and V pointers will be passed directly to forward_inner
         offs_n = kv_start + tl.arange(0, BLOCK_N)
diff --git a/torch/_inductor/kernel/flex/templates/flex_backwards.py.jinja b/torch/_inductor/kernel/flex/templates/flex_backwards.py.jinja
index d753a97cd2d50..41f4d8c8f1316 100644
--- a/torch/_inductor/kernel/flex/templates/flex_backwards.py.jinja
+++ b/torch/_inductor/kernel/flex/templates/flex_backwards.py.jinja
@@ -90,12 +90,21 @@
         stride_kv_num_blks_h = {{stride("KV_NUM_BLKS", 1)}}
         stride_kv_idx_h = {{stride("KV_IDX", 1)}}
         stride_kv_idx_m = {{stride("KV_IDX", 2)}}
+        {%- if HAS_FULL_BLOCKS %}
+        stride_full_kv_num_blks_h = {{stride("FULL_KV_NUM_BLKS", 1)}}
+        stride_full_kv_idx_h = {{stride("FULL_KV_IDX", 1)}}
+        stride_full_kv_idx_m = {{stride("FULL_KV_IDX", 2)}}
+        {%- endif %}
 
         sparse_idx_hq2 = off_hq2 % SPARSE_HQ
         sparse_hz_offset = sparse_idx_z * SPARSE_HQ + sparse_idx_hq2
 
         sparse_kv_num_blks_offset = sparse_hz_offset * stride_kv_num_blks_h + off_pid_mask
         sparse_kv_idx_offset = sparse_hz_offset * stride_kv_idx_h + off_pid_mask * stride_kv_idx_m  # noqa: B950
+        {%- if HAS_FULL_BLOCKS %}
+        sparse_full_kv_num_blks_offset = sparse_hz_offset * stride_full_kv_num_blks_h + off_pid_mask
+        sparse_full_kv_idx_offset = sparse_hz_offset * stride_full_kv_idx_h + off_pid_mask * stride_full_kv_idx_m  # noqa: B950
+        {%- endif %}
 
         # Offset Q, DQ, DO, DELTA & LSE. These inputs are offsetted by query heads.
         q_adj2 = (stride_qh * off_hq2 + stride_qz * off_zq).to(tl.int64)
@@ -190,9 +199,9 @@
         if HAS_FULL_BLOCKS:
             # ~~~~~~~~~~~ partial unmasked blocks ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
             # FULL_KV_IDX and FULL_KV_NUM_BLKS are always contiguous.
-            kv_indices = FULL_KV_IDX + sparse_kv_idx_offset
+            kv_indices = FULL_KV_IDX + sparse_full_kv_idx_offset
             kv_start = tl.load(kv_indices) * SPARSE_KV_BLOCK_SIZE # first kv block we're loading
-            sparse_kv_num_blocks = tl.load(FULL_KV_NUM_BLKS + sparse_kv_num_blks_offset)
+            sparse_kv_num_blocks = tl.load(FULL_KV_NUM_BLKS + sparse_full_kv_num_blks_offset)
 
             offs_n2 = kv_start + tl.arange(0, BLOCK_N2)
             dq = bwd_dq_inner(
@@ -223,6 +232,11 @@
         stride_q_num_blks_h = {{stride("Q_NUM_BLKS", 1)}}
         stride_q_idx_h = {{stride("Q_IDX", 1)}}
         stride_q_idx_n = {{stride("Q_IDX", 2)}}
+        {%- if HAS_FULL_BLOCKS %}
+        stride_full_q_num_blks_h = {{stride("FULL_Q_NUM_BLKS", 1)}}
+        stride_full_q_idx_h = {{stride("FULL_Q_IDX", 1)}}
+        stride_full_q_idx_n = {{stride("FULL_Q_IDX", 2)}}
+        {%- endif %}
 
 
         dv = tl.zeros([BLOCK_N1, V_HEAD_DIM_ROUNDED], dtype=tl.float32)
@@ -281,6 +295,10 @@
 
             sparse_q_num_blks_offset = sparse_hz_offset * stride_q_num_blks_h + pid_mask
             sparse_q_idx_offset = sparse_hz_offset * stride_q_idx_h + pid_mask * stride_q_idx_n  # noqa: B950
+            {%- if HAS_FULL_BLOCKS %}
+            sparse_full_q_num_blks_offset = sparse_hz_offset * stride_full_q_num_blks_h + pid_mask
+            sparse_full_q_idx_offset = sparse_hz_offset * stride_full_q_idx_h + pid_mask * stride_full_q_idx_n  # noqa: B950
+            {%- endif %}
 
             # ~~~~~~~~~~~~~~~ fully unmasked blocks ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
             # Q_IDX and Q_NUM_BLKS are always contiguous.
@@ -321,9 +339,9 @@
             if HAS_FULL_BLOCKS:
                 # ~~~~~~~~~~~~~~~ fully unmasked blocks ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
                 # FULL_Q_IDX and FULL_Q_NUM_BLKS are always contiguous.
-                q_indices = FULL_Q_IDX + sparse_q_idx_offset
+                q_indices = FULL_Q_IDX + sparse_full_q_idx_offset
                 q_start = tl.load(q_indices) * SPARSE_Q_BLOCK_SIZE # first q block we're loading
-                sparse_q_num_blocks = tl.load(FULL_Q_NUM_BLKS + sparse_q_num_blks_offset)
+                sparse_q_num_blocks = tl.load(FULL_Q_NUM_BLKS + sparse_full_q_num_blks_offset)
 
                 offs_m1 = q_start + tl.arange(0, BLOCK_M1)
                 dk, dv = bwd_dkdv_inner(
diff --git a/torch/_inductor/kernel/mm.py b/torch/_inductor/kernel/mm.py
index 112f201d52f1d..13653d41e66d5 100644
--- a/torch/_inductor/kernel/mm.py
+++ b/torch/_inductor/kernel/mm.py
@@ -1,7 +1,7 @@
 # mypy: allow-untyped-defs
 import functools
 import logging
-from typing import Any, Optional, Union
+from typing import Any
 
 import torch
 from torch._dynamo.utils import counters
@@ -100,6 +100,14 @@
     source=load_kernel_template("triton_persistent_tma_mm"),
 )
 
+# Non-TMA Triton template for persistent MM
+# used on AMD
+persistent_mm_template = TritonTemplate(
+    name="mm_persistent",
+    grid=persistent_mm_grid,
+    source=load_kernel_template("triton_persistent_mm"),
+)
+
 
 scaled_mm_device_tma_epilogue_scaling_template = TritonTemplate(
     name="scaled_mm_device_tma_epilogue_scaling",
@@ -359,7 +367,6 @@ def _to_dtype(x):
                 return ops.to_dtype(x, mat1.dtype, use_compute_types=False)
 
             args = [make_pointwise(_to_dtype)(x) for x in args]
-
         mul_pointwise = make_pointwise(ops.dot)(*args)
         dot_reduction = make_reduction("dot")(mul_pointwise, 1)
 
@@ -396,7 +403,7 @@ def _to_dtype(x):
         aten_handler = aten_mm_dtype
         aten_extra_kwargs = {"out_dtype": out_dtype}
 
-    templates_to_use: list[Union[ExternKernelChoice, KernelTemplate]] = []
+    templates_to_use: list[ExternKernelChoice | KernelTemplate] = []
     kwarg_overrides: dict[str, dict[str, Any]] = {}
     if use_aten_gemm_kernels():
         templates_to_use.append(aten_handler)
@@ -419,18 +426,17 @@ def _to_dtype(x):
         if is_exhaustive or not use_decompose_k_choice(m, n, k, threshold_multiple=2):
             templates_to_use.append(mm_template)
 
-            if use_triton_blackwell_tma_template(mat1, mat2, output_layout=layout):
+            if use_triton_blackwell_tma_template(
+                mat1, mat2, output_layout=layout, add_guards=True
+            ):
                 templates_to_use.append(blackwell_ws_persistent_device_tma_mm_template)
-            elif use_triton_tma_template(mat1, mat2, output_layout=layout):
-                templates_to_use.append(persistent_tma_mm_template)
-
-            if (
-                inductor_config.is_fbcode()
-                and inductor_config.triton.enable_tlx_templates
+            elif use_triton_tma_template(
+                mat1, mat2, output_layout=layout, add_guards=True
             ):
-                from torch._inductor.fb.tlx_templates.mm_templates import append_tlx
-
-                templates_to_use = append_tlx(templates_to_use)
+                if torch.version.hip is None:
+                    templates_to_use.append(persistent_tma_mm_template)
+                else:
+                    templates_to_use.append(persistent_mm_template)
 
         templates_to_use.append(mm_contiguous_subgraph_template)
 
@@ -540,13 +546,14 @@ def _to_dtype(x):
     ):
         return box
 
-    return autotune_select_algorithm(
+    node, _ = autotune_select_algorithm(
         name,
         choices,
         kernel_inputs.nodes(),
         layout,
         best_config_future=best_config_future,
     )
+    return node
 
 
 @register_lowering(aten._int_mm, type_promotion_kind=None)
@@ -576,7 +583,7 @@ def tuned_int_mm(mat1, mat2, *, layout=None):
     kernel_inputs = MMKernelInputs([mat1, mat2], out_dtype=torch.int32)
 
     # Collect all templates for unified call
-    templates_to_use: list[Union[ExternKernelChoice, KernelTemplate]] = []
+    templates_to_use: list[ExternKernelChoice | KernelTemplate] = []
     if use_aten_gemm_kernels():
         templates_to_use.append(aten__int_mm)
 
@@ -595,7 +602,8 @@ def tuned_int_mm(mat1, mat2, *, layout=None):
             choices, layout, kernel_inputs.nodes(), fuseable=True, non_fuseable=True
         )
 
-    return autotune_select_algorithm(name, choices, kernel_inputs.nodes(), layout)
+    node, _ = autotune_select_algorithm(name, choices, kernel_inputs.nodes(), layout)
+    return node
 
 
 @register_lowering(aten.addmm, type_promotion_kind=None)
@@ -625,6 +633,10 @@ def tuned_addmm(inp, mat1, mat2, *, alpha=1, beta=1, layout=None):
     kernel_inputs = MMKernelInputs(
         [inp_expanded, mat1, mat2], scalars=dict(alpha=alpha, beta=beta)
     )
+    kernel_inputs_aten = MMKernelInputs(
+        [inp, mat1, mat2], scalars=dict(alpha=alpha, beta=beta)
+    )
+
     choices: list[ChoiceCaller] = []
 
     # below is for getting an overview logging info of inductor mms
@@ -641,33 +653,47 @@ def tuned_addmm(inp, mat1, mat2, *, alpha=1, beta=1, layout=None):
     if (not is_nonzero) or (
         not (inductor_config.max_autotune or inductor_config.max_autotune_gemm)
     ):
-        # TODO(coconutruben): combine this with the main flow of addmm through
-        # a subgraph or something as inp vs inp_expanded causes some slight numeric
-        # differences
-        kernel_inputs = MMKernelInputs(
-            [inp, mat1, mat2], scalars=dict(alpha=alpha, beta=beta)
-        )
         choices.extend(
             V.choices.get_template_configs(
-                kernel_inputs,
+                kernel_inputs_aten,
                 [aten_addmm],
                 name,
             )
         )
-        return autotune_select_algorithm(name, choices, kernel_inputs.nodes(), layout)
+        node, _ = autotune_select_algorithm(
+            name, choices, kernel_inputs.nodes(), layout
+        )
+        return node
+
+    templates_to_use: list[ExternKernelChoice | KernelTemplate] = []
 
-    # Collect all templates for unified call
-    templates_to_use: list[Union[ExternKernelChoice, KernelTemplate]] = []
     if use_aten_gemm_kernels():
-        templates_to_use.extend([aten_bias_addmm, aten_addmm])
+        aten_templates: list[ExternKernelChoice | KernelTemplate] = [aten_addmm]
+        if (
+            inp.get_stride()[0] == 0
+            and len(inp.get_size()) == 2
+            and inductor_config.triton.autotune_cublasLt
+            and not V.graph.cpp_wrapper  # bias_addmm only has a Python implementation
+        ):
+            aten_templates.append(aten_bias_addmm)
+
+        # On ROCm, ATen choices use original bias input; non-ROCm keeps unified inputs.
+        choices.extend(
+            V.choices.get_template_configs(kernel_inputs_aten, aten_templates, name)
+        )
 
     if is_nonzero and use_triton_template(layout, check_max_autotune=False):
         templates_to_use.append(mm_template)
 
-        if use_triton_blackwell_tma_template(mat1, mat2, output_layout=layout):
+        if use_triton_blackwell_tma_template(
+            mat1, mat2, output_layout=layout, add_guards=True
+        ):
             templates_to_use.append(blackwell_ws_persistent_device_tma_mm_template)
-        elif use_triton_tma_template(mat1, mat2, output_layout=layout):
-            templates_to_use.append(persistent_tma_mm_template)
+        elif use_triton_tma_template(mat1, mat2, output_layout=layout, add_guards=True):
+            if torch.version.hip is None:
+                templates_to_use.append(persistent_tma_mm_template)
+            else:
+                templates_to_use.append(persistent_mm_template)
 
         templates_to_use.append(addmm_contiguous_subgraph_template)
 
@@ -689,6 +715,7 @@ def tuned_addmm(inp, mat1, mat2, *, alpha=1, beta=1, layout=None):
             kernel_inputs.nodes(reorder=[1, 2, 0]),
             alpha=alpha,
             beta=beta,
+            input_reorder=[2, 0, 1],
         )
 
     if is_nonzero and use_ck_gemm_template(layout, m, n, k):
@@ -713,7 +740,8 @@ def tuned_addmm(inp, mat1, mat2, *, alpha=1, beta=1, layout=None):
             has_bias=True,
         )
 
-    return autotune_select_algorithm(name, choices, kernel_inputs.nodes(), layout)
+    node, _ = autotune_select_algorithm(name, choices, kernel_inputs.nodes(), layout)
+    return node
 
 
 @register_lowering(aten._sparse_semi_structured_mm, type_promotion_kind=None)
@@ -760,9 +788,10 @@ def tuned_sparse_semi_structured_mm(
             choices, layout, [mat1, mat2, mat1_meta], fuseable=True, non_fuseable=True
         )
 
-    return autotune_select_algorithm(
+    node, _ = autotune_select_algorithm(
         "sparse_semi_structured_mm", choices, (mat1, mat1_meta, mat2), layout
     )
+    return node
 
 
 scaling_pairs = [
@@ -920,7 +949,7 @@ def tuned_scaled_mm(
     choices: list[ChoiceCaller] = []
 
     # Collect all templates for unified call
-    templates_to_use: list[Union[ExternKernelChoice, KernelTemplate]] = []
+    templates_to_use: list[ExternKernelChoice | KernelTemplate] = []
     kwarg_overrides = {}
 
     if use_aten_gemm_kernels():
@@ -947,7 +976,10 @@ def tuned_scaled_mm(
 
         # TODO (paulzhan): There is no template that exists for bias and TMA
         # Don't run tma template currently if bias exist
-        if use_triton_tma_template(mat_a, mat_b, output_layout=layout) and not bias:
+        if (
+            use_triton_tma_template(mat_a, mat_b, output_layout=layout, add_guards=True)
+            and not bias
+        ):
             overriders["SCALE_RECIPE_A"] = scale_option_a.value
             overriders["SCALE_RECIPE_B"] = scale_option_b.value
 
@@ -975,7 +1007,9 @@ def tuned_scaled_mm(
                 )
 
         if (
-            use_triton_blackwell_tma_template(mat_a, mat_b, output_layout=layout)
+            use_triton_blackwell_tma_template(
+                mat_a, mat_b, output_layout=layout, add_guards=True
+            )
             and not bias
         ):
             templates_to_use.append(blackwell_ws_persistent_device_tma_mm_template)
@@ -1012,7 +1046,8 @@ def tuned_scaled_mm(
 
     # Early return for MX variants
     if scale_a.dtype != torch.float32:
-        return autotune_select_algorithm(name, choices, input_nodes, layout)
+        node, _ = autotune_select_algorithm(name, choices, input_nodes, layout)
+        return node
 
     if (
         is_nonzero
@@ -1029,11 +1064,12 @@ def tuned_scaled_mm(
     if is_nonzero and use_ck_gemm_template(layout, m, n, k):
         CKGemmTemplate.add_ck_gemm_choices(choices, layout, kernel_inputs.nodes())
 
-    return autotune_select_algorithm(name, choices, kernel_inputs.nodes(), layout)
+    node, _ = autotune_select_algorithm(name, choices, kernel_inputs.nodes(), layout)
+    return node
 
 
 @functools.cache
-def _is_sm7x_or_older_gpu(index: Optional[int]) -> bool:
+def _is_sm7x_or_older_gpu(index: int | None) -> bool:
     props = torch.cuda.get_device_properties(index or 0)
     return props.major <= 7
 
@@ -1053,7 +1089,7 @@ def mm_autoheuristic(
     input_nodes,
     ops,
     precondition,
-    top_k: Optional[int] = None,
+    top_k: int | None = None,
     always_included=None,
 ):
     m, n, k = get_size_hints(mat1, mat2, m, n, k)
diff --git a/torch/_inductor/kernel/mm_grouped.py b/torch/_inductor/kernel/mm_grouped.py
index c607c5e21f0f9..3b35127745ea3 100644
--- a/torch/_inductor/kernel/mm_grouped.py
+++ b/torch/_inductor/kernel/mm_grouped.py
@@ -1,7 +1,7 @@
 # mypy: allow-untyped-defs
 import logging
 from dataclasses import asdict, dataclass
-from typing import Any, Optional
+from typing import Any
 
 import torch
 from torch._dynamo.utils import counters
@@ -146,7 +146,7 @@ def early_config_prune(g, m, dtsize, configs, named_args):
 def grouped_mm_args(
     mat1: TensorBox,
     mat2: TensorBox,
-    offs: Optional[TensorBox],
+    offs: TensorBox | None,
     layout=None,
     out_dtype=None,
 ):
@@ -216,9 +216,9 @@ def grouped_mm_args(
 def can_use_triton_kernel(
     mat_a: TensorBox,
     mat_b: TensorBox,
-    offs: Optional[TensorBox],
-    bias: Optional[TensorBox],
-    scale_result: Optional[TensorBox],
+    offs: TensorBox | None,
+    bias: TensorBox | None,
+    scale_result: TensorBox | None,
 ) -> bool:
     if not (
         torch.cuda.is_available()
@@ -271,14 +271,14 @@ def _tuned_grouped_mm_common(
     kernel_template: TritonTemplate,
     mat_a: TensorBox,
     mat_b: TensorBox,
-    scale_a: Optional[TensorBox] = None,
-    scale_b: Optional[TensorBox] = None,
-    offs: Optional[TensorBox] = None,
-    bias: Optional[TensorBox] = None,
-    scale_result: Optional[TensorBox] = None,
-    out_dtype: Optional[torch.dtype] = None,
-    use_fast_accum: Optional[bool] = None,
-    layout: Optional[Layout] = None,
+    scale_a: TensorBox | None = None,
+    scale_b: TensorBox | None = None,
+    offs: TensorBox | None = None,
+    bias: TensorBox | None = None,
+    scale_result: TensorBox | None = None,
+    out_dtype: torch.dtype | None = None,
+    use_fast_accum: bool | None = None,
+    layout: Layout | None = None,
 ) -> TensorBox:
     assert (scale_a is None) == (scale_b is None)
     assert scale_result is None or scale_a is not None
@@ -448,19 +448,20 @@ def _tuned_grouped_mm_common(
         input_gen_fns[input_offs_idx] = lambda x: create_offsets(
             x, a_is_2d, b_is_2d, m, n, k, alignment
         )
-    return autotune_select_algorithm(
+    node, _ = autotune_select_algorithm(
         algorithm_name, choices, input_nodes, layout, input_gen_fns=input_gen_fns
     )
+    return node
 
 
 @register_lowering(aten._grouped_mm.default, type_promotion_kind=None)
 def tuned_grouped_mm(
     mat_a: TensorBox,
     mat_b: TensorBox,
-    offs: Optional[TensorBox] = None,
-    bias: Optional[TensorBox] = None,
-    out_dtype: Optional[torch.dtype] = None,
-    layout: Optional[Layout] = None,
+    offs: TensorBox | None = None,
+    bias: TensorBox | None = None,
+    out_dtype: torch.dtype | None = None,
+    layout: Layout | None = None,
 ) -> TensorBox:
     """Auto-tuning for _grouped_mm() operator."""
 
@@ -488,12 +489,12 @@ def tuned_scaled_grouped_mm(
     mat_b: TensorBox,
     scale_a: TensorBox,
     scale_b: TensorBox,
-    offs: Optional[TensorBox] = None,
-    bias: Optional[TensorBox] = None,
-    scale_result: Optional[TensorBox] = None,
-    out_dtype: Optional[torch.dtype] = None,
+    offs: TensorBox | None = None,
+    bias: TensorBox | None = None,
+    scale_result: TensorBox | None = None,
+    out_dtype: torch.dtype | None = None,
     use_fast_accum: bool = False,
-    layout: Optional[Layout] = None,
+    layout: Layout | None = None,
 ) -> TensorBox:
     """Auto-tuning for _scaled_grouped_mm() operator."""
 
diff --git a/torch/_inductor/kernel/mm_plus_mm.py b/torch/_inductor/kernel/mm_plus_mm.py
index aef8dfb2168f4..4bbd909d4702f 100644
--- a/torch/_inductor/kernel/mm_plus_mm.py
+++ b/torch/_inductor/kernel/mm_plus_mm.py
@@ -1,7 +1,7 @@
 # mypy: allow-untyped-defs
 
 import logging
-from typing import TYPE_CHECKING, Union
+from typing import TYPE_CHECKING
 
 import torch
 
@@ -160,7 +160,7 @@ def tuned_mm_plus_mm(mat1, mat2, mat3, mat4, *, layout=None):
     choices: list[ChoiceCaller] = []
 
     # Collect all templates for unified call
-    templates_to_use: list[Union[ExternKernelChoice, KernelTemplate]] = []
+    templates_to_use: list[ExternKernelChoice | KernelTemplate] = []
     if use_aten_gemm_kernels():
         templates_to_use.append(aten_mm_plus_mm)
 
@@ -172,6 +172,7 @@ def tuned_mm_plus_mm(mat1, mat2, mat3, mat4, *, layout=None):
         V.choices.get_template_configs(kernel_inputs, templates_to_use, "mm_plus_mm")
     )
 
-    return autotune_select_algorithm(
+    node, _ = autotune_select_algorithm(
         "mm_plus_mm", choices, kernel_inputs.nodes(), layout1
     )
+    return node
diff --git a/torch/_inductor/kernel/templates/cutedsl_mm_grouped.py.jinja b/torch/_inductor/kernel/templates/cutedsl_mm_grouped.py.jinja
index 989f297c5f80f..02285013f165d 100644
--- a/torch/_inductor/kernel/templates/cutedsl_mm_grouped.py.jinja
+++ b/torch/_inductor/kernel/templates/cutedsl_mm_grouped.py.jinja
@@ -3,7 +3,7 @@ from torch._inductor.runtime.runtime_utils import ceildiv
 from cutlass.utils import TensorMapUpdateMode
 {{gen_defines()}}
 # ---- Import GroupedGemm implementation, copied on PyTorch build from Cutlass repository: cutlass/examples/python/CuTeDSL/blackwell/grouped_gemm.py ----
-from torch._inductor.kernel.vendored_templates.cutedsl_grouped_gemm import (
+from torch._inductor.kernel.vendored_templates.cutedsl.kernels.cutedsl_grouped_gemm import (
     GroupedGemmKernel,
 )
 
diff --git a/torch/_inductor/kernel/templates/triton_blackwell_ws_persistent_device_tma_mm.py.jinja b/torch/_inductor/kernel/templates/triton_blackwell_ws_persistent_device_tma_mm.py.jinja
index b5d915129b580..f0a363725ae13 100644
--- a/torch/_inductor/kernel/templates/triton_blackwell_ws_persistent_device_tma_mm.py.jinja
+++ b/torch/_inductor/kernel/templates/triton_blackwell_ws_persistent_device_tma_mm.py.jinja
@@ -43,17 +43,24 @@
         )
         offs_am = pid_m * BLOCK_M
         offs_bn = pid_n * BLOCK_N
+        offs_am_desc = offs_am.to(tl.int32)
+        offs_bn_desc = offs_bn.to(tl.int32)
 
         accumulator = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
         for ki in range(k_tiles):
             offs_k = ki * BLOCK_K
+            offs_k_desc = offs_k.to(tl.int32)
             a = tl.load_tensor_descriptor(
                 a_desc,
-                [offs_am, offs_k] if A_ROW_MAJOR else [offs_k, offs_am],
+                [offs_am_desc, offs_k_desc]
+                if A_ROW_MAJOR
+                else [offs_k_desc, offs_am_desc],
             )
             b = tl.load_tensor_descriptor(
                 b_desc,
-                [offs_k, offs_bn] if B_ROW_MAJOR else [offs_bn, offs_k],
+                [offs_k_desc, offs_bn_desc]
+                if B_ROW_MAJOR
+                else [offs_bn_desc, offs_k_desc],
             )
             accumulator += tl.dot(
                 a if A_ROW_MAJOR else a.T,
diff --git a/torch/_inductor/kernel/templates/triton_bmm.py.jinja b/torch/_inductor/kernel/templates/triton_bmm.py.jinja
index c7e3e574e8de7..94cec7f806301 100644
--- a/torch/_inductor/kernel/templates/triton_bmm.py.jinja
+++ b/torch/_inductor/kernel/templates/triton_bmm.py.jinja
@@ -2,6 +2,7 @@
     M = {{size("A", -2)}}
     N = {{size("B", -1)}}
     K = {{size("A", -1)}}
+    BATCH = {{size("A", 0)}}
 
     stride_aq = {{stride("A", 0)}}
     stride_am = {{stride("A", 1)}}
@@ -25,22 +26,26 @@
     tl.assume(pid_m >= 0)
     tl.assume(pid_n >= 0)
 
-    rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
-    rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
-    if (stride_am == 1 and stride_ak == M) or (stride_am == K and stride_ak == 1):
+    rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M).to(INDEX_DTYPE)
+    rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N).to(INDEX_DTYPE)
+    if ((stride_am == 1 and stride_ak == M) or (stride_am == K and stride_ak == 1)) and (M >= BLOCK_M and K > 1):
         ram = tl.max_contiguous(tl.multiple_of(rm % M, BLOCK_M), BLOCK_M)
     else:
         ram = rm % M
-    if (stride_bk == 1 and stride_bn == K) or (stride_bk == N and stride_bn == 1):
+    if ((stride_bk == 1 and stride_bn == K) or (stride_bk == N and stride_bn == 1)) and (N >= BLOCK_N and K > 1):
         rbn = tl.max_contiguous(tl.multiple_of(rn % N, BLOCK_N), BLOCK_N)
     else:
         rbn = rn % N
 
-    rk = tl.arange(0, BLOCK_K)
+    rk = tl.arange(0, BLOCK_K).to(INDEX_DTYPE)
 
-    idx_q = tl.program_id(1).to(INDEX_DTYPE)  # batch dimension for BMM
-    A = A + (ram[:, None] * stride_am + rk[None, :] * stride_ak + idx_q*stride_aq)
-    B = B + (rk[:, None] * stride_bk + rbn[None, :] * stride_bn + idx_q*stride_bq)
+    # Reconstruct batch index from grid_y/grid_z split (handles batch > 65535)
+    idx_q = (tl.program_id(1) + tl.program_id(2) * tl.num_programs(1)).to(INDEX_DTYPE)
+    # Clamp to valid range for safe pointer arithmetic; out-of-bounds CTAs are
+    # masked off at the store below.
+    idx_q_clamped = tl.minimum(idx_q, BATCH - 1)
+    A = A + (ram[:, None] * stride_am + rk[None, :] * stride_ak + idx_q_clamped*stride_aq)
+    B = B + (rk[:, None] * stride_bk + rbn[None, :] * stride_bn + idx_q_clamped*stride_bq)
 
     acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=ACC_TYPE)
     for k in range(K, 0, -BLOCK_K):
@@ -55,12 +60,14 @@
         B += BLOCK_K * stride_bk
 
     # rematerialize rm and rn to save registers
-    rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
-    rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
-    idx_q = tl.program_id(1).to(INDEX_DTYPE)  # batch dimension for BMM
+    rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M).to(INDEX_DTYPE)
+    rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N).to(INDEX_DTYPE)
+    idx_q = (tl.program_id(1) + tl.program_id(2) * tl.num_programs(1)).to(INDEX_DTYPE)
     idx_m = rm[:, None]
     idx_n = rn[None, :]
-    mask = (idx_m < M) & (idx_n < N)
+    mask = (idx_m < M) & (idx_n < N) & (idx_q < BATCH)
 
+    # cast accumulator to output dtype
+    acc = acc.to(OUT_DTYPE)
     # inductor generates a suffix
     {{store_output(("idx_q", "idx_m", "idx_n"), "acc", "mask", val_shape=("BLOCK_M", "BLOCK_N"))}}
diff --git a/torch/_inductor/kernel/templates/triton_depthwise_conv.py.jinja b/torch/_inductor/kernel/templates/triton_depthwise_conv.py.jinja
new file mode 100644
index 0000000000000..312177c49dda3
--- /dev/null
+++ b/torch/_inductor/kernel/templates/triton_depthwise_conv.py.jinja
@@ -0,0 +1,64 @@
+{{def_kernel("X", "W")}}
+    # Depthwise conv1d - channels-last (NLC) layout
+    # Each group has exactly 1 input channel (groups == in_channels == out_channels)
+    # 3D tiling: BLOCK_N x BLOCK_L x BLOCK_C
+    # Matches the hand-written NLC kernel from depthwise_conv1d_benchmark.py
+
+    BATCH = {{size("X", 0)}}
+    CHANNELS = {{size("X", 1)}}
+    IN_L = {{size("X", 2)}}
+    OUT_L = {{size(None, 2)}}
+
+    stride_xn = {{stride("X", 0)}}
+    stride_xc = {{stride("X", 1)}}
+    stride_xl = {{stride("X", 2)}}
+
+    stride_wc = {{stride("W", 0)}}
+    stride_wk = {{stride("W", 2)}}
+
+    # Grid: (cdiv(BATCH, BLOCK_N), cdiv(OUT_L, BLOCK_L), cdiv(CHANNELS, BLOCK_C))
+    pid_n = tl.program_id(0).to(INDEX_DTYPE)
+    pid_l = tl.program_id(1).to(INDEX_DTYPE)
+    pid_c = tl.program_id(2).to(INDEX_DTYPE)
+
+    n_start = pid_n * BLOCK_N
+    l_start = pid_l * BLOCK_L
+    c_start = pid_c * BLOCK_C
+
+    n_offs = n_start + tl.arange(0, BLOCK_N)
+    l_offs = l_start + tl.arange(0, BLOCK_L)
+    c_offs = c_start + tl.arange(0, BLOCK_C)
+    n_mask = n_offs < BATCH
+    l_mask = l_offs < OUT_L
+    c_mask = c_offs < CHANNELS
+
+    n3 = n_offs[:, None, None]
+    nm3 = n_mask[:, None, None]
+    c3 = c_offs[None, None, :]
+    cm3 = c_mask[None, None, :]
+
+    in_base = n3 * stride_xn + c3 * stride_xc
+
+    acc = tl.zeros((BLOCK_N, BLOCK_L, BLOCK_C), dtype=tl.float32)
+
+    for k in range(KERNEL_SIZE):
+        wk = tl.load(
+            W + c_offs * stride_wc + k * stride_wk,
+            mask=c_mask,
+            other=0.0,
+        ).to(tl.float32)
+        l_in = l_offs * CONV_STRIDE - PADDING + k
+        mask_in = (l_in >= 0) & (l_in < IN_L)
+        x_vals = tl.load(
+            X + in_base + l_in[None, :, None] * stride_xl,
+            mask=nm3 & mask_in[None, :, None] & cm3,
+            other=0.0,
+        ).to(tl.float32)
+        acc += x_vals * wk[None, None, :]
+
+    mask = nm3 & l_mask[None, :, None] & cm3
+    idx_n = n_offs[:, None, None]
+    idx_c = c_offs[None, None, :]
+    idx_l = l_offs[None, :, None]
+
+    {{store_output(("idx_n", "idx_c", "idx_l"), "acc", "mask", val_shape=("BLOCK_N", "BLOCK_L", "BLOCK_C"))}}
diff --git a/torch/_inductor/kernel/templates/triton_epilogue_scaled_mm.py.jinja b/torch/_inductor/kernel/templates/triton_epilogue_scaled_mm.py.jinja
index 56ef18b7a91e3..cd3a6bdd4913a 100644
--- a/torch/_inductor/kernel/templates/triton_epilogue_scaled_mm.py.jinja
+++ b/torch/_inductor/kernel/templates/triton_epilogue_scaled_mm.py.jinja
@@ -98,17 +98,26 @@
             offs_bn = pid_n * BLOCK_N
 
         offs_k = ki * BLOCK_K
+        offs_am_desc = offs_am.to(tl.int32)
+        offs_bn_desc = offs_bn.to(tl.int32)
+        offs_k_desc = offs_k.to(tl.int32)
 
         {%- if TMA_EXPERIMENTAL_API %}
         a = tl._experimental_descriptor_load(
-            a_desc_ptr, [offs_am, offs_k], [BLOCK_M, BLOCK_K],  A.dtype.element_ty
+            a_desc_ptr,
+            [offs_am_desc, offs_k_desc],
+            [BLOCK_M, BLOCK_K],
+            A.dtype.element_ty,
         )
         b = tl._experimental_descriptor_load(
-            b_desc_ptr, [offs_bn, offs_k], [BLOCK_N, BLOCK_K],  B.dtype.element_ty
+            b_desc_ptr,
+            [offs_bn_desc, offs_k_desc],
+            [BLOCK_N, BLOCK_K],
+            B.dtype.element_ty,
         )
         {%- else %}
-        a = tl.load_tensor_descriptor(a_desc, [offs_am, offs_k])
-        b = tl.load_tensor_descriptor(b_desc, [offs_bn, offs_k])
+        a = tl.load_tensor_descriptor(a_desc, [offs_am_desc, offs_k_desc])
+        b = tl.load_tensor_descriptor(b_desc, [offs_bn_desc, offs_k_desc])
         {%- endif %}
         if USE_FAST_ACCUM:
             accumulator = tl.dot(a, b.T, accumulator)
diff --git a/torch/_inductor/kernel/templates/triton_main_loop_scaled_mm.py.jinja b/torch/_inductor/kernel/templates/triton_main_loop_scaled_mm.py.jinja
index 91f73e8742ac6..0c2e51a1ec2f2 100644
--- a/torch/_inductor/kernel/templates/triton_main_loop_scaled_mm.py.jinja
+++ b/torch/_inductor/kernel/templates/triton_main_loop_scaled_mm.py.jinja
@@ -45,13 +45,16 @@
 
         offs_am = pid_m * BLOCK_M
         offs_bn = pid_n * BLOCK_N
+        offs_am_desc = offs_am.to(tl.int32)
+        offs_bn_desc = offs_bn.to(tl.int32)
 
         accumulator = tl.zeros((BLOCK_M, BLOCK_N), dtype=ACC_TYPE)
         for ki in range(0, k_tiles):
             offs_k = ki * BLOCK_K
+            offs_k_desc = offs_k.to(tl.int32)
 
-            a = tl.load_tensor_descriptor(a_desc, [offs_am, offs_k])
-            b = tl.load_tensor_descriptor(b_desc, [offs_bn, offs_k])
+            a = tl.load_tensor_descriptor(a_desc, [offs_am_desc, offs_k_desc])
+            b = tl.load_tensor_descriptor(b_desc, [offs_bn_desc, offs_k_desc])
 
             {%- if SCALE_RECIPE_A == 5 %}  # ScalingType.Blockwise128x128
             scale_a_block = blockwise128x128_scaling(
@@ -80,7 +83,7 @@
             )
             {%- endif %}
 
-            {%- if SCALE_RECIPE_A == 5 %}  # ScalingType.Blockwise128x128
+            {%- if SCALE_RECIPE_B == 5 %}  # ScalingType.Blockwise128x128
             scale_b_block = blockwise128x128_scaling(
                 pid_n,
                 B_inverse_scale,
diff --git a/torch/_inductor/kernel/templates/triton_mm.py.jinja b/torch/_inductor/kernel/templates/triton_mm.py.jinja
index 2da348f3e767c..fcb22c94eada1 100644
--- a/torch/_inductor/kernel/templates/triton_mm.py.jinja
+++ b/torch/_inductor/kernel/templates/triton_mm.py.jinja
@@ -24,8 +24,8 @@
     tl.assume(pid_m >= 0)
     tl.assume(pid_n >= 0)
 
-    rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
-    rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
+    rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M).to(INDEX_DTYPE)
+    rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N).to(INDEX_DTYPE)
     if ((stride_am == 1 and stride_ak == M) or (stride_am == K and stride_ak == 1)) and (M >= BLOCK_M and K > 1):
         offs_a_m = tl.max_contiguous(tl.multiple_of(rm % M, BLOCK_M), BLOCK_M)
     else:
@@ -34,7 +34,7 @@
         offs_b_n = tl.max_contiguous(tl.multiple_of(rn % N, BLOCK_N), BLOCK_N)
     else:
         offs_b_n = rn % N
-    offs_k = tl.arange(0, BLOCK_K)
+    offs_k = tl.arange(0, BLOCK_K).to(INDEX_DTYPE)
     acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=ACC_TYPE)
 
     for k_idx in range(0, tl.cdiv(K, BLOCK_K)):
@@ -62,8 +62,8 @@
         {% endif %}
 
     # rematerialize rm and rn to save registers
-    rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
-    rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
+    rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M).to(INDEX_DTYPE)
+    rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N).to(INDEX_DTYPE)
     idx_m = rm[:, None]
     idx_n = rn[None, :]
     mask = (idx_m < M) & (idx_n < N)
diff --git a/torch/_inductor/kernel/templates/triton_mm_grouped.py.jinja b/torch/_inductor/kernel/templates/triton_mm_grouped.py.jinja
index 46c40d99aa502..52049582bf21d 100644
--- a/torch/_inductor/kernel/templates/triton_mm_grouped.py.jinja
+++ b/torch/_inductor/kernel/templates/triton_mm_grouped.py.jinja
@@ -294,10 +294,10 @@ def do_mma(a, b, accumulator):
 {%- endif %}
                     accumulator = do_mma(a, b, accumulator)
 {%- else %}
-                offs_am = tile_m_idx * BLOCK_M + tl.arange(0, BLOCK_M)
-                offs_bn = tile_n_idx * BLOCK_N + tl.arange(0, BLOCK_N)
+                offs_am = tile_m_idx * BLOCK_M + tl.arange(0, BLOCK_M).to(INDEX_DTYPE)
+                offs_bn = tile_n_idx * BLOCK_N + tl.arange(0, BLOCK_N).to(INDEX_DTYPE)
                 for k_block_offset in range(0, k_size, BLOCK_K):
-                    block_offs_k = k_block_offset + tl.arange(0, BLOCK_K)
+                    block_offs_k = k_block_offset + tl.arange(0, BLOCK_K).to(INDEX_DTYPE)
                     offs_k = block_offs_k + k_start_offset
                     a_ptrs = (
                         a_ptr
@@ -328,8 +328,8 @@ def do_mma(a, b, accumulator):
                     b_ptrs += BLOCK_K
 {%- endif %}
 
-                offs_am = tile_m_idx * BLOCK_M + tl.arange(0, BLOCK_M)
-                offs_bn = tile_n_idx * BLOCK_N + tl.arange(0, BLOCK_N)
+                offs_am = tile_m_idx * BLOCK_M + tl.arange(0, BLOCK_M).to(INDEX_DTYPE)
+                offs_bn = tile_n_idx * BLOCK_N + tl.arange(0, BLOCK_N).to(INDEX_DTYPE)
 {%- if SCALED %}
                 scale_a = tl.load(
                     scale_a_ptr
diff --git a/torch/_inductor/kernel/templates/triton_mm_rocm.py.jinja b/torch/_inductor/kernel/templates/triton_mm_rocm.py.jinja
index 42b99c70d5cbd..6ad00438cf66d 100644
--- a/torch/_inductor/kernel/templates/triton_mm_rocm.py.jinja
+++ b/torch/_inductor/kernel/templates/triton_mm_rocm.py.jinja
@@ -24,8 +24,8 @@
     tl.assume(pid_m >= 0)
     tl.assume(pid_n >= 0)
 
-    rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
-    rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
+    rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M).to(INDEX_DTYPE)
+    rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N).to(INDEX_DTYPE)
     if (stride_am == 1 and stride_ak == M) or (stride_am == K and stride_ak == 1):
         offs_a_m = tl.max_contiguous(tl.multiple_of(rm % M, BLOCK_M), BLOCK_M)
     else:
@@ -34,7 +34,7 @@
         offs_b_n = tl.max_contiguous(tl.multiple_of(rn % N, BLOCK_N), BLOCK_N)
     else:
         offs_b_n = rn % N
-    offs_k = tl.arange(0, BLOCK_K)
+    offs_k = tl.arange(0, BLOCK_K).to(INDEX_DTYPE)
     acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=ACC_TYPE)
 
     for k_idx in range(0, tl.cdiv(K, BLOCK_K)):
@@ -61,8 +61,8 @@
         {% endif %}
 
     # rematerialize rm and rn to save registers
-    rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
-    rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
+    rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M).to(INDEX_DTYPE)
+    rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N).to(INDEX_DTYPE)
     idx_m = rm[:, None]
     idx_n = rn[None, :]
     mask = (idx_m < M) & (idx_n < N)
diff --git a/torch/_inductor/kernel/templates/triton_persistent_mm.py.jinja b/torch/_inductor/kernel/templates/triton_persistent_mm.py.jinja
new file mode 100644
index 0000000000000..7450533be5adf
--- /dev/null
+++ b/torch/_inductor/kernel/templates/triton_persistent_mm.py.jinja
@@ -0,0 +1,75 @@
+{{def_kernel("A", "B")}}
+    M = {{size("A", 0)}}
+    N = {{size("B", 1)}}
+    K = {{size("A", 1)}}
+    if M * N == 0:
+        # early exit due to zero-size input(s)
+        return
+    stride_am = {{stride("A", 0)}}
+    stride_ak = {{stride("A", 1)}}
+    stride_bk = {{stride("B", 0)}}
+    stride_bn = {{stride("B", 1)}}
+
+    # persistent kernel: each CTA processes multiple tiles
+    start_pid = tl.program_id(0).to(INDEX_DTYPE)
+    grid_m = tl.cdiv(M, BLOCK_M)
+    grid_n = tl.cdiv(N, BLOCK_N)
+    num_tiles = grid_m * grid_n
+    width = GROUP_M * grid_n
+
+    for tile_id in tl.range(start_pid, num_tiles, NUM_SMS):
+
+        # re-order program ID for better L2 performance
+        group_id = tile_id // width
+        group_size = min(grid_m - group_id * GROUP_M, GROUP_M)
+        pid_m = group_id * GROUP_M + (tile_id % group_size)
+        pid_n = (tile_id % width) // (group_size)
+        tl.assume(pid_m >= 0)
+        tl.assume(pid_n >= 0)
+
+        rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
+        rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
+        if ((stride_am == 1 and stride_ak == M) or (stride_am == K and stride_ak == 1)) and (M >= BLOCK_M and K > 1):
+            offs_a_m = tl.max_contiguous(tl.multiple_of(rm % M, BLOCK_M), BLOCK_M)
+        else:
+            offs_a_m = rm % M
+        if ((stride_bk == 1 and stride_bn == K) or (stride_bk == N and stride_bn == 1)) and (N >= BLOCK_N and K > 1):
+            offs_b_n = tl.max_contiguous(tl.multiple_of(rn % N, BLOCK_N), BLOCK_N)
+        else:
+            offs_b_n = rn % N
+        offs_k = tl.arange(0, BLOCK_K)
+        acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=ACC_TYPE)
+
+        for k_idx in range(0, tl.cdiv(K, BLOCK_K)):
+            {% if not EVEN_K %}
+            a_mask = offs_k[None, :] < (K - k_idx * BLOCK_K)
+            b_mask = offs_k[:, None] < (K - k_idx * BLOCK_K)
+            {% endif %}
+            a_k_idx_vals = offs_k[None, :] + (k_idx * BLOCK_K)
+            b_k_idx_vals = offs_k[:, None] + (k_idx * BLOCK_K)
+
+            idx_m = offs_a_m[:, None]
+            idx_n = a_k_idx_vals
+            {{load_input("A", "a", ("idx_m", "idx_n"), mask=None if EVEN_K else "a_mask",
+                         indent_width=12, index_shape=("BLOCK_M", "BLOCK_K"))}}
+
+            idx_m = b_k_idx_vals
+            idx_n = offs_b_n[None, :]
+            {{load_input("B", "b", ("idx_m", "idx_n"), mask=None if EVEN_K else "b_mask",
+                         indent_width=12, index_shape=("BLOCK_K", "BLOCK_N"))}}
+
+            {% if USE_FAST_ACCUM %}
+            acc = tl.dot(a, b, acc, allow_tf32=ALLOW_TF32, out_dtype=ACC_TYPE)
+            {% else %}
+            acc += tl.dot(a, b, allow_tf32=ALLOW_TF32, out_dtype=ACC_TYPE)
+            {% endif %}
+
+        # rematerialize rm and rn to save registers
+        rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
+        rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
+        idx_m = rm[:, None]
+        idx_n = rn[None, :]
+        mask = (idx_m < M) & (idx_n < N)
+
+        # inductor generates a suffix
+        {{store_output(("idx_m", "idx_n"), "acc", "mask", indent_width=8, val_shape=("BLOCK_M", "BLOCK_N"))}}
diff --git a/torch/_inductor/kernel/templates/triton_persistent_tma_mm.py.jinja b/torch/_inductor/kernel/templates/triton_persistent_tma_mm.py.jinja
index b27ef29478d8b..cb6b4a1440812 100644
--- a/torch/_inductor/kernel/templates/triton_persistent_tma_mm.py.jinja
+++ b/torch/_inductor/kernel/templates/triton_persistent_tma_mm.py.jinja
@@ -66,31 +66,34 @@
 
         rm = pid_m * BLOCK_M
         rn = pid_n * BLOCK_N
+        rm_desc = rm.to(tl.int32)
+        rn_desc = rn.to(tl.int32)
         acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=ACC_TYPE)
 
         for rk in tl.range(0, K, BLOCK_K):
+            rk_desc = rk.to(tl.int32)
 
             {%- if TMA_EXPERIMENTAL_API %}
             a = tl._experimental_descriptor_load(
                 a_desc_ptr,
-                [rm, rk] if A_ROW_MAJOR else [rk, rm],
+                [rm_desc, rk_desc] if A_ROW_MAJOR else [rk_desc, rm_desc],
                 [BLOCK_M, BLOCK_K] if A_ROW_MAJOR else [BLOCK_K, BLOCK_M],
                 A.dtype.element_ty,
             )
             b = tl._experimental_descriptor_load(
                 b_desc_ptr,
-                [rk, rn] if B_ROW_MAJOR else [rn, rk],
+                [rk_desc, rn_desc] if B_ROW_MAJOR else [rn_desc, rk_desc],
                 [BLOCK_K, BLOCK_N] if B_ROW_MAJOR else [BLOCK_N, BLOCK_K],
                 B.dtype.element_ty,
             )
             {%- else %}
             a = tl.load_tensor_descriptor(
                 a_desc,
-                [rm, rk] if A_ROW_MAJOR else [rk, rm],
+                [rm_desc, rk_desc] if A_ROW_MAJOR else [rk_desc, rm_desc],
             )
             b = tl.load_tensor_descriptor(
                 b_desc,
-                [rk, rn] if B_ROW_MAJOR else [rn, rk],
+                [rk_desc, rn_desc] if B_ROW_MAJOR else [rn_desc, rk_desc],
             )
             {%- endif %}
             if USE_FAST_ACCUM:
diff --git a/test/dynamo_expected_failures/CPython313-test_exception_variations-ExceptStarTestCases.test_try_except_else_finally b/torch/_inductor/kernel/vendored_templates/__init__.py
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_exception_variations-ExceptStarTestCases.test_try_except_else_finally
rename to torch/_inductor/kernel/vendored_templates/__init__.py
diff --git a/torch/_inductor/kernel/vendored_templates/cutedsl/__init__.py b/torch/_inductor/kernel/vendored_templates/cutedsl/__init__.py
new file mode 100644
index 0000000000000..d7dc51ed6f633
--- /dev/null
+++ b/torch/_inductor/kernel/vendored_templates/cutedsl/__init__.py
@@ -0,0 +1 @@
+# CuTeDSL vendored kernels and wrappers for PyTorch Inductor.
diff --git a/torch/_inductor/kernel/vendored_templates/cutedsl/dense_blockscaled_gemm_persistent.py b/torch/_inductor/kernel/vendored_templates/cutedsl/dense_blockscaled_gemm_persistent.py
new file mode 100644
index 0000000000000..91188fb97e107
--- /dev/null
+++ b/torch/_inductor/kernel/vendored_templates/cutedsl/dense_blockscaled_gemm_persistent.py
@@ -0,0 +1,2156 @@
+# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-3-Clause
+
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions are met:
+
+# 1. Redistributions of source code must retain the above copyright notice, this
+# list of conditions and the following disclaimer.
+
+# 2. Redistributions in binary form must reproduce the above copyright notice,
+# this list of conditions and the following disclaimer in the documentation
+# and/or other materials provided with the distribution.
+
+# 3. Neither the name of the copyright holder nor the names of its
+# contributors may be used to endorse or promote products derived from
+# this software without specific prior written permission.
+
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
+import argparse
+from typing import Tuple, Type, Union
+
+import cuda.bindings.driver as cuda
+
+import cutlass
+import cutlass.cute as cute
+import cutlass.pipeline as pipeline
+import cutlass.torch as cutlass_torch
+import cutlass.utils as utils
+import cutlass.utils.blackwell_helpers as sm100_utils
+import cutlass.utils.blockscaled_layout as blockscaled_utils
+from cutlass.cute.nvgpu import cpasync, tcgen05
+from cutlass.cute.runtime import from_dlpack
+from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait
+
+import torch
+
+
+"""
+This example provides an experimental implementation of the SM100 batched dense blockscaled GEMM kernel, please note that the APIs and implementation details related to this kernel may change in future releases.
+
+A high-performance persistent batched dense blockscaled GEMM example for the NVIDIA Blackwell SM100 architecture
+using CUTE DSL.
+- Matrix A is MxKxL, L is batch dimension, A can be row-major("K") or column-major("M") for MXF8 input type and can only be row-major("K") for MXF4/NVF4 input type
+- Matrix B is NxKxL, L is batch dimension, B can be row-major("N") or column-major("K") for MXF8 input type and can only be row-major("K") for MXF4/NVF4 input type
+- Matrix C is MxNxL, L is batch dimension, C can be row-major("N") or column-major("M")
+- Matrix SFA layout is filled internally according to A shape and BlockScaledBasicChunk, which has M×ceil_div(K, sf_vec_size)×L elements respectively
+- Matrix SFB layout is filled internally according to B shape and BlockScaledBasicChunk, which has N×ceil_div(K, sf_vec_size)×L elements respectively
+
+This GEMM kernel supports the following features:
+    - Utilizes Tensor Memory Access (TMA) for efficient memory operations
+    - Utilizes Blackwell's tcgen05.mma for matrix multiply-accumulate (MMA) operations (including 2cta mma instructions)
+    - Implements TMA multicast with cluster to reduce L2 memory traffic
+    - Support persistent tile scheduling to better overlap memory load/store with mma between tiles
+    - Support warp specialization to avoid explicit pipelining between mainloop load and mma
+
+This GEMM works as follows:
+1. DMA warp: Load A and B matrices from global memory (GMEM) to shared memory (SMEM) using TMA operations.
+2. MMA warp:
+    - Load scale factor A/B from shared memory (SMEM) to tensor memory (TMEM) using tcgen05.cp instruction.
+    - Perform matrix multiply-accumulate (MMA) operations using tcgen05.mma instruction.
+3. EPILOGUE warp:
+    - Load completed accumulator from tensor memory (TMEM) to registers (RMEM) using tcgen05.ld.
+    - Type convert C matrix to output type.
+    - Optionally store C matrix from registers (RMEM) to shared memory (SMEM) to global memory (GMEM) with TMA operations,
+      or directly store C matrix from registers (RMEM) to global memory (GMEM) without TMA operations.
+    - Optionally accept an elementwise lambda function epilogue_op to apply to the output tensor:
+      e.g., relu can set epilogue_op = lambda x: cute.where(x > 0, x, cute.full_like(x, 0))
+
+SM100 tcgen05.mma.kind.block_scale instructions operate as follows:
+- Read matrix A from SMEM
+- Read matrix B from SMEM
+- Read scalefactor A from TMEM
+- Read scalefactor B from TMEM
+- Write accumulator to TMEM
+The accumulator in TMEM must then be loaded to registers before writing back to GMEM.
+
+Input arguments to this example is shown below:
+
+.. code-block:: bash
+
+    python examples/blackwell/dense_blockscaled_gemm_persistent.py             \
+      --ab_dtype Float4E2M1FN --sf_dtype Float8E8M0FNU --sf_vec_size 16        \
+      --c_dtype Float16                                                        \
+      --mma_tiler_mn 256,128 --cluster_shape_mn 2,1                            \
+      --mnkl 8192,8192,1024,1
+
+To collect performance with NCU profiler:
+
+.. code-block:: bash
+
+    ncu python examples/blackwell/dense_blockscaled_gemm_persistent.py         \
+      --ab_dtype Float4E2M1FN --sf_dtype Float8E8M0FNU --sf_vec_size 16        \
+      --c_dtype Float16                                                        \
+      --mma_tiler_mn 256,128 --cluster_shape_mn 2,1                            \
+      --mnkl 8192,8192,1024,1                                                  \
+      --warmup_iterations 1 --iterations 10 --skip_ref_check
+
+
+Constraints:
+* Supported input data types: mxf8, mxf4, nvf4
+  see detailed valid dtype combinations in below Sm100BlockScaledPersistentDenseGemmKernel class documentation
+* A/B tensor must have the same data type, mixed data type is not supported (e.g., mxf8 x mxf4)
+* Mma tiler M must be 128 or 256(use_2cta_instrs)
+* Mma tiler N must be 64/128/192/256
+* Cluster shape M/N must be positive and power of 2, total cluster size <= 16
+* Cluster shape M must be multiple of 2 if Mma tiler M is 256(use_2cta_instrs)
+* The contiguous dimension of A/B/C tensors must be at least 16 bytes aligned,
+  i.e, number of elements is a multiple of 16 and 32 for Float8 and Float4, respectively.
+"""
+
+
+class Sm100BlockScaledPersistentDenseGemmKernel:
+    """This class implements batched matrix multiplication (C = A x SFA x B x SFB) with support for various data types
+    and architectural features specific to Blackwell GPUs with persistent tile scheduling and warp specialization.
+
+    :param sf_vec_size: Scalefactor vector size.
+    :type sf_vec_size: int
+    :param mma_tiler_mn: Shape of the Matrix Multiply-Accumulate (MMA) tile (M,N)
+    :type mma_tiler_mn: Tuple[int, int]
+    :param cluster_shape_mn: Cluster dimensions (M,N) for parallel processing
+    :type cluster_shape_mn: Tuple[int, int]
+
+    :note: In current version, A and B tensor must have the same data type
+        - i.e., Float8E4M3FN for A and Float8E5M2 for B is not supported
+
+    :note: Supported combinations of A/B data types, SF data typs and SF vector size:
+        - MXF8: A/B: Float8E5M2/Float8E4M3FN + SF: Float8E8M0FNU + sf_vec_size: 32
+        - MXF4: A/B: Float4E2M1FN + SF: Float8E8M0FNU + sf_vec_size: 32
+        - NVF4: A/B: Float4E2M1FN + SF: Float8E8M0FNU/Float8E4M3FN + sf_vec_size: 16
+
+    :note: Supported accumulator data types:
+        - Float32
+
+    :note: Supported C data types:
+        - Float32
+        - Float16/BFloat16
+        - Float8E4M3FN/Float8E5M2
+    :note: Constraints:
+        - MMA tiler M must be 128 or 256 (use_2cta_instrs)
+        - MMA tiler N must be 64/128/192/256
+        - Cluster shape M must be multiple of 2 if Mma tiler M is 256
+        - Cluster shape M/N must be positive and power of 2, total cluster size <= 16
+        - Also, Cluster shape M/N must be <= 4 for scale factor multicasts due to limited size of scale factors
+
+    Example:
+        >>> # xdoctest: +SKIP(requires cutlass)
+        >>> gemm = Sm100BlockScaledPersistentDenseGemmKernel(
+        ...     sf_vec_size=16, mma_tiler_mn=(256, 128), cluster_shape_mn=(2, 1)
+        ... )
+        >>> gemm(
+        ...     a_tensor,
+        ...     b_tensor,
+        ...     sfa_tensor,
+        ...     sfb_tensor,
+        ...     c_tensor,
+        ...     max_active_clusters,
+        ...     stream,
+        ... )
+    """
+
+    def __init__(
+        self,
+        sf_vec_size: int,
+        mma_tiler_mn: Tuple[int, int],
+        cluster_shape_mn: Tuple[int, int],
+    ):
+        """Initializes the configuration for a Blackwell dense GEMM kernel.
+
+        This configuration includes several key aspects:
+
+        1.  MMA Instruction Settings (tcgen05):
+            - acc_dtype: Data types for MMA accumulator, always set to Float32
+            - sf_vec_size: Scalefactor A/B vector size.
+            - mma_tiler_mn: The (M, N) shape of the MMA instruction tiler.
+
+        2.  Cluster Shape:
+            - cluster_shape_mn: The (ClusterM, ClusterN) shape of the CTA cluster.
+
+        :param sf_vec_size: Scalefactor vector size.
+        :type sf_vec_size: int
+        :param mma_tiler_mn: Tuple (M, N) shape of the MMA instruction.
+        :type mma_tiler_mn: Tuple[int, int]
+        :param cluster_shape_mn: Tuple (ClusterM, ClusterN) shape of the cluster.
+        :type cluster_shape_mn: Tuple[int, int]
+        """
+
+        self.acc_dtype = cutlass.Float32
+        self.sf_vec_size = sf_vec_size
+        self.use_2cta_instrs = mma_tiler_mn[0] == 256
+        self.cluster_shape_mn = cluster_shape_mn
+        # K dimension is deferred in _setup_attributes
+        self.mma_tiler = (*mma_tiler_mn, 1)
+
+        self.cta_group = (
+            tcgen05.CtaGroup.TWO if self.use_2cta_instrs else tcgen05.CtaGroup.ONE
+        )
+
+        self.occupancy = 1
+        # Set specialized warp ids
+        self.epilog_warp_id = (
+            0,
+            1,
+            2,
+            3,
+        )
+        self.mma_warp_id = 4
+        self.tma_warp_id = 5
+        self.threads_per_cta = 32 * len(
+            (self.mma_warp_id, self.tma_warp_id, *self.epilog_warp_id)
+        )
+        # Set barrier id for epilogue sync and tmem ptr sync
+        self.epilog_sync_barrier = pipeline.NamedBarrier(
+            barrier_id=1,
+            num_threads=32 * len(self.epilog_warp_id),
+        )
+        self.tmem_alloc_barrier = pipeline.NamedBarrier(
+            barrier_id=2,
+            num_threads=32 * len((self.mma_warp_id, *self.epilog_warp_id)),
+        )
+        self.smem_capacity = utils.get_smem_capacity_in_bytes("sm_100")
+        SM100_TMEM_CAPACITY_COLUMNS = 512
+        self.num_tmem_alloc_cols = SM100_TMEM_CAPACITY_COLUMNS
+
+    def _setup_attributes(self):
+        """Set up configurations that are dependent on GEMM inputs
+
+        This method configures various attributes based on the input tensor properties
+        (data types, leading dimensions) and kernel settings:
+        - Configuring tiled MMA
+        - Computing MMA/cluster/tile shapes
+        - Computing cluster layout
+        - Computing multicast CTAs for A/B/SFA/SFB
+        - Computing epilogue subtile
+        - Setting up A/B/SFA/SFB/C stage counts in shared memory
+        - Computing A/B/SFA/SFB/C shared memory layout
+        """
+        # Compute mma instruction shapes
+        # (MMA_Tile_Shape_M, MMA_Tile_Shape_N, MMA_Inst_Shape_K)
+        self.mma_inst_shape_mn = (
+            self.mma_tiler[0],
+            self.mma_tiler[1],
+        )
+        # (CTA_Tile_Shape_M, Round_Up(MMA_Tile_Shape_N, 128), MMA_Inst_Shape_K)
+        self.mma_inst_shape_mn_sfb = (
+            self.mma_inst_shape_mn[0] // (2 if self.use_2cta_instrs else 1),
+            cute.round_up(self.mma_inst_shape_mn[1], 128),
+        )
+
+        tiled_mma = sm100_utils.make_blockscaled_trivial_tiled_mma(
+            self.a_dtype,
+            self.a_major_mode,
+            self.b_major_mode,
+            self.sf_dtype,
+            self.sf_vec_size,
+            self.cta_group,
+            self.mma_inst_shape_mn,
+        )
+
+        tiled_mma_sfb = sm100_utils.make_blockscaled_trivial_tiled_mma(
+            self.a_dtype,
+            self.a_major_mode,
+            self.b_major_mode,
+            self.sf_dtype,
+            self.sf_vec_size,
+            cute.nvgpu.tcgen05.CtaGroup.ONE,
+            self.mma_inst_shape_mn_sfb,
+        )
+
+        # Compute mma/cluster/tile shapes
+        mma_inst_shape_k = cute.size(tiled_mma.shape_mnk, mode=[2])
+        mma_inst_tile_k = 4
+        self.mma_tiler = (
+            self.mma_inst_shape_mn[0],
+            self.mma_inst_shape_mn[1],
+            mma_inst_shape_k * mma_inst_tile_k,
+        )
+        self.mma_tiler_sfb = (
+            self.mma_inst_shape_mn_sfb[0],
+            self.mma_inst_shape_mn_sfb[1],
+            mma_inst_shape_k * mma_inst_tile_k,
+        )
+        self.cta_tile_shape_mnk = (
+            self.mma_tiler[0] // cute.size(tiled_mma.thr_id.shape),
+            self.mma_tiler[1],
+            self.mma_tiler[2],
+        )
+        self.cta_tile_shape_mnk_sfb = (
+            self.mma_tiler_sfb[0] // cute.size(tiled_mma.thr_id.shape),
+            self.mma_tiler_sfb[1],
+            self.mma_tiler_sfb[2],
+        )
+
+        # Compute cluster layout
+        self.cluster_layout_vmnk = cute.tiled_divide(
+            cute.make_layout((*self.cluster_shape_mn, 1)),
+            (tiled_mma.thr_id.shape,),
+        )
+        self.cluster_layout_sfb_vmnk = cute.tiled_divide(
+            cute.make_layout((*self.cluster_shape_mn, 1)),
+            (tiled_mma_sfb.thr_id.shape,),
+        )
+
+        # Compute number of multicast CTAs for A/B
+        self.num_mcast_ctas_a = cute.size(self.cluster_layout_vmnk.shape[2])
+        self.num_mcast_ctas_b = cute.size(self.cluster_layout_vmnk.shape[1])
+        self.num_mcast_ctas_sfb = cute.size(self.cluster_layout_sfb_vmnk.shape[1])
+        self.is_a_mcast = self.num_mcast_ctas_a > 1
+        self.is_b_mcast = self.num_mcast_ctas_b > 1
+        self.is_sfb_mcast = self.num_mcast_ctas_sfb > 1
+
+        # Compute epilogue subtile
+        self.epi_tile = sm100_utils.compute_epilogue_tile_shape(
+            self.cta_tile_shape_mnk,
+            self.use_2cta_instrs,
+            self.c_layout,
+            self.c_dtype,
+        )
+        self.epi_tile_n = cute.size(self.epi_tile[1])
+
+        # Setup A/B/C stage count in shared memory and ACC stage count in tensor memory
+        self.num_acc_stage, self.num_ab_stage, self.num_c_stage = self._compute_stages(
+            tiled_mma,
+            self.mma_tiler,
+            self.a_dtype,
+            self.b_dtype,
+            self.epi_tile,
+            self.c_dtype,
+            self.c_layout,
+            self.sf_dtype,
+            self.sf_vec_size,
+            self.smem_capacity,
+            self.occupancy,
+        )
+
+        # Compute A/B/SFA/SFB/C shared memory layout
+        self.a_smem_layout_staged = sm100_utils.make_smem_layout_a(
+            tiled_mma,
+            self.mma_tiler,
+            self.a_dtype,
+            self.num_ab_stage,
+        )
+        self.b_smem_layout_staged = sm100_utils.make_smem_layout_b(
+            tiled_mma,
+            self.mma_tiler,
+            self.b_dtype,
+            self.num_ab_stage,
+        )
+        self.sfa_smem_layout_staged = blockscaled_utils.make_smem_layout_sfa(
+            tiled_mma,
+            self.mma_tiler,
+            self.sf_vec_size,
+            self.num_ab_stage,
+        )
+        self.sfb_smem_layout_staged = blockscaled_utils.make_smem_layout_sfb(
+            tiled_mma,
+            self.mma_tiler,
+            self.sf_vec_size,
+            self.num_ab_stage,
+        )
+        self.c_smem_layout_staged = sm100_utils.make_smem_layout_epi(
+            self.c_dtype,
+            self.c_layout,
+            self.epi_tile,
+            self.num_c_stage,
+        )
+
+        # Overlap and double buffer accumulator when num_acc_stage == 1 for cta_tile_n = 256 case
+        self.overlapping_accum = self.num_acc_stage == 1
+
+        # Compute number of TMEM columns for SFA/SFB/Accumulator
+        sf_atom_mn = 32
+        self.num_sfa_tmem_cols = (
+            self.cta_tile_shape_mnk[0] // sf_atom_mn
+        ) * mma_inst_tile_k
+        self.num_sfb_tmem_cols = (
+            self.cta_tile_shape_mnk_sfb[1] // sf_atom_mn
+        ) * mma_inst_tile_k
+        self.num_sf_tmem_cols = self.num_sfa_tmem_cols + self.num_sfb_tmem_cols
+        self.num_accumulator_tmem_cols = (
+            self.cta_tile_shape_mnk[1] * self.num_acc_stage
+            if not self.overlapping_accum
+            else self.cta_tile_shape_mnk[1] * 2 - self.num_sf_tmem_cols
+        )
+
+        # Only when overlapping_accum is enabled, we need to release accumulator buffer early in epilogue
+        self.iter_acc_early_release_in_epilogue = (
+            self.num_sf_tmem_cols // self.epi_tile_n
+        )
+
+    @cute.jit
+    def __call__(
+        self,
+        a_tensor: cute.Tensor,
+        b_tensor: cute.Tensor,
+        sfa_tensor: cute.Tensor,
+        sfb_tensor: cute.Tensor,
+        c_tensor: cute.Tensor,
+        max_active_clusters: cutlass.Constexpr,
+        stream: cuda.CUstream,
+        epilogue_op: cutlass.Constexpr = lambda x: x,
+    ):
+        """Execute the GEMM operation in steps:
+        - Setup static attributes before smem/grid/tma computation
+        - Setup TMA load/store atoms and tensors
+        - Compute grid size with regard to hardware constraints
+        - Define shared storage for kernel
+        - Launch the kernel synchronously
+
+        :param a_tensor: Input tensor A
+        :type a_tensor: cute.Tensor
+        :param b_tensor: Input tensor B
+        :type b_tensor: cute.Tensor
+        :param sfa_tensor: Scale factor tensor A
+        :type sfa_tensor: cute.Tensor
+        :param sfb_tensor: Scale factor tensor B
+        :type sfb_tensor: cute.Tensor
+        :param c_tensor: Output tensor C
+        :type c_tensor: cute.Tensor
+        :param max_active_clusters: Maximum number of active clusters
+        :type max_active_clusters: cutlass.Constexpr
+        :param stream: CUDA stream for asynchronous execution
+        :type stream: cuda.CUstream
+        :param epilogue_op: Optional elementwise lambda function to apply to the output tensor
+        :type epilogue_op: cutlass.Constexpr
+        :raises TypeError: If input data types are incompatible with the MMA instruction.
+        """
+
+        # Convert from torch convention to CuTe convention.
+        # CUTLASS API passes tensors as A:(L,M,K) B:(L,K,N) C:(L,M,N)
+        # but CuTe DSL kernels expect A:(M,K,L) B:(N,K,L) C:(M,N,L).
+        def add_batch_mode(tensor: cute.Tensor) -> cute.Tensor:
+            return cute.make_tensor(
+                tensor.iterator,
+                cute.prepend(tensor.layout, cute.make_layout(1), up_to_rank=3),
+            )
+
+        a_tensor = add_batch_mode(a_tensor)
+        b_tensor = add_batch_mode(b_tensor)
+        c_tensor = add_batch_mode(c_tensor)
+
+        a_tensor = cute.make_tensor(
+            a_tensor.iterator, cute.select(a_tensor.layout, [1, 2, 0])
+        )
+        b_tensor = cute.make_tensor(
+            b_tensor.iterator, cute.select(b_tensor.layout, [2, 1, 0])
+        )
+        c_tensor = cute.make_tensor(
+            c_tensor.iterator, cute.select(c_tensor.layout, [1, 2, 0])
+        )
+
+        # Setup static attributes before smem/grid/tma computation
+        self.a_dtype: Type[cutlass.Numeric] = a_tensor.element_type
+        self.b_dtype: Type[cutlass.Numeric] = b_tensor.element_type
+        self.sf_dtype: Type[cutlass.Numeric] = sfa_tensor.element_type
+        self.c_dtype: Type[cutlass.Numeric] = c_tensor.element_type
+        self.a_major_mode = utils.LayoutEnum.from_tensor(a_tensor).mma_major_mode()
+        self.b_major_mode = utils.LayoutEnum.from_tensor(b_tensor).mma_major_mode()
+        self.c_layout = utils.LayoutEnum.from_tensor(c_tensor)
+
+        # Check if input data types are compatible with MMA instruction
+        if cutlass.const_expr(self.a_dtype != self.b_dtype):
+            raise TypeError(f"Type must match: {self.a_dtype} != {self.b_dtype}")
+
+        # Setup attributes that dependent on gemm inputs
+        self._setup_attributes()
+
+        # Setup sfa/sfb tensor by filling A/B tensor to scale factor atom layout
+        # ((Atom_M, Rest_M),(Atom_K, Rest_K),RestL)
+        sfa_layout = blockscaled_utils.tile_atom_to_shape_SF(
+            a_tensor.shape, self.sf_vec_size
+        )
+        sfa_tensor = cute.make_tensor(sfa_tensor.iterator, sfa_layout)
+
+        # ((Atom_N, Rest_N),(Atom_K, Rest_K),RestL)
+        sfb_layout = blockscaled_utils.tile_atom_to_shape_SF(
+            b_tensor.shape, self.sf_vec_size
+        )
+        sfb_tensor = cute.make_tensor(sfb_tensor.iterator, sfb_layout)
+
+        tiled_mma = sm100_utils.make_blockscaled_trivial_tiled_mma(
+            self.a_dtype,
+            self.a_major_mode,
+            self.b_major_mode,
+            self.sf_dtype,
+            self.sf_vec_size,
+            self.cta_group,
+            self.mma_inst_shape_mn,
+        )
+
+        tiled_mma_sfb = sm100_utils.make_blockscaled_trivial_tiled_mma(
+            self.a_dtype,
+            self.a_major_mode,
+            self.b_major_mode,
+            self.sf_dtype,
+            self.sf_vec_size,
+            cute.nvgpu.tcgen05.CtaGroup.ONE,
+            self.mma_inst_shape_mn_sfb,
+        )
+        atom_thr_size = cute.size(tiled_mma.thr_id.shape)
+
+        # Setup TMA load for A
+        a_op = sm100_utils.cluster_shape_to_tma_atom_A(
+            self.cluster_shape_mn, tiled_mma.thr_id
+        )
+        a_smem_layout = cute.slice_(self.a_smem_layout_staged, (None, None, None, 0))
+        tma_atom_a, tma_tensor_a = cute.nvgpu.make_tiled_tma_atom_A(
+            a_op,
+            a_tensor,
+            a_smem_layout,
+            self.mma_tiler,
+            tiled_mma,
+            self.cluster_layout_vmnk.shape,
+        )
+
+        # Setup TMA load for B
+        b_op = sm100_utils.cluster_shape_to_tma_atom_B(
+            self.cluster_shape_mn, tiled_mma.thr_id
+        )
+        b_smem_layout = cute.slice_(self.b_smem_layout_staged, (None, None, None, 0))
+        tma_atom_b, tma_tensor_b = cute.nvgpu.make_tiled_tma_atom_B(
+            b_op,
+            b_tensor,
+            b_smem_layout,
+            self.mma_tiler,
+            tiled_mma,
+            self.cluster_layout_vmnk.shape,
+        )
+
+        # Setup TMA load for SFA
+        sfa_op = sm100_utils.cluster_shape_to_tma_atom_A(
+            self.cluster_shape_mn, tiled_mma.thr_id
+        )
+        sfa_smem_layout = cute.slice_(
+            self.sfa_smem_layout_staged, (None, None, None, 0)
+        )
+        tma_atom_sfa, tma_tensor_sfa = cute.nvgpu.make_tiled_tma_atom_A(
+            sfa_op,
+            sfa_tensor,
+            sfa_smem_layout,
+            self.mma_tiler,
+            tiled_mma,
+            self.cluster_layout_vmnk.shape,
+            internal_type=cutlass.Int16,
+        )
+
+        # Setup TMA load for SFB
+        sfb_op = sm100_utils.cluster_shape_to_tma_atom_SFB(
+            self.cluster_shape_mn, tiled_mma.thr_id
+        )
+        sfb_smem_layout = cute.slice_(
+            self.sfb_smem_layout_staged, (None, None, None, 0)
+        )
+        tma_atom_sfb, tma_tensor_sfb = cute.nvgpu.make_tiled_tma_atom_B(
+            sfb_op,
+            sfb_tensor,
+            sfb_smem_layout,
+            self.mma_tiler_sfb,
+            tiled_mma_sfb,
+            self.cluster_layout_sfb_vmnk.shape,
+            internal_type=cutlass.Int16,
+        )
+
+        if cutlass.const_expr(self.cta_tile_shape_mnk[1] == 192):
+            x = tma_tensor_sfb.stride[0][1]
+            y = cute.ceil_div(tma_tensor_sfb.shape[0][1], 4)
+
+            new_shape = (
+                (tma_tensor_sfb.shape[0][0], ((2, 2), y)),
+                tma_tensor_sfb.shape[1],
+                tma_tensor_sfb.shape[2],
+            )
+            # Use right multiplication for ScaledBasis (3 * x instead of x * 3)
+            x_times_3 = 3 * x
+            new_stride = (
+                (tma_tensor_sfb.stride[0][0], ((x, x), x_times_3)),
+                tma_tensor_sfb.stride[1],
+                tma_tensor_sfb.stride[2],
+            )
+            tma_tensor_sfb_new_layout = cute.make_layout(new_shape, stride=new_stride)
+            tma_tensor_sfb = cute.make_tensor(
+                tma_tensor_sfb.iterator, tma_tensor_sfb_new_layout
+            )
+
+        a_copy_size = cute.size_in_bytes(self.a_dtype, a_smem_layout)
+        b_copy_size = cute.size_in_bytes(self.b_dtype, b_smem_layout)
+        sfa_copy_size = cute.size_in_bytes(self.sf_dtype, sfa_smem_layout)
+        sfb_copy_size = cute.size_in_bytes(self.sf_dtype, sfb_smem_layout)
+        self.num_tma_load_bytes = (
+            a_copy_size + b_copy_size + sfa_copy_size + sfb_copy_size
+        ) * atom_thr_size
+
+        # Setup TMA store for C
+        epi_smem_layout = cute.slice_(self.c_smem_layout_staged, (None, None, 0))
+        tma_atom_c, tma_tensor_c = cpasync.make_tiled_tma_atom(
+            cpasync.CopyBulkTensorTileS2GOp(),
+            c_tensor,
+            epi_smem_layout,
+            self.epi_tile,
+        )
+
+        # Compute grid size
+        self.tile_sched_params, grid = self._compute_grid(
+            c_tensor,
+            self.cta_tile_shape_mnk,
+            self.cluster_shape_mn,
+            max_active_clusters,
+        )
+
+        self.buffer_align_bytes = 1024
+
+        # Define shared storage for kernel
+        @cute.struct
+        class SharedStorage:
+            ab_full_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage]
+            ab_empty_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage]
+            acc_full_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage]
+            acc_empty_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage]
+            tmem_dealloc_mbar_ptr: cutlass.Int64
+            tmem_holding_buf: cutlass.Int32
+            # (EPI_TILE_M, EPI_TILE_N, STAGE)
+            sC: cute.struct.Align[
+                cute.struct.MemRange[
+                    self.c_dtype,
+                    cute.cosize(self.c_smem_layout_staged.outer),
+                ],
+                self.buffer_align_bytes,
+            ]
+            # (MMA, MMA_M, MMA_K, STAGE)
+            sA: cute.struct.Align[
+                cute.struct.MemRange[
+                    self.a_dtype, cute.cosize(self.a_smem_layout_staged.outer)
+                ],
+                self.buffer_align_bytes,
+            ]
+            # (MMA, MMA_N, MMA_K, STAGE)
+            sB: cute.struct.Align[
+                cute.struct.MemRange[
+                    self.b_dtype, cute.cosize(self.b_smem_layout_staged.outer)
+                ],
+                self.buffer_align_bytes,
+            ]
+            # (MMA, MMA_M, MMA_K, STAGE)
+            sSFA: cute.struct.Align[
+                cute.struct.MemRange[
+                    self.sf_dtype, cute.cosize(self.sfa_smem_layout_staged)
+                ],
+                self.buffer_align_bytes,
+            ]
+            # (MMA, MMA_N, MMA_K, STAGE)
+            sSFB: cute.struct.Align[
+                cute.struct.MemRange[
+                    self.sf_dtype, cute.cosize(self.sfb_smem_layout_staged)
+                ],
+                self.buffer_align_bytes,
+            ]
+
+        self.shared_storage = SharedStorage
+
+        # Launch the kernel synchronously
+        self.kernel(
+            tiled_mma,
+            tiled_mma_sfb,
+            tma_atom_a,
+            tma_tensor_a,
+            tma_atom_b,
+            tma_tensor_b,
+            tma_atom_sfa,
+            tma_tensor_sfa,
+            tma_atom_sfb,
+            tma_tensor_sfb,
+            tma_atom_c,
+            tma_tensor_c,
+            self.cluster_layout_vmnk,
+            self.cluster_layout_sfb_vmnk,
+            self.a_smem_layout_staged,
+            self.b_smem_layout_staged,
+            self.sfa_smem_layout_staged,
+            self.sfb_smem_layout_staged,
+            self.c_smem_layout_staged,
+            self.epi_tile,
+            self.tile_sched_params,
+            epilogue_op,
+        ).launch(
+            grid=grid,
+            block=[self.threads_per_cta, 1, 1],
+            cluster=(*self.cluster_shape_mn, 1),
+            stream=stream,
+            min_blocks_per_mp=1,
+        )
+        return
+
+    # GPU device kernel
+    @cute.kernel
+    def kernel(
+        self,
+        tiled_mma: cute.TiledMma,
+        tiled_mma_sfb: cute.TiledMma,
+        tma_atom_a: cute.CopyAtom,
+        mA_mkl: cute.Tensor,
+        tma_atom_b: cute.CopyAtom,
+        mB_nkl: cute.Tensor,
+        tma_atom_sfa: cute.CopyAtom,
+        mSFA_mkl: cute.Tensor,
+        tma_atom_sfb: cute.CopyAtom,
+        mSFB_nkl: cute.Tensor,
+        tma_atom_c: cute.CopyAtom,
+        mC_mnl: cute.Tensor,
+        cluster_layout_vmnk: cute.Layout,
+        cluster_layout_sfb_vmnk: cute.Layout,
+        a_smem_layout_staged: cute.ComposedLayout,
+        b_smem_layout_staged: cute.ComposedLayout,
+        sfa_smem_layout_staged: cute.Layout,
+        sfb_smem_layout_staged: cute.Layout,
+        c_smem_layout_staged: Union[cute.Layout, cute.ComposedLayout],
+        epi_tile: cute.Tile,
+        tile_sched_params: utils.PersistentTileSchedulerParams,
+        epilogue_op: cutlass.Constexpr,
+    ):
+        """
+        GPU device kernel performing the Persistent batched GEMM computation.
+        """
+        warp_idx = cute.arch.warp_idx()
+        warp_idx = cute.arch.make_warp_uniform(warp_idx)
+
+        #
+        # Prefetch tma desc
+        #
+        if warp_idx == self.tma_warp_id:
+            cpasync.prefetch_descriptor(tma_atom_a)
+            cpasync.prefetch_descriptor(tma_atom_b)
+            cpasync.prefetch_descriptor(tma_atom_sfa)
+            cpasync.prefetch_descriptor(tma_atom_sfb)
+            cpasync.prefetch_descriptor(tma_atom_c)
+
+        use_2cta_instrs = cute.size(tiled_mma.thr_id.shape) == 2
+
+        #
+        # Setup cta/thread coordinates
+        #
+        # Coords inside cluster
+        bidx, bidy, bidz = cute.arch.block_idx()
+        mma_tile_coord_v = bidx % cute.size(tiled_mma.thr_id.shape)
+        is_leader_cta = mma_tile_coord_v == 0
+        cta_rank_in_cluster = cute.arch.make_warp_uniform(
+            cute.arch.block_idx_in_cluster()
+        )
+        block_in_cluster_coord_vmnk = cluster_layout_vmnk.get_flat_coord(
+            cta_rank_in_cluster
+        )
+        block_in_cluster_coord_sfb_vmnk = cluster_layout_sfb_vmnk.get_flat_coord(
+            cta_rank_in_cluster
+        )
+        # Coord inside cta
+        tidx, _, _ = cute.arch.thread_idx()
+
+        #
+        # Alloc and init: a+b full/empty, accumulator full/empty, tensor memory dealloc barrier
+        #
+        smem = utils.SmemAllocator()
+        storage = smem.allocate(self.shared_storage)
+
+        # Initialize mainloop ab_pipeline (barrier) and states
+        ab_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread)
+        num_tma_producer = self.num_mcast_ctas_a + self.num_mcast_ctas_b - 1
+        ab_pipeline_consumer_group = pipeline.CooperativeGroup(
+            pipeline.Agent.Thread, num_tma_producer
+        )
+        ab_pipeline = pipeline.PipelineTmaUmma.create(
+            barrier_storage=storage.ab_full_mbar_ptr.data_ptr(),
+            num_stages=self.num_ab_stage,
+            producer_group=ab_pipeline_producer_group,
+            consumer_group=ab_pipeline_consumer_group,
+            tx_count=self.num_tma_load_bytes,
+            cta_layout_vmnk=cluster_layout_vmnk,
+            defer_sync=True,
+        )
+
+        # Initialize acc_pipeline (barrier) and states
+        acc_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread)
+        num_acc_consumer_threads = len(self.epilog_warp_id) * (
+            2 if use_2cta_instrs else 1
+        )
+        acc_pipeline_consumer_group = pipeline.CooperativeGroup(
+            pipeline.Agent.Thread, num_acc_consumer_threads
+        )
+        acc_pipeline = pipeline.PipelineUmmaAsync.create(
+            barrier_storage=storage.acc_full_mbar_ptr.data_ptr(),
+            num_stages=self.num_acc_stage,
+            producer_group=acc_pipeline_producer_group,
+            consumer_group=acc_pipeline_consumer_group,
+            cta_layout_vmnk=cluster_layout_vmnk,
+            defer_sync=True,
+        )
+
+        # Tensor memory dealloc barrier init
+        tmem = utils.TmemAllocator(
+            storage.tmem_holding_buf,
+            barrier_for_retrieve=self.tmem_alloc_barrier,
+            allocator_warp_id=self.epilog_warp_id[0],
+            is_two_cta=use_2cta_instrs,
+            two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar_ptr,
+        )
+
+        # Cluster arrive after barrier init
+        pipeline_init_arrive(cluster_shape_mn=self.cluster_shape_mn, is_relaxed=True)
+
+        #
+        # Setup smem tensor A/B/SFA/SFB/C
+        #
+        # (EPI_TILE_M, EPI_TILE_N, STAGE)
+        sC = storage.sC.get_tensor(
+            c_smem_layout_staged.outer, swizzle=c_smem_layout_staged.inner
+        )
+        # (MMA, MMA_M, MMA_K, STAGE)
+        sA = storage.sA.get_tensor(
+            a_smem_layout_staged.outer, swizzle=a_smem_layout_staged.inner
+        )
+        # (MMA, MMA_N, MMA_K, STAGE)
+        sB = storage.sB.get_tensor(
+            b_smem_layout_staged.outer, swizzle=b_smem_layout_staged.inner
+        )
+        # (MMA, MMA_M, MMA_K, STAGE)
+        sSFA = storage.sSFA.get_tensor(sfa_smem_layout_staged)
+        # (MMA, MMA_N, MMA_K, STAGE)
+        sSFB = storage.sSFB.get_tensor(sfb_smem_layout_staged)
+
+        #
+        # Compute multicast mask for A/B/SFA/SFB buffer full
+        #
+        a_full_mcast_mask = None
+        b_full_mcast_mask = None
+        sfa_full_mcast_mask = None
+        sfb_full_mcast_mask = None
+        if cutlass.const_expr(self.is_a_mcast or self.is_b_mcast or use_2cta_instrs):
+            a_full_mcast_mask = cpasync.create_tma_multicast_mask(
+                cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=2
+            )
+            b_full_mcast_mask = cpasync.create_tma_multicast_mask(
+                cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=1
+            )
+            sfa_full_mcast_mask = cpasync.create_tma_multicast_mask(
+                cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=2
+            )
+            sfb_full_mcast_mask = cpasync.create_tma_multicast_mask(
+                cluster_layout_sfb_vmnk, block_in_cluster_coord_sfb_vmnk, mcast_mode=1
+            )
+
+        #
+        # Local_tile partition global tensors
+        #
+        # (bM, bK, RestM, RestK, RestL)
+        gA_mkl = cute.local_tile(
+            mA_mkl, cute.slice_(self.mma_tiler, (None, 0, None)), (None, None, None)
+        )
+        # (bN, bK, RestN, RestK, RestL)
+        gB_nkl = cute.local_tile(
+            mB_nkl, cute.slice_(self.mma_tiler, (0, None, None)), (None, None, None)
+        )
+        # (bM, bK, RestM, RestK, RestL)
+        gSFA_mkl = cute.local_tile(
+            mSFA_mkl, cute.slice_(self.mma_tiler, (None, 0, None)), (None, None, None)
+        )
+        # (bN, bK, RestN, RestK, RestL)
+        gSFB_nkl = cute.local_tile(
+            mSFB_nkl,
+            cute.slice_(self.mma_tiler_sfb, (0, None, None)),
+            (None, None, None),
+        )
+        # (bM, bN, RestM, RestN, RestL)
+        gC_mnl = cute.local_tile(
+            mC_mnl, cute.slice_(self.mma_tiler, (None, None, 0)), (None, None, None)
+        )
+        k_tile_cnt = cute.size(gA_mkl, mode=[3])
+
+        #
+        # Partition global tensor for TiledMMA_A/B/C
+        #
+        thr_mma = tiled_mma.get_slice(mma_tile_coord_v)
+        thr_mma_sfb = tiled_mma_sfb.get_slice(mma_tile_coord_v)
+        # (MMA, MMA_M, MMA_K, RestM, RestK, RestL)
+        tCgA = thr_mma.partition_A(gA_mkl)
+        # (MMA, MMA_N, MMA_K, RestN, RestK, RestL)
+        tCgB = thr_mma.partition_B(gB_nkl)
+        # (MMA, MMA_M, MMA_K, RestM, RestK, RestL)
+        tCgSFA = thr_mma.partition_A(gSFA_mkl)
+        # (MMA, MMA_N, MMA_K, RestN, RestK, RestL)
+        tCgSFB = thr_mma_sfb.partition_B(gSFB_nkl)
+        # (MMA, MMA_M, MMA_N, RestM, RestN, RestL)
+        tCgC = thr_mma.partition_C(gC_mnl)
+
+        #
+        # Partition global/shared tensor for TMA load A/B
+        #
+        # TMA load A partition_S/D
+        a_cta_layout = cute.make_layout(
+            cute.slice_(cluster_layout_vmnk, (0, 0, None, 0)).shape
+        )
+        # ((atom_v, rest_v), STAGE)
+        # ((atom_v, rest_v), RestM, RestK, RestL)
+        tAsA, tAgA = cpasync.tma_partition(
+            tma_atom_a,
+            block_in_cluster_coord_vmnk[2],
+            a_cta_layout,
+            cute.group_modes(sA, 0, 3),
+            cute.group_modes(tCgA, 0, 3),
+        )
+        # TMA load B partition_S/D
+        b_cta_layout = cute.make_layout(
+            cute.slice_(cluster_layout_vmnk, (0, None, 0, 0)).shape
+        )
+        # ((atom_v, rest_v), STAGE)
+        # ((atom_v, rest_v), RestN, RestK, RestL)
+        tBsB, tBgB = cpasync.tma_partition(
+            tma_atom_b,
+            block_in_cluster_coord_vmnk[1],
+            b_cta_layout,
+            cute.group_modes(sB, 0, 3),
+            cute.group_modes(tCgB, 0, 3),
+        )
+
+        #  TMA load SFA partition_S/D
+        sfa_cta_layout = a_cta_layout
+        # ((atom_v, rest_v), STAGE)
+        # ((atom_v, rest_v), RestM, RestK, RestL)
+        tAsSFA, tAgSFA = cute.nvgpu.cpasync.tma_partition(
+            tma_atom_sfa,
+            block_in_cluster_coord_vmnk[2],
+            sfa_cta_layout,
+            cute.group_modes(sSFA, 0, 3),
+            cute.group_modes(tCgSFA, 0, 3),
+        )
+        tAsSFA = cute.filter_zeros(tAsSFA)
+        tAgSFA = cute.filter_zeros(tAgSFA)
+
+        # TMA load SFB partition_S/D
+        sfb_cta_layout = cute.make_layout(
+            cute.slice_(cluster_layout_sfb_vmnk, (0, None, 0, 0)).shape
+        )
+        # ((atom_v, rest_v), STAGE)
+        # ((atom_v, rest_v), RestN, RestK, RestL)
+        tBsSFB, tBgSFB = cute.nvgpu.cpasync.tma_partition(
+            tma_atom_sfb,
+            block_in_cluster_coord_sfb_vmnk[1],
+            sfb_cta_layout,
+            cute.group_modes(sSFB, 0, 3),
+            cute.group_modes(tCgSFB, 0, 3),
+        )
+        tBsSFB = cute.filter_zeros(tBsSFB)
+        tBgSFB = cute.filter_zeros(tBgSFB)
+
+        #
+        # Partition shared/tensor memory tensor for TiledMMA_A/B/C
+        #
+        # (MMA, MMA_M, MMA_K, STAGE)
+        tCrA = tiled_mma.make_fragment_A(sA)
+        # (MMA, MMA_N, MMA_K, STAGE)
+        tCrB = tiled_mma.make_fragment_B(sB)
+        # (MMA, MMA_M, MMA_N)
+        acc_shape = tiled_mma.partition_shape_C(self.mma_tiler[:2])
+        if cutlass.const_expr(self.overlapping_accum):
+            num_acc_stage_overlapped = 2
+            tCtAcc_fake = tiled_mma.make_fragment_C(
+                cute.append(acc_shape, num_acc_stage_overlapped)
+            )
+            # (MMA, MMA_M, MMA_N, STAGE)
+            tCtAcc_fake = cute.make_tensor(
+                tCtAcc_fake.iterator,
+                cute.make_layout(
+                    tCtAcc_fake.shape,
+                    stride=(
+                        tCtAcc_fake.stride[0],
+                        tCtAcc_fake.stride[1],
+                        tCtAcc_fake.stride[2],
+                        (256 - self.num_sf_tmem_cols) * tCtAcc_fake.stride[0][1],
+                    ),
+                ),
+            )
+        else:
+            # (MMA, MMA_M, MMA_N, STAGE)
+            tCtAcc_fake = tiled_mma.make_fragment_C(
+                cute.append(acc_shape, self.num_acc_stage)
+            )
+
+        #
+        # Cluster wait before tensor memory alloc
+        #
+        pipeline_init_wait(cluster_shape_mn=self.cluster_shape_mn)
+
+        #
+        # Specialized TMA load warp
+        #
+        if warp_idx == self.tma_warp_id:
+            #
+            # Persistent tile scheduling loop
+            #
+            tile_sched = utils.StaticPersistentTileScheduler.create(
+                tile_sched_params, cute.arch.block_idx(), cute.arch.grid_dim()
+            )
+            work_tile = tile_sched.initial_work_tile_info()
+
+            ab_producer_state = pipeline.make_pipeline_state(
+                pipeline.PipelineUserType.Producer, self.num_ab_stage
+            )
+
+            while work_tile.is_valid_tile:
+                # Get tile coord from tile scheduler
+                cur_tile_coord = work_tile.tile_idx
+                mma_tile_coord_mnl = (
+                    cur_tile_coord[0] // cute.size(tiled_mma.thr_id.shape),
+                    cur_tile_coord[1],
+                    cur_tile_coord[2],
+                )
+
+                #
+                # Slice to per mma tile index
+                #
+                # ((atom_v, rest_v), RestK)
+                tAgA_slice = tAgA[
+                    (None, mma_tile_coord_mnl[0], None, mma_tile_coord_mnl[2])
+                ]
+                # ((atom_v, rest_v), RestK)
+                tBgB_slice = tBgB[
+                    (None, mma_tile_coord_mnl[1], None, mma_tile_coord_mnl[2])
+                ]
+
+                # ((atom_v, rest_v), RestK)
+                tAgSFA_slice = tAgSFA[
+                    (None, mma_tile_coord_mnl[0], None, mma_tile_coord_mnl[2])
+                ]
+
+                slice_n = mma_tile_coord_mnl[1]
+                if cutlass.const_expr(self.cta_tile_shape_mnk[1] == 64):
+                    slice_n = mma_tile_coord_mnl[1] // 2
+                # ((atom_v, rest_v), RestK)
+                tBgSFB_slice = tBgSFB[(None, slice_n, None, mma_tile_coord_mnl[2])]
+
+                # Peek (try_wait) AB buffer empty for k_tile = prefetch_k_tile_cnt
+                ab_producer_state.reset_count()
+                peek_ab_empty_status = cutlass.Boolean(1)
+                if ab_producer_state.count < k_tile_cnt:
+                    peek_ab_empty_status = ab_pipeline.producer_try_acquire(
+                        ab_producer_state
+                    )
+                #
+                # Tma load loop
+                #
+                for k_tile in cutlass.range(0, k_tile_cnt, 1, unroll=1):
+                    # Conditionally wait for AB buffer empty
+                    ab_pipeline.producer_acquire(
+                        ab_producer_state, peek_ab_empty_status
+                    )
+
+                    # TMA load A/B/SFA/SFB
+                    cute.copy(
+                        tma_atom_a,
+                        tAgA_slice[(None, ab_producer_state.count)],
+                        tAsA[(None, ab_producer_state.index)],
+                        tma_bar_ptr=ab_pipeline.producer_get_barrier(ab_producer_state),
+                        mcast_mask=a_full_mcast_mask,
+                    )
+                    cute.copy(
+                        tma_atom_b,
+                        tBgB_slice[(None, ab_producer_state.count)],
+                        tBsB[(None, ab_producer_state.index)],
+                        tma_bar_ptr=ab_pipeline.producer_get_barrier(ab_producer_state),
+                        mcast_mask=b_full_mcast_mask,
+                    )
+                    cute.copy(
+                        tma_atom_sfa,
+                        tAgSFA_slice[(None, ab_producer_state.count)],
+                        tAsSFA[(None, ab_producer_state.index)],
+                        tma_bar_ptr=ab_pipeline.producer_get_barrier(ab_producer_state),
+                        mcast_mask=sfa_full_mcast_mask,
+                    )
+                    cute.copy(
+                        tma_atom_sfb,
+                        tBgSFB_slice[(None, ab_producer_state.count)],
+                        tBsSFB[(None, ab_producer_state.index)],
+                        tma_bar_ptr=ab_pipeline.producer_get_barrier(ab_producer_state),
+                        mcast_mask=sfb_full_mcast_mask,
+                    )
+
+                    # Peek (try_wait) AB buffer empty for k_tile = prefetch_k_tile_cnt + k_tile + 1
+                    ab_producer_state.advance()
+                    peek_ab_empty_status = cutlass.Boolean(1)
+                    if ab_producer_state.count < k_tile_cnt:
+                        peek_ab_empty_status = ab_pipeline.producer_try_acquire(
+                            ab_producer_state
+                        )
+
+                #
+                # Advance to next tile
+                #
+                tile_sched.advance_to_next_work()
+                work_tile = tile_sched.get_current_work()
+
+            #
+            # Wait A/B buffer empty
+            #
+            ab_pipeline.producer_tail(ab_producer_state)
+
+        #
+        # Specialized MMA warp
+        #
+        if warp_idx == self.mma_warp_id:
+            #
+            # Bar sync for retrieve tensor memory ptr from shared mem
+            #
+            tmem.wait_for_alloc()
+
+            #
+            # Retrieving tensor memory ptr and make accumulator/SFA/SFB tensor
+            #
+            acc_tmem_ptr = tmem.retrieve_ptr(self.acc_dtype)
+            # Make accumulator tmem tensor
+            # (MMA, MMA_M, MMA_N, STAGE)
+            tCtAcc_base = cute.make_tensor(acc_tmem_ptr, tCtAcc_fake.layout)
+
+            # Make SFA tmem tensor
+            sfa_tmem_ptr = cute.recast_ptr(
+                acc_tmem_ptr + self.num_accumulator_tmem_cols,
+                dtype=self.sf_dtype,
+            )
+            # (MMA, MMA_M, MMA_K)
+            tCtSFA_layout = blockscaled_utils.make_tmem_layout_sfa(
+                tiled_mma,
+                self.mma_tiler,
+                self.sf_vec_size,
+                cute.slice_(sfa_smem_layout_staged, (None, None, None, 0)),
+            )
+            tCtSFA = cute.make_tensor(sfa_tmem_ptr, tCtSFA_layout)
+
+            # Make SFB tmem tensor
+            sfb_tmem_ptr = cute.recast_ptr(
+                acc_tmem_ptr + self.num_accumulator_tmem_cols + self.num_sfa_tmem_cols,
+                dtype=self.sf_dtype,
+            )
+            # (MMA, MMA_N, MMA_K)
+            tCtSFB_layout = blockscaled_utils.make_tmem_layout_sfb(
+                tiled_mma,
+                self.mma_tiler,
+                self.sf_vec_size,
+                cute.slice_(sfb_smem_layout_staged, (None, None, None, 0)),
+            )
+            tCtSFB = cute.make_tensor(sfb_tmem_ptr, tCtSFB_layout)
+            #
+            # Partition for S2T copy of SFA/SFB
+            #
+            (
+                tiled_copy_s2t_sfa,
+                tCsSFA_compact_s2t,
+                tCtSFA_compact_s2t,
+            ) = self.mainloop_s2t_copy_and_partition(sSFA, tCtSFA)
+            (
+                tiled_copy_s2t_sfb,
+                tCsSFB_compact_s2t,
+                tCtSFB_compact_s2t,
+            ) = self.mainloop_s2t_copy_and_partition(sSFB, tCtSFB)
+
+            #
+            # Persistent tile scheduling loop
+            #
+            tile_sched = utils.StaticPersistentTileScheduler.create(
+                tile_sched_params, cute.arch.block_idx(), cute.arch.grid_dim()
+            )
+            work_tile = tile_sched.initial_work_tile_info()
+
+            ab_consumer_state = pipeline.make_pipeline_state(
+                pipeline.PipelineUserType.Consumer, self.num_ab_stage
+            )
+            acc_producer_state = pipeline.make_pipeline_state(
+                pipeline.PipelineUserType.Producer, self.num_acc_stage
+            )
+
+            while work_tile.is_valid_tile:
+                # Get tile coord from tile scheduler
+                cur_tile_coord = work_tile.tile_idx
+                mma_tile_coord_mnl = (
+                    cur_tile_coord[0] // cute.size(tiled_mma.thr_id.shape),
+                    cur_tile_coord[1],
+                    cur_tile_coord[2],
+                )
+
+                # Get accumulator stage index
+                if cutlass.const_expr(self.overlapping_accum):
+                    acc_stage_index = acc_producer_state.phase ^ 1
+                else:
+                    acc_stage_index = acc_producer_state.index
+
+                # Set tensor memory buffer for current tile
+                # (MMA, MMA_M, MMA_N)
+                tCtAcc = tCtAcc_base[(None, None, None, acc_stage_index)]
+
+                # Peek (try_wait) AB buffer full for k_tile = 0
+                ab_consumer_state.reset_count()
+                peek_ab_full_status = cutlass.Boolean(1)
+                if ab_consumer_state.count < k_tile_cnt and is_leader_cta:
+                    peek_ab_full_status = ab_pipeline.consumer_try_wait(
+                        ab_consumer_state
+                    )
+
+                #
+                # Wait for accumulator buffer empty
+                #
+                if is_leader_cta:
+                    acc_pipeline.producer_acquire(acc_producer_state)
+
+                tCtSFB_mma = tCtSFB
+                if cutlass.const_expr(self.cta_tile_shape_mnk[1] == 192):
+                    # If this is an ODD tile, shift the TMEM start address for cta_tile_shape_n=192 case by two words (ignores first 64 columns of SFB)
+                    offset = (
+                        cutlass.Int32(2)
+                        if mma_tile_coord_mnl[1] % 2 == 1
+                        else cutlass.Int32(0)
+                    )
+                    shifted_ptr = cute.recast_ptr(
+                        acc_tmem_ptr
+                        + self.num_accumulator_tmem_cols
+                        + self.num_sfa_tmem_cols
+                        + offset,
+                        dtype=self.sf_dtype,
+                    )
+                    tCtSFB_mma = cute.make_tensor(shifted_ptr, tCtSFB_layout)
+                elif cutlass.const_expr(self.cta_tile_shape_mnk[1] == 64):
+                    # Move in increments of 64 columns of SFB
+                    offset = cutlass.Int32((mma_tile_coord_mnl[1] % 2) * 2)
+                    shifted_ptr = cute.recast_ptr(
+                        acc_tmem_ptr
+                        + self.num_accumulator_tmem_cols
+                        + self.num_sfa_tmem_cols
+                        + offset,
+                        dtype=self.sf_dtype,
+                    )
+                    tCtSFB_mma = cute.make_tensor(shifted_ptr, tCtSFB_layout)
+
+                #
+                # Reset the ACCUMULATE field for each tile
+                #
+                tiled_mma.set(tcgen05.Field.ACCUMULATE, False)
+
+                #
+                # Mma mainloop
+                #
+                for k_tile in range(k_tile_cnt):
+                    if is_leader_cta:
+                        # Conditionally wait for AB buffer full
+                        ab_pipeline.consumer_wait(
+                            ab_consumer_state, peek_ab_full_status
+                        )
+
+                        #  Copy SFA/SFB from smem to tmem
+                        s2t_stage_coord = (
+                            None,
+                            None,
+                            None,
+                            None,
+                            ab_consumer_state.index,
+                        )
+                        tCsSFA_compact_s2t_staged = tCsSFA_compact_s2t[s2t_stage_coord]
+                        tCsSFB_compact_s2t_staged = tCsSFB_compact_s2t[s2t_stage_coord]
+                        cute.copy(
+                            tiled_copy_s2t_sfa,
+                            tCsSFA_compact_s2t_staged,
+                            tCtSFA_compact_s2t,
+                        )
+                        cute.copy(
+                            tiled_copy_s2t_sfb,
+                            tCsSFB_compact_s2t_staged,
+                            tCtSFB_compact_s2t,
+                        )
+
+                        # tCtAcc += tCrA * tCrSFA * tCrB * tCrSFB
+                        num_kblocks = cute.size(tCrA, mode=[2])
+                        for kblock_idx in cutlass.range(num_kblocks, unroll_full=True):
+                            kblock_coord = (
+                                None,
+                                None,
+                                kblock_idx,
+                                ab_consumer_state.index,
+                            )
+
+                            # Set SFA/SFB tensor to tiled_mma
+                            sf_kblock_coord = (None, None, kblock_idx)
+                            tiled_mma.set(
+                                tcgen05.Field.SFA,
+                                tCtSFA[sf_kblock_coord].iterator,
+                            )
+                            tiled_mma.set(
+                                tcgen05.Field.SFB,
+                                tCtSFB_mma[sf_kblock_coord].iterator,
+                            )
+
+                            cute.gemm(
+                                tiled_mma,
+                                tCtAcc,
+                                tCrA[kblock_coord],
+                                tCrB[kblock_coord],
+                                tCtAcc,
+                            )
+
+                            # Enable accumulate on tCtAcc after first kblock
+                            tiled_mma.set(tcgen05.Field.ACCUMULATE, True)
+
+                        # Async arrive AB buffer empty
+                        ab_pipeline.consumer_release(ab_consumer_state)
+
+                    # Peek (try_wait) AB buffer full for k_tile = k_tile + 1
+                    ab_consumer_state.advance()
+                    peek_ab_full_status = cutlass.Boolean(1)
+                    if ab_consumer_state.count < k_tile_cnt:
+                        if is_leader_cta:
+                            peek_ab_full_status = ab_pipeline.consumer_try_wait(
+                                ab_consumer_state
+                            )
+
+                #
+                # Async arrive accumulator buffer full
+                #
+                if is_leader_cta:
+                    acc_pipeline.producer_commit(acc_producer_state)
+                acc_producer_state.advance()
+
+                #
+                # Advance to next tile
+                #
+                tile_sched.advance_to_next_work()
+                work_tile = tile_sched.get_current_work()
+
+            #
+            # Wait for accumulator buffer empty
+            #
+            acc_pipeline.producer_tail(acc_producer_state)
+        #
+        # Specialized epilogue warps
+        #
+        if warp_idx < self.mma_warp_id:
+            #
+            # Alloc tensor memory buffer
+            #
+            tmem.allocate(self.num_tmem_alloc_cols)
+
+            #
+            # Bar sync for retrieve tensor memory ptr from shared memory
+            #
+            tmem.wait_for_alloc()
+
+            #
+            # Retrieving tensor memory ptr and make accumulator tensor
+            #
+            acc_tmem_ptr = tmem.retrieve_ptr(self.acc_dtype)
+            # (MMA, MMA_M, MMA_N, STAGE)
+            tCtAcc_base = cute.make_tensor(acc_tmem_ptr, tCtAcc_fake.layout)
+
+            #
+            # Partition for epilogue
+            #
+            epi_tidx = tidx
+            (
+                tiled_copy_t2r,
+                tTR_tAcc_base,
+                tTR_rAcc,
+            ) = self.epilog_tmem_copy_and_partition(
+                epi_tidx, tCtAcc_base, tCgC, epi_tile, use_2cta_instrs
+            )
+
+            tTR_rC = cute.make_rmem_tensor(tTR_rAcc.shape, self.c_dtype)
+            tiled_copy_r2s, tRS_rC, tRS_sC = self.epilog_smem_copy_and_partition(
+                tiled_copy_t2r, tTR_rC, epi_tidx, sC
+            )
+            (
+                tma_atom_c,
+                bSG_sC,
+                bSG_gC_partitioned,
+            ) = self.epilog_gmem_copy_and_partition(
+                epi_tidx, tma_atom_c, tCgC, epi_tile, sC
+            )
+
+            #
+            # Persistent tile scheduling loop
+            #
+            tile_sched = utils.StaticPersistentTileScheduler.create(
+                tile_sched_params, cute.arch.block_idx(), cute.arch.grid_dim()
+            )
+            work_tile = tile_sched.initial_work_tile_info()
+
+            acc_consumer_state = pipeline.make_pipeline_state(
+                pipeline.PipelineUserType.Consumer, self.num_acc_stage
+            )
+
+            # Threads/warps participating in tma store pipeline
+            c_producer_group = pipeline.CooperativeGroup(
+                pipeline.Agent.Thread,
+                32 * len(self.epilog_warp_id),
+            )
+            c_pipeline = pipeline.PipelineTmaStore.create(
+                num_stages=self.num_c_stage,
+                producer_group=c_producer_group,
+            )
+
+            while work_tile.is_valid_tile:
+                # Get tile coord from tile scheduler
+                cur_tile_coord = work_tile.tile_idx
+                mma_tile_coord_mnl = (
+                    cur_tile_coord[0] // cute.size(tiled_mma.thr_id.shape),
+                    cur_tile_coord[1],
+                    cur_tile_coord[2],
+                )
+
+                #
+                # Slice to per mma tile index
+                #
+                # ((ATOM_V, REST_V), EPI_M, EPI_N)
+                bSG_gC = bSG_gC_partitioned[
+                    (
+                        None,
+                        None,
+                        None,
+                        *mma_tile_coord_mnl,
+                    )
+                ]
+
+                # Get accumulator stage index
+                if cutlass.const_expr(self.overlapping_accum):
+                    acc_stage_index = acc_consumer_state.phase
+                    reverse_subtile = (
+                        cutlass.Boolean(True)
+                        if acc_stage_index == 0
+                        else cutlass.Boolean(False)
+                    )
+                else:
+                    acc_stage_index = acc_consumer_state.index
+
+                # Set tensor memory buffer for current tile
+                # (T2R, T2R_M, T2R_N, EPI_M, EPI_M)
+                tTR_tAcc = tTR_tAcc_base[
+                    (None, None, None, None, None, acc_stage_index)
+                ]
+
+                #
+                # Wait for accumulator buffer full
+                #
+                acc_pipeline.consumer_wait(acc_consumer_state)
+
+                tTR_tAcc = cute.group_modes(tTR_tAcc, 3, cute.rank(tTR_tAcc))
+                bSG_gC = cute.group_modes(bSG_gC, 1, cute.rank(bSG_gC))
+
+                #
+                # Store accumulator to global memory in subtiles
+                #
+                subtile_cnt = cute.size(tTR_tAcc.shape, mode=[3])
+                num_prev_subtiles = tile_sched.num_tiles_executed * subtile_cnt
+                for subtile_idx in cutlass.range(subtile_cnt):
+                    real_subtile_idx = subtile_idx
+                    if cutlass.const_expr(self.overlapping_accum):
+                        if reverse_subtile:
+                            real_subtile_idx = (
+                                self.cta_tile_shape_mnk[1] // self.epi_tile_n
+                                - 1
+                                - subtile_idx
+                            )
+                    #
+                    # Load accumulator from tensor memory buffer to register
+                    #
+                    tTR_tAcc_mn = tTR_tAcc[(None, None, None, real_subtile_idx)]
+                    cute.copy(tiled_copy_t2r, tTR_tAcc_mn, tTR_rAcc)
+
+                    #
+                    # Async arrive accumulator buffer empty ealier when overlapping_accum is enabled
+                    #
+                    if cutlass.const_expr(self.overlapping_accum):
+                        if subtile_idx == self.iter_acc_early_release_in_epilogue:
+                            # Fence for TMEM load
+                            cute.arch.fence_view_async_tmem_load()
+                            with cute.arch.elect_one():
+                                acc_pipeline.consumer_release(acc_consumer_state)
+                            acc_consumer_state.advance()
+
+                    #
+                    # Convert to C type
+                    #
+                    acc_vec = tiled_copy_r2s.retile(tTR_rAcc).load()
+                    acc_vec = epilogue_op(acc_vec.to(self.c_dtype))
+                    tRS_rC.store(acc_vec)
+
+                    #
+                    # Store C to shared memory
+                    #
+                    c_buffer = (num_prev_subtiles + real_subtile_idx) % self.num_c_stage
+                    cute.copy(
+                        tiled_copy_r2s,
+                        tRS_rC,
+                        tRS_sC[(None, None, None, c_buffer)],
+                    )
+                    # Fence and barrier to make sure shared memory store is visible to TMA store
+                    cute.arch.fence_proxy(
+                        cute.arch.ProxyKind.async_shared,
+                        space=cute.arch.SharedSpace.shared_cta,
+                    )
+                    self.epilog_sync_barrier.arrive_and_wait()
+
+                    #
+                    # TMA store C to global memory
+                    #
+                    if warp_idx == self.epilog_warp_id[0]:
+                        cute.copy(
+                            tma_atom_c,
+                            bSG_sC[(None, c_buffer)],
+                            bSG_gC[(None, real_subtile_idx)],
+                        )
+                        # Fence and barrier to make sure shared memory store is visible to TMA store
+                        c_pipeline.producer_commit()
+                        c_pipeline.producer_acquire()
+                    self.epilog_sync_barrier.arrive_and_wait()
+
+                #
+                # Async arrive accumulator buffer empty
+                #
+                if cutlass.const_expr(not self.overlapping_accum):
+                    with cute.arch.elect_one():
+                        acc_pipeline.consumer_release(acc_consumer_state)
+                    acc_consumer_state.advance()
+
+                #
+                # Advance to next tile
+                #
+                tile_sched.advance_to_next_work()
+                work_tile = tile_sched.get_current_work()
+
+            #
+            # Dealloc the tensor memory buffer
+            #
+            tmem.relinquish_alloc_permit()
+            self.epilog_sync_barrier.arrive_and_wait()
+            tmem.free(acc_tmem_ptr)
+            #
+            # Wait for C store complete
+            #
+            c_pipeline.producer_tail()
+
+    def mainloop_s2t_copy_and_partition(
+        self,
+        sSF: cute.Tensor,
+        tSF: cute.Tensor,
+    ) -> Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor]:
+        """
+        Make tiledCopy for smem to tmem load for scale factor tensor, then use it to partition smem memory (source) and tensor memory (destination).
+
+        :param sSF: The scale factor tensor in smem
+        :type sSF: cute.Tensor
+        :param tSF: The scale factor tensor in tmem
+        :type tSF: cute.Tensor
+
+        :return: A tuple containing (tiled_copy_s2t, tCsSF_compact_s2t, tCtSF_compact_s2t) where:
+            - tiled_copy_s2t: The tiled copy operation for smem to tmem load for scale factor tensor(s2t)
+            - tCsSF_compact_s2t: The partitioned scale factor tensor in smem
+            - tSF_compact_s2t: The partitioned scale factor tensor in tmem
+        :rtype: Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor]
+        """
+        # (MMA, MMA_MN, MMA_K, STAGE)
+        tCsSF_compact = cute.filter_zeros(sSF)
+        # (MMA, MMA_MN, MMA_K)
+        tCtSF_compact = cute.filter_zeros(tSF)
+
+        # Make S2T CopyAtom and tiledCopy
+        copy_atom_s2t = cute.make_copy_atom(
+            tcgen05.Cp4x32x128bOp(self.cta_group),
+            self.sf_dtype,
+        )
+        tiled_copy_s2t = tcgen05.make_s2t_copy(copy_atom_s2t, tCtSF_compact)
+        thr_copy_s2t = tiled_copy_s2t.get_slice(0)
+
+        # ((ATOM_V, REST_V), Rest_Tiler, MMA_MN, MMA_K, STAGE)
+        tCsSF_compact_s2t_ = thr_copy_s2t.partition_S(tCsSF_compact)
+        # ((ATOM_V, REST_V), Rest_Tiler, MMA_MN, MMA_K, STAGE)
+        tCsSF_compact_s2t = tcgen05.get_s2t_smem_desc_tensor(
+            tiled_copy_s2t, tCsSF_compact_s2t_
+        )
+        # ((ATOM_V, REST_V), Rest_Tiler, MMA_MN, MMA_K)
+        tCtSF_compact_s2t = thr_copy_s2t.partition_D(tCtSF_compact)
+
+        return tiled_copy_s2t, tCsSF_compact_s2t, tCtSF_compact_s2t
+
+    def epilog_tmem_copy_and_partition(
+        self,
+        tidx: cutlass.Int32,
+        tAcc: cute.Tensor,
+        gC_mnl: cute.Tensor,
+        epi_tile: cute.Tile,
+        use_2cta_instrs: Union[cutlass.Boolean, bool],
+    ) -> Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor]:
+        """
+        Make tiledCopy for tensor memory load, then use it to partition tensor memory (source) and register array (destination).
+
+        :param tidx: The thread index in epilogue warp groups
+        :type tidx: cutlass.Int32
+        :param tAcc: The accumulator tensor to be copied and partitioned
+        :type tAcc: cute.Tensor
+        :param gC_mnl: The global tensor C
+        :type gC_mnl: cute.Tensor
+        :param epi_tile: The epilogue tiler
+        :type epi_tile: cute.Tile
+        :param use_2cta_instrs: Whether use_2cta_instrs is enabled
+        :type use_2cta_instrs: bool
+
+        :return: A tuple containing (tiled_copy_t2r, tTR_tAcc, tTR_rAcc) where:
+            - tiled_copy_t2r: The tiled copy operation for tmem to register copy(t2r)
+            - tTR_tAcc: The partitioned accumulator tensor
+            - tTR_rAcc: The accumulated tensor in register used to hold t2r results
+        :rtype: Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor]
+        """
+        # Make tiledCopy for tensor memory load
+        copy_atom_t2r = sm100_utils.get_tmem_load_op(
+            self.cta_tile_shape_mnk,
+            self.c_layout,
+            self.c_dtype,
+            self.acc_dtype,
+            epi_tile,
+            use_2cta_instrs,
+        )
+        # (EPI_TILE_M, EPI_TILE_N, EPI_M, EPI_N, STAGE)
+        tAcc_epi = cute.flat_divide(
+            tAcc[((None, None), 0, 0, None)],
+            epi_tile,
+        )
+        # (EPI_TILE_M, EPI_TILE_N)
+        tiled_copy_t2r = tcgen05.make_tmem_copy(
+            copy_atom_t2r, tAcc_epi[(None, None, 0, 0, 0)]
+        )
+
+        thr_copy_t2r = tiled_copy_t2r.get_slice(tidx)
+        # (T2R, T2R_M, T2R_N, EPI_M, EPI_M, STAGE)
+        tTR_tAcc = thr_copy_t2r.partition_S(tAcc_epi)
+
+        # (EPI_TILE_M, EPI_TILE_N, EPI_M, EPI_N, RestM, RestN, RestL)
+        gC_mnl_epi = cute.flat_divide(
+            gC_mnl[((None, None), 0, 0, None, None, None)], epi_tile
+        )
+        # (T2R, T2R_M, T2R_N, EPI_M, EPI_N, RestM, RestN, RestL)
+        tTR_gC = thr_copy_t2r.partition_D(gC_mnl_epi)
+        # (T2R, T2R_M, T2R_N)
+        tTR_rAcc = cute.make_rmem_tensor(
+            tTR_gC[(None, None, None, 0, 0, 0, 0, 0)].shape, self.acc_dtype
+        )
+        return tiled_copy_t2r, tTR_tAcc, tTR_rAcc
+
+    def epilog_smem_copy_and_partition(
+        self,
+        tiled_copy_t2r: cute.TiledCopy,
+        tTR_rC: cute.Tensor,
+        tidx: cutlass.Int32,
+        sC: cute.Tensor,
+    ) -> Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor]:
+        """
+        Make tiledCopy for shared memory store, then use it to partition register array (source) and shared memory (destination).
+
+        :param tiled_copy_t2r: The tiled copy operation for tmem to register copy(t2r)
+        :type tiled_copy_t2r: cute.TiledCopy
+        :param tTR_rC: The partitioned accumulator tensor
+        :type tTR_rC: cute.Tensor
+        :param tidx: The thread index in epilogue warp groups
+        :type tidx: cutlass.Int32
+        :param sC: The shared memory tensor to be copied and partitioned
+        :type sC: cute.Tensor
+        :type sepi: cute.Tensor
+
+        :return: A tuple containing (tiled_copy_r2s, tRS_rC, tRS_sC) where:
+            - tiled_copy_r2s: The tiled copy operation for register to smem copy(r2s)
+            - tRS_rC: The partitioned tensor C (register source)
+            - tRS_sC: The partitioned tensor C (smem destination)
+        :rtype: Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor]
+        """
+        copy_atom_r2s = sm100_utils.get_smem_store_op(
+            self.c_layout, self.c_dtype, self.acc_dtype, tiled_copy_t2r
+        )
+        tiled_copy_r2s = cute.make_tiled_copy_D(copy_atom_r2s, tiled_copy_t2r)
+        # (R2S, R2S_M, R2S_N, PIPE_D)
+        thr_copy_r2s = tiled_copy_r2s.get_slice(tidx)
+        tRS_sC = thr_copy_r2s.partition_D(sC)
+        # (R2S, R2S_M, R2S_N)
+        tRS_rC = tiled_copy_r2s.retile(tTR_rC)
+        return tiled_copy_r2s, tRS_rC, tRS_sC
+
+    def epilog_gmem_copy_and_partition(
+        self,
+        tidx: cutlass.Int32,
+        atom: Union[cute.CopyAtom, cute.TiledCopy],
+        gC_mnl: cute.Tensor,
+        epi_tile: cute.Tile,
+        sC: cute.Tensor,
+    ) -> Tuple[cute.CopyAtom, cute.Tensor, cute.Tensor]:
+        """Make tiledCopy for global memory store, then use it to:
+        partition shared memory (source) and global memory (destination) for TMA store version.
+
+        :param tidx: The thread index in epilogue warp groups
+        :type tidx: cutlass.Int32
+        :param atom: The copy_atom_c to be used for TMA store version, or tiled_copy_t2r for none TMA store version
+        :type atom: cute.CopyAtom or cute.TiledCopy
+        :param gC_mnl: The global tensor C
+        :type gC_mnl: cute.Tensor
+        :param epi_tile: The epilogue tiler
+        :type epi_tile: cute.Tile
+        :param sC: The shared memory tensor to be copied and partitioned
+        :type sC: cute.Tensor
+
+        :return: A tuple containing (tma_atom_c, bSG_sC, bSG_gC) where:
+            - tma_atom_c: The TMA copy atom
+            - bSG_sC: The partitioned shared memory tensor C
+            - bSG_gC: The partitioned global tensor C
+        :rtype: Tuple[cute.CopyAtom, cute.Tensor, cute.Tensor]
+        """
+        # (EPI_TILE_M, EPI_TILE_N, EPI_M, EPI_N, RestM, RestN, RestL)
+        gC_epi = cute.flat_divide(
+            gC_mnl[((None, None), 0, 0, None, None, None)], epi_tile
+        )
+
+        tma_atom_c = atom
+        sC_for_tma_partition = cute.group_modes(sC, 0, 2)
+        gC_for_tma_partition = cute.group_modes(gC_epi, 0, 2)
+        # ((ATOM_V, REST_V), EPI_M, EPI_N)
+        # ((ATOM_V, REST_V), EPI_M, EPI_N, RestM, RestN, RestL)
+        bSG_sC, bSG_gC = cpasync.tma_partition(
+            tma_atom_c,
+            0,
+            cute.make_layout(1),
+            sC_for_tma_partition,
+            gC_for_tma_partition,
+        )
+        return tma_atom_c, bSG_sC, bSG_gC
+
+    @staticmethod
+    def _compute_stages(
+        tiled_mma: cute.TiledMma,
+        mma_tiler_mnk: Tuple[int, int, int],
+        a_dtype: Type[cutlass.Numeric],
+        b_dtype: Type[cutlass.Numeric],
+        epi_tile: cute.Tile,
+        c_dtype: Type[cutlass.Numeric],
+        c_layout: utils.LayoutEnum,
+        sf_dtype: Type[cutlass.Numeric],
+        sf_vec_size: int,
+        smem_capacity: int,
+        occupancy: int,
+    ) -> Tuple[int, int, int]:
+        """Computes the number of stages for A/B/C operands based on heuristics.
+
+        :param tiled_mma: The tiled MMA object defining the core computation.
+        :type tiled_mma: cute.TiledMma
+        :param mma_tiler_mnk: The shape (M, N, K) of the MMA tiler.
+        :type mma_tiler_mnk: tuple[int, int, int]
+        :param a_dtype: Data type of operand A.
+        :type a_dtype: type[cutlass.Numeric]
+        :param b_dtype: Data type of operand B.
+        :type b_dtype: type[cutlass.Numeric]
+        :param epi_tile: The epilogue tile shape.
+        :type epi_tile: cute.Tile
+        :param c_dtype: Data type of operand C (output).
+        :type c_dtype: type[cutlass.Numeric]
+        :param c_layout: Layout enum of operand C.
+        :type c_layout: utils.LayoutEnum
+        :param sf_dtype: Data type of Scale factor.
+        :type sf_dtype: type[cutlass.Numeric]
+        :param sf_vec_size: Scale factor vector size.
+        :type sf_vec_size: int
+        :param smem_capacity: Total available shared memory capacity in bytes.
+        :type smem_capacity: int
+        :param occupancy: Target number of CTAs per SM (occupancy).
+        :type occupancy: int
+
+        :return: A tuple containing the computed number of stages for:
+                 (ACC stages, A/B operand stages, C stages)
+        :rtype: tuple[int, int, int]
+        """
+        # ACC stages
+        num_acc_stage = 1 if mma_tiler_mnk[1] == 256 else 2
+
+        # Default C stages
+        num_c_stage = 2
+
+        # Calculate smem layout and size for one stage of A, B, SFA, SFB and C
+        a_smem_layout_stage_one = sm100_utils.make_smem_layout_a(
+            tiled_mma,
+            mma_tiler_mnk,
+            a_dtype,
+            1,  # a tmp 1 stage is provided
+        )
+        b_smem_layout_staged_one = sm100_utils.make_smem_layout_b(
+            tiled_mma,
+            mma_tiler_mnk,
+            b_dtype,
+            1,  # a tmp 1 stage is provided
+        )
+        sfa_smem_layout_staged_one = blockscaled_utils.make_smem_layout_sfa(
+            tiled_mma,
+            mma_tiler_mnk,
+            sf_vec_size,
+            1,  # a tmp 1 stage is provided
+        )
+        sfb_smem_layout_staged_one = blockscaled_utils.make_smem_layout_sfb(
+            tiled_mma,
+            mma_tiler_mnk,
+            sf_vec_size,
+            1,  # a tmp 1 stage is provided
+        )
+
+        c_smem_layout_staged_one = sm100_utils.make_smem_layout_epi(
+            c_dtype,
+            c_layout,
+            epi_tile,
+            1,
+        )
+
+        ab_bytes_per_stage = (
+            cute.size_in_bytes(a_dtype, a_smem_layout_stage_one)
+            + cute.size_in_bytes(b_dtype, b_smem_layout_staged_one)
+            + cute.size_in_bytes(sf_dtype, sfa_smem_layout_staged_one)
+            + cute.size_in_bytes(sf_dtype, sfb_smem_layout_staged_one)
+        )
+        mbar_helpers_bytes = 1024
+        c_bytes_per_stage = cute.size_in_bytes(c_dtype, c_smem_layout_staged_one)
+        c_bytes = c_bytes_per_stage * num_c_stage
+
+        # Calculate A/B/SFA/SFB stages:
+        # Start with total smem per CTA (capacity / occupancy)
+        # Subtract reserved bytes and initial C stages bytes
+        # Divide remaining by bytes needed per A/B/SFA/SFB stage
+        num_ab_stage = (
+            smem_capacity // occupancy - (mbar_helpers_bytes + c_bytes)
+        ) // ab_bytes_per_stage
+
+        # Refine epilogue stages:
+        # Calculate remaining smem after allocating for A/B/SFA/SFB stages and reserved bytes
+        # Add remaining unused smem to epilogue
+        num_c_stage += (
+            smem_capacity
+            - occupancy * ab_bytes_per_stage * num_ab_stage
+            - occupancy * (mbar_helpers_bytes + c_bytes)
+        ) // (occupancy * c_bytes_per_stage)
+
+        return num_acc_stage, num_ab_stage, num_c_stage
+
+    @staticmethod
+    def _compute_grid(
+        c: cute.Tensor,
+        cta_tile_shape_mnk: Tuple[int, int, int],
+        cluster_shape_mn: Tuple[int, int],
+        max_active_clusters: cutlass.Constexpr,
+    ) -> Tuple[utils.PersistentTileSchedulerParams, Tuple[int, int, int]]:
+        """Use persistent tile scheduler to compute the grid size for the output tensor C.
+
+        :param c: The output tensor C
+        :type c: cute.Tensor
+        :param cta_tile_shape_mnk: The shape (M, N, K) of the CTA tile.
+        :type cta_tile_shape_mnk: tuple[int, int, int]
+        :param cluster_shape_mn: Shape of each cluster in M, N dimensions.
+        :type cluster_shape_mn: tuple[int, int]
+        :param max_active_clusters: Maximum number of active clusters.
+        :type max_active_clusters: cutlass.Constexpr
+
+        :return: A tuple containing:
+            - tile_sched_params: Parameters for the persistent tile scheduler.
+            - grid: Grid shape for kernel launch.
+        :rtype: Tuple[utils.PersistentTileSchedulerParams, tuple[int, int, int]]
+        """
+        c_shape = cute.slice_(cta_tile_shape_mnk, (None, None, 0))
+        gc = cute.zipped_divide(c, tiler=c_shape)
+        num_ctas_mnl = gc[(0, (None, None, None))].shape
+        cluster_shape_mnl = (*cluster_shape_mn, 1)
+
+        tile_sched_params = utils.PersistentTileSchedulerParams(
+            num_ctas_mnl, cluster_shape_mnl
+        )
+        grid = utils.StaticPersistentTileScheduler.get_grid_shape(
+            tile_sched_params, max_active_clusters
+        )
+
+        return tile_sched_params, grid
+
+    @staticmethod
+    def is_valid_dtypes_and_scale_factor_vec_size(
+        ab_dtype: Type[cutlass.Numeric],
+        sf_dtype: Type[cutlass.Numeric],
+        sf_vec_size: int,
+        c_dtype: Type[cutlass.Numeric],
+    ) -> bool:
+        """
+        Check if the dtypes and sf_vec_size are valid combinations
+
+        :param ab_dtype: The data type of the A and B operands
+        :type ab_dtype: Type[cutlass.Numeric]
+        :param sf_dtype: The data type of the scale factor
+        :type sf_dtype: Type[cutlass.Numeric]
+        :param sf_vec_size: The vector size of the scale factor
+        :type sf_vec_size: int
+        :param c_dtype: The data type of the output tensor
+        :type c_dtype: Type[cutlass.Numeric]
+
+        :return: True if the dtypes and sf_vec_size are valid, False otherwise
+        :rtype: bool
+        """
+        is_valid = True
+
+        # Check valid ab_dtype
+        if ab_dtype not in {
+            cutlass.Float4E2M1FN,
+            cutlass.Float8E5M2,
+            cutlass.Float8E4M3FN,
+        }:
+            is_valid = False
+
+        # Check valid sf_vec_size
+        if sf_vec_size not in {16, 32}:
+            is_valid = False
+
+        # Check valid sf_dtype
+        if sf_dtype not in {cutlass.Float8E8M0FNU, cutlass.Float8E4M3FN}:
+            is_valid = False
+
+        # Check valid sf_dtype and sf_vec_size combinations
+        if sf_dtype == cutlass.Float8E4M3FN and sf_vec_size == 32:
+            is_valid = False
+        if ab_dtype in {cutlass.Float8E5M2, cutlass.Float8E4M3FN} and sf_vec_size == 16:
+            is_valid = False
+
+        # Check valid c_dtype
+        if c_dtype not in {
+            cutlass.Float32,
+            cutlass.Float16,
+            cutlass.BFloat16,
+            cutlass.Float8E5M2,
+            cutlass.Float8E4M3FN,
+        }:
+            is_valid = False
+
+        return is_valid
+
+    @staticmethod
+    def is_valid_layouts(
+        ab_dtype: Type[cutlass.Numeric],
+        c_dtype: Type[cutlass.Numeric],
+        a_major: str,
+        b_major: str,
+        c_major: str,
+    ) -> bool:
+        """
+        Check if layouts and dtypes are valid combinations
+
+        :param ab_dtype: The data type of the A and B operands
+        :type ab_dtype: Type[cutlass.Numeric]
+        :param c_dtype: The data type of the output tensor
+        :type c_dtype: Type[cutlass.Numeric]
+        :param a_major: The major dimension of the A tensor
+        :type a_major: str
+        :param b_major: The major dimension of the B tensor
+        :type b_major: str
+        :param c_major: The major dimension of the C tensor
+        :type c_major: str
+
+        :return: True if the layouts are valid, False otherwise
+        :rtype: bool
+        """
+        is_valid = True
+
+        if ab_dtype is cutlass.Float4E2M1FN and not (a_major == "k" and b_major == "k"):
+            is_valid = False
+        return is_valid
+
+    @staticmethod
+    def is_valid_mma_tiler_and_cluster_shape(
+        mma_tiler_mn: Tuple[int, int],
+        cluster_shape_mn: Tuple[int, int],
+    ) -> bool:
+        """
+        Check if the mma tiler and cluster shape are valid
+
+        :param mma_tiler_mn: The (M, N) shape of the MMA instruction tiler
+        :type mma_tiler_mn: Tuple[int, int]
+        :param cluster_shape_mn: The (ClusterM, ClusterN) shape of the CTA cluster
+        :type cluster_shape_mn: Tuple[int, int]
+
+        :return: True if the mma tiler and cluster shape are valid, False otherwise
+        :rtype: bool
+        """
+        is_valid = True
+        # Skip invalid mma tile shape
+        if mma_tiler_mn[0] not in [128, 256]:
+            is_valid = False
+        if mma_tiler_mn[1] not in [64, 128, 192, 256]:
+            is_valid = False
+        # Skip illegal cluster shape
+        if cluster_shape_mn[0] % (2 if mma_tiler_mn[0] == 256 else 1) != 0:
+            is_valid = False
+        # Skip invalid cluster shape
+        is_power_of_2 = lambda x: x > 0 and (x & (x - 1)) == 0
+        if (
+            cluster_shape_mn[0] * cluster_shape_mn[1] > 16
+            or cluster_shape_mn[0] <= 0
+            or cluster_shape_mn[1] <= 0
+            # Special cluster shape check for scale factor multicasts.
+            # Due to limited size of scale factors, we can't multicast among more than 4 CTAs.
+            or cluster_shape_mn[0] > 4
+            or cluster_shape_mn[1] > 4
+            or not is_power_of_2(cluster_shape_mn[0])
+            or not is_power_of_2(cluster_shape_mn[1])
+        ):
+            is_valid = False
+        return is_valid
+
+    @staticmethod
+    def is_valid_tensor_alignment(
+        m: int,
+        n: int,
+        k: int,
+        l: int,
+        ab_dtype: Type[cutlass.Numeric],
+        c_dtype: Type[cutlass.Numeric],
+        a_major: str,
+        b_major: str,
+        c_major: str,
+    ) -> bool:
+        """
+        Check if the tensor alignment is valid
+
+        :param m: The number of rows in the A tensor
+        :type m: int
+        :param n: The number of columns in the B tensor
+        :type n: int
+        :param k: The number of columns in the A tensor
+        :type k: int
+        :param l: The number of columns in the C tensor
+        :type l: int
+        :param ab_dtype: The data type of the A and B operands
+        :type ab_dtype: Type[cutlass.Numeric]
+        :param c_dtype: The data type of the output tensor
+        :type c_dtype: Type[cutlass.Numeric]
+        :param a_major: The major axis of the A tensor
+        :type a_major: str
+        :param b_major: The major axis of the B tensor
+        :type b_major: str
+        :param c_major: The major axis of the C tensor
+        :type c_major: str
+
+        :return: True if the problem shape is valid, False otherwise
+        :rtype: bool
+        """
+        is_valid = True
+
+        def check_contigous_16B_alignment(dtype, is_mode0_major, tensor_shape):
+            major_mode_idx = 0 if is_mode0_major else 1
+            num_major_elements = tensor_shape[major_mode_idx]
+            num_contiguous_elements = 16 * 8 // dtype.width
+            return num_major_elements % num_contiguous_elements == 0
+
+        if (
+            not check_contigous_16B_alignment(ab_dtype, a_major == "m", (m, k, l))
+            or not check_contigous_16B_alignment(ab_dtype, b_major == "n", (n, k, l))
+            or not check_contigous_16B_alignment(c_dtype, c_major == "m", (m, n, l))
+        ):
+            is_valid = False
+        return is_valid
+
+    @staticmethod
+    def can_implement(
+        ab_dtype: Type[cutlass.Numeric],
+        sf_dtype: Type[cutlass.Numeric],
+        sf_vec_size: int,
+        c_dtype: Type[cutlass.Numeric],
+        mma_tiler_mn: Tuple[int, int],
+        cluster_shape_mn: Tuple[int, int],
+        m: int,
+        n: int,
+        k: int,
+        l: int,
+        a_major: str,
+        b_major: str,
+        c_major: str,
+    ) -> bool:
+        """
+        Check if the gemm can be implemented
+
+        :param ab_dtype: The data type of the A and B operands
+        :type ab_dtype: Type[cutlass.Numeric]
+        :param sf_dtype: The data type of the scale factor tensor
+        :type sf_dtype: Type[cutlass.Numeric]
+        :param sf_vec_size: The vector size
+        :type sf_vec_size: int
+        :param c_dtype: The data type of the output tensor
+        :type c_dtype: Type[cutlass.Numeric]
+        :param mma_tiler_mn: The (M, N) shape of the MMA instruction tiler
+        :type mma_tiler_mn: Tuple[int, int]
+        :param cluster_shape_mn: The (ClusterM, ClusterN) shape of the CTA cluster
+        :type cluster_shape_mn: Tuple[int, int]
+        :param m: The number of rows in the A tensor
+        :type m: int
+        :param n: The number of columns in the B tensor
+        :type n: int
+        :param k: The number of columns in the A tensor
+        :type k: int
+        :param l: The number of columns in the C tensor
+        :type l: int
+        :param a_major: The major axis of the A tensor
+        :type a_major: str
+        :param b_major: The major axis of the B tensor
+        :type b_major: str
+        :param c_major: The major axis of the C tensor
+        :type c_major: str
+
+        :return: True if the gemm can be implemented, False otherwise
+        :rtype: bool
+        """
+        can_implement = True
+        # Skip unsupported types
+        if not Sm100BlockScaledPersistentDenseGemmKernel.is_valid_dtypes_and_scale_factor_vec_size(
+            ab_dtype, sf_dtype, sf_vec_size, c_dtype
+        ):
+            can_implement = False
+        # Skip unsupported layouts
+        if not Sm100BlockScaledPersistentDenseGemmKernel.is_valid_layouts(
+            ab_dtype, c_dtype, a_major, b_major, c_major
+        ):
+            can_implement = False
+        # Skip invalid mma tile shape and cluster shape
+        if not Sm100BlockScaledPersistentDenseGemmKernel.is_valid_mma_tiler_and_cluster_shape(
+            mma_tiler_mn, cluster_shape_mn
+        ):
+            can_implement = False
+        # Skip illegal problem shape for load/store alignment
+        if not Sm100BlockScaledPersistentDenseGemmKernel.is_valid_tensor_alignment(
+            m, n, k, l, ab_dtype, c_dtype, a_major, b_major, c_major
+        ):
+            can_implement = False
+        return can_implement
+
+
+@cute.jit
+def cvt_sf_MKL_to_M32x4xrm_K4xrk_L(
+    sf_ref_tensor: cute.Tensor,
+    sf_mma_tensor: cute.Tensor,
+):
+    """Convert scale factor tensor from MKL layout to mma specification M(32x4xrest_m)xK(4xrest_k)xL layout"""
+    # sf_mma_tensor has flatten shape (32, 4, rest_m, 4, rest_k, l)
+    # group to ((32, 4, rest_m), (4, rest_k), l)
+    sf_mma_tensor = cute.group_modes(sf_mma_tensor, 0, 3)
+    sf_mma_tensor = cute.group_modes(sf_mma_tensor, 1, 3)
+    for i in cutlass.range(cute.size(sf_ref_tensor)):
+        mkl_coord = sf_ref_tensor.layout.get_hier_coord(i)
+        sf_mma_tensor[mkl_coord] = sf_ref_tensor[mkl_coord]
diff --git a/torch/_inductor/kernel/vendored_templates/cutedsl/wrappers/__init__.py b/torch/_inductor/kernel/vendored_templates/cutedsl/wrappers/__init__.py
new file mode 100644
index 0000000000000..583f1f04b7c97
--- /dev/null
+++ b/torch/_inductor/kernel/vendored_templates/cutedsl/wrappers/__init__.py
@@ -0,0 +1,2 @@
+# CuTeDSL Cutlass API registrations for PyTorch Inductor.
+from . import dense_blockscaled_gemm_kernel  # noqa: F401
diff --git a/torch/_inductor/kernel/vendored_templates/cutedsl/wrappers/dense_blockscaled_gemm_kernel.py b/torch/_inductor/kernel/vendored_templates/cutedsl/wrappers/dense_blockscaled_gemm_kernel.py
new file mode 100644
index 0000000000000..91a7da1c8497e
--- /dev/null
+++ b/torch/_inductor/kernel/vendored_templates/cutedsl/wrappers/dense_blockscaled_gemm_kernel.py
@@ -0,0 +1,449 @@
+"""Kernel wrapper for dense_blockscaled_gemm_persistent vendored template."""
+
+from __future__ import annotations
+
+import itertools
+import logging
+from collections.abc import Callable, Generator  # noqa: TC003
+
+import cutlass_api
+from cutlass_api.arguments import GemmArguments  # noqa: TC002
+from cutlass_api.artifact import CompiledArtifact
+from cutlass_api.library import ScaleMode, ScaleSwizzleMode
+from cutlass_api.metadata import (
+    DenseTensorAttributes,
+    GemmOperandsMetadata,
+    KernelMetadata,
+    ScaledTensorAttributes,
+    Sm100DesignMetadata,
+)
+from cutlass_api.providers.cutedsl.kernel import CuteDslKernel
+from cutlass_api.providers.cutedsl.utils import get_max_active_clusters
+from cutlass_api.status import Status
+from cutlass_api.utils import (
+    ceil_div,
+    round_up,
+    strides_to_layout_string,
+    to_cuda_stream,
+    tuple_to_string,
+)
+
+
+log = logging.getLogger(__name__)
+
+
+try:
+    from ..dense_blockscaled_gemm_persistent import (  # pyrefly: ignore[missing-import]
+        Sm100BlockScaledPersistentDenseGemmKernel as BlockScaledGemmKernelImpl,
+    )
+except ImportError:
+    BlockScaledGemmKernelImpl = None  # type: ignore[misc, assignment]
+
+
+class VendoredDenseBlockScaledGemmKernel(CuteDslKernel):
+    """Wrapper for vendored dense blockscaled GEMM template for SM100 GPUs."""
+
+    def __init__(self, metadata: KernelMetadata):
+        super().__init__(metadata)
+
+        self.sf_vec_size = metadata.operands.A.mode[-1]
+        mma_tiler_mn = (metadata.design.tile_shape[0], metadata.design.tile_shape[1])
+        cluster_shape_mn = (
+            metadata.design.cluster_shape[0],
+            metadata.design.cluster_shape[1],
+        )
+
+        self.impl = BlockScaledGemmKernelImpl(  # pyrefly: ignore[not-callable]
+            self.sf_vec_size,
+            mma_tiler_mn,
+            cluster_shape_mn,
+        )
+        self.cluster_shape_mn = cluster_shape_mn
+
+    @staticmethod
+    def _major_modes(args):
+        """Extract major modes from arguments or operand metadata."""
+        import cutlass.utils as utils
+        from cutlass.cute.nvgpu.tcgen05 import OperandMajorMode
+
+        if args.A.stride[-2:].index(1) == 1:
+            a_major_mode = (OperandMajorMode.K, "k")
+        else:
+            a_major_mode = (OperandMajorMode.MN, "m")
+
+        if args.B.stride[-2:].index(1) == 0:
+            b_major_mode = (OperandMajorMode.K, "k")
+        else:
+            b_major_mode = (OperandMajorMode.MN, "n")
+
+        if args.out.stride[-2:].index(1) == 1:
+            out_layout = (utils.LayoutEnum.ROW_MAJOR, "n")
+        else:
+            out_layout = (utils.LayoutEnum.COL_MAJOR, "m")
+
+        return a_major_mode, b_major_mode, out_layout
+
+    def compile(self, args: GemmArguments, cc: int | None = None) -> CompiledArtifact:
+        import cutlass.cute as cute
+
+        stream = cute.runtime.make_fake_stream()
+        max_active_clusters = get_max_active_clusters(self.cluster_shape_mn)
+
+        # pyrefly: ignore[missing-attribute]
+        compiled_kernel = self.cute_compile(
+            self.impl,
+            args.A.tensor,
+            args.B.tensor,
+            args.A.scale.tensor,
+            args.B.scale.tensor,
+            args.out.tensor,
+            max_active_clusters,
+            stream,
+        )
+
+        return CompiledArtifact(compiled_kernel, self)
+
+    def _run(
+        self,
+        args: GemmArguments,
+        compiled_artifact: CompiledArtifact,
+        stream,
+        workspace=None,
+    ) -> None:
+        import torch
+
+        stream = to_cuda_stream(stream)
+        compiled_gemm = compiled_artifact.compiled_obj
+
+        # TVM FFI needs a torch.cuda.Stream, not a raw int handle
+        if isinstance(stream, int):
+            stream = torch.cuda.ExternalStream(stream)
+
+        self.cute_run(  # pyrefly: ignore[missing-attribute]
+            compiled_gemm,
+            args.A.tensor,
+            args.B.tensor,
+            args.A.scale.tensor,
+            args.B.scale.tensor,
+            args.out.tensor,
+            stream,
+        )
+
+    def get_workspace_size(self, args: GemmArguments) -> int:
+        return 0
+
+    def _supports(self, args: GemmArguments) -> Status:
+        import cutlass
+
+        m, n = args.out.shape[-2:]
+        k = args.A.shape[-1]
+        if args.A.element_type == cutlass.Float4E2M1FN:
+            k = k * 2
+        l = args.A.shape[0] if len(args.A.shape) == 3 else 1
+
+        expected_sf_k = round_up(ceil_div(k, self.sf_vec_size), 4)
+        expected_sfa_elements = l * m * expected_sf_k
+        expected_sfb_elements = l * n * expected_sf_k
+
+        if args.A.scale.numel() != expected_sfa_elements:
+            return Status.fail(
+                f"Scale factor A for tensor A of shape {args.A.shape} must have "
+                f"{expected_sfa_elements} elements, got {args.A.scale.numel()}."
+            )
+        if args.B.scale.numel() != expected_sfb_elements:
+            return Status.fail(
+                f"Scale factor B for tensor B of shape {args.B.shape} must have "
+                f"{expected_sfb_elements} elements, got {args.B.scale.numel()}."
+            )
+
+        return Status.success()
+
+    @staticmethod
+    def _is_valid_dtype_combo(ab_dtype, sf_dtype, sf_vec_size, out_dtype) -> bool:
+        """Validate dtype/scale-factor/vec-size combinations.
+
+        Matches constraints from the upstream kernel:
+          MXF8: Float8E5M2/Float8E4M3FN + Float8E8M0FNU, sf_vec_size=32
+          MXF4: Float4E2M1FN + Float8E8M0FNU, sf_vec_size=32
+          NVF4: Float4E2M1FN + Float8E8M0FNU/Float8E4M3FN, sf_vec_size=16
+        """
+        import cutlass
+
+        if ab_dtype not in {
+            cutlass.Float4E2M1FN,
+            cutlass.Float8E5M2,
+            cutlass.Float8E4M3FN,
+        }:
+            return False
+        if sf_vec_size not in {16, 32}:
+            return False
+        if sf_dtype not in {cutlass.Float8E8M0FNU, cutlass.Float8E4M3FN}:
+            return False
+        # Float8E4M3FN as SF only valid with sf_vec_size=16 (NVF4)
+        if sf_dtype == cutlass.Float8E4M3FN and sf_vec_size == 32:
+            return False
+        # Float8 AB types require sf_vec_size=32 (MXF8)
+        if ab_dtype in {cutlass.Float8E5M2, cutlass.Float8E4M3FN} and sf_vec_size == 16:
+            return False
+        if out_dtype not in {
+            cutlass.Float32,
+            cutlass.Float16,
+            cutlass.BFloat16,
+            cutlass.Float8E5M2,
+            cutlass.Float8E4M3FN,
+        }:
+            return False
+        return True
+
+    @staticmethod
+    def _is_valid_layouts(ab_dtype, a_major, b_major) -> bool:
+        """Float4E2M1FN (MXF4/NVF4) requires row-major (K-major) A and B."""
+        import cutlass
+
+        if ab_dtype is cutlass.Float4E2M1FN and not (a_major == "k" and b_major == "k"):
+            return False
+        return True
+
+    @staticmethod
+    def _valid_operands(operands: GemmOperandsMetadata, sf_vec_size: int) -> bool:
+        import cutlass
+
+        if operands.A.dtype != operands.B.dtype:
+            return False
+        if operands.A.scale.dtype != operands.B.scale.dtype:
+            return False
+        if operands.accumulator_type != cutlass.Float32:
+            return False
+
+        ab_dtype = operands.A.dtype
+        sf_dtype = operands.A.scale.dtype
+        out_dtype = operands.out.dtype
+
+        if not VendoredDenseBlockScaledGemmKernel._is_valid_dtype_combo(
+            ab_dtype, sf_dtype, sf_vec_size, out_dtype
+        ):
+            return False
+
+        (_, a_major), (_, b_major), _ = VendoredDenseBlockScaledGemmKernel._major_modes(
+            operands
+        )
+        if not VendoredDenseBlockScaledGemmKernel._is_valid_layouts(
+            ab_dtype, a_major, b_major
+        ):
+            return False
+
+        return True
+
+    @staticmethod
+    def _metadata_operand_combinations() -> Generator[GemmOperandsMetadata, None, None]:
+        import cutlass
+
+        ab_dtypes = [cutlass.Float8E5M2, cutlass.Float8E4M3FN, cutlass.Float4E2M1FN]
+        out_dtypes = [
+            cutlass.Float32,
+            cutlass.Float16,
+            cutlass.BFloat16,
+            cutlass.Float8E4M3FN,
+            cutlass.Float8E5M2,
+        ]
+        sf_dtypes = [cutlass.Float8E8M0FNU, cutlass.Float8E4M3FN]
+        scale_modes = [ScaleMode.Blockwise1x16, ScaleMode.Blockwise1x32]
+
+        row_major = (0, 0, 1)
+        col_major = (0, 1, 0)
+        alignment_bytes = 16
+
+        def major_str_a(stride):
+            return "k" if stride == row_major else "m"
+
+        def major_str_b(stride):
+            return "n" if stride == row_major else "k"
+
+        for ab_dtype, sf_dtype, scale_mode, out_dtype in itertools.product(
+            ab_dtypes, sf_dtypes, scale_modes, out_dtypes
+        ):
+            sf_vec_size = scale_mode[-1]
+            if not VendoredDenseBlockScaledGemmKernel._is_valid_dtype_combo(
+                ab_dtype, sf_dtype, sf_vec_size, out_dtype
+            ):
+                continue
+
+            for stride_A, stride_B, stride_out in itertools.product(
+                [row_major, col_major], repeat=3
+            ):
+                a_major = major_str_a(stride_A)
+                b_major = major_str_b(stride_B)
+
+                if not VendoredDenseBlockScaledGemmKernel._is_valid_layouts(
+                    ab_dtype, a_major, b_major
+                ):
+                    continue
+
+                ab_div = alignment_bytes * 8 // ab_dtype.width
+                out_div = alignment_bytes * 8 // out_dtype.width
+                sf_div = alignment_bytes * 8 // sf_dtype.width
+
+                yield GemmOperandsMetadata(
+                    A=ScaledTensorAttributes(
+                        base=DenseTensorAttributes(
+                            dtype=ab_dtype,
+                            stride=stride_A,
+                            divisibility=ab_div,
+                        ),
+                        scale=DenseTensorAttributes(
+                            dtype=sf_dtype,
+                            stride=None,
+                            divisibility=sf_div,
+                        ),
+                        mode=scale_mode,
+                        swizzle=ScaleSwizzleMode.Swizzle32x4x4,
+                    ),
+                    B=ScaledTensorAttributes(
+                        base=DenseTensorAttributes(
+                            dtype=ab_dtype,
+                            stride=stride_B,
+                            divisibility=ab_div,
+                        ),
+                        scale=DenseTensorAttributes(
+                            dtype=sf_dtype,
+                            stride=None,
+                            divisibility=sf_div,
+                        ),
+                        mode=scale_mode,
+                        swizzle=ScaleSwizzleMode.Swizzle32x4x4,
+                    ),
+                    out=DenseTensorAttributes(
+                        dtype=out_dtype,
+                        stride=stride_out,
+                        divisibility=out_div,
+                    ),
+                    accumulator_type=cutlass.Float32,
+                )
+
+    @staticmethod
+    def _valid_metadata(metadata: KernelMetadata) -> bool:
+        scale_vec = metadata.operands.A.mode
+
+        if len(scale_vec) > 1:
+            for i in range(len(scale_vec) - 1):
+                if scale_vec[i] != 1:
+                    return False
+
+        sf_vec_size = scale_vec[-1]
+        if not VendoredDenseBlockScaledGemmKernel._valid_operands(
+            metadata.operands, sf_vec_size
+        ):
+            return False
+
+        design = metadata.design
+        if not isinstance(design, Sm100DesignMetadata):
+            return False
+
+        cm, cn, _ = design.cluster_shape
+        if cm <= 0 or cn <= 0:
+            return False
+        if cm * cn > 16:
+            return False
+        if cm & (cm - 1) != 0 or cn & (cn - 1) != 0:
+            return False
+        # SF multicast constraint: cluster dims <=4
+        if cm > 4 or cn > 4:
+            return False
+
+        tile_m, tile_n, _ = design.tile_shape
+        if tile_m not in [128, 256]:
+            return False
+        if tile_n not in [64, 128, 192, 256]:
+            return False
+        use_2cta = tile_m == 256
+        if use_2cta and cm % 2 != 0:
+            return False
+
+        if metadata.epilogue is not None:
+            return False
+
+        return True
+
+    @staticmethod
+    def generate_kernels(
+        metadata_filter: Callable[[KernelMetadata], bool],
+        epilogue_args=None,
+        cc: int | None = None,
+    ) -> list[VendoredDenseBlockScaledGemmKernel]:
+        if cc is not None and cc not in [100, 101, 103]:
+            return []
+        if epilogue_args is not None:
+            return []
+
+        design_params = {
+            "use_2cta_mma": [True],
+            "tile_shape": [
+                (M, N, 256) for M in [128, 256] for N in [64, 128, 192, 256]
+            ],
+            "cluster_shape": [(M, N, 1) for M in [1, 2, 4] for N in [1, 2, 4]],
+            "use_tma_store": [True],
+        }
+
+        param_names = list(design_params.keys())
+        param_values = [design_params[name] for name in param_names]
+
+        kernel_list = []
+
+        for (
+            operands
+        ) in VendoredDenseBlockScaledGemmKernel._metadata_operand_combinations():
+            # pyrefly: ignore[no-matching-overload]
+            for values in itertools.product(*param_values):
+                design = Sm100DesignMetadata(**dict(zip(param_names, values)))
+
+                kernel_name = (
+                    "inductor_vendored.DenseBlockScaledGemmKernel_sm100_"
+                    "{layout}_A{A}_B{B}_out{out}_SFA{SFA}_SFB{SFB}_"
+                    "acc{acc}_scale{scale_mode}_swizzle{scale_swizzle}_"
+                    "{num_cta}cta_cluster{cluster}_tile{tile}"
+                    "{_tma_store}"
+                ).format(
+                    layout=strides_to_layout_string(
+                        operands.A.stride,
+                        operands.B.stride,
+                        operands.out.stride,
+                    ),
+                    A=operands.A.dtype,
+                    B=operands.B.dtype,
+                    out=operands.out.dtype,
+                    SFA=operands.A.scale.dtype,
+                    SFB=operands.B.scale.dtype,
+                    acc=operands.accumulator_type,
+                    scale_mode=operands.A.mode,
+                    scale_swizzle=operands.A.swizzle,
+                    num_cta="2" if design.use_2cta_mma else "1",
+                    cluster=tuple_to_string(design.cluster_shape),
+                    tile=tuple_to_string(design.tile_shape),
+                    _tma_store="_tma_store" if design.use_tma_store else "",
+                )
+
+                metadata = KernelMetadata(
+                    operands=operands,
+                    design=design,
+                    kernel_name=kernel_name,
+                    kernel_class=VendoredDenseBlockScaledGemmKernel,
+                    min_cc=100,
+                    epilogue=None,
+                )
+
+                if VendoredDenseBlockScaledGemmKernel._valid_metadata(metadata):
+                    if metadata_filter is None or metadata_filter(metadata):
+                        kernel_list.append(VendoredDenseBlockScaledGemmKernel(metadata))
+
+        log.debug(
+            "Generated %d DenseBlockScaledGemmKernel configurations",
+            len(kernel_list),
+        )
+        return kernel_list
+
+
+# Only register if kernel implementation is available
+if BlockScaledGemmKernelImpl is not None:
+    cutlass_api.providers.cutedsl.CuTeDSLProvider.register(
+        VendoredDenseBlockScaledGemmKernel
+    )
diff --git a/torch/_inductor/kernel_inputs.py b/torch/_inductor/kernel_inputs.py
index 5fe24859b32b9..596d6d688d1a8 100644
--- a/torch/_inductor/kernel_inputs.py
+++ b/torch/_inductor/kernel_inputs.py
@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 from abc import ABC, abstractmethod
-from typing import Any, Optional, TYPE_CHECKING, Union
+from typing import Any, TYPE_CHECKING
 
 import torch
 import torch._inductor.config
@@ -27,8 +27,8 @@ class KernelInputs(ABC):
     def __init__(
         self,
         input_nodes: list[Any],
-        scalars: Optional[dict[str, Union[float, int]]] = None,
-        out_dtype: Optional[torch.dtype] = None,
+        scalars: dict[str, float | int] | None = None,
+        out_dtype: torch.dtype | None = None,
     ):
         """
         Initialize with a tuple of input nodes.
@@ -38,12 +38,12 @@ def __init__(
             out_dtype: Optional output dtype to store
         """
         self._input_nodes = input_nodes
-        self._device_name: Optional[str] = None
+        self._device_name: str | None = None
         self._scalars = scalars if scalars is not None else {}
         self._out_dtype = out_dtype
         assert len(input_nodes) > 0, "Expected at least one input node"
 
-    def nodes(self, reorder: Optional[Sequence[int]] = None) -> list[Any]:
+    def nodes(self, reorder: Sequence[int] | None = None) -> list[Any]:
         """
         Return the stored input nodes, optionally reordered.
 
@@ -72,7 +72,7 @@ def count(self) -> int:
         return len(self._input_nodes)
 
     @property
-    def device_type(self) -> Optional[str]:
+    def device_type(self) -> str | None:
         """
         Get the device type of the first node.
 
@@ -91,7 +91,7 @@ def device(self) -> torch.device:
         """
         return self._input_nodes[0].get_device()
 
-    def device_name(self) -> Optional[str]:
+    def device_name(self) -> str | None:
         """
         Get the device name information.
 
@@ -122,10 +122,7 @@ def shapes_hinted(self) -> tuple[tuple[int, ...], ...]:
             A tuple of shape tuples with integer hints for each input node
         """
         return tuple(
-            V.graph.sizevars.size_hints(
-                node.get_size(),
-                fallback=torch._inductor.config.unbacked_symint_fallback,
-            )
+            V.graph.sizevars.optimization_hints(node.get_size())
             for node in self._input_nodes
         )
 
@@ -146,10 +143,7 @@ def strides_hinted(self) -> tuple[tuple[int, ...], ...]:
             A tuple of stride tuples with integer hints for each input node
         """
         return tuple(
-            V.graph.sizevars.size_hints(
-                node.get_stride(),
-                fallback=torch._inductor.config.unbacked_symint_fallback,
-            )
+            V.graph.sizevars.optimization_hints(node.get_stride())
             for node in self._input_nodes
         )
 
@@ -183,7 +177,7 @@ def out_dtype(self) -> torch.dtype:
             The output dtype
         """
 
-    def get_scalar(self, name: str) -> Union[float, int]:
+    def get_scalar(self, name: str) -> float | int:
         """
         Get the scalar value for a given name.
 
@@ -216,8 +210,8 @@ class MMKernelInputs(KernelInputs):
     def __init__(
         self,
         input_nodes: list[Any],
-        scalars: Optional[dict[str, Union[float, int]]] = None,
-        out_dtype: Optional[torch.dtype] = None,
+        scalars: dict[str, float | int] | None = None,
+        out_dtype: torch.dtype | None = None,
         mat1_idx: int = -2,
         mat2_idx: int = -1,
     ):
diff --git a/torch/_inductor/kernel_template_choice.py b/torch/_inductor/kernel_template_choice.py
index 8f90157c6c1a0..cd68cfc62ab4c 100644
--- a/torch/_inductor/kernel_template_choice.py
+++ b/torch/_inductor/kernel_template_choice.py
@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import Any, Optional, TYPE_CHECKING, Union
+from typing import Any, TYPE_CHECKING
 
 from .template_heuristics.params import DictKernelTemplateParams
 
@@ -25,7 +25,7 @@ class KernelTemplateChoice:
 
     def __init__(
         self,
-        template: Union[KernelTemplate, ExternKernelChoice],
+        template: KernelTemplate | ExternKernelChoice,
         params: KernelTemplateParams,
         extra_kwargs: dict[str, Any],
         layout: Layout,
@@ -39,7 +39,7 @@ def __init__(
         self.annotations: dict[str, Any] = {"ktc": self}
 
     @property
-    def choice(self) -> Optional[ChoiceCaller]:
+    def choice(self) -> ChoiceCaller | None:
         """
         Lazily evaluate and return the ChoiceCaller for this template choice.
 
@@ -65,7 +65,7 @@ def choice(self) -> Optional[ChoiceCaller]:
 
 
 def make_ktc_generator(
-    template: Union[KernelTemplate, ExternKernelChoice],
+    template: KernelTemplate | ExternKernelChoice,
     cs: Generator[KernelTemplateParams, None, None],
     extra_kwargs: dict[str, Any],
     overrides: dict[str, Any],
diff --git a/torch/_inductor/lookup_table/choices.py b/torch/_inductor/lookup_table/choices.py
index 46e54180114ab..03fdab47e11e7 100644
--- a/torch/_inductor/lookup_table/choices.py
+++ b/torch/_inductor/lookup_table/choices.py
@@ -3,7 +3,7 @@
 import copy
 import logging
 from functools import lru_cache
-from typing import Any, Optional, TYPE_CHECKING, Union
+from typing import Any, TYPE_CHECKING
 
 import torch
 from torch._inductor import config
@@ -40,7 +40,7 @@ def _get_lookup_table(self) -> dict[str, list[dict[str, Any]]]:
 
     @staticmethod
     @lru_cache
-    def _get_device_key(device: torch.device) -> Optional[str]:
+    def _get_device_key(device: torch.device) -> str | None:
         """
         Generate a device key for lookup table indexing.
         For CPU devices, returns None.
@@ -89,7 +89,7 @@ def _generate_kernel_inputs_key(kernel_inputs: KernelInputs) -> str:
 
     def make_lookup_key(
         self, kernel_inputs: KernelInputs, op_name: str, include_device: bool = False
-    ) -> Optional[str]:
+    ) -> str | None:
         """
         Create a flattened lookup key from kernel inputs and operation name.
         Override this method to customize key generation.
@@ -124,7 +124,7 @@ def make_lookup_key(
 
     def make_lookup_key_variants(
         self, kernel_inputs: KernelInputs, op_name: str
-    ) -> tuple[Optional[str], Optional[str]]:
+    ) -> tuple[str | None, str | None]:
         """
         Generate both device-specific and device-agnostic lookup keys.
         Override this method to customize key variant generation.
@@ -147,7 +147,7 @@ def make_lookup_key_variants(
     def _entry_is_valid(
         cfg: dict[str, Any],
         template_id: str,
-        template_hash_map: Optional[dict[str, Optional[str]]],
+        template_hash_map: dict[str, str | None] | None,
     ) -> bool:
         """
         Check if a config entry is valid based on template hash validation.
@@ -208,7 +208,7 @@ def lookup_template_configs(
         kernel_inputs: KernelInputs,
         op_name: str,
         template_uids: list[str],
-        template_hash_map: Optional[dict[str, Optional[str]]] = None,
+        template_hash_map: dict[str, str | None] | None = None,
     ) -> dict[str, list[dict[str, Any]]]:
         """
         Unified function to look up template configurations for multiple templates.
@@ -313,9 +313,9 @@ def _finalize_template_configs(
         self,
         template_choices: dict[str, Generator[KernelTemplateChoice, None, None]],
         kernel_inputs: KernelInputs,
-        templates: list[Union[KernelTemplate, ExternKernelChoice]],
+        templates: list[KernelTemplate | ExternKernelChoice],
         op_name: str,
-        kwarg_overrides: Optional[dict[str, dict[str, Any]]] = None,
+        kwarg_overrides: dict[str, dict[str, Any]] | None = None,
     ) -> list[KernelTemplateChoice]:
         """Check lookup table for hits, use those if found, otherwise fall back to parent."""
         # 1. Collect template src_hashes for validation
@@ -364,9 +364,9 @@ def _fallback(
         self,
         template_choices: dict[str, Generator[KernelTemplateChoice, None, None]],
         kernel_inputs: KernelInputs,
-        templates: list[Union[KernelTemplate, ExternKernelChoice]],
+        templates: list[KernelTemplate | ExternKernelChoice],
         op_name: str,
-        kwarg_overrides: Optional[dict[str, dict[str, Any]]] = None,
+        kwarg_overrides: dict[str, dict[str, Any]] | None = None,
     ) -> list[KernelTemplateChoice]:
         """Fallback to parent if no lookup table or no matches."""
         # NOTE: this is broken out, so that subclasses are able to override this
@@ -383,7 +383,7 @@ def _fallback(
     def _create_lookup_choices(
         self,
         lookup_results: dict[str, list[dict[str, Any]]],
-        templates: list[Union[KernelTemplate, ExternKernelChoice]],
+        templates: list[KernelTemplate | ExternKernelChoice],
         kernel_inputs: KernelInputs,
         op_name: str,
     ) -> list[KernelTemplateChoice]:
@@ -404,11 +404,14 @@ def _create_lookup_choices(
                 continue
 
             # For each lookup config, create a KTC with the override kwargs
+            # Start from base params (which include derived options like OUT_DTYPE)
+            # and let the lookup config override them
+            base_kwargs = base_ktc.params.to_kwargs()
             for c in configs:
+                merged = {**base_kwargs, **c}
                 lookup_ktc = KernelTemplateChoice(
                     template=base_ktc.template,
-                    # use the ones from the lookup table
-                    params=DictKernelTemplateParams(c),
+                    params=DictKernelTemplateParams(merged),
                     extra_kwargs=base_ktc.extra_kwargs,
                     layout=base_ktc.layout,
                     inputs=base_ktc.inputs,
diff --git a/torch/_inductor/loop_body.py b/torch/_inductor/loop_body.py
index 3921aa955a836..1ed47c0c2a32d 100644
--- a/torch/_inductor/loop_body.py
+++ b/torch/_inductor/loop_body.py
@@ -6,13 +6,14 @@
 import itertools
 import re
 from enum import auto, Enum
-from typing import Any, NamedTuple, Optional, TYPE_CHECKING, TypeVar
+from typing import Any, NamedTuple, TYPE_CHECKING, TypeVar
 
 import sympy
 
 import torch.fx
 from torch._dynamo.utils import identity
 from torch.fx.proxy import Scope, TracerBase
+from torch.utils._sympy.functions import Mod
 from torch.utils._sympy.symbol import SymT
 
 from . import config, dependencies
@@ -73,8 +74,8 @@ def __init__(self):
 
 class MemoryEntry(NamedTuple):
     index_name: str  # LoopBody.indexing_exprs[index_name]
-    buffer_name: Optional[str]
-    mode: Optional[str]  # V.ops.store(..., mode=mode)
+    buffer_name: str | None
+    mode: str | None  # V.ops.store(..., mode=mode)
 
 
 class MemoryUsageType(Enum):
@@ -268,7 +269,7 @@ def new_body(*indices: Sequence[sympy.Expr]) -> Any:
             reduce_idx = index[len(iter_size) :]
 
             new_iter_idx = list(iter_idx)
-            new_iter_idx[dimension] = iter_idx[dimension] % original_range
+            new_iter_idx[dimension] = Mod(iter_idx[dimension], original_range)
 
             return old_body(new_iter_idx, reduce_idx)
 
@@ -427,8 +428,8 @@ def add_index_expr(
         self,
         expr: sympy.Expr,
         mtype: MemoryUsageType,
-        buffer_name: Optional[str] = None,
-        mode: Optional[str] = None,
+        buffer_name: str | None = None,
+        mode: str | None = None,
     ):
         name = self.indexing_exprs_name.get(expr)
         if not name:
@@ -532,7 +533,12 @@ def __init__(self, body: LoopBody, fn: Callable[..., Any], args: list[Any]):
         from .index_propagation import IndexPropagation
 
         handler: Any = CountOps(
-            CaptureIndexing(proxy_ops, body, tracer),
+            CaptureIndexing(
+                # pyrefly: ignore[bad-argument-type]
+                proxy_ops,
+                body,
+                tracer,
+            ),
             body.op_counts,
         )
         if config.constant_and_index_propagation:
@@ -689,8 +695,8 @@ def bucketize(
         boundary_indices: T,
         indexing_dtype: torch.dtype,
         right: bool,
-        sorter: Optional[tuple[str, sympy.Expr]] = None,
-        sorter_indices: Optional[T] = None,
+        sorter: tuple[str, sympy.Expr] | None = None,
+        sorter_indices: T | None = None,
     ) -> T:
         """
         See [Note: Inductor bucketize op]
diff --git a/torch/_inductor/lowering.py b/torch/_inductor/lowering.py
index 0c4c9638c036b..0b1c3f3593204 100644
--- a/torch/_inductor/lowering.py
+++ b/torch/_inductor/lowering.py
@@ -9,11 +9,11 @@
 import math
 import operator
 import os
-import textwrap
+import sys
 import warnings
 from collections import defaultdict
 from collections.abc import Callable, Collection, Iterable, Sequence
-from typing import Any, cast, Optional, TYPE_CHECKING, TypeGuard, TypeVar, Union
+from typing import Any, cast, TYPE_CHECKING, TypeGuard, TypeVar
 from typing_extensions import ParamSpec
 from unittest.mock import patch
 
@@ -27,6 +27,7 @@
 from torch._higher_order_ops.associative_scan import associative_scan_op
 from torch._higher_order_ops.triton_kernel_wrap import triton_kernel_wrapper_mutation
 from torch._library.fake_class_registry import FakeScriptObject
+from torch._library.opaque_object import is_opaque_value
 from torch._library.utils import get_layout_constraint_tag
 from torch._prims_common import (
     canonicalize_dim,
@@ -111,11 +112,11 @@
 )
 
 log = logging.getLogger(__name__)
-lowerings: dict[Union[Callable[..., Any], str], Callable[..., Any]] = {}
+lowerings: dict[Callable[..., Any] | str, Callable[..., Any]] = {}
+# User-registered lowerings that take priority over built-in lowerings.
+user_lowerings: dict[torch._ops.OpOverload, Callable[..., Any]] = {}
 # Use maybe_layout_constraints to access this dict, we lazily register tag-based layout constraints
-_maybe_layout_constraints: dict[
-    torch._ops.OpOverload, Optional[Callable[..., Any]]
-] = {}
+_maybe_layout_constraints: dict[torch._ops.OpOverload, Callable[..., Any] | None] = {}
 fallbacks = OrderedSet[torch._ops.OpOverload]()
 aten = torch.ops.aten
 tr_c10d = torch.ops.tr_c10d
@@ -167,7 +168,7 @@ def group_foreach_args(
     return out
 
 
-def maybe_layout_constraints(fn: Callable[..., Any]) -> Optional[Callable[..., Any]]:
+def maybe_layout_constraints(fn: Callable[..., Any]) -> Callable[..., Any] | None:
     """Get layout constraints. Returns None if there are no layout constraints."""
     if not isinstance(fn, torch._ops.OpOverload):
         # Only OpOverloads have layout constraints.
@@ -183,7 +184,7 @@ def maybe_layout_constraints(fn: Callable[..., Any]) -> Optional[Callable[..., A
 
 def tag_to_layout_constraint(
     tag: torch._C.Tag,
-) -> Optional[Callable[..., tuple[Any, Any]]]:
+) -> Callable[..., tuple[Any, Any]] | None:
     if tag == torch._C.Tag.needs_exact_strides:
         return constrain_to_fake_tensors
     if tag == torch._C.Tag.needs_contiguous_strides:  # type: ignore[attr-defined]
@@ -201,12 +202,10 @@ def assert_nyi(cond: bool, msg: str) -> None:
 
 
 def add_needs_realized_inputs(
-    fn: Union[
-        Collection[Union[torch._ops.OpOverload, torch._ops.OpOverloadPacket]],
-        torch._ops.OpOverload,
-        torch._ops.OpOverloadPacket,
-    ],
-) -> Optional[list[Any]]:
+    fn: Collection[torch._ops.OpOverload | torch._ops.OpOverloadPacket]
+    | torch._ops.OpOverload
+    | torch._ops.OpOverloadPacket,
+) -> list[Any] | None:
     if isinstance(fn, (list, set, tuple, OrderedSet)):  # noqa: set_linter
         # pyrefly: ignore [bad-argument-type]
         return [add_needs_realized_inputs(x) for x in fn]
@@ -220,7 +219,7 @@ def add_needs_realized_inputs(
 
 
 def add_layout_constraint(
-    fn: Union[torch._ops.OpOverloadPacket, torch._ops.OpOverload],
+    fn: torch._ops.OpOverloadPacket | torch._ops.OpOverload,
     constraint: Callable[..., tuple[Any, Any]],
 ) -> None:
     if isinstance(fn, torch._ops.OpOverloadPacket):
@@ -274,7 +273,7 @@ def add_layout_constraint(
 }
 
 
-def decode_dtype(dtype: Union[int, torch.dtype]) -> torch.dtype:
+def decode_dtype(dtype: int | torch.dtype) -> torch.dtype:
     if not isinstance(dtype, int):
         return dtype
     assert dtype in DTYPE_ID_LOOKUP, f"id {dtype} missing from DTYPE_ID_LOOKUP"
@@ -283,7 +282,7 @@ def decode_dtype(dtype: Union[int, torch.dtype]) -> torch.dtype:
     return dtype
 
 
-def is_integer_type(x: Any) -> TypeGuard[Union[TensorBox, sympy.Expr, int]]:
+def is_integer_type(x: Any) -> TypeGuard[TensorBox | sympy.Expr | int]:
     if isinstance(x, TensorBox):
         return is_integer_dtype(x.get_dtype()) or is_boolean_dtype(x.get_dtype())
     elif isinstance(x, sympy.Expr):
@@ -292,7 +291,7 @@ def is_integer_type(x: Any) -> TypeGuard[Union[TensorBox, sympy.Expr, int]]:
         return isinstance(x, int)
 
 
-def is_boolean_type(x: Any) -> TypeGuard[Union[TensorBox, bool]]:
+def is_boolean_type(x: Any) -> TypeGuard[TensorBox | bool]:
     if isinstance(x, TensorBox):
         return is_boolean_dtype(x.get_dtype())
     else:
@@ -336,7 +335,7 @@ def get_overloads(aten_fn):
 
 
 def in_namespace(
-    op: Union[Any, torch._ops.OpOverloadPacket, torch._ops.OpOverload], namespace: str
+    op: Any | torch._ops.OpOverloadPacket | torch._ops.OpOverload, namespace: str
 ) -> bool:
     if isinstance(op, torch._ops.OpOverloadPacket):
         return namespace in op._qualified_op_name
@@ -353,7 +352,7 @@ def maybe_copy_cpu_scalar(x: TensorBox, device: torch.device) -> TensorBox:
         x.get_size()
     ):
         return x
-    size = [V.graph.sizevars.size_hint_or_throw(s) for s in x.get_size()]
+    size = V.graph.sizevars.guarding_hints_or_throw(x.get_size())
     cur_device = x.get_device()
     if (
         cur_device is not None
@@ -369,7 +368,7 @@ def transform_args(
     args: list[Any],
     kwargs: dict[str, Any],
     broadcast: bool,
-    type_promotion_kind: Optional[ELEMENTWISE_TYPE_PROMOTION_KIND],
+    type_promotion_kind: ELEMENTWISE_TYPE_PROMOTION_KIND | None,
     convert_input_to_bool: bool,
 ) -> tuple[list[Any], dict[str, Any]]:
     """
@@ -477,9 +476,9 @@ def _register_lowering(
     aten_fn,
     decomp_fn: Callable[..., Any],
     broadcast: bool,
-    type_promotion_kind: Optional[ELEMENTWISE_TYPE_PROMOTION_KIND],
+    type_promotion_kind: ELEMENTWISE_TYPE_PROMOTION_KIND | None,
     convert_input_to_bool: bool,
-    lowering_dict: dict[Union[Callable[..., Any], str], Callable[..., Any]],
+    lowering_dict: dict[Callable[..., Any] | str, Callable[..., Any]],
 ):
     """
     Add a lowering to lowerings dict
@@ -529,9 +528,8 @@ def wrapped(*args, **kwargs):
 def register_lowering(
     aten_fn,
     broadcast=False,
-    type_promotion_kind: Optional[
-        ELEMENTWISE_TYPE_PROMOTION_KIND
-    ] = ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT,
+    type_promotion_kind: ELEMENTWISE_TYPE_PROMOTION_KIND
+    | None = ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT,
     convert_input_to_bool=False,
     lowering_dict=lowerings,
 ) -> Callable[[Callable[_P, _T]], Callable[_P, _T]]:
@@ -601,13 +599,26 @@ def const_func(x):
 
         return [const_func(x) for x in inputs]
     ex = next(x for x in inputs if isinstance(x, (TensorBox, ExpandView, ir.Constant)))
+    tensor_dtype = ex.get_dtype()
+
+    # Round scalar to tensor's dtype for comparison ops to match eager
+    if override_return_dtype == torch.bool and tensor_dtype in (
+        torch.bfloat16,
+        torch.float16,
+    ):
+        _round_scalar = lambda v: torch.tensor(v, dtype=tensor_dtype).item()  # noqa: E731
+    else:
+        _round_scalar = lambda v: v  # noqa: E731
+
     out = []
     for x in inputs:
         if isinstance(x, (int, float)):
             out.append(
                 ExpandView.create(
                     ir.Constant(
-                        value=x, dtype=ex.get_dtype(), device=ex.get_device_or_error()
+                        value=_round_scalar(x),
+                        dtype=tensor_dtype,
+                        device=ex.get_device_or_error(),
                     ),
                     list(ex.get_size()),
                 )
@@ -616,7 +627,7 @@ def const_func(x):
             out.append(
                 ExpandView.create(
                     IndexingConstant(
-                        index=x, dtype=ex.get_dtype(), device=ex.get_device_or_error()
+                        index=x, dtype=tensor_dtype, device=ex.get_device_or_error()
                     ),
                     list(ex.get_size()),
                 )
@@ -627,14 +638,45 @@ def const_func(x):
     return out
 
 
+def _add_with_alpha_fma(a, b, alpha):
+    """Compute a + alpha * b using FMA for CUDA floating-point precision."""
+    dtype = get_promoted_dtype(
+        a,
+        b,
+        type_promotion_kind=ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT,
+    )
+    a_loader = a.make_loader()
+    b_loader = b.make_loader()
+
+    def inner_fn(idx):
+        a_val = a_loader(idx)
+        b_val = b_loader(idx)
+        if isinstance(alpha, sympy.Basic):
+            alpha_expr = ops.index_expr(alpha, dtype)
+        else:
+            alpha_expr = ops.constant(alpha, dtype)
+        return ops.fma(b_val, alpha_expr, a_val)
+
+    return Pointwise.create(
+        device=a.get_device(),
+        dtype=dtype,
+        inner_fn=inner_fn,
+        ranges=a.get_size(),
+    )
+
+
 def make_pointwise(
     fn,
     override_return_dtype=None,
     override_device=None,
     override_fn_when_input_bool=None,
     allow_alpha=False,
+    use_fma_for_alpha=False,
     triton_fallback=None,
 ):
+    """Wraps a pointwise fn and returns a function representing the pointwise in
+    the define-by-run IR."""
+
     def inner(*inputs: TensorBox, alpha=None):
         if triton_fallback is not None and any(
             isinstance(inp, IRNode) and is_triton(inp) for inp in inputs
@@ -645,6 +687,18 @@ def inner(*inputs: TensorBox, alpha=None):
         inputs = promote_constants(inputs, override_return_dtype)
         if allow_alpha:
             if alpha is not None and alpha != 1:
+                # Use FMA for add-with-alpha on CUDA floating-point.
+                # Eager CUDA computes a + alpha * b as fma(b, alpha, a).
+                if use_fma_for_alpha and isinstance(inputs[0], IRNode):
+                    inp_device = inputs[0].get_device()
+                    if (
+                        inputs[0].get_dtype().is_floating_point
+                        and not torch.version.hip
+                        and inp_device is not None
+                        and inp_device.type == "cuda"
+                    ):
+                        return _add_with_alpha_fma(inputs[0], inputs[1], alpha)
+
                 # pyrefly: ignore [bad-assignment]
                 inputs = list(inputs)
                 # pyrefly: ignore [unsupported-operation]
@@ -716,8 +770,25 @@ def inner_fn(index):
     return inner
 
 
-def make_foreach_pointwise(pw_fn, allow_alpha=False):
-    def inner(*inputs: list[list[TensorBox]], alpha=1):
+def make_foreach_pointwise(pw_fn, allow_alpha=False, scalar_kwarg="alpha"):
+    def inner(*inputs: list[list[TensorBox]], alpha=1, value=1):
+        # For ops like addcmul/addcdiv, the scalar `value` arrives as a
+        # positional arg (not keyword) due to the ATen schema. Extract it
+        # from the end of inputs if present.
+        # pyrefly: ignore [bad-assignment]
+        inputs = list(inputs)
+        if (
+            scalar_kwarg == "value"
+            and inputs
+            and not isinstance(inputs[-1], (list, tuple))
+        ):
+            # pyrefly: ignore [missing-attribute]
+            scalar_val = inputs.pop()
+        elif scalar_kwarg == "value":
+            scalar_val = value
+        else:
+            scalar_val = alpha
+
         realize_outputs = (
             len(V.graph.current_node.users) == 0
             or V.graph.current_node.target in inplace_foreach_ops
@@ -746,7 +817,7 @@ def inner(*inputs: list[list[TensorBox]], alpha=1):
 
         def apply_fn(args):
             if allow_alpha:
-                return pw_fn(*args, alpha=alpha)
+                return pw_fn(*args, **{scalar_kwarg: scalar_val})
             else:
                 return pw_fn(*args)
 
@@ -787,13 +858,27 @@ def foreach_group_loop(groups, num_outputs, apply_fn, realize_outputs):
     return outputs
 
 
-def to_dtype(x: TensorBox, dtype: torch.dtype, copy: bool = False):
+def to_dtype(
+    x: TensorBox, dtype: torch.dtype, copy: bool = False, use_compute_types: bool = True
+):
     src_dtype = x.get_dtype()
     if src_dtype == dtype:
         return clone(x) if copy else x
 
     def _to_dtype(x):
-        return ops.to_dtype(x, dtype, src_dtype=src_dtype)
+        result = ops.to_dtype(
+            x,
+            dtype,
+            src_dtype=src_dtype,
+            use_compute_types=use_compute_types,
+        )
+        low_pr_fp = (torch.bfloat16, torch.float16)
+        if not use_compute_types and dtype in low_pr_fp:
+            # Upcast back to compute type so fused consumers see a compute-type
+            # value. Without this, a raw low-precision value gets a redundant
+            # downcast from the consumer's input emulation.
+            result = ops.to_dtype(result, dtype)
+        return result
 
     return make_pointwise(_to_dtype, override_return_dtype=dtype)(x)
 
@@ -856,7 +941,13 @@ def _convert_element_type(x: TensorBox, dtype: torch.dtype):
             return fallback_handler(
                 prims.convert_element_type.default, add_to_fallback_set=False
             )(x, dtype)
-    return to_dtype(x, dtype, copy=True)
+    src_dtype = x.get_dtype()
+    low_pr_fp = (torch.bfloat16, torch.float16)
+    use_compute_types = not (
+        config.emulate_precision_casts
+        and (src_dtype in low_pr_fp or dtype in low_pr_fp)
+    )
+    return to_dtype(x, dtype, copy=True, use_compute_types=use_compute_types)
 
 
 def to_dtype_bitcast(x: TensorBox, dtype: torch.dtype, *, copy=False):
@@ -867,6 +958,9 @@ def to_dtype_bitcast(x: TensorBox, dtype: torch.dtype, *, copy=False):
     def _get_primitive_bitwidth(dtype):
         if dtype.is_floating_point:
             return torch.finfo(dtype).bits
+        elif dtype == torch.bool:
+            # torch.iinfo doesn't support bool; bools are stored as uint8 (8 bits)
+            return 8
         else:
             return torch.iinfo(dtype).bits
 
@@ -909,6 +1003,7 @@ def register_pointwise(
     override_return_dtype=None,
     override_fn_when_input_bool=None,
     allow_alpha=False,
+    use_fma_for_alpha=False,
     triton_fallback=None,
 ):
     """A pointwise function that maps ops.{name} to inputs"""
@@ -927,6 +1022,7 @@ def register_pointwise(
         override_return_dtype=override_return_dtype,
         override_fn_when_input_bool=override_fn_when_input_bool,
         allow_alpha=allow_alpha,
+        use_fma_for_alpha=use_fma_for_alpha,
         triton_fallback=triton_fallback,
     )
     fn = register_lowering(
@@ -1022,8 +1118,11 @@ def register_foreach_pointwise(
     aten_fn,
     pointwise_lowering_fn,
     allow_alpha=False,
+    scalar_kwarg="alpha",
 ):
-    fn = make_foreach_pointwise(pointwise_lowering_fn, allow_alpha=allow_alpha)
+    fn = make_foreach_pointwise(
+        pointwise_lowering_fn, allow_alpha=allow_alpha, scalar_kwarg=scalar_kwarg
+    )
     fn = _register_foreach_lowering(aten_fn, fn)
     return fn
 
@@ -1183,7 +1282,7 @@ def expand(x, sizes):
         return x
 
     if not free_unbacked_symbols(x.get_size()):
-        x_size_product = V.graph.sizevars.size_hint_or_throw(
+        x_size_product = V.graph.sizevars.guarding_hint_or_throw(
             sympy_product(x.get_size())
         )
         # TODO: It would be better to realize the input if any of its sizes
@@ -1193,7 +1292,7 @@ def expand(x, sizes):
         if x_size_product > 0 and not free_unbacked_symbols(sizes):
             # maybe realize input before broadcasting it
             x.mark_reuse(
-                V.graph.sizevars.size_hint_or_throw(sympy_product(sizes))
+                V.graph.sizevars.guarding_hint_or_throw(sympy_product(sizes))
                 // x_size_product
             )
     return TensorBox(ExpandView.create(x.data, tuple(sizes)))
@@ -1252,13 +1351,16 @@ def inner_fn(index):
                     index[i] = ModularIndexing(index[i], 1, old_size[i])
         return x_loader(index)
 
+    # TODO Laith is there better check
     if not free_unbacked_symbols(old_size) and not free_unbacked_symbols(new_size):
-        old_size_product = V.graph.sizevars.size_hint_or_throw(sympy_product(old_size))
+        old_size_product = V.graph.sizevars.guarding_hint_or_throw(
+            sympy_product(old_size)
+        )
         if old_size_product > 0:
             # maybe realize the input but skip for unbacked symints since it'll
             # choke on the size hint.
             x.mark_reuse(
-                V.graph.sizevars.size_hint_or_throw(sympy_product(new_size))
+                V.graph.sizevars.guarding_hint_or_throw(sympy_product(new_size))
                 // old_size_product
             )
 
@@ -1285,8 +1387,10 @@ def permute(x, dims):
     return TensorBox(PermuteView.create(x.data, tuple(dims)))
 
 
+# Note: logic in this function need to be always synchronized with
+# slice_forward in fake implementation.
 @register_lowering(aten.slice, type_promotion_kind=None)
-def slice_(x, dim=0, start=0, end=2**63, step=1, clamp=True):
+def slice_(x, dim=0, start=0, end=sys.maxsize, step=1, clamp=True):
     """
     Lowers a slice call, creating ExternKernels for the output size & storage offset symbols,
     if the indices are unbacked and appropriate semantics aren't known.
@@ -1323,30 +1427,78 @@ def compute_slice_index(index, size, default=None):
         fn = lambda x: V.graph.sizevars.guard_or_false(x)  # noqa: E731
         index = sympy.expand(index)
         size = sympy.expand(size)
-        if fn(sympy.Ge(index, 0)) and fn(sympy.Le(index, size)):
+        if fn(sympy.And(sympy.Ge(index, 0), sympy.Le(index, size))):
             return index
-        elif fn(sympy.Lt(index, 0)) and fn(sympy.Ge(index, -size)):
+        elif fn(sympy.And(sympy.Lt(index, 0), sympy.Ge(index, -size))):
             return index + size
         elif fn(sympy.Gt(index, size)):
             return size
         elif fn(sympy.Lt(index, -size)):
             return 0
+        elif fn(sympy.Ge(index, 0)):
+            # If index >= 0, the resolved index is at most min(index, size).
+            return sympy.Min(index, size)
+        elif fn(sympy.Lt(index, 0)):
+            # If index < 0, wrap and clamp: the resolved index is at least 0.
+            return sympy.Max(index + size, 0)
         return None
 
     start_index, end_index = None, None
+    # ambiguous_slice=False means we know what semantics this slice call follows,
+    # and don't need to generate an extern kernel to represent the output size.
+    # This is assumed True for clamp=False
+    # (meant to follow standard indexing semantics: 0 <= index < size)
     ambiguous_slice = clamp
     if ambiguous_slice:
         start_index = compute_slice_index(start, size, 0)
-        end_index = compute_slice_index(end, size, size)
+        # Special case: if end is maxsize (unbounded), use size directly
+        # This matches the logic in fake_impls.py
+        if end is not None and V.graph.sizevars.statically_known_equals(
+            end, sys.maxsize
+        ):
+            end_index = size
+        else:
+            end_index = compute_slice_index(end, size, size)
         if start_index is not None and end_index is not None:
             start, end = start_index, end_index
             ambiguous_slice = False
 
-    # ambiguous_slice=False means we know what semantics this slice call follows,
-    # and don't need to generate an extern kernel to represent the output size.
-    # This is assumed True for clamp=False
-    # (meant to follow standard indexing semantics: 0 <= index < size)
     if not ambiguous_slice:
+        # Even though the bounds are resolvable now, the FX node may have
+        # allocated unbacked symbols for the slice output size because dynamo
+        # couldn't prove the bounds at trace time (constraints may have been
+        # learned after tracing the slice). We still need to define those
+        # symbols so the assertion new_unbacked_defs >= renamed_unbacked_bindings
+        # passes. Register a DynamicSliceSize operation to define the size symbol.
+        # Note: storage_offset bindings should not appear here because
+        # a resolved start_index means the offset is computable directly
+        # (base_offset + start * stride), so dynamo wouldn't allocate an
+        # unbacked symbol for it.
+        # Note: current_node may be None when slice_ is called from template
+        # rendering (e.g. cpp_template_kernel.slice_nd) rather than FX graph
+        # lowering, so we handle that.
+        current_node = V.graph.current_node
+        node_unbacked_bindings = resolve_unbacked_bindings(
+            V.graph.sizevars.shape_env,
+            current_node.meta.get("unbacked_bindings", {})
+            if current_node is not None
+            else {},
+        )
+        if node_unbacked_bindings:
+            for sym, keypath in node_unbacked_bindings.items():
+                if keypath == (CallMethodKey("size"), pytree.SequenceKey(dim)):
+                    b_size = ir.DynamicSliceSize(sym, start, end, step, size)
+                    b_size.name = V.graph.register_buffer(b_size)
+                    V.graph.register_operation(b_size)
+                elif keypath == (CallMethodKey("storage_offset"),):
+                    # Not handled yet — would require materializing the
+                    # tensor layout. Unlikely to be hit because a resolved
+                    # start_index means the offset is computable directly.
+                    raise AssertionError(
+                        "Unexpected storage_offset unbacked binding when both "
+                        "start and end indices are resolved"
+                    )
+
         return TensorBox(
             ir.SliceView.create(x.data, dim, start, end, step, clamp=clamp)
         )  # go to SliceView/ReinterpretView
@@ -1608,7 +1760,7 @@ def quantized_decomposed_dequantize_per_channel(
     quant_max: int,
     dtype: torch.dtype,
     *,
-    out_dtype: Optional[torch.dtype] = None,
+    out_dtype: torch.dtype | None = None,
 ) -> TensorBox:
     assert len(scales.get_size()) == 1, "expect scales 1 dim"
     assert len(zero_points.get_size()) == 1, "expect zero_points 1 dim"
@@ -1699,7 +1851,7 @@ def quantized_decomposed_dequantize_per_tensor_default(
     quant_max: int,
     dtype: torch.dtype,
     *,
-    out_dtype: Optional[torch.dtype] = None,
+    out_dtype: torch.dtype | None = None,
 ) -> TensorBox:
     assert input.get_dtype() == dtype, (
         f"Expecting input to have dtype {dtype}, but got dtype: {input.get_dtype()}"
@@ -1754,6 +1906,8 @@ def quantized_decomposed_quantize_per_tensor_tensor(
     scale_loader = scale.make_loader()
     zero_point_loader = zero_point.make_loader()
 
+    device = input.get_device()
+
     def inner_fn(idx):
         input = input_loader(idx)
         _scale = scale_loader((0,) if len(scale.get_size()) == 1 else ())
@@ -1762,6 +1916,9 @@ def inner_fn(idx):
             _scale = ops.to_dtype(_scale, torch.float32)
         if zero_point.dtype != torch.float32:
             _zero_point = ops.to_dtype(_zero_point, torch.float32)
+        if device and device.type == "cpu":
+            val = ops.fma(input, ops.reciprocal(_scale), _zero_point)
+            return ops.round_to_int(val, dtype)
         val = ops.round(input * ops.reciprocal(_scale)) + _zero_point
         qmin, qmax = _create_constants(quant_min, quant_max, dtype=torch.float32)
         clamped = ops.minimum(ops.maximum(val, qmin), qmax)
@@ -1786,7 +1943,7 @@ def quantized_decomposed_dequantize_per_tensor_tensor(
     quant_max: int,
     dtype: torch.dtype,
     *,
-    out_dtype: Optional[torch.dtype] = None,
+    out_dtype: torch.dtype | None = None,
 ) -> TensorBox:
     assert len(scale.get_size()) == 0 or (
         len(scale.get_size()) == 1 and scale.get_size()[0] == 1
@@ -1827,6 +1984,7 @@ def inner_fn(idx):
 
 @register_lowering(aten.cat)
 def cat(inputs, dim=0):
+    """Lower aten.cat, choosing between pointwise_cat and ConcatKernel."""
     cpu_device = inputs[0].get_device().type == "cpu"
     if cpu_device and all(
         input.get_dtype() in [torch.int8, torch.uint8] for input in inputs
@@ -1848,7 +2006,7 @@ def cat(inputs, dim=0):
     )
     inputs = [to_dtype(inp, dtype) for inp in inputs]
 
-    def unwrap_tensor(x: Union[TensorBox, ir.StorageBox]) -> ir.IRNode:
+    def unwrap_tensor(x: TensorBox | ir.StorageBox) -> ir.IRNode:
         if isinstance(x, TensorBox):
             if isinstance(x.data, ir.BaseView):
                 return x.data.unwrap_view()
@@ -1941,10 +2099,50 @@ def additional_pointwise_ops(op: torch._ops.OpOverload):
 
         # horizontal fuse in case all inputs will require a copy kernel anyway.
         # only horizontally fuse pointwise kernels
+
+        # Skip pointwise_cat when any cat input has a fusible (pointwise)
+        # multi-consumer — ConcatKernel + NonOwningLayout avoids redundant
+        # reads. Also skip when input is an unrealized Pointwise with
+        # multiple consumers to avoid recomputation (e.g. pad-as-cat).
+        def any_input_has_multi_consumers() -> bool:
+            current_node = V.current_node
+            if current_node is None:
+                return False
+            fx_args = current_node.args[0]
+            if isinstance(fx_args, (list, tuple)):
+                input_nodes = fx_args
+            elif isinstance(fx_args, torch.fx.Node):
+                input_nodes = [fx_args]
+            else:
+                return False
+
+            def is_unrealized_pointwise(x):
+                if isinstance(x, (TensorBox, ir.StorageBox)):
+                    return is_unrealized_pointwise(unwrap_tensor(x))
+                return isinstance(x, ir.Pointwise)
+
+            for arg, ir_input in zip(input_nodes, inputs):
+                if not hasattr(arg, "users") or len(arg.users) <= 1:
+                    continue
+                # input will be computed multiple times because other consumers
+                # (eg. pointwise) will also inline it. So we should realize-in-place via ConcatKernel
+                if any(is_pointwise_use(u) for u in arg.users if u is not current_node):
+                    return True
+                # If input is an unrealized Pointwise with multiple consumers, pointwise_cat
+                # will inline input without realizing it to memory, causing separate
+                # realization cost for input. So we should realize-in-place via ConcatKernel
+                if is_unrealized_pointwise(ir_input):
+                    return True
+            return False
+
+        has_multi_consumers = any_input_has_multi_consumers()
+
         horizontal_fuse_cat = all(
             should_lower_cat_input(inp) for inp in inputs
         ) and not any(can_fuse_reduction(t) for t in inputs)
-        if fuse_pointwise_use or (horizontal_fuse_cat and not fusable_reduction):
+        if not has_multi_consumers and (
+            fuse_pointwise_use or (horizontal_fuse_cat and not fusable_reduction)
+        ):
             return pointwise_cat(inputs, dim)
 
     return TensorBox(ir.ConcatKernel.create(inputs, dim))
@@ -1965,14 +2163,16 @@ def diagonal(input, offset: int = 0, dim1: int = 0, dim2: int = 1):
     if offset_negative:
         diag_size = V.graph.sizevars.evaluate_max(
             V.graph.sizevars.evaluate_min(
-                original_shape[dim1] + offset, original_shape[dim2]
+                original_shape[dim1] + offset,
+                original_shape[dim2],
             ),
             0,  # type: ignore[arg-type]
         )
     else:
         diag_size = V.graph.sizevars.evaluate_max(
             V.graph.sizevars.evaluate_min(
-                original_shape[dim1], original_shape[dim2] - offset
+                original_shape[dim1],
+                original_shape[dim2] - offset,
             ),
             0,  # type: ignore[arg-type]
         )
@@ -2138,9 +2338,9 @@ def unfold(x, dimension, size, step):
     sizevars.check_lt(0, step)  # type: ignore[arg-type]
 
     new_dim_size = FloorDiv(dim_size - size, step) + 1
-    if sizevars.size_hint_or_throw(dim_size) > 0:
+    if sizevars.guarding_hint_or_throw(dim_size) > 0:
         x.mark_reuse(
-            sizevars.size_hint_or_throw(CeilDiv(new_dim_size * size, dim_size))
+            sizevars.guarding_hint_or_throw(CeilDiv(new_dim_size * size, dim_size))
         )
 
     out_size = [*sizes[:dim], new_dim_size, *sizes[dim + 1 :], size]
@@ -2195,7 +2395,7 @@ def fallback_handler(kernel, add_to_fallback_set=True):
 
     def handler(*args, **kwargs):
         def wrap_tensors(x):
-            return TensorBox.create(x) if isinstance(x, ir.IRNode) else x
+            return x.wrap_for_lowering() if isinstance(x, ir.IRNode) else x
 
         return pytree.tree_map(
             wrap_tensors, ir.FallbackKernel.create(kernel, *args, **kwargs)
@@ -2303,16 +2503,15 @@ def check_skip_condition(inp_out_node, is_output):
     return check_skip_condition(node, is_output=True)
 
 
-def make_fallback(op, layout_constraint=None, warn=True, override_decomp=False):
-    # When emulate_precision_casts is enabled, we skip decomposing addcmul ops
-    # to use the inductor lowering which preserves FMA semantics.
-    # For _foreach_addcdiv, we use the native CUDA kernel.
-    skip_decomp_for_precision = config.emulate_precision_casts and op in {
-        aten.addcmul,
-        aten._foreach_addcmul.Scalar,
-        aten._foreach_addcdiv.Scalar,
-    }
-    assert op not in decompositions or override_decomp or skip_decomp_for_precision, (
+def make_fallback(
+    op,
+    layout_constraint=None,
+    warn=True,
+    override_decomp=False,
+    get_decomp_fn=None,
+):
+    check_decomps = get_decomp_fn() if get_decomp_fn is not None else decompositions
+    assert op not in check_decomps or override_decomp, (
         f"both a fallback and a decomp for same op: {op}"
     )
     if (
@@ -2467,10 +2666,15 @@ def warn_triton_random():
 fallback_randn_default = fallback_handler(aten.randn.default)
 fallback_randn_generator = fallback_handler(aten.randn.generator)
 make_fallback(aten.randint)
+make_fallback(aten.rand_like, override_decomp=True)
+make_fallback(aten.randn_like, override_decomp=True)
+make_fallback(aten.randint_like, override_decomp=True)
 
 # TODO: mlazos reevaluate if we want to codegen something different
 make_fallback(torch.ops.streams.record_event.default)
 make_fallback(torch.ops.streams.wait_event.default)
+make_fallback(torch.ops.streams.synchronize_event.default)
+make_fallback(torch.ops.streams.synchronize_device.default)
 
 
 @register_lowering(aten.rand)
@@ -2523,8 +2727,35 @@ def inner_fn(_):
     )
 
 
+def get_threads_per_round(device: torch.device):
+    if not isinstance(device, torch.device):
+        device = torch.device(device)
+
+    if device.type == "cuda":
+        idx = device.index
+        if idx is None:
+            idx = torch.cuda.current_device()
+
+        prop = torch.cuda.get_device_properties(idx)
+        threads_per_round = (
+            prop.multi_processor_count * prop.max_threads_per_multi_processor
+        )
+    else:
+        _CPU_GRAIN_SIZE = 32768
+        threads_per_round = _CPU_GRAIN_SIZE
+
+    return threads_per_round
+
+
 @register_lowering(inductor_prims.random, type_promotion_kind=None)
-def inductor_random(size: list[int], seed: TensorBox, mode: str, *, offset: int = 0):
+def inductor_random(
+    size: list[int],
+    seed: TensorBox,
+    mode: str,
+    *,
+    offset: int = 0,
+    align_dtype: torch.dtype = torch.float32,
+):
     assert not config.fallback_random
     assert mode in ("rand", "randn")
     size = [*size]
@@ -2535,11 +2766,33 @@ def inductor_random(size: list[int], seed: TensorBox, mode: str, *, offset: int
     ).make_indexer()
     seed_loader = seed.make_loader()
 
-    def inner_fn(index):
-        return getattr(ops, mode)(
-            seed_loader([]),
-            ops.index_expr(random_pos(index), torch.int32),
-        )
+    if config.align_random_eager and device.type == "cuda":
+        threads_per_round = get_threads_per_round(device)
+
+        def _vec_from_dtype(dt: torch.dtype) -> int:
+            if dt in (torch.float16, torch.bfloat16):
+                return 8
+            return 4
+
+        vec = _vec_from_dtype(align_dtype)
+
+        def inner_fn(index):
+            rng_seed = seed_loader([0])
+            base_offset = seed_loader([1])
+            return ops.rand_eager(
+                rng_seed,
+                base_offset,
+                threads_per_round,
+                ops.index_expr(random_pos(index), torch.int32),
+                vec=int(vec),
+            )
+    else:
+
+        def inner_fn(index):
+            return getattr(ops, mode)(
+                seed_loader([]),
+                ops.index_expr(random_pos(index), torch.int32),
+            )
 
     result = Pointwise.create(
         device=device,
@@ -2551,6 +2804,10 @@ def inner_fn(index):
     return result
 
 
+make_fallback(inductor_prims.rand_eager_offset)
+make_fallback(inductor_prims.rand_eager_offsets)
+
+
 @register_lowering(inductor_prims.randint, type_promotion_kind=None)
 def inductor_randint(
     low: int, high: int, size: list[int], seed: TensorBox, *, offset: int = 0
@@ -2607,8 +2864,8 @@ def searchsorted(
     *,
     out_int32: bool = False,
     right: bool = False,
-    side: Optional[str] = None,
-    sorter: Optional[TensorBox] = None,
+    side: str | None = None,
+    sorter: TensorBox | None = None,
 ) -> TensorBox:
     validate_bucketize = lambda tb: V.graph.has_feature(  # noqa: E731
         tb, BackendFeature.BUCKETIZE
@@ -2759,16 +3016,24 @@ def inner_fn(index):
     return result
 
 
+def _is_tensor_irnode(x):
+    return isinstance(x, ir.IRNode) and not isinstance(
+        x, (ir.NonTensorObj, ir.OpaqueMultiOutput)
+    )
+
+
 def require_dense(_, *args, **kwargs):
     args, kwargs = pytree.tree_map_only(
-        ir.IRNode, ir.ExternKernel.require_stride1, (args, kwargs)
+        _is_tensor_irnode, ir.ExternKernel.require_stride1, (args, kwargs)
     )
     return args, kwargs
 
 
 def require_contiguous(_, *args, **kwargs):
     args, kwargs = pytree.tree_map_only(
-        ir.IRNode, ir.ExternKernel.require_contiguous, (args, kwargs)
+        _is_tensor_irnode,
+        ir.ExternKernel.require_contiguous,
+        (args, kwargs),
     )
     return args, kwargs
 
@@ -2777,14 +3042,18 @@ def require_contiguous_strides(_, *args, **kwargs):
     # TODO: combine this with require_contiguous after
     # https://github.com/pytorch/pytorch/pull/148235 lands.
     args, kwargs = pytree.tree_map_only(
-        ir.IRNode, ir.ExternKernel.require_contiguous_strides, (args, kwargs)
+        _is_tensor_irnode,
+        ir.ExternKernel.require_contiguous_strides,
+        (args, kwargs),
     )
     return args, kwargs
 
 
 def require_channels_last(_, *args, **kwargs):
     args, kwargs = pytree.tree_map_only(
-        ir.IRNode, ir.ExternKernel.require_channels_last, (args, kwargs)
+        _is_tensor_irnode,
+        ir.ExternKernel.require_channels_last,
+        (args, kwargs),
     )
     return args, kwargs
 
@@ -2792,7 +3061,7 @@ def require_channels_last(_, *args, **kwargs):
 def constrain_to_fake_tensor(arg, fake_arg):
     if fake_arg is None:
         return arg
-    if isinstance(fake_arg, FakeScriptObject):
+    if isinstance(fake_arg, FakeScriptObject) or is_opaque_value(fake_arg):
         return arg
     if isinstance(arg, ir.IRNode):
         return ir.ExternKernel.require_exact_strides(arg, fake_arg.stride())
@@ -2816,9 +3085,12 @@ def constrain_to_fake_tensors(args, kwargs, fake_args, fake_kwargs):
 
 def constrain_to_fx_strides(fx_node, *args, **kwargs):
     def apply_constraint(arg, fx_arg):
-        if isinstance(arg, ir.IRNode):
+        if _is_tensor_irnode(arg):
+            fake_val = fx_arg.meta.get("val")
+            if not isinstance(fake_val, torch.Tensor):
+                return arg
             stride_order = ir.get_stride_order(
-                fx_arg.meta["val"].stride(), V.graph.sizevars.shape_env
+                fake_val.stride(), V.graph.sizevars.shape_env
             )
             return ir.ExternKernel.require_stride_order(arg, stride_order)
         if isinstance(arg, dict):
@@ -2832,11 +3104,17 @@ def apply_constraint(arg, fx_arg):
     return args, kwargs
 
 
+# native_dropout uses empty_like(input) internally, so bernoulli_ consumes
+# RNG values in the input's stride order. Constrain input strides to match
+# the FX graph (i.e. eager) so the dropout mask is identical.
+add_layout_constraint(aten.native_dropout.default, constrain_to_fx_strides)
+
+
 def sdpa_constraint(fx_node, *args, **kwargs):
-    # sdpa requires dense last dimension]
+    """Apply stride constraints to SDPA inputs, ensuring dense last dimension."""
 
     def apply_constraint(idx, arg, fx_arg):
-        if not isinstance(arg, ir.IRNode):
+        if not _is_tensor_irnode(arg):
             return arg
 
         meta_val = fx_arg.meta["val"]
@@ -2862,7 +3140,32 @@ def apply_constraint(idx, arg, fx_arg):
             # Check https://github.com/pytorch/pytorch/issues/138772
             stride_order = (3, 1, 2, 0)
 
-        if not meta_val.is_cuda:
+        # Cache keyed by (id(arg), arg_name, stride_order) to avoid
+        # duplicate copy_input when the same tensor feeds multiple SDPA
+        # positions (e.g., key=value).  Including arg_name handles
+        # mutation: mark_buffer_mutated() renames the buffer in place,
+        # so a mutated tensor has the same id but a different name,
+        # causing a cache miss.
+        cache_key = None
+        if config.cache_sdpa_constraint:
+            arg_name = arg.maybe_get_name()
+            cache_key = (
+                id(arg),
+                arg_name,
+                tuple(stride_order) if stride_order else None,
+            )
+            if cache_key in V.graph.sdpa_constraint_cache:
+                return V.graph.sdpa_constraint_cache[cache_key]
+
+        result = _apply_constraint_inner(
+            idx, arg, meta_val, meta_stride_expr, stride_order
+        )
+        if cache_key is not None:
+            V.graph.sdpa_constraint_cache[cache_key] = result
+        return result
+
+    def _apply_constraint_inner(idx, arg, meta_val, meta_stride_expr, stride_order):
+        if not (meta_val.is_cuda or meta_val.is_xpu):
             return ir.ExternKernel.require_stride_order(arg, stride_order)
 
         # This is the minimum alignment required by SDPA kernels for attention_bias.
@@ -2921,12 +3224,14 @@ def apply_constraint(idx, arg, fx_arg):
                 # we can make them expanded by setting the stride equal to 0
                 if i in expanded_dims:
                     if V.graph.sizevars.statically_known_equals(
-                        out_strides[i + 1] % ALIGNMENT, 0
+                        Mod(out_strides[i + 1], ALIGNMENT), 0
                     ):
                         out_strides[i] = 0
                         continue
 
-                if not V.graph.sizevars.statically_known_equals(stride % ALIGNMENT, 0):
+                if not V.graph.sizevars.statically_known_equals(
+                    Mod(stride, ALIGNMENT), 0
+                ):
                     stride = ceildiv(stride, ALIGNMENT) * ALIGNMENT
 
                 out_strides[i] = stride
@@ -2981,7 +3286,7 @@ def is_aligned(x):
 make_fallback(aten._pdist_forward, require_contiguous)  # Has decomp. Needs benchmarks
 make_fallback(aten.soft_margin_loss_backward, warn=False)  # py_impl?
 make_fallback(aten._fused_rms_norm, warn=False)  # (MPS-only and faster than decomp)
-if torch.xpu.is_available():
+if torch.xpu._is_compiled():
     make_fallback(
         aten.embedding_dense_backward, warn=False
     )  # (XPU-only and faster than decomp)
@@ -3055,12 +3360,6 @@ def is_aligned(x):
 # 5) Impossible (missing triton/CPU features)
 
 # Sorting / Sorting-like
-make_fallback(aten.sort)
-make_fallback(aten.sort.stable)
-make_fallback(aten.kthvalue)
-make_fallback(aten.topk)
-make_fallback(aten.mode)
-make_fallback(aten.median)
 make_fallback(aten.nanmedian)
 make_fallback(aten.randperm)
 # see: https://github.com/pytorch/pytorch/pull/121354
@@ -3103,6 +3402,7 @@ def is_aligned(x):
 make_fallback(torch._prims.rng_prims.run_and_save_rng_state)
 make_fallback(torch._prims.rng_prims.run_with_rng_state)
 make_fallback(torch._prims.rng_prims.graphsafe_run_with_rng_state)
+make_fallback(torch._prims.rng_prims.run_dtensor_rng_op)
 
 
 # Implemented / Half implemented
@@ -3268,6 +3568,30 @@ def fn(index):
     )
 
 
+@register_lowering(aten.arange.start_step, type_promotion_kind=None)
+def arange_start_step(
+    start,
+    end,
+    step=1,
+    *,
+    dtype=None,
+    device=None,
+    layout=None,
+    pin_memory=None,
+    requires_grad=False,
+):
+    assert dtype is not None
+    length = ceildiv(end - start, step)
+    return iota(
+        length,
+        start=start,
+        step=step,
+        dtype=dtype,
+        device=device if device is not None else "cpu",
+        requires_grad=requires_grad,
+    )
+
+
 @register_lowering(aten.select_scatter, type_promotion_kind=None)
 def select_scatter(x, src, dim: int, index: int):
     src = to_dtype(src, x.get_dtype())
@@ -3379,7 +3703,7 @@ def _unwrap(x):
     return x
 
 
-@register_lowering([torch.tensor, aten.scalar_tensor])
+@register_lowering([torch.tensor, aten.scalar_tensor, prims.scalar_tensor])
 def tensor(data, *, dtype=None, device=None, layout=None, pin_memory=False):
     assert_nyi(layout in (None, torch.strided), f"layout={layout}")
     assert_nyi(not pin_memory, "pin_memory")
@@ -3712,7 +4036,7 @@ def new_empty_strided(
 
 @register_lowering(prims.copy_strided.default)
 def copy_strided(x, stride):
-    stride = [V.graph.sizevars.size_hint_or_throw(s) for s in stride]
+    stride = V.graph.sizevars.guarding_hints_or_throw(stride)
     stride_order = sorted(range(len(stride)), key=stride.__getitem__)
     return ir.ExternKernel.require_stride_order(x, stride_order)
 
@@ -4117,7 +4441,6 @@ def indice_slice_from_randperm(indice):
         None,
         check=check,
     )
-
     values = expand(values, expected_vals_size)
     # all guards are set above during broadcast_tensors and expand
 
@@ -4214,7 +4537,7 @@ def scatter_fallback(
     index,
     src,
     *,
-    reduce: Optional[str] = None,
+    reduce: str | None = None,
     include_self: bool = True,
 ):
     src_is_tensor = isinstance(src, TensorBox)
@@ -4242,7 +4565,7 @@ def scatter_fallback(
 
 
 @register_lowering(aten.scatter_, type_promotion_kind=None)
-def scatter_(self, dim: int, index, src, *, reduce: Optional[str] = None):
+def scatter_(self, dim: int, index, src, *, reduce: str | None = None):
     assert reduce in (None, "add", "multiply")
     if reduce is None:
         op_overload = getattr(aten.scatter_, V.graph.current_node.target._overloadname)  # type: ignore[union-attr]
@@ -4394,7 +4717,7 @@ def backend_reduce_str(reduce):
 def upsample_nearestnd(
     x,
     output_size,
-    scales_x: tuple[Optional[float], ...],
+    scales_x: tuple[float | None, ...],
     n: int = 2,
     exact: bool = False,
 ):
@@ -4439,25 +4762,25 @@ def fn(idx):
 
 
 @register_lowering(aten.upsample_nearest1d.default)
-def upsample_nearest1d(x, output_size, scales: Optional[float] = None):
+def upsample_nearest1d(x, output_size, scales: float | None = None):
     return upsample_nearestnd(x, output_size, (scales,), n=1)
 
 
 @register_lowering(aten._upsample_nearest_exact1d.default)
-def _upsample_nearest_exact1d(x, output_size, scales: Optional[float] = None):
+def _upsample_nearest_exact1d(x, output_size, scales: float | None = None):
     return upsample_nearestnd(x, output_size, (scales,), n=1, exact=True)
 
 
 @register_lowering(aten.upsample_nearest2d.default)
 def upsample_nearest2d(
-    x, output_size, scales_h: Optional[float] = None, scales_w: Optional[float] = None
+    x, output_size, scales_h: float | None = None, scales_w: float | None = None
 ):
     return upsample_nearestnd(x, output_size, (scales_h, scales_w), n=2)
 
 
 @register_lowering(aten._upsample_nearest_exact2d.default)
 def _upsample_nearest_exact2d(
-    x, output_size, scales_h: Optional[float] = None, scales_w: Optional[float] = None
+    x, output_size, scales_h: float | None = None, scales_w: float | None = None
 ):
     return upsample_nearestnd(x, output_size, (scales_h, scales_w), n=2, exact=True)
 
@@ -4466,9 +4789,9 @@ def _upsample_nearest_exact2d(
 def upsample_nearest3d(
     x,
     output_size,
-    scales_d: Optional[float] = None,
-    scales_h: Optional[float] = None,
-    scales_w: Optional[float] = None,
+    scales_d: float | None = None,
+    scales_h: float | None = None,
+    scales_w: float | None = None,
 ):
     return upsample_nearestnd(x, output_size, (scales_d, scales_h, scales_w), n=3)
 
@@ -4477,9 +4800,9 @@ def upsample_nearest3d(
 def _upsample_nearest_exact3d(
     x,
     output_size,
-    scales_d: Optional[float] = None,
-    scales_h: Optional[float] = None,
-    scales_w: Optional[float] = None,
+    scales_d: float | None = None,
+    scales_h: float | None = None,
+    scales_w: float | None = None,
 ):
     return upsample_nearestnd(
         x, output_size, (scales_d, scales_h, scales_w), n=3, exact=True
@@ -4514,7 +4837,7 @@ def loader(idx):
 
 def inplace_constant_pad_nd(
     x: TensorBox, padding: Sequence[int], fill_value: float
-) -> Optional[TensorBox]:
+) -> TensorBox | None:
     """
     This optimization changes the semantics of padding from 'clone'
     style to 'view' style.
@@ -4607,6 +4930,59 @@ def _padding_can_be_fused():
     return resized_x
 
 
+def _pad_as_cat(
+    x: TensorBox, padding: Sequence[int], fill_value: float
+) -> TensorBox | None:
+    """Decompose right-pad into cat([x, fill], dim) and delegate to cat lowering.
+
+    The cat lowering already has heuristics for choosing between pointwise_cat
+    (fusion) and ConcatKernel (memory planning / zero-copy).  By routing through
+    cat() we reuse those heuristics rather than duplicating them here.
+    """
+    # Bail out for symbolic padding, dynamic shapes
+    if not all(isinstance(p, int) for p in padding):
+        return None
+
+    sizes = x.get_size()
+    ndim = len(sizes)
+    pad_pairs = list(zip(padding[::2], padding[1::2]))
+
+    # Only support single-dimension right-pad
+    pad_dim = None
+    pad_amount = None
+    for i, (left, right) in enumerate(pad_pairs):
+        if left != 0:
+            return None
+        if right > 0:
+            if pad_dim is not None:
+                return None  # multi-dim pad
+            pad_dim = ndim - 1 - i  # padding format is reversed dim order
+            pad_amount = right
+        elif right < 0:
+            return None  # trim, not pad
+
+    if pad_dim is None:
+        return None
+
+    # CPU cat always uses ConcatKernel (no pointwise_cat), which adds
+    # extra kernel launches for the fill.  Skip pad-as-cat on CPU.
+    device = x.get_device()
+    if device is not None and device.type == "cpu":
+        return None
+
+    # Build the fill tensor for the padding region
+    pad_shape = list(sizes)
+    pad_shape[pad_dim] = pad_amount
+    dtype = x.get_dtype()
+    fill_value_typed = dtype_to_type(dtype)(fill_value)
+    pad_tensor = tensor_constructor(fill_value_typed)(
+        pad_shape, dtype=dtype, device=device
+    )
+
+    counters["inductor"]["pad_rewritten_as_cat"] += 1
+    return cat([x, pad_tensor], pad_dim)
+
+
 @register_lowering(aten.constant_pad_nd, type_promotion_kind=None)
 def constant_pad_nd(x, padding, fill_value=0):
     assert (len(padding) % 2) == 0
@@ -4619,6 +4995,10 @@ def constant_pad_nd(x, padding, fill_value=0):
             return out
             # fall through if can not inplace the padding
 
+    out = _pad_as_cat(x, padding, fill_value)
+    if out is not None:
+        return out
+
     sizes = x.get_size()
 
     bounds = list(reversed(list(zip(padding[::2], padding[1::2]))))
@@ -4663,7 +5043,7 @@ def offset_fn(index):
     )
 
 
-def range_mask_low(i: sympy.Expr, low: Union[sympy.Expr, int]):
+def range_mask_low(i: sympy.Expr, low: sympy.Expr | int):
     return ops.ge(
         ops.index_expr(i, torch.int64),
         ops.index_expr(sympy.Integer(low), torch.int64),
@@ -4698,7 +5078,7 @@ def load(index):
 
         mask = functools.reduce(
             ops.and_,
-            # pyrefly: ignore [no-matching-overload]
+            # pyrefly: ignore [bad-argument-type, no-matching-overload]
             [range_mask(ih[i], h[i] + padding_h[i], -padding_h[i]) for i in range(dim)],
         )
         return (
@@ -4733,13 +5113,13 @@ def pooling_size(x, i, kernel_size, stride, padding, ceil_mode, *, dilation=None
             + 2 * (stride[i] - 1),
             stride[i],
         )
-        if V.graph.sizevars.size_hint((x_alt - 1) * stride[i] - x - padding[i]) >= 0:
+        if V.graph.sizevars.guard_or_false(
+            sympy.Ge((x_alt - 1) * stride[i] - x - padding[i], 0)
+        ):
             # Sliding windows must start within the input or left padding
             x_alt -= 1  # type: ignore[assignment]
-            V.graph.sizevars.check_leq(0, x_alt * stride[i] - x - padding[i])  # type: ignore[arg-type]
-        if V.graph.sizevars.size_hint(x_out - x_alt) == 0:
+        if V.graph.sizevars.guard_or_false(sympy.Eq(x_out, x_alt)):
             # ceil mode is actually a no-op, lets guard on that
-            V.graph.sizevars.check_equals(x_out, x_alt)
             ceil_mode = False
         else:
             x_out = x_alt
@@ -4893,10 +5273,10 @@ def _low_memory_max_pool_with_offsets(
 
 def _pool_offsets_to_indices(
     offsets: TensorBox,
-    kernel_size: Sequence[Union[int, torch.SymInt]],
-    input_size: Sequence[Union[int, torch.SymInt]],
+    kernel_size: Sequence[int | torch.SymInt],
+    input_size: Sequence[int | torch.SymInt],
     increments_to_index: Callable[
-        [Sequence[Union[int, torch.SymInt]], Sequence[Union[int, torch.SymInt]]],
+        [Sequence[int | torch.SymInt], Sequence[int | torch.SymInt]],
         torch._inductor.virtualized.OpsValue,
     ],
 ) -> TensorBox:
@@ -5029,7 +5409,7 @@ def max_pool2d_with_indices_backward(
     # we will read this many times, so make sure it is computed
     grad_output.realize_hint()
     gO_stride = grad_output.maybe_get_stride()
-    x_stride: Optional[Sequence[Any]]
+    x_stride: Sequence[Any] | None
     if isinstance(x, TensorBox) and isinstance(x.data.data, Pointwise):  # type: ignore[attr-defined]
         data = x.data.data  # type: ignore[attr-defined]
         device = data.get_device()
@@ -5613,7 +5993,7 @@ def fn(idx):
         device=x.get_device(),
         dtype=x.get_dtype(),
         inner_fn=fn,
-        # pyrefly: ignore [no-matching-overload]
+        # pyrefly: ignore [bad-argument-type, no-matching-overload]
         ranges=list(input_size),
     )
 
@@ -6281,6 +6661,15 @@ def loader(index, reduction_index):
 
 def make_reduction(reduction_type: ReductionType, override_return_dtype=None):
     def inner(x, axis=None, keepdims=False, *, dtype=None):
+        # For argmax/argmin on boolean tensors, cast to int32 first to ensure
+        # correct comparison in Triton. See https://github.com/pytorch/pytorch/issues/174069
+        # Only apply on Triton backend; MPS handles bool comparisons natively.
+        if (
+            reduction_type in ("argmax", "argmin")
+            and x.get_dtype() == torch.bool
+            and is_triton(x)
+        ):
+            x = to_dtype(x, torch.int32)
         kwargs = _make_reduction_inner(
             x,
             axis=axis,
@@ -6358,7 +6747,8 @@ def var_mean_sum_(x, axis, correction, keepdim, return_mean):
 
 
 def use_two_step_variance(x, axis, keepdim):
-    # Instead of unrolling welford, just unroll the simpler two-step var
+    # two-step algorithm can get better performance in small reductions size
+    # while it can accumulate more numerical error than Welford algorithm.
     axis = _validate_reduction_axis(x, axis)
     kwargs = _make_reduction_inner(
         x, axis=axis, keepdims=keepdim, dtype=None, override_return_dtype=None
@@ -6366,9 +6756,16 @@ def use_two_step_variance(x, axis, keepdim):
 
     ranges = kwargs["ranges"]
     reduction_numel = sympy_product(kwargs["reduction_ranges"])
+    device = x.get_device()
+    if not (device and device.type == "cpu"):
+        threshold = config.unroll_reductions_threshold
+    else:
+        # 1024 is a default value to pass all the UTs about accuracy.
+        # A larger threshold can still get performance benefits.
+        threshold = config.cpp.use_two_step_variance_threshold
     return (
         isinstance(reduction_numel, sympy.Integer)
-        and int(reduction_numel) < config.unroll_reductions_threshold
+        and int(reduction_numel) <= threshold
         and sympy_product(ranges) != 1
     )
 
@@ -6429,7 +6826,10 @@ def var_mean_helper_(x, *, axis, correction, keepdim, return_mean):
     )
     output = (
         var_mean_sum_(**kwargs)
-        if use_two_step_variance(x, axis=axis, keepdim=keepdim)
+        if (
+            config.mtia.disable_welford_reduction
+            or use_two_step_variance(x, axis=axis, keepdim=keepdim)
+        )
         else var_mean_welford_(**kwargs)
     )
     output = tuple(to_dtype(x, out_dtype, copy=False) for x in output)
@@ -6591,6 +6991,11 @@ def truncdiv(a, b):
     return ops.truncdiv(a, b)
 
 
+@make_pointwise
+def _div_rn(a, b):
+    return ops.div_rn(a, b)
+
+
 @register_lowering(aten.div, broadcast=True)
 def div_mode(a, b, rounding_mode=None):
     both_integer = is_integer_type(a) and is_integer_type(b)
@@ -6600,7 +7005,11 @@ def div_mode(a, b, rounding_mode=None):
     # see the discussion at https://github.com/triton-lang/triton/issues/605
     if rounding_mode == "floor":
         assert not both_boolean, "floordiv operands can not be boolean at the same time"
-        return floordiv(a, b) if both_integer else floor(div(a, b))
+        # Use div_rn (IEEE round-to-nearest) instead of truediv here because
+        # Triton's default division uses an approximate reciprocal, which can
+        # produce a result slightly below the true quotient and cause floor()
+        # to round down by one.
+        return floordiv(a, b) if both_integer else floor(_div_rn(a, b))
     if rounding_mode == "trunc":
         assert not both_boolean, "truncdiv operands can not be boolean at the same time"
         return truncdiv(a, b) if both_integer else trunc(div(a, b))
@@ -6617,7 +7026,7 @@ def mul(a, b):
         return make_pointwise(fn)(a, b)
 
 
-def get_constant_value(x: ir.IRNode) -> Optional[ir.Constant]:
+def get_constant_value(x: ir.IRNode) -> ir.Constant | None:
     """Try convert an arbitrary IR node into an ir.Constant value"""
 
     # First try unwrapping the IRNode to see if it is already an ir.Constant
@@ -6886,7 +7295,10 @@ def reduce_min(x, dim=None, keepdim=False):
 )
 
 add = register_pointwise(
-    aten.add, allow_alpha=True, override_fn_when_input_bool="logical_or"
+    aten.add,
+    allow_alpha=True,
+    use_fma_for_alpha=True,
+    override_fn_when_input_bool="logical_or",
 )
 
 sort_fallback = fallback_handler(aten.sort.stable, add_to_fallback_set=False)
@@ -6904,11 +7316,18 @@ def sort_stable(x, *, stable=None, dim=-1, descending=False):
         return clone(x), _full(0, device, torch.int64, shape)
 
     dim_size = shape[dim] if len(shape) else 1
-    if not V.graph.sizevars.statically_known_lt(dim_size, torch.iinfo(torch.int16).max):
+    # Use int32 indices when decompose_sort_ops is enabled, allowing sort
+    # dimensions up to 2^31-1.  Default int16 keeps register pressure low
+    # on GPU where the bitonic network holds all indices in-block.
+    if config.triton.decompose_sort_ops:
+        idx_dtype = torch.int32
+    else:
+        idx_dtype = torch.int16
+    if not V.graph.sizevars.statically_known_lt(dim_size, torch.iinfo(idx_dtype).max):
         return sort_fallback(x, stable=stable, dim=dim, descending=descending)
 
     indices = iota(
-        dim_size, start=0, step=1, dtype=torch.int16, device=device, requires_grad=False
+        dim_size, start=0, step=1, dtype=idx_dtype, device=device, requires_grad=False
     )
     view_shape = [1] * len(shape)
     if len(shape):
@@ -6937,6 +7356,199 @@ def sort(x, dim=-1, descending=False):
     return sort_stable(x, stable=False, dim=dim, descending=descending)
 
 
+# Sort-based op lowerings
+# When config.triton.decompose_sort_ops is enabled, decompose into sort-based
+# ops so Inductor generates Triton kernels via ir.Sort.
+# Otherwise, fall back to ATen eager.
+topk_fallback = fallback_handler(aten.topk.default, add_to_fallback_set=False)
+kthvalue_fallback = fallback_handler(aten.kthvalue.default, add_to_fallback_set=False)
+median_fallback = fallback_handler(aten.median.default, add_to_fallback_set=False)
+median_dim_fallback = fallback_handler(aten.median.dim, add_to_fallback_set=False)
+mode_fallback = fallback_handler(aten.mode.default, add_to_fallback_set=False)
+
+# sort/sort.stable already have register_lowering above (sort_stable, sort).
+# They use ir.Sort directly and fall back when the dimension is too large.
+# When decompose_sort_ops is enabled, the size limit is lifted (int32 indices).
+
+
+@register_lowering(aten.median.default, type_promotion_kind=None)
+def median_default(self):
+    if not config.triton.decompose_sort_ops:
+        return median_fallback(self)
+    size = self.get_size()
+    numel = functools.reduce(operator.mul, size, sympy.Integer(1))
+    flat = view(self, [numel])
+    sorted_vals, _ = sort_stable(flat, dim=0)
+    k = (numel - 1) // 2
+    return select(sorted_vals, 0, k)
+
+
+@register_lowering(aten.median.dim, type_promotion_kind=None)
+def median_dim(self, dim, keepdim=False):
+    if not config.triton.decompose_sort_ops:
+        return median_dim_fallback(self, dim, keepdim)
+    shape = self.get_size()
+    ndim = len(shape)
+    if ndim == 0:
+        return clone(self), _full(0, self.get_device(), torch.int64, shape)
+    dim = canonicalize_dim(ndim, dim)
+    sorted_vals, sorted_idxs = sort_stable(self, stable=True, dim=dim)
+    n = shape[dim]
+    k = (n - 1) // 2
+    values = select(sorted_vals, dim, k)
+    indices = select(sorted_idxs, dim, k)
+    if keepdim:
+        values = unsqueeze(values, dim)
+        indices = unsqueeze(indices, dim)
+    return values, indices
+
+
+@register_lowering(aten.mode.default, type_promotion_kind=None)
+def mode_default(self, dim=-1, keepdim=False):
+    """Lower aten.mode via sort-based decomposition or fallback."""
+    if not config.triton.decompose_sort_ops:
+        return mode_fallback(self, dim, keepdim)
+    shape = self.get_size()
+    ndim = len(shape)
+    device = self.get_device()
+    if ndim == 0:
+        return clone(self), _full(0, device, torch.int64, shape)
+    dim = canonicalize_dim(ndim, dim)
+    sorted_vals, sorted_idxs = sort_stable(self, stable=True, dim=dim)
+    n = shape[dim]
+
+    # Position indices along dim: [0, 1, ..., n-1]
+    positions = iota(
+        n, start=0, step=1, dtype=torch.int64, device=device, requires_grad=False
+    )
+    pos_view_shape = [sympy.Integer(1)] * ndim
+    pos_view_shape[dim] = n
+    positions = view(positions, pos_view_shape)
+    positions = expand(positions, shape)
+
+    # Shift positions by -1, clamp to 0 for position 0
+    positions_loader0 = positions.make_loader()
+
+    def prev_pos_fn(idx):
+        return ops.maximum(
+            ops.sub(positions_loader0(idx), ops.constant(1, torch.int64)),
+            ops.constant(0, torch.int64),
+        )
+
+    prev_positions = Pointwise.create(
+        device=decode_device(device),
+        dtype=torch.int64,
+        inner_fn=prev_pos_fn,
+        ranges=shape,
+    )
+
+    # Gather shifted values and compare for run boundaries
+    shifted_vals = gather(sorted_vals, dim, prev_positions)
+
+    sorted_loader = sorted_vals.make_loader()
+    shifted_loader = shifted_vals.make_loader()
+    positions_loader = positions.make_loader()
+
+    # is_boundary = (sorted != shifted) | (position == 0)
+    def is_boundary_fn(idx):
+        return ops.or_(
+            ops.ne(sorted_loader(idx), shifted_loader(idx)),
+            ops.eq(positions_loader(idx), ops.constant(0, torch.int64)),
+        )
+
+    is_boundary = Pointwise.create(
+        device=decode_device(device),
+        dtype=torch.bool,
+        inner_fn=is_boundary_fn,
+        ranges=shape,
+    )
+
+    # boundary_pos = where(is_boundary, position, -1)
+    is_boundary_loader = is_boundary.make_loader()
+    positions_loader2 = positions.make_loader()
+
+    def boundary_pos_fn(idx):
+        return ops.where(
+            is_boundary_loader(idx),
+            positions_loader2(idx),
+            ops.constant(-1, torch.int64),
+        )
+
+    boundary_pos = Pointwise.create(
+        device=decode_device(device),
+        dtype=torch.int64,
+        inner_fn=boundary_pos_fn,
+        ranges=shape,
+    )
+
+    # Propagate boundary positions forward with cummax
+    last_boundary, _ = cummax(boundary_pos, dim)
+
+    # run_len = position - last_boundary + 1
+    positions_loader3 = positions.make_loader()
+    last_boundary_loader = last_boundary.make_loader()
+
+    def run_len_fn(idx):
+        return ops.add(
+            ops.sub(positions_loader3(idx), last_boundary_loader(idx)),
+            ops.constant(1, torch.int64),
+        )
+
+    run_len = Pointwise.create(
+        device=decode_device(device),
+        dtype=torch.int64,
+        inner_fn=run_len_fn,
+        ranges=shape,
+    )
+
+    # argmax returns first maximum -> end of leftmost longest run
+    max_pos = reduce_argmax(run_len, axis=dim, keepdims=True)
+    mode_vals = gather(sorted_vals, dim, max_pos)
+    mode_idxs = gather(sorted_idxs, dim, max_pos)
+
+    if not keepdim:
+        mode_vals = squeeze(mode_vals, dim)
+        mode_idxs = squeeze(mode_idxs, dim)
+
+    return mode_vals, mode_idxs
+
+
+@register_lowering(aten.topk.default, type_promotion_kind=None)
+def topk(self, k, dim=-1, largest=True, sorted=True):
+    if not config.triton.decompose_sort_ops:
+        return topk_fallback(self, k, dim, largest, sorted)
+    shape = self.get_size()
+    ndim = len(shape)
+    if ndim == 0:
+        return clone(self), _full(0, self.get_device(), torch.int64, shape)
+    dim = canonicalize_dim(ndim, dim)
+    sorted_vals, sorted_idxs = sort_stable(
+        self, stable=True, dim=dim, descending=largest
+    )
+    values = slice_(sorted_vals, dim, 0, k)
+    indices = slice_(sorted_idxs, dim, 0, k)
+    return values, indices
+
+
+@register_lowering(aten.kthvalue.default, type_promotion_kind=None)
+def kthvalue(self, k, dim=-1, keepdim=False):
+    if not config.triton.decompose_sort_ops:
+        return kthvalue_fallback(self, k, dim, keepdim)
+    shape = self.get_size()
+    ndim = len(shape)
+    if ndim == 0:
+        return clone(self), _full(0, self.get_device(), torch.int64, shape)
+    dim = canonicalize_dim(ndim, dim)
+    sorted_vals, sorted_idxs = sort_stable(self, stable=True, dim=dim)
+    # k is 1-based
+    values = select(sorted_vals, dim, k - 1)
+    indices = select(sorted_idxs, dim, k - 1)
+    if keepdim:
+        values = unsqueeze(values, dim)
+        indices = unsqueeze(indices, dim)
+    return values, indices
+
+
 def register_pointwise_numeric(op, name=None, triton_fallback=None):
     return register_pointwise(
         op,
@@ -6980,13 +7592,7 @@ def addcmul(self, tensor1, tensor2, *, value=1):
     to force rounding of the product before the FMA. This prevents Triton's
     compiler from fusing the multiplication with the FMA, matching eager's
     rounding behavior.
-
-    When emulate_precision_casts is False, we return NotImplemented to use the
-    decomposition instead.
     """
-    if not config.emulate_precision_casts:
-        return NotImplemented
-
     dtype = get_promoted_dtype(
         self,
         tensor1,
@@ -6998,8 +7604,14 @@ def addcmul(self, tensor1, tensor2, *, value=1):
     t1_loader = tensor1.make_loader()
     t2_loader = tensor2.make_loader()
 
-    # FMA is only available for floating-point types on non-AMD GPUs
-    use_fma = dtype.is_floating_point and not torch.version.hip
+    # FMA/mul_rn/div_rn are only available for floating-point types on CUDA (non-AMD)
+    device = self.get_device()
+    use_fma = (
+        dtype.is_floating_point
+        and not torch.version.hip
+        and device is not None
+        and device.type in ["cuda", "xpu"]
+    )
 
     def inner_fn(idx):
         self_val = self_loader(idx)
@@ -7039,32 +7651,86 @@ def inner_fn(idx):
     )
 
 
-def _foreach_addcmul_scalar(self, tensor1, tensor2, value=1):
+@register_lowering(aten.addcdiv, broadcast=True)
+def addcdiv(self, tensor1, tensor2, *, value=1):
     """
-    Foreach version of addcmul with scalar value parameter.
-    Uses foreach_group_loop for consistent grouping behavior.
+    Computes self + value * (tensor1 / tensor2) using FMA for better precision.
+
+    Matches eager CUDA kernel order: self + value * (tensor1 / tensor2)
+    This is computed as: fma(value, tensor1 / tensor2, self)
+
+    For value=1: self + tensor1 / tensor2 (no FMA needed, just add the division)
+    For value!=1: fma(value, div_rn(tensor1, tensor2), self)
+
+    Note: FMA is only used for floating-point types on non-AMD GPUs. For integer types,
+    we fall back to regular arithmetic since FMA doesn't support integers.
 
-    When emulate_precision_casts is False, we return NotImplemented to use the
-    decomposition instead.
+    We use div_rn (round-to-nearest division) to force proper rounding, preventing
+    Triton from fusing operations in ways that change the rounding behavior.
     """
-    if not config.emulate_precision_casts:
-        return NotImplemented
+    dtype = get_promoted_dtype(
+        self,
+        tensor1,
+        tensor2,
+        type_promotion_kind=ELEMENTWISE_TYPE_PROMOTION_KIND.INT_TO_FLOAT,
+    )
 
-    realize_outputs = (
-        len(V.graph.current_node.users) == 0
-        or V.graph.current_node.target in inplace_foreach_ops
-        or cur_node_has_non_foreach_users()
+    self_loader = self.make_loader()
+    t1_loader = tensor1.make_loader()
+    t2_loader = tensor2.make_loader()
+
+    # FMA/mul_rn/div_rn are only available for floating-point types on CUDA (non-AMD)
+    device = self.get_device()
+    use_fma = (
+        dtype.is_floating_point
+        and not torch.version.hip
+        and device is not None
+        and device.type in ["cuda", "xpu"]
     )
 
-    groups = group_foreach_args(zip(self, tensor1, tensor2))
+    def inner_fn(idx):
+        self_val = self_loader(idx)
+        t1_val = t1_loader(idx)
+        t2_val = t2_loader(idx)
 
-    def apply_fn(args):
-        return addcmul(*args, value=value)
+        # Compute tensor1 / tensor2 first
+        # Use div_rn for round-to-nearest division on CUDA to match eager behavior
+        if use_fma:
+            t1_div_t2 = ops.div_rn(t1_val, t2_val)
+        else:
+            t1_div_t2 = ops.truediv(t1_val, t2_val)
 
-    return foreach_group_loop(groups, len(self), apply_fn, realize_outputs)
+        if value == 1:
+            # For value=1, just add the division result (no FMA needed)
+            return ops.add(self_val, t1_div_t2)
 
+        # Use index_expr for sympy expressions (e.g., from .item()), constant otherwise
+        if isinstance(value, sympy.Basic):
+            value_expr = ops.index_expr(value, dtype)
+        else:
+            value_expr = ops.constant(value, dtype)
 
-_register_foreach_lowering(aten._foreach_addcmul.Scalar, _foreach_addcmul_scalar)
+        if use_fma:
+            # Use FMA for floating-point types for better precision
+            return ops.fma(value_expr, t1_div_t2, self_val)
+        else:
+            # Fall back to regular arithmetic for integer types
+            return ops.add(self_val, ops.mul(value_expr, t1_div_t2))
+
+    return Pointwise.create(
+        device=self.get_device(),
+        dtype=dtype,
+        inner_fn=inner_fn,
+        ranges=self.get_size(),
+    )
+
+
+_foreach_addcmul_scalar = register_foreach_pointwise(
+    aten._foreach_addcmul.Scalar, addcmul, allow_alpha=True, scalar_kwarg="value"
+)
+_foreach_addcdiv_scalar = register_foreach_pointwise(
+    aten._foreach_addcdiv.Scalar, addcdiv, allow_alpha=True, scalar_kwarg="value"
+)
 
 
 register_pointwise_numeric_ldf64(aten.cos)
@@ -7225,6 +7891,7 @@ def make_triton_fallback(op):
 register_foreach_pointwise(aten._foreach_clamp_max.Scalar, minimum)
 register_foreach_pointwise(aten._foreach_reciprocal, reciprocal)
 register_foreach_pointwise(aten._foreach_sign, sign)
+register_foreach_pointwise(aten._foreach_clone, clone)
 foreach_copy = register_foreach_pointwise(aten._foreach_copy, copy)
 
 
@@ -7267,6 +7934,16 @@ def fn(*args, **kwargs):
 register_foreach_inplace(
     aten._foreach_copy_.default, aten._foreach_copy.default, foreach_copy
 )
+register_foreach_inplace(
+    aten._foreach_addcmul_.Scalar,
+    aten._foreach_addcmul.Scalar,
+    _foreach_addcmul_scalar,
+)
+register_foreach_inplace(
+    aten._foreach_addcdiv_.Scalar,
+    aten._foreach_addcdiv.Scalar,
+    _foreach_addcdiv_scalar,
+)
 
 
 def register_inplace(aten_op, outplace_op):
@@ -7344,7 +8021,11 @@ def sym_numel(a):
 
 
 @register_lowering(torch.sym_sum)
-def sym_sum(args):
+def sym_sum(*args):
+    # sym_sum can be called as sym_sum([a, b]) or sym_sum(a, b).
+    # Normalize to a flat list before summing.
+    if len(args) == 1 and isinstance(args[0], (list, tuple)):
+        args = args[0]
     return sympy.Add(*args)
 
 
@@ -7405,9 +8086,6 @@ def resize(x, size, *, memory_format=None):
     dtype = x.get_dtype()
     device = x.get_device_or_error()
 
-    if isinstance(x.data, ir.BaseView):
-        x.data = x.data.unwrap_view()
-
     if (
         torch.are_deterministic_algorithms_enabled()
         and torch.utils.deterministic.fill_uninitialized_memory  # type: ignore[attr-defined]
@@ -7425,15 +8103,18 @@ def resize(x, size, *, memory_format=None):
     if V.graph.sizevars.statically_known_equals(old_numel, 0):  # type: ignore[arg-type]
         return full(size, uninitialized_val, dtype=dtype, device=device)
 
-    x_flat = as_strided(
-        x,
-        [
-            old_numel,
-        ],
-        [
-            1,
-        ],
+    strides = x.maybe_get_stride()
+    has_overlapping = strides is not None and any(
+        V.graph.sizevars.statically_known_equals(s, 0) for s in strides
     )
+    if has_overlapping:
+        # overlapping: provide a contiguous logical view
+        x_flat = view(x, [old_numel])
+    else:
+        # non-overlapping: keep storage order
+        if isinstance(x.data, ir.BaseView):
+            x.data = x.data.unwrap_view()
+        x_flat = as_strided(x, [old_numel], [1])
     flat_loader = x_flat.make_loader()
     out_stride = ir.FlexibleLayout.stride_ordered_for_memory_format(size, memory_format)
     out_indexer = ir.FixedLayout(device, dtype, size, out_stride).make_indexer()
@@ -7481,7 +8162,7 @@ def triton_kernel_wrap_(
 @register_lowering(torch.ops.higher_order.cond, type_promotion_kind=None)
 def cond(
     pred, true_fn, false_fn, operands
-) -> list[Union[ir.TensorBox, ir.ShapeAsConstantBuffer]]:
+) -> list[ir.TensorBox | ir.ShapeAsConstantBuffer]:
     # TODO: when graph_partition is enabled, skip - partitioning handles control flow
     # we run into memory cleanup issue
     if any(isinstance(x, IRNode) and is_triton(x) for x in [pred, *operands]):
@@ -7525,6 +8206,9 @@ def invoke_subgraph(subgraph_fn: ir.Subgraph, identifier: str, *operands):
     return list(map(TensorBox.create, result))  # type: ignore[call-overload]
 
 
+_MISSING = object()
+
+
 def process_subgraph_nodes(graph_module: torch.fx.GraphModule, args: list[Any]):
     """Process nodes from a FX graph by executing them through V.graph.
 
@@ -7534,7 +8218,7 @@ def process_subgraph_nodes(graph_module: torch.fx.GraphModule, args: list[Any]):
     - Other nodes are executed via V.graph.run_node
 
     """
-    output = None
+    output = _MISSING
 
     for i, node in enumerate(graph_module.graph.nodes):
         if node.op == "placeholder":
@@ -7554,7 +8238,7 @@ def process_subgraph_nodes(graph_module: torch.fx.GraphModule, args: list[Any]):
             finally:
                 V.graph.current_node = saved_current_node
 
-    if output is None:
+    if output is _MISSING:
         raise RuntimeError("No output node found in graph")
 
     return output
@@ -7593,7 +8277,7 @@ def control_deps_op_lowering(additional_deps, subgraph_fn, *args):
     # Process subgraph nodes using the shared helper
     output = process_subgraph_nodes(subgraph_fn.graph_module, list(args))
 
-    assert output is not None and additional_deps
+    assert additional_deps
 
     # some operators, like wait_tensor, just return their input,
     # so its more robust to add dep to the operation itself,
@@ -7723,7 +8407,7 @@ def with_effects(token, op, *args, **kwargs):
     else:
 
         def wrap_tensors(x):
-            return TensorBox.create(x) if isinstance(x, ir.IRNode) else x
+            return x.wrap_for_lowering() if isinstance(x, ir.IRNode) else x
 
         result = pytree.tree_map(
             wrap_tensors, ir.FallbackKernel.create(op, *args, **kwargs)
@@ -7833,15 +8517,8 @@ def prepare_softmax_online(x, dim):
         #
         # TODO: does inference need split online_softmax_reduce?
 
-        warnings.warn(
-            textwrap.dedent(
-                """
-            Online softmax is disabled on the fly since Inductor decides to
-            split the reduction. Cut an issue to PyTorch if this is an
-            important use case and you want to speed it up with online
-            softmax.
-            """
-            )
+        log.debug(
+            "Online softmax is disabled on the fly since Inductor decides to split the reduction."
         )
         amax = reduce_amax(x, dim, keepdims=True)
         exp = lowerings[aten.exp](sub(x, amax))
@@ -7892,6 +8569,43 @@ def cvt_e8m0_rceil_lowering(inp):
     return to_dtype(result, torch.uint8)
 
 
+@register_lowering(
+    torch._higher_order_ops.inline_asm_elementwise, type_promotion_kind=None
+)
+def lower_inline_asm_elementwise(
+    *inputs, asm_str, constraints, dtype, is_pure=True, pack=1
+):
+    inputs = broadcast_tensors(*inputs)
+
+    input_dtypes = tuple(inp.get_dtype() for inp in inputs)
+    loaders = [inp.make_loader() for inp in inputs]
+
+    def inner_fn(idx):
+        vals = tuple(loader(idx) for loader in loaders)
+        result = ops.inline_asm_elementwise(
+            *vals,
+            asm=asm_str,
+            constraints=constraints,
+            dtype=dtype,
+            is_pure=is_pure,
+            pack=pack,
+            input_dtypes=input_dtypes,
+        )
+        # Inductor computes in fp32 for bf16/fp16. Upcast so fused downstream
+        # ops (reductions, etc.) see fp32 values. The Pointwise's storage dtype
+        # handles the final downcast on store.
+        if dtype in (torch.float16, torch.bfloat16):
+            result = ops.to_dtype(result, torch.float32)
+        return result
+
+    return ir.Pointwise.create(
+        device=inputs[0].get_device(),
+        dtype=dtype,
+        inner_fn=inner_fn,
+        ranges=list(inputs[0].get_size()),
+    )
+
+
 # populate lowerings defined in kernel/*
 from . import kernel
 
diff --git a/torch/_inductor/memory.py b/torch/_inductor/memory.py
index 4f587b23cda0c..bcd047e40f5ad 100644
--- a/torch/_inductor/memory.py
+++ b/torch/_inductor/memory.py
@@ -4,7 +4,7 @@
 import dataclasses
 import heapq
 import logging
-from typing import Optional, TYPE_CHECKING, TypedDict, Union
+from typing import TYPE_CHECKING, TypedDict
 
 import torch
 from torch._environment import is_fbcode
@@ -61,8 +61,8 @@ def __post_init__(self) -> None:
 class MemoryPlanningInfoForNode:
     index: int = 0
     size: int = 0
-    pred_buffers: OrderedSet[Union[SchedulerBuffer, FreeableInputBuffer]] = (
-        dataclasses.field(default_factory=OrderedSet)
+    pred_buffers: OrderedSet[SchedulerBuffer | FreeableInputBuffer] = dataclasses.field(
+        default_factory=OrderedSet
     )
     pred_nodes: OrderedSet[BaseSchedulerNode] = dataclasses.field(
         default_factory=OrderedSet
@@ -205,7 +205,7 @@ def _compute_and_update_buf_size(
             )
             return size_alloc
         else:
-            buf_size = V.graph.sizevars.size_hint(
+            buf_size = V.graph.sizevars.optimization_hint(
                 sched_buf.node.get_numel(), fallback=0
             ) * get_dtype_size(sched_buf.node.get_dtype())
             sched_buf_to_size[sched_buf.get_name()] = (
@@ -336,7 +336,7 @@ def assign_memory_planning_info_for_scheduler_nodes(
 # map each scheduler buffer to its size, start step, and end step
 @dataclasses.dataclass
 class BufferInfo:
-    buffer: Union[SchedulerBuffer, FreeableInputBuffer]
+    buffer: SchedulerBuffer | FreeableInputBuffer
     size_alloc: int
     size_free: int
     start_step: int
@@ -350,7 +350,7 @@ def compute_memory_timeline(
 ) -> tuple[
     list[BufferInfo],
     dict[BaseSchedulerNode, int],
-    dict[Union[FreeableInputBuffer, SchedulerBuffer], BaseSchedulerNode],
+    dict[FreeableInputBuffer | SchedulerBuffer, BaseSchedulerNode],
 ]:
     """
     Compute buffer allocation and deallocation sizes and map their
@@ -366,14 +366,14 @@ def compute_memory_timeline(
     # get buffers' size and liveliness information
     buf_info_list: list[BufferInfo] = []
     buf_to_snode_last_use: dict[
-        Union[FreeableInputBuffer, SchedulerBuffer], BaseSchedulerNode
+        FreeableInputBuffer | SchedulerBuffer, BaseSchedulerNode
     ] = {}
 
     def _get_end_step_and_snode(
-        buf: Union[FreeableInputBuffer, SchedulerBuffer],
-    ) -> tuple[int, Optional[BaseSchedulerNode]]:
+        buf: FreeableInputBuffer | SchedulerBuffer,
+    ) -> tuple[int, BaseSchedulerNode | None]:
         max_step: int = -1
-        max_step_snode: Optional[BaseSchedulerNode] = None
+        max_step_snode: BaseSchedulerNode | None = None
         succ_nodes = buf.mpi_buffer.succ_nodes
         if succ_nodes:
             for succ_node in succ_nodes:
@@ -484,7 +484,7 @@ def estimate_peak_memory_allocfree(
     int,
     list[tuple[int, int]],
     dict[BaseSchedulerNode, SNodeMemory],
-    dict[Union[FreeableInputBuffer, SchedulerBuffer], BaseSchedulerNode],
+    dict[FreeableInputBuffer | SchedulerBuffer, BaseSchedulerNode],
 ]:
     """
     Alternative version of estimate_peak_memory, that respects the fact,
@@ -571,7 +571,7 @@ class BufferInfo(TypedDict):
         outdegree: int
 
     node_info: dict[BaseSchedulerNode, NodeInfo] = dict()
-    buf_info: dict[Union[SchedulerBuffer, FreeableInputBuffer], BufferInfo] = dict()
+    buf_info: dict[SchedulerBuffer | FreeableInputBuffer, BufferInfo] = dict()
 
     # compute nodes' number of unmet dependencies (for schedulability)
     # initialize the list of nodes ready to be scheduled
diff --git a/torch/_inductor/metrics.py b/torch/_inductor/metrics.py
index 4df5084adde76..47c3eac2b7e58 100644
--- a/torch/_inductor/metrics.py
+++ b/torch/_inductor/metrics.py
@@ -7,7 +7,7 @@
 import re
 from dataclasses import dataclass
 from functools import lru_cache
-from typing import Optional, TYPE_CHECKING, Union
+from typing import TYPE_CHECKING
 
 from torch._inductor import config
 from torch._inductor.utils import get_benchmark_name
@@ -146,9 +146,7 @@ class MetricTable:
 
     num_rows_added: int = 0
 
-    def add_row(
-        self, row_fn: Callable[[], dict[str, Optional[Union[str, float]]]]
-    ) -> None:
+    def add_row(self, row_fn: Callable[[], dict[str, str | float | None]]) -> None:
         if self.table_name not in enabled_metric_tables():
             return
 
@@ -308,7 +306,7 @@ def _parse_kernel_line_of_code(proper_kernel_fn_code: str) -> int:
     return len(proper_kernel_fn_code.splitlines())
 
 
-def _parse_size_hints(kernel_module_code: str, kernel_category: str) -> Optional[str]:
+def _parse_size_hints(kernel_module_code: str, kernel_category: str) -> str | None:
     if kernel_category == "foreach":
         # foreach kernel does not have size_hints
         return None
@@ -317,9 +315,7 @@ def _parse_size_hints(kernel_module_code: str, kernel_category: str) -> Optional
     return m.group(1)
 
 
-def _parse_reduction_hint(
-    kernel_category: str, kernel_module_code: str
-) -> Optional[str]:
+def _parse_reduction_hint(kernel_category: str, kernel_module_code: str) -> str | None:
     if kernel_category not in ("reduction", "persistent_reduction"):
         return None
     m = re.search(r"reduction_hint=ReductionHint\.(\w*),", kernel_module_code)
@@ -349,7 +345,7 @@ def _parse_proper_kernel_fn_code(kernel_fn_code: str) -> str:
     return kernel_fn_code[start_pos:]
 
 
-def _parse_numel(proper_kernel_fn_code: str, numel_arg_name: str) -> Optional[int]:
+def _parse_numel(proper_kernel_fn_code: str, numel_arg_name: str) -> int | None:
     m = re.search(f"{numel_arg_name} = ([\\d]+)", proper_kernel_fn_code)
     if m:
         return int(m.group(1))
@@ -359,7 +355,7 @@ def _parse_numel(proper_kernel_fn_code: str, numel_arg_name: str) -> Optional[in
 
 def _parse_kernel_args_num_gb(
     kernel_fn_code: str, kernel_category: str
-) -> Optional[float]:
+) -> float | None:
     """
     inductor meta looks like:
         inductor_meta={... 'mutated_arg_names': [], 'no_x_dim': False, 'kernel_num_gb': 2.0},
diff --git a/torch/_inductor/mkldnn_ir.py b/torch/_inductor/mkldnn_ir.py
index 4c09332df8eb7..091e29c10e31e 100644
--- a/torch/_inductor/mkldnn_ir.py
+++ b/torch/_inductor/mkldnn_ir.py
@@ -38,8 +38,8 @@ def _prepare_convolution_fusion_create(
     dilation: Sequence[int],
     groups: int,
     transposed: bool = False,
-    output_padding: Optional[Sequence[int]] = None,
-    quantize_args: Optional[list["TensorBox"]] = None,
+    output_padding: Sequence[int] | None = None,
+    quantize_args: list["TensorBox"] | None = None,
     other: Optional["TensorBox"] = None,
 ):
     """
@@ -117,8 +117,8 @@ def _original_deconv_weight_size(
         bias.realize()
     with V.graph.fake_mode:
         # TODO <Leslie> cleaned up the fake_tensor trace as Linear implementation
-        x_fake = ir_node_to_tensor(x, guard_shape=True)
-        weight_fake = ir_node_to_tensor(weight, guard_shape=True)
+        x_fake = ir_node_to_tensor(x)
+        weight_fake = ir_node_to_tensor(weight)
         dims = len(x_fake.size()) - 2
         assert 0 < len(padding) <= dims
         assert 0 < len(dilation) <= dims
@@ -230,7 +230,7 @@ def _prepare_linear_fusion_create(
     x: "TensorBox",
     weight: "TensorBox",
     bias: "TensorBox",
-    quantize_args: Optional[list["TensorBox"]] = None,
+    quantize_args: list["TensorBox"] | None = None,
     other: Optional["TensorBox"] = None,
     binary_sum: bool = False,
 ):
@@ -333,7 +333,7 @@ def create(
         dilation_: list[int],
         groups: int,
         attr,
-        scalars: Optional[list[Any]],
+        scalars: list[Any] | None,
         algorithm,
     ):
         (
@@ -395,10 +395,10 @@ def create(
         dilation_: list[int],
         groups: int,
         binary_attr: str,
-        binary_alpha: Optional[float],
-        unary_attr: Optional[str],
-        unary_scalars: Optional[list[Any]],
-        unary_algorithm: Optional[str],
+        binary_alpha: float | None,
+        unary_attr: str | None,
+        unary_scalars: list[Any] | None,
+        unary_algorithm: str | None,
     ):
         (
             inputs,
@@ -473,10 +473,10 @@ def create(
         dilation_: list[int],
         groups: int,
         binary_attr: str,
-        binary_alpha: Optional[float],
-        unary_attr: Optional[str],
-        unary_scalars: Optional[list[Any]],
-        unary_algorithm: Optional[str],
+        binary_alpha: float | None,
+        unary_attr: str | None,
+        unary_scalars: list[Any] | None,
+        unary_algorithm: str | None,
     ):
         (
             inputs,
@@ -543,7 +543,7 @@ def create(
         dilation_: list[int],
         groups_: int,
         attr,
-        scalars: Optional[list[Any]],
+        scalars: list[Any] | None,
         algorithm,
     ):
         transposed = True
diff --git a/torch/_inductor/mkldnn_lowerings.py b/torch/_inductor/mkldnn_lowerings.py
index f14732d0eeaca..e193fe425dfe6 100644
--- a/torch/_inductor/mkldnn_lowerings.py
+++ b/torch/_inductor/mkldnn_lowerings.py
@@ -1,6 +1,5 @@
 # mypy: allow-untyped-defs
 import functools
-from typing import Optional, Union
 
 import torch
 import torch.utils._pytree as pytree
@@ -38,9 +37,9 @@ def create_int8_compensation(
 ) -> tuple[
     bool,
     ir.TensorBox,
-    Optional[ir.TensorBox],
+    ir.TensorBox | None,
 ]:
-    x_w_scale: Optional[ir.TensorBox] = None
+    x_w_scale: ir.TensorBox | None = None
     use_int8_fast_compensation_path = all(
         isinstance(item, ir.TensorBox)
         and item.get_name() in V.graph.constants
@@ -81,10 +80,10 @@ def codegen_int8_gemm_template_compensation(
     use_int8_fast_compensation_path: bool,
     input: OpsValue,
     _weight_compo: OpsValue,
-    _x_scale: Optional[OpsValue],
-    _x_zp: Optional[OpsValue],
-    _w_scale: Optional[OpsValue],
-    _x_w_scale: Optional[OpsValue],
+    _x_scale: OpsValue | None,
+    _x_zp: OpsValue | None,
+    _w_scale: OpsValue | None,
+    _x_w_scale: OpsValue | None,
 ) -> OpsValue:
     if use_int8_fast_compensation_path:
         temp = ops.sub(
@@ -166,7 +165,7 @@ def grouped_gemm_lowering(
     )
 
     assert len(choices) != 0
-    result = autotune_select_algorithm(
+    result, _ = autotune_select_algorithm(
         "grouped_gemm",
         choices,
         input_nodes,
@@ -224,7 +223,7 @@ def register_onednn_fusion_ops():
             kernel_creator=mkldnn_ir.QLinearPointwiseBinaryPT2E.create,
         )
         cpu_needs_realized_inputs: list[
-            Union[torch._ops.OpOverload, torch._ops.OpOverloadPacket]
+            torch._ops.OpOverload | torch._ops.OpOverloadPacket
         ] = [
             torch.ops.mkldnn._convolution_pointwise,
             torch.ops.mkldnn._convolution_pointwise_,
@@ -387,7 +386,7 @@ def epilogue_creator(buf):
             input_gen_fns = {
                 1: lambda x: V.graph.constants[x.get_name()],
             }
-            result = autotune_select_algorithm(
+            result, _ = autotune_select_algorithm(
                 "linear_unary",
                 choices,
                 [x, w] if b is None else [x, w, b],
@@ -451,7 +450,7 @@ def epilogue_creator(buf):
             input_gen_fns = {
                 2: lambda x: V.graph.constants[x.get_name()],
             }
-            result = autotune_select_algorithm(
+            result, _ = autotune_select_algorithm(
                 "linear_binary",
                 choices,
                 [x, y, w] if b is None else [x, y, w, b],
@@ -975,7 +974,7 @@ def inner_fn_requant(index, scale, zero_point):
             ):
                 input_gen_fns[2] = lambda x: V.graph.constants[x.get_name()]
 
-            result = autotune_select_algorithm(
+            result, _ = autotune_select_algorithm(
                 "qlinear_unary",
                 choices,
                 [x, x_scale, x_zp, packed_weight, w_scale, w_zp]
@@ -1305,7 +1304,7 @@ def inner_fn_requant(index, scale, zero_point):
             }
             if bias is not None:
                 input_gen_fns[7] = lambda x: V.graph.constants[x.get_name()]  # For bias
-            result = autotune_select_algorithm(
+            result, _ = autotune_select_algorithm(
                 "qlinear_binary",
                 choices,
                 [x, x_scale, x_zp, packed_weight, w_scale, w_zp, x2]
@@ -1350,7 +1349,7 @@ def mkl_packed_linear(
                 x: TensorBox,
                 packed_w: TensorBox,
                 orig_w: TensorBox,
-                b: Optional[TensorBox],
+                b: TensorBox | None,
                 batch_size,
                 *,
                 layout=None,
@@ -1385,7 +1384,8 @@ def mkl_packed_linear(
                     1: lambda x: V.graph.constants[x.get_name()],
                     2: lambda x: V.graph.constants[x.get_name()],
                 }
-                result: TensorBox = autotune_select_algorithm(
+                result: TensorBox  # annotation on separate line since tuple unpacking doesn't support inline annotation
+                result, _ = autotune_select_algorithm(
                     "packed_linear",
                     choices,
                     [x, packed_w, orig_w],
diff --git a/torch/_inductor/mock_cache.py b/torch/_inductor/mock_cache.py
index 3d9c58f1db8bd..49c41c90a5d70 100644
--- a/torch/_inductor/mock_cache.py
+++ b/torch/_inductor/mock_cache.py
@@ -6,7 +6,7 @@
 import dataclasses
 import sys
 import threading
-from typing import Any, Optional, TYPE_CHECKING
+from typing import Any, TYPE_CHECKING
 from typing_extensions import override, Self
 from unittest.mock import patch
 
@@ -137,7 +137,7 @@ def wrapper() -> MockBackend:
         return wrapper
 
     @override
-    def _get(self, key: str) -> Optional[Any]:
+    def _get(self, key: str) -> Any | None:
         stat = global_stats.get_stat(self._name)
         if key in stat.cache:
             stat += Stats(num_get_hit=1)
@@ -267,8 +267,8 @@ def __enter__(self) -> Self:
 
     def __exit__(
         self,
-        exc_type: Optional[type[BaseException]],
-        exc_value: Optional[BaseException],
-        traceback: Optional[TracebackType],
+        exc_type: type[BaseException] | None,
+        exc_value: BaseException | None,
+        traceback: TracebackType | None,
     ) -> None:
         self._stack.__exit__(exc_type, exc_value, traceback)
diff --git a/torch/_inductor/ops_handler.py b/torch/_inductor/ops_handler.py
index 290ac983e0b4a..510924f214c37 100644
--- a/torch/_inductor/ops_handler.py
+++ b/torch/_inductor/ops_handler.py
@@ -6,16 +6,7 @@
 import re
 import warnings
 from io import StringIO
-from typing import (
-    Any,
-    Generic,
-    Literal,
-    NamedTuple,
-    Optional,
-    TYPE_CHECKING,
-    TypeVar,
-    Union,
-)
+from typing import Any, Generic, Literal, NamedTuple, TYPE_CHECKING, TypeVar
 from unittest.mock import patch
 
 import sympy
@@ -42,7 +33,7 @@
     "atomic_cas",
     "atomic_xchg",
 ]
-StoreMode = Optional[Union[AtomicMode, Literal["tma"]]]
+StoreMode = AtomicMode | Literal["tma"] | None
 ReductionType = Literal[
     "argmax",
     "argmin",
@@ -95,7 +86,7 @@ class OpsHandler(Generic[T]):
     all the metaprogramming has run.
     """
 
-    def constant(self, value: Union[bool, float, int], dtype: torch.dtype) -> T:
+    def constant(self, value: bool | float | int, dtype: torch.dtype) -> T:
         """Produces a scalar constant of type dtype."""
         raise NotImplementedError
 
@@ -107,6 +98,12 @@ def rand(self, seed: T, offset: T) -> T:
         """Computes inductor_prims.random with mode="rand".  offset has dtype int32."""
         raise NotImplementedError
 
+    def rand_eager(
+        self, seed: T, base_offset: T, threads_per_round: T, tid: T, vec: T
+    ) -> T:
+        """Computes inductor_prims.random with mode="rand_eager".  offset has dtype int32."""
+        raise NotImplementedError
+
     def randn(self, seed: T, offset: T) -> T:
         """Computes inductor_prims.random with mode="randn".  offset has dtype int32."""
         raise NotImplementedError
@@ -147,7 +144,7 @@ def to_dtype(
         self,
         x: T,
         dtype: torch.dtype,
-        src_dtype: Optional[torch.dtype] = None,
+        src_dtype: torch.dtype | None = None,
         use_compute_types: bool = True,
     ) -> T:
         """
@@ -251,7 +248,7 @@ def reduction(
         src_dtype: torch.dtype,
         reduction_type: ReductionType,
         value: T,
-    ) -> Union[T, tuple[T, ...]]:
+    ) -> T | tuple[T, ...]:
         """
         Perform a 'reduction_type' reduction on 'value' of dtype 'src_dtype',
         using 'dtype' as the accumulation dtype for the reduction.  The result
@@ -305,8 +302,8 @@ def bucketize(
         boundary_indices: T,
         indexing_dtype: torch.dtype,
         right: bool,
-        sorter: Optional[tuple[str, sympy.Expr]] = None,
-        sorter_indices: Optional[T] = None,
+        sorter: tuple[str, sympy.Expr] | None = None,
+        sorter_indices: T | None = None,
     ) -> T:
         # See [Note: Inductor bucketize op]
         raise NotImplementedError
@@ -699,6 +696,11 @@ def truediv(self, x0: T, x1: T) -> T:
         do Python-style (int, int) -> float division, use int_truediv"""
         raise NotImplementedError
 
+    def div_rn(self, x0: T, x1: T) -> T:
+        """Division with round-to-nearest rounding mode.  Used for matching
+        eager CUDA semantics where division uses IEEE round-to-nearest."""
+        raise NotImplementedError
+
     def int_truediv(self, x0: T, x1: T) -> T:
         """True division between integers.  This is NOT the same as promoting
         to float and doing integer division, there is a bespoke algorithm for
@@ -735,10 +737,11 @@ def inline_asm_elementwise(
         self,
         *inputs: T,
         asm: str,
-        constraints: Optional[str] = None,
+        constraints: str | None = None,
         dtype: torch.dtype = torch.float32,
         is_pure: bool = True,
         pack: int = 1,
+        input_dtypes: tuple[torch.dtype, ...] | None = None,
     ) -> T:
         raise NotImplementedError
 
@@ -850,6 +853,7 @@ def masked(mask, body, other) -> None:
         return None
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def frexp(x) -> tuple[None, None]:
         return (None, None)
 
@@ -862,6 +866,7 @@ def sort(dtypes, values, stable, descending) -> tuple[None, ...]:
         return (None,) * len(values)
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def indirect_indexing(index_var, size, check=True, wrap_neg=True) -> sympy.Symbol:
         return sympy.S.Zero
 
@@ -959,6 +964,7 @@ def masked(mask, body, other) -> str:
         return f"ops.masked({mask}, {body()}, {other})"
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def frexp(x):
         return (f"ops.frexp({x})[0]", f"ops.frexp({x})[1]")
 
@@ -977,6 +983,7 @@ def sort(dtypes, values, stable, descending):
         )
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def indirect_indexing(index_var, size, check=True, wrap_neg=True) -> sympy.Symbol:
         return sympy_index_symbol(str(index_var))
 
@@ -1034,8 +1041,8 @@ def reduction(
         dtype: torch.dtype,
         src_dtype: torch.dtype,
         reduction_type: ReductionType,
-        value: Union[str, tuple[str, ...]],
-    ) -> Union[str, tuple[str, ...]]:
+        value: str | tuple[str, ...],
+    ) -> str | tuple[str, ...]:
         line = self.parent_handler.reduction(dtype, src_dtype, reduction_type, value)
         num_values = reduction_num_outputs(reduction_type)
         varnames = [f"tmp{next(self.var_counter)}" for _ in range(num_values)]
@@ -1123,8 +1130,8 @@ def bucketize(
         boundary_indices: T,
         indexing_dtype: torch.dtype,
         right: bool,
-        sorter: Optional[tuple[str, sympy.Expr]] = None,
-        sorter_indices: Optional[T] = None,
+        sorter: tuple[str, sympy.Expr] | None = None,
+        sorter_indices: T | None = None,
     ) -> T:
         """
         See [Note: Inductor bucketize op]
@@ -1152,7 +1159,7 @@ def getvalue(self):
 
 
 class ExtractConstantsHandler(NoopHandler):
-    def __init__(self, device: Optional[torch.device]):
+    def __init__(self, device: torch.device | None):
         self.device = device
 
     def constant(self, value: Any, dtype: torch.dtype) -> torch._inductor.ir.Constant:
@@ -1172,7 +1179,7 @@ class SimpleCSEHandler(WrapperHandler):
 
     def __init__(self, inner: Any):
         super().__init__(inner)
-        self.cse_cache: dict[str, Union[Any, tuple[Any, ...]]] = {}
+        self.cse_cache: dict[str, Any | tuple[Any, ...]] = {}
         self.mock = MockHandler()
 
     def indirect_indexing(self, *args, **kwargs) -> sympy.Expr:
diff --git a/torch/_inductor/output_code.py b/torch/_inductor/output_code.py
index 37bf8fe47dff7..6d1279d22883a 100644
--- a/torch/_inductor/output_code.py
+++ b/torch/_inductor/output_code.py
@@ -26,7 +26,7 @@
 import logging
 import os
 from functools import partial
-from typing import Any, Optional, TYPE_CHECKING, TypeAlias, Union
+from typing import Any, TYPE_CHECKING, TypeAlias
 
 import torch
 from torch._dynamo.utils import counters, get_runtime_metrics_context
@@ -50,6 +50,7 @@
     output_node,
     set_tracing_context_output_strides,
 )
+from torch._opaque_base import OpaqueBase
 from torch.fx._graph_pickler import _node_metadata_key_filter_safe, _ops_filter_safe
 from torch.utils._ordered_set import OrderedSet
 from torch.utils._python_dispatch import is_in_torch_dispatch_mode
@@ -78,13 +79,13 @@ class OutputCode:
     # TODO: Remove underscores here
 
     # None if the output is not remote cacheable
-    _fx_graph_cache_key: Optional[str] = dataclasses.field(default=None, init=False)
-    _fx_graph_cache_debug_lines: Optional[list[str]] = dataclasses.field(
+    _fx_graph_cache_key: str | None = dataclasses.field(default=None, init=False)
+    _fx_graph_cache_debug_lines: list[str] | None = dataclasses.field(
         default=None, init=False
     )
 
     # How long it took to compile this OutputCode, end to end
-    _time_taken_ns: Optional[int] = dataclasses.field(default=None, init=False)
+    _time_taken_ns: int | None = dataclasses.field(default=None, init=False)
 
     def __call__(self, inputs: Sequence[Any]) -> Any:
         raise NotImplementedError(type(self))
@@ -149,7 +150,7 @@ def complex_memory_overlap(t: torch.Tensor) -> bool:
 
 def maybe_handle_backward_generation(
     compiled_graph: CompiledFxGraph,
-    boxed_forward_device_index: Optional[BoxedDeviceIndex],
+    boxed_forward_device_index: BoxedDeviceIndex | None,
 ) -> None:
     assert compiled_graph.current_callable is not None
     is_backward = compiled_graph.fx_kwargs["is_backward"]
@@ -178,7 +179,7 @@ def compiled_artifact(new_inputs: list[Any]) -> Callable[..., Any]:
 def prepare_cudagraph_post_compile(
     compiled_graph: CompiledFxGraph,
     example_inputs: Sequence[InputType],
-    boxed_forward_device_index: Optional[BoxedDeviceIndex],
+    boxed_forward_device_index: BoxedDeviceIndex | None,
 ) -> None:
     if not config.triton.cudagraph_trees:
         # Force specialize all inputs so that CUDA graphs will work
@@ -196,8 +197,8 @@ def cudagraph_post_compile(
     example_inputs: Sequence[InputType],
     compiled_graph: CompiledFxGraph,
     cudagraphs: BoxedBool,
-    constants: dict[str, Union[torch.Tensor, type]],
-    boxed_forward_device_index: Optional[BoxedDeviceIndex],
+    constants: dict[str, torch.Tensor | type],
+    boxed_forward_device_index: BoxedDeviceIndex | None,
 ) -> None:
     """
     Checks for any reasons not to run cudagraphs and then
@@ -226,23 +227,24 @@ def cudagraph_post_compile(
 
         placeholders = cached_info.placeholders
         stack_traces = cached_info.stack_traces
+        assert stack_traces is not None, (
+            "stack_traces should not be None in cudagraph_post_compile"
+        )
 
         prepare_cudagraph_post_compile(
             compiled_graph, example_inputs, boxed_forward_device_index
         )
 
-        from .compile_fx import cudagraphify
-
         current_callable = compiled_graph.current_callable
         assert current_callable is not None
         # Filter to only tensor constants (exclude opaque value type classes)
         tensor_constants = {
             k: v for k, v in constants.items() if isinstance(v, torch.Tensor)
         }
-        compiled_graph.current_callable = cudagraphify(
-            current_callable,
-            static_input_idxs=static_input_idxs or (),
-            device_index=next(iter(compiled_graph.device_idxs)),
+
+        device_index = next(iter(compiled_graph.device_idxs))
+        cudagraphify_kwargs = dict(
+            device_index=device_index,
             stack_traces=stack_traces,
             is_backward=is_backward,
             is_inference=is_inference,
@@ -251,6 +253,23 @@ def cudagraph_post_compile(
             mutated_input_idxs=tuple(compiled_graph.mutated_input_idxs),
         )
 
+        policy = config.cudagraph_policy
+        if policy is not None:
+            compiled_graph.current_callable = policy.cudagraphify(
+                current_callable,
+                example_inputs,
+                static_input_idxs or (),
+                **cudagraphify_kwargs,
+            )
+        else:
+            from .compile_fx import cudagraphify
+
+            compiled_graph.current_callable = cudagraphify(
+                current_callable,
+                static_input_idxs=static_input_idxs or (),
+                **cudagraphify_kwargs,
+            )
+
     else:
         BoxedBool.disable(cudagraphs)
         maybe_handle_backward_generation(compiled_graph, boxed_forward_device_index)
@@ -272,8 +291,8 @@ def cudagraph_partition_post_compile(
     example_inputs: Sequence[InputType],
     compiled_graph: CompiledFxGraph,
     cudagraphs: BoxedBool,
-    constants: dict[str, Union[torch.Tensor, type]],
-    boxed_forward_device_index: Optional[BoxedDeviceIndex],
+    constants: dict[str, torch.Tensor | type],
+    boxed_forward_device_index: BoxedDeviceIndex | None,
 ) -> None:
     """
     Cudagraphify each partition functions, which first prepares the necessary
@@ -303,8 +322,6 @@ def cudagraph_partition_post_compile(
         maybe_handle_backward_generation(compiled_graph, boxed_forward_device_index)
         return
 
-    from .compile_fx import cudagraphify
-
     assert compiled_graph.current_callable is not None
     assert compiled_graph.recursively_apply_fns is not None
     is_inference = compiled_graph.fx_kwargs["is_inference"]
@@ -318,6 +335,9 @@ def cudagraph_partition_post_compile(
         k: v for k, v in constants.items() if isinstance(v, torch.Tensor)
     }
 
+    assert compiled_graph.cudagraph_info.stack_traces is not None, (
+        "stack_traces should not be None in cudagraph_partition_post_compile"
+    )
     graph_metadata = CudagraphMetadata(
         compiled_graph.cudagraph_info.placeholders,
         static_input_idxs,
@@ -330,6 +350,8 @@ def cudagraph_partition_post_compile(
         compiled_graph, example_inputs, boxed_forward_device_index
     )
 
+    from .compile_fx import cudagraphify
+
     # cudagraphify each partition function, assuming every graph partition function
     # is cudagraphable. Non-cudagraphable ops (e.g., cpu ops) are inlined into
     # `call` function and not included in partition functions.
@@ -367,14 +389,28 @@ def maybe_realign_inputs(
     we didn't end up running cudagraphs. Mutates
     `compiled_graph.current_callable` if cudagraphs
     was run. Otherwise, does nothing.
+
+    Non-mutated inputs are handled by deferred alignment copies
+    in the generated code. Only mutated inputs need the wrapper
+    for writeback.
     """
     if not ran_cudagraphs:
-        assert compiled_graph.current_callable is not None
-        new_callable = align_inputs_from_check_idxs(
-            compiled_graph.current_callable, inputs_to_check, mutated_inputs_idxs
-        )
-        if new_callable is not compiled_graph.current_callable:
-            compiled_graph.current_callable = new_callable
+        check_idxs = inputs_to_check
+        if compiled_graph._defers_input_alignment:
+            # Non-mutated inputs are handled by deferred alignment copies
+            # in the generated Python code. Only mutated inputs need the wrapper
+            # for writeback. Backends that don't emit deferred copies (cpp_wrapper,
+            # FXIR) need the full wrapper.
+            check_idxs = [i for i in inputs_to_check if i in mutated_inputs_idxs]
+        if check_idxs:
+            assert compiled_graph.current_callable is not None
+            new_callable = align_inputs_from_check_idxs(
+                compiled_graph.current_callable,
+                check_idxs,
+                mutated_inputs_idxs,
+            )
+            if new_callable is not compiled_graph.current_callable:
+                compiled_graph.current_callable = new_callable
 
 
 class CompiledFxGraphConstants:
@@ -395,7 +431,7 @@ class CompiledFxGraphConstants:
     the value of constants directly off of the original saved object.
     """
 
-    def unwrap(self, g: CompiledFxGraph) -> dict[str, Union[torch.Tensor, type]]:
+    def unwrap(self, g: CompiledFxGraph) -> dict[str, torch.Tensor | type]:
         assert g.constants is not None
         return {**g.constants, **g.opaque_value_type_classes}
 
@@ -412,7 +448,7 @@ class CompiledFxGraphConstantsWithGm(CompiledFxGraphConstants):
     def __init__(self, gm: torch.fx.GraphModule) -> None:
         self.gm = gm
 
-    def unwrap(self, g: CompiledFxGraph) -> dict[str, Union[torch.Tensor, type]]:
+    def unwrap(self, g: CompiledFxGraph) -> dict[str, torch.Tensor | type]:
         frozen_params = {
             name: getattr(self.gm, orig_name)
             for name, orig_name in g.frozen_param_names.items()
@@ -428,26 +464,26 @@ class CompiledFxGraph(OutputCode):
     to support FxGraph caching.
     """
 
-    current_callable: Optional[Callable[..., Any]]
-    recursively_apply_fns: Optional[Callable[..., Any]]
-    compiled_fn_runner: Optional[Any]
+    current_callable: Callable[..., Any] | None
+    recursively_apply_fns: Callable[..., Any] | None
+    compiled_fn_runner: Any | None
     cache_key: str
     source_code: str = dataclasses.field(repr=False)  # Do not display source_code
     runnable_graph_str: str = dataclasses.field(repr=False)  # Do not display graph
     inductor_post_grad_graph_str: str = dataclasses.field(
         repr=False
     )  # Do not display graph
-    cache_linemap: Optional[list[tuple[int, str]]]
+    cache_linemap: list[tuple[int, str]] | None
     device_types: OrderedSet[str]
     device_idxs: OrderedSet[int]
     mutated_inputs: OrderedSet[str]
     mutated_input_idxs: OrderedSet[int]
-    constants: Optional[dict[str, torch.Tensor]]
+    constants: dict[str, torch.Tensor] | None
     frozen_param_names: dict[str, str]
     torchbind_constants: dict[str, torch._C.ScriptObject | FakeScriptObject]
     opaque_value_type_classes: dict[str, type]
-    output_strides: Optional[list[Optional[tuple[_StrideExprStr, ...]]]]
-    disabled_cudagraphs_reason: Optional[str]
+    output_strides: list[tuple[_StrideExprStr, ...] | None] | None
+    disabled_cudagraphs_reason: str | None
     metrics_deltas: metrics.CachedMetricsDeltas
     counter_deltas: Counter[str]
     # This is a string representation of an expression we serialize
@@ -455,38 +491,46 @@ class CompiledFxGraph(OutputCode):
     # context in order to verify the validity of serving a cached
     # fx graph. The expression must be generated by:
     # ShapeEnv.produce_guards_expression()
-    guards_expr: Optional[str]
-    inductor_provenance_mapping_str: Optional[str]
-    inductor_provenance_stack_traces_str: Optional[str]
-
-    cudagraph_info: Optional[CudagraphCachedInfo]
-    partition_maps: Optional[list[GraphPartitionMap]]
+    guards_expr: str | None
+    extern_libs_key: str | None
+    inductor_provenance_mapping_str: str | None
+    inductor_provenance_stack_traces_str: str | None
+
+    cudagraph_info: CudagraphCachedInfo | None
+    partition_maps: list[GraphPartitionMap] | None
+    compile_region_name: str | None
     fx_kwargs: _CompileFxKwargs
     inputs_to_check: Sequence[int]
 
-    _boxed_call: Optional[bool] = None
-    _triton_bundle: Optional[TritonBundle] = None
+    _boxed_call: bool | None = None
+    _triton_bundle: TritonBundle | None = None
     _wrap_compiled_regions: bool = False
+    _defers_input_alignment: bool = False
+    # Metadata-stripped copy of the FX graph for fake tensor propagation.
+    # Running this graph under FakeTensorMode re-derives output shapes
+    # (including aliasing) from the input shapes.
+    _original_gm: torch.fx.GraphModule | None = None
 
     def __init__(
         self,
-        current_callable: Optional[Callable[..., Any]],
+        current_callable: Callable[..., Any] | None,
         graph: GraphLowering,
         gm: torch.fx.GraphModule,
-        output_strides: list[Optional[tuple[_StrideExprStr, ...]]],
-        disabled_cudagraphs_reason: Optional[str],
+        output_strides: list[tuple[_StrideExprStr, ...] | None],
+        disabled_cudagraphs_reason: str | None,
         metrics_deltas: metrics.CachedMetricsDeltas,
         counter_deltas: Counter[str],
         cudagraphs: BoxedBool,
         example_inputs: Sequence[InputType],
         static_input_idxs: Sequence[int],
+        compile_region_name: str | None,
         fx_kwargs: _CompileFxKwargs,
         inputs_to_check: Sequence[int],
         runnable_graph_str: str,
         inductor_post_grad_graph_str: str,
-        compiled_fn_runner: Optional[Any] = None,
-        inductor_provenance_mapping_str: Optional[str] = None,
-        inductor_provenance_stack_traces_str: Optional[str] = None,
+        compiled_fn_runner: Any | None = None,
+        inductor_provenance_mapping_str: str | None = None,
+        inductor_provenance_stack_traces_str: str | None = None,
     ) -> None:
         self.current_callable = current_callable
         self.compiled_fn_runner = compiled_fn_runner
@@ -536,8 +580,10 @@ def __init__(
         self.metrics_deltas = metrics_deltas
         self.counter_deltas = counter_deltas
         self.guards_expr = None
+        self.extern_libs_key = None
         self.cudagraph_info = None
         self.partition_maps = graph.partition_maps
+        self._defers_input_alignment = getattr(graph, "_defers_input_alignment", False)
         self.fx_kwargs = {}
         self.inputs_to_check = ()
 
@@ -586,7 +632,15 @@ def __init__(
                     (not complex_memory_overlap_inputs, "complex memory overlap"),
                     (
                         all(
-                            isinstance(t, (torch.Tensor, torch.SymInt, torch.Generator))
+                            isinstance(
+                                t,
+                                (
+                                    torch.Tensor,
+                                    torch.SymInt,
+                                    torch.Generator,
+                                    OpaqueBase,
+                                ),
+                            )
                             for t in example_inputs
                         ),
                         "non-Tensor inputs",
@@ -595,7 +649,10 @@ def __init__(
                 output = output_node(gm)
                 # output args are tuple of first argument
                 assert len(output.args) == 1
-                stack_traces = [
+                # Use stack traces captured on the output node before
+                # post-grad passes, which may strip stack_trace from
+                # individual arg nodes.
+                stack_traces = output.meta.get("output_stack_traces") or [
                     (arg.stack_trace if isinstance(arg, torch.fx.node.Node) else None)
                     for arg in output.args[0]  # type: ignore[union-attr]
                 ]
@@ -606,6 +663,7 @@ def __init__(
                 )
 
         self.cudagraph_info = cudagraph_info
+        self.compile_region_name = compile_region_name
         self.inputs_to_check = inputs_to_check
         self.fx_kwargs = fx_kwargs
 
@@ -616,6 +674,17 @@ def __init__(
         # This is set at compile time to avoid runtime overhead
         self._wrap_compiled_regions = config.wrap_inductor_compiled_regions
 
+        if self._wrap_compiled_regions:
+            # Store a metadata-stripped copy of the FX graph. Running this
+            # under FakeTensorMode re-derives output shapes and aliasing
+            # from the input fake tensors.
+            import copy
+
+            gm_copy = copy.deepcopy(gm)
+            for node in gm_copy.graph.nodes:
+                node.meta.clear()
+            self._original_gm = gm_copy
+
     def __del__(self) -> None:
         if self.compiled_fn_runner is not None:
             # For torch._inductor.config.graph_partition = True,
@@ -695,6 +764,17 @@ def post_compile(
         assert graph_kwargs["is_backward"] is not None
         is_backward = graph_kwargs["is_backward"]
         cudagraphs: BoxedBool = graph_kwargs["cudagraphs"]
+
+        # When a CUDAGraphPolicy is set and it says not to wrap this
+        # inner CompiledFxGraph (e.g. because wrapping happens at the
+        # outer level via policy.wrap_output), disable cudagraphs for
+        # this graph so the rest of post_compile (input realignment,
+        # _wrap_compiled_regions) still runs normally.
+        policy = config.cudagraph_policy
+        if policy is not None and not policy.should_wrap(self):
+            counters["inductor"]["cudagraph_skips"] += 1
+            BoxedBool.disable(cudagraphs)
+
         if cudagraphs:
             # It's possible that cudagraphs is enabled, but was disabled
             # during a previous compilation we're loading from the cache.
@@ -720,9 +800,12 @@ def post_compile(
                         "boxed_forward_device_index", None
                     )
 
-                if config.graph_partition:
-                    # with graph_partition=True, we skip some cudagraph checks if it's supported
-                    # with partition. So we have to use cudagraph_partition_post_compile.
+                if config.graph_partition and policy is None:
+                    # With graph_partition=True, we skip some cudagraph checks
+                    # if it's supported with partition, so we use
+                    # cudagraph_partition_post_compile.  When a CUDAGraphPolicy
+                    # is active, we use cudagraph_post_compile instead so the
+                    # policy controls wrapping via policy.cudagraphify().
                     cudagraph_partition_post_compile(
                         example_inputs,
                         self,
@@ -751,11 +834,24 @@ def post_compile(
         # Apply inductor_compiled_code HOP wrapper if configured
         # This is done in post_compile to ensure it works with cached artifacts
         if self._wrap_compiled_regions and self.current_callable is not None:
+            from torch._higher_order_ops.wrap import InductorCompiledCallable
+
             original_callable = self.current_callable
 
+            inductor_callable = InductorCompiledCallable(
+                original_callable,
+                self._original_gm,
+                compile_region_name=self.compile_region_name,
+            )
+
             def wrapped_callable(inputs):
                 if is_in_torch_dispatch_mode():
-                    return inductor_compiled_code(original_callable, inputs)
+                    kwargs = (
+                        {"name": self.compile_region_name}
+                        if self.compile_region_name is not None
+                        else {}
+                    )
+                    return inductor_compiled_code(inductor_callable, inputs, **kwargs)
                 else:
                     return original_callable(inputs)
 
@@ -772,6 +868,9 @@ def prepare_for_serialization(self) -> None:
         self.current_callable = None
         self.recursively_apply_fns = None
         self.compiled_fn_runner = None
+        # Note: _original_gm is already picklable (metadata stripped at creation)
+        # Note: _serialized_fx_graph is already in serializable form (SerializedGraphModule)
+        # so it doesn't need to be cleared
 
     def write_to_disk(self) -> str:
         from torch._dynamo.utils import counters
@@ -821,9 +920,9 @@ class CompiledAOTI(OutputCode):
     Class holding an AOTInductor compiled so.
     """
 
-    filename: Union[str, list[Union[str, Weights]], torch.fx.GraphModule]
+    filename: str | list[str | Weights] | torch.fx.GraphModule
     device_type: str
-    current_callable: Optional[Callable[..., Any]] = None
+    current_callable: Callable[..., Any] | None = None
     _cached_files: dict[str, bytes] = dataclasses.field(default_factory=dict)
 
     def __post_init__(self):
@@ -867,6 +966,15 @@ def __post_init__(self):
                     True,
                 ).run  # type: ignore[attr-defined]
             )  # type: ignore[attr-defined]
+        elif self.device_type.startswith("xpu"):
+            current_callable = (
+                torch._C._aoti.AOTIModelContainerRunnerXpu(  # type: ignore[call-arg]
+                    current_callable,
+                    1,
+                    self.device_type,
+                    "",
+                ).run  # type: ignore[attr-defined]
+            )  # type: ignore[attr-defined]
         elif self.device_type == "cpu":
             current_callable = (
                 torch._C._aoti.AOTIModelContainerRunnerCpu(  # type: ignore[call-arg]
@@ -954,10 +1062,10 @@ class RegionalOutputCode(OutputCode):
     """
 
     # The serialized graph module as bytes (using GraphPickler)
-    _serialized_graph_module: Optional[bytes] = None
+    _serialized_graph_module: bytes | None = None
 
     # The actual graph module (cleared during serialization)
-    _graph_module: Optional[torch.nn.Module] = None
+    _graph_module: torch.nn.Module | None = None
 
     # Optional filter for ops during serialization
     _ops_filter: Callable[[str], bool] | None = None
@@ -966,7 +1074,7 @@ def __init__(
         self,
         graph_module: torch.fx.GraphModule,
         ops_filter: Callable[[str], bool] = _ops_filter_safe,
-        node_metadata_key_filter: Optional[Callable[[str], bool]] = (
+        node_metadata_key_filter: Callable[[str], bool] | None = (
             _node_metadata_key_filter_safe
         ),
     ):
diff --git a/torch/_inductor/pattern_matcher.py b/torch/_inductor/pattern_matcher.py
index 40248fd7997a2..5b363cc1e50a2 100644
--- a/torch/_inductor/pattern_matcher.py
+++ b/torch/_inductor/pattern_matcher.py
@@ -51,7 +51,7 @@
 from collections import defaultdict
 from collections.abc import Callable, Collection, Generator, Iterable, Mapping, Sequence
 from pathlib import Path
-from typing import Any, NoReturn, Optional, Protocol, TypeVar, Union
+from typing import Any, NoReturn, Protocol, TypeVar
 from typing_extensions import Self, TypeIs
 
 import torch
@@ -85,11 +85,41 @@
 prims = torch.ops.prims
 
 Constant = Any
-NodeOrConstant = Union[Constant, torch.fx.Node]
+NodeOrConstant = Constant | torch.fx.Node
 
 backend = os.environ.get("TORCHINDUCTOR_PATTERN_MATCH_BACKEND", "inductor")
 
 
+_debug_nodes_cache: bool | OrderedSet[str] | None = None
+_debug_nodes_env_value_cache: str | None = None
+
+
+def _should_debug_node(node_name: str) -> bool:
+    def _get_debug_nodes() -> bool | OrderedSet[str]:
+        global _debug_nodes_cache, _debug_nodes_env_value_cache
+
+        def parse_debug_env(env_value: str | None) -> bool | OrderedSet[str]:
+            if not env_value:
+                return False
+            if env_value == "all":
+                return True
+            return OrderedSet(env_value.split(","))
+
+        current_env = os.environ.get("TORCHINDUCTOR_PATTERN_MATCH_DEBUG")
+
+        # Recompute only if env changed
+        if current_env != _debug_nodes_env_value_cache or _debug_nodes_cache is None:
+            _debug_nodes_cache = parse_debug_env(current_env)
+            _debug_nodes_env_value_cache = current_env
+
+        return _debug_nodes_cache
+
+    debug_nodes = _get_debug_nodes()
+    if isinstance(debug_nodes, bool):
+        return debug_nodes
+    return node_name in debug_nodes
+
+
 class SearchFn(Protocol):
     __name__: str
 
@@ -102,14 +132,14 @@ def __call__(self, *args: Any, **kwargs: Any) -> Any: ...
 
 class TraceFn(Protocol):
     def __call__(
-        self, fn: Union[SearchFn, ReplaceFn], *args: Any, **kwargs: Any
+        self, fn: SearchFn | ReplaceFn, *args: Any, **kwargs: Any
     ) -> torch.fx.GraphModule: ...
 
 
 T = TypeVar("T")
 
 # What's a better name for this?
-FnsType = Union[torch.fx.node.Target, str]
+FnsType = torch.fx.node.Target | str
 
 
 class Multiple:
@@ -127,11 +157,9 @@ def _transfer_meta(
 ) -> None:
     from torch.fx.traceback import NodeSource, NodeSourceAction
 
-    # transfer metadata after pattern matching occurs.
-    # skip "val" and "tensor_meta" because this info is too specific; it's unlikely
-    # to remain accurate after pattern matching has occurred.
+    # Transfer metadata after pattern matching occurs.
+    # Copies _COPY_META_FIELDS, stack_trace, and (if missing) val/tensor_meta.
     if config.trace.provenance_tracking_level == 1:
-        # We handle "from_node" field of the node meta specially to record that the new node comes from the old_node.
         new_from_node = new_meta.get("from_node", []).copy()
         new_from_node.append(NodeSource(old_node, pass_name, NodeSourceAction.REPLACE))
         new_meta.update(
@@ -148,6 +176,13 @@ def _transfer_meta(
         )
     if "stack_trace" in old_node.meta:
         new_meta["stack_trace"] = old_node.meta["stack_trace"]
+    # Copy val/tensor_meta only when the new node doesn't already have them
+    # (e.g. from tracing the replacement graph). Don't overwrite if present
+    # since the replacement's own val is more accurate.
+    if "val" not in new_meta and "val" in old_node.meta:
+        new_meta["val"] = old_node.meta["val"]
+    if "tensor_meta" not in new_meta and "tensor_meta" in old_node.meta:
+        new_meta["tensor_meta"] = old_node.meta["tensor_meta"]
 
 
 class Match:
@@ -168,14 +203,14 @@ class Match:
     nodes: list[torch.fx.Node]
     targets: dict[_TargetExpr, torch.fx.node.Target]
     ctx: MatchContext
-    replacement_graph: Optional[torch.fx.GraphModule]
+    replacement_graph: torch.fx.GraphModule | None
 
     def __init__(
         self,
         ctx: MatchContext,
         pattern: PatternExpr,
-        args: Optional[Sequence[Any]] = None,
-        kwargs: Optional[dict[str, Any]] = None,
+        args: Sequence[Any] | None = None,
+        kwargs: dict[str, Any] | None = None,
     ) -> None:
         super().__init__()
         self.pattern = pattern
@@ -217,7 +252,7 @@ def erase_nodes(self) -> None:
             if not n._erased and not n.users:
                 graph.erase_node(n)
 
-    def output_nodes(self) -> list[Optional[torch.fx.Node]]:
+    def output_nodes(self) -> list[torch.fx.Node | None]:
         return [
             (self.ctx.pattern_to_node[p] if p is not None else None)
             for p in self.ctx.outputs
@@ -237,7 +272,7 @@ def replace_by_example(
         self,
         replacement_fn: ReplaceFn,
         args: Sequence[Any],
-        trace_fn: Optional[TraceFn] = None,
+        trace_fn: TraceFn | None = None,
         run_functional_passes: bool = True,
     ) -> None:
         """Replace with a graph generated by tracing the replacement_fn.
@@ -331,6 +366,17 @@ def record(node: torch.fx.Node, val: Any) -> None:
                     if "val" in arg.meta
                     else arg.meta["example_value"],
                 )
+                fake_mode = torch._dynamo.utils.detect_fake_mode(example_vals)
+                if fake_mode is not None:
+
+                    def _convert_to_fake_mode(it):
+                        if isinstance(it, FakeTensor) and fake_mode is not None:
+                            return fake_mode.from_tensor(it)
+                        return it
+
+                    example_vals = torch.fx.node.map_aggregate(
+                        example_vals, _convert_to_fake_mode
+                    )
                 replacement = trace_fn(replacement_fn, example_vals)
             if len(self.nodes) == 1:
                 for n in replacement.graph.nodes:
@@ -376,7 +422,7 @@ def __bool__(self) -> bool:
         return False
 
 
-MatchResult = Union[Match, FailedMatch]
+MatchResult = Match | FailedMatch
 
 
 def is_match(m: MatchResult) -> TypeIs[Match]:
@@ -392,15 +438,15 @@ class MatchContext:
     Internal state needed while running PatternExpr._match().
     """
 
-    outputs: list[Optional[PatternExpr]]
-    pattern_to_node: dict[PatternExpr, Optional[torch.fx.Node]]
+    outputs: list[PatternExpr | None]
+    pattern_to_node: dict[PatternExpr, torch.fx.Node | None]
     graph: torch.fx.Graph
     exclusive_node_set: list[NodeOrConstant]
 
     def __init__(
         self,
-        outputs: list[Optional[PatternExpr]],
-        pattern_to_node: Optional[dict[PatternExpr, torch.fx.Node]] = None,
+        outputs: list[PatternExpr | None],
+        pattern_to_node: dict[PatternExpr, torch.fx.Node] | None = None,
         *,
         graph: torch.fx.Graph,
     ) -> None:
@@ -451,7 +497,7 @@ def __repr__(self) -> str:
 
     def find_anchor_nodes(
         self, ctx: MatchContext, searched: OrderedSet[torch.fx.Node]
-    ) -> Generator[Optional[torch.fx.Node], None, None]:
+    ) -> Generator[torch.fx.Node | None, None, None]:
         if self in ctx.pattern_to_node:
             yield ctx.pattern_to_node[self]
 
@@ -544,7 +590,7 @@ class _TargetExpr(PatternExpr):
     fns_set: OrderedSet[FnsType]
 
     def __init__(
-        self, fns: Union[FnsType, Sequence[FnsType]], users: Union[Multiple, int] = 1
+        self, fns: FnsType | Sequence[FnsType], users: Multiple | int = 1
     ) -> None:
         super().__init__()
         fns = [fns] if callable(fns) or isinstance(fns, str) else list(fns)
@@ -590,7 +636,7 @@ def has_multiple_users(self) -> bool:
 
     def find_anchor_nodes(
         self, ctx: MatchContext, searched: OrderedSet[torch.fx.Node]
-    ) -> Generator[Optional[torch.fx.Node], None, None]:
+    ) -> Generator[torch.fx.Node | None, None, None]:
         raise NotImplementedError
 
     def _match_fns(self, node: torch.fx.Node) -> bool:
@@ -627,9 +673,9 @@ class _TargetArgsExpr(_TargetExpr):
 
     def __init__(
         self,
-        fns: Union[torch.fx.node.Target, str, Sequence[Any]],
+        fns: torch.fx.node.Target | str | Sequence[Any],
         *args: Any,
-        _users: Union[int, Multiple] = 1,
+        _users: int | Multiple = 1,
         **kwargs: Any,
     ) -> None:
         super().__init__(fns, _users)
@@ -647,7 +693,7 @@ def __init__(
     @staticmethod
     def simple_flatten(
         args: Sequence[Any], kwargs: Mapping[Any, Any]
-    ) -> tuple[Sequence[Any], Union[_SimpleSpec, pytree.TreeSpec]]:
+    ) -> tuple[Sequence[Any], _SimpleSpec | pytree.TreeSpec]:
         values = (*args, *kwargs.values())
         spec = (len(args), *kwargs.keys())
         return values, spec
@@ -655,7 +701,7 @@ def simple_flatten(
     @staticmethod
     def pytree_flatten(
         args: Sequence[Any], kwargs: Mapping[Any, Any]
-    ) -> tuple[Sequence[Any], Union[_SimpleSpec, pytree.TreeSpec]]:
+    ) -> tuple[Sequence[Any], _SimpleSpec | pytree.TreeSpec]:
         type_mapping: dict[type, type] = {
             immutable_list: tuple,
             list: tuple,
@@ -763,7 +809,7 @@ def _match(self, node: torch.fx.Node, ctx: MatchContext) -> MatchResult:
 
     def find_anchor_nodes(
         self, ctx: MatchContext, searched: OrderedSet[torch.fx.Node]
-    ) -> Generator[Optional[torch.fx.Node], None, None]:
+    ) -> Generator[torch.fx.Node | None, None, None]:
         """
         This is used when we are matching a pattern with multiple outputs.
         There is a partial match (stored in ctx) and we want to walk
@@ -902,9 +948,9 @@ def pattern_eq(self, other: Any) -> bool:
 
 
 class MultiOutputPattern(PatternExpr):
-    outputs: list[Optional[PatternExpr]]
+    outputs: list[PatternExpr | None]
 
-    def __init__(self, outputs: Sequence[Optional[PatternExpr]]) -> None:
+    def __init__(self, outputs: Sequence[PatternExpr | None]) -> None:
         super().__init__()
         assert isinstance(outputs[0], _TargetExpr)
         assert all(x is None or isinstance(x, PatternExpr) for x in outputs), outputs
@@ -912,7 +958,7 @@ def __init__(self, outputs: Sequence[Optional[PatternExpr]]) -> None:
         self.op = outputs[0].op
 
     @property
-    def fns(self) -> Union[Callable[..., Any], str, Sequence[Any]]:
+    def fns(self) -> Callable[..., Any] | str | Sequence[Any]:
         # This cast is checked above in __init__()
         output = typing.cast(_TargetExpr, self.outputs[0])
         return output.fns
@@ -1085,8 +1131,8 @@ def apply(self, match: Match, graph: torch.fx.Graph, node: torch.fx.Node) -> Non
 
     def register(
         self,
-        pass_dicts: Union[_PassDictsType, Sequence[_PassDictsType]],
-        target: Union[torch.fx.node.Target, None] = None,
+        pass_dicts: _PassDictsType | Sequence[_PassDictsType],
+        target: torch.fx.node.Target | None = None,
         prepend: bool = False,
     ) -> None:
         if target is None:
@@ -1113,7 +1159,7 @@ def apply(self, match: Match, graph: torch.fx.Graph, node: torch.fx.Node) -> Non
         handler = functools.wraps(self.handler)(functools.partial(self.handler, match))
         with graph.inserting_before(node):
             replacement = graph.call_function(handler, tuple(match.args), match.kwargs)
-            replacement.meta.update(node.meta)
+            _transfer_meta(replacement.meta, node)
             node.replace_all_uses_with(replacement)
         assert match.nodes[-1] is node
         match.erase_nodes()
@@ -1139,13 +1185,15 @@ class ReplacementPatternEntry(PatternEntry):
     """
 
     normalize_args: Callable[..., list[Any]]
+    pattern_name: str | None = None  # Unique identifier for per-pattern telemetry
 
     @staticmethod
     def replace_with_graph(
         match: Match,
         graph: torch.fx.Graph,
-        replacement_graph: Union[torch.fx.Graph, torch.fx.GraphModule],
+        replacement_graph: torch.fx.Graph | torch.fx.GraphModule,
         args: Sequence[torch.fx.Node],
+        pass_name: str | None = None,
     ) -> None:
         """
         Inserts the replacement graph into the toplevel graph at the match
@@ -1170,18 +1218,14 @@ def run_node(self, node: torch.fx.Node) -> Any:
                     _transfer_meta(
                         new_meta=result.meta,
                         old_node=node,
-                        pass_name="Interpreter_Replacer",
+                        pass_name=pass_name or "",
                     )
                     # This function copy-pastes the replacement graph into
                     # the graph. If the replacement graph had any eager_input_vals,
-                    # or val/tensor_meta, we propagate those over.
+                    # we propagate those over (val/tensor_meta are handled by
+                    # _transfer_meta above).
                     if "eager_input_vals" in node.meta:
                         result.meta["eager_input_vals"] = node.meta["eager_input_vals"]
-                    if "val" in node.meta and "val" not in result.meta:
-                        result.meta["val"] = node.meta["val"]
-                        if isinstance(node.meta["val"], torch.Tensor):
-                            assert "tensor_meta" in node.meta
-                            result.meta["tensor_meta"] = node.meta["tensor_meta"]
                     return result
                 if node.op == "get_attr":
                     # If the replacement graph contains a HOP, the subgraphs of the HOP are "get_attr" nodes.
@@ -1266,8 +1310,8 @@ def maybe_getitem(node: torch.fx.Node) -> Any:
                 return node.args[1]
 
             def replace(
-                old: Union[torch.fx.Node, None],
-                new: Union[torch.fx.Node, Sequence[torch.fx.Node], None],
+                old: torch.fx.Node | None,
+                new: torch.fx.Node | Sequence[torch.fx.Node] | None,
             ) -> None:
                 def filter_nodes_in_newly_added_nodes(node: torch.fx.Node) -> bool:
                     # Do not replace the use of a node if it is being used by
@@ -1287,8 +1331,7 @@ def filter_nodes_in_newly_added_nodes(node: torch.fx.Node) -> bool:
                         graph.erase_node(old)
                     return
                 if isinstance(new, torch.fx.Node):
-                    if "val" not in new.meta:
-                        new.meta.update(old.meta)
+                    _transfer_meta(new.meta, old, pass_name=pass_name or "")
 
                     # Preserve the recompute tags in the replacement graph. We
                     # look at the recompute tags of the original output node to
@@ -1358,6 +1401,16 @@ def filter_nodes_in_newly_added_nodes(node: torch.fx.Node) -> bool:
 
         match.erase_nodes()
 
+        # Remove dead replacement nodes so they don't inflate user counts
+        # in later lowering heuristics (e.g. should_realize_on_reuse).
+        for node in reversed(added_replacement_nodes):
+            if (
+                not node.users
+                and not node.is_impure()
+                and not isinstance(node.target, torch._ops.HigherOrderOperator)
+            ):
+                graph.erase_node(node)
+
     def apply(self, match: Match, graph: torch.fx.Graph, node: torch.fx.Node) -> None:
         assert match.replacement_graph is not None
         self.replace_with_graph(
@@ -1365,6 +1418,7 @@ def apply(self, match: Match, graph: torch.fx.Graph, node: torch.fx.Node) -> Non
             graph,
             match.replacement_graph,
             self.normalize_args(*match.args, **match.kwargs),
+            pass_name=self.pattern_name or "replace_with_graph",
         )
 
 
@@ -1382,8 +1436,8 @@ def log_trace_failure(search_fn: Callable[..., Any], e: RuntimeError) -> None:
 
 def check_and_add_duplicate_pattern(
     pattern: PatternExpr,
-    graph: Optional[torch.fx.Graph],
-    seen_patterns: dict[str, list[Optional[str]]],
+    graph: torch.fx.Graph | None,
+    seen_patterns: dict[str, list[str | None]],
     skip_duplicates: bool = False,
 ) -> bool:
     """
@@ -1425,14 +1479,16 @@ def check_and_add_duplicate_pattern(
 def register_replacement(
     search_fn: SearchFn,
     replace_fn: ReplaceFn,
-    example_inputs: Iterable[Any],
+    example_inputs: list[Any] | tuple[Any, ...],
     trace_fn: TraceFn,
-    pass_dicts: Union[_PassDictsType, Sequence[_PassDictsType]],
+    pass_dicts: _PassDictsType | Sequence[_PassDictsType],
     extra_check: Callable[[Match], bool] = _return_true,
-    scalar_workaround: Union[dict[str, Union[float, int]], None] = None,
+    scalar_workaround: dict[str, float | int] | None = None,
     exclusive_arg_names: Sequence[str] = (),
-    search_fn_pattern: Union[PatternExpr, None] = None,
+    search_fn_pattern: PatternExpr | None = None,
     skip_duplicates: bool = False,
+    pattern_name: str | None = None,
+    get_decomp_fn: Callable[..., dict[Any, Callable[..., Any]]] = select_decomp_table,
 ) -> bool:
     """
     Create a replacement rule based on example functions that get traced
@@ -1456,6 +1512,11 @@ def register_replacement(
         replace_argnames = [*inspect.signature(replace_fn).parameters.keys()]
         replace_fn = _wrap_bound_method(replace_fn, replace_argnames)
 
+    if not isinstance(example_inputs, (list, tuple)):
+        raise TypeError(
+            f"example_inputs must be a list or tuple, got {type(example_inputs)}"
+        )
+
     def check_fn(match: Match) -> bool:
         """
         Often shapes get burned into the pattern, so our initial match ran with
@@ -1519,12 +1580,15 @@ def check_fn(match: Match) -> bool:
                     # Later, when we actually do the replacement, the symbolic shape
                     # sizes will get re-traced and added to the graph.
 
-                    def search_fn_new(*args_new: Any) -> Any:
+                    def search_fn_new(*args_new: Any, **_: Any) -> Any:
                         return search_fn(*args_new[len(args_new) - len(args) :])
 
                     try:
-                        # pyrefly: ignore [bad-argument-type]
-                        specific_graph = trace_fn(search_fn_new, sym_args + args)
+                        specific_graph = trace_fn(
+                            search_fn_new,
+                            sym_args + args,
+                            get_decomp_fn=get_decomp_fn,
+                        )
                     except RuntimeError as e:
                         log_trace_failure(search_fn, e)
                         return False
@@ -1550,7 +1614,9 @@ def search_fn_new(*args_new: Any) -> Any:
                     argnames = sym_arg_names + argnames
                 else:
                     try:
-                        specific_graph = trace_fn(search_fn, args)
+                        specific_graph = trace_fn(
+                            search_fn, args, get_decomp_fn=get_decomp_fn
+                        )
                     except RuntimeError as e:
                         log_trace_failure(search_fn, e)
                         return False
@@ -1566,7 +1632,7 @@ def search_fn_new(*args_new: Any) -> Any:
             assert node is not None
             specific_pattern_match = specific_pattern.match(node)
 
-            if os.environ.get("TORCHINDUCTOR_PATTERN_MATCH_DEBUG") == node.name:
+            if _should_debug_node(node.name):
                 log.warning(
                     "Specific pattern match: %s%s %s %s",
                     node,
@@ -1577,7 +1643,9 @@ def search_fn_new(*args_new: Any) -> Any:
 
             if is_match(specific_pattern_match) and extra_check(specific_pattern_match):
                 # trace the pattern using the shapes from the user program
-                match.replacement_graph = trace_fn(replace_fn, args)
+                match.replacement_graph = trace_fn(
+                    replace_fn, args, get_decomp_fn=get_decomp_fn
+                )
                 if len(match.nodes) == 1:
                     for n in match.replacement_graph.graph.nodes:
                         _transfer_meta(
@@ -1615,6 +1683,7 @@ def normalize_args(**kwargs: Any) -> list[Any]:
                 trace_fn,
                 scalar_workaround,
                 exclusive_arg_names,
+                get_decomp_fn=get_decomp_fn,
             )
         else:
             pattern = search_fn_pattern
@@ -1636,6 +1705,7 @@ def normalize_args(**kwargs: Any) -> list[Any]:
             pattern=pattern,
             extra_check=check_fn,
             normalize_args=normalize_args,
+            pattern_name=pattern_name,
         )
         pattern.register(pass_dicts)
         return pattern.pattern  # type: ignore[return-value]
@@ -1649,7 +1719,7 @@ def _serialize_pattern(
     search_fn: SearchFn,
     example_inputs: Sequence[Any],
     trace_fn: TraceFn,
-    scalar_workaround: Union[dict[str, Union[float, int]], None],
+    scalar_workaround: dict[str, float | int] | None,
 ) -> PatternExpr:
     def get_file_template() -> str:
         auto_generated_msg = textwrap.dedent(
@@ -1746,11 +1816,12 @@ def gen_register_replacement(
     replace_fn: ReplaceFn,
     example_inputs: Iterable[Any],
     trace_fn: TraceFn,
-    pass_dicts: Union[_PassDictsType, Sequence[_PassDictsType]],
+    pass_dicts: _PassDictsType | Sequence[_PassDictsType],
     extra_check: Callable[[Match], bool] = _return_true,
-    scalar_workaround: Union[dict[str, Union[float, int]], None] = None,
+    scalar_workaround: dict[str, float | int] | None = None,
     exclusive_arg_names: Sequence[str] = (),
     skip_duplicates: bool = False,
+    get_decomp_fn: Callable[..., dict[Any, Callable[..., Any]]] = select_decomp_table,
 ) -> None:
     # Make sure the example_inputs is materialized.
     example_inputs = tuple(example_inputs)
@@ -1793,6 +1864,8 @@ def gen_register_replacement(
         exclusive_arg_names,
         search_fn_pattern=pat,
         skip_duplicates=skip_duplicates,
+        pattern_name=unique_name,
+        get_decomp_fn=get_decomp_fn,
     )
 
 
@@ -1801,8 +1874,9 @@ def gen_pattern_and_search_gm(
     search_fn: SearchFn,
     example_inputs: Sequence[Any],
     trace_fn: TraceFn,
-    scalar_workaround: Union[dict[str, Union[float, int]], None] = None,
+    scalar_workaround: dict[str, float | int] | None = None,
     exclusive_arg_names: Sequence[str] = (),
+    get_decomp_fn: Callable[..., dict[Any, Callable[..., Any]]] = select_decomp_table,
 ) -> tuple[PatternExpr, torch.fx.GraphModule]:
     argnames = [*inspect.signature(search_fn).parameters.keys()]
 
@@ -1818,7 +1892,7 @@ def gen_pattern_and_search_gm(
             flat_inputs.append(example_inputs[input_idx])
             input_idx += 1
 
-    search_gm = trace_fn(search_fn, flat_inputs)
+    search_gm = trace_fn(search_fn, flat_inputs, get_decomp_fn=get_decomp_fn)
     return (
         fx_to_pattern(
             search_gm,
@@ -1835,7 +1909,7 @@ def gen_pattern(
     search_fn: SearchFn,
     example_inputs: Sequence[Any],
     trace_fn: TraceFn,
-    scalar_workaround: Union[dict[str, Union[float, int]], None] = None,
+    scalar_workaround: dict[str, float | int] | None = None,
     exclusive_arg_names: Sequence[str] = (),
 ) -> PatternExpr:
     return gen_pattern_and_search_gm(
@@ -1979,10 +2053,14 @@ def wrapper(*args: Any, **kwargs: Any) -> Any:
 
 
 class PatternMatcherPass:
+    """
+    Registry of patterns to match and replace in FX graphs.
+    """
+
     def __init__(
         self,
-        pass_name: Optional[str] = None,
-        subsystem: Optional[str] = None,
+        pass_name: str | None = None,
+        subsystem: str | None = None,
     ) -> None:
         super().__init__()
         self.patterns: defaultdict[
@@ -1995,12 +2073,13 @@ def __init__(
         # of the graph used to generate them. Because we ignore certain patterns
         # in searching, but not in matching, use the graph to distinguish if two equivalent
         # searches are actually different.
-        self.seen_patterns: dict[str, list[Optional[str]]] = defaultdict(list)
+        self.seen_patterns: dict[str, list[str | None]] = defaultdict(list)
 
     def __getitem__(self, item: tuple[str, torch.fx.node.Target]) -> list[PatternEntry]:
         return self.patterns[item]
 
-    def apply(self, gm: Union[torch.fx.GraphModule, torch.fx.Graph]) -> int:
+    def apply(self, gm: torch.fx.GraphModule | torch.fx.Graph) -> int:
+        """Apply all registered patterns to the graph, returning the number of matches."""
         if not self.patterns:
             return 0
         if isinstance(gm, torch.fx.GraphModule):
@@ -2055,7 +2134,19 @@ def apply(self, gm: Union[torch.fx.GraphModule, torch.fx.Graph]) -> int:
                         != 1
                     ):
                         continue
-                    if os.environ.get("TORCHINDUCTOR_PATTERN_MATCH_DEBUG") == node.name:
+                    # pattern match crosses stream boundary - discard
+                    if (
+                        is_match(m)
+                        and len(
+                            OrderedSet(
+                                n.meta.get("custom", {}).get("stream", 0)
+                                for n in m.nodes
+                            )
+                        )
+                        != 1
+                    ):
+                        continue
+                    if _should_debug_node(node.name):
                         log.warning("%s%s %s %s", node, node.args, m, entry.pattern)
 
                     if is_match(m) and guard_or_false(entry.extra_check(m)):
@@ -2063,6 +2154,19 @@ def apply(self, gm: Union[torch.fx.GraphModule, torch.fx.Graph]) -> int:
                         entry.apply(m, graph, node)
                         counters[backend]["pattern_matcher_count"] += 1
                         counters[backend]["pattern_matcher_nodes"] += len(m.nodes)
+
+                        # Track per-pattern counts when debug mode is active
+                        if os.environ.get("TORCHINDUCTOR_PATTERN_MATCH_DEBUG"):
+                            if getattr(entry, "pattern_name", None):
+                                pattern_name = entry.pattern_name
+                            else:
+                                # Fallback: use pattern class name + operation target
+                                pattern_class = entry.pattern.__class__.__name__
+                                target = str(node.target)
+                                pattern_name = f"{pattern_class}_{target}"
+
+                            pattern_key = f"{backend}_pattern_matcher_per_pattern"
+                            counters[pattern_key][pattern_name] += 1
         return count
 
     def clear(self) -> None:
@@ -2074,10 +2178,10 @@ def _not_implemented(*args: Any, **kwargs: Any) -> NoReturn:
 
 
 def fx_to_pattern(
-    gm: Union[torch.fx.GraphModule, torch.fx.Graph],
+    gm: torch.fx.GraphModule | torch.fx.Graph,
     ignore_types: Sequence[type[Any]] = (),
     argnames: Sequence[str] = (),
-    scalar_workaround: Union[dict[str, Union[float, int]], None] = None,
+    scalar_workaround: dict[str, float | int] | None = None,
     exclusive_arg_names: Sequence[str] = (),
 ) -> PatternExpr:
     """
@@ -2091,8 +2195,8 @@ def fx_to_pattern(
     assert len(inv_scalar_workaround) == len(scalar_workaround)
 
     def process_arg(
-        x: T, ignore_types_override: Optional[Sequence[type[Any]]] = None
-    ) -> Union[T, KeywordArg, Ignored]:
+        x: T, ignore_types_override: Sequence[type[Any]] | None = None
+    ) -> T | KeywordArg | Ignored:
         current_ignore_types = (
             ignore_types_override if ignore_types_override is not None else ignore_types
         )
@@ -2120,7 +2224,7 @@ def placeholder(
             target: str,  # type: ignore[override]
             args: Sequence[Any],
             kwargs: Mapping[str, Any],
-        ) -> Union[ExclusiveKeywordArg, KeywordArg]:
+        ) -> ExclusiveKeywordArg | KeywordArg:
             n = next(argnum)
             if n < len(argnames):
                 name = argnames[n]
@@ -2148,10 +2252,10 @@ def call_function(
 
                 def process_arg_fn_impl(
                     x: T,
-                    ignore_types_override: Optional[Sequence[type[Any]]] = tuple(
+                    ignore_types_override: Sequence[type[Any]] | None = tuple(
                         t for t in ignore_types if t is not int
                     ),
-                ) -> Union[T, KeywordArg, Ignored]:
+                ) -> T | KeywordArg | Ignored:
                     return process_arg(x, ignore_types_override)
 
                 process_arg_fn = process_arg_fn_impl
@@ -2189,15 +2293,12 @@ def fwd_only(
     args: Sequence[Any],
     *,
     run_functional_passes: bool = True,
-    get_decomp_fn: Optional[Callable[..., Any]] = None,
+    get_decomp_fn: Callable[..., Any] = select_decomp_table,
 ) -> torch.fx.GraphModule:
     """Build a normalized inference graph, for use with fx_to_pattern"""
     # TODO - look into using aot autograd, asserting no mutating ops here
     with enable_python_dispatcher(), preserve_node_meta():
-        decompositions = (
-            get_decomp_fn() if get_decomp_fn is not None else select_decomp_table()
-        )
-        gm = make_fx(fn, decompositions, tracing_mode="real")(*args)
+        gm = make_fx(fn, get_decomp_fn(), tracing_mode="real")(*args)
 
     from .fx_passes.post_grad import remove_noop_ops
 
@@ -2218,9 +2319,14 @@ def fwd_only(
 
 
 @torch.enable_grad()
-def joint_fwd_bwd(fn: Callable[..., Any], args: Sequence[Any]) -> torch.fx.GraphModule:
+def joint_fwd_bwd(
+    fn: Callable[..., Any],
+    args: Sequence[Any],
+    *,
+    get_decomp_fn: Callable[..., Any] = select_decomp_table,
+) -> torch.fx.GraphModule:
     """Build a normalized training graph, for use with fx_to_pattern"""
-    gm: Optional[torch.fx.GraphModule] = None
+    gm: torch.fx.GraphModule | None = None
 
     def record_joint_graph(
         joint_graph: torch.fx.GraphModule, inputs: Sequence[Any], **kwargs: Any
@@ -2236,7 +2342,7 @@ def record_joint_graph(
             # pyrefly: ignore[bad-argument-type]
             lambda gm, example_inputs: make_boxed_func(gm),
             partition_fn=record_joint_graph,
-            decompositions=select_decomp_table(),
+            decompositions=get_decomp_fn(),
             keep_inference_input_mutations=True,
             enable_log=False,
         )(*args)
@@ -2299,16 +2405,29 @@ def stable_topological_sort(graph: torch.fx.Graph) -> None:
     assert not waiting and len(ready) == len(graph.nodes)
 
 
-def init_once_fakemode(fn: Callable[..., Any]) -> Callable[[], Any]:
+def init_once_fakemode(fn: Callable[..., Any]) -> Callable[..., Any]:
     """Wrapper around lazy init functions in fx_passes/"""
 
+    _fn_params = inspect.signature(fn).parameters
+
     @functools.cache
     @functools.wraps(fn)
-    def lazy_init(input_device: Optional[torch.device] = None) -> Any:
+    def lazy_init(
+        input_device: Any | None = None,
+        get_decomp_fn: Callable[
+            ..., dict[Any, Callable[..., Any]]
+        ] = select_decomp_table,
+    ) -> Any:
         counters_ref = counters[backend].copy()
 
+        kwargs: dict[str, Any] = {}
+        if "input_device" in _fn_params:
+            kwargs["input_device"] = input_device
+        if "get_decomp_fn" in _fn_params:
+            kwargs["get_decomp_fn"] = get_decomp_fn
+
         with torch._guards.tracing(None), unset_fake_temporarily(), FakeTensorMode():
-            result = fn(input_device)
+            result = fn(**kwargs)
 
         # clear view matches encountered during tracing
         counters[backend] = counters_ref
@@ -2347,7 +2466,7 @@ def run_node(self, old_node: torch.fx.Node) -> torch.fx.Node:
 
 
 def get_arg_value(
-    node: torch.fx.Node, arg_number: int, kwarg_name: Optional[str] = None
+    node: torch.fx.Node, arg_number: int, kwarg_name: str | None = None
 ) -> Any:
     if len(node.args) > arg_number:
         return node.args[arg_number]
diff --git a/torch/_inductor/quantized_lowerings.py b/torch/_inductor/quantized_lowerings.py
index 5b6f8c12309b8..3a62f5690c7c2 100644
--- a/torch/_inductor/quantized_lowerings.py
+++ b/torch/_inductor/quantized_lowerings.py
@@ -90,9 +90,10 @@ def _mul_epilogue(buf: torch.Tensor) -> Any:
                 epilogue_creator=_mul_epilogue,  # type: ignore[arg-type]
             )
 
-        return autotune_select_algorithm(
+        node, _ = autotune_select_algorithm(
             "_weight_int8pack_mm", choices, [mat1, mat2, scale], aten_layout
         )
+        return node
 
     @register_lowering(aten._weight_int4pack_mm_for_cpu, type_promotion_kind=None)  # type: ignore[misc]
     def int4pack_mm_cpu(
@@ -157,13 +158,14 @@ def get_example_weight(x: torch._inductor.ir.IRNode) -> torch.Tensor:
             2: lambda x: V.graph.constants[x.get_name()],  # group size
         }
 
-        return autotune_select_algorithm(
+        node, _ = autotune_select_algorithm(
             "_weight_int4pack_mm_for_cpu",
             choices,
             [mat1, mat2, group_size, qScaleAndZeros],
             aten_layout,
             input_gen_fns=input_gen_fns,
         )
+        return node
 
     lowering.make_fallback(aten._dyn_quant_matmul_4bit)
     lowering.make_fallback(aten._dyn_quant_pack_4bit_weight)
diff --git a/torch/_inductor/remote_cache.py b/torch/_inductor/remote_cache.py
index d9a2d4af9d1be..9856fdbde2cbc 100644
--- a/torch/_inductor/remote_cache.py
+++ b/torch/_inductor/remote_cache.py
@@ -10,7 +10,7 @@
 import sys
 import typing
 from abc import abstractmethod
-from typing import Any, Generic, Optional, TypeAlias, TypeVar, Union
+from typing import Any, Generic, TypeAlias, TypeVar, Union
 from typing_extensions import override
 
 from torch._dynamo.utils import dynamo_timed
@@ -73,14 +73,14 @@ def __init__(self) -> None:
         self._name = f"backend:{type(self).__name__}"
 
     @abstractmethod
-    def _get(self, key: str) -> Optional[_T]:
+    def _get(self, key: str) -> _T | None:
         pass
 
     @abstractmethod
     def _put(self, key: str, data: _T) -> None:
         pass
 
-    def get(self, key: str) -> Optional[_T]:
+    def get(self, key: str) -> _T | None:
         try:
             value = self._get(key)
             cache_stats.get(self._name, value)
@@ -109,8 +109,8 @@ def decode(self, data: _U) -> _T:
         pass
 
 
-JsonDataTy = Optional[
-    Union[int, float, str, bool, dict[str, "JsonDataTy"], list["JsonDataTy"]]
+JsonDataTy: TypeAlias = Union[  # noqa: UP007
+    int, float, str, bool, dict[str, "JsonDataTy"], list["JsonDataTy"], None
 ]
 
 
@@ -154,7 +154,7 @@ def decode(self, data: _T) -> _T:
 # use the concrete type of the RemoteCache as the reported cache. See
 # RemoteFxGraphCache below as an example.
 class RemoteCache(Generic[_T]):
-    backend_override_cls: Optional[Callable[[], RemoteCacheBackend[Any]]] = None
+    backend_override_cls: Callable[[], RemoteCacheBackend[Any]] | None = None
 
     def __init__(
         self, backend: RemoteCacheBackend[_U], serde: RemoteCacheSerde[_T, _U]
@@ -169,7 +169,7 @@ def __init__(
 
     # See if the cache contains `key`. Returns `None` if the value is not
     # present in the cache.
-    def get(self, key: str) -> Optional[_T]:
+    def get(self, key: str) -> _T | None:
         with _WaitCounter("pytorch.remote_cache.get").guard():
             sample = self._create_sample()
             try:
@@ -203,16 +203,16 @@ def put(self, key: str, value: _T) -> None:
                 self._log_sample(sample)
 
     # Used to convert data from the cache into structured data.
-    def _decode(self, data: _U, sample: Optional[Sample]) -> _T:  # type: ignore[override]
+    def _decode(self, data: _U, sample: Sample | None) -> _T:  # type: ignore[override]
         return self.serde.decode(data)  # type: ignore[arg-type]
 
     # Used to convert structured data into data for the cache.
-    def _encode(self, value: _T, sample: Optional[Sample]) -> object:  # returns _U
+    def _encode(self, value: _T, sample: Sample | None) -> object:  # returns _U
         return self.serde.encode(value)
 
     # Get structured data from the cache.
     # Separate from `get` so that it can be overridden.
-    def _get(self, key: str, sample: Optional[Sample]) -> Optional[_T]:
+    def _get(self, key: str, sample: Sample | None) -> _T | None:
         if data := self._backend_get(key):
             return self._decode(data, sample)
         return None
@@ -225,7 +225,7 @@ def _backend_get(self, key: str) -> object:
 
     # Put structured data into the cache.
     # Separate from `put` so that it can be overridden.
-    def _put(self, key: str, value: _T, sample: Optional[Sample]) -> None:
+    def _put(self, key: str, value: _T, sample: Sample | None) -> None:
         data = self._encode(value, sample)
         self._backend_put(key, data)
 
@@ -237,11 +237,11 @@ def _backend_put(self, key: str, data: object) -> None:
 
     # Create a logging Sample - used with internal loggers to monitor cache
     # effectiveness.
-    def _create_sample(self) -> Optional[Sample]:
+    def _create_sample(self) -> Sample | None:
         return None
 
     # Write the logging Sample to the logger.
-    def _log_sample(self, sample: Optional[Sample]) -> None:
+    def _log_sample(self, sample: Sample | None) -> None:
         pass
 
 
@@ -251,7 +251,7 @@ class RedisRemoteCacheBackend(RemoteCacheBackend[bytes]):
     """
 
     # pyrefly: ignore [missing-attribute]
-    _redis: Optional[redis.Redis] = None
+    _redis: redis.Redis | None = None
 
     def __init__(self, cache_id: str) -> None:
         super().__init__()
@@ -267,7 +267,7 @@ def __init__(self, cache_id: str) -> None:
             )
 
     @override
-    def _get(self, key: str) -> Optional[bytes]:
+    def _get(self, key: str) -> bytes | None:
         if not self._redis:
             # Either redis wasn't found or we already had some trouble...
             return None
@@ -320,12 +320,12 @@ def _get_key(self, key: str) -> str:
         return self._key_fmt.format(key=key)
 
     @override
-    def _get(self, key: str, sample: Optional[Sample]) -> Optional[JsonDataTy]:
+    def _get(self, key: str, sample: Sample | None) -> JsonDataTy | None:
         key = self._get_key(key)
         return super()._get(key, sample)
 
     @override
-    def _put(self, key: str, value: JsonDataTy, sample: Optional[Sample]) -> None:
+    def _put(self, key: str, value: JsonDataTy, sample: Sample | None) -> None:
         key = self._get_key(key)
         super()._put(key, value, sample)
 
@@ -355,7 +355,7 @@ def create_cache(
     is_fbcode: bool,
     fb_cache_cls: str,
     oss_cache_cls: str,
-) -> Optional[RemoteCache[JsonDataTy]]:
+) -> RemoteCache[JsonDataTy] | None:
     try:
         if is_fbcode:
             import torch._inductor.fb.remote_cache
@@ -397,7 +397,7 @@ def miss(self, name: str, count: int = 1) -> None:
     def hit(self, name: str, count: int = 1) -> None:
         self._stats[name].hit += count
 
-    def get(self, name: str, value: Optional[object]) -> None:
+    def get(self, name: str, value: object | None) -> None:
         if value is None:
             self.miss(name)
         else:
diff --git a/torch/_inductor/rocm_multiarch_utils.py b/torch/_inductor/rocm_multiarch_utils.py
index a1a6103e10915..8a713d4ca033d 100644
--- a/torch/_inductor/rocm_multiarch_utils.py
+++ b/torch/_inductor/rocm_multiarch_utils.py
@@ -3,14 +3,18 @@
 Compile LLVM IR to multi-arch bundles that HIP can load automatically.
 """
 
+import logging
 import os
+import re
 import subprocess
-from typing import Optional
 
 import torch
 from torch.utils.cpp_extension import _join_rocm_home, ROCM_HOME
 
 
+log = logging.getLogger(__name__)
+
+
 def get_rocm_compiler() -> str:
     """
     Get path to ROCm's clang compiler.
@@ -69,19 +73,21 @@ def get_rocm_bundler() -> str:
 
 
 def get_rocm_target_archs() -> list[str]:
-    """
-    Get target architectures from environment or config.
-    Returns: List of architecture strings (e.g., ['gfx90a', 'gfx942'])
-    """
-    # Check PYTORCH_ROCM_ARCH environment variable
     env_archs = os.environ.get("PYTORCH_ROCM_ARCH", "").strip()
     if env_archs:
         archs = [arch.strip() for arch in env_archs.replace(";", ",").split(",")]
         archs = [arch for arch in archs if arch]
         if archs:
+            # Ensure current device arch is included
+            if torch.cuda.is_available():
+                for dev_idx in range(torch.cuda.device_count()):
+                    current_arch = torch.cuda.get_device_properties(
+                        dev_idx
+                    ).gcnArchName.split(":")[0]
+                    if current_arch not in archs:
+                        archs.append(current_arch)
             return archs
 
-    # Try to get from inductor config
     try:
         from torch._inductor import config
 
@@ -96,6 +102,43 @@ def get_rocm_target_archs() -> list[str]:
     return torch.cuda.get_arch_list()
 
 
+def _sanitize_llvm_ir_for_rocm(llvm_ir_path: str) -> str:
+    """
+    Sanitize LLVM IR to be compatible with ROCm's clang.
+
+    Triton's LLVM (upstream) may emit attributes and metadata that ROCm's
+    older clang does not yet support. Only strips attributes confirmed to
+    cause parse errors — preserves all others to maintain correct codegen.
+
+    Currently strips:
+        - nocreateundeforpoison: function attribute (upstream LLVM, not in ROCm)
+        - dwarfAddressSpace: debug metadata field (upstream LLVM, not in ROCm)
+
+    Returns:
+        Path to sanitized .ll file, or original path if no changes needed.
+    """
+    with open(llvm_ir_path) as f:
+        content = f.read()
+
+    sanitized = content
+    sanitized = re.sub(r"\bnocreateundeforpoison\b\s*", "", sanitized)
+    sanitized = re.sub(r",\s*dwarfAddressSpace:\s*\d+", "", sanitized)
+
+    if sanitized == content:
+        return llvm_ir_path
+
+    sanitized_path = llvm_ir_path + ".sanitized.ll"
+    with open(sanitized_path, "w") as f:
+        f.write(sanitized)
+
+    log.debug(
+        "Sanitized LLVM IR for ROCm clang compatibility: %s -> %s",
+        llvm_ir_path,
+        sanitized_path,
+    )
+    return sanitized_path
+
+
 def compile_llvm_ir_to_code_object(
     llvm_ir_path: str, output_path: str, target_arch: str
 ) -> bool:
@@ -120,6 +163,9 @@ def compile_llvm_ir_to_code_object(
     except RuntimeError:
         return False
 
+    # Sanitize LLVM IR to remove attributes unsupported by ROCm's clang
+    llvm_ir_path = _sanitize_llvm_ir_for_rocm(llvm_ir_path)
+
     # Using clang and not hipcc since we are not compiling source code
     # Instead we use the LLVM IR (.ll) provided by triton
     cmd = [
@@ -213,7 +259,7 @@ def create_multiarch_bundle(code_objects: dict, output_bundle_path: str) -> bool
 
 
 def compile_multiarch_bundle_from_llvm_ir(
-    llvm_ir_path: str, output_bundle_path: str, target_archs: Optional[list[str]] = None
+    llvm_ir_path: str, output_bundle_path: str, target_archs: list[str] | None = None
 ) -> bool:
     """
     Complete workflow: LLVM IR → multiple code objects → bundle.
diff --git a/torch/_inductor/runtime/autotune_cache.py b/torch/_inductor/runtime/autotune_cache.py
index f6c0531001273..5a162342faa52 100644
--- a/torch/_inductor/runtime/autotune_cache.py
+++ b/torch/_inductor/runtime/autotune_cache.py
@@ -286,6 +286,10 @@ def save(
             "time_taken_ms": time_taken_ns // 1000000,  # Convert from NS to MS
             "triton_cache_hash": triton_cache_hash,
         }
+        # Save extra_options if present on the config. This allows third-party
+        # backends to store custom tuned options alongside the standard config.
+        if extra_options := getattr(config, "extra_options", None):
+            data["extra_options"] = extra_options
         if HAS_WARP_SPEC:
             data.update(
                 {
@@ -553,6 +557,10 @@ def _load_cached_autotuning(
 
     best_config.pop("triton_cache_hash", None)
 
+    # Extract extra_options if present. This allows third-party backends
+    # to restore custom tuned options from the cache.
+    extra_options = best_config.pop("extra_options", None)
+
     if inductor_meta.get("coordinate_descent_tuning") and best_config.pop(
         "found_by_coordesc", False
     ):
@@ -580,6 +588,9 @@ def _load_cached_autotuning(
         triton_config = Config(best_config, **config_args)
         # pyrefly: ignore [missing-attribute]
         triton_config.found_by_coordesc = True
+        # Restore extra_options (may be None if not used by backend)
+        # pyrefly: ignore [missing-attribute]
+        triton_config.extra_options = extra_options
         return triton_config
 
     matching_configs = [
@@ -595,7 +606,11 @@ def _load_cached_autotuning(
     if len(matching_configs) != 1:
         return None
 
-    return matching_configs[0]
+    matched_config = matching_configs[0]
+    # Restore extra_options (may be None if not used by backend)
+    # pyrefly: ignore [missing-attribute]
+    matched_config.extra_options = extra_options
+    return matched_config
 
 
 class _LocalAutotuneCacheBackend(RemoteCacheBackend[bytes]):
diff --git a/torch/_inductor/runtime/benchmarking.py b/torch/_inductor/runtime/benchmarking.py
index 6bee8efee7c1f..cb7c5dcaefa8e 100644
--- a/torch/_inductor/runtime/benchmarking.py
+++ b/torch/_inductor/runtime/benchmarking.py
@@ -5,7 +5,7 @@
 from functools import cached_property, wraps
 from itertools import chain
 from statistics import median
-from typing import Any, Concatenate, Optional, Union
+from typing import Any, Concatenate
 from typing_extensions import ParamSpec, Self, TypeVar
 
 import torch
@@ -27,6 +27,40 @@
 T = TypeVar("T")
 
 
+# Device-type → benchmarking function registry.
+# Keys must match torch.device.type (e.g., "cpu", "cuda", "mps", "xpu", ...).
+# Values are callables with signature:
+#   fn(self: Benchmarker, _callable: Callable[..., Any], *, warmup: int, rep: int, **kwargs) -> Any
+_BENCHMARK_DISPATCH: dict[str, Callable[..., Any]] = {}
+
+
+def register_benchmarker(
+    device_type: str,
+    fn: Callable[..., Any],
+    *,
+    override: bool = False,
+) -> None:
+    """
+    Register a device-type specific benchmarker.
+
+    Args:
+        device_type: torch.device.type string (e.g., "cuda", "cpu", "mps", "xpu").
+        fn: callable(self, _callable, *, warmup, rep, **kwargs) -> Any
+        override: allow overriding an existing registration.
+    """
+    if not isinstance(device_type, str) or not device_type:
+        raise ValueError(
+            "device_type must be a non-empty string matching torch.device.type"
+        )
+    if not callable(fn):
+        raise TypeError("fn must be callable")
+    if not override and device_type in _BENCHMARK_DISPATCH:
+        raise ValueError(
+            f"Benchmarker for device_type '{device_type}' already registered"
+        )
+    _BENCHMARK_DISPATCH[device_type] = fn
+
+
 def may_distort_benchmarking_result(fn: Callable[..., Any]) -> Callable[..., Any]:
     from torch._inductor import config
 
@@ -107,7 +141,7 @@ def __init__(self: Self) -> None:
         pass
 
     def infer_device(self, *fn_args: Any, **fn_kwargs: Any) -> torch.device:
-        inferred_device: Optional[torch.device] = None
+        inferred_device: torch.device | None = None
         for arg_or_kwarg in chain(fn_args, fn_kwargs.values()):
             # Some callables take nested structures as arguments so use the
             # flattened form to find any tensors
@@ -134,9 +168,9 @@ def infer_device(self, *fn_args: Any, **fn_kwargs: Any) -> torch.device:
     def benchmark(
         self: Self,
         fn: Callable[..., Any],
-        fn_args: Optional[tuple[Any, ...]] = None,
-        fn_kwargs: Optional[dict[str, Any]] = None,
-        device: Optional[Union[str, torch.device]] = None,
+        fn_args: tuple[Any, ...] | None = None,
+        fn_kwargs: dict[str, Any] | None = None,
+        device: str | torch.device | None = None,
         **kwargs: Any,
     ) -> float:
         """Benchmark `fn(*fn_args, *fn_kwargs)` and return the runtime, in milliseconds (the
@@ -167,7 +201,7 @@ def benchmark(
         Returns:
         - The runtime of `fn(*fn_args, **fn_kwargs)`, in milliseconds.
         """
-        inferred_device: Optional[torch.device] = None
+        inferred_device: torch.device | None = None
         if device is not None:
             inferred_device = (
                 torch.device(device) if isinstance(device, str) else device
@@ -189,20 +223,33 @@ def benchmark(
 
         # No need to wrap if the callable takes no arguments
         if len(fn_args) == 0 and len(fn_kwargs) == 0:
-            _callable = fn
+            # Keep a true zero-arg callable type to satisfy type checkers.
+            def _callable() -> Any:
+                return fn()
         else:
-            _callable = lambda: fn(*fn_args, **fn_kwargs)  # noqa: E731
+            _args = fn_args
+            _kwargs = fn_kwargs
+
+            def _callable() -> Any:
+                return fn(*_args, **_kwargs)
 
         warmup = kwargs.pop("warmup", inductor_config.inductor_default_autotune_warmup)
         rep = kwargs.pop("rep", inductor_config.inductor_default_autotune_rep)
 
         # Surfacing all kernels during autotuning is super noisy; filtering these out.
         with DebugMode._benchmarking_inductor():
+            # First, try a registered device-specific benchmarker
+            benchmark_fn: Callable[..., Any] | None = _BENCHMARK_DISPATCH.get(
+                inferred_device.type
+            )
+            if benchmark_fn is not None:
+                return benchmark_fn(self, _callable, warmup=warmup, rep=rep, **kwargs)
+
+            # Backward-compatible default:
+            # - CPU  -> CPU benchmark path
+            # - else -> GPU benchmark path (legacy behavior retained for non-CPU)
             if inferred_device == torch.device("cpu"):
                 return self.benchmark_cpu(_callable, warmup=warmup, rep=rep, **kwargs)
-            # TODO(nmacchioni): For non-CPU functions we default to using the GPU-specific benchmarking
-            # implementation which was written specifically with CUDA devices in mind, we may want to
-            # explore alternate implementations for other device types.
             return self.benchmark_gpu(_callable, warmup=warmup, rep=rep, **kwargs)
 
     @time_and_count
@@ -244,6 +291,48 @@ def run_for(ms: int) -> list[float]:
     def benchmark_gpu(self: Self, *args: Any, **kwargs: Any) -> float:
         raise NotImplementedError
 
+    @time_and_count
+    def benchmark_gpu_with_cuda_graph(
+        self: Self,
+        _callable: Callable[[], Any],
+        **kwargs: Any,
+    ) -> float:
+        """Benchmark a GPU callable using CUDA graph capture and replay.
+
+        This captures the callable into a CUDA graph and benchmarks the graph replay,
+        which eliminates kernel launch overhead for fair comparison between different
+        implementations.
+        """
+        # Warmup
+        _callable()
+        torch.cuda.synchronize()
+
+        # Capture into CUDA graph
+        cuda_graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(cuda_graph, capture_error_mode="thread_local"):
+            _callable()
+        torch.cuda.synchronize()
+
+        return self.benchmark_gpu(cuda_graph.replay, **kwargs)
+
+
+# Make built-in defaults explicit via the registry
+def _default_cpu_bench(self, f, *, warmup, rep, **kw):
+    return self.benchmark_cpu(f, warmup=warmup, rep=rep, **kw)
+
+
+def _default_cuda_bench(self, f, *, warmup, rep, **kw):
+    return self.benchmark_gpu(f, warmup=warmup, rep=rep, **kw)
+
+
+def _default_xpu_bench(self, f, *, warmup, rep, **kw):
+    return self.benchmark_gpu(f, warmup=warmup, rep=rep, **kw)
+
+
+register_benchmarker("cpu", _default_cpu_bench, override=True)
+register_benchmarker("cuda", _default_cuda_bench, override=True)
+register_benchmarker("xpu", _default_xpu_bench, override=True)
+
 
 class TritonBenchmarker(Benchmarker):
     @cached_property
@@ -287,11 +376,23 @@ def benchmark_gpu(
         for kwarg in list(kwargs.keys()):
             if kwarg not in do_bench_params:
                 del kwargs[kwarg]
-        if "quantiles" in kwargs:
-            return self.triton_do_bench(_callable, **kwargs)[0]
-        elif "return_mode" in kwargs:
-            return self.triton_do_bench(_callable, **kwargs)
-        return self.triton_do_bench(_callable, **kwargs, return_mode="median")
+        try:
+            if "quantiles" in kwargs:
+                return self.triton_do_bench(_callable, **kwargs)[0]
+            elif "return_mode" in kwargs:
+                return self.triton_do_bench(_callable, **kwargs)
+            return self.triton_do_bench(_callable, **kwargs, return_mode="median")
+        except Exception as e:
+            # ErrorInvalidConfiguration
+            # Return inf to skip this config during autotuning
+            error_str = str(e).lower()
+            if "invalid configuration" in error_str:
+                logger.warning(
+                    "Skipping benchmark due to invalid configuration error: %s",
+                    error_str,
+                )
+                return float("inf")
+            raise
 
 
 class InductorBenchmarker(TritonBenchmarker):  # noqa: docstring_linter
diff --git a/torch/_inductor/runtime/caching/interfaces.py b/torch/_inductor/runtime/caching/interfaces.py
index 27322454b5b19..2fa20649d32be 100644
--- a/torch/_inductor/runtime/caching/interfaces.py
+++ b/torch/_inductor/runtime/caching/interfaces.py
@@ -310,6 +310,7 @@ def __call__(
     ) -> Callable[_P, Callable[[_R], _EncodedR | DeferredRecording[_R, _EncodedR]]]: ...
 
 
+# pyrefly: ignore [variance-mismatch]
 class ResultDecoderFactory(Protocol[_P, _R, _EncodedR]):
     """Protocol for custom result decoder factories.
 
diff --git a/torch/_inductor/runtime/compile_tasks.py b/torch/_inductor/runtime/compile_tasks.py
index 11801eac92584..8e19fe5d0d60c 100644
--- a/torch/_inductor/runtime/compile_tasks.py
+++ b/torch/_inductor/runtime/compile_tasks.py
@@ -51,6 +51,69 @@ def _set_triton_ptxas_path() -> None:
         warnings.warn(f"{ptxas} exists but is not an executable")
 
 
+def _set_triton_libdevice_path() -> None:
+    """
+    Use the CUDA toolkit's libdevice instead of Triton's bundled version.
+    This ensures Triton's pow matches CUDA's powf for bitwise precision.
+    Gated by config.eager_numerics.use_pytorch_libdevice.
+    """
+    from torch._inductor import config
+
+    if not config.eager_numerics.use_pytorch_libdevice:
+        return
+
+    _set_triton_libdevice_path_impl()
+
+
+def _set_triton_libdevice_path_impl() -> None:
+    try:
+        from triton import knobs
+    except ImportError:
+        return
+
+    env_path = os.environ.get("TRITON_LIBDEVICE_PATH")
+    if env_path is not None:
+        knobs.nvidia.libdevice_path = env_path
+        return
+
+    if knobs.nvidia.libdevice_path is not None:
+        return
+
+    try:
+        from torch.utils.cpp_extension import CUDA_HOME
+
+        if CUDA_HOME is None:
+            warnings.warn(
+                "CUDA_HOME not set; using Triton's bundled libdevice which may "
+                "cause minor precision differences in pow operations. "
+                "To fix: set TRITON_LIBDEVICE_PATH to your CUDA toolkit's libdevice, "
+                "e.g., export TRITON_LIBDEVICE_PATH=/usr/local/cuda/nvvm/libdevice/libdevice.10.bc",
+                stacklevel=3,
+            )
+            return
+        libdevice = Path(CUDA_HOME) / "nvvm" / "libdevice" / "libdevice.10.bc"
+        if libdevice.is_file():
+            knobs.nvidia.libdevice_path = str(libdevice)
+            # Also set env var so subprocess compile workers inherit it
+            os.environ["TRITON_LIBDEVICE_PATH"] = str(libdevice)
+        else:
+            warnings.warn(
+                f"CUDA libdevice not found at {libdevice}; using Triton's bundled "
+                "libdevice which may cause minor precision differences in pow operations. "
+                "To fix: set TRITON_LIBDEVICE_PATH to your CUDA toolkit's libdevice, "
+                "e.g., export TRITON_LIBDEVICE_PATH=/usr/local/cuda/nvvm/libdevice/libdevice.10.bc",
+                stacklevel=3,
+            )
+    except ImportError:
+        warnings.warn(
+            "torch.utils.cpp_extension not available; using Triton's bundled "
+            "libdevice which may cause minor precision differences in pow operations. "
+            "To fix: set TRITON_LIBDEVICE_PATH to your CUDA toolkit's libdevice, "
+            "e.g., export TRITON_LIBDEVICE_PATH=/usr/local/cuda/nvvm/libdevice/libdevice.10.bc",
+            stacklevel=3,
+        )
+
+
 def _worker_compile_triton(
     load_kernel: Callable[[], CachingAutotuner],
     extra_env: dict[str, str],
@@ -58,6 +121,15 @@ def _worker_compile_triton(
 ) -> tuple[CachingAutotuner, int]:
     _set_triton_ptxas_path()
     os.environ.update(extra_env)
+    # Set libdevice path if passed via env from main process
+    libdevice_path = extra_env.get("TRITON_LIBDEVICE_PATH")
+    if libdevice_path:
+        try:
+            from triton import knobs
+
+            knobs.nvidia.libdevice_path = libdevice_path
+        except ImportError:
+            pass
     from torch._inductor import config
 
     with config.patch(extra_config):
diff --git a/torch/_inductor/runtime/halide_helpers.py b/torch/_inductor/runtime/halide_helpers.py
index 3b02892eb2eb7..2b0161ead97ae 100644
--- a/torch/_inductor/runtime/halide_helpers.py
+++ b/torch/_inductor/runtime/halide_helpers.py
@@ -101,6 +101,40 @@ def rand(seed, offset, n_rounds=PHILOX_N_ROUNDS_DEFAULT):
     return _uint_to_uniform_float(source)
 
 
+def rand_eager_kernel(seed, offset_blocks, tid, VEC, n_rounds=PHILOX_N_ROUNDS_DEFAULT):
+    inv = hl.cast(hl.Float(32), 1.0 / 4294967296.0)  # 2^-32
+    half = hl.cast(hl.Float(32), 0.5) * inv
+
+    tid_u64 = hl.cast(hl.UInt(64), tid)
+    VEC_u64 = hl.cast(hl.UInt(64), VEC)
+    subseq = tid_u64 // VEC_u64
+    which4 = (tid_u64 % VEC_u64) // hl.cast(hl.UInt(64), 4)
+    lane = tid_u64 % hl.cast(hl.UInt(64), 4)
+
+    offblk = hl.cast(hl.UInt(64), offset_blocks) + which4
+
+    c0 = hl.cast(hl.UInt(32), offblk & hl.cast(hl.UInt(64), 0xFFFFFFFF))
+    c1 = hl.cast(
+        hl.UInt(32),
+        (offblk >> hl.cast(hl.UInt(64), 32)) & hl.cast(hl.UInt(64), 0xFFFFFFFF),
+    )
+    c2 = hl.cast(hl.UInt(32), subseq & hl.cast(hl.UInt(64), 0xFFFFFFFF))
+    c3 = hl.cast(
+        hl.UInt(32),
+        (subseq >> hl.cast(hl.UInt(64), 32)) & hl.cast(hl.UInt(64), 0xFFFFFFFF),
+    )
+
+    u0, u1, u2, u3 = halide_philox(seed, c0, c1, c2, c3, n_rounds)
+
+    v01 = hl.select(lane == hl.cast(hl.UInt(64), 0), u0, u1)
+    v23 = hl.select(lane == hl.cast(hl.UInt(64), 2), u2, u3)
+    rand_int = hl.select(
+        (lane == hl.cast(hl.UInt(64), 0)) | (lane == hl.cast(hl.UInt(64), 1)), v01, v23
+    )
+
+    return hl.cast(hl.Float(32), 1.0) - (hl.cast(hl.Float(32), rand_int) * inv + half)
+
+
 def randn(seed, offset):
     i1, i2, _, _ = randint4x(seed, offset, PHILOX_N_ROUNDS_DEFAULT)
     u1 = _uint_to_uniform_float(i1)
diff --git a/torch/_inductor/runtime/hints.py b/torch/_inductor/runtime/hints.py
index a9ddf91e9a59c..1552a0764aa08 100644
--- a/torch/_inductor/runtime/hints.py
+++ b/torch/_inductor/runtime/hints.py
@@ -47,6 +47,7 @@ class TileHint(Enum):
         def AttrsDescriptorWrapper(
             divisible_by_16=None,
             equal_to_1=None,
+            pointer_range_32=None,
         ):
             # Prepare the arguments for AttrsDescriptor
             kwargs = {
@@ -69,6 +70,7 @@ def AttrsDescriptorWrapper(
         def AttrsDescriptorWrapper(
             divisible_by_16=None,
             equal_to_1=None,
+            pointer_range_32=None,
         ):
             # Prepare the arguments for AttrsDescriptor
             kwargs = {
@@ -88,23 +90,30 @@ def AttrsDescriptorWrapper(
         def AttrsDescriptorWrapper(
             divisible_by_16=None,
             equal_to_1=None,
+            pointer_range_32=None,
         ):
             # pyrefly: ignore [not-iterable]
-            return {(x,): [["tt.divisibility", 16]] for x in divisible_by_16}
+            # Build attr dict merging divisibility and pointer_range per arg index,
+            # since a single arg can carry both attributes.
+            result = {(x,): [["tt.divisibility", 16]] for x in (divisible_by_16 or ())}
+            for x in pointer_range_32 or ():
+                key = (x,)
+                if key in result:
+                    result[key].append(["tt.pointer_range", 32])
+                else:
+                    result[key] = [["tt.pointer_range", 32]]
+            return result
 
 else:
     # Define a namedtuple as a fallback when AttrsDescriptor is not available
     AttrsDescriptorWrapper = collections.namedtuple(  # type: ignore[no-redef, name-match]
         # pyrefly: ignore [invalid-argument]
         "AttrsDescriptor",
-        ["divisible_by_16", "equal_to_1"],
-        defaults=[(), ()],
+        ["divisible_by_16", "equal_to_1", "pointer_range_32"],
+        defaults=[(), (), ()],
     )
 
 
-_NUM_THREADS_PER_WARP = 32
-
-
 class HeuristicType(Enum):
     PERSISTENT_REDUCTION = auto()
     POINTWISE = auto()
@@ -138,6 +147,14 @@ class DeviceProperties(typing.NamedTuple):
     max_threads_per_block: int | None = None
     warp_size: int | None = None
 
+    @property
+    def warp_size_or_default(self) -> int:
+        if self.warp_size is not None:
+            return self.warp_size
+        if self.type in ("cuda", "hip"):
+            raise RuntimeError(f"{self.type} device properties must report warp_size")
+        return 32
+
     @classmethod
     @functools.cache
     def create(cls, device) -> DeviceProperties:
@@ -171,10 +188,21 @@ def create(cls, device) -> DeviceProperties:
                 props, "max_threads_per_multi_processor", None
             ),
             max_threads_per_block=getattr(props, "max_threads_per_block", 1024),
-            warp_size=getattr(props, "warp_size", 32 if device_type != "cpu" else None),
+            warp_size=getattr(props, "warp_size", None),
         )
 
 
+def get_warp_size(device) -> int:
+    """Return the wave/warp size in threads for the given device.
+
+    Reads from torch.cuda.get_device_properties(device).warp_size via the cached
+    DeviceProperties.create(). Correct on both AMD (64 for CDNA/gfx9, 32 for
+    RDNA/gfx10+) and NVIDIA (always 32). Missing cuda/hip warp_size metadata is
+    treated as an error rather than silently falling back.
+    """
+    return DeviceProperties.create(device).warp_size_or_default
+
+
 class HalideInputSpec(typing.NamedTuple):
     ctype: str
     name: str
diff --git a/torch/_inductor/runtime/runtime_utils.py b/torch/_inductor/runtime/runtime_utils.py
index c9b039e9760f9..632c39889f8e1 100644
--- a/torch/_inductor/runtime/runtime_utils.py
+++ b/torch/_inductor/runtime/runtime_utils.py
@@ -193,6 +193,27 @@ def compile_mps_shader(source: str) -> Any:
         raise SyntaxError(f"failed to compile {source} with {err.msg}") from err
 
 
+def compile_mps_shaders(
+    kernels: list[tuple[str, str, list[str]]],
+) -> dict[str, Any]:
+    """Compile a batch of Metal kernels into one library.
+
+    Args:
+        kernels: list of (kernel_name, metal_source, headers) tuples.
+            headers are bare names resolved as <c10/metal/{name}.h>.
+
+    Returns:
+        dict mapping each kernel_name to its compiled function handle.
+    """
+    from torch.utils._ordered_set import OrderedSet
+
+    all_headers = sorted(OrderedSet(h for _, _, hs in kernels for h in hs))
+    header_src = "\n".join(f"#include <c10/metal/{h}.h>" for h in all_headers)
+    body_src = "\n".join(src for _, src, _ in kernels)
+    lib = compile_mps_shader(header_src + "\n" + body_src)
+    return {name: getattr(lib, name) for name, _, _ in kernels}
+
+
 def torch_dtype_to_jax_runtime(dtype: torch.dtype) -> Any:
     """
     Map PyTorch dtype to actual JAX dtype object at runtime.
@@ -247,47 +268,6 @@ def torch_dtype_to_jax(dtype: torch.dtype) -> str:
     return f"jnp.{dtype_name}"
 
 
-def pallas_partial_reduce(reduce_fn: Any, v: Any, pw_numel: int, red_numel: int) -> Any:
-    """
-    Helper for partial reductions in Pallas kernels.
-
-    Reorders axes and reduces, returning result with keepdims-style shape
-    for proper in-kernel broadcasting.
-
-    Args:
-        reduce_fn: The reduction function to apply (e.g., jnp.sum, jnp.max)
-        v: The input array to reduce
-        pw_numel: The number of pointwise elements
-        red_numel: The number of reduction elements
-
-    Returns:
-        Reduced array with keepdims-style shape
-    """
-    import jax.numpy as jnp  # pyrefly: ignore [import-error, missing-import]
-
-    shape = tuple(v.shape)
-    # Find contiguous axes whose product = red_numel (search from right)
-    red_axes = None
-    for i in range(len(shape) - 1, -1, -1):
-        prod = 1
-        for j in range(i, -1, -1):
-            prod *= shape[j]
-            if prod == red_numel:
-                red_axes = list(range(j, i + 1))
-                break
-        if red_axes is not None:
-            break
-    if red_axes is None:
-        red_axes = [len(shape) - 1]
-    # Build output shape with 1s for reduced dimensions (keepdims style)
-    out_shape = tuple(1 if i in red_axes else s for i, s in enumerate(shape))
-    # Move pointwise axes to front, reduction axes to back
-    pw_axes = [i for i in range(len(shape)) if i not in red_axes]
-    reordered = jnp.moveaxis(v, pw_axes, list(range(len(pw_axes))))
-    result = reduce_fn(reordered.reshape(pw_numel, red_numel), axis=-1)
-    return result.reshape(out_shape)
-
-
 def pallas_gpu_pad_inputs(inputs: list[Any], alignment: int = 128) -> list[Any]:
     """Flatten and pad each input JAX array to a multiple of alignment."""
     import jax.numpy as jnp  # pyrefly: ignore [import-error, missing-import]
@@ -349,3 +329,501 @@ def pallas_gpu_unpad_results(
             orig_numel = math.prod(shape)
             unpadded.append(res[:orig_numel].reshape(shape))
     return unpadded[0] if len(unpadded) == 1 else tuple(unpadded)
+
+
+# ---------------------------------------------------------------------------
+# Pallas CPU / TPU tiling helpers
+# ---------------------------------------------------------------------------
+# TPU alignment: last dim must be full or a multiple of 128,
+#                second-to-last dim must be full or a multiple of 8.
+_TPU_ALIGN_LAST = 128
+_TPU_ALIGN_SECOND_LAST = 8
+
+
+def _pallas_tile_size(
+    dim: int, alignment: int, max_tile: int = 1024, is_tpu: bool = False
+) -> int:
+    """Pick the largest aligned tile size <= max_tile for *dim*.
+
+    If *dim* is already <= alignment the full dimension is used (no tiling
+    on this axis).
+    """
+    if dim == 0:
+        # Tile size >= 1 avoids division by zero in JAX's _pad_to_block_dimension
+        return alignment if is_tpu else 1
+
+    if is_tpu:
+        # On TPU, Mosaic requires block dimensions to perfectly align to hardware
+        # registers (128 for inner, 8 for outer). We MUST pad the block spec up to
+        # the next alignment boundary (Mosaic handles the OOB masking).
+        if dim <= alignment:
+            return alignment
+        t = min(max_tile, dim)
+        # Use ceiling division to ensure the tile covers the dimension or aligns upward
+        t = ((t + alignment - 1) // alignment) * alignment
+        return t
+
+    if dim <= alignment:
+        return dim
+    t = min(max_tile, dim)
+    t = (t // alignment) * alignment
+    return max(alignment, t)
+
+
+def pallas_permute(x, perm):
+    """Permute array x according to perm, working around Mosaic TPU bugs.
+
+    Mosaic's jnp.permute_dims produces corrupted output for certain 3D+
+    permutations at large sizes.  Decomposes into: loop over smallest output
+    dim via unrolled Python loop + stack, with recursive sub-permutation.
+    At the 2D base case, uses jnp.permute_dims (which works on Mosaic).
+
+    Callers should ensure tiles are small enough via tiling so that the
+    Mosaic bug doesn't trigger (threshold ~ 256K elements).
+    """
+    import jax.numpy as jnp  # pyrefly: ignore [import-error, missing-import]
+
+    ndim = len(perm)
+    if ndim <= 2:
+        return jnp.permute_dims(x, perm)
+    if perm == tuple(range(ndim)):
+        return x  # identity
+
+    shape = x.shape
+    out_shape = tuple(shape[p] for p in perm)
+
+    # Choose the output dimension with the smallest size to loop over.
+    loop_out_dim = min(range(ndim), key=lambda d: out_shape[d])
+    loop_size = out_shape[loop_out_dim]
+    loop_in_dim = perm[loop_out_dim]
+
+    # Build the sub-permutation for the remaining (ndim-1) dimensions.
+    remaining_in_dims = [d for d in range(ndim) if d != loop_in_dim]
+    remaining_out_perm_raw = [perm[d] for d in range(ndim) if d != loop_out_dim]
+    dim_map = {old: new for new, old in enumerate(remaining_in_dims)}
+    sub_perm = tuple(dim_map[d] for d in remaining_out_perm_raw)
+
+    # Unrolled loop: extract slices with static indices, apply sub-perm,
+    # then stack along the loop output dimension.
+    slices = []
+    # pyrefly: ignore [bad-argument-type]
+    for i in range(loop_size):
+        idx: list[Any] = [slice(None)] * ndim
+        idx[loop_in_dim] = i
+        slc = x[tuple(idx)]  # static index, removes loop_in_dim
+        slc_p = pallas_permute(slc, sub_perm)
+        slices.append(slc_p)
+
+    return jnp.stack(slices, axis=loop_out_dim)
+
+
+# Maximum tile elements for permutation kernels.  Mosaic's
+# jnp.permute_dims produces corrupted output above ~1.5M elements for
+# certain 3D+ permutations.  Keep tiles well under that threshold.
+_PERMUTE_MAX_TILE_ELEMS = 1 << 18  # ~256K elements
+
+
+def pallas_compute_tiling(
+    ref_shape: tuple[int, ...],
+    transpose: bool = False,  # constrain tile size for pallas_permute VMEM
+    skip_last_n: int = 0,
+    exact_only: bool = False,
+    is_tpu: bool = False,
+    permutations: list[tuple[int, ...]] | None = None,
+    max_grid_product: int | None = None,
+) -> tuple[tuple[int, ...], tuple[int, ...], dict[int, int]]:
+    """Compute tile shape, grid and axis→grid-dim mapping for CPU/TPU.
+
+    Always uses TPU-compatible alignment (last dim multiple of 128,
+    second-to-last multiple of 8) so that the same generated kernel works
+    on both CPU-interpret and real TPU.
+
+    *transpose* constrains tile sizes to keep VMEM usage within safe
+    limits for pallas_permute (gather-based permutation).
+
+    *skip_last_n* prevents tiling the last N dimensions (used when those
+    dims correspond to internal reduction ranges that must remain full).
+
+    *exact_only* restricts tiling to dimensions that divide evenly by the
+    tile size (no remainder blocks).  Required on TPU where Mosaic needs
+    block shapes to match the XLA memory layout.
+
+    Returns ``(tile_shape, grid, axis_to_grid)`` where *axis_to_grid*
+    maps each tiled reference-shape axis index to its position in the
+    grid tuple.
+
+    When no dimension benefits from tiling the grid is ``(1,)`` and the
+    tile covers the full tensor.
+    """
+    nd = len(ref_shape)
+    if nd == 0:
+        return (), (1,), {}
+
+    # Effective number of dims eligible for tiling
+    tileable_nd = nd - skip_last_n
+
+    tile = list(ref_shape)
+    grid_parts: list[int] = []
+    axis_to_grid: dict[int, int] = {}  # ref axis → grid dim
+
+    # Pick alignment based on the physical position of the axis in the
+    # tensor, not its position in the tileable subset.  The TPU requires
+    # the physical last dim to be a multiple of 128 and the physical
+    # second-to-last dim to be a multiple of 8.
+    #
+    # When permutations are present, an output dim `ax` maps to input
+    # dim `perm[ax]`.  If `perm[ax]` is the input's last dim, that
+    # output dim must also satisfy _TPU_ALIGN_LAST.  Compute the
+    # strictest alignment across all buffers.
+    def _input_align(input_ax: int, input_nd: int) -> int:
+        return _TPU_ALIGN_LAST if input_ax == input_nd - 1 else _TPU_ALIGN_SECOND_LAST
+
+    def _align(ax: int) -> int:
+        out_align = _TPU_ALIGN_LAST if ax == nd - 1 else _TPU_ALIGN_SECOND_LAST
+        if permutations:
+            for perm in permutations:
+                if perm is not None and len(perm) == nd:
+                    in_ax = perm[ax]
+                    in_align = _input_align(in_ax, len(perm))
+                    out_align = max(out_align, in_align)
+        return out_align
+
+    def _can_tile_ax(ax: int, dim: int, t: int) -> bool:
+        """Check if tiling dim to t is valid."""
+        if t >= dim:
+            # For TPU padding, we allow tiles >= dimension for the aligned axes
+            if _align(ax) == _TPU_ALIGN_LAST or _align(ax) == _TPU_ALIGN_SECOND_LAST:
+                return True
+            return False
+        if exact_only and dim % t != 0:
+            if _align(ax) == _TPU_ALIGN_LAST or _align(ax) == _TPU_ALIGN_SECOND_LAST:
+                # TPU DMA `#tpu.element_window` natively masks out-of-bounds remainder tiles
+                return True
+            return False
+        return True
+
+    # Second-to-last tileable dim (added first so it becomes grid dim 0)
+    if tileable_nd >= 2:
+        ax = tileable_nd - 2
+        t = _pallas_tile_size(ref_shape[ax], _align(ax), is_tpu=is_tpu)
+        if _can_tile_ax(ax, ref_shape[ax], t):
+            tile[ax] = t
+            axis_to_grid[ax] = len(grid_parts)
+            grid_parts.append((ref_shape[ax] + t - 1) // t)
+
+    # Last tileable dim
+    if tileable_nd >= 1:
+        ax = tileable_nd - 1
+        t = _pallas_tile_size(ref_shape[ax], _align(ax), is_tpu=is_tpu)
+        if _can_tile_ax(ax, ref_shape[ax], t):
+            tile[ax] = t
+            axis_to_grid[ax] = len(grid_parts)
+            grid_parts.append((ref_shape[ax] + t - 1) // t)
+
+    # When transpose is active, tile dimensions to keep the total tile
+    # elements below the safe threshold for Mosaic's permute_dims.
+    # Mosaic pads the last dimension to _TPU_ALIGN_LAST (128 for f32),
+    # so when the last dim is small (e.g. 2), VMEM usage is much higher
+    # than the logical tile size.  We compute an effective element count
+    # that accounts for this padding.
+    if transpose:
+        tile_elems = 1
+        for t in tile:
+            tile_elems *= t
+        last_dim = tile[nd - 1] if nd > 0 else 1
+        padded_last = max(last_dim, _TPU_ALIGN_LAST)
+        # pallas_permute decomposes ndim>2 permutations by looping over
+        # the smallest dim and recursing, creating (ndim-1)-D intermediates
+        # at each level.  Mosaic pads the last physical dim of each
+        # intermediate to _TPU_ALIGN_LAST (128), which can cause VMEM OOM
+        # when a small dim (e.g. 4) ends up as the last dim.
+        #
+        # For ndim >= 3, compute effective_max from the worst-case VMEM
+        # of the (ndim-1)-D intermediate created by pallas_permute.
+        # The intermediate removes the loop dim (smallest output dim) and
+        # may have a small dim as its last physical dim -> padded to 128.
+        if nd >= 4:
+            # For 4D+, pallas_permute creates (nd-1)-D intermediates
+            # that can have small dims as their last physical dim,
+            # which Mosaic pads to _TPU_ALIGN_LAST (128).  This can
+            # cause VMEM OOM.  (3D pallas_permute only creates 2D
+            # intermediates which Mosaic handles efficiently.)
+            sorted_dims = sorted(tile)
+            # The loop dim is the smallest; intermediate has the rest.
+            inter_dims = sorted_dims[1:]  # (nd-1) dims, sorted
+            smallest_inter = inter_dims[0]
+            # Padded intermediate: smallest remaining dim padded to 128.
+            padded_inter_elems = 1
+            for d in inter_dims:
+                padded_inter_elems *= d
+            pad_factor = max(smallest_inter, _TPU_ALIGN_LAST) // max(smallest_inter, 1)
+            padded_inter_elems *= pad_factor
+            padded_inter_bytes = padded_inter_elems * 4
+            # Conservative VMEM budget: 2MB per intermediate buffer.
+            # pallas_permute creates multiple live intermediates (input
+            # slices, sub-permutation results, stack temps).  Empirically,
+            # 2MB per worst-case intermediate avoids OOM on TPU v4 (64MB VMEM).
+            vmem_per_buf = 2 * 1024 * 1024  # 2MB
+            if padded_inter_bytes > vmem_per_buf:
+                scale = vmem_per_buf / padded_inter_bytes
+                effective_max = max(int(tile_elems * scale), 1024)
+            else:
+                effective_max = _PERMUTE_MAX_TILE_ELEMS
+        else:
+            # For ndim <= 3: pallas_permute creates 2D intermediates
+            # (fine) but jnp.stack creates a 3D result with the same
+            # shape as the output.  If the last dim is small, Mosaic
+            # pads it to _TPU_ALIGN_LAST (128), which can exhaust VMEM.
+            # Empirically: 8MB per padded buffer keeps us safe on TPU v4
+            # (64MB VMEM, multiple buffers live simultaneously).
+            vmem_per_buf = 8 * 1024 * 1024  # 8MB
+            max_nonlast_product = vmem_per_buf // (padded_last * 4)
+            effective_max = min(
+                max_nonlast_product * last_dim,
+                _PERMUTE_MAX_TILE_ELEMS,
+            )
+        effective_max = min(effective_max, _PERMUTE_MAX_TILE_ELEMS)
+        effective_max = max(effective_max, 1024)
+
+        if tile_elems > effective_max:
+            # Collect tileable dims.  For ndim >= 4 with VMEM-constrained
+            # intermediates, include ALL tileable dims (even the last)
+            # since we may need to shrink every dim to fit in VMEM.
+            # For ndim <= 3, exclude the last dim (it's often small
+            # and already at minimum).
+            if nd >= 4:
+                tileable_axes = [
+                    ax for ax in range(tileable_nd - 1, -1, -1) if ref_shape[ax] > 1
+                ]
+            else:
+                tileable_axes = [
+                    ax for ax in range(tileable_nd - 2, -1, -1) if ref_shape[ax] > 1
+                ]
+            # Distribute reduction proportionally across dims.
+            reduction = tile_elems / effective_max
+            n = len(tileable_axes)
+            per_dim = reduction ** (1.0 / max(n, 1))
+
+            for ax in tileable_axes:
+                if tile_elems <= effective_max:
+                    break
+                dim = ref_shape[ax]
+                cur_t = tile[ax]
+                align = _align(ax)
+                # Target tile: proportional reduction, but also check
+                # against the remaining budget for subsequent dims.
+                target = int(cur_t / per_dim)
+                remaining_budget = effective_max * cur_t // tile_elems
+                target = min(target, remaining_budget)
+                target = (target // align) * align
+                target = max(target, align)
+                # Find largest aligned divisor of dim <= target.
+                best_t = cur_t
+                for candidate in range(target, 0, -align):
+                    if candidate >= cur_t:
+                        continue
+                    if dim % candidate != 0:
+                        continue
+                    best_t = candidate
+                    break
+                if best_t < cur_t:
+                    if ax in axis_to_grid:
+                        grid_parts[axis_to_grid[ax]] = dim // best_t
+                    else:
+                        axis_to_grid[ax] = len(grid_parts)
+                        grid_parts.append(dim // best_t)
+                    tile_elems = tile_elems * best_t // cur_t
+                    tile[ax] = best_t
+
+    grid = tuple(grid_parts) if grid_parts else (1,)
+
+    # Enforce max_grid_product by scaling up tiles on tiled axes.
+    if max_grid_product is not None and grid_parts:
+        grid_product = 1
+        for g in grid_parts:
+            grid_product *= g
+        if grid_product > max_grid_product:
+            # Sort tiled axes by their grid contribution (descending)
+            # and increase tiles to reduce the grid.
+            tiled_axes = sorted(
+                axis_to_grid.keys(),
+                key=lambda ax: grid_parts[axis_to_grid[ax]],
+                reverse=True,
+            )
+            for ax in tiled_axes:
+                if grid_product <= max_grid_product:
+                    break
+                gi = axis_to_grid[ax]
+                cur_grid = grid_parts[gi]
+                cur_tile = tile[ax]
+                dim = ref_shape[ax]
+                align = _align(ax)
+                # Target grid for this axis: proportional share of reduction
+                target_grid = max(1, int(cur_grid * max_grid_product / grid_product))
+                target_tile = (dim + target_grid - 1) // target_grid
+                # Round up to alignment
+                target_tile = ((target_tile + align - 1) // align) * align
+                target_tile = min(target_tile, dim)
+                # Find smallest aligned value >= target_tile that divides dim
+                # (or just use the rounded-up value if exact_only is False)
+                best_tile = dim  # fallback: full dim
+                if exact_only:
+                    for candidate in range(target_tile, dim + 1, align):
+                        if dim % candidate == 0:
+                            best_tile = candidate
+                            break
+                else:
+                    best_tile = target_tile
+                if best_tile > cur_tile:
+                    new_grid = (dim + best_tile - 1) // best_tile
+                    grid_product = grid_product * new_grid // cur_grid
+                    grid_parts[gi] = new_grid
+                    tile[ax] = best_tile
+            grid = tuple(grid_parts)
+
+    return tuple(tile), grid, axis_to_grid
+
+
+def pallas_make_block_spec(
+    buf_shape: tuple[int, ...],
+    ref_shape: tuple[int, ...],
+    tile_shape: tuple[int, ...],
+    axis_to_grid: dict[int, int],
+    n_grid: int,
+    permutation: tuple[int, ...] | None = None,
+    is_output: bool = False,
+) -> Any:
+    """Build a ``pl.BlockSpec`` for *buf_shape* given tiling of *ref_shape*.
+
+    Lower-ndim buffers are right-aligned with the reference shape (numpy
+    broadcast rules).  Dimensions that match a tiled reference dimension
+    are tiled; broadcast dimensions (size 1 or absent) are kept full.
+
+    When *buf_nd > ref_nd* (reduction inputs), we find an alignment offset
+    so the ref dims map into the buffer.  Extra dims are kept at full size
+    with index 0 in the index_map.
+
+    When *permutation* is given (a tuple mapping ref axis -> buf axis),
+    the buffer dimensions are accessed in permuted order relative to the
+    reference shape.  For example, permutation=(1, 0) swaps the last two.
+
+    When *is_output* is True and *buf_nd < ref_nd*, left-alignment is used
+    as a fallback (for reduction outputs whose trailing dims were reduced).
+    """
+    from jax.experimental import (  # pyrefly: ignore [import-error, missing-import]
+        pallas as pl,
+    )
+
+    buf_nd = len(buf_shape)
+    ref_nd = len(ref_shape)
+
+    if buf_nd == 0:
+        # Scalar — untouched regardless of grid shape.
+        return pl.BlockSpec((1,), _make_index_map([], 1, n_grid))
+
+    bs = list(buf_shape)
+    tiled_pairs: list[tuple[int, int]] = []
+
+    if buf_nd > ref_nd:
+        # Reduction input: find alignment offset k where ref dims map into buf.
+        align_k = 0
+        for k in range(buf_nd - ref_nd + 1):
+            ok = True
+            for i in range(ref_nd):
+                if ref_shape[i] == 1:
+                    continue
+                if buf_shape[k + i] != ref_shape[i]:
+                    ok = False
+                    break
+            if ok:
+                align_k = k
+                break
+
+        for ref_ax, grid_dim in axis_to_grid.items():
+            buf_ax = align_k + ref_ax
+            if 0 <= buf_ax < buf_nd and buf_shape[buf_ax] == ref_shape[ref_ax]:
+                bs[buf_ax] = tile_shape[ref_ax]
+                tiled_pairs.append((buf_ax, grid_dim))
+
+    elif permutation is not None and buf_nd == ref_nd:
+        # Permuted buffer: map ref axes through the permutation.
+        for ref_ax, grid_dim in axis_to_grid.items():
+            buf_ax = permutation[ref_ax]
+            if 0 <= buf_ax < buf_nd and buf_shape[buf_ax] == ref_shape[ref_ax]:
+                bs[buf_ax] = tile_shape[ref_ax]
+                tiled_pairs.append((buf_ax, grid_dim))
+
+    else:
+        # Standard right-alignment, with left-alignment fallback for
+        # reduction outputs (e.g. sum(dim=-1) on (10,10) → (10,)).
+        for ref_ax, grid_dim in axis_to_grid.items():
+            buf_ax = ref_ax - (ref_nd - buf_nd)
+            if 0 <= buf_ax < buf_nd and buf_shape[buf_ax] == ref_shape[ref_ax]:
+                bs[buf_ax] = tile_shape[ref_ax]
+                tiled_pairs.append((buf_ax, grid_dim))
+            elif (
+                is_output
+                and buf_nd < ref_nd
+                and 0 <= ref_ax < buf_nd
+                and buf_shape[ref_ax] == ref_shape[ref_ax]
+            ):
+                # Left-aligned match for output: buf dim i matches ref dim i
+                # (reduction output whose trailing dims were reduced away)
+                bs[ref_ax] = tile_shape[ref_ax]
+                tiled_pairs.append((ref_ax, grid_dim))
+
+    return pl.BlockSpec(
+        tuple(bs),
+        _make_index_map(tiled_pairs, buf_nd, n_grid),
+    )
+
+
+def _make_index_map(
+    tiled_pairs: list[tuple[int, int]],
+    buf_nd: int,
+    n_grid: int,
+) -> Any:
+    """Return an index_map callable for ``pl.BlockSpec``.
+
+    *tiled_pairs* is a list of ``(buf_axis, grid_dim)`` indicating which
+    buffer axes receive a grid index.  All other axes return 0 (full block).
+
+    All returned values are explicitly ``jnp.int32`` so that TPU Mosaic
+    lowering (which rejects 64-bit types) works when ``jax_enable_x64`` is
+    active.  The casts are created inside the function body (not captured)
+    to satisfy JAX's "index_map must not capture constants" rule.
+    """
+    import jax.numpy as jnp  # pyrefly: ignore [import-error, missing-import]
+
+    # Pre-build the mapping so the returned lambda is a plain lookup.
+    mapping = dict(tiled_pairs)
+
+    if n_grid == 0 or (n_grid == 1 and not mapping):
+        return lambda _i: tuple(jnp.int32(0) for _ in range(buf_nd))
+
+    def index_map(*grid_args):
+        return tuple(
+            jnp.int32(grid_args[mapping[d]]) if d in mapping else jnp.int32(0)
+            for d in range(buf_nd)
+        )
+
+    return index_map
+
+
+def pallas_ensure_nonzero_rank(x: torch.Tensor) -> torch.Tensor:
+    if len(x.shape) == 0:
+        return x.reshape((1,))
+    return x
+
+
+def pallas_make_block_spec_non_tiled(shape: tuple[int, ...]) -> Any:
+    import jax.numpy as jnp  # pyrefly: ignore [import-error, missing-import]
+    from jax.experimental import (  # pyrefly: ignore [import-error, missing-import]
+        pallas as pl,
+    )
+
+    nonzero_rank_shape = shape if len(shape) > 0 else (1,)
+    return pl.BlockSpec(
+        nonzero_rank_shape,
+        lambda i: [jnp.int32(i)] * len(nonzero_rank_shape),
+    )
diff --git a/torch/_inductor/runtime/static_triton_launcher.py b/torch/_inductor/runtime/static_triton_launcher.py
index 5aaaf4fcf1974..4502d7a48f037 100644
--- a/torch/_inductor/runtime/static_triton_launcher.py
+++ b/torch/_inductor/runtime/static_triton_launcher.py
@@ -111,6 +111,7 @@ def needs_scratch_arg(scratch_name: str, param_name: str) -> bool:
         # pyrefly: ignore [missing-attribute]
         self.arg_tys = self.arg_ty_from_signature(kernel.src)
         self.function: int | None = None  # Loaded by load_kernel(on the parent process)
+        self.module: int | None = None  # Owns the HIP/CUDA module loaded for function
         num_ctas = 1
         if hasattr(kernel, "num_ctas"):
             num_ctas = kernel.num_ctas
@@ -141,13 +142,29 @@ def load_kernel(self, device: int) -> None:
 
         assert hasattr(self, "cubin_path")
         assert self.cubin_path is not None
-        (self.function, self.n_regs, self.n_spills) = self.C_impl._load_kernel(
-            self.cubin_path, self.name, self.shared, device
+        (self.module, self.function, self.n_regs, self.n_spills) = (
+            self.C_impl._load_kernel(self.cubin_path, self.name, self.shared, device)
         )
         # Don't need the cubin path anymore now that we've loaded
         self.cubin_path = None
         self.cubin_raw = None
 
+    def close(self) -> None:
+        mod = self.module
+        if mod is None:
+            return
+        # Clear Python-visible handles first so repeated cleanup is harmless even if
+        # the driver reports an error while unloading.
+        self.module = None
+        self.function = None
+        self.C_impl._unload_kernel(mod)
+
+    def __del__(self) -> None:
+        try:
+            self.close()
+        except Exception:
+            pass
+
     @staticmethod
     @functools.lru_cache
     def type_mappings() -> dict[str, str]:
@@ -224,6 +241,7 @@ def __getstate__(self) -> dict[str, Any]:
         # Remove objects that are no longer valid for pickling
         state = self.__dict__.copy()
         state["function"] = None
+        state["module"] = None
         # Cubin paths aren't consistent across processes, so we clear
         # and reload them.
         state["cubin_path"] = None
@@ -337,6 +355,26 @@ def __init__(self, kernel: CompiledKernel) -> None:
         self.cubin_raw = kernel.asm.get("zebin", None)
         super().__init__(kernel)
 
+    def load_kernel(self, device: int) -> None:
+        if self.function is not None:
+            return
+
+        assert hasattr(self, "cubin_path")
+        assert self.cubin_path is not None
+        # The XPU static launcher returns a PyCapsule for the loaded SYCL kernel,
+        # not a separate module/function pair like the CUDA/HIP launcher.
+        (self.function, self.n_regs, self.n_spills) = self.C_impl._load_kernel(
+            self.cubin_path, self.name, self.shared, device
+        )
+        self.module = None
+        self.cubin_path = None
+        self.cubin_raw = None
+
+    def close(self) -> None:
+        self.module = None
+        # Drop the PyCapsule reference so its destructor can release sycl::kernel.
+        self.function = None
+
 
 def statically_launched_kernel_by_device(
     kernel: CompiledKernel, device_type: str = "cuda"
diff --git a/torch/_inductor/runtime/triton_compat.py b/torch/_inductor/runtime/triton_compat.py
index 49ceacb50bc3d..cd3d53f139c03 100644
--- a/torch/_inductor/runtime/triton_compat.py
+++ b/torch/_inductor/runtime/triton_compat.py
@@ -91,6 +91,13 @@ def _triton_config_has(param_name: str) -> bool:
             triton_key,  # type: ignore[attr-defined,no-redef]
         )
 
+    try:
+        from triton.runtime.errors import IntelGPUError
+    except ImportError:
+
+        class IntelGPUError(Exception):  # type: ignore[no-redef]
+            pass
+
     builtins_use_semantic_kwarg = (
         "_semantic" in inspect.signature(triton.language.core.view).parameters
     )
@@ -106,6 +113,9 @@ class OutOfResources(Exception):  # type: ignore[no-redef]
     class PTXASError(Exception):  # type: ignore[no-redef]
         pass
 
+    class IntelGPUError(Exception):  # type: ignore[no-redef]
+        pass
+
     Config = object
     CompiledKernel = object
     KernelInterface = object
@@ -138,17 +148,6 @@ class JITFunction:  # type: ignore[no-redef]
     HAS_TRITON = False
 
 
-def cc_warp_size(cc: str | int) -> int:
-    if torch.version.hip:
-        cc_str = str(cc)
-        if "gfx10" in cc_str or "gfx11" in cc_str:
-            return 32
-        else:
-            return 64
-    else:
-        return 32
-
-
 try:
     autograd_profiler = torch.autograd.profiler
 except AttributeError:  # Compile workers only have a mock version of torch
@@ -163,6 +162,7 @@ class autograd_profiler:  # type: ignore[no-redef]
     "OutOfResources",
     "KernelInterface",
     "PTXASError",
+    "IntelGPUError",
     "ASTSource",
     "GPUTarget",
     "tl",
@@ -170,7 +170,6 @@ class autograd_profiler:  # type: ignore[no-redef]
     "libdevice",
     "math",
     "triton",
-    "cc_warp_size",
     "knobs",
     "triton_key",
 ]
diff --git a/torch/_inductor/runtime/triton_helpers.py b/torch/_inductor/runtime/triton_helpers.py
index 5cfa4786c0791..ed08cc38b037e 100644
--- a/torch/_inductor/runtime/triton_helpers.py
+++ b/torch/_inductor/runtime/triton_helpers.py
@@ -34,10 +34,27 @@ def set_driver_to_cpu():
     )
 
 
+def _is_backend_active(name, backend):
+    if backend.driver.is_active():
+        return True
+    # Triton may fail to detect the GPU in subprocess workers when using
+    # ctypes-based driver detection (triton-lang/triton#9578). Fall back
+    # to torch's own device checks which are more reliable in these environments.
+    if name == "nvidia":
+        import torch
+
+        return torch.cuda.is_available() and torch.version.hip is None
+    if name == "amd":
+        import torch
+
+        return torch.cuda.is_available() and torch.version.hip is not None
+    return False
+
+
 def set_driver_to_gpu():
     driver = triton.runtime.driver
     for name, backend in triton.backends.backends.items():
-        if backend.driver.is_active() and name != "cpu":
+        if _is_backend_active(name, backend) and name != "cpu":
             # After https://github.com/triton-lang/triton/commit/b844d519bc5e86edf00fe6b3c6c2d1badcd509a4,
             # `driver.active` can be of `LazyProxy` type and the sign of this - `_obj` attribute.
             if (
@@ -88,6 +105,23 @@ def remainder_integer(a, b):
     return tl.where((remainder != 0) & ((a < 0) != (b < 0)), remainder + b, remainder)
 
 
+@triton.jit
+def pow_integer(base, exponent):
+    # Triton has no exact integer pow primitive; use repeated squaring for
+    # nonnegative integer exponents so integral scalar pow does not round
+    # through libdevice.pow before casting back to int.
+    exponent_dtype: tl.constexpr = tl.core.get_int_dtype(
+        exponent.dtype.primitive_bitwidth, signed=False
+    )
+    exp = exponent.to(exponent_dtype)
+    result = tl.full(base.shape, 1, base.dtype)
+    for _ in tl.static_range(exponent_dtype.primitive_bitwidth):
+        result = tl.where((exp & 1) != 0, result * base, result)
+        exp = exp >> 1
+        base = base * base
+    return result
+
+
 @triton.jit
 def is_floating(x):
     return promote_to_tensor(x).dtype.is_floating()
@@ -249,6 +283,34 @@ def device_assert_then(cond, msg, r):
     return r
 
 
+@triton.jit
+def rand_eager_kernel(seed, offset_blocks, tid: tl.tensor, VEC: tl.constexpr):
+    inv = 1.0 / 4294967296.0
+    half = inv * 0.5
+
+    tid_u64 = tid.to(tl.uint64)
+
+    subseq = tid_u64 // VEC
+    which4 = (tid_u64 % VEC) // 4
+    lane = tid_u64 % 4
+
+    offblk = offset_blocks.to(tl.uint64) + which4
+
+    u0, u1, u2, u3 = tl.philox(
+        seed,
+        (offblk & 0xFFFFFFFF).to(tl.uint32),
+        ((offblk >> 32) & 0xFFFFFFFF).to(tl.uint32),
+        (subseq & 0xFFFFFFFF).to(tl.uint32),
+        ((subseq >> 32) & 0xFFFFFFFF).to(tl.uint32),
+    )
+
+    v01 = tl.where(lane == 0, u0, u1)
+    v23 = tl.where(lane == 2, u2, u3)
+    rand_int = tl.where((lane == 0) | (lane == 1), v01, v23)
+
+    return 1.0 - (rand_int.to(tl.float32) * inv + half)
+
+
 @triton.jit
 def randint64(seed, offset, low, high):
     r0, r1, _r2, _r3 = tl.randint4x(seed, offset)
@@ -342,6 +404,9 @@ def bucketize_binary_search(
         else:
             is_above = values > bucket_upper_bound
 
+        if is_floating(values):
+            is_above = is_above | (values != values)
+
         low = tl.where(is_above & mask, mid + 1, low)
         high = tl.where(is_above, high, mid)
 
@@ -679,7 +744,7 @@ def select_one(x, mask, dim, keep_dims=False):
     idtype = tl.core.get_int_dtype(x.dtype.primitive_bitwidth, signed=False)
     ix = x.to(idtype, bitcast=True)
     iy = tl.sum(ix * mask, dim, keep_dims=keep_dims)
-    return iy.to(x.dtype, bitcast=True)
+    return iy.to(idtype).to(x.dtype, bitcast=True)
 
 
 @triton.jit
@@ -759,3 +824,27 @@ def if_mask(mask: Any, val, *, _builder: object = None) -> tl.constexpr:
     if isinstance(mask, tl.constexpr) and mask.value is None:
         return tl.constexpr(None)
     return val
+
+
+@triton.jit
+def inline_asm_pack(x, pack: tl.constexpr):
+    """Ravel to 1D and pad (via join with zeros) so numel is divisible by pack."""
+    result = x.ravel()
+    # Only pad when the block size is smaller than pack. When block >= pack
+    # the numel is already divisible by pack (both are powers of 2).
+    n_pad: tl.constexpr = _log2(pack) - _log2(result.numel)
+    for _ in tl.static_range(n_pad):
+        result = tl.reshape(
+            tl.join(result, tl.zeros_like(result)), (result.shape[0] * 2,)
+        )
+    return result
+
+
+@triton.jit
+def inline_asm_unpack(x, orig, pack: tl.constexpr):
+    """Unpad and reshape back to orig's shape."""
+    result = x
+    n_pad: tl.constexpr = _log2(pack) - _log2(orig.numel)
+    for _ in tl.static_range(n_pad):
+        result, _ = tl.split(tl.reshape(result, (result.shape[0] // 2, 2)))
+    return tl.reshape(result, orig.shape)
diff --git a/torch/_inductor/runtime/triton_heuristics.py b/torch/_inductor/runtime/triton_heuristics.py
index 61bb640f5a072..8c2ee879c384d 100644
--- a/torch/_inductor/runtime/triton_heuristics.py
+++ b/torch/_inductor/runtime/triton_heuristics.py
@@ -4,6 +4,7 @@
 import builtins
 import copy
 import dataclasses
+import enum
 import functools
 import hashlib
 import inspect
@@ -18,7 +19,7 @@
 import threading
 import time
 from collections import namedtuple
-from typing import Any, Generic, Literal, TYPE_CHECKING, TypeVar, Union
+from typing import Any, Generic, Literal, TYPE_CHECKING, TypeVar
 
 import torch
 from torch._dynamo.utils import counters, set_feature_use
@@ -41,7 +42,6 @@
 from .benchmarking import benchmarker
 from .coordinate_descent_tuner import CoordescTuner
 from .hints import (
-    _NUM_THREADS_PER_WARP,
     AutotuneHint,
     DeviceProperties,
     HeuristicType,
@@ -74,11 +74,11 @@
 from .triton_compat import (
     ASTSource,
     autograd_profiler,
-    cc_warp_size,
     CompiledKernel,
     Config,
     GPUTarget,
     HAS_WARP_SPEC,
+    IntelGPUError,
     KernelInterface,
     knobs,
     OutOfResources,
@@ -88,6 +88,13 @@
 from .triton_helpers import get_constexprs
 
 
+class BenchmarkFailureReason(enum.Enum):
+    """Reasons why a triton config benchmark may return float('inf')."""
+
+    REGISTER_SPILLING = "register_spilling"
+    INVALID_CONFIG = "invalid_config"
+
+
 class InductorConfig(Config):
     """Inductor-specific Triton config with additional control flags"""
 
@@ -107,9 +114,9 @@ class NoTritonConfigsError(RuntimeError):
 
     LauncherType = Any
 
-_KernelType = Union[
-    CompiledKernel, StaticallyLaunchedCudaKernel, StaticallyLaunchedXpuKernel
-]
+_KernelType = (
+    CompiledKernel | StaticallyLaunchedCudaKernel | StaticallyLaunchedXpuKernel
+)
 _T = TypeVar("_T", bound=_KernelType)
 
 log = logging.getLogger(__name__)
@@ -168,6 +175,13 @@ def autotune_hints_to_configs(
     configs: list[Config] = []
     for hint in hints:
         if hint == AutotuneHint.ONE_ELEMENT_PER_THREAD:
+            if device_props.warp_size is None:
+                log.debug(
+                    "Skipping %s autotune hint because device %s does not report warp_size",
+                    AutotuneHint.ONE_ELEMENT_PER_THREAD,
+                    device_props.type,
+                )
+                continue
             if len(size_hints) == 1:
                 xyz_options = ((block_size // 4, None, None),)
             elif len(size_hints) == 2:
@@ -178,13 +192,13 @@ def autotune_hints_to_configs(
                     (1, block_size // 4, 1),
                     (1, 1, block_size // 4),
                 )
+            warp_size = device_props.warp_size_or_default
             configs.extend(
                 triton_config(
                     size_hints,
                     *xyz,
-                    num_elements_per_warp=(
-                        device_props.warp_size if device_props.warp_size else 32
-                    ),
+                    num_elements_per_warp=warp_size,
+                    warp_size=warp_size,
                 )
                 for xyz in xyz_options
             )
@@ -381,6 +395,7 @@ def __init__(
         self.compile_results: list[CompileResult[_KernelType]] = []
         self.launchers: list[LauncherType] = []
         self.lock = threading.Lock()
+        self.benchmark_failure_reasons: dict[Any, BenchmarkFailureReason] = {}
         if os.getenv("TRITON_CACHE_DIR") is None:
             os.environ["TRITON_CACHE_DIR"] = triton_cache_dir(
                 self.triton_meta.get("device", 0)
@@ -431,6 +446,37 @@ def __init__(
         # Mode for launch grid calculation
         self.grid_mode: Literal["python", "cpp"] = "python"
 
+    @staticmethod
+    def _close_compiled_kernel(kernel: Any) -> None:
+        if kernel is None:
+            return
+        close = getattr(kernel, "close", None)
+        if close is not None:
+            close()
+            return
+        module = getattr(kernel, "module", None)
+        if module is not None:
+            # Transitional fallback for Triton versions before CompiledKernel.close().
+            # Their __del__ unloads the module and clears kernel.module, so repeated
+            # calls from our idempotent benchmark cleanup paths are harmless.
+            delete = getattr(kernel, "__del__", None)
+            if delete is not None:
+                delete()
+
+    def release_benchmark_artifacts(self) -> None:
+        for launcher in self.launchers:
+            kernel = getattr(launcher, "__self__", None)
+            self._close_compiled_kernel(kernel)
+
+        for result in self.compile_results:
+            self._close_compiled_kernel(getattr(result, "kernel", None))
+
+        self.launchers = []
+        self.compile_results = []
+        self.benchmark_failure_reasons.clear()
+        self._cached_launcher = None
+        self._debug_call = None
+
     def is_statically_launchable(self):
         """
         Checks if every compiled kernel is statically launchable, which
@@ -458,20 +504,22 @@ def recheck_autotune_cache(
         )
         self.autotune_cache_info = autotune_cache_info
         # I.e. there was an autotune cache hit
-        if len(cached_configs) == 1 and len(configs) > 1:
+        if len(cached_configs) == 1:
             best_config = cached_configs[0]
+            found_by_coordesc = getattr(best_config, "found_by_coordesc", False)
             # Grab the best compiled config, if it's in the list of available ones
             best_config_hash = triton_config_to_hashable(best_config)
 
             for compile_result in self.compile_results:
                 if triton_config_to_hashable(compile_result.config) == best_config_hash:
+                    compile_result.config.found_by_coordesc = found_by_coordesc
                     self.compile_results = [compile_result]
                     return
 
             # If the best config isn't in our list of compile results,
             # it's likely because it was found by coordesc after the cache
             # already saved
-            if best_config.found_by_coordesc:
+            if found_by_coordesc:
                 with dynamo_timed("CachingAutotuner.slow_precompile_config"):
                     if self.fn.fn is None:
                         self.fn = reload_kernel_from_src().fn
@@ -520,7 +568,7 @@ def _precompile_worker(self):
         for c in self.configs:
             try:
                 compile_results.append(self._precompile_config(c))
-            except (OutOfResources, PTXASError) as e:
+            except (OutOfResources, PTXASError, IntelGPUError) as e:
                 exc = e
         if len(compile_results) == 0:
             raise NoTritonConfigsError(
@@ -544,12 +592,13 @@ def _dynamic_scale_rblock(self):
             and device_prop.major
             and (device_prop.major >= 8 or torch.version.hip)
             and device_prop.regs_per_multiprocessor is not None
+            and device_prop.warp_size is not None
         ):
             assert device_prop.regs_per_multiprocessor
             assert device_prop.max_threads_per_multi_processor
             assert device_prop.multi_processor_count
             seen_config_hashes: OrderedSet[Hashable] | None = None
-            warp_size = device_prop.warp_size or 32
+            warp_size = device_prop.warp_size
             for result in self.compile_results:
                 triton_config = result.config
                 compiled_binary = result.kernel
@@ -588,14 +637,14 @@ def _dynamic_scale_rblock(self):
                 nreg_per_warp = nreg * warp_size
                 nreg_per_block = nreg_per_warp * triton_config.num_warps
 
-                # Previously we set max_blocks_per_sm to 'max_threads_per_multi_processo / (32 * num_warps)'
+                # Previously we set max_blocks_per_sm to 'max_threads_per_multi_processor / (warp_size * num_warps)'
                 # The formula below is a tighter upper bound since we have the assumption that
                 #   nreg > device_prop.regs_per_multiprocessor // device_prop.max_threads_per_multi_processor
                 # due to the if condition above and:
                 #   regs_per_multiprocessor / nreg_per_block
-                #   = regs_per_multiprocessor / (nreg * 32 * num_warps)
-                #   < regs_per_multiprocessor / ((regs_per_multiprocessor / max_threads_per_multi_processor) * 32 * num_warps)
-                #   = max_threads_per_multi_processor / (32 * num_warps)
+                #   = regs_per_multiprocessor / (nreg * warp_size * num_warps)
+                #   < regs_per_multiprocessor / ((regs_per_multiprocessor / max_threads_per_multi_processor) * warp_size * num_warps)
+                #   = max_threads_per_multi_processor / (warp_size * num_warps)
                 # Using a tighter upper bound can reveal more optimization opportunities.
                 max_blocks_per_sm = max(
                     device_prop.regs_per_multiprocessor // nreg_per_block, 1
@@ -642,6 +691,20 @@ def _dynamic_scale_rblock(self):
 
             self._make_launchers()
 
+    def compile_by_disabling_pipelining(self, config):
+        # self.fn.fn is dropped by prepare_for_pickle() after the initial
+        # compile; reload it so triton.compile can access the source.
+        if self.fn.fn is None:
+            assert callable(self._reload_kernel)
+            self.fn = self._reload_kernel().fn
+        cfg = copy.deepcopy(config)
+        cfg.num_stages = 1
+        if "NUM_STAGES" in cfg.kwargs:
+            cfg.kwargs["NUM_STAGES"] = 1
+        result = self._precompile_config(cfg)
+        self.compile_results = [result]
+        return result.make_launcher()
+
     def _make_launchers(self):
         if len(self.launchers) == len(self.compile_results):
             return
@@ -649,22 +712,83 @@ def _make_launchers(self):
         from torch._dynamo.device_interface import DeviceGuard
 
         device_interface = self.get_device_interface()
-
-        # load binary to the correct device
-        with DeviceGuard(device_interface, self.triton_meta["device"]):
-            launchers = []
+        launchers = []
+        exc = None
+        try:
+            # Load each binary on the correct device.
+            with DeviceGuard(device_interface, self.triton_meta["device"]):
+                for result in self.compile_results:
+                    try:
+                        launchers.append(result.make_launcher())
+                    except (
+                        OutOfResources,
+                        PTXASError,
+                        torch.cuda.OutOfMemoryError,
+                        IntelGPUError,
+                    ) as e:
+                        exc = e
+                if len(launchers) == 0:
+                    result = self.compile_results[-1]
+                    config = result.config
+                    if (
+                        isinstance(exc, (OutOfResources, torch.cuda.OutOfMemoryError))
+                        and (
+                            config.num_stages > 1
+                            or config.kwargs.get("NUM_STAGES", 1) > 1
+                        )
+                        and self.inductor_meta.get("dynamic_disable_pipelining", True)
+                    ):
+                        self.launchers = [
+                            self.compile_by_disabling_pipelining(config)
+                        ]
+                        return
+                    raise RuntimeError(
+                        f"No valid triton configs. {type(exc).__name__}: {exc}"
+                    )
+            self.launchers = launchers
+        finally:
+            # Do not retain a failed launcher's exception and traceback. The
+            # traceback keeps its frame chain alive and can pin Triton's large
+            # benchmark cache-flush buffer until cyclic GC runs.
             exc = None
-            for result in self.compile_results:
-                try:
-                    launchers.append(result.make_launcher())
 
-                except (OutOfResources, PTXASError, torch.cuda.OutOfMemoryError) as e:
-                    exc = e
-        if len(launchers) == 0:
-            raise RuntimeError(f"No valid triton configs. {type(exc).__name__}: {exc}")
-        self.launchers = launchers
+    def _prune_compile_results_to_launcher(self, launcher: LauncherType) -> None:
+        if not self.compile_results:
+            return
+
+        launcher_config = launcher.config  # type: ignore[attr-defined]
+        for result in self.compile_results:
+            if result.config is launcher_config:
+                self.compile_results = [result]
+                return
+
+        launcher_config_hash = triton_config_to_hashable(launcher_config)
+        for result in self.compile_results:
+            if triton_config_to_hashable(result.config) == launcher_config_hash:
+                self.compile_results = [result]
+                return
+
+        raise AssertionError(
+            f"Autotuned launcher config does not match any compile result: {launcher_config}"
+        )
+
+    def _ensure_kernel_loaded(self) -> None:
+        """Reload the kernel in the parent process if needed.
+
+        When this autotuner was compiled in a worker subprocess and
+        unpickled into the parent, ``prepare_for_pickle`` cleared
+        ``self.fn.fn`` so the worker's JITFunction wouldn't follow
+        across the pickle. Any path that needs the live function in
+        the parent (coordesc, dynamic_scale_rblock, kernel_autotune
+        metrics) must reload first via the stashed ``_reload_kernel``
+        callback. No-op when the function is already present.
+        """
+        if self.fn.fn is None:
+            assert hasattr(self, "_reload_kernel")
+            assert callable(self._reload_kernel)
+            self.fn = self._reload_kernel().fn
 
-    def prepare_for_pickle(self) -> tuple[Any, Any, Any, Any, Any, Any]:
+    def prepare_for_pickle(self) -> tuple[Any, ...]:
         """Drop stuff from triton.JITFunction that does not pickle.
         This must be called after precompile so that these things are no longer needed.
         Returns a tuple of old values
@@ -676,18 +800,18 @@ def prepare_for_pickle(self) -> tuple[Any, Any, Any, Any, Any, Any]:
             self.fn.repr,
             self.launchers,
             getattr(self.fn, "_hash_lock", None),
+            self.benchmark_failure_reasons,
         )
         self.fn.fn = None
         self.fn.__globals__ = None
         self.fn.used_global_vals = None
         self.fn.repr = _ConstRepr(self.fn.repr(self.fn))
         self.launchers = []
+        self.benchmark_failure_reasons = {}
         self.fn._hash_lock = None
         return old_values
 
-    def restore_after_unpickle(
-        self, old_values: tuple[Any, Any, Any, Any, Any, Any] | None
-    ) -> None:
+    def restore_after_unpickle(self, old_values: tuple[Any, ...] | None) -> None:
         if old_values:
             (
                 self.fn.fn,
@@ -696,6 +820,7 @@ def restore_after_unpickle(
                 self.fn.repr,
                 self.launchers,
                 self.fn._hash_lock,
+                self.benchmark_failure_reasons,
             ) = old_values
         else:
             # even if we don't need/have specific values, we do need the
@@ -819,6 +944,8 @@ def _create_compile_options(
                 options["waves_per_eu"] = compile_meta["waves_per_eu"]
             if "matrix_instr_nonkdim" in compile_meta:
                 options["matrix_instr_nonkdim"] = compile_meta["matrix_instr_nonkdim"]
+            if "kpack" in compile_meta:
+                options["kpack"] = compile_meta["kpack"]
 
         if self.device_props.type == "xpu" and XPU_KERNEL_FORMAT == "zebin":
             options["generate_native_code"] = True
@@ -858,7 +985,7 @@ def _precompile_config(self, cfg: Config) -> CompileResult[_KernelType]:
         target = GPUTarget(
             compile_meta["device_type"],
             arch,
-            cc_warp_size(compile_meta["cc"]),
+            self.device_props.warp_size_or_default,
         )
 
         options = self._create_compile_options(cfg, compile_meta)
@@ -925,7 +1052,7 @@ def _precompile_config(self, cfg: Config) -> CompileResult[_KernelType]:
         return TritonCompileResult(binary, cfg, compile_meta, self.inductor_meta)
 
     def bench(self, launcher, *args, with_profiler=False, **kwargs):
-        """Measure the performance of a given launcher"""
+        """Measure the performance of a given launcher."""
         # we don't skip configs with spilled registers when auto-tuning custom
         # (user-written) Triton kernels, as (i) we don't have any knowledge or
         # control over the kernel code; (ii) there is empirical evidence that
@@ -942,6 +1069,9 @@ def bench(self, launcher, *args, with_profiler=False, **kwargs):
                 launcher.config,
                 launcher.n_spills,
             )
+            self.benchmark_failure_reasons[launcher] = (
+                BenchmarkFailureReason.REGISTER_SPILLING
+            )
             return float("inf")
 
         device_interface = self.get_device_interface()
@@ -970,7 +1100,14 @@ def kernel_call():
                             stream=stream,
                         )
                     except Exception:
-                        log.error("Failed during launch %s: ", kernel_name)
+                        log.error(
+                            "Failed during launch %s with config: %s (num_warps=%s, num_stages=%s, kwargs=%s)",
+                            kernel_name,
+                            launcher.config,
+                            launcher.config.num_warps,
+                            launcher.config.num_stages,
+                            launcher.config.kwargs,
+                        )
                         raise
 
             else:
@@ -981,7 +1118,14 @@ def kernel_call():
                         stream=stream,
                     )
                 except Exception:
-                    log.error("Failed during launch %s: ", kernel_name)
+                    log.error(
+                        "Failed during launch %s with config: %s (num_warps=%s, num_stages=%s, kwargs=%s)",
+                        kernel_name,
+                        launcher.config,
+                        launcher.config.num_warps,
+                        launcher.config.num_stages,
+                        launcher.config.kwargs,
+                    )
                     raise
             self.restore_args_from_cpu(cpu_copies)
 
@@ -996,11 +1140,19 @@ def kernel_call():
             if self.device_props.type == "cpu"
             else {"rep": 40, "is_vetted_benchmarking": True}
         )
-        return benchmarker.benchmark(
+        result = benchmarker.benchmark(
             fn=kernel_call,
             device=self.device_props.type,
             **benchmark_kwargs,  # type: ignore[arg-type]
         )
+        # benchmarker.benchmark() only returns float("inf") when catching an
+        # "invalid configuration" exception - all other exceptions are re-raised.
+        # Therefore, if result is inf here, it must be due to invalid config.
+        if result == float("inf"):
+            self.benchmark_failure_reasons[launcher] = (
+                BenchmarkFailureReason.INVALID_CONFIG
+            )
+        return result
 
     def copy_args_to_cpu_if_needed(self, *args, **kwargs):
         """
@@ -1015,13 +1167,22 @@ def copy_args_to_cpu_if_needed(self, *args, **kwargs):
 
         copies = {}
         try:
-            budget = torch.cuda.max_memory_allocated() - torch.cuda.memory_allocated()
+            if torch.accelerator.current_accelerator() is None:
+                # No initialized accelerator; skip memory-optimized path
+                return {}
+            budget = (
+                torch.accelerator.max_memory_allocated()
+                - torch.accelerator.memory_allocated()
+            )
         except RuntimeError:
             # Possibly a custom CUDA allocator, see https://github.com/pytorch/pytorch/issues/163257
             return {}
 
         def maybe_copy(name, arg):
-            if name in self.mutated_arg_names and arg.is_cuda:
+            if name in self.mutated_arg_names and arg.device.type in (
+                "cuda",
+                "xpu",
+            ):
                 nonlocal budget
                 assert isinstance(arg, torch.Tensor)
                 required_storage_length = compute_required_storage_length(
@@ -1117,6 +1278,41 @@ def prepare_arg(name, arg):
     def clone_args(self, *args, **kwargs) -> tuple[list[Any], dict[str, Any]]:
         return self.maybe_clone_args(OrderedSet(), *args, **kwargs)
 
+    @staticmethod
+    def _close_static_launcher(launcher: LauncherType) -> None:
+        if not getattr(launcher, "_is_static", False):
+            return
+        runner = getattr(launcher, "__globals__", {}).get("runner")
+        kernel = getattr(runner, "__self__", None)
+        close = getattr(kernel, "close", None)
+        if close is not None:
+            close()
+
+    def _release_static_launchers_except(self, keep_launcher: LauncherType) -> None:
+        for launcher in self.launchers:
+            if launcher is not keep_launcher:
+                self._close_static_launcher(launcher)
+        if not getattr(keep_launcher, "_is_static", False):
+            return
+        keep_hash = getattr(keep_launcher, "cache_hash", None)
+        if keep_hash is None:
+            return
+        keep_results = [
+            result
+            for result in self.compile_results
+            if isinstance(result, StaticTritonCompileResult)
+            and triton_hash_to_path_key(result.kernel.hash) == keep_hash
+        ]
+        if len(keep_results) == 1:
+            for result in self.compile_results:
+                if result is not keep_results[0] and isinstance(
+                    result, StaticTritonCompileResult
+                ):
+                    close = getattr(result.kernel, "close", None)
+                    if close is not None:
+                        close()
+            self.compile_results = keep_results
+
     def benchmark_all_configs(self, *args, **kwargs):
         with (
             dynamo_timed(
@@ -1135,10 +1331,21 @@ def benchmark_all_configs(self, *args, **kwargs):
             #     str(self.compile_id),
             # ),
         ):
-            timings = {
-                launcher: self.bench(launcher, *args, **kwargs)
-                for launcher in self.launchers
-            }
+            timings = {}
+            best_launcher = None
+            best_timing = float("inf")
+            for launcher in self.launchers:
+                timing = self.bench(launcher, *args, **kwargs)
+                timings[launcher] = timing
+                # Close losing static launchers eagerly so exhaustive autotuning
+                # keeps only the current winner and candidate modules loaded.
+                if best_launcher is None or timing < best_timing:
+                    if best_launcher is not None:
+                        self._close_static_launcher(best_launcher)
+                    best_launcher = launcher
+                    best_timing = timing
+                else:
+                    self._close_static_launcher(launcher)
 
             for k, v in timings.items():
                 self.coordesc_tuner.cache_benchmark_result(k.config, v)
@@ -1175,7 +1382,51 @@ def autotune_to_one_config(self, *args, **kwargs):
         start_time = time.time_ns()
         timings = self.benchmark_all_configs(*args, **kwargs)
         benchmark_time_taken_ns = time.time_ns() - start_time
-        self.launchers = [builtins.min(timings, key=timings.get)]
+
+        # Check if any configs failed (have inf timing) and log which one was selected
+        failed_launchers = [
+            launcher for launcher, timing in timings.items() if timing == float("inf")
+        ]
+        if failed_launchers:
+            valid_timings = [(k, v) for k, v in timings.items() if v != float("inf")]
+            if valid_timings:
+                best_launcher, best_time = min(valid_timings, key=lambda x: x[1])
+
+                # Count failures by reason
+                spill_count = sum(
+                    1
+                    for launcher in failed_launchers
+                    if self.benchmark_failure_reasons.get(launcher)
+                    == BenchmarkFailureReason.REGISTER_SPILLING
+                )
+                invalid_config_count = sum(
+                    1
+                    for launcher in failed_launchers
+                    if self.benchmark_failure_reasons.get(launcher)
+                    == BenchmarkFailureReason.INVALID_CONFIG
+                )
+
+                reason_parts = []
+                if spill_count > 0:
+                    reason_parts.append(f"{spill_count} register spilling")
+                if invalid_config_count > 0:
+                    reason_parts.append(f"{invalid_config_count} invalid config")
+                reason_str = ", ".join(reason_parts) if reason_parts else "unknown"
+
+                log.info(
+                    "Skipped %d/%d configs for %s (%s). Selected: %s (%.4f ms)",
+                    len(failed_launchers),
+                    len(timings),
+                    self.fn.__name__,
+                    reason_str,
+                    best_launcher.config,
+                    best_time,
+                )
+
+        best_launcher = builtins.min(timings, key=timings.get)
+        self._release_static_launchers_except(best_launcher)
+        self.launchers = [best_launcher]
+        self._prune_compile_results_to_launcher(best_launcher)
         self.autotune_time_taken_ns = (
             self.precompile_time_taken_ns + benchmark_time_taken_ns
         )
@@ -1192,13 +1443,166 @@ def autotune_to_one_config(self, *args, **kwargs):
             launcher.shared,
         )
 
+        TritonBundler.put_winner(launcher.cache_hash)
+
         if self.save_cache_hook:
             self.save_cache_hook(
                 launcher.config,
                 self.autotune_time_taken_ns,
+                found_by_coordesc=self.inductor_meta.get(
+                    "coordinate_descent_tuning", False
+                ),
                 triton_cache_hash=launcher.cache_hash,
             )
 
+    def _combo_sequential_autotune(self, launcher, *args, **kwargs):
+        """
+        Chain block-size decisions for combo kernels: tune one group at a time,
+        each step building on the previous winner.
+
+        Phase 1: Tune block sizes with warps/stages fixed from the base config.
+        Phase 2: Re-tune warps/stages with finalized block sizes.
+        """
+        combo_tuning_groups = self.inductor_meta.get("combo_tuning_groups")
+        if not combo_tuning_groups:
+            return launcher
+
+        if self.fn.fn is None:
+            assert hasattr(self, "_reload_kernel")
+            self.fn = self._reload_kernel().fn
+
+        best_config = launcher.config
+        current_kwargs = dict(best_config.kwargs)
+        base_num_warps = best_config.num_warps
+        base_num_stages = best_config.num_stages
+
+        start_time = time.time_ns()
+        best_time = self.bench(launcher, *args, **kwargs)
+        counters["inductor"]["combo_autotune_bench"] += 1
+        log.debug(
+            "  Phase 1 baseline: %s warps=%d time=%f",
+            dict(current_kwargs),
+            base_num_warps,
+            best_time,
+        )
+
+        # Phase 1: Tune block sizes per sub-kernel (largest first).
+        # warps/stages stay fixed at base config values.
+        for gi, group in enumerate(combo_tuning_groups):
+            member_indices = group["member_indices"]
+            cfgs = group["configs"]
+            skip_rblock = group["skip_rblock"]
+
+            if len(cfgs) <= 1:
+                log.debug("  Phase 1 group %d SK%s: 1 config, skip", gi, member_indices)
+                continue
+
+            log.debug(
+                "  Phase 1 group %d SK%s: trying %d configs, current_kwargs=%s",
+                gi,
+                member_indices,
+                len(cfgs),
+                dict(current_kwargs),
+            )
+            for ci, cfg in enumerate(cfgs):
+                trial_kwargs = dict(current_kwargs)
+                for idx in member_indices:
+                    for key, value in cfg.kwargs.items():
+                        if skip_rblock and key.startswith("R") and "BLOCK" in key:
+                            continue
+                        trial_kwargs[f"{key}_{idx}"] = value
+
+                if trial_kwargs == current_kwargs:
+                    log.debug("    cfg[%d] skip (same as current)", ci)
+                    continue
+
+                trial_config = triton.Config(
+                    trial_kwargs,
+                    num_warps=base_num_warps,
+                    num_stages=base_num_stages,
+                )
+
+                with self.lock:
+                    trial_launcher = self._precompile_config(
+                        trial_config
+                    ).make_launcher()
+                trial_time = self.bench(trial_launcher, *args, **kwargs)
+                counters["inductor"]["combo_autotune_bench"] += 1
+
+                improved = trial_time < best_time
+                log.debug(
+                    "    cfg[%d] trial=%s time=%f%s",
+                    ci,
+                    dict(trial_kwargs),
+                    trial_time,
+                    " (BETTER)" if improved else "",
+                )
+                if improved:
+                    best_time = trial_time
+                    launcher = trial_launcher
+                    current_kwargs = trial_kwargs
+
+            log.debug(
+                "  Phase 1 group %d winner: current_kwargs=%s",
+                gi,
+                dict(current_kwargs),
+            )
+
+        # Phase 2: Re-tune num_warps/num_stages with finalized block sizes.
+        # Block sizes are now optimal — find the best warp/stage pair for them.
+        warp_stage_candidates = self.inductor_meta.get("combo_warp_stage_candidates")
+        log.debug(
+            "  Phase 2: blocks=%s, trying %d warp/stage pairs",
+            dict(current_kwargs),
+            len(warp_stage_candidates),
+        )
+        best_warps = launcher.config.num_warps
+        best_stages = launcher.config.num_stages
+        for num_warps, num_stages in warp_stage_candidates:
+            if num_warps == best_warps and num_stages == best_stages:
+                log.debug(
+                    "    warps=%d stages=%d skip (same as current)",
+                    num_warps,
+                    num_stages,
+                )
+                continue
+
+            trial_config = triton.Config(
+                dict(current_kwargs),
+                num_warps=num_warps,
+                num_stages=num_stages,
+            )
+            with self.lock:
+                trial_launcher = self._precompile_config(trial_config).make_launcher()
+            trial_time = self.bench(trial_launcher, *args, **kwargs)
+            counters["inductor"]["combo_autotune_bench"] += 1
+
+            improved = trial_time < best_time
+            log.debug(
+                "    warps=%d stages=%d time=%f%s",
+                num_warps,
+                num_stages,
+                trial_time,
+                " (BETTER)" if improved else "",
+            )
+            if improved:
+                best_time = trial_time
+                launcher = trial_launcher
+                best_warps = num_warps
+                best_stages = num_stages
+
+        log.debug(
+            "Combo sequential autotune for %s: best config %s, time %f",
+            self.fn.__name__,
+            launcher.config,
+            best_time,
+        )
+        launcher.config.found_by_combo_autotune = True
+        self.autotune_time_taken_ns += time.time_ns() - start_time
+        if self.save_cache_hook:
+            self.save_cache_hook(launcher.config, self.autotune_time_taken_ns)
+        return launcher
+
     def save_gpu_kernel(self, stream, launcher):
         key = self.inductor_meta.get("kernel_name", None)  # unique kernel name
         assert key is not None, "kernel_name can not be None"
@@ -1228,14 +1632,6 @@ def save_gpu_kernel(self, stream, launcher):
             "global_scratch": launcher.global_scratch,
             "profile_scratch": launcher.profile_scratch,
         }
-        if self.device_props.type == "xpu":
-            # On the XPU backend, threads_per_warp is not always 32.
-            # For Intel GEMM Triton kernels, it can be 16.
-            # This information must be preserved so that the Cpp wrapper
-            # can launch the kernel with the correct configuration.
-            params["threads_per_warp"] = getattr(
-                launcher.bin.metadata, "threads_per_warp", 32
-            )
 
         from torch._inductor import config
         from torch._inductor.codecache import CudaKernelParamCache
@@ -1376,11 +1772,14 @@ def benchmark_one_config(config):
                 best_config
             ).make_launcher()
 
+        winner = config2launcher[best_config]
+        TritonBundler.put_winner(winner.cache_hash)
+
         fn_hash = generate_lookup_hash_from_source_code(
             str(self.size_hints), self.fn.src
         )
         log.debug("Function hash %s has best config %s", fn_hash, best_config)
-        return config2launcher[best_config]
+        return winner
 
     def get_profiler_kwargs(self, stream, launcher):
         kernel_kwargs_str = ",".join(
@@ -1447,6 +1846,23 @@ def alloc_fn(size: int, align: int, stream: int | None):
             if len(self.launchers) > 1:
                 self.autotune_to_one_config(*args, **kwargs)
 
+        if self.inductor_meta.get("combo_tuning_groups") and not getattr(
+            self.launchers[0].config, "found_by_combo_autotune", False
+        ):
+            with dynamo_timed(
+                "CachingAutotuner.combo_sequential_autotune",
+                log_pt2_compile_event=False,
+                metadata={"kernel_name": self.inductor_meta.get("kernel_name")},
+                dynamo_compile_column_us="runtime_triton_autotune_time_us",
+                compile_id=self.compile_id,
+                is_backward=self.is_backward,
+                log_waitcounter=True,
+                waitcounter_name_override="triton_autotuner",
+            ):
+                self.launchers = [
+                    self._combo_sequential_autotune(self.launchers[0], *args, **kwargs)
+                ]
+
         if not getattr(
             self.launchers[0].config, "found_by_coordesc", False
         ) and self.inductor_meta.get("coordinate_descent_tuning", False):
@@ -1455,6 +1871,11 @@ def alloc_fn(size: int, align: int, stream: int | None):
             ]
 
         (launcher,) = self.launchers
+        # Ensure the final launcher is marked as a winner for bundle filtering.
+        # For multi-config autotuning and coordesc, put_winner was already called
+        # (this is an idempotent set-add). For single-config kernels that skip
+        # autotuning entirely, this is the only call site that records the winner.
+        TritonBundler.put_winner(launcher.cache_hash)
         if launcher.store_cubin and (not benchmark_run or not self.cuda_kernel_saved):
             self.save_gpu_kernel(stream, launcher)
 
@@ -1753,7 +2174,7 @@ def check_can_launch() -> _KernelType:
             result = check_can_launch()
             return result
         except CannotStaticallyLaunchKernel as e:
-            log.info("Bypassing StaticallyLaunchedCudaKernel due to %s", str(e))  # noqa: G200
+            log.info("Bypassing StaticallyLaunchedCudaKernel due to %s", e)  # noqa: G200
             if torch._inductor.config.strict_static_triton_launcher:
                 raise e
             return None
@@ -2349,10 +2770,19 @@ def check_max_block(cfg: dict[str, int]):
             )
 
 
-def _num_warps(num_warps, max_num_warps=8, min_num_warps=2, register_intensive=False):
-    # On AMD GPU each warp has 64 lanes which is double the size on NV GPU,
-    # therefore using half the number of warps here correspondingly.
-    if torch.version.hip:
+def _num_warps(
+    num_warps,
+    max_num_warps=8,
+    min_num_warps=2,
+    register_intensive=False,
+    *,
+    warp_size: int = 32,
+):
+    # On wave64 AMD GPUs (CDNA / gfx9) each warp has 64 lanes, double NVIDIA
+    # and RDNA, so use half the number of warps to keep total thread count
+    # comparable. RDNA (wave32) follows the NVIDIA path.
+    if warp_size == 64:
+
         max_num_warps = (max_num_warps + 1) // 2
         min_num_warps = (min_num_warps + 1) // 2
     # persistent reduction is register intensive
@@ -2361,13 +2791,10 @@ def _num_warps(num_warps, max_num_warps=8, min_num_warps=2, register_intensive=F
     return next_power_of_2(min(max(num_warps, min_num_warps), max_num_warps))
 
 
-def _check_max_grid_x(size_hints, x, num_warps):
+def _check_max_grid_x(size_hints, x, num_warps, *, warp_size: int = 32):
     # Check if maxGridSize is exceeded - if so then must scale XBLOCK further
     max_grid_x = 2147483647
     max_block_x = TRITON_MAX_BLOCK["X"]
-    warp_size = (
-        64 if torch.version.hip else 32
-    )  # TODO: query warp size once #129663 is merged
     num_blocks = (size_hints["x"] + x - 1) // x
 
     if torch.version.hip:
@@ -2403,6 +2830,9 @@ def triton_config(
     num_warps=None,
     matrix_instr=None,
     waves_per_eu=None,
+    kpack=None,
+    *,
+    warp_size: int = 32,
 ) -> Config:
     """
     Construct a pointwise triton config with some adjustment heuristics
@@ -2462,7 +2892,9 @@ def triton_config(
     # Calculate num_warps if they are not hard passed to config
     if num_warps is None:
         num_warps = _num_warps(
-            conditional_product(x, y, z) // num_elements_per_warp, min_num_warps=1
+            conditional_product(x, y, z) // num_elements_per_warp,
+            min_num_warps=1,
+            warp_size=warp_size,
         )
     # we are going to arrive at 2 warps only if bs was too small due to
     # numel being too small. However to workaround some ptx bugs we still
@@ -2479,11 +2911,11 @@ def triton_config(
     # Increase x to satisfy min_elem_per_thread requirements.
     block_size = max(
         conditional_product(x, y, z),
-        min_elem_per_thread * _NUM_THREADS_PER_WARP * num_warps,
+        min_elem_per_thread * warp_size * num_warps,
     )
     x *= math.ceil(block_size / conditional_product(x, y, z))
 
-    x, _num_blocks = _check_max_grid_x(size_hints, x, num_warps)
+    x, _num_blocks = _check_max_grid_x(size_hints, x, num_warps, warp_size=warp_size)
     x = min(x, size_hints["x"])
 
     cfg = {"XBLOCK": x}
@@ -2500,6 +2932,8 @@ def triton_config(
             config.kwargs["matrix_instr_nonkdim"] = matrix_instr
         if waves_per_eu is not None:
             config.kwargs["waves_per_eu"] = waves_per_eu
+        if kpack is not None:
+            config.kwargs["kpack"] = kpack
 
     return config
 
@@ -2553,6 +2987,8 @@ def triton_config_reduction(
     dynamic_scale_rblock=True,
     reduction_hint=None,
     min_num_warps=None,
+    *,
+    warp_size: int = 32,
 ) -> Config:
     """
     Construct a reduction triton config with some adjustment heuristics
@@ -2594,10 +3030,13 @@ def total_numel() -> int:
         _num_warps_func = _num_warps
 
     num_warps = _num_warps_func(
-        num_warps, max_num_warps=max_num_warps, register_intensive=register_intensive
+        num_warps,
+        max_num_warps=max_num_warps,
+        register_intensive=register_intensive,
+        warp_size=warp_size,
     )
 
-    x, _num_blocks = _check_max_grid_x(size_hints, x, num_warps)
+    x, _num_blocks = _check_max_grid_x(size_hints, x, num_warps, warp_size=warp_size)
 
     for prefix in sorted(rnumels):
         while total_numel() > target:
@@ -2643,8 +3082,12 @@ def _handle_combo_kernel_per_subkernel_blocks(
     Handle per-subkernel config generation for combo kernels.
 
     Each sub-kernel gets its own block sizes (XBLOCK_0, XBLOCK_1, etc.) generated
-    using the same heuristics as standalone Triton kernels. The final config uses
-    the maximum num_warps and num_stages across all sub-kernels.
+    using the same heuristics as standalone Triton kernels.
+
+    Returns base configs that vary (num_warps, num_stages) with all blocks at
+    heuristic defaults. Stores per-subkernel candidate configs in
+    inductor_meta["combo_tuning_groups"] for sequential chained autotuning
+    in CachingAutotuner._combo_sequential_autotune().
 
     Returns:
         List of configs if combo kernel with combo_grid_meta and per-subkernel
@@ -2664,12 +3107,14 @@ def _handle_combo_kernel_per_subkernel_blocks(
     all_num_stages: list[int] = []
     unique_warp_stage_pairs: OrderedSet[tuple[int, int]] = OrderedSet()
 
+    combo_tuning_groups: list[dict[str, Any]] = []
+
     for i in range(num_kernels):
         subkernel_heuristic = combo_meta[f"heuristic_{i}"]
         size_hints_i = combo_meta[f"size_hints_{i}"]
 
         if subkernel_heuristic == "pointwise":
-            cfg = pointwise(
+            cfgs = pointwise(
                 size_hints_i,
                 triton_meta=triton_meta,
                 tile_hint=TileHint.SQUARE
@@ -2679,31 +3124,32 @@ def _handle_combo_kernel_per_subkernel_blocks(
                 min_elem_per_thread=min_elem_per_thread,
                 inductor_meta=inductor_meta_clean,
                 return_configs=True,
-            )[0]
+            )
             skip_rblock = False
         elif subkernel_heuristic == "reduction":
-            cfg = reduction(
+            cfgs = reduction(
                 size_hints_i,
-                reduction_hint=reduction_hint,
+                reduction_hint=ReductionHint[combo_meta[f"reduction_hint_{i}"]],
                 triton_meta=triton_meta,
                 filename=filename,
                 inductor_meta=inductor_meta_clean,
                 return_configs=True,
-            )[0]
+            )
             skip_rblock = False
         elif subkernel_heuristic == "persistent_reduction":
-            cfg = persistent_reduction(
+            cfgs = persistent_reduction(
                 size_hints_i,
-                reduction_hint=reduction_hint,
+                reduction_hint=ReductionHint[combo_meta[f"reduction_hint_{i}"]],
                 triton_meta=triton_meta,
                 filename=filename,
                 inductor_meta=inductor_meta_clean,
                 return_configs=True,
-            )[0]
+            )
             skip_rblock = True  # persistent reduction embeds RBLOCK in kernel body
         else:
             raise ValueError(f"Unknown heuristic: {subkernel_heuristic}")
 
+        cfg = cfgs[0]
         for key, value in cfg.kwargs.items():
             if skip_rblock and key.startswith("R") and "BLOCK" in key:
                 continue
@@ -2711,22 +3157,52 @@ def _handle_combo_kernel_per_subkernel_blocks(
 
         all_num_warps.append(cfg.num_warps)
         all_num_stages.append(cfg.num_stages)
-        unique_warp_stage_pairs.add((cfg.num_warps, cfg.num_stages))
+        for c in cfgs:
+            unique_warp_stage_pairs.add((c.num_warps, c.num_stages))
+
+        combo_tuning_groups.append(
+            {
+                "member_indices": [i],
+                "configs": cfgs,
+                "skip_rblock": skip_rblock,
+                "size_hints": size_hints_i,
+            }
+        )
 
     unique_warp_stage_pairs.add((max(all_num_warps), max(all_num_stages)))
 
+    # Largest sub-kernels tuned first — they dominate runtime and get most freedom
+    combo_tuning_groups.sort(
+        key=lambda g: -functools.reduce(operator.mul, g["size_hints"].values())
+    )
+    inductor_meta["combo_tuning_groups"] = combo_tuning_groups
+    # Candidates for num_warps/num_stages re-tuning after block sizes are finalized
+    inductor_meta["combo_warp_stage_candidates"] = list(unique_warp_stage_pairs)
+
+    # Single base config: max warps/stages, all blocks at heuristic defaults.
+    # Block sizes are tuned in _combo_sequential_autotune, then num_warps/num_stages
+    # are re-tuned at the end with finalized block sizes.
+    base_num_warps = max(all_num_warps)
+    base_num_stages = max(all_num_stages)
     return [
         triton.Config(
             combined_kwargs,
-            num_warps=num_warps,
-            num_stages=num_stages,
+            num_warps=base_num_warps,
+            num_stages=base_num_stages,
         )
-        for num_warps, num_stages in unique_warp_stage_pairs
     ]
 
 
 def triton_config_tiled_reduction(
-    size_hints, x, y, r, num_stages=1, register_intensive=False, waves_per_eu=None
+    size_hints,
+    x,
+    y,
+    r,
+    num_stages=1,
+    register_intensive=False,
+    waves_per_eu=None,
+    *,
+    warp_size: int = 32,
 ):
     """
     Construct a tile reduction triton config with some adjustment
@@ -2757,9 +3233,12 @@ def total_numel() -> int:
         y *= 2
 
     cfg = _get_config({"x": x, "y": y, **rnumels})
-    num_warps = _num_warps(total_numel() // 256, min_num_warps=1)
+    num_warps = _num_warps(total_numel() // 256, min_num_warps=1, warp_size=warp_size)
     num_warps = _num_warps(
-        num_warps, max_num_warps=16, register_intensive=register_intensive
+        num_warps,
+        max_num_warps=16,
+        register_intensive=register_intensive,
+        warp_size=warp_size,
     )
     check_config(cfg, xnumel=size_hints["x"], ynumel=size_hints["y"])
     check_max_block(cfg)
@@ -2859,8 +3338,9 @@ def pointwise(
         triton_meta["device"],
     )
 
+    warp_size = triton_meta["device"].warp_size_or_default
     triton_config_with_settings = functools.partial(
-        triton_config, min_elem_per_thread=min_elem_per_thread
+        triton_config, min_elem_per_thread=min_elem_per_thread, warp_size=warp_size
     )
 
     configs = None
@@ -2880,25 +3360,24 @@ def pointwise(
             ]
             # Additional configs appended for ROCm builds
             if torch.version.hip:
-                if inductor_meta.get("max_autotune_pointwise"):
-                    configs.extend(
-                        [
-                            triton_config_with_settings(
-                                size_hints, TRITON_MAX_BLOCK["X"], waves_per_eu=2
-                            ),
-                            triton_config_with_settings(
-                                size_hints,
-                                4096,  # wrt: better than the max_block for some kernel
-                            ),
-                            triton_config_with_settings(
-                                size_hints,
-                                2048,
-                                num_warps=8,
-                                num_stages=2,
-                                waves_per_eu=1,  # 20% improvement
-                            ),
-                        ]
-                    )
+                configs.extend(
+                    [
+                        triton_config_with_settings(
+                            size_hints, TRITON_MAX_BLOCK["X"], waves_per_eu=2
+                        ),
+                        triton_config_with_settings(
+                            size_hints,
+                            4096,  # wrt: better than the max_block for some kernel
+                        ),
+                        triton_config_with_settings(
+                            size_hints,
+                            2048,
+                            num_warps=8,
+                            num_stages=2,
+                            waves_per_eu=1,  # 20% improvement
+                        ),
+                    ]
+                )
                 if inductor_meta.get("atomic_add_found"):
                     configs.extend(
                         [
@@ -2970,7 +3449,9 @@ def pointwise(
                 )
     if len(size_hints) == 3:
         if not (
-            inductor_meta.get("max_autotune_pointwise") or torch.xpu.is_available()
+            inductor_meta.get("max_autotune")
+            or inductor_meta.get("max_autotune_pointwise")
+            or torch.xpu.is_available()
         ):
             configs = [triton_config_with_settings(size_hints, 16, 16, 16)]
         else:
@@ -3058,6 +3539,16 @@ def _config_helper(bmm=False, persistent=False):
 triton_native_persistent_bmm_configs = _config_helper(bmm=True, persistent=True)
 
 
+def _get_tiling_scores(
+    inductor_meta: dict[str, Any],
+    size_hints: dict[str, int],
+) -> dict[str, float]:
+    """
+    Retrieve the tiling scores, providing suitable defaults if they are missing.
+    """
+    return inductor_meta.get("tiling_scores") or dict.fromkeys(size_hints, 1)
+
+
 def _reduction_configs(
     *,
     size_hints: dict[str, int],
@@ -3076,10 +3567,14 @@ def _reduction_configs(
     )
 
     register_intensive = False
-    MAX_R0_BLOCK = 2048
     loads_and_red = inductor_meta.get("num_load", 0) + inductor_meta.get(
         "num_reduction", 0
     )
+
+    device_major = triton_meta["device"].major
+    warp_size = triton_meta["device"].warp_size_or_default
+    # Prefer smaller MAX_R0_BLOCK for Blackwell
+    MAX_R0_BLOCK = 1024 if device_major is not None and device_major >= 10 else 2048
     if size_hints["x"] >= 1024 and loads_and_red >= 10:
         # A heuristics to reduce R0_BLOCK if a kernel potentially need many registers.
         # Consider load and reduction since load need move data into registers and
@@ -3121,16 +3616,17 @@ def make_config(
     ):
         # For 3D case with tiling scores, create an adapted version
         if "y" in size_hints:
-            assert "tiling_scores" in inductor_meta
+            tiling_scores = _get_tiling_scores(inductor_meta, size_hints)
             return adapt_config_for_tiling(
                 size_hints,
-                inductor_meta["tiling_scores"],
+                tiling_scores,
                 x,
                 r,
                 num_warps=num_warps,
                 num_stages=num_stages,
                 register_intensive=register_intensive,
                 waves_per_eu=waves_per_eu,
+                warp_size=warp_size,
             )
         else:
             # For other cases, use the original function
@@ -3144,6 +3640,7 @@ def make_config(
                 waves_per_eu=waves_per_eu,
                 dynamic_scale_rblock=dynamic_scale_rblock,
                 reduction_hint=reduction_hint,
+                warp_size=warp_size,
             )
 
     def outer_config_opt():
@@ -3253,12 +3750,25 @@ def outer_config_opt():
     ]
 
     if torch.version.hip:
-        result_configs.extend(
-            [
-                make_config(1024, 8, num_warps=4, num_stages=1, waves_per_eu=2),
-                make_config(512, 8, num_warps=4, num_stages=1, waves_per_eu=1),
+        hip_configs = [
+            make_config(1024, 8, num_warps=4, num_stages=1, waves_per_eu=2),
+            make_config(512, 8, num_warps=4, num_stages=1, waves_per_eu=1),
+        ]
+        result_configs.extend(hip_configs)
+
+        # Filter ALL configs (not just HIP-specific ones) when a combo kernel
+        # has a persistent sub-kernel with a large hardcoded R0_BLOCK.  The
+        # persistent tile size (XBLOCK * max_persistent_rblock) causes
+        # pathological ROCm compilation times (e.g. 64 * 1024 = 64K elements
+        # → 60+ min triton.compile).  Use the same 4096-element threshold as
+        # _persistent_reduction_configs.
+        max_persistent_rblock = inductor_meta.get("max_persistent_rblock", 0)
+        if max_persistent_rblock > 0:
+            result_configs = [
+                c
+                for c in result_configs
+                if c.kwargs.get("XBLOCK", 0) * max_persistent_rblock <= 4096
             ]
-        )
 
     return result_configs
 
@@ -3282,6 +3792,7 @@ def match_target_block_product(
         # just assume even score with no minimum block size
         min_block_size = 1
         tiling_scores = dict.fromkeys(tiling_scores.keys(), target_block_product)
+        total_score = target_block_product * len(tiling_scores)
 
     # First, give each coalescing dimension at least min_block_size
     block_sizes = {}
@@ -3328,6 +3839,8 @@ def adapt_config_for_tiling(
     register_intensive=False,
     persistent_reduction=False,
     waves_per_eu=None,
+    *,
+    warp_size: int = 32,
 ) -> Config:
     """
     Create an adapted configuration based on tiling scores,
@@ -3347,6 +3860,7 @@ def adapt_config_for_tiling(
         num_stages=num_stages,
         register_intensive=register_intensive,
         waves_per_eu=waves_per_eu,
+        warp_size=warp_size,
     )
 
 
@@ -3541,9 +4055,7 @@ def cooperative_reduction(
     # the GPU, we want to create as many CTAs as possible, while keeping things
     # in powers of 2.
     target = last_power_of_2(triton_meta["device"].multi_processor_count)
-    split = max(1, min(target // xnumel, TRITON_MAX_RSPLIT))
-    assert rnumel >= split
-    assert split <= TRITON_MAX_RSPLIT
+    split = max(1, min((rnumel, target // xnumel, TRITON_MAX_RSPLIT)))
     if inductor_meta["persistent_reduction"]:
         configs = _persistent_reduction_configs(
             {"x": xnumel, "r0_": rnumel // split},
@@ -3583,6 +4095,7 @@ def _persistent_reduction_configs(
     rnumel = get_total_reduction_numel(size_hints)
 
     MAX_PERSISTENT_BLOCK_NUMEL = 4096
+    warp_size = triton_meta["device"].warp_size_or_default
 
     if triton_meta.get("native_matmul"):
         if len(size_hints) == 3:
@@ -3615,6 +4128,7 @@ def _persistent_reduction_configs(
                 rnumel,
                 register_intensive=True,
                 reduction_hint=reduction_hint,
+                warp_size=warp_size,
             )
             for xblock in xblock_vals
             if xblock == 1
@@ -3622,8 +4136,8 @@ def _persistent_reduction_configs(
         ]
     else:
         configs = []
-        assert "tiling_scores" in inductor_meta
-        x_y_scores = {dim: inductor_meta["tiling_scores"][dim] for dim in ("x", "y")}
+        tiling_scores = _get_tiling_scores(inductor_meta, size_hints)
+        x_y_scores = {dim: tiling_scores[dim] for dim in ("x", "y")}
         for target_block_size in xblock_vals:
             if target_block_size * rnumel > MAX_PERSISTENT_BLOCK_NUMEL:
                 continue
@@ -3633,7 +4147,11 @@ def _persistent_reduction_configs(
             )
             configs.append(
                 triton_config_tiled_reduction(
-                    size_hints, block_sizes["x"], block_sizes["y"], rnumel
+                    size_hints,
+                    block_sizes["x"],
+                    block_sizes["y"],
+                    rnumel,
+                    warp_size=warp_size,
                 )
             )
 
@@ -3642,6 +4160,7 @@ def _persistent_reduction_configs(
             size_hints,
             2 * (256 // rnumel) if rnumel <= 256 else 1,
             rnumel,
+            warp_size=warp_size,
         )
     ]
 
@@ -3677,6 +4196,7 @@ def _persistent_reduction_configs(
                         num_warps=num_warps,
                         min_num_warps=min_num_warps,
                         reduction_hint=reduction_hint,
+                        warp_size=warp_size,
                     )
                 ]
 
@@ -3776,7 +4296,10 @@ def persistent_reduction(
             # With large rnumel, we have higher chance of out-of-shared memory
             # To avoid adding too much autotuning overhead, we just constrain NUM_STAGES
             # if rnumel is large
-            MAX_NUM_STAGES = 2 if rnumel_hint > 8192 else 3
+            if inductor_meta.get("mix_order_reduction_allow_multi_stages", True):
+                MAX_NUM_STAGES = 2 if rnumel_hint > 8192 else 3
+            else:
+                MAX_NUM_STAGES = 1
             c.kwargs["NUM_STAGES"] = min(max(num_iters // 4, 1), MAX_NUM_STAGES)
 
             if rnumel_hint <= 1024:
@@ -3795,7 +4318,8 @@ def persistent_reduction(
                 # more warps for larger rows
                 new_configs.append(c)
 
-                if max_autotune_enabled and c.num_warps < 32:
+                max_warps_limit = 16 if torch.version.hip else 32
+                if max_autotune_enabled and c.num_warps < max_warps_limit:
                     newc = copy.deepcopy(c)
                     newc.num_warps *= 2
                     new_configs.append(newc)
@@ -4011,7 +4535,7 @@ class GridExpr:
     def __post_init__(self) -> None:
         assert self.mode in ("python", "cpp")
 
-    def generate(self, meta: dict[str, int]) -> None:
+    def generate(self, meta: dict[str, int], is_lazy: bool = False) -> None:
         raise NotImplementedError
 
     def ceildiv(self, numel: str | int, block: int | str | None) -> str | int:
@@ -4033,7 +4557,10 @@ def maximum(self, seq: list[int | str]) -> int | str:
             return items[0]
         if self.mode == "python":
             return f"max({', '.join(map(str, items))})"
-        return functools.reduce(lambda x, y: f"std::max({x}, {y})", items)
+        # Cast int constants to (long) to avoid type deduction errors with std::max
+        # when mixing long variables with int literals
+        cpp_items = [f"(long){x}" if isinstance(x, int) else str(x) for x in items]
+        return functools.reduce(lambda x, y: f"std::max({x}, {y})", cpp_items)
 
     def summation(self, seq: list[int | str]) -> int | str:
         """Codegen for sum function with constant folding, constants are represented as int"""
@@ -4089,34 +4616,68 @@ def eval_slow(self, meta: dict[str, int]) -> tuple[int, int, int]:
         exec(f"grid_2 = {self.z_grid}", scope)
         return scope["grid_0"], scope["grid_1"], scope["grid_2"]
 
+    def generate_lazy(self, kernel_name: str) -> None:
+        """
+        Creates a GridExpr for lazy compile, where config values are not known
+        at codegen time and are instead referenced by variable names.
+        """
+        meta: dict[str, Any] = {
+            "XBLOCK": f"{kernel_name}_result.xblock",
+            "YBLOCK": f"{kernel_name}_result.yblock",
+            "ZBLOCK": f"{kernel_name}_result.zblock",
+            "R0_BLOCK": f"{kernel_name}_result.r0block",
+            "RSPLIT": f"{kernel_name}_result.rsplit",
+            "RSPLIT_SIZE": f"{kernel_name}_result.rsplit_size",
+        }
+        # assertions are done based on real values, so we can skip here
+        self.generate(meta, is_lazy=True)
+
+    @classmethod
+    def from_meta_lazy(
+        cls,
+        inductor_meta: dict[str, Any] | None,
+        kernel_name: str,
+    ) -> GridExpr:
+        """Factory method for lazy compile mode."""
+        assert inductor_meta is not None, (
+            "inductor_meta must be specified for lazy compile"
+        )
+        grid_type = inductor_meta.get("grid_type", None)
+        assert grid_type is not None, "grid_type must be specified for lazy compile"
+        grid_cls = globals()[grid_type]
+        assert issubclass(grid_cls, GridExpr)
+        grid = grid_cls(inductor_meta=inductor_meta, mode="cpp")
+        grid.generate_lazy(kernel_name)
+        return grid
+
 
 class Grid1D(GridExpr):
-    def generate(self, meta: dict[str, int]) -> None:
+    def generate(self, meta: dict[str, int], is_lazy: bool = False) -> None:
         self.x_grid = self.ceildiv("xnumel", meta.get("XBLOCK"))
 
 
 class Grid2D(GridExpr):
-    def generate(self, meta: dict[str, int]) -> None:
+    def generate(self, meta: dict[str, int], is_lazy: bool = False) -> None:
         self.x_grid = self.ceildiv("xnumel", meta.get("XBLOCK"))
         self.y_grid = self.ceildiv("ynumel", meta.get("YBLOCK"))
 
 
 class Grid3D(GridExpr):
-    def generate(self, meta: dict[str, int]) -> None:
+    def generate(self, meta: dict[str, int], is_lazy: bool = False) -> None:
         self.x_grid = self.ceildiv("xnumel", meta.get("XBLOCK"))
         self.y_grid = self.ceildiv("ynumel", meta.get("YBLOCK"))
         self.z_grid = self.ceildiv("znumel", meta.get("ZBLOCK"))
 
 
 class BatchMatmulGrid3D(GridExpr):
-    def generate(self, meta: dict[str, int]) -> None:
+    def generate(self, meta: dict[str, int], is_lazy: bool = False) -> None:
         self.z_grid = self.ceildiv("xnumel", meta.get("XBLOCK"))
         self.y_grid = self.ceildiv("ynumel", meta.get("YBLOCK"))
         self.x_grid = self.ceildiv("znumel", meta.get("ZBLOCK"))
 
 
 class Grid2DWithYZOverflow(GridExpr):
-    def generate(self, meta: dict[str, int]) -> None:
+    def generate(self, meta: dict[str, int], is_lazy: bool = False) -> None:
         self.x_grid = self.ceildiv("xnumel", meta.get("XBLOCK"))
         self.prefix.extend(
             [
@@ -4128,29 +4689,35 @@ def generate(self, meta: dict[str, int]) -> None:
                 ),
             ]
         )
-        self.y_grid = self.ceildiv("y_grid_raw_", "y_grid_div_")
+        ceildiv_expr = self.ceildiv("y_grid_raw_", "y_grid_div_")
+        if self.mode == "python":
+            self.y_grid = f"(0 if y_grid_div_ == 0 else {ceildiv_expr})"
+        else:
+            self.y_grid = f"(y_grid_div_ == 0 ? 0 : {ceildiv_expr})"
         self.z_grid = "y_grid_div_"
 
 
 class MixOrderReductionGrid(GridExpr):
-    def generate(self, meta: dict[str, int]) -> None:
+    def generate(self, meta: dict[str, int], is_lazy: bool = False) -> None:
         split_size = meta.get("RSPLIT_SIZE")
         xblock = meta.get("XBLOCK")
-        assert split_size, "Missing RSPLIT_SIZE"
-        assert xblock, "Missing XBLOCK"
-        assert split_size % xblock == 0, f"{split_size=}, {xblock=}"
+        if not is_lazy:
+            assert split_size, "Missing RSPLIT_SIZE"
+            assert xblock, "Missing XBLOCK"
+            assert split_size % xblock == 0, f"{split_size=}, {xblock=}"
         self.x_grid = self.ceildiv("xnumel", split_size)
 
 
 class CooperativeReductionGrid(GridExpr):
-    def generate(self, meta: dict[str, int]) -> None:
+    def generate(self, meta: dict[str, int], is_lazy: bool = False) -> None:
         self.x_grid = str(meta["RSPLIT"])
         self.y_grid = self.ceildiv("xnumel", meta.get("XBLOCK"))
 
 
 class SplitScanGrid(GridExpr):
-    def generate(self, meta: dict[str, int]) -> None:
-        assert meta.get("XBLOCK", 1) == 1
+    def generate(self, meta: dict[str, int], is_lazy: bool = False) -> None:
+        if not is_lazy:
+            assert meta.get("XBLOCK", 1) == 1
         self.x_grid = self.ceildiv("r0_numel", meta.get("R0_BLOCK"))
         self.y_grid = "xnumel"
 
@@ -4165,12 +4732,12 @@ def setup_grid_as_args() -> dict[str, Any]:
             "extra_launcher_args": ["_grid_0", "_grid_1", "_grid_2"],
         }
 
-    def generate(self, meta: dict[str, int]) -> None:
+    def generate(self, meta: dict[str, int], is_lazy: bool = False) -> None:
         self.x_grid, self.y_grid, self.z_grid = self.inductor_meta["fixed_grid"]
 
 
 class PrecomputedGrid(GridExpr):
-    def generate(self, meta: dict[str, int]) -> None:
+    def generate(self, meta: dict[str, int], is_lazy: bool = False) -> None:
         for candidate in self.inductor_meta["precomputed_grids"]:
             if all(meta.get(k) == v for k, v in candidate["config"].items()):
                 self.x_grid, self.y_grid, self.z_grid = candidate[self.mode]
@@ -4181,7 +4748,7 @@ def generate(self, meta: dict[str, int]) -> None:
 
 
 class ComboKernelGrid(GridExpr):
-    def generate(self, meta: dict[str, int]):
+    def generate(self, meta: dict[str, int], is_lazy: bool = False) -> None:
         combo_meta = self.inductor_meta["combo_grid_meta"]
         if combo_meta["default_config"]:
             meta = {**combo_meta["default_config"], **meta}
@@ -4213,7 +4780,11 @@ def generate(self, meta: dict[str, int]):
                     ),
                 ]
             )
-            self.y_grid = self.ceildiv("y_grid_raw_", "y_grid_div_")
+            ceildiv_expr = self.ceildiv("y_grid_raw_", "y_grid_div_")
+            if self.mode == "python":
+                self.y_grid = f"(0 if y_grid_div_ == 0 else {ceildiv_expr})"
+            else:
+                self.y_grid = f"(y_grid_div_ == 0 ? 0 : {ceildiv_expr})"
             self.z_grid = "y_grid_div_"
 
     def combo_x_grid(
@@ -4244,7 +4815,7 @@ def combo_x_grid(
 class SequentialFlattenComboKernelGrid(GridExpr):
     """Flattened grid: (sum of x*y blocks, 1, 1) for per-subkernel with flattened dispatch."""
 
-    def generate(self, meta: dict[str, int]):
+    def generate(self, meta: dict[str, int], is_lazy: bool = False) -> None:
         combo_meta = self.inductor_meta["combo_grid_meta"]
         if combo_meta["default_config"]:
             meta = {**combo_meta["default_config"], **meta}
diff --git a/torch/_inductor/runtime/triton_lazy_compile.py b/torch/_inductor/runtime/triton_lazy_compile.py
new file mode 100644
index 0000000000000..35c28d9d4c37f
--- /dev/null
+++ b/torch/_inductor/runtime/triton_lazy_compile.py
@@ -0,0 +1,243 @@
+"""
+Lazy Triton kernel compilation for C++ wrapper.
+
+This module provides functionality to compile and autotune Triton kernels at runtime
+when the C++ wrapper is used without autotune_at_compile_time.
+
+The workflow is:
+1. At model initialization: Call start_kernel_compile() for all kernels to start
+   parallel compilation using multi-process async_compile
+2. At kernel execution time: Call run_triton_kernel_with_autotune() which waits
+   for the specific kernel to be ready, then runs it with autotuning
+"""
+
+from __future__ import annotations
+
+import dataclasses
+import logging
+import re
+from typing import Any
+
+from .triton_heuristics import CachingAutotuner
+
+
+log = logging.getLogger(__name__)
+
+
+@dataclasses.dataclass
+class TritonKernelCompileResult:
+    cubin_path: str
+    mangled_name: str
+    num_warps: int
+    shared_mem: int
+    xblock: int
+    yblock: int
+    zblock: int
+    r0block: int
+    rsplit: int
+    rsplit_size: int
+    config_index: int | None
+    global_scratch: int | None
+    profile_scratch: int | None
+
+
+_async_compile: Any = None
+
+
+def _get_async_compile() -> Any:
+    """Get or create the shared AsyncCompile instance."""
+    global _async_compile
+    if _async_compile is None:
+        from torch._inductor.async_compile import AsyncCompile
+
+        _async_compile = AsyncCompile()
+    return _async_compile
+
+
+def _wrap_tma_args(args: list[Any], kernel_fn: CachingAutotuner) -> list[Any]:
+    """Wrap tensor args with TMA descriptors where the signature requires them."""
+    signature = kernel_fn.triton_meta.get("signature", {})
+    sig_items = list(signature.items())
+
+    # Track args index separately from sig_items index since the signature
+    # may include constexpr entries that are not present in args.
+    tma_indices = []
+    arg_idx = 0
+    for name, sig_type in sig_items:
+        if isinstance(sig_type, str) and sig_type == "constexpr":
+            continue
+        if isinstance(sig_type, str) and sig_type == "nvTmaDesc":
+            raise RuntimeError(
+                f"nvTmaDesc (experimental TMA API) is not supported in lazy compile "
+                f"for arg '{name}'. Use the stable tensordesc API instead."
+            )
+        if isinstance(sig_type, str) and sig_type.startswith("tensordesc<"):
+            tma_indices.append((arg_idx, name, sig_type))
+        arg_idx += 1
+
+    if not tma_indices:
+        return args
+
+    from triton.tools.tensor_descriptor import TensorDescriptor
+
+    wrapped = list(args)
+    for arg_idx, name, sig_type in tma_indices:
+        if arg_idx >= len(wrapped):
+            raise RuntimeError(
+                f"TMA arg index {arg_idx} for '{name}' exceeds arg count {len(wrapped)}"
+            )
+        tensor = wrapped[arg_idx]
+        # Parse block_shape from tensordesc<dtype[dim0, dim1, ...]>
+        match = re.match(r"tensordesc<[^[]*\[([^\]]*)\]", sig_type)
+        if match:
+            block_shape = [int(x.strip()) for x in match.group(1).split(",")]
+            wrapped[arg_idx] = TensorDescriptor.from_tensor(tensor, block_shape)
+
+    return wrapped
+
+
+def start_kernel_compile(
+    pending_kernels: dict[str, Any], kernel_name: str, kernel_source: str
+) -> None:
+    """
+    This function is called from C++ at model initialization time for each kernel.
+    It starts the compilation in a background process but does NOT wait for it.
+    The actual kernel execution happens later in run_triton_kernel_with_autotune().
+
+    The pending_kernels dict is per-module, created in C++ and passed through
+    to avoid global state collisions across compiled modules.
+    """
+    if kernel_name in pending_kernels:
+        return
+
+    async_compile = _get_async_compile()  # noqa: F841 (used by eval below)
+
+    # Evaluate the kernel source to get the Future or CachingAutotuner
+    # The kernel_source is like: async_compile.triton('name', '''...''', ...)
+    kernel_obj = eval(kernel_source.strip())  # noqa: S307
+
+    pending_kernels[kernel_name] = kernel_obj
+
+
+def run_triton_kernel_with_autotune(
+    pending_kernels: dict[str, Any],
+    kernel_name: str,
+    stream: Any,
+    args: list[Any],
+) -> TritonKernelCompileResult:
+    """
+    Run a Triton kernel with full autotuning using actual tensor arguments.
+    """
+    from torch._inductor.codecache import CodeCacheFuture, CudaKernelParamCache
+    from torch._inductor.runtime.triton_heuristics import config_to_dict
+
+    if kernel_name not in pending_kernels:
+        raise RuntimeError(f"Kernel {kernel_name} not found in pending kernels.")
+    kernel_obj = pending_kernels[kernel_name]
+
+    if isinstance(kernel_obj, CodeCacheFuture):
+        kernel_fn = kernel_obj.result()
+    elif isinstance(kernel_obj, CachingAutotuner):
+        kernel_fn = kernel_obj
+    else:
+        raise RuntimeError(f"Unexpected kernel object type: {type(kernel_obj)}")
+
+    assert isinstance(kernel_fn, CachingAutotuner)
+
+    inductor_meta = kernel_fn.inductor_meta
+    inductor_meta["store_cubin"] = True
+
+    # For TMA kernels, wrap tensor args with TMA descriptors
+    args = _wrap_tma_args(args, kernel_fn)
+
+    # Run the kernel with the provided arguments
+    # This will trigger autotuning if there are multiple configs
+    kernel_fn.run(*args, stream=stream)
+    if not kernel_fn.launchers:
+        raise RuntimeError("Kernel run did not produce any launchers")
+    launcher = kernel_fn.launchers[0]
+
+    cached_params: dict[str, Any] | None = CudaKernelParamCache.get(kernel_name)
+    if cached_params is None:
+        raise RuntimeError(f"Failed to get cached params for kernel {kernel_name}")
+
+    from torch._inductor.codecache import get_cpp_wrapper_cubin_path_name
+
+    cubin_path_name = get_cpp_wrapper_cubin_path_name()
+    for key_name in (cubin_path_name, "mangled_name", "num_warps", "shared_mem"):
+        if key_name not in cached_params:
+            raise RuntimeError(
+                f"{key_name} not found in cached params for {kernel_name}"
+            )
+    cubin_path = cached_params[cubin_path_name]
+    mangled_name = cached_params["mangled_name"]
+    num_warps = cached_params["num_warps"]
+    shared_mem = cached_params["shared_mem"]
+
+    config = config_to_dict(launcher.config) if launcher.config else {}
+
+    # For combo/foreach kernels, the autotuned config may have empty kwargs
+    # (e.g., the foreach heuristic only tunes num_warps, not XBLOCK).
+    # In that case, use the default_config from combo_grid_meta
+    combo_grid_meta = inductor_meta.get("combo_grid_meta") if inductor_meta else None
+    default_config = combo_grid_meta.get("default_config") if combo_grid_meta else None
+    if default_config:
+        config = {**default_config, **config}
+
+    xblock = config.get("XBLOCK", 128)
+    yblock = config.get("YBLOCK", 1)
+    zblock = config.get("ZBLOCK", 1)
+    r0block = config.get("R0_BLOCK", 1)
+    rsplit = config.get("RSPLIT", 1)
+    rsplit_size = config.get("RSPLIT_SIZE", 1)
+
+    config_index = None
+    grid_type = inductor_meta.get("grid_type") if inductor_meta else None
+    if grid_type == "PrecomputedGrid" and inductor_meta:
+        # PrecomputedGrid selects one of precomputed_grids. We use config_index
+        # to remember which grid is chosen.
+        precomputed_grids = inductor_meta.get("precomputed_grids", [])
+        for idx, entry in enumerate(precomputed_grids):
+            entry_config = entry.get("config", {})
+            if all(config.get(k) == v for k, v in entry_config.items()):
+                config_index = idx
+                break
+
+    global_scratch: int | None = cached_params.get("global_scratch")
+    profile_scratch: int | None = cached_params.get("profile_scratch")
+
+    log.debug(
+        "Successfully autotuned Triton kernel: cubin_path=%s, mangled_name=%s, "
+        "num_warps=%d, shared_mem=%d, xblock=%d, yblock=%d, zblock=%d, r0block=%d, "
+        "rsplit=%d, rsplit_size=%d, config_index=%s, global_scratch=%s, profile_scratch=%s",
+        cubin_path,
+        mangled_name,
+        num_warps,
+        shared_mem,
+        xblock,
+        yblock,
+        zblock,
+        r0block,
+        rsplit,
+        rsplit_size,
+        config_index,
+        global_scratch,
+        profile_scratch,
+    )
+
+    result = TritonKernelCompileResult(
+        cubin_path=cubin_path,
+        mangled_name=mangled_name,
+        num_warps=num_warps,
+        shared_mem=shared_mem,
+        xblock=xblock,
+        yblock=yblock,
+        zblock=zblock,
+        r0block=r0block,
+        rsplit=rsplit,
+        rsplit_size=rsplit_size,
+        config_index=config_index,
+        global_scratch=global_scratch,
+        profile_scratch=profile_scratch,
+    )
+    return result
diff --git a/torch/_inductor/scheduler.py b/torch/_inductor/scheduler.py
index 6e1323ca942a3..1a59e9214ae7d 100644
--- a/torch/_inductor/scheduler.py
+++ b/torch/_inductor/scheduler.py
@@ -16,18 +16,22 @@
 import typing
 from collections import Counter, defaultdict
 from concurrent.futures import as_completed, Future
-from typing import Any, Generic, Optional, TYPE_CHECKING, TypeAlias, TypeVar, Union
+from typing import Any, Generic, TYPE_CHECKING, TypeAlias, TypeVar
 from typing_extensions import ParamSpec
 
 from torch.utils._ordered_set import OrderedSet
 
-from .ir import ComputedBuffer
+from .ir import ComputedBuffer, Pointwise
 
 
 if TYPE_CHECKING:
     from collections.abc import Callable, Iterator, Sequence
     from types import ModuleType
 
+    from torch._inductor.codegen.wrapper import EnterCudaStreamContextLine
+
+    from .codegen.wrapper import PythonWrapperCodegen
+
 import sympy
 
 import torch
@@ -38,6 +42,7 @@
 from torch._inductor.codecache import LambdaFuture, PyCodeCache
 from torch._inductor.ir import TritonTemplateCallerBase
 from torch._inductor.metrics import get_metric_table, is_metric_table_enabled
+from torch._inductor.stream_utils import get_stream_name
 from torch.fx.experimental.symbolic_shapes import free_symbols
 from torch.utils._sympy.symbol import free_symbol_is_type, symbol_is_type, SymT
 from torch.utils._triton import has_triton
@@ -62,7 +67,7 @@
 )
 from .loop_body import LoopBody
 from .memory import MemoryPlanningInfoForBuffer, MemoryPlanningInfoForNode
-from .runtime.hints import ReductionHint
+from .runtime.hints import DeviceProperties, ReductionHint
 from .runtime.runtime_utils import green_text, red_text
 from .sizevars import SimplifyIndexing
 from .utils import (
@@ -103,9 +108,9 @@
 
 @dataclasses.dataclass
 class FusionResult:
-    should_fuse: Optional[bool] = None
-    callable_fn: Optional[Callable[[], bool]] = None
-    future: Optional[LambdaFuture] = None
+    should_fuse: bool | None = None
+    callable_fn: Callable[[], bool] | None = None
+    future: LambdaFuture | None = None
 
     def __post_init__(self):
         assert (self.should_fuse is not None) ^ (self.callable_fn is not None), (
@@ -118,7 +123,7 @@ def fuse(cls, should_fuse: bool):
 
     @classmethod
     def from_callable(
-        cls, callable_fn: Callable[[], bool], future: Optional[LambdaFuture] = None
+        cls, callable_fn: Callable[[], bool], future: LambdaFuture | None = None
     ):
         return FusionResult(callable_fn=callable_fn, future=future)
 
@@ -128,7 +133,7 @@ class PendingFusion:
     callable_fn: Callable[[], bool]
     node1: BaseSchedulerNode
     node2: BaseSchedulerNode
-    future: Optional[LambdaFuture] = None
+    future: LambdaFuture | None = None
 
     def get_fusion_nodes(self) -> tuple[BaseSchedulerNode, BaseSchedulerNode]:
         return (self.node1, self.node2)
@@ -356,6 +361,9 @@ def can_fuse(cls, node1: BaseSchedulerNode, node2: BaseSchedulerNode) -> bool:
         if not V.graph.sizevars.statically_known_leq(ncol, 1024 * 16):
             return False
 
+        if MixOrderReduction.is_split_reduction(contiguous_node):
+            return False
+
         # Other reduction types like max/min is not supported yet.
         # There are no real use case as well.
         out = all(
@@ -419,7 +427,7 @@ def is_contiguous_load(cls, buf: str, parent_node: BaseSchedulerNode) -> bool:
 class SchedulerBuffer:
     scheduler: Scheduler
     node: ir.Buffer
-    defining_op: Optional[BaseSchedulerNode]
+    defining_op: BaseSchedulerNode | None
     users: list[NodeUser] = dataclasses.field(default_factory=list)
     mpi_buffer: MemoryPlanningInfoForBuffer = dataclasses.field(
         default_factory=MemoryPlanningInfoForBuffer
@@ -474,7 +482,7 @@ def allocate(self) -> None:
             hasattr(V.kernel, "args")
             and self.get_name() in V.kernel.inplace_update_buffers
         ):
-            input_buffer: Union[ir.DonatedBuffer, ir.Buffer]
+            input_buffer: ir.DonatedBuffer | ir.Buffer
             input_buffer_name = V.kernel.inplace_update_buffers[self.get_name()]
             if input_buffer_name in self.scheduler.name_to_donated_buffer:
                 input_buffer = self.scheduler.name_to_donated_buffer[
@@ -519,13 +527,13 @@ def get_mutations(self) -> Sequence[str]:
         assert self.node is not None
         return self.node.get_mutation_names()
 
-    def get_device(self) -> Optional[torch.device]:
+    def get_device(self) -> torch.device | None:
         return self.node.get_output_spec().get_device()
 
 
 @dataclasses.dataclass
 class SchedulerDonatedBuffer(SchedulerBuffer):
-    defining_op: Optional[BaseSchedulerNode] = None
+    defining_op: BaseSchedulerNode | None = None
 
 
 class BaseSchedulerNode:
@@ -537,14 +545,16 @@ class BaseSchedulerNode:
     # in `self.scheduler.nodes`, then for this FusedSchedulerNode, .min_order is 1 and .max_order is 3.
     # For non-"grouped" nodes (i.e. regular SchedulerNode),
     # .min_order = .max_order = X if this node is X-th node in `self.scheduler.nodes`.
+    min_input_distance: int
+    max_input_distance: int
     min_order: int
     max_order: int
     mpi_node: MemoryPlanningInfoForNode
     mutation_renames: dict[str, str]
-    node: Optional[ir.Operation] = None
+    node: ir.Operation | None = None
     outputs: list[SchedulerBuffer]
     outputs_by_name: dict[str, SchedulerBuffer]
-    override_estimated_runtime: Optional[float] = None
+    override_estimated_runtime: float | None = None
     read_writes: dependencies.ReadWrites
     unmet_dependencies: OrderedSet[Dep]
     written: bool = False
@@ -558,6 +568,8 @@ def __init__(self, scheduler: Scheduler) -> None:
     def _init_from_node(self, node: ir.Operation) -> None:
         self.node = node
         self.ancestors = OrderedSet()
+        self.min_input_distance = 0
+        self.max_input_distance = 0
         self.last_usage = OrderedSet[
             str
         ]()  # buffers that won't be used after this kernel
@@ -592,6 +604,8 @@ def debug_str(self) -> str:
 {name}.writes = {pformat(self.read_writes.writes)}
 {name}.unmet_dependencies = {pformat(self.unmet_dependencies)}
 {name}.met_dependencies = {pformat(self.read_writes.reads - self.unmet_dependencies)}
+{name}.min_input_distance = {self.min_input_distance}
+{name}.max_input_distance = {self.max_input_distance}
 {name}.outputs = [
         """
         )
@@ -773,7 +787,7 @@ def get_outputs(self) -> Sequence[SchedulerBuffer]:
     def get_output(self, buf_name: str) -> SchedulerBuffer:
         return self.outputs_by_name[buf_name]
 
-    def get_device(self) -> Optional[torch.device]:
+    def get_device(self) -> torch.device | None:
         assert self.node is not None
         return self.node.get_device()
 
@@ -878,11 +892,14 @@ def single_index_in_fused_node(buf_to_be_inplaced: SchedulerBuffer) -> bool:
                 or buf_node.get_inputs_that_alias_output()
                 or buf_node.get_mutation_names()
                 or buf.get_name() in V.graph.removed_buffers
+                # CommBufferLayout buffer must keep its P2P allocation.
+                # Do not allow in-place reuse into or from a P2P buffer.
+                or isinstance(buf_node.get_output_spec(), ir.CommBufferLayout)
             ):
                 continue
 
             for read in self.read_writes.reads:
-                input_buf: Optional[Union[SchedulerBuffer, SchedulerDonatedBuffer]]
+                input_buf: SchedulerBuffer | SchedulerDonatedBuffer | None
                 if read.name in self.scheduler.name_to_donated_buffer:
                     input_buf = self.scheduler.name_to_donated_buffer[read.name]
                 else:
@@ -899,8 +916,13 @@ def single_index_in_fused_node(buf_to_be_inplaced: SchedulerBuffer) -> bool:
                         for x in input_buf.users
                         if x.node.get_name() not in inconsequential_nodes
                     ]
+                    has_cross_stream_hazard = self.scheduler.has_cross_stream_hazard(
+                        read.name, self
+                    )
+
                     if (
-                        len(remaining_uses) == 1
+                        not has_cross_stream_hazard
+                        and len(remaining_uses) == 1
                         and remaining_uses[0].can_inplace
                         and remaining_uses[0].node is self
                         and input_buf.node is not None
@@ -910,6 +932,7 @@ def single_index_in_fused_node(buf_to_be_inplaced: SchedulerBuffer) -> bool:
                                 ir.NoneLayout,
                                 ir.MultiOutputLayout,
                                 ir.MutationLayoutSHOULDREMOVE,
+                                ir.CommBufferLayout,
                             ),
                         )
                         and not (
@@ -1103,7 +1126,7 @@ def is_materialized(buf: str, snodes: Sequence[BaseSchedulerNode]) -> bool:
 
         for buf_name in reads | writes:
             buf_accessed_elems = sum(node_numel for dep in buf_accesses[buf_name])
-            buf: Union[ir.Buffer, ir.TensorBox, ir.TorchBindObject]
+            buf: ir.Buffer | ir.TensorBox | ir.TorchBindObject
             if buf_name in V.graph.name_to_buffer:
                 buf = V.graph.name_to_buffer[buf_name]
             elif buf_name in V.graph.graph_inputs:
@@ -1112,7 +1135,7 @@ def is_materialized(buf: str, snodes: Sequence[BaseSchedulerNode]) -> bool:
                 continue
 
             def get_buf_bytes(
-                buf: Optional[Union[ir.Buffer, ir.TensorBox, ir.TorchBindObject]],
+                buf: ir.Buffer | ir.TensorBox | ir.TorchBindObject | None,
             ) -> int:
                 if not buf:
                     return 0
@@ -1269,7 +1292,7 @@ def _get_estimated_runtime(self) -> float:
         ms = ns / 1e6
         return ms
 
-    def get_template_node(self) -> Optional[ir.TemplateBuffer]:
+    def get_template_node(self) -> ir.TemplateBuffer | None:
         return None
 
     def get_template_node_or_throw(self) -> ir.TemplateBuffer:
@@ -1308,7 +1331,9 @@ def get_estimate_runtime_cache_key_from_snode(snode: BaseSchedulerNode) -> str:
     flat_args, flat_args_pytree_spec = pytree.tree_flatten((args, kwargs))
 
     def _is_tensor_ir(x) -> bool:  # type: ignore[no-untyped-def]
-        return isinstance(x, ir.IRNode) and not isinstance(x, ir.GeneratorState)
+        return isinstance(x, ir.IRNode) and not isinstance(
+            x, (ir.GeneratorState, ir.OpaqueObjectState)
+        )
 
     cache_key = str(
         (python_kernel_name,)
@@ -1317,7 +1342,7 @@ def _is_tensor_ir(x) -> bool:  # type: ignore[no-untyped-def]
     return cache_key
 
 
-def _get_mm_like_fn(snode: BaseSchedulerNode) -> Optional[Callable[[Any], Any]]:
+def _get_mm_like_fn(snode: BaseSchedulerNode) -> Callable[[Any], Any] | None:
     if not isinstance(snode, ExternKernelSchedulerNode):
         return None
     mms_fns = {
@@ -1333,7 +1358,7 @@ def _get_mm_like_fn(snode: BaseSchedulerNode) -> Optional[Callable[[Any], Any]]:
     return mms_fns[python_kernel_name]
 
 
-def maybe_estimate_runtime_benchmark(snode: BaseSchedulerNode) -> Optional[float]:
+def maybe_estimate_runtime_benchmark(snode: BaseSchedulerNode) -> float | None:
     bench_fn = None
     args_kwargs_fn = None
     if config.runtime_estimations_mms_benchmark:
@@ -1470,6 +1495,13 @@ def __init__(self, scheduler: Scheduler, node: ir.Operation) -> None:
         self._init_from_node(node)
         self.set_read_writes(node.get_read_writes())
 
+        if isinstance(node, ir.UserDefinedTritonKernel) and node.can_fuse_epilogue():
+            numel = math.prod(node.mutable_args[0].shape)
+            rnumel = 1
+            device = node.get_device_or_error()
+            # pyrefly: ignore [bad-assignment]
+            self.group = (device, (numel, rnumel))
+
     def debug_str_extra(self) -> str:
         return f"{self.get_name()}.node.kernel = {getattr(self.node, 'python_kernel_name', None)}"
 
@@ -1480,6 +1512,19 @@ def has_side_effects(self) -> bool:
         assert self.node is not None
         return hasattr(self.node, "has_side_effects") and self.node.has_side_effects()
 
+    def get_ranges(self) -> Sequence[Sequence[sympy.Expr]]:
+        if (
+            isinstance(self.node, ir.UserDefinedTritonKernel)
+            and self.node.can_fuse_epilogue()
+        ):
+            numel = math.prod(self.node.mutable_args[0].shape)
+            return ([numel], [])
+        return ([], [])
+
+    def codegen(self, wrapper: PythonWrapperCodegen) -> None:
+        assert isinstance(self.node, ir.ExternKernel)
+        return self.node.codegen(wrapper)
+
 
 class NopKernelSchedulerNode(BaseSchedulerNode):
     def __init__(self, scheduler: Scheduler, node: ir.Operation) -> None:
@@ -1500,7 +1545,7 @@ class SchedulerNode(BaseSchedulerNode):
     def __init__(
         self,
         scheduler: Scheduler,
-        node: Union[ir.ComputedBuffer, ir.TemplateBuffer],
+        node: ir.ComputedBuffer | ir.TemplateBuffer,
     ) -> None:
         super().__init__(scheduler)
         self._init_from_node(node)
@@ -1508,8 +1553,8 @@ def __init__(
 
     def _compute_attrs(
         self,
-        extra_indexing_constraints: Optional[tuple[dict[Any, Any], list[Any]]] = None,
-        recompute_sizes_body_func: Optional[Callable[_P, _T]] = None,
+        extra_indexing_constraints: tuple[dict[Any, Any], list[Any]] | None = None,
+        recompute_sizes_body_func: Callable[_P, _T] | None = None,
     ) -> None:
         assert isinstance(self.node, (ir.ComputedBuffer, ir.TemplateBuffer))
         self._sizes, body = self.node.simplify_and_reorder(
@@ -1541,13 +1586,20 @@ def _compute_attrs(
 
     def recompute_size_and_body(
         self,
-        extra_indexing_constraints: Optional[tuple[dict[Any, Any], list[Any]]] = None,
-        recompute_sizes_body_func: Optional[Callable[..., Any]] = None,
+        extra_indexing_constraints: tuple[dict[Any, Any], list[Any]] | None = None,
+        recompute_sizes_body_func: Callable[..., Any] | None = None,
     ) -> None:
+        fake_deps: OrderedSet[Dep] = OrderedSet(
+            dep for dep in self.read_writes.reads if isinstance(dep, (WeakDep, StarDep))
+        )
         self._compute_attrs(
             extra_indexing_constraints=extra_indexing_constraints,
             recompute_sizes_body_func=recompute_sizes_body_func,
         )
+        if fake_deps:
+            self.set_read_writes(
+                self.read_writes.with_read(fake_deps).rename(self.mutation_renames)
+            )
 
     def refresh_dependencies(
         self, normalize: bool, need_clear_tiling_cache: bool
@@ -1713,7 +1765,7 @@ def is_split_scan(self) -> bool:
     def is_template(self) -> bool:
         return isinstance(self.node, ir.TemplateBuffer)
 
-    def get_template_node(self) -> Optional[ir.TemplateBuffer]:
+    def get_template_node(self) -> ir.TemplateBuffer | None:
         return self.node if isinstance(self.node, ir.TemplateBuffer) else None
 
     def run(self, *index_vars: Sequence[sympy.Expr]) -> None:
@@ -1824,7 +1876,7 @@ def has_side_effects(self) -> bool:
 
 
 def refresh_group_node_dependencies(
-    group_snode: Union[FusedSchedulerNode, GroupedSchedulerNode],
+    group_snode: FusedSchedulerNode | GroupedSchedulerNode,
 ) -> None:
     snodes = group_snode.snodes
     group_snode.set_read_writes(
@@ -1842,7 +1894,7 @@ def refresh_group_node_dependencies(
 
 
 def init_group_node(
-    group_snode: Union[FusedSchedulerNode, GroupedSchedulerNode],
+    group_snode: FusedSchedulerNode | GroupedSchedulerNode,
     scheduler: Scheduler,
     snodes: list[BaseSchedulerNode],
 ) -> None:
@@ -1858,6 +1910,12 @@ def init_group_node(
 
     group_snode.min_order = min(x.min_order for x in group_snode.snodes)
     group_snode.max_order = max(x.max_order for x in group_snode.snodes)
+    group_snode.min_input_distance = min(
+        x.min_input_distance for x in group_snode.snodes
+    )
+    group_snode.max_input_distance = max(
+        x.max_input_distance for x in group_snode.snodes
+    )
     group_snode.outputs_by_name = {
         buf.get_name(): buf for buf in group_snode.get_outputs()
     }
@@ -1879,28 +1937,7 @@ def fuse(
         assert node1.scheduler is node2.scheduler
         assert isinstance(node1, (SchedulerNode, FusedSchedulerNode))
         if node1.is_template() and isinstance(node2, ExternKernelSchedulerNode):
-            # Fuse multi outputs template and its outputs
-            #   * Node1 has memorydep of MultiOutput in reads
-            #   * Node2 has StarDep of MultiOutput in writes
-            # Rewrite the Node2' StarDep to MemoryDep, because calculate score_fusion_memory
-            # of the template node and its epilogue requires the same type of dependencies
-            assert isinstance(node2.node, MultiOutput)
-            assert len(node2.read_writes.writes) == 1
-            assert isinstance(next(iter(node2.read_writes.writes)), StarDep)
-            name = next(iter(node2.read_writes.writes)).name
-            template_nodes = [node for node in node1.get_nodes() if node.is_template()]
-            assert len(template_nodes) == 1
-            template_node = template_nodes[0]
-            assert len(template_node.read_writes.writes) == 1
-            write = next(iter(template_node.read_writes.writes))
-            assert isinstance(write, MemoryDep)
-            node2.read_writes.writes = OrderedSet(
-                [
-                    MemoryDep(
-                        name, write.index, write.var_names, write.size, write.mode
-                    ),
-                ]
-            )
+            assert isinstance(node2.node, ir.MultiOutput)
         else:
             assert isinstance(node2, (SchedulerNode, FusedSchedulerNode))
         nodes = list(itertools.chain(node1.get_nodes(), node2.get_nodes()))
@@ -1947,7 +1984,8 @@ def reorder_loops_by_dep_pair(
             return False
         self_sizes = None
         for snode in self.snodes:
-            assert isinstance(snode, SchedulerNode)
+            if not isinstance(snode, SchedulerNode):
+                return False
             if self_sizes is not None and tuple(self_sizes) != tuple(snode._sizes[0]):
                 loop_ordering_log.debug(
                     "Can not reorder fused node due to different sizes"
@@ -2062,7 +2100,7 @@ def is_template(self) -> bool:
         return any(x.is_template() for x in self.snodes)
 
     @cache_on_self
-    def get_template_node(self) -> Optional[ir.TemplateBuffer]:
+    def get_template_node(self) -> ir.TemplateBuffer | None:
         for node in self.snodes:
             if node.is_template():
                 return node.get_template_node()
@@ -2097,6 +2135,8 @@ def debug_str(self) -> str:
 {name}.writes = {pformat(self.read_writes.writes)}
 {name}.unmet_dependencies = {pformat(self.unmet_dependencies)}
 {name}.met_dependencies = {pformat(self.read_writes.reads - self.unmet_dependencies)}
+{name}.min_input_distance = {self.min_input_distance}
+{name}.max_input_distance = {self.max_input_distance}
 {name}.outputs = [
             """
         )
@@ -2153,6 +2193,13 @@ def sub_node_can_fuse(
         if not self.scheduler.can_fuse(node1, node2, allow_mix_order_reduction=False):
             return False
 
+        # Since node1 is from the current mix order reduction, if node1 is
+        # contiguous, the fused node should also be contiguous.
+        if MixOrderReduction.is_contiguous_node(
+            node1
+        ) and not MixOrderReduction.is_contiguous_node(node2):
+            return False
+
         def _get_ancestors(nodes: tuple[BaseSchedulerNode, ...]) -> OrderedSet[str]:
             out = OrderedSet()
             return out.union(*(n.ancestors for n in nodes))
@@ -2206,6 +2253,67 @@ def fuse_with(self, other: BaseSchedulerNode):
                 return FusedMixOrderReductions(self.node1, fused_node)
 
 
+class FusedExternTritonKernelSchedulerNode(FusedSchedulerNode):
+    def __init__(
+        self,
+        scheduler: Scheduler,
+        kernel_node: ExternKernelSchedulerNode,
+        fused_epilogue: SchedulerNode,
+    ) -> None:
+        assert isinstance(kernel_node.node, ir.UserDefinedTritonKernel)
+        snodes = typing.cast(list[BaseSchedulerNode], [kernel_node, fused_epilogue])
+        super().__init__(scheduler, snodes)
+        self.kernel_node = kernel_node
+        self.fused_epilogue = fused_epilogue
+        self.min_order = self.kernel_node.min_order
+        self.outputs = fused_epilogue.outputs
+
+    @classmethod
+    def epilogue_fuse(
+        cls,
+        node1: ExternKernelSchedulerNode,
+        node2: SchedulerNode,
+    ) -> FusedSchedulerNode:
+        scheduler = node1.scheduler
+        # this unmet dependency is the buffer which is mutated
+        # after fusion, we don't need this buffer anymore,
+        # because the kernel directly writes to the output buffer of the epilogue
+        assert len(node1.unmet_dependencies) == 1
+        original_mutated_buffer = scheduler.name_to_buf[
+            next(iter(node1.unmet_dependencies)).name
+        ]
+        original_mutated_buffer.users.remove(NodeUser(node1))
+        return FusedExternTritonKernelSchedulerNode(scheduler, node1, node2)
+
+    def codegen(self, wrapper: PythonWrapperCodegen) -> None:
+        assert isinstance(self.fused_epilogue.node, ir.ComputedBuffer)
+        assert isinstance(self.kernel_node.node, ir.UserDefinedTritonKernel)
+        assert self.kernel_node.node.can_fuse_epilogue()
+        numel = math.prod(self.kernel_node.node.mutable_args[0].shape)
+        from torch._inductor.codegen.simd import SIMDScheduling
+
+        tiling, _ = SIMDScheduling.get_tiling_and_scores([self.fused_epilogue], numel)
+
+        from torch._inductor.codegen.simd_kernel_features import SIMDKernelFeatures
+
+        kernel_features = SIMDKernelFeatures([self.fused_epilogue], numel)
+
+        from torch._inductor.codegen.triton import FusedUserDefinedTritonKernel
+
+        fused_user_kernel = FusedUserDefinedTritonKernel(tiling, kernel_features, self)
+        new_kernel_src = fused_user_kernel.codegen()
+
+        return self.kernel_node.node.codegen_with_epilogue_fusion(
+            wrapper, (self.fused_epilogue.node, new_kernel_src)
+        )
+
+    def is_extern(self) -> bool:
+        return True
+
+    def get_ranges(self) -> Sequence[Sequence[sympy.Expr]]:
+        return self.kernel_node.get_ranges()
+
+
 class ForeachKernelSchedulerNode(FusedSchedulerNode):
     """
     This is a schedular node that consists of a set of scheduler nodes that
@@ -2214,7 +2322,7 @@ class ForeachKernelSchedulerNode(FusedSchedulerNode):
 
     def get_consumer_subnode_for(
         self, producer: BaseSchedulerNode
-    ) -> Optional[BaseSchedulerNode]:
+    ) -> BaseSchedulerNode | None:
         for buf in producer.get_outputs():
             if buf.get_name() in self.read_to_node:
                 return self.read_to_node[buf.get_name()]
@@ -2223,7 +2331,7 @@ def get_consumer_subnode_for(
 
     def get_producer_subnode_for(
         self, consumer: BaseSchedulerNode
-    ) -> Optional[BaseSchedulerNode]:
+    ) -> BaseSchedulerNode | None:
         producers = OrderedSet[BaseSchedulerNode]()
         for rd in consumer.read_writes.reads:
             if rd.name not in self.scheduler.name_to_buf:
@@ -2357,8 +2465,8 @@ def __init__(
         scheduler: Scheduler,
         snodes: list[BaseSchedulerNode],
         use_custom_partition_algo: bool,
-        prev_node_1: Optional[BaseSchedulerNode] = None,
-        prev_node_2: Optional[BaseSchedulerNode] = None,
+        prev_node_1: BaseSchedulerNode | None = None,
+        prev_node_2: BaseSchedulerNode | None = None,
         enable_autotune: bool = False,
     ) -> None:
         self.read_to_node = {}
@@ -2398,6 +2506,12 @@ def __init__(
 
             self.min_order = min([prev_node_1.min_order, prev_node_2.min_order])
             self.max_order = max([prev_node_1.max_order, prev_node_2.max_order])
+            self.min_input_distance = min(
+                prev_node_1.min_input_distance, prev_node_2.min_input_distance
+            )
+            self.max_input_distance = max(
+                prev_node_1.max_input_distance, prev_node_2.max_input_distance
+            )
 
             if prev_node_1.is_foreach():
                 assert isinstance(prev_node_1, ForeachKernelSchedulerNode)
@@ -2512,7 +2626,7 @@ def _default_group_nodes_for_combo_kernels(
         )
         for nodes in sorted_nodes:
             # Group nodes by device first to avoid mixed-device fusion
-            device_groups: dict[Optional[torch.device], list[BaseSchedulerNode]] = (
+            device_groups: dict[torch.device | None, list[BaseSchedulerNode]] = (
                 defaultdict(list)
             )
             for node in nodes:
@@ -2525,14 +2639,20 @@ def _default_group_nodes_for_combo_kernels(
                     continue
                 device_groups[device].append(node)
 
-            # Chunk each device group separately
+            # Sub-group by stream to avoid mixing nodes across stream
+            # boundaries.  When multi-stream scheduling is inactive every
+            # node maps to DEFAULT_STREAM_IDX so this is a no-op.
             for device_nodes in device_groups.values():
-                grouped_nodes.extend(
-                    [
-                        device_nodes[i : i + max_num_nodes]
-                        for i in range(0, len(device_nodes), max_num_nodes)
-                    ]
-                )
+                stream_groups: dict[int, list[BaseSchedulerNode]] = defaultdict(list)
+                for node in device_nodes:
+                    stream_groups[scheduler.node_to_stream.get(node, 0)].append(node)
+                for stream_nodes in stream_groups.values():
+                    grouped_nodes.extend(
+                        [
+                            stream_nodes[i : i + max_num_nodes]
+                            for i in range(0, len(stream_nodes), max_num_nodes)
+                        ]
+                    )
         return grouped_nodes
 
     group_algorithm_for_combo_kernels: Callable[
@@ -2676,7 +2796,7 @@ def estimate_flops(self) -> int | None:
     def get_nodes(self) -> Sequence[BaseSchedulerNode]:
         return self.snodes
 
-    def get_device(self) -> Optional[torch.device]:
+    def get_device(self) -> torch.device | None:
         return self.snodes[0].get_device() if self.snodes else None
 
     @classmethod
@@ -2760,11 +2880,11 @@ def _replace_operation_buffer(
 
 
 def _estimate_fused_epilogue_runtime(node1, node2, epilogue_runtime) -> float:
+    template_write_bytes = node1.get_write_buffer_sizes()
+    epilogue_read_bytes = node2.get_read_buffer_sizes()
+    extra_bytes = epilogue_read_bytes - template_write_bytes
     # If no extra memory read by epilogue, assume epilogue is free
     # if extra memory is read by epilogue, add to minimum choice
-    total_read_bytes = node2.get_read_buffer_sizes()
-    template_write_bytes = node1.get_write_buffer_sizes()
-    extra_bytes = total_read_bytes - template_write_bytes
     extra_bytes_ratio = extra_bytes / template_write_bytes
 
     # Smoothly approaches 1 as extra_bytes_ratio increases
@@ -2772,9 +2892,70 @@ def _estimate_fused_epilogue_runtime(node1, node2, epilogue_runtime) -> float:
     return extra_memory_ratio * epilogue_runtime
 
 
+def _occupancy_before_and_after_fusion(
+    unfused_n_regs: int,
+    fused_n_regs: int,
+    fused_n_spills: int,
+    num_warps: int,
+    device_props: DeviceProperties,
+) -> tuple[int, int]:
+    if fused_n_spills >= 8:
+        return 0, -1
+
+    # # Need device info to calculate occupancy
+    regs_per_sm = device_props.regs_per_multiprocessor
+    warp_size = device_props.warp_size
+    if regs_per_sm is None or warp_size is None:
+        return 1, 1  # Can't calculate, allow fusion
+
+    if not num_warps:
+        raise AssertionError("expected num_warps to be truthy")
+    threads_per_block = num_warps * warp_size
+
+    regs_per_block_unfused = unfused_n_regs * threads_per_block
+    regs_per_block_fused = fused_n_regs * threads_per_block
+
+    blocks_unfused = regs_per_sm // regs_per_block_unfused
+    blocks_fused = regs_per_sm // regs_per_block_fused
+
+    return blocks_unfused, blocks_fused
+
+
+def _fuse_epilogue(
+    ms1: float,
+    ms2: float,
+    unfused_n_regs: int,
+    fused_n_regs: int,
+    fused_n_spills: int,
+    num_warps: int,
+    device_props: DeviceProperties,
+) -> bool:
+    """
+    Determine whether to fuse an epilogue into a GEMM template.
+    """
+    MIN_ACCEPTED_OCCUPANCY = 4
+    REGRESSED_OCCUPANCY_RATIO = 0.5
+
+    # Check occupancy impact
+    blocks_unfused, blocks_fused = _occupancy_before_and_after_fusion(
+        unfused_n_regs, fused_n_regs, fused_n_spills, num_warps, device_props
+    )
+
+    epilogue_dominated_with_sufficient_occupancy = ms2 > 2 * ms1 and blocks_fused > 1
+
+    # fuse if no major register spills
+    # Occupancy can decrease but if memory bound/epilogue dominated
+    # optimistically fuse
+    return blocks_fused != -1 and (
+        blocks_fused >= MIN_ACCEPTED_OCCUPANCY
+        or blocks_fused / blocks_unfused > REGRESSED_OCCUPANCY_RATIO
+        or epilogue_dominated_with_sufficient_occupancy
+    )
+
+
 @dataclasses.dataclass
 class NodeUser:
-    node: Union[BaseSchedulerNode, OutputNode]
+    node: BaseSchedulerNode | OutputNode
     can_inplace: bool = False
 
     # A weak user must be scheduled after a given node, but doesn't actually
@@ -2848,12 +3029,36 @@ def get_scheduler_node_symbol_uses(
     return free_symbol_uses
 
 
+def _is_epilogue_fusion_enabled(template_node: BaseSchedulerNode) -> bool:
+    """Check per-template flag, fall back to global config."""
+    tb = template_node.get_template_node()
+    if tb is not None and tb.allow_epilogue_fusion is not None:
+        return tb.allow_epilogue_fusion
+    return config.epilogue_fusion
+
+
+def _is_prologue_fusion_enabled(template_node: BaseSchedulerNode) -> bool:
+    """Check per-template flag, fall back to global config."""
+    tb = template_node.get_template_node()
+    if tb is not None and tb.allow_prologue_fusion is not None:
+        return tb.allow_prologue_fusion
+    return config.prologue_fusion
+
+
 def is_epilogue_fusion(node1: BaseSchedulerNode, node2: BaseSchedulerNode):
-    return node1.is_template() and config.epilogue_fusion and not node2.is_template()
+    return (
+        node1.is_template()
+        and not node2.is_template()
+        and _is_epilogue_fusion_enabled(node1)
+    )
 
 
 def is_prologue_fusion(node1: BaseSchedulerNode, node2: BaseSchedulerNode):
-    return node2.is_template() and config.prologue_fusion and not node1.is_template()
+    return (
+        node2.is_template()
+        and not node1.is_template()
+        and _is_prologue_fusion_enabled(node2)
+    )
 
 
 def is_template_fusion(node1: BaseSchedulerNode, node2: BaseSchedulerNode):
@@ -2890,8 +3095,8 @@ def _init(self, nodes: list[ir.Operation]) -> None:
             ]
         )
         self.nodes = [self.create_scheduler_node(n) for n in nodes]
-        self.previous_node: Optional[BaseSchedulerNode] = None
-        self.current_node: Optional[BaseSchedulerNode] = None
+        self.previous_node: BaseSchedulerNode | None = None
+        self.current_node: BaseSchedulerNode | None = None
         self.update_zero_dim_cpu_tensor()
         # some new constants could have been created above
         self.available_buffer_names.update(V.graph.constants.keys())
@@ -2899,7 +3104,7 @@ def _init(self, nodes: list[ir.Operation]) -> None:
             node.prune_deps()
 
         # See [Note: Graph Partition Device Contexts]
-        self.default_device_context: Optional[torch.device] = None
+        self.default_device_context: torch.device | None = None
 
         self.name_to_donated_buffer: dict[str, SchedulerDonatedBuffer] = (
             self.get_donated_buffers()
@@ -2946,6 +3151,7 @@ def _init(self, nodes: list[ir.Operation]) -> None:
         self.dead_node_elimination()
         self.name_to_fused_node = {n.get_name(): n for n in self.nodes}
         self.compute_ancestors()
+        self.compute_input_distances()
 
         # pyrefly: ignore [bad-assignment]
         metrics.ir_nodes_pre_fusion += len(self.nodes)
@@ -2965,10 +3171,28 @@ def _init(self, nodes: list[ir.Operation]) -> None:
             distributed_autotune.schedule(self)
             self.compute_ancestors()
 
+        # Stream assignments must be populated BEFORE fusion
+        # to prevent fusing nodes across stream boundaries
+        self.node_to_stream: dict[BaseSchedulerNode, int] = {}
+        self.buff_to_stream: dict[str, int] = {}
+        self._multi_stream_nodes: bool = False
+        # Maps stream_idx → user_object_index for retrieving user stream objects
+        self.stream_idx_to_user_obj_idx: dict[int, int] = {}
+        self._populate_stream_assignments()
+
         self.nodes = self.fuse_nodes(self.nodes)
         if config._post_fusion_custom_pass is not None:
             self.nodes = config._post_fusion_custom_pass(self.nodes)
 
+        if any(
+            isinstance(node, FusedExternTritonKernelSchedulerNode)
+            for node in self.nodes
+        ):
+            # if a user triton kernel has been epilogue-fused,
+            # there is likely an opportunity to prune an NopKernel
+            # (which is originally used to generate the buffer which the triton kernel writes to)
+            self.dead_node_elimination()
+
         self.merge_loops()
         self.finalize_multi_template_buffers()
         if (
@@ -2984,6 +3208,10 @@ def _init(self, nodes: list[ir.Operation]) -> None:
             ):
                 self.create_combo_kernel_nodes(num_ck_nodes=None)
 
+        # torch.cond can contain arbitrary subgraphs, which can contain collectives
+        # reordering these can cause a nccl hang
+        self._enforce_conditional_ordering()
+
         # Peak memory pass and overlap pass must run last, otherwise
         # other reordering passes could undo their effects.
         if config.reorder_for_peak_memory:
@@ -3077,6 +3305,9 @@ def _init(self, nodes: list[ir.Operation]) -> None:
         # for debug attribution
         self.origin_to_index: dict[torch.fx.Node, int] = {}
 
+        # The only source of which stream context we are currently in during the codegen phase.
+        self._current_stream_ctx: EnterCudaStreamContextLine | None = None
+
         get_metric_table("graph_stats").add_row(
             lambda: {
                 "graph_id": self.post_grad_graph_id,
@@ -3100,12 +3331,87 @@ def get_donated_buffers(self) -> dict[str, SchedulerDonatedBuffer]:
                 )
         return name_to_donated_buf
 
+    def _populate_stream_assignments(self) -> None:
+        """Populate node_to_stream and buff_to_stream from IR node stream_idx.
+
+        Reads the stream_idx field set on IR nodes during lowering to determine
+        which stream each scheduler node should run on. This field is propagated
+        from 'custom.stream' FX node metadata via IRNode.current_stream_idx().
+        """
+        from .stream_constants import DEFAULT_STREAM_IDX
+
+        # Map user_object_index to stream index (1-indexed for side streams)
+        user_obj_to_stream_idx: dict[int, int] = {}
+        stream_idx_counter = itertools.count(1)  # 0 is reserved for default stream
+
+        for node in self.nodes:
+            stream_idx = DEFAULT_STREAM_IDX
+
+            if node.node is not None:
+                user_obj_idx = node.node.get_stream_idx()
+                if user_obj_idx is not None:
+                    if user_obj_idx not in user_obj_to_stream_idx:
+                        new_stream_idx = next(stream_idx_counter)
+                        user_obj_to_stream_idx[user_obj_idx] = new_stream_idx
+                        self.stream_idx_to_user_obj_idx[new_stream_idx] = user_obj_idx
+                    stream_idx = user_obj_to_stream_idx[user_obj_idx]
+
+            self.node_to_stream[node] = stream_idx
+
+            # Also populate buff_to_stream for all buffers produced by this node.
+            # Mutation renames are resolved at lookup time via get_buf_stream.
+            for buf in node.get_buffer_names():
+                self.buff_to_stream[buf] = stream_idx
+
+        # Propagate a device to device-less nodes (e.g. record_event,
+        # wait_event) so they naturally enter the device guard in the
+        # main codegen loop instead of requiring special-case handling.
+        if any(s != DEFAULT_STREAM_IDX for s in self.node_to_stream.values()):
+            device = next(
+                (n.get_device() for n in self.nodes if n.get_device() is not None), None
+            )
+            if device is not None:
+                for node in self.nodes:
+                    ir_node = node.node
+                    if (
+                        node.get_device() is None
+                        and isinstance(ir_node, ir.Buffer)
+                        and isinstance(ir_node.layout, ir.NoneLayout)
+                    ):
+                        # pyrefly: ignore [bad-assignment]
+                        ir_node.layout = ir.NoneLayout(device=device)
+
+        # Check if we have any nodes on non-default streams
+        self._multi_stream_nodes = any(
+            stream_idx != DEFAULT_STREAM_IDX
+            for stream_idx in self.node_to_stream.values()
+        )
+
+    def _has_multi_stream_nodes(self) -> bool:
+        """Check if any nodes are assigned to non-default streams."""
+        return self._multi_stream_nodes
+
+    def get_buf_stream(self, buf_name: str) -> int:
+        """Return the stream index for a buffer, resolving mutation renames."""
+        real = self.mutation_renames.get(buf_name, buf_name)
+        return self.buff_to_stream.get(real, self.buff_to_stream.get(buf_name, 0))
+
+    def has_cross_stream_hazard(self, buf_name: str, node: BaseSchedulerNode) -> bool:
+        """True if buf_name was produced on a different stream than node.
+
+        Resolves mutation renames so that mutated buffers inherit the
+        stream of their original definition.
+        """
+        if not self._has_multi_stream_nodes():
+            return False
+        return self.get_buf_stream(buf_name) != self.node_to_stream.get(node, 0)
+
     @property
-    def current_device(self) -> Optional[torch.device]:
+    def current_device(self) -> torch.device | None:
         return V.graph.current_device
 
     @current_device.setter
-    def current_device(self, device: Optional[torch.device]) -> None:
+    def current_device(self, device: torch.device | None) -> None:
         V.graph.current_device = device
 
     def debug_draw_graph(self) -> None:
@@ -3188,8 +3494,8 @@ class DedupList(Generic[_T]):
 
             def __init__(
                 self,
-                items: Optional[list[_T]] = None,
-                membership: Optional[OrderedSet[_T]] = None,
+                items: list[_T] | None = None,
+                membership: OrderedSet[_T] | None = None,
             ) -> None:
                 self.items = items or []
                 self.membership = membership or OrderedSet()
@@ -3253,7 +3559,7 @@ def rename(n: str) -> str:
         def add_user(
             # pyrefly: ignore [not-a-type]
             used_by_name: str,
-            user_node: Union[BaseSchedulerNode, OutputNode],
+            user_node: BaseSchedulerNode | OutputNode,
             can_inplace: bool = False,
             is_weak: bool = False,
         ) -> None:
@@ -3261,8 +3567,8 @@ def add_user(
                 NodeUser(user_node, can_inplace, is_weak)
             )
 
-        # pyrefly: ignore [not-a-type]
-        unbacked_symbol_to_origin_node: dict[sympy.Symbol, Optional[str]] = {}
+        # pyrefly: ignore [not-a-type, unsupported-operation]
+        unbacked_symbol_to_origin_node: dict[sympy.Symbol, str | None] = {}
 
         # NB: None means that the dependency is on an input.  Don't actually
         # generate a dependency because if we do, Inductor will start trying
@@ -3563,7 +3869,7 @@ def can_eliminate_user(user: NodeUser) -> bool:
         for node in self.nodes:
             node.prune_weak_deps()
 
-    def mode_requires_synchronization(self, mode: Optional[str]) -> bool:
+    def mode_requires_synchronization(self, mode: str | None) -> bool:
         """Check if store mode requires cross-thread synchronization."""
         return mode is not None  # Currently all non-None modes need sync
 
@@ -3594,6 +3900,17 @@ def visit(n: BaseSchedulerNode) -> None:
             visit(node)
         return result
 
+    def _enforce_conditional_ordering(self) -> None:
+        conditional_nodes = [
+            n for n in self.nodes if isinstance(n.node, ir.Conditional)
+        ]
+        for i in range(1, len(conditional_nodes)):
+            mutating_buf = next(iter(conditional_nodes[i].get_buffer_names()))
+            prev_buf = next(iter(conditional_nodes[i - 1].get_buffer_names()))
+            conditional_nodes[i].add_fake_dep(
+                WeakDep(prev_buf, mutating_buf=mutating_buf, is_fake=True)
+            )
+
     def _get_unmet_dep_nodes(self, snode: BaseSchedulerNode) -> list[BaseSchedulerNode]:
         unmet_deps: OrderedSet[str] = OrderedSet()
         if isinstance(
@@ -3660,6 +3977,36 @@ def compute_ancestors(self) -> None:
             node.min_order = order
             node.max_order = order
 
+    def compute_input_distances(self) -> None:
+        """
+        Populate each node's min/max_input_distance with the depth from graph
+        inputs, measured as dependency hops before fusion. Nodes whose
+        dependencies are all satisfied by graph inputs/constants have depth 0.
+        """
+        name_to_min_distance: dict[str, int] = {}
+        name_to_max_distance: dict[str, int] = {}
+        for node in self.nodes:
+            if not node.unmet_dependencies:
+                min_dist = 0
+                max_dist = 0
+            else:
+                dep_min_dists = [
+                    name_to_min_distance[self.name_to_buf[dep.name].defining_op_name()]
+                    + 1
+                    for dep in node.unmet_dependencies
+                ]
+                dep_max_dists = [
+                    name_to_max_distance[self.name_to_buf[dep.name].defining_op_name()]
+                    + 1
+                    for dep in node.unmet_dependencies
+                ]
+                min_dist = min(dep_min_dists)
+                max_dist = max(dep_max_dists)
+            name_to_min_distance[node.get_name()] = min_dist
+            name_to_max_distance[node.get_name()] = max_dist
+            node.min_input_distance = min_dist
+            node.max_input_distance = max_dist
+
     def merge_loops(self) -> None:
         if not config.loop_ordering_after_fusion:
             return
@@ -3755,7 +4102,7 @@ def generate_kernel_code_from_nodes(
         self,
         nodes: Sequence[BaseSchedulerNode],
         benchmark_kernel: bool,
-        hint_override: Optional[int] = None,
+        hint_override: int | None = None,
     ) -> str:
         """
         Benchmark fused list of nodes and return the execution time
@@ -3797,7 +4144,12 @@ def _has_layout_conflict_for_template(
         for inp in multi_node.inputs:
             # pyrefly: ignore [missing-attribute]
             inp_name = inp.get_name()
-            if not getattr(inp, "layout", None) or inp_name not in constraints:
+            # View has its own fixed layout that is not constrained
+            if (
+                not getattr(inp, "layout", None)
+                or inp_name not in constraints
+                or isinstance(inp, ir.ReinterpretView)
+            ):
                 continue
 
             layout = inp.layout
@@ -3876,7 +4228,7 @@ def finalize_multi_template_buffers(self) -> None:
                 ):
                     # pyrefly: ignore [unbound-name]
                     if config.multi_kernel_hints:
-                        callers: dict[Optional[int], TritonTemplateCallerBase] = {}
+                        callers: dict[int | None, TritonTemplateCallerBase] = {}
                         callers[None] = min_node_unfused
 
                         # pyrefly: ignore [unbound-name]
@@ -3959,8 +4311,8 @@ def _any_atomic_add(self, node_list: Sequence[BaseSchedulerNode]) -> bool:
         )
 
     def compile_kernel(
-        self, nodes: Sequence[BaseSchedulerNode], hint_override: Optional[int] = None
-    ) -> tuple[Optional[LambdaFuture], ModuleType]:
+        self, nodes: Sequence[BaseSchedulerNode], hint_override: int | None = None
+    ) -> tuple[LambdaFuture | None, ModuleType]:
         src_code = self.generate_kernel_code_from_nodes(
             nodes, benchmark_kernel=True, hint_override=hint_override
         )
@@ -4055,9 +4407,9 @@ def log_fusion(ms_fused: float, ms1: float, ms2: float) -> None:
                 return FusionResult.fuse(False)
 
             hint_override_best_fusion_choice: dict[
-                Optional[int], TritonTemplateCallerBase
+                int | None, TritonTemplateCallerBase
             ] = {}
-            future_choices: list[tuple[Any, Optional[LambdaFuture], ModuleType]] = []
+            future_choices: list[tuple[Any, LambdaFuture | None, ModuleType]] = []
             for hint_override in config.multi_kernel_hints:
                 choice_timings = multi_node.choice_timings(hint_override)
                 for choice, _ in sorted(choice_timings.items(), key=lambda x: x[1]):
@@ -4076,7 +4428,7 @@ def log_fusion(ms_fused: float, ms1: float, ms2: float) -> None:
                         )
 
                 min_ms_fused = float("inf")
-                ms_fused_choice: Optional[TritonTemplateCallerBase] = None
+                ms_fused_choice: TritonTemplateCallerBase | None = None
                 new_timings = {}
                 for choice, future, mod_fused in future_choices:
                     try:
@@ -4087,7 +4439,7 @@ def log_fusion(ms_fused: float, ms1: float, ms2: float) -> None:
                             fusion_log.debug(  # noqa: G200
                                 "Exception in compiling %s: %s",
                                 "prologue" if not epilogue_fusion else "epilogue",
-                                str(e),
+                                e,
                             )
                         continue
                     with multi_node.swap_as_triton_caller(choice):
@@ -4142,7 +4494,9 @@ def log_fusion(ms_fused: float, ms1: float, ms2: float) -> None:
                 ms2_fused = _estimate_fused_epilogue_runtime(node1, node2, ms2)
 
             # Start compiling choices in parallel
-            future_choices: list[tuple[Any, Optional[LambdaFuture], ModuleType]] = []
+            from torch._inductor.codegen.simd import CantSplit
+
+            future_choices: list[tuple[Any, LambdaFuture | None, ModuleType]] = []
             triton_choices = 0
             for choice, unfused_time in choice_timings_iter:
                 if not isinstance(choice, TritonTemplateCallerBase):
@@ -4168,9 +4522,14 @@ def log_fusion(ms_fused: float, ms1: float, ms2: float) -> None:
                     break
 
                 with multi_node.swap_as_triton_caller(choice):
-                    future_choices.append(
-                        (choice, *self.compile_kernel(node_list_fused))
-                    )
+                    try:
+                        future_choices.append(
+                            (choice, *self.compile_kernel(node_list_fused))
+                        )
+                    except CantSplit:
+                        # Epilogue node ranges may be incompatible with the
+                        # template kernel's tiling groups — skip this choice.
+                        continue
 
             if len(future_choices) == 0:
                 return FusionResult.fuse(False)
@@ -4186,6 +4545,13 @@ def benchmark_when_ready() -> bool:
                     choice_timings = multi_node.choice_timings()
                     min_choice, ms1 = multi_node.get_min_choice()
 
+                    # Some choices can fail to benchmark, inf timing
+                    future_choices = [
+                        fut_choice
+                        for fut_choice in future_choices
+                        if fut_choice[0] in choice_timings
+                    ]
+
                     future_choices = sorted(
                         future_choices,
                         key=lambda x: choice_timings[x[0]],
@@ -4208,7 +4574,7 @@ def benchmark_when_ready() -> bool:
                             fusion_log.debug(  # noqa: G200
                                 "Exception in compiling %s: %s",
                                 "prologue" if not epilogue_fusion else "epilogue",
-                                str(e),
+                                e,
                             )
                         continue
 
@@ -4230,16 +4596,28 @@ def benchmark_when_ready() -> bool:
                             or ms2 + ms1 > choice_timings[choice] + ms2_fused
                         )
 
-                        if (
-                            res
+                        if res and fusible_choice:
+                            choice.precompile()
                             # pyrefly: ignore [missing-attribute]
-                            and len(res.launchers) == 1
+                            assert res.launchers and choice.n_regs
                             # pyrefly: ignore [bad-index]
-                            and res.launchers[0].n_spills <= 8
-                            and fusible_choice
-                        ):
-                            ms_fused_choice = choice
-                            break
+                            compiled_kernel = res.launchers[0]
+                            # pyrefly: ignore [missing-attribute]
+                            fused_n_regs = compiled_kernel.n_regs
+                            # pyrefly: ignore [missing-attribute]
+                            fused_n_spills = compiled_kernel.n_spills
+                            should_fuse_epilogue = _fuse_epilogue(
+                                ms1,
+                                ms2,
+                                choice.n_regs,
+                                fused_n_regs,
+                                fused_n_spills,
+                                choice.bmreq.num_warps,
+                                DeviceProperties.create(device),
+                            )
+                            if should_fuse_epilogue:
+                                ms_fused_choice = choice
+                                break
 
                 if bench_epilogue:
                     log_fusion(min_ms_fused, ms1, ms2)
@@ -4257,8 +4635,9 @@ def benchmark_when_ready() -> bool:
                         # pyrefly: ignore [missing-attribute]
                         multi_node.finalize_as_triton_caller(ms_fused_choice)
 
-                    # pyrefly: ignore [missing-attribute]
-                    multi_node._choice_timings[None] = new_timings
+                    if bench_epilogue:
+                        # pyrefly: ignore [missing-attribute]
+                        multi_node._choice_timings[None] = new_timings
                     return True
                 else:
                     return False
@@ -4368,6 +4747,13 @@ def fuse_two_nodes(
         fused_nodes.remove(node2)
         fused_nodes.add(node3)
         self.name_to_fused_node.update({n.get_name(): node3 for n in node3.get_nodes()})
+
+        # Propagate stream assignment to the fused node so that subsequent
+        # fusion rounds still respect stream boundaries.
+        stream1 = self.node_to_stream.get(node1)
+        if stream1 is not None:
+            self.node_to_stream[node3] = stream1
+
         return node3
 
     def fuse_if_speedup(
@@ -4575,12 +4961,12 @@ def _handle_template_overlap(
         )
         new_possible_fusions = []
         for n1, n2 in possible_fusions:
-            if is_prologue_fusion(n1, n2) and n1 in epilogue_template_nodes:
+            if is_prologue_fusion(n1, n2) and n2 in epilogue_template_nodes:
                 deferred_prologue_fusions.append((n1, n2))
             else:
                 new_possible_fusions.append((n1, n2))
 
-        possible_fusions = new_possible_fusions
+        return new_possible_fusions
 
     def fuse_nodes_once(
         self,
@@ -4620,12 +5006,10 @@ def fuse_nodes_once(
             is_reorder_round,
         )
 
-        if (
-            (config.max_autotune_gemm or config.max_autotune)
-            and config.prologue_fusion
-            and config.epilogue_fusion
-        ):
-            self._handle_template_overlap(possible_fusions, deferred_prologue_fusions)
+        if config.max_autotune_gemm or config.max_autotune:
+            possible_fusions = self._handle_template_overlap(
+                possible_fusions, deferred_prologue_fusions
+            )
 
         self._try_fusion_pairs(
             possible_fusions,
@@ -4653,7 +5037,7 @@ def fuse_nodes_once(
         nodes = self.topological_sort_schedule(nodes)
         return nodes
 
-    def create_combo_kernel_nodes(self, num_ck_nodes: Optional[int] = None) -> None:
+    def create_combo_kernel_nodes(self, num_ck_nodes: int | None = None) -> None:
         """
         Groups parallel nodes
         """
@@ -4691,6 +5075,11 @@ def create_combo_kernel_nodes(self, num_ck_nodes: Optional[int] = None) -> None:
             self.name_to_fused_node.update(
                 {n.get_name(): group_snode for n in group_snode.get_nodes()}
             )
+            # Propagate stream assignment so codegen can place the combo
+            # kernel in the correct stream context.
+            stream = self.node_to_stream.get(node_list[0])
+            if stream is not None:
+                self.node_to_stream[group_snode] = stream
         self.nodes = sorted(fused_nodes, key=lambda x: x.min_order)
         self.nodes = self.topological_sort_schedule(self.nodes)
         log.info(
@@ -4943,7 +5332,7 @@ def decide_fusion_fail_reason(
         self,
         node1: BaseSchedulerNode,
         node2: BaseSchedulerNode,
-        common_buf_names: Union[tuple[str, ...], OrderedSet[str]],
+        common_buf_names: tuple[str, ...] | OrderedSet[str],
     ) -> str:
         """
         Try to decide reasons why fusion fail due to no shared memory even though
@@ -5192,7 +5581,9 @@ def shared_data_after_reordering_loop(
             ):
                 candidates.append(
                     (
-                        V.graph.sizevars.size_hint(lhs_dep.get_numel(), fallback=0),
+                        V.graph.sizevars.optimization_hint(
+                            lhs_dep.get_numel(), fallback=0
+                        ),
                         lhs_dep,
                         rhs_dep,
                     )
@@ -5239,11 +5630,17 @@ def unfusable_node(self, node: BaseSchedulerNode) -> bool:
         """
         Is this node unfusable under any conditions.
         """
-        return (
-            isinstance(node, (ExternKernelSchedulerNode, NopKernelSchedulerNode))
-            and not node.is_template()
-            and not is_output_of_multi_outputs_template(node.node)
-        )
+        if isinstance(node, NopKernelSchedulerNode):
+            return not node.is_template() and not is_output_of_multi_outputs_template(
+                node.node
+            )
+        if isinstance(node, ExternKernelSchedulerNode):
+            if isinstance(node.node, ir.UserDefinedTritonKernel):
+                return not node.node.can_fuse_epilogue()
+            return not node.is_template() and not is_output_of_multi_outputs_template(
+                node.node
+            )
+        return False
 
     def check_prologue_fusion_heuristics_fusable(
         self,
@@ -5293,8 +5690,10 @@ def check_prologue_fusion_heuristics_fusable(
         def low_prec_fp(dtype: torch.dtype) -> bool:
             return dtype.itemsize <= 2 and dtype.is_floating_point
 
+        template_buf = template_node.get_template_node_or_throw()
         if (
-            low_prec_fp(template_node.get_template_node_or_throw().dtype)
+            not template_buf.is_multi_outputs_template()
+            and low_prec_fp(template_buf.dtype)
             and not prologue_node.can_codegen_in_low_precision()
         ):
             why(
@@ -5306,7 +5705,7 @@ def low_prec_fp(dtype: torch.dtype) -> bool:
 
     def get_expand_dim_for_pointwise_nodes(
         self, node1: BaseSchedulerNode, node2: BaseSchedulerNode
-    ) -> Optional[tuple[int, SchedulerNode, sympy.Expr]]:
+    ) -> tuple[int, SchedulerNode, sympy.Expr] | None:
         """
         Fusing two small pointwise nodes significantly reduces kernel overhead
         and launch overhead. However, slightly different sizes would prevent fusion.
@@ -5364,7 +5763,7 @@ def get_expand_dim_for_pointwise_nodes(
         # race condition
         def has_reusable_buffer(node: BaseSchedulerNode) -> bool:
             for read in node.read_writes.reads:
-                input_buf: Optional[Union[SchedulerBuffer, SchedulerDonatedBuffer]]
+                input_buf: SchedulerBuffer | SchedulerDonatedBuffer | None
                 if read.name in self.name_to_donated_buffer:
                     input_buf = self.name_to_donated_buffer[read.name]
                 else:
@@ -5416,6 +5815,13 @@ def can_fuse(
         if node1 is node2:
             return False
 
+        # Prevent fusion across stream boundaries
+        if self._has_multi_stream_nodes():
+            stream1 = self.node_to_stream.get(node1)
+            stream2 = self.node_to_stream.get(node2)
+            if stream1 is not None and stream2 is not None and stream1 != stream2:
+                return False
+
         if isinstance(node1, FusedMixOrderReductions):
             return node1.can_fuse_with(node2)
         if isinstance(node2, FusedMixOrderReductions):
@@ -5435,12 +5841,67 @@ def can_fuse(
         ):
             why("grouped node must not be fused with other nodes")
             return False
-        if (
-            isinstance(node1, (ExternKernelSchedulerNode, NopKernelSchedulerNode))
-            and not node1.is_template()
-        ):
-            why("node1 is extern or nop")
+        if isinstance(node1, NopKernelSchedulerNode) and not node1.is_template():
+            why("node1 is nop")
             return False
+
+        if isinstance(node1, ExternKernelSchedulerNode):
+            if not isinstance(node1.node, ir.UserDefinedTritonKernel):
+                why("node1 is extern but not a triton kernel")
+                return False
+
+            if not node1.node.can_fuse_epilogue():
+                why("node1's triton kernel doesn't support epilogue fusion")
+                return False
+
+            if not isinstance(node2, SchedulerNode):
+                why("node1 is extern but node2 is not SchedulerNode")
+                return False
+            if not isinstance(node2.node, ComputedBuffer):
+                why("node1 is extern but node2.node is not SchedulerNode")
+                return False
+            if not isinstance(node2.node.data, Pointwise):
+                why("node1 is extern but node2.node.data is not Pointwise")
+                return False
+
+            assert len(node1.node.mutation_outputs) == 1
+            written_buffer_name = node1.node.mutation_outputs[0].name
+
+            # The epilogue can only read from the output buffer.
+            # Any other tensor/s would require additional load expressions.
+            if any(dep.name != written_buffer_name for dep in node2.read_writes.reads):
+                why("epilogue reads from buffers other than the mutated output")
+                return False
+
+            # the epilogue depends on expressions which may not available in the user triton kernel
+            # (e.g. indexing exprs used not in a load)
+            node2_inner_fn_free_symbols = node2.node.data.inner_fn_free_symbols()
+            for symbol in node2_inner_fn_free_symbols:
+                usages = node2.node.data.collect_inner_fn_symbol_usage(symbol)
+                if any(usage != "load" for usage in usages):
+                    return False
+
+            # should be true now because we checked `can_fuse_epilogue`
+            assert len(node1.node.mutable_args) == 1
+            if node1.node.mutable_args[0].layout != node2.node.layout:
+                why("node1 and node2 uses different buf layouts")
+                return False
+
+            def _is_other_node_that_references_mutation_buffer(
+                other_node: BaseSchedulerNode,
+            ):
+                return (
+                    (other_node is not node1)
+                    and (other_node is not node2)
+                    and written_buffer_name in other_node.used_buffer_names()
+                )
+
+            if any(
+                _is_other_node_that_references_mutation_buffer(node)
+                for node in self.nodes
+            ):
+                return False
+
         if (
             isinstance(node2, (ExternKernelSchedulerNode, NopKernelSchedulerNode))
             and not node2.is_template()
@@ -5453,7 +5914,7 @@ def can_fuse(
             return False
 
         if node2.is_template():
-            if not config.prologue_fusion:
+            if not _is_prologue_fusion_enabled(node2):
                 why("prologue fusion turned off")
                 return False
 
@@ -5462,11 +5923,10 @@ def can_fuse(
                 return False
 
             template = node2.get_template_node_or_throw()
-            if not isinstance(template, ir.TritonTemplateBuffer):
-                why("prologue fusion only supported for TritonTemplates")
-                return False
-
             allowed_prologue_inps = template.get_allowed_prologue_inps()
+            if not allowed_prologue_inps:
+                why("template has no allowed prologue inputs")
+                return False
 
             unsupported_prologue_args = (
                 OrderedSet(inp.get_name() for inp in template.inputs)  # type: ignore[union-attr]
@@ -5510,13 +5970,21 @@ def can_fuse(
             if not self.check_prologue_fusion_heuristics_fusable(node1, node2, why):
                 return False
 
-        if node1.is_template() and (
-            node2.has_aliasing_or_mutation()
-            or node2.is_reduction()
-            or not config.epilogue_fusion
-        ):
-            why("template epilogue not satisfied")
-            return False
+        if node1.is_template():
+            if (
+                node2.has_aliasing_or_mutation()
+                or node2.is_reduction()
+                or not _is_epilogue_fusion_enabled(node1)
+            ):
+                why("template epilogue not satisfied")
+                return False
+            template_buf = node1.get_template_node()
+            assert template_buf is not None
+            if template_buf.is_multi_outputs_template() and not isinstance(
+                node2.node, ir.ComputedBuffer
+            ):
+                why("multi-output template epilogue requires ComputedBuffer")
+                return False
 
         if (node1.get_buffer_names() & V.graph.no_fuse_buffer_names) or (
             node2.get_buffer_names() & V.graph.no_fuse_buffer_names
@@ -5606,14 +6074,22 @@ def can_fuse_vertical(
             remaining_deps_by_name[name].append(dep)
 
         for cd in node1.read_writes.writes:
-            if not isinstance(cd, MemoryDep):
+            if not isinstance(cd, MemoryDep) and not isinstance(cd, StarDep):
                 continue
             remaining = remaining_deps_by_name.get(
                 self.mutation_renames.get(cd.name, cd.name)
             )
             if remaining:
                 for rd in remaining:
-                    if self.fusable_read_and_write(rd, cd):
+                    if isinstance(cd, MemoryDep) and self.fusable_read_and_write(
+                        rd, cd
+                    ):
+                        remaining.remove(rd)  # noqa: B909
+                    elif isinstance(
+                        cd, StarDep
+                    ) and self.fusable_stardep_write_and_read_on_empty_tensor(
+                        rd, cd, node1.node
+                    ):
                         remaining.remove(rd)  # noqa: B909
 
         remaining_deps = OrderedSet(
@@ -5661,6 +6137,12 @@ def fusable_weak_dep(
         if free_symbol_is_type(write.index, SymT.TMP):
             return False
 
+        # Non-injective scatter: range vars absent from write index mean
+        # multiple iterations hit the same location. Can't fuse the reader
+        # in or it will see partially-written state between iterations.
+        if not OrderedSet(write.var_names) <= write.index.free_symbols:
+            return False
+
         real_name = self.mutation_real_name[weak_dep.mutating_buf]
         relevant_reading_nodes = [node1]
         if isinstance(node1, ForeachKernelSchedulerNode):
@@ -5728,6 +6210,22 @@ def fusable_read_and_write(self, read: Dep, write: MemoryDep) -> bool:
                 return True
         return False
 
+    # on tensors that are "empty" (i.e. with undefined values),
+    # we relax the conditions for fusion and additionally allow matching a writing StarDep with any read dep.
+    # This makes use of the fact that `f(UB) = UB`.
+    def fusable_stardep_write_and_read_on_empty_tensor(
+        self, read: Dep, write: StarDep, writing_node: ir.Operation | None
+    ) -> bool:
+        if not isinstance(writing_node, ir.UserDefinedTritonKernel):
+            return False
+        if not writing_node.can_fuse_epilogue():
+            return False
+        read_name = self.mutation_renames.get(read.name, read.name)
+        write_name = self.mutation_renames.get(write.name, write.name)
+        if isinstance(write, StarDep) and read_name == write_name:
+            return True
+        return False
+
     def dep_size_hint(self, dep: Dep, count_bytes: bool = True) -> int:
         return V.graph.get_dep_size_hint(dep, count_bytes)
 
@@ -5738,18 +6236,30 @@ def score_fusion_memory(
         count_bytes: bool = True,
         return_is_mix_order_reduction: bool = False,
         allow_mix_order_reduction: bool = True,
-    ) -> int | tuple[int, bool]:
+    ) -> int | tuple[int, int, bool]:
         """
         The first term in our fusion score that estimates number of saved
         memory operations.
+
+        This function scores fusion candidates based on shared memory access patterns.
+        Higher scores indicate better fusion candidates.
+
+        Scoring strategy:
+        1. If nodes share exact memory deps (same buffer + same indexing), return
+           the sum of shared dep sizes (original behavior).
+        2. If no exact matches (score == 0), check for same-buffer reads with
+           different indexing (e.g., split operations reading different slices).
+           - Give bonus if nodes read from exactly the same set of buffers
+           - Score based on overlap ratio: common_buffer_size / total_read_size
+           - High overlap (>50%) suggests good cache locality benefit from fusion
         """
 
-        def _construct_return_value(score, is_mix_order_reduction):
-            return (
-                (score, is_mix_order_reduction)
-                if return_is_mix_order_reduction
-                else score
-            )
+        def _construct_return_value(
+            score, buffer_overlap_score, is_mix_order_reduction
+        ):
+            if return_is_mix_order_reduction:
+                return (score, buffer_overlap_score, is_mix_order_reduction)
+            return score + buffer_overlap_score
 
         if allow_mix_order_reduction and MixOrderReduction.can_fuse(node1, node2):
             # The fusion score for mix order reduction only count
@@ -5757,7 +6267,43 @@ def _construct_return_value(score, is_mix_order_reduction):
             # sharing the same amount of numels go first; but make
             # fusions only share weight/bias go later.
             score = MixOrderReduction.get_fusion_score(node1, node2)
-            return _construct_return_value(score, True)
+            return _construct_return_value(score, 0, True)
+
+        # For UserDefinedTritonKernel, the write deps are StarDep that won't
+        # match the epilogue's MemoryDep via set intersection.  For templates,
+        # a view/reshape between the template output and epilogue can produce
+        # different index expressions that don't match via set intersection.
+        # Fall back to name-based matching so that the fusion score reflects
+        # the actual shared buffers.
+        if (
+            (
+                isinstance(node1.node, ir.UserDefinedTritonKernel)
+                and node1.node.can_fuse_epilogue()
+            )
+            or node1.is_template()
+            or node2.is_template()
+        ):
+            node1_deps = node1.read_writes.reads | node1.read_writes.writes
+            node2_deps = node2.read_writes.reads | node2.read_writes.writes
+
+            def _match(dep1: Dep, dep2: Dep):
+                if dep1 == dep2:
+                    return True
+                if isinstance(dep1, (StarDep, MemoryDep)) and isinstance(
+                    dep2, (StarDep, MemoryDep)
+                ):
+                    return dep1.name == dep2.name
+                return False
+
+            score = 0
+            for node1_dep in node1_deps:
+                for node2_dep in node2_deps:
+                    if _match(node1_dep, node2_dep):
+                        score += max(
+                            self.dep_size_hint(node1_dep), self.dep_size_hint(node2_dep)
+                        )
+
+            return _construct_return_value(score, 0, False)
 
         node1_dep_len = len(node1.read_writes.reads) + len(node1.read_writes.writes)
         node2_dep_len = len(node2.read_writes.reads) + len(node2.read_writes.writes)
@@ -5774,14 +6320,210 @@ def _construct_return_value(score, is_mix_order_reduction):
             ]
 
             return _construct_return_value(
-                sum(self.dep_size_hint(dep, count_bytes) for dep in deps), False
+                sum(self.dep_size_hint(dep, count_bytes) for dep in deps), 0, False
             )
 
         common_memory_deps = (node1.read_writes.reads | node1.read_writes.writes) & (
             node2.read_writes.reads | node2.read_writes.writes
         )
-        return _construct_return_value(
-            sum(self.dep_size_hint(dep) for dep in common_memory_deps), False
+
+        score = sum(self.dep_size_hint(dep) for dep in common_memory_deps)
+
+        # If no exact dep matches, check for same-buffer reads with different indexing.
+        # This handles cases like split operations that read different slices of the
+        # same buffer - they should fuse for cache locality benefits.
+        buffer_overlap_score = 0
+        if score == 0 and self._can_use_buffer_overlap_scoring(node1, node2):
+            buffer_overlap_score = self._score_fusion_memory_by_buffer_overlap(
+                node1, node2
+            )
+
+        return _construct_return_value(score, buffer_overlap_score, False)
+
+    def _can_use_buffer_overlap_scoring(
+        self,
+        node1: BaseSchedulerNode,
+        node2: BaseSchedulerNode,
+    ) -> bool:
+        """
+        Check if buffer overlap scoring should be used for this node pair.
+
+        Buffer overlap scoring handles split/cat patterns where nodes read from
+        the same buffer at different indices. We skip it when:
+        - Either node is a reduction (different memory access patterns)
+        - Either node is a template
+        - Both nodes are prologue/epilogue candidates for the same template,
+          because horizontal fusion would prevent them from being absorbed
+          into the template kernel. For example, in:
+            q = a[:64, :]; k = a[64:, :]
+            return mm(q + 2, k - 2)
+          "q + 2" and "k - 2" both read from `a` and would get a high overlap
+          score, but fusing them horizontally prevents prologue fusion into mm
+          (resulting in 2 kernels instead of 1).
+
+        We allow buffer overlap scoring when:
+        - The node outputs are not actually in the template's allowed_prologue_inps,
+          meaning they can't be prologue-fused anyway, so horizontal fusion doesn't
+          prevent any optimization opportunity.
+        """
+        if node1.is_reduction() or node2.is_reduction():
+            return False
+        if node1.is_template() or node2.is_template():
+            return False
+
+        if config.max_autotune or config.max_autotune_gemm:
+            node1_outputs = node1.get_outputs()
+            node2_outputs = node2.get_outputs()
+
+            # Early return if either node has no outputs
+            if not node1_outputs or not node2_outputs:
+                return True
+
+            node1_output_names = OrderedSet(buf.get_name() for buf in node1_outputs)
+            node2_output_names = OrderedSet(buf.get_name() for buf in node2_outputs)
+
+            # Find templates that consume node1's outputs via buffer users
+            # and check if the outputs are actually prologue-fusable
+            node1_prologue_eligible_template_users: OrderedSet[BaseSchedulerNode] = (
+                OrderedSet()
+            )
+            for buf in node1_outputs:
+                for user in buf.users:
+                    if (
+                        isinstance(user.node, BaseSchedulerNode)
+                        and user.node.is_template()
+                        and _is_prologue_fusion_enabled(user.node)
+                    ):
+                        # Check if this output is actually in the template's
+                        # allowed_prologue_inps. If not, fusing horizontally
+                        # won't prevent any prologue fusion opportunity.
+                        template_node = user.node.get_template_node()
+                        if template_node is not None and isinstance(
+                            template_node, ir.TritonTemplateBuffer
+                        ):
+                            allowed_inps = template_node.get_allowed_prologue_inps()
+                            if node1_output_names & allowed_inps:
+                                node1_prologue_eligible_template_users.add(user.node)
+                        else:
+                            # Conservative: assume it could be a prologue candidate
+                            node1_prologue_eligible_template_users.add(user.node)
+
+            # Check if any of node1's prologue-eligible template users also consume
+            # node2's outputs in a prologue-eligible way
+            if node1_prologue_eligible_template_users:
+                for buf in node2_outputs:
+                    for user in buf.users:
+                        if (
+                            isinstance(user.node, BaseSchedulerNode)
+                            and user.node.is_template()
+                            and user.node in node1_prologue_eligible_template_users
+                        ):
+                            # Also verify node2's output is prologue-eligible
+                            template_node = user.node.get_template_node()
+                            if template_node is not None and isinstance(
+                                template_node, ir.TritonTemplateBuffer
+                            ):
+                                allowed_inps = template_node.get_allowed_prologue_inps()
+                                if node2_output_names & allowed_inps:
+                                    return False
+                            else:
+                                # Conservative: block fusion
+                                return False
+
+            for node in (node1, node2):
+                for dep in node.read_writes.reads:
+                    producer = self.name_to_fused_node.get(dep.name)
+                    if (
+                        producer is not None
+                        and producer.is_template()
+                        and _is_epilogue_fusion_enabled(producer)
+                    ):
+                        return False
+
+        return True
+
+    def _score_fusion_memory_by_buffer_overlap(
+        self, node1: BaseSchedulerNode, node2: BaseSchedulerNode
+    ) -> int:
+        """
+        Score fusion based on buffer name overlap when exact dep matching fails.
+
+        This handles the split/cat fusion case where nodes read from the same buffer
+        but at different indices (e.g., different slices from a split operation).
+
+        Scoring logic:
+        - If nodes read from exactly the same buffers: high bonus (encourages fusion)
+        - For common buffers: score based on overlap ratio
+          - overlap_ratio = common_buffer_size /
+            max(node1_total_reads, node2_total_reads)
+          - If overlap_ratio > threshold (e.g., 0.5): give proportional score
+          - If overlap_ratio < threshold: minimal/no score (not worth fusing)
+
+        Note on dynamic shapes:
+        - When deps have unbacked symbols (dynamic shapes), dep_size_hint returns 0
+        - In this case, we use count * 10 as a proxy for size
+        - This ensures fusion still works for models with dynamic batch sizes
+
+        Note on multiple deps from same buffer:
+        - A node may have multiple MemoryDep entries for the same buffer name
+          (e.g., 4 split reads from arg0_1 at different indices)
+        - We sum ALL dep sizes for each buffer, not just take max
+        - This ensures overlap ratio is calculated correctly when nodes read
+          multiple slices from the same underlying buffer
+        """
+        # Fallback size when dep_size_hint returns 0 (e.g., unbacked symbols)
+        FALLBACK_DEP_SIZE = 10
+
+        def get_dep_size(dep: Dep) -> int:
+            size = self.dep_size_hint(dep)
+            return size if size > 0 else FALLBACK_DEP_SIZE
+
+        node1_read_names = OrderedSet(dep.name for dep in node1.read_writes.reads)
+        node2_read_names = OrderedSet(dep.name for dep in node2.read_writes.reads)
+
+        # Early exit if no common buffer names
+        common_names = node1_read_names & node2_read_names
+
+        if not common_names:
+            return 0
+
+        # Calculate total read sizes for each node (sum of ALL deps)
+        node1_total_read_size = sum(
+            get_dep_size(dep) for dep in node1.read_writes.reads
+        )
+        node2_total_read_size = sum(
+            get_dep_size(dep) for dep in node2.read_writes.reads
+        )
+
+        max_total_read_size = max(node1_total_read_size, node2_total_read_size)
+        if max_total_read_size == 0:
+            return 0
+
+        # Calculate total reads from common buffers for each node
+        # Sum ALL deps for each common buffer name
+        # (handles multiple reads from same buffer)
+        node1_common_read_size = sum(
+            get_dep_size(dep)
+            for dep in node1.read_writes.reads
+            if dep.name in common_names
+        )
+        node2_common_read_size = sum(
+            get_dep_size(dep)
+            for dep in node2.read_writes.reads
+            if dep.name in common_names
+        )
+
+        # Use max of the two as the common buffer size estimate
+        # This represents how much data is being read from shared buffers
+        common_read_buffer_size = max(node1_common_read_size, node2_common_read_size)
+
+        # Calculate overlap ratio
+        overlap_ratio = common_read_buffer_size / max_total_read_size
+        # Scale score by overlap ratio and common buffer size
+        # Higher overlap = higher score
+        # Larger common buffer = higher score (more cache benefit)
+        return (
+            common_read_buffer_size if overlap_ratio >= config.min_overlap_ratio else 0
         )
 
     def get_possible_fusions_with_highest_priority(
@@ -5850,7 +6592,7 @@ def free_buffers(self) -> None:
                 inp = V.graph.graph_inputs[name]
                 if isinstance(inp, ir.TorchBindObject):
                     V.graph.wrapper_code.codegen_free(inp)
-                elif isinstance(inp, ir.GeneratorState):
+                elif isinstance(inp, (ir.GeneratorState, ir.OpaqueObjectState)):
                     continue
                 else:
                     storage = inp.data
@@ -5866,8 +6608,14 @@ def flush(self) -> None:
             backend.flush()
         self.free_buffers()
 
-    def codegen_extern_call(self, scheduler_node: ExternKernelSchedulerNode) -> None:
-        assert isinstance(scheduler_node, ExternKernelSchedulerNode)
+    def codegen_extern_call(
+        self,
+        scheduler_node: BaseSchedulerNode,
+    ) -> None:
+        assert isinstance(
+            scheduler_node,
+            (ExternKernelSchedulerNode, FusedExternTritonKernelSchedulerNode),
+        )
         # 'decide_inplace_update' stores the inplace update decisions in
         # the current kernel from where 'allocate' retrieve those decisions.
         # We have to make sure there is a non-NULL kernel handler to store
@@ -5876,9 +6624,7 @@ def codegen_extern_call(self, scheduler_node: ExternKernelSchedulerNode) -> None
         with V.set_kernel_handler(Kernel(increase_kernel_count=False)):
             scheduler_node.decide_inplace_update()
             scheduler_node.mark_run()
-        node = scheduler_node.node
-        assert isinstance(node, ir.ExternKernel), f"{type(node)=}"
-        node.codegen(V.graph.wrapper_code)
+        scheduler_node.codegen(V.graph.wrapper_code)
         self.free_buffers()
 
     def create_backend(self, device: torch.device) -> BaseScheduling:
@@ -5902,7 +6648,7 @@ def create_backend(self, device: torch.device) -> BaseScheduling:
 
         return device_scheduling(self)
 
-    def get_backend(self, device: Optional[torch.device]) -> BaseScheduling:
+    def get_backend(self, device: torch.device | None) -> BaseScheduling:
         assert device is not None
         if device not in self.backends:
             self.backends[device] = self.create_backend(device)
@@ -5939,7 +6685,7 @@ def can_buffer_be_removed_through_fusion(
             and name not in self.mutation_real_name
         )
 
-    def should_partition(self, node: BaseSchedulerNode) -> Optional[str]:
+    def should_partition(self, node: BaseSchedulerNode) -> str | None:
         """
         Return the reason why we should partition the inductor graph on this node,
         or None if the node is cudagraphable.
@@ -6040,7 +6786,7 @@ def _get_cudagraph_unsafe_unbacked_symints(self) -> OrderedSet[sympy.Symbol]:
 
     def _uses_cudagraph_unsafe_unbacked_symint(
         self, node: BaseSchedulerNode
-    ) -> Optional[str]:
+    ) -> str | None:
         unsafe_symints = self._get_cudagraph_unsafe_unbacked_symints()
         if not unsafe_symints:
             return None
@@ -6057,12 +6803,12 @@ def _uses_cudagraph_unsafe_unbacked_symint(
 
     def get_name_to_nodes(
         self,
-    ) -> dict[str, Union[ir.IRNode, ir.TorchBindObject, sympy.Expr]]:
+    ) -> dict[str, ir.IRNode | ir.TorchBindObject | sympy.Expr]:
         """
         Return a mapping from name strings to the corresponding graph inputs or
         base scheduler node outputs.
         """
-        name_to_node: dict[str, Union[ir.IRNode, ir.TorchBindObject, sympy.Expr]] = {}
+        name_to_node: dict[str, ir.IRNode | ir.TorchBindObject | sympy.Expr] = {}
         name_to_node.update(V.graph.graph_inputs)
 
         for node in self.nodes:
@@ -6115,7 +6861,7 @@ def compute_graph_partition_maps(
     def get_graph_partition_symbol_inputs(
         self,
         partition: PartitionType,
-        input_nodes: dict[str, Union[ir.IRNode, ir.TorchBindObject, sympy.Expr]],
+        input_nodes: dict[str, ir.IRNode | ir.TorchBindObject | sympy.Expr],
     ) -> OrderedSet[sympy.Symbol]:
         """
         Returns all symbol inputs which are required to be in scope to successfully
@@ -6126,7 +6872,7 @@ def get_graph_partition_symbol_inputs(
         """
 
         def get_input_node_symbols(
-            node: Union[ir.IRNode, sympy.Expr, ir.TorchBindObject],
+            node: ir.IRNode | sympy.Expr | ir.TorchBindObject,
         ) -> OrderedSet[sympy.Symbol]:
             """
             Gets symbols used in input node shapes, strides, and offsets.
@@ -6510,6 +7256,27 @@ def graph_partition(
             partitions.append(cur_partition)
             skip_cudagraphs.append(skip_cudagraph)
 
+        # Apply minimum partition size threshold: if a cudagraph-eligible partition
+        # has fewer kernels than the threshold, mark it as non-cudagraphable
+        min_size = config.triton.cudagraph_min_partition_size
+        if min_size > 0:
+            for i, (partition, skip) in enumerate(zip(partitions, skip_cudagraphs)):
+                if not skip:
+                    # Count kernels excluding NopKernelSchedulerNode
+                    kernel_count = sum(
+                        1
+                        for n in partition
+                        if not isinstance(n, NopKernelSchedulerNode)
+                    )
+                    if kernel_count < min_size:
+                        skip_cudagraphs[i] = True
+                        cudagraphs_log.debug(
+                            "Partition %d has %d kernels, below minimum size %d, skipping cudagraph",
+                            i,
+                            kernel_count,
+                            min_size,
+                        )
+
         signatures = self.get_graph_partition_signature(
             partitions=partitions, skip_cudagraphs=skip_cudagraphs
         )
@@ -6764,6 +7531,10 @@ def _codegen(self, nodes: list[BaseSchedulerNode]) -> None:
         if self.default_device_context and config.triton.autotune_at_compile_time:
             V.graph.wrapper_code.write_get_raw_stream_header()
 
+        # Register non-mutated inputs that need alignment checks.
+        # Deferred to just before the first kernel that reads each input.
+        V.graph.wrapper_code.register_alignment_check_inputs()
+
         for node in nodes:
             if log.isEnabledFor(logging.DEBUG):
                 try:
@@ -6780,6 +7551,12 @@ def _codegen(self, nodes: list[BaseSchedulerNode]) -> None:
 
             self.enter_context(node)
 
+            # pyrefly: ignore [unbound-name]
+            if config.size_asserts:
+                V.graph.wrapper_code.codegen_deferred_input_asserts(
+                    dep.name for dep in node.read_writes.reads
+                )
+
             if device := node.get_device():
                 if (
                     device != self.current_device
@@ -6791,11 +7568,43 @@ def _codegen(self, nodes: list[BaseSchedulerNode]) -> None:
                     if self.current_device and device_need_guard(
                         self.current_device.type
                     ):
+                        # Exit stream context before exiting device guard
+                        if self.current_stream_idx is not None:
+                            self.generate_stream_ctx_exit()
                         V.graph.wrapper_code.codegen_device_guard_exit()
                     self.current_device = device
                     if device_need_guard(device.type):
                         assert device.index is not None, "device should have an index"
-                        V.graph.wrapper_code.codegen_device_guard_enter(device.index)
+                        # Compute num_streams if we have multi-stream nodes
+                        num_streams = 1
+                        if self._has_multi_stream_nodes():
+                            # Count unique streams (excluding default stream 0)
+                            unique_streams = OrderedSet(self.node_to_stream.values())
+                            num_streams = (
+                                max(unique_streams) + 1 if unique_streams else 1
+                            )
+                        V.graph.wrapper_code.codegen_device_guard_enter(
+                            device.index,
+                            num_streams,
+                            self.stream_idx_to_user_obj_idx,
+                        )
+
+            # Handle stream context switching for multi-stream scheduling.
+            # This runs for all nodes (including device-less sync ops like
+            # record_event/wait_event) so they are placed inside the correct
+            # stream context. Only switch when inside a device guard (i.e.
+            # current_device is set), since stream variables are declared there.
+            if self._has_multi_stream_nodes() and self.current_device is not None:
+                self.generate_stream_ctx_switching(node)
+
+            # Emit deferred alignment copies for inputs first used by this
+            # node.  This runs *after* stream context switching so the copy
+            # executes on the same stream as the consuming kernel.
+            # TODO: inputs read on multiple streams should be copied in the
+            # prologue instead, to avoid cross-stream races.
+            V.graph.wrapper_code.codegen_deferred_alignment_copies(
+                dep.name for dep in node.read_writes.reads
+            )
 
             self.current_node = node
             self.buffer_names_to_free.update(node.last_usage)
@@ -6809,7 +7618,6 @@ def _codegen(self, nodes: list[BaseSchedulerNode]) -> None:
                     template_node, epilogue, prologue
                 )
             elif node.is_extern():
-                node = typing.cast(ExternKernelSchedulerNode, node)
                 self.codegen_extern_call(node)
             elif node.is_foreach():
                 node = typing.cast(ForeachKernelSchedulerNode, node)
@@ -6817,8 +7625,12 @@ def _codegen(self, nodes: list[BaseSchedulerNode]) -> None:
                 backend_ = self.get_backend(device)
                 from .codegen.cuda_combined_scheduling import CUDACombinedScheduling
                 from .codegen.simd import SIMDScheduling
+                from .codegen.xpu.xpu_combined_scheduling import XPUCombinedScheduling
 
-                if isinstance(backend_, (SIMDScheduling, CUDACombinedScheduling)):
+                if isinstance(
+                    backend_,
+                    (SIMDScheduling, CUDACombinedScheduling, XPUCombinedScheduling),
+                ):
                     backend = backend_
                 else:
                     raise AssertionError(f"{type(self)=}")
@@ -6870,7 +7682,7 @@ def _codegen(self, nodes: list[BaseSchedulerNode]) -> None:
 
     def benchmark_combo_kernel(
         self, node_list: Sequence[BaseSchedulerNode], node_benchmark_results
-    ) -> tuple[float, float, list[Optional[str]]]:
+    ) -> tuple[float, float, list[str | None]]:
         """
         Benchmark fused list of nodes and return the execution time
         in milliseconds on randomly generated inputs.
@@ -6982,9 +7794,68 @@ def update_zero_dim_cpu_tensor(self) -> None:
                     ):
                         V.graph.zero_dim_cpu_tensor_list.add(read.name)
 
+    @property
+    def current_stream_idx(self) -> int | None:
+        """CUDA Stream index that current scheduler node assigned to."""
+        if self._current_stream_ctx is not None:
+            return self._current_stream_ctx.stream_idx
+        else:
+            return None
+
+    @property
+    def current_stream_name(self) -> str | None:
+        """CUDA Stream name that current scheduler node assigned to."""
+        if (stream_idx := self.current_stream_idx) is not None:
+            return get_stream_name(stream_idx)
+        else:
+            return None
+
+    def generate_stream_ctx_enter(self, node: BaseSchedulerNode) -> None:
+        """Code-gen to enter the Stream context assigned to node."""
+        assert not isinstance(node, NopKernelSchedulerNode)
+        node_stream = self.node_to_stream[node]
+        self._current_stream_ctx = V.graph.wrapper_code.codegen_cuda_stream_enter(
+            stream_idx=node_stream,
+        )
+
+    def generate_stream_ctx_exit(self) -> None:
+        """Code-gen to exit from the current Stream context."""
+        assert self._current_stream_ctx is not None
+        V.graph.wrapper_code.codegen_cuda_stream_exit()
+        self._current_stream_ctx = None
+
+    def generate_stream_ctx_switching(self, node: BaseSchedulerNode) -> None:
+        """Generate stream entering and exiting to properly run node in a multi-stream scenario.
+
+        Stream context switching is only generated if ``node``'s assigned stream is different from
+        the previous node's stream. NopKernelSchedulerNodes have stream=None and inherit the
+        enclosing stream context (or do nothing if no context is active yet).
+        """
+        assert node in self.node_to_stream
+        stream = (
+            None
+            if isinstance(node, NopKernelSchedulerNode)
+            else self.node_to_stream[node]
+        )
+        if self.current_stream_idx == stream:
+            # Covers: same stream as current (no switch needed), and both None
+            # (nop node before any stream context — nothing to do).
+            return
+        elif self.current_stream_idx is not None and stream is None:
+            # Don't generate ctx switching. Memory planning code (e.g., delete buffers) on current
+            # node goes to previous stream ctx.
+            return
+        elif self.current_stream_idx is None and stream is not None:
+            # Enter new ctx, update current stream status.
+            self.generate_stream_ctx_enter(node)
+        else:
+            # Switching from previous stream ctx to the new stream ctx.
+            self.generate_stream_ctx_exit()
+            self.generate_stream_ctx_enter(node)
+
 
 class BaseScheduling:  # noqa: docstring_linter
-    def __init__(self, scheduler: Optional[Scheduler]):
+    def __init__(self, scheduler: Scheduler | None):
         super().__init__()
         self.scheduler = scheduler
 
@@ -7021,7 +7892,21 @@ def can_fuse_multi_outputs_template(
         In this context, we verify whether node1 represents the Multi-Output Template
         and node2 corresponds to one of its outputs. If so, we further check if
         backend supports this fusion.
+
         """
+        template_buf = node1.get_template_node()
+        if not isinstance(template_buf, ir.TemplateBuffer):
+            return False
+        if not template_buf.is_multi_outputs_template():
+            return False
+
+        if isinstance(node2.node, ir.MultiOutput):
+            return (
+                len(node2.node.inputs) == 1
+                and isinstance(node2.node.inputs[0], ir.IRNode)
+                and node2.node.inputs[0].get_name() == template_buf.get_name()
+            )
+
         return False
 
     def fuse(
@@ -7036,6 +7921,11 @@ def fuse(
             return FusedMixOrderReductions(node1, node2)
         elif isinstance(node1, FusedMixOrderReductions):
             return node1.fuse_with(node2)
+        elif isinstance(node1, ExternKernelSchedulerNode) and isinstance(
+            node2, SchedulerNode
+        ):
+            assert isinstance(node1.node, ir.UserDefinedTritonKernel)
+            return FusedExternTritonKernelSchedulerNode.epilogue_fuse(node1, node2)
         else:
             return FusedSchedulerNode.fuse(node1, node2)
 
@@ -7052,7 +7942,7 @@ def codegen_template(
         template_node: BaseSchedulerNode,
         epilogue_nodes: Sequence[BaseSchedulerNode],
         prologue_nodes: Sequence[BaseSchedulerNode],
-    ) -> Optional[str]:
+    ) -> str | None:
         """
         Given a template node, generate a kernel.
 
@@ -7065,14 +7955,14 @@ def generate_kernel_code_from_nodes(
         self,
         nodes: Sequence[BaseSchedulerNode],
         benchmark_kernel: bool,
-        hint_override: Optional[int] = None,
+        hint_override: int | None = None,
     ) -> str:
         """
         Generate a kernel given a list of pre-fused nodes.
         """
         raise NotImplementedError
 
-    def codegen_node(self, node: Union[FusedSchedulerNode, SchedulerNode]) -> None:
+    def codegen_node(self, node: FusedSchedulerNode | SchedulerNode) -> None:
         """
         Generate a kernel given a list of pre-fused nodes.
         """
@@ -7127,7 +8017,7 @@ def get_fusion_pair_priority(
 
     def benchmark_combo_kernel(
         self, node_list: Sequence[BaseSchedulerNode], node_benchmark_results
-    ) -> tuple[float, float, list[Optional[str]]]:
+    ) -> tuple[float, float, list[str | None]]:
         """
         Benchmark the list of nodes to combine and return the execution time
         and memory copy time in milliseconds on randomly generated inputs.
@@ -7137,7 +8027,7 @@ def benchmark_combo_kernel(
     def codegen_comment(
         self,
         node_schedule: Sequence[BaseSchedulerNode],
-        kernel_name: Optional[str] = None,
+        kernel_name: str | None = None,
     ) -> None:
         if kernel_name:
             from torch._inductor.debug import set_kernel_post_grad_provenance_tracing
diff --git a/torch/_inductor/select_algorithm.py b/torch/_inductor/select_algorithm.py
index 9c02233204dbb..7468ca1688fb7 100644
--- a/torch/_inductor/select_algorithm.py
+++ b/torch/_inductor/select_algorithm.py
@@ -14,12 +14,13 @@
 import sys
 import textwrap
 import time
+from collections import defaultdict
 from collections.abc import Callable, Sequence
 from concurrent.futures import as_completed, ThreadPoolExecutor
 from io import StringIO
 from pathlib import Path
 from types import ModuleType
-from typing import Any, cast, NamedTuple, Optional, TYPE_CHECKING, Union
+from typing import Any, cast, NamedTuple, Optional, TYPE_CHECKING
 from typing_extensions import Self
 from unittest.mock import patch
 
@@ -131,7 +132,7 @@ class BenchmarkTensors:
     """Represents a set of inputs and outputs for autotuning with a template"""
 
     input_tensors: list[torch.Tensor]
-    output_tensor: Optional[torch.Tensor]
+    output_tensor: torch.Tensor | None
 
     def unpack(self):
         return self.input_tensors, self.output_tensor
@@ -149,7 +150,7 @@ class AutotuneArgs:
 
     triton: BenchmarkTensors
     extern: BenchmarkTensors
-    expected: Optional[torch.Tensor] = None
+    expected: torch.Tensor | None = None
 
     def get_benchmark_tensors(self, extern=False) -> BenchmarkTensors:
         """Returns the inputs and output tensors for a given choice."""
@@ -163,7 +164,7 @@ def from_choice_args(
         example_inputs_extern: list[torch.Tensor],
         out: torch.Tensor,
         out_extern: torch.Tensor,
-        expected: Optional[torch.Tensor] = None,
+        expected: torch.Tensor | None = None,
     ) -> Self:
         """Factory method to create AutotuneInputs from separate inputs/outputs"""
         return cls(
@@ -187,12 +188,10 @@ class PartialRender:
 
     HookFn = Callable[[], str]
 
-    def __init__(
-        self, code: str, replacement_hooks: dict[str, Optional[HookFn]]
-    ) -> None:
+    def __init__(self, code: str, replacement_hooks: dict[str, HookFn | None]) -> None:
         super().__init__()
         self._code: str = code
-        self.replacement_hooks: dict[str, Optional[PartialRender.HookFn]] = (
+        self.replacement_hooks: dict[str, PartialRender.HookFn | None] = (
             replacement_hooks
         )
 
@@ -210,6 +209,94 @@ def code(self) -> str:
         )
         return self._code
 
+    def _replace_placeholder(self, hook_key: str, result: str) -> str:
+        """Replace all occurrences of a placeholder in the code with the hook result.
+
+        During Jinja template rendering, hooks like ``def_kernel()``,
+        ``load_input()``, and ``store_output()`` emit placeholder strings
+        (e.g. ``<DEF_KERNEL>``, ``<LOAD_INPUT_A>``, ``<STORE_OUTPUT_0>``)
+        into the rendered code and register corresponding hook functions in
+        ``self.replacement_hooks``.  When ``finalize_hook()`` is called for
+        each key, the hook function runs and this method replaces every
+        occurrence of the placeholder with the result.
+
+        Three replacement modes based on how each placeholder appears in
+        the code:
+
+        **Non-empty result, whole-line placeholder** — indent-propagating
+        substitution.
+            When the placeholder is the only non-whitespace content on a
+            line, the entire line is replaced.  The placeholder line's
+            leading whitespace (indent) is prepended to each result line
+            that has no leading whitespace of its own; lines that are
+            already indented are kept as-is.  This handles both hooks that
+            return results at uniform indent 0 (e.g.
+            ``ExternalTritonTemplateKernel``) and hooks that use the
+            ``strip()`` convention (first line un-indented, subsequent
+            lines pre-indented to the target level).
+
+        **Non-empty result, inline placeholder** — direct substitution
+        (``str.replace``).
+            When the placeholder appears mid-line (e.g.
+            ``{{gen_argdefs()}},``), all occurrences on that line are
+            replaced with the result verbatim via ``str.replace``.
+
+        **Empty result** — line removal.
+            Every line whose only non-whitespace content is the placeholder
+            is removed.  This handles cases like:
+
+            - ``<DEF_KERNEL>`` returning ``""`` for
+              ``ExternalTritonTemplateKernel`` (Helion), which emits its
+              own kernel definition outside the template system.
+            - ``<LOAD_INPUT_A>`` returning ``""`` when prologue fusion
+              replaces the explicit load with inline computation from a
+              fused producer.
+            - ``<STORE_OUTPUT_0>`` returning ``""`` for
+              ``ExternalTritonTemplateKernel`` outputs that have no
+              matching epilogue consumer.
+        """
+        if hook_key not in self._code:
+            return self._code
+
+        # Empty result — remove every line that contains only the placeholder
+        if not (result and result.strip()):
+            lines = self._code.split("\n")
+            return "\n".join(line for line in lines if line.strip() != hook_key)
+
+        # Non-empty result — line-by-line replacement
+        lines = self._code.split("\n")
+        new_lines = []
+        for line in lines:
+            if line.strip() == hook_key:
+                # Whole-line placeholder: decide how to indent the result.
+                indent = line[: len(line) - len(line.lstrip())]
+                result_lines = result.strip("\n").split("\n")
+                non_empty = [rl for rl in result_lines if rl.strip()]
+                all_unindented = bool(non_empty) and all(
+                    not rl[0].isspace() for rl in non_empty
+                )
+                if all_unindented:
+                    # Result is at uniform indent 0 (e.g.
+                    # ExternalTritonTemplateKernel hooks) — apply the
+                    # placeholder indent to every non-empty line.
+                    indented = [
+                        indent + rl if rl.strip() else rl for rl in result_lines
+                    ]
+                    new_lines.append("\n".join(indented).rstrip())
+                else:
+                    # Result has internal indentation (e.g. hooks using the
+                    # strip() convention, or <DEF_KERNEL> with function
+                    # bodies) — fall back to str.replace which prepends
+                    # the placeholder line's whitespace to the first line
+                    # only.
+                    new_lines.append(line.replace(hook_key, result))
+            elif hook_key in line:
+                # Inline placeholder: simple substitution
+                new_lines.append(line.replace(hook_key, result))
+            else:
+                new_lines.append(line)
+        return "\n".join(new_lines)
+
     def finalize_hook(self, hook_key: str, strict: bool = True) -> None:
         """
         Finalize a hook by name.
@@ -228,7 +315,7 @@ def finalize_hook(self, hook_key: str, strict: bool = True) -> None:
 
         hook = self.replacement_hooks[hook_key]
         assert hook is not None, f"Hook key {hook_key} can only be called once"
-        self._code = self._code.replace(hook_key, hook())
+        self._code = self._replace_placeholder(hook_key, hook())
 
         self.replacement_hooks[hook_key] = None
 
@@ -264,19 +351,24 @@ def finalize_all(self) -> str:
 @dataclasses.dataclass()
 class SubgraphInfo:
     body: IndentedBuffer
-    template_mask: Optional[str] = None
-    template_out_shape: Optional[Union[str, tuple[str]]] = None
+    template_mask: str | None = None
+    template_out_shape: str | tuple[str] | None = None
     compute: IndentedBuffer = dataclasses.field(default_factory=IndentedBuffer)
     indexing_code: IndentedBuffer = dataclasses.field(default_factory=IndentedBuffer)
     loads: IndentedBuffer = dataclasses.field(default_factory=IndentedBuffer)
     stores: IndentedBuffer = dataclasses.field(default_factory=IndentedBuffer)
-    ops_handler: Optional[V.WrapperHandler] = None  # type: ignore[name-defined]
+    ops_handler: V.WrapperHandler | None = None  # type: ignore[name-defined]
     cse: Optional["CSE[Any]"] = None
 
     # only copied over if not None
-    range_trees: Optional[list["IterationRangesRoot"]] = None
-    range_tree_nodes: Optional[dict[sympy.Symbol, "IterationRangesEntry"]] = None
-    numels: Optional[dict[str, sympy.Expr]] = None
+    range_trees: list["IterationRangesRoot"] | None = None
+    range_tree_nodes: dict[sympy.Symbol, "IterationRangesEntry"] | None = None
+    numels: dict[str, sympy.Expr] | None = None
+
+    # Mapping from original range-tree root variable names (e.g. "xindex")
+    # to renamed prologue variables (e.g. "_prologue_x_xindex").  Used by
+    # prologue hooks to apply text-level renames in a structured way.
+    root_var_renames: dict[str, str] = dataclasses.field(default_factory=dict)
 
     def __post_init__(self):
         self.only_copy_if_non_none_fields = (
@@ -300,8 +392,8 @@ def __init__(
         kernel,
         subgraph_number: int,
         fixed_inputs: dict[str, Any],
-        mask: Optional[str],
-        input_shapes: Optional[dict[str, tuple[str, ...]]] = None,
+        mask: str | None,
+        input_shapes: dict[str, tuple[str, ...]] | None = None,
     ):
         super().__init__(V.ops)
         self.name = f"PlaceholderSubstitution_{subgraph_number}"
@@ -414,11 +506,13 @@ def __init__(
         prefix_args=0,
         suffix_args=0,
         epilogue_fn=identity,
-        subgraphs: Optional[list[ir.ComputedBuffer]] = None,
-        workspace_arg: Optional[WorkspaceArg] = None,
+        subgraphs: list[ir.ComputedBuffer] | None = None,
+        workspace_arg: WorkspaceArg | None = None,
         prologue_loads_all_inputs=False,
-        hint_override: Optional[int] = None,
-        triton_meta: Optional[dict[str, object]] = None,
+        hint_override: int | None = None,
+        triton_meta: dict[str, object] | None = None,
+        always_freeze_layout: bool = False,
+        index_dtype_override: str | None = None,
     ) -> None:
         if tma_store:
             pass
@@ -488,9 +582,10 @@ def __init__(
         # pyrefly: ignore [invalid-type-var]
         self.epilogue_fn = epilogue_fn
         self.render_hooks = {}  # type: ignore[var-annotated]
-        self.triton_meta: Optional[dict[str, object]] = triton_meta
+        self.triton_meta: dict[str, object] | None = triton_meta
+        self._index_dtype_override = index_dtype_override
         # For Templated Attention this can be a list of ir.Subgraph
-        self.subgraphs: Optional[list[ir.ComputedBuffer]] = subgraphs
+        self.subgraphs: list[ir.ComputedBuffer] | None = subgraphs
 
         # Some templates use extra global memory as a workspace
         self.workspace_arg = workspace_arg
@@ -520,9 +615,10 @@ def __init__(
         self.indexing_code: IndentedBuffer = FakeIndentedBuffer()
         self.loads: IndentedBuffer = FakeIndentedBuffer()
         self.stores: IndentedBuffer = FakeIndentedBuffer()
-        self.template_mask: Optional[str] = None
-        self.template_out_shape: Optional[Union[str, tuple[str]]] = None
-        self.ops_handler: Optional[V.WrapperHandler] = None  # type: ignore[name-defined]
+        self.template_mask: str | None = None
+        self.template_out_shape: str | tuple[str] | None = None
+        self.ops_handler: V.WrapperHandler | None = None  # type: ignore[name-defined]
+        self.root_var_renames: dict[str, str] = {}
 
         # When caching is enabled, the generated code is not dependent on the input nodes names, or
         # symbolic sizes names.
@@ -532,7 +628,7 @@ def __init__(
         # input names or symbol names, we do a record and replay method.
         # During template expansions we record all function calls that change input_dependent_preserved_state
         # and replay them on a cache hit to regenerate them.
-        self.cached_replay_events: Optional[RecordedEventsType] = None
+        self.cached_replay_events: RecordedEventsType | None = None
 
         # Update each time an input is marked frozen, used to replay the freezing of inputs on a cache hit.
         self.frozen_layouts_cnt = 0
@@ -541,12 +637,23 @@ def __init__(
         # by adding all inputs.
         self.prologue_loads_all_inputs = prologue_loads_all_inputs
 
+        # When always_freeze_layout is True, get_stride_and_maybe_freeze_layout will
+        # always freeze the layout immediately, bypassing layout constraints.
+        # This is used by FlexAttention templates which require frozen layouts.
+        self.always_freeze_layout = always_freeze_layout
+
         # Extra functions to be exposed during partial template rendering.
         self.extra_template_env_fns: list[Callable[..., Any]] = []
 
         # Tracking for intermediate variables
         self.tmp_var_ctr = itertools.count()
 
+    @property
+    def index_dtype(self) -> str:
+        if self._index_dtype_override is not None:
+            return self._index_dtype_override
+        return super().index_dtype
+
     def _gen_tmp_var(self) -> str:
         return f"_tmp_var{next(self.tmp_var_ctr)}"
 
@@ -626,6 +733,75 @@ def create_subgraph_body(self, body_name: str, clear_cse: bool = False):
         with self.set_subgraph_body(body_name):
             yield
 
+    def _make_independent_subgraph(self, subgraph_name, numel, **extra_fields):
+        """Create a subgraph with fresh independent range trees.
+
+        Used by external template backends for epilogue/prologue hooks
+        that need their own range tree state.
+        """
+        groups = {"x": V.graph.sizevars.simplify(numel), "r0_": sympy.S.One}
+        self.subgraph_bodies[subgraph_name] = SubgraphInfo(
+            body=IndentedBuffer(),
+            cse=self.cse.clone(),
+            range_trees=self.construct_range_trees(
+                pid_cache=None,
+                inside_reduction=False,
+                is_reduction=False,
+                numels=groups,
+                no_x_dim=False,
+            ),
+            range_tree_nodes={},
+            numels=groups,
+            **extra_fields,
+        )
+
+    def _setup_contiguous_index_state(
+        self,
+        indices: list[str],
+        index_symbols: list[sympy.Symbol],
+        lengths: list[sympy.Expr],
+        mask: str | None,
+        xindex_name: str = "xindex",
+    ) -> sympy.Expr:
+        """Set up range trees, contiguous index, mask, and template indices.
+
+        Common boilerplate shared by store_output, _setup_epilogue_hook,
+        and load_input.  Returns the contiguous index expression.
+        """
+        for name, range_tree_entry in zip(
+            indices, self.range_trees[0].construct_entries(lengths)
+        ):
+            range_tree_entry.set_name(name)
+        contiguous_index = sympy_dot(
+            ir.FlexibleLayout.contiguous_strides(lengths), index_symbols
+        )
+        contiguous_index = self.rename_indexing(contiguous_index)
+        self.body.writeline(f"{xindex_name} = " + texpr(contiguous_index))
+        self.range_trees[0].lookup(sympy.S.One, sympy_product(lengths)).set_name(
+            xindex_name
+        )
+        self.template_mask = mask
+        self.template_indices = indices
+        return contiguous_index
+
+    def _make_codegen_hook(
+        self, subgraph_name: str, indent_width: int = 0
+    ) -> Callable[[], str]:
+        """Create a hook closure that codegen's a subgraph body."""
+
+        def hook():
+            with self.set_subgraph_body(subgraph_name):
+                self.codegen_body()
+                self.cse.invalidate(OrderedSet())
+                result = self.body.getvalue()
+                if indent_width:
+                    result = textwrap.indent(
+                        textwrap.dedent(result), " " * indent_width
+                    )
+                return result.strip()
+
+        return hook
+
     def need_numel_args(self):
         return False
 
@@ -819,10 +995,12 @@ def hook():
 
         return self._register_hook("<DEF_KERNEL>", hook)
 
-    def size(self, name: Optional[str], index: int):
+    def size(self, name: str | None, index: int):
         """
         Hook called from template code to get the size of an arg.
         Will add needed args to pass it in if it is dynamic.
+        Automatically wraps with tl.full([], ..., dtype=INDEX_DTYPE) when
+        int64 indexing is needed to prevent overflow in size arithmetic.
         """
         assert isinstance(index, int)
         if name is None:
@@ -830,7 +1008,10 @@ def size(self, name: Optional[str], index: int):
         else:
             assert isinstance(name, str)
             val = self.named_input_nodes[name].get_size()[index]
-        return texpr(self.rename_indexing(val))
+        result = texpr(self.rename_indexing(val))
+        if self.index_dtype == "tl.int64":
+            return f"tl.full([], {result}, dtype=INDEX_DTYPE)"
+        return result
 
     def stride(self, name, index=None):
         """
@@ -881,9 +1062,9 @@ def contiguous_strides(x):
     def modification(
         self,
         subgraph_number: int,
-        output_name: Optional[str],
-        mask: Optional[str] = None,
-        input_shapes: Optional[dict[str, tuple[str, ...]]] = None,
+        output_name: str | None,
+        mask: str | None = None,
+        input_shapes: dict[str, tuple[str, ...]] | None = None,
         **fixed_inputs,
     ) -> str:
         """This creates a modification function for a subgraph.
@@ -938,11 +1119,11 @@ def load_input(
         self,
         input_name: str,
         output_name: str,
-        indices: Union[list[Any], tuple[Any]],
-        mask: Optional[str] = None,
-        other: Optional[Union[float, int]] = 0.0,
+        indices: list[Any] | tuple[Any],
+        mask: str | None = None,
+        other: float | int | None = 0.0,
         indent_width: int = 4,
-        index_shape: Optional[tuple[str]] = None,
+        index_shape: tuple[str] | None = None,
     ):
         """Loads an input and applies any necessary preprocessing or masking.
 
@@ -986,42 +1167,16 @@ def load_input(
             lengths = [V.graph.sizevars.simplify(s) for s in input_node.get_size()]
             assert len(indices) == len(lengths)
 
-            index_symbols = [sympy.Symbol(x, integer=True) for x in indices]
-            assert len(indices) == len(lengths)
-
-            # glue to make generated code use same indexing from template
-
-            # TODO (from reviewers as well)
-            # in codegen_template,
-            # prologue_node.codegen(kernel.split_and_set_ranges(prologue_node.get_ranges()))
-            # the ranges need to reflect the group of the prologue input or it will error
-            # not sure if there is any difference between original range_tree_entry in
-            # and new one from correct lengths/groups... both actually seem to work
-            for name, range_tree_entry in zip(
-                indices, self.range_trees[0].construct_entries(lengths)
-            ):
-                range_tree_entry.set_name(name)
-            contiguous_index = sympy_dot(
-                ir.FlexibleLayout.contiguous_strides(lengths), index_symbols
-            )
-            contiguous_index = self.rename_indexing(contiguous_index)
-            self.body.writeline("xindex = " + texpr(contiguous_index))
-
-            xindex_range_root = self.range_trees[0].lookup(
-                sympy.Integer(1), sympy_product(lengths)
+            # MM templates use out-of-bounds wrapping (e.g. `rm % M`) so no mask
+            # is needed on the load.  Pass "None" when mask is unset to override
+            # the mask that would otherwise be inherited.
+            contiguous_index = self._setup_contiguous_index_state(
+                indices,
+                index_symbols,
+                lengths,
+                mask=mask if mask is not None else "None",
             )
-            xindex_range_root.set_name("xindex")
-
-            # Note - ["None" override_mask]
-            # MM Templates work by taking out of bounds index values and wrapping them around to 0
-            # so that no mask is required on the load: offs_a_m = `rm % M`
-            # We should to override the mask to be "None" instead of inheriting the mask that would
-            # have been loaded otherwise.
-            # We are using "None" for clarity in output code, but
-            # we could alternatively emit `xmask = tl.full([xindex.shape], True, tl.int1)`
-            self.template_mask = mask if mask is not None else "None"
             self.template_out_shape = index_shape if index_shape else "xindex"
-            self.template_indices = indices
             self.cse.invalidate(OrderedSet())
 
             template_mask = self.template_mask
@@ -1118,7 +1273,10 @@ def hook():
                     assert load_code is not None
                     self.body.writeline(load_code)
 
-                return textwrap.indent(self.body.getvalue(), " " * indent_width).strip()
+                result = self.body.getvalue()
+                if indent_width:
+                    result = textwrap.indent(result, " " * indent_width)
+                return result.strip()
 
         return self._register_hook(hook_key, hook)
 
@@ -1130,7 +1288,7 @@ def _generate_index_from_tma_index(
         block_size: str,
         dim: int,
         num_dims: int,
-        block_name: Optional[str] = None,
+        block_name: str | None = None,
     ) -> list[str]:
         """
         Generate the logic to compute the regular tl.load index from the provided
@@ -1194,11 +1352,11 @@ def _generated_mask_for_tma(
 
     def store_output(
         self,
-        indices: Union[list[Any], tuple[Any]],
+        indices: list[Any] | tuple[Any],
         val: str,
-        mask: Optional[str] = None,
+        mask: str | None = None,
         indent_width: int = 4,
-        val_shape: Optional[tuple[str]] = None,
+        val_shape: tuple[str] | None = None,
         block_indexing: bool = False,
     ):
         """Stores the final output and appends any epilogue fusions if the buffer hasn't been optimized away.
@@ -1331,21 +1489,9 @@ def store_output(
                     self.body.writeline(line)
             else:
                 assert not self.tma_store, "TMA store requires block indexing"
-                # glue to make generated code use same indexing from template
-                for name, range_tree_entry in zip(
-                    indices, self.range_trees[0].construct_entries(lengths)
-                ):
-                    range_tree_entry.set_name(name)
-                contiguous_index = sympy_dot(
-                    ir.FlexibleLayout.contiguous_strides(lengths), index_symbols
+                contiguous_index = self._setup_contiguous_index_state(
+                    indices, index_symbols, lengths, mask
                 )
-                contiguous_index = self.rename_indexing(contiguous_index)
-                self.body.writeline("xindex = " + texpr(contiguous_index))
-                self.range_trees[0].lookup(
-                    sympy.S.One, sympy_product(lengths)
-                ).set_name("xindex")
-                self.template_mask = mask
-                self.template_indices = indices
                 output_index = self.output_node.get_layout().make_indexer()(
                     index_symbols
                 )
@@ -1385,15 +1531,9 @@ def store_output(
             )
             self.codegen_body()
 
-        def hook():
-            with self.set_subgraph_body(subgraph_name):
-                # more stuff might have been added since the codegen_body above
-                self.codegen_body()
-                self.cse.invalidate(OrderedSet())
-
-                return textwrap.indent(self.body.getvalue(), " " * indent_width).strip()
-
-        return self._register_hook(subgraph_name, hook)
+        return self._register_hook(
+            subgraph_name, self._make_codegen_hook(subgraph_name, indent_width)
+        )
 
     def _register_hook(
         self,
@@ -1490,7 +1630,7 @@ def indexing(
         copy_shape=None,
         override_mask=None,
         block_ptr=False,
-        tma_compatibility_checker: Optional[TMACompatibilityChecker] = None,
+        tma_compatibility_checker: TMACompatibilityChecker | None = None,
     ):
         """
         Override the default indexing to use our custom mask and force
@@ -1522,7 +1662,7 @@ def additional_call_args_and_types(self):
         return ((), ())
 
     def call_kernel(
-        self, name: str, node: Optional[ir.IRNode] = None, deallocate_ws: bool = True
+        self, name: str, node: ir.IRNode | None = None, deallocate_ws: bool = True
     ):
         wrapper = V.graph.wrapper_code
         _, call_args, _, arg_types = self.args.python_argdefs()
@@ -1546,21 +1686,31 @@ def call_kernel(
 
         if self.workspace_arg is not None:
             wrapper.generate_workspace_allocation(self.workspace_arg)
+
+        # Use FixedGrid which properly handles grid values passed as arguments
+        inductor_meta = FixedGrid.setup_grid_as_args() if additional_call_args else None
         wrapper.generate_kernel_call(
             name,
             call_args,
             arg_types=arg_types,
             triton_meta=self.triton_meta,
+            inductor_meta=inductor_meta,
             triton=True,
         )
+        self._emit_post_kernel_code(wrapper, name)
         if self.workspace_arg is not None:
             wrapper.generate_workspace_deallocation(self.workspace_arg)
 
+    def _emit_post_kernel_code(self, wrapper, kernel_name: str) -> None:
+        """Hook for subclasses to emit code after kernel call, before workspace dealloc."""
+        pass  # noqa: PIE790
+
     def kernel_benchmark_extra_args(self) -> list[str]:
+        # Grid args are only used for benchmarking, not correctness
         return [
             str(x)
             for x in self.grid_fn(
-                *V.graph.sizevars.size_hints(self.call_sizes), self.meta
+                *V.graph.sizevars.optimization_hints(self.call_sizes), self.meta
             )
         ]
 
@@ -1575,14 +1725,22 @@ def get_stride_and_maybe_freeze_layout(self, node) -> list[int]:
         Scheduler falls back to aten if layout constraint violated. If no aten,
         freeze right away.
         """
+
         # realizing for safety
         ir.ExternKernel.realize_input(node)
         layout = node.data.layout
         node_name = node.get_name()
 
-        if isinstance(layout, ir.FlexibleLayout):
-            if not use_aten_gemm_kernels():
-                # No ExternKernel fallback available, freeze immediately
+        # For ReinterpretView, the view's strides are already determined by its layout.
+        # We skip constraint tracking because node.get_name() returns the underlying
+        # buffer name, not the view's identity, so constraints would be incorrectly
+        # associated with the underlying buffer rather than the view.
+        if isinstance(layout, ir.FlexibleLayout) and not isinstance(
+            node, ir.ReinterpretView
+        ):
+            if not use_aten_gemm_kernels() or self.always_freeze_layout:
+                # No ExternKernel fallback available, or always_freeze_layout is set
+                # (e.g., for FlexAttention templates), freeze immediately
                 node.data.freeze_layout()
             else:
                 # Compute what strides WOULD be if frozen, without actually freezing
@@ -1601,6 +1759,559 @@ def get_stride_and_maybe_freeze_layout(self, node) -> list[int]:
         # Already frozen or not a FlexibleLayout, just return current strides
         return node.get_stride()
 
+    def _compute_fusion_metadata(
+        self, scheduling, epilogue_nodes, prologue_nodes, buf_name_to_prologue_group
+    ):
+        """Prepare epilogue/prologue routing before render().
+
+        Default: trivial routing — all epilogues broadcast to every subgraph,
+        none unfused, no prologue source tracking.  Override in subclasses
+        for per-output routing.
+        """
+        self._epilogue_nodes_by_subgraph: defaultdict[int, list[Any]] = defaultdict(
+            lambda: epilogue_nodes
+        )
+        self._unfused_epilogues: list[Any] = []
+        self._prologue_sources: dict[str, frozenset[str]] = {}
+
+    def codegen_template_body(
+        self,
+        scheduling,
+        template_node,
+        epilogue_nodes,
+        prologue_nodes,
+        buf_name_to_prologue_group,
+        prologue_preserves_zero_mask_fn,
+        render,
+    ) -> str:
+        """Generate template source code with fused prologues and epilogues.
+
+        Returns the final source code string.
+        """
+        self._compute_fusion_metadata(
+            scheduling, epilogue_nodes, prologue_nodes, buf_name_to_prologue_group
+        )
+        with self:
+            partial_code = render()
+
+            num_store_subgraphs = self.get_store_output_count()
+            for i in range(num_store_subgraphs):
+                subgraph_name = self._get_store_output_subgraph_name(i)
+                with self.set_subgraph_body(subgraph_name):
+                    for node in self._epilogue_nodes_by_subgraph[i]:
+                        node.codegen(self.split_and_set_ranges(node.get_ranges()))
+                    self.cse.invalidate(OrderedSet())
+
+            self.codegen_prologues_in_subgraphs(
+                buf_name_to_prologue_group, prologue_preserves_zero_mask_fn
+            )
+
+        partial_code = self._finalize_partial_render(partial_code)
+
+        # Template hooks must be finalised after kernel.remove_kernel_local_buffers
+        # is called (this is called when the kernel context is exited above), and when
+        # the kernel handler is set (as below). This is because the hooks may add
+        # DeferredLine type lines, which preclude lines involving buffers that have
+        # been removed
+
+        # finalize must be called after adding epilogue above
+        with V.set_kernel_handler(self):
+            if isinstance(partial_code, str):
+                src_code = partial_code
+            else:
+                # This is used to calculate flops in TritonTemplateKernels
+                with ir.IRNode.current_origins(template_node.node.origins):
+                    partial_code.finalize_hook("<DEF_KERNEL>")
+                partial_code.finalize_hook("<ARGDEFS>", strict=False)
+
+                # TODO: Maybe unify CUTLASSTemplateKernel to also use PartialRender for flexible epilogue fusion.
+
+                for input_name in self.named_input_nodes:
+                    subgraph_name = f"<LOAD_INPUT_{input_name}>"
+                    partial_code.finalize_hook(subgraph_name, strict=False)
+
+                num_store_subgraphs = self.get_store_output_count()
+                for i in range(num_store_subgraphs):
+                    subgraph_name = self._get_store_output_subgraph_name(i)
+                    partial_code.finalize_hook(subgraph_name)
+
+                # Ensure all hooks are finalized before the kernel is defined.
+                # Note: some of these hooks may have been registered by a kernel subclass
+                src_code = partial_code.finalize_remaining()
+
+        return src_code
+
+    def codegen_prologues_in_subgraphs(
+        self, buf_name_to_prologue_group, prologue_preserves_zero_mask_fn
+    ):
+        """Run prologue codegen in each load-input subgraph body."""
+        for input_name, buffer in self.named_input_nodes.items():
+            subgraph_name = f"<LOAD_INPUT_{input_name}>"
+            prologue_group = buf_name_to_prologue_group.get(buffer.get_name(), [])
+            if not prologue_group:
+                continue
+            can_codegen_without_upcast = all(
+                p_n.can_codegen_without_upcasts() for p_n in prologue_group
+            )
+            with config.patch(
+                "triton.codegen_upcast_to_fp32", not can_codegen_without_upcast
+            ):
+                with self.set_subgraph_body(subgraph_name):
+                    for prologue_node in prologue_group:
+                        if (
+                            len(prologue_node.get_buffer_names()) == 1
+                            and len(prologue_group) == 1
+                        ):
+                            if prologue_preserves_zero_mask_fn(prologue_node):
+                                self.prologue_fused_inputs_preserve_zero |= (
+                                    prologue_node.get_buffer_names()
+                                )
+                        prologue_node.codegen(
+                            self.split_and_set_ranges(prologue_node.get_ranges())
+                        )
+                    self.cse.invalidate(OrderedSet())
+
+    def _finalize_partial_render(
+        self, partial_code: str | PartialRender
+    ) -> str | PartialRender:
+        """Hook to intercept or replace the PartialRender before hook finalization.
+
+        Called after Inductor has populated subgraph buffers (epilogue stores,
+        prologue loads) but before ``finalize_hook`` resolves placeholders.
+        Subclasses may return a replacement ``PartialRender`` — e.g. to capture
+        the rendered hook outputs and supply entirely new source code (as
+        ``ExternalTritonTemplateKernel`` does for fusion-aware autotuning).
+        """
+        return partial_code
+
+
+class ExternalTritonTemplateKernel(TritonTemplateKernel):
+    """TritonTemplateKernel variant for external template backends (e.g. Helion).
+
+    Orchestrates prologue/epilogue fusion by running Inductor's codegen ops
+    (store, load, indexing) on fused scheduler nodes, capturing the generated
+    code into subgraph buffers.  The backend's ``render()`` callable handles
+    hook setup and AST generation, while the standard ``codegen_template_body``
+    handles hook finalization and subgraph codegen.
+
+    Also implements ``call_kernel`` and ``emit_kernel_override`` for
+    post-codegen emission of the fused kernel.
+
+    Subclasses TritonTemplateKernel for subgraph body management and
+    template indexing support.
+    """
+
+    def __init__(self, template_buffer: "ir.TemplateBuffer") -> None:
+        class _RealOutputNode:
+            def get_size(self) -> list:
+                return list(template_buffer.get_size())
+
+            def get_layout(self):
+                return template_buffer.get_layout()
+
+            def get_name(self) -> str:
+                return template_buffer.get_name()
+
+        # Pass dummy values for TritonTemplateKernel params that are only
+        # relevant for standalone Triton kernel codegen (grid, warps, etc.).
+        super().__init__(
+            kernel_name="",
+            input_nodes=(),
+            output_node=_RealOutputNode(),
+            defines={},
+            num_stages=0,
+            num_warps=1,
+            grid_fn=None,
+            meta={},
+            call_sizes=[],
+            hint_override=None,
+        )
+        self._template_buffer = template_buffer
+        # Extra inputs needed by fused ops beyond the template's own I/O
+        self._extra_inputs: dict[str, str] = {}
+        # Prologue primary source buffers, populated by load_input
+        self._prologue_source_buffers: dict[str, str | None] = {}
+        # Simplified epilogue interface: {output_param: epilogue_idx}
+        self._epilogue_idx_by_param: dict[str, int] = {}
+        # Output params that must keep their original tl.store
+        self._epilogue_keep_store: OrderedSet[str] = OrderedSet()
+        # Store target buffers: {buf_name: param_name}
+        self._extra_store_targets: dict[str, str] = {}
+        # Prologue variable names per input param
+        self._prologue_vars: dict[str, dict[str, str]] = {}
+        # Import lines for emit_kernel_override, populated by external render
+        self._kernel_imports: list[str] = []
+        # Call emission state, populated by _setup_fusion_hooks / external render
+        self._call_preamble: list[str] = []
+        self._call_args: list[str] = []
+        # Epilogues that could not be fused into the kernel
+        self._unfused_epilogues: list[Any] = []
+        # Reference to the scheduler, set by _compute_fusion_metadata;
+        # used in call_kernel() to codegen unfused epilogue nodes
+        self._scheduling_ref: Any = None
+
+    def _finalize_partial_render(
+        self, partial_code: str | PartialRender
+    ) -> str | PartialRender:
+        # Capture hook outputs.
+        hook_outputs: dict[str, str] = {}
+        with V.set_kernel_handler(self):
+            for key, hook in self.render_hooks.items():
+                hook_outputs[key] = hook()
+
+        # Delegate to template buffer.
+        result = self._template_buffer._finalize_codegen(hook_outputs)
+        if result is None:
+            return partial_code
+
+        # Apply structured result to kernel state.
+        self._kernel_imports = result.imports
+        self._call_preamble = result.call_preamble
+        self._call_args = result.call_args
+
+        # Freeze hooks and clear subgraph buffers.
+        if hook_outputs:
+            self.render_hooks = {
+                k: (lambda captured=v: captured) for k, v in hook_outputs.items()
+            }
+            for info in self.subgraph_bodies.values():
+                info.loads = IndentedBuffer()
+                info.compute = IndentedBuffer()
+                info.stores = IndentedBuffer()
+                info.indexing_code = IndentedBuffer()
+
+        return PartialRender(result.source, self.render_hooks)
+
+    def get_unfused_epilogues(self) -> list[Any]:
+        return self._unfused_epilogues
+
+    def _compute_fusion_metadata(
+        self, scheduling, epilogue_nodes, prologue_nodes, buf_name_to_prologue_group
+    ):
+        """Compute fusion metadata for external backends.
+
+        Determines eligible epilogues/prologues, builds epilogue specs,
+        and computes prologue sources — all before render().
+
+        Hook setup (_setup_epilogue_hook / _setup_prologue_hook) cannot
+        happen here because it requires V.kernel context, which is only
+        active during codegen_template_body → render().
+        """
+        self._scheduling_ref = scheduling
+        from torch._inductor.dependencies import MemoryDep
+
+        tb = self._template_buffer
+        self._eligible_epilogues = self._find_eligible_epilogues(
+            epilogue_nodes, tb.epilogue_fusable_outputs
+        )
+        self._epilogue_nodes_by_subgraph = defaultdict(
+            list,
+            {i: [sn] for i, (sn, _, _, _) in enumerate(self._eligible_epilogues)},
+        )
+        fused_ids = OrderedSet(id(sn) for sn, _, _, _ in self._eligible_epilogues)
+        self._unfused_epilogues = [
+            n
+            for n in epilogue_nodes
+            if id(n) not in fused_ids and not isinstance(n.node, ir.MultiOutput)
+        ]
+        self._prologue_sources = {
+            buf_name: frozenset(
+                d.name for d in pro_node.read_writes.reads if isinstance(d, MemoryDep)
+            )
+            for buf_name, pro_nodes in buf_name_to_prologue_group.items()
+            for pro_node in pro_nodes
+        }
+
+        # Build simplified epilogue interface: _epilogue_idx_by_param,
+        # _epilogue_keep_store, and _extra_store_targets.
+        from torch._inductor.codegen.common import RemovedArg
+
+        scheduler = V.graph.scheduler
+        epilogues = self._eligible_epilogues
+
+        # Compute fused node names for buffer removability
+        fused_node_names = None
+        if scheduler is not None:
+            all_store_names = OrderedSet([tb.get_name()])
+            all_store_names.update(tb._multi_output_children)
+            all_store_names.update(st for _, _, _, st in epilogues if st)
+            fused_node_names = OrderedSet(
+                scheduler.name_to_buf[n].defining_op_name()
+                for n in all_store_names
+                if n in scheduler.name_to_buf
+            )
+
+        # Pre-register store_target buffers so we know their param names
+        for _, _, _, store_target in epilogues:
+            if (
+                store_target is not None
+                and store_target not in self.args.output_buffers
+            ):
+                self.args.output(store_target)
+
+        # Build per-epilogue metadata
+        for i, (_, output_buf, output_param, store_target) in enumerate(epilogues):
+            self._epilogue_idx_by_param[output_param] = i
+
+            can_remove = (
+                store_target is not None
+                and fused_node_names is not None
+                and scheduler.can_buffer_be_removed_through_fusion(
+                    output_buf, fused_node_names
+                )
+            )
+
+            store_target_param_raw = (
+                self.args.output_buffers.get(store_target)
+                if store_target is not None
+                else None
+            )
+            store_target_param = (
+                None
+                if isinstance(store_target_param_raw, RemovedArg)
+                else store_target_param_raw
+            )
+
+            if store_target_param is not None:
+                self._extra_store_targets[store_target] = store_target_param
+                if can_remove:
+                    self.removed_buffers.add(output_buf)
+                else:
+                    self._epilogue_keep_store.add(output_param)
+
+    def _setup_fusion_hooks(self):
+        """Set up epilogue/prologue render hooks and mark prologue buffers.
+
+        Must be called during render() (inside V.kernel context), after
+        _compute_fusion_metadata has run.
+        """
+        tb = self._template_buffer
+
+        # Mark prologue buffers on the kernel.
+        for pro_buf, source_bufs in self._prologue_sources.items():
+            self.store_buffer_names.add(pro_buf)
+            if not source_bufs:
+                self.removed_buffers.add(pro_buf)
+
+        # Set up epilogue hooks.
+        epilogues = self._eligible_epilogues
+        for epilogue_idx in range(len(tb.epilogue_fusable_outputs)):
+            epi = epilogues[epilogue_idx] if epilogue_idx < len(epilogues) else None
+            self._setup_epilogue_hook(
+                output_buf=epi[1] if epi else None,
+                output_param=epi[2] if epi else None,
+            )
+
+        # Set up prologue hooks.
+        for param_name in tb._named_inputs:
+            self._setup_prologue_hook(
+                param_name, prologue_sources=self._prologue_sources
+            )
+
+        # Register no-op <DEF_KERNEL> hook (standard path requires it).
+        self.render_hooks["<DEF_KERNEL>"] = lambda: ""
+
+    def _find_eligible_epilogues(self, epilogue_nodes, output_param_mapping):
+        """Compute fusion eligibility and register extra inputs.
+
+        Returns list of eligible epilogue tuples:
+            [(snode, output_buf, output_param, store_target), ...]
+        """
+        from torch._inductor.dependencies import MemoryDep
+
+        # Filter eligible epilogues
+        epilogues = []
+        for epilogue_node in epilogue_nodes:
+            if isinstance(epilogue_node.node, ir.MultiOutput):
+                continue
+            dep_names = OrderedSet(
+                d.name
+                for d in epilogue_node.read_writes.reads
+                if isinstance(d, MemoryDep) and d.name in output_param_mapping
+            )
+            if len(dep_names) != 1:
+                continue
+            output_buf = next(iter(dep_names))
+            epilogue_writes = epilogue_node.read_writes.writes
+            raw_st = next(iter(epilogue_writes)).name if epilogue_writes else None
+            epilogues.append(
+                (
+                    epilogue_node,
+                    output_buf,
+                    output_param_mapping[output_buf],
+                    raw_st if raw_st != output_buf else None,
+                )
+            )
+
+        # Register extra inputs needed by fused epilogues
+        for snode, _, _, _ in epilogues:
+            for dep in snode.read_writes.reads:
+                if isinstance(dep, MemoryDep) and dep.name not in output_param_mapping:
+                    if dep.name not in self._extra_inputs:
+                        param = f"_extra_input_{len(self._extra_inputs)}"
+                        self._extra_inputs[dep.name] = param
+                        self.args.input_buffers[dep.name] = param
+
+        return epilogues
+
+    def _setup_epilogue_hook(self, output_buf=None, output_param=None):
+        store_idx = next(self.store_output_ctr)
+        subgraph_name = self._get_store_output_subgraph_name(store_idx)
+        if output_buf is None:
+            self.subgraph_bodies[subgraph_name] = SubgraphInfo(body=IndentedBuffer())
+            self.render_hooks[subgraph_name] = lambda: ""
+            return
+
+        n_dims = len(self._template_buffer.get_size())
+        indices = [f"x_epilogue{store_idx}_{d}" for d in range(n_dims)]
+        val = f"_kernel_val_{store_idx}"
+        mask = f"_tile_mask_{store_idx}"
+
+        buf = output_buf
+        node = V.graph.get_buffer(buf) if buf else None
+        output_size = (
+            list(node.get_size())
+            if node is not None
+            else list(self.output_node.get_size())
+        )
+        self._make_independent_subgraph(subgraph_name, sympy_product(output_size))
+        with self.set_subgraph_body(subgraph_name):
+            indices = list(map(OpOverrides.paren, indices))
+            index_symbols = [sympy.Symbol(x, integer=True) for x in indices]
+            lengths = [V.graph.sizevars.simplify(s) for s in output_size]
+            assert len(indices) == len(lengths)
+            self.template_out = val
+            self._setup_contiguous_index_state(
+                indices,
+                index_symbols,
+                lengths,
+                mask,
+                xindex_name=f"x_epilogue{store_idx}_index",
+            )
+            self.template_out_shape = val
+
+            # Set up CSE state for epilogue codegen
+            block_shape = tuple(
+                f"{rt.prefix.upper()}BLOCK"
+                for rt in self.range_trees
+                if not rt.is_reduction
+            )
+            if not block_shape:
+                block_shape = ("XBLOCK",)
+            self.cse.store_cache[buf] = self.cse.namedvar(
+                val, dtype=torch.float32, shape=block_shape
+            )
+            assert output_param is not None
+            self.args.output_buffers[buf] = output_param
+
+        self.render_hooks[subgraph_name] = self._make_codegen_hook(subgraph_name)
+
+    def _setup_prologue_hook(self, input_name, prologue_sources=None):
+        ir_node = self._template_buffer._named_inputs.get(input_name)
+        if ir_node is None:
+            return
+        self.named_input_nodes[input_name] = ir_node
+        input_buf = ir_node.get_name()
+        source_bufs = (
+            prologue_sources.get(input_buf) if prologue_sources is not None else None
+        )
+        if source_bufs is None:
+            return
+        source_buffer = next(iter(source_bufs)) if source_bufs else None
+        self._prologue_source_buffers[input_name] = source_buffer
+        if source_buffer is not None:
+            self.args.input_buffers[source_buffer] = input_name
+        for src in source_bufs:
+            if src != source_buffer:
+                if src not in self._extra_inputs:
+                    self._extra_inputs[src] = f"_extra_input_{len(self._extra_inputs)}"
+                self.args.input_buffers[src] = self._extra_inputs[src]
+
+        subgraph_name = f"<LOAD_INPUT_{input_name}>"
+        result_var = f"_prologue_{input_name}_result"
+
+        # Compute prologue variable names once and store them.
+        renames = {
+            "xindex": f"_prologue_{input_name}_xindex",
+            "xmask": f"_prologue_{input_name}_xmask",
+        }
+        self._prologue_vars[input_name] = {
+            "xindex": renames["xindex"],
+            "xmask": renames["xmask"],
+            "result": result_var,
+        }
+
+        class _CaptureStoreHandler(V.WrapperHandler):  # type: ignore[name-defined]
+            def store(self, name, index, value, mode=None):
+                V.kernel.store_buffer_names.add(name)
+                V.kernel.cse.store_cache[name] = value
+                V.kernel.compute.writeline(f"{result_var} = {value}")
+
+        self._make_independent_subgraph(
+            subgraph_name,
+            sympy_product(ir_node.get_size()),
+            ops_handler=_CaptureStoreHandler,
+            root_var_renames=renames,
+        )
+
+        def hook(_name=subgraph_name, _input=input_name, _self=self):
+            with _self.set_subgraph_body(_name):
+                _self.codegen_body()
+                _self.cse.invalidate(OrderedSet())
+                body = _self.body.getvalue()
+            # Rename range-tree root variables (xindex/xmask) to avoid collisions
+            # across prologue subgraphs.  tmp/x0 names are already unique (shared
+            # kernel-level counters).  Read renames from the subgraph info.
+            subgraph = _self.subgraph_bodies[_name]
+            for orig, renamed in subgraph.root_var_renames.items():
+                body = re.sub(rf"\b{orig}\b", renamed, body)
+            return body.rstrip()
+
+        self.render_hooks[f"<LOAD_INPUT_{input_name}>"] = hook
+
+    def call_kernel(self, name, node=None, deallocate_ws=True):
+        """Emit the kernel call, multi-output unpacking, and unfused epilogues."""
+        tb = self._template_buffer
+        wrapper = V.graph.wrapper_code
+        for line in self._call_preamble:
+            wrapper.writeline(line)
+        output_name = tb.get_name()
+        wrapper.writeline(f"{output_name} = {name}({', '.join(self._call_args)})")
+        # Unpack multi-output children from the kernel result
+        for mo_name, mo in sorted(tb._multi_output_children.items()):
+            if mo_name not in self.removed_buffers:
+                idx_str = output_name
+                for _, idx in mo.indices:
+                    idx_str = f"{idx_str}[{idx}]"
+                wrapper.writeline(f"{mo_name} = {idx_str}")
+        # Unfused epilogues are codegen'd separately after the kernel call
+        for epi_node in self._unfused_epilogues:
+            self._scheduling_ref.codegen_node(epi_node)
+
+    def emit_kernel_override(
+        self,
+        wrapper,
+        src_code,
+        kernel_name,
+        node_schedule,
+        kernel_path,
+        get_kernel_metadata,
+    ):
+        origins, detailed = get_kernel_metadata(node_schedule, wrapper)
+        wrapper.header.writeline(f"# kernel path: {kernel_path}\n{origins}\n{detailed}")
+        for imp in self._kernel_imports:
+            wrapper.add_import_once(imp)
+        # Hooks may reference standard Inductor runtime modules
+        wrapper.add_import_once(
+            "from torch._inductor.runtime import triton_helpers, triton_heuristics"
+        )
+        wrapper.add_import_once(
+            "from torch._inductor.runtime.triton_helpers import libdevice, math as tl_math"
+        )
+        wrapper.header.splice(src_code, strip=True)
+        wrapper.header.writeline("")
+        return True
+
 
 @functools.cache
 def _jinja2_env():
@@ -1657,19 +2368,19 @@ def make_key(
         call_sizes: Sequence[sympy.core.symbol.Symbol],
         prefix_args: int,
         suffix_args: int,
-        epilogue_fn: Optional[Callable[..., Any]],
-        epilogue_fn_hash: Optional[str],
+        epilogue_fn: Callable[..., Any] | None,
+        epilogue_fn_hash: str | None,
         tma_store: bool,
-        transpose_discontiguous_tensor_descriptors_override: Optional[bool],
-        subgraphs: Optional[list[ir.Buffer]],  # has to be none to cache
-        workspace_arg: Optional[WorkspaceArg],  # has to be none to cache
+        transpose_discontiguous_tensor_descriptors_override: bool | None,
+        subgraphs: list[ir.Buffer] | None,  # has to be none to cache
+        workspace_arg: WorkspaceArg | None,  # has to be none to cache
         layout: ir.Layout,
         num_consumer_groups: int,
         num_buffers_warp_spec: int,
         kwargs: dict[str, Any],
-        hint_override: Optional[int] = None,
-        triton_meta: Optional[dict[str, Any]] = None,
-    ) -> Optional[str]:
+        hint_override: int | None = None,
+        triton_meta: dict[str, Any] | None = None,
+    ) -> str | None:
         def layout_key(layout: ir.Layout) -> str:
             assert not isinstance(layout, ir.FlexibleLayout)
             return repr(
@@ -1726,7 +2437,7 @@ def has_flexible_layout() -> bool:
             }
         )
 
-    def get_entry(self, cache_key: Optional[str]) -> Optional[GeneratedCodeCacheEntry]:
+    def get_entry(self, cache_key: str | None) -> GeneratedCodeCacheEntry | None:
         if cache_key is None:
             return None
 
@@ -1739,7 +2450,7 @@ def get_entry(self, cache_key: Optional[str]) -> Optional[GeneratedCodeCacheEntr
 
     def put_entry(
         self,
-        cache_key: Optional[str],
+        cache_key: str | None,
         code: str,
         extra: str,
         events: list[Any],
@@ -1768,6 +2479,7 @@ def __init__(
         debug=False,
         cache_codegen_enabled_for_template=False,
         prologue_loads_all_inputs=False,
+        always_freeze_layout: bool = False,
     ) -> None:
         super().__init__(name, hash=hashlib.sha256(source.encode("utf-8")).hexdigest())
         self.grid = grid
@@ -1781,6 +2493,10 @@ def __init__(
         # When prologue_loads_all_inputs is true, prologue_supported_inputs is populated during def_kernel
         # by adding all inputs.
         self.prologue_loads_all_inputs = prologue_loads_all_inputs
+        # When always_freeze_layout is True, the kernel will always freeze layouts
+        # immediately instead of using layout constraints. This is used by
+        # FlexAttention templates which require frozen layouts.
+        self.always_freeze_layout = always_freeze_layout
 
     # When this flag is on, we ensure that the cached results and the generated result if cache
     # was not used are the same.
@@ -1793,7 +2509,7 @@ def uid(self) -> str:
 
     def maybe_append_choice(
         self, choices: list[Any], **kwargs: Any
-    ) -> Optional[NotImplementedError]:
+    ) -> NotImplementedError | None:
         """
         Maybe generates a new ChoiceCaller and appends it into existing choices.
         Returns None if success, otherwise returns the error.
@@ -1825,20 +2541,20 @@ def generate_and_load(
         call_sizes: Sequence[sympy.core.symbol.Symbol],
         prefix_args: int,
         suffix_args: int,
-        epilogue_fn: Optional[Callable[..., Any]],
-        epilogue_fn_hash: Optional[str],
-        subgraphs: Optional[list[ir.Buffer]],
-        workspace_arg: Optional[WorkspaceArg],
+        epilogue_fn: Callable[..., Any] | None,
+        epilogue_fn_hash: str | None,
+        subgraphs: list[ir.Buffer] | None,
+        workspace_arg: WorkspaceArg | None,
         num_consumer_groups: int,
         num_buffers_warp_spec: int,
         layout: ir.Layout,
         kwargs: dict[str, Any],
         generate_with_caching,
-        hint_override: Optional[int] = None,
+        hint_override: int | None = None,
         tma_store: bool = False,
-        transpose_discontiguous_tensor_descriptors_override: Optional[bool] = None,
-        triton_meta: Optional[dict[str, Any]] = None,
-    ) -> Optional[GenerateAndLoadResult]:
+        transpose_discontiguous_tensor_descriptors_override: bool | None = None,
+        triton_meta: dict[str, Any] | None = None,
+    ) -> GenerateAndLoadResult | None:
         """Generate the python code and load it into the current process"""
         caching_enabled = (
             generate_with_caching
@@ -1902,6 +2618,8 @@ def generate_and_load(
             "epilogue_fn": epilogue_fn,
             "subgraphs": subgraphs,
             "prologue_loads_all_inputs": self.prologue_loads_all_inputs,
+            "always_freeze_layout": self.always_freeze_layout,
+            "index_dtype_override": index_dtype,
         }
 
         if HAS_WARP_SPEC:
@@ -1925,7 +2643,7 @@ def make_kernel():
                 **kernel_options,
             )
 
-        def generate_code(kernel) -> Optional[tuple[str, str]]:
+        def generate_code(kernel) -> tuple[str, str] | None:
             def make_extra() -> str:
                 extra_parts = [
                     f"{kwarg}={repr(kwargs[kwarg])}" for kwarg in sorted(kwargs.keys())
@@ -1979,8 +2697,8 @@ def maybe_test_cache(code: str, extra: str, kernel):
                     ), "Generated code cache results in wrong output"
 
         # Generate code, extra.
-        code: Optional[str] = None
-        extra: Optional[str] = None
+        code: str | None = None
+        extra: str | None = None
         with (
             patch.object(V.graph, "get_dtype", self._fake_get_dtype(fake_out)),
             V.graph.set_current_device(layout.device),
@@ -2005,7 +2723,7 @@ def maybe_test_cache(code: str, extra: str, kernel):
 
         assert code is not None and extra is not None
 
-        mod = PyCodeCache.load(code, extra)
+        mod = PyCodeCache.load(code, extra, set_sys_modules=False)
 
         input_call_args = tuple(kernel.args.input_buffers.keys())
         prologue_supported_inputs = kernel.prologue_supported_inputs.copy()
@@ -2033,17 +2751,17 @@ def generate(  # type: ignore[override]
         num_buffers_warp_spec: int = 0,
         prefix_args: int = 0,
         suffix_args: int = 0,
-        epilogue_fn: Optional[Callable[..., Any]] = identity,
-        epilogue_fn_hash: Optional[str] = None,
-        subgraphs: Optional[list[ir.Buffer]] = None,
-        mutated_inputs: Optional[list[ir.IRNode]] = None,
-        call_sizes: Optional[Sequence[sympy.core.symbol.Symbol]] = None,
-        workspace_arg: Optional[WorkspaceArg] = None,
+        epilogue_fn: Callable[..., Any] | None = identity,
+        epilogue_fn_hash: str | None = None,
+        subgraphs: list[ir.Buffer] | None = None,
+        mutated_inputs: list[ir.IRNode] | None = None,
+        call_sizes: Sequence[sympy.core.symbol.Symbol] | None = None,
+        workspace_arg: WorkspaceArg | None = None,
         generate_with_caching=False,
-        hint_override: Optional[int] = None,
+        hint_override: int | None = None,
         tma_store: bool = False,
-        transpose_discontiguous_tensor_descriptors_override: Optional[bool] = None,
-        triton_meta: Optional[dict[str, Any]] = None,
+        transpose_discontiguous_tensor_descriptors_override: bool | None = None,
+        triton_meta: dict[str, Any] | None = None,
         **kwargs,
     ):
         """This function generates a TritonTemplateCaller
@@ -2132,11 +2850,12 @@ def generate(  # type: ignore[override]
 
         # Extract workspace metadata for async autotuning (don't create tensor here
         # as it can't be pickled for subprocess communication)
-        workspace_size_bytes: Optional[int] = None
+        workspace_size_bytes: int | None = None
         workspace_zero_fill = False
         workspace_args = []
         if workspace_arg is not None:
-            workspace_size_bytes = workspace_arg.count
+            ws_count = V.graph.sizevars.optimization_hint(workspace_arg.count)
+            workspace_size_bytes = ws_count * get_dtype_size(workspace_arg.dtype)
             workspace_zero_fill = (
                 workspace_arg.zero_mode != WorkspaceZeroMode.UNINITIALIZED
             )
@@ -2145,7 +2864,7 @@ def generate(  # type: ignore[override]
 
         options = result.kernel_options
 
-        def make_kernel_render(out_node, hint_override: Optional[int] = None):
+        def make_kernel_render(out_node, hint_override: int | None = None):
             assert result is not None
             # Create a new unique name for the workspace arg buffer for each render
             # to prevent buffer reuse of the same workspace arg
@@ -2210,6 +2929,21 @@ def make_kernel_render(out_node, hint_override: Optional[int] = None):
             output_tensor_meta=TensorMeta.from_irnodes(layout),
         )
 
+        # Convolution-specific parameters to include in logging
+        CONV_TUNABLE_KEYS = [
+            "KERNEL_H",
+            "KERNEL_W",
+            "KERNEL_D",
+            "STRIDE_H",
+            "STRIDE_W",
+            "STRIDE_D",
+            "PADDING_H",
+            "PADDING_W",
+            "PADDING_D",
+            "GROUPS",
+            "UNROLL",
+        ]
+
         return TritonTemplateCaller(
             kernel_hash_name,
             codegen_input_nodes,
@@ -2233,11 +2967,13 @@ def make_kernel_render(out_node, hint_override: Optional[int] = None):
                 "matrix_instr_nonkdim": kwargs.get("matrix_instr_nonkdim", 0),
                 "waves_per_eu": kwargs.get("waves_per_eu", 0),
                 "kpack": kwargs.get("kpack", 2),
+                "epilogue_subtile": kwargs.get("EPILOGUE_SUBTILE", 0),
                 **{
                     k: kwargs[k]
                     for k in AlgorithmSelectorCache.FLEX_ATTENTION_TUNABLE_KEYS
                     if k in kwargs
                 },
+                **{k: kwargs[k] for k in CONV_TUNABLE_KEYS if k in kwargs},
             },
             mutated_inputs=mutated_inputs,
             workspace_arg=workspace_arg,
@@ -2247,6 +2983,16 @@ def make_kernel_render(out_node, hint_override: Optional[int] = None):
 
 
 class ExternKernelChoice:
+    """Represents an external kernel that can participate in autotuning.
+
+    Each instance is registered as a singleton by name in ``_registry`` and
+    on the ``extern_kernels`` module so that codegen can emit a stable
+    reference.  Use ``lookup(name)`` to retrieve an existing instance
+    before creating a new one to avoid duplicate registrations.
+    """
+
+    _registry: dict[str, "ExternKernelChoice"] = {}
+
     def __init__(
         self,
         kernel,
@@ -2275,6 +3021,11 @@ def __init__(
         self.src_hash = None
         # By default GraphModule is None for extern kernels if not set
         self.gm = None
+        ExternKernelChoice._registry[name] = self
+
+    @classmethod
+    def lookup(cls, name: str) -> Optional["ExternKernelChoice"]:
+        return cls._registry.get(name)
 
     def to_callable(self):
         return getattr(extern_kernels, self.name)
@@ -2305,7 +3056,11 @@ def bind(
     ):
         self.ordered_kwargs_for_cpp_kernel = ordered_kwargs_for_cpp_kernel
         return ExternKernelCaller(
-            self, input_nodes, layout, kwargs, has_out_variant=self.has_out_variant
+            self,
+            input_nodes,
+            layout,
+            kwargs,
+            has_out_variant=self.has_out_variant,
         )
 
     @property
@@ -2313,7 +3068,7 @@ def uid(self) -> str:
         # unique by prefixing with aten
         return f"aten::{self.name}"
 
-    def choice_or_none(self, **kwargs: Any) -> Optional[ChoiceCaller]:
+    def choice_or_none(self, **kwargs: Any) -> ChoiceCaller | None:
         """
         Maybe generates a new ChoiceCaller and returns it, or None if generation fails.
 
@@ -2327,7 +3082,7 @@ def choice_or_none(self, **kwargs: Any) -> Optional[ChoiceCaller]:
 
     def maybe_append_choice(
         self, choices: list[Any], **kwargs: Any
-    ) -> Optional[NotImplementedError]:
+    ) -> NotImplementedError | None:
         # convenience function to match the Template interface, so that
         # templates and ExternKernelChoice can be treated the same when
         # generating choice callers
@@ -2340,6 +3095,10 @@ def maybe_append_choice(
 
 
 class TritonTemplateCaller(ir.TritonTemplateCallerBase):
+    """
+    Represents a ChoiceCaller for a TritonTemplate
+    """
+
     def __init__(
         self,
         name,
@@ -2348,13 +3107,11 @@ def __init__(
         make_kernel_render,
         description,
         bmreq,
-        log_info: Optional[
-            dict[str, Union[PrimitiveInfoType, list[PrimitiveInfoType]]]
-        ] = None,
+        log_info: dict[str, PrimitiveInfoType | list[PrimitiveInfoType]] | None = None,
         mutated_inputs=None,
-        workspace_arg: Optional[WorkspaceArg] = None,
-        allowed_prologue_inps: Optional[OrderedSet[str]] = None,
-        hint_override: Optional[int] = None,
+        workspace_arg: WorkspaceArg | None = None,
+        allowed_prologue_inps: OrderedSet[str] | None = None,
+        hint_override: int | None = None,
     ) -> None:
         super().__init__(name, input_nodes, layout, description)
         self.make_kernel_render = make_kernel_render
@@ -2376,17 +3133,25 @@ def __init__(
         )
         self.hint_override = hint_override
 
+        self.n_regs = None
+
     def benchmark(self, *args, out):
         assert self.bmreq is not None
-        if config.profile_bandwidth_with_do_bench_using_profiling:
+        if (
+            config.profile_bandwidth_with_do_bench_using_profiling
+            and not self._benchmark_with_cudagraphs
+        ):
             algo = self.bmreq.make_run_fn(*args, out=out)
             return do_bench_using_profiling(algo)
+        self.bmreq.benchmark_with_cudagraphs = self._benchmark_with_cudagraphs
         return self.bmreq.benchmark(*args, out=out)
 
     def precompile(self):
         assert self.bmreq is not None
         self.bmreq.precompile()
 
+        self.n_regs = self.bmreq.n_regs
+
     def __str__(self) -> str:
         return f"TritonTemplateCaller({self.bmreq.module_path}, {self.description})"
 
@@ -2414,7 +3179,7 @@ def output_node(self):
             buffer.annotations["ktc"] = self.annotations["ktc"]
         return ir.TensorBox.create(buffer)
 
-    def info_dict(self) -> dict[str, Union[PrimitiveInfoType, list[PrimitiveInfoType]]]:
+    def info_dict(self) -> dict[str, PrimitiveInfoType | list[PrimitiveInfoType]]:
         """Information returned here is logged to the autotune log file when that is enabled."""
         return self.log_info
 
@@ -2454,7 +3219,7 @@ def __init__(
         self.kwargs = kwargs or {}
         self.has_out_variant = has_out_variant
         self.gm = choice.gm
-        self.bmreq: Optional[BenchmarkRequest] = None
+        self.bmreq: BenchmarkRequest | None = None
 
         from torch._inductor.autotune_process import (
             ExternKernelBenchmarkRequest,
@@ -2476,8 +3241,8 @@ def __init__(
             if not device:
                 device = torch.device("cpu")
 
-        self.input_tensor_meta: Union[list[TensorMeta], TensorMeta]
-        self.output_tensor_meta: Union[list[TensorMeta], TensorMeta]
+        self.input_tensor_meta: list[TensorMeta] | TensorMeta
+        self.output_tensor_meta: list[TensorMeta] | TensorMeta
         self.input_tensor_meta, self.output_tensor_meta = [], []
         if device.type == "cpu":
             benchmark_cls = ExternKernelCPUBenchmarkRequest
@@ -2506,7 +3271,9 @@ def __str__(self) -> str:
         return f"ExternKernelCaller({self.choice.call_name()})"
 
     def benchmark(self, *args, out):
-        # pyrefly: ignore [missing-attribute]
+        assert self.bmreq is not None
+        # pyrefly: ignore[missing-attribute]
+        self.bmreq.benchmark_with_cudagraphs = self._benchmark_with_cudagraphs
         return self.bmreq.benchmark(*args, out=out)
 
     def benchmark_collective(self, *args, out):
@@ -2565,7 +3332,7 @@ def output_node(self):
             inner.annotations["ktc"] = self.annotations["ktc"]
         return ir.TensorBox.create(inner)
 
-    def info_dict(self) -> dict[str, Union[PrimitiveInfoType, list[PrimitiveInfoType]]]:
+    def info_dict(self) -> dict[str, PrimitiveInfoType | list[PrimitiveInfoType]]:
         """Information returned here is logged to the autotune log file when that is enabled."""
         return {
             "backend": "extern",
@@ -2577,7 +3344,7 @@ def autoheuristic_id(self):
 
 
 @functools.cache
-def get_mm_log_filename() -> Optional[str]:
+def get_mm_log_filename() -> str | None:
     mm_file_name = os.environ.get("TORCHINDUCTOR_MM_LOGGING_FILE", None)
     if not mm_file_name:
         return None
@@ -2589,7 +3356,7 @@ def get_mm_log_filename() -> Optional[str]:
 
 
 @functools.cache
-def get_flex_attention_log_filename() -> Optional[str]:
+def get_flex_attention_log_filename() -> str | None:
     flex_attention_file_name = os.environ.get(
         "TORCHINDUCTOR_FLEX_ATTENTION_LOGGING_FILE", None
     )
@@ -2599,6 +3366,15 @@ def get_flex_attention_log_filename() -> Optional[str]:
     return str(Path(flex_attention_file_name).with_suffix(".json"))
 
 
+@functools.cache
+def get_conv_log_filename() -> str | None:
+    conv_file_name = os.environ.get("TORCHINDUCTOR_CONV_LOGGING_FILE", None)
+    if not conv_file_name:
+        return None
+
+    return str(Path(conv_file_name).with_suffix(".json"))
+
+
 def append_to_log(filename, data):
     lock_file = filename.replace(".json", ".lock")
     lock = FileLock(lock_file)
@@ -2712,11 +3488,7 @@ def get_num_workers() -> int:
     if "TORCHINDUCTOR_COMPILE_THREADS" in os.environ:
         return int(os.environ["TORCHINDUCTOR_COMPILE_THREADS"])
 
-    cpu_count = (
-        len(os.sched_getaffinity(0))
-        if hasattr(os, "sched_getaffinity")
-        else os.cpu_count()
-    )
+    cpu_count = torch._utils.cpu_count()
     assert cpu_count
 
     # Divide the number of CPUs by the number of GPUs for distributed workloads
@@ -2747,6 +3519,27 @@ def create_precompile_key(
     )
 
 
+@dataclasses.dataclass(frozen=True, slots=True)
+class PrecompileFunction:
+    """Callable precompile wrapper that carries its cache key for scoped cleanup."""
+
+    fn: Callable[[], dict[ChoiceCaller, float]]
+    precompile_key: str | None = None
+
+    def __call__(self) -> dict[ChoiceCaller, float]:
+        return self.fn()
+
+
+def _benchmark_request_for_choice(choice: ChoiceCaller) -> Any | None:
+    bmreq = getattr(choice, "_bmreq", None)
+    if bmreq is not None:
+        return bmreq
+    try:
+        return vars(choice).get("bmreq")
+    except TypeError:
+        return None
+
+
 # Args to FeedbackFunctions
 # timings: mapping from choices to the benchmark time
 # name: name of the op
@@ -2800,6 +3593,102 @@ def filter_choices_by_desc_regex(choices: list[ChoiceCaller]) -> list[ChoiceCall
     return choices
 
 
+def _classify_kernel_operation(
+    name: str, choices: list[ChoiceCaller], input_nodes
+) -> str:
+    """
+    Classify the operation type for logging and filtering purposes.
+    Returns one of: "mm", "conv", "flex", or "other"
+
+    This is more robust than simple string matching as it:
+    1. Checks template types from choices
+    2. Uses input shape patterns
+    3. Falls back to exact name matching (not substring)
+    """
+    # First, try to classify from choice types
+    if choices:
+        for choice in choices:
+            if isinstance(choice, TritonTemplateCaller):
+                # Extract template name (e.g., "mm" from "mm_1", "convolution2d" from "convolution2d_3")
+                template_name = choice.name.rsplit("_", 1)[0]
+
+                # Check known template patterns
+                if template_name in (
+                    "mm",
+                    "bmm",
+                    "mm_persistent_tma",
+                    "grouped_mm",
+                    "scaled_grouped_mm",
+                    "mm_plus_mm",
+                    "blackwell_ws_persistent_device_tma",
+                    "scaled_mm_device_tma_main_loop_scaling",
+                ):
+                    return "mm"
+                elif template_name in ("convolution2d", "convolution3d"):
+                    return "conv"
+                elif template_name.startswith("flex_"):
+                    return "flex"
+
+            elif isinstance(choice, ExternKernelChoice):
+                # Check extern kernel names
+                choice_name = choice.name
+                if choice_name in (
+                    "mm",
+                    "bmm",
+                    "addmm",
+                    "baddbmm",
+                    "_int_mm",
+                    "_scaled_mm",
+                ):
+                    return "mm"
+                elif "conv" in choice_name:
+                    return "conv"
+
+    # Second, use input shape heuristics for additional validation
+    if len(input_nodes) >= 2:
+        try:
+            input_0_shape = input_nodes[0].get_size()
+            input_1_shape = input_nodes[1].get_size()
+
+            # Matrix multiplication patterns
+            if len(input_0_shape) == 2 and len(input_1_shape) == 2:
+                return "mm"
+            elif len(input_0_shape) == 3 and len(input_1_shape) == 3:
+                return "mm"  # bmm
+
+            # Convolution patterns: input NCHW/NCDHW, weight OIHW/OIDHW
+            elif len(input_0_shape) in (4, 5) and len(input_1_shape) in (4, 5):
+                # Could be conv or flex_attention, prefer template name if available
+                if len(input_0_shape) == 4 and len(input_1_shape) == 4:
+                    # Check if it looks like conv (channel dims match)
+                    # Conv: input[N,C,H,W] @ weight[O,C,kH,kW] where input[1] == weight[1]
+                    try:
+                        if input_0_shape[1] == input_1_shape[1]:
+                            return "conv"
+                    except (IndexError, TypeError):
+                        pass
+
+        except (ValueError, IndexError, AttributeError):
+            pass
+
+    # Last resort: exact name matching (not substring to avoid false positives)
+    name_lower = name.lower()
+    if name_lower in ("mm", "bmm", "addmm", "baddbmm"):
+        return "mm"
+    elif name_lower in (
+        "convolution",
+        "convolution2d",
+        "convolution3d",
+        "conv2d",
+        "conv3d",
+    ):
+        return "conv"
+    elif name_lower.startswith("flex_"):
+        return "flex"
+
+    return "other"
+
+
 class AlgorithmSelectorCache(PersistentCache):
     """
     A persistent cache for algorithm selection results used in autotuning of GEMMs
@@ -2837,7 +3726,7 @@ def __init__(self, *args, **kwargs) -> None:
         # no guarantee that the first lowering for a given key will also be the
         # first to benchmark it. share a single precompilation function for all lowerings
         # of a particular key
-        self.precompile_cache: dict[str, Callable[[], dict[ChoiceCaller, float]]] = {}
+        self.precompile_cache: dict[str, PrecompileFunction] = {}
         # cache for prescreening results to ensure deterministic candidate selection
         self.prescreening_cache: dict[str, OrderedSet[str]] = {}
         # list of callbacks that are called after benchmarking
@@ -2882,12 +3771,13 @@ def __call__(
         # corresponding ir.Buffer. if passed for a given
         # arg, the function will be called instead of
         # generating a random torch.Tensor for benchmarking.
-        input_gen_fns: Optional[dict[int, Callable[[ir.Buffer], torch.Tensor]]] = None,
+        input_gen_fns: dict[int, Callable[[ir.Buffer], torch.Tensor]] | None = None,
         precompilation_timeout_seconds: int = 60 * 60,
         return_multi_template=False,
         best_config_future=None,
-        return_choice=False,  # TODO: return_choice is temporary and will be refactored soon
         is_collective=False,
+        min_speedup_threshold: float = 1.0,  # Only pick non-fallback if faster by this ratio
+        benchmark_with_cudagraphs: bool = False,  # Use CUDA graphs for ExternKernelCaller benchmarking
     ):
         from .codegen.cutlass.kernel import CUTLASSTemplateCaller
 
@@ -2895,6 +3785,11 @@ def __call__(
         for preprocessing_fn in self.preprocessing_fns:
             choices = preprocessing_fn(choices)
 
+        # Apply benchmark_with_cudagraphs to all choices
+        if benchmark_with_cudagraphs:
+            for choice in choices:
+                choice._benchmark_with_cudagraphs = True
+
         # Templates selected with input_gen_fns require specific input data to avoid IMA
         # Passing custom input gen fns to benchmark_fusion NYI, so skip deferred template selection
         # TODO(jgong5): support multi-template on CPU C++ backend
@@ -2905,27 +3800,26 @@ def __call__(
 
         # TODO - assert that we have not mutating kernels here
 
-        if mm_file_name := get_mm_log_filename():
-            M, K = input_nodes[-2].get_size()[:2]
-            N = input_nodes[-1].get_size()[-1]
-            append_to_log(mm_file_name, {"invoke": str((M, K, N))})
-
         if len(choices) == 0:
             raise self.create_no_valid_choices(name, "No choices exist for backend.")
-        log.debug("Max autotune selects from %s choices.", str(len(choices)))
+        log.debug("Max autotune selects from %s choices.", len(choices))
 
         if len(choices) == 1:
             if not isinstance(choices[0], CUTLASSTemplateCaller):
                 # CUTLASSTemplateCaller still needs to go through the autotuning process to retrieve workspace size.
-                return choices[0].output_node()
+                node = choices[0].output_node()
+                return node, choices[0]
 
         if config.deterministic:
-            return self.pick_deterministic_choice(choices).output_node()
+            choice = self.pick_deterministic_choice(choices)
+            node = choice.output_node()
+            return node, choice
 
         inputs_key = create_inputs_key(input_nodes)
 
-        if config.autotune_in_subproc:
-            # Initialize the suprocess pool so it will warmup early.
+        has_cutlass = any(isinstance(c, CUTLASSTemplateCaller) for c in choices)
+        if config.autotune_in_subproc or has_cutlass:
+            # Warmup the subprocess pool early so it's ready for benchmarking
             torch._inductor.autotune_process.get_tuning_process_pool()
 
         precompile_fn = self.make_precompile_fn(
@@ -2940,9 +3834,6 @@ def __call__(
                 assert not config.benchmark_epilogue_fusion, (
                     "Benchmarking epilogues will cause gpu contention with pipelined autotuning"
                 )
-                assert all(not isinstance(c, SubgraphChoiceCaller) for c in choices), (
-                    "Pipelined autotuning not compatible yet with subgraph choices"
-                )
                 extern_kernels = [
                     c for c in choices if AlgorithmSelectorCache._is_extern(c)
                 ]
@@ -2968,19 +3859,26 @@ def __call__(
                 else:
                     precompile_future = None
 
-                def get_timings(hint_override: Optional[int] = None):
+                def get_timings(hint_override: int | None = None):
                     assert not hint_override, (
                         "Hint not supported with pipelined autotuning"
                     )
                     # Await precompilation future, thread pool
                     precompile_start_ts = time.time()
+                    final_choices = choices
                     if precompile_future:
-                        precompile_future.result()
+                        try:
+                            precompile_future.result()
+                        except NoValidChoicesError:
+                            log.error(
+                                "Runtime error for autotuning triton choices, defaulting to extern kernels.",
+                            )
+                            final_choices = extern_kernels
                     precompile_elapse = time.time() - precompile_start_ts
 
                     # Await autotuning in subproc pool
                     autotune_start_ts = time.time()
-                    results = AsyncAutotuner.get_results(choices, inputs_key)
+                    results = AsyncAutotuner.get_results(final_choices, inputs_key)
                     autotune_wait_ts = time.time() - autotune_start_ts
                     AlgorithmSelectorCache.log_results(
                         name,
@@ -2993,7 +3891,7 @@ def get_timings(hint_override: Optional[int] = None):
                     return results
             else:
 
-                def get_timings(hint_override: Optional[int] = None):
+                def get_timings(hint_override: int | None = None):
                     filtered_choices = [
                         c
                         for c in choices
@@ -3035,14 +3933,18 @@ def get_timings(hint_override: Optional[int] = None):
                 if isinstance(c, TritonTemplateCaller):
                     allowed_prologue_inps |= c.allowed_prologue_inps
 
-            return torch._inductor.ir.TensorBox.create(
-                torch._inductor.ir.MultiTemplateBuffer(
-                    layout,
-                    input_nodes,
-                    get_timings,
-                    choices,
-                    allowed_prologue_inps,
-                )
+            # No single winning choice yet; selection is deferred to benchmark fusion
+            return (
+                torch._inductor.ir.TensorBox.create(
+                    torch._inductor.ir.MultiTemplateBuffer(
+                        layout,
+                        input_nodes,
+                        get_timings,
+                        choices,
+                        allowed_prologue_inps,
+                    )
+                ),
+                None,
             )
 
         timings = self.do_autotuning(
@@ -3070,27 +3972,77 @@ def get_timings(hint_override: Optional[int] = None):
                         "Autotuning returned empty timings, falling back to first `ExternKernelCaller`: %s",
                         node,
                     )
-                    if return_choice:
-                        return node, choice
-                    return node
+                    return node, choice
             node = choices[0].output_node()
             choice = choices[0]
             log.debug(
                 "Autotuning returned empty timings, falling back to first choice: %s",
                 node,
             )
-            if return_choice:
-                return node, choice
-            return node
+            return node, choice
 
         # if we got any timings at all, pick the best of those
-        choice = min(timings, key=timings.__getitem__)
+        best_choice = min(timings, key=timings.__getitem__)
+        best_time = timings[best_choice]
+
+        # All benchmarks failed; fall back to ATen if available (#171094)
+        if math.isinf(best_time):
+            extern_choices = [c for c in choices if isinstance(c, ExternKernelCaller)]
+            if extern_choices:
+                best_choice = extern_choices[0]
+                log.warning(
+                    "All autotuning benchmarks failed (timing=inf). Falling back to ExternKernelCaller: %s",
+                    getattr(best_choice, "name", "<unknown>"),
+                )
+            else:
+                log.warning(
+                    "All autotuning benchmarks failed (timing=inf) and no ExternKernelCaller fallback available. "
+                    "Selected kernel %s may cause runtime errors.",
+                    getattr(best_choice, "name", "<unknown>"),
+                )
+
+        # Apply min_speedup_threshold: only pick non-fallback if it beats fallback by threshold
+        elif min_speedup_threshold > 1.0:
+
+            def is_fallback(c: ChoiceCaller) -> bool:
+                return isinstance(c, ExternKernelCaller) and getattr(
+                    c.choice, "use_fallback_kernel", False
+                )
+
+            fallback_choices = [c for c in timings if is_fallback(c)]
+            if fallback_choices and not is_fallback(best_choice):
+                fallback_time = min(timings[c] for c in fallback_choices)
+                speedup = fallback_time / best_time if best_time > 0 else 0
+
+                if speedup < min_speedup_threshold:
+                    # Best choice doesn't beat fallback by enough, use fallback instead
+                    log.debug(
+                        "Best choice %s speedup %.2fx < threshold %.2fx, using fallback",
+                        best_choice.name,
+                        speedup,
+                        min_speedup_threshold,
+                    )
+                    best_choice = min(fallback_choices, key=lambda c: timings[c])
+
+        best_choice = V.choices.override_best_choice(best_choice, timings)
+
+        # Test-only: force choosing decomposition (non-fallback) if available
+        if config.test_configs.force_custom_op_decomposition:
+
+            def is_fallback(c: ChoiceCaller) -> bool:
+                return isinstance(c, ExternKernelCaller) and getattr(
+                    c.choice, "use_fallback_kernel", False
+                )
+
+            non_fallback_choices = [c for c in timings if not is_fallback(c)]
+            if non_fallback_choices:
+                best_choice = min(non_fallback_choices, key=lambda c: timings[c])
+
+        choice = best_choice
         node = choice.output_node()
 
         log.debug("Autotuning selected choice: %s", node)
-        if return_choice:
-            return node, choice
-        return node
+        return node, choice
 
     def benchmark(
         self,
@@ -3098,8 +4050,9 @@ def benchmark(
         input_nodes,
         layout,
         input_gen_fns,
-        hint_override: Optional[int] = None,
+        hint_override: int | None = None,
         is_collective=False,
+        precompile_key: str | None = None,
     ):
         counters["inductor"]["select_algorithm_autotune"] += 1
         # TODO(nmacchioni): remove this layer of abstraction
@@ -3114,7 +4067,32 @@ def benchmark(
         )
         # `benchmark_fn(choices)` will execute each choice, and return a dict[choice, timing] which
         # maps each choice to its runtime, calculated by the specified benchmarker, in milliseconds
-        return benchmark_fn(choices)
+        try:
+            return benchmark_fn(choices)
+        finally:
+            # Safety net for failures before individual BenchmarkRequest cleanup
+            # can run; request-level cleanup remains the authoritative owner cleanup.
+            evict_paths = OrderedSet(
+                [
+                    path
+                    for choice in choices
+                    for bmreq in (_benchmark_request_for_choice(choice),)
+                    for path in (getattr(bmreq, "module_path", None),)
+                    if path is not None
+                ]
+            )
+            if evict_paths:
+                for path in evict_paths:
+                    PyCodeCache.modules_no_attr.pop(path, None)
+                    PyCodeCache.linemaps.pop(path, None)
+                PyCodeCache.modules[:] = [
+                    module
+                    for module in PyCodeCache.modules
+                    if getattr(module, "__file__", None) not in evict_paths
+                ]
+
+            if precompile_key is not None:
+                self.precompile_cache.pop(precompile_key, None)
 
     def autotune(
         self,
@@ -3123,8 +4101,9 @@ def autotune(
         layout,
         input_gen_fns,
         choices,
-        hint_override: Optional[int] = None,
+        hint_override: int | None = None,
         is_collective=False,
+        precompile_key: str | None = None,
     ):
         log.debug("Starting autotuning")
 
@@ -3141,6 +4120,7 @@ def autotune(
                 input_gen_fns,
                 hint_override=hint_override,
                 is_collective=is_collective,
+                precompile_key=precompile_key,
             )
             if config.max_autotune_report_choices_stats:
                 _log_autotune_choices_stats(
@@ -3157,7 +4137,7 @@ def do_autotuning(
         inputs_key,
         choices,
         precompile_fn,
-        hint_override: Optional[int] = None,
+        hint_override: int | None = None,
         best_config_future=None,
         is_collective=False,
     ):
@@ -3190,6 +4170,8 @@ def do_autotuning(
             NoValidChoicesError: When all choices fail to compile or benchmark, or when all
                 timing results are non-finite.
         """
+        precompile_key = getattr(precompile_fn, "precompile_key", None)
+
         if log.isEnabledFor(logging.DEBUG) and not use_pipelined_autotuning():
             # Log shape information for debugging timeout issues
             sizevars = V.graph.sizevars
@@ -3237,7 +4219,7 @@ def do_autotuning(
         candidates = self.prescreen_choices(
             choices, name, inputs_key, self.prescreening_cache
         )
-        prescreening_elapse: Optional[float] = None
+        prescreening_elapse: float | None = None
         if candidates:
             prescreening_start_ts = time.time()
             timings = self.lookup(
@@ -3307,6 +4289,7 @@ def track_has_autotuned(choices):
                 choices,
                 hint_override=hint_override,
                 is_collective=is_collective,
+                precompile_key=precompile_key,
             )
 
         timings = self.lookup(
@@ -3397,7 +4380,7 @@ def make_precompile_fn(
         choices,
         name: str,
         inputs_key: str,
-        precompilation_timeout_seconds: Optional[int] = 60 * 60,
+        precompilation_timeout_seconds: int | None = 60 * 60,
     ) -> Callable[[], dict[ChoiceCaller, float]]:
         """
         Returns a function that precompiles the given choices.
@@ -3592,9 +4575,10 @@ def wait_on_futures() -> dict[ChoiceCaller, float]:
                     precompile_times[choice] = elapsed_times[future]
             return precompile_times
 
-        self.precompile_cache[precompile_key] = wait_on_futures
+        precompile_fn = PrecompileFunction(wait_on_futures, precompile_key)
+        self.precompile_cache[precompile_key] = precompile_fn
 
-        return wait_on_futures
+        return precompile_fn
 
     @classmethod
     def get_inputs(
@@ -3602,8 +4586,8 @@ def get_inputs(
         choices: Sequence[ChoiceCaller],
         input_nodes: list[ir.IRNode],
         layout: ir.Layout,
-        input_gen_fns: Optional[dict[int, Callable[[ir.Buffer], torch.Tensor]]],
-        hint_override: Optional[int] = None,
+        input_gen_fns: dict[int, Callable[[ir.Buffer], torch.Tensor]] | None,
+        hint_override: int | None = None,
     ) -> AutotuneArgs:
         """
         Factory method to create AutotuneArgs from a list of ChoiceCallers.
@@ -3645,7 +4629,7 @@ def get_inputs(
                         hint_override=hint_override,
                     )
                     strides = V.graph.sizevars.optimization_hints_with_override(
-                        input_node.get_stride(),
+                        get_strides_with_layout_constraints(input_node),
                         hint_override=hint_override,
                     )
                     storage_offset = V.graph.sizevars.optimization_hint_with_override(
@@ -3662,12 +4646,21 @@ def get_inputs(
 
                 if needed_size > current_size:
                     # Create a new base tensor with sufficient storage
-                    new_base = torch.randn(
-                        needed_size,
-                        dtype=base.dtype,
-                        device=base.device,
-                        requires_grad=base.requires_grad,
-                    )
+                    if base.dtype == torch.float4_e2m1fn_x2:
+                        new_base = torch.randint(
+                            0,
+                            256,
+                            (needed_size,),
+                            dtype=torch.uint8,
+                            device=base.device,
+                        ).view(torch.float4_e2m1fn_x2)
+                    else:
+                        new_base = torch.randn(
+                            needed_size,
+                            dtype=base.dtype,
+                            device=base.device,
+                            requires_grad=base.requires_grad,
+                        )
                     base = new_base.as_strided(
                         base.size(), base.stride(), base.storage_offset()
                     )
@@ -3679,7 +4672,9 @@ def get_inputs(
 
         # Also check the output tensor for storage size
         out_base = out if out._base is None else out._base
-        out_offset = V.graph.sizevars.size_hint(layout.offset)
+        # Only used for benchmarking tensor setup, not correctness.
+        # Offset is almost always 0; use that as fallback.
+        out_offset = V.graph.sizevars.optimization_hint(layout.offset, fallback=0)
         needed_out_size = torch._prims_common.compute_required_storage_length(
             out.size(), out.stride(), out_offset
         )
@@ -3687,12 +4682,21 @@ def get_inputs(
 
         if needed_out_size > current_out_size:
             # Create a new base tensor with sufficient storage
-            new_out_base = torch.randn(
-                needed_out_size,
-                dtype=out_base.dtype,
-                device=out_base.device,
-                requires_grad=out_base.requires_grad,
-            )
+            if out_base.dtype == torch.float4_e2m1fn_x2:
+                new_out_base = torch.randint(
+                    0,
+                    256,
+                    (needed_out_size,),
+                    dtype=torch.uint8,
+                    device=out_base.device,
+                ).view(torch.float4_e2m1fn_x2)
+            else:
+                new_out_base = torch.randn(
+                    needed_out_size,
+                    dtype=out_base.dtype,
+                    device=out_base.device,
+                    requires_grad=out_base.requires_grad,
+                )
             out_base = new_out_base.as_strided(
                 out_base.size(), out_base.stride(), out_base.storage_offset()
             )
@@ -3722,18 +4726,26 @@ def benchmark_choice(
         benchmark_tensors = autotune_args.get_benchmark_tensors(cls._is_extern(choice))
         inputs, output = benchmark_tensors.unpack()
         output.zero_()
-        result = choice.benchmark(*inputs, out=output)
-        device_type = next(
-            (tensor.device.type for tensor in inputs if is_gpu(tensor.device.type)),
-            "cuda",
-        )
-        device_interface = get_interface_for_device(device_type)
-        if device_interface.is_available():
-            device_interface.synchronize()  # shake out any CUDA errors
-
-        if VERIFY and autotune_args.expected is not None:
-            autotune_args.verify(**VERIFY)
-        return result
+        try:
+            result = choice.benchmark(*inputs, out=output)
+            device_type = next(
+                (tensor.device.type for tensor in inputs if is_gpu(tensor.device.type)),
+                "cuda",
+            )
+            device_interface = get_interface_for_device(device_type)
+            if device_interface.is_available():
+                device_interface.synchronize()  # shake out any CUDA errors
+
+            if VERIFY and autotune_args.expected is not None:
+                autotune_args.verify(**VERIFY)
+            return result
+        finally:
+            bmreq = _benchmark_request_for_choice(choice)
+            cleanup_run_fn = getattr(bmreq, "cleanup_run_fn", None)
+            if cleanup_run_fn is not None:
+                # In-process autotuning owns the loaded benchmark module, so clean it
+                # immediately after each choice. The outer benchmark cleanup is a fallback.
+                cleanup_run_fn()
 
     @classmethod
     def _run_collective_benchmark(
@@ -3868,6 +4880,23 @@ def benchmark_choices(
         """
         Benchmark a list of choices and return timing dict.
         """
+        from torch._inductor.codegen.cutlass.kernel import CUTLASSTemplateCaller
+
+        # Reorder only when CUTLASS is present: benchmark ATen/cuBLAS before
+        # CUTLASS to collect valid fallback timings before any CUTLASS kernel
+        # can corrupt the CUDA context (#171094)
+        if any(isinstance(c, CUTLASSTemplateCaller) for c in choices):
+
+            def choice_priority(c: ChoiceCaller) -> int:
+                if isinstance(c, ExternKernelCaller):
+                    return 0  # ATen/cuBLAS first
+                elif isinstance(c, CUTLASSTemplateCaller):
+                    return 2  # CUTLASS last
+                else:
+                    return 1  # Triton and others in the middle
+
+            choices = sorted(choices, key=choice_priority)
+
         if is_collective:
             import torch.distributed as dist
 
@@ -3892,8 +4921,6 @@ def benchmark_choices(
                 else:
                     timing = cls.benchmark_choice(choice, autotune_args)
             except CUDACompileError:
-                from torch._inductor.codegen.cutlass.kernel import CUTLASSTemplateCaller
-
                 if not isinstance(choice, CUTLASSTemplateCaller):
                     log.exception(
                         "CUDA compilation error during autotuning: \n%s. \nIgnoring this choice."
@@ -3903,8 +4930,6 @@ def benchmark_choices(
                 log.warning("Not yet implemented", exc_info=True)
                 timing = float("inf")
             except RuntimeError as e:
-                from torch._inductor.codegen.cutlass.kernel import CUTLASSTemplateCaller
-
                 msg = str(e)
                 if "invalid argument" in msg:
                     msg += "\n\nThis may mean this GPU is too small for max_autotune mode.\n\n"
@@ -3954,6 +4979,23 @@ def benchmark_choices(
                 timings.update({c: float("inf") for c in choices if c not in timings})
                 break
 
+            # Skip remaining CUTLASS choices after first failure (#171094)
+            if not math.isfinite(timing) and isinstance(choice, CUTLASSTemplateCaller):
+                has_valid_fallback = any(
+                    math.isfinite(t) and not isinstance(c, CUTLASSTemplateCaller)
+                    for c, t in timings.items()
+                )
+                if has_valid_fallback:
+                    log.warning(
+                        "CUTLASS choice %s failed during benchmarking. "
+                        "Skipping remaining CUTLASS choices to avoid CUDA context corruption.",
+                        getattr(choice, "name", "<unknown>"),
+                    )
+                    for c in choices:
+                        if c not in timings and isinstance(c, CUTLASSTemplateCaller):
+                            timings[c] = float("inf")
+                    break
+
         return timings
 
     @classmethod
@@ -3962,8 +5004,8 @@ def benchmark_in_current_process(
         choices: Sequence[ChoiceCaller],
         input_nodes: list[ir.IRNode],
         layout: ir.Layout,
-        input_gen_fns: Optional[dict[int, Callable[[ir.Buffer], torch.Tensor]]],
-        hint_override: Optional[int] = None,
+        input_gen_fns: dict[int, Callable[[ir.Buffer], torch.Tensor]] | None,
+        hint_override: int | None = None,
         is_collective=False,
     ) -> dict[ChoiceCaller, float]:
         inputs = cls.get_inputs(
@@ -3981,20 +5023,27 @@ def benchmark_in_sub_process(
         choices: Sequence[ChoiceCaller],
         input_nodes: list[ir.IRNode],
         layout: ir.Layout,
-        input_gen_fns: Optional[dict[int, Callable[[ir.Buffer], torch.Tensor]]],
-        hint_override: Optional[int] = None,
+        input_gen_fns: dict[int, Callable[[ir.Buffer], torch.Tensor]] | None,
+        hint_override: int | None = None,
     ):
         from . import autotune_process
+        from .codegen.cutlass.kernel import CUTLASSTemplateCaller
 
-        # only benchmark triton kernel in sub process for now.
-        # ATen/Extern kernel are still benchmarked in the current process.
+        # ATen/Extern kernels are safe to benchmark in the current process.
         extern = [c for c in choices if cls._is_extern(c)]
-        triton = [c for c in choices if not cls._is_extern(c)]
+        non_cutlass = [
+            c
+            for c in choices
+            if not cls._is_extern(c) and not isinstance(c, CUTLASSTemplateCaller)
+        ]
+        cutlass = [c for c in choices if isinstance(c, CUTLASSTemplateCaller)]
 
         timings = cls.benchmark_in_current_process(
             extern, input_nodes, layout, input_gen_fns, hint_override=hint_override
         )
-        timings.update(autotune_process.benchmark_in_sub_process(triton))  # type: ignore[arg-type]
+        # Order Triton before CUTLASS so valid Triton timings are collected
+        # before any CUTLASS kernel can crash the subprocess (#171094)
+        timings.update(autotune_process.benchmark_in_sub_process(non_cutlass + cutlass))  # type: ignore[arg-type]
         return timings
 
     @classmethod
@@ -4003,15 +5052,19 @@ def make_benchmark_fn(
         choices: Sequence[ChoiceCaller],
         input_nodes: list[ir.IRNode],
         layout: ir.Layout,
-        input_gen_fns: Optional[dict[int, Callable[[ir.Buffer], torch.Tensor]]],
-        hint_override: Optional[int] = None,
+        input_gen_fns: dict[int, Callable[[ir.Buffer], torch.Tensor]] | None,
+        hint_override: int | None = None,
         is_collective=False,
     ):
+        from .codegen.cutlass.kernel import CUTLASSTemplateCaller
+
         if DEBUG:
             print(f"{len(choices)} tuning requests:")
 
+        has_cutlass = any(isinstance(c, CUTLASSTemplateCaller) for c in choices)
+
         # Collective ops must use current process
-        if is_collective or not config.autotune_in_subproc:
+        if is_collective:
             return functools.partial(
                 cls.benchmark_in_current_process,
                 input_nodes=input_nodes,
@@ -4020,7 +5073,9 @@ def make_benchmark_fn(
                 hint_override=hint_override,
                 is_collective=is_collective,
             )
-        else:
+        # CUTLASS kernels can cause sticky CUDA errors that permanently corrupt
+        # the CUDA context, so always benchmark them in a subprocess (#171094)
+        elif config.autotune_in_subproc or has_cutlass:
             return functools.partial(
                 cls.benchmark_in_sub_process,
                 input_nodes=input_nodes,
@@ -4028,6 +5083,14 @@ def make_benchmark_fn(
                 input_gen_fns=input_gen_fns,
                 hint_override=hint_override,
             )
+        else:
+            return functools.partial(
+                cls.benchmark_in_current_process,
+                input_nodes=input_nodes,
+                layout=layout,
+                input_gen_fns=input_gen_fns,
+                hint_override=hint_override,
+            )
 
     @staticmethod
     def prescreen_choices(
@@ -4087,6 +5150,10 @@ def prescreen_choices(
         if len(candidates) < 10:
             return []
 
+        # Include ATen in prescreening as fallback (#171094)
+        extern_choices = [c for c in choices if isinstance(c, ExternKernelCaller)]
+        candidates = extern_choices + candidates
+
         return candidates  # type: ignore[return-value]
 
     @staticmethod
@@ -4177,6 +5244,120 @@ def prune_choices_postscreen(
         )
         return pruned_choices
 
+    @staticmethod
+    def maybe_log_mm_results(
+        name: str, input_nodes: list[ir.IRNode], timings: dict[ChoiceCaller, float]
+    ) -> None:
+        """Log matrix multiplication autotuning results."""
+        mm_filename = get_mm_log_filename()
+        if not mm_filename:
+            return
+
+        # Classify operation to ensure it's actually an MM operation
+        choices_list = list(timings.keys())
+        operation_type = _classify_kernel_operation(name, choices_list, input_nodes)
+        if operation_type != "mm":
+            return
+
+        if len(input_nodes) < 2:
+            return
+
+        M, K = input_nodes[-2].get_size()[:2]
+        N = input_nodes[-1].get_size()[-1]
+
+        def get_choice_info(choice):
+            if isinstance(choice, ExternKernelCaller):
+                return {"type": "cublas", "time": timings[choice]}
+
+            if isinstance(choice, TritonTemplateCaller):
+                info = choice.info_dict()
+                tile = info["tile_shape"]
+
+                tile_vals = eval(tile)  # type: ignore[arg-type]
+                BLOCK_M = tile_vals[0]
+                BLOCK_K = tile_vals[1]
+                BLOCK_N = tile_vals[2]
+
+                return {
+                    "type": "triton",
+                    "time": timings[choice],
+                    "BLOCK_M": BLOCK_M,
+                    "BLOCK_K": BLOCK_K,
+                    "BLOCK_N": BLOCK_N,
+                    "num_stages": info["num_stages"],
+                    "num_warps": info["num_warps"],
+                    "waves_per_eu": info.get("waves_per_eu", 0),
+                    "matrix_instr_nonkdim": info.get("matrix_instr_nonkdim", 0),
+                    "kpack": info.get("kpack", 2),
+                }
+            return None
+
+        out_dict = {
+            str((M, K, N)): [get_choice_info(choice) for choice in timings],
+            "kernel_type": name,
+        }
+
+        append_to_log(mm_filename, out_dict)
+
+    @staticmethod
+    def maybe_log_conv_results(
+        name: str, input_nodes: list[ir.IRNode], timings: dict[ChoiceCaller, float]
+    ) -> None:
+        """Log convolution autotuning results."""
+        conv_filename = get_conv_log_filename()
+        if not conv_filename:
+            return
+
+        # Classify operation to ensure it's actually a conv operation
+        choices_list = list(timings.keys())
+        operation_type = _classify_kernel_operation(name, choices_list, input_nodes)
+        if operation_type != "conv":
+            return
+
+        if len(input_nodes) < 2:
+            return
+
+        x_size = input_nodes[0].get_size()
+        w_size = input_nodes[1].get_size()
+
+        def get_conv_choice_info(choice):
+            if choice not in timings:
+                return None
+            info = choice.info_dict()
+
+            # Start with timing and backend type
+            result = {
+                "time": timings[choice],
+                "backend": info.get("backend", "unknown"),
+            }
+
+            # Add all parameters from info_dict
+            for key, value in info.items():
+                if key != "backend":  # Already added
+                    # Convert non-serializable types to strings
+                    try:
+                        import json
+
+                        json.dumps(value)  # Test if serializable
+                        result[key] = value
+                    except (TypeError, ValueError):
+                        result[key] = str(value)
+
+            return result
+
+        out_dict = {
+            "input_shape": str(x_size),
+            "weight_shape": str(w_size),
+            "choices": [
+                get_conv_choice_info(choice)
+                for choice in timings
+                if get_conv_choice_info(choice) is not None
+            ],
+            "kernel_type": name,
+        }
+
+        append_to_log(conv_filename, out_dict)
+
     @staticmethod
     def get_flex_attention_choice_info(
         choice: ChoiceCaller, timings: dict[ChoiceCaller, float]
@@ -4204,7 +5385,10 @@ def maybe_log_flex_attention_results(
         name: str, input_nodes: list[ir.IRNode], timings: dict[ChoiceCaller, float]
     ) -> None:
         flex_attention_filename = get_flex_attention_log_filename()
-        if not flex_attention_filename or "flex_attention" not in name:
+        # Support both flex_attention and flex_decoding
+        if not flex_attention_filename or (
+            "flex_attention" not in name and "flex_decoding" not in name
+        ):
             return
 
         if len(input_nodes) < 3:
@@ -4214,34 +5398,65 @@ def maybe_log_flex_attention_results(
         key_size = input_nodes[1].get_size()
         value_size = input_nodes[2].get_size()
 
-        B = query_size[0]
-        Hq = query_size[1]
-        seq_len_q = query_size[2]
-        qk_head_dim = query_size[3]
-        Hkv = key_size[1]
-        seq_len_kv = key_size[2]
-        v_head_dim = value_size[3]
-
-        kernel_type = "backward" if "backward" in name else "forward"
-        dims_key = str(
-            (
-                kernel_type,
-                B,
-                Hq,
-                Hkv,
-                seq_len_q,
-                seq_len_kv,
-                qk_head_dim,
-                v_head_dim,
-            )
+        # Handle both 4D (forward/backward) and 5D (decode) tensor formats
+        # 4D: [B, H, seq_len, head_dim]
+        # 5D: [B, H, 1, 1, head_dim] (decode mode has extra dimension)
+        if len(query_size) == 5:
+            # Decode mode with 5D tensors
+            B = query_size[0]
+            Hq = query_size[1]
+            # query_size[2] and query_size[3] are both 1 for decode
+            seq_len_q = query_size[2]  # This will be 1
+            qk_head_dim = query_size[4]  # Head dim is at index 4 for 5D
+            Hkv = key_size[1]
+            seq_len_kv = key_size[2]
+            v_head_dim = value_size[4] if len(value_size) == 5 else value_size[3]
+        else:
+            # Forward/backward mode with 4D tensors
+            B = query_size[0]
+            Hq = query_size[1]
+            seq_len_q = query_size[2]
+            qk_head_dim = query_size[3]
+            Hkv = key_size[1]
+            seq_len_kv = key_size[2]
+            v_head_dim = value_size[3]
+
+        kernel_type = (
+            "backward"
+            if "backward" in name
+            else ("decode" if "decoding" in name else "forward")
         )
 
+        # Create shape info dictionary
+        shape_info = {
+            "kernel_type": kernel_type,
+            "B": int(B),
+            "Hq": int(Hq),
+            "Hkv": int(Hkv),
+            "seq_len_q": int(seq_len_q),
+            "seq_len_kv": int(seq_len_kv),
+            "qk_head_dim": int(qk_head_dim),
+            "v_head_dim": int(v_head_dim),
+        }
+
         sorted_choices = sorted(timings, key=timings.__getitem__)
+
+        # Include shape info in each choice
+        choices_with_shapes = []
+        for choice in sorted_choices:
+            choice_info = AlgorithmSelectorCache.get_flex_attention_choice_info(
+                choice, timings
+            )
+            # Merge shape info with choice info
+            choice_info.update(shape_info)
+            choices_with_shapes.append(choice_info)
+
         out_dict = {
-            dims_key: [
-                AlgorithmSelectorCache.get_flex_attention_choice_info(choice, timings)
-                for choice in sorted_choices
-            ]
+            "query_shape": str(query_size),
+            "key_shape": str(key_size),
+            "value_shape": str(value_size),
+            "kernel_type": kernel_type,
+            "choices": choices_with_shapes,
         }
         append_to_log(flex_attention_filename, out_dict)
 
@@ -4252,8 +5467,8 @@ def log_results(
         timings: dict[ChoiceCaller, float],
         elapse: float,
         precompile_elapse: float,
-        prescreening_elapse: Optional[float] = None,
-        hint_override: Optional[int] = None,
+        prescreening_elapse: float | None = None,
+        hint_override: int | None = None,
         is_collective: bool = False,
     ):
         """Log the autotuning results, currently only handles mm and flex. Log Collective op autotuning result"""
@@ -4295,7 +5510,9 @@ def log_results(
             ]
         )
 
-        strides = ", ".join([str(n.get_stride()) for n in input_nodes])
+        strides = ", ".join(
+            [str(get_strides_with_layout_constraints(n)) for n in input_nodes]
+        )
         dtypes = ", ".join([str(n.get_dtype()) for n in input_nodes])
         if config.autotune_num_choices_displayed == 0:
             return
@@ -4306,41 +5523,9 @@ def log_results(
 
         best = top_k[0]
 
-        def get_choice_info(choice):
-            if isinstance(choice, torch._inductor.select_algorithm.ExternKernelCaller):
-                return {"type": "cublas", "time": timings[choice]}
-
-            assert isinstance(
-                choice, torch._inductor.select_algorithm.TritonTemplateCaller
-            )
-
-            info = choice.info_dict()
-            tile = info["tile_shape"]
-
-            tile_vals = eval(tile)  # type: ignore[arg-type]
-            BLOCK_M = tile_vals[0]
-            BLOCK_K = tile_vals[1]
-            BLOCK_N = tile_vals[2]
-
-            return {
-                "type": "triton",
-                "time": timings[choice],
-                "BLOCK_M": BLOCK_M,
-                "BLOCK_K": BLOCK_K,
-                "BLOCK_N": BLOCK_N,
-                "num_stages": info["num_stages"],
-                "num_warps": info["num_warps"],
-            }
-
-        mm_filename = get_mm_log_filename()
-        if mm_filename and "mm" in name:
-            M, K = input_nodes[-2].get_size()[:2]
-            N = input_nodes[-1].get_size()[-1]
-
-            out_dict = {str((M, K, N)): [get_choice_info(choice) for choice in timings]}
-
-            append_to_log(mm_filename, out_dict)
-
+        # Log autotuning results for each operation type
+        AlgorithmSelectorCache.maybe_log_mm_results(name, input_nodes, timings)
+        AlgorithmSelectorCache.maybe_log_conv_results(name, input_nodes, timings)
         AlgorithmSelectorCache.maybe_log_flex_attention_results(
             name, input_nodes, timings
         )
@@ -4378,7 +5563,7 @@ def get_choice_info(choice):
         )
 
     @staticmethod
-    def benchmark_example_value(node, hint_override: Optional[int] = None):
+    def benchmark_example_value(node, hint_override: int | None = None):
         """
         Convert an ir.Buffer into a concrete torch.Tensor we can use for
         benchmarking.
@@ -4386,6 +5571,9 @@ def benchmark_example_value(node, hint_override: Optional[int] = None):
         if isinstance(node, ir.Layout):
             node = ir.Buffer(name="fake", layout=node)
         # triton templates want the base tensor.
+        # Preserve the original dtype before unwrapping, since dtype views
+        # (e.g. uint8 -> float4_e2m1fn_x2) would be lost by unwrap_view.
+        original_dtype = node.get_dtype()
         if isinstance(node, ir.BaseView):
             node = node.unwrap_view()
 
@@ -4393,13 +5581,13 @@ def benchmark_example_value(node, hint_override: Optional[int] = None):
         # stride is large enough. The V.graph.get_allocation_size takes this into account.
         # So we need call as_strided in the end to 'view' the tensor with the correct
         # sizes/strides
-        return AlgorithmSelectorCache.generate_example_value(
+        result = AlgorithmSelectorCache.generate_example_value(
             V.graph.sizevars.optimization_hints_with_override(
                 node.get_size(),
                 hint_override=hint_override,
             ),
             V.graph.sizevars.optimization_hints_with_override(
-                node.get_stride(),
+                get_strides_with_layout_constraints(node),
                 hint_override=hint_override,
             ),
             node.get_device(),
@@ -4413,6 +5601,10 @@ def benchmark_example_value(node, hint_override: Optional[int] = None):
                 hint_override=hint_override,
             ),
         )
+        # Restore dtype if it was lost by unwrapping a dtype view
+        if result.dtype != original_dtype:
+            result = result.view(original_dtype)
+        return result
 
     @staticmethod
     def generate_example_value(
@@ -4449,7 +5641,9 @@ def key_of(node):
             node.get_device().type,
             str(node.get_dtype()),
             *sizevars.optimization_hints(node.get_size()),
-            *V.graph.sizevars.optimization_hints(node.get_stride()),
+            *V.graph.sizevars.optimization_hints(
+                get_strides_with_layout_constraints(node)
+            ),
             sizevars.optimization_hint(node.get_layout().offset),
         )
 
@@ -4474,7 +5668,7 @@ def clear_preprocessing_fns(self, clear_defaults: bool = False):
             self._register_default_preprocessing_fns()
 
 
-_ALGORITHM_SELECTOR_CACHE: Optional[AlgorithmSelectorCache] = None
+_ALGORITHM_SELECTOR_CACHE: AlgorithmSelectorCache | None = None
 
 
 def get_algorithm_selector_cache() -> AlgorithmSelectorCache:
@@ -4491,7 +5685,7 @@ def autotune_select_algorithm(*args, **kwargs):
     if "return_multi_template" not in kwargs:
         kwargs["return_multi_template"] = (
             torch._inductor.config.benchmark_epilogue_fusion
-            or use_pipelined_autotuning()
+            or torch._inductor.config.pipeline_max_autotune_gemm
         )
 
     if "precompilation_timeout_seconds" not in kwargs:
@@ -4554,6 +5748,15 @@ def realize_inputs(*args):
     return [realize_inputs(x) for x in args]
 
 
+def get_strides_with_layout_constraints(node):
+    if (
+        not isinstance(node, ir.ReinterpretView)
+        and node.get_name() in V.graph.buffer_layout_constraints
+    ):
+        return V.graph.buffer_layout_constraints[node.get_name()].stride
+    return node.get_stride()
+
+
 class SymbolicGridFn:
     """
     Wrapper around a grid function that allows either int or sympy inputs.
@@ -4622,7 +5825,7 @@ def _log_autotune_choices_stats(
     if not timings:
         return None
 
-    metadata: dict[str, Union[int, float, str]] = {
+    metadata: dict[str, int | float | str] = {
         "num_choices": len(timings),
         "num_triton_choices": len(
             [c for c in timings if isinstance(c, TritonTemplateCaller)]
diff --git a/torch/_inductor/shape_propagation.py b/torch/_inductor/shape_propagation.py
index e11e802fbe81f..2c5ac290cb31c 100644
--- a/torch/_inductor/shape_propagation.py
+++ b/torch/_inductor/shape_propagation.py
@@ -1,6 +1,6 @@
 import functools
 from collections.abc import Callable, Sequence
-from typing import Optional, Protocol, Union
+from typing import Protocol
 
 import sympy
 
@@ -9,7 +9,7 @@
 from .virtualized import OpsValue, V
 
 
-BlockShapeType = Optional[Sequence[Union[int, str]]]
+BlockShapeType = Sequence[int | str] | None
 
 
 class ShapeVar(Protocol):
@@ -17,7 +17,7 @@ class ShapeVar(Protocol):
     def shape(self) -> BlockShapeType: ...
 
 
-ShapeArg = Union[ShapeVar, torch.types.Number, str, OpsValue, torch.dtype]
+ShapeArg = ShapeVar | torch.types.Number | str | OpsValue | torch.dtype
 
 # Inputs need to be cacheable (e.g., not a CSEVar) in order for the cache to be effective
 # So first decompose CSEVars -> tuple before calling this
@@ -34,9 +34,7 @@ def get_broadcasted_shape(a: BlockShapeType, b: BlockShapeType) -> BlockShapeTyp
         return get_broadcasted_shape(a, (*[1] * (len(a) - len(b)), *b))
     else:
 
-        def _get_broadcasted_dim(
-            d1: Union[int, str], d2: Union[int, str]
-        ) -> Union[int, str]:
+        def _get_broadcasted_dim(d1: int | str, d2: int | str) -> int | str:
             if str(d1) == "1":
                 return d2
             elif str(d2) == "1":
@@ -100,21 +98,19 @@ def reduction(
         dtype: torch.dtype,
         src_dtype: torch.dtype,
         reduction_type: str,
-        value: Union[ShapeArg, tuple[ShapeArg, ...]],
-    ) -> Union[BlockShapeType, tuple[BlockShapeType, ...]]:
+        value: ShapeArg | tuple[ShapeArg, ...],
+    ) -> BlockShapeType | tuple[BlockShapeType, ...]:
         raise NotImplementedError
 
     @staticmethod
-    def store(
-        name: str, index: int, value: ShapeArg, mode: Optional[str] = None
-    ) -> None:
+    def store(name: str, index: int, value: ShapeArg, mode: str | None = None) -> None:
         return None
 
     @staticmethod
     def to_dtype(
         value: ShapeVar,
         dtype: torch.dtype,
-        src_dtype: Optional[torch.dtype] = None,
+        src_dtype: torch.dtype | None = None,
         use_compute_types: bool = True,
     ) -> BlockShapeType:
         return value.shape
@@ -138,7 +134,7 @@ def load_seed(name: str, offset: int) -> BlockShapeType:
     @staticmethod
     def indirect_indexing(
         var: ShapeArg,
-        size: Union[sympy.Expr, int],
+        size: sympy.Expr | int,
         check: bool = True,
         wrap_neg: bool = True,
     ) -> None:
diff --git a/torch/_inductor/sizevars.py b/torch/_inductor/sizevars.py
index 740e87f230f37..df4f3632d6b5e 100644
--- a/torch/_inductor/sizevars.py
+++ b/torch/_inductor/sizevars.py
@@ -2,25 +2,29 @@
 import functools
 import itertools
 import logging
-import sys
-from collections import defaultdict
 from collections.abc import Callable, Iterable, Sequence
-from typing import Any, cast, Optional, Union
+from typing import Any, cast
 
 import sympy
 from sympy import Expr
 
 from torch import SymInt
+from torch.fx.experimental._size_hinting import (
+    _guarding_hint_or_throw_base,
+    _maybe_realize_expr,
+    _optimization_hint_base,
+)
 from torch.fx.experimental.symbolic_shapes import (
     free_symbols,
-    has_free_unbacked_symbols,
+    free_unbacked_symbols,
+    IterateExprs,
     ShapeEnv,
+    SymNode,
 )
 from torch.utils._ordered_set import OrderedSet
-from torch.utils._sympy.functions import FloorDiv, ModularIndexing
-from torch.utils._sympy.numbers import int_oo
+from torch.utils._sympy.functions import FloorDiv, Mod, ModularIndexing
 from torch.utils._sympy.symbol import symbol_is_type, SymT
-from torch.utils._sympy.value_ranges import bound_sympy, IntInfinity, ValueRanges
+from torch.utils._sympy.value_ranges import IntInfinity, ValueRanges
 
 from . import config
 from .runtime.runtime_utils import is_power_of_2
@@ -36,12 +40,31 @@
 
 log = logging.getLogger(__name__)
 
+# Symbols created by CppTemplateKernel.slice_nd → parse_expr_with_index_symbols
+# and are internal to C++ GEMM codegen (not tracked by ShapeEnv).
+_GEMM_TEMPLATE_SYMBOL_NAMES = OrderedSet(
+    [
+        "m_start",
+        "m_end",
+        "n_start",
+        "n_end",
+        "k_start",
+        "k_end",
+        "m_size",
+        "n_size",
+        "m_offset",
+        "m_start_unsliced",
+        "m_end_unsliced",
+        "m_size_unsliced",
+    ]
+)
+
 
 def statically_known_true(
     shape_env: ShapeEnv,
-    expr: Union[sympy.Basic, bool],
-    axioms: Optional[tuple[sympy.Expr]] = None,
-    var_to_range: Optional[tuple[tuple[sympy.Symbol, ValueRanges[Any]]]] = None,
+    expr: sympy.Basic | bool,
+    axioms: tuple[sympy.Expr] | None = None,
+    var_to_range: tuple[tuple[sympy.Symbol, ValueRanges[Any]]] | None = None,
 ) -> bool:
     if expr in (True, False):
         return bool(expr)
@@ -83,7 +106,7 @@ def __init__(self, shape_env=None) -> None:
         self.backed_var_to_val = self.shape_env.backed_var_to_val
         self.var_to_hint_override = self.shape_env.var_to_hint_override
         self.replacements: dict[sympy.Symbol, Expr] = self.shape_env.replacements
-        self.unbacked_replacements: Optional[dict[Expr, Expr]] = None
+        self.unbacked_replacements: dict[Expr, Expr] | None = None
         # Maps of dynamic sizes that have to be precomputed on the host to the kernel args.
         # The basic idea is if we have some complicated sympy expression
         # f(s0), we may choose to precompute it on the host and then replace
@@ -198,19 +221,29 @@ def remove_zero_terms(base, divisor):
 
             for v in base.free_symbols:
                 if v in var_ranges:
-                    # var smaller than divisor can be removed
-                    # if the rest is guaranteed to be multiple of divisor
                     rest = sympy.Wild("_rest", exclude=[v])
                     m = base.match(v + rest)
                     if m and v not in m[rest].free_symbols:
+                        # v can be removed if it doesn't affect the FloorDiv.
+                        # rest is always a multiple of gcd(rest, divisor), so
+                        # rest % divisor is also a multiple of that gcd. The
+                        # worst case is rest % divisor == divisor - gcd, so
+                        # adding v is safe when v < gcd.
                         gcd = sympy.gcd(m[rest], divisor)
-                        if gcd == divisor:
-                            if statically_known(v < divisor):
-                                base = m[rest]
+                        if statically_known(v < gcd):
+                            base = m[rest]
             return base
 
         def visit_indexing_div(base, divisor):
-            return FloorDiv(remove_zero_terms(base, divisor), divisor)
+            base = remove_zero_terms(base, divisor)
+            if statically_known(base >= 0) and statically_known(base < divisor):
+                return sympy.S.Zero
+            # FloorDiv(ModularIndexing(b, d1, m), d2) = ModularIndexing(b, d1*d2, m//d2)
+            if isinstance(base, ModularIndexing) and isinstance(divisor, sympy.Integer):
+                b, d1, m = base.args
+                if m % divisor == 0:
+                    return ModularIndexing(b, d1 * divisor, FloorDiv(m, divisor))
+            return FloorDiv(base, divisor)
 
         def visit_modular_indexing(base, divisor, modulus):
             base = remove_zero_terms(base, divisor)
@@ -330,16 +363,14 @@ def prune(index):
     # asked questions can be answered without guarding otherwise they return False.
     # Those are similar to statically_known_true in symbolic_shapes.py but operate on sympy
     # expressions instead of symnodes.
-    def statically_known_true(self, expr: Union[sympy.Basic, bool]) -> bool:
+    def statically_known_true(self, expr: sympy.Basic | bool) -> bool:
         """
         Returns true if an expression is always true (symbolically or via guards),
         false otherwise. Never add guards, or throw data dependent errors.
         """
         return statically_known_true(self.shape_env, expr)
 
-    def statically_known_equals(
-        self, left: Union[Expr, int], right: Union[Expr, int]
-    ) -> bool:
+    def statically_known_equals(self, left: Expr | int, right: Expr | int) -> bool:
         """
         Returns a bool indicating if it is sound to optimize as if left and right are equal.
         """
@@ -355,36 +386,84 @@ def statically_known_list_equals(
             self.statically_known_equals(l, r) for l, r in zip(left, right)
         )
 
-    def statically_known_leq(self, left: Expr, right: Union[Expr, int]) -> bool:
+    def statically_known_leq(self, left: Expr, right: Expr | int) -> bool:
         """
         Returns a bool indicating if it is sound to optimize as if left is less than or equal to right.
         """
         expr = left <= right
         return self.statically_known_true(expr)
 
-    def statically_known_geq(self, left: Expr, right: Union[Expr, int]) -> bool:
+    def statically_known_geq(self, left: Expr, right: Expr | int) -> bool:
         """
         Returns a bool indicating if it is sound to optimize as if left is greater than or equal to right.
         """
         expr = left >= right
         return self.statically_known_true(expr)
 
-    def statically_known_lt(self, left: Expr, right: Union[Expr, int]) -> bool:
+    def statically_known_lt(self, left: Expr, right: Expr | int) -> bool:
         """
         Returns a bool indicating if it is sound to optimize as if left is less than right.
         """
         expr = left < right
         return self.statically_known_true(expr)
 
-    def statically_known_gt(self, left: Expr, right: Union[Expr, int]) -> bool:
+    def statically_known_gt(self, left: Expr, right: Expr | int) -> bool:
         """
         Returns a bool indicating if it is sound to optimize as if left is greater than right.
         """
         expr = left > right
         return self.statically_known_true(expr)
 
+    def _is_multiple_of(self, numerator: Expr, denominator: int) -> bool:
+        """
+        Structural divisibility check: returns True only if numerator is
+        provably a multiple of denominator.  Recurses over sympy expression
+        structure before falling back to statically_known_true.
+        """
+        # Rule 1 — concrete value
+        if isinstance(numerator, (int, sympy.Integer)):
+            return int(numerator) % denominator == 0
+
+        # Rule 2 — product: any factor divisible → product divisible
+        if isinstance(numerator, sympy.Mul):
+            for factor in numerator.args:
+                if self._is_multiple_of(factor, denominator):
+                    return True
+            # Also check if combined constant factors are divisible
+            const = 1
+            for factor in numerator.args:
+                if isinstance(factor, (int, sympy.Integer)):
+                    const *= int(factor)
+            if const != 1 and const % denominator == 0:
+                return True
+
+        # Rule 3 — sum: all terms divisible → sum divisible
+        if isinstance(numerator, sympy.Add):
+            if all(self._is_multiple_of(term, denominator) for term in numerator.args):
+                return True
+
+        # Rule 4 — FloorDiv(a, b): if a is multiple of b*n
+        if isinstance(numerator, FloorDiv):
+            a, b = numerator.args
+            if isinstance(b, (int, sympy.Integer)):
+                if self._is_multiple_of(a, int(b) * denominator):
+                    return True
+
+        # Rule 5 — Mod(a, b): Mod(a,b) = a - b*floor(a/b), so if both a and b
+        # are multiples of n, then Mod(a,b) is too.
+        if isinstance(numerator, (Mod, sympy.Mod)):
+            a, b = numerator.args
+            if self._is_multiple_of(a, denominator) and self._is_multiple_of(
+                b, denominator
+            ):
+                return True
+
+        # Rule 6 — axiom fallback: ask ShapeEnv
+        expr = sympy.Eq(Mod(numerator, denominator), 0)
+        return self.statically_known_true(expr)
+
     def statically_known_multiple_of(
-        self, numerator: Expr, denominator: Union[Expr, int]
+        self, numerator: Expr, denominator: Expr | int
     ) -> bool:
         """
         Return a bool indicating if it is sound to optimize for the numerator being a multiple of the denominator.
@@ -395,7 +474,11 @@ def statically_known_multiple_of(
         if len(free_symbols(numerator)) > 20:
             return False
 
-        expr = sympy.Eq(numerator % denominator, 0)
+        if isinstance(denominator, (int, sympy.Integer)):
+            return self._is_multiple_of(numerator, int(denominator))
+
+        # For symbolic denominators, fall back to direct sympy check
+        expr = sympy.Eq(Mod(numerator, denominator), 0)
         return self.statically_known_true(expr)  # type: ignore[arg-type]
 
     def statically_known_power_of_2(self, expr: Expr) -> bool:
@@ -484,9 +567,9 @@ def guard_or_true(self, left):
     # which does the wrong thing if a or b is a sympy expression
     def evaluate_expr(
         self,
-        left: Union[Expr, sympy.logic.boolalg.Boolean],
+        left: Expr | sympy.logic.boolalg.Boolean,
         size_oblivious: bool = False,
-        fallback_value: Optional[bool] = None,
+        fallback_value: bool | None = None,
     ) -> bool:
         assert isinstance(left, (Expr, sympy.logic.boolalg.Boolean)), type(left)
         return self.shape_env.evaluate_expr(
@@ -503,42 +586,57 @@ def is_size_one_or_false(self, size: Expr) -> bool:
         return self.guard_or_false(sympy.Eq(size, 1))
 
     def evaluate_min(self, left: Expr, right: Expr) -> Expr:
-        """return the smaller of left and right, and guard on that choice"""
+        """Return the smaller of left and right, and guard on that choice."""
         if isinstance(left, Expr):
             left = sympy_subs(left, self.inv_precomputed_replacements)  # type: ignore[arg-type]
         if isinstance(right, Expr):
             right = sympy_subs(right, self.inv_precomputed_replacements)  # type: ignore[arg-type]
-        try:
-            lv = self.size_hint_or_throw(left)
-            rv = self.size_hint_or_throw(right)
-        except TypeError:  # unbacked symints
-            if left == right or self.statically_known_leq(left, right):
-                return left
-            if self.statically_known_leq(right, left):
-                return right
-            gcd = sympy.gcd(left, right)
-            if left == gcd:  # handle `min(10*u0, u0)` etc
-                return left
-            if right == gcd:
-                return right
-            raise TypeError(
-                f"evaluate_min({left}, {right}) with unbacked symints"
-            ) from None
-        if lv <= rv:
-            self.check_leq(left, right)
+        if self.guard_or_false(sympy.Le(left, right)):
             return left
-        else:
-            self.check_leq(right, left)
+        if self.guard_or_false(sympy.Le(right, left)):
             return right
 
+        # GCD fallback: if gcd(a, b) == a then a divides b, implying a <= b.
+        #
+        # TODO: This is NOT always sound for unbacked symints.  It can
+        # produce wrong results when:
+        #   - inputs can be negative: gcd(u0, 10*u0) = u0, returns u0,
+        #     but if u0 < 0 then u0 > 10*u0 (e.g. u0=-1: min(-1,-10) = -10)
+        #   - a factor can be zero: gcd(u0, u0*u1) = u0, returns u0,
+        #     but if u1=0 then u0*u1=0 < u0 (e.g. u0=5,u1=0: min(5,0) = 0)
+        # TODO shall we add a runtime assertion at least.
+        gcd = sympy.gcd(left, right)
+        if left == gcd:
+            return left
+        if right == gcd:
+            return right
+
+        # Min/Max fallback: we can prove Min(a, b) <= c when any arg <= c, but
+        # sympy doesn't simplify this yet. So, evaluate it here. Same for Max.
+        for lhs, rhs in [(left, right), (right, left)]:
+
+            def le_rhs(a: Expr) -> bool:
+                return self.guard_or_false(sympy.Le(a, rhs))
+
+            # Min(Min(a, b), c) ==> Min(a, b) if (a <= c) or (b <= c).
+            if isinstance(lhs, sympy.Min) and any(le_rhs(a) for a in lhs.args):
+                return lhs
+            # Min(Max(a, b), c) ==> Max(a, b) if (a <= c) and (b <= c).
+            if isinstance(lhs, sympy.Max) and all(le_rhs(a) for a in lhs.args):
+                return lhs
+
+        raise TypeError(
+            f"evaluate_min({left}, {right}) with unbacked symints"
+        ) from None
+
     def evaluate_max(self, left: Expr, right: Expr) -> Expr:
-        """return the larger of left and right, and guard on that choice"""
+        """Return the larger of left and right, and guard on that choice."""
         # Always choose the opposite of eval min for consistency
         # This means min(a, b) and max(a, b) produce the same guards
         min_val = self.evaluate_min(left, right)
         return right if min_val is left else left
 
-    def guard_int(self, expr: Union[Expr, int]) -> int:
+    def guard_int(self, expr: Expr | int) -> int:
         """
         Similar to guard_int in symbolic_shapes.py, except this function works with SymPy
         expressions instead of SymNodes. It extracts the value represented by expr from shapeEnv
@@ -547,11 +645,11 @@ def guard_int(self, expr: Union[Expr, int]) -> int:
         """
         if isinstance(expr, int):
             return expr
-        val = self.size_hint_or_throw(expr)
+        val = self.guarding_hint_or_throw(expr)
         self.check_equals(expr, sympy.Integer(val))
         return int(val)
 
-    def guard_int_seq(self, left: Sequence[Union[Expr, int]]) -> list[int]:
+    def guard_int_seq(self, left: Sequence[Expr | int]) -> list[int]:
         """
         Apply guard_int on a sequence of inputs.
         """
@@ -562,14 +660,18 @@ def remove_precomputed_replacements(self, expr: Expr) -> Expr:
             return sympy_subs(expr, self.inv_precomputed_replacements)  # type: ignore[arg-type]
         return expr
 
-    def symbolic_hint(
+    def replace_backed_symbols_with_hints(
         self,
-        expr: Union[Expr, int],
-        # Only flip this flag if you don't plan on guarding/adding runtime
-        # asserts based on this value and promise to only use this value
-        # in a heuristic nature.
-        use_user_provided_hint_override: bool = False,
-    ) -> Union[Expr, int]:
+        expr: Expr | int,
+    ) -> Expr | int:
+        """
+        Replace all backed symbols in an expression with their concrete hint values.
+
+        This function substitutes backed symbolic variables with their known concrete
+        values from the shape environment, while leaving unbacked symbols (data-dependent
+        values) untouched. Not recommended to be used except as last resort!
+        was added for existing usages compatibility.
+        """
         if isinstance(expr, int):
             return expr
         # Substitute all hints into expr, but leave unbacked symints alone
@@ -577,21 +679,22 @@ def symbolic_hint(
         if not isinstance(expr, Expr):
             assert isinstance(expr, int)
             return expr
+
+        expr = self.remove_precomputed_replacements(expr)
+        expr = sympy_subs(expr, self.backed_var_to_val)
+        if isinstance(expr, Expr):
+            expr = expr.expand(identity=True)
+
         free_symbols = expr.free_symbols
         if not free_symbols:
             try:
-                return int(expr)  # type: ignore[return-value]
+                return int(expr)
             except TypeError:
                 return expr  # inf/nan/I
 
-        expr = self.remove_precomputed_replacements(expr)
-
-        if use_user_provided_hint_override:
-            expr = sympy_subs(expr, self.var_to_hint_override)
-
-        return sympy_subs(expr, self.backed_var_to_val)
+        return expr
 
-    def to_symint_or_int(self, expr: Union[Expr, int]) -> Union[SymInt, int]:
+    def to_symint_or_int(self, expr: Expr | int) -> SymInt | int:
         """Convert a sympy expression to SymInt, or return int as is."""
         if isinstance(expr, int):
             return expr
@@ -607,96 +710,56 @@ def to_symint_or_int(self, expr: Union[Expr, int]) -> Union[SymInt, int]:
         expr = self.remove_precomputed_replacements(expr)
 
         try:
-            hint = self.size_hint(expr)
+            hint = self.guarding_hint_or_throw(expr)
         except Exception:
             hint = None
         node = SymNode(expr, self.shape_env, int, hint)
         return SymInt(node)
 
-    def to_symints_or_ints(
-        self, exprs: Sequence[Union[Expr, int]]
-    ) -> list[Union[SymInt, int]]:
+    def to_symints_or_ints(self, exprs: Sequence[Expr | int]) -> list[SymInt | int]:
         """Convert a sequence of sympy expressions to SymInts, or return ints as is."""
         return [self.to_symint_or_int(e) for e in exprs]
 
-    def size_hint(
-        self,
-        expr: Union[Expr, int],
-        *,
-        fallback: Optional[int] = None,
-    ) -> int:
-        if isinstance(expr, SymInt):
-            raise TypeError(
-                "wrong API usage!, use size_hint from torch.fx.experimental.symbolic_shapes or pass sympy expressions instead"
-            )
-
-        out = self.symbolic_hint(
-            expr,
-            use_user_provided_hint_override=fallback is not None,
-        )
-        if not isinstance(out, (int, sympy.Integer)) and fallback is not None:
-            # Use the provided heuristic fallback hint
-            unbacked_sym_vrs = {
-                s: self.shape_env.var_to_range.get(s, None) for s in out.free_symbols
-            }
-            if all(vr is not None for vr in unbacked_sym_vrs.values()):
-                hint_vr = bound_sympy(out, unbacked_sym_vrs)  # type: ignore[arg-type]
-                if isinstance(hint_vr.lower, (int, sympy.Integer)):
-                    fallback = max(fallback, int(hint_vr.lower))
-                if isinstance(hint_vr.upper, (int, sympy.Integer)):
-                    fallback = min(fallback, int(hint_vr.upper))
-            return fallback
+    def guarding_hint_or_throw(self, expr: Expr | int) -> int:
+        """
+        Return a concrete integer hint for an expression that is safe to use for guarding.
 
-        try:
-            return int(out)
-        except Exception:
-            log.debug("failed on: %s", out)
-            raise
+        This method evaluates the expression using only backed-symbols hints. Unlike
+        optimization_hint(), this method does NOT use heuristics or fallback values
+        for unbacked symbols.
 
-    def size_hint_or_throw(self, expr: Union[Expr, int]) -> int:
-        # Like size_hint but there's no fallback for unbacked symints, so it throws.
-        out = self.symbolic_hint(expr)
-        try:
-            return int(out)
-        except Exception:
-            log.debug("failed on: %s", out, exc_info=True)
-            raise
+        Use this method when you need a hint value that will be used for guarding decision.
 
-    def _maybe_realize_expr(self, expr: Expr, fallback: Optional[int]) -> Optional[int]:
-        """
-        Handle special sympy values in optimization hints.
+        Args:
+            expr: A sympy expression or integer to evaluate.
 
         Returns:
-            - Raises ValueError for complex numbers
-            - sys.maxsize for positive infinity
-            - -sys.maxsize for negative infinity
-            - fallback for NaN
-            - None if no special handling needed
-        """
-        if isinstance(expr, (int, sympy.Integer, float, sympy.Float, sympy.Rational)):
-            return int(expr)
+            The concrete integer value of the expression based on backed symbol hints.
 
-        if isinstance(expr, Expr):
-            if expr.has(sympy.I):
-                raise ValueError(
-                    f"optimization_hint received a complex expression: {expr}. "
-                    "Tensor dimensions cannot be complex numbers."
-                )
-            if expr in (int_oo, sympy.oo):
-                return sys.maxsize
-            if expr in (-int_oo, -sympy.oo):
-                return -sys.maxsize
-            if fallback is not None and expr is sympy.nan or expr.has(sympy.nan):
-                return fallback
+        Raises:
+            GuardOnDataDependentSymNode: If the expression contains unbacked symbols
+            (data-dependent values) that cannot be resolved to concrete values.
 
-        return None
+        See Also:
+            optimization_hint: For cases where fallback/heuristic values are acceptable
+                for unbacked symbols.
+        """
+        if isinstance(expr, SymNode):
+            raise TypeError(
+                f"guarding_hint_or_throw expects a sympy Expr or int, not {type(expr)}. "
+                "Use expr.expr to extract the sympy expression from a SymNode."
+            )
+        return _guarding_hint_or_throw_base(
+            self.shape_env, expr, self.inv_precomputed_replacements
+        )
 
-    def optimization_hint(
-        self, expr: Union[Expr, int], fallback: Optional[int] = None
-    ) -> int:
+    def optimization_hint(self, expr: Expr | int, fallback: int | None = None) -> int:
         """
         Return a concrete integer hint for an expression.
 
+        This function should be used for non-guarding based optimizations. If you
+        want a hint that you can guard on, use the guarding_hint API instead.
+
         This function will hint unbacked symbols using user provided optimization
         hints. If not provided, fallback will be used along with some heuristics
         that try to maximize consistency with the shape environment.
@@ -707,71 +770,14 @@ def optimization_hint(
         - Infinity (int_oo, sympy.oo): returns sys.maxsize.
         - NaN (sympy.nan): returns the fallback value.
         """
-        # Read config at call time to respect runtime patches (e.g., in tests)
-        if fallback is None:
-            fallback = config.unbacked_symint_fallback
-        assert fallback is not None
-        simplified = self.simplify(expr)
-        result = self._maybe_realize_expr(simplified, fallback)
-        if result is not None:
-            return result
-
-        # remove precomputed_replacements
-        expr = self.remove_precomputed_replacements(simplified)
-        original = expr
-
-        # replace all backed symbols with their backed hints,
-        # unbacked with optimizations hints if exists.
-        expr = sympy_subs(expr, self.backed_var_to_val)
-        expr = sympy_subs(expr, self.var_to_hint_override)
-
-        result = self._maybe_realize_expr(expr, fallback)
-        if result is not None:
-            return result
-
-        # Assign values to remaining unbacked symbols using a heuristic
-        # tries to maximize consistency with shape environment.
-        assert has_free_unbacked_symbols(expr), expr
-
-        # Make sure to substitute with the factored version
-        # e.g. 10*(s0 + u0) instead of 10*s0 + 10*u0
-        # TODO optimize _sub_unbacked_exprs
-        expr = self._sub_unbacked_exprs(sympy.factor(original))
-
-        expr = sympy_subs(expr, self.backed_var_to_val)
-        expr = sympy_subs(expr, self.var_to_hint_override)
-
-        # For multiple expressions that depend on an unbacked symint,
-        # we want to compute them consistently for a size hint we have chosen.
-        # So, recursively compute expressions via size hints of contained symbols.
-        # For example: u1 * u2 - 10 ==> fallback * fallback - 10
-
-        assert isinstance(expr, Expr), type(expr)
-        free_symbols = expr.free_symbols
-
-        # Constrain fallback per-symbol based on var_to_range bounds
-        size_dict = {}
-        for s in free_symbols:
-            sym_fallback = fallback
-            vr = self.shape_env.var_to_range.get(s, None)
-            if vr is not None:
-                if isinstance(vr.lower, (int, sympy.Integer)):
-                    sym_fallback = max(sym_fallback, int(vr.lower))
-                if isinstance(vr.upper, (int, sympy.Integer)):
-                    sym_fallback = min(sym_fallback, int(vr.upper))
-            size_dict[s] = sym_fallback
-
-        final_result = expr.subs(size_dict)
-
-        final_result = self._maybe_realize_expr(final_result, fallback)
-        assert final_result is not None, final_result
-
-        return int(final_result)
+        return _optimization_hint_base(
+            self.shape_env, expr, self.inv_precomputed_replacements, fallback
+        )
 
     def optimization_hints(
         self,
-        exprs: Iterable[Union[Expr, int]],
-        fallback: Optional[int] = None,
+        exprs: Iterable[Expr | int],
+        fallback: int | None = None,
     ) -> tuple[int, ...]:
         """
         Like optimization_hint but for a sequence of expressions.
@@ -782,10 +788,21 @@ def optimization_hints(
             fallback = config.unbacked_symint_fallback
         return tuple(self.optimization_hint(x, fallback=fallback) for x in exprs)
 
+    def all_unbacked_explicitly_hinted(self, exprs: IterateExprs) -> bool:
+        """
+        Return True if every unbacked symbol in *exprs* has an explicit
+        user-provided hint in var_to_hint_override.  If there are no
+        unbacked symbols at all, returns True (vacuously).
+        *exprs* can be a single expression or any iterable accepted by
+        free_unbacked_symbols (list, tuple, etc.).
+        """
+        unbacked = free_unbacked_symbols(exprs)
+        return unbacked.issubset(self.var_to_hint_override.keys())
+
     def optimization_hint_with_override(
         self,
-        expr: Union[Expr, int],
-        hint_override: Optional[int],
+        expr: Expr | int,
+        hint_override: int | None,
     ) -> int:
         r"""Return a concrete integer hint for an expression, with optional override.
         This is used in dynamic dispatch scenarios where callers may want to
@@ -803,7 +820,7 @@ def optimization_hint_with_override(
         Returns:
             int: A concrete integer hint for the expression.
         """
-        simplified = self._maybe_realize_expr(self.simplify(expr), None)
+        simplified = _maybe_realize_expr(self.simplify(expr), None)
 
         if simplified is not None:
             return simplified
@@ -817,8 +834,8 @@ def optimization_hint_with_override(
 
     def optimization_hints_with_override(
         self,
-        exprs: Iterable[Union[Expr, int]],
-        hint_override: Optional[int],
+        exprs: Iterable[Expr | int],
+        hint_override: int | None,
     ) -> tuple[int, ...]:
         """
         Like optimization_hint_with_override but for a sequence of expressions.
@@ -828,26 +845,11 @@ def optimization_hints_with_override(
             self.optimization_hint_with_override(e, hint_override) for e in exprs
         )
 
-    def size_hints(
+    def guarding_hints_or_throw(
         self,
-        exprs: Iterable[Union[Expr, int]],
-        *,
-        fallback: Optional[int] = None,
+        exprs: Iterable[Expr | int],
     ) -> tuple[int, ...]:
-        return tuple(
-            self.size_hint(
-                x,
-                fallback=fallback,
-            )
-            for x in exprs
-        )
-
-    def size_hints_or_throw(
-        self,
-        exprs: Iterable[Union[Expr, int]],
-    ) -> tuple[int, ...]:
-        # Like size_hints but there's no fallback for unbacked symints, so it throws.
-        return tuple(self.size_hint_or_throw(x) for x in exprs)
+        return tuple(self.guarding_hint_or_throw(x) for x in exprs)
 
     def _lru_cache(self, fn, maxsize=None):
         """
@@ -873,7 +875,7 @@ def make_stride_vars_cache(self):
         def stride_vars(
             index: Expr,
             vars: Sequence[sympy.Symbol],
-            support_vars: Optional[Sequence[sympy.Symbol]] = None,
+            support_vars: Sequence[sympy.Symbol] | None = None,
         ) -> list[Expr]:
             if not support_vars:
                 support_vars = vars
@@ -921,204 +923,25 @@ def _stride_vars(
                 )
         return strides
 
-    def _get_unbacked_replacements(self) -> dict[Expr, Expr]:
-        if self.unbacked_replacements is not None:
-            return self.unbacked_replacements
-
-        class CanonicalExprFinder:
-            """
-            Purpose:
-            A disjoint-set/union-find data structure that can return the
-            "canonical" expression for a group of equivalent expressions.
-            - The canonical expression must come from the input eq_graph.
-            - The heuristics used to choose a leader determines which
-            expression becomes the canonical expression.
-
-            Problem:
-            Given any unbacked expression, we should be able to find a size_hint
-            for the unbacked expression, that adheres to the ShapeEnv's deferred
-            runtime assertions. Otherwise, we may generate conflicting size hints.
-            In other words, even though we know u0 + s0 == u2, we may generate
-            size hints, such that, size_hint(u0 + s0) != size_hint(u2).
-            NOTE: At this time, only deferred runtime asserts that are equalities
-            (i.e. Eq(lhs, rhs)) are considered in this data structure.
-
-            Examples:
-            - u0 + u1 == 9000, then find_expr(u0 + u1) == find_expr(9000)
-            - u0 + u1 == s9, then find_expr(u0 + u1) == find_expr(s9)
-            - u0 + s0 == u10, then find_expr(u0 + s0) == find_expr(u10)
-
-            Inputs:
-            - equality_graph: An adjacency set of expressions where the edge
-            connects two expressions that are found equal to each other. The
-            edges are sourced from ShapeEnv's deferred_runtime_asserts.
-
-            Usage:
-            - Call union_expr(a, b) to merge a & b into a single set which
-            shares the same canonical expression.
-            - Call find_expr(x) to find the canonical expression for x.
-            """
-
-            def __init__(self, eq_graph: dict[Expr, OrderedSet[Expr]]):
-                self.eq_graph = eq_graph
-                self.expressions = list(eq_graph.keys())
-                self.reverse_expressions = {
-                    expr: i for i, expr in enumerate(self.expressions)
-                }
-                # Each node is its own leader/parent initially
-                self.leader = list(range(len(self.expressions)))
-                # Track size for union-by-size
-                self.size = [1] * len(self.expressions)
-
-                # Takes each edge from the undirected graph and starts merging them.
-                self._build_canonical_expr_mapping()
-
-            def _build_canonical_expr_mapping(self):
-                for expr, edges in self.eq_graph.items():
-                    for adj in edges:
-                        self.union_expr(expr, adj)
-
-            def union_expr(self, a: Expr, b: Expr):
-                return self.union(
-                    self.reverse_expressions[a], self.reverse_expressions[b]
-                )
-
-            def union(self, a: int, b: int):
-                rootA = self.find(a)
-                rootB = self.find(b)
-                if rootA == rootB:
-                    return False  # already connected
-                leader, other = self.choose_leader(rootA, rootB)
-                self.leader[other] = leader
-                self.size[leader] += self.size[other]
-                return True
-
-            def find_expr(self, expr: Expr):
-                parent = self.find(self.reverse_expressions[expr])
-                return self.expressions[parent]
-
-            def find(self, x: int):
-                # Path compression
-                if self.leader[x] != x:
-                    self.leader[x] = self.find(self.leader[x])
-                return self.leader[x]
-
-            def choose_leader(self, a: int, b: int):
-                """
-                The leader will become the canonical expression.
-                Returns a (leader, follower) tuple.
-
-                Here are the heuristics used for choosing a leader:
-                1. Backed expression or constants preferred over unbacked expr
-                2. Simpler sub-expr when one contains the other
-                3. Higher frequency across equalities from deferred runtime assertions
-                4. Size of the set
-                5. Fallback to sympy.Basic.compare
-                """
-
-                def _choose(x: int, y: int) -> bool:
-                    lhs, rhs = self.expressions[x], self.expressions[y]
-
-                    # Prefer replacing unbacked exprs with backed expressions/constants.
-                    # Examples:
-                    # u0 + s3 ==> s0 + s1, then leader is s0 + s1
-                    # u2 ==> 300, then leader is 300
-                    any_unbacked_lhs = has_free_unbacked_symbols(lhs)
-                    any_unbacked_rhs = has_free_unbacked_symbols(rhs)
-                    if any_unbacked_lhs != any_unbacked_rhs:
-                        return bool(any_unbacked_rhs)
-
-                    # Handles cases where LHS contains the RHS. In other words,
-                    # RHS is a sub-expression of LHS. For example:
-                    # s1 * Max(2, u0) ==> Max(2, u0), then leader is Max(2, u0)
-                    if lhs.has(rhs):
-                        return False
-                    elif rhs.has(lhs):
-                        return True
-
-                    # Prefer expressions that come up more often.
-                    degrees_lhs = len(self.eq_graph[lhs])
-                    degrees_rhs = len(self.eq_graph[rhs])
-                    if degrees_lhs != degrees_rhs:
-                        return degrees_lhs > degrees_rhs
-
-                    # Try to apply union-by-size optimization to flatten the
-                    # leader trees.
-                    if self.size[x] != self.size[y]:
-                        return self.size[x] > self.size[y]
-
-                    # Fallback to sympy.Basic.compare for a deterministic ordering.
-                    return lhs.compare(rhs) == -1
-
-                if _choose(a, b):
-                    return a, b
-                return b, a
-
-        # Build an undirected graph using ShapeEnv's deferred runtime assertions.
-        self.equality_graph: dict[Expr, OrderedSet[Expr]] = defaultdict(OrderedSet)
-        for assertions in self.shape_env.deferred_runtime_asserts.values():
-            for assertion in assertions:
-                if not isinstance(assertion.expr, sympy.Equality):
-                    # We're ignoring other relationals for now. If you need to
-                    # account for relationals, then you may need a solver solution.
-                    continue
-                lhs = sympy.sympify(assertion.expr.lhs)  # sympify helps with ints
-                rhs = sympy.sympify(assertion.expr.rhs)
-                self.equality_graph[lhs].add(rhs)
-                self.equality_graph[rhs].add(lhs)
-
-        # Use the undirected graph to create a DSU data structure, so we can
-        # query for a "canonical" expression.
-        uf = CanonicalExprFinder(self.equality_graph)
-
-        # Start building the unbacked replacements mapping using CanonicalExprFinder
-        # The mapping is from Expr to its "canonical" Expr.
-        self.unbacked_replacements = {}
-        for expr in self.equality_graph:
-            canonical_expr = uf.find_expr(expr)
-            if expr != canonical_expr:
-                self.unbacked_replacements[expr] = canonical_expr
-
-        return self.unbacked_replacements
-
-    @functools.lru_cache  # noqa: B019
-    def _sub_unbacked_exprs(self, expr: Expr) -> Expr:
-        # it's fine to cache this fn since self is a singleton
-        replacements = self._get_unbacked_replacements()
-
-        # consider making this threshold configurable
-        sub_cnt_limit = 30
-        sub_cnt = 0
-        while sub_cnt < sub_cnt_limit:
-            new_expr = expr.subs(replacements)
-            if new_expr == expr:
-                return new_expr
-            expr = sympy.factor(new_expr)
-            sub_cnt += 1
-
-        log.warning("Substitution limit (%d) reached w/ %s", sub_cnt_limit, expr)
-        return expr
-
     def offset_var(self, index: Expr, vars: Sequence[sympy.Symbol]) -> Expr:
         """Extract offset part of an indexing expression"""
         index = self.simplify(index)
         return sympy_subs(index, {v: sympy.S.Zero for v in vars if v != 0})
 
+    # Return stride optimizaitons hints,
+    # only used for optimizations.
     def stride_hints(
         self,
         index: Expr,
         vars: Sequence[sympy.Symbol],
-        support_vars: Optional[Sequence[sympy.Symbol]] = None,
+        support_vars: Sequence[sympy.Symbol] | None = None,
     ) -> list[int]:
         for v in index.free_symbols:
             if symbol_is_type(v, SymT.INDIRECT):  # type: ignore[attr-defined]
                 index = sympy_subs(index, {v: 0})  # type: ignore[dict-item]
         result = []
         for s in self.stride_vars(index, vars, support_vars):
-            try:
-                result.append(self.size_hint_or_throw(s))
-            except TypeError:
-                result.append(0)
+            result.append(self.optimization_hint(s, fallback=0))
         return result
 
     def stride_order(self, index: Expr, vars: list[sympy.Symbol]) -> list[int]:
@@ -1191,7 +1014,7 @@ def _check_args(x, div, mod, is_first):
             if not _check_args(x2, div2, mod2, False):
                 return index
 
-            if mod2 % mod != 0:
+            if Mod(mod2, mod) != 0:
                 return index
 
             return ModularIndexing(x2, 1, mod)
@@ -1200,7 +1023,7 @@ def _check_args(x, div, mod, is_first):
 
     def expand_floor_div(
         self, index: sympy.Expr
-    ) -> Union[bool, tuple[sympy.Expr, sympy.Expr]]:
+    ) -> bool | tuple[sympy.Expr, sympy.Expr]:
         """
         Expand the FloorDiv to the entire expression so that the expression may
         be simplified.
diff --git a/torch/_inductor/standalone_compile.py b/torch/_inductor/standalone_compile.py
index b4535b89ad9c7..bd8876a17499e 100644
--- a/torch/_inductor/standalone_compile.py
+++ b/torch/_inductor/standalone_compile.py
@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import contextlib
 import copy
 import logging
 import os
@@ -7,7 +8,7 @@
 import shutil
 from abc import ABC, abstractmethod
 from contextlib import AbstractContextManager, nullcontext
-from typing import Any, Literal, Optional, TYPE_CHECKING
+from typing import Any, Literal, TYPE_CHECKING
 
 import torch.fx
 from torch._dynamo.aot_compile_types import BundledAOTAutogradSerializableCallable
@@ -53,7 +54,7 @@ class CompiledArtifact(ABC):
     def __init__(
         self,
         compiled_fn: Callable[..., Any],
-        artifacts: Optional[tuple[bytes, CacheInfo]],
+        artifacts: tuple[bytes, CacheInfo] | None,
     ):
         self._compiled_fn = compiled_fn
         self._artifacts = artifacts
@@ -113,7 +114,7 @@ class CacheCompiledArtifact(CompiledArtifact):
     def __init__(
         self,
         compiled_fn: Callable[..., Any],
-        artifacts: Optional[tuple[bytes, CacheInfo]],
+        artifacts: tuple[bytes, CacheInfo] | None,
     ):
         self._compiled_fn = compiled_fn
         self._artifacts = artifacts
@@ -371,35 +372,22 @@ def load(
             return AOTCompiledArtifact.deserialize(artifact)
 
 
-def standalone_compile(
-    gm: GraphModule,
-    example_inputs: Sequence[InputType],
-    *,
-    dynamic_shapes: Any,
-    options: Any,
-    aot: bool = False,  # AOT mode, which uses BundledAOTAutogradCache
-) -> CompiledArtifact:
-    """
-    Implementation of torch.inductor.standalone_compile
-    """
-    from torch.compiler._cache import CacheArtifactManager
+def _resolve_ignore_shape_env(dynamic_shapes: Any):
+    # tells compile_fx to ignore the shape_envs on the ambient context
+    # and the graph_module.
+    return dynamic_shapes == "from_example_inputs"
 
-    from .compile_fx import compile_fx
 
-    ignore_shape_env = False
+def _resolve_fake_mode(gm: GraphModule, dynamic_shapes: Any) -> FakeTensorMode:
     if dynamic_shapes == "from_example_inputs":
-        fake_mode = FakeTensorMode(shape_env=ShapeEnv())
-        # tells compile_fx to ignore the shape_envs on the ambient context
-        # and the graph_module.
-        ignore_shape_env = True
+        return FakeTensorMode(shape_env=ShapeEnv())
     elif dynamic_shapes == "from_tracing_context":
         # Reuse fake_mode from the TracingContext.
         # NB: The TracingContext only exists if we're currently in a torch.compile backend.
         context = torch._guards.TracingContext.get()
         assert context.fake_mode is not None
-        fake_mode = context.fake_mode
+        return context.fake_mode
     elif dynamic_shapes == "from_graph":
-        fake_mode = FakeTensorMode(shape_env=ShapeEnv())
         # Strategy: find a FakeTensor in the graph output, grab its FakeTensorMode.
         # The graph passed to standalone_compile must be an Inductor-approved graph,
         # which means that there is at least one Tensor output and the output node
@@ -408,39 +396,64 @@ def standalone_compile(
         assert last_node.op == "output"
         assert len(last_node.args) == 1
 
-        def handle_node(node: torch.fx.Node) -> None:
-            nonlocal fake_mode
-            if "example_value" in node.meta:
-                maybe_tensor = node.meta["example_value"]
-                if isinstance(maybe_tensor, torch._subclasses.fake_tensor.FakeTensor):
-                    fake_mode = maybe_tensor.fake_mode
-
         # If gm came from Dynamo, then last_node.args[0] is always a list,
         # even in single-Tensor returns.
         #
         # It's possible to get into a situation where last_node.args[0]
         # is a Node (and not a list!). This happens if you call split_module
         # on the graph. We allow for this case since it is common.
-        if isinstance(last_node.args[0], torch.fx.Node):
-            handle_node(last_node.args[0])
-        else:
-            for node in last_node.args[0]:
-                handle_node(node)
+        nodes = (
+            [last_node.args[0]]
+            if isinstance(last_node.args[0], torch.fx.Node)
+            else last_node.args[0]
+        )
+        for node in nodes:
+            if "example_value" in node.meta:
+                maybe_tensor = node.meta["example_value"]
+                if isinstance(maybe_tensor, torch._subclasses.fake_tensor.FakeTensor):
+                    return maybe_tensor.fake_mode
 
+        return FakeTensorMode(shape_env=ShapeEnv())
     else:
         raise ValueError(
             f"standalone_compile got unsupported `dynamic_shapes` value: dynamic_shapes={dynamic_shapes}."
         )
 
-    context = torch._guards.TracingContext(fake_mode)
+
+@contextlib.contextmanager
+def _standalone_context(gm: GraphModule, dynamic_shapes: Any, aot: bool):
+    from torch.compiler._cache import CacheArtifactManager
+
+    fake_mode = _resolve_fake_mode(gm, dynamic_shapes)
+    tracing_context = torch._guards.TracingContext(fake_mode)
     with (
-        torch._guards.tracing(context),
+        torch._guards.tracing(tracing_context),
         CacheArtifactManager.with_fresh_cache(),
         config.patch("triton.autotune_at_compile_time", True),
         torch._functorch.config.patch("bundled_autograd_cache", aot),
     ):
-        # compile_fx can mutate gm
-        gm = copy.deepcopy(gm)
+        yield
+
+
+def standalone_compile(
+    gm: GraphModule,
+    example_inputs: Sequence[InputType],
+    *,
+    dynamic_shapes: Any,
+    options: Any,
+    aot: bool = False,  # AOT mode, which uses BundledAOTAutogradCache
+    donate_graph_module: bool = False,
+) -> CompiledArtifact:
+    """
+    Implementation of torch.inductor.standalone_compile
+    """
+    from .compile_fx import compile_fx
+
+    ignore_shape_env = _resolve_ignore_shape_env(dynamic_shapes)
+    with _standalone_context(gm, dynamic_shapes, aot):
+        # compile_fx takes ownership of gm and may mutate it on cache miss.
+        if not donate_graph_module:
+            gm = copy.deepcopy(gm)
         compiled_fn = compile_fx(
             gm, example_inputs, ignore_shape_env=ignore_shape_env, **options
         )
@@ -459,3 +472,20 @@ def handle_node(node: torch.fx.Node) -> None:
             )
 
     return CacheCompiledArtifact(compiled_fn, artifacts)
+
+
+def autograd_cache_key(
+    graph,
+    example_inputs,
+    dynamic_shapes: Any,
+    aot: bool = False,  # AOT mode, which uses BundledAOTAutogradCache
+):
+    from . import compile_fx
+
+    ignore_shape_env = _resolve_ignore_shape_env(dynamic_shapes)
+    with _standalone_context(graph, dynamic_shapes, aot):
+        return compile_fx.autograd_cache_key(
+            graph,
+            example_inputs,
+            ignore_shape_env=ignore_shape_env,
+        )
diff --git a/torch/_inductor/stream_constants.py b/torch/_inductor/stream_constants.py
new file mode 100644
index 0000000000000..40ca199461746
--- /dev/null
+++ b/torch/_inductor/stream_constants.py
@@ -0,0 +1,5 @@
+"""Constants for CUDA stream naming in Inductor codegen."""
+
+DEFAULT_STREAM: str = "default_stream"
+DEFAULT_STREAM_IDX: int = 0
+STREAM_NAME_TEMPLATE: str = "stream{stream_idx:d}"
diff --git a/torch/_inductor/stream_utils.py b/torch/_inductor/stream_utils.py
new file mode 100644
index 0000000000000..9e356753eda92
--- /dev/null
+++ b/torch/_inductor/stream_utils.py
@@ -0,0 +1,39 @@
+"""Stream utilities for Inductor codegen."""
+
+from __future__ import annotations
+
+import functools
+
+from torch._inductor.stream_constants import (
+    DEFAULT_STREAM,
+    DEFAULT_STREAM_IDX,
+    STREAM_NAME_TEMPLATE,
+)
+
+
+__all__ = [
+    "DEFAULT_STREAM",
+    "DEFAULT_STREAM_IDX",
+    "STREAM_NAME_TEMPLATE",
+    "get_raw_stream_name",
+    "get_stream_name",
+]
+
+
+@functools.lru_cache
+def get_stream_name(stream_idx: int) -> str:
+    """Generate CUDA Stream name from stream index number.
+
+    Args:
+        stream_idx: Non-negative index number. 0 refers to the default stream, others refer to side
+            streams.
+    """
+    if stream_idx == 0:
+        return DEFAULT_STREAM
+    else:
+        return STREAM_NAME_TEMPLATE.format(stream_idx=stream_idx)
+
+
+def get_raw_stream_name(device_idx: int) -> str:
+    """Generate variable name for a raw stream handle on the given device."""
+    return f"raw_stream{device_idx}"
diff --git a/torch/_inductor/subgraph_lowering.py b/torch/_inductor/subgraph_lowering.py
index aa1b4d2db025d..b8938e8c778af 100644
--- a/torch/_inductor/subgraph_lowering.py
+++ b/torch/_inductor/subgraph_lowering.py
@@ -5,7 +5,7 @@
 from collections.abc import Callable, Generator
 from contextlib import contextmanager
 from dataclasses import dataclass
-from typing import Any, Optional, TypeVar, Union
+from typing import Any, cast, TypeVar
 from typing_extensions import ParamSpec
 
 import torch
@@ -14,7 +14,7 @@
 from . import ir
 from .exc import SubgraphLoweringException
 from .graph import GraphLowering
-from .ops_handler import SimpleCSEHandler
+from .ops_handler import OpsHandler, SimpleCSEHandler
 from .virtualized import ops, V, WrapperHandler
 
 
@@ -22,8 +22,8 @@
 _P = ParamSpec("_P")
 
 OpOverload = torch._ops.OpOverload
-LoweringDict = dict[Union[OpOverload, str], Callable[..., Any]]
-TargetType = Union[Callable[..., Any], str]
+LoweringDict = dict[OpOverload | str, Callable[..., Any]]
+TargetType = Callable[..., Any] | str
 
 
 class PointwiseSubgraphLowering(torch.fx.Interpreter):
@@ -32,12 +32,12 @@ class PointwiseSubgraphLowering(torch.fx.Interpreter):
     lowering object. Errors if buffers are created unexpectedly
     """
 
-    graph_outputs: Optional[list[ir.IRNode]]
+    graph_outputs: list[ir.IRNode] | None
     root_graph: GraphLowering
-    _current_op: Optional[TargetType]
+    _current_op: TargetType | None
     # For backwards of buffer_grads with scatters we allow mutations
-    allowed_mutations: Optional[OrderedSet[OpOverload]]
-    additional_lowerings: Optional[LoweringDict]
+    allowed_mutations: OrderedSet[OpOverload] | None
+    additional_lowerings: LoweringDict | None
     buffers: list[ir.Buffer]
     mutated_buffers: OrderedSet[str]
 
@@ -45,8 +45,8 @@ def __init__(
         self,
         gm: torch.fx.GraphModule,
         root_graph_lowering: GraphLowering,
-        allowed_mutations: Optional[OrderedSet[OpOverload]] = None,
-        additional_lowerings: Optional[LoweringDict] = None,
+        allowed_mutations: OrderedSet[OpOverload] | None = None,
+        additional_lowerings: LoweringDict | None = None,
     ) -> None:
         super().__init__(gm)
         self.graph_outputs = None
@@ -137,7 +137,7 @@ class InputDescriptor:
 class TracingOpsHandler(WrapperHandler):
     def __init__(self, tracer: torch.fx.Tracer, num_inputs: int) -> None:
         parent = tracer.create_proxy("placeholder", "ops", (), {})
-        super().__init__(parent)
+        super().__init__(cast(OpsHandler[Any], parent))
         self.tracer = tracer
 
         self.placeholders = [
@@ -158,9 +158,7 @@ def lower_pointwise_subgraph(
     subgraph: ir.Subgraph, inputs: list[InputDescriptor]
 ) -> Callable[_P, Any]:
     # Lower subgraph to ir.Pointwise nodes
-    def fake_inner_fn(
-        loop_idx: int, input_idx: int
-    ) -> Union[ir.Expr, ir.TensorBox, None]:
+    def fake_inner_fn(loop_idx: int, input_idx: int) -> ir.Expr | ir.TensorBox | None:
         return ops.placeholder(input_idx)
 
     graph_inputs = [
diff --git a/torch/_inductor/template_heuristics/aten.py b/torch/_inductor/template_heuristics/aten.py
index d4508b3185869..ad496c36a40e4 100644
--- a/torch/_inductor/template_heuristics/aten.py
+++ b/torch/_inductor/template_heuristics/aten.py
@@ -86,5 +86,9 @@ def _get_template_configs_impl(
         nodes = kernel_inputs.nodes()
         # for addmm, bias is the first input
         bias = nodes[0]
-        if bias.get_stride()[0] == 0 and inductor_config.triton.autotune_cublasLt:
-            yield dict()
+        assert (
+            len(bias.get_size()) == 2
+            and bias.get_stride()[0] == 0
+            and inductor_config.triton.autotune_cublasLt
+        )
+        yield from super()._get_template_configs_impl(kernel_inputs, op_name)
diff --git a/torch/_inductor/template_heuristics/nv_universal_gemm.py b/torch/_inductor/template_heuristics/nv_universal_gemm.py
index 494a63d438b26..adbdfae241df7 100644
--- a/torch/_inductor/template_heuristics/nv_universal_gemm.py
+++ b/torch/_inductor/template_heuristics/nv_universal_gemm.py
@@ -23,10 +23,10 @@
 autotuning_log = getArtifactLogger(__name__, "autotuning")
 
 # Type alias for kernel config key tuple.
-# Currently matches on (tile_m, tile_n, tile_k, cluster_m, cluster_n).
-# #TODO(nikhilap) When cutlass_api adds support for stages/split_k, extend this tuple and
-# update the _make_config_key_* helper functions below.
-ConfigKey = tuple[int, int, int, int, int]
+# Currently matches on (tile_m, tile_n, cluster_m, cluster_n).
+# tile_k excluded because nvMatmulHeuristics and cutlass_api use it to mean different things.
+# TODO(nikhilap): Extend config key for stages/split_k https://github.com/pytorch/pytorch/issues/177578
+ConfigKey = tuple[int, int, int, int]
 
 
 @dataclass
@@ -48,21 +48,20 @@ class HeuristicConfig:
 
 def _make_config_key_from_heuristic(cfg: HeuristicConfig) -> ConfigKey:
     """Build config key from HeuristicConfig returned by nvMatmulHeuristics."""
-    return (cfg.tile_m, cfg.tile_n, cfg.tile_k, cfg.cluster_m, cfg.cluster_n)
+    return (cfg.tile_m, cfg.tile_n, cfg.cluster_m, cfg.cluster_n)
 
 
 def _make_config_key_from_kernel_design(design) -> ConfigKey | None:
     """Build config key from cutlass_api kernel metadata.design."""
     if (
         hasattr(design, "tile_shape")
-        and len(design.tile_shape) >= 3
+        and len(design.tile_shape) >= 2
         and hasattr(design, "cluster_shape")
         and len(design.cluster_shape) >= 2
     ):
         return (
             design.tile_shape[0],
             design.tile_shape[1],
-            design.tile_shape[2],
             design.cluster_shape[0],
             design.cluster_shape[1],
         )
@@ -74,7 +73,6 @@ def _make_config_key_from_heuristics_kernel(kernel) -> ConfigKey:
     return (
         kernel.cta[0],
         kernel.cta[1],
-        kernel.cta[2],
         kernel.cluster[0],
         kernel.cluster[1],
     )
@@ -268,16 +266,20 @@ def _get_heuristic_configs(
             torch.bfloat16: "T",
             torch.float8_e4m3fn: "Q",
             torch.float8_e5m2: "R",
+            torch.float4_e2m1fn_x2: "F4",
         }
         a_char = dtype_to_cublas.get(dtype_a, "H")
         b_char = dtype_to_cublas.get(dtype_b or dtype_a, a_char)
         out_char = dtype_to_cublas.get(out_dtype or dtype_a, a_char)
         acc_char = dtype_to_cublas.get(accumulator_type, "S")
 
-        # 3-letter {input}{compute}{output}: used when A=B (standard GEMM).
-        # 5-letter {A}{B}{C}{compute}{D}: used when A≠B or for FP8/FP4
-        # (nvMatmulHeuristics discovery sets only have 5-letter entries for these).
-        if a_char != b_char or a_char in ("Q", "R", "O"):
+        # nvMatmulHeuristics precision string formats:
+        # - 3-letter {A}{B}{out}: used for standard GEMM and multi-char tokens (F4, BF)
+        # - 5-letter {A}{B}{C}{compute}{D}: used for single-char FP8 types (Q, R, O)
+        has_multichar = any(len(c) > 1 for c in (a_char, b_char, out_char))
+        if has_multichar:
+            precision = f"{a_char}{b_char}{out_char}"
+        elif a_char != b_char or a_char in ("Q", "R", "O"):
             precision = f"{a_char}{b_char}{out_char}{acc_char}{out_char}"
         else:
             precision = f"{a_char}{acc_char}{out_char}"
diff --git a/torch/_inductor/template_heuristics/registry.py b/torch/_inductor/template_heuristics/registry.py
index c03998608db70..b5606fce08a58 100644
--- a/torch/_inductor/template_heuristics/registry.py
+++ b/torch/_inductor/template_heuristics/registry.py
@@ -10,7 +10,7 @@
 
 import contextlib
 import logging
-from typing import Any, Optional, TYPE_CHECKING, Union
+from typing import Any, TYPE_CHECKING
 
 from .base import TemplateConfigHeuristics
 
@@ -21,7 +21,7 @@
 
 # Module-wide registry for template heuristics
 _TEMPLATE_HEURISTIC_REGISTRY: dict[
-    tuple[Union[str, None], ...], type[TemplateConfigHeuristics]
+    tuple[str | None, ...], type[TemplateConfigHeuristics]
 ] = {}
 
 # Manual cache for successful lookups only (fallback instances are not cached)
@@ -32,9 +32,9 @@
 
 def register_template_heuristic(
     template_name: str,
-    device_type: Union[str, None],
+    device_type: str | None,
     register: bool = True,
-    op_name: Optional[str] = None,
+    op_name: str | None = None,
 ) -> Any:
     """
     Decorator to register template heuristic classes.
@@ -60,7 +60,7 @@ def decorator(
         cls: type[TemplateConfigHeuristics],
     ) -> type[TemplateConfigHeuristics]:
         if register:
-            key: tuple[Union[str, None], ...] = (template_name, device_type, op_name)
+            key: tuple[str | None, ...] = (template_name, device_type, op_name)
             _TEMPLATE_HEURISTIC_REGISTRY[key] = cls
             log.info(
                 f"Registered template heuristic: {cls.__name__} for '{template_name=}', '{device_type=}', '{op_name=}'"  # noqa: G004
diff --git a/torch/_inductor/template_heuristics/tlx.py b/torch/_inductor/template_heuristics/tlx.py
index 83381687cd5de..e75457aceb1cf 100644
--- a/torch/_inductor/template_heuristics/tlx.py
+++ b/torch/_inductor/template_heuristics/tlx.py
@@ -3,5 +3,3 @@
 
 if config.is_fbcode():
     import torch._inductor.fb.tlx_templates.registry  # noqa: F401  # type: ignore[import-not-used]
-
-# TODO. Move the registry to this file once the TLX template is more complete.
diff --git a/torch/_inductor/template_heuristics/triton.py b/torch/_inductor/template_heuristics/triton.py
index 96457817ff5e3..ae0998c7b5759 100644
--- a/torch/_inductor/template_heuristics/triton.py
+++ b/torch/_inductor/template_heuristics/triton.py
@@ -6,13 +6,14 @@
 import os
 from functools import partial
 from threading import Lock
-from typing import Any, Optional, TYPE_CHECKING
+from typing import Any, TYPE_CHECKING
 
 import sympy
 
 import torch
 from torch._inductor.template_heuristics.triton_addmm import AddMMConfigMixin
 from torch.utils._ordered_set import OrderedSet
+from torch.utils._sympy.functions import Mod
 from torch.utils._triton import has_triton_stable_tma_api
 
 from .. import config, config as inductor_config
@@ -22,14 +23,17 @@
     get_scaling_options,
     get_tile_size,
     mm_template,
+    persistent_mm_template,
     persistent_tma_mm_template,
     scaled_mm_device_tma_epilogue_scaling_template,
     scaled_mm_device_tma_main_loop_scaling_template,
 )
 from ..kernel.mm_plus_mm import mm_plus_mm_template
 from ..kernel_inputs import KernelInputs, MMKernelInputs
+from ..runtime.hints import DeviceProperties
 from ..utils import (
     get_backend_num_stages,
+    get_default_kpack,
     get_num_sms,
     get_tma_workspace_arg,
     TMA_DESCRIPTOR_SIZE,
@@ -61,7 +65,7 @@ class BaseConfig:
     block_k: int
     num_stages: int
     num_warps: int
-    hint_override: Optional[int] = dataclasses.field(kw_only=True, default=None)
+    hint_override: int | None = dataclasses.field(kw_only=True, default=None)
 
 
 @dataclasses.dataclass
@@ -76,6 +80,21 @@ class GemmConfig(BaseConfig):
 ConvConfig = BaseConfig
 
 
+@dataclasses.dataclass
+class DepthwiseConvConfig:
+    """
+    Configuration for depthwise conv1d Triton template.
+    Uses BLOCK_N x BLOCK_L x BLOCK_C tiling (channels-last NLC layout).
+    Matches the hand-written NLC kernel from depthwise_conv1d_benchmark.py.
+    """
+
+    block_n: int
+    block_l: int
+    block_c: int
+    num_stages: int
+    num_warps: int
+
+
 @dataclasses.dataclass
 class BlackwellGPUGemmConfig(GemmConfig):
     """
@@ -161,7 +180,7 @@ class ROCmGemmConfig(GemmConfig):
 
     matrix_instr_nonkdim: int = 16
     waves_per_eu: int = 0
-    kpack: int = 2
+    kpack: int = 1
 
 
 @dataclasses.dataclass
@@ -172,7 +191,7 @@ class ROCmConvConfig(ConvConfig):
 
     matrix_instr_nonkdim: int = 16
     waves_per_eu: int = 0
-    kpack: int = 2
+    kpack: int = 1
 
 
 @dataclasses.dataclass
@@ -183,7 +202,7 @@ class ROCmFlexConfig(FlexConfig):
 
     matrix_instr_nonkdim: int = 0
     waves_per_eu: int = 0
-    kpack: int = 2
+    kpack: int = 1
 
 
 @dataclasses.dataclass
@@ -194,7 +213,7 @@ class ROCmFlexBwDConfig(FlexBwDConfig):
 
     matrix_instr_nonkdim: int = 0
     waves_per_eu: int = 0
-    kpack: int = 2
+    kpack: int = 1
 
 
 @dataclasses.dataclass
@@ -205,7 +224,7 @@ class ROCmFlexDecodeConfig(FlexDecodeConfig):
 
     matrix_instr_nonkdim: int = 0
     waves_per_eu: int = 0
-    kpack: int = 2
+    kpack: int = 1
 
 
 class BaseHeuristicSingleton(type):
@@ -654,6 +673,63 @@ def __init__(self) -> None:
             ConvConfig(128, 256, 128, 2, 8),
         ]
 
+        # Depthwise conv1d configs: BLOCK_N x BLOCK_L x BLOCK_C tiling
+        # Derived from autotuning results on H100 for depthwise conv1d
+        # channels-last (NLC) layout with shape x=[3072, 128, 202]
+        # Matches _nlc_autotune_configs from depthwise_conv1d_benchmark.py
+        self.depthwise_conv_configs: list[DepthwiseConvConfig] = [
+            # BLOCK_C=32, BLOCK_L=32
+            DepthwiseConvConfig(
+                block_n=16, block_l=32, block_c=32, num_stages=4, num_warps=8
+            ),
+            DepthwiseConvConfig(
+                block_n=16, block_l=32, block_c=32, num_stages=4, num_warps=4
+            ),
+            DepthwiseConvConfig(
+                block_n=32, block_l=32, block_c=32, num_stages=5, num_warps=8
+            ),
+            DepthwiseConvConfig(
+                block_n=32, block_l=32, block_c=32, num_stages=4, num_warps=4
+            ),
+            # BLOCK_C=32, BLOCK_L=64
+            DepthwiseConvConfig(
+                block_n=16, block_l=64, block_c=32, num_stages=4, num_warps=8
+            ),
+            DepthwiseConvConfig(
+                block_n=16, block_l=64, block_c=32, num_stages=4, num_warps=4
+            ),
+            DepthwiseConvConfig(
+                block_n=32, block_l=64, block_c=32, num_stages=3, num_warps=8
+            ),
+            # BLOCK_C=32, BLOCK_L=256
+            DepthwiseConvConfig(
+                block_n=16, block_l=256, block_c=32, num_stages=5, num_warps=8
+            ),
+            DepthwiseConvConfig(
+                block_n=16, block_l=256, block_c=32, num_stages=4, num_warps=4
+            ),
+            DepthwiseConvConfig(
+                block_n=32, block_l=256, block_c=32, num_stages=3, num_warps=8
+            ),
+            # BLOCK_C=64
+            DepthwiseConvConfig(
+                block_n=16, block_l=32, block_c=64, num_stages=4, num_warps=8
+            ),
+            DepthwiseConvConfig(
+                block_n=16, block_l=32, block_c=64, num_stages=4, num_warps=4
+            ),
+            DepthwiseConvConfig(
+                block_n=16, block_l=64, block_c=64, num_stages=3, num_warps=8
+            ),
+            # BLOCK_C=128
+            DepthwiseConvConfig(
+                block_n=16, block_l=32, block_c=128, num_stages=3, num_warps=8
+            ),
+            DepthwiseConvConfig(
+                block_n=16, block_l=32, block_c=128, num_stages=3, num_warps=4
+            ),
+        ]
+
         self.flex_attn_fwd_autotune_configs: list[FlexConfig] = [
             FlexConfig(128, 64, 3, 4),
             FlexConfig(128, 128, 3, 4),
@@ -707,6 +783,12 @@ def __init__(self) -> None:
             for num_warps in [2, 4, 8]
         ]
 
+    def _get_extra_config_key_and_kwargs(
+        self, conf: BaseConfig
+    ) -> tuple[tuple[int | None, ...], dict[str, Any]]:
+        """Hook for subclasses to extend config dedup key and kwargs."""
+        return (), {}
+
     def _finalize_mm_configs(
         self,
         configs: list[BaseConfig],
@@ -714,7 +796,7 @@ def _finalize_mm_configs(
         """
         Finalizes configs after scaling, applying additional constraints.
         """
-        used: OrderedSet[tuple[Optional[int], ...]] = OrderedSet()
+        used: OrderedSet[tuple[int | None, ...]] = OrderedSet()
 
         max_mm_configs = config.test_configs.max_mm_configs
 
@@ -723,7 +805,7 @@ def _finalize_mm_configs(
             num_warps = min(conf.num_warps, conf.block_m * conf.block_n // 256)
 
             # Construct key for finding duplicate configs
-            key: tuple[Optional[int], ...] = (
+            key: tuple[int | None, ...] = (
                 conf.block_m,
                 conf.block_n,
                 conf.block_k,
@@ -741,16 +823,8 @@ def _finalize_mm_configs(
             if isinstance(conf, BlackwellGPUGemmConfig):
                 key += (conf.epilogue_subtile, conf.warp_specialize, conf.flatten)
 
-            # Add TlxGemmConfig specific fields to key if present
-            if config.is_fbcode() and config.triton.enable_tlx_templates:
-                from torch._inductor.fb.tlx_templates.registry import (
-                    get_tlx_config_key_and_kwargs,
-                )
-
-                tlx_key_fields, tlx_kwargs = get_tlx_config_key_and_kwargs(conf)
-                key += tlx_key_fields
-            else:
-                tlx_kwargs = {}
+            extra_key, extra_kwargs = self._get_extra_config_key_and_kwargs(conf)
+            key += extra_key
 
             if key not in used and (
                 max_mm_configs is None or len(used) < max_mm_configs
@@ -771,8 +845,7 @@ def _finalize_mm_configs(
                     kwargs["WARP_SPECIALIZE"] = conf.warp_specialize
                     kwargs["FLATTEN"] = conf.flatten
 
-                # Add TlxGemmConfig specific fields if present
-                kwargs.update(tlx_kwargs)
+                kwargs.update(extra_kwargs)
 
                 yield self.triton_config(conf.num_stages, num_warps, **kwargs)
 
@@ -785,7 +858,7 @@ def _scale_mm_configs(
         scale: float,
         has_int8_tensor: bool,
         exclude: Callable[[sympy.Integer, sympy.Integer, sympy.Integer], bool],
-        hint_override: Optional[int] = None,
+        hint_override: int | None = None,
     ) -> list[BaseConfig]:
         """
         Scales and filters matrix multiplication configs based on input size.
@@ -892,7 +965,7 @@ def _get_exceeding_shared_memory_checker(
         self,
         has_sm_layout_conversion: bool,
         layout_conversion_byte_size: int,
-    ) -> Optional[Callable[[BaseConfig, int], bool]]:
+    ) -> Callable[[BaseConfig, int], bool] | None:
         """
         Returns a function that checks whether a given configuration exceeds the available shared memory for the device.
         based on the config's theoretical maximum shared memory used.
@@ -947,11 +1020,22 @@ def _prune_reg_spill_configs(
         self,
         configs: list[BaseConfig],
     ) -> list[BaseConfig]:
+        try:
+            device = torch.cuda.current_device()
+            props = torch.cuda.get_device_properties(device)
+            warp_size = props.warp_size
+        except (RuntimeError, AttributeError, AssertionError):
+            return configs
+
         pruned_configs = []
+        # NVIDIA-oriented approximation for the per-thread register limit.
+        # TODO: generalize this threshold for ROCm devices.
+        NUM_REG = 255
         for gemm_config in configs:
-            NUM_REG = 255
             acc_regs = math.ceil(
-                gemm_config.block_m * gemm_config.block_n / (gemm_config.num_warps * 32)
+                gemm_config.block_m
+                * gemm_config.block_n
+                / (gemm_config.num_warps * warp_size)
             )
             # Lower bound for register spillage, if exceeds the kernel will certainly spill
             if acc_regs > NUM_REG:
@@ -1019,6 +1103,21 @@ def get_conv_configs(self) -> partial[Generator[TritonConfig, None, None]]:
             self.preprocess_mm_configs, configs=self.conv_configs, op_name="conv"
         )
 
+    def get_depthwise_conv_configs(self) -> list[TritonConfig]:
+        """Return TritonConfig list for depthwise conv1d autotuning."""
+        return [
+            TritonConfig(
+                {
+                    "BLOCK_N": cfg.block_n,
+                    "BLOCK_L": cfg.block_l,
+                    "BLOCK_C": cfg.block_c,
+                },
+                num_stages=cfg.num_stages,
+                num_warps=cfg.num_warps,
+            )
+            for cfg in self.depthwise_conv_configs
+        ]
+
     # Flex attn helpers
     def get_flex_attn_fwd_configs(self, head_dim: int, dtype: Any) -> list[FlexConfig]:
         flex_attn_fwd_configs: list[FlexConfig] = []
@@ -1195,6 +1294,7 @@ def __init__(self) -> None:
         self.h100_default_flex_config = {
             (torch.float32, 64): FlexConfig(128, 32, 3, 4),
             (torch.float32, 128): FlexConfig(32, 64, 3, 4),
+            (torch.float32, 192): FlexConfig(32, 64, 1, 8),
             (torch.float32, 256): FlexConfig(32, 32, 3, 4),
             (torch.bfloat16, 64): FlexConfig(128, 128, 3, 4),
             (torch.bfloat16, 128): FlexConfig(128, 64, 3, 8),
@@ -1460,23 +1560,26 @@ def __init__(self) -> None:
             for group_m in [4, 8, 16]
             for matrix_instr_nonkdim in [0, 16]
             for waves_per_eu in [0, 2]
-            for kpack in [2]
+            for kpack in [1, 2]
         ]
 
+        # Architecture-aware default kpack for flex configs
+        default_kpack = get_default_kpack()
+
         self.default_flex_config = {
-            (torch.float32, 64): ROCmFlexConfig(128, 32, 1, 4),
-            (torch.float32, 128): ROCmFlexConfig(128, 32, 1, 4),
-            (torch.float32, 256): ROCmFlexConfig(64, 16, 1, 4),
-            (torch.bfloat16, 64): ROCmFlexConfig(128, 64, 1, 8),
-            (torch.bfloat16, 128): ROCmFlexConfig(128, 64, 1, 8),
-            (torch.bfloat16, 256): ROCmFlexConfig(32, 64, 1, 8),
-            (torch.float16, 64): ROCmFlexConfig(128, 64, 1, 8),
-            (torch.float16, 128): ROCmFlexConfig(128, 64, 1, 8),
-            (torch.float16, 256): ROCmFlexConfig(32, 64, 1, 4),
+            (torch.float32, 64): ROCmFlexConfig(128, 32, 1, 4, kpack=default_kpack),
+            (torch.float32, 128): ROCmFlexConfig(128, 32, 1, 4, kpack=default_kpack),
+            (torch.float32, 256): ROCmFlexConfig(64, 16, 1, 4, kpack=default_kpack),
+            (torch.bfloat16, 64): ROCmFlexConfig(128, 64, 2, 4, kpack=default_kpack),
+            (torch.bfloat16, 128): ROCmFlexConfig(128, 64, 2, 4, kpack=default_kpack),
+            (torch.bfloat16, 256): ROCmFlexConfig(32, 64, 2, 4, kpack=default_kpack),
+            (torch.float16, 64): ROCmFlexConfig(128, 64, 2, 8, kpack=default_kpack),
+            (torch.float16, 128): ROCmFlexConfig(128, 64, 2, 8, kpack=default_kpack),
+            (torch.float16, 256): ROCmFlexConfig(32, 64, 2, 4, kpack=default_kpack),
         }
 
         self.flex_attn_fwd_autotune_configs: list[FlexConfig] = [
-            ROCmFlexConfig(BLOCK1, BLOCK2, 1, w)
+            ROCmFlexConfig(BLOCK1, BLOCK2, 1, w, kpack=default_kpack)
             for BLOCK1 in [16, 64, 128]
             for BLOCK2 in [16, 32, 64, 128]
             for w in [4, 8]
@@ -1484,7 +1587,9 @@ def __init__(self) -> None:
 
         self.flex_attn_bwd_autotune_configs: list[FlexBwDConfig] = [
             # See Note: flex bwd configs
-            ROCmFlexBwDConfig(BLOCK1, BLOCK2, BLOCK2, BLOCK1, 1, w, mfma)
+            ROCmFlexBwDConfig(
+                BLOCK1, BLOCK2, BLOCK2, BLOCK1, 1, w, mfma, kpack=default_kpack
+            )
             for BLOCK1 in [16, 32, 64]
             for BLOCK2 in [32, 64, 128]
             for w in ([4, 8] if BLOCK1 >= 128 or BLOCK2 >= 128 else [4])
@@ -1493,22 +1598,23 @@ def __init__(self) -> None:
         ]
 
         self.flex_decode_autotune_configs: list[FlexDecodeConfig] = [
-            ROCmFlexDecodeConfig(32, 1, 4),
-            ROCmFlexDecodeConfig(64, 1, 4),
-            ROCmFlexDecodeConfig(128, 1, 4),
-            ROCmFlexDecodeConfig(32, 1, 8),
-            ROCmFlexDecodeConfig(64, 1, 8),
-            ROCmFlexDecodeConfig(128, 1, 8),
+            ROCmFlexDecodeConfig(32, 1, 4, kpack=default_kpack),
+            ROCmFlexDecodeConfig(64, 1, 4, kpack=default_kpack),
+            ROCmFlexDecodeConfig(128, 1, 4, kpack=default_kpack),
+            ROCmFlexDecodeConfig(32, 1, 8, kpack=default_kpack),
+            ROCmFlexDecodeConfig(64, 1, 8, kpack=default_kpack),
+            ROCmFlexDecodeConfig(128, 1, 8, kpack=default_kpack),
         ]
 
         self.exhaustive_flex_attn_fwd_configs: list[FlexConfig] = [
-            ROCmFlexConfig(BLOCK_M, BLOCK_N, num_stages, num_warps, mfma, wpeu)
+            ROCmFlexConfig(BLOCK_M, BLOCK_N, num_stages, num_warps, mfma, wpeu, kpack)
             for BLOCK_M in [16, 32, 64, 128]
             for BLOCK_N in [32, 64, 128]
             for num_stages in [1, 2]
             for num_warps in [2, 4, 8]
             for mfma in [0, 16]
             for wpeu in [0, int(8 // num_warps)]
+            for kpack in [1, 2]
         ]
 
         self.exhaustive_flex_attn_bwd_configs: list[FlexBwDConfig] = [
@@ -1522,6 +1628,7 @@ def __init__(self) -> None:
                 num_warps,
                 mfma,
                 wpeu,
+                kpack,
             )
             for BLOCK_M1 in [16, 32, 64, 128]
             for BLOCK_N1 in [16, 32, 64, 128]
@@ -1531,17 +1638,21 @@ def __init__(self) -> None:
             for num_warps in [2, 4, 8]
             for mfma in [0, 16]
             for wpeu in [0, int(8 // num_warps)]
+            for kpack in [1, 2]
             if BLOCK_N1 % BLOCK_M1 == 0
             and BLOCK_M2 % BLOCK_N2 == 0  # kernel static assertions
         ]
 
         self.exhaustive_flex_decode_configs: list[FlexDecodeConfig] = [
-            ROCmFlexDecodeConfig(block_n, num_stages, num_warps, mfma, wpeu, kpack=2)
+            ROCmFlexDecodeConfig(
+                block_n, num_stages, num_warps, mfma, wpeu, kpack=kpack
+            )
             for block_n in [16, 32, 64, 128]
             for num_stages in [1, 2]
             for num_warps in [2, 4, 8]
             for mfma in [0, 16]
             for wpeu in [0, int(8 // num_warps)]
+            for kpack in [1, 2]
         ]
 
     def _prune_exhaustive_configs(
@@ -1554,8 +1665,11 @@ def _prune_exhaustive_configs(
             c
             for c in configs
             if not (
-                getattr(c, "matrix_instr_nonkdim", 0) == 2
-                and getattr(c, "kpack", 0) == 2
+                (
+                    getattr(c, "matrix_instr_nonkdim", 0) == 2
+                    and getattr(c, "kpack", 0) == 2
+                )
+                or (c.block_k <= 16 and getattr(c, "kpack", 0) == 2)
             )
         ]
         return pruned_configs
@@ -1584,9 +1698,11 @@ def _finalize_mm_configs(
             conf.num_warps = min(conf.num_warps, conf.block_m * conf.block_n // 256)
 
             # Defaults for AMD triton backend kern args if not set
-            matrix_instr_nonkdim = getattr(conf, "matrix_instr_nonkdim", 16)
-            waves_per_eu = getattr(conf, "waves_per_eu", 0)
-            kpack = getattr(conf, "kpack", 2)
+            matrix_instr_nonkdim: int = getattr(conf, "matrix_instr_nonkdim", 16)
+            waves_per_eu: int = getattr(conf, "waves_per_eu", 0)
+            # Use explicit kpack if set, otherwise determine optimal value based on
+            # architecture and BLOCK_K
+            kpack: int = getattr(conf, "kpack", get_default_kpack(conf.block_k))
 
             if matrix_instr_nonkdim != 0 and (
                 conf.block_m % matrix_instr_nonkdim != 0
@@ -1644,19 +1760,20 @@ def get_flex_attn_fwd_configs(self, head_dim: int, dtype: Any) -> list[FlexConfi
                 return self.exhaustive_flex_attn_fwd_configs
             flex_attn_fwd_configs += self.flex_attn_fwd_autotune_configs
 
+        default_kpack = get_default_kpack()
         if head_dim <= 256:
             if dtype == torch.float32:
-                default_config = ROCmFlexConfig(64, 64, 1, 4)
+                default_config = ROCmFlexConfig(64, 64, 1, 4, kpack=default_kpack)
             else:
-                default_config = ROCmFlexConfig(128, 64, 1, 8)
+                default_config = ROCmFlexConfig(128, 64, 2, 8, kpack=default_kpack)
             default_config = self.default_flex_config.get(
                 (dtype, head_dim), default_config
             )
         else:
             if dtype == torch.float32:
-                default_config = ROCmFlexConfig(32, 16, 1, 4)
+                default_config = ROCmFlexConfig(32, 16, 1, 4, kpack=default_kpack)
             else:
-                default_config = ROCmFlexConfig(64, 32, 1, 4)
+                default_config = ROCmFlexConfig(64, 32, 2, 4, kpack=default_kpack)
 
         if default_config not in flex_attn_fwd_configs:
             flex_attn_fwd_configs.append(default_config)
@@ -1673,17 +1790,28 @@ def get_flex_attn_bwd_configs(
                 return self.exhaustive_flex_attn_bwd_configs
             flex_attn_bwd_configs += self.flex_attn_bwd_autotune_configs
 
+        default_kpack = get_default_kpack()
         if dtype == torch.float32:
-            default_config = ROCmFlexBwDConfig(16, 16, 16, 16, 1, 4)
+            default_config = ROCmFlexBwDConfig(
+                16, 16, 16, 16, 1, 4, kpack=default_kpack
+            )
         elif head_dim <= 256:
             if head_dim == 64:
-                default_config = ROCmFlexBwDConfig(64, 64, 64, 64, 1, 4)
+                default_config = ROCmFlexBwDConfig(
+                    64, 64, 64, 64, 1, 4, kpack=default_kpack
+                )
             elif head_dim == 128:
-                default_config = ROCmFlexBwDConfig(64, 128, 128, 64, 1, 8)
+                default_config = ROCmFlexBwDConfig(
+                    64, 128, 128, 64, 1, 4, kpack=default_kpack
+                )
             else:
-                default_config = ROCmFlexBwDConfig(64, 64, 64, 64, 1, 4)
+                default_config = ROCmFlexBwDConfig(
+                    64, 64, 64, 64, 1, 4, kpack=default_kpack
+                )
         else:
-            default_config = ROCmFlexBwDConfig(16, 16, 16, 16, 1, 4)
+            default_config = ROCmFlexBwDConfig(
+                16, 16, 16, 16, 1, 4, kpack=default_kpack
+            )
 
         if default_config not in flex_attn_bwd_configs:
             flex_attn_bwd_configs.append(default_config)
@@ -1700,7 +1828,8 @@ def get_flex_decode_configs(
                 return self.exhaustive_flex_decode_configs
             flex_decode_configs += self.flex_decode_autotune_configs
 
-        default_config = ROCmFlexDecodeConfig(64, 1, 4)
+        default_kpack = get_default_kpack()
+        default_config = ROCmFlexDecodeConfig(64, 1, 4, kpack=default_kpack)
 
         if default_config not in flex_decode_configs:
             flex_decode_configs.append(default_config)
@@ -1715,6 +1844,10 @@ class XPUConfigHeuristic(BaseConfigHeuristic):
 
     def __init__(self) -> None:
         super().__init__()
+        self.mm_configs = self.mm_configs + [
+            GemmConfig(32, 64, 128, 2, 2),
+            GemmConfig(64, 64, 32, 2, 8),
+        ]
         self.xpu_default_flex_config = {
             (torch.float32, 64): FlexConfig(128, 32, 1, 16),
             (torch.float32, 128): FlexConfig(128, 32, 1, 16),
@@ -1953,7 +2086,7 @@ def _convert_config_to_template_kwargs(
         Moved from mm_common.mm_options.
         """
         # Calculate EVEN_K symbolic. (It isn't worth guarding on this)
-        even_k_symbolic = (k % triton_config.kwargs["BLOCK_K"]) == 0
+        even_k_symbolic = sympy.Eq(Mod(k, triton_config.kwargs["BLOCK_K"]), 0)
         even_k_symbolic = V.graph.sizevars.statically_known_true(even_k_symbolic)
 
         # Build options dict
@@ -1962,6 +2095,7 @@ def _convert_config_to_template_kwargs(
             EVEN_K=even_k_symbolic,
             USE_FAST_ACCUM=False,  # Option for _scaled_mm
             ACC_TYPE=self._get_acc_type(out_dtype),
+            OUT_DTYPE=self._get_out_dtype(out_dtype),
             num_stages=triton_config.num_stages,
             num_warps=triton_config.num_warps,
             **triton_config.kwargs,
@@ -1974,6 +2108,11 @@ def _convert_config_to_template_kwargs(
 
         return options_dict
 
+    @staticmethod
+    def _dtype_to_triton(dtype: torch.dtype) -> str:
+        """Convert a torch dtype to a triton type string."""
+        return f"tl.{dtype}".replace("torch.", "")
+
     def _get_acc_type(self, dtype: torch.dtype) -> str:
         """
         Get accumulator type for the given dtype.
@@ -1981,7 +2120,11 @@ def _get_acc_type(self, dtype: torch.dtype) -> str:
         """
         if dtype in (torch.float16, torch.bfloat16):
             return "tl.float32"
-        return f"tl.{dtype}".replace("torch.", "")
+        return self._dtype_to_triton(dtype)
+
+    def _get_out_dtype(self, dtype: torch.dtype) -> str:
+        """Get output dtype as a triton type string."""
+        return self._dtype_to_triton(dtype)
 
 
 # INT8 specific mixin to filter correctly
@@ -2204,7 +2347,10 @@ def adjust_kernel_inputs(
         if bias:
             nodes.append(bias)
         return MMKernelInputs(
-            nodes, mat1_idx=kernel_inputs._mat1_idx, mat2_idx=kernel_inputs._mat2_idx
+            nodes,
+            mat1_idx=kernel_inputs._mat1_idx,
+            mat2_idx=kernel_inputs._mat2_idx,
+            out_dtype=kernel_inputs._out_dtype,
         )
 
     def _get_template_configs_impl(
@@ -2446,6 +2592,33 @@ def __init__(self) -> None:
         self.mm_configs = self.persistent_mm_configs
 
 
+@register_template_heuristic(
+    persistent_mm_template.uid,
+    "cuda",
+    register=torch.version.hip is not None,
+)
+class PersistentMMTemplateConfigHeuristic(
+    MMTemplateConfigMixin,
+    ROCmConfigHeuristic,  # type: ignore[misc]
+):
+    """Persistent MM template heuristic (no TMA, standard pointer loads)"""
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.mm_configs = self.persistent_mm_configs
+
+    def _get_template_configs_impl(
+        self,
+        kernel_inputs: KernelInputs,
+        op_name: str,
+        **kwargs,
+    ) -> Generator[dict[str, Any], None, None]:
+        for template_kwargs in super()._get_template_configs_impl(
+            kernel_inputs, op_name, **kwargs
+        ):
+            yield {**template_kwargs, "NUM_SMS": get_num_sms()}
+
+
 @register_template_heuristic(
     blackwell_ws_persistent_device_tma_mm_template.uid,
     "cuda",
@@ -2597,7 +2770,8 @@ def _get_template_configs_impl(
 
 
 @register_template_heuristic(
-    blackwell_ws_persistent_device_tma_mm_template.uid,  # regular Blackwell MM template + scaling epilogue from ScaledMMConfigMixin
+    # regular Blackwell MM template + scaling epilogue from ScaledMMConfigMixin
+    blackwell_ws_persistent_device_tma_mm_template.uid,
     "cuda",
     register=torch.version.hip is None,
     op_name="scaled_mm",
@@ -2683,6 +2857,18 @@ class ROCmAddMMTemplateConfigHeuristic(AddMMConfigMixin, ROCmMMTemplateConfigHeu
     """Addmm specific mixin for ROCm"""
 
 
+@register_template_heuristic(
+    persistent_mm_template.uid,
+    "cuda",
+    register=torch.version.hip is not None,
+    op_name="addmm",
+)
+class ROCmAddMMPersistentTemplateConfigHeuristic(
+    AddMMConfigMixin, PersistentMMTemplateConfigHeuristic
+):
+    """Addmm specific mixin for persistent MM on ROCm"""
+
+
 # TODO(coconutruben): deprecate once autoheuristic is deprecated
 @register_template_heuristic("mm-ah", "cuda", register=torch.version.hip is not None)
 class ROCmMMAHTemplateConfigHeuristic(MMTemplateConfigMixin, ROCmConfigHeuristic):
@@ -2708,11 +2894,10 @@ def __init__(self) -> None:
         super().__init__()
         # Override mm_configs to use scaled_mm_configs
         self.mm_configs = self.scaled_mm_configs
-        # NOTE: overriding exhaustive configs here to be the same as mm_configs
-        # as we haven't validated exhaustive support here yet
-        # TODO(coconutruben): remove this once we have validated exhaustive support
-        # for scaled_mm
-        self.exhaustive_configs = self.scaled_mm_configs
+
+    def _filter_configs(self, configs: list[BaseConfig]) -> list[BaseConfig]:
+        configs = [c for c in configs if c.block_k >= 32]
+        return super()._filter_configs(configs)
 
 
 @register_template_heuristic(
diff --git a/torch/_inductor/test_case.py b/torch/_inductor/test_case.py
index efdef48884cef..9b65101b5d5c0 100644
--- a/torch/_inductor/test_case.py
+++ b/torch/_inductor/test_case.py
@@ -1,6 +1,5 @@
 import contextlib
 import os
-from typing import Union
 
 from torch._dynamo.test_case import (
     run_tests as dynamo_run_tests,
@@ -11,7 +10,7 @@
 from torch._inductor.utils import fresh_cache
 
 
-def run_tests(needs: Union[str, tuple[str, ...]] = ()) -> None:
+def run_tests(needs: str | tuple[str, ...] = ()) -> None:
     dynamo_run_tests(needs)
 
 
diff --git a/torch/_inductor/tiling_utils.py b/torch/_inductor/tiling_utils.py
index e59010a2a963b..80acf14b6eeab 100644
--- a/torch/_inductor/tiling_utils.py
+++ b/torch/_inductor/tiling_utils.py
@@ -2,7 +2,7 @@
 import itertools
 from collections import Counter, defaultdict
 from collections.abc import Callable
-from typing import Literal, Optional, overload, TYPE_CHECKING, TypeVar, Union
+from typing import Literal, overload, TYPE_CHECKING, TypeVar, Union
 
 import sympy
 
@@ -33,7 +33,7 @@
     from torch._inductor.scheduler import FusedSchedulerNode, SchedulerNode
 
 
-def solve_for_zero(expr: sympy.Expr) -> Optional[sympy.Expr]:
+def solve_for_zero(expr: sympy.Expr) -> sympy.Expr | None:
     """
     Given an expr with a single free symbol, solve for a constant relation that would make
     this expression 0.
@@ -54,7 +54,7 @@ def solve_for_zero(expr: sympy.Expr) -> Optional[sympy.Expr]:
     return out[1]
 
 
-def solve_for_tiling(expr: sympy.Expr) -> Optional[sympy.Expr]:
+def solve_for_tiling(expr: sympy.Expr) -> sympy.Expr | None:
     """
     Giving an expr with a single free symbol, try to find a tiling that would
     make the expression coalesced with respect to that symbol.
@@ -71,7 +71,7 @@ def solve_for_tiling(expr: sympy.Expr) -> Optional[sympy.Expr]:
 
     free_symbol = next(iter(expr.free_symbols))
 
-    def _solve_simple_expr(expr: sympy.Expr) -> Optional[sympy.Expr]:
+    def _solve_simple_expr(expr: sympy.Expr) -> sympy.Expr | None:
         assert not expr.has(ModularIndexing) and not expr.has(FloorDiv)
         if len(expr.free_symbols) != 1:
             return None
@@ -119,7 +119,7 @@ def _solve_simple_expr(expr: sympy.Expr) -> Optional[sympy.Expr]:
     def indexing_div_rep(
         x: sympy.Expr,
         y: sympy.Expr,
-        z: Optional[sympy.Expr] = None,
+        z: sympy.Expr | None = None,
     ) -> sympy.Expr:
         return x / y
 
@@ -146,7 +146,7 @@ def indexing_div_rep(
 
 def find_broadcast_var(
     index: sympy.Expr, var_ranges: dict[sympy.Expr, int]
-) -> Optional[sympy.Expr]:
+) -> sympy.Expr | None:
     """
     Try to find the variable that this index is broadcast over.
     A broadcast pattern is one where consecutive values of a variable
@@ -181,7 +181,7 @@ def find_broadcast_var(
 
 def find_coalesced_var(
     index: sympy.Expr, var_ranges: dict[sympy.Expr, int]
-) -> Optional[sympy.Expr]:
+) -> sympy.Expr | None:
     """
     Try to find the symbol which coalesces this index
     """
@@ -243,7 +243,7 @@ def get_pw_red_splits(
     pointwise_numel: sympy.Expr,
     red_numel: sympy.Expr,
     none_if_not_divisible: Literal[True],
-) -> Optional[tuple[VarsAndRanges, VarsAndRanges]]: ...
+) -> tuple[VarsAndRanges, VarsAndRanges] | None: ...
 
 
 @overload
@@ -260,8 +260,13 @@ def get_pw_red_splits(
     pointwise_numel: sympy.Expr,
     red_numel: sympy.Expr,
     none_if_not_divisible: bool = False,
-) -> Optional[tuple[VarsAndRanges, VarsAndRanges]]:
-    if n.is_reduction() or sympy_product(n._body.sizes[0]) == pointwise_numel:
+) -> tuple[VarsAndRanges, VarsAndRanges] | None:
+    # nb: use statically_known_equals here to mimic scheduler.
+    # TODO : store type of split/broadcast on fused node itself,
+    # instead of re-deriving it.
+    if n.is_reduction() or V.graph.sizevars.statically_known_equals(
+        sympy_product(n._body.sizes[0]), pointwise_numel
+    ):
         # pyrefly: ignore [bad-return]
         return (
             (n._body.iter_vars, n._body.sizes[0]),
@@ -393,7 +398,7 @@ def add_combined_split_options(
         # if for whatever reason we couldn't split above, return default split
         return ((self.pointwise_numel,), (self.red_numel,))
 
-    def try_split(self, pw: Split, red: Split) -> Optional[tuple[Split, Split]]:
+    def try_split(self, pw: Split, red: Split) -> tuple[Split, Split] | None:
         """
         See if this split is compatible, and potentially returning a longer split
         than the input.
@@ -498,7 +503,7 @@ def apply_var_mapping(
 
 def extract_normalized_read_writes(
     node: Union["FusedSchedulerNode", "SchedulerNode"],
-) -> Optional[FusedNormalizedReadsWrites]:
+) -> FusedNormalizedReadsWrites | None:
     """Extracts index variables, reduce variables, read/write expressions, and variable ranges from a fused node."""
     reads: dict[sympy.Expr, OrderedSet[str]] = defaultdict(OrderedSet)
     writes: dict[sympy.Expr, OrderedSet[str]] = defaultdict(OrderedSet)
@@ -636,14 +641,14 @@ def get_score(
     return V.graph.sizevars.optimization_hint(sympy_product(var_sizes))
 
 
-def try_get_buf_size(buf_name: str) -> Optional[int]:
+def try_get_buf_size(buf_name: str) -> int | None:
     buf = V.graph.try_get_buffer(buf_name)
     if not buf:
         return None
     return V.graph.sizevars.optimization_hint(sympy_product(buf.get_size()))
 
 
-def get_hint(v: Union[sympy.Expr, int]) -> int:
+def get_hint(v: sympy.Expr | int) -> int:
     if isinstance(v, int):
         return v
     else:
@@ -672,12 +677,12 @@ class CoalesceVarAnalysis:
 
     norm_read_writes: FusedNormalizedReadsWrites
 
-    suggested_split: Optional[VarTiling] = None
+    suggested_split: VarTiling | None = None
 
 
 def analyze_memory_coalescing(
     fused_node: Union["FusedSchedulerNode", "SchedulerNode"],
-) -> Optional[CoalesceVarAnalysis]:
+) -> CoalesceVarAnalysis | None:
     """
     Find variables that coalesce the reads and writes and score the total size.
 
@@ -804,7 +809,7 @@ def analyze_memory_coalescing(
             norm_read_writes=norm_read_writes,
         )
 
-    best_tiling: Optional[tuple[sympy.Expr, int]] = None
+    best_tiling: tuple[sympy.Expr, int] | None = None
     best_tiling_score = 0
 
     for var, tiling_counter in tiling_scores.items():
diff --git a/torch/_inductor/triton_bundler.py b/torch/_inductor/triton_bundler.py
index c9267325b44db..ed095541d9a2e 100644
--- a/torch/_inductor/triton_bundler.py
+++ b/torch/_inductor/triton_bundler.py
@@ -5,11 +5,11 @@
 import shutil
 import uuid
 from pathlib import Path
-from typing import Optional
 
 from torch._dynamo.utils import counters, dynamo_timed, set_feature_use
 from torch._utils_internal import justknobs_check
 from torch.utils._filelock import FileLock
+from torch.utils._ordered_set import OrderedSet
 
 from .runtime.runtime_utils import triton_cache_dir
 from .utils import _IS_WINDOWS, GPU_KERNEL_BIN_EXTS
@@ -105,8 +105,9 @@ class TritonBundler:
     - TritonBundler.read_and_emit is called when a cache entry is read
     """
 
-    _entries: Optional[list[TritonBundleEntry]] = None
-    _static_autotuners: Optional[list[StaticallyLaunchedAutotuner]] = None
+    _entries: list[TritonBundleEntry] | None = None
+    _static_autotuners: list[StaticallyLaunchedAutotuner] | None = None
+    _winners: OrderedSet[str] | None = None
 
     # __grp__kernel_name.json contains metadata with source code paths
     # we use this as sentinel value for search and replace
@@ -141,6 +142,7 @@ def begin_compile(cls) -> None:
         assert cls._entries is None
         cls._entries = []
         cls._static_autotuners = []
+        cls._winners = OrderedSet()
 
     @classmethod
     def end_compile(cls) -> None:
@@ -151,6 +153,7 @@ def end_compile(cls) -> None:
         log.debug("TritonBundler.end_compile is called")
         cls._entries = None
         cls._static_autotuners = None
+        cls._winners = None
 
     @classmethod
     def put(cls, kernel_hash: str, device: int) -> None:
@@ -163,6 +166,17 @@ def put(cls, kernel_hash: str, device: int) -> None:
                 TritonBundleEntry(kernel_hash, device, triton_cache_dir(device))
             )
 
+    @classmethod
+    def put_winner(cls, kernel_hash: str) -> None:
+        """
+        Marks a kernel hash as a winning autotuning config. Only winning
+        kernels are included in the bundle by collect(). If no winners are
+        recorded (e.g. single-config kernels that skip autotuning), all
+        entries are bundled.
+        """
+        if cls._winners is not None:
+            cls._winners.add(kernel_hash)
+
     @classmethod
     def put_static_autotuner(cls, key: str, kernel: "CachingAutotuner") -> None:  # type: ignore[name-defined] # noqa: F821
         from torch._inductor import config
@@ -204,7 +218,7 @@ def collect_static_autotuners(
 
     @classmethod
     def load_autotuners(
-        cls, static_autotuners: Optional[list[StaticallyLaunchedAutotuner]]
+        cls, static_autotuners: list[StaticallyLaunchedAutotuner] | None
     ) -> list[str]:
         """
         Load statically launchable CachingAutotuners into async_compile.CompiledTritonKernels
@@ -245,7 +259,7 @@ def load_autotuners(
     @classmethod
     def collect(
         cls,
-    ) -> tuple[TritonBundle, Optional[TritonBundlerMetadata]]:
+    ) -> tuple[TritonBundle, TritonBundlerMetadata | None]:
         """
         This is the main function called when a cache write happens. This function
         converts all the previously remembered kernels into bundled format so that
@@ -263,9 +277,17 @@ def collect(
         with dynamo_timed(key="TritonBundler.collect", log_pt2_compile_event=True):
             entries = cls._entries
             if entries is not None:
+                # Only bundle winning autotuning configs. If _winners is
+                # non-empty, skip entries whose kernel_hash is not a winner.
+                # When _winners is empty (single-config kernels, or no
+                # autotuning ran), bundle everything.
+                winners = cls._winners
                 result: list[TritonKernelArtifacts] = []
                 kernel_names: list[str] = []
                 for entry in entries:
+                    if winners and entry.kernel_hash not in winners:
+                        log.debug("Skipping non-winning kernel %s", entry.kernel_hash)
+                        continue
                     artifacts: list[TritonKernelArtifact] = []
                     path = os.path.join(entry.directory, entry.kernel_hash)
                     if not os.path.exists(path):
@@ -324,7 +346,7 @@ def collect(
             return TritonBundle([], []), None
 
     @staticmethod
-    def read_and_emit(bundle: TritonBundle) -> Optional[TritonBundlerMetadata]:
+    def read_and_emit(bundle: TritonBundle) -> TritonBundlerMetadata | None:
         """
         This is the main function called when a cache read happens. This function
         converts the bundled format back into individual files and writes them
diff --git a/torch/_inductor/utils.py b/torch/_inductor/utils.py
index c9c60aa71eb0e..b613819215670 100644
--- a/torch/_inductor/utils.py
+++ b/torch/_inductor/utils.py
@@ -23,6 +23,7 @@
 import textwrap
 import time
 import unittest
+import warnings
 from collections.abc import (
     Callable,
     Collection,
@@ -42,13 +43,11 @@
     Generic,
     Literal,
     NamedTuple,
-    Optional,
     Protocol,
     TYPE_CHECKING,
     TypeAlias,
     TypeGuard,
     TypeVar,
-    Union,
 )
 from typing_extensions import dataclass_transform, ParamSpec, Self
 from unittest import mock
@@ -70,6 +69,7 @@
     "inductor_autotune_lookup_table",
 ]
 
+from torch.fx.experimental._size_hinting import _sympy_subs
 from torch.fx.experimental.symbolic_shapes import (
     free_symbols,
     free_unbacked_symbols,
@@ -138,7 +138,7 @@ def get_gpu_type() -> str:
 
 _T = TypeVar("_T")
 VarRanges = dict[sympy.Expr, sympy.Expr]
-InputType = Optional[Union[torch.Tensor, int, torch.SymInt]]
+InputType = torch.Tensor | int | torch.SymInt | None
 
 XPU_KERNEL_FORMAT = (
     "spv" if _IS_WINDOWS else os.getenv("TORCHINDUCTOR_XPU_KERNEL_FORMAT", "zebin")
@@ -202,7 +202,7 @@ class align(sympy.Function):
     is_integer = True
 
     @classmethod
-    def eval(cls, value: sympy.Expr) -> Optional[sympy.Expr]:
+    def eval(cls, value: sympy.Expr) -> sympy.Expr | None:
         if isinstance(value, (int, sympy.Integer)):
             return _align(int(value))
         if _is_aligned(value):
@@ -220,8 +220,8 @@ class GraphPartitionMap:
 
     # map partition input/output indices to graph input/output indices. None indicates
     # a partition input/output is not a graph input/output.
-    input_index_mapping: list[Optional[int]]
-    output_index_mapping: list[Optional[int]]
+    input_index_mapping: list[int | None]
+    output_index_mapping: list[int | None]
 
     # name of constants read/written by the graph partition
     constant_names: list[str]
@@ -396,22 +396,18 @@ def _do_bench_using_profiling(
             and event.name != "Context Sync"
         ]
     )
-    if len(filtered_events) % n_repeat != 0:
-        raise RuntimeError(
-            "Failed to divide all profiling events into #repeat groups. "
-            "#%s events: %d, #repeats: %s",
-            device_type,
-            len(filtered_events),
-            n_repeat,
-        )
-    num_event_per_group = len(filtered_events) / n_repeat
+    # Filter out cache.zero_() events by name pattern instead of relying on
+    # positional grouping. cache.zero_() may not always generate an event
+    # making positional assumptions unreliable.
+    # The kernel name contains "FillFunctor" when generated by cache.zero_().
     actual_events = EventList(
-        [
-            event
-            for i, event in enumerate(filtered_events)
-            if i % num_event_per_group != 0
-        ]
+        [event for event in filtered_events if "FillFunctor" not in event.name]
     )
+    if len(actual_events) == 0:
+        raise RuntimeError(
+            f"Failed to capture any events after filtering cache clearing events. "
+            f"{device_type} events: {len(filtered_events)}, repeats: {n_repeat}"
+        )
     actual_events._build_tree()
     actual_events = actual_events.key_averages()
 
@@ -439,7 +435,7 @@ def has_torchvision_roi_align() -> bool:
         return False
 
 
-def decode_device(device: Union[Optional[torch.device], str]) -> torch.device:
+def decode_device(device: torch.device | None | str) -> torch.device:
     if device is None:
         return torch.tensor(0.0).device  # default device
     if isinstance(device, str):
@@ -463,9 +459,7 @@ def unique(it: Iterable[_T]) -> ValuesView[_T]:
     return {id(x): x for x in it}.values()
 
 
-def ceildiv(
-    number: Union[int, sympy.Expr], denom: Union[int, sympy.Expr]
-) -> Union[int, sympy.Expr]:
+def ceildiv(number: int | sympy.Expr, denom: int | sympy.Expr) -> int | sympy.Expr:
     if isinstance(number, sympy.Expr) or isinstance(denom, sympy.Expr):
         return CeilDiv(sympy.sympify(number), sympy.sympify(denom))
     # TODO: There is a bug in a call to this function, to repro:
@@ -477,7 +471,7 @@ def ceildiv(
     return runtime_ceildiv(number, denom)
 
 
-def _type_of(key: Optional[torch.dtype]) -> str:
+def _type_of(key: torch.dtype | None) -> str:
     # Use the function here to get rid of dependencies on the Triton during the codegen.
     # Refer to Triton implementation here:
     # https://github.com/triton-lang/triton/blob/98b5945d2aef679e00ebca8e07c35c3658ec76de/python/triton/runtime/jit.py#L238
@@ -493,6 +487,8 @@ def _type_of(key: Optional[torch.dtype]) -> str:
         "float8e4b15x4": "fp8e4b15x4",
         "float8_e4m3fn": "fp8e4nv",
         "float8_e5m2": "fp8e5",
+        "float8_e4m3fnuz": "fp8e4b8",
+        "float8_e5m2fnuz": "fp8e5b16",
         # TODO: remove when support is added in triton
         # https://github.com/triton-lang/triton/issues/6054
         "float8_e8m0fnu": "u8",
@@ -516,7 +512,7 @@ def _type_of(key: Optional[torch.dtype]) -> str:
 
 
 def convert_shape_to_inductor(
-    lst: Iterable[Union[int, torch.SymInt]],
+    lst: Iterable[int | torch.SymInt],
 ) -> list[sympy.Expr]:
     """
     Gets the shape and stride of a tensor. For non-symbolic tensors, this is
@@ -526,7 +522,7 @@ def convert_shape_to_inductor(
     return [sympy.sympify(i) for i in lst]
 
 
-def convert_symint_to_expr(val: Union[int, torch.SymInt]) -> Union[int, sympy.Expr]:
+def convert_symint_to_expr(val: int | torch.SymInt) -> int | sympy.Expr:
     """
     Convert SymInt to sympy.Expr, leave int as is.
 
@@ -538,7 +534,7 @@ def convert_symint_to_expr(val: Union[int, torch.SymInt]) -> Union[int, sympy.Ex
     return val
 
 
-def convert_to_symint(i: Union[int, sympy.Expr]) -> Union[int, torch.SymInt]:
+def convert_to_symint(i: int | sympy.Expr) -> int | torch.SymInt:
     """
     Like convert_shape_to_symint, but operates on a single expression.
     """
@@ -556,8 +552,8 @@ def convert_to_symint(i: Union[int, sympy.Expr]) -> Union[int, torch.SymInt]:
 
 
 def convert_shape_to_symint(
-    lst: Iterable[Union[int, sympy.Expr]],
-) -> list[Union[int, torch.SymInt]]:
+    lst: Iterable[int | sympy.Expr],
+) -> list[int | torch.SymInt]:
     """
     Takes a list of shapes from Inductor and converts them into symints (or just
     ints if all shapes are static).
@@ -678,7 +674,7 @@ def cmp(a: int, b: int) -> int:
     return int(a > b) - int(a < b)
 
 
-def pad_listlike(x: Union[int, Sequence[int]], size: int) -> Sequence[int]:
+def pad_listlike(x: int | Sequence[int], size: int) -> Sequence[int]:
     if isinstance(x, int):
         return [x] * size
     if len(x) == 1:
@@ -802,7 +798,7 @@ def clear_cache(self: Any) -> None:
 
 
 def aggregate_origins(
-    node_schedule: Union[Sequence[BaseSchedulerNode], ExternKernel],
+    node_schedule: Sequence[BaseSchedulerNode] | ExternKernel,
 ) -> OrderedSet[Node]:
     from . import ir
 
@@ -879,7 +875,7 @@ def get_origin_meta_str(origin):
 
 
 def get_kernel_metadata(
-    node_schedule: Union[Sequence[BaseSchedulerNode], ExternKernel],
+    node_schedule: Sequence[BaseSchedulerNode] | ExternKernel,
     wrapper: PythonWrapperCodegen,
 ) -> tuple[str, str]:
     """
@@ -958,7 +954,7 @@ def get_kernel_metadata(
             from .virtualized import V
 
             def get_buffer_info(
-                buffer: Union[ir.TensorBox, ir.Buffer, ir.TorchBindObject], rw_name: str
+                buffer: ir.TensorBox | ir.Buffer | ir.TorchBindObject, rw_name: str
             ) -> tuple[str, ir.Layout | None]:
                 if isinstance(buffer, ir.TensorBox) and isinstance(
                     buffer.data, ir.StorageBox
@@ -1023,9 +1019,17 @@ def stringfy_layout(layout: ir.Layout | None) -> str:
                         all_writes.append("%" + output_name)
 
         for node in inductor_nodes:
-            detailed_metadata.append(
-                f"{wrapper.comment}   {node.format_node(include_tensor_metadata=True)}"
-            )
+            formatted_node = node.format_node(include_tensor_metadata=True)
+            if formatted_node is not None and torch.version.hip:
+                # AMDGCN asm strings can contain newlines, which propagate
+                # into format_node() output.  Split so every line gets the
+                # comment prefix; otherwise bare newlines break the wrapper.
+                detailed_metadata.extend(
+                    f"{wrapper.comment}   {line}"
+                    for line in formatted_node.splitlines()
+                )
+            else:
+                detailed_metadata.append(f"{wrapper.comment}   {formatted_node}")
 
         detailed_metadata.append(f"{wrapper.comment}   return {','.join(all_writes)}")
 
@@ -1034,7 +1038,7 @@ def stringfy_layout(layout: ir.Layout | None) -> str:
 
 def dominated_nodes(
     initial_queue: Iterable[torch.fx.Node],
-    skip_filter: Optional[Callable[[Any], bool]] = None,
+    skip_filter: Callable[[Any], bool] | None = None,
 ) -> OrderedSet[torch.fx.Node]:
     """Returns the set of nodes whose values depend on those within initial_queue"""
     initial_queue = list(initial_queue)
@@ -1179,27 +1183,10 @@ def sympy_subs(expr: sympy.Expr, replacements: dict[sympy.Expr, Any]) -> sympy.E
     When the passed replacement symbol v is a string, it is converted to a symbol with name v that
     have the same replaced expression integer and nonnegative properties.
     """
-
-    def to_symbol(
-        replaced: sympy.Expr, replacement: Union[sympy.Expr, str]
-    ) -> sympy.Symbol:
-        assert isinstance(replaced, sympy.Expr)
-        if isinstance(replacement, str):
-            return sympy.Symbol(
-                replacement,
-                integer=replaced.is_integer,  # type: ignore[attr-defined]
-                nonnegative=replaced.is_nonnegative,  # type: ignore[attr-defined]
-            )
-        else:
-            return replacement
-
-    # xreplace is faster than subs, but is way more picky
-    return sympy.sympify(expr).xreplace(
-        {k: to_symbol(k, v) for k, v in replacements.items()}
-    )
+    return _sympy_subs(expr, replacements)
 
 
-def is_symbolic(a: Any) -> TypeGuard[Union[torch.SymInt, torch.Tensor]]:
+def is_symbolic(a: Any) -> TypeGuard[torch.SymInt | torch.Tensor]:
     return isinstance(a, torch.SymInt) or (
         isinstance(a, torch.Tensor) and a._has_symbolic_sizes_strides
     )
@@ -1231,7 +1218,7 @@ def any_is_symbolic(*args: Any) -> bool:
 
 def get_first_incompatible_cudagraph_node(
     gm: torch.fx.GraphModule,
-) -> Optional[torch.fx.Node]:
+) -> torch.fx.Node | None:
     from torch.fx.experimental.symbolic_shapes import free_unbacked_symbols
 
     for node in gm.graph.nodes:
@@ -1349,8 +1336,8 @@ def _set_env(key: str, value: str) -> Iterator[None]:
 
 @contextlib.contextmanager
 def fresh_cache(
-    cache_entries: Optional[dict[str, Any]] = None,
-    dir: Optional[str] = None,
+    cache_entries: dict[str, Any] | None = None,
+    dir: str | None = None,
     delete: bool = True,
 ) -> Iterator[None]:
     """
@@ -1430,7 +1417,7 @@ def argsort(seq: Sequence[Any], *, reverse: bool = False) -> list[int]:
 
 def argsort_sym(
     shape_env: ShapeEnv,
-    seq: Sequence[Union[int, torch.SymInt, sympy.Expr]],
+    seq: Sequence[int | torch.SymInt | sympy.Expr],
     *,
     reverse: bool = False,
 ) -> list[int]:
@@ -1438,7 +1425,7 @@ def cmp(a: tuple[int, sympy.Expr], b: tuple[int, sympy.Expr]) -> int:
         a_idx, a_val = a
         b_idx, b_val = b
 
-        def evaluate(expr: Union[bool, torch.SymInt, sympy.Expr]) -> bool:
+        def evaluate(expr: bool | torch.SymInt | sympy.Expr) -> bool:
             if isinstance(expr, bool):
                 return expr
             return shape_env.evaluate_expr(expr, size_oblivious=True)
@@ -1490,7 +1477,7 @@ class IndentedBuffer:
     tabwidth = 4
 
     def __init__(self, initial_indent: int = 0) -> None:
-        self._lines: list[Union[DeferredLineBase, LineContext, str]] = []
+        self._lines: list[DeferredLineBase | LineContext | str] = []
         self._indent = initial_indent
 
     @contextlib.contextmanager
@@ -1545,6 +1532,9 @@ def getrawvalue(self) -> str:
                 buf.write("\n")
         return buf.getvalue()
 
+    def get_lines_ref(self):
+        return self._lines
+
     def clear(self) -> None:
         self._lines.clear()
 
@@ -1557,7 +1547,7 @@ def prefix(self) -> str:
     def newline(self) -> None:
         self.writeline("\n")
 
-    def writeline(self, line: Union[LineContext, DeferredLineBase, str]) -> None:
+    def writeline(self, line: LineContext | DeferredLineBase | str) -> None:
         if isinstance(line, LineContext):
             self._lines.append(line)
         elif isinstance(line, DeferredLineBase):
@@ -1567,9 +1557,7 @@ def writeline(self, line: Union[LineContext, DeferredLineBase, str]) -> None:
         else:
             self._lines.append("")
 
-    def writelines(
-        self, lines: Sequence[Union[LineContext, DeferredLineBase, str]]
-    ) -> None:
+    def writelines(self, lines: Sequence[LineContext | DeferredLineBase | str]) -> None:
         for line in lines:
             self.writeline(line)
 
@@ -1590,9 +1578,7 @@ def do_indent(self, offset: int = 1) -> None:
     def do_unindent(self, offset: int = 1) -> None:
         self._indent -= offset
 
-    def splice(
-        self, other_code: Union[IndentedBuffer, str], strip: bool = False
-    ) -> None:
+    def splice(self, other_code: IndentedBuffer | str, strip: bool = False) -> None:
         if isinstance(other_code, IndentedBuffer):
             dedent = float("inf")
 
@@ -1632,7 +1618,7 @@ def __add__(self, other: Self) -> IndentedBuffer:
         res.writelines(other._lines)
         return res
 
-    def contains(self, new_line: Union[DeferredLineBase, LineContext, str]) -> bool:
+    def contains(self, new_line: DeferredLineBase | LineContext | str) -> bool:
         return new_line in self._lines
 
 
@@ -1668,7 +1654,7 @@ def __init__(self, line: str):
             line = ""
         self.line = line
 
-    def __call__(self) -> Union[str, None]:
+    def __call__(self) -> str | None:
         """Returns either self.line or None to indicate the line has been 'unwritten'"""
         raise NotImplementedError
 
@@ -1682,7 +1668,7 @@ def with_prefix(self, prefix: str) -> Self:
     def lstrip(self) -> Self:
         return self._new_line(self.line.lstrip())
 
-    def __getitem__(self, index: Union[int, slice]) -> Self:
+    def __getitem__(self, index: int | slice) -> Self:
         return self._new_line(self.line[index])
 
     def __bool__(self) -> bool:
@@ -1708,7 +1694,7 @@ def _new_line(self, line: str) -> DelayReplaceLine:
 
 
 @functools.cache
-def is_big_gpu(index_or_device: Union[int, torch.device] = 0) -> bool:
+def is_big_gpu(index_or_device: int | torch.device = 0) -> bool:
     if isinstance(index_or_device, torch.device):
         device = index_or_device
     else:
@@ -1765,7 +1751,7 @@ def get_num_sms() -> int:
 def get_tma_workspace_arg(
     num_tma_descriptors: int,
     device: torch.device,
-    num_programs: Optional[int] = None,
+    num_programs: int | None = None,
 ) -> WorkspaceArg:
     """Builds and returns a WorkspaceArg for the device side TMA workspace buffer."""
     from .codegen.common import WorkspaceArg, WorkspaceZeroMode
@@ -1782,6 +1768,14 @@ def get_tma_workspace_arg(
     )
 
 
+def get_default_kpack(block_k: int = 16) -> int:
+    if not torch.version.hip:
+        return 0
+    if "gfx942" in torch.cuda.get_device_properties(0).gcnArchName and block_k <= 16:
+        return 1
+    return 2
+
+
 def _use_template_for_gpu(
     layout: Layout, allowed_layout_dtypes: list[torch.dtype]
 ) -> bool:
@@ -1840,7 +1834,7 @@ def use_triton_template(
 
 
 def can_use_tma(
-    *matrices: IRNode, output_layout: Optional[Layout] = None, add_guards: bool = False
+    *matrices: IRNode, output_layout: Layout | None = None, add_guards: bool = False
 ) -> bool:
     """
     Return True iff *all* supplied tensors satisfy the CUDA TMA constraints
@@ -1860,10 +1854,10 @@ def can_use_tma(
 
     from .virtualized import V
 
-    def _aligned(expr_bytes: Union[int, sympy.Expr]) -> bool:
+    def _aligned(expr_bytes: int | sympy.Expr) -> bool:
         return V.graph.sizevars.statically_known_multiple_of(expr_bytes, TMA_ALIGNMENT)
 
-    def _is_tma_compatible_layout(layout: Optional[Layout]) -> bool:
+    def _is_tma_compatible_layout(layout: Layout | None) -> bool:
         if layout is None:
             return True
         sizes = layout.size
@@ -1908,8 +1902,12 @@ def _is_tma_compatible(
             sizes_i = V.graph.sizevars.guard_int_seq(sizes)
             strides_i = V.graph.sizevars.guard_int_seq(strides)
         else:
-            sizes_i = [V.graph.sizevars.symbolic_hint(s) for s in sizes]
-            strides_i = [V.graph.sizevars.symbolic_hint(st) for st in strides]
+            sizes_i = [
+                V.graph.sizevars.replace_backed_symbols_with_hints(s) for s in sizes
+            ]
+            strides_i = [
+                V.graph.sizevars.replace_backed_symbols_with_hints(st) for st in strides
+            ]
 
         # Find the single contiguous ("inner") dim
         inner = [
@@ -1946,14 +1944,16 @@ def _is_tma_compatible_xpu(
     ) -> bool:
         # Make sure the last dimension is contiguous
         last_stride = strides[-1]
-        last_stride_hint = V.graph.sizevars.symbolic_hint(last_stride)
+        last_stride_hint = V.graph.sizevars.replace_backed_symbols_with_hints(
+            last_stride
+        )
         if not V.graph.sizevars.statically_known_equals(last_stride_hint, 1):
             return False
 
         # Triton's type of index is uint32, so all dimensions must fit in uint32
         MAX_UINT32 = 2**32 - 1
         for size in sizes:
-            size_hint = V.graph.sizevars.symbolic_hint(size)
+            size_hint = V.graph.sizevars.replace_backed_symbols_with_hints(size)
             if V.graph.sizevars.statically_known_gt(size_hint, MAX_UINT32):
                 return False
 
@@ -1966,15 +1966,53 @@ def _is_tma_compatible_xpu(
     )
 
 
+def _descriptor_shape_fits_in_int32(
+    sizes: Sequence[sympy.Expr], add_guards: bool = False
+) -> bool:
+    int32_max = torch.iinfo(torch.int32).max
+    conditions = []
+    for size in sizes:
+        if isinstance(size, (int, sympy.Integer)):
+            if size > int32_max:
+                return False
+        else:
+            conditions.append(sympy.Le(size, int32_max))
+
+    if not conditions:
+        return True
+
+    from .virtualized import V
+
+    condition = conditions[0] if len(conditions) == 1 else sympy.And(*conditions)
+    return (
+        V.graph.sizevars.guard_or_false(condition)
+        if add_guards
+        else V.graph.sizevars.statically_known_true(condition)
+    )
+
+
 def use_triton_tma_template(
     *matrices: IRNode, output_layout: Layout, add_guards: bool = False
 ) -> bool:
+    if not config.triton.enable_persistent_tma_matmul:
+        return False
+    if not all(len(m.get_size()) == 2 for m in matrices):
+        return False
+    if not all(
+        _descriptor_shape_fits_in_int32(m.get_size(), add_guards=add_guards)
+        for m in matrices
+    ):
+        return False
+    if config.triton.enable_template_tma_store and not _descriptor_shape_fits_in_int32(
+        output_layout.size, add_guards=add_guards
+    ):
+        return False
+    # On AMD (HIP), TMA is not available but we still use non-TMA persistent
+    # kernels, so skip the TMA compatibility checks.
+    if torch.version.hip is not None:
+        return True
     layout = output_layout if config.triton.enable_template_tma_store else None
-    return (
-        all(len(m.get_size()) == 2 for m in matrices)
-        and can_use_tma(*matrices, output_layout=layout, add_guards=add_guards)
-        and config.triton.enable_persistent_tma_matmul
-    )
+    return can_use_tma(*matrices, output_layout=layout, add_guards=add_guards)
 
 
 def use_triton_blackwell_tma_template(
@@ -2022,9 +2060,36 @@ def ensure_nv_universal_gemm_available() -> bool:
     in the same interpreter to retry the import.
     """
     try:
-        return importlib.util.find_spec("cutlass_api") is not None
+        available = importlib.util.find_spec("cutlass_api") is not None
     except ImportError:
         return False
+    if available:
+        _ensure_fp4_dtype_registered()
+    return available
+
+
+def _ensure_fp4_dtype_registered():
+    """Patch cutlass_api to handle torch.float4_e2m1fn_x2 -> cutlass.Float4E2M1FN.
+
+    NOTE: cutlass_api doesn't natively map this dtype. We patch the lookup function
+    in-place so all callers (including TensorWrapper) pick up the change.
+    Remove once cutlass_api adds native FP4 support.
+    """
+    import cutlass_api.utils
+
+    try:
+        cutlass_api.utils.cutlass_type_from_torch_type(torch.float4_e2m1fn_x2)
+    except (KeyError, AttributeError):
+        import cutlass
+
+        _orig = cutlass_api.utils.cutlass_type_from_torch_type
+
+        def _patched(dtype):
+            if dtype == torch.float4_e2m1fn_x2:
+                return cutlass.Float4E2M1FN
+            return _orig(dtype)
+
+        cutlass_api.utils.cutlass_type_from_torch_type = _patched
 
 
 @functools.lru_cache(maxsize=1)
@@ -2049,9 +2114,9 @@ def use_blackwell_cutedsl_grouped_mm(
     layout: Layout,
     a_is_2d: bool,
     b_is_2d: bool,
-    offs: Optional[Any],
-    bias: Optional[Any],
-    scale_result: Optional[Any],
+    offs: Any | None,
+    bias: Any | None,
+    scale_result: Any | None,
 ) -> bool:
     """
     Returns True if we can use the blackwell kernel for grouped mm.
@@ -2111,6 +2176,21 @@ def use_blackwell_cutedsl_grouped_mm(
 def use_cutlass_template(layout: Layout, m: int, n: int, k: int) -> bool:
     from .virtualized import V
 
+    # TODO: Enable CUTLASS in non-AOT cpp_wrapper mode. The CUTLASS
+    # codegen (CUDATemplateKernel.call_kernel) already has cpp_wrapper-aware
+    # arg handling, but the other half is missing: the non-triton branch of
+    # CppWrapperGpu._generate_kernel_call_helper unconditionally emits
+    # `kernels.<name>(...)`, and that AOTInductorModelKernels struct only
+    # exists in AOT mode. Fixing this requires adding dlopen/dlsym loading
+    # for the compiled CUTLASS .so, similar to how the triton branch uses
+    # static CUfunction + loadKernel for non-AOT mode.
+    if V.graph.cpp_wrapper and not V.graph.aot_mode:
+        warnings.warn(
+            "CUTLASS backend is not supported with non-AOT cpp_wrapper mode. "
+            "Skipping CUTLASS backend.",
+        )
+        return False
+
     gemm_size = V.graph.sizevars.optimization_hint(m * n * k, fallback=-1)
     if gemm_size <= 0 or gemm_size < config.cutlass.cutlass_backend_min_gemm_size:
         return False
@@ -2148,8 +2228,8 @@ def use_nv_universal_gemm_template(
     k: _IntLike,
     mat_a: IRNode,
     mat_b: IRNode,
-    offs: Optional[IRNode] = None,
-    g: Optional[_IntLike] = None,
+    offs: IRNode | None = None,
+    g: _IntLike | None = None,
 ) -> bool:
     """
     Return True if we can use the NVIDIA Universal GEMM Template.
@@ -2221,7 +2301,7 @@ def _use_cutlass_for_op(op_name: str) -> bool:
     return op_name.upper() in [x.strip() for x in enabled_ops.split(",")]
 
 
-_IntLike: TypeAlias = Union[int, sympy.Expr]
+_IntLike: TypeAlias = int | sympy.Expr
 
 
 @functools.cache
@@ -2333,7 +2413,7 @@ def _rocm_native_device_arch_name(device: str) -> str:
 
 @functools.cache
 def try_import_ck_lib() -> tuple[
-    Optional[str], Callable[[], list[Any]], Callable[[], list[Any]], type[Any]
+    str | None, Callable[[], list[Any]], Callable[[], list[Any]], type[Any]
 ]:
     try:
         import ck4inductor  # type: ignore[import]
@@ -2429,7 +2509,7 @@ def _use_template_for_cpu(layout: Layout) -> bool:
 
 
 def use_cpp_bmm_template(
-    layout: Layout, mat1: Union[ReinterpretView, Buffer], mat2: IRNode
+    layout: Layout, mat1: ReinterpretView | Buffer, mat2: IRNode
 ) -> bool:
     from .ir import Layout
 
@@ -2460,7 +2540,7 @@ def use_cpp_gemm_template(
     mat2_transposed: bool = False,
     require_constant_mat2: bool = True,
     is_woq_int4: bool = False,
-    q_group_size: Optional[int] = None,
+    q_group_size: int | None = None,
 ) -> bool:
     from . import ir
     from .codegen.cpp_micro_gemm import create_micro_gemm
@@ -2566,7 +2646,12 @@ def run_and_get_kernels(
     result, source_codes = run_and_get_code(fn, *args, **kwargs)
     kernels = []
     for code in source_codes:
-        kernels.extend(re.findall(r"'''.*?'''", code, re.DOTALL))
+        if config.cpp_wrapper and config.triton.autotune_at_compile_time is not True:
+            # With lazy Triton kernel compilation, kernel sources are embedded
+            # inside C++ R"TRITON(...)TRITON" raw strings.
+            kernels.extend(re.findall(r'R"TRITON\((.*?)\)TRITON"', code, re.DOTALL))
+        else:
+            kernels.extend(re.findall(r"'''.*?'''", code, re.DOTALL))
         if remove_quote:
             kernels = [kernel[3:-3] for kernel in kernels]
     return result, kernels
@@ -2686,7 +2771,7 @@ def override_lowering(
 
 
 def add_scheduler_init_hook(
-    pre_fn: Callable[..., Any], post_fn: Optional[Callable[..., Any]] = None
+    pre_fn: Callable[..., Any], post_fn: Callable[..., Any] | None = None
 ) -> Any:
     """
     Add hook functions to be called at the beginning and end of Scheduler.__init__.
@@ -2718,7 +2803,7 @@ def developer_warning(msg: str) -> None:
         log.info(msg)
 
 
-def get_benchmark_name() -> Optional[str]:
+def get_benchmark_name() -> str | None:
     """
     An experimental API used only when config.benchmark_kernel is true.
 
@@ -2859,14 +2944,16 @@ def get_gpu_shared_memory() -> int:
 
 def get_max_numwarps() -> int:
     if torch.cuda.is_available():
-        warp_size = torch.cuda.get_device_properties().warp_size
-        # pyrefly: ignore [missing-attribute]
-        max_threads_per_block = torch.cuda.get_device_properties().max_threads_per_block
-    else:
-        # Defaults
-        warp_size = 32
-        max_threads_per_block = 1024
-    return max_threads_per_block // warp_size
+        device = torch.device("cuda", torch.cuda.current_device())
+        props = DeviceProperties.create(device)
+        warp_size = props.warp_size_or_default
+        max_threads_per_block = props.max_threads_per_block
+        if max_threads_per_block is None:
+            raise AssertionError("expected max_threads_per_block to be set")
+        return max_threads_per_block // warp_size
+
+    log.debug("CUDA is not available; defaulting max num warps to 32")
+    return 32
 
 
 def is_welford_reduction(reduction_type: str) -> bool:
@@ -2959,19 +3046,20 @@ def pass_execution_and_save(
         )
 
 
-def is_multi_outputs_template(input_buf: Optional[Union[Buffer, Operation]]) -> bool:
+def is_multi_outputs_template(input_buf: Buffer | Operation | None) -> bool:
     """
     Check if input buffer is a multi-outputs template buffer
     """
     from . import ir
 
-    return isinstance(input_buf, ir.CppTemplateBuffer) and isinstance(
-        input_buf.layout, ir.MultiOutputLayout
+    return (
+        isinstance(input_buf, ir.TemplateBuffer)
+        and input_buf.is_multi_outputs_template()
     )
 
 
 def is_output_of_multi_outputs_template(
-    input_buf: Optional[Union[Buffer, Operation]],
+    input_buf: Buffer | Operation | None,
 ) -> bool:
     """
     Check if input buffer is a output of multi-outputs template buffer
@@ -2986,8 +3074,8 @@ def is_output_of_multi_outputs_template(
 
 
 def is_collective(
-    node: Optional[Union[Node, Operation]],
-    op: Optional[torch._ops.OperatorBase] = None,
+    node: Node | Operation | None,
+    op: torch._ops.OperatorBase | None = None,
 ) -> bool:
     if node is None:
         return False
@@ -3023,7 +3111,7 @@ def is_collective(
     )
 
 
-def is_wait(node: Optional[Union[IRNode, Operation]]) -> bool:
+def is_wait(node: IRNode | Operation | None) -> bool:
     from . import ir
 
     return type(node) is ir._WaitKernel
@@ -3031,7 +3119,7 @@ def is_wait(node: Optional[Union[IRNode, Operation]]) -> bool:
 
 def contains_collective(
     snode: BaseSchedulerNode,
-    filter_fn: Optional[Callable[[BaseSchedulerNode], bool]] = None,
+    filter_fn: Callable[[BaseSchedulerNode], bool] | None = None,
 ) -> bool:
     from torch._inductor.scheduler import GroupedSchedulerNode
 
@@ -3051,8 +3139,8 @@ def contains_wait(snode: BaseSchedulerNode) -> bool:
 
 
 def is_fallback_op(
-    node: Optional[Operation],
-    op: Union[torch._ops.OpOverload, Collection[torch._ops.OpOverload]],
+    node: Operation | None,
+    op: torch._ops.OpOverload | Collection[torch._ops.OpOverload],
 ) -> bool:
     from . import ir
 
@@ -3157,6 +3245,22 @@ def is_saved_tensor(x: Node) -> bool:
     return len(static_arg_idxs)
 
 
+def get_static_bw_input_idxs(fx_g: torch.fx.GraphModule) -> list[int]:
+    """
+    Returns indices of backward graph inputs that are always at fixed
+    addresses: primals (parameters/buffers/user inputs saved for backward).
+    Excludes saved activations which may not be at fixed addresses when
+    the forward is partitioned for CUDA graphs.
+    """
+    static_idxs = []
+    for idx, n in enumerate(fx_g.graph.nodes):
+        if n.op != "placeholder":
+            break
+        if n.name.startswith("primals_"):
+            static_idxs.append(idx)
+    return static_idxs
+
+
 @dataclasses.dataclass
 class BoxedBool:
     value: bool
@@ -3165,7 +3269,7 @@ def __bool__(self) -> bool:
         return self.value
 
     @staticmethod
-    def disable(obj: Any) -> Union[BoxedBool, bool]:
+    def disable(obj: Any) -> BoxedBool | bool:
         if isinstance(obj, BoxedBool):
             obj.value = False
             return obj
@@ -3182,9 +3286,9 @@ def define_kernel(
         self: PythonWrapperCodegen,
         kernel_name: str,
         kernel_code: str,
-        metadata: Optional[str] = None,
+        metadata: str | None = None,
         gpu: bool = True,
-        cpp_definition: Optional[str] = None,
+        cpp_definition: str | None = None,
     ) -> Any:
         kernel_list.append(kernel_code)
         return orig_define_kernel(
@@ -3199,7 +3303,7 @@ def get_cloned_parameter_buffer_name(name: str) -> str:
     return name + "__original__"
 
 
-def is_gpu(device: Optional[str]) -> bool:
+def is_gpu(device: str | None) -> bool:
     return device in GPU_TYPES
 
 
@@ -3223,7 +3327,7 @@ def needs_fallback_due_to_atomic_add_limitations(dtype: torch.dtype) -> bool:
 
 def use_scatter_fallback(
     op_overload: torch._ops.OpOverload,
-    reduction_type: Optional[str],
+    reduction_type: str | None,
     self_dtype: torch.dtype,
     src_dtype: torch.dtype,
     src_device_type: str,
@@ -3357,7 +3461,7 @@ def run_and_get_cpp_code(
     return result, s
 
 
-def shape_env_from_inputs(inputs: Sequence[InputType]) -> Optional[ShapeEnv]:
+def shape_env_from_inputs(inputs: Sequence[InputType]) -> ShapeEnv | None:
     fake_mode = detect_fake_mode(inputs)
 
     # TODO(voz): It would be nice to enable this assert, but there are lots of tests that
@@ -3425,7 +3529,7 @@ def clone_preserve_strides(x: torch.Tensor) -> torch.Tensor:
 def copy_misaligned_inputs(
     new_inputs: list[InputType],
     check_inputs_idxs: Sequence[int],
-    return_pair_idxs: Optional[OrderedSet[int]] = None,
+    return_pair_idxs: OrderedSet[int] | None = None,
 ) -> tuple[list[torch.Tensor], list[torch.Tensor]]:
     """
     Clones misaligned tensors which we inferred were aligned. Returns a tuple of [old_tensors], [new_tensors] for every
@@ -3471,11 +3575,17 @@ def remove_unaligned_input_idxs(
 
 
 def expr_fits_within_32bit(e: sympy.Expr) -> bool:
+    """Check if an expression fits within 32-bit integer range.
+
+    NOTE: This function intentionally does not install guards. Callers are
+    responsible for guarding (e.g. via check_leq) when they decide to use
+    32-bit indexing based on this result.
+    """
     from .virtualized import V
 
     int_max = torch.iinfo(torch.int32).max
-    size_hint = V.graph.sizevars.size_hint
-    has_hint = V.graph.sizevars.shape_env.has_hint
+    guarding_hint_or_throw = V.graph.sizevars.guarding_hint_or_throw
+    has_guarding_hint = V.graph.sizevars.shape_env.has_guarding_hint
 
     if config.assume_32bit_indexing:
         V.graph.sizevars.check_leq(e, int_max)  # type: ignore[arg-type]
@@ -3506,7 +3616,7 @@ def expr_fits_within_32bit(e: sympy.Expr) -> bool:
             return False
 
     # Otherwise, the hint MUST exist and be in range
-    return has_hint(e) and size_hint(e) <= int_max
+    return has_guarding_hint(e) and guarding_hint_or_throw(e) <= int_max
 
 
 def set_tracing_context_output_strides(
@@ -3526,7 +3636,7 @@ def set_tracing_context_output_strides(
                 if ctx := torch._guards.TracingContext.try_get():
                     fakify_first_call = ctx.fakify_first_call
 
-                def map_expr(e: Any) -> Union[float, int, SymInt, SymFloat, SymBool]:
+                def map_expr(e: Any) -> float | int | SymInt | SymFloat | SymBool:
                     if shape_env is None:
                         return int(e)
                     if fakify_first_call:
@@ -3638,7 +3748,7 @@ def boolean_ops() -> tuple[str, ...]:
 @dataclasses.dataclass
 class OpDtypeRule:
     type_promotion_kind: ELEMENTWISE_TYPE_PROMOTION_KIND
-    override_return_dtype: Optional[torch.dtype]
+    override_return_dtype: torch.dtype | None
 
 
 op_dtype_propagation_rules: dict[str, OpDtypeRule] = {}
@@ -3647,7 +3757,7 @@ class OpDtypeRule:
 def register_op_dtype_propagation_rules(
     name: str,
     type_promotion_kind: ELEMENTWISE_TYPE_PROMOTION_KIND,
-    override_return_dtype: Optional[torch.dtype],
+    override_return_dtype: torch.dtype | None,
 ) -> None:
     op_dtype_propagation_rules[name] = OpDtypeRule(
         type_promotion_kind, override_return_dtype
@@ -3661,7 +3771,7 @@ def register_op_requires_libdevice_fp64(name: str) -> None:
     op_requires_libdevice_fp64.add(name)
 
 
-def get_current_backend(device_type: Optional[str] = None) -> str:
+def get_current_backend(device_type: str | None = None) -> str:
     from torch._inductor.virtualized import V
 
     if not device_type:
@@ -3672,6 +3782,8 @@ def get_current_backend(device_type: Optional[str] = None) -> str:
         return "mps"
     elif device_type == "xpu":
         return config.xpu_backend
+    elif device_type == "tpu":
+        return config.tpu_backend
     else:
         return config.cuda_backend
 
@@ -3714,7 +3826,7 @@ def __setitem__(self, key: KeyType, value: ValType) -> None:
     def __contains__(self, key: object) -> bool:
         return key in self.new_items or key in self.original_dict
 
-    def get(self, key: KeyType, default: Optional[ValType] = None) -> Optional[ValType]:  # type: ignore[override]
+    def get(self, key: KeyType, default: ValType | None = None) -> ValType | None:  # type: ignore[override]
         if key in self.new_items:
             return self.new_items[key]
         return self.original_dict.get(key, default)
@@ -3740,7 +3852,7 @@ def __delitem__(self, key: KeyType) -> None:
 
 
 @dataclass_transform(frozen_default=True)
-def ir_dataclass(cls: Optional[type[Any]] = None, /, *, frozen: bool = True) -> Any:
+def ir_dataclass(cls: type[Any] | None = None, /, *, frozen: bool = True) -> Any:
     def wrap(cls: _T) -> _T:
         return dataclasses.dataclass(cls, kw_only=True, frozen=frozen)  # type: ignore[call-overload]
 
@@ -3749,7 +3861,7 @@ def wrap(cls: _T) -> _T:
     return wrap(cls)
 
 
-def get_donated_idxs() -> Optional[list[int]]:
+def get_donated_idxs() -> list[int] | None:
     tracing_context = torch._guards.TracingContext.try_get()
     if tracing_context is not None and tracing_context.fw_metadata:
         return tracing_context.fw_metadata.bw_donated_idxs
@@ -4234,7 +4346,7 @@ class CUDAGraphWrapperMetadata:
 
 # only incremented by user call of mark_step_begin
 class CUDAGraphWrapper:
-    wrapper: Optional[CUDAGraphWrapperType] = None
+    wrapper: CUDAGraphWrapperType | None = None
 
 
 # A customized partition wrappers from users. Interface should be:
@@ -4266,11 +4378,12 @@ def snode_args_kwargs(snode: BaseSchedulerNode) -> tuple[list[Any], dict[str, An
 
     def _is_tensor_ir(x) -> bool:  # type: ignore[no-untyped-def]
         return isinstance(x, torch._inductor.ir.IRNode) and not isinstance(
-            x, torch._inductor.ir.GeneratorState
+            x,
+            (torch._inductor.ir.GeneratorState, torch._inductor.ir.OpaqueObjectState),
         )
 
     flat_args = [
-        torch._inductor.ir.ir_node_to_tensor(a, guard_shape=False)
+        torch._inductor.ir.ir_node_to_tensor(a, replace_symbols_with_hints=True)
         if _is_tensor_ir(a)
         else a
         for a in flat_args
@@ -4359,6 +4472,9 @@ def should_fallback_by_default(node: torch.fx.Node) -> bool:
         "torch.ops._c10d_functional_autograd.all_gather_into_tensor.default",
         "torch.ops._c10d_functional_autograd.reduce_scatter_tensor.default",
         "torch.ops._c10d_functional_autograd.all_to_all_single.default",
+        "torch.ops._c10d_functional.isend.default",
+        "torch.ops._c10d_functional.irecv.default",
+        "torch.ops._c10d_functional.batch_p2p_ops.default",
     ]
 )
 
@@ -4395,7 +4511,7 @@ def _infer_scale_swizzle_impl(
     mat_dtype: torch.dtype,
     scale_dtype: torch.dtype,
     eq_fn: Callable[[Any, Any], bool],
-) -> tuple[Optional[Any], Optional[Any]]:
+) -> tuple[Any | None, Any | None]:
     """
     Core implementation for scale/swizzle inference.
     """
@@ -4470,7 +4586,7 @@ def _infer_scale_swizzle_impl(
 
 def infer_scale_swizzle(
     mat: torch.Tensor, scale: torch.Tensor
-) -> tuple[Optional[Any], Optional[Any]]:
+) -> tuple[Any | None, Any | None]:
     """
     Infer the scaling type and swizzle mode from matrix and scale tensor shapes/dtypes.
 
@@ -4502,7 +4618,7 @@ def infer_scale_swizzle_ir(
     mat: Buffer,
     scale: Buffer,
     transpose: bool = False,
-) -> tuple[Optional[Any], Optional[Any]]:
+) -> tuple[Any | None, Any | None]:
     """
     Infer the scaling type and swizzle mode for IR nodes (used during graph lowering).
 
diff --git a/torch/_inductor/virtualized.py b/torch/_inductor/virtualized.py
index f45e372e2b3a3..c083821dba77f 100644
--- a/torch/_inductor/virtualized.py
+++ b/torch/_inductor/virtualized.py
@@ -10,8 +10,8 @@
     from .virtualized import V
 
 Various handlers are accessible as attributes on this module; for example,
-you might access ``V.graph.sizevars.size_hint`` to resolve a size hint associated with
-a number.
+you might access ``V.graph.sizevars.optimzations_hint`` to resolve a hint
+associated with a symbolic expression.
 
 There are a few distinct usage patterns for virtualized global variables:
 
@@ -59,7 +59,7 @@
 
 from contextlib import AbstractContextManager, contextmanager
 from threading import local
-from typing import Any, cast, Generic, TYPE_CHECKING, TypeVar, Union
+from typing import Any, cast, Generic, TYPE_CHECKING, TypeVar
 
 from torch.utils._ordered_set import OrderedSet
 
@@ -122,7 +122,7 @@ class Virtualized(Generic[T]):
     store other things, like booleans.
     """
 
-    def __init__(self, vname: str, default: Union[Callable[[], T], type[NullHandler]]):
+    def __init__(self, vname: str, default: Callable[[], T] | type[NullHandler]):
         self._vname = vname
         self._key: str = f"__torchinductor_{vname}"
         self._default = default
@@ -208,6 +208,18 @@ def get_index_dtype_as_torch_dtype(self):
 )
 
 
+def _active_user_lowering_ops_default() -> OrderedSet[Any]:
+    """Default factory for active_user_lowering_ops - returns persisted empty set."""
+    rv: OrderedSet[Any] = OrderedSet()
+    setattr(threadlocal, _active_user_lowering_ops._key, rv)
+    return rv
+
+
+_active_user_lowering_ops: Virtualized[OrderedSet[Any]] = Virtualized(
+    "active_user_lowering_ops", _active_user_lowering_ops_default
+)
+
+
 def _choices_default():
     """
     Lazy init the global choices handler
@@ -339,6 +351,7 @@ def _wrap(x):
         return OpsValue(x)
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def indirect_indexing(index, size, check=True, wrap_neg=True):
         # Returns a sympy value, not IR value
         index = OpsWrapper._unwrap(index)
@@ -381,6 +394,12 @@ class _V:
     get_distributed_autotune_state: Callable[[], Any] = (
         _distributed_autotune_state._get_handler
     )
+    set_active_user_lowering_ops: Callable[[Any], Any] = (
+        _active_user_lowering_ops._set_handler
+    )
+    get_active_user_lowering_ops: Callable[[], OrderedSet[Any]] = (
+        _active_user_lowering_ops._get_handler
+    )
 
     @property
     def ops(self) -> OpsHandler[Any]:
@@ -444,5 +463,10 @@ def choices(self) -> InductorChoices:
     def distributed_autotune_state(self):
         return _distributed_autotune_state._get_handler()
 
+    @property
+    def active_user_lowering_ops(self) -> OrderedSet[Any]:
+        """Set of ops currently being lowered via user_lowerings (for recursion guard)."""
+        return _active_user_lowering_ops._get_handler()
+
 
 V = _V()
diff --git a/torch/_inductor/wrapper_benchmark.py b/torch/_inductor/wrapper_benchmark.py
index 56adde809079f..2b727ccdd407c 100644
--- a/torch/_inductor/wrapper_benchmark.py
+++ b/torch/_inductor/wrapper_benchmark.py
@@ -4,7 +4,7 @@
 from collections import defaultdict
 from dataclasses import dataclass
 from types import ModuleType
-from typing import Any, Optional, Protocol
+from typing import Any, Protocol
 
 import torch
 from torch.autograd import DeviceType
@@ -73,7 +73,7 @@ def get_triton_kernel(mod: ModuleType):  # type: ignore[no-untyped-def]
 
 
 def benchmark_all_kernels(
-    benchmark_name: str, benchmark_all_configs: Optional[dict[Any, Any]]
+    benchmark_name: str, benchmark_all_configs: dict[Any, Any] | None
 ) -> None:
     """
     An experimental API used only when config.benchmark_kernel is true.
@@ -109,9 +109,9 @@ def benchmark_all_kernels(
 
         def get_info_str(
             ms: float,
-            n_regs: Optional[Any],
-            n_spills: Optional[Any],
-            shared: Optional[Any],
+            n_regs: Any | None,
+            n_spills: Any | None,
+            shared: Any | None,
             prefix: str = "",
         ) -> str:
             if not any(x is None for x in [n_regs, n_spills, shared]):
diff --git a/torch/_jit_internal.py b/torch/_jit_internal.py
index 449c46de09c1e..d28a7011b4cce 100644
--- a/torch/_jit_internal.py
+++ b/torch/_jit_internal.py
@@ -10,6 +10,7 @@
 import collections
 import contextlib
 import enum
+import functools
 import inspect
 import io
 import pickle
@@ -376,8 +377,7 @@ def get_closure(fn):
 # annotations on `eg``, but starting in Python 4.0, they will represented as
 # strings and no longer present. Furthermore, since the body of `eg` does
 # not reference those names, they do not appear in the list of closed over
-# variables. In Python 2.x, type annotations are in comments, leading to a
-# similar situation where their definitions are not available. We anticipate
+# variables. We anticipate
 # that most users will not run into this issue because their modules and
 # functions will be defined at a global scope like MyGlobalClass. In cases
 # where they are not, it is possible to work around issues by declaring the
@@ -409,6 +409,7 @@ def __getattr__(self, key: str) -> Any:
     return createResolutionCallbackFromEnv(closure_lookup())
 
 
+@functools.cache
 def can_compile_class(cls) -> bool:
     # If any of the functions on a type don't have a code object, this type can't
     # be compiled and is probably a builtin / bound from C
@@ -1097,7 +1098,24 @@ def get_class_name_lineno(method) -> tuple[str, int]:
 
 
 def _overload_method(func):
-    _check_overload_body(func)
+    try:
+        _check_overload_body(func)
+    except IndentationError:
+        # CPython 3.13.8 has a bug (https://github.com/python/cpython/issues/139783)
+        # where inspect.getsourcelines() returns truncated source when a decorator
+        # is followed by a comment, causing ast.parse() to fail with IndentationError.
+        # Fixed in 3.13.9. Swallow the error on affected versions; re-raise otherwise.
+        if sys.version_info[:3] == (3, 13, 8):
+            import warnings
+
+            warnings.warn(
+                "Skipping overload body check due to a known CPython 3.13.8 bug "
+                "(https://github.com/python/cpython/issues/139783). "
+                "Consider upgrading to Python 3.13.9+.",
+                stacklevel=2,
+            )
+        else:
+            raise
     qual_name = _qualified_name(func)
     global _overloaded_methods
     class_name_map = _overloaded_methods.get(qual_name)
diff --git a/torch/_lazy/device_context.py b/torch/_lazy/device_context.py
index 49f33cf7f7c6d..c5e9d94266fe2 100644
--- a/torch/_lazy/device_context.py
+++ b/torch/_lazy/device_context.py
@@ -1,5 +1,5 @@
 import threading
-from typing import Any, Optional
+from typing import Any
 
 import torch._C._lazy
 
@@ -12,7 +12,7 @@ def __init__(self, device: str) -> None:
         self.device = device
 
 
-def get_device_context(device: Optional[str] = None) -> DeviceContext:
+def get_device_context(device: str | None = None) -> DeviceContext:
     if device is None:
         device = torch._C._lazy._get_default_device_type()
     else:
diff --git a/torch/_library/_out_variant.py b/torch/_library/_out_variant.py
index 8de7c817d661a..9e6960fb08940 100644
--- a/torch/_library/_out_variant.py
+++ b/torch/_library/_out_variant.py
@@ -8,6 +8,26 @@
 log = logging.getLogger(__name__)
 
 
+# Manual registry for ops whose out variant is not discoverable via
+# to_out_variant() (e.g. flat _out naming instead of .out overload).
+_manual_out_variant_registry: dict[torch._ops.OpOverload, torch._ops.OpOverload] = {}
+
+
+def register_out_variant(
+    functional_op: torch._ops.OpOverload,
+    out_op: torch._ops.OpOverload,
+) -> None:
+    """Register a functional op -> out variant mapping."""
+    _manual_out_variant_registry[functional_op] = out_op
+
+
+def lookup_manual_out_variant(
+    op: torch._ops.OpOverload,
+) -> torch._ops.OpOverload | None:
+    """Return the manually registered out variant for op, or None."""
+    return _manual_out_variant_registry.get(op)
+
+
 def _is_functional(schema: torch._C.FunctionSchema) -> bool:
     """
     A schema is functional if no argument is written to and the name doesn't
@@ -62,6 +82,12 @@ def _has_valid_out_variant_returns(
     return True
 
 
+def get_out_arg_names(out_op: torch._ops.OpOverload) -> list[str]:
+    """Get the names of out arguments for an out variant op."""
+    schema = out_op._schema
+    return [arg.name for arg in schema.arguments if _is_mutable_arg(arg)]
+
+
 def to_out_variant(op: torch._ops.OpOverload) -> torch._ops.OpOverload | None:
     """
     Given a functional operator overload, return its corresponding out variant.
diff --git a/torch/_library/autograd.py b/torch/_library/autograd.py
index 455efd83b3e64..7926c1a8e38a0 100644
--- a/torch/_library/autograd.py
+++ b/torch/_library/autograd.py
@@ -2,7 +2,7 @@
 import dataclasses
 from collections.abc import Callable
 from dataclasses import dataclass
-from typing import Any, Optional, Protocol
+from typing import Any, Protocol
 
 from torch import _C, _ops, autograd, Tensor
 from torch.utils import _pytree
@@ -11,14 +11,14 @@
 
 
 class InfoProtocol(Protocol):
-    _backward_fn: Optional[Callable]
-    _setup_context_fn: Optional[Callable]
+    _backward_fn: Callable | None
+    _setup_context_fn: Callable | None
 
 
 @dataclasses.dataclass
 class Info:
-    _backward_fn: Optional[Callable]
-    _setup_context_fn: Optional[Callable]
+    _backward_fn: Callable | None
+    _setup_context_fn: Callable | None
 
 
 def make_autograd_impl(op: _ops.OpOverload, info: InfoProtocol) -> Callable:
@@ -129,8 +129,8 @@ def supports_tensorlist(cls: Any) -> Any:
     @dataclass
     class Metadata:
         input_spec: _pytree.TreeSpec
-        output_spec: Optional[_pytree.TreeSpec] = None
-        result_is_tuple: Optional[bool] = None
+        output_spec: _pytree.TreeSpec | None = None
+        result_is_tuple: bool | None = None
 
     def new_forward(ctx, *args):
         metadata = args[-1]
diff --git a/torch/_library/custom_ops.py b/torch/_library/custom_ops.py
index c09a07bbdb138..d491d83743c6d 100644
--- a/torch/_library/custom_ops.py
+++ b/torch/_library/custom_ops.py
@@ -6,7 +6,7 @@
 import weakref
 from collections.abc import Callable, Iterable, Sequence
 from contextlib import contextmanager
-from typing import Any, Optional, overload, Union
+from typing import Any, overload, Union
 
 import torch
 from torch import _C, _ops, Tensor
@@ -17,7 +17,7 @@
 from .effects import EffectType
 
 
-device_types_t = Optional[Union[str, Sequence[str]]]
+device_types_t = str | Sequence[str] | None
 log = logging.getLogger(__name__)
 
 
@@ -27,10 +27,10 @@ def custom_op(
     fn: None = None,
     /,
     *,
-    mutates_args: Union[str, Iterable[str]],
+    mutates_args: str | Iterable[str],
     device_types: device_types_t = None,
-    schema: Optional[str] = None,
-    tags: Optional[Sequence[_C.Tag]] = None,
+    schema: str | None = None,
+    tags: Sequence[_C.Tag] | None = None,
 ) -> Callable[[Callable[..., object]], "CustomOpDef"]: ...
 
 
@@ -40,23 +40,23 @@ def custom_op(
     fn: Callable[..., object],
     /,
     *,
-    mutates_args: Union[str, Iterable[str]],
+    mutates_args: str | Iterable[str],
     device_types: device_types_t = None,
-    schema: Optional[str] = None,
-    tags: Optional[Sequence[_C.Tag]] = None,
+    schema: str | None = None,
+    tags: Sequence[_C.Tag] | None = None,
 ) -> "CustomOpDef": ...
 
 
 @exposed_in("torch.library")
 def custom_op(
     name: str,
-    fn: Optional[Callable] = None,
+    fn: Callable | None = None,
     /,
     *,
-    mutates_args: Union[str, Iterable[str]],
+    mutates_args: str | Iterable[str],
     device_types: device_types_t = None,
-    schema: Optional[str] = None,
-    tags: Optional[Sequence[_C.Tag]] = None,
+    schema: str | None = None,
+    tags: Sequence[_C.Tag] | None = None,
 ) -> Union[Callable[[Callable[..., object]], "CustomOpDef"], "CustomOpDef"]:
     """Wraps a function into custom operator.
 
@@ -90,6 +90,22 @@ def custom_op(
             have a specific reason not to.
             Example: "(Tensor x, int y) -> (Tensor, Tensor)".
 
+    The following types are supported for the wrapped function's input parameters:
+
+        - Scalars: ``int``, ``float``, ``bool``, ``str``, ``torch.types.Number``
+        - Tensors: ``torch.Tensor``
+        - Enums/devices: ``torch.dtype``, ``torch.device``
+        - Flat list of the same type: ``list[torch.Tensor]``,
+          ``list[int]``, ``list[float]``, ``list[bool]``,
+          ``list[torch.types.Number]``
+        - Optionals: ``Optional`` of any of the above scalar/tensor types
+        - Types registered via :func:`torch.library.register_opaque_type`
+
+    The following types are supported for the return value:
+
+        ``torch.Tensor``, ``list[torch.Tensor]``, ``int``, ``float``,
+        ``bool``, ``torch.types.Number``.
+
     .. note::
         We recommend not passing in a ``schema`` arg and instead letting us infer
         it from the type annotations. It is error-prone to write your own schema.
@@ -142,6 +158,18 @@ def custom_op(
         >>>     return torch.ones(3)
         >>>
         >>> bar("cpu")
+        >>>
+        >>> # Example of a custom op with list inputs
+        >>> @custom_op("mylib::weighted_sum", mutates_args=())
+        >>> def weighted_sum(
+        >>>     tensors: list[Tensor],
+        >>>     weights: list[float],
+        >>> ) -> Tensor:
+        >>>     return sum(t * w for t, w in zip(tensors, weights))
+        >>>
+        >>> x = torch.randn(3)
+        >>> y = torch.randn(3)
+        >>> out = weighted_sum([x, y], [0.3, 0.7])
 
     """
 
@@ -192,7 +220,7 @@ def __init__(
         name: str,
         schema: str,
         fn: Callable,
-        tags: Optional[Sequence[_C.Tag]] = None,
+        tags: Sequence[_C.Tag] | None = None,
     ) -> None:
         # Fields used to interface with the PyTorch dispatcher
         self._namespace = namespace
@@ -202,14 +230,14 @@ def __init__(
 
         self._init_fn = fn
 
-        self._backend_fns: dict[Union[str, None], Callable] = {}
-        self._abstract_fn: Optional[Callable] = None
-        self._setup_context_fn: Optional[Callable] = None
-        self._backward_fn: Optional[Callable] = None
+        self._backend_fns: dict[str | None, Callable] = {}
+        self._abstract_fn: Callable | None = None
+        self._setup_context_fn: Callable | None = None
+        self._backward_fn: Callable | None = None
         self._torch_dispatch_fns: dict[type, Callable] = {}
-        self._vmap_fn: Optional[Callable] = None
-        self._autocast_cuda_dtype: Optional[_dtype] = None
-        self._autocast_cpu_dtype: Optional[_dtype] = None
+        self._vmap_fn: Callable | None = None
+        self._autocast_cuda_dtype: _dtype | None = None
+        self._autocast_cpu_dtype: _dtype | None = None
 
         self._lib = get_library_allowing_overwrite(self._namespace, self._name)
         self._register_to_dispatcher(self._tags)
@@ -295,7 +323,7 @@ def set_kernel_enabled(self, device_type: str, enabled: bool = True):
                 self._disabled_kernel.discard(device_type)
 
     def register_kernel(
-        self, device_types: device_types_t, fn: Optional[Callable] = None, /
+        self, device_types: device_types_t, fn: Callable | None = None, /
     ) -> Callable:
         """Register an implementation for a device type for this operator.
 
@@ -337,7 +365,7 @@ def register_kernel(
 
         def inner(fn):
             if device_types is None or isinstance(device_types, str):
-                dtypes: list[Union[str, None]] = [device_types]
+                dtypes: list[str | None] = [device_types]
             else:
                 dtypes = list(device_types)
             for device_type in dtypes:
@@ -473,11 +501,11 @@ def register_fake(self, fn: Callable, /) -> Callable:
         self._abstract_fn = fn
         return fn
 
-    def register_effect(self, effect: Optional[EffectType]) -> None:
+    def register_effect(self, effect: EffectType | None) -> None:
         self._lib._register_effectful_op(self._qualname, effect)
 
     def register_torch_dispatch(
-        self, torch_dispatch_class: Any, fn: Optional[Callable] = None, /
+        self, torch_dispatch_class: Any, fn: Callable | None = None, /
     ) -> Callable:
         r"""Registers a torch_dispatch rule for the given operator and ``torch_dispatch_class``.
 
@@ -512,7 +540,7 @@ def register_autograd(
         backward: Callable,
         /,
         *,
-        setup_context: Optional[Callable] = None,
+        setup_context: Callable | None = None,
     ) -> None:
         r"""Register a backward formula for this custom op.
 
@@ -686,6 +714,7 @@ def backend_select(keyset, *args, **kwargs):
                     f"{self._name} does not have a kernel registered for {device}. "
                     "Please use register_kernel to do so."
                 )
+            # pyrefly: ignore [bad-argument-type]
             dispatch_key = _C._dispatch_key_for_device(device)
             dispatch_key = getattr(_C.DispatchKey, dispatch_key)
             return self._opoverload.redispatch(
@@ -699,7 +728,7 @@ def __call__(self, *args, **kwargs):
 
     def register_vmap(
         self,
-        func: Optional[Callable] = None,
+        func: Callable | None = None,
     ):
         r"""Register a vmap implementation to support :func:`torch.vmap` for this custom op.
 
@@ -945,8 +974,8 @@ def get_library_allowing_overwrite(
 
 
 def _maybe_get_opdef(
-    op: Union[CustomOpDef, _ops.OpOverload, str],
-) -> Optional[CustomOpDef]:
+    op: CustomOpDef | _ops.OpOverload | str,
+) -> CustomOpDef | None:
     if isinstance(op, CustomOpDef):
         return op
     if isinstance(op, _ops.OpOverload):
diff --git a/torch/_library/effects.py b/torch/_library/effects.py
index c723d14a00103..5aa91a5bdba31 100644
--- a/torch/_library/effects.py
+++ b/torch/_library/effects.py
@@ -1,5 +1,4 @@
 from enum import Enum
-from typing import Optional
 
 import torch
 
@@ -32,7 +31,7 @@ def __init__(self, qualname: str):
         self._set_default_effect()
 
     def _set_default_effect(self) -> None:
-        self._effect: Optional[EffectType] = None
+        self._effect: EffectType | None = None
 
         # If the op contains a ScriptObject input, we want to mark it as having effects
         namespace, opname = torch._library.utils.parse_namespace(self.qualname)
@@ -63,14 +62,14 @@ def _set_default_effect(self) -> None:
                     return
 
     @property
-    def effect(self) -> Optional[EffectType]:
+    def effect(self) -> EffectType | None:
         return self._effect
 
     @effect.setter
     def effect(self, _):
         raise RuntimeError("Unable to directly set kernel.")
 
-    def register(self, effect: Optional[EffectType]) -> RegistrationHandle:
+    def register(self, effect: EffectType | None) -> RegistrationHandle:
         """Register an effect
 
         Returns a RegistrationHandle that one can use to de-register this
diff --git a/torch/_library/fake_class_registry.py b/torch/_library/fake_class_registry.py
index 32bd8ad223a45..662e03b4b063a 100644
--- a/torch/_library/fake_class_registry.py
+++ b/torch/_library/fake_class_registry.py
@@ -1,7 +1,7 @@
 # mypy: allow-untyped-defs
 import copy
 import logging
-from typing import Any, Optional, Protocol, Union
+from typing import Any, Protocol
 
 import torch
 from torch._library.utils import parse_namespace
@@ -13,24 +13,32 @@
 
 class FakeScriptObject:
     def __init__(
-        self, wrapped_obj: Any, script_class_name: str, x: Optional[torch.ScriptObject]
+        self, wrapped_obj: Any, script_class_name: str, x: torch.ScriptObject | None
     ):
         # Use object.__setattr__ to bypass our custom __setattr__ during initialization
         object.__setattr__(self, "wrapped_obj", wrapped_obj)
         object.__setattr__(self, "script_class_name", script_class_name)
-        try:
-            with _disable_current_modes():
-                real_obj = copy.deepcopy(x)
-        except (RuntimeError, TypeError) as e:
-            log.warning(  # noqa: G200
-                "Unable to deepcopy the custom object %s due to %s. "
-                "Defaulting to the user given object. This might be "
-                "dangerous as side effects may be directly applied "
-                "to the object.",
-                script_class_name,
-                str(e),
-            )
-            real_obj = x
+
+        from torch._library.opaque_object import is_opaque_type
+
+        # We dont want to deepcopy when tracing with opaque objects because
+        # if a mutation happens intentionally (Ex. caching in device mesh)
+        # then we want it to be recorded on the real object
+        real_obj = x
+        if not is_opaque_type(type(x)):
+            try:
+                with _disable_current_modes():
+                    real_obj = copy.deepcopy(x)
+            except (RuntimeError, TypeError) as e:
+                log.warning(  # noqa: G200
+                    "Unable to deepcopy the custom object %s due to %s. "
+                    "Defaulting to the user given object. This might be "
+                    "dangerous as side effects may be directly applied "
+                    "to the object.",
+                    script_class_name,
+                    e,
+                )
+
         object.__setattr__(self, "real_obj", real_obj)
 
     def __getattribute__(self, name):
@@ -56,13 +64,84 @@ def __setattr__(self, name, value):
             "(but note that this is more difficult)."
         )
 
+    def __getitem__(self, key):
+        # This is needed for DeviceMesh support
+        return self.real_obj[key]
+
+    def __eq__(self, other):
+        if self is other:
+            return True
+        # Get real_obj without triggering custom __getattribute__
+        self_real = object.__getattribute__(self, "real_obj")
+        if isinstance(other, FakeScriptObject):
+            other_real = object.__getattribute__(other, "real_obj")
+            # For reference types, identity check first
+            if self_real is other_real:
+                return True
+            # Fall back to equality check
+            return self_real == other_real
+        # Compare with the real object directly
+        return self_real == other
+
+    def __ne__(self, other):
+        return not self.__eq__(other)
+
+    def __hash__(self):
+        # Use real_obj's hash if available, otherwise use object id
+        real_obj = object.__getattribute__(self, "real_obj")
+        try:
+            return hash(real_obj)
+        except TypeError:
+            # Object is not hashable, use identity-based hash
+            return id(real_obj)
+
+    def __deepcopy__(self, memo: dict[int, Any]) -> "FakeScriptObject":
+        if id(self) in memo:
+            return memo[id(self)]
+        new_obj = FakeScriptObject.__new__(FakeScriptObject)
+        memo[id(self)] = new_obj
+        object.__setattr__(
+            new_obj, "wrapped_obj", copy.deepcopy(self.wrapped_obj, memo)
+        )
+        object.__setattr__(new_obj, "script_class_name", self.script_class_name)
+        # Disable dispatch modes during deepcopy of real_obj and attribute
+        # access to prevent tensor operations (e.g. storage cloning, property
+        # access on DeviceMesh) from going through proxy tracing or
+        # functionalization.
+        with _disable_current_modes():
+            new_real_obj = copy.deepcopy(self.real_obj, memo)
+            object.__setattr__(new_obj, "real_obj", new_real_obj)
+            for name, value in self.__dict__.items():
+                if name not in ("wrapped_obj", "script_class_name", "real_obj"):
+                    if isinstance(value, FakeScriptMethod):
+                        object.__setattr__(
+                            new_obj,
+                            name,
+                            FakeScriptMethod(new_obj, value.method_name, value.schema),
+                        )
+                    else:
+                        if hasattr(new_real_obj, name):
+                            object.__setattr__(
+                                new_obj, name, getattr(new_real_obj, name)
+                            )
+                        else:
+                            object.__setattr__(new_obj, name, value)
+        return new_obj
+
+
+def maybe_unwrap_fake_script_object(obj: Any) -> Any:
+    """If obj is a FakeScriptObject, return the underlying real object."""
+    if isinstance(obj, FakeScriptObject):
+        return obj.real_obj
+    return obj
+
 
 class FakeScriptMethod:
     def __init__(
         self,
         self_fake_obj: FakeScriptObject,
         method_name: str,
-        schema: Optional[torch.FunctionSchema],
+        schema: torch.FunctionSchema | None,
     ):
         self.self_fake_obj = self_fake_obj
         self.method_name = method_name
@@ -152,9 +231,8 @@ def tracing_with_real(x: torch.ScriptObject) -> bool:
 def maybe_to_fake_obj(
     fake_mode,
     x: Any,
-) -> Union[FakeScriptObject, torch.ScriptObject]:
+) -> FakeScriptObject | torch.ScriptObject:
     import torch.utils._pytree as pytree
-    from torch.utils._python_dispatch import _disable_current_modes
 
     # When tracing with real mode, people should implement meta kernels that can
     # handle the case of real script object + fake tensor inputs.
@@ -168,7 +246,6 @@ def maybe_to_fake_obj(
         is_opaque_type,
         OpaqueTypeStr,
     )
-    from torch._subclasses.fake_tensor import unset_fake_temporarily
 
     x_type = type(x)
     if is_opaque_type(x_type):
@@ -180,7 +257,7 @@ def maybe_to_fake_obj(
         if opaque_info is None:
             raise AssertionError(f"opaque_info for type {x_type} must not be None")
         for attr_name in opaque_info.members:
-            with unset_fake_temporarily():
+            with _disable_current_modes():
                 if not hasattr(x, attr_name):
                     raise TypeError(
                         f"Opaque object of type '{type_name}' was specified to have member "
@@ -252,7 +329,7 @@ def maybe_to_fake_obj(
             real_attr = getattr(x, name)  # type: ignore[attr-defined]
 
             # real attr sometimes is not torch.ScriptMethod thus doesn't have schema e.g. __init___ or __eq__
-            method_schema: Optional[torch.FunctionSchema] = None
+            method_schema: torch.FunctionSchema | None = None
             if isinstance(real_attr, torch.ScriptMethod):
                 method_schema = real_attr.schema  # type: ignore[attr-defined]
 
@@ -269,7 +346,7 @@ def maybe_to_fake_obj(
     return fake_x_wrapped
 
 
-def register_fake_class(qualname, fake_class: Optional[HasStaticMethodFromReal] = None):
+def register_fake_class(qualname, fake_class: HasStaticMethodFromReal | None = None):
     r"""Register a fake implementation for this class.
 
     It's in the same spirit of registering a fake implementation for
@@ -369,7 +446,7 @@ def has_fake_class(full_qualname) -> bool:
     return global_fake_class_registry.has_impl(full_qualname)
 
 
-def find_fake_class(full_qualname) -> Optional[Any]:
+def find_fake_class(full_qualname) -> Any | None:
     if not has_fake_class(full_qualname):
         return None
     return global_fake_class_registry.get_impl(full_qualname)
diff --git a/torch/_library/fake_profile.py b/torch/_library/fake_profile.py
index 8949c775f7669..322b1ef9018e8 100644
--- a/torch/_library/fake_profile.py
+++ b/torch/_library/fake_profile.py
@@ -4,7 +4,7 @@
 import os
 from collections.abc import Callable, Generator
 from dataclasses import dataclass
-from typing import Any, Optional, Union
+from typing import Any, Optional
 
 import torch
 from torch._library.custom_ops import _maybe_get_opdef
@@ -37,20 +37,20 @@ def maybe_from_tensor(t: Any) -> Optional["TensorMetadata"]:
 
 @dataclass(frozen=True)
 class OpProfile:
-    args_profile: tuple[Optional[TensorMetadata]]
-    out_profile: Union[TensorMetadata, tuple[TensorMetadata]]
+    args_profile: tuple[TensorMetadata | None]
+    out_profile: TensorMetadata | tuple[TensorMetadata]
 
 
 def _generate_fake_kernel(op_name: str, op_profile: set[OpProfile]) -> Callable:
-    def _match_args(args_profile: tuple[Optional[TensorMetadata]], args: Any) -> bool:
+    def _match_args(args_profile: tuple[TensorMetadata | None], args: Any) -> bool:
         return all(
             TensorMetadata.maybe_from_tensor(arg) == args_profile[i]
             for i, arg in enumerate(args)
         )
 
     def _generate_res(
-        out_profile: Union[TensorMetadata, tuple[TensorMetadata]],
-    ) -> Union[torch.Tensor, list[torch.Tensor]]:
+        out_profile: TensorMetadata | tuple[TensorMetadata],
+    ) -> torch.Tensor | list[torch.Tensor]:
         ctx = torch.library.get_ctx()
 
         def _generate_tensor_out(t: TensorMetadata) -> torch.Tensor:
@@ -285,7 +285,7 @@ def deserialize_op_profile(data: dict) -> OpProfile:
             deserialize_tensor_metadata(arg) for arg in data["args_profile"]
         )
         out_profile_data = data["out_profile"]
-        out_profile: Union[tuple[TensorMetadata], TensorMetadata] = (
+        out_profile: tuple[TensorMetadata] | TensorMetadata = (
             tuple(deserialize_tensor_metadata(out) for out in out_profile_data)  # type: ignore[assignment]
             if isinstance(out_profile_data, list)
             else deserialize_tensor_metadata(out_profile_data)
diff --git a/torch/_library/infer_schema.py b/torch/_library/infer_schema.py
index 312e9a260d8c1..ae4ee9bd41a31 100644
--- a/torch/_library/infer_schema.py
+++ b/torch/_library/infer_schema.py
@@ -3,13 +3,16 @@
 import inspect
 import typing
 from types import GenericAlias
-from typing import Optional, Union
 
 import torch
 from torch import device, dtype, Tensor, types
 from torch.utils._exposed_in import exposed_in
 
-from .opaque_object import _OPAQUE_TYPES, is_opaque_type
+from .opaque_object import (
+    _resolve_opaque_type_info,
+    is_opaque_reference_type,
+    is_opaque_type,
+)
 
 
 # This is used as a negative test for
@@ -23,7 +26,7 @@ def infer_schema(
     /,
     *,
     mutates_args,
-    op_name: Optional[str] = None,
+    op_name: str | None = None,
 ) -> str:
     r"""Parses the schema of a given function with type hints. The schema is inferred from the
     function's type hints, and can be used to define a new operator.
@@ -79,7 +82,7 @@ def convert_type_string(annotation_type: str):
             )
 
     def unstringify_types(
-        tys: tuple[Union[type[object], str], ...],
+        tys: tuple[type[object] | str, ...],
     ) -> tuple[tuple[typing.Any, ...], bool]:
         res = []
         changed = False
@@ -92,7 +95,7 @@ def unstringify_types(
         else:
             return tys, False  # type: ignore[return-value]
 
-    def unstringify_type(ty: Union[type[object], str]) -> tuple[typing.Any, bool]:
+    def unstringify_type(ty: type[object] | str) -> tuple[typing.Any, bool]:
         # Dig through a generic type and if it contains a stringified type
         # convert that to a real type. The second return value indicates if the
         # type contained a string or not.
@@ -128,7 +131,7 @@ def unstringify_type(ty: Union[type[object], str]) -> tuple[typing.Any, bool]:
         schema_type = None
         if annotation_type not in SUPPORTED_PARAM_TYPES:
             if is_opaque_type(annotation_type):
-                schema_type = _OPAQUE_TYPES[annotation_type].class_name
+                schema_type = _resolve_opaque_type_info(annotation_type).class_name  # type: ignore[union-attr]
             elif annotation_type == torch._C.ScriptObject:
                 error_fn(
                     f"Parameter {name}'s type cannot be inferred from the schema "
@@ -217,19 +220,19 @@ def unstringify_type(ty: Union[type[object], str]) -> tuple[typing.Any, bool]:
 
 
 def derived_types(
-    base_type: Union[type, typing._SpecialForm],
+    base_type: type | typing._SpecialForm,
     cpp_type: str,
     list_base: bool,
     optional_base_list: bool,
     optional_list_base: bool,
 ):
-    result: list[tuple[Union[type, typing._SpecialForm, GenericAlias], str]] = [
+    result: list[tuple[type | typing._SpecialForm | GenericAlias, str]] = [
         (base_type, cpp_type),
         # pyrefly: ignore [not-a-type]
-        (typing.Optional[base_type], f"{cpp_type}?"),
+        (typing.Optional[base_type], f"{cpp_type}?"),  # noqa: UP045
     ]
 
-    def derived_seq_types(typ: Union[type, typing._SpecialForm]):
+    def derived_seq_types(typ: type | typing._SpecialForm):
         return (
             typing.Sequence[typ],  # type: ignore[valid-type]  # noqa: UP006
             typing.List[typ],  # type: ignore[valid-type]  # noqa: UP006
@@ -245,18 +248,18 @@ def derived_seq_types(typ: Union[type, typing._SpecialForm]):
         result.extend(
             (seq_typ, f"{cpp_type}?[]")
             # pyrefly: ignore [not-a-type]
-            for seq_typ in derived_seq_types(typing.Optional[base_type])
+            for seq_typ in derived_seq_types(typing.Optional[base_type])  # noqa: UP045
         )
     if optional_list_base:
         result.extend(
-            (typing.Optional[seq_typ], f"{cpp_type}[]?")
+            (typing.Optional[seq_typ], f"{cpp_type}[]?")  # noqa: UP045
             for seq_typ in derived_seq_types(base_type)
         )
     return result
 
 
 def get_supported_param_types():
-    data: list[tuple[Union[type, typing._SpecialForm], str, bool, bool, bool]] = [
+    data: list[tuple[type | typing._SpecialForm, str, bool, bool, bool]] = [
         # (python type, schema type, type[] variant, type?[] variant, type[]? variant
         (Tensor, "Tensor", True, True, False),
         (int, "SymInt", True, False, True),
@@ -300,6 +303,8 @@ def parse_return(annotation, error_fn):
     origin = typing.get_origin(annotation)
     if origin is not tuple:
         if annotation not in SUPPORTED_RETURN_TYPES:
+            if is_opaque_reference_type(annotation):
+                return _resolve_opaque_type_info(annotation).class_name  # type: ignore[union-attr]
             error_fn(
                 f"Return has unsupported type {annotation}. "
                 f"The valid types are: {SUPPORTED_RETURN_TYPES}."
@@ -309,12 +314,18 @@ def parse_return(annotation, error_fn):
 
     args = typing.get_args(annotation)
     for arg in args:
-        if arg not in SUPPORTED_RETURN_TYPES:
+        if arg not in SUPPORTED_RETURN_TYPES and not is_opaque_reference_type(arg):
             error_fn(
                 f"Return has unsupported type {annotation}. "
                 f"The valid types are: {SUPPORTED_RETURN_TYPES}."
             )
-    output_ty = ", ".join([SUPPORTED_RETURN_TYPES[arg] for arg in args])
+
+    def _return_type_str(arg):
+        if ty := SUPPORTED_RETURN_TYPES.get(arg):
+            return ty
+        return _resolve_opaque_type_info(arg).class_name  # type: ignore[union-attr]
+
+    output_ty = ", ".join(_return_type_str(arg) for arg in args)
 
     # use (()) to represent tuple with single element
     if len(args) == 1:
@@ -353,4 +364,4 @@ def tuple_to_list(tuple_type: type[tuple]) -> type[list]:
     elif len(type_args) == 2 and type_args[1] is Ellipsis:
         return list[type_args[0]]  # type: ignore[valid-type]
     else:
-        return list[typing.Union[tuple(type_args)]]  # type: ignore[misc, return-value]
+        return list[typing.Union[tuple(type_args)]]  # type: ignore[misc, return-value]  # noqa: UP007
diff --git a/torch/_library/opaque_object.py b/torch/_library/opaque_object.py
index b71f6d6ec8250..302e00f4d3bd6 100644
--- a/torch/_library/opaque_object.py
+++ b/torch/_library/opaque_object.py
@@ -35,19 +35,28 @@
 through `register_opaque_type(MyClass, typ="value")`.
 """
 
+import logging
 from collections.abc import Callable
 from dataclasses import dataclass
 from enum import Enum
-from typing import Any, Literal, NewType, Optional
+from typing import Any, Literal, NewType, TYPE_CHECKING, TypeAlias
 from typing_extensions import TypeIs
 from weakref import WeakKeyDictionary
 
 import torch
 from torch._opaque_base import OpaqueBase, OpaqueBaseMeta  # noqa: F401
 
+
+if TYPE_CHECKING:
+    from torch.fx import Proxy
+    from torch.fx.experimental.proxy_tensor import PythonKeyTracer
+
 from .fake_class_registry import register_fake_class
 
 
+log = logging.getLogger(__name__)
+
+
 class MemberType(Enum):
     """
     Defines how a member (attribute/property/method) of an opaque object is handled
@@ -77,6 +86,15 @@ def __obj_unflatten__(cls, flattened_ctx: dict[str, Any]) -> None:
 
 OpaqueType = NewType("OpaqueType", torch._C.ScriptObject)
 
+# Type for reconstruct_fn: called by PythonKeyTracer.create_arg when make_fx
+# encounters an untracked opaque reference (e.g. a backward closure capture).
+# Should derive the object from existing graph inputs or return None to fall
+# back to get_attr.  Args: (obj, get_tracked_proxy, tracer).
+ReconstructFn: TypeAlias = Callable[
+    [OpaqueBase, Callable[[OpaqueBase], "Proxy | None"], "PythonKeyTracer"],
+    "Proxy | None",
+]
+
 
 @dataclass
 class _OpaqueTypeInfo:
@@ -87,6 +105,7 @@ class _OpaqueTypeInfo:
     ]  # Callable that takes the object and returns list of values to guard on
     members: dict[str, MemberType]  # Maps member name to how it should be handled
     hoist: bool
+    reconstruct_fn: ReconstructFn | None
 
 
 # Mapping of type -> (string name, reference/value type)
@@ -95,6 +114,19 @@ class _OpaqueTypeInfo:
 _OPAQUE_TYPES_BY_NAME: dict[str, _OpaqueTypeInfo] = {}
 
 
+def _resolve_opaque_type_info(cls: Any) -> _OpaqueTypeInfo | None:
+    if cls in _OPAQUE_TYPES:
+        return _OPAQUE_TYPES[cls]
+    if not isinstance(cls, type):
+        return None
+
+    # Allow subclasses too
+    for parent in cls.__mro__[1:]:
+        if parent in _OPAQUE_TYPES:
+            return _OPAQUE_TYPES[parent]
+    return None
+
+
 def get_opaque_type_name(cls: Any) -> str:
     """
     Gets the registered opaque type name for a given class.
@@ -108,12 +140,13 @@ def get_opaque_type_name(cls: Any) -> str:
     Raises:
         ValueError: If the class is not registered as an opaque type.
     """
-    if cls not in _OPAQUE_TYPES:
+    info = _resolve_opaque_type_info(cls)
+    if info is None:
         raise ValueError(
             f"Class {cls} is not registered as an opaque type. "
             f"Call register_opaque_type({cls.__name__}) first."
         )
-    return _OPAQUE_TYPES[cls].class_name
+    return info.class_name
 
 
 def register_opaque_type(
@@ -123,6 +156,7 @@ def register_opaque_type(
     hoist=False,
     guard_fn: Any = None,
     members: dict[str, MemberType] | None = None,
+    reconstruct_fn: ReconstructFn | None = None,
 ) -> None:
     """
     Registers the given type as an opaque type which allows this to be consumed
@@ -169,7 +203,9 @@ def register_opaque_type(
             "registered as a pytree. Opaque objects must be pytree leaves."
         )
 
-    if not isinstance(cls, OpaqueBaseMeta):
+    # Value types store the real object directly during tracing (no
+    # FakeScriptObject wrapper), so they don't need OpaqueBaseMeta.
+    if typ != "value" and not isinstance(cls, OpaqueBaseMeta):
         raise TypeError(
             f"Opaque type {cls} must subclass torch._opaque_base.OpaqueBase "
             "or 'metaclass=torch._opaque_base.OpaqueBaseMeta'. "
@@ -184,7 +220,8 @@ def register_opaque_type(
         )
 
     if typ == "value":
-        if cls.__eq__ is object.__eq__:  # type: ignore[comparison-overlap]
+        # Enums use identity-based equality (singletons), which is fine for guarding.
+        if not issubclass(cls, Enum) and cls.__eq__ is object.__eq__:  # type: ignore[comparison-overlap]
             raise TypeError(
                 f"Value-type opaque object of type {cls} is "
                 "expected to have a non-default `__eq__` "
@@ -202,13 +239,14 @@ def register_opaque_type(
                 "for FakeTensor caching."
             )
 
-        if not hasattr(cls, "__fx_repr__"):
+        # Enums are special-cased in get_opaque_obj_repr.
+        if not issubclass(cls, Enum) and not hasattr(cls, "__fx_repr__"):
             raise TypeError(
                 f"Value-type opaque object of type {cls} is "
                 "expected to have a `__fx_repr__` method "
                 "implementation as we will use this to reconstruct "
                 "the object in the FX codegen. __fx_repr__ should return "
-                "a tuple of (repr_string, set_of_types)."
+                "a tuple of (repr_string, dict[str, type])."
             )
 
         if guard_fn is not None:
@@ -221,43 +259,64 @@ def register_opaque_type(
     # Generate a fully qualified name by combining module and qualname
     name = f"{cls.__module__}.{cls.__qualname__}"
 
-    type_info = _OpaqueTypeInfo(name, typ, guard_fn, members or {}, hoist)
+    type_info = _OpaqueTypeInfo(
+        name, typ, guard_fn, members or {}, hoist, reconstruct_fn
+    )
     _OPAQUE_TYPES[cls] = type_info
     _OPAQUE_TYPES_BY_NAME[name] = type_info
 
     torch._C._register_opaque_type(name)
 
 
+# Enums are always opaque value types.
+register_opaque_type(Enum, typ="value")
+
+
 def is_opaque_value(value: object) -> TypeIs[OpaqueType]:
     return is_opaque_type(type(value))
 
 
 def should_hoist(cls: Any) -> bool:
-    if cls not in _OPAQUE_TYPES:
+    info = _resolve_opaque_type_info(cls)
+    if info is None:
         return False
-    return _OPAQUE_TYPES[cls].hoist
+    return info.hoist
+
+
+def get_reconstruct_fn(cls: type[OpaqueBase]) -> ReconstructFn | None:
+    info = _resolve_opaque_type_info(cls)
+    if info is None:
+        return None
+    return info.reconstruct_fn
 
 
 def has_members(cls: Any) -> bool:
-    if cls not in _OPAQUE_TYPES:
+    info = _resolve_opaque_type_info(cls)
+    if info is None:
         return False
-    return len(_OPAQUE_TYPES[cls].members) > 0
+    return len(info.members) > 0
 
 
-def is_opaque_type(cls: Any) -> bool:
+def is_opaque_type(cls: type[Any] | str) -> bool:
     """
     Checks if the given type is an opaque type.
+    Also returns True for subclasses of registered opaque types.
     """
     if isinstance(cls, str):
         return torch._C._is_opaque_type_registered(cls)
 
-    if cls not in _OPAQUE_TYPES:
+    if not isinstance(cls, type):
+        log.warning("Passed invalid type `%s` to is_opaque_type, returning False", cls)
+        return False
+
+    info = _resolve_opaque_type_info(cls)
+    if info is None:
         return False
 
-    return torch._C._is_opaque_type_registered(_OPAQUE_TYPES[cls].class_name)
+    return torch._C._is_opaque_type_registered(info.class_name)
 
 
-def is_opaque_value_type(cls: Any) -> bool:
+def is_opaque_value_type(cls: type[Any] | str) -> bool:
     """
     Checks if the given type is an opaque **value** type.
     See Note [Opaque Objects] for more information.
@@ -268,7 +327,10 @@ def is_opaque_value_type(cls: Any) -> bool:
     if isinstance(cls, str):
         return _OPAQUE_TYPES_BY_NAME[cls].opaque_typ == "value"
 
-    return _OPAQUE_TYPES[cls].opaque_typ == "value"
+    info = _resolve_opaque_type_info(cls)
+    if info is None:
+        return False
+    return info.opaque_typ == "value"
 
 
 def is_opaque_reference_type(cls: Any) -> bool:
@@ -282,7 +344,10 @@ def is_opaque_reference_type(cls: Any) -> bool:
     if isinstance(cls, str):
         return _OPAQUE_TYPES_BY_NAME[cls].opaque_typ == "reference"
 
-    return _OPAQUE_TYPES[cls].opaque_typ == "reference"
+    info = _resolve_opaque_type_info(cls)
+    if info is None:
+        return False
+    return info.opaque_typ == "reference"
 
 
 def get_opaque_obj_repr(obj: Any) -> tuple[str, dict[str, type]]:
@@ -298,6 +363,12 @@ def get_opaque_obj_repr(obj: Any) -> tuple[str, dict[str, type]]:
     For example, if repr_string is "Foo(bar=Bar(1))", the dict should be:
         {"Foo": Foo, "Bar": Bar}
     """
+
+    # Enums are special cased
+    if isinstance(obj, Enum):
+        cls = type(obj)
+        return f"{cls.__name__}.{obj.name}", {cls.__name__: cls}
+
     if not hasattr(obj, "__fx_repr__"):
         raise TypeError(
             f"Value-type opaque object of type {obj} is "
@@ -324,17 +395,17 @@ def get_opaque_obj_repr(obj: Any) -> tuple[str, dict[str, type]]:
     return repr_str, globals_dict
 
 
-def get_opaque_obj_info(cls: Any) -> Optional[_OpaqueTypeInfo]:
+def get_opaque_obj_info(cls: Any) -> _OpaqueTypeInfo | None:
     if not is_opaque_type(cls):
         return None
 
     if isinstance(cls, str):
         return _OPAQUE_TYPES_BY_NAME[cls]
 
-    return _OPAQUE_TYPES[cls]
+    return _resolve_opaque_type_info(cls)
 
 
-def get_member_type(cls: Any, member_name: str) -> Optional[MemberType]:
+def get_member_type(cls: Any, member_name: str) -> MemberType | None:
     """
     Get the MemberType for a specific member of an opaque object class.
 
diff --git a/torch/_library/simple_registry.py b/torch/_library/simple_registry.py
index 466f6cc68e52b..e3670b36358b2 100644
--- a/torch/_library/simple_registry.py
+++ b/torch/_library/simple_registry.py
@@ -1,5 +1,5 @@
 from collections.abc import Callable
-from typing import Any, Optional
+from typing import Any
 
 from .effects import EffectHolder
 from .fake_impl import FakeImplHolder
@@ -79,13 +79,13 @@ def deregister() -> None:
 
         return RegistrationHandle(deregister)
 
-    def find(self, torch_dispatch_class: type) -> Optional[Callable[..., Any]]:
+    def find(self, torch_dispatch_class: type) -> Callable[..., Any] | None:
         return self._data.get(torch_dispatch_class, None)
 
 
 def find_torch_dispatch_rule(
     op: Any, torch_dispatch_class: type
-) -> Optional[Callable[..., Any]]:
+) -> Callable[..., Any] | None:
     return singleton.find(op.__qualname__).torch_dispatch_rules.find(
         torch_dispatch_class
     )
diff --git a/torch/_library/triton.py b/torch/_library/triton.py
index 9fc2768d36aab..f9bef08b889e7 100644
--- a/torch/_library/triton.py
+++ b/torch/_library/triton.py
@@ -4,7 +4,7 @@
 import logging
 import threading
 from collections.abc import Callable, Generator, Iterable
-from typing import Any, Optional, Union
+from typing import Any
 
 from torch.utils._exposed_in import exposed_in
 
@@ -296,11 +296,11 @@ def resolve_names_to_kernels(
 @exposed_in("torch.library")
 def triton_op(
     name: str,
-    fn: Optional[Callable] = None,
+    fn: Callable | None = None,
     /,
     *,
-    mutates_args: Union[str, Iterable[str]],
-    schema: Optional[str] = None,
+    mutates_args: str | Iterable[str],
+    schema: str | None = None,
 ) -> Callable:
     """Create a custom operator whose implementation is backed by 1+ triton kernels.
 
diff --git a/torch/_library/utils.py b/torch/_library/utils.py
index 9d5c2ae45bad5..961594ddfa984 100644
--- a/torch/_library/utils.py
+++ b/torch/_library/utils.py
@@ -3,7 +3,7 @@
 import inspect
 import sys
 from collections.abc import Callable, Iterable, Iterator
-from typing import Any, Literal, Optional, overload, Union
+from typing import Any, Literal, overload
 
 import torch
 import torch.utils._pytree as pytree
@@ -350,7 +350,7 @@ def has_tensor_arg(schema: _C.FunctionSchema) -> bool:
     )
 
 
-def get_device_arg_index(schema: _C.FunctionSchema) -> Union[int, None]:
+def get_device_arg_index(schema: _C.FunctionSchema) -> int | None:
     """
     Given a schema, returns the id of the `device: torch.device` argument.
     If it does not exist, returns None.
@@ -551,7 +551,7 @@ def get_layout_constraint_tag(
 @overload
 def get_layout_constraint_tag(
     fn: Any, *, with_default: Literal[False]
-) -> Optional[_C.Tag]: ...
+) -> _C.Tag | None: ...
 
 
 def get_layout_constraint_tag(fn, *, with_default=True):
@@ -587,8 +587,8 @@ def get_layout_constraint_tag(fn, *, with_default=True):
 def is_impure(
     op: Callable,
     *,
-    args: Optional[tuple[Any, ...]] = None,
-    kwargs: Optional[dict[str, Any]] = None,
+    args: tuple[Any, ...] | None = None,
+    kwargs: dict[str, Any] | None = None,
     impure_random: bool = True,
 ) -> bool:
     """
@@ -634,6 +634,9 @@ def is_impure(
         if _get_effect(op) is not None:
             return True
 
+        if op in _side_effectful_functions:
+            return True
+
         return False
 
     # Impure since it mutates RNG state
diff --git a/torch/_logging/_internal.py b/torch/_logging/_internal.py
index 97debeaa529d1..df9d54f11c12c 100644
--- a/torch/_logging/_internal.py
+++ b/torch/_logging/_internal.py
@@ -18,7 +18,7 @@
 from collections import defaultdict
 from collections.abc import Callable, Sequence
 from dataclasses import dataclass, field
-from typing import Any, Generic, Optional, Union
+from typing import Any, Generic, Optional
 from typing_extensions import ParamSpec
 from weakref import WeakSet
 
@@ -105,7 +105,7 @@ def is_log(self, alias):
         return alias in self.log_alias_to_log_qnames
 
     # register a log with an alias
-    def register_log(self, alias, log_qnames: Union[str, list[str]]) -> None:
+    def register_log(self, alias, log_qnames: str | list[str]) -> None:
         if isinstance(log_qnames, str):
             log_qnames = [log_qnames]
         self.log_alias_to_log_qnames[alias] = log_qnames
@@ -209,18 +209,18 @@ def clear(self) -> None:
 
 def set_logs(
     *,
-    all: Optional[int] = None,
-    dynamo: Optional[int] = None,
-    aot: Optional[int] = None,
-    autograd: Optional[int] = None,
-    dynamic: Optional[int] = None,
-    inductor: Optional[int] = None,
-    distributed: Optional[int] = None,
-    c10d: Optional[int] = None,
-    ddp: Optional[int] = None,
-    fsdp: Optional[int] = None,
-    dtensor: Optional[int] = None,
-    onnx: Optional[int] = None,
+    all: int | None = None,
+    dynamo: int | None = None,
+    aot: int | None = None,
+    autograd: int | None = None,
+    dynamic: int | None = None,
+    inductor: int | None = None,
+    distributed: int | None = None,
+    c10d: int | None = None,
+    ddp: int | None = None,
+    fsdp: int | None = None,
+    dtensor: int | None = None,
+    onnx: int | None = None,
     bytecode: bool = False,
     aot_graphs: bool = False,
     aot_joint_graph: bool = False,
@@ -248,8 +248,8 @@ def set_logs(
     onnx_diagnostics: bool = False,
     fusion: bool = False,
     overlap: bool = False,
-    export: Optional[int] = None,
-    modules: Optional[dict[str, Union[int, bool]]] = None,
+    export: int | None = None,
+    modules: dict[str, int | bool] | None = None,
     cudagraphs: bool = False,
     sym_node: bool = False,
     compiled_autograd: bool = False,
@@ -944,7 +944,7 @@ def make_module_path_relative(abs_path: str) -> str:
 # apply custom formats to artifacts when necessary
 class TorchLogsFormatter(logging.Formatter):
     def __init__(
-        self, *, trace: bool = False, trace_id_filter: Optional[set[str]] = None
+        self, *, trace: bool = False, trace_id_filter: set[str] | None = None
     ) -> None:
         super().__init__()
         self._is_trace = trace
@@ -1024,11 +1024,7 @@ def format(self, record):
         if self._is_trace:
             if s != "":
                 raise AssertionError(f"expected empty string for trace, got {s!r}")
-            try:
-                r = f"{prefix} {json.dumps(record.metadata)}"
-            except TypeError:
-                log.warning("failing metadata: %r", record.metadata)
-                raise
+            r = f"{prefix} {json.dumps(record.metadata, default=repr)}"
             if record.payload is not None:
                 r += "".join(f"\n\t{l}" for l in record.payload.split("\n"))
             return r
@@ -1162,6 +1158,13 @@ def _init_logs(log_file_name=None) -> None:
     # configuration
     trace_dir_name = os.environ.get(TRACE_ENV_VAR, None)
 
+    # If TORCH_COMPILE_DEBUG=1 is set but no TORCH_TRACE, automatically use
+    # the torch_compile_debug directory for trace logs (to simplify tlparse usage)
+    if trace_dir_name is None and os.environ.get("TORCH_COMPILE_DEBUG", "0") == "1":
+        import torch._dynamo.config as dynamo_config
+
+        trace_dir_name = os.path.join(dynamo_config.debug_dir_root, "tlparse")
+
     if dtrace_dir_name := os.environ.get(DTRACE_ENV_VAR, None):
         GET_DTRACE_STRUCTURED = True
         trace_dir_name = dtrace_dir_name
@@ -1186,13 +1189,14 @@ def _init_logs(log_file_name=None) -> None:
 class LazyTraceHandler(logging.StreamHandler):
     """Like FileHandler, but the file is allocated lazily only upon the first log message"""
 
-    def __init__(self, root_dir: Optional[str]) -> None:
+    def __init__(self, root_dir: str | None) -> None:
         # This is implemented in the same way that delay is implemented on
         # FileHandler
         self.root_dir = root_dir
         logging.Handler.__init__(self)
         self.stream = None
         self._builtin_open = open
+        self._pending_log_version = False
 
     # cloned from FileHandler in cpython
     def close(self) -> None:
@@ -1260,12 +1264,40 @@ def emit(self, record) -> None:
                     delete=False,
                 )
                 log.info("LazyTraceHandler: logging to %s", self.stream.name)
+                # Log tlparse path via inductor logger so it shows when
+                # TORCH_LOGS="inductor" is enabled
+                inductor_log = logging.getLogger("torch._inductor")
+                inductor_log.info("tlparse raw data: %s", self.stream.name)
+                self._pending_log_version = True
             else:
                 # We go poof, remove and no-op
                 trace_log.removeHandler(self)
                 return
         if self.stream:
             super().emit(record)
+            if self._pending_log_version:
+                self._pending_log_version = False
+                _log_torch_version()
+
+
+def _log_torch_version() -> None:
+    import torch
+    from torch._environment import is_fbcode
+    from torch._utils_internal import get_torch_source_version
+
+    version_info: dict[str, object] = {
+        "pytorch_version": torch.__version__,
+        "commit": get_torch_source_version(),
+        "oss": not is_fbcode(),
+    }
+
+    trace_structured(
+        "artifact",
+        metadata_fn=lambda: {"name": "torch_version", "encoding": "json"},
+        payload_fn=lambda: version_info,
+        suppress_context=True,
+        expect_trace_id=False,
+    )
 
 
 @functools.cache
@@ -1353,7 +1385,7 @@ def add_structured_logging_overhead(time_spent: float) -> None:
         structured_logging_overhead[key] += time_spent
 
 
-def get_structured_logging_overhead() -> Optional[float]:
+def get_structured_logging_overhead() -> float | None:
     key = None
     if (trace_id := torch._guards.CompileContext.current_trace_id()) is not None:
         frame_id = trace_id.compile_id.frame_id
@@ -1368,8 +1400,8 @@ def get_structured_logging_overhead() -> Optional[float]:
 def trace_structured_artifact(
     name: str,  # this will go in metadata
     encoding: str,
-    payload_fn: Callable[[], Optional[Union[str, object]]] = lambda: None,
-    compile_id: Optional[CompileId] = None,
+    payload_fn: Callable[[], str | object | None] = lambda: None,
+    compile_id: CompileId | None = None,
 ) -> None:
     trace_structured(
         "artifact",
@@ -1386,13 +1418,13 @@ def trace_structured(
     name: str,
     # NB: metadata expected to be dict so adding more info is forward compatible
     # Tuple[str, int] is a special case for string interning
-    metadata_fn: Callable[[], Union[dict[str, Any], tuple[str, int]]] = dict,
+    metadata_fn: Callable[[], dict[str, Any] | tuple[str, int]] = dict,
     *,
-    payload_fn: Callable[[], Optional[Union[str, object]]] = lambda: None,
+    payload_fn: Callable[[], str | object | None] = lambda: None,
     suppress_context: bool = False,
     expect_trace_id: bool = True,  # Whether or not we expect to have a current trace id
     record_logging_overhead: bool = True,  # Whether or not to record the time spent on structured logging
-    compile_id: Optional[CompileId] = None,  # Optional if unavailable in the trace
+    compile_id: CompileId | None = None,  # Optional if unavailable in the trace
 ) -> None:
     """
     metadata is an arbitrary JSON compatible struct, but it's expected to not be
@@ -1490,9 +1522,9 @@ def dtrace_structured(
     name: str,
     # NB: metadata expected to be dict so adding more info is forward compatible
     # Tuple[str, int] is a special case for string interning
-    metadata_fn: Callable[[], Union[dict[str, Any], tuple[str, int]]] = dict,
+    metadata_fn: Callable[[], dict[str, Any] | tuple[str, int]] = dict,
     *,
-    payload_fn: Callable[[], Optional[Union[str, object]]] = lambda: None,
+    payload_fn: Callable[[], str | object | None] = lambda: None,
     suppress_context: bool = False,
     expect_trace_id: bool = False,  # Whether or not we expect to have a current trace id
     record_logging_overhead: bool = True,  # Whether or not to record the time spent on structured logging
diff --git a/torch/_logging/scribe.py b/torch/_logging/scribe.py
index 2feb814d4a2c7..a457b549e4f1e 100644
--- a/torch/_logging/scribe.py
+++ b/torch/_logging/scribe.py
@@ -1,5 +1,5 @@
 from collections.abc import Callable
-from typing import TypeAlias, Union
+from typing import TypeAlias
 
 
 try:
@@ -7,9 +7,9 @@
         make_scribe_logger,
     )
 except ImportError:
-    TAtom: TypeAlias = Union[int, float, bool, str]
-    TField: TypeAlias = Union[TAtom, list[TAtom]]
-    TLazyField: TypeAlias = Union[TField, Callable[[], TField]]
+    TAtom: TypeAlias = int | float | bool | str
+    TField: TypeAlias = TAtom | list[TAtom]
+    TLazyField: TypeAlias = TField | Callable[[], TField]
 
     def make_scribe_logger(name: str, thrift_src: str) -> Callable[..., None]:
         def inner(**kwargs: TLazyField) -> None:
diff --git a/torch/_logging/structured.py b/torch/_logging/structured.py
index e6dd36a6c6968..c31c83ab787bf 100644
--- a/torch/_logging/structured.py
+++ b/torch/_logging/structured.py
@@ -6,7 +6,7 @@
 import os
 import traceback
 from collections.abc import Sequence
-from typing import Any, Optional
+from typing import Any
 
 import torch._logging._internal
 
@@ -17,7 +17,7 @@
 DUMPED_FILES: set[str] = set()
 
 
-def intern_string(s: Optional[str]) -> int:
+def intern_string(s: str | None) -> int:
     if s is None:
         return -1
 
diff --git a/torch/_meta_registrations.py b/torch/_meta_registrations.py
index 504ea6b705f08..d53858bf40faf 100644
--- a/torch/_meta_registrations.py
+++ b/torch/_meta_registrations.py
@@ -255,14 +255,9 @@ def _exec_fft(out, self, out_sizes, dim, *, forward):
     for d in dim:
         is_transformed_dim[d] = True
 
-    # std::partition
-    left, right = [], []
-    for d in dim_permute:
-        if not is_transformed_dim[d]:
-            left.append(d)
-        else:
-            right.append(d)
-    dim_permute = left + right
+    # std::partition + std::copy(dim.begin(), dim.end(), batch_end)
+    left = [d for d in dim_permute if not is_transformed_dim[d]]
+    dim_permute = left + list(dim)
     batch_end = len(left)
 
     self_strides = self.stride()
@@ -300,8 +295,9 @@ def _exec_fft(out, self, out_sizes, dim, *, forward):
 def _sort_dims(self: Tensor, dim: list[int], exclude_last: bool = False):
     sorted_dims = list(dim)
     self_strides = self.stride()
-    sorted_dims[: len(sorted_dims) - int(exclude_last)].sort(
-        key=lambda i: self_strides[i]
+    end = len(sorted_dims) - int(exclude_last)
+    sorted_dims[:end] = sorted(
+        sorted_dims[:end], key=lambda i: self_strides[i], reverse=True
     )
     return sorted_dims
 
@@ -315,9 +311,36 @@ def meta_fft_c2c(self, dim, normalization, forward):
     if not dim:
         return self.clone()
 
-    sorted_dims = _sort_dims(self, dim)
-    out = self.new_empty(self.size())
-    return _exec_fft(out, self, self.size(), sorted_dims, forward=forward)
+    if device_hint(self) == "cpu" and not torch.backends.mkl.is_available():
+        return self.new_empty(self.size())
+
+    out_sizes = self.size()
+    output = self.new_empty(out_sizes)
+    if device_hint(self) != "cuda":
+        sorted_dims = _sort_dims(self, dim)
+        return _exec_fft(output, self, out_sizes, sorted_dims, forward=forward)
+
+    # Match _fft_c2c_cufft, which re-sorts the remaining dimensions after each
+    # staged transform because _exec_fft restrides the output in place.
+    sorted_dims = list(dim)
+    working_tensor = self
+    while True:
+        strides = working_tensor.stride()
+        sorted_dims.sort(key=lambda i: strides[i], reverse=True)
+        max_dims = min(cufft_max_ndim, len(sorted_dims))
+        last_dims = sorted_dims[len(sorted_dims) - max_dims :]
+
+        _exec_fft(output, working_tensor, out_sizes, last_dims, forward=forward)
+        sorted_dims = sorted_dims[: len(sorted_dims) - max_dims]
+
+        if not sorted_dims:
+            return output
+
+        if working_tensor is self:
+            working_tensor = output
+            output = self.new_empty(out_sizes)
+        else:
+            output, working_tensor = working_tensor, output
 
 
 cufft_max_ndim = 3
@@ -385,6 +408,14 @@ def meta_fft_r2c(self, dim, normalization, onesided):
 
         return output
 
+    elif torch.backends.mkl.is_available():
+        # _fft_r2c_mkl in aten/src/ATen/native/mkl/SpectralOps.cpp
+        sorted_dims = _sort_dims(self, dim, exclude_last=True)
+        output = self.new_empty(
+            out_sizes, dtype=utils.corresponding_complex_dtype(self.dtype)
+        )
+        return _exec_fft(output, self, out_sizes, sorted_dims, forward=True)
+
     else:
         return self.new_empty(
             out_sizes, dtype=utils.corresponding_complex_dtype(self.dtype)
@@ -460,6 +491,89 @@ def meta_rand_default(size, *, dtype=None, layout=None, device=None, pin_memory=
     )
 
 
+@register_meta(aten._philox_key_split.default)
+def meta_philox_key_split(key, num_splits):
+    torch._check(
+        key.dim() >= 1 and key.shape[-1] == 2,
+        lambda: f"_philox_key_split: key must have shape (*batch, 2), got shape {key.shape}",
+    )
+    torch._check(
+        key.dtype == torch.uint64,
+        lambda: f"_philox_key_split: key must have dtype uint64, got {key.dtype}",
+    )
+    torch._check(
+        num_splits > 0,
+        lambda: f"_philox_key_split: num_splits must be positive, got {num_splits}",
+    )
+    batch_sizes = key.shape[:-1]
+    return key.new_empty((num_splits, *batch_sizes, 2))
+
+
+@register_meta(aten._philox_key_fold_in.default)
+def meta_philox_key_fold_in(key, data):
+    torch._check(
+        key.dim() >= 1 and key.shape[-1] == 2,
+        lambda: f"_philox_key_fold_in: key must have shape (*batch, 2), got shape {key.shape}",
+    )
+    torch._check(
+        key.dtype == torch.uint64,
+        lambda: f"_philox_key_fold_in: key must have dtype uint64, got {key.dtype}",
+    )
+    return torch.empty_like(key)
+
+
+def _check_philox_distribution_args(op_name, self, key):
+    torch._check(
+        self.dtype.is_floating_point,
+        lambda: f"{op_name}: self must be a floating point tensor, got {self.dtype}",
+    )
+    torch._check(
+        key.dtype == torch.uint64,
+        lambda: f"{op_name}: key must have dtype uint64, got {key.dtype}",
+    )
+    torch._check(
+        self.device == key.device,
+        lambda: (
+            f"{op_name}: self and key must be on the same device, "
+            f"got {self.device} and {key.device}"
+        ),
+    )
+    torch._check(
+        key.dim() >= 1 and key.shape[-1] == 2,
+        lambda: (
+            f"{op_name}: key must have shape (2,) or (*batch, 2), got shape {key.shape}"
+        ),
+    )
+    if key.dim() > 1:
+        torch._check(
+            key.dim() == self.dim() + 1,
+            lambda: (
+                f"{op_name}: batched key must have ndim == output ndim + 1, "
+                f"got key shape {key.shape} with output shape {self.shape}"
+            ),
+        )
+        key_batch = key.shape[: self.dim()]
+        torch._check(
+            all(ks == 1 or ks == ss for ks, ss in zip(key_batch, self.shape)),
+            lambda: (
+                f"{op_name}: key batch shape {list(key_batch)} "
+                f"is not broadcastable with output shape {self.shape}"
+            ),
+        )
+
+
+@register_meta(aten._philox_normal_.default)
+def meta_philox_normal_(self, key, mean=0.0, std=1.0):
+    _check_philox_distribution_args("_philox_normal_", self, key)
+    return self
+
+
+@register_meta(aten._philox_uniform_.default)
+def meta_philox_uniform_(self, key, low=0.0, high=1.0):
+    _check_philox_distribution_args("_philox_uniform_", self, key)
+    return self
+
+
 @register_meta([aten._fft_c2r.default, aten._fft_c2r.out])
 @out_wrapper()
 def meta_fft_c2r(self: Tensor, dim: list[int], normalization: int, lastdim: int):
@@ -483,12 +597,13 @@ def meta_fft_c2r(self: Tensor, dim: list[int], normalization: int, lastdim: int)
         else:
             # First complete any C2C transforms
             if len(dim) > 1:
-                temp = meta_fft_c2c(self, dim[:-1], 0, lastdim)  # fft_norm_mode::none
+                temp = meta_fft_c2c(self, dim[:-1], 0, forward=False)
             else:
                 temp = self.clone(memory_format=torch.contiguous_format)
             return _exec_fft(output, temp, out_sizes, [dim[-1]], forward=False)
 
-    else:
+    elif torch.backends.mkl.is_available():
+        # _fft_c2r_mkl in aten/src/ATen/native/mkl/SpectralOps.cpp
         input = self
         if len(dim) > 1:
             c2c_dims = dim[:-1]
@@ -500,6 +615,11 @@ def meta_fft_c2r(self: Tensor, dim: list[int], normalization: int, lastdim: int)
         out = self.new_empty(out_sizes, dtype=toRealValueType(self.dtype))
         return _exec_fft(out, input, out_sizes, dim, forward=False)
 
+    else:
+        out_sizes = list(self.size())
+        out_sizes[dim[-1]] = lastdim
+        return self.new_empty(out_sizes, dtype=toRealValueType(self.dtype))
+
 
 @register_meta(aten.copy_.default)
 def meta_copy_(self, src, non_blocking=False):
@@ -521,8 +641,13 @@ def meta_copy_(self, src, non_blocking=False):
 
     if isinstance(src, Tensor):
         intermediate = src.to(self, non_blocking)
-        if self.size() != intermediate.size():
-            aten.expand_copy.default(intermediate, self.size())
+        # Validate broadcast compatibility. We call the _refs expand
+        # directly rather than aten.expand_copy.default because the
+        # aten dispatch path hits C++ that cannot handle unbacked
+        # symints. The refs version performs the same checks using
+        # sym_or(x == 1, requested_length == x) which gracefully
+        # handles unbacked symints.
+        torch._refs.expand(intermediate, self.size())
     return self
 
 
@@ -582,10 +707,11 @@ def meta_sparse_structured_linear(
             raise AssertionError(
                 f"out_dtype is only supported for i8i8->i32 linear operator, got input.dtype={input.dtype}, out_dtype={out_dtype}"
             )
-    output = input.new_empty(
+    output = input.new_empty_strided(
         output_sizes,
+        transposed_strides,
         dtype=input.dtype if out_dtype is None else out_dtype,
-    ).as_strided(output_sizes, transposed_strides)
+    )
 
     return output
 
@@ -876,7 +1002,7 @@ def make_dep_token(
     pin_memory=None,
     memory_format=None,
 ):
-    return torch.empty(0, device="meta")
+    return torch.empty((), device="meta")
 
 
 @register_meta(aten.sym_constrain_range.default)
@@ -1529,7 +1655,7 @@ def _linalg_svd_meta(
     batch_dims = list(A.shape[:-2])
     m = A.shape[-2]
     n = A.shape[-1]
-    k = min(m, n)
+    k = torch.sym_min(m, n)
 
     if compute_uv:
         U_shape = batch_dims + [m, m if full_matrices else k]
@@ -2420,8 +2546,10 @@ def _formula_transposed(ln: int, p: int, d: int, k: int, s: int, op: int) -> int
         out_channels = groups * weight.shape[1]
     else:
         out_channels = weight.shape[0]
-        if weight.shape[1] * groups != input_tensor.shape[1]:
-            raise RuntimeError("Invalid channel dimensions")
+        torch._check(
+            weight.shape[1] * groups == input_tensor.shape[1],
+            lambda: "Invalid channel dimensions",
+        )
 
     ret_shape = [input_tensor.shape[0], out_channels]
     if isinstance(stride, IntLike):
@@ -2566,11 +2694,18 @@ def meta_conv(
         output_padding if is_transposed else None,
     )
 
+    from torch.fx.experimental.symbolic_shapes import guard_or_false
+
     input_channels_dim = 1
     output_channels_dim = 1
-    if input_tensor.size(input_channels_dim) == 0:
+    if guard_or_false(input_tensor.size(input_channels_dim) == 0):
         shape_out[output_channels_dim] = 0
 
+    # Memory format is left as contiguous: meta tensors have no device info,
+    # so _select_conv_backend returns Overrideable and the correct format
+    # cannot be determined here.  The FakeTensor path (torch.compile, export)
+    # intercepts via a register_op_impl in fake_impls.py before reaching this
+    # kernel and uses FakeTensor.fake_device for an accurate answer.
     out = input_tensor.new_empty(shape_out)
     return out
 
@@ -2863,6 +2998,13 @@ def check_dim_size(tensor, dim, dim_size, size):
     )
 
 
+@register_meta(aten.quantize_per_tensor)
+def meta_quantize_per_tensor(
+    input: torch.Tensor, scale: float, zero_point: int, dtype: torch.dtype
+) -> torch.Tensor:
+    return torch.empty_like(input)
+
+
 @register_meta(aten.avg_pool2d.default)
 def meta_avg_pool2d(
     input,
@@ -3445,8 +3587,8 @@ def meta_complex(real, imag):
 @register_meta([aten.nonzero_static.default, aten.nonzero_static.out])
 @out_wrapper()
 def nonzero_static(self, *, size, fill_value: int = -1):
-    # The impl of xpu nonzero_static is different with cuda but aligned with cpu
-    if device_hint(self) in ("cpu", "xpu"):
+    # The impl of nonzero_static on xpu and mps differs from cuda but aligned with cpu
+    if device_hint(self) in ("cpu", "mps", "xpu"):
         return self.new_empty((size, self.dim()), dtype=torch.long)
     else:
         return torch.empty_strided(
@@ -3629,18 +3771,13 @@ def meta_convolution_backward(
     backend_grad_weight = None
     backend_grad_bias = None
 
-    # Backend layout expectation: GPU backends (CUDA via cudnn_conv_suggest_memory_format,
-    # MPS via mps_conv_use_channels_last) return channels_last outputs when either input
-    # tensor is channels_last. This must be matched here to avoid stride assertion failures
-    # in inductor when the predicted strides don't match actual backend output strides.
+    # All GPU backends compute output memory format via
+    # determine_backend_memory_format(input, weight, backend) — which calls
+    # cudnn_conv_suggest_memory_format(input, weight), mps_conv_use_channels_last(input, weight),
+    # etc. The format depends only on input and weight, NOT on grad_output.
+    # Both grad_input and grad_weight use this same backend_memory_format.
     # See: https://github.com/pytorch/pytorch/issues/171622
-    #
-    # Memory format inference rules (matching backend behavior):
-    #   - grad_input format: derived from grad_output and weight
-    #   - grad_weight format: derived from input and grad_output
     def _conv_memory_format(t1, t2):
-        # Match the logic in cudnn_conv_suggest_memory_format and mps_conv_use_channels_last:
-        # Use channels_last if either tensor suggests it
         fmt1 = suggest_memory_format(t1)
         fmt2 = suggest_memory_format(t2)
         if fmt1 == torch.channels_last or fmt2 == torch.channels_last:
@@ -3649,13 +3786,12 @@ def _conv_memory_format(t1, t2):
             return torch.channels_last_3d
         return torch.contiguous_format
 
+    memory_format = _conv_memory_format(input_, weight_)
     if output_mask[0]:
-        memory_format = _conv_memory_format(grad_output_, weight_)
         backend_grad_input = grad_output_.new_empty(input_.size()).to(
             memory_format=memory_format
         )
     if output_mask[1]:
-        memory_format = _conv_memory_format(input_, grad_output_)
         backend_grad_weight = grad_output_.new_empty(weight_.size()).to(
             memory_format=memory_format
         )
@@ -3765,8 +3901,8 @@ def meta__int_mm(a, b):
     torch._check(a.dim() == 2, lambda: "a must be a 2D tensor")
     torch._check(b.dim() == 2, lambda: "b must be a 2D tensor")
     torch._check(
-        a.dtype is torch.int8,
-        lambda: f"expected self to be int8, got {a.dtype}",
+        a.dtype in [torch.int8, torch.uint8],
+        lambda: f"expected self to be int8 or uint8, got {a.dtype}",
     )
     torch._check(
         b.dtype is torch.int8,
@@ -4540,6 +4676,10 @@ def common_meta_baddbmm_bmm(batch1, batch2, is_bmm, self_baddbmm=None, out_dtype
 
     torch._check(batch1.dim() == 3, lambda: "batch1 must be a 3D tensor")
     torch._check(batch2.dim() == 3, lambda: "batch2 must be a 3D tensor")
+    torch._check(
+        batch1.dtype == batch2.dtype,
+        lambda: f"expected scalar type {batch1.dtype} but found {batch2.dtype}",
+    )
 
     batch1_sizes = batch1.size()
     batch2_sizes = batch2.size()
@@ -4578,12 +4718,14 @@ def common_meta_baddbmm_bmm(batch1, batch2, is_bmm, self_baddbmm=None, out_dtype
     return output
 
 
-@register_meta(aten.bmm.default)
+@register_meta([aten.bmm.default, aten.bmm.out])
+@out_wrapper(exact_dtype=True)
 def meta_bmm(self, mat2):
     return common_meta_baddbmm_bmm(self, mat2, True)
 
 
-@register_meta(aten.bmm.dtype)
+@register_meta([aten.bmm.dtype, aten.bmm.dtype_out])
+@out_wrapper(exact_dtype=True)
 def meta_bmm_dtype(self, mat2, out_dtype):
     return common_meta_baddbmm_bmm(self, mat2, True, out_dtype=out_dtype)
 
@@ -4986,7 +5128,8 @@ def unpack(name, val):
     return nInputPlane, outputHeight, outputWidth
 
 
-@register_meta(aten.max_pool2d_with_indices_backward.default)
+@register_meta(aten.max_pool2d_with_indices_backward)
+@out_wrapper("grad_input")
 def meta_max_pool2d_with_indices_backward(
     grad_output,
     self,
@@ -5579,11 +5722,22 @@ def meta_zeros(
 
 @register_meta(aten.select_scatter.default)
 def meta_select_scatter(self, src, dim, index):
-    return utils.clone_preserve_strides(self)
+    return _scatter_meta_output(self)
 
 
 @register_meta(aten.slice_scatter.default)
 def meta_slice_scatter(self, src, dim=0, start=None, end=None, step=1):
+    return _scatter_meta_output(self)
+
+
+def _scatter_meta_output(self):
+    from torch.fx.experimental.symbolic_shapes import free_unbacked_symbols
+
+    # Match clone_preserve_strides() in aten/native/TensorShape.cpp: overlapping
+    # bases cannot preserve their logical strides because the scatter writes would
+    # alias, so eager falls back to clone().
+    if not free_unbacked_symbols(self) and torch._debug_has_internal_overlap(self) == 1:
+        return self.clone()
     return utils.clone_preserve_strides(self)
 
 
@@ -5945,14 +6099,23 @@ def meta__scaled_dot_product_fused_attention_overrideable(
     return_debug_mask: bool = False,
     scale: float | None = None,
 ):
-    B = query.size(0)
-    H_Q = query.size(1)
-    S_Q = query.size(2)
-    S_KV = key.size(2)
+    # Explicitly handle 3D (H, S, D) and 4D (B, H, S, D) inputs,
+    # matching the C++ runtime in aten_mtia_ops.cpp.
+    B, H_Q, S_Q = 0, 0, 0
+    if query.dim() == 4:
+        B, H_Q, S_Q, _ = query.size()
+    elif query.dim() == 3:
+        H_Q, S_Q, _ = query.size()
+        B = 1
+    else:
+        raise RuntimeError("query must be 3D or 4D")
+    S_KV = key.size(-2)
     D_V = value.size(-1)
 
-    res_shape = (B, H_Q, S_Q, D_V)
-    res = alloc_with_matching_layout(query, res_shape)
+    # Preserve input dimensionality for the output shape
+    out_shape = list(query.shape)
+    out_shape[-1] = D_V
+    res = alloc_with_matching_layout(query, tuple(out_shape))
 
     logsum_exp = torch.empty(
         (B, H_Q, S_Q),
@@ -5977,6 +6140,34 @@ def meta__scaled_dot_product_fused_attention_overrideable(
     )
 
 
+@register_meta(aten._scaled_dot_product_fused_attention_overrideable_backward)
+def meta__scaled_dot_product_fused_attention_overrideable_backward(
+    grad_out: Tensor,
+    query: Tensor,
+    key: Tensor,
+    value: Tensor,
+    attn_bias: Tensor,
+    grad_input_mask: list[bool],
+    out: Tensor,
+    logsumexp: Tensor,
+    cum_seq_q: Tensor,
+    cum_seq_k: Tensor,
+    max_q: int,
+    max_k: int,
+    dropout_p: float,
+    is_causal: bool,
+    philox_seed: Tensor,
+    philox_offset: Tensor,
+    *,
+    scale: float | None = None,
+):
+    grad_q = torch.empty_like(query)
+    grad_k = torch.empty_like(key)
+    grad_v = torch.empty_like(value)
+    grad_attn_bias = torch.empty_like(attn_bias) if attn_bias is not None else None
+    return grad_q, grad_k, grad_v, grad_attn_bias
+
+
 @register_meta(
     [
         aten._scaled_dot_product_flash_attention_backward,
@@ -6090,6 +6281,7 @@ def meta__scaled_dot_product_attention_math_for_mps(
     is_causal: bool = False,
     dropout_mask: Tensor | None = None,
     scale: float | None = None,
+    enable_gqa: bool = False,
 ) -> tuple[Tensor, Tensor]:
     def ensure_4d(x):
         if x.dim() == 3:
@@ -6106,33 +6298,26 @@ def ensure_4d(x):
     k_, _ = ensure_4d(key)
     v_, _ = ensure_4d(value)
 
-    batch_size, num_head, q_size, head_size = q_.shape
-    _, k_size, max_seq_length, _ = k_.shape
-
-    def sdpa_vector_fast_mps():
-        out = q_.new_empty(q_.shape)
-        if unsqueezed:
-            out = out.view_as(query)
+    batch_size, num_head, q_size, _ = q_.shape
+    _, _, max_seq_length, value_head_size = v_.shape
 
+    def sdpa_general_mps():
+        out = q_.new_empty((batch_size, num_head, q_size, value_head_size))
         attn = q_.new_empty((batch_size, num_head, q_size, max_seq_length))
         if unsqueezed:
             if query.dim() == 3:
+                out = out.squeeze(0)
                 attn = attn.squeeze(0)
             else:
-                shape = list(query.shape[:-3]) + attn.shape[1:4]
-                attn = attn.view(shape)
+                out_shape = list(query.shape[:-3]) + list(out.shape[1:4])
+                attn_shape = list(query.shape[:-3]) + list(attn.shape[1:4])
+                out = out.view(out_shape)
+                attn = attn.view(attn_shape)
         return out, attn
 
-    def sdpa_vector_2pass_mps():
-        blocks = 32
-        out = q_.new_empty(q_.shape)
-        intermediate = q_.new_empty((batch_size, num_head, q_size, blocks, head_size))
-        return out, intermediate
-
-    if (max_seq_length >= 1024) or (k_size < q_size and max_seq_length >= 4096):
-        return sdpa_vector_2pass_mps()
-    else:
-        return sdpa_vector_fast_mps()
+    # sdpa_vector_2pass_mps and sdpa_vector_fast_mps are intentionally left out.
+    # See https://github.com/pytorch/pytorch/issues/177603 for additional context.
+    return sdpa_general_mps()
 
 
 @register_meta([aten._scaled_dot_product_efficient_attention])
@@ -6290,6 +6475,8 @@ def meta__flash_attention_forward(
     window_size_right: int | None = None,
     seqused_k: Tensor | None = None,
     alibi_slopes: Tensor | None = None,
+    block_table: Tensor | None = None,
+    num_splits: int | None = None,
 ):
     # NB: there are two underlying paths:
     # 1. normal dense path; expect 4D inputs of shape (batch_size, seqlen, num_heads, head_dim)
@@ -6349,6 +6536,48 @@ def meta__flash_attention_forward(
     )
 
 
+@register_meta([aten._flash_attention_forward_no_dropout_inplace.default])
+def meta__flash_attention_forward_no_dropout_inplace(
+    out: Tensor,
+    query: Tensor,
+    key: Tensor,
+    value: Tensor,
+    cum_seq_q: Tensor | None,
+    cum_seq_k: Tensor | None,
+    max_q: int,
+    max_k: int,
+    dropout_p: float,
+    is_causal: bool,
+    return_debug_mask: bool,
+    scale: float | None = None,
+    window_size_left: int | None = None,
+    window_size_right: int | None = None,
+    seqused_k: Tensor | None = None,
+    alibi_slopes: Tensor | None = None,
+    block_table: Tensor | None = None,
+    num_splits: int | None = None,
+):
+    _, logsumexp, _, _, _ = meta__flash_attention_forward(
+        query,
+        key,
+        value,
+        cum_seq_q,
+        cum_seq_k,
+        max_q,
+        max_k,
+        dropout_p,
+        is_causal,
+        return_debug_mask,
+        scale,
+        window_size_left,
+        window_size_right,
+        seqused_k,
+        alibi_slopes,
+        block_table,
+    )
+    return logsumexp
+
+
 @register_meta([aten._flash_attention_forward.quantized])
 def meta__flash_attention_forward_quantized(
     query: Tensor,
@@ -6566,7 +6795,11 @@ def is_fp8_or_fp4_type(dtype):
         lambda: f"Expected both inputs to be fp8 or fp4 types but got self.dtype={self.dtype} and mat2.dtype={mat2.dtype}",
     )
 
-    if device_hint(self) == "cuda" or device_hint(self) == "xpu":
+    if (
+        device_hint(self) == "cuda"
+        or device_hint(self) == "xpu"
+        or device_hint(self) == "cpu"
+    ):
 
         def is_row_major(stride):
             return stride[0] > stride[1] and stride[1] == 1
@@ -6577,22 +6810,23 @@ def is_col_major(stride):
         def has_zero_dim(tensor_2d):
             return tensor_2d.size(0) == 0 or tensor_2d.size(1) == 0
 
-        torch._check(
-            is_row_major(self.stride()) or has_zero_dim(self),
-            lambda: f"self must be row_major, got stride {self.stride()}",
-        )
-        torch._check(
-            is_col_major(mat2.stride()) or has_zero_dim(mat2),
-            lambda: f"mat2 must be col_major, got stride {mat2.stride()}",
-        )
-        torch._check(
-            self.size(1) % 16 == 0,
-            lambda: f"Expected self.size(1) to be divisible by 16, but got self.size(1)={self.size(1)}",
-        )
-        torch._check(
-            mat2.size(0) % 16 == 0 and mat2.size(1) % 16 == 0,
-            lambda: f"Expected both dimensions of mat2 to be divisible by 16 but got {mat2.shape}",
-        )
+        if device_hint(self) != "cpu":
+            torch._check(
+                is_row_major(self.stride()) or has_zero_dim(self),
+                lambda: f"self must be row_major, got stride {self.stride()}",
+            )
+            torch._check(
+                is_col_major(mat2.stride()) or has_zero_dim(mat2),
+                lambda: f"mat2 must be col_major, got stride {mat2.stride()}",
+            )
+            torch._check(
+                self.size(1) % 16 == 0,
+                lambda: f"Expected self.size(1) to be divisible by 16, but got self.size(1)={self.size(1)}",
+            )
+            torch._check(
+                mat2.size(0) % 16 == 0 and mat2.size(1) % 16 == 0,
+                lambda: f"Expected both dimensions of mat2 to be divisible by 16 but got {mat2.shape}",
+            )
 
         # determine scaling type and check input dimensions (refer to Blas.cpp op)
 
@@ -6627,6 +6861,8 @@ def has_zero_dim(tensor_2d):
                 _k = _k * 2
             else:
                 block_size_k = 32
+                if self.dtype == torch.float4_e2m1fn_x2:
+                    _k = _k * 2
 
             block_size_mn = 128
 
@@ -6805,7 +7041,11 @@ def is_fp4_type(dtype):
             SwizzleType.NO_SWIZZLE,
         ]
 
-    if device_hint(self) == "cuda" or device_hint(self) == "xpu":
+    if (
+        device_hint(self) == "cuda"
+        or device_hint(self) == "xpu"
+        or device_hint(self) == "cpu"
+    ):
 
         def is_row_major(stride):
             return stride[0] > stride[1] and stride[1] == 1
@@ -6816,22 +7056,23 @@ def is_col_major(stride):
         def has_zero_dim(tensor_2d):
             return tensor_2d.size(0) == 0 or tensor_2d.size(1) == 0
 
-        torch._check(
-            is_row_major(self.stride()) or has_zero_dim(self),
-            lambda: f"self must be row_major, got stride {self.stride()}",
-        )
-        torch._check(
-            is_col_major(mat2.stride()) or has_zero_dim(mat2),
-            lambda: f"mat2 must be col_major, got stride {mat2.stride()}",
-        )
-        torch._check(
-            self.size(1) % 16 == 0,
-            lambda: f"Expected self.size(1) to be divisible by 16, but got self.size(1)={self.size(1)}",
-        )
-        torch._check(
-            mat2.size(0) % 16 == 0 and mat2.size(1) % 16 == 0,
-            lambda: f"Expected both dimensions of mat2 to be divisible by 16 but got {mat2.shape}",
-        )
+        if device_hint(self) != "cpu":
+            torch._check(
+                is_row_major(self.stride()) or has_zero_dim(self),
+                lambda: f"self must be row_major, got stride {self.stride()}",
+            )
+            torch._check(
+                is_col_major(mat2.stride()) or has_zero_dim(mat2),
+                lambda: f"mat2 must be col_major, got stride {mat2.stride()}",
+            )
+            torch._check(
+                self.size(1) % 16 == 0,
+                lambda: f"Expected self.size(1) to be divisible by 16, but got self.size(1)={self.size(1)}",
+            )
+            torch._check(
+                mat2.size(0) % 16 == 0 and mat2.size(1) % 16 == 0,
+                lambda: f"Expected both dimensions of mat2 to be divisible by 16 but got {mat2.shape}",
+            )
 
         def is_tensorwise(recipe_a: list[ScalingType], recipe_b: list[ScalingType]):
             return (
@@ -7767,7 +8008,11 @@ def meta_histc(input, bins=100, min=0, max=0):
 
 
 @register_meta(
-    [aten._upsample_bilinear2d_aa.default, aten._upsample_bicubic2d_aa.default]
+    [
+        aten._upsample_bilinear2d_aa.default,
+        aten._upsample_bicubic2d_aa.default,
+        aten._upsample_lanczos2d_aa.default,
+    ]
 )
 def meta_upsample_bimode2d_aa(
     input,
@@ -7788,7 +8033,12 @@ def meta_upsample_bimode2d_aa(
     )
 
 
-@register_meta([aten._upsample_bilinear2d_aa_backward.default])
+@register_meta(
+    [
+        aten._upsample_bilinear2d_aa_backward.default,
+        aten._upsample_lanczos2d_aa_backward.default,
+    ]
+)
 def meta_upsample_bimode2d_aa_backward(
     grad_output,
     output_size,
diff --git a/torch/_native/README.md b/torch/_native/README.md
new file mode 100644
index 0000000000000..f746c94848265
--- /dev/null
+++ b/torch/_native/README.md
@@ -0,0 +1,236 @@
+# Native Ops (and DSLs)
+
+The `torch._native` directory provides a place for PyTorch native ops written in python and DSLs, along with utilities to help facilitate this.
+
+# Creating & Registering a native op
+
+**All native ops must be registered to the dispatcher**
+
+When writing native ops, they are required to interact meaningfully with torch's dispatcher, and thus must be registered correctly.
+
+As a further clarification, ops cannot be labelled as `CompositeImplicitAutograd` in `native_functions.yaml`, as-in the op must have an explicit autograd function registered, or at minimum an explicit implementation registered for the same backend as being overridden/added.
+
+Further, we expect the overriding function to take as the first argument `dispatch_keys` (of type `torch.DispatchKeySet`), which is necessary for fallback implementations.
+
+## A Note on Imports
+
+All registrations will happen at the end of `import torch`. It is expected at that point that **no DSL runtime library is loaded by registration code** - this means that the runtime(s) must only be imported lazily. We can still check the presence of a module, and get it's version without importing, but special care must be taken when writing op kernels to not import DSLs too early. An illustrative example is below, using `triton`:
+
+First, we're going to write the registration function, and a top-level call, being very careful to not pull in the `triton` package early:
+
+```
+# torch/_native/ops/test_op/triton_impl.py
+
+# NOTE: no triton import in the file
+
+from ... import triton_utils as tu
+
+def calling_fn(dispatch_key, ...):
+    # Lazily import the kernels (and triton) on first call
+    from .triton_kernels import outer_fn
+    torch.library.wrap_triton(outer_fn)(...)
+
+def register_to_dispatch():
+    tu.register_op_override(
+        "aten",
+        "_scaled_mm_v2",
+        "CUDA",
+        calling_fn,
+    )
+
+```
+
+Note the lazy import of kernels inside `calling_fn` - this function isn't called during registration (it just needs to be defined), and on first call it'll import the `triton_kernels` module, with all its dependencies. This kernels module can import anything it wants internally, our restriction on no DSL imports during registration has been met. Quick example of `triton_kernels.py` below:
+
+```
+# torch/_native/ops/test_op/triton_kernels.py
+import triton
+
+@triton.jit
+def inner_fn(...) -> ...:
+    # depends internally on triton, triton.language
+    pass
+
+@triton.jit
+def outer_fn(...) -> ...:
+    # depends internally on triton, triton.language
+    inner_fn(...)
+```
+
+
+## Registering Implementations to Existing Operators
+
+There are 2 options when interacting with an existing operator:
+1. Replace the operator for **all** cases with a new implementation
+2. Replace **some subset of functionality** of a given operator with a new implementation, falling-back to the original implementation otherwise.
+
+Both cases are very similar, with 2) only requiring an extra step to obtain the original implementation, and logic to determine which implementation should be run for a given case.
+
+### Replacing an Operator
+
+This follows a simple and standard path, with a good example being the implementation of [FlashAttention v4 (FAv4)](https://github.com/pytorch/pytorch/blob/1f66f34cda5b5ad02d231b90fa0c0de2cb4e02d1/torch/nn/attention/_fa4.py#L67) in torch.
+
+The following example replaces the implementation of `aten._scaled_grouped_mm_v2` on `CUDA` devices:
+
+```
+from ... import cutedsl_utils as cu
+
+def my_impl(dispatch_key, ...) -> ...:
+    """
+    Replacement implementation
+    """
+    pass
+
+# Override the symbol `aten._scaled_grouped_mm_v2` in this example with the implementation in `my_impl`,
+# noting the function signatures must match.
+# Replacing an operator completely (no fallback) requires passing the `unconditional_override=True` flag.
+def register_kernel_override():
+    tu.register_op_override(
+        "aten",
+        "_scaled_grouped_mm_v2",
+        "CUDA",
+        my_impl,
+        unconditional_override=True
+    )
+```
+
+### Replacing a Subset of Calls
+
+This time we only want to override the behavior of a subset of `aten._scaled_grouped_mm_v2` calls, and choose whether to invoke our implementation or the original depending on some input arguments. Note that the core of the example -- creating a `torch.library.Library`, and registering our function using `lib.impl(...)` are the same as in [Replacing an Operator](#Replacing-an-Operator).
+
+```
+from ... import cutedsl_utils as cu
+
+def my_impl(...) -> ...:
+    """
+    Replacement implementation - laxy import actual implementation and run it
+    """
+    from .my_impl_kernel import my_kernel
+    return my_kernel(...)
+
+
+# Note the dispatch_keys argument here - this must be passed as the first argument
+# to the fallback kernel.
+# Also note that we need the `fallback_kernel` argument to be specified, as
+# we this function gets called every operator invocation - the "fallback" is
+# actually the currently registered function, which is... this one.
+def enable_my_impl(dispatch_keys, arg1, arg2, *args, fallback_kernel, **kwargs):
+    # determine if we want to call our implementation
+    if arg1 == ... and arg2 == ...:
+        return my_impl(arg1, arg2, *args, **kwargs)
+    else:
+        # Call the fallback
+        return fallback_kernel(dispatch_keys,
+                               arg1, arg2, *args, **kwargs)
+
+# Override the symbol `aten._scaled_grouped_mm_v2` in this example with the implementation in `my_impl`,
+# only when the check-method `enable_my_impl` returns `True`
+def register_kernel_override():
+
+    # Get the original implementation for fallback purposes
+    fallback_kernel = torch.library.get_kernel("aten::_scaled_grouped_mm_v2", "CUDA")
+
+    # partially-specialize our function so that we're not grabbing the
+    # fallback every invocation.
+    fn = functools.partial(enable_my_impl, fallback_kernel=fallback_kernel)
+
+    # Same as before
+    cu.register_op_override(
+        "aten",
+        "_scaled_grouped_mm_v2",
+        "CUDA",
+        fn,
+    )
+```
+
+## Registering a New Operator
+
+We currently don't have a use for this functionality, please come talk to us if you want it!
+
+# Adding a new DSL
+
+Adding a new DSL is as simple as adding a single helper utils file, then writing your op.
+
+Some universal utilities are provided in `common_utils.py`:
+* `check_native_jit_disabled() -> bool` : returns `True` if native DSL ops have been globally disabled.
+* `_available_version(package: str) -> tuple[int,int,int] | None` : Gets the installed version of `package` without importing it
+* `_unavailable_reasons(deps: list[tuple[str,str]]) -> str | None` : For a list of `(package_name, module_name)` pairs, check (without importing) if the module is available, returning a string describing the mitigation step if it isn't.
+
+## DSL utils file spec
+
+A DSL utils file, named `$dsl_utils.py` (i.e. `cutedsl_utils.py` for `$dsl=cutedsl`) requires three methods to be implemented.
+
+`runtime_available() -> bool` : tell the user if the runtime is available - note that this needs to be available during init, and must be fork-safe. Packages should also not be imported at this time - rely on `importlib.util.find_spec(package_name)` or similar to get the necessary information without importing.
+
+`runtime_version() -> tuple[int, int, int]` : return the `(major, minor, update)` version of the installed package.
+
+
+```
+register_op_override(
+    lib_symbol: str,
+    op_symbol: str,
+    dispatch_key: str,
+    implementation_fn: _OpOverrideFn,
+    *,
+    allow_multiple_override: bool = False,
+    unconditional_override: bool = False,) -> None
+```
+Register a given implementation to a library - `lib_symbol = "aten"` for most cases, `op_symbol` refers to the library method you wish to override (ex. `"_scaled_grouped_mm_v2"` from above), and dispatch key will generally be one of `("CPU", "CUDA")` depending on what backend you're overriding. For all arguments, please see the comments for `register_op_override` in [registry.py](registry.py).
+
+`deregister_op_overrides() -> None` : De-register all operators that are currently registered by this DSL. Note that `torch._native.registry` has a `deregister_op_overrides` method to enable this in a centralized fashion.
+
+An example of an implementation of this spec can be found in [cutedsl_utils.py](cutedsl_utils.py), but please talk to us if you're planning on adding a new DSL.
+
+## Registration Orders and You
+
+Currently the registration order (both in general and per-op) is set by the order of imports in `torch/_native/ops/__init__.py`, noting that registration acts as a stack, in that **the last registered override for an op is the first that will be called**. If you wish to exercise control of the override ordering, please utilize one of the methods below.
+
+### User-Ordering Functions
+
+We allow for user-defined ordering functions of the form:
+
+```
+from torch._native.registry import _OverrideNode
+
+def ordering_fn(
+    op_symbol: str,
+    dispatch_key: str,
+    graph: list[_OverrideNode],
+) -> list[_OverrideNode]
+```
+
+In other words, a function that takes some context and a graph describing the override order, and returning a modified graph.
+
+**NOTE**: Graphs are described as lists of the private class `_OverrideNode` -- while this graph re-ordering functionality is public, it is both experimental and intended for advanced users only. The `_OverrideNode` class is to be used very carefully, and may change in the future.
+
+This functionality can used by either setting the environment variable `TORCH_PYTHON_NATIVE_USER_GRAPH_ORDER_FN` to an importable python function with the above signature, or by adding the following to your top-level script, post `import torch`:
+
+```
+torch._native.reorder_graphs_from_user_function(
+    my_ordering_fn,
+    reregister_overrides=True,
+)
+```
+
+Both methods are equivalent in functionality, but the environment-variable version is a little more efficient in that torch doesn't have to register **all** ops, before disabling/re-registering again based on the user-passed function.
+
+**NOTE**: The passed ordering function can be destructive in nature - one can disable an op completely by returning `[]` for a given graph, indicating that no overrides exist / are allowed. **There is currently no supported way to retrieve the original graphs - they are considered gone for the lifetime of the process**.
+
+An example user-ordering function is demonstrated below:
+
+```
+def example_ordering_fn(op_symbol, dispatch_key, nodes):
+    out_nodes = []
+
+    # disable overrides for these symbols completely
+    if op_symbol in ["_scaled_mm_v2", "add"]:
+         return []
+
+    # Only keep triton overrides otherwise
+    for node in nodes:
+        if node.dsl_name != 'triton':
+            continue
+        out_nodes.append(node)
+
+    return out_nodes
+```
diff --git a/torch/_native/__init__.py b/torch/_native/__init__.py
new file mode 100644
index 0000000000000..a5e1925944ab9
--- /dev/null
+++ b/torch/_native/__init__.py
@@ -0,0 +1,53 @@
+import os
+from functools import cache
+from typing import cast
+
+# This handles collecting registration of all native ops
+# Also need to import DSL utils to make sure DSL registration is ok
+from . import cutedsl_utils, dsl_registry, ops, registry, triton_utils
+
+
+@cache
+def get_user_ordering_fn() -> registry.UserOrderingFn | None:
+    """
+    Get a user-supplied graph-ordering function if specified.
+
+    Pass in a `package.submodule.fn` string to the env variable
+    `TORCH_PYTHON_NATIVE_USER_GRAPH_ORDER_FN` that implements the
+    calling API described in `torch/_native/README.md`. This function
+    must be part of an importable path.
+
+    Return either the imported function or `None`
+    """
+    env_var = os.getenv("TORCH_PYTHON_NATIVE_USER_GRAPH_ORDER_FN")
+
+    if not env_var:
+        return None
+
+    try:
+        import importlib
+
+        # Split into "package.submodule.fn_name
+        module_name, fn_name = env_var.rsplit(".", 1)
+
+        module = importlib.import_module(module_name)
+        fn = getattr(module, fn_name)
+
+        if not callable(fn):
+            raise TypeError(f"{env_var} does not describe a callable function")
+
+        # Cast needed: getattr returns object, but we've verified fn is callable with correct signature
+        return cast(registry.UserOrderingFn, fn)
+    except Exception as e:
+        raise ValueError(
+            f"Could not resolve {env_var} into an importable & callable function"
+        ) from e
+
+
+user_order_fn = get_user_ordering_fn()
+if user_order_fn:
+    registry.reorder_graphs_from_user_function(user_order_fn)
+
+
+# Actually perform all registrations
+registry._register_all_overrides()
diff --git a/torch/_native/common_utils.py b/torch/_native/common_utils.py
new file mode 100644
index 0000000000000..40a07b253cb8d
--- /dev/null
+++ b/torch/_native/common_utils.py
@@ -0,0 +1,64 @@
+import importlib
+import importlib.metadata
+import os
+from functools import cache
+
+from torch._vendor.packaging import version as _packaging_version
+
+
+@cache
+def check_native_jit_disabled() -> bool:
+    """
+    Single point to check if native DSL ops are disabled globally,
+    checked via:
+    TORCH_DISABLE_NATIVE_JIT=1
+    """
+    return int(os.getenv("TORCH_DISABLE_NATIVE_JIT", 0)) == 1
+
+
+def _unavailable_reason(deps: list[tuple[str, str]]) -> None | str:
+    """
+    Check availability of required packages - cuteDSL & deps,
+    informing user what (if anything) is missing
+
+    NOTE: Doesn't actually import anything.
+    """
+    for package_name, module_name in deps:
+        # Note this doesn't actually import the packages
+        if importlib.util.find_spec(module_name) is None:
+            return (
+                f"missing optional dependency `{package_name}` "
+                f"(importlib.util.find_spec({package_name}) failed)"
+            )
+    return None
+
+
+def _available_version(package: str) -> _packaging_version.Version | None:
+    """
+    Get the installed version of a package as (major, minor, patch).
+
+    Handles pre-release suffixes like "0.7.0rc1" or "3.1.0.post1" by
+    stripping non-numeric tails from each component. Returns None on
+    parse failure.
+    """
+    try:
+        version = importlib.metadata.version(package)
+    except importlib.metadata.PackageNotFoundError:
+        return None
+
+    try:
+        v = _packaging_version.parse(version)
+    except _packaging_version.InvalidVersion:
+        return None
+
+    return v
+
+
+@cache
+def check_native_version_skip() -> bool:
+    """
+    Single point to check if native DSL version gating should be skipped,
+    checked via:
+    TORCH_NATIVE_SKIP_VERSION_CHECK=1
+    """
+    return int(os.getenv("TORCH_NATIVE_SKIP_VERSION_CHECK", 0)) == 1
diff --git a/torch/_native/cutedsl_utils.py b/torch/_native/cutedsl_utils.py
new file mode 100644
index 0000000000000..203c8812bc379
--- /dev/null
+++ b/torch/_native/cutedsl_utils.py
@@ -0,0 +1,132 @@
+import functools
+import logging
+import sys
+from typing import cast
+
+from torch._vendor.packaging.version import Version
+
+from ..backends import cuda as _cuda
+from .common_utils import (
+    _available_version,
+    _unavailable_reason,
+    check_native_jit_disabled,
+    check_native_version_skip,
+)
+from .dsl_registry import dsl_registry, DSLModuleProtocol
+from .registry import (
+    _OpFn,
+    deregister_op_overrides as _deregister_op_overrides_impl,
+    register_op_override as _register_op_override_impl,
+)
+
+
+log = logging.getLogger(__name__)
+
+
+_CUTEDSL_DSL_NAME = "cutedsl"
+_CUTEDSL_REQUIRED_VERSIONS: set[Version] = {
+    # Current version - Note Version.from_part(release=(4.4.1)) is better
+    #                   but > v26 of packaging.
+    Version(f"{4}.{4}.{1}"),
+    Version(f"{4}.{4}.{2}"),
+}
+
+
+@functools.cache
+def _check_runtime_available() -> tuple[bool, Version | None]:
+    """
+    Check if cutedsl (and deps) are available.
+
+    NOTE: Doesn't import at this point
+    """
+    # Skip all checks if running on CPU-only binary
+    if not _cuda.is_built():
+        return (False, None)
+
+    deps = [
+        ("nvidia_cutlass_dsl", "cutlass"),
+        ("apache_tvm_ffi", "tvm_ffi"),
+    ]
+    reason = _unavailable_reason(deps)
+    if reason is None:
+        available = True
+        version = _available_version("nvidia_cutlass_dsl")
+    else:
+        log.warning(
+            "CuTeDSL operators require optional Python packages "
+            "`nvidia-cutlass-dsl` and `apache-tvm-ffi`; "
+            "%s",
+            reason,
+        )
+        available = False
+        version = None
+    return available, version
+
+
+def runtime_available() -> bool:
+    available, _ = _check_runtime_available()
+    return available
+
+
+def runtime_version() -> None | Version:
+    _, version = _check_runtime_available()
+    return version
+
+
+@functools.cache
+def _version_is_ok() -> bool:
+    _, version = _check_runtime_available()
+    if check_native_version_skip() or (version in _CUTEDSL_REQUIRED_VERSIONS):
+        return True
+
+    log.warning(
+        "cutedsl version %s is not known-good (ok: %s); "
+        "set TORCH_NATIVE_SKIP_VERSION_CHECK=1 to override",
+        version,
+        _CUTEDSL_REQUIRED_VERSIONS,
+    )
+    return False
+
+
+def deregister_op_overrides() -> None:
+    """
+    Deregister all ops through cuteDSL
+    """
+    _deregister_op_overrides_impl(disable_dsl_names=_CUTEDSL_DSL_NAME)
+
+
+def register_op_override(
+    lib_symbol: str,
+    op_symbol: str,
+    dispatch_key: str,
+    impl: _OpFn,
+    *,
+    allow_multiple_override: bool = False,
+    unconditional_override: bool = False,
+) -> None:
+    """
+    See torch/_native/registry.py for the underlying implementation
+    and arguments. This is a thin, DSL-checking wrapper over
+    _register_op_override_impl
+    """
+    available, version = _check_runtime_available()
+    if (not available) or check_native_jit_disabled():
+        return
+
+    if not _version_is_ok():
+        return
+
+    _register_op_override_impl(
+        _CUTEDSL_DSL_NAME,
+        lib_symbol,
+        op_symbol,
+        dispatch_key,
+        impl,
+        allow_multiple_override=allow_multiple_override,
+        unconditional_override=unconditional_override,
+    )
+
+
+# Register this DSL module with the registry
+# Note: Import-time registration ensures DSL is available when module is loaded
+dsl_registry.register_dsl("cutedsl", cast(DSLModuleProtocol, sys.modules[__name__]))
diff --git a/torch/_native/dsl_registry.py b/torch/_native/dsl_registry.py
new file mode 100644
index 0000000000000..b2fa082ee0ae6
--- /dev/null
+++ b/torch/_native/dsl_registry.py
@@ -0,0 +1,153 @@
+# Owner(s): ["module: dsl-native-ops"]
+
+import functools
+import logging
+from typing import Protocol
+
+from torch._vendor.packaging.version import Version
+
+from .registry import _OpFn
+
+
+log = logging.getLogger(__name__)
+
+
+class DSLModuleProtocol(Protocol):
+    """Complete interface for DSL utility modules"""
+
+    def runtime_available(self) -> bool: ...
+    def runtime_version(self) -> Version | None: ...
+
+    def deregister_op_overrides(self) -> None: ...
+
+    def register_op_override(
+        self,
+        lib_symbol: str,
+        op_symbol: str,
+        dispatch_key: str,
+        impl: _OpFn,
+        *,
+        allow_multiple_override: bool = False,
+        unconditional_override: bool = False,
+    ) -> None: ...
+
+
+class DSLRegistry:
+    """Registry for DSL modules - calls their existing API functions dynamically"""
+
+    def __init__(self):
+        self._dsl_modules: dict[str, DSLModuleProtocol] = {}
+
+    def _validate_dsl_name(self, name: str) -> None:
+        """Validate DSL name at runtime"""
+        if not isinstance(name, str):
+            raise TypeError(f"DSL name must be string, got {type(name).__name__}")
+
+        if not name.strip():
+            raise ValueError("DSL name cannot be empty or whitespace")
+
+    def register_dsl(self, name: str, dsl_module: DSLModuleProtocol) -> None:
+        """Register a DSL module with required interface"""
+        # Runtime validation for name and module interface
+        self._validate_dsl_name(name)
+
+        # Validate that module implements the protocol
+        required_methods = [
+            "runtime_available",
+            "runtime_version",
+            "register_op_override",
+            "deregister_op_overrides",
+        ]
+        missing_methods = [
+            method for method in required_methods if not hasattr(dsl_module, method)
+        ]
+        if missing_methods:
+            raise TypeError(
+                f"DSL module '{name}' missing required methods: {missing_methods}"
+            )
+
+        # Handle duplicate registration case
+        if name in self._dsl_modules:
+            existing_module = self._dsl_modules[name]
+            if existing_module is dsl_module:
+                # Same module re-registering - this is OK (import-time registration)
+                log.debug(
+                    "DSL '%s' re-registered with same module",
+                    name,
+                )
+                return
+            else:
+                # Different module object but same name - warn and allow (for testing)
+                # This can happen when tests import modules directly
+                log.warning(
+                    "DSL '%s' re-registered with different module object (possibly from test imports)",
+                    name,
+                )
+                # Continue to allow the registration
+
+        # No cast needed - already properly typed
+        self._dsl_modules[name] = dsl_module
+
+        # Clear caches to prevent stale results after registration
+        self.is_dsl_available.cache_clear()
+        self.get_dsl_version.cache_clear()
+        self.list_available_dsls.cache_clear()
+        self.list_all_dsls.cache_clear()
+
+        log.info("Successfully registered DSL: %s", name)
+
+    @functools.cache  # noqa: B019
+    def is_dsl_available(self, dsl_name: str) -> bool:
+        """Check if DSL is available by calling its runtime_available()"""
+        dsl_module = self._dsl_modules.get(dsl_name)
+        if dsl_module is None:
+            return False
+        try:
+            return dsl_module.runtime_available()
+        except ImportError:
+            log.debug("DSL %s import error", dsl_name, exc_info=True)
+            return False
+        except Exception:
+            log.exception("Error checking availability for DSL %s", dsl_name)
+            return False
+
+    @functools.cache  # noqa: B019
+    def get_dsl_version(self, dsl_name: str) -> Version | None:
+        """Get DSL version by calling its runtime_version()"""
+        dsl_module = self._dsl_modules.get(dsl_name)
+        if dsl_module is None:
+            return None
+        try:
+            return dsl_module.runtime_version()
+        except Exception:
+            log.debug("Error getting version for DSL %s", dsl_name, exc_info=True)
+            return None
+
+    @functools.cache  # noqa: B019
+    def list_available_dsls(self) -> tuple[str, ...]:
+        """Get names of currently available DSLs"""
+        available = []
+        for name in self._dsl_modules:
+            if self.is_dsl_available(name):  # Use cached method
+                available.append(name)
+        return tuple(available)
+
+    @functools.cache  # noqa: B019
+    def list_all_dsls(self) -> tuple[str, ...]:
+        """Get all registered DSL names (available or not)"""
+        return tuple(self._dsl_modules.keys())
+
+    def get_dsl_module(self, name: str) -> DSLModuleProtocol | None:
+        """Get a registered DSL module by name.
+
+        Args:
+            name: Name of the DSL to retrieve.
+
+        Returns:
+            The DSL module if registered, None otherwise.
+        """
+        return self._dsl_modules.get(name)
+
+
+# Global registry instance
+dsl_registry = DSLRegistry()
diff --git a/test/dynamo_expected_failures/CPython313-test_exception_variations-ExceptStarTestCases.test_try_except_finally b/torch/_native/ops/__init__.py
similarity index 100%
rename from test/dynamo_expected_failures/CPython313-test_exception_variations-ExceptStarTestCases.test_try_except_finally
rename to torch/_native/ops/__init__.py
diff --git a/torch/_native/registry.py b/torch/_native/registry.py
new file mode 100644
index 0000000000000..0924e51f6d4a3
--- /dev/null
+++ b/torch/_native/registry.py
@@ -0,0 +1,807 @@
+import logging
+from collections.abc import Callable, Iterable
+from dataclasses import dataclass, field
+from typing import Concatenate, ParamSpec, TypeVar
+
+import torch.library
+
+
+__all__ = [
+    "UserOrderingFn",
+    "register_op_override",
+    "reorder_graphs_from_user_function",
+    "reenable_op_overrides",
+    "deregister_op_overrides",
+    "get_dsl_operations",
+]
+
+log = logging.getLogger(__name__)
+
+P = ParamSpec("P")
+R = TypeVar("R")
+
+_OpOverrideFn = Callable[Concatenate[torch.DispatchKeySet, P], R]
+_OpReplaceFn = Callable[P, R]
+
+_OpFn = _OpOverrideFn | _OpReplaceFn
+
+
+@dataclass
+class _OverrideNode:
+    """Track function override data."""
+
+    dsl_name: str
+    op_symbol: str
+    dispatch_key: str
+    override_fn: _OpFn
+    unconditional_override: bool = False
+    active: bool = True
+
+
+UserOrderingFn = Callable[[str, str, list[_OverrideNode]], list[_OverrideNode]]
+
+
+@dataclass
+class _FilterState:
+    """Manages filtering state for override nodes."""
+
+    _dsl_names: set[str] = field(default_factory=set)
+    _op_symbols: set[str] = field(default_factory=set)
+    _dispatch_keys: set[str] = field(default_factory=set)
+
+    def check_enabled(self, node: _OverrideNode) -> bool:
+        """
+        Check if a node is enabled based on current filter state.
+
+        Args:
+            node: The override node to check
+
+        Returns:
+            bool: True if the node should be enabled, False if filtered out
+        """
+        if node.dsl_name in self._dsl_names:
+            return False
+
+        if node.op_symbol in self._op_symbols:
+            return False
+
+        if node.dispatch_key in self._dispatch_keys:
+            return False
+
+        return True
+
+    def update(
+        self,
+        dsl_names: str | Iterable[str] | None,
+        op_symbols: str | Iterable[str] | None,
+        dispatch_keys: str | Iterable[str] | None,
+        remove_keys: bool = False,
+    ) -> None:
+        """
+        Update filter sets as (current | new) or (current ~ new).
+
+        Args:
+            dsl_names: DSL names to add/remove from filter
+            op_symbols: Operation symbols to add/remove from filter
+            dispatch_keys: Dispatch keys to add/remove from filter
+            remove_keys: If True, remove keys from filter; if False, add them
+
+        Note:
+            Uses set.discard as it doesn't raise an exception if the element
+            wasn't in the set to begin with.
+        """
+        if remove_keys:
+            self._dsl_names -= set(_resolve_iterable(dsl_names))
+            self._op_symbols -= set(_resolve_iterable(op_symbols))
+            self._dispatch_keys -= set(_resolve_iterable(dispatch_keys))
+        else:
+            self._dsl_names |= set(_resolve_iterable(dsl_names))
+            self._op_symbols |= set(_resolve_iterable(op_symbols))
+            self._dispatch_keys |= set(_resolve_iterable(dispatch_keys))
+
+    def build_disable_key_set(self) -> set[tuple[str, str]]:
+        """
+        Build a set of dictionary keys based on the current filter state.
+
+        Returns:
+            set[tuple[str, str]]: Set of (op_symbol, dispatch_key) tuples
+        """
+        return _build_key_set(
+            self._dsl_names,
+            self._op_symbols,
+            self._dispatch_keys,
+        )
+
+    def __str__(self) -> str:
+        """Return string representation of filter state."""
+        s = ""
+        s += "Filter State:\n"
+        s += "  === DSL: ===\n"
+        for i, dsl in enumerate(self._dsl_names):
+            s += f"    {i}: {dsl}\n"
+        s += "  === OP SYMBOL: ===\n"
+        for i, op in enumerate(self._op_symbols):
+            s += f"    {i}: {op}\n"
+        s += "  === DISPATCH KEYS: ===\n"
+        for i, key in enumerate(self._dispatch_keys):
+            s += f"    {i}: {key}\n"
+
+        return s
+
+
+# Store the global override filtering state
+_filter_state: _FilterState = _FilterState()
+
+# Store torch.library.Library instances
+_libs: dict[tuple[str, str], torch.library.Library] = {}
+
+# store graph structures
+_GraphsType = dict[tuple[str, str], list[_OverrideNode]]
+_graphs: _GraphsType = {}
+
+_MappingType = dict[str, list[tuple[str, str]]]
+
+# map a {dsl, op, dispatch_key} to keys to all graphs that contain it
+_dsl_name_to_lib_graph: _MappingType = {}
+_dispatch_key_to_lib_graph: _MappingType = {}
+_op_symbol_to_lib_graph: _MappingType = {}
+
+
+def _build_key_set(
+    dsl_names: str | Iterable[str] | None,
+    op_symbols: str | Iterable[str] | None,
+    dispatch_keys: str | Iterable[str] | None,
+) -> set[tuple[str, str]]:
+    """
+    Build a set of dictionary keys based on filter criteria.
+
+    Args:
+        dsl_names: DSL names to include in key set
+        op_symbols: Operation symbols to include in key set
+        dispatch_keys: Dispatch keys to include in key set
+
+    Returns:
+        set[tuple[str, str]]: Set of (op_symbol, dispatch_key) tuples
+    """
+    key_set: set[tuple[str, str]] = set()
+
+    def _append_to_set(
+        entries: str | Iterable[str] | None, graph_lib_dict: _MappingType
+    ) -> None:
+        """Helper to add matching keys from graph_lib_dict to key_set."""
+        resolved_entries = _resolve_iterable(entries)
+
+        for entry in resolved_entries:
+            if entry in graph_lib_dict:
+                for key in graph_lib_dict[entry]:
+                    key_set.add(key)
+
+    _append_to_set(dsl_names, _dsl_name_to_lib_graph)
+    _append_to_set(op_symbols, _op_symbol_to_lib_graph)
+    _append_to_set(dispatch_keys, _dispatch_key_to_lib_graph)
+
+    return key_set
+
+
+def _print_override_graphs(*, print_inactive: bool = False) -> None:
+    """
+    Print all override graphs for debugging purposes.
+
+    Args:
+        print_inactive: Whether to print inactive nodes
+    """
+    for (op, key), node_list in _graphs.items():
+        print(f"{op=}, {key=}")
+
+        for i, node in enumerate(node_list):
+            if node.active or print_inactive:
+                s: str = f"    {i}: {node.dsl_name=}, {node.unconditional_override=}"
+                if print_inactive:
+                    s += f" {node.active=}"
+
+                print(s)
+
+
+def _get_or_create_library(op_symbol: str, dispatch_key: str) -> torch.library.Library:
+    """
+    Get or create a torch.library.Library instance for the given key.
+
+    Args:
+        op_symbol: The operation symbol
+        dispatch_key: The dispatch key
+
+    Returns:
+        torch.library.Library: The library instance
+    """
+    global _libs
+
+    key = (op_symbol, dispatch_key)
+    if key not in _libs:
+        _libs[key] = torch.library.Library("aten", "IMPL", dispatch_key)
+
+    return _libs[key]
+
+
+def _register_node_impl(
+    lib: torch.library.Library, node: _OverrideNode, dispatch_key: str
+) -> None:
+    """
+    Register a single node implementation with the library.
+
+    Args:
+        lib: The torch.library.Library instance
+        node: The override node to register
+        dispatch_key: The dispatch key for registration
+    """
+    lib.impl(
+        node.op_symbol,
+        node.override_fn,
+        dispatch_key,
+        with_keyset=not node.unconditional_override,
+        allow_override=True,
+    )
+
+
+def _resolve_iterable(iterable: str | Iterable[str] | None) -> Iterable[str]:
+    """
+    Resolve various input types to a consistent iterable of strings.
+
+    Args:
+        iterable: String, iterable of strings, or None
+
+    Returns:
+        Iterable[str]: Consistent iterable output
+    """
+    if iterable is None:
+        return []
+
+    if not isinstance(iterable, Iterable) or isinstance(iterable, str):
+        return (iterable,)
+
+    return iterable
+
+
+def reenable_op_overrides(
+    *,
+    enable_dsl_names: str | list[str] | None = None,
+    enable_op_symbols: str | list[str] | None = None,
+    enable_dispatch_keys: str | list[str] | None = None,
+) -> None:
+    """
+    Re-enable overrides by removing them from filter state and reregistering.
+
+    Args:
+        enable_dsl_names: DSL names to re-enable
+        enable_op_symbols: Operation symbols to re-enable
+        enable_dispatch_keys: Dispatch keys to re-enable
+
+    Note:
+        This function uses reverse filter state management (removing from
+        filters to enable).
+    """
+    log.info(
+        "Re-registering ops by dsl: %s, op_symbol: %s, dispatch_key: %s",
+        enable_dsl_names,
+        enable_op_symbols,
+        enable_dispatch_keys,
+    )
+
+    # Update the filters - note `remove_keys=True` because
+    # we are removing keys from the filters (vs. adding them)
+    _filter_state.update(
+        enable_dsl_names,
+        enable_op_symbols,
+        enable_dispatch_keys,
+        remove_keys=True,
+    )
+
+    # Get the set of keys that need to be reprocessed
+    key_set: set[tuple[str, str]] = _build_key_set(
+        enable_dsl_names,
+        enable_op_symbols,
+        enable_dispatch_keys,
+    )
+
+    # Process each affected graph with updated filter state
+    for key in key_set:
+        op_symbol, dispatch_key = key
+
+        if key in _graphs:
+            # Note: We don't need to cleanup and recreate the library here
+            # since we're just updating the registration with new filter state
+            _register_overrides_from_graph(
+                op_symbol, dispatch_key, _graphs[key], filter_state=_filter_state
+            )
+
+
+def deregister_op_overrides(
+    *,
+    disable_dsl_names: str | list[str] | None = None,
+    disable_op_symbols: str | list[str] | None = None,
+    disable_dispatch_keys: str | list[str] | None = None,
+) -> None:
+    """
+    De-register overrides by updating filter state and reregistering graphs.
+
+    Args:
+        disable_dsl_names: DSL names to disable
+        disable_op_symbols: Operation symbols to disable
+        disable_dispatch_keys: Dispatch keys to disable
+
+    Note:
+        This function uses filter state management to selectively disable
+        operations.
+    """
+    log.info(
+        "De-registering ops by dsl: %s, op_symbol: %s, dispatch_key: %s",
+        disable_dsl_names,
+        disable_op_symbols,
+        disable_dispatch_keys,
+    )
+
+    # Update filter state to disable specified entries
+    _filter_state.update(disable_dsl_names, disable_op_symbols, disable_dispatch_keys)
+
+    # Get the set of keys that need to be reprocessed
+    key_set: set[tuple[str, str]] = _filter_state.build_disable_key_set()
+
+    # Process each affected graph with filter state
+    for key in key_set:
+        op_symbol, dispatch_key = key
+
+        if key in _graphs:
+            _cleanup_and_reregister_graph(
+                op_symbol,
+                dispatch_key,
+                _graphs[key],
+                filter_state=_filter_state,
+            )
+
+
+def get_dsl_operations(dsl_name: str) -> list[str]:
+    """Get list of operations registered by a specific DSL.
+
+    Args:
+        dsl_name: Name of the DSL to query.
+
+    Returns:
+        Sorted list of operation names registered by the DSL.
+    """
+    operations = set()
+    for (op_symbol, _), nodes in _graphs.items():
+        for node in nodes:
+            if node.dsl_name == dsl_name:
+                operations.add(op_symbol)
+                break
+    return sorted(operations)
+
+
+def _update_registration_maps(
+    dsl_name: str,
+    op_symbol: str,
+    dispatch_key: str,
+    key: tuple[str, str],
+) -> None:
+    """
+    Update the registration mapping dictionaries.
+
+    Args:
+        dsl_name: The DSL name
+        op_symbol: The operation symbol
+        dispatch_key: The dispatch key
+        key: The dictionary key tuple
+    """
+    global _dsl_name_to_lib_graph
+    global _op_symbol_to_lib_graph
+    global _dispatch_key_to_lib_graph
+
+    def _get_new_entry_or_append(
+        registration: dict[str, list[tuple[str, str]]],
+        symbol: str,
+        key: tuple[str, str],
+    ) -> None:
+        """Helper to add key to registration list or create new entry."""
+        entry_list = registration.get(symbol)
+
+        if entry_list is None:
+            entry_list = [key]
+            registration[symbol] = entry_list
+        else:
+            entry_list.append(key)
+
+    _get_new_entry_or_append(_dsl_name_to_lib_graph, dsl_name, key)
+    _get_new_entry_or_append(_op_symbol_to_lib_graph, op_symbol, key)
+    _get_new_entry_or_append(_dispatch_key_to_lib_graph, dispatch_key, key)
+
+
+def register_op_override(
+    backend: str,
+    lib_symbol: str,
+    op_symbol: str,
+    dispatch_key: str,
+    impl: _OpOverrideFn | _OpReplaceFn,
+    *,
+    allow_multiple_override: bool = False,
+    unconditional_override: bool = False,
+) -> None:
+    """
+    Register a passed override function to the dispatcher.
+
+    Actually a graph-building operation; real registration happens later.
+
+    Args:
+        backend: The backend name (DSL name)
+        lib_symbol: Library you're overriding symbols in (must be "aten")
+        op_symbol: Name of the operation you're overriding
+        dispatch_key: Dispatch key to override
+        impl: Implementation function for the override
+        allow_multiple_override: Allow overriding an existing override
+        unconditional_override: Implementation doesn't have a fallback and
+            doesn't require torch.DispatchKeySet as the first argument
+
+    Raises:
+        ValueError: If lib_symbol is not "aten"
+    """
+    if lib_symbol != "aten":
+        raise ValueError(f'Unsupported lib_symbol (must be "aten", got: "{lib_symbol}"')
+
+    key = (op_symbol, dispatch_key)
+
+    global _graphs
+    op_graph = _graphs.get(key, [])
+
+    op_graph.append(
+        _OverrideNode(
+            dsl_name=backend,
+            op_symbol=op_symbol,
+            dispatch_key=dispatch_key,
+            override_fn=impl,
+            unconditional_override=unconditional_override,
+        )
+    )
+    _graphs[key] = op_graph
+    # Build additional maps helpful for de-registration
+    _update_registration_maps(backend, op_symbol, dispatch_key, key=key)
+
+
+def _should_reregister_graph(
+    original_graph: list[_OverrideNode],
+    new_graph: list[_OverrideNode],
+    *,
+    force_reregister: bool = False,
+) -> bool:
+    """
+    Determine if a graph needs reregistration based on changes.
+
+    Args:
+        original_graph: The original graph before modification
+        new_graph: The graph after modification
+        force_reregister: If True, always reregister regardless of changes
+
+    Returns:
+        bool: True if reregistration is needed
+    """
+    if force_reregister:
+        return True
+
+    # Check if the graph structure has changed
+    return original_graph != new_graph
+
+
+def _cleanup_and_reregister_graph(
+    op_symbol: str,
+    dispatch_key: str,
+    graph: list[_OverrideNode],
+    *,
+    filter_state: _FilterState | None = None,
+) -> None:
+    """
+    Clean up existing library and reregister a graph.
+
+    This is the common pattern used across reorder, deregister, and reenable operations.
+
+    Args:
+        op_symbol: The operation symbol
+        dispatch_key: The dispatch key
+        graph: The graph to register
+        filter_state: Optional filter state for conditional registration
+    """
+    key = (op_symbol, dispatch_key)
+
+    # Remove existing library if it exists
+    if key in _libs:
+        del _libs[key]
+
+    # Only create a library if the graph has nodes
+    # Empty graphs (disabled operations) shouldn't get libraries
+    if graph:
+        _register_overrides_from_graph(
+            op_symbol,
+            dispatch_key,
+            graph,
+            filter_state=filter_state,
+        )
+
+
+def _apply_graph_transformation(
+    transformation_fn: UserOrderingFn,
+    *,
+    keys_to_process: set[tuple[str, str]] | None = None,
+    reregister_overrides: bool = False,
+    filter_state: _FilterState | None = None,
+) -> None:
+    """
+    Apply a transformation function to graphs and optionally reregister.
+
+    This is the core pattern used by reorder_graphs_from_user_function and
+    can be reused for other graph transformation operations.
+
+    Args:
+        transformation_fn: Function to transform each graph
+        keys_to_process: Keys to process, or None for all graphs
+        reregister_overrides: Whether to reregister changed graphs
+        filter_state: Optional filter state for conditional registration
+
+    Note:
+        If transformation_fn raises an exception for a specific graph, that graph
+        will be skipped and processing will continue with remaining graphs.
+    """
+    global _graphs
+
+    # Determine which graphs to process
+    target_keys = (
+        keys_to_process if keys_to_process is not None else set(_graphs.keys())
+    )
+
+    # Process each graph
+    for op_symbol, dispatch_key in list(target_keys):
+        if (op_symbol, dispatch_key) not in _graphs:
+            continue  # Skip if graph doesn't exist
+
+        original_graph = list(_graphs[(op_symbol, dispatch_key)])
+
+        # Apply the transformation with error handling
+        try:
+            new_graph = transformation_fn(op_symbol, dispatch_key, original_graph)
+        except (TypeError, ValueError, AttributeError, RuntimeError):
+            log.warning(
+                "Graph transformation failed for %s/%s. Preserving original graph.",
+                op_symbol,
+                dispatch_key,
+                exc_info=True,
+            )
+            continue
+        except Exception:
+            log.exception(
+                "Unexpected error in graph transformation for %s/%s. Preserving original graph.",
+                op_symbol,
+                dispatch_key,
+            )
+            continue
+
+        # Validate that the transformation returned a valid result
+        if not isinstance(new_graph, list):
+            log.warning(
+                "Graph transformation returned invalid type %s for %s/%s. Expected list. Preserving original graph.",
+                type(new_graph).__name__,
+                op_symbol,
+                dispatch_key,
+            )
+            continue
+
+        # Update the graph
+        _graphs[(op_symbol, dispatch_key)] = new_graph
+
+        # Reregister if needed
+        if reregister_overrides and _should_reregister_graph(
+            original_graph, new_graph, force_reregister=False
+        ):
+            _cleanup_and_reregister_graph(
+                op_symbol,
+                dispatch_key,
+                new_graph,
+                filter_state=filter_state,
+            )
+
+
+def _register_overrides_from_graph(
+    op_symbol: str,
+    dispatch_key: str,
+    graph: list[_OverrideNode],
+    *,
+    filter_state: _FilterState | None = None,
+) -> None:
+    """
+    Register all overrides in a single graph.
+
+    Args:
+        op_symbol: The operation symbol
+        dispatch_key: The dispatch key
+        graph: List of override nodes to register
+        filter_state: Optional filter state for conditional registration
+    """
+    key = (op_symbol, dispatch_key)
+    lib = _get_or_create_library(*key)
+
+    for node in graph:
+        enable = True
+        if filter_state:
+            enable = filter_state.check_enabled(node)
+
+        if enable:
+            _register_node_impl(lib, node, dispatch_key)
+            node.active = True
+        else:
+            node.active = False
+
+
+def _register_all_overrides() -> None:
+    """
+    Perform all registration calls from previously-built override graphs.
+    """
+    for key, graph in _graphs.items():
+        op_symbol, dispatch_key = key
+
+        _register_overrides_from_graph(
+            op_symbol,
+            dispatch_key,
+            graph,
+        )
+
+
+def reorder_graphs_from_user_function(
+    fn: UserOrderingFn,
+    *,
+    reregister_overrides: bool = False,
+) -> None:
+    """
+    Reorder override graphs using a user-provided ordering function.
+
+    Args:
+        fn: User-provided function that takes (op_symbol, dispatch_key, graph)
+            and returns a reordered graph
+        reregister_overrides: Whether to reregister graphs that have changed
+
+    Note:
+        This function uses the common graph transformation pattern and can serve
+        as an example for other graph manipulation operations.
+    """
+    _apply_graph_transformation(
+        transformation_fn=fn,
+        reregister_overrides=reregister_overrides,
+    )
+
+
+def _apply_graph_filter(
+    filter_fn: Callable[[str, str, _OverrideNode], bool],
+    *,
+    reregister_overrides: bool = False,
+) -> None:
+    """
+    Apply a filter function to remove nodes from graphs.
+
+    This is a convenience function that uses the graph transformation pattern
+    to filter out unwanted nodes.
+
+    Args:
+        filter_fn: Function that takes (op_symbol, dispatch_key, node) and
+            returns True to keep the node, False to remove it
+        reregister_overrides: Whether to reregister modified graphs
+
+    Example:
+        # Remove all nodes with "deprecated" in the DSL name
+        _apply_graph_filter(
+            lambda op, dk, node: "deprecated" not in node.dsl_name,
+            reregister_overrides=True
+        )
+
+    Note:
+        If filter_fn raises an exception for a specific graph, the original
+        graph will be preserved and processing will continue.
+    """
+
+    def filtering_transformation(
+        op_symbol: str, dispatch_key: str, graph: list[_OverrideNode]
+    ) -> list[_OverrideNode]:
+        """Apply filter_fn to graph with error handling."""
+        try:
+            return [node for node in graph if filter_fn(op_symbol, dispatch_key, node)]
+        except (TypeError, ValueError, AttributeError, RuntimeError):
+            log.warning(
+                "Graph transformation failed for %s/%s. Preserving original graph.",
+                op_symbol,
+                dispatch_key,
+                exc_info=True,
+            )
+            return graph
+        except Exception:
+            log.exception(
+                "Unexpected error in graph transformation for %s/%s. Preserving original graph.",
+                op_symbol,
+                dispatch_key,
+            )
+            return graph
+
+    _apply_graph_transformation(
+        transformation_fn=filtering_transformation,
+        reregister_overrides=reregister_overrides,
+    )
+
+
+def _apply_selective_reordering(
+    condition_fn: Callable[[str, str], bool],
+    ordering_fn: UserOrderingFn,
+    *,
+    reregister_overrides: bool = False,
+) -> None:
+    """
+    Apply reordering only to graphs that match a condition.
+
+    This allows for more targeted reordering operations.
+
+    Args:
+        condition_fn: Function that takes (op_symbol, dispatch_key) and
+            returns True if the graph should be reordered
+        ordering_fn: Ordering function to apply to matching graphs
+        reregister_overrides: Whether to reregister modified graphs
+
+    Example:
+        # Only reorder CUDA operations
+        _apply_selective_reordering(
+            condition_fn=lambda op, dk: dk == "CUDA",
+            ordering_fn=lambda op, dk, g: sorted(g, key=lambda n: n.dsl_name),
+            reregister_overrides=True
+        )
+
+    Note:
+        If condition_fn or ordering_fn raises an exception for a specific graph,
+        the original graph will be preserved and processing will continue.
+    """
+
+    def conditional_transformation(
+        op_symbol: str, dispatch_key: str, graph: list[_OverrideNode]
+    ) -> list[_OverrideNode]:
+        """Apply ordering_fn conditionally based on condition_fn result."""
+        try:
+            should_reorder = condition_fn(op_symbol, dispatch_key)
+        except (TypeError, ValueError, AttributeError, RuntimeError):
+            log.warning(
+                "Graph transformation failed for %s/%s. Preserving original graph.",
+                op_symbol,
+                dispatch_key,
+                exc_info=True,
+            )
+            return graph
+        except Exception:
+            log.exception(
+                "Unexpected error in graph transformation for %s/%s. Preserving original graph.",
+                op_symbol,
+                dispatch_key,
+            )
+            return graph
+
+        if should_reorder:
+            try:
+                return ordering_fn(op_symbol, dispatch_key, graph)
+            except (TypeError, ValueError, AttributeError, RuntimeError):
+                log.warning(
+                    "Graph transformation failed for %s/%s. Preserving original graph.",
+                    op_symbol,
+                    dispatch_key,
+                    exc_info=True,
+                )
+                return graph
+            except Exception:
+                log.exception(
+                    "Unexpected error in graph transformation for %s/%s. Preserving original graph.",
+                    op_symbol,
+                    dispatch_key,
+                )
+                return graph
+
+        return graph  # Return unchanged if condition doesn't match
+
+    _apply_graph_transformation(
+        transformation_fn=conditional_transformation,
+        reregister_overrides=reregister_overrides,
+    )
diff --git a/torch/_native/triton_utils.py b/torch/_native/triton_utils.py
new file mode 100644
index 0000000000000..f875e9aacc303
--- /dev/null
+++ b/torch/_native/triton_utils.py
@@ -0,0 +1,131 @@
+import functools
+import logging
+import sys
+from typing import cast
+
+from torch._vendor.packaging.version import Version
+
+from ..backends import cuda as _cuda
+from .common_utils import (
+    _available_version,
+    _unavailable_reason,
+    check_native_jit_disabled,
+    check_native_version_skip,
+)
+from .dsl_registry import dsl_registry, DSLModuleProtocol
+from .registry import (
+    _OpFn,
+    deregister_op_overrides as _deregister_op_overrides_impl,
+    register_op_override as _register_op_override_impl,
+)
+
+
+log = logging.getLogger(__name__)
+
+
+_TRITON_DSL_NAME = "triton"
+_TRITON_REQUIRED_VERSION_MAJOR = 3
+_TRITON_MINIMUM_VERSION_MINOR = 6
+
+
+@functools.cache
+def _check_runtime_available() -> tuple[bool, Version | None]:
+    """
+    Check if triton is available
+
+    NOTE: must not import at this point
+    """
+    # Skip all checks if running on CPU-only binary
+    if not _cuda.is_built():
+        return (False, None)
+
+    deps = [
+        ("triton", "triton"),
+    ]
+    reason = _unavailable_reason(deps)
+    if reason is None:
+        available = True
+        version = _available_version("triton")
+    else:
+        log.warning("triton native DSL ops require: `triton` %s", reason)
+        available = False
+        version = None
+    return available, version
+
+
+def runtime_available() -> bool:
+    available, _ = _check_runtime_available()
+    return available
+
+
+def runtime_version() -> None | Version:
+    _, version = _check_runtime_available()
+    return version
+
+
+@functools.cache
+def _version_is_sufficient() -> bool:
+    _, version = _check_runtime_available()
+
+    if version is None:
+        return False
+
+    # Either exact version, or same major
+    major_ok = version.major == _TRITON_REQUIRED_VERSION_MAJOR
+    minor_ok = version.minor >= _TRITON_MINIMUM_VERSION_MINOR
+
+    if (major_ok and minor_ok) or check_native_version_skip():
+        return True
+
+    log.warning(
+        "triton version %s is not sufficient (>= (%s.%s.*)); "
+        "set TORCH_NATIVE_SKIP_VERSION_CHECK=1 to override",
+        version,
+        _TRITON_REQUIRED_VERSION_MAJOR,
+        _TRITON_MINIMUM_VERSION_MINOR,
+    )
+    return False
+
+
+def deregister_op_overrides() -> None:
+    """
+    Deregister all ops through triton
+    """
+    _deregister_op_overrides_impl(disable_dsl_names=_TRITON_DSL_NAME)
+
+
+def register_op_override(
+    lib_symbol: str,
+    op_symbol: str,
+    dispatch_key: str,
+    impl: _OpFn,
+    *,
+    allow_multiple_override: bool = False,
+    unconditional_override: bool = False,
+) -> None:
+    """
+    See torch/_native/registry.py for the underlying implementation
+    and arguments. This is a thin, DSL-checking wrapper over
+    _register_op_override_impl
+    """
+    available, version = _check_runtime_available()
+    if (not available) or check_native_jit_disabled():
+        return
+
+    if not _version_is_sufficient():
+        return
+
+    _register_op_override_impl(
+        _TRITON_DSL_NAME,
+        lib_symbol,
+        op_symbol,
+        dispatch_key,
+        impl,
+        allow_multiple_override=allow_multiple_override,
+        unconditional_override=unconditional_override,
+    )
+
+
+# Register this DSL module with the registry
+# Note: Import-time registration ensures DSL is available when module is loaded
+dsl_registry.register_dsl("triton", cast(DSLModuleProtocol, sys.modules[__name__]))
diff --git a/torch/_numpy/_funcs_impl.py b/torch/_numpy/_funcs_impl.py
index f3fa2968aa20f..269c8a34c28f0 100644
--- a/torch/_numpy/_funcs_impl.py
+++ b/torch/_numpy/_funcs_impl.py
@@ -13,7 +13,7 @@
 import builtins
 import itertools
 import operator
-from typing import Optional, TYPE_CHECKING
+from typing import TYPE_CHECKING
 
 import torch
 
@@ -43,7 +43,7 @@ def copy(
 def copyto(
     dst: NDArray,
     src: ArrayLike,
-    casting: Optional[CastingModes] = "same_kind",
+    casting: CastingModes | None = "same_kind",
     where: NotImplementedType = None,
 ):
     (src,) = _util.typecast_tensors((src,), dst.dtype, casting=casting)
@@ -103,7 +103,7 @@ def _concat_cast_helper(tensors, out=None, dtype=None, casting="same_kind"):
 
 
 def _concatenate(
-    tensors, axis=0, out=None, dtype=None, casting: Optional[CastingModes] = "same_kind"
+    tensors, axis=0, out=None, dtype=None, casting: CastingModes | None = "same_kind"
 ):
     # pure torch implementation, used below and in cov/corrcoef below
     tensors, axis = _util.axis_none_flatten(*tensors, axis=axis)
@@ -114,9 +114,9 @@ def _concatenate(
 def concatenate(
     ar_tuple: Sequence[ArrayLike],
     axis=0,
-    out: Optional[OutArray] = None,
-    dtype: Optional[DTypeLike] = None,
-    casting: Optional[CastingModes] = "same_kind",
+    out: OutArray | None = None,
+    dtype: DTypeLike | None = None,
+    casting: CastingModes | None = "same_kind",
 ):
     _concat_check(ar_tuple, dtype, out=out)
     result = _concatenate(ar_tuple, axis=axis, out=out, dtype=dtype, casting=casting)
@@ -126,8 +126,8 @@ def concatenate(
 def vstack(
     tup: Sequence[ArrayLike],
     *,
-    dtype: Optional[DTypeLike] = None,
-    casting: Optional[CastingModes] = "same_kind",
+    dtype: DTypeLike | None = None,
+    casting: CastingModes | None = "same_kind",
 ):
     _concat_check(tup, dtype, out=None)
     tensors = _concat_cast_helper(tup, dtype=dtype, casting=casting)
@@ -140,8 +140,8 @@ def vstack(
 def hstack(
     tup: Sequence[ArrayLike],
     *,
-    dtype: Optional[DTypeLike] = None,
-    casting: Optional[CastingModes] = "same_kind",
+    dtype: DTypeLike | None = None,
+    casting: CastingModes | None = "same_kind",
 ):
     _concat_check(tup, dtype, out=None)
     tensors = _concat_cast_helper(tup, dtype=dtype, casting=casting)
@@ -151,8 +151,8 @@ def hstack(
 def dstack(
     tup: Sequence[ArrayLike],
     *,
-    dtype: Optional[DTypeLike] = None,
-    casting: Optional[CastingModes] = "same_kind",
+    dtype: DTypeLike | None = None,
+    casting: CastingModes | None = "same_kind",
 ):
     # XXX: in numpy 1.24 dstack does not have dtype and casting keywords
     # but {h,v}stack do.  Hence add them here for consistency.
@@ -164,8 +164,8 @@ def dstack(
 def column_stack(
     tup: Sequence[ArrayLike],
     *,
-    dtype: Optional[DTypeLike] = None,
-    casting: Optional[CastingModes] = "same_kind",
+    dtype: DTypeLike | None = None,
+    casting: CastingModes | None = "same_kind",
 ):
     # XXX: in numpy 1.24 column_stack does not have dtype and casting keywords
     # but row_stack does. (because row_stack is an alias for vstack, really).
@@ -178,10 +178,10 @@ def column_stack(
 def stack(
     arrays: Sequence[ArrayLike],
     axis=0,
-    out: Optional[OutArray] = None,
+    out: OutArray | None = None,
     *,
-    dtype: Optional[DTypeLike] = None,
-    casting: Optional[CastingModes] = "same_kind",
+    dtype: DTypeLike | None = None,
+    casting: CastingModes | None = "same_kind",
 ):
     _concat_check(arrays, dtype, out=out)
 
@@ -210,7 +210,7 @@ def _split_helper(tensor, indices_or_sections, axis, strict=False):
         # NB: drop split=..., it only applies to split_helper_int
         return _split_helper_list(tensor, list(indices_or_sections), axis)
     else:
-        raise TypeError("split_helper: ", type(indices_or_sections))
+        raise TypeError(f"split_helper: {type(indices_or_sections)}")
 
 
 def _split_helper_int(tensor, indices_or_sections, axis, strict=False):
@@ -301,7 +301,7 @@ def linspace(
     num=50,
     endpoint=True,
     retstep=False,
-    dtype: Optional[DTypeLike] = None,
+    dtype: DTypeLike | None = None,
     axis=0,
 ):
     if axis != 0 or retstep or not endpoint:
@@ -317,7 +317,7 @@ def geomspace(
     stop: ArrayLike,
     num=50,
     endpoint=True,
-    dtype: Optional[DTypeLike] = None,
+    dtype: DTypeLike | None = None,
     axis=0,
 ):
     if axis != 0 or not endpoint:
@@ -338,7 +338,7 @@ def logspace(
     num=50,
     endpoint=True,
     base=10.0,
-    dtype: Optional[DTypeLike] = None,
+    dtype: DTypeLike | None = None,
     axis=0,
 ):
     if axis != 0 or not endpoint:
@@ -347,10 +347,10 @@ def logspace(
 
 
 def arange(
-    start: Optional[ArrayLikeOrScalar] = None,
-    stop: Optional[ArrayLikeOrScalar] = None,
-    step: Optional[ArrayLikeOrScalar] = 1,
-    dtype: Optional[DTypeLike] = None,
+    start: ArrayLikeOrScalar | None = None,
+    stop: ArrayLikeOrScalar | None = None,
+    step: ArrayLikeOrScalar | None = 1,
+    dtype: DTypeLike | None = None,
     *,
     like: NotImplementedType = None,
 ):
@@ -392,7 +392,7 @@ def arange(
 
 def empty(
     shape,
-    dtype: Optional[DTypeLike] = None,
+    dtype: DTypeLike | None = None,
     order: NotImplementedType = "C",
     *,
     like: NotImplementedType = None,
@@ -408,7 +408,7 @@ def empty(
 
 def empty_like(
     prototype: ArrayLike,
-    dtype: Optional[DTypeLike] = None,
+    dtype: DTypeLike | None = None,
     order: NotImplementedType = "K",
     subok: NotImplementedType = False,
     shape=None,
@@ -422,7 +422,7 @@ def empty_like(
 def full(
     shape,
     fill_value: ArrayLike,
-    dtype: Optional[DTypeLike] = None,
+    dtype: DTypeLike | None = None,
     order: NotImplementedType = "C",
     *,
     like: NotImplementedType = None,
@@ -439,7 +439,7 @@ def full(
 def full_like(
     a: ArrayLike,
     fill_value,
-    dtype: Optional[DTypeLike] = None,
+    dtype: DTypeLike | None = None,
     order: NotImplementedType = "K",
     subok: NotImplementedType = False,
     shape=None,
@@ -453,7 +453,7 @@ def full_like(
 
 def ones(
     shape,
-    dtype: Optional[DTypeLike] = None,
+    dtype: DTypeLike | None = None,
     order: NotImplementedType = "C",
     *,
     like: NotImplementedType = None,
@@ -465,7 +465,7 @@ def ones(
 
 def ones_like(
     a: ArrayLike,
-    dtype: Optional[DTypeLike] = None,
+    dtype: DTypeLike | None = None,
     order: NotImplementedType = "K",
     subok: NotImplementedType = False,
     shape=None,
@@ -478,7 +478,7 @@ def ones_like(
 
 def zeros(
     shape,
-    dtype: Optional[DTypeLike] = None,
+    dtype: DTypeLike | None = None,
     order: NotImplementedType = "C",
     *,
     like: NotImplementedType = None,
@@ -490,7 +490,7 @@ def zeros(
 
 def zeros_like(
     a: ArrayLike,
-    dtype: Optional[DTypeLike] = None,
+    dtype: DTypeLike | None = None,
     order: NotImplementedType = "K",
     subok: NotImplementedType = False,
     shape=None,
@@ -531,12 +531,12 @@ def _xy_helper_corrcoef(x_tensor, y_tensor=None, rowvar=True):
 
 def corrcoef(
     x: ArrayLike,
-    y: Optional[ArrayLike] = None,
+    y: ArrayLike | None = None,
     rowvar=True,
     bias=None,
     ddof=None,
     *,
-    dtype: Optional[DTypeLike] = None,
+    dtype: DTypeLike | None = None,
 ):
     if bias is not None or ddof is not None:
         # deprecated in NumPy
@@ -559,14 +559,14 @@ def corrcoef(
 
 def cov(
     m: ArrayLike,
-    y: Optional[ArrayLike] = None,
+    y: ArrayLike | None = None,
     rowvar=True,
     bias=False,
     ddof=None,
-    fweights: Optional[ArrayLike] = None,
-    aweights: Optional[ArrayLike] = None,
+    fweights: ArrayLike | None = None,
+    aweights: ArrayLike | None = None,
     *,
-    dtype: Optional[DTypeLike] = None,
+    dtype: DTypeLike | None = None,
 ):
     m = _xy_helper_corrcoef(m, y, rowvar)
 
@@ -629,7 +629,7 @@ def correlate(a: ArrayLike, v: ArrayLike, mode="valid"):
 # ### logic & element selection ###
 
 
-def bincount(x: ArrayLike, /, weights: Optional[ArrayLike] = None, minlength=0):
+def bincount(x: ArrayLike, /, weights: ArrayLike | None = None, minlength=0):
     if x.numel() == 0:
         # edge case allowed by numpy
         x = x.new_empty(0, dtype=int)
@@ -642,8 +642,8 @@ def bincount(x: ArrayLike, /, weights: Optional[ArrayLike] = None, minlength=0):
 
 def where(
     condition: ArrayLike,
-    x: Optional[ArrayLikeOrScalar] = None,
-    y: Optional[ArrayLikeOrScalar] = None,
+    x: ArrayLikeOrScalar | None = None,
+    y: ArrayLikeOrScalar | None = None,
     /,
 ):
     if (x is None) != (y is None):
@@ -748,7 +748,7 @@ def meshgrid(*xi: ArrayLike, copy=True, sparse=False, indexing="xy"):
     return list(output)  # match numpy, return a list
 
 
-def indices(dimensions, dtype: Optional[DTypeLike] = int, sparse=False):
+def indices(dimensions, dtype: DTypeLike | None = int, sparse=False):
     # https://github.com/numpy/numpy/blob/v1.24.0/numpy/core/numeric.py#L1691-L1791
     dimensions = tuple(dimensions)
     N = len(dimensions)
@@ -809,7 +809,7 @@ def tri(
     N,
     M=None,
     k=0,
-    dtype: Optional[DTypeLike] = None,
+    dtype: DTypeLike | None = None,
     *,
     like: NotImplementedType = None,
 ):
@@ -879,7 +879,7 @@ def take(
     a: ArrayLike,
     indices: ArrayLike,
     axis=None,
-    out: Optional[OutArray] = None,
+    out: OutArray | None = None,
     mode: NotImplementedType = "raise",
 ):
     (a,), axis = _util.axis_none_flatten(a, axis=axis)
@@ -930,7 +930,7 @@ def put_along_axis(arr: ArrayLike, indices: ArrayLike, values: ArrayLike, axis):
 def choose(
     a: ArrayLike,
     choices: Sequence[ArrayLike],
-    out: Optional[OutArray] = None,
+    out: OutArray | None = None,
     mode: NotImplementedType = "raise",
 ):
     # First, broadcast elements of `choices`
@@ -983,9 +983,9 @@ def flatnonzero(a: ArrayLike):
 
 def clip(
     a: ArrayLike,
-    min: Optional[ArrayLike] = None,
-    max: Optional[ArrayLike] = None,
-    out: Optional[OutArray] = None,
+    min: ArrayLike | None = None,
+    max: ArrayLike | None = None,
+    out: OutArray | None = None,
 ):
     return torch.clamp(a, min, max)
 
@@ -1041,8 +1041,8 @@ def trace(
     offset=0,
     axis1=0,
     axis2=1,
-    dtype: Optional[DTypeLike] = None,
-    out: Optional[OutArray] = None,
+    dtype: DTypeLike | None = None,
+    out: OutArray | None = None,
 ):
     result = torch.diagonal(a, offset, dim1=axis1, dim2=axis2).sum(-1, dtype=dtype)
     return result
@@ -1052,7 +1052,7 @@ def eye(
     N,
     M=None,
     k=0,
-    dtype: Optional[DTypeLike] = None,
+    dtype: DTypeLike | None = None,
     order: NotImplementedType = "C",
     *,
     like: NotImplementedType = None,
@@ -1066,7 +1066,7 @@ def eye(
     return z
 
 
-def identity(n, dtype: Optional[DTypeLike] = None, *, like: NotImplementedType = None):
+def identity(n, dtype: DTypeLike | None = None, *, like: NotImplementedType = None):
     return torch.eye(n, dtype=dtype)
 
 
@@ -1171,7 +1171,7 @@ def tensordot(a: ArrayLike, b: ArrayLike, axes=2):
     return torch.tensordot(a, b, dims=axes)
 
 
-def dot(a: ArrayLike, b: ArrayLike, out: Optional[OutArray] = None):
+def dot(a: ArrayLike, b: ArrayLike, out: OutArray | None = None):
     dtype = _dtypes_impl.result_type_impl(a, b)
     is_bool = dtype == torch.bool
     if is_bool:
@@ -1214,7 +1214,7 @@ def inner(a: ArrayLike, b: ArrayLike, /):
     return result
 
 
-def outer(a: ArrayLike, b: ArrayLike, out: Optional[OutArray] = None):
+def outer(a: ArrayLike, b: ArrayLike, out: OutArray | None = None):
     return torch.outer(a, b)
 
 
@@ -1410,7 +1410,7 @@ def argsort(a: ArrayLike, axis=-1, kind=None, order: NotImplementedType = None):
 
 
 def searchsorted(
-    a: ArrayLike, v: ArrayLike, side="left", sorter: Optional[ArrayLike] = None
+    a: ArrayLike, v: ArrayLike, side="left", sorter: ArrayLike | None = None
 ):
     if a.dtype.is_complex:
         raise NotImplementedError(f"searchsorted with dtype={a.dtype}")
@@ -1515,8 +1515,8 @@ def diff(
     a: ArrayLike,
     n=1,
     axis=-1,
-    prepend: Optional[ArrayLike] = None,
-    append: Optional[ArrayLike] = None,
+    prepend: ArrayLike | None = None,
+    append: ArrayLike | None = None,
 ):
     axis = _util.normalize_axis_index(axis, a.ndim)
 
@@ -1731,7 +1731,7 @@ def gradient(f: ArrayLike, *varargs, axis=None, edge_order=1):
 # ### Type/shape etc queries ###
 
 
-def round(a: ArrayLike, decimals=0, out: Optional[OutArray] = None):
+def round(a: ArrayLike, decimals=0, out: OutArray | None = None):
     if a.is_floating_point():
         result = torch.round(a, decimals=decimals)
     elif a.is_complex():
@@ -1793,11 +1793,11 @@ def isrealobj(x: ArrayLike):
     return not torch.is_complex(x)
 
 
-def isneginf(x: ArrayLike, out: Optional[OutArray] = None):
+def isneginf(x: ArrayLike, out: OutArray | None = None):
     return torch.isneginf(x)
 
 
-def isposinf(x: ArrayLike, out: Optional[OutArray] = None):
+def isposinf(x: ArrayLike, out: OutArray | None = None):
     return torch.isposinf(x)
 
 
@@ -1890,7 +1890,7 @@ def histogram(
     bins: ArrayLike = 10,
     range=None,
     normed=None,
-    weights: Optional[ArrayLike] = None,
+    weights: ArrayLike | None = None,
     density=None,
 ):
     if normed is not None:
@@ -1933,9 +1933,9 @@ def histogram2d(
     x,
     y,
     bins=10,
-    range: Optional[ArrayLike] = None,
+    range: ArrayLike | None = None,
     normed=None,
-    weights: Optional[ArrayLike] = None,
+    weights: ArrayLike | None = None,
     density=None,
 ):
     # vendored from https://github.com/numpy/numpy/blob/v1.24.0/numpy/lib/twodim_base.py#L655-L821
@@ -1958,9 +1958,9 @@ def histogram2d(
 def histogramdd(
     sample,
     bins=10,
-    range: Optional[ArrayLike] = None,
+    range: ArrayLike | None = None,
     normed=None,
-    weights: Optional[ArrayLike] = None,
+    weights: ArrayLike | None = None,
     density=None,
 ):
     # have to normalize manually because `sample` interpretation differs
diff --git a/torch/_numpy/_ndarray.py b/torch/_numpy/_ndarray.py
index d024f98588d2d..e55c119373c86 100644
--- a/torch/_numpy/_ndarray.py
+++ b/torch/_numpy/_ndarray.py
@@ -415,6 +415,10 @@ def imag(self):
     def imag(self, value):
         self.tensor.imag = asarray(value).tensor
 
+    @property
+    def flat(self):
+        return self.ravel()
+
     # ctors
     def astype(self, dtype, order="K", casting="unsafe", subok=True, copy=True):
         if order != "K":
diff --git a/torch/_numpy/_normalizations.py b/torch/_numpy/_normalizations.py
index 24908001a926b..ce6c25216f905 100644
--- a/torch/_numpy/_normalizations.py
+++ b/torch/_numpy/_normalizations.py
@@ -16,8 +16,8 @@
 
 
 ArrayLike = typing.TypeVar("ArrayLike")
-Scalar = typing.Union[int, float, complex, bool]
-ArrayLikeOrScalar = typing.Union[ArrayLike, Scalar]
+Scalar = int | float | complex | bool
+ArrayLikeOrScalar = ArrayLike | Scalar
 
 DTypeLike = typing.TypeVar("DTypeLike")
 AxisLike = typing.TypeVar("AxisLike")
@@ -135,15 +135,21 @@ def normalize_casting(arg, parm=None):  # codespell:ignore
     "ArrayLike": normalize_array_like,
     "ArrayLikeOrScalar": normalize_array_like_or_scalar,
     "Optional[ArrayLike]": normalize_optional_array_like,
+    "ArrayLike | None": normalize_optional_array_like,
     "Sequence[ArrayLike]": normalize_seq_array_like,
     "Optional[ArrayLikeOrScalar]": normalize_optional_array_like_or_scalar,
+    "ArrayLikeOrScalar | None": normalize_optional_array_like_or_scalar,
     "Optional[NDArray]": normalize_ndarray,
+    "NDArray | None": normalize_ndarray,
     "Optional[OutArray]": normalize_outarray,
+    "OutArray | None": normalize_outarray,
     "NDArray": normalize_ndarray,
     "Optional[DTypeLike]": normalize_dtype,
+    "DTypeLike | None": normalize_dtype,
     "AxisLike": normalize_axis_like,
     "NotImplementedType": normalize_not_implemented,
     "Optional[CastingModes]": normalize_casting,
+    "CastingModes | None": normalize_casting,
 }
 
 
diff --git a/torch/_numpy/_reductions_impl.py b/torch/_numpy/_reductions_impl.py
index fe5795d6ab6c2..da0a212654055 100644
--- a/torch/_numpy/_reductions_impl.py
+++ b/torch/_numpy/_reductions_impl.py
@@ -9,7 +9,7 @@
 from __future__ import annotations
 
 import functools
-from typing import Optional, TYPE_CHECKING
+from typing import TYPE_CHECKING
 
 import torch
 
@@ -73,7 +73,7 @@ def count_nonzero(a: ArrayLike, axis: AxisLike = None, *, keepdims: KeepDims = F
 def argmax(
     a: ArrayLike,
     axis: AxisLike = None,
-    out: Optional[OutArray] = None,
+    out: OutArray | None = None,
     *,
     keepdims: KeepDims = False,
 ):
@@ -93,7 +93,7 @@ def argmax(
 def argmin(
     a: ArrayLike,
     axis: AxisLike = None,
-    out: Optional[OutArray] = None,
+    out: OutArray | None = None,
     *,
     keepdims: KeepDims = False,
 ):
@@ -113,7 +113,7 @@ def argmin(
 def any(
     a: ArrayLike,
     axis: AxisLike = None,
-    out: Optional[OutArray] = None,
+    out: OutArray | None = None,
     keepdims: KeepDims = False,
     *,
     where: NotImplementedType = None,
@@ -127,7 +127,7 @@ def any(
 def all(
     a: ArrayLike,
     axis: AxisLike = None,
-    out: Optional[OutArray] = None,
+    out: OutArray | None = None,
     keepdims: KeepDims = False,
     *,
     where: NotImplementedType = None,
@@ -141,7 +141,7 @@ def all(
 def amax(
     a: ArrayLike,
     axis: AxisLike = None,
-    out: Optional[OutArray] = None,
+    out: OutArray | None = None,
     keepdims: KeepDims = False,
     initial: NotImplementedType = None,
     where: NotImplementedType = None,
@@ -159,7 +159,7 @@ def amax(
 def amin(
     a: ArrayLike,
     axis: AxisLike = None,
-    out: Optional[OutArray] = None,
+    out: OutArray | None = None,
     keepdims: KeepDims = False,
     initial: NotImplementedType = None,
     where: NotImplementedType = None,
@@ -177,7 +177,7 @@ def amin(
 def ptp(
     a: ArrayLike,
     axis: AxisLike = None,
-    out: Optional[OutArray] = None,
+    out: OutArray | None = None,
     keepdims: KeepDims = False,
 ):
     return a.amax(axis) - a.amin(axis)
@@ -187,8 +187,8 @@ def ptp(
 def sum(
     a: ArrayLike,
     axis: AxisLike = None,
-    dtype: Optional[DTypeLike] = None,
-    out: Optional[OutArray] = None,
+    dtype: DTypeLike | None = None,
+    out: OutArray | None = None,
     keepdims: KeepDims = False,
     initial: NotImplementedType = None,
     where: NotImplementedType = None,
@@ -209,8 +209,8 @@ def sum(
 def prod(
     a: ArrayLike,
     axis: AxisLike = None,
-    dtype: Optional[DTypeLike] = None,
-    out: Optional[OutArray] = None,
+    dtype: DTypeLike | None = None,
+    out: OutArray | None = None,
     keepdims: KeepDims = False,
     initial: NotImplementedType = None,
     where: NotImplementedType = None,
@@ -231,8 +231,8 @@ def prod(
 def mean(
     a: ArrayLike,
     axis: AxisLike = None,
-    dtype: Optional[DTypeLike] = None,
-    out: Optional[OutArray] = None,
+    dtype: DTypeLike | None = None,
+    out: OutArray | None = None,
     keepdims: KeepDims = False,
     *,
     where: NotImplementedType = None,
@@ -249,8 +249,8 @@ def mean(
 def std(
     a: ArrayLike,
     axis: AxisLike = None,
-    dtype: Optional[DTypeLike] = None,
-    out: Optional[OutArray] = None,
+    dtype: DTypeLike | None = None,
+    out: OutArray | None = None,
     ddof=0,
     keepdims: KeepDims = False,
     *,
@@ -267,8 +267,8 @@ def std(
 def var(
     a: ArrayLike,
     axis: AxisLike = None,
-    dtype: Optional[DTypeLike] = None,
-    out: Optional[OutArray] = None,
+    dtype: DTypeLike | None = None,
+    out: OutArray | None = None,
     ddof=0,
     keepdims: KeepDims = False,
     *,
@@ -289,8 +289,8 @@ def var(
 def cumsum(
     a: ArrayLike,
     axis: AxisLike = None,
-    dtype: Optional[DTypeLike] = None,
-    out: Optional[OutArray] = None,
+    dtype: DTypeLike | None = None,
+    out: OutArray | None = None,
 ):
     if dtype == torch.bool:
         dtype = _dtypes_impl.default_dtypes().int_dtype
@@ -306,8 +306,8 @@ def cumsum(
 def cumprod(
     a: ArrayLike,
     axis: AxisLike = None,
-    dtype: Optional[DTypeLike] = None,
-    out: Optional[OutArray] = None,
+    dtype: DTypeLike | None = None,
+    out: OutArray | None = None,
 ):
     if dtype == torch.bool:
         dtype = _dtypes_impl.default_dtypes().int_dtype
@@ -380,7 +380,7 @@ def quantile(
     a: ArrayLike,
     q: ArrayLike,
     axis: AxisLike = None,
-    out: Optional[OutArray] = None,
+    out: OutArray | None = None,
     overwrite_input=False,
     method="linear",
     keepdims: KeepDims = False,
@@ -423,7 +423,7 @@ def percentile(
     a: ArrayLike,
     q: ArrayLike,
     axis: AxisLike = None,
-    out: Optional[OutArray] = None,
+    out: OutArray | None = None,
     overwrite_input=False,
     method="linear",
     keepdims: KeepDims = False,
@@ -449,7 +449,7 @@ def percentile(
 def median(
     a: ArrayLike,
     axis=None,
-    out: Optional[OutArray] = None,
+    out: OutArray | None = None,
     overwrite_input=False,
     keepdims: KeepDims = False,
 ):
diff --git a/torch/_numpy/_ufuncs.py b/torch/_numpy/_ufuncs.py
index 0543aad2f7a6b..56d93936b441c 100644
--- a/torch/_numpy/_ufuncs.py
+++ b/torch/_numpy/_ufuncs.py
@@ -2,8 +2,6 @@
 
 from __future__ import annotations
 
-from typing import Optional
-
 import torch
 
 from . import _binary_ufuncs_impl, _dtypes_impl, _unary_ufuncs_impl, _util
@@ -76,12 +74,12 @@ def wrapped(
         x1: ArrayLikeOrScalar,
         x2: ArrayLikeOrScalar,
         /,
-        out: Optional[OutArray] = None,
+        out: OutArray | None = None,
         *,
         where: NotImplementedType = True,
-        casting: Optional[CastingModes] = "same_kind",
+        casting: CastingModes | None = "same_kind",
         order: NotImplementedType = "K",
-        dtype: Optional[DTypeLike] = None,
+        dtype: DTypeLike | None = None,
         subok: NotImplementedType = False,
         signature: NotImplementedType = None,
         extobj: NotImplementedType = None,
@@ -123,11 +121,11 @@ def matmul(
     x1: ArrayLike,
     x2: ArrayLike,
     /,
-    out: Optional[OutArray] = None,
+    out: OutArray | None = None,
     *,
-    casting: Optional[CastingModes] = "same_kind",
+    casting: CastingModes | None = "same_kind",
     order: NotImplementedType = "K",
-    dtype: Optional[DTypeLike] = None,
+    dtype: DTypeLike | None = None,
     subok: NotImplementedType = False,
     signature: NotImplementedType = None,
     extobj: NotImplementedType = None,
@@ -150,12 +148,12 @@ def ldexp(
     x1: ArrayLikeOrScalar,
     x2: ArrayLikeOrScalar,
     /,
-    out: Optional[OutArray] = None,
+    out: OutArray | None = None,
     *,
     where: NotImplementedType = True,
-    casting: Optional[CastingModes] = "same_kind",
+    casting: CastingModes | None = "same_kind",
     order: NotImplementedType = "K",
-    dtype: Optional[DTypeLike] = None,
+    dtype: DTypeLike | None = None,
     subok: NotImplementedType = False,
     signature: NotImplementedType = None,
     extobj: NotImplementedType = None,
@@ -189,15 +187,15 @@ def ldexp(
 def divmod(
     x1: ArrayLike,
     x2: ArrayLike,
-    out1: Optional[OutArray] = None,
-    out2: Optional[OutArray] = None,
+    out1: OutArray | None = None,
+    out2: OutArray | None = None,
     /,
-    out: tuple[Optional[OutArray], Optional[OutArray]] = (None, None),
+    out: tuple[OutArray | None, OutArray | None] = (None, None),
     *,
     where: NotImplementedType = True,
-    casting: Optional[CastingModes] = "same_kind",
+    casting: CastingModes | None = "same_kind",
     order: NotImplementedType = "K",
-    dtype: Optional[DTypeLike] = None,
+    dtype: DTypeLike | None = None,
     subok: NotImplementedType = False,
     signature: NotImplementedType = None,
     extobj: NotImplementedType = None,
@@ -297,12 +295,12 @@ def deco_unary_ufunc(torch_func):
     def wrapped(
         x: ArrayLike,
         /,
-        out: Optional[OutArray] = None,
+        out: OutArray | None = None,
         *,
         where=True,
-        casting: Optional[CastingModes] = "same_kind",
+        casting: CastingModes | None = "same_kind",
         order="K",
-        dtype: Optional[DTypeLike] = None,
+        dtype: DTypeLike | None = None,
         subok: NotImplementedType = False,
         signature=None,
         extobj=None,
diff --git a/torch/_numpy/random.py b/torch/_numpy/random.py
index a3d4a1c73241f..9e42cc51f4591 100644
--- a/torch/_numpy/random.py
+++ b/torch/_numpy/random.py
@@ -12,7 +12,6 @@
 
 import functools
 from math import sqrt
-from typing import Optional
 
 import torch
 
@@ -155,7 +154,7 @@ def randint(low, high=None, size=None):
 
 @deco_stream
 @normalizer
-def choice(a: ArrayLike, size=None, replace=True, p: Optional[ArrayLike] = None):
+def choice(a: ArrayLike, size=None, replace=True, p: ArrayLike | None = None):
     # https://stackoverflow.com/questions/59461811/random-choice-with-pytorch
     if a.numel() == 1:
         a = torch.arange(a)
diff --git a/torch/_opaque_base.py b/torch/_opaque_base.py
index c37115c3af1f7..d08048713e3d0 100644
--- a/torch/_opaque_base.py
+++ b/torch/_opaque_base.py
@@ -1,11 +1,15 @@
 class OpaqueBaseMeta(type):
     def __instancecheck__(cls, instance):
-        from torch._library.fake_class_registry import FakeScriptObject
+        if super().__instancecheck__(instance):
+            return True
 
-        return super().__instancecheck__(instance) or (
-            isinstance(instance, FakeScriptObject)
-            and super().__instancecheck__(instance.real_obj)
-        )
+        if hasattr(instance, "real_obj"):
+            from torch._library.fake_class_registry import FakeScriptObject
+
+            if isinstance(instance, FakeScriptObject):
+                return super().__instancecheck__(instance.real_obj)
+
+        return False
 
 
 class OpaqueBase(metaclass=OpaqueBaseMeta):
diff --git a/torch/_ops.py b/torch/_ops.py
index a27785854f46e..b2372e0266744 100644
--- a/torch/_ops.py
+++ b/torch/_ops.py
@@ -280,7 +280,9 @@ class HigherOrderOperator(OperatorBase, abc.ABC):
     # If you're creating a new HigherOrderOperator, please do not change the
     # default. Adding operators to the global torch.ops namespace is a bad
     # practice due to name collisions.
-    def __init__(self, name, *, cacheable=False):
+    def __init__(
+        self, name, *, cacheable=False, supports_training_input_mutation=False
+    ):
         super().__init__()
         if type(self) is HigherOrderOperator:
             raise RuntimeError(
@@ -294,6 +296,10 @@ def __init__(self, name, *, cacheable=False):
         self._ns = "higher_order"
         self.__module__ = "torch.ops.higher_order"
         self._cacheable = cacheable
+        # When True, allows mutating inputs that don't require grad during
+        # training. Mutations must be handled by auto_functionalize before
+        # reaching autograd.
+        self._supports_training_input_mutation = supports_training_input_mutation
 
         self.non_fallthrough_keys = torch._C._dispatch_keyset_full()
 
@@ -331,7 +337,7 @@ def maybe_run_autograd(*args: _P.args, **kwargs: _P.kwargs) -> _T:
 
             from torch._higher_order_ops.utils import _has_gen_schema
 
-            if _has_gen_schema(self):
+            if not self._supports_training_input_mutation and _has_gen_schema(self):
                 schema = self.gen_schema(*args, **kwargs)
                 if any(arg.is_write for arg in schema.arguments):
                     raise RuntimeError(
diff --git a/torch/_prims/rng_prims.py b/torch/_prims/rng_prims.py
index a37f58c6a4e5f..60039060769fb 100644
--- a/torch/_prims/rng_prims.py
+++ b/torch/_prims/rng_prims.py
@@ -394,5 +394,110 @@ def impl_functional(ctx, op, *args, rng_state=None, **kwargs):
 graphsafe_run_with_rng_state = register_graphsafe_run_with_rng_state_op()
 
 
+def register_run_dtensor_rng_op():
+    """
+    Register a higher-order operator for DTensor distributed random operations.
+    Takes pre-computed integer offsets (start_offset_incr, end_offset_incr), an op,
+    and args. Internally adjusts the RNG state using the offsets before/after
+    running the op.
+
+    The offsets are computed at trace time from the DTensorSpec, so the compiled graph
+    contains only plain integers — no DTensorSpec or DeviceMesh objects.
+    """
+
+    class RunDTensorRngOp(HigherOrderOperator):
+        def __init__(self):
+            super().__init__("run_dtensor_rng_op", cacheable=True)
+
+        def __call__(self, start_offset_incr, end_offset_incr, op, *args, **kwargs):
+            # pyrefly: ignore [missing-attribute]
+            return super().__call__(
+                start_offset_incr, end_offset_incr, op, *args, **kwargs
+            )
+
+    run_dtensor_rng_op = RunDTensorRngOp()
+
+    run_dtensor_rng_op.py_impl(DispatchKey.Autograd)(
+        autograd_not_implemented(run_dtensor_rng_op, deferred_error=True)
+    )
+
+    @run_dtensor_rng_op.py_impl(DispatchKey.CUDA)
+    def impl_cuda(start_offset_incr, end_offset_incr, op, *args, **kwargs):
+        from torch.distributed.tensor._random import _PhiloxState
+
+        state = _PhiloxState(torch.cuda.get_rng_state())
+        old_offset = state.offset.clone()
+        state.offset = old_offset + start_offset_incr
+
+        device = (
+            args[0].device
+            if args and hasattr(args[0], "device")
+            else torch.device(f"cuda:{torch.cuda.current_device()}")
+        )
+        with torch.random.fork_rng(devices=[device], device_type="cuda"):
+            torch.cuda.set_rng_state(state.state)
+            try:
+                out = op(*args, **kwargs)
+            finally:
+                state.offset = old_offset + end_offset_incr
+
+        torch.cuda.set_rng_state(state.state)
+        return out
+
+    @run_dtensor_rng_op.py_impl(DispatchKey.BackendSelect)
+    def impl_backend_select(start_offset_incr, end_offset_incr, op, *args, **kwargs):
+        device = get_device(args, kwargs)
+        if device != "cuda":
+            raise RuntimeError(
+                f"run_dtensor_rng_op only supports CUDA device, got {device}. "
+                f"This operator is designed for distributed random operations on CUDA."
+            )
+        return impl_cuda(start_offset_incr, end_offset_incr, op, *args, **kwargs)
+
+    @run_dtensor_rng_op.py_impl(FakeTensorMode)
+    def impl_fake_tensor_mode(
+        mode, start_offset_incr, end_offset_incr, op, *args, **kwargs
+    ):
+        with mode:
+            return op(*args, **kwargs)
+
+    @run_dtensor_rng_op.py_impl(ProxyTorchDispatchMode)
+    def impl_proxy_dispatch_mode(
+        mode, start_offset_incr, end_offset_incr, op, *args, **kwargs
+    ):
+        with disable_proxy_modes_tracing():
+            out = run_dtensor_rng_op(
+                start_offset_incr, end_offset_incr, op, *args, **kwargs
+            )
+        proxy_args = pytree.tree_map(
+            mode.tracer.unwrap_proxy, (start_offset_incr, end_offset_incr, op, *args)
+        )
+        proxy_kwargs = pytree.tree_map(mode.tracer.unwrap_proxy, kwargs)
+        out_proxy = mode.tracer.create_proxy(
+            "call_function", run_dtensor_rng_op, proxy_args, proxy_kwargs
+        )
+        return track_tensor_tree(out, out_proxy, constant=None, tracer=mode.tracer)
+
+    @run_dtensor_rng_op.py_functionalize_impl
+    def impl_functional(ctx, start_offset_incr, end_offset_incr, op, *args, **kwargs):
+        unwrapped_args = ctx.unwrap_tensors(args)
+        unwrapped_kwargs = ctx.unwrap_tensors(kwargs)
+
+        with ctx.redispatch_to_next():
+            out = run_dtensor_rng_op(
+                start_offset_incr,
+                end_offset_incr,
+                op,
+                *unwrapped_args,
+                **unwrapped_kwargs,
+            )
+            return ctx.wrap_tensors(out)
+
+    return run_dtensor_rng_op
+
+
+run_dtensor_rng_op = register_run_dtensor_rng_op()
+
+
 def register_rng_prims():
     register_philox_rand()
diff --git a/torch/_prims_common/__init__.py b/torch/_prims_common/__init__.py
index 398d2eb10d181..4705b6986eb04 100644
--- a/torch/_prims_common/__init__.py
+++ b/torch/_prims_common/__init__.py
@@ -45,16 +45,16 @@ def __rmul__(self, other: Any) -> typing.Self: ...
     _IntLikeT = TypeVar("_IntLikeT", bound=_WorksWithInt)
 
 
-ShapeType: TypeAlias = Union[torch.Size, list[int], tuple[int, ...]]
-StrideType: TypeAlias = Union[list[int], tuple[int, ...]]
-DimsType: TypeAlias = Union[int, list[int], tuple[int, ...]]
-DimsSequenceType: TypeAlias = Union[list[int], tuple[int, ...]]
+ShapeType: TypeAlias = torch.Size | list[int] | tuple[int, ...]
+StrideType: TypeAlias = list[int] | tuple[int, ...]
+DimsType: TypeAlias = int | list[int] | tuple[int, ...]
+DimsSequenceType: TypeAlias = list[int] | tuple[int, ...]
 # TODO: Type[torch.SymInt], Type[torch.SymFloat]
-NumberTypeType: TypeAlias = Union[type[bool], type[int], type[float], type[complex]]
+NumberTypeType: TypeAlias = type[bool] | type[int] | type[float] | type[complex]
 # TODO: This needs a lot more type annotations
 # NumberType = Union[bool, int, float, complex, torch.SymInt, torch.SymFloat]
-NumberType: TypeAlias = Union[bool, int, float, complex]
-RealNumberType: TypeAlias = Union[bool, int, float]
+NumberType: TypeAlias = bool | int | float | complex
+RealNumberType: TypeAlias = bool | int | float
 
 Number = (bool, int, float, complex, torch.SymInt, torch.SymFloat, torch.SymBool)
 # I don't call it Integral because numbers.Integral includes bool, but IntLike
@@ -65,7 +65,7 @@ def __rmul__(self, other: Any) -> typing.Self: ...
 BoolLike = (bool, torch.SymBool)
 IntWithoutSymInt = int
 FloatWithoutSymFloat = float
-DeviceLikeType: TypeAlias = Union[str, torch.device, int]
+DeviceLikeType: TypeAlias = str | torch.device | int
 Tensor = torch.Tensor
 
 
@@ -101,8 +101,8 @@ def __rmul__(self, other: Any) -> typing.Self: ...
 
 TensorLikeType = torch.Tensor
 TensorLike = torch.Tensor
-TensorSequenceType: TypeAlias = Union[list[TensorLikeType], tuple[TensorLikeType, ...]]
-TensorOrNumberLikeType: TypeAlias = Union[TensorLikeType, NumberType]
+TensorSequenceType: TypeAlias = list[TensorLikeType] | tuple[TensorLikeType, ...]
+TensorOrNumberLikeType: TypeAlias = TensorLikeType | NumberType
 
 CustomOutParamAnnotation = "__custom_out_param__"
 
@@ -214,7 +214,7 @@ def _check_strides_helper(
     only_cuda=True,
     significant_only=True,
     allow_rhs_unbacked=False,
-) -> tuple[bool, Optional[int]]:
+) -> tuple[bool, int | None]:
     # NOTE: only on CUDA because CPU elementwise strides are incorrect in PyTorch
     # See https://github.com/pytorch/pytorch/issues/77553
     # Only compares strides that are "meaningful" -- strides for dimensions with length > 1
@@ -235,7 +235,7 @@ def _check_strides_helper(
 
 def check_significant_strides(
     a: TensorLikeType, b: TensorLikeType, *, only_cuda=True, allow_rhs_unbacked=False
-) -> tuple[bool, Optional[int]]:
+) -> tuple[bool, int | None]:
     return _check_strides_helper(
         a,
         b,
@@ -247,7 +247,7 @@ def check_significant_strides(
 
 def check_all_strides(
     a: TensorLikeType, b: TensorLikeType, *, only_cuda=True
-) -> tuple[bool, Optional[int]]:
+) -> tuple[bool, int | None]:
     return _check_strides_helper(a, b, only_cuda=only_cuda, significant_only=False)
 
 
@@ -967,7 +967,7 @@ def check_same_shape(*args, allow_cpu_scalar_tensors: bool):
 
 # Acquires a common shape, if it exists, from one or more tensor arguments,
 # filtering number arguments
-def extract_shape(*args, allow_cpu_scalar_tensors: bool) -> Optional[ShapeType]:
+def extract_shape(*args, allow_cpu_scalar_tensors: bool) -> ShapeType | None:
     shape = None
     scalar_shape = None
 
@@ -994,7 +994,7 @@ def extract_shape(*args, allow_cpu_scalar_tensors: bool) -> Optional[ShapeType]:
 # Extracts dimensions that might be passed either as a list/tuple or as varargs.
 # A typical case is Tensor.permute .
 def extract_dims_from_varargs(
-    dims: Union[DimsSequenceType, tuple[DimsSequenceType, ...]],
+    dims: DimsSequenceType | tuple[DimsSequenceType, ...],
 ) -> DimsSequenceType:
     if dims and isinstance(dims[0], Sequence):
         if len(dims) != 1:
@@ -1008,7 +1008,7 @@ def extract_dims_from_varargs(
 
 
 def extract_shape_from_varargs(
-    shape: Union[ShapeType, tuple[ShapeType]],
+    shape: ShapeType | tuple[ShapeType],
     validate=True,
 ) -> tuple[int, ...]:
     """
@@ -1248,7 +1248,7 @@ def type_to_dtype(typ: type) -> torch.dtype:
     raise ValueError(f"Invalid type {typ}!")
 
 
-def get_dtype(x: Union[torch.Tensor, NumberType]):
+def get_dtype(x: torch.Tensor | NumberType):
     if isinstance(x, torch.Tensor):
         return x.dtype
     else:
@@ -1309,9 +1309,9 @@ def get_higher_type(a: type, b: type) -> type:
 #   are not ordered relative to each other, the next
 #   higher datatype
 def get_higher_dtype(
-    a: Optional[Union[torch.dtype, TensorLikeType, NumberType]],
-    b: Optional[Union[torch.dtype, TensorLikeType, NumberType]],
-) -> Optional[torch.dtype]:
+    a: torch.dtype | TensorLikeType | NumberType | None,
+    b: torch.dtype | TensorLikeType | NumberType | None,
+) -> torch.dtype | None:
     """
     Computes the "lowest" datatype that is weakly
     "higher" than both a and b.
@@ -1328,8 +1328,8 @@ def get_higher_dtype(
         )
 
     def _extract_dtype(
-        x: Optional[Union[torch.dtype, TensorLikeType, NumberType]],
-    ) -> Optional[torch.dtype]:
+        x: torch.dtype | TensorLikeType | NumberType | None,
+    ) -> torch.dtype | None:
         if x is None:
             return None
         if isinstance(x, torch.dtype):
@@ -1547,7 +1547,7 @@ class RETURN_TYPE(Enum):
 
 # TODO: when NumberType contains the sym types, can simplify this
 def number_type(
-    x: Union[NumberType, torch.SymInt, torch.SymFloat, torch.SymBool],
+    x: NumberType | torch.SymInt | torch.SymFloat | torch.SymBool,
 ) -> type:
     if isinstance(x, torch.SymInt):
         return int
@@ -1687,7 +1687,7 @@ def elementwise_dtypes(
 
     def _find_highest_dtype_filtered(
         args, filter, *, float_as_complex=False
-    ) -> Optional[torch.dtype]:
+    ) -> torch.dtype | None:
         zero_dim_tensor_dtype = None
         one_plus_dim_tensor_dtype = None
         for x in args:
@@ -1758,8 +1758,8 @@ def _find_highest_dtype_filtered(
 def reduction_dtypes(
     arg,
     output_dtype_kind: REDUCTION_OUTPUT_TYPE_KIND,
-    dtype: Optional[torch.dtype] = None,
-) -> tuple[torch.dtype, Optional[torch.dtype]]:
+    dtype: torch.dtype | None = None,
+) -> tuple[torch.dtype, torch.dtype | None]:
     # even though some reductions, like amin or amax, don't strictly require type promotion,
     # all the math ops (including comparisons) are still defined only for a computation type,
     # so promotion will still happen. We are doing it explicitly here
@@ -1786,7 +1786,7 @@ def reduction_dtypes(
 # batched_matrix_contiguous_strides and contiguous_strides
 def make_contiguous_strides_for(
     shape: ShapeType, row_major: bool = True
-) -> tuple[Union[_IntLikeT, int], ...]:
+) -> tuple[_IntLikeT | int, ...]:
     """
     Returns the strides of a contiguous tensor if row_major
     If row_major=True, it returns the strides of a contiguous batch of Fortran-contiguous matrices
@@ -1799,7 +1799,7 @@ def make_contiguous_strides_for(
 
     from torch.fx.experimental.symbolic_shapes import is_nested_int
 
-    multiplier: Union[_IntLikeT, int] = 1
+    multiplier: _IntLikeT | int = 1
     strides = []
     for l in reversed(shape):
         strides.append(multiplier)
@@ -1819,14 +1819,14 @@ def make_contiguous_strides_for(
 
 def make_channels_last_1d_strides_for(
     shape: Sequence[_IntLikeT],
-) -> tuple[Union[_IntLikeT, int], ...]:
+) -> tuple[_IntLikeT | int, ...]:
     torch._check(
         len(shape) == 3,
         lambda: "Only tensors of rank 3 can use the channels_last_1d memory format",
     )
 
-    multiplier: Union[_IntLikeT, int] = 1
-    strides: list[Union[_IntLikeT, int]] = [0] * 3
+    multiplier: _IntLikeT | int = 1
+    strides: list[_IntLikeT | int] = [0] * 3
     for idx in (1, -1, 0):
         # NOTE: intentionally divergence from make_contiguous_strides_for
         # This is consistent with eager
@@ -1838,15 +1838,15 @@ def make_channels_last_1d_strides_for(
 
 def make_channels_last_2d_strides_for(
     shape: Sequence[_IntLikeT],
-) -> tuple[Union[_IntLikeT, int], ...]:
+) -> tuple[_IntLikeT | int, ...]:
     # TODO: maybe inform the user of channels_last_3d if rank of the tensor is 5?
     torch._check(
         len(shape) == 4,
         lambda: "Only tensors of rank 4 can use the channels_last memory format",
     )
 
-    multiplier: Union[_IntLikeT, int] = 1
-    strides: list[Union[_IntLikeT, int]] = [0] * 4
+    multiplier: _IntLikeT | int = 1
+    strides: list[_IntLikeT | int] = [0] * 4
     for idx in (1, -1, -2, 0):
         # NOTE: intentionally divergence from make_contiguous_strides_for
         # This is consistent with eager
@@ -1858,14 +1858,14 @@ def make_channels_last_2d_strides_for(
 
 def make_channels_last_3d_strides_for(
     shape: Sequence[_IntLikeT],
-) -> tuple[Union[_IntLikeT, int], ...]:
+) -> tuple[_IntLikeT | int, ...]:
     torch._check(
         len(shape) == 5,
         lambda: "Only tensors of rank 5 can use the channels_last_3d memory format",
     )
 
-    multiplier: Union[_IntLikeT, int] = 1
-    strides: list[Union[_IntLikeT, int]] = [0] * 5
+    multiplier: _IntLikeT | int = 1
+    strides: list[_IntLikeT | int] = [0] * 5
     for idx in (1, -1, -2, -3, 0):
         # NOTE: intentionally divergence from make_contiguous_strides_for
         # This is consistent with eager
@@ -1877,7 +1877,7 @@ def make_channels_last_3d_strides_for(
 
 def make_channels_last_strides_for(
     shape: Sequence[_IntLikeT],
-) -> tuple[Union[_IntLikeT, int], ...]:
+) -> tuple[_IntLikeT | int, ...]:
     ndim = len(shape) if isinstance(shape, Sequence) else 1
     if ndim == 3:
         return make_channels_last_1d_strides_for(shape)
@@ -1912,7 +1912,7 @@ def validate_no_repeating_dims(dims: Sequence):
         raise RuntimeError("duplicate value in the list of dims")
 
 
-def reduction_dims(shape: ShapeType, dims: Optional[Sequence]) -> tuple[int, ...]:
+def reduction_dims(shape: ShapeType, dims: Sequence | None) -> tuple[int, ...]:
     if dims is None:
         return tuple(range(len(shape)))
     dims = tuple(canonicalize_dim(len(shape), idx) for idx in dims)
@@ -1921,8 +1921,8 @@ def reduction_dims(shape: ShapeType, dims: Optional[Sequence]) -> tuple[int, ...
 
 
 def set_correction(
-    unbiased: Optional[bool] = None,
-    correction: Optional[NumberType] = None,
+    unbiased: bool | None = None,
+    correction: NumberType | None = None,
 ) -> float:
     if correction is not None and unbiased is not None:
         raise RuntimeError("cannot specify both correction and unbiased arguments")
@@ -2128,15 +2128,15 @@ def get_aten_op(fn: Callable, name: str):
     return getattr(torch._ops.ops.aten, f"{module}{name}")
 
 
-def dtype_or_default(dtype: Optional[torch.dtype]) -> torch.dtype:
+def dtype_or_default(dtype: torch.dtype | None) -> torch.dtype:
     return dtype if dtype is not None else torch.get_default_dtype()
 
 
-def device_or_default(device: Optional[DeviceLikeType]) -> DeviceLikeType:
+def device_or_default(device: DeviceLikeType | None) -> DeviceLikeType:
     return device if device is not None else torch.device("cpu")
 
 
-def layout_or_default(layout: Optional[torch.layout]) -> torch.layout:
+def layout_or_default(layout: torch.layout | None) -> torch.layout:
     return layout if layout is not None else torch.strided
 
 
diff --git a/torch/_prims_common/wrappers.py b/torch/_prims_common/wrappers.py
index 2ac116b3ca685..17083fb30a9f4 100644
--- a/torch/_prims_common/wrappers.py
+++ b/torch/_prims_common/wrappers.py
@@ -5,7 +5,7 @@
 from collections.abc import Callable, Sequence
 from functools import wraps
 from types import GenericAlias
-from typing import NamedTuple, Optional, overload, TypeVar, Union
+from typing import NamedTuple, overload, TypeVar
 from typing_extensions import ParamSpec
 
 import torch
@@ -20,6 +20,7 @@
     TensorLikeType,
 )
 from torch.utils import _pytree as pytree
+from torch.utils._inspect import _fast_bind
 from torch.utils._pytree import tree_flatten, tree_unflatten
 
 
@@ -116,7 +117,7 @@ def __init__(
         self,
         *,
         type_promotion_kind: ELEMENTWISE_TYPE_PROMOTION_KIND,
-        type_promoting_args: Optional[Sequence[str]] = None,
+        type_promoting_args: Sequence[str] | None = None,
     ):
         self.type_promoting_arg_names = type_promoting_args
         self.type_promotion_kind = type_promotion_kind
@@ -129,7 +130,7 @@ def __call__(self, fn: Callable) -> Callable:
         @torch._disable_dynamo
         @wraps(fn)
         def _fn(*args, **kwargs):
-            bound = sig.bind(*args, **kwargs)
+            bound = _fast_bind(sig, *args, **kwargs)
             type_promoting_args = tuple(
                 bound.arguments[x]
                 for x in self.type_promoting_arg_names  # type: ignore[union-attr]
@@ -188,7 +189,7 @@ def _resize_output_check(out: TensorLikeType, shape: ShapeType):
 def _maybe_resize_out(
     out: TensorLikeType,
     shape: ShapeType,
-    memory_format: Optional[torch.memory_format] = None,
+    memory_format: torch.memory_format | None = None,
 ):
     if _resize_output_check(out, shape):
         return out.resize_(shape, memory_format=memory_format)
@@ -483,7 +484,7 @@ def _autograd_impl(*args, **kwargs):
 # TODO: this wrapper is currently untested
 def elementwise_unary_scalar_wrapper(
     fn: Callable[_P, _T],
-) -> Callable[_P, Union[_T, NumberType]]:
+) -> Callable[_P, _T | NumberType]:
     """
     Allows unary operators that accept tensors to work with Python numbers.
     """
diff --git a/torch/_refs/__init__.py b/torch/_refs/__init__.py
index c97593a9a6585..7c089a68a57ec 100644
--- a/torch/_refs/__init__.py
+++ b/torch/_refs/__init__.py
@@ -50,6 +50,7 @@
     elementwise_unary_scalar_wrapper,
     out_wrapper,
 )
+from torch.testing._internal.common_dtype import highest_precision_float
 
 
 # Experimental module containing prototype Python references for existing
@@ -387,9 +388,9 @@ def handle_noncontiguous_outputs(input_tlist, output):
 def _broadcast_shapes(*_shapes):
     from torch.fx.experimental.symbolic_shapes import (
         guard_or_false,
-        has_hint,
+        guarding_hint_or_throw,
+        has_guarding_hint,
         is_nested_int,
-        size_hint,
     )
 
     backed_so = torch.fx.experimental._config.backed_size_oblivious
@@ -406,12 +407,11 @@ def _broadcast_shapes(*_shapes):
     for shape in shapes:
         if not isinstance(shape, Sequence):
             raise RuntimeError(
-                "Input shapes should be of type ints, a tuple of ints, or a list of ints, got ",
-                shape,
+                f"Input shapes should be of type ints, a tuple of ints, or a list of ints, got {shape}"
             )
 
     # Computes common shape
-    common_shape: list[Union[int, torch.SymInt]] = [
+    common_shape: list[int | torch.SymInt] = [
         1,
     ] * reduce(max, (len(shape) for shape in shapes))
     for arg_idx, shape in enumerate(shapes):
@@ -434,9 +434,13 @@ def _broadcast_shapes(*_shapes):
                 #           specialize(s0) to be 1.
                 #      s0:4, s1:1 ==>
                 #           specialize(s1) to be 1.
-                if backed_so and has_hint(shape[idx]) and has_hint(common_shape[idx]):
-                    a = size_hint(shape[idx])
-                    b = size_hint(common_shape[idx])
+                if (
+                    backed_so
+                    and has_guarding_hint(shape[idx])
+                    and has_guarding_hint(common_shape[idx])
+                ):
+                    a = guarding_hint_or_throw(shape[idx])
+                    b = guarding_hint_or_throw(common_shape[idx])
                     if a == 1 and b != 1:
                         torch._check(shape[idx] == 1)
                     if b == 1 and a != 1:
@@ -886,7 +890,7 @@ def log10(a):
 def log_softmax(
     a: TensorLikeType,
     dim: int,
-    dtype: Optional[torch.dtype] = None,
+    dtype: torch.dtype | None = None,
 ) -> TensorLikeType:
     result_dtype = dtype or a.dtype
     computation_dtype = utils.get_computation_dtype(result_dtype)
@@ -921,9 +925,9 @@ def logsumexp(
 @out_wrapper()
 def nan_to_num(
     a: TensorLikeType,
-    nan: Optional[NumberType] = 0.0,
-    posinf: Optional[NumberType] = None,
-    neginf: Optional[NumberType] = None,
+    nan: NumberType | None = 0.0,
+    posinf: NumberType | None = None,
+    neginf: NumberType | None = None,
 ) -> TensorLikeType:
     if not isinstance(a, TensorLike):
         raise AssertionError(f"a must be TensorLike, got {type(a)}")
@@ -1149,8 +1153,8 @@ def inner(prim: Callable):
             type_promotion_kind=type_promotion_kind,
         )
         def _ref(
-            a: Union[Tensor, NumberType],
-            b: Union[Tensor, NumberType],
+            a: Tensor | NumberType,
+            b: Tensor | NumberType,
         ) -> Tensor:
             torch._check_value(
                 supports_lhs_python_scalar or not isinstance(a, Number),
@@ -1185,6 +1189,16 @@ def _ref(
     return inner
 
 
+def _binary_op_dtype(
+    a: TensorLikeType | NumberType, b: TensorLikeType | NumberType
+) -> torch.dtype:
+    if isinstance(a, TensorLike):
+        return a.dtype
+    if isinstance(b, TensorLike):
+        return b.dtype
+    return utils.type_to_dtype(type(a))
+
+
 # Add has its own implementation because it has an alpha argument
 @register_decomposition(aten.add)
 @out_wrapper()
@@ -1193,10 +1207,10 @@ def _ref(
     type_promotion_kind=ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT,
 )
 def add(
-    a: Union[TensorLikeType, NumberType],
-    b: Union[TensorLikeType, NumberType],
+    a: TensorLikeType | NumberType,
+    b: TensorLikeType | NumberType,
     *,
-    alpha: Optional[NumberType] = None,
+    alpha: NumberType | None = None,
 ):
     """
     Reference implementation of torch.add
@@ -1205,7 +1219,7 @@ def add(
     a, b = _maybe_broadcast(a, b)
 
     if alpha is not None:
-        dtype = a.dtype if isinstance(a, TensorLike) else b.dtype  # type: ignore[union-attr]
+        dtype = _binary_op_dtype(a, b)
         python_type = utils.dtype_to_type(dtype)
         if python_type is not bool and not utils.is_weakly_lesser_type(
             type(alpha), python_type
@@ -1269,9 +1283,7 @@ def bitwise_xor(a: TensorLikeType, b: TensorLikeType) -> TensorLikeType:
     type_promotion_kind=ELEMENTWISE_TYPE_PROMOTION_KIND.INT_TO_FLOAT,
     supports_lhs_python_scalar=False,
 )
-def copysign(
-    a: Union[TensorLikeType, NumberType], b: Union[TensorLikeType, NumberType]
-):
+def copysign(a: TensorLikeType | NumberType, b: TensorLikeType | NumberType):
     if isinstance(b, Number) and isinstance(a, Tensor):
         # pyrefly: ignore [bad-argument-type]
         b = scalar_tensor(b, dtype=a.dtype, device=a.device)
@@ -1288,10 +1300,10 @@ def copysign(
 @register_decomposition(aten.div)
 @out_wrapper()
 def div(
-    a: Union[TensorLikeType, NumberType],
-    b: Union[TensorLikeType, NumberType],
+    a: TensorLikeType | NumberType,
+    b: TensorLikeType | NumberType,
     *,
-    rounding_mode: Optional[str] = None,
+    rounding_mode: str | None = None,
 ):
     """
     Reference implementation of torch.div
@@ -1319,8 +1331,8 @@ def eq(a: TensorLikeType, b: TensorLikeType) -> TensorLikeType:
     type_promotion_kind=ELEMENTWISE_TYPE_PROMOTION_KIND.BOOL_TO_LONG,
 )
 def pow(
-    a: Union[TensorLikeType, NumberType],
-    b: Union[TensorLikeType, NumberType],
+    a: TensorLikeType | NumberType,
+    b: TensorLikeType | NumberType,
 ) -> TensorLikeType:
     if not (isinstance(a, TensorLikeType) or isinstance(b, TensorLikeType)):
         raise AssertionError("at least one of a or b must be TensorLikeType")
@@ -1347,8 +1359,8 @@ def pow(
 # CompositeImplicitAutograd - don't register decomp
 @out_wrapper()
 def float_power(
-    a: Union[TensorLikeType, NumberType],
-    b: Union[TensorLikeType, NumberType],
+    a: TensorLikeType | NumberType,
+    b: TensorLikeType | NumberType,
 ) -> Tensor:
     if isinstance(a, Number) and isinstance(b, Number):
         raise ValueError(
@@ -1409,9 +1421,7 @@ def float_power(
     supports_two_python_scalars=True,
     should_register_decomposition=False,
 )
-def floor_divide(
-    a: Union[TensorLikeType, NumberType], b: Union[TensorLikeType, NumberType]
-):
+def floor_divide(a: TensorLikeType | NumberType, b: TensorLikeType | NumberType):
     # Wrap scalars because some references only accept tensor arguments.
     if isinstance(a, Number) and isinstance(b, Number):
         # pyrefly: ignore [bad-argument-type]
@@ -1818,8 +1828,8 @@ def remainder(a: TensorLikeType, b: TensorLikeType) -> TensorLikeType:
 @register_decomposition(aten.rsub)
 @out_wrapper()
 def rsub(
-    a: Union[TensorLikeType, NumberType],
-    b: Union[TensorLikeType, NumberType],
+    a: TensorLikeType | NumberType,
+    b: TensorLikeType | NumberType,
     alpha: NumberType = 1,
 ):
     if isinstance(a, Number):
@@ -1838,8 +1848,8 @@ def rsub(
     type_promotion_kind=ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT,
 )
 def sub(
-    a: Union[TensorLikeType, NumberType],
-    b: Union[TensorLikeType, NumberType],
+    a: TensorLikeType | NumberType,
+    b: TensorLikeType | NumberType,
     *,
     alpha: NumberType = 1,
 ):
@@ -1859,7 +1869,7 @@ def sub(
         )
 
     if alpha != 1:
-        dtype = a.dtype if isinstance(a, TensorLike) else b.dtype  # type: ignore[union-attr]
+        dtype = _binary_op_dtype(a, b)
         python_type = utils.dtype_to_type(dtype)
         if not utils.is_weakly_lesser_type(type(alpha), python_type):
             msg = f"alpha argument of type {type(alpha)} cannot be safely cast to type {python_type}!"
@@ -1892,7 +1902,7 @@ def true_divide(a: TensorLikeType, b: TensorLikeType) -> TensorLikeType:
     type_promoting_args=("a", "b"),
     type_promotion_kind=ELEMENTWISE_TYPE_PROMOTION_KIND.INT_TO_FLOAT,
 )
-def xlogy(a: Union[TensorLikeType, NumberType], b: Union[TensorLikeType, NumberType]):
+def xlogy(a: TensorLikeType | NumberType, b: TensorLikeType | NumberType):
     torch._check(
         isinstance(a, TensorLike) or isinstance(b, TensorLike),
         lambda: 'Expected either argument a or b to be a Tensor"',
@@ -1920,9 +1930,7 @@ def xlogy(a: Union[TensorLikeType, NumberType], b: Union[TensorLikeType, NumberT
     aten_op=None,  # CompositeImplicitAutograd
     supports_two_python_scalars=True,
 )
-def trunc_divide(
-    a: Union[TensorLikeType, NumberType], b: Union[TensorLikeType, NumberType]
-):
+def trunc_divide(a: TensorLikeType | NumberType, b: TensorLikeType | NumberType):
     dtype = utils.get_dtype(a)
     if utils.is_integer_dtype(dtype):
         return prims.div(a, b)
@@ -1997,13 +2005,14 @@ def addcmul(
 )
 def clamp(
     a: TensorLikeType,
-    min: Optional[TensorOrNumberLikeType] = None,
-    max: Optional[TensorOrNumberLikeType] = None,
+    min: TensorOrNumberLikeType | None = None,
+    max: TensorOrNumberLikeType | None = None,
 ) -> TensorLikeType:
     # NOTE: grad behavior with implementation `where` is not consistent on `nan`
     if min is None and max is None:
         msg = "clamp called but both min and max are none!"
         raise ValueError(msg)
+
     if min is not None:
         a_isnan = torch.isnan(a)
         condition = torch.bitwise_or(torch.ge(a, min), a_isnan)  # type: ignore[arg-type]
@@ -2025,7 +2034,7 @@ def clamp(
 @out_wrapper()
 def clamp_min(
     self: TensorLikeType,
-    min: Optional[TensorOrNumberLikeType] = None,
+    min: TensorOrNumberLikeType | None = None,
 ) -> TensorLikeType:
     return torch.clamp(self, min=min)  # type: ignore[arg-type]
 
@@ -2034,7 +2043,7 @@ def clamp_min(
 @out_wrapper()
 def clamp_max(
     self: TensorLikeType,
-    max: Optional[TensorOrNumberLikeType] = None,
+    max: TensorOrNumberLikeType | None = None,
 ) -> TensorLikeType:
     return torch.clamp(self, max=max)  # type: ignore[arg-type]
 
@@ -2045,7 +2054,11 @@ def clamp_max(
 
 
 # https://pytorch.org/docs/stable/generated/torch.where.html
-# TODO: implement where.default
+@register_decomposition(aten.where.default)
+def _where_default(pred: Tensor) -> tuple[Tensor, ...]:
+    return torch.nonzero(pred, as_tuple=True)
+
+
 @register_decomposition(aten.where.self)
 @register_decomposition(aten.where.ScalarSelf)
 @register_decomposition(aten.where.ScalarOther)
@@ -2058,8 +2071,8 @@ def clamp_max(
 )
 def where(
     pred: Tensor,
-    a: Optional[TensorOrNumberLikeType] = None,
-    b: Optional[TensorOrNumberLikeType] = None,
+    a: TensorOrNumberLikeType | None = None,
+    b: TensorOrNumberLikeType | None = None,
 ):
     """ """
 
@@ -2111,12 +2124,12 @@ def item(a: TensorLikeType) -> NumberType:
 # fast path when `to` returns an alias to input. This mimics the same function in aten
 def _to_will_alias(
     a: TensorLikeType,
-    device: Optional[DeviceLikeType] = None,
-    dtype: Optional[torch.dtype] = None,
-    copy: Optional[bool] = None,
-    layout: Optional[torch.layout] = None,
-    memory_format: Optional[torch.memory_format] = None,
-    pin_memory: Optional[bool] = False,
+    device: DeviceLikeType | None = None,
+    dtype: torch.dtype | None = None,
+    copy: bool | None = None,
+    layout: torch.layout | None = None,
+    memory_format: torch.memory_format | None = None,
+    pin_memory: bool | None = False,
     non_blocking: bool = False,  # not using non_blocking
 ) -> bool:
     return (
@@ -2145,7 +2158,7 @@ def _to_device(
     dtype: torch.dtype,
     non_blocking: bool = False,
     copy: bool = False,
-    memory_format: Optional[torch.memory_format] = None,
+    memory_format: torch.memory_format | None = None,
 ) -> dict[str, Any]:
     kwargs = {
         "device": device,
@@ -2163,7 +2176,7 @@ def _to_device_str(
     dtype: torch.dtype,
     non_blocking: bool = False,
     copy: bool = False,
-    memory_format: Optional[torch.memory_format] = None,
+    memory_format: torch.memory_format | None = None,
 ) -> dict[str, Any]:
     kwargs = {
         "device": torch.device(device),
@@ -2180,7 +2193,7 @@ def _to_dtype(
     dtype: torch.dtype,
     non_blocking: bool = False,
     copy: bool = False,
-    memory_format: Optional[torch.memory_format] = None,
+    memory_format: torch.memory_format | None = None,
 ) -> dict[str, Any]:
     kwargs = {
         "dtype": dtype,
@@ -2196,7 +2209,7 @@ def _to_other(
     other: Tensor,
     non_blocking: bool = False,
     copy: bool = False,
-    memory_format: Optional[torch.memory_format] = None,
+    memory_format: torch.memory_format | None = None,
 ) -> dict[str, Any]:
     device = other.device
     dtype = other.dtype
@@ -2285,10 +2298,10 @@ def _reduction(
     *,
     has_identity: bool = True,
     accepts_dim_tuple: bool = True,  # to handle min/argmin that accept single dim only
-    dims: Optional[DimsType] = None,
+    dims: DimsType | None = None,
     keepdims: bool = False,
-    dtype: Optional[torch.dtype] = None,  # should be specified for ops that support it
-    out: Optional[Tensor] = None,
+    dtype: torch.dtype | None = None,  # should be specified for ops that support it
+    out: Tensor | None = None,
     output_dtype_kind: REDUCTION_OUTPUT_TYPE_KIND,
 ) -> TensorLikeType:  # it is usually SAME, but I want
     # ref writers to actually think about what to put here
@@ -2383,7 +2396,7 @@ def _fn(*args, out=None, **kwargs):
 @out_wrapper()
 def all(
     a: TensorLikeType,
-    dim: Optional[DimsType] = None,
+    dim: DimsType | None = None,
     keepdim: bool = False,
 ) -> TensorLikeType:
     result = torch.logical_not(torch.any(torch.logical_not(a), dim, keepdim=keepdim))
@@ -2398,7 +2411,7 @@ def all(
 @out_wrapper()
 def any(
     a: TensorLikeType,
-    dim: Optional[DimsType] = None,
+    dim: DimsType | None = None,
     keepdim: bool = False,
 ) -> TensorLikeType:
     a_ = _maybe_convert_to_dtype(a, torch.bool)
@@ -2417,11 +2430,11 @@ def any(
 @register_decomposition([aten.sum.dim_IntList, aten.sum.IntList_out])
 def sum(
     a: TensorLikeType,
-    dim: Union[Optional[int], Optional[list[int]]] = None,
+    dim: int | None | list[int] = None,
     keepdim: bool = False,
     *,
-    dtype: Optional[torch.dtype] = None,
-    out: Optional[Tensor] = None,
+    dtype: torch.dtype | None = None,
+    out: Tensor | None = None,
 ) -> TensorLikeType:
     if dtype is None:
         if out is not None:
@@ -2469,11 +2482,11 @@ def sum_to_size(
 @register_decomposition(aten.prod)
 def prod(
     a: TensorLikeType,
-    dim: Union[Optional[int], Optional[list[int]]] = None,
+    dim: int | None | list[int] = None,
     keepdim: bool = False,
     *,
     dtype=None,
-    out: Optional[Tensor] = None,
+    out: Tensor | None = None,
 ) -> TensorLikeType:
     if dtype is None:
         if out is not None:
@@ -2499,10 +2512,10 @@ def prod(
 @register_decomposition(aten.amin)
 def amin(
     a: TensorLikeType,
-    dim: Optional[DimsType] = None,
+    dim: DimsType | None = None,
     keepdim: bool = False,
     *,
-    out: Optional[Tensor] = None,
+    out: Tensor | None = None,
 ) -> TensorLikeType:
     # reduces over all dimensions if dim=() is passed
     if dim == () or dim == []:
@@ -2523,10 +2536,10 @@ def amin(
 @register_decomposition(aten.amax)
 def amax(
     a: TensorLikeType,
-    dim: Optional[DimsType] = None,
+    dim: DimsType | None = None,
     keepdim: bool = False,
     *,
-    out: Optional[Tensor] = None,
+    out: Tensor | None = None,
 ) -> TensorLikeType:
     # reduces over all dimensions if dim=() is passed
     if dim == () or dim == []:
@@ -2558,11 +2571,11 @@ def _dim_var_dispatch(dim=None, unbiased=None):
 @out_wrapper()
 def var(
     a: TensorLikeType,
-    dim: Optional[DimsType] = None,
-    unbiased: Optional[bool] = None,
+    dim: DimsType | None = None,
+    unbiased: bool | None = None,
     keepdim: bool = False,
     *,
-    correction: Optional[NumberType] = None,
+    correction: NumberType | None = None,
 ) -> TensorLikeType:
     dim, unbiased = _dim_var_dispatch(dim, unbiased)
     correction = utils.set_correction(unbiased, correction)
@@ -2587,11 +2600,11 @@ def var(
 @out_wrapper()
 def std(
     a: TensorLikeType,
-    dim: Union[Optional[int], Optional[list[int]]] = None,
-    unbiased: Optional[bool] = None,
+    dim: int | None | list[int] = None,
+    unbiased: bool | None = None,
     keepdim: bool = False,
     *,
-    correction: Optional[NumberType] = None,
+    correction: NumberType | None = None,
 ) -> TensorLikeType:
     dim, unbiased = _dim_var_dispatch(dim, unbiased)
     correction = utils.set_correction(unbiased, correction)
@@ -2610,7 +2623,7 @@ def std(
 @register_decomposition(aten.mean)
 def mean(
     a: TensorLikeType,
-    dim: Optional[DimsType] = None,
+    dim: DimsType | None = None,
     keepdim: bool = False,
     *,
     dtype=None,
@@ -2658,11 +2671,11 @@ def mean(
 @out_wrapper("out0", "out1")
 def std_mean(
     a: TensorLikeType,
-    dim: Optional[DimsType] = None,
+    dim: DimsType | None = None,
     *,
-    unbiased: Optional[bool] = None,
+    unbiased: bool | None = None,
     keepdim: bool = False,
-    correction: Optional[NumberType] = None,
+    correction: NumberType | None = None,
 ):
     dim, unbiased = _dim_var_dispatch(dim, unbiased)
     correction = utils.set_correction(unbiased, correction)
@@ -2685,11 +2698,11 @@ def std_mean(
 @out_wrapper("out0", "out1")
 def var_mean(
     a: TensorLikeType,
-    dim: Optional[DimsType] = None,
-    unbiased: Optional[bool] = None,
+    dim: DimsType | None = None,
+    unbiased: bool | None = None,
     keepdim: bool = False,
     *,
-    correction: Optional[NumberType] = None,
+    correction: NumberType | None = None,
 ):
     dim, unbiased = _dim_var_dispatch(dim, unbiased)
     v = var(a, dim, unbiased, keepdim, correction=correction)
@@ -2763,8 +2776,8 @@ def addr(
 
 # CompositeImplicitAutograd - don't register decomp
 def atleast_1d(
-    arg: Union[TensorLikeType, Sequence[TensorLikeType]], *args: TensorLikeType
-) -> Union[TensorLikeType, tuple[TensorLikeType, ...]]:
+    arg: TensorLikeType | Sequence[TensorLikeType], *args: TensorLikeType
+) -> TensorLikeType | tuple[TensorLikeType, ...]:
     """Reference implementation of :func:`torch.atleast_1d`."""
     if not args and isinstance(arg, collections.abc.Sequence):
         args_ = arg
@@ -2789,8 +2802,8 @@ def _unsqueeze_atleast(
 
 # CompositeImplicitAutograd - don't register decomp
 def atleast_2d(
-    arg: Union[TensorLikeType, Sequence[TensorLikeType]], *args: TensorLikeType
-) -> Union[TensorLikeType, tuple[TensorLikeType, ...]]:
+    arg: TensorLikeType | Sequence[TensorLikeType], *args: TensorLikeType
+) -> TensorLikeType | tuple[TensorLikeType, ...]:
     """Reference implementation of :func:`torch.atleast_2d`."""
     if not args and isinstance(arg, collections.abc.Sequence):
         args_ = arg
@@ -2805,8 +2818,8 @@ def atleast_2d(
 
 # CompositeImplicitAutograd - don't register decomp
 def atleast_3d(
-    arg: Union[TensorLikeType, Sequence[TensorLikeType]], *args: TensorLikeType
-) -> Union[TensorLikeType, tuple[TensorLikeType, ...]]:
+    arg: TensorLikeType | Sequence[TensorLikeType], *args: TensorLikeType
+) -> TensorLikeType | tuple[TensorLikeType, ...]:
     """Reference implementation of :func:`torch.atleast_3d`."""
     if not args and isinstance(arg, collections.abc.Sequence):
         args_ = arg
@@ -2823,7 +2836,7 @@ def as_strided(
     a: TensorLikeType,
     size: ShapeType,
     stride: StrideType,
-    storage_offset: Optional[int] = None,
+    storage_offset: int | None = None,
 ) -> TensorLikeType:
     storage_offset_int = (
         storage_offset if storage_offset is not None else a.storage_offset()
@@ -2838,7 +2851,7 @@ def as_strided_scatter(
     src: TensorLikeType,
     size: ShapeType,
     stride: StrideType,
-    storage_offset: Optional[int] = None,
+    storage_offset: int | None = None,
 ) -> TensorLikeType:
     storage_offset_int = 0 if storage_offset is None else storage_offset
     return prims.as_strided_scatter(input, src, size, stride, storage_offset_int)
@@ -3122,8 +3135,8 @@ def dstack(tensors: TensorSequenceType) -> TensorLikeType:
 def expand(a: Tensor, *shape, implicit: bool = False) -> Tensor:
     from torch.fx.experimental.symbolic_shapes import (
         guard_or_false,
-        has_hint,
-        size_hint,
+        guarding_hint_or_throw,
+        has_guarding_hint,
         sym_or,
     )
 
@@ -3167,9 +3180,13 @@ def expand(a: Tensor, *shape, implicit: bool = False) -> Tensor:
             #            The non-broadcast path is picked
             #      x:1, requested_length:4 ==>
             #           specialize(x) to be 1.
-            if backed_so and has_hint(x) and has_hint(requested_length):
-                x_hint = size_hint(x)
-                requested_hint = size_hint(requested_length)
+            if (
+                backed_so
+                and has_guarding_hint(x)
+                and has_guarding_hint(requested_length)
+            ):
+                x_hint = guarding_hint_or_throw(x)
+                requested_hint = guarding_hint_or_throw(requested_length)
                 if x_hint == 1 and requested_hint != 1:
                     torch._check(x == 1)
 
@@ -3264,7 +3281,7 @@ def flipud(a: TensorLikeType) -> TensorLikeType:
 
 # CompositeImplicitAutograd - don't register decomp
 def narrow(
-    a: TensorLikeType, dim: int, start: Union[int, TensorLikeType], length: int
+    a: TensorLikeType, dim: int, start: int | TensorLikeType, length: int
 ) -> TensorLikeType:
     # Supports Tensor overload that was added for XLA:
     # https://github.com/pytorch/pytorch/issues/31558
@@ -3338,8 +3355,8 @@ def _unsqueeze_multiple(x: TensorLikeType, dimensions: list[int]) -> TensorLikeT
 @register_decomposition(aten.native_group_norm.default)
 def native_group_norm(
     input: Tensor,
-    weight: Optional[Tensor],
-    bias: Optional[Tensor],
+    weight: Tensor | None,
+    bias: Tensor | None,
     batch_size: int,
     num_channels: int,
     flattened_inner_size: int,
@@ -3412,8 +3429,8 @@ def native_group_norm(
 def native_layer_norm(
     input: Tensor,
     normalized_shape: ShapeType,
-    weight: Optional[Tensor],
-    bias: Optional[Tensor],
+    weight: Tensor | None,
+    bias: Tensor | None,
     eps: float,
 ) -> tuple[Tensor, Tensor, Tensor]:
     from torch.fx.experimental.symbolic_shapes import sym_eq
@@ -3459,6 +3476,10 @@ def native_layer_norm(
         + ", but got input of size "
         + str(input.shape),
     )
+    torch._check(
+        not input.is_complex(),
+        lambda: "native_layer_norm does not support complex inputs",
+    )
 
     input = contiguous(input)
     if weight is not None:
@@ -3544,15 +3565,15 @@ def renorm(
 def stft(
     input: Tensor,
     n_fft: int,
-    hop_length: Optional[int] = None,
-    win_length: Optional[int] = None,
-    window: Optional[Tensor] = None,
+    hop_length: int | None = None,
+    win_length: int | None = None,
+    window: Tensor | None = None,
     center: bool = True,
     pad_mode: str = "reflect",
     normalized: bool = False,
-    onesided: Optional[bool] = None,
-    return_complex: Optional[bool] = None,
-    align_to_window: Optional[bool] = None,
+    onesided: bool | None = None,
+    return_complex: bool | None = None,
+    align_to_window: bool | None = None,
 ) -> Tensor:
     torch._check(
         window is None or window.device == input.device,
@@ -3661,13 +3682,13 @@ def stft(
 def istft(
     input: Tensor,
     n_fft: int,
-    hop_length: Optional[int] = None,
-    win_length: Optional[int] = None,
-    window: Optional[Tensor] = None,
+    hop_length: int | None = None,
+    win_length: int | None = None,
+    window: Tensor | None = None,
     center: bool = True,
     normalized: bool = False,
-    onesided: Optional[bool] = None,
-    length: Optional[int] = None,
+    onesided: bool | None = None,
+    length: int | None = None,
     return_complex=False,
 ) -> Tensor:
     torch._check(
@@ -4186,7 +4207,7 @@ def stack(tensors: TensorSequenceType, dim: int = 0) -> TensorLikeType:
 def softmax(
     a: TensorLikeType,
     dim: int,
-    dtype: Optional[torch.dtype] = None,
+    dtype: torch.dtype | None = None,
 ) -> TensorLikeType:
     result_dtype = dtype or a.dtype
     computation_dtype = utils.get_computation_dtype(result_dtype)
@@ -4269,14 +4290,14 @@ def index_copy_(x: TensorLike, dim: int, index: TensorLike, tensor: TensorLike):
 @register_decomposition(aten.index_fill)
 @out_wrapper()
 def index_fill(
-    x: TensorLike, dim: int, index: TensorLike, value: Union[NumberType, TensorLike]
+    x: TensorLike, dim: int, index: TensorLike, value: NumberType | TensorLike
 ):
     return _index_fill(x, dim, index, value, inplace=False)
 
 
 @register_decomposition(aten.index_fill_)
 def index_fill_(
-    x: TensorLike, dim: int, index: TensorLike, value: Union[NumberType, TensorLike]
+    x: TensorLike, dim: int, index: TensorLike, value: NumberType | TensorLike
 ):
     return _index_fill(x, dim, index, value, inplace=True)
 
@@ -4285,7 +4306,7 @@ def _index_fill(
     x: TensorLike,
     dim: int,
     index: TensorLike,
-    value: Union[NumberType, TensorLike],
+    value: NumberType | TensorLike,
     *,
     inplace: bool,
 ):
@@ -4322,11 +4343,12 @@ def _index_fill(
         if zero_dim:
             # The clone is necessary so that it returns a fresh tensor rather than a view
             out = out.squeeze(0).clone()
-        # index_fill preserves the strides. index_copy always returns contiguous tensors
-        if out.stride() != x.stride():
-            new_out = torch.empty_like(x)
-            new_out.copy_(out)
-            out = new_out
+        # index_fill preserves the strides for non-overlapping-and-dense inputs
+        # (matching clone(Preserve) behavior). index_copy always returns contiguous tensors.
+        if out.stride() != x.stride() and utils.is_non_overlapping_and_dense_or_false(
+            x
+        ):
+            out = prims.copy_strided(out, x.stride())
         return out
 
 
@@ -4364,11 +4386,11 @@ def index_select(x: TensorLike, dim: int, index: TensorLike):
         return torch.empty_like(x).index_copy(0, index, x.expand_as(index))
 
     idx = (slice(None),) * dim + (index,)
-    return x[idx].contiguous()
+    return x[idx].contiguous(memory_format=utils.suggest_memory_format(x))
 
 
 @register_decomposition(aten.squeeze.dims)
-def squeeze(a: TensorLikeType, dim: Optional[DimsType] = None) -> TensorLikeType:
+def squeeze(a: TensorLikeType, dim: DimsType | None = None) -> TensorLikeType:
     from torch.fx.experimental.symbolic_shapes import guard_or_false
 
     if dim is None:
@@ -4435,7 +4457,7 @@ def split_with_sizes(
 # CompositeImplicitAutograd - don't register decomp
 def tensor_split(
     a: TensorLikeType,
-    indices_or_sections: Union[Tensor, DimsType],
+    indices_or_sections: Tensor | DimsType,
     dim: int = 0,
 ) -> tuple[TensorLikeType, ...]:
     _dim = utils.canonicalize_dim(a.ndim, dim)
@@ -4868,8 +4890,8 @@ def _cumsumprod_common(
     a: TensorLikeType,
     dim: int,
     *,
-    dtype: Optional[torch.dtype] = None,
-    out: Optional[Tensor] = None,
+    dtype: torch.dtype | None = None,
+    out: Tensor | None = None,
 ) -> TensorLikeType:
     # We implement all the kwargs of a reduction. ATen just handles dtype
     # nb. This decomposition may not be as efficient as a backend-specific implementation
@@ -4891,8 +4913,8 @@ def cumsum(
     a: TensorLikeType,
     dim: int,
     *,
-    dtype: Optional[torch.dtype] = None,
-    out: Optional[Tensor] = None,
+    dtype: torch.dtype | None = None,
+    out: Tensor | None = None,
 ) -> TensorLikeType:
     return _cumsumprod_common(func=sum, init=0, a=a, dim=dim, dtype=dtype, out=out)
 
@@ -4902,8 +4924,8 @@ def cumprod(
     a: TensorLikeType,
     dim: int,
     *,
-    dtype: Optional[torch.dtype] = None,
-    out: Optional[Tensor] = None,
+    dtype: torch.dtype | None = None,
+    out: Tensor | None = None,
 ) -> TensorLikeType:
     return _cumsumprod_common(func=prod, init=1, a=a, dim=dim, dtype=dtype, out=out)
 
@@ -4933,7 +4955,7 @@ def view(a: TensorLikeType, *shape: ShapeType | tuple[ShapeType]) -> TensorLikeT
     # tuple[ShapeType | tuple[ShapeType], ...] but the function expects
     # Union[ShapeType, tuple[ShapeType]].
     shape_tuple = utils.extract_shape_from_varargs(
-        cast(Union[ShapeType, tuple[ShapeType]], shape), validate=False
+        cast(ShapeType | tuple[ShapeType], shape), validate=False
     )
     if torch.fx.experimental._config.backed_size_oblivious or _view_has_unbacked_input(
         a,
@@ -4957,7 +4979,7 @@ def ravel(a: TensorLikeType) -> TensorLikeType:
 # missing ref impl. for aten.gather
 @out_wrapper()
 def take_along_dim(
-    a: torch.Tensor, indices: torch.Tensor, dim: Optional[int] = None
+    a: torch.Tensor, indices: torch.Tensor, dim: int | None = None
 ) -> torch.Tensor:
     torch._check(
         a.ndim == indices.ndim,
@@ -4999,9 +5021,9 @@ def take_along_dim(
 @out_wrapper()
 def empty(
     *shape,
-    dtype: Optional[torch.dtype] = None,
+    dtype: torch.dtype | None = None,
     layout: torch.layout = torch.strided,
-    device: Optional[DeviceLikeType] = None,
+    device: DeviceLikeType | None = None,
     requires_grad: bool = False,
     pin_memory: bool = False,
     memory_format: torch.memory_format = torch.contiguous_format,
@@ -5041,9 +5063,9 @@ def empty(
 def empty_permuted(
     shape,
     physical_layout,
-    dtype: Optional[torch.dtype] = None,
+    dtype: torch.dtype | None = None,
     layout: torch.layout = torch.strided,
-    device: Optional[DeviceLikeType] = None,
+    device: DeviceLikeType | None = None,
     requires_grad: bool = False,
     pin_memory: bool = False,
 ) -> TensorLikeType:
@@ -5062,9 +5084,9 @@ def new_empty(
     a: TensorLikeType,
     size: ShapeType,
     *,
-    dtype: Optional[torch.dtype] = None,
-    layout: Optional[torch.layout] = None,
-    device: Optional[DeviceLikeType] = None,
+    dtype: torch.dtype | None = None,
+    layout: torch.layout | None = None,
+    device: DeviceLikeType | None = None,
     pin_memory: bool = False,
 ) -> TensorLikeType:
     dtype = a.dtype if dtype is None else dtype
@@ -5087,9 +5109,9 @@ def new_empty_strided(
     size: ShapeType,
     stride: StrideType,
     *,
-    dtype: Optional[torch.dtype] = None,
-    layout: Optional[torch.layout] = None,
-    device: Optional[DeviceLikeType] = None,
+    dtype: torch.dtype | None = None,
+    layout: torch.layout | None = None,
+    device: DeviceLikeType | None = None,
     pin_memory: bool = False,
 ) -> TensorLikeType:
     """
@@ -5114,9 +5136,9 @@ def new_empty_strided(
 @out_wrapper()
 def zeros(
     *size,
-    dtype: Optional[torch.dtype] = None,
+    dtype: torch.dtype | None = None,
     layout: torch.layout = torch.strided,
-    device: Optional[DeviceLikeType] = None,
+    device: DeviceLikeType | None = None,
     pin_memory: bool = False,
     requires_grad: bool = False,
 ) -> TensorLikeType:
@@ -5142,9 +5164,9 @@ def new_zeros(
     a: TensorLikeType,
     size: ShapeType,
     *,
-    dtype: Optional[torch.dtype] = None,
-    layout: Optional[torch.layout] = None,
-    device: Optional[DeviceLikeType] = None,
+    dtype: torch.dtype | None = None,
+    layout: torch.layout | None = None,
+    device: DeviceLikeType | None = None,
     pin_memory: bool = False,
     requires_grad: bool = False,
 ) -> TensorLikeType:
@@ -5167,9 +5189,9 @@ def new_zeros(
 @out_wrapper()
 def ones(
     *size,
-    dtype: Optional[torch.dtype] = None,
+    dtype: torch.dtype | None = None,
     layout: torch.layout = torch.strided,
-    device: Optional[DeviceLikeType] = None,
+    device: DeviceLikeType | None = None,
     pin_memory: bool = False,
     requires_grad: bool = False,
 ) -> TensorLikeType:
@@ -5195,9 +5217,9 @@ def new_ones(
     a: TensorLikeType,
     size: ShapeType,
     *,
-    dtype: Optional[torch.dtype] = None,
-    layout: Optional[torch.layout] = None,
-    device: Optional[DeviceLikeType] = None,
+    dtype: torch.dtype | None = None,
+    layout: torch.layout | None = None,
+    device: DeviceLikeType | None = None,
     pin_memory: bool = False,
     requires_grad: bool = False,
 ) -> TensorLikeType:
@@ -5223,9 +5245,9 @@ def new_full(
     size: ShapeType,
     fill_value: NumberType,
     *,
-    dtype: Optional[torch.dtype] = None,
-    layout: Optional[torch.layout] = None,
-    device: Optional[DeviceLikeType] = None,
+    dtype: torch.dtype | None = None,
+    layout: torch.layout | None = None,
+    device: DeviceLikeType | None = None,
     pin_memory: bool = False,
 ) -> TensorLikeType:
     dtype = a.dtype if dtype is None else dtype
@@ -5246,7 +5268,7 @@ def new_full(
 def empty_out(
     size: TensorLikeType,
     out: TensorLikeType,
-    memory_format: Optional[torch.memory_format] = None,
+    memory_format: torch.memory_format | None = None,
 ) -> TensorLikeType:
     return out
 
@@ -5256,9 +5278,9 @@ def empty_out(
 def empty_like(
     a: TensorLikeType,
     *,
-    dtype: Optional[torch.dtype] = None,
-    device: Optional[DeviceLikeType] = None,
-    layout: Optional[torch.layout] = None,
+    dtype: torch.dtype | None = None,
+    device: DeviceLikeType | None = None,
+    layout: torch.layout | None = None,
     pin_memory: bool = False,
     requires_grad: bool = False,
     memory_format: torch.memory_format = torch.preserve_format,
@@ -5298,12 +5320,12 @@ def empty_like(
 @out_wrapper()
 def arange(
     start: NumberType = 0,
-    end: Optional[NumberType] = None,
+    end: NumberType | None = None,
     step: NumberType = 1,
     *,
-    dtype: Optional[torch.dtype] = None,
+    dtype: torch.dtype | None = None,
     layout: torch.layout = torch.strided,
-    device: Optional[DeviceLikeType] = None,
+    device: DeviceLikeType | None = None,
     pin_memory: bool = False,
     requires_grad: bool = False,
 ) -> TensorLikeType:
@@ -5404,7 +5426,7 @@ def is_finite(x):
     type_promoting_args=("start", "end", "weight"),
     type_promotion_kind=ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT,
 )
-def lerp(start: Tensor, end: Tensor, weight: Union[Tensor, NumberType]):
+def lerp(start: Tensor, end: Tensor, weight: Tensor | NumberType):
     inputs = [start, end]
     if isinstance(weight, Number):
         weight = start.new_full((), weight)  # type: ignore[arg-type]
@@ -5433,12 +5455,12 @@ def lerp(start: Tensor, end: Tensor, weight: Union[Tensor, NumberType]):
 @register_decomposition(aten.linspace)
 @out_wrapper()
 def linspace(
-    start: Union[NumberType, TensorLikeType],
-    end: Union[NumberType, TensorLikeType],
+    start: NumberType | TensorLikeType,
+    end: NumberType | TensorLikeType,
     steps: NumberType,
     *,
-    dtype: Optional[torch.dtype] = None,
-    device: Optional[DeviceLikeType] = None,
+    dtype: torch.dtype | None = None,
+    device: DeviceLikeType | None = None,
     layout: torch.layout = torch.strided,
     pin_memory: bool = False,
     requires_grad: bool = False,
@@ -5448,13 +5470,13 @@ def linspace(
             start.dim() == 0,
             lambda: "linspace only supports 0-dimensional start and end tensors",
         )
-        start = _maybe_convert_to_dtype(start, torch.float64)
+        start = _maybe_convert_to_dtype(start, highest_precision_float(device))
     if isinstance(end, TensorLikeType):
         torch._check(
             end.dim() == 0,
             lambda: "linspace only supports 0-dimensional start and end tensors",
         )
-        end = _maybe_convert_to_dtype(end, torch.float64)
+        end = _maybe_convert_to_dtype(end, highest_precision_float(device))
 
     if builtins.any(isinstance(arg, complex) for arg in (start, end, steps)):
         default_complex_dtype = utils.corresponding_complex_dtype(
@@ -5526,13 +5548,13 @@ def linspace(
 @register_decomposition(aten.logspace)
 @out_wrapper()
 def logspace(
-    start: Union[NumberType, TensorLikeType],
-    end: Union[NumberType, TensorLikeType],
+    start: NumberType | TensorLikeType,
+    end: NumberType | TensorLikeType,
     steps: NumberType,
     base: NumberType = 10,
     *,
-    dtype: Optional[torch.dtype] = None,
-    device: Optional[DeviceLikeType] = None,
+    dtype: torch.dtype | None = None,
+    device: DeviceLikeType | None = None,
     layout: torch.layout = torch.strided,
     pin_memory: bool = False,
     requires_grad: bool = False,
@@ -5598,7 +5620,7 @@ def meshgrid(*tensors: TensorLikeType, indexing: str):
 
 @register_decomposition(aten.meshgrid)  # type: ignore[misc]
 def meshgrid(
-    *tensors: Union[TensorLikeType, list[TensorLikeType], tuple[TensorLikeType]],
+    *tensors: TensorLikeType | list[TensorLikeType] | tuple[TensorLikeType],
     indexing: str,
 ) -> list[TensorLikeType]:
     # This ref simultaneously handles two overloads (see stubs above)
@@ -5670,8 +5692,8 @@ def meshgrid(
 # CompositeImplicitAutograd - don't register decomp
 def movedim(
     input: TensorLikeType,
-    source: Union[int, DimsSequenceType],
-    destination: Union[int, DimsSequenceType],
+    source: int | DimsSequenceType,
+    destination: int | DimsSequenceType,
 ) -> TensorLikeType:
     """
     Reference implementation of torch.movedim
@@ -5734,11 +5756,11 @@ def movedim(
 @register_decomposition(aten.empty_strided)
 @out_wrapper()
 def empty_strided(
-    shape: Union[ShapeType, tuple[ShapeType]],
+    shape: ShapeType | tuple[ShapeType],
     strides: StrideType,
     *,
-    dtype: Optional[torch.dtype] = None,
-    device: Optional[DeviceLikeType] = None,
+    dtype: torch.dtype | None = None,
+    device: DeviceLikeType | None = None,
     layout: torch.layout = torch.strided,
     requires_grad: bool = False,
     pin_memory: bool = False,
@@ -5771,11 +5793,11 @@ def _strength_reduce_integer(val: int) -> torch.dtype:
 @out_wrapper()
 def eye(
     n: int,
-    m: Optional[int] = None,
+    m: int | None = None,
     *,
-    dtype: Optional[torch.dtype] = None,
+    dtype: torch.dtype | None = None,
     layout: torch.layout = torch.strided,
-    device: Optional[DeviceLikeType] = None,
+    device: DeviceLikeType | None = None,
     pin_memory: bool = False,
     requires_grad: bool = False,  # TODO: unused
 ) -> TensorLikeType:
@@ -5818,9 +5840,9 @@ def full(
     shape: ShapeType,
     fill_value: NumberType,
     *,
-    dtype: Optional[torch.dtype] = None,
+    dtype: torch.dtype | None = None,
     layout: torch.layout = torch.strided,
-    device: Optional[DeviceLikeType] = None,
+    device: DeviceLikeType | None = None,
     pin_memory: bool = False,
     requires_grad: bool = False,
 ) -> TensorLikeType:
@@ -5845,9 +5867,9 @@ def full_like(
     a: TensorLikeType,
     fill_value: NumberType,
     *,
-    dtype: Optional[torch.dtype] = None,
-    layout: Optional[torch.layout] = None,
-    device: Optional[DeviceLikeType] = None,
+    dtype: torch.dtype | None = None,
+    layout: torch.layout | None = None,
+    device: DeviceLikeType | None = None,
     pin_memory: bool = False,
     requires_grad: bool = False,
     memory_format: torch.memory_format = torch.preserve_format,
@@ -5869,9 +5891,9 @@ def full_like(
 def zeros_like(
     a: TensorLikeType,
     *,
-    dtype: Optional[torch.dtype] = None,
-    layout: Optional[torch.layout] = None,
-    device: Optional[DeviceLikeType] = None,
+    dtype: torch.dtype | None = None,
+    layout: torch.layout | None = None,
+    device: DeviceLikeType | None = None,
     pin_memory: bool = False,
     requires_grad: bool = False,
     memory_format: torch.memory_format = torch.preserve_format,
@@ -5893,9 +5915,9 @@ def zeros_like(
 def ones_like(
     a: TensorLikeType,
     *,
-    dtype: Optional[torch.dtype] = None,
-    layout: Optional[torch.layout] = None,
-    device: Optional[DeviceLikeType] = None,
+    dtype: torch.dtype | None = None,
+    layout: torch.layout | None = None,
+    device: DeviceLikeType | None = None,
     pin_memory: bool = False,
     requires_grad: bool = False,
     memory_format: torch.memory_format = torch.preserve_format,
@@ -5916,9 +5938,9 @@ def ones_like(
 @out_wrapper()
 def randn(
     *shape,
-    dtype: Optional[torch.dtype] = None,
-    device: Optional[DeviceLikeType] = None,
-    layout: Optional[torch.layout] = None,
+    dtype: torch.dtype | None = None,
+    device: DeviceLikeType | None = None,
+    layout: torch.layout | None = None,
     requires_grad: bool = False,
     pin_memory: bool = False,
 ) -> TensorLikeType:
@@ -5942,9 +5964,9 @@ def randn(
 def scalar_tensor(
     a: NumberType,
     *,
-    dtype: Optional[torch.dtype] = None,
+    dtype: torch.dtype | None = None,
     layout: torch.layout = torch.strided,
-    device: Optional[DeviceLikeType] = None,
+    device: DeviceLikeType | None = None,
     pin_memory: bool = False,
 ) -> TensorLikeType:
     utils.check_layout(layout)
@@ -5961,8 +5983,8 @@ def scalar_tensor(
 
 def _uniform_helper(
     shape: ShapeType,
-    low: Union[bool, int, float] = 0.0,
-    high: Union[bool, int, float] = 1.0,
+    low: bool | int | float = 0.0,
+    high: bool | int | float = 1.0,
     *,
     stride: ShapeType,
     dtype: torch.dtype,
@@ -6083,11 +6105,11 @@ def equal(a: TensorLikeType, b: TensorLikeType) -> bool:
 @out_wrapper(exact_dtype=True)
 def norm(
     input: TensorLikeType,
-    p: Optional[Union[float, str]] = "fro",
-    dim: Optional[DimsType] = None,
+    p: float | str | None = "fro",
+    dim: DimsType | None = None,
     keepdim: bool = False,
     *,
-    dtype: Optional[torch.dtype] = None,
+    dtype: torch.dtype | None = None,
 ) -> TensorLikeType:
     # In these cases we compute the "Frobenius norm"
     if (
@@ -6118,8 +6140,8 @@ def trace(self: TensorLikeType) -> TensorLikeType:
 
 def _make_r_binary_op(base_op):
     def rop(
-        a: Union[TensorLikeType, NumberType],
-        b: Union[TensorLikeType, NumberType],
+        a: TensorLikeType | NumberType,
+        b: TensorLikeType | NumberType,
     ) -> TensorLikeType:
         return base_op(b, a)
 
@@ -6560,7 +6582,7 @@ def deg2rad(self: TensorLikeType):
 
 @register_decomposition(aten.count_nonzero)
 @out_wrapper()
-def count_nonzero(self, dim: Optional[DimsType] = None):
+def count_nonzero(self, dim: DimsType | None = None):
     return (self != 0).sum(dim)
 
 
@@ -6848,7 +6870,7 @@ def _infer_scalar_type(obj):
 # Analogous to recursive_store
 # xref: recursive_store in torch/csrc/utils/tensor_new.cpp
 def _recursive_build(
-    scalarType: torch.dtype, obj: Union[TensorOrNumberLikeType, TensorSequenceType]
+    scalarType: torch.dtype, obj: TensorOrNumberLikeType | TensorSequenceType
 ):
     if isinstance(obj, Tensor) and obj.numel() == 1:
         return obj.detach().to(dtype=scalarType, device="cpu", copy=True).view(())
diff --git a/torch/_refs/fft.py b/torch/_refs/fft.py
index 5bc074e55c74e..8de041aa56b22 100644
--- a/torch/_refs/fft.py
+++ b/torch/_refs/fft.py
@@ -1,6 +1,6 @@
 import math
 from collections.abc import Iterable, Sequence
-from typing import Literal, NamedTuple, Optional, Union
+from typing import Literal, NamedTuple
 
 import torch
 import torch._prims as prims
@@ -35,7 +35,7 @@
     "ifftshift",
 ]
 
-NormType = Union[None, Literal["forward", "backward", "ortho"]]
+NormType = None | Literal["forward", "backward", "ortho"]
 _NORM_VALUES = {None, "forward", "backward", "ortho"}
 aten = torch._ops.ops.aten
 
@@ -67,7 +67,7 @@ def _promote_type_fft(
         dtype = torch.get_default_dtype()
 
     allowed_types = [torch.float32, torch.float64]
-    maybe_support_half = device.type in ["cuda", "meta"]
+    maybe_support_half = device.type in ["cuda", "meta", "xpu"]
 
     if maybe_support_half:
         allowed_types.append(torch.float16)
@@ -121,7 +121,7 @@ def _resize_fft_input(
 def _fft_c2r(
     func_name: str,
     input: TensorLikeType,
-    n: Optional[int],
+    n: int | None,
     dim: int,
     norm: NormType,
     forward: bool,
@@ -148,7 +148,7 @@ def _fft_c2r(
 def _fft_r2c(
     func_name: str,
     input: TensorLikeType,
-    n: Optional[int],
+    n: int | None,
     dim: int,
     norm: NormType,
     forward: bool,
@@ -177,7 +177,7 @@ def _fft_r2c(
 def _fft_c2c(
     func_name: str,
     input: TensorLikeType,
-    n: Optional[int],
+    n: int | None,
     dim: int,
     norm: NormType,
     forward: bool,
@@ -204,7 +204,7 @@ def _fft_c2c(
 @out_wrapper()
 def fft(
     input: TensorLikeType,
-    n: Optional[int] = None,
+    n: int | None = None,
     dim: int = -1,
     norm: NormType = None,
 ) -> TensorLikeType:
@@ -218,7 +218,7 @@ def fft(
 @out_wrapper()
 def ifft(
     input: TensorLikeType,
-    n: Optional[int] = None,
+    n: int | None = None,
     dim: int = -1,
     norm: NormType = None,
 ) -> TensorLikeType:
@@ -232,7 +232,7 @@ def ifft(
 @out_wrapper()
 def rfft(
     input: TensorLikeType,
-    n: Optional[int] = None,
+    n: int | None = None,
     dim: int = -1,
     norm: NormType = None,
 ) -> TensorLikeType:
@@ -243,7 +243,7 @@ def rfft(
 @out_wrapper()
 def irfft(
     input: TensorLikeType,
-    n: Optional[int] = None,
+    n: int | None = None,
     dim: int = -1,
     norm: NormType = None,
 ) -> TensorLikeType:
@@ -254,7 +254,7 @@ def irfft(
 @out_wrapper()
 def hfft(
     input: TensorLikeType,
-    n: Optional[int] = None,
+    n: int | None = None,
     dim: int = -1,
     norm: NormType = None,
 ) -> TensorLikeType:
@@ -265,7 +265,7 @@ def hfft(
 @out_wrapper()
 def ihfft(
     input: TensorLikeType,
-    n: Optional[int] = None,
+    n: int | None = None,
     dim: int = -1,
     norm: NormType = None,
 ) -> TensorLikeType:
@@ -278,7 +278,7 @@ class _ShapeAndDims(NamedTuple):
 
 
 def _canonicalize_fft_shape_and_dim_args(
-    input: TensorLikeType, shape: Optional[ShapeType], dim: Optional[DimsType]
+    input: TensorLikeType, shape: ShapeType | None, dim: DimsType | None
 ) -> _ShapeAndDims:
     """Convert the shape and dim arguments into a canonical form where neither are optional"""
     input_dim = input.ndim
@@ -365,8 +365,8 @@ def _fftn_c2c(
 @out_wrapper()
 def fftn(
     input: TensorLikeType,
-    s: Optional[ShapeType] = None,
-    dim: Optional[DimsType] = None,
+    s: ShapeType | None = None,
+    dim: DimsType | None = None,
     norm: NormType = None,
 ) -> TensorLikeType:
     (shape, dim) = _canonicalize_fft_shape_and_dim_args(input, s, dim)
@@ -378,8 +378,8 @@ def fftn(
 @out_wrapper()
 def ifftn(
     input: TensorLikeType,
-    s: Optional[ShapeType] = None,
-    dim: Optional[DimsType] = None,
+    s: ShapeType | None = None,
+    dim: DimsType | None = None,
     norm: NormType = None,
 ) -> TensorLikeType:
     (shape, dim) = _canonicalize_fft_shape_and_dim_args(input, s, dim)
@@ -391,8 +391,8 @@ def ifftn(
 @out_wrapper()
 def rfftn(
     input: TensorLikeType,
-    s: Optional[ShapeType] = None,
-    dim: Optional[DimsType] = None,
+    s: ShapeType | None = None,
+    dim: DimsType | None = None,
     norm: NormType = None,
 ) -> TensorLikeType:
     torch._check(
@@ -410,8 +410,8 @@ def rfftn(
 @out_wrapper()
 def ihfftn(
     input: TensorLikeType,
-    s: Optional[ShapeType] = None,
-    dim: Optional[DimsType] = None,
+    s: ShapeType | None = None,
+    dim: DimsType | None = None,
     norm: NormType = None,
 ) -> TensorLikeType:
     torch._check(
@@ -443,8 +443,8 @@ class _CanonicalizeC2rReturn(NamedTuple):
 def _canonicalize_fft_c2r_shape_and_dim_args(
     fname: str,
     input: TensorLikeType,
-    s: Optional[ShapeType],
-    dim: Optional[DimsType],
+    s: ShapeType | None,
+    dim: DimsType | None,
 ) -> _CanonicalizeC2rReturn:
     """Canonicalize shape and dim arguments for n-dimensional c2r transforms,
     as well as calculating the last_dim_size which is shape[dim[-1]] for the output"""
@@ -472,8 +472,8 @@ def _canonicalize_fft_c2r_shape_and_dim_args(
 @out_wrapper()
 def irfftn(
     input: TensorLikeType,
-    s: Optional[ShapeType] = None,
-    dim: Optional[DimsType] = None,
+    s: ShapeType | None = None,
+    dim: DimsType | None = None,
     norm: NormType = None,
 ) -> TensorLikeType:
     shape, dim, last_dim_size = _canonicalize_fft_c2r_shape_and_dim_args(
@@ -489,8 +489,8 @@ def irfftn(
 @out_wrapper()
 def hfftn(
     input: TensorLikeType,
-    s: Optional[ShapeType] = None,
-    dim: Optional[DimsType] = None,
+    s: ShapeType | None = None,
+    dim: DimsType | None = None,
     norm: NormType = None,
 ) -> TensorLikeType:
     shape, dim, last_dim_size = _canonicalize_fft_c2r_shape_and_dim_args(
@@ -510,8 +510,8 @@ def hfftn(
 @out_wrapper()
 def fft2(
     input: TensorLikeType,
-    s: Optional[ShapeType] = None,
-    dim: Optional[DimsType] = (-2, -1),
+    s: ShapeType | None = None,
+    dim: DimsType | None = (-2, -1),
     norm: NormType = None,
 ) -> TensorLikeType:
     return torch.fft.fftn(input, s=s, dim=dim, norm=norm)
@@ -521,8 +521,8 @@ def fft2(
 @out_wrapper()
 def ifft2(
     input: TensorLikeType,
-    s: Optional[ShapeType] = None,
-    dim: Optional[DimsType] = (-2, -1),
+    s: ShapeType | None = None,
+    dim: DimsType | None = (-2, -1),
     norm: NormType = None,
 ) -> TensorLikeType:
     return torch.fft.ifftn(input, s=s, dim=dim, norm=norm)
@@ -532,8 +532,8 @@ def ifft2(
 @out_wrapper()
 def rfft2(
     input: TensorLikeType,
-    s: Optional[ShapeType] = None,
-    dim: Optional[DimsType] = (-2, -1),
+    s: ShapeType | None = None,
+    dim: DimsType | None = (-2, -1),
     norm: NormType = None,
 ) -> TensorLikeType:
     return torch.fft.rfftn(input, s=s, dim=dim, norm=norm)
@@ -543,8 +543,8 @@ def rfft2(
 @out_wrapper()
 def irfft2(
     input: TensorLikeType,
-    s: Optional[ShapeType] = None,
-    dim: Optional[DimsType] = (-2, -1),
+    s: ShapeType | None = None,
+    dim: DimsType | None = (-2, -1),
     norm: NormType = None,
 ) -> TensorLikeType:
     return torch.fft.irfftn(input, s=s, dim=dim, norm=norm)
@@ -554,8 +554,8 @@ def irfft2(
 @out_wrapper()
 def hfft2(
     input: TensorLikeType,
-    s: Optional[ShapeType] = None,
-    dim: Optional[DimsType] = (-2, -1),
+    s: ShapeType | None = None,
+    dim: DimsType | None = (-2, -1),
     norm: NormType = None,
 ) -> TensorLikeType:
     return torch.fft.hfftn(input, s=s, dim=dim, norm=norm)
@@ -565,14 +565,14 @@ def hfft2(
 @out_wrapper()
 def ihfft2(
     input: TensorLikeType,
-    s: Optional[ShapeType] = None,
-    dim: Optional[DimsType] = (-2, -1),
+    s: ShapeType | None = None,
+    dim: DimsType | None = (-2, -1),
     norm: NormType = None,
 ) -> TensorLikeType:
     return torch.fft.ihfftn(input, s=s, dim=dim, norm=norm)
 
 
-def _default_alldims(dim: Optional[DimsType], x: TensorLikeType) -> list[int]:
+def _default_alldims(dim: DimsType | None, x: TensorLikeType) -> list[int]:
     """Convert Optional[DimsType] to a simple list, defaulting to all dimensions"""
     if dim is None:
         return list(range(x.ndim))
@@ -583,14 +583,14 @@ def _default_alldims(dim: Optional[DimsType], x: TensorLikeType) -> list[int]:
 
 
 @register_decomposition(aten.fft_fftshift)
-def fftshift(input: TensorLikeType, dim: Optional[DimsType] = None) -> TensorLikeType:
+def fftshift(input: TensorLikeType, dim: DimsType | None = None) -> TensorLikeType:
     dims = _default_alldims(dim, input)
     shift = [input.shape[d] // 2 for d in dims]
     return torch.roll(input, shift, dims)
 
 
 @register_decomposition(aten.fft_ifftshift)
-def ifftshift(input: TensorLikeType, dim: Optional[DimsType] = None) -> TensorLikeType:
+def ifftshift(input: TensorLikeType, dim: DimsType | None = None) -> TensorLikeType:
     dims = _default_alldims(dim, input)
     shift = [(input.shape[d] + 1) // 2 for d in dims]
     return torch.roll(input, shift, dims)
diff --git a/torch/_refs/linalg/__init__.py b/torch/_refs/linalg/__init__.py
index b1fadc2fafcba..a1059aea29e33 100644
--- a/torch/_refs/linalg/__init__.py
+++ b/torch/_refs/linalg/__init__.py
@@ -37,7 +37,7 @@
 ]
 
 
-def _check_norm_dtype(dtype: Optional[torch.dtype], x_dtype: torch.dtype, fn_name: str):
+def _check_norm_dtype(dtype: torch.dtype | None, x_dtype: torch.dtype, fn_name: str):
     """
     Checks related to the dtype kwarg in `linalg.*norm` functions
     """
@@ -99,7 +99,7 @@ def diagonal(
 
 
 def _check_vector_norm_args(
-    x: TensorLikeType, ord: Union[float, int] = 2, dim: Optional[DimsType] = None
+    x: TensorLikeType, ord: float | int = 2, dim: DimsType | None = None
 ):
     from torch.fx.experimental.symbolic_shapes import sym_or
 
@@ -130,11 +130,11 @@ def _check_vector_norm_args(
 @out_wrapper(exact_dtype=True)
 def vector_norm(
     x: TensorLikeType,
-    ord: Union[float, int] = 2,
-    dim: Optional[DimsType] = None,
+    ord: float | int = 2,
+    dim: DimsType | None = None,
     keepdim: bool = False,
     *,
-    dtype: Optional[torch.dtype] = None,
+    dtype: torch.dtype | None = None,
 ) -> Tensor:
     from torch.fx.experimental.symbolic_shapes import guard_or_false
 
@@ -169,7 +169,7 @@ def vector_norm(
         if dim == []:
             dim = None
 
-        if (dim is None and x.numel() == 1) or (
+        if (dim is None and guard_or_false(x.numel() == 1)) or (
             dim is not None
             and (x.ndim > 0 and all(guard_or_false(x.shape[d] == 1) for d in dim))
         ):
@@ -208,11 +208,11 @@ def _inverse_permutation(perm):
 @out_wrapper(exact_dtype=True)
 def matrix_norm(
     A: TensorLikeType,
-    ord: Union[float, str] = "fro",
+    ord: float | str = "fro",
     dim: DimsType = (-2, -1),
     keepdim: bool = False,
     *,
-    dtype: Optional[torch.dtype] = None,
+    dtype: torch.dtype | None = None,
 ) -> TensorLikeType:
     # shape
     check_is_matrix(A, "linalg.matrix_norm")
@@ -309,11 +309,11 @@ def _max_min_wrapper(A, dim):
 @out_wrapper(exact_dtype=True)
 def norm(
     A: TensorLikeType,
-    ord: Optional[Union[float, str]] = None,
-    dim: Optional[DimsType] = None,
+    ord: float | str | None = None,
+    dim: DimsType | None = None,
     keepdim: bool = False,
     *,
-    dtype: Optional[torch.dtype] = None,
+    dtype: torch.dtype | None = None,
 ) -> TensorLikeType:
     if dim is not None:
         if isinstance(dim, Dim):
diff --git a/torch/_refs/nn/functional/__init__.py b/torch/_refs/nn/functional/__init__.py
index 4f17b96a01d7e..a697559dfbd18 100644
--- a/torch/_refs/nn/functional/__init__.py
+++ b/torch/_refs/nn/functional/__init__.py
@@ -175,7 +175,7 @@ def _fn(*args: _P.args, **kwargs: _P.kwargs) -> _T:
     type_promotion_kind=ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT,
 )
 def celu(
-    a: TensorLikeType, alpha: Optional[NumberType] = None, inplace: bool = False
+    a: TensorLikeType, alpha: NumberType | None = None, inplace: bool = False
 ) -> TensorLikeType:
     """
     Reference implementation of torch.nn.functional.celu
@@ -322,8 +322,8 @@ def channel_shuffle(input: TensorLikeType, groups: int) -> TensorLikeType:
 def group_norm(
     input: Tensor,
     num_groups: int,
-    weight: Optional[Tensor] = None,
-    bias: Optional[Tensor] = None,
+    weight: Tensor | None = None,
+    bias: Tensor | None = None,
     eps: float = 1e-5,
 ) -> Tensor:
     """
@@ -362,8 +362,8 @@ def group_norm(
 def layer_norm(
     input: Tensor,
     normalized_shape: ShapeType,
-    weight: Optional[Tensor] = None,
-    bias: Optional[Tensor] = None,
+    weight: Tensor | None = None,
+    bias: Tensor | None = None,
     eps: float = 1e-5,
 ) -> Tensor:
     """
@@ -439,9 +439,9 @@ def selu(a: TensorLikeType, inplace: bool = False) -> TensorLikeType:
 # CompositeImplicitAutograd - don't register decomp
 def softmax(
     a: TensorLikeType,
-    dim: Optional[int] = None,
+    dim: int | None = None,
     _stacklevel: int = 3,  # for compat when using TorchRefsMode(strict=True)
-    dtype: Optional[torch.dtype] = None,
+    dtype: torch.dtype | None = None,
 ) -> TensorLikeType:
     # The error is for compat with regular PyTorch, which has this behavior
     # deprecated.  For PrimTorch, it's fine to drop support for deprecated
@@ -454,9 +454,9 @@ def softmax(
 # CompositeImplicitAutograd - don't register decomp
 def softmin(
     a: TensorLikeType,
-    dim: Optional[int] = None,
+    dim: int | None = None,
     _stacklevel: int = 3,  # for compat when using TorchRefsMode(strict=True)
-    dtype: Optional[torch.dtype] = None,
+    dtype: torch.dtype | None = None,
 ) -> TensorLikeType:
     # The error is for compat with regular PyTorch, which has this behavior
     # deprecated.  For PrimTorch, it's fine to drop support for deprecated
@@ -476,7 +476,7 @@ def softmin(
 )
 def softplus(
     a: TensorLikeType,
-    beta: Optional[NumberType] = None,
+    beta: NumberType | None = None,
     threshold: NumberType = 20,
     inplace: bool = False,
 ) -> TensorLikeType:
@@ -562,9 +562,7 @@ def _check_reduction_value(reduction: str):
 
 # This helper function maps depreciated arguments, "size_average" and "reduce"
 # to their corresponding "reduction" string argument
-def _get_string_reduction_arg(
-    *, size_average: Optional[bool], reduce: Optional[bool]
-) -> str:
+def _get_string_reduction_arg(*, size_average: bool | None, reduce: bool | None) -> str:
     if size_average is None:
         size_average = True
     if reduce is None:
@@ -586,8 +584,8 @@ def _get_string_reduction_arg(
 def l1_loss(
     input: TensorLikeType,
     target: TensorLikeType,
-    size_average: Optional[bool] = None,
-    reduce: Optional[bool] = None,
+    size_average: bool | None = None,
+    reduce: bool | None = None,
     reduction: str = "mean",
 ) -> TensorLikeType:
     """
@@ -610,8 +608,8 @@ def l1_loss(
 def smooth_l1_loss(
     input: TensorLikeType,
     target: TensorLikeType,
-    size_average: Optional[bool] = None,
-    reduce: Optional[bool] = None,
+    size_average: bool | None = None,
+    reduce: bool | None = None,
     reduction: str = "mean",
     beta: float = 1.0,
 ) -> TensorLikeType:
@@ -640,9 +638,9 @@ def smooth_l1_loss(
 # CompositeImplicitAutograd - don't register decomp
 def log_softmax(
     a: TensorLikeType,
-    dim: Optional[int] = None,
+    dim: int | None = None,
     _stacklevel: int = 3,  # for compat when using TorchRefsMode(strict=True)
-    dtype: Optional[torch.dtype] = None,
+    dtype: torch.dtype | None = None,
 ) -> TensorLikeType:
     # The error is for compat with regular PyTorch, which has this behavior
     # deprecated.  For PrimTorch, it's fine to drop support for deprecated
@@ -678,8 +676,8 @@ def margin_ranking_loss(
 def mse_loss(
     input: TensorLikeType,
     target: TensorLikeType,
-    size_average: Optional[bool] = None,
-    reduce: Optional[bool] = None,
+    size_average: bool | None = None,
+    reduce: bool | None = None,
     reduction: str = "mean",
 ) -> TensorLikeType:
     if size_average is not None or reduce is not None:
@@ -712,7 +710,7 @@ def hinge_embedding_loss(
 def _nll_loss_nd(
     input: TensorLikeType,
     target: TensorLikeType,
-    weight: Optional[TensorLikeType],
+    weight: TensorLikeType | None,
     reduction: str,
     ignore_index: int,
 ) -> TensorLikeType:
@@ -797,10 +795,10 @@ def _nll_loss_nd(
 def nll_loss(
     input: TensorLikeType,
     target: TensorLikeType,
-    weight: Optional[TensorLikeType] = None,
-    size_average: Optional[bool] = None,
+    weight: TensorLikeType | None = None,
+    size_average: bool | None = None,
     ignore_index: int = -100,
-    reduce: Optional[bool] = None,
+    reduce: bool | None = None,
     reduction: str = "mean",
 ) -> TensorLikeType:
     """
@@ -880,7 +878,7 @@ def nll_loss(
 def huber_loss(
     input: TensorLikeType,
     target: TensorLikeType,
-    reduction: Union[str, int] = "mean",
+    reduction: str | int = "mean",
     delta: float = 1.0,
 ) -> TensorLikeType:
     """
@@ -925,7 +923,7 @@ def tanhshrink(a: TensorLikeType) -> TensorLikeType:
 def threshold(
     a: TensorLikeType,
     threshold: NumberType,
-    value: Union[bool, int, float],
+    value: bool | int | float,
     inplace: bool = False,
 ) -> TensorLikeType:
     """
@@ -948,8 +946,8 @@ def triplet_margin_loss(
     p: float = 2,
     eps: float = 1e-6,
     swap: bool = False,
-    size_average: Optional[bool] = None,
-    reduce: Optional[bool] = None,
+    size_average: bool | None = None,
+    reduce: bool | None = None,
     reduction: str = "mean",
 ) -> TensorLikeType:
     if size_average is not None or reduce is not None:
@@ -981,9 +979,8 @@ def _triplet_margin_with_distance_loss(
     positive: TensorLikeType,
     negative: TensorLikeType,
     *,
-    distance_function: Optional[
-        Callable[[TensorLikeType, TensorLikeType], TensorLikeType]
-    ] = None,
+    distance_function: Callable[[TensorLikeType, TensorLikeType], TensorLikeType]
+    | None = None,
     margin: float = 1.0,
     swap: bool = False,
     reduction: str = "mean",
@@ -1097,9 +1094,9 @@ def poisson_nll_loss(
     target: TensorLikeType,
     log_input: bool = True,
     full: bool = False,
-    size_average: Optional[bool] = None,
+    size_average: bool | None = None,
     eps: float = 1e-8,
-    reduce: Optional[bool] = None,
+    reduce: bool | None = None,
     reduction: str = "mean",
 ) -> TensorLikeType:
     """
diff --git a/torch/_refs/special/__init__.py b/torch/_refs/special/__init__.py
index dd85fd49fc9b9..b2ae9da46b09a 100644
--- a/torch/_refs/special/__init__.py
+++ b/torch/_refs/special/__init__.py
@@ -132,7 +132,7 @@ def log_ndtr(a: TensorLikeType) -> TensorLikeType:
     type_promoting_args=("self",),
     type_promotion_kind=utils.ELEMENTWISE_TYPE_PROMOTION_KIND.INT_TO_FLOAT,
 )
-def logit(self: TensorLikeType, eps: Optional[float] = None) -> TensorLikeType:
+def logit(self: TensorLikeType, eps: float | None = None) -> TensorLikeType:
     if eps is None:
         eps = -1.0
     lo = eps
@@ -147,7 +147,7 @@ def logit(self: TensorLikeType, eps: Optional[float] = None) -> TensorLikeType:
     type_promoting_args=("a", "b"),
     type_promotion_kind=ELEMENTWISE_TYPE_PROMOTION_KIND.INT_TO_FLOAT,
 )
-def xlog1py(a: Union[TensorLikeType, NumberType], b: Union[TensorLikeType, NumberType]):
+def xlog1py(a: TensorLikeType | NumberType, b: TensorLikeType | NumberType):
     torch._check(
         isinstance(a, TensorLike) or isinstance(b, TensorLike),
         lambda: 'Expected either argument a or b to be a Tensor"',
@@ -210,7 +210,7 @@ def ndtri(a: TensorLikeType) -> TensorLikeType:
 def log_softmax(
     a: TensorLikeType,
     dim: int,
-    dtype: Optional[torch.dtype] = None,
+    dtype: torch.dtype | None = None,
 ) -> TensorLikeType:
     return torch.log_softmax(a=a, dim=dim, dtype=dtype)  # type: ignore[call-overload]
 
@@ -220,7 +220,7 @@ def log_softmax(
 def softmax(
     a: TensorLikeType,
     dim: int,
-    dtype: Optional[torch.dtype] = None,
+    dtype: torch.dtype | None = None,
 ) -> TensorLikeType:
     return torch.softmax(a=a, dim=dim, dtype=dtype)  # type: ignore[call-overload]
 
diff --git a/torch/_strobelight/cli_function_profiler.py b/torch/_strobelight/cli_function_profiler.py
index a63a49c3938a1..886bc5039f3b0 100644
--- a/torch/_strobelight/cli_function_profiler.py
+++ b/torch/_strobelight/cli_function_profiler.py
@@ -9,7 +9,7 @@
 from collections.abc import Callable, Sequence
 from threading import Lock
 from timeit import default_timer as timer
-from typing import Any, Optional, TypeVar
+from typing import Any, TypeVar
 from typing_extensions import ParamSpec
 
 
@@ -35,14 +35,14 @@ class StrobelightCLIProfilerError(Exception):
     """
 
 
-def _pid_namespace_link(pid: Optional[int] = None) -> str:
+def _pid_namespace_link(pid: int | None = None) -> str:
     """Returns the link to the process's namespace, example: pid:[4026531836]"""
     PID_NAMESPACE_PATH = "/proc/{}/ns/pid"
     pid = pid or os.getpid()
     return os.readlink(PID_NAMESPACE_PATH.format(pid))
 
 
-def _pid_namespace(pid: Optional[int] = None) -> int:
+def _pid_namespace(pid: int | None = None) -> int:
     """Returns the process's namespace id"""
     pid = pid or os.getpid()
     link = _pid_namespace_link(pid)
@@ -78,8 +78,8 @@ def __init__(
         run_user_name: str = "pytorch-strobelight-ondemand",
         timeout_wait_for_running_sec: int = 60,
         timeout_wait_for_finished_sec: int = 60,
-        recorded_env_variables: Optional[list[str]] = None,
-        sample_tags: Optional[list[str]] = None,
+        recorded_env_variables: list[str] | None = None,
+        sample_tags: list[str] | None = None,
         stack_max_len: int = 127,
         async_stack_max_len: int = 127,
     ):
@@ -91,8 +91,8 @@ def __init__(
         self.timeout_wait_for_finished_sec = timeout_wait_for_finished_sec
         # Results of the most recent run.
         # Tracks the strobelight run id of the most recent run
-        self.current_run_id: Optional[int] = None
-        self.profile_result: Optional[list[str]] = None
+        self.current_run_id: int | None = None
+        self.profile_result: list[str] | None = None
         self.sample_tags = sample_tags
 
     def _run_async(self) -> None:
@@ -257,7 +257,7 @@ def _start_strobelight(self) -> bool:
 
     def profile(
         self, work_function: Callable[_P, _R], *args: _P.args, **kwargs: _P.kwargs
-    ) -> Optional[_R]:
+    ) -> _R | None:
         self.current_run_id = None
         self.profile_result = None
 
@@ -304,16 +304,16 @@ def profile(
 # @strobelight(profiler = StrobelightFunctionProfiler(stop_at_error=True,..))
 # @strobelight(stop_at_error=True,...)
 def strobelight(
-    profiler: Optional[StrobelightCLIFunctionProfiler] = None, **kwargs: Any
-) -> Callable[[Callable[_P, _R]], Callable[_P, Optional[_R]]]:
+    profiler: StrobelightCLIFunctionProfiler | None = None, **kwargs: Any
+) -> Callable[[Callable[_P, _R]], Callable[_P, _R | None]]:
     if not profiler:
         profiler = StrobelightCLIFunctionProfiler(**kwargs)
 
     def strobelight_inner(
         work_function: Callable[_P, _R],
-    ) -> Callable[_P, Optional[_R]]:
+    ) -> Callable[_P, _R | None]:
         @functools.wraps(work_function)
-        def wrapper_function(*args: _P.args, **kwargs: _P.kwargs) -> Optional[_R]:
+        def wrapper_function(*args: _P.args, **kwargs: _P.kwargs) -> _R | None:
             # pyrefly: ignore [bad-argument-type]
             return profiler.profile(work_function, *args, **kwargs)
 
diff --git a/torch/_strobelight/compile_time_profiler.py b/torch/_strobelight/compile_time_profiler.py
index 436f9a2c8b594..a990e49f23322 100644
--- a/torch/_strobelight/compile_time_profiler.py
+++ b/torch/_strobelight/compile_time_profiler.py
@@ -7,7 +7,7 @@
 import subprocess
 from datetime import datetime
 from socket import gethostname
-from typing import Any, Optional
+from typing import Any
 
 from torch._strobelight.cli_function_profiler import StrobelightCLIFunctionProfiler
 
@@ -87,15 +87,15 @@ class StrobelightCompileTimeProfiler:
     enabled: bool = False
 
     # A regex that can be used to filter out what frames to profile. ex: "1/.*"
-    frame_id_filter: Optional[str] = os.environ.get("COMPILE_STROBELIGHT_FRAME_FILTER")
+    frame_id_filter: str | None = os.environ.get("COMPILE_STROBELIGHT_FRAME_FILTER")
 
     # A unique identifier that is used as the run_user_name in the strobelight profile to
     # associate all compile time profiles together.
-    identifier: Optional[str] = None
+    identifier: str | None = None
 
-    current_phase: Optional[str] = None
+    current_phase: str | None = None
 
-    profiler: Optional[Any] = None
+    profiler: Any | None = None
 
     max_stack_length: int = int(
         os.environ.get("COMPILE_STROBELIGHT_MAX_STACK_LENGTH", 500)
diff --git a/torch/_subclasses/_fake_tensor_utils.py b/torch/_subclasses/_fake_tensor_utils.py
index 1c8831ef39bbe..61a71a5a1d3c7 100644
--- a/torch/_subclasses/_fake_tensor_utils.py
+++ b/torch/_subclasses/_fake_tensor_utils.py
@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
-from typing import Optional, TYPE_CHECKING, Union
+from typing import TYPE_CHECKING
 
 import torch
 from torch import SymInt
@@ -26,8 +26,8 @@ class _DeconstructedSymNode:
     # n.b. keep the same protocol as SymNode
     _expr: sympy.Expr
     pytype: type
-    _hint: Optional[Union[int, float, bool]]
-    constant: Optional[Union[int, float, bool]]
+    _hint: int | float | bool | None
+    constant: int | float | bool | None
     fx_node: torch.fx.Node
 
     @staticmethod
@@ -35,9 +35,10 @@ def from_node(node: SymNode) -> _DeconstructedSymNode:
         return _DeconstructedSymNode(
             node._expr,
             node.pytype,
+            # pyrefly: ignore[bad-argument-type]
             node._hint,
             node.constant,
-            # pyrefly: ignore [bad-argument-type]
+            # pyrefly: ignore[bad-argument-type]
             node.fx_node,
         )
 
@@ -125,10 +126,10 @@ class _PySymInputStub:
     #                          the cache to avoid cyclic ShapeEnv references.
     #   _InputBackref: This is a back-reference to a previous _PySymInputStub in
     #                  the key.
-    value: Union[PySymType, _DeconstructedSymType, _InputBackref]
+    value: PySymType | _DeconstructedSymType | _InputBackref
 
     def __init__(
-        self, value: Union[PySymType, _DeconstructedSymType, _InputBackref]
+        self, value: PySymType | _DeconstructedSymType | _InputBackref
     ) -> None:
         # For inputs (values in the `key`) we need to keep the PySymType intact
         # - this way if we need to reuse it as an output we can properly copy
@@ -183,9 +184,9 @@ class _SymIntOutputStub:
 
     # This is either an `int` which represents the index in the key to copy the
     # SymNode from or it's the deconstructed SymNode itself.
-    value: Union[int, _DeconstructedSymNode]
+    value: int | _DeconstructedSymNode
 
-    def __init__(self, value: SymInt, key_path: Optional[int]) -> None:
+    def __init__(self, value: SymInt, key_path: int | None) -> None:
         if key_path is None:
             self.value = _DeconstructedSymNode.from_node(value.node)
         else:
@@ -234,9 +235,9 @@ class _CacheKeyState:
     # ShapeEnv on the FakeTensorMode - but for SymNodes we MUST have a
     # ShapeEnv. So as we scan if we see a SymNode (with a ShapeEnv) we record it
     # here.
-    shape_env: Optional[ShapeEnv]
+    shape_env: ShapeEnv | None
 
-    def __init__(self, shape_env: Optional[ShapeEnv] = None) -> None:
+    def __init__(self, shape_env: ShapeEnv | None = None) -> None:
         self.sym_node_lookup = {}
         self.known_symbols = set()
         self.shape_env = shape_env
diff --git a/torch/_subclasses/fake_impls.py b/torch/_subclasses/fake_impls.py
index 1d1d7f4afaf88..e7e66bdd57acd 100644
--- a/torch/_subclasses/fake_impls.py
+++ b/torch/_subclasses/fake_impls.py
@@ -6,7 +6,7 @@
 import operator
 import sys
 from functools import reduce
-from typing import Any, cast as typing_cast, TYPE_CHECKING, TypeVar, Union
+from typing import Any, cast as typing_cast, TYPE_CHECKING, TypeVar
 from typing_extensions import ParamSpec
 
 import torch
@@ -47,7 +47,7 @@
     from torch.types import IntLikeType
 
 
-FakeTensorLike = Union[FakeTensor, torch.Tensor]
+FakeTensorLike = FakeTensor | torch.Tensor
 _P = ParamSpec("_P")
 _R = TypeVar("_R")
 _T = TypeVar("_T")
@@ -426,9 +426,9 @@ def _unique(
 
         if dim is None:
             if unique_consecutive:
-                arg.unique_consecutive_memo = nnz
+                arg.unique_consecutive_memo = nnz  # pyrefly: ignore[bad-assignment]
             else:
-                arg.unique_memo = nnz
+                arg.unique_memo = nnz  # pyrefly: ignore[bad-assignment]
 
     if dim is None:
         # pyrefly: ignore[no-matching-overload]
@@ -439,15 +439,19 @@ def _unique(
 
     return_if_dim_and_cpu = dim is not None and arg.fake_device == torch.device("cpu")
     if return_inverse or return_if_dim_and_cpu:
-        inverse = arg.new_empty(arg.shape if dim is None else (arg.shape[dim],))
+        inverse = arg.new_empty(
+            arg.shape if dim is None else (arg.shape[dim],), dtype=torch.int64
+        )
     else:
-        inverse = arg.new_empty(0)
+        inverse = arg.new_empty(0, dtype=torch.int64)
     ret.append(inverse)
 
     if return_counts or return_if_dim_and_cpu:
-        counts = arg.new_empty(ret[0].shape if dim is None else (ret[0].shape[dim],))
+        counts = arg.new_empty(
+            ret[0].shape if dim is None else (ret[0].shape[dim],), dtype=torch.int64
+        )
     else:
-        counts = arg.new_empty(0)
+        counts = arg.new_empty(0, dtype=torch.int64)
     ret.append(counts)
 
     return tuple(ret)
@@ -641,14 +645,14 @@ def maybe_guard_or_true(x: Any) -> Any:
 def _view_has_unbacked_input(
     a: torch.Tensor, shape: ShapeType | tuple[ShapeType]
 ) -> bool:
-    from torch.fx.experimental.symbolic_shapes import has_hint
+    from torch.fx.experimental.symbolic_shapes import has_guarding_hint
 
     shape = utils.extract_shape_from_varargs(shape, validate=False)
 
     return (
-        any(not has_hint(s) for s in a.size())
-        or any(not has_hint(s) for s in a.stride())
-        or any(not has_hint(s) for s in shape)
+        any(not has_guarding_hint(s) for s in a.size())
+        or any(not has_guarding_hint(s) for s in a.stride())
+        or any(not has_guarding_hint(s) for s in shape)
     )
 
 
@@ -656,6 +660,7 @@ def _view_unbacked_meta(
     a: torch.Tensor,
     shape: ShapeType | tuple[ShapeType],
     size_oblivious_enabled: bool = True,
+    allow_copy: bool = False,
 ) -> torch.Tensor:
     from torch._prims import view_of
     from torch.fx.experimental.symbolic_shapes import guard_or_false, sym_eq
@@ -719,7 +724,16 @@ def _view_unbacked_meta(
         torch.fx.experimental._config.backed_size_oblivious
         or _view_has_unbacked_input(a, shape)
     ):
-        return _view_unbacked_meta(a, shape, size_oblivious_enabled=False)
+        return _view_unbacked_meta(
+            a, shape, size_oblivious_enabled=False, allow_copy=allow_copy
+        )
+
+    # When allow_copy=True (i.e., view_copy), define unbacked semantics
+    # as "materialize": clone the input to break aliasing, then reshape.
+    if allow_copy:
+        strides = make_contiguous_strides_for(shape)
+        # pyrefly: ignore[bad-return]
+        return a.clone(memory_format=torch.contiguous_format).as_strided(shape, strides)
 
     msg = f"Cannot view a tensor with shape {a.shape} and strides {a.stride()} as a tensor with shape {shape}!"
     raise ValueError(msg)
@@ -755,14 +769,18 @@ def _view_meta(
     func: OpOverload,
     a: FakeTensor,
     *shape: Any,
+    allow_copy: bool = False,
 ) -> FakeTensor:
     if torch.fx.experimental._config.backed_size_oblivious or _view_has_unbacked_input(
         a, shape
     ):
-        return typing_cast(FakeTensor, _view_unbacked_meta(a, shape))
+        return typing_cast(
+            FakeTensor, _view_unbacked_meta(a, shape, allow_copy=allow_copy)
+        )
     else:
         return typing_cast(
-            FakeTensor, torch._refs._reshape_view_helper(a, *shape, allow_copy=False)
+            FakeTensor,
+            torch._refs._reshape_view_helper(a, *shape, allow_copy=allow_copy),
         )
 
 
@@ -774,7 +792,10 @@ def _view_meta_copy(
     *shape: IntLikeType,
     out: FakeTensor | None = None,
 ) -> FakeTensor:
-    result = _view_meta(fake_mode, func, a, *shape)
+    # view_copy is the non-aliasing counterpart of view. Eager may succeed on
+    # cases where a pure view is impossible (e.g. expand -> flatten) by
+    # materializing the result. Match eager by allowing copy-if-needed in meta.
+    result = _view_meta(fake_mode, func, a, *shape, allow_copy=True)
 
     if out is not None:
         return result
@@ -888,7 +909,7 @@ def nonzero(fake_mode: FakeTensorMode, func: OpOverload, arg: FakeTensor) -> Fak
 
             _constrain_range_for_size(nnz, max=maxval)
 
-        arg.nonzero_memo = nnz
+        arg.nonzero_memo = nnz  # pyrefly: ignore[bad-assignment]
     return arg.new_empty_strided((nnz, arg.dim()), (1, nnz), dtype=torch.int64)  # type: ignore[return]
 
 
@@ -944,6 +965,11 @@ def _compute_slice_index(size: IntLikeType, index: IntLikeType) -> IntLikeType |
         return 0
     elif guard_or_false(index > size):
         return size
+    elif guard_or_false(index >= 0):
+        return torch.sym_min(index, size)
+    elif guard_or_false(index < 0):
+        return torch.sym_max(index + size, 0)
+
     return None
 
 
@@ -988,6 +1014,12 @@ def slice_forward(
             new_size = (end_index - start_index + step - 1) // step
         elif guard_or_false(start_index >= end_index):
             new_size = 0
+        else:
+            # Both indices are resolved but we can't statically determine their
+            # ordering (e.g., when they involve Min/Max). Compute the size via
+            # max(end - start, 0) to avoid creating an unbacked symint.
+            diff = torch.sym_max(end_index - start_index, 0)
+            new_size = (diff + step - 1) // step
 
     # create unbacked if case unknown
     if new_size is None:
@@ -1312,68 +1344,53 @@ def conv(
     _, new_kwargs = _normalize_function_or_error(
         func, args=args, kwargs=kwargs, normalize_to_only_use_kwargs=True
     )
-    device = new_kwargs["input"].fake_device
+    input_ = new_kwargs["input"]
+    weight = new_kwargs["weight"]
+    device = input_.fake_device
     # need to re-enable mode so the tensors report fake device
     with fake_mode:
-        # if the input is unsqueezed is done in Convolution.cpp we get segfault
-        k = new_kwargs["weight"].ndim
-        batch = new_kwargs["input"].shape[0]
+        # if the input is unsqueezed in Convolution.cpp we get segfault
+        k = weight.ndim
 
         # Avoid importing sympy at a module level
-        from torch.fx.experimental.symbolic_shapes import has_hint
+        from torch.fx.experimental.symbolic_shapes import has_guarding_hint
 
-        if not has_hint(batch):
+        all_hinted = all(has_guarding_hint(s) for s in input_.shape) and all(
+            has_guarding_hint(s) for s in weight.shape
+        )
+
+        if not all_hinted:
             # TODO: We can make this a little more faithful with best effort
             # channels last detection (but only if it's statically obvious!)
             mem_fmt = None
         else:
-            if func is aten.convolution.default:
-                conv_backend = torch._C._select_conv_backend(**new_kwargs)
-            else:
-                conv_backend = torch._C._select_conv_backend(
-                    new_kwargs["input"],
-                    new_kwargs["weight"],
-                    bias=None,
-                    stride=new_kwargs["stride"],
-                    padding=new_kwargs["padding"],
-                    dilation=new_kwargs["dilation"],
-                    transposed=new_kwargs["transposed"],
-                    output_padding=new_kwargs["output_padding"],
-                    groups=new_kwargs["groups"],
-                    bias_sizes=new_kwargs["bias_sizes"],
-                )
+            # convolution has "bias" but not "bias_sizes"; convolution_backward
+            # has "bias_sizes" but not "bias". .get() handles both with one call.
+            bias = new_kwargs.get("bias")
+            select_kwargs: dict[str, object] = dict(
+                stride=new_kwargs["stride"],
+                padding=new_kwargs["padding"],
+                dilation=new_kwargs["dilation"],
+                transposed=new_kwargs["transposed"],
+                output_padding=new_kwargs["output_padding"],
+                groups=new_kwargs["groups"],
+                bias=bias,
+            )
+            if bias is None:
+                select_kwargs["bias_sizes"] = new_kwargs.get("bias_sizes")
+            conv_backend = torch._C._select_conv_backend(
+                input_, weight, **select_kwargs
+            )
             # Expand 1d -> 2d.
             # Note: Avoid expanding before calling _select_conv_backend,
             # as the function handles 2D expansion internally.
-            if (
-                k == 3
-                and not new_kwargs["input"].is_mkldnn
-                and not new_kwargs["input"].is_xpu
-            ):
+            if k == 3 and not input_.is_mkldnn and not input_.is_xpu:
                 # Note: Using input.to(memory_format=contiguous) does not work.
-                new_kwargs["input"] = new_kwargs["input"].contiguous().unsqueeze(2)
-                new_kwargs["weight"] = new_kwargs["weight"].unsqueeze(2)
-                if len(new_kwargs["stride"]) == 1:
-                    new_kwargs["stride"].insert(0, 1)
-                    new_kwargs["padding"].insert(0, 0)
-                    new_kwargs["dilation"].insert(0, 1)
-                    new_kwargs["output_padding"].insert(0, 0)
+                input_ = input_.contiguous().unsqueeze(2)
+                weight = weight.unsqueeze(2)
             mem_fmt = torch._C._conv_determine_backend_memory_format(
-                new_kwargs["input"], new_kwargs["weight"], conv_backend
+                input_, weight, conv_backend
             )
-            # revert 2d -> 1d
-            if (
-                k == 3
-                and not new_kwargs["input"].is_mkldnn
-                and not new_kwargs["input"].is_xpu
-            ):
-                new_kwargs["input"] = new_kwargs["input"].squeeze(2)
-                new_kwargs["weight"] = new_kwargs["weight"].squeeze(2)
-                if len(new_kwargs["stride"]) == 2:
-                    new_kwargs["stride"].pop(0)
-                    new_kwargs["padding"].pop(0)
-                    new_kwargs["dilation"].pop(0)
-                    new_kwargs["output_padding"].pop(0)
 
     def convert(
         t: torch.Tensor | None, mem_fmt: torch.memory_format | None
diff --git a/torch/_subclasses/fake_tensor.py b/torch/_subclasses/fake_tensor.py
index 20dc1376373a9..5375659868527 100644
--- a/torch/_subclasses/fake_tensor.py
+++ b/torch/_subclasses/fake_tensor.py
@@ -14,16 +14,7 @@
 import weakref
 from collections import defaultdict
 from dataclasses import dataclass
-from typing import (
-    Any,
-    cast,
-    Literal,
-    Optional,
-    TYPE_CHECKING,
-    TypeGuard,
-    TypeVar,
-    Union,
-)
+from typing import Any, cast, Literal, TYPE_CHECKING, TypeGuard, TypeVar, Union
 from typing_extensions import Self
 from weakref import ReferenceType
 
@@ -34,6 +25,7 @@
 from torch._library.fake_class_registry import FakeScriptObject
 from torch._library.fake_profile import MissingOpProfile
 from torch._logging import dtrace_structured
+from torch._opaque_base import OpaqueBase
 from torch._prims_common import suggest_memory_format
 from torch._subclasses.meta_utils import (
     assert_eq,
@@ -42,7 +34,7 @@
     is_sparse_compressed,
     MetaConverter,
 )
-from torch._utils import render_call
+from torch._utils import _is_privateuse1_backend_available, render_call
 from torch.fx.immutable_collections import immutable_dict
 from torch.fx.operator_schemas import normalize_function
 from torch.multiprocessing.reductions import StorageWeakRef
@@ -65,7 +57,6 @@
     from types import TracebackType
 
     from torch._guards import Source
-    from torch._library.opaque_object import OpaqueType
     from torch._ops import OpOverload
     from torch.fx.experimental.symbolic_shapes import ShapeEnv, SymbolicContext
 
@@ -95,6 +86,19 @@
 RECURSION_COUNT = 0
 
 
+# Check if device type supports device index
+@functools.cache
+def _is_indexed_device_type(device_type: str) -> bool:
+    return device_type in [
+        "cuda",
+        "hpu",
+        "xpu",
+        "mps",
+        "mtia",
+        torch._C._get_privateuse1_backend_name(),
+    ]
+
+
 # Small helper that increments recursion count, and
 # resets it when the object goes out of scope.  Useful
 # if you don't want to increase indentation which is
@@ -178,8 +182,8 @@ def disable_fake_tensor_cache(fake_mode: FakeTensorMode) -> Generator[None, None
 
 
 def get_plain_tensors(
-    subclass: Tensor, *, out: list[Tensor | int | SymInt | OpaqueType]
-) -> list[Tensor | int | SymInt | OpaqueType]:
+    subclass: Tensor, *, out: list[Tensor | int | SymInt | OpaqueBase]
+) -> list[Tensor | int | SymInt | OpaqueBase]:
     # This function is used in Runtime, do not add redundant asserts
     todo = [subclass]
     while todo:
@@ -201,12 +205,22 @@ def is_fake(x: object) -> TypeGuard[Tensor]:
         return True
     if is_traceable_wrapper_subclass(x):
         attrs, _ = type(x).__tensor_flatten__(x)
-        flattened_tensors = [getattr(x, attr) for attr in attrs]
-        all_fake = all(is_fake(x) for x in flattened_tensors)
-        any_fake = any(is_fake(x) for x in flattened_tensors)
-        if all_fake != any_fake:
-            raise AssertionError("got mixed fake and real tensors!")
-        return all_fake
+        got_fake: bool | None = None
+        for attr in attrs:
+            match getattr(x, attr):
+                case Tensor() as v:
+                    fake = is_fake(v)
+                    if got_fake is None:
+                        got_fake = fake
+                    elif got_fake != fake:
+                        raise AssertionError("got mixed fake and real tensors!")
+                case OpaqueBase():
+                    pass
+                case unexpected:
+                    raise AssertionError(
+                        f"expected Tensor or OpaqueBase, got {type(unexpected)}"
+                    )
+        return got_fake or False
     elif isinstance(x, FunctionalTensor):
         return is_fake(x.elem)
     elif isinstance(x, Tensor) and torch._is_functional_tensor(x):
@@ -219,20 +233,29 @@ def is_fake(x: object) -> TypeGuard[Tensor]:
     return False
 
 
-def maybe_get_fake_mode(t: object) -> Optional[FakeTensorMode]:
+def maybe_get_fake_mode(t: object) -> FakeTensorMode | None:
     from torch._subclasses.functional_tensor import FunctionalTensor
 
     if isinstance(t, FakeTensor):
         return t.fake_mode
     if is_traceable_wrapper_subclass(t):
         inner_tensor_names, _ = t.__tensor_flatten__()
-        modes = [
-            maybe_get_fake_mode(getattr(t, t_name)) for t_name in inner_tensor_names
-        ]
-        m = modes[0]
-        if not all(m is x for x in modes):
-            raise AssertionError("All fake tensor modes must be the same")
-        return m
+        mode: FakeTensorMode | None = None
+        for t_name in inner_tensor_names:
+            match getattr(t, t_name):
+                case Tensor() as v:
+                    m = maybe_get_fake_mode(v)
+                    if mode is None:
+                        mode = m
+                    elif mode is not m:
+                        raise AssertionError("All fake tensor modes must be the same")
+                case OpaqueBase():
+                    pass
+                case unexpected:
+                    raise AssertionError(
+                        f"expected Tensor or OpaqueBase, got {type(unexpected)}"
+                    )
+        return mode
     elif isinstance(t, FunctionalTensor):
         return maybe_get_fake_mode(t.elem)
     elif isinstance(t, Tensor) and torch._is_functional_tensor(t):
@@ -343,7 +366,7 @@ def invalidate_constant_aliases(self, tensor: Tensor) -> None:
 
         del self.constant_storage_mapping[weak_st]
 
-    def _get_memo(self, t: Tensor) -> Optional[FakeTensor]:
+    def _get_memo(self, t: Tensor) -> FakeTensor | None:
         tid = self.meta_converter.describer.lookup_tensor.get(t)
         if tid is None:
             return None
@@ -364,10 +387,10 @@ def from_real_tensor(
         fake_mode: FakeTensorMode,
         t: Tensor,
         make_constant: bool = False,
-        shape_env: Optional[ShapeEnv] = None,
+        shape_env: ShapeEnv | None = None,
         *,
-        source: Optional[Source] = None,
-        symbolic_context: Optional[SymbolicContext] = None,
+        source: Source | None = None,
+        symbolic_context: SymbolicContext | None = None,
         trace: bool = True,
     ) -> FakeTensor:
         # see note [Tensor Fakification and Symbol Caching]
@@ -401,7 +424,7 @@ def from_real_tensor(
         # caller to explicitly specify the device in case outer and inner tensors
         # have different devices.
         def mk_fake_tensor(
-            make_meta_t: Callable[[], object], device: Union[torch.device, str]
+            make_meta_t: Callable[[], object], device: torch.device | str
         ) -> FakeTensor:
             # NB: don't use in_kernel_invocation_manager. to
             # ensure FakeTensor can internally do constant computation
@@ -514,8 +537,8 @@ def from_meta_and_device(
         fake_mode: FakeTensorMode,
         t: Tensor,
         device: torch.device,
-        pytype: Optional[type[torch.Tensor]] = None,
-        dispatch_keys: Optional[torch.DispatchKeySet] = None,
+        pytype: type[torch.Tensor] | None = None,
+        dispatch_keys: torch.DispatchKeySet | None = None,
     ) -> FakeTensor:
         if t.device.type != "meta":
             raise AssertionError(
@@ -536,7 +559,7 @@ def from_meta_and_device(
 @functools.cache
 def init_gpu_context(device: torch.device) -> None:
     # Backward will error with cuda Fake Tensors if no cuda tensors have been initialized first
-    if torch.cuda.is_available() or torch.xpu.is_available():
+    if torch.accelerator.current_accelerator(True) == device.type:
         (
             torch.empty(1, device=device)
             if torch.version.hip is None
@@ -619,8 +642,8 @@ def _memo_epoch(self, obj: FakeTensor) -> str:
         return f"_{self._name}_epoch"
 
     def __get__(
-        self, obj: FakeTensor, objtype: Optional[type[FakeTensor]] = None
-    ) -> Optional[Union[torch.SymInt, torch.SymFloat]]:
+        self, obj: FakeTensor, objtype: type[FakeTensor] | None = None
+    ) -> torch.SymInt | torch.SymFloat | None:
         if (r := getattr(obj, self._memo(obj))) is None:
             return None
 
@@ -641,7 +664,9 @@ def __get__(
         return r
 
     def __set__(
-        self, obj: FakeTensor, value: Optional[Union[torch.SymInt, torch.SymFloat]]
+        self,
+        obj: FakeTensor,
+        value: torch.SymInt | torch.SymFloat | torch.SymBool | int | float | None,
     ) -> None:
         if value is None:
             setattr(obj, self._memo(obj), None)
@@ -663,10 +688,10 @@ class FakeTensor(Tensor):
     which tracks devices that would have been used.
     """
 
-    fake_device: torch.device
+    _fake_device: torch.device
     fake_mode: FakeTensorMode
-    constant: Optional[Tensor]
-    real_tensor: Optional[Tensor]
+    constant: Tensor | None
+    real_tensor: Tensor | None
 
     # TODO: Generalize this as needed, e.g., into a trie of memos, if
     # you do something like x[0].item()  (x[0] is fresh each time, so
@@ -686,8 +711,8 @@ class FakeTensor(Tensor):
     # FakeTensor doesn't fully emulate the original tensor's Python type
     # and dispatch key set, therefore sometimes we want to track them
     # separately.
-    pytype: Optional[type[Tensor]]
-    dispatch_keys: Optional[torch.DispatchKeySet]
+    pytype: type[Tensor] | None
+    dispatch_keys: torch.DispatchKeySet | None
 
     # Indicates to our torch_dispatch dispatching infra that
     # this is an "infra" mode with lower dispatching precedence.
@@ -705,6 +730,14 @@ def device(self) -> torch.device:
     def device(self, _: torch.device) -> None:
         raise NotImplementedError
 
+    @property
+    def fake_device(self) -> torch.device:
+        return self._fake_device
+
+    @fake_device.setter
+    def fake_device(self, device: torch.device) -> None:
+        self._fake_device = self._normalize_fake_device(device)
+
     # Note: [Fake Tensor Dispatch Keys]
     # In order to model the behavior of device-specific autocast
     # and autograd logic, we update the dispatch keys of FakeTensors
@@ -732,16 +765,31 @@ def names(self) -> list[str]:
     def names(self, _: list[str]) -> None:
         raise NotImplementedError
 
+    @staticmethod
+    def _normalize_fake_device(device: torch.device) -> torch.device:
+        """Normalize device by initializing GPU context and setting device index."""
+        if device.type in ("cuda", "xpu"):
+            init_gpu_context(device)
+
+        if _is_indexed_device_type(device.type) and device.index is None:
+            if device.type != "mps" and getattr(torch, device.type).is_initialized():
+                device = torch.device(
+                    f"{device.type}:{getattr(torch, device.type).current_device()}"
+                )
+            else:
+                device = torch.device(f"{device.type}:0")
+        return device
+
     @staticmethod
     def __new__(
         cls,
         fake_mode: FakeTensorMode,
         elem: Tensor,
         device: torch.device,
-        constant: Optional[Tensor] = None,
-        real_tensor: Optional[Tensor] = None,
-        pytype: Optional[type[Tensor]] = None,
-        dispatch_keys: Optional[torch.DispatchKeySet] = None,
+        constant: Tensor | None = None,
+        real_tensor: Tensor | None = None,
+        pytype: type[Tensor] | None = None,
+        dispatch_keys: torch.DispatchKeySet | None = None,
     ) -> Self:
         self = Tensor._make_subclass(
             cls,
@@ -772,29 +820,6 @@ def __new__(
                 raise AssertionError(
                     "device.type must not be 'meta' when allow_meta is False"
                 )
-        # normalize device.
-        if device.type in ["cuda", "xpu"]:
-            init_gpu_context(device)
-
-        if (
-            device.type
-            in [
-                "cuda",
-                "hpu",
-                "xpu",
-                "mps",
-                "mtia",
-                torch._C._get_privateuse1_backend_name(),
-            ]
-            and device.index is None
-        ):
-            if device.type != "mps" and getattr(torch, device.type).is_initialized():
-                device = torch.device(
-                    f"{device.type}:{getattr(torch, device.type).current_device()}"
-                )
-            else:
-                device = torch.device(f"{device.type}:0")
-        # pyrefly: ignore [read-only]
         self.fake_device = device
         self.fake_mode = fake_mode
         self.constant = constant
@@ -941,16 +966,25 @@ def _find_common_device(
             aten._foreach_copy.default,
         )
 
+        # These in-place ops keep the destination tensor's device even if the
+        # rhs was explicitly constructed on meta.
+        meta_rhs_mixed_device_fns = ordered_set(
+            aten.add_.Tensor,
+        )
+
         # list of ops not using zero dim cpu tensor logic to align with the eager mode.
         bypass_zero_dim_cpu_tensor_check_ops = ordered_set(
             aten.nextafter.default,
         )
 
-        def check_cpu_device(device: torch.device) -> bool:
+        def is_device_cpu(device: torch.device) -> bool:
             return device.type == "cpu"
 
+        def is_device_meta(device: torch.device) -> bool:
+            return device.type == "meta"
+
         def cpu_zero_dim(t: Tensor) -> bool:
-            return check_cpu_device(t.device) and t.dim() == 0
+            return is_device_cpu(t.device) and t.dim() == 0
 
         def merge_devices(t: object) -> None:
             nonlocal common_device
@@ -989,7 +1023,11 @@ def merge_devices(t: object) -> None:
             # device must be cpu in this case we will return from here without
             # throwing an error
             if func in mixed_device_fns:
-                if any(map(check_cpu_device, (common_device, t.device))):
+                if any(map(is_device_cpu, (common_device, t.device))):
+                    return
+
+            if func in meta_rhs_mixed_device_fns:
+                if any(map(is_device_meta, (common_device, t.device))):
                     return
 
             # if prefer_device_type is set, prefer that device type over others
@@ -1032,7 +1070,7 @@ def merge_devices(t: object) -> None:
     def get_nested_int(
         self,
         *,
-        coeff: Union[int, torch.SymInt] = 1,
+        coeff: int | torch.SymInt = 1,
     ) -> torch.SymInt:
         if self.nested_int_memo is None:
             self.nested_int_memo = self.fake_mode.create_symbolic_nested_int(
@@ -1068,18 +1106,18 @@ class TensorMetadata:
     stride: tuple[_MetadataIntLike, ...]
     device: torch.device
     layout: torch.layout
-    memory_format: Optional[torch.memory_format]
+    memory_format: torch.memory_format | None
     storage_offset: _MetadataIntLike
-    storage_bytes: Optional[_MetadataIntLike]
+    storage_bytes: _MetadataIntLike | None
     requires_grad: bool
     is_quantized: bool
     is_conj: bool
     is_neg: bool
     is_inference: bool
     is_sparse: bool  # read: is sparse COO
-    is_coalesced: Optional[bool]
-    dense_dim: Optional[int]
-    sparse_dim: Optional[int]
+    is_coalesced: bool | None
+    dense_dim: int | None
+    sparse_dim: int | None
 
     def _flatten_into(
         self,
@@ -1187,10 +1225,10 @@ class _DispatchCacheEntryOutputInfo:
     None or integer)
     """
 
-    inplace_idx: Optional[int]
-    metadata: Optional[TensorMetadata]
-    view_idx: Optional[int]
-    constant_value: Optional[Any] = SingletonConstant
+    inplace_idx: int | None
+    metadata: TensorMetadata | None
+    view_idx: int | None
+    constant_value: Any | None = SingletonConstant
 
 
 @dataclass(frozen=True, slots=True)
@@ -1217,7 +1255,7 @@ class _DispatchCacheBypassEntry:
 
 
 if TYPE_CHECKING:
-    _DispatchCacheEntry = Union[_DispatchCacheValidEntry, _DispatchCacheBypassEntry]
+    _DispatchCacheEntry = _DispatchCacheValidEntry | _DispatchCacheBypassEntry
 
 
 @dataclass(frozen=True, slots=True)
@@ -1260,8 +1298,8 @@ class FakeTensorMode(TorchDispatchMode):
     epoch: int = 0
     in_kernel_invocation: bool = False
     static_shapes: bool
-    shape_env: Optional[ShapeEnv]
-    _stack: Optional[str]
+    shape_env: ShapeEnv | None
+    _stack: str | None
     allow_meta: bool
 
     # NestedTensor uses a tensor_id_counter to uniquely identify offsets.
@@ -1279,8 +1317,8 @@ def __init__(
         *,
         allow_fallback_kernels: bool = True,
         allow_non_fake_inputs: bool = False,
-        shape_env: Optional[ShapeEnv] = None,
-        static_shapes: Optional[bool] = None,
+        shape_env: ShapeEnv | None = None,
+        static_shapes: bool | None = None,
         # TODO: This is a temporary measure, see
         # https://github.com/pytorch/pytorch/pull/126245#discussion_r1604185748
         # We're currently solely using this to impede population of
@@ -1349,9 +1387,7 @@ def __init__(
         # in_kernel_invocation
         # If another fake mode was already active when we enter, we also stash it here.
         # That way when we exit, we know to re-enable the previous fake mode.
-        self.enter_stack: list[
-            tuple[bool, Optional[TorchDispatchMode], Optional[bool]]
-        ] = []
+        self.enter_stack: list[tuple[bool, TorchDispatchMode | None, bool | None]] = []
 
         self.shape_env = shape_env
 
@@ -1402,6 +1438,7 @@ def avoid_device_init(self) -> bool:
         return not (
             torch.cuda.is_available()
             or (hasattr(torch, "hpu") and torch.hpu.is_available())
+            or _is_privateuse1_backend_available()
         )
 
     @property
@@ -1726,7 +1763,7 @@ def _validate_cache_key(
     def _prep_args_for_hash(
         self,
         result: list[object],
-        args: Union[Mapping[str, object], Sequence[object], Iterable[object]],
+        args: Mapping[str, object] | Sequence[object] | Iterable[object],
         state: _CacheKeyState,
         id_hashed_objects: list[object],
     ) -> None:
@@ -1801,7 +1838,7 @@ def _validate_output_for_cache_entry(
         func: OpOverload,
         args: Sequence[object],
         kwargs: Mapping[str, object],
-        output: Optional[FakeTensor],
+        output: FakeTensor | None,
     ) -> None:
         # Is this even possible? According to the signature this can be None but
         # not `int`. So either the signature is a lie or (part of) this line is
@@ -1920,7 +1957,7 @@ def _make_cache_entry(
         func: OpOverload,
         args: Sequence[object],
         kwargs: Mapping[str, object],
-        output: Optional[FakeTensor],
+        output: FakeTensor | None,
     ) -> _DispatchCacheValidEntry:
         """
         Make a cache entry object for the given 'output' Tensor. Raises
@@ -2020,7 +2057,7 @@ def _get_output_tensor_from_cache_entry(
         key: _DispatchCacheKey,
         func: OpOverload,
         args: Sequence[object],
-    ) -> Optional[FakeTensor]:
+    ) -> FakeTensor | None:
         if (
             entry.inplace_idx is None
             and entry.metadata is None
@@ -2046,9 +2083,7 @@ def _get_output_tensor_from_cache_entry(
         if is_sparse_any(metadata):
             raise AssertionError("Sparse tensors are not supported in cache")
 
-        def check_value(
-            value: _MetadataIntLike, state: _CacheKeyState
-        ) -> Union[IntLikeType]:
+        def check_value(value: _MetadataIntLike, state: _CacheKeyState) -> IntLikeType:
             if isinstance(value, _SymIntOutputStub):
                 if state.shape_env is None:
                     raise AssertionError(
@@ -2105,7 +2140,7 @@ def _output_from_cache_entry(
         key: _DispatchCacheKey,
         func: OpOverload,
         args: Sequence[object],
-    ) -> Union[Optional[FakeTensor], tuple[Optional[FakeTensor], ...]]:
+    ) -> FakeTensor | None | tuple[FakeTensor | None, ...]:
         """
         Create a new FakeTensor from the cache entry.
         """
@@ -2125,7 +2160,7 @@ def _output_from_cache_entry(
 
     def _crosscheck_cache_output(
         self,
-        output: Union[Optional[FakeTensor], tuple[Optional[FakeTensor], ...]],
+        output: FakeTensor | None | tuple[FakeTensor | None, ...],
         func: OpOverload,
         types: Sequence[type],
         args: Sequence[object],
@@ -2210,7 +2245,7 @@ def dispatch(
 
     def _maybe_infer_fake(
         self, func: OpOverload, path: KeyPath, fake: object, real: object
-    ) -> tuple[Optional[object], bool]:
+    ) -> tuple[object | None, bool]:
         """
         Helper to cross-check fake/real output properties & values,
         and create new fake vals if mismatched.
@@ -2328,7 +2363,7 @@ def _maybe_infer_fake_kernel_from_pytree_out(
         real_in: object,
         fake_out: object,
         real_out: object,
-    ) -> Optional[object]:
+    ) -> object | None:
         """
         Helper to cross-check fake/real output properties & values,
         and create new fake vals if mismatched, but at the kernel level.
@@ -2403,7 +2438,7 @@ def _dispatch_impl(
         types: Sequence[type],
         args: Sequence[object],
         kwargs: Mapping[str, object],
-    ) -> Optional[FakeTensor]:
+    ) -> FakeTensor | None:
         from torch._higher_order_ops.utils import registered_hop_fake_fns
 
         flat_args, args_spec = pytree.tree_flatten((args, kwargs))
@@ -2587,7 +2622,7 @@ def _dispatch_impl(
 
         def maybe_to_real_tensor(
             t: T,
-        ) -> Optional[Union[T, Tensor, torch._C.ScriptObject]]:
+        ) -> T | Tensor | torch._C.ScriptObject | None:
             if isinstance(t, FakeTensor):
                 return t.real_tensor
             elif isinstance(t, py_sym_types):
@@ -2708,9 +2743,7 @@ def go(t: object, real_t: Tensor) -> None:
                                 "self.shape_env must not be None for symbolic Eq"
                             )
 
-                        self.shape_env.set_real_tensor_prop_unbacked_vals(
-                            s, int(real_t)
-                        )
+                        self.shape_env.set_real_tensor_prop_unbacked_vals(s, real_t)
 
             if real_out is not nil:
                 # cross check fake/real outputs, and optionally override fake kernel mismatches
@@ -2871,13 +2904,15 @@ def go(t: object, real_t: Tensor) -> None:
         # and then afterwards wrapping them to a FakeTensor
         for run_impl_check, op_impl in op_implementations_checks:
             if run_impl_check(func):
+                # pyrefly: ignore [bad-argument-count]
                 op_impl_out = op_impl(self, func, *args, **kwargs)
                 if op_impl_out is not NotImplemented:
+                    # pyrefly: ignore [bad-return]
                     return maybe_propagate_real_tensors(op_impl_out)
 
         def maybe_run_unsafe_fallback(
-            error: Optional[RuntimeError] = None,
-        ) -> Optional[FakeTensor]:
+            error: RuntimeError | None = None,
+        ) -> FakeTensor | None:
             # We infer the meta of a custom ops that return None to just
             # return None. custom ops are not allowed to mutate metadata
             # of their inputs, so this is safe.
@@ -2955,7 +2990,7 @@ def validate_and_convert_non_fake_tensors(
         """
         flat_arg_fake_tensors: list[FakeTensor] = []
 
-        def validate(x: T) -> Union[T, FakeTensor]:
+        def validate(x: T) -> T | FakeTensor:
             if not isinstance(x, Tensor):
                 return x
 
@@ -2973,7 +3008,9 @@ def validate(x: T) -> Union[T, FakeTensor]:
                 )
                 if not allow_non_fake_inputs:
                     if isinstance(x, FakeTensor) and x.fake_mode is not self:
-                        raise AssertionError("Mixing fake modes NYI")
+                        raise AssertionError(
+                            f"Mixing fake modes NYI x.fake_mode={x.fake_mode} vs self={self}"
+                        )
                     args, kwargs = pytree.tree_unflatten(flat_args, args_spec)
                     raise AssertionError(
                         f"Please convert all Tensors to FakeTensors first or instantiate FakeTensorMode "
@@ -3003,7 +3040,7 @@ def wrap_meta_outputs_with_default_device_logic(
         common_device = None
         has_scalar_only_inputs = False
 
-        def wrap(e: T) -> Union[T, FakeTensor]:
+        def wrap(e: T) -> T | FakeTensor:
             nonlocal common_device
             nonlocal has_scalar_only_inputs
 
@@ -3040,8 +3077,8 @@ def wrap(e: T) -> Union[T, FakeTensor]:
         return tree_map(wrap, r)
 
     def create_symbolic_nested_int(
-        self, *, nt_tensor_id: Optional[int] = None
-    ) -> torch.SymInt:
+        self, *, nt_tensor_id: int | None = None
+    ) -> IntLikeType:
         # See Note: [Creating symbolic nested int]
         # Returned nested int always has coeff=1; multiply the result by coeff if needed
         import torch.nested._internal.nested_tensor
@@ -3133,12 +3170,12 @@ def from_tensor(
         self,
         tensor: Tensor,
         *,
-        static_shapes: Optional[bool] = None,
-        source: Optional[Source] = None,
-        symbolic_context: Optional[SymbolicContext] = None,
+        static_shapes: bool | None = None,
+        source: Source | None = None,
+        symbolic_context: SymbolicContext | None = None,
         trace: bool = True,
     ) -> FakeTensor:
-        shape_env: Optional[ShapeEnv] = self.shape_env
+        shape_env: ShapeEnv | None = self.shape_env
         if static_shapes is None:
             static_shapes = self.static_shapes
         if static_shapes:
@@ -3242,7 +3279,7 @@ def run_fallback_kernel(
     # REAL compute (not with meta device)
     with no_dispatch():
 
-        def to_real_tensor(e: T) -> Union[T, Tensor]:
+        def to_real_tensor(e: T) -> T | Tensor:
             if fake_mode.is_our_fake(e):
                 out = torch.zeros_like(e, device=e.fake_device)
                 if e.is_sparse:
@@ -3268,7 +3305,7 @@ def to_real_tensor(e: T) -> Union[T, Tensor]:
     # not be set up, bc of conversion to device, unless we can reuse an
     # input impl
 
-    def map_out(e: T) -> Union[T, FakeTensor]:
+    def map_out(e: T) -> T | FakeTensor:
         if id(e) not in inp_impls and (
             isinstance(e, Tensor)
             and not is_sparse_any(e)
@@ -3315,7 +3352,7 @@ def __torch_function__(
         func: OpOverload,
         types: Sequence[type],
         args: Sequence[object] = (),
-        kwargs: Optional[Mapping[str, object]] = None,
+        kwargs: Mapping[str, object] | None = None,
     ) -> FakeTensor:
         kwargs = kwargs if kwargs else {}
 
diff --git a/torch/_subclasses/functional_tensor.py b/torch/_subclasses/functional_tensor.py
index 46c9403c36f84..17d45e2e3687e 100644
--- a/torch/_subclasses/functional_tensor.py
+++ b/torch/_subclasses/functional_tensor.py
@@ -1,11 +1,12 @@
 from __future__ import annotations
 
 import contextlib
+import functools
 import warnings
 import weakref
 from abc import ABC, abstractmethod
 from contextlib import AbstractContextManager
-from typing import Any, Optional, TYPE_CHECKING, Union
+from typing import Any, TYPE_CHECKING
 from typing_extensions import Self
 
 
@@ -35,6 +36,66 @@
 not_implemented_log = torch._logging.getArtifactLogger(__name__, "not_implemented")
 
 
+def _has_unrecognized_tensor_types(types: Sequence[type]) -> bool:
+    unrecognized_types = [
+        t
+        for t in types
+        if t not in (torch.Tensor, torch._subclasses.FakeTensor, FunctionalTensor)
+    ]
+    if unrecognized_types:
+        not_implemented_log.debug(
+            "FunctionalTensor unrecognized subclass(es): %s", unrecognized_types
+        )
+    return bool(unrecognized_types)
+
+
+@functools.lru_cache(maxsize=512)
+def _can_decompose_fast(
+    func: OpOverload, export: bool, pre_dispatch: bool
+) -> bool | None:
+    """Fast path for _can_decompose that depends only on (func, export, pre_dispatch).
+
+    Returns True/False for a definitive answer, or None to fall through
+    to the slow path (autograd_would_have_decomposed).
+    """
+    if export and func is torch.ops.aten.dropout.default:
+        return False
+
+    from torch._decomp import _should_decompose_because_unsafe_op
+
+    if _should_decompose_because_unsafe_op(func):
+        return True
+
+    alias_info_present = any(arg.alias_info for arg in func._schema.arguments)
+    if alias_info_present or func._schema.is_mutable:
+        return True
+
+    if export:
+        if pre_dispatch:
+            if func.namespace not in ("aten", "prim") and func._can_decompose():
+                warnings.warn(
+                    f"At pre-dispatch tracing, we assume that any custom op marked with "
+                    f"CompositeImplicitAutograd and have functional schema are safe to not decompose. "
+                    f"Found {func} to be one such op.",
+                    stacklevel=3,
+                )
+            return False
+        return True
+
+    return None
+
+
+def _assert_functionalize_not_active(msg: str) -> None:
+    is_included = torch._C._dispatch_tls_is_dispatch_key_included(
+        torch._C.DispatchKey.Functionalize
+    )
+    is_excluded = torch._C._dispatch_tls_is_dispatch_key_excluded(
+        torch._C.DispatchKey.Functionalize
+    )
+    if not is_excluded and is_included:
+        raise AssertionError(msg)
+
+
 # NOTE Some special handling for tensor conversion during export is needed.
 # Normally, when tracing through the model with tensor.to(), the maybe-aliasing
 # relationship between input and output tensors will be baked into the graph.
@@ -66,7 +127,7 @@ class FunctionalTensor(torch.Tensor):
     This class is a lightweight python shim around the C++ functionalization logic.
 
     FunctionalTensor is required to be used with a corresponding
-    FunctionalTensormode active, because it relies
+    FunctionalTensorMode active, because it relies
     on using the mode for dispatch (which can properly handle factory functions).
     """
 
@@ -88,25 +149,27 @@ class FunctionalTensor(torch.Tensor):
 
     # These are all aten ops that correspond to metadata queries.
     # We want FunctionalTensor to be able to handle them directly.
-    metadata_fns = [
-        torch.ops.aten.is_contiguous.default,
-        torch.ops.aten.is_contiguous.memory_format,
-        torch.ops.aten.is_strides_like_format.default,
-        torch.ops.aten.is_non_overlapping_and_dense.default,
-        torch.ops.aten.size.default,
-        torch.ops.aten.sym_size.default,
-        torch.ops.aten.stride.default,
-        torch.ops.aten.sym_stride.default,
-        torch.ops.aten.storage_offset.default,
-        torch.ops.aten.sym_storage_offset.default,
-        torch.ops.aten.numel.default,
-        torch.ops.aten.sym_numel.default,
-        torch.ops.aten.dim.default,
-        torch.ops.prim.device.default,
-    ]
+    metadata_fns = frozenset(
+        {
+            torch.ops.aten.is_contiguous.default,
+            torch.ops.aten.is_contiguous.memory_format,
+            torch.ops.aten.is_strides_like_format.default,
+            torch.ops.aten.is_non_overlapping_and_dense.default,
+            torch.ops.aten.size.default,
+            torch.ops.aten.sym_size.default,
+            torch.ops.aten.stride.default,
+            torch.ops.aten.sym_stride.default,
+            torch.ops.aten.storage_offset.default,
+            torch.ops.aten.sym_storage_offset.default,
+            torch.ops.aten.numel.default,
+            torch.ops.aten.sym_numel.default,
+            torch.ops.aten.dim.default,
+            torch.ops.prim.device.default,
+        }
+    )
 
     # Used by auto_functionalize to determine base of tensors during inference mode.
-    _inference_mode_base: Optional[FunctionalTensor] = None
+    _inference_mode_base: FunctionalTensor | None = None
 
     def __new__(cls, elem: torch.Tensor, mode: FunctionalTensorMode) -> Self:
         if not torch._is_functional_tensor(elem):
@@ -182,15 +245,7 @@ def __torch_dispatch__(  # type: ignore[override]
         args: tuple[Any, ...] = (),
         kwargs: dict[str, Any] | None = None,
     ) -> Any:
-        unrecognized_types = [
-            t
-            for t in types
-            if t not in [torch.Tensor, torch._subclasses.FakeTensor, FunctionalTensor]
-        ]
-        if unrecognized_types:
-            not_implemented_log.debug(
-                "FunctionalTensor unrecognized subclass(es): %s", unrecognized_types
-            )
+        if _has_unrecognized_tensor_types(types):
             return NotImplemented
 
         if kwargs is None:
@@ -366,12 +421,12 @@ def __init__(
         self._allow_token_discovery = _allow_token_discovery
 
         self._storage_to_base: weakref.WeakKeyDictionary[
-            torch.storage.UntypedStorage, Optional[FunctionalTensor]
+            torch.storage.UntypedStorage, FunctionalTensor | None
         ] = weakref.WeakKeyDictionary()
 
     # No-op if FunctionalTensorMode is already in use
     def __enter__(self) -> Self:
-        def _get_prev_mode() -> Optional[FunctionalTensorMode]:
+        def _get_prev_mode() -> FunctionalTensorMode | None:
             if self._dispatch_key == torch._C.DispatchKey.PreDispatch:
                 return _get_dispatch_mode_pre_dispatch(
                     torch._C._TorchDispatchModeKey.FUNCTIONAL
@@ -389,9 +444,9 @@ def _get_prev_mode() -> Optional[FunctionalTensorMode]:
 
     def __exit__(
         self,
-        exc_type: Optional[type[BaseException]],
-        exc_val: Optional[BaseException],
-        exc_tb: Optional[TracebackType],
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: TracebackType | None,
     ) -> None:
         is_on_stack = self.enter_stack.pop()
         if is_on_stack:
@@ -407,69 +462,12 @@ def __torch_dispatch__(
         if kwargs is None:
             kwargs = {}
 
-        unrecognized_types = [
-            t
-            for t in types
-            if not issubclass(t, torch._subclasses.FakeTensor)
-            and t not in [torch.Tensor, FunctionalTensor]
-        ]
-
-        if unrecognized_types:
-            not_implemented_log.debug(
-                "FunctionalTensor unrecognized subclass(es): %s", unrecognized_types
-            )
+        if _has_unrecognized_tensor_types(types):
             return NotImplemented
 
-        def _can_decompose(func: OpOverload) -> bool:
-            # See https://github.com/pytorch/pytorch/pull/115258#issuecomment-1900755832
-            # Never decompose dropout in export
-            if self.export and func is torch.ops.aten.dropout.default:
-                return False
-
-            # We unconditionally decompose ops that are maybe aliasing or mutating ops
-            from torch._decomp import _should_decompose_because_unsafe_op
-
-            if _should_decompose_because_unsafe_op(func):
-                return True
-
-            # (1) we unconditionally decompose maybe-aliasing or maybe-mutating ops,
-            # because we must know statically of an op mutates or aliasing in order to functionalize it properly
-            # (2) for mutating ops that have CompositeImplicit decomps, we choose to decompose them today.
-            # In theory, we could walk this back and avoid decomposing them later if we need to.
-            alias_info_present = any(arg.alias_info for arg in func._schema.arguments)
-            if alias_info_present or func._schema.is_mutable:
-                return True
-
-            # If we are here, it means we are seeing functional composite op.
-            # For pre-dispatch IR, we don't want to decompose this op
-            # For post-dispatch IR, we do want to decompose this op. it is fine
-            # to decompose here even if you want to preserve a CIA in post-dispatch export
-            # because we already override decompose behaviour so it will do the
-            # right thing.
-            if self.export:
-                if self.pre_dispatch:
-                    # If it is CIA custom op, we warn that we are assuming this op is indeed functional.
-                    if func.namespace not in ["aten", "prim"] and func._can_decompose():
-                        warnings.warn(
-                            f"At pre-dispatch tracing, we assume that any custom op marked with "
-                            f"CompositeImplicitAutograd and have functional schema are safe to not decompose. "
-                            f"Found {func} to be one such op.",
-                            stacklevel=2,
-                        )
-                    return False
-                return True
-
-            # in normal torch.compile IR, we only decompose an op if autograd
-            # would have decomposed it (NB: autograd may have been skipped if
-            # we are in inference mode)
-            # TODO: the flatten here can potentially be deduped with the
-            # unwrapping pytree_map later
-            flat_args_kwargs, _ = pytree.tree_flatten((args, kwargs))
-            return autograd_would_have_decomposed(func, flat_args_kwargs)
-
         if (
             func not in FunctionalTensor.metadata_fns
-            and _can_decompose(func)
+            and self._can_decompose(func, args, kwargs)
             # Not all funcs from __torch_dispatch__ are actual dispatcher ops,
             # e.g. prim.device
             and torch._C._dispatch_has_kernel(func.name())
@@ -536,16 +534,9 @@ def unwrap(x: FunctionalTensor) -> torch.Tensor:
         # Expectation: functionalization should not **already** be enabled above our mode.
         # Why would that be bad? when we return a FunctionalTensor here, we don't want functionalization
         # to run above this mode and further wrap that output in **another** C++ FunctionalTensorWrapper.
-        is_included = torch._C._dispatch_tls_is_dispatch_key_included(
-            torch._C.DispatchKey.Functionalize
+        _assert_functionalize_not_active(
+            "Functionalization should not already be enabled above this mode"
         )
-        is_excluded = torch._C._dispatch_tls_is_dispatch_key_excluded(
-            torch._C.DispatchKey.Functionalize
-        )
-        if not is_excluded and is_included:
-            raise AssertionError(
-                "Functionalization should not already be enabled above this mode"
-            )
         include_to_set = (
             torch._C._dispatch_tls_local_include_set()
             | torch._C.DispatchKeySet(torch._C.DispatchKey.Functionalize)
@@ -588,34 +579,7 @@ def unwrap(x: FunctionalTensor) -> torch.Tensor:
                             torch.Tensor, wrap, outs_unwrapped
                         )
                     else:
-                        # Note: [Functionalization View Replay Annotation]
-                        # When functionalization encounters a mutation, it handles aliases by lazily regenerating the aliases
-                        # at the first time they are next used.
-                        # This is a problem when plumbing user annotations during tracing. We want the view ops from view replay
-                        # to have the same annotation that the user specified on the original views. But view replay in
-                        # functionalization happens the next time the alias is used (e.g. second_op(alias_with_pending_mutation)),
-                        # so when we regenerate views before calling into second_op, those views will end up getting the metadata
-                        # for second_op!
-                        #
-                        # Instead, we need to remember the node metadata from the original views, and ensure that this node metadata
-                        # is globally set when we lazily perform view replay.
-                        # The globally set metadata will be used to populate the fx node created for the replayed operation.
-                        if m := torch._C._get_dispatch_mode(
-                            torch._C._TorchDispatchModeKey.PROXY
-                        ):
-                            for a in pytree.tree_leaves([args, kwargs]):
-                                if not isinstance(a, FunctionalTensor):
-                                    continue
-                                unwrapped = torch._from_functional_tensor(a.elem)
-                                try:
-                                    tracker_entry = m.tracer.tensor_tracker[unwrapped]
-                                except KeyError:
-                                    raise RuntimeError(
-                                        f"cannot find {unwrapped} in tensor_tracker"
-                                    ) from None
-                                curr_node = tracker_entry.proxy.node
-                                with fx_traceback.set_current_replay_node(curr_node):
-                                    torch._sync(a)
+                        self._sync_view_replay_annotations(args, kwargs)
 
                         # When we dispatch to the C++ functionalization kernel, we might need to jump back to the
                         # PreDispatch mode stack afterwards, to handle any other PreDispatch modes underneath
@@ -638,16 +602,9 @@ def unwrap(x: FunctionalTensor) -> torch.Tensor:
                     torch._disable_functionalization()
                     torch._functionalize_enable_reapply_views(old_apply_views)  # type: ignore[attr-defined]
 
-        is_included = torch._C._dispatch_tls_is_dispatch_key_included(
-            torch._C.DispatchKey.Functionalize
+        _assert_functionalize_not_active(
+            "Functionalization should not already be enabled above this mode after dispatch"
         )
-        is_excluded = torch._C._dispatch_tls_is_dispatch_key_excluded(
-            torch._C.DispatchKey.Functionalize
-        )
-        if not is_excluded and is_included:
-            raise AssertionError(
-                "Functionalization should not already be enabled above this mode after dispatch"
-            )
 
         if (
             # If no outputs are our functional subclass, then don't try to fix up aliasing
@@ -675,6 +632,59 @@ def unwrap(x: FunctionalTensor) -> torch.Tensor:
         # If none of our inputs were wrapped, then we have no FunctionalTensor outputs that we need to fix up storages for.
         return return_and_correct_aliasing(func, args, kwargs, outs_wrapped)
 
+    def _sync_view_replay_annotations(
+        self,
+        args: tuple[Any, ...],
+        kwargs: dict[str, Any],
+    ) -> None:
+        """Sync FunctionalTensor args so view replay uses correct fx node metadata.
+
+        When functionalization encounters a mutation, it handles aliases by lazily
+        regenerating them at the first time they are next used. This is a problem when
+        plumbing user annotations during tracing: we want view ops from view replay to
+        have the same annotation the user specified on the original views. But view
+        replay happens the next time the alias is used (e.g.
+        second_op(alias_with_pending_mutation)), so the regenerated views would get the
+        metadata for second_op instead.
+
+        To fix this, we remember the node metadata from the original views and globally
+        set it when we lazily perform view replay. The globally set metadata will be
+        used to populate the fx node created for the replayed operation.
+        """
+        m = torch._C._get_dispatch_mode(torch._C._TorchDispatchModeKey.PROXY)
+        if m is not None:
+            for a in pytree.tree_leaves([args, kwargs]):
+                if not isinstance(a, FunctionalTensor):
+                    continue
+                unwrapped = torch._from_functional_tensor(a.elem)
+                try:
+                    tracker_entry = m.tracer.tensor_tracker[unwrapped]
+                except KeyError:
+                    raise RuntimeError(
+                        f"cannot find {unwrapped} in tensor_tracker"
+                    ) from None
+                curr_node = tracker_entry.proxy.node
+                with fx_traceback.set_current_replay_node(curr_node):
+                    torch._sync(a)
+
+    def _can_decompose(
+        self,
+        func: OpOverload,
+        args: tuple[Any, ...],
+        kwargs: dict[str, Any],
+    ) -> bool:
+        result = _can_decompose_fast(func, self.export, self.pre_dispatch)
+        if result is not None:
+            return result
+
+        # in normal torch.compile IR, we only decompose an op if autograd
+        # would have decomposed it (NB: autograd may have been skipped if
+        # we are in inference mode)
+        # TODO: the flatten here can potentially be deduped with the
+        # unwrapping pytree_map later
+        flat_args_kwargs, _ = pytree.tree_flatten((args, kwargs))
+        return autograd_would_have_decomposed(func, flat_args_kwargs)
+
     @classmethod
     def is_infra_mode(cls) -> bool:
         return True
@@ -746,26 +756,22 @@ def functionalize(self, inner_f: Callable[..., Any]) -> Callable[..., Any]:
     def redispatch_to_next(self) -> AbstractContextManager[None]:
         pass
 
-    @abstractmethod
     def replace(self, input_tensor: torch.Tensor, output_tensor: torch.Tensor) -> None:
-        pass
+        torch._functionalize_replace(input_tensor, output_tensor)
 
-    @abstractmethod
     def commit_update(self, tensor: torch.Tensor) -> None:
-        pass
+        torch._functionalize_commit_update(tensor)
 
-    @abstractmethod
     def sync(self, tensor: torch.Tensor) -> None:
-        pass
+        torch._functionalize_sync(tensor)
 
-    @abstractmethod
     def mark_mutation_hidden_from_autograd(self, tensor: torch.Tensor) -> None:
-        pass
+        torch._functionalize_mark_mutation_hidden_from_autograd(tensor)
 
 
 class PythonFunctionalizeAPI(BaseFunctionalizeAPI):
     def __init__(
-        self, mode: Optional[FunctionalTensorMode] = None, pre_dispatch: bool = False
+        self, mode: FunctionalTensorMode | None = None, pre_dispatch: bool = False
     ) -> None:
         super().__init__()
         self.mode = mode if mode else FunctionalTensorMode()
@@ -778,7 +784,7 @@ def wrap_tensors(self, args: tuple[Any]) -> tuple[Any]:
             )
 
     def unwrap_tensors(
-        self, args: Union[torch.Tensor, tuple[torch.Tensor, ...], list[torch.Tensor]]
+        self, args: torch.Tensor | tuple[torch.Tensor, ...] | list[torch.Tensor]
     ) -> Any:
         return torch.utils._pytree.tree_map_only(
             FunctionalTensor, FunctionalTensor.from_functional, args
@@ -796,35 +802,30 @@ def redispatch_to_next(self) -> AbstractContextManager[None]:
         # directly instead of globally setting it.
         return contextlib.nullcontext()
 
-    def replace(self, input_tensor: torch.Tensor, output_tensor: torch.Tensor) -> None:
-        if not isinstance(input_tensor, FunctionalTensor):
+    @staticmethod
+    def _check_cast_functional(tensor: torch.Tensor, name: str) -> FunctionalTensor:
+        if not isinstance(tensor, FunctionalTensor):
             raise AssertionError(
-                f"input_tensor must be a FunctionalTensor, got {type(input_tensor)}"
+                f"{name} must be a FunctionalTensor, got {type(tensor)}"
             )
+        return tensor
+
+    def replace(self, input_tensor: torch.Tensor, output_tensor: torch.Tensor) -> None:
+        ft = self._check_cast_functional(input_tensor, "input_tensor")
         if isinstance(output_tensor, FunctionalTensor):
             raise AssertionError("output_tensor must not be a FunctionalTensor")
-        input_tensor.replace_(output_tensor)
+        ft.replace_(output_tensor)
 
     def commit_update(self, tensor: torch.Tensor) -> None:
-        if not isinstance(tensor, FunctionalTensor):
-            raise AssertionError(
-                f"tensor must be a FunctionalTensor, got {type(tensor)}"
-            )
-        tensor.commit_update()
+        self._check_cast_functional(tensor, "tensor").commit_update()
 
     def sync(self, tensor: torch.Tensor) -> None:
-        if not isinstance(tensor, FunctionalTensor):
-            raise AssertionError(
-                f"tensor must be a FunctionalTensor, got {type(tensor)}"
-            )
-        tensor.sync()
+        self._check_cast_functional(tensor, "tensor").sync()
 
     def mark_mutation_hidden_from_autograd(self, tensor: torch.Tensor) -> None:
-        if not isinstance(tensor, FunctionalTensor):
-            raise AssertionError(
-                f"tensor must be a FunctionalTensor, got {type(tensor)}"
-            )
-        tensor.mark_mutation_hidden_from_autograd()
+        self._check_cast_functional(
+            tensor, "tensor"
+        ).mark_mutation_hidden_from_autograd()
 
 
 class CppFunctionalizeAPI(BaseFunctionalizeAPI):
@@ -834,8 +835,8 @@ def wrap_tensors(self, args: tuple[Any, ...]) -> tuple[Any, ...]:
         return _wrap_all_tensors_to_functional(args, level=0)
 
     def unwrap_tensors(
-        self, args: Union[torch.Tensor, tuple[torch.Tensor, ...]]
-    ) -> Union[torch.Tensor, tuple[torch.Tensor, ...]]:
+        self, args: torch.Tensor | tuple[torch.Tensor, ...]
+    ) -> torch.Tensor | tuple[torch.Tensor, ...]:
         from torch._functorch.eager_transforms import (
             _unwrap_all_tensors_from_functional,
         )
@@ -851,18 +852,6 @@ def redispatch_to_next(self) -> AbstractContextManager[None]:
             torch._C.DispatchKeySet(torch._C.DispatchKey.Functionalize)
         )
 
-    def replace(self, input_tensor: torch.Tensor, output_tensor: torch.Tensor) -> None:
-        torch._functionalize_replace(input_tensor, output_tensor)
-
-    def commit_update(self, tensor: torch.Tensor) -> None:
-        torch._functionalize_commit_update(tensor)
-
-    def sync(self, tensor: torch.Tensor) -> None:
-        torch._functionalize_sync(tensor)
-
-    def mark_mutation_hidden_from_autograd(self, tensor: torch.Tensor) -> None:
-        torch._functionalize_mark_mutation_hidden_from_autograd(tensor)
-
 
 class FunctorchFunctionalizeAPI(BaseFunctionalizeAPI):
     def __init__(self, interpreter: FunctionalizeInterpreter) -> None:
@@ -874,8 +863,8 @@ def wrap_tensors(self, args: tuple[Any]) -> tuple[Any]:
         return _wrap_all_tensors_to_functional(args, level=self.interpreter.level())
 
     def unwrap_tensors(
-        self, args: Union[torch.Tensor, tuple[torch.Tensor, ...]]
-    ) -> Union[torch.Tensor, tuple[torch.Tensor, ...]]:
+        self, args: torch.Tensor | tuple[torch.Tensor, ...]
+    ) -> torch.Tensor | tuple[torch.Tensor, ...]:
         from torch._functorch.eager_transforms import (
             _unwrap_all_tensors_from_functional,
         )
@@ -898,18 +887,6 @@ def functionalize(self, inner_f: Callable) -> Callable:
     def redispatch_to_next(self) -> AbstractContextManager[None]:
         return self.interpreter.lower()
 
-    def replace(self, input_tensor: torch.Tensor, output_tensor: torch.Tensor) -> None:
-        torch._functionalize_replace(input_tensor, output_tensor)
-
-    def commit_update(self, tensor: torch.Tensor) -> None:
-        torch._functionalize_commit_update(tensor)
-
-    def sync(self, tensor: torch.Tensor) -> None:
-        torch._functionalize_sync(tensor)
-
-    def mark_mutation_hidden_from_autograd(self, tensor: torch.Tensor) -> None:
-        torch._functionalize_mark_mutation_hidden_from_autograd(tensor)
-
 
 def mb_unwrap_functional_tensor(tensor: torch.Tensor) -> torch.Tensor:
     if isinstance(tensor, FunctionalTensor):
diff --git a/torch/_subclasses/meta_utils.py b/torch/_subclasses/meta_utils.py
index 7f5650e07edb2..3d2490d5e7600 100644
--- a/torch/_subclasses/meta_utils.py
+++ b/torch/_subclasses/meta_utils.py
@@ -14,12 +14,10 @@
     ClassVar,
     Generic,
     NewType,
-    Optional,
     Protocol,
     TYPE_CHECKING,
     TypeGuard,
     TypeVar,
-    Union,
 )
 from typing_extensions import override, TypedDict, TypeIs, Unpack
 
@@ -40,6 +38,7 @@
 )
 from torch._dispatch.python import enable_python_dispatcher
 from torch._logging import trace_structured
+from torch._opaque_base import OpaqueBase
 from torch.utils._mode_utils import no_dispatch
 from torch.utils._python_dispatch import is_traceable_wrapper_subclass
 from torch.utils.weak import WeakIdKeyDictionary
@@ -56,6 +55,7 @@
     # Import the following modules during type checking to enable code intelligence features,
     # Do not import unconditionally, as they import sympy and importing sympy is very slow
     from torch.fx.experimental.symbolic_shapes import ShapeEnv, SymbolicContext
+    from torch.types import IntLikeType
 
 
 def _is_fake_tensor(t: object) -> TypeIs[FakeTensor]:
@@ -71,7 +71,7 @@ def _is_fake_tensor(t: object) -> TypeIs[FakeTensor]:
 _TensorT_cov = TypeVar("_TensorT_cov", bound=torch.Tensor, covariant=True)
 
 
-def safe_is_leaf(t: Union[MetaTensorDesc[Any], torch.Tensor]) -> bool:
+def safe_is_leaf(t: MetaTensorDesc[Any] | torch.Tensor) -> bool:
     try:
         return t.is_leaf
     except RuntimeError:
@@ -79,7 +79,7 @@ def safe_is_leaf(t: Union[MetaTensorDesc[Any], torch.Tensor]) -> bool:
         return False
 
 
-def safe_grad(t: _TensorLikeT) -> Optional[_TensorLikeT]:
+def safe_grad(t: _TensorLikeT) -> _TensorLikeT | None:
     with torch._logging.hide_warnings(torch._logging._internal.safe_grad_filter):
         # pyrefly: ignore [bad-return]
         return t.grad
@@ -116,7 +116,7 @@ def disable_inference_mode_for_fake_prop() -> Generator[None, None, None]:
 
 def assert_metadata_eq(
     assert_eq: Callable[[object, object], None],
-    m1: Union[MetaTensorDesc[Any], torch.Tensor],
+    m1: MetaTensorDesc[Any] | torch.Tensor,
     m2: torch.Tensor,
     *,
     skip_symbolic: bool = False,
@@ -367,6 +367,7 @@ def describe_tensor(
                 pass
 
         attrs = None
+        opaque_attrs = None
         ctx = None
         type_v = None
         if is_traceable_wrapper_subclass_v:
@@ -375,10 +376,23 @@ def describe_tensor(
                     "Traceable wrapper subclass must have __tensor_flatten__ method"
                 )
             raw_attrs, ctx = t.__tensor_flatten__()
-            attrs = {
-                attr: self.describe_tensor(getattr(t, attr), trace=trace)
-                for attr in raw_attrs
-            }
+            attrs = {}
+            opaque_attrs = {}
+            for attr in raw_attrs:
+                inner = getattr(t, attr)
+                match inner:
+                    case torch.Tensor():
+                        attrs[attr] = self.describe_tensor(inner, trace=trace)
+                    case OpaqueBase():
+                        from torch._library.fake_class_registry import (
+                            maybe_unwrap_fake_script_object,
+                        )
+
+                        opaque_attrs[attr] = maybe_unwrap_fake_script_object(inner)
+                    case _:
+                        raise AssertionError(
+                            f"expected Tensor or OpaqueBase, got {type(inner)}"
+                        )
             type_v = type(t)
 
         from torch.nested._internal.nested_tensor import _tensor_symint_registry
@@ -485,6 +499,7 @@ def describe_tensor(
             view_func=view_func,
             # pyrefly: ignore [bad-argument-type]
             attrs=attrs,
+            opaque_attrs=opaque_attrs if opaque_attrs else None,
             ctx=ctx,
             type=type_v,
             # NB: even if functorch is enabled, don't actually save the
@@ -510,7 +525,7 @@ class MetaStorageDesc:
     size: int
     # NB: this is only populated with copy_data True, it is not directly
     # serializable in JSON, you want to do something special here anyway
-    data: Optional[torch.UntypedStorage]
+    data: torch.UntypedStorage | None
 
     def as_json(self, describer_id: _DescriberId) -> dict[str, object]:
         return {
@@ -527,8 +542,8 @@ def apply(
         self,
         t: _TensorT,
         new_base: _TensorT,
-        symint_visitor_fn: Optional[Callable[[int], int]] = None,
-        tensor_visitor_fn: Optional[Callable[[torch.Tensor], _TensorT]] = None,
+        symint_visitor_fn: Callable[[int], int] | None = None,
+        tensor_visitor_fn: Callable[[torch.Tensor], _TensorT] | None = None,
     ) -> _TensorT: ...
 
     @staticmethod
@@ -546,8 +561,8 @@ def apply(
         self,
         t: torch.Tensor,
         new_base: torch.Tensor,
-        symint_visitor_fn: Optional[Callable[[int], int]] = None,
-        tensor_visitor_fn: Optional[Callable[[torch.Tensor], FakeTensor]] = None,
+        symint_visitor_fn: Callable[[int], int] | None = None,
+        tensor_visitor_fn: Callable[[torch.Tensor], FakeTensor] | None = None,
     ) -> FakeTensor:
         return torch._subclasses.fake_tensor.FakeTensor._view_func_unsafe(
             # pyrefly: ignore [bad-argument-type]
@@ -563,8 +578,8 @@ class _CustomViewFunc(ViewFunc[_TensorT], Generic[_TensorT]):
     func: Callable[
         [
             torch.Tensor,
-            Optional[Callable[[int], int]],
-            Optional[Callable[[torch.Tensor], _TensorT]],
+            Callable[[int], int] | None,
+            Callable[[torch.Tensor], _TensorT] | None,
         ],
         _TensorT,
     ]
@@ -574,8 +589,8 @@ def apply(
         self,
         t: torch.Tensor,
         new_base: torch.Tensor,
-        symint_visitor_fn: Optional[Callable[[int], int]] = None,
-        tensor_visitor_fn: Optional[Callable[[torch.Tensor], _TensorT]] = None,
+        symint_visitor_fn: Callable[[int], int] | None = None,
+        tensor_visitor_fn: Callable[[torch.Tensor], _TensorT] | None = None,
     ) -> _TensorT:
         # ignore `t`
         return self.func(new_base, symint_visitor_fn, tensor_visitor_fn)
@@ -588,12 +603,12 @@ def apply(
 #   def mk(arg: Callable[[], torch.Tensor], device: Optional[Union[torch.device, str]] = None)
 class _MetaTensorCallback(Protocol, Generic[_TensorT_cov]):
     def __call__(
-        self, arg: Callable[[], torch.Tensor], /, *, device: Union[torch.device, str]
+        self, arg: Callable[[], torch.Tensor], /, *, device: torch.device | str
     ) -> _TensorT_cov: ...
 
 
 class _MetaTensorCallbackKwargs(TypedDict, total=False):
-    device: Union[torch.device, str]
+    device: torch.device | str
 
 
 # A callback where the device may not be provided (is optional).
@@ -647,35 +662,38 @@ class MetaTensorDesc(Generic[_TensorT]):
     # We eagerly symbolicize the associated nested int for e.g. offsets / lengths
     # metadata if that offsets is already associated with a nested int.
     # See test_construct_from_jagged_with_input_offsets_mixed_case.
-    nested_int: Optional[int] = None
+    nested_int: int | None = None
     is_traceable_wrapper_subclass: bool = False
     is_functional: bool = False
     is_conj: bool = False
     is_neg: bool = False
     is_parameter: bool = False
-    stride: Optional[tuple[int, ...]] = None
+    stride: tuple[int, ...] | None = None
     storage_offset: int = 0
     # NB: We have a choice whether or not to store the id or a direct pointer
     # to the data structure.  For ease of use, we store the data structure,
     # but this means that when we serialize, we have to swizzle these pointers
     # back into ids (so we have accurate aliasing relationships)
-    storage: Optional[MetaStorageDesc] = None
-    sparse_dim: Optional[int] = None  # is_sparse, is_sparse_compressed
-    dense_dim: Optional[int] = None  # is_sparse, is_sparse_compressed
-    is_coalesced: Optional[bool] = None  # is_sparse
-    crow_indices: Optional[MetaTensorDesc[Any]] = None  # is_sparse_compressed
-    col_indices: Optional[MetaTensorDesc[Any]] = None  # is_sparse_compressed
-    ccol_indices: Optional[MetaTensorDesc[Any]] = None  # is_sparse_compressed
-    row_indices: Optional[MetaTensorDesc[Any]] = None  # is_sparse_compressed
-    values: Optional[MetaTensorDesc[Any]] = None  # is_sparse_compressed
-    unwrapped: Optional[MetaTensorDesc[Any]] = None  # is_functorch_wrapped
-    bdim: Optional[int] = None  # is_functorch_wrapped
-    base: Optional[MetaTensorDesc[Any]] = None  # is_view
-    attrs: Optional[dict[str, MetaTensorDesc[Any]]] = (
-        None  # is_traceable_wrapper_subclass
+    storage: MetaStorageDesc | None = None
+    sparse_dim: int | None = None  # is_sparse, is_sparse_compressed
+    dense_dim: int | None = None  # is_sparse, is_sparse_compressed
+    is_coalesced: bool | None = None  # is_sparse
+    crow_indices: MetaTensorDesc[Any] | None = None  # is_sparse_compressed
+    col_indices: MetaTensorDesc[Any] | None = None  # is_sparse_compressed
+    ccol_indices: MetaTensorDesc[Any] | None = None  # is_sparse_compressed
+    row_indices: MetaTensorDesc[Any] | None = None  # is_sparse_compressed
+    values: MetaTensorDesc[Any] | None = None  # is_sparse_compressed
+    unwrapped: MetaTensorDesc[Any] | None = None  # is_functorch_wrapped
+    bdim: int | None = None  # is_functorch_wrapped
+    base: MetaTensorDesc[Any] | None = None  # is_view
+    attrs: dict[str, MetaTensorDesc[Any]] | None = None  # is_traceable_wrapper_subclass
+    # A Tensor subclass containing opaque references is almost certainly NOT
+    # serializable.
+    opaque_attrs: dict[str, OpaqueBase] | None = (
+        None  # non-tensor attrs from __tensor_flatten__
     )
-    creation_meta: Optional[CreationMeta] = None
-    grad: Optional[MetaTensorDesc[Any]] = None
+    creation_meta: CreationMeta | None = None
+    grad: MetaTensorDesc[Any] | None = None
 
     # Everything below is NOT serializable, need some more work
 
@@ -693,21 +711,21 @@ class MetaTensorDesc(Generic[_TensorT]):
         "nested_int",
     }
 
-    ctx: Optional[object] = None  # is_traceable_wrapper_subclass
-    type: Optional[type] = None  # is_traceable_wrapper_subclass
-    fake_mode: Optional[FakeTensorMode] = None
-    view_func: Optional[ViewFunc[Any]] = None
+    ctx: object | None = None  # is_traceable_wrapper_subclass
+    type: type | None = None  # is_traceable_wrapper_subclass
+    fake_mode: FakeTensorMode | None = None
+    view_func: ViewFunc[Any] | None = None
     # level looks serializable, but actually it is meaningless without
     # the functorch_stack below
-    level: Optional[int] = None  # is_functorch_wrapped
-    current_level: Optional[int] = None
-    functorch_stack: Optional[list[CInterpreter]] = None
-    autograd_meta_from: Optional[torch.Tensor] = None
+    level: int | None = None  # is_functorch_wrapped
+    current_level: int | None = None
+    functorch_stack: list[CInterpreter] | None = None
+    autograd_meta_from: torch.Tensor | None = None
 
     # This is only populated on copy_data, and typically is not used at all,
     # except for some of our meta-ification paths that don't properly use
     # storage (pro-tip: you should use storage)
-    data: Optional[torch.Tensor] = None
+    data: torch.Tensor | None = None
 
     # Faithfully serializing functorch tensors will not be too difficult.
     # We only need to consider grad/vmap interpreters, and their internal
@@ -774,18 +792,77 @@ def shape(self) -> tuple[int, ...]:
 # FakeTensor as src, we MUST NOT run the copy/clone operation.  A better way
 # to do this would be to not use no_dispatch and instead just disable fake
 # tensor mode only (allowing for subclass dispatch to occur)
-def _safe_copy(dst: torch.Tensor, src: Optional[torch.Tensor]) -> None:
+def _safe_copy(dst: torch.Tensor, src: torch.Tensor | None) -> None:
     if type(src) is not torch.Tensor:
         return
     dst.copy_(src)
 
 
-def _safe_clone(src: torch.Tensor) -> Optional[torch.Tensor]:
+def _safe_clone(src: torch.Tensor) -> torch.Tensor | None:
     if type(src) is not torch.Tensor:
         return None
     return src.clone()
 
 
+def _grad_context_compatible(
+    symbolic_context: torch.fx.experimental.symbolic_shapes.SymbolicContext,
+    grad_desc: MetaTensorDesc[torch.Tensor],
+) -> bool:
+    """Check if a symbolic_context is compatible with a grad tensor.
+
+    Returns False when the view base structure in symbolic_context doesn't
+    match the grad, which means we need a fresh symbolic context.  This
+    happens in FSDP2 where param._local_tensor is a view of an N-D padded
+    base while grad._local_tensor is a view of a 1-D flat gradient buffer.
+
+    We check at both the outer level and the inner (subclass attr) level.
+    """
+    from torch.fx.experimental.symbolic_shapes import (
+        StatelessSymbolicContext,
+        SubclassSymbolicContext,
+    )
+
+    def _view_base_compatible(
+        ctx: StatelessSymbolicContext[Any, Any],
+        grad_t: MetaTensorDesc[torch.Tensor],
+    ) -> bool:
+        vbc = ctx.view_base_context
+        if grad_t.is_view and vbc is None:
+            return False
+        if not grad_t.is_view and vbc is not None:
+            return False
+        if (
+            grad_t.is_view
+            and vbc is not None
+            and isinstance(vbc, StatelessSymbolicContext)
+            and grad_t.base is not None
+            and len(vbc.dynamic_sizes) != grad_t.base.ndim
+        ):
+            return False
+        return True
+
+    if not isinstance(symbolic_context, StatelessSymbolicContext):
+        return True
+
+    # Check outer level
+    if not _view_base_compatible(symbolic_context, grad_desc):
+        return False
+
+    # Check inner (subclass) level
+    if isinstance(symbolic_context, SubclassSymbolicContext):
+        if grad_desc.attrs is None:
+            return False
+        for attr, inner_ctx in symbolic_context.inner_contexts.items():
+            if attr not in grad_desc.attrs:
+                return False
+            if isinstance(
+                inner_ctx, StatelessSymbolicContext
+            ) and not _view_base_compatible(inner_ctx, grad_desc.attrs[attr]):
+                return False
+
+    return True
+
+
 # This is a class for converting multiple tensors into meta tensors which
 # share the same view/storage structure.  The operation model is you allocate
 # one of these, and then call it repeatedly on all the tensors you want to
@@ -818,7 +895,7 @@ def __init__(self, *, copy_data: bool = False) -> None:
     def successful(self) -> bool:
         return self.hit > 0 and self.miss == 0
 
-    def get_tensor_memo(self, t: MetaTensorDesc[Any]) -> Optional[torch.Tensor]:
+    def get_tensor_memo(self, t: MetaTensorDesc[Any]) -> torch.Tensor | None:
         return self.tensor_memo.get(t.id, None)
 
     def _checked_get_tensor_memo(self, t: MetaTensorDesc[Any]) -> _TensorT:
@@ -830,7 +907,7 @@ def _checked_get_tensor_memo(self, t: MetaTensorDesc[Any]) -> _TensorT:
     def set_tensor_memo(self, t: MetaTensorDesc[Any], v: _TensorT) -> None:
         self.tensor_memo[t.id] = v
 
-    def get_storage_memo(self, s: MetaStorageDesc) -> Optional[torch.UntypedStorage]:
+    def get_storage_memo(self, s: MetaStorageDesc) -> torch.UntypedStorage | None:
         return self.storage_memo.get(s.id, None)
 
     def set_storage_memo(self, s: MetaStorageDesc, v: torch.UntypedStorage) -> None:
@@ -870,7 +947,7 @@ def _checked_cast_tensor_t(cls, t: torch.Tensor) -> _TensorT:
     def _identity_callable(
         cls,
         t: Callable[[], torch.Tensor],
-        device: Optional[Union[torch.device, str]] = None,
+        device: torch.device | str | None = None,
     ) -> _TensorT:
         return cls._checked_cast_tensor_t(t())
 
@@ -883,6 +960,73 @@ def _backward_error(cls, t: _TensorT) -> _TensorT:
         err = errfn(t)
         return typing.cast(_TensorT, err)
 
+    def _empty_create_subclass(
+        self,
+        t: MetaTensorDesc[Any],
+        outer_size: tuple[int, ...] | None,
+        outer_stride: tuple[int, ...] | None,
+        shape_env: ShapeEnv | None,
+        symbolic_context: torch.fx.experimental.symbolic_shapes.SymbolicContext | None,
+        callback: _MetaTensorCallbackOptDevice[_TensorT],
+        source: torch._guards.Source,
+    ) -> _TensorT:
+        from torch._dynamo.source import AttrSource
+        from torch.fx.experimental.symbolic_shapes import SubclassSymbolicContext
+
+        # We are hitting plain meta_desc tensor so actually
+        # create a tensor here.
+        if t.attrs is None:
+            return self.meta_tensor(
+                t,
+                shape_env,
+                callback,
+                source,
+                symbolic_context,
+            )
+
+        inner_tensors: dict[str, torch.Tensor | OpaqueBase] = {}
+        for attr, meta_tensor_desc in t.attrs.items():
+            current_context = None
+            if symbolic_context is not None:
+                if not isinstance(symbolic_context, SubclassSymbolicContext):
+                    raise AssertionError(
+                        f"Expected SubclassSymbolicContext, got {type(symbolic_context)}"
+                    )
+                if attr not in symbolic_context.inner_contexts:
+                    raise AssertionError(
+                        f"tensor attr {attr!r} missing from inner_contexts"
+                    )
+                if (
+                    current_context_ := symbolic_context.inner_contexts[attr]
+                ) is not None:
+                    current_context = _checked_cast(
+                        torch.fx.experimental.symbolic_shapes.SymbolicContext,
+                        current_context_,
+                    )
+
+            current_source = AttrSource(source, attr)
+            inner_callback = functools.partial(callback, device=meta_tensor_desc.device)
+            new_empty_tensor = self._empty_create_subclass(
+                meta_tensor_desc,
+                meta_tensor_desc.size,
+                meta_tensor_desc.stride,
+                shape_env,
+                current_context,
+                inner_callback,
+                current_source,
+            )
+            inner_tensors[attr] = new_empty_tensor
+
+        # Pass through opaque (non-tensor) attrs
+        if t.opaque_attrs:
+            inner_tensors.update(t.opaque_attrs)
+
+        if t.type is None:
+            raise AssertionError("t.type must not be None for subclass")
+        return t.type.__tensor_unflatten__(  # type: ignore[attr-defined]
+            inner_tensors, t.ctx, outer_size, outer_stride
+        )
+
     # This function assumes that it's possible to do the conversion
     # NB: name here is used in a conventional way by Dynamo; it corresponds
     # precisely to the Source.name of the tensor we're fakeifying and
@@ -892,10 +1036,10 @@ def _backward_error(cls, t: _TensorT) -> _TensorT:
     def meta_tensor(
         self,
         t: MetaTensorDesc[Any],
-        shape_env: Optional[ShapeEnv],
+        shape_env: ShapeEnv | None,
         callback_: _MetaTensorCallback[_TensorT],
-        source: Optional[Source],
-        symbolic_context: Optional[SymbolicContext],
+        source: Source | None,
+        symbolic_context: SymbolicContext | None,
     ) -> _TensorT:
         callback: _MetaTensorCallbackOptDevice[_TensorT] = functools.partial(
             callback_, device=t.device
@@ -949,10 +1093,9 @@ def meta_tensor(
         def sym_sizes_strides_storage_offset(
             t: MetaTensorDesc[Any],
             src: torch._guards.Source,
-            symbolic_context: Optional[
-                torch.fx.experimental.symbolic_shapes.SymbolicContext
-            ] = symbolic_context,
-        ) -> tuple[tuple[int, ...], tuple[int, ...], int]:
+            symbolic_context: torch.fx.experimental.symbolic_shapes.SymbolicContext
+            | None = symbolic_context,
+        ) -> tuple[tuple[IntLikeType, ...], tuple[IntLikeType, ...], IntLikeType]:
             # local import to prevent circular import
             from torch.fx.experimental.symbolic_shapes import is_symbolic
 
@@ -1004,9 +1147,8 @@ def sym_sizes_strides_storage_offset(
         def empty_create(
             inner_t: MetaTensorDesc[Any],
             inner_src: torch._guards.Source,
-            symbolic_context: Optional[
-                torch.fx.experimental.symbolic_shapes.SymbolicContext
-            ] = symbolic_context,
+            symbolic_context: torch.fx.experimental.symbolic_shapes.SymbolicContext
+            | None = symbolic_context,
         ) -> torch.Tensor:
             (
                 inner_sizes,
@@ -1024,14 +1166,12 @@ def empty_create(
         # symbolic context.
         def empty_create_subclass(
             t: MetaTensorDesc[Any],
-            outer_size: tuple[int, ...],
-            outer_stride: tuple[int, ...],
-            symbolic_context: Optional[
-                torch.fx.experimental.symbolic_shapes.SymbolicContext
-            ] = symbolic_context,
-            source: Optional[torch._guards.Source] = source,
+            outer_size: tuple[IntLikeType, ...],
+            outer_stride: tuple[IntLikeType, ...],
+            symbolic_context: torch.fx.experimental.symbolic_shapes.SymbolicContext
+            | None = symbolic_context,
+            source: torch._guards.Source | None = source,
         ) -> _TensorT:
-            from torch._dynamo.source import AttrSource
             from torch.fx.experimental.symbolic_shapes import SubclassSymbolicContext
 
             if t.attrs is None:
@@ -1061,67 +1201,18 @@ def empty_create_subclass(
                     f"Expected SubclassSymbolicContext or None, got {type(symbolic_context)}"
                 )
 
-            def _empty_create_subclass(
-                t: MetaTensorDesc[Any],
-                outer_size: Optional[tuple[int, ...]],
-                outer_stride: Optional[tuple[int, ...]],
-                symbolic_context: Optional[
-                    torch.fx.experimental.symbolic_shapes.SymbolicContext
-                ],
-                callback: _MetaTensorCallbackOptDevice[_TensorT],
-                source: torch._guards.Source,
-            ) -> _TensorT:
-                # We are hitting plain meta_desc tensor so actually
-                # create a tensor here.
-                if t.attrs is None:
-                    return self.meta_tensor(
-                        t,
-                        shape_env,
-                        callback,
-                        source,
-                        symbolic_context,
-                    )
-
-                inner_tensors = {}
-                for attr, meta_tensor_desc in t.attrs.items():
-                    current_context = None
-                    if symbolic_context is not None:
-                        if not isinstance(symbolic_context, SubclassSymbolicContext):
-                            raise AssertionError(
-                                f"Expected SubclassSymbolicContext, got {type(symbolic_context)}"
-                            )
-                        if (
-                            current_context_ := symbolic_context.inner_contexts[attr]
-                        ) is not None:
-                            current_context = _checked_cast(
-                                torch.fx.experimental.symbolic_shapes.SymbolicContext,
-                                current_context_,
-                            )
-
-                    current_source = AttrSource(source, attr)
-                    inner_callback = functools.partial(
-                        callback, device=meta_tensor_desc.device
-                    )
-                    new_empty_tensor = _empty_create_subclass(
-                        meta_tensor_desc,
-                        meta_tensor_desc.size,
-                        meta_tensor_desc.stride,
-                        current_context,
-                        inner_callback,
-                        current_source,
-                    )
-                    inner_tensors[attr] = new_empty_tensor
-
-                if t.type is None:
-                    raise AssertionError("t.type must not be None for subclass")
-                return t.type.__tensor_unflatten__(  # type: ignore[attr-defined]
-                    inner_tensors, t.ctx, outer_size, outer_stride
-                )
-
             if source is None:
                 raise AssertionError("source must not be None")
-            sub = _empty_create_subclass(
-                t, outer_size, outer_stride, symbolic_context, callback, source
+            sub = self._empty_create_subclass(
+                t,
+                # pyrefly: ignore[bad-argument-type]
+                outer_size,
+                # pyrefly: ignore[bad-argument-type]
+                outer_stride,
+                shape_env,
+                symbolic_context,
+                callback,
+                source,
             )
 
             # NB: Purposefully guard here to simplify the inner / outer symbols.
@@ -1147,7 +1238,7 @@ def _empty_create_subclass(
         def all_dynamic_symbolic_context(
             t: MetaTensorDesc[Any],
             source: torch._guards.Source,
-            shape_env: Optional[torch.fx.experimental.symbolic_shapes.ShapeEnv],
+            shape_env: torch.fx.experimental.symbolic_shapes.ShapeEnv | None,
             callback: _MetaTensorCallback[_TensorT],
         ) -> torch.fx.experimental.symbolic_shapes.SymbolicContext:
             from torch._dynamo.source import AttrSource
@@ -1157,9 +1248,9 @@ def all_dynamic_symbolic_context(
                 SubclassSymbolicContext,
             )
 
-            view_base_context: Optional[
-                torch.fx.experimental.symbolic_shapes.SymbolicContext
-            ] = None
+            view_base_context: (
+                torch.fx.experimental.symbolic_shapes.SymbolicContext | None
+            ) = None
             if t.is_view:
                 if t.base is None:
                     raise AssertionError("t.base must not be None for view tensor")
@@ -1231,9 +1322,8 @@ def all_dynamic_symbolic_context(
         def view_from_base(
             base: _TensorT,
             t: MetaTensorDesc[Any],
-            shape_env: Optional[
-                torch.fx.experimental.symbolic_shapes.ShapeEnv
-            ] = shape_env,
+            shape_env: torch.fx.experimental.symbolic_shapes.ShapeEnv
+            | None = shape_env,
         ) -> _TensorT:
             with enable_python_dispatcher():
                 # fake-ify t's metadata according to the outer symbolic context
@@ -1258,7 +1348,7 @@ def view_from_base(
                     sym_eq,
                 )
 
-                def symint_visitor_fn(s: int) -> int:
+                def symint_visitor_fn(s: int) -> IntLikeType:
                     nonlocal symbolic_context
                     from torch.fx.experimental.symbolic_shapes import DimDynamic
 
@@ -1312,15 +1402,17 @@ def symint_visitor_fn(s: int) -> int:
                     )
                     attrs, _ = fake_t.__tensor_flatten__()  # type: ignore[attr-defined]
                     for attr in attrs:
-                        real_to_fake_mapping[t.attrs[attr].id] = getattr(fake_t, attr)
+                        if attr in t.attrs:
+                            real_to_fake_mapping[t.attrs[attr].id] = getattr(
+                                fake_t, attr
+                            )
 
                 def tensor_visitor_fn(
                     visited_t: torch.Tensor,
                     # These arguments are never passed, we just use them to close
                     # over these relevant values
-                    shape_env: Optional[
-                        torch.fx.experimental.symbolic_shapes.ShapeEnv
-                    ] = shape_env,
+                    shape_env: torch.fx.experimental.symbolic_shapes.ShapeEnv
+                    | None = shape_env,
                     callback: _MetaTensorCallbackOptDevice[_TensorT] = callback,
                 ) -> torch.Tensor:
                     # It's possible to close over an undefined tensor (e.g. NJT's lengths).
@@ -1363,7 +1455,11 @@ def tensor_visitor_fn(
                 # NB: we do NOT suppress guards here, we need to remove ephemeral
                 # sources
                 fake_t = t.view_func.apply(
-                    t, base, symint_visitor_fn, tensor_visitor_fn
+                    t,
+                    base,
+                    # pyrefly: ignore[bad-argument-type]
+                    symint_visitor_fn,
+                    tensor_visitor_fn,
                 )
 
                 # Ensure the output has symbolic shapes according to the outer symbolic context.
@@ -1719,22 +1815,21 @@ def _to_fake_tensor(t: MetaTensorDesc[Any]) -> _TensorT:
                         if symbolic_context.view_base_context is not None:
                             base_symbolic_context = symbolic_context.view_base_context
 
-                    base = self.meta_tensor(
-                        t.base,
-                        shape_env,
-                        callback,
-                        torch._dynamo.source.AttrSource(source, "_base"),
-                        base_symbolic_context,
-                    )
-                    # If the base tensor has unbacked symbols (e.g., from mark_unbacked),
-                    # we need to bind them now. Otherwise they'll be left pending and
-                    # compute_unbacked_bindings will fail when called for the view
-                    # (since the view may have concrete shapes, not the unbacked symbols).
-                    from torch.fx.experimental.symbolic_shapes import (
-                        compute_unbacked_bindings,
-                    )
-
-                    compute_unbacked_bindings(shape_env, base)
+                    # We don't track the example values of base tensors.
+                    # Allocation of unbacked symbols can happen when mark_unbacked
+                    # is called on a base tensor that have unbacked_indices.
+                    with (
+                        shape_env.ignore_fresh_unbacked_symbols()
+                        if shape_env
+                        else contextlib.nullcontext()
+                    ):
+                        base = self.meta_tensor(
+                            t.base,
+                            shape_env,
+                            callback,
+                            torch._dynamo.source.AttrSource(source, "_base"),
+                            base_symbolic_context,
+                        )
 
                     def is_c_of_r(
                         complex_dtype: torch.dtype, real_dtype: torch.dtype
@@ -1995,15 +2090,28 @@ def is_c_of_r(
                 if t.grad is not None:
                     from torch._dynamo.source import AttrSource
 
-                    # TODO: Use a valid grad-specific symbolic context instead of recycling
-                    # the one from t. This isn't correct if e.g. t._is_view() != t.grad._is_view().
+                    grad_source = AttrSource(source, "grad")
+                    grad_symbolic_context = symbolic_context
+
+                    # The param's symbolic_context may be incompatible with the
+                    # grad when they have different view base dimensionalities.
+                    # This happens in FSDP2 where param._local_tensor is a view
+                    # of an N-D padded base but grad._local_tensor is a view of
+                    # a 1-D flat gradient buffer.  Build a fresh context only in
+                    # that case to preserve FX graph cache consistency.
+                    if shape_env is not None and symbolic_context is not None:
+                        if not _grad_context_compatible(symbolic_context, t.grad):
+                            grad_symbolic_context = all_dynamic_symbolic_context(
+                                t.grad, grad_source, shape_env, callback
+                            )
+
                     # pyrefly: ignore [unbound-name]
                     r.grad = self.meta_tensor(
                         t.grad,
                         shape_env,
                         callback,
-                        AttrSource(source, "grad"),
-                        symbolic_context,
+                        grad_source,
+                        grad_symbolic_context,
                     )
                 # pyrefly: ignore [unbound-name]
                 torch._C._set_conj(r, t.is_conj)
@@ -2046,11 +2154,11 @@ def is_c_of_r(
     def __call__(
         self,
         t: torch.Tensor,
-        shape_env: Optional[ShapeEnv] = None,
+        shape_env: ShapeEnv | None = None,
         *,
-        callback: Optional[_MetaTensorCallback[_TensorT]] = None,
-        source: Optional[Source] = None,
-        symbolic_context: Optional[SymbolicContext] = None,
+        callback: _MetaTensorCallback[_TensorT] | None = None,
+        source: Source | None = None,
+        symbolic_context: SymbolicContext | None = None,
         # Controls whether or not we should dump the tensor metadata to structured logs
         # when source is not None.  Because we refakify after Dynamo is done,
         # we don't want to dump info again from AOTAutograd, it is redundant.
diff --git a/torch/_tensor.py b/torch/_tensor.py
index f836599e69549..77436cc17d887 100644
--- a/torch/_tensor.py
+++ b/torch/_tensor.py
@@ -43,7 +43,7 @@ def wrapped(self: _TensorLike, *args: _P.args, **kwargs: _P.kwargs) -> "Tensor":
             # See https://github.com/pytorch/pytorch/issues/75462
             sargs = self, *args
             if has_torch_function(sargs):
-                return handle_torch_function(wrapped, sargs, *sargs, **kwargs)
+                return handle_torch_function(wrapped, sargs, self, *args, **kwargs)
             return f(self, *args, **kwargs)
         except TypeError:
             return NotImplemented
@@ -900,7 +900,7 @@ def norm(
         keepdim=False,
         dtype=None,
     ):
-        r"""See :func:`torch.norm`"""
+        r"""See :func:`torch.linalg.norm`"""
         if has_torch_function_unary(self):
             return handle_torch_function(
                 Tensor.norm, (self,), self, p=p, dim=dim, keepdim=keepdim, dtype=dtype
@@ -1201,9 +1201,9 @@ def __len__(self):
         return self.shape[0]
 
     def __iter__(self):
-        # NB: we use 'imap' and not 'map' here, so that in Python 2 we get a
-        # generator and don't eagerly perform all the indexes.  This could
-        # save us work, and also helps keep trace ordering deterministic
+        # NB: we use 'imap' and not 'map' here, so that we get a generator
+        # and don't eagerly perform all the indexes.  This could save us
+        # work, and also helps keep trace ordering deterministic
         # (e.g., if you zip(*hiddens), the eager map will force all the
         # indexes of hiddens[0] before hiddens[1], while the generator
         # map will interleave them.)
@@ -1294,10 +1294,10 @@ def __cuda_array_interface__(self):
         """
         if has_torch_function_unary(self):
             # TODO mypy doesn't support @property, see: https://github.com/python/mypy/issues/6185
-            return handle_torch_function(
+            return handle_torch_function(  # pyrefly: ignore [bad-argument-count]
                 Tensor.__cuda_array_interface__.__get__,  # type: ignore[attr-defined]
                 (self,),
-                self,
+                self,  # pyrefly: ignore [bad-argument-type]
             )
 
         # raise AttributeError for unsupported tensors, so that
@@ -1364,8 +1364,6 @@ def refine_names(self, *names):  # pyrefly: ignore  # bad-override
         :attr:`names` to the same length as ``self.dim()`` using names from the
         corresponding indices of ``self.names``.
 
-        Python 2 does not support Ellipsis but one may use a string literal
-        instead (``'...'``).
 
         Args:
             names (iterable of str): The desired names of the output tensor. May
@@ -1408,8 +1406,6 @@ def align_to(self, *names):  # pyrefly: ignore  # bad-override
         that are not mentioned in :attr:`names`, in the order that they appear
         in :attr:`self`.
 
-        Python 2 does not support Ellipsis but one may use a string literal
-        instead (``'...'``).
 
         Args:
             names (iterable of str): The desired dimension ordering of the
diff --git a/torch/_torch_docs.py b/torch/_torch_docs.py
index 062f321872f8a..a95fac98b309e 100644
--- a/torch/_torch_docs.py
+++ b/torch/_torch_docs.py
@@ -339,7 +339,7 @@ def merge_dicts(*dicts):
     r"""
 index_copy(input: Tensor, dim: int, index: Tensor, source: Tensor, *, out: Optional[Tensor]) -> Tensor
 
-See :meth:`~Tensor.index_add_` for function description.
+See :meth:`~Tensor.index_copy_` for function description.
 """,
 )
 
@@ -1235,7 +1235,7 @@ def merge_dicts(*dicts):
 add_docstr(
     torch.asarray,
     r"""
-asarray(obj: Any, *, dtype: Optional[dtype], device: Optional[DeviceLikeType], copy: Optional[bool] = None, requires_grad: bool = False) -> Tensor # noqa: B950
+asarray(obj: Any, *, dtype: Optional[dtype], device: Optional[DeviceLikeType], copy: Optional[bool] = None, requires_grad: Optional[bool] = None) -> Tensor # noqa: B950
 
 Converts :attr:`obj` to a tensor.
 
@@ -1249,13 +1249,15 @@ def merge_dicts(*dicts):
 6. a sequence of scalars
 
 When :attr:`obj` is a tensor, NumPy array, or DLPack capsule the returned tensor will,
-by default, not require a gradient, have the same datatype as :attr:`obj`, be on the
-same device, and share memory with it. These properties can be controlled with the
-:attr:`dtype`, :attr:`device`, :attr:`copy`, and :attr:`requires_grad` keyword arguments.
-If the returned tensor is of a different datatype, on a different device, or a copy is
-requested then it will not share its memory with :attr:`obj`. If :attr:`requires_grad`
-is ``True`` then the returned tensor will require a gradient, and if :attr:`obj` is
-also a tensor with an autograd history then the returned tensor will have the same history.
+by default, have the same requires_grad as :attr:`obj` (defaulting to False), have the
+same datatype, be on the same device, and share memory with it. These properties can be
+controlled with the :attr:`dtype`, :attr:`device`, :attr:`copy`, and
+:attr:`requires_grad` keyword arguments. If the returned tensor is of a different
+datatype, on a different device, or a copy is requested then it will not share its
+memory with :attr:`obj`. If :attr:`requires_grad` is ``True`` (or ``None``, and
+:attr:`obj` was a tensor with requires_grad set), then the returned tensor will require
+a gradient, and if :attr:`obj` is also a tensor with an autograd history then the
+returned tensor will have the same history.
 
 When :attr:`obj` is not a tensor, NumPy array, or DLPack capsule but implements Python's
 buffer protocol then the buffer is interpreted as an array of bytes grouped according to
@@ -1298,10 +1300,10 @@ def merge_dicts(*dicts):
            Default: ``None``, which causes the device of :attr:`obj` to be used. Or, if
            :attr:`obj` is a Python sequence, the current default device will be used.
     requires_grad (bool, optional): whether the returned tensor requires grad.
-           Default: ``False``, which causes the returned tensor not to require a gradient.
-           If ``True``, then the returned tensor will require a gradient, and if :attr:`obj`
-           is also a tensor with an autograd history then the returned tensor will have
-           the same history.
+           Default: ``None``, which causes requires_grad for the returned tensor to be
+           inferred from :attr:`obj`. If ``True``, then the returned tensor will require
+           a gradient, and if :attr:`obj` is also a tensor with an autograd history then
+           the returned tensor will have the same history.
 
 Example::
 
@@ -1320,13 +1322,17 @@ def merge_dicts(*dicts):
     >>> b
     tensor([3., 4., 5.], grad_fn=<AddBackward0>)
     >>> # Shares memory with tensor 'b', with no grad
-    >>> c = torch.asarray(b)
+    >>> c = torch.asarray(b, requires_grad=False)
     >>> c
     tensor([3., 4., 5.])
     >>> # Shares memory with tensor 'b', retaining autograd history
     >>> d = torch.asarray(b, requires_grad=True)
     >>> d
     tensor([3., 4., 5.], grad_fn=<AddBackward0>)
+    >>> # Shares memory with tensor 'b', retaining autograd history
+    >>> e = torch.asarray(b)
+    >>> e
+    tensor([3., 4., 5.], grad_fn=<AddBackward0>)
 
     >>> array = numpy.array([1, 2, 3])
     >>> # Shares memory with array 'array'
@@ -6743,7 +6749,7 @@ def merge_dicts(*dicts):
 add_docstr(
     torch.amax,
     r"""
-amax(input, dim, keepdim=False, *, out=None) -> Tensor
+amax(input, dim=None, keepdim=False, *, out=None) -> Tensor
 
 Returns the maximum value of each slice of the :attr:`input` tensor in the given
 dimension(s) :attr:`dim`.
@@ -7367,7 +7373,7 @@ def merge_dicts(*dicts):
 add_docstr(
     torch.amin,
     r"""
-amin(input, dim, keepdim=False, *, out=None) -> Tensor
+amin(input, dim=None, keepdim=False, *, out=None) -> Tensor
 
 Returns the minimum value of each slice of the :attr:`input` tensor in the given
 dimension(s) :attr:`dim`.
@@ -11344,7 +11350,7 @@ def merge_dicts(*dicts):
 Args:
     {input}
     indices (LongTensor): the indices into :attr:`input`. Must have long dtype.
-    dim (int, optional): dimension to select along. Default: 0
+    dim (int, optional): dimension to select along. Default: `None`.
 
 Keyword args:
     {out}
@@ -13634,6 +13640,24 @@ def merge_dicts(*dicts):
 """,
 )
 
+add_docstr(
+    torch.Stream.is_capturing,
+    r"""
+Stream.is_capturing() -> bool
+
+Return true if this stream is currently recording work for graph capture.
+
+Returns:
+    bool: A boolean indicating if the stream is capturing.
+
+Example::
+
+    >>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_CUDA)
+    >>> s_cuda = torch.Stream(device='cuda')
+    >>> s_cuda.is_capturing()
+""",
+)
+
 
 add_docstr(
     torch.Stream.wait_event,
diff --git a/torch/_utils.py b/torch/_utils.py
index e2aa2778f0651..0c823d52f2f1f 100644
--- a/torch/_utils.py
+++ b/torch/_utils.py
@@ -703,17 +703,6 @@ def _take_tensors(tensors, size_limit):
             yield buf
 
 
-# annotation decorator to get annotations in a way that is compatible
-# with both Python 2 and 3
-def annotate(ret, **kwargs):
-    def dec(fun):
-        fun.__annotations__ = dict(kwargs)
-        fun.__annotations__["return"] = ret
-        return fun
-
-    return dec
-
-
 def render_call(fn, args, kwargs):
     str_fn = torch.overrides.resolve_name(fn)
     if str_fn is None:
@@ -785,6 +774,19 @@ def reraise(self):
         raise exception
 
 
+def cpu_count() -> int | None:
+    """Return the number of CPUs available to the current process.
+
+    Prefers ``os.sched_getaffinity`` (respects cgroups / taskset) and
+    falls back to ``os.cpu_count``.
+    """
+    # os.process_cpu_count was added in CPython 3.13, see
+    # https://docs.python.org/3/library/os.html#os.process_cpu_count
+    if hasattr(os, "sched_getaffinity"):
+        return len(os.sched_getaffinity(0))
+    return os.cpu_count()
+
+
 def _get_available_device_type():
     if torch.cuda.is_available():
         return "cuda"
@@ -1128,6 +1130,32 @@ def try_import(module_name: str) -> ModuleType | None:
 }
 
 
+def _chunk_or_narrow_cat(
+    tensor: "torch.Tensor",
+    num_chunks: int,
+    narrow_dim: int,
+    cat_dim: int = 0,
+) -> "torch.Tensor":
+    """
+    Splits tensor along narrow_dim into num_chunks and concatenates along cat_dim.
+    Uses torch.chunk in eager mode, but torch.narrow under tracing to be unbacked-symint safe.
+    """
+    if torch.distributed.is_available():
+        from torch.distributed._functional_collectives import _are_we_tracing
+        from torch.fx.experimental.symbolic_shapes import has_free_unbacked_symbols
+
+        # TODO(pianpwk): remove the unbacked symbols check and fix AsyncTP pattern matching
+        # for test_micro_pipeline_tp.py.
+        if _are_we_tracing() and has_free_unbacked_symbols(tensor):
+            chunk_size = tensor.size(narrow_dim) // num_chunks
+            chunks = [
+                torch.narrow(tensor, narrow_dim, i * chunk_size, chunk_size)
+                for i in range(num_chunks)
+            ]
+            return torch.cat(chunks, dim=cat_dim)
+    return torch.cat(torch.chunk(tensor, num_chunks, dim=narrow_dim), dim=cat_dim)
+
+
 def _maybe_view_chunk_cat(
     res: "torch.Tensor", group_size: int, gather_dim: int
 ) -> "torch.Tensor":
@@ -1176,9 +1204,11 @@ def _maybe_view_chunk_cat(
     # Optimization: Can use view instead of split+cat when:
     # 1. res.shape[0] == group_size (invariant after all_gather)
     # 2. All dims between 0 and gather_dim (exclusive) have size 1
+    from torch.fx.experimental.symbolic_shapes import guard_or_false
+
     numel_between = math.prod(shape[1:gather_dim]) if gather_dim > 1 else 1
 
-    if shape[0] == group_size and numel_between == 1:
+    if guard_or_false(shape[0] == group_size) and guard_or_false(numel_between == 1):
         # View optimization: reshape to collapse dim 0 into gather_dim
         final_shape = (
             [1]  # Dim 0 becomes 1
@@ -1191,7 +1221,7 @@ def _maybe_view_chunk_cat(
         # General case: fall back to split + cat
         # This is better than torch.flatten as cat can be vectorized, whereas
         # the contiguous kernel is always bad.
-        return torch.cat(torch.chunk(res, group_size, dim=0), dim=gather_dim)
+        return _chunk_or_narrow_cat(res, group_size, narrow_dim=0, cat_dim=gather_dim)
 
 
 class _Frame(TypedDict):
@@ -1351,3 +1381,17 @@ def _augment_memory_snapshot_stack_traces(
                 _augment_frames(trace_entry["frames"])
 
     return snapshot_dict
+
+
+def _is_privateuse1_backend_available():
+    """
+    Determines whether the privateuse1 backend is registered and available.
+
+    Returns:
+        Return True if the privateuse1 backend is registered and available.
+    """
+    privateuse1_backend_name = torch._C._get_privateuse1_backend_name()
+    privateuse1_backend_module = getattr(torch, privateuse1_backend_name, None)
+    return (
+        is_available := getattr(privateuse1_backend_module, "is_available", None)
+    ) and is_available()
diff --git a/torch/_utils_internal.py b/torch/_utils_internal.py
index ec7c0a32f52b5..75ed0a754caed 100644
--- a/torch/_utils_internal.py
+++ b/torch/_utils_internal.py
@@ -375,3 +375,10 @@ def find_compile_subproc_binary() -> str | None:
     Allows overriding the binary used for subprocesses
     """
     return None
+
+
+def get_torch_source_version() -> str:
+    """Return the source commit hash for the current PyTorch build."""
+    import torch.version as torch_version
+
+    return getattr(torch_version, "git_version", "")
diff --git a/torch/accelerator/__init__.py b/torch/accelerator/__init__.py
index 4b21450006e2e..0e7528e71a527 100644
--- a/torch/accelerator/__init__.py
+++ b/torch/accelerator/__init__.py
@@ -9,8 +9,10 @@
 import torch
 
 from ._utils import _device_t, _get_device_index
+from .graphs import Graph
 from .memory import (
     empty_cache,
+    empty_host_cache,
     get_memory_info,
     max_memory_allocated,
     max_memory_reserved,
@@ -23,6 +25,7 @@
 
 
 __all__ = [
+    "Graph",
     "current_accelerator",
     "current_device_idx",  # deprecated
     "current_device_index",
@@ -31,6 +34,7 @@
     "device_count",
     "device_index",
     "empty_cache",
+    "empty_host_cache",
     "get_memory_info",
     "is_available",
     "max_memory_allocated",
@@ -165,7 +169,10 @@ def get_device_capability(device: _device_t = None, /) -> dict[str, Any]:
 
     Returns:
         dict[str, Any]: A dictionary containing device capability information. The dictionary includes:
-            - ``supported_dtypes`` (set(torch.dtype)): Set of PyTorch data types supported by the device
+            - ``supported_dtypes`` (set(torch.dtype)): Set of PyTorch data types for which
+              tensors can be allocated on the accelerator and type conversion across
+              supported dtypes are supported. Any operator support outside of that
+              is not guaranteed
 
     Examples:
         >>> # xdoctest: +SKIP("requires cuda")
diff --git a/torch/accelerator/graphs.py b/torch/accelerator/graphs.py
new file mode 100644
index 0000000000000..f02e91f253a45
--- /dev/null
+++ b/torch/accelerator/graphs.py
@@ -0,0 +1,177 @@
+import gc
+from typing import Literal
+from typing_extensions import Self
+
+import torch
+from torch._C import _acceleratorGraph
+
+
+class Graph(_acceleratorGraph):
+    r"""
+    Wrapper around an :ref:`accelerator<accelerators>` graph that supports capture and replay.
+
+    A graph captures a sequence of operations and their dependencies, allowing them to be
+    replayed efficiently with reduced overhead. This class can be used as a context manager
+    to automatically capture operations on the current stream.
+
+    Arguments:
+        keep_graph (bool, optional): If ``False``, the underlying graph is destroyed and the
+            executable graph is instantiated on the GPU at the end of ``capture_end``.
+            If ``True``, the underlying graph is preserved after ``capture_end``. In this case,
+            the executable graph is not instantiated automatically; it must be explicitly created
+            by calling ``instantiate``, or it will be instantiated on the first call to ``replay``.
+            Defaults to ``False``.
+        pool (tuple[int, int], optional): Memory pool identifier for this graph. Multiple graphs
+            can share the same pool by passing the same identifier, which can reduce memory overhead.
+            Defaults to ``None``.
+        capture_error_mode (Literal["default", "global", "thread_local", "relaxed"], optional):
+            Specifies the behavior of graph capture. The exact semantics are backend-specific.
+            ``"default"``: backend-defined default capture behavior.
+            ``"global"``: potentially unsafe API calls are prohibited. Errors may occur if capture
+            in the current thread affects other threads.
+            ``"thread_local"``: potentially unsafe API calls are prohibited. Errors occur only if
+            capture in the current thread affects itself.
+            ``"relaxed"``: the current thread is allowed to make potentially unsafe API calls, except
+            for calls that inherently conflict with stream capture.
+            Default: ``"default"``.
+
+    Example::
+
+        >>> # xdoctest: +SKIP
+        >>> x = torch.zeros([2000], device=0)
+
+        >>> stream = torch.Stream()
+        >>> graph = torch.accelerator.Graph()
+        >>> with stream, graph:
+        ...     x += 1
+
+        >>> graph.replay()
+    """
+
+    def __new__(
+        cls,
+        keep_graph: bool = False,
+        *,
+        pool: tuple[int, int] | None = None,
+        capture_error_mode: Literal[
+            "default", "global", "thread_local", "relaxed"
+        ] = "default",
+    ) -> Self:
+        return super().__new__(cls, keep_graph)
+
+    def __init__(
+        self,
+        keep_graph: bool = False,
+        *,
+        pool: tuple[int, int] | None = None,
+        capture_error_mode: Literal[
+            "default", "global", "thread_local", "relaxed"
+        ] = "default",
+    ) -> None:
+        super().__init__(keep_graph)
+        self.graph_pool = pool
+        self.capture_error_mode = capture_error_mode
+
+    # pyrefly: ignore [bad-override]
+    def capture_begin(self) -> None:
+        r"""
+        Begin graph capture on the current stream.
+
+        All operations on the current stream after this call will be recorded into the graph until
+        ``capture_end`` is called, using the memory pool and capture error mode provided at construction time.
+        """
+        super().capture_begin(
+            pool=self.graph_pool, capture_error_mode=self.capture_error_mode
+        )
+
+    def capture_end(self) -> None:
+        r"""
+        End graph capture on the current stream of the current device.
+
+        After this call, the graph can be replayed via ``replay``.
+        """
+        super().capture_end()
+
+    def instantiate(self) -> None:
+        r"""
+        Instantiate the underlying graph. Will be called by ``capture_end``
+        if ``keep_graph=False``, or by ``replay`` if ``keep_graph=True`` and
+        ``instantiate`` has not already been explicitly called.
+        """
+        super().instantiate()
+
+    def replay(self) -> None:
+        r"""Replay the work captured by this graph."""
+        super().replay()
+
+    def reset(self) -> None:
+        r"""
+        Delete the graph currently held by this instance.
+
+        After this call, the graph can be recaptured. Set :attr:`graph_pool` or
+        :attr:`capture_error_mode` beforehand to use different settings on the next capture.
+        """
+        super().reset()
+
+    def pool(self) -> tuple[int, int]:
+        r"""
+        Return an opaque token representing the id of this graph's memory pool.
+
+        This id can optionally be passed to another graph's ``capture_begin``,
+        which hints the other graph may share the same memory pool.
+
+        Example::
+            >>> # xdoctest: +SKIP
+            >>> g1 = torch.accelerator.Graph()
+            >>> g1.capture_begin()
+            >>> # ... operations ...
+            >>> g1.capture_end()
+
+            >>> # Share g1's memory pool with a new graph
+            >>> pool_id = g1.pool()
+            >>> g2 = torch.accelerator.Graph(pool=pool_id)
+        """
+        return super().pool()
+
+    def enable_debug_mode(self) -> None:
+        r"""Enable debugging mode for ``debug_dump``."""
+        return super().enable_debug_mode()
+
+    def debug_dump(self, path: str) -> None:
+        r"""
+        Dump the captured graph to a file for debugging purposes if the debugging is
+        enabled via ``enable_debug_mode``.
+
+        Arguments:
+            path (str): Path to dump the graph to.
+
+        Example::
+            >>> # xdoctest: +SKIP
+            >>> s = torch.Stream()
+            >>> g = torch.accelerator.Graph()
+            >>> g.enable_debug_mode()
+
+            >>> with s, g:
+            >>> # ... operations ...
+
+            >>> # Dump captured graph to a file "graph_dump.dot"
+            >>> g.debug_dump("graph_dump.dot")
+        """
+        return super().debug_dump(path)
+
+    def __enter__(self) -> None:
+        torch.accelerator.synchronize()
+        if torch.compiler.config.force_cudagraph_gc:
+            # We previously always ran garbage collection here. While this can help
+            # reclaim accelerator device memory held by dead Python cycles, it is
+            # very expensive, especially when performing multiple graph captures in sequence.
+            gc.collect()
+        torch.accelerator.empty_cache()
+        torch.accelerator.empty_host_cache()
+        self.capture_begin()
+
+    def __exit__(self, *exc_info: object) -> None:
+        self.capture_end()
+
+
+__all__ = ["Graph"]
diff --git a/torch/accelerator/memory.py b/torch/accelerator/memory.py
index a326af6dbcc1b..08998d221a58b 100644
--- a/torch/accelerator/memory.py
+++ b/torch/accelerator/memory.py
@@ -8,6 +8,7 @@
 
 __all__ = [
     "empty_cache",
+    "empty_host_cache",
     "get_memory_info",
     "max_memory_allocated",
     "max_memory_reserved",
@@ -31,6 +32,16 @@ def empty_cache() -> None:
     torch._C._accelerator_emptyCache()
 
 
+def empty_host_cache() -> None:
+    r"""Release all unoccupied cached host (pinned) memory currently held by the host caching
+    allocator so that it can be used by other applications.
+
+    .. note:: This function is a no-op if the memory allocator for the current
+        :ref:`accelerator <accelerators>` has not been initialized.
+    """
+    torch._C._accelerator_emptyHostCache()
+
+
 def memory_stats(device_index: _device_t = None, /) -> OrderedDict[str, Any]:
     r"""Return a dictionary of accelerator device memory allocator statistics for a given device index.
 
diff --git a/torch/amp/autocast_mode.py b/torch/amp/autocast_mode.py
index 3cb2d6eb5dc6b..892b9eeef646c 100644
--- a/torch/amp/autocast_mode.py
+++ b/torch/amp/autocast_mode.py
@@ -2,7 +2,7 @@
 import collections
 import functools
 import warnings
-from typing import Any, Optional
+from typing import Any
 
 import torch
 from torch.types import _dtype
@@ -222,9 +222,9 @@ def forward(self, x):
     def __init__(
         self,
         device_type: str,
-        dtype: Optional[_dtype] = None,
+        dtype: _dtype | None = None,
         enabled: bool = True,
-        cache_enabled: Optional[bool] = None,
+        cache_enabled: bool | None = None,
     ):
         if not isinstance(device_type, str):
             raise ValueError(
@@ -382,9 +382,21 @@ def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any):  # type: ignore[ov
     def __call__(self, func):
         if torch._jit_internal.is_scripting():
             return func
+        if not callable(func):
+            raise TypeError(
+                f"autocast()(func) requires a callable, but got {type(func).__name__}. "
+                f"Did you mean to use autocast as a context manager? For example:\n"
+                f"    with torch.autocast(device_type=...):\n"
+                f"        output = model(input)"
+            )
         return autocast_decorator(self, func)
 
 
+# Subclass to distinguish autocast variables created by _enter_autocast (and not managed by a with statement)
+class _UnmanagedAutocast(autocast):
+    pass
+
+
 # These functions aren't meant for public usage.
 # They are what we trace into a graph during pre_dispatch tracing
 # when we encounter an autocast context manager.
@@ -394,7 +406,7 @@ def _enter_autocast(*vals):
         return torch.overrides.handle_torch_function(
             torch.amp._enter_autocast, [], *vals
         )
-    mode = torch.amp.autocast(*vals)
+    mode = _UnmanagedAutocast(*vals)
     mode.__enter__()
     return mode
 
@@ -442,7 +454,7 @@ def custom_fwd(
     fwd=None,
     *,
     device_type: str,
-    cast_inputs: Optional[_dtype] = None,
+    cast_inputs: _dtype | None = None,
 ):
     """
     Create a helper decorator for ``forward`` methods of custom autograd functions.
diff --git a/torch/amp/grad_scaler.py b/torch/amp/grad_scaler.py
index 25900f88eb1b3..b17d84e72ae87 100644
--- a/torch/amp/grad_scaler.py
+++ b/torch/amp/grad_scaler.py
@@ -5,7 +5,7 @@
 import warnings
 from collections import abc, defaultdict
 from enum import Enum
-from typing import Any, cast, Optional, overload, TYPE_CHECKING, Union
+from typing import Any, cast, overload, TYPE_CHECKING
 
 import torch
 
@@ -147,13 +147,13 @@ def __init__(
 
             self._init_scale = init_scale
             # self._scale will be lazily initialized during the first call to scale()
-            self._scale: Optional[torch.Tensor] = None
+            self._scale: torch.Tensor | None = None
             self._growth_factor = growth_factor
             self._backoff_factor = backoff_factor
             self._growth_interval = growth_interval
             self._init_growth_tracker = 0
             # self._growth_tracker will be lazily initialized during the first call to scale()
-            self._growth_tracker: Optional[torch.Tensor] = None
+            self._growth_tracker: torch.Tensor | None = None
             self._per_optimizer_states: dict[int, dict[str, Any]] = defaultdict(
                 _refresh_per_optimizer_state
             )
@@ -192,8 +192,8 @@ def scale(self, outputs: Iterable[torch.Tensor]) -> Iterable[torch.Tensor]: ...
 
     def scale(
         self,
-        outputs: Union[torch.Tensor, Iterable[torch.Tensor]],
-    ) -> Union[torch.Tensor, Iterable[torch.Tensor]]:
+        outputs: torch.Tensor | Iterable[torch.Tensor],
+    ) -> torch.Tensor | Iterable[torch.Tensor]:
         """
         Multiplies ('scales') a tensor or list of tensors by the scale factor.
 
@@ -221,7 +221,7 @@ def scale(
             _MultiDeviceReplicator
         ] = []  # holds a reference that can be overwritten by apply_scale
 
-        def apply_scale(val: Union[torch.Tensor, Iterable[torch.Tensor]]):
+        def apply_scale(val: torch.Tensor | Iterable[torch.Tensor]):
             if isinstance(val, torch.Tensor):
                 if len(stash) == 0:
                     if self._scale is None:
@@ -366,15 +366,15 @@ def _maybe_opt_step(
         optimizer_state: dict[str, Any],
         *args: Any,
         **kwargs: Any,
-    ) -> Optional[float]:
-        retval: Optional[float] = None
+    ) -> float | None:
+        retval: float | None = None
         if not sum(v.item() for v in optimizer_state["found_inf_per_device"].values()):
             retval = optimizer.step(*args, **kwargs)
         return retval
 
     def step(
         self, optimizer: torch.optim.Optimizer, *args: Any, **kwargs: Any
-    ) -> Optional[float]:
+    ) -> float | None:
         """Invoke ``unscale_(optimizer)`` followed by parameter update, if gradients are not infs/NaN.
 
         :meth:`step` carries out the following two operations:
@@ -413,7 +413,7 @@ def step(
                 "step() has already been called since the last update()."
             )
 
-        retval: Optional[float] = None
+        retval: float | None = None
 
         if getattr(optimizer, "_step_supports_amp_scaling", False):
             # This optimizer has customized scale-handling logic, so we can call optimizer.step() directly.
@@ -481,7 +481,7 @@ def step(
 
         return retval
 
-    def update(self, new_scale: Optional[Union[float, torch.Tensor]] = None) -> None:
+    def update(self, new_scale: float | torch.Tensor | None = None) -> None:
         """Update the scale factor.
 
         If any optimizer steps were skipped the scale is multiplied by ``backoff_factor``
@@ -558,7 +558,7 @@ def update(self, new_scale: Optional[Union[float, torch.Tensor]] = None) -> None
         # To prepare for next iteration, clear the data collected from optimizers this iteration.
         self._per_optimizer_states = defaultdict(_refresh_per_optimizer_state)
 
-    def _get_scale_async(self) -> Optional[torch.Tensor]:
+    def _get_scale_async(self) -> torch.Tensor | None:
         return self._scale
 
     def get_scale(self) -> float:
diff --git a/torch/ao/nn/qat/modules/embedding_ops.py b/torch/ao/nn/qat/modules/embedding_ops.py
index 7db96312bc684..532ef186bf30b 100644
--- a/torch/ao/nn/qat/modules/embedding_ops.py
+++ b/torch/ao/nn/qat/modules/embedding_ops.py
@@ -50,7 +50,6 @@ def __init__(
             scale_grad_by_freq,
             sparse,
             _weight,
-            # pyrefly: ignore [bad-argument-type]
             **factory_kwargs,
         )
         if not qconfig:
diff --git a/torch/ao/nn/quantized/dynamic/modules/rnn.py b/torch/ao/nn/quantized/dynamic/modules/rnn.py
index 36e9c08b4f7f9..b1e186ef130fb 100644
--- a/torch/ao/nn/quantized/dynamic/modules/rnn.py
+++ b/torch/ao/nn/quantized/dynamic/modules/rnn.py
@@ -525,7 +525,6 @@ class LSTM(RNNBase):
         >>> output, (hn, cn) = rnn(input, (h0, c0))
     """
 
-    # pyrefly: ignore [bad-override]
     _FLOAT_MODULE = nn.LSTM
 
     __overloads__ = {"forward": ["forward_packed", "forward_tensor"]}
@@ -814,7 +813,6 @@ class GRU(RNNBase):
         >>> output, hn = rnn(input, h0)
     """
 
-    # pyrefly: ignore [bad-override]
     _FLOAT_MODULE = nn.GRU
 
     __overloads__ = {"forward": ["forward_packed", "forward_tensor"]}
diff --git a/torch/ao/nn/quantized/modules/activation.py b/torch/ao/nn/quantized/modules/activation.py
index 3ecf1d5c9a1e2..669a9978c27ba 100644
--- a/torch/ao/nn/quantized/modules/activation.py
+++ b/torch/ao/nn/quantized/modules/activation.py
@@ -67,9 +67,7 @@ class Hardswish(torch.nn.Hardswish):
     def __init__(self, scale, zero_point, device=None, dtype=None):
         factory_kwargs = {"device": device, "dtype": dtype}
         super().__init__()
-        # pyrefly: ignore [bad-argument-type]
         self.register_buffer("scale", torch.tensor(scale, **factory_kwargs))
-        # pyrefly: ignore [bad-argument-type]
         self.register_buffer("zero_point", torch.tensor(zero_point, **factory_kwargs))
 
     def forward(self, input):
@@ -140,9 +138,7 @@ def __init__(
     ) -> None:
         factory_kwargs = {"device": device, "dtype": dtype}
         super().__init__(negative_slope, inplace)
-        # pyrefly: ignore [bad-argument-type]
         self.register_buffer("scale", torch.tensor(scale, **factory_kwargs))
-        # pyrefly: ignore [bad-argument-type]
         self.register_buffer("zero_point", torch.tensor(zero_point, **factory_kwargs))
 
     def forward(self, input):
@@ -230,7 +226,6 @@ def from_reference(cls, mod, scale, zero_point):
 
 
 class MultiheadAttention(torch.ao.nn.quantizable.MultiheadAttention):
-    # pyrefly: ignore [bad-override]
     _FLOAT_MODULE = torch.ao.nn.quantizable.MultiheadAttention
 
     def _get_name(self):
diff --git a/torch/ao/nn/quantized/modules/batchnorm.py b/torch/ao/nn/quantized/modules/batchnorm.py
index f1e6779c08b1f..ad648743084e0 100644
--- a/torch/ao/nn/quantized/modules/batchnorm.py
+++ b/torch/ao/nn/quantized/modules/batchnorm.py
@@ -12,9 +12,7 @@ def __init__(
     ) -> None:
         factory_kwargs = {"device": device, "dtype": dtype}
         super().__init__(num_features, eps, momentum, True, True, **factory_kwargs)
-        # pyrefly: ignore [bad-argument-type]
         self.register_buffer("scale", torch.tensor(1.0, **factory_kwargs))
-        # pyrefly: ignore [bad-argument-type]
         self.register_buffer("zero_point", torch.tensor(0, **factory_kwargs))
 
     @staticmethod
diff --git a/torch/ao/nn/quantized/modules/conv.py b/torch/ao/nn/quantized/modules/conv.py
index fa134438d1759..83e0b95e421fa 100644
--- a/torch/ao/nn/quantized/modules/conv.py
+++ b/torch/ao/nn/quantized/modules/conv.py
@@ -1,7 +1,7 @@
 # mypy: allow-untyped-defs
 r"""Quantized convolution modules."""
 
-from typing import ClassVar, Literal, Optional
+from typing import ClassVar, Literal
 
 import torch
 import torch.ao.nn.intrinsic as nni
@@ -384,10 +384,10 @@ class Conv1d(_ConvNd):
     """
 
     _FLOAT_MODULE: ClassVar[type[nn.Conv1d]] = nn.Conv1d
-    _NNIQAT_CONV_BN_MODULE: ClassVar[Optional[type[nn.Module]]] = nniqat.ConvBn1d
-    _NNI_CONV_RELU_MODULE: ClassVar[Optional[type[nn.Module]]] = nni.ConvReLU1d
-    _NNI_CONV_ADD_MODULE: ClassVar[Optional[type[nn.Module]]] = None
-    _NNI_CONV_ADD_RELU_MODULE: ClassVar[Optional[type[nn.Module]]] = None
+    _NNIQAT_CONV_BN_MODULE: ClassVar[type[nn.Module] | None] = nniqat.ConvBn1d
+    _NNI_CONV_RELU_MODULE: ClassVar[type[nn.Module] | None] = nni.ConvReLU1d
+    _NNI_CONV_ADD_MODULE: ClassVar[type[nn.Module] | None] = None
+    _NNI_CONV_ADD_RELU_MODULE: ClassVar[type[nn.Module] | None] = None
 
     def __init__(
         self,
@@ -430,7 +430,7 @@ def __init__(
     def _get_name(self):
         return "QuantizedConv1d"
 
-    def set_weight_bias(self, w: torch.Tensor, b: Optional[torch.Tensor]) -> None:
+    def set_weight_bias(self, w: torch.Tensor, b: torch.Tensor | None) -> None:
         if self.padding_mode == "zeros":
             self._packed_params = torch.ops.quantized.conv1d_prepack(
                 w, b, self.stride, self.padding, self.dilation, self.groups
@@ -518,8 +518,8 @@ class Conv2d(_ConvNd):
     """
 
     _FLOAT_MODULE: ClassVar[type[nn.Conv2d]] = nn.Conv2d
-    _NNIQAT_CONV_BN_MODULE: ClassVar[Optional[type[nn.Module]]] = nniqat.ConvBn2d
-    _NNI_CONV_RELU_MODULE: ClassVar[Optional[type[nn.Module]]] = nni.ConvReLU2d
+    _NNIQAT_CONV_BN_MODULE: ClassVar[type[nn.Module] | None] = nniqat.ConvBn2d
+    _NNI_CONV_RELU_MODULE: ClassVar[type[nn.Module] | None] = nni.ConvReLU2d
     _NNI_CONV_ADD_MODULE: ClassVar[type[nni.ConvAdd2d]] = nni.ConvAdd2d
     _NNI_CONV_ADD_RELU_MODULE: ClassVar[type[nni.ConvAddReLU2d]] = nni.ConvAddReLU2d
 
@@ -562,7 +562,7 @@ def __init__(
     def _get_name(self):
         return "QuantizedConv2d"
 
-    def set_weight_bias(self, w: torch.Tensor, b: Optional[torch.Tensor]) -> None:
+    def set_weight_bias(self, w: torch.Tensor, b: torch.Tensor | None) -> None:
         if self.padding_mode == "zeros":
             self._packed_params = torch.ops.quantized.conv2d_prepack(
                 w, b, self.stride, self.padding, self.dilation, self.groups
@@ -648,10 +648,10 @@ class Conv3d(_ConvNd):
     """
 
     _FLOAT_MODULE: ClassVar[type[nn.Conv3d]] = nn.Conv3d
-    _NNIQAT_CONV_BN_MODULE: ClassVar[Optional[type[nn.Module]]] = nniqat.ConvBn3d
-    _NNI_CONV_RELU_MODULE: ClassVar[Optional[type[nn.Module]]] = nni.ConvReLU3d
-    _NNI_CONV_ADD_MODULE: ClassVar[Optional[type[nn.Module]]] = None
-    _NNI_CONV_ADD_RELU_MODULE: ClassVar[Optional[type[nn.Module]]] = None
+    _NNIQAT_CONV_BN_MODULE: ClassVar[type[nn.Module] | None] = nniqat.ConvBn3d
+    _NNI_CONV_RELU_MODULE: ClassVar[type[nn.Module] | None] = nni.ConvReLU3d
+    _NNI_CONV_ADD_MODULE: ClassVar[type[nn.Module] | None] = None
+    _NNI_CONV_ADD_RELU_MODULE: ClassVar[type[nn.Module] | None] = None
 
     def __init__(
         self,
@@ -694,7 +694,7 @@ def __init__(
     def _get_name(self):
         return "QuantizedConv3d"
 
-    def set_weight_bias(self, w: torch.Tensor, b: Optional[torch.Tensor]) -> None:
+    def set_weight_bias(self, w: torch.Tensor, b: torch.Tensor | None) -> None:
         if self.padding_mode == "zeros":
             self._packed_params = torch.ops.quantized.conv3d_prepack(
                 w, b, self.stride, self.padding, self.dilation, self.groups
@@ -960,7 +960,7 @@ def __init__(
     def _get_name(self):
         return "QuantizedConvTranspose1d"
 
-    def set_weight_bias(self, w: torch.Tensor, b: Optional[torch.Tensor]) -> None:
+    def set_weight_bias(self, w: torch.Tensor, b: torch.Tensor | None) -> None:
         self._packed_params = torch.ops.quantized.conv_transpose1d_prepack(
             w,
             b,
@@ -1083,7 +1083,7 @@ def __init__(
     def _get_name(self):
         return "QuantizedConvTranspose2d"
 
-    def set_weight_bias(self, w: torch.Tensor, b: Optional[torch.Tensor]) -> None:
+    def set_weight_bias(self, w: torch.Tensor, b: torch.Tensor | None) -> None:
         self._packed_params = torch.ops.quantized.conv_transpose2d_prepack(
             w,
             b,
@@ -1208,7 +1208,7 @@ def __init__(
     def _get_name(self):
         return "QuantizedConvTranspose3d"
 
-    def set_weight_bias(self, w: torch.Tensor, b: Optional[torch.Tensor]) -> None:
+    def set_weight_bias(self, w: torch.Tensor, b: torch.Tensor | None) -> None:
         self._packed_params = torch.ops.quantized.conv_transpose3d_prepack(
             w,
             b,
diff --git a/torch/ao/nn/quantized/modules/normalization.py b/torch/ao/nn/quantized/modules/normalization.py
index fa335b4699db5..4db2ac6e928f4 100644
--- a/torch/ao/nn/quantized/modules/normalization.py
+++ b/torch/ao/nn/quantized/modules/normalization.py
@@ -37,14 +37,11 @@ def __init__(
             normalized_shape,
             eps=eps,
             elementwise_affine=elementwise_affine,
-            # pyrefly: ignore [bad-argument-type]
             **factory_kwargs,
         )
         self.weight = weight
         self.bias = bias
-        # pyrefly: ignore [bad-argument-type]
         self.register_buffer("scale", torch.tensor(scale, **factory_kwargs))
-        # pyrefly: ignore [bad-argument-type]
         self.register_buffer("zero_point", torch.tensor(zero_point, **factory_kwargs))
 
     def forward(self, input):
@@ -116,9 +113,7 @@ def __init__(
         super().__init__(num_groups, num_channels, eps, affine, **factory_kwargs)
         self.weight = weight
         self.bias = bias
-        # pyrefly: ignore [bad-argument-type]
         self.register_buffer("scale", torch.tensor(scale, **factory_kwargs))
-        # pyrefly: ignore [bad-argument-type]
         self.register_buffer("zero_point", torch.tensor(zero_point, **factory_kwargs))
 
     def forward(self, input):
@@ -180,9 +175,7 @@ def __init__(
         )
         self.weight = weight
         self.bias = bias
-        # pyrefly: ignore [bad-argument-type]
         self.register_buffer("scale", torch.tensor(scale, **factory_kwargs))
-        # pyrefly: ignore [bad-argument-type]
         self.register_buffer("zero_point", torch.tensor(zero_point, **factory_kwargs))
 
     def forward(self, input):
@@ -249,9 +242,7 @@ def __init__(
         )
         self.weight = weight
         self.bias = bias
-        # pyrefly: ignore [bad-argument-type]
         self.register_buffer("scale", torch.tensor(scale, **factory_kwargs))
-        # pyrefly: ignore [bad-argument-type]
         self.register_buffer("zero_point", torch.tensor(zero_point, **factory_kwargs))
 
     def forward(self, input):
@@ -318,9 +309,7 @@ def __init__(
         )
         self.weight = weight
         self.bias = bias
-        # pyrefly: ignore [bad-argument-type]
         self.register_buffer("scale", torch.tensor(scale, **factory_kwargs))
-        # pyrefly: ignore [bad-argument-type]
         self.register_buffer("zero_point", torch.tensor(zero_point, **factory_kwargs))
 
     def forward(self, input):
diff --git a/torch/ao/nn/quantized/reference/modules/conv.py b/torch/ao/nn/quantized/reference/modules/conv.py
index d3bddd6cfe8a9..3e69270211eed 100644
--- a/torch/ao/nn/quantized/reference/modules/conv.py
+++ b/torch/ao/nn/quantized/reference/modules/conv.py
@@ -1,5 +1,5 @@
 # mypy: allow-untyped-defs
-from typing import Any, Literal, Optional
+from typing import Any, Literal
 
 import torch
 import torch.nn as nn
@@ -26,7 +26,7 @@ class _ConvNd(torch.nn.modules.conv._ConvNd, ReferenceQuantizedModule):
     this is useful when user want to use this module in other backends like Glow.
     """
 
-    __annotations__ = {"bias": Optional[torch.Tensor]}
+    __annotations__ = {"bias": torch.Tensor | None}
     _IS_REFERENCE = True
 
     @staticmethod
diff --git a/torch/ao/nn/quantized/reference/modules/rnn.py b/torch/ao/nn/quantized/reference/modules/rnn.py
index 0df79ccd4fa66..fd5179e7f3ad0 100644
--- a/torch/ao/nn/quantized/reference/modules/rnn.py
+++ b/torch/ao/nn/quantized/reference/modules/rnn.py
@@ -663,7 +663,6 @@ def forward(self, input, hx=None):  # noqa: F811
         if isinstance(orig_input, PackedSequence):
             output_packed = PackedSequence(
                 output,
-                # pyrefly: ignore [bad-argument-type]
                 batch_sizes,
                 sorted_indices,
                 unsorted_indices,
@@ -825,7 +824,6 @@ def forward(self, input, hx=None):  # noqa: F811
         if isinstance(orig_input, PackedSequence):
             output_packed = PackedSequence(
                 output,
-                # pyrefly: ignore [bad-argument-type]
                 batch_sizes,
                 sorted_indices,
                 unsorted_indices,
diff --git a/torch/ao/ns/_numeric_suite_fx.py b/torch/ao/ns/_numeric_suite_fx.py
index 7ab8db5576f5e..1861d0160db15 100644
--- a/torch/ao/ns/_numeric_suite_fx.py
+++ b/torch/ao/ns/_numeric_suite_fx.py
@@ -228,7 +228,6 @@ def forward(self, x):
         if isinstance(x, torch.Tensor):
             self.stats.append(x.detach())
         elif isinstance(x, tuple) and len(x) == 2 and len(x[1]) == 2:
-            # pyrefly: ignore [bad-index]
             new_res = (x[0].detach(), (x[1][0].detach(), x[1][1].detach()))
             self.stats_rnn.append(new_res)
         return x
diff --git a/torch/ao/ns/fx/graph_matcher.py b/torch/ao/ns/fx/graph_matcher.py
index 0e85cf9a91ba2..e8dbd98e31375 100644
--- a/torch/ao/ns/fx/graph_matcher.py
+++ b/torch/ao/ns/fx/graph_matcher.py
@@ -83,7 +83,7 @@ def __next__(self) -> NSSubgraph:
                 )
                 if is_match:
                     # navigate to the base node
-                    # pyrefly: ignore [bad-assignment]
+                    # pyrefly: ignore [bad-assignment, non-convergent-recursion]
                     for rev_fusion_idx in range(len(_reverse_fusion_ops) - 1):
                         # pyrefly: ignore [bad-argument-type]
                         self.seen_nodes.add(cur_start_node)
diff --git a/torch/ao/ns/fx/mappings.py b/torch/ao/ns/fx/mappings.py
index 6b27660f0e36d..275291789f1c5 100644
--- a/torch/ao/ns/fx/mappings.py
+++ b/torch/ao/ns/fx/mappings.py
@@ -376,19 +376,16 @@ def get_base_name_to_sets_of_related_ops() -> dict[str, set[NSNodeTargetType]]:
         if config.fused_module is not None:
             # case 1: pattern fuses a pattern of ops into an op
             # example: nn.Conv1d, nn.ReLU fused into nni.ConvReLU1d
-            # pyrefly: ignore [bad-argument-type]
             new_connections.append((first_element, config.fused_module))
 
         if config.qat_module is not None:
             # case 2: pattern swaps a module into a QAT module
             # example: nni.ConvReLU1d swapped into nniqat.ConvReLU1d
-            # pyrefly: ignore [bad-argument-type]
             new_connections.append((first_element, config.qat_module))
 
         if config.reference_quantized_module is not None:
             # case 3: reference version of floating point module, such as
             # nn.Conv2d and nnqr.Conv2d
-            # pyrefly: ignore [bad-argument-type]
             new_connections.append((first_element, config.reference_quantized_module))
 
     #
diff --git a/torch/ao/ns/fx/n_shadows_utils.py b/torch/ao/ns/fx/n_shadows_utils.py
index 8a730abc83869..41df14e962674 100644
--- a/torch/ao/ns/fx/n_shadows_utils.py
+++ b/torch/ao/ns/fx/n_shadows_utils.py
@@ -97,7 +97,6 @@ def fetch_attr(target: str):
                 # pyrefly: ignore [unbound-name]
                 node.traced_result = result
 
-            # pyrefly: ignore [unsupported-operation]
             # pyrefly: ignore [unbound-name]
             env[node.name] = result
 
@@ -404,10 +403,9 @@ def _add_placeholder(
                         cur_name_idx += 1
                         setattr(gm, mod_name, new_arg)
                         new_arg_placeholder = gm.placeholder(mod_name)  # type: ignore[operator]
-                        # pyrefly: ignore [missing-attribute]
                         cur_args_copy.append(new_arg_placeholder)
                     elif isinstance(arg, (float, int, torch.dtype)):
-                        # pyrefly: ignore [missing-attribute]
+                        # pyrefly: ignore [bad-argument-type]
                         cur_args_copy.append(arg)
                     else:
                         raise AssertionError(f"arg of type {type(arg)} not handled yet")
diff --git a/torch/ao/ns/fx/ns_types.py b/torch/ao/ns/fx/ns_types.py
index 134fd485130e0..9a25c0a9eefca 100644
--- a/torch/ao/ns/fx/ns_types.py
+++ b/torch/ao/ns/fx/ns_types.py
@@ -1,6 +1,6 @@
 import enum
 from collections.abc import Callable
-from typing import Any, NamedTuple, Union
+from typing import Any, NamedTuple
 
 from torch.fx.graph import Node
 
@@ -63,4 +63,4 @@ class NSSubgraph(NamedTuple):
 # `F.conv1d` for a `call_function` conv node
 # `nn.Conv1d` for a `call_module` node calling the forward of a `nn.Conv1d` module
 # `'sigmoid'` for a `call_method` node calling `x.sigmoid()`
-NSNodeTargetType = Union[Callable, str]
+NSNodeTargetType = Callable | str
diff --git a/torch/ao/ns/fx/pattern_utils.py b/torch/ao/ns/fx/pattern_utils.py
index 0ece07ee3743b..a885e0166c825 100644
--- a/torch/ao/ns/fx/pattern_utils.py
+++ b/torch/ao/ns/fx/pattern_utils.py
@@ -1,5 +1,5 @@
 from collections.abc import Callable
-from typing import Any, Union
+from typing import Any
 
 import torch
 import torch.nn as nn
@@ -36,17 +36,17 @@ def get_type_a_related_to_b(
     return type_a_related_to_b
 
 
-NSFusionElType = Union[
-    Callable,  # call_function or call_module type, example: F.linear or nn.Conv2d
-    str,  # call_method name, example: "dequantize"
-    tuple[
+NSFusionElType = (
+    Callable  # call_function or call_module type, example: F.linear or nn.Conv2d
+    | str  # call_method name, example: "dequantize"
+    | tuple[
         str, Any
-    ],  # call_method name and first argument, example: ("to", torch.float16)
-]
-NSFusionType = Union[
-    tuple[NSFusionElType, NSFusionElType],
-    tuple[NSFusionElType, NSFusionElType, NSFusionElType, NSFusionElType],
-]
+    ]  # call_method name and first argument, example: ("to", torch.float16)
+)
+NSFusionType = (
+    tuple[NSFusionElType, NSFusionElType]
+    | tuple[NSFusionElType, NSFusionElType, NSFusionElType, NSFusionElType]
+)
 
 
 def get_reversed_fusions() -> list[tuple[NSFusionType, int]]:
diff --git a/torch/ao/pruning/_experimental/data_sparsifier/benchmarks/dlrm_utils.py b/torch/ao/pruning/_experimental/data_sparsifier/benchmarks/dlrm_utils.py
index e2b31e0e563bf..c16f2342318e4 100644
--- a/torch/ao/pruning/_experimental/data_sparsifier/benchmarks/dlrm_utils.py
+++ b/torch/ao/pruning/_experimental/data_sparsifier/benchmarks/dlrm_utils.py
@@ -20,17 +20,12 @@ class SparseDLRM(DLRM_Net):
     """
 
     def forward(self, dense_x, lS_o, lS_i):
-        # pyrefly: ignore [missing-attribute]
         x = self.apply_mlp(dense_x, self.bot_l)  # dense features
-        # pyrefly: ignore [missing-attribute]
         ly = self.apply_emb(lS_o, lS_i, self.emb_l, self.v_W_l)  # apply embedding bag
-        # pyrefly: ignore [missing-attribute]
         z = self.interact_features(x, ly)
 
         z = z.to_sparse_coo()
-        # pyrefly: ignore [missing-attribute]
         z = torch.mm(z, self.top_l[0].weight.T).add(self.top_l[0].bias)
-        # pyrefly: ignore [missing-attribute]
         for layer in self.top_l[1:]:
             z = layer(z)
 
diff --git a/torch/ao/pruning/_experimental/data_sparsifier/quantization_utils.py b/torch/ao/pruning/_experimental/data_sparsifier/quantization_utils.py
index b727635d08151..097337aa8b42c 100644
--- a/torch/ao/pruning/_experimental/data_sparsifier/quantization_utils.py
+++ b/torch/ao/pruning/_experimental/data_sparsifier/quantization_utils.py
@@ -92,6 +92,7 @@ def post_training_sparse_quantize(
 
         # quantize
         for _, emb_module in embedding_modules:
+            # pyrefly: ignore [bad-argument-type]
             emb_module.qconfig = torch.ao.quantization.float_qparams_weight_only_qconfig
 
         torch.ao.quantization.prepare(model, inplace=True)
@@ -100,6 +101,7 @@ def post_training_sparse_quantize(
     else:
         # quantize
         for _, emb_module in embedding_modules:
+            # pyrefly: ignore [bad-argument-type]
             emb_module.qconfig = torch.ao.quantization.float_qparams_weight_only_qconfig
 
         torch.ao.quantization.prepare(model, inplace=True)
diff --git a/torch/ao/pruning/_experimental/pruner/FPGM_pruner.py b/torch/ao/pruning/_experimental/pruner/FPGM_pruner.py
index 1a89de12bd934..f5405253f9cb6 100644
--- a/torch/ao/pruning/_experimental/pruner/FPGM_pruner.py
+++ b/torch/ao/pruning/_experimental/pruner/FPGM_pruner.py
@@ -69,7 +69,6 @@ def _compute_distance(self, t):
         dist_matrix = self.dist_fn(t_flatten)
 
         # more similar with other filter indicates large in the sum of row
-        # pyrefly: ignore [bad-argument-type]
         distance = torch.sum(torch.abs(dist_matrix), 1)
 
         return distance
diff --git a/torch/ao/pruning/sparsifier/weight_norm_sparsifier.py b/torch/ao/pruning/sparsifier/weight_norm_sparsifier.py
index 0fd0368f15674..7ef1913f2725f 100644
--- a/torch/ao/pruning/sparsifier/weight_norm_sparsifier.py
+++ b/torch/ao/pruning/sparsifier/weight_norm_sparsifier.py
@@ -235,7 +235,6 @@ def update_mask(  # type: ignore[call-override, override]
             ww = self.norm_fn(getattr(module, tensor_name))
             tensor_mask = self._make_tensor_mask(
                 data=ww,
-                # pyrefly: ignore [missing-attribute]
                 input_shape=ww.shape,
                 sparsity_level=sparsity_level,
                 sparse_block_shape=sparse_block_shape,
diff --git a/torch/ao/quantization/__init__.py b/torch/ao/quantization/__init__.py
index e305801cba875..357a0bfece028 100644
--- a/torch/ao/quantization/__init__.py
+++ b/torch/ao/quantization/__init__.py
@@ -2,7 +2,6 @@
 
 import sys
 from collections.abc import Callable
-from typing import Optional, Union
 from typing_extensions import TypeAliasType
 
 import torch
diff --git a/torch/ao/quantization/backend_config/utils.py b/torch/ao/quantization/backend_config/utils.py
index 4d486a0611293..ed0c7f92bde4c 100644
--- a/torch/ao/quantization/backend_config/utils.py
+++ b/torch/ao/quantization/backend_config/utils.py
@@ -280,7 +280,9 @@ def _get_pattern_in_reversed_nested_tuple_format(
         (a, b, c) = config.pattern
         return (c, (b, a))
     else:
-        raise ValueError("Expected a tuple with 2 or 3 elements, got: ", config.pattern)
+        raise ValueError(
+            f"Expected a tuple with 2 or 3 elements, got: {config.pattern}"
+        )
 
 
 def _get_fuser_method_in_reversed_nested_tuple_format(
@@ -306,7 +308,7 @@ def _get_fuser_method_in_reversed_nested_tuple_format(
     if config._pattern_complex_format is not None:
         return config.fuser_method
     if not isinstance(config.pattern, tuple):
-        raise ValueError("Expected pattern to be a tuple, got: ", config.pattern)
+        raise ValueError(f"Expected pattern to be a tuple, got: {config.pattern}")
 
     # Pattern is specified in the simple tuple format, need to convert
     if len(config.pattern) == 2:
@@ -314,4 +316,6 @@ def _get_fuser_method_in_reversed_nested_tuple_format(
     elif len(config.pattern) == 3:
         return _reverse3(config.fuser_method)
     else:
-        raise ValueError("Expected a tuple with 2 or 3 elements, got: ", config.pattern)
+        raise ValueError(
+            f"Expected a tuple with 2 or 3 elements, got: {config.pattern}"
+        )
diff --git a/torch/ao/quantization/fx/_equalize.py b/torch/ao/quantization/fx/_equalize.py
index a7346daa283a5..95fc679d79ae4 100644
--- a/torch/ao/quantization/fx/_equalize.py
+++ b/torch/ao/quantization/fx/_equalize.py
@@ -3,6 +3,7 @@
 import warnings
 from collections import namedtuple
 from typing import Any
+from typing_extensions import TypeIs
 
 import torch
 import torch.ao.nn.intrinsic as nni
@@ -327,7 +328,9 @@ def node_supports_equalization(node: Node, modules) -> bool:
     return False
 
 
-def is_equalization_observer(observer: nn.Module) -> bool:
+def is_equalization_observer(
+    observer: nn.Module,
+) -> TypeIs[_InputEqualizationObserver | _WeightEqualizationObserver]:
     return isinstance(
         observer, (_InputEqualizationObserver, _WeightEqualizationObserver)
     )
diff --git a/torch/ao/quantization/fx/prepare.py b/torch/ao/quantization/fx/prepare.py
index 9bff92acaa0b2..71793a026fa9b 100644
--- a/torch/ao/quantization/fx/prepare.py
+++ b/torch/ao/quantization/fx/prepare.py
@@ -185,7 +185,6 @@ def _is_input_arg_dtype_supported_by_backend(
     if isinstance(arg, (list, tuple)):  # noqa: UP038
         return all(
             _is_input_arg_dtype_supported_by_backend(
-                # pyrefly: ignore [bad-argument-type]
                 a,
                 node,
                 qconfig,
@@ -740,7 +739,6 @@ def _maybe_insert_input_observer_for_arg_or_kwarg(
         for inner_arg in arg:
             new_inner_arg = _maybe_insert_input_observer_for_arg_or_kwarg(
                 node,
-                # pyrefly: ignore [bad-argument-type]
                 inner_arg,
                 qconfig,
                 model,
@@ -1160,7 +1158,6 @@ def _recursive_maybe_replace_node_with_obs(
         elif isinstance(maybe_node, (list, tuple)):  # noqa: UP038
             results = [
                 _recursive_maybe_replace_node_with_obs(
-                    # pyrefly: ignore [bad-argument-type]
                     inner_node,
                     model,
                     named_modules,
diff --git a/torch/ao/quantization/fx/quantize_handler.py b/torch/ao/quantization/fx/quantize_handler.py
index fe5ffb43f9cb3..0bd8d7fe3a174 100644
--- a/torch/ao/quantization/fx/quantize_handler.py
+++ b/torch/ao/quantization/fx/quantize_handler.py
@@ -155,11 +155,9 @@ def _get_pattern_to_quantize_handlers(
         num_tensor_args_to_observation_type = (
             config._num_tensor_args_to_observation_type
         )
-        # pyrefly: ignore [unsupported-operation]
         pattern_to_quantize_handlers[pattern] = _get_quantize_handler_cls(
             observation_type, dtype_configs, num_tensor_args_to_observation_type
         )
-    # pyrefly: ignore [bad-return]
     return pattern_to_quantize_handlers
 
 
diff --git a/torch/ao/quantization/fx/tracer.py b/torch/ao/quantization/fx/tracer.py
index 2c1635936845a..89aff43df28f4 100644
--- a/torch/ao/quantization/fx/tracer.py
+++ b/torch/ao/quantization/fx/tracer.py
@@ -31,7 +31,7 @@ def __init__(
         # We can change this if there is a use case that configures
         # qconfig using top level module type
         self.scope = Scope("", None)
-        self.record_stack_traces = True
+        self.record_stack_traces = not torch.fx.config.do_not_emit_stack_traces
 
     def is_leaf_module(self, m: torch.nn.Module, module_qualified_name: str) -> bool:
         return (
diff --git a/torch/ao/quantization/observer.py b/torch/ao/quantization/observer.py
index eb16d5da00c07..858f44d8708cd 100644
--- a/torch/ao/quantization/observer.py
+++ b/torch/ao/quantization/observer.py
@@ -144,7 +144,7 @@ def _with_callable_args(cls_or_self, **kwargs):
     return r.with_callable_args(**kwargs)
 
 
-ABC: Any = ABCMeta("ABC", (object,), {})  # compatible with Python 2 *and* 3:
+ABC: Any = ABCMeta("ABC", (object,), {})
 
 
 class ObserverBase(ABC, nn.Module):
@@ -1158,7 +1158,6 @@ def _non_linear_param_search(self) -> tuple[torch.Tensor, torch.Tensor]:
             # find the left and right bins between the quantile bounds
             l = start_bin
             r = end_bin
-            # pyrefly: ignore [bad-assignment]
             while l < end_bin and cSum[l] < next_alpha * total:
                 l = l + 1
             while r > start_bin and cSum[r] > next_beta * total:
diff --git a/torch/ao/quantization/qconfig.py b/torch/ao/quantization/qconfig.py
index dc083c9a08696..956d2c6a07ab7 100644
--- a/torch/ao/quantization/qconfig.py
+++ b/torch/ao/quantization/qconfig.py
@@ -2,7 +2,7 @@
 import copy
 import warnings
 from collections import namedtuple
-from typing import Any, Union
+from typing import Any
 from typing_extensions import deprecated, TypeAliasType
 
 import torch
@@ -621,9 +621,9 @@ def configure_constructor_to_put_obs_on_module_device(original_constructor):
     return QConfig(activation, weight)
 
 
-_ObserverOrFakeQuantizeConstructor = Union[
-    _PartialWrapper, type[ObserverBase], type[FakeQuantizeBase]
-]
+_ObserverOrFakeQuantizeConstructor = (
+    _PartialWrapper | type[ObserverBase] | type[FakeQuantizeBase]
+)
 
 
 def _obs_or_fq_ctr_equals(
diff --git a/torch/autograd/__init__.py b/torch/autograd/__init__.py
index 034e1512ddea4..53b5f5846d3c3 100644
--- a/torch/autograd/__init__.py
+++ b/torch/autograd/__init__.py
@@ -10,7 +10,7 @@
 
 import warnings
 from collections.abc import Sequence
-from typing import cast, Optional, Union
+from typing import cast
 
 import torch
 from torch import _vmap_internals
@@ -53,12 +53,12 @@
     "variable",
 ]
 
-_OptionalTensor = Optional[torch.Tensor]
-_ShapeorNestedShape = Union[_size, Sequence[_size], torch.Tensor]
+_OptionalTensor = torch.Tensor | None
+_ShapeorNestedShape = _size | Sequence[_size] | torch.Tensor
 
 
 def _calculate_shape(
-    output: Union[torch.Tensor, graph.GradientEdge],
+    output: torch.Tensor | graph.GradientEdge,
     grad: torch.Tensor,
     is_grads_batched: bool,
 ) -> tuple[_ShapeorNestedShape, _ShapeorNestedShape]:
@@ -87,7 +87,7 @@ def _calculate_shape(
 
 
 def _make_grads(
-    outputs: Union[Sequence[torch.Tensor], Sequence[graph.GradientEdge]],
+    outputs: Sequence[torch.Tensor] | Sequence[graph.GradientEdge],
     grads: Sequence[_OptionalTensor],
     is_grads_batched: bool,
 ) -> tuple[_OptionalTensor, ...]:
@@ -95,7 +95,7 @@ def _make_grads(
 
     for out, grad in zip(outputs, grads):
         # pyrefly: ignore [redundant-cast]
-        out = cast(Union[torch.Tensor, graph.GradientEdge], out)
+        out = cast(torch.Tensor | graph.GradientEdge, out)
         out_size = None
         out_device = None
 
@@ -143,7 +143,7 @@ def _make_grads(
                 shape_matches = expect_true(sym_eq(out_size, first_grad.size()))
 
             if not shape_matches:
-                out = cast(Union[torch.Tensor, graph.GradientEdge], out)  # type: ignore[redundant-cast]
+                out = cast(torch.Tensor | graph.GradientEdge, out)  # type: ignore[redundant-cast]
                 out_shape, grad_shape = _calculate_shape(
                     out, first_grad, is_grads_batched
                 )
@@ -241,7 +241,7 @@ def _make_grads(
 
 
 def _tensor_or_tensors_to_tuple(
-    tensors: Optional[_TensorOrOptionalTensors], length: int
+    tensors: _TensorOrOptionalTensors | None, length: int
 ) -> tuple[_OptionalTensor, ...]:
     if tensors is None:
         return (None,) * length
@@ -252,11 +252,11 @@ def _tensor_or_tensors_to_tuple(
 
 def backward(
     tensors: _TensorOrTensorsOrGradEdge,
-    grad_tensors: Optional[_TensorOrOptionalTensors] = None,
-    retain_graph: Optional[bool] = None,
+    grad_tensors: _TensorOrOptionalTensors | None = None,
+    retain_graph: bool | None = None,
     create_graph: bool = False,
-    grad_variables: Optional[_TensorOrOptionalTensors] = None,
-    inputs: Optional[_TensorOrTensorsOrGradEdge] = None,
+    grad_variables: _TensorOrOptionalTensors | None = None,
+    inputs: _TensorOrTensorsOrGradEdge | None = None,
 ) -> None:
     r"""Compute the sum of gradients of given tensors with respect to graph leaves.
 
@@ -336,7 +336,7 @@ def backward(
                 "use `grad_tensors`."
             )
 
-    inputs_tuple: tuple[Union[torch.Tensor, graph.GradientEdge], ...]
+    inputs_tuple: tuple[torch.Tensor | graph.GradientEdge, ...]
     if inputs is None:
         inputs_tuple = ()
     elif isinstance(inputs, (torch.Tensor, graph.GradientEdge)):
@@ -347,9 +347,7 @@ def backward(
             raise RuntimeError("`inputs` argument to `backward()` cannot be empty.")
 
     if is_tensor_like(tensors) or isinstance(tensors, graph.GradientEdge):
-        tensors = cast(
-            Union[tuple[torch.Tensor], tuple[graph.GradientEdge]], (tensors,)
-        )
+        tensors = cast(tuple[torch.Tensor] | tuple[graph.GradientEdge], (tensors,))
     else:
         # pyrefly: ignore [bad-argument-type]
         tensors = tuple(tensors)
@@ -392,11 +390,11 @@ def backward(
 def grad(
     outputs: _TensorOrTensorsOrGradEdge,
     inputs: _TensorOrTensorsOrGradEdge,
-    grad_outputs: Optional[_TensorOrOptionalTensors] = None,
-    retain_graph: Optional[bool] = None,
+    grad_outputs: _TensorOrOptionalTensors | None = None,
+    retain_graph: bool | None = None,
     create_graph: bool = False,
     only_inputs: bool = True,
-    allow_unused: Optional[bool] = None,
+    allow_unused: bool | None = None,
     is_grads_batched: bool = False,
     materialize_grads: bool = False,
 ) -> tuple[torch.Tensor, ...]:
@@ -464,7 +462,7 @@ def grad(
         allow_unused = materialize_grads
     if is_tensor_like(outputs) or isinstance(outputs, graph.GradientEdge):
         outputs = cast(
-            Union[Sequence[torch.Tensor], Sequence[graph.GradientEdge]], (outputs,)
+            Sequence[torch.Tensor] | Sequence[graph.GradientEdge], (outputs,)
         )
     else:
         # pyrefly: ignore [bad-argument-type]
@@ -549,6 +547,7 @@ def vjp(gO):
         result = tuple(
             output
             if output is not None
+            # pyrefly: ignore [bad-argument-type]
             else torch.zeros_like(input, requires_grad=create_graph)
             for (output, input) in zip(result, inputs)
         )
diff --git a/torch/autograd/forward_ad.py b/torch/autograd/forward_ad.py
index 8fcb64beba3b4..ce3eb94266271 100644
--- a/torch/autograd/forward_ad.py
+++ b/torch/autograd/forward_ad.py
@@ -1,6 +1,6 @@
 # mypy: allow-untyped-defs
 import os
-from typing import Any, NamedTuple, Optional
+from typing import Any, NamedTuple
 
 import torch
 
@@ -135,7 +135,7 @@ class UnpackedDualTensor(NamedTuple):
     """
 
     primal: torch.Tensor
-    tangent: Optional[torch.Tensor]
+    tangent: torch.Tensor | None
 
 
 def unpack_dual(tensor, *, level=None):
diff --git a/torch/autograd/function.py b/torch/autograd/function.py
index 00375eb4f1afd..4c639fd79b280 100644
--- a/torch/autograd/function.py
+++ b/torch/autograd/function.py
@@ -5,7 +5,7 @@
 import warnings
 from collections import OrderedDict
 from collections.abc import Callable
-from typing import Any, Concatenate, Optional, TypeVar
+from typing import Any, Concatenate, TypeVar
 from typing_extensions import deprecated, ParamSpec
 
 import torch
@@ -747,7 +747,7 @@ def _unflatten(input, proto):
     # unflatten a list or tuple input into a nested list/tuple structure
     # specified by proto
     def unflatten_helper(input, proto):
-        res: list[Optional[torch.Tensor]] = []
+        res: list[torch.Tensor | None] = []
         if hasattr(proto, "_jit_wrap"):
             return proto._jit_wrap(input)
         if not isinstance(proto, (list, tuple)):
diff --git a/torch/autograd/grad_mode.py b/torch/autograd/grad_mode.py
index 6044ba9ac9232..0775d43994283 100644
--- a/torch/autograd/grad_mode.py
+++ b/torch/autograd/grad_mode.py
@@ -1,5 +1,5 @@
 # mypy: allow-untyped-defs
-from typing import Any, Union
+from typing import Any
 
 import torch
 from torch.utils._contextlib import (
@@ -376,11 +376,15 @@ class _force_original_view_tracking(_DecoratorContextManager):
 
     def __init__(self, mode: bool) -> None:
         self.prev = torch._C._is_view_replay_enabled()
-        torch._C._set_view_replay_enabled(mode)
         self.mode = mode
+        torch._C._set_view_replay_enabled(mode)
+
+    def __call__(self, orig_func: F) -> F:
+        torch._C._set_view_replay_enabled(self.prev)
+        return super().__call__(orig_func)
 
     def __enter__(self) -> None:
-        pass
+        torch._C._set_view_replay_enabled(self.mode)
 
     def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None:
         torch._C._set_view_replay_enabled(self.prev)
@@ -412,7 +416,7 @@ class _unsafe_preserve_version_counter(_DecoratorContextManager):
 
     """
 
-    def __init__(self, tensors: Union[torch.Tensor, tuple[torch.Tensor, ...]]) -> None:
+    def __init__(self, tensors: torch.Tensor | tuple[torch.Tensor, ...]) -> None:
         self.tensors = (tensors,) if isinstance(tensors, torch.Tensor) else tensors
         if not isinstance(self.tensors, tuple):
             raise AssertionError("Expected tensors to be a tuple")
diff --git a/torch/autograd/gradcheck.py b/torch/autograd/gradcheck.py
index 2f6d111948b5d..4d5a4585b03bd 100644
--- a/torch/autograd/gradcheck.py
+++ b/torch/autograd/gradcheck.py
@@ -4,7 +4,6 @@
 import warnings
 from collections.abc import Callable, Iterable
 from itertools import product
-from typing import Optional, Union
 from typing_extensions import deprecated
 
 import torch
@@ -83,7 +82,7 @@ def _allocate_jacobians_with_outputs(
 
 
 def _iter_tensors(
-    x: Union[torch.Tensor, Iterable[torch.Tensor]], only_requiring_grad: bool = False
+    x: torch.Tensor | Iterable[torch.Tensor], only_requiring_grad: bool = False
 ) -> Iterable[torch.Tensor]:
     if is_tensor_like(x):
         # mypy doesn't narrow type of `x` to torch.Tensor
@@ -438,7 +437,7 @@ def _combine_jacobian_cols(
 
 
 def _prepare_input(
-    input: torch.Tensor, maybe_perturbed_input: Optional[torch.Tensor], fast_mode=False
+    input: torch.Tensor, maybe_perturbed_input: torch.Tensor | None, fast_mode=False
 ) -> torch.Tensor:
     # Prepares the inputs to be passed into the function while including the new
     # modified input.
@@ -899,7 +898,7 @@ def _get_analytical_jacobian(inputs, outputs, input_idx, output_idx):
 
 def _compute_analytical_jacobian_rows(
     vjp_fn, sample_output
-) -> list[list[Optional[torch.Tensor]]]:
+) -> list[list[torch.Tensor | None]]:
     # Computes Jacobian row-by-row by projecting `vjp_fn` = v^T J on standard basis
     # vectors: vjp_fn(e) = e^T J is a corresponding row of the Jacobian.
     # NB: this function does not assume vjp_fn(v) to return tensors with the same
@@ -910,7 +909,7 @@ def _compute_analytical_jacobian_rows(
     )
     flat_grad_out = grad_out_base.view(-1)
     # jacobians_rows[i][j] is the Jacobian jth row for the ith input
-    jacobians_rows: list[list[Optional[torch.Tensor]]] = []
+    jacobians_rows: list[list[torch.Tensor | None]] = []
     for j in range(flat_grad_out.numel()):
         flat_grad_out.zero_()
         flat_grad_out[j] = 1.0  # projection for jth row of Jacobian
@@ -926,9 +925,9 @@ def _compute_analytical_jacobian_rows(
 
 def _get_analytical_vjps_wrt_specific_output(
     vjp_fn, sample_output, v
-) -> list[list[Optional[torch.Tensor]]]:
+) -> list[list[torch.Tensor | None]]:
     grad_inputs = vjp_fn(v.reshape(sample_output.shape))
-    vjps: list[list[Optional[torch.Tensor]]] = [
+    vjps: list[list[torch.Tensor | None]] = [
         [vjp.clone() if isinstance(vjp, torch.Tensor) else None] for vjp in grad_inputs
     ]
     return vjps
@@ -1998,7 +1997,7 @@ def _fast_gradcheck(
 # the '...' first argument of Callable can be replaced with VarArg(Tensor).
 # For now, we permit any input.
 def gradcheck(
-    func: Callable[..., Union[_TensorOrTensors]],  # See Note [VarArg of Tensors]
+    func: Callable[..., _TensorOrTensors],  # See Note [VarArg of Tensors]
     inputs: _TensorOrTensors,
     *,
     eps: float = 1e-6,
@@ -2013,7 +2012,7 @@ def gradcheck(
     check_forward_ad: bool = False,
     check_backward_ad: bool = True,
     fast_mode: bool = False,
-    masked: Optional[bool] = None,
+    masked: bool | None = None,
 ) -> bool:  # noqa: D400,D205
     r"""Check gradients computed via small finite differences against analytical
     gradients wrt tensors in :attr:`inputs` that are of floating point or complex type
@@ -2168,7 +2167,7 @@ def _gradcheck_helper(
 def gradgradcheck(
     func: Callable[..., _TensorOrTensors],  # See Note [VarArg of Tensors]
     inputs: _TensorOrTensors,
-    grad_outputs: Optional[_TensorOrOptionalTensors] = None,
+    grad_outputs: _TensorOrOptionalTensors | None = None,
     *,
     eps: float = 1e-6,
     atol: float = 1e-5,
diff --git a/torch/autograd/graph.py b/torch/autograd/graph.py
index 31d987ae45361..0c285fbe7b7c4 100644
--- a/torch/autograd/graph.py
+++ b/torch/autograd/graph.py
@@ -1,5 +1,6 @@
 import abc
 import contextlib
+import contextvars
 import functools
 import logging
 import threading
@@ -208,7 +209,7 @@ class GradientEdge(NamedTuple):
     output_nr: int
     # This token can be used to ensure the graph stays alive when it cannot be
     # done via the node field
-    ownership_token: Optional[Node] = None
+    ownership_token: Node | None = None
 
 
 def get_gradient_edge(tensor: torch.Tensor) -> GradientEdge:
@@ -239,7 +240,7 @@ def get_gradient_edge(tensor: torch.Tensor) -> GradientEdge:
     return GradientEdge(grad_fn, tensor.output_nr, ownership_token=token)
 
 
-def increment_version(tensor: Union[torch.Tensor, Iterable[torch.Tensor]]) -> None:
+def increment_version(tensor: torch.Tensor | Iterable[torch.Tensor]) -> None:
     """Update autograd metadata tracking whether the given Tensor was modified in place.
 
     This is to enable more accurate error checking within the autograd engine.
@@ -355,9 +356,18 @@ class save_on_cpu(saved_tensors_hooks):
     Use this context-manager to trade compute for GPU memory usage (e.g.
     when your model doesn't fit in GPU memory during training).
 
+    .. warning::
+
+        When ``pin_memory=True``, the GPU to CPU copy during packing is
+        asynchronous. Accessing saved tensors on CPU (e.g. via
+        ``grad_fn._saved_self``) before the CUDA stream has finished may
+        yield incorrect data. Call :func:`torch.cuda.synchronize` first
+        if you need to read them.
+
     Args:
         pin_memory (bool): If ``True`` tensors will be saved to CPU pinned memory
-                           during packing and copied to GPU asynchronously during unpacking.
+                           during packing and copied to GPU asynchronously during both
+                           packing and unpacking.
                            Defaults to ``False``.
                            Also see :ref:`cuda-memory-pinning`.
 
@@ -391,13 +401,14 @@ def __init__(self, pin_memory: bool = False, device_type: str = "cuda") -> None:
         def pack_to_cpu(tensor: torch.Tensor) -> tuple[torch.device, torch.Tensor]:
             if not pin_memory:
                 return (tensor.device, tensor.cpu())
+            is_pinnable = device_module.is_available() and not tensor.is_sparse
             packed = torch.empty(
                 tensor.size(),
                 dtype=tensor.dtype,
                 layout=tensor.layout,
-                pin_memory=(device_module.is_available() and not tensor.is_sparse),
+                pin_memory=is_pinnable,
             )
-            packed.copy_(tensor)
+            packed.copy_(tensor, non_blocking=is_pinnable)
             return (tensor.device, packed)
 
         def unpack_from_cpu(packed: tuple[torch.device, torch.Tensor]) -> torch.Tensor:
@@ -469,10 +480,8 @@ def __setstate__(self, state: tuple[RemovableHandle, ...]) -> None:
 
 def register_multi_grad_hook(
     tensors: Sequence[torch.Tensor],
-    fn: Union[
-        Callable[[Sequence[Optional[torch.Tensor]]], None],
-        Callable[[torch.Tensor], None],
-    ],
+    fn: Callable[[Sequence[torch.Tensor | None]], None]
+    | Callable[[torch.Tensor], None],
     *,
     mode: Literal["all", "any"] = "all",
 ) -> RemovableHandle:
@@ -532,7 +541,7 @@ def register_multi_grad_hook(
     if mode == "all":
         count: dict[int, int] = {}
         nb_calls = None
-        buffer: dict[int, list[Optional[torch.Tensor]]] = {}
+        buffer: dict[int, list[torch.Tensor | None]] = {}
 
         grad_fns = list(map(_get_grad_fn_or_grad_acc, tensors))
         len_tensors = len(tensors)
@@ -563,7 +572,7 @@ def inner_hook(grad: torch.Tensor) -> None:
                 if nb_calls is None:
                     raise AssertionError("Expected nb_calls to be set")
                 if curr_count == nb_calls - 1:
-                    fn = cast(Callable[[Sequence[Optional[torch.Tensor]]], None], fn)
+                    fn = cast(Callable[[Sequence[torch.Tensor | None]], None], fn)
                     fn(buffer[id])
                     del count[id]
                     del buffer[id]
@@ -660,7 +669,7 @@ def pack_hook(tensor: torch.Tensor) -> _Handle:
             tid = _get_tid(tensor)
             sid = _get_sid(tensor)
             # Tensors saved for backward have an entry in _tid_to_weakhandle
-            handle: Optional[_Handle] = None
+            handle: _Handle | None = None
 
             # Save aliasing information
             ctx.sid_to_tid[sid].add(tid)
@@ -702,7 +711,7 @@ def __torch_dispatch__(
         func: "OpOverload",
         types: Iterable[type],
         args: tuple[Any, ...] = (),
-        kwargs: Optional[dict[Any, Any]] = None,
+        kwargs: dict[Any, Any] | None = None,
     ) -> Any:
         kwargs = kwargs or {}
 
@@ -810,7 +819,7 @@ def allow_mutation_on_saved_tensors() -> Generator[
 
 
 def _register_logging_hooks_on_whole_graph(
-    t_outputs: Sequence[Union[torch.Tensor, GradientEdge]],
+    t_outputs: Sequence[torch.Tensor | GradientEdge],
 ) -> Callable[[], None]:
     grad_fns = list(map(_get_grad_fn_or_grad_acc, t_outputs))
 
@@ -834,7 +843,7 @@ def iter_graph(roots: list[Node]) -> Iterator[Node]:
 
             yield node
 
-    def fmt(t: Optional[torch.Tensor]) -> str:
+    def fmt(t: torch.Tensor | None) -> str:
         # Avoid circular import
         from torch.utils._dtype_abbrs import dtype_abbrs
 
@@ -842,7 +851,7 @@ def fmt(t: Optional[torch.Tensor]) -> str:
             return "None"
         return f"{dtype_abbrs[t.dtype]}[{', '.join(map(str, t.shape))}]"
 
-    def prehook(grad_outputs: Sequence[Optional[torch.Tensor]]) -> None:
+    def prehook(grad_outputs: Sequence[torch.Tensor | None]) -> None:
         node = torch._C._current_autograd_node()
         grad_outputs_str = f"[{','.join(fmt(t) for t in grad_outputs)}]"
         log_str = f"Executing: {node} with grad_outputs: {grad_outputs_str}"
@@ -858,13 +867,17 @@ def unregister_hooks() -> None:
 
 
 def _engine_run_backward(
-    t_outputs: Sequence[Union[torch.Tensor, GradientEdge]],
+    t_outputs: Sequence[torch.Tensor | GradientEdge],
     *args: Any,
     **kwargs: Any,
 ) -> tuple[torch.Tensor, ...]:
     attach_logging_hooks = log.getEffectiveLevel() <= logging.DEBUG
     if attach_logging_hooks:
         unregister_hooks = _register_logging_hooks_on_whole_graph(t_outputs)
+
+    # Need to save the context so compiler config will be visible in device threads
+    torch._C._stash_obj_in_tls("context", contextvars.copy_context())
+
     try:
         return Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
             t_outputs, *args, **kwargs
@@ -872,3 +885,4 @@ def _engine_run_backward(
     finally:
         if attach_logging_hooks:
             unregister_hooks()  # type: ignore[possibly-undefined]
+        torch._C._stash_obj_in_tls("context", None)
diff --git a/torch/autograd/profiler.py b/torch/autograd/profiler.py
index a43249fa566b6..d067c8fc4d0a4 100644
--- a/torch/autograd/profiler.py
+++ b/torch/autograd/profiler.py
@@ -225,7 +225,8 @@ def __init__(
         experimental_config=None,
         acc_events=False,
         custom_trace_id_callback=None,
-        post_processing_timeout_s: Optional[float] = None,
+        post_processing_timeout_s: float | None = None,
+        activity_filters: dict[ProfilerActivity, set[str]] | None = None,
     ):
         self.enabled: bool = enabled
         if not self.enabled:
@@ -238,12 +239,12 @@ def __init__(
                 FutureWarning,
                 stacklevel=2,
             )
-            self.use_device: Optional[str] = "cuda"
+            self.use_device: str | None = "cuda"
         else:
             self.use_device = use_device
         # TODO Consider changing _function_events into data structure with size cap
-        self._function_events: Optional[EventList] = None
-        self._old_function_events: Optional[EventList] = None
+        self._function_events: EventList | None = None
+        self._old_function_events: EventList | None = None
         # Function event processing is done lazily
         self._needs_processing = False
         self.entered = False
@@ -258,12 +259,13 @@ def __init__(
         if experimental_config is None:
             experimental_config = _ExperimentalConfig()
         self.experimental_config = experimental_config
-        self.kineto_results: Optional[_ProfilerResult] = None
+        self.kineto_results: _ProfilerResult | None = None
         self.profiling_start_time_ns = 0
         self.profiling_end_time_ns = 0
         self._stats = _ProfilerStats()
         self.custom_trace_id_callback = custom_trace_id_callback
         self.post_processing_timeout_s = post_processing_timeout_s
+        self.activity_filters = activity_filters or {}
         self.trace_id = ""
         if not self.use_cpu:
             if not use_kineto:
@@ -327,17 +329,18 @@ def __init__(
                 )
             self.kineto_activities.add(ProfilerActivity.HPU)
         elif self.use_device is not None and self.use_device != "privateuseone":
-            if (
-                not use_kineto
-                or ProfilerActivity.PrivateUse1 not in _supported_activities()
-            ):
+            if use_kineto:
+                # Native tracing mode: use KINETO_PRIVATEUSE1 with registered IActivityProfiler
+                self.profiler_kind = ProfilerState.KINETO_PRIVATEUSE1
+                if ProfilerActivity.PrivateUse1 in _supported_activities():
+                    self.kineto_activities.add(ProfilerActivity.PrivateUse1)
+            else:
+                # Marker-only mode: use fallback state
                 if not self.use_cpu:
                     raise AssertionError(
-                        "Legacy custombackend profiling requires use_cpu=True"
+                        "Legacy privateuse1 profiling requires use_cpu=True"
                     )
                 self.profiler_kind = ProfilerState.KINETO_PRIVATEUSE1_FALLBACK
-            else:
-                self.kineto_activities.add(ProfilerActivity.PrivateUse1)
 
         if len(self.kineto_activities) == 0:
             raise AssertionError("No activities specified for the profiler")
@@ -387,7 +390,11 @@ def __enter__(self):
     def _prepare_trace(self):
         self.entered = True
         t0 = perf_counter_ns()
-        _prepare_profiler(self.config(create_trace_id=True), self.kineto_activities)
+        _prepare_profiler(
+            self.config(create_trace_id=True),
+            self.kineto_activities,
+            activity_filter=self.activity_filters,
+        )
         t1 = perf_counter_ns()
         self._stats.profiler_prepare_call_duration_us = int((t1 - t0) / 1000)
 
@@ -478,9 +485,6 @@ def _ensure_function_events(self):
                 self._function_events.append(evt)
             self._old_function_events = None
 
-        if self._function_events is None:
-            raise RuntimeError("Profiler didn't finish running")
-
     @property
     def function_events(self):
         if self._function_events is None or self._needs_processing:
@@ -576,11 +580,12 @@ def self_cpu_time_total(self):
         return self._function_events.self_cpu_time_total
 
     def _parse_kineto_results(
-        self, result: _ProfilerResult, timeout_s: Optional[float] = None
+        self, result: _ProfilerResult, timeout_s: float | None = None
     ):
         # result.events() has most of the events - PyTorch op-level and device-level events
 
         timeout_ns = int(timeout_s * 1e9) if timeout_s is not None else None
+        result_events = result.events()
         if timeout_ns is not None and timeout_ns < 0:
             raise ValueError("timeout_s must be non-negative")
         start_time_ns = perf_counter_ns()
@@ -597,10 +602,10 @@ def _check_timeout() -> bool:
 
         trace_start_ns = result.trace_start_ns()
         mem_records = [
-            [evt, False] for evt in result.events() if evt.name() == MEMORY_EVENT_NAME
+            [evt, False] for evt in result_events if evt.name() == MEMORY_EVENT_NAME
         ]
         oom_records = [
-            evt for evt in result.events() if evt.name() == OUT_OF_MEMORY_EVENT_NAME
+            evt for evt in result_events if evt.name() == OUT_OF_MEMORY_EVENT_NAME
         ]
         mem_records_acc = MemRecordsAcc(mem_records)
 
@@ -634,7 +639,7 @@ def _device_memory_usage(mem_record):
         frontend_function_events = []
         device_corr_map: dict[int, list[FunctionEvent]] = {}
         max_evt_id = 0
-        for kineto_event in result.events():
+        for kineto_event in result_events:
             if _check_timeout():
                 break
 
@@ -645,14 +650,13 @@ def _device_memory_usage(mem_record):
                 continue
             rel_start_ns = kineto_event.start_ns() - trace_start_ns
             rel_end_ns = kineto_event.end_ns() - trace_start_ns
-            abs_end_ns = kineto_event.end_ns()
 
             cpu_memory_usage = 0
             device_memory_usage = 0
             if kineto_event.device_type() == DeviceType.CPU:
                 # find the corresponding memory allocation events
                 for mem_record in mem_records_acc.in_interval(
-                    kineto_event.start_ns(), abs_end_ns
+                    kineto_event.start_ns(), kineto_event.end_ns()
                 ):
                     cpu_memory_usage += _cpu_memory_usage(mem_record[0])
                     device_memory_usage += _device_memory_usage(mem_record[0])
@@ -690,7 +694,21 @@ def _device_memory_usage(mem_record):
                 device_resource_id=kineto_event.device_resource_id(),
                 flops=kineto_event.flops(),
                 is_user_annotation=kineto_event.is_user_annotation(),
+                is_python_function=kineto_event.is_python_function(),
+                activity_type=kineto_event.activity_type(),
                 metadata_json=kineto_event.metadata_json(),
+                extra_meta=kineto_event.extra_meta() or None,
+                flow_id=kineto_event.flow_id(),
+                flow_type=kineto_event.flow_type(),
+                flow_start=kineto_event.flow_start(),
+                external_id=kineto_event.external_id(),
+                linked_correlation_id=kineto_event.linked_correlation_id(),
+                structured_input_shapes=kineto_event.structured_input_shapes(),
+                structured_input_strides=kineto_event.structured_input_strides(),
+                input_dtypes=kineto_event.dtypes(),
+                python_id=kineto_event.python_id(),
+                python_parent_id=kineto_event.python_parent_id(),
+                python_module_id=kineto_event.python_module_id(),
             )
             max_evt_id = max(max_evt_id, fe.id)
             if fe.device_type == DeviceType.CPU and not fe.is_async:
@@ -742,7 +760,7 @@ def _device_memory_usage(mem_record):
                         # parents and children
                         f_evt.thread = fe.thread
 
-        def createFunctionEventForMemoryEvents(evt):
+        def _create_function_event_for_memory_events(evt):
             rel_start_ns = evt.start_ns() - trace_start_ns
             fe = FunctionEvent(
                 id=max_evt_id,
@@ -773,7 +791,7 @@ def createFunctionEventForMemoryEvents(evt):
 
             if not mem_record[1]:
                 max_evt_id += 1
-                fe = createFunctionEventForMemoryEvents(mem_record[0])
+                fe = _create_function_event_for_memory_events(mem_record[0])
                 all_function_events.append(fe)
 
         for oom_record in oom_records:
@@ -781,7 +799,7 @@ def createFunctionEventForMemoryEvents(evt):
                 break
 
             max_evt_id += 1
-            fe = createFunctionEventForMemoryEvents(oom_record)
+            fe = _create_function_event_for_memory_events(oom_record)
             all_function_events.append(fe)
 
         if timed_out:
@@ -838,9 +856,9 @@ class record_function(_ContextDecorator):
 
     """
 
-    def __init__(self, name: str, args: Optional[str] = None):
+    def __init__(self, name: str, args: str | None = None):
         self.name: str = name
-        self.args: Optional[str] = args
+        self.args: str | None = args
         # Whether or not we should run record function's end callbacks when exiting.
         self.run_callbacks_on_exit: bool = True
         # TODO: TorchScript ignores standard type annotation here
diff --git a/torch/autograd/profiler_util.py b/torch/autograd/profiler_util.py
index 38f5d9153883f..5461434f170f9 100644
--- a/torch/autograd/profiler_util.py
+++ b/torch/autograd/profiler_util.py
@@ -1,10 +1,12 @@
 # mypy: allow-untyped-defs
 import bisect
 import itertools
+import json
 import math
 from collections import defaultdict, namedtuple
+from collections.abc import Callable
 from operator import attrgetter
-from typing import Any, Optional
+from typing import Any, NamedTuple
 from typing_extensions import deprecated
 
 import torch
@@ -13,6 +15,7 @@
 
 __all__ = [
     "EventList",
+    "EventMetadata",
     "FormattedTimesMixin",
     "Interval",
     "Kernel",
@@ -549,6 +552,106 @@ def elapsed_us(self):
 Kernel = namedtuple("Kernel", ["name", "device", "duration"])
 
 
+class EventMetadata(NamedTuple):
+    # Kernel fields
+    registers_per_thread: int | None
+    shared_memory: int | None
+    grid: list[int] | None
+    block: list[int] | None
+    priority: int | None
+    blocks_per_sm: float | None
+    warps_per_sm: float | None
+    occupancy: float | None
+    queued: int | None
+    graph_id: int | None
+    graph_node_id: int | None
+    stream: int | None
+    context: int | None
+    # Memory fields
+    bytes: int | None
+    bandwidth_gb_s: float | None
+    # NCCL fields
+    collective_name: str | None
+    dtype: str | None
+    in_msg_nelems: int | None
+    out_msg_nelems: int | None
+    in_split_size: str | None
+    out_split_size: str | None
+    global_rank_start: int | None
+    global_rank_stride: int | None
+    group_size: int | None
+    process_group_name: str | None
+    process_group_desc: str | None
+    group_ranks: str | None
+    rank: int | None
+    src_rank: int | None
+    dst_rank: int | None
+    seq: int | None
+    is_async: bool | None
+
+
+def _to_int_list(v: str) -> list[int]:
+    return json.loads(v)
+
+
+def _to_str(v: str) -> str:
+    return v.strip('"')
+
+
+def _to_bool(v: str) -> bool:
+    return v in ("1", "true")
+
+
+# Kineto key → (EventMetadata field name, converter from string)
+_EVENT_METADATA_KEYS: dict[str, tuple[str, Callable[[str], Any]]] = {
+    "registers per thread": ("registers_per_thread", int),
+    "shared memory": ("shared_memory", int),
+    "grid": ("grid", _to_int_list),
+    "block": ("block", _to_int_list),
+    "priority": ("priority", int),
+    "blocks per SM": ("blocks_per_sm", float),
+    "warps per SM": ("warps_per_sm", float),
+    "est. achieved occupancy %": ("occupancy", float),
+    "queued": ("queued", int),
+    "graph id": ("graph_id", int),
+    "graph node id": ("graph_node_id", int),
+    "stream": ("stream", int),
+    "context": ("context", int),
+    "bytes": ("bytes", int),
+    "memory bandwidth (GB/s)": ("bandwidth_gb_s", float),
+    "Collective name": ("collective_name", _to_str),
+    "dtype": ("dtype", _to_str),
+    "In msg nelems": ("in_msg_nelems", int),
+    "Out msg nelems": ("out_msg_nelems", int),
+    "In split size": ("in_split_size", _to_str),
+    "Out split size": ("out_split_size", _to_str),
+    "Global rank start": ("global_rank_start", int),
+    "Global rank stride": ("global_rank_stride", int),
+    "Group size": ("group_size", int),
+    "Process Group Name": ("process_group_name", _to_str),
+    "Process Group Description": ("process_group_desc", _to_str),
+    "Process Group Ranks": ("group_ranks", _to_str),
+    "Rank": ("rank", int),
+    "Src Rank": ("src_rank", int),
+    "Dst Rank": ("dst_rank", int),
+    "Seq": ("seq", int),
+    "Is asynchronized op": ("is_async", _to_bool),
+}
+
+
+def _build_metadata(extra_meta):
+    fields: dict[str, Any] = {}
+    any_populated = False
+    for kineto_key, (field_name, convert) in _EVENT_METADATA_KEYS.items():
+        v = extra_meta.get(kineto_key)
+        if v is not None:
+            fields[field_name] = convert(v)
+            any_populated = True
+        else:
+            fields[field_name] = None
+    return EventMetadata(**fields) if any_populated else None
+
+
 class FunctionEvent(FormattedTimesMixin):
     """Profiling information about a single function.
 
@@ -573,7 +676,10 @@ class FunctionEvent(FormattedTimesMixin):
         count (int): Number of times this event was called (usually 1).
         cpu_children (List[FunctionEvent]): Direct CPU child operations.
         cpu_parent (FunctionEvent): Direct CPU parent operation.
-        input_shapes (Tuple[int, ...]): Shapes of input tensors (requires record_shapes=true).
+        input_shapes (List[List[int]]): Shapes of input tensors (requires record_shapes=True).
+            For plain tensor inputs, each entry is a list of dimensions (e.g. ``[16, 16]``).
+            TensorList inputs are represented as an empty list ``[]``; use
+            ``structured_input_shapes`` to get per-element shapes for TensorList inputs.
         concrete_inputs (List[Any]): Concrete input values (requires record_shapes=true).
         kwinputs (Dict[str, Any]): Keyword arguments (requires record_shapes=true).
         stack (List[str]): Python stack trace where the operation was called (requires with_stack=true).
@@ -590,7 +696,15 @@ class FunctionEvent(FormattedTimesMixin):
         is_legacy (bool): Whether this is from the legacy profiler.
         flops (int): Estimated floating point operations.
         is_user_annotation (bool): Whether this is a user-annotated region.
-        metadata_json (str): Additional metadata in JSON format.
+        metadata_json (str): Deprecated. Use event_metadata instead.
+        event_metadata (EventMetadata): Additional metadata in structured format.
+        structured_input_shapes (List[List[int] | List[List[int]]]): Like ``input_shapes``
+            but distinguishes TensorList inputs.  Plain tensor inputs are ``List[int]``;
+            TensorList inputs are ``List[List[int]]`` containing one shape per tensor in the list.
+            Matches the ``"Input Dims"`` field in the Chrome trace JSON.
+        structured_input_strides (List[List[int] | List[List[int]]]): Strides of input
+            tensors in the same format as ``structured_input_shapes`` (requires
+            record_shapes=True).
 
     Properties:
         cpu_time_total (float): Total CPU time in microseconds.
@@ -637,7 +751,21 @@ def __init__(
         concrete_inputs=None,
         kwinputs=None,
         is_user_annotation=False,
+        is_python_function=False,
+        activity_type=None,
         metadata_json=None,
+        flow_id=None,
+        flow_type=None,
+        flow_start=None,
+        external_id=0,
+        linked_correlation_id=0,
+        extra_meta=None,
+        structured_input_shapes=None,
+        structured_input_strides=None,
+        input_dtypes=None,
+        python_id=-1,
+        python_parent_id=-1,
+        python_module_id=-1,
     ):
         self.id: int = id
         self.node_id: int = node_id
@@ -648,11 +776,11 @@ def __init__(
         self.trace_name: str = trace_name
         self.time_range: Interval = Interval(start_us, end_us)
         self.thread: int = thread
-        self.fwd_thread: Optional[int] = fwd_thread
+        self.fwd_thread: int | None = fwd_thread
         self.kernels: list[Kernel] = []
         self.count: int = 1
         self.cpu_children: list[FunctionEvent] = []
-        self.cpu_parent: Optional[FunctionEvent] = None
+        self.cpu_parent: FunctionEvent | None = None
         # pyrefly: ignore [bad-assignment]
         self.input_shapes: tuple[int, ...] = input_shapes
         # pyrefly: ignore [bad-assignment]
@@ -662,7 +790,7 @@ def __init__(
         # pyrefly: ignore [bad-assignment]
         self.stack: list = stack
         self.scope: int = scope
-        self.use_device: Optional[str] = use_device
+        self.use_device: str | None = use_device
         self.cpu_memory_usage: int = cpu_memory_usage
         self.device_memory_usage: int = device_memory_usage
         self.is_async: bool = is_async
@@ -674,12 +802,31 @@ def __init__(
             thread if device_resource_id is None else device_resource_id
         )
         self.is_legacy: bool = is_legacy
-        self.flops: Optional[int] = flops
-        self.is_user_annotation: Optional[bool] = is_user_annotation
+        self.flops: int | None = flops
+        self.is_user_annotation: bool | None = is_user_annotation
+        self.is_python_function: bool = is_python_function
+        self.activity_type: str | None = activity_type
         self.self_cpu_percent = -1
         self.total_cpu_percent = -1
         self.total_device_percent = -1
-        self.metadata_json = metadata_json
+        self._metadata_json = metadata_json
+        self.flow_id: int | None = flow_id
+        self.flow_type: int | None = flow_type
+        self.flow_start: bool | None = flow_start
+        self.external_id: int = external_id
+        self.linked_correlation_id: int = linked_correlation_id
+        self.event_metadata: EventMetadata | None = (
+            _build_metadata(extra_meta) if extra_meta else None
+        )
+        # pyrefly: ignore [bad-assignment]
+        self.structured_input_shapes: list = structured_input_shapes
+        # pyrefly: ignore [bad-assignment]
+        self.structured_input_strides: list = structured_input_strides
+        # pyrefly: ignore [bad-assignment]
+        self.input_dtypes: list[str] = input_dtypes
+        self.python_id: int = python_id
+        self.python_parent_id: int = python_parent_id
+        self.python_module_id: int = python_module_id
 
     def append_kernel(self, name, device, duration):
         if self.device_type != DeviceType.CPU:
@@ -733,6 +880,14 @@ def self_device_memory_usage(self):
             child.device_memory_usage for child in self.cpu_children
         )
 
+    @property
+    @deprecated(
+        "`metadata_json` is deprecated. Use `event_metadata` instead.",
+        category=FutureWarning,
+    )
+    def metadata_json(self):
+        return self._metadata_json
+
     @property
     @deprecated(
         "`self_cuda_memory_usage` is deprecated. Use `self_device_memory_usage` instead.",
@@ -885,26 +1040,26 @@ class FunctionEventAvg(FormattedTimesMixin):
     """
 
     def __init__(self) -> None:
-        self.key: Optional[str] = None
+        self.key: str | None = None
         self.count: int = 0
         self.node_id: int = 0
         self.is_async: bool = False
         self.is_remote: bool = False
-        self.use_device: Optional[str] = None
+        self.use_device: str | None = None
         self.cpu_time_total: int = 0
         self.device_time_total: int = 0
         self.self_cpu_time_total: int = 0
         self.self_device_time_total: int = 0
-        self.input_shapes: Optional[list[list[int]]] = None
-        self.overload_name: Optional[str] = None
-        self.stack: Optional[list] = None
-        self.scope: Optional[int] = None
+        self.input_shapes: list[list[int]] | None = None
+        self.overload_name: str | None = None
+        self.stack: list | None = None
+        self.scope: int | None = None
         self.cpu_memory_usage: int = 0
         self.device_memory_usage: int = 0
         self.self_cpu_memory_usage: int = 0
         self.self_device_memory_usage: int = 0
-        self.cpu_children: Optional[list[FunctionEvent]] = None
-        self.cpu_parent: Optional[FunctionEvent] = None
+        self.cpu_children: list[FunctionEvent] | None = None
+        self.cpu_parent: FunctionEvent | None = None
         self.device_type: DeviceType = DeviceType.CPU
         self.is_legacy: bool = False
         self.flops: int = 0
@@ -1160,7 +1315,6 @@ def _build_table(
     if append_node_id:
         headers.append("Node ID")
 
-    # Have to use a list because nonlocal is Py3 only...
     SPACING_SIZE = 2
     row_format_lst = [""]
     header_sep_lst = [""]
@@ -1221,7 +1375,6 @@ def auto_scale_flops(flops):
     line_length = line_length_lst[0]
     add_column = None  # type: ignore[assignment]
 
-    # Have to use a list because nonlocal is Py3 only...
     result = []
 
     def append(s):
diff --git a/torch/backends/__init__.py b/torch/backends/__init__.py
index f54a3fd6820c7..f07a4797c64da 100644
--- a/torch/backends/__init__.py
+++ b/torch/backends/__init__.py
@@ -135,5 +135,6 @@ class GenericModule(PropModule):
     nnpack as nnpack,
     openmp as openmp,
     opt_einsum as opt_einsum,
+    python_native as python_native,
     quantized as quantized,
 )
diff --git a/torch/backends/_nnapi/prepare.py b/torch/backends/_nnapi/prepare.py
index 4519507d1cfed..c83eac663e92f 100644
--- a/torch/backends/_nnapi/prepare.py
+++ b/torch/backends/_nnapi/prepare.py
@@ -1,5 +1,6 @@
 # mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
+
 from typing import Optional
 
 import torch
@@ -20,7 +21,7 @@ class NnapiModule(torch.nn.Module):
     """
 
     # _nnapi.Compilation is defined
-    comp: Optional[torch.classes._nnapi.Compilation]  # type: ignore[name-defined]
+    comp: Optional[torch.classes._nnapi.Compilation]  # type: ignore[name-defined]  # noqa: UP045
     weights: list[torch.Tensor]
     out_templates: list[torch.Tensor]
 
diff --git a/torch/backends/_nnapi/serializer.py b/torch/backends/_nnapi/serializer.py
index 6ce2fa2766fd2..0be813c7610d1 100644
--- a/torch/backends/_nnapi/serializer.py
+++ b/torch/backends/_nnapi/serializer.py
@@ -6,7 +6,7 @@
 import operator
 import struct
 import sys
-from typing import NamedTuple, Optional
+from typing import NamedTuple
 
 import torch
 
@@ -1004,8 +1004,8 @@ def add_list_construct(self, node):
             )
         output = node.outputsAt(0)
         ctype = output.type()
-        const_vals: Optional[list] = []
-        tensors: Optional[list] = []
+        const_vals: list | None = []
+        tensors: list | None = []
         for inp in node.inputs():
             if const_vals is not None and inp in self.constants:
                 _, val = self.get_constant_value(inp)
diff --git a/torch/backends/cuda/__init__.py b/torch/backends/cuda/__init__.py
index d62c2b05a1ea1..a7bbcd14ce1a7 100644
--- a/torch/backends/cuda/__init__.py
+++ b/torch/backends/cuda/__init__.py
@@ -1,6 +1,6 @@
 # mypy: allow-untyped-defs
 import contextlib
-from typing import Any, Union
+from typing import Any
 from typing_extensions import deprecated
 
 import torch
@@ -15,6 +15,7 @@
     "preferred_linalg_library",
     "preferred_blas_library",
     "preferred_rocm_fa_library",
+    "is_ck_sdpa_available",
     "cufft_plan_cache",
     "matmul",
     "SDPAParams",
@@ -215,7 +216,7 @@ def __setattr__(self, name, value):
 
 
 def preferred_linalg_library(
-    backend: Union[None, str, torch._C._LinalgBackend] = None,
+    backend: None | str | torch._C._LinalgBackend = None,
 ) -> torch._C._LinalgBackend:
     r"""
     Override the heuristic PyTorch uses to choose between cuSOLVER and MAGMA for CUDA linear algebra operations.
@@ -286,7 +287,7 @@ def preferred_linalg_library(
 
 
 def preferred_blas_library(
-    backend: Union[None, str, torch._C._BlasBackend] = None,
+    backend: None | str | torch._C._BlasBackend = None,
 ) -> torch._C._BlasBackend:
     r"""
     Override the library PyTorch uses for BLAS operations. Choose between cuBLAS, cuBLASLt, and CK [ROCm-only].
@@ -341,7 +342,7 @@ def preferred_blas_library(
 
 
 def preferred_rocm_fa_library(
-    backend: Union[None, str, torch._C._ROCmFABackend] = None,
+    backend: None | str | torch._C._ROCmFABackend = None,
 ) -> torch._C._ROCmFABackend:
     r"""
     [ROCm-only]
@@ -385,6 +386,17 @@ def preferred_rocm_fa_library(
 SDPAParams.__name__ = "SDPAParams"
 
 
+def is_ck_sdpa_available() -> bool:
+    r"""
+    .. warning:: This flag is beta and subject to change.
+
+    Returns whether composable_kernel may be used as the backend for
+    scaled-dot-product-attention.
+    """
+    # pyrefly: ignore [missing-attribute]
+    return torch._C._is_ck_sdpa_available()
+
+
 def flash_sdp_enabled():
     r"""
     .. warning:: This flag is beta and subject to change.
diff --git a/torch/backends/cudnn/__init__.py b/torch/backends/cudnn/__init__.py
index a7c2ef9b75060..e2494aa5ff8ef 100644
--- a/torch/backends/cudnn/__init__.py
+++ b/torch/backends/cudnn/__init__.py
@@ -29,7 +29,7 @@
 #
 # to globally disable CuDNN/MIOpen
 
-__cudnn_version: Optional[int] = None
+__cudnn_version: int | None = None
 
 if _cudnn is not None:
 
diff --git a/torch/backends/cudnn/rnn.py b/torch/backends/cudnn/rnn.py
index a17dfbdef59e8..724e4a08138b1 100644
--- a/torch/backends/cudnn/rnn.py
+++ b/torch/backends/cudnn/rnn.py
@@ -46,8 +46,6 @@ def get(self):
         return self.inner
 
     def __getstate__(self):
-        # Note: can't return {}, because python2 won't call __setstate__
-        # if the value evaluates to False
         return "<unserializable>"
 
     def __setstate__(self, state):
diff --git a/torch/backends/cusparselt/__init__.py b/torch/backends/cusparselt/__init__.py
index 3e9b9df2acf14..6b0006d91637d 100644
--- a/torch/backends/cusparselt/__init__.py
+++ b/torch/backends/cusparselt/__init__.py
@@ -14,8 +14,8 @@
 except ImportError:
     _cusparselt = None  # type: ignore[assignment]
 
-__cusparselt_version: Optional[int] = None
-__MAX_ALG_ID: Optional[int] = None
+__cusparselt_version: int | None = None
+__MAX_ALG_ID: int | None = None
 
 if _cusparselt is not None:
 
@@ -39,7 +39,7 @@ def _init() -> bool:
         return False
 
 
-def version() -> Optional[int]:
+def version() -> int | None:
     """Return the version of cuSPARSELt"""
     if not _init():
         return None
@@ -51,7 +51,7 @@ def is_available() -> bool:
     return torch._C._has_cusparselt
 
 
-def get_max_alg_id() -> Optional[int]:
+def get_max_alg_id() -> int | None:
     if not _init():
         return None
     return __MAX_ALG_ID
diff --git a/torch/backends/mps/__init__.py b/torch/backends/mps/__init__.py
index 5c3c507428cff..9e1f7b811c03d 100644
--- a/torch/backends/mps/__init__.py
+++ b/torch/backends/mps/__init__.py
@@ -60,7 +60,7 @@ def get_core_count() -> int:
     return torch._C._mps_get_core_count()
 
 
-_lib: Optional[_Library] = None
+_lib: _Library | None = None
 
 
 def _init() -> None:
diff --git a/torch/backends/python_native/__init__.py b/torch/backends/python_native/__init__.py
new file mode 100644
index 0000000000000..9e67571e148a1
--- /dev/null
+++ b/torch/backends/python_native/__init__.py
@@ -0,0 +1,387 @@
+"""
+User-facing API for controlling DSL operation overrides.
+
+The torch.backends.python_native module provides control over DSL (Domain Specific Language)
+operation overrides defined in torch._native. This allows users to selectively enable or disable
+high-performance implementations from various DSLs like Triton and CuteDSL.
+
+The module supports both coarse-grained control (entire DSLs) and fine-grained control
+(individual operations or dispatch keys). All control operations support context managers
+for temporary state changes.
+
+Example usage::
+
+    import torch.backends.python_native as pn
+
+    # DSL-level control
+    pn.triton.enabled = False  # Disable all triton ops
+    pn.cutedsl.enabled = True  # Enable all cutedsl ops
+
+    # Individual operation control
+    pn.disable_operations("scaled_mm")  # Disable specific op across all DSLs
+    pn.enable_operations("scaled_mm")  # Re-enable specific op
+
+    # Context manager support
+    with pn.triton.disabled():
+        result = some_computation()  # Triton ops disabled here
+
+    # Query capabilities
+    print(pn.available_dsls)  # ['triton', 'cutedsl']
+    print(pn.get_dsl_operations("triton"))  # Operations for triton
+"""
+
+import functools
+import sys
+import types
+from contextlib import contextmanager
+
+from torch.backends import ContextProp, flags_frozen, PropModule
+
+
+@contextmanager
+def _preserve_filter_state():
+    """Context manager to save and restore registry filter state."""
+    filter_state = _get_filter_state()
+
+    # Save original state
+    original_state = (
+        set(filter_state._dsl_names),
+        set(filter_state._op_symbols),
+        set(filter_state._dispatch_keys),
+    )
+
+    try:
+        yield filter_state
+    finally:
+        # Restore original state
+        filter_state._dsl_names.clear()
+        filter_state._op_symbols.clear()
+        filter_state._dispatch_keys.clear()
+
+        filter_state._dsl_names.update(original_state[0])
+        filter_state._op_symbols.update(original_state[1])
+        filter_state._dispatch_keys.update(original_state[2])
+
+
+def _get_dsl_registry():
+    """Lazy import to avoid circular imports."""
+    from torch._native.dsl_registry import dsl_registry
+
+    return dsl_registry
+
+
+def _get_registry_functions():
+    """Lazy import of registry functions."""
+    from torch._native.registry import (
+        _filter_state,
+        _graphs,
+        deregister_op_overrides,
+        reenable_op_overrides,
+    )
+
+    return deregister_op_overrides, reenable_op_overrides, _graphs, _filter_state
+
+
+def _get_filter_state():
+    """Direct access to filter state."""
+    return _get_registry_functions()[3]
+
+
+def _get_dsl_module(dsl_name: str):
+    """Get the registered DSL module for direct control.
+
+    Uses the DSL registry to dynamically look up DSL modules instead of
+    hard-coding the mapping. This makes the function automatically extensible
+    for new DSLs without code changes.
+
+    Args:
+        dsl_name (str): Name of the DSL to retrieve.
+
+    Returns:
+        DSLModuleProtocol: The registered DSL module.
+
+    Raises:
+        ValueError: If the DSL is not registered.
+    """
+    registry = _get_dsl_registry()
+
+    # Use the public API to get the DSL module
+    dsl_module = registry.get_dsl_module(dsl_name)
+    if dsl_module is not None:
+        return dsl_module
+    else:
+        raise ValueError(
+            f"Unknown DSL: {dsl_name}. Available DSLs: {registry.list_all_dsls()}"
+        )
+
+
+class DSLController:
+    """Controller for a specific DSL."""
+
+    def __init__(self, dsl_name: str):
+        self._dsl_name = dsl_name
+
+    @property
+    def name(self) -> str:
+        return self._dsl_name
+
+    @property
+    def available(self) -> bool:
+        """Check if DSL runtime is available."""
+        registry = _get_dsl_registry()
+        return registry.is_dsl_available(self._dsl_name)
+
+    @property
+    def version(self):
+        """Get DSL version."""
+        registry = _get_dsl_registry()
+        return registry.get_dsl_version(self._dsl_name)
+
+    @property
+    def enabled(self) -> bool:
+        """Check if DSL is currently enabled."""
+        filter_state = _get_filter_state()
+        return self._dsl_name not in filter_state._dsl_names
+
+    @enabled.setter
+    def enabled(self, value: bool):
+        """Enable or disable the DSL."""
+        if flags_frozen():
+            raise RuntimeError(
+                f"not allowed to set {self._dsl_name} DSL flags "
+                "after disable_global_flags; please use flags() context manager instead"
+            )
+        if value:
+            self.enable()
+        else:
+            self.disable()
+
+    def disable(self):
+        """Disable all operations for this DSL."""
+        dsl_module = _get_dsl_module(self._dsl_name)
+        dsl_module.deregister_op_overrides()
+
+    def enable(self):
+        """Re-enable all operations for this DSL."""
+        reenable_op_overrides = _get_registry_functions()[1]
+        reenable_op_overrides(enable_dsl_names=self._dsl_name)
+
+    @contextmanager
+    def disabled(self):
+        """Context manager to temporarily disable DSL."""
+        original_state = self.enabled
+        try:
+            self.disable()
+            yield
+        finally:
+            if original_state:
+                self.enable()
+
+    def __repr__(self):
+        status = "available" if self.available else "unavailable"
+        enabled_status = "enabled" if self.enabled else "disabled"
+        return f"DSLController({self._dsl_name}, {status}, {enabled_status})"
+
+
+class PythonNativeModule(PropModule):
+    """Main module for python_native DSL control."""
+
+    def __init__(self, original_module):
+        super().__init__(original_module, original_module.__name__)
+
+    @property
+    def available_dsls(self) -> list[str]:
+        """Get list of available DSLs."""
+        registry = _get_dsl_registry()
+        result = registry.list_available_dsls()
+        return list(result) if not isinstance(result, list) else result
+
+    @property
+    def all_dsls(self) -> list[str]:
+        """Get list of all registered DSLs."""
+        registry = _get_dsl_registry()
+        result = registry.list_all_dsls()
+        return list(result) if not isinstance(result, list) else result
+
+    def get_dsl_operations(self, dsl_name: str) -> list[str]:
+        """Get list of operations registered by a specific DSL.
+
+        Args:
+            dsl_name (str): Name of the DSL to query (e.g., 'triton', 'cutedsl').
+
+        Returns:
+            list[str]: Sorted list of operation names registered by the DSL.
+
+        Example::
+
+            ops = torch.backends.python_native.get_dsl_operations("triton")
+            print(ops)  # ['triton_to_mxfp8_dim0', ...]
+        """
+        from torch._native.registry import get_dsl_operations
+
+        return get_dsl_operations(dsl_name)
+
+    def disable_operations(self, *op_symbols: str):
+        """Disable specific operations across all DSLs.
+
+        Args:
+            *op_symbols (str): Names of operations to disable.
+
+        Example::
+
+            # Disable scaled matrix multiply across all DSLs
+            torch.backends.python_native.disable_operations("scaled_mm")
+
+            # Disable multiple operations
+            torch.backends.python_native.disable_operations(
+                "scaled_mm", "flash_attention"
+            )
+        """
+        deregister_op_overrides = _get_registry_functions()[0]
+        deregister_op_overrides(disable_op_symbols=list(op_symbols))
+
+    def enable_operations(self, *op_symbols: str):
+        """Re-enable specific operations across all DSLs.
+
+        Args:
+            *op_symbols (str): Names of operations to re-enable.
+
+        Example::
+
+            # Re-enable previously disabled operations
+            torch.backends.python_native.enable_operations(
+                "scaled_mm", "flash_attention"
+            )
+        """
+        reenable_op_overrides = _get_registry_functions()[1]
+        reenable_op_overrides(enable_op_symbols=list(op_symbols))
+
+    def disable_dispatch_keys(self, *dispatch_keys: str):
+        """Disable operations at specific dispatch keys.
+
+        Args:
+            *dispatch_keys (str): Dispatch keys to disable (e.g., 'CUDA', 'CPU').
+
+        Example::
+
+            # Disable all native operations on CUDA
+            torch.backends.python_native.disable_dispatch_keys("CUDA")
+        """
+        deregister_op_overrides = _get_registry_functions()[0]
+        deregister_op_overrides(disable_dispatch_keys=list(dispatch_keys))
+
+    def enable_dispatch_keys(self, *dispatch_keys: str):
+        """Re-enable operations at specific dispatch keys.
+
+        Args:
+            *dispatch_keys (str): Dispatch keys to re-enable (e.g., 'CUDA', 'CPU').
+
+        Example::
+
+            # Re-enable native operations on CUDA
+            torch.backends.python_native.enable_dispatch_keys("CUDA")
+        """
+        reenable_op_overrides = _get_registry_functions()[1]
+        reenable_op_overrides(enable_dispatch_keys=list(dispatch_keys))
+
+    @contextmanager
+    def operations_disabled(self, *op_symbols: str):
+        """Context manager to temporarily disable operations.
+
+        Args:
+            *op_symbols (str): Names of operations to temporarily disable.
+
+        Example::
+
+            with torch.backends.python_native.operations_disabled("scaled_mm"):
+                # scaled_mm is disabled across all DSLs
+                result = model(input)
+            # scaled_mm is automatically re-enabled here
+        """
+        filter_state = _get_filter_state()
+        previously_disabled_ops = {
+            op for op in op_symbols if op in filter_state._op_symbols
+        }
+
+        self.disable_operations(*op_symbols)
+        try:
+            yield
+        finally:
+            # Only re-enable operations that weren't already disabled
+            ops_to_reenable = [
+                op for op in op_symbols if op not in previously_disabled_ops
+            ]
+            if ops_to_reenable:
+                self.enable_operations(*ops_to_reenable)
+
+    @functools.lru_cache(maxsize=16)  # noqa: B019
+    def _get_dsl_controller(self, name: str) -> "DSLController":
+        """Get or create a DSL controller (cached)."""
+        return DSLController(name)
+
+    def _get_registry_functions(self):
+        """Expose registry functions for testing."""
+        return _get_registry_functions()
+
+    def is_operation_disabled(self, op_symbol: str) -> bool:
+        """Check if an operation is currently disabled."""
+        filter_state = _get_filter_state()
+        return op_symbol in filter_state._op_symbols
+
+    def is_dsl_disabled(self, dsl_name: str) -> bool:
+        """Check if a DSL is currently disabled."""
+        filter_state = _get_filter_state()
+        return dsl_name in filter_state._dsl_names
+
+    def __getattr__(self, name: str):
+        """Dynamic attribute access for DSL controllers."""
+        # Skip dunder attributes to avoid triggering DSL registry lookups
+        # during torch initialization. inspect.getmodule() calls
+        # hasattr(module, '__file__') which would otherwise cause a circular
+        # import through _get_dsl_registry() while torch is still loading.
+        if name.startswith("__") and name.endswith("__"):
+            raise AttributeError(f"module '{self.__name__}' has no attribute '{name}'")
+
+        if name in self.all_dsls:
+            return self._get_dsl_controller(name)
+
+        # Expose private functions for testing
+        if name == "_get_dsl_module":
+            return _get_dsl_module
+        if name == "_get_registry_functions":
+            return self._get_registry_functions
+        if name == "_get_filter_state":
+            return _get_filter_state
+
+        raise AttributeError(f"module '{self.__name__}' has no attribute '{name}'")
+
+    def __dir__(self):
+        """Return available attributes including DSL names."""
+        attrs = set(super().__dir__())
+        attrs.update(
+            {
+                "available_dsls",
+                "all_dsls",
+                "get_dsl_operations",
+                "disable_operations",
+                "enable_operations",
+                "disable_dispatch_keys",
+                "enable_dispatch_keys",
+                "operations_disabled",
+                "is_operation_disabled",
+                "is_dsl_disabled",
+            }
+        )
+
+        # Add DSL names
+        try:
+            attrs.update(self.all_dsls)
+        except Exception:
+            # If registry not available yet, skip DSL names
+            pass
+
+        return sorted(attrs)
+
+
+# Replace the current module with our enhanced version
+sys.modules[__name__] = PythonNativeModule(sys.modules[__name__])
diff --git a/torch/backends/xeon/run_cpu.py b/torch/backends/xeon/run_cpu.py
index 0f9aec7152425..cd9502baf4846 100644
--- a/torch/backends/xeon/run_cpu.py
+++ b/torch/backends/xeon/run_cpu.py
@@ -249,8 +249,8 @@ def numa_aware_check(self, core_list):
                 "Numa Aware: cores:%s on different NUMA nodes:%s. To avoid \
 this behavior, please use --ncores-per-instance knob to make sure number of cores is divisible by --ncores-per-\
 instance. Alternatively, please use --skip-cross-node-cores knob.",
-                str(core_list),
-                str(numa_ids),
+                core_list,
+                numa_ids,
             )
         if len(numa_ids) == 0:
             raise RuntimeError(
diff --git a/torch/compiler/__init__.py b/torch/compiler/__init__.py
index 31ff3b6a31ab5..0495b3ac61084 100644
--- a/torch/compiler/__init__.py
+++ b/torch/compiler/__init__.py
@@ -1,18 +1,15 @@
 # mypy: allow-untyped-defs
+import contextlib
 import io
 from collections.abc import Callable
-from dataclasses import dataclass
-from typing import Any, Optional, TYPE_CHECKING, TypeVar, Union
+from typing import Any, TypeVar
 from typing_extensions import ParamSpec
 
 import torch
 from torch._higher_order_ops.invoke_subgraph import NestedCompileRegionOptions
 
 from . import config
-
-
-if TYPE_CHECKING:
-    from ._cache import CacheInfo
+from ._cache import CacheInfo
 
 
 __all__ = [
@@ -265,7 +262,7 @@ def set_stance(
     stance: str = "default",
     *,
     skip_guard_eval_unsafe: bool = False,
-    force_backend: Union[str, Callable[..., Any], None] = None,
+    force_backend: str | Callable[..., Any] | None = None,
 ):
     """
     Set the current stance of the compiler.
@@ -437,6 +434,7 @@ def wrap_numpy(fn):
 
 _is_compiling_flag: bool = False
 _is_exporting_flag: bool = False
+_is_non_strict_tracing_flag: bool = False
 
 
 def is_compiling() -> bool:
@@ -461,6 +459,60 @@ def is_compiling() -> bool:
         return _is_compiling_flag
 
 
+def _is_non_strict_tracing() -> bool:
+    """
+    Indicates whether we are inside a non-strict make_fx-based tracing session.
+    """
+    return _is_non_strict_tracing_flag
+
+
+@contextlib.contextmanager
+def _non_strict_tracing_context():
+    """Context manager that sets the non-strict tracing flag."""
+    global _is_non_strict_tracing_flag
+    old = _is_non_strict_tracing_flag
+    try:
+        _is_non_strict_tracing_flag = True
+        yield
+    finally:
+        _is_non_strict_tracing_flag = old
+
+
+@contextlib.contextmanager
+def _patch_autograd_grad():
+    """Patch autograd.grad for non-strict make_fx tracing.
+
+    This patch installs autograd hooks so traced backward nodes preserve
+    stack trace, seq_nr, and autograd_backward metadata before delegating to
+    the real torch.autograd.grad.
+    """
+    import functools
+
+    import torch.autograd
+    from torch._functorch._aot_autograd.logging_utils import (
+        setup_stacktrace_preservation_hooks_from_tensors,
+    )
+
+    _orig_grad = torch.autograd.grad
+
+    @functools.wraps(_orig_grad)
+    def _patched_grad(outputs, inputs, *args, **kwargs):
+        if not _is_non_strict_tracing():
+            raise AssertionError(
+                "_patch_autograd_grad() must be used under "
+                "_non_strict_tracing_context()"
+            )
+
+        setup_stacktrace_preservation_hooks_from_tensors(outputs)
+        return _orig_grad(outputs, inputs, *args, **kwargs)
+
+    torch.autograd.grad = _patched_grad
+    try:
+        yield
+    finally:
+        torch.autograd.grad = _orig_grad
+
+
 def is_dynamo_compiling() -> bool:
     """
     Indicates whether a graph is traced via TorchDynamo.
@@ -497,7 +549,7 @@ def is_exporting() -> bool:
     return _is_exporting_flag
 
 
-def save_cache_artifacts() -> Optional[tuple[bytes, "CacheInfo"]]:
+def save_cache_artifacts() -> tuple[bytes, CacheInfo] | None:
     """
     Serializes all the cache artifacts that were created during the compilation
 
@@ -516,7 +568,7 @@ def save_cache_artifacts() -> Optional[tuple[bytes, "CacheInfo"]]:
     return CacheArtifactManager.serialize()
 
 
-def load_cache_artifacts(serialized_artifacts: bytes) -> Optional["CacheInfo"]:
+def load_cache_artifacts(serialized_artifacts: bytes) -> CacheInfo | None:
     """
     Hot loads cache artifacts that were previously serialized via
     save_cache_artifacts
@@ -663,7 +715,11 @@ def skip_all_guards_unsafe(guard_entries):
 
 
 def nested_compile_region(
-    fn=None, options: Optional[NestedCompileRegionOptions] = None
+    fn=None,
+    *,
+    options: NestedCompileRegionOptions | None = None,
+    max_reuse_entries: int = 8,
+    reuse_hash_fn=None,
 ):
     """
     Tells **``torch.compile``** that the marked set of operations forms a nested
@@ -694,6 +750,17 @@ def nested_compile_region(
         options: Optional backend to use for compiling the subgraph.
             Warning: this is an experimental feature under development and
             not ready for use yet.
+        max_reuse_entries: Maximum number of reuse cache entries per function
+            before raising an error. If this limit is hit, guards keep failing
+            across invocations and hierarchical compilation is not effective.
+        reuse_hash_fn: Optional callable that takes the same ``*args, **kwargs``
+            as the wrapped function and returns an integer hash key. When
+            provided, Dynamo traces this function to obtain a constant integer
+            and uses it as the cache key for subgraph reuse, bypassing the
+            automatic fingerprint/guard machinery. Two calls that produce the
+            same hash key reuse the same cached subgraph. The hash function
+            must be fully traceable (no graph breaks) and must return a
+            constant integer.
     """
 
     if options is not None:
@@ -708,7 +775,12 @@ def nested_compile_region(
         mark_compile_region as _mark_compile_region,
     )
 
-    return _mark_compile_region(fn, options=options)
+    return _mark_compile_region(
+        fn,
+        options=options,
+        max_reuse_entries=max_reuse_entries,
+        reuse_hash_fn=reuse_hash_fn,
+    )
 
 
 def load_compiled_function(
diff --git a/torch/compiler/_cache.py b/torch/compiler/_cache.py
index f715e21f46b0a..6252148dae33e 100644
--- a/torch/compiler/_cache.py
+++ b/torch/compiler/_cache.py
@@ -6,7 +6,7 @@
 from collections.abc import Generator
 from contextlib import contextmanager
 from itertools import chain
-from typing import Any, Optional
+from typing import Any
 
 from torch.utils._appending_byte_serializer import (
     AppendingByteSerializer,
@@ -242,7 +242,7 @@ def record_artifact(
         artifact = CacheArtifactFactory.encode_create(artifact_type, key, content)
         if artifact in cls._seen_artifacts:
             return
-        log.debug("Recording %s", str(artifact))
+        log.debug("Recording %s", artifact)
         cls._new_cache_artifacts[artifact_type].append(artifact)
         cls._seen_artifacts.add(artifact)
 
@@ -254,7 +254,7 @@ def need_serialize(cls) -> bool:
         return len(cls._new_cache_artifacts) != 0
 
     @classmethod
-    def serialize(cls) -> Optional[tuple[bytes, CacheInfo]]:
+    def serialize(cls) -> tuple[bytes, CacheInfo] | None:
         """
         Converts the "mega" list into portable format
         """
@@ -280,7 +280,7 @@ def serialize(cls) -> Optional[tuple[bytes, CacheInfo]]:
         return None
 
     @staticmethod
-    def deserialize(serialized_artifacts: bytes) -> Optional[CacheArtifactsResult]:
+    def deserialize(serialized_artifacts: bytes) -> CacheArtifactsResult | None:
         """
         Converts the portable format back into CacheArtifacts
         """
diff --git a/torch/compiler/config.py b/torch/compiler/config.py
index e507ddc18052e..3fcd4d4b037c8 100644
--- a/torch/compiler/config.py
+++ b/torch/compiler/config.py
@@ -13,7 +13,6 @@
 """
 
 import sys
-from typing import Optional
 
 from torch.utils._config_module import Config, install_config_module
 
@@ -45,7 +44,7 @@
 # FB-internal note: you do NOT have to specify this explicitly specify this if
 # you run on MAST, we will automatically default this to
 # mast:MAST_JOB_NAME:MAST_JOB_VERSION.
-job_id: Optional[str] = Config(
+job_id: str | None = Config(
     env_name_default=["TORCH_COMPILE_JOB_ID", "TORCH_COMPILE_STICKY_PGO_KEY"],
     default=None,
 )
@@ -75,10 +74,10 @@
 consistent profiles across all ranks.
 """
 
-pgo_extra_read_key: Optional[str] = Config(
+pgo_extra_read_key: str | None = Config(
     env_name_default="TORCH_COMPILE_STICKY_PGO_READ", default=None
 )
-pgo_extra_write_key: Optional[str] = Config(
+pgo_extra_write_key: str | None = Config(
     env_name_default="TORCH_COMPILE_STICKY_PGO_WRITE", default=None
 )
 """
@@ -214,7 +213,7 @@
 output shapes during tracing.
 """
 
-log_file_name: Optional[str] = Config(alias="torch._dynamo.config.log_file_name")
+log_file_name: str | None = Config(alias="torch._dynamo.config.log_file_name")
 """
 Specifies a file path for TorchDynamo-specific logging output. When set, internal
 TorchDynamo debug information is written to this file rather than stdout. This is
diff --git a/torch/csrc/DeviceAccelerator.cpp b/torch/csrc/DeviceAccelerator.cpp
index c6ffa893d95ae..c3caf43e805b4 100644
--- a/torch/csrc/DeviceAccelerator.cpp
+++ b/torch/csrc/DeviceAccelerator.cpp
@@ -1,5 +1,6 @@
 #include <c10/core/AllocatorConfig.h>
 #include <torch/csrc/DeviceAccelerator.h>
+#include <torch/csrc/Exceptions.h>
 #include <torch/csrc/utils/device_lazy_init.h>
 
 namespace torch::accelerator {
@@ -100,6 +101,15 @@ void initModule(PyObject* module) {
 
   m.def("_accelerator_emptyCache", []() { at::accelerator::emptyCache(); });
 
+  m.def("_accelerator_emptyHostCache", []() {
+    const auto device_type = at::accelerator::getAccelerator(true).value();
+    if (torch::utils::is_device_lazy_init_supported(device_type) &&
+        !torch::utils::is_device_initialized(device_type)) {
+      return;
+    }
+    at::accelerator::emptyHostCache();
+  });
+
   m.def("_accelerator_getDeviceStats", [](c10::DeviceIndex device_index) {
     using c10::CachingAllocator::Stat;
     using c10::CachingAllocator::StatArray;
@@ -164,9 +174,72 @@ void initModule(PyObject* module) {
     return at::accelerator::getMemoryInfo(device_index);
   });
 
+  m.def("_accelerator_getAllocatorSettings", []() {
+    return c10::CachingAllocator::getAllocatorSettings();
+  });
+
   m.def("_accelerator_setAllocatorSettings", [](std::string env) {
     c10::CachingAllocator::setAllocatorSettings(env);
   });
+
+  // Accelerator Graph class binding
+  py::class_<at::accelerator::Graph, std::shared_ptr<at::accelerator::Graph>>(
+      m, "_acceleratorGraph")
+      .def(py::init<bool>(), py::arg("keep_graph") = false)
+      .def(
+          "capture_begin",
+          [](at::accelerator::Graph& self,
+             std::optional<c10::MempoolId_t> pool_opt,
+             const std::string& capture_error_mode) {
+            c10::MempoolId_t pool = pool_opt.has_value()
+                ? pool_opt.value()
+                : c10::MempoolId_t{0, 0};
+            at::GraphCaptureMode capture_mode = at::GraphCaptureMode::Default;
+            if (capture_error_mode == "default") {
+              capture_mode = at::GraphCaptureMode::Default;
+            } else if (capture_error_mode == "global") {
+              capture_mode = at::GraphCaptureMode::Global;
+            } else if (capture_error_mode == "thread_local") {
+              capture_mode = at::GraphCaptureMode::ThreadLocal;
+            } else if (capture_error_mode == "relaxed") {
+              capture_mode = at::GraphCaptureMode::Relaxed;
+            } else {
+              TORCH_CHECK(
+                  false,
+                  "Unknown capture error mode. Expected `default`, `global`, `thread_local`, or `relaxed`, got ",
+                  capture_error_mode);
+            }
+            return self.capture_begin(pool, capture_mode);
+          },
+          py::arg("pool") = std::nullopt,
+          py::arg("capture_error_mode") = "default",
+          py::call_guard<py::gil_scoped_release>())
+      .def(
+          "capture_end",
+          torch::wrap_pybind_function_no_gil(
+              &at::accelerator::Graph::capture_end))
+      .def(
+          "instantiate",
+          torch::wrap_pybind_function_no_gil(
+              &at::accelerator::Graph::instantiate))
+      .def(
+          "replay",
+          torch::wrap_pybind_function_no_gil(&at::accelerator::Graph::replay))
+      .def(
+          "reset",
+          torch::wrap_pybind_function_no_gil(&at::accelerator::Graph::reset))
+      .def(
+          "pool",
+          torch::wrap_pybind_function_no_gil(&at::accelerator::Graph::pool))
+      .def(
+          "enable_debug_mode",
+          torch::wrap_pybind_function_no_gil(
+              &::at::accelerator::Graph::enable_debug_mode))
+      .def(
+          "debug_dump",
+          torch::wrap_pybind_function_no_gil(
+              &::at::accelerator::Graph::debug_dump),
+          py::arg("path"));
 }
 
 } // namespace torch::accelerator
diff --git a/torch/csrc/Dtype.cpp b/torch/csrc/Dtype.cpp
index bff17ca0cbc79..5f1c064586e33 100644
--- a/torch/csrc/Dtype.cpp
+++ b/torch/csrc/Dtype.cpp
@@ -9,7 +9,6 @@
 #include <cstring>
 
 PyObject* THPDtype_New(at::ScalarType scalar_type, const std::string& name) {
-  HANDLE_TH_ERRORS
   AT_ASSERT(name.length() < DTYPE_NAME_LEN);
   auto type = &THPDtypeType;
   auto self = THPObjectPtr{type->tp_alloc(type, 0)};
@@ -19,7 +18,6 @@ PyObject* THPDtype_New(at::ScalarType scalar_type, const std::string& name) {
   self_->scalar_type = scalar_type;
   std::strncpy(self_->name, name.c_str(), DTYPE_NAME_LEN);
   return self.release();
-  END_HANDLE_TH_ERRORS
 }
 
 static PyObject* THPDtype_is_floating_point(THPDtype* self, PyObject* noargs) {
diff --git a/torch/csrc/DynamicTypes.cpp b/torch/csrc/DynamicTypes.cpp
index 9db1903eec33a..d5d211d5b8fc5 100644
--- a/torch/csrc/DynamicTypes.cpp
+++ b/torch/csrc/DynamicTypes.cpp
@@ -113,7 +113,7 @@ std::tuple<at::Storage, at::ScalarType, bool> createStorageGetType(
       "'");
 
   auto storage = THPStorage_Unpack(untyped_storage_obj);
-  return std::make_tuple(storage, scalar_type, is_typed_storage);
+  return std::make_tuple(std::move(storage), scalar_type, is_typed_storage);
 }
 
 at::Storage createStorage(PyObject* obj) {
diff --git a/torch/csrc/Exceptions.h b/torch/csrc/Exceptions.h
index a9c7f1bdca93e..21d52f79360c9 100644
--- a/torch/csrc/Exceptions.h
+++ b/torch/csrc/Exceptions.h
@@ -1,3 +1,4 @@
+// @allow-raw-throw
 #pragma once
 
 #include <exception>
@@ -6,7 +7,6 @@
 #include <system_error>
 
 #include <ATen/detail/FunctionTraits.h>
-#include <c10/util/C++17.h>
 #include <c10/util/Exception.h>
 #include <c10/util/StringUtil.h>
 #include <pybind11/pybind11.h>
diff --git a/torch/csrc/Module.cpp b/torch/csrc/Module.cpp
index 56bce2ba966ed..f3aee904e530c 100644
--- a/torch/csrc/Module.cpp
+++ b/torch/csrc/Module.cpp
@@ -2,6 +2,7 @@
 #include <fmt/core.h>
 #include <sys/types.h>
 #include <torch/csrc/python_headers.h>
+#include <torch/csrc/utils/pythoncapi_compat.h>
 #include <csignal>
 #include <optional>
 
@@ -19,7 +20,6 @@
 
 #include <ATen/Parallel.h>
 #include <ATen/Utils.h>
-#include <ATen/core/Vitals.h>
 #include <ATen/dlpack.h>
 #include <ATen/native/ConvUtils.h>
 #include <ATen/native/ForeachUtils.h>
@@ -1779,7 +1779,56 @@ static PyObject* THCPModule_ensureCUDADeviceGuardSet(
   END_HANDLE_TH_ERRORS
 }
 
+struct TorchModuleState {
+  PyObject* log_api_usage_seen; // dict used by _log_api_usage_once
+};
+
+static int torchmodule_traverse(PyObject* mod, visitproc visit, void* arg) {
+  auto* state = static_cast<TorchModuleState*>(PyModule_GetState(mod));
+  Py_VISIT(state->log_api_usage_seen);
+  return 0;
+}
+
+static int torchmodule_clear(PyObject* mod) {
+  auto* state = static_cast<TorchModuleState*>(PyModule_GetState(mod));
+  Py_CLEAR(state->log_api_usage_seen);
+  return 0;
+}
+
+static void torchmodule_free(void* mod) {
+  torchmodule_clear((PyObject*)mod);
+}
+
+// Thread-safe in free-threaded Python.
+static PyObject* LogAPIUsageOnceFromPython(PyObject* self, PyObject* event) {
+  auto* state = static_cast<TorchModuleState*>(PyModule_GetState(self));
+  PyObject* api_usage_seen = state->log_api_usage_seen;
+
+  int found = PyDict_Contains(api_usage_seen, event);
+  if (found < 0) {
+    return nullptr;
+  } else if (found != 0) {
+    Py_RETURN_NONE;
+  }
+
+  const char* event_str = PyUnicode_AsUTF8(event);
+  if (!event_str) {
+    return nullptr;
+  }
+
+  // Returns 0 if we inserted, 1 if already present, -1 on error.
+  int rc = PyDict_SetDefaultRef(api_usage_seen, event, Py_None, nullptr);
+  if (rc < 0) {
+    return nullptr;
+  }
+  if (rc == 0) {
+    c10::LogAPIUsage(event_str);
+  }
+  Py_RETURN_NONE;
+}
+
 static std::initializer_list<PyMethodDef> TorchMethods = {
+    {"_log_api_usage_once", LogAPIUsageOnceFromPython, METH_O, nullptr},
     {"_initExtension", THPModule_initExtension, METH_O, nullptr},
     {"_autograd_init", THPAutograd_initExtension, METH_NOARGS, nullptr},
     {"_add_docstr", THPModule_addDocStr, METH_VARARGS, nullptr},
@@ -2121,16 +2170,6 @@ void initModule(PyObject* module);
 
 static std::vector<PyMethodDef> methods;
 
-// In Python we can't use the trick of C10_LOG_API_USAGE_ONCE
-// Guaranteed to be invoked from Python under GIL, no locking on map needed
-static void LogAPIUsageOnceFromPython(const std::string& event) {
-  static std::unordered_set<std::string> seen;
-  if (!seen.count(event)) {
-    seen.insert(event);
-    c10::LogAPIUsage(event);
-  }
-}
-
 static void LogAPIUsageMetadataFromPython(
     const std::string& event,
     const std::map<std::string, std::string>& metadata_map) {
@@ -2249,9 +2288,22 @@ PyObject* initModule() {
 #endif
 
   static struct PyModuleDef torchmodule = {
-      PyModuleDef_HEAD_INIT, "torch._C", nullptr, -1, methods.data()};
+      PyModuleDef_HEAD_INIT,
+      "torch._C",
+      nullptr,
+      sizeof(TorchModuleState),
+      methods.data(),
+      nullptr, // m_slots
+      torchmodule_traverse,
+      torchmodule_clear,
+      torchmodule_free};
   module = PyModule_Create(&torchmodule);
   ASSERT_TRUE(module);
+
+  auto* mod_state = static_cast<TorchModuleState*>(PyModule_GetState(module));
+  mod_state->log_api_usage_seen = PyDict_New();
+  ASSERT_TRUE(mod_state->log_api_usage_seen);
+
 #ifdef Py_GIL_DISABLED
   PyUnstable_Module_SetGIL(module, Py_MOD_GIL_NOT_USED);
 #endif
@@ -2367,7 +2419,7 @@ PyObject* initModule() {
 #endif
   ASSERT_TRUE(set_module_attr("_has_cudnn", has_cudnn));
 
-#if defined(USE_CUSPARSELT)
+#if defined(USE_CUSPARSELT) || defined(USE_ROCM)
   PyObject* has_cusparselt = Py_True;
 #else
   PyObject* has_cusparselt = Py_False;
@@ -2399,20 +2451,8 @@ PyObject* initModule() {
   auto py_module = py::reinterpret_borrow<py::module>(module);
   py_module.def("_initCrashHandler", &_initCrashHandler);
   py_module.def("_demangle", &c10::demangle);
-  py_module.def("_log_api_usage_once", &LogAPIUsageOnceFromPython);
   py_module.def("_log_api_usage_metadata", &LogAPIUsageMetadataFromPython);
 
-  py_module.def("vitals_enabled", &at::vitals::torchVitalEnabled);
-  py_module.def(
-      "set_vital",
-      [](const std::string& vital,
-         const std::string& attr,
-         const std::string& value) {
-        return at::vitals::VitalsAPI.setVital(vital, attr, value);
-      });
-  py_module.def(
-      "read_vitals", []() { return at::vitals::VitalsAPI.readVitals(); });
-
   py_module.def(
       "init_num_threads",
       torch::wrap_pybind_function(at::init_num_threads),
@@ -2739,6 +2779,14 @@ Call this whenever a new thread is created in order to propagate values from
     return at::globalContext().getROCmFAPreferredBackend();
   });
 
+  py_module.def("_is_ck_sdpa_available", []() {
+#ifdef USE_ROCM
+    return at::globalContext().ckSupported() && at::globalContext().hasCKSDPA();
+#else
+    return false;
+#endif
+  });
+
   py_module.def(
       "_set_sm_carveout_experimental", [](std::optional<int32_t> val) {
         at::globalContext()._setSMCarveout_EXPERIMENTAL(val);
@@ -2778,6 +2826,7 @@ Call this whenever a new thread is created in order to propagate values from
 
   py_module.def(
       "_stash_obj_in_tls", [](const std::string& key, py::handle arg) {
+        Py_INCREF(arg.ptr());
         at::impl::ThreadLocalPythonObjects::get_state().set(
             key,
             std::make_shared<c10::SafePyObject>(arg.ptr(), getPyInterpreter()));
diff --git a/torch/csrc/Storage.cpp b/torch/csrc/Storage.cpp
index 58e37d2c236b0..12ab48bf7da07 100644
--- a/torch/csrc/Storage.cpp
+++ b/torch/csrc/Storage.cpp
@@ -254,7 +254,6 @@ static PyObject* THPStorage_pynew(
           "but one of the items was of type ",
           THPUtils_typename(item.get()),
           " instead of int");
-      return nullptr;
     }
   }
   return self;
@@ -302,7 +301,6 @@ static PyObject* THPStorage_get(THPStorage* self, PyObject* index) {
           step,
           ", but only a step of "
           "1 is supported");
-      return nullptr;
     }
 
     const auto& storage = THPStorage_Unpack(self);
@@ -351,7 +349,6 @@ static int THPStorage_set(THPStorage* self, PyObject* index, PyObject* value) {
         "can only set storage content with a int types, but got ",
         THPUtils_typename(value),
         " instead");
-    return -1;
   }
 
   uint8_t rvalue = THPByteUtils_unpackReal(value);
@@ -374,7 +371,6 @@ static int THPStorage_set(THPStorage* self, PyObject* index, PyObject* value) {
           step,
           ", but only a step of "
           "1 is supported");
-      return 0;
     }
     // TODO: check the bounds only once
     // TODO: fill?
@@ -384,7 +380,6 @@ static int THPStorage_set(THPStorage* self, PyObject* index, PyObject* value) {
   }
   TORCH_CHECK(
       false, "can't index a " THPStorageStr " with ", THPUtils_typename(index));
-  return -1;
   END_HANDLE_TH_ERRORS_RET(-1)
 }
 
diff --git a/torch/csrc/StorageMethods.cpp b/torch/csrc/StorageMethods.cpp
index 178f735802fb7..edabf0efabf2b 100644
--- a/torch/csrc/StorageMethods.cpp
+++ b/torch/csrc/StorageMethods.cpp
@@ -519,6 +519,35 @@ static PyObject* THPStorage_setFromFile(PyObject* self, PyObject* args) {
   END_HANDLE_TH_ERRORS
 }
 
+static PyObject* THPStorage_swapDataPtr(PyObject* self, PyObject* other) {
+  HANDLE_TH_ERRORS
+  THPStorage_assertNotNull(self);
+  TORCH_CHECK(
+      THPStorage_Check(other),
+      "_swap_data_ptr_ expects an UntypedStorage, but got ",
+      THPUtils_typename(other));
+  THPStorage_assertNotNull(other);
+  auto& self_storage = THPStorage_Unpack(self);
+  auto& other_storage = THPStorage_Unpack(other);
+  TORCH_CHECK(
+      self_storage.device() == other_storage.device(),
+      "_swap_data_ptr_: storages must be on the same device, got ",
+      self_storage.device(),
+      " and ",
+      other_storage.device());
+  size_t self_nbytes = self_storage.nbytes();
+  size_t other_nbytes = other_storage.nbytes();
+  TORCH_CHECK(
+      self_nbytes == other_nbytes || self_nbytes == 0 || other_nbytes == 0,
+      "_swap_data_ptr_: storages must have the same nbytes or one must have 0, got ",
+      self_nbytes,
+      " and ",
+      other_nbytes);
+  self_storage.swap_data_ptr(const_cast<c10::Storage&>(other_storage));
+  Py_RETURN_NONE;
+  END_HANDLE_TH_ERRORS
+}
+
 static PyObject* THPStorage__setCdata(PyObject* _self, PyObject* new_cdata) {
   HANDLE_TH_ERRORS
   auto self = reinterpret_cast<THPStorage*>(_self);
@@ -630,6 +659,7 @@ static PyMethodDef THPStorage_methods[] = {
      castPyCFunctionWithKeywords(THPStorage_fromFile),
      METH_VARARGS | METH_KEYWORDS | METH_STATIC,
      nullptr},
+    {"_swap_data_ptr_", THPStorage_swapDataPtr, METH_O, nullptr},
     {"_set_cdata", THPStorage__setCdata, METH_O, nullptr},
     {"_byteswap", THPStorage_byteswap, METH_VARARGS, nullptr},
     {"_fix_weakref", THPStorage_fix_weakref, METH_NOARGS, nullptr},
diff --git a/torch/csrc/Stream.cpp b/torch/csrc/Stream.cpp
index fc1ca916cfbed..c9f824a1d60f1 100644
--- a/torch/csrc/Stream.cpp
+++ b/torch/csrc/Stream.cpp
@@ -119,6 +119,7 @@ PyObject* THPStream_Wrap(const c10::Stream& stream) {
 
 static void THPStream_dealloc(THPStream* self) {
   PyObject_ClearWeakRefs((PyObject*)self);
+  Py_CLEAR(self->context);
   Py_TYPE(self)->tp_free(reinterpret_cast<PyObject*>(self));
 }
 
@@ -168,6 +169,19 @@ static PyObject* THPStream_synchronize(PyObject* _self, PyObject* noargs) {
   END_HANDLE_TH_ERRORS
 }
 
+static PyObject* THPStream_is_capturing(PyObject* _self, PyObject* noargs) {
+  HANDLE_TH_ERRORS
+  auto self = reinterpret_cast<THPStream*>(_self);
+
+  return PyBool_FromLong(c10::Stream::unpack3(
+                             self->stream_id,
+                             static_cast<c10::DeviceIndex>(self->device_index),
+                             static_cast<c10::DeviceType>(self->device_type))
+                             .is_capturing());
+
+  END_HANDLE_TH_ERRORS
+}
+
 static PyObject* THPStream_wait_event(PyObject* _self, PyObject* _event) {
   HANDLE_TH_ERRORS {
     auto self = reinterpret_cast<THPStream*>(_self);
@@ -277,38 +291,71 @@ static PyObject* THPStream_enter(PyObject* _self, PyObject* unused) {
   auto self = reinterpret_cast<THPStream*>(_self);
   c10::DeviceType stream_device_type =
       static_cast<c10::DeviceType>(self->device_type);
+
   // No operation is performed if the stream does not belong to an accelerator.
   if (C10_UNLIKELY(!at::accelerator::isAccelerator(stream_device_type))) {
     Py_INCREF(_self);
     return _self;
   }
+
+  // Note [Reentrant Stream Context Manager]
+  //
+  // We maintain a stack of context entries to support nested/reentrant
+  // stream context managers. Each entry records the previously active
+  // stream and device so that they can be restored in __exit__.
+  //
+  // The stack is stored as a Python list where each entry is either:
+  //   - Py_None: no-op enter (stream was already current);
+  //   - dict:    {_ctx_stream, _ctx_device_index} saved before switching.
+  //
+  // self->context is initialized lazily as a PyList on first __enter__.
+  if (!self->context) {
+    auto list = THPObjectPtr(PyList_New(0));
+    if (!list) {
+      throw python_error();
+    }
+    self->context = list.release();
+  }
+
   c10::DeviceIndex cur_device_idx = at::accelerator::getDeviceIndex();
   c10::DeviceIndex stream_device_idx =
       static_cast<c10::DeviceIndex>(self->device_index);
+  c10::Stream cur_stream = at::accelerator::getCurrentStream(stream_device_idx);
+
+  // If the stream is already current, push None as a no-op sentinel.
+  if (cur_stream.id() == self->stream_id &&
+      cur_stream.device_index() == stream_device_idx) {
+    if (PyList_Append(self->context, Py_None) < 0) {
+      throw python_error();
+    }
+    Py_INCREF(_self);
+    return _self;
+  }
+
   // If the stream is not on the current device, switch the current device to
   // the device of the stream.
   if (stream_device_idx != cur_device_idx) {
     at::accelerator::setDeviceIndex(stream_device_idx);
   }
-  c10::Stream cur_stream = at::accelerator::getCurrentStream(stream_device_idx);
   at::accelerator::setCurrentStream(c10::Stream::unpack3(
       self->stream_id, stream_device_idx, stream_device_type));
-  // Save the current device index and previous stream to the context.
+
+  // Save the current device index and previous stream as a dict on the stack.
   auto ctx_device_index =
       THPObjectPtr(THPUtils_packDeviceIndex(cur_device_idx));
   auto ctx_stream = THPObjectPtr(THPStream_Wrap(cur_stream));
-  TORCH_CHECK(!(self->context), "Stream's context should not be initialized.");
   auto dict = THPObjectPtr(PyDict_New());
   if (!dict) {
     throw python_error();
   }
-  self->context = dict.release();
   if (PyDict_SetItemString(
-          self->context, "_ctx_device_index", ctx_device_index.get()) < 0) {
+          dict.get(), "_ctx_device_index", ctx_device_index.get()) < 0) {
+    throw python_error();
+  }
+  if (PyDict_SetItemString(dict.get(), "_ctx_stream", ctx_stream.get()) < 0) {
     throw python_error();
   }
-  if (PyDict_SetItemString(self->context, "_ctx_stream", ctx_stream.get()) <
-      0) {
+  if (PyList_Append(self->context, dict.get()) < 0) {
     throw python_error();
   }
   Py_INCREF(_self);
@@ -319,19 +366,34 @@ static PyObject* THPStream_enter(PyObject* _self, PyObject* unused) {
 static PyObject* THPStream_exit(PyObject* _self, PyObject* unused) {
   HANDLE_TH_ERRORS
   auto self = reinterpret_cast<THPStream*>(_self);
+
   // No operation is performed if the stream does not belong to an accelerator.
   if (C10_UNLIKELY(!at::accelerator::isAccelerator(
           static_cast<c10::DeviceType>(self->device_type)))) {
     Py_RETURN_NONE;
   }
+
+  // Pop the top entry from the stack.
+  Py_ssize_t stack_size = PyList_Size(self->context);
+  TORCH_INTERNAL_ASSERT(stack_size > 0, "Stream context stack is empty.");
+  PyObject* top = PyList_GET_ITEM(self->context, stack_size - 1);
+
+  // Sentinel: this __enter__ was a no-op, nothing to restore.
+  if (top == Py_None) {
+    if (PyList_SetSlice(self->context, stack_size - 1, stack_size, nullptr) <
+        0) {
+      throw python_error();
+    }
+    Py_RETURN_NONE;
+  }
+
   PyObject* py_stream = nullptr;
-  if (PyDict_GetItemStringRef(self->context, "_ctx_stream", &py_stream) < 0) {
+  if (PyDict_GetItemStringRef(top, "_ctx_stream", &py_stream) < 0) {
     throw python_error();
   }
   auto ctx_stream = THPObjectPtr(py_stream);
   PyObject* py_device_index = nullptr;
-  if (PyDict_GetItemStringRef(
-          self->context, "_ctx_device_index", &py_device_index) < 0) {
+  if (PyDict_GetItemStringRef(top, "_ctx_device_index", &py_device_index) < 0) {
     throw python_error();
   }
   auto ctx_device_index = THPObjectPtr(py_device_index);
@@ -342,6 +404,7 @@ static PyObject* THPStream_exit(PyObject* _self, PyObject* unused) {
       ctx_device_index.get(),
       "ctx_device_index should be present on the context dict.");
   auto prev_device_index = THPUtils_unpackDeviceIndex(ctx_device_index.get());
+
   at::accelerator::setCurrentStream(c10::Stream::unpack3(
       prev_stream->stream_id,
       static_cast<c10::DeviceIndex>(prev_stream->device_index),
@@ -350,7 +413,9 @@ static PyObject* THPStream_exit(PyObject* _self, PyObject* unused) {
   if (static_cast<c10::DeviceIndex>(self->device_index) != prev_device_index) {
     at::accelerator::setDeviceIndex(prev_device_index);
   }
-  Py_CLEAR(self->context);
+  if (PyList_SetSlice(self->context, stack_size - 1, stack_size, nullptr) < 0) {
+    throw python_error();
+  }
   Py_RETURN_NONE;
   END_HANDLE_TH_ERRORS
 }
@@ -426,6 +491,7 @@ static const std::initializer_list<PyGetSetDef> THPStream_properties = {
 static const std::initializer_list<PyMethodDef> THPStream_methods = {
     {"query", THPStream_query, METH_NOARGS, nullptr},
     {"synchronize", THPStream_synchronize, METH_NOARGS, nullptr},
+    {"is_capturing", THPStream_is_capturing, METH_NOARGS, nullptr},
     {"wait_event", THPStream_wait_event, METH_O, nullptr},
     {"wait_stream", THPStream_wait_stream, METH_O, nullptr},
     {"record_event",
diff --git a/torch/csrc/api/include/torch/nn/functional/activation.h b/torch/csrc/api/include/torch/nn/functional/activation.h
index 49de1c8af63f3..350ec4f20d2b1 100644
--- a/torch/csrc/api/include/torch/nn/functional/activation.h
+++ b/torch/csrc/api/include/torch/nn/functional/activation.h
@@ -918,9 +918,10 @@ inline std::tuple<Tensor, Tensor> multi_head_attention_forward(
       // average attention weights over heads
       attn_output_weights = attn_output_weights.sum(/*dim=*/1) / num_heads;
     }
-    return std::make_tuple(attn_output, attn_output_weights);
+    return std::make_tuple(
+        std::move(attn_output), std::move(attn_output_weights));
   } else {
-    return std::make_tuple(attn_output, Tensor());
+    return std::make_tuple(std::move(attn_output), Tensor());
   }
 }
 } // namespace detail
diff --git a/torch/csrc/api/include/torch/nn/module.h b/torch/csrc/api/include/torch/nn/module.h
index 4eff40199ff43..66ec084f1f429 100644
--- a/torch/csrc/api/include/torch/nn/module.h
+++ b/torch/csrc/api/include/torch/nn/module.h
@@ -122,10 +122,12 @@ class TORCH_API Module : public std::enable_shared_from_this<Module> {
   ///
   /// \rst
   /// .. code-block:: cpp
-  ///   MyModule module;
-  ///   module->apply([](nn::Module& module) {
-  ///     std::cout << module.name() << std::endl;
-  ///   });
+  ///
+  ///    MyModule module;
+  ///    module->apply([](nn::Module& module) {
+  ///      std::cout << module.name() << std::endl;
+  ///    });
+  ///
   /// \endrst
   void apply(const ModuleApplyFunction& function);
 
@@ -134,10 +136,12 @@ class TORCH_API Module : public std::enable_shared_from_this<Module> {
   ///
   /// \rst
   /// .. code-block:: cpp
-  ///   MyModule module;
-  ///   module->apply([](const nn::Module& module) {
-  ///     std::cout << module.name() << std::endl;
-  ///   });
+  ///
+  ///    MyModule module;
+  ///    module->apply([](const nn::Module& module) {
+  ///      std::cout << module.name() << std::endl;
+  ///    });
+  ///
   /// \endrst
   void apply(const ConstModuleApplyFunction& function) const;
 
@@ -149,10 +153,12 @@ class TORCH_API Module : public std::enable_shared_from_this<Module> {
   ///
   /// \rst
   /// .. code-block:: cpp
-  ///   MyModule module;
-  ///   module->apply([](const std::string& key, nn::Module& module) {
-  ///     std::cout << key << ": " << module.name() << std::endl;
-  ///   });
+  ///
+  ///    MyModule module;
+  ///    module->apply([](const std::string& key, nn::Module& module) {
+  ///      std::cout << key << ": " << module.name() << std::endl;
+  ///    });
+  ///
   /// \endrst
   void apply(
       const NamedModuleApplyFunction& function,
@@ -166,10 +172,12 @@ class TORCH_API Module : public std::enable_shared_from_this<Module> {
   ///
   /// \rst
   /// .. code-block:: cpp
-  ///   MyModule module;
-  ///   module->apply([](const std::string& key, const nn::Module& module) {
-  ///     std::cout << key << ": " << module.name() << std::endl;
-  ///   });
+  ///
+  ///    MyModule module;
+  ///    module->apply([](const std::string& key, const nn::Module& module) {
+  ///      std::cout << key << ": " << module.name() << std::endl;
+  ///    });
+  ///
   /// \endrst
   void apply(
       const ConstNamedModuleApplyFunction& function,
@@ -180,10 +188,12 @@ class TORCH_API Module : public std::enable_shared_from_this<Module> {
   ///
   /// \rst
   /// .. code-block:: cpp
-  ///   MyModule module;
-  ///   module->apply([](const std::shared_ptr<nn::Module>& module) {
-  ///     std::cout << module->name() << std::endl;
-  ///   });
+  ///
+  ///    MyModule module;
+  ///    module->apply([](const std::shared_ptr<nn::Module>& module) {
+  ///      std::cout << module->name() << std::endl;
+  ///    });
+  ///
   /// \endrst
   void apply(const ModulePointerApplyFunction& function) const;
 
@@ -196,11 +206,13 @@ class TORCH_API Module : public std::enable_shared_from_this<Module> {
   ///
   /// \rst
   /// .. code-block:: cpp
-  ///   MyModule module;
-  ///   module->apply([](const std::string& key,
-  ///                    const std::shared_ptr<nn::Module>& module) {
-  ///     std::cout << key << ": " << module->name() << std::endl;
-  ///   });
+  ///
+  ///    MyModule module;
+  ///    module->apply([](const std::string& key,
+  ///                     const std::shared_ptr<nn::Module>& module) {
+  ///      std::cout << key << ": " << module->name() << std::endl;
+  ///    });
+  ///
   /// \endrst
   void apply(
       const NamedModulePointerApplyFunction& function,
@@ -310,18 +322,20 @@ class TORCH_API Module : public std::enable_shared_from_this<Module> {
   /// Attempts to cast this `Module` to the given `ModuleType`.
   ///
   /// This method is useful when calling `apply()`.
+  ///
   /// \rst
   /// .. code-block:: cpp
   ///
-  ///   void initialize_weights(nn::Module& module) {
-  ///     torch::NoGradGuard no_grad;
-  ///     if (auto* linear = module.as<nn::Linear>()) {
-  ///       linear->weight.normal_(0.0, 0.02);
-  ///     }
-  ///   }
+  ///    void initialize_weights(nn::Module& module) {
+  ///      torch::NoGradGuard no_grad;
+  ///      if (auto* linear = module.as<nn::Linear>()) {
+  ///        linear->weight.normal_(0.0, 0.02);
+  ///      }
+  ///    }
+  ///
+  ///    MyModule module;
+  ///    module->apply(initialize_weights);
   ///
-  ///   MyModule module;
-  ///   module->apply(initialize_weights);
   /// \endrst
   template <typename ModuleType>
   typename ModuleType::ContainedType* as() noexcept;
@@ -329,17 +343,20 @@ class TORCH_API Module : public std::enable_shared_from_this<Module> {
   /// Attempts to cast this `Module` to the given `ModuleType`.
   ///
   /// This method is useful when calling `apply()`.
+  ///
   /// \rst
   /// .. code-block:: cpp
-  ///   void initialize_weights(nn::Module& module) {
-  ///     torch::NoGradGuard no_grad;
-  ///     if (auto* linear = module.as<nn::Linear>()) {
-  ///       linear->weight.normal_(0.0, 0.02);
-  ///     }
-  ///   }
   ///
-  ///   MyModule module;
-  ///   module->apply(initialize_weights);
+  ///    void initialize_weights(nn::Module& module) {
+  ///      torch::NoGradGuard no_grad;
+  ///      if (auto* linear = module.as<nn::Linear>()) {
+  ///        linear->weight.normal_(0.0, 0.02);
+  ///      }
+  ///    }
+  ///
+  ///    MyModule module;
+  ///    module->apply(initialize_weights);
+  ///
   /// \endrst
   template <typename ModuleType>
   const typename ModuleType::ContainedType* as() const noexcept;
@@ -347,18 +364,20 @@ class TORCH_API Module : public std::enable_shared_from_this<Module> {
   /// Attempts to cast this `Module` to the given `ModuleType`.
   ///
   /// This method is useful when calling `apply()`.
+  ///
   /// \rst
   /// .. code-block:: cpp
   ///
-  ///   void initialize_weights(nn::Module& module) {
-  ///     torch::NoGradGuard no_grad;
-  ///     if (auto* linear = module.as<nn::Linear>()) {
-  ///       linear->weight.normal_(0.0, 0.02);
-  ///     }
-  ///   }
+  ///    void initialize_weights(nn::Module& module) {
+  ///      torch::NoGradGuard no_grad;
+  ///      if (auto* linear = module.as<nn::Linear>()) {
+  ///        linear->weight.normal_(0.0, 0.02);
+  ///      }
+  ///    }
+  ///
+  ///    MyModule module;
+  ///    module.apply(initialize_weights);
   ///
-  ///   MyModule module;
-  ///   module.apply(initialize_weights);
   /// \endrst
   template <
       typename ModuleType,
@@ -368,18 +387,20 @@ class TORCH_API Module : public std::enable_shared_from_this<Module> {
   /// Attempts to cast this `Module` to the given `ModuleType`.
   ///
   /// This method is useful when calling `apply()`.
+  ///
   /// \rst
   /// .. code-block:: cpp
   ///
-  ///   void initialize_weights(nn::Module& module) {
-  ///     torch::NoGradGuard no_grad;
-  ///     if (auto* linear = module.as<nn::Linear>()) {
-  ///       linear->weight.normal_(0.0, 0.02);
-  ///     }
-  ///   }
+  ///    void initialize_weights(nn::Module& module) {
+  ///      torch::NoGradGuard no_grad;
+  ///      if (auto* linear = module.as<nn::Linear>()) {
+  ///        linear->weight.normal_(0.0, 0.02);
+  ///      }
+  ///    }
+  ///
+  ///    MyModule module;
+  ///    module.apply(initialize_weights);
   ///
-  ///   MyModule module;
-  ///   module.apply(initialize_weights);
   /// \endrst
   template <
       typename ModuleType,
@@ -437,7 +458,7 @@ class TORCH_API Module : public std::enable_shared_from_this<Module> {
   ///
   /// A buffer is intended to be state in your module that does not record
   /// gradients, such as running statistics. Registering it makes it available
-  /// to methods such as `buffers()`, `clone()` or `to().
+  /// to methods such as `buffers()`, `clone()` or `to()`.
   ///
   /// \rst
   /// .. code-block:: cpp
@@ -487,8 +508,7 @@ class TORCH_API Module : public std::enable_shared_from_this<Module> {
   /// Replaces a registered submodule with this `Module`.
   ///
   /// This takes care of the registration, if you used submodule members, you
-  /// should
-  //  assign the submodule as well, i.e. use as
+  /// should assign the submodule as well, i.e. use as
   ///     module->submodule_ = module->replace_module("linear",
   ///     torch::nn::Linear(3, 4));
   /// It only works when a module of the name is already registered.
@@ -504,8 +524,7 @@ class TORCH_API Module : public std::enable_shared_from_this<Module> {
   /// This method deals with `ModuleHolder`s.
   ///
   /// This takes care of the registration, if you used submodule members, you
-  /// should
-  //  assign the submodule as well, i.e. use as
+  /// should assign the submodule as well, i.e. use as
   ///     module->submodule_ = module->replace_module("linear", linear_holder);
   /// It only works when a module of the name is already registered.
   ///
@@ -683,15 +702,15 @@ std::shared_ptr<ModuleType> Module::replace_module(
 
 template <typename... Ts>
 void Module::to_impl(Ts&&... ts) {
-  // First call `to()` on every child module.
+  /// First call `to()` on every child module.
   for (auto& child : children_) {
     child.value()->to(ts...);
   }
-  // Then move every parameter to the new dtype/device.
+  /// Then move every parameter to the new dtype/device.
   for (auto& parameter : named_parameters(/*recurse=*/false)) {
     parameter->set_data(parameter->to(ts...));
   }
-  // Then move every buffer to the new dtype/device.
+  /// Then move every buffer to the new dtype/device.
   for (auto& buffer : named_buffers(/*recurse=*/false)) {
     buffer->set_data(buffer->to(ts...));
   }
diff --git a/torch/csrc/api/include/torch/optim/optimizer.h b/torch/csrc/api/include/torch/optim/optimizer.h
index e2af5b74ee599..32f3e36f46d57 100644
--- a/torch/csrc/api/include/torch/optim/optimizer.h
+++ b/torch/csrc/api/include/torch/optim/optimizer.h
@@ -12,6 +12,7 @@
 #include <iterator>
 #include <memory>
 #include <string>
+#include <type_traits>
 #include <vector>
 
 // Forward declarations confuse Doxygen
@@ -66,12 +67,209 @@ class TORCH_API OptimizerOptions {
   virtual void set_lr(const double lr);
 };
 
+// Forward declarations for optimizer option types
+struct SGDOptions;
+struct AdamOptions;
+struct AdamWOptions;
+struct AdagradOptions;
+struct RMSpropOptions;
+struct LBFGSOptions;
+
+/**
+ * OptimizerCloneableOptions provides parameter group inheritance functionality
+ * for PyTorch C++ optimizer options. When creating parameter groups with
+ * partial options (e.g., AdamOptions().weight_decay(0.1)), fields not
+ * explicitly set by the user inherit from the optimizer's default values,
+ * while explicitly set fields are preserved.
+ *
+ * This enables Python-like behavior in C++:
+ * ```cpp
+ * // Python equivalent:
+ * // optimizer = Adam([{'params': params1, 'weight_decay': 0.1}], lr=0.01)
+ * // Result: weight_decay=0.1 preserved, lr=0.01 inherited
+ *
+ * AdamOptions defaults;
+ * defaults.lr(0.01).weight_decay(0.05);
+ *
+ * std::vector<OptimizerParamGroup> groups;
+ * groups.emplace_back(params1, std::make_unique<AdamOptions>(
+ *     AdamOptions().weight_decay(0.1)));  // Only weight_decay specified
+ *
+ * Adam optimizer(groups, defaults);
+ * // Result: group inherits lr=0.01, preserves weight_decay=0.1
+ * ```
+ *
+ * **Implementation**: Uses SFINAE-based field detection and constructor-default
+ * comparison to distinguish explicitly set fields from default values.
+ * Fields that match constructor defaults are inherited; others are preserved.
+ */
 template <typename Derived>
 class OptimizerCloneableOptions : public OptimizerOptions {
  private:
   std::unique_ptr<OptimizerOptions> clone() const override {
     return std::make_unique<Derived>(static_cast<const Derived&>(*this));
   }
+
+  // SFINAE field detection - detects optimizer fields using public accessor
+  // methods
+  template <class T, class Enable = void>
+  struct _has_lr : std::false_type {};
+  template <class T>
+  struct _has_lr<T, std::void_t<decltype(std::declval<const T&>().lr())>>
+      : std::true_type {};
+
+  template <class T, class Enable = void>
+  struct _has_momentum : std::false_type {};
+  template <class T>
+  struct _has_momentum<
+      T,
+      std::void_t<decltype(std::declval<const T&>().momentum())>>
+      : std::true_type {};
+
+  template <class T, class Enable = void>
+  struct _has_weight_decay : std::false_type {};
+  template <class T>
+  struct _has_weight_decay<
+      T,
+      std::void_t<decltype(std::declval<const T&>().weight_decay())>>
+      : std::true_type {};
+
+  template <class T, class Enable = void>
+  struct _has_dampening : std::false_type {};
+  template <class T>
+  struct _has_dampening<
+      T,
+      std::void_t<decltype(std::declval<const T&>().dampening())>>
+      : std::true_type {};
+
+  template <class T, class Enable = void>
+  struct _has_nesterov : std::false_type {};
+  template <class T>
+  struct _has_nesterov<
+      T,
+      std::void_t<decltype(std::declval<const T&>().nesterov())>>
+      : std::true_type {};
+
+  template <class T, class Enable = void>
+  struct _has_betas : std::false_type {};
+  template <class T>
+  struct _has_betas<T, std::void_t<decltype(std::declval<const T&>().betas())>>
+      : std::true_type {};
+
+  template <class T, class Enable = void>
+  struct _has_eps : std::false_type {};
+  template <class T>
+  struct _has_eps<T, std::void_t<decltype(std::declval<const T&>().eps())>>
+      : std::true_type {};
+
+  template <class T, class Enable = void>
+  struct _has_amsgrad : std::false_type {};
+  template <class T>
+  struct _has_amsgrad<
+      T,
+      std::void_t<decltype(std::declval<const T&>().amsgrad())>>
+      : std::true_type {};
+
+  // Optimizer-specific field detection
+  template <class T, class Enable = void>
+  struct _has_lr_decay : std::false_type {};
+  template <class T>
+  struct _has_lr_decay<
+      T,
+      std::void_t<decltype(std::declval<const T&>().lr_decay())>>
+      : std::true_type {};
+
+  template <class T, class Enable = void>
+  struct _has_alpha : std::false_type {};
+  template <class T>
+  struct _has_alpha<T, std::void_t<decltype(std::declval<const T&>().alpha())>>
+      : std::true_type {};
+
+  template <class T, class Enable = void>
+  struct _has_centered : std::false_type {};
+  template <class T>
+  struct _has_centered<
+      T,
+      std::void_t<decltype(std::declval<const T&>().centered())>>
+      : std::true_type {};
+
+  template <class T, class Enable = void>
+  struct _has_initial_accumulator_value : std::false_type {};
+  template <class T>
+  struct _has_initial_accumulator_value<
+      T,
+      std::void_t<
+          decltype(std::declval<const T&>().initial_accumulator_value())>>
+      : std::true_type {};
+
+  // LBFGS-specific fields with appropriate types
+  template <class T, class Enable = void>
+  struct _has_max_iter : std::false_type {};
+  template <class T>
+  struct _has_max_iter<
+      T,
+      std::void_t<decltype(std::declval<const T&>().max_iter())>>
+      : std::true_type {};
+
+  template <class T, class Enable = void>
+  struct _has_max_eval : std::false_type {};
+  template <class T>
+  struct _has_max_eval<
+      T,
+      std::void_t<decltype(std::declval<const T&>().max_eval())>>
+      : std::true_type {};
+
+  template <class T, class Enable = void>
+  struct _has_tolerance_grad : std::false_type {};
+  template <class T>
+  struct _has_tolerance_grad<
+      T,
+      std::void_t<decltype(std::declval<const T&>().tolerance_grad())>>
+      : std::true_type {};
+
+  template <class T, class Enable = void>
+  struct _has_tolerance_change : std::false_type {};
+  template <class T>
+  struct _has_tolerance_change<
+      T,
+      std::void_t<decltype(std::declval<const T&>().tolerance_change())>>
+      : std::true_type {};
+
+  template <class T, class Enable = void>
+  struct _has_history_size : std::false_type {};
+  template <class T>
+  struct _has_history_size<
+      T,
+      std::void_t<decltype(std::declval<const T&>().history_size())>>
+      : std::true_type {};
+
+  template <class T, class Enable = void>
+  struct _has_line_search_fn : std::false_type {};
+  template <class T>
+  struct _has_line_search_fn<
+      T,
+      std::void_t<decltype(std::declval<const T&>().line_search_fn())>>
+      : std::true_type {};
+
+  /**
+   * Merges user-specified options with optimizer defaults using
+   * constructor-default comparison to detect explicitly set fields.
+   *
+   * Algorithm:
+   * 1. Start with optimizer defaults as base
+   * 2. Create fresh constructor instance for comparison
+   * 3. If user_value != constructor_default → user explicitly set it → preserve
+   * 4. If user_value == constructor_default → user didn't set it → inherit from
+   * defaults
+   *
+   * Implementation is in optimizer.cpp to anchor vtable/typeinfo.
+   */
+  void _merge_by_comparison(
+      const Derived& defaults,
+      const Derived& user_options);
+
+  // Friend class for controlled access to private _merge_by_comparison method
+  friend class Optimizer;
 };
 
 /// Stores parameters in the param_group and stores a pointer to the
@@ -186,6 +384,43 @@ class TORCH_API Optimizer {
   /// Deserializes the optimizer state from the given `archive`.
   virtual void load(serialize::InputArchive& archive);
 
+ private:
+  /// Helper function to try merging for a specific optimizer type
+  template <typename OptimizerType>
+  static bool _try_merge_optimizer_type(
+      std::unique_ptr<OptimizerOptions>& final_options,
+      const OptimizerOptions& user_options,
+      const OptimizerOptions& defaults) {
+    auto* typed_final = dynamic_cast<OptimizerType*>(final_options.get());
+    auto* typed_user = dynamic_cast<const OptimizerType*>(&user_options);
+    auto* typed_defaults = dynamic_cast<const OptimizerType*>(&defaults);
+
+    if (typed_final && typed_user && typed_defaults) {
+      typed_final->_merge_by_comparison(*typed_defaults, *typed_user);
+      return true;
+    }
+    return false;
+  }
+
+  /// Simple variadic dispatch helper - try all optimizer types in one call
+  template <typename... OptimizerTypes>
+  static void _try_merge_all_optimizer_types(
+      std::unique_ptr<OptimizerOptions>& final_options,
+      const OptimizerOptions& user_options,
+      const OptimizerOptions& defaults) {
+    // Try each optimizer type until one succeeds - much cleaner than manual
+    // chain
+    (void)(_try_merge_optimizer_type<OptimizerTypes>(
+               final_options, user_options, defaults) ||
+           ...);
+  }
+
+  /// Convenience function with all known PyTorch optimizers
+  static void _try_merge_all_optimizers(
+      std::unique_ptr<OptimizerOptions>& final_options,
+      const OptimizerOptions& user_options,
+      const OptimizerOptions& defaults);
+
  protected:
   std::vector<OptimizerParamGroup> param_groups_;
   ska::flat_hash_map<void*, std::unique_ptr<OptimizerParamState>> state_;
diff --git a/torch/csrc/api/src/nn/modules/rnn.cpp b/torch/csrc/api/src/nn/modules/rnn.cpp
index 7ee864bc8ea94..c03d46e5b99e2 100644
--- a/torch/csrc/api/src/nn/modules/rnn.cpp
+++ b/torch/csrc/api/src/nn/modules/rnn.cpp
@@ -619,7 +619,7 @@ std::tuple<Tensor, std::tuple<Tensor, Tensor>> LSTMImpl::forward_helper(
          max_batch_size,
          options.hidden_size()},
         torch::dtype(input.dtype()).device(input.device()));
-    hx = std::make_tuple(h_zeros, c_zeros);
+    hx = std::make_tuple(std::move(h_zeros), std::move(c_zeros));
   } else {
     hx = hx_opt.value();
     // Each batch of the hidden state should match the input sequence that
@@ -632,7 +632,7 @@ std::tuple<Tensor, std::tuple<Tensor, Tensor>> LSTMImpl::forward_helper(
   if (!batch_sizes.defined()) {
     result = torch::lstm(
         input,
-        {std::get<0>(hx), std::get<1>(hx)},
+        {std::move(std::get<0>(hx)), std::move(std::get<1>(hx))},
         flat_weights_,
         options.bias(),
         options.num_layers(),
@@ -644,7 +644,7 @@ std::tuple<Tensor, std::tuple<Tensor, Tensor>> LSTMImpl::forward_helper(
     result = torch::lstm(
         input,
         batch_sizes,
-        {std::get<0>(hx), std::get<1>(hx)},
+        {std::move(std::get<0>(hx)), std::move(std::get<1>(hx))},
         flat_weights_,
         options.bias(),
         options.num_layers(),
@@ -652,10 +652,11 @@ std::tuple<Tensor, std::tuple<Tensor, Tensor>> LSTMImpl::forward_helper(
         this->is_training(),
         options.bidirectional());
   }
-  auto output = std::get<0>(result);
-  auto hidden = std::make_tuple(std::get<1>(result), std::get<2>(result));
+  auto output = std::move(std::get<0>(result));
+  auto hidden = std::make_tuple(
+      std::move(std::get<1>(result)), std::move(std::get<2>(result)));
 
-  return std::make_tuple(output, hidden);
+  return std::make_tuple(std::move(output), std::move(hidden));
 }
 
 std::tuple<Tensor, std::tuple<Tensor, Tensor>> LSTMImpl::forward(
diff --git a/torch/csrc/api/src/optim/lbfgs.cpp b/torch/csrc/api/src/optim/lbfgs.cpp
index 2241d8e0964b6..4da31ee903377 100644
--- a/torch/csrc/api/src/optim/lbfgs.cpp
+++ b/torch/csrc/api/src/optim/lbfgs.cpp
@@ -179,7 +179,7 @@ std::tuple<double, Tensor> LBFGS::_directional_evaluate(
   }
   auto flat_grad = _gather_flat_grad();
   _set_param(x);
-  return std::make_tuple(loss, flat_grad);
+  return std::make_tuple(loss, std::move(flat_grad));
 }
 
 static double _cubic_interpolate(
diff --git a/torch/csrc/api/src/optim/optimizer.cpp b/torch/csrc/api/src/optim/optimizer.cpp
index 24c40089bb5ea..674e07307d958 100644
--- a/torch/csrc/api/src/optim/optimizer.cpp
+++ b/torch/csrc/api/src/optim/optimizer.cpp
@@ -1,10 +1,195 @@
 #include <torch/optim/optimizer.h>
 
+#include <torch/csrc/autograd/generated/variable_factories.h>
+#include <torch/types.h>
+
+// Include complete type definitions for all optimizers to enable dynamic_cast
+#include <torch/optim/adagrad.h>
+#include <torch/optim/adam.h>
+#include <torch/optim/adamw.h>
+#include <torch/optim/lbfgs.h>
+#include <torch/optim/rmsprop.h>
+#include <torch/optim/sgd.h>
+
+#include <string>
 #include <utility>
 #include <vector>
 
 namespace torch::optim {
 
+// Implementation of OptimizerCloneableOptions<Derived>::_merge_by_comparison
+// Moved here to anchor vtable/typeinfo for template instantiations
+template <typename Derived>
+void OptimizerCloneableOptions<Derived>::_merge_by_comparison(
+    const Derived& defaults,
+    const Derived& user_options) {
+  auto* result = static_cast<Derived*>(this);
+  *result = defaults; // Start with optimizer defaults
+
+  // Create constructor defaults instance for comparison
+  Derived constructor_defaults = []() {
+    if constexpr (std::is_default_constructible_v<Derived>) {
+      return Derived{};
+    } else {
+      // Handle optimizers requiring constructor parameters
+      if constexpr (std::is_same_v<Derived, SGDOptions>) {
+        return Derived(1e-3);
+      } else if constexpr (std::is_same_v<Derived, AdagradOptions>) {
+        return Derived(1e-2);
+      } else if constexpr (std::is_same_v<Derived, RMSpropOptions>) {
+        return Derived(1e-2);
+      } else if constexpr (std::is_same_v<Derived, LBFGSOptions>) {
+        return Derived(1);
+      } else {
+        return Derived{};
+      }
+    }
+  }();
+
+  // Merge fields: preserve user-set values, inherit defaults for unset values
+
+  if constexpr (OptimizerCloneableOptions<Derived>::_has_lr<Derived>::value) {
+    if (user_options.lr() != constructor_defaults.lr()) {
+      result->lr(user_options.lr());
+    }
+  }
+  if constexpr (OptimizerCloneableOptions<Derived>::_has_momentum<
+                    Derived>::value) {
+    if (user_options.momentum() != constructor_defaults.momentum()) {
+      result->momentum(user_options.momentum());
+    }
+  }
+  if constexpr (OptimizerCloneableOptions<Derived>::_has_weight_decay<
+                    Derived>::value) {
+    if (user_options.weight_decay() != constructor_defaults.weight_decay()) {
+      result->weight_decay(user_options.weight_decay());
+    }
+  }
+  if constexpr (OptimizerCloneableOptions<Derived>::_has_dampening<
+                    Derived>::value) {
+    if (user_options.dampening() != constructor_defaults.dampening()) {
+      result->dampening(user_options.dampening());
+    }
+  }
+  if constexpr (OptimizerCloneableOptions<Derived>::_has_nesterov<
+                    Derived>::value) {
+    if (user_options.nesterov() != constructor_defaults.nesterov()) {
+      result->nesterov(user_options.nesterov());
+    }
+  }
+  if constexpr (OptimizerCloneableOptions<Derived>::_has_betas<
+                    Derived>::value) {
+    if (user_options.betas() != constructor_defaults.betas()) {
+      result->betas(user_options.betas());
+    }
+  }
+  if constexpr (OptimizerCloneableOptions<Derived>::_has_eps<Derived>::value) {
+    if (user_options.eps() != constructor_defaults.eps()) {
+      result->eps(user_options.eps());
+    }
+  }
+  if constexpr (OptimizerCloneableOptions<Derived>::_has_amsgrad<
+                    Derived>::value) {
+    if (user_options.amsgrad() != constructor_defaults.amsgrad()) {
+      result->amsgrad(user_options.amsgrad());
+    }
+  }
+
+  // Optimizer-specific fields - automatically detected and handled
+  if constexpr (OptimizerCloneableOptions<Derived>::_has_lr_decay<
+                    Derived>::value) {
+    if (user_options.lr_decay() != constructor_defaults.lr_decay()) {
+      result->lr_decay(user_options.lr_decay());
+    }
+  }
+  if constexpr (OptimizerCloneableOptions<Derived>::_has_alpha<
+                    Derived>::value) {
+    if (user_options.alpha() != constructor_defaults.alpha()) {
+      result->alpha(user_options.alpha());
+    }
+  }
+  if constexpr (OptimizerCloneableOptions<Derived>::_has_centered<
+                    Derived>::value) {
+    if (user_options.centered() != constructor_defaults.centered()) {
+      result->centered(user_options.centered());
+    }
+  }
+  if constexpr (OptimizerCloneableOptions<
+                    Derived>::_has_initial_accumulator_value<Derived>::value) {
+    if (user_options.initial_accumulator_value() !=
+        constructor_defaults.initial_accumulator_value()) {
+      result->initial_accumulator_value(
+          user_options.initial_accumulator_value());
+    }
+  }
+
+  // LBFGS-specific fields with appropriate types
+  if constexpr (OptimizerCloneableOptions<Derived>::_has_max_iter<
+                    Derived>::value) {
+    if (user_options.max_iter() != constructor_defaults.max_iter()) {
+      result->max_iter(user_options.max_iter());
+    }
+  }
+  if constexpr (OptimizerCloneableOptions<Derived>::_has_max_eval<
+                    Derived>::value) {
+    if (user_options.max_eval() != constructor_defaults.max_eval()) {
+      result->max_eval(user_options.max_eval());
+    }
+  }
+  if constexpr (OptimizerCloneableOptions<Derived>::_has_tolerance_grad<
+                    Derived>::value) {
+    if (user_options.tolerance_grad() !=
+        constructor_defaults.tolerance_grad()) {
+      result->tolerance_grad(user_options.tolerance_grad());
+    }
+  }
+  if constexpr (OptimizerCloneableOptions<Derived>::_has_tolerance_change<
+                    Derived>::value) {
+    if (user_options.tolerance_change() !=
+        constructor_defaults.tolerance_change()) {
+      result->tolerance_change(user_options.tolerance_change());
+    }
+  }
+  if constexpr (OptimizerCloneableOptions<Derived>::_has_history_size<
+                    Derived>::value) {
+    if (user_options.history_size() != constructor_defaults.history_size()) {
+      result->history_size(user_options.history_size());
+    }
+  }
+  if constexpr (OptimizerCloneableOptions<Derived>::_has_line_search_fn<
+                    Derived>::value) {
+    if (user_options.line_search_fn() !=
+        constructor_defaults.line_search_fn()) {
+      result->line_search_fn(user_options.line_search_fn());
+    }
+  }
+}
+
+// Explicit template instantiations to anchor vtable/typeinfo
+// These instantiations ensure the compiler generates the full class definition
+// and vtable for each OptimizerCloneableOptions<T> specialization
+template class OptimizerCloneableOptions<SGDOptions>;
+template class OptimizerCloneableOptions<AdamOptions>;
+template class OptimizerCloneableOptions<AdamWOptions>;
+template class OptimizerCloneableOptions<AdagradOptions>;
+template class OptimizerCloneableOptions<RMSpropOptions>;
+template class OptimizerCloneableOptions<LBFGSOptions>;
+
+// Simple implementation using variadic template helper
+void Optimizer::_try_merge_all_optimizers(
+    std::unique_ptr<OptimizerOptions>& final_options,
+    const OptimizerOptions& user_options,
+    const OptimizerOptions& defaults) {
+  // Clean one-liner replaces the entire repetitive dispatch chain
+  _try_merge_all_optimizer_types<
+      SGDOptions,
+      AdamOptions,
+      AdamWOptions,
+      AdagradOptions,
+      RMSpropOptions,
+      LBFGSOptions>(final_options, user_options, defaults);
+}
+
 bool OptimizerParamGroup::has_options() const {
   return options_ != nullptr;
 }
@@ -97,9 +282,20 @@ void Optimizer::add_param_group(const OptimizerParamGroup& param_group) {
   TORCH_INTERNAL_ASSERT(defaults_ != nullptr);
   OptimizerParamGroup param_group_(param_group.params());
   if (!param_group.has_options()) {
+    // No options provided - use defaults directly
     param_group_.set_options(defaults_->clone());
   } else {
-    param_group_.set_options(param_group.options().clone());
+    // Options provided - merge user's explicit settings with defaults for
+    // parameter group inheritance This enables Python-C++ API parity by
+    // honoring user intent while inheriting missing parameters
+    auto final_options = defaults_->clone();
+
+    // Simple variadic dispatch - try all known optimizer types
+    _try_merge_all_optimizers(final_options, param_group.options(), *defaults_);
+
+    // If no merging was done (custom optimizer), final_options already contains
+    // defaults
+    param_group_.set_options(std::move(final_options));
   }
   for (const auto& p : param_group_.params()) {
     TORCH_CHECK(
diff --git a/torch/csrc/autograd/FunctionsManual.cpp b/torch/csrc/autograd/FunctionsManual.cpp
index 31fcec3deaf40..e1761b000d8c2 100644
--- a/torch/csrc/autograd/FunctionsManual.cpp
+++ b/torch/csrc/autograd/FunctionsManual.cpp
@@ -17,6 +17,7 @@
 #include <ATen/core/Reduction.h>
 #include <ATen/core/grad_mode.h>
 #include <ATen/native/Activation.h>
+#include <ATen/native/GridSamplerUtils.h>
 #include <ATen/native/IndexingUtils.h>
 #include <ATen/native/LinearAlgebraUtils.h>
 #include <ATen/native/SparseTensorUtils.h>
@@ -43,6 +44,8 @@ using at::OptionalIntArrayRef;
 using at::Scalar;
 using at::Tensor;
 using at::TensorList;
+using at::native::detail::GridSamplerInterpolation;
+using at::native::detail::GridSamplerPadding;
 
 const char* kCudnnDoubleBackwardMsg =
     "Double backwards is not supported for CuDNN RNNs due to limitations in the CuDNN API. To run double backwards, please disable the CuDNN backend temporarily while running the forward pass of your RNN. For example: \nwith torch.backends.cudnn.flags(enabled=False):\n    output = model(inputs)";
@@ -223,6 +226,67 @@ Tensor amaxamin_jvp(
   return at::where(mask, dx, 0.).sum(dim, keepdim) / mask.sum(dim, keepdim);
 }
 
+// Builds the dims vector for aminmax: all dims when dim is nullopt,
+// or the single wrapped dim otherwise.
+static std::vector<int64_t> _aminmax_dims(
+    const Tensor& self,
+    std::optional<int64_t> dim_opt) {
+  if (dim_opt.has_value()) {
+    return {at::maybe_wrap_dim(*dim_opt, self.dim())};
+  }
+  std::vector<int64_t> dims(self.dim());
+  std::iota(dims.begin(), dims.end(), 0);
+  return dims;
+}
+
+Tensor aminmax_backward(
+    const Tensor& self,
+    std::optional<int64_t> dim,
+    bool keepdim,
+    const Tensor& grad_min,
+    const Tensor& grad_max,
+    const Tensor& min,
+    const Tensor& max) {
+  auto dims_vec = _aminmax_dims(self, dim);
+  IntArrayRef dims(dims_vec);
+  Tensor result;
+
+  if (grad_min.defined()) {
+    auto grad_min_expanded = restore_reduced_dims(grad_min, dims, keepdim);
+    auto min_mask = (self == restore_reduced_dims(min, dims, keepdim));
+    result = scale_grad_by_count(grad_min_expanded, min_mask, dims);
+  }
+
+  if (grad_max.defined()) {
+    auto grad_max_expanded = restore_reduced_dims(grad_max, dims, keepdim);
+    auto max_mask = (self == restore_reduced_dims(max, dims, keepdim));
+    auto grad_max_result =
+        scale_grad_by_count(grad_max_expanded, max_mask, dims);
+
+    if (result.defined()) {
+      if (!areAnyTensorSubclassLike({result, grad_max_result})) {
+        result.add_(grad_max_result);
+      } else {
+        result = result + grad_max_result;
+      }
+    } else {
+      result = grad_max_result;
+    }
+  }
+
+  return result;
+}
+
+Tensor aminmax_jvp(
+    const Tensor& self_p,
+    const Tensor& self_t,
+    const Tensor& result,
+    std::optional<int64_t> dim,
+    bool keepdim) {
+  auto dims_vec = _aminmax_dims(self_p, dim);
+  return amaxamin_jvp(self_p, self_t, result, IntArrayRef(dims_vec), keepdim);
+}
+
 std::tuple<Tensor, Tensor> _euclidean_dist_backward(
     const Tensor& grad,
     const Tensor& x1,
@@ -989,6 +1053,7 @@ Tensor unbind_backward_nested(
     int64_t dim,
     const at::TensorOptions& options) {
   std::vector<Tensor> grads_tensors;
+  grads_tensors.reserve(grads.size());
   for (int64_t i : c10::irange(static_cast<int64_t>(grads.size()))) {
     if (grads[i].defined()) {
       grads_tensors.push_back(static_cast<Tensor>(grads[i]));
@@ -2185,6 +2250,7 @@ Tensor _nested_split_with_sizes_backward(
   // it's possible some of the grads are not defined (represents tensors of all
   // 0s). Since at::cat can't handle those, let's define them
   std::vector<Tensor> grads_all_defined;
+  grads_all_defined.reserve(grads.size());
   for (int64_t i : c10::irange(static_cast<int64_t>(grads.size()))) {
     if (grads[i].defined()) {
       grads_all_defined.push_back(static_cast<Tensor>(grads[i]));
@@ -2251,7 +2317,6 @@ Tensor error_for_max_pool2d_double_backward() { // This is mps-only.
       "max_pool2d with `return_indices=False` is not infinitely differentiable.",
       " If you want to calculate higher order derivatives, e.g. second order,",
       " set `return_indices=True`.");
-  return Tensor();
 }
 
 Tensor glu_double_backward(
@@ -4002,7 +4067,7 @@ std::tuple<Tensor, Tensor> linalg_lstsq_backward(
     B_grad = B_grad_X;
   }
 
-  return std::make_tuple(A_grad, B_grad);
+  return std::make_tuple(std::move(A_grad), std::move(B_grad));
 }
 
 std::tuple<Tensor, Tensor> linalg_qr_jvp(
@@ -4312,10 +4377,20 @@ Tensor linalg_det_backward(
     const Tensor& LU,
     const Tensor& pivots) {
   at::NoTF32Guard disable_tf32;
-  // A.numel() == 0 necessary for the singular case
-  if (!grad.defined() || A.sym_numel() == 0) {
+  if (!grad.defined()) {
     return {};
   }
+  if (A.sym_numel() == 0) {
+    return at::zeros_like(A);
+  }
+
+  // Special case handling for 1 x 1 matrix, to ensure mathematically correct.
+  // d(det)/dA = 1, so gradient = grad * ones_like(A)
+  // See #80761
+  if (A.sym_size(-2) == 1 && A.sym_size(-1) == 1) {
+    // For batched 1x1 matrices, broadcast grad to match A's shape
+    return grad.unsqueeze(-1).unsqueeze(-1).expand_as(A);
+  }
 
   // The gradient G is the matrix solving
   // A.mH G = det(A).conj() * grad * I
@@ -4590,7 +4665,7 @@ std::tuple<Tensor, Tensor> linalg_solve_triangular_backward(
     Tensor G_A = left ? -at::matmul(G_B, X_H) : -at::matmul(X_H, G_B);
     G_A = upper ? G_A.triu(static_cast<int>(unitriangular))
                 : G_A.tril(-static_cast<int>(unitriangular));
-    return std::make_tuple(G_A, B_requires_grad ? G_B : Tensor{});
+    return std::make_tuple(std::move(G_A), B_requires_grad ? G_B : Tensor{});
   } else {
     return std::make_tuple(Tensor{}, G_B);
   }
@@ -5132,7 +5207,7 @@ std::tuple<Tensor, Tensor> infinitely_differentiable_native_rms_norm_backward(
     }
   }
 
-  return std::make_tuple(dX, dgamma);
+  return std::make_tuple(std::move(dX), std::move(dgamma));
 }
 
 std::tuple<Tensor, Tensor, Tensor>
@@ -5220,7 +5295,7 @@ infinitely_differentiable_native_group_norm_backward(
     dbeta = db.sum(0).reshape_as(toNonOptTensor(gamma));
   }
 
-  return std::make_tuple(dX, dgamma, dbeta);
+  return std::make_tuple(std::move(dX), std::move(dgamma), std::move(dbeta));
 }
 
 std::tuple<Tensor, Tensor, Tensor> _trilinear_backward(
@@ -5754,7 +5829,8 @@ std::tuple<Tensor, Tensor, Tensor> ormqr_backward(
   Tensor self_grad, tau_grad, other_grad;
 
   if (!grad.defined()) {
-    return std::make_tuple(self_grad, tau_grad, other_grad);
+    return std::make_tuple(
+        std::move(self_grad), std::move(tau_grad), std::move(other_grad));
   }
 
   const auto self_requires_grad = grad_output_mask[0];
@@ -5813,7 +5889,7 @@ std::tuple<Tensor, Tensor> polar_backward(
     auto result_mul_1_j = result * Scalar(c10::complex<double>{0.0, 1.0});
     grad_angle = at::real(grad_conj * result_mul_1_j);
   }
-  return std::make_tuple(grad_abs, grad_angle);
+  return std::make_tuple(std::move(grad_abs), std::move(grad_angle));
 }
 
 Tensor i1_backward(
@@ -7199,7 +7275,7 @@ std::tuple<Tensor, Tensor> scatter_reduce_backward(
     grad_self = grad_self.scatter(dim, index, 0);
   }
 
-  return std::make_tuple(grad_self, grad_src);
+  return std::make_tuple(std::move(grad_self), std::move(grad_src));
 }
 
 Tensor _to_copy_backward(
@@ -7295,7 +7371,7 @@ std::tuple<Tensor, Tensor> index_reduce_backward(
     grad_self = grad_self.index_fill(dim, index, 0);
   }
 
-  return std::make_tuple(grad_self, grad_src);
+  return std::make_tuple(std::move(grad_self), std::move(grad_src));
 }
 
 Tensor take_backward(
@@ -7453,7 +7529,14 @@ mkldnn_rnn_layer_differentiable_backward(
   }
 
   auto cat_layer_dx = at::cat(layer_dx, 0);
-  return std::make_tuple(cat_layer_dx, dWx, dWh, db, db, dprev_h, dprev_c);
+  return std::make_tuple(
+      std::move(cat_layer_dx),
+      std::move(dWx),
+      std::move(dWh),
+      db,
+      db,
+      std::move(dprev_h),
+      std::move(dprev_c));
 }
 
 Tensor values_backward(const Tensor& grad, const Tensor& self) {
@@ -7485,4 +7568,697 @@ Tensor values_backward(const Tensor& grad, const Tensor& self) {
   return grad_self;
 }
 
+// ==================== grid_sampler double backward ====================
+
+namespace {
+
+// Bound integer tap coordinate per padding mode (for bicubic taps).
+// Returns a kLong tensor in [0, size-1].
+static Tensor gs_bound_coord(
+    const Tensor& idx,
+    int64_t size,
+    GridSamplerPadding padding_mode,
+    bool align_corners) {
+  if (padding_mode != GridSamplerPadding::Reflection) {
+    return idx.clamp(0, size - 1);
+  }
+  // Reflection: mirrors the PyTorch reflect_coordinates convention.
+  if (size <= 1) {
+    return at::zeros_like(idx);
+  }
+  auto x = idx.to(at::kDouble);
+  double span, min_v;
+  if (align_corners) {
+    span = size - 1;
+    min_v = 0.0;
+  } else {
+    span = size;
+    min_v = -0.5;
+  }
+  // Native formula: in = idx - min_v (= idx + 0.5 for no-align_corners).
+  // For integer idx this is always a half-integer, so fold result + min_v is
+  // an exact integer — no rounding artifacts.
+  auto in = x - min_v;
+  auto in_abs = in.abs();
+  auto even = at::fmod(at::floor(in_abs / span), 2.0) < 0.5;
+  auto rem = at::fmod(in_abs, span);
+  auto r = at::where(even, rem, span - rem);
+  return (r + min_v).round().clamp(0, size - 1).to(at::kLong);
+}
+
+// Compute unnormalized source coordinate and per-position chain-rule multiplier
+// d(source_coord)/d(grid_coord), accounting for unnormalization and padding.
+static std::pair<Tensor, Tensor> gs_compute_coords(
+    const Tensor& coord,
+    int64_t size,
+    GridSamplerPadding padding_mode,
+    bool align_corners) {
+  double unnorm_scale = align_corners ? (size - 1) / 2.0 : size / 2.0;
+  Tensor ix = align_corners ? (coord + 1) * unnorm_scale
+                            : (coord + 1) * unnorm_scale - 0.5;
+  Tensor padding_grad;
+  if (padding_mode == GridSamplerPadding::Zeros) {
+    padding_grad = at::ones_like(ix);
+  } else if (padding_mode == GridSamplerPadding::Border) {
+    // clip_coordinates_set_grad: grad = 1 iff strictly inside (0, size-1)
+    padding_grad =
+        ((ix > 0) & (ix < static_cast<double>(size - 1))).to(ix.dtype());
+    ix = ix.clamp(0, size - 1);
+  } else {
+    TORCH_CHECK_NOT_IMPLEMENTED(
+        padding_mode == GridSamplerPadding::Reflection,
+        "Unknown padding mode: ",
+        static_cast<int64_t>(padding_mode));
+    double twice_low = align_corners ? 0.0 : -1.0;
+    double twice_high = align_corners ? 2.0 * (size - 1) : 2.0 * size - 1.0;
+    if (twice_high <= twice_low) {
+      auto z = at::zeros_like(ix);
+      return {z, z};
+    }
+    double min_v = twice_low / 2.0;
+    double span = (twice_high - twice_low) / 2.0;
+    auto in = ix - min_v;
+    auto in_abs = in.abs();
+    auto reflect_sign = at::where(in >= 0, 1.0, -1.0);
+    auto even = at::fmod(at::floor(in_abs / span), 2.0) < 0.5;
+    auto flip_sign = at::where(even, 1.0, -1.0);
+    auto rem = in_abs % span;
+    auto ix_refl = at::where(even, rem + min_v, span - rem + min_v);
+    auto clip_grad = ((ix_refl > 0) & (ix_refl < static_cast<double>(size - 1)))
+                         .to(ix.dtype());
+    padding_grad = reflect_sign * flip_sign * clip_grad;
+    ix = ix_refl.clamp(0, size - 1);
+  }
+  return {std::move(ix), padding_grad * unnorm_scale};
+}
+
+static Tensor gs_accum_sumprod_k(const Tensor& values, const Tensor& basis) {
+  return (values * basis.unsqueeze(1)).sum(-1);
+}
+
+static Tensor gs_outer_last(const Tensor& a, const Tensor& b) {
+  auto shape = a.sizes().vec();
+  shape.back() *= b.size(-1);
+  return (b.unsqueeze(-1) * a.unsqueeze(-2)).reshape(shape);
+}
+
+// Multi-tap gather for 2D: h_idx/w_idx [N, Ho, Wo, K] → [N, C, Ho, Wo, K]
+static Tensor gs_gather2d_multi(
+    const Tensor& input,
+    const Tensor& h_idx,
+    const Tensor& w_idx,
+    bool zeros_oob) {
+  auto N = input.size(0), C = input.size(1);
+  auto H = input.size(2), W = input.size(3);
+  auto out_H = h_idx.size(1), out_W = h_idx.size(2), K = h_idx.size(3);
+  auto flat = (h_idx.clamp(0, H - 1) * W + w_idx.clamp(0, W - 1))
+                  .reshape({N, 1, out_H * out_W * K})
+                  .expand({N, C, out_H * out_W * K});
+  auto result = input.reshape({N, C, H * W})
+                    .gather(2, flat)
+                    .reshape({N, C, out_H, out_W, K});
+  if (zeros_oob) {
+    auto mask = (h_idx >= 0) & (h_idx < H) & (w_idx >= 0) & (w_idx < W);
+    result = result * mask.unsqueeze(1).to(result.dtype());
+  }
+  return result;
+}
+
+// Multi-tap scatter for 2D: values [N,C,Ho,Wo], weights [N,Ho,Wo,K] → [N,C,H,W]
+static Tensor gs_scatter2d_multi(
+    const Tensor& values,
+    const Tensor& weights,
+    const Tensor& h_idx,
+    const Tensor& w_idx,
+    int64_t H,
+    int64_t W,
+    bool zeros_oob) {
+  auto N = values.size(0), C = values.size(1);
+  auto out_H = values.size(2), out_W = values.size(3);
+  auto K = h_idx.size(3);
+  auto flat = (h_idx.clamp(0, H - 1) * W + w_idx.clamp(0, W - 1))
+                  .reshape({N, 1, out_H * out_W * K})
+                  .expand({N, C, out_H * out_W * K});
+  auto weighted = values.unsqueeze(-1) * weights.unsqueeze(1);
+  if (zeros_oob) {
+    auto mask = (h_idx >= 0) & (h_idx < H) & (w_idx >= 0) & (w_idx < W);
+    weighted = weighted * mask.unsqueeze(1).to(weighted.dtype());
+  }
+  return at::zeros({N, C, H * W}, values.options())
+      .scatter_add(2, flat, weighted.reshape({N, C, out_H * out_W * K}))
+      .reshape({N, C, H, W});
+}
+
+// Multi-tap bounded gather for bicubic 2D: h_idx/w_idx [N, Ho, Wo, K] → [N, C,
+// Ho, Wo, K]
+static Tensor gs_gather2d_bc_multi(
+    const Tensor& input,
+    const Tensor& h_idx,
+    const Tensor& w_idx,
+    GridSamplerPadding padding_mode,
+    bool align_corners) {
+  auto N = input.size(0), C = input.size(1);
+  auto H = input.size(2), W = input.size(3);
+  auto out_H = h_idx.size(1), out_W = h_idx.size(2), K = h_idx.size(3);
+  auto flat = (gs_bound_coord(h_idx, H, padding_mode, align_corners) * W +
+               gs_bound_coord(w_idx, W, padding_mode, align_corners))
+                  .reshape({N, 1, out_H * out_W * K})
+                  .expand({N, C, out_H * out_W * K});
+  auto result = input.reshape({N, C, H * W})
+                    .gather(2, flat)
+                    .reshape({N, C, out_H, out_W, K});
+  if (padding_mode == GridSamplerPadding::Zeros) {
+    auto mask = (h_idx >= 0) & (h_idx < H) & (w_idx >= 0) & (w_idx < W);
+    result = result * mask.unsqueeze(1).to(result.dtype());
+  }
+  return result;
+}
+
+// Multi-tap bounded scatter for bicubic 2D: values [N,C,Ho,Wo], weights
+// [N,Ho,Wo,K] → [N,C,H,W]
+static Tensor gs_scatter2d_bc_multi(
+    const Tensor& values,
+    const Tensor& weights,
+    const Tensor& h_idx,
+    const Tensor& w_idx,
+    int64_t H,
+    int64_t W,
+    GridSamplerPadding padding_mode,
+    bool align_corners) {
+  auto N = values.size(0), C = values.size(1);
+  auto out_H = values.size(2), out_W = values.size(3);
+  auto K = h_idx.size(3);
+  auto flat = (gs_bound_coord(h_idx, H, padding_mode, align_corners) * W +
+               gs_bound_coord(w_idx, W, padding_mode, align_corners))
+                  .reshape({N, 1, out_H * out_W * K})
+                  .expand({N, C, out_H * out_W * K});
+  auto weighted = values.unsqueeze(-1) * weights.unsqueeze(1);
+  if (padding_mode == GridSamplerPadding::Zeros) {
+    auto mask = (h_idx >= 0) & (h_idx < H) & (w_idx >= 0) & (w_idx < W);
+    weighted = weighted * mask.unsqueeze(1).to(weighted.dtype());
+  }
+  return at::zeros({N, C, H * W}, values.options())
+      .scatter_add(2, flat, weighted.reshape({N, C, out_H * out_W * K}))
+      .reshape({N, C, H, W});
+}
+
+// Multi-tap gather for 3D: d_idx/h_idx/w_idx [N, Do, Ho, Wo, K] → [N, C, Do,
+// Ho, Wo, K]
+static Tensor gs_gather3d_multi(
+    const Tensor& input,
+    const Tensor& d_idx,
+    const Tensor& h_idx,
+    const Tensor& w_idx,
+    bool zeros_oob) {
+  auto N = input.size(0), C = input.size(1);
+  auto D = input.size(2), H = input.size(3), W = input.size(4);
+  auto out_D = d_idx.size(1), out_H = d_idx.size(2), out_W = d_idx.size(3),
+       K = d_idx.size(4);
+  auto flat = ((d_idx.clamp(0, D - 1) * H + h_idx.clamp(0, H - 1)) * W +
+               w_idx.clamp(0, W - 1))
+                  .reshape({N, 1, out_D * out_H * out_W * K})
+                  .expand({N, C, out_D * out_H * out_W * K});
+  auto result = input.reshape({N, C, D * H * W})
+                    .gather(2, flat)
+                    .reshape({N, C, out_D, out_H, out_W, K});
+  if (zeros_oob) {
+    auto mask = (d_idx >= 0) & (d_idx < D) & (h_idx >= 0) & (h_idx < H) &
+        (w_idx >= 0) & (w_idx < W);
+    result = result * mask.unsqueeze(1).to(result.dtype());
+  }
+  return result;
+}
+
+// Multi-tap scatter for 3D: values [N,C,Do,Ho,Wo], weights [N,Do,Ho,Wo,K] →
+// [N,C,D,H,W]
+static Tensor gs_scatter3d_multi(
+    const Tensor& values,
+    const Tensor& weights,
+    const Tensor& d_idx,
+    const Tensor& h_idx,
+    const Tensor& w_idx,
+    int64_t D,
+    int64_t H,
+    int64_t W,
+    bool zeros_oob) {
+  auto N = values.size(0), C = values.size(1);
+  auto out_D = values.size(2), out_H = values.size(3), out_W = values.size(4);
+  auto K = d_idx.size(4);
+  auto flat = ((d_idx.clamp(0, D - 1) * H + h_idx.clamp(0, H - 1)) * W +
+               w_idx.clamp(0, W - 1))
+                  .reshape({N, 1, out_D * out_H * out_W * K})
+                  .expand({N, C, out_D * out_H * out_W * K});
+  auto weighted = values.unsqueeze(-1) * weights.unsqueeze(1);
+  if (zeros_oob) {
+    auto mask = (d_idx >= 0) & (d_idx < D) & (h_idx >= 0) & (h_idx < H) &
+        (w_idx >= 0) & (w_idx < W);
+    weighted = weighted * mask.unsqueeze(1).to(weighted.dtype());
+  }
+  return at::zeros({N, C, D * H * W}, values.options())
+      .scatter_add(2, flat, weighted.reshape({N, C, out_D * out_H * out_W * K}))
+      .reshape({N, C, D, H, W});
+}
+
+} // anonymous namespace
+
+std::tuple<Tensor, Tensor, Tensor> grid_sampler_2d_double_backward(
+    const Tensor& ggI,
+    const Tensor& ggGrid,
+    const Tensor& grad_output,
+    const Tensor& input,
+    const Tensor& grid,
+    int64_t interpolation_mode,
+    int64_t padding_mode,
+    bool align_corners,
+    std::array<bool, 3> output_mask) {
+  Tensor d_grad_output, d_input, d_grid;
+  const auto interpolation =
+      static_cast<GridSamplerInterpolation>(interpolation_mode);
+  const auto padding_mode_enum = static_cast<GridSamplerPadding>(padding_mode);
+
+  // ggI -> d_grad_output: gather ggI at grid positions = grid_sampler_2d(ggI,
+  // grid)
+  if (output_mask[0] && ggI.defined()) {
+    d_grad_output = at::grid_sampler_2d(
+        ggI, grid, interpolation_mode, padding_mode, align_corners);
+  }
+  // ggI -> d_grid: same structure as grad_grid but with ggI as "input"
+  if (output_mask[2] && ggI.defined()) {
+    d_grid = std::get<1>(at::grid_sampler_2d_backward(
+        grad_output,
+        ggI,
+        grid,
+        interpolation_mode,
+        padding_mode,
+        align_corners,
+        {false, true}));
+  }
+  // d_input from ggI is 0: grad_input has no dependence on input.
+  // For nearest, grad_grid = 0, so all ggGrid contributions vanish.
+  if (!ggGrid.defined() || interpolation == GridSamplerInterpolation::Nearest) {
+    return {std::move(d_grad_output), std::move(d_input), std::move(d_grid)};
+  }
+  auto H = input.size(2), W = input.size(3);
+
+  auto [ix, gix_mult] = gs_compute_coords(
+      grid.select(-1, 0), W, padding_mode_enum, align_corners);
+  auto [iy, giy_mult] = gs_compute_coords(
+      grid.select(-1, 1), H, padding_mode_enum, align_corners);
+
+  auto x0 = at::floor(ix).to(at::kLong);
+  auto y0 = at::floor(iy).to(at::kLong);
+  auto fx = ix - x0.to(ix.dtype());
+  auto fy = iy - y0.to(iy.dtype());
+
+  bool zeros_oob = (padding_mode_enum == GridSamplerPadding::Zeros);
+  auto ggG_x = ggGrid.select(-1, 0) * gix_mult; // (N, out_H, out_W)
+  auto ggG_y = ggGrid.select(-1, 1) * giy_mult;
+
+  if (interpolation == GridSamplerInterpolation::Bilinear) {
+    auto x1 = x0 + 1;
+    auto y1 = y0 + 1;
+    // Stack all 4 tap indices: [N, Ho, Wo, 4] order: nw, ne, sw, se
+    auto h_stk = at::stack({y0, y0, y1, y1}, -1);
+    auto w_stk = at::stack({x0, x1, x0, x1}, -1);
+
+    Tensor I; // [N, C, Ho, Wo, 4]
+    if (output_mask[0] || output_mask[1] || output_mask[2]) {
+      I = gs_gather2d_multi(input, h_stk, w_stk, zeros_oob);
+    }
+
+    Tensor dw_dx, dw_dy;
+    if (output_mask[0] || output_mask[1]) {
+      // Derivative weights [N, Ho, Wo, 4]: d(w_k)/d(ix) and d(w_k)/d(iy)
+      auto omx = 1.0 - fx, omy = 1.0 - fy;
+      dw_dx = at::stack({-omy, omy, -fy, fy}, -1);
+      dw_dy = at::stack({-omx, -fx, omx, fx}, -1);
+    }
+
+    if (output_mask[0]) {
+      auto sum_dx = gs_accum_sumprod_k(I, dw_dx);
+      auto sum_dy = gs_accum_sumprod_k(I, dw_dy);
+      auto contrib = sum_dx * ggG_x.unsqueeze(1);
+      contrib.addcmul_(sum_dy, ggG_y.unsqueeze(1));
+      d_grad_output = d_grad_output.defined() ? d_grad_output + contrib
+                                              : std::move(contrib);
+    }
+
+    if (output_mask[1]) {
+      // Combined per-tap weight [N, Ho, Wo, 4]
+      auto w_in = ggG_x.unsqueeze(-1) * dw_dx;
+      w_in.addcmul_(ggG_y.unsqueeze(-1), dw_dy);
+      d_input =
+          gs_scatter2d_multi(grad_output, w_in, h_stk, w_stk, H, W, zeros_oob);
+    }
+
+    if (output_mask[2]) {
+      // cross second derivative: coeffs [nw,ne,sw,se] = [+1,-1,-1,+1]
+      auto I_nw = I.select(-1, 0);
+      auto I_ne = I.select(-1, 1);
+      auto I_sw = I.select(-1, 2);
+      auto I_se = I.select(-1, 3);
+      auto cross = (grad_output * (I_nw - I_ne - I_sw + I_se)).sum(1);
+      auto d_grid_x = gix_mult * ggG_y * cross;
+      auto d_grid_y = giy_mult * ggG_x * cross;
+      auto contrib = at::stack({std::move(d_grid_x), std::move(d_grid_y)}, -1);
+      d_grid = d_grid.defined() ? d_grid + contrib : std::move(contrib);
+    }
+  } else if (interpolation == GridSamplerInterpolation::Bicubic) {
+    // Native bicubic backward uses raw unnormalized coordinates with a constant
+    // gx_mult/gy_mult (the unnormalize scale only), applying padding
+    // exclusively when sampling each tap value.  Using gs_compute_coords here
+    // would be wrong: for border/reflection modes it zeroes gix_mult at
+    // boundaries, making ggG_x/y = 0 where native still has nonzero
+    // sensitivity.
+    double x_scale = align_corners ? (W - 1) / 2.0 : W / 2.0;
+    double y_scale = align_corners ? (H - 1) / 2.0 : H / 2.0;
+    auto x_raw = align_corners ? (grid.select(-1, 0) + 1) * x_scale
+                               : (grid.select(-1, 0) + 1) * x_scale - 0.5;
+    auto y_raw = align_corners ? (grid.select(-1, 1) + 1) * y_scale
+                               : (grid.select(-1, 1) + 1) * y_scale - 0.5;
+    auto x0_bc = at::floor(x_raw).to(at::kLong);
+    auto y0_bc = at::floor(y_raw).to(at::kLong);
+    auto ggG_x_bc = ggGrid.select(-1, 0) * x_scale;
+    auto ggG_y_bc = ggGrid.select(-1, 1) * y_scale;
+    auto fx = x_raw - x0_bc.to(x_raw.dtype());
+    auto fy = y_raw - y0_bc.to(y_raw.dtype());
+
+    // Cubic interpolation coefficients and derivatives w.r.t. fractional
+    // offset. For t in [0,1], the four corners are at offsets {-1, 0, 1, 2}
+    // from base.
+    constexpr double A = -0.75;
+    auto fx1 = fx + 1.0, fx2 = 1.0 - fx, fx3 = 2.0 - fx;
+    auto fy1 = fy + 1.0, fy2 = 1.0 - fy, fy3 = 2.0 - fy;
+
+    auto cx_t = at::stack(
+        {(((A * fx1) - (5 * A)) * fx1 + (8 * A)) * fx1 - (4 * A),
+         (((A + 2) * fx) - (A + 3)) * fx.square() + 1,
+         (((A + 2) * fx2) - (A + 3)) * fx2.square() + 1,
+         (((A * fx3) - (5 * A)) * fx3 + (8 * A)) * fx3 - (4 * A)},
+        -1);
+    auto dcx_t = at::stack(
+        {(((3 * A) * fx1) - (10 * A)) * fx1 + (8 * A),
+         (((3 * (A + 2)) * fx) - (2 * (A + 3))) * fx,
+         -((((3 * (A + 2)) * fx2) - (2 * (A + 3))) * fx2),
+         (((-3 * A) * fx3) + (10 * A)) * fx3 - (8 * A)},
+        -1);
+    auto cy_t = at::stack(
+        {(((A * fy1) - (5 * A)) * fy1 + (8 * A)) * fy1 - (4 * A),
+         (((A + 2) * fy) - (A + 3)) * fy.square() + 1,
+         (((A + 2) * fy2) - (A + 3)) * fy2.square() + 1,
+         (((A * fy3) - (5 * A)) * fy3 + (8 * A)) * fy3 - (4 * A)},
+        -1);
+    auto dcy_t = at::stack(
+        {(((3 * A) * fy1) - (10 * A)) * fy1 + (8 * A),
+         (((3 * (A + 2)) * fy) - (2 * (A + 3))) * fy,
+         -((((3 * (A + 2)) * fy2) - (2 * (A + 3))) * fy2),
+         (((-3 * A) * fy3) + (10 * A)) * fy3 - (8 * A)},
+        -1);
+
+    // Stack coefficients into [N, Ho, Wo, 4] tensors for vectorized outer
+    // products. Tap indices [N, Ho, Wo, 16]: flat index k=j*4+i → h=y0+(j-1),
+    // w=x0+(i-1)
+    auto offs = at::arange(-1, 3, x0_bc.options());
+    auto xt = x0_bc.unsqueeze(-1) + offs; // [N, Ho, Wo, 4]
+    auto yt = y0_bc.unsqueeze(-1) + offs;
+    // Outer product: h_idx[...,j,i]=yt[j], w_idx[...,j,i]=xt[i]
+    auto w_idx_bc = xt.unsqueeze(-2)
+                        .expand({-1, -1, -1, 4, 4})
+                        .reshape({-1, xt.size(1), xt.size(2), 16});
+    auto h_idx_bc = yt.unsqueeze(-1)
+                        .expand({-1, -1, -1, 4, 4})
+                        .reshape({-1, yt.size(1), yt.size(2), 16});
+
+    // Gather all 16 taps at once: [N, C, Ho, Wo, 16]
+    auto I_all = gs_gather2d_bc_multi(
+        input, h_idx_bc, w_idx_bc, padding_mode_enum, align_corners);
+
+    Tensor B_dx, B_dy;
+    if (output_mask[0] || output_mask[1]) {
+      // 1D weight outer products → [N, Ho, Wo, 16] where k=j*4+i
+      // B_dx[k] = dcx[i]*cy[j],  B_dy[k] = cx[i]*dcy[j]
+      B_dx = gs_outer_last(dcx_t, cy_t);
+      B_dy = gs_outer_last(cx_t, dcy_t);
+    }
+
+    if (output_mask[0]) {
+      auto sum_dx = gs_accum_sumprod_k(I_all, B_dx);
+      auto sum_dy = gs_accum_sumprod_k(I_all, B_dy);
+      auto contrib = sum_dx * ggG_x_bc.unsqueeze(1);
+      contrib.addcmul_(sum_dy, ggG_y_bc.unsqueeze(1));
+      d_grad_output = d_grad_output.defined() ? d_grad_output + contrib
+                                              : std::move(contrib);
+    }
+
+    if (output_mask[1]) {
+      auto w_in = ggG_x_bc.unsqueeze(-1) * B_dx;
+      w_in.addcmul_(ggG_y_bc.unsqueeze(-1), B_dy);
+      d_input = gs_scatter2d_bc_multi(
+          grad_output,
+          w_in,
+          h_idx_bc,
+          w_idx_bc,
+          H,
+          W,
+          padding_mode_enum,
+          align_corners);
+    }
+
+    // d_grid from ggGrid: non-zero for bicubic (second derivative of cubic
+    // weights).
+    if (output_mask[2]) {
+      auto d2cx_t = at::stack(
+          {((6 * A) * fx1) - (10 * A),
+           ((6 * (A + 2)) * fx) - (2 * (A + 3)),
+           ((6 * (A + 2)) * fx2) - (2 * (A + 3)),
+           ((6 * A) * fx3) - (10 * A)},
+          -1);
+      auto d2cy_t = at::stack(
+          {((6 * A) * fy1) - (10 * A),
+           ((6 * (A + 2)) * fy) - (2 * (A + 3)),
+           ((6 * (A + 2)) * fy2) - (2 * (A + 3)),
+           ((6 * A) * fy3) - (10 * A)},
+          -1);
+      Tensor dot_dx2, dot_dy2, dot_dxdy;
+      for (const auto j : c10::irange(4)) {
+        auto cy_j = cy_t.select(-1, j);
+        auto dcy_j = dcy_t.select(-1, j);
+        auto d2cy_j = d2cy_t.select(-1, j);
+        for (const auto i : c10::irange(4)) {
+          auto tap_dot = (grad_output * I_all.select(-1, j * 4 + i)).sum(1);
+          auto dx2_term = tap_dot * d2cx_t.select(-1, i) * cy_j;
+          auto dy2_term = tap_dot * cx_t.select(-1, i) * d2cy_j;
+          auto dxdy_term = tap_dot * dcx_t.select(-1, i) * dcy_j;
+          dot_dx2 =
+              dot_dx2.defined() ? dot_dx2.add_(dx2_term) : std::move(dx2_term);
+          dot_dy2 =
+              dot_dy2.defined() ? dot_dy2.add_(dy2_term) : std::move(dy2_term);
+          dot_dxdy = dot_dxdy.defined() ? dot_dxdy.add_(dxdy_term)
+                                        : std::move(dxdy_term);
+        }
+      }
+      auto ggrid_d_grid_x = ggG_x_bc * dot_dx2;
+      ggrid_d_grid_x.addcmul_(ggG_y_bc, dot_dxdy);
+      ggrid_d_grid_x.mul_(x_scale);
+      auto ggrid_d_grid_y = ggG_x_bc * dot_dxdy;
+      ggrid_d_grid_y.addcmul_(ggG_y_bc, dot_dy2);
+      ggrid_d_grid_y.mul_(y_scale);
+      auto ggrid_d_grid =
+          at::stack({std::move(ggrid_d_grid_x), std::move(ggrid_d_grid_y)}, -1);
+      d_grid =
+          d_grid.defined() ? d_grid + ggrid_d_grid : std::move(ggrid_d_grid);
+    }
+  } else {
+    TORCH_CHECK_NOT_IMPLEMENTED(
+        false,
+        "grid_sampler_2d_double_backward does not support interpolation mode ",
+        interpolation_mode);
+  }
+  return {std::move(d_grad_output), std::move(d_input), std::move(d_grid)};
+}
+
+std::tuple<Tensor, Tensor, Tensor> grid_sampler_3d_double_backward(
+    const Tensor& ggI,
+    const Tensor& ggGrid,
+    const Tensor& grad_output,
+    const Tensor& input,
+    const Tensor& grid,
+    int64_t interpolation_mode,
+    int64_t padding_mode,
+    bool align_corners,
+    std::array<bool, 3> output_mask) {
+  Tensor d_grad_output, d_input, d_grid;
+  const auto interpolation =
+      static_cast<GridSamplerInterpolation>(interpolation_mode);
+  const auto padding_mode_enum = static_cast<GridSamplerPadding>(padding_mode);
+
+  if (output_mask[0] && ggI.defined()) {
+    d_grad_output = at::grid_sampler_3d(
+        ggI, grid, interpolation_mode, padding_mode, align_corners);
+  }
+  if (output_mask[2] && ggI.defined()) {
+    d_grid = std::get<1>(at::grid_sampler_3d_backward(
+        grad_output,
+        ggI,
+        grid,
+        interpolation_mode,
+        padding_mode,
+        align_corners,
+        {false, true}));
+  }
+  if (!ggGrid.defined() || interpolation == GridSamplerInterpolation::Nearest) {
+    return {std::move(d_grad_output), std::move(d_input), std::move(d_grid)};
+  }
+  TORCH_CHECK_NOT_IMPLEMENTED(
+      interpolation == GridSamplerInterpolation::Bilinear,
+      "grid_sampler_3d double backward not implemented for interpolation_mode=",
+      interpolation_mode);
+
+  // Bilinear 3D: ggGrid contributions.
+  auto D = input.size(2), H = input.size(3), W = input.size(4);
+
+  auto [ix, gix_mult] = gs_compute_coords(
+      grid.select(-1, 0), W, padding_mode_enum, align_corners);
+  auto [iy, giy_mult] = gs_compute_coords(
+      grid.select(-1, 1), H, padding_mode_enum, align_corners);
+  auto [iz, giz_mult] = gs_compute_coords(
+      grid.select(-1, 2), D, padding_mode_enum, align_corners);
+
+  auto x0 = at::floor(ix).to(at::kLong);
+  auto y0 = at::floor(iy).to(at::kLong);
+  auto z0 = at::floor(iz).to(at::kLong);
+  auto x1 = x0 + 1, y1 = y0 + 1, z1 = z0 + 1;
+  auto fx = ix - x0.to(ix.dtype());
+  auto fy = iy - y0.to(iy.dtype());
+  auto fz = iz - z0.to(iz.dtype());
+
+  bool zeros_oob = (padding_mode_enum == GridSamplerPadding::Zeros);
+  auto ggG_x = ggGrid.select(-1, 0) * gix_mult;
+  auto ggG_y = ggGrid.select(-1, 1) * giy_mult;
+  auto ggG_z = ggGrid.select(-1, 2) * giz_mult;
+
+  // Stack all 8 tap indices [N, Do, Ho, Wo, 8]: tnw,tne,tsw,tse,bnw,bne,bsw,bse
+  auto d_stk = at::stack({z0, z0, z0, z0, z1, z1, z1, z1}, -1);
+  auto h_stk = at::stack({y0, y0, y1, y1, y0, y0, y1, y1}, -1);
+  auto w_stk = at::stack({x0, x1, x0, x1, x0, x1, x0, x1}, -1);
+
+  Tensor I; // [N, C, Do, Ho, Wo, 8]
+  if (output_mask[0] || output_mask[1] || output_mask[2]) {
+    I = gs_gather3d_multi(input, d_stk, h_stk, w_stk, zeros_oob);
+  }
+
+  Tensor dw_dx, dw_dy, dw_dz;
+  if (output_mask[0] || output_mask[1]) {
+    // Trilinear derivative weights [N, Do, Ho, Wo, 8]
+    auto omx = 1.0 - fx, omy = 1.0 - fy, omz = 1.0 - fz;
+    dw_dx = at::stack(
+        {-omy * omz,
+         omy * omz,
+         -fy * omz,
+         fy * omz,
+         -omy * fz,
+         omy * fz,
+         -fy * fz,
+         fy * fz},
+        -1);
+    dw_dy = at::stack(
+        {-omx * omz,
+         -fx * omz,
+         omx * omz,
+         fx * omz,
+         -omx * fz,
+         -fx * fz,
+         omx * fz,
+         fx * fz},
+        -1);
+    dw_dz = at::stack(
+        {-omx * omy,
+         -fx * omy,
+         -omx * fy,
+         -fx * fy,
+         omx * omy,
+         fx * omy,
+         omx * fy,
+         fx * fy},
+        -1);
+  }
+
+  if (output_mask[0]) {
+    auto sum_dx = gs_accum_sumprod_k(I, dw_dx);
+    auto sum_dy = gs_accum_sumprod_k(I, dw_dy);
+    auto sum_dz = gs_accum_sumprod_k(I, dw_dz);
+    auto d_gO = sum_dx * ggG_x.unsqueeze(1);
+    d_gO.addcmul_(sum_dy, ggG_y.unsqueeze(1));
+    d_gO.addcmul_(sum_dz, ggG_z.unsqueeze(1));
+    d_grad_output =
+        d_grad_output.defined() ? d_grad_output + d_gO : std::move(d_gO);
+  }
+
+  if (output_mask[1]) {
+    auto w_in = ggG_x.unsqueeze(-1) * dw_dx;
+    w_in.addcmul_(ggG_y.unsqueeze(-1), dw_dy);
+    w_in.addcmul_(ggG_z.unsqueeze(-1), dw_dz);
+    d_input = gs_scatter3d_multi(
+        grad_output, w_in, d_stk, h_stk, w_stk, D, H, W, zeros_oob);
+  }
+
+  // ggGrid -> d_grid: cross second derivatives for 3D bilinear.
+  // d²output/(dix diy) = (1-fz)*(I_tnw-I_tne-I_tsw+I_tse) + fz*(I_bnw-...)
+  // d²output/(dix diz) = (1-fy)*(I_tnw-I_tne-I_bnw+I_bne) + fy*(I_tsw-...)
+  // d²output/(diy diz) = (1-fx)*(I_tnw-I_tsw-I_bnw+I_bsw) + fx*(I_tne-...)
+  if (output_mask[2]) {
+    auto fx_ = fx.unsqueeze(1), fy_ = fy.unsqueeze(1), fz_ = fz.unsqueeze(1);
+    auto I_tnw = I.select(-1, 0);
+    auto I_tne = I.select(-1, 1);
+    auto I_tsw = I.select(-1, 2);
+    auto I_tse = I.select(-1, 3);
+    auto I_bnw = I.select(-1, 4);
+    auto I_bne = I.select(-1, 5);
+    auto I_bsw = I.select(-1, 6);
+    auto I_bse = I.select(-1, 7);
+    auto top_xy = I_tnw - I_tne;
+    top_xy.sub_(I_tsw);
+    top_xy.add_(I_tse);
+    auto bottom_xy = I_bnw - I_bne;
+    bottom_xy.sub_(I_bsw);
+    bottom_xy.add_(I_bse);
+    auto d2_xy = (1 - fz_) * top_xy;
+    d2_xy.add_(fz_ * bottom_xy);
+
+    auto top_xz = I_tnw - I_tne;
+    top_xz.sub_(I_bnw);
+    top_xz.add_(I_bne);
+    auto bottom_xz = I_tsw - I_tse;
+    bottom_xz.sub_(I_bsw);
+    bottom_xz.add_(I_bse);
+    auto d2_xz = (1 - fy_) * top_xz;
+    d2_xz.addcmul_(fy_, bottom_xz);
+
+    auto top_yz = I_tnw - I_tsw;
+    top_yz.sub_(I_bnw);
+    top_yz.add_(I_bsw);
+    auto bottom_yz = I_tne - I_tse;
+    bottom_yz.sub_(I_bne);
+    bottom_yz.add_(I_bse);
+    auto d2_yz = (1 - fx_) * top_yz;
+    d2_yz.addcmul_(fx_, bottom_yz);
+    auto dot_xy = (grad_output * d2_xy).sum(1);
+    auto dot_xz = (grad_output * d2_xz).sum(1);
+    auto dot_yz = (grad_output * d2_yz).sum(1);
+    auto d_grid_x = ggG_y * dot_xy;
+    d_grid_x.addcmul_(ggG_z, dot_xz);
+    d_grid_x.mul_(gix_mult);
+    auto d_grid_y = ggG_x * dot_xy;
+    d_grid_y.addcmul_(ggG_z, dot_yz);
+    d_grid_y.mul_(giy_mult);
+    auto d_grid_z = ggG_x * dot_xz;
+    d_grid_z.addcmul_(ggG_y, dot_yz);
+    d_grid_z.mul_(giz_mult);
+    auto contrib = at::stack(
+        {std::move(d_grid_x), std::move(d_grid_y), std::move(d_grid_z)}, -1);
+    d_grid = d_grid.defined() ? d_grid + contrib : std::move(contrib);
+  }
+  return {std::move(d_grad_output), std::move(d_input), std::move(d_grid)};
+}
+
 } // namespace torch::autograd::generated::details
diff --git a/torch/csrc/autograd/FunctionsManual.h b/torch/csrc/autograd/FunctionsManual.h
index 2df0a9af28e50..ecc6bc8bdb5dd 100644
--- a/torch/csrc/autograd/FunctionsManual.h
+++ b/torch/csrc/autograd/FunctionsManual.h
@@ -736,6 +736,26 @@ std::tuple<Tensor, Tensor, Tensor> batchnorm_double_backward(
     const std::optional<Tensor>& save_mean,
     const std::optional<Tensor>& save_invstd,
     std::array<bool, 3> output_mask);
+std::tuple<Tensor, Tensor, Tensor> grid_sampler_2d_double_backward(
+    const Tensor& ggI,
+    const Tensor& ggGrid,
+    const Tensor& grad_output,
+    const Tensor& input,
+    const Tensor& grid,
+    int64_t interpolation_mode,
+    int64_t padding_mode,
+    bool align_corners,
+    std::array<bool, 3> output_mask);
+std::tuple<Tensor, Tensor, Tensor> grid_sampler_3d_double_backward(
+    const Tensor& ggI,
+    const Tensor& ggGrid,
+    const Tensor& grad_output,
+    const Tensor& input,
+    const Tensor& grid,
+    int64_t interpolation_mode,
+    int64_t padding_mode,
+    bool align_corners,
+    std::array<bool, 3> output_mask);
 std::tuple<Tensor, Tensor> _euclidean_dist_backward(
     const Tensor& grad,
     const Tensor& x1,
@@ -820,6 +840,20 @@ Tensor amaxamin_jvp(
     const Tensor& result,
     IntArrayRef dim,
     bool keepdim);
+Tensor aminmax_backward(
+    const at::Tensor& self,
+    std::optional<int64_t> dim,
+    bool keepdim,
+    const at::Tensor& grad_min,
+    const at::Tensor& grad_max,
+    const at::Tensor& min,
+    const at::Tensor& max);
+Tensor aminmax_jvp(
+    const Tensor& self_p,
+    const Tensor& self_t,
+    const Tensor& result,
+    std::optional<int64_t> dim,
+    bool keepdim);
 std::tuple<Tensor, Tensor, Tensor> layer_norm_double_backward(
     const Tensor& input,
     const std::optional<Tensor>& gamma,
diff --git a/torch/csrc/autograd/VariableTypeManual.cpp b/torch/csrc/autograd/VariableTypeManual.cpp
index c2c4dffee66eb..76dda79d4185f 100644
--- a/torch/csrc/autograd/VariableTypeManual.cpp
+++ b/torch/csrc/autograd/VariableTypeManual.cpp
@@ -479,11 +479,10 @@ static Tensor _fw_primal(
   std::function<at::Tensor(const at::Tensor&)> rev_func = nullptr;
   if (!self.unsafeGetTensorImpl()->support_as_strided()) {
     func = std::make_unique<ViewViewFunc>(self.sym_sizes());
-    rev_func = [=](const at::Tensor& input_view) {
+    rev_func = [=](const at::Tensor& input_view) -> at::Tensor {
       TORCH_INTERNAL_ASSERT(
           false,
           "Reverse view_func for _fw_primal() is not currently supported");
-      return Tensor();
     };
   }
   auto result = as_view(
@@ -512,11 +511,10 @@ static Tensor _make_dual(
   std::function<at::Tensor(const at::Tensor&)> rev_func = nullptr;
   if (!primal.unsafeGetTensorImpl()->support_as_strided()) {
     func = std::make_unique<ViewViewFunc>(primal.sym_sizes());
-    rev_func = [=](const at::Tensor& input_view) {
+    rev_func = [=](const at::Tensor& input_view) -> at::Tensor {
       TORCH_INTERNAL_ASSERT(
           false,
           "Reverse view_func for _make_dual() is not currently supported");
-      return Tensor();
     };
   }
   auto result = as_view(
diff --git a/torch/csrc/autograd/autograd_not_implemented_fallback.cpp b/torch/csrc/autograd/autograd_not_implemented_fallback.cpp
index 386a8a9df534d..efe91535a5a36 100644
--- a/torch/csrc/autograd/autograd_not_implemented_fallback.cpp
+++ b/torch/csrc/autograd/autograd_not_implemented_fallback.cpp
@@ -618,7 +618,7 @@ static void autogradNotImplementedInplaceOrViewFallbackImpl(
       "input and the first output (the output can be a vector of tensors). Please change the "
       "order of your operator's parameters so that this is the case.");
   const bool is_view = aliased_input_idx.has_value();
-  size_t aliased_input_idx_val;
+  size_t aliased_input_idx_val = 0;
 
   // Save inputs before we redispatch down
   torch::jit::Stack non_tensor_stack;
@@ -657,13 +657,13 @@ static void autogradNotImplementedInplaceOrViewFallbackImpl(
          "which does not have a derivative implemented is forbidden.");
     auto erroring_view_func = std::make_unique<ErroringViewFunc>(error_msg);
 
-    const auto erroring_rev_view_func = [op_name = op_name](const at::Tensor&) {
+    const auto erroring_rev_view_func =
+        [op_name = op_name](const at::Tensor&) -> at::Tensor {
       TORCH_CHECK(
           false,
           "Accessing the reverse view for ",
           op_name,
           " which does not have a derivative implemented is forbidden.");
-      return at::Tensor();
     };
 
     if (aliased_output_iv.isTensorList()) {
diff --git a/torch/csrc/autograd/functions/tensor.cpp b/torch/csrc/autograd/functions/tensor.cpp
index ca57f6488ce1a..d03e673f108fb 100644
--- a/torch/csrc/autograd/functions/tensor.cpp
+++ b/torch/csrc/autograd/functions/tensor.cpp
@@ -254,6 +254,7 @@ variable_list CopySlices::apply_with_saved(
     update_exec_info();
 
     std::vector<bool> needs_input_grad;
+    needs_input_grad.reserve(num_outputs());
     for (const auto i : c10::irange(num_outputs())) {
       needs_input_grad.emplace_back(task_should_compute_output(i));
     }
diff --git a/torch/csrc/autograd/init.cpp b/torch/csrc/autograd/init.cpp
index 28c1a4308bf99..504b02bebd79a 100644
--- a/torch/csrc/autograd/init.cpp
+++ b/torch/csrc/autograd/init.cpp
@@ -293,6 +293,10 @@ PyObject* THPAutograd_initExtension(PyObject* _unused, PyObject* unused) {
       .def(
           "linked_correlation_id",
           [](const KinetoEvent& e) { return e.linkedCorrelationId(); })
+      .def("flow_id", [](const KinetoEvent& e) { return e.flowId(); })
+      .def("flow_type", [](const KinetoEvent& e) { return e.flowType(); })
+      .def("flow_start", [](const KinetoEvent& e) { return e.flowStart(); })
+      .def("external_id", [](const KinetoEvent& e) { return e.externalId(); })
       // compute flops
       .def("flops", [](const KinetoEvent& e) { return e.flops(); })
       // Whether this is async event or not
@@ -307,14 +311,58 @@ PyObject* THPAutograd_initExtension(PyObject* _unused, PyObject* unused) {
                 e.activityType() ==
                 (uint8_t)libkineto::ActivityType::GPU_USER_ANNOTATION;
           })
+      .def(
+          "is_python_function",
+          [](const KinetoEvent& e) { return e.isPythonFunction(); })
       .def("nbytes", [](const KinetoEvent& e) { return e.nBytes(); })
       // whether the event is hidden
       .def(
           "is_hidden_event",
           [](const KinetoEvent& e) { return e.isHiddenEvent(); })
       // KinetoEvent metadata
-      .def("metadata_json", [](const KinetoEvent& e) {
-        return e.metadataJson();
+      .def(
+          "metadata_json",
+          [](const KinetoEvent& e) { return e.metadataJson(); })
+      .def(
+          "activity_type",
+          [](const KinetoEvent& e) {
+            return libkineto::toString(
+                static_cast<libkineto::ActivityType>(e.activityType()));
+          })
+      .def("extra_meta", [](const KinetoEvent& e) { return e.extraMeta(); })
+      // Like shapes/strides, but also contains TensorList input shapes.
+      .def(
+          "structured_input_shapes",
+          [](const KinetoEvent& e) {
+            py::list result;
+            for (const auto& s : e.structuredInputShapes()) {
+              if (std::holds_alternative<std::vector<int64_t>>(s)) {
+                result.append(std::get<std::vector<int64_t>>(s));
+              } else {
+                result.append(std::get<std::vector<std::vector<int64_t>>>(s));
+              }
+            }
+            return result;
+          })
+      .def(
+          "structured_input_strides",
+          [](const KinetoEvent& e) {
+            py::list result;
+            for (const auto& s : e.structuredInputStrides()) {
+              if (std::holds_alternative<std::vector<int64_t>>(s)) {
+                result.append(std::get<std::vector<int64_t>>(s));
+              } else {
+                result.append(std::get<std::vector<std::vector<int64_t>>>(s));
+              }
+            }
+            return result;
+          })
+      .def("python_id", [](const KinetoEvent& e) { return e.pythonId(); })
+      .def(
+          "python_parent_id",
+          [](const KinetoEvent& e) { return e.pythonParentId(); })
+      .def("python_module_id", [](const KinetoEvent& e) {
+        return e.pythonModuleId();
       });
 
   m.def("_soft_assert_raises", &setSoftAssertRaises);
@@ -339,6 +387,9 @@ PyObject* THPAutograd_initExtension(PyObject* _unused, PyObject* unused) {
   m.def(
       "_prepare_profiler",
       prepareProfiler,
+      py::arg("config"),
+      py::arg("activities"),
+      py::arg("activity_filter") = torch::autograd::profiler::ActivityFilter{},
       py::call_guard<py::gil_scoped_release>());
   m.def(
       "_toggle_collection_dynamic",
@@ -582,6 +633,11 @@ PyObject* THPAutograd_initExtension(PyObject* _unused, PyObject* unused) {
             false,
             "Trying to create a SavedTensor object from Python is forbidden.");
       }))
+      .def(
+          "unpack",
+          [](const torch::autograd::SavedVariable& s) -> at::Tensor {
+            return s.unpack();
+          })
       .def(
           "register_hooks",
           [](torch::autograd::SavedVariable& s,
@@ -621,6 +677,18 @@ PyObject* THPAutograd_initExtension(PyObject* _unused, PyObject* unused) {
             return py::reinterpret_borrow<py::function>(unpack_ptr);
           });
 
+  m.def(
+      "_make_saved_tensor",
+      [](const at::Tensor& tensor,
+         bool is_output,
+         bool is_inplace_on_view) -> torch::autograd::SavedVariable {
+        return torch::autograd::SavedVariable(
+            tensor, is_output, is_inplace_on_view);
+      },
+      py::arg("tensor"),
+      py::arg("is_output"),
+      py::arg("is_inplace_on_view") = false);
+
   torch::autograd::profiler::python_tracer::init();
   Py_RETURN_TRUE;
 }
diff --git a/torch/csrc/autograd/profiler_kineto.cpp b/torch/csrc/autograd/profiler_kineto.cpp
index 0ce905fff7e71..2c952a9005c21 100644
--- a/torch/csrc/autograd/profiler_kineto.cpp
+++ b/torch/csrc/autograd/profiler_kineto.cpp
@@ -27,6 +27,7 @@
 #include <ApproximateClock.h>
 #include <libkineto.h>
 #include <time_since_epoch.h>
+#include <torch/csrc/profiler/standalone/privateuse1_profiler.h>
 
 #ifndef _MSC_VER
 // TODO: TO be removed, once this properly works from libkineto
@@ -78,7 +79,8 @@ using torch::profiler::impl::variantShapesToStr;
 inline bool isKinetoCompatibleState(ProfilerState state) {
   return state == ProfilerState::KINETO ||
       state == ProfilerState::KINETO_GPU_FALLBACK ||
-      state == ProfilerState::KINETO_PRIVATEUSE1_FALLBACK;
+      state == ProfilerState::KINETO_PRIVATEUSE1_FALLBACK ||
+      state == ProfilerState::KINETO_PRIVATEUSE1;
 }
 
 // Helper function to check if ProfilerState is valid for disabling profiler
@@ -267,10 +269,10 @@ struct AddGenericMetadata : public MetadataBase {
     if (arg_data.hasData) {
       if (get_record_concrete_inputs_enabled()) {
         addMetadata("Input Dims", variantShapesToStr(arg_data.shapes));
-        addMetadata("Input Strides", variantShapesToStr(arg_data.strides));
       } else {
         addMetadata("Input Dims", shapesToStr(arg_data.shapesForKinetoEvent));
       }
+      addMetadata("Input Strides", variantShapesToStr(arg_data.strides));
       addMetadata("Input type", strListToStr(arg_data.dtypes));
       if (!arg_data.concreteInputs.empty()) {
         addMetadata(
@@ -461,16 +463,7 @@ struct KinetoThreadLocalState : public ProfilerStateBase {
     auto records_and_trace =
         recordQueue.getRecords(std::move(converter), startTime, end_time);
 
-    materializeOpEvents(records_and_trace.first);
-
-    // `kinetoEvents` does not include Python events. Instead it exposes them
-    // via the `stacks` property.
-    kinetoEvents.erase(
-        std::remove_if(
-            kinetoEvents.begin(),
-            kinetoEvents.end(),
-            [](const auto& i) { return i.isPythonFunction(); }),
-        kinetoEvents.end());
+    materializeOpEvents(records_and_trace.first, end_time);
 
     return std::move(records_and_trace.second);
   }
@@ -482,25 +475,34 @@ struct KinetoThreadLocalState : public ProfilerStateBase {
     }
   }
 
-  void materializeOpEvents(std::vector<std::shared_ptr<Result>>& events) {
+  void materializeOpEvents(
+      std::vector<std::shared_ptr<Result>>& events,
+      int64_t trace_end_ns) {
     for (auto& e : events) {
       if (e->parent_.expired() && e->deviceType() == c10::DeviceType::CPU) {
         eventTree.push_back(e);
       }
 
-      if (e->finished_) {
+      // Unfinished events automatically have end time set to trace end time
+      if (!e->finished_) {
         e->visit(c10::overloaded(
-            [this](ExtraFields<EventType::TorchOp>& i) { invokeCallback(i); },
-            [this](ExtraFields<EventType::Backend>& i) { invokeCallback(i); },
+            [trace_end_ns](ExtraFields<EventType::TorchOp>& i) {
+              i.end_time_ns_ = trace_end_ns;
+            },
             [](auto&) {}));
+      }
 
-        kinetoEvents.emplace_back(e, config_.experimental_config.verbose);
-        AddTensorboardFields add_tb(e, kinetoEvents.back());
-        AddGenericMetadata add_generic(e, &config_);
+      e->visit(c10::overloaded(
+          [this](ExtraFields<EventType::TorchOp>& i) { invokeCallback(i); },
+          [this](ExtraFields<EventType::Backend>& i) { invokeCallback(i); },
+          [](auto&) {}));
 
-        // It is not safe to use the activity after post processing.
-        e->kineto_activity_ = nullptr;
-      }
+      kinetoEvents.emplace_back(e, config_.experimental_config.verbose);
+      AddTensorboardFields add_tb(e, kinetoEvents.back());
+      AddGenericMetadata add_generic(e, &config_);
+
+      // It is not safe to use the activity after post processing.
+      e->kineto_activity_ = nullptr;
     }
   }
 
@@ -640,11 +642,23 @@ void reportBackendEventToActiveKinetoProfiler(
 
 void prepareProfiler(
     const torch::profiler::impl::ProfilerConfig& config,
-    const std::set<torch::profiler::impl::ActivityType>& activities) {
+    const std::set<torch::profiler::impl::ActivityType>& activities,
+    const ActivityFilter& activity_filter) {
   if (config.state == ProfilerState::NVTX ||
       config.state == ProfilerState::ITT) {
     return;
   }
+
+  // Forward registered PrivateUse1 profiler factory to Kineto.
+  // Only for KINETO_PRIVATEUSE1 state where backend provides its own
+  // IActivityProfiler.
+#ifdef USE_KINETO
+  if (config.state == ProfilerState::KINETO_PRIVATEUSE1) {
+    torch::profiler::impl::PrivateUse1ProfilerRegistry::instance()
+        .onKinetoInit();
+  }
+#endif // USE_KINETO
+
   TORCH_CHECK(
       isKinetoCompatibleState(config.state),
       "Supported only in Kineto profiler");
@@ -654,7 +668,8 @@ void prepareProfiler(
           c10::get_privateuse1_backend() != "privateuseone"),
       activities,
       config.experimental_config,
-      config.trace_id);
+      config.trace_id,
+      activity_filter);
 
   if (!config.experimental_config.performance_events.empty()) {
     /* For now only CPU activity is supported */
@@ -861,10 +876,16 @@ std::unique_ptr<ProfilerResult> disableProfiler() {
   profiler_state_info_ptr = nullptr;
 
   auto state_ptr = ProfilerStateBase::pop();
+  if (!state_ptr) {
+    LOG(WARNING)
+        << "disableProfiler called but no active profiling session found. "
+        << "This can happen if profiling was cancelled during warmup.";
+    return std::make_unique<ProfilerResult>();
+  }
   const auto& config = state_ptr->config();
   TORCH_CHECK(
-      state_ptr && isValidDisableState(config.state),
-      "Can't disable Kineto profiler when it's not running");
+      isValidDisableState(config.state),
+      "Can't disable Kineto profiler: config is not in a valid disable state");
 
   state_ptr->removeCallback();
 
@@ -932,6 +953,8 @@ KinetoEvent::KinetoEvent(
   result->visit_if_base<ExtraFields<EventType::TorchOp>>([&](const auto& op) {
     auto arg_data = parseArgData(op.inputs_, op.concrete_inputs_);
     shapes_ = std::move(arg_data.shapesForKinetoEvent);
+    structured_input_shapes_ = std::move(arg_data.shapes);
+    structured_input_strides_ = std::move(arg_data.strides);
     dtypes_ = std::move(arg_data.dtypes);
     concrete_inputs_ = std::move(arg_data.concreteInputs);
     kwinputs_ = std::move(op.kwinputs_);
@@ -944,6 +967,45 @@ bool KinetoEvent::isPythonFunction() const {
   return out;
 }
 
+int64_t KinetoEvent::pythonId() const {
+  int64_t out{-1};
+  result_->visit_if_base<PyExtraFieldsBase>(
+      [&](const auto& i) { out = static_cast<int64_t>(i.id_); });
+  return out;
+}
+
+int64_t KinetoEvent::pythonParentId() const {
+  int64_t out{-1};
+  // Walk the python parent pointers up to find the next event of type
+  // PyExtraFieldsBase
+  result_->visit_if_base<PyExtraFieldsBase>([&](const auto&) {
+    auto parent = result_->parent_.lock();
+    while (parent) {
+      parent->visit_if_base<PyExtraFieldsBase>(
+          [&](const auto& j) { out = static_cast<int64_t>(j.id_); });
+      if (out >= 0) {
+        break;
+      }
+      parent = parent->parent_.lock();
+    }
+  });
+  return out;
+}
+
+int64_t KinetoEvent::pythonModuleId() const {
+  int64_t out{-1};
+  // Returns the module id for PyCall events (python function calls to
+  // nn.Module)
+  result_->visit(c10::overloaded(
+      [&](const ExtraFields<EventType::PyCall>& py_call) {
+        if (py_call.module_.has_value()) {
+          out = static_cast<int64_t>(py_call.module_->id_);
+        }
+      },
+      [](const auto&) {}));
+  return out;
+}
+
 bool KinetoEvent::hasShapes() const {
   return !shapes_.empty();
 }
@@ -952,6 +1014,16 @@ const c10::ArrayRef<std::vector<int64_t>> KinetoEvent::shapes() const {
   return shapes_;
 }
 
+const c10::ArrayRef<torch::profiler::impl::shape> KinetoEvent::
+    structuredInputShapes() const {
+  return structured_input_shapes_;
+}
+
+const c10::ArrayRef<torch::profiler::impl::shape> KinetoEvent::
+    structuredInputStrides() const {
+  return structured_input_strides_;
+}
+
 bool KinetoEvent::hasTypes() const {
   return !dtypes_.empty();
 }
@@ -1063,7 +1135,6 @@ int64_t KinetoEvent::privateuse1ElapsedUs() const {
   }
   return (int64_t)torch::profiler::impl::privateuse1Stubs()->elapsed(
       &privateuse1_event_start, &privateuse1_event_end);
-  return -1;
 }
 
 void KinetoEvent::getPerfEventCounters(std::vector<uint64_t>& in) const {
@@ -1092,6 +1163,40 @@ std::string KinetoEvent::metadataJson() const {
       [](const auto&) -> std::string { return std::string(""); }));
 }
 
+int64_t KinetoEvent::externalId() const {
+  // Mirrors libkineto::ChromeTraceLogger::handleActivity() "External id" logic.
+  // libkineto::ChromeTraceLogger checks op.linkedActivity() != nullptr; here we
+  // check linkedCorrelationId() > 0, which is equivalent because PyTorch
+  // correlation IDs are monotonically increasing from 1 (a valid linked
+  // activity always has a non-zero correlation ID).
+  uint64_t linked = linkedCorrelationId();
+  if (linked > 0) {
+    return static_cast<int64_t>(linked);
+  }
+
+  // Orphaned GPU activities (no linked CPU op) in these types should not get
+  // an External id, to avoid incorrect cross-linking in trace viewers.
+  auto type = static_cast<libkineto::ActivityType>(activityType());
+  if (type != libkineto::ActivityType::GPU_MEMCPY &&
+      type != libkineto::ActivityType::GPU_MEMSET &&
+      type != libkineto::ActivityType::CONCURRENT_KERNEL &&
+      type != libkineto::ActivityType::CUDA_RUNTIME &&
+      type != libkineto::ActivityType::CUDA_DRIVER &&
+      type != libkineto::ActivityType::PRIVATEUSE1_RUNTIME &&
+      type != libkineto::ActivityType::PRIVATEUSE1_DRIVER) {
+    return static_cast<int64_t>(result_->visit(c10::overloaded(
+        [](const ExtraFields<EventType::TorchOp>& e) -> uint64_t {
+          return e.correlation_id_;
+        },
+        [](const ExtraFields<EventType::Kineto>& e) -> uint64_t {
+          return e.correlation_id_;
+        },
+        [](const auto&) -> uint64_t { return 0; })));
+  }
+
+  return 0;
+}
+
 #define FORWARD_FROM_RESULT(method_name, result_expr)                        \
   decltype(std::declval<KinetoEvent>().method_name())                        \
   KinetoEvent::method_name() const {                                         \
@@ -1133,7 +1238,16 @@ TYPED_ATTR(TorchOp, fwdThreadId, e.sequence_number_ >= 0 ? e.forward_tid_ : 0)
 TYPED_ATTR(TorchOp, scope, static_cast<uint8_t>(e.scope_))
 TYPED_ATTR(TorchOp, hasModuleHierarchy, !e.jit_modules_.empty())
 TYPED_ATTR(TorchOp, isAsync, e.is_async_)
-TYPED_ATTR(TorchOp, extraMeta, e.extra_meta_)
+
+extra_meta_t KinetoEvent::extraMeta() const {
+  extra_meta_t out;
+  result_->visit(c10::overloaded(
+      [&](const ExtraFields<EventType::TorchOp>& e) { out = e.extra_meta_; },
+      [&](const ExtraFields<EventType::Kineto>& e) { out = e.extra_meta_; },
+      [](const auto&) {}));
+  return out;
+}
+
 TYPED_ATTR(TorchOp, fallbackStart, e.device_fallback_.device_event_start_)
 TYPED_ATTR(TorchOp, fallbackEnd, e.device_fallback_.device_event_end_)
 TYPED_ATTR(
@@ -1151,6 +1265,30 @@ TYPED_ATTR(Kineto, linkedCorrelationId, [&]() {
 #undef TYPED_ATTR
 #undef TYPED_ATTR_WITH_DEFAULT
 
+// Flow fields exist on both TorchOp and Kineto event types.
+uint32_t KinetoEvent::flowId() const {
+  return result_->visit(c10::overloaded(
+      [](const ExtraFields<EventType::TorchOp>& e) { return e.flow.id; },
+      [](const ExtraFields<EventType::Kineto>& e) { return e.flow.id; },
+      [](const auto&) -> uint32_t { return 0; }));
+}
+uint32_t KinetoEvent::flowType() const {
+  return result_->visit(c10::overloaded(
+      [](const ExtraFields<EventType::TorchOp>& e) { return e.flow.type; },
+      [](const ExtraFields<EventType::Kineto>& e) { return e.flow.type; },
+      [](const auto&) -> uint32_t { return 0; }));
+}
+bool KinetoEvent::flowStart() const {
+  return result_->visit(c10::overloaded(
+      [](const ExtraFields<EventType::TorchOp>& e) {
+        return static_cast<bool>(e.flow.start);
+      },
+      [](const ExtraFields<EventType::Kineto>& e) {
+        return static_cast<bool>(e.flow.start);
+      },
+      [](const auto&) { return false; }));
+}
+
 ProfilerResult::ProfilerResult(
     uint64_t start_time,
     std::vector<KinetoEvent> events,
diff --git a/torch/csrc/autograd/profiler_kineto.h b/torch/csrc/autograd/profiler_kineto.h
index dbb4febce78b9..4c7db2cae8aa7 100644
--- a/torch/csrc/autograd/profiler_kineto.h
+++ b/torch/csrc/autograd/profiler_kineto.h
@@ -58,6 +58,10 @@ struct TORCH_API KinetoEvent {
   bool isAsync() const;
   uint64_t correlationId() const;
   uint64_t linkedCorrelationId() const;
+  uint32_t flowId() const;
+  uint32_t flowType() const;
+  bool flowStart() const;
+  int64_t externalId() const;
   int64_t deviceResourceId() const;
   std::string backend() const;
   bool isPythonFunction() const;
@@ -67,6 +71,14 @@ struct TORCH_API KinetoEvent {
   extra_meta_t extraMeta() const;
   std::string metadataJson() const;
 
+  const c10::ArrayRef<torch::profiler::impl::shape> structuredInputShapes()
+      const;
+  const c10::ArrayRef<torch::profiler::impl::shape> structuredInputStrides()
+      const;
+  int64_t pythonId() const;
+  int64_t pythonParentId() const;
+  int64_t pythonModuleId() const;
+
  private:
   torch::profiler::impl::ProfilerVoidEventStub fallbackStart() const;
   torch::profiler::impl::ProfilerVoidEventStub fallbackEnd() const;
@@ -79,6 +91,8 @@ struct TORCH_API KinetoEvent {
   std::vector<std::string> dtypes_;
   std::vector<c10::IValue> concrete_inputs_;
   std::unordered_map<std::string, c10::IValue> kwinputs_;
+  std::vector<torch::profiler::impl::shape> structured_input_shapes_;
+  std::vector<torch::profiler::impl::shape> structured_input_strides_;
 };
 
 // Consolidating events returned directly from Kineto
@@ -179,9 +193,13 @@ TORCH_API void enableProfilerWithEventPostProcess(
 
 TORCH_API std::unique_ptr<ProfilerResult> disableProfiler();
 
+using ActivityFilter = std::unordered_map<
+    torch::profiler::impl::ActivityType,
+    std::unordered_set<std::string>>;
 TORCH_API void prepareProfiler(
     const torch::profiler::impl::ProfilerConfig& config,
-    const std::set<torch::profiler::impl::ActivityType>& activities);
+    const std::set<torch::profiler::impl::ActivityType>& activities,
+    const ActivityFilter& activity_filter = {});
 
 TORCH_API void toggleCollectionDynamic(
     const bool enable,
diff --git a/torch/csrc/autograd/profiler_python.cpp b/torch/csrc/autograd/profiler_python.cpp
index a45935ecb2995..9e1c2e8e18118 100644
--- a/torch/csrc/autograd/profiler_python.cpp
+++ b/torch/csrc/autograd/profiler_python.cpp
@@ -7,6 +7,7 @@
 #include <memory>
 #include <queue>
 #include <string>
+#include <unordered_map>
 #include <utility>
 #include <vector>
 
@@ -18,6 +19,8 @@
 #include <c10/util/ApproximateClock.h>
 #include <c10/util/Exception.h>
 #include <c10/util/Logging.h>
+#include <c10/util/ScopeExit.h>
+#include <c10/util/Semaphore.h>
 #include <c10/util/flat_hash_map.h>
 #include <c10/util/irange.h>
 #include <torch/csrc/autograd/python_variable.h>
@@ -551,13 +554,13 @@ struct TraceKeyCacheState {
 // ============================================================================
 // == Core CPython data types =================================================
 // ============================================================================
-// PyObject that allows different threads to record events without colliding.
-// It is passed as the second argument when enabling tracing via
-// `PyEval_SetProfile`.
-struct ThreadLocalResults;
+// PyObject passed as the second argument when enabling tracing via
+// `PyEval_SetProfile`. A single shared instance is used for all threads;
+// the callback resolves per-thread state via PyThreadState_Get().
+class PythonTracer;
 struct TraceContext {
   PyObject_HEAD
-  ThreadLocalResults* thread_local_results_;
+  PythonTracer* tracer_;
 };
 
 // CPython boilerplate to define `TraceContext` as a proper python object.
@@ -604,43 +607,39 @@ static PyTypeObject TraceContextType = {
     nullptr /* tp_free */
 };
 
-class gil_and_restore_thread {
+#if IS_PYTHON_3_14_PLUS
+extern "C" void _PyEval_StopTheWorld(PyInterpreterState*);
+extern "C" void _PyEval_StartTheWorld(PyInterpreterState*);
+
+class StopTheWorldGuard {
  public:
-  gil_and_restore_thread() : initial_thread_state_{PyThreadState_Get()} {}
-  ~gil_and_restore_thread() {
-    PyThreadState_Swap(initial_thread_state_);
-
-    // `gil_scoped_acquire` is a bit fragile in on-demand mode:
-    // https://github.com/pytorch/pytorch/pull/91684#issuecomment-1413154458
-    if (!Py_IsInitialized()) {
-      gil_.disarm();
-    }
+  explicit StopTheWorldGuard(PyInterpreterState* interp) : interp_(interp) {
+    _PyEval_StopTheWorld(interp_);
   }
-
-  PyThreadState* initial_thread_state() const {
-    return initial_thread_state_;
+  ~StopTheWorldGuard() {
+    _PyEval_StartTheWorld(interp_);
   }
+  StopTheWorldGuard(const StopTheWorldGuard&) = delete;
+  StopTheWorldGuard& operator=(const StopTheWorldGuard&) = delete;
 
  private:
-  pybind11::gil_scoped_acquire gil_;
-  PyThreadState* initial_thread_state_;
+  PyInterpreterState* interp_;
 };
+#else
+class StopTheWorldGuard {
+ public:
+  explicit StopTheWorldGuard(PyInterpreterState*) {}
+  StopTheWorldGuard(const StopTheWorldGuard&) = delete;
+  StopTheWorldGuard& operator=(const StopTheWorldGuard&) = delete;
+};
+#endif
 
 // ============================================================================
 // == Thread local cache ======================================================
 // ============================================================================
-class PythonTracer;
 struct ThreadLocalResults {
-  ThreadLocalResults(
-      PyThreadState* thread_state,
-      ValueCache* value_cache,
-      PythonTracer* active_tracer)
-      : thread_state_{thread_state},
-        ctx_{(TraceContext*)TraceContextType.tp_alloc(&TraceContextType, 0)},
-        value_cache_{value_cache},
-        active_tracer_{active_tracer} {
-    ctx_->thread_local_results_ = this;
-  }
+  ThreadLocalResults(PythonTracer* active_tracer)
+      : active_tracer_{active_tracer} {}
 
   ThreadLocalResults() = delete;
   ThreadLocalResults(const ThreadLocalResults&) = delete;
@@ -648,30 +647,18 @@ struct ThreadLocalResults {
   ThreadLocalResults& operator=(const ThreadLocalResults&) = delete;
   ThreadLocalResults& operator=(const ThreadLocalResults&&) = delete;
 
-  ~ThreadLocalResults() {
-    // Currently, there is a bug in Profiler when using Python 3.12 that causes
-    // a segfault when decrementing the refcount of a TraceContext during
-    // on-demand. We are purposefully allowing for a small leak in this
-    // situation to avoid the segfault. This should be fixed in the future.
-#if PY_MAJOR_VERSION < 3 || (PY_MAJOR_VERSION == 3 && PY_MINOR_VERSION < 12)
-    Py_DECREF((PyObject*)ctx_);
-#endif
-  }
-
   template <CallType C, EventType E, typename Ephemeral, typename... Args>
   TraceKey intern(Ephemeral ephemeral, Args... args) {
     static_assert(
         Config<C>::event_type == E,
         "ThreadLocalResults.intern called from the wrong typed context.");
     auto callsite = Callsite<C>(std::forward<Args>(args)...);
-    return std::get<C>(trace_keys_).intern(callsite, ephemeral, *value_cache_);
+    return std::get<C>(trace_keys_).intern(callsite, ephemeral, value_cache_);
   }
 
   static constexpr size_t BLOCK_SIZE = 1024;
 
-  PyThreadState* thread_state_;
-  TraceContext* ctx_;
-  ValueCache* value_cache_;
+  ValueCache value_cache_;
   PythonTracer* active_tracer_;
   CallTypeHelper<TraceKeyCacheState>::tuple_type trace_keys_;
   AppendOnlyList<c10::approx_time_t, BLOCK_SIZE> exit_times_;
@@ -679,6 +666,12 @@ struct ThreadLocalResults {
 
   int active_frames_{0};
   int remaining_start_frames_{0};
+
+  // Guards against teardown racing with in-flight callbacks.
+  // pyProfileFn acquires this on entry and releases on exit.
+  // PythonTracer::stop() acquires each thread's semaphore after
+  // clearing the profiling callback to ensure all callbacks have finished.
+  c10::Semaphore profile_sem{1};
 };
 
 // ============================================================================
@@ -733,7 +726,10 @@ class PythonTracer final : public python_tracer::PythonTracerBase {
       PyObject* arg,
       bool start_frame = false);
 
+  ThreadLocalResults* findThreadLocalResults(PyThreadState* tstate) const;
+
   const std::vector<PyThreadState*> interpreterThreads() const;
+  void setprofileAllThreads(Py_tracefunc func, PyObject* arg) const;
 
   std::atomic<bool> active_lock_{false};
   bool active_{false};
@@ -743,10 +739,12 @@ class PythonTracer final : public python_tracer::PythonTracerBase {
   PyInterpreterState* interpreter_{nullptr};
   PyCodeObject* module_call_code_;
   PyCodeObject* optimizer_hook_;
+  TraceContext* shared_ctx_{nullptr};
 
   std::vector<StartFrame> start_frames_;
   std::deque<ThreadLocalResults> thread_local_results_;
-  ValueCache value_cache_;
+  std::unordered_map<PyThreadState*, ThreadLocalResults*>
+      thread_local_results_map_;
 
 #if IS_PYTHON_3_12
   friend PyObject* c_call_callback(
@@ -833,12 +831,14 @@ static PyObject* c_call_callback(
             PyExc_SystemError, "Missing frame when calling profile function.");
         return NULL;
       }
-      Py_INCREF(frame);
-      auto& local_results =
-          *reinterpret_cast<TraceContext*>(tstate->c_profileobj)
-               ->thread_local_results_;
-      local_results.active_tracer_->recordCCall(local_results, frame, func);
-      Py_DECREF(frame);
+      auto* tracer =
+          reinterpret_cast<TraceContext*>(tstate->c_profileobj)->tracer_;
+      auto* local_results = tracer->findThreadLocalResults(tstate);
+      if (local_results) {
+        Py_INCREF(frame);
+        local_results->active_tracer_->recordCCall(*local_results, frame, func);
+        Py_DECREF(frame);
+      }
     }
   }
   Py_RETURN_NONE;
@@ -965,6 +965,12 @@ static void unregisterMonitoringCallback() {
 }
 #endif
 
+ThreadLocalResults* PythonTracer::findThreadLocalResults(
+    PyThreadState* tstate) const {
+  auto it = thread_local_results_map_.find(tstate);
+  return it != thread_local_results_map_.end() ? it->second : nullptr;
+}
+
 const std::vector<PyThreadState*> PythonTracer::interpreterThreads() const {
   pybind11::gil_scoped_acquire gil;
   std::vector<PyThreadState*> out;
@@ -978,6 +984,19 @@ const std::vector<PyThreadState*> PythonTracer::interpreterThreads() const {
   return out;
 }
 
+void PythonTracer::setprofileAllThreads(Py_tracefunc func, PyObject* arg)
+    const {
+#if IS_PYTHON_3_13_PLUS
+  PyEval_SetProfileAllThreads(func, arg);
+#else
+  for (const auto thread_state : interpreterThreads()) {
+    if (_PyEval_SetProfile(thread_state, func, arg) < 0) {
+      PyErr_WriteUnraisable(nullptr);
+    }
+  }
+#endif
+}
+
 // we are only registering on main thread while holding GIL so this should be
 // safe
 static PyObject* py_gc_callback = nullptr;
@@ -1015,56 +1034,65 @@ PythonTracer::PythonTracer(torch::profiler::impl::RecordQueue* queue)
     return;
   }
 
-  gil_and_restore_thread gil;
-  interpreter_ = PyInterpreterState_Get();
-
-  if (!gil.initial_thread_state()) {
-    TORCH_WARN("PyThreadState_Get returned NULL");
-    return;
-  }
-
-  // Register the tracer in each thread.
-  for (const auto thread_state : interpreterThreads()) {
-    PyThreadState_Swap(thread_state);
-
-    thread_local_results_.emplace_back(thread_state, &value_cache_, this);
-    auto& tls = thread_local_results_.back();
-    auto* ctx = tls.ctx_;
+#if defined(Py_GIL_DISABLED) && !defined(IS_PYTHON_3_14_PLUS)
+  TORCH_WARN(
+      "The PyTorch profiler is not thread-safe on Python 3.13t. "
+      "Please use Python 3.14t or later.");
+#endif
 
-    // When we begin profiling there are already frames on the Python
-    // interpreter stack. To ensure a complete trace, we must push calls
-    // to all the prior frames onto our event stack. (We stop at depth=128)
+  pybind11::gil_scoped_acquire gil;
+  interpreter_ = PyInterpreterState_Get();
 
-    std::vector<THPFrameObjectPtr> current_stack;
-    auto frame = PyEval_GetFrame();
-    Py_XINCREF(frame);
+  // Shared context passed as the profile arg to all threads.
+  shared_ctx_ = (TraceContext*)TraceContextType.tp_alloc(&TraceContextType, 0);
+  shared_ctx_->tracer_ = this;
+
+  // Enable profiling on all threads. setprofileAllThreads handles its own
+  // synchronization (stop-the-world on free-threaded builds). The callback
+  // returns early because findThreadLocalResults returns nullptr for threads
+  // we haven't set up yet.
+  // Note: This profile will not compose with other CPython profilers, and
+  // cannot be round tripped via `sys.settrace(sys.gettrace())`
+  setprofileAllThreads(PythonTracer::pyProfileFn, (PyObject*)shared_ctx_);
+
+  // Capture existing frames on each thread's stack.
+  {
+    StopTheWorldGuard stw(interpreter_);
+    for (const auto thread_state : interpreterThreads()) {
+      thread_local_results_.emplace_back(this);
+      auto& tls = thread_local_results_.back();
+      thread_local_results_map_[thread_state] = &tls;
+
+      // When we begin profiling there are already frames on the Python
+      // interpreter stack. To ensure a complete trace, we must push calls
+      // to all the prior frames onto our event stack. (We stop at depth=128)
+
+      // NB: `PyThreadState_GetFrame` returns a strong reference.
+      std::vector<THPFrameObjectPtr> current_stack;
+      auto frame = PyThreadState_GetFrame(thread_state);
+
+      size_t depth = 0; // Make sure we can't infinite loop.
+      while (frame != nullptr) {
+        current_stack.emplace_back(frame);
+        if (++depth == 128) {
+          break;
+        }
 
-    size_t depth = 0; // Make sure we can't infinite loop.
-    while (frame != nullptr) {
-      current_stack.emplace_back(frame);
-      if (++depth == 128) {
-        break;
+        // NB: `PyFrame_GetBack` returns a strong reference.
+        frame = PyFrame_GetBack(frame);
       }
 
-      // NB: `PyFrame_GetBack` returns a strong reference.
-      frame = PyFrame_GetBack(frame);
-    }
+      for (auto it = current_stack.rbegin(); it != current_stack.rend(); it++) {
+        recordPyCall(tls, it->get(), true);
+        auto frame_refcount = Py_REFCNT(it->get());
 
-    for (auto it = current_stack.rbegin(); it != current_stack.rend(); it++) {
-      recordPyCall(tls, it->get(), true);
-      auto frame_refcount = Py_REFCNT(it->get());
+        // We hold one reference in `current_stack`, and the interpreter holds
+        // another.
+        TORCH_INTERNAL_ASSERT(frame_refcount >= 2, frame_refcount);
+      }
 
-      // We hold one reference in `current_stack`, and the interpreter holds
-      // another.
-      TORCH_INTERNAL_ASSERT(frame_refcount >= 2, frame_refcount);
+      tls.remaining_start_frames_ = tls.active_frames_;
     }
-
-    tls.remaining_start_frames_ = tls.active_frames_;
-
-    // Note:
-    //   This profile will not compose with other CPython profilers, and
-    //   cannot be round tripped via `sys.settrace(sys.gettrace())`
-    PyEval_SetProfile(PythonTracer::pyProfileFn, (PyObject*)ctx);
   }
 #if IS_PYTHON_3_12
   registerMonitoringCallback();
@@ -1092,6 +1120,7 @@ void unregister_gc_callback() {
     PySequence_DelItem(callbacks, idx);
   } else {
     // Not found, maybe already removed
+    PyErr_Clear();
   }
   Py_DECREF(callbacks);
   Py_DECREF(gc_module);
@@ -1134,16 +1163,23 @@ void PythonTracer::register_gc_callback() {
 }
 
 void PythonTracer::stop() {
-  gil_and_restore_thread gil;
+  pybind11::gil_scoped_acquire gil;
   if (gc_callback_registered_) {
     unregister_gc_callback();
     gc_callback_registered_ = false;
   }
   if (active_) {
-    for (const auto thread_state : interpreterThreads()) {
-      if (thread_state->c_profilefunc == &PythonTracer::pyProfileFn) {
-        PyThreadState_Swap(thread_state);
-        PyEval_SetProfile(nullptr, nullptr);
+    setprofileAllThreads(nullptr, nullptr);
+
+    // Wait for any in-flight pyProfileFn callbacks to finish. Threads inside
+    // pyProfileFn hold their thread's profile_sem. They may have temporarily
+    // released the GIL or parked mid-callback due to a stop-the-world event.
+    // Acquiring each semaphore here blocks until those callbacks complete.
+    {
+      pybind11::gil_scoped_release release;
+      for (auto& tls : thread_local_results_) {
+        tls.profile_sem.acquire();
+        tls.profile_sem.release();
       }
     }
 
@@ -1158,7 +1194,7 @@ void PythonTracer::stop() {
 }
 
 void PythonTracer::restart() {
-  gil_and_restore_thread gil;
+  pybind11::gil_scoped_acquire gil;
   active_ = active_lock_.compare_exchange_strong(active_, true);
   if (!active_) {
     TORCH_WARN(
@@ -1166,14 +1202,7 @@ void PythonTracer::restart() {
         "Refusing to register profile functions.");
     return;
   }
-  int cur_thread = 0;
-  for (const auto thread_state : interpreterThreads()) {
-    if (thread_state->c_profilefunc == nullptr) {
-      auto* ctx = thread_local_results_[cur_thread].ctx_;
-      PyThreadState_Swap(thread_state);
-      PyEval_SetProfile(PythonTracer::pyProfileFn, (PyObject*)ctx);
-    }
-  }
+  setprofileAllThreads(PythonTracer::pyProfileFn, (PyObject*)shared_ctx_);
 #if IS_PYTHON_3_12
   registerMonitoringCallback();
 #endif
@@ -1185,6 +1214,10 @@ PythonTracer::~PythonTracer() {
     TORCH_WARN("`PythonTracer::stop()` was not called.");
     stop();
   }
+  if (Py_IsInitialized() && !Py_IsFinalizing()) {
+    pybind11::gil_scoped_acquire gil;
+    Py_XDECREF((PyObject*)shared_ctx_);
+  }
 }
 
 void PythonTracer::recordPyCall(
@@ -1207,13 +1240,13 @@ void PythonTracer::recordPyCall(
       auto locals = THPObjectPtr(PyFrame_GetLocals(frame));
 
 #if PY_MAJOR_VERSION < 3 || (PY_MAJOR_VERSION == 3 && PY_MINOR_VERSION < 13)
-      auto self = THPObjectPtr(PyDict_GetItemString(locals, "self"));
+      auto self =
+          THPObjectPtr(Py_XNewRef(PyDict_GetItemString(locals, "self")));
 #else
       // In Python-3.13+ `PyFrame_GetLocals()` returns instance of
       // PyFrameLocalsProxy_Type See PEP 667 for more info
       auto self = THPObjectPtr(PyMapping_GetItemString(locals, "self"));
 #endif
-      Py_INCREF(self.get());
       auto back = THPFrameObjectPtr(PyFrame_GetBack(frame));
       TORCH_INTERNAL_ASSERT(back != nullptr);
       return tls.intern<CallType::PyModuleCall, E>(
@@ -1221,11 +1254,11 @@ void PythonTracer::recordPyCall(
     } else if (code.get() == optimizer_hook_) {
       auto locals = THPObjectPtr(PyFrame_GetLocals(frame));
 #if PY_MAJOR_VERSION < 3 || (PY_MAJOR_VERSION == 3 && PY_MINOR_VERSION < 13)
-      auto self = THPObjectPtr(PyDict_GetItemString(locals, "self"));
+      auto self =
+          THPObjectPtr(Py_XNewRef(PyDict_GetItemString(locals, "self")));
 #else
       auto self = THPObjectPtr(PyMapping_GetItemString(locals, "self"));
 #endif
-      Py_INCREF(self.get());
       auto back = THPFrameObjectPtr(PyFrame_GetBack(frame));
       TORCH_INTERNAL_ASSERT(back != nullptr);
       return tls.intern<CallType::PyOptimizerCall, E>(
@@ -1279,12 +1312,14 @@ class PostProcess {
   PostProcess(
       std::function<c10::time_t(c10::approx_time_t)> time_converter,
       std::deque<ThreadLocalResults>& tls,
-      const ValueCache& value_cache,
       c10::time_t end_time_ns)
       : end_time_{end_time_ns}, time_converter_{std::move(time_converter)} {
     for (size_t python_tid : c10::irange(tls.size())) {
       CallTypeHelper<TraceKeyCacheState>::map(
-          tls[python_tid].trace_keys_, *this, value_cache, python_tid);
+          tls[python_tid].trace_keys_,
+          *this,
+          tls[python_tid].value_cache_,
+          python_tid);
 
       addExits<EventType::PyCall>(tls[python_tid].exit_times_, python_tid);
       addExits<EventType::PyCCall>(tls[python_tid].c_exit_times_, python_tid);
@@ -1445,12 +1480,11 @@ std::vector<std::shared_ptr<Result>> PythonTracer::getEvents(
     std::function<c10::time_t(c10::approx_time_t)> time_converter,
     std::vector<python_tracer::CompressedEvent>& enters,
     c10::time_t end_time_ns) {
-  value_cache_.trimPrefixes();
+  for (auto& tls : thread_local_results_) {
+    tls.value_cache_.trimPrefixes();
+  }
   PostProcess post_process(
-      std::move(time_converter),
-      thread_local_results_,
-      value_cache_,
-      end_time_ns);
+      std::move(time_converter), thread_local_results_, end_time_ns);
   post_process.set_start_frames(start_frames_, enters);
   auto out = post_process.run(enters);
 
@@ -1493,17 +1527,15 @@ static void toggle_memory_tracing(bool enable) {
   }
   // Call the function with arguments
   PyObject* args = PyTuple_New(6);
-  PyTuple_SetItem(args, 0, enable ? PyUnicode_FromString("all") : Py_None);
+  PyTuple_SetItem(
+      args, 0, enable ? PyUnicode_FromString("all") : Py_NewRef(Py_None));
   PyTuple_SetItem(args, 1, PyUnicode_FromString("all")); // context
   PyTuple_SetItem(args, 2, PyUnicode_FromString("all")); // stacks
   PyTuple_SetItem(args, 3, THPUtils_packInt64(100000)); // max_entries
-  PyTuple_SetItem(args, 4, Py_None); // device (None)
+  PyTuple_SetItem(args, 4, Py_NewRef(Py_None)); // device (None)
   PyTuple_SetItem(args, 5, PyBool_FromLong(0)); // clear_history (False)
-  PyObject* result = PyObject_Call(snapshot_func.get(), args, nullptr);
+  THPObjectPtr result(PyObject_Call(snapshot_func.get(), args, nullptr));
   Py_DECREF(args);
-  if (result == nullptr) {
-    return;
-  }
 }
 
 void PythonMemoryTracer::start() {
@@ -1522,14 +1554,9 @@ void PythonMemoryTracer::export_memory_history(const std::string& path) {
   if (!snapshot_func) {
     return;
   }
-  PyObject* py_filename = PyUnicode_FromString(path.c_str());
-  // Call the function with arguments (e.g., a file path)
-  PyObject* args = PyTuple_Pack(1, py_filename);
-  PyObject* result = PyObject_Call(snapshot_func.get(), args, nullptr);
-  Py_DECREF(args);
-  if (result == nullptr) {
-    return;
-  }
+  THPObjectPtr py_filename(PyUnicode_FromString(path.c_str()));
+  THPObjectPtr result(
+      PyObject_CallOneArg(snapshot_func.get(), py_filename.get()));
 }
 
 void PythonMemoryTracer::stop() {
@@ -1544,36 +1571,47 @@ int PythonTracer::pyProfileFn(
     PyFrameObject* frame,
     int what,
     PyObject* arg) {
-  auto& local_results =
-      *reinterpret_cast<TraceContext*>(obj)->thread_local_results_;
+  HANDLE_TH_ERRORS
+  auto* tracer = reinterpret_cast<TraceContext*>(obj)->tracer_;
+  auto* local_results = tracer->findThreadLocalResults(PyThreadState_Get());
+  if (C10_UNLIKELY(!local_results)) {
+    return 0;
+  }
+  bool acquired = local_results->profile_sem.tryAcquire();
+  TORCH_INTERNAL_ASSERT(acquired, "pyProfileFn: profile_sem unexpectedly held");
+  // RAII release: ensures the semaphore is released on both normal
+  // return and C++ exception paths (e.g. from pybind11 in ValueCache::store).
+  auto release_sem =
+      c10::make_scope_exit([&]() { local_results->profile_sem.release(); });
   switch (what) {
     case PyTrace_CALL:
-      local_results.active_tracer_->recordPyCall(local_results, frame, false);
+      local_results->active_tracer_->recordPyCall(*local_results, frame, false);
       break;
 
     case PyTrace_C_CALL:
-      local_results.active_tracer_->recordCCall(local_results, frame, arg);
+      local_results->active_tracer_->recordCCall(*local_results, frame, arg);
       break;
 
     case PyTrace_RETURN:
-      local_results.exit_times_.emplace_back(c10::getApproximateTime());
-      local_results.active_frames_--;
-      if (local_results.active_frames_ <
-          local_results.remaining_start_frames_) {
-        local_results.remaining_start_frames_ = local_results.active_frames_;
+      local_results->exit_times_.emplace_back(c10::getApproximateTime());
+      local_results->active_frames_--;
+      if (local_results->active_frames_ <
+          local_results->remaining_start_frames_) {
+        local_results->remaining_start_frames_ = local_results->active_frames_;
       }
       break;
 
     case PyTrace_C_EXCEPTION:
     case PyTrace_C_RETURN:
-      if (local_results.active_frames_ >
-          local_results.remaining_start_frames_) {
-        local_results.c_exit_times_.emplace_back(c10::getApproximateTime());
-        local_results.active_frames_--;
+      if (local_results->active_frames_ >
+          local_results->remaining_start_frames_) {
+        local_results->c_exit_times_.emplace_back(c10::getApproximateTime());
+        local_results->active_frames_--;
       }
       break;
   }
   return 0;
+  END_HANDLE_TH_ERRORS_RET(-1)
 }
 
 std::unique_ptr<python_tracer::PythonTracerBase> getTracer(
diff --git a/torch/csrc/autograd/python_cpp_function.cpp b/torch/csrc/autograd/python_cpp_function.cpp
index 6787df7080ad8..e4e68442114df 100644
--- a/torch/csrc/autograd/python_cpp_function.cpp
+++ b/torch/csrc/autograd/python_cpp_function.cpp
@@ -326,16 +326,11 @@ void registerCppFunction(const std::type_info& type, PyTypeObject* pytype) {
 }
 
 bool THPCppFunction_Check(PyObject* obj) {
-  THPObjectPtr type = THPObjectPtr(PyObject_Type(obj));
-  if ((PyTypeObject*)type.get() == get_default_type()) {
-    return true;
-  }
-  if (cpp_function_types_set.find((PyTypeObject*)type.get()) ==
-      cpp_function_types_set.end()) {
-    return false;
-  } else {
+  PyTypeObject* type = Py_TYPE(obj);
+  if (type == get_default_type()) {
     return true;
   }
+  return cpp_function_types_set.contains(type);
 }
 
 static PyObject* callRegisterFn(PyObject* dict, PyObject* hook) {
diff --git a/torch/csrc/autograd/python_function.cpp b/torch/csrc/autograd/python_function.cpp
index 3319ef4d83c90..2a7ee2b19a045 100644
--- a/torch/csrc/autograd/python_function.cpp
+++ b/torch/csrc/autograd/python_function.cpp
@@ -1,5 +1,7 @@
 #include <torch/csrc/autograd/python_function.h>
 
+#include <atomic>
+
 #include <ATen/ATen.h>
 #include <ATen/SequenceNumber.h>
 #include <c10/util/irange.h>
@@ -1097,6 +1099,7 @@ void _trace_post_record(
   }
 
   std::vector<torch::jit::Value*> trace_outputs;
+  trace_outputs.reserve(static_cast<size_t>(std::max(0, num_outputs)));
   for (const auto i : c10::irange(num_outputs)) {
     PyObject* obj = PyTuple_GET_ITEM(output_objects, i);
     if (THPVariable_Check(obj)) {
@@ -1125,9 +1128,10 @@ void _trace_post_record(
   // to the original tuple type.
   if (!unpack_output) {
     std::vector<at::TypePtr> new_tuple_values;
+    new_tuple_values.reserve(num_outputs);
     for (const auto i : c10::irange(num_outputs)) {
       auto ptr = node->outputs()[i]->type();
-      new_tuple_values.push_back(ptr);
+      new_tuple_values.push_back(std::move(ptr));
     }
     auto tuple_type = at::TupleType::create(std::move(new_tuple_values));
     // The i-th tuple element receives a new tensor type with element type and
@@ -1302,11 +1306,14 @@ THPObjectPtr make_ctx_input_output_tuple(
   return result;
 }
 
-static PyObject* THPFunction_setup_context = nullptr;
-
 static PyObject* get_base_setup_context() {
-  if (THPFunction_setup_context != nullptr) {
-    return THPFunction_setup_context;
+  // NOTE: THPFunction_setup_context is intentionally leaked and never freed.
+  static std::atomic<PyObject*> THPFunction_setup_context = nullptr;
+
+  PyObject* setup_context =
+      THPFunction_setup_context.load(std::memory_order_acquire);
+  if (setup_context != nullptr) {
+    return setup_context;
   }
 
   auto module = THPObjectPtr(PyImport_ImportModule("torch.autograd.function"));
@@ -1320,11 +1327,17 @@ static PyObject* get_base_setup_context() {
 
   // setup_context gets "leaked" - we return a new reference and hold onto it
   // forever.
-  auto setup_context = PyObject_GetAttrString(function, "setup_context");
+  setup_context = PyObject_GetAttrString(function, "setup_context");
   if (!setup_context)
     return nullptr;
-  THPFunction_setup_context = setup_context;
-  return THPFunction_setup_context;
+
+  PyObject* expected = nullptr;
+  if (!THPFunction_setup_context.compare_exchange_strong(
+          expected, setup_context, std::memory_order_acq_rel)) {
+    Py_DECREF(setup_context);
+    return expected;
+  }
+  return setup_context;
 }
 
 PyObject* THPFunction_apply(PyObject* cls, PyObject* inputs) {
diff --git a/torch/csrc/autograd/python_torch_functions_manual.cpp b/torch/csrc/autograd/python_torch_functions_manual.cpp
index 79739b6e459d2..51a9fa41de5f2 100644
--- a/torch/csrc/autograd/python_torch_functions_manual.cpp
+++ b/torch/csrc/autograd/python_torch_functions_manual.cpp
@@ -330,7 +330,7 @@ static PyObject* THPVariable_asarray(
   HANDLE_TH_ERRORS
   static PythonArgParser parser(
       {
-          "asarray(PyObject* obj, *, ScalarType? dtype=None, Device? device=None, bool? copy=None, bool requires_grad=False)",
+          "asarray(PyObject* obj, *, ScalarType? dtype=None, Device? device=None, bool? copy=None, bool? requires_grad=None)",
       },
       /*traceable=*/false);
 
@@ -347,7 +347,7 @@ static PyObject* THPVariable_asarray(
     auto dtype = r.scalartypeOptional(1);
     auto device = r.deviceOptional(2);
     auto copy = r.toBoolOptional(3);
-    auto requires_grad = r.toBool(4);
+    auto requires_grad = r.toBoolOptional(4);
     return wrap(torch::utils::asarray(obj, dtype, device, copy, requires_grad));
   }
 
diff --git a/torch/csrc/autograd/python_variable.cpp b/torch/csrc/autograd/python_variable.cpp
index 92b54a9a6bc56..45f204e9cd589 100644
--- a/torch/csrc/autograd/python_variable.cpp
+++ b/torch/csrc/autograd/python_variable.cpp
@@ -227,6 +227,14 @@ std::pair<py::object, py::dict> parseIValuesToPyArgsKwargs(
       }
       return false;
     };
+    auto matchList = [&](c10::TypeKind kind) {
+      const auto& t = arg.real_type();
+      if (auto list_t = t->cast<c10::ListType>()) {
+        if (list_t->getElementType()->kind() == kind)
+          return true;
+      }
+      return false;
+    };
     if (argument.isNone()) {
       return py::none();
     } else if (match(c10::ScalarTypeType::Kind)) {
@@ -239,6 +247,31 @@ std::pair<py::object, py::dict> parseIValuesToPyArgsKwargs(
           reinterpret_cast<PyObject*>(obj));
     } else if (match(c10::MemoryFormatType::Kind)) {
       return py::cast(static_cast<c10::MemoryFormat>(argument.toInt()));
+    } else if (matchList(c10::ScalarTypeType::Kind)) {
+      const auto& list = argument.toListRef();
+      py::list result(list.size());
+      for (const auto i : c10::irange(list.size())) {
+        auto* obj = getTHPDtype(static_cast<c10::ScalarType>(list[i].toInt()));
+        result[i] = py::reinterpret_borrow<py::object>(
+            reinterpret_cast<PyObject*>(obj));
+      }
+      return result;
+    } else if (matchList(c10::LayoutType::Kind)) {
+      const auto& list = argument.toListRef();
+      py::list result(list.size());
+      for (const auto i : c10::irange(list.size())) {
+        auto* obj = getTHPLayout(static_cast<c10::Layout>(list[i].toInt()));
+        result[i] = py::reinterpret_borrow<py::object>(
+            reinterpret_cast<PyObject*>(obj));
+      }
+      return result;
+    } else if (matchList(c10::MemoryFormatType::Kind)) {
+      const auto& list = argument.toListRef();
+      py::list result(list.size());
+      for (const auto i : c10::irange(list.size())) {
+        result[i] = py::cast(static_cast<c10::MemoryFormat>(list[i].toInt()));
+      }
+      return result;
     } else {
       return torch::jit::toPyObject(argument);
     }
@@ -1087,9 +1120,10 @@ struct IValueOrDTensorSpec {
   py::object dtensor_spec;
 
   bool operator==(const IValueOrDTensorSpec& rhs) const {
-    return dtensor_spec
-        ? (rhs.dtensor_spec && dtensor_spec.equal(rhs.dtensor_spec))
-        : (iv == rhs.iv);
+    if (dtensor_spec) {
+      return rhs.dtensor_spec && dtensor_spec.equal(rhs.dtensor_spec);
+    }
+    return !rhs.dtensor_spec && iv == rhs.iv;
   }
 };
 
@@ -1162,18 +1196,6 @@ class NativeOpSchema {
   // have no guarantees about its lifetime. This class is cheap anyway.
   c10::OperatorHandle op_;
   std::size_t hash_;
-  // Subtle point: consider clamp.Tensor(Tensor self, Tensor?
-  // min=None, Tensor? max=None). The invocations clamp(t1, None, t2)
-  // and clamp(t1, t2, None) have the same comparison key (t1, t2)
-  // because we drop non-static non-tensor args from comparison. The
-  // only way we happen to be able to tell them apart is that we omit
-  // trailing defaulted arguments from the args tuple passed to
-  // __torch_dispatch__ (and hence to DTensor dispatch as well), so
-  // they have different args_schema_len_.
-  //
-  // I am preserving this existing behavior, but I suspect we should
-  // make an algorithm change to be less brittle, such as including
-  // None defaults for Tensor arguments in the comparison.
   std::size_t args_schema_len_;
   // There is no particular justification for the choice of 8
   // here. Feel free to change it.
@@ -1842,7 +1864,7 @@ static bool DTensor_OpSchema_recompute_comparison_key_impl(
   size_t idx = 0;
   for (const auto& e : args_schema) {
     if (idx >= native_info.static_argnum ||
-        arg_type_tensor_or_tensor_list_like(e)) {
+        arg_type_tensor_or_tensor_list_like(e) || e.is_none()) {
       if (PyList_Check(e.ptr())) {
         args_to_hash.push_back(
             py::reinterpret_steal<py::object>(PyList_AsTuple(e.ptr())));
@@ -2156,6 +2178,19 @@ static py::object try_find_mesh_from_args(
       return py::reinterpret_borrow<py::object>(
           py_tensor.attr(dtensor_interned_strings.device_mesh));
     }
+    if (argument_it->isList()) {
+      const auto list = argument_it->toList();
+      // need to expand list and use first DTensor
+      for (const auto& item : list) {
+        const auto [item_flavor, item_py_tensor] =
+            check_for_dtensor_or_tensor(item);
+        if (item_flavor == TensorFlavor::EXACTLY_DTENSOR ||
+            item_flavor == TensorFlavor::DTENSOR_SUBCLASS) {
+          return py::reinterpret_borrow<py::object>(
+              item_py_tensor.attr(dtensor_interned_strings.device_mesh));
+        }
+      }
+    }
   }
   TORCH_CHECK_VALUE(
       false, "Cannot find device mesh from args for op : ", op.operator_name());
@@ -2258,6 +2293,10 @@ static bool dtensor_spec_has_symints(py::handle spec) {
   return contains_any_symint(shape);
 }
 
+// set to false to bail out of the C++ fast path for ops that need pytree
+// flattening (e.g. future pytree custom ops)
+static constexpr bool kRunPytreeWithFastPath = true;
+
 static std::optional<std::pair<NativeOpSchema, /*ComputeMesh*/ py::object>>
 create_native_op_schema(
     const c10::OperatorHandle& op,
@@ -2267,7 +2306,7 @@ create_native_op_schema(
   // operating on IValues instead of Python stuff.
 
   py::object runtime_schema_info = get_runtime_schema_info_for_op(py_op);
-  if (runtime_schema_info &&
+  if (!kRunPytreeWithFastPath && runtime_schema_info &&
       checked_istrue(py::handle(runtime_schema_info)
                          .attr(dtensor_interned_strings.needs_pytree)
                          .ptr())) {
@@ -2275,7 +2314,6 @@ create_native_op_schema(
     // now since only a minority of ops need it.
     return std::nullopt;
   }
-
   OperatorArgsKwargsView args_kwargs(op, *stack);
   auto native_info = unpack_runtime_schema_info(
       py::handle(runtime_schema_info), args_kwargs.num_positional_args());
@@ -2288,7 +2326,9 @@ create_native_op_schema(
   const auto handle_non_dtensor_arg =
       [&comparison_key, &comparison_key_hash, &native_info](
           size_t idx, c10::IValue arg) {
-        if (idx >= native_info.static_argnum) {
+        bool is_none_or_undefined =
+            arg.isNone() || (arg.isTensor() && !arg.toTensor().defined());
+        if (idx >= native_info.static_argnum || is_none_or_undefined) {
           if (arg.isList()) {
             const auto& list = arg.toList();
             if (list.empty()) {
@@ -2319,58 +2359,105 @@ create_native_op_schema(
     comparison_key.emplace_back(std::move(arg));
   };
 
-  Py_ssize_t idx = 0;
+  const auto handle_non_tensor_or_undefined =
+      [&comparison_key, &comparison_key_hash](c10::IValue arg) {
+        // We reach here when arg is TensorFlavor::NON_TENSOR
+        // (not a Tensor at all or undefined Tensor)
+        // We coerce undefined Tensor to None, just as we do when
+        // converting IValues to PyObject. (same behaviour as
+        // handle_non_dtensor_arg)
+        if (arg.isTensor() && !arg.toTensor().defined()) {
+          arg = c10::IValue();
+        }
+        comparison_key_hash =
+            c10::hash_combine(comparison_key_hash, c10::IValue::hash(arg));
+        comparison_key.emplace_back(std::move(arg));
+      };
+
   const bool allow_implicit_replication =
       at::get_dtensor_allow_implicit_replication();
+
+  const auto handle_exactly_dtensor = [&](const auto py_tensor) {
+    py::object spec = py_tensor.attr(dtensor_interned_strings._spec);
+    if (dtensor_spec_has_symints(spec)) {
+      // Symints are unhashable, so we can't use the cache for
+      // sharding propagation. bail out to slow path.
+      return true;
+    }
+    handle_dtensor_arg(std::move(spec));
+    if (compute_mesh.is_none()) {
+      compute_mesh = py::reinterpret_borrow<py::object>(
+          py_tensor.attr(dtensor_interned_strings.device_mesh));
+    }
+    return false;
+  };
+
+  const auto handle_exactly_tensor = [&](const auto py_tensor) {
+    if (compute_mesh.is_none()) {
+      compute_mesh = try_find_mesh_from_args(op, args_kwargs);
+    }
+    handle_dtensor_arg(try_replicate_spec_for_scalar_tensor(
+        allow_implicit_replication, py_op, py_tensor, compute_mesh));
+  };
+
+  Py_ssize_t idx = 0;
+
   for (auto argument_it = args_kwargs.args_begin();
        argument_it != args_kwargs.args_end();
        ++argument_it) {
     const auto& arg = *argument_it;
     const auto [tensor_flavor, py_tensor] = check_for_dtensor_or_tensor(arg);
+
     switch (tensor_flavor) {
       case TensorFlavor::EXACTLY_DTENSOR:
       case TensorFlavor::DTENSOR_SUBCLASS: {
-        py::object spec = py_tensor.attr(dtensor_interned_strings._spec);
-        if (dtensor_spec_has_symints(spec)) {
-          // Symints are unhashable, so we can't use the cache for
-          // sharding propagation. bail out to slow path.
+        bool is_symint = handle_exactly_dtensor(py_tensor);
+        if (is_symint) {
           return std::nullopt;
         }
-        handle_dtensor_arg(std::move(spec));
-        if (compute_mesh.is_none()) {
-          compute_mesh = py::reinterpret_borrow<py::object>(
-              py_tensor.attr(dtensor_interned_strings.device_mesh));
-        }
         break;
       }
       case TensorFlavor::EXACTLY_TENSOR:
       case TensorFlavor::NON_DTENSOR_TENSOR_SUBCLASS: {
-        if (compute_mesh.is_none()) {
-          compute_mesh = try_find_mesh_from_args(op, args_kwargs);
-        }
-        handle_dtensor_arg(try_replicate_spec_for_scalar_tensor(
-            allow_implicit_replication, py_op, py_tensor, compute_mesh));
+        handle_exactly_tensor(py_tensor);
         break;
       }
       case TensorFlavor::NON_TENSOR: {
         // Check if this is a list/tuple that might contain DTensors (e.g.,
         // torch.cat)
-        if (arg.isList() && compute_mesh.is_none()) {
+        if (arg.isList()) {
           const auto list = arg.toList();
+          comparison_key_hash = hash_combine(comparison_key_hash, list.size());
+          comparison_key.emplace_back(static_cast<int64_t>(list.size()));
+          // pytree unflattening
           for (const auto& item : list) {
             const auto [item_flavor, item_py_tensor] =
                 check_for_dtensor_or_tensor(item);
             if (item_flavor == TensorFlavor::EXACTLY_DTENSOR ||
                 item_flavor == TensorFlavor::DTENSOR_SUBCLASS) {
-              compute_mesh = py::reinterpret_borrow<py::object>(
-                  item_py_tensor.attr(dtensor_interned_strings.device_mesh));
-              break;
+              bool is_symint = handle_exactly_dtensor(item_py_tensor);
+              if (is_symint) {
+                return std::nullopt;
+              }
+            } else if (
+                item_flavor == TensorFlavor::EXACTLY_TENSOR ||
+                item_flavor == TensorFlavor::NON_DTENSOR_TENSOR_SUBCLASS) {
+              handle_exactly_tensor(item_py_tensor);
+            } else { // non-tensor
+              // Use handle_non_dtensor_arg to respect static_argnum.
+              // Non-tensor items in lists (e.g., ScalarList args to
+              // foreach ops) should only be included in the cache key
+              // if the list's argument index is >= static_argnum.
+              // Otherwise, step-varying scalars (like AdamW bias
+              // corrections) cause unbounded cache growth.
+              handle_non_dtensor_arg(idx, item);
             }
           }
+        } else {
+          // non DTensor/Tensor args (i.e. int/float/bool), just add to
+          // local_args
+          handle_non_dtensor_arg(idx, arg);
         }
-        // non DTensor/Tensor args (i.e. int/float/bool), just add to
-        // local_args
-        handle_non_dtensor_arg(idx, arg);
         break;
       }
       default:
@@ -2397,6 +2484,18 @@ create_native_op_schema(
   }
 
   if (native_info.static_kwargkey && !native_info.static_kwargkey.is_none()) {
+    // Only kwargs named in static_kwargkey affect sharding propagation and
+    // belong in the cache key. The Python comparison key
+    // (DTensor_OpSchema_recompute_comparison_key_impl) already filters this
+    // way; the C++ fast path must match. Without this filter, step-varying
+    // scalar kwargs (e.g. the `value` arg of addcdiv_ used by AdamW bias
+    // corrections) cause unbounded cache growth.
+    c10::SmallVector<std::string, 2> static_kwarg_names;
+    for (const auto& key :
+         py::reinterpret_borrow<py::list>(native_info.static_kwargkey)) {
+      static_kwarg_names.push_back(py::cast<std::string>(key));
+    }
+
     // Separator to disambiguate kwargs from args in comparison and hashing.
     static constexpr int64_t kwargs_separator = 0x0011223344556677LL;
     comparison_key.emplace_back(static_cast<int64_t>(kwargs_separator));
@@ -2405,47 +2504,66 @@ create_native_op_schema(
     for (auto argument_it = args_kwargs.kwargs_begin();
          argument_it != args_kwargs.kwargs_end();
          ++argument_it) {
+      const auto underlying_index = argument_it.underlying_index();
+      const auto [tensor_flavor, py_tensor] =
+          check_for_dtensor_or_tensor(*argument_it);
+      // Skip non-tensor kwargs not listed in static_kwargkey.
+      if (tensor_flavor == TensorFlavor::NON_TENSOR) {
+        const auto& kwarg_name =
+            op.schema().arguments()[underlying_index].name();
+        if (std::find(
+                static_kwarg_names.begin(),
+                static_kwarg_names.end(),
+                kwarg_name) == static_kwarg_names.end()) {
+          continue;
+        }
+      }
+
       // Rather than hash/compare the string key, we can just use the
       // index of the kwarg in the schema!
-      const auto underlying_index = argument_it.underlying_index();
       comparison_key.emplace_back(c10::IValue(underlying_index));
       comparison_key_hash = hash_combine(
           comparison_key_hash, c10::IValue::hash(comparison_key.back().iv));
-      const auto [tensor_flavor, py_tensor] =
-          check_for_dtensor_or_tensor(*argument_it);
       switch (tensor_flavor) {
         case TensorFlavor::EXACTLY_DTENSOR:
         case TensorFlavor::DTENSOR_SUBCLASS: {
-          handle_dtensor_arg(py_tensor.attr(dtensor_interned_strings._spec));
-          if (compute_mesh.is_none()) {
-            compute_mesh = py::reinterpret_borrow<py::object>(
-                py_tensor.attr(dtensor_interned_strings.device_mesh));
+          bool is_symint = handle_exactly_dtensor(py_tensor);
+          if (is_symint) {
+            return std::nullopt;
           }
           break;
         }
         case TensorFlavor::EXACTLY_TENSOR:
         case TensorFlavor::NON_DTENSOR_TENSOR_SUBCLASS: {
-          handle_dtensor_arg(try_replicate_spec_for_scalar_tensor(
-              allow_implicit_replication, py_op, py_tensor, compute_mesh));
+          handle_exactly_tensor(py_tensor);
           break;
         }
         case TensorFlavor::NON_TENSOR: {
-          // Check if this is a list/tuple that might contain DTensors (e.g.,
-          // torch.cat)
-          if (argument_it->isList() && compute_mesh.is_none()) {
+          if (argument_it->isList()) {
             const auto list = argument_it->toList();
+            comparison_key_hash =
+                hash_combine(comparison_key_hash, list.size());
+            comparison_key.emplace_back(static_cast<int64_t>(list.size()));
             for (const auto& item : list) {
               const auto [item_flavor, item_py_tensor] =
                   check_for_dtensor_or_tensor(item);
               if (item_flavor == TensorFlavor::EXACTLY_DTENSOR ||
                   item_flavor == TensorFlavor::DTENSOR_SUBCLASS) {
-                compute_mesh = py::reinterpret_borrow<py::object>(
-                    item_py_tensor.attr(dtensor_interned_strings.device_mesh));
-                break;
+                bool is_symint = handle_exactly_dtensor(item_py_tensor);
+                if (is_symint) {
+                  return std::nullopt;
+                }
+              } else if (
+                  item_flavor == TensorFlavor::EXACTLY_TENSOR ||
+                  item_flavor == TensorFlavor::NON_DTENSOR_TENSOR_SUBCLASS) {
+                handle_exactly_tensor(item_py_tensor);
+              } else { // non-tensor
+                handle_non_tensor_or_undefined(item);
               }
             }
+          } else {
+            handle_non_dtensor_arg(native_info.static_argnum, *argument_it);
           }
-          handle_non_dtensor_arg(native_info.static_argnum, *argument_it);
           break;
         }
         default:
@@ -2488,6 +2606,13 @@ static PyObject* clear_DTensor_sharding_propagator_cache(
   Py_RETURN_NONE;
 }
 
+static PyObject* reinit_DTensor_dispatch_logger(
+    PyObject* self,
+    PyObject* noargs) {
+  dtensor_dispatch_logger_initialized = false;
+  Py_RETURN_NONE;
+}
+
 using getter = PyObject* (*)(PyObject*, void*);
 using setter = int (*)(PyObject*, PyObject*, void*);
 
@@ -3482,6 +3607,10 @@ static PyMethodDef extra_dtensor_functions[] = {
      clear_DTensor_sharding_propagator_cache,
      METH_NOARGS,
      nullptr},
+    {"_reinit_DTensor_dispatch_logger",
+     reinit_DTensor_dispatch_logger,
+     METH_NOARGS,
+     nullptr},
     {nullptr}};
 
 struct THPVariableMeta {
diff --git a/torch/csrc/autograd/variable.cpp b/torch/csrc/autograd/variable.cpp
index 89135f9aa9a22..1e86192c184fb 100644
--- a/torch/csrc/autograd/variable.cpp
+++ b/torch/csrc/autograd/variable.cpp
@@ -118,9 +118,8 @@ ViewInfo ViewInfo::chain(
             "Attempted to chain views when the parent view has no view_func() and "
             "does not support as_strided(). This is not supported.";
         view_func = std::make_unique<ErroringViewFunc>(error_msg);
-        rev_view_func = [=](const at::Tensor& root_view) {
+        rev_view_func = [=](const at::Tensor& root_view) -> at::Tensor {
           TORCH_CHECK(false, error_msg);
-          return root_view;
         };
       }
     }
diff --git a/torch/csrc/cuda/CUDAPluggableAllocator.cpp b/torch/csrc/cuda/CUDAPluggableAllocator.cpp
index bb08bcaf9bedd..1fb6795eae6fd 100644
--- a/torch/csrc/cuda/CUDAPluggableAllocator.cpp
+++ b/torch/csrc/cuda/CUDAPluggableAllocator.cpp
@@ -297,6 +297,14 @@ void CUDAPluggableAllocator::attachOutOfMemoryObserver(
       "If you need it, please file an issue describing your use case.");
 }
 
+void CUDAPluggableAllocator::attachOomRejectionObserver(
+    c10::cuda::CUDACachingAllocator::OomRejectionObserver observer) {
+  TORCH_CHECK_NOT_IMPLEMENTED(
+      false,
+      "CUDAPluggableAllocator does not yet support attachOomRejectionObserver. "
+      "If you need it, please file an issue describing your use case.");
+}
+
 void CUDAPluggableAllocator::attachAllocatorTraceTracker(
     c10::cuda::CUDACachingAllocator::AllocatorTraceTracker tracker) {
   TORCH_CHECK(
diff --git a/torch/csrc/cuda/CUDAPluggableAllocator.h b/torch/csrc/cuda/CUDAPluggableAllocator.h
index 5434dad237608..91b2179038a19 100644
--- a/torch/csrc/cuda/CUDAPluggableAllocator.h
+++ b/torch/csrc/cuda/CUDAPluggableAllocator.h
@@ -123,6 +123,8 @@ struct TORCH_CUDA_CPP_API CUDAPluggableAllocator
       const std::vector<std::string>& skip_actions) override;
   void attachOutOfMemoryObserver(
       c10::cuda::CUDACachingAllocator::OutOfMemoryObserver observer) override;
+  void attachOomRejectionObserver(
+      c10::cuda::CUDACachingAllocator::OomRejectionObserver observer) override;
   void attachAllocatorTraceTracker(
       c10::cuda::CUDACachingAllocator::AllocatorTraceTracker tracker) override;
   std::shared_ptr<c10::cuda::CUDACachingAllocator::AllocatorState>
diff --git a/torch/csrc/cuda/Graph.cpp b/torch/csrc/cuda/Graph.cpp
index 2a551ae28e96d..99d17f4aba736 100644
--- a/torch/csrc/cuda/Graph.cpp
+++ b/torch/csrc/cuda/Graph.cpp
@@ -112,5 +112,19 @@ void THCPGraph_init(PyObject* module) {
             // compile error.
             return reinterpret_cast<uintptr_t>(graph_exec);
           },
-          py::call_guard<py::gil_scoped_release>());
+          py::call_guard<py::gil_scoped_release>())
+      .def_static(
+          "get_currently_capturing_graph",
+          torch::wrap_pybind_function_no_gil(
+              &::at::cuda::CUDAGraph::get_currently_capturing_graph),
+          py::return_value_policy::reference)
+      .def(
+          "begin_capture_to_if_node",
+          torch::wrap_pybind_function_no_gil(
+              &::at::cuda::CUDAGraph::begin_capture_to_if_node),
+          py::arg("scalar_cuda_pred_tensor"))
+      .def(
+          "end_capture_to_conditional_node",
+          torch::wrap_pybind_function_no_gil(
+              &::at::cuda::CUDAGraph::end_capture_to_conditional_node));
 }
diff --git a/torch/csrc/cuda/GreenContext.cpp b/torch/csrc/cuda/GreenContext.cpp
index 52cca598211eb..d420600b0968f 100644
--- a/torch/csrc/cuda/GreenContext.cpp
+++ b/torch/csrc/cuda/GreenContext.cpp
@@ -6,13 +6,50 @@
 
 void THCPGreenContext_init(PyObject* module) {
   auto m = py::handle(module).cast<py::module>();
+
+  py::enum_<at::cuda::WorkqueueScope>(m, "_WorkqueueScope")
+      .value("device_ctx", at::cuda::WorkqueueScope::DeviceCtx)
+      .value("balanced", at::cuda::WorkqueueScope::Balanced);
+
   py::class_<at::cuda::GreenContext>(m, "_CUDAGreenContext")
-      .def_static("create", &::at::cuda::GreenContext::create)
+      .def_static(
+          "create",
+          [](std::optional<uint32_t> device_id,
+             std::optional<uint32_t> num_sms,
+             std::optional<std::string> workqueue_scope,
+             std::optional<uint32_t> workqueue_concurrency_limit) {
+            std::optional<int32_t> scope;
+            if (workqueue_scope.has_value()) {
+              const auto& s = *workqueue_scope;
+              if (s == "device_ctx") {
+                scope =
+                    static_cast<int32_t>(at::cuda::WorkqueueScope::DeviceCtx);
+              } else if (s == "balanced") {
+                scope =
+                    static_cast<int32_t>(at::cuda::WorkqueueScope::Balanced);
+              } else {
+                throw std::invalid_argument(
+                    "workqueue_scope must be 'device_ctx' or 'balanced', got '" +
+                    s + "'");
+              }
+            }
+            return at::cuda::GreenContext::create(
+                device_id, num_sms, scope, workqueue_concurrency_limit);
+          },
+          py::kw_only(),
+          py::arg("device_id") = py::none(),
+          py::arg("num_sms") = py::none(),
+          py::arg("workqueue_scope") = py::none(),
+          py::arg("workqueue_concurrency_limit") = py::none())
+      .def_static(
+          "max_workqueue_concurrency",
+          &at::cuda::GreenContext::max_workqueue_concurrency,
+          py::arg("device_id") = py::none())
       .def("set_context", &::at::cuda::GreenContext::setContext)
       .def("pop_context", &::at::cuda::GreenContext::popContext)
       .def("Stream", [](at::cuda::GreenContext& self) {
         auto s = self.Stream();
-        cudaStream_t raw = self.Stream();
+        cudaStream_t raw = s.stream();
         auto ptr_val = reinterpret_cast<uintptr_t>(raw);
 
         py::object torch_cuda = py::module::import("torch.cuda");
diff --git a/torch/csrc/cuda/Module.cpp b/torch/csrc/cuda/Module.cpp
index ac982dc5d33d0..efa34b48921ba 100644
--- a/torch/csrc/cuda/Module.cpp
+++ b/torch/csrc/cuda/Module.cpp
@@ -3,6 +3,7 @@
 #include <ATen/cuda/CUDAConfig.h>
 #include <ATen/detail/CUDAHooksInterface.h>
 #include <ATen/native/ConvUtils.h>
+#include <ATen/native/RNN.h>
 #include <c10/core/Device.h>
 #include <c10/core/TensorImpl.h>
 #include <c10/util/Exception.h>
@@ -297,6 +298,11 @@ PyObject* THCPModule_cudaCachingAllocator_raw_alloc(
     return nullptr;
   }
   auto size = PyLong_AsSsize_t(size_o);
+  TORCH_CHECK_VALUE(
+      size >= 0,
+      "Invalid memory size: ",
+      size,
+      ". caching_allocator_alloc requires a non-negative size.");
   cudaStream_t stream = static_cast<cudaStream_t>(PyLong_AsVoidPtr(stream_o));
   void* mem = nullptr;
   {
@@ -408,6 +414,23 @@ PyObject* THCPModule_cudaJiteratorCompileAndLaunchKernel(
   END_HANDLE_TH_ERRORS
 }
 
+PyObject* THCPModule_allocationTraceback(PyObject* _unused, PyObject* obj) {
+  HANDLE_TH_ERRORS
+  void* ptr = PyLong_AsVoidPtr(obj);
+  if (!ptr && PyErr_Occurred()) {
+    return nullptr;
+  }
+  auto ctx = c10::cuda::CUDACachingAllocator::getContextForPointer(ptr);
+  if (!ctx) {
+    Py_RETURN_NONE;
+  }
+  auto* tb = getCapturedTracebackFromContext(ctx);
+  std::vector<CapturedTraceback*> tbs = {tb};
+  auto frames = py_symbolize(tbs);
+  return frames.at(0).release().ptr();
+  END_HANDLE_TH_ERRORS
+}
+
 PyObject* THCPModule_cudaCachingAllocator_raw_delete(
     PyObject* _unused,
     PyObject* obj) {
@@ -434,6 +457,14 @@ PyObject* THCPModule_cudaCachingAllocator_enable(
   END_HANDLE_TH_ERRORS
 }
 
+PyObject* THCPModule_cudaCachingAllocator_is_enabled(
+    PyObject* _unused,
+    PyObject* noargs) {
+  HANDLE_TH_ERRORS
+  return PyBool_FromLong(c10::cuda::CUDACachingAllocator::isEnabled());
+  END_HANDLE_TH_ERRORS
+}
+
 PyObject* THCPModule_getAllocatorBackend(PyObject* _unused, PyObject* noargs) {
   HANDLE_TH_ERRORS
   return THPUtils_packString(c10::cuda::CUDACachingAllocator::name());
@@ -469,41 +500,6 @@ PyObject* THCPModule_cudaSleep(PyObject* _unused, PyObject* cycles) {
   END_HANDLE_TH_ERRORS
 }
 
-// We need to ensure that as long as a thread will NEVER loose the GIL as long
-// as it holds the CUDA mutex. Otherwise another thread might be scheduled and
-// try to e.g. allocate a new tensor which will cause a deadlock. It's enough to
-// have a single global, because it can be only set once (cudaMutex is not
-// recursive) by the thread that owns the mutex (obviously there can be only one
-// such thread).
-static PyGILState_STATE cudaMutexGILState;
-
-PyObject* THCPModule_cudaLockMutex(PyObject* module, PyObject* noargs) {
-  auto mutex = c10::cuda::getFreeMutex();
-  // This has to be a busy loop because we **absolutely need to** hold the GIL
-  // or it's a recipe for a deadlock otherwise (if we let other Python threads
-  // run while we have the cudaMutex, but not the GIL, they might try to e.g.
-  // free a CUDA tensor and acquire the cudaMutex without giving up the GIL,
-  // because it happens deep within THC).
-  while (true) {
-    if (mutex->try_lock())
-      break;
-    {
-      pybind11::gil_scoped_release no_gil;
-      std::this_thread::sleep_for(std::chrono::microseconds(10));
-    }
-  }
-
-  cudaMutexGILState = PyGILState_Ensure();
-  Py_RETURN_NONE;
-}
-
-PyObject* THCPModule_cudaUnlockMutex(PyObject* module, PyObject* noargs) {
-  auto mutex = c10::cuda::getFreeMutex();
-  PyGILState_Release(cudaMutexGILState);
-  mutex->unlock();
-  Py_RETURN_NONE;
-}
-
 PyObject* THCPModule_hasPrimaryContext(PyObject* _unused, PyObject* arg) {
   HANDLE_TH_ERRORS
   TORCH_CHECK(
@@ -610,6 +606,7 @@ PyObject* THCPModule_memoryStats(PyObject* _unused, PyObject* arg) {
   result["num_sync_all_streams"] = stats.num_sync_all_streams;
   result["num_device_alloc"] = stats.num_device_alloc;
   result["num_device_free"] = stats.num_device_free;
+  result["num_oom_rejections"] = stats.num_oom_rejections;
   result["allocation"] = statArrayToDict(stats.allocation);
   result["segment"] = statArrayToDict(stats.segment);
   result["active"] = statArrayToDict(stats.active);
@@ -776,6 +773,7 @@ PyObject* THCPModule_memorySnapshot(PyObject* _unused, PyObject* arg) {
   py::str time_us_s = "time_us";
   py::str compile_context_s = "compile_context";
   py::str user_metadata_s = "user_metadata";
+  py::str pool_id_s = "pool_id";
 
   py::list empty_frames;
   std::vector<CapturedTraceback*> to_gather_frames;
@@ -895,6 +893,7 @@ PyObject* THCPModule_memorySnapshot(PyObject* _unused, PyObject* arg) {
       trace_entry[time_us_s] = te.time_.t_;
       trace_entry[compile_context_s] = te.compile_context_;
       trace_entry[user_metadata_s] = te.user_metadata_;
+      trace_entry[pool_id_s] = te.mempool_;
       trace.append(trace_entry);
     }
     traces.append(trace);
@@ -1387,6 +1386,13 @@ static void registerCudaPluggableAllocator(PyObject* module) {
         storage_impl->release_data_and_set_meta_custom_data_ptr_error_msg_(s);
       });
 
+  m.def(
+      "_clear_storage_data_ptr_access_error_msg", [](size_t storage_impl_ptr) {
+        // NOLINTNEXTLINE(performance-no-int-to-ptr)
+        c10::StorageImpl* storage_impl = (c10::StorageImpl*)storage_impl_ptr;
+        storage_impl->clear_data_ptr_access_error_msg_();
+      });
+
   m.def("_has_Standard_Deleter", [](size_t storage_impl_ptr) {
     // NOLINTNEXTLINE(performance-no-int-to-ptr)
     c10::StorageImpl* storage_impl = (c10::StorageImpl*)storage_impl_ptr;
@@ -1607,6 +1613,20 @@ static PyObject* THCPModule_clearBlasWorkspaces_wrap(
   END_HANDLE_TH_ERRORS
 }
 
+static PyObject* THCPModule_cudnnClearDropoutState_wrap(
+    PyObject* self,
+    PyObject* noargs) {
+  HANDLE_TH_ERRORS
+#if defined(USE_ROCM)
+  // On ROCm, RNNs dispatch to MIOpen, which caches its dropout state buffer in
+  // thread-local storage separate from the cuDNN path.
+  at::native::_miopen_clear_dropout_state();
+#endif
+  at::native::_cudnn_clear_dropout_state();
+  Py_RETURN_NONE;
+  END_HANDLE_TH_ERRORS
+}
+
 PyObject* THCPModule_rocm_is_backward_pass(
     PyObject* _unused,
     PyObject* noargs) {
@@ -2050,6 +2070,10 @@ static struct PyMethodDef _THCPModule_methods[] = {
      THCPModule_clearBlasWorkspaces_wrap,
      METH_NOARGS,
      nullptr},
+    {"_cudnn_clear_dropout_state",
+     THCPModule_cudnnClearDropoutState_wrap,
+     METH_NOARGS,
+     nullptr},
     {"_cuda_isCurrentStreamCapturing",
      THCPModule_isCurrentStreamCapturing_wrap,
      METH_NOARGS,
@@ -2086,6 +2110,10 @@ static struct PyMethodDef _THCPModule_methods[] = {
      METH_O,
      nullptr},
     {"_cuda_memorySnapshot", THCPModule_memorySnapshot, METH_O, nullptr},
+    {"_cuda_allocationTraceback",
+     THCPModule_allocationTraceback,
+     METH_O,
+     nullptr},
     {"_cuda_attach_out_of_memory_observer",
      THCPModule_attachOutOfMemoryObserver,
      METH_O,
@@ -2116,6 +2144,10 @@ static struct PyMethodDef _THCPModule_methods[] = {
      THCPModule_cudaCachingAllocator_enable,
      METH_O,
      nullptr},
+    {"_cuda_cudaCachingAllocator_is_enabled",
+     THCPModule_cudaCachingAllocator_is_enabled,
+     METH_NOARGS,
+     nullptr},
     {"_cuda_getAllocatorBackend",
      THCPModule_getAllocatorBackend,
      METH_NOARGS,
@@ -2123,8 +2155,6 @@ static struct PyMethodDef _THCPModule_methods[] = {
     {"_cuda_synchronize", THCPModule_cudaSynchronize, METH_NOARGS, nullptr},
     {"_cuda_ipc_collect", THCPModule_cudaIPCCollect, METH_NOARGS, nullptr},
     {"_cuda_sleep", THCPModule_cudaSleep, METH_O, nullptr},
-    {"_cuda_lock_mutex", THCPModule_cudaLockMutex, METH_NOARGS, nullptr},
-    {"_cuda_unlock_mutex", THCPModule_cudaUnlockMutex, METH_NOARGS, nullptr},
     {"_cuda_set_sync_debug_mode",
      THCPModule_cudaSetSyncDebugMode,
      METH_O,
diff --git a/torch/csrc/cuda/memory_snapshot.cpp b/torch/csrc/cuda/memory_snapshot.cpp
index 685e2fa38a9a7..9ce2b9bf383c6 100644
--- a/torch/csrc/cuda/memory_snapshot.cpp
+++ b/torch/csrc/cuda/memory_snapshot.cpp
@@ -315,6 +315,7 @@ std::string _memory_snapshot_pickled() {
   IValue time_us_s = "time_us";
   IValue compile_contexts_s = "compile_context";
   IValue user_metadata_s = "user_metadata";
+  IValue pool_id_s = "pool_id";
 
   auto empty_frames = new_list();
 
@@ -439,6 +440,7 @@ std::string _memory_snapshot_pickled() {
         frame_dict.push_back(trace_entry);
       }
       trace_entry.insert(time_us_s, te.time_.t_);
+      trace_entry.insert(pool_id_s, std::tuple<int64_t, int64_t>(te.mempool_));
       trace.push_back(trace_entry);
     }
     traces.push_back(trace);
diff --git a/torch/csrc/distributed/c10d/Backend.hpp b/torch/csrc/distributed/c10d/Backend.hpp
index 8ce9dc1d207dd..bb23e13a2d8e7 100644
--- a/torch/csrc/distributed/c10d/Backend.hpp
+++ b/torch/csrc/distributed/c10d/Backend.hpp
@@ -1,6 +1,8 @@
 #pragma once
 
 #include <memory>
+#include <string>
+#include <unordered_map>
 #include <utility>
 #include <vector>
 
@@ -509,6 +511,26 @@ class TORCH_API Backend : public torch::CustomClassHolder {
   // normal shutdown.
   virtual void shutdown() {}
 
+  // APIs related to memory offload
+  virtual void suspend() {
+    TORCH_CHECK(
+        false,
+        c10::str("Backend ", getBackendName(), " does not support suspend"));
+  }
+
+  virtual void resume() {
+    TORCH_CHECK(
+        false,
+        c10::str("Backend ", getBackendName(), " does not support resume"));
+  }
+
+  virtual std::unordered_map<std::string, uint64_t> getMemoryStats() {
+    TORCH_CHECK(
+        false,
+        c10::str(
+            "Backend ", getBackendName(), " does not support getMemoryStats"));
+  }
+
  protected:
   // Implementations of this interface need to call this to setup
   // appropriate logging etc.
diff --git a/torch/csrc/distributed/c10d/FlightRecorderDetail.hpp b/torch/csrc/distributed/c10d/FlightRecorderDetail.hpp
index 28647b8c50f5a..50aa63ed748d1 100644
--- a/torch/csrc/distributed/c10d/FlightRecorderDetail.hpp
+++ b/torch/csrc/distributed/c10d/FlightRecorderDetail.hpp
@@ -89,13 +89,13 @@ typename FlightRecorder<EventType>::TraceIdentifier FlightRecorder<EventType>::
   if (!enabled_) {
     return TraceIdentifier{std::nullopt, std::nullopt};
   }
+  auto traceback =
+      torch::CapturedTraceback::gather(true, true, capture_cpp_stack_);
+  std::lock_guard<std::mutex> guard(mutex_);
   if (all_pg_status_.find(pg_id) == all_pg_status_.end()) {
     // Current pg_status is not in FR.
     all_pg_status_[pg_id] = std::move(pg_status);
   }
-  auto traceback =
-      torch::CapturedTraceback::gather(true, true, capture_cpp_stack_);
-  std::lock_guard<std::mutex> guard(mutex_);
 
   TORCH_CHECK(
       reset_epoch_start_idx_.find(reset_epoch_) !=
diff --git a/torch/csrc/distributed/c10d/Functional.cpp b/torch/csrc/distributed/c10d/Functional.cpp
index 1284676dae015..fd8112d2b9dba 100644
--- a/torch/csrc/distributed/c10d/Functional.cpp
+++ b/torch/csrc/distributed/c10d/Functional.cpp
@@ -4,6 +4,8 @@
 #include <torch/csrc/distributed/c10d/Functional.hpp>
 #include <torch/csrc/distributed/c10d/GroupRegistry.hpp>
 #include <torch/csrc/distributed/c10d/ProcessGroup.hpp>
+#include <torch/csrc/jit/frontend/schema_type_parser.h>
+#include <torch/custom_class_detail.h>
 #include <utility>
 
 namespace {
@@ -31,7 +33,6 @@ c10d::ReduceOp to_reduce_op(const std::string& reduce_op) {
 at::Tensor allocate_all_gather_output(
     const at::Tensor& input,
     int64_t group_size) {
-  TORCH_CHECK(input.is_contiguous());
   auto output_size = input.sizes().vec();
   if (output_size.empty()) {
     output_size.push_back(group_size);
@@ -46,7 +47,6 @@ at::Tensor allocate_all_gather_output(
 at::Tensor allocate_reduce_scatter_output(
     const at::Tensor& input,
     const int64_t group_size) {
-  TORCH_CHECK(input.is_contiguous());
   auto output_size = input.sizes().vec();
   if (output_size[0] % group_size != 0) {
     LOG(WARNING) << "The first dimension of the reduce_scatter input ("
@@ -67,13 +67,11 @@ at::Tensor& all_reduce_(
     at::Tensor& input,
     // NOLINTNEXTLINE(performance-unnecessary-value-param)
     std::string reduce_op,
-    // NOLINTNEXTLINE(performance-unnecessary-value-param)
-    std::string group_name) {
+    c10::intrusive_ptr<c10d::ProcessGroup> group) {
   c10d::AllreduceOptions opts;
   opts.reduceOp = to_reduce_op(reduce_op);
 
   std::vector<at::Tensor> inputs{input};
-  auto group = c10d::resolve_process_group(group_name);
   auto work = group->allreduce(inputs, opts);
   c10d::register_work(input, work);
   return input;
@@ -82,7 +80,7 @@ at::Tensor& all_reduce_(
 at::Tensor all_reduce(
     const at::Tensor& input,
     std::string reduce_op,
-    std::string group_name) {
+    c10::intrusive_ptr<c10d::ProcessGroup> group) {
   if (input.is_complex()) {
     TORCH_CHECK(
         // TODO - ideally use 'to_reduce_op' helper but it currently errors on
@@ -95,21 +93,44 @@ at::Tensor all_reduce(
   }
   auto input_real = input.is_complex() ? at::view_as_real(input) : input;
   auto output = input_real.clone(at::MemoryFormat::Contiguous);
-  auto output_ret =
-      all_reduce_(output, std::move(reduce_op), std::move(group_name));
+  auto output_ret = all_reduce_(output, std::move(reduce_op), std::move(group));
   return input.is_complex() ? at::view_as_complex(output_ret) : output_ret;
 }
 
+at::Tensor& all_reduce_(
+    at::Tensor& input,
+    std::string reduce_op,
+    std::string group_name) {
+  auto group = c10d::resolve_process_group(std::move(group_name));
+  return all_reduce_(input, std::move(reduce_op), std::move(group));
+}
+
+at::Tensor all_reduce(
+    const at::Tensor& input,
+    std::string reduce_op,
+    std::string group_name) {
+  auto group = c10d::resolve_process_group(std::move(group_name));
+  return all_reduce(input, std::move(reduce_op), std::move(group));
+}
+
 std::vector<at::Tensor> all_reduce_coalesced_(
     std::vector<at::Tensor> inputs,
     // NOLINTNEXTLINE(performance-unnecessary-value-param)
     std::string reduce_op,
     // NOLINTNEXTLINE(performance-unnecessary-value-param)
     std::string group_name) {
+  auto group = c10d::resolve_process_group(std::move(group_name));
+  return all_reduce_coalesced_(inputs, std::move(reduce_op), std::move(group));
+}
+
+std::vector<at::Tensor> all_reduce_coalesced_(
+    std::vector<at::Tensor> inputs,
+    // NOLINTNEXTLINE(performance-unnecessary-value-param)
+    std::string reduce_op,
+    c10::intrusive_ptr<c10d::ProcessGroup> group) {
   c10d::AllreduceCoalescedOptions opts;
   opts.reduceOp = to_reduce_op(reduce_op);
 
-  auto group = c10d::resolve_process_group(group_name);
   auto work = group->allreduce_coalesced(inputs, opts);
   for (const auto& tensor : inputs) {
     c10d::register_work(tensor, work);
@@ -122,28 +143,33 @@ std::vector<at::Tensor> all_reduce_coalesced(
     std::vector<at::Tensor> inputs,
     std::string reduce_op,
     std::string group_name) {
+  auto group = c10d::resolve_process_group(std::move(group_name));
+  return all_reduce_coalesced(inputs, std::move(reduce_op), std::move(group));
+}
+
+std::vector<at::Tensor> all_reduce_coalesced(
+    std::vector<at::Tensor> inputs,
+    std::string reduce_op,
+    c10::intrusive_ptr<c10d::ProcessGroup> group) {
   std::vector<at::Tensor> outputs;
   outputs.reserve(inputs.size());
   for (const auto& tensor : inputs) {
     outputs.push_back(tensor.clone(at::MemoryFormat::Contiguous));
   }
-  return all_reduce_coalesced_(
-      outputs, std::move(reduce_op), std::move(group_name));
+  return all_reduce_coalesced_(outputs, std::move(reduce_op), std::move(group));
 }
 
 std::vector<at::Tensor> all_gather_into_tensor_coalesced(
     std::vector<at::Tensor> inputs,
     int64_t group_size,
-    // NOLINTNEXTLINE(performance-unnecessary-value-param)
-    std::string group_name) {
+    c10::intrusive_ptr<c10d::ProcessGroup> group) {
   std::vector<at::Tensor> outputs;
   outputs.reserve(inputs.size());
-  for (const auto& tensor : inputs) {
-    TORCH_CHECK(tensor.is_contiguous());
+  for (auto& tensor : inputs) {
+    tensor = tensor.contiguous();
     outputs.push_back(allocate_all_gather_output(tensor, group_size));
   }
 
-  auto group = c10d::resolve_process_group(group_name);
   auto work = group->allgather_into_tensor_coalesced(outputs, inputs);
   for (const auto& tensor : outputs) {
     c10d::register_work(tensor, work);
@@ -154,29 +180,53 @@ std::vector<at::Tensor> all_gather_into_tensor_coalesced(
 at::Tensor all_gather_into_tensor(
     const at::Tensor& input,
     int64_t group_size,
-    std::string group_name) {
-  TORCH_CHECK(input.is_contiguous());
+    c10::intrusive_ptr<c10d::ProcessGroup> group) {
   auto real_input = input.is_complex() ? at::view_as_real(input) : input;
   std::vector<at::Tensor> inputs{real_input};
-  auto output = all_gather_into_tensor_coalesced(
-      inputs, group_size, std::move(group_name))[0];
+  auto output =
+      all_gather_into_tensor_coalesced(inputs, group_size, std::move(group))[0];
   return input.is_complex() ? at::view_as_complex(output) : output;
 }
 
 at::Tensor& all_gather_into_tensor_out(
     at::Tensor& input,
     int64_t group_size,
-    const std::string& group_name,
+    c10::intrusive_ptr<c10d::ProcessGroup> group,
     at::Tensor& output) {
-  TORCH_CHECK(input.is_contiguous());
+  auto contig_input = input.contiguous();
   c10d::AllgatherOptions opts;
 
-  auto group = c10d::resolve_process_group(group_name);
-  auto work = group->_allgather_base(output, input, opts);
+  auto work = group->_allgather_base(output, contig_input, opts);
   c10d::register_work(output, work);
   return output;
 }
 
+std::vector<at::Tensor> all_gather_into_tensor_coalesced(
+    std::vector<at::Tensor> inputs,
+    int64_t group_size,
+    std::string group_name) {
+  auto group = c10d::resolve_process_group(std::move(group_name));
+  return all_gather_into_tensor_coalesced(inputs, group_size, std::move(group));
+}
+
+at::Tensor all_gather_into_tensor(
+    const at::Tensor& input,
+    int64_t group_size,
+    std::string group_name) {
+  auto group = c10d::resolve_process_group(std::move(group_name));
+  return all_gather_into_tensor(input, group_size, std::move(group));
+}
+
+at::Tensor& all_gather_into_tensor_out(
+    at::Tensor& input,
+    int64_t group_size,
+    const std::string& group_name,
+    at::Tensor& output) {
+  auto group = c10d::resolve_process_group(std::move(group_name));
+  return all_gather_into_tensor_out(
+      input, group_size, std::move(group), output);
+}
+
 std::vector<at::Tensor> reduce_scatter_tensor_coalesced(
     std::vector<at::Tensor> inputs,
     // NOLINTNEXTLINE(performance-unnecessary-value-param)
@@ -184,16 +234,26 @@ std::vector<at::Tensor> reduce_scatter_tensor_coalesced(
     int64_t group_size,
     // NOLINTNEXTLINE(performance-unnecessary-value-param)
     std::string group_name) {
+  auto group = c10d::resolve_process_group(std::move(group_name));
+  return reduce_scatter_tensor_coalesced(
+      inputs, std::move(reduce_op), group_size, std::move(group));
+}
+
+std::vector<at::Tensor> reduce_scatter_tensor_coalesced(
+    std::vector<at::Tensor> inputs,
+    // NOLINTNEXTLINE(performance-unnecessary-value-param)
+    std::string reduce_op,
+    int64_t group_size,
+    c10::intrusive_ptr<c10d::ProcessGroup> group) {
   c10d::ReduceScatterOptions opts;
   opts.reduceOp = to_reduce_op(reduce_op);
   std::vector<at::Tensor> outputs;
   outputs.reserve(inputs.size());
-  for (const auto& tensor : inputs) {
-    TORCH_CHECK(tensor.is_contiguous());
+  for (auto& tensor : inputs) {
+    tensor = tensor.contiguous();
     outputs.push_back(allocate_reduce_scatter_output(tensor, group_size));
   }
 
-  auto group = c10d::resolve_process_group(group_name);
   auto work = group->reduce_scatter_tensor_coalesced(outputs, inputs, opts);
   for (const auto& tensor : outputs) {
     c10d::register_work(tensor, work);
@@ -206,13 +266,11 @@ static std::vector<at::Tensor> reduce_scatter_tensor_coalesced_out(
     // NOLINTNEXTLINE(performance-unnecessary-value-param)
     std::string reduce_op,
     int64_t group_size,
-    // NOLINTNEXTLINE(performance-unnecessary-value-param)
-    std::string group_name,
+    c10::intrusive_ptr<c10d::ProcessGroup> group,
     std::vector<at::Tensor>& outputs) {
   c10d::ReduceScatterOptions opts;
   opts.reduceOp = to_reduce_op(reduce_op);
 
-  auto group = c10d::resolve_process_group(std::move(group_name));
   auto work = group->reduce_scatter_tensor_coalesced(outputs, inputs, opts);
   for (const auto& tensor : outputs) {
     c10d::register_work(tensor, work);
@@ -225,16 +283,25 @@ at::Tensor reduce_scatter_tensor(
     std::string reduce_op,
     int64_t group_size,
     std::string group_name) {
-  TORCH_CHECK(input.is_contiguous());
+  auto group = c10d::resolve_process_group(std::move(group_name));
+  return reduce_scatter_tensor(
+      input, std::move(reduce_op), group_size, std::move(group));
+}
+
+at::Tensor reduce_scatter_tensor(
+    const at::Tensor& input,
+    std::string reduce_op,
+    int64_t group_size,
+    c10::intrusive_ptr<c10d::ProcessGroup> group) {
   if (input.is_complex()) {
     auto real_input = at::view_as_real(input);
-    std::vector<at::Tensor> inputs{real_input};
+    std::vector<at::Tensor> inputs{std::move(real_input)};
     return at::view_as_complex(reduce_scatter_tensor_coalesced(
-        inputs, std::move(reduce_op), group_size, std::move(group_name))[0]);
+        inputs, std::move(reduce_op), group_size, std::move(group))[0]);
   }
   std::vector<at::Tensor> inputs{input};
   return reduce_scatter_tensor_coalesced(
-      inputs, std::move(reduce_op), group_size, std::move(group_name))[0];
+      inputs, std::move(reduce_op), group_size, std::move(group))[0];
 }
 
 at::Tensor reduce_scatter_tensor_out(
@@ -243,10 +310,21 @@ at::Tensor reduce_scatter_tensor_out(
     int64_t group_size,
     std::string group_name,
     at::Tensor& output) {
-  TORCH_CHECK(input.is_contiguous());
-  if (input.is_complex()) {
+  auto group = c10d::resolve_process_group(std::move(group_name));
+  return reduce_scatter_tensor_out(
+      input, std::move(reduce_op), group_size, std::move(group), output);
+}
+
+at::Tensor reduce_scatter_tensor_out(
+    const at::Tensor& input,
+    std::string reduce_op,
+    int64_t group_size,
+    c10::intrusive_ptr<c10d::ProcessGroup> group,
+    at::Tensor& output) {
+  auto contig_input = input.contiguous();
+  if (contig_input.is_complex()) {
     TORCH_CHECK(output.is_complex())
-    auto real_input = at::view_as_real(input);
+    auto real_input = at::view_as_real(contig_input);
     std::vector<at::Tensor> inputs{std::move(real_input)};
     auto real_output = at::view_as_real(output);
     std::vector<at::Tensor> outputs{std::move(real_output)};
@@ -254,17 +332,13 @@ at::Tensor reduce_scatter_tensor_out(
         inputs,
         std::move(reduce_op),
         group_size,
-        std::move(group_name),
+        std::move(group),
         outputs)[0]);
   }
-  std::vector<at::Tensor> inputs{std::move(input)};
+  std::vector<at::Tensor> inputs{std::move(contig_input)};
   std::vector<at::Tensor> outputs{std::move(output)};
   return reduce_scatter_tensor_coalesced_out(
-      inputs,
-      std::move(reduce_op),
-      group_size,
-      std::move(group_name),
-      outputs)[0];
+      inputs, std::move(reduce_op), group_size, std::move(group), outputs)[0];
 }
 
 at::Tensor all_to_all_single(
@@ -273,6 +347,16 @@ at::Tensor all_to_all_single(
     c10::SymIntArrayRef _input_split_sizes,
     // NOLINTNEXTLINE(performance-unnecessary-value-param)
     std::string group_name) {
+  auto group = c10d::resolve_process_group(std::move(group_name));
+  return all_to_all_single(
+      input, _output_split_sizes, _input_split_sizes, std::move(group));
+}
+
+at::Tensor all_to_all_single(
+    const at::Tensor& input,
+    c10::SymIntArrayRef _output_split_sizes,
+    c10::SymIntArrayRef _input_split_sizes,
+    c10::intrusive_ptr<ProcessGroup> group) {
   std::vector<int64_t> output_split_sizes;
   std::vector<int64_t> input_split_sizes;
   output_split_sizes.reserve(_output_split_sizes.size());
@@ -284,17 +368,16 @@ at::Tensor all_to_all_single(
     input_split_sizes.emplace_back(size.expect_int());
   }
 
-  TORCH_CHECK(input.is_contiguous());
-  std::vector<int64_t> output_sizes = input.sizes().vec();
+  auto contig_input = input.contiguous();
+  std::vector<int64_t> output_sizes = contig_input.sizes().vec();
   output_sizes[0] = std::accumulate(
       output_split_sizes.begin(), output_split_sizes.end(), int64_t(0));
-  auto output = input.new_empty(output_sizes);
+  auto output = contig_input.new_empty(output_sizes);
 
-  auto group = c10d::resolve_process_group(std::move(group_name));
   auto work = group->alltoall_base(
       output,
       // NOLINTNEXTLINE(cppcoreguidelines-pro-type-const-cast)
-      const_cast<at::Tensor&>(input),
+      const_cast<at::Tensor&>(contig_input),
       output_split_sizes,
       input_split_sizes);
   c10d::register_work(output, work);
@@ -303,12 +386,19 @@ at::Tensor all_to_all_single(
 
 // NOLINTNEXTLINE(performance-unnecessary-value-param)
 at::Tensor& broadcast_(at::Tensor& input, int64_t src, std::string group_name) {
+  auto group = c10d::resolve_process_group(std::move(group_name));
+  return broadcast_(input, src, std::move(group));
+}
+
+at::Tensor& broadcast_(
+    at::Tensor& input,
+    int64_t src,
+    c10::intrusive_ptr<c10d::ProcessGroup> group) {
   c10d::BroadcastOptions opts;
   opts.rootRank = src;
   auto input_real = input.is_complex() ? at::view_as_real(input) : input;
   std::vector<at::Tensor> inputs{input_real};
 
-  auto group = c10d::resolve_process_group(group_name);
   auto work = group->broadcast(inputs, opts);
   c10d::register_work(input, work);
   return input;
@@ -318,79 +408,296 @@ at::Tensor broadcast(
     const at::Tensor& input,
     int64_t src,
     std::string group_name) {
+  auto group = c10d::resolve_process_group(std::move(group_name));
+  return broadcast(input, src, std::move(group));
+}
+
+at::Tensor broadcast(
+    const at::Tensor& input,
+    int64_t src,
+    c10::intrusive_ptr<c10d::ProcessGroup> group) {
   auto output = input.clone(at::MemoryFormat::Contiguous);
-  return broadcast_(output, src, std::move(group_name));
+  return broadcast_(output, src, std::move(group));
+}
+
+at::Tensor isend(
+    at::Tensor& input,
+    int64_t dst,
+    int64_t tag,
+    std::string group_name) {
+  auto group = c10d::resolve_process_group(group_name);
+  std::vector<at::Tensor> input_wrap = {input};
+  auto work = group->send(input_wrap, dst, tag);
+  c10d::register_work(input, work);
+  auto placeholder = at::empty({0}, input.options());
+  c10d::register_work(placeholder, work);
+  return placeholder;
+}
+
+at::Tensor irecv(
+    at::Tensor& output,
+    int64_t src,
+    int64_t tag,
+    std::string group_name) {
+  auto group = c10d::resolve_process_group(group_name);
+  std::vector<at::Tensor> output_wrap = {output};
+  auto work = group->recv(output_wrap, src, tag);
+  c10d::register_work(output, work);
+  return output;
+}
+
+std::vector<at::Tensor> batch_p2p_ops(
+    std::vector<std::string> op_list,
+    std::vector<int64_t> peer_list,
+    std::vector<int64_t> tag_list,
+    std::vector<at::Tensor> tensors,
+    std::string group_name) {
+  const uint64_t N = op_list.size();
+  TORCH_CHECK(tensors.size() == N, "");
+  TORCH_CHECK(peer_list.size() == N, "");
+  TORCH_CHECK(tag_list.size() == N, "");
+  if (N == 0)
+    return {at::Tensor()};
+  auto group = c10d::resolve_process_group(group_name);
+  auto device = tensors[0].device().type();
+  auto backend = group->getBackend(device);
+  bool should_coalesce = backend->supportsCoalescing();
+  if (should_coalesce) {
+    group->startCoalescing(device);
+  }
+  std::vector<c10::intrusive_ptr<c10d::Work>> works;
+  std::vector<at::Tensor> result_tensors;
+  works.reserve(N);
+  result_tensors.reserve(N);
+  for (uint32_t i = 0; i < N; ++i) {
+    c10::intrusive_ptr<c10d::Work> work;
+    at::Tensor t = tensors[i];
+    std::vector<at::Tensor> tt{t};
+    if (op_list[i] == "isend") {
+      work = group->send(
+          tt,
+          static_cast<int64_t>(peer_list[i]),
+          static_cast<int64_t>(tag_list[i]));
+      auto placeholder = at::empty({0}, t.options());
+      if (work) {
+        c10d::register_work(t, work);
+        c10d::register_work(placeholder, work);
+        works.push_back(std::move(work));
+      }
+      result_tensors.push_back(std::move(placeholder));
+    } else if (op_list[i] == "irecv") {
+      work = group->recv(
+          tt,
+          static_cast<int64_t>(peer_list[i]),
+          static_cast<int64_t>(tag_list[i]));
+      if (work) {
+        c10d::register_work(t, work);
+        works.push_back(std::move(work));
+      }
+      result_tensors.push_back(std::move(t));
+    } else {
+      TORCH_CHECK(false, "Unsupported async op " + op_list[i]);
+    }
+  }
+  if (should_coalesce) {
+    auto work = group->endCoalescing(device);
+    if (!work)
+      TORCH_CHECK(
+          false,
+          "The coalesced work object returned from group->endCoalescing() is empty");
+    for (auto tensor : result_tensors) {
+      c10d::register_work(tensor, work);
+    }
+  }
+  return result_tensors;
 }
 
 } // namespace c10d
 
+namespace {
+
+c10::intrusive_ptr<c10d::ProcessGroup> get_process_group(
+    const c10::IValue& group_name,
+    const char* func_name) {
+  if (group_name.isString()) {
+    return c10d::resolve_process_group(group_name.toStringRef());
+  } else if (group_name.isCapsule()) {
+    return c10::static_intrusive_pointer_cast<c10d::ProcessGroup>(
+        group_name.toCapsule());
+  } else {
+    TORCH_CHECK(
+        false,
+        func_name,
+        "(): argument 'group_name' must be either a string (group name) "
+        "or a ProcessGroup object, but got ",
+        group_name.type()->str());
+  }
+}
+
+// all_to_all_single_dispatch is kept as a named function because it is
+// referenced via decltype inside the AllToAllSingle autograd class.
+at::Tensor all_to_all_single_dispatch(
+    const at::Tensor& input,
+    c10::SymIntArrayRef output_split_sizes,
+    c10::SymIntArrayRef input_split_sizes,
+    const c10::IValue& group_name) {
+  return c10d::all_to_all_single(
+      input,
+      output_split_sizes,
+      input_split_sizes,
+      get_process_group(group_name, "all_to_all_single"));
+}
+
+} // namespace
+
 TORCH_LIBRARY(_c10d_functional, m) {
   m.def(
-      "all_reduce(Tensor input, str reduce_op, str group_name) -> Tensor",
+      "all_reduce(Tensor input, str reduce_op, Any group_name) -> Tensor",
       torch::dispatch(
-          c10::DispatchKey::CompositeExplicitAutograd, c10d::all_reduce),
+          c10::DispatchKey::CompositeExplicitAutograd,
+          [](const at::Tensor& input,
+             std::string reduce_op,
+             const c10::IValue& group) {
+            return c10d::all_reduce(
+                input,
+                std::move(reduce_op),
+                get_process_group(group, "all_reduce"));
+          }),
       {at::Tag::pt2_compliant_tag});
 
   m.def(
-      "all_reduce_(Tensor(a!) input, str reduce_op, str group_name) -> Tensor(a!)",
+      "all_reduce_(Tensor(a!) input, str reduce_op, Any group_name) -> Tensor(a!)",
       torch::dispatch(
-          c10::DispatchKey::CompositeExplicitAutograd, c10d::all_reduce_),
+          c10::DispatchKey::CompositeExplicitAutograd,
+          [](at::Tensor& input, std::string reduce_op, const c10::IValue& group)
+              -> at::Tensor& {
+            return c10d::all_reduce_(
+                input,
+                std::move(reduce_op),
+                get_process_group(group, "all_reduce_"));
+          }),
       {at::Tag::pt2_compliant_tag});
 
   m.def(
-      "all_reduce_coalesced(Tensor[] inputs, str reduce_op, str group_name) -> Tensor[]",
+      "all_reduce_coalesced(Tensor[] inputs, str reduce_op, Any group_name) -> Tensor[]",
       torch::dispatch(
           c10::DispatchKey::CompositeExplicitAutograd,
-          c10d::all_reduce_coalesced),
+          [](std::vector<at::Tensor> inputs,
+             std::string reduce_op,
+             const c10::IValue& group) {
+            return c10d::all_reduce_coalesced(
+                inputs,
+                std::move(reduce_op),
+                get_process_group(group, "all_reduce_coalesced"));
+          }),
       {at::Tag::pt2_compliant_tag});
 
   m.def(
-      "all_reduce_coalesced_(Tensor[](a!) inputs, str reduce_op, str group_name) -> Tensor[](a!)",
+      "all_reduce_coalesced_(Tensor[](a!) inputs, str reduce_op, Any group_name) -> Tensor[](a!)",
       torch::dispatch(
           c10::DispatchKey::CompositeExplicitAutograd,
-          c10d::all_reduce_coalesced_),
+          [](std::vector<at::Tensor> inputs,
+             std::string reduce_op,
+             const c10::IValue& group) {
+            return c10d::all_reduce_coalesced_(
+                inputs,
+                std::move(reduce_op),
+                get_process_group(group, "all_reduce_coalesced_"));
+          }),
       {at::Tag::pt2_compliant_tag});
 
   m.def(
-      "all_gather_into_tensor_out(Tensor input, int group_size, str group_name, *, Tensor(a!) out) -> Tensor(a!)",
+      "all_gather_into_tensor_out(Tensor input, int group_size, Any group_name, *, Tensor(a!) out) -> Tensor(a!)",
       torch::dispatch(
           c10::DispatchKey::CompositeExplicitAutograd,
-          c10d::all_gather_into_tensor_out),
+          [](at::Tensor& input,
+             int64_t group_size,
+             const c10::IValue& group,
+             at::Tensor& output) -> at::Tensor& {
+            return c10d::all_gather_into_tensor_out(
+                input,
+                group_size,
+                get_process_group(group, "all_gather_into_tensor_out"),
+                output);
+          }),
       {at::Tag::pt2_compliant_tag, at::Tag::needs_contiguous_strides});
 
   m.def(
-      "all_gather_into_tensor(Tensor input, int group_size, str group_name) -> Tensor",
+      "all_gather_into_tensor(Tensor input, int group_size, Any group_name) -> Tensor",
       torch::dispatch(
           c10::DispatchKey::CompositeExplicitAutograd,
-          c10d::all_gather_into_tensor),
+          [](const at::Tensor& input,
+             int64_t group_size,
+             const c10::IValue& group) {
+            return c10d::all_gather_into_tensor(
+                input,
+                group_size,
+                get_process_group(group, "all_gather_into_tensor"));
+          }),
       {at::Tag::pt2_compliant_tag, at::Tag::needs_contiguous_strides});
 
   m.def(
-      "all_gather_into_tensor_coalesced(Tensor[] inputs, int group_size, str group_name) -> Tensor[]",
+      "all_gather_into_tensor_coalesced(Tensor[] inputs, int group_size, Any group_name) -> Tensor[]",
       torch::dispatch(
           c10::DispatchKey::CompositeExplicitAutograd,
-          c10d::all_gather_into_tensor_coalesced),
+          [](std::vector<at::Tensor> inputs,
+             int64_t group_size,
+             const c10::IValue& group) {
+            return c10d::all_gather_into_tensor_coalesced(
+                inputs,
+                group_size,
+                get_process_group(group, "all_gather_into_tensor_coalesced"));
+          }),
       {at::Tag::pt2_compliant_tag, at::Tag::needs_contiguous_strides});
 
   m.def(
-      "reduce_scatter_tensor(Tensor input, str reduce_op, int group_size, str group_name) -> Tensor",
+      "reduce_scatter_tensor(Tensor input, str reduce_op, int group_size, Any group_name) -> Tensor",
       torch::dispatch(
           c10::DispatchKey::CompositeExplicitAutograd,
-          c10d::reduce_scatter_tensor),
+          [](const at::Tensor& input,
+             std::string reduce_op,
+             int64_t group_size,
+             const c10::IValue& group) {
+            return c10d::reduce_scatter_tensor(
+                input,
+                std::move(reduce_op),
+                group_size,
+                get_process_group(group, "reduce_scatter_tensor"));
+          }),
       {at::Tag::pt2_compliant_tag, at::Tag::needs_contiguous_strides});
 
   m.def(
-      "reduce_scatter_tensor_out(Tensor input, str reduce_op, int group_size, str group_name, *, Tensor(a!) out) -> Tensor(a!)",
+      "reduce_scatter_tensor_out(Tensor input, str reduce_op, int group_size, Any group_name, *, Tensor(a!) out) -> Tensor(a!)",
       torch::dispatch(
           c10::DispatchKey::CompositeExplicitAutograd,
-          c10d::reduce_scatter_tensor_out),
+          [](const at::Tensor& input,
+             std::string reduce_op,
+             int64_t group_size,
+             const c10::IValue& group,
+             at::Tensor& output) {
+            return c10d::reduce_scatter_tensor_out(
+                input,
+                std::move(reduce_op),
+                group_size,
+                get_process_group(group, "reduce_scatter_tensor_out"),
+                output);
+          }),
       {at::Tag::pt2_compliant_tag, at::Tag::needs_contiguous_strides});
 
   m.def(
-      "reduce_scatter_tensor_coalesced(Tensor[] inputs, str reduce_op, int group_size, str group_name) -> Tensor[]",
+      "reduce_scatter_tensor_coalesced(Tensor[] inputs, str reduce_op, int group_size, Any group_name) -> Tensor[]",
       torch::dispatch(
           c10::DispatchKey::CompositeExplicitAutograd,
-          c10d::reduce_scatter_tensor_coalesced),
+          [](std::vector<at::Tensor> inputs,
+             std::string reduce_op,
+             int64_t group_size,
+             const c10::IValue& group) {
+            return c10d::reduce_scatter_tensor_coalesced(
+                inputs,
+                std::move(reduce_op),
+                group_size,
+                get_process_group(group, "reduce_scatter_tensor_coalesced"));
+          }),
       {at::Tag::pt2_compliant_tag, at::Tag::needs_contiguous_strides});
 
   m.def(
@@ -398,21 +705,31 @@ TORCH_LIBRARY(_c10d_functional, m) {
       "Tensor input, "
       "SymInt[] output_split_sizes, "
       "SymInt[] input_split_sizes, "
-      "str group_name) -> Tensor",
+      "Any group_name) -> Tensor",
       torch::dispatch(
-          c10::DispatchKey::CompositeExplicitAutograd, c10d::all_to_all_single),
+          c10::DispatchKey::CompositeExplicitAutograd,
+          all_to_all_single_dispatch),
       {at::Tag::pt2_compliant_tag, at::Tag::needs_contiguous_strides});
 
   m.def(
-      "broadcast(Tensor input, int src, str group_name) -> Tensor",
+      "broadcast(Tensor input, int src, Any group_name) -> Tensor",
       torch::dispatch(
-          c10::DispatchKey::CompositeExplicitAutograd, c10d::broadcast),
+          c10::DispatchKey::CompositeExplicitAutograd,
+          [](const at::Tensor& input, int64_t src, const c10::IValue& group) {
+            return c10d::broadcast(
+                input, src, get_process_group(group, "broadcast"));
+          }),
       {at::Tag::pt2_compliant_tag});
 
   m.def(
-      "broadcast_(Tensor(a!) input, int src, str group_name) -> Tensor(a!)",
+      "broadcast_(Tensor(a!) input, int src, Any group_name) -> Tensor(a!)",
       torch::dispatch(
-          c10::DispatchKey::CompositeExplicitAutograd, c10d::broadcast_),
+          c10::DispatchKey::CompositeExplicitAutograd,
+          [](at::Tensor& input, int64_t src, const c10::IValue& group)
+              -> at::Tensor& {
+            return c10d::broadcast_(
+                input, src, get_process_group(group, "broadcast_"));
+          }),
       {at::Tag::pt2_compliant_tag});
 
   m.def(
@@ -420,8 +737,130 @@ TORCH_LIBRARY(_c10d_functional, m) {
       torch::dispatch(
           c10::DispatchKey::CompositeExplicitAutograd, c10d::wait_tensor),
       {at::Tag::pt2_compliant_tag});
+  m.def(
+      "isend(Tensor tensor, int dst, int tag, str group_name) -> Tensor",
+      torch::dispatch(c10::DispatchKey::CompositeExplicitAutograd, c10d::isend),
+      {at::Tag::pt2_compliant_tag});
+
+  m.def(
+      "irecv(Tensor tensor, int src, int tag, str group_name) -> Tensor",
+      torch::dispatch(c10::DispatchKey::CompositeExplicitAutograd, c10d::irecv),
+      {at::Tag::pt2_compliant_tag});
+
+  m.def(
+      "batch_p2p_ops(str[] op_list, int[] peer_list,"
+      "int[] tag_list, Tensor[] tensors, str group_name)"
+      "-> Tensor[]",
+      torch::dispatch(
+          c10::DispatchKey::CompositeExplicitAutograd, c10d::batch_p2p_ops),
+      {at::Tag::pt2_compliant_tag});
 }
 
+namespace {
+class BatchP2P : public torch::autograd::Function<BatchP2P> {
+ public:
+  static torch::autograd::Variable forward(
+      torch::autograd::AutogradContext* ctx,
+      std::vector<std::string> op_list,
+      std::vector<int64_t> peer_list,
+      std::vector<int64_t> tag_list,
+      std::vector<at::Tensor> tensors,
+      std::string group_name) {
+    auto work = c10::Dispatcher::singleton()
+                    .findSchemaOrThrow("_c10d_functional::batch_p2p_ops", "")
+                    .typed<decltype(c10d::batch_p2p_ops)>()
+                    .call(op_list, peer_list, tag_list, tensors, group_name);
+    return {at::Tensor()};
+  }
+
+  static torch::autograd::variable_list backward(
+      torch::autograd::AutogradContext* ctx,
+      torch::autograd::variable_list& grad_out_list) {
+    return {at::Tensor(), at::Tensor(), at::Tensor(), at::Tensor()};
+  }
+};
+} // namespace
+
+namespace {
+class Isend : public torch::autograd::Function<Isend> {
+ public:
+  static torch::autograd::Variable forward(
+      torch::autograd::AutogradContext* ctx,
+      at::Tensor& input,
+      int64_t dst,
+      int64_t tag,
+      std::string group_name) {
+    ctx->saved_data["dst"] = dst;
+    ctx->saved_data["tag"] = tag;
+    ctx->saved_data["group_name"] = group_name;
+
+    auto group = c10d::resolve_process_group(group_name);
+    ctx->saved_data["src"] = group->getRank();
+    auto work = c10::Dispatcher::singleton()
+                    .findSchemaOrThrow("_c10d_functional::isend", "")
+                    .typed<decltype(c10d::isend)>()
+                    .call(input, dst, tag, group_name);
+    return at::Tensor();
+  }
+
+  static torch::autograd::variable_list backward(
+      torch::autograd::AutogradContext* ctx,
+      torch::autograd::variable_list& grad_out_list) {
+    const std::string& group_name = ctx->saved_data["group_name"].toStringRef();
+    const int64_t src = ctx->saved_data["dst"].toInt();
+    const int64_t tag = ctx->saved_data["tag"].toInt();
+
+    DCHECK(grad_out_list.size() == 1);
+    auto& grad_out = grad_out_list[0];
+    auto out = c10::Dispatcher::singleton()
+                   .findSchemaOrThrow("_c10d_functional::irecv", "")
+                   .typed<decltype(c10d::irecv)>()
+                   .call(grad_out, src, tag, group_name);
+    return {out, at::Tensor(), at::Tensor(), at::Tensor()};
+  }
+};
+} // namespace
+
+namespace {
+class Irecv : public torch::autograd::Function<Irecv> {
+ public:
+  static torch::autograd::Variable forward(
+      torch::autograd::AutogradContext* ctx,
+      at::Tensor& input,
+      int64_t src,
+      int64_t tag,
+      std::string group_name) {
+    ctx->saved_data["src"] = src;
+    ctx->saved_data["tag"] = tag;
+    ctx->saved_data["group_name"] = group_name;
+
+    auto group = c10d::resolve_process_group(group_name);
+    ctx->saved_data["my_rank"] = group->getRank();
+    auto work = c10::Dispatcher::singleton()
+                    .findSchemaOrThrow("_c10d_functional::irecv", "")
+                    .typed<decltype(c10d::irecv)>()
+                    .call(input, src, tag, group_name);
+    return at::Tensor();
+  }
+
+  static torch::autograd::variable_list backward(
+      torch::autograd::AutogradContext* ctx,
+      torch::autograd::variable_list& grad_out_list) {
+    const std::string& group_name = ctx->saved_data["group_name"].toStringRef();
+    const int64_t dst = ctx->saved_data["my_rank"].toInt();
+    const int64_t tag = ctx->saved_data["tag"].toInt();
+
+    DCHECK(grad_out_list.size() == 1);
+    auto& grad_out = grad_out_list[0];
+    auto out = c10::Dispatcher::singleton()
+                   .findSchemaOrThrow("_c10d_functional::isend", "")
+                   .typed<decltype(c10d::isend)>()
+                   .call(grad_out, dst, tag, group_name);
+    return {out, at::Tensor(), at::Tensor(), at::Tensor()};
+  }
+};
+} // namespace
+
 namespace {
 class AllToAllSingle : public torch::autograd::Function<AllToAllSingle> {
  public:
@@ -432,8 +871,7 @@ class AllToAllSingle : public torch::autograd::Function<AllToAllSingle> {
       at::SymIntArrayRef output_split_sizes,
       // NOLINTNEXTLINE(performance-unnecessary-value-param)
       at::SymIntArrayRef input_split_sizes,
-      // NOLINTNEXTLINE(performance-unnecessary-value-param)
-      std::string group_name) {
+      const c10::IValue& group_name) {
     // swap sizes for backwards pass
     ctx->saved_data["output_split_sizes"] = input_split_sizes.vec();
     ctx->saved_data["input_split_sizes"] = output_split_sizes.vec();
@@ -441,7 +879,7 @@ class AllToAllSingle : public torch::autograd::Function<AllToAllSingle> {
 
     return c10::Dispatcher::singleton()
         .findSchemaOrThrow("_c10d_functional::all_to_all_single", "")
-        .typed<decltype(c10d::all_to_all_single)>()
+        .typed<decltype(all_to_all_single_dispatch)>()
         .call(input, output_split_sizes, input_split_sizes, group_name);
   }
 
@@ -452,7 +890,7 @@ class AllToAllSingle : public torch::autograd::Function<AllToAllSingle> {
         ctx->saved_data["output_split_sizes"].toSymIntVector();
     std::vector<c10::SymInt> input_split_sizes =
         ctx->saved_data["input_split_sizes"].toSymIntVector();
-    const std::string& group_name = ctx->saved_data["group_name"].toStringRef();
+    auto group_name = ctx->saved_data["group_name"];
 
     DCHECK(grad_out_list.size() == 1);
     auto grad_out = grad_out_list[0].contiguous();
@@ -460,7 +898,7 @@ class AllToAllSingle : public torch::autograd::Function<AllToAllSingle> {
     auto out =
         c10::Dispatcher::singleton()
             .findSchemaOrThrow("_c10d_functional::all_to_all_single", "")
-            .typed<decltype(c10d::all_to_all_single)>()
+            .typed<decltype(all_to_all_single_dispatch)>()
             .call(grad_out, output_split_sizes, input_split_sizes, group_name);
 
     // do an explicit wait to avoid cuda stream issues
@@ -478,7 +916,7 @@ at::Tensor all_to_all_single_autograd(
     const at::Tensor& input,
     at::SymIntArrayRef output_split_sizes,
     at::SymIntArrayRef input_split_sizes,
-    const std::string& group_name) {
+    const c10::IValue& group_name) {
   return AllToAllSingle::apply(
       input, output_split_sizes, input_split_sizes, group_name);
 }
@@ -491,7 +929,7 @@ class ReduceScatterTensor
       const at::Tensor& input,
       const std::string& reduce_op,
       int64_t group_size,
-      const std::string& group_name) {
+      const c10::IValue& group_name) {
     TORCH_CHECK(reduce_op == "sum", "Only sum reduce op is supported");
 
     ctx->saved_data["group_size"] = group_size;
@@ -499,7 +937,8 @@ class ReduceScatterTensor
 
     return c10::Dispatcher::singleton()
         .findSchemaOrThrow("_c10d_functional::reduce_scatter_tensor", "")
-        .typed<decltype(c10d::reduce_scatter_tensor)>()
+        .typed<at::Tensor(
+            const at::Tensor&, std::string, int64_t, const c10::IValue&)>()
         .call(input, reduce_op, group_size, group_name);
   }
 
@@ -507,7 +946,7 @@ class ReduceScatterTensor
       torch::autograd::AutogradContext* ctx,
       const torch::autograd::variable_list& grad_out_list) {
     const int64_t group_size = ctx->saved_data["group_size"].toInt();
-    const std::string& group_name = ctx->saved_data["group_name"].toStringRef();
+    auto group_name = ctx->saved_data["group_name"];
 
     DCHECK(grad_out_list.size() == 1);
     const auto& grad_out = grad_out_list[0];
@@ -515,7 +954,7 @@ class ReduceScatterTensor
     auto out =
         c10::Dispatcher::singleton()
             .findSchemaOrThrow("_c10d_functional::all_gather_into_tensor", "")
-            .typed<decltype(c10d::all_gather_into_tensor)>()
+            .typed<at::Tensor(const at::Tensor&, int64_t, const c10::IValue&)>()
             .call(grad_out, group_size, group_name);
 
     // do an explicit wait to avoid cuda stream issues
@@ -538,7 +977,7 @@ at::Tensor reduce_scatter_tensor_autograd(
     const at::Tensor& input,
     const std::string& reduce_op,
     int64_t group_size,
-    const std::string& group_name) {
+    const c10::IValue& group_name) {
   return ReduceScatterTensor::apply(input, reduce_op, group_size, group_name);
 }
 
@@ -549,13 +988,13 @@ class AllGatherIntoTensor
       torch::autograd::AutogradContext* ctx,
       const at::Tensor& input,
       int64_t group_size,
-      const std::string& group_name) {
+      const c10::IValue& group_name) {
     ctx->saved_data["group_size"] = group_size;
     ctx->saved_data["group_name"] = group_name;
 
     return c10::Dispatcher::singleton()
         .findSchemaOrThrow("_c10d_functional::all_gather_into_tensor", "")
-        .typed<decltype(c10d::all_gather_into_tensor)>()
+        .typed<at::Tensor(const at::Tensor&, int64_t, const c10::IValue&)>()
         .call(input, group_size, group_name);
   }
 
@@ -563,7 +1002,7 @@ class AllGatherIntoTensor
       torch::autograd::AutogradContext* ctx,
       const torch::autograd::variable_list& grad_out_list) {
     const int64_t group_size = ctx->saved_data["group_size"].toInt();
-    const std::string& group_name = ctx->saved_data["group_name"].toStringRef();
+    auto group_name = ctx->saved_data["group_name"];
 
     DCHECK(grad_out_list.size() == 1);
     const auto& grad_out = grad_out_list[0];
@@ -571,7 +1010,8 @@ class AllGatherIntoTensor
     auto out =
         c10::Dispatcher::singleton()
             .findSchemaOrThrow("_c10d_functional::reduce_scatter_tensor", "")
-            .typed<decltype(c10d::reduce_scatter_tensor)>()
+            .typed<at::Tensor(
+                const at::Tensor&, std::string, int64_t, const c10::IValue&)>()
             .call(grad_out, "sum", group_size, group_name);
 
     // do an explicit wait to avoid cuda stream issues
@@ -592,7 +1032,7 @@ class AllGatherIntoTensor
 at::Tensor all_gather_into_tensor_autograd(
     const at::Tensor& input,
     int64_t group_size,
-    const std::string& group_name) {
+    const c10::IValue& group_name) {
   return AllGatherIntoTensor::apply(input, group_size, group_name);
 }
 
@@ -604,15 +1044,23 @@ TORCH_LIBRARY(_c10d_functional_autograd, m) {
       "Tensor input, "
       "SymInt[] output_split_sizes, "
       "SymInt[] input_split_sizes, "
-      "str group_name) -> Tensor",
-      torch::dispatch(c10::DispatchKey::Autograd, ::all_to_all_single_autograd),
+      "Any group_name) -> Tensor",
+      torch::dispatch(
+          c10::DispatchKey::Autograd,
+          [](const at::Tensor& input,
+             c10::SymIntArrayRef output_split_sizes,
+             c10::SymIntArrayRef input_split_sizes,
+             const c10::IValue& group) {
+            return all_to_all_single_autograd(
+                input, output_split_sizes, input_split_sizes, group);
+          }),
       {at::Tag::pt2_compliant_tag});
   m.def(
       "reduce_scatter_tensor("
       "Tensor input, "
       "str reduce_op, "
       "int group_size, "
-      "str group_name) -> Tensor",
+      "Any group_name) -> Tensor",
       torch::dispatch(
           c10::DispatchKey::Autograd, ::reduce_scatter_tensor_autograd),
       {at::Tag::pt2_compliant_tag});
@@ -620,7 +1068,7 @@ TORCH_LIBRARY(_c10d_functional_autograd, m) {
       "all_gather_into_tensor("
       "Tensor input, "
       "int group_size, "
-      "str group_name) -> Tensor",
+      "Any group_name) -> Tensor",
       torch::dispatch(
           c10::DispatchKey::Autograd, ::all_gather_into_tensor_autograd),
       {at::Tag::pt2_compliant_tag});
@@ -633,8 +1081,7 @@ at::Tensor shard_dim_alltoall(
     const at::Tensor& input,
     int64_t gather_dim,
     int64_t shard_dim,
-    const std::string& group_name) {
-  auto group = c10d::resolve_process_group(group_name);
+    c10::intrusive_ptr<c10d::ProcessGroup> group) {
   auto group_size = group->getSize();
   std::vector<int64_t> input_sizes = input.sizes().vec();
   std::vector<int64_t> output_sizes = input.sizes().vec();
@@ -676,13 +1123,24 @@ at::Tensor shard_dim_alltoall(
   return input.is_complex() ? at::view_as_complex(output).view(output_sizes)
                             : output.view(output_sizes);
 }
+
 } // namespace
 
 // DTensor comm op registry
 TORCH_LIBRARY(_dtensor, m) {
   m.def(
-      "shard_dim_alltoall(Tensor input, int gather_dim, int shard_dim, str group_name) -> Tensor",
+      "shard_dim_alltoall(Tensor input, int gather_dim, int shard_dim, Any group_name) -> Tensor",
       torch::dispatch(
-          c10::DispatchKey::CompositeExplicitAutograd, ::shard_dim_alltoall),
+          c10::DispatchKey::CompositeExplicitAutograd,
+          [](const at::Tensor& input,
+             int64_t gather_dim,
+             int64_t shard_dim,
+             const c10::IValue& group) {
+            return shard_dim_alltoall(
+                input,
+                gather_dim,
+                shard_dim,
+                get_process_group(group, "shard_dim_alltoall"));
+          }),
       {at::Tag::pt2_compliant_tag});
 }
diff --git a/torch/csrc/distributed/c10d/Functional.hpp b/torch/csrc/distributed/c10d/Functional.hpp
index 9c0ccbe1b0f2c..72bac44406b1b 100644
--- a/torch/csrc/distributed/c10d/Functional.hpp
+++ b/torch/csrc/distributed/c10d/Functional.hpp
@@ -9,11 +9,27 @@ C10_EXPORT at::Tensor& all_reduce_(
     std::string reduce_op,
     std::string group_name);
 
+C10_EXPORT at::Tensor& all_reduce_(
+    at::Tensor& input,
+    std::string reduce_op,
+    c10::intrusive_ptr<ProcessGroup> group);
+
 C10_EXPORT at::Tensor all_reduce(
     const at::Tensor& input,
     std::string reduce_op,
     std::string group_name);
 
+C10_EXPORT at::Tensor all_reduce(
+    const at::Tensor& input,
+    std::string reduce_op,
+    c10::intrusive_ptr<ProcessGroup> group);
+
+C10_EXPORT std::vector<at::Tensor> all_reduce_coalesced_(
+    std::vector<at::Tensor> inputs,
+    // NOLINTNEXTLINE(performance-unnecessary-value-param)
+    std::string reduce_op,
+    c10::intrusive_ptr<ProcessGroup> group);
+
 C10_EXPORT std::vector<at::Tensor> all_reduce_coalesced_(
     std::vector<at::Tensor> inputs,
     // NOLINTNEXTLINE(performance-unnecessary-value-param)
@@ -21,6 +37,11 @@ C10_EXPORT std::vector<at::Tensor> all_reduce_coalesced_(
     // NOLINTNEXTLINE(performance-unnecessary-value-param)
     std::string group_name);
 
+C10_EXPORT std::vector<at::Tensor> all_reduce_coalesced(
+    std::vector<at::Tensor> inputs,
+    std::string reduce_op,
+    c10::intrusive_ptr<ProcessGroup> group);
+
 C10_EXPORT std::vector<at::Tensor> all_reduce_coalesced(
     // NOLINTNEXTLINE(performance-unnecessary-value-param)
     std::vector<at::Tensor> inputs,
@@ -33,17 +54,40 @@ C10_EXPORT std::vector<at::Tensor> all_gather_into_tensor_coalesced(
     // NOLINTNEXTLINE(performance-unnecessary-value-param)
     std::string group_name);
 
+C10_EXPORT std::vector<at::Tensor> all_gather_into_tensor_coalesced(
+    std::vector<at::Tensor> inputs,
+    int64_t group_size,
+    c10::intrusive_ptr<ProcessGroup> group);
+
 C10_EXPORT at::Tensor all_gather_into_tensor(
     const at::Tensor& input,
     int64_t group_size,
     std::string group_name);
 
+C10_EXPORT at::Tensor all_gather_into_tensor(
+    const at::Tensor& input,
+    int64_t group_size,
+    c10::intrusive_ptr<ProcessGroup> group);
+
 C10_EXPORT at::Tensor& all_gather_into_tensor_out(
     at::Tensor& input,
     int64_t group_size,
     const std::string& group_name,
     at::Tensor& output);
 
+C10_EXPORT at::Tensor& all_gather_into_tensor_out(
+    at::Tensor& input,
+    int64_t group_size,
+    c10::intrusive_ptr<ProcessGroup> group,
+    at::Tensor& output);
+
+C10_EXPORT std::vector<at::Tensor> reduce_scatter_tensor_coalesced(
+    std::vector<at::Tensor> inputs,
+    // NOLINTNEXTLINE(performance-unnecessary-value-param)
+    std::string reduce_op,
+    int64_t group_size,
+    c10::intrusive_ptr<ProcessGroup> group);
+
 C10_EXPORT std::vector<at::Tensor> reduce_scatter_tensor_coalesced(
     std::vector<at::Tensor> inputs,
     // NOLINTNEXTLINE(performance-unnecessary-value-param)
@@ -52,12 +96,25 @@ C10_EXPORT std::vector<at::Tensor> reduce_scatter_tensor_coalesced(
     // NOLINTNEXTLINE(performance-unnecessary-value-param)
     std::string group_name);
 
+C10_EXPORT at::Tensor reduce_scatter_tensor(
+    const at::Tensor& input,
+    std::string reduce_op,
+    int64_t group_size,
+    c10::intrusive_ptr<ProcessGroup> group);
+
 C10_EXPORT at::Tensor reduce_scatter_tensor(
     const at::Tensor& input,
     std::string reduce_op,
     int64_t group_size,
     std::string group_name);
 
+C10_EXPORT at::Tensor reduce_scatter_tensor_out(
+    const at::Tensor& input,
+    std::string reduce_op,
+    int64_t group_size,
+    c10::intrusive_ptr<ProcessGroup> group,
+    at::Tensor& output);
+
 C10_EXPORT at::Tensor reduce_scatter_tensor_out(
     const at::Tensor& input,
     std::string reduce_op,
@@ -72,6 +129,17 @@ C10_EXPORT at::Tensor all_to_all_single(
     // NOLINTNEXTLINE(performance-unnecessary-value-param)
     std::string group_name);
 
+C10_EXPORT at::Tensor all_to_all_single(
+    const at::Tensor& input,
+    at::SymIntArrayRef output_split_sizes,
+    at::SymIntArrayRef input_split_sizes,
+    c10::intrusive_ptr<ProcessGroup> group);
+
+C10_EXPORT at::Tensor& broadcast_(
+    at::Tensor& input,
+    int64_t src,
+    c10::intrusive_ptr<ProcessGroup> group);
+
 C10_EXPORT at::Tensor& broadcast_(
     at::Tensor& input,
     int64_t src,
@@ -80,6 +148,30 @@ C10_EXPORT at::Tensor& broadcast_(
 C10_EXPORT at::Tensor broadcast(
     const at::Tensor& input,
     int64_t src,
+    c10::intrusive_ptr<ProcessGroup> group);
+
+C10_EXPORT at::Tensor broadcast(
+    const at::Tensor& input,
+    int64_t src,
+    std::string group_name);
+
+C10_EXPORT at::Tensor isend(
+    at::Tensor& send_buf,
+    int64_t dst,
+    int64_t tag,
+    std::string group_name);
+
+C10_EXPORT at::Tensor irecv(
+    at::Tensor& recv_buf,
+    int64_t src,
+    int64_t tag,
+    std::string group_name);
+
+C10_EXPORT std::vector<at::Tensor> batch_p2p_ops(
+    std::vector<std::string> op_list,
+    std::vector<int64_t> peer_list,
+    std::vector<int64_t> tag_list,
+    std::vector<at::Tensor> tensors_for_op,
     std::string group_name);
 
 } // namespace c10d
diff --git a/torch/csrc/distributed/c10d/GlooDeviceFactory.cpp b/torch/csrc/distributed/c10d/GlooDeviceFactory.cpp
index 25448dbc9f690..aefb79cfa3932 100644
--- a/torch/csrc/distributed/c10d/GlooDeviceFactory.cpp
+++ b/torch/csrc/distributed/c10d/GlooDeviceFactory.cpp
@@ -195,22 +195,25 @@ std::shared_ptr<::gloo::transport::Device> makeGlooDevice(
         transportName.value(), interfaceName, hostName, lazyInit);
   }
 
-#ifdef __linux__
+#if defined(__linux__)
+
   return GlooDeviceRegistry()->Create(
       "LINUX", interfaceName, hostName, lazyInit);
-#endif
 
-#ifdef __APPLE__
+#elif defined(__APPLE__)
+
   return GlooDeviceRegistry()->Create(
       "APPLE", interfaceName, hostName, lazyInit);
-#endif
 
-#ifdef _WIN32
+#elif defined(_WIN32)
+
   return GlooDeviceRegistry()->Create(
       "WIN32", interfaceName, hostName, lazyInit);
-#endif
+#else
 
   return nullptr;
+
+#endif
 }
 } // anonymous namespace
 
diff --git a/torch/csrc/distributed/c10d/NCCLUtils.cpp b/torch/csrc/distributed/c10d/NCCLUtils.cpp
index e1996ad376b16..182aa1975e891 100644
--- a/torch/csrc/distributed/c10d/NCCLUtils.cpp
+++ b/torch/csrc/distributed/c10d/NCCLUtils.cpp
@@ -595,6 +595,54 @@ std::string NCCLComm::repr() const {
   return c10::str((void*)ncclComm_);
 }
 
+void NCCLComm::suspend() {
+#ifdef NCCL_HAS_COMM_OFFLOAD
+  LockType lock(mutex_);
+  at::cuda::OptionalCUDAGuard gpuGuard(deviceIndex_);
+  auto comm = getNcclComm();
+  C10D_NCCL_CHECK(ncclCommSuspend(comm, NCCL_SUSPEND_MEM), std::nullopt);
+#else
+  TORCH_CHECK(false, "suspend() requires NCCL 2.29.7 or later");
+#endif
+}
+
+void NCCLComm::resume() {
+#ifdef NCCL_HAS_COMM_OFFLOAD
+  LockType lock(mutex_);
+  at::cuda::OptionalCUDAGuard gpuGuard(deviceIndex_);
+  auto comm = getNcclComm();
+  C10D_NCCL_CHECK(ncclCommResume(comm), std::nullopt);
+#else
+  TORCH_CHECK(false, "resume() requires NCCL 2.29.7 or later");
+#endif
+}
+
+std::unordered_map<std::string, uint64_t> NCCLComm::getMemoryStats() {
+#ifdef NCCL_HAS_COMM_OFFLOAD
+  LockType lock(mutex_);
+  at::cuda::OptionalCUDAGuard gpuGuard(deviceIndex_);
+  auto comm = getNcclComm();
+  uint64_t suspend, suspended, persist, total;
+  C10D_NCCL_CHECK(
+      ncclCommMemStats(comm, ncclStatGpuMemSuspend, &suspend), std::nullopt);
+  C10D_NCCL_CHECK(
+      ncclCommMemStats(comm, ncclStatGpuMemSuspended, &suspended),
+      std::nullopt);
+  C10D_NCCL_CHECK(
+      ncclCommMemStats(comm, ncclStatGpuMemPersist, &persist), std::nullopt);
+  C10D_NCCL_CHECK(
+      ncclCommMemStats(comm, ncclStatGpuMemTotal, &total), std::nullopt);
+  return {
+      {"suspend", suspend},
+      {"suspended", suspended},
+      {"persist", persist},
+      {"total", total},
+  };
+#else
+  TORCH_CHECK(false, "getMemoryStats() requires NCCL 2.29.7 or later");
+#endif
+}
+
 #if (defined(IS_NCCLX) || defined(USE_ROCM)) && defined(NCCL_COMM_DUMP)
 std::unordered_map<std::string, std::string> NCCLComm::ncclCommDump() {
   std::unordered_map<std::string, std::string> dump;
diff --git a/torch/csrc/distributed/c10d/NCCLUtils.hpp b/torch/csrc/distributed/c10d/NCCLUtils.hpp
index 4b0f38ce70ff1..f6bd900dc6dc2 100644
--- a/torch/csrc/distributed/c10d/NCCLUtils.hpp
+++ b/torch/csrc/distributed/c10d/NCCLUtils.hpp
@@ -94,6 +94,10 @@ static_assert(
 #define NCCL_HAS_COMM_SHRINK
 #endif
 
+#if NCCL_VERSION_CODE >= NCCL_VERSION(2, 29, 7)
+#define NCCL_HAS_COMM_OFFLOAD
+#endif
+
 // Macro to throw on a non-successful NCCL return value.
 #define C10D_NCCL_CHECK(cmd, failureReason)                                   \
   do {                                                                        \
@@ -376,6 +380,13 @@ class NCCLComm {
 
   std::string repr() const;
 
+  // APIs related to memory offload (require NCCL 2.29.7+ at runtime)
+  void suspend();
+
+  void resume();
+
+  std::unordered_map<std::string, uint64_t> getMemoryStats();
+
   friend class ProcessGroupNCCL;
 
  protected:
diff --git a/torch/csrc/distributed/c10d/NanCheck.cpp b/torch/csrc/distributed/c10d/NanCheck.cpp
new file mode 100644
index 0000000000000..f8392af91370a
--- /dev/null
+++ b/torch/csrc/distributed/c10d/NanCheck.cpp
@@ -0,0 +1,55 @@
+#include <ATen/Dispatch.h>
+#include <ATen/NumericUtils.h>
+#include <ATen/Parallel.h>
+#include <torch/csrc/distributed/c10d/NanCheck.hpp>
+#include <torch/library.h>
+
+namespace c10d {
+
+namespace {
+
+void check_for_nan_cpu(const at::Tensor& tensor) {
+  if (!tensor.is_floating_point()) {
+    return;
+  }
+  if (tensor.numel() == 0) {
+    return;
+  }
+
+  AT_DISPATCH_FLOATING_TYPES_AND4(
+      at::ScalarType::Half,
+      at::ScalarType::BFloat16,
+      at::ScalarType::Float8_e4m3fn,
+      at::ScalarType::Float8_e5m2,
+      tensor.scalar_type(),
+      "check_for_nan_cpu",
+      [&] {
+        auto* data = tensor.data_ptr<scalar_t>();
+        auto numel = tensor.numel();
+        std::atomic<bool> found{false};
+        at::parallel_for(0, numel, 1024, [&](int64_t begin, int64_t end) {
+          for (int64_t i = begin; i < end; i++) {
+            if (at::_isnan(data[i])) {
+              found.store(true, std::memory_order_relaxed);
+              return;
+            }
+          }
+        });
+        TORCH_CHECK(!found.load(), "NaN found in input tensor.");
+      });
+}
+
+TORCH_LIBRARY_IMPL(c10d, CPU, m) {
+  m.impl("check_for_nan", check_for_nan_cpu);
+}
+
+} // namespace
+
+void checkForNan(const at::Tensor& tensor) {
+  static auto op = c10::Dispatcher::singleton()
+                       .findSchemaOrThrow("c10d::check_for_nan", "")
+                       .typed<void(const at::Tensor&)>();
+  op.call(tensor);
+}
+
+} // namespace c10d
diff --git a/torch/csrc/distributed/c10d/NanCheck.cu b/torch/csrc/distributed/c10d/NanCheck.cu
index 2506ccd1ad094..d39bfbbf0d7ea 100644
--- a/torch/csrc/distributed/c10d/NanCheck.cu
+++ b/torch/csrc/distributed/c10d/NanCheck.cu
@@ -1,14 +1,12 @@
-#ifdef USE_C10D_NCCL
-
 #include <ATen/Dispatch.h>
 #include <ATen/cuda/CUDAContext.h>
 #include <c10/cuda/CUDAGuard.h>
-#include <torch/torch.h>
-#include <algorithm>
-#include <torch/csrc/distributed/c10d/NanCheck.hpp>
 #include <stdint.h>
+#include <torch/library.h>
+#include <algorithm>
 
 namespace c10d {
+namespace {
 
 // CUDA kernel to check if data has NAN, device side assert
 // is raised if NAN is found
@@ -21,11 +19,6 @@ union BytePack16 {
 
 typedef union BytePack16 BytePack;
 
-// AMD HIP doesn't define `__trap()`, using `assert` instead
-#ifdef USE_ROCM
-#define __trap() assert(0)
-#endif
-
 //// Start of templated functions for checking NaNs inside a BytePack
 
 // (i) General implementation (aka fallback)
@@ -36,9 +29,9 @@ template <typename T, int EltPerPack>
 struct CheckBytePack {
   static __device__ __forceinline__ void check(BytePack* tmp) {
     T* data = (T*)tmp;
-    #pragma unroll 8
+#pragma unroll 8
     for (int i = 0; i < EltPerPack; i++) {
-      if (isnan(data[i])) __trap();
+      CUDA_KERNEL_ASSERT(!isnan(data[i]));
     }
   }
 };
@@ -47,10 +40,10 @@ struct CheckBytePack {
 // EltPerPack = 16 / 8 = 2
 
 template <typename T>
-struct CheckBytePack<T, /*EltPerPack*/2> {
+struct CheckBytePack<T, /*EltPerPack*/ 2> {
   static __device__ __forceinline__ void check(BytePack* tmp) {
     T* data = (T*)tmp;
-    if (isnan(data[0]) || isnan(data[1])) __trap();
+    CUDA_KERNEL_ASSERT(!isnan(data[0]) && !isnan(data[1]));
   }
 };
 
@@ -58,24 +51,26 @@ struct CheckBytePack<T, /*EltPerPack*/2> {
 // EltPerPack = 16 / 4 = 4
 
 template <typename T>
-struct CheckBytePack<T, /*EltPerPack*/4> {
+struct CheckBytePack<T, /*EltPerPack*/ 4> {
   static __device__ __forceinline__ void check(BytePack* tmp) {
     T* data = (T*)tmp;
-    if (isnan(data[0]) || isnan(data[1]) || isnan(data[2]) || isnan(data[3])) __trap();
+    CUDA_KERNEL_ASSERT(
+        !isnan(data[0]) && !isnan(data[1]) && !isnan(data[2]) &&
+        !isnan(data[3]));
   }
 };
 
-// (iv) Template specialization for 2-byte data types, e.g. float16, bfloat16, half.
-// EltPerPack = 16 / 2 = 8
+// (iv) Template specialization for 2-byte data types, e.g. float16, bfloat16,
+// half. EltPerPack = 16 / 2 = 8
 
 template <typename T>
-struct CheckBytePack<T, /*EltPerPack*/8> {
+struct CheckBytePack<T, /*EltPerPack*/ 8> {
   static __device__ __forceinline__ void check(BytePack* tmp) {
     T* data = (T*)tmp;
-    if (isnan(data[0]) || isnan(data[1]) || isnan(data[2]) || isnan(data[3]) ||
-        isnan(data[4]) || isnan(data[5]) || isnan(data[6]) || isnan(data[7])) {
-          __trap();
-    }
+    CUDA_KERNEL_ASSERT(
+        !isnan(data[0]) && !isnan(data[1]) && !isnan(data[2]) &&
+        !isnan(data[3]) && !isnan(data[4]) && !isnan(data[5]) &&
+        !isnan(data[6]) && !isnan(data[7]));
   }
 };
 
@@ -83,7 +78,7 @@ struct CheckBytePack<T, /*EltPerPack*/8> {
 // EltPerPack = 16 / 1 = 16
 
 // We want to check 8 x FP8 simultaneously, hence this template definition.
-template<typename T>
+template <typename T>
 struct HasNanFP8x8 {
   static __device__ __forceinline__ bool check(uint64_t fp8x8) = delete;
   /*
@@ -93,8 +88,8 @@ struct HasNanFP8x8 {
     static_assert(
       false,
       "You should never call this template definition because it is empty. You "
-      "can follow the example of Float8_e4m3fn below to implement the check for "
-      "your new datatype."
+      "can follow the example of Float8_e4m3fn below to implement the check for
+  " "your new datatype."
     );
   }
   */
@@ -112,7 +107,7 @@ struct HasNanFP8x8 {
 // (3) Only leave the 8th bit by masking with 0x80.
 // (4) If any x[i] is nan, then the whole x != 0.
 
-template<>
+template <>
 struct HasNanFP8x8<c10::Float8_e4m3fn> {
   static __device__ __forceinline__ bool check(uint64_t fp8x8) {
     auto t = fp8x8 & 0x7F7F7F7F7F7F7F7FULL;
@@ -127,7 +122,7 @@ struct HasNanFP8x8<c10::Float8_e4m3fn> {
 // This case does not overflow: 0x7c + 0x03 == 0x7f but adding 0x03 to anything
 // greater than 0x7c will overflow.
 
-template<>
+template <>
 struct HasNanFP8x8<c10::Float8_e5m2> {
   static __device__ __forceinline__ bool check(uint64_t fp8x8) {
     auto t = fp8x8 & 0x7F7F7F7F7F7F7F7FULL;
@@ -137,17 +132,17 @@ struct HasNanFP8x8<c10::Float8_e5m2> {
   }
 };
 
-template<typename T>
-struct CheckBytePack<T, /*EltPerPack*/16> {
+template <typename T>
+struct CheckBytePack<T, /*EltPerPack*/ 16> {
   static __device__ __forceinline__ void check(BytePack* tmp) {
-    if (HasNanFP8x8<T>::check(tmp->ul[0]) || HasNanFP8x8<T>::check(tmp->ul[1]))
-        __trap();
+    CUDA_KERNEL_ASSERT(
+        !HasNanFP8x8<T>::check(tmp->ul[0]) &&
+        !HasNanFP8x8<T>::check(tmp->ul[1]));
   }
 };
 
 //// End of templated functions for checking NaNs inside a BytePack
 
-
 // Fast-path check routine:
 // each thread will load and check 8 BytePacks in this routine
 
@@ -158,21 +153,22 @@ template <typename T>
 __device__ __forceinline__ void checkChunk(BytePack* ptr) {
   BytePack tmp[UNROLL];
   int nWorkers = blockDim.x * gridDim.x;
-  // First load values from global memory into tmp buffer
-  #pragma unroll 8
+// First load values from global memory into tmp buffer
+#pragma unroll 8
   for (int j = 0; j < UNROLL; j++) {
     tmp[j] = ptr[nWorkers * j];
   }
-  // Then check each BytePack in the tmp buffer
-  #pragma unroll 8
+// Then check each BytePack in the tmp buffer
+#pragma unroll 8
   for (int j = 0; j < UNROLL; j++) {
-    CheckBytePack<T, sizeof(BytePack)/sizeof(T)>::check(tmp + j);
+    CheckBytePack<T, sizeof(BytePack) / sizeof(T)>::check(tmp + j);
   }
   // Note: we separate the check from the load for efficient loading
 }
 
 // Align address of `ptr` up, to the alignment of `T`
-#define ALIGN_UP(ptr, T) (((uintptr_t)ptr + sizeof(T) - 1) / sizeof(T) * sizeof(T))
+#define ALIGN_UP(ptr, T) \
+  (((uintptr_t)ptr + sizeof(T) - 1) / sizeof(T) * sizeof(T))
 
 // This is the host-facing kernel
 
@@ -189,7 +185,7 @@ __global__ void checkForNaN(T* data, size_t size) {
   // Read memory by T (slow). One iter is enough bc the number of threads would
   // be bigger than `preProcElts`
   if (offset < preProcElts) {
-    if (isnan(data[offset])) __trap();
+    CUDA_KERNEL_ASSERT(!isnan(data[offset]));
   }
   // We have processes this amount of data
   size -= preProcElts;
@@ -218,14 +214,12 @@ __global__ void checkForNaN(T* data, size_t size) {
   // TODO: merge this tail check with head check to make them concurrent
   if (threadIdx.x < size % EltPerPack) {
     T* tailPtr = (T*)(ptr + sizeInBP);
-    if (isnan(tailPtr[threadIdx.x])) __trap();
+    CUDA_KERNEL_ASSERT(!isnan(tailPtr[threadIdx.x]));
   }
 }
 
-// CHECK if a Tensor contains NAN in any of its element
-void checkForNan(const at::Tensor& tensor, at::cuda::CUDAStream& stream) {
-  // skip check for non float types
-  if (!torch::is_floating_point(tensor)) {
+void check_for_nan_cuda(const at::Tensor& tensor) {
+  if (!tensor.is_floating_point()) {
     return;
   }
   const size_t maxNumThreadsPerBlock = 512;
@@ -241,6 +235,8 @@ void checkForNan(const at::Tensor& tensor, at::cuda::CUDAStream& stream) {
       maxNumBlocks,
       (tensor.numel() + numThreadsPerBlock - 1) / numThreadsPerBlock);
 
+  auto stream = at::cuda::getCurrentCUDAStream(tensor.device().index());
+
   AT_DISPATCH_FLOATING_TYPES_AND4(
       at::ScalarType::Half,
       at::ScalarType::BFloat16,
@@ -255,6 +251,9 @@ void checkForNan(const at::Tensor& tensor, at::cuda::CUDAStream& stream) {
       });
 }
 
-} // namespace c10d
+TORCH_LIBRARY_IMPL(c10d, CUDA, m) {
+  m.impl("check_for_nan", check_for_nan_cuda);
+}
 
-#endif // USE_C10D_NCCL
+} // namespace
+} // namespace c10d
diff --git a/torch/csrc/distributed/c10d/NanCheck.hpp b/torch/csrc/distributed/c10d/NanCheck.hpp
index cc9a5867c3dd4..57b0b211a24c6 100644
--- a/torch/csrc/distributed/c10d/NanCheck.hpp
+++ b/torch/csrc/distributed/c10d/NanCheck.hpp
@@ -1,16 +1,12 @@
 #pragma once
 
-#ifdef USE_C10D_NCCL
-
 #include <ATen/ATen.h>
-#include <c10/cuda/CUDAStream.h>
+#include <torch/csrc/Export.h>
 
 namespace c10d {
 
-// Check for NaNs in a tensor on a given stream. If any are found, throw a
-// device-side error.
-void checkForNan(const at::Tensor& tensor, at::cuda::CUDAStream& stream);
+// Check for NaNs in a tensor. If any are found, throw an error.
+// Dispatches to device-specific implementations via the c10d::check_for_nan op.
+TORCH_API void checkForNan(const at::Tensor& tensor);
 
 } // namespace c10d
-
-#endif // USE_C10D_NCCL
diff --git a/torch/csrc/distributed/c10d/Ops.cpp b/torch/csrc/distributed/c10d/Ops.cpp
index 0ded9d4cc733d..cd191225e78d4 100644
--- a/torch/csrc/distributed/c10d/Ops.cpp
+++ b/torch/csrc/distributed/c10d/Ops.cpp
@@ -59,6 +59,7 @@ TORCH_LIBRARY(c10d, m) {
       "recv_(Tensor[] tensors, __torch__.torch.classes.c10d.ProcessGroup process_group, int src, int tag) -> __torch__.torch.classes.c10d.Work");
   m.def(
       "recv_any_source_(Tensor[] tensors, __torch__.torch.classes.c10d.ProcessGroup process_group, int tag) -> __torch__.torch.classes.c10d.Work");
+  m.def("check_for_nan(Tensor tensor) -> ()");
 }
 } // namespace
 
diff --git a/torch/csrc/distributed/c10d/ParamCommsUtils.cpp b/torch/csrc/distributed/c10d/ParamCommsUtils.cpp
index ca505d1f1d266..9f126564ae0b7 100644
--- a/torch/csrc/distributed/c10d/ParamCommsUtils.cpp
+++ b/torch/csrc/distributed/c10d/ParamCommsUtils.cpp
@@ -18,7 +18,8 @@ ParamCommsDebugInfo::ParamCommsDebugInfo(
     std::vector<int64_t> outSplitSizes,
     int globalRankStart,
     int globalRankStride,
-    int worldSize)
+    int worldSize,
+    bool isAsynchronizedOp)
     : pgName_(std::move(pgName)),
       rank_(rank),
       worldSize_(worldSize),
@@ -29,7 +30,8 @@ ParamCommsDebugInfo::ParamCommsDebugInfo(
       inputSplitSizes_(std::move(inSplitSizes)),
       outputSplitSizes_(std::move(outSplitSizes)),
       globalRankStart_(globalRankStart),
-      globalRankStride_(globalRankStride) {
+      globalRankStride_(globalRankStride),
+      isAsynchronizedOp_(isAsynchronizedOp) {
   if (globalRankStride > 0) {
     for (int i = 0; i < worldSize; i++) {
       groupRanks_.push_back(globalRankStart + i * globalRankStride);
diff --git a/torch/csrc/distributed/c10d/ParamCommsUtils.hpp b/torch/csrc/distributed/c10d/ParamCommsUtils.hpp
index 678c98e91a0b9..ad38fcba68ded 100644
--- a/torch/csrc/distributed/c10d/ParamCommsUtils.hpp
+++ b/torch/csrc/distributed/c10d/ParamCommsUtils.hpp
@@ -5,6 +5,7 @@
 #include <c10/macros/Macros.h>
 #include <c10/util/ThreadLocalDebugInfo.h>
 #include <string>
+#include <tuple>
 #include <vector>
 
 namespace torch {
@@ -23,7 +24,8 @@ class TORCH_API ParamCommsDebugInfo : public c10::DebugInfoBase {
       std::vector<int64_t> outSplitSizes,
       int globalRankStart,
       int globalRankStride,
-      int worldSize);
+      int worldSize,
+      bool isAsynchronizedOp = true);
 
   ~ParamCommsDebugInfo() override = default;
 
@@ -79,6 +81,23 @@ class TORCH_API ParamCommsDebugInfo : public c10::DebugInfoBase {
     return groupRanks_;
   }
 
+  bool isAsynchronizedOp() const {
+    return isAsynchronizedOp_;
+  }
+
+  int64_t getSequenceNumber() const {
+    return sequenceNumber_;
+  }
+
+  bool getIsP2P() const {
+    return isP2P_;
+  }
+
+  void setSequenceInfo(int64_t seqNum, bool isP2P) {
+    sequenceNumber_ = seqNum;
+    isP2P_ = isP2P;
+  }
+
  private:
   std::tuple<std::string, std::string> pgName_; // <group_name, group_desc>
   int rank_{};
@@ -92,8 +111,25 @@ class TORCH_API ParamCommsDebugInfo : public c10::DebugInfoBase {
   int globalRankStart_{};
   int globalRankStride_{};
   std::vector<int64_t> groupRanks_;
+  bool isAsynchronizedOp_{};
+  int64_t sequenceNumber_{-1};
+  bool isP2P_{false};
 };
 
+// Helper to set sequence info from tuple-typed seq arguments (NCCL backend).
+// No-op fallback for backends that pass non-tuple seq types (e.g., XPU/XCCL).
+template <typename A, typename B>
+inline void maybeSetSequenceInfo(
+    const std::shared_ptr<ParamCommsDebugInfo>& info,
+    const std::tuple<A, B>& seq) {
+  info->setSequenceInfo(std::get<0>(seq), std::get<1>(seq));
+}
+
+template <typename T>
+inline void maybeSetSequenceInfo(
+    const std::shared_ptr<ParamCommsDebugInfo>&,
+    const T&) {}
+
 #define RECORD_PARAM_COMMS(                                                    \
     seq,                                                                       \
     pgName,                                                                    \
@@ -118,7 +154,9 @@ class TORCH_API ParamCommsDebugInfo : public c10::DebugInfoBase {
       outSplitSizes,                                                           \
       globalRankStart,                                                         \
       globalRankStride,                                                        \
-      worldSize);                                                              \
+      worldSize,                                                               \
+      false);                                                                  \
+  torch::maybeSetSequenceInfo(paramCommsInfo, seq);                            \
   c10::DebugInfoGuard g(c10::DebugInfoKind::PARAM_COMMS_INFO, paramCommsInfo); \
   std::initializer_list<const c10::IValue> paramList = {                       \
       seq,                                                                     \
@@ -129,11 +167,44 @@ class TORCH_API ParamCommsDebugInfo : public c10::DebugInfoBase {
       outSplitSizes,                                                           \
       globalRankStart,                                                         \
       globalRankStride,                                                        \
-      worldSize};                                                              \
+      worldSize,                                                               \
+      false};                                                                  \
   c10::ArrayRef<const c10::IValue> paramInputs(paramList);                     \
   RECORD_FUNCTION(at::kParamCommsCallName, paramInputs);
 
-#define RECORD_PARAM_COMMS_DATA(                                               \
+#define RECORD_PARAM_COMMS_DATA(         \
+    seq,                                 \
+    pgName,                              \
+    InputTensors,                        \
+    OutputTensors,                       \
+    rank,                                \
+    collName,                            \
+    inNelems,                            \
+    outNelems,                           \
+    dType,                               \
+    inSplitSizes,                        \
+    outSplitSizes,                       \
+    globalRankStart,                     \
+    globalRankStride,                    \
+    worldSize)                           \
+  RECORD_PARAM_COMMS_DATA_WITH_ASYNC_OP( \
+      seq,                               \
+      pgName,                            \
+      InputTensors,                      \
+      OutputTensors,                     \
+      rank,                              \
+      collName,                          \
+      inNelems,                          \
+      outNelems,                         \
+      dType,                             \
+      inSplitSizes,                      \
+      outSplitSizes,                     \
+      globalRankStart,                   \
+      globalRankStride,                  \
+      worldSize,                         \
+      true);
+
+#define RECORD_PARAM_COMMS_DATA_WITH_ASYNC_OP(                                 \
     seq,                                                                       \
     pgName,                                                                    \
     InputTensors,                                                              \
@@ -147,7 +218,8 @@ class TORCH_API ParamCommsDebugInfo : public c10::DebugInfoBase {
     outSplitSizes,                                                             \
     globalRankStart,                                                           \
     globalRankStride,                                                          \
-    worldSize)                                                                 \
+    worldSize,                                                                 \
+    isAsyncOp)                                                                 \
   auto paramCommsInfo = std::make_shared<torch::ParamCommsDebugInfo>(          \
       pgName,                                                                  \
       rank,                                                                    \
@@ -159,7 +231,9 @@ class TORCH_API ParamCommsDebugInfo : public c10::DebugInfoBase {
       outSplitSizes,                                                           \
       globalRankStart,                                                         \
       globalRankStride,                                                        \
-      worldSize);                                                              \
+      worldSize,                                                               \
+      isAsyncOp);                                                              \
+  torch::maybeSetSequenceInfo(paramCommsInfo, seq);                            \
   c10::DebugInfoGuard g(c10::DebugInfoKind::PARAM_COMMS_INFO, paramCommsInfo); \
   std::initializer_list<const c10::IValue> paramList = {                       \
       c10::IValue(InputTensors),                                               \
@@ -171,7 +245,8 @@ class TORCH_API ParamCommsDebugInfo : public c10::DebugInfoBase {
       outSplitSizes,                                                           \
       globalRankStart,                                                         \
       globalRankStride,                                                        \
-      worldSize};                                                              \
+      worldSize,                                                               \
+      isAsyncOp};                                                              \
   c10::ArrayRef<const c10::IValue> paramInputs(paramList);                     \
   RECORD_FUNCTION_WITH_INPUTS_OUTPUTS(                                         \
       at::kParamCommsCallName,                                                 \
diff --git a/torch/csrc/distributed/c10d/ProcessGroup.cpp b/torch/csrc/distributed/c10d/ProcessGroup.cpp
index 14d4be3a13f71..bdd7a0d442706 100644
--- a/torch/csrc/distributed/c10d/ProcessGroup.cpp
+++ b/torch/csrc/distributed/c10d/ProcessGroup.cpp
@@ -51,10 +51,11 @@ std::string opTypeToString(OpType opType) {
       return "COALESCED";
     case OpType::_ALLREDUCE_SPARSE:
       return "_ALLREDUCE_SPARSE";
+    case OpType::REDUCE_SCATTER_TENSOR_COALESCED:
+      return "REDUCE_SCATTER_TENSOR_COALESCED";
     default:
       TORCH_INTERNAL_ASSERT(false, "Unknown op type!");
   }
-  return "UNKNOWN";
 }
 
 bool isP2POp(OpType opType, bool batchP2P /*= false*/) {
diff --git a/torch/csrc/distributed/c10d/ProcessGroup.hpp b/torch/csrc/distributed/c10d/ProcessGroup.hpp
index 8c4a657fd7eed..95204e583caed 100644
--- a/torch/csrc/distributed/c10d/ProcessGroup.hpp
+++ b/torch/csrc/distributed/c10d/ProcessGroup.hpp
@@ -793,7 +793,7 @@ class TORCH_API ProcessGroup : public torch::CustomClassHolder {
 
   virtual c10::intrusive_ptr<Work> barrier(
       const BarrierOptions& opts = BarrierOptions()) {
-    static at::Tensor tensor;
+    at::Tensor tensor;
     // TODO: if nccl was specified then use it
     auto device = opts.device;
     if (device.has_value()) {
diff --git a/torch/csrc/distributed/c10d/ProcessGroupGloo.cpp b/torch/csrc/distributed/c10d/ProcessGroupGloo.cpp
index 9f60a1f0cbf26..c81516b10bad7 100644
--- a/torch/csrc/distributed/c10d/ProcessGroupGloo.cpp
+++ b/torch/csrc/distributed/c10d/ProcessGroupGloo.cpp
@@ -2468,6 +2468,242 @@ c10::intrusive_ptr<Work> ProcessGroupGloo::alltoall_base(
   return work;
 }
 
+namespace {
+
+class AsyncAlltoallListWork : public ProcessGroupGloo::AsyncWork {
+ public:
+  AsyncAlltoallListWork(
+      std::shared_ptr<gloo::Context> context,
+      std::vector<at::Tensor>& outputTensors,
+      std::vector<at::Tensor>& inputTensors,
+      uint32_t tag,
+      uint64_t seq,
+      std::chrono::milliseconds timeout)
+      : ProcessGroupGloo::AsyncWork(
+            std::move(context),
+            {outputTensors},
+            OpType::ALLTOALL,
+            seq,
+            timeout,
+            "gloo:all_to_all",
+            inputTensors),
+        outputTensors(outputTensors),
+        inputTensors(inputTensors),
+        tag(tag) {}
+
+  std::vector<at::Tensor> outputTensors;
+  std::vector<at::Tensor> inputTensors;
+  const uint32_t tag;
+
+  void alltoall(
+      std::vector<at::Tensor>& outputTensors,
+      std::vector<at::Tensor>& inputTensors) {
+    const auto scalarType = inputTensors[0].scalar_type();
+    gloo::AlltoallOptions opts(context_);
+    opts.setTag(tag);
+    opts.setTimeout(getTimeout());
+
+    // Flatten input tensors into a single buffer
+    at::Tensor flatInputTensor = flattenDenseTensors(inputTensors);
+    GENERATE_ALL_TYPES(scalarType, setInput, opts, flatInputTensor);
+
+    // Allocate flat output tensor with same total size
+    at::Tensor flatOutputTensor = newLikeFlat(outputTensors);
+    GENERATE_ALL_TYPES(scalarType, setOutput, opts, flatOutputTensor);
+
+    // Perform the all-to-all operation
+    gloo::alltoall(opts);
+
+    // Unflatten output into individual tensors
+    for (const auto i : c10::irange(outputTensors.size())) {
+      outputTensors[i].copy_(flatOutputTensor[static_cast<int64_t>(i)]);
+    }
+  }
+
+  const std::vector<at::Tensor> getInputTensors() override {
+    return inputTensors;
+  }
+
+  const std::vector<at::Tensor> getOutputTensors() override {
+    return outputTensors;
+  }
+
+  void run() override {
+    alltoall(outputTensors, inputTensors);
+  }
+};
+
+class AsyncAlltoallListCUDAWork : public AsyncAlltoallListWork {
+ public:
+  AsyncAlltoallListCUDAWork(
+      const std::shared_ptr<gloo::Context>& context,
+      std::vector<at::Tensor>& outputTensors,
+      std::vector<at::Tensor>& inputTensors,
+      uint32_t tag,
+      uint64_t seq,
+      std::chrono::milliseconds timeout)
+      : AsyncAlltoallListWork(
+            context,
+            outputTensors,
+            inputTensors,
+            tag,
+            seq,
+            timeout) {
+    initializeStreamsEvents(inputTensors, inputStreams, inputEvents);
+    initializeStreamsEvents(outputTensors, outputStreams, outputEvents);
+
+    // Kick off copy from CUDA tensors to pinned CPU tensors.
+    tmpInputs.reserve(inputTensors.size());
+    c10::OptionalStreamGuard guard;
+    for (const auto i : c10::irange(inputTensors.size())) {
+      guard.reset_stream(inputStreams[i]);
+      tmpInputs.push_back(
+          pinnedLike(inputTensors[i]).copy_(inputTensors[i], true));
+    }
+
+    tmpOutputs.reserve(outputTensors.size());
+    for (const auto i : c10::irange(outputTensors.size())) {
+      guard.reset_stream(outputStreams[i]);
+      tmpOutputs.push_back(pinnedLike(outputTensors[i]));
+    }
+  }
+
+  void run() override {
+    // Synchronize with copy operations.
+    for (const auto i : c10::irange(inputTensors.size())) {
+      inputStreams[i].synchronize();
+    }
+    for (const auto i : c10::irange(outputTensors.size())) {
+      outputStreams[i].synchronize();
+    }
+
+    // Run alltoall on host side tensors.
+    alltoall(tmpOutputs, tmpInputs);
+
+    // Kick off copy back to the CUDA tensors.
+    c10::OptionalStreamGuard guard;
+    for (const auto i : c10::irange(outputTensors.size())) {
+      guard.reset_stream(outputStreams[i]);
+      outputTensors[i].copy_(tmpOutputs[i], /* non_blocking */ true);
+      outputEvents[i].record(outputStreams[i]);
+    }
+  }
+
+  void synchronize() override {
+    // Synchronize with the copy back to CUDA tensors.
+    for (const auto i : c10::irange(outputTensors.size())) {
+      c10::Device device = outputTensors[i].device();
+      outputEvents[i].block(
+          c10::impl::VirtualGuardImpl(device.type()).getStream(device));
+    }
+  }
+
+  std::vector<at::Tensor> tmpInputs;
+  std::vector<c10::Stream> inputStreams;
+  std::vector<c10::Event> inputEvents;
+
+  std::vector<at::Tensor> tmpOutputs;
+  std::vector<c10::Stream> outputStreams;
+  std::vector<c10::Event> outputEvents;
+};
+
+} // namespace
+
+c10::intrusive_ptr<Work> ProcessGroupGloo::alltoall(
+    std::vector<at::Tensor>& outputTensors,
+    std::vector<at::Tensor>& inputTensors,
+    const AllToAllOptions& opts) {
+  static auto invalidArgument = [](const std::string& msg) {
+    TORCH_CHECK(false, "ProcessGroupGloo::alltoall: " + msg);
+  };
+
+  // Validate input and output tensor lists
+  if (inputTensors.size() != static_cast<size_t>(getSize())) {
+    invalidArgument(
+        "input tensor list size " + std::to_string(inputTensors.size()) +
+        " does not match world size " + std::to_string(getSize()));
+  }
+
+  if (outputTensors.size() != static_cast<size_t>(getSize())) {
+    invalidArgument(
+        "output tensor list size " + std::to_string(outputTensors.size()) +
+        " does not match world size " + std::to_string(getSize()));
+  }
+
+  assertDense(invalidArgument, inputTensors);
+  assertDense(invalidArgument, outputTensors);
+
+  // Check that all tensors are on the same device
+  assertSameDevice(invalidArgument, inputTensors);
+  assertSameDevice(invalidArgument, outputTensors);
+
+  // Check that all input tensors have the same type and size
+  const auto& options = inputTensors[0].options();
+  const auto& sizes = inputTensors[0].sizes();
+  assertTypeAndSizesMatch(invalidArgument, inputTensors, options, sizes);
+
+  // Check that all output tensors have the same type and size
+  const auto& outputOptions = outputTensors[0].options();
+  const auto& outputSizes = outputTensors[0].sizes();
+  assertTypeAndSizesMatch(
+      invalidArgument, outputTensors, outputOptions, outputSizes);
+
+  // Check input and output tensors have compatible types
+  if (!options.type_equal(outputOptions)) {
+    invalidArgument("input and output tensors must have the same type");
+  }
+
+  // Check input and output tensors have compatible sizes
+  if (sizes != outputSizes) {
+    invalidArgument("input and output tensors must have the same size");
+  }
+
+  // Check device type
+  const auto& device = inputTensors[0].device();
+  TORCH_CHECK(
+      outputTensors[0].device() == device,
+      "input and output tensors must be on the same device");
+
+  switch (device.type()) {
+    case at::kCPU:
+      break;
+    case at::kCUDA:
+      // If the user gave us a CUDA tensor then CUDA must be loaded.
+      TORCH_INTERNAL_ASSERT(at::hasCUDA());
+      break;
+    default:
+      invalidArgument(c10::str("unsupported device type ", device.type()));
+  }
+
+  c10::intrusive_ptr<AsyncAlltoallListWork> work;
+  auto tag = nextTag();
+  auto context = getContext(tag);
+  ++seq_;
+
+  if (device.type() == at::kCPU) {
+    work = c10::make_intrusive<AsyncAlltoallListWork>(
+        std::move(context),
+        outputTensors,
+        inputTensors,
+        tag,
+        seq_,
+        opts.timeout);
+  } else if (device.type() == at::kCUDA) {
+    work = c10::make_intrusive<AsyncAlltoallListCUDAWork>(
+        std::move(context),
+        outputTensors,
+        inputTensors,
+        tag,
+        seq_,
+        opts.timeout);
+  } else {
+    TORCH_CHECK(false, "Invalid backend");
+  }
+
+  enqueue(work);
+  return work;
+}
+
 static at::Tensor& checkSingleTensor(std::vector<at::Tensor>& tensors) {
   if (tensors.size() != 1) {
     TORCH_CHECK(false, "ProcessGroupGloo::send takes a single tensor");
diff --git a/torch/csrc/distributed/c10d/ProcessGroupGloo.hpp b/torch/csrc/distributed/c10d/ProcessGroupGloo.hpp
index f0b37ffa5ca40..7dba096df1bbd 100644
--- a/torch/csrc/distributed/c10d/ProcessGroupGloo.hpp
+++ b/torch/csrc/distributed/c10d/ProcessGroupGloo.hpp
@@ -401,6 +401,11 @@ class TORCH_API ProcessGroupGloo : public Backend {
       std::vector<int64_t>& inputCounts,
       const AllToAllOptions& opts = AllToAllOptions()) override;
 
+  c10::intrusive_ptr<Work> alltoall(
+      std::vector<at::Tensor>& outputTensors,
+      std::vector<at::Tensor>& inputTensors,
+      const AllToAllOptions& opts = AllToAllOptions()) override;
+
   c10::intrusive_ptr<Work> send(
       std::vector<at::Tensor>& tensors,
       int dstRank,
diff --git a/torch/csrc/distributed/c10d/ProcessGroupMPI.cpp b/torch/csrc/distributed/c10d/ProcessGroupMPI.cpp
index f7a5f181f99f5..c51066f2a0b0a 100644
--- a/torch/csrc/distributed/c10d/ProcessGroupMPI.cpp
+++ b/torch/csrc/distributed/c10d/ProcessGroupMPI.cpp
@@ -67,7 +67,15 @@ bool cudaAwareMpiCheck() {
   } else {
     return false;
   }
-#else // !defined(MPIX_CUDA_AWARE_SUPPORT)
+// Recognize that Cray MPICH is CUDA-aware (used on Cray/HPE supercomputers)
+#elif defined(MPIX_GPU_SUPPORT_CUDA)
+  const char* cray_gpu_support = std::getenv("MPICH_GPU_SUPPORT_ENABLED");
+  if (cray_gpu_support != nullptr && std::string(cray_gpu_support) == "1") {
+    return true;
+  } else {
+    return false;
+  }
+#else // !defined(MPIX_CUDA_AWARE_SUPPORT) && !defined(MPIX_GPU_SUPPORT_CUDA)
   return false;
 #endif // MPIX_CUDA_AWARE_SUPPORT
 }
diff --git a/torch/csrc/distributed/c10d/ProcessGroupNCCL.cpp b/torch/csrc/distributed/c10d/ProcessGroupNCCL.cpp
index 827a8c1b13db7..06e9546de50fa 100644
--- a/torch/csrc/distributed/c10d/ProcessGroupNCCL.cpp
+++ b/torch/csrc/distributed/c10d/ProcessGroupNCCL.cpp
@@ -11,8 +11,10 @@
 #include <utility>
 
 #include <ATen/cuda/CUDAContext.h>
+#include <ATen/cuda/CUDAGraph.h>
 #include <c10/core/DeviceType.h>
 #include <c10/cuda/CUDAAllocatorConfig.h>
+#include <c10/cuda/CUDAException.h>
 #include <c10/cuda/CUDAGraphsC10Utils.h>
 #include <c10/cuda/CUDAGuard.h>
 #include <c10/util/Exception.h>
@@ -131,7 +133,6 @@ ncclRedOpRAII getNcclReduceOp(
           C10_THROW_ERROR(
               TypeError,
               "PreMulSum Data type must be half, float, bfloat16 or double");
-          return ncclRedOp_t{};
       }
 #else
       C10_THROW_ERROR(ValueError, "PreMulSum requires NCCL>=2.11.1");
@@ -221,6 +222,65 @@ std::string getExceptionMsgFromExceptionPtr(
   }
 }
 
+#ifdef USE_ROCM
+// Indicates that we're in the watchdog's event-query phase. This allows ROCm
+// workaround behavior to be applied only to watchdog-side queries, while
+// preserving existing behavior for user/main-thread `WorkNCCL::isCompleted()`
+// and `wait()` calls.
+thread_local bool g_in_rocm_watchdog_event_query_context = false;
+
+struct RocmWatchdogEventQueryContextGuard {
+  RocmWatchdogEventQueryContextGuard()
+      : previous_(g_in_rocm_watchdog_event_query_context) {
+    g_in_rocm_watchdog_event_query_context = true;
+  }
+  ~RocmWatchdogEventQueryContextGuard() {
+    g_in_rocm_watchdog_event_query_context = previous_;
+  }
+
+ private:
+  bool previous_;
+};
+#endif // USE_ROCM
+
+#ifdef USE_ROCM
+// Watchdog-side cudaEventQuery workaround for HIP runtimes without the
+// capture-mode fix.
+// TODO: Remove once all supported runtimes include
+// https://github.com/ROCm/rocm-systems/pull/3176
+bool queryEventWithRocmWatchdogCaptureWorkaround(
+    const std::shared_ptr<at::cuda::CUDAEvent>& event) {
+  if (!event->isCreated()) {
+    return true;
+  }
+
+  // Must unconditionally return false here during watchdog + active capture:
+  // on affected HIP runtimes, even calling cudaEventQuery from the watchdog
+  // thread while another thread has GLOBAL capture active can invalidate that
+  // capture and cause downstream failures. Skip the query entirely and report
+  // "not complete yet"; the watchdog will re-poll once capture ends. Timeout
+  // enforcement is also deferred during this window (see the
+  // is_graph_capture_active() gate in the watchdog loop).
+  if (g_in_rocm_watchdog_event_query_context &&
+      at::cuda::is_graph_capture_active()) {
+    return false;
+  }
+
+  const cudaError_t err =
+      C10_CUDA_ERROR_HANDLED(cudaEventQuery(event->event()));
+  if (err == cudaSuccess) {
+    return true;
+  } else if (err != cudaErrorNotReady) {
+    C10_CUDA_CHECK(err);
+  } else {
+    // ignore and clear the error if not ready
+    (void)cudaGetLastError();
+  }
+
+  return false;
+}
+#endif // USE_ROCM
+
 inline void errorIfCapturingNonCapturableNCCL(c10::cuda::CaptureStatus status) {
   // parentheses avoid some compiler warnings
   static const uint64_t min_version =
@@ -644,7 +704,11 @@ bool ProcessGroupNCCL::WorkNCCL::startedGPUExecutionInternal() const {
     return false;
   }
   // Checking the work's corresponding CUDA event's status
+#ifdef USE_ROCM
+  if (!queryEventWithRocmWatchdogCaptureWorkaround(ncclStartEvent_)) {
+#else
   if (!ncclStartEvent_->query()) {
+#endif
     return false;
   }
   return true;
@@ -657,7 +721,11 @@ bool ProcessGroupNCCL::WorkNCCL::finishedGPUExecutionInternal() const {
   // hang if another thread is holding the CUDA global context lock. For
   // example, when doing a `cudaDeviceSynchronize` or even
   // `cudaStreamSynchronize`.
+#ifdef USE_ROCM
+  if (!queryEventWithRocmWatchdogCaptureWorkaround(ncclEndEvent_)) {
+#else
   if (!ncclEndEvent_->query()) {
+#endif
     return false;
   }
   return true;
@@ -892,7 +960,7 @@ ProcessGroupNCCL::ProcessGroupNCCL(
       options_(std::move(options)),
       terminateProcessGroup_(false),
       local_id_(process_group_id++),
-      intraNodeComm_(initIntraNodeComm()) {
+      intraNodeComm_(nullptr) {
   TORCH_CHECK_WITH(
       ValueError,
       at::cuda::getNumGPUs() != 0,
@@ -922,6 +990,7 @@ ProcessGroupNCCL::ProcessGroupNCCL(
   PrefixStore* prefixStore = dynamic_cast<PrefixStore*>(store_.get());
   globalStore_ =
       prefixStore ? prefixStore->getUnderlyingNonPrefixStore() : store_;
+  debugInfoPipeFile_ = getCvarString({"TORCH_NCCL_DEBUG_INFO_PIPE_FILE"}, "");
   auto desyncDebug = getCvarBool(TORCH_NCCL_DESYNC_DEBUG, false) ||
       (dist_debug_level_ >= DebugLevel::Detail);
 #ifdef ENABLE_NCCL_ERROR_CHECKING
@@ -1197,7 +1266,10 @@ c10::intrusive_ptr<intra_node_comm::IntraNodeComm> ProcessGroupNCCL::
     return nullptr;
   }
   auto prefixStore = c10::make_intrusive<PrefixStore>("IntraNodeComm", store_);
-  auto comm = c10::make_intrusive<IntraNodeComm>(prefixStore, rank_, size_);
+  const std::string groupName =
+      options_->group_name.empty() ? "0" : options_->group_name;
+  auto comm = c10::make_intrusive<IntraNodeComm>(
+      prefixStore, rank_, size_, std::nullopt, groupName);
   if (comm->rendezvous()) {
     return comm;
   } else {
@@ -1778,7 +1850,8 @@ void ProcessGroupNCCL::HeartbeatMonitor::runLoop() {
     // DumpPipe is one per-trainer process, and its convenient to name them
     // after 'global' ranks in the system, So we assume processgroup (uid)==0 is
     // the global PG and has globally unique rank ids across trainers.
-    dumpPipe.emplace(pg_->globalRank());
+    dumpPipe.emplace(
+        pg_->globalRank(), pg_->debugInfoPipeFile_, pg_->traceBufferSize_);
   }
   while (true) {
     // This won't have any lock since this lock is only used here.
@@ -2008,12 +2081,17 @@ void ProcessGroupNCCL::HeartbeatMonitor::runLoop() {
   //
   // Or we get stuck in destructors, we will sleep for some time before calling
   // std::abort() to kill the whole process.
-  if ((pg_->terminateProcessGroup_.load() || shouldDump_.load()) &&
-      !terminateHeartbeatMonitorThread_.load()) {
-    std::this_thread::sleep_for(std::chrono::seconds(heartbeatTimeoutInSec_));
-    LOG(INFO)
-        << pg_->logPrefix() << "slept for " << heartbeatTimeoutInSec_
-        << " because we want to wait longer to verify there is indeed a watchdog hang.";
+  if (pg_->terminateProcessGroup_.load() || shouldDump_.load()) {
+    for (int t = 0; t < heartbeatTimeoutInSec_; ++t) {
+      if (terminateHeartbeatMonitorThread_.load()) {
+        if (t > 0)
+          LOG(INFO)
+              << pg_->logPrefix() << "slept for " << t
+              << " seconds because we want to wait longer to verify there is indeed a watchdog hang.";
+        break;
+      }
+      std::this_thread::sleep_for(std::chrono::seconds(1));
+    }
   }
 
   // At this point, we either already sleep for another `heartbeatTimeoutInSec_`
@@ -2284,9 +2362,23 @@ void ProcessGroupNCCL::Watchdog::runLoop() {
         }
       }
 
-      // Then check if work has timed out
-      // Skip if work has encountered an error
-      bool timedout = !work.exception() && work.checkTimeout();
+      // Then check if work has timed out.
+      // Skip if work has encountered an error.
+
+      bool timedout = false;
+#ifdef USE_ROCM
+      // On ROCm, watchdog event queries may be intentionally skipped during
+      // active graph capture to avoid HIP runtime capture invalidation.
+      // In that window, timeout checks can report false positives for
+      // otherwise-complete work, so we defer timeout enforcement.
+      // TODO: Remove once all supported HIP runtimes include:
+      // https://github.com/ROCm/clr/pull/3176
+      if (!at::cuda::is_graph_capture_active()) {
+        timedout = !work.exception() && work.checkTimeout();
+      }
+#else
+      timedout = !work.exception() && work.checkTimeout();
+#endif
 
       // Report desync state in case of timeout (if TORCH_NCCL_DESYNC_DEBUG is
       // turned on; otherwise, run() is no-op)
@@ -2351,7 +2443,11 @@ void ProcessGroupNCCL::Watchdog::runLoop() {
       // allow watchdog to do an event query on a side thread
       at::cuda::CUDAGuard device_guard(work.ncclEndEvent_->device_index());
       at::cuda::CUDAStreamCaptureModeGuard g{cudaStreamCaptureModeThreadLocal};
-
+#ifdef USE_ROCM
+      // Mark this thread/scope as watchdog event-query context so the ROCm
+      // workaround applies only here (not to main-thread wait()/isCompleted()).
+      RocmWatchdogEventQueryContextGuard watchdog_event_query_context_guard;
+#endif
       // a work could be started but not completed, so we should not update
       // lastStartedSeq and lastStartedOpName if the work state is checked
       // multiple times after the start
@@ -3270,6 +3366,30 @@ uint64_t ProcessGroupNCCL::getCommSplitCounter() const {
   return ret;
 }
 
+void ProcessGroupNCCL::suspend() {
+  auto device = at::Device(at::kCUDA, guessDeviceId());
+  std::string deviceKey = getKeyFromDevice(device);
+  auto ncclComm = getNCCLComm(deviceKey);
+  TORCH_CHECK(ncclComm != nullptr, "NCCL communicator not initialized.");
+  ncclComm->suspend();
+}
+
+void ProcessGroupNCCL::resume() {
+  auto device = at::Device(at::kCUDA, guessDeviceId());
+  std::string deviceKey = getKeyFromDevice(device);
+  auto ncclComm = getNCCLComm(deviceKey);
+  TORCH_CHECK(ncclComm != nullptr, "NCCL communicator not initialized.");
+  ncclComm->resume();
+}
+
+std::unordered_map<std::string, uint64_t> ProcessGroupNCCL::getMemoryStats() {
+  auto device = at::Device(at::kCUDA, guessDeviceId());
+  std::string deviceKey = getKeyFromDevice(device);
+  auto ncclComm = getNCCLComm(deviceKey);
+  TORCH_CHECK(ncclComm != nullptr, "NCCL communicator not initialized.");
+  return ncclComm->getMemoryStats();
+}
+
 namespace {
 
 // Check validity of tensor
@@ -3742,8 +3862,9 @@ c10::intrusive_ptr<Work> ProcessGroupNCCL::collective(
   }
 
   if (nanCheck) {
+    at::cuda::CUDAStreamGuard guard(ncclStream);
     for (const auto& input : inputs) {
-      checkForNan(input, ncclStream);
+      checkForNan(input);
     }
   }
 
@@ -4220,7 +4341,8 @@ c10::intrusive_ptr<Work> ProcessGroupNCCL::pointToPoint(
   // Only check for NaN for send ops, for recv ops `tensor` can be a random
   // placeholder
   if (enableNanCheck_ && opType == OpType::SEND) {
-    checkForNan(tensor, ncclStream);
+    at::cuda::CUDAStreamGuard guard(ncclStream);
+    checkForNan(tensor);
   }
 
   if (!coalescing_state_) {
@@ -4492,18 +4614,23 @@ c10::intrusive_ptr<Work> ProcessGroupNCCL::allreduce(
   }
   check_gpu_single_tensor(tensor);
 
-  if (intraNodeComm_ != nullptr && opts.reduceOp == ReduceOp::SUM) {
+  if (opts.reduceOp == ReduceOp::SUM) {
     using namespace intra_node_comm;
-    auto algo = intraNodeComm_->selectAllReduceAlgo(tensor);
-    if (algo != intra_node_comm::AllReduceAlgo::NONE) {
-      intraNodeComm_->allReduce(tensor, algo);
-      return c10::make_intrusive<IntraNodeCommWork>();
+    if (intraNodeComm_ == nullptr && IntraNodeComm::isEnabled()) {
+      intraNodeComm_ = initIntraNodeComm();
+    }
+    if (intraNodeComm_ != nullptr) {
+      auto algo = intraNodeComm_->selectAllReduceAlgo(tensor);
+      if (algo != intra_node_comm::AllReduceAlgo::NONE) {
+        intraNodeComm_->allReduce(tensor, algo);
+        return c10::make_intrusive<IntraNodeCommWork>();
+      }
     }
   }
   TORCH_CHECK(
       !isUnsupportedFloat8(tensor.scalar_type()),
       "Unsupported Float8 type for NCCL reduction");
-  RECORD_PARAM_COMMS_DATA(
+  RECORD_PARAM_COMMS_DATA_WITH_ASYNC_OP(
       std::make_tuple(
           static_cast<int64_t>(seqCollective_) + 1,
           false), // seq + 1 to match collective
@@ -4519,7 +4646,8 @@ c10::intrusive_ptr<Work> ProcessGroupNCCL::allreduce(
       std::vector<int64_t>(), // outSplitSizes
       globalRankStart_, // globalRankStart_
       globalRankStride_, // globalRankStride_
-      this->getSize()); // worldSize
+      this->getSize(), // worldSize
+      opts.asyncOp); // is asynchronized op
 
   // avoidRecordStreams_ note: collective() will stash tensors.
   return allreduce_impl(tensor, "nccl:all_reduce", opts);
@@ -4533,7 +4661,7 @@ c10::intrusive_ptr<Work> ProcessGroupNCCL::allreduce_coalesced(
       !isUnsupportedFloat8(tensors.back().scalar_type()),
       "Unsupported Float8 type for NCCL reduction");
 
-  RECORD_PARAM_COMMS_DATA(
+  RECORD_PARAM_COMMS_DATA_WITH_ASYNC_OP(
       std::make_tuple(
           static_cast<int64_t>(seqCollective_) + 1,
           false), // seq + 1 to match collective and assume only one collective
@@ -4551,7 +4679,8 @@ c10::intrusive_ptr<Work> ProcessGroupNCCL::allreduce_coalesced(
       std::vector<int64_t>(), // outSplitSizes
       globalRankStart_, // globalRankStart_
       globalRankStride_, // globalRankStride_
-      this->getSize()); // worldSize
+      this->getSize(), // worldSize
+      opts.asyncOp); // is asynchronized op
 
   // avoidRecordStreams_ note: collective() will stash tensors.
   return collectiveCoalesced(
@@ -4588,7 +4717,7 @@ c10::intrusive_ptr<Work> ProcessGroupNCCL::broadcast(
   }
   check_gpu_single_tensor(tensor);
 
-  RECORD_PARAM_COMMS_DATA(
+  RECORD_PARAM_COMMS_DATA_WITH_ASYNC_OP(
       std::make_tuple(
           static_cast<int64_t>(seqCollective_) + 1,
           false), // seq + 1 to match collective
@@ -4604,7 +4733,8 @@ c10::intrusive_ptr<Work> ProcessGroupNCCL::broadcast(
       std::vector<int64_t>(), // outSplitSizes
       globalRankStart_, // globalRankStart_
       globalRankStride_, // globalRankStride_
-      this->getSize()); // worldSize
+      this->getSize(), // worldSize
+      opts.asyncOp); // is asynchronized op
 
   const auto root = opts.rootRank + opts.rootTensor;
   bool nanCheck = (root == rank_);
@@ -4685,7 +4815,7 @@ c10::intrusive_ptr<Work> ProcessGroupNCCL::reduce(
     tensor = at::view_as_real(tensor);
   }
   check_gpu_single_tensor(tensor);
-  RECORD_PARAM_COMMS_DATA(
+  RECORD_PARAM_COMMS_DATA_WITH_ASYNC_OP(
       std::make_tuple(
           static_cast<int64_t>(seqCollective_) + 1,
           false), // seq + 1 to match collective
@@ -4701,7 +4831,8 @@ c10::intrusive_ptr<Work> ProcessGroupNCCL::reduce(
       std::vector<int64_t>(), // outSplitSizes
       globalRankStart_, // globalRankStart_
       globalRankStride_, // globalRankStride_
-      this->getSize()); // worldSize
+      this->getSize(), // worldSize
+      opts.asyncOp); // is asynchronized op
 
   // avoidRecordStreams_ note: collective() will stash tensors.
   return collective(
@@ -4781,7 +4912,7 @@ c10::intrusive_ptr<Work> ProcessGroupNCCL::allgather(
   check_gpu_single_tensor(inputTensor);
   auto outputTensors_ = outputTensors.back();
 
-  RECORD_PARAM_COMMS_DATA(
+  RECORD_PARAM_COMMS_DATA_WITH_ASYNC_OP(
       std::make_tuple(
           static_cast<int64_t>(seqCollective_) + 1,
           false), // seq + 1 to match collective
@@ -4798,7 +4929,8 @@ c10::intrusive_ptr<Work> ProcessGroupNCCL::allgather(
       std::vector<int64_t>(), // outSplitSize
       globalRankStart_, // globalRankStart_
       globalRankStride_, // globalRankStride_
-      this->getSize()); // worldSize
+      this->getSize(), // worldSize
+      opts.asyncOp); // is asynchronized op
 
   bool same_size = check_same_size(outputTensors_);
   if (same_size) {
@@ -4881,7 +5013,7 @@ c10::intrusive_ptr<Work> ProcessGroupNCCL::allgather_into_tensor_coalesced(
     std::vector<at::Tensor>& outputs,
     std::vector<at::Tensor>& inputs,
     const AllgatherOptions& opts) {
-  RECORD_PARAM_COMMS_DATA(
+  RECORD_PARAM_COMMS_DATA_WITH_ASYNC_OP(
       std::make_tuple(
           static_cast<int64_t>(seqCollective_) + 1,
           false), // seq + 1 to match collective and assume only one collective
@@ -4898,7 +5030,8 @@ c10::intrusive_ptr<Work> ProcessGroupNCCL::allgather_into_tensor_coalesced(
       std::vector<int64_t>(), // outSplitSizes
       globalRankStart_, // globalRankStart_
       globalRankStride_, // globalRankStride_
-      this->getSize()); // worldSize
+      this->getSize(), // worldSize
+      opts.asyncOp); // is asynchronized op
 
   return collectiveCoalesced(
       inputs,
@@ -4932,7 +5065,7 @@ c10::intrusive_ptr<Work> ProcessGroupNCCL::reduce_scatter(
       !isUnsupportedFloat8(outputTensor.scalar_type()),
       "Unsupported Float8 type for NCCL reduction");
 
-  RECORD_PARAM_COMMS_DATA(
+  RECORD_PARAM_COMMS_DATA_WITH_ASYNC_OP(
       std::make_tuple(
           static_cast<int64_t>(seqCollective_) + 1,
           false), // seq + 1 to match collective
@@ -4948,7 +5081,8 @@ c10::intrusive_ptr<Work> ProcessGroupNCCL::reduce_scatter(
       std::vector<int64_t>(), // outSplitSizes
       globalRankStart_, // globalRankStart_
       globalRankStride_, // globalRankStride_
-      this->getSize()); // worldSize
+      this->getSize(), // worldSize
+      opts.asyncOp); // is asynchronized op
 
   bool same_size = check_same_size(inputTensors_);
   if (same_size) {
@@ -5044,7 +5178,7 @@ c10::intrusive_ptr<Work> ProcessGroupNCCL::_reduce_scatter_base(
   TORCH_CHECK(
       !isUnsupportedFloat8(tensor.scalar_type()),
       "Unsupported Float8 type for NCCL reduction");
-  RECORD_PARAM_COMMS_DATA(
+  RECORD_PARAM_COMMS_DATA_WITH_ASYNC_OP(
       std::make_tuple(
           static_cast<int64_t>(seqCollective_) + 1,
           false), // seq + 1 to match collective
@@ -5060,7 +5194,8 @@ c10::intrusive_ptr<Work> ProcessGroupNCCL::_reduce_scatter_base(
       std::vector<int64_t>(), // outSplitSizes
       globalRankStart_, // globalRankStart_
       globalRankStride_, // globalRankStride_
-      this->getSize()); // worldSize
+      this->getSize(), // worldSize
+      opts.asyncOp); // is asynchronized op
 
   // avoidRecordStreams_ note: collective() will stash inputs and outputs.
   // Note 2: for asyncOp = false, we don't want to record streams because we
@@ -5112,7 +5247,7 @@ c10::intrusive_ptr<Work> ProcessGroupNCCL::reduce_scatter_tensor_coalesced(
       !isUnsupportedFloat8(inputs.back().scalar_type()),
       "Unsupported Float8 type for NCCL reduction");
 
-  RECORD_PARAM_COMMS_DATA(
+  RECORD_PARAM_COMMS_DATA_WITH_ASYNC_OP(
       std::make_tuple(
           static_cast<int64_t>(seqCollective_) + 1,
           false), // seq + 1 to match collective and assume only one collective
@@ -5129,7 +5264,8 @@ c10::intrusive_ptr<Work> ProcessGroupNCCL::reduce_scatter_tensor_coalesced(
       std::vector<int64_t>(), // outSplitSizes
       globalRankStart_, // globalRankStart_
       globalRankStride_, // globalRankStride_
-      this->getSize()); // worldSize
+      this->getSize(), // worldSize
+      opts.asyncOp); // is asynchronized op
 
   return collectiveCoalesced(
       inputs,
@@ -5269,7 +5405,7 @@ c10::intrusive_ptr<Work> ProcessGroupNCCL::alltoall_base(
   check_gpu_single_tensor(outputTensor);
   check_gpu_single_tensor(inputTensor);
   if (outputSplitSizes.empty() && inputSplitSizes.empty()) {
-    RECORD_PARAM_COMMS_DATA(
+    RECORD_PARAM_COMMS_DATA_WITH_ASYNC_OP(
         std::make_tuple(
             static_cast<int64_t>(seqCollective_) + 1,
             false), // seq + 1 to match collective
@@ -5285,7 +5421,8 @@ c10::intrusive_ptr<Work> ProcessGroupNCCL::alltoall_base(
         std::vector<int64_t>(), // outSplitSizes
         globalRankStart_, // globalRankStart_
         globalRankStride_, // globalRankStride_
-        this->getSize()); // worldSize
+        this->getSize(), // worldSize
+        opts.asyncOp); // is asynchronized op
 
     // avoidRecordStreams_ note: collective() will stash inputTensors and
     // outputTensors.
@@ -5307,7 +5444,7 @@ c10::intrusive_ptr<Work> ProcessGroupNCCL::alltoall_base(
     c10d::checkSplitSizes(inputSplitSizes, inputTensor, size_);
     c10d::checkSplitSizes(outputSplitSizes, outputTensor, size_);
 
-    RECORD_PARAM_COMMS_DATA(
+    RECORD_PARAM_COMMS_DATA_WITH_ASYNC_OP(
         std::make_tuple(
             static_cast<int64_t>(seqCollective_) + 1,
             false), // seq + 1 to match collective
@@ -5323,7 +5460,8 @@ c10::intrusive_ptr<Work> ProcessGroupNCCL::alltoall_base(
         outputSplitSizes, // outSplitSizes
         globalRankStart_, // globalRankStart_
         globalRankStride_, // globalRankStride_
-        this->getSize()); // worldSize
+        this->getSize(), // worldSize
+        opts.asyncOp); // is asynchronized op
 
     // avoidRecordStreams_ note: collective() will stash inputTensors and
     // outputTensors.
@@ -5387,7 +5525,7 @@ c10::intrusive_ptr<Work> ProcessGroupNCCL::alltoall(
     outSplitSizes.push_back(outputTensors[r].numel());
   }
 
-  RECORD_PARAM_COMMS_DATA(
+  RECORD_PARAM_COMMS_DATA_WITH_ASYNC_OP(
       std::make_tuple(
           static_cast<int64_t>(seqCollective_) + 1,
           false), // seq + 1 to match collective
@@ -5403,7 +5541,8 @@ c10::intrusive_ptr<Work> ProcessGroupNCCL::alltoall(
       outSplitSizes, // outSplitSizes
       globalRankStart_, // globalRankStart_
       globalRankStride_, // globalRankStride_
-      this->getSize()); // worldSize
+      this->getSize(), // worldSize
+      opts.asyncOp); // is asynchronized op
 
   return collective(
       inputTensors,
@@ -5432,7 +5571,7 @@ c10::intrusive_ptr<Work> ProcessGroupNCCL::send(
   auto tensor = tensors.back();
   check_gpu_single_tensor(tensor, true);
 
-  RECORD_PARAM_COMMS_DATA(
+  RECORD_PARAM_COMMS_DATA_WITH_ASYNC_OP(
       std::make_tuple(
           static_cast<int64_t>(seqP2P_) + (coalescing_state_ & CoalP2P ? 0 : 1),
           true), // the 1st p2p in coalesced range sets coalescing_state_ and
@@ -5449,7 +5588,8 @@ c10::intrusive_ptr<Work> ProcessGroupNCCL::send(
       std::vector<int64_t>(), // outSplitSizes
       globalRankStart_, // globalRankStart_
       globalRankStride_, // globalRankStride_
-      this->getSize()); // worldSize
+      this->getSize(), // worldSize
+      true); // is asynchronized op
 
   auto ret = pointToPoint(
       tensor,
@@ -5480,7 +5620,7 @@ c10::intrusive_ptr<Work> ProcessGroupNCCL::recv(
   auto tensor = tensors.back();
   check_gpu_single_tensor(tensor, true);
 
-  RECORD_PARAM_COMMS_DATA(
+  RECORD_PARAM_COMMS_DATA_WITH_ASYNC_OP(
       std::make_tuple(
           static_cast<int64_t>(seqP2P_) + (coalescing_state_ & CoalP2P ? 0 : 1),
           true), // the 1st p2p in coalesced range sets coalescing_state_ and
@@ -5497,7 +5637,8 @@ c10::intrusive_ptr<Work> ProcessGroupNCCL::recv(
       std::vector<int64_t>(), // outSplitSizes
       globalRankStart_, // globalRankStart_
       globalRankStride_, // globalRankStride_
-      this->getSize()); // worldSize
+      this->getSize(), // worldSize
+      true); // is asynchronized op
 
   auto ret = pointToPoint(
       tensor,
@@ -5589,7 +5730,7 @@ c10::intrusive_ptr<Work> ProcessGroupNCCL::gather(
     outputs.emplace_back();
   }
 
-  RECORD_PARAM_COMMS_DATA(
+  RECORD_PARAM_COMMS_DATA_WITH_ASYNC_OP(
       std::make_tuple(
           static_cast<int64_t>(seqCollective_) + 1,
           false), // seq + 1 to match collective
@@ -5605,7 +5746,8 @@ c10::intrusive_ptr<Work> ProcessGroupNCCL::gather(
       std::vector<int64_t>(), // outSplitSize
       globalRankStart_, // globalRankStart_
       globalRankStride_, // globalRankStride_
-      this->getSize()); // worldSize
+      this->getSize(), // worldSize
+      opts.asyncOp); // is asynchronized op
 
   // avoidRecordStreams_ note: collective() will stash inputTensors and
   // outputs, which == outputTensors[0] on the root rank where it matters.
@@ -5677,7 +5819,7 @@ c10::intrusive_ptr<Work> ProcessGroupNCCL::scatter(
     inputs.emplace_back();
   }
 
-  RECORD_PARAM_COMMS_DATA(
+  RECORD_PARAM_COMMS_DATA_WITH_ASYNC_OP(
       std::make_tuple(
           static_cast<int64_t>(seqCollective_) + 1,
           false), // seq + 1 to match collective
@@ -5693,7 +5835,8 @@ c10::intrusive_ptr<Work> ProcessGroupNCCL::scatter(
       std::vector<int64_t>(), // outSplitSize
       globalRankStart_, // globalRankStart_
       globalRankStride_, // globalRankStride_
-      this->getSize()); // worldSize
+      this->getSize(), // worldSize
+      opts.asyncOp); // is asynchronized op
 
   // avoidRecordStreams_ note: collective() will stash outputTensors and
   // inputs, which == inputTensors[0] on the root rank where it matters.
@@ -5747,7 +5890,7 @@ c10::intrusive_ptr<Work> ProcessGroupNCCL::_allgather_base(
         "output tensor size must be equal to world_size times input tensor size");
   }
 
-  RECORD_PARAM_COMMS_DATA(
+  RECORD_PARAM_COMMS_DATA_WITH_ASYNC_OP(
       std::make_tuple(
           static_cast<int64_t>(seqCollective_) + 1,
           false), // seq + 1 to match collective
@@ -5763,7 +5906,8 @@ c10::intrusive_ptr<Work> ProcessGroupNCCL::_allgather_base(
       std::vector<int64_t>(), // outSplitSize
       globalRankStart_, // globalRankStart_
       globalRankStride_, // globalRankStride_
-      this->getSize()); // worldSize
+      this->getSize(), // worldSize
+      opts.asyncOp); // is asynchronized op
 
   // avoidRecordStreams_ note: collective() will stash inputs and outputs.
   // Note 2: for asyncOp = false, we don't want to record streams because we
diff --git a/torch/csrc/distributed/c10d/ProcessGroupNCCL.hpp b/torch/csrc/distributed/c10d/ProcessGroupNCCL.hpp
index 12aeb49660f6c..0069c556ca038 100644
--- a/torch/csrc/distributed/c10d/ProcessGroupNCCL.hpp
+++ b/torch/csrc/distributed/c10d/ProcessGroupNCCL.hpp
@@ -195,16 +195,10 @@ static std::vector<std::string> TORCH_NCCL_USE_TENSOR_REGISTER_ALLOCATOR_HOOK =
 
 #if defined(__linux__)
 struct DumpPipe {
-  DumpPipe(int rank) {
-    std::string fileStem =
-        getCvarString({"TORCH_NCCL_DEBUG_INFO_PIPE_FILE"}, "");
-    // NOTE: This default value (2000) is duplicated in FlightRecorder.hpp.
-    // Keep in sync. See FlightRecorder.hpp for details.
-    if (fileStem.empty() ||
-        getCvarInt({"TORCH_NCCL_TRACE_BUFFER_SIZE"}, 2000) <= 0) {
+  DumpPipe(int rank, const std::string& fileStem, int traceBufferSize) {
+    if (fileStem.empty() || traceBufferSize <= 0) {
       return;
     }
-    TORCH_CHECK(!fileStem.empty(), "TORCH_NCCL_DEBUG_INFO_PIPE_FILE is empty");
     std::string filename = c10::str(fileStem, rank, ".pipe");
     TORCH_CHECK(
         unlink(filename.c_str()) != -1 || errno == ENOENT,
@@ -1057,6 +1051,13 @@ class TORCH_API ProcessGroupNCCL : public Backend {
 
   void setEnableNanCheck(bool enableNanCheck);
 
+  // APIs related to memory offload (require NCCL 2.29.7+ at runtime)
+  void suspend() override;
+
+  void resume() override;
+
+  std::unordered_map<std::string, uint64_t> getMemoryStats() override;
+
  protected:
   uint64_t getWatchdogHeartbt() const;
 
@@ -1357,6 +1358,9 @@ class TORCH_API ProcessGroupNCCL : public Backend {
   // Size of ring buffer where we store NCCL Traces for debugging.
   int traceBufferSize_;
 
+  // Stores TORCH_NCCL_DEBUG_INFO_PIPE_FILE
+  std::string debugInfoPipeFile_;
+
   // We gate the cudaEventCache so that we can roll it out gradually.
   std::atomic<bool> cudaEventCacheEnabled_;
 
diff --git a/torch/csrc/distributed/c10d/ProcessGroupWrapper.cpp b/torch/csrc/distributed/c10d/ProcessGroupWrapper.cpp
index 9dd9c6d0a6c35..b19750fb146ba 100644
--- a/torch/csrc/distributed/c10d/ProcessGroupWrapper.cpp
+++ b/torch/csrc/distributed/c10d/ProcessGroupWrapper.cpp
@@ -416,6 +416,13 @@ c10::intrusive_ptr<Work> ProcessGroupWrapper::allreduce_coalesced(
   return backend_->allreduce_coalesced(tensors, opts);
 }
 
+c10::intrusive_ptr<Work> ProcessGroupWrapper::allreduce_sparse(
+    std::vector<at::Tensor>& tensors,
+    const AllreduceOptions& opts) {
+  runCollectiveChecks(OpType::_ALLREDUCE_SPARSE, tensors);
+  return backend_->allreduce_sparse(tensors, opts);
+}
+
 c10::intrusive_ptr<Work> ProcessGroupWrapper::reduce(
     std::vector<at::Tensor>& tensors,
     const ReduceOptions& opts) {
@@ -456,6 +463,13 @@ c10::intrusive_ptr<Work> ProcessGroupWrapper::allgather_coalesced(
   return backend_->allgather_coalesced(outputTensorLists, inputTensors, opts);
 }
 
+c10::intrusive_ptr<Work> ProcessGroupWrapper::allgather_into_tensor_coalesced(
+    std::vector<at::Tensor>& outputs,
+    std::vector<at::Tensor>& inputs,
+    const AllgatherOptions& opts) {
+  return backend_->allgather_into_tensor_coalesced(outputs, inputs, opts);
+}
+
 c10::intrusive_ptr<Work> ProcessGroupWrapper::gather(
     std::vector<std::vector<at::Tensor>>& outputTensors,
     std::vector<at::Tensor>& inputTensors,
@@ -558,6 +572,18 @@ c10::intrusive_ptr<Work> ProcessGroupWrapper::_reduce_scatter_base(
   return backend_->_reduce_scatter_base(outputBuffer, inputBuffer, opts);
 }
 
+c10::intrusive_ptr<Work> ProcessGroupWrapper::reduce_scatter_tensor_coalesced(
+    std::vector<at::Tensor>& outputs,
+    std::vector<at::Tensor>& inputs,
+    const ReduceScatterOptions& opts) {
+  // NOTE: We don't enforce shape checking for reduce_scatter_tensor_coalesced
+  // because the implementation itself does not enforce it we have tests that
+  // use inconsistent shapes, see python implementation in distributed_c10d for
+  // details.
+  runCollectiveChecks(OpType::REDUCE_SCATTER_TENSOR_COALESCED, {});
+  return backend_->reduce_scatter_tensor_coalesced(outputs, inputs, opts);
+}
+
 void ProcessGroupWrapper::startCoalescing() {
   return backend_->startCoalescing();
 }
@@ -578,6 +604,21 @@ bool ProcessGroupWrapper::supportsTimeEstimation() const {
   return backend_->supportsTimeEstimation();
 }
 
+bool ProcessGroupWrapper::supportsShrinking() const {
+  return backend_->supportsShrinking();
+}
+
+c10::intrusive_ptr<Backend> ProcessGroupWrapper::shrink(
+    const std::vector<int64_t>& ranks_to_exclude,
+    int shrink_flags,
+    const c10::intrusive_ptr<Options>& opts_override) {
+  return backend_->shrink(ranks_to_exclude, shrink_flags, opts_override);
+}
+
+void ProcessGroupWrapper::setTimeout(std::chrono::milliseconds timeout) {
+  backend_->setTimeout(timeout);
+}
+
 c10::intrusive_ptr<Backend::Options> ProcessGroupWrapper::getBackendOptions() {
   return backend_->getBackendOptions();
 }
@@ -596,6 +637,27 @@ bool ProcessGroupWrapper::supportsTensorAlloc(c10::DeviceIndex deviceIdx) {
   return backend_->supportsTensorAlloc(deviceIdx);
 }
 
+void ProcessGroupWrapper::abort() {
+  backend_->abort();
+}
+
+void ProcessGroupWrapper::shutdown() {
+  backend_->shutdown();
+}
+
+void ProcessGroupWrapper::suspend() {
+  backend_->suspend();
+}
+
+void ProcessGroupWrapper::resume() {
+  backend_->resume();
+}
+
+std::unordered_map<std::string, uint64_t> ProcessGroupWrapper::
+    getMemoryStats() {
+  return backend_->getMemoryStats();
+}
+
 ErrorType ProcessGroupWrapper::getError() {
   return backend_->getError();
 }
@@ -604,6 +666,34 @@ void ProcessGroupWrapper::eagerConnectSingleDevice(at::Device device) {
   backend_->eagerConnectSingleDevice(device);
 }
 
+void ProcessGroupWrapper::registerOnCompletionHook(
+    std::function<void(std::shared_ptr<WorkInfo>)>&& hook) {
+  backend_->registerOnCompletionHook(std::move(hook));
+}
+
+void ProcessGroupWrapper::waitForPendingWorks() {
+  backend_->waitForPendingWorks();
+}
+
+void ProcessGroupWrapper::enableCollectivesTiming() {
+  backend_->enableCollectivesTiming();
+}
+
+c10::intrusive_ptr<Backend> ProcessGroupWrapper::split(
+    const c10::intrusive_ptr<Store>& store,
+    const std::vector<int>& ranks,
+    const c10::intrusive_ptr<Options>& opts) {
+  return backend_->split(store, ranks, opts);
+}
+
+c10::intrusive_ptr<Backend> ProcessGroupWrapper::merge(
+    const c10::intrusive_ptr<Store>& store,
+    const c10::intrusive_ptr<Options>& opts,
+    const int& rank,
+    const int& size) {
+  return backend_->merge(store, opts, rank, size);
+}
+
 c10::intrusive_ptr<Backend> ProcessGroupWrapper::getWrappedPg() const {
   return backend_;
 }
diff --git a/torch/csrc/distributed/c10d/ProcessGroupWrapper.hpp b/torch/csrc/distributed/c10d/ProcessGroupWrapper.hpp
index 1d70599dce025..b41166eed5535 100644
--- a/torch/csrc/distributed/c10d/ProcessGroupWrapper.hpp
+++ b/torch/csrc/distributed/c10d/ProcessGroupWrapper.hpp
@@ -23,8 +23,6 @@ class TORCH_API ProcessGroupWrapper : public Backend {
       const c10::intrusive_ptr<Backend>& backend,
       c10::intrusive_ptr<Backend> glooBackend);
 
-  const std::string getBackendName() const override;
-
   c10::intrusive_ptr<Work> broadcast(
       std::vector<at::Tensor>& data,
       const BroadcastOptions& opts = BroadcastOptions()) override;
@@ -33,6 +31,10 @@ class TORCH_API ProcessGroupWrapper : public Backend {
       std::vector<at::Tensor>& data,
       const AllreduceOptions& opts = AllreduceOptions()) override;
 
+  c10::intrusive_ptr<Work> allreduce_sparse(
+      std::vector<at::Tensor>& tensors,
+      const AllreduceOptions& opts = AllreduceOptions()) override;
+
   c10::intrusive_ptr<Work> allreduce_coalesced(
       std::vector<at::Tensor>& tensors,
       const AllreduceCoalescedOptions& opts =
@@ -61,6 +63,11 @@ class TORCH_API ProcessGroupWrapper : public Backend {
       std::vector<at::Tensor>& inputTensors,
       const AllgatherOptions& opts = AllgatherOptions()) override;
 
+  c10::intrusive_ptr<Work> allgather_into_tensor_coalesced(
+      std::vector<at::Tensor>& outputs,
+      std::vector<at::Tensor>& inputs,
+      const AllgatherOptions& opts = AllgatherOptions()) override;
+
   c10::intrusive_ptr<Work> gather(
       std::vector<std::vector<at::Tensor>>& outputTensors,
       std::vector<at::Tensor>& inputTensors,
@@ -76,6 +83,16 @@ class TORCH_API ProcessGroupWrapper : public Backend {
       std::vector<std::vector<at::Tensor>>& inputTensors,
       const ReduceScatterOptions& opts = ReduceScatterOptions()) override;
 
+  c10::intrusive_ptr<Work> _reduce_scatter_base(
+      at::Tensor& inputBuffer,
+      at::Tensor& outputBuffer,
+      const ReduceScatterOptions& opts = ReduceScatterOptions()) override;
+
+  c10::intrusive_ptr<Work> reduce_scatter_tensor_coalesced(
+      std::vector<at::Tensor>& outputs,
+      std::vector<at::Tensor>& inputs,
+      const ReduceScatterOptions& opts = ReduceScatterOptions()) override;
+
   c10::intrusive_ptr<Work> alltoall_base(
       at::Tensor& outputTensor,
       at::Tensor& inputTensor,
@@ -118,24 +135,46 @@ class TORCH_API ProcessGroupWrapper : public Backend {
 
   c10::intrusive_ptr<Work> barrier(
       const BarrierOptions& opts = BarrierOptions()) override;
+  void registerOnCompletionHook(
+      std::function<void(std::shared_ptr<WorkInfo>)>&& hook) override;
 
-  c10::intrusive_ptr<Work> _reduce_scatter_base(
-      at::Tensor& outputBuffer,
-      at::Tensor& inputBuffer,
-      const ReduceScatterOptions& opts) override;
+  void waitForPendingWorks() override;
+  void enableCollectivesTiming() override;
 
-  void startCoalescing() override;
+  c10::intrusive_ptr<Backend> split(
+      const c10::intrusive_ptr<Store>& store,
+      const std::vector<int>& ranks,
+      const c10::intrusive_ptr<Options>& opts) override;
 
-  c10::intrusive_ptr<Work> endCoalescing() override;
+  c10::intrusive_ptr<Backend> merge(
+      const c10::intrusive_ptr<Store>& store,
+      const c10::intrusive_ptr<Options>& opts,
+      const int& rank,
+      const int& size) override;
 
   // Forward methods to wrapped backend
   bool supportsSplitting() const override;
   bool supportsCoalescing() const override;
   bool supportsTimeEstimation() const override;
+  bool supportsShrinking() const override;
+  c10::intrusive_ptr<Backend> shrink(
+      const std::vector<int64_t>& ranks_to_exclude,
+      int shrink_flags = 0,
+      const c10::intrusive_ptr<Options>& opts_override = nullptr) override;
+  void setTimeout(std::chrono::milliseconds timeout) override;
+  void startCoalescing() override;
+  c10::intrusive_ptr<Work> endCoalescing() override;
+  const std::string getBackendName() const override;
   c10::intrusive_ptr<Options> getBackendOptions() override;
   std::shared_ptr<c10::Allocator> getMemAllocator() override;
   at::Tensor allocateTensor(long size, at::TensorOptions options = {}) override;
   bool supportsTensorAlloc(c10::DeviceIndex deviceIdx) override;
+  void abort() override;
+  void shutdown() override;
+  void suspend() override;
+  void resume() override;
+  std::unordered_map<std::string, uint64_t> getMemoryStats() override;
+
   ErrorType getError() override;
   void eagerConnectSingleDevice(at::Device device) override;
 
diff --git a/torch/csrc/distributed/c10d/Store.cpp b/torch/csrc/distributed/c10d/Store.cpp
index e7fc0705f734c..76a9bd3e1a7bf 100644
--- a/torch/csrc/distributed/c10d/Store.cpp
+++ b/torch/csrc/distributed/c10d/Store.cpp
@@ -66,4 +66,20 @@ bool Store::hasExtendedApi() const {
   return false;
 }
 
+void Store::barrier(
+    const std::string& key,
+    int64_t world_size,
+    const std::chrono::milliseconds& timeout) {
+  // Default implementation using add() + wait() pattern.
+  // Subclasses can override this with optimized implementations.
+  std::string numMembersKey = key + "/num_members";
+  std::string lastMemberKey = key + "/last_member";
+
+  int64_t idx = add(numMembersKey, 1);
+  if (idx == world_size) {
+    set(lastMemberKey, std::vector<uint8_t>{'1'});
+  }
+  wait({lastMemberKey}, timeout);
+}
+
 } // namespace c10d
diff --git a/torch/csrc/distributed/c10d/Store.hpp b/torch/csrc/distributed/c10d/Store.hpp
index 9a037c65ee7c2..ed29d8e298aba 100644
--- a/torch/csrc/distributed/c10d/Store.hpp
+++ b/torch/csrc/distributed/c10d/Store.hpp
@@ -119,6 +119,19 @@ class TORCH_API Store : public torch::CustomClassHolder {
         NotImplementedError, "listKeys support is not implemented.");
   }
 
+  // Barrier operation that blocks until world_size workers have reached it.
+  // This is an optimized operation that combines increment and wait into a
+  // single operation, reducing network round trips compared to using
+  // separate add() and wait() calls.
+  virtual void barrier(
+      const std::string& key,
+      int64_t world_size,
+      const std::chrono::milliseconds& timeout);
+
+  void barrier(const std::string& key, int64_t world_size) {
+    barrier(key, world_size, timeout_);
+  }
+
  protected:
   std::chrono::milliseconds timeout_;
 };
diff --git a/torch/csrc/distributed/c10d/TCPStore.cpp b/torch/csrc/distributed/c10d/TCPStore.cpp
index 9f566032b5b3c..b5c52eb659093 100644
--- a/torch/csrc/distributed/c10d/TCPStore.cpp
+++ b/torch/csrc/distributed/c10d/TCPStore.cpp
@@ -747,6 +747,54 @@ std::vector<std::string> TCPStore::listKeys() {
   return keys;
 }
 
+void TCPStore::barrier(
+    const std::string& key,
+    int64_t world_size,
+    const std::chrono::milliseconds& timeout) {
+  STATIC_SCOPED_WAIT_COUNTER(pytorch.wait_counter.TCPStore__barrier);
+  const std::lock_guard<std::mutex> lock(activeOpLock_);
+
+  detail::SendBuffer buffer(*client_, detail::QueryType::BARRIER);
+  buffer.appendString(keyPrefix_ + key);
+  buffer.appendValue<int64_t>(world_size);
+  buffer.flush();
+
+  auto response_opt =
+      client_->receiveValueWithTimeout<detail::WaitResponseType>(timeout);
+  if (response_opt.has_value()) {
+    if (response_opt != detail::WaitResponseType::STOP_WAITING) {
+      TORCH_CHECK_WITH(
+          DistStoreError, false, "Stop_waiting response is expected");
+    }
+    return;
+  }
+
+  // Timeout occurred - send cancel and handle response
+  {
+    detail::SendBuffer cancelBuffer(*client_, detail::QueryType::CANCEL_WAIT);
+    cancelBuffer.flush();
+  }
+
+  auto response = client_->receiveValue<detail::WaitResponseType>();
+  // This can happen if the server responds before we cancel
+  if (response != detail::WaitResponseType::WAIT_CANCELED) {
+    if (response != detail::WaitResponseType::STOP_WAITING) {
+      TORCH_CHECK_WITH(
+          DistStoreError, false, "Stop_waiting response is expected");
+    }
+    // Wait for the cancel acknowledgment
+    response = client_->receiveValue<detail::WaitResponseType>();
+    if (response != detail::WaitResponseType::WAIT_CANCELED) {
+      TORCH_CHECK_WITH(
+          DistStoreError, false, "wait_canceled response is expected");
+    }
+  }
+
+  C10_THROW_ERROR(
+      DistStoreError,
+      fmt::format("barrier timeout after {}ms, key: {}", timeout.count(), key));
+}
+
 bool TCPStore::hasExtendedApi() const {
   return true;
 }
diff --git a/torch/csrc/distributed/c10d/TCPStore.hpp b/torch/csrc/distributed/c10d/TCPStore.hpp
index 09d7ae111c57a..5d708929dbdf8 100644
--- a/torch/csrc/distributed/c10d/TCPStore.hpp
+++ b/torch/csrc/distributed/c10d/TCPStore.hpp
@@ -123,6 +123,11 @@ class TORCH_API TCPStore : public Store {
 
   std::vector<std::string> listKeys() override;
 
+  void barrier(
+      const std::string& key,
+      int64_t world_size,
+      const std::chrono::milliseconds& timeout) override;
+
   // Waits for all workers to join.
   void waitForWorkers();
 
diff --git a/torch/csrc/distributed/c10d/TCPStoreBackend.cpp b/torch/csrc/distributed/c10d/TCPStoreBackend.cpp
index dd25729a6ee13..f718886445f07 100644
--- a/torch/csrc/distributed/c10d/TCPStoreBackend.cpp
+++ b/torch/csrc/distributed/c10d/TCPStoreBackend.cpp
@@ -79,6 +79,7 @@ class TCPStoreMasterDaemon : public BackgroundThread {
   void multiSetHandler(int socket);
   void cancelWaitHandler(int socket);
   void listKeysHandler(int socket);
+  void barrierHandler(int socket);
   void addMiscellaneousSocket(int socket);
   void removeMiscellaneousSocket(int socket);
   bool isMiscellaneousSocket(int socket);
@@ -220,27 +221,11 @@ void TCPStoreMasterDaemon::queryFds(std::vector<struct pollfd>& fds) {
 
 void TCPStoreMasterDaemon::clearSocketWaitState(int socket) {
   // Remove all the tracking state of the close FD
-  for (auto it = waitingSockets_.begin(); it != waitingSockets_.end();) {
-    for (auto vecIt = it->second.begin(); vecIt != it->second.end();) {
-      if (*vecIt == socket) {
-        vecIt = it->second.erase(vecIt);
-      } else {
-        ++vecIt;
-      }
-    }
-    if (it->second.empty()) {
-      it = waitingSockets_.erase(it);
-    } else {
-      ++it;
-    }
-  }
-  for (auto it = keysAwaited_.begin(); it != keysAwaited_.end();) {
-    if (it->first == socket) {
-      it = keysAwaited_.erase(it);
-    } else {
-      ++it;
-    }
-  }
+  std::erase_if(waitingSockets_, [&](auto& entry) {
+    std::erase(entry.second, socket);
+    return entry.second.empty();
+  });
+  keysAwaited_.erase(socket);
 }
 
 // query communicates with the worker. The format
@@ -298,6 +283,8 @@ void TCPStoreMasterDaemon::query(int socket) {
     cancelWaitHandler(socket);
   } else if (qt == QueryType::LIST_KEYS) {
     listKeysHandler(socket);
+  } else if (qt == QueryType::BARRIER) {
+    barrierHandler(socket);
   } else {
     TORCH_CHECK(false, "Unexpected query type");
   }
@@ -492,6 +479,34 @@ void TCPStoreMasterDaemon::listKeysHandler(int socket) {
   }
 }
 
+void TCPStoreMasterDaemon::barrierHandler(int socket) {
+  std::string key = tcputil::recvString(socket);
+  int64_t worldSize = tcputil::recvValue<int64_t>(socket);
+
+  // Atomically increment the barrier counter
+  auto it = tcpStore_.find(key);
+  int64_t count = 1;
+  if (it != tcpStore_.end()) {
+    auto buf = reinterpret_cast<const char*>(it->second.data());
+    auto len = it->second.size();
+    count = std::stoll(std::string(buf, len)) + 1;
+  }
+  auto countStr = std::to_string(count);
+  tcpStore_[key] = std::vector<uint8_t>(countStr.begin(), countStr.end());
+
+  if (count >= worldSize) {
+    // All workers have arrived, notify this client
+    tcputil::sendValue<WaitResponseType>(
+        socket, WaitResponseType::STOP_WAITING);
+    // Wake up all previously waiting clients
+    wakeupWaitingClients(key);
+  } else {
+    // Register this client to wait for remaining workers
+    waitingSockets_[key].push_back(socket);
+    keysAwaited_[socket] = 1;
+  }
+}
+
 bool TCPStoreMasterDaemon::checkKeys(
     const std::vector<std::string>& keys) const {
   return std::all_of(keys.begin(), keys.end(), [this](const std::string& s) {
diff --git a/torch/csrc/distributed/c10d/TCPStoreBackend.hpp b/torch/csrc/distributed/c10d/TCPStoreBackend.hpp
index d176ccb702838..9c2bccf2c4cd3 100644
--- a/torch/csrc/distributed/c10d/TCPStoreBackend.hpp
+++ b/torch/csrc/distributed/c10d/TCPStoreBackend.hpp
@@ -37,6 +37,7 @@ enum class QueryType : uint8_t {
   QUEUE_POP,
   QUEUE_LEN,
   LIST_KEYS,
+  BARRIER,
 };
 
 enum class CheckResponseType : uint8_t { READY, NOT_READY };
diff --git a/torch/csrc/distributed/c10d/TCPStoreLibUvBackend.cpp b/torch/csrc/distributed/c10d/TCPStoreLibUvBackend.cpp
index 7427848b8445b..f39f829953fd6 100644
--- a/torch/csrc/distributed/c10d/TCPStoreLibUvBackend.cpp
+++ b/torch/csrc/distributed/c10d/TCPStoreLibUvBackend.cpp
@@ -685,6 +685,11 @@ class LibUVStoreDaemon : public BackgroundThread {
   int64_t queueLen(const std::string& queueName);
   std::vector<std::string> listKeys();
 
+  bool barrier(
+      const std::string& key,
+      int64_t worldSize,
+      const c10::intrusive_ptr<UvHandle>& client);
+
   void registerClient(const c10::intrusive_ptr<UvHandle>& client);
   void unregisterClient(const c10::intrusive_ptr<UvHandle>& client);
   void clearClientWaitState(const c10::intrusive_ptr<UvHandle>& client);
@@ -827,6 +832,10 @@ class UvClient : public UvTcpSocket {
             if (!parse_list_keys_command())
               return;
             break;
+          case QueryType::BARRIER:
+            if (!parse_barrier_command())
+              return;
+            break;
           default:
             C10D_DEBUG(
                 "Client sent invalid command. client:{} command:{}",
@@ -1182,6 +1191,30 @@ class UvClient : public UvTcpSocket {
     return true;
   }
 
+  bool parse_barrier_command() {
+    std::string key;
+    if (!stream.read_key(key))
+      return false;
+
+    int64_t worldSize = 0;
+    if (!stream.read_value(worldSize))
+      return false;
+
+    C10D_TRACE(
+        "barrier key:{} worldSize:{} address:{}",
+        key,
+        worldSize,
+        this->address());
+
+    if (store->barrier(key, worldSize, iptr())) {
+      StreamWriter sw(iptr());
+      sw.write1(static_cast<uint8_t>(WaitResponseType::STOP_WAITING));
+      sw.send();
+    }
+
+    return true;
+  }
+
  public:
   explicit UvClient(uv_loop_t* loop, LibUVStoreDaemon* store)
       : UvTcpSocket(loop), store(store) {}
@@ -1372,20 +1405,10 @@ void LibUVStoreDaemon::clearClientWaitState(
     return;
   }
   keysAwaited_.erase(client);
-  for (auto it = waitingSockets_.begin(); it != waitingSockets_.end();) {
-    for (auto vecIt = it->second.begin(); vecIt != it->second.end();) {
-      if (*vecIt == client) {
-        vecIt = it->second.erase(vecIt);
-      } else {
-        ++vecIt;
-      }
-    }
-    if (it->second.empty()) {
-      it = waitingSockets_.erase(it);
-    } else {
-      ++it;
-    }
-  }
+  std::erase_if(waitingSockets_, [&](auto& entry) {
+    std::erase(entry.second, client);
+    return entry.second.empty();
+  });
 }
 
 void LibUVStoreDaemon::set(
@@ -1570,6 +1593,33 @@ std::vector<std::string> LibUVStoreDaemon::listKeys() {
   return keys;
 }
 
+bool LibUVStoreDaemon::barrier(
+    const std::string& key,
+    int64_t worldSize,
+    const c10::intrusive_ptr<UvHandle>& client) {
+  // Atomically increment the barrier counter
+  auto it = tcpStore_.find(key);
+  int64_t count = 1;
+  if (it != tcpStore_.end()) {
+    auto buf = reinterpret_cast<const char*>(it->second.data());
+    auto len = it->second.size();
+    count = std::stoll(std::string(buf, len)) + 1;
+  }
+  auto countStr = std::to_string(count);
+  tcpStore_[key] = std::vector<uint8_t>(countStr.begin(), countStr.end());
+
+  if (count >= worldSize) {
+    // Wake up all previously waiting clients
+    wakeupWaitingClients(key);
+    return true; // Caller should send STOP_WAITING to this client
+  } else {
+    // Register this client to wait for remaining workers
+    waitingSockets_[key].push_back(client);
+    keysAwaited_[client] = 1;
+    return false;
+  }
+}
+
 #endif
 
 std::unique_ptr<BackgroundThread> create_libuv_tcpstore_backend(
diff --git a/torch/csrc/distributed/c10d/Types.cpp b/torch/csrc/distributed/c10d/Types.cpp
index 300d21780bdb0..bf718a2fb9d61 100644
--- a/torch/csrc/distributed/c10d/Types.cpp
+++ b/torch/csrc/distributed/c10d/Types.cpp
@@ -16,7 +16,6 @@ bool isComplexViewAsRealAllowed(const ReduceOp& reduceOp) {
     default:
       return false;
   }
-  return false;
 }
 
 } // namespace c10d
diff --git a/torch/csrc/distributed/c10d/Work.cpp b/torch/csrc/distributed/c10d/Work.cpp
index cdec9185ce537..86adb9770faed 100644
--- a/torch/csrc/distributed/c10d/Work.cpp
+++ b/torch/csrc/distributed/c10d/Work.cpp
@@ -7,13 +7,44 @@
 
 namespace c10d {
 
+namespace {
+// Raw pointer avoids thread_local destructor issues in forked
+// processes and dynamically-loaded libraries.
+thread_local std::string* comm_profiling_name = nullptr;
+static_assert(
+    std::is_trivially_destructible_v<decltype(comm_profiling_name)>,
+    "comm_profiling_name must be trivially destructible — a non-trivial "
+    "destructor (e.g. std::string) causes deadlocks after fork() in "
+    "dlopen'd libraries via __cxa_thread_atexit.");
+} // namespace
+
+void set_comm_profiling_name(const std::string& name) {
+  if (!comm_profiling_name) {
+    comm_profiling_name = new std::string(name);
+  } else {
+    *comm_profiling_name = name;
+  }
+}
+
+const std::string& get_comm_profiling_name() {
+  if (comm_profiling_name) {
+    return *comm_profiling_name;
+  }
+  static const std::string empty;
+  return empty;
+}
+
 Work::Work(
     int rank,
     OpType opType,
     const char* profilingTitle,
     const std::optional<std::vector<at::Tensor>>& inputTensors)
     : rank_(rank), opType_(opType) {
-  if (profilingTitle != nullptr) {
+  // comm_profiling_name is thread-local; take a local copy so the
+  // RecordFunction owns the string (the TLS can be mutated after we return).
+  const bool use_tls_name =
+      comm_profiling_name != nullptr && !comm_profiling_name->empty();
+  if (use_tls_name || profilingTitle != nullptr) {
     auto recordingFunction =
         std::make_shared<at::RecordFunction>(at::RecordScope::USER_SCOPE);
     if (recordingFunction->isActive()) {
@@ -29,9 +60,17 @@ Work::Work(
           inputs.emplace_back(tensor);
         }
       }
-      recordingFunction->before(
-          profilingTitle,
-          c10::ArrayRef<const c10::IValue>(inputs.data(), inputs.size()));
+      if (use_tls_name) {
+        recordingFunction->before(
+            std::string(*comm_profiling_name),
+            c10::ArrayRef<const c10::IValue>(inputs.data(), inputs.size()));
+      } else {
+        // const char* overload — pointer is a string literal with static
+        // lifetime
+        recordingFunction->before(
+            profilingTitle,
+            c10::ArrayRef<const c10::IValue>(inputs.data(), inputs.size()));
+      }
       std::function<void()> end_handler = [recordingFunction]() {
         recordingFunction->end();
       };
diff --git a/torch/csrc/distributed/c10d/Work.hpp b/torch/csrc/distributed/c10d/Work.hpp
index 2eeea75330fde..f7e4317c2ca29 100644
--- a/torch/csrc/distributed/c10d/Work.hpp
+++ b/torch/csrc/distributed/c10d/Work.hpp
@@ -31,6 +31,7 @@ enum class OpType : std::uint8_t {
   _REDUCE_SCATTER_BASE = 16,
   COALESCED = 17,
   _ALLREDUCE_SPARSE = 18,
+  REDUCE_SCATTER_TENSOR_COALESCED = 19,
   UNKNOWN = 100,
 };
 
@@ -182,4 +183,7 @@ struct TORCH_API WorkInfo {
   std::chrono::duration<float> activeDuration;
 };
 
+TORCH_API void set_comm_profiling_name(const std::string& name);
+TORCH_API const std::string& get_comm_profiling_name();
+
 } // namespace c10d
diff --git a/torch/csrc/distributed/c10d/cuda/AsyncMM.cu b/torch/csrc/distributed/c10d/cuda/AsyncMM.cu
index a7e9ddaa17b1e..2bec14bad98dd 100644
--- a/torch/csrc/distributed/c10d/cuda/AsyncMM.cu
+++ b/torch/csrc/distributed/c10d/cuda/AsyncMM.cu
@@ -79,7 +79,7 @@ at::Tensor async_input_mm_impl(
           cutlass::epilogue::collective::EpilogueTileAuto,
           ElementAccumulator,
           ElementAccumulator,
-          ElementC,
+          void, // Indicate there is no beta scaling to save register
           LayoutC,
           AlignmentC,
           ElementC,
@@ -151,7 +151,7 @@ at::Tensor async_input_mm_impl(
           stride_B,
       },
       {{},
-       reinterpret_cast<ElementC*>(out.data_ptr<at::BFloat16>()),
+       nullptr,
        stride_C,
        reinterpret_cast<ElementC*>(out.data_ptr<at::BFloat16>()),
        stride_C},
diff --git a/torch/csrc/distributed/c10d/cuda/StreamBlock.cu b/torch/csrc/distributed/c10d/cuda/StreamBlock.cu
index db4a118a25e59..79f295238765e 100644
--- a/torch/csrc/distributed/c10d/cuda/StreamBlock.cu
+++ b/torch/csrc/distributed/c10d/cuda/StreamBlock.cu
@@ -3,7 +3,7 @@
 #include <c10/cuda/CUDAException.h>
 #include <c10/cuda/CUDAGuard.h>
 #include <cuda_runtime.h>
-#include <torch/csrc/distributed/c10d/symm_mem/CUDASymmetricMemory-inl.h>
+#include <torch/csrc/distributed/c10d/symm_mem/CUDASymmetricMemory-inl.cuh>
 #include <torch/csrc/distributed/c10d/cuda/StreamBlock.cuh>
 
 #ifndef AT_PER_OPERATOR_HEADERS
diff --git a/torch/csrc/distributed/c10d/exception.h b/torch/csrc/distributed/c10d/exception.h
index 49c6dfd36b69a..3ec988c2aa749 100644
--- a/torch/csrc/distributed/c10d/exception.h
+++ b/torch/csrc/distributed/c10d/exception.h
@@ -1,3 +1,4 @@
+// @allow-raw-throw
 // Copyright (c) Facebook, Inc. and its affiliates.
 // All rights reserved.
 //
diff --git a/torch/csrc/distributed/c10d/init.cpp b/torch/csrc/distributed/c10d/init.cpp
index b70c36d5b8a6c..2815676a17543 100644
--- a/torch/csrc/distributed/c10d/init.cpp
+++ b/torch/csrc/distributed/c10d/init.cpp
@@ -52,7 +52,7 @@
 #include <torch/csrc/distributed/c10d/symm_mem/SymmetricMemory.hpp>
 
 #ifdef USE_NVSHMEM
-#include <torch/csrc/distributed/c10d/symm_mem/nvshmem_extension.cuh>
+#include <torch/csrc/distributed/c10d/symm_mem/nvshmem_extension.hpp>
 #endif
 
 #include <torch/csrc/distributed/c10d/comm.hpp>
@@ -209,7 +209,7 @@ class PythonStore : public ::c10d::Store {
   using ::c10d::Store::Store;
 
   // Note: this function manually calls the Python-side overload
-  // for this function instead of using the PYBIND11_OVERLOAD_XYZ
+  // for this function instead of using the PYBIND11_OVERRIDE_XYZ
   // macros. This is done so that we can call the Python-side
   // function with a std::string instead of a std::vector<uint8_t>.
   void set(const std::string& key, const std::vector<uint8_t>& value) override {
@@ -222,7 +222,7 @@ class PythonStore : public ::c10d::Store {
   }
 
   // Note: this function manually calls the Python-side overload
-  // for this function instead of using the PYBIND11_OVERLOAD_XYZ
+  // for this function instead of using the PYBIND11_OVERRIDE_XYZ
   // macros. This is done so that the Python-side function can
   // return a py::bytes instead of a std::vector<uint8_t>.
   std::vector<uint8_t> get(const std::string& key) override {
@@ -239,7 +239,7 @@ class PythonStore : public ::c10d::Store {
   }
 
   // Note: this function manually calls the Python-side overload
-  // for this function instead of using the PYBIND11_OVERLOAD_XYZ
+  // for this function instead of using the PYBIND11_OVERRIDE_XYZ
   // macros. This is done so that the Python-side function can
   // return a py::bytes instead of a std::vector<uint8_t>.
   std::vector<uint8_t> compareSet(
@@ -260,37 +260,37 @@ class PythonStore : public ::c10d::Store {
   }
 
   int64_t add(const std::string& key, int64_t value) override {
-    PYBIND11_OVERLOAD_PURE(int64_t, ::c10d::Store, add, key, value);
+    PYBIND11_OVERRIDE_PURE(int64_t, ::c10d::Store, add, key, value);
   }
 
   int64_t getNumKeys() override {
-    PYBIND11_OVERLOAD_PURE(int64_t, ::c10d::Store, getNumKeys);
+    PYBIND11_OVERRIDE_PURE(int64_t, ::c10d::Store, getNumKeys);
   }
 
   bool deleteKey(const std::string& key) override {
-    PYBIND11_OVERLOAD_PURE(bool, ::c10d::Store, deleteKey, key);
+    PYBIND11_OVERRIDE_PURE(bool, ::c10d::Store, deleteKey, key);
   }
 
   bool check(const std::vector<std::string>& keys) override {
-    PYBIND11_OVERLOAD_PURE(bool, ::c10d::Store, check, keys);
+    PYBIND11_OVERRIDE_PURE(bool, ::c10d::Store, check, keys);
   }
 
   void wait(const std::vector<std::string>& keys) override {
-    PYBIND11_OVERLOAD_PURE(void, ::c10d::Store, wait, keys);
+    PYBIND11_OVERRIDE_PURE(void, ::c10d::Store, wait, keys);
   }
 
   void wait(
       const std::vector<std::string>& keys,
       const std::chrono::milliseconds& timeout) override {
-    PYBIND11_OVERLOAD_PURE(void, ::c10d::Store, wait, keys, timeout);
+    PYBIND11_OVERRIDE_PURE(void, ::c10d::Store, wait, keys, timeout);
   }
 
   c10::intrusive_ptr<Store> clone() override {
-    PYBIND11_OVERLOAD_PURE(c10::intrusive_ptr<Store>, ::c10d::Store, clone);
+    PYBIND11_OVERRIDE_PURE(c10::intrusive_ptr<Store>, ::c10d::Store, clone);
   }
 
   // Note: this function manually calls the Python-side overload
-  // for this function instead of using the PYBIND11_OVERLOAD_XYZ
+  // for this function instead of using the PYBIND11_OVERRIDE_XYZ
   // macros. This is done so that we can call the Python-side
   // function with a std::string instead of a std::vector<uint8_t>.
   void append(const std::string& key, const std::vector<uint8_t>& value)
@@ -339,7 +339,7 @@ class PythonStore : public ::c10d::Store {
   }
 
   bool hasExtendedApi() const override {
-    PYBIND11_OVERLOAD_NAME(
+    PYBIND11_OVERRIDE_NAME(
         bool, ::c10d::Store, "has_extended_api", hasExtendedApi);
   }
 };
@@ -574,7 +574,8 @@ An enum-like class for built-in communication hooks: ``ALLREDUCE`` and ``FP16_CO
                  int64_t first_bucket_bytes_cap,
                  bool skip_all_reduce_unused_params,
                  bool use_python_reducer,
-                 std::vector<int64_t> bucket_bytes_cap_list) {
+                 std::vector<int64_t> bucket_bytes_cap_list,
+                 bool batched_grad_copy) {
                 // gil_scoped_release is not safe as a call_guard in init.
                 // https://github.com/pybind/pybind11/issues/5473
                 py::gil_scoped_release nogil{};
@@ -591,7 +592,8 @@ An enum-like class for built-in communication hooks: ``ALLREDUCE`` and ``FP16_CO
                     first_bucket_bytes_cap,
                     skip_all_reduce_unused_params,
                     use_python_reducer,
-                    std::move(bucket_bytes_cap_list));
+                    std::move(bucket_bytes_cap_list),
+                    batched_grad_copy);
               }),
           py::arg("params"),
           py::arg("bucket_indices"),
@@ -606,7 +608,8 @@ An enum-like class for built-in communication hooks: ``ALLREDUCE`` and ``FP16_CO
           py::arg("first_bucket_bytes_cap") = ::c10d::kDefaultFirstBucketBytes,
           py::arg("skip_all_reduce_unused_params") = false,
           py::arg("use_python_reducer") = false,
-          py::arg("bucket_bytes_cap_list") = std::vector<int64_t>())
+          py::arg("bucket_bytes_cap_list") = std::vector<int64_t>(),
+          py::arg("batched_grad_copy") = false)
       .def(
           "prepare_for_forward",
           &::c10d::Reducer::prepare_for_forward,
@@ -887,7 +890,7 @@ This class does not support ``__members__`` property.)");
             return ::c10d::ReduceOp(self);
           })
       .def(py::pickle(
-          [](const ::c10d::ReduceOp& r) {
+          [](const ::c10d::ReduceOp& r) -> py::tuple {
             // __getstate__
             if (r.op_ != ::c10d::ReduceOp::RedOpType::PREMUL_SUM) {
               return py::make_tuple(r.op_, py::none());
@@ -1140,6 +1143,10 @@ This class does not support ``__members__`` property.)");
           &::c10d::symmetric_memory::has_multicast_support)
       .def_static("set_backend", &::c10d::symmetric_memory::set_backend)
       .def_static("get_backend", &::c10d::symmetric_memory::get_backend)
+      .def_static(
+          "is_symm_mem_tensor",
+          &::c10d::symmetric_memory::is_symm_mem_tensor,
+          py::arg("tensor"))
       .def_property_static(
           "signal_pad_size",
           [](py::object /* self */) {
@@ -1685,7 +1692,41 @@ Returns a list of all keys in the store.
           .def(
               "has_extended_api",
               &::c10d::Store::hasExtendedApi,
-              R"(Returns true if the store supports extended operations.)");
+              R"(Returns true if the store supports extended operations.)")
+          .def(
+              "barrier",
+              [](::c10d::Store& store,
+                 const std::string& key,
+                 int64_t world_size,
+                 const std::optional<std::chrono::milliseconds>& timeout) {
+                if (timeout.has_value()) {
+                  store.barrier(key, world_size, *timeout);
+                } else {
+                  store.barrier(key, world_size);
+                }
+              },
+              py::call_guard<py::gil_scoped_release>(),
+              py::arg("key"),
+              py::arg("world_size"),
+              py::arg("timeout") = py::none(),
+              R"(
+Barrier operation that blocks until ``world_size`` workers have called it
+with the same ``key``. If ``timeout`` is not specified, the store's default
+timeout is used.
+
+Arguments:
+    key (str): The unique key for this barrier instance.
+    world_size (int): The number of workers that must call barrier before it unblocks.
+    timeout (timedelta, optional): Time to wait before throwing an exception. Defaults to store timeout.
+
+Example::
+    >>> import torch.distributed as dist
+    >>> from datetime import timedelta
+    >>> store = dist.TCPStore("127.0.0.1", 0, 1, True, timedelta(seconds=30))
+    >>> # This will return immediately since world_size=1
+    >>> store.barrier("my_barrier", 1)
+    >>> store.barrier("my_barrier2", 1, timedelta(seconds=10))
+)");
 
   intrusive_ptr_class_<::c10d::FileStore>(
       module,
@@ -2073,9 +2114,17 @@ communication mechanism.
           py::arg("rank"),
           py::arg("world_size"));
 
+  // Use OpaqueBase as the metaclass to allow isinstance(fake_obj, ProcessGroup)
+  // to work.
+  py::object opaque_base_module = py::module_::import("torch._opaque_base");
+  py::object opaque_base = opaque_base_module.attr("OpaqueBaseMeta");
+
   auto processGroup =
       intrusive_ptr_no_gil_destructor_trampoline_class_<
-          ::c10d::ProcessGroup, ::c10d::PyProcessGroup>(module, "ProcessGroup",
+          ::c10d::ProcessGroup, ::c10d::PyProcessGroup>(
+          module,
+          "ProcessGroup",
+          py::metaclass(opaque_base),
           R"(A ProcessGroup is a communication primitive that allows for
           collective operations across a group of processes.
 
@@ -2716,6 +2765,13 @@ The hook must have the following signature:
   module.def("_set_process_group", &::c10d::setProcessGroup);
   module.def("_current_process_group", &::c10d::currentProcessGroup);
 
+  // Thread local comm profiling name
+  module.def(
+      "_set_comm_profiling_name",
+      &::c10d::set_comm_profiling_name,
+      py::arg("name"));
+  module.def("_get_comm_profiling_name", &::c10d::get_comm_profiling_name);
+
   py::enum_<::c10d::ProcessGroup::BackendType>(
       processGroup,
       "BackendType",
@@ -3114,7 +3170,15 @@ The hook must have the following signature:
               py::arg("device"),
               py::call_guard<py::gil_scoped_release>())
           .def_property_readonly(
-              "mem_allocator", &::c10d::Backend::getMemAllocator);
+              "mem_allocator", &::c10d::Backend::getMemAllocator)
+          .def("suspend", &::c10d::Backend::suspend)
+          .def("resume", &::c10d::Backend::resume)
+          .def("memory_stats", &::c10d::Backend::getMemoryStats, R"(
+            Get the memory statistics of the backend.
+
+            Returns:
+              A dictionary containing the memory statistics.
+            )");
 
   // base Backend::Options binding
   // TODO: Maybe we can consider how to merge this with
@@ -3673,6 +3737,9 @@ Example::
       .value("_REDUCE_SCATTER_BASE", ::c10d::OpType::_REDUCE_SCATTER_BASE)
       .value("COALESCED", ::c10d::OpType::COALESCED)
       .value("_ALLREDUCE_SPARSE", ::c10d::OpType::_ALLREDUCE_SPARSE)
+      .value(
+          "REDUCE_SCATTER_TENSOR_COALESCED",
+          ::c10d::OpType::REDUCE_SCATTER_TENSOR_COALESCED)
       .value("UNKNOWN", ::c10d::OpType::UNKNOWN);
 
   py::enum_<::c10d::WorkResult>(module, "WorkResult")
diff --git a/torch/csrc/distributed/c10d/reducer.cpp b/torch/csrc/distributed/c10d/reducer.cpp
index 4007e08227e96..6b56e09639214 100644
--- a/torch/csrc/distributed/c10d/reducer.cpp
+++ b/torch/csrc/distributed/c10d/reducer.cpp
@@ -95,7 +95,8 @@ Reducer::Reducer(
     int64_t first_bucket_bytes_cap,
     bool skip_all_reduce_unused_params,
     bool use_python_reducer,
-    std::vector<int64_t> bucket_bytes_cap_list)
+    std::vector<int64_t> bucket_bytes_cap_list,
+    bool batched_grad_copy)
     : params_(std::move(params)),
       process_group_(std::move(process_group)),
       expect_sparse_gradients_(std::move(expect_sparse_gradients)),
@@ -105,6 +106,7 @@ Reducer::Reducer(
       has_marked_unused_parameters_(false),
       find_unused_parameters_(find_unused_parameters),
       gradient_as_bucket_view_(gradient_as_bucket_view),
+      batched_grad_copy_(batched_grad_copy),
       local_used_map_reduced_(false),
       num_iterations_(0),
       num_bwd_calls_(0),
@@ -372,51 +374,66 @@ void Reducer::mark_variable_ready_dense(size_t variable_index) {
       // to bucket_view. If grad has already been set as views of buckets in
       // previous iterations, no copy is needed.
       if (!grad.is_alias_of(bucket_view)) {
-        if (comm_hook_ == nullptr) {
-          auto wrapped = at::native::wrapped_scalar_tensor(1. / div_factor_);
-          if (!grad.requires_grad()) {
-            // Divides while copying into the bucket view to save one scan over
-            // all the input parameters.
-            RECORD_FUNCTION(
-                "torch::distributed::reducer::mul_out",
-                std::vector<c10::IValue>({bucket_view}))
-            at::mul_out(bucket_view, grad, wrapped);
+        if (batched_grad_copy_ && !grad.requires_grad()) {
+          // Defer the copy — will be batched with _foreach_copy_ + flat div_
+          // when bucket.pending == 0.
+          bucket.deferred_copy_indices.push_back(
+              bucket_index.intra_bucket_index);
+        } else {
+          if (comm_hook_ == nullptr) {
+            auto wrapped = at::native::wrapped_scalar_tensor(1. / div_factor_);
+            if (!grad.requires_grad()) {
+              // Divides while copying into the bucket view to save one scan
+              // over all the input parameters.
+              RECORD_FUNCTION(
+                  "torch::distributed::reducer::mul_out",
+                  std::vector<c10::IValue>({bucket_view}))
+              at::mul_out(bucket_view, grad, wrapped);
+            } else {
+              // If DDP is running with create_graph=True, gradients
+              // require_grad themselves in order to compute higher order
+              // derivatives. However, DDP will not sync up these gradients
+              // currently (see
+              // https://github.com/pytorch/pytorch/issues/63812).
+              C10_LOG_EVERY_N(WARNING, 1000)
+                  << "Using DistributedDataParallel with create_graph=True "
+                  << " is not well-supported. The higher-order gradient will "
+                  << " not be synchronized across ranks, and backpropagation "
+                  << " through all_reduce operations will not occur. If you require "
+                  << " DDP to work with higher-order gradients for your use case, "
+                  << " please ping https://github.com/pytorch/pytorch/issues/63929";
+              if (batched_grad_copy_) {
+                C10_LOG_EVERY_N(WARNING, 1000)
+                    << "batched_grad_copy is incompatible with "
+                    << "create_graph=True and has been bypassed.";
+              }
+              auto div_result = at::mul(grad, wrapped);
+              RECORD_FUNCTION(
+                  "torch::distributed::reducer::copy_",
+                  std::vector<c10::IValue>({bucket_view}))
+              bucket_view.copy_(div_result);
+            }
           } else {
-            // If DDP is running with create_graph=True, gradients require_grad
-            // themselves in order to compute higher order derivatives. However,
-            // DDP will not sync up these gradients currently (see
-            // https://github.com/pytorch/pytorch/issues/63812).
-            C10_LOG_EVERY_N(WARNING, 1000)
-                << "Using DistributedDataParallel with create_graph=True "
-                << " is not well-supported. The higher-order gradient will "
-                << " not be synchronized across ranks, and backpropagation "
-                << " through all_reduce operations will not occur. If you require "
-                << " DDP to work with higher-order gradients for your use case, "
-                << " please ping https://github.com/pytorch/pytorch/issues/63929";
-            auto div_result = at::mul(grad, wrapped);
             RECORD_FUNCTION(
                 "torch::distributed::reducer::copy_",
                 std::vector<c10::IValue>({bucket_view}))
-            bucket_view.copy_(div_result);
+            bucket_view.copy_(grad);
           }
-        } else {
-          RECORD_FUNCTION(
-              "torch::distributed::reducer::copy_",
-              std::vector<c10::IValue>({bucket_view}))
-          bucket_view.copy_(grad);
-        }
 
-        if (gradient_as_bucket_view_) {
-          // Let grad point to bucket_view buffer.
-          grad = bucket_view;
-          // The grad is modified and need to be written back.
-          return true;
+          if (gradient_as_bucket_view_) {
+            grad = bucket_view;
+            return true;
+          }
         }
       } else {
         // If grad and bucket view point to the same storage, no need to copy.
-        if (comm_hook_ == nullptr) {
-          bucket_view.div_(div_factor_);
+        if (!batched_grad_copy_) {
+          if (comm_hook_ == nullptr) {
+            bucket_view.div_(div_factor_);
+          }
         }
+        // When batched_grad_copy_ is enabled, div_ is deferred to a single
+        // flat bucket div_ in flush_deferred_copies.
       }
     } else {
       // Gradient is undefined. When find_unused_parameters=True, ensure it is
@@ -433,7 +450,6 @@ void Reducer::mark_variable_ready_dense(size_t variable_index) {
       }
       bucket_view.zero_();
     }
-    // The grad is not modified and doesn't need to be written back.
     return false;
   });
 }
@@ -630,7 +646,11 @@ void Reducer::delay_all_reduce() {
   }
 
   // launch all reduces for all buckets
-  for (auto& bucket : buckets_) {
+  for (const auto bucket_index : c10::irange(buckets_.size())) {
+    auto& bucket = buckets_[bucket_index];
+    if (batched_grad_copy_) {
+      flush_deferred_copies(bucket, bucket_index);
+    }
     all_reduce_bucket(bucket);
   }
 
@@ -907,6 +927,11 @@ void Reducer::mark_variable_ready(size_t variable_index) {
 
   // Check if this was the final gradient for this bucket.
   if (--bucket.pending == 0) {
+    // When batched_grad_copy_ is enabled, flush deferred copies and perform
+    // a single div on the flat bucket tensor instead of per-variable ops.
+    if (batched_grad_copy_) {
+      flush_deferred_copies(bucket, bucket_index.bucket_index);
+    }
     mark_bucket_ready(bucket_index.bucket_index);
   }
 
@@ -1431,6 +1456,7 @@ void Reducer::reset_bucket_counting() {
 
   for (auto& bucket : buckets_) {
     bucket.pending = bucket.variables.size();
+    bucket.deferred_copy_indices.clear();
   }
 
   if (static_graph_) {
@@ -1812,6 +1838,53 @@ void Reducer::runGradCallbackForVariable(
 #endif
 }
 
+void Reducer::flush_deferred_copies(Bucket& bucket, size_t bucket_index) {
+  // Sparse gradients are already divided in mark_variable_ready_sparse and
+  // communicated independently — skip to avoid double division.
+  if (bucket.expect_sparse_gradient) {
+    return;
+  }
+  if (!bucket.deferred_copy_indices.empty()) {
+    std::vector<at::Tensor> dsts;
+    std::vector<at::Tensor> srcs;
+    dsts.reserve(bucket.deferred_copy_indices.size());
+    srcs.reserve(bucket.deferred_copy_indices.size());
+    for (auto idx : bucket.deferred_copy_indices) {
+      auto grad = bucket.variables[idx].grad();
+      TORCH_INTERNAL_ASSERT(
+          grad.defined(),
+          "Gradient became undefined between defer and flush for variable ",
+          idx,
+          " in bucket ",
+          bucket_index,
+          ". This indicates a bug — gradients should not be modified during backward.");
+      dsts.push_back(bucket.bucket_views_in[idx]);
+      srcs.push_back(grad);
+    }
+    at::_foreach_copy_(dsts, srcs);
+
+    // Re-alias grads to bucket views if gradient_as_bucket_view
+    if (gradient_as_bucket_view_) {
+      for (auto idx : bucket.deferred_copy_indices) {
+        auto& variable = bucket.variables[idx];
+        auto& bucket_view = bucket.bucket_views_in[idx];
+        runGradCallbackForVariable(variable, [&](auto& grad) {
+          grad = bucket_view;
+          return true;
+        });
+      }
+    }
+    bucket.deferred_copy_indices.clear();
+  }
+  // Single div on the entire flat bucket tensor.
+  // This also divides regions zeroed for undefined gradients, which is a no-op
+  // (0 / div_factor_ == 0) but avoids the complexity of tracking whether any
+  // variable in the bucket had a defined grad.
+  if (comm_hook_ == nullptr) {
+    bucket.gradients.div_(div_factor_);
+  }
+}
+
 #ifndef _WIN32
 void Reducer::RpcContext::set(ContextPtr&& new_context_ptr) {
   // We should set 'new_context_ptr' even if it's nullptr. That means the
diff --git a/torch/csrc/distributed/c10d/reducer.hpp b/torch/csrc/distributed/c10d/reducer.hpp
index 3badec59594fa..73730ead6de03 100644
--- a/torch/csrc/distributed/c10d/reducer.hpp
+++ b/torch/csrc/distributed/c10d/reducer.hpp
@@ -60,7 +60,8 @@ class TORCH_API Reducer {
       int64_t first_bucket_bytes_cap,
       bool skip_all_reduce_unused_params,
       bool use_python_reducer,
-      std::vector<int64_t> bucket_bytes_cap_list);
+      std::vector<int64_t> bucket_bytes_cap_list,
+      bool batched_grad_copy = false);
 
   ~Reducer() noexcept(false);
 
@@ -236,6 +237,8 @@ class TORCH_API Reducer {
   // NOLINTNEXTLINE(cppcoreguidelines-non-private-member-variables-in-classes)
   const bool gradient_as_bucket_view_;
   // NOLINTNEXTLINE(cppcoreguidelines-non-private-member-variables-in-classes)
+  const bool batched_grad_copy_;
+  // NOLINTNEXTLINE(cppcoreguidelines-non-private-member-variables-in-classes)
   std::vector<size_t> unused_parameters_;
   // Previous iteration's unused params, used for checking if unused parameters
   // change between iterations. Only filled during the first backwards call.
@@ -312,6 +315,12 @@ class TORCH_API Reducer {
 #endif
   void runGradCallbackForVariable(at::Tensor& variable, const GradCallback& cb);
 
+  // Flushes deferred grad-to-bucket copies for a single bucket when
+  // batched_grad_copy_ is enabled. Called from mark_variable_ready (when
+  // bucket.pending == 0) and from delay_all_reduce (after all variables
+  // are marked ready).
+  void flush_deferred_copies(Bucket& bucket, size_t bucket_index);
+
   // This function is called inside `initialize_buckets()`. It initializes both
   // `bucket_views_in` and `bucket_views_out` with views for each variable's
   // gradient into the bucket's flattened `gradients` tensor. Views serve as
@@ -402,6 +411,11 @@ class TORCH_API Reducer {
     // done on different CUDA streams. We record an event for every copy
     // so that we can synchronize with them prior to kicking off the reduction.
     // std::vector<at::cuda::CUDAEvent> events;
+
+    // Intra-bucket indices of variables whose grad-to-bucket copies are
+    // deferred for batching. Flushed as _foreach_copy_ + flat div_ when
+    // pending == 0. Only used when batched_grad_copy is enabled.
+    std::vector<size_t> deferred_copy_indices;
   };
 
   std::vector<Bucket> buckets_;
diff --git a/torch/csrc/distributed/c10d/symm_mem/CUDASymmetricMemory-inl.h b/torch/csrc/distributed/c10d/symm_mem/CUDASymmetricMemory-inl.cuh
similarity index 100%
rename from torch/csrc/distributed/c10d/symm_mem/CUDASymmetricMemory-inl.h
rename to torch/csrc/distributed/c10d/symm_mem/CUDASymmetricMemory-inl.cuh
diff --git a/torch/csrc/distributed/c10d/symm_mem/CUDASymmetricMemory.cu b/torch/csrc/distributed/c10d/symm_mem/CUDASymmetricMemory.cu
index c15ebbcd457b3..576e59dad2f85 100644
--- a/torch/csrc/distributed/c10d/symm_mem/CUDASymmetricMemory.cu
+++ b/torch/csrc/distributed/c10d/symm_mem/CUDASymmetricMemory.cu
@@ -1,6 +1,7 @@
 #include <torch/csrc/distributed/c10d/GroupRegistry.hpp>
+#include <torch/csrc/distributed/c10d/ParamCommsUtils.hpp>
 #include <torch/csrc/distributed/c10d/cuda/utils.hpp>
-#include <torch/csrc/distributed/c10d/symm_mem/CUDASymmetricMemory-inl.h>
+#include <torch/csrc/distributed/c10d/symm_mem/CUDASymmetricMemory-inl.cuh>
 #include <torch/csrc/distributed/c10d/symm_mem/CUDASymmetricMemory.hpp>
 #include <torch/csrc/distributed/c10d/symm_mem/CUDASymmetricMemoryUtils.hpp>
 
@@ -9,6 +10,7 @@
 #include <ATen/cuda/PeerToPeerAccess.h>
 #include <c10/cuda/CUDACachingAllocator.h>
 #include <c10/cuda/CUDAGuard.h>
+#include <c10/util/env.h>
 #include <c10/util/error.h>
 
 #include <sys/socket.h>
@@ -79,7 +81,8 @@ CUDAPeerAllocInfo::CUDAPeerAllocInfo(
     size_t buffer_size,
     int local_device_idx,
     int rank,
-    int world_size)
+    int world_size,
+    std::string group_name)
     : alloc_refs_(std::move(alloc_refs)),
       buffers_(std::move(buffers)),
       signal_pads_(std::move(signal_pads)),
@@ -88,7 +91,8 @@ CUDAPeerAllocInfo::CUDAPeerAllocInfo(
       buffer_size_(buffer_size),
       local_device_idx_(local_device_idx),
       rank_(rank),
-      world_size_(world_size) {
+      world_size_(world_size),
+      group_name_(std::move(group_name)) {
   const size_t arr_size = sizeof(void*) * world_size_;
   buffers_dev_ = reinterpret_cast<void**>(
       c10::cuda::CUDACachingAllocator::raw_alloc(arr_size));
@@ -210,7 +214,21 @@ static __global__ void barrier_kernel(
 
 void CUDASymmetricMemory::barrier(int channel, size_t timeout_ms) {
   check_channel(channel, world_size_);
-  c10::cuda::CUDAGuard guard(local_device_idx_);
+  auto pg = c10d::resolve_process_group(pai_->group_name_);
+  RECORD_PARAM_COMMS(
+      static_cast<int64_t>(0),
+      std::make_tuple(pg->getGroupName(), pg->getGroupDesc()),
+      rank_,
+      "symm_mem::barrier",
+      0,
+      0,
+      at::kByte,
+      std::vector<int64_t>(),
+      std::vector<int64_t>(),
+      -1,
+      -1,
+      world_size_);
+  c10::cuda::CUDAGuard device_guard(local_device_idx_);
   barrier_kernel<<<
       1,
       max(at::cuda::warp_size(), world_size_),
@@ -252,7 +270,21 @@ void CUDASymmetricMemory::put_signal(
     int channel,
     size_t timeout_ms) {
   check_channel(channel, world_size_);
-  c10::cuda::CUDAGuard guard(local_device_idx_);
+  auto pg = c10d::resolve_process_group(pai_->group_name_);
+  RECORD_PARAM_COMMS(
+      static_cast<int64_t>(0),
+      std::make_tuple(pg->getGroupName(), pg->getGroupDesc()),
+      rank_,
+      "symm_mem::put_signal",
+      0,
+      0,
+      at::kByte,
+      std::vector<int64_t>(),
+      std::vector<int64_t>(),
+      -1,
+      -1,
+      world_size_);
+  c10::cuda::CUDAGuard device_guard(local_device_idx_);
   put_signal_kernel<<<
       1,
       at::cuda::warp_size(),
@@ -300,7 +332,21 @@ void CUDASymmetricMemory::wait_signal(
     int channel,
     size_t timeout_ms) {
   check_channel(channel, world_size_);
-  c10::cuda::CUDAGuard guard(local_device_idx_);
+  auto pg = c10d::resolve_process_group(pai_->group_name_);
+  RECORD_PARAM_COMMS(
+      static_cast<int64_t>(0),
+      std::make_tuple(pg->getGroupName(), pg->getGroupDesc()),
+      rank_,
+      "symm_mem::wait_signal",
+      0,
+      0,
+      at::kByte,
+      std::vector<int64_t>(),
+      std::vector<int64_t>(),
+      -1,
+      -1,
+      world_size_);
+  c10::cuda::CUDAGuard device_guard(local_device_idx_);
   wait_signal_kernel<<<
       1,
       at::cuda::warp_size(),
@@ -453,9 +499,25 @@ struct RendezvousRequest {
   size_t buffer_size;
   size_t signal_pad_offset;
   bool has_multicast_support;
+  int clique_id;
   char hostname[HOST_NAME_MAX + 1];
 };
 
+static std::string import_err_msg(
+    int rank,
+    int peer,
+    const std::vector<RendezvousRequest>& reqs) {
+  std::ostringstream oss;
+  oss << ". Rank " << rank << " (host: " << reqs[rank].hostname
+      << ", device: " << reqs[rank].device_idx << ", fabric_info: {"
+      << at::cuda::get_nvml_fabric_info(reqs[rank].device_idx)
+      << "}) failed to import memory from rank " << peer
+      << " (host: " << reqs[peer].hostname
+      << ", device: " << reqs[peer].device_idx << ", NCCL_MNNVL_CLIQUE_ID: "
+      << c10::utils::get_env("NCCL_MNNVL_CLIQUE_ID").value_or("unset") << ").";
+  return oss.str();
+}
+
 void validate_rendezvous_requests(
     const std::vector<RendezvousRequest>& reqs,
     int world_size) {
@@ -466,7 +528,8 @@ void validate_rendezvous_requests(
   // Use (hostname, device_idx) pair to uniquely identify each allocation.
   std::set<std::pair<std::string, int>> device_host_pairs;
   for (auto req : reqs) {
-    device_host_pairs.insert(std::make_pair(std::string(req.hostname), req.device_idx));
+    device_host_pairs.insert(
+        std::make_pair(std::string(req.hostname), req.device_idx));
   }
   if (!allow_overlapping_devices() &&
       device_host_pairs.size() < (size_t)world_size) {
@@ -484,6 +547,35 @@ void validate_rendezvous_requests(
   }
 }
 
+// All ranks must be in the same NVLink domain (same clique_id). Detect
+// mismatches early before the import fails with an opaque CUDA error.
+static void validate_nvlink_fabric_support(
+    const std::vector<RendezvousRequest>& reqs,
+    int world_size) {
+  std::unordered_set<int> clique_ids;
+  for (const auto& req : reqs) {
+    if (req.clique_id >= 0) {
+      clique_ids.insert(req.clique_id);
+    }
+  }
+  if (clique_ids.size() > 1) {
+    std::ostringstream oss;
+    oss << "CUDASymmetricMemory::rendezvous: "
+        << "ranks have mismatched NVLink clique_ids. "
+        << "All ranks using fabric handles must be in the same NVLink domain. "
+        << "Per-rank info: ";
+    for (int r = 0; r < world_size; ++r) {
+      if (r > 0) {
+        oss << ", ";
+      }
+      oss << "rank " << r << " (host: " << reqs[r].hostname
+          << ", device: " << reqs[r].device_idx
+          << ", clique_id: " << reqs[r].clique_id << ")";
+    }
+    TORCH_CHECK(false, oss.str());
+  }
+}
+
 static bool check_group_multicast_support(
     const std::vector<RendezvousRequest>& reqs) {
   std::vector<size_t> ranks_with_multicast_support;
@@ -678,11 +770,13 @@ c10::intrusive_ptr<CUDAPeerAllocInfo> make_peer_alloc_info(
       .block_size = block->block_size,
       .buffer_size = block->buffer_size,
       .signal_pad_offset = block->signal_pad_offset,
-      .has_multicast_support = device_has_multicast_support(block->device_idx)};
+      .has_multicast_support = device_has_multicast_support(block->device_idx),
+      .clique_id = at::cuda::get_fabric_clique_id(block->device_idx)};
 
   // Populate hostname field for host identification
   gethostname(local_req.hostname, sizeof(local_req.hostname));
   auto reqs = storeExchange.all_gather(store, rank, world_size, local_req);
+  validate_nvlink_fabric_support(reqs, world_size);
   validate_rendezvous_requests(reqs, world_size);
 
   std::vector<int> pids(world_size);
@@ -715,15 +809,19 @@ c10::intrusive_ptr<CUDAPeerAllocInfo> make_peer_alloc_info(
     // note how in one case it's directly imported_handles[r] and in another
     // &(imported_handles[r]) so can't do with just type definitions
     if constexpr (!use_fabric_handle) {
-      C10_CUDA_DRIVER_CHECK(driver_api->cuMemImportFromShareableHandle_(
-          &handles[r],
-          (void*)(uintptr_t)imported_handles[r],
-          CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR));
+      C10_CUDA_DRIVER_CHECK_MSG(
+          driver_api->cuMemImportFromShareableHandle_(
+              &handles[r],
+              (void*)(uintptr_t)imported_handles[r],
+              CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR),
+          import_err_msg(rank, r, reqs));
     } else {
-      C10_CUDA_DRIVER_CHECK(driver_api->cuMemImportFromShareableHandle_(
-          &handles[r],
-          (void*)&(imported_handles[r]),
-          CU_MEM_HANDLE_TYPE_FABRIC));
+      C10_CUDA_DRIVER_CHECK_MSG(
+          driver_api->cuMemImportFromShareableHandle_(
+              &handles[r],
+              (void*)&(imported_handles[r]),
+              CU_MEM_HANDLE_TYPE_FABRIC),
+          import_err_msg(rank, r, reqs));
     }
 #elif defined(USE_ROCM)
     C10_CUDA_CHECK(hipMemImportFromShareableHandle(
@@ -785,7 +883,8 @@ c10::intrusive_ptr<CUDAPeerAllocInfo> make_peer_alloc_info(
       block->buffer_size,
       block->device_idx,
       rank,
-      world_size);
+      world_size,
+      group_name);
 
   return pai;
 }
@@ -888,6 +987,10 @@ c10::intrusive_ptr<Block> CUDASymmetricMemoryAllocator::find_block_covering(void
   return alloc_it->second;
 }
 
+bool CUDASymmetricMemoryAllocator::has_allocation(void* ptr) {
+  return find_block(ptr) != nullptr;
+}
+
 struct RegisterCUDASymmetricMemoryAllocator {
   RegisterCUDASymmetricMemoryAllocator() {
     auto allocator = c10::make_intrusive<CUDASymmetricMemoryAllocator>();
diff --git a/torch/csrc/distributed/c10d/symm_mem/CUDASymmetricMemory.hpp b/torch/csrc/distributed/c10d/symm_mem/CUDASymmetricMemory.hpp
index 0ed09c6cc5dbf..5636bca810f6a 100644
--- a/torch/csrc/distributed/c10d/symm_mem/CUDASymmetricMemory.hpp
+++ b/torch/csrc/distributed/c10d/symm_mem/CUDASymmetricMemory.hpp
@@ -84,7 +84,8 @@ class CUDAPeerAllocInfo : public c10::intrusive_ptr_target {
       size_t buffer_size,
       int local_device_idx,
       int rank,
-      int world_size);
+      int world_size,
+      std::string group_name);
 
  private:
   std::vector<c10::intrusive_ptr<AllocationRef>> alloc_refs_;
@@ -98,6 +99,7 @@ class CUDAPeerAllocInfo : public c10::intrusive_ptr_target {
   int world_size_;
   void** buffers_dev_;
   void** signal_pads_dev_;
+  std::string group_name_;
 
   friend class CUDASymmetricMemory;
 };
@@ -135,6 +137,7 @@ class CUDASymmetricMemoryAllocator : public SymmetricMemoryAllocator {
       void* ptr,
       const std::optional<std::string>& group_name) override;
   bool has_multicast_support(int device_idx) override;
+  bool has_allocation(void* ptr) override;
   c10::DeviceType supported_device_type() override;
   std::string name() override;
 
diff --git a/torch/csrc/distributed/c10d/symm_mem/CUDASymmetricMemoryOps.cu b/torch/csrc/distributed/c10d/symm_mem/CUDASymmetricMemoryOps.cu
index 5c6f08f44db99..da8c1611f13a0 100644
--- a/torch/csrc/distributed/c10d/symm_mem/CUDASymmetricMemoryOps.cu
+++ b/torch/csrc/distributed/c10d/symm_mem/CUDASymmetricMemoryOps.cu
@@ -17,7 +17,8 @@
 
 #include <torch/csrc/distributed/c10d/cuda/AsyncMM.cuh>
 #include <torch/csrc/distributed/c10d/GroupRegistry.hpp>
-#include <torch/csrc/distributed/c10d/symm_mem/CUDASymmetricMemory-inl.h>
+#include <torch/csrc/distributed/c10d/ParamCommsUtils.hpp>
+#include <torch/csrc/distributed/c10d/symm_mem/CUDASymmetricMemory-inl.cuh>
 #include <torch/csrc/distributed/c10d/symm_mem/CUDASymmetricMemory.hpp>
 
 #if defined(USE_ROCM) || (defined(CUDART_VERSION) && CUDART_VERSION >= 12030)
@@ -163,6 +164,20 @@ at::Tensor multimem_all_reduce_(
     const at::Tensor& input,
     std::string reduce_op,
     std::string group_name) {
+  auto pg = c10d::resolve_process_group(group_name);
+  RECORD_PARAM_COMMS(
+      static_cast<int64_t>(0),
+      std::make_tuple(pg->getGroupName(), pg->getGroupDesc()),
+      pg->getRank(),
+      "symm_mem::multimem_all_reduce",
+      input.numel(),
+      input.numel(),
+      input.scalar_type(),
+      std::vector<int64_t>(),
+      std::vector<int64_t>(),
+      -1,
+      -1,
+      pg->getSize());
   TORCH_CHECK(
       input.is_contiguous(), "multimem_all_reduce_: input must be contiguous.");
   TORCH_CHECK(
@@ -247,6 +262,20 @@ at::Tensor multimem_one_shot_reduce_out(
     int64_t root,
     std::string group_name,
     at::Tensor out) {
+  auto pg = c10d::resolve_process_group(group_name);
+  RECORD_PARAM_COMMS(
+      static_cast<int64_t>(0),
+      std::make_tuple(pg->getGroupName(), pg->getGroupDesc()),
+      pg->getRank(),
+      "symm_mem::multimem_one_shot_reduce",
+      input.numel(),
+      out.numel(),
+      input.scalar_type(),
+      std::vector<int64_t>(),
+      std::vector<int64_t>(),
+      -1,
+      -1,
+      pg->getSize());
   TORCH_CHECK(
       input.is_contiguous(),
       "multimem_one_shot_reduce: input must be contiguous.");
@@ -361,6 +390,20 @@ at::Tensor multimem_all_gather_out(
     const at::Tensor& input,
     std::string group_name,
     at::Tensor out) {
+  auto pg = c10d::resolve_process_group(group_name);
+  RECORD_PARAM_COMMS(
+      static_cast<int64_t>(0),
+      std::make_tuple(pg->getGroupName(), pg->getGroupDesc()),
+      pg->getRank(),
+      "symm_mem::multimem_all_gather",
+      input.numel(),
+      out.numel(),
+      input.scalar_type(),
+      std::vector<int64_t>(),
+      std::vector<int64_t>(),
+      -1,
+      -1,
+      pg->getSize());
   auto symm_mem = c10d::symmetric_memory::rendezvous(out, group_name);
   TORCH_CHECK(
       symm_mem != nullptr,
@@ -475,6 +518,20 @@ at::Tensor one_shot_all_reduce_out_impl(
     std::string reduce_op,
     std::string group_name,
     at::Tensor out) {
+  auto pg = c10d::resolve_process_group(group_name);
+  RECORD_PARAM_COMMS(
+      static_cast<int64_t>(0),
+      std::make_tuple(pg->getGroupName(), pg->getGroupDesc()),
+      pg->getRank(),
+      "symm_mem::one_shot_all_reduce",
+      input.numel(),
+      out.numel(),
+      input.scalar_type(),
+      std::vector<int64_t>(),
+      std::vector<int64_t>(),
+      -1,
+      -1,
+      pg->getSize());
   TORCH_CHECK(
       input.is_contiguous(), "one_shot_all_reduce: input must be contiguous.");
   TORCH_CHECK(
@@ -732,6 +789,20 @@ at::Tensor two_shot_all_reduce_impl(
     std::optional<at::Tensor> output,
     std::string reduce_op,
     std::string group_name) {
+  auto pg = c10d::resolve_process_group(group_name);
+  RECORD_PARAM_COMMS(
+      static_cast<int64_t>(0),
+      std::make_tuple(pg->getGroupName(), pg->getGroupDesc()),
+      pg->getRank(),
+      "symm_mem::two_shot_all_reduce",
+      input.numel(),
+      input.numel(),
+      input.scalar_type(),
+      std::vector<int64_t>(),
+      std::vector<int64_t>(),
+      -1,
+      -1,
+      pg->getSize());
   TORCH_CHECK(
       input.is_contiguous(), "two_shot_all_reduce: input must be contiguous.");
   TORCH_CHECK(
@@ -856,6 +927,20 @@ at::Tensor reduce_scatter_out(
     std::string group_name,
     bool split_last_dim,
     at::Tensor output) {
+  auto pg = c10d::resolve_process_group(group_name);
+  RECORD_PARAM_COMMS(
+      static_cast<int64_t>(0),
+      std::make_tuple(pg->getGroupName(), pg->getGroupDesc()),
+      pg->getRank(),
+      "symm_mem::reduce_scatter",
+      input.numel(),
+      output.numel(),
+      input.scalar_type(),
+      std::vector<int64_t>(),
+      std::vector<int64_t>(),
+      -1,
+      -1,
+      pg->getSize());
   TORCH_CHECK(
       input.is_contiguous(), "reduce_scatter: input must be contiguous.");
   TORCH_CHECK(
diff --git a/torch/csrc/distributed/c10d/symm_mem/CUDASymmetricMemoryTypes.hpp b/torch/csrc/distributed/c10d/symm_mem/CUDASymmetricMemoryTypes.hpp
index 7c255fa283ec9..cec3bfc8d05d6 100644
--- a/torch/csrc/distributed/c10d/symm_mem/CUDASymmetricMemoryTypes.hpp
+++ b/torch/csrc/distributed/c10d/symm_mem/CUDASymmetricMemoryTypes.hpp
@@ -2,6 +2,10 @@
 
 #include <cstddef>
 #include <cstdint>
+#include <string>
+#include <utility>
+
+#include <c10/util/hash.h>
 
 #if defined(USE_ROCM)
 #include <hip/hip_runtime_api.h>
@@ -9,6 +13,14 @@
 
 namespace c10d::symmetric_memory {
 
+// Key type for the symmetric memory map. `void*` for tensor storage ptr,
+// `std::string` for group name.
+using SymmMemKey = std::pair<void*, std::string>;
+// Hash function for the symmetric memory map. c10::hash has a std::pair
+// specialization (line 323-329 of hash.h) that delegates to the tuple hasher
+// which combines hashes of each element.
+using SymmMemKeyHash = c10::hash<SymmMemKey>;
+
 // Covers NVL72
 constexpr int max_cuda_p2p_domain_size = 72;
 // Maximum number of channels
diff --git a/torch/csrc/distributed/c10d/symm_mem/CUDASymmetricMemoryUtils.cpp b/torch/csrc/distributed/c10d/symm_mem/CUDASymmetricMemoryUtils.cpp
index 98a272468a84d..5c728e0ee0658 100644
--- a/torch/csrc/distributed/c10d/symm_mem/CUDASymmetricMemoryUtils.cpp
+++ b/torch/csrc/distributed/c10d/symm_mem/CUDASymmetricMemoryUtils.cpp
@@ -251,12 +251,6 @@ void map_block(
       0,
       reinterpret_cast<hipMemGenericAllocationHandle_t>(handle),
       0ULL));
-  C10_CUDA_CHECK(hipMemMap(
-      *ptr,
-      size,
-      0,
-      reinterpret_cast<hipMemGenericAllocationHandle_t>(handle),
-      0ULL));
 
   hipMemAccessDesc desc;
   desc.location.type = hipMemLocationTypeDevice;
diff --git a/torch/csrc/distributed/c10d/symm_mem/NCCLSymmetricMemory.cu b/torch/csrc/distributed/c10d/symm_mem/NCCLSymmetricMemory.cu
index 0925e5079b3d8..02d3e35d19924 100644
--- a/torch/csrc/distributed/c10d/symm_mem/NCCLSymmetricMemory.cu
+++ b/torch/csrc/distributed/c10d/symm_mem/NCCLSymmetricMemory.cu
@@ -2,14 +2,15 @@
 
 #ifdef NCCL_HAS_SYMMEM_SUPPORT
 
+#include <algorithm>
 #include <vector_types.h>
 #include <torch/csrc/distributed/c10d/GroupRegistry.hpp>
 #include <torch/csrc/distributed/c10d/NCCLUtils.hpp>
 #include <torch/csrc/distributed/c10d/ProcessGroupNCCL.hpp>
 #include <torch/csrc/distributed/c10d/cuda/utils.hpp>
-#include <torch/csrc/distributed/c10d/symm_mem/CUDASymmetricMemory-inl.h>
-#include <torch/csrc/distributed/c10d/symm_mem/CUDASymmetricMemoryUtils.hpp>
+#include <torch/csrc/distributed/c10d/symm_mem/CUDASymmetricMemory-inl.cuh>
 #include <torch/csrc/distributed/c10d/symm_mem/CUDASymmetricMemoryTypes.hpp>
+#include <torch/csrc/distributed/c10d/symm_mem/CUDASymmetricMemoryUtils.hpp>
 #include <torch/csrc/distributed/c10d/symm_mem/NCCLSymmetricMemory.hpp>
 #include <torch/csrc/distributed/c10d/symm_mem/nccl_devcomm_manager.hpp>
 
@@ -18,6 +19,13 @@
 #include <c10/cuda/CUDACachingAllocator.h>
 #include <c10/cuda/CUDAGuard.h>
 #include <c10/util/error.h>
+#include <mutex>
+#include <c10/util/flat_hash_map.h>
+#include <c10/util/hash.h>
+
+#if !defined(USE_ROCM) && defined(PYTORCH_C10_DRIVER_API_SUPPORTED)
+#include <c10/cuda/driver_api.h>
+#endif
 
 namespace c10d {
 namespace symmetric_memory {
@@ -30,8 +38,10 @@ struct NCCLAllocation {
   void* ptr;
   size_t buffer_size;
   int device_idx;
+  std::mutex mutex;
   // Map of group name to peer alloc info
-  std::unordered_map<std::string, c10::intrusive_ptr<NCCLPeerAllocInfo>> peer_alloc_infos_;
+  ska::flat_hash_map<std::string, c10::intrusive_ptr<NCCLPeerAllocInfo>>
+      peer_alloc_infos_;
 
   NCCLAllocation(void* ptr, size_t buffer_size, int device_idx)
       : ptr(ptr), buffer_size(buffer_size), device_idx(device_idx) {}
@@ -50,6 +60,68 @@ struct NCCLAllocation {
   }
 };
 
+namespace {
+
+// Base allocation ptr -> owning NCCL allocation metadata.
+using NCCLAllocMap = ska::flat_hash_map<void*, std::unique_ptr<NCCLAllocation>>;
+// (Tensor storage/data ptr, group name) -> cached SymmetricMemory handle.
+using NCCLSymmMemMap = ska::flat_hash_map<
+    SymmMemKey,
+    c10::intrusive_ptr<NCCLSymmetricMemory>,
+    SymmMemKeyHash>;
+// Base allocation ptr -> cached `(tensor ptr, group)` keys derived from it.
+using NCCLSymmMemKeysByAlloc =
+    ska::flat_hash_map<void*, ska::flat_hash_set<SymmMemKey, SymmMemKeyHash>>;
+
+bool pointer_in_allocation(void* ptr, const NCCLAllocation& allocation) {
+  auto ptr_int = reinterpret_cast<uintptr_t>(ptr);
+  auto base_ptr = reinterpret_cast<uintptr_t>(allocation.ptr);
+  return ptr_int >= base_ptr && ptr_int < base_ptr + allocation.buffer_size;
+}
+
+NCCLAllocMap::iterator find_allocation_covering_linear(
+    void* ptr,
+    NCCLAllocMap& allocations) {
+  return std::find_if(
+      allocations.begin(),
+      allocations.end(),
+      [&](const auto& entry) {
+        return pointer_in_allocation(ptr, *entry.second);
+      });
+}
+
+NCCLAllocMap::iterator find_allocation_covering(
+    void* ptr,
+    NCCLAllocMap& allocations) {
+  auto alloc_it = allocations.find(ptr);
+  if (alloc_it != allocations.end()) {
+    return alloc_it;
+  }
+#if !defined(USE_ROCM) && defined(PYTORCH_C10_DRIVER_API_SUPPORTED)
+  auto driver_api = c10::cuda::DriverAPI::get();
+  CUdeviceptr base_ptr = 0;
+  // Recover the CUDA allocation base for interior pointers before falling
+  // back to the linear scan below when the direct lookup cannot help.
+  auto status = driver_api->cuMemGetAddressRange_(
+      &base_ptr,
+      nullptr,
+      reinterpret_cast<CUdeviceptr>(ptr));
+  if (status == CUDA_SUCCESS) {
+    alloc_it = allocations.find(reinterpret_cast<void*>(base_ptr));
+    if (alloc_it != allocations.end()) {
+      return alloc_it;
+    }
+  }
+#else
+  // No driver API support here, so fall through to the linear scan below.
+#endif
+  return find_allocation_covering_linear(ptr, allocations);
+}
+
+} // namespace
+
+// Before NCCL 2.29, we can use device-side APIs to get peer pointers.
+#if NCCL_VERSION_CODE < NCCL_VERSION(2, 29, 0)
 #ifdef NCCL_HAS_SYMMEM_DEVICE_SUPPORT
 static __global__ void build_ptr_dev(
   ncclWindow_t  handle,
@@ -63,7 +135,8 @@ static __global__ void build_ptr_dev(
       buffer[peer] = ncclGetLsaPointer(handle, offset, peer);
   }
 }
-#endif
+#endif // NCCL_HAS_SYMMEM_DEVICE_SUPPORT
+#endif // NCCL_VERSION_CODE < NCCL_VERSION(2, 29, 0)
 
 class NCCLPeerAllocInfo : public c10::intrusive_ptr_target {
  public:
@@ -81,10 +154,10 @@ class NCCLPeerAllocInfo : public c10::intrusive_ptr_target {
     auto* ncclPg = dynamic_cast<c10d::ProcessGroupNCCL*>(
         group->getBackend(c10::DeviceType::CUDA).get());
     TORCH_CHECK(ncclPg != nullptr, "backend must be a NCCL process group");
-    ncclComm_t comm = reinterpret_cast<ncclComm_t>(ncclPg->getCommPtr());
+    comm_ = reinterpret_cast<ncclComm_t>(ncclPg->getCommPtr());
 
     C10D_NCCL_CHECK(
-      ncclCommWindowRegister(comm, allocation->ptr, buffer_size_, &buffer_win_, NCCL_WIN_COLL_SYMMETRIC),
+      ncclCommWindowRegister(comm_, allocation->ptr, buffer_size_, &buffer_win_, NCCL_WIN_COLL_SYMMETRIC),
       c10::str(
           "Failed to window register segment with ptr ",
           allocation->ptr,
@@ -93,27 +166,26 @@ class NCCLPeerAllocInfo : public c10::intrusive_ptr_target {
           " on rank ",
           rank_));
 
-    void* signal_pad_ptr;
     const size_t signal_pad_size = get_signal_pad_size();
     C10D_NCCL_CHECK(
-        ncclMemAlloc(&signal_pad_ptr, signal_pad_size), "ncclMemAlloc failed");
+        ncclMemAlloc(&signal_pad_ptr_, signal_pad_size), "ncclMemAlloc failed");
     C10D_NCCL_CHECK(
-    ncclCommWindowRegister(comm, signal_pad_ptr, signal_pad_size, &signal_handle_, NCCL_WIN_COLL_SYMMETRIC),
+    ncclCommWindowRegister(comm_, signal_pad_ptr_, signal_pad_size, &signal_handle_, NCCL_WIN_COLL_SYMMETRIC),
     c10::str(
         "Failed to window register segment with ptr ",
-        signal_pad_ptr,
+        signal_pad_ptr_,
         ", size ",
         signal_pad_size,
         " on rank ",
         rank_));
 
-    // Starting from NCCL 2.28, we can use device communicators and get peer pointers
 #ifdef NCCL_HAS_SYMMEM_DEVICE_SUPPORT
-    // Create NCCL device communicator if it doesn't exist. Skip if it already exists.
-    auto& mr = NCCLDevCommManager::get(c10::Device(c10::DeviceType::CUDA, device_idx_));
-    // Each CTA will need a separate barrier. Assume `symm_max_nblocks` as a starting point.
-    mr.try_emplace_devcomm(group_name_, comm, /*LSA*/ symm_max_nblocks, /*GIN*/ symm_max_nblocks);
+    // Register the host-side communicator for device communicator management.
+    // `ncclDevCommCreate` will require it.
+    auto& manager = NCCLDevCommManager::get(c10::Device(c10::DeviceType::CUDA, device_idx_));
+    manager.register_comm(group_name_, comm_);
 
+    // Starting from NCCL 2.28, we can get peer pointers.
     const size_t arr_size = sizeof(void*) * world_size_;
     buffers_dev_ = reinterpret_cast<void**>(
         c10::cuda::CUDACachingAllocator::raw_alloc(arr_size));
@@ -122,6 +194,10 @@ class NCCLPeerAllocInfo : public c10::intrusive_ptr_target {
     buffers_.resize(world_size_);
     signal_pads_.resize(world_size_);
 
+    // Fill out the peer pointer array
+#if NCCL_VERSION_CODE < NCCL_VERSION(2, 29, 0)
+    // Lack of host-side API to get peer pointers, so we get them inside a
+    // kernel and copy the result to host.
     int threads = std::min(128, world_size_);
     auto stream = at::cuda::getCurrentCUDAStream();
     build_ptr_dev<<<1, threads, 0, stream>>>(buffer_win_, 0, buffers_dev_, world_size_);
@@ -139,15 +215,40 @@ class NCCLPeerAllocInfo : public c10::intrusive_ptr_target {
       signal_pads_dev_,  // src (device)
       arr_size,
       cudaMemcpyDeviceToHost));
-#endif
+#else
+  // Starting from NCCL 2.29, we can use host-side APIs to get peer pointers.
+  for (int i = 0; i < world_size_; i++) {
+    // If peer is not accessible within LSA domain, `ncclGetPeerDevicePointer`
+    // returns nullptr.
+    C10D_NCCL_CHECK(
+      ncclGetPeerDevicePointer(buffer_win_, 0, i, &buffers_[i]),
+      "ncclGetPeerDevicePointer failed");
+    C10D_NCCL_CHECK(
+      ncclGetPeerDevicePointer(signal_handle_, 0, i, &signal_pads_[i]),
+      "ncclGetPeerDevicePointer failed");
+  }
+  // Copy the peer access pointers to device arrays.
+  C10_CUDA_CHECK(cudaMemcpy(
+    buffers_dev_,  // dst (device)
+    buffers_.data(),  // src (host)
+    arr_size,
+    cudaMemcpyHostToDevice));
+  C10_CUDA_CHECK(cudaMemcpy(
+    signal_pads_dev_,  // dst (device)
+    signal_pads_.data(),  // src (host)
+    arr_size,
+    cudaMemcpyHostToDevice));
 
-#if NCCL_VERSION_CODE >= NCCL_VERSION(2, 29, 0)
   // Starting from NCCL 2.29, we can use `ncclGetLsaMultimemDevicePointer`
+  // to get multicast address.
   void* mc_addr = nullptr;
+  // Skip CHECK on purpose to improve fault tolerance since some machine's
+  // Fabric Manager may be in bad NVLink Sharp state.
   if (ncclGetLsaMultimemDevicePointer(buffer_win_, 0, &mc_addr) == ncclSuccess) {
     mc_addr_ = mc_addr;
   }
-#endif
+#endif // NCCL_VERSION_CODE < NCCL_VERSION(2, 29, 0)
+#endif // NCCL_HAS_SYMMEM_DEVICE_SUPPORT
   }
 
   // Exact copy is not needed / supported
@@ -155,7 +256,40 @@ class NCCLPeerAllocInfo : public c10::intrusive_ptr_target {
   NCCLPeerAllocInfo& operator=(const NCCLPeerAllocInfo& other) = delete;
   NCCLPeerAllocInfo(NCCLPeerAllocInfo&& other) = default;
   NCCLPeerAllocInfo& operator=(NCCLPeerAllocInfo&& other) = default;
-  ~NCCLPeerAllocInfo() = default;
+
+  ~NCCLPeerAllocInfo() {
+    if (is_finalizing()) {
+      return;
+    }
+    c10::cuda::CUDAGuard guard(device_idx_);
+    if (buffer_win_ != nullptr) {
+      auto res = ncclCommWindowDeregister(comm_, buffer_win_);
+      if (res != ncclSuccess) {
+        LOG(WARNING) << "ncclCommWindowDeregister failed for buffer_win: "
+                     << ncclGetErrorString(res);
+      }
+    }
+    if (signal_handle_ != nullptr) {
+      auto res = ncclCommWindowDeregister(comm_, signal_handle_);
+      if (res != ncclSuccess) {
+        LOG(WARNING) << "ncclCommWindowDeregister failed for signal_handle: "
+                     << ncclGetErrorString(res);
+      }
+    }
+    if (signal_pad_ptr_ != nullptr) {
+      auto res = ncclMemFree(signal_pad_ptr_);
+      if (res != ncclSuccess) {
+        LOG(WARNING) << "ncclMemFree failed for signal_pad: "
+                     << ncclGetErrorString(res);
+      }
+    }
+    if (buffers_dev_ != nullptr) {
+      c10::cuda::CUDACachingAllocator::raw_delete(buffers_dev_);
+    }
+    if (signal_pads_dev_ != nullptr) {
+      c10::cuda::CUDACachingAllocator::raw_delete(signal_pads_dev_);
+    }
+  }
 
  private:
   size_t buffer_size_;
@@ -164,14 +298,15 @@ class NCCLPeerAllocInfo : public c10::intrusive_ptr_target {
   int world_size_;
   std::vector<void*> buffers_;
   std::vector<void*> signal_pads_;
-  void** buffers_dev_;
-  void** signal_pads_dev_;
+  void** buffers_dev_{nullptr};
+  void** signal_pads_dev_{nullptr};
   std::string group_name_;
-  ncclWindow_t buffer_win_;
-  ncclWindow_t signal_handle_;
+  ncclWindow_t buffer_win_{nullptr};
+  ncclWindow_t signal_handle_{nullptr};
+  void* signal_pad_ptr_{nullptr};
   // Multicast address
   void* mc_addr_{nullptr};
-
+  ncclComm_t comm_{nullptr};
   friend class NCCLSymmetricMemory;
 };
 
@@ -318,18 +453,32 @@ class NCCLSymmetricMemoryAllocator : public SymmetricMemoryAllocator {
     // TODO: we might need to use a roundup or mempool for mem allocation.
     void* ptr;
     C10D_NCCL_CHECK(ncclMemAlloc(&ptr, size), "ncclMemAlloc");
-    // TODO: thread safety
-    allocations_.try_emplace(
-        ptr, std::make_unique<NCCLAllocation>(ptr, size, device_idx));
+    {
+      std::lock_guard<std::mutex> lock(mutex_);
+      allocations_.emplace(
+          ptr, std::make_unique<NCCLAllocation>(ptr, size, device_idx));
+    }
     return ptr;
   }
 
   void free(void* ptr) override {
-    // TODO: thread safety
-    allocations_.erase(ptr);
+    std::lock_guard<std::mutex> lock(mutex_);
+    auto alloc_it = allocations_.find(ptr);
+    if (alloc_it == allocations_.end()) {
+      return;
+    }
+    auto cache_keys_it = symm_mem_keys_by_alloc_.find(ptr);
+    if (cache_keys_it != symm_mem_keys_by_alloc_.end()) {
+      for (const auto& key : cache_keys_it->second) {
+        symm_mems_.erase(key);
+      }
+      symm_mem_keys_by_alloc_.erase(cache_keys_it);
+    }
+    allocations_.erase(alloc_it);
   };
 
   size_t get_alloc_size(void* ptr) override {
+    std::lock_guard<std::mutex> lock(mutex_);
     auto it = allocations_.find(ptr);
     if (it == allocations_.end()) {
       TORCH_CHECK(
@@ -342,42 +491,68 @@ class NCCLSymmetricMemoryAllocator : public SymmetricMemoryAllocator {
       void* ptr,
       const std::optional<std::string>& group_name) override {
     TORCH_CHECK(group_name.has_value(), "group_name must be provided");
-    // Then search allocation covering the ptr, get the base address
-    auto alloc_it = std::find_if(
-        allocations_.begin(), allocations_.end(), [&](const auto& pair) {
-          auto& allocation = pair.second;
-          auto ptr_int = reinterpret_cast<uintptr_t>(ptr);
-          auto base_ptr = reinterpret_cast<uintptr_t>(allocation->ptr);
-          return ptr_int >= base_ptr &&
-              ptr_int < base_ptr + allocation->buffer_size;
-        });
-    TORCH_CHECK(
-        alloc_it != allocations_.end(),
-        "Pointer not within any SymmetricMemory allocation, "
-        "is the tensor allocated from SymmetricMemory?");
-
-    auto& allocation = alloc_it->second;
+    NCCLAllocation* allocation;
+    SymmMemKey key{ptr, *group_name};
+    {
+      std::lock_guard<std::mutex> lock(mutex_);
+      auto it = symm_mems_.find(key);
+      if (it != symm_mems_.end()) {
+        return it->second;
+      }
+
+      // Find the allocation covering the ptr under the allocator lock.
+      // We grab a raw pointer to the NCCLAllocation so we can release the
+      // allocator lock before doing expensive per-allocation work.
+      auto alloc_it = find_allocation_covering(ptr, allocations_);
+      TORCH_CHECK(
+          alloc_it != allocations_.end(),
+          "Pointer not within any SymmetricMemory allocation, "
+          "is the tensor allocated from SymmetricMemory?");
+      allocation = alloc_it->second.get();
+    }
 
-    // Get or create peer alloc info for the group
+    // Get or create peer alloc info for the group under the per-allocation
+    // lock. This serializes concurrent rendezvous on the same allocation
+    // for different groups (e.g., forward vs backward).
+    std::lock_guard<std::mutex> alloc_lock(allocation->mutex);
     auto& peer_alloc_infos = allocation->peer_alloc_infos_;
-    auto pai_it = peer_alloc_infos.find(*group_name);
-    if (pai_it == peer_alloc_infos.end()) {
-      // Never rendezvoused with this group before, create a new peer alloc info
-      pai_it = peer_alloc_infos.emplace_hint(
-          pai_it,
-          *group_name,
-          c10::make_intrusive<NCCLPeerAllocInfo>(allocation.get(), *group_name));
+    auto& pai = peer_alloc_infos[*group_name];
+    if (!pai) {
+      pai = c10::make_intrusive<NCCLPeerAllocInfo>(allocation, *group_name);
     }
-
-    auto& pai = pai_it->second;
-    size_t offset = reinterpret_cast<uintptr_t>(ptr) - reinterpret_cast<uintptr_t>(allocation->ptr);
-    return c10::make_intrusive<NCCLSymmetricMemory>(pai, offset);
+    size_t offset =
+        reinterpret_cast<uintptr_t>(ptr) -
+        reinterpret_cast<uintptr_t>(allocation->ptr);
+    // Create the SymmetricMemory handle.
+    auto symm_mem = c10::make_intrusive<NCCLSymmetricMemory>(pai, offset);
+    {
+      std::lock_guard<std::mutex> lock(mutex_);
+      // Insert the SymmetricMemory handle into the map (cache), keyed by the
+      // (Tensor storage ptr, group name) pair.
+      auto [it, inserted] = symm_mems_.emplace(key, symm_mem);
+      if (!inserted) {
+        // This condition should rarely happen, only when another thread happens
+        // to be concurrently rendezvousing with the same allocation for the
+        // same group.  For safety, we return the existing SymmetricMemory
+        // handle and discard the new one.
+        return it->second;
+      }
+      // There is no more use of `key`; we can move it into the per-allocation
+      // key set to avoid an extra copy.
+      symm_mem_keys_by_alloc_[allocation->ptr].insert(std::move(key));
+    }
+    return symm_mem;
   }
 
   bool has_multicast_support(int device_idx) override {
     return device_has_multicast_support(device_idx);
   }
 
+  bool has_allocation(void* ptr) override {
+    std::lock_guard<std::mutex> lock(mutex_);
+    return find_allocation_covering(ptr, allocations_) != allocations_.end();
+  }
+
   c10::DeviceType supported_device_type() override {
     return c10::DeviceType::CUDA;
   }
@@ -387,7 +562,10 @@ class NCCLSymmetricMemoryAllocator : public SymmetricMemoryAllocator {
   }
 
  private:
-  std::unordered_map<void*, std::unique_ptr<NCCLAllocation>> allocations_;
+  std::mutex mutex_;
+  NCCLAllocMap allocations_;
+  NCCLSymmMemMap symm_mems_;
+  NCCLSymmMemKeysByAlloc symm_mem_keys_by_alloc_;
 };
 
 struct RegisterNCCLSymmetricMemoryAllocator {
diff --git a/torch/csrc/distributed/c10d/symm_mem/NVSHMEMSymmetricMemory.cu b/torch/csrc/distributed/c10d/symm_mem/NVSHMEMSymmetricMemory.cpp
similarity index 88%
rename from torch/csrc/distributed/c10d/symm_mem/NVSHMEMSymmetricMemory.cu
rename to torch/csrc/distributed/c10d/symm_mem/NVSHMEMSymmetricMemory.cpp
index 8bfaa1ad0fac7..7a3e853197030 100644
--- a/torch/csrc/distributed/c10d/symm_mem/NVSHMEMSymmetricMemory.cu
+++ b/torch/csrc/distributed/c10d/symm_mem/NVSHMEMSymmetricMemory.cpp
@@ -1,9 +1,9 @@
-#include <torch/csrc/distributed/c10d/cuda/utils.hpp>
 #include <torch/csrc/distributed/c10d/GroupRegistry.hpp>
-#include <torch/csrc/distributed/c10d/symm_mem/CUDASymmetricMemory-inl.h>
+#include <torch/csrc/distributed/c10d/cuda/utils.hpp>
+#include <torch/csrc/distributed/c10d/symm_mem/CUDASymmetricMemoryTypes.hpp>
 #include <torch/csrc/distributed/c10d/symm_mem/CUDASymmetricMemoryUtils.hpp>
 #include <torch/csrc/distributed/c10d/symm_mem/SymmetricMemory.hpp>
-#include <torch/csrc/distributed/c10d/symm_mem/nvshmem_extension.cuh>
+#include <torch/csrc/distributed/c10d/symm_mem/nvshmem_extension.hpp>
 #include <torch/csrc/distributed/c10d/symm_mem/nvshmem_team_manager.hpp>
 
 #include <ATen/ceil_div.h>
@@ -11,13 +11,18 @@
 #include <c10/cuda/CUDACachingAllocator.h>
 #include <c10/cuda/CUDAGuard.h>
 #include <c10/util/error.h>
+#include <c10/util/flat_hash_map.h>
+
+#include <mutex>
 
 // Starting from NVSHMEM 3.3.9, nvshmem_host.h exists so that we can cleanly
 // include only the nvshmem host library headers:
 // #include <nvshmem_host.h>
 // It translates into the following two lines:
+#if !defined(USE_ROCM)
 #include <host/nvshmem_api.h>
 #include <host/nvshmemx_api.h>
+#endif
 // For maximum compatibility, we use the "host/" style for now.
 
 namespace c10d {
@@ -52,7 +57,9 @@ struct NVSHMEMAllocation {
 };
 
 // A map from group name to rank-to-global rank mapping
-static std::unordered_map<std::string, std::vector<int>> rank_to_global_rank_map{};
+static std::mutex rank_map_mutex;
+static std::unordered_map<std::string, std::vector<int>>
+    rank_to_global_rank_map{};
 // A map from group name to rank-to-global rank device array
 static std::unordered_map<std::string, int*> rank_to_global_rank_dev_map{};
 
@@ -78,6 +85,7 @@ class NVSHMEMPeerAllocInfo : public c10::intrusive_ptr_target {
 
     // Exchange rank to global rank mapping for this group.
     // If it is already available, skip the exchange.
+    std::lock_guard<std::mutex> rank_map_lock(rank_map_mutex);
     auto it = rank_to_global_rank_map.find(group_name);
     if (it == rank_to_global_rank_map.end()) {
       auto global_group = resolve_process_group("0");
@@ -91,11 +99,13 @@ class NVSHMEMPeerAllocInfo : public c10::intrusive_ptr_target {
                   << ", group_name: " << group_name
                   << ", exchanged_n_times: " << exchanged_n_times;
       }
-      it = rank_to_global_rank_map.emplace_hint(it, group_name, rank_to_global_rank);
+      it = rank_to_global_rank_map.emplace_hint(
+          it, group_name, rank_to_global_rank);
 
       // Emplace a device array of rank to global rank mapping
-      auto rank_to_global_rank_dev = reinterpret_cast<int*>(
-          c10::cuda::CUDACachingAllocator::raw_alloc(sizeof(int) * world_size_));
+      auto rank_to_global_rank_dev =
+          reinterpret_cast<int*>(c10::cuda::CUDACachingAllocator::raw_alloc(
+              sizeof(int) * world_size_));
       AT_CUDA_CHECK(cudaMemcpy(
           rank_to_global_rank_dev,
           rank_to_global_rank.data(),
@@ -140,12 +150,14 @@ class NVSHMEMPeerAllocInfo : public c10::intrusive_ptr_target {
         arr_size,
         cudaMemcpyHostToDevice));
 
+#if !defined(USE_ROCM) // Multi-cast is not supported on ROCm yet
     // Initialize multicast address
     // On unsupported platforms, this API returns a nullptr.
     auto device = c10::Device(c10::DeviceType::CUDA, allocation->device_idx);
     auto& team_manager = c10d::nvshmem_extension::TeamManager::get(device);
     auto team = team_manager.get_team(group_name, rank_to_global_rank);
     mc_addr_ = nvshmemx_mc_ptr(team, base_ptr_);
+#endif
   }
 
  private:
@@ -255,6 +267,7 @@ class NVSHMEMSymmetricMemory : public SymmetricMemory {
   }
 
   const std::vector<int>& get_rank_to_global_rank() override {
+    std::lock_guard<std::mutex> lock(rank_map_mutex);
     auto it = rank_to_global_rank_map.find(group_name_);
     if (it == rank_to_global_rank_map.end()) {
       TORCH_CHECK(false, "Group name not found in rank_to_global_rank_map");
@@ -263,6 +276,7 @@ class NVSHMEMSymmetricMemory : public SymmetricMemory {
   };
 
   int* get_rank_to_global_rank_dev() override {
+    std::lock_guard<std::mutex> lock(rank_map_mutex);
     auto it = rank_to_global_rank_dev_map.find(group_name_);
     if (it == rank_to_global_rank_dev_map.end()) {
       TORCH_CHECK(false, "Group name not found in rank_to_global_rank_dev_map");
@@ -339,9 +353,11 @@ static void initialize_nvshmem_with_store(
   is_initialized = true;
 
   // Print version
+#if !defined(USE_ROCM)
   int major, minor;
   ::nvshmem_info_get_version(&major, &minor);
   LOG(INFO) << "NVSHMEM is available, version: " << major << '.' << minor;
+#endif
 }
 
 class NVSHMEMSymmetricMemoryAllocator : public SymmetricMemoryAllocator {
@@ -358,23 +374,27 @@ class NVSHMEMSymmetricMemoryAllocator : public SymmetricMemoryAllocator {
 
     // NVSHMEM needs to be initialized with the global group
     auto group = resolve_process_group("0");
-    initialize_nvshmem_with_store(group->getStore(), group->getRank(), group->getSize(), device_idx);
+    initialize_nvshmem_with_store(
+        group->getStore(), group->getRank(), group->getSize(), device_idx);
 
     auto ptr = nvshmem_malloc(size);
     // If size is 0 (which is legal allocation request) we shouldn't error out
     TORCH_CHECK(ptr != nullptr || size == 0, "nvshmem_malloc failed");
-    // TODO: thread safety
-    allocations_.try_emplace(
-        ptr, std::make_unique<NVSHMEMAllocation>(ptr, size, device_idx));
+    {
+      std::lock_guard<std::mutex> lock(mutex_);
+      allocations_.try_emplace(
+          ptr, std::make_unique<NVSHMEMAllocation>(ptr, size, device_idx));
+    }
     return ptr;
   }
 
   void free(void* ptr) override {
-    // TODO: thread safety
+    std::lock_guard<std::mutex> lock(mutex_);
     allocations_.erase(ptr);
   };
 
   size_t get_alloc_size(void* ptr) override {
+    std::lock_guard<std::mutex> lock(mutex_);
     auto it = allocations_.find(ptr);
     if (it == allocations_.end()) {
       TORCH_CHECK(
@@ -387,8 +407,9 @@ class NVSHMEMSymmetricMemoryAllocator : public SymmetricMemoryAllocator {
       void* ptr,
       const std::optional<std::string>& group_name) override {
     TORCH_CHECK(group_name.has_value());
+    std::lock_guard<std::mutex> lock(mutex_);
     {
-      auto it = symm_mems_.find(std::make_tuple(ptr, *group_name));
+      auto it = symm_mems_.find(SymmMemKey{ptr, *group_name});
       if (it != symm_mems_.end()) {
         return it->second;
       }
@@ -413,7 +434,7 @@ class NVSHMEMSymmetricMemoryAllocator : public SymmetricMemoryAllocator {
 
     // Search again using allocation base ptr (which is the key we use for
     // caching, see below)
-    auto it = symm_mems_.find(std::make_tuple(allocation->ptr, *group_name));
+    auto it = symm_mems_.find(SymmMemKey{allocation->ptr, *group_name});
     c10::intrusive_ptr<NVSHMEMSymmetricMemory> symm_mem;
     if (it != symm_mems_.end()) {
       // Base allocation has been rendezvoused
@@ -425,7 +446,7 @@ class NVSHMEMSymmetricMemoryAllocator : public SymmetricMemoryAllocator {
     }
 
     // Cache rendezvous using allocation's base address as key
-    symm_mems_[std::make_tuple(allocation->ptr, *group_name)] = symm_mem;
+    symm_mems_[SymmMemKey{allocation->ptr, *group_name}] = symm_mem;
 
     // TODO: change the `ptr` below to `tensor.data_ptr()` when adding support
     // for user slice/view operations. For MemPool support,
@@ -446,6 +467,18 @@ class NVSHMEMSymmetricMemoryAllocator : public SymmetricMemoryAllocator {
     return device_has_multicast_support(device_idx);
   }
 
+  bool has_allocation(void* ptr) override {
+    std::lock_guard<std::mutex> lock(mutex_);
+    auto alloc_it = std::find_if(
+        allocations_.begin(), allocations_.end(), [&](const auto& pair) {
+          auto ptr_int = reinterpret_cast<uintptr_t>(ptr);
+          auto base_ptr = reinterpret_cast<uintptr_t>(pair.second->ptr);
+          return ptr_int >= base_ptr &&
+              ptr_int < base_ptr + pair.second->buffer_size;
+        });
+    return alloc_it != allocations_.end();
+  }
+
   c10::DeviceType supported_device_type() override {
     return c10::DeviceType::CUDA;
   }
@@ -455,10 +488,12 @@ class NVSHMEMSymmetricMemoryAllocator : public SymmetricMemoryAllocator {
   }
 
  private:
+  std::mutex mutex_;
   std::unordered_map<void*, std::unique_ptr<NVSHMEMAllocation>> allocations_;
-  std::map<
-      std::tuple<void*, std::string>,
-      c10::intrusive_ptr<NVSHMEMSymmetricMemory>>
+  ska::flat_hash_map<
+      SymmMemKey,
+      c10::intrusive_ptr<NVSHMEMSymmetricMemory>,
+      SymmMemKeyHash>
       symm_mems_;
 };
 
diff --git a/torch/csrc/distributed/c10d/symm_mem/SymmetricMemory.cpp b/torch/csrc/distributed/c10d/symm_mem/SymmetricMemory.cpp
index 366a6d8031a9a..279c578547e15 100644
--- a/torch/csrc/distributed/c10d/symm_mem/SymmetricMemory.cpp
+++ b/torch/csrc/distributed/c10d/symm_mem/SymmetricMemory.cpp
@@ -4,6 +4,7 @@
 #include <torch/custom_class.h>
 
 #include <atomic>
+#include <mutex>
 
 namespace {
 
@@ -37,16 +38,19 @@ class AllocatorMap {
   void register_allocator(
       c10::DeviceType device_type,
       c10::intrusive_ptr<SymmetricMemoryAllocator> allocator) {
+    std::lock_guard<std::mutex> lock(mutex_);
     map_[device_type] = std::move(allocator);
   }
 
   void register_availability(
       const std::string& name,
       c10::intrusive_ptr<SymmetricMemoryAllocator> allocator) {
+    std::lock_guard<std::mutex> lock(mutex_);
     avail_map_[name] = std::move(allocator);
   }
 
   void set_backend(const std::string& name) {
+    std::lock_guard<std::mutex> lock(mutex_);
     auto it = avail_map_.find(name);
     TORCH_CHECK(
         it != avail_map_.end(),
@@ -62,10 +66,11 @@ class AllocatorMap {
       }
       TORCH_CHECK(!in_use_, "Backend can not be changed after use.");
     }
-    register_allocator(device_type, it->second);
+    map_[device_type] = it->second;
   }
 
   std::optional<std::string> get_backend(c10::DeviceType device_type) {
+    std::lock_guard<std::mutex> lock(mutex_);
     auto it = map_.find(device_type);
     if (it == map_.end()) {
       return std::nullopt;
@@ -75,6 +80,7 @@ class AllocatorMap {
 
   c10::intrusive_ptr<SymmetricMemoryAllocator> get_allocator(
       c10::DeviceType device_type) {
+    std::lock_guard<std::mutex> lock(mutex_);
     auto it = map_.find(device_type);
     TORCH_CHECK(
         it != map_.end(),
@@ -85,6 +91,7 @@ class AllocatorMap {
   }
 
   bool has_allocator(c10::DeviceType device_type) {
+    std::lock_guard<std::mutex> lock(mutex_);
     auto it = map_.find(device_type);
     return it != map_.end();
   }
@@ -96,6 +103,7 @@ class AllocatorMap {
  private:
   AllocatorMap() = default;
 
+  std::mutex mutex_;
   std::unordered_map<
       c10::DeviceType,
       c10::intrusive_ptr<SymmetricMemoryAllocator>>
@@ -113,9 +121,11 @@ class AllocatorMap {
   bool in_use_ = false;
 };
 
+static std::mutex group_info_mutex;
 static std::unordered_map<std::string, GroupInfo> group_info_map{};
 
 // Data structures for tracking persistent allocations
+static std::mutex persistent_alloc_mutex;
 static std::unordered_map<uint64_t, void*> alloc_id_to_dev_ptr{};
 static std::unordered_map<uint64_t, c10::weak_intrusive_ptr<c10::StorageImpl>>
     alloc_id_to_storage{};
@@ -127,6 +137,7 @@ static at::Tensor empty_strided_p2p_persistent(
     c10::Device device,
     const std::optional<std::string>& group_name,
     uint64_t alloc_id) {
+  std::lock_guard<std::mutex> lock(persistent_alloc_mutex);
   // Make the allocation fails if a previous allocation with the same alloc_id
   // is still active.
   auto storage = alloc_id_to_storage.find(alloc_id);
@@ -227,6 +238,7 @@ void set_group_info(
     int rank,
     int world_size,
     c10::intrusive_ptr<Store> store) {
+  std::lock_guard<std::mutex> lock(group_info_mutex);
   TORCH_CHECK(group_info_map.find(group_name) == group_info_map.end());
   GroupInfo group_info;
   group_info.rank = rank;
@@ -236,6 +248,7 @@ void set_group_info(
 }
 
 GroupInfo& get_group_info(const std::string& group_name) {
+  std::lock_guard<std::mutex> lock(group_info_mutex);
   TORCH_CHECK(
       group_info_map.find(group_name) != group_info_map.end(),
       "get_group_info: no group info associated with the group name ",
@@ -282,6 +295,14 @@ TORCH_API c10::intrusive_ptr<SymmetricMemory> rendezvous(
   return allocator->rendezvous(tensor.storage().data_ptr().get(), group_name);
 }
 
+TORCH_API bool is_symm_mem_tensor(const at::Tensor& tensor) {
+  if (!has_allocator(tensor.device().type())) {
+    return false;
+  }
+  auto allocator = get_allocator(tensor.device().type());
+  return allocator->has_allocation(tensor.storage().data_ptr().get());
+}
+
 TORCH_API bool has_multicast_support(
     c10::DeviceType device_type,
     int device_idx) {
@@ -538,6 +559,8 @@ TORCH_LIBRARY_FRAGMENT(symm_mem, m) {
   m.def("nccl_get(Tensor(a!) tensor, int peer) -> ()");
   m.def("nccl_wait_for_signal(Tensor sigpad, int signal) -> ()");
   m.def("nccl_put_with_signal(Tensor(a) tensor, int signal, int peer) -> ()");
+  m.def(
+      "nccl_reduce_scatter_offset(Tensor input, Tensor(a!)[] out, str group_name, int dim, int[]? offsets=None, int[]? dst_ranks=None, str red_op='sum') -> ()");
   m.def(
       "nvshmem_all_to_all(Tensor input, Tensor(a!) out, str group_name) -> Tensor(a!)");
   m.def(
diff --git a/torch/csrc/distributed/c10d/symm_mem/SymmetricMemory.hpp b/torch/csrc/distributed/c10d/symm_mem/SymmetricMemory.hpp
index 4156e24be1f26..a509568381e33 100644
--- a/torch/csrc/distributed/c10d/symm_mem/SymmetricMemory.hpp
+++ b/torch/csrc/distributed/c10d/symm_mem/SymmetricMemory.hpp
@@ -113,6 +113,9 @@ class SymmetricMemoryAllocator : public c10::intrusive_ptr_target {
   virtual bool has_multicast_support(int device_idx) = 0;
   virtual c10::DeviceType supported_device_type() = 0;
   virtual std::string name() = 0;
+  virtual bool has_allocation(void* ptr) {
+    return false;
+  }
 };
 
 C10_EXPORT bool is_finalizing();
@@ -191,6 +194,8 @@ TORCH_API bool has_multicast_support(
     c10::DeviceType device_type,
     int device_idx);
 
+TORCH_API bool is_symm_mem_tensor(const at::Tensor& tensor);
+
 TORCH_API void set_backend(const std::string& name);
 
 TORCH_API std::optional<std::string> get_backend(c10::Device device);
diff --git a/torch/csrc/distributed/c10d/symm_mem/intra_node_comm.cpp b/torch/csrc/distributed/c10d/symm_mem/intra_node_comm.cpp
index f62577e701847..a84c0dcebebfe 100644
--- a/torch/csrc/distributed/c10d/symm_mem/intra_node_comm.cpp
+++ b/torch/csrc/distributed/c10d/symm_mem/intra_node_comm.cpp
@@ -3,7 +3,10 @@
 #include <torch/csrc/distributed/c10d/symm_mem/intra_node_comm.hpp>
 
 #if defined(USE_ROCM)
-#include <rocm_smi/rocm_smi.h>
+#include <amd_smi/amdsmi.h>
+#include <dlfcn.h>
+#include <cstdlib>
+#include <string>
 #endif
 
 namespace c10d::intra_node_comm {
@@ -20,7 +23,7 @@ static int intraNodeCommIdx = 0;
  * Query the nvlink connection among devices.
  */
 static NvlMesh getNvlMesh(const std::vector<int>& rankToDeviceIdx) {
-#if !defined(USE_RCOM)
+#if !defined(USE_ROCM)
   auto connectivity = detect_dma_connectivity(c10::DeviceType::CUDA, "nvlink");
   NvlMesh nvlMesh = {};
   for (size_t srcRank = 0; srcRank < kMaxDevices; ++srcRank) {
@@ -35,23 +38,126 @@ static NvlMesh getNvlMesh(const std::vector<int>& rankToDeviceIdx) {
   }
   return nvlMesh;
 #else
+  // Load libamd_smi at runtime to avoid linking it into torch_hip (double-load
+  // with Python amdsmi causes bus errors). Types/constants from amdsmi.h only.
+  struct AmdsmiApi {
+    amdsmi_status_t (*init)(uint64_t);
+    amdsmi_status_t (*get_socket_handles)(uint32_t*, amdsmi_socket_handle*);
+    amdsmi_status_t (*get_processor_handles)(
+        amdsmi_socket_handle,
+        uint32_t*,
+        amdsmi_processor_handle*);
+    amdsmi_status_t (*is_P2P_accessible)(
+        amdsmi_processor_handle,
+        amdsmi_processor_handle,
+        bool*);
+  };
+  static void* amdsmi_handle = nullptr;
+  static AmdsmiApi amdsmi = {};
+  static bool amdsmi_resolved = false;
+
+  if (!amdsmi_resolved) {
+    amdsmi_resolved = true;
+    const char* rocm = std::getenv("ROCM_PATH");
+    std::string path =
+        rocm ? std::string(rocm) + "/lib/libamd_smi.so" : "libamd_smi.so";
+    amdsmi_handle = dlopen(path.c_str(), RTLD_NOW | RTLD_LOCAL);
+    if (!amdsmi_handle) {
+      amdsmi_handle = dlopen("libamd_smi.so", RTLD_NOW | RTLD_LOCAL);
+    }
+    if (!amdsmi_handle) {
+      LOG(ERROR) << "IntraNodeComm:: getNvlMesh: dlopen libamd_smi.so failed: "
+                 << dlerror();
+      return {};
+    }
+    amdsmi.init = reinterpret_cast<decltype(amdsmi.init)>(
+        dlsym(amdsmi_handle, "amdsmi_init"));
+    amdsmi.get_socket_handles =
+        reinterpret_cast<decltype(amdsmi.get_socket_handles)>(
+            dlsym(amdsmi_handle, "amdsmi_get_socket_handles"));
+    amdsmi.get_processor_handles =
+        reinterpret_cast<decltype(amdsmi.get_processor_handles)>(
+            dlsym(amdsmi_handle, "amdsmi_get_processor_handles"));
+    amdsmi.is_P2P_accessible =
+        reinterpret_cast<decltype(amdsmi.is_P2P_accessible)>(
+            dlsym(amdsmi_handle, "amdsmi_is_P2P_accessible"));
+    if (!amdsmi.init || !amdsmi.get_socket_handles ||
+        !amdsmi.get_processor_handles || !amdsmi.is_P2P_accessible) {
+      LOG(ERROR) << "IntraNodeComm:: getNvlMesh: dlsym amdsmi failed";
+      return {};
+    }
+  }
+
   NvlMesh nvlMesh = {};
   const auto worldSize = rankToDeviceIdx.size();
-  // For each device, loop over devices connected to it
+
+  uint32_t socket_count = 0;
+  amdsmi_status_t ret = amdsmi.get_socket_handles(&socket_count, nullptr);
+  if (ret == AMDSMI_STATUS_NOT_INIT) {
+    ret = amdsmi.init(AMDSMI_INIT_AMD_GPUS);
+    if (ret != AMDSMI_STATUS_SUCCESS) {
+      LOG(ERROR) << "IntraNodeComm:: getNvlMesh: amdsmi_init failed, ret="
+                 << static_cast<int>(ret);
+      return {};
+    }
+    socket_count = 0;
+    ret = amdsmi.get_socket_handles(&socket_count, nullptr);
+  }
+  if (ret != AMDSMI_STATUS_SUCCESS) {
+    LOG(ERROR)
+        << "IntraNodeComm:: getNvlMesh: amdsmi_get_socket_handles failed, ret="
+        << static_cast<int>(ret);
+    return {};
+  }
+
+  std::vector<amdsmi_socket_handle> socket_handles(socket_count);
+  ret = amdsmi.get_socket_handles(&socket_count, &socket_handles[0]);
+  if (ret != AMDSMI_STATUS_SUCCESS) {
+    LOG(ERROR)
+        << "IntraNodeComm:: getNvlMesh: amdsmi_get_socket_handles (buffer) failed, ret="
+        << static_cast<int>(ret);
+    return {};
+  }
+
+  std::vector<amdsmi_processor_handle> processor_handles;
+  for (size_t i = 0; i < socket_count; ++i) {
+    uint32_t device_count = 0;
+    ret =
+        amdsmi.get_processor_handles(socket_handles[i], &device_count, nullptr);
+    if (ret != AMDSMI_STATUS_SUCCESS) {
+      LOG(ERROR)
+          << "IntraNodeComm:: getNvlMesh: amdsmi_get_processor_handles (count) failed, ret="
+          << static_cast<int>(ret);
+      return {};
+    }
+    std::vector<amdsmi_processor_handle> _processor_handles(device_count);
+    ret = amdsmi.get_processor_handles(
+        socket_handles[i], &device_count, &_processor_handles[0]);
+    if (ret != AMDSMI_STATUS_SUCCESS) {
+      LOG(ERROR)
+          << "IntraNodeComm:: getNvlMesh: amdsmi_get_processor_handles (buffer) failed, ret="
+          << static_cast<int>(ret);
+      return {};
+    }
+    processor_handles.insert(
+        processor_handles.end(),
+        _processor_handles.begin(),
+        _processor_handles.end());
+  }
+
   for (size_t idx = 0; idx < worldSize; ++idx) {
     for (size_t link = 0; link < kMaxDevices; ++link) {
       if (idx == link)
         continue;
-
       bool conn = false;
-      auto ret = rsmi_is_P2P_accessible(idx, link, &conn);
-      if (ret != RSMI_STATUS_SUCCESS) {
+      ret = amdsmi.is_P2P_accessible(
+          processor_handles[idx], processor_handles[link], &conn);
+      if (ret != AMDSMI_STATUS_SUCCESS) {
         LOG(ERROR)
-            << "IntraNodeComm: getNvlMesh: rsmi_is_P2P_accessible returned error ret="
-            << ret;
+            << "IntraNodeComm: getNvlMesh: amdsmi_is_P2P_accessible failed, ret="
+            << static_cast<int>(ret);
         return {};
       }
-
       if (conn) {
         nvlMesh[idx][link] += 1;
       }
@@ -88,11 +194,13 @@ IntraNodeComm::IntraNodeComm(
     c10::intrusive_ptr<c10d::Store> store,
     size_t rank,
     size_t worldSize,
-    std::optional<size_t> bufferSize)
+    std::optional<size_t> bufferSize,
+    std::string groupName)
     : store_(std::move(store)),
       rank_(rank),
       worldSize_(worldSize),
-      bufferSize_(bufferSize.has_value() ? *bufferSize : kDefaultBufferSize) {}
+      bufferSize_(bufferSize.has_value() ? *bufferSize : kDefaultBufferSize),
+      groupName_(std::move(groupName)) {}
 
 IntraNodeComm::~IntraNodeComm() {
   if (!isInitialized_) {
@@ -171,14 +279,6 @@ bool IntraNodeComm::rendezvous() {
   gethostname(devInfo.hostname, sizeof(devInfo.hostname));
   devInfo.deviceIdx = deviceIdx_;
 
-#if defined(USE_ROCM)
-  auto ret = rsmi_init(0);
-  if (ret != RSMI_STATUS_SUCCESS) {
-    LOG(ERROR) << "IntraNodeComm:: rendezvous failed in rsmi_init, ret=" << ret;
-    return false;
-  }
-#endif
-
   auto peerDevInfos =
       storeAllGather(store_, "handshake-0", rank_, worldSize_, devInfo);
 
@@ -214,11 +314,13 @@ bool IntraNodeComm::rendezvous() {
     return false;
   }
 
-  auto groupName = "IntraNodeComm" + std::to_string(intraNodeCommIdx++);
+  const std::string name = groupName_.empty()
+      ? "IntraNodeComm" + std::to_string(intraNodeCommIdx++)
+      : groupName_;
   set_group_info(
-      groupName, static_cast<int>(rank_), static_cast<int>(worldSize_), store_);
+      name, static_cast<int>(rank_), static_cast<int>(worldSize_), store_);
   auto allocator = get_allocator(c10::DeviceType::CUDA);
-  symmetricMemoryPtr_ = allocator->alloc(bufferSize_, deviceIdx_, groupName);
+  symmetricMemoryPtr_ = allocator->alloc(bufferSize_, deviceIdx_, name);
   symmetricMemory_ = allocator->rendezvous(symmetricMemoryPtr_, std::nullopt);
   isInitialized_ = true;
   return true;
diff --git a/torch/csrc/distributed/c10d/symm_mem/intra_node_comm.cu b/torch/csrc/distributed/c10d/symm_mem/intra_node_comm.cu
index 6a6a6520e36ba..1a3fb6c5c0a25 100644
--- a/torch/csrc/distributed/c10d/symm_mem/intra_node_comm.cu
+++ b/torch/csrc/distributed/c10d/symm_mem/intra_node_comm.cu
@@ -1,6 +1,6 @@
 #include <torch/csrc/distributed/c10d/symm_mem/intra_node_comm.hpp>
 
-#include <torch/csrc/distributed/c10d/symm_mem/CUDASymmetricMemory-inl.h>
+#include <torch/csrc/distributed/c10d/symm_mem/CUDASymmetricMemory-inl.cuh>
 
 namespace c10d {
 namespace intra_node_comm {
@@ -23,7 +23,7 @@ static void checkInput(const at::Tensor& input, int deviceIdx) {
 }
 
 bool isIntraNodeCommSupported() {
-#if defined(USE_ROCM) || (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ < 800))
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ < 800)
   return false;
 #else
   return true;
@@ -46,7 +46,7 @@ at::Tensor IntraNodeComm::oneShotAllReduce(
       at::TensorOptions().dtype(input.dtype()).device(input.device()));
 
   symmMemTensor.copy_(input);
-  op.call(symmMemTensor, "sum", "", input);
+  op.call(symmMemTensor, "sum", groupName_, input);
   return input;
 }
 
@@ -65,7 +65,7 @@ at::Tensor IntraNodeComm::twoShotAllReduce(
       at::TensorOptions().dtype(input.dtype()).device(input.device()));
 
   symmMemTensor.copy_(input);
-  op.call(symmMemTensor, "sum", "");
+  op.call(symmMemTensor, "sum", groupName_);
   input.copy_(symmMemTensor);
   return input;
 }
diff --git a/torch/csrc/distributed/c10d/symm_mem/intra_node_comm.hpp b/torch/csrc/distributed/c10d/symm_mem/intra_node_comm.hpp
index 7b5e8ff999c5d..8b2a425c47ce2 100644
--- a/torch/csrc/distributed/c10d/symm_mem/intra_node_comm.hpp
+++ b/torch/csrc/distributed/c10d/symm_mem/intra_node_comm.hpp
@@ -33,7 +33,8 @@ class TORCH_API IntraNodeComm : public c10::intrusive_ptr_target {
       c10::intrusive_ptr<c10d::Store> store,
       size_t rank,
       size_t worldSize,
-      std::optional<size_t> bufferSize = std::nullopt);
+      std::optional<size_t> bufferSize = std::nullopt,
+      std::string groupName = "");
 
   ~IntraNodeComm() override;
 
@@ -67,6 +68,7 @@ class TORCH_API IntraNodeComm : public c10::intrusive_ptr_target {
   size_t rank_;
   size_t worldSize_;
   size_t bufferSize_;
+  std::string groupName_;
 
   /**
    * Members initialized after rendezvous
diff --git a/torch/csrc/distributed/c10d/symm_mem/macros.hpp b/torch/csrc/distributed/c10d/symm_mem/macros.hpp
new file mode 100644
index 0000000000000..18bcc0559b8c4
--- /dev/null
+++ b/torch/csrc/distributed/c10d/symm_mem/macros.hpp
@@ -0,0 +1,24 @@
+// Macros for type dispatch and common utilities for symmetric memory
+#pragma once
+
+#include <ATen/ATen.h>
+
+// Convert ATen floating point types to NV floating point types
+// at::kBFloat16 -> __nv_bfloat16
+// at::kHalf -> __half
+// Float is the same.
+
+#define AT_DISPATCH_CASE_CONVERT(enum_type, scalar_type, ...) \
+  case enum_type: {                                           \
+    AT_PRIVATE_CHECK_SELECTIVE_BUILD(enum_type);              \
+    using scalar_t = scalar_type;                             \
+    return __VA_ARGS__();                                     \
+  }
+
+#define AT_DISPATCH_NV_FLOATS(scalar_type, name, ...)                      \
+  AT_DISPATCH_SWITCH(                                                      \
+      scalar_type,                                                         \
+      name,                                                                \
+      AT_DISPATCH_CASE_CONVERT(at::kBFloat16, __nv_bfloat16, __VA_ARGS__); \
+      AT_DISPATCH_CASE_CONVERT(at::kHalf, __half, __VA_ARGS__);            \
+      AT_DISPATCH_CASE(at::kFloat, __VA_ARGS__));
diff --git a/torch/csrc/distributed/c10d/symm_mem/nccl_dev_cap.hpp b/torch/csrc/distributed/c10d/symm_mem/nccl_dev_cap.hpp
index 63b3d452b7e4f..e6a2e1c366c28 100644
--- a/torch/csrc/distributed/c10d/symm_mem/nccl_dev_cap.hpp
+++ b/torch/csrc/distributed/c10d/symm_mem/nccl_dev_cap.hpp
@@ -10,11 +10,19 @@
 #endif
 
 #if NCCL_VERSION_CODE >= NCCL_VERSION(2, 28, 0)
+#if !defined(USE_ROCM)
 #define NCCL_HAS_SYMMEM_DEVICE_SUPPORT
 #include <nccl_device.h>
 #endif
+#endif
 
-#if NCCL_VERSION_CODE >= NCCL_VERSION(2, 29, 0)
+#if defined(NCCL_HAS_SYMMEM_DEVICE_SUPPORT) && \
+    NCCL_VERSION_CODE >= NCCL_VERSION(2, 29, 0)
 #define NCCL_HAS_ONE_SIDED_API
 #endif
+
+#if defined(NCCL_HAS_SYMMEM_DEVICE_SUPPORT) && \
+    NCCL_VERSION_CODE >= NCCL_VERSION(2, 29, 7)
+#define NCCL_DEVICE_HAS_REDUCE_COPY
+#endif
 #endif // USE_NCCL
diff --git a/torch/csrc/distributed/c10d/symm_mem/nccl_devcomm_manager.hpp b/torch/csrc/distributed/c10d/symm_mem/nccl_devcomm_manager.hpp
index c3a7da8f4148f..437f956d1f939 100644
--- a/torch/csrc/distributed/c10d/symm_mem/nccl_devcomm_manager.hpp
+++ b/torch/csrc/distributed/c10d/symm_mem/nccl_devcomm_manager.hpp
@@ -3,7 +3,11 @@
 #include <c10/cuda/CUDAException.h>
 #include <c10/cuda/CUDAGuard.h>
 #include <c10/util/Exception.h>
+#include <c10/util/Logging.h>
 #include <torch/csrc/distributed/c10d/symm_mem/nccl_dev_cap.hpp>
+#include <functional>
+#include <mutex>
+#include <optional>
 #include <string>
 #include <unordered_map>
 
@@ -11,13 +15,30 @@
 
 namespace c10d::symmetric_memory {
 
-// Manage all the NCCL device communicator business. Singleton.
+// Manages NCCL device communicators for symmetric memory operations.
+// This is a singleton class that maintains a registry of device communicators
+// organized by process group name and an optional key (typically the caller
+// function name). This allows different functions within the same process group
+// to use different device communicators, which is useful for concurrent
+// collective operations.
+//
+// The registry uses a two-level map structure:
+// - First level: keyed by process group name
+// - Second level: keyed by an optional key (defaults to caller function name)
+//
+// Device communicators are stored by value in the registry, but methods return
+// references wrapped in std::optional for safe access.
 class NCCLDevCommManager {
  public:
   // Constructor
+  // @param device The CUDA device this manager is associated with
   explicit NCCLDevCommManager(const c10::Device device) : device_(device) {}
 
-  // Get single, global manager.
+  // Get the singleton instance for the given device.
+  // This ensures there's only one manager per device. If called with a
+  // different device than the one used to create the singleton, it will throw.
+  // @param device The CUDA device to get the manager for
+  // @return Reference to the singleton manager instance
   static NCCLDevCommManager& get(const c10::Device device) {
     static NCCLDevCommManager manager(device);
     TORCH_CHECK_VALUE(
@@ -26,91 +47,196 @@ class NCCLDevCommManager {
     return manager;
   }
 
-  // Get an NCCL device communicator for a group.
-  ncclDevComm& get_devcomm(const std::string& group_name) {
-    auto it = group_to_comms_.find(group_name);
-    if (it == group_to_comms_.end()) {
-      TORCH_CHECK(
-          false,
-          "NCCL device communicator for group ",
-          group_name,
-          " not found. Have you rendezvoused any tensor with this group?");
+  // Get an NCCL device communicator for a group, for the caller function.  By
+  // default, we search for the device communicator using the caller function
+  // name as the key.  If you previously registered a device communicator with a
+  // different key, you should provide that key instead.
+  // Returns std::nullopt if the device communicator is not found.
+  // Example:
+  // void foo(const std::string& group_name) {
+  //   // Try to get first.
+  //   auto devcomm_opt = get_devcomm(group_name);
+  //   if (!devcomm_opt) {
+  //     // Not found, create then register.
+  //     ncclDevComm devcomm = ncclDevCommCreate(...);
+  //     devcomm_opt = register_devcomm(group_name, devcomm);
+  //   }
+  //   ncclDevComm& devcomm_ref = *devcomm_opt;
+  //   // Use devcomm_ref
+  // }
+  std::optional<std::reference_wrapper<ncclDevComm>> get_devcomm(
+      const std::string& group_name,
+      const std::string& key = __builtin_FUNCTION()) {
+    std::lock_guard<std::mutex> lock(mutex_);
+    // First, look up the group in the registry
+    auto group_it = devcomm_registry_.find(group_name);
+    if (group_it == devcomm_registry_.end()) {
+      return std::nullopt;
+    }
+    // Then, look up the key within that group's map
+    auto key_it = group_it->second.find(key);
+    if (key_it == group_it->second.end()) {
+      return std::nullopt;
     }
-    return it->second.second;
+    // Return a reference wrapper to the device communicator
+    // Using reference_wrapper because std::optional cannot hold references
+    // directly
+    return std::make_optional(std::ref(key_it->second));
   }
 
-  // Get a host-side communicator for a group.
+  // Get a host-side NCCL communicator for a group.
+  // This is the regular host-side communicator, not the device communicator.
+  // @param group_name The process group name
+  // @return The host-side NCCL communicator
+  // @throws TORCH_CHECK if the communicator is not found
   ncclComm_t get_comm(const std::string& group_name) {
-    auto it = group_to_comms_.find(group_name);
-    if (it == group_to_comms_.end()) {
+    std::lock_guard<std::mutex> lock(mutex_);
+    auto it = group_to_comm_.find(group_name);
+    if (it == group_to_comm_.end()) {
       TORCH_CHECK(
           false,
           "NCCL host communicator for group ",
           group_name,
           " not found. Have you rendezvoused any tensor with this group?");
     }
-    return it->second.first;
+    return it->second;
   }
 
-  // Create device communicator if it doesn't exist. Skip if it already exists.
-  void try_emplace_devcomm(
+  // Register a device communicator for a group. If `key` is not
+  // specified, we use the caller function name as the default `key`, to
+  // distinguish between different collective functions within the same group.
+  // You can provide your own `key` if your function uses two different
+  // device communicators on the same group at the same time, for example,
+  // when concurrent collective operations are used.
+  // Returns a reference to the newly registered device communicator.
+  // @throws TORCH_CHECK if the device communicator is already registered for
+  //         the given group and key combination.
+  // Example:
+  // void foo(const std::string& group_name) {
+  //   // Try to get first.
+  //   auto devcomm_opt = get_devcomm(group_name);
+  //   if (!devcomm_opt) {
+  //     // Not found, create then register.
+  //     ncclDevComm devcomm = ncclDevCommCreate(...);
+  //     devcomm_opt = register_devcomm(group_name, devcomm);
+  //   }
+  //   ncclDevComm& devcomm_ref = *devcomm_opt;
+  //   // Use devcomm_ref
+  // }
+  // void bar(const std::string& group_name) {
+  //   ncclDevComm devcomm0 = ncclDevCommCreate(...);
+  //   ncclDevComm devcomm1 = ncclDevCommCreate(...);
+  //   // You can provide your own `key` if you want to, for example, to
+  //   // distinguish between concurrent collective operations.
+  //   register_devcomm(group_name, devcomm0, "bar0");
+  //   register_devcomm(group_name, devcomm1, "bar1");
+  // }
+  std::optional<std::reference_wrapper<ncclDevComm>> register_devcomm(
       const std::string& group_name,
-      ncclComm_t comm,
-      int lsa_barrier_count,
-      int gin_barrier_count) {
-    auto it = group_to_comms_.find(group_name);
-    if (it != group_to_comms_.end()) {
-      return;
+      ncclDevComm devcomm,
+      const std::string& key = __builtin_FUNCTION()) {
+    std::lock_guard<std::mutex> lock(mutex_);
+    // Ensure the group exists in the registry, creating an empty map if needed
+    auto [group_it, inserted] = devcomm_registry_.try_emplace(
+        group_name, std::unordered_map<std::string, ncclDevComm>());
+    auto& group_map = group_it->second;
+    // Try to insert the device communicator with the given key
+    // Use std::move to avoid copying the device communicator
+    auto [key_it, key_inserted] =
+        group_map.try_emplace(key, std::move(devcomm));
+    if (!key_inserted) {
+      // Already registered - this is a programming error, so throw
+      TORCH_CHECK(
+          false,
+          "NCCL device communicator for group ",
+          group_name,
+          " with key ",
+          key,
+          " already registered.");
     }
-    c10::cuda::CUDAGuard guard(device_);
-    ncclDevComm devComm;
-
-    // Initializer available from NCCL 2.29
-#ifdef NCCL_DEV_COMM_REQUIREMENTS_INITIALIZER
-    ncclDevCommRequirements reqs = NCCL_DEV_COMM_REQUIREMENTS_INITIALIZER;
-#else
-    // In 2.28, we can set it to zero
-    ncclDevCommRequirements reqs;
-    memset(&reqs, 0, sizeof(ncclDevCommRequirements));
-#endif
-
-    // Specifies the number of barriers to allocate, LSA for NVLink domain, GIN
-    // for network.
-    reqs.lsaBarrierCount = lsa_barrier_count;
-    reqs.railGinBarrierCount = gin_barrier_count;
+    // Return a reference to the newly registered device communicator
+    return std::make_optional(std::ref(key_it->second));
+  }
 
-    C10D_NCCL_CHECK(
-        ncclDevCommCreate(comm, &reqs, &devComm), "ncclDevCommCreate failed");
-    // Cache the device communicator for future reuse
-    // TODO (kwen2501):
-    // Cache devComm not just based on group name, but also on requirements.
-    group_to_comms_.emplace(group_name, std::make_pair(comm, devComm));
+  // Register a host-side NCCL communicator for a group.
+  // This should be called before registering any device communicators for the
+  // same group, as device communicators need the host communicator for cleanup.
+  // @param group_name The process group name
+  // @param comm The host-side NCCL communicator to register
+  // @throws TORCH_CHECK if the group is already registered with a different
+  // communicator.
+  void register_comm(const std::string& group_name, ncclComm_t comm) {
+    std::lock_guard<std::mutex> lock(mutex_);
+    auto [it, inserted] = group_to_comm_.try_emplace(group_name, comm);
+    // If the communicator is already registered, check if it is the same one.
+    // If not, throw an error.
+    TORCH_CHECK(
+        inserted || it->second == comm, // this is just a pointer comparison
+        "NCCL host communicator for group ",
+        group_name,
+        " already registered.");
   }
 
+  // Destructor: Clean up all registered device communicators.
+  // This is a best-effort cleanup. If the CUDA context has already been
+  // destroyed, the cleanup will be skipped. All errors are caught and ignored
+  // to prevent exceptions from propagating during destruction.
   ~NCCLDevCommManager() noexcept {
     // Best effort to destroy the device communicators. Skip if CUDA context has
     // exited.
     try {
       c10::cuda::CUDAGuard guard(device_);
       // Make sure all kernels have completed before destroying the device
-      // communicator.
+      // communicator. This is important to ensure no kernels are still using
+      // the device communicator when we destroy it.
       C10_CUDA_CHECK(cudaDeviceSynchronize());
-      for (auto& [_, comm_pair] : group_to_comms_) {
-        auto& [comm, devcomm] = comm_pair;
-        ncclDevCommDestroy(comm, &devcomm);
+      // Iterate through all groups and their device communicators
+      for (auto& [group_name, group_map] : devcomm_registry_) {
+        // Find the host communicator for the group.
+        // Device communicators need the host communicator for destruction.
+        auto comm_it = group_to_comm_.find(group_name);
+        if (comm_it != group_to_comm_.end()) {
+          // Destroy each device communicator in this group
+          for (auto& [_, devcomm] : group_map) {
+            // Destroy the device communicator using the host communicator
+            ncclDevCommDestroy(comm_it->second, &devcomm);
+          }
+        }
       }
     } catch (...) {
-      // Ignore the error
-      std::cerr << "Failed to destroy the NCCL device communicator, skipping\n";
+      // Ignore the error - we're in a destructor and can't throw
+      // Log a warning for debugging purposes
+      LOG(WARNING)
+          << "Failed to destroy the NCCL device communicator, skipping";
     }
   }
 
  private:
-  // Device where the NCCL device communicator manager is created
+  // Device where the NCCL device communicator manager is created.
+  // The manager is device-specific and cannot be used across multiple devices.
   const c10::Device device_;
-  // A map from group name to NCCL device communicator for that group.
-  std::unordered_map<std::string, std::pair<ncclComm_t, ncclDevComm>>
-      group_to_comms_;
+
+  // Mutex to protect the registry maps.
+  std::mutex mutex_;
+
+  // A map from process group name to the host-side NCCL communicator.
+  // The host communicator is required for creating and destroying device
+  // communicators. It should be registered before any device communicators
+  // for the same group.
+  std::unordered_map<std::string, ncclComm_t> group_to_comm_;
+
+  // A two-level map for device communicators:
+  // - First level: keyed by process group name
+  // - Second level: keyed by an optional key (defaults to caller function name
+  //   via __builtin_FUNCTION())
+  //
+  // This structure allows multiple device communicators per process group,
+  // which is useful when different functions need separate device communicators
+  // for concurrent operations. The key defaults to the caller's function name,
+  // but can be customized for cases where a single function needs multiple
+  // device communicators.
+  std::unordered_map<std::string, std::unordered_map<std::string, ncclDevComm>>
+      devcomm_registry_;
 };
 
 } // namespace c10d::symmetric_memory
diff --git a/torch/csrc/distributed/c10d/symm_mem/nccl_extension.cu b/torch/csrc/distributed/c10d/symm_mem/nccl_extension.cu
index 5fffa98e527c2..2f37255bc24fb 100644
--- a/torch/csrc/distributed/c10d/symm_mem/nccl_extension.cu
+++ b/torch/csrc/distributed/c10d/symm_mem/nccl_extension.cu
@@ -2,7 +2,7 @@
 #include <ATen/native/cuda/MemoryAccess.cuh>
 #include <torch/csrc/distributed/c10d/NCCLUtils.hpp>
 #include <torch/csrc/distributed/c10d/symm_mem/nccl_dev_cap.hpp>
-#include <torch/csrc/distributed/c10d/symm_mem/nccl_extension.cuh>
+#include <torch/csrc/distributed/c10d/symm_mem/nccl_extension.hpp>
 #include <torch/csrc/distributed/c10d/symm_mem/nccl_devcomm_manager.hpp>
 #include <torch/csrc/distributed/c10d/symm_mem/NCCLSymmetricMemory.hpp>
 
@@ -408,6 +408,7 @@ TORCH_LIBRARY_IMPL(symm_mem, CUDA, m) {
   // API that uses internal signal mechanism and accepts handle
   m.impl("nccl_put_signal",
       torch::CppFunction::makeFromBoxedFunction<&nccl_put_signal_boxed>());
+  m.impl("nccl_reduce_scatter_offset", c10d::nccl_extension::nccl_reduce_scatter_offset);
 }
 
 // Use CompositeExplicitAutograd as key since ops do not accept tensor as input
diff --git a/torch/csrc/distributed/c10d/symm_mem/nccl_extension.cuh b/torch/csrc/distributed/c10d/symm_mem/nccl_extension.cuh
deleted file mode 100644
index be566b8deffce..0000000000000
--- a/torch/csrc/distributed/c10d/symm_mem/nccl_extension.cuh
+++ /dev/null
@@ -1,18 +0,0 @@
-#pragma once
-
-#include <c10/macros/Macros.h>
-#include <ATen/ATen.h>
-
-namespace c10d::nccl_extension {
-
-TORCH_API bool is_nccl_symmem_available();
-
-TORCH_API void nccl_put(at::Tensor& tensor, const int64_t peer);
-
-TORCH_API void nccl_get(at::Tensor& tensor, const int64_t peer);
-
-TORCH_API void nccl_wait_for_signal(at::Tensor& sigpad, int64_t signal);
-
-TORCH_API void nccl_put_with_signal(at::Tensor& tensor, int64_t signal, int64_t peer);
-
-} // namespace c10d::nccl_extension
diff --git a/torch/csrc/distributed/c10d/symm_mem/nccl_extension.hpp b/torch/csrc/distributed/c10d/symm_mem/nccl_extension.hpp
new file mode 100644
index 0000000000000..5d779f1564726
--- /dev/null
+++ b/torch/csrc/distributed/c10d/symm_mem/nccl_extension.hpp
@@ -0,0 +1,33 @@
+#pragma once
+
+#include <ATen/ATen.h>
+#include <c10/macros/Macros.h>
+
+namespace c10d::nccl_extension {
+
+TORCH_API bool is_nccl_symmem_available();
+
+TORCH_API void nccl_put(at::Tensor& tensor, const int64_t peer);
+
+TORCH_API void nccl_get(at::Tensor& tensor, const int64_t peer);
+
+TORCH_API void nccl_wait_for_signal(at::Tensor& sigpad, int64_t signal);
+
+TORCH_API void nccl_put_with_signal(
+    at::Tensor& tensor,
+    int64_t signal,
+    int64_t peer);
+
+// Simultaneously reduce N blocks of a 2-D input tensor from a shared symmetric
+// memory buffer, routing each to a specific destination rank. Blocks are
+// described by inclusive-prefix-sum offsets along `dim` (0 or 1); all blocks
+// must have equal size.
+TORCH_API void nccl_reduce_scatter_offset(
+    const at::Tensor& input,
+    at::TensorList out,
+    const std::string& group_name,
+    int64_t dim,
+    std::optional<at::IntArrayRef> offsets,
+    std::optional<at::IntArrayRef> dst_ranks,
+    const std::string& red_op);
+} // namespace c10d::nccl_extension
diff --git a/torch/csrc/distributed/c10d/symm_mem/nvshmem_extension.cu b/torch/csrc/distributed/c10d/symm_mem/nvshmem_extension.cu
index 9b9bcea522900..b766443612578 100644
--- a/torch/csrc/distributed/c10d/symm_mem/nvshmem_extension.cu
+++ b/torch/csrc/distributed/c10d/symm_mem/nvshmem_extension.cu
@@ -3,9 +3,10 @@
 #include <c10/cuda/CUDAGuard.h>
 
 #include <torch/csrc/distributed/c10d/symm_mem/env.hpp>
-#include <torch/csrc/distributed/c10d/symm_mem/nvshmem_extension.cuh>
+#include <torch/csrc/distributed/c10d/symm_mem/macros.hpp>
+#include <torch/csrc/distributed/c10d/symm_mem/nvshmem_extension.hpp>
 #include <torch/csrc/distributed/c10d/symm_mem/nvshmem_team_manager.hpp>
-#include <torch/csrc/distributed/c10d/symm_mem/CUDASymmetricMemory-inl.h>
+#include <torch/csrc/distributed/c10d/symm_mem/CUDASymmetricMemory-inl.cuh>
 #include <torch/csrc/distributed/c10d/symm_mem/CUDASymmetricMemoryUtils.hpp>
 #include <torch/csrc/distributed/c10d/symm_mem/SymmetricMemory.hpp>
 
@@ -912,20 +913,6 @@ __global__ void tile_reduce_kernel(
 #endif
 }
 
-#define AT_DISPATCH_CASE_CONVERT(enum_type, scalar_type, ...)               \
-  case enum_type: {                                                         \
-    AT_PRIVATE_CHECK_SELECTIVE_BUILD(enum_type);                            \
-    using scalar_t = scalar_type;                                           \
-    return __VA_ARGS__();                                                   \
-  }
-
-#define AT_DISPATCH_NVSHMEM_FLOATS(scalar_type, name, ...)                  \
-  AT_DISPATCH_SWITCH(                                                       \
-      scalar_type, name,                                                    \
-      AT_DISPATCH_CASE_CONVERT(at::kBFloat16, __nv_bfloat16, __VA_ARGS__);  \
-      AT_DISPATCH_CASE_CONVERT(at::kHalf, __half, __VA_ARGS__);             \
-      AT_DISPATCH_CASE(at::kFloat, __VA_ARGS__));
-
 void tile_reduce(
     at::Tensor& in_tile,
     at::Tensor& out_tile,
@@ -974,7 +961,7 @@ void tile_reduce(
       &root,
       &teams_dev};
 
-  AT_DISPATCH_NVSHMEM_FLOATS(in_tile.scalar_type(), "tile_reduce", [&]() {
+  AT_DISPATCH_NV_FLOATS(in_tile.scalar_type(), "tile_reduce", [&]() {
     nvshmemx_collective_launch(
         (const void*)tile_reduce_kernel<scalar_t>,
         dim3(nblocks),
@@ -1057,7 +1044,7 @@ void multi_root_tile_reduce(
       &root,
       &teams_dev};
 
-  AT_DISPATCH_NVSHMEM_FLOATS(out_tile.scalar_type(), "multi_root_tile_reduce", [&]() {
+  AT_DISPATCH_NV_FLOATS(out_tile.scalar_type(), "multi_root_tile_reduce", [&]() {
     nvshmemx_collective_launch(
         (const void*)tile_reduce_kernel<scalar_t>,
         dim3(nblocks),
diff --git a/torch/csrc/distributed/c10d/symm_mem/nvshmem_extension.cuh b/torch/csrc/distributed/c10d/symm_mem/nvshmem_extension.hpp
similarity index 85%
rename from torch/csrc/distributed/c10d/symm_mem/nvshmem_extension.cuh
rename to torch/csrc/distributed/c10d/symm_mem/nvshmem_extension.hpp
index 50b9e268cba7a..86f945dded460 100644
--- a/torch/csrc/distributed/c10d/symm_mem/nvshmem_extension.cuh
+++ b/torch/csrc/distributed/c10d/symm_mem/nvshmem_extension.hpp
@@ -1,7 +1,7 @@
 #pragma once
 
-#include <c10/macros/Macros.h>
 #include <ATen/ATen.h>
+#include <c10/macros/Macros.h>
 
 #define NVSHMEM_CHECK(stmt, msg)                                             \
   do {                                                                       \
@@ -25,11 +25,21 @@ TORCH_API void nvshmem_put(at::Tensor& tensor, const int64_t peer);
 
 TORCH_API void nvshmem_get(at::Tensor& tensor, const int64_t peer);
 
-at::Tensor nvshmem_broadcast(at::Tensor& input, const int64_t root, const std::string& group_name);
-
-TORCH_API void nvshmem_wait_for_signal(at::Tensor& sigpad, int64_t signal, int64_t peer);
-
-TORCH_API void nvshmem_put_with_signal(at::Tensor& tensor, at::Tensor& sigpad, int64_t signal, int64_t peer);
+at::Tensor nvshmem_broadcast(
+    at::Tensor& input,
+    const int64_t root,
+    const std::string& group_name);
+
+TORCH_API void nvshmem_wait_for_signal(
+    at::Tensor& sigpad,
+    int64_t signal,
+    int64_t peer);
+
+TORCH_API void nvshmem_put_with_signal(
+    at::Tensor& tensor,
+    at::Tensor& sigpad,
+    int64_t signal,
+    int64_t peer);
 
 at::Tensor nvshmem_all_to_all(
     at::Tensor& input,
diff --git a/torch/csrc/distributed/c10d/symm_mem/nvshmem_team_manager.hpp b/torch/csrc/distributed/c10d/symm_mem/nvshmem_team_manager.hpp
index 0deda8e8f6575..947795c9350f3 100644
--- a/torch/csrc/distributed/c10d/symm_mem/nvshmem_team_manager.hpp
+++ b/torch/csrc/distributed/c10d/symm_mem/nvshmem_team_manager.hpp
@@ -12,8 +12,12 @@
 // include only the nvshmem host library headers:
 // #include <nvshmem_host.h>
 // It translates into the following two lines:
+#if !defined(USE_ROCM)
 #include <host/nvshmem_api.h>
 #include <host/nvshmemx_api.h>
+#else
+#include <rocshmem/rocshmem.hpp>
+#endif
 // For maximum compatibility, we use the "host/" style for now.
 
 namespace c10d::nvshmem_extension {
@@ -167,4 +171,4 @@ class TeamManager {
   std::unordered_map<std::string, nvshmem_team_t*> team_pool_devptrs_;
 };
 
-} // namespace c10d::nvshmem_extension
\ No newline at end of file
+} // namespace c10d::nvshmem_extension
diff --git a/torch/csrc/distributed/c10d/symm_mem/ops/nccl_reduce_scatter_offset.cu b/torch/csrc/distributed/c10d/symm_mem/ops/nccl_reduce_scatter_offset.cu
new file mode 100644
index 0000000000000..6cd13189c36bb
--- /dev/null
+++ b/torch/csrc/distributed/c10d/symm_mem/ops/nccl_reduce_scatter_offset.cu
@@ -0,0 +1,380 @@
+#include <c10/cuda/CUDAGuard.h>
+#include <ATen/Dispatch.h>
+#include <ATen/cuda/CUDAContext.h>
+#include <torch/csrc/distributed/c10d/NCCLUtils.hpp>
+#include <torch/csrc/distributed/c10d/symm_mem/macros.hpp>
+#include <torch/csrc/distributed/c10d/symm_mem/nccl_dev_cap.hpp>
+#include <torch/csrc/distributed/c10d/symm_mem/nccl_extension.hpp>
+#include <torch/csrc/distributed/c10d/symm_mem/nccl_devcomm_manager.hpp>
+#include <torch/csrc/distributed/c10d/symm_mem/NCCLSymmetricMemory.hpp>
+
+// Simultaneously reduce N blocks of a 2-D input tensor from a symmetric memory
+// buffer, routing each block to a specific destination rank (dst_ranks[i]).
+// Only the destination rank writes the reduced value to a contiguous output
+// tensor, with the same shape as the owned block.
+//
+// The `dim` argument controls which dimension is sharded (0 or 1):
+//   dim=1 (column sharding): each block spans input[:, offsets[i-1]:offsets[i]]
+//   dim=0 (row sharding):    each block spans input[offsets[i-1]:offsets[i], :]
+//
+// Blocks are described by inclusive-prefix-sum offsets along `dim`.
+// For each j, out[j] must have the same shape across all ranks (i.e. the j-th
+// owned block on every rank must have equal size); different j's may differ.
+//
+// If offsets is nullopt, input.size(dim) is divided equally into group_size blocks.
+// If dst_ranks is nullopt, blocks are distributed round-robin across ranks.
+//
+// Ownership must be balanced: every rank must own the same number of blocks
+// (N % group_size == 0 and dst_ranks distributes evenly).
+
+namespace c10d::nccl_extension {
+
+using namespace c10d::symmetric_memory;
+
+// Kernel requires device-side API: ncclLsaReduceSum.
+#ifdef NCCL_DEVICE_HAS_REDUCE_COPY
+
+// Naming conventions in this file:
+// "BLOCK" means tensor block (as opposed to CUDA block);
+// "CTA" means CUDA block;
+// "RS" means Reduce Scatter;
+// "slot" means which tensor block a CTA is assigned to.
+
+constexpr int RS_MAX_BLOCKS = 64;           // max total blocks being scattered (N)
+constexpr int RS_MAX_BLOCKS_PER_RANK = 16;  // max blocks owned by a single rank
+constexpr int RS_MAX_CTAS_PER_BLOCK = 16;   // max CTAs assigned to one block
+// Threads per CTA; defaults to a medium value to fit medium-width blocks.
+constexpr int RS_THREADS_PER_CTA = 128;
+// Total LSA barrier slots needed: one per CTA across all owned blocks.
+constexpr int RS_MAX_CTA_COUNT = (RS_MAX_BLOCKS_PER_RANK * RS_MAX_CTAS_PER_BLOCK);
+
+// Per-slot data passed to the kernel in a single struct to avoid multiple
+// kernel arguments.  Indexed by owned slot (0..n_owned-1).
+struct ReduceScatterOffsetsInfo {
+  size_t byte_offsets[RS_MAX_BLOCKS_PER_RANK]; // byte offset into the NCCL window
+  void* dst_ptrs[RS_MAX_BLOCKS_PER_RANK];      // output pointer (contiguous)
+  uint16_t dst_block_size[RS_MAX_BLOCKS_PER_RANK]; // per-slot size along the sharding dim
+  uint16_t ctas_offset[RS_MAX_BLOCKS_PER_RANK]; // inclusive prefix sum of per-slot CTA counts
+  uint8_t cta_slot[RS_MAX_CTA_COUNT];          // slot index for each flat CTA
+  int n_owned;
+};
+
+// Grid: 1D, total_ctas = sum of per-slot CTA counts (info.ctas_offset[n_owned]).
+// Each CTA belongs to one slot; blockIdx.x is the flat CTA index used as the
+// LSA barrier index, ensuring all ranks assign the same index to each logical
+// (slot, local_block) pair (because owned_sizes[j] is consistent across ranks).
+//
+// UseMultimem=true: uses ncclMultimemReduceSum for hardware reduction via
+// NVLink multicast; requires devcomm created with lsaMultimem=true.
+// UseMultimem=false: uses ncclLsaReduceSum (software reduce via LSA reads).
+template <typename T, bool UseMultimem>
+__global__ void reduce_scatter_offset_kernel(
+    ncclWindow_t window,
+    ReduceScatterOffsetsInfo info,
+    int fixed_dim_size,   // input.size(1-dim): constant across all slots
+    bool col_sharded,     // true when dim==1
+    int64_t outer_stride, // row stride of the input buffer (in elements)
+    ncclDevComm devComm) {
+  // cta_slot maps the flat CTA index to its owned slot.
+  const int slot = info.cta_slot[blockIdx.x];
+  // ctas_offset is an inclusive prefix sum, so slot_start is the flat index
+  // of the first CTA assigned to this slot.
+  const int slot_start = slot > 0 ? info.ctas_offset[slot - 1] : 0;
+  // local_block is this CTA's position within its slot (0-based row tile index).
+  const int local_block = static_cast<int>(blockIdx.x) - slot_start;
+  // Number of CTAs sharing this slot; used as the row-loop stride.
+  const int ctas_for_slot = info.ctas_offset[slot] - slot_start;
+  const ncclCoopCta coop{};
+
+  // One LSA barrier per CTA; all ranks must call both syncs unconditionally.
+  ncclLsaBarrierSession<ncclCoopCta> bar{
+      coop,
+      devComm,
+      ncclTeamLsa(devComm),
+      devComm.lsaBarrier,
+      blockIdx.x};
+  // Acquire: wait until all peers have written their data into the window.
+  bar.sync(coop, cuda::memory_order_acquire);
+
+  const size_t base_byte_offset = info.byte_offsets[slot]; // start of this block in the window
+  T* dst_base = reinterpret_cast<T*>(info.dst_ptrs[slot]); // start of out[slot]
+  const int block_size = info.dst_block_size[slot]; // size along the sharding dim
+  const int rows = col_sharded ? fixed_dim_size : block_size;
+  const int cols = col_sharded ? block_size : fixed_dim_size;
+
+  // Each CTA handles a strided subset of rows; the reduce reads from all peers
+  // and writes cols elements starting at dst_row.
+  for (int row = local_block; row < rows; row += ctas_for_slot) {
+    const size_t row_offset =
+        base_byte_offset +
+        static_cast<size_t>(row * outer_stride) * sizeof(T);
+    T* dst_row = dst_base + row * cols;
+    if constexpr (UseMultimem) {
+      ncclMultimemReduceSum(
+          coop, window, row_offset, dst_row, cols, devComm.lsaMultimem);
+    } else {
+      ncclLsaReduceSum(coop, window, row_offset, dst_row, cols, devComm);
+    }
+  }
+
+  // Release: signal peers that we are done reading window memory.
+  bar.sync(coop, cuda::memory_order_release);
+}
+
+#endif // NCCL_DEVICE_HAS_REDUCE_COPY
+
+// Host entry point.  Validates arguments, resolves defaults, builds the
+// per-slot ReduceScatterOffsetsInfo, and launches the kernel.
+// See file-level comment for semantics.
+void nccl_reduce_scatter_offset(
+    const at::Tensor& input,
+    at::TensorList out,
+    const std::string& group_name,
+    int64_t dim,
+    std::optional<at::IntArrayRef> offsets,
+    std::optional<at::IntArrayRef> dst_ranks,
+    const std::string& red_op) {
+#ifdef NCCL_DEVICE_HAS_REDUCE_COPY
+  TORCH_CHECK(
+      red_op == "sum",
+      "nccl_reduce_scatter_offset: only red_op='sum' is supported, got '", red_op, "'");
+
+  TORCH_CHECK(
+      input.dim() == 2,
+      "nccl_reduce_scatter_offset: input must be 2-D");
+  TORCH_CHECK(
+      dim == 0 || dim == 1,
+      "nccl_reduce_scatter_offset: dim must be 0 or 1, got ", dim);
+  TORCH_CHECK(
+      input.stride(-1) == 1,
+      "nccl_reduce_scatter_offset: innermost dimension must be contiguous "
+      "(stride[-1] == 1)");
+
+  // rendezvous retrieves the symmetric memory handle; the tensor must have
+  // been allocated via empty_strided_p2p with the NCCL backend.
+  auto symm_mem = c10d::symmetric_memory::rendezvous(input, group_name);
+  TORCH_CHECK(
+      symm_mem != nullptr,
+      "nccl_reduce_scatter_offset: input must be allocated via NCCL symmetric "
+      "memory (use empty_strided_p2p with NCCL backend)");
+
+  auto* nccl_hdl = dynamic_cast<NCCLSymmetricMemory*>(symm_mem.get());
+  TORCH_CHECK(
+      nccl_hdl != nullptr,
+      "nccl_reduce_scatter_offset: requires NCCL symmetric memory backend");
+
+  c10::cuda::CUDAGuard guard(input.device());
+  auto stream = at::cuda::getCurrentCUDAStream();
+  auto device = input.device();
+
+  auto& manager = c10d::symmetric_memory::NCCLDevCommManager::get(device);
+  // Get the host-side communicator.
+  ncclComm_t comm = manager.get_comm(group_name);
+
+  const bool use_multimem = nccl_hdl->has_multicast_support();
+
+  // The devcomm is cached per (group, key); create it on first use.
+  // lsaBarrierCount must cover the maximum number of concurrent CTAs.
+  // lsaMultimem is set when the allocation has multicast support, so that
+  // devComm.lsaMultimem is valid for ncclMultimemReduceSum in the kernel.
+  static constexpr char const kDevcommKey[] = "nccl_reduce_scatter_offset";
+  auto devcomm_opt = manager.get_devcomm(group_name, kDevcommKey);
+  if (!devcomm_opt) {
+    ncclDevCommRequirements reqs = NCCL_DEV_COMM_REQUIREMENTS_INITIALIZER;
+    reqs.lsaBarrierCount = RS_MAX_CTA_COUNT;
+    reqs.lsaMultimem = use_multimem;
+    ncclDevComm devcomm;
+    C10D_NCCL_CHECK(
+        ncclDevCommCreate(comm, &reqs, &devcomm),
+        "ncclDevCommCreate failed in nccl_reduce_scatter_offset");
+    // Cache the device communicator.
+    devcomm_opt = manager.register_devcomm(group_name, devcomm, kDevcommKey);
+  }
+  ncclDevComm& devcomm = devcomm_opt->get();
+
+  const int my_rank = devcomm.rank;
+  const int group_size = devcomm.nRanks;
+
+  // Determine n_blocks: from offsets if given, else group_size (equal-size default).
+  const int n_blocks = offsets.has_value()
+      ? static_cast<int>(offsets->size())
+      : group_size;
+  TORCH_CHECK(
+      n_blocks > 0,
+      "nccl_reduce_scatter_offset: must have at least one block");
+
+  // Fill dst_ranks default: round-robin across ranks.
+  std::vector<int64_t> dst_ranks_vec;
+  at::IntArrayRef effective_dst_ranks;
+  if (dst_ranks.has_value()) {
+    effective_dst_ranks = *dst_ranks;
+  } else {
+    dst_ranks_vec.resize(n_blocks);
+    for (int i = 0; i < n_blocks; i++) {
+      dst_ranks_vec[i] = i % group_size;
+    }
+    effective_dst_ranks = at::IntArrayRef(dst_ranks_vec);
+  }
+
+  // Fill offsets default: divide input.size(dim) equally among group_size blocks.
+  std::vector<int64_t> offsets_vec;
+  at::IntArrayRef effective_offsets;
+  if (offsets.has_value()) {
+    effective_offsets = *offsets;
+    TORCH_CHECK(
+        effective_offsets[n_blocks - 1] <= input.size(dim),
+        "nccl_reduce_scatter_offset: offsets exceed input size along dim ", dim);
+  } else {
+    const int64_t total = input.size(dim);
+    TORCH_CHECK(
+        total % group_size == 0,
+        "nccl_reduce_scatter_offset: input.size(", dim, ")=", total,
+        " must be divisible by group size (", group_size, ")");
+    const int64_t block_size = total / group_size;
+    offsets_vec.resize(n_blocks);
+    for (int i = 0; i < n_blocks; i++) {
+      offsets_vec[i] = (i + 1) * block_size;
+    }
+    effective_offsets = at::IntArrayRef(offsets_vec);
+  }
+
+  TORCH_CHECK(
+      n_blocks <= RS_MAX_BLOCKS,
+      "nccl_reduce_scatter_offset: too many blocks: ", n_blocks,
+      " (max ", RS_MAX_BLOCKS, ")");
+  TORCH_CHECK(
+      static_cast<int>(effective_dst_ranks.size()) == n_blocks,
+      "nccl_reduce_scatter_offset: dst_ranks.size() must match offsets.size()");
+
+  const int64_t outer_stride = input.stride(0);
+
+  // Collect owned blocks (in order).
+  std::vector<int> owned_indices;
+  for (int i = 0; i < n_blocks; i++) {
+    if (static_cast<int>(effective_dst_ranks[i]) == my_rank) {
+      owned_indices.push_back(i);
+    }
+  }
+  const int n_owned = static_cast<int>(owned_indices.size());
+  TORCH_CHECK(
+      n_owned * group_size == n_blocks,
+      "nccl_reduce_scatter_offset: dst_ranks must distribute blocks evenly "
+      "(rank owns ", n_owned, "/", n_blocks, ", group_size=", group_size, ")");
+  TORCH_CHECK(
+      n_owned <= RS_MAX_BLOCKS_PER_RANK,
+      "nccl_reduce_scatter_offset: too many owned blocks: ", n_owned,
+      " (max ", RS_MAX_BLOCKS_PER_RANK, ")");
+  // Balance is guaranteed above (n_owned * group_size == n_blocks), so
+  // rank_counter[r] never exceeds n_owned during the owned_sizes loop.
+
+  // For each j, out[j] must have the same shape across all ranks.  That means
+  // all blocks that are the j-th owned block on their respective rank must have
+  // equal size.  Different j's may differ in size.
+  //
+  // Compute the size for each j by iterating all blocks in order, tracking
+  // how many blocks each rank has seen so far (= the j-index for that block).
+  std::vector<int64_t> owned_sizes(n_owned, -1);
+  {
+    std::vector<int> rank_counter(group_size, 0);
+    for (int i = 0; i < n_blocks; i++) {
+      const int r = static_cast<int>(effective_dst_ranks[i]);
+      const int j = rank_counter[r]++;
+      const int64_t sz =
+          effective_offsets[i] - (i > 0 ? effective_offsets[i - 1] : 0);
+      if (owned_sizes[j] < 0) {
+        owned_sizes[j] = sz;
+      } else {
+        TORCH_CHECK(
+            sz == owned_sizes[j],
+            "nccl_reduce_scatter_offset: all output at position j=", j,
+            " must have equal size across all ranks");
+      }
+    }
+  }
+
+  TORCH_CHECK(
+      static_cast<int>(out.size()) == n_owned,
+      "nccl_reduce_scatter_offset: out.size() must be ", n_owned);
+  for (int j = 0; j < n_owned; j++) {
+    // dim=1: out[j] shape is (input.size(0), owned_sizes[j])
+    // dim=0: out[j] shape is (owned_sizes[j], input.size(1))
+    const int64_t exp0 = dim == 1 ? input.size(0) : owned_sizes[j];
+    const int64_t exp1 = dim == 1 ? owned_sizes[j] : input.size(1);
+    TORCH_CHECK(
+        out[j].size(0) == exp0 && out[j].size(1) == exp1,
+        "nccl_reduce_scatter_offset: out[", j, "] must have shape (",
+        exp0, ", ", exp1, ")");
+    TORCH_CHECK(
+        out[j].is_contiguous(),
+        "nccl_reduce_scatter_offset: out[", j, "] must be contiguous");
+    TORCH_CHECK(
+        out[j].scalar_type() == input.scalar_type(),
+        "nccl_reduce_scatter_offset: out[", j, "] must have the same dtype as input");
+  }
+
+  // Per-slot CTA count: sized for each slot independently.  owned_sizes[j] is
+  // consistent across ranks, so ctas_offset is identical on every rank, which
+  // guarantees all ranks launch the same total CTA count and agree on the
+  // flat barrier index for each (slot, local_block) pair.
+  const bool col_sharded = (dim == 1);
+  const int fixed_dim_size = static_cast<int>(col_sharded ? input.size(0) : input.size(1));
+  const int unroll = 4 * 16 / static_cast<int>(input.element_size());
+  const int elems_per_cta = RS_THREADS_PER_CTA * unroll;
+  const size_t window_base_offset = nccl_hdl->get_offset();
+
+  // Build the per-slot info struct.
+  // For dim=1: byte_offsets encodes the column-block start within the window.
+  // For dim=0: byte_offsets encodes the row-block start within the window.
+  ReduceScatterOffsetsInfo info;
+  info.n_owned = n_owned;
+  for (int j = 0; j < n_owned; j++) {
+    const int i = owned_indices[j];
+    const int64_t block_start = (i > 0 ? effective_offsets[i - 1] : 0);
+    const size_t elem_offset = col_sharded
+        ? static_cast<size_t>(input.storage_offset() + block_start)
+        : static_cast<size_t>(input.storage_offset()) +
+              static_cast<size_t>(block_start) * outer_stride;
+    info.byte_offsets[j] = window_base_offset + elem_offset * input.element_size();
+    info.dst_ptrs[j] = out[j].data_ptr();
+    info.dst_block_size[j] = static_cast<uint16_t>(owned_sizes[j]);
+    const int numel_j = static_cast<int>(owned_sizes[j]) * fixed_dim_size;
+    const int ctas_j = std::max(1, std::min(
+        (numel_j + elems_per_cta - 1) / elems_per_cta, RS_MAX_CTAS_PER_BLOCK));
+    info.ctas_offset[j] = static_cast<uint16_t>((j > 0 ? info.ctas_offset[j - 1] : 0) + ctas_j);
+    const int slot_start = j > 0 ? info.ctas_offset[j - 1] : 0;
+    for (int k = slot_start; k < info.ctas_offset[j]; ++k) {
+      info.cta_slot[k] = static_cast<uint8_t>(j);
+    }
+  }
+  const int total_ctas = info.ctas_offset[n_owned - 1];
+
+  auto window = nccl_hdl->get_window();
+  TORCH_CHECK(window != nullptr, "nccl_reduce_scatter_offset: NCCL window is null");
+
+  // Each owned (slot, local_block) pair gets one CTA; the flat CTA index is
+  // the LSA barrier index.  All ranks launch the same total_ctas because
+  // owned_sizes[j] is consistent, so every rank's ctas_offset is identical.
+  AT_DISPATCH_NV_FLOATS(
+      input.scalar_type(),
+      "nccl_reduce_scatter_offset",
+      [&]() {
+        if (use_multimem) {
+          reduce_scatter_offset_kernel<scalar_t, true>
+              <<<total_ctas, RS_THREADS_PER_CTA, 0, stream>>>(
+                  window, info, fixed_dim_size, col_sharded, outer_stride, devcomm);
+          C10_CUDA_KERNEL_LAUNCH_CHECK();
+        } else {
+          reduce_scatter_offset_kernel<scalar_t, false>
+              <<<total_ctas, RS_THREADS_PER_CTA, 0, stream>>>(
+                  window, info, fixed_dim_size, col_sharded, outer_stride, devcomm);
+          C10_CUDA_KERNEL_LAUNCH_CHECK();
+        }
+      });
+#else
+  TORCH_CHECK(
+      false,
+      "nccl_reduce_scatter_offset requires NCCL >= 2.29.7 with reduce copy support");
+#endif // NCCL_DEVICE_HAS_REDUCE_COPY
+}
+
+} // namespace c10d::nccl_extension
diff --git a/torch/csrc/distributed/c10d/symm_mem/rocshmem_extension.cu b/torch/csrc/distributed/c10d/symm_mem/rocshmem_extension.cu
new file mode 100644
index 0000000000000..4fc6b70863fbc
--- /dev/null
+++ b/torch/csrc/distributed/c10d/symm_mem/rocshmem_extension.cu
@@ -0,0 +1,904 @@
+// ROCm implementation of the NVSHMEM symmetric memory extension ops.
+//
+// This is a separate file from nvshmem_extension.cu (rather than a hipified
+// copy) for the following reasons:
+//
+// 1. API differences: NVSHMEM and rocSHMEM device APIs diverge enough that
+//    #ifdef'ing would be more noise than signal. Key differences include:
+//    - nvshmemx_collective_launch (grid-wide sync) has no rocSHMEM equivalent;
+//      ROCm uses regular hip kernel launches with host-side barriers instead.
+//    - nvshmemx_getmem_nbi_block → rocshmem_getmem_nbi_wg (workgroup scope).
+//
+// 2. Missing features: rocSHMEM does not yet support tiled communication
+//    (nvshmemx::Tensor, nvshmemx::tile_sum_reduce_block, etc.), so the
+//    tile_reduce and multi_root_tile_reduce ops are not included here.
+//
+// 3. Offset writeback: without grid-wide sync, multi-block kernels cannot
+//    safely write output offsets in-kernel (race with blocks still reading
+//    source_offsets). A separate writeOutputOffsets kernel runs after the data
+//    exchange completes.
+
+#include <hip/hip_runtime.h>
+#include <algorithm>
+#include <cstdlib>
+#include <vector>
+#include <ATen/ceil_div.h>
+#include <c10/hip/HIPGuard.h>
+
+#include <torch/csrc/distributed/c10d/symm_mem/env.hpp>
+#include <torch/csrc/distributed/c10d/symm_mem/nvshmem_extension.hpp>
+#include <torch/csrc/distributed/c10d/symm_mem/nvshmem_team_manager.hpp>
+#include <torch/csrc/distributed/c10d/symm_mem/CUDASymmetricMemory-inl.cuh>
+#include <torch/csrc/distributed/c10d/symm_mem/CUDASymmetricMemoryUtils.hpp>
+#include <torch/csrc/distributed/c10d/symm_mem/SymmetricMemory.hpp>
+
+#include <ATen/hip/cub.cuh>
+
+#include <c10/hip/HIPException.h>
+#include <rocshmem/rocshmem.hpp>
+
+using namespace rocshmem;
+namespace c10d::nvshmem_extension {
+
+#define THREADS_PER_BLOCK 512
+#define WARP_SIZE 64
+
+namespace {
+
+bool parse_rocshmem_version_ge(
+    const char* version,
+    unsigned min_major,
+    unsigned min_minor,
+    unsigned min_patch) {
+  if (version == nullptr) {
+    return false;
+  }
+  char* end = nullptr;
+  unsigned long major = std::strtoul(version, &end, 10);
+  if (end == version || *end != '.') {
+    return false;
+  }
+  version = end + 1;
+  unsigned long minor = std::strtoul(version, &end, 10);
+  if (end == version || *end != '.') {
+    return false;
+  }
+  version = end + 1;
+  unsigned long patch = std::strtoul(version, &end, 10);
+  if (end == version) {
+    return false;
+  }
+  if (major > min_major) {
+    return true;
+  }
+  if (major < min_major) {
+    return false;
+  }
+  if (minor > min_minor) {
+    return true;
+  }
+  if (minor < min_minor) {
+    return false;
+  }
+  return patch >= min_patch;
+}
+
+} // namespace
+
+extern "C" void rocshmem_init() __attribute__((weak));
+
+bool is_nvshmem_available() {
+  static const bool ok =
+      parse_rocshmem_version_ge(rocshmem::VERSION, 3, 3, 0);
+  return ok;
+}
+
+void nvshmemx_cumodule_init(uintptr_t module) {
+  auto hipmodule = reinterpret_cast<hipModule_t>(module);
+  NVSHMEM_CHECK(
+    rocshmem_hipmodule_init(hipmodule),
+    "rocshmem_hipmodule_init failed");
+}
+
+at::Tensor nvshmem_broadcast(at::Tensor& input, const int64_t root, const std::string& group_name) {
+  auto input_hdl = c10d::symmetric_memory::rendezvous(input, group_name);
+  int rank = input_hdl->get_rank();
+  void* buffer_ptr = input.mutable_data_ptr();
+  auto buffer_size = input.numel() * input.element_size();
+  auto& team_manager = TeamManager::get(input.device());
+  auto team = team_manager.get_team(group_name, input_hdl->get_rank_to_global_rank());
+  int team_size = rocshmem_team_n_pes(team);
+  TORCH_CHECK(root < team_size, "root must be smaller than group size");
+
+  auto stream = at::cuda::getCurrentCUDAStream();
+  rocshmem_broadcastmem_on_stream(team, buffer_ptr, buffer_ptr, buffer_size, root, stream);
+  return input;
+}
+
+void nvshmem_put(at::Tensor& tensor, const int64_t peer) {
+  // TODO: support non-contiguous tensors
+  TORCH_CHECK(tensor.is_contiguous(),
+      "put op currently supports contiguous tensors only");
+  // TODO: rendezvous should remember the group name
+  auto hdl = c10d::symmetric_memory::rendezvous(tensor, "0");
+  auto rank = hdl->get_rank();
+  void* buffer_ptr = hdl->get_buffer_ptrs()[rank];
+  auto buffer_size = tensor.numel() * tensor.element_size();
+  TORCH_CHECK(peer < hdl->get_world_size(), "peer must be smaller than world size");
+
+  c10::cuda::CUDAGuard guard(tensor.device());
+  auto stream = at::cuda::getCurrentCUDAStream();
+  rocshmem_putmem_on_stream(buffer_ptr, tensor.data_ptr(), buffer_size, peer, stream);
+}
+
+void nvshmem_wait_for_signal(at::Tensor& sigpad, int64_t signal, int64_t peer) {
+  c10::cuda::CUDAGuard guard(sigpad.device());
+  auto stream = at::cuda::getCurrentCUDAStream();
+  rocshmem_signal_wait_until_on_stream(static_cast<uint64_t*>(sigpad.data_ptr()), ROCSHMEM_CMP_EQ, signal, stream);
+}
+
+void nvshmem_put_with_signal(at::Tensor& tensor, at::Tensor& sigpad, int64_t signal, int64_t peer) {
+  auto buffer_size = tensor.numel() * tensor.element_size();
+
+  c10::cuda::CUDAGuard guard(tensor.device());
+  auto stream = at::cuda::getCurrentCUDAStream();
+  rocshmem_putmem_signal_on_stream(
+    tensor.mutable_data_ptr(),
+    tensor.mutable_data_ptr(),
+    buffer_size,
+    static_cast<uint64_t*>(sigpad.mutable_data_ptr()),
+    signal,
+    ROCSHMEM_SIGNAL_SET,
+    peer,
+    stream);
+}
+
+void nvshmem_get(at::Tensor& tensor, const int64_t peer) {
+  // TODO: support non-contiguous tensors
+  TORCH_CHECK(tensor.is_contiguous(),
+      "get op currently supports contiguous tensors only");
+  // TODO: rendezvous should remember the group name
+  auto hdl = c10d::symmetric_memory::rendezvous(tensor, "0");
+  auto rank = hdl->get_rank();
+  void* buffer_ptr = hdl->get_buffer_ptrs()[rank];
+  auto buffer_size = tensor.numel() * tensor.element_size();
+  TORCH_CHECK(peer < hdl->get_world_size(), "peer must be smaller than world size");
+
+  c10::cuda::CUDAGuard guard(tensor.device());
+  auto stream = at::cuda::getCurrentCUDAStream();
+  rocshmem_getmem_on_stream(tensor.mutable_data_ptr(), buffer_ptr, buffer_size, peer, stream);
+}
+
+at::Tensor nvshmem_all_to_all(
+    at::Tensor& input,
+    at::Tensor& out,
+    std::string group_name) {
+  auto input_hdl = c10d::symmetric_memory::rendezvous(input, group_name);
+  auto out_hdl = c10d::symmetric_memory::rendezvous(out, group_name);
+  int rank = input_hdl->get_rank();
+  int world_size = input_hdl->get_world_size();
+  auto& team_manager = TeamManager::get(input.device());
+  auto team = team_manager.get_team(group_name, input_hdl->get_rank_to_global_rank());
+
+  void* input_ptr = input.data_ptr();
+  void* output_ptr = out.mutable_data_ptr();
+  TORCH_CHECK(input.is_contiguous() && out.is_contiguous());
+  TORCH_CHECK_EQ(input.numel(), out.numel());
+  TORCH_CHECK_EQ(input.dtype(), out.dtype());
+  TORCH_CHECK_EQ(input.numel() % world_size, 0);
+  auto buffer_size = input.numel() * input.element_size();
+  size_t bytes_per_rank = buffer_size / world_size;
+
+  auto stream = at::cuda::getCurrentCUDAStream(input.device().index());
+  rocshmem_alltoallmem_on_stream(team, output_ptr, input_ptr, bytes_per_rank, stream);
+  return out;
+}
+
+// This is an exclusive prefix sum function that calculates read (or write) offsets for each peer.
+__device__ int64_t prefixSum(int64_t *odata, int64_t *idata, int n) {
+  // Specialize BlockScan for a 1D block of threads, of type int64_t.
+  // - `BLOCK_SCAN_WARP_SCANS` is a low-latency scan algorithm (instead of high
+  // throughput which we don't need here).
+  // - `at_cuda_detail::cub` is torch's cub wrapper, see #55292.
+  using BlockScanT = ROCM_HIPCUB(at_cuda_detail::cub)::BlockScan<int64_t,
+        THREADS_PER_BLOCK, ROCM_HIPCUB(at_cuda_detail::cub)::BLOCK_SCAN_WARP_SCANS>;
+  // Allocate shared memory for BlockScan
+  __shared__ typename BlockScanT::TempStorage temp_storage;
+
+  // TODO: currently it is assumed that the number of PE's is smaller than
+  // `THREADS_PER_BLOCK`
+  CUDA_KERNEL_ASSERT(n <= THREADS_PER_BLOCK);
+
+  // Obtain input item for each thread
+  int tid = threadIdx.x;
+  int64_t thread_data = (tid < n) ? idata[tid] : 0;
+
+  // Collectively compute the block-wide exclusive prefix sum
+  int64_t block_aggregate;
+  BlockScanT(temp_storage).ExclusiveSum(thread_data, thread_data, block_aggregate);
+
+  // Store the result
+  odata[tid] = thread_data;
+  return block_aggregate;
+}
+
+static int get_a2a_nblocks(size_t size, int world_size, bool intra_node) {
+  // Check user setting first
+  int num_blocks = c10d::symmetric_memory::getenv_nblocks();
+  if (num_blocks > 0) {  // set by user
+    return num_blocks;
+  }
+  // 16B per thread, 8 loops
+  constexpr size_t chunk_size = 16 * THREADS_PER_BLOCK * 8;
+  num_blocks = at::ceil_div(size, chunk_size);
+  // Allow kernel to target even number of blocks per peer
+  num_blocks = at::round_up(num_blocks, world_size);
+  const int max_blocks = intra_node ? 64 : 16;
+  return ::min(num_blocks, max_blocks);
+}
+
+// ROCm-only offset writeback kernel.
+//
+// On ROCm, allToAllV is a regular multi-block kernel with no grid-wide barrier.
+// Writing source_offsets in-kernel can race with other blocks still reading it for
+// remote gets. We therefore compute output offsets in a separate kernel after
+// allToAllV has completed on the stream.
+__global__ void writeOutputOffsets1d(int64_t* out_splits_offsets, int world_size) {
+  auto output_splits = out_splits_offsets;
+  auto output_offsets = out_splits_offsets + world_size;
+  int tid = threadIdx.x;
+
+  CUDA_KERNEL_ASSERT(world_size <= THREADS_PER_BLOCK);
+  __shared__ int64_t peer_offsets[THREADS_PER_BLOCK];
+  prefixSum(peer_offsets, output_splits, world_size);
+  __syncthreads();
+
+  if (tid < world_size) {
+    output_offsets[tid] = peer_offsets[tid];
+  }
+}
+// This kernel is used to exchange output splits and source offsets between peers.
+// `in_out_splits` is of size (3, npes) and contains:
+// - input splits (IN)
+// - output splits (OUT) and
+// - source offsets (OUT).
+__global__ void exchangeSplitAndOffset(int64_t* input_splits, int64_t* out_splits_offsets, rocshmem_team_t team) {
+  CUDA_KERNEL_ASSERT(team != ROCSHMEM_TEAM_INVALID);
+  int mype = rocshmem_team_my_pe(team);
+  int npes = rocshmem_team_n_pes(team);
+  auto output_splits = out_splits_offsets;
+  auto source_offsets = out_splits_offsets + npes;
+  int tid = threadIdx.x;
+
+  CUDA_KERNEL_ASSERT(npes <= THREADS_PER_BLOCK);
+  __shared__ int64_t peer_offsets[THREADS_PER_BLOCK];
+
+  // Scan input splits to get the source offsets
+  prefixSum(peer_offsets, input_splits, npes);
+  __syncthreads();;
+
+  // Use 1 block to do the exchange
+  if (tid < npes) {
+    // tid is peer index within team, but put calls require global rank
+    int peer_global = rocshmem_team_translate_pe(team, tid, ROCSHMEM_TEAM_WORLD);
+    rocshmem_int64_p(source_offsets + mype, peer_offsets[tid], peer_global);
+    rocshmem_int64_p(output_splits + mype, input_splits[tid], peer_global);
+  }
+  rocshmem_barrier_wg();
+}
+
+// This kernel is used to do the actual data exchange.
+// `in_out_splits` has the same definition as in `exchangeSplitAndOffset`.
+// `stride` is the stride at dim 0, unit in byte.
+__global__ void allToAllV(void *send_data, void *recv_data, int64_t* out_splits_offsets, size_t stride, rocshmem_team_t team) {
+  CUDA_KERNEL_ASSERT(team != ROCSHMEM_TEAM_INVALID);
+  int mype = rocshmem_team_my_pe(team);
+  int npes = rocshmem_team_n_pes(team);
+  auto output_splits = out_splits_offsets;
+  auto source_offsets = out_splits_offsets + npes;
+  int bid = blockIdx.x;
+  int tid = threadIdx.x;
+  int blocks_per_peer = max(gridDim.x / npes, 1);
+
+  // Calculate the output offsets
+  CUDA_KERNEL_ASSERT(npes <= THREADS_PER_BLOCK);
+  __shared__ int64_t peer_offsets[THREADS_PER_BLOCK];
+  prefixSum(peer_offsets, output_splits, npes);
+  __syncthreads();
+
+  // Target a different peer based on bid
+  for (int i = bid / blocks_per_peer; i < npes; i += gridDim.x / blocks_per_peer) {
+    int peer = (mype + i) % npes;
+    auto peer_global = rocshmem_team_translate_pe(team, peer, ROCSHMEM_TEAM_WORLD);
+    // Total amount from `peer`
+    auto peer_size = output_splits[peer] * stride;
+    // Amount to get from `peer` in this block
+    auto block_size = peer_size / blocks_per_peer;
+    // Being lazy here, we should handle the residual if the division is not exact
+    CUDA_KERNEL_ASSERT(block_size * blocks_per_peer == peer_size);
+    // This block's offset in the data from `peer`
+    auto block_offset = block_size * (bid % blocks_per_peer);
+    auto source_offset = source_offsets[peer] * stride + block_offset;
+    auto write_offset = peer_offsets[peer] * stride + block_offset;
+    rocshmem_getmem_nbi_wg(
+      (char*)recv_data + write_offset,
+      (char*)send_data + source_offset,
+      block_size,
+      peer_global);
+  }
+  rocshmem_quiet();
+}
+
+void all_to_all_vdev(
+    at::Tensor& input,
+    at::Tensor& out,
+    at::Tensor& in_splits,
+    at::Tensor& out_splits_offsets,
+    std::string group_name) {
+  /* Perform AllToAllv operation using NVSHMEM, with split information provided on device.
+   * Step 1: Rendezvous tensors so all ranks have symmetric (device) pointers.
+   * Step 2: Launch exchangeSplitAndOffset kernel to exchange per-rank split counts
+   *         and compute source offsets (prefix sum); uses team barrier.
+   * Step 3: Launch allToAllV kernel to copy data between peers according to
+   *         the exchanged splits/offsets.
+   * Arguments:
+   *  - `input` is the send buffer
+   *  - `out` is the receive buffer
+   *  - `in_splits`: 1D[npes] num of elements this rank sends to each peer
+   *  - `out_splits_offsets`:2D (2, npes). row0 = output splits, row1 = output offsets
+   */
+  auto input_hdl = c10d::symmetric_memory::rendezvous(input, group_name);
+  auto out_hdl = c10d::symmetric_memory::rendezvous(out, group_name);
+  auto in_splits_hdl = c10d::symmetric_memory::rendezvous(in_splits, group_name);
+  auto out_splits_offsets_hdl = c10d::symmetric_memory::rendezvous(out_splits_offsets, group_name);
+  int world_size = input_hdl->get_world_size();
+
+  void* input_ptr = input.data_ptr();
+  void* output_ptr = out.mutable_data_ptr();
+  int64_t* in_splits_ptr = (int64_t*)(in_splits.const_data_ptr());
+  int64_t* out_splits_offsets_ptr = (int64_t*)(out_splits_offsets.mutable_data_ptr());
+
+  TORCH_CHECK_EQ(input.device(), out.device());
+  auto device = input.device();
+  c10::cuda::CUDAGuard guard(device);
+  auto& team_manager = TeamManager::get(device);
+  auto team = team_manager.get_team(group_name, input_hdl->get_rank_to_global_rank());
+  auto stream = at::cuda::getCurrentCUDAStream(device.index());
+
+  exchangeSplitAndOffset<<<dim3(1), dim3(THREADS_PER_BLOCK), 0, stream>>>(
+      in_splits_ptr, out_splits_offsets_ptr, team);
+  C10_CUDA_KERNEL_LAUNCH_CHECK();
+  C10_CUDA_CHECK(hipStreamSynchronize(stream));
+  rocshmem::rocshmem_barrier_all();
+  // CTA Tuning
+  auto input_size = input.numel() * input.element_size();
+  int num_blocks = get_a2a_nblocks(
+    input_size,
+    input_hdl->get_world_size(),
+    input_hdl->world_within_direct_access());
+
+  // Stride at dim 0 (assuming input is contiguous, TODO)
+  size_t stride_bytes = input.stride(0) * input.element_size();
+
+  allToAllV<<<dim3(num_blocks), dim3(THREADS_PER_BLOCK), 0, stream>>>(
+      input_ptr, output_ptr, out_splits_offsets_ptr,
+      stride_bytes, team);
+  C10_CUDA_KERNEL_LAUNCH_CHECK();
+  C10_CUDA_CHECK(hipStreamSynchronize(stream));
+  // `allToAllV` reads source_offsets while fetching remote shards. Since ROCm has
+  // no grid-wide sync here, writing output offsets in the same kernel can race
+  // with those reads. Write output offsets in a follow-up kernel instead.
+  writeOutputOffsets1d<<<dim3(1), dim3(THREADS_PER_BLOCK), 0, stream>>>(
+      out_splits_offsets_ptr, world_size);
+  C10_CUDA_KERNEL_LAUNCH_CHECK();
+}
+
+// Start of `all_to_all_vdev_2d`
+
+// This is an warp-scope, exclusive prefix sum. When called by a block of
+// threads, each warp will perform an independent prefix sum, concurrently.
+// Returns the sum of all elements in the warp.
+// `NUM_WARPS` is the number of warps participating the concurrent prefix sum.
+template <int NUM_WARPS>
+__device__ int64_t prefixSum_warp(int64_t *odata, int64_t *idata, int n) {
+  CUDA_KERNEL_ASSERT(n <= WARP_SIZE);
+
+  // Specialize WarpScan for type int
+  using WarpScan = ROCM_HIPCUB(at_cuda_detail::cub)::WarpScan<int64_t>;
+  // Allocate WarpScan shared memory for N warps
+  __shared__ typename WarpScan::TempStorage temp_storage[NUM_WARPS];
+
+  int warp_id = threadIdx.x / WARP_SIZE;
+  if (warp_id >= NUM_WARPS) {
+    return 0;
+  }
+
+  // Obtain input item for each thread
+  int tid = threadIdx.x % WARP_SIZE;
+  int64_t thread_data = (tid < n) ? idata[tid] : 0;
+
+  // Total sum of all elements in the warp
+  int64_t warp_aggregate;
+  // Compute the warp-wide exclusive prefix sum
+  WarpScan(temp_storage[warp_id]).ExclusiveSum(thread_data, thread_data, warp_aggregate);
+
+  // Store only valid lanes to avoid out-of-bounds writes when n < WARP_SIZE.
+  if (tid < n) {
+    odata[tid] = thread_data;
+  }
+  return warp_aggregate;
+}
+
+// This is for abstracting a thread-group-scope, exclusive prefix sum.
+// Since we use warp-scope prefix sum, the thread group size is limited to warp size.
+#define A2AV_TILE_SIZE WARP_SIZE
+
+
+__global__ void writeOutputOffsets_2d(
+    int64_t* out_splits_offsets,
+    int minor_size,
+    int major_size,
+    int64_t major_align) {
+  int nsplits = minor_size * major_size;
+  auto output_splits = out_splits_offsets;
+  auto source_offsets = out_splits_offsets + nsplits;
+  int tid = threadIdx.x;
+
+  constexpr int NUM_TILES = THREADS_PER_BLOCK / A2AV_TILE_SIZE;
+  int tileId = tid / A2AV_TILE_SIZE;
+  int laneId = tid % A2AV_TILE_SIZE;
+  __shared__ int64_t tile_prefix_sums[NUM_TILES][A2AV_TILE_SIZE];
+  int nsplits_per_tile = min(minor_size, nsplits - tileId * minor_size);
+
+  __shared__ int64_t len_per_tile[NUM_TILES];
+  if (nsplits_per_tile > 0) {
+    int64_t my_tile_len = prefixSum_warp<NUM_TILES>(tile_prefix_sums[tileId], output_splits + tileId * minor_size, nsplits_per_tile);
+    if (laneId == A2AV_TILE_SIZE - 1) {
+      if (major_align != 0) {
+        auto aligned_len = (my_tile_len + major_align - 1) / major_align * major_align;
+        len_per_tile[tileId] = max(aligned_len, major_align);
+      } else {
+        len_per_tile[tileId] = my_tile_len;
+      }
+    }
+  }
+  __syncthreads();
+
+  __shared__ int64_t start_offset_per_tile[WARP_SIZE];
+  prefixSum_warp<1>(start_offset_per_tile, len_per_tile, NUM_TILES);
+  __syncthreads();
+
+  tile_prefix_sums[tileId][laneId] += start_offset_per_tile[tileId];
+  __syncthreads();
+
+  if (tid < nsplits) {
+    source_offsets[tid] = tile_prefix_sums[tid / minor_size][tid % minor_size];
+  }
+}
+
+// `exchangeSplitAndOffset_2d` is used to exchange output splits and source
+// offsets between peers.
+
+/* Arguments:
+ * `in_splits_offsets`: input splits and offsets (optional), of size (2, nsplits), or (1, nsplits) if no offsets are provided.
+ * `out_splits_offsets`: output splits and offsets, of size (2, nsplits).
+ * `mype`: the rank of the current PE.
+ * `npes`: the number of PEs.
+ * `ne`: the number of experts.
+ * `input_dim0`: the size of dim 0 of the input tensor.
+ * `rank_is_row_in` is a boolean flag indicating whether the input has ranks as row or experts as row.
+*/
+
+/* Template parameters:
+ * `HAS_IN_OFFSETS` is a boolean flag indicating whether `in_splits_offsets` has offsets (2nd row) or not.
+*/
+
+template <bool HAS_IN_OFFSETS>
+__global__ void exchangeSplitAndOffset_2d(int64_t* in_splits_offsets, int64_t* out_splits_offsets, rocshmem_team_t team, int ne, size_t input_dim0, bool rank_is_row_in) {
+  CUDA_KERNEL_ASSERT(team != ROCSHMEM_TEAM_INVALID);
+  int mype = rocshmem_team_my_pe(team);
+  int npes = rocshmem_team_n_pes(team);
+  int nsplits = npes * ne;
+  auto input_splits = in_splits_offsets;
+  auto output_splits = out_splits_offsets;
+  // Borrowing the space below as a temporary exchange pad.
+  auto source_offsets = out_splits_offsets + nsplits;
+  int tid = threadIdx.x;
+
+  int64_t* input_offsets = nullptr;
+  if (HAS_IN_OFFSETS) {
+    // input offset are provided, so we can use them directly
+    input_offsets = in_splits_offsets + nsplits;
+  } else {
+    // input offset are not provided, so we need to calculate them.
+    // Scan input splits to get the source offsets
+    __shared__ int64_t peer_offsets[THREADS_PER_BLOCK];
+    auto sum_of_splits = prefixSum(peer_offsets, input_splits, nsplits);
+    __syncthreads();;
+    CUDA_KERNEL_ASSERT(sum_of_splits <= input_dim0 && "sum of splits is larger than input dim\n");
+    // Redirect the input splits to the calculated result
+    input_offsets = peer_offsets;
+  }
+
+  // Use 1 block to do the exchange
+  if (tid < nsplits) {
+    int peer, e, dst_offset;
+    if (rank_is_row_in) {
+      peer = tid / ne;
+      e = tid % ne;
+      dst_offset = e * npes + mype;
+    } else {  // expert is row in input
+      peer = tid % npes;
+      e = tid / npes;
+      dst_offset = mype * ne + e;
+    }
+    // This does a transpose from rank-major order to expert-major order
+    // (or vice versa).
+    auto split_val = input_splits[tid];
+    CUDA_KERNEL_ASSERT(split_val >= 0 && "split value is negative\n");
+    auto peer_global = rocshmem_team_translate_pe(team, peer, ROCSHMEM_TEAM_WORLD);
+    rocshmem_int64_p(source_offsets + dst_offset, input_offsets[tid], peer_global);
+    rocshmem_int64_p(output_splits + dst_offset, split_val, peer_global);
+  }
+  rocshmem_barrier_wg();
+}
+
+
+// This kernel is used to do the actual data exchange.
+// `in_out_splits` has the same definition as in `exchangeSplitAndOffset`.
+// `stride` is the stride at dim 0, unit in byte.
+// For meaning of `mype` and `npes`, see the docstring of `all_to_all_vdev_2d`.
+// `major_align` is the alignment at dim 0, unit in element. If 0, no alignment is needed.
+
+// `rank_is_row_out` is a boolean flag indicating whether the output has ranks as rows or experts as rows.
+// In dispatch case, rank_is_row_out = false, major_size = ne, minor_size = npes.
+// In combine case, rank_is_row_out = true, major_size = npes, minor_size = ne.
+
+__global__ void allToAllV_2d(void *send_data, void *recv_data, int64_t* in_splits, int64_t* out_splits_offsets, size_t stride, int minor_size, int major_size, int64_t major_align, bool rank_is_row_out, rocshmem_team_t team) {
+  int nsplits = minor_size * major_size;
+  auto output_splits = out_splits_offsets;
+  auto source_offsets = out_splits_offsets + nsplits;
+  int bid = blockIdx.x;
+  int tid = threadIdx.x;
+
+  // Split the thread block into tiles
+  constexpr int NUM_TILES = THREADS_PER_BLOCK / A2AV_TILE_SIZE;
+  int tileId = tid / A2AV_TILE_SIZE;
+  int laneId = tid % A2AV_TILE_SIZE;
+  // Each tile calculates its own prefix sum
+  __shared__ int64_t tile_prefix_sums[NUM_TILES][A2AV_TILE_SIZE];
+  // A tile takes care of minor_size worth of splits
+  int nsplits_per_tile = min(minor_size, nsplits - tileId * minor_size);
+  // TODO: currently it is assumed that the number of PE's is smaller than
+  // `A2AV_TILE_SIZE` bc the warp-scope prefix sum can only handle up to
+  // WARP_SIZE elements
+  CUDA_KERNEL_ASSERT(minor_size <= A2AV_TILE_SIZE && "minor_size is too large\n");
+  // Similarly, the number of experts per rank is also assumed to be smaller
+  // than `NUM_TILES`
+  CUDA_KERNEL_ASSERT(major_size <= NUM_TILES && "major_size is too large\n");
+
+  // Total length of each tile
+  __shared__ int64_t len_per_tile[NUM_TILES];
+  // When `nsplits` is small, not every tile gets data to sum. They can skip
+  // this local prefix sum.
+  if (nsplits_per_tile > 0) {
+    // Each tile calculates its own prefix sum, return value is the sum of all elements in the tile.
+    int64_t my_tile_len = prefixSum_warp<NUM_TILES>(tile_prefix_sums[tileId], output_splits + tileId * minor_size, nsplits_per_tile);
+    // Last thread in each tile does the up aligning.
+    if (laneId == A2AV_TILE_SIZE - 1) {
+      if (major_align != 0) {  // Needs alignment
+        auto aligned_len = (my_tile_len + major_align - 1) / major_align * major_align;
+        // In case `aligned_len` is 0, we set it to `major_align` to avoid an
+        // empty bin, bc cutlass currently does not support it. See
+        // https://github.com/pytorch/pytorch/issues/152668.
+        len_per_tile[tileId] = max(aligned_len, major_align);
+      } else {  // 0 means alignment not needed
+        len_per_tile[tileId] = my_tile_len;
+      }
+    }
+  }
+  __syncthreads();
+
+  // Starting offset of each tile
+  __shared__ int64_t start_offset_per_tile[NUM_TILES];
+  // Prefix sum again to get the tiles' start offsets.
+  // `NUM_TILES` is typically not greater than 32, because 32 tiles * 32 threads
+  // = 1024 threads, and this kernel is launched within 1024 threads. Thus, we
+  // can use warp-scope prefix sum.
+  static_assert(NUM_TILES <= WARP_SIZE);
+  // Only 1 warp is needed
+  prefixSum_warp<1>(start_offset_per_tile, len_per_tile, NUM_TILES);
+  __syncthreads();
+
+  // Add tile offset to every element in the tile
+  tile_prefix_sums[tileId][laneId] += start_offset_per_tile[tileId];
+  __syncthreads();
+
+  // Target a different e based on bid
+  for (int eid = bid; eid < nsplits; eid += gridDim.x) {
+    int row = eid / minor_size;
+    int col = eid % minor_size;
+    // Amount from `peer` for `e`
+    auto peer_size = output_splits[eid] * stride;
+    auto source_offset = source_offsets[eid] * stride;
+    auto e_offset = tile_prefix_sums[row][col];
+    auto write_offset = e_offset * stride;
+    auto peer_global = rocshmem_team_translate_pe(team, rank_is_row_out ? row : col, ROCSHMEM_TEAM_WORLD);
+    rocshmem_getmem_nbi_wg(
+      (char*)recv_data + write_offset,
+      (char*)send_data + source_offset,
+      peer_size,
+      peer_global);  // peer's global index
+  }
+  rocshmem_quiet();
+}
+
+void all_to_all_vdev_2d(
+    at::Tensor& input,
+    at::Tensor& out,
+    at::Tensor& in_splits,
+    at::Tensor& out_splits_offsets,
+    std::string group_name,
+    std::optional<int64_t> major_align) {
+  /* Perform a 2D AllToAllv shuffle operation using NVSHMEM, with split information provided on device.
+   * Arguments:
+   *  - `input` is the input tensor
+   *  - `out` is the output tensor
+   *  - `in_out_splits` is a 2D tensor of size (3, `world_size` * `ne`). In the
+        scenario of Mixture-of-Experts models, `ne` is the number of experts per
+        rank. The rows of `in_out_splits` are (in order):
+        input splits (IN)
+        output splits (OUT) and
+        output offsets (OUT).
+   *  - `group_name` is the name of the group to use for the collective operation.
+   *  - `major_align` is the alignment of the "major dimension" of the output
+        sequence. See below for details.
+
+   *  A 2D AllToAllv shuffle is illustrated below:
+        (world_size = 2, ne = 2, total number of experts = 4)
+        Source: |       Rank 0      |       Rank 1      |
+                | c0 | c1 | c2 | c3 | d0 | d1 | d2 | d3 |
+
+        Dest  : |       Rank 0      |       Rank 1      |
+                | c0 | d0 | c1 | d1 | c2 | d2 | c3 | d3 |
+        where each `c_i` / `d_i` are slices of the `input` tensor, targeting
+        expert `i`, with length indicated by input splits (in
+        `in_out_splits[0]`).  That is, the 2D AllToAllv shuffle achieves a
+        transpose from rank-major order at input to expert-major order at
+        output.
+
+   *  If `major_align` is not 1, the output offsets of c1, c2, c3 will be
+      up-aligned to this value. For example, if c0 has length 5 and d0 has
+      length 7 (making a total of 12), and if the `major_align` is set to 16,
+      the output offset of c1 will be 16. Similar for c2 and c3. This value has
+      no effect on the offset of the minor dimension, i.e.  d0, d1, d2 and d3.
+      Note: since cutlass does not support empty bins, we set the aligned length
+      to `major_align` if it is 0. See
+      https://github.com/pytorch/pytorch/issues/152668.
+  */
+  auto input_hdl = c10d::symmetric_memory::rendezvous(input, group_name);
+  auto out_hdl = c10d::symmetric_memory::rendezvous(out, group_name);
+  auto in_splits_hdl = c10d::symmetric_memory::rendezvous(in_splits, group_name);
+  auto out_splits_offsets_hdl = c10d::symmetric_memory::rendezvous(out_splits_offsets, group_name);
+  int world_size = input_hdl->get_world_size();
+  // TODO: world_size is currently limited by the number of elements in a WarpScan.
+  TORCH_CHECK(world_size <= A2AV_TILE_SIZE, "world_size must be smaller than A2AV_TILE_SIZE", A2AV_TILE_SIZE);
+
+  // If `major_align` is not provided, use 1 as the default value.
+  int64_t major_align_val = major_align.value_or(1);
+  TORCH_CHECK(major_align_val > 0, "major_align must be positive");
+
+  void* input_ptr = input.data_ptr();
+  void* output_ptr = out.mutable_data_ptr();
+  int64_t* in_splits_ptr = (int64_t*)(in_splits.data_ptr());
+  int64_t* out_splits_offsets_ptr = (int64_t*)(out_splits_offsets.mutable_data_ptr());
+
+  // Shape checks
+  TORCH_CHECK(in_splits.is_contiguous()
+      && out_splits_offsets.is_contiguous()
+      && input.is_contiguous()
+      && out.is_contiguous(),
+      "input, out, in_splits and out_splits_offsets must be contiguous");
+  auto in_split_shape = in_splits.sizes();
+  auto out_split_shape = out_splits_offsets.sizes();
+  TORCH_CHECK(out_split_shape.size() == 2
+      && out_split_shape[0] == 2
+      && out_split_shape[1] == in_split_shape[0]
+      && in_split_shape[0] % world_size == 0,
+      "out_splits_offsets must be 2D with 2 rows, "
+      "each row must be a multiple of world_size");
+
+  // Consistency checks
+  TORCH_CHECK(input.dtype() == out.dtype()
+      && input.stride(0) == out.stride(0),
+      "input and out must have the same dtype and same stride at dim 0");
+  TORCH_CHECK(in_splits.scalar_type() == at::kLong
+      && out_splits_offsets.scalar_type() == at::kLong,
+      "splits and offsets must be int64");
+
+  // Number of experts per rank
+  int ne = in_split_shape[0] / world_size;
+  constexpr int NUM_TILES = THREADS_PER_BLOCK / A2AV_TILE_SIZE;
+  TORCH_CHECK(ne <= NUM_TILES, "Number of experts must be smaller than NUM_TILES", NUM_TILES);
+
+  // Set device context for getting the stream and launching kernels below
+  auto device = input.device();
+  TORCH_CHECK(device.type() == at::DeviceType::CUDA &&
+      out.device() == device &&
+      in_splits.device() == device &&
+      out_splits_offsets.device() == device,
+      "all tensor arguments must be on the same CUDA device");
+  c10::cuda::CUDAGuard guard(device);
+  auto stream = at::cuda::getCurrentCUDAStream();
+  auto& team_manager = TeamManager::get(device);
+  auto team = team_manager.get_team(group_name, input_hdl->get_rank_to_global_rank());
+
+  // Exchange output splits and source offsets
+  auto input_dim0 = input.size(0);
+  bool rank_is_row_in = true;
+  exchangeSplitAndOffset_2d<false><<<dim3(1), dim3(THREADS_PER_BLOCK), 0, stream>>>(
+      in_splits_ptr, out_splits_offsets_ptr, team,
+      ne, input_dim0, rank_is_row_in);
+  C10_CUDA_KERNEL_LAUNCH_CHECK();
+  C10_CUDA_CHECK(hipStreamSynchronize(stream));
+  rocshmem::rocshmem_barrier_all();
+  // CTA Tuning
+  // Naive for now, use 1 block per expert.
+  // Total number of blocks is limited to 64 (intra-node) or 8 (inter-node).
+  int num_blocks = ::min(world_size * ne, world_size > 8 ? 8 : 64);
+
+  // Stride at dim 0
+  size_t stride_bytes = input.stride(0) * input.element_size();
+  bool rank_is_row_out = !rank_is_row_in;
+
+  allToAllV_2d<<<dim3(num_blocks), dim3(THREADS_PER_BLOCK), 0, stream>>>(
+      input_ptr, output_ptr,
+      in_splits_ptr, out_splits_offsets_ptr,
+      stride_bytes, world_size,
+      ne, major_align_val, rank_is_row_out, team);
+  C10_CUDA_KERNEL_LAUNCH_CHECK();
+  // allToAllV_2d uses a regular multi-block launch with no grid-wide sync, so
+  // in-kernel writeback can race with other blocks still reading source_offsets.
+  writeOutputOffsets_2d<<<dim3(1), dim3(THREADS_PER_BLOCK), 0, stream>>>(
+      out_splits_offsets_ptr, world_size, ne, major_align_val);
+  C10_CUDA_KERNEL_LAUNCH_CHECK();
+  C10_CUDA_CHECK(hipStreamSynchronize(stream));
+}
+
+void all_to_all_vdev_2d_offset(
+    at::Tensor& input,
+    at::Tensor& out,
+    at::Tensor& in_splits_offsets,
+    at::Tensor& out_splits_offsets,
+    std::string group_name) {
+  /* Perform a 2D AllToAllv shuffle operation, with input split and offset
+   * information provided on device. The input offsets are not required to be
+   * exact prefix sum of the input splits, i.e. paddings are allowed between the
+   * split chunks. The paddings, however, will not be transferred to peer
+   * ranks.
+
+   * In Mixture of Experts models, this operation can be used to combine tokens
+   * processed by experts on parallel ranks. This operation can be viewed as an
+   * "reverse" operation to the `all_to_all_vdev_2d` operation (which shuffles
+   * tokens to experts).
+
+   * Arguments:
+   *  - `input` is the input tensor
+   *  - `out` is the output tensor
+   *  - `in_splits_offsets` is a 2D tensor of size (2, `ne` * `world_size`). In the
+        scenario of Mixture-of-Experts models, `ne` is the number of experts per
+        rank. The rows of `in_splits_offsets` are (in order):
+        input splits (IN) and
+        input offsets (IN)
+   *  - `out_splits_offsets` is a 2D tensor of size (2, `world_size` * `ne`). The
+        rows are (in order):
+        output splits (OUT) and
+        output offsets (OUT).
+   *  - `group_name` is the name of the group to use for the collective operation.
+  */
+  auto input_hdl = c10d::symmetric_memory::rendezvous(input, group_name);
+  auto out_hdl = c10d::symmetric_memory::rendezvous(out, group_name);
+  auto out_splits_offsets_hdl = c10d::symmetric_memory::rendezvous(out_splits_offsets, group_name);
+  auto in_splits_offsets_hdl = c10d::symmetric_memory::rendezvous(in_splits_offsets, group_name);
+  int rank = input_hdl->get_rank();
+  int world_size = input_hdl->get_world_size();
+  constexpr int NUM_TILES = THREADS_PER_BLOCK / A2AV_TILE_SIZE;
+  TORCH_CHECK(world_size <= NUM_TILES, "world_size must be smaller than NUM_TILES", NUM_TILES);
+
+  int64_t major_align_val = 0;
+
+  void* input_ptr = input.data_ptr();
+  void* output_ptr = out.mutable_data_ptr();
+  int64_t* out_splits_offsets_ptr = (int64_t*)(out_splits_offsets.mutable_data_ptr());
+  int64_t* in_splits_offsets_ptr = (int64_t*)(in_splits_offsets.data_ptr());
+
+  // Shape checks
+  TORCH_CHECK(out_splits_offsets.is_contiguous()
+      && in_splits_offsets.is_contiguous()
+      && input.is_contiguous()
+      && out.is_contiguous(),
+      "input, out, in_splits_offsets and out_splits_offsets must be contiguous");
+  auto out_split_shape = out_splits_offsets.sizes();
+  auto in_split_shape = in_splits_offsets.sizes();
+  TORCH_CHECK(in_split_shape.size() == 2
+      && in_split_shape[0] == 2
+      && in_split_shape[1] % world_size == 0,
+      "in_splits_offsets must be 2D with 2 rows, "
+      "each row must be a multiple of world_size");
+
+  // Consistency checks
+  TORCH_CHECK(input.dtype() == out.dtype()
+      && input.stride(0) == out.stride(0),
+      "input and out must have the same dtype and same stride at dim 0");
+  TORCH_CHECK(out_splits_offsets.scalar_type() == at::kLong
+      && in_splits_offsets.scalar_type() == at::kLong,
+      "splits and offsets must be int64");
+
+  // Number of experts per rank
+  int ne = in_split_shape[1] / world_size;
+  // TODO: number of experts is currently limited by the number of elements in a WarpScan.
+  TORCH_CHECK(ne <= A2AV_TILE_SIZE, "Number of experts must be smaller than A2AV_TILE_SIZE", A2AV_TILE_SIZE);
+
+  // Set device context for getting the stream and launching kernels below
+  auto device = input.device();
+  TORCH_CHECK(device.type() == at::DeviceType::CUDA &&
+      out.device() == device &&
+      in_splits_offsets.device() == device &&
+      out_splits_offsets.device() == device,
+      "all tensor arguments must be on the same CUDA device");
+  c10::cuda::CUDAGuard guard(device);
+  auto stream = at::cuda::getCurrentCUDAStream();
+  auto& team_manager = TeamManager::get(device);
+  auto team = team_manager.get_team(group_name, input_hdl->get_rank_to_global_rank());
+
+  // Exchange output splits and source offsets
+  auto input_dim0 = input.size(0);
+  bool rank_is_row_in = false;
+  exchangeSplitAndOffset_2d<true><<<dim3(1), dim3(THREADS_PER_BLOCK), 0, stream>>>(
+      in_splits_offsets_ptr,
+      out_splits_offsets_ptr,
+      team,
+      ne, input_dim0, rank_is_row_in);
+  C10_CUDA_KERNEL_LAUNCH_CHECK();
+  C10_CUDA_CHECK(hipStreamSynchronize(stream));
+  rocshmem::rocshmem_barrier_all();
+  // CTA Tuning
+  // Naive for now, use 1 block per expert.
+  // Total number of blocks is limited to 64 (intra-node) or 8 (inter-node).
+  int num_blocks = ::min(world_size * ne, world_size > 8 ? 8 : 64);
+
+  // Stride at dim 0
+  size_t stride_bytes = input.stride(0) * input.element_size();
+  bool rank_is_row_out = !rank_is_row_in;
+
+  allToAllV_2d<<<dim3(num_blocks), dim3(THREADS_PER_BLOCK), 0, stream>>>(
+      input_ptr,
+      output_ptr,
+      in_splits_offsets_ptr,
+      out_splits_offsets_ptr,
+      stride_bytes,
+      ne,
+      world_size,
+      major_align_val,
+      rank_is_row_out,
+      team);
+  C10_CUDA_KERNEL_LAUNCH_CHECK();
+  writeOutputOffsets_2d<<<dim3(1), dim3(THREADS_PER_BLOCK), 0, stream>>>(
+      out_splits_offsets_ptr, ne, world_size, major_align_val);
+  C10_CUDA_KERNEL_LAUNCH_CHECK();
+  C10_CUDA_CHECK(hipStreamSynchronize(stream));
+}
+
+} // namespace c10d::nvshmem_extension
+
+TORCH_LIBRARY_IMPL(symm_mem, CUDA, m) {
+  m.impl("nvshmem_broadcast", c10d::nvshmem_extension::nvshmem_broadcast);
+  m.impl("nvshmem_put", c10d::nvshmem_extension::nvshmem_put);
+  m.impl("nvshmem_get", c10d::nvshmem_extension::nvshmem_get);
+  m.impl("nvshmem_wait_for_signal", c10d::nvshmem_extension::nvshmem_wait_for_signal);
+  m.impl("nvshmem_put_with_signal", c10d::nvshmem_extension::nvshmem_put_with_signal);
+  m.impl("nvshmem_all_to_all", c10d::nvshmem_extension::nvshmem_all_to_all);
+  m.impl("all_to_all_vdev", c10d::nvshmem_extension::all_to_all_vdev);
+  m.impl("all_to_all_vdev_2d", c10d::nvshmem_extension::all_to_all_vdev_2d);
+  m.impl("all_to_all_vdev_2d_offset", c10d::nvshmem_extension::all_to_all_vdev_2d_offset);
+}
diff --git a/torch/csrc/distributed/python_placement.cpp b/torch/csrc/distributed/python_placement.cpp
index 697e118157eb9..533a29e39e19c 100644
--- a/torch/csrc/distributed/python_placement.cpp
+++ b/torch/csrc/distributed/python_placement.cpp
@@ -21,86 +21,112 @@ This class is not meant to be used directly, mainly served as a typing stub.
 void initPlacementBindings(PyObject* module) {
   auto py_module = py::reinterpret_borrow<py::module>(module);
   auto distributed_module = py_module.def_submodule("_distributed");
-  py::class_<Placement>(
-      distributed_module, "Placement", placement_class_docstring)
-      .def(py::init<>()) // Allow construction of Python subclasses.
-      .def(
-          "is_partial",
-          &Placement::is_partial,
-          py::arg("reduce_op") = py::none())
-      .def("is_replicate", &Placement::is_replicate)
-      .def("is_shard", &Placement::is_shard, py::arg("dim") = py::none());
-  py::class_<Shard, Placement>(distributed_module, "Shard")
-      .def(py::init<int64_t>(), py::arg("dim"))
-      .def_readonly("dim", &Shard::dim)
-      .def("is_shard", &Shard::is_shard, py::arg("dim") = py::none())
-      .def(
-          "__eq__",
-          [](const Shard& lhs, const Shard& rhs) { return lhs == rhs; },
-          py::is_operator())
-      // Note: we need to use dicts for pickling to match the old
-      // dataclasses.
-      .def(py::pickle(
-          [](const Shard& shard) { return py::dict("dim"_a = shard.dim); },
-          [](const py::dict& d) {
-            return Shard(py::cast<int64_t>(d["dim"]));
-          }));
-  py::class_<StridedShard, Placement>(distributed_module, "StridedShard")
-      .def(
-          py::init<int64_t, int64_t>(),
-          py::arg("dim"),
-          py::kw_only(),
-          py::arg("split_factor"))
-      .def_readonly("dim", &StridedShard::dim)
-      .def_readonly("split_factor", &StridedShard::split_factor)
-      .def(
-          "__eq__",
-          [](const StridedShard& lhs, const StridedShard& rhs) {
-            return lhs == rhs;
-          },
-          py::is_operator())
-      .def(py::pickle(
-          [](const StridedShard& shard) {
-            return py::dict(
-                "dim"_a = shard.dim, "split_factor"_a = shard.split_factor);
-          },
-          [](const py::dict& d) {
-            return StridedShard(
-                py::cast<int64_t>(d["dim"]),
-                py::cast<int64_t>(d["split_factor"]));
-          }));
-  py::class_<Replicate, Placement>(distributed_module, "Replicate")
-      .def(py::init())
-      .def("is_replicate", &Replicate::is_replicate)
-      .def(
-          "__eq__",
-          [](const Replicate& lhs, const Replicate& rhs) { return lhs == rhs; },
-          py::is_operator())
-      .def(py::pickle(
-          // I observed SIGSEGV when trying to use None as the
-          // pickled state, though AFAICT that matches the
-          // behavior of
-          // object().__reduce__().
-          // test_placement_types.test_type_identification will repro if an
-          // enterprising reader wants to get this fixed.
-          [](const Replicate& repl) { return py::dict(); },
-          [](const py::dict&) { return Replicate(); }));
-  py::class_<Partial, Placement>(distributed_module, "Partial")
-      .def(py::init<>())
-      .def(py::init<std::optional<std::string>>(), py::arg("reduce_op"))
-      .def_readonly("reduce_op", &Partial::reduce_op)
-      .def(
-          "is_partial", &Partial::is_partial, py::arg("reduce_op") = py::none())
-      .def(
-          "__eq__",
-          [](const Partial& lhs, const Partial& rhs) { return lhs == rhs; },
-          py::is_operator())
-      .def(py::pickle(
-          [](const Partial& part) {
-            return py::dict("reduce_op"_a = part.reduce_op);
-          },
-          [](const py::dict& d) {
-            return Partial(py::cast<std::string>(d["reduce_op"]));
-          }));
+
+  // Use OpaqueBase as the metaclass to allow isinstance(fake_obj, Placement) to
+  // work.
+  py::object opaque_base_module = py::module_::import("torch._opaque_base");
+  py::object opaque_base = opaque_base_module.attr("OpaqueBaseMeta");
+
+  auto placement_cls =
+      py::class_<Placement>(
+          distributed_module,
+          "Placement",
+          py::metaclass(opaque_base),
+          placement_class_docstring)
+          .def(py::init<>()) // Allow construction of Python subclasses.
+          .def(
+              "is_partial",
+              &Placement::is_partial,
+              py::arg("reduce_op") = py::none())
+          .def("is_replicate", &Placement::is_replicate)
+          .def("is_shard", &Placement::is_shard, py::arg("dim") = py::none());
+
+  auto shard_cls =
+      py::class_<Shard, Placement>(
+          distributed_module, "Shard", py::metaclass(opaque_base))
+          .def(py::init<int64_t>(), py::arg("dim"))
+          .def_readonly("dim", &Shard::dim)
+          .def("is_shard", &Shard::is_shard, py::arg("dim") = py::none())
+          .def(
+              "__eq__",
+              [](const Shard& lhs, const Shard& rhs) { return lhs == rhs; },
+              py::is_operator())
+          // Note: we need to use dicts for pickling to match the old
+          // dataclasses.
+          .def(py::pickle(
+              [](const Shard& shard) { return py::dict("dim"_a = shard.dim); },
+              [](const py::dict& d) {
+                return Shard(py::cast<int64_t>(d["dim"]));
+              }));
+
+  auto strided_shard_cls =
+      py::class_<StridedShard, Placement>(
+          distributed_module, "StridedShard", py::metaclass(opaque_base))
+          .def(
+              py::init<int64_t, int64_t>(),
+              py::arg("dim"),
+              py::kw_only(),
+              py::arg("split_factor"))
+          .def_readonly("dim", &StridedShard::dim)
+          .def_readonly("split_factor", &StridedShard::split_factor)
+          .def(
+              "__eq__",
+              [](const StridedShard& lhs, const StridedShard& rhs) {
+                return lhs == rhs;
+              },
+              py::is_operator())
+          .def(py::pickle(
+              [](const StridedShard& shard) {
+                return py::dict(
+                    "dim"_a = shard.dim, "split_factor"_a = shard.split_factor);
+              },
+              [](const py::dict& d) {
+                return StridedShard(
+                    py::cast<int64_t>(d["dim"]),
+                    py::cast<int64_t>(d["split_factor"]));
+              }));
+
+  auto replicate_cls =
+      py::class_<Replicate, Placement>(
+          distributed_module, "Replicate", py::metaclass(opaque_base))
+          .def(py::init())
+          .def("is_replicate", &Replicate::is_replicate)
+          .def(
+              "__eq__",
+              [](const Replicate& lhs, const Replicate& rhs) {
+                return lhs == rhs;
+              },
+              py::is_operator())
+          .def(py::pickle(
+              // I observed SIGSEGV when trying to use None as the
+              // pickled state, though AFAICT that matches the
+              // behavior of
+              // object().__reduce__().
+              // test_placement_types.test_type_identification will repro if an
+              // enterprising reader wants to get this fixed.
+              [](const Replicate& repl) { return py::dict(); },
+              [](const py::dict&) { return Replicate(); }));
+
+  auto partial_cls =
+      py::class_<Partial, Placement>(
+          distributed_module, "Partial", py::metaclass(opaque_base))
+          .def(py::init<>())
+          .def(py::init<std::optional<std::string>>(), py::arg("reduce_op"))
+          .def_readonly("reduce_op", &Partial::reduce_op)
+          .def(
+              "is_partial",
+              &Partial::is_partial,
+              py::arg("reduce_op") = py::none())
+          .def(
+              "__eq__",
+              [](const Partial& lhs, const Partial& rhs) { return lhs == rhs; },
+              py::is_operator())
+          .def(py::pickle(
+              [](const Partial& part) {
+                return py::dict("reduce_op"_a = part.reduce_op);
+              },
+              [](const py::dict& d) {
+                return Partial(py::cast<std::string>(d["reduce_op"]));
+              }));
 }
 } // namespace torch::distributed
diff --git a/torch/csrc/distributed/rpc/init.cpp b/torch/csrc/distributed/rpc/init.cpp
index ce85ee4f5c4ba..09a8ca0555b55 100644
--- a/torch/csrc/distributed/rpc/init.cpp
+++ b/torch/csrc/distributed/rpc/init.cpp
@@ -393,28 +393,18 @@ PyObject* rpc_init(PyObject* _unused, PyObject* noargs) {
           .def(
               py::pickle(
                   /* __getstate__ */
-                  [](const PyRRef& /* unused */) {
+                  [](const PyRRef& /* unused */) -> py::tuple {
                     TORCH_CHECK(
                         false,
                         "Can not pickle rref in python pickler, rref can only be "
                         "pickled when using RPC");
-                    // Note that this return has no meaning since we always
-                    // throw, it's only here to satisfy Pybind API's
-                    // requirement.
-                    return py::make_tuple();
                   },
                   /* __setstate__ */
-                  [](py::tuple /* unused */) { // NOLINT
+                  [](py::tuple /* unused */) -> std::nullptr_t { // NOLINT
                     TORCH_CHECK(
                         false,
                         "Can not unpickle rref in python pickler, rref can only be "
                         "unpickled when using RPC");
-                    // Note that this return has no meaning since we always
-                    // throw, it's only here to satisfy PyBind's API
-                    // requirement.
-                    return PyRRef(
-                        py::cast<py::none>(Py_None),
-                        py::cast<py::none>(Py_None));
                   }),
               py::call_guard<py::gil_scoped_release>())
           .def(
diff --git a/torch/csrc/distributed/rpc/tensorpipe_agent.cpp b/torch/csrc/distributed/rpc/tensorpipe_agent.cpp
index e20f8730b6ec0..84358a8d2cb13 100644
--- a/torch/csrc/distributed/rpc/tensorpipe_agent.cpp
+++ b/torch/csrc/distributed/rpc/tensorpipe_agent.cpp
@@ -1265,23 +1265,14 @@ void TensorPipeAgent::updateGroupMembership(
     workerNameToURL_.erase(name);
 
     // remove reverse device maps that are no longer used
-    for (auto it = reverseDeviceMaps_.begin();
-         it != reverseDeviceMaps_.end();) {
-      if (reverseDeviceMaps.find(it->first) == reverseDeviceMaps.end()) {
-        it = reverseDeviceMaps_.erase(it);
-      } else {
-        it++;
-      }
-    }
+    std::erase_if(reverseDeviceMaps_, [&reverseDeviceMaps](const auto& kv) {
+      return !reverseDeviceMaps.contains(kv.first);
+    });
 
     // remove devices that are no longer used
-    for (auto it = devices_.begin(); it != devices_.end();) {
-      if (std::find(devices.begin(), devices.end(), *it) == devices.end()) {
-        it = devices_.erase(it);
-      } else {
-        it++;
-      }
-    }
+    std::erase_if(devices_, [&](const auto& d) {
+      return std::find(devices.begin(), devices.end(), d) == devices.end();
+    });
   }
 }
 std::unordered_map<std::string, std::string> TensorPipeAgent::getMetrics() {
diff --git a/torch/csrc/dynamo/cpython_defs.c b/torch/csrc/dynamo/cpython_defs.c
index e0cb3bfe29607..b879c3d9e82c5 100644
--- a/torch/csrc/dynamo/cpython_defs.c
+++ b/torch/csrc/dynamo/cpython_defs.c
@@ -28,6 +28,12 @@ void init_THPCaches() {}
 
 #if IS_PYTHON_3_11_PLUS
 
+// Rename opcode table symbols to avoid multiple definition conflict
+// with the identical tables in libpython at link time.
+#define _PyOpcode_Caches _torch_PyOpcode_Caches
+#define _PyOpcode_Jump _torch_PyOpcode_Jump
+#define _PyOpcode_Deopt _torch_PyOpcode_Deopt
+
 #define Py_BUILD_CORE
 #define NEED_OPCODE_TABLES // To get _PyOpcode_Deopt, _PyOpcode_Caches
 
@@ -42,6 +48,9 @@ void init_THPCaches() {}
 
 #undef NEED_OPCODE_TABLES
 #undef Py_BUILD_CORE
+#undef _PyOpcode_Deopt
+#undef _PyOpcode_Jump
+#undef _PyOpcode_Caches
 
 // As a simple way to reduce the impact of ABI changes on the CPython side, this
 // check forces us to manually re-check that the function didn't change on the
@@ -497,8 +506,8 @@ const uint8_t* THP_PyOpcode_Caches = NULL;
 int THP_PyOpcode_Caches_size = 0;
 void init_THPCaches() {
 #if IS_PYTHON_3_11_PLUS
-  THP_PyOpcode_Caches = _PyOpcode_Caches;
-  THP_PyOpcode_Caches_size = sizeof(_PyOpcode_Caches) / sizeof(uint8_t);
+  THP_PyOpcode_Caches = _torch_PyOpcode_Caches;
+  THP_PyOpcode_Caches_size = sizeof(_torch_PyOpcode_Caches) / sizeof(uint8_t);
 #endif
 }
 
diff --git a/torch/csrc/dynamo/eval_frame.c b/torch/csrc/dynamo/eval_frame.c
index 58cb48de664d5..cb83ac9316e30 100644
--- a/torch/csrc/dynamo/eval_frame.c
+++ b/torch/csrc/dynamo/eval_frame.c
@@ -727,6 +727,48 @@ static int clear_state(PyObject* module) {
 
 bool is_skip_guard_eval_unsafe = false;
 
+// -1 means inactive, >= 0 means active with that many compiled frames.
+int fullgraph_compiled_frame_count = -1;
+
+// When true and fullgraph_compiled_frame_count > 0, sub-frames under fullgraph
+// compilation will error (via get_fail_callback) instead of being silently
+// skipped.
+bool fullgraph_error_on_nested_compile = false;
+
+// Set the fullgraph compiled frame counter and return the old value.
+// If setting to >= 0 (activating) and already active, no-op.
+static PyObject* set_fullgraph_compiled_frame_count_py(
+    PyObject* dummy,
+    PyObject* arg) {
+  long val = PyLong_AsLong(arg);
+  if (val == -1 && PyErr_Occurred()) {
+    return NULL;
+  }
+  int old = fullgraph_compiled_frame_count;
+  if (val >= 0 && old >= 0) {
+    // Already active, no-op.
+  } else {
+    fullgraph_compiled_frame_count = (int)val;
+  }
+  return PyLong_FromLong(old);
+}
+
+// Set fullgraph_error_on_nested_compile and return the old value.
+static PyObject* set_fullgraph_error_on_nested_compile_py(
+    PyObject* dummy,
+    PyObject* arg) {
+  if (arg != Py_False && arg != Py_True) {
+    PyErr_SetString(PyExc_TypeError, "expected True/False");
+    return NULL;
+  }
+  bool old = fullgraph_error_on_nested_compile;
+  fullgraph_error_on_nested_compile = arg == Py_True;
+  if (old) {
+    Py_RETURN_TRUE;
+  }
+  Py_RETURN_FALSE;
+}
+
 static PyMethodDef _methods[] = {
     {"set_eval_frame", set_eval_frame_py, METH_O, NULL},
     {"set_skip_guard_eval_unsafe", set_skip_guard_eval_unsafe, METH_O, NULL},
@@ -740,6 +782,14 @@ static PyMethodDef _methods[] = {
     {"set_guard_error_hook", set_guard_error_hook, METH_O, NULL},
     {"set_guard_complete_hook", set_guard_complete_hook, METH_O, NULL},
     {"raise_sigtrap", raise_sigtrap, METH_NOARGS, NULL},
+    {"set_fullgraph_compiled_frame_count",
+     set_fullgraph_compiled_frame_count_py,
+     METH_O,
+     NULL},
+    {"set_fullgraph_error_on_nested_compile",
+     set_fullgraph_error_on_nested_compile_py,
+     METH_O,
+     NULL},
     {NULL, NULL, 0, NULL}};
 
 static struct PyModuleDef _module = {
diff --git a/torch/csrc/dynamo/eval_frame.h b/torch/csrc/dynamo/eval_frame.h
index 870603262ddb6..703b9fa306e1f 100644
--- a/torch/csrc/dynamo/eval_frame.h
+++ b/torch/csrc/dynamo/eval_frame.h
@@ -32,6 +32,8 @@ THPPyInterpreterFrame* THPPyInterpreterFrame_New(
     THP_EVAL_API_FRAME_OBJECT* frame);
 
 extern bool is_skip_guard_eval_unsafe;
+extern int fullgraph_compiled_frame_count;
+extern bool fullgraph_error_on_nested_compile;
 
 void clear_old_frame_if_python_312_plus(
     PyThreadState* tstate,
diff --git a/torch/csrc/dynamo/eval_frame_cpp.cpp b/torch/csrc/dynamo/eval_frame_cpp.cpp
index a624d35e510f0..9ac630c0091fd 100644
--- a/torch/csrc/dynamo/eval_frame_cpp.cpp
+++ b/torch/csrc/dynamo/eval_frame_cpp.cpp
@@ -5,16 +5,253 @@
 #include <torch/csrc/dynamo/debug_macros.h>
 #include <torch/csrc/dynamo/eval_frame.h>
 #include <torch/csrc/dynamo/eval_frame_cpp.h>
+#include <torch/csrc/dynamo/extra_state.h>
 #include <torch/csrc/dynamo/framelocals_mapping.h>
+#include <torch/csrc/dynamo/stackref_bridge.h>
 #include <torch/csrc/utils/python_compat.h>
 
+#include <algorithm>
+#include <optional>
+#include <unordered_set>
+
 extern "C" {
 extern PyObject* guard_complete_hook;
 }
 
+// Bytecode debugger callback - stored as raw PyObject* to avoid
+// shutdown issues with static py::object destructor running after Python
+// finalizes.
+namespace {
+PyObject* bytecode_debugger_callback_obj = nullptr;
+std::unordered_set<PyCodeObject*> breakpoint_code_objects;
+
+// RAII guard that calls __exit__ on a Python context manager when destroyed.
+struct DebugContextGuard {
+  py::object ctx;
+
+  explicit DebugContextGuard(py::object c) : ctx(std::move(c)) {
+    ctx.attr("__enter__")();
+  }
+
+  ~DebugContextGuard() {
+    // Save any pending Python exception (e.g. KeyboardInterrupt from the
+    // debugger's 'q' command) so calling __exit__ doesn't clobber it.
+    PyObject *exc_type, *exc_value, *exc_tb;
+    PyErr_Fetch(&exc_type, &exc_value, &exc_tb);
+    try {
+      ctx.attr("__exit__")(py::none(), py::none(), py::none());
+    } catch (py::error_already_set& e) {
+      e.restore();
+      PyErr_Clear();
+    }
+    if (exc_type != nullptr) {
+      PyErr_Restore(exc_type, exc_value, exc_tb);
+    }
+  }
+
+  DebugContextGuard(const DebugContextGuard&) = delete;
+  DebugContextGuard& operator=(const DebugContextGuard&) = delete;
+};
+
+} // namespace
+
+void set_bytecode_debugger_callback(py::object callback) {
+  if (callback.is_none()) {
+    Py_XSETREF(bytecode_debugger_callback_obj, nullptr);
+  } else {
+    Py_XSETREF(bytecode_debugger_callback_obj, callback.inc_ref().ptr());
+  }
+}
+
+py::object get_bytecode_debugger_callback() {
+  if (bytecode_debugger_callback_obj == nullptr) {
+    return py::none();
+  }
+  return py::reinterpret_borrow<py::object>(
+      py::handle(bytecode_debugger_callback_obj));
+}
+
+void register_breakpoint_code(py::object code) {
+  breakpoint_code_objects.insert((PyCodeObject*)code.ptr());
+}
+
+// NullStackValue singleton for representing NULL stack values
+NullStackValue& NullStackValue::get_singleton() {
+  static NullStackValue instance;
+  return instance;
+}
+
+py::object get_null_stack_value() {
+  return py::cast(
+      NullStackValue::get_singleton(), py::return_value_policy::reference);
+}
+
+// Caller must provide the expected stack depth because frame->stacktop is -1
+// during INSTRUCTION monitoring callbacks. In CPython 3.12+, the interpreter
+// keeps the stack pointer in a local variable and only saves it to the frame
+// via _PyFrame_SetStackPointer() before certain callbacks (e.g., line tracing).
+// However, INSTRUMENTED_INSTRUCTION calls _Py_call_instrumentation_instruction
+// without saving the stack pointer first, leaving stacktop as -1.
+// See: cpython/Include/internal/pycore_frame.h (_PyFrame_GetStackPointer)
+py::list _get_frame_value_stack_with_depth(
+    const py::handle& frame_obj,
+    int depth) {
+  TORCH_CHECK_TYPE(PyFrame_Check(frame_obj.ptr()), "expected a frame object!");
+
+  py::list result;
+  if (depth <= 0) {
+    return result;
+  }
+
+#if IS_PYTHON_3_11_PLUS
+  PyFrameObject* frame = (PyFrameObject*)frame_obj.ptr();
+  _PyInterpreterFrame* iframe = frame->f_frame;
+  if (iframe == nullptr) {
+    return result;
+  }
+
+  PyCodeObject* code = F_CODE(iframe);
+  if (code == nullptr) {
+    return result;
+  }
+
+  int nlocalsplus = code->co_nlocalsplus;
+  int stacksize = code->co_stacksize;
+
+  // Clamp depth to valid range
+  if (depth > stacksize) {
+    depth = stacksize;
+  }
+
+  // When stacktop/stackpointer is available, use it as the authoritative
+  // depth. The caller's tracked depth can lag behind (e.g. after a CALL
+  // instruction pops arguments but the effect hasn't been applied to the
+  // Python-side tracker yet).
+#if IS_PYTHON_3_14_PLUS
+  if (iframe->stackpointer != nullptr) {
+    int actual_depth =
+        (int)(iframe->stackpointer - (iframe->localsplus + nlocalsplus));
+    if (actual_depth >= 0) {
+      depth = std::min(actual_depth, depth);
+    }
+  }
+#else
+  bool have_stack_pointer = false;
+  if (iframe->stacktop > 0) {
+    int actual_depth = iframe->stacktop - nlocalsplus;
+    if (actual_depth >= 0) {
+      depth = std::min(actual_depth, depth);
+      have_stack_pointer = true;
+    }
+  }
+#endif
+
+#if IS_PYTHON_3_14_PLUS
+  if (iframe->stackpointer == nullptr) {
+    return result;
+  }
+  _PyStackRef* stack_base = iframe->localsplus + nlocalsplus;
+  for (int i = 0; i < depth; i++) {
+    PyObject* obj = THP_PyStackRef_AsPyObjectBorrow(&stack_base[i]);
+    if (obj == nullptr) {
+      result.append(get_null_stack_value());
+    } else {
+      result.append(py::reinterpret_borrow<py::object>(py::handle(obj)));
+    }
+  }
+#else
+  int stack_start = nlocalsplus;
+  for (int i = 0; i < depth; i++) {
+    PyObject* obj = iframe->localsplus[stack_start + i];
+    if (obj == nullptr) {
+      result.append(get_null_stack_value());
+    } else if (!have_stack_pointer && Py_REFCNT(obj) <= 0) {
+      // Without a reliable stack pointer (current frame, stacktop == -1),
+      // the caller's tracked depth may overestimate. Stop at entries that
+      // look like freed objects to avoid dereferencing stale pointers.
+      break;
+    } else {
+      result.append(py::reinterpret_borrow<py::object>(py::handle(obj)));
+    }
+  }
+#endif
+
+#else
+  // Python 3.10 and earlier - use f_valuestack
+  PyFrameObject* frame = (PyFrameObject*)frame_obj.ptr();
+  if (frame->f_valuestack == nullptr) {
+    return result;
+  }
+
+  PyCodeObject* code = frame->f_code;
+  if (code == nullptr) {
+    return result;
+  }
+
+  int stacksize = code->co_stacksize;
+  if (depth > stacksize) {
+    depth = stacksize;
+  }
+
+  for (int i = 0; i < depth; i++) {
+    PyObject* obj = frame->f_valuestack[i];
+    if (obj == nullptr) {
+      result.append(get_null_stack_value());
+    } else {
+      result.append(py::reinterpret_borrow<py::object>(py::handle(obj)));
+    }
+  }
+#endif // IS_PYTHON_3_11_PLUS
+
+  return result;
+}
+
 static constexpr const char* cache_lookup_profiler_str =
     "TorchDynamo Cache Lookup";
 
+// Cache the random module to avoid calling py::module_::import("random") at
+// arbitrary points during execution. torch.package overrides the import
+// machinery, and importing "random" inside a package archive context would fail
+// if random isn't in the extern list.
+// Stored as raw PyObject* (leaked ref) to avoid destructor running after Python
+// finalizes, same pattern as bytecode_debugger_callback_obj above.
+static PyObject* random_module = nullptr;
+static py::handle get_random_module() {
+  if (random_module == nullptr) {
+    random_module = py::module_::import("random").release().ptr();
+  }
+  return py::handle(random_module);
+}
+
+// Use RAII to save/restore global state across the dynamo callback
+class PreserveGlobalState {
+  py::object random_state;
+
+ public:
+  PreserveGlobalState() {
+    this->random_state = get_random_module().attr("getstate")();
+  }
+  PreserveGlobalState(const PreserveGlobalState&) = delete;
+  PreserveGlobalState(PreserveGlobalState&&) = delete;
+  PreserveGlobalState& operator=(const PreserveGlobalState&) = delete;
+  PreserveGlobalState& operator=(PreserveGlobalState&&) = delete;
+  ~PreserveGlobalState() {
+    try {
+      get_random_module().attr("setstate")(this->random_state);
+    } catch (py::error_already_set& e) {
+      try {
+        e.restore();
+      } catch (...) {
+        // Intentionally return to silence empty catch linter.
+        // We can't propagate exceptions since we are in a destructor.
+        return;
+      }
+    } catch (...) {
+      return;
+    }
+  }
+};
+
 // Remember to update the type signature for DynamoCallbackFn.__call__ in
 // torch/_dynamo/types.py if this function's signature changes.
 static py::object dynamo_call_callback(
@@ -55,9 +292,9 @@ static py::handle _callback_from_action(
 static int32_t c_recursion_limit = -1;
 
 void dynamo_set_c_recursion_limit(int32_t limit) {
-  if (limit < 1 && limit != -1) {
-    throw std::range_error("recursion limit must be >= 1, or -1 to reset");
-  }
+  TORCH_CHECK_VALUE(
+      limit >= 1 || limit == -1,
+      "recursion limit must be >= 1, or -1 to reset");
   c_recursion_limit = limit;
 }
 
@@ -100,18 +337,6 @@ struct CRecursionLimitRAII {
 
 #endif
 
-EvalFrameOverride eval_frame_override = EvalFrameOverride::NONE;
-
-EvalFrameOverride get_eval_frame_override() {
-  return eval_frame_override;
-}
-
-EvalFrameOverride set_eval_frame_override(EvalFrameOverride override) {
-  EvalFrameOverride prev = eval_frame_override;
-  eval_frame_override = override;
-  return prev;
-}
-
 // frame and callback are borrowed references.
 // Returns new reference.
 PyObject* dynamo__custom_eval_frame(
@@ -157,7 +382,7 @@ PyObject* dynamo__custom_eval_frame(
     // immediately skip the frame, and (2) even if it did, this would only
     // be profitable if there was tensor code in the unwinding code.  Seems
     // unlikely.
-    DEBUG_TRACE("throw %s", get_frame_name(frame));
+    DEBUG_TRACE("throw %s", get_frame_name(frame)); // @allow-raw-throw
     return dynamo_eval_frame_default(tstate, frame, throw_flag);
   }
 
@@ -191,32 +416,49 @@ PyObject* dynamo__custom_eval_frame(
   // original frame, we are responsible for clearing it - via
   // clear_old_frame_if_python_312_plus.
   auto eval_custom = [&]() {
-    // If we're attempting to run dynamo-generated code and eval frame override
-    // is set to SKIP, then we should set the callback to None to skip.
-    // If the override is set to ERROR, then we call
-    // torch._dynamo.convert_frame.get_fail_callback, which patches
-    // convert_frame.compile_frame with a function that errors unconditionally.
-    // This means Dynamo will error if it attempts to trace into the frame
-    // (Python-level skips pre-trace are permissible).
-    if (!recursive_callback.is_none() &&
-        !recursive_callback.is(py::bool_(false))) {
-      if (eval_frame_override == EvalFrameOverride::SKIP) {
-        recursive_callback = py::none();
-      } else if (eval_frame_override == EvalFrameOverride::ERROR) {
-        if (!convert_frame_get_fail_callback) {
-          convert_frame_get_fail_callback =
-              py::module_::import("torch._dynamo.convert_frame")
-                  .attr("get_fail_callback");
-          auto atexit = py::module_::import("atexit");
-          atexit.attr("register")(py::cpp_function(
-              []() { convert_frame_get_fail_callback = std::nullopt; }));
+    if (fullgraph_compiled_frame_count >= 0) {
+      fullgraph_compiled_frame_count++;
+      // Under fullgraph, disable or error Dynamo for sub-frames of compiled
+      // code. If fullgraph_error_on_nested_compile is set, wrap the callback
+      // with get_fail_callback so compilation attempts error. Otherwise, set
+      // callback to None to skip sub-frames entirely.
+      if (!recursive_callback.is_none() &&
+          !recursive_callback.is(py::bool_(false))) {
+        if (fullgraph_error_on_nested_compile) {
+          if (!convert_frame_get_fail_callback) {
+            convert_frame_get_fail_callback =
+                py::module_::import("torch._dynamo.convert_frame")
+                    .attr("get_fail_callback");
+            auto atexit = py::module_::import("atexit");
+            atexit.attr("register")(py::cpp_function(
+                []() { convert_frame_get_fail_callback = std::nullopt; }));
+          }
+          recursive_callback =
+              convert_frame_get_fail_callback.value()(recursive_callback);
+        } else {
+          recursive_callback = py::none();
         }
-        recursive_callback =
-            convert_frame_get_fail_callback.value()(recursive_callback);
       }
     }
     eval_frame_callback_set(recursive_callback.ptr());
     DEBUG_NULL_CHECK(cached_code);
+    // Auto-activate debugger for code objects with breakpoints.
+    // DebugContextGuard calls __enter__ on construction and __exit__ on
+    // destruction, so the debug session is scoped to this eval_custom call.
+    std::optional<DebugContextGuard> debug_guard;
+    if (breakpoint_code_objects.count(cached_code) &&
+        bytecode_debugger_callback_obj == nullptr) {
+      auto ctx = py::module_::import("torch._dynamo.bytecode_debugger")
+                     .attr("_DebugContext")();
+      ctx.attr("_stop_at_new_code") = false;
+      debug_guard.emplace(std::move(ctx));
+    }
+    // Call bytecode debugger callback if set, to allow instruction-level
+    // debugging of the Dynamo-generated code
+    py::object debugger_cb = get_bytecode_debugger_callback();
+    if (!debugger_cb.is_none()) {
+      debugger_cb(py::handle((PyObject*)cached_code));
+    }
     eval_result = dynamo_eval_custom_code(
         tstate, frame, cached_code, trace_annotation, throw_flag);
     if (!callback.is(recursive_callback)) {
@@ -354,6 +596,7 @@ PyObject* dynamo__custom_eval_frame(
       fail();
       return eval_result;
     }
+    PreserveGlobalState preserve_global_state;
     callback_result = dynamo_call_callback(
         callback, frame, locals.get(), cache_entry, frame_state);
     new_strategy =
diff --git a/torch/csrc/dynamo/eval_frame_cpp.h b/torch/csrc/dynamo/eval_frame_cpp.h
index 2d76d6008fee5..34ab9676185e1 100644
--- a/torch/csrc/dynamo/eval_frame_cpp.h
+++ b/torch/csrc/dynamo/eval_frame_cpp.h
@@ -4,8 +4,11 @@
 #include <torch/csrc/dynamo/eval_frame.h>
 #include <torch/csrc/dynamo/extra_state.h>
 #include <torch/csrc/dynamo/framelocals_mapping.h>
+
 #ifdef __cplusplus
 
+#include <torch/csrc/utils/pybind.h>
+
 extern "C" {
 
 #endif
@@ -26,13 +29,22 @@ int32_t dynamo_get_c_recursion_limit();
 
 } // extern "C"
 
-// Used to override the Dynamo callback for fullgraph=True'd compiled objects
-enum class EvalFrameOverride {
-  NONE, // Run regular set callback
-  SKIP, // skip frames recursively
-  ERROR, // error if Dynamo attempts to trace code
+// Bytecode debugger callback functions
+void set_bytecode_debugger_callback(py::object callback);
+py::object get_bytecode_debugger_callback();
+
+// Breakpoint code object tracking
+void register_breakpoint_code(py::object code);
+
+// NullStackValue - sentinel class for representing NULL values on Python stack
+class NullStackValue {
+ public:
+  static NullStackValue& get_singleton();
 };
 
-EvalFrameOverride set_eval_frame_override(EvalFrameOverride override);
+py::object get_null_stack_value();
+py::list _get_frame_value_stack_with_depth(
+    const py::handle& frame_obj,
+    int depth);
 
 #endif
diff --git a/torch/csrc/dynamo/extra_state.cpp b/torch/csrc/dynamo/extra_state.cpp
index b890c2848011b..878404096740b 100644
--- a/torch/csrc/dynamo/extra_state.cpp
+++ b/torch/csrc/dynamo/extra_state.cpp
@@ -229,9 +229,9 @@ CacheEntry* create_cache_entry(
 }
 
 py::list _debug_get_cache_entry_list(const py::handle& code_obj) {
-  if (!py::isinstance(code_obj, py::module::import("types").attr("CodeType"))) {
-    throw py::type_error("expected a code object!");
-  }
+  TORCH_CHECK_TYPE(
+      py::isinstance(code_obj, py::module::import("types").attr("CodeType")),
+      "expected a code object!");
   PyCodeObject* code = (PyCodeObject*)code_obj.ptr();
   ExtraState* extra = get_extra_state(code);
   py::list result;
@@ -252,9 +252,9 @@ PrecompileEntry::PrecompileEntry(py::object gm, py::object c)
 }
 
 void _reset_precompile_entries(const py::handle& code_obj) {
-  if (!py::isinstance(code_obj, py::module::import("types").attr("CodeType"))) {
-    throw py::type_error("expected a code object!");
-  }
+  TORCH_CHECK_TYPE(
+      py::isinstance(code_obj, py::module::import("types").attr("CodeType")),
+      "expected a code object!");
   PyCodeObject* code = (PyCodeObject*)code_obj.ptr();
   ExtraState* extra = get_extra_state(code);
   py::list result;
@@ -267,9 +267,9 @@ void _load_precompile_entry(
     const py::handle& code_obj,
     py::object guard_manager,
     py::object dynamo_code) {
-  if (!py::isinstance(code_obj, py::module::import("types").attr("CodeType"))) {
-    throw py::type_error("expected a code object!");
-  }
+  TORCH_CHECK_TYPE(
+      py::isinstance(code_obj, py::module::import("types").attr("CodeType")),
+      "expected a code object!");
   PyCodeObject* code = (PyCodeObject*)code_obj.ptr();
   ExtraState* extra = get_extra_state(code);
   py::list result;
@@ -290,9 +290,9 @@ void _set_lru_cache(py::object boolean) {
 }
 
 py::list _debug_get_precompile_entries(const py::handle& code_obj) {
-  if (!py::isinstance(code_obj, py::module::import("types").attr("CodeType"))) {
-    throw py::type_error("expected a code object!");
-  }
+  TORCH_CHECK_TYPE(
+      py::isinstance(code_obj, py::module::import("types").attr("CodeType")),
+      "expected a code object!");
   PyCodeObject* code = (PyCodeObject*)code_obj.ptr();
   ExtraState* extra = get_extra_state(code);
   py::list result;
diff --git a/torch/csrc/dynamo/guards.cpp b/torch/csrc/dynamo/guards.cpp
index 833e7667009df..9d97d7ccf2268 100644
--- a/torch/csrc/dynamo/guards.cpp
+++ b/torch/csrc/dynamo/guards.cpp
@@ -6,6 +6,7 @@
 #define PY_SSIZE_T_CLEAN
 #include <ATen/EmptyTensor.h>
 #include <ATen/SparseCsrTensorUtils.h>
+#include <c10/util/Synchronized.h>
 #include <c10/util/flat_hash_map.h>
 #include <fmt/format.h>
 #include <torch/csrc/autograd/grad_mode.h>
@@ -537,13 +538,15 @@ PyObject* TensorGuards_check_verbose(
     PyObject* item = PyTuple_GET_ITEM(args, i);
     if (Py_TYPE(item) != checks[i].pytype) {
       std::stringstream fail_reason;
-      PyObject* type_str = PyObject_Str(PyObject_Type(item));
+      PyObject* type_str =
+          PyObject_Str(reinterpret_cast<PyObject*>(Py_TYPE(item)));
       fail_reason << "expected type of '" << tensor_check_names[i]
                   << "' to be a tensor type, ";
       if (!type_str) {
         fail_reason << "but found a different type";
       } else {
         fail_reason << "' but found " << PyUnicode_AsUTF8(type_str);
+        Py_DECREF(type_str);
       }
       return Py_BuildValue("s", fail_reason.str().c_str());
     }
@@ -874,20 +877,26 @@ static PyObject* check_obj_id(PyObject* dummy, PyObject* args) {
 
 #if IS_PYTHON_3_12_PLUS
 
-static std::unordered_map<PyObject*, uint64_t> dict_version_map;
+struct DictVersionState {
+  std::unordered_map<PyObject*, uint64_t> map;
+  uint64_t next_id = 1;
+};
+
+static c10::Synchronized<DictVersionState> dict_version_state;
 static int dict_version_watcher_id;
 static int dict_recursive_tag_watcher_id;
-static uint64_t global_dict_version_id = 1;
 static int dict_version_watch_callback(
     PyDict_WatchEvent event,
     PyObject* dict,
     PyObject* key,
     PyObject* new_value) noexcept {
-  if (event == PyDict_EVENT_DEALLOCATED) {
-    dict_version_map.erase(dict);
-  } else if (event != PyDict_EVENT_CLONED) {
-    dict_version_map[dict] = global_dict_version_id++;
-  }
+  dict_version_state.withLock([&](DictVersionState& state) {
+    if (event == PyDict_EVENT_DEALLOCATED) {
+      state.map.erase(dict);
+    } else if (event != PyDict_EVENT_CLONED) {
+      state.map[dict] = state.next_id++;
+    }
+  });
   return 0;
 }
 
@@ -899,10 +908,13 @@ static uint64_t get_dict_version_unchecked(PyObject* dict) {
   TORCH_CHECK(
       !PyDict_Watch(dict_version_watcher_id, dict),
       "failed to add version watcher to dict!");
-  if (!dict_version_map.count(dict)) {
-    dict_version_map[dict] = global_dict_version_id++;
-  }
-  return dict_version_map[dict];
+  return dict_version_state.withLock([&](DictVersionState& state) -> uint64_t {
+    auto [it, inserted] = state.map.try_emplace(dict, state.next_id);
+    if (inserted) {
+      state.next_id++;
+    }
+    return it->second;
+  });
 
 #else
 
@@ -911,16 +923,11 @@ static uint64_t get_dict_version_unchecked(PyObject* dict) {
 #endif
 }
 
-static PyObject* dict_version(PyObject* dummy, PyObject* args) {
-  // Retrieves the version of a dictionary.
-  PyObject* obj = nullptr;
-  if (!PyArg_ParseTuple(args, "O", &obj)) {
-    return nullptr;
-  }
-  if (!PyDict_Check(obj)) {
-    return nullptr;
-  }
+static PyObject* dict_version(PyObject* dummy, PyObject* obj) {
+  HANDLE_TH_ERRORS
+  TORCH_CHECK(PyDict_Check(obj), "dict_version expects a dict");
   return THPUtils_packUInt64(get_dict_version_unchecked(obj));
+  END_HANDLE_TH_ERRORS
 }
 
 static PyObject* assert_size_stride(PyObject* dummy, PyObject* args) {
@@ -1054,6 +1061,48 @@ static PyObject* assert_alignment(PyObject* dummy, PyObject* args) {
   Py_RETURN_TRUE;
 }
 
+static PyObject* copy_misaligned(PyObject* dummy, PyObject* item) {
+  /*
+   * If the tensor's data pointer is not 16-byte aligned, return a
+   * clone that preserves strides. Otherwise return the original
+   * tensor (new reference).  Implemented in C++ so the aligned
+   * fast-path is just a pointer check with minimal Python overhead.
+   *
+   * NOTE: kAlignment is hardcoded to match torch._inductor.utils.ALIGNMENT.
+   * If alignment requirements ever change or become per-platform, this
+   * constant must be updated (or turned into a parameter).
+   */
+  constexpr size_t kAlignment = 16;
+
+  if (!THPVariable_CheckExact(item) && !THPVariable_Check(item)) {
+    PyErr_SetString(PyExc_TypeError, "expected Tensor()");
+    return nullptr;
+  }
+
+  at::Tensor tensor = THPVariable_Unpack(item);
+
+  if (reinterpret_cast<uintptr_t>(tensor.data_ptr()) % kAlignment == 0) {
+    // Already aligned – return the original tensor.
+    Py_INCREF(item);
+    return item;
+  }
+
+  // Misaligned – clone while preserving strides.
+  // Same logic as torch._inductor.utils.clone_preserve_strides.
+  int64_t needed_size = 0;
+  if (tensor.numel() > 0) {
+    auto sizes = tensor.sizes();
+    auto strides = tensor.strides();
+    for (int64_t i = 0; i < tensor.dim(); ++i) {
+      needed_size += (sizes[i] - 1) * strides[i];
+    }
+    needed_size += 1;
+  }
+  at::Tensor flat = at::as_strided(tensor, {needed_size}, {1}).clone();
+  at::Tensor result = at::as_strided(flat, tensor.sizes(), tensor.strides());
+  return THPVariable_Wrap(std::move(result));
+}
+
 template <typename T>
 static void unwrap_size_tuple(PyObject* obj, T& output) {
   TORCH_CHECK(PyTuple_CheckExact(obj));
@@ -1176,7 +1225,8 @@ static PyMethodDef _methods[] = {
     {"check_obj_id", check_obj_id, METH_VARARGS, nullptr},
     {"assert_size_stride", assert_size_stride, METH_VARARGS, nullptr},
     {"assert_alignment", assert_alignment, METH_VARARGS, nullptr},
-    {"dict_version", dict_version, METH_VARARGS, nullptr},
+    {"copy_misaligned", copy_misaligned, METH_O, nullptr},
+    {"dict_version", dict_version, METH_O, nullptr},
     {"_empty_strided_cpu", _empty_strided_cpu, METH_VARARGS, nullptr},
     {"_empty_strided_cpu_pinned",
      _empty_strided_cpu_pinned,
@@ -1637,7 +1687,9 @@ class DictGuardManager;
 // the container can grow large over the lifetime of the process.  That’s
 // acceptable: lookup is by pointer (hash/equals = identity) and each entry
 // stores only lightweight pointers.
-std::unordered_map<PyObject*, std::list<GuardManager*>> dict_to_guard_managers;
+using DictToGuardManagersMap =
+    std::unordered_map<PyObject*, std::list<GuardManager*>>;
+c10::Synchronized<DictToGuardManagersMap> dict_to_guard_managers;
 
 /**
  * Base class for the leaf guard in the GuardManager hierarchy.
@@ -1683,7 +1735,6 @@ class LeafGuard {
   // is not exposed to Python and can only be called from C++.
   virtual bool check_nopybind(PyObject* value) = 0;
   virtual bool check_nopybind(FrameLocalsMapping* map) {
-    // throw std::runtime_error("fallback to python");
     // Could fallback to running check on the Python dict (lazily constructed)
     return check_nopybind((PyObject*)map->to_dict());
   }
@@ -1726,7 +1777,8 @@ class LAMBDA_GUARD : public LeafGuard {
     if (py::isinstance<py::function>(guard_check_fn)) {
       _guard_check_fn = py::cast<py::function>(std::move(guard_check_fn));
     } else {
-      throw py::type_error("LAMBDA_GUARD expects (callable, str)");
+      throw py::type_error(
+          "LAMBDA_GUARD expects (callable, str)"); // @allow-raw-throw
     }
   }
 
@@ -2154,7 +2206,8 @@ class GLOBAL_STATE : public LeafGuard {
         owner_(std::move(initial_state)),
         _guard((GlobalStateGuard*)owner_.ptr()) {
     if (!PyObject_TypeCheck(owner_.ptr(), &GlobalStateGuardType)) {
-      throw py::type_error("GLOBAL_STATE expects a GlobalStateGuard");
+      throw py::type_error(
+          "GLOBAL_STATE expects a GlobalStateGuard"); // @allow-raw-throw
     }
   }
 
@@ -2536,11 +2589,13 @@ class SYMBOLIC_SHAPE_GUARD : public RelationalGuard {
     _nargs_float = PyLong_AsSize_t(nargs_float.ptr());
     _nargs = _nargs_int + _nargs_float;
     if (PyErr_Occurred()) {
+      // @allow-raw-throw
       throw py::value_error(
           "SYMBOLIC_SHAPE_GUARD expected a non-negative number of arguments.");
     }
     uintptr_t addr = PyLong_AsUnsignedLongLong(py_addr.ptr());
     if (PyErr_Occurred()) {
+      // @allow-raw-throw
       throw py::value_error(
           "SYMBOLIC_SHAPE_GUARD expected an address to a C function.");
     }
@@ -2687,7 +2742,7 @@ class DICT_VERSION : public LeafGuard {
             std::move(verbose_code_parts),
             std::move(user_stack)) {
     if (!PyDict_Check(value.ptr())) {
-      throw py::type_error("DICT_VERSION expects a dict");
+      throw py::type_error("DICT_VERSION expects a dict"); // @allow-raw-throw
     }
     _tag = get_dict_version_unchecked(value.ptr());
   }
@@ -2724,6 +2779,7 @@ void stop_recording_dict_pointers(
 bool is_recording_dict_pointers(RootGuardManager* root);
 void record_dict_pointer(RootGuardManager* root, PyObject* dict_pointer);
 void record_tensor_pointer(RootGuardManager* root, PyObject* tensor_pointer);
+void record_tensor_metadata(RootGuardManager* root, PyObject* tensor_pointer);
 
 GuardManager* clone_guard_manager(
     GuardManager* from,
@@ -2734,12 +2790,53 @@ void add_relational_guard_resetter_to_cloned_root(
     std::shared_ptr<RelationalGuard> guard);
 std::shared_ptr<RelationalGuard> get_no_tensor_aliasing_guard(
     RootGuardManager* _root);
+const LocalState& get_local_state(RootGuardManager* root);
 // std::string get_compile_id(RootGuardManager* root);
 
 struct WeakEntry {
   PyObject* wr; // weakref
   PyObject* cap; // capsule whose m_self is used by the callback
 };
+
+// Convert concrete sizes/strides to the optional<SymInt> vectors that
+// TensorCheck expects.  All dimensions are treated as static (no nullopt).
+inline std::vector<std::optional<c10::SymInt>> to_opt_symint(
+    c10::IntArrayRef vals) {
+  std::vector<std::optional<c10::SymInt>> out;
+  out.reserve(vals.size());
+  for (auto v : vals) {
+    out.emplace_back(c10::SymInt(v));
+  }
+  return out;
+}
+
+// Build a TensorCheck that validates all concrete metadata (dispatch key,
+// dtype, device, requires_grad, sizes, strides) for the dict-tag fast path.
+inline TensorCheck make_tensor_check(
+    const LocalState& state,
+    const at::Tensor& tensor) {
+  auto layout = tensor.layout();
+  bool sparse = layout == c10::kSparseCsr || layout == c10::kSparseCsc ||
+      layout == c10::kSparseBsc || layout == c10::kSparseBsr;
+  // Sparse layouts don't support strides; use nullopt per dim so
+  // TensorCheck skips stride comparison for each dimension.
+  auto strides = sparse
+      ? std::vector<std::optional<c10::SymInt>>(tensor.dim(), std::nullopt)
+      : to_opt_symint(tensor.strides());
+  return TensorCheck(
+      state,
+      /*pt=*/nullptr,
+      tensor,
+      tensor.key_set(),
+      to_opt_symint(tensor.sizes()),
+      std::move(strides));
+}
+
+struct RecordedTensorMetadata {
+  PyObject* tensor_ptr;
+  TensorCheck check;
+};
+
 /**
  * Base class representing a pair of accessor and the associated guard
  * manager. The accessor defines how to access the child value from the
@@ -2782,7 +2879,6 @@ class GuardAccessor {
   // subtree on immutable dict getitems.
   virtual bool check_nopybind(PyObject* obj, bool matches_dict_tag = false) = 0;
   virtual bool check_nopybind(FrameLocalsMapping* map, bool matches_dict_tag) {
-    // throw std::runtime_error("fallback to python");
     // Could fallback to running check on the Python dict (lazily constructed)
     return check_nopybind((PyObject*)map->to_dict(), matches_dict_tag);
   }
@@ -2919,6 +3015,7 @@ class GuardManager {
       if (PyCapsule_IsValid(e.cap, "GuardManager*")) {
         PyCapsule_SetName(e.cap, "DeadGuardManager");
       }
+      Py_DECREF(e.cap);
       Py_CLEAR(e.wr); // kills weakref (may remove callback)
     }
     _tag_safe_entries.clear();
@@ -3009,8 +3106,21 @@ class GuardManager {
     _tensor_pointers[value] = tensor_pointers;
   }
 
+  void stash_tensor_metadata(
+      PyObject* value,
+      std::vector<RecordedTensorMetadata>&& tensor_metadata) {
+    _tensor_metadata_pointers[value] = std::move(tensor_metadata);
+  }
+
   void disable_recursive_dict_tag_optimization() {
-    unwatch_all_saved_dict_pointers();
+    dict_to_guard_managers.withLock([&](DictToGuardManagersMap& map) {
+      disable_recursive_dict_tag_optimization(map);
+    });
+  }
+
+  // Caller must hold dict_to_guard_managers lock.
+  void disable_recursive_dict_tag_optimization(DictToGuardManagersMap& map) {
+    unwatch_all_saved_dict_pointers(map);
     _disable_dict_tag_matching = true;
   }
 
@@ -3135,6 +3245,22 @@ class GuardManager {
     return true;
   }
 
+  bool check_tensor_metadata_fast(PyObject* value) {
+    auto it = _tensor_metadata_pointers.find(value);
+    if (it == _tensor_metadata_pointers.end()) {
+      return true;
+    }
+    for (auto& recorded_tensor : it->second) {
+      if (!THPVariable_Check(recorded_tensor.tensor_ptr) ||
+          !recorded_tensor.check.check(
+              get_local_state(_root),
+              THPVariable_Unpack(recorded_tensor.tensor_ptr))) {
+        return false;
+      }
+    }
+    return true;
+  }
+
   bool check_no_tensor_aliasing_guards_fast(PyObject* value) {
     std::shared_ptr<RelationalGuard> no_tensor_aliasing_guard =
         get_no_tensor_aliasing_guard(_root);
@@ -3161,15 +3287,17 @@ class GuardManager {
     // For a `tag_safe_root`, the input pointer called `value`, the object the
     // guard is inspecting, serves as a proxy for the entire nested dictionary
     // structure beneath that node.  If this `value` pointer is one we have
-    // already recorded, then verifying each dictionary’s tag is sufficient to
-    // prove that nothing inside the subtree has changed.
+    // already recorded, then verifying each dictionary’s tag plus the cached
+    // tensor metadata is sufficient to prove that nothing inside the subtree
+    // has changed.
     //
     // Runtime flow
     // -------------
     // 1) Previously‑seen `value` pointer
     //    • Look up the current `value` pointer in our cache.
-    //    • If found, perform a recursive tag comparison on the cached subtree.
-    //      All tags match means guard passes with no further traversal.
+    //    • If found, perform a recursive tag comparison on the cached subtree
+    //      and revalidate recorded tensor metadata.
+    //      All checks passing means guard passes with no further traversal.
     //
     // 2) First‑time `value` pointer
     //    • Enter recording mode; walk the subtree, each tag safe root collects
@@ -3211,7 +3339,8 @@ class GuardManager {
         if (_dict_pointers.find(value) != _dict_pointers.end()) {
           // Check for fast path
           // if (is_weakref_valid(value) && check_dict_pointer_tags(value)) {
-          if (check_dict_pointer_tags(value)) {
+          if (check_dict_pointer_tags(value) &&
+              check_tensor_metadata_fast(value)) {
             if (check_no_tensor_aliasing_guards_fast(value)) {
               return true;
             } else {
@@ -3241,6 +3370,10 @@ class GuardManager {
         } else if (_has_no_tensor_aliasing_guard) {
           record_tensor_pointer(_root, value);
         }
+        // Tensor metadata can mutate in-place without changing dict tags.
+        if (_is_immutable && THPVariable_Check(value)) {
+          record_tensor_metadata(_root, value);
+        }
       }
     }
 
@@ -3323,6 +3456,7 @@ class GuardManager {
       return false;
     }
     // These will be decrefed in destructor
+    Py_INCREF(capsule);
     _tag_safe_entries.push_back({wr, capsule});
     return true;
   }
@@ -3353,26 +3487,16 @@ class GuardManager {
         PyErr_Clear();
         return false;
       }
-      dict_to_guard_managers[dict_pointer].push_back(this);
+      dict_to_guard_managers.withLock([&](DictToGuardManagersMap& map) {
+        map[dict_pointer].push_back(this);
+      });
     }
 #endif
     return true;
   }
 
-  void unwatch_all_saved_dict_pointers() {
-    /*
-    We may have recorded hundreds/thousands of dict pointers for the recursive
-    dict-tag optimisation. If any of those dicts mutates, we want to disable the
-    optimisation and then unwatch as many dict pointers as we can.
-
-    Be careful: the same dict pointer can be recorded by multiple GuardManagers.
-    So the flow is:
-
-      1) Remove *this* GuardManager from dict_to_guard_managers[dict_pointer].
-      2) If the list for that dict becomes empty, then:
-          - PyDict_Unwatch(dict_recursive_tag_watcher_id, dict_pointer)
-          - erase the dict_pointer entry from dict_to_guard_managers.
-    */
+  // Caller must hold dict_to_guard_managers lock.
+  void unwatch_all_saved_dict_pointers(DictToGuardManagersMap& map) {
 #if IS_PYTHON_3_12_PLUS
     if (!_disable_dict_tag_matching) {
       for (auto& value_stashed_pointers : _dict_pointers) {
@@ -3381,20 +3505,15 @@ class GuardManager {
         for (auto& stashed_pointer : stashed_pointers) {
           PyObject* dict_pointer = stashed_pointer.first;
 
-          // Delete the guard manager from the dict_to_guard_managers
           auto it = std::find(
-              dict_to_guard_managers[dict_pointer].begin(),
-              dict_to_guard_managers[dict_pointer].end(),
-              this);
-          if (it != dict_to_guard_managers[dict_pointer].end()) {
-            dict_to_guard_managers[dict_pointer].erase(it);
+              map[dict_pointer].begin(), map[dict_pointer].end(), this);
+          if (it != map[dict_pointer].end()) {
+            map[dict_pointer].erase(it);
           }
 
-          // Unwatch the dict pointer if this was the last guard manager
-          // watching it.
-          if (dict_to_guard_managers[dict_pointer].empty()) {
+          if (map[dict_pointer].empty()) {
             PyDict_Unwatch(dict_recursive_tag_watcher_id, dict_pointer);
-            dict_to_guard_managers.erase(dict_pointer);
+            map.erase(dict_pointer);
           }
         }
       }
@@ -3589,7 +3708,7 @@ class GuardManager {
 
   void add_permitted_leaf_guard(std::shared_ptr<LeafGuard> leaf_guard) {
     // Selectively called for permitted guards. This is used by DictGuardManager
-    // which overrides the add_leaf_guard manager to throw runtime error.
+    // which overrides the add_leaf_guard manager to raise a runtime error.
     GuardManager::add_leaf_guard(std::move(leaf_guard));
   }
 
@@ -3645,6 +3764,8 @@ class GuardManager {
   std::unordered_map<PyObject*, std::vector<std::pair<PyObject*, uint64_t>>>
       _dict_pointers;
   std::unordered_map<PyObject*, std::vector<PyObject*>> _tensor_pointers;
+  std::unordered_map<PyObject*, std::vector<RecordedTensorMetadata>>
+      _tensor_metadata_pointers;
   std::vector<WeakEntry> _tag_safe_entries;
 
   // 3.12+ related helper
@@ -3913,6 +4034,7 @@ class RootGuardManager : public GuardManager {
     _current_tag_safe_root = nullptr;
     _recorded_dict_pointers.clear();
     _recorded_tensor_pointers.clear();
+    _recorded_tensor_metadata.clear();
   }
 
   void stop_recording_dict_pointers(PyObject* value, bool result) {
@@ -3922,6 +4044,8 @@ class RootGuardManager : public GuardManager {
           value, _recorded_dict_pointers);
       _current_tag_safe_root->stash_tensor_pointers(
           value, _recorded_tensor_pointers);
+      _current_tag_safe_root->stash_tensor_metadata(
+          value, std::move(_recorded_tensor_metadata));
     }
     reset_dict_tag_recording_variables();
   }
@@ -3939,6 +4063,13 @@ class RootGuardManager : public GuardManager {
     _recorded_tensor_pointers.push_back(tensor_pointer);
   }
 
+  void record_tensor_metadata(PyObject* tensor_pointer) {
+    _recorded_tensor_metadata.push_back(RecordedTensorMetadata{
+        tensor_pointer,
+        make_tensor_check(_local_state, THPVariable_Unpack(tensor_pointer)),
+    });
+  }
+
  public:
   // Local state for TENSOR_MATCH guards.
   LocalState _local_state;
@@ -3995,6 +4126,7 @@ class RootGuardManager : public GuardManager {
   GuardManager* _current_tag_safe_root{nullptr};
   std::vector<std::pair<PyObject*, uint64_t>> _recorded_dict_pointers;
   std::vector<PyObject*> _recorded_tensor_pointers;
+  std::vector<RecordedTensorMetadata> _recorded_tensor_metadata;
 };
 
 /*
@@ -4356,15 +4488,18 @@ static int dict_recursive_tag_watch_callback(
     PyObject* key,
     PyObject* new_value) noexcept {
   if (event != PyDict_EVENT_CLONED) {
-    auto it = dict_to_guard_managers.find(dict);
-    if (it != dict_to_guard_managers.end()) {
-      auto guard_managers = it->second;
-      for (auto& guard_manager : guard_managers) {
-        if (guard_manager) {
-          guard_manager->disable_recursive_dict_tag_optimization();
+    dict_to_guard_managers.withLock([&](DictToGuardManagersMap& map) {
+      auto it = map.find(dict);
+      if (it != map.end()) {
+        // Copy the list — unwatch_all_saved_dict_pointers may mutate it.
+        auto guard_managers = it->second;
+        for (auto& guard_manager : guard_managers) {
+          if (guard_manager) {
+            guard_manager->disable_recursive_dict_tag_optimization(map);
+          }
         }
       }
-    }
+    });
   }
   return 0; // keep watching
 }
@@ -4410,7 +4545,7 @@ std::unique_ptr<GuardManager> make_guard_manager(
       return std::make_unique<DictGuardManager>(
           root, std::move(source), example_value);
     } else {
-      throw py::type_error("Invalid guard manager enum");
+      throw py::type_error("Invalid guard manager enum"); // @allow-raw-throw
     }
   }
   return std::make_unique<GuardManager>(root, std::move(source), example_value);
@@ -4441,11 +4576,19 @@ void record_tensor_pointer(RootGuardManager* root, PyObject* tensor_pointer) {
   root->record_tensor_pointer(tensor_pointer);
 }
 
+void record_tensor_metadata(RootGuardManager* root, PyObject* tensor_pointer) {
+  root->record_tensor_metadata(tensor_pointer);
+}
+
 std::shared_ptr<RelationalGuard> get_no_tensor_aliasing_guard(
     RootGuardManager* _root) {
   return _root->get_no_tensor_aliasing_guard();
 }
 
+const LocalState& get_local_state(RootGuardManager* root) {
+  return root->_local_state;
+}
+
 // std::string get_compile_id(RootGuardManager* root) {
 //   return root->get_compile_id();
 // }
@@ -4594,13 +4737,15 @@ class TENSOR_MATCH : public LeafGuard {
 
     if (Py_TYPE(value) != _tensor_check->pytype) {
       std::stringstream fail_reason;
-      PyObject* type_str = PyObject_Str(PyObject_Type(value));
+      PyObject* type_str =
+          PyObject_Str(reinterpret_cast<PyObject*>(Py_TYPE(value)));
       fail_reason << "expected type of '" << _tensor_name
                   << "' to be a tensor type, ";
       if (!type_str) {
         fail_reason << "but found a different type";
       } else {
         fail_reason << "' but found " << PyUnicode_AsUTF8(type_str);
+        Py_DECREF(type_str);
       }
       return GuardDebugInfo(false, fail_reason.str(), 0);
     }
@@ -4956,7 +5101,8 @@ class FrameLocalsGuardAccessor : public GuardAccessor {
             guard_manager_enum),
         _key(key[0].ptr()),
         _framelocals_idx(key[1].cast<int>()),
-        _is_immutable_object(is_immutable_object(example_value)) {}
+        _is_immutable_object(is_immutable_object(example_value)),
+        _is_tensor(THPVariable_Check(example_value.ptr())) {}
 
   // Run as a result of calling run_root_guard_manager/check_nopybind
   // NB: Intentional duplication between check_nopybind and
@@ -4964,8 +5110,7 @@ class FrameLocalsGuardAccessor : public GuardAccessor {
   bool check_nopybind(
       FrameLocalsMapping* obj,
       bool matches_dict_tag = false) override { // borrowed ref
-    if (matches_dict_tag && _is_immutable_object) {
-      // immutable object and dict tag matches, we can skip the guard subtree.
+    if (matches_dict_tag && _is_immutable_object && !_is_tensor) {
       return true;
     }
 
@@ -4988,8 +5133,7 @@ class FrameLocalsGuardAccessor : public GuardAccessor {
         PyDict_Check(obj),
         "FrameLocalsGuardAccessor check expected dict() input");
 
-    if (matches_dict_tag && _is_immutable_object) {
-      // immutable object and dict tag matches, we can skip the guard subtree.
+    if (matches_dict_tag && _is_immutable_object && !_is_tensor) {
       return true;
     }
 
@@ -5046,6 +5190,7 @@ class FrameLocalsGuardAccessor : public GuardAccessor {
     to->_key = _key;
     to->_framelocals_idx = _framelocals_idx;
     to->_is_immutable_object = _is_immutable_object;
+    to->_is_tensor = _is_tensor;
   }
 
  private:
@@ -5055,6 +5200,7 @@ class FrameLocalsGuardAccessor : public GuardAccessor {
   // If immutable object and dict tag matches, we can skip the guard subtree and
   // return true.
   bool _is_immutable_object{false};
+  bool _is_tensor{false};
 };
 
 /**
@@ -5077,18 +5223,15 @@ class DictGetItemGuardAccessor : public GuardAccessor {
             example_value,
             guard_manager_enum),
         _key(key.ptr()),
-        _is_immutable_object(is_immutable_object(example_value)) {}
+        _is_immutable_object(is_immutable_object(example_value)),
+        _is_tensor(THPVariable_Check(example_value.ptr())) {}
 
   // NB: Intentional duplication between check_nopybind and
   // check_verbose_nopybind.
   bool check_nopybind(PyObject* obj, bool matches_dict_tag = false) override {
-    if (matches_dict_tag && _is_immutable_object &&
+    if (matches_dict_tag && _is_immutable_object && !_is_tensor &&
         !is_recording_dict_pointers(get_guard_manager()->get_root()) &&
         _guard_manager->has_no_accessors()) {
-      // immutable object and dict tag matches, we can skip the guard subtree.
-      // NB: We only skip the subtree if there are no accessors in the subtree.
-      // This is specifically for tensors which are used in symbolic shape C++
-      // guards, and therefore have accessors on the tensor GuardManager itself.
       return true;
     }
 
@@ -5135,6 +5278,7 @@ class DictGetItemGuardAccessor : public GuardAccessor {
   void clone_visitor(DictGetItemGuardAccessor* to) {
     to->_key = _key;
     to->_is_immutable_object = _is_immutable_object;
+    to->_is_tensor = _is_tensor;
   }
 
  private:
@@ -5143,6 +5287,7 @@ class DictGetItemGuardAccessor : public GuardAccessor {
   // If immutable object and dict tag matches, we can skip the guard subtree and
   // return true.
   bool _is_immutable_object{false};
+  bool _is_tensor{false};
 };
 
 /**
@@ -5933,22 +6078,47 @@ class TypeDictGuardAccessor : public GuardAccessor {
 
   // NB: Intentional duplication between check_nopybind and
   // check_verbose_nopybind.
+  //
+  // In CPython 3.12+, types with Py_TPFLAGS_MANAGED_DICT (e.g. enum.Enum)
+  // can have tp_dict=NULL because the dict is lazily materialized. Use
+  // PyType_GetDict() which handles this transparently. It returns a new
+  // reference, so we must decref after use.
   bool check_nopybind(PyObject* obj, bool matches_dict_tag = false)
       override { // borrowed ref
+#if PY_VERSION_HEX >= 0x030C0000
+    PyObject* x = PyType_GetDict((PyTypeObject*)obj); // new ref
+    if (x == nullptr) {
+      return false;
+    }
+    bool result = _guard_manager->check_nopybind(x);
+    Py_DECREF(x);
+    return result;
+#else
     PyObject* x = ((PyTypeObject*)obj)->tp_dict; // borrowed ref
     if (x == nullptr) {
       return false;
     }
     return _guard_manager->check_nopybind(x);
+#endif
   }
 
   GuardDebugInfo check_verbose_nopybind(
       PyObject* obj) override { // borrowed ref
+#if PY_VERSION_HEX >= 0x030C0000
+    PyObject* x = PyType_GetDict((PyTypeObject*)obj); // new ref
+    if (x == nullptr) {
+      return GuardDebugInfo(false, "null type dict on " + repr(), 0);
+    }
+    auto result = _guard_manager->check_verbose_nopybind(x);
+    Py_DECREF(x);
+    return result;
+#else
     PyObject* x = ((PyTypeObject*)obj)->tp_dict; // borrowed ref
     if (x == nullptr) {
       return GuardDebugInfo(false, "null type dict on " + repr(), 0);
     }
     return _guard_manager->check_verbose_nopybind(x);
+#endif
   }
 
   std::string repr() const override {
@@ -6877,6 +7047,30 @@ PyObject* torch_c_dynamo_guards_init() {
     return nullptr;
   }
 
+  // Expose THPVariable_Wrap for cpp_wrapper inductor, for fbcode
+  {
+    using WrapFn = PyObject* (*)(const at::TensorBase&);
+    WrapFn fn = &THPVariable_Wrap;
+    if (PyModule_AddObject(
+            m,
+            "_torchinductor_thp_variable_wrap",
+            PyLong_FromVoidPtr(reinterpret_cast<void*>(fn))) < 0) {
+      return nullptr;
+    }
+  }
+
+  // Expose THPUtils_unpackInt for cpp_wrapper inductor, for fbcode
+  {
+    using UnpackFn = int32_t (*)(PyObject*);
+    UnpackFn fn = &THPUtils_unpackInt;
+    if (PyModule_AddObject(
+            m,
+            "_torchinductor_thputils_unpack_int",
+            PyLong_FromVoidPtr(reinterpret_cast<void*>(fn))) < 0) {
+      return nullptr;
+    }
+  }
+
   auto py_m = py::handle(m).cast<py::module>();
   py::class_<GuardDebugInfo, std::unique_ptr<GuardDebugInfo>>(
       py_m, "GuardDebugInfo")
@@ -8107,7 +8301,7 @@ PyObject* torch_c_dynamo_guards_init() {
   py_m.def("install_symbolic_shape_guard", install_symbolic_shape_guard);
   py_m.def("profile_guard_manager", profile_guard_manager);
 
-// initialize dict_version_map watcher for 3.12
+// initialize dict version watcher for 3.12
 #if IS_PYTHON_3_12_PLUS
 
   dict_version_watcher_id = PyDict_AddWatcher(dict_version_watch_callback);
diff --git a/torch/csrc/dynamo/init.cpp b/torch/csrc/dynamo/init.cpp
index 244d79bd15b55..92ad88c771f01 100644
--- a/torch/csrc/dynamo/init.cpp
+++ b/torch/csrc/dynamo/init.cpp
@@ -158,10 +158,238 @@ PyObject* _is_valid_var_name(
   return result.release();
 }
 
+// Slot bit position definitions (each int64_t has independent bit positions)
+
+enum class PySequenceSlotBit : int64_t {
+  SQ_LENGTH = 0,
+  SQ_CONCAT = 1,
+  SQ_REPEAT = 2,
+  SQ_ITEM = 3,
+  SQ_CONTAINS = 4,
+  SQ_ASS_ITEM = 5,
+  SQ_INPLACE_CONCAT = 6,
+  SQ_INPLACE_REPEAT = 7,
+};
+
+enum class PyMappingSlotBit : int64_t {
+  MP_LENGTH = 0,
+  MP_SUBSCRIPT = 1,
+  MP_ASS_SUBSCRIPT = 2,
+};
+
+enum class PyNumberSlotBit : int64_t {
+  NB_ADD = 0,
+  NB_SUBTRACT = 1,
+  NB_MULTIPLY = 2,
+  NB_REMAINDER = 3,
+  NB_POWER = 4,
+  NB_NEGATIVE = 5,
+  NB_POSITIVE = 6,
+  NB_ABSOLUTE = 7,
+  NB_BOOL = 8,
+  NB_INVERT = 9,
+  NB_LSHIFT = 10,
+  NB_RSHIFT = 11,
+  NB_AND = 12,
+  NB_XOR = 13,
+  NB_OR = 14,
+  NB_INT = 15,
+  NB_FLOAT = 16,
+  NB_INPLACE_ADD = 17,
+  NB_INPLACE_SUBTRACT = 18,
+  NB_INPLACE_MULTIPLY = 19,
+  NB_INPLACE_REMAINDER = 20,
+  NB_INPLACE_POWER = 21,
+  NB_INPLACE_LSHIFT = 22,
+  NB_INPLACE_RSHIFT = 23,
+  NB_INPLACE_AND = 24,
+  NB_INPLACE_XOR = 25,
+  NB_INPLACE_OR = 26,
+  NB_FLOOR_DIVIDE = 27,
+  NB_TRUE_DIVIDE = 28,
+  NB_INPLACE_FLOOR_DIVIDE = 29,
+  NB_INPLACE_TRUE_DIVIDE = 30,
+  NB_INDEX = 31,
+  NB_MATRIX_MULTIPLY = 32,
+  NB_INPLACE_MATRIX_MULTIPLY = 33,
+};
+
+enum class PyTypeSlotBit : int64_t {
+  TP_HASH = 0,
+  TP_ITER = 1,
+  TP_ITERNEXT = 2,
+  TP_CALL = 3,
+  TP_REPR = 4,
+  TP_RICHCOMPARE = 5,
+  TP_GETATTRO = 6,
+  TP_SETATTRO = 7,
+  TP_DESCR_GET = 8,
+  TP_DESCR_SET = 9,
+};
+
+int64_t get_pysequence_slots(PyTypeObject* type) {
+  int64_t slots = 0;
+  if (PyType_GetSlot(type, Py_sq_length) != nullptr)
+    slots |= (1LL << static_cast<int>(PySequenceSlotBit::SQ_LENGTH));
+  if (PyType_GetSlot(type, Py_sq_concat) != nullptr)
+    slots |= (1LL << static_cast<int>(PySequenceSlotBit::SQ_CONCAT));
+  if (PyType_GetSlot(type, Py_sq_repeat) != nullptr)
+    slots |= (1LL << static_cast<int>(PySequenceSlotBit::SQ_REPEAT));
+  if (PyType_GetSlot(type, Py_sq_item) != nullptr)
+    slots |= (1LL << static_cast<int>(PySequenceSlotBit::SQ_ITEM));
+  if (PyType_GetSlot(type, Py_sq_contains) != nullptr)
+    slots |= (1LL << static_cast<int>(PySequenceSlotBit::SQ_CONTAINS));
+  if (PyType_GetSlot(type, Py_sq_ass_item) != nullptr)
+    slots |= (1LL << static_cast<int>(PySequenceSlotBit::SQ_ASS_ITEM));
+  if (PyType_GetSlot(type, Py_sq_inplace_concat) != nullptr)
+    slots |= (1LL << static_cast<int>(PySequenceSlotBit::SQ_INPLACE_CONCAT));
+  if (PyType_GetSlot(type, Py_sq_inplace_repeat) != nullptr)
+    slots |= (1LL << static_cast<int>(PySequenceSlotBit::SQ_INPLACE_REPEAT));
+  return slots;
+}
+
+int64_t get_pymapping_slots(PyTypeObject* type) {
+  int64_t slots = 0;
+  if (PyType_GetSlot(type, Py_mp_length) != nullptr)
+    slots |= (1LL << static_cast<int>(PyMappingSlotBit::MP_LENGTH));
+  if (PyType_GetSlot(type, Py_mp_subscript) != nullptr)
+    slots |= (1LL << static_cast<int>(PyMappingSlotBit::MP_SUBSCRIPT));
+  if (PyType_GetSlot(type, Py_mp_ass_subscript) != nullptr)
+    slots |= (1LL << static_cast<int>(PyMappingSlotBit::MP_ASS_SUBSCRIPT));
+  return slots;
+}
+
+int64_t get_pynumber_slots(PyTypeObject* type) {
+  int64_t slots = 0;
+  if (PyType_GetSlot(type, Py_nb_add) != nullptr)
+    slots |= (1LL << static_cast<int>(PyNumberSlotBit::NB_ADD));
+  if (PyType_GetSlot(type, Py_nb_subtract) != nullptr)
+    slots |= (1LL << static_cast<int>(PyNumberSlotBit::NB_SUBTRACT));
+  if (PyType_GetSlot(type, Py_nb_multiply) != nullptr)
+    slots |= (1LL << static_cast<int>(PyNumberSlotBit::NB_MULTIPLY));
+  if (PyType_GetSlot(type, Py_nb_remainder) != nullptr)
+    slots |= (1LL << static_cast<int>(PyNumberSlotBit::NB_REMAINDER));
+  if (PyType_GetSlot(type, Py_nb_power) != nullptr)
+    slots |= (1LL << static_cast<int>(PyNumberSlotBit::NB_POWER));
+  if (PyType_GetSlot(type, Py_nb_negative) != nullptr)
+    slots |= (1LL << static_cast<int>(PyNumberSlotBit::NB_NEGATIVE));
+  if (PyType_GetSlot(type, Py_nb_positive) != nullptr)
+    slots |= (1LL << static_cast<int>(PyNumberSlotBit::NB_POSITIVE));
+  if (PyType_GetSlot(type, Py_nb_absolute) != nullptr)
+    slots |= (1LL << static_cast<int>(PyNumberSlotBit::NB_ABSOLUTE));
+  if (PyType_GetSlot(type, Py_nb_bool) != nullptr)
+    slots |= (1LL << static_cast<int>(PyNumberSlotBit::NB_BOOL));
+  if (PyType_GetSlot(type, Py_nb_invert) != nullptr)
+    slots |= (1LL << static_cast<int>(PyNumberSlotBit::NB_INVERT));
+  if (PyType_GetSlot(type, Py_nb_lshift) != nullptr)
+    slots |= (1LL << static_cast<int>(PyNumberSlotBit::NB_LSHIFT));
+  if (PyType_GetSlot(type, Py_nb_rshift) != nullptr)
+    slots |= (1LL << static_cast<int>(PyNumberSlotBit::NB_RSHIFT));
+  if (PyType_GetSlot(type, Py_nb_and) != nullptr)
+    slots |= (1LL << static_cast<int>(PyNumberSlotBit::NB_AND));
+  if (PyType_GetSlot(type, Py_nb_xor) != nullptr)
+    slots |= (1LL << static_cast<int>(PyNumberSlotBit::NB_XOR));
+  if (PyType_GetSlot(type, Py_nb_or) != nullptr)
+    slots |= (1LL << static_cast<int>(PyNumberSlotBit::NB_OR));
+  if (PyType_GetSlot(type, Py_nb_int) != nullptr)
+    slots |= (1LL << static_cast<int>(PyNumberSlotBit::NB_INT));
+  if (PyType_GetSlot(type, Py_nb_float) != nullptr)
+    slots |= (1LL << static_cast<int>(PyNumberSlotBit::NB_FLOAT));
+  if (PyType_GetSlot(type, Py_nb_inplace_add) != nullptr)
+    slots |= (1LL << static_cast<int>(PyNumberSlotBit::NB_INPLACE_ADD));
+  if (PyType_GetSlot(type, Py_nb_inplace_subtract) != nullptr)
+    slots |= (1LL << static_cast<int>(PyNumberSlotBit::NB_INPLACE_SUBTRACT));
+  if (PyType_GetSlot(type, Py_nb_inplace_multiply) != nullptr)
+    slots |= (1LL << static_cast<int>(PyNumberSlotBit::NB_INPLACE_MULTIPLY));
+  if (PyType_GetSlot(type, Py_nb_inplace_remainder) != nullptr)
+    slots |= (1LL << static_cast<int>(PyNumberSlotBit::NB_INPLACE_REMAINDER));
+  if (PyType_GetSlot(type, Py_nb_inplace_power) != nullptr)
+    slots |= (1LL << static_cast<int>(PyNumberSlotBit::NB_INPLACE_POWER));
+  if (PyType_GetSlot(type, Py_nb_inplace_lshift) != nullptr)
+    slots |= (1LL << static_cast<int>(PyNumberSlotBit::NB_INPLACE_LSHIFT));
+  if (PyType_GetSlot(type, Py_nb_inplace_rshift) != nullptr)
+    slots |= (1LL << static_cast<int>(PyNumberSlotBit::NB_INPLACE_RSHIFT));
+  if (PyType_GetSlot(type, Py_nb_inplace_and) != nullptr)
+    slots |= (1LL << static_cast<int>(PyNumberSlotBit::NB_INPLACE_AND));
+  if (PyType_GetSlot(type, Py_nb_inplace_xor) != nullptr)
+    slots |= (1LL << static_cast<int>(PyNumberSlotBit::NB_INPLACE_XOR));
+  if (PyType_GetSlot(type, Py_nb_inplace_or) != nullptr)
+    slots |= (1LL << static_cast<int>(PyNumberSlotBit::NB_INPLACE_OR));
+  if (PyType_GetSlot(type, Py_nb_floor_divide) != nullptr)
+    slots |= (1LL << static_cast<int>(PyNumberSlotBit::NB_FLOOR_DIVIDE));
+  if (PyType_GetSlot(type, Py_nb_true_divide) != nullptr)
+    slots |= (1LL << static_cast<int>(PyNumberSlotBit::NB_TRUE_DIVIDE));
+  if (PyType_GetSlot(type, Py_nb_inplace_floor_divide) != nullptr)
+    slots |=
+        (1LL << static_cast<int>(PyNumberSlotBit::NB_INPLACE_FLOOR_DIVIDE));
+  if (PyType_GetSlot(type, Py_nb_inplace_true_divide) != nullptr)
+    slots |= (1LL << static_cast<int>(PyNumberSlotBit::NB_INPLACE_TRUE_DIVIDE));
+  if (PyType_GetSlot(type, Py_nb_index) != nullptr)
+    slots |= (1LL << static_cast<int>(PyNumberSlotBit::NB_INDEX));
+  if (PyType_GetSlot(type, Py_nb_matrix_multiply) != nullptr)
+    slots |= (1LL << static_cast<int>(PyNumberSlotBit::NB_MATRIX_MULTIPLY));
+  if (PyType_GetSlot(type, Py_nb_inplace_matrix_multiply) != nullptr)
+    slots |=
+        (1LL << static_cast<int>(PyNumberSlotBit::NB_INPLACE_MATRIX_MULTIPLY));
+  return slots;
+}
+
+int64_t get_pytype_slots(PyTypeObject* type) {
+  int64_t slots = 0;
+  if (PyType_GetSlot(type, Py_tp_hash) != nullptr)
+    slots |= (1LL << static_cast<int>(PyTypeSlotBit::TP_HASH));
+  if (PyType_GetSlot(type, Py_tp_iter) != nullptr)
+    slots |= (1LL << static_cast<int>(PyTypeSlotBit::TP_ITER));
+  if (PyType_GetSlot(type, Py_tp_iternext) != nullptr)
+    slots |= (1LL << static_cast<int>(PyTypeSlotBit::TP_ITERNEXT));
+  if (PyType_GetSlot(type, Py_tp_call) != nullptr)
+    slots |= (1LL << static_cast<int>(PyTypeSlotBit::TP_CALL));
+  if (PyType_GetSlot(type, Py_tp_repr) != nullptr)
+    slots |= (1LL << static_cast<int>(PyTypeSlotBit::TP_REPR));
+  if (PyType_GetSlot(type, Py_tp_richcompare) != nullptr)
+    slots |= (1LL << static_cast<int>(PyTypeSlotBit::TP_RICHCOMPARE));
+  if (PyType_GetSlot(type, Py_tp_getattro) != nullptr)
+    slots |= (1LL << static_cast<int>(PyTypeSlotBit::TP_GETATTRO));
+  if (PyType_GetSlot(type, Py_tp_setattro) != nullptr)
+    slots |= (1LL << static_cast<int>(PyTypeSlotBit::TP_SETATTRO));
+  if (PyType_GetSlot(type, Py_tp_descr_get) != nullptr)
+    slots |= (1LL << static_cast<int>(PyTypeSlotBit::TP_DESCR_GET));
+  if (PyType_GetSlot(type, Py_tp_descr_set) != nullptr)
+    slots |= (1LL << static_cast<int>(PyTypeSlotBit::TP_DESCR_SET));
+  return slots;
+}
+
+PyObject* _get_type_slots(
+    PyObject* self,
+    PyObject* const* args,
+    Py_ssize_t nargs) {
+  if (!_checkParamCount(nargs, 1)) {
+    return nullptr;
+  }
+
+  PyObject* arg = args[0];
+  PyTypeObject* type = PyType_Check(arg) ? (PyTypeObject*)arg : Py_TYPE(arg);
+
+  int64_t seq_slots = get_pysequence_slots(type);
+  int64_t map_slots = get_pymapping_slots(type);
+  int64_t num_slots = get_pynumber_slots(type);
+  int64_t type_slots = get_pytype_slots(type);
+
+  PyObject* tuple = PyTuple_New(4);
+  if (tuple == nullptr) {
+    return nullptr;
+  }
+  PyTuple_SetItem(tuple, 0, PyLong_FromLongLong(seq_slots));
+  PyTuple_SetItem(tuple, 1, PyLong_FromLongLong(map_slots));
+  PyTuple_SetItem(tuple, 2, PyLong_FromLongLong(num_slots));
+  PyTuple_SetItem(tuple, 3, PyLong_FromLongLong(type_slots));
+  return tuple;
+}
+
 #define PYC_FN(x) ((PyCFunction)(void (*)()) & x)
 
 void _register_functions(PyObject* mod) {
-  static std::array<PyMethodDef, 3> fns = {
+  static std::array<PyMethodDef, 4> fns = {
       PyMethodDef{
           "strip_function_call",
           PYC_FN(_strip_function_call),
@@ -172,6 +400,8 @@ void _register_functions(PyObject* mod) {
           PYC_FN(_is_valid_var_name),
           METH_FASTCALL,
           nullptr},
+      PyMethodDef{
+          "get_type_slots", PYC_FN(_get_type_slots), METH_FASTCALL, nullptr},
       PyMethodDef{nullptr, nullptr, 0, nullptr},
   };
   PyModule_AddFunctions(mod, fns.data());
@@ -182,7 +412,7 @@ void _register_functions(PyObject* mod) {
 void initDynamoBindings(PyObject* torch) {
   PyObject* dynamo = PyModule_Create(&_module);
   if (dynamo == nullptr || PyModule_AddObject(torch, "_dynamo", dynamo) != 0) {
-    throw python_error();
+    throw python_error(); // @allow-raw-throw
   }
 #ifdef Py_GIL_DISABLED
   PyUnstable_Module_SetGIL(dynamo, Py_MOD_GIL_NOT_USED);
@@ -191,23 +421,23 @@ void initDynamoBindings(PyObject* torch) {
   PyObject* eval_frame = torch_c_dynamo_eval_frame_init();
   if (eval_frame == nullptr ||
       PyModule_AddObject(dynamo, "eval_frame", eval_frame) != 0) {
-    throw python_error();
+    throw python_error(); // @allow-raw-throw
   }
 
   PyObject* utils = torch_c_dynamo_utils_init();
   if (utils == nullptr || PyModule_AddObject(dynamo, "utils", utils) != 0) {
-    throw python_error();
+    throw python_error(); // @allow-raw-throw
   }
 
   PyObject* guards = torch_c_dynamo_guards_init();
   if (guards == nullptr || PyModule_AddObject(dynamo, "guards", guards) != 0) {
-    throw python_error();
+    throw python_error(); // @allow-raw-throw
   }
 
   PyObject* compiled_autograd = torch_c_dynamo_compiled_autograd_init();
   if (compiled_autograd == nullptr ||
       PyModule_AddObject(dynamo, "compiled_autograd", compiled_autograd) != 0) {
-    throw python_error();
+    throw python_error(); // @allow-raw-throw
   }
 
   auto m = py::handle(eval_frame).cast<py::module>();
@@ -252,6 +482,17 @@ void initDynamoBindings(PyObject* torch) {
   m.def("_load_precompile_entry", &_load_precompile_entry);
   m.def("_debug_get_precompile_entries", &_debug_get_precompile_entries);
   m.def("_set_lru_cache", &_set_lru_cache);
+  m.def(
+      "_get_frame_value_stack_with_depth", &_get_frame_value_stack_with_depth);
+  m.def("set_bytecode_debugger_callback", &set_bytecode_debugger_callback);
+  m.def("get_bytecode_debugger_callback", &get_bytecode_debugger_callback);
+  m.def("register_breakpoint_code", &register_breakpoint_code);
+
+  // NullStackValue - sentinel for NULL stack values
+  py::class_<NullStackValue>(m, "NullStackValue")
+      .def("__repr__", [](const NullStackValue&) { return "<NULL>"; });
+  m.attr("NULL_STACK_VALUE") = get_null_stack_value();
+
   py::bind_vector<std::vector<uint8_t>>(m, "VectorUInt8");
   init_THPCaches();
   if (THP_PyOpcode_Caches != nullptr) {
@@ -264,12 +505,77 @@ void initDynamoBindings(PyObject* torch) {
   m.def("code_framelocals_names", &code_framelocals_names);
   _register_functions(dynamo);
 
-  py::enum_<EvalFrameOverride>(m, "_EvalFrameOverride")
-      .value("NONE", EvalFrameOverride::NONE)
-      .value("SKIP", EvalFrameOverride::SKIP)
-      .value("ERROR", EvalFrameOverride::ERROR);
-
-  m.def("set_eval_frame_override", &set_eval_frame_override);
+  auto dynamo_module = py::handle(dynamo).cast<py::module>();
+  dynamo_module.def("has_slot", [](int64_t slots, py::object slot_bit_obj) {
+    // Convert slot_bit to int - handle both int and pybind11 enums
+    int64_t slot_bit = py::cast<int64_t>(slot_bit_obj.attr("__index__")());
+    return (slots & (1LL << slot_bit)) != 0;
+  });
+  py::enum_<PySequenceSlotBit>(dynamo_module, "PySequenceSlots")
+      .value("SQ_LENGTH", PySequenceSlotBit::SQ_LENGTH)
+      .value("SQ_CONCAT", PySequenceSlotBit::SQ_CONCAT)
+      .value("SQ_REPEAT", PySequenceSlotBit::SQ_REPEAT)
+      .value("SQ_ITEM", PySequenceSlotBit::SQ_ITEM)
+      .value("SQ_CONTAINS", PySequenceSlotBit::SQ_CONTAINS)
+      .value("SQ_ASS_ITEM", PySequenceSlotBit::SQ_ASS_ITEM)
+      .value("SQ_INPLACE_CONCAT", PySequenceSlotBit::SQ_INPLACE_CONCAT)
+      .value("SQ_INPLACE_REPEAT", PySequenceSlotBit::SQ_INPLACE_REPEAT);
+
+  py::enum_<PyMappingSlotBit>(dynamo_module, "PyMappingSlots")
+      .value("MP_LENGTH", PyMappingSlotBit::MP_LENGTH)
+      .value("MP_SUBSCRIPT", PyMappingSlotBit::MP_SUBSCRIPT)
+      .value("MP_ASS_SUBSCRIPT", PyMappingSlotBit::MP_ASS_SUBSCRIPT);
+
+  py::enum_<PyNumberSlotBit>(dynamo_module, "PyNumberSlots")
+      .value("NB_ADD", PyNumberSlotBit::NB_ADD)
+      .value("NB_SUBTRACT", PyNumberSlotBit::NB_SUBTRACT)
+      .value("NB_MULTIPLY", PyNumberSlotBit::NB_MULTIPLY)
+      .value("NB_REMAINDER", PyNumberSlotBit::NB_REMAINDER)
+      .value("NB_POWER", PyNumberSlotBit::NB_POWER)
+      .value("NB_NEGATIVE", PyNumberSlotBit::NB_NEGATIVE)
+      .value("NB_POSITIVE", PyNumberSlotBit::NB_POSITIVE)
+      .value("NB_ABSOLUTE", PyNumberSlotBit::NB_ABSOLUTE)
+      .value("NB_BOOL", PyNumberSlotBit::NB_BOOL)
+      .value("NB_INVERT", PyNumberSlotBit::NB_INVERT)
+      .value("NB_LSHIFT", PyNumberSlotBit::NB_LSHIFT)
+      .value("NB_RSHIFT", PyNumberSlotBit::NB_RSHIFT)
+      .value("NB_AND", PyNumberSlotBit::NB_AND)
+      .value("NB_XOR", PyNumberSlotBit::NB_XOR)
+      .value("NB_OR", PyNumberSlotBit::NB_OR)
+      .value("NB_INT", PyNumberSlotBit::NB_INT)
+      .value("NB_FLOAT", PyNumberSlotBit::NB_FLOAT)
+      .value("NB_INPLACE_ADD", PyNumberSlotBit::NB_INPLACE_ADD)
+      .value("NB_INPLACE_SUBTRACT", PyNumberSlotBit::NB_INPLACE_SUBTRACT)
+      .value("NB_INPLACE_MULTIPLY", PyNumberSlotBit::NB_INPLACE_MULTIPLY)
+      .value("NB_INPLACE_REMAINDER", PyNumberSlotBit::NB_INPLACE_REMAINDER)
+      .value("NB_INPLACE_POWER", PyNumberSlotBit::NB_INPLACE_POWER)
+      .value("NB_INPLACE_LSHIFT", PyNumberSlotBit::NB_INPLACE_LSHIFT)
+      .value("NB_INPLACE_RSHIFT", PyNumberSlotBit::NB_INPLACE_RSHIFT)
+      .value("NB_INPLACE_AND", PyNumberSlotBit::NB_INPLACE_AND)
+      .value("NB_INPLACE_XOR", PyNumberSlotBit::NB_INPLACE_XOR)
+      .value("NB_INPLACE_OR", PyNumberSlotBit::NB_INPLACE_OR)
+      .value("NB_FLOOR_DIVIDE", PyNumberSlotBit::NB_FLOOR_DIVIDE)
+      .value("NB_TRUE_DIVIDE", PyNumberSlotBit::NB_TRUE_DIVIDE)
+      .value(
+          "NB_INPLACE_FLOOR_DIVIDE", PyNumberSlotBit::NB_INPLACE_FLOOR_DIVIDE)
+      .value("NB_INPLACE_TRUE_DIVIDE", PyNumberSlotBit::NB_INPLACE_TRUE_DIVIDE)
+      .value("NB_INDEX", PyNumberSlotBit::NB_INDEX)
+      .value("NB_MATRIX_MULTIPLY", PyNumberSlotBit::NB_MATRIX_MULTIPLY)
+      .value(
+          "NB_INPLACE_MATRIX_MULTIPLY",
+          PyNumberSlotBit::NB_INPLACE_MATRIX_MULTIPLY);
+
+  py::enum_<PyTypeSlotBit>(dynamo_module, "PyTypeSlots")
+      .value("TP_HASH", PyTypeSlotBit::TP_HASH)
+      .value("TP_ITER", PyTypeSlotBit::TP_ITER)
+      .value("TP_ITERNEXT", PyTypeSlotBit::TP_ITERNEXT)
+      .value("TP_CALL", PyTypeSlotBit::TP_CALL)
+      .value("TP_REPR", PyTypeSlotBit::TP_REPR)
+      .value("TP_RICHCOMPARE", PyTypeSlotBit::TP_RICHCOMPARE)
+      .value("TP_GETATTRO", PyTypeSlotBit::TP_GETATTRO)
+      .value("TP_SETATTRO", PyTypeSlotBit::TP_SETATTRO)
+      .value("TP_DESCR_GET", PyTypeSlotBit::TP_DESCR_GET)
+      .value("TP_DESCR_SET", PyTypeSlotBit::TP_DESCR_SET);
 }
 
 } // namespace torch::dynamo
diff --git a/torch/csrc/functorch/init.cpp b/torch/csrc/functorch/init.cpp
index 32e781ce43056..57669b84bcc3f 100644
--- a/torch/csrc/functorch/init.cpp
+++ b/torch/csrc/functorch/init.cpp
@@ -19,6 +19,7 @@
 #include <ATen/functorch/PlumbingHelper.h>
 #include <ATen/functorch/TensorWrapper.h>
 #include <c10/core/AutogradState.h>
+#include <c10/core/InferenceMode.h>
 
 #include <iostream>
 
@@ -246,13 +247,36 @@ static RandomnessType get_randomness_enum(const std::string& randomness) {
 static int64_t _grad_increment_nesting() {
   // See NOTE [grad and vjp interaction with no_grad]
   bool prev_grad_mode = c10::GradMode::is_enabled();
+  // When inference_mode is on, new tensors lack autograd dispatch keys
+  // (TensorImpl strips them in its constructor). Toggle the flag off so
+  // tensors created inside the transform can participate in autograd.
+  // Uses AutogradState::set_inference_mode — not the InferenceMode RAII
+  // guard, which would clobber grad_mode and fw_grad_mode.
+  bool prev_inference_mode = c10::InferenceMode::is_enabled();
+  if (prev_inference_mode) {
+    auto state = c10::AutogradState::get_tls_state();
+    state.set_inference_mode(false);
+    c10::AutogradState::set_tls_state(state);
+  }
   return initAndPushDynamicLayer(
-      TransformType::Grad, std::nullopt, std::nullopt, prev_grad_mode);
+      TransformType::Grad,
+      std::nullopt,
+      std::nullopt,
+      prev_grad_mode,
+      std::nullopt,
+      std::nullopt,
+      prev_inference_mode);
 }
 
 static int64_t _grad_decrement_nesting() {
   auto layer = popDynamicLayerAndDeleteMetadata();
   TORCH_INTERNAL_ASSERT(layer.key() == TransformType::Grad);
+  auto& meta = std::get<GradInterpreterMeta>(layer.interpreter().meta());
+  if (meta.prevInferenceMode_) {
+    auto state = c10::AutogradState::get_tls_state();
+    state.set_inference_mode(true);
+    c10::AutogradState::set_tls_state(state);
+  }
   return layer.layerId();
 }
 
@@ -260,17 +284,31 @@ static int64_t _jvp_increment_nesting() {
   // See NOTE [grad and vjp interaction with no_grad]
   bool prev_fwd_grad_mode =
       c10::AutogradState::get_tls_state().get_fw_grad_mode();
+  bool prev_inference_mode = c10::InferenceMode::is_enabled();
+  if (prev_inference_mode) {
+    auto state = c10::AutogradState::get_tls_state();
+    state.set_inference_mode(false);
+    c10::AutogradState::set_tls_state(state);
+  }
   return initAndPushDynamicLayer(
       TransformType::Jvp,
       std::nullopt,
       std::nullopt,
       std::nullopt,
-      prev_fwd_grad_mode);
+      prev_fwd_grad_mode,
+      std::nullopt,
+      prev_inference_mode);
 }
 
 static int64_t _jvp_decrement_nesting() {
   auto layer = popDynamicLayerAndDeleteMetadata();
   TORCH_INTERNAL_ASSERT(layer.key() == TransformType::Jvp);
+  auto& meta = std::get<JvpInterpreterMeta>(layer.interpreter().meta());
+  if (meta.prevInferenceMode_) {
+    auto state = c10::AutogradState::get_tls_state();
+    state.set_inference_mode(true);
+    c10::AutogradState::set_tls_state(state);
+  }
   return layer.layerId();
 }
 
@@ -416,9 +454,15 @@ static void dump_local_tls() {
 namespace {
 
 // Pop the DynamicLayer stack until it's at the given depth.
+// Used by Dynamo for error-recovery cleanup of the transform stack.
+//
+// NB: we peek at .back() to determine the type, then call the
+// type-specific decrement helper which does the actual pop.
+// Do NOT pop before the switch — the helpers call
+// popDynamicLayerAndDeleteMetadata() internally.
 void popDynamicLayerStackToDepth(size_t depth) {
   while (at::functorch::getDynamicLayerStack().size() > depth) {
-    const auto top = popDynamicLayer();
+    const auto& top = at::functorch::getDynamicLayerStack().back();
     switch (top.key()) {
       case at::functorch::TransformType::Vmap:
         _vmap_decrement_nesting();
diff --git a/torch/csrc/inductor/aoti_eager/kernel_holder.cpp b/torch/csrc/inductor/aoti_eager/kernel_holder.cpp
index 5a8956c5c2354..6e8d40f31e7df 100644
--- a/torch/csrc/inductor/aoti_eager/kernel_holder.cpp
+++ b/torch/csrc/inductor/aoti_eager/kernel_holder.cpp
@@ -157,14 +157,16 @@ std::vector<ParameterMetadata> unpack_input_parameters(
 AOTIPythonKernelHolder::AOTIPythonKernelHolder(
     c10::DispatchKey dispatch_key,
     std::string_view ns,
-    std::string_view op_name_with_overload)
+    std::string_view op_name_with_overload,
+    bool dynamic)
     : dispatch_key_(dispatch_key),
       ns_(std::string(ns)),
       op_name_with_overload_(std::string(op_name_with_overload)),
       device_(c10::dispatchKeyToDeviceType(dispatch_key_), 0),
-      pyinterpreter_(getPyInterpreter()) {
+      pyinterpreter_(getPyInterpreter()),
+      dynamic_(dynamic) {
   auto device_name = c10::DeviceTypeName(device_.type());
-  auto registered_aoti_runner = getAOTIModelRunnerRegistry();
+  auto& registered_aoti_runner = getAOTIModelRunnerRegistry();
   TORCH_CHECK(
       device_.type() == c10::DeviceType::CUDA ||
           device_.type() == c10::DeviceType::XPU ||
@@ -196,11 +198,15 @@ bool AOTIPythonKernelHolder::cache_lookup(
     const torch::jit::Stack* stack,
     AOTIKernelMetadata& aoti_kernel_metadata) {
   TORCH_CHECK_NOT_IMPLEMENTED(
-      op.schema().returns().size() == 1,
-      "Not implemented for operations that return either multiple values or no value.");
-  TORCH_CHECK_NOT_IMPLEMENTED(
-      op.schema().returns()[0].type()->isSubtypeOf(c10::TensorType::get()),
-      "Not implemented for operations that return a non-Tensor value.");
+      op.schema().returns().size() >= 1,
+      "Not implemented for operations that return no value.");
+  for (const auto& ret : op.schema().returns()) {
+    TORCH_CHECK_NOT_IMPLEMENTED(
+        ret.type()->isSubtypeOf(c10::TensorType::get()),
+        "Not implemented for operations that return a non-Tensor value. "
+        "Got return type: ",
+        ret.type()->str());
+  }
 
   auto inputs_metadata =
       unpack_input_parameters(op.schema().arguments(), *stack);
@@ -405,7 +411,7 @@ void AOTIPythonKernelHolder::init_aoti_kernel_cache() {
 std::shared_ptr<AOTIModelContainerRunner> AOTIPythonKernelHolder::
     load_aoti_model_runner(const std::string& so_path) {
   auto device_name = c10::DeviceTypeName(device_.type());
-  auto registered_aoti_runner = getAOTIModelRunnerRegistry();
+  auto& registered_aoti_runner = getAOTIModelRunnerRegistry();
   TORCH_CHECK(
       device_.type() == c10::DeviceType::CUDA ||
           device_.type() == c10::DeviceType::XPU ||
@@ -431,6 +437,13 @@ std::shared_ptr<AOTIModelContainerRunner> AOTIPythonKernelHolder::
   } else if (device_.type() == c10::DeviceType::CPU) {
     return std::make_shared<AOTIModelContainerRunnerCpu>(so_path);
   } else {
+    TORCH_CHECK(
+        registered_aoti_runner.find(device_name) !=
+            registered_aoti_runner.end(),
+        "AOTI eager mode: no runner registered for device ",
+        device_name,
+        ". External device backends must register an AOTI model runner factory "
+        "using RegisterAOTIModelRunner before using AOTI eager mode.");
     auto aoti_model_runer_fn = registered_aoti_runner[device_name];
     return aoti_model_runer_fn(
         so_path, 1, device_name, "", /*run_single_threaded=*/false);
@@ -448,6 +461,16 @@ void AOTIPythonKernelHolder::cache_miss(
       kernel != nullptr,
       "Unsupported device: ",
       c10::DeviceTypeName(device_.type()));
+
+  // Populate the in-memory cache so subsequent calls hit cache_hit.
+  auto inputs_metadata =
+      unpack_input_parameters(op.schema().arguments(), *stack);
+  AOTIKernelMetadata aoti_kernel_metadata;
+  aoti_kernel_metadata.parameter_metadata_list_ = std::move(inputs_metadata);
+  aoti_kernel_metadata.kernel_runner_ = kernel;
+  aoti_kernel_metadata.is_dynamic_ = dynamic_;
+  aoti_kernel_cache_.push_back(std::move(aoti_kernel_metadata));
+
   auto inputs = unpack_tensors(op.schema().arguments(), *stack, device_);
   auto outputs = kernel->run(inputs);
   torch::jit::drop(*stack, op.schema().arguments().size());
@@ -513,7 +536,7 @@ std::string AOTIPythonKernelHolder::produce_aoti_kernel_lib(
       py::str(ns_str).ptr(),
       py::str(op_name_with_overload_).ptr(),
       py::str(c10::DeviceTypeName(device_.type(), true)).ptr(),
-      py::bool_(false).ptr(),
+      py::bool_(dynamic_).ptr(),
       op_py_func.ptr(),
       args_kwargs.first.ptr(),
       args_kwargs.second.ptr(),
diff --git a/torch/csrc/inductor/aoti_eager/kernel_holder.h b/torch/csrc/inductor/aoti_eager/kernel_holder.h
index 1575481148a07..131ed9f5d9218 100644
--- a/torch/csrc/inductor/aoti_eager/kernel_holder.h
+++ b/torch/csrc/inductor/aoti_eager/kernel_holder.h
@@ -21,6 +21,11 @@ struct AOTIKernelMetadata {
   std::vector<ParameterMetadata> parameter_metadata_list_;
   // AOTI model runner to run the AOTI kernel
   std::shared_ptr<AOTIModelContainerRunner> kernel_runner_;
+  // Whether this kernel was compiled with dynamic shapes. When true, cache
+  // matching skips exact size/stride comparison on tensors and matches by
+  // dtype/device/rank instead, allowing a single compiled kernel to serve
+  // multiple input shapes.
+  bool is_dynamic_{false};
   AOTIKernelMetadata() : kernel_runner_(nullptr) {}
 
   // Check whether the given parameter metadata list is the same as the
@@ -32,10 +37,16 @@ struct AOTIKernelMetadata {
     }
 
     for (size_t i = 0; i < parameter_metadata_list_.size(); ++i) {
-      if (parameter_metadata_list_[i] == parameter_metadata_list[i]) {
-        continue;
+      if (is_dynamic_) {
+        // Dynamic shapes: match by type/dtype/device/rank, skip exact sizes
+        if (!parameter_metadata_list_[i].dynamic_check(
+                parameter_metadata_list[i])) {
+          return false;
+        }
       } else {
-        return false;
+        if (!(parameter_metadata_list_[i] == parameter_metadata_list[i])) {
+          return false;
+        }
       }
     }
 
@@ -64,12 +75,15 @@ class AOTIPythonKernelHolder : public c10::OperatorKernel {
   c10::impl::PyInterpreter* pyinterpreter_;
   // Cache the produced kernels by AOTI and its metadata
   std::vector<AOTIKernelMetadata> aoti_kernel_cache_;
+  // Whether to compile with dynamic shapes support
+  bool dynamic_;
 
  public:
   AOTIPythonKernelHolder(
       c10::DispatchKey dispatch_key,
       std::string_view ns,
-      std::string_view op_name_with_overload);
+      std::string_view op_name_with_overload,
+      bool dynamic = false);
 
   void operator()(
       const c10::OperatorHandle& op,
diff --git a/torch/csrc/inductor/aoti_eager/kernel_meta_info.cpp b/torch/csrc/inductor/aoti_eager/kernel_meta_info.cpp
index 25cd32b6b52fe..911665f89f676 100644
--- a/torch/csrc/inductor/aoti_eager/kernel_meta_info.cpp
+++ b/torch/csrc/inductor/aoti_eager/kernel_meta_info.cpp
@@ -92,6 +92,15 @@ bool TensorMetadata::operator==(const TensorMetadata& other) const {
   }
 }
 
+bool TensorMetadata::dynamic_check(const TensorMetadata& other) const {
+  // Match by dtype, device, and rank (number of dimensions) but skip
+  // exact sizes/strides so one compiled kernel serves multiple shapes.
+  return this->dtype_ == other.dtype_ && this->device_ == other.device_ &&
+      this->dispatch_key_set_ == other.dispatch_key_set_ &&
+      this->requires_grad_ == other.requires_grad_ &&
+      this->sizes_.size() == other.sizes_.size();
+}
+
 std::ostream& operator<<(
     std::ostream& stream,
     const TensorMetadata& tensor_metadata) {
@@ -210,5 +219,44 @@ bool ParameterMetadata::equal_to(const c10::Scalar& scalar) const {
   return false;
 }
 
+bool ParameterMetadata::dynamic_check(const ParameterMetadata& other) const {
+  if (tag_ != other.tag_ || order_ != other.order_) {
+    return false;
+  }
+
+  switch (tag_) {
+    case TENSOR: {
+      const auto& self_tm = std::get<TensorMetadata>(value_);
+      const auto& other_tm = std::get<TensorMetadata>(other.value_);
+      return self_tm.dynamic_check(other_tm);
+    }
+    case TENSOR_LIST: {
+      const auto& self_list = std::get<std::vector<TensorMetadata>>(value_);
+      const auto& other_list =
+          std::get<std::vector<TensorMetadata>>(other.value_);
+      if (self_list.size() != other_list.size()) {
+        return false;
+      }
+      for (size_t i = 0; i < self_list.size(); ++i) {
+        if (!self_list[i].dynamic_check(other_list[i])) {
+          return false;
+        }
+      }
+      return true;
+    }
+    // Non-tensor parameters use exact matching even in dynamic mode
+    case SCALAR:
+      return equal_to(std::get<c10::Scalar>(other.value_));
+    case STRING:
+      return std::get<std::string>(value_) ==
+          std::get<std::string>(other.value_);
+    case DEVICE:
+      return std::get<c10::Device>(value_) ==
+          std::get<c10::Device>(other.value_);
+    default:
+      return false;
+  }
+}
+
 } // namespace torch::inductor
 #endif
diff --git a/torch/csrc/inductor/aoti_eager/kernel_meta_info.h b/torch/csrc/inductor/aoti_eager/kernel_meta_info.h
index 15560377a0a69..5b52bfaff4998 100644
--- a/torch/csrc/inductor/aoti_eager/kernel_meta_info.h
+++ b/torch/csrc/inductor/aoti_eager/kernel_meta_info.h
@@ -69,6 +69,10 @@ struct TensorMetadata {
 
   // Compare two TensorMetadata objects
   bool operator==(const TensorMetadata& other) const;
+
+  // Dynamic-shape-aware comparison: matches by dtype/device/rank but
+  // skips exact sizes/strides comparison.
+  bool dynamic_check(const TensorMetadata& other) const;
 };
 
 // ParameterTag is to represent the type of the input parameters of a aten
@@ -132,6 +136,11 @@ struct ParameterMetadata {
 
   bool operator==(const ParameterMetadata& other) const;
 
+  // Dynamic-shape-aware comparison: matches by type/dtype/device/rank but
+  // skips exact size/stride comparison for tensors, so a single compiled
+  // kernel can serve multiple input shapes.
+  bool dynamic_check(const ParameterMetadata& other) const;
+
  private:
   // Helper function to compare two ParameterMetadata objects with the same
   // SCALAR tag.
diff --git a/torch/csrc/inductor/aoti_runner/model_container_runner.cpp b/torch/csrc/inductor/aoti_runner/model_container_runner.cpp
index 4930ecad6c452..31881ffe6745a 100644
--- a/torch/csrc/inductor/aoti_runner/model_container_runner.cpp
+++ b/torch/csrc/inductor/aoti_runner/model_container_runner.cpp
@@ -21,6 +21,8 @@
 
 namespace torch::inductor {
 
+AOTIModelContainerRunner::AOTIModelContainerRunner() = default;
+
 AOTIModelContainerRunner::AOTIModelContainerRunner(
     const std::string& model_so_path,
     size_t num_models,
@@ -126,9 +128,13 @@ consider rebuild your model with the latest AOTInductor.");
 }
 
 AOTIModelContainerRunner::~AOTIModelContainerRunner() {
-  AOTIRuntimeError result = delete_func_(container_handle_);
-  TORCH_CHECK(
-      result == AOTI_RUNTIME_SUCCESS, "AOTInductorModelContainerDelete failed");
+  // Custom device implementations don't set delete_func_
+  if (delete_func_ != nullptr) {
+    AOTIRuntimeError result = delete_func_(container_handle_);
+    TORCH_CHECK(
+        result == AOTI_RUNTIME_SUCCESS,
+        "AOTInductorModelContainerDelete failed");
+  }
 }
 
 std::vector<at::Tensor> AOTIModelContainerRunner::run_impl(
diff --git a/torch/csrc/inductor/aoti_runner/model_container_runner.h b/torch/csrc/inductor/aoti_runner/model_container_runner.h
index e9bf0daf7267c..b8119217ba7d6 100644
--- a/torch/csrc/inductor/aoti_runner/model_container_runner.h
+++ b/torch/csrc/inductor/aoti_runner/model_container_runner.h
@@ -15,7 +15,6 @@ using TensorConstantMap = std::unordered_map<std::string, at::Tensor*>;
 
 class TORCH_API AOTIModelContainerRunner {
  public:
-  AOTIModelContainerRunner() = delete;
   AOTIModelContainerRunner(const AOTIModelContainerRunner& other) = delete;
   AOTIModelContainerRunner(AOTIModelContainerRunner&& other) = delete;
   AOTIModelContainerRunner& operator=(const AOTIModelContainerRunner& other) =
@@ -67,6 +66,10 @@ class TORCH_API AOTIModelContainerRunner {
       const std::string& cubin_dir,
       const bool run_single_threaded);
 
+  // Default constructor for custom device implementations that don't
+  // use .so files. Derived classes must override run_impl().
+  AOTIModelContainerRunner();
+
   virtual std::vector<at::Tensor> run_impl(
       std::vector<AtenTensorHandle>& input_handles,
       void* stream_handle);
@@ -107,7 +110,7 @@ class TORCH_API AOTIModelContainerRunner {
 
   AOTInductorModelContainerHandle container_handle_ = nullptr;
 
-  AOTIProxyExecutorHandle proxy_executor_handle_;
+  AOTIProxyExecutorHandle proxy_executor_handle_ = nullptr;
 
  private:
   std::unique_ptr<torch::aot_inductor::ProxyExecutor> proxy_executor_;
diff --git a/torch/csrc/inductor/aoti_runtime/arrayref_tensor.h b/torch/csrc/inductor/aoti_runtime/arrayref_tensor.h
index 6ca660093c67d..3f551677c4e8a 100644
--- a/torch/csrc/inductor/aoti_runtime/arrayref_tensor.h
+++ b/torch/csrc/inductor/aoti_runtime/arrayref_tensor.h
@@ -39,6 +39,8 @@ inline bool is_contiguous_strides_for_shape(
 template <typename T>
 class ArrayRefTensor {
  public:
+  using value_type = T;
+
   ArrayRefTensor() = default;
 
   explicit ArrayRefTensor(
diff --git a/torch/csrc/inductor/aoti_runtime/arrayref_tensor_conversion.h b/torch/csrc/inductor/aoti_runtime/arrayref_tensor_conversion.h
new file mode 100644
index 0000000000000..9a45c4171575b
--- /dev/null
+++ b/torch/csrc/inductor/aoti_runtime/arrayref_tensor_conversion.h
@@ -0,0 +1,85 @@
+#pragma once
+
+// Zero-copy conversion utilities between ArrayRefTensor<T> (C++ template) and
+// AOTInductorArrayRefTensor (plain C struct).
+//
+// These helpers allow the host process to marshal ArrayRefTensor objects into
+// the C-compatible AOTInductorArrayRefTensor descriptors before calling into a
+// DSO, and to unmarshal the descriptors back after the call.  Because only
+// C types cross the DSO boundary, the host and DSO can be linked against
+// different C++ standard libraries (e.g. libc++ vs libstdc++) without ABI
+// conflicts.
+//
+// IMPORTANT: Both sides share the same underlying data buffers -- no copies
+// are made.  The caller must ensure the data remains valid for the lifetime
+// of the descriptor.
+
+#include <torch/csrc/inductor/aoti_runtime/arrayref_tensor.h>
+#include <torch/csrc/inductor/aoti_runtime/interface.h>
+
+#include <cassert>
+#include <cstring>
+#include <stdexcept>
+#include <type_traits>
+
+namespace torch::aot_inductor {
+
+inline void validate_arrayref_tensor_ndim(int32_t ndim) {
+  if (ndim < 0 || ndim > AOTI_ARRAYREF_TENSOR_MAX_DIMS) {
+    throw std::runtime_error(
+        "AOTInductorArrayRefTensor ndim exceeds AOTI_ARRAYREF_TENSOR_MAX_DIMS");
+  }
+}
+
+// -------------------------------------------------------------------------
+// ArrayRefTensor<T> --> AOTInductorArrayRefTensor  (zero-copy)
+// -------------------------------------------------------------------------
+template <typename T>
+inline void arrayref_tensor_to_c(
+    const ArrayRefTensor<T>& src,
+    AOTInductorArrayRefTensor& dst) {
+  const auto sizes = src.sizes();
+  const auto strides = src.strides();
+  dst.data = const_cast<void*>(static_cast<const void*>(src.data()));
+  dst.numel = static_cast<int64_t>(src.numel());
+  dst.ndim = static_cast<int32_t>(sizes.size());
+  dst.dtype = aoti_torch_dtype<std::remove_const_t<T>>();
+  dst.device_type = src.device_type();
+  dst.device_idx = src.device_idx();
+
+  validate_arrayref_tensor_ndim(dst.ndim);
+  assert(dst.ndim <= AOTI_ARRAYREF_TENSOR_MAX_DIMS);
+  std::memcpy(dst.sizes, sizes.data(), dst.ndim * sizeof(int64_t));
+  std::memcpy(dst.strides, strides.data(), dst.ndim * sizeof(int64_t));
+  const int32_t remaining = AOTI_ARRAYREF_TENSOR_MAX_DIMS - dst.ndim;
+  std::memset(dst.sizes + dst.ndim, 0, remaining * sizeof(int64_t));
+  std::memset(dst.strides + dst.ndim, 0, remaining * sizeof(int64_t));
+  std::memset(dst.reserved, 0, sizeof(dst.reserved));
+}
+
+template <typename T>
+inline AOTInductorArrayRefTensor arrayref_tensor_to_c(
+    const ArrayRefTensor<T>& src) {
+  AOTInductorArrayRefTensor dst;
+  arrayref_tensor_to_c(src, dst);
+  return dst;
+}
+
+// -------------------------------------------------------------------------
+// AOTInductorArrayRefTensor --> ArrayRefTensor<T>  (zero-copy)
+// -------------------------------------------------------------------------
+template <typename T>
+inline ArrayRefTensor<T> c_to_arrayref_tensor(
+    const AOTInductorArrayRefTensor& src) {
+  validate_arrayref_tensor_ndim(src.ndim);
+  return ArrayRefTensor<T>(
+      MiniArrayRef<T>(
+          static_cast<T*>(const_cast<void*>(src.data)),
+          static_cast<size_t>(src.numel)),
+      MiniArrayRef<const int64_t>(src.sizes, static_cast<size_t>(src.ndim)),
+      MiniArrayRef<const int64_t>(src.strides, static_cast<size_t>(src.ndim)),
+      src.device_type,
+      src.device_idx);
+}
+
+} // namespace torch::aot_inductor
diff --git a/torch/csrc/inductor/aoti_runtime/interface.h b/torch/csrc/inductor/aoti_runtime/interface.h
index ffccdd94e5be2..a47bb3e4c67ff 100644
--- a/torch/csrc/inductor/aoti_runtime/interface.h
+++ b/torch/csrc/inductor/aoti_runtime/interface.h
@@ -35,6 +35,52 @@ struct AOTInductorConstantMapEntry {
   AtenTensorHandle handle;
 };
 
+// ---------------------------------------------------------------------------
+// C-compatible tensor descriptor for crossing the DSO boundary.
+//
+// This struct carries the same information as ArrayRefTensor<T> but uses only
+// C-compatible types so the host process and DSO can be built with different
+// C++ standard libraries (e.g. libc++ vs libstdc++).  All pointer fields
+// reference memory owned by the caller; no copies are made.
+//
+// Maximum supported number of dimensions.  8 covers all practical AOTI
+// models; tensors with more dims should fall back to the AtenTensorHandle
+// interface.
+// ---------------------------------------------------------------------------
+#define AOTI_ARRAYREF_TENSOR_MAX_DIMS 8
+
+struct AOTInductorArrayRefTensor {
+  // Pointer to the raw data buffer.  Not owned.
+  void* data;
+
+  // Number of elements in the data buffer (product of sizes for contiguous
+  // tensors).
+  int64_t numel;
+
+  // Static-size arrays for shape metadata.  Only the first `ndim` entries
+  // are meaningful.
+  int64_t sizes[AOTI_ARRAYREF_TENSOR_MAX_DIMS];
+  int64_t strides[AOTI_ARRAYREF_TENSOR_MAX_DIMS];
+
+  // Number of dimensions (0 <= ndim <= AOTI_ARRAYREF_TENSOR_MAX_DIMS).
+  int32_t ndim;
+
+  // Torch dtype encoded as int32_t (same encoding as aoti_torch_dtype_*()).
+  int32_t dtype;
+
+  // Device information.
+  int32_t device_type;
+  int32_t device_idx;
+
+  // Reserved for future extension.  Zero-initialize and do not read — a
+  // newer reader must tolerate zeros, and an older reader must ignore them.
+  int64_t reserved[4];
+};
+
+static_assert(
+    sizeof(AOTInductorArrayRefTensor) == 192,
+    "changing the size of AOTInductorArrayRefTensor breaks ABI compatibility!");
+
 // TODO: Deprecate this API. This was kept for BC compatibility.
 // Please use AOTInductorModelContainerCreateWithDevice instead.
 AOTI_API AOTIRuntimeError AOTInductorModelContainerCreate(
@@ -265,4 +311,25 @@ AOTI_API AOTIRuntimeError AOTInductorModelContainerGetCallSpec(
     const char** in_spec,
     const char** out_spec);
 
+// ---------------------------------------------------------------------------
+// C-ABI-safe variant of AOTInductorModelRunMinimalArrayrefInterface.
+//
+// Instead of passing std::tuple<ArrayRefTensor<T>...>& (which encodes C++
+// standard library types into the ABI), this function accepts flat C arrays
+// of AOTInductorArrayRefTensor descriptors.  The descriptors reference the
+// same underlying data buffers -- no copies are made.
+//
+// The host process marshals its ArrayRefTensor<T> objects into
+// AOTInductorArrayRefTensor descriptors, calls into the DSO through this
+// pure-C interface, and then unmarshals the output descriptors back.
+// Because only C types cross the DSO boundary, the host and DSO can be
+// built with different C++ standard libraries (e.g. libc++ vs libstdc++).
+// ---------------------------------------------------------------------------
+AOTI_API AOTIRuntimeError AOTInductorModelRunMinimalArrayrefInterfaceV2(
+    AOTInductorModelHandle model_handle,
+    int32_t num_inputs,
+    const AOTInductorArrayRefTensor* inputs,
+    int32_t num_outputs,
+    AOTInductorArrayRefTensor* outputs);
+
 } // extern "C"
diff --git a/torch/csrc/inductor/aoti_runtime/kernel_context_tls.h b/torch/csrc/inductor/aoti_runtime/kernel_context_tls.h
index 1001dac9cc68d..fc6e61abb185d 100644
--- a/torch/csrc/inductor/aoti_runtime/kernel_context_tls.h
+++ b/torch/csrc/inductor/aoti_runtime/kernel_context_tls.h
@@ -14,8 +14,9 @@ struct KernelContext {
   std::string compressed_python_stack;
 
   KernelContext(std::string name, std::string stack)
-      : kernel_name(std::move(name)), python_stack(std::move(stack)) {
-    compressed_python_stack = compress_python_stack(python_stack);
+      : kernel_name(std::move(name)) {
+    python_stack = trim_stack(stack);
+    compressed_python_stack = compress_stack(python_stack);
   }
 
   KernelContext(const KernelContext&) = default;
@@ -24,28 +25,84 @@ struct KernelContext {
   KernelContext& operator=(KernelContext&&) = default;
 
  private:
-  static std::string compress_python_stack(const std::string& stack) {
+  // Strip leading and trailing newlines from stack:
+  // - finds first and last non-newline characters
+  // - outputs substring between them (inclusive)
+  // - returns empty string if stack contains only newlines
+  static std::string trim_stack(const std::string& stack) {
+    std::string res_stack;
+    auto beg = stack.find_first_not_of('\n');
+    if (beg == std::string::npos) {
+      return res_stack;
+    }
+    auto end = stack.find_last_not_of('\n');
+    if (end == std::string::npos) {
+      res_stack = stack.substr(beg);
+    } else {
+      res_stack = stack.substr(beg, end - beg + 1);
+    }
+    return res_stack;
+  }
+
+  // Compress stack into compact format:
+  // - blank lines are treated as stack separators
+  // - lines with 0 or 2 spaces of indentation are parsed as filename, fileline
+  //   and function
+  // - lines with 4 spaces of indentation are parsed as the code snippet
+  // - outputs function[code], filename and fileline (at newline each)
+  // - returns empty string on any parse error
+  static std::string compress_stack(const std::string& stack) {
+    std::string res_stack;
     namespace fs = std::filesystem;
-    char func[129];
-    char path[1025];
-    uint32_t line;
-    int ret;
-    std::string compressed_stack;
+    char function[1025];
+    char filename[1025];
+    uint32_t fileline;
+    int ret, n, ws;
+    const char* p;
     std::stringstream stream{stack};
-    std::string str;
-    std::string fmt = "File \"%1024[^\"]\", line %u, in %128[^\n]\n";
-    while (std::getline(stream, str)) {
-      ret = sscanf(str.c_str(), fmt.c_str(), path, &line, func);
-      if (ret == 3) {
-        compressed_stack += func;
-        compressed_stack += ' ';
-        compressed_stack += fs::path{path}.filename();
-        compressed_stack += ':';
-        compressed_stack += std::to_string(line);
-        compressed_stack += '\n';
+    std::string line;
+    std::string fmt = "File \"%1024[^\"]\", line %u, in %1024[^\n]\n%n";
+    while (std::getline(stream, line)) {
+      // check if new stack
+      if (line.empty()) {
+        res_stack += '\n';
+        continue;
+      }
+      p = line.c_str();
+      ws = 0;
+      while (*p == ' ') {
+        ++p;
+        ++ws;
+      }
+      // check if new file
+      if (ws != 0 && ws != 2) {
+        return {};
+      }
+      ret = sscanf(p, fmt.c_str(), filename, &fileline, function, &n);
+      if (ret != 3) {
+        return {};
+      }
+      if (!std::getline(stream, line)) {
+        return {};
+      }
+      p = line.c_str();
+      ws = 0;
+      while (*p == ' ') {
+        ++p;
+        ++ws;
+      }
+      // check if command
+      if (ws != 4) {
+        return {};
       }
+      res_stack += std::string{function} + '[' + std::string{p} + ']';
+      res_stack += '\n';
+      res_stack += fs::path{filename}.filename();
+      res_stack += '\n';
+      res_stack += std::to_string(fileline);
+      res_stack += '\n';
     }
-    return compressed_stack;
+    return res_stack;
   }
 };
 
diff --git a/torch/csrc/inductor/aoti_runtime/model_container.h b/torch/csrc/inductor/aoti_runtime/model_container.h
index 6a82ae61d3e5f..b594a812237bf 100644
--- a/torch/csrc/inductor/aoti_runtime/model_container.h
+++ b/torch/csrc/inductor/aoti_runtime/model_container.h
@@ -172,9 +172,9 @@ class AOTInductorModelContainer {
           /* use_inactive = */ false,
           /* validate_full_update = */ false);
       const_folded = ConstantState::FOLDED;
-    } else if (constant_folded_ != ConstantState::FOLDED) {
+    } else if (const_folded != ConstantState::FOLDED) {
       throw std::runtime_error(
-          "Unknown constant state: " + toStringConstantState(constant_folded_));
+          "Unknown constant state: " + toStringConstantState(const_folded));
     }
 
     model->run_single_threaded(
diff --git a/torch/csrc/inductor/aoti_runtime/sycl_runtime_wrappers.h b/torch/csrc/inductor/aoti_runtime/sycl_runtime_wrappers.h
index 56bb478396004..b85b69449edc2 100644
--- a/torch/csrc/inductor/aoti_runtime/sycl_runtime_wrappers.h
+++ b/torch/csrc/inductor/aoti_runtime/sycl_runtime_wrappers.h
@@ -134,8 +134,13 @@ static std::unique_ptr<sycl::kernel> _createKernel(
     uint32_t numWarps,
     uint32_t sharedMemory,
     void** params,
-    sycl::queue* queuePtr,
-    uint32_t threadsPerWarp) {
+    sycl::queue* queuePtr) {
+  uint32_t threadsPerWarp = kernelPtr->get_info<
+      sycl::info::kernel_device_specific::compile_sub_group_size>(
+      queuePtr->get_device());
+  if (threadsPerWarp == 0) {
+    threadsPerWarp = 32; // default to 32 if not set
+  }
   std::string kernelName =
       kernelPtr->get_info<sycl::info::kernel::function_name>();
   uint32_t numParams = kernelPtr->get_info<sycl::info::kernel::num_args>();
diff --git a/torch/csrc/inductor/aoti_torch/c/shim.h b/torch/csrc/inductor/aoti_torch/c/shim.h
index 2eda2b218e705..f719828ba00d6 100644
--- a/torch/csrc/inductor/aoti_torch/c/shim.h
+++ b/torch/csrc/inductor/aoti_torch/c/shim.h
@@ -3,6 +3,7 @@
 
 #include <torch/csrc/inductor/aoti_torch/c/macros.h>
 #include <torch/csrc/inductor/aoti_torch/c/shim_deprecated.h>
+#include <torch/headeronly/util/Exception.h>
 
 // This header defines a stable C API for certain ATen functionality in
 // libtorch. The AOTInductor compiled model.so will only refer to this header
@@ -63,6 +64,10 @@ AOTI_TORCH_EXPORT int32_t aoti_torch_dtype_float8_e5m2();
 AOTI_TORCH_EXPORT int32_t aoti_torch_dtype_float8_e4m3fn();
 AOTI_TORCH_EXPORT int32_t aoti_torch_dtype_float8_e5m2fnuz();
 AOTI_TORCH_EXPORT int32_t aoti_torch_dtype_float8_e4m3fnuz();
+#if TORCH_FEATURE_VERSION >= TORCH_VERSION_2_12_0
+AOTI_TORCH_EXPORT int32_t aoti_torch_dtype_float8_e8m0fnu();
+AOTI_TORCH_EXPORT int32_t aoti_torch_dtype_float4_e2m1fn_x2();
+#endif // TORCH_FEATURE_VERSION >= TORCH_VERSION_2_12_0
 AOTI_TORCH_EXPORT int32_t aoti_torch_dtype_bfloat16();
 AOTI_TORCH_EXPORT int32_t aoti_torch_dtype_float16();
 AOTI_TORCH_EXPORT int32_t aoti_torch_dtype_float32();
@@ -629,24 +634,24 @@ AOTI_TORCH_EXPORT void aoti_torch_check(
     const char* msg);
 
 #ifdef STRIP_ERROR_MESSAGES
-#define AOTI_TORCH_CHECK(cond, ...)              \
-  if (!(cond)) {                                 \
-    aoti_torch_check(                            \
-        false,                                   \
-        __func__,                                \
-        __FILE__,                                \
-        static_cast<uint32_t>(__LINE__),         \
-        TORCH_CHECK_MSG(cond, "", __VA_ARGS__)); \
+#define AOTI_TORCH_CHECK(cond, ...)                  \
+  if (!(cond)) {                                     \
+    aoti_torch_check(                                \
+        false,                                       \
+        __func__,                                    \
+        __FILE__,                                    \
+        static_cast<uint32_t>(__LINE__),             \
+        STD_TORCH_CHECK_MSG(cond, "", __VA_ARGS__)); \
   }
 #else
-#define AOTI_TORCH_CHECK(cond, ...)                \
-  if (!(cond)) {                                   \
-    aoti_torch_check(                              \
-        false,                                     \
-        __func__,                                  \
-        __FILE__,                                  \
-        static_cast<uint32_t>(__LINE__),           \
-        TORCH_CHECK_MSG(cond, "", ##__VA_ARGS__)); \
+#define AOTI_TORCH_CHECK(cond, ...)                    \
+  if (!(cond)) {                                       \
+    aoti_torch_check(                                  \
+        false,                                         \
+        __func__,                                      \
+        __FILE__,                                      \
+        static_cast<uint32_t>(__LINE__),               \
+        STD_TORCH_CHECK_MSG(cond, "", ##__VA_ARGS__)); \
   }
 #endif
 
diff --git a/torch/csrc/inductor/aoti_torch/generated/c_shim_cpu.h b/torch/csrc/inductor/aoti_torch/generated/c_shim_cpu.h
index 2f0c28df53506..21ee7438165dc 100644
--- a/torch/csrc/inductor/aoti_torch/generated/c_shim_cpu.h
+++ b/torch/csrc/inductor/aoti_torch/generated/c_shim_cpu.h
@@ -30,6 +30,7 @@ AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu__fft_r2c(AtenTensorHandle self,
 AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu__fused_moving_avg_obs_fq_helper(AtenTensorHandle self, AtenTensorHandle observer_on, AtenTensorHandle fake_quant_on, AtenTensorHandle running_min, AtenTensorHandle running_max, AtenTensorHandle scale, AtenTensorHandle zero_point, double averaging_const, int64_t quant_min, int64_t quant_max, int64_t ch_axis, int32_t per_row_fake_quant, int32_t symmetric_quant, AtenTensorHandle* ret0, AtenTensorHandle* ret1);
 AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu__fused_moving_avg_obs_fq_helper_functional(AtenTensorHandle self, AtenTensorHandle observer_on, AtenTensorHandle fake_quant_on, AtenTensorHandle running_min, AtenTensorHandle running_max, AtenTensorHandle scale, AtenTensorHandle zero_point, double averaging_const, int64_t quant_min, int64_t quant_max, int64_t ch_axis, int32_t per_row_fake_quant, int32_t symmetric_quant, AtenTensorHandle* ret0, AtenTensorHandle* ret1, AtenTensorHandle* ret2, AtenTensorHandle* ret3, AtenTensorHandle* ret4, AtenTensorHandle* ret5);
 AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu__fused_rms_norm(AtenTensorHandle input, const int64_t* normalized_shape, int64_t normalized_shape_len_, AtenTensorHandle* weight, double* eps, AtenTensorHandle* ret0, AtenTensorHandle* ret1);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu__grouped_mm(AtenTensorHandle self, AtenTensorHandle mat2, AtenTensorHandle* offs, AtenTensorHandle* bias, int32_t* out_dtype, AtenTensorHandle* ret0);
 AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu__histogramdd_from_bin_cts(AtenTensorHandle self, const int64_t* bins, int64_t bins_len_, const double** range, int64_t range_len_, AtenTensorHandle* weight, int32_t density, AtenTensorHandle* ret0);
 AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu__int_mm_out(AtenTensorHandle out, AtenTensorHandle self, AtenTensorHandle mat2);
 AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu__pdist_backward(AtenTensorHandle grad, AtenTensorHandle self, double p, AtenTensorHandle pdist, AtenTensorHandle* ret0);
@@ -120,12 +121,18 @@ AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_pow_Tensor_Scalar(AtenTensorHand
 AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_pow_Tensor_Tensor(AtenTensorHandle self, AtenTensorHandle exponent, AtenTensorHandle* ret0);
 AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_rand(const int64_t* size, int64_t size_len_, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, AtenTensorHandle* ret0);
 AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_rand_generator(const int64_t* size, int64_t size_len_, AtenGeneratorHandle* generator, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_rand_like(AtenTensorHandle self, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, int32_t* memory_format, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_rand_like_generator(AtenTensorHandle self, AtenGeneratorHandle* generator, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, int32_t* memory_format, AtenTensorHandle* ret0);
 AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_randint(int64_t high, const int64_t* size, int64_t size_len_, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, AtenTensorHandle* ret0);
 AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_randint_generator(int64_t high, const int64_t* size, int64_t size_len_, AtenGeneratorHandle* generator, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, AtenTensorHandle* ret0);
 AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_randint_low(int64_t low, int64_t high, const int64_t* size, int64_t size_len_, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, AtenTensorHandle* ret0);
 AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_randint_low_out(AtenTensorHandle out, int64_t low, int64_t high, const int64_t* size, int64_t size_len_);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_randint_like(AtenTensorHandle self, int64_t high, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, int32_t* memory_format, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_randint_like_low_dtype(AtenTensorHandle self, int64_t low, int64_t high, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, int32_t* memory_format, AtenTensorHandle* ret0);
 AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_randn(const int64_t* size, int64_t size_len_, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, AtenTensorHandle* ret0);
 AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_randn_generator(const int64_t* size, int64_t size_len_, AtenGeneratorHandle* generator, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_randn_like(AtenTensorHandle self, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, int32_t* memory_format, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_randn_like_generator(AtenTensorHandle self, AtenGeneratorHandle* generator, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, int32_t* memory_format, AtenTensorHandle* ret0);
 AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_randperm(int64_t n, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, AtenTensorHandle* ret0);
 AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_repeat_interleave_Tensor(AtenTensorHandle repeats, int64_t* output_size, AtenTensorHandle* ret0);
 AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_replication_pad1d_backward(AtenTensorHandle grad_output, AtenTensorHandle self, const int64_t* padding, int64_t padding_len_, AtenTensorHandle* ret0);
diff --git a/torch/csrc/inductor/aoti_torch/generated/c_shim_cuda.h b/torch/csrc/inductor/aoti_torch/generated/c_shim_cuda.h
index bb57cabe53093..413a393f6fa7d 100644
--- a/torch/csrc/inductor/aoti_torch/generated/c_shim_cuda.h
+++ b/torch/csrc/inductor/aoti_torch/generated/c_shim_cuda.h
@@ -29,11 +29,15 @@ AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda__embedding_bag_per_sample_weigh
 AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda__fft_c2c(AtenTensorHandle self, const int64_t* dim, int64_t dim_len_, int64_t normalization, int32_t forward, AtenTensorHandle* ret0);
 AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda__fft_r2c(AtenTensorHandle self, const int64_t* dim, int64_t dim_len_, int64_t normalization, int32_t onesided, AtenTensorHandle* ret0);
 AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda__flash_attention_backward(AtenTensorHandle grad_out, AtenTensorHandle query, AtenTensorHandle key, AtenTensorHandle value, AtenTensorHandle out, AtenTensorHandle logsumexp, AtenTensorHandle cum_seq_q, AtenTensorHandle cum_seq_k, int64_t max_q, int64_t max_k, double dropout_p, int32_t is_causal, AtenTensorHandle rng_state, AtenTensorHandle unused, double* scale, int64_t* window_size_left, int64_t* window_size_right, AtenTensorHandle* ret0, AtenTensorHandle* ret1, AtenTensorHandle* ret2);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda__flash_attention_forward_v2(AtenTensorHandle query, AtenTensorHandle key, AtenTensorHandle value, AtenTensorHandle* cum_seq_q, AtenTensorHandle* cum_seq_k, int64_t max_q, int64_t max_k, double dropout_p, int32_t is_causal, int32_t return_debug_mask, double* scale, int64_t* window_size_left, int64_t* window_size_right, AtenTensorHandle* seqused_k, AtenTensorHandle* alibi_slopes, AtenTensorHandle* block_table, int64_t* num_splits, AtenTensorHandle* ret0, AtenTensorHandle* ret1, AtenTensorHandle* ret2, AtenTensorHandle* ret3, AtenTensorHandle* ret4);
 AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda__flash_attention_forward(AtenTensorHandle query, AtenTensorHandle key, AtenTensorHandle value, AtenTensorHandle* cum_seq_q, AtenTensorHandle* cum_seq_k, int64_t max_q, int64_t max_k, double dropout_p, int32_t is_causal, int32_t return_debug_mask, double* scale, int64_t* window_size_left, int64_t* window_size_right, AtenTensorHandle* seqused_k, AtenTensorHandle* alibi_slopes, AtenTensorHandle* ret0, AtenTensorHandle* ret1, AtenTensorHandle* ret2, AtenTensorHandle* ret3, AtenTensorHandle* ret4);
 AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda__flash_attention_forward_quantized(AtenTensorHandle query, AtenTensorHandle key, AtenTensorHandle value, AtenTensorHandle* cum_seq_q, AtenTensorHandle* cum_seq_k, int64_t max_q, int64_t max_k, double dropout_p, int32_t is_causal, int32_t return_debug_mask, AtenTensorHandle* q_descale, AtenTensorHandle* k_descale, AtenTensorHandle* v_descale, double* scale, int64_t* window_size_left, int64_t* window_size_right, AtenTensorHandle* seqused_k, AtenTensorHandle* alibi_slopes, AtenTensorHandle* ret0, AtenTensorHandle* ret1, AtenTensorHandle* ret2, AtenTensorHandle* ret3, AtenTensorHandle* ret4);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda__flash_attention_forward_no_dropout_inplace_v2(AtenTensorHandle out, AtenTensorHandle query, AtenTensorHandle key, AtenTensorHandle value, AtenTensorHandle* cum_seq_q, AtenTensorHandle* cum_seq_k, int64_t max_q, int64_t max_k, double dropout_p, int32_t is_causal, int32_t return_debug_mask, double* scale, int64_t* window_size_left, int64_t* window_size_right, AtenTensorHandle* seqused_k, AtenTensorHandle* alibi_slopes, AtenTensorHandle* block_table, int64_t* num_splits, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda__flash_attention_forward_no_dropout_inplace(AtenTensorHandle out, AtenTensorHandle query, AtenTensorHandle key, AtenTensorHandle value, AtenTensorHandle* cum_seq_q, AtenTensorHandle* cum_seq_k, int64_t max_q, int64_t max_k, double dropout_p, int32_t is_causal, int32_t return_debug_mask, double* scale, int64_t* window_size_left, int64_t* window_size_right, AtenTensorHandle* seqused_k, AtenTensorHandle* alibi_slopes, AtenTensorHandle* block_table, AtenTensorHandle* ret0);
 AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda__fused_moving_avg_obs_fq_helper(AtenTensorHandle self, AtenTensorHandle observer_on, AtenTensorHandle fake_quant_on, AtenTensorHandle running_min, AtenTensorHandle running_max, AtenTensorHandle scale, AtenTensorHandle zero_point, double averaging_const, int64_t quant_min, int64_t quant_max, int64_t ch_axis, int32_t per_row_fake_quant, int32_t symmetric_quant, AtenTensorHandle* ret0, AtenTensorHandle* ret1);
 AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda__fused_moving_avg_obs_fq_helper_functional(AtenTensorHandle self, AtenTensorHandle observer_on, AtenTensorHandle fake_quant_on, AtenTensorHandle running_min, AtenTensorHandle running_max, AtenTensorHandle scale, AtenTensorHandle zero_point, double averaging_const, int64_t quant_min, int64_t quant_max, int64_t ch_axis, int32_t per_row_fake_quant, int32_t symmetric_quant, AtenTensorHandle* ret0, AtenTensorHandle* ret1, AtenTensorHandle* ret2, AtenTensorHandle* ret3, AtenTensorHandle* ret4, AtenTensorHandle* ret5);
 AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda__fused_rms_norm(AtenTensorHandle input, const int64_t* normalized_shape, int64_t normalized_shape_len_, AtenTensorHandle* weight, double* eps, AtenTensorHandle* ret0, AtenTensorHandle* ret1);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda__grouped_mm(AtenTensorHandle self, AtenTensorHandle mat2, AtenTensorHandle* offs, AtenTensorHandle* bias, int32_t* out_dtype, AtenTensorHandle* ret0);
 AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda__int_mm_out(AtenTensorHandle out, AtenTensorHandle self, AtenTensorHandle mat2);
 AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda__pdist_backward(AtenTensorHandle grad, AtenTensorHandle self, double p, AtenTensorHandle pdist, AtenTensorHandle* ret0);
 AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda__pdist_forward(AtenTensorHandle self, double p, AtenTensorHandle* ret0);
@@ -131,12 +135,18 @@ AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_pow_Tensor_Scalar(AtenTensorHan
 AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_pow_Tensor_Tensor(AtenTensorHandle self, AtenTensorHandle exponent, AtenTensorHandle* ret0);
 AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_rand(const int64_t* size, int64_t size_len_, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, AtenTensorHandle* ret0);
 AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_rand_generator(const int64_t* size, int64_t size_len_, AtenGeneratorHandle* generator, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_rand_like(AtenTensorHandle self, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, int32_t* memory_format, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_rand_like_generator(AtenTensorHandle self, AtenGeneratorHandle* generator, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, int32_t* memory_format, AtenTensorHandle* ret0);
 AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_randint(int64_t high, const int64_t* size, int64_t size_len_, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, AtenTensorHandle* ret0);
 AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_randint_generator(int64_t high, const int64_t* size, int64_t size_len_, AtenGeneratorHandle* generator, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, AtenTensorHandle* ret0);
 AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_randint_low(int64_t low, int64_t high, const int64_t* size, int64_t size_len_, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, AtenTensorHandle* ret0);
 AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_randint_low_out(AtenTensorHandle out, int64_t low, int64_t high, const int64_t* size, int64_t size_len_);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_randint_like(AtenTensorHandle self, int64_t high, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, int32_t* memory_format, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_randint_like_low_dtype(AtenTensorHandle self, int64_t low, int64_t high, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, int32_t* memory_format, AtenTensorHandle* ret0);
 AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_randn(const int64_t* size, int64_t size_len_, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, AtenTensorHandle* ret0);
 AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_randn_generator(const int64_t* size, int64_t size_len_, AtenGeneratorHandle* generator, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_randn_like(AtenTensorHandle self, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, int32_t* memory_format, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_randn_like_generator(AtenTensorHandle self, AtenGeneratorHandle* generator, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, int32_t* memory_format, AtenTensorHandle* ret0);
 AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_randperm(int64_t n, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, AtenTensorHandle* ret0);
 AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_repeat_interleave_Tensor(AtenTensorHandle repeats, int64_t* output_size, AtenTensorHandle* ret0);
 AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_replication_pad1d_backward(AtenTensorHandle grad_output, AtenTensorHandle self, const int64_t* padding, int64_t padding_len_, AtenTensorHandle* ret0);
diff --git a/torch/csrc/inductor/aoti_torch/generated/c_shim_mps.h b/torch/csrc/inductor/aoti_torch/generated/c_shim_mps.h
index 592aeecc68f2b..f5fe2cb852d4f 100644
--- a/torch/csrc/inductor/aoti_torch/generated/c_shim_mps.h
+++ b/torch/csrc/inductor/aoti_torch/generated/c_shim_mps.h
@@ -23,7 +23,9 @@ AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps__fft_c2c(AtenTensorHandle self,
 AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps__fft_r2c(AtenTensorHandle self, const int64_t* dim, int64_t dim_len_, int64_t normalization, int32_t onesided, AtenTensorHandle* ret0);
 AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps__fused_moving_avg_obs_fq_helper_functional(AtenTensorHandle self, AtenTensorHandle observer_on, AtenTensorHandle fake_quant_on, AtenTensorHandle running_min, AtenTensorHandle running_max, AtenTensorHandle scale, AtenTensorHandle zero_point, double averaging_const, int64_t quant_min, int64_t quant_max, int64_t ch_axis, int32_t per_row_fake_quant, int32_t symmetric_quant, AtenTensorHandle* ret0, AtenTensorHandle* ret1, AtenTensorHandle* ret2, AtenTensorHandle* ret3, AtenTensorHandle* ret4, AtenTensorHandle* ret5);
 AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps__fused_rms_norm(AtenTensorHandle input, const int64_t* normalized_shape, int64_t normalized_shape_len_, AtenTensorHandle* weight, double* eps, AtenTensorHandle* ret0, AtenTensorHandle* ret1);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps__grouped_mm(AtenTensorHandle self, AtenTensorHandle mat2, AtenTensorHandle* offs, AtenTensorHandle* bias, int32_t* out_dtype, AtenTensorHandle* ret0);
 AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps__histogramdd_from_bin_cts(AtenTensorHandle self, const int64_t* bins, int64_t bins_len_, const double** range, int64_t range_len_, AtenTensorHandle* weight, int32_t density, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps__scaled_dot_product_attention_math_for_mps_v2(AtenTensorHandle query, AtenTensorHandle key, AtenTensorHandle value, AtenTensorHandle* attn_mask, double dropout_p, int32_t is_causal, AtenTensorHandle* dropout_mask, double* scale, int32_t enable_gqa, AtenTensorHandle* ret0, AtenTensorHandle* ret1);
 AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps__scaled_dot_product_attention_math_for_mps(AtenTensorHandle query, AtenTensorHandle key, AtenTensorHandle value, AtenTensorHandle* attn_mask, double dropout_p, int32_t is_causal, AtenTensorHandle* dropout_mask, double* scale, AtenTensorHandle* ret0, AtenTensorHandle* ret1);
 AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps__scaled_dot_product_fused_attention_overrideable(AtenTensorHandle query, AtenTensorHandle key, AtenTensorHandle value, AtenTensorHandle* attn_bias, double dropout_p, int32_t is_causal, int32_t return_debug_mask, double* scale, AtenTensorHandle* ret0, AtenTensorHandle* ret1, AtenTensorHandle* ret2, AtenTensorHandle* ret3, int64_t* ret4, int64_t* ret5, AtenTensorHandle* ret6, AtenTensorHandle* ret7, AtenTensorHandle* ret8);
 AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps__scaled_dot_product_fused_attention_overrideable_backward(AtenTensorHandle grad_out, AtenTensorHandle query, AtenTensorHandle key, AtenTensorHandle value, AtenTensorHandle attn_bias, const int32_t* grad_input_mask, int64_t grad_input_mask_len_, AtenTensorHandle out, AtenTensorHandle logsumexp, AtenTensorHandle cum_seq_q, AtenTensorHandle cum_seq_k, int64_t max_q, int64_t max_k, double dropout_p, int32_t is_causal, AtenTensorHandle philox_seed, AtenTensorHandle philox_offset, double* scale, AtenTensorHandle* ret0, AtenTensorHandle* ret1, AtenTensorHandle* ret2, AtenTensorHandle* ret3);
@@ -66,6 +68,7 @@ AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_histc(AtenTensorHandle self, int
 AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_histogram_bin_ct(AtenTensorHandle self, int64_t bins, const double** range, int64_t range_len_, AtenTensorHandle* weight, int32_t density, AtenTensorHandle* ret0, AtenTensorHandle* ret1);
 AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_index_Tensor(AtenTensorHandle self, const AtenTensorHandle** indices, int64_t indices_len_, AtenTensorHandle* ret0);
 AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_index_put(AtenTensorHandle self, const AtenTensorHandle** indices, int64_t indices_len_, AtenTensorHandle values, int32_t accumulate, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_index_reduce(AtenTensorHandle self, int64_t dim, AtenTensorHandle index, AtenTensorHandle source, const char* reduce, int32_t include_self, AtenTensorHandle* ret0);
 AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_kthvalue(AtenTensorHandle self, int64_t k, int64_t dim, int32_t keepdim, AtenTensorHandle* ret0, AtenTensorHandle* ret1);
 AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_logcumsumexp(AtenTensorHandle self, int64_t dim, AtenTensorHandle* ret0);
 AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_lu_unpack(AtenTensorHandle LU_data, AtenTensorHandle LU_pivots, int32_t unpack_data, int32_t unpack_pivots, AtenTensorHandle* ret0, AtenTensorHandle* ret1, AtenTensorHandle* ret2);
@@ -86,6 +89,7 @@ AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_nanmedian(AtenTensorHandle self,
 AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_narrow(AtenTensorHandle self, int64_t dim, int64_t start, int64_t length, AtenTensorHandle* ret0);
 AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_native_dropout(AtenTensorHandle input, double p, int32_t* train, AtenTensorHandle* ret0, AtenTensorHandle* ret1);
 AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_nonzero(AtenTensorHandle self, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_nonzero_static(AtenTensorHandle self, int64_t size, int64_t fill_value, AtenTensorHandle* ret0);
 AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_normal_functional(AtenTensorHandle self, double mean, double std, AtenGeneratorHandle* generator, AtenTensorHandle* ret0);
 AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_pad(AtenTensorHandle self, const int64_t* pad, int64_t pad_len_, const char* mode, double* value, AtenTensorHandle* ret0);
 AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_permute(AtenTensorHandle self, const int64_t* dims, int64_t dims_len_, AtenTensorHandle* ret0);
@@ -95,12 +99,18 @@ AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_pow_Tensor_Scalar(AtenTensorHand
 AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_pow_Tensor_Tensor(AtenTensorHandle self, AtenTensorHandle exponent, AtenTensorHandle* ret0);
 AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_rand(const int64_t* size, int64_t size_len_, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, AtenTensorHandle* ret0);
 AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_rand_generator(const int64_t* size, int64_t size_len_, AtenGeneratorHandle* generator, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_rand_like(AtenTensorHandle self, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, int32_t* memory_format, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_rand_like_generator(AtenTensorHandle self, AtenGeneratorHandle* generator, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, int32_t* memory_format, AtenTensorHandle* ret0);
 AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_randint(int64_t high, const int64_t* size, int64_t size_len_, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, AtenTensorHandle* ret0);
 AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_randint_generator(int64_t high, const int64_t* size, int64_t size_len_, AtenGeneratorHandle* generator, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, AtenTensorHandle* ret0);
 AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_randint_low(int64_t low, int64_t high, const int64_t* size, int64_t size_len_, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, AtenTensorHandle* ret0);
 AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_randint_low_out(AtenTensorHandle out, int64_t low, int64_t high, const int64_t* size, int64_t size_len_);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_randint_like(AtenTensorHandle self, int64_t high, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, int32_t* memory_format, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_randint_like_low_dtype(AtenTensorHandle self, int64_t low, int64_t high, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, int32_t* memory_format, AtenTensorHandle* ret0);
 AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_randn(const int64_t* size, int64_t size_len_, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, AtenTensorHandle* ret0);
 AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_randn_generator(const int64_t* size, int64_t size_len_, AtenGeneratorHandle* generator, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_randn_like(AtenTensorHandle self, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, int32_t* memory_format, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_randn_like_generator(AtenTensorHandle self, AtenGeneratorHandle* generator, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, int32_t* memory_format, AtenTensorHandle* ret0);
 AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_randperm(int64_t n, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, AtenTensorHandle* ret0);
 AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_repeat_interleave_Tensor(AtenTensorHandle repeats, int64_t* output_size, AtenTensorHandle* ret0);
 AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_replication_pad1d_backward(AtenTensorHandle grad_output, AtenTensorHandle self, const int64_t* padding, int64_t padding_len_, AtenTensorHandle* ret0);
diff --git a/torch/csrc/inductor/aoti_torch/generated/c_shim_xpu.h b/torch/csrc/inductor/aoti_torch/generated/c_shim_xpu.h
index 20b6e85b5f1a9..ea0cbfdc5659d 100644
--- a/torch/csrc/inductor/aoti_torch/generated/c_shim_xpu.h
+++ b/torch/csrc/inductor/aoti_torch/generated/c_shim_xpu.h
@@ -14,6 +14,7 @@ extern "C" {
 AOTI_TORCH_EXPORT AOTITorchError aoti_torch_xpu__addmm_activation(AtenTensorHandle self, AtenTensorHandle mat1, AtenTensorHandle mat2, double beta, double alpha, int32_t use_gelu, AtenTensorHandle* ret0);
 AOTI_TORCH_EXPORT AOTITorchError aoti_torch_xpu__fused_moving_avg_obs_fq_helper_functional(AtenTensorHandle self, AtenTensorHandle observer_on, AtenTensorHandle fake_quant_on, AtenTensorHandle running_min, AtenTensorHandle running_max, AtenTensorHandle scale, AtenTensorHandle zero_point, double averaging_const, int64_t quant_min, int64_t quant_max, int64_t ch_axis, int32_t per_row_fake_quant, int32_t symmetric_quant, AtenTensorHandle* ret0, AtenTensorHandle* ret1, AtenTensorHandle* ret2, AtenTensorHandle* ret3, AtenTensorHandle* ret4, AtenTensorHandle* ret5);
 AOTI_TORCH_EXPORT AOTITorchError aoti_torch_xpu__fused_rms_norm(AtenTensorHandle input, const int64_t* normalized_shape, int64_t normalized_shape_len_, AtenTensorHandle* weight, double* eps, AtenTensorHandle* ret0, AtenTensorHandle* ret1);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_xpu__grouped_mm(AtenTensorHandle self, AtenTensorHandle mat2, AtenTensorHandle* offs, AtenTensorHandle* bias, int32_t* out_dtype, AtenTensorHandle* ret0);
 AOTI_TORCH_EXPORT AOTITorchError aoti_torch_xpu__int_mm_out(AtenTensorHandle out, AtenTensorHandle self, AtenTensorHandle mat2);
 AOTI_TORCH_EXPORT AOTITorchError aoti_torch_xpu__scaled_dot_product_flash_attention(AtenTensorHandle query, AtenTensorHandle key, AtenTensorHandle value, double dropout_p, int32_t is_causal, int32_t return_debug_mask, double* scale, AtenTensorHandle* ret0, AtenTensorHandle* ret1, AtenTensorHandle* ret2, AtenTensorHandle* ret3, int64_t* ret4, int64_t* ret5, AtenTensorHandle* ret6, AtenTensorHandle* ret7, AtenTensorHandle* ret8);
 AOTI_TORCH_EXPORT AOTITorchError aoti_torch_xpu__scaled_dot_product_flash_attention_backward(AtenTensorHandle grad_out, AtenTensorHandle query, AtenTensorHandle key, AtenTensorHandle value, AtenTensorHandle out, AtenTensorHandle logsumexp, AtenTensorHandle cum_seq_q, AtenTensorHandle cum_seq_k, int64_t max_q, int64_t max_k, double dropout_p, int32_t is_causal, AtenTensorHandle philox_seed, AtenTensorHandle philox_offset, double* scale, AtenTensorHandle* ret0, AtenTensorHandle* ret1, AtenTensorHandle* ret2);
@@ -53,12 +54,18 @@ AOTI_TORCH_EXPORT AOTITorchError aoti_torch_xpu_permute(AtenTensorHandle self, c
 AOTI_TORCH_EXPORT AOTITorchError aoti_torch_xpu_polar(AtenTensorHandle abs, AtenTensorHandle angle, AtenTensorHandle* ret0);
 AOTI_TORCH_EXPORT AOTITorchError aoti_torch_xpu_rand(const int64_t* size, int64_t size_len_, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, AtenTensorHandle* ret0);
 AOTI_TORCH_EXPORT AOTITorchError aoti_torch_xpu_rand_generator(const int64_t* size, int64_t size_len_, AtenGeneratorHandle* generator, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_xpu_rand_like(AtenTensorHandle self, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, int32_t* memory_format, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_xpu_rand_like_generator(AtenTensorHandle self, AtenGeneratorHandle* generator, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, int32_t* memory_format, AtenTensorHandle* ret0);
 AOTI_TORCH_EXPORT AOTITorchError aoti_torch_xpu_randint(int64_t high, const int64_t* size, int64_t size_len_, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, AtenTensorHandle* ret0);
 AOTI_TORCH_EXPORT AOTITorchError aoti_torch_xpu_randint_generator(int64_t high, const int64_t* size, int64_t size_len_, AtenGeneratorHandle* generator, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, AtenTensorHandle* ret0);
 AOTI_TORCH_EXPORT AOTITorchError aoti_torch_xpu_randint_low(int64_t low, int64_t high, const int64_t* size, int64_t size_len_, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, AtenTensorHandle* ret0);
 AOTI_TORCH_EXPORT AOTITorchError aoti_torch_xpu_randint_low_out(AtenTensorHandle out, int64_t low, int64_t high, const int64_t* size, int64_t size_len_);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_xpu_randint_like(AtenTensorHandle self, int64_t high, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, int32_t* memory_format, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_xpu_randint_like_low_dtype(AtenTensorHandle self, int64_t low, int64_t high, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, int32_t* memory_format, AtenTensorHandle* ret0);
 AOTI_TORCH_EXPORT AOTITorchError aoti_torch_xpu_randn(const int64_t* size, int64_t size_len_, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, AtenTensorHandle* ret0);
 AOTI_TORCH_EXPORT AOTITorchError aoti_torch_xpu_randn_generator(const int64_t* size, int64_t size_len_, AtenGeneratorHandle* generator, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_xpu_randn_like(AtenTensorHandle self, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, int32_t* memory_format, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_xpu_randn_like_generator(AtenTensorHandle self, AtenGeneratorHandle* generator, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, int32_t* memory_format, AtenTensorHandle* ret0);
 AOTI_TORCH_EXPORT AOTITorchError aoti_torch_xpu_randperm(int64_t n, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, AtenTensorHandle* ret0);
 AOTI_TORCH_EXPORT AOTITorchError aoti_torch_xpu_reshape(AtenTensorHandle self, const int64_t* shape, int64_t shape_len_, AtenTensorHandle* ret0);
 AOTI_TORCH_EXPORT AOTITorchError aoti_torch_xpu_resize_as_(AtenTensorHandle self, AtenTensorHandle the_template, int32_t* memory_format);
diff --git a/torch/csrc/inductor/aoti_torch/generated_enum_converters.h b/torch/csrc/inductor/aoti_torch/generated_enum_converters.h
new file mode 100644
index 0000000000000..14b97a7f4bc26
--- /dev/null
+++ b/torch/csrc/inductor/aoti_torch/generated_enum_converters.h
@@ -0,0 +1,123 @@
+// @generated by update_schema.py
+// clang-format off
+
+#pragma once
+
+#include <c10/core/Layout.h>
+#include <c10/core/MemoryFormat.h>
+#include <c10/core/ScalarType.h>
+#include <c10/util/Exception.h>
+
+// Converter functions from serialized enum values (torch._export.serde.schema)
+// to c10 enums. The serialized format has different enum values than c10.
+
+namespace torch::aot_inductor {
+
+inline c10::ScalarType convertSerializedScalarType(int serialized_value) {
+  constexpr int kInvalid = -1;
+  constexpr int kScalarTypeMap[] = {
+      kInvalid, // 0
+      static_cast<int>(c10::ScalarType::Byte), // 1
+      static_cast<int>(c10::ScalarType::Char), // 2
+      static_cast<int>(c10::ScalarType::Short), // 3
+      static_cast<int>(c10::ScalarType::Int), // 4
+      static_cast<int>(c10::ScalarType::Long), // 5
+      static_cast<int>(c10::ScalarType::Half), // 6
+      static_cast<int>(c10::ScalarType::Float), // 7
+      static_cast<int>(c10::ScalarType::Double), // 8
+      static_cast<int>(c10::ScalarType::ComplexHalf), // 9
+      static_cast<int>(c10::ScalarType::ComplexFloat), // 10
+      static_cast<int>(c10::ScalarType::ComplexDouble), // 11
+      static_cast<int>(c10::ScalarType::Bool), // 12
+      static_cast<int>(c10::ScalarType::BFloat16), // 13
+      kInvalid, // 14
+      kInvalid, // 15
+      kInvalid, // 16
+      kInvalid, // 17
+      kInvalid, // 18
+      kInvalid, // 19
+      kInvalid, // 20
+      kInvalid, // 21
+      kInvalid, // 22
+      kInvalid, // 23
+      kInvalid, // 24
+      kInvalid, // 25
+      kInvalid, // 26
+      kInvalid, // 27
+      static_cast<int>(c10::ScalarType::UInt16), // 28
+      static_cast<int>(c10::ScalarType::Float8_e4m3fn), // 29
+      static_cast<int>(c10::ScalarType::Float8_e5m2), // 30
+      static_cast<int>(c10::ScalarType::Float8_e4m3fnuz), // 31
+      static_cast<int>(c10::ScalarType::Float8_e5m2fnuz), // 32
+      static_cast<int>(c10::ScalarType::Float8_e8m0fnu), // 33
+      static_cast<int>(c10::ScalarType::UInt32), // 34
+      static_cast<int>(c10::ScalarType::UInt64), // 35
+  };
+  constexpr int kMapSize = sizeof(kScalarTypeMap) / sizeof(kScalarTypeMap[0]);
+
+  TORCH_CHECK(
+      serialized_value >= 0 && serialized_value < kMapSize,
+      "Serialized ScalarType value out of range: ",
+      serialized_value);
+  int result = kScalarTypeMap[serialized_value];
+  TORCH_CHECK(
+      result != kInvalid,
+      "Invalid serialized ScalarType value: ",
+      serialized_value);
+  return static_cast<c10::ScalarType>(result);
+}
+
+
+inline c10::Layout convertSerializedLayout(int serialized_value) {
+  constexpr int kInvalid = -1;
+  constexpr int kLayoutMap[] = {
+      kInvalid, // 0
+      static_cast<int>(c10::Layout::Sparse), // 1
+      static_cast<int>(c10::Layout::SparseCsr), // 2
+      static_cast<int>(c10::Layout::SparseCsc), // 3
+      static_cast<int>(c10::Layout::SparseBsr), // 4
+      static_cast<int>(c10::Layout::SparseBsc), // 5
+      static_cast<int>(c10::Layout::Mkldnn), // 6
+      static_cast<int>(c10::Layout::Strided), // 7
+  };
+  constexpr int kMapSize = sizeof(kLayoutMap) / sizeof(kLayoutMap[0]);
+
+  TORCH_CHECK(
+      serialized_value >= 0 && serialized_value < kMapSize,
+      "Serialized Layout value out of range: ",
+      serialized_value);
+  int result = kLayoutMap[serialized_value];
+  TORCH_CHECK(
+      result != kInvalid,
+      "Invalid serialized Layout value: ",
+      serialized_value);
+  return static_cast<c10::Layout>(result);
+}
+
+
+inline c10::MemoryFormat convertSerializedMemoryFormat(int serialized_value) {
+  constexpr int kInvalid = -1;
+  constexpr int kMemoryFormatMap[] = {
+      kInvalid, // 0
+      static_cast<int>(c10::MemoryFormat::Contiguous), // 1
+      static_cast<int>(c10::MemoryFormat::ChannelsLast), // 2
+      static_cast<int>(c10::MemoryFormat::ChannelsLast3d), // 3
+      static_cast<int>(c10::MemoryFormat::Preserve), // 4
+  };
+  constexpr int kMapSize = sizeof(kMemoryFormatMap) / sizeof(kMemoryFormatMap[0]);
+
+  TORCH_CHECK(
+      serialized_value >= 0 && serialized_value < kMapSize,
+      "Serialized MemoryFormat value out of range: ",
+      serialized_value);
+  int result = kMemoryFormatMap[serialized_value];
+  TORCH_CHECK(
+      result != kInvalid,
+      "Invalid serialized MemoryFormat value: ",
+      serialized_value);
+  return static_cast<c10::MemoryFormat>(result);
+}
+
+} // namespace torch::aot_inductor
+
+// clang-format on
diff --git a/torch/csrc/inductor/aoti_torch/oss_proxy_executor.cpp b/torch/csrc/inductor/aoti_torch/oss_proxy_executor.cpp
index b0a8a09ffc3e4..5c92df4b9455a 100644
--- a/torch/csrc/inductor/aoti_torch/oss_proxy_executor.cpp
+++ b/torch/csrc/inductor/aoti_torch/oss_proxy_executor.cpp
@@ -5,6 +5,7 @@
 
 #include <c10/util/Exception.h>
 #include <c10/util/FileSystem.h>
+#include <torch/csrc/inductor/aoti_torch/generated_enum_converters.h>
 #include <torch/csrc/inductor/aoti_torch/oss_proxy_executor.h>
 #include <torch/csrc/jit/serialization/pickle.h>
 
@@ -19,6 +20,10 @@ bool has_key(
   return map.find(key) != map.end();
 }
 
+using torch::aot_inductor::convertSerializedLayout;
+using torch::aot_inductor::convertSerializedMemoryFormat;
+using torch::aot_inductor::convertSerializedScalarType;
+
 } // namespace
 
 namespace torch::aot_inductor {
@@ -164,7 +169,8 @@ void OSSProxyExecutor::prefill_stack_with_static_arguments(
           index,
           " but got ",
           serialized_arg_type);
-      stack.at(index) = serialized_arg_val.get<c10::ScalarType>();
+      stack.at(index) =
+          convertSerializedScalarType(serialized_arg_val.get<int>());
       break;
     }
     case c10::TypeKind::MemoryFormatType: {
@@ -176,7 +182,8 @@ void OSSProxyExecutor::prefill_stack_with_static_arguments(
           index,
           " but got ",
           serialized_arg_type);
-      stack.at(index) = serialized_arg_val.get<c10::MemoryFormat>();
+      stack.at(index) =
+          convertSerializedMemoryFormat(serialized_arg_val.get<int>());
       break;
     }
     case c10::TypeKind::LayoutType: {
@@ -188,7 +195,7 @@ void OSSProxyExecutor::prefill_stack_with_static_arguments(
           index,
           " but got ",
           serialized_arg_type);
-      stack.at(index) = serialized_arg_val.get<c10::Layout>();
+      stack.at(index) = convertSerializedLayout(serialized_arg_val.get<int>());
       break;
     }
     case c10::TypeKind::DeviceObjType: {
@@ -393,6 +400,32 @@ void OSSProxyExecutor::prefill_stack_with_static_arguments(
       }
       break;
     }
+    case c10::TypeKind::AnyType: {
+      // For Any type, dispatch based on the serialized type
+      if (serialized_arg_type == "as_string") {
+        stack.at(index) = serialized_arg_val.get<std::string>();
+      } else if (serialized_arg_type == "as_int") {
+        dynamic_args.emplace_back(index, DynamicArgType::IntType, 1);
+      } else if (serialized_arg_type == "as_float") {
+        stack.at(index) = serialized_arg_val.get<double>();
+      } else if (serialized_arg_type == "as_bool") {
+        stack.at(index) = serialized_arg_val.get<bool>();
+      } else if (serialized_arg_type == "as_tensor") {
+        dynamic_args.emplace_back(index, DynamicArgType::TensorType, 1);
+      } else if (serialized_arg_type == "as_none") {
+        stack.at(index) = c10::IValue{};
+      } else {
+        TORCH_CHECK(
+            false,
+            "Unsupported serialized type ",
+            serialized_arg_type,
+            " for Any type argument ",
+            index,
+            " in extern kernel ",
+            op_kernel->target_);
+      }
+      break;
+    }
     // TODO: handle the other input types
     default:
       TORCH_CHECK(
diff --git a/torch/csrc/inductor/aoti_torch/shim_common.cpp b/torch/csrc/inductor/aoti_torch/shim_common.cpp
index d6db06af5f2cc..d6c22b096effe 100644
--- a/torch/csrc/inductor/aoti_torch/shim_common.cpp
+++ b/torch/csrc/inductor/aoti_torch/shim_common.cpp
@@ -99,6 +99,8 @@ AOTI_TORCH_DTYPE_IMPL(float8_e5m2, Float8_e5m2)
 AOTI_TORCH_DTYPE_IMPL(float8_e4m3fn, Float8_e4m3fn)
 AOTI_TORCH_DTYPE_IMPL(float8_e5m2fnuz, Float8_e5m2fnuz)
 AOTI_TORCH_DTYPE_IMPL(float8_e4m3fnuz, Float8_e4m3fnuz)
+AOTI_TORCH_DTYPE_IMPL(float8_e8m0fnu, Float8_e8m0fnu)
+AOTI_TORCH_DTYPE_IMPL(float4_e2m1fn_x2, Float4_e2m1fn_x2)
 AOTI_TORCH_DTYPE_IMPL(bfloat16, BFloat16)
 AOTI_TORCH_DTYPE_IMPL(float16, Half)
 AOTI_TORCH_DTYPE_IMPL(float32, Float)
diff --git a/torch/csrc/inductor/array_ref_impl.h b/torch/csrc/inductor/array_ref_impl.h
index 8cfbc12fb2c3d..e36a3653009ad 100644
--- a/torch/csrc/inductor/array_ref_impl.h
+++ b/torch/csrc/inductor/array_ref_impl.h
@@ -1,6 +1,7 @@
 #pragma once
 
 #include <torch/csrc/inductor/aoti_runtime/arrayref_tensor.h>
+#include <torch/csrc/inductor/aoti_runtime/arrayref_tensor_conversion.h>
 #include <torch/csrc/inductor/aoti_runtime/scalar_to_tensor.h>
 #include <torch/csrc/inductor/aoti_runtime/thread_local.h>
 #include <torch/csrc/inductor/aoti_torch/utils.h>
diff --git a/torch/csrc/inductor/cpp_wrapper/common.h b/torch/csrc/inductor/cpp_wrapper/common.h
index a2eebfcc86032..39a6fbbd2cf24 100644
--- a/torch/csrc/inductor/cpp_wrapper/common.h
+++ b/torch/csrc/inductor/cpp_wrapper/common.h
@@ -6,15 +6,30 @@
 #include <utility>
 
 #include <Python.h>
+#if __has_include(<pybind11/gil_simple.h>)
 #include <pybind11/gil_simple.h>
+#else
+// pybind11 < 3.0: gil_simple.h does not exist yet.
+#define PYBIND11_SIMPLE_GIL_MANAGEMENT
+#include <pybind11/gil.h>
+// Provide the _simple aliases so generated code works with either version.
+namespace pybind11 {
+using gil_scoped_acquire_simple = gil_scoped_acquire;
+using gil_scoped_release_simple = gil_scoped_release;
+} // namespace pybind11
+#endif
 
+// Required for custom op dispatch via the stable ABI
+#include <torch/csrc/stable/library.h>
+
+#ifdef TORCH_INDUCTOR_PRECOMPILE_HEADERS
 // Include some often-used cpp_wrapper headers, for precompiling.
 #include <c10/util/BFloat16.h>
 #include <torch/csrc/Device.h>
 #include <torch/csrc/DynamicTypes.h>
-#include <torch/csrc/stable/library.h>
 #include <torch/csrc/utils/pythoncapi_compat.h>
 #include <torch/csrc/utils/tensor_memoryformats.h>
+#endif
 
 namespace py = pybind11; // NOLINT(misc-unused-alias-decls)
 
diff --git a/torch/csrc/inductor/cpp_wrapper/cuda.h b/torch/csrc/inductor/cpp_wrapper/cuda.h
index 782a2b677276a..eb9888356c6ff 100644
--- a/torch/csrc/inductor/cpp_wrapper/cuda.h
+++ b/torch/csrc/inductor/cpp_wrapper/cuda.h
@@ -1,4 +1,10 @@
 #pragma once
-
 #include <torch/csrc/inductor/cpp_wrapper/common.h>
 #include <torch/csrc/inductor/cpp_wrapper/device_internal/cuda.h>
+#include <torch/csrc/inductor/cpp_wrapper/lazy_triton_compile.h>
+
+#ifdef TORCH_INDUCTOR_PRECOMPILE_HEADERS
+#include <ATen/cuda/EmptyTensor.h>
+#include <c10/cuda/CUDAGuard.h>
+#include <c10/cuda/CUDAStream.h>
+#endif
diff --git a/torch/csrc/inductor/cpp_wrapper/lazy_triton_compile.h b/torch/csrc/inductor/cpp_wrapper/lazy_triton_compile.h
new file mode 100644
index 0000000000000..27f6e81198385
--- /dev/null
+++ b/torch/csrc/inductor/cpp_wrapper/lazy_triton_compile.h
@@ -0,0 +1,201 @@
+#pragma once
+
+#include <string>
+
+#include <torch/csrc/inductor/aoti_torch/utils.h>
+#include <torch/csrc/inductor/cpp_wrapper/common.h>
+#if defined(USE_XPU)
+#include <torch/csrc/inductor/cpp_wrapper/device_internal/xpu.h>
+#else
+#include <torch/csrc/inductor/cpp_wrapper/device_internal/cuda.h>
+#endif
+
+struct LazyKernelCompileResult {
+  std::string cubin_path;
+  std::string mangled_name;
+  int num_warps;
+  int shared_mem;
+  int xblock;
+  int yblock;
+  int zblock;
+  int r0block;
+  int rsplit;
+  int rsplit_size;
+  int config_index;
+  int global_scratch;
+  int profile_scratch;
+};
+
+static PyObject* (*_THPVariable_Wrap)(const at::TensorBase&) = nullptr;
+static int32_t (*_THPUtils_unpackInt)(PyObject*) = nullptr;
+
+// Cached module and function references
+static PyObject* triton_lazy_compile_module = nullptr;
+static PyObject* start_kernel_compile = nullptr;
+static PyObject* run_triton_kernel_with_autotune = nullptr;
+
+// Per-module dict for pending kernel compile results (avoids global state
+// collisions when multiple compiled modules produce kernels with the same
+// name).
+static PyObject* _module_pending_kernels = nullptr;
+
+static inline void loadLazyCompileFuncs() {
+  if (triton_lazy_compile_module == nullptr) {
+    triton_lazy_compile_module =
+        PyImport_ImportModule("torch._inductor.runtime.triton_lazy_compile");
+    AOTI_TORCH_CHECK(
+        triton_lazy_compile_module, "Failed to import triton_lazy_compile");
+
+    start_kernel_compile = PyObject_GetAttrString(
+        triton_lazy_compile_module, "start_kernel_compile");
+    AOTI_TORCH_CHECK(
+        start_kernel_compile, "Failed to get start_kernel_compile");
+
+    run_triton_kernel_with_autotune = PyObject_GetAttrString(
+        triton_lazy_compile_module, "run_triton_kernel_with_autotune");
+    AOTI_TORCH_CHECK(
+        run_triton_kernel_with_autotune,
+        "Failed to get run_triton_kernel_with_autotune");
+
+    RAIIPyObject guards_mod = PyImport_ImportModule("torch._C._dynamo.guards");
+    AOTI_TORCH_CHECK(guards_mod, "Failed to import torch._C._dynamo.guards");
+
+    RAIIPyObject wrap_addr =
+        PyObject_GetAttrString(guards_mod, "_torchinductor_thp_variable_wrap");
+    AOTI_TORCH_CHECK(
+        wrap_addr, "Failed to get _torchinductor_thp_variable_wrap");
+    _THPVariable_Wrap = reinterpret_cast<decltype(_THPVariable_Wrap)>(
+        PyLong_AsVoidPtr(wrap_addr));
+    AOTI_TORCH_CHECK(_THPVariable_Wrap, "THPVariable_Wrap not resolved");
+
+    RAIIPyObject unpack_addr = PyObject_GetAttrString(
+        guards_mod, "_torchinductor_thputils_unpack_int");
+    AOTI_TORCH_CHECK(
+        unpack_addr, "Failed to get _torchinductor_thputils_unpack_int");
+    _THPUtils_unpackInt = reinterpret_cast<decltype(_THPUtils_unpackInt)>(
+        PyLong_AsVoidPtr(unpack_addr));
+    AOTI_TORCH_CHECK(_THPUtils_unpackInt, "THPUtils_unpackInt not resolved");
+  }
+}
+
+static inline std::string getStringAttr(PyObject* obj, const char* attr) {
+  RAIIPyObject val = PyObject_GetAttrString(obj, attr);
+  AOTI_TORCH_CHECK(val, "Failed to get attribute");
+  return PyUnicode_AsUTF8(val);
+}
+
+static inline int getIntAttr(PyObject* obj, const char* attr) {
+  RAIIPyObject val = PyObject_GetAttrString(obj, attr);
+  AOTI_TORCH_CHECK(val, "Failed to get attribute");
+  return _THPUtils_unpackInt(val);
+}
+
+static inline int getOptionalIntAttr(
+    PyObject* obj,
+    const char* attr,
+    int sentinel = -1) {
+  RAIIPyObject val = PyObject_GetAttrString(obj, attr);
+  AOTI_TORCH_CHECK(val, "Failed to get attribute");
+  return (val.get() != Py_None) ? _THPUtils_unpackInt(val) : sentinel;
+}
+
+static inline LazyKernelCompileResult extractCompileResult(PyObject* result) {
+  LazyKernelCompileResult compile_result;
+  compile_result.cubin_path = getStringAttr(result, "cubin_path");
+  compile_result.mangled_name = getStringAttr(result, "mangled_name");
+  compile_result.num_warps = getIntAttr(result, "num_warps");
+  compile_result.shared_mem = getIntAttr(result, "shared_mem");
+  compile_result.xblock = getIntAttr(result, "xblock");
+  compile_result.yblock = getIntAttr(result, "yblock");
+  compile_result.zblock = getIntAttr(result, "zblock");
+  compile_result.r0block = getIntAttr(result, "r0block");
+  compile_result.rsplit = getIntAttr(result, "rsplit");
+  compile_result.rsplit_size = getIntAttr(result, "rsplit_size");
+  compile_result.config_index = getOptionalIntAttr(result, "config_index");
+  compile_result.global_scratch = getOptionalIntAttr(result, "global_scratch");
+  compile_result.profile_scratch =
+      getOptionalIntAttr(result, "profile_scratch");
+  return compile_result;
+}
+
+template <typename T>
+static inline PyObject* convertArgToPython(const T& arg) {
+  using DecayedT = std::decay_t<T>;
+  if constexpr (std::is_same_v<DecayedT, AtenTensorHandle>) {
+    at::Tensor* tensor_ptr =
+        torch::aot_inductor::tensor_handle_to_tensor_pointer(arg);
+    return _THPVariable_Wrap(*tensor_ptr);
+  } else if constexpr (std::is_same_v<
+                           DecayedT,
+                           torch::aot_inductor::RAIIAtenTensorHandle>) {
+    at::Tensor* tensor_ptr =
+        torch::aot_inductor::tensor_handle_to_tensor_pointer(arg.get());
+    return _THPVariable_Wrap(*tensor_ptr);
+  } else if constexpr (std::is_same_v<DecayedT, bool>) {
+    PyObject* py_arg = arg ? Py_True : Py_False;
+    Py_INCREF(py_arg);
+    return py_arg;
+  } else if constexpr (std::is_integral_v<DecayedT>) {
+    return PyLong_FromLongLong(static_cast<long long>(arg));
+  } else if constexpr (std::is_floating_point_v<DecayedT>) {
+    return PyFloat_FromDouble(static_cast<double>(arg));
+  } else {
+    AOTI_TORCH_CHECK(false, "Invalid input type to convertArgToPython");
+  }
+}
+
+template <typename... Args>
+static inline LazyKernelCompileResult runTritonKernelWithAutotune(
+    PyObject* pending_kernels,
+    const std::string& kernel_name,
+    void* stream,
+    const Args&... kernel_args) {
+  py::gil_scoped_acquire_simple acquire;
+
+  constexpr size_t num_args = sizeof...(Args);
+  RAIIPyObject py_args_list = PyList_New(num_args);
+  AOTI_TORCH_CHECK(py_args_list, "Failed to create args list");
+
+  size_t idx = 0;
+  auto add_arg = [&py_args_list, &idx](PyObject* py_arg) {
+    AOTI_TORCH_CHECK(py_arg, "Failed to convert argument");
+    PyList_SetItem(py_args_list, idx++, py_arg);
+  };
+  // Use array pack-expansion instead of a fold expression to avoid
+  // hitting the compiler's expression-nesting limit when there are
+  // hundreds of kernel arguments (e.g. combo kernels).
+  int dummy[] = {0, (add_arg(convertArgToPython(kernel_args)), 0)...};
+  (void)dummy;
+
+  RAIIPyObject call_args = PyTuple_Pack(
+      4,
+      pending_kernels,
+      PyUnicode_FromString(kernel_name.c_str()),
+      PyLong_FromVoidPtr(stream),
+      py_args_list.get());
+  AOTI_TORCH_CHECK(call_args, "Failed to create call args");
+
+  RAIIPyObject result =
+      PyObject_CallObject(run_triton_kernel_with_autotune, call_args);
+  AOTI_TORCH_CHECK(result, "Failed to run kernel with autotuning");
+
+  return extractCompileResult(result);
+}
+
+static inline void startKernelCompile(
+    PyObject* pending_kernels,
+    const std::string& kernel_name,
+    const std::string& kernel_source) {
+  py::gil_scoped_acquire_simple acquire;
+
+  RAIIPyObject py_name = PyUnicode_FromString(kernel_name.c_str());
+  RAIIPyObject py_source = PyUnicode_FromString(kernel_source.c_str());
+  AOTI_TORCH_CHECK(py_name && py_source, "Failed to create Python args");
+
+  RAIIPyObject call_args =
+      PyTuple_Pack(3, pending_kernels, py_name.get(), py_source.get());
+  AOTI_TORCH_CHECK(call_args, "Failed to create call args");
+
+  RAIIPyObject result = PyObject_CallObject(start_kernel_compile, call_args);
+  AOTI_TORCH_CHECK(result, "Failed to start kernel compilation");
+}
diff --git a/torch/csrc/inductor/cpp_wrapper/xpu.h b/torch/csrc/inductor/cpp_wrapper/xpu.h
index e26dea0f3b6e2..160788dfa1bab 100644
--- a/torch/csrc/inductor/cpp_wrapper/xpu.h
+++ b/torch/csrc/inductor/cpp_wrapper/xpu.h
@@ -2,3 +2,4 @@
 
 #include <torch/csrc/inductor/cpp_wrapper/common.h>
 #include <torch/csrc/inductor/cpp_wrapper/device_internal/xpu.h>
+#include <torch/csrc/inductor/cpp_wrapper/lazy_triton_compile.h>
diff --git a/torch/csrc/inductor/inductor_ops.cpp b/torch/csrc/inductor/inductor_ops.cpp
index 9723e27e6ba8a..96a41a78c00bc 100644
--- a/torch/csrc/inductor/inductor_ops.cpp
+++ b/torch/csrc/inductor/inductor_ops.cpp
@@ -108,4 +108,11 @@ TORCH_LIBRARY_FRAGMENT(inductor, m) {
       {at::Tag::pt2_compliant_tag});
 }
 
+TORCH_LIBRARY_FRAGMENT(inductor_prims, m) {
+  m.def(
+      "inductor_reserve_rng_state(Generator? generator, SymInt increment) "
+      "-> (Tensor, Tensor, Tensor)",
+      {at::Tag::pt2_compliant_tag});
+}
+
 } // namespace torch::inductor
diff --git a/torch/csrc/inductor/inductor_ops_gpu.cpp b/torch/csrc/inductor/inductor_ops_gpu.cpp
new file mode 100644
index 0000000000000..2133dbcb4f796
--- /dev/null
+++ b/torch/csrc/inductor/inductor_ops_gpu.cpp
@@ -0,0 +1,89 @@
+#include <c10/core/SymInt.h>
+#include <torch/csrc/inductor/inductor_ops.h>
+#include <torch/library.h>
+#include <tuple>
+
+#ifndef AT_PER_OPERATOR_HEADERS
+#include <ATen/Functions.h>
+#else
+#include <ATen/ops/from_blob.h>
+#include <ATen/ops/scalar_tensor.h>
+#include <ATen/ops/zeros.h>
+#endif
+
+#if defined(USE_CUDA) || defined(USE_ROCM)
+#include <ATen/cuda/CUDAGeneratorImpl.h>
+#endif
+
+namespace torch::inductor {
+using namespace at;
+
+#if defined(USE_CUDA) || defined(USE_ROCM)
+
+// Reserves RNG state for Inductor with CUDA Graph support.
+//
+// This function allows Inductor to reserve a specific amount of RNG offset
+// (increment) for a kernel. It is designed to be safe for CUDA Graph capture
+// by explicitly handling the internal generator state via public APIs.
+//
+// Behavior:
+// - Graph Mode: Advances the generator state and returns pointers (wrapped as
+// tensors) to the extragraph state. These tensors effectively point to the
+// GPU memory that will be updated by `replay_prologue`.
+// - Eager Mode: Advances the generator state and returns concrete values
+// wrapped in 1D tensors to maintain shape consistency.
+//
+// -param gen The CUDA generator to use.
+// -param increment The number of RNG values to reserve.
+// -return A tuple of (Seed Tensor, Offset Tensor, Intragraph Offset CPU
+// Tensor).
+static std::tuple<Tensor, Tensor, Tensor> inductor_reserve_rng_state_impl(
+    const Generator& generator,
+    c10::SymInt increment) {
+  auto* gen_impl = at::check_generator<at::CUDAGeneratorImpl>(generator);
+
+  const auto dev_opts =
+      at::TensorOptions().dtype(at::kLong).device(generator.device());
+  const auto cpu_opts = at::TensorOptions().dtype(at::kLong).device(at::kCPU);
+
+  int64_t inc = increment.expect_int();
+  const at::PhiloxCudaState st =
+      gen_impl->philox_cuda_state(static_cast<uint64_t>(inc));
+
+  if (st.captured_) {
+    auto seed_t = at::from_blob(
+        static_cast<void*>(st.seed_.ptr), {1}, [](void*) {}, dev_opts);
+    auto off_t = at::from_blob(
+        static_cast<void*>(st.offset_.ptr), {1}, [](void*) {}, dev_opts);
+    auto intra_t =
+        at::scalar_tensor(static_cast<int64_t>(st.offset_intragraph_), cpu_opts)
+            .unsqueeze(0);
+    return {seed_t, off_t, intra_t};
+  }
+
+  auto seed_t = at::scalar_tensor(static_cast<int64_t>(st.seed_.val), dev_opts)
+                    .unsqueeze(0);
+  auto off_t = at::scalar_tensor(static_cast<int64_t>(st.offset_.val), dev_opts)
+                   .unsqueeze(0);
+  auto intra_t = at::zeros({1}, cpu_opts);
+  return {seed_t, off_t, intra_t};
+}
+
+TORCH_LIBRARY_IMPL(inductor_prims, BackendSelect, m) {
+  m.impl(
+      "inductor_reserve_rng_state", TORCH_FN(inductor_reserve_rng_state_impl));
+}
+
+TORCH_LIBRARY_IMPL(inductor_prims, CUDA, m) {
+  m.impl(
+      "inductor_reserve_rng_state", TORCH_FN(inductor_reserve_rng_state_impl));
+}
+
+TORCH_LIBRARY_IMPL(inductor_prims, HIP, m) {
+  m.impl(
+      "inductor_reserve_rng_state", TORCH_FN(inductor_reserve_rng_state_impl));
+}
+
+#endif
+
+} // namespace torch::inductor
diff --git a/torch/csrc/inductor/static_launcher/cuda.cpp b/torch/csrc/inductor/static_launcher/cuda.cpp
index a2378b7c1a248..44cfe7d2853ef 100644
--- a/torch/csrc/inductor/static_launcher/cuda.cpp
+++ b/torch/csrc/inductor/static_launcher/cuda.cpp
@@ -9,9 +9,13 @@
 #include <torch/csrc/utils/python_numbers.h>
 #include <filesystem>
 #include <optional>
+#include <utility>
 
 #if defined(USE_ROCM)
 #include <hip/hip_runtime_api.h>
+#include <fstream>
+#include <iterator>
+#include <vector>
 #endif
 
 /**
@@ -101,7 +105,20 @@ CUdeviceptr getPointer(PyObject* obj) {
 
 #define SHARED_MEM_STATIC_MAX 49152 // 48 KB
 
-CUfunction loadKernel(
+#if defined(USE_ROCM)
+std::vector<char> readKernelImage(const std::string& filePath) {
+  std::ifstream file(filePath, std::ios::binary);
+  TORCH_CHECK(file, "Failed to open kernel image: ", filePath);
+
+  auto begin = std::istreambuf_iterator<char>(file);
+  auto end = std::istreambuf_iterator<char>();
+  std::vector<char> image(begin, end);
+  TORCH_CHECK(!image.empty(), "Kernel image is empty: ", filePath);
+  return image;
+}
+#endif
+
+std::pair<CUmodule, CUfunction> loadKernel(
     std::string filePath,
     const std::string& funcName,
     uint32_t sharedMemBytes,
@@ -116,7 +133,10 @@ CUfunction loadKernel(
   CUfunction func = nullptr;
 
 #if defined(USE_ROCM)
-  AT_CUDA_DRIVER_CHECK(hipModuleLoad(&mod, filePath.c_str()));
+  // Unlike cuModuleLoad, hipModuleLoad keeps a file descriptor for the loaded
+  // HSACO. Load from memory to avoid retaining one FD per static launcher.
+  auto image = readKernelImage(filePath);
+  AT_CUDA_DRIVER_CHECK(hipModuleLoadData(&mod, image.data()));
   AT_CUDA_DRIVER_CHECK(hipModuleGetFunction(&func, mod, funcName.c_str()));
   int shared_optin = 0;
   AT_CUDA_DRIVER_CHECK(hipDeviceGetAttribute(
@@ -199,7 +219,7 @@ CUfunction loadKernel(
         shared_optin - shared_static));
 #endif
   }
-  return func;
+  return {mod, func};
 }
 
 inline void launchKernel(
@@ -332,10 +352,12 @@ void parseKernelArgs(
 }
 
 /* Load the CUDA kernel into memory (called during torch.compile), and
-  return a pointer to it (along with nregs and nspills).
+  return pointers to its owning module and function (along with nregs
+  and nspills).  The module must remain loaded while the function may
+  be launched.
   Called in python as:
-  (function, n_regs, n_spills) = load_kernel(cubin_path, func_name,
-  sharedMemBytes)
+  (module, function, n_regs, n_spills) = load_kernel(
+      cubin_path, func_name, sharedMemBytes)
 */
 PyObject* load_kernel(PyObject* self, PyObject* args) {
   HANDLE_TH_ERRORS
@@ -368,8 +390,7 @@ PyObject* load_kernel(PyObject* self, PyObject* args) {
   }
 #endif
 
-  CUfunction func = nullptr;
-  func = loadKernel(filePath, funcName, sharedMemBytes, device);
+  auto [mod, func] = loadKernel(filePath, funcName, sharedMemBytes, device);
 
 #if defined(USE_ROCM)
   AT_CUDA_DRIVER_CHECK(
@@ -385,9 +406,13 @@ PyObject* load_kernel(PyObject* self, PyObject* args) {
 
 #endif
   n_spills /= 4;
-  // Return a tuple of CUFunction, n_regs, n_spills
+  // Return a tuple of CUmodule, CUfunction, n_regs, n_spills.
   return Py_BuildValue(
-      "(Kii)", reinterpret_cast<uint64_t>(func), n_regs, n_spills);
+      "(KKii)",
+      reinterpret_cast<uint64_t>(mod),
+      reinterpret_cast<uint64_t>(func),
+      n_regs,
+      n_spills);
   END_HANDLE_TH_ERRORS
 }
 
@@ -556,7 +581,25 @@ PyObject* launch_kernel(PyObject* self, PyObject* args) {
   END_HANDLE_TH_ERRORS
 }
 
-std::array<PyMethodDef, 2> StaticCudaLauncherMethods = {
+PyObject* unload_kernel(PyObject* self, PyObject* args) {
+  HANDLE_TH_ERRORS
+  uint64_t mod_ptr = 0;
+  if (!PyArg_ParseTuple(args, "K", &mod_ptr)) {
+    return nullptr;
+  }
+  CUmodule mod = reinterpret_cast<CUmodule>(mod_ptr);
+  if (mod) {
+#if defined(USE_ROCM)
+    AT_CUDA_DRIVER_CHECK(hipModuleUnload(mod));
+#else
+    AT_CUDA_DRIVER_CHECK(nvrtc().cuModuleUnload(mod));
+#endif
+  }
+  Py_RETURN_NONE;
+  END_HANDLE_TH_ERRORS
+}
+
+std::array<PyMethodDef, 3> StaticCudaLauncherMethods = {
     PyMethodDef{
         "_launch_kernel",
         launch_kernel,
@@ -566,7 +609,12 @@ std::array<PyMethodDef, 2> StaticCudaLauncherMethods = {
         "_load_kernel",
         load_kernel,
         METH_VARARGS,
-        "Load CUDA kernel from cubin file"}};
+        "Load CUDA kernel from cubin file"},
+    PyMethodDef{
+        "_unload_kernel",
+        unload_kernel,
+        METH_VARARGS,
+        "Unload CUDA/HIP module loaded by _load_kernel"}};
 
 // Define a minimal type for StaticCudaLauncher.
 // We don't implement __new__ or __init__ because we're using it only as a
diff --git a/torch/csrc/jit/backends/backend_exception.h b/torch/csrc/jit/backends/backend_exception.h
index 807ef38e28305..b8437d305e1b2 100644
--- a/torch/csrc/jit/backends/backend_exception.h
+++ b/torch/csrc/jit/backends/backend_exception.h
@@ -1,3 +1,4 @@
+// @allow-raw-throw
 #pragma once
 #include <c10/util/Exception.h>
 
diff --git a/torch/csrc/jit/backends/coreml/objc/PTMCoreMLBackend.mm b/torch/csrc/jit/backends/coreml/objc/PTMCoreMLBackend.mm
index 7823d066bafc2..96828ab39e5d4 100644
--- a/torch/csrc/jit/backends/coreml/objc/PTMCoreMLBackend.mm
+++ b/torch/csrc/jit/backends/coreml/objc/PTMCoreMLBackend.mm
@@ -213,8 +213,9 @@ bool is_available() override {
 #elif TARGET_OS_MAC
     NSOperatingSystemVersion supportedVer = {10, 13, 0};
     return [[NSProcessInfo processInfo] isOperatingSystemAtLeastVersion:supportedVer];
-#endif
+#else
     return false;
+#endif
   }
 };
 
diff --git a/torch/csrc/jit/codegen/fuser/cpu/fused_kernel.cpp b/torch/csrc/jit/codegen/fuser/cpu/fused_kernel.cpp
index 4cd8cedca1d2e..d8fd2ce6f890b 100644
--- a/torch/csrc/jit/codegen/fuser/cpu/fused_kernel.cpp
+++ b/torch/csrc/jit/codegen/fuser/cpu/fused_kernel.cpp
@@ -256,7 +256,7 @@ static const std::string compile_string =
 #ifndef __PPC64__
 //  "-march=native "
 #endif
-    "-std=c++17 -fPIC ${fopenmp} -shared \"${cpp_file}\" -o \"${so_file}\" -lm";
+    "-std=c++20 -fPIC ${fopenmp} -shared \"${cpp_file}\" -o \"${so_file}\" -lm";
 #endif
 static void runCompiler(
     const std::string& cpp_file,
diff --git a/torch/csrc/jit/codegen/fuser/cuda/fused_kernel.cpp b/torch/csrc/jit/codegen/fuser/cuda/fused_kernel.cpp
index cba3aff9efa64..66855cae43420 100644
--- a/torch/csrc/jit/codegen/fuser/cuda/fused_kernel.cpp
+++ b/torch/csrc/jit/codegen/fuser/cuda/fused_kernel.cpp
@@ -122,7 +122,7 @@ FusedKernelCUDA::FusedKernelCUDA(
       &program, code_.c_str(), nullptr, 0, nullptr, nullptr));
 
 #if defined(USE_ROCM)
-  std::vector<const char*> args = {"--std=c++17"};
+  std::vector<const char*> args = {"--std=c++20"};
   args.push_back("-hip-pch");
 #else
   const std::string compute = std::string("--gpu-architecture=") +
@@ -140,7 +140,7 @@ FusedKernelCUDA::FusedKernelCUDA(
 #endif
       std::to_string(major) + std::to_string(minor);
   const std::vector<const char*> args = {
-      "--std=c++17", compute.c_str(), "-default-device"};
+      "--std=c++20", compute.c_str(), "-default-device"};
 #endif
   const auto result =
       nvrtc().nvrtcCompileProgram(program, args.size(), args.data());
diff --git a/torch/csrc/jit/codegen/fuser/cuda/resource_strings.h b/torch/csrc/jit/codegen/fuser/cuda/resource_strings.h
index 16ccc5002f9ab..c38338ef86241 100644
--- a/torch/csrc/jit/codegen/fuser/cuda/resource_strings.h
+++ b/torch/csrc/jit/codegen/fuser/cuda/resource_strings.h
@@ -268,8 +268,26 @@ typedef __half half;
 #define BF16_UINT32_DEF ""
 #endif
 
+// NOTE [ ROCm JIT bfloat16 symbol collision ]
+// This file is excluded from hipify (see tools/amd_build/build_amd.py), so it
+// ships its own self-contained __nv_bfloat16 type plus the
+// __float2bfloat16/__bfloat162float conversions for the runtime-compiled fuser
+// kernel. Starting with ROCm 7.x, hiprtc auto-injects a runtime header
+// (hiprtc_runtime.h) that also defines __float2bfloat16 returning
+// __hip_bfloat16. Since C++ cannot overload on return type alone, the two
+// definitions collide. Rather than gate on ROCM_VERSION (a build-time value
+// that does not reliably track whether the *runtime* hiprtc actually provides
+// these symbols), we keep PyTorch's own round-to-nearest-even implementation
+// and rename our private definitions via the preprocessor so they cannot
+// clash. hiprtc's definitions remain (parsed earlier in the translation unit)
+// but go unused. These #defines must stay in effect for the rest of the kernel
+// source -- do NOT #undef them, since the emitted kernel body calls these
+// conversions after this header is spliced in.
 constexpr auto bfloat16_support_literal =
     R"(
+#define __internal_float2bfloat16 __pt_jit_internal_float2bfloat16
+#define __float2bfloat16 __pt_jit_float2bfloat16
+#define __bfloat162float __pt_jit_bfloat162float
 #ifndef __align__
 #define __align__(x) __attribute__((aligned(x)))
 #endif
diff --git a/torch/csrc/jit/frontend/schema_type_parser.cpp b/torch/csrc/jit/frontend/schema_type_parser.cpp
index dd8cbe98b84a3..960fc9c260666 100644
--- a/torch/csrc/jit/frontend/schema_type_parser.cpp
+++ b/torch/csrc/jit/frontend/schema_type_parser.cpp
@@ -7,6 +7,7 @@
 #include <torch/csrc/jit/frontend/lexer.h>
 #include <torch/csrc/jit/frontend/parse_string_literal.h>
 #include <torch/custom_class.h>
+#include <mutex>
 #include <string>
 #include <unordered_set>
 
@@ -43,12 +44,18 @@ using c10::VarType;
 
 namespace torch::jit {
 
+static std::mutex& getOpaqueTypesMutex() {
+  static std::mutex opaque_types_mutex;
+  return opaque_types_mutex;
+}
+
 static std::unordered_set<std::string>& getOpaqueTypes() {
   static std::unordered_set<std::string> global_opaque_types;
   return global_opaque_types;
 }
 
 void registerOpaqueType(const std::string& type_name) {
+  std::lock_guard<std::mutex> lock(getOpaqueTypesMutex());
   auto& global_opaque_types = getOpaqueTypes();
   auto [_, inserted] = global_opaque_types.insert(type_name);
   if (!inserted) {
@@ -58,11 +65,13 @@ void registerOpaqueType(const std::string& type_name) {
 }
 
 void unregisterOpaqueType(const std::string& type_name) {
+  std::lock_guard<std::mutex> lock(getOpaqueTypesMutex());
   auto& global_opaque_types = getOpaqueTypes();
   global_opaque_types.erase(type_name);
 }
 
 bool isRegisteredOpaqueType(const std::string& type_name) {
+  std::lock_guard<std::mutex> lock(getOpaqueTypesMutex());
   auto& global_opaque_types = getOpaqueTypes();
   return global_opaque_types.find(type_name) != global_opaque_types.end();
 }
@@ -99,6 +108,9 @@ TypePtr SchemaTypeParser::parseBaseType() {
       {"Any", c10::TypeFactory::get<c10::AnyType>()},
       {"AnyClassType", c10::TypeFactory::get<c10::AnyClassType>()},
       {"AnyEnumType", c10::TypeFactory::get<c10::AnyEnumType>()},
+      // PyObjectType::get() used directly because PyObjectType is excluded
+      // from FORALL_DYNAMIC_TYPES (not supported on xplat/mobile)
+      {"PyObject", c10::PyObjectType::get()},
   };
   auto tok = L.cur();
   if (!L.nextIf(TK_NONE) && !L.nextIf(TK_NONE_TYPE)) {
diff --git a/torch/csrc/jit/ir/ir.cpp b/torch/csrc/jit/ir/ir.cpp
index 9b85b0909826f..902d8ffad9278 100644
--- a/torch/csrc/jit/ir/ir.cpp
+++ b/torch/csrc/jit/ir/ir.cpp
@@ -929,12 +929,7 @@ void Value::replaceAllUsesAfterNodeWith(const Node* node, Value* newValue) {
     }
   });
 
-  uses_.erase(
-      std::remove_if(
-          uses_.begin(),
-          uses_.end(),
-          [&node](const Use& u) { return u.user->isAfter(node); }),
-      uses_.end());
+  std::erase_if(uses_, [&node](const Use& u) { return u.user->isAfter(node); });
 }
 
 void Value::replaceAllUsesDominatedByNodeWith(
@@ -947,12 +942,8 @@ void Value::replaceAllUsesDominatedByNodeWith(
     }
   });
 
-  uses_.erase(
-      std::remove_if(
-          uses_.begin(),
-          uses_.end(),
-          [&node](const Use& u) { return u.user->isDominatedBy(node); }),
-      uses_.end());
+  std::erase_if(
+      uses_, [&node](const Use& u) { return u.user->isDominatedBy(node); });
 }
 
 static size_t findArgument(
@@ -1218,7 +1209,6 @@ bool Node::hasSideEffects() const {
       return true;
   }
   TORCH_INTERNAL_ASSERT(false, "Unhandled AliasAnalysisKind case");
-  return false; // silence compiler warning
 }
 
 // Assign this node a topological position, to facilitate fast isBefore() and
diff --git a/torch/csrc/jit/ir/ir.h b/torch/csrc/jit/ir/ir.h
index 57c6d9ef89d41..7c6573ba8c9e1 100644
--- a/torch/csrc/jit/ir/ir.h
+++ b/torch/csrc/jit/ir/ir.h
@@ -68,18 +68,10 @@ using ::c10::TypeKind;
 
 using ::c10::fmap;
 
-namespace prim {
-using namespace ::c10::prim;
-}
-namespace attr {
-using namespace ::c10::attr;
-}
-namespace aten {
-using namespace ::c10::aten;
-}
-namespace cuda {
-using namespace ::c10::cuda;
-} // namespace cuda
+namespace prim = ::c10::prim;
+namespace attr = ::c10::attr;
+namespace aten = ::c10::aten;
+namespace cuda = ::c10::cuda;
 
 struct Function;
 struct GraphFunction;
diff --git a/torch/csrc/jit/mobile/interpreter.cpp b/torch/csrc/jit/mobile/interpreter.cpp
index a0e0959d6033d..9a0d976df8042 100644
--- a/torch/csrc/jit/mobile/interpreter.cpp
+++ b/torch/csrc/jit/mobile/interpreter.cpp
@@ -387,7 +387,6 @@ bool InterpreterState::run(Stack& stack) {
     //    }
     //  }
   }
-  return false;
 }
 
 IValue& InterpreterState::reg(size_t reg) {
diff --git a/torch/csrc/jit/mobile/type_parser.cpp b/torch/csrc/jit/mobile/type_parser.cpp
index f9287a5eb7040..a4d476070c12d 100644
--- a/torch/csrc/jit/mobile/type_parser.cpp
+++ b/torch/csrc/jit/mobile/type_parser.cpp
@@ -158,7 +158,6 @@ TypePtr TypeParser::parse() {
         " is not supported in the parser, ",
         "or the token is in wrong format.");
   }
-  return nullptr;
 }
 
 // NamedTuple custom type will be following structure:
@@ -243,7 +242,6 @@ TypePtr TypeParser::parseCustomType() {
       TORCH_CHECK(
           false, "Can't find definition for the type: ", qualified_name);
     }
-    return nullptr;
   }
 }
 
diff --git a/torch/csrc/jit/passes/autocast.cpp b/torch/csrc/jit/passes/autocast.cpp
index 79ead2c5ee6c3..7b638b85f9da1 100644
--- a/torch/csrc/jit/passes/autocast.cpp
+++ b/torch/csrc/jit/passes/autocast.cpp
@@ -133,8 +133,6 @@ std::optional<AutocastScope> parseAutocast(
     //
     TORCH_CHECK(false, "Unsupported autocast syntax");
   }
-
-  return std::nullopt;
 }
 
 void castTensorInputs(
diff --git a/torch/csrc/jit/passes/dead_code_elimination.cpp b/torch/csrc/jit/passes/dead_code_elimination.cpp
index 4548c6b2ccce5..b2dffc35983b3 100644
--- a/torch/csrc/jit/passes/dead_code_elimination.cpp
+++ b/torch/csrc/jit/passes/dead_code_elimination.cpp
@@ -9,10 +9,6 @@
 
 namespace torch::jit {
 
-namespace prim {
-using namespace ::c10::prim;
-}
-
 class DeadCodeEliminator {
  public:
   explicit DeadCodeEliminator(
diff --git a/torch/csrc/jit/passes/inliner.cpp b/torch/csrc/jit/passes/inliner.cpp
index 9c06f748e43a9..14943a47d679c 100644
--- a/torch/csrc/jit/passes/inliner.cpp
+++ b/torch/csrc/jit/passes/inliner.cpp
@@ -6,10 +6,6 @@
 
 namespace torch::jit {
 
-namespace prim {
-using namespace ::c10::prim;
-}
-
 GraphFunction* tryToGraphFunction(Node* n) {
   if (n->kind() == prim::CallFunction) {
     AT_ASSERT(n->input(0)->node()->kind() == prim::Constant);
diff --git a/torch/csrc/jit/passes/onnx/constant_map.cpp b/torch/csrc/jit/passes/onnx/constant_map.cpp
index 60699a1e75ef4..c332ea4204eb9 100644
--- a/torch/csrc/jit/passes/onnx/constant_map.cpp
+++ b/torch/csrc/jit/passes/onnx/constant_map.cpp
@@ -120,28 +120,25 @@ std::optional<std::vector<int64_t>> ConstantValueMap::
     GetShapeInto1DInt64VectorWithOneUnknown(const std::string& value_name) {
   if (ConstantValueMap::HasShape(value_name)) {
     auto shape_size = ConstantValueMap::GetShape(value_name).value();
-    std::vector<int64_t> shape_value;
     if (shape_size.isComplete()) {
-      shape_value =
-          ConstantValueMap::GetCompleteShapeInto1DInt64Vector(shape_size);
-      return shape_value;
-    } else {
-      size_t count_unknown = 0;
-      auto shape_size_sizes = shape_size.sizes();
-      if (shape_size_sizes.has_value()) {
-        auto shape_symbol_list = shape_size_sizes.value();
-        for (const auto& v : shape_symbol_list) {
-          if (v.is_static()) {
-            shape_value.emplace_back(v.static_size());
-          } else {
-            shape_value.emplace_back(-1);
-            count_unknown += 1;
-          }
-        }
-        if (count_unknown == 1) {
-          return shape_value;
+      return ConstantValueMap::GetCompleteShapeInto1DInt64Vector(shape_size);
+    }
+    size_t count_unknown = 0;
+    auto shape_size_sizes = shape_size.sizes();
+    if (shape_size_sizes.has_value()) {
+      std::vector<int64_t> shape_value;
+      auto shape_symbol_list = shape_size_sizes.value();
+      for (const auto& v : shape_symbol_list) {
+        if (v.is_static()) {
+          shape_value.emplace_back(v.static_size());
+        } else {
+          shape_value.emplace_back(-1);
+          count_unknown += 1;
         }
       }
+      if (count_unknown == 1) {
+        return shape_value;
+      }
     }
   }
   return std::nullopt;
diff --git a/torch/csrc/jit/passes/onnx/deduplicate_initializers.cpp b/torch/csrc/jit/passes/onnx/deduplicate_initializers.cpp
index e5179cdf5f668..aca21d2c9a059 100644
--- a/torch/csrc/jit/passes/onnx/deduplicate_initializers.cpp
+++ b/torch/csrc/jit/passes/onnx/deduplicate_initializers.cpp
@@ -2,53 +2,119 @@
 #include <torch/csrc/jit/passes/onnx/deduplicate_initializers.h>
 #include <torch/csrc/jit/passes/onnx/helper.h>
 
+#include <c10/util/hash.h>
 #include <c10/util/irange.h>
 
+#include <functional>
+#include <unordered_set>
+
 namespace torch::jit {
 
 namespace onnx {
 using namespace ::c10::onnx;
 }
 
+struct HashValue {
+  HashValue(ValueToParamPairMap& valsToParamsMap, bool compare_by_ptr)
+      : valsToParamsMap_(valsToParamsMap), compare_by_ptr_(compare_by_ptr) {}
+
+  size_t operator()(Value* v) const {
+    auto t = valsToParamsMap_.find(v)->second.second.toTensor();
+
+    if (compare_by_ptr_) {
+      // Hash by metadata + data pointer
+      return at::get_hash(
+          t.sizes(), t.strides(), t.has_storage() ? t.data_ptr() : 0);
+    }
+
+    // Hash by metadata + first element value (if exists). This is a fast
+    // approximation of hashing by the whole tensor value, which can be
+    // expensive for large tensors.
+    double first_elem_double = 0.0;
+    int64_t first_elem_int64 = 0;
+    uint64_t first_elem_uint64 = 0;
+
+    if (t.numel() > 0 && t.has_storage()) {
+      auto scalar = t.reshape(-1)[0].item();
+
+      if (scalar.isFloatingPoint()) {
+        first_elem_double = scalar.to<double>();
+      } else if (scalar.isIntegral(/*includeBool=*/true)) {
+        if (scalar.isUnsigned()) {
+          first_elem_uint64 = scalar.to<uint64_t>();
+        } else {
+          first_elem_int64 = scalar.to<int64_t>();
+        }
+      }
+    }
+
+    return at::get_hash(
+        first_elem_double,
+        first_elem_int64,
+        first_elem_uint64,
+        t.sizes(),
+        t.strides());
+  }
+
+ private:
+  ValueToParamPairMap& valsToParamsMap_;
+  bool compare_by_ptr_;
+};
+
+struct CompareValue {
+  CompareValue(std::function<bool(Value*, Value*)> is_same_tensor_as)
+      : is_same_tensor_as_(is_same_tensor_as) {}
+
+  bool operator()(Value* v1, Value* v2) const {
+    return is_same_tensor_as_(v1, v2);
+  }
+
+ private:
+  std::function<bool(Value*, Value*)> is_same_tensor_as_;
+};
+
+// forward declaration
+static bool DeduplicateInitializersByDataPtr(at::Tensor& t1, at::Tensor& t2);
+
 static void DeduplicateInitializers(
     std::shared_ptr<Graph>& g,
     ValueToParamPairMap& valsToParamsMap,
     bool (*comp)(at::Tensor&, at::Tensor&)) {
-  auto is_same_tensor_as = [&valsToParamsMap, comp](Value* v1) {
-    return [&valsToParamsMap, v1, comp](Value* v2) {
-      if ((valsToParamsMap.find(v1) == valsToParamsMap.end()) ||
-          (valsToParamsMap.find(v2) == valsToParamsMap.end())) {
-        return false;
-      }
-      auto iv1 = valsToParamsMap.find(v1)->second.second;
-      auto iv2 = valsToParamsMap.find(v2)->second.second;
-      if (!iv1.isTensor() || !iv2.isTensor()) {
-        return false;
-      }
-      auto t1 = iv1.toTensor();
-      auto t2 = iv2.toTensor();
-      return comp(t1, t2);
-    };
+  auto is_same_tensor_as = [&valsToParamsMap, comp](Value* v1, Value* v2) {
+    auto t1 = valsToParamsMap.find(v1)->second.second.toTensor();
+    auto t2 = valsToParamsMap.find(v2)->second.second.toTensor();
+    return comp(t1, t2);
   };
-  std::vector<Value*> uniqueVals;
+
+  bool compare_by_ptr = comp == &DeduplicateInitializersByDataPtr;
+  std::unordered_set<Value*, HashValue, CompareValue> uniqueVals(
+      0,
+      HashValue(valsToParamsMap, compare_by_ptr),
+      CompareValue(is_same_tensor_as));
   std::vector<size_t> inputsIndicesToRemove;
   auto b = g->block();
 
   for (auto i : c10::irange(b->inputs().size())) {
     auto v = g->inputs().at(i);
-    if (valsToParamsMap.find(v) == valsToParamsMap.end()) {
-      // Skip model inputs
+    auto vals_to_param_it = valsToParamsMap.find(v);
+
+    // Skip parameters without initializers
+    if (vals_to_param_it == valsToParamsMap.end()) {
+      continue;
+    }
+
+    // Skip non-tensors
+    if (!vals_to_param_it->second.second.isTensor()) {
       continue;
     }
-    auto it = std::find_if(
-        uniqueVals.begin(), uniqueVals.end(), is_same_tensor_as(v));
-    if (it == uniqueVals.end()) {
-      uniqueVals.emplace_back(v);
-    } else {
+
+    auto it = uniqueVals.insert(v);
+    if (!it.second) {
+      // Same value already exists
       inputsIndicesToRemove.emplace_back(i);
       auto id_node = g->create(onnx::Identity);
       id_node->insertAfter(g->block()->param_node());
-      id_node->addInput(*it);
+      id_node->addInput(*it.first);
       id_node->output()->copyMetadata(v);
       id_node->copyMetadata(g->block()->param_node());
       v->replaceAllUsesWith(id_node->output());
diff --git a/torch/csrc/jit/passes/onnx/function_extraction.cpp b/torch/csrc/jit/passes/onnx/function_extraction.cpp
index c82fd3d63c9af..eee0251f4153b 100644
--- a/torch/csrc/jit/passes/onnx/function_extraction.cpp
+++ b/torch/csrc/jit/passes/onnx/function_extraction.cpp
@@ -710,14 +710,8 @@ void FunctionExtractor::ConvertScopeToFunction(
       ctx_nlist.insert(last_n_it, func_n);
 
       // remove replaced nodes from list
-      ctx_nlist.erase(
-          std::remove_if(
-              ctx_nlist.begin(),
-              ctx_nlist.end(),
-              [&old_nodes](Node* n) {
-                return old_nodes.find(n) != old_nodes.end();
-              }),
-          ctx_nlist.end());
+      std::erase_if(
+          ctx_nlist, [&old_nodes](Node* n) { return old_nodes.contains(n); });
 
       GRAPH_DEBUG("Parent total nodes after remove: ", ctx_nlist.size());
 
diff --git a/torch/csrc/jit/passes/onnx/helper.cpp b/torch/csrc/jit/passes/onnx/helper.cpp
index 3897a8d5cae5e..c2ee1d7dc285f 100644
--- a/torch/csrc/jit/passes/onnx/helper.cpp
+++ b/torch/csrc/jit/passes/onnx/helper.cpp
@@ -39,14 +39,8 @@ void eraseUnusedBlockInputs(Block* b) {
 }
 
 void eraseUnusedValuesFromMap(ValueToParamPairMap& valsToParamsMap) {
-  auto it = valsToParamsMap.begin();
-  while (it != valsToParamsMap.end()) {
-    if (!it->first->hasUses()) {
-      it = valsToParamsMap.erase(it);
-    } else {
-      ++it;
-    }
-  }
+  std::erase_if(
+      valsToParamsMap, [](const auto& pr) { return !pr.first->hasUses(); });
 }
 
 void buildParamsMapFromValueToParamsMap(
@@ -101,7 +95,6 @@ std::optional<at::ScalarType> ONNXTypeToATenType(int32_t onnx_type) {
           onnx_type,
           " is an unexpected tensor scalar type");
   }
-  return std::optional<at::ScalarType>{};
 }
 
 Node* addNodeToBlock(Block* block, Symbol kind, ArrayRef<Value*> inputs) {
diff --git a/torch/csrc/jit/passes/shape_analysis.cpp b/torch/csrc/jit/passes/shape_analysis.cpp
index 7493667a2f027..a329fd2e40d0b 100644
--- a/torch/csrc/jit/passes/shape_analysis.cpp
+++ b/torch/csrc/jit/passes/shape_analysis.cpp
@@ -106,10 +106,6 @@ void PropertyPropBase::setUnshapedType(Node* node) {
   }
 }
 
-namespace prim {
-using namespace ::c10::prim;
-}
-
 #define SHAPE_ASSERT(cond) \
   if (!(cond))             \
   throw propagation_error()
diff --git a/torch/csrc/jit/passes/subgraph_rewrite.h b/torch/csrc/jit/passes/subgraph_rewrite.h
index 4e4f4a04b68e3..451d61fcde6ec 100644
--- a/torch/csrc/jit/passes/subgraph_rewrite.h
+++ b/torch/csrc/jit/passes/subgraph_rewrite.h
@@ -18,8 +18,18 @@
 
 namespace torch::jit {
 
+/** Rewrite pattern descriptor.
+ *
+ * This structure is used in the implementation of `SubgraphRewriter` and
+ * is not supposed to be used externally.
+ */
+struct RewritePatternDescr {
+  std::string pattern;
+  std::string replacement;
+  std::unordered_map<std::string, std::string> value_name_map;
+};
+
 // Forward declarations.
-struct RewritePatternDescr;
 struct Match;
 
 using MatchFilter = std::function<
@@ -101,15 +111,4 @@ class TORCH_API SubgraphRewriter {
   bool overlapsWithPreviousMatches(const Match* match);
 };
 
-/** Rewrite pattern descriptor.
- *
- * This structure is used in the implementation of `SubgraphRewriter` and
- * is not supposed to be used externally.
- */
-struct RewritePatternDescr {
-  std::string pattern;
-  std::string replacement;
-  std::unordered_map<std::string, std::string> value_name_map;
-};
-
 } // namespace torch::jit
diff --git a/torch/csrc/jit/passes/symbolic_shape_cache.cpp b/torch/csrc/jit/passes/symbolic_shape_cache.cpp
index 0cca03d6f74d0..11c79dedcd0e4 100644
--- a/torch/csrc/jit/passes/symbolic_shape_cache.cpp
+++ b/torch/csrc/jit/passes/symbolic_shape_cache.cpp
@@ -92,7 +92,7 @@ ShapeCacheKey get_cache_key(
     std::unordered_map<int64_t, int64_t>& ss_map,
     bool deep_copy = true) {
   CanonicalArgVec canonical_args = cannonicalizeVec(arg_vec, ss_map, deep_copy);
-  return std::make_tuple(schema->operator_name(), canonical_args);
+  return std::make_tuple(schema->operator_name(), std::move(canonical_args));
 }
 
 } // namespace
diff --git a/torch/csrc/jit/passes/symbolic_shape_runtime_fusion.cpp b/torch/csrc/jit/passes/symbolic_shape_runtime_fusion.cpp
index c8c6953f3447f..3d858bb39d9dc 100644
--- a/torch/csrc/jit/passes/symbolic_shape_runtime_fusion.cpp
+++ b/torch/csrc/jit/passes/symbolic_shape_runtime_fusion.cpp
@@ -215,12 +215,11 @@ static void moveConstantTensorsOutOfSubgraph(
     const std::shared_ptr<Graph>& tensorexpr_graph) {
   auto parent = tensorexpr_graph_node->owningGraph();
 
-  auto env = [&](Value* v) {
+  auto env = [&](Value* v) -> Value* {
     TORCH_INTERNAL_ASSERT(
         false,
         "this should never happen since constant nodes do not have any inputs",
         v->debugName());
-    return v;
   };
 
   WithInsertPoint wip(tensorexpr_graph_node);
diff --git a/torch/csrc/jit/passes/utils/subgraph_utils.cpp b/torch/csrc/jit/passes/utils/subgraph_utils.cpp
index 6f92e821e5b44..f1b6a1f875134 100644
--- a/torch/csrc/jit/passes/utils/subgraph_utils.cpp
+++ b/torch/csrc/jit/passes/utils/subgraph_utils.cpp
@@ -516,7 +516,6 @@ void unmergeNode(Node* n, Node* subgraphNode) {
         false,
         "all inputs should've been mapped. Couldn't map %",
         v->debugName());
-    return v;
   };
 
   for (auto i : c10::irange(subgraph->outputs().size())) {
diff --git a/torch/csrc/jit/python/init.cpp b/torch/csrc/jit/python/init.cpp
index 931d80ffd6a68..67c6349f64efd 100644
--- a/torch/csrc/jit/python/init.cpp
+++ b/torch/csrc/jit/python/init.cpp
@@ -1796,7 +1796,7 @@ void initJITBindings(PyObject* module) {
 
   m.def(
       "_jit_get_operation",
-      [](const std::string& op_name) {
+      [](const std::string& op_name) -> py::tuple {
         try {
           auto symbol = Symbol::fromQualString(op_name);
           const auto sortedOps = getAllSortedOperatorsFor(symbol);
@@ -1843,7 +1843,7 @@ void initJITBindings(PyObject* module) {
       "_maybe_call_torch_function_for_op_packet",
       [](py::handle op_overload_packet,
          const py::args& args,
-         const py::kwargs& kwargs) {
+         const py::kwargs& kwargs) -> py::tuple {
         py::list ns_method =
             op_overload_packet.attr("_qualified_op_name").attr("split")("::");
         auto res = _maybe_handle_torch_function(
@@ -2183,20 +2183,12 @@ void initJITBindings(PyObject* module) {
       .def(
           py::pickle(
               /* __getstate__ */
-              [](const PythonFutureWrapper& /* unused */) {
+              [](const PythonFutureWrapper& /* unused */) -> py::tuple {
                 TORCH_CHECK(false, "Can not pickle torch.futures.Future");
-                // Note that this return has no meaning since we always
-                // throw, it's only here to satisfy Pybind API's
-                // requirement.
-                return py::make_tuple();
               },
               /* __setstate__ */
-              [](const py::tuple& /* unused */) {
+              [](const py::tuple& /* unused */) -> std::nullptr_t {
                 TORCH_CHECK(false, "Can not unpickle torch.futures.Future");
-                // Note that this return has no meaning since we always
-                // throw, it's only here to satisfy PyBind's API
-                // requirement.
-                return nullptr;
               }),
           py::call_guard<py::gil_scoped_release>());
 
@@ -2220,20 +2212,12 @@ void initJITBindings(PyObject* module) {
       .def(
           py::pickle(
               /* __getstate__ */
-              [](const PythonAwaitWrapper& /* unused */) {
+              [](const PythonAwaitWrapper& /* unused */) -> py::tuple {
                 TORCH_CHECK(false, "Can not pickle torch.jit._Await");
-                // Note that this return has no meaning since we always
-                // throw, it's only here to satisfy Pybind API's
-                // requirement.
-                return py::make_tuple();
               },
               /* __setstate__ */
-              [](const py::tuple& /* unused */) {
+              [](const py::tuple& /* unused */) -> std::nullptr_t {
                 TORCH_CHECK(false, "Can not unpickle torch.jit._Await");
-                // Note that this return has no meaning since we always
-                // throw, it's only here to satisfy PyBind's API
-                // requirement.
-                return nullptr;
               }),
           py::call_guard<py::gil_scoped_release>());
 
diff --git a/torch/csrc/jit/python/pybind_utils.cpp b/torch/csrc/jit/python/pybind_utils.cpp
index cf422102f7dec..bbcc31515ce3e 100644
--- a/torch/csrc/jit/python/pybind_utils.cpp
+++ b/torch/csrc/jit/python/pybind_utils.cpp
@@ -543,9 +543,15 @@ IValue toIValue(py::handle obj, const TypePtr& type, std::optional<int32_t> N) {
     }
     case TypeKind::CapsuleType: {
 #ifdef USE_DISTRIBUTED
-      // Handle ProcessGroup custom class as a capsule
+      // Handle ProcessGroup custom class as a capsule.  FakeScriptObject
+      // (used during Dynamo tracing with CooR) passes py::isinstance via
+      // OpaqueBaseMeta but cannot be cast directly; unwrap real_obj first.
       if (py::isinstance<c10d::ProcessGroup>(obj)) {
-        auto cpp_obj = obj.cast<c10::intrusive_ptr<c10d::ProcessGroup>>();
+        py::handle target = obj;
+        if (py::hasattr(obj, "real_obj")) {
+          target = obj.attr("real_obj");
+        }
+        auto cpp_obj = target.cast<c10::intrusive_ptr<c10d::ProcessGroup>>();
         return IValue::make_capsule(cpp_obj);
       }
 #endif
@@ -1006,6 +1012,18 @@ std::optional<InferredType> detail::_tryToInferTypeImpl(py::handle input) {
   if (py::isinstance<c10d::ProcessGroup>(input)) {
     return InferredType(CapsuleType::get());
   }
+  // During Dynamo tracing with compile-on-one-rank (CooR), opaque reference
+  // types like ProcessGroup are wrapped in FakeScriptObject.  Python-level
+  // isinstance() sees through the wrapper (via OpaqueBaseMeta), but the C++
+  // py::isinstance above does too — yet the subsequent pybind11 cast would
+  // fail because FakeScriptObject is not a C++ bound object.  Detect this
+  // case by checking for the wrapped real_obj attribute.
+  if (py::hasattr(input, "real_obj")) {
+    py::object real = input.attr("real_obj");
+    if (py::isinstance<c10d::ProcessGroup>(real)) {
+      return InferredType(CapsuleType::get());
+    }
+  }
 #endif
 
   return std::nullopt;
diff --git a/torch/csrc/jit/python/pybind_utils.h b/torch/csrc/jit/python/pybind_utils.h
index 378bf4636fd2e..37e49c2952620 100644
--- a/torch/csrc/jit/python/pybind_utils.h
+++ b/torch/csrc/jit/python/pybind_utils.h
@@ -23,7 +23,6 @@
 #include <torch/csrc/jit/runtime/operator.h>
 #include <torch/csrc/utils/pybind.h>
 #include <torch/csrc/utils/python_arg_parser.h>
-#include <torch/csrc/utils/six.h>
 #ifdef USE_DISTRIBUTED
 #include <torch/csrc/distributed/rpc/py_rref.h>
 #include <torch/csrc/distributed/rpc/rref_impl.h>
@@ -429,6 +428,14 @@ inline InferredType tryToInferType(py::handle input) {
     return InferredType(IntType::get());
   }
 
+  // Check for types registered in _tryToInferTypeImpl (e.g. ProcessGroup)
+  // before falling through to the expensive inspect.isclass / JIT compilation
+  // path below.
+  auto ty = detail::_tryToInferTypeImpl(input);
+  if (ty.has_value()) {
+    return ty.value();
+  }
+
   auto enum_type = py::module::import("enum").attr("Enum");
   py::bool_ isEnumValue = py::isinstance(input, enum_type);
   if (py::cast<bool>(isEnumValue)) {
@@ -518,11 +525,6 @@ inline InferredType tryToInferType(py::handle input) {
     return InferredType("Cannot infer concrete type of torch.nn.Module");
   }
 
-  auto ty = detail::_tryToInferTypeImpl(input);
-  if (ty.has_value()) {
-    return ty.value();
-  }
-
   // Try container types
   return tryToInferContainerType(input, false);
 }
@@ -554,7 +556,7 @@ inline InferredType tryToInferPrimitiveType(py::handle input) {
 inline InferredType tryToInferContainerType(
     py::handle input,
     bool primitiveTypeOnly = false) {
-  if (six::isTuple(input)) {
+  if (PyTuple_Check(input.ptr())) {
     py::tuple tuple = py::cast<py::tuple>(input);
     std::vector<TypePtr> element_types;
     element_types.reserve(tuple.size());
diff --git a/torch/csrc/jit/python/python_arg_flatten.cpp b/torch/csrc/jit/python/python_arg_flatten.cpp
index b71f21d043a31..fde4b35f37560 100644
--- a/torch/csrc/jit/python/python_arg_flatten.cpp
+++ b/torch/csrc/jit/python/python_arg_flatten.cpp
@@ -1,7 +1,7 @@
 #include <c10/util/irange.h>
 #include <torch/csrc/jit/python/python_arg_flatten.h>
 #include <torch/csrc/utils/python_strings.h>
-#include <torch/csrc/utils/six.h>
+#include <torch/csrc/utils/structseq.h>
 
 #include <torch/csrc/autograd/grad_mode.h>
 
@@ -44,7 +44,7 @@ py::object cast_handle_sequence(std::vector<py::handle> objs) {
 
 void flatten_rec(PyObject* obj, ParsedArgs& args) {
   auto& structure = args.desc.structure;
-  if (six::isTuple(obj)) {
+  if (PyTuple_Check(obj)) {
     structure.push_back(D::TupleOpen);
     for (auto item : py::reinterpret_borrow<py::tuple>(obj))
       flatten_rec(item.ptr(), args);
diff --git a/torch/csrc/jit/runtime/graph_executor.cpp b/torch/csrc/jit/runtime/graph_executor.cpp
index 4bdab8c5dcb22..fa37ccab96d06 100644
--- a/torch/csrc/jit/runtime/graph_executor.cpp
+++ b/torch/csrc/jit/runtime/graph_executor.cpp
@@ -635,6 +635,9 @@ struct GraphExecutorImpl : public GraphExecutorImplBase {
   const ExecutionPlan& getPlanFor(
       Stack& stack,
       std::optional<size_t> remaining_bailout_depth) override {
+    if (FLAGS_torch_jit_input_independent_optimization) {
+      return getInputIndependentPlan();
+    }
     return getGraphExecutorOptimize() ? getOrCompile(stack)
                                       : getOrCompileFallback();
   }
@@ -778,6 +781,38 @@ struct GraphExecutorImpl : public GraphExecutorImplBase {
     return ExecutionPlan(opt_graph, function_name_);
   }
 
+  const ExecutionPlan& getInputIndependentPlan() override {
+    std::lock_guard<std::mutex> lock(compile_mutex);
+    if (!input_independent_plan_) {
+      auto opt_graph = graph->copy();
+      Inline(*opt_graph);
+      LowerGradOf(*opt_graph);
+      specializeAutogradZero(opt_graph);
+      LowerSimpleTuples(opt_graph);
+      ConstantPooling(opt_graph);
+      runRequiredPasses(opt_graph);
+      ConstantPropagation(opt_graph);
+      // Skip PropagateInputShapes and PropagateRequiresGrad since they need
+      // actual input data.
+      runOptimization(opt_graph);
+
+      // Input-independent passes from runNondiffOptimization. Skipped:
+      // FuseTensorExprs/FuseGraph (need specialized tensor types).
+      for (const auto& passPair : getCustomPrePasses()) {
+        passPair.first(opt_graph);
+      }
+      DecomposeOps(opt_graph);
+      BatchMM(opt_graph);
+      for (const auto& passPair : getCustomPostPasses()) {
+        passPair.first(opt_graph);
+      }
+
+      EliminateDeadCode(opt_graph);
+      input_independent_plan_ = ExecutionPlan(opt_graph, function_name_);
+    }
+    return *input_independent_plan_;
+  }
+
   ~GraphExecutorImpl() override = default;
 
   // NOLINTNEXTLINE(cppcoreguidelines-non-private-member-variables-in-classes)
@@ -787,6 +822,10 @@ struct GraphExecutorImpl : public GraphExecutorImplBase {
   // NOLINTNEXTLINE(cppcoreguidelines-non-private-member-variables-in-classes)
   ExecutionPlan fallback;
 
+  // Cached plan from getOptimizedPlan() -- uses only input-independent passes.
+  // NOLINTNEXTLINE(cppcoreguidelines-non-private-member-variables-in-classes)
+  std::optional<ExecutionPlan> input_independent_plan_;
+
   // Mapping from argument configurations to optimized versions of the graph
   // that are specialized to the spec.
   // NOLINTNEXTLINE(cppcoreguidelines-non-private-member-variables-in-classes)
@@ -842,6 +881,10 @@ const ExecutionPlan& GraphExecutor::getPlanFor(
   return pImpl->getPlanFor(inputs, remaining_bailout_depth);
 }
 
+const ExecutionPlan& GraphExecutor::getInputIndependentPlan() {
+  return pImpl->getInputIndependentPlan();
+}
+
 GraphExecutorState GraphExecutor::getDebugState() {
   return pImpl->getDebugState();
 }
diff --git a/torch/csrc/jit/runtime/graph_executor.h b/torch/csrc/jit/runtime/graph_executor.h
index d1039216de3ea..94ed97f1784e8 100644
--- a/torch/csrc/jit/runtime/graph_executor.h
+++ b/torch/csrc/jit/runtime/graph_executor.h
@@ -87,6 +87,11 @@ struct TORCH_API GraphExecutor {
   const ExecutionPlan& getPlanFor(
       Stack& inputs,
       std::optional<size_t> remaining_bailout_depth = std::nullopt);
+  // Returns an optimized execution plan without requiring input arguments.
+  // Runs input-independent optimization passes (e.g. inlining, constant
+  // propagation, peephole, CSE) but skips profiling-based specializations
+  // that require runtime type/shape information.
+  const ExecutionPlan& getInputIndependentPlan();
   GraphExecutorState getDebugState();
 
   void debugFlushCompilationCache();
diff --git a/torch/csrc/jit/runtime/graph_executor_impl.h b/torch/csrc/jit/runtime/graph_executor_impl.h
index 70069ac1907b0..78c7a7b79d982 100644
--- a/torch/csrc/jit/runtime/graph_executor_impl.h
+++ b/torch/csrc/jit/runtime/graph_executor_impl.h
@@ -79,6 +79,11 @@ struct GraphExecutorImplBase {
   virtual const ExecutionPlan& getPlanFor(
       Stack& stack,
       std::optional<size_t> remaining_bailout_depth = std::nullopt) = 0;
+  // Returns an optimized execution plan without requiring input arguments.
+  // Runs input-independent optimization passes (e.g. inlining, constant
+  // propagation, peephole, CSE) but skips profiling-based specializations
+  // that require runtime type/shape information.
+  virtual const ExecutionPlan& getInputIndependentPlan() = 0;
   virtual GraphExecutorState getDebugState() = 0;
   virtual ~GraphExecutorImplBase() = default;
 
diff --git a/torch/csrc/jit/runtime/profiling_graph_executor_impl.cpp b/torch/csrc/jit/runtime/profiling_graph_executor_impl.cpp
index 680244b363c36..5ae1ec45050a0 100644
--- a/torch/csrc/jit/runtime/profiling_graph_executor_impl.cpp
+++ b/torch/csrc/jit/runtime/profiling_graph_executor_impl.cpp
@@ -56,6 +56,13 @@ C10_DEFINE_bool(
     false,
     "fuse on 12 dynamic compilations")
 
+C10_DEFINE_bool(
+    torch_jit_input_independent_optimization,
+    false,
+    "If set, getPlanFor will use input-independent optimization passes only, "
+    "skipping profiling-based specializations that require runtime type/shape "
+    "information. Useful for predictor nets and AOT compilation scenarios.")
+
 C10_DEFINE_bool(
     torch_jit_release_profiling_graph_after_optimization,
     false,
@@ -111,8 +118,9 @@ static FusionStrategy getInitialStrategy() {
 // TODO remove ifdef
 #ifdef FBCODE_CAFFE2
   return {{FusionBehavior::STATIC, 20}};
-#endif
+#else
   return mixed;
+#endif
 }
 
 // defer initial value so that we can load in gflags
@@ -714,10 +722,59 @@ const ExecutionPlan& ProfilingGraphExecutorImpl::getPlanFor(
     }
     return *optimized_plan_;
   }
+  if (FLAGS_torch_jit_input_independent_optimization) {
+    return getInputIndependentPlanImpl();
+  }
   // if depth is not set, use
   return getOptimizedPlanFor(stack, remaining_bailout_depth);
 }
 
+const ExecutionPlan& ProfilingGraphExecutorImpl::getInputIndependentPlan() {
+  std::lock_guard<std::mutex> lock(compile_mutex);
+  return getInputIndependentPlanImpl();
+}
+
+const ExecutionPlan&
+ProfilingGraphExecutorImpl::getInputIndependentPlanImpl() {
+  if (optimized_plan_) {
+    return *optimized_plan_;
+  }
+
+  auto copy = graph->copy();
+  if (!getGraphExecutorOptimize() || !getProfilingMode()) {
+    LowerGradOf(*copy);
+    RemoveExpands(copy);
+  } else {
+    // Run all input-independent optimizations. This includes
+    // runProfilingInsensitiveOptimizations (inlining, constant propagation,
+    // CSE, peephole, etc.) followed by runPreAutodiffPassPipeline which adds
+    // loop unrolling, list mutation removal, and additional rounds of
+    // peephole + constant propagation. Profiling-dependent passes (type
+    // specialization, fusion, autodiff guards) are skipped since they
+    // require runtime type/shape information from actual inputs.
+    runProfilingInsensitiveOptimizations(copy);
+    runPreAutodiffPassPipeline(copy);
+
+    // Run additional input-independent passes from runNoGradOptimizations
+    // and runFinalOptimizations. These operate on graph structure (node
+    // patterns, alias analysis) and do not require profiled type/shape info.
+    // Skipped: RemoveProfileNodesAndSpecializeTypes, FuseTensorExprs,
+    // FuseGraph (these need profiled tensor types for specialization/fusion).
+    for (const auto& passPair : getCustomPrePasses()) {
+      passPair.first(copy);
+    }
+    BatchMM(copy);
+    for (const auto& passPair : getCustomPostPasses()) {
+      passPair.first(copy);
+    }
+    AddIfThenElseOp(copy);
+    EliminateDeadCode(copy);
+  }
+  optimized_plan_ = ExecutionPlan(copy, function_name_);
+  time_optimized_plan_created_ = getNowInSecs();
+  return *optimized_plan_;
+}
+
 GraphExecutorState ProfilingGraphExecutorImpl::getDebugState() {
   GraphExecutorState state;
   TORCH_INTERNAL_ASSERT(optimized_plan_);
diff --git a/torch/csrc/jit/runtime/profiling_graph_executor_impl.h b/torch/csrc/jit/runtime/profiling_graph_executor_impl.h
index da437bb456e92..422a7646cef8a 100644
--- a/torch/csrc/jit/runtime/profiling_graph_executor_impl.h
+++ b/torch/csrc/jit/runtime/profiling_graph_executor_impl.h
@@ -7,6 +7,7 @@ TORCH_DECLARE_bool(torch_jit_static_then_dynamic);
 
 TORCH_DECLARE_bool(torch_jit_always_dynamic);
 
+C10_DECLARE_bool(torch_jit_input_independent_optimization);
 C10_DECLARE_bool(torch_jit_release_profiling_graph_after_optimization);
 C10_DECLARE_int32(torch_jit_release_profiling_graph_delay_in_seconds);
 C10_DECLARE_int64(torch_jit_num_profiled_runs);
@@ -24,6 +25,7 @@ struct TORCH_API ProfilingGraphExecutorImpl : public GraphExecutorImplBase {
   const ExecutionPlan& getPlanFor(
       Stack& stack,
       std::optional<size_t> remaining_bailout_depth) override;
+  const ExecutionPlan& getInputIndependentPlan() override;
   GraphExecutorState getDebugState() override;
   ~ProfilingGraphExecutorImpl() override = default;
 
@@ -37,6 +39,8 @@ struct TORCH_API ProfilingGraphExecutorImpl : public GraphExecutorImplBase {
   const ExecutionPlan& getOptimizedPlanFor(
       Stack& stack,
       std::optional<size_t> remaining_bailout_depth);
+  // Input-independent optimization, assumes compile_mutex is held.
+  const ExecutionPlan& getInputIndependentPlanImpl();
   void runProfilingInsensitiveOptimizations(std::shared_ptr<Graph>& graph);
   void runProfilingOptimizations(
       std::shared_ptr<Graph>& graph,
diff --git a/torch/csrc/jit/runtime/register_ops_utils.cpp b/torch/csrc/jit/runtime/register_ops_utils.cpp
index adbfdb46f0932..dfd6fba02de4f 100644
--- a/torch/csrc/jit/runtime/register_ops_utils.cpp
+++ b/torch/csrc/jit/runtime/register_ops_utils.cpp
@@ -132,7 +132,6 @@ void checkDoubleInRange(double a) {
       a < double(std::numeric_limits<int64_t>::min())) {
     throw c10::Error(
         "Cannot convert float " + std::to_string(a) + " to integer");
-    return;
   }
 }
 
diff --git a/torch/csrc/jit/runtime/register_prim_ops.cpp b/torch/csrc/jit/runtime/register_prim_ops.cpp
index 1f168d24e8adf..e512853a57567 100644
--- a/torch/csrc/jit/runtime/register_prim_ops.cpp
+++ b/torch/csrc/jit/runtime/register_prim_ops.cpp
@@ -583,56 +583,6 @@ static const std::vector<OperatorGeneratorArgs> opGenArgs{
           push(stack, arg.element_size());
         },
         aliasAnalysisFromSchema()),
-    OperatorGeneratorArgs(
-        TORCH_SELECTIVE_SCHEMA("aten::numel(Tensor self) -> int"),
-        [](Stack& stack) {
-          at::Tensor arg = pop(stack).toTensor();
-          push(stack, arg.numel());
-        },
-        aliasAnalysisFromSchema()),
-    OperatorGeneratorArgs(
-        TORCH_SELECTIVE_SCHEMA("aten::dim(Tensor self) -> int"),
-        dim,
-        aliasAnalysisFromSchema()),
-    OperatorGeneratorArgs(
-        TORCH_SELECTIVE_SCHEMA("aten::get_device(Tensor self) -> int"),
-        [](Stack& stack) {
-          RECORD_FUNCTION("get_device", c10::ArrayRef<const c10::IValue>{});
-          auto result =
-              at::get_device((std::move(peek(stack, 0, 1))).toTensor());
-          drop(stack, 1);
-          pack(stack, result);
-        },
-        aliasAnalysisFromSchema()),
-    OperatorGeneratorArgs(
-        TORCH_SELECTIVE_SCHEMA("aten::storage_offset(Tensor self) -> int"),
-        [](Stack& stack) {
-          RECORD_FUNCTION("storage_offset", c10::ArrayRef<const c10::IValue>{});
-          auto result =
-              ((std::move(peek(stack, 0, 1))).toTensor()).storage_offset();
-          drop(stack, 1);
-          pack(stack, result);
-        },
-        aliasAnalysisFromSchema()),
-    OperatorGeneratorArgs(
-        TORCH_SELECTIVE_SCHEMA("aten::is_contiguous(Tensor self) -> bool"),
-        [](Stack& stack) {
-          RECORD_FUNCTION("is_contiguous", c10::ArrayRef<const c10::IValue>{});
-          auto result =
-              ((std::move(peek(stack, 0, 1))).toTensor()).is_contiguous();
-          drop(stack, 1);
-          pack(stack, result);
-        },
-        aliasAnalysisFromSchema()),
-    OperatorGeneratorArgs(
-        TORCH_SELECTIVE_SCHEMA(
-            "aten::is_contiguous.memory_format(Tensor self, MemoryFormat memory_format) -> bool"),
-        [](Stack& stack) {
-          auto memory_format = pop(stack).toMemoryFormat();
-          auto t = pop(stack).toTensor();
-          push(stack, t.is_contiguous(memory_format));
-        },
-        aliasAnalysisFromSchema()),
     OperatorGeneratorArgs(
         // NB: intentionally suffixed with extra _format to prevent tests for
         // "_like" suffix from triggering on this
diff --git a/torch/csrc/jit/runtime/simple_graph_executor_impl.cpp b/torch/csrc/jit/runtime/simple_graph_executor_impl.cpp
index fd908b48ee043..4671ab0fdcaf7 100644
--- a/torch/csrc/jit/runtime/simple_graph_executor_impl.cpp
+++ b/torch/csrc/jit/runtime/simple_graph_executor_impl.cpp
@@ -28,6 +28,22 @@ const ExecutionPlan& SimpleGraphExecutorImpl::getPlanFor(
   return *execution_plan_;
 }
 
+const ExecutionPlan& SimpleGraphExecutorImpl::getInputIndependentPlan() {
+  std::lock_guard<std::mutex> lock(compile_mutex);
+  return getInputIndependentPlanImpl();
+}
+
+const ExecutionPlan& SimpleGraphExecutorImpl::getInputIndependentPlanImpl() {
+  if (execution_plan_) {
+    return *execution_plan_;
+  }
+  auto copy = graph->copy();
+  runNooptPassPipeline(copy);
+  execution_plan_ = ExecutionPlan(copy, function_name_);
+
+  return *execution_plan_;
+}
+
 GraphExecutorState SimpleGraphExecutorImpl::getDebugState() {
   GraphExecutorState state;
   TORCH_INTERNAL_ASSERT(execution_plan_);
diff --git a/torch/csrc/jit/runtime/simple_graph_executor_impl.h b/torch/csrc/jit/runtime/simple_graph_executor_impl.h
index e1ebed46ede80..d972432b6883e 100644
--- a/torch/csrc/jit/runtime/simple_graph_executor_impl.h
+++ b/torch/csrc/jit/runtime/simple_graph_executor_impl.h
@@ -13,10 +13,13 @@ struct TORCH_API SimpleGraphExecutorImpl : public GraphExecutorImplBase {
   const ExecutionPlan& getPlanFor(
       Stack& stack,
       std::optional<size_t> remaining_bailout_depth) override;
+  const ExecutionPlan& getInputIndependentPlan() override;
   GraphExecutorState getDebugState() override;
   ~SimpleGraphExecutorImpl() override = default;
 
  private:
+  const ExecutionPlan& getInputIndependentPlanImpl();
+
   std::optional<ExecutionPlan> execution_plan_;
 };
 
diff --git a/torch/csrc/jit/runtime/static/impl.h b/torch/csrc/jit/runtime/static/impl.h
index 24f8f01d7547e..c8a880aaadd12 100644
--- a/torch/csrc/jit/runtime/static/impl.h
+++ b/torch/csrc/jit/runtime/static/impl.h
@@ -241,7 +241,6 @@ class TORCH_API StaticRuntimeMetadata : public torch::CustomClassHolder {
 /// @endcode
 ///
 class MemoryPlanner;
-class StaticNodeInfo;
 class ProcessedNode;
 class StaticRuntime;
 
@@ -296,6 +295,36 @@ class TORCH_API ProcessedFunction {
   size_t num_outputs_{0};
 };
 
+class TORCH_API StaticNodeInfo {
+ public:
+  StaticNodeInfo(
+      Node* n,
+      ProcessedFunction* fn,
+      ProcessedNodeInputs inputs,
+      uint16_t outputs_offset);
+
+  Node* node() const {
+    return node_;
+  }
+
+  size_t num_outputs() const {
+    DCHECK(fn_ != nullptr);
+    return fn_->num_outputs();
+  }
+
+  bool has_out_variant() const {
+    return fn_->kind() == ProcessedFunction::Kind::kOutVariant;
+  }
+
+ private:
+  friend class ProcessedNode;
+
+  Node* node_;
+  const ProcessedFunction* fn_;
+  ProcessedNodeInputs inputs_;
+  uint16_t outputs_offset_;
+};
+
 // A `BlockInfo` instance stores all of the shared state that each
 // `BlockRunner` will need to access. Most of this information is
 // read-only and shared between threads.
@@ -815,36 +844,6 @@ class TORCH_API BlockRunner {
   std::vector<ProcessedNode> nodes_;
 };
 
-class TORCH_API StaticNodeInfo {
- public:
-  StaticNodeInfo(
-      Node* n,
-      ProcessedFunction* fn,
-      ProcessedNodeInputs inputs,
-      uint16_t outputs_offset);
-
-  Node* node() const {
-    return node_;
-  }
-
-  size_t num_outputs() const {
-    DCHECK(fn_ != nullptr);
-    return fn_->num_outputs();
-  }
-
-  bool has_out_variant() const {
-    return fn_->kind() == ProcessedFunction::Kind::kOutVariant;
-  }
-
- private:
-  friend class ProcessedNode;
-
-  Node* node_;
-  const ProcessedFunction* fn_;
-  ProcessedNodeInputs inputs_;
-  uint16_t outputs_offset_;
-};
-
 inline size_t BlockInfo::num_nodes() const {
   return nodes_.size();
 }
diff --git a/torch/csrc/jit/serialization/export.cpp b/torch/csrc/jit/serialization/export.cpp
index 59ed5281db6bc..c1d35d6069d9f 100644
--- a/torch/csrc/jit/serialization/export.cpp
+++ b/torch/csrc/jit/serialization/export.cpp
@@ -1446,7 +1446,6 @@ void check_onnx_proto(const std::string& proto_string) {
   onnx::ModelProto model;
   if (!ParseProtoFromBytes(&model, proto_string.c_str(), proto_string.size())) {
     throw std::runtime_error("Invalid ONNX proto string.");
-    return;
   }
   // 1. baseline check
   // These two checks prevent broken graph being generated
diff --git a/torch/csrc/jit/tensorexpr/codegen.cpp b/torch/csrc/jit/tensorexpr/codegen.cpp
index 04034554e25ed..967cfea3a55cc 100644
--- a/torch/csrc/jit/tensorexpr/codegen.cpp
+++ b/torch/csrc/jit/tensorexpr/codegen.cpp
@@ -21,6 +21,16 @@ CodeGen::CodeGen(
   allocIntermediateBufs();
 }
 
+CodeGen::CodeGen(const CodeGen& rhs) = default;
+
+CodeGen::CodeGen(CodeGen&& rhs) = default;
+
+CodeGen::~CodeGen() = default;
+
+CodeGen& CodeGen::operator=(const CodeGen& rhs) = default;
+
+CodeGen& CodeGen::operator=(CodeGen&& rhs) = default;
+
 RegisterCodeGenList& RegisterCodeGenList::GetInstance() {
   static RegisterCodeGenList codegen_list;
   return codegen_list;
@@ -93,7 +103,6 @@ void* CodeGen::argToPtr(const BufferArg& bufferArg, const CallArg& callArg) {
     default:
       throw unsupported_dtype();
   }
-  return nullptr;
 }
 
 void CodeGen::call_with_numel(void** args, int64_t numel) {
diff --git a/torch/csrc/jit/tensorexpr/codegen.h b/torch/csrc/jit/tensorexpr/codegen.h
index cad930b58bd93..f333906339610 100644
--- a/torch/csrc/jit/tensorexpr/codegen.h
+++ b/torch/csrc/jit/tensorexpr/codegen.h
@@ -26,7 +26,15 @@ class TORCH_API CodeGen {
       at::Device device = at::kCPU,
       std::string kernel_func_name = "func");
 
-  virtual ~CodeGen() = default;
+  CodeGen(const CodeGen& rhs);
+
+  CodeGen(CodeGen&& rhs);
+
+  virtual ~CodeGen();
+
+  CodeGen& operator=(const CodeGen& rhs);
+
+  CodeGen& operator=(CodeGen&& rhs);
 
   StmtPtr stmt() const {
     return stmt_;
diff --git a/torch/csrc/jit/tensorexpr/cuda_codegen.cpp b/torch/csrc/jit/tensorexpr/cuda_codegen.cpp
index c787ccd88ddcf..a54e1eae6fc75 100644
--- a/torch/csrc/jit/tensorexpr/cuda_codegen.cpp
+++ b/torch/csrc/jit/tensorexpr/cuda_codegen.cpp
@@ -1280,7 +1280,7 @@ void CudaCodeGen::CompileToNVRTC(
       &program, code.c_str(), nullptr, 0, nullptr, nullptr));
 
 #if defined(USE_ROCM)
-  std::vector<const char*> args = {"--std=c++17"};
+  std::vector<const char*> args = {"--std=c++20"};
   args.push_back("-hip-pch");
 #else
   const std::string compute = std::string("--gpu-architecture=") +
@@ -1298,7 +1298,7 @@ void CudaCodeGen::CompileToNVRTC(
 #endif
       std::to_string(major) + std::to_string(minor);
   const std::vector<const char*> args = {
-      "--std=c++17", compute.c_str(), "-default-device"};
+      "--std=c++20", compute.c_str(), "-default-device"};
 #endif
 
   auto result = nvrtc().nvrtcCompileProgram(
diff --git a/torch/csrc/jit/tensorexpr/eval.cpp b/torch/csrc/jit/tensorexpr/eval.cpp
index 1709ef2bbff5a..9c8bd175a50e6 100644
--- a/torch/csrc/jit/tensorexpr/eval.cpp
+++ b/torch/csrc/jit/tensorexpr/eval.cpp
@@ -20,7 +20,6 @@ int64_t InterpValue::intValue() const {
   AT_FORALL_INT_TYPES(TYPE_CASE);
 #undef TYPE_CASE
   throw unsupported_dtype();
-  return 0;
 }
 
 template <typename T>
@@ -662,7 +661,6 @@ class SimpleIREvaluatorImpl : public IRVisitor {
       default:
         throw unsupported_dtype();
     }
-    return {};
   }
 
   void check_bounds_throw(int64_t idx, int64_t bound, const BufPtr& buf) {
diff --git a/torch/csrc/jit/tensorexpr/ir_simplifier.cpp b/torch/csrc/jit/tensorexpr/ir_simplifier.cpp
index 53eca369cc5ac..7d918009e2421 100644
--- a/torch/csrc/jit/tensorexpr/ir_simplifier.cpp
+++ b/torch/csrc/jit/tensorexpr/ir_simplifier.cpp
@@ -37,8 +37,8 @@ static inline ExprPtr newBinaryOpOfType(
     case IRNodeType::kRshift:
       return alloc<Rshift>(lhs, rhs);
     default:
-      LOG(FATAL) << "unsupported expr_type: " << static_cast<int>(expr_type);
-      return nullptr;
+      TORCH_INTERNAL_ASSERT(
+          false, "unsupported expr_type: ", static_cast<int>(expr_type));
   }
 }
 
diff --git a/torch/csrc/jit/tensorexpr/kernel.cpp b/torch/csrc/jit/tensorexpr/kernel.cpp
index d696d29bf733e..84d78a665e62d 100644
--- a/torch/csrc/jit/tensorexpr/kernel.cpp
+++ b/torch/csrc/jit/tensorexpr/kernel.cpp
@@ -330,8 +330,9 @@ bool mkldnnPrepackedConvIsSupportedJit(const torch::jit::Node* node) {
       _pair_int(*pad),
       _pair_int(*dilation),
       groups->toInt());
-#endif
+#else
   return false;
+#endif
 }
 
 bool isConv2d(const Node* node) {
diff --git a/torch/csrc/jit/tensorexpr/loopnest.cpp b/torch/csrc/jit/tensorexpr/loopnest.cpp
index 1bdae4ca7ae90..e0a7a64396c04 100644
--- a/torch/csrc/jit/tensorexpr/loopnest.cpp
+++ b/torch/csrc/jit/tensorexpr/loopnest.cpp
@@ -2830,7 +2830,6 @@ LoopNest::AccessResult LoopNest::cacheAccesses(
       if (reduceOp) {
         throw std::runtime_error(
             "can only cache accesses used by at most a single reduceOp");
-        return {nullptr, nullptr};
       }
 
       reduceOp = ro;
@@ -2842,7 +2841,6 @@ LoopNest::AccessResult LoopNest::cacheAccesses(
   auto bounds_it = consumer_bounds_info.find(producer);
   if (bounds_it == consumer_bounds_info.end()) {
     throw std::runtime_error("consumer does not use the Tensor produced");
-    return {nullptr, nullptr};
   }
 
   TORCH_INTERNAL_ASSERT(
diff --git a/torch/csrc/jit/tensorexpr/operators/conv2d.cpp b/torch/csrc/jit/tensorexpr/operators/conv2d.cpp
index 2dd46335e09f8..3b0b21f94699c 100644
--- a/torch/csrc/jit/tensorexpr/operators/conv2d.cpp
+++ b/torch/csrc/jit/tensorexpr/operators/conv2d.cpp
@@ -341,8 +341,9 @@ bool mkldnnPrepackedConvIsSupported(
       input.dims[0] * input.dims[1] * input.dims[2] * input.dims[3] > 20480;
   GRAPH_DEBUG("mkldnnPrepackedConvIsSupported: ", use_mkldnn);
   return use_mkldnn;
-#endif
+#else
   return false;
+#endif
 }
 
 Tensor computeConv2d(
diff --git a/torch/csrc/jit/tensorexpr/operators/misc.cpp b/torch/csrc/jit/tensorexpr/operators/misc.cpp
index f633923723535..266199bd7216e 100644
--- a/torch/csrc/jit/tensorexpr/operators/misc.cpp
+++ b/torch/csrc/jit/tensorexpr/operators/misc.cpp
@@ -152,8 +152,6 @@ ExprHandle demoteOutput(
     default:
       throw unsupported_dtype();
   }
-
-  return e;
 }
 
 std::optional<TensorInfo> getTensorInfo(const BufHandle& b) {
diff --git a/torch/csrc/jit/tensorexpr/registerizer.cpp b/torch/csrc/jit/tensorexpr/registerizer.cpp
index 9ad44e31a3873..acf216c51b530 100644
--- a/torch/csrc/jit/tensorexpr/registerizer.cpp
+++ b/torch/csrc/jit/tensorexpr/registerizer.cpp
@@ -161,17 +161,12 @@ AccessHashMap& Scope::getAccessMapByBuf(const BufPtr& b) {
 }
 
 void Scope::filterClosed() {
-  closedAccesses_.erase(
-      std::remove_if(
-          closedAccesses_.begin(),
-          closedAccesses_.end(),
-          [](auto info) {
-            return info->store_cost()->isConstant() &&
-                immediateAs<int>(info->store_cost()) <= 1 &&
-                info->load_cost()->isConstant() &&
-                immediateAs<int>(info->load_cost()) <= 1;
-          }),
-      closedAccesses_.end());
+  std::erase_if(closedAccesses_, [](auto info) {
+    return info->store_cost()->isConstant() &&
+        immediateAs<int>(info->store_cost()) <= 1 &&
+        info->load_cost()->isConstant() &&
+        immediateAs<int>(info->load_cost()) <= 1;
+  });
 }
 
 // RegisterizerAnalysis
diff --git a/torch/csrc/jit/tensorexpr/types.cpp b/torch/csrc/jit/tensorexpr/types.cpp
index 57f6c1c9ec342..1873035014fc0 100644
--- a/torch/csrc/jit/tensorexpr/types.cpp
+++ b/torch/csrc/jit/tensorexpr/types.cpp
@@ -115,7 +115,6 @@ std::string Dtype::ToCppString() const {
     default:
       throw unsupported_dtype();
   }
-  return "invalid";
 }
 
 } // namespace torch::jit::tensorexpr
diff --git a/torch/csrc/jit/testing/file_check.cpp b/torch/csrc/jit/testing/file_check.cpp
index 0e792934472d1..27ef7557837e1 100644
--- a/torch/csrc/jit/testing/file_check.cpp
+++ b/torch/csrc/jit/testing/file_check.cpp
@@ -144,8 +144,6 @@ size_t assertFindRegex(
       extra_msg(ss);
     }
     throw std::runtime_error(ss.str());
-
-    return std::string::npos;
   }
   return pos;
 }
diff --git a/torch/csrc/lazy/core/ir.cpp b/torch/csrc/lazy/core/ir.cpp
index 3fc49babfc464..8201617d5abf6 100644
--- a/torch/csrc/lazy/core/ir.cpp
+++ b/torch/csrc/lazy/core/ir.cpp
@@ -96,6 +96,16 @@ Node::Node(OpKind op, Shape shape, size_t num_outputs) : Node(op, num_outputs) {
   shapes_.push_back(std::move(shape));
 }
 
+Node::Node(const Node& rhs) = default;
+
+Node::Node(Node&& rhs) = default;
+
+Node::~Node() = default;
+
+Node& Node::operator=(const Node& rhs) = default;
+
+Node& Node::operator=(Node&& rhs) = default;
+
 // Retrieves the full shape of the IR Node.
 c10::ArrayRef<Shape> Node::shapes() const {
   return shapes_;
diff --git a/torch/csrc/lazy/core/ir.h b/torch/csrc/lazy/core/ir.h
index b55745fc96ae2..6e1a6f6f162a0 100644
--- a/torch/csrc/lazy/core/ir.h
+++ b/torch/csrc/lazy/core/ir.h
@@ -101,7 +101,15 @@ class TORCH_API Node {
   // Construct node with shape and no operands
   Node(OpKind op, Shape shape, size_t num_outputs = 1);
 
-  virtual ~Node() = default;
+  Node(const Node& rhs);
+
+  Node(Node&& rhs);
+
+  virtual ~Node();
+
+  Node& operator=(const Node& rhs);
+
+  Node& operator=(Node&& rhs);
 
   const OpKind& op() const {
     return op_;
diff --git a/torch/csrc/lazy/python/init.cpp b/torch/csrc/lazy/python/init.cpp
index 4807aa6a4c7d1..07281f7e444c7 100644
--- a/torch/csrc/lazy/python/init.cpp
+++ b/torch/csrc/lazy/python/init.cpp
@@ -272,8 +272,6 @@ void initLazyBindings(PyObject* module) {
 #else
         TORCH_CHECK(
             false, "TorchScript backend not yet supported in FBCODE builds");
-        return std::make_pair(
-            std::vector<int64_t>(), std::vector<at::IValue>());
 #endif // !(defined(FBCODE_CAFFE2) || defined(OVRSOURCE))
       });
   // TODO(shunting) revisit this part for XLA
@@ -303,13 +301,13 @@ void initLazyBindings(PyObject* module) {
         for (torch::jit::IValue elem : stack) {
           result.push_back(elem.toTensor());
         }
+        return result;
 #else
         TORCH_CHECK(
             false, "TorchScript backend not yet supported in FBCODE builds");
 #endif // !(defined(FBCODE_CAFFE2) || defined(OVRSOURCE))
-        return result;
       });
-  lazy_ts_backend.def("_get_latest_computation_graph", []() {
+  lazy_ts_backend.def("_get_latest_computation_graph", []() -> std::string {
 #if !(defined(FBCODE_CAFFE2) || defined(OVRSOURCE))
     auto computation = LazyGraphExecutor::Get()
                            ->GetComputationCache()
@@ -321,7 +319,6 @@ void initLazyBindings(PyObject* module) {
 #else
     TORCH_CHECK(
         false, "TorchScript backend not yet supported in FBCODE builds");
-    return "";
 #endif // !(defined(FBCODE_CAFFE2) || defined(OVRSOURCE))
   });
 
diff --git a/torch/csrc/mtia/Module.cpp b/torch/csrc/mtia/Module.cpp
index a5cfbfbcd4e8b..0b6fe1e51e230 100644
--- a/torch/csrc/mtia/Module.cpp
+++ b/torch/csrc/mtia/Module.cpp
@@ -179,6 +179,14 @@ void initModule(PyObject* module) {
     return at::detail::getMTIAHooks().deviceCount();
   });
 
+  m.def("_mtia_getDevice", []() {
+    return at::detail::getMTIAHooks().getCurrentDevice();
+  });
+
+  m.def("_mtia_setDevice", [](c10::DeviceIndex device_index) {
+    at::detail::getMTIAHooks().setCurrentDevice(device_index);
+  });
+
   m.def("_mtia_resetPeakMemoryStats", [](c10::DeviceIndex device_index) {
     at::detail::getMTIAHooks().resetPeakMemoryStats(device_index);
   });
diff --git a/torch/csrc/profiler/collection.cpp b/torch/csrc/profiler/collection.cpp
index 617316617fc67..1be2d80310910 100644
--- a/torch/csrc/profiler/collection.cpp
+++ b/torch/csrc/profiler/collection.cpp
@@ -10,6 +10,7 @@
 #include <utility>
 
 #include <fmt/format.h>
+#include <nlohmann/json.hpp>
 
 #ifdef USE_KINETO
 #include <libkineto.h>
@@ -1096,6 +1097,18 @@ class TransferEvents {
     for (const auto* activity : trace_activities_) {
       auto e = toResult(activity);
       if (e) {
+        // Flow data for Kineto events is already set during
+        // resultFromActivity(). TorchOp events need it copied here because
+        // their Result is created during RecordFunction callbacks, before
+        // flow data exists on the GenericTraceActivity.
+        e->visit(c10::overloaded(
+            [&](ExtraFields<EventType::TorchOp>& i) {
+              i.flow = {
+                  /*id=*/static_cast<uint32_t>(activity->flowId()),
+                  /*type=*/static_cast<uint32_t>(activity->flowType()),
+                  /*start=*/activity->flowStart()};
+            },
+            [](auto&) {}));
         if (config_.experimental_config.expose_kineto_event_metadata) {
           e->visit(c10::overloaded(
               [&](ExtraFields<EventType::TorchOp>& i) {
@@ -1105,6 +1118,27 @@ class TransferEvents {
                 i.metadata_json_ = activity->metadataJson();
               },
               [](auto&) { return; }));
+          // Parse metadataJson() into extra_meta_ so events() exposes
+          // Kineto metadata as typed fields without export_chrome_trace().
+          // Python schemas (profiler_util.py) are the single SOT for
+          // which keys to expose and how to type-convert them.
+          e->visit(c10::overloaded(
+              [&](ExtraFields<EventType::Kineto>& i) {
+                auto json_str = activity->metadataJson();
+                if (!json_str.empty()) {
+                  auto j = nlohmann::json::parse(
+                      "{" + json_str + "}", nullptr, false);
+                  if (!j.is_discarded()) {
+                    for (auto& [key, val] : j.items()) {
+                      i.extra_meta_.emplace(
+                          key,
+                          val.is_string() ? val.get<std::string>()
+                                          : val.dump());
+                    }
+                  }
+                }
+              },
+              [](auto&) {}));
         }
         const auto* linked_activity = activity->linkedActivity();
         if (linked_activity) {
diff --git a/torch/csrc/profiler/collection.h b/torch/csrc/profiler/collection.h
index d66eb630a47d9..bd46814f8be9b 100644
--- a/torch/csrc/profiler/collection.h
+++ b/torch/csrc/profiler/collection.h
@@ -132,6 +132,14 @@ using extra_args_t = std::unordered_map<std::string, c10::IValue>;
 using extra_meta_t = std::unordered_map<std::string, std::string>;
 using kwinputs_t = std::unordered_map<std::string, c10::IValue>;
 
+// Mirrors `libkineto::GenericTraceActivity::Flow`. Used during post processing
+// to embed Kineto events into the broader profiler tree structure.
+struct Flow {
+  uint32_t id{0};
+  uint32_t type{0};
+  uint32_t start{0};
+};
+
 struct FallbackPair {
   ProfilerVoidEventStub device_event_start_ = nullptr;
   ProfilerVoidEventStub device_event_end_ = nullptr;
@@ -179,6 +187,7 @@ struct ExtraFields<EventType::TorchOp> : TorchOpBasicFields {
   bool allow_tf32_cublas_;
   std::unique_ptr<perf_counters_t> perf_event_counters_;
   std::string metadata_json_;
+  Flow flow;
 };
 
 template <>
@@ -354,16 +363,6 @@ struct ExtraFields<EventType::PyCCall> : public PyExtraFieldsBase {
 
 template <>
 struct ExtraFields<EventType::Kineto> {
-  // Mirrors `libkineto::GenericTraceActivity::Flow`. This information is used
-  // during post processing to properly embed Kineto events into the broader
-  // profiler tree structure. End users are not generally expected to use these
-  // fields directly, but they are available for debugging.
-  struct Flow {
-    uint32_t id{0};
-    uint32_t type{0};
-    uint32_t start{0};
-  };
-
   std::string name_;
   int64_t duration_ns_{0};
   uint64_t correlation_id_{0};
@@ -371,6 +370,7 @@ struct ExtraFields<EventType::Kineto> {
   Flow flow;
   std::weak_ptr<Result> linked_activity_;
   std::string metadata_json_;
+  extra_meta_t extra_meta_;
 };
 
 struct TORCH_API Result : public std::enable_shared_from_this<Result> {
diff --git a/torch/csrc/profiler/data_flow.cpp b/torch/csrc/profiler/data_flow.cpp
index a9f98930f8c66..e18d88960fbba 100644
--- a/torch/csrc/profiler/data_flow.cpp
+++ b/torch/csrc/profiler/data_flow.cpp
@@ -131,15 +131,10 @@ void calculateUniqueTensorIDs(
         tensor_set.insert(t.allocation_id_ref_.get().value());
       }
     }
-    tensors.erase(
-        std::remove_if(
-            tensors.begin(),
-            tensors.end(),
-            [&tensor_set](const auto& i) {
-              auto it = tensor_set.find(i.allocation_id_ref_.get().value());
-              return it == tensor_set.end();
-            }),
-        tensors.end());
+    std::erase_if(tensors, [&tensor_set](const auto& i) {
+      auto it = tensor_set.find(i.allocation_id_ref_.get().value());
+      return it == tensor_set.end();
+    });
   }
 
   // Handle the case that the storage of a TensorImpl changed.
diff --git a/torch/csrc/profiler/kineto_shim.cpp b/torch/csrc/profiler/kineto_shim.cpp
index b93bdeee0cd06..fa232e1a01016 100644
--- a/torch/csrc/profiler/kineto_shim.cpp
+++ b/torch/csrc/profiler/kineto_shim.cpp
@@ -1,5 +1,6 @@
 #include <torch/csrc/profiler/collection.h>
 #include <torch/csrc/profiler/kineto_shim.h>
+
 #include <type_traits>
 
 #ifdef USE_KINETO
@@ -17,55 +18,88 @@ namespace profiler::impl::kineto {
 
 #ifdef USE_KINETO
 namespace {
-const std::set<libkineto::ActivityType> kCpuTypes{
-    libkineto::ActivityType::CPU_OP,
-    libkineto::ActivityType::CPU_INSTANT_EVENT,
-    libkineto::ActivityType::USER_ANNOTATION,
-    libkineto::ActivityType::EXTERNAL_CORRELATION,
-    libkineto::ActivityType::XPU_RUNTIME,
-    libkineto::ActivityType::XPU_DRIVER,
-    libkineto::ActivityType::CUDA_RUNTIME,
-    libkineto::ActivityType::CUDA_DRIVER,
-    libkineto::ActivityType::PYTHON_FUNCTION,
-    libkineto::ActivityType::PRIVATEUSE1_RUNTIME,
-    libkineto::ActivityType::PRIVATEUSE1_DRIVER,
+
+using ActivityTypeMap =
+    std::unordered_map<libkineto::ActivityType, std::string>;
+
+// clang-format off
+const ActivityTypeMap kCpuTypes{
+    {libkineto::ActivityType::CPU_OP,                "CPU_OP"},
+    {libkineto::ActivityType::CPU_INSTANT_EVENT,     "CPU_INSTANT_EVENT"},
+    {libkineto::ActivityType::USER_ANNOTATION,       "USER_ANNOTATION"},
+    {libkineto::ActivityType::EXTERNAL_CORRELATION,  "EXTERNAL_CORRELATION"},
+    {libkineto::ActivityType::XPU_RUNTIME,           "XPU_RUNTIME"},
+    {libkineto::ActivityType::XPU_DRIVER,            "XPU_DRIVER"},
+    {libkineto::ActivityType::CUDA_RUNTIME,          "CUDA_RUNTIME"},
+    {libkineto::ActivityType::CUDA_DRIVER,           "CUDA_DRIVER"},
+    {libkineto::ActivityType::PYTHON_FUNCTION,       "PYTHON_FUNCTION"},
+    {libkineto::ActivityType::PRIVATEUSE1_RUNTIME,   "PRIVATEUSE1_RUNTIME"},
+    {libkineto::ActivityType::PRIVATEUSE1_DRIVER,    "PRIVATEUSE1_DRIVER"},
 };
 
-const std::set<libkineto::ActivityType> kCudaTypes = {
-    libkineto::ActivityType::GPU_MEMCPY,
-    libkineto::ActivityType::GPU_MEMSET,
-    libkineto::ActivityType::GPU_USER_ANNOTATION,
-    libkineto::ActivityType::CONCURRENT_KERNEL,
+const ActivityTypeMap kCudaTypes{
+    {libkineto::ActivityType::GPU_MEMCPY,            "GPU_MEMCPY"},
+    {libkineto::ActivityType::GPU_MEMSET,            "GPU_MEMSET"},
+    {libkineto::ActivityType::GPU_USER_ANNOTATION,   "GPU_USER_ANNOTATION"},
+    {libkineto::ActivityType::CONCURRENT_KERNEL,     "CONCURRENT_KERNEL"},
     // CUDA_RUNTIME and CUDA_DRIVER appear in both kCpuTypes and kCudaTypes.
-    libkineto::ActivityType::CUDA_RUNTIME,
-    libkineto::ActivityType::CUDA_DRIVER,
-    libkineto::ActivityType::OVERHEAD,
+    {libkineto::ActivityType::CUDA_RUNTIME,          "CUDA_RUNTIME"},
+    {libkineto::ActivityType::CUDA_DRIVER,           "CUDA_DRIVER"},
+    {libkineto::ActivityType::OVERHEAD,              "OVERHEAD"},
 };
-const std::set<libkineto::ActivityType> kXpuTypes = {
-    libkineto::ActivityType::GPU_MEMCPY,
-    libkineto::ActivityType::GPU_MEMSET,
-    libkineto::ActivityType::CONCURRENT_KERNEL,
+
+const ActivityTypeMap kXpuTypes{
+    {libkineto::ActivityType::GPU_MEMCPY,            "GPU_MEMCPY"},
+    {libkineto::ActivityType::GPU_MEMSET,            "GPU_MEMSET"},
+    {libkineto::ActivityType::CONCURRENT_KERNEL,     "CONCURRENT_KERNEL"},
     // XPU_RUNTIME and XPU_DRIVER appear in both kCpuTypes and kXpuTypes.
-    libkineto::ActivityType::XPU_RUNTIME,
-    libkineto::ActivityType::XPU_DRIVER,
+    {libkineto::ActivityType::XPU_RUNTIME,           "XPU_RUNTIME"},
+    {libkineto::ActivityType::XPU_DRIVER,            "XPU_DRIVER"},
 };
-const std::set<libkineto::ActivityType> kMtiaTypes = {
-    libkineto::ActivityType::MTIA_CCP_EVENTS,
-    libkineto::ActivityType::MTIA_RUNTIME,
-    libkineto::ActivityType::MTIA_INSIGHT,
+
+const ActivityTypeMap kMtiaTypes{
+    {libkineto::ActivityType::MTIA_CCP_EVENTS,       "MTIA_CCP_EVENTS"},
+    {libkineto::ActivityType::MTIA_RUNTIME,          "MTIA_RUNTIME"},
+    {libkineto::ActivityType::MTIA_INSIGHT,          "MTIA_INSIGHT"},
+    {libkineto::ActivityType::MTIA_COUNTERS,         "MTIA_COUNTERS"},
 };
-const std::set<libkineto::ActivityType> hpuTypes = {
-    libkineto::ActivityType::HPU_OP,
+
+const ActivityTypeMap kHpuTypes{
+    {libkineto::ActivityType::HPU_OP,                "HPU_OP"},
 };
-const std::set<libkineto::ActivityType> kPrivateUse1Types = {
-    libkineto::ActivityType::GPU_MEMCPY,
-    libkineto::ActivityType::GPU_MEMSET,
-    libkineto::ActivityType::GPU_USER_ANNOTATION,
-    libkineto::ActivityType::CONCURRENT_KERNEL,
+
+const ActivityTypeMap kPrivateUse1Types{
+    {libkineto::ActivityType::GPU_MEMCPY,            "GPU_MEMCPY"},
+    {libkineto::ActivityType::GPU_MEMSET,            "GPU_MEMSET"},
+    {libkineto::ActivityType::GPU_USER_ANNOTATION,   "GPU_USER_ANNOTATION"},
+    {libkineto::ActivityType::CONCURRENT_KERNEL,     "CONCURRENT_KERNEL"},
     // PRIVATEUSE1_RUNTIME appears in both kCpuTypes and kPrivateUse1Types.
-    libkineto::ActivityType::PRIVATEUSE1_RUNTIME,
-    libkineto::ActivityType::PRIVATEUSE1_DRIVER,
+    {libkineto::ActivityType::PRIVATEUSE1_RUNTIME,   "PRIVATEUSE1_RUNTIME"},
+    {libkineto::ActivityType::PRIVATEUSE1_DRIVER,    "PRIVATEUSE1_DRIVER"},
 };
+// clang-format on
+
+// Given a named activity type map and a set of requested name strings,
+// return the matching subset of activity types.
+std::unordered_set<libkineto::ActivityType> filterActivities(
+    const ActivityTypeMap& defaults,
+    const std::unordered_set<std::string>& requested) {
+  std::unordered_set<libkineto::ActivityType> result;
+  for (const auto& name : requested) {
+    bool found = false;
+    for (const auto& [type, type_name] : defaults) {
+      if (type_name == name) {
+        result.insert(type);
+        found = true;
+        break;
+      }
+    }
+    TORCH_CHECK(
+        found, "Unknown or non-member activity type name: '", name, "'");
+  }
+  return result;
+}
+
 } // namespace
 #endif // USE_KINETO
 
@@ -189,7 +223,9 @@ class ExperimentalConfigWrapper {
 
     // Add CPU activities if we are measuring per kernel ranges
     if (config_.profiler_measure_per_kernel) {
-      k_activities.insert(kCpuTypes.begin(), kCpuTypes.end());
+      for (const auto& [type, name] : kCpuTypes) {
+        k_activities.insert(type);
+      }
     }
 
     const size_t num_metrics = config_.profiler_metrics.size();
@@ -260,7 +296,8 @@ void prepareTrace(
     const bool cpuOnly,
     const ActivitySet& activities,
     const torch::profiler::impl::ExperimentalConfig& config,
-    const std::string& trace_id) {
+    const std::string& trace_id,
+    const ActivityFilter& activity_filter) {
 #ifdef USE_KINETO
   libkineto::api().resetKinetoTLS();
   if (!libkineto::api().isProfilerRegistered()) {
@@ -273,18 +310,34 @@ void prepareTrace(
   }
 
   std::set<libkineto::ActivityType> k_activities;
+
+  // Helper: insert activity types for a group, applying the filter if present.
+  auto insertActivities = [&](torch::autograd::profiler::ActivityType group,
+                              const ActivityTypeMap& defaults) {
+    auto filter_it = activity_filter.find(group);
+    if (filter_it != activity_filter.end()) {
+      auto filtered = filterActivities(defaults, filter_it->second);
+      k_activities.insert(filtered.begin(), filtered.end());
+    } else {
+      for (const auto& [type, name] : defaults) {
+        k_activities.insert(type);
+      }
+    }
+  };
+
   bool has_cpu_activity =
       activities.count(torch::autograd::profiler::ActivityType::CPU);
 
   if (has_cpu_activity) {
-    k_activities.insert(kCpuTypes.begin(), kCpuTypes.end());
+    insertActivities(torch::autograd::profiler::ActivityType::CPU, kCpuTypes);
   }
   if (activities.count(torch::autograd::profiler::ActivityType::XPU)) {
-    k_activities.insert(kXpuTypes.begin(), kXpuTypes.end());
+    insertActivities(torch::autograd::profiler::ActivityType::XPU, kXpuTypes);
   }
   if (activities.count(torch::autograd::profiler::ActivityType::MTIA)) {
     if (config.custom_profiler_config.empty()) {
-      k_activities.insert(kMtiaTypes.begin(), kMtiaTypes.end());
+      insertActivities(
+          torch::autograd::profiler::ActivityType::MTIA, kMtiaTypes);
     } else {
       if (config.custom_profiler_config.find("disable_runtime_events") ==
           std::string::npos) {
@@ -304,13 +357,19 @@ void prepareTrace(
       } else {
         LOG(INFO) << "Disabling MTIA insight events";
       }
+      if (config.custom_profiler_config.find("disable_counter_events") ==
+          std::string::npos) {
+        k_activities.insert(libkineto::ActivityType::MTIA_COUNTERS);
+      } else {
+        LOG(INFO) << "Disabling MTIA counter events";
+      }
     }
   }
   if (activities.count(torch::autograd::profiler::ActivityType::HPU)) {
-    k_activities.insert(hpuTypes.begin(), hpuTypes.end());
+    insertActivities(torch::autograd::profiler::ActivityType::HPU, kHpuTypes);
   }
   if (activities.count(torch::autograd::profiler::ActivityType::CUDA)) {
-    k_activities.insert(kCudaTypes.begin(), kCudaTypes.end());
+    insertActivities(torch::autograd::profiler::ActivityType::CUDA, kCudaTypes);
     if (config.enable_cuda_sync_events || get_cuda_sync_enabled()) {
       LOG(INFO) << "Enabling CUDA Sync Events";
       k_activities.insert(libkineto::ActivityType::CUDA_SYNC);
@@ -320,7 +379,9 @@ void prepareTrace(
     k_activities.insert(libkineto::ActivityType::COLLECTIVE_COMM);
   }
   if (activities.count(torch::autograd::profiler::ActivityType::PrivateUse1)) {
-    k_activities.insert(kPrivateUse1Types.begin(), kPrivateUse1Types.end());
+    insertActivities(
+        torch::autograd::profiler::ActivityType::PrivateUse1,
+        kPrivateUse1Types);
   }
 
   ExperimentalConfigWrapper configWrap(config);
@@ -436,6 +497,7 @@ c10::DeviceType deviceTypeFromActivity(libkineto::ActivityType activity_type) {
     // TODO: T151322015
     case libkineto::ActivityType::MTIA_CCP_EVENTS:
     case libkineto::ActivityType::MTIA_INSIGHT:
+    case libkineto::ActivityType::MTIA_COUNTERS:
       return device_type_privateuse1_or(c10::DeviceType::MTIA);
     case libkineto::ActivityType::HPU_OP:
       return c10::DeviceType::HPU;
diff --git a/torch/csrc/profiler/kineto_shim.h b/torch/csrc/profiler/kineto_shim.h
index c4efd7785b795..4f9bdc6770507 100644
--- a/torch/csrc/profiler/kineto_shim.h
+++ b/torch/csrc/profiler/kineto_shim.h
@@ -108,11 +108,15 @@ struct ActivityTraceWrapper {
 };
 
 using ActivitySet = std::set<torch::autograd::profiler::ActivityType>;
+using ActivityFilter = std::unordered_map<
+    torch::autograd::profiler::ActivityType,
+    std::unordered_set<std::string>>;
 void prepareTrace(
     const bool cpuOnly,
     const ActivitySet& activities,
     const torch::profiler::impl::ExperimentalConfig& config,
-    const std::string& trace_id = "");
+    const std::string& trace_id = "",
+    const ActivityFilter& activity_filter = {});
 
 void toggleCollectionDynamic(const bool enable);
 void startTrace();
diff --git a/torch/csrc/profiler/orchestration/observer.h b/torch/csrc/profiler/orchestration/observer.h
index 3b59466e6060f..f60e0194d48a5 100644
--- a/torch/csrc/profiler/orchestration/observer.h
+++ b/torch/csrc/profiler/orchestration/observer.h
@@ -38,6 +38,7 @@ enum class C10_API_ENUM ProfilerState {
   KINETO, // use libkineto
   KINETO_GPU_FALLBACK, // use CUDA events when CUPTI is not available
   KINETO_PRIVATEUSE1_FALLBACK, // use PrivateUse1 events
+  KINETO_PRIVATEUSE1, // use Kineto with registered IActivityProfiler
   KINETO_ONDEMAND, // run the profiler in on-demand mode
   NUM_PROFILER_STATES, // must be the last one
 };
diff --git a/torch/csrc/profiler/python/combined_traceback.cpp b/torch/csrc/profiler/python/combined_traceback.cpp
index 4261d226752a3..330f67ab65ffd 100644
--- a/torch/csrc/profiler/python/combined_traceback.cpp
+++ b/torch/csrc/profiler/python/combined_traceback.cpp
@@ -157,10 +157,23 @@ struct PythonTraceback : public CapturedTraceback::Python {
       return result;
     }
 
-    // Get the traceback from the metadata dict
+    // Get the traceback from the metadata dict.
+    // This runs from a CUDA allocator callback, so a Python exception may
+    // already be pending (e.g. the forward function just raised). The compat
+    // shim for PyDict_GetItemRef on Python < 3.13 uses PyErr_Occurred() to
+    // distinguish "not found" from "error", so a stale pending exception would
+    // be misread as a lookup failure and then cleared, destroying the real
+    // exception. Save/restore the exception state to avoid that.
     py::gil_scoped_acquire gil;
+
+    PyObject* exc_type = nullptr;
+    PyObject* exc_value = nullptr;
+    PyObject* exc_tb = nullptr;
+    PyErr_Fetch(&exc_type, &exc_value, &exc_tb);
+
     PyObject* dict = metadata->dict();
     if (!dict || !PyDict_Check(dict)) {
+      PyErr_Restore(exc_type, exc_value, exc_tb);
       return result;
     }
 
@@ -170,11 +183,13 @@ struct PythonTraceback : public CapturedTraceback::Python {
             torch::autograd::PyAnomalyMetadata::ANOMALY_TRACE_KEY,
             &traceback) < 0) {
       PyErr_Clear();
+      PyErr_Restore(exc_type, exc_value, exc_tb);
       return result;
     }
 
     if (!traceback || !PyList_Check(traceback)) {
       Py_XDECREF(traceback);
+      PyErr_Restore(exc_type, exc_value, exc_tb);
       return result;
     }
 
@@ -192,6 +207,7 @@ struct PythonTraceback : public CapturedTraceback::Python {
     }
 
     Py_DECREF(traceback);
+    PyErr_Restore(exc_type, exc_value, exc_tb);
     return result;
   }
 };
diff --git a/torch/csrc/profiler/python/init.cpp b/torch/csrc/profiler/python/init.cpp
index d42af26c6b6a0..68e033449bb22 100644
--- a/torch/csrc/profiler/python/init.cpp
+++ b/torch/csrc/profiler/python/init.cpp
@@ -242,6 +242,10 @@ PyObject* RecordFunctionFast_enter(PyObject* selfGeneric, PyObject* unused) {
 
     // parse through kwargs if they exist
     if (self->keyword_values != nullptr && profiler_need_input) {
+      const auto kw_size = PyDict_Size(self->keyword_values);
+      if (kw_size > 0) {
+        kwargs.reserve(static_cast<size_t>(kw_size));
+      }
       Py_ssize_t pos = 0;
       PyObject *key = nullptr, *value = nullptr;
       while (PyDict_Next(self->keyword_values, &pos, &key, &value)) {
@@ -347,7 +351,8 @@ void initPythonBindings(PyObject* module) {
       .value("KINETO_GPU_FALLBACK", ProfilerState::KINETO_GPU_FALLBACK)
       .value(
           "KINETO_PRIVATEUSE1_FALLBACK",
-          ProfilerState::KINETO_PRIVATEUSE1_FALLBACK);
+          ProfilerState::KINETO_PRIVATEUSE1_FALLBACK)
+      .value("KINETO_PRIVATEUSE1", ProfilerState::KINETO_PRIVATEUSE1);
 
   py::enum_<ActiveProfilerType>(m, "ActiveProfilerType")
       .value("NONE", ActiveProfilerType::NONE)
diff --git a/torch/csrc/profiler/standalone/execution_trace_observer.cpp b/torch/csrc/profiler/standalone/execution_trace_observer.cpp
index 3df77bb1a8bff..379c065b22b07 100644
--- a/torch/csrc/profiler/standalone/execution_trace_observer.cpp
+++ b/torch/csrc/profiler/standalone/execution_trace_observer.cpp
@@ -48,6 +48,7 @@ constexpr auto kETGlobalRankStride = "global_rank_stride";
 constexpr auto kETGroupSize = "pg_size";
 constexpr auto kETProcessGroupName = "pg_name";
 constexpr auto kETProcessGroupDesc = "pg_desc";
+constexpr auto kETIsAsynchronizedOp = "is_asynchronized_op";
 #endif // USE_DISTRIBUTED
 
 namespace torch::profiler::impl {
@@ -485,7 +486,10 @@ convertIValue(
         itemsize,
         device_str);
     return std::make_tuple(
-        tensor_shape, tensor_stride, tensor_type, tensor_value);
+        std::move(tensor_shape),
+        std::move(tensor_stride),
+        std::move(tensor_type),
+        std::move(tensor_value));
   } else if (val.isTuple()) {
     const auto& val_tuple = val.toTupleRef().elements();
     size_t tuple_size = val_tuple.size();
@@ -493,6 +497,10 @@ convertIValue(
     std::vector<std::string> stride_array;
     std::vector<std::string> type_array;
     std::vector<std::string> value_array;
+    shape_array.reserve(tuple_size);
+    stride_array.reserve(tuple_size);
+    type_array.reserve(tuple_size);
+    value_array.reserve(tuple_size);
     for (const auto j : c10::irange(tuple_size)) {
       auto tuple = convertIValue(
           ob,
@@ -504,17 +512,17 @@ convertIValue(
           val_tuple[j],
           false,
           maxArrayLen);
-      shape_array.push_back(std::get<0>(tuple));
-      stride_array.push_back(std::get<1>(tuple));
-      type_array.push_back(std::get<2>(tuple));
-      value_array.push_back(std::get<3>(tuple));
+      shape_array.push_back(std::move(std::get<0>(tuple)));
+      stride_array.push_back(std::move(std::get<1>(tuple)));
+      type_array.push_back(std::move(std::get<2>(tuple)));
+      value_array.push_back(std::move(std::get<3>(tuple)));
     }
     type = type + vectorToString(type_array);
     std::string tensor_type = baseType ? fmt::format("\"{}\"", type) : type;
     return std::make_tuple(
         vectorToString(shape_array),
         vectorToString(stride_array),
-        tensor_type,
+        std::move(tensor_type),
         vectorToString(value_array));
   } else if (val.isList()) {
     const auto& val_list = val.toList();
@@ -523,6 +531,11 @@ convertIValue(
     std::vector<std::string> stride_array;
     std::vector<std::string> type_array;
     std::vector<std::string> value_array;
+    const size_t effective_list_size = std::min(list_size, maxArrayLen + 1);
+    shape_array.reserve(effective_list_size);
+    stride_array.reserve(effective_list_size);
+    type_array.reserve(effective_list_size);
+    value_array.reserve(effective_list_size);
     for (const auto j : c10::irange(list_size)) {
       auto tuple = convertIValue(
           ob,
@@ -534,10 +547,10 @@ convertIValue(
           val_list.get(j),
           false,
           maxArrayLen);
-      shape_array.push_back(std::get<0>(tuple));
-      stride_array.push_back(std::get<1>(tuple));
-      type_array.push_back(std::get<2>(tuple));
-      value_array.push_back(std::get<3>(tuple));
+      shape_array.push_back(std::move(std::get<0>(tuple)));
+      stride_array.push_back(std::move(std::get<1>(tuple)));
+      type_array.push_back(std::move(std::get<2>(tuple)));
+      value_array.push_back(std::move(std::get<3>(tuple)));
       if (j >= maxArrayLen) {
         LOG(WARNING) << "list size=" << val_list.size()
                      << " exceeded maxArrayLen=" << maxArrayLen;
@@ -549,7 +562,7 @@ convertIValue(
     return std::make_tuple(
         vectorToString(shape_array),
         vectorToString(stride_array),
-        tensor_type,
+        std::move(tensor_type),
         vectorToString(value_array));
   } else {
     std::string tensor_shape = "[]";
@@ -558,7 +571,10 @@ convertIValue(
     std::string tensor_value = getScalarValue(val);
 
     return std::make_tuple(
-        tensor_shape, tensor_stride, tensor_type, tensor_value);
+        std::move(tensor_shape),
+        std::move(tensor_stride),
+        std::move(tensor_type),
+        std::move(tensor_value));
   }
 }
 
@@ -667,6 +683,8 @@ inline std::string getCommsNodeAttrs(const RecordFunction& fn) { // NOLINT
 
   addAttr(kGroupSize, kETGroupSize, "uint64");
 
+  addAttr(kIsAsynchronizedOp, kETIsAsynchronizedOp, "string");
+
 #endif // USE_DISTRIBUTED
 
   // XXX consider using as string stream?
diff --git a/torch/csrc/profiler/standalone/privateuse1_profiler.cpp b/torch/csrc/profiler/standalone/privateuse1_profiler.cpp
new file mode 100644
index 0000000000000..3bd9c2c38ea3d
--- /dev/null
+++ b/torch/csrc/profiler/standalone/privateuse1_profiler.cpp
@@ -0,0 +1,73 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#ifdef USE_KINETO
+
+#include <c10/util/Exception.h>
+#include <torch/csrc/profiler/standalone/privateuse1_profiler.h>
+
+#include <libkineto.h>
+
+namespace torch::profiler::impl {
+
+PrivateUse1ProfilerRegistry& PrivateUse1ProfilerRegistry::instance() {
+  static PrivateUse1ProfilerRegistry registry;
+  return registry;
+}
+
+void PrivateUse1ProfilerRegistry::registerFactory(
+    PrivateUse1ProfilerFactory factory) {
+  std::lock_guard<std::mutex> lock(mutex_);
+
+  if (factory_) {
+    TORCH_WARN("PrivateUse1 profiler factory already registered, overwriting");
+  }
+
+  factory_ = std::move(factory);
+
+  // If Kineto was already initialized, register immediately
+  if (kineto_initialized_ && !registered_with_kineto_) {
+    registerWithKineto();
+  }
+}
+
+bool PrivateUse1ProfilerRegistry::hasFactory() const {
+  std::lock_guard<std::mutex> lock(mutex_);
+  return factory_ != nullptr;
+}
+
+bool PrivateUse1ProfilerRegistry::isRegisteredWithKineto() const {
+  std::lock_guard<std::mutex> lock(mutex_);
+  return registered_with_kineto_;
+}
+
+void PrivateUse1ProfilerRegistry::registerWithKineto() {
+  // Note: Caller must hold mutex_
+  if (!factory_ || registered_with_kineto_) {
+    return;
+  }
+
+  // Register the factory with Kineto's activity profiler
+  // Kineto will call the factory to create the profiler instance
+  libkineto::api().registerProfilerFactory(factory_);
+  registered_with_kineto_ = true;
+}
+
+void PrivateUse1ProfilerRegistry::onKinetoInit() {
+  std::lock_guard<std::mutex> lock(mutex_);
+  kineto_initialized_ = true;
+
+  // If a factory was registered before Kineto init, register it now
+  if (factory_ && !registered_with_kineto_) {
+    registerWithKineto();
+  }
+}
+
+} // namespace torch::profiler::impl
+
+#endif // USE_KINETO
diff --git a/torch/csrc/profiler/standalone/privateuse1_profiler.h b/torch/csrc/profiler/standalone/privateuse1_profiler.h
new file mode 100644
index 0000000000000..63e2d9bf982c2
--- /dev/null
+++ b/torch/csrc/profiler/standalone/privateuse1_profiler.h
@@ -0,0 +1,113 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#pragma once
+
+#ifdef USE_KINETO
+
+#include <functional>
+#include <memory>
+#include <mutex>
+#include <type_traits>
+
+#include <torch/csrc/Export.h>
+
+#include <IActivityProfiler.h>
+
+namespace torch::profiler::impl {
+
+// Factory function type that creates an IActivityProfiler instance
+using PrivateUse1ProfilerFactory =
+    std::function<std::unique_ptr<libkineto::IActivityProfiler>()>;
+
+// Registry for PrivateUse1 activity profiler factories.
+//
+// This registry allows custom accelerator backends to register their
+// IActivityProfiler implementation with Kineto, enabling full profiling
+// integration without modifying Kineto code.
+//
+// Usage:
+//   1. Backend implements libkineto::IActivityProfiler
+//   2. Backend uses REGISTER_PRIVATEUSE1_PROFILER macro to register
+//   3. PyTorch forwards the factory to Kineto during initialization
+//
+// Example:
+//   class MyAcceleratorProfiler : public libkineto::IActivityProfiler {
+//     const std::string& name() const override { return name_; }
+//     const std::set<libkineto::ActivityType>& availableActivities() const
+//     override; std::unique_ptr<libkineto::IActivityProfilerSession>
+//     configure(...) override;
+//   private:
+//     std::string name_{"my_accelerator"};
+//   };
+//
+//   REGISTER_PRIVATEUSE1_PROFILER(MyAcceleratorProfiler)
+//
+class TORCH_API PrivateUse1ProfilerRegistry {
+ public:
+  static PrivateUse1ProfilerRegistry& instance();
+
+  // Register a factory function for creating the PrivateUse1 profiler.
+  // This should be called during static initialization.
+  void registerFactory(PrivateUse1ProfilerFactory factory);
+
+  // Check if a factory has been registered.
+  bool hasFactory() const;
+
+  // Check if the factory has been registered with Kineto.
+  // Useful for testing to verify the registration logic.
+  bool isRegisteredWithKineto() const;
+
+  // Register the factory with Kineto's activity profiler.
+  // This is called internally when Kineto is ready.
+  // Safe to call multiple times - will only register once.
+  void registerWithKineto();
+
+  // Mark that Kineto has been initialized.
+  // If a factory was registered before Kineto init, it will be forwarded.
+  void onKinetoInit();
+
+ private:
+  PrivateUse1ProfilerRegistry() = default;
+
+  mutable std::mutex mutex_;
+  PrivateUse1ProfilerFactory factory_;
+  bool registered_with_kineto_ = false;
+  bool kineto_initialized_ = false;
+};
+
+// Helper struct for static registration via macro.
+// Enforces at compile-time that ProfilerClass inherits from
+// libkineto::IActivityProfiler.
+template <typename ProfilerClass>
+struct RegisterPrivateUse1Profiler {
+  static_assert(
+      std::is_base_of_v<libkineto::IActivityProfiler, ProfilerClass>,
+      "ProfilerClass must inherit from libkineto::IActivityProfiler. "
+      "Please ensure your profiler class implements the IActivityProfiler interface.");
+
+  RegisterPrivateUse1Profiler() {
+    PrivateUse1ProfilerRegistry::instance().registerFactory(
+        []() -> std::unique_ptr<libkineto::IActivityProfiler> {
+          return std::make_unique<ProfilerClass>();
+        });
+  }
+};
+
+// Macro for registering a PrivateUse1 activity profiler.
+// The profiler class must implement libkineto::IActivityProfiler.
+//
+// Usage:
+//   REGISTER_PRIVATEUSE1_PROFILER(MyAcceleratorProfiler)
+#define REGISTER_PRIVATEUSE1_PROFILER(ProfilerClass)                         \
+  static ::torch::profiler::impl::RegisterPrivateUse1Profiler<ProfilerClass> \
+      privateuse1_profiler_register_##ProfilerClass
+
+} // namespace torch::profiler::impl
+
+#endif // USE_KINETO
diff --git a/torch/csrc/profiler/unwind/action.h b/torch/csrc/profiler/unwind/action.h
index 5a982cfd046a0..7dc6c6c0b8166 100644
--- a/torch/csrc/profiler/unwind/action.h
+++ b/torch/csrc/profiler/unwind/action.h
@@ -11,7 +11,8 @@ enum {
   A_REG_PLUS_DATA_DEREF = 0x3 // exp = *(REG[reg] + data0)
 };
 
-// register numbers in dwarf info
+// DWARF register numbers — architecture-specific
+#if defined(__x86_64__)
 enum {
   D_UNDEFINED = -1,
   D_RBP = 6,
@@ -19,6 +20,28 @@ enum {
   D_RIP = 16,
   D_REG_SIZE = 17,
 };
+static constexpr int D_FRAME_PTR = D_RBP;
+static constexpr int D_STACK_PTR = D_RSP;
+static constexpr int D_RET_ADDR = D_RIP;
+static constexpr int D_EXPECTED_RA_REG = 16;
+#elif defined(__aarch64__)
+enum {
+  D_UNDEFINED = -1,
+  D_FP = 29,
+  D_LR = 30,
+  D_SP = 31,
+  D_REG_SIZE = 32,
+};
+static constexpr int D_FRAME_PTR = D_FP;
+static constexpr int D_STACK_PTR = D_SP;
+static constexpr int D_RET_ADDR = D_LR;
+static constexpr int D_EXPECTED_RA_REG = 30;
+#else
+enum {
+  D_UNDEFINED = -1,
+  D_REG_SIZE = 1,
+};
+#endif
 
 struct Action {
   uint8_t kind = A_UNDEFINED;
diff --git a/torch/csrc/profiler/unwind/dwarf_enums.h b/torch/csrc/profiler/unwind/dwarf_enums.h
index 91af24b34e1f9..3effefd1f0083 100644
--- a/torch/csrc/profiler/unwind/dwarf_enums.h
+++ b/torch/csrc/profiler/unwind/dwarf_enums.h
@@ -30,6 +30,7 @@ enum {
   DW_CFA_advance_loc1 = 0x02,
   DW_CFA_advance_loc2 = 0x03,
   DW_CFA_advance_loc4 = 0x04,
+  DW_CFA_offset_extended = 0x05,
   DW_CFA_restore_extended = 0x06,
   DW_CFA_undefined = 0x07,
   DW_CFA_register = 0x09,
diff --git a/torch/csrc/profiler/unwind/fde.h b/torch/csrc/profiler/unwind/fde.h
index ffb06b5ab1f46..70b829a0f8023 100644
--- a/torch/csrc/profiler/unwind/fde.h
+++ b/torch/csrc/profiler/unwind/fde.h
@@ -63,8 +63,10 @@ struct FDE {
     } else {
       ra_register_ = static_cast<int64_t>(LC.readULEB128());
     }
-    // we assume this in the state
-    TORCH_INTERNAL_ASSERT(ra_register_ == 16, "unexpected number of registers");
+    TORCH_INTERNAL_ASSERT(
+        ra_register_ == D_EXPECTED_RA_REG,
+        "unexpected ra register: ",
+        ra_register_);
     if (augmentation_string_ && *augmentation_string_ == 'z') {
       augmentation_length_ = static_cast<int64_t>(LC.readULEB128());
       Lexer A(LC.loc());
@@ -271,6 +273,11 @@ struct FDE {
             auto delta = L.read<uint32_t>();
             return advance_loc(delta);
           }
+          case DW_CFA_offset_extended: {
+            auto reg = L.readULEB128();
+            auto off = L.readULEB128();
+            return offset(reg, off);
+          }
           case DW_CFA_restore_extended: {
             auto reg = L.readULEB128();
             return restore(reg);
diff --git a/torch/csrc/profiler/unwind/unwind.cpp b/torch/csrc/profiler/unwind/unwind.cpp
index bdb610627c897..b1696e34f24ce 100644
--- a/torch/csrc/profiler/unwind/unwind.cpp
+++ b/torch/csrc/profiler/unwind/unwind.cpp
@@ -3,32 +3,32 @@
 #include <c10/util/env.h>
 #include <torch/csrc/profiler/unwind/unwind.h>
 
-#if !defined(__linux__) || !defined(__x86_64__) || !defined(__has_include) || \
-    !__has_include("ext/stdio_filebuf.h")
+#if !defined(__linux__) || !(defined(__x86_64__) || defined(__aarch64__)) || \
+    !defined(__has_include) || !__has_include("ext/stdio_filebuf.h")
 namespace torch::unwind {
 std::vector<void*> unwind() {
   TORCH_WARN_ONCE(
-      "record_context_cpp is not support on non-linux non-x86_64 platforms");
+      "record_context_cpp is not supported on this platform (requires linux x86_64 or aarch64)");
   return {};
 }
 
 std::optional<std::pair<std::string, uint64_t>> libraryFor(void* addr) {
   TORCH_WARN_ONCE(
-      "record_context_cpp is not support on non-linux non-x86_64 platforms");
+      "record_context_cpp is not supported on this platform (requires linux x86_64 or aarch64)");
   return {};
 }
 
 #ifndef FBCODE_CAFFE2
 std::vector<Frame> symbolize(const std::vector<void*>& frames, Mode mode) {
   TORCH_WARN_ONCE(
-      "record_context_cpp is not support on non-linux non-x86_64 platforms");
+      "record_context_cpp is not supported on this platform (requires linux x86_64 or aarch64)");
   return {};
 }
 #endif
 
 Stats stats() {
   TORCH_WARN_ONCE(
-      "record_context_cpp is not support on non-linux non-x86_64 platforms");
+      "record_context_cpp is not supported on this platform (requires linux x86_64 or aarch64)");
   return {};
 }
 
@@ -41,8 +41,10 @@ Stats stats() {
 #include <elf.h>
 #include <link.h>
 #include <linux/limits.h>
+#include <pthread.h>
 #include <algorithm>
 #include <climits>
+#include <cstring>
 #include <vector>
 
 #include <c10/util/irange.h>
@@ -53,7 +55,14 @@ Stats stats() {
 #include <torch/csrc/profiler/unwind/unwinder.h>
 #include <shared_mutex>
 
+#if defined(__aarch64__)
+extern "C" void unwind_c(
+    std::vector<void*>* result,
+    uintptr_t fp,
+    uintptr_t lr);
+#else
 extern "C" void unwind_c(std::vector<void*>* result, int64_t rsp, int64_t rbp);
+#endif
 extern "C" void unwind_entry(std::vector<void*>* result);
 
 namespace torch::unwind {
@@ -98,7 +107,8 @@ struct LibraryInfo {
     void* fde_data = eh_frame_hdr_.entryForAddr(addr);
     FDE fde(fde_data, name().c_str(), load_bias());
     TableState state = fde.readUpTo(addr);
-    return Unwinder(state.cfa, state.registers[D_RIP], state.registers[D_RBP]);
+    return Unwinder(
+        state.cfa, state.registers[D_RET_ADDR], state.registers[D_FRAME_PTR]);
   }
   const std::string& name() const {
     return name_;
@@ -501,6 +511,111 @@ Stats stats() {
 
 } // namespace torch::unwind
 
+#if defined(__aarch64__)
+// aarch64 uses frame-pointer chain walking instead of DWARF unwinding.
+// Each frame has: *(FP) = caller's FP, *(FP+8) = saved LR (return address).
+// This is simpler and avoids issues with tail calls producing stale x30
+// values in DWARF-based unwinding.  GCC/Clang on aarch64 emit frame
+// pointers by default even at -O2.
+//
+// External libraries (CPython, libc, CUDA runtime) may be built without
+// frame pointers, making x29 an arbitrary callee-saved value.  We obtain
+// the current thread's stack bounds and reject any fp outside that range
+// to avoid dereferencing garbage pointers.
+//
+// No cache_mutex_ needed: frame-pointer walking reads only the stack,
+// unlike the x86 path which queries the DWARF FDE cache.
+
+// Stack bounds are essentially immutable over a thread's lifetime (base and
+// size are fixed at pthread_create for spawned threads).  On glibc aarch64
+// the main thread is the pathological case: pthread_getattr_np() parses
+// /proc/self/maps on every call because the main-thread stack isn't recorded
+// in TLS at pthread_create.  Cache once per thread to avoid that parse on
+// every unwind.  This mirrors the process-global caches the x86 unwinder
+// uses for /proc-derived data (library list, exe path); here the data is
+// per-thread, so thread_local is the right scope.
+//
+// Edge case: main thread's stack can grow up to RLIMIT_STACK, and this cache
+// freezes the bounds at first observation.  If the stack later grows past
+// cached lo, unwind_c will terminate early on those frames rather than
+// follow them - truncated backtrace, not incorrect.  In practice main-thread
+// stack reaches steady-state depth during init, long before heavy unwinding.
+namespace {
+struct StackBounds {
+  uintptr_t lo = 0;
+  uintptr_t hi = 0;
+  bool initialized = false;
+};
+thread_local StackBounds tls_stack_bounds;
+} // namespace
+
+static bool get_stack_bounds(uintptr_t& lo, uintptr_t& hi) {
+  auto& b = tls_stack_bounds;
+  if (!b.initialized) {
+    pthread_attr_t attr;
+    if (pthread_getattr_np(pthread_self(), &attr) != 0) {
+      return false;
+    }
+    void* base = nullptr;
+    size_t size = 0;
+    int rc = pthread_attr_getstack(&attr, &base, &size);
+    pthread_attr_destroy(&attr);
+    if (rc != 0) {
+      return false;
+    }
+    b.lo = reinterpret_cast<uintptr_t>(base);
+    b.hi = b.lo + size;
+    b.initialized = true;
+  }
+  lo = b.lo;
+  hi = b.hi;
+  return true;
+}
+
+extern "C" C10_USED void unwind_c(
+    std::vector<void*>* result,
+    uintptr_t fp,
+    uintptr_t lr) {
+  // NOLINTNEXTLINE(performance-no-int-to-ptr)
+  result->push_back((void*)lr);
+
+  uintptr_t stack_lo = 0, stack_hi = 0;
+  if (!get_stack_bounds(stack_lo, stack_hi)) {
+    return;
+  }
+
+  constexpr int kMaxFrames = 4096;
+  int depth = 0;
+  while (fp != 0 && (fp & 0xF) == 0 && depth++ < kMaxFrames) {
+    if (fp < stack_lo || fp + 16 > stack_hi) {
+      break;
+    }
+    uintptr_t saved_lr;
+    std::memcpy(
+        &saved_lr, reinterpret_cast<const void*>(fp + 8), sizeof(saved_lr));
+    if (saved_lr == 0) {
+      break;
+    }
+    // NOLINTNEXTLINE(performance-no-int-to-ptr)
+    result->push_back((void*)saved_lr);
+    uintptr_t next_fp;
+    std::memcpy(&next_fp, reinterpret_cast<const void*>(fp), sizeof(next_fp));
+    if (next_fp <= fp) {
+      break;
+    }
+    fp = next_fp;
+  }
+}
+
+// x0 already holds the result pointer.
+// Pass FP (x29) and LR (x30), then tail-call unwind_c.
+__asm__(
+    ".global unwind_entry\n"
+    "unwind_entry:\n"
+    "mov x1, x29\n"
+    "mov x2, x30\n"
+    "b unwind_c\n");
+#else
 extern "C" C10_USED void unwind_c(
     std::vector<void*>* result,
     int64_t rsp,
@@ -508,17 +623,17 @@ extern "C" C10_USED void unwind_c(
   std::shared_lock lock(torch::unwind::cache_mutex_);
   torch::unwind::UnwindState state{};
   // NOLINTNEXTLINE(performance-no-int-to-ptr)
-  state.rip = *(int64_t*)rsp;
+  state.pc = *(int64_t*)rsp;
   // +8 because we saved rsp after the return address was already pushed
   // to the stack
-  state.rsp = rsp + 8;
-  state.rbp = rbp;
+  state.sp = rsp + 8;
+  state.fp = rbp;
   torch::unwind::unwind_cache.checkRefresh(lock);
-  while (true) { // unwind for _start sets rip as being undefined
+  while (true) {
     // NOLINTNEXTLINE(performance-no-int-to-ptr)
-    result->push_back((void*)state.rip);
+    result->push_back((void*)state.pc);
     const torch::unwind::Unwinder& uw =
-        torch::unwind::unwind_cache.unwinderFor(state.rip, lock);
+        torch::unwind::unwind_cache.unwinderFor(state.pc, lock);
     if (uw.terminator()) {
       if (uw.isUnknown()) {
         result->push_back(nullptr);
@@ -529,8 +644,7 @@ extern "C" C10_USED void unwind_c(
   }
 }
 
-// calling convention puts the first three pointer/int64_t arguments in
-// rdi rsi rdx (all caller-saved)
+// x86-64 calling convention: rdi rsi rdx (all caller-saved)
 // rdi already holds the pointer to the result vector
 // we add arguments for current rsp and rbp and then tail call
 // into unwind_c
@@ -540,5 +654,6 @@ __asm__(
     "mov %rsp, %rsi;\n"
     "mov %rbp, %rdx;\n"
     "jmp unwind_c;\n");
+#endif
 
 #endif
diff --git a/torch/csrc/profiler/unwind/unwinder.h b/torch/csrc/profiler/unwind/unwinder.h
index d673f47af8db2..3836203688530 100644
--- a/torch/csrc/profiler/unwind/unwinder.h
+++ b/torch/csrc/profiler/unwind/unwinder.h
@@ -6,30 +6,32 @@
 
 namespace torch::unwind {
 
+// Architecture-neutral names: pc (program counter / return address),
+// fp (frame pointer: x86 RBP, aarch64 x29), sp (stack pointer).
 struct UnwindState {
-  int64_t rip, rbp, rsp;
+  int64_t pc, fp, sp;
 };
 
 struct Unwinder {
-  Unwinder(Action rsp, Action rip, Action rbp)
-      : kind_(rip.kind == A_UNDEFINED ? END : STANDARD),
-        reg_(rsp.reg),
-        off_(rsp.data),
-        rip_off_(rip.data),
-        rbp_off_(
-            rbp.kind == A_UNDEFINED ? std::numeric_limits<int64_t>::max()
-                                    : rbp.data),
-        deref_(rsp.kind == A_REG_PLUS_DATA_DEREF) {
-    check(rsp.reg == D_RSP || rsp.reg == D_RBP);
-    check(rip.kind == A_UNDEFINED || rip.kind == A_LOAD_CFA_OFFSET);
-    if (rsp.kind == A_REG_PLUS_DATA) {
-      check(rbp.kind == A_LOAD_CFA_OFFSET || rbp.kind == A_UNDEFINED);
-    } else if (rsp.kind == A_REG_PLUS_DATA_DEREF) {
-      if (rbp.kind == A_REG_PLUS_DATA_DEREF) {
-        check(rbp.reg == rsp.reg);
-        rbp_off_ -= rsp.data;
+  Unwinder(Action cfa, Action ret, Action fp)
+      : kind_(ret.kind == A_UNDEFINED ? END : STANDARD),
+        reg_(cfa.reg),
+        off_(cfa.data),
+        ret_off_(ret.data),
+        fp_off_(
+            fp.kind == A_UNDEFINED ? std::numeric_limits<int64_t>::max()
+                                   : fp.data),
+        deref_(cfa.kind == A_REG_PLUS_DATA_DEREF) {
+    check(cfa.reg == D_STACK_PTR || cfa.reg == D_FRAME_PTR);
+    check(ret.kind == A_UNDEFINED || ret.kind == A_LOAD_CFA_OFFSET);
+    if (cfa.kind == A_REG_PLUS_DATA) {
+      check(fp.kind == A_LOAD_CFA_OFFSET || fp.kind == A_UNDEFINED);
+    } else if (cfa.kind == A_REG_PLUS_DATA_DEREF) {
+      if (fp.kind == A_REG_PLUS_DATA_DEREF) {
+        check(fp.reg == cfa.reg);
+        fp_off_ -= cfa.data;
       } else {
-        check(rbp.kind == A_UNDEFINED);
+        check(fp.kind == A_UNDEFINED);
       }
     } else {
       check(false);
@@ -53,28 +55,28 @@ struct Unwinder {
   }
   UnwindState run(const UnwindState& cur) const {
     UnwindState r = cur;
-    r.rsp = (reg_ == D_RSP ? cur.rsp : cur.rbp) + off_;
-    r.rbp = rbp_off_ == std::numeric_limits<int64_t>::max()
-        ? cur.rbp
+    r.sp = (reg_ == D_STACK_PTR ? cur.sp : cur.fp) + off_;
+    r.fp = fp_off_ == std::numeric_limits<int64_t>::max()
+        ? cur.fp
         // NOLINTNEXTLINE(performance-no-int-to-ptr)
-        : *(int64_t*)(r.rsp + rbp_off_);
+        : *(int64_t*)(r.sp + fp_off_);
     if (deref_) {
       // NOLINTNEXTLINE(performance-no-int-to-ptr)
-      r.rsp = *(int64_t*)r.rsp;
+      r.sp = *(int64_t*)r.sp;
     }
     // NOLINTNEXTLINE(performance-no-int-to-ptr)
-    r.rip = *(int64_t*)(r.rsp + rip_off_);
+    r.pc = *(int64_t*)(r.sp + ret_off_);
 
     return r;
   }
 
  private:
-  Unwinder() : kind_(UNKNOWN), reg_(0), off_(0), rip_off_(0), rbp_off_(0) {}
+  Unwinder() : kind_(UNKNOWN), reg_(0), off_(0), ret_off_(0), fp_off_(0) {}
   enum Kind { STANDARD, END, UNKNOWN } kind_;
   uint32_t reg_;
   int64_t off_;
-  int64_t rip_off_;
-  int64_t rbp_off_;
+  int64_t ret_off_;
+  int64_t fp_off_;
   bool deref_{false};
 };
 
diff --git a/torch/csrc/profiler/util.cpp b/torch/csrc/profiler/util.cpp
index cdec85182ce1e..541b6b79c56ce 100644
--- a/torch/csrc/profiler/util.cpp
+++ b/torch/csrc/profiler/util.cpp
@@ -509,8 +509,16 @@ std::unordered_map<std::string, std::string> saveNcclMeta(
         map.emplace(kP2pSrc, std::to_string(groupRanks[rank]));
       }
     }
+
+    auto seqNum = debugInfo->getSequenceNumber();
+    if (seqNum >= 0) {
+      map.emplace(kSeqNum, std::to_string(seqNum));
+    }
   }
 
+  map.emplace(
+      kIsAsynchronizedOp, std::to_string(debugInfo->isAsynchronizedOp()));
+
   if (get_record_tensor_addrs_enabled()) {
     std::vector<std::string> addressList;
     if (config.introspectInputs) {
diff --git a/torch/csrc/profiler/util.h b/torch/csrc/profiler/util.h
index f2ae57fa0e591..eed931cbf5a9c 100644
--- a/torch/csrc/profiler/util.h
+++ b/torch/csrc/profiler/util.h
@@ -201,8 +201,10 @@ constexpr auto kGroupRanks = "Process Group Ranks";
 constexpr auto kRank = "Rank";
 constexpr auto kP2pSrc = "Src Rank";
 constexpr auto kP2pDst = "Dst Rank";
+constexpr auto kSeqNum = "Seq";
 constexpr auto kInTensorsStart = "Input Tensors start";
 constexpr auto kOutTensorsStart = "Output Tensors start";
+constexpr auto kIsAsynchronizedOp = "Is asynchronized op";
 #endif // USE_DISTRIBUTED
 
 } // namespace torch::profiler::impl
diff --git a/torch/csrc/shim_common.cpp b/torch/csrc/shim_common.cpp
index d6c3cd2b41e0d..4b5233955a18f 100644
--- a/torch/csrc/shim_common.cpp
+++ b/torch/csrc/shim_common.cpp
@@ -667,7 +667,8 @@ AOTI_TORCH_EXPORT AOTITorchError torch_from_blob(
     int32_t layout,
     const uint8_t* opaque_metadata,
     int64_t opaque_metadata_size,
-    void (*deleter)(void*)) {
+    void (*deleter_callback)(void* data, void* ctx),
+    void* deleter_ctx) {
   AOTI_TORCH_CONVERT_EXCEPTION_TO_ERROR_CODE({
     c10::IntArrayRef sizes(sizes_ptr, ndim);
     c10::IntArrayRef strides(strides_ptr, ndim);
@@ -676,11 +677,14 @@ AOTI_TORCH_EXPORT AOTITorchError torch_from_blob(
         static_cast<c10::ScalarType>(dtype));
     at::Tensor tensor;
     if (data != nullptr) {
-      if (deleter != nullptr) {
+      if (deleter_callback != nullptr) {
+        auto wrapped_deleter = [deleter_callback, deleter_ctx](void* data) {
+          deleter_callback(data, deleter_ctx);
+        };
         tensor = at::for_blob(data, sizes)
                      .strides(strides)
                      .storage_offset(storage_offset)
-                     .deleter(deleter)
+                     .deleter(wrapped_deleter)
                      .options(options)
                      .make_tensor();
       } else {
@@ -696,3 +700,45 @@ AOTI_TORCH_EXPORT AOTITorchError torch_from_blob(
     *ret_new_tensor = torch::aot_inductor::new_tensor_handle(std::move(tensor));
   });
 }
+
+// Tag getter shims — see torch/csrc/stable/c/shim.h
+#define TORCH_TAG_IMPL(name, enum_val) \
+  int32_t torch_tag_##name() {         \
+    return (int32_t)at::Tag::enum_val; \
+  }
+
+TORCH_TAG_IMPL(core, core)
+TORCH_TAG_IMPL(cudagraph_unsafe, cudagraph_unsafe)
+TORCH_TAG_IMPL(data_dependent_output, data_dependent_output)
+TORCH_TAG_IMPL(dynamic_output_shape, dynamic_output_shape)
+TORCH_TAG_IMPL(flexible_layout, flexible_layout)
+TORCH_TAG_IMPL(generated, generated)
+TORCH_TAG_IMPL(inplace_view, inplace_view)
+TORCH_TAG_IMPL(maybe_aliasing_or_mutating, maybe_aliasing_or_mutating)
+TORCH_TAG_IMPL(needs_contiguous_strides, needs_contiguous_strides)
+TORCH_TAG_IMPL(needs_exact_strides, needs_exact_strides)
+TORCH_TAG_IMPL(needs_fixed_stride_order, needs_fixed_stride_order)
+TORCH_TAG_IMPL(nondeterministic_bitwise, nondeterministic_bitwise)
+TORCH_TAG_IMPL(nondeterministic_seeded, nondeterministic_seeded)
+TORCH_TAG_IMPL(out_variant, out_variant)
+TORCH_TAG_IMPL(pointwise, pointwise)
+TORCH_TAG_IMPL(pt2_compliant_tag, pt2_compliant_tag)
+TORCH_TAG_IMPL(reduction, reduction)
+TORCH_TAG_IMPL(view_copy, view_copy)
+#undef TORCH_TAG_IMPL
+
+AOTI_TORCH_EXPORT AOTITorchError torch_library_def_with_tags(
+    TorchLibraryHandle self,
+    const char* schema,
+    const int32_t* tags,
+    int32_t num_tags) {
+  AOTI_TORCH_CONVERT_EXCEPTION_TO_ERROR_CODE({
+    std::vector<at::Tag> tag_vec;
+    tag_vec.reserve(num_tags);
+    for (int32_t i = 0; i < num_tags; i++) {
+      tag_vec.push_back(static_cast<at::Tag>(tags[i]));
+    }
+    reinterpret_cast<torch::Library*>(self)->def(
+        torch::schema(schema), tag_vec, torch::_RegisterOrVerify::REGISTER);
+  });
+}
diff --git a/torch/csrc/stable/c/shim.h b/torch/csrc/stable/c/shim.h
index f8a21e6c570f6..3844a3d5cc21d 100644
--- a/torch/csrc/stable/c/shim.h
+++ b/torch/csrc/stable/c/shim.h
@@ -165,8 +165,9 @@ AOTI_TORCH_EXPORT int32_t torch_dtype_float8_e8m0fnu();
 AOTI_TORCH_EXPORT int32_t torch_dtype_float4_e2m1fn_x2();
 
 // Creates a tensor from an existing data blob with an optional deleter.
-// The deleter is called with the data pointer when the tensor's storage
-// is deallocated.
+// The deleter receives both the data pointer and a caller-supplied context
+// pointer, which allows passing capturing lambdas across the C ABI boundary
+// by heap-allocating the callable and passing it as deleter_ctx.
 AOTI_TORCH_EXPORT AOTITorchError torch_from_blob(
     void* data,
     int64_t ndim,
@@ -180,10 +181,48 @@ AOTI_TORCH_EXPORT AOTITorchError torch_from_blob(
     int32_t layout,
     const uint8_t* opaque_metadata,
     int64_t opaque_metadata_size,
-    void (*deleter)(void*));
+    void (*deleter)(void* data, void* ctx),
+    void* deleter_ctx);
 
 #endif // TORCH_FEATURE_VERSION >= TORCH_VERSION_2_11_0
 
+/**
+ * The beginning of all shims added in 2.12.0 onwards.
+ */
+#if TORCH_FEATURE_VERSION >= TORCH_VERSION_2_12_0
+
+// Tag getter functions for ABI-stable tag passing.  By hiding these behind
+// functions, the precise enum ordinal is NOT part of the ABI contract.
+AOTI_TORCH_EXPORT int32_t torch_tag_core();
+AOTI_TORCH_EXPORT int32_t torch_tag_cudagraph_unsafe();
+AOTI_TORCH_EXPORT int32_t torch_tag_data_dependent_output();
+AOTI_TORCH_EXPORT int32_t torch_tag_dynamic_output_shape();
+AOTI_TORCH_EXPORT int32_t torch_tag_flexible_layout();
+AOTI_TORCH_EXPORT int32_t torch_tag_generated();
+AOTI_TORCH_EXPORT int32_t torch_tag_inplace_view();
+AOTI_TORCH_EXPORT int32_t torch_tag_maybe_aliasing_or_mutating();
+AOTI_TORCH_EXPORT int32_t torch_tag_needs_contiguous_strides();
+AOTI_TORCH_EXPORT int32_t torch_tag_needs_exact_strides();
+AOTI_TORCH_EXPORT int32_t torch_tag_needs_fixed_stride_order();
+AOTI_TORCH_EXPORT int32_t torch_tag_nondeterministic_bitwise();
+AOTI_TORCH_EXPORT int32_t torch_tag_nondeterministic_seeded();
+AOTI_TORCH_EXPORT int32_t torch_tag_out_variant();
+AOTI_TORCH_EXPORT int32_t torch_tag_pointwise();
+AOTI_TORCH_EXPORT int32_t torch_tag_pt2_compliant_tag();
+AOTI_TORCH_EXPORT int32_t torch_tag_reduction();
+AOTI_TORCH_EXPORT int32_t torch_tag_view_copy();
+
+// Stable corollary to torch::Library method m.def() with tags.
+// Tags are passed as int32_t values obtained from torch_tag_*() getters,
+// not raw enum ordinals, so the ABI is stable across versions.
+AOTI_TORCH_EXPORT AOTITorchError torch_library_def_with_tags(
+    TorchLibraryHandle self,
+    const char* schema,
+    const int32_t* tags,
+    int32_t num_tags);
+
+#endif // TORCH_FEATURE_VERSION >= TORCH_VERSION_2_12_0
+
 #ifdef __cplusplus
 } // extern "C"
 #endif
diff --git a/torch/csrc/stable/c/shim_function_versions.txt b/torch/csrc/stable/c/shim_function_versions.txt
index 9d04283991ce9..264188ea08957 100644
--- a/torch/csrc/stable/c/shim_function_versions.txt
+++ b/torch/csrc/stable/c/shim_function_versions.txt
@@ -26,3 +26,24 @@ torch_parse_device_string: TORCH_VERSION_2_10_0
 torch_dtype_float4_e2m1fn_x2: TORCH_VERSION_2_11_0
 torch_dtype_float8_e8m0fnu: TORCH_VERSION_2_11_0
 torch_from_blob: TORCH_VERSION_2_11_0
+aoti_torch_dtype_float4_e2m1fn_x2: TORCH_VERSION_2_12_0
+aoti_torch_dtype_float8_e8m0fnu: TORCH_VERSION_2_12_0
+torch_library_def_with_tags: TORCH_VERSION_2_12_0
+torch_tag_core: TORCH_VERSION_2_12_0
+torch_tag_cudagraph_unsafe: TORCH_VERSION_2_12_0
+torch_tag_data_dependent_output: TORCH_VERSION_2_12_0
+torch_tag_dynamic_output_shape: TORCH_VERSION_2_12_0
+torch_tag_flexible_layout: TORCH_VERSION_2_12_0
+torch_tag_generated: TORCH_VERSION_2_12_0
+torch_tag_inplace_view: TORCH_VERSION_2_12_0
+torch_tag_maybe_aliasing_or_mutating: TORCH_VERSION_2_12_0
+torch_tag_needs_contiguous_strides: TORCH_VERSION_2_12_0
+torch_tag_needs_exact_strides: TORCH_VERSION_2_12_0
+torch_tag_needs_fixed_stride_order: TORCH_VERSION_2_12_0
+torch_tag_nondeterministic_bitwise: TORCH_VERSION_2_12_0
+torch_tag_nondeterministic_seeded: TORCH_VERSION_2_12_0
+torch_tag_out_variant: TORCH_VERSION_2_12_0
+torch_tag_pointwise: TORCH_VERSION_2_12_0
+torch_tag_pt2_compliant_tag: TORCH_VERSION_2_12_0
+torch_tag_reduction: TORCH_VERSION_2_12_0
+torch_tag_view_copy: TORCH_VERSION_2_12_0
diff --git a/torch/csrc/stable/library.h b/torch/csrc/stable/library.h
index 891918a19ca7e..f12f0d768acfc 100644
--- a/torch/csrc/stable/library.h
+++ b/torch/csrc/stable/library.h
@@ -5,6 +5,7 @@
 
 #include <torch/csrc/inductor/aoti_torch/c/shim.h>
 #include <torch/csrc/stable/c/shim.h>
+#include <torch/headeronly/core/enum_tag.h>
 #include <torch/headeronly/macros/Macros.h>
 #include <torch/headeronly/util/Metaprogramming.h>
 
@@ -97,6 +98,21 @@ class StableLibrary final {
     aoti_torch_library_def(lib_, schema);
     return *this;
   }
+
+#if TORCH_FEATURE_VERSION >= TORCH_VERSION_2_12_0
+  // corresponds to a limited, stable version of torch::library::def() with tags
+  StableLibrary& def(const char* schema, const std::vector<at::Tag>& tags) {
+    std::vector<int32_t> tag_ints;
+    tag_ints.reserve(tags.size());
+    for (auto t : tags) {
+      tag_ints.push_back(
+          torch::stable::detail::to<int32_t>(torch::stable::detail::from(t)));
+    }
+    torch_library_def_with_tags(
+        lib_, schema, tag_ints.data(), static_cast<int32_t>(tag_ints.size()));
+    return *this;
+  }
+#endif // TORCH_FEATURE_VERSION >= TORCH_VERSION_2_12_0
 };
 
 class StableTorchLibraryInit final {
diff --git a/torch/csrc/stable/ops.h b/torch/csrc/stable/ops.h
index 19c109404cb5b..dbede4faba49e 100644
--- a/torch/csrc/stable/ops.h
+++ b/torch/csrc/stable/ops.h
@@ -5,6 +5,7 @@
 #include <cstdint>
 #include <optional>
 #include <string>
+#include <type_traits>
 
 #include <torch/csrc/inductor/aoti_torch/generated/c_shim_aten.h>
 #include <torch/csrc/stable/c/shim.h>
@@ -722,26 +723,28 @@ inline torch::stable::Tensor from_blob(
 ///
 /// This is the same as the from_blob function above, but allows specifying a
 /// custom deleter function that will be called when the tensor's storage is
-/// deallocated.
+/// deallocated. Accepts both plain function pointers and capturing lambdas.
+///
 /// Minimum compatible version: PyTorch 2.11.
 ///
+/// @tparam F The callable type. Must be invocable with (void*).
 /// @param data Pointer to the data buffer.
 /// @param sizes The size of each dimension of the tensor.
 /// @param strides The stride for each dimension.
 /// @param device The device where the data resides.
 /// @param dtype The scalar type of the data.
-/// @param deleter Function to call when the tensor is deallocated. May be
-///                nullptr if no cleanup is needed.
+/// @param deleter Callable to invoke when the tensor is deallocated.
 /// @param storage_offset The offset into the data buffer. Defaults to 0.
 /// @param layout The memory layout. Defaults to Strided.
 /// @return A tensor backed by the provided data.
+template <class F, std::enable_if_t<std::is_invocable_v<F, void*>, int> = 0>
 inline torch::stable::Tensor from_blob(
     void* data,
     torch::headeronly::IntHeaderOnlyArrayRef sizes,
     torch::headeronly::IntHeaderOnlyArrayRef strides,
     torch::stable::Device device,
     torch::headeronly::ScalarType dtype,
-    DeleterFnPtr deleter,
+    F deleter,
     int64_t storage_offset = 0,
     torch::headeronly::Layout layout = torch::headeronly::Layout::Strided) {
   auto shim_dtype =
@@ -750,21 +753,53 @@ inline torch::stable::Tensor from_blob(
       torch::stable::detail::from(device.type()));
   auto shim_layout =
       torch::stable::detail::to<int32_t>(torch::stable::detail::from(layout));
+
   AtenTensorHandle ath;
-  TORCH_ERROR_CODE_CHECK(torch_from_blob(
-      data,
-      sizes.size(),
-      sizes.data(),
-      strides.data(),
-      storage_offset,
-      shim_dtype,
-      shim_device_type,
-      device.index(),
-      &ath,
-      shim_layout,
-      nullptr,
-      0,
-      deleter));
+  if constexpr (std::is_convertible_v<F, DeleterFnPtr>) {
+    // Simple function pointer: pass it as ctx, no heap allocation.
+    auto deleter_callback = [](void* data, void* ctx) {
+      auto fn = reinterpret_cast<DeleterFnPtr>(ctx);
+      fn(data);
+    };
+    TORCH_ERROR_CODE_CHECK(torch_from_blob(
+        data,
+        sizes.size(),
+        sizes.data(),
+        strides.data(),
+        storage_offset,
+        shim_dtype,
+        shim_device_type,
+        device.index(),
+        &ath,
+        shim_layout,
+        nullptr,
+        0,
+        deleter_callback,
+        reinterpret_cast<void*>(static_cast<DeleterFnPtr>(deleter))));
+  } else {
+    // Capturing lambda: heap-allocate and type-erase.
+    F* heap_allocated_deleter = new F(std::move(deleter));
+    auto deleter_callback = [](void* data, void* ctx) {
+      F* func = static_cast<F*>(ctx);
+      (*func)(data);
+      delete func;
+    };
+    TORCH_ERROR_CODE_CHECK(torch_from_blob(
+        data,
+        sizes.size(),
+        sizes.data(),
+        strides.data(),
+        storage_offset,
+        shim_dtype,
+        shim_device_type,
+        device.index(),
+        &ath,
+        shim_layout,
+        nullptr,
+        0,
+        deleter_callback,
+        static_cast<void*>(heap_allocated_deleter)));
+  }
   return torch::stable::Tensor(ath);
 }
 #endif // TORCH_FEATURE_VERSION >= TORCH_VERSION_2_11_0
diff --git a/torch/csrc/stable/stableivalue_conversions.h b/torch/csrc/stable/stableivalue_conversions.h
index 23d97b6b7afb1..eece32c7eb33d 100644
--- a/torch/csrc/stable/stableivalue_conversions.h
+++ b/torch/csrc/stable/stableivalue_conversions.h
@@ -8,6 +8,7 @@
 #include <torch/headeronly/core/Layout.h>
 #include <torch/headeronly/core/MemoryFormat.h>
 #include <torch/headeronly/core/ScalarType.h>
+#include <torch/headeronly/core/enum_tag.h>
 #include <torch/headeronly/macros/Macros.h>
 #include <torch/headeronly/util/Deprecated.h>
 #include <torch/headeronly/util/Exception.h>
@@ -773,6 +774,126 @@ struct ToImpl<std::string> {
 
 #endif // TORCH_FEATURE_VERSION >= TORCH_VERSION_2_10_0
 
+// =============================================================================
+// FROM/TO CONVERSIONS requiring TORCH_FEATURE_VERSION >= TORCH_VERSION_2_12_0
+// =============================================================================
+#if TORCH_FEATURE_VERSION >= TORCH_VERSION_2_12_0
+
+// Specialization for torch::headeronly::Tag => StableIValue
+// Uses shim getter functions so the integer representation is resolved at
+// runtime from libtorch, not baked in at extension compile time.
+using torch::headeronly::Tag;
+template <>
+struct FromImpl<Tag> {
+  static StableIValue call(
+      Tag val,
+      [[maybe_unused]] uint64_t extension_build_version,
+      [[maybe_unused]] bool is_internal) {
+    switch (val) {
+      case Tag::core:
+        return torch::stable::detail::from(torch_tag_core());
+      case Tag::cudagraph_unsafe:
+        return torch::stable::detail::from(torch_tag_cudagraph_unsafe());
+      case Tag::data_dependent_output:
+        return torch::stable::detail::from(torch_tag_data_dependent_output());
+      case Tag::dynamic_output_shape:
+        return torch::stable::detail::from(torch_tag_dynamic_output_shape());
+      case Tag::flexible_layout:
+        return torch::stable::detail::from(torch_tag_flexible_layout());
+      case Tag::generated:
+        return torch::stable::detail::from(torch_tag_generated());
+      case Tag::inplace_view:
+        return torch::stable::detail::from(torch_tag_inplace_view());
+      case Tag::maybe_aliasing_or_mutating:
+        return torch::stable::detail::from(
+            torch_tag_maybe_aliasing_or_mutating());
+      case Tag::needs_contiguous_strides:
+        return torch::stable::detail::from(
+            torch_tag_needs_contiguous_strides());
+      case Tag::needs_exact_strides:
+        return torch::stable::detail::from(torch_tag_needs_exact_strides());
+      case Tag::needs_fixed_stride_order:
+        return torch::stable::detail::from(
+            torch_tag_needs_fixed_stride_order());
+      case Tag::nondeterministic_bitwise:
+        return torch::stable::detail::from(
+            torch_tag_nondeterministic_bitwise());
+      case Tag::nondeterministic_seeded:
+        return torch::stable::detail::from(torch_tag_nondeterministic_seeded());
+      case Tag::out_variant:
+        return torch::stable::detail::from(torch_tag_out_variant());
+      case Tag::pointwise:
+        return torch::stable::detail::from(torch_tag_pointwise());
+      case Tag::pt2_compliant_tag:
+        return torch::stable::detail::from(torch_tag_pt2_compliant_tag());
+      case Tag::reduction:
+        return torch::stable::detail::from(torch_tag_reduction());
+      case Tag::view_copy:
+        return torch::stable::detail::from(torch_tag_view_copy());
+      default:
+        STD_TORCH_CHECK(
+            false,
+            "Not yet supported Tag, please file an issue describing your use case.");
+    }
+  }
+};
+
+// Specialization for StableIValue => torch::headeronly::Tag
+template <>
+struct ToImpl<Tag> {
+  static Tag call(
+      StableIValue val,
+      [[maybe_unused]] uint64_t extension_build_version,
+      [[maybe_unused]] bool is_internal) {
+    int32_t shim_tag = torch::stable::detail::to<int32_t>(val);
+    if (shim_tag == torch_tag_core()) {
+      return Tag::core;
+    } else if (shim_tag == torch_tag_cudagraph_unsafe()) {
+      return Tag::cudagraph_unsafe;
+    } else if (shim_tag == torch_tag_data_dependent_output()) {
+      return Tag::data_dependent_output;
+    } else if (shim_tag == torch_tag_dynamic_output_shape()) {
+      return Tag::dynamic_output_shape;
+    } else if (shim_tag == torch_tag_flexible_layout()) {
+      return Tag::flexible_layout;
+    } else if (shim_tag == torch_tag_generated()) {
+      return Tag::generated;
+    } else if (shim_tag == torch_tag_inplace_view()) {
+      return Tag::inplace_view;
+    } else if (shim_tag == torch_tag_maybe_aliasing_or_mutating()) {
+      return Tag::maybe_aliasing_or_mutating;
+    } else if (shim_tag == torch_tag_needs_contiguous_strides()) {
+      return Tag::needs_contiguous_strides;
+    } else if (shim_tag == torch_tag_needs_exact_strides()) {
+      return Tag::needs_exact_strides;
+    } else if (shim_tag == torch_tag_needs_fixed_stride_order()) {
+      return Tag::needs_fixed_stride_order;
+    } else if (shim_tag == torch_tag_nondeterministic_bitwise()) {
+      return Tag::nondeterministic_bitwise;
+    } else if (shim_tag == torch_tag_nondeterministic_seeded()) {
+      return Tag::nondeterministic_seeded;
+    } else if (shim_tag == torch_tag_out_variant()) {
+      return Tag::out_variant;
+    } else if (shim_tag == torch_tag_pointwise()) {
+      return Tag::pointwise;
+    } else if (shim_tag == torch_tag_pt2_compliant_tag()) {
+      return Tag::pt2_compliant_tag;
+    } else if (shim_tag == torch_tag_reduction()) {
+      return Tag::reduction;
+    } else if (shim_tag == torch_tag_view_copy()) {
+      return Tag::view_copy;
+    } else {
+      STD_TORCH_CHECK(
+          false,
+          "Not yet supported Tag ",
+          std::to_string(shim_tag),
+          ", please file an issue describing your use case.");
+    }
+  }
+};
+
+#endif // TORCH_FEATURE_VERSION >= TORCH_VERSION_2_12_0
+
 // =============================================================================
 //  end to helpers for converting between StableIValue and T
 // =============================================================================
diff --git a/torch/csrc/stable/version.h b/torch/csrc/stable/version.h
index 57e2e11643fec..eeb4f9152546a 100644
--- a/torch/csrc/stable/version.h
+++ b/torch/csrc/stable/version.h
@@ -28,3 +28,4 @@
 
 #define TORCH_VERSION_2_10_0 (((0ULL + 2) << 56) | ((0ULL + 10) << 48))
 #define TORCH_VERSION_2_11_0 (((0ULL + 2) << 56) | ((0ULL + 11) << 48))
+#define TORCH_VERSION_2_12_0 (((0ULL + 2) << 56) | ((0ULL + 12) << 48))
diff --git a/torch/csrc/utils/device_lazy_init.cpp b/torch/csrc/utils/device_lazy_init.cpp
index b1338fc43e913..fc5c7a47b8b82 100644
--- a/torch/csrc/utils/device_lazy_init.cpp
+++ b/torch/csrc/utils/device_lazy_init.cpp
@@ -12,24 +12,21 @@
 namespace torch::utils {
 namespace {
 
-std::array<bool, at::COMPILE_TIME_MAX_DEVICE_TYPES> is_initialized{};
-std::array<bool, at::COMPILE_TIME_MAX_DEVICE_TYPES> is_in_bad_fork{};
-std::array<bool, at::COMPILE_TIME_MAX_DEVICE_TYPES> at_fork_registered{};
+std::array<std::atomic<bool>, at::COMPILE_TIME_MAX_DEVICE_TYPES>
+    is_initialized{};
+std::array<std::atomic<bool>, at::COMPILE_TIME_MAX_DEVICE_TYPES>
+    is_in_bad_fork{};
+std::array<std::atomic<bool>, at::COMPILE_TIME_MAX_DEVICE_TYPES>
+    at_fork_registered{};
 c10::once_flag at_fork_register_once{};
 
 } // anonymous namespace
 
 bool is_device_initialized(at::DeviceType device_type) {
-  pybind11::gil_scoped_acquire g;
   return is_initialized[static_cast<int>(device_type)];
 }
 
 void device_lazy_init(at::DeviceType device_type) {
-  pybind11::gil_scoped_acquire g;
-  // Protected by the GIL.  We don't use call_once because under ASAN it
-  // has a buggy implementation that deadlocks if an instance throws an
-  // exception.  In any case, call_once isn't necessary, because we
-  // have taken a lock.
   if (is_device_initialized(device_type)) {
     return;
   }
@@ -40,6 +37,10 @@ void device_lazy_init(at::DeviceType device_type) {
     return;
   }
 
+  // Don't use call_once because under ASAN it has a buggy implementation that
+  // deadlocks if an instance throws an exception and Python _lazy_init()
+  // functions are idempotent.
+  pybind11::gil_scoped_acquire g;
   std::string module_name = "torch." + at::DeviceTypeName(device_type, true);
   auto module = THPObjectPtr(PyImport_ImportModule(module_name.c_str()));
   if (!module) {
@@ -50,6 +51,7 @@ void device_lazy_init(at::DeviceType device_type) {
     auto has_lazy_init_method =
         PyObject_HasAttrString(module.get(), "_lazy_init") == 1;
     if (!has_lazy_init_method) {
+      is_initialized[static_cast<int>(device_type)] = true;
       return;
     }
   }
diff --git a/torch/csrc/utils/generated_serialization_types.h b/torch/csrc/utils/generated_serialization_types.h
index 93f469d8390f0..995e2405cb4da 100644
--- a/torch/csrc/utils/generated_serialization_types.h
+++ b/torch/csrc/utils/generated_serialization_types.h
@@ -1,5 +1,5 @@
 // @generated by update_schema.py
-// checksum<<5b9bcd54d4ce78ce8bf3ec2e2cad62f19889537e14223d9e9adcfba1f2edc82c>>
+// checksum<<eb9645ae5f539caace53ada70fb262097ce780e6edf685149dafc5a5abe8d31b>>
 // clang-format off
 
 #pragma once
@@ -288,6 +288,9 @@ enum class ScalarType {
   FLOAT8E5M2 = 30,
   FLOAT8E4M3FNUZ = 31,
   FLOAT8E5M2FNUZ = 32,
+  FLOAT8E8M0FNU = 33,
+  UINT32 = 34,
+  UINT64 = 35,
 };
 
 inline std::string_view printEnum(const ScalarType& e) {
@@ -311,6 +314,9 @@ inline std::string_view printEnum(const ScalarType& e) {
     case ScalarType::FLOAT8E5M2: return "FLOAT8E5M2";
     case ScalarType::FLOAT8E4M3FNUZ: return "FLOAT8E4M3FNUZ";
     case ScalarType::FLOAT8E5M2FNUZ: return "FLOAT8E5M2FNUZ";
+    case ScalarType::FLOAT8E8M0FNU: return "FLOAT8E8M0FNU";
+    case ScalarType::UINT32: return "UINT32";
+    case ScalarType::UINT64: return "UINT64";
     default:
       throw std::runtime_error("Unknown enum value");
   }
@@ -336,6 +342,9 @@ inline void parseEnum(std::string_view s, ScalarType& t) {
   if (s == "FLOAT8E5M2") { t = ScalarType::FLOAT8E5M2; return; }
   if (s == "FLOAT8E4M3FNUZ") { t = ScalarType::FLOAT8E4M3FNUZ; return; }
   if (s == "FLOAT8E5M2FNUZ") { t = ScalarType::FLOAT8E5M2FNUZ; return; }
+  if (s == "FLOAT8E8M0FNU") { t = ScalarType::FLOAT8E8M0FNU; return; }
+  if (s == "UINT32") { t = ScalarType::UINT32; return; }
+  if (s == "UINT64") { t = ScalarType::UINT64; return; }
   throw std::runtime_error("Unknown enum value: " + std::string{s});
 }
 
@@ -1232,11 +1241,11 @@ class Argument {
 
  public:
   enum class Tag {
-    AS_NONE, AS_TENSOR, AS_TENSORS, AS_INT, AS_INTS, AS_FLOAT, AS_FLOATS, AS_STRING, AS_STRINGS, AS_SYM_INT, AS_SYM_INTS, AS_SCALAR_TYPE, AS_MEMORY_FORMAT, AS_LAYOUT, AS_DEVICE, AS_BOOL, AS_BOOLS, AS_SYM_BOOL, AS_SYM_BOOLS, AS_GRAPH, AS_OPTIONAL_TENSORS, AS_CUSTOM_OBJ, AS_OPERATOR, AS_SYM_FLOAT, AS_SYM_FLOATS, AS_OPTIONAL_TENSOR, AS_COMPLEX, AS_NESTED_TENSORS, AS_INT_LISTS, AS_STRING_TO_ARGUMENT
+    AS_NONE, AS_TENSOR, AS_TENSORS, AS_INT, AS_INTS, AS_FLOAT, AS_FLOATS, AS_STRING, AS_STRINGS, AS_SYM_INT, AS_SYM_INTS, AS_SCALAR_TYPE, AS_MEMORY_FORMAT, AS_LAYOUT, AS_DEVICE, AS_BOOL, AS_BOOLS, AS_SYM_BOOL, AS_SYM_BOOLS, AS_GRAPH, AS_OPTIONAL_TENSORS, AS_CUSTOM_OBJ, AS_OPERATOR, AS_SYM_FLOAT, AS_SYM_FLOATS, AS_OPTIONAL_TENSOR, AS_COMPLEX, AS_NESTED_TENSORS, AS_INT_LISTS, AS_STRING_TO_ARGUMENT, AS_FLOAT_LISTS
   };
 
  private:
-  std::variant<Void, bool, TensorArgument, std::vector<TensorArgument>, int64_t, std::vector<int64_t>, F64, std::vector<F64>, std::string, std::vector<std::string>, SymIntArgument, std::vector<SymIntArgument>, ScalarType, MemoryFormat, Layout, Device, bool, std::vector<bool>, SymBoolArgument, std::vector<SymBoolArgument>, GraphArgument, std::vector<OptionalTensorArgument>, CustomObjArgument, std::string, SymFloatArgument, std::vector<SymFloatArgument>, OptionalTensorArgument, ComplexValue, std::vector<std::vector<TensorArgument>>, std::vector<std::vector<int64_t>>, std::unordered_map<std::string, ForwardRef<Argument>>> variant_;
+  std::variant<Void, bool, TensorArgument, std::vector<TensorArgument>, int64_t, std::vector<int64_t>, F64, std::vector<F64>, std::string, std::vector<std::string>, SymIntArgument, std::vector<SymIntArgument>, ScalarType, MemoryFormat, Layout, Device, bool, std::vector<bool>, SymBoolArgument, std::vector<SymBoolArgument>, GraphArgument, std::vector<OptionalTensorArgument>, CustomObjArgument, std::string, SymFloatArgument, std::vector<SymFloatArgument>, OptionalTensorArgument, ComplexValue, std::vector<std::vector<TensorArgument>>, std::vector<std::vector<int64_t>>, std::unordered_map<std::string, ForwardRef<Argument>>, std::vector<std::vector<F64>>> variant_;
   Tag tag_;
 
  public:
@@ -1514,6 +1523,15 @@ class Argument {
     tag_ = Tag::AS_STRING_TO_ARGUMENT;
   }
 
+  const std::vector<std::vector<F64>>& get_as_float_lists() const {
+    return std::get<31>(variant_);
+  }
+
+  void set_as_float_lists(std::vector<std::vector<F64>> def) {
+    variant_.emplace<31>(std::move(def));
+    tag_ = Tag::AS_FLOAT_LISTS;
+  }
+
   friend void to_json(nlohmann::json& nlohmann_json_j, const Argument& nlohmann_json_t) {
 
     if (nlohmann_json_t.tag_ == Tag::AS_NONE) {
@@ -1636,6 +1654,10 @@ class Argument {
       nlohmann_json_j["as_string_to_argument"] = nlohmann_json_t.get_as_string_to_argument();
       return;
     }
+    if (nlohmann_json_t.tag_ == Tag::AS_FLOAT_LISTS) {
+      nlohmann_json_j["as_float_lists"] = nlohmann_json_t.get_as_float_lists();
+      return;
+    }
   }
 
   friend void from_json(const nlohmann::json& nlohmann_json_j, Argument& nlohmann_json_t) {
@@ -1790,6 +1812,11 @@ class Argument {
       nlohmann_json_t.tag_ = Tag::AS_STRING_TO_ARGUMENT;
       return;
     }
+    if (nlohmann_json_j.contains("as_float_lists")) {
+      nlohmann_json_t.variant_.emplace<31>(nlohmann_json_j.at("as_float_lists").template get<std::vector<std::vector<F64>>>());
+      nlohmann_json_t.tag_ = Tag::AS_FLOAT_LISTS;
+      return;
+    }
   }
 };
 
@@ -1825,6 +1852,7 @@ inline std::string_view printEnum(const Argument::Tag& e) {
     case Argument::Tag::AS_NESTED_TENSORS: return "AS_NESTED_TENSORS";
     case Argument::Tag::AS_INT_LISTS: return "AS_INT_LISTS";
     case Argument::Tag::AS_STRING_TO_ARGUMENT: return "AS_STRING_TO_ARGUMENT";
+    case Argument::Tag::AS_FLOAT_LISTS: return "AS_FLOAT_LISTS";
     default:
       throw std::runtime_error("Unknown enum value");
   }
@@ -1861,6 +1889,7 @@ inline void parseEnum(std::string_view s, Argument::Tag& t) {
   if (s == "AS_NESTED_TENSORS") { t = Argument::Tag::AS_NESTED_TENSORS; return; }
   if (s == "AS_INT_LISTS") { t = Argument::Tag::AS_INT_LISTS; return; }
   if (s == "AS_STRING_TO_ARGUMENT") { t = Argument::Tag::AS_STRING_TO_ARGUMENT; return; }
+  if (s == "AS_FLOAT_LISTS") { t = Argument::Tag::AS_FLOAT_LISTS; return; }
   throw std::runtime_error("Unknown enum value: " + std::string{s});
 }
 
diff --git a/torch/csrc/utils/python_arg_parser.cpp b/torch/csrc/utils/python_arg_parser.cpp
index 999bd00b3bcd6..a9f52cafb65ad 100644
--- a/torch/csrc/utils/python_arg_parser.cpp
+++ b/torch/csrc/utils/python_arg_parser.cpp
@@ -302,8 +302,22 @@ static py::object maybe_get_registered_torch_dispatch_rule(
 static bool is_dtensor(PyObject* obj) {
 #ifdef USE_DISTRIBUTED
   const py::handle dtensor = get_dtensor_class();
-  return (PyObject*)Py_TYPE(obj) == dtensor.ptr() ||
-      py::isinstance(py::handle(obj), dtensor);
+  if ((PyObject*)Py_TYPE(obj) == dtensor.ptr()) {
+    return true;
+  }
+  if (!py::isinstance(py::handle(obj), dtensor)) {
+    return false;
+  }
+  // DTensor subclass: only use the C++ fast path if it does not override
+  // __torch_dispatch__. Subclasses with a custom override should fall
+  // through to the normal Python dispatch path.
+  // Compare via __func__ because @classmethod descriptors create new bound
+  // method objects on each attr access, making direct identity checks fail.
+  static py::object base_td =
+      dtensor.attr("__torch_dispatch__").attr("__func__");
+  py::object sub_td =
+      py::type::handle_of(obj).attr("__torch_dispatch__").attr("__func__");
+  return sub_td.is(base_td);
 #else
   return false;
 #endif
@@ -755,9 +769,12 @@ auto handle_torch_function_indexing(
   }
   py::object func =
       PyObject_FastGetAttrString(THPVariableClass, (char*)func_name);
-  py::object args = (val == nullptr)
-      ? py::make_tuple(py::handle(self), py::handle(index))
-      : py::make_tuple(py::handle(self), py::handle(index), py::handle(val));
+  py::tuple args;
+  if (val == nullptr) {
+    args = py::make_tuple(py::handle(self), py::handle(index));
+  } else {
+    args = py::make_tuple(py::handle(self), py::handle(index), py::handle(val));
+  }
   return handle_torch_function_no_python_arg_parser(
       overridable_args,
       args.ptr(),
@@ -878,7 +895,7 @@ bool is_tensor_and_append_overloaded(
 static bool is_scalar_list(
     PyObject* obj,
     std::vector<PyObject*>* overloaded_args = nullptr) {
-  auto tuple = six::isTuple(obj);
+  auto tuple = PyTuple_Check(obj);
   if (!(tuple || PyList_Check(obj))) {
     return false;
   }
@@ -909,7 +926,7 @@ bool is_tensor_list_and_append_overloaded(
     std::vector<PyObject*>* overloaded_args,
     size_t argnum,
     bool throw_error) {
-  auto tuple = six::isTuple(obj);
+  auto tuple = PyTuple_Check(obj);
   if (!(tuple || PyList_Check(obj))) {
     return false;
   }
@@ -950,7 +967,7 @@ static bool is_float_or_symfloat(PyObject* obj) {
 static bool is_float_or_complex_list(
     PyObject* obj,
     std::vector<PyObject*>* overloaded_args = nullptr) {
-  auto tuple = six::isTuple(obj);
+  auto tuple = PyTuple_Check(obj);
   if (!(tuple || PyList_Check(obj))) {
     return false;
   }
diff --git a/torch/csrc/utils/python_arg_parser.h b/torch/csrc/utils/python_arg_parser.h
index 55a5902450951..c72d3046a90dc 100644
--- a/torch/csrc/utils/python_arg_parser.h
+++ b/torch/csrc/utils/python_arg_parser.h
@@ -65,7 +65,7 @@
 #include <torch/csrc/utils/python_numbers.h>
 #include <torch/csrc/utils/python_strings.h>
 #include <torch/csrc/utils/python_symnode.h>
-#include <torch/csrc/utils/six.h>
+#include <torch/csrc/utils/structseq.h>
 
 #include <ATen/DeviceAccelerator.h>
 #include <ATen/PythonTorchFunctionTLS.h>
@@ -125,8 +125,6 @@ enum class ParameterType {
   DISPATCH_KEY_SET
 };
 
-struct FunctionParameter;
-struct FunctionSignature;
 struct PythonArgs;
 
 // Contains bound Python arguments in declaration order
@@ -137,6 +135,81 @@ struct ParsedArgs {
   PyObject* args[N];
 };
 
+// FunctionParameter is a single formal parameter of a Python function.
+// It is immutable once constructed.
+struct FunctionParameter {
+  FunctionParameter(const std::string& fmt, bool keyword_only);
+
+  bool check(
+      PyObject* obj,
+      std::vector<PyObject*>& overloaded_args,
+      int argnum,
+      int64_t* failed_idx = nullptr);
+
+  bool _check(
+      PyObject* obj,
+      std::vector<PyObject*>& overloaded_args,
+      int argnum,
+      int64_t* failed_idx = nullptr);
+
+  void set_default_str(const std::string& str);
+  TORCH_PYTHON_API std::string type_name() const;
+
+  ParameterType type_;
+  bool optional{false};
+  bool allow_none{false};
+  bool keyword_only;
+  bool allow_numbers_as_tensors = false;
+  int size{0};
+  std::string name;
+  // having this as a raw PyObject * will presumably leak it, but these are only
+  // held by static objects anyway, and Py_Finalize can already be called when
+  // this is destructed.
+  PyObject* python_name;
+  // NOLINTNEXTLINE(cppcoreguidelines-avoid-magic-numbers)
+  at::SmallVector<PyObject*, 5> numpy_python_names;
+  at::Scalar default_scalar;
+  std::vector<int64_t> default_intlist;
+  std::string default_string;
+  union {
+    bool default_bool;
+    int64_t default_int;
+    double default_double;
+    // NOLINTNEXTLINE(cppcoreguidelines-avoid-c-arrays,modernize-avoid-c-arrays)
+    double default_complex[2]; // see Scalar
+    at::ScalarType default_scalartype;
+    at::Layout default_layout;
+  };
+  std::string default_value;
+};
+
+// FunctionSignature represents a single valid signature for a Python function.
+// It is immutable once constructed. The contained data can be concurrently
+// accessed by multiple calls.
+struct FunctionSignature {
+  explicit FunctionSignature(const std::string& fmt, int index);
+
+  bool parse(
+      PyObject* self,
+      PyObject* args,
+      PyObject* kwargs,
+      // NOLINTNEXTLINE(cppcoreguidelines-avoid-c-arrays,modernize-avoid-c-arrays)
+      PyObject* dst[],
+      std::vector<PyObject*>& overloaded_args,
+      bool raise_exception);
+
+  std::string toString() const;
+
+  std::string name;
+  std::vector<FunctionParameter> params;
+  size_t min_args{0};
+  size_t max_args{0};
+  size_t max_pos_args{0};
+  int index;
+  bool hidden{false};
+  bool deprecated{false};
+};
+
 // A PythonArgParser contains a list of valid signatures. Instances are
 // typically global variables and should be immutable.
 struct PYBIND11_EXPORT PythonArgParser {
@@ -181,33 +254,6 @@ struct PYBIND11_EXPORT PythonArgParser {
   bool traceable;
 };
 
-// FunctionSignature represents a single valid signature for a Python function.
-// It is immutable once constructed. The contained data can be concurrently
-// accessed by multiple calls.
-struct FunctionSignature {
-  explicit FunctionSignature(const std::string& fmt, int index);
-
-  bool parse(
-      PyObject* self,
-      PyObject* args,
-      PyObject* kwargs,
-      // NOLINTNEXTLINE(cppcoreguidelines-avoid-c-arrays,modernize-avoid-c-arrays)
-      PyObject* dst[],
-      std::vector<PyObject*>& overloaded_args,
-      bool raise_exception);
-
-  std::string toString() const;
-
-  std::string name;
-  std::vector<FunctionParameter> params;
-  size_t min_args{0};
-  size_t max_args{0};
-  size_t max_pos_args{0};
-  int index;
-  bool hidden{false};
-  bool deprecated{false};
-};
-
 // PythonArgs contains bound Python arguments for an actual invocation
 // along with references to the matched signature.
 struct TORCH_PYTHON_API PythonArgs {
@@ -311,54 +357,6 @@ struct TORCH_PYTHON_API PythonArgs {
   at::Scalar scalar_slow(PyObject* arg);
 };
 
-// FunctionParameter is a single formal parameter of a Python function.
-// It is immutable once constructed.
-struct FunctionParameter {
-  FunctionParameter(const std::string& fmt, bool keyword_only);
-
-  bool check(
-      PyObject* obj,
-      std::vector<PyObject*>& overloaded_args,
-      int argnum,
-      int64_t* failed_idx = nullptr);
-
-  bool _check(
-      PyObject* obj,
-      std::vector<PyObject*>& overloaded_args,
-      int argnum,
-      int64_t* failed_idx = nullptr);
-
-  void set_default_str(const std::string& str);
-  TORCH_PYTHON_API std::string type_name() const;
-
-  ParameterType type_;
-  bool optional{false};
-  bool allow_none{false};
-  bool keyword_only;
-  bool allow_numbers_as_tensors = false;
-  int size{0};
-  std::string name;
-  // having this as a raw PyObject * will presumably leak it, but these are only
-  // held by static objects anyway, and Py_Finalize can already be called when
-  // this is destructed.
-  PyObject* python_name;
-  // NOLINTNEXTLINE(cppcoreguidelines-avoid-magic-numbers)
-  at::SmallVector<PyObject*, 5> numpy_python_names;
-  at::Scalar default_scalar;
-  std::vector<int64_t> default_intlist;
-  std::string default_string;
-  union {
-    bool default_bool;
-    int64_t default_int;
-    double default_double;
-    // NOLINTNEXTLINE(cppcoreguidelines-avoid-c-arrays,modernize-avoid-c-arrays)
-    double default_complex[2]; // see Scalar
-    at::ScalarType default_scalartype;
-    at::Layout default_layout;
-  };
-  std::string default_value;
-};
-
 template <int N>
 inline PythonArgs PythonArgParser::parse(
     PyObject* self,
@@ -422,14 +420,14 @@ inline at::Scalar PythonArgs::scalar(int i) {
 inline std::vector<at::Scalar> PythonArgs::scalarlist(int i) {
   if (!args[i])
     return std::vector<at::Scalar>();
-  auto tuple = six::isTuple(args[i]);
-  THPObjectPtr arg = six::maybeAsTuple(args[i]);
+  auto tuple = PyTuple_Check(args[i]);
+  TORCH_INTERNAL_ASSERT_DEBUG_ONLY(tuple || PyList_Check(args[i]));
   // NOLINTNEXTLINE(bugprone-branch-clone)
-  auto size = tuple ? PyTuple_GET_SIZE(arg.get()) : PyList_GET_SIZE(arg.get());
+  auto size = tuple ? PyTuple_GET_SIZE(args[i]) : PyList_GET_SIZE(args[i]);
   std::vector<at::Scalar> res(size);
   for (const auto idx : c10::irange(size)) {
-    PyObject* obj = tuple ? PyTuple_GET_ITEM(arg.get(), idx)
-                          : PyList_GET_ITEM(arg.get(), idx);
+    PyObject* obj =
+        tuple ? PyTuple_GET_ITEM(args[i], idx) : PyList_GET_ITEM(args[i], idx);
     res[idx] = scalar_slow(obj);
   }
   return res;
@@ -452,14 +450,14 @@ inline std::optional<at::Scalar> PythonArgs::scalarOptional(int i) {
 inline std::vector<at::Tensor> PythonArgs::tensorlist(int i) {
   if (!args[i])
     return std::vector<at::Tensor>();
-  auto tuple = six::isTuple(args[i]);
-  THPObjectPtr arg = six::maybeAsTuple(args[i]);
+  auto tuple = PyTuple_Check(args[i]);
+  TORCH_INTERNAL_ASSERT_DEBUG_ONLY(tuple || PyList_Check(args[i]));
   // NOLINTNEXTLINE(bugprone-branch-clone)
-  auto size = tuple ? PyTuple_GET_SIZE(arg.get()) : PyList_GET_SIZE(arg.get());
+  auto size = tuple ? PyTuple_GET_SIZE(args[i]) : PyList_GET_SIZE(args[i]);
   std::vector<at::Tensor> res(size);
   for (const auto idx : c10::irange(size)) {
-    PyObject* obj = tuple ? PyTuple_GET_ITEM(arg.get(), idx)
-                          : PyList_GET_ITEM(arg.get(), idx);
+    PyObject* obj =
+        tuple ? PyTuple_GET_ITEM(args[i], idx) : PyList_GET_ITEM(args[i], idx);
     // This is checked by the argument parser so it's safe to cast without
     // checking if this is a tensor first
     res[idx] = THPVariable_Unpack(obj);
@@ -471,15 +469,15 @@ inline torch::List<std::optional<at::Tensor>> PythonArgs::
     list_of_optional_tensors(int i) {
   if (!args[i])
     return torch::List<std::optional<at::Tensor>>();
-  auto tuple = six::isTuple(args[i]);
-  THPObjectPtr arg = six::maybeAsTuple(args[i]);
+  auto tuple = PyTuple_Check(args[i]);
+  TORCH_INTERNAL_ASSERT_DEBUG_ONLY(tuple || PyList_Check(args[i]));
   // NOLINTNEXTLINE(bugprone-branch-clone)
-  auto size = tuple ? PyTuple_GET_SIZE(arg.get()) : PyList_GET_SIZE(arg.get());
+  auto size = tuple ? PyTuple_GET_SIZE(args[i]) : PyList_GET_SIZE(args[i]);
   torch::List<std::optional<at::Tensor>> res;
   res.reserve(size);
   for (const auto idx : c10::irange(size)) {
-    PyObject* obj = tuple ? PyTuple_GET_ITEM(arg.get(), idx)
-                          : PyList_GET_ITEM(arg.get(), idx);
+    PyObject* obj =
+        tuple ? PyTuple_GET_ITEM(args[i], idx) : PyList_GET_ITEM(args[i], idx);
     // This is checked by the argument parser so it's safe to cast without
     // checking if this is a tensor first
     res.push_back(THPVariable_Unpack(obj));
@@ -492,18 +490,18 @@ inline std::array<at::Tensor, N> PythonArgs::tensorlist_n(int i) {
   auto res = std::array<at::Tensor, N>();
   if (!args[i])
     return res;
-  auto tuple = six::isTuple(args[i]);
-  THPObjectPtr arg = six::maybeAsTuple(args[i]);
+  auto tuple = PyTuple_Check(args[i]);
+  TORCH_INTERNAL_ASSERT_DEBUG_ONLY(tuple || PyList_Check(args[i]));
   // NOLINTNEXTLINE(bugprone-branch-clone)
-  auto size = tuple ? PyTuple_GET_SIZE(arg.get()) : PyList_GET_SIZE(arg.get());
+  auto size = tuple ? PyTuple_GET_SIZE(args[i]) : PyList_GET_SIZE(args[i]);
   if (size != N) {
     TORCH_CHECK_TYPE(
         false,
         fmt::format("expected tuple of {} elements but got {}", N, size));
   }
   for (const auto idx : c10::irange(size)) {
-    PyObject* obj = tuple ? PyTuple_GET_ITEM(arg.get(), idx)
-                          : PyList_GET_ITEM(arg.get(), idx);
+    PyObject* obj =
+        tuple ? PyTuple_GET_ITEM(args[i], idx) : PyList_GET_ITEM(args[i], idx);
     // This is checked by the argument parser so it's safe to cast without
     // checking if this is a tensor first
     res[idx] = THPVariable_Unpack(obj);
diff --git a/torch/csrc/utils/python_dispatch.cpp b/torch/csrc/utils/python_dispatch.cpp
index 0be7c4a3970af..c141aa928307a 100644
--- a/torch/csrc/utils/python_dispatch.cpp
+++ b/torch/csrc/utils/python_dispatch.cpp
@@ -18,6 +18,7 @@
 #include <torch/csrc/jit/python/pybind_utils.h>
 #include <torch/csrc/utils/tensor_new.h>
 
+#include <c10/util/Synchronized.h>
 #include <c10/util/flat_hash_map.h>
 #include <torch/csrc/inductor/aoti_eager/kernel_holder.h>
 #include <torch/csrc/utils/python_raii.h>
@@ -30,15 +31,18 @@ namespace py = pybind11;
 namespace torch::impl::dispatch {
 
 // Global storage for leaked Python filenames to ensure they remain valid
-// for the lifetime of Library objects
-static std::vector<std::string> leaked_python_filenames_;
+// for the lifetime of Library objects. We use unique_ptr<string> rather than
+// plain string so that c_str() pointers handed to Library objects remain valid
+// when the vector reallocates.
+static c10::Synchronized<std::vector<std::unique_ptr<std::string>>>
+    leaked_python_filenames_;
 
 // NB: I'd like to index this on OperatorHandle, but I can't, as I can't
 // guarantee that the main interpreter has finish doing all registrations before
 // the other interpreters start banging on it
-static ska::flat_hash_map<
+static c10::Synchronized<ska::flat_hash_map<
     c10::OperatorName,
-    ska::flat_hash_map<c10::DispatchKey, std::shared_ptr<c10::SafePyObject>>>
+    ska::flat_hash_map<c10::DispatchKey, std::shared_ptr<c10::SafePyObject>>>>
     python_registrations_;
 
 static torch::Library::Kind parseKind(const std::string& k) {
@@ -424,10 +428,12 @@ void initDispatchBindings(PyObject* module) {
                           std::make_unique<PythonKernelHolder>(
                               func, dispatch, with_keyset))),
                   register_or_verify());
-              python_registrations_[lib._resolve(name)].insert_or_assign(
-                  dispatch,
-                  std::make_shared<c10::SafePyObject>(
-                      func.release().ptr(), getPyInterpreter()));
+              python_registrations_.withLock([&](auto& regs) {
+                regs[lib._resolve(name)].insert_or_assign(
+                    dispatch,
+                    std::make_shared<c10::SafePyObject>(
+                        func.release().ptr(), getPyInterpreter()));
+              });
             }
             END_HANDLE_TH_ERRORS_PYBIND
           },
@@ -512,8 +518,11 @@ void initDispatchBindings(PyObject* module) {
         HANDLE_TH_ERRORS
         // Store the file string in global storage to ensure it remains valid
         // for the lifetime of the Library object
-        leaked_python_filenames_.emplace_back(file);
-        const char* leaked_file = leaked_python_filenames_.back().c_str();
+        const char* leaked_file =
+            leaked_python_filenames_.withLock([&](auto& filenames) {
+              filenames.push_back(std::make_unique<std::string>(file));
+              return filenames.back()->c_str();
+            });
 
         return std::make_unique<torch::Library>(
             parseKind(kind),
@@ -534,7 +543,7 @@ void initDispatchBindings(PyObject* module) {
 
   m.def(
       "_dispatch_clear_leaked_python_filenames",
-      []() { leaked_python_filenames_.clear(); },
+      []() { leaked_python_filenames_.withLock([](auto& f) { f.clear(); }); },
       "Clear the global storage of leaked Python filenames. "
       "WARNING: Only call this if you're sure no Library objects are still using the filenames.");
 
@@ -1077,7 +1086,8 @@ void python_op_registration_trampoline_impl(
   auto arguments = torch::jit::pop(*stack, op.schema().arguments().size());
   py::gil_scoped_acquire g;
   auto args_kwargs = parseIValuesToPyArgsKwargs(op, arguments);
-  const auto& func = python_registrations_[op.operator_name()][key];
+  auto func = python_registrations_.withLock(
+      [&](auto& regs) { return regs[op.operator_name()][key]; });
   TORCH_INTERNAL_ASSERT(func != nullptr);
   auto* pyobj = func->ptr(getPyInterpreter());
   TORCH_INTERNAL_ASSERT(pyobj != nullptr);
diff --git a/torch/csrc/utils/six.h b/torch/csrc/utils/six.h
deleted file mode 100644
index 9671e5156b9d4..0000000000000
--- a/torch/csrc/utils/six.h
+++ /dev/null
@@ -1,52 +0,0 @@
-#pragma once
-
-#include <pybind11/pybind11.h>
-#include <torch/csrc/utils/object_ptr.h>
-#include <torch/csrc/utils/pybind.h>
-#include <torch/csrc/utils/structseq.h>
-
-namespace six {
-
-// Usually instances of PyStructSequence is also an instance of tuple
-// but in some py2 environment it is not, so we have to manually check
-// the name of the type to determine if it is a namedtupled returned
-// by a pytorch operator.
-
-inline bool isStructSeq(pybind11::handle input) {
-  return pybind11::cast<std::string>(pybind11::type::handle_of(input).attr(
-             "__module__")) == "torch.return_types";
-}
-
-inline bool isStructSeq(PyObject* obj) {
-  return isStructSeq(pybind11::handle(obj));
-}
-
-inline bool isTuple(pybind11::handle input) {
-  if (PyTuple_Check(input.ptr())) {
-    return true;
-  }
-  return false;
-}
-
-inline bool isTuple(PyObject* obj) {
-  return isTuple(pybind11::handle(obj));
-}
-
-// maybeAsTuple: if the input is a structseq, then convert it to a tuple
-//
-// On Python 3, structseq is a subtype of tuple, so these APIs could be used
-// directly. But on Python 2, structseq is not a subtype of tuple, so we need to
-// manually create a new tuple object from structseq.
-inline THPObjectPtr maybeAsTuple(PyStructSequence* obj) {
-  Py_INCREF(obj);
-  return THPObjectPtr((PyObject*)obj);
-}
-
-inline THPObjectPtr maybeAsTuple(PyObject* obj) {
-  if (isStructSeq(obj))
-    return maybeAsTuple((PyStructSequence*)obj);
-  Py_INCREF(obj);
-  return THPObjectPtr(obj);
-}
-
-} // namespace six
diff --git a/torch/csrc/utils/structseq.cpp b/torch/csrc/utils/structseq.cpp
index 2e804aa44bad9..ded49128cf284 100644
--- a/torch/csrc/utils/structseq.cpp
+++ b/torch/csrc/utils/structseq.cpp
@@ -12,7 +12,7 @@
  * https://github.com/python/cpython#copyright-and-license-information
  */
 
-#include <torch/csrc/utils/six.h>
+#include <torch/csrc/utils/object_ptr.h>
 #include <torch/csrc/utils/structseq.h>
 #include <sstream>
 
@@ -20,19 +20,12 @@
 
 namespace torch::utils {
 
-// NOTE: The built-in repr method from PyStructSequence was updated in
-// https://github.com/python/cpython/commit/c70ab02df2894c34da2223fc3798c0404b41fd79
-// so this function might not be required in Python 3.8+.
 PyObject* returned_structseq_repr(PyStructSequence* obj) {
   PyTypeObject* typ = Py_TYPE(obj);
-  THPObjectPtr tup = six::maybeAsTuple(obj);
-  if (tup == nullptr) {
-    return nullptr;
-  }
+  Py_ssize_t num_elements = PyTuple_GET_SIZE(obj);
 
   std::stringstream ss;
   ss << typ->tp_name << "(\n";
-  Py_ssize_t num_elements = Py_SIZE(obj);
 
   for (Py_ssize_t i = 0; i < num_elements; i++) {
     const char* cname = typ->tp_members[i].name;
@@ -46,11 +39,7 @@ PyObject* returned_structseq_repr(PyStructSequence* obj) {
       return nullptr;
     }
 
-    PyObject* val = PyTuple_GetItem(tup.get(), i);
-    if (val == nullptr) {
-      return nullptr;
-    }
-
+    PyObject* val = PyTuple_GET_ITEM(obj, i);
     auto repr = THPObjectPtr(PyObject_Repr(val));
     if (repr == nullptr) {
       return nullptr;
diff --git a/torch/csrc/utils/tensor_dtypes.cpp b/torch/csrc/utils/tensor_dtypes.cpp
index 39df9be68868a..938c7bc7558d0 100644
--- a/torch/csrc/utils/tensor_dtypes.cpp
+++ b/torch/csrc/utils/tensor_dtypes.cpp
@@ -20,17 +20,15 @@ void initializeDtypes() {
 
   for (at::ScalarType scalarType : all_scalar_types) {
     auto [primary_name, legacy_name] = c10::getDtypeNames(scalarType);
-    PyObject* dtype = THPDtype_New(scalarType, primary_name);
-    torch::registerDtypeObject((THPDtype*)dtype, scalarType);
-    Py_INCREF(dtype);
-    if (PyModule_AddObject(torch_module.get(), primary_name.c_str(), dtype) !=
-        0) {
+    THPObjectPtr dtype(THPDtype_New(scalarType, primary_name));
+    torch::registerDtypeObject((THPDtype*)dtype.get(), scalarType);
+    if (PyModule_AddObjectRef(
+            torch_module.get(), primary_name.c_str(), dtype.get()) != 0) {
       throw python_error();
     }
     if (!legacy_name.empty()) {
-      Py_INCREF(dtype);
-      if (PyModule_AddObject(torch_module.get(), legacy_name.c_str(), dtype) !=
-          0) {
+      if (PyModule_AddObjectRef(
+              torch_module.get(), legacy_name.c_str(), dtype.get()) != 0) {
         throw python_error();
       }
     }
diff --git a/torch/csrc/utils/tensor_layouts.cpp b/torch/csrc/utils/tensor_layouts.cpp
index d0bccbcf9106f..3b1e04cb08a30 100644
--- a/torch/csrc/utils/tensor_layouts.cpp
+++ b/torch/csrc/utils/tensor_layouts.cpp
@@ -6,49 +6,36 @@
 
 namespace torch::utils {
 
-#define REGISTER_LAYOUT(layout, LAYOUT)                                     \
-  PyObject* layout##_layout =                                               \
-      THPLayout_New(at::Layout::LAYOUT, "torch." #layout);                  \
-  Py_INCREF(layout##_layout);                                               \
-  if (PyModule_AddObject(torch_module, "" #layout, layout##_layout) != 0) { \
-    throw python_error();                                                   \
-  }                                                                         \
-  registerLayoutObject((THPLayout*)layout##_layout, at::Layout::LAYOUT);
+static void registerLayout(
+    PyObject* torch_module,
+    at::Layout layout,
+    const char* name,
+    const char* qualified_name) {
+  THPObjectPtr obj(THPLayout_New(layout, qualified_name));
+  if (PyModule_AddObjectRef(torch_module, name, obj.get()) != 0) {
+    throw python_error();
+  }
+  registerLayoutObject((THPLayout*)obj.get(), layout);
+}
 
 void initializeLayouts() {
   auto torch_module = THPObjectPtr(PyImport_ImportModule("torch"));
   if (!torch_module)
     throw python_error();
 
-  PyObject* strided_layout =
-      THPLayout_New(at::Layout::Strided, "torch.strided");
-  Py_INCREF(strided_layout);
-  if (PyModule_AddObject(torch_module, "strided", strided_layout) != 0) {
-    throw python_error();
-  }
-  registerLayoutObject((THPLayout*)strided_layout, at::Layout::Strided);
-
-  PyObject* sparse_coo_layout =
-      THPLayout_New(at::Layout::Sparse, "torch.sparse_coo");
-  Py_INCREF(sparse_coo_layout);
-  if (PyModule_AddObject(torch_module, "sparse_coo", sparse_coo_layout) != 0) {
-    throw python_error();
-  }
-  registerLayoutObject((THPLayout*)sparse_coo_layout, at::Layout::Sparse);
-
-  REGISTER_LAYOUT(sparse_csr, SparseCsr)
-  REGISTER_LAYOUT(sparse_csc, SparseCsc)
-  REGISTER_LAYOUT(sparse_bsr, SparseBsr)
-  REGISTER_LAYOUT(sparse_bsc, SparseBsc)
-
-  PyObject* mkldnn_layout = THPLayout_New(at::Layout::Mkldnn, "torch._mkldnn");
-  Py_INCREF(mkldnn_layout);
-  if (PyModule_AddObject(torch_module, "_mkldnn", mkldnn_layout) != 0) {
-    throw python_error();
-  }
-  registerLayoutObject((THPLayout*)mkldnn_layout, at::Layout::Mkldnn);
-
-  REGISTER_LAYOUT(jagged, Jagged);
+  registerLayout(torch_module, at::Layout::Strided, "strided", "torch.strided");
+  registerLayout(
+      torch_module, at::Layout::Sparse, "sparse_coo", "torch.sparse_coo");
+  registerLayout(
+      torch_module, at::Layout::SparseCsr, "sparse_csr", "torch.sparse_csr");
+  registerLayout(
+      torch_module, at::Layout::SparseCsc, "sparse_csc", "torch.sparse_csc");
+  registerLayout(
+      torch_module, at::Layout::SparseBsr, "sparse_bsr", "torch.sparse_bsr");
+  registerLayout(
+      torch_module, at::Layout::SparseBsc, "sparse_bsc", "torch.sparse_bsc");
+  registerLayout(torch_module, at::Layout::Mkldnn, "_mkldnn", "torch._mkldnn");
+  registerLayout(torch_module, at::Layout::Jagged, "jagged", "torch.jagged");
 }
 
 } // namespace torch::utils
diff --git a/torch/csrc/utils/tensor_memoryformats.cpp b/torch/csrc/utils/tensor_memoryformats.cpp
index c1a3ff326493a..0ba5c8bd842c5 100644
--- a/torch/csrc/utils/tensor_memoryformats.cpp
+++ b/torch/csrc/utils/tensor_memoryformats.cpp
@@ -31,14 +31,12 @@ void initializeMemoryFormats() {
 
   auto add_memory_format = [&](at::MemoryFormat format, const char* name) {
     std::string module_name = "torch.";
-    PyObject* memory_format = THPMemoryFormat_New(format, module_name + name);
-    Py_INCREF(memory_format);
-    if (PyModule_AddObject(torch_module, name, memory_format) != 0) {
-      Py_DECREF(memory_format);
+    THPObjectPtr memory_format(THPMemoryFormat_New(format, module_name + name));
+    if (PyModule_AddObjectRef(torch_module, name, memory_format.get()) != 0) {
       throw python_error();
     }
-    Py_INCREF(memory_format);
-    memory_format_registry[static_cast<size_t>(format)] = memory_format;
+    memory_format_registry[static_cast<size_t>(format)] =
+        memory_format.release();
   };
 
   add_memory_format(at::MemoryFormat::Preserve, "preserve_format");
diff --git a/torch/csrc/utils/tensor_new.cpp b/torch/csrc/utils/tensor_new.cpp
index 8b758df7224a9..97ae280dfe332 100644
--- a/torch/csrc/utils/tensor_new.cpp
+++ b/torch/csrc/utils/tensor_new.cpp
@@ -1382,26 +1382,16 @@ static void _validate_sparse_compressed_tensor_args_template(
     ARG_SIZE,
     ARGS_COUNT
   };
-  static std::string sig;
-  switch (required_layout) {
-    case c10::Layout::SparseCsr:
-      sig =
-          "_validate_sparse_csr_tensor(PyObject* crow_indices, PyObject* col_indices, PyObject* values, IntArrayRef size)";
-      break;
-    case c10::Layout::SparseCsc:
-      sig =
-          "_validate_sparse_csc_tensor(PyObject* ccol_indices, PyObject* row_indices, PyObject* values, IntArrayRef size)";
-      break;
-    case c10::Layout::SparseBsr:
-      sig =
-          "_validate_sparse_bsr_tensor(PyObject* crow_indices, PyObject* col_indices, PyObject* values, IntArrayRef size)";
-      break;
-    case c10::Layout::SparseBsc:
-      sig =
-          "_validate_sparse_bsc_tensor(PyObject* ccol_indices, PyObject* row_indices, PyObject* values, IntArrayRef size)";
-      break;
-    default:;
-  }
+  constexpr const char* sig = [] {
+    if constexpr (required_layout == c10::Layout::SparseCsr)
+      return "_validate_sparse_csr_tensor(PyObject* crow_indices, PyObject* col_indices, PyObject* values, IntArrayRef size)";
+    else if constexpr (required_layout == c10::Layout::SparseCsc)
+      return "_validate_sparse_csc_tensor(PyObject* ccol_indices, PyObject* row_indices, PyObject* values, IntArrayRef size)";
+    else if constexpr (required_layout == c10::Layout::SparseBsr)
+      return "_validate_sparse_bsr_tensor(PyObject* crow_indices, PyObject* col_indices, PyObject* values, IntArrayRef size)";
+    else if constexpr (required_layout == c10::Layout::SparseBsc)
+      return "_validate_sparse_bsc_tensor(PyObject* ccol_indices, PyObject* row_indices, PyObject* values, IntArrayRef size)";
+  }();
   static PythonArgParser parser({sig});
 
   ParsedArgs<ARGS_COUNT> parsed_args;
@@ -1744,7 +1734,7 @@ Tensor asarray(
     std::optional<ScalarType> dtype,
     std::optional<Device> device,
     std::optional<bool> copy,
-    bool requires_grad) {
+    std::optional<bool> requires_grad) {
   Tensor tensor;
 
   bool force_copy = copy.value_or(false);
@@ -1761,9 +1751,17 @@ Tensor asarray(
   if (THPVariable_Check(obj)) {
     tensor = THPVariable_Unpack(obj);
   }
+  bool return_requires_grad =
+      requires_grad.value_or(tensor.defined() ? tensor.requires_grad() : false);
+  if (return_requires_grad && !requires_grad) {
+    TORCH_WARN_ONCE(
+        "torch.asarray: unspecified requires_grad now defaults to obj.requires_grad "
+        "instead of False. Pass requires_grad=False explicitly to get the old behavior "
+        "and silence this warning.")
+  }
 
 #ifdef USE_NUMPY
-  if (is_numpy_available()) {
+  if (!tensor.defined() && is_numpy_available()) {
     // Check whether 'obj' is a NumPy Array or Scalar.
     bool is_numpy_array = PyArray_Check(obj);
     bool is_numpy_scalar = PyArray_CheckScalar(obj);
@@ -1809,7 +1807,8 @@ Tensor asarray(
 
   // Check whether 'obj' implements the buffer protocol
   if (!tensor.defined() && PyObject_CheckBuffer(obj) != 0) {
-    tensor = tensor_frombuffer(obj, dtype_unwrapped, -1, 0, requires_grad);
+    tensor =
+        tensor_frombuffer(obj, dtype_unwrapped, -1, 0, return_requires_grad);
   }
 
   if (tensor.defined()) {
@@ -1857,10 +1856,10 @@ Tensor asarray(
 
     // Setting 'requires_grad' when the tensor is not a leaf does not work.
     // Whenever that happens, we have to use 'detach'.
-    if (!tensor.is_leaf() && !requires_grad) {
+    if (!tensor.is_leaf() && !return_requires_grad) {
       tensor = tensor.detach();
     } else {
-      tensor.set_requires_grad(requires_grad);
+      tensor.set_requires_grad(return_requires_grad);
     }
   } else {
     // Undefined tensor means it does not implement neither DLPack nor
@@ -1881,7 +1880,7 @@ Tensor asarray(
         /* copy_variables = */ false,
         /* copy_numpy = */ false,
         /* type_inference = */ !dtype.has_value());
-    tensor.set_requires_grad(requires_grad);
+    tensor.set_requires_grad(return_requires_grad);
   }
 
   return tensor;
diff --git a/torch/csrc/utils/tensor_new.h b/torch/csrc/utils/tensor_new.h
index 8ae71fcde4cfb..903645c40dc1d 100644
--- a/torch/csrc/utils/tensor_new.h
+++ b/torch/csrc/utils/tensor_new.h
@@ -132,5 +132,5 @@ at::Tensor asarray(
     std::optional<c10::ScalarType> dtype,
     std::optional<c10::Device> device,
     std::optional<bool> copy,
-    bool requires_grad);
+    std::optional<bool> requires_grad);
 } // namespace torch::utils
diff --git a/torch/csrc/utils/tensor_qschemes.cpp b/torch/csrc/utils/tensor_qschemes.cpp
index f85d091bd57a0..6c86be1ff7f8a 100644
--- a/torch/csrc/utils/tensor_qschemes.cpp
+++ b/torch/csrc/utils/tensor_qschemes.cpp
@@ -20,13 +20,12 @@ void initializeQSchemes() {
 
   for (const auto i : c10::irange(at::COMPILE_TIME_NUM_QSCHEMES)) {
     auto qscheme = static_cast<at::QScheme>(i);
-    PyObject* qscheme_obj = THPQScheme_New(qscheme, toString(qscheme));
-    thp_qscheme_array[static_cast<int>(qscheme)] = qscheme_obj;
-    Py_INCREF(qscheme_obj);
-    if (PyModule_AddObject(
-            torch_module, toString(qscheme).c_str(), qscheme_obj) != 0) {
+    THPObjectPtr qscheme_obj(THPQScheme_New(qscheme, toString(qscheme)));
+    if (PyModule_AddObjectRef(
+            torch_module, toString(qscheme).c_str(), qscheme_obj.get()) != 0) {
       throw python_error();
     }
+    thp_qscheme_array[static_cast<int>(qscheme)] = qscheme_obj.release();
   }
 }
 
diff --git a/torch/csrc/utils/tensor_types.cpp b/torch/csrc/utils/tensor_types.cpp
index 620086f9ad50d..02414c8494c58 100644
--- a/torch/csrc/utils/tensor_types.cpp
+++ b/torch/csrc/utils/tensor_types.cpp
@@ -6,7 +6,6 @@
 #include <torch/csrc/autograd/generated/VariableType.h>
 #include <torch/csrc/tensor/python_tensor.h>
 
-#include <algorithm>
 #include <sstream>
 #include <unordered_map>
 
@@ -76,20 +75,21 @@ std::string type_to_string(const at::DeprecatedTypeProperties& type) {
   return ss.str();
 }
 
+using TypeMap = std::unordered_map<std::string, at::DeprecatedTypeProperties*>;
+
+static TypeMap build_type_map(
+    const std::vector<at::DeprecatedTypeProperties*>& types) {
+  TypeMap m;
+  m.reserve(types.size());
+  for (auto type : types)
+    m.emplace(type_to_string(*type), type);
+  return m;
+}
+
 at::TensorOptions options_from_string(const std::string& str) {
-  static std::string cuda_prefix("torch.cuda.");
-  static std::string xpu_prefix("torch.xpu.");
-  static std::string privateUser_prefix(
-      std::string(parse_privateuseone_backend()) + ".");
-  static std::unordered_map<std::string, at::DeprecatedTypeProperties*> cpu_map;
-  static std::unordered_map<std::string, at::DeprecatedTypeProperties*> xpu_map;
-  static std::unordered_map<std::string, at::DeprecatedTypeProperties*>
-      cuda_map;
-  static std::unordered_map<std::string, at::DeprecatedTypeProperties*>
-      privateUser1_map;
-
-  const std::unordered_map<std::string, at::DeprecatedTypeProperties*>* map =
-      nullptr;
+  static const std::string privateUser_prefix =
+      std::string(parse_privateuseone_backend()) + ".";
+  const TypeMap* map = nullptr;
 
   if (str == "torch.Tensor") {
     auto backend =
@@ -98,46 +98,21 @@ at::TensorOptions options_from_string(const std::string& str) {
     return getDeprecatedTypeProperties(backend, scalar_type).options();
   }
 
-  if (std::mismatch(cuda_prefix.begin(), cuda_prefix.end(), str.begin())
-          .first == cuda_prefix.end()) {
-    // torch.cuda. is prefix of str
-    static bool cuda_once [[maybe_unused]] = []() {
-      for (auto type : autograd::VariableType::allCUDATypes()) {
-        cuda_map.emplace(type_to_string(*type), type);
-      }
-      return true;
-    }();
+  if (str.starts_with("torch.cuda.")) {
+    static const auto cuda_map =
+        build_type_map(autograd::VariableType::allCUDATypes());
     map = &cuda_map;
-  } else if (
-      std::mismatch(xpu_prefix.begin(), xpu_prefix.end(), str.begin()).first ==
-      xpu_prefix.end()) {
-    // torch.xpu. is prefix of str
-    static bool xpu_once [[maybe_unused]] = []() {
-      for (auto type : autograd::VariableType::allXPUTypes()) {
-        xpu_map.emplace(type_to_string(*type), type);
-      }
-      return true;
-    }();
+  } else if (str.starts_with("torch.xpu.")) {
+    static const auto xpu_map =
+        build_type_map(autograd::VariableType::allXPUTypes());
     map = &xpu_map;
-  } else if (
-      std::mismatch(
-          privateUser_prefix.begin(), privateUser_prefix.end(), str.begin())
-          .first == privateUser_prefix.end()) {
-    // torch.foo. foo is privateUser1 name
-    static bool privateUser1_once [[maybe_unused]] = []() {
-      for (auto type : autograd::VariableType::allPrivateUser1Types()) {
-        privateUser1_map.emplace(type_to_string(*type), type);
-      }
-      return true;
-    }();
+  } else if (str.starts_with(privateUser_prefix)) {
+    static const auto privateUser1_map =
+        build_type_map(autograd::VariableType::allPrivateUser1Types());
     map = &privateUser1_map;
   } else {
-    static bool cpu_once [[maybe_unused]] = []() {
-      for (auto type : autograd::VariableType::allCPUTypes()) {
-        cpu_map.emplace(type_to_string(*type), type);
-      }
-      return true;
-    }();
+    static const auto cpu_map =
+        build_type_map(autograd::VariableType::allCPUTypes());
     map = &cpu_map;
   }
 
diff --git a/torch/csrc/xpu/Module.cpp b/torch/csrc/xpu/Module.cpp
index 0834b7d3e4f7c..c3dd26809ab5c 100644
--- a/torch/csrc/xpu/Module.cpp
+++ b/torch/csrc/xpu/Module.cpp
@@ -336,6 +336,8 @@ static void registerXpuDeviceProperties(PyObject* module) {
       ._(gpu_eu_count)                                           \
       ._(max_work_group_size)                                    \
       ._(max_num_sub_groups)                                     \
+      ._(memory_clock_rate)                                      \
+      ._(memory_bus_width)                                       \
       ._(sub_group_sizes)                                        \
       ._(local_mem_size)                                         \
       ._(has_fp16)                                               \
@@ -371,7 +373,9 @@ static void registerXpuDeviceProperties(PyObject* module) {
                    << prop.global_mem_size / (1024ull * 1024)
                    << "MB, local_mem_size=" << prop.local_mem_size / 1024ull
                    << "KB, max_compute_units=" << prop.max_compute_units
-                   << ", gpu_eu_count=" << prop.gpu_eu_count
+                   << ", memory_clock_rate=" << prop.memory_clock_rate
+                   << "MHz, memory_bus_width=" << prop.memory_bus_width
+                   << "-bit, gpu_eu_count=" << prop.gpu_eu_count
                    << ", gpu_subslice_count=" << gpu_subslice_count(prop)
                    << ", max_work_group_size=" << prop.max_work_group_size
                    << ", max_num_sub_groups=" << prop.max_num_sub_groups
diff --git a/torch/cuda/__init__.py b/torch/cuda/__init__.py
index 894760ebd67b5..e36e2b44acc5d 100644
--- a/torch/cuda/__init__.py
+++ b/torch/cuda/__init__.py
@@ -18,7 +18,7 @@
 import warnings
 from collections.abc import Callable
 from functools import lru_cache
-from typing import Any, cast, NewType, Optional, TYPE_CHECKING, Union
+from typing import Any, cast, NewType, Optional, TYPE_CHECKING
 
 import torch
 import torch._C
@@ -52,7 +52,9 @@
 _is_in_bad_fork = getattr(torch._C, "_cuda_isInBadFork", lambda: False)
 
 _HAS_PYNVML = False
+_HAS_AMDSMI = False
 _PYNVML_ERR = None
+_AMDSMI_ERR = None
 try:
     from torch import version as _version
 
@@ -104,8 +106,13 @@ def __enter__(self) -> None:
                 def __exit__(self, type: Any, value: Any, traceback: Any) -> None:
                     ctypes.CDLL = self.original_CDLL  # type: ignore[misc]
 
-            with _amdsmi_cdll_hook():
-                import amdsmi  # type: ignore[import]
+            try:
+                with _amdsmi_cdll_hook():
+                    import amdsmi  # type: ignore[import]
+                _HAS_AMDSMI = True
+            except ModuleNotFoundError as err:
+                _AMDSMI_ERR = err
+                raise
 
         _HAS_PYNVML = True
     except ModuleNotFoundError:
@@ -245,7 +252,7 @@ class _CompatInterval:
     also allows excluding specific versions from the range.
     """
 
-    def __init__(self, start, exclude: Optional[set[int]] = None):
+    def __init__(self, start, exclude: set[int] | None = None):
         self.major, self.minor = start // 10, start % 10
         self.exclude = set() if exclude is None else exclude
 
@@ -288,7 +295,7 @@ def __str__(self):
 # - The keys in dict correspond to known sm versions but the values
 #   are merely rules based on sm compatibility guarantees for NVIDIA
 #   devices while accounting for incompatibility of iGPU and dGPU.
-DEVICE_REQUIREMENT: dict[int, Union[_CompatSet, _CompatInterval]] = {
+DEVICE_REQUIREMENT: dict[int, _CompatSet | _CompatInterval] = {
     50: _CompatInterval(start=50, exclude={53}),
     52: _CompatInterval(start=52, exclude={53}),
     53: _CompatSet({53}),
@@ -312,11 +319,13 @@ def __str__(self):
 }
 
 
-# TORCH_CUDA_ARCH_LIST for PyTorch releases
+# TORCH_CUDA_ARCH_LIST for PyTorch releases (union of x86_64 and aarch64
+# build matrices). Keep in sync with .ci/manywheel/build_cuda.sh and
+# .github/scripts/generate_binary_build_matrix.py:CUDA_ARCHES.
 PYTORCH_RELEASES_CODE_CC: dict[str, set[int]] = {
-    "12.6": {50, 60, 70, 80, 86, 90},
-    "12.8": {70, 80, 86, 90, 100, 120},
+    "12.6": {50, 60, 70, 75, 80, 86, 90},
     "13.0": {75, 80, 86, 90, 100, 110, 120},
+    "13.2": {75, 80, 86, 90, 100, 110, 120},
 }
 
 
@@ -361,7 +370,11 @@ def _check_capability():
     if torch.version.cuda is None:  # on ROCm we don't want this check
         return
 
-    code_ccs = [_extract_arch_version(cc) for cc in get_arch_list()]
+    arch_list = get_arch_list()
+    if len(arch_list) == 0:
+        return
+
+    code_ccs = [_extract_arch_version(cc) for cc in arch_list]
     for d in range(device_count()):
         major, minor = get_device_capability(d)
         device_cc = 10 * major + minor
@@ -404,9 +417,16 @@ def is_initialized():
 
 
 def _lazy_call(callable, **kwargs):
+    # Do not invoke user callbacks while holding _initialization_lock
+    # they may call back into _lazy_call.
+    if is_initialized():
+        callable()
+        return
+
+    run_now = False
     with _initialization_lock:
         if is_initialized():
-            callable()
+            run_now = True
         else:
             # TODO(torch_deploy): this accesses linecache, which attempts to read the
             # file system to get traceback info. Patch linecache or do something
@@ -420,6 +440,9 @@ def _lazy_call(callable, **kwargs):
                 # Don't store the actual traceback to avoid memory cycle
                 _queued_calls.append((callable, traceback.format_stack()))
 
+    if run_now:
+        callable()
+
 
 _lazy_call(_check_capability)
 _lazy_call(_check_cubins)
@@ -879,7 +902,7 @@ def parse_list_with_prefix(lst: str, prefix: str) -> list[str]:
 
 
 def _raw_device_count_amdsmi() -> int:
-    if not _HAS_PYNVML:  # If amdsmi is not available
+    if not _HAS_AMDSMI:
         return -1
     try:
         amdsmi.amdsmi_init()
@@ -913,7 +936,7 @@ def _raw_device_count_nvml() -> int:
 def _raw_device_uuid_amdsmi() -> list[str] | None:
     from ctypes import byref, c_int, c_void_p, CDLL, create_string_buffer
 
-    if not _HAS_PYNVML:  # If amdsmi is not available
+    if not _HAS_AMDSMI:
         return None
     try:
         amdsmi.amdsmi_init()
@@ -1108,9 +1131,15 @@ def device_count() -> int:
         return 0
     if _cached_device_count is not None:
         return _cached_device_count
-    # bypass _device_count_nvml() if rocm (not supported)
-    nvml_count = _device_count_amdsmi() if torch.version.hip else _device_count_nvml()
-    r = torch._C._cuda_getDeviceCount() if nvml_count < 0 else nvml_count
+    if _initialized or hasattr(_tls, "is_initializing"):
+        r = torch._C._cuda_getDeviceCount()
+    else:
+        # bypass _device_count_nvml() if rocm (not supported)
+        if torch.version.hip:
+            nvml_count = _device_count_amdsmi()
+        else:
+            nvml_count = _device_count_nvml()
+        r = torch._C._cuda_getDeviceCount() if nvml_count < 0 else nvml_count
     # NB: Do not cache the device count prior to CUDA initialization, because
     # the number of devices can change due to changes to CUDA_VISIBLE_DEVICES
     # setting prior to CUDA initialization.
@@ -1296,11 +1325,10 @@ def _get_pynvml_handler(device: Device = None):
 
 
 def _get_amdsmi_handler(device: Device = None):
-    if not _HAS_PYNVML:
+    if not _HAS_AMDSMI:
         raise ModuleNotFoundError(
             "amdsmi does not seem to be installed or it can't be imported."
-            # pyrefly: ignore [invalid-inheritance]
-        ) from _PYNVML_ERR
+        ) from _AMDSMI_ERR
     try:
         amdsmi.amdsmi_init()
     except amdsmi.AmdSmiException as e:
@@ -1312,11 +1340,49 @@ def _get_amdsmi_handler(device: Device = None):
     return handle
 
 
+_cached_hip_to_amdsmi: dict[int, int] | None = None
+
+
+def _get_amdsmi_device_index_from_hip_index(device: int) -> int:
+    r"""Return amdsmi index from HIP device index. They are not always the same.
+
+    Assume amdsmi_init() already completes successfully."""
+    global _cached_hip_to_amdsmi
+    if _cached_hip_to_amdsmi is None:
+        amdsmi_handles = amdsmi.amdsmi_get_processor_handles()
+
+        def gen():
+            for amdsmi_idx, handle in enumerate(amdsmi_handles):
+                info = amdsmi.amdsmi_get_gpu_enumeration_info(handle)
+                if "hip_id" in info:
+                    yield info["hip_id"], amdsmi_idx
+
+        _cached_hip_to_amdsmi = dict(gen())
+        if not _cached_hip_to_amdsmi and len(amdsmi_handles) > 1:
+            warnings.warn(
+                "Cannot translate HIP ID to AMD SMI ID due to"
+                " lack of translation information prior to ROCm 6.4."
+                " Functions that rely on amdsmi"
+                " (e.g. temperature()) may operate on wrong devices."
+            )
+    if device not in _cached_hip_to_amdsmi:
+        warnings.warn(
+            f"Cannot translate HIP ID {device} to AMD SMI ID due to"
+            " undetected HIP ID from amdsmi."
+            " amdsmi_get_gpu_enumeration_info() only report these HIP IDs"
+            f" {list(_cached_hip_to_amdsmi.keys())}."
+            " Functions that rely on amdsmi"
+            " (e.g. temperature()) may operate on wrong devices."
+        )
+    return _cached_hip_to_amdsmi.get(device, device)
+
+
 def _get_amdsmi_device_index(device: Device) -> int:
     r"""Return the amdsmi index of the device, taking visible_devices into account."""
     idx = _get_device_index(device, optional=True)
     visible_devices = _parse_visible_devices()
-    if type(visible_devices[0]) is str:
+    visible_device_is_str = type(visible_devices[0]) is str
+    if visible_device_is_str:
         uuids = _raw_device_uuid_amdsmi()
         if uuids is None:
             raise RuntimeError("Can't get device UUIDs")
@@ -1329,7 +1395,10 @@ def _get_amdsmi_device_index(device: Device) -> int:
         raise RuntimeError(
             f"device {idx} is not visible (HIP_VISIBLE_DEVICES={visible_devices})"
         )
-    return idx_map[idx]
+    if visible_device_is_str:
+        return idx_map[idx]
+    else:
+        return _get_amdsmi_device_index_from_hip_index(idx_map[idx])
 
 
 def _get_amdsmi_device_memory_used(device: Device = None) -> int:
@@ -1927,6 +1996,7 @@ def _compile_kernel(
     "amp",
     "caching_allocator_alloc",
     "caching_allocator_delete",
+    "caching_allocator_disabled",
     "caching_allocator_enable",
     "can_device_access_peer",
     "check_error",
diff --git a/torch/cuda/_annotate_cuda_graph_trace.py b/torch/cuda/_annotate_cuda_graph_trace.py
new file mode 100644
index 0000000000000..92a0c665a027c
--- /dev/null
+++ b/torch/cuda/_annotate_cuda_graph_trace.py
@@ -0,0 +1,362 @@
+#!/usr/bin/env python3
+"""Post-process a profiler trace to add CUDA graph kernel annotations.
+
+Reads a profiler trace (gzipped or plain JSON) and a kernel annotations
+pickle, matches kernel events by their graph node id, and writes an
+annotated trace with the annotation fields added to each kernel event's
+args (displayed alongside grid/block size in trace viewers).
+
+The annotations pickle is auto-discovered from the trace file's parent
+directory (one level up, matching the rank from the trace filename).
+
+Usage:
+    python -m torch.cuda._annotate_cuda_graph_trace <trace_file> [-a <annotations_pkl>] [-o <output_file>]
+
+Examples:
+    # Auto-discover annotations pickle from trace location
+    python -m torch.cuda._annotate_cuda_graph_trace \\
+        traces/step_000000000014/000000.*.pt.trace.json.gz
+
+    # Explicit annotations pickle
+    python -m torch.cuda._annotate_cuda_graph_trace trace.json.gz -a annotations.pkl
+"""
+
+import argparse
+import gzip
+import json
+import pickle
+import re
+import sys
+from collections import defaultdict
+from pathlib import Path
+from typing import Any
+
+
+_WORK_CATEGORIES = {"kernel", "gpu_memcpy", "gpu_memset"}
+
+
+def _move_overlapping_to_stream(
+    trace: dict, default_stream: int = 7, overlap_stream: int = 8
+) -> int:
+    """Move graphed kernels that overlap with their predecessor to a separate stream.
+
+    Perfetto cannot display overlapping (non-nested) events on the same
+    stream -- they get hidden.  This pass detects graphed kernel events on
+    *default_stream* whose start timestamp falls before the previous
+    kernel's end, and moves them to *overlap_stream* so they're visible.
+
+    Returns the number of events moved.
+    """
+    graphed_on_default = [
+        e
+        for e in trace["traceEvents"]
+        if e.get("cat") == "kernel"
+        and e.get("tid") == default_stream
+        and e.get("args", {}).get("graph node id", 0) != 0
+    ]
+    graphed_on_default.sort(key=lambda e: e["ts"])
+
+    moved = 0
+    prev_end = 0.0
+    for event in graphed_on_default:
+        ts = event["ts"]
+        dur = event.get("dur", 0)
+        if ts < prev_end:
+            event["tid"] = overlap_stream
+            event.get("args", {})["stream"] = overlap_stream
+            moved += 1
+        else:
+            prev_end = ts + dur
+
+    return moved
+
+
+def _fix_overlapping_timestamps(trace: dict, max_adjust_us: float = 1.0) -> int:
+    """Clamp graphed kernel/memcpy timestamps so they don't overlap on the same stream.
+
+    CUPTI can produce slightly overlapping timestamps for consecutive graphed
+    events, causing Perfetto to hide events that sit entirely "under" their
+    neighbours.  This pass sorts graphed work events per stream and ensures
+    each event starts at or after the previous event's end.
+
+    Overlaps larger than *max_adjust_us* are flagged as warnings and left
+    unchanged, since they likely indicate a real issue rather than CUPTI
+    timestamp jitter.
+
+    Returns the number of events adjusted.
+    """
+    per_stream: dict[int, list[dict]] = defaultdict(list)
+    for event in trace["traceEvents"]:
+        if (
+            event.get("cat") in _WORK_CATEGORIES
+            and event.get("args", {}).get("graph node id", 0) != 0
+        ):
+            per_stream[event.get("tid")].append(event)
+
+    adjusted = 0
+    for tid, events in per_stream.items():
+        events.sort(key=lambda e: e["ts"])
+        prev_end = 0.0
+        for event in events:
+            ts = event["ts"]
+            dur = event.get("dur", 0)
+            if ts < prev_end:
+                overlap = prev_end - ts
+                if overlap > max_adjust_us:
+                    print(
+                        f"WARNING: large overlap {overlap:.3f}us on stream {tid} "
+                        f"for {event.get('name', '?')[:60]}, skipping adjustment"
+                    )
+                else:
+                    event["ts"] = prev_end
+                    adjusted += 1
+            prev_end = event["ts"] + dur
+
+    return adjusted
+
+
+def annotate_trace(
+    trace: dict,
+    annotations: dict[int, list[Any]],
+    default_stream: int = 7,
+) -> int:
+    """Add annotation fields to kernel events matching the annotations dict.
+
+    Each annotation entry is a list (from nested ``mark_kernels`` scopes).
+    Fields from all annotations are merged into the event args; if multiple
+    annotations define the same key, later entries in the list win.
+
+    For graphed events (graph_node_id != 0), reassigns ``tid`` and
+    ``args["stream"]`` to the stream recorded in annotations, or to
+    *default_stream* if there is no annotation.  Also moves the
+    corresponding ``ac2g`` flow-finish events to the new tid so that
+    CPU-to-GPU correlation arrows are preserved.
+
+    Removes ``gpu_user_annotation`` events and orphaned ``ac2g`` events
+    from streams that have no kernel or memcpy events after reassignment,
+    since CUPTI replicates these onto every stream during graph replay.
+
+    Returns the number of events annotated.
+    """
+    # Build an index of ac2g 'f' events keyed by (tid, ts) so we can
+    # move them together with the kernel events they correspond to.
+    ac2g_f_index: dict[tuple, list] = {}
+    for event in trace["traceEvents"]:
+        if event.get("cat") == "ac2g" and event.get("ph") == "f":
+            key = (event.get("tid"), event.get("ts"))
+            ac2g_f_index.setdefault(key, []).append(event)
+
+    annotated = 0
+    for event in trace.get("traceEvents", []):
+        args = event.get("args", {})
+        graph_node_id = args.get("graph node id")
+        if graph_node_id is None or graph_node_id == 0:
+            continue
+        stream_id = None
+        if graph_node_id in annotations:
+            for ann in annotations[graph_node_id]:
+                if isinstance(ann, dict):
+                    for key, value in ann.items():
+                        args[key] = str(value)
+                    if "stream" in ann:
+                        stream_id = int(ann["stream"])
+                else:
+                    args["annotation"] = str(ann)
+            annotated += 1
+
+        # Reassign stream: use annotated stream if available, else default
+        if stream_id is None:
+            stream_id = default_stream
+        old_key = (event.get("tid"), event.get("ts"))
+        event["tid"] = stream_id
+        args["stream"] = stream_id
+
+        # Move the matching ac2g 'f' event(s) to the same new tid
+        for ac2g_event in ac2g_f_index.get(old_key, ()):
+            ac2g_event["tid"] = stream_id
+
+    # Remove gpu_user_annotation events and ac2g flow-finish events from
+    # streams that have no real kernel/memcpy/memset work -- these are
+    # noise replicated by CUPTI onto every stream during graph replay.
+    tids_with_work = set()
+    for event in trace["traceEvents"]:
+        if event.get("cat") in _WORK_CATEGORIES:
+            tids_with_work.add(event.get("tid"))
+
+    def _is_noise(event: dict) -> bool:
+        cat = event.get("cat")
+        if cat == "gpu_user_annotation":
+            return event.get("tid") not in tids_with_work
+        if cat == "ac2g" and event.get("ph") == "f":
+            return event.get("tid") not in tids_with_work
+        return False
+
+    original_count = len(trace["traceEvents"])
+    trace["traceEvents"] = [
+        event for event in trace["traceEvents"] if not _is_noise(event)
+    ]
+    removed = original_count - len(trace["traceEvents"])
+    if removed:
+        print(f"Removed {removed} noise events from empty streams")
+
+    # Clean up metadata: remove thread_name / thread_sort_index entries
+    # for noise streams that have no real (non-metadata) events, and add
+    # thread_name entries for our new annotation streams.
+    all_tids_in_trace = {
+        e.get("tid") for e in trace["traceEvents"] if e.get("ph") != "M"
+    }
+    # Find the GPU process pid from existing thread_name metadata
+    gpu_pid = 0
+    for event in trace["traceEvents"]:
+        if (
+            event.get("ph") == "M"
+            and event.get("name") == "thread_name"
+            and str(event.get("args", {}).get("name", "")).startswith("stream ")
+        ):
+            gpu_pid = event.get("pid", 0)
+            break
+
+    # Remove metadata entries for tids with no non-metadata events
+    trace["traceEvents"] = [
+        event
+        for event in trace["traceEvents"]
+        if event.get("ph") != "M" or event.get("tid") in all_tids_in_trace
+    ]
+
+    # Add thread_name metadata for new annotation tids that lack one
+    existing_thread_names = {
+        e.get("tid")
+        for e in trace["traceEvents"]
+        if e.get("ph") == "M" and e.get("name") == "thread_name"
+    }
+    for tid in sorted(tids_with_work - existing_thread_names):
+        trace["traceEvents"].append(
+            {
+                "ph": "M",
+                "pid": gpu_pid,
+                "tid": tid,
+                "name": "thread_name",
+                "args": {"name": f"stream {tid}"},
+            }
+        )
+
+    return annotated
+
+
+def load_trace(path: Path) -> dict:
+    if path.suffix == ".gz" or path.name.endswith(".json.gz"):
+        with gzip.open(path, "rt") as f:
+            return json.load(f)
+    else:
+        with open(path) as f:
+            return json.load(f)
+
+
+def save_trace(trace: dict, path: Path) -> None:
+    if path.suffix == ".gz" or path.name.endswith(".json.gz"):
+        with gzip.open(path, "wt") as f:
+            json.dump(trace, f)
+    else:
+        with open(path, "w") as f:
+            json.dump(trace, f)
+
+
+def _find_annotations_pkl(trace_file: Path) -> Path | None:
+    """Auto-discover the annotations pickle from the trace file location.
+
+    Trace files live in e.g. ``traces/step_000000000014/000000.<id>.pt.trace.json.gz``
+    where the leading digits are the rank. The pickle lives one level up:
+    ``traces/kernel_annotations_rank0_*.pkl``.
+    """
+    match = re.match(r"^(\d+)", trace_file.name)
+    if not match:
+        return None
+    rank = int(match.group(1))
+
+    traces_dir = trace_file.parent.parent
+    candidates = sorted(traces_dir.glob(f"kernel_annotations_rank{rank}_*.pkl"))
+    if candidates:
+        return candidates[0]
+    return None
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(
+        description="Annotate a profiler trace with CUDA graph kernel annotations."
+    )
+    parser.add_argument(
+        "trace_file", type=Path, help="Input trace file (.json or .json.gz)"
+    )
+    parser.add_argument(
+        "-a",
+        "--annotations",
+        type=Path,
+        default=None,
+        help="Kernel annotations pickle file. Auto-discovered if omitted.",
+    )
+    parser.add_argument(
+        "-o",
+        "--output",
+        type=Path,
+        default=None,
+        help="Output file path. Defaults to <trace_file>.annotated.<ext>",
+    )
+    parser.add_argument(
+        "--default-stream",
+        type=int,
+        default=7,
+        help="Stream ID to assign to unannotated graphed events (default: 7).",
+    )
+    args = parser.parse_args()
+
+    annotations_pkl = args.annotations
+    if annotations_pkl is None:
+        annotations_pkl = _find_annotations_pkl(args.trace_file)
+        if annotations_pkl is None:
+            print(
+                f"Could not auto-discover annotations pickle for {args.trace_file}. "
+                f"Use -a to specify it explicitly.",
+                file=sys.stderr,
+            )
+            sys.exit(1)
+        print(f"Auto-discovered annotations: {annotations_pkl}")
+
+    with open(annotations_pkl, "rb") as f:
+        annotations = pickle.load(f)
+    print(f"Loaded {len(annotations)} kernel annotations")
+
+    trace = load_trace(args.trace_file)
+    total_events = len(trace.get("traceEvents", []))
+    print(f"Loaded trace with {total_events} events")
+
+    count = annotate_trace(trace, annotations, default_stream=args.default_stream)
+    print(f"Annotated {count} kernel events")
+
+    overlap_moved = _move_overlapping_to_stream(
+        trace, default_stream=args.default_stream
+    )
+    if overlap_moved:
+        print(f"Moved {overlap_moved} overlapping events to stream 8")
+
+    ts_fixed = _fix_overlapping_timestamps(trace)
+    if ts_fixed:
+        print(f"Fixed {ts_fixed} overlapping graphed event timestamps")
+
+    output = args.output
+    if output is None:
+        name = args.trace_file.name
+        if name.endswith(".json.gz"):
+            output = args.trace_file.with_name(
+                name.replace(".json.gz", ".annotated.json.gz")
+            )
+        elif name.endswith(".json"):
+            output = args.trace_file.with_suffix(".annotated.json")
+        else:
+            output = args.trace_file.with_suffix(args.trace_file.suffix + ".annotated")
+
+    save_trace(trace, output)
+    print(f"Saved annotated trace to {output}")
+
+
+if __name__ == "__main__":
+    main()
diff --git a/torch/cuda/_graph_annotations.py b/torch/cuda/_graph_annotations.py
new file mode 100644
index 0000000000000..16a895686cf41
--- /dev/null
+++ b/torch/cuda/_graph_annotations.py
@@ -0,0 +1,434 @@
+"""Annotate CUDA graph kernel nodes during capture.
+
+During CUDA graph capture, ``mark_kernels`` uses ``cudaGraphGetNodes``
+to count nodes before and after the wrapped region.  Nodes at indices
+``[before, after)`` are the ones added within the scope.  Each kernel
+or memcpy node found is annotated by its ``toolsId`` so it can later
+be matched to profiler trace events.
+
+The annotations can be pickled and later merged into a Chrome profiler
+trace using ``torch.cuda._annotate_cuda_graph_trace``.
+
+Requires ``cuda.bindings`` package and a CUDA driver that supports
+``cudaGraphNodeGetToolsId`` (CUDA >= 13.1 or appropriate cuda-compat).
+When unavailable, ``mark_kernels`` silently becomes a no-op.
+
+Usage during capture::
+
+    from torch.cuda._graph_annotations import (
+        enable_annotations,
+        mark_kernels,
+        resolve_pending_annotations,
+        remap_to_exec_graph,
+    )
+
+    enable_annotations()
+
+    with torch.cuda.graph(graph):
+        with mark_kernels("phase_A"):
+            y = workload_a(x)
+        with mark_kernels("phase_B"):
+            z = workload_b(y)
+        resolve_pending_annotations()
+
+    remap_to_exec_graph(graph)
+"""
+
+from collections import defaultdict
+from contextlib import contextmanager
+from logging import getLogger
+from typing import Any
+
+import torch
+from torch.cuda._utils import _check_cuda_bindings, _HAS_CUDA_BINDINGS
+
+
+try:
+    from cuda.bindings import (  # pyrefly: ignore[missing-import]
+        runtime as _cuda_runtime,
+    )
+except ImportError:
+    _cuda_runtime = None  # type: ignore[assignment]
+
+
+logger = getLogger(__name__)
+
+
+# Tri-state: None = not probed, True = available, False = unavailable.
+# Deferred to first use to avoid premature CUDA initialization.
+_tools_id_available: bool | None = None
+
+# Global kill switch. When False, mark_kernels and mark_stream are no-ops.
+_annotations_enabled: bool = False
+
+
+def enable_annotations() -> None:
+    """Enable kernel annotation recording."""
+    global _annotations_enabled
+    _annotations_enabled = True
+
+
+def disable_annotations() -> None:
+    """Disable kernel annotation recording."""
+    global _annotations_enabled
+    _annotations_enabled = False
+
+
+def _is_tools_id_unavailable() -> bool:
+    """Return True if we already know cudaGraphNodeGetToolsId is missing."""
+    if not _HAS_CUDA_BINDINGS:
+        return True
+    if _tools_id_available is False:
+        return True
+    if not hasattr(_cuda_runtime, "cudaGraphNodeGetToolsId"):
+        return True
+    return False
+
+
+def _get_tools_id(node: Any) -> int | None:
+    """Return the toolsId for a graph node, or None if unavailable."""
+    global _tools_id_available
+    if _tools_id_available is None:
+        try:
+            tools_id = _check_cuda_bindings(
+                _cuda_runtime.cudaGraphNodeGetToolsId(  # pyrefly: ignore[missing-attribute]
+                    node
+                )
+            )
+        except Exception:
+            _tools_id_available = False
+            logger.info(
+                "cudaGraphNodeGetToolsId not available; "
+                "CUDA graph kernel annotations will be disabled"
+            )
+            return None
+        _tools_id_available = True
+        return tools_id
+    return _check_cuda_bindings(
+        _cuda_runtime.cudaGraphNodeGetToolsId(  # pyrefly: ignore[missing-attribute]
+            node
+        )
+    )
+
+
+def _get_capture_graph(stream: Any) -> Any:
+    """Return the graph handle for the active capture, or None."""
+    status, _id, graph, _deps, _edge_data, _num_deps = _check_cuda_bindings(
+        _cuda_runtime.cudaStreamGetCaptureInfo(  # pyrefly: ignore[missing-attribute]
+            stream
+        )
+    )
+    if (
+        status
+        != _cuda_runtime.cudaStreamCaptureStatus.cudaStreamCaptureStatusActive  # pyrefly: ignore[missing-attribute]
+    ):
+        return None
+    return graph
+
+
+def _get_node_count(graph: Any) -> int:
+    """Return the number of nodes currently in the graph."""
+    _, num = _check_cuda_bindings(
+        _cuda_runtime.cudaGraphGetNodes(  # pyrefly: ignore[missing-attribute]
+            graph, numNodes=0
+        )
+    )
+    return num
+
+
+# toolsId -> list of annotation objects.
+_kernel_annotations: defaultdict[int, list[Any]] = defaultdict(list)
+
+# Node types we annotate. Initialized lazily to avoid touching cuda.bindings
+# at import time.
+_ANNOTATABLE_TYPES: set[Any] | None = None
+
+
+def _get_annotatable_types() -> set[Any]:
+    global _ANNOTATABLE_TYPES
+    if _ANNOTATABLE_TYPES is None:
+        _ANNOTATABLE_TYPES = {
+            _cuda_runtime.cudaGraphNodeType.cudaGraphNodeTypeKernel,  # pyrefly: ignore[missing-attribute]
+            _cuda_runtime.cudaGraphNodeType.cudaGraphNodeTypeMemcpy,  # pyrefly: ignore[missing-attribute]
+        }
+    return _ANNOTATABLE_TYPES
+
+
+# Pending scopes: (annotation, start_node_index, end_node_index).
+_pending_scopes: list[tuple[Any, int, int]] = []
+
+# Graph handle saved during capture for post-capture resolution.
+_capture_graph: Any = None
+
+# Capture graph ID saved by resolve_pending_annotations for remap_to_exec_graph.
+_last_capture_graph_id: int | None = None
+
+
+@contextmanager  # type: ignore[arg-type]
+def mark_kernels(annotation: str | dict[str, Any]):
+    """Context manager that records node index ranges for later annotation.
+
+    During capture, calls ``cudaGraphGetNodes`` to count graph nodes before
+    and after the scope.  Nodes at indices ``[before, after)`` were added
+    inside the scope.  After capture, ``resolve_pending_annotations``
+    enumerates all nodes and annotates kernel/memcpy nodes in those ranges.
+
+    Must be called inside an active ``torch.cuda.graph()`` capture.  If the
+    current stream is not capturing, or if ``cudaGraphNodeGetToolsId`` is not
+    available, the context manager is a no-op.
+
+    Args:
+        annotation: Arbitrary object appended to the annotation list for
+            every kernel/memcpy node whose index falls within this scope.
+    """
+    if not _annotations_enabled or _is_tools_id_unavailable():
+        yield
+        return
+
+    if isinstance(annotation, str):
+        annotation = {"str": annotation}
+
+    stream = _cuda_runtime.cudaStream_t(  # pyrefly: ignore[missing-attribute]
+        init_value=torch.cuda.current_stream().cuda_stream
+    )
+    graph = _get_capture_graph(stream)
+    if graph is None:
+        yield
+        return
+
+    global _capture_graph
+    _capture_graph = graph
+
+    start_count = _get_node_count(graph)
+
+    yield
+
+    end_count = _get_node_count(graph)
+
+    if end_count > start_count:
+        _pending_scopes.append((annotation, start_count, end_count))
+
+
+def resolve_pending_annotations() -> None:
+    """Resolve pending scope index ranges into kernel annotations.
+
+    Enumerates all graph nodes and annotates kernel/memcpy nodes whose
+    indices fall within recorded scope ranges. Must be called while still
+    inside the ``torch.cuda.graph()`` capture context.
+    """
+    global _capture_graph
+    if not _pending_scopes:
+        _capture_graph = None
+        return
+
+    # Get a fresh graph handle from the active capture.
+    stream = _cuda_runtime.cudaStream_t(  # pyrefly: ignore[missing-attribute]
+        init_value=torch.cuda.current_stream().cuda_stream
+    )
+    graph = _get_capture_graph(stream)
+    if graph is None:
+        graph = _capture_graph
+    if graph is None:
+        logger.warning("resolve_pending_annotations: no graph handle available")
+        _pending_scopes.clear()
+        return
+
+    try:
+        num = _get_node_count(graph)
+        if num == 0:
+            _pending_scopes.clear()
+            _capture_graph = None
+            return
+
+        nodes, num = _check_cuda_bindings(
+            _cuda_runtime.cudaGraphGetNodes(  # pyrefly: ignore[missing-attribute]
+                graph, numNodes=num
+            )
+        )
+
+        # Save capture graph ID for remap_to_exec_graph.
+        global _last_capture_graph_id
+        if num > 0:
+            first_tid = _get_tools_id(nodes[0])
+            _last_capture_graph_id = (first_tid >> 32) if first_tid else None
+
+        annotatable = _get_annotatable_types()
+
+        # Sort by (start, -end, -append_index). The append index encodes
+        # nesting depth: inner context managers exit first, so they are
+        # appended to _pending_scopes first (smaller index). Using
+        # -append_index as tiebreaker ensures that for same-range scopes
+        # the outer scope (larger index) sorts first and is pushed onto
+        # the stack first, leaving the inner scope on top.
+        sorted_scopes = sorted(
+            (
+                (ann, start, end, i)
+                for i, (ann, start, end) in enumerate(_pending_scopes)
+            ),
+            key=lambda s: (s[1], -s[2], -s[3]),
+        )
+        scope_ptr = 0
+        active_stack: list[tuple[int, Any]] = []  # (end_idx, annotation)
+
+        for i in range(num):
+            # Pop scopes whose range ended.
+            while active_stack and active_stack[-1][0] <= i:
+                active_stack.pop()
+
+            # Push scopes that start at or before this index.
+            while scope_ptr < len(sorted_scopes) and sorted_scopes[scope_ptr][1] <= i:
+                ann, _start_idx, end_idx, _idx = sorted_scopes[scope_ptr]
+                if end_idx > i:
+                    active_stack.append((end_idx, ann))
+                scope_ptr += 1
+
+            if not active_stack:
+                continue
+
+            node = nodes[i]
+            node_type = _check_cuda_bindings(
+                _cuda_runtime.cudaGraphNodeGetType(  # pyrefly: ignore[missing-attribute]
+                    node
+                )
+            )
+            if node_type not in annotatable:
+                continue
+
+            tools_id = _get_tools_id(node)
+            if tools_id is None:
+                logger.warning(
+                    "resolve_pending_annotations: toolsId unavailable, aborting"
+                )
+                _pending_scopes.clear()
+                _capture_graph = None
+                return
+
+            if len(active_stack) == 1:
+                _kernel_annotations[tools_id].append(active_stack[0][1])
+            else:
+                # Merge all active scopes into one dict. Inner scopes sit
+                # on top of the stack. Iterating reversed (inner first)
+                # with setdefault lets the inner scope's values win for
+                # overlapping keys (e.g. name, stream) while outer scopes
+                # fill in any missing keys.
+                merged: dict[str, Any] = {}
+                for _, ann in reversed(active_stack):
+                    if isinstance(ann, dict):
+                        for ak, av in ann.items():
+                            merged.setdefault(ak, av)
+                    else:
+                        merged.setdefault("name", ann)
+                _kernel_annotations[tools_id].append(merged)
+    except Exception:
+        logger.exception("resolve_pending_annotations failed")
+    finally:
+        _pending_scopes.clear()
+        _capture_graph = None
+
+
+def remap_to_exec_graph(torch_cuda_graph: torch.cuda.CUDAGraph) -> None:
+    """Remap annotation keys from capture graph ID to exec graph ID.
+
+    During capture, toolsId encodes the capture graph's ID in the upper
+    32 bits. After instantiation, the profiler uses the exec graph's ID.
+    This function rewrites the keys so annotations match the trace.
+
+    Must be called after the ``torch.cuda.graph()`` context exits.
+    """
+    if not _kernel_annotations:
+        return
+
+    exec_handle = _cuda_runtime.cudaGraphExec_t(  # pyrefly: ignore[missing-attribute]
+        init_value=torch_cuda_graph.raw_cuda_graph_exec()
+    )
+    exec_graph_id = _check_cuda_bindings(
+        _cuda_runtime.cudaGraphExecGetId(  # pyrefly: ignore[missing-attribute]
+            exec_handle
+        )
+    )
+
+    # Only remap annotations from the most recent capture graph.
+    # Previously remapped annotations (from earlier captures) keep their
+    # correct exec graph IDs.
+    capture_graph_id = _last_capture_graph_id
+    remapped: dict[int, list[Any]] = {}
+    for tools_id, ann_list in _kernel_annotations.items():
+        graph_id = tools_id >> 32
+        if capture_graph_id is not None and graph_id != capture_graph_id:
+            # Belongs to a different graph — keep as-is.
+            remapped[tools_id] = ann_list
+            continue
+        node_id = tools_id & 0xFFFFFFFF
+        new_tools_id = (exec_graph_id << 32) | node_id
+        if new_tools_id in remapped:
+            remapped[new_tools_id].extend(ann_list)
+        else:
+            remapped[new_tools_id] = list(ann_list)
+
+    _kernel_annotations.clear()
+    _kernel_annotations.update(remapped)
+
+
+def get_kernel_annotations() -> dict[int, list[Any]]:
+    """Return the current kernel annotation map (toolsId -> annotations)."""
+    return _kernel_annotations
+
+
+def clear_kernel_annotations() -> None:
+    """Clear all recorded kernel annotations and pending scopes."""
+    global _capture_graph
+    _kernel_annotations.clear()
+    _pending_scopes.clear()
+    _capture_graph = None
+
+
+# Counter-based stream ID registry. IDs start at 60 (above the highest
+# observed non-graphed CUDA stream ID) so every assigned lane is visually
+# distinct in Perfetto and doesn't collide with real streams.
+_stream_id_counter: int = 60
+_stream_id_map: dict[int, int] = {}
+
+
+def _get_stream_id(stream: torch.cuda.Stream) -> int:
+    """Return a small, stable stream ID for the given CUDA stream."""
+    global _stream_id_counter
+    key = stream.cuda_stream
+    if key not in _stream_id_map:
+        _stream_id_map[key] = _stream_id_counter
+        _stream_id_counter += 1
+    return _stream_id_map[key]
+
+
+def get_stream_for_pg(pg_key: str) -> int:
+    """Return a unique stream ID for the given process group key."""
+    global _stream_id_counter
+    if pg_key not in _stream_id_map:
+        _stream_id_map[pg_key] = _stream_id_counter  # type: ignore[assignment]
+        _stream_id_counter += 1
+    return _stream_id_map[pg_key]  # type: ignore[return-value]
+
+
+@contextmanager  # type: ignore[arg-type]
+def mark_stream(stream: torch.cuda.Stream, annotation: str | dict[str, Any]):
+    """Switch to stream, inject its ID into annotation, and mark kernels.
+
+    If *stream* is already the current stream, no stream switch or stream ID
+    injection happens — the kernels stay on whatever stream is active (which
+    keeps the trace faithful when e.g. FSDP uses the current stream for
+    copy-in instead of a separate one).
+    """
+    if not _annotations_enabled:
+        with torch.cuda.stream(stream):
+            yield
+        return
+    if stream.cuda_stream == torch.cuda.current_stream().cuda_stream:
+        with mark_kernels(annotation):
+            yield
+    else:
+        if isinstance(annotation, str):
+            annotation = {"str": annotation}
+        if isinstance(annotation, dict):
+            annotation["stream"] = _get_stream_id(stream)
+        with torch.cuda.stream(stream):
+            with mark_kernels(annotation):
+                yield
diff --git a/torch/cuda/_memory_viz.py b/torch/cuda/_memory_viz.py
index f58dd988851bd..b072721b0ab23 100644
--- a/torch/cuda/_memory_viz.py
+++ b/torch/cuda/_memory_viz.py
@@ -97,17 +97,26 @@ def format_flamegraph(flamegraph_lines, flamegraph_script=None):
         import urllib.request
 
         print(f"Downloading flamegraph.pl to: {flamegraph_script}")
-        with tempfile.NamedTemporaryFile(mode="wb", suffix=".pl") as f:
+        # Download to a temp file next to the target, then atomically move it
+        # into place so a concurrent reader never sees a partial file. The temp
+        # file lives in the same directory (same filesystem) so os.replace is
+        # atomic, and is created with delete=False + manual cleanup: moving a
+        # delete=True NamedTemporaryFile out from under its context manager makes
+        # __exit__ fail to unlink the now-missing file.
+        fd, tmp_name = tempfile.mkstemp(
+            suffix=".pl", dir=os.path.dirname(flamegraph_script) or None
+        )
+        os.close(fd)
+        try:
             urllib.request.urlretrieve(
                 "https://raw.githubusercontent.com/brendangregg/FlameGraph/master/flamegraph.pl",
-                f.name,
+                tmp_name,
             )
-            try:
-                os.chmod(f.name, 0o755)
-                os.rename(f.name, flamegraph_script)
-            except OSError:  # noqa: B001,E722
-                # Ok to skip, the file will be removed by tempfile
-                pass
+            os.chmod(tmp_name, 0o755)
+            os.replace(tmp_name, flamegraph_script)
+        finally:
+            if os.path.exists(tmp_name):
+                os.remove(tmp_name)
     args = [flamegraph_script, "--countname", "bytes"]
     with subprocess.Popen(
         args, stdin=subprocess.PIPE, stdout=subprocess.PIPE, encoding="utf-8"
diff --git a/torch/cuda/_utils.py b/torch/cuda/_utils.py
index d842e8b56ef41..bede63233f38e 100644
--- a/torch/cuda/_utils.py
+++ b/torch/cuda/_utils.py
@@ -4,6 +4,17 @@
 
 import torch
 
+
+try:
+    from cuda.bindings import (  # pyrefly: ignore[missing-import]
+        runtime as _cuda_bindings_runtime,
+    )
+
+    _HAS_CUDA_BINDINGS = True
+except ImportError:
+    _cuda_bindings_runtime = None  # type: ignore[assignment]
+    _HAS_CUDA_BINDINGS = False
+
 # The _get_device_index has been moved to torch.utils._get_device_index
 from torch._utils import _get_device_index as _torch_get_device_index
 
@@ -59,6 +70,35 @@ def _check_cuda(result: int) -> None:
     raise RuntimeError(f"CUDA error: {error_message}")
 
 
+def _check_cuda_bindings(result: Any) -> Any:
+    """Check a cuda.bindings (cuda-python) call result for errors.
+
+    All cuda.bindings runtime calls return ``(error, *outputs)``.  This
+    helper unpacks the tuple, raises on non-success, and returns the
+    outputs (``None`` for zero outputs, scalar for one, tuple otherwise).
+    """
+    if not _HAS_CUDA_BINDINGS:
+        raise RuntimeError("cuda.bindings is not available")
+    err, *out = result
+    if (
+        err
+        != _cuda_bindings_runtime.cudaError_t.cudaSuccess  # pyrefly: ignore[missing-attribute]
+    ):
+        _, err_str = (
+            _cuda_bindings_runtime.cudaGetErrorString(  # pyrefly: ignore[missing-attribute]
+                err
+            )
+        )
+        if isinstance(err_str, bytes):
+            err_str = err_str.decode()
+        raise RuntimeError(f"CUDA error: {err} ({err_str})")
+    if len(out) == 0:
+        return None
+    if len(out) == 1:
+        return out[0]
+    return out
+
+
 def _get_hiprtc_library() -> ctypes.CDLL:
     try:
         # pyrefly: ignore [import-error, missing-import]
@@ -435,11 +475,15 @@ def set_shared_memory_config(self, shared_mem_bytes: int) -> None:
         device_props = torch.cuda.get_device_properties()
         # HIP doesn't have shared_memory_per_block_optin in device properties, so we hard-code it here
         if torch.version.hip:
-            # navi, CDNA1-CDNA3 allows a max of 64KB shared memory
-            # CDNA4 allows a max of 160KB shared memory
-            max_shared_mem = (
-                65536 if device_props.gcnArchName != "gfx950" else 160 * 1024
-            )
+            # navi, CDNA1-CDNA3 allows a max of 64KB shared memory,
+            # CDNA4 (gfx950) 160KB, and CDNA5 (gfx1250) 320KB.
+            gcn_arch = device_props.gcnArchName.split(":", 1)[0]
+            if gcn_arch == "gfx950":
+                max_shared_mem = 160 * 1024
+            elif gcn_arch == "gfx1250":
+                max_shared_mem = 320 * 1024
+            else:
+                max_shared_mem = 65536
         else:
             max_shared_mem = getattr(
                 device_props, "shared_memory_per_block_optin", 49152
diff --git a/torch/cuda/graphs.py b/torch/cuda/graphs.py
index 43d559ddc2440..c6ff29236c25a 100644
--- a/torch/cuda/graphs.py
+++ b/torch/cuda/graphs.py
@@ -197,6 +197,12 @@ class graph:
             may be unsafe. "global" will error on actions in other threads, "thread_local" will only error for
             actions in the current thread, and "relaxed" will not error on actions. Do NOT change this setting
             unless you're familiar with `cudaStreamCaptureMode <https://docs.nvidia.com/cuda/cuda-runtime-api/group__CUDART__STREAM.html#group__CUDART__STREAM_1g9d0535d93a214cbf126835257b16ba85>`_
+        enable_annotations (bool, optional): If ``True``, enables kernel annotation
+            recording on entry and automatically calls
+            :func:`~torch.cuda._graph_annotations.resolve_pending_annotations` before
+            the capture ends.  Annotations are **not** cleared on exit so that multiple
+            graphs in the same workload can accumulate annotations.
+            Requires ``cuda.bindings`` package and cuda-compat >= 13.1 or CUDA driver >= 13.1.
 
     .. note::
         For effective memory sharing, if you pass a ``pool`` used by a previous capture and the previous capture
@@ -217,6 +223,7 @@ def __init__(
         pool: _POOL_HANDLE | None = None,
         stream: torch.cuda.Stream | None = None,
         capture_error_mode: str = "global",
+        enable_annotations: bool = False,
     ):
         # Lazy-init of default_capture_stream helps avoid circular-import errors.
         # Not thread safe, but graphs already have the general (explicitly documented)
@@ -233,6 +240,7 @@ def __init__(
         self.stream_ctx = torch.cuda.stream(self.capture_stream)
         self.cuda_graph = cuda_graph
         self.capture_error_mode = capture_error_mode
+        self._enable_annotations = enable_annotations
 
     def __enter__(self) -> None:
         # Free as much memory as we can for the graph
@@ -250,6 +258,11 @@ def __enter__(self) -> None:
         # pyrefly: ignore [missing-attribute]
         torch._C._host_emptyCache()
 
+        if self._enable_annotations:
+            from torch.cuda._graph_annotations import enable_annotations as _enable_ann
+
+            _enable_ann()
+
         # Stackoverflow seems comfortable with this pattern
         # https://stackoverflow.com/questions/26635684/calling-enter-and-exit-manually#39172487
         self.stream_ctx.__enter__()
@@ -262,8 +275,18 @@ def __enter__(self) -> None:
         )
 
     def __exit__(self, *args: object) -> None:
+        if self._enable_annotations:
+            from torch.cuda._graph_annotations import resolve_pending_annotations
+
+            resolve_pending_annotations()
+
         self.cuda_graph.capture_end()
         self.stream_ctx.__exit__(*args)
+
+        if self._enable_annotations:
+            from torch.cuda._graph_annotations import remap_to_exec_graph
+
+            remap_to_exec_graph(self.cuda_graph)
         # returning None should propagate exceptions from either capture_end or stream_ctx.__exit__()
 
 
@@ -426,7 +449,10 @@ def make_graphed_callables(
     # Hopefully prevents cudnn benchmarking and other lazy-initialization cuda work
     # from ending up in any captures.
     torch.cuda.synchronize()
-    with torch.cuda.stream(torch.cuda.Stream()):
+    # hipBLASLt handles are per-(device, stream) on ROCm and lazily created.
+    # Need to use the same stream for warmup and capture to avoid capture errors.
+    stream = torch.cuda.Stream()
+    with torch.cuda.stream(stream):
         for func, args, static_input_surface in zip(
             callables, _sample_args, per_callable_static_input_surfaces
         ):
@@ -459,7 +485,7 @@ def make_graphed_callables(
     per_callable_static_outputs = []
     per_callable_output_unflatten_spec = []
     for func, args, fwd_graph in zip(callables, _sample_args, fwd_graphs):
-        with torch.cuda.graph(fwd_graph, pool=mempool):
+        with torch.cuda.graph(fwd_graph, stream=stream, pool=mempool):
             func_outputs = func(*args)
 
         flatten_outputs, spec = torch.utils._pytree.tree_flatten(func_outputs)
@@ -483,7 +509,7 @@ def make_graphed_callables(
         outputs_grad = tuple(o for o in static_outputs if o.requires_grad)
         grad_inputs = None
         if len(outputs_grad) > 0:
-            with torch.cuda.graph(bwd_graph, pool=mempool):
+            with torch.cuda.graph(bwd_graph, stream=stream, pool=mempool):
                 grad_inputs = torch.autograd.grad(
                     outputs=outputs_grad,
                     inputs=tuple(i for i in static_input_surface if i.requires_grad),
diff --git a/torch/cuda/green_contexts.py b/torch/cuda/green_contexts.py
index c08d2c061e404..c7aab1b749db1 100644
--- a/torch/cuda/green_contexts.py
+++ b/torch/cuda/green_contexts.py
@@ -1,6 +1,10 @@
 import torch
 
 
+__all__ = [
+    "GreenContext",
+]
+
 _GreenContext = object
 SUPPORTED = False
 
@@ -19,18 +23,59 @@ class GreenContext(_GreenContext):
     """
 
     @staticmethod
-    def create(num_sms: int, device_id: int = 0) -> _GreenContext:
+    def create(
+        *,
+        num_sms: int | None = None,
+        workqueue_scope: str | None = None,
+        workqueue_concurrency_limit: int | None = None,
+        device_id: int | None = None,
+    ) -> _GreenContext:
         r"""Create a CUDA green context.
 
+        At least one of ``num_sms`` or ``workqueue_scope`` must be specified.
+        Both can be combined to partition SMs and configure workqueues in the
+        same green context.
+
         Arguments:
-            num_sms (int): The number of SMs to use in the green context.
+            num_sms (int, optional): The number of SMs to use in the green
+                context. When ``None``, SMs are not partitioned.
+            workqueue_scope (str, optional): Workqueue sharing scope. One of
+                ``"device_ctx"`` (shared across all contexts, default driver
+                behaviour) or ``"balanced"`` (non-overlapping workqueues with
+                other balanced green contexts). When ``None``, no workqueue
+                configuration is applied.
+            workqueue_concurrency_limit (int, optional): Maximum number of
+                concurrent stream-ordered workloads for the workqueue. Requires
+                ``workqueue_scope`` to be set.
             device_id (int, optional): The device index of green context.
+                When ``None``, the current device is used.
+        """
+        if not SUPPORTED:
+            raise RuntimeError("PyTorch was not built with Green Context support!")
+        return _GreenContext.create(  # type: ignore[attr-defined]
+            device_id=device_id,
+            num_sms=num_sms,
+            workqueue_scope=workqueue_scope,
+            workqueue_concurrency_limit=workqueue_concurrency_limit,
+        )
+
+    @staticmethod
+    def max_workqueue_concurrency(device_id: int | None = None) -> int:
+        r"""Return the maximum workqueue concurrency limit for the device.
+
+        This queries the device for the default number of concurrent
+        stream-ordered workloads supported by workqueue configuration
+        resources.
+
+        Arguments:
+            device_id (int, optional): The device index to query. When
+                ``None``, the current device is used.
         """
         if not SUPPORTED:
             raise RuntimeError("PyTorch was not built with Green Context support!")
-        return _GreenContext.create(num_sms, device_id)  # type: ignore[attr-defined]
+        return _GreenContext.max_workqueue_concurrency(device_id=device_id)  # type: ignore[attr-defined]
 
-    # Note that these functions are bypassed by we define them here
+    # Note that these functions are bypassed but we define them here
     # for Sphinx documentation purposes
     def set_context(self) -> None:  # pylint: disable=useless-parent-delegation
         r"""Make the green context the current context."""
@@ -42,6 +87,6 @@ def pop_context(self) -> None:  # pylint: disable=useless-parent-delegation
         """
         return super().pop_context()  # type: ignore[misc]
 
-    def Stream(self) -> torch.Stream:
+    def Stream(self) -> "torch.cuda.Stream":
         r"""Return the CUDA Stream used by the green context."""
         return super().Stream()
diff --git a/torch/cuda/memory.py b/torch/cuda/memory.py
index 26bd7eb35e9f0..90ab843c139f4 100644
--- a/torch/cuda/memory.py
+++ b/torch/cuda/memory.py
@@ -32,6 +32,7 @@
 __all__ = [
     "caching_allocator_alloc",
     "caching_allocator_delete",
+    "caching_allocator_disabled",
     "caching_allocator_enable",
     "get_per_process_memory_fraction",
     "set_per_process_memory_fraction",
@@ -98,15 +99,6 @@ def _host_allocator():
     return torch._C._cuda_cudaHostAllocator()
 
 
-@contextlib.contextmanager
-def _free_mutex():
-    torch._C._cuda_lock_mutex()
-    try:
-        yield
-    finally:
-        torch._C._cuda_unlock_mutex()
-
-
 def caching_allocator_alloc(size, device: "Device" = None, stream=None):
     r"""Perform a memory allocation using the CUDA memory allocator.
 
@@ -166,6 +158,17 @@ def caching_allocator_enable(value: bool = True) -> None:
         torch._C._cuda_cudaCachingAllocator_enable(value)
 
 
+@contextlib.contextmanager
+def caching_allocator_disabled():
+    r"""Context manager that temporarily disables the CUDA caching allocator."""
+    prev = torch._C._cuda_cudaCachingAllocator_is_enabled()
+    caching_allocator_enable(False)
+    try:
+        yield
+    finally:
+        caching_allocator_enable(prev)
+
+
 def set_per_process_memory_fraction(fraction, device: "Device" = None) -> None:
     r"""Set memory fraction for a process.
 
@@ -277,6 +280,8 @@ def memory_stats(device: "Device" = None) -> dict[str, Any]:
       cuMemMap and cudaMalloc.
     - ``"num_device_free"``: number of CUDA free calls. This includes both cuMemUnmap
       and cudaFree.
+    - ``"num_oom_rejections"``: number of allocations preemptively rejected by the
+        throw_on_cudamalloc_oom + per_process_memory_fraction policy.
 
     The caching allocator can be configured via ENV to not split blocks larger than a
     defined size (see Memory Management section of the Cuda Semantics documentation).
@@ -989,6 +994,21 @@ def _record_memory_history_impl(
 _record_memory_history.__signature__ = signature(_record_memory_history_impl)  # type: ignore[attr-defined]
 
 
+def _allocation_traceback(data_ptr: int) -> list[dict[str, Any]] | None:
+    r"""Return the allocation traceback for a currently-allocated CUDA pointer,
+    or ``None`` if the pointer is not found or recording was not enabled.
+
+    ``data_ptr`` must be the base address of a recorded allocation, e.g. the
+    result of ``tensor.untyped_storage().data_ptr()``. Passing the result of
+    ``tensor.data_ptr()`` for a view with a nonzero ``storage_offset`` may
+    return ``None`` even if the underlying allocation is still live.
+
+    Requires :func:`_record_memory_history` to have been called with
+    ``context != None`` beforehand.
+    """
+    return torch._C._cuda_allocationTraceback(data_ptr)
+
+
 def _snapshot(device: "Device" = None, augment_with_fx_traces=False):
     """Save a snapshot of CUDA memory state at the time it was called.
 
@@ -1012,6 +1032,9 @@ class Segment(TypedDict):
             total_size: int  #  cudaMalloc'd size of segment
             stream: int
             segment_type: Literal["small", "large"]  # 'large' (>1MB)
+            segment_pool_id: Tuple[
+                int, int
+            ]  # id of the memory pool owning this segment
             allocated_size: int  # size of memory in use
             active_size: int  # size of memory in use or in active_awaiting_free state
             blocks: List[Block]
@@ -1070,6 +1093,7 @@ class TraceEntry(TypedDict):
             stream: int
             device_free: int  # only present for OOM, the amount of
             # memory cuda still reports to be free
+            pool_id: Tuple[int, int]  # id of the memory pool for this entry
 
     Args:
         device: Device to capture snapshot for. If None, captures for current device.
diff --git a/torch/custom_class.h b/torch/custom_class.h
index 4a3d1a39db31e..b31a21a41b1f4 100644
--- a/torch/custom_class.h
+++ b/torch/custom_class.h
@@ -6,7 +6,6 @@
 #include <ATen/core/class_type.h>
 #include <ATen/core/op_registration/infer_schema.h>
 #include <ATen/core/stack.h>
-#include <c10/util/C++17.h>
 #include <c10/util/Metaprogramming.h>
 #include <c10/util/TypeList.h>
 #include <c10/util/TypeTraits.h>
diff --git a/torch/distributed/__init__.py b/torch/distributed/__init__.py
index 24add6a4d8aeb..4f24b45494e6e 100644
--- a/torch/distributed/__init__.py
+++ b/torch/distributed/__init__.py
@@ -169,7 +169,8 @@ def breakpoint(rank: int = 0, skip: int = 0, timeout_s=3600):
     # stubs as necessary.
     # We cannot define stubs directly because they confuse pyre
 
-    class _ProcessGroupStub:
+    class _Stub:
         pass
 
-    sys.modules["torch.distributed"].ProcessGroup = _ProcessGroupStub  # type: ignore[attr-defined]
+    sys.modules["torch.distributed"].GroupName = _Stub  # type: ignore[attr-defined]
+    sys.modules["torch.distributed"].ProcessGroup = _Stub  # type: ignore[attr-defined]
diff --git a/torch/distributed/_composable/checkpoint_activation.py b/torch/distributed/_composable/checkpoint_activation.py
index 93ae14110ef79..37ed17b03b5c6 100644
--- a/torch/distributed/_composable/checkpoint_activation.py
+++ b/torch/distributed/_composable/checkpoint_activation.py
@@ -115,7 +115,8 @@ def forward_hook(module: nn.Module, inputs: tuple[Any, ...], output: Any) -> Any
         if checkpoint.state(module).enable_hook:
             try:
                 gen = checkpoint.state(module)._ac_generator
-                assert gen is not None
+                if gen is None:
+                    raise AssertionError
                 next(gen)
             except StopIteration:
                 pass
diff --git a/torch/distributed/_composable/contract.py b/torch/distributed/_composable/contract.py
index 283232d46f2ea..1bf3727aa63d2 100644
--- a/torch/distributed/_composable/contract.py
+++ b/torch/distributed/_composable/contract.py
@@ -119,7 +119,6 @@ def wrapper(
             all_orig_named_buffers: list[dict[str, torch.Tensor]] = []
             all_orig_named_modules: list[dict[str, nn.Module]] = []
 
-            # pyrefly: ignore [bad-assignment]
             for module in modules:
                 default_all_state: dict[Callable, _State] = OrderedDict()
                 default_registry: dict[str, RegistryItem] = OrderedDict()
@@ -146,11 +145,8 @@ def wrapper(
                 all_state.setdefault(func, state)
                 registry.setdefault(func.__name__, registry_item)
 
-                # pyrefly: ignore [missing-attribute]
                 all_orig_named_params.append(OrderedDict(module.named_parameters()))
-                # pyrefly: ignore [missing-attribute]
                 all_orig_named_buffers.append(OrderedDict(module.named_buffers()))
-                # pyrefly: ignore [missing-attribute]
                 all_orig_named_modules.append(OrderedDict(module.named_modules()))
 
             updated = func(inp_module, *args, **kwargs)
@@ -165,13 +161,9 @@ def wrapper(
             all_new_named_params: list[dict[str, nn.Parameter]] = []
             all_new_named_buffers: list[dict[str, torch.Tensor]] = []
             all_new_named_modules: list[dict[str, nn.Module]] = []
-            # pyrefly: ignore [bad-assignment]
             for module in updated_modules:
-                # pyrefly: ignore [missing-attribute]
                 all_new_named_params.append(OrderedDict(module.named_parameters()))
-                # pyrefly: ignore [missing-attribute]
                 all_new_named_buffers.append(OrderedDict(module.named_buffers()))
-                # pyrefly: ignore [missing-attribute]
                 all_new_named_modules.append(OrderedDict(module.named_modules()))
 
             num_orig_modules = len(all_orig_named_modules)
diff --git a/torch/distributed/_composable/replicate.py b/torch/distributed/_composable/replicate.py
index fd3f6fb565cfd..70f5a21f8f09c 100644
--- a/torch/distributed/_composable/replicate.py
+++ b/torch/distributed/_composable/replicate.py
@@ -66,7 +66,8 @@ def _collect_params(
     def lazy_init(self) -> None:
         @torch._disable_dynamo(recursive=True)
         def _lazy_init():
-            assert self._init_args is not None
+            if self._init_args is None:
+                raise AssertionError
             self.init(*self._init_args, **self._init_kwargs)
             self.register_comm_hook()
             self._init_args = ()
diff --git a/torch/distributed/_composable/replicate_with_fsdp.py b/torch/distributed/_composable/replicate_with_fsdp.py
index ba8a54452058f..e0c9125fca52d 100644
--- a/torch/distributed/_composable/replicate_with_fsdp.py
+++ b/torch/distributed/_composable/replicate_with_fsdp.py
@@ -15,9 +15,11 @@
 from torch.distributed.fsdp._fully_shard._fsdp_init import (
     _apply_to_module,
     _get_device_from_mesh,
+    _get_mesh_info,
     _get_modules_and_states,
     _init_default_mesh,
     _init_param_group,
+    _validate_mesh as _validate_mesh_common,
     _validate_module as _validate_module_common,
 )
 from torch.distributed.fsdp._fully_shard._fsdp_state import FSDPState, FSDPStateContext
@@ -30,6 +32,7 @@
 
 
 if TYPE_CHECKING:
+    from torch.distributed.fsdp._fully_shard._fsdp_api import DataParallelMeshDims
     from torch.distributed.tensor import DeviceMesh
 
 
@@ -85,6 +88,7 @@ def replicate(
     mp_policy: MixedPrecisionPolicy = ...,
     offload_policy: OffloadPolicy = ...,
     ignored_params: set[nn.Parameter] | None = ...,
+    dp_mesh_dims: DataParallelMeshDims | None = ...,
 ) -> ReplicateModule: ...
 
 
@@ -97,6 +101,7 @@ def replicate(
     mp_policy: MixedPrecisionPolicy = ...,
     offload_policy: OffloadPolicy = ...,
     ignored_params: set[nn.Parameter] | None = ...,
+    dp_mesh_dims: DataParallelMeshDims | None = ...,
 ) -> list[ReplicateModule]: ...
 
 
@@ -108,6 +113,7 @@ def replicate(
     mp_policy: MixedPrecisionPolicy = MixedPrecisionPolicy(),
     offload_policy: OffloadPolicy = OffloadPolicy(),
     ignored_params: set[nn.Parameter] | None = None,
+    dp_mesh_dims: DataParallelMeshDims | None = None,
 ):
     r"""Replicates a module
 
@@ -122,8 +128,17 @@ def replicate(
     torch._C._log_api_usage_once("torch.distributed._composable.replicate_with_fsdp")
     _validate_module(module)
     mesh = mesh or _init_default_mesh(mesh_dim_names=("replicate",))
-    _validate_mesh(mesh)
-    mesh_info = DDPMeshInfo(mesh, replicate_mesh_dim=0)
+    if dp_mesh_dims is not None:
+        _validate_mesh_common(mesh, dp_mesh_dims)
+        mesh_info = _get_mesh_info(mesh, dp_mesh_dims)
+        if not isinstance(mesh_info, DDPMeshInfo):
+            raise ValueError(
+                "replicate() with dp_mesh_dims requires replicate-only "
+                "dims (no shard dims). Use fully_shard() for sharding."
+            )
+    else:
+        _validate_mesh(mesh)
+        mesh_info = DDPMeshInfo(mesh, replicate_mesh_dim=0)
     device = _get_device_from_mesh(mesh)
     # managed_modules (3rd return) and buffers (5th return) are unused:
     # - managed_modules: FSDP uses this to set Dynamo-specific attributes
diff --git a/torch/distributed/_functional_collectives.py b/torch/distributed/_functional_collectives.py
index fcfa8a2f43a1e..733147622ad85 100644
--- a/torch/distributed/_functional_collectives.py
+++ b/torch/distributed/_functional_collectives.py
@@ -1,8 +1,9 @@
 # mypy: allow-untyped-defs
 import contextlib
+import math
 import sys
 import warnings
-from typing import Any, cast, TYPE_CHECKING, Union
+from typing import Any, cast, TYPE_CHECKING
 
 import torch
 import torch.distributed as dist
@@ -90,14 +91,17 @@ def is_torchdynamo_compiling():  # type: ignore[misc]
 
 The different types will be desugared to a canonical format
 """
-RANK_TYPES = Union[
-    list[int],
-    list[list[int]],
-    dist.ProcessGroup,
-    DeviceMesh,
-    tuple["dist.tensor.DeviceMesh", int],
-    c10d.GroupName,
-]
+RANK_TYPES = (
+    list[int]
+    | list[list[int]]
+    | dist.ProcessGroup
+    | DeviceMesh
+    | tuple["dist.tensor.DeviceMesh", int]
+    | c10d.GroupName
+)
+
+
+from torch._utils import _chunk_or_narrow_cat  # noqa: F401
 
 
 """
@@ -147,8 +151,10 @@ def broadcast(self: torch.Tensor, src: int, group: RANK_TYPES, tag: str = ""):
         group (ProcessGroup or List[int]): The process group to work on.
         tag (str, optional): A unique identifier for the collective. Default: empty string
     """
-    group_name = _resolve_group_name(group, tag)
-    tensor = torch.ops._c10d_functional.broadcast(self, src, group_name)
+    group = _resolve_group(group, tag)
+    tensor = torch.ops._c10d_functional.broadcast(
+        self, src, _group_or_group_name(group)
+    )
     return _maybe_wrap_tensor(tensor)
 
 
@@ -169,8 +175,10 @@ def all_reduce(self: torch.Tensor, reduceOp: str, group: RANK_TYPES, tag: str =
     :: N.B. If you pass a PG or a 1D list to perform a MPMD collective, the compiler won't be able to recover
     that information and perform collective algebraic optimization. Use other forms of input for that.
     """
-    group_name = _resolve_group_name(group, tag)
-    tensor = torch.ops._c10d_functional.all_reduce(self, reduceOp.lower(), group_name)
+    group = _resolve_group(group, tag)
+    tensor = torch.ops._c10d_functional.all_reduce(
+        self, reduceOp.lower(), _group_or_group_name(group)
+    )
     return _maybe_wrap_tensor(tensor)
 
 
@@ -196,21 +204,22 @@ def all_gather_tensor(
     :: N.B. If you pass a PG or a 1D list to perform a MPMD collective, the compiler won't be able to recover
     that information and perform collective algebraic optimization. Use other forms of input for that.
     """
-    if not self.is_contiguous():
-        raise AssertionError("Tensor must be contiguous for all_gather_tensor")
-    group_name = _resolve_group_name(group, tag)
-    group_size = c10d._get_group_size_by_name(group_name)
+    group = _resolve_group(group, tag)
+    group_size = c10d._get_group_size_by_name(group)
     tensor = torch.ops._c10d_functional.all_gather_into_tensor(
-        self, group_size, group_name
+        self, group_size, _group_or_group_name(group)
     )
     res = _maybe_wrap_tensor(tensor)
-    # TODO this should be done inside AsyncCollectiveTensor to delay the wait() call
     if gather_dim != 0:
-        # torch.cat access the data so we already need to wait here, first do wait
-        # and then chunk + cat avoid us going through ACT dispatching logic again
+        # Check if _maybe_view_chunk_cat can use the view optimization.
+        # If not, it will use torch.cat which needs the data anyway, so
+        # wait early to avoid AsyncCollectiveTensor dispatch overhead.
         if isinstance(res, AsyncCollectiveTensor):
-            res = res.wait()  # type: ignore[attr-defined]
-
+            shape = list(res.shape)
+            numel_between = math.prod(shape[1:gather_dim]) if gather_dim > 1 else 1
+            can_use_view = shape[0] == group_size and numel_between == 1
+            if not can_use_view:
+                res = res.wait()
         res = _maybe_view_chunk_cat(res, group_size, gather_dim)
     return res
 
@@ -231,20 +240,24 @@ def all_gather_tensor_autograd(
 
     See all_gather_tensor for more details on usage.
     """
-    group_name = _resolve_group_name(group, tag)
-    group_size = c10d._get_group_size_by_name(group_name)
+    group = _resolve_group(group, tag)
+    group_size = c10d._get_group_size_by_name(group)
 
     tensor = torch.ops._c10d_functional_autograd.all_gather_into_tensor(
-        self, group_size, group_name
+        self, group_size, _group_or_group_name(group)
     )
     res = _FromTorchTensor.apply(tensor)
-    # TODO this should be done inside AsyncCollectiveTensor to delay the wait() call
     if gather_dim != 0:
-        # torch.cat access the data so we already need to wait here, first do wait
-        # and then chunk + cat avoid us going through ACT dispatching logic again
+        # Check if _maybe_view_chunk_cat can use the view optimization.
+        # If not, it will use torch.cat which needs the data anyway, so
+        # wait early to avoid AsyncCollectiveTensor dispatch overhead.
         if isinstance(res, AsyncCollectiveTensor):
-            res = res.wait()  # type: ignore[attr-defined]
-        res = torch.cat(torch.chunk(res, group_size, dim=0), dim=gather_dim)
+            shape = list(res.shape)
+            numel_between = math.prod(shape[1:gather_dim]) if gather_dim > 1 else 1
+            can_use_view = shape[0] == group_size and numel_between == 1
+            if not can_use_view:
+                res = res.wait()
+        res = _maybe_view_chunk_cat(res, group_size, gather_dim)
     return res
 
 
@@ -270,22 +283,21 @@ def reduce_scatter_tensor(
     :: N.B. If you pass a PG or a 1D list to perform a MPMD collective, the compiler won't be able to recover
     that information and perform collective algebraic optimization. Use other forms of input for that.
     """
-    group_name = _resolve_group_name(group, tag)
-    group_size = c10d._get_group_size_by_name(group_name)
+    group = _resolve_group(group, tag)
+    group_size = c10d._get_group_size_by_name(group)
 
     if self.size(scatter_dim) % group_size != 0:
         raise AssertionError(
             f"input dimension 0 ({self.size(0)} must be a multiple of group_size {group_size})"
         )
     if scatter_dim != 0:
-        tensor_list = torch.chunk(self, group_size, dim=scatter_dim)
-        self = torch.cat(tensor_list)
+        self = _chunk_or_narrow_cat(self, group_size, narrow_dim=scatter_dim, cat_dim=0)
 
     tensor = torch.ops._c10d_functional.reduce_scatter_tensor(
         self,
         reduceOp.lower(),
         group_size,
-        group_name,  # type: ignore[possibly-undefined]
+        _group_or_group_name(group),
     )
     res = _maybe_wrap_tensor(tensor)
     return res
@@ -310,22 +322,21 @@ def reduce_scatter_tensor_autograd(
     See reduce_scatter_tensor for more details on usage.
     """
 
-    group_name = _resolve_group_name(group, tag)
-    group_size = c10d._get_group_size_by_name(group_name)
+    group = _resolve_group(group, tag)
+    group_size = c10d._get_group_size_by_name(group)
 
     if self.size(scatter_dim) % group_size != 0:
         raise AssertionError(
             f"input dimension 0 ({self.size(0)} must be a multiple of group_size {group_size}"
         )
     if scatter_dim != 0:
-        tensor_list = torch.chunk(self, group_size, dim=scatter_dim)
-        self = torch.cat(tensor_list)
+        self = _chunk_or_narrow_cat(self, group_size, narrow_dim=scatter_dim, cat_dim=0)
 
     tensor = torch.ops._c10d_functional_autograd.reduce_scatter_tensor(
         self,
         reduceOp.lower(),
         group_size,
-        group_name,  # type: ignore[possibly-undefined]
+        _group_or_group_name(group),
     )
     res = _FromTorchTensor.apply(tensor)
     return res
@@ -350,11 +361,11 @@ def all_reduce_coalesced(
     :: N.B. If you pass a PG or a 1D list to perform a MPMD collective, the compiler won't be able to recover
     that information and perform collective algebraic optimization. Use other forms of input for that.
     """
-    group_name = _resolve_group_name(group, tag)
+    group = _resolve_group(group, tag)
     tensor_list = torch.ops._c10d_functional.all_reduce_coalesced(  # type: ignore[attr-defined]
         self,
         reduceOp.lower(),
-        group_name,
+        _group_or_group_name(group),
     )
     return list(map(_maybe_wrap_tensor, tensor_list))
 
@@ -378,12 +389,12 @@ def all_gather_into_tensor_coalesced(
     :: N.B. If you pass a PG or a 1D list to perform a MPMD collective, the compiler won't be able to recover
     that information and perform collective algebraic optimization. Use other forms of input for that.
     """
-    group_name = _resolve_group_name(group, tag)
-    group_size = c10d._get_group_size_by_name(group_name)
+    group = _resolve_group(group, tag)
+    group_size = c10d._get_group_size_by_name(group)
     tensor_list = torch.ops._c10d_functional.all_gather_into_tensor_coalesced(  # type: ignore[attr-defined]
         self,
         group_size,
-        group_name,
+        _group_or_group_name(group),
     )
     return list(map(_maybe_wrap_tensor, tensor_list))
 
@@ -410,8 +421,8 @@ def reduce_scatter_tensor_coalesced(
     :: N.B. If you pass a PG or a 1D list to perform a MPMD collective, the compiler won't be able to recover
     that information and perform collective algebraic optimization. Use other forms of input for that.
     """
-    group_name = _resolve_group_name(group, tag)
-    group_size = c10d._get_group_size_by_name(group_name)
+    group = _resolve_group(group, tag)
+    group_size = c10d._get_group_size_by_name(group)
 
     if len(scatter_dim) != len(inputs):
         raise AssertionError(
@@ -430,7 +441,7 @@ def reduce_scatter_tensor_coalesced(
         inputs,
         reduceOp.lower(),
         group_size,
-        group_name,  # type: ignore[possibly-undefined]
+        _group_or_group_name(group),
     )
 
     return list(map(_maybe_wrap_tensor, tensor_list))
@@ -488,8 +499,8 @@ def all_to_all_single(
             raise AssertionError(
                 f"All input_split_sizes must be int or SymInt, got {input_split_sizes}"
             )
-    group_name = _resolve_group_name(group, tag)
-    group_size = c10d._get_group_size_by_name(group_name)
+    group = _resolve_group(group, tag)
+    group_size = c10d._get_group_size_by_name(group)
     if output_split_sizes is None or input_split_sizes is None:
         if not (output_split_sizes is None and input_split_sizes is None):
             raise AssertionError(
@@ -502,7 +513,7 @@ def all_to_all_single(
         self,
         output_split_sizes,
         input_split_sizes,
-        group_name,
+        _group_or_group_name(group),
     )
     return _maybe_wrap_tensor(tensor)
 
@@ -530,8 +541,8 @@ def all_to_all_single_autograd(
                 f"All input_split_sizes must be int or SymInt, got {input_split_sizes}"
             )
 
-    group_name = _resolve_group_name(group, tag)
-    group_size = c10d._get_group_size_by_name(group_name)
+    group = _resolve_group(group, tag)
+    group_size = c10d._get_group_size_by_name(group)
     if output_split_sizes is None or input_split_sizes is None:
         if not (output_split_sizes is None and input_split_sizes is None):
             raise AssertionError(
@@ -544,7 +555,7 @@ def all_to_all_single_autograd(
         self,
         output_split_sizes,
         input_split_sizes,
-        group_name,
+        _group_or_group_name(group),
     )
     return _FromTorchTensor.apply(tensor)
 
@@ -1178,26 +1189,31 @@ def cast_listint(x):
     return (tag, rankset, group_size)
 
 
-def _resolve_group_name(group: RANK_TYPES, tag: str = "") -> c10d.GroupName:
+def _resolve_group(
+    group: RANK_TYPES, tag: str = ""
+) -> dist.ProcessGroup | c10d.GroupName:
     """
-    Given group in RANK_TYPES, return the group name.
+    Given group in RANK_TYPES, return a ProcessGroup or group name.
     """
     # `tag` will be deprecated. See details in:
     # https://github.com/pytorch/pytorch/issues/93173#issuecomment-1907095208
     if isinstance(group, dist.ProcessGroup):
-        return group.group_name
+        return group
     elif isinstance(group, str):
         # In some cases Dynamo doesn't like tracing through NewType constructors
         # - so use a cast instead (the actual newtype representation is
         # literally the underlying type so this is fine). I haven't been able to
         # reproduce it in isolation (see T247631668).
         # pyrefly: ignore [redundant-cast]
-        return cast(c10d.GroupName, group)  # c10d.GroupName(group)
+        group_name = cast(c10d.GroupName, group)  # c10d.GroupName(group)
+        return group_name
     elif isinstance(group, DeviceMesh):
         if group.ndim != 1:
             raise AssertionError(
                 "Only 1D mesh is supported, pass in (DeviceMesh, int) together if mesh > 1D"
             )
+        if dist.config.compile_on_one_rank:
+            return torch.ops._dtensor.mesh_get_process_group(group, 0)
         return group._dim_group_names[0]
     elif isinstance(group, tuple):
         if (
@@ -1207,9 +1223,13 @@ def _resolve_group_name(group: RANK_TYPES, tag: str = "") -> c10d.GroupName:
         ):
             dmesh = group[0]
             dim = group[1]
+            if dist.config.compile_on_one_rank:
+                return torch.ops._dtensor.mesh_get_process_group(dmesh, dim)
             return dmesh._dim_group_names[dim]
         else:
-            raise ValueError("Invalid tuple for group must be (DeviceMesh, int)")
+            raise ValueError(
+                f"Invalid tuple for group must be (DeviceMesh, int). Instead got {(type(group[0]), type(group[1]))}"
+            )
     elif isinstance(group, list):
         if not is_torchdynamo_compiling():
             warnings.warn(
@@ -1219,12 +1239,26 @@ def _resolve_group_name(group: RANK_TYPES, tag: str = "") -> c10d.GroupName:
                 FutureWarning,
                 stacklevel=3,
             )
-        # pyrefly: ignore [redundant-cast]
-        return c10d._resolve_group_name_by_ranks_and_tag(cast(list[int], group), tag)
+        return c10d._resolve_group_name_by_ranks_and_tag(
+            # pyrefly: ignore [redundant-cast]
+            cast(list[int], group),
+            tag,
+        )
     else:
         raise ValueError(f"Unsupported group type: {type(group)}, {group}")
 
 
+def _resolve_group_name(group: RANK_TYPES, tag: str = "") -> c10d.GroupName:
+    """
+    Given group in RANK_TYPES, return the group name.
+    """
+    group = _resolve_group(group, tag)
+    if isinstance(group, str):
+        return c10d.GroupName(group)
+    else:
+        return group.group_name
+
+
 class _FromTorchTensor(torch.autograd.Function):
     """
     _FromTorchTensor allows autograd to propagate from a normal Tensor to an
@@ -1392,6 +1426,21 @@ def _wait_tensor_meta(self, *args):
     return torch.empty_like(self)
 
 
+def _isend_meta(self, *args):
+    return torch.empty(0, dtype=self.dtype, device=self.device)
+
+
+def _irecv_meta(self, *args):
+    return torch.empty_like(self)
+
+
+def _batch_p2p_ops_meta(op_list, peer_list, tag_list, tensors, group_name):
+    return [
+        t if op == "irecv" else torch.empty(0, dtype=t.dtype, device=t.device)
+        for op, t in zip(op_list, tensors)
+    ]
+
+
 def _all_gather_into_tensor_meta(shard, tag, rankset, group_size):
     return _make_all_gather_out_tensor(shard, group_size)
 
@@ -1491,6 +1540,10 @@ def _reduce_scatter_tensor_coalesced_native_meta(
 lib_impl.impl("all_reduce_coalesced", _all_reduce_coalesced_meta, "Meta")
 lib_impl.impl("all_reduce_coalesced_", _all_reduce_coalesced__meta, "Meta")
 lib_impl.impl("wait_tensor", _wait_tensor_meta, "Meta")
+lib_impl.impl("isend", _isend_meta, "Meta")
+lib_impl.impl("irecv", _irecv_meta, "Meta")
+lib_impl.impl("batch_p2p_ops", _batch_p2p_ops_meta, "Meta")
+
 lib_impl.impl(
     "all_gather_into_tensor_out", _all_gather_into_tensor_out_native_meta, "Meta"
 )
@@ -1513,9 +1566,17 @@ def _reduce_scatter_tensor_coalesced_native_meta(
 lib_impl.impl("broadcast", _broadcast_meta, "Meta")
 lib_impl.impl("broadcast_", _broadcast__meta, "Meta")
 
-# mark these ops has side effect so that they won't be removed by DCE
+# Mark these ops as side effectful so that DCE does not remove communication
+# whose result tensors are ignored by user code.
 torch.fx.node.has_side_effect(torch.ops._c10d_functional.wait_tensor.default)  # type: ignore[has-type]
 torch.fx.node.has_side_effect(torch.ops._c10d_functional.wait_tensor)  # type: ignore[has-type]
+torch.fx.node.has_side_effect(torch.ops._c10d_functional.isend.default)  # type: ignore[has-type]
+torch.fx.node.has_side_effect(torch.ops._c10d_functional.isend)  # type: ignore[has-type]
+torch.fx.node.has_side_effect(torch.ops._c10d_functional.irecv.default)  # type: ignore[has-type]
+torch.fx.node.has_side_effect(torch.ops._c10d_functional.irecv)  # type: ignore[has-type]
+torch.fx.node.has_side_effect(torch.ops._c10d_functional.batch_p2p_ops.default)  # type: ignore[has-type]
+torch.fx.node.has_side_effect(torch.ops._c10d_functional.batch_p2p_ops)  # type: ignore[has-type]
+
 
 # Register legacy ops for backward compatibility
 # TODO(yifu): remove these in functional collective beta release
@@ -1526,6 +1587,9 @@ def _reduce_scatter_tensor_coalesced_native_meta(
     "all_reduce(Tensor self, str reduceOp, str tag, int[] ranks, int group_size) -> Tensor",
     "all_reduce_coalesced(Tensor[] self, str reduceOp, str tag, int[] ranks, int group_size) -> Tensor[]",
     "wait_tensor(Tensor self) -> Tensor",
+    "isend(Tensor self, int dst, int tag, str group_name) -> Tensor",
+    "irecv(Tensor self, int src, int tag, str group_name) -> Tensor",
+    "batch_p2p_ops(str[] op_list, int[] peer_list, int[] tag_list, Tensor[] tensors, str group_name) -> Tensor[]",
     "all_gather_into_tensor(Tensor shard, str tag, int[] ranks, int group_size) -> Tensor",
     "all_gather_into_tensor_coalesced(Tensor[] input, str tag, int[] ranks, int group_size) -> Tensor[]",
     "reduce_scatter_tensor(Tensor input, str reduceOp, str tag, int[] ranks, int group_size) -> Tensor",
@@ -1690,25 +1754,181 @@ def all_gather_inplace(
     return tensor_list
 
 
-from torch.distributed.distributed_c10d import (
+def isend_inplace(
+    tensor: torch.Tensor,
+    dst: int,
+    tag: int = 0,
+    group: dist.ProcessGroup | None = None,
+    group_dst: int = -1,
+):
+    if group is None:
+        group = dist.group.WORLD
+    if group is None:
+        raise AssertionError("group cannot be None")
+    if group_dst != -1:
+        if dst is not None:
+            raise ValueError(
+                "Cannot specify both 'dst' and 'group_dst' args as per eager impl"
+            )
+        global_dst = c10d.get_global_rank(group, group_dst)
+    else:
+        global_dst = dst
+
+    group_name = _resolve_group_name(group)
+    tensor = torch.ops._c10d_functional.isend(tensor, global_dst, tag, group_name)
+    if _are_we_tracing():
+        return tensor
+    return _maybe_wrap_tensor(tensor)
+
+
+def irecv_inplace(
+    tensor: torch.Tensor,
+    src: int,
+    tag: int = 0,
+    group: dist.ProcessGroup | None = None,
+    group_src: int = -1,
+):
+    if group is None:
+        group = dist.group.WORLD
+    if group is None:
+        raise AssertionError("group cannot be None")
+    if group_src != -1:
+        if src is not None:
+            raise ValueError(
+                "Cannot specify both 'src' and 'group_src' args as per eager impl"
+            )
+        global_src = c10d.get_global_rank(group, group_src)
+    else:
+        global_src = src
+    group_name = _resolve_group_name(group)
+    tensor = torch.ops._c10d_functional.irecv(tensor, global_src, tag, group_name)
+    return _maybe_wrap_tensor(tensor)
+
+
+def batch_p2p_ops_inplace(
+    op_list: list[str],
+    peer_list: list[int],
+    tag_list: list[int],
+    tensors: list[torch.Tensor],
+    group_name: RANK_TYPES,
+):
+    if not dist.is_initialized():
+        raise AssertionError("torch.distributed must be initialized")
+    if group_name is None or group_name == "":
+        group_name = c10d._get_default_group()
+    resolved = _resolve_group(group_name)
+    group_name = resolved if isinstance(resolved, str) else resolved.group_name
+    tensors = torch.ops._c10d_functional.batch_p2p_ops(
+        op_list, peer_list, tag_list, tensors, group_name
+    )
+    if _are_we_tracing():
+        return [
+            _maybe_wrap_tensor(t) if op == "irecv" else t
+            for op, t in zip(op_list, tensors)
+        ]
+    return list(map(_maybe_wrap_tensor, tensors))
+
+
+def _group_or_group_name(
+    group: dist.ProcessGroup | c10d.GroupName,
+) -> dist.ProcessGroup | c10d.GroupName:
+    if isinstance(group, str):
+        return group
+    elif dist.config.compile_on_one_rank:
+        return group
+    else:
+        return group.group_name
+
+
+from torch.distributed.distributed_c10d import (  # pyrefly: ignore  # deprecated; pyrefly: ignore [deprecated]
     _all_gather_base as legacy_all_gather_base,
     _reduce_scatter_base as legacy_reduce_scatter_base,
     all_gather as legacy_all_gather,
     all_gather_into_tensor as legacy_allgather,
     all_reduce as legacy_allreduce,
     all_to_all_single as legacy_all_to_all_single,
+    batch_isend_irecv as legacy_batch_p2p_ops,
+    irecv as legacy_irecv,
+    isend as legacy_isend,
     reduce_scatter_tensor as legacy_reducescatter,
 )
 
 
+# Dynamo remaps dist.* collectives to these wrappers via traceable_collective_remaps.
+# Each wrapper calls the in-place functional collective and returns None,
+# matching the return type of the original dist.* APIs when async_op=False.
+# async_op=True already graph-breaks in CollectiveFunctionRewriteVariable.
+# These must be module-level def statements (not closures from a decorator factory)
+# because _traceable_collectives_source resolves Dynamo guard sources by looking up
+# fn.__name__ as a module attribute — a def's __name__ matches its variable name
+# automatically, whereas a closure's would not.
+def _remapped_allgather(*args, **kwargs):
+    if not _are_we_tracing():
+        raise AssertionError("_remapped_allgather should only be called during tracing")
+    all_gather_tensor_inplace(*args, **kwargs)
+
+
+def _remapped_reducescatter(*args, **kwargs):
+    if not _are_we_tracing():
+        raise AssertionError(
+            "_remapped_reducescatter should only be called during tracing"
+        )
+    reduce_scatter_tensor_inplace(*args, **kwargs)
+
+
+def _remapped_allreduce(*args, **kwargs):
+    if not _are_we_tracing():
+        raise AssertionError("_remapped_allreduce should only be called during tracing")
+    all_reduce_inplace(*args, **kwargs)
+
+
+def _remapped_all_to_all_single(*args, **kwargs):
+    if not _are_we_tracing():
+        raise AssertionError(
+            "_remapped_all_to_all_single should only be called during tracing"
+        )
+    all_to_all_inplace(*args, **kwargs)
+
+
+def _remapped_all_gather(*args, **kwargs):
+    if not _are_we_tracing():
+        raise AssertionError(
+            "_remapped_all_gather should only be called during tracing"
+        )
+    all_gather_inplace(*args, **kwargs)
+
+
+def _remapped_isend(*args, **kwargs):
+    if not _are_we_tracing():
+        raise AssertionError("_remapped_isend should only be called during tracing")
+    return isend_inplace(*args, **kwargs)
+
+
+def _remapped_irecv(*args, **kwargs):
+    if not _are_we_tracing():
+        raise AssertionError("_remapped_irecv should only be called during tracing")
+    return irecv_inplace(*args, **kwargs)
+
+
+def _remapped_batch_p2p_ops(*args, **kwargs):
+    if not _are_we_tracing():
+        raise AssertionError(
+            "_remapped_batch_p2p_ops should only be called during tracing"
+        )
+    return batch_p2p_ops_inplace(*args, **kwargs)
+
+
 # This dict should contain sets of functions that dynamo is allowed to remap.
 # Functions in this set should accept the same args/kwargs 1:1 as their mapping.
 traceable_collective_remaps = {
-    legacy_allgather: all_gather_tensor_inplace,  # type: ignore[has-type]
-    legacy_reducescatter: reduce_scatter_tensor_inplace,  # type: ignore[has-type]
-    legacy_allreduce: all_reduce_inplace,  # type: ignore[has-type]
-    legacy_all_to_all_single: all_to_all_inplace,  # type: ignore[has-type]
-    legacy_all_gather: all_gather_inplace,  # type: ignore[has-type]
-    legacy_reduce_scatter_base: reduce_scatter_tensor_inplace,  # type: ignore[has-type]
-    legacy_all_gather_base: all_gather_tensor_inplace,  # type: ignore[has-type]
+    legacy_allgather: _remapped_allgather,
+    legacy_reducescatter: _remapped_reducescatter,
+    legacy_allreduce: _remapped_allreduce,
+    legacy_all_to_all_single: _remapped_all_to_all_single,
+    legacy_all_gather: _remapped_all_gather,
+    legacy_reduce_scatter_base: _remapped_reducescatter,
+    legacy_all_gather_base: _remapped_allgather,
+    legacy_isend: _remapped_isend,
+    legacy_irecv: _remapped_irecv,
+    legacy_batch_p2p_ops: _remapped_batch_p2p_ops,
 }
diff --git a/torch/distributed/_functional_collectives_impl.py b/torch/distributed/_functional_collectives_impl.py
index fcb659b74bc05..3d6979a330469 100644
--- a/torch/distributed/_functional_collectives_impl.py
+++ b/torch/distributed/_functional_collectives_impl.py
@@ -115,3 +115,23 @@ def _all_to_all_single(
 
 def _wait_tensor(tensor: torch.Tensor) -> torch.Tensor:
     return torch.ops._c10d_functional.wait_tensor(tensor)
+
+
+def _isend(tensor: torch.Tensor, dst: int, tag: str, group_name):
+    return torch.ops._c10d_functional.isend(tensor, dst, tag, group_name)
+
+
+def _irecv(tensor: torch.Tensor, src: int, tag: str, group_name):
+    return torch.ops._c10d_functional.irecv(tensor, src, tag, group_name)
+
+
+def _batch_p2p_ops(
+    op_list: list[str],
+    peer_list: list[int],
+    tag_list: list[int],
+    tensors: list[torch.Tensor],
+    group_name: str,
+):
+    return torch.ops._c10d_functional.batch_p2p_ops(
+        op_list, peer_list, tag_list, tensors, group_name
+    )
diff --git a/torch/distributed/_local_tensor/__init__.py b/torch/distributed/_local_tensor/__init__.py
index 408d546dfe954..d769da74d7613 100644
--- a/torch/distributed/_local_tensor/__init__.py
+++ b/torch/distributed/_local_tensor/__init__.py
@@ -302,16 +302,16 @@ def _combine_any_rank_results(rank_results: dict[int, Any]) -> Any:
         return _combine_int_rank_results(rank_results)
 
     if isinstance(any_v, torch.device):
-        assert all(v.type == any_v.type for v in rank_results.values()), (
-            "device type should be the same"
-        )
+        if not all(v.type == any_v.type for v in rank_results.values()):
+            raise AssertionError("device type should be the same")
         # Just use the first device - the device type is what matters,
         # and LocalTensorMode runs on a single physical device anyway
         return any_v
 
-    assert all(v == any_v for v in rank_results.values()), (
-        "Non Tensor or int rank results must be equal for all ranks"
-    )
+    if not all(v == any_v for v in rank_results.values()):
+        raise AssertionError(
+            "Non Tensor or int rank results must be equal for all ranks"
+        )
 
     return any_v
 
@@ -363,11 +363,13 @@ def _for_each_rank_run_func(
     default_value: Tensor | None = None
     for r in sorted(ranks):
         if use_per_rank_rng:
-            assert lm is not None
+            if lm is None:
+                raise AssertionError
             if r in lm._per_rank_rng_states:
                 _set_rng_state(*lm._per_rank_rng_states[r])
         else:
-            assert global_rng_state is not None
+            if global_rng_state is None:
+                raise AssertionError
             _set_rng_state(*global_rng_state)
 
         rank_flat_args = [_map_to_rank_local_val(a, r) for a in flat_args]
@@ -380,7 +382,8 @@ def _for_each_rank_run_func(
         flat_rank_rets[r] = rank_ret
 
         if use_per_rank_rng:
-            assert lm is not None
+            if lm is None:
+                raise AssertionError
             lm._per_rank_rng_states[r] = _get_rng_state()
 
         if default_value is None and func is torch.ops.aten.split.Tensor:
@@ -548,22 +551,26 @@ def ne(self, other: "int | LocalIntNode | ConstantIntNode") -> bool | SymBool:
 
     def ge(self, other: "int | LocalIntNode | ConstantIntNode") -> bool | SymBool:
         r = {self._local_ints[r] >= _int_on_rank(other, r) for r in self._local_ints}
-        assert len(r) == 1, (self, other)
+        if len(r) != 1:
+            raise AssertionError((self, other))
         return torch._C._get_constant_bool_symnode(next(iter(r)))
 
     def le(self, other: "int | LocalIntNode | ConstantIntNode") -> bool | SymBool:
         r = {self._local_ints[r] <= _int_on_rank(other, r) for r in self._local_ints}
-        assert len(r) == 1, (self, other)
+        if len(r) != 1:
+            raise AssertionError((self, other))
         return torch._C._get_constant_bool_symnode(next(iter(r)))
 
     def gt(self, other: "int | LocalIntNode | ConstantIntNode") -> bool | SymBool:
         r = {self._local_ints[r] > _int_on_rank(other, r) for r in self._local_ints}
-        assert len(r) == 1, (self, other)
+        if len(r) != 1:
+            raise AssertionError((self, other))
         return torch._C._get_constant_bool_symnode(next(iter(r)))
 
     def lt(self, other: "int | LocalIntNode | ConstantIntNode") -> bool | SymBool:
         r = {self._local_ints[r] < _int_on_rank(other, r) for r in self._local_ints}
-        assert len(r) == 1, (self, other)
+        if len(r) != 1:
+            raise AssertionError((self, other))
         return torch._C._get_constant_bool_symnode(next(iter(r)))
 
     def wrap_int(self, num: int) -> "LocalIntNode | ConstantIntNode":
@@ -634,7 +641,8 @@ def set_rng_state(self, state):
         """
         if isinstance(state, LocalTensor):
             lm = enabled_local_tensor_mode()
-            assert lm is not None
+            if lm is None:
+                raise AssertionError
 
             # Similar to get_rng_state but in reverse: we need to convert from
             # device-specific tensor format to full state tuple format.
@@ -689,7 +697,8 @@ def _set_pre_op_offset(self, state, spec) -> None:
         from torch.distributed.tensor.placement_types import Shard
 
         lm = enabled_local_tensor_mode()
-        assert lm is not None
+        if lm is None:
+            raise AssertionError
 
         state._per_rank_offsets = {}
 
@@ -740,7 +749,8 @@ def _set_post_op_offset(self, state, spec, old_offset) -> None:
         from torch.distributed.tensor._ops.utils import prod
 
         lm = enabled_local_tensor_mode()
-        assert lm is not None
+        if lm is None:
+            raise AssertionError
 
         dtensor_shape = spec.shape
         numel = prod(dtensor_shape)
@@ -767,7 +777,8 @@ def _set_post_op_offset(self, state, spec, old_offset) -> None:
     def _distribute_region(self, spec, generator=None):
         """Context manager for LocalTensor mode distribute region."""
         lm = enabled_local_tensor_mode()
-        assert lm is not None
+        if lm is None:
+            raise AssertionError
 
         # get base state
         if generator is not None:
@@ -787,7 +798,8 @@ def _distribute_region(self, spec, generator=None):
                 any_rank_cpu, any_rank_cuda = any_rank_state
 
                 if self._device.type == "cuda":
-                    assert self._device.index in any_rank_cuda
+                    if self._device.index not in any_rank_cuda:
+                        raise AssertionError
                     any_rank_device_state = any_rank_cuda[self._device.index]
                 else:
                     any_rank_device_state = any_rank_cpu
@@ -867,15 +879,19 @@ def _compute_local_tensor_meta(
     # Assert that all tensors have the same dtype, layout and dispatch keys. Due
     # to uneven sharding, it is possible that tensors will have different shapes.
     for local_tensor in it:
-        assert dtype == local_tensor.dtype, (
-            "Tensors representing LocalTensor shards must have the same dtype"
-        )
-        assert layout == local_tensor.layout, (
-            "Tensors representing LocalTensor shards must have the same layout"
-        )
-        assert extra_dispatch_keys == _get_extra_dispatch_keys(local_tensor), (
-            "Tensors representing LocalTensor shards must have the same set of extra dispatch keys"
-        )
+        if dtype != local_tensor.dtype:
+            raise AssertionError(
+                "Tensors representing LocalTensor shards must have the same dtype"
+            )
+        if layout != local_tensor.layout:
+            raise AssertionError(
+                "Tensors representing LocalTensor shards must have the same layout"
+            )
+        if extra_dispatch_keys != _get_extra_dispatch_keys(local_tensor):
+            raise AssertionError(
+                "Tensors representing LocalTensor shards must have the "
+                "same set of extra dispatch keys"
+            )
 
     # Compute shape/stride.  We allow for non-SPMD'ness here
     local_shapes: dict[int, dict[int, int]] = defaultdict(dict)  # dim => rank => size
@@ -989,7 +1005,6 @@ def __deepcopy__(self, memo: dict[Any, Any] | None) -> "LocalTensor":
     def __repr__(self) -> str:  # type: ignore[override]
         parts = []
         for k, v in self._local_tensors.items():
-            # pyrefly: ignore [bad-argument-type]
             parts.append(f"  {k}: {v}")
         tensors_str = ",\n".join(parts)
         return f"LocalTensor(\n{tensors_str}\n)"
@@ -1017,9 +1032,10 @@ def __tensor_unflatten__(
         outer_size: torch.Size,
         outer_stride: tuple[int, ...],
     ) -> "LocalTensor":
-        assert flatten_spec is not None, (
-            "Expecting spec to be not None from `__tensor_flatten__` return value!"
-        )
+        if flatten_spec is None:
+            raise AssertionError(
+                "Expecting spec to be not None from `__tensor_flatten__` return value!"
+            )
         local_tensors = {
             _from_local_tensor_attr(a): t for a, t in inner_tensors.items()
         }
@@ -1046,9 +1062,8 @@ def __torch_dispatch__(  # type: ignore[override]
                 local_tensor = arg
                 break
 
-        assert local_tensor is not None, (
-            "At least one of the arguments must be a LocalTensor"
-        )
+        if local_tensor is None:
+            raise AssertionError("At least one of the arguments must be a LocalTensor")
 
         # Check for unrecognized tensor subclasses (but allow regular tensors and scalars)
         has_unrecognized_types = _check_for_subclass(flat_args)
@@ -1119,9 +1134,8 @@ def reconcile(self) -> torch.Tensor:
 
         # Force all local tensor shards across ranks to be the same
         equal_obj = self._equal_local_tensors()
-        assert isinstance(equal_obj, torch.Tensor), (
-            "LocalTensor shards must be the same to reconcile"
-        )
+        if not isinstance(equal_obj, torch.Tensor):
+            raise AssertionError("LocalTensor shards must be the same to reconcile")
         cl = equal_obj.clone().detach()
         cl.requires_grad_(self.requires_grad)
         return cl
@@ -1232,7 +1246,8 @@ def __init__(self, ranks: int | frozenset[int]):
             # assume is world size
             self.ranks = frozenset(range(ranks))
         else:
-            assert isinstance(ranks, frozenset)
+            if not isinstance(ranks, frozenset):
+                raise AssertionError
             self.ranks = ranks
         self._disable = True
         # Used to store the patched DeviceMesh methods
@@ -1242,9 +1257,12 @@ def __init__(self, ranks: int | frozenset[int]):
         self._per_rank_rng_states: dict[
             int, tuple[torch.Tensor, dict[int, torch.Tensor]]
         ] = {}
-        # Cache for get_coordinate results, keyed by mesh id
-        # Protected by _coordinate_cache_lock for thread safety in MPMD contexts
-        self._coordinate_cache: dict[int, list[SymInt]] = {}
+        # Cache for get_coordinate results, keyed by the inputs used to compute them:
+        # (ndim, flattened rank map, layout).
+        # Protected by _coordinate_cache_lock for thread safety in MPMD contexts.
+        self._coordinate_cache: dict[
+            tuple[int, tuple, torch.distributed._mesh_layout._MeshLayout], list[SymInt]
+        ] = {}
         self._coordinate_cache_lock = threading.Lock()
 
     def __enter__(self) -> "LocalTensorMode":
@@ -1317,9 +1335,11 @@ def __torch_dispatch__(
         # For LocalTensors, verify they have compatible ranks
         for a in flat_args:
             if isinstance(a, LocalTensor):
-                assert a._ranks <= self.ranks, (
-                    f"Input LocalTensor {a} must be configured for a subset of the LocalTensorMode ranks {self.ranks}"
-                )
+                if a._ranks > self.ranks:
+                    raise AssertionError(
+                        f"Input LocalTensor {a} must be configured for a "
+                        f"subset of the LocalTensorMode ranks {self.ranks}"
+                    )
 
         if func.overloadpacket == torch.ops.aten.dim:
             return len(args[0]._size)
@@ -1458,7 +1478,8 @@ def _any_local_rng_state(self) -> tuple[torch.Tensor, dict[int, torch.Tensor]]:
         return self._per_rank_rng_states[next(iter(self.ranks))]
 
     def _patch_device_mesh(self) -> None:
-        assert self._old_device_mesh_methods is None
+        if self._old_device_mesh_methods is not None:
+            raise AssertionError
         saved = {}
         for name in _PATCHED_DEVICE_MESH_METHODS:
             saved[name] = getattr(DeviceMesh, name)
@@ -1468,7 +1489,8 @@ def _patch_device_mesh(self) -> None:
 
     def _unpatch_device_mesh(self) -> None:
         saved, self._old_device_mesh_methods = self._old_device_mesh_methods, None
-        assert saved is not None
+        if saved is None:
+            raise AssertionError
         for name, value in saved.items():
             setattr(DeviceMesh, name, value)
 
@@ -1568,18 +1590,20 @@ def get_coordinate(self: DeviceMesh) -> list[SymInt] | None:
         # rank (therefore below we are patching get_rank method). We are trying to
         # limit the invasiveness of local tensor.
         lm = enabled_local_tensor_mode()
-        assert lm is not None, "Unexpectedly not in LocalTensorMode"
+        if lm is None:
+            raise AssertionError("Unexpectedly not in LocalTensorMode")
 
+        # Include all attributes used below in the cache key
+        cache_key = (self.ndim, self._flatten_rank_map, self._layout)
         # Check cache first (fast path without lock)
-        mesh_id = id(self)
-        if mesh_id in lm._coordinate_cache:
-            return lm._coordinate_cache[mesh_id]
+        if cache_key in lm._coordinate_cache:
+            return lm._coordinate_cache[cache_key]
 
         # Acquire lock for thread safety in MPMD contexts
         with lm._coordinate_cache_lock:
             # Double-check after acquiring lock
-            if mesh_id in lm._coordinate_cache:
-                return lm._coordinate_cache[mesh_id]
+            if cache_key in lm._coordinate_cache:
+                return lm._coordinate_cache[cache_key]
 
             coords: list[dict[int, int]] = [{} for _ in range(self.ndim)]
             # Clone rank_map to avoid "Cannot set version_counter for inference tensor"
@@ -1588,13 +1612,14 @@ def get_coordinate(self: DeviceMesh) -> list[SymInt] | None:
             for r in lm.ranks:
                 rank_tensor = self._layout.remap_to_tensor(rank_map)
                 rank_coords = (rank_tensor == r).nonzero().tolist()
-                assert len(rank_coords) == 1
+                if len(rank_coords) != 1:
+                    raise AssertionError
                 for d, c in enumerate(rank_coords[0][1:]):
                     coords[d][r] = c
 
             out = [torch.SymInt(LocalIntNode(c)) for c in coords]
             # Cache the result
-            lm._coordinate_cache[mesh_id] = out
+            lm._coordinate_cache[cache_key] = out
             # The output contains coordinates for each of the ranks with respect to
             # their meshes formed from root mesh and selecting the same dimensions
             # as the current mesh.
@@ -1608,19 +1633,22 @@ def _is_current_rank_part_of_mesh(self: DeviceMesh) -> bool:
     @staticmethod
     def _sym_get_coordinate(self: DeviceMesh, index: int) -> int:
         my_coordinate = self.get_coordinate()
-        assert my_coordinate is not None
+        if my_coordinate is None:
+            raise AssertionError
         return my_coordinate[index]
 
     @staticmethod
     def get_rank(self) -> int | SymInt:
         lm = enabled_local_tensor_mode()
-        assert lm is not None, "Unexpectedly not in LocalTensorMode"
+        if lm is None:
+            raise AssertionError("Unexpectedly not in LocalTensorMode")
         return torch.SymInt(LocalIntNode(local_ints={r: r for r in lm.ranks}))
 
     @staticmethod
     def get_local_rank(self, mesh_dim: int | str | None = None) -> int | SymInt:
         lm = enabled_local_tensor_mode()
-        assert lm is not None, "Unexpectedly not in LocalTensorMode"
+        if lm is None:
+            raise AssertionError("Unexpectedly not in LocalTensorMode")
 
         if self.ndim > 1 and mesh_dim is None:
             raise RuntimeError(
@@ -1637,7 +1665,8 @@ def get_local_rank(self, mesh_dim: int | str | None = None) -> int | SymInt:
         # get_coordinate returns a list of SymInt, one per mesh dimension
         # We need to extract the coordinate for the specified mesh_dim
         coords = _LocalDeviceMesh.get_coordinate(self)
-        assert coords is not None
+        if coords is None:
+            raise AssertionError
         return coords[mesh_dim]
 
 
@@ -1736,6 +1765,42 @@ def wrapper(*args: _P.args, **kwargs: _P.kwargs) -> _R:
     return wrapper
 
 
+def rank_map(cb: Callable[[int], Tensor]) -> Tensor:
+    """
+    Creates a tensor by mapping a callback over the current rank.
+
+    Under LocalTensorMode, calls cb(rank) for each simulated rank and returns
+    a LocalTensor. In real distributed (no LocalTensorMode), calls
+    cb(dist.get_rank()) and returns a plain Tensor.
+    """
+    lm = enabled_local_tensor_mode()
+    if lm is not None:
+        return lm.rank_map(cb)
+    else:
+        return cb(dist.get_rank())
+
+
+def tensor_map(tensor: Tensor, cb: Callable[[int, Tensor], Tensor | None]) -> Tensor:
+    """
+    Transforms a tensor by mapping a callback over the current rank and its
+    local shard.
+
+    Under LocalTensorMode, calls cb(rank, shard) for each simulated rank and
+    returns a LocalTensor. In real distributed (no LocalTensorMode), calls
+    cb(dist.get_rank(), tensor) and returns the result directly.
+    """
+    lm = enabled_local_tensor_mode()
+    if lm is not None:
+        if not isinstance(tensor, LocalTensor):
+            raise AssertionError(f"Expected LocalTensor, got {type(tensor)}")
+        return lm.tensor_map(tensor, cb)
+    else:
+        r = cb(dist.get_rank(), tensor)
+        if r is None:
+            raise AssertionError("callback returned None")
+        return r
+
+
 def maybe_disable_local_tensor_mode() -> contextlib.AbstractContextManager:
     """
     Context manager that disables LocalTensorMode for the duration of the context.
@@ -1871,7 +1936,8 @@ def __init__(
 
     def __enter__(self) -> "LocalRunnerMode":
         global _LOCAL_RUNNER_MODE
-        assert _LOCAL_RUNNER_MODE is None, "LocalRunnerMode is already running"
+        if _LOCAL_RUNNER_MODE is not None:
+            raise AssertionError("LocalRunnerMode is already running")
         _LOCAL_RUNNER_MODE = self
 
         global _PROCESS_MODE
@@ -1913,9 +1979,8 @@ def _release_run_lock(self) -> None:
         self._run_lock.release()
 
     def _assert_holds_run_lock(self) -> None:
-        assert self._run_id == LocalRunnerMode.runner_context.id, (
-            "Calling thread does not hold the run lock"
-        )
+        if self._run_id != LocalRunnerMode.runner_context.id:
+            raise AssertionError("Calling thread does not hold the run lock")
 
     def _get_recv_object(self, src: int, dst: int) -> object | None:
         peers = [src] if src != -1 else list(self._ranks)
@@ -1928,7 +1993,8 @@ def _get_recv_object(self, src: int, dst: int) -> object | None:
         return None
 
     def _signal_send(self, src: int, dst: int, obj: object) -> None:
-        assert obj is not None, "Cannot signal None"
+        if obj is None:
+            raise AssertionError("Cannot signal None")
         # Only a single thread a time executes so it is safe to mutate
         # read objects queue (executing thread is already holding the lock)
         self._recv_objects[dst][src].put(obj)
@@ -1950,7 +2016,8 @@ def _wait_recv(self, src: int, dst: int, post: Callable[[object], None]) -> None
     @staticmethod
     def current() -> "LocalRunnerMode":
         global _LOCAL_RUNNER_MODE
-        assert _LOCAL_RUNNER_MODE is not None, "LocalRunnerMode is not enabled"
+        if _LOCAL_RUNNER_MODE is None:
+            raise AssertionError("LocalRunnerMode is not enabled")
         return _LOCAL_RUNNER_MODE
 
 
@@ -1965,9 +2032,8 @@ class _LocalPhiloxState:
     """
 
     def __init__(self, state: torch.Tensor):
-        assert isinstance(state, LocalTensor), (
-            "_LocalPhiloxState requires a LocalTensor"
-        )
+        if not isinstance(state, LocalTensor):
+            raise AssertionError("_LocalPhiloxState requires a LocalTensor")
         self._local_tensor = state
         self._per_rank_states = {
             rank: local_state.to("cpu")
@@ -2048,7 +2114,8 @@ def apply_to_local_tensor_mode(self, device_handle) -> None:
         if not enabled_local_tensor_mode():
             return
 
-        assert hasattr(self, "_per_rank_offsets")
+        if not hasattr(self, "_per_rank_offsets"):
+            raise AssertionError
 
         for rank in sorted(self._per_rank_states.keys()):
             offset_value = self._per_rank_offsets[rank]
diff --git a/torch/distributed/_local_tensor/_c10d.py b/torch/distributed/_local_tensor/_c10d.py
index 7d326e8827347..9d812eea01bfc 100644
--- a/torch/distributed/_local_tensor/_c10d.py
+++ b/torch/distributed/_local_tensor/_c10d.py
@@ -44,10 +44,9 @@ def _indices_to_layout(indices: list[int]) -> tuple[tuple[int, ...], tuple[int,
     diffs = [indices[i] - indices[i - 1] for i in range(1, len(indices))]
     last_stride = _gcd_list(diffs)
 
-    assert last_stride != 0, (
-        # This case should not be reached if indices are unique and sorted.
-        "Cannot determine stride; indices may not be unique."
-    )
+    # This case should not be reached if indices are unique and sorted.
+    if last_stride == 0:
+        raise AssertionError("Cannot determine stride; indices may not be unique.")
 
     # Identify the starting index of each "row" in the last dimension.
     # An index starts a new row if the preceding index (index - stride) is not present.
@@ -58,10 +57,11 @@ def _indices_to_layout(indices: list[int]) -> tuple[tuple[int, ...], tuple[int,
             higher_dim_indices.append(index)
 
     # From the number of rows, we can deduce the shape of the last dimension.
-    assert len(indices) % len(higher_dim_indices) == 0, (
-        "Indices do not form a regular grid. "
-        f"Found {len(higher_dim_indices)} subgroups for {len(indices)} total elements."
-    )
+    if len(indices) % len(higher_dim_indices) != 0:
+        raise AssertionError(
+            "Indices do not form a regular grid. "
+            f"Found {len(higher_dim_indices)} subgroups for {len(indices)} total elements."
+        )
     last_shape = len(indices) // len(higher_dim_indices)
 
     # Recurse on the higher-dimensional indices (the start of each row).
@@ -84,10 +84,12 @@ def _prepare_collective_groups(
     )
 
     ranks = torch.distributed.get_process_group_ranks(process_group)
-    assert ranks
+    if not ranks:
+        raise AssertionError
     # TODO: We can handle permutations but the layout inference algorithm will
     # lose the permutation so we will have to reapply it
-    assert ranks == sorted(ranks), ranks
+    if ranks != sorted(ranks):
+        raise AssertionError(ranks)
     offset = ranks[0]
     ranks = [r - offset for r in ranks]
 
@@ -114,7 +116,8 @@ def _local_functional_all_gather_into_tensor(
         _resolve_process_group(group_name)
     )
 
-    assert isinstance(tensor, LocalTensor), "Input tensor must be a LocalTensor"
+    if not isinstance(tensor, LocalTensor):
+        raise AssertionError("Input tensor must be a LocalTensor")
     output_local_tensors: dict[int, torch.Tensor] = {}
 
     for group_offset in group_offsets:
@@ -148,7 +151,8 @@ def _local_functional_reduce_scatter_tensor(
         _resolve_process_group(group_name)
     )
 
-    assert isinstance(tensor, LocalTensor), "Input tensor must be a LocalTensor"
+    if not isinstance(tensor, LocalTensor):
+        raise AssertionError("Input tensor must be a LocalTensor")
     output_local_tensors: dict[int, torch.Tensor] = {}
 
     for group_offset in group_offsets:
@@ -191,7 +195,8 @@ def _local_functional_shard_dim_alltoall(
         _resolve_process_group(group_name)
     )
 
-    assert isinstance(tensor, LocalTensor), "Input tensor must be a LocalTensor"
+    if not isinstance(tensor, LocalTensor):
+        raise AssertionError("Input tensor must be a LocalTensor")
     output_local_tensors: dict[int, torch.Tensor] = {}
 
     for group_offset in group_offsets:
@@ -239,7 +244,8 @@ def _local_functional_all_to_all_single(
         _resolve_process_group(group_name)
     )
 
-    assert isinstance(tensor, LocalTensor), "Input tensor must be a LocalTensor"
+    if not isinstance(tensor, LocalTensor):
+        raise AssertionError("Input tensor must be a LocalTensor")
 
     split_local_sizes: dict[int, list[int]] = {}
     for input_split_size in input_split_sizes:
@@ -293,8 +299,10 @@ def _local_broadcast_(
     # "int root_rank, int root_tensor, bool async_op=True, int timeout=-1) -> (Tensor[], __torch__.torch.classes.c10d.Work)"
     from . import LocalTensor
 
-    assert len(tensors) == 1
-    assert root_tensor == 0
+    if len(tensors) != 1:
+        raise AssertionError
+    if root_tensor != 0:
+        raise AssertionError
     tensor = tensors[0]
 
     ranks, group_offsets, offset = _prepare_collective_groups(process_group_so)
@@ -303,7 +311,8 @@ def _local_broadcast_(
     # the same relative to others
     relative_root_rank = root_rank - offset
 
-    assert isinstance(tensor, LocalTensor), "Input tensor must be a LocalTensor"
+    if not isinstance(tensor, LocalTensor):
+        raise AssertionError("Input tensor must be a LocalTensor")
 
     for group_offset in group_offsets:
         # For the tensors in this group [group_offset + r for r in ranks]
@@ -353,7 +362,8 @@ def _local_reduce(
     if reduce_op == ReduceOp.AVG or reduce_op == "avg":
         return functools.reduce(operator.add, tensors) / len(tensors)
     else:
-        assert op is not None
+        if op is None:
+            raise AssertionError
         return functools.reduce(op, tensors)
 
 
@@ -370,13 +380,15 @@ def _local_all_reduce_(
     # "int timeout=-1) -> (Tensor[], __torch__.torch.classes.c10d.Work)");
     from . import LocalTensor
 
-    assert len(tensors) == 1
+    if len(tensors) != 1:
+        raise AssertionError
     tensor = tensors[0]
     reduce_op = reduce_op_so.op()  # type: ignore[attr-defined]
 
     ranks, group_offsets, _offset = _prepare_collective_groups(process_group_so)
 
-    assert isinstance(tensor, LocalTensor), "Input tensor must be a LocalTensor"
+    if not isinstance(tensor, LocalTensor):
+        raise AssertionError("Input tensor must be a LocalTensor")
 
     for group_offset in group_offsets:
         # For the tensors in this group [group_offset + r for r in ranks]
@@ -423,7 +435,8 @@ def _local_allreduce_coalesced_(
 
         # For each tensor, perform the reduction operation
         for tensor in tensors:
-            assert isinstance(tensor, LocalTensor), "Input tensor must be a LocalTensor"
+            if not isinstance(tensor, LocalTensor):
+                raise AssertionError("Input tensor must be a LocalTensor")
             if not all(rank in tensor._local_tensors for rank in group_ranks):
                 continue
             # Collect tensors from the specified ranks in this group
@@ -468,12 +481,10 @@ def _local_reduce_scatter_tensor_coalesced_(
 
         # For each tensor, perform the reduction operation
         for input_tensor, output_tensor in zip(input_tensors, output_tensors):
-            assert isinstance(input_tensor, LocalTensor), (
-                "Input tensor must be a LocalTensor"
-            )
-            assert isinstance(output_tensor, LocalTensor), (
-                "Output tensor must be a LocalTensor"
-            )
+            if not isinstance(input_tensor, LocalTensor):
+                raise AssertionError("Input tensor must be a LocalTensor")
+            if not isinstance(output_tensor, LocalTensor):
+                raise AssertionError("Output tensor must be a LocalTensor")
             if not all(rank in input_tensor._local_tensors for rank in group_ranks):
                 continue
             if not all(rank in output_tensor._local_tensors for rank in group_ranks):
@@ -513,8 +524,10 @@ def _local_allgather_base_(
 
     ranks, group_offsets, _offset = _prepare_collective_groups(process_group_so)
 
-    assert isinstance(output_tensor, LocalTensor), "Output tensor must be a LocalTensor"
-    assert isinstance(input_tensor, LocalTensor), "Input tensor must be a LocalTensor"
+    if not isinstance(output_tensor, LocalTensor):
+        raise AssertionError("Output tensor must be a LocalTensor")
+    if not isinstance(input_tensor, LocalTensor):
+        raise AssertionError("Input tensor must be a LocalTensor")
 
     for group_offset in group_offsets:
         group_ranks = [group_offset + r for r in ranks]
@@ -555,8 +568,10 @@ def _local_reduce_scatter_base_(  # type: ignore[no-untyped-def]
     reduce_op = reduce_op_so.op()  # type: ignore[attr-defined]
     ranks, group_offsets, _offset = _prepare_collective_groups(process_group_so)
 
-    assert isinstance(output_tensor, LocalTensor), "Output tensor must be a LocalTensor"
-    assert isinstance(input_tensor, LocalTensor), "Input tensor must be a LocalTensor"
+    if not isinstance(output_tensor, LocalTensor):
+        raise AssertionError("Output tensor must be a LocalTensor")
+    if not isinstance(input_tensor, LocalTensor):
+        raise AssertionError("Input tensor must be a LocalTensor")
 
     for group_offset in group_offsets:
         group_ranks = [group_offset + r for r in ranks]
@@ -598,8 +613,10 @@ def _local_all_gather_(
 
     from . import LocalTensor
 
-    assert len(output_tensors) == 1
-    assert len(input_tensors) == 1
+    if len(output_tensors) != 1:
+        raise AssertionError
+    if len(input_tensors) != 1:
+        raise AssertionError
 
     input_tensor = input_tensors[0]
     # pyrefly: ignore [bad-assignment]
@@ -608,9 +625,8 @@ def _local_all_gather_(
     ranks, group_offsets, _offset = _prepare_collective_groups(process_group_so)
 
     for i in range(len(output_tensors)):
-        assert isinstance(output_tensors[i], LocalTensor), (
-            "Output tensor must be a LocalTensor"
-        )
+        if not isinstance(output_tensors[i], LocalTensor):
+            raise AssertionError("Output tensor must be a LocalTensor")
 
     for group_offset in group_offsets:
         # For the tensors in this group [group_offset + r for r in ranks]
@@ -647,9 +663,11 @@ def _local_allgather_into_tensor_coalesced_(
 
     # Each output tensor should be sized to hold all gathered inputs
     # outputs[i] will contain all inputs[i] from all ranks
-    assert len(output_tensors) == len(input_tensors), (
-        f"Number of outputs ({len(output_tensors)}) must match number of inputs ({len(input_tensors)})"
-    )
+    if len(output_tensors) != len(input_tensors):
+        raise AssertionError(
+            f"Number of outputs ({len(output_tensors)}) must match "
+            f"number of inputs ({len(input_tensors)})"
+        )
 
     for group_offset in group_offsets:
         # For the tensors in this group [group_offset + r for r in ranks]
@@ -658,12 +676,10 @@ def _local_allgather_into_tensor_coalesced_(
 
         # For each input/output pair
         for input_tensor, output_tensor in zip(input_tensors, output_tensors):
-            assert isinstance(input_tensor, LocalTensor), (
-                "Input tensor must be a LocalTensor"
-            )
-            assert isinstance(output_tensor, LocalTensor), (
-                "Output tensor must be a LocalTensor"
-            )
+            if not isinstance(input_tensor, LocalTensor):
+                raise AssertionError("Input tensor must be a LocalTensor")
+            if not isinstance(output_tensor, LocalTensor):
+                raise AssertionError("Output tensor must be a LocalTensor")
 
             if not all(rank in input_tensor._local_tensors for rank in group_ranks):
                 continue
@@ -718,8 +734,10 @@ def _local_scatter_(
 
     from . import LocalTensor
 
-    assert len(output_tensors) == 1
-    assert len(input_tensors) == 1
+    if len(output_tensors) != 1:
+        raise AssertionError
+    if len(input_tensors) != 1:
+        raise AssertionError
     output_tensor = output_tensors[0]
     # pyrefly: ignore [bad-assignment]
     input_tensors = input_tensors[0]
@@ -730,8 +748,10 @@ def _local_scatter_(
     # the same relative to others
     relative_root_rank = root_rank - offset
 
-    assert isinstance(output_tensor, LocalTensor), "Output tensor must be a LocalTensor"
-    assert len(ranks) == len(input_tensors), (ranks, input_tensors)
+    if not isinstance(output_tensor, LocalTensor):
+        raise AssertionError("Output tensor must be a LocalTensor")
+    if len(ranks) != len(input_tensors):
+        raise AssertionError((ranks, input_tensors))
 
     for group_offset in group_offsets:
         # For the tensors in this group [group_offset + r for r in ranks]
@@ -743,7 +763,8 @@ def _local_scatter_(
         # Root rank scatters its input tensors to all ranks in this group
         for i, rank in enumerate(group_ranks):
             input_tensor = input_tensors[i]
-            assert isinstance(input_tensor, LocalTensor)
+            if not isinstance(input_tensor, LocalTensor):
+                raise AssertionError
             # Each rank i gets the i-th input tensor from the root
             source_tensor = input_tensor._local_tensors[
                 group_offset + relative_root_rank
@@ -770,10 +791,11 @@ def _local_alltoall_(
 
     ranks, group_offsets, _offset = _prepare_collective_groups(process_group_so)
 
-    assert len(input_tensors) == len(output_tensors) == len(ranks), (
-        f"Number of input tensors ({len(input_tensors)}), "
-        f"output tensors ({len(output_tensors)}), and ranks ({len(ranks)}) must match"
-    )
+    if not (len(input_tensors) == len(output_tensors) == len(ranks)):
+        raise AssertionError(
+            f"Number of input tensors ({len(input_tensors)}), "
+            f"output tensors ({len(output_tensors)}), and ranks ({len(ranks)}) must match"
+        )
 
     for group_offset in group_offsets:
         # For the tensors in this group [group_offset + r for r in ranks]
@@ -783,18 +805,16 @@ def _local_alltoall_(
         # In alltoall, rank i sends input_tensors[j] to rank j and receives into output_tensors[i] from rank j
         for i, rank_i in enumerate(group_ranks):
             output_tensor = output_tensors[i]
-            assert isinstance(output_tensor, LocalTensor), (
-                "Output tensor must be a LocalTensor"
-            )
+            if not isinstance(output_tensor, LocalTensor):
+                raise AssertionError("Output tensor must be a LocalTensor")
 
             if not all(rank in output_tensor._local_tensors for rank in group_ranks):
                 continue
 
             for j, rank_j in enumerate(group_ranks):
                 input_tensor = input_tensors[j]
-                assert isinstance(input_tensor, LocalTensor), (
-                    "Input tensor must be a LocalTensor"
-                )
+                if not isinstance(input_tensor, LocalTensor):
+                    raise AssertionError("Input tensor must be a LocalTensor")
 
                 if not all(rank in input_tensor._local_tensors for rank in group_ranks):
                     continue
@@ -824,8 +844,10 @@ def _local_alltoall_base_(
 
     ranks, group_offsets, _offset = _prepare_collective_groups(process_group_so)
 
-    assert isinstance(input_tensor, LocalTensor), "Input tensor must be a LocalTensor"
-    assert isinstance(output_tensor, LocalTensor), "Output tensor must be a LocalTensor"
+    if not isinstance(input_tensor, LocalTensor):
+        raise AssertionError("Input tensor must be a LocalTensor")
+    if not isinstance(output_tensor, LocalTensor):
+        raise AssertionError("Output tensor must be a LocalTensor")
     # Convert split sizes to lists if they aren't already
     if output_split_sizes is not None:
         output_split_sizes = list(output_split_sizes)
@@ -909,7 +931,8 @@ def _local_barrier(
     # Barrier is a synchronization primitive - in local simulation,
     # we don't need to do any actual work since all "ranks" are in the same process
     # Just validate that the tensor is a LocalTensor
-    assert isinstance(tensor, LocalTensor)
+    if not isinstance(tensor, LocalTensor):
+        raise AssertionError
 
     # In a real distributed setting, barrier would synchronize all processes
     # In local simulation, this is essentially a no-op since all ranks are local
@@ -933,7 +956,8 @@ def _local_monitored_barrier_(
     # Monitored barrier is a synchronization primitive with monitoring - in local simulation,
     # we don't need to do any actual work since all "ranks" are in the same process
     # Just validate that the tensor is a LocalTensor
-    assert isinstance(tensor, LocalTensor)
+    if not isinstance(tensor, LocalTensor):
+        raise AssertionError
 
     # In a real distributed setting, monitored barrier would synchronize all processes
     # and provide monitoring capabilities. In local simulation, this is essentially a no-op
@@ -952,10 +976,12 @@ def _local_send(
 
     from . import LocalRunnerMode, LocalTensor
 
-    assert len(tensors) == 1
+    if len(tensors) != 1:
+        raise AssertionError
     tensor = tensors[0]
 
-    assert isinstance(tensor, LocalTensor), "Input tensor must be a Tensor"
+    if not isinstance(tensor, LocalTensor):
+        raise AssertionError("Input tensor must be a Tensor")
     src = int(tensor.__src_rank__)
 
     LocalRunnerMode.current()._signal_send(src, dst, tensor._local_tensors[src])
@@ -975,16 +1001,20 @@ def _local_recv_(
     # "int src, int tag) -> __torch__.torch.classes.c10d.Work";
     from . import LocalRunnerMode, LocalTensor
 
-    assert len(tensors) == 1
+    if len(tensors) != 1:
+        raise AssertionError
     tensor = tensors[0]
 
-    assert isinstance(tensor, LocalTensor), "Input tensor must be a Tensor"
+    if not isinstance(tensor, LocalTensor):
+        raise AssertionError("Input tensor must be a Tensor")
     dst = int(tensor.__src_rank__)
 
     def _recv_and_store(timeout: timedelta) -> bool:
         def _wait_and_store(obj: object) -> None:
-            assert isinstance(obj, torch.Tensor), "Expected to receive a Tensor"
-            assert isinstance(tensor, LocalTensor), "Input tensor must be a Tensor"
+            if not isinstance(obj, torch.Tensor):
+                raise AssertionError("Expected to receive a Tensor")
+            if not isinstance(tensor, LocalTensor):
+                raise AssertionError("Input tensor must be a Tensor")
             tensor._local_tensors[dst] = obj
 
         LocalRunnerMode.current()._wait_recv(src, dst, _wait_and_store)
@@ -1030,9 +1060,11 @@ def local_p2p_op(
 
     from . import LocalIntNode
 
-    assert isinstance(dst.node, LocalIntNode), (
-        "Expected 'dst' to be a LocalIntNode where the value is the destination rank and key is the source rank"
-    )
+    if not isinstance(dst.node, LocalIntNode):
+        raise AssertionError(
+            "Expected 'dst' to be a LocalIntNode where the value is the "
+            "destination rank and key is the source rank"
+        )
 
     w = []
     for s, d in dst.node._local_ints.items():
diff --git a/torch/distributed/_mesh_layout.py b/torch/distributed/_mesh_layout.py
index 4838853b212ca..f6babf0f98cf7 100644
--- a/torch/distributed/_mesh_layout.py
+++ b/torch/distributed/_mesh_layout.py
@@ -305,9 +305,12 @@ def remap_to_tensor(self, rank_map: torch.Tensor) -> torch.Tensor:
         Returns:
             torch.Tensor: A tensor representing the actual device allocation from rank_map
         """
-        assert rank_map.ndim == 1
-        assert rank_map.is_contiguous()
-        assert rank_map.numel() >= self.cosize()
+        if rank_map.ndim != 1:
+            raise AssertionError
+        if not rank_map.is_contiguous():
+            raise AssertionError
+        if rank_map.numel() < self.cosize():
+            raise AssertionError
 
         complement_layout = self.complement(rank_map.numel())
 
diff --git a/torch/distributed/_ops/device_mesh.py b/torch/distributed/_ops/device_mesh.py
index b4e60b0ea1ade..61d4312c3bb04 100644
--- a/torch/distributed/_ops/device_mesh.py
+++ b/torch/distributed/_ops/device_mesh.py
@@ -59,5 +59,47 @@ def _runtime_compute_coordinate_on_dim_impl(full_mesh: torch.Tensor, index: int)
     rank = torch.distributed.get_rank()
     mesh = DeviceMesh._get_mesh_tensor_from_full_mesh(full_mesh)
     mesh_coords = DeviceMesh._compute_coordinates_from_mesh(mesh, rank)
-    assert mesh_coords is not None
+    if mesh_coords is None:
+        raise AssertionError
     return mesh_coords[index]
+
+
+def _get_flattened_submesh_impl(mesh: DeviceMesh, mesh_dims: list[int]) -> DeviceMesh:
+    from torch.distributed.tensor._redistribute import (
+        _get_flattened_mesh_by_layout_impl,
+    )
+
+    result = _get_flattened_mesh_by_layout_impl(mesh, tuple(mesh_dims))
+    if result is None:
+        raise ValueError(f"No flattened mesh found for mesh_dims={mesh_dims} on {mesh}")
+    return result
+
+
+@torch.library.custom_op("device_mesh::_get_flattened_submesh", mutates_args=())
+def _get_flattened_submesh(mesh: DeviceMesh, mesh_dims: list[int]) -> DeviceMesh:
+    return _get_flattened_submesh_impl(mesh, mesh_dims)
+
+
+@_get_flattened_submesh.register_fake
+def _get_flattened_submesh_fake(mesh: DeviceMesh, mesh_dims: list[int]) -> DeviceMesh:
+    return _get_flattened_submesh_impl(mesh, mesh_dims)
+
+
+def _get_submesh_impl(mesh: DeviceMesh, mesh_dims: list[int]) -> DeviceMesh:
+    all_dim_names = mesh._mesh_dim_names
+    if all_dim_names is None:
+        raise ValueError(f"Cannot slice mesh without dim names: {mesh}")
+    dim_names = tuple(all_dim_names[i] for i in mesh_dims)
+    if len(dim_names) == 1:
+        return mesh[dim_names[0]]
+    return mesh[dim_names]
+
+
+@torch.library.custom_op("device_mesh::_get_submesh", mutates_args=())
+def _get_submesh(mesh: DeviceMesh, mesh_dims: list[int]) -> DeviceMesh:
+    return _get_submesh_impl(mesh, mesh_dims)
+
+
+@_get_submesh.register_fake
+def _get_submesh_fake(mesh: DeviceMesh, mesh_dims: list[int]) -> DeviceMesh:
+    return _get_submesh_impl(mesh, mesh_dims)
diff --git a/torch/distributed/_pycute/int_tuple.py b/torch/distributed/_pycute/int_tuple.py
index 65ebf9f0c6060..830810ab48789 100644
--- a/torch/distributed/_pycute/int_tuple.py
+++ b/torch/distributed/_pycute/int_tuple.py
@@ -91,10 +91,12 @@ def product(a: IntTuple) -> int:
 
 def inner_product(a: IntTuple, b: IntTuple) -> int:
     if is_tuple(a) and is_tuple(b):  # tuple tuple
-        assert len(a) == len(b)
+        if len(a) != len(b):
+            raise AssertionError
         return sum(inner_product(x, y) for x, y in zip(a, b))
     else:  # "int" "int"
-        assert not is_tuple(a) and not is_tuple(b)
+        if is_tuple(a) or is_tuple(b):
+            raise AssertionError
         return a * b
 
 
@@ -108,7 +110,8 @@ def tuple_max(a: IntTuple) -> int:
 def elem_scale(a: IntTuple, b: IntTuple) -> IntTuple:
     if is_tuple(a):
         if is_tuple(b):  # tuple tuple
-            assert len(a) == len(b)
+            if len(a) != len(b):
+                raise AssertionError
             return tuple(elem_scale(x, y) for x, y in zip(a, b))
         else:  # tuple "int"
             raise AssertionError("Invalid combination: tuple with int")
@@ -123,7 +126,8 @@ def elem_scale(a: IntTuple, b: IntTuple) -> IntTuple:
 def shape_div(a: IntTuple, b: IntTuple) -> IntTuple:
     if is_tuple(a):
         if is_tuple(b):  # tuple tuple
-            assert len(a) == len(b)
+            if len(a) != len(b):
+                raise AssertionError
             return tuple(shape_div(x, y) for x, y in zip(a, b))
         else:  # tuple "int"
             # r = [shape_div(a[0],b)] + [shape_div(a[i],b := shape_div(b, product(a[i-1]))) for i in range(1,len(a))]
@@ -136,7 +140,8 @@ def shape_div(a: IntTuple, b: IntTuple) -> IntTuple:
         if is_tuple(b):  # "int" tuple
             return shape_div(a, product(b))
         else:  # "int" "int"
-            assert a % b == 0 or b % a == 0
+            if not (a % b == 0 or b % a == 0):
+                raise AssertionError
             return (a + b - 1) // b
 
 
@@ -145,7 +150,8 @@ def suffix_product(a: IntTuple, init: IntTuple = 1) -> IntTuple:
     # TODO: With all these length asserts, may want to create a zip_strict wrapper.
     if is_tuple(a):
         if is_tuple(init):  # tuple tuple
-            assert len(a) == len(init)
+            if len(a) != len(init):
+                raise AssertionError
             return tuple(suffix_product(x, i) for x, i in zip(a, init))
         else:  # tuple "int"
             # Process from right to left for lexicographic ordering
@@ -174,16 +180,19 @@ def idx2crd(idx: IntTuple, shape: IntTuple, stride: IntTuple | None = None) -> I
 
     if is_tuple(idx):
         if is_tuple(shape) and is_tuple(stride):  # tuple tuple tuple
-            assert len(idx) == len(shape) and len(stride) == len(shape)
+            if not (len(idx) == len(shape) and len(stride) == len(shape)):
+                raise AssertionError
             return tuple(idx2crd(i, s, d) for i, s, d in zip(idx, shape, stride))
         else:  # tuple "int" "int"
             raise AssertionError("Invalid combination: tuple with int stride")
     else:
         if is_tuple(shape) and is_tuple(stride):  # "int" tuple tuple
-            assert len(shape) == len(stride)
+            if len(shape) != len(stride):
+                raise AssertionError
             return tuple(idx2crd(idx, s, d) for s, d in zip(shape, stride))
         else:  # "int" "int" "int"
-            assert not is_tuple(shape) and not is_tuple(stride)
+            if is_tuple(shape) or is_tuple(stride):
+                raise AssertionError
             return (idx // stride) % shape  # all are ints after type checks
 
 
@@ -195,7 +204,8 @@ def crd2idx(
 
     if is_tuple(crd):
         if is_tuple(shape) and is_tuple(stride):  # tuple tuple tuple
-            assert len(crd) == len(shape) and len(stride) == len(shape)
+            if not (len(crd) == len(shape) and len(stride) == len(shape)):
+                raise AssertionError
             return sum(crd2idx(c, s, d) for c, s, d in zip(crd, shape, stride))
         else:  # tuple "int" "int"
             raise AssertionError(f"Invalid combination: crd={crd}, shape={shape}")
@@ -204,7 +214,8 @@ def crd2idx(
             crd = 0
 
         if is_tuple(shape) and is_tuple(stride):  # "int" tuple tuple
-            assert len(shape) == len(stride)
+            if len(shape) != len(stride):
+                raise AssertionError
             result = 0
             # Process from right to left for lexicographic ordering
             for i in range(len(shape) - 1, 0, -1):
@@ -214,7 +225,8 @@ def crd2idx(
                 result += crd2idx(crd, shape[0], stride[0])
             return result
         else:  # "int" "int" "int"
-            assert not is_tuple(shape) and not is_tuple(stride)
+            if is_tuple(shape) or is_tuple(stride):
+                raise AssertionError
             return crd * stride  # all are ints after type checks
 
 
@@ -224,17 +236,20 @@ def crd2crd(
 ) -> IntTuple:
     if is_tuple(crd):
         if is_tuple(dst_shape):  # tuple tuple
-            assert len(crd) == len(dst_shape)
+            if len(crd) != len(dst_shape):
+                raise AssertionError
             return tuple(crd2crd(x, y) for x, y in zip(crd, dst_shape))
         else:  # tuple "int"
             # Ambiguous unless we have src_shape
-            assert src_shape is not None
+            if src_shape is None:
+                raise AssertionError
             return crd2idx(crd, src_shape)
     else:
         if is_tuple(dst_shape):  # "int" tuple
             return idx2crd(crd, dst_shape)
         else:  # "int" "int"
-            assert crd < dst_shape
+            if crd >= dst_shape:
+                raise AssertionError
             return crd
 
 
@@ -242,7 +257,8 @@ def crd2crd(
 def slice_(crd: tuple | int | None, trg: tuple | int) -> tuple | int:
     if is_tuple(crd):
         if is_tuple(trg):  # tuple tuple
-            assert len(crd) == len(trg)
+            if len(crd) != len(trg):
+                raise AssertionError
             # match C++ behavior of `filter_tuple` using `tuple_cat(...)`
             return tuple(
                 chain(
diff --git a/torch/distributed/_pycute/layout.py b/torch/distributed/_pycute/layout.py
index d18547daec2bd..d14bb2d0ee24c 100644
--- a/torch/distributed/_pycute/layout.py
+++ b/torch/distributed/_pycute/layout.py
@@ -112,7 +112,8 @@ def __getitem__(self, i: int) -> Self:
         if is_tuple(self.shape):
             return Layout(self.shape[i], self.stride[i])  # type: ignore[index]
         else:
-            assert i == 0
+            if i != 0:
+                raise AssertionError
             return Layout(self.shape, self.stride)
 
     # size(layout)   Size of the domain
@@ -162,7 +163,8 @@ def cosize(layout: Layout) -> int:
 # Layout coalesce -- flatten and combine as many modes as possible while preserving the int-to-int function
 def coalesce(layout: Layout, profile: LayoutProfile = None) -> Layout:
     if is_tuple(profile):
-        assert len(layout) >= len(profile)
+        if len(layout) < len(profile):
+            raise AssertionError
         return make_layout(
             # pyrefly: ignore [bad-argument-type]
             chain(
@@ -204,7 +206,8 @@ def coalesce(layout: Layout, profile: LayoutProfile = None) -> Layout:
 # Layout filter -- replace all stride-0 modes with size-1 and then coalesce to remove them
 def filter(layout: Layout, profile: LayoutProfile = None) -> Layout:
     if is_tuple(profile):
-        assert len(layout) >= len(profile)
+        if len(layout) < len(profile):
+            raise AssertionError
         return make_layout(
             # pyrefly: ignore [bad-argument-type]
             chain(
@@ -235,7 +238,8 @@ def composition(layoutA: Layout, layoutB: LayoutInput) -> Layout:
     elif is_int(layoutB):
         return composition(layoutA, Layout(layoutB))
     elif is_tuple(layoutB):
-        assert len(layoutA) >= len(layoutB)
+        if len(layoutA) < len(layoutB):
+            raise AssertionError
         return make_layout(
             # pyrefly: ignore [bad-argument-type]
             chain(
@@ -265,7 +269,8 @@ def composition(layoutA: Layout, layoutB: LayoutInput) -> Layout:
         for curr_shape, curr_stride in zip(
             reversed(flatten(flat_A.shape)[1:]), reversed(flatten(flat_A.stride)[1:])
         ):
-            assert curr_shape % rest_stride == 0 or rest_stride % curr_shape == 0  # type: ignore[operator]
+            if not (curr_shape % rest_stride == 0 or rest_stride % curr_shape == 0):  # type: ignore[operator]
+                raise AssertionError
             new_shape = min(max(1, curr_shape // rest_stride), rest_shape)  # type: ignore[operator]
 
             if new_shape != 1:
@@ -310,7 +315,8 @@ def complement(layout: LayoutOrIntTuple, max_idx: int = 1) -> Layout:
 
         in_bound = current_idx <= shape * stride
         # To support symbolic value which can't be evaluated now
-        assert (type(in_bound) is not bool) or in_bound
+        if (type(in_bound) is bool) and not in_bound:
+            raise AssertionError
 
         result_shape.append(stride // current_idx)
         result_stride.append(current_idx)
@@ -372,7 +378,8 @@ def logical_divide(layoutA: Layout, layoutB: LayoutInput) -> Layout:
     elif is_int(layoutB):
         return logical_divide(layoutA, Layout(layoutB))
     elif is_tuple(layoutB):
-        assert len(layoutA) >= len(layoutB)
+        if len(layoutA) < len(layoutB):
+            raise AssertionError
         return make_layout(
             # pyrefly: ignore [bad-argument-type]
             chain(
@@ -398,7 +405,8 @@ def logical_product(layoutA: Layout, layoutB: LayoutInput) -> Layout:
     elif is_int(layoutB):
         return logical_divide(layoutA, Layout(layoutB))
     elif is_tuple(layoutB):
-        assert len(layoutA) >= len(layoutB)
+        if len(layoutA) < len(layoutB):
+            raise AssertionError
         return make_layout(
             # pyrefly: ignore [bad-argument-type]
             chain(
@@ -425,7 +433,8 @@ def hier_unzip(
     if layoutB is None:
         return make_layout(Layout(1, 0), layoutA)
     elif is_tuple(layoutB):
-        assert len(layoutA) >= len(layoutB)
+        if len(layoutA) < len(layoutB):
+            raise AssertionError
         # A layout with shape ((A,a),(B,b),(C,c))
         split = make_layout(
             hier_unzip(splitter, layoutA[i], layoutB[i])  # type: ignore[arg-type]
diff --git a/torch/distributed/_shard/sharded_tensor/_ops/_common.py b/torch/distributed/_shard/sharded_tensor/_ops/_common.py
index c5cbc67442933..433f6b810cb25 100644
--- a/torch/distributed/_shard/sharded_tensor/_ops/_common.py
+++ b/torch/distributed/_shard/sharded_tensor/_ops/_common.py
@@ -43,7 +43,7 @@ def decorator_sharded_func(wrapped_func):
         def wrapper(types, args=(), kwargs=None, pg=None):
             _basic_validation(op, args, kwargs)
 
-            # pyrefly: ignore [bad-index, index-error]
+            # pyrefly: ignore [bad-index]
             st = args[0]
             if kwargs is None:
                 kwargs = {}
@@ -93,7 +93,7 @@ def _register_sharded_op_on_local_shards(
     @_sharded_op_impl(op)
     @_sharded_op_common(op, early_stop_func, extra_check)
     def sharded_tensor_op_on_local_shards(types, args=(), kwargs=None, pg=None):
-        # pyrefly: ignore [bad-index, index-error]
+        # pyrefly: ignore [bad-index]
         st = args[0]
         st_metadata = st.metadata()
         local_shards = st.local_shards()
diff --git a/torch/distributed/_shard/sharded_tensor/_ops/init.py b/torch/distributed/_shard/sharded_tensor/_ops/init.py
index 453634272f5bf..58365997b8154 100644
--- a/torch/distributed/_shard/sharded_tensor/_ops/init.py
+++ b/torch/distributed/_shard/sharded_tensor/_ops/init.py
@@ -149,7 +149,7 @@ def tensor_creation_op(types, args=(), kwargs=None, pg=None):
         if kwargs is None:
             kwargs = {}
 
-        # pyrefly: ignore [bad-index, index-error]
+        # pyrefly: ignore [bad-index]
         st = args[0]
 
         new_st = creation_op(st.sharding_spec(), st.size(), *args[1:], **kwargs)  # type: ignore[operator]
diff --git a/torch/distributed/_shard/sharded_tensor/_ops/tensor_ops.py b/torch/distributed/_shard/sharded_tensor/_ops/tensor_ops.py
index 5aa0cb3131ca3..7ce0823458bc2 100644
--- a/torch/distributed/_shard/sharded_tensor/_ops/tensor_ops.py
+++ b/torch/distributed/_shard/sharded_tensor/_ops/tensor_ops.py
@@ -40,7 +40,7 @@
 # the device property on each rank
 @_sharded_op_impl(torch.Tensor.device.__get__)
 def tensor_device(types, args=(), kwargs=None, pg=None):
-    # pyrefly: ignore [bad-index, index-error]
+    # pyrefly: ignore [bad-index]
     self_st = args[0]
     # Validate types
     if not isinstance(self_st, ShardedTensor):
@@ -57,7 +57,7 @@ def tensor_device(types, args=(), kwargs=None, pg=None):
 
 @_sharded_op_impl(torch.Tensor.is_meta.__get__)  # type: ignore[attr-defined]
 def st_is_meta(types, args=(), kwargs=None, pg=None):
-    # pyrefly: ignore [bad-index, index-error]
+    # pyrefly: ignore [bad-index]
     return args[0].local_tensor().is_meta
 
 
@@ -198,7 +198,7 @@ def sharded_detach(args, kwargs, pg):
 
 @_sharded_op_impl(torch.Tensor.requires_grad_)
 def tensor_requires_grad_set(types, args=(), kwargs=None, pg=None):
-    # pyrefly: ignore [bad-index, index-error]
+    # pyrefly: ignore [bad-index]
     self_st = args[0]
     # Validate types
     if not isinstance(self_st, ShardedTensor):
diff --git a/torch/distributed/_shard/sharded_tensor/api.py b/torch/distributed/_shard/sharded_tensor/api.py
index a8e8677d6ae7c..7b2f39018693c 100644
--- a/torch/distributed/_shard/sharded_tensor/api.py
+++ b/torch/distributed/_shard/sharded_tensor/api.py
@@ -299,9 +299,7 @@ def _post_init(self):
         if self._init_rrefs:
             with _sharded_tensor_lock:
                 global _sharded_tensor_current_id, _sharded_tensor_map
-                # pyrefly: ignore [bad-assignment]
                 self._sharded_tensor_id = _sharded_tensor_current_id
-                # pyrefly: ignore [unsupported-operation]
                 _sharded_tensor_map[self._sharded_tensor_id] = weakref.ref(self)
                 _sharded_tensor_current_id += 1
 
@@ -438,7 +436,8 @@ def shard_size(shard_md):
         # collect sizes
         for shard_md in self.metadata().shards_metadata:
             shard_rank = cast(_remote_device, shard_md.placement).rank()
-            assert shard_rank is not None
+            if shard_rank is None:
+                raise AssertionError
 
             shard_placement[shard_md] = (shard_rank, rank_sizes[shard_rank])
             rank_sizes[shard_rank] += shard_size(shard_md)
@@ -446,7 +445,8 @@ def shard_size(shard_md):
 
         gather_list: list[torch.Tensor] | None
         if rank == dst:
-            assert out is not None
+            if out is None:
+                raise AssertionError
             if enforce_dtype:
                 # enforce_dtype is deprecated.  Do it for backward compatibility.
                 dtype = out.dtype
@@ -487,7 +487,8 @@ def shard_size(shard_md):
             return
         # In _validate_output_tensor_for_gather, we raise if out == None and rank == dst
         out = cast(torch.Tensor, out)
-        assert gather_list is not None
+        if gather_list is None:
+            raise AssertionError
 
         full_size = self.metadata().size
         dims = len(full_size)
@@ -592,12 +593,13 @@ def cuda(
 
         if device is not None:
             device = torch.device(device) if isinstance(device, str) else device
-            assert (
+            if not (
                 isinstance(device, torch.device)
                 and device.index == torch.cuda.current_device()
-            ), (
-                """Only device without device id (e.g. "cpu" or "cuda") is expected for ShardedTensor!"""
-            )
+            ):
+                raise AssertionError(
+                    """Only device without device id (e.g. "cpu" or "cuda") is expected for ShardedTensor!"""
+                )
 
         current_device = torch.device(torch.cuda.current_device())
         # returns a copy of ShardedTensor on CUDA current device
@@ -784,22 +786,24 @@ def _init_from_local_shards(
         if recalc_metadata:
             # for recalc use cases, we only support rw for now, limit the blast radius
             # will modify here once we support more sharding type
-            assert (
+            if not (
                 len(local_shards) > 0
                 and len(global_sharded_tensor_metadata.shards_metadata) > current_rank
-            ), (
-                f"# for metadata recalculation, local_shards must be larger than 0 "
-                f"actual:{len(local_shards)}, # glb metadata must be greater than any rank id, "
-                f"# metadata:{len(global_sharded_tensor_metadata.shards_metadata)}, rank id:{current_rank}"
-            )
+            ):
+                raise AssertionError(
+                    f"# for metadata recalculation, local_shards must be larger than 0 "
+                    f"actual:{len(local_shards)}, # glb metadata must be greater than any rank id, "
+                    f"# metadata:{len(global_sharded_tensor_metadata.shards_metadata)}, rank id:{current_rank}"
+                )
             local_md = [
                 shard_md
                 for shard_md in global_sharded_tensor_metadata.shards_metadata
                 if shard_md.placement.rank() == current_rank
             ]
-            assert len(local_md) == 1, (
-                f"should has and only has one metadata for local rank, actual:{local_md}"
-            )
+            if len(local_md) != 1:
+                raise AssertionError(
+                    f"should has and only has one metadata for local rank, actual:{local_md}"
+                )
             local_shards[0].metadata = local_md[0]
         tensor_properties = global_sharded_tensor_metadata.tensor_properties
 
@@ -1010,7 +1014,8 @@ def _raise_if_mismatch(expected, actual, prop_name, rank, is_property=False):
             shard_meta = shard.metadata
             local_shard_tensor = shard.tensor
             placement = shard_meta.placement
-            assert placement is not None, "Must specify placement for `Shard`!"
+            if placement is None:
+                raise AssertionError("Must specify placement for `Shard`!")
             rank = placement.rank()
             local_device = placement.device()
 
diff --git a/torch/distributed/_shard/sharded_tensor/utils.py b/torch/distributed/_shard/sharded_tensor/utils.py
index 9a11590d1a3bf..d65b61b2e113c 100644
--- a/torch/distributed/_shard/sharded_tensor/utils.py
+++ b/torch/distributed/_shard/sharded_tensor/utils.py
@@ -92,7 +92,8 @@ def _flatten_tensor_size(size) -> torch.Size:
 
 def _raise_if_mismatch(expected, actual, prop_name, ranks, is_local=True):
     if is_local:
-        assert isinstance(ranks, int)
+        if not isinstance(ranks, int):
+            raise AssertionError
         if expected != actual:
             raise ValueError(
                 f"Local shards' tensor {prop_name} property need to be the same on rank:{ranks}! "
@@ -101,7 +102,8 @@ def _raise_if_mismatch(expected, actual, prop_name, ranks, is_local=True):
             )
     else:
         # compare failure check across ranks, ranks list should have two rank
-        assert len(ranks) == 2
+        if len(ranks) != 2:
+            raise AssertionError
         if expected != actual:
             raise ValueError(
                 f"ShardedTensor {prop_name} property does not match from different ranks! "
@@ -116,7 +118,8 @@ def build_metadata_from_local_shards(
     current_rank: int,
     pg: c10d.ProcessGroup,
 ) -> ShardedTensorMetadata:
-    assert len(local_shards) > 0, "must have local shards!"
+    if len(local_shards) <= 0:
+        raise AssertionError("must have local shards!")
     local_shard_metadatas: list[ShardMetadata] = []
 
     first_shard_dtype = local_shards[0].tensor.dtype
diff --git a/torch/distributed/_shard/sharding_spec/chunk_sharding_spec.py b/torch/distributed/_shard/sharding_spec/chunk_sharding_spec.py
index 4d7b11b7c16c5..f91a1241cbeb6 100644
--- a/torch/distributed/_shard/sharding_spec/chunk_sharding_spec.py
+++ b/torch/distributed/_shard/sharding_spec/chunk_sharding_spec.py
@@ -182,8 +182,10 @@ def shard(
 
         # each rank should have local_tensor and local_metadata initialized if we build
         # the metadata list in a correct way.
-        assert local_tensor is not None
-        assert local_metadata is not None
+        if local_tensor is None:
+            raise AssertionError
+        if local_metadata is None:
+            raise AssertionError
 
         # Scatter the shards to all ranks in the pg
         # scatter takes the global rank as ``src``
@@ -200,7 +202,8 @@ def shard(
         if current_rank == src_rank:
             tensors_to_scatter_ = []
             for t in tensors_to_scatter:
-                assert isinstance(t, torch.Tensor)
+                if not isinstance(t, torch.Tensor):
+                    raise AssertionError
                 tensors_to_scatter_.append(t)
 
         dist.scatter(
diff --git a/torch/distributed/_shard/sharding_spec/chunk_sharding_spec_ops/_common.py b/torch/distributed/_shard/sharding_spec/chunk_sharding_spec_ops/_common.py
index 1d7b0be3c35cd..eecafaf8a4ec1 100644
--- a/torch/distributed/_shard/sharding_spec/chunk_sharding_spec_ops/_common.py
+++ b/torch/distributed/_shard/sharding_spec/chunk_sharding_spec_ops/_common.py
@@ -58,7 +58,7 @@ def _register_sharded_op_on_local_tensor(
     @custom_sharding_spec_op(ChunkShardingSpec, op)
     @_sharded_op_common(op, early_stop_func, extra_check)
     def sharded_tensor_op_on_local_tensor(types, args=(), kwargs=None, pg=None):
-        # pyrefly: ignore [bad-index, index-error]
+        # pyrefly: ignore [bad-index]
         st = args[0]
         sharding_spec = st.sharding_spec()
         if len(st.local_shards()) != 1:
diff --git a/torch/distributed/_state_dict_utils.py b/torch/distributed/_state_dict_utils.py
index 989ddd2fb52da..d2ac6be5f45f1 100644
--- a/torch/distributed/_state_dict_utils.py
+++ b/torch/distributed/_state_dict_utils.py
@@ -4,7 +4,7 @@
 import math
 import weakref
 from collections.abc import Callable, Mapping, MutableMapping
-from typing import Any, cast, NamedTuple, TYPE_CHECKING, Union
+from typing import Any, cast, NamedTuple, TYPE_CHECKING
 
 import torch
 import torch.cuda._pin_memory_utils as pin_memory_utils
@@ -731,7 +731,7 @@ def _distribute_state_dict(
 # These APIs are from torch.distributed.checkpoint.
 # TODO: We should consolidate the code here as some not all modules can depend on
 # DCP.
-PATH_ITEM = Union[str, int]
+PATH_ITEM = str | int
 OBJ_PATH = tuple[PATH_ITEM, ...]
 FLATTEN_MAPPING = dict[str, OBJ_PATH]
 STATE_DICT_TYPE = dict[str, Any]
diff --git a/torch/distributed/_symmetric_memory/__init__.py b/torch/distributed/_symmetric_memory/__init__.py
index 2e07c37b48175..1a6c19b85385e 100644
--- a/torch/distributed/_symmetric_memory/__init__.py
+++ b/torch/distributed/_symmetric_memory/__init__.py
@@ -179,16 +179,20 @@ def _pipelined_multi_all_gather_and_consume(
     backend_stream.wait_stream(torch.cuda.current_stream())
 
     for x, y in zip(shard, ag_out):
-        assert x.is_contiguous(), (
-            "_pipelined_all_gather_and_consume: all tensors "
-            "in `shard` must be contiguous"
-        )
-        assert y.is_contiguous(), (
-            "_pipelined_all_gather_and_consume: all tensors "
-            "in `ag_out` must be contiguous"
-        )
-        assert x.shape[0] * group_size == y.shape[0]
-        assert x.shape[1:] == y.shape[1:]
+        if not x.is_contiguous():
+            raise AssertionError(
+                "_pipelined_all_gather_and_consume: all tensors "
+                "in `shard` must be contiguous"
+            )
+        if not y.is_contiguous():
+            raise AssertionError(
+                "_pipelined_all_gather_and_consume: all tensors "
+                "in `ag_out` must be contiguous"
+            )
+        if x.shape[0] * group_size != y.shape[0]:
+            raise AssertionError
+        if x.shape[1:] != y.shape[1:]:
+            raise AssertionError
 
     def copy_shard(dst: list[torch.Tensor], src: list[torch.Tensor]) -> None:
         for d, s in zip(dst, src):
@@ -347,7 +351,8 @@ def _pipelined_produce_and_all2all(
     backend_stream.wait_stream(torch.cuda.current_stream())
 
     def get_p2p_buf(rank: int, idx: int) -> torch.Tensor:
-        assert idx in (0, 1)
+        if idx not in (0, 1):
+            raise AssertionError
         offset = 0 if idx == 0 else out_chunks[0].numel()
         return symm_mem.get_buffer(
             rank, out_chunks[0].shape, out_chunks[0].dtype, offset
@@ -604,7 +609,8 @@ def unflatten(t: torch.Tensor) -> torch.Tensor:
 
     # Computing block-wise matmul along the first dim of A
     if scale_mode == _ScaleMode.ROW_WISE_SHARDED:
-        assert A_scale is not None
+        if A_scale is None:
+            raise AssertionError
         A_scale_shard = A_scale.movedim(gather_dim, 0).flatten(0, -2)
         A_scale_flat = A_scale_shard.new_empty(
             A_scale_shard.shape[0] * group.size(),
@@ -629,7 +635,8 @@ def row_wise_sharded_consumer(shard: list[torch.Tensor], rank: int) -> None:
             return_A,
         )
     elif scale_mode == _ScaleMode.ROW_WISE_REPLICATED:
-        assert A_scale is not None
+        if A_scale is None:
+            raise AssertionError
         A_scale_shards = (
             A_scale.movedim(gather_dim, 0).flatten(0, -2).chunk(group.size())
         )
@@ -653,11 +660,13 @@ def row_wise_replicated_consumer(shard: torch.Tensor, rank: int) -> None:
         )
     else:
         if scale_mode == _ScaleMode.TENSOR_WISE:
-            assert A_scale is not None
+            if A_scale is None:
+                raise AssertionError
             for kwargs in kwargs_list:
                 kwargs["scale_a"] = A_scale
         else:
-            assert scale_mode == _ScaleMode.UNSCALED
+            if scale_mode != _ScaleMode.UNSCALED:
+                raise AssertionError
 
         def default_consumer(shard: torch.Tensor, rank: int) -> None:
             for idx, (B, kwargs) in enumerate(zip(Bs, kwargs_list)):
@@ -1049,7 +1058,8 @@ def _fused_all_gather_scaled_matmul_fallback(
     elif scale_mode == _ScaleMode.ROW_WISE_REPLICATED:
         A_scale = A_scale.movedim(gather_dim, 0).flatten(0, -2)
     else:
-        assert scale_mode == _ScaleMode.TENSOR_WISE
+        if scale_mode != _ScaleMode.TENSOR_WISE:
+            raise AssertionError
 
     def scaled_matmul(
         A: torch.Tensor,
@@ -1161,7 +1171,8 @@ def _fused_all_gather_scaled_matmul(
             group_name,
             True,
         )
-        assert A is not None
+        if A is None:
+            raise AssertionError
         return A, res
 
 
@@ -1736,7 +1747,8 @@ def _low_contention_reduce_scatter_with_symm_mem_input(
     rank = symm_mem.rank
     world_size = symm_mem.world_size
 
-    assert tensor.shape[0] % world_size == 0
+    if tensor.shape[0] % world_size != 0:
+        raise AssertionError
     a2a_res = torch.empty_like(tensor)
     chunks = a2a_res.chunk(world_size)
 
@@ -1774,7 +1786,8 @@ def _low_contention_reduce_scatter_with_workspace(
     rank = workspace.rank
     world_size = workspace.world_size
 
-    assert tensor.shape[0] % world_size == 0
+    if tensor.shape[0] % world_size != 0:
+        raise AssertionError
     chunks = tensor.chunk(world_size)
 
     _get_backend_stream().wait_stream(torch.cuda.current_stream())
@@ -1995,7 +2008,9 @@ def is_nvshmem_available() -> bool:
     r"""
     is_nvshmem_available() -> bool
 
-    Check if NVSHMEM is available in current build and on current system.
+    Check if NVSHMEM (CUDA) or rocSHMEM (ROCm) is available in the current
+    build and usable at runtime. On ROCm, rocSHMEM ``VERSION`` must be at
+    least 3.3.0 (see ``rocshmem/rocshmem.hpp``).
     """
     try:
         from torch._C._distributed_c10d import _is_nvshmem_available
@@ -2182,8 +2197,100 @@ def wait_signal(hdl: _SymmetricMemory, peer: int) -> None:
         raise ValueError(f"wait_signal: unsupported backend: {backend}")
 
 
+def reduce_scatter_offset(
+    input: torch.Tensor,
+    out: list[torch.Tensor],
+    group: str,
+    *,
+    dim: int,
+    offsets: list[int] | None = None,
+    dst_ranks: list[int] | None = None,
+    red_op: str = "sum",
+) -> None:
+    r"""
+    reduce_scatter_offset(input, out, group, *, dim, offsets, dst_ranks, red_op='sum') -> None
+
+    Simultaneously reduce N blocks of a 2-D ``input`` tensor from a symmetric
+    memory buffer, routing each block to a specific destination rank.  Only
+    ``dst_ranks[i]`` writes the reduced result for block ``i``; the result is
+    written to a contiguous output tensor, with the same shape as block ``i``.
+
+    The ``dim`` argument controls which dimension is sharded:
+
+    - ``dim=0`` (row sharding): block ``i`` spans
+      ``input[offsets[i-1] : offsets[i], :]``.  Each ``out[j]`` has shape
+      ``(size_j, input.size(1))``.
+    - ``dim=1`` (column sharding): block ``i`` spans
+      ``input[:, offsets[i-1] : offsets[i]]``.  Each ``out[j]`` has shape
+      ``(input.size(0), size_j)``.
+
+    Blocks are described by ``offsets``, an inclusive prefix-sum of block sizes
+    along ``dim`` (first block starts at index 0 by convention).  Block offsets
+    can be even or uneven; when uneven, the following condition must be met: for
+    each ``j``, the ``j``-th owned block must have the same size across all
+    ranks (so that ``out[j]`` has a uniform shape); different ``j``'s may
+    differ.
+
+    Args:
+        input (Tensor): 2-D tensor allocated via symmetric memory (innermost
+            dimension must be contiguous).
+        out (list[Tensor]): Output tensors for this rank's owned blocks.  Must
+            have length equal to the number of blocks owned by this rank (i.e.
+            the count of ``i`` where ``dst_ranks[i] == my_rank``).  Each
+            ``out[j]`` must be contiguous with the same dtype as ``input``.
+        group (str): The name of the ``ProcessGroup`` to perform the operation on.
+        dim (int): Dimension along which blocks are defined (0 or 1).
+        offsets (list[int] | None): Inclusive prefix-sum of block sizes along
+            ``dim``, length N.  If not provided, ``input.size(dim)`` is divided
+            into equal-size blocks based on the size of the ``group``.
+        dst_ranks (list[int] | None): Destination rank for each block.  If not
+            provided, blocks are distributed round-robin across ranks.
+        red_op (str): Reduction operation; currently only ``'sum'`` is supported.
+
+    Example::
+
+        >>> # doctest: +SKIP
+        >>> # Each rank holds a Grouped GEMM gradient buffer in symmetric memory.
+        >>> # The buffer has W experts laid out as equal column blocks; each expert
+        >>> # is reduced to a specific rank (dst_ranks[i] == i % world_size).
+        >>> buf = symm_mem.empty(H, W * C, dtype=torch.bfloat16, device="cuda")
+        >>> symm_mem.rendezvous(buf, group=group_name)
+        >>> offsets = [i * C for i in range(1, W + 1)]  # inclusive prefix-sum
+        >>> dst_ranks = [i % world_size for i in range(W)]
+        >>> n_owned = sum(r == rank for r in dst_ranks)
+        >>> out = [torch.empty(H, C, dtype=torch.bfloat16, device="cuda") for _ in range(n_owned)]
+        >>> symm_mem.reduce_scatter_offset(buf, out, group_name, dim=1, offsets=offsets, dst_ranks=dst_ranks)
+    """
+    backend = get_backend(input.device)
+    if backend == "NCCL":
+        torch.ops.symm_mem.nccl_reduce_scatter_offset(
+            input, out, group, dim, offsets, dst_ranks, red_op
+        )
+    else:
+        raise NotImplementedError(
+            f"reduce_scatter_offset: unsupported backend: {backend}"
+        )
+
+
+def is_symm_mem_tensor(tensor: torch.Tensor) -> bool:
+    r"""
+    is_symm_mem_tensor(tensor) -> bool
+
+    Returns ``True`` if ``tensor`` was allocated via symmetric memory
+    (i.e. via :func:`torch.distributed._symmetric_memory.empty` or
+    :meth:`_SymmetricMemory.empty_strided_p2p`).
+
+    This is a non-collective, O(1) check.
+
+    Args:
+        tensor (:class:`torch.Tensor`): the tensor to check.
+    """
+    return _SymmetricMemory.is_symm_mem_tensor(tensor)
+
+
 __all__ = [
     "empty",
+    "is_symm_mem_tensor",
     "rendezvous",
     "is_nvshmem_available",
     "set_backend",
@@ -2191,4 +2298,5 @@ def wait_signal(hdl: _SymmetricMemory, peer: int) -> None:
     "set_signal_pad_size",
     "get_signal_pad_size",
     "get_mem_pool",
+    "reduce_scatter_offset",
 ]
diff --git a/torch/distributed/_tools/common_utils.py b/torch/distributed/_tools/common_utils.py
index 0188a4aa08440..2e3869cce0888 100644
--- a/torch/distributed/_tools/common_utils.py
+++ b/torch/distributed/_tools/common_utils.py
@@ -1,6 +1,7 @@
 import warnings
 
 import torch
+from torch._opaque_base import OpaqueBase
 from torch.utils._python_dispatch import is_traceable_wrapper_subclass
 
 
@@ -19,8 +20,17 @@ def get_untyped_storages(t: torch.Tensor) -> set[torch.UntypedStorage]:
     while len(unflattened_tensors) > 0:
         obj = unflattened_tensors.pop()
         if is_traceable_wrapper_subclass(obj):
-            attrs, _ = obj.__tensor_flatten__()  # type: ignore[attr-defined]
-            unflattened_tensors.extend([getattr(obj, attr) for attr in attrs])
+            attrs, _ = obj.__tensor_flatten__()
+            for attr in attrs:
+                match getattr(obj, attr):
+                    case torch.Tensor() as v:
+                        unflattened_tensors.append(v)
+                    case OpaqueBase():
+                        pass
+                    case unexpected:
+                        raise AssertionError(
+                            f"expected Tensor or OpaqueBase, got {type(unexpected)}"
+                        )
         else:
             if not hasattr(obj, "untyped_storage"):
                 warnings.warn(
diff --git a/torch/distributed/_tools/fake_collectives.py b/torch/distributed/_tools/fake_collectives.py
index 4ee1ff93307d7..f8893f3b973e7 100644
--- a/torch/distributed/_tools/fake_collectives.py
+++ b/torch/distributed/_tools/fake_collectives.py
@@ -59,6 +59,9 @@
     _c10d_functional.reduce_scatter_tensor_coalesced.default,
     _c10d_functional_autograd.reduce_scatter_tensor.default,
     _c10d_functional.broadcast_.default,
+    _c10d_functional.isend.default,
+    _c10d_functional.irecv.default,
+    _c10d_functional.batch_p2p_ops.default,
     _dtensor.shard_dim_alltoall.default,
 }
 
@@ -117,9 +120,15 @@ class CollectiveOp:
         _c10d_functional_autograd.reduce_scatter_tensor.default,
         _c10d_functional.all_to_all_single.default,
         _c10d_functional_autograd.all_to_all_single.default,
+        _c10d_functional.isend.default,
+        _c10d_functional.irecv.default,
         _dtensor.shard_dim_alltoall.default,
     }
 
+    PG_ARG_5 = {
+        _c10d_functional.batch_p2p_ops.default,
+    }
+
     WK_ARG_1 = {
         c10d.broadcast_.default,
         c10d.allreduce_.default,
@@ -129,6 +138,8 @@ class CollectiveOp:
         c10d._allgather_base_.default,
         c10d.scatter_.default,
         c10d.alltoall_.default,
+        _c10d_functional.isend.default,
+        _c10d_functional.irecv.default,
     }
 
     WK = {
@@ -212,6 +223,8 @@ def get_process_group(func, args) -> ProcessGroup:  # type: ignore[no-untyped-de
             return _resolve_process_group(args[2])
         if func in CollectiveOp.PG_ARG_4:
             return _resolve_process_group(args[3])
+        if func in CollectiveOp.PG_ARG_5:
+            return _resolve_process_group(args[4])
         raise TypeError(f"Func {func} not found in {collective_ops}")
 
     @staticmethod
@@ -244,6 +257,8 @@ def get_comm_tensor_size(func, res, args, kwargs) -> int:  # type: ignore[no-unt
         if func in CollectiveOp.COMM_TENSOR_ARG_0_AND_RES:
             # TODO(@sanketpurandare) - Confirm size computation
             return args[0].untyped_storage().nbytes() + res.untyped_storage().nbytes()
+        if func is _c10d_functional.batch_p2p_ops.default:
+            return CollectiveOp.sum_tensors(args[3])
         raise TypeError(f"Unknown function: {func} in {collective_ops}")
 
     @staticmethod
diff --git a/torch/distributed/_tools/fsdp2_mem_tracker.py b/torch/distributed/_tools/fsdp2_mem_tracker.py
index 216802b721c9e..1a10223f03ae4 100644
--- a/torch/distributed/_tools/fsdp2_mem_tracker.py
+++ b/torch/distributed/_tools/fsdp2_mem_tracker.py
@@ -165,7 +165,8 @@ def __init__(
         optm: torch.optim.Optimizer | None = None,
     ) -> None:
         super().__init__()
-        assert isinstance(mod, FSDPModule), "FSDPMemTracker only supports FSDP modules"
+        if not isinstance(mod, FSDPModule):
+            raise AssertionError("FSDPMemTracker only supports FSDP modules")
         self._root_mod = mod
         self._optm = optm
         self._fsdp_mod_to_saved_methods: WeakIdKeyDictionary = WeakIdKeyDictionary()
@@ -212,7 +213,8 @@ def inner(
         ) -> tuple[tuple[Unpack[_Ts]], dict[str, Any]]:
             self._fsdp_state = _FSDPState.PRE_FW
             mod_fqn = self._mod_tracker.get_known_fqn(fsdp_mod)
-            assert mod_fqn is not None
+            if mod_fqn is None:
+                raise AssertionError
             if fsdp_mod not in self.memory_tracking:
                 mod_stat = _FSDPModMemStats(mod_fqn)
                 self.memory_tracking[fsdp_mod] = mod_stat
@@ -255,6 +257,7 @@ def inner(
                 state = _FSDPModState.AFT_PRE_FW
             mod_stat.snapshots.setdefault(state, []).append(self.get_tracker_snapshot())
             self._fsdp_state = _FSDPState.FW
+            # pyrefly: ignore [bad-return]
             return args, kwargs
 
         return inner
diff --git a/torch/distributed/_tools/ilp_utils.py b/torch/distributed/_tools/ilp_utils.py
index c6738070d7a2a..15a80ed554d88 100644
--- a/torch/distributed/_tools/ilp_utils.py
+++ b/torch/distributed/_tools/ilp_utils.py
@@ -211,7 +211,8 @@ def parse_module_info(module_info: ModuleInfo) -> Graph:
     mod_stats = module_info["mod_stats"]
     fw_pre_order = module_info["mod_order"]["fw_pre_order"]
     # assertion and number of nodes
-    assert len(mod_stats) == len(fw_pre_order)
+    if len(mod_stats) != len(fw_pre_order):
+        raise AssertionError
     n_nodes = len(mod_stats)
 
     # create graph
diff --git a/torch/distributed/_tools/mem_tracker.py b/torch/distributed/_tools/mem_tracker.py
index f43af62e6af30..52aa888989355 100644
--- a/torch/distributed/_tools/mem_tracker.py
+++ b/torch/distributed/_tools/mem_tracker.py
@@ -142,7 +142,6 @@ def __init__(
         self.size = size
         self.element_size = element_size
         self.reftype = reftype
-        # pyrefly: ignore [read-only]
         self.device = device
         self.mem_consumed = self._calculate_mem_consumed()
 
@@ -419,12 +418,14 @@ def _update_snap(
             dev_snap[_TOTAL_KEY] -= winfo.mem_consumed
             maybe_zero = True
         elif u_type == _UpdateType.REF:
-            assert old_reftype is not None
+            if old_reftype is None:
+                raise AssertionError
             # Adjust memory consumption between two reference types within the same device.
             dev_snap[old_reftype] -= winfo.mem_consumed
             dev_snap[winfo.reftype] += winfo.mem_consumed
         elif u_type == _UpdateType.SIZE:
-            assert old_mem_consumed is not None
+            if old_mem_consumed is None:
+                raise AssertionError
             # Adjust the memory consumed for a reference type due to a change in size.
             change = winfo.mem_consumed - old_mem_consumed
             dev_snap[winfo.reftype] += change
@@ -646,7 +647,8 @@ def _pre_fw_hook(self, module: nn.Module, inputs: Any) -> None:
         #         used multiple times in the same iteration, which we allow and track.
         # For Case 1 and 3, we also initialize the ``local_peak`` and ``PEAK_FW`` snapshot for the module.
         mod_name = self._mod_tracker.get_known_fqn(module)
-        assert mod_name is not None
+        if mod_name is None:
+            raise AssertionError
         if module not in self.memory_tracking:
             mod_stats = _ModMemStats(mod_name)
             param_mem, buffer_mem = self._track_module_params_and_buffers(
@@ -919,7 +921,7 @@ def __torch_dispatch__(self, func, types, args=(), kwargs=None):  # type: ignore
         ):
             # N.B: This is a hacky way to override the Meta IMPL of wait_tensor. The original impl returns
             # a new tensor which does not happen in eager mode, when a wait_tensor is called.
-            # pyrefly: ignore [bad-index, index-error]
+            # pyrefly: ignore [bad-index]
             res = args[0]
         else:
             res = func(*args, **kwargs or {})
diff --git a/torch/distributed/_tools/memory_tracker.py b/torch/distributed/_tools/memory_tracker.py
index a3896249f93d1..b421267b25e1c 100644
--- a/torch/distributed/_tools/memory_tracker.py
+++ b/torch/distributed/_tools/memory_tracker.py
@@ -108,7 +108,8 @@ def start_monitor(self, root_module: nn.Module) -> None:
             # h3 = m.register_backward_hook(self._create_backward_hook(name))
             self._hooks.extend([h1, h2])
         self._device_module.empty_cache()
-        assert getattr(self, "profile_mode", None) is None
+        if getattr(self, "profile_mode", None) is not None:
+            raise AssertionError
         self.profile_mode = MemoryProfileDispatchMode(self)
         self.profile_mode.__enter__()
 
@@ -126,7 +127,8 @@ def stop(self) -> None:
         for h in self._hooks:
             h.remove()
         self._hooks.clear()
-        assert getattr(self, "profile_mode", None) is not None
+        if getattr(self, "profile_mode", None) is None:
+            raise AssertionError
         self.profile_mode.__exit__(None, None, None)
         self.profile_mode = None
 
diff --git a/torch/distributed/_tools/runtime_estimator.py b/torch/distributed/_tools/runtime_estimator.py
index caf399cf6a802..806bdc6911c67 100644
--- a/torch/distributed/_tools/runtime_estimator.py
+++ b/torch/distributed/_tools/runtime_estimator.py
@@ -198,9 +198,10 @@ def _benchmark_estimate(cls, func, args, kwargs) -> tuple[Any, float]:  # type:
             Tuple[Any, float]: A tuple containing the result of the function and
                 the mean operation time in milliseconds.
         """
-        assert isinstance(cls.fake_mode, FakeTensorMode), (
-            "Initialize/Assign FakeTensorMode before using this function"
-        )
+        if not isinstance(cls.fake_mode, FakeTensorMode):
+            raise AssertionError(
+                "Initialize/Assign FakeTensorMode before using this function"
+            )
         mean_op_time = 0.0
         if func._overloadpacket not in _VIEW_OPS:
             try:
@@ -232,9 +233,10 @@ def _roofline_estimate(cls, func, args, kwargs) -> tuple[Any, float]:  # type: i
             Tuple[Any, float]: A tuple containing the result of the function and
                 the mean operation time in milliseconds.
         """
-        assert torch.cuda.is_available(), (
-            "Roofline estimation needs to access CUDA capabilities to make estimations"
-        )
+        if not torch.cuda.is_available():
+            raise AssertionError(
+                "Roofline estimation needs to access CUDA capabilities to make estimations"
+            )
 
         # Roofline Cost Model Explanation
 
@@ -357,9 +359,10 @@ def __call__(self, estimate_mode_type: str) -> Self:
 
     def __enter__(self) -> Self:
         fake_mode = active_fake_mode()
-        assert isinstance(fake_mode, FakeTensorMode), (
-            "No FakeTensorMode found, designed to used under FakeTensorMode"
-        )
+        if not isinstance(fake_mode, FakeTensorMode):
+            raise AssertionError(
+                "No FakeTensorMode found, designed to used under FakeTensorMode"
+            )
         RuntimeEstimator.fake_mode = fake_mode
         self.total_runtime = 0.0
         self.mod_runtimes = defaultdict(lambda: defaultdict(lambda: 0.0))
diff --git a/torch/distributed/_tools/sac_estimator.py b/torch/distributed/_tools/sac_estimator.py
index c43de8c2b9167..3bf4c2394d4db 100644
--- a/torch/distributed/_tools/sac_estimator.py
+++ b/torch/distributed/_tools/sac_estimator.py
@@ -243,7 +243,8 @@ def _pre_fw_hook(self, mod: nn.Module, inputs: Any) -> None:
         # Tracks module FQN, force store random flag, and ``SACModMetadata``
         # Initializes metadata for non-leaf modules, marks leaf modules
         mod_fqn = self._mod_tracker.get_known_fqn(mod)
-        assert mod_fqn is not None
+        if mod_fqn is None:
+            raise AssertionError
         num_children = sum(1 for _ in mod.children())
         if num_children > 0:
             force_store_random = self._get_force_store_random(inputs)
@@ -261,7 +262,8 @@ def _post_fw_hook(self, mod: nn.Module, inputs: Any, outputs: Any) -> None:
         #    - ``SACStats`` using the module's metadata and force store random flag
         #    - ``SACGreedyOrderMeta`` using the computed SAC statistics
         mod_fqn = self._mod_tracker.get_known_fqn(mod)
-        assert mod_fqn is not None
+        if mod_fqn is None:
+            raise AssertionError
         if mod_fqn in self._leaf_modules:
             return
         else:
@@ -374,7 +376,8 @@ def _get_inplace_metadata(
                             if i >= acm_stats.start_idx:
                                 mod_op_parent_idxs[mod_fqn] = i
                         else:
-                            assert mod_fqn == "Global"
+                            if mod_fqn != "Global":
+                                raise AssertionError
                             mod_op_parent_idxs[mod_fqn] = i
         # 7. If no parent tensor is found, then it's probably an inplace op on the arguments
         # so one can just store the current-op idx as parent idx
@@ -405,9 +408,10 @@ def __torch_dispatch__(  # type: ignore[no-untyped-def]
                     out_storages_cpu.update(get_untyped_storages(o))
 
         # Check if there's more than 1 CUDA device
-        assert len(cuda_devices) <= 1, (
-            f"{func.__name__}'s output has more than 1 CUDA devices {cuda_devices}"
-        )
+        if len(cuda_devices) > 1:
+            raise AssertionError(
+                f"{func.__name__}'s output has more than 1 CUDA devices {cuda_devices}"
+            )
 
         # 2. Get the memory consumed by output
         nbytes_cuda = sum(
@@ -448,9 +452,8 @@ def __torch_dispatch__(  # type: ignore[no-untyped-def]
             if acm_stats := self._sac_mod_metadata.get(mod_fqn, None):
                 acm_stats.sac_metadata.append(acm)
             else:
-                assert mod_fqn == "Global", (
-                    f"Module {mod_fqn} not found in AC Mod Stats"
-                )
+                if mod_fqn != "Global":
+                    raise AssertionError(f"Module {mod_fqn} not found in AC Mod Stats")
                 self._sac_metadata.append(acm)
 
         return out
@@ -653,9 +656,11 @@ def save_prediction_graph(
             save_prediction_graph(tradeoff_pwlf, x, y, filename)
         # 9. Obtain the slopes, intercepts and breakpoints of the fitted piecewise linear functions
         slopes = tradeoff_pwlf.calc_slopes().tolist()
-        assert isinstance(tradeoff_pwlf.intercepts, np.ndarray) and isinstance(
-            tradeoff_pwlf.fit_breaks, np.ndarray
-        )
+        if not (
+            isinstance(tradeoff_pwlf.intercepts, np.ndarray)
+            and isinstance(tradeoff_pwlf.fit_breaks, np.ndarray)
+        ):
+            raise AssertionError
         intercepts = tradeoff_pwlf.intercepts.tolist()
         fit_breaks = tradeoff_pwlf.fit_breaks.tolist()
         return SACTradeOffStats(
@@ -943,9 +948,8 @@ def __call__(self, estimate_mode_type: str) -> Self:
 
     def __enter__(self) -> Self:  # type: ignore[no-untyped-def]
         fake_mode = active_fake_mode()
-        assert isinstance(fake_mode, FakeTensorMode), (
-            "SAC Estimator should be called in FakeTensorMode"
-        )
+        if not isinstance(fake_mode, FakeTensorMode):
+            raise AssertionError("SAC Estimator should be called in FakeTensorMode")
         RuntimeEstimator.fake_mode = fake_mode
         self._mod_tracker.register_user_hooks(
             pre_fw_hook=self._pre_fw_hook,
diff --git a/torch/distributed/algorithms/ddp_comm_hooks/ddp_zero_hook.py b/torch/distributed/algorithms/ddp_comm_hooks/ddp_zero_hook.py
index fa8c865c89151..b9bc2a1efb0d9 100644
--- a/torch/distributed/algorithms/ddp_comm_hooks/ddp_zero_hook.py
+++ b/torch/distributed/algorithms/ddp_comm_hooks/ddp_zero_hook.py
@@ -39,9 +39,10 @@ def _perform_local_step(
     """
     overlap_info = zero._overlap_info
     bucket_index = bucket.index()
-    assert len(zero.optim.param_groups) == 1, (
-        "Overlapping DDP with ZeRO only supports a single parameter group"
-    )
+    if len(zero.optim.param_groups) != 1:
+        raise AssertionError(
+            "Overlapping DDP with ZeRO only supports a single parameter group"
+        )
 
     # Construct the `gradients` input for the local optimizer step, which
     # expects `None` in a list position to indicate that the corresponding
@@ -50,9 +51,10 @@ def _perform_local_step(
     gradients: list[torch.Tensor | None] = [
         _NO_PARAM_UPDATE for _ in range(num_local_optim_params)
     ]
-    assert bucket_index in overlap_info.offsets, (
-        f"Bucket index {bucket_index} was not assigned to rank {rank}"
-    )
+    if bucket_index not in overlap_info.offsets:
+        raise AssertionError(
+            f"Bucket index {bucket_index} was not assigned to rank {rank}"
+        )
     gradients_offset = overlap_info.offsets[bucket_index]
     bucket_assignment = zero._bucket_assignments_per_rank[rank][bucket_index]
     bucket_offset = bucket_assignment.offset
@@ -78,19 +80,20 @@ def _broadcast_bucket(
             :class:`ZeroRedundancyOptimizer` instance.
     """
     overlap_info = zero._overlap_info
-    assert len(overlap_info.assigned_ranks_per_bucket) > bucket_index, (
-        "`assigned_ranks_per_bucket` is not fully constructed"
-    )
+    if len(overlap_info.assigned_ranks_per_bucket) <= bucket_index:
+        raise AssertionError("`assigned_ranks_per_bucket` is not fully constructed")
     # Sort to ensure the same ordering across ranks
     assigned_ranks = sorted(overlap_info.assigned_ranks_per_bucket[bucket_index])
-    assert len(assigned_ranks) > 0, (
-        f"Bucket {bucket_index} should be assigned to at least one rank"
-    )
+    if len(assigned_ranks) <= 0:
+        raise AssertionError(
+            f"Bucket {bucket_index} should be assigned to at least one rank"
+        )
     for assigned_rank in assigned_ranks:
         bucket_assignments = zero._bucket_assignments_per_rank[assigned_rank]
         if bucket_index in bucket_assignments:
             send_tensor = bucket_assignments[bucket_index].tensor
-            assert send_tensor is not None
+            if send_tensor is None:
+                raise AssertionError
             overlap_info.broadcast_handles.append(
                 dist.broadcast(
                     send_tensor,
@@ -119,7 +122,8 @@ def _save_ddp_bucket_info(
     """
     overlap_info = zero._overlap_info
     bucket_params = bucket.parameters()
-    assert len(bucket_params) > 0, "Empty bucket"
+    if len(bucket_params) <= 0:
+        raise AssertionError("Empty bucket")
 
     # Save the parameters in the bucket
     overlap_info.params_per_bucket.append(bucket_params)
@@ -128,7 +132,8 @@ def _save_ddp_bucket_info(
         bucket_size = 0
         for param in bucket_params:
             bucket_size += param.numel()
-        assert overlap_info.total_size is not None
+        if overlap_info.total_size is None:
+            raise AssertionError
         overlap_info.total_size += bucket_size
 
 
@@ -154,7 +159,8 @@ def _hook_with_zero_step_setup(
     """
     # Proceed as normal until the DDP buckets have been rebuilt
     if not ddp_ref()._has_rebuilt_buckets:  # type: ignore[union-attr]
-        assert zero._overlap_info.status == _OverlapStatus.UNINITIALIZED
+        if zero._overlap_info.status != _OverlapStatus.UNINITIALIZED:
+            raise AssertionError
         return
 
     bucket_index = bucket.index()
@@ -271,10 +277,10 @@ def hook_with_zero_fn(
         bucket_index = bucket.index()
         rank = zero.global_rank
 
-        assert overlap_info.status == _OverlapStatus.INITIALIZED
-        assert len(overlap_info.assigned_ranks_per_bucket) > bucket_index, (
-            "`assigned_ranks_per_bucket` is not fully constructed"
-        )
+        if overlap_info.status != _OverlapStatus.INITIALIZED:
+            raise AssertionError
+        if len(overlap_info.assigned_ranks_per_bucket) <= bucket_index:
+            raise AssertionError("`assigned_ranks_per_bucket` is not fully constructed")
         assigned_to_bucket = (
             rank in overlap_info.assigned_ranks_per_bucket[bucket_index]
         )
@@ -287,11 +293,11 @@ def hook_with_zero_fn(
         # Check that buckets are indexed incrementally starting from 0 in the
         # order of their autograd hooks firing
         if len(overlap_info.bucket_indices_seen) > 0:
-            assert overlap_info.bucket_indices_seen[-1] == bucket_index - 1, (
-                "Bucket indices are not in incremental order"
-            )
+            if overlap_info.bucket_indices_seen[-1] != bucket_index - 1:
+                raise AssertionError("Bucket indices are not in incremental order")
         else:
-            assert bucket_index == 0, "Bucket indices do not start from 0"
+            if bucket_index != 0:
+                raise AssertionError("Bucket indices do not start from 0")
         overlap_info.bucket_indices_seen.append(bucket_index)
 
         # Directly return the future without any optimizer computation if this
@@ -311,10 +317,11 @@ def hook_with_zero_fn(
             if rank in assigned_ranks:
                 # Wait on the bucket's all-reduce future to ensure correct
                 # gradients
-                assert bucket_index in overlap_info.bucket_index_to_future, (
-                    f"All-reduce future for bucket {bucket_index} not saved "
-                    f"on rank {rank}"
-                )
+                if bucket_index not in overlap_info.bucket_index_to_future:
+                    raise AssertionError(
+                        f"All-reduce future for bucket {bucket_index} not saved "
+                        f"on rank {rank}"
+                    )
                 allreduce_future = overlap_info.bucket_index_to_future[bucket_index]
                 allreduce_future.wait()
 
diff --git a/torch/distributed/algorithms/ddp_comm_hooks/powerSGD_hook.py b/torch/distributed/algorithms/ddp_comm_hooks/powerSGD_hook.py
index f1e95d12514ed..2e3514edef063 100644
--- a/torch/distributed/algorithms/ddp_comm_hooks/powerSGD_hook.py
+++ b/torch/distributed/algorithms/ddp_comm_hooks/powerSGD_hook.py
@@ -22,7 +22,8 @@ def _orthogonalize(matrices, epsilon=0):
 
     QR factorization doesn't work with half-precision, but it is usually faster with a rank > 2.
     """
-    assert len(matrices.shape) == 3 and matrices.shape[2] <= matrices.shape[1]
+    if not (len(matrices.shape) == 3 and matrices.shape[2] <= matrices.shape[1]):
+        raise AssertionError
 
     num_matrices = matrices.shape[0]
     rank = matrices.shape[2]
@@ -631,7 +632,8 @@ def decompress(fut):
 
         if state.use_error_feedback:
             # Memorize the local errors.
-            assert input_tensor_cp is not None
+            if input_tensor_cp is None:
+                raise AssertionError
             state.error_dict[bucket_index] = input_tensor_cp - input_tensor
         if not state.warm_start:
             state.p_memory_dict.clear()
@@ -844,7 +846,8 @@ def decompress(fut):
 
         if state.use_error_feedback:
             # Memorize the local errors.
-            assert input_tensor_cp is not None
+            if input_tensor_cp is None:
+                raise AssertionError
             state.error_dict[bucket_index] = input_tensor_cp - input_tensor
         # Removing this seemingly unnecessary sync somehow may cause failures.
         # See: https://github.com/pytorch/pytorch/pull/54838
diff --git a/torch/distributed/algorithms/join.py b/torch/distributed/algorithms/join.py
index 52d0c52fbfb59..2c56b1fdd9687 100644
--- a/torch/distributed/algorithms/join.py
+++ b/torch/distributed/algorithms/join.py
@@ -183,7 +183,8 @@ def __init__(
 
     def _set_joinable_configs(self) -> None:
         r"""Set the :class:`_JoinConfig` of each participating :class:`Joinable`."""
-        assert len(self._joinables) > 0
+        if len(self._joinables) <= 0:
+            raise AssertionError
         is_first_joinable = True
         for joinable in self._joinables:
             joinable._join_config = _JoinConfig(
@@ -320,10 +321,11 @@ def notify_join_context(joinable: Joinable):
             manager that the process has not yet joined if ``joinable`` is the
             first one passed into the context manager; ``None`` otherwise.
         """
-        assert hasattr(joinable, "_join_config"), (
-            f"Check that the {type(joinable)} constructor calls the "
-            "``Joinable`` constructor"
-        )
+        if not hasattr(joinable, "_join_config"):
+            raise AssertionError(
+                f"Check that the {type(joinable)} constructor calls the "
+                "``Joinable`` constructor"
+            )
 
         join_config = joinable._join_config
         # First joinable is responsible for the collective communications
diff --git a/torch/distributed/benchmarks/benchmark_ddp_rpc.py b/torch/distributed/benchmarks/benchmark_ddp_rpc.py
index e0279966695bc..40394c088dcf5 100644
--- a/torch/distributed/benchmarks/benchmark_ddp_rpc.py
+++ b/torch/distributed/benchmarks/benchmark_ddp_rpc.py
@@ -71,7 +71,8 @@ def forward(self, indices, offsets):
             emb_lookups_cat = torch.cat(emb_lookups, dim=1)
 
         # Make sure combined PS dimension is always bigger or equal than the FC input
-        assert NUM_PS * EMBEDDING_DIM >= 512
+        if NUM_PS * EMBEDDING_DIM < 512:
+            raise AssertionError
         dim_normalizer = int(NUM_PS * EMBEDDING_DIM / 512)
         emb_lookups_reshaped = emb_lookups_cat.reshape(  # type: ignore[possibly-undefined]
             # pyrefly: ignore [unbound-name]
@@ -108,7 +109,8 @@ def _print_cont(msg):
 
 def _run_printable(cmd):
     proc = subprocess.run(shlex.split(cmd), capture_output=True, check=False)  # type: ignore[call-overload]
-    assert proc.returncode == 0
+    if proc.returncode != 0:
+        raise AssertionError
 
     buffer = io.BytesIO()
     torch.save(proc.stdout.decode("utf-8"), buffer)
diff --git a/torch/distributed/checkpoint/_async_process_executor.py b/torch/distributed/checkpoint/_async_process_executor.py
index 569f385a6fc08..110ea622ca4de 100644
--- a/torch/distributed/checkpoint/_async_process_executor.py
+++ b/torch/distributed/checkpoint/_async_process_executor.py
@@ -85,12 +85,10 @@ def get_master_addr_and_port() -> tuple[str, int]:
             if self.use_prefix_store:
                 master_addr = os.environ.get("MASTER_ADDR")
                 master_port = os.environ.get("MASTER_PORT")
-                assert master_addr is not None, (
-                    "DCP needs MASTER_ADDR to use prefix store"
-                )
-                assert master_port is not None, (
-                    "DCP needs MASTER_PORT to use prefix store"
-                )
+                if master_addr is None:
+                    raise AssertionError("DCP needs MASTER_ADDR to use prefix store")
+                if master_port is None:
+                    raise AssertionError("DCP needs MASTER_PORT to use prefix store")
                 master_port = int(master_port)
             else:
                 master_addr = os.environ.get("MASTER_ADDR")
diff --git a/torch/distributed/checkpoint/_consolidate_hf_safetensors.py b/torch/distributed/checkpoint/_consolidate_hf_safetensors.py
index 00d4315501267..cb06faa962add 100644
--- a/torch/distributed/checkpoint/_consolidate_hf_safetensors.py
+++ b/torch/distributed/checkpoint/_consolidate_hf_safetensors.py
@@ -5,7 +5,6 @@
 import json
 import logging
 import math
-import mmap
 import os
 import struct
 import time
@@ -199,17 +198,17 @@ def _write_metadata(
             output_data.metadata_size = f.tell()
 
 
-def _read_tensor_data_mmap(
-    file_path: str,
+def _read_tensor_data(
+    f,
     start_offset: int,
     end_offset: int,
     metadata_size: int,
 ) -> bytes:
     """
-    Read tensor data from a safetensors file using memory mapping for efficiency.
+    Read a specific byte range of tensor data from an open safetensors file.
 
     Args:
-        file_path: Path to the safetensors file
+        f: An open file object (handle) for the safetensors file
         start_offset: Start offset of tensor data within the data section
         end_offset: End offset of tensor data within the data section
         metadata_size: Size of the metadata header
@@ -217,12 +216,11 @@ def _read_tensor_data_mmap(
     Returns:
         Raw tensor data as bytes
     """
-    # Use mmap for efficient access
-    with open(file_path, "rb") as f:
-        with mmap.mmap(f.fileno(), 0, access=mmap.ACCESS_READ) as mm:
-            absolute_start = metadata_size + start_offset
-            absolute_end = metadata_size + end_offset
-            return bytes(mm[absolute_start:absolute_end])
+    absolute_start = metadata_size + start_offset
+    length = end_offset - start_offset
+
+    f.seek(absolute_start)
+    return f.read(length)
 
 
 def _process_output_file(
@@ -231,7 +229,7 @@ def _process_output_file(
     input_files_data: dict[str, _InputFileData],
 ) -> None:
     """
-    Process a single output file by writing tensor data from input files using memory mapping.
+    Process a single output file by writing tensor data from input files using direct reads.
 
     This function is designed to be run in parallel for different output files.
 
@@ -245,55 +243,71 @@ def _process_output_file(
         output_data.fqn_data.items(), key=lambda x: x[1].offset_in_file
     )
 
-    with open(output_file, "r+b") as output_stream:
-        output_stream.seek(0, os.SEEK_END)
-        # Process each tensor in sequential output order
-        for tensor_fqn, tensor_fqn_data in sorted_tensors:
-            full_tensor_mv = memoryview(
-                bytearray(
-                    math.prod(tensor_fqn_data.shape_in_file)
-                    * tensor_fqn_data.dtype_size
+    file_handles = {}
+    dcp_metadata = {}
+    for safetensors_file, file_data in input_files_data.items():
+        dcp_metadata[safetensors_file] = _get_dcp_custom_metadata(file_data.metadata)
+
+    try:
+        # Open all input files for reading
+        for safetensors_file in input_files_data:
+            file_handles[safetensors_file] = open(safetensors_file, "rb")  # noqa: SIM115
+
+        with open(output_file, "r+b") as output_stream:
+            output_stream.seek(0, os.SEEK_END)
+            # Process each tensor in sequential output order
+            for tensor_fqn, tensor_fqn_data in sorted_tensors:
+                full_tensor_mv = memoryview(
+                    bytearray(
+                        math.prod(tensor_fqn_data.shape_in_file)
+                        * tensor_fqn_data.dtype_size
+                    )
                 )
-            )
 
-            # Process each input safetensors file
-            for safetensors_file in input_files_data:
-                file_metadata = input_files_data[safetensors_file].metadata
-                input_metadata_size = input_files_data[safetensors_file].metadata_size
+                # Process each input safetensors file
+                for safetensors_file in input_files_data:
+                    file_metadata = input_files_data[safetensors_file].metadata
+                    input_metadata_size = input_files_data[
+                        safetensors_file
+                    ].metadata_size
 
-                if tensor_fqn not in file_metadata:
-                    continue
+                    if tensor_fqn not in file_metadata:
+                        continue
 
-                metadata = file_metadata[tensor_fqn]
+                    metadata = file_metadata[tensor_fqn]
 
-                data_offsets = metadata[DATA_OFFSETS_KEY]
+                    data_offsets = metadata[DATA_OFFSETS_KEY]
 
-                # Use memory mapping to read tensor data efficiently
-                data_to_write = _read_tensor_data_mmap(
-                    safetensors_file,
-                    data_offsets[0],
-                    data_offsets[1],
-                    input_metadata_size,
-                )
+                    # Use explicit reads to fetch tensor data efficiently
+                    data_to_write = _read_tensor_data(
+                        file_handles[safetensors_file],
+                        data_offsets[0],
+                        data_offsets[1],
+                        input_metadata_size,
+                    )
 
-                # Get the offsets of this tensor shard within the full tensor
-                # pyrefly: ignore [unsupported-operation]
-                fqn_custom_metadata = _get_dcp_custom_metadata(file_metadata)[
-                    tensor_fqn
-                ]  # type: ignore[index]
-                offsets_of_tensor_being_read = fqn_custom_metadata[SAVED_OFFSETS_KEY]  # type: ignore[index]
-
-                # Write this tensor shard to the appropriate position in the output file
-                _write_sub_tensor_to_file_optimized(
-                    full_tensor_mv,
-                    data_to_write,
-                    tensor_fqn_data.dtype_size,  # Size of each element in bytes
-                    tensor_fqn_data.shape_in_file,  # Full tensor shape
-                    offsets_of_tensor_being_read,  # Where this shard belongs in the full tensor
-                    metadata[SHAPE_KEY],  # Shape of this shard
-                )
+                    # Get the offsets of this tensor shard within the full tensor
+                    # pyrefly: ignore [unsupported-operation]
+                    fqn_custom_metadata = dcp_metadata[safetensors_file][tensor_fqn]  # type: ignore[index]
+                    offsets_of_tensor_being_read = fqn_custom_metadata[
+                        SAVED_OFFSETS_KEY
+                    ]  # type: ignore[index]
+
+                    # Write this tensor shard to the appropriate position in the output file
+                    _write_sub_tensor_to_file_optimized(
+                        full_tensor_mv,
+                        data_to_write,
+                        tensor_fqn_data.dtype_size,  # Size of each element in bytes
+                        tensor_fqn_data.shape_in_file,  # Full tensor shape
+                        offsets_of_tensor_being_read,  # Where this shard belongs in the full tensor
+                        metadata[SHAPE_KEY],  # Shape of this shard
+                    )
+
+                output_stream.write(full_tensor_mv)
 
-            output_stream.write(full_tensor_mv)
+    finally:
+        for f in file_handles.values():
+            f.close()
 
 
 def _write_data(
@@ -724,7 +738,8 @@ def consolidate_safetensors_files_on_every_rank(
         if rank == 0:
             # Merge all output_files_data from all ranks
             all_output_files_data: dict[str, _OutputFileData] = {}
-            assert gathered_output_files_data is not None
+            if gathered_output_files_data is None:
+                raise AssertionError
             for rank_data in gathered_output_files_data:
                 all_output_files_data.update(rank_data)
 
diff --git a/torch/distributed/checkpoint/_experimental/staging.py b/torch/distributed/checkpoint/_experimental/staging.py
index 0843a7c58eed3..1e8e9aa35ea2b 100644
--- a/torch/distributed/checkpoint/_experimental/staging.py
+++ b/torch/distributed/checkpoint/_experimental/staging.py
@@ -147,14 +147,12 @@ def __init__(
         self._staging_stream = None
 
         if self._config.use_async_staging:
-            # pyrefly: ignore [bad-assignment]
             self._staging_executor = ThreadPoolExecutor(max_workers=1)
             if torch.accelerator.is_available():
                 # Note: stream needs to be initialized on the main thread after default cuda
                 # stream is setup/used to avoid the risk of accidentally reusing the main
                 # compute stream or in other cases kernels actually launching from the
                 # main thread.
-                # pyrefly: ignore [bad-assignment]
                 self._staging_stream = torch.Stream()
 
         if self._config.use_non_blocking_copy:
diff --git a/torch/distributed/checkpoint/_extension.py b/torch/distributed/checkpoint/_extension.py
index ab08a8a8b4b2b..f7ecb421f0d97 100644
--- a/torch/distributed/checkpoint/_extension.py
+++ b/torch/distributed/checkpoint/_extension.py
@@ -94,7 +94,6 @@ def is_available() -> bool:
         return zstandard is not None or pyzstd is not None
 
     @staticmethod
-    # pyrefly: ignore [bad-override]
     def from_descriptor(version: str) -> "ZStandard":
         if version.partition(".")[0] != "1":
             raise ValueError(f"Unknown extension {version=}")
diff --git a/torch/distributed/checkpoint/_pg_transport.py b/torch/distributed/checkpoint/_pg_transport.py
index b43c7ef86b3e8..46267c40eddc7 100644
--- a/torch/distributed/checkpoint/_pg_transport.py
+++ b/torch/distributed/checkpoint/_pg_transport.py
@@ -225,7 +225,6 @@ def __init__(
         self._work: list[Work] = []
         self._pg = pg
         self._timeout = timeout
-        # pyrefly: ignore [read-only]
         self._device = device
         self._state_dict = state_dict
 
diff --git a/torch/distributed/checkpoint/_state_dict_stager.py b/torch/distributed/checkpoint/_state_dict_stager.py
index 064fa567f7289..1a906a5c009ec 100644
--- a/torch/distributed/checkpoint/_state_dict_stager.py
+++ b/torch/distributed/checkpoint/_state_dict_stager.py
@@ -146,12 +146,13 @@ def _stage_untyped_storage(
         # Check if we've already cached this storage
         if storage in self._cached_storage_mapping:
             cached_storage = self._cached_storage_mapping[storage]
-            assert cached_storage.size() == storage.size(), (
-                "For async checkpointing,  We cache storages in DRAM and reuse them."
-                "Cached storage size does not match original storage size."
-                "This should never happen as we track the original storage weakref "
-                "and clean up the cache storage. Please report this to PyTorch Distributed Checkpointing."
-            )
+            if cached_storage.size() != storage.size():
+                raise AssertionError(
+                    "For async checkpointing,  We cache storages in DRAM and reuse them."
+                    "Cached storage size does not match original storage size."
+                    "This should never happen as we track the original storage weakref "
+                    "and clean up the cache storage. Please report this to PyTorch Distributed Checkpointing."
+                )
             # Reuse cached storage but update with new data
             cached_storage.copy_(storage, non_blocking=non_blocking)
             return cached_storage
diff --git a/torch/distributed/checkpoint/_traverse.py b/torch/distributed/checkpoint/_traverse.py
index 81507fadf26f9..cf0c00a48dda9 100644
--- a/torch/distributed/checkpoint/_traverse.py
+++ b/torch/distributed/checkpoint/_traverse.py
@@ -2,7 +2,8 @@
 # ruff: noqa: F821
 # flake8: noqa: F821
 from collections.abc import Callable, Collection, Mapping, MutableMapping
-from typing import cast, TypeVar, Union
+from typing import cast, TypeVar
+from typing_extensions import TypeIs
 
 import torch
 from torch.distributed._shard.sharded_tensor.api import ShardedTensor
@@ -10,7 +11,7 @@
 from torch.distributed.tensor import DTensor
 
 
-PATH_ITEM = Union[str, int]
+PATH_ITEM = str | int
 OBJ_PATH = tuple[PATH_ITEM, ...]
 T = TypeVar("T")
 
@@ -20,7 +21,7 @@
 __all__ = ["traverse_state_dict", "set_element", "get_element", "print_tensor"]
 
 
-def _keep_visiting_tensors(value: STATE_DICT_ITEM) -> bool:
+def _keep_visiting_tensors(value: STATE_DICT_ITEM) -> TypeIs[torch.Tensor]:
     return isinstance(value, torch.Tensor)
 
 
diff --git a/torch/distributed/checkpoint/api.py b/torch/distributed/checkpoint/api.py
index 4aa4854db2358..4d9817f235ca4 100644
--- a/torch/distributed/checkpoint/api.py
+++ b/torch/distributed/checkpoint/api.py
@@ -8,7 +8,15 @@
 
 
 def _wrap_exception(exc: BaseException) -> WRAPPED_EXCEPTION:
-    return (exc, tb.extract_tb(exc.__traceback__))
+    summary = tb.extract_tb(exc.__traceback__)
+    # Python 3.13+ stores bytecode objects in FrameSummary._code,
+    # which cannot be pickled. Clear them so gather_object succeeds
+    # and the real exception is reported instead of a misleading
+    # "cannot pickle code objects" TypeError.
+    for frame in summary:
+        if hasattr(frame, "_code"):
+            object.__setattr__(frame, "_code", None)
+    return (exc, summary)
 
 
 def _is_wrapped_exception(obj: Any) -> bool:
diff --git a/torch/distributed/checkpoint/default_planner.py b/torch/distributed/checkpoint/default_planner.py
index aaa6126d5d041..760981a7bc12f 100644
--- a/torch/distributed/checkpoint/default_planner.py
+++ b/torch/distributed/checkpoint/default_planner.py
@@ -123,7 +123,12 @@ def create_local_plan(self) -> SavePlan:
                 )
                 return SavePlan([], usable=False)
             else:
-                SavePlanner._cached_save_plan[self._cached_plans_key] = plan
+                # Store the plan as pending. It will be promoted to the
+                # class-level cache in finish_plan after the global plan
+                # has succeeded. This avoids a stale local cache when
+                # the global plan fails (e.g. validation error) but the
+                # local cache was already populated.
+                self._pending_local_plan = plan
 
         return self.plan
 
@@ -146,8 +151,12 @@ def _create_global_plan(
             merged_mappings = dict(ChainMap(*planner_data_dict))
             metadata = dataclasses.replace(metadata, planner_data=merged_mappings)
 
-        if not _validate_global_plan(global_plan, metadata):
-            raise ValueError("Failed to validate global plan")
+        validation_errors = _validate_global_plan(global_plan, metadata)
+        if validation_errors:
+            error_summary = "; ".join(validation_errors)
+            if len(error_summary) > 500:
+                error_summary = error_summary[:500] + "... (truncated)"
+            raise ValueError(f"Failed to validate global plan: {error_summary}")
 
         return global_plan, metadata
 
@@ -164,12 +173,13 @@ def _create_global_plan_with_caching(
             # Case 1: If the plans are not cached, the cache will be hydrated with the
             # all_plans, global_plans (Deduped), and metadata.
 
-            # Cache the original all_plans
-            SavePlanner._cached_all_plans[self._cached_plans_key] = all_plans
+            # First create and validate the global plan. Only cache everything
+            # after success to avoid partial cache state
             global_plan, metadata = self._create_global_plan(all_plans)
-            # Cache the deduped and validated global_plan
+
+            # Cache all plans atomically after successful validation
+            SavePlanner._cached_all_plans[self._cached_plans_key] = all_plans
             SavePlanner._cached_global_plan[self._cached_plans_key] = global_plan
-            # Cache the metadata
             SavePlanner._cached_metadata[self._cached_plans_key] = metadata
             # If plans are not cached, global_plan delta will be the same as global plan.
             return global_plan, global_plan, metadata
@@ -248,6 +258,16 @@ def finish_plan(self, new_plan: SavePlan) -> SavePlan:
         if self._enable_plan_caching:
             finished_plan = self._finish_plan_with_caching(new_plan)
 
+            # Promote the pending local plan to the class-level cache now
+            # that the global plan has succeeded and we are finalizing.
+            # This ensures the local cache is only populated after a
+            # successful end-to-end checkpoint plan creation.
+            if hasattr(self, "_pending_local_plan"):
+                SavePlanner._cached_save_plan[self._cached_plans_key] = (
+                    self._pending_local_plan
+                )
+                del self._pending_local_plan
+
         self.plan = finished_plan
         return self.plan
 
@@ -628,8 +648,9 @@ def _check_box_bounds(
     return True
 
 
-def _validate_global_plan(global_plan: list[SavePlan], metadata: Metadata) -> bool:
-    all_good = True
+def _validate_global_plan(global_plan: list[SavePlan], metadata: Metadata) -> list[str]:
+    """Validate the global plan and return a list of error messages (empty if valid)."""
+    errors: list[str] = []
     for key, value in metadata.state_dict_metadata.items():
         if isinstance(value, BytesStorageMetadata):
             continue
@@ -640,16 +661,12 @@ def _validate_global_plan(global_plan: list[SavePlan], metadata: Metadata) -> bo
         for chunk in chunks:
             # Compute the volume
             if not _check_box_bounds(value.size, chunk):
-                logger.warning(
-                    """
-                        key:%s has out of bounds chunk:
-                        tensor-size:%s chunk: %s
-                    """,
-                    key,
-                    value.size,
-                    chunk,
+                msg = (
+                    f"key:{key} has out of bounds chunk: "
+                    f"tensor-size:{value.size} chunk: {chunk}"
                 )
-                all_good = False
+                logger.warning(msg)
+                errors.append(msg)
             chunks_volume += math.prod(chunk.sizes)
 
         if len(chunks) > 1:
@@ -675,28 +692,20 @@ def _validate_global_plan(global_plan: list[SavePlan], metadata: Metadata) -> bo
                 for _, other_idx in active:
                     other = chunks[other_idx]
                     if _check_box_overlap(current, other):
-                        logger.warning(
-                            "key:%s has overlapping chunks: %s %s",
-                            key,
-                            current,
-                            other,
-                        )
-                        all_good = False
+                        msg = f"key:{key} has overlapping chunks: {current} {other}"
+                        logger.warning(msg)
+                        errors.append(msg)
 
                 insort(active, (end, idx))
 
         # Check whether combined chunk cover the whole tensor
         tensor_volume = math.prod(value.size)
         if len(global_plan) > 1 and chunks_volume != tensor_volume:
-            logger.warning(
-                """
-                    key:%s invalid fill tensor-volume:
-                    %s chunks-volume: %s
-                """,
-                key,
-                tensor_volume,
-                chunks_volume,
+            msg = (
+                f"key:{key} invalid fill tensor-volume: "
+                f"{tensor_volume} chunks-volume: {chunks_volume}"
             )
-            all_good = False
+            logger.warning(msg)
+            errors.append(msg)
 
-    return all_good
+    return errors
diff --git a/torch/distributed/checkpoint/filesystem.py b/torch/distributed/checkpoint/filesystem.py
index d693790356a9d..831fa16394563 100644
--- a/torch/distributed/checkpoint/filesystem.py
+++ b/torch/distributed/checkpoint/filesystem.py
@@ -756,7 +756,7 @@ def _write_data(
             return fut
 
     def finish(self, metadata: Metadata, results: list[list[WriteResult]]) -> None:
-        metadata = dataclasses.replace(metadata, version=CURRENT_DCP_VERSION)
+        metadata.version = CURRENT_DCP_VERSION
 
         storage_md = {}
         for wr_list in results:
diff --git a/torch/distributed/checkpoint/metadata.py b/torch/distributed/checkpoint/metadata.py
index ef7a3e95b19ed..7855b94df33ce 100644
--- a/torch/distributed/checkpoint/metadata.py
+++ b/torch/distributed/checkpoint/metadata.py
@@ -3,7 +3,7 @@
 from collections.abc import Sequence
 from dataclasses import dataclass, field
 from enum import Enum
-from typing import Any, Union
+from typing import Any
 
 import torch
 from torch.distributed.checkpoint.stateful import StatefulT
@@ -122,7 +122,7 @@ class BytesStorageMetadata:
     pass
 
 
-STORAGE_TYPES = Union[TensorStorageMetadata, BytesStorageMetadata]
+STORAGE_TYPES = TensorStorageMetadata | BytesStorageMetadata
 STATE_DICT_TYPE = dict[str, StatefulT | Any]
 
 
diff --git a/torch/distributed/checkpoint/quantized_hf_storage.py b/torch/distributed/checkpoint/quantized_hf_storage.py
index 464052d99062a..af9e5c2e7da0a 100644
--- a/torch/distributed/checkpoint/quantized_hf_storage.py
+++ b/torch/distributed/checkpoint/quantized_hf_storage.py
@@ -230,9 +230,8 @@ def _dequantize_tensor_mxfp4(
 
         scales = scales.to(torch.int32) - 127
 
-        assert blocks.shape[:-1] == scales.shape, (
-            f"{blocks.shape=} does not match {scales.shape=}"
-        )
+        if blocks.shape[:-1] != scales.shape:
+            raise AssertionError(f"{blocks.shape=} does not match {scales.shape=}")
 
         lut = torch.tensor(FP4_VALUES, dtype=self.target_dtype, device=blocks.device)
 
@@ -413,7 +412,8 @@ def _read_quantized_tensor_with_block_alignment(
 
                 # Calculate which groups we need based on the requested range in dim 2.
                 # Ensure the reequest is in 3D.
-                assert len(req.storage_offsets) == 3
+                if len(req.storage_offsets) != 3:
+                    raise AssertionError
 
                 # Positions in dequantized space.
                 dim2_start_deq = req.storage_offsets[2]
diff --git a/torch/distributed/checkpoint/staging.py b/torch/distributed/checkpoint/staging.py
index 78231aeebcf5b..04c16cdec1705 100644
--- a/torch/distributed/checkpoint/staging.py
+++ b/torch/distributed/checkpoint/staging.py
@@ -182,14 +182,12 @@ def __init__(
         self._staging_executor = None
         self._staging_stream = None
         if self._config.use_async_staging:
-            # pyrefly: ignore [bad-assignment]
             self._staging_executor = ThreadPoolExecutor(max_workers=1)
             if torch.accelerator.is_available():
                 # Note: stream needs to be initialized on the main thread after default cuda
                 # stream is setup/used to avoid the risk of accidentally reusing the main
                 # compute stream or in other cases kernels actually launching from the
                 # main thread.
-                # pyrefly: ignore [bad-assignment]
                 self._staging_stream = torch.Stream()
 
         if self._config.use_non_blocking_copy:
@@ -361,7 +359,6 @@ def __init__(
     ):
         self._pg = pg
         self._timeout = timeout
-        # pyrefly: ignore [read-only]
         self._device = device
         self._transport = PGTransport(pg, timeout, device, None)
 
diff --git a/torch/distributed/checkpoint/state_dict.py b/torch/distributed/checkpoint/state_dict.py
index e1ee6cdb2e904..7d4428afa25db 100644
--- a/torch/distributed/checkpoint/state_dict.py
+++ b/torch/distributed/checkpoint/state_dict.py
@@ -6,7 +6,7 @@
 from collections.abc import Callable, Generator, Iterable
 from dataclasses import asdict, dataclass, field
 from itertools import chain
-from typing import Any, cast, no_type_check, Union
+from typing import Any, cast, no_type_check
 
 import torch
 import torch.distributed as dist
@@ -66,10 +66,10 @@
 _STATE = "state"
 
 FQNS_T = set[str]
-PrimitiveType = Union[DTensor, ShardedTensor, torch.Tensor, int, float, str]
-ValueType = Union[
-    PrimitiveType, list[PrimitiveType], tuple[PrimitiveType], dict[str, "ValueType"]
-]
+PrimitiveType = DTensor | ShardedTensor | torch.Tensor | int | float | str
+ValueType = (
+    PrimitiveType | list[PrimitiveType] | tuple[PrimitiveType] | dict[str, "ValueType"]
+)
 DictValueType = dict[str, ValueType]
 ListDictValueType = list[DictValueType]
 OptimizerStateType = dict[str, DictValueType | ListDictValueType]
@@ -200,7 +200,6 @@ def _get_fqns(
                 return {f"{prefix}{fqn}" for fqn in flat_param._fqns}
             curr_obj = getattr(curr_obj, FSDP_WRAPPED_MODULE)
             if curr_obj_name != FSDP_WRAPPED_MODULE:
-                # pyrefly: ignore [bad-argument-type]
                 fqn_obj_names.append(curr_obj_name)
                 curr_obj = getattr(curr_obj, curr_obj_name)
         elif isinstance(curr_obj, torch._dynamo.eval_frame.OptimizedModule):
@@ -218,7 +217,6 @@ def _get_fqns(
                 ):
                     if hasattr(curr_obj, removed_fqn):
                         curr_obj = getattr(curr_obj, removed_fqn)
-            # pyrefly: ignore [bad-argument-type]
             fqn_obj_names.append(curr_obj_name)
             if curr_obj_name == nn.modules.module._EXTRA_STATE_KEY_SUFFIX:
                 if i != len(obj_names) - 1:
@@ -799,7 +797,8 @@ def _reconstruct_nested_dict(
                 if part not in current:
                     current[part] = {}
                 # Move deeper into the nested structure
-                assert isinstance(current[part], dict)
+                if not isinstance(current[part], dict):
+                    raise AssertionError
                 current = current[part]
 
             # Set the value at the final level using the last part as the key
@@ -845,6 +844,7 @@ def _reconstruct_nested_dict(
                     continue
 
                 # Reconstruct state for this parameter
+                # pyrefly: ignore [unsupported-operation]
                 state[fqn] = {}
                 for state_name in optim.state[param]:
                     flattened_state_key = f"{_STATE}.{fqn}.{state_name}"
@@ -854,11 +854,13 @@ def _reconstruct_nested_dict(
                         reconstructed_value = _reconstruct_nested_dict(
                             flattened_state_key, state_dict
                         )
+                        # pyrefly: ignore [bad-index]
                         cast(DictValueType, state[fqn])[state_name] = (
                             reconstructed_value
                         )
                     else:
                         # Existing keys mean no nesting, directly use the value.
+                        # pyrefly: ignore [bad-index]
                         cast(DictValueType, state[fqn])[state_name] = state_dict[
                             flattened_state_key
                         ]
diff --git a/torch/distributed/checkpoint/state_dict_loader.py b/torch/distributed/checkpoint/state_dict_loader.py
index 62acd1bfef359..955ccbcdcf3bc 100644
--- a/torch/distributed/checkpoint/state_dict_loader.py
+++ b/torch/distributed/checkpoint/state_dict_loader.py
@@ -228,11 +228,15 @@ def local_step():
         nonlocal metadata
 
         # Use global metadata if available, otherwise fallback to rank local metadata
+        global_metadata_exc: Exception | None = None
+        rank_metadata_exc: Exception | None = None
         try:
             metadata = storage_reader.read_metadata()
-        except Exception:
-            logger.info(
-                "Global metadata is not found. Falling back to rank local metadata."
+        except Exception as e:
+            global_metadata_exc = e
+            logger.warning(
+                "Global metadata is not found. Falling back to rank local metadata.",
+                exc_info=True,
             )
 
         if (
@@ -242,13 +246,23 @@ def local_step():
             try:
                 metadata = storage_reader.read_metadata(rank=distW.rank)  # noqa: F841
                 use_collectives = False
-            except Exception:
-                logger.info("Rank local metadata is not found.")
+            except Exception as e:
+                rank_metadata_exc = e
+                logger.warning("Rank local metadata is not found.", exc_info=True)
 
         if planner is None:
             raise AssertionError("planner is None")
         if metadata is None:
-            raise AssertionError("metadata is None")
+            error_parts = ["metadata is None"]
+            if global_metadata_exc is not None:
+                error_parts.append(
+                    f"global metadata read failed: {global_metadata_exc}"
+                )
+            if rank_metadata_exc is not None:
+                error_parts.append(
+                    f"rank local metadata read failed: {rank_metadata_exc}"
+                )
+            raise AssertionError("; ".join(error_parts))
         planner.set_up_planner(state_dict, metadata, distW.is_coordinator)
 
         if (
diff --git a/torch/distributed/checkpoint/utils.py b/torch/distributed/checkpoint/utils.py
index d5e19eee111a0..4f14bce83b09e 100644
--- a/torch/distributed/checkpoint/utils.py
+++ b/torch/distributed/checkpoint/utils.py
@@ -261,6 +261,7 @@ def all_reduce(
         final_result = self.broadcast_object(result)
         if isinstance(final_result, CheckpointException):
             raise final_result
+        # pyrefly: ignore [redundant-cast]
         return cast(R, final_result)
 
     def all_gather(
@@ -310,6 +311,7 @@ def broadcast(
         final_result = self.broadcast_object(result)
         if isinstance(final_result, CheckpointException):
             raise final_result
+        # pyrefly: ignore [redundant-cast]
         return cast(T, final_result)
 
     def barrier(self) -> None:
diff --git a/torch/distributed/collective_utils.py b/torch/distributed/collective_utils.py
index 54cf4d2354288..c207fd9163af5 100644
--- a/torch/distributed/collective_utils.py
+++ b/torch/distributed/collective_utils.py
@@ -189,7 +189,7 @@ def all_gather(
             raise RuntimeError(  # type: ignore[misc]
                 error_msg,
                 exception_list,
-            ) from exception_list[0]  # pyrefly: ignore [bad-raise, invalid-inheritance]
+            ) from exception_list[0]  # pyrefly: ignore [bad-raise]
         return ret_list
     else:
         if not sync_obj.success:
@@ -274,13 +274,10 @@ def _summarize_ranks(ranks: Iterable[int]) -> str:
     result = []
     for r in ranges:
         if len(r) == 1:
-            # pyrefly: ignore [bad-argument-type]
             result.append(f"{r.start}")
         elif r.step == 1:
-            # pyrefly: ignore [bad-argument-type]
             result.append(f"{r.start}:{r.stop}")
         else:
-            # pyrefly: ignore [bad-argument-type]
             result.append(f"{r.start}:{r.stop}:{r.step}")
     return ",".join(result)
 
diff --git a/torch/distributed/debug/__init__.py b/torch/distributed/debug/__init__.py
index 3145a3ab13ba1..a5a66db1ef90a 100644
--- a/torch/distributed/debug/__init__.py
+++ b/torch/distributed/debug/__init__.py
@@ -32,6 +32,7 @@ def start_debug_server(
     dump_interval: float = 60.0,
     enabled_dumps: set[str] | None = None,
     handlers: list["DebugHandler"] | None = None,
+    fetch_timeout: float = 60.0,
 ) -> None:
     """
     Start the debug server stack on all workers. The frontend debug server is
@@ -71,11 +72,16 @@ def start_debug_server(
         handlers (list[DebugHandler] | None): List of debug handlers to use. If None,
             uses the default handlers. See torch.distributed.debug._handlers for
             the default handlers.
+        fetch_timeout (float): Timeout in seconds for fetching data from individual
+            workers. Defaults to 60. Workers that don't respond within this time
+            will be reported as unavailable.
     """
     global _WORKER_SERVER, _DEBUG_SERVER_PROC
 
-    assert _WORKER_SERVER is None, "debug server already started"
-    assert _DEBUG_SERVER_PROC is None, "debug server already started"
+    if _WORKER_SERVER is not None:
+        raise AssertionError("debug server already started")
+    if _DEBUG_SERVER_PROC is not None:
+        raise AssertionError("debug server already started")
 
     logger.info("Starting debug server on port %d", port)
 
@@ -104,6 +110,7 @@ def start_debug_server(
             "dump_interval": dump_interval,
             "enabled_dumps": enabled_dumps,
             "handlers": handlers,
+            "fetch_timeout": fetch_timeout,
         }
 
         if start_method is not None:
@@ -125,8 +132,10 @@ def stop_debug_server() -> None:
     """
     global _WORKER_SERVER, _DEBUG_SERVER_PROC
 
-    assert _DEBUG_SERVER_PROC is not None
-    assert _WORKER_SERVER is not None
+    if _DEBUG_SERVER_PROC is None:
+        raise AssertionError
+    if _WORKER_SERVER is None:
+        raise AssertionError
 
     logger.info("Stopping debug server")
 
diff --git a/torch/distributed/debug/_debug_handlers.py b/torch/distributed/debug/_debug_handlers.py
index 1c11060e701b2..902f78e2e863f 100644
--- a/torch/distributed/debug/_debug_handlers.py
+++ b/torch/distributed/debug/_debug_handlers.py
@@ -7,6 +7,7 @@
 from torch.distributed.debug._frontend import (
     DebugHandler,
     fetch_all,
+    format_fetch_summary,
     format_json,
     NavLink,
     Response,
@@ -74,6 +75,7 @@
     <h1>{% block title %}{{ title }}{% endblock %}</h1>
 {% endblock %}
 {% block content %}
+    {% if fetch_summary %}<pre>{{ fetch_summary }}</pre>{% endif %}
     <h2>Groups</h2>
     {{ groups | safe }}
     <h2>Memberships</h2>
@@ -192,14 +194,18 @@ def nav_links(self) -> list[NavLink]:
         return [NavLink("/stacks", "Python Stack Traces")]
 
     def _handle(self, req: HTTPRequestHandler) -> bytes:
-        addrs, resps = fetch_all("dump_traceback")
+        addrs, resps = fetch_all("dump_traceback", timeout=self.fetch_timeout)
         return req.frontend.render_template(
             "raw_resp.html", title="Stacks", addrs=addrs, resps=resps
         )
 
     def dump(self) -> str | None:
-        addrs, resps = fetch_all("dump_traceback")
+        addrs, resps = fetch_all("dump_traceback", timeout=self.fetch_timeout)
         parts: list[str] = []
+        summary = format_fetch_summary(addrs, resps)
+        if summary:
+            parts.append(summary)
+            parts.append("")
         for i, (addr, resp) in enumerate(zip(addrs, resps)):
             parts.append(f"=== Rank {i}: {addr} ===")
             parts.append(
@@ -222,7 +228,10 @@ def templates(self) -> dict[str, str]:
         return {"pyspy_dump.html": PYSPY_DUMP_TEMPLATE}
 
     def _handle(self, req: HTTPRequestHandler) -> bytes:
-        addrs, resps = fetch_all("pyspy_dump", req.get_raw_query())
+        query = req.get_raw_query()
+        if "nonblocking" not in query:
+            query = f"nonblocking=1&{query}" if query else "nonblocking=1"
+        addrs, resps = fetch_all("pyspy_dump", query, timeout=self.fetch_timeout)
         return req.frontend.render_template(
             "pyspy_dump.html",
             addrs=addrs,
@@ -230,8 +239,14 @@ def _handle(self, req: HTTPRequestHandler) -> bytes:
         )
 
     def dump(self) -> str | None:
-        addrs, resps = fetch_all("pyspy_dump")
+        addrs, resps = fetch_all(
+            "pyspy_dump", "nonblocking=1", timeout=self.fetch_timeout
+        )
         parts: list[str] = []
+        summary = format_fetch_summary(addrs, resps)
+        if summary:
+            parts.append(summary)
+            parts.append("")
         for i, (addr, resp) in enumerate(zip(addrs, resps)):
             parts.append(f"=== Rank {i}: {addr} ===")
             parts.append(
@@ -272,7 +287,8 @@ def _build_db(addrs: list[str], resps: list[Response]) -> Database:
 
         details = {}
         for rank, resp in enumerate(resps):
-            resp.raise_for_status()
+            if resp.status_code != 200:
+                continue
             dump = {
                 "rank": rank,
                 "host_name": addrs[rank],
@@ -282,6 +298,9 @@ def _build_db(addrs: list[str], resps: list[Response]) -> Database:
                 dump["entries"] = []
             details[f"rank{rank}.json"] = dump
 
+        if not details:
+            raise RuntimeError("All workers failed to respond")
+
         version = next(iter(details.values()))["version"]
         # pyrefly: ignore [bad-argument-type]
         return build_db(details, args, version)
@@ -293,6 +312,7 @@ def _render_tables(
         return server.render_template(
             "fr_trace.html",
             title="FlightRecorder",
+            fetch_summary=format_fetch_summary(addrs, resps),
             groups=tabulate(db.groups, headers=Group._fields, tablefmt="html"),
             memberships=tabulate(
                 db.memberships, headers=Membership._fields, tablefmt="html"
@@ -304,11 +324,11 @@ def _render_tables(
         )
 
     def _handle_fr_trace(self, req: HTTPRequestHandler) -> bytes:
-        addrs, resps = fetch_all("fr_trace_json")
+        addrs, resps = fetch_all("fr_trace_json", timeout=self.fetch_timeout)
         return self._render_tables(req.frontend, addrs, list(resps))
 
     def _handle_fr_trace_json(self, req: HTTPRequestHandler) -> bytes:
-        addrs, resps = fetch_all("fr_trace_json")
+        addrs, resps = fetch_all("fr_trace_json", timeout=self.fetch_timeout)
         return req.frontend.render_template(
             "json_resp.html",
             title="FlightRecorder",
@@ -317,11 +337,15 @@ def _handle_fr_trace_json(self, req: HTTPRequestHandler) -> bytes:
         )
 
     def _handle_fr_trace_nccl(self, req: HTTPRequestHandler) -> bytes:
-        addrs, resps = fetch_all("dump_nccl_trace_json", "onlyactive=true")
+        addrs, resps = fetch_all(
+            "dump_nccl_trace_json", "onlyactive=true", timeout=self.fetch_timeout
+        )
         return self._render_tables(req.frontend, addrs, list(resps))
 
     def _handle_fr_trace_nccl_json(self, req: HTTPRequestHandler) -> bytes:
-        addrs, resps = fetch_all("dump_nccl_trace_json", "onlyactive=true")
+        addrs, resps = fetch_all(
+            "dump_nccl_trace_json", "onlyactive=true", timeout=self.fetch_timeout
+        )
         return req.frontend.render_template(
             "json_resp.html",
             title="FlightRecorder NCCL",
@@ -332,8 +356,12 @@ def _handle_fr_trace_nccl_json(self, req: HTTPRequestHandler) -> bytes:
     def dump(self) -> str | None:
         parts = []
 
-        addrs, resps = fetch_all("fr_trace_json")
-        db = self._build_db(addrs, list(resps))
+        addrs, resps = fetch_all("fr_trace_json", timeout=self.fetch_timeout)
+        summary = format_fetch_summary(addrs, resps)
+        if summary:
+            parts.append(summary)
+            parts.append("")
+        db = self._build_db(addrs, resps)
         parts.extend(
             [
                 "=== FR Trace ===",
@@ -350,9 +378,9 @@ def dump(self) -> str | None:
 
         try:
             nccl_addrs, nccl_resps = fetch_all(
-                "dump_nccl_trace_json", "onlyactive=true"
+                "dump_nccl_trace_json", "onlyactive=true", timeout=self.fetch_timeout
             )
-            nccl_db = self._build_db(nccl_addrs, list(nccl_resps))
+            nccl_db = self._build_db(nccl_addrs, nccl_resps)
             parts.extend(
                 [
                     "",
@@ -400,7 +428,9 @@ def templates(self) -> dict[str, str]:
 
     def _handle(self, req: HTTPRequestHandler) -> bytes:
         duration = req.get_query_arg("duration", default=1.0, type=float)
-        addrs, resps = fetch_all("torch_profile", f"duration={duration}")
+        addrs, resps = fetch_all(
+            "torch_profile", f"duration={duration}", timeout=self.fetch_timeout
+        )
         return req.frontend.render_template("profile.html", addrs=addrs, resps=resps)
 
 
@@ -412,14 +442,18 @@ def nav_links(self) -> list[NavLink]:
         return [NavLink("/wait_counters", "Wait Counters")]
 
     def _handle(self, req: HTTPRequestHandler) -> bytes:
-        addrs, resps = fetch_all("wait_counter_values")
+        addrs, resps = fetch_all("wait_counter_values", timeout=self.fetch_timeout)
         return req.frontend.render_template(
             "json_resp.html", title="Wait Counters", addrs=addrs, resps=resps
         )
 
     def dump(self) -> str | None:
-        addrs, resps = fetch_all("wait_counter_values")
+        addrs, resps = fetch_all("wait_counter_values", timeout=self.fetch_timeout)
         parts: list[str] = []
+        summary = format_fetch_summary(addrs, resps)
+        if summary:
+            parts.append(summary)
+            parts.append("")
         for i, (addr, resp) in enumerate(zip(addrs, resps)):
             parts.append(f"=== Rank {i}: {addr} ===")
             if resp.status_code == 200:
@@ -461,12 +495,93 @@ def dump_filename(self) -> str:
         return "tcpstore"
 
 
+class TorchCommsFlightRecorderHandler(DebugHandler):
+    """Handler for TorchComms FlightRecorder trace data."""
+
+    def routes(self) -> list[Route]:
+        return [
+            Route("/torchcomms_fr_trace", self._handle_torchcomms_fr_trace),
+            Route("/torchcomms_fr_trace_json", self._handle_torchcomms_fr_trace_json),
+        ]
+
+    def nav_links(self) -> list[NavLink]:
+        return [
+            NavLink("/torchcomms_fr_trace", "TorchComms FR"),
+            NavLink("/torchcomms_fr_trace_json", "(JSON)"),
+        ]
+
+    def templates(self) -> dict[str, str]:
+        return {"fr_trace.html": FR_TRACE_TEMPLATE}
+
+    def _handle_torchcomms_fr_trace(self, req: HTTPRequestHandler) -> bytes:
+        addrs, resps = fetch_all(
+            "torchcomms_fr_trace_json", "onlyactive=true", timeout=self.fetch_timeout
+        )
+        return self._render_tables(req.frontend, addrs, list(resps))
+
+    def _handle_torchcomms_fr_trace_json(self, req: HTTPRequestHandler) -> bytes:
+        addrs, resps = fetch_all(
+            "torchcomms_fr_trace_json", "onlyactive=true", timeout=self.fetch_timeout
+        )
+        return req.frontend.render_template(
+            "json_resp.html",
+            title="TorchComms FlightRecorder",
+            addrs=addrs,
+            resps=resps,
+        )
+
+    def _render_tables(
+        self, server: FrontendServer, addrs: list[str], resps: list[Response]
+    ) -> bytes:
+        db = FlightRecorderHandler._build_db(addrs, resps)
+        return server.render_template(
+            "fr_trace.html",
+            title="TorchComms FlightRecorder",
+            fetch_summary=format_fetch_summary(addrs, resps),
+            groups=tabulate(db.groups, headers=Group._fields, tablefmt="html"),
+            memberships=tabulate(
+                db.memberships, headers=Membership._fields, tablefmt="html"
+            ),
+            collectives=tabulate(
+                db.collectives, headers=Collective._fields, tablefmt="html"
+            ),
+            ncclcalls=tabulate(db.ncclcalls, headers=NCCLCall._fields, tablefmt="html"),
+        )
+
+    def dump(self) -> str | None:
+        addrs, resps = fetch_all("torchcomms_fr_trace_json", timeout=self.fetch_timeout)
+        parts: list[str] = []
+        summary = format_fetch_summary(addrs, resps)
+        if summary:
+            parts.append(summary)
+            parts.append("")
+        db = FlightRecorderHandler._build_db(addrs, resps)
+        parts.extend(
+            [
+                "=== TorchComms FR Trace ===",
+                "--- Groups ---",
+                tabulate(db.groups, headers=Group._fields, tablefmt="plain"),
+                "--- Memberships ---",
+                tabulate(db.memberships, headers=Membership._fields, tablefmt="plain"),
+                "--- Collectives ---",
+                tabulate(db.collectives, headers=Collective._fields, tablefmt="plain"),
+                "--- NCCL Calls ---",
+                tabulate(db.ncclcalls, headers=NCCLCall._fields, tablefmt="plain"),
+            ]
+        )
+        return "\n".join(parts)
+
+    def dump_filename(self) -> str:
+        return "torchcomms_fr_trace"
+
+
 def default_handlers() -> list[DebugHandler]:
     return [
         IndexHandler(),
         StacksHandler(),
         PySpyHandler(),
         FlightRecorderHandler(),
+        TorchCommsFlightRecorderHandler(),
         ProfilerHandler(),
         WaitCountersHandler(),
         TCPStoreHandler(),
diff --git a/torch/distributed/debug/_frontend.py b/torch/distributed/debug/_frontend.py
index 68361dc447722..5a8779ef7424e 100644
--- a/torch/distributed/debug/_frontend.py
+++ b/torch/distributed/debug/_frontend.py
@@ -6,7 +6,7 @@
 import threading
 import time
 from abc import ABC, abstractmethod
-from collections.abc import Callable, Iterable
+from collections.abc import Callable
 from concurrent.futures import ThreadPoolExecutor
 from dataclasses import dataclass
 from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
@@ -19,6 +19,8 @@
 
 logger: logging.Logger = logging.getLogger(__name__)
 
+_DEFAULT_FETCH_TIMEOUT: float = 60.0
+
 
 # ---------------------------------------------------------------------------
 # Base types
@@ -51,6 +53,8 @@ class Route:
 
 
 class DebugHandler(ABC):
+    fetch_timeout: float = _DEFAULT_FETCH_TIMEOUT
+
     @abstractmethod
     def routes(self) -> list[Route]: ...
 
@@ -72,49 +76,66 @@ def dump_filename(self) -> str:
 # ---------------------------------------------------------------------------
 
 
-def fetch_thread_pool(urls: list[str]) -> Iterable[Response]:
+def fetch_thread_pool(urls: list[str], timeout: float) -> list[Response]:
     # late import for optional dependency
     import requests
 
     max_workers = 20
 
     def get(url: str) -> Response:
-        resp = requests.post(url)
-        return Response(resp.status_code, resp.text)
+        try:
+            resp = requests.post(url, timeout=timeout)
+            return Response(resp.status_code, resp.text)
+        except requests.exceptions.Timeout as e:
+            return Response(408, f"Timeout: {e}")
+        except requests.exceptions.ConnectionError as e:
+            return Response(503, f"ConnectionError: {e}")
+        except Exception as e:
+            return Response(502, f"{type(e).__name__}: {e}")
 
     with ThreadPoolExecutor(max_workers=max_workers) as executor:
-        resps = executor.map(get, urls)
+        resps = list(executor.map(get, urls))
 
     return resps
 
 
-def fetch_aiohttp(urls: list[str]) -> Iterable[Response]:
+def fetch_aiohttp(urls: list[str], timeout: float) -> list[Response]:
     # late import for optional dependency
     # pyrefly: ignore [missing-import]
     import aiohttp
 
     async def fetch(session: aiohttp.ClientSession, url: str) -> Response:
-        async with session.post(url) as resp:
-            text = await resp.text()
-            return Response(resp.status, text)
-
-    async def gather(urls: list[str]) -> Iterable[Response]:
-        async with aiohttp.ClientSession() as session:
-            return await asyncio.gather(*[fetch(session, url) for url in urls])
+        try:
+            async with session.post(url) as resp:
+                text = await resp.text()
+                return Response(resp.status, text)
+        except asyncio.TimeoutError as e:
+            return Response(408, f"TimeoutError: {e}")
+        except aiohttp.ClientError as e:
+            return Response(503, f"{type(e).__name__}: {e}")
+        except Exception as e:
+            return Response(502, f"{type(e).__name__}: {e}")
+
+    async def gather(urls: list[str]) -> list[Response]:
+        client_timeout = aiohttp.ClientTimeout(total=timeout)
+        async with aiohttp.ClientSession(timeout=client_timeout) as session:
+            return list(await asyncio.gather(*[fetch(session, url) for url in urls]))
 
     return asyncio.run(gather(urls))
 
 
-def fetch_all(endpoint: str, args: str = "") -> tuple[list[str], Iterable[Response]]:
+def fetch_all(
+    endpoint: str, args: str = "", *, timeout: float
+) -> tuple[list[str], list[Response]]:
     store = tcpstore_client()
     keys = [f"rank{r}" for r in range(get_world_size())]
     addrs = store.multi_get(keys)
     addrs = [f"{addr.decode()}/handler/{endpoint}?{args}" for addr in addrs]
 
     try:
-        resps = fetch_aiohttp(addrs)
+        resps = fetch_aiohttp(addrs, timeout=timeout)
     except ImportError:
-        resps = fetch_thread_pool(addrs)
+        resps = fetch_thread_pool(addrs, timeout=timeout)
 
     return addrs, resps
 
@@ -124,6 +145,19 @@ def format_json(blob: str):
     return json.dumps(parsed, indent=2)
 
 
+def format_fetch_summary(addrs: list[str], resps: list[Response]) -> str | None:
+    """Return a summary string if any workers failed, or None if all succeeded."""
+    failed = [(i, r) for i, r in enumerate(resps) if r.status_code != 200]
+    if not failed:
+        return None
+    total = len(addrs)
+    ok = total - len(failed)
+    lines = [f"PARTIAL DATA: {ok}/{total} workers responded"]
+    for rank, resp in failed:
+        lines.append(f"  Rank {rank}: {resp.text}")
+    return "\n".join(lines)
+
+
 # ---------------------------------------------------------------------------
 # Template constants
 # ---------------------------------------------------------------------------
@@ -428,7 +462,11 @@ def main(
     dump_interval: float,
     handlers: list[DebugHandler],
     enabled_dumps: set[str],
+    fetch_timeout: float = 60.0,
 ) -> None:
+    for handler in handlers:
+        handler.fetch_timeout = fetch_timeout
+
     logger.setLevel(logging.INFO)
 
     server = FrontendServer(port=port, handlers=handlers)
diff --git a/torch/distributed/device_mesh.py b/torch/distributed/device_mesh.py
index e18ec608967ea..99b41eb2ad529 100644
--- a/torch/distributed/device_mesh.py
+++ b/torch/distributed/device_mesh.py
@@ -4,14 +4,16 @@
 import os
 import threading
 import warnings
-from collections.abc import Iterator
+from collections.abc import Callable, Iterator
 from itertools import zip_longest
-from typing import Optional, TYPE_CHECKING, Union
+from typing import Any, TYPE_CHECKING
 
 import torch
+from torch._opaque_base import OpaqueBase
 from torch.distributed import is_available
 from torch.distributed._mesh_layout import _MeshLayout
 from torch.distributed._pycute import IntTuple, is_int, suffix_product
+from torch.types import IntLikeType
 from torch.utils._typing_utils import not_none
 
 
@@ -148,7 +150,7 @@ def _get_device_handle(device_type: str = "cuda"):
         """
         return getattr(torch, device_type, None)
 
-    class DeviceMesh:
+    class DeviceMesh(OpaqueBase):
         """
         DeviceMesh represents a mesh of devices, where layout of devices could be
         represented as a n-d dimension array, and each value of the n-d dimensional
@@ -201,7 +203,8 @@ class DeviceMesh:
         _rank_map: torch.Tensor
         _mesh_dim_names: tuple[str, ...] | None
         _layout: _MeshLayout
-        _root_mesh: Optional["DeviceMesh"] = None
+        _root_mesh: "DeviceMesh | None" = None
+        _thread_id: int | None
         # Record flatten mesh name to its flattened mesh in root mesh.
         _flatten_mapping: dict[str, "DeviceMesh"]
         # Registry mapping group names to ProcessGroup objects (to avoid C++ lookup)
@@ -210,7 +213,7 @@ class DeviceMesh:
         def __init__(
             self,
             device_type: str,
-            mesh: Union[torch.Tensor, "ArrayLike"] | None = None,
+            mesh: "torch.Tensor | ArrayLike | None" = None,
             *,
             mesh_dim_names: tuple[str, ...] | None = None,
             backend_override: tuple[BackendConfig, ...] | None = None,
@@ -218,7 +221,7 @@ def __init__(
             _rank: int | None = None,
             _layout: _MeshLayout | None = None,
             _rank_map: torch.Tensor | None = None,
-            _root_mesh: Optional["DeviceMesh"] = None,
+            _root_mesh: "DeviceMesh | None" = None,
         ) -> None:
             # no-op in OSS, logs API usage metrics in meta-internal runs
             torch._C._log_api_usage_once(
@@ -244,15 +247,19 @@ def __init__(
                         "The mesh argument is required except for PRIVATE USAGE ONLY!"
                     )
 
-            assert _layout.check_non_overlap(), (
-                "Please use a non-overlapping layout when creating a DeviceMesh."
-            )
-            assert _rank_map.ndim == 1, "The rank map must be 1-dimensional"
-            assert _rank_map.is_contiguous(), "The rank map must be contiguous"
-            assert _rank_map.numel() >= _layout.cosize(), (
-                f"The rank map contains {_rank_map.numel()} element, "
-                f"which isn't large enough for layout {_layout}"
-            )
+            if not _layout.check_non_overlap():
+                raise AssertionError(
+                    "Please use a non-overlapping layout when creating a DeviceMesh."
+                )
+            if _rank_map.ndim != 1:
+                raise AssertionError("The rank map must be 1-dimensional")
+            if not _rank_map.is_contiguous():
+                raise AssertionError("The rank map must be contiguous")
+            if _rank_map.numel() < _layout.cosize():
+                raise AssertionError(
+                    f"The rank map contains {_rank_map.numel()} element, "
+                    f"which isn't large enough for layout {_layout}"
+                )
 
             self._device_type = device_type
             self._layout = _layout
@@ -319,7 +326,6 @@ def __init__(
                             target_registry[name] = pg
 
                 if is_initialized() and get_backend() == "threaded":
-                    # pyrefly: ignore [bad-assignment]
                     self._thread_id = threading.get_ident()
 
                 # Now that the process group is initialized, we can get the rank
@@ -330,6 +336,8 @@ def __init__(
 
                 self._coordinate_on_dim = self._compute_coordinate_on_dim()
 
+            self._hash: int | None = None
+
         @staticmethod
         def _compute_coordinates_from_mesh(
             mesh_tensor: torch.Tensor,
@@ -541,10 +549,12 @@ def _init_one_process_group(
                     getattr(default_group, "bound_device_id", None) is not None
                     or dist_config.use_torchcomms
                 )
-                and torch.cuda.is_available()
+                and torch.accelerator.is_available()
                 and (
                     backend is None
-                    or default_group._get_backend(torch.device("cuda")).name()
+                    or default_group._get_backend(
+                        torch.accelerator.current_accelerator()  # pyrefly: ignore[bad-argument-type]
+                    ).name()
                     == backend
                 )
             ):
@@ -607,9 +617,10 @@ def _init_process_groups(
                 )
             # Filter out None values. If any are None then they should all be None.
             dim_non_none_group_names = [n for n in dim_group_names if n is not None]
-            assert not dim_non_none_group_names or len(dim_non_none_group_names) == len(
+            if dim_non_none_group_names and len(dim_non_none_group_names) != len(
                 dim_group_names
-            )
+            ):
+                raise AssertionError
             return dim_non_none_group_names
 
         def _get_root_mesh(self) -> "DeviceMesh":
@@ -637,19 +648,21 @@ def __repr__(self) -> str:
                 device_mesh_repr += f", Mesh: {self.mesh.tolist()}"
             return f"{device_mesh_repr})"
 
+        def _hash_key(self) -> tuple[Any, ...]:
+            """Return the tuple used for hashing. Used by both __hash__ and _stable_hash."""
+            return (
+                self._flatten_rank_map,
+                self._layout,
+                self._device_type,
+                self._mesh_dim_names,
+                self._thread_id,
+            )
+
         def __hash__(self):
             # lazily compute hash
             self._hash = getattr(self, "_hash", None)
             if not self._hash:
-                self._hash = hash(
-                    (
-                        self._flatten_rank_map,
-                        self._layout,
-                        self._device_type,
-                        self._mesh_dim_names,
-                        self._thread_id,
-                    )
-                )
+                self._hash = hash(self._hash_key())
             return self._hash
 
         def __eq__(self, other: object) -> bool:
@@ -665,6 +678,17 @@ def __eq__(self, other: object) -> bool:
                 and self._thread_id == other._thread_id
             )
 
+        def _stable_hash(self) -> str:
+            """
+            Return a stable hash for AOT autograd caching.
+            [See note: Tensor subclass stable hashing for AOT autograd cache]
+            """
+            import hashlib
+
+            return hashlib.blake2b(
+                repr(self._hash_key()).encode(), digest_size=16
+            ).hexdigest()
+
         def __getitem__(self, mesh_dim_names: str | tuple[str, ...]) -> "DeviceMesh":
             """
             Slice the current DeviceMesh based on the mesh_dim_names given to create a submesh.
@@ -722,14 +746,22 @@ def __getitem__(self, mesh_dim_names: str | tuple[str, ...]) -> "DeviceMesh":
                 sliced_mesh_layout = self._get_slice_mesh_layout(mesh_dim_names)
                 # When using FakeTensorMode to trace the model, `_create_sub_mesh()` will
                 # fail as it will require a real tensor to manipulate.
-                # `unset_fake_temporarily()` will allow us to materialize the tensors
-                # within `_create_sub_mesh`, which should not affect modling.
+                # `unset_fake_temporarily()` and `disable_proxy_modes_tracing()`
+                # will allow us to materialize the tensors within
+                # `_create_sub_mesh`, which should not affect modling.
                 #
                 # Note that this should be orthogonal to torch.compile(). But whether
                 # we can compile device_mesh `slicing` (no graph break) is not verified
                 # yet and need a follow-up,
                 # TODO: compiler + device_mesh slicing.
-                with torch._subclasses.fake_tensor.unset_fake_temporarily():
+                from torch.fx.experimental.proxy_tensor import (
+                    disable_proxy_modes_tracing,
+                )
+
+                with (
+                    torch._subclasses.fake_tensor.unset_fake_temporarily(),
+                    disable_proxy_modes_tracing(),
+                ):
                     submesh = self._create_sub_mesh(sliced_mesh_layout, mesh_dim_names)
                 return submesh
 
@@ -799,12 +831,11 @@ def _create_sub_mesh(
             root_mesh = self._get_root_mesh()
             slice_dim_group_name = []
             if len(self._dim_group_names) > 0:
-                assert len(self._dim_group_names) == len(
-                    not_none(self._mesh_dim_names)
-                ), (
-                    "The number of dim_group_names and mesh_dim_names "
-                    "should have the same length if the rank is in the mesh."
-                )
+                if len(self._dim_group_names) != len(not_none(self._mesh_dim_names)):
+                    raise AssertionError(
+                        "The number of dim_group_names and mesh_dim_names "
+                        "should have the same length if the rank is in the mesh."
+                    )
                 for name in submesh_dim_names:
                     if name in not_none(self._mesh_dim_names):
                         slice_dim_group_name.append(
@@ -835,18 +866,15 @@ def _create_sub_mesh(
 
         def _create_flatten_mesh(
             self,
-            mesh_dim_name: str | None = None,
+            mesh_dim_name: str,
             backend_override: BackendConfig = (None, None),
-        ) -> "DeviceMesh":
+        ) -> None:
+            """
+            Creates a flattened mesh and stores it in the root mesh's _flatten_mapping.
+            This is a helper method for _flatten() that performs the actual mesh creation.
+            """
             root_mesh = self._get_root_mesh()
 
-            if not mesh_dim_name:
-                mesh_dim_name = "_".join(not_none(self._mesh_dim_names))
-
-            # Flatten a 1D device mesh into its original mesh_dim_name will return itself.
-            if self.ndim == 1 and mesh_dim_name in not_none(self._mesh_dim_names):
-                return self
-
             # Check whether the mesh_dim_name for flattened mesh is valid.
             invalid_dim_names = not_none(root_mesh._mesh_dim_names)
             if mesh_dim_name in invalid_dim_names:
@@ -859,13 +887,15 @@ def _create_flatten_mesh(
             flattened_mesh_layout = self._layout.coalesce()
             if len(flattened_mesh_layout) > 1:
                 flattened_mesh_layout = flattened_mesh_layout.nest()
-            # Quick return if the flatten mesh has been created before.
+
+            # Check if the flatten mesh has been created before.
             if mesh_dim_name in root_mesh._flatten_mapping:
                 if (
                     flattened_mesh_layout
                     == root_mesh._flatten_mapping[mesh_dim_name]._layout
                 ):
-                    return root_mesh._flatten_mapping[mesh_dim_name]
+                    # Already exists with same layout, nothing to do
+                    return
                 else:
                     raise ValueError(
                         f"Flatten mesh with mesh_dim_name {mesh_dim_name} has been created before, "
@@ -882,8 +912,6 @@ def _create_flatten_mesh(
             )
             root_mesh._flatten_mapping[mesh_dim_name] = res_flattened_mesh
 
-            return res_flattened_mesh
-
         def _get_root_mesh_dim(self) -> int | None:
             """
             Returns the index of the mesh dim in the root mesh.
@@ -1027,7 +1055,7 @@ def _get_all_submeshes(self, mesh_dim_name: str) -> list["DeviceMesh"]:
         def from_group(
             group: ProcessGroup | list[ProcessGroup],
             device_type: str,
-            mesh: Union[torch.Tensor, "ArrayLike"] | None = None,
+            mesh: "torch.Tensor | ArrayLike | None" = None,
             *,
             mesh_dim_names: tuple[str, ...] | None = None,
         ) -> "DeviceMesh":
@@ -1210,13 +1238,18 @@ def get_coordinate(self) -> tuple[int, ...] | None:
             """
             return self._coordinate_on_dim
 
-        def _sym_get_coordinate(self, index: int) -> int:
+        def _sym_get_coordinate(self, index: int) -> IntLikeType:
             import torch.distributed.config as config
             from torch._guards import detect_fake_mode
 
-            if not detect_fake_mode() or not config.compile_on_one_rank:
+            if (
+                not config.compile_on_one_rank
+                or not (fake_mode := detect_fake_mode())
+                or not fake_mode.shape_env
+            ):
                 # This is only valid when the current rank is part of the mesh.
-                assert self._coordinate_on_dim is not None
+                if self._coordinate_on_dim is None:
+                    raise AssertionError
                 return self._coordinate_on_dim[index]
 
             # This will cause the ops to be registered - so don't let RUFF
@@ -1259,6 +1292,14 @@ def _flatten(
                     "Cannot flatten a DeviceMesh without mesh_dim_names!"
                 )
 
+            # Compute mesh_dim_name first if not provided
+            if not mesh_dim_name:
+                mesh_dim_name = "_".join(not_none(self._mesh_dim_names))
+
+            # Flatten a 1D device mesh into its original mesh_dim_name will return itself.
+            if self.ndim == 1 and mesh_dim_name in not_none(self._mesh_dim_names):
+                return self
+
             if backend_override is not None:
                 (backend_override_tuple,) = _normalize_backend_override(
                     {0: backend_override}, 1
@@ -1266,7 +1307,9 @@ def _flatten(
             else:
                 backend_override_tuple = (None, None)
 
-            return self._create_flatten_mesh(mesh_dim_name, backend_override_tuple)
+            self._create_flatten_mesh(mesh_dim_name, backend_override_tuple)
+
+            return self._get_root_mesh()._flatten_mapping[mesh_dim_name]
 
         def _create_unflatten_mesh(
             self,
@@ -1354,7 +1397,7 @@ def _unflatten(
                 raise ValueError(
                     f"dim {dim} specified in `_unflatten` is out of range {self.ndim}"
                 )
-            elif isinstance(dim, str) and dim in not_none(self.mesh_dim_names):
+            elif isinstance(dim, str) and dim not in not_none(self.mesh_dim_names):
                 raise ValueError(
                     f"dim {dim} specified in `_unflatten` is not in {self.mesh_dim_names}"
                 )
@@ -1511,14 +1554,14 @@ def init_device_mesh(
         if mesh_dim_names is not None:
             if len(set(mesh_dim_names)) != len(mesh_dim_names):
                 raise RuntimeError(
-                    "Each mesh_dim_name must be unique.",
-                    f"Found repeated mesh_dim_name in mesh_dim_names {mesh_dim_names}",
+                    "Each mesh_dim_name must be unique. "
+                    f"Found repeated mesh_dim_name in mesh_dim_names {mesh_dim_names}"
                 )
 
             if len(mesh_shape) != len(mesh_dim_names):
                 raise RuntimeError(
-                    "mesh_shape and mesh_dim_names should have same length!",
-                    f"Found len(mesh_dim_names): {len(mesh_dim_names)} and len(mesh_shape):{len(mesh_shape)}.",
+                    "mesh_shape and mesh_dim_names should have same length! "
+                    f"Found len(mesh_dim_names): {len(mesh_dim_names)} and len(mesh_shape):{len(mesh_shape)}."
                 )
 
         if backend_override is not None:
@@ -1551,3 +1594,154 @@ def init_device_mesh(
         )
 
         return device_mesh
+
+
+_distributed_opaque_types_registered = False
+
+
+def _device_mesh_reconstruct_fn(
+    mesh: "OpaqueBase",
+    get_tracked_proxy: Callable[["OpaqueBase"], "torch.fx.Proxy | None"],
+    tracer: Any,
+) -> "torch.fx.Proxy | None":
+    """Reconstruct a DeviceMesh submesh from a tracked ancestor mesh.
+
+    Called by PythonKeyTracer when make_fx encounters a DeviceMesh that isn't
+    tracked (e.g. a submesh captured by a backward closure). Looks for any
+    tracked mesh that shares the same root and contains the target dim names,
+    then emits a call_function node that derives the submesh via _get_submesh.
+    """
+    if not isinstance(mesh, DeviceMesh):
+        raise AssertionError("DeviceMesh expected")
+
+    root_mesh = mesh._get_root_mesh()
+
+    # Only submeshes can be reconstructed; root meshes must already be tracked.
+    if mesh is root_mesh:
+        return None
+
+    dim_names = mesh._mesh_dim_names
+    if dim_names is None:
+        return None
+
+    # Ensure the custom ops are registered
+    from torch.distributed._ops import device_mesh as _dm_ops  # noqa: F401
+
+    # Try the root mesh first (original path).
+    ancestor_proxy = get_tracked_proxy(root_mesh)
+    ancestor_dim_names = root_mesh._mesh_dim_names
+
+    # If root isn't tracked, search for any tracked DeviceMesh that shares
+    # the same root AND contains all our dim names. This handles the case
+    # where e.g. a concatenated (fsdp, tp) mesh is a graph input (from
+    # DTensor.__tensor_flatten__) but neither root nor the individual
+    # submeshes are tracked directly.
+    if ancestor_proxy is None:
+        from torch._library.fake_class_registry import FakeScriptObject
+
+        for tracked_obj, proxy in tracer.opaque_tracker.items():
+            real_obj = (
+                tracked_obj.real_obj
+                if isinstance(tracked_obj, FakeScriptObject)
+                else tracked_obj
+            )
+            if not isinstance(real_obj, DeviceMesh) or real_obj is mesh:
+                continue
+            if real_obj._get_root_mesh() is not root_mesh:
+                continue
+            tracked_dim_names = real_obj._mesh_dim_names
+            if tracked_dim_names is None:
+                continue
+            if all(n in tracked_dim_names for n in dim_names):
+                ancestor_proxy = proxy
+                ancestor_dim_names = tracked_dim_names
+                break
+
+    if ancestor_proxy is None or ancestor_dim_names is None:
+        return None
+
+    # Convert our dim names to indices into the ancestor mesh's dim names
+    mesh_dims = [ancestor_dim_names.index(n) for n in dim_names]
+
+    # Dispatch through the custom op with proxy mode active so that
+    # meta["val"] is set and the result is tracked in opaque_tracker.
+    return torch.ops.device_mesh._get_submesh(ancestor_proxy, mesh_dims)
+
+
+def _register_distributed_opaque_types():
+    """
+    Register DeviceMesh as an opaque type for torch.compile.
+    This must happen before any custom ops that use DeviceMesh in their schema.
+    Called lazily to avoid circular import issues.
+    """
+    global _distributed_opaque_types_registered
+    if _distributed_opaque_types_registered:
+        return
+    _distributed_opaque_types_registered = True
+
+    from torch._library.opaque_object import MemberType, register_opaque_type
+
+    register_opaque_type(
+        ProcessGroup,
+        typ="reference",
+        members={
+            "size": MemberType.USE_REAL,
+            "rank": MemberType.USE_REAL,
+            "_get_backend_name": MemberType.USE_REAL,
+            "group_name": MemberType.USE_REAL,
+            "group_desc": MemberType.USE_REAL,
+            "__eq__": MemberType.USE_REAL,
+            "__ne__": MemberType.USE_REAL,
+        },
+    )
+
+    register_opaque_type(
+        DeviceMesh,
+        typ="reference",
+        reconstruct_fn=_device_mesh_reconstruct_fn,
+        guard_fn=lambda obj: [
+            obj._flatten_rank_map,
+            obj._layout,
+            obj._device_type,
+            obj._mesh_dim_names,
+            obj._thread_id,
+        ],
+        members={
+            # USE_REAL: Evaluate these with the real object at compile time
+            # and bake the result as a constant
+            "_flatten_rank_map": MemberType.USE_REAL,
+            "_layout": MemberType.USE_REAL,
+            "_device_type": MemberType.USE_REAL,
+            "_mesh_dim_names": MemberType.USE_REAL,
+            "_thread_id": MemberType.USE_REAL,
+            "get_rank": MemberType.USE_REAL,
+            "size": MemberType.USE_REAL,
+            "get_coordinate": MemberType.USE_REAL,
+            "get_local_rank": MemberType.USE_REAL,
+            "__eq__": MemberType.USE_REAL,
+            "__ne__": MemberType.USE_REAL,
+            "ndim": MemberType.USE_REAL,
+            "shape": MemberType.USE_REAL,
+            "mesh_dim_names": MemberType.USE_REAL,
+            "get_group": MemberType.INLINED,
+            "_pg_registry": MemberType.INLINED,
+            "_hash": MemberType.USE_REAL,
+            "_create_flatten_mesh": MemberType.USE_REAL,
+            "_coordinate_on_dim": MemberType.USE_REAL,
+            "_dim_group_names": MemberType.USE_REAL,
+            "_flatten_mapping": MemberType.USE_REAL,
+            "_rank_map": MemberType.USE_REAL,
+            "_root_mesh": MemberType.USE_REAL,
+            "device_type": MemberType.USE_REAL,
+            "mesh": MemberType.USE_REAL,
+            "_flatten": MemberType.INLINED,
+            "_unflatten": MemberType.USE_REAL,
+            "_is_current_rank_part_of_mesh": MemberType.USE_REAL,
+            "_sym_get_coordinate": MemberType.USE_REAL,
+            "_get_mesh_dim_by_name": MemberType.USE_REAL,
+            "_get_root_mesh": MemberType.INLINED,
+            "__getitem__": MemberType.INLINED,
+            "_get_slice_mesh_layout": MemberType.INLINED,
+            "_create_sub_mesh": MemberType.INLINED,
+        },
+    )
diff --git a/torch/distributed/distributed_c10d.py b/torch/distributed/distributed_c10d.py
index f14a24db8c190..05986c191a233 100644
--- a/torch/distributed/distributed_c10d.py
+++ b/torch/distributed/distributed_c10d.py
@@ -134,6 +134,7 @@
     "get_node_local_rank",
     "split_group",
     "shrink_group",
+    "record_comm",
 ]
 
 _MPI_AVAILABLE = True
@@ -143,8 +144,18 @@
 _XCCL_AVAILABLE = True
 
 try:
-    # pyrefly: ignore [import-error, missing-import]
-    from torchcomms._comms import _BackendWrapper, new_comm
+    try:
+        # pyrefly: ignore [missing-import]
+        from torchcomms._comms import _BackendWrapper
+    except ImportError:
+        # pyrefly: ignore [missing-import]
+        from torchcomms._backend_wrapper import _BackendWrapper
+
+    # pyrefly: ignore [missing-import]
+    from torchcomms import new_comm
+
+    # pyrefly: ignore [missing-import]
+    from torchcomms.hooks import FlightRecorderHook
 
     _TORCHCOMM_AVAILABLE = True
 except ImportError:
@@ -364,8 +375,12 @@ def register_backend(
 
         if devices is not None:
             for device in devices:
-                if device not in Backend.default_device_backend_map:
+                current = Backend.default_device_backend_map.get(device)
+                # Allow remapping from fake backend to actual backend (e.g., HPU from fake to HCCL)
+                # but prevent fake backend from claiming devices
+                if current is None or (current == "fake" and name.lower() != "fake"):
                     Backend.default_device_backend_map[device] = name.lower()
+
         Backend.backend_type_map[name.lower()] = ProcessGroup.BackendType.CUSTOM
 
         # Update device capability matrix in Backend class
@@ -1137,9 +1152,11 @@ def _get_group_size(group: ProcessGroup | None) -> int:
     return group.size()
 
 
-def _get_group_size_by_name(group_name: GroupName) -> int:
-    group = _resolve_process_group(group_name)
-    return group.size()
+def _get_group_size_by_name(group_name: GroupName | ProcessGroup) -> int:
+    if isinstance(group_name, str):
+        # pyrefly: ignore[bad-argument-type]  # pyrefly bug
+        group_name = _resolve_process_group(group_name)
+    return group_name.size()
 
 
 def _resolve_group_name_by_ranks_and_tag(ranks: list[int], tag: str) -> GroupName:
@@ -2053,10 +2070,15 @@ def _new_process_group_helper(
             comm = new_comm(
                 backend_str, torch_device, name=group_name, store=backend_prefix_store
             )
+            buffer_size = os.environ.get(
+                "TORCH_FR_BUFFER_SIZE",
+                os.environ.get("TORCH_NCCL_TRACE_BUFFER_SIZE", "0"),
+            )
+            recorder = FlightRecorderHook(max_entries=int(buffer_size))
+            recorder.register_with_comm(comm)
             # Keep a reference so the comm outlives this function scope.
             _world.comms.append(comm)
             group_name = GroupName(group_name)
-            # pyrefly: ignore [bad-assignment]
             backend_class = _BackendWrapper(comm)
             backend_type = ProcessGroup.BackendType.CUSTOM
         elif backend_str == Backend.MPI:
@@ -2539,6 +2561,18 @@ def isend(
         None, if not part of the group
 
     """
+    relevant_args = (tensor,)
+    if has_torch_function(relevant_args):
+        return handle_torch_function(
+            isend,
+            relevant_args,
+            tensor,
+            dst=dst,
+            group=group,
+            tag=tag,
+            group_dst=group_dst,
+        )
+
     group = _group_or_default_group(group)
     group_dst = _canonicalize_group_rank(group, dst, group_dst)
     _check_single_tensor(tensor, "tensor")
@@ -2581,6 +2615,18 @@ def irecv(
         None, if not part of the group
 
     """
+    relevant_args = (tensor,)
+    if has_torch_function(relevant_args):
+        return handle_torch_function(
+            irecv,
+            relevant_args,
+            tensor,
+            src=src,
+            group=group,
+            tag=tag,
+            group_src=group_src,
+        )
+
     _check_single_tensor(tensor, "tensor")
     if _rank_not_in_group(group):
         _warn_not_in_group("irecv")
@@ -2621,6 +2667,18 @@ def send(
         group_dst (int, optional): Destination rank on ``group``.  Invalid to specify both ``dst`` and ``group_dst``.
 
     """
+    relevant_args = (tensor,)
+    if has_torch_function(relevant_args):
+        return handle_torch_function(
+            send,
+            relevant_args,
+            tensor,
+            dst=dst,
+            group=group,
+            tag=tag,
+            group_dst=group_dst,
+        )
+
     group = _group_or_default_group(group)
     group_dst = _canonicalize_group_rank(group, dst, group_dst)
     _check_not_self_rank(group, group_dst, "destination")
@@ -2651,11 +2709,24 @@ def recv(
             the default process group will be used.
         tag (int, optional): Tag to match recv with remote send
         group_src (int, optional): Destination rank on ``group``.  Invalid to specify both ``src`` and ``group_src``.
+
     Returns:
         Sender rank
         -1, if not part of the group
 
     """
+    relevant_args = (tensor,)
+    if has_torch_function(relevant_args):
+        return handle_torch_function(
+            recv,
+            relevant_args,
+            tensor,
+            src=src,
+            group=group,
+            tag=tag,
+            group_src=group_src,
+        )
+
     work = irecv(tensor, src=src, group=group, tag=tag, group_src=group_src)
     if work is None:
         return -1
@@ -2749,6 +2820,12 @@ def _coalescing_manager(
         # - coalesced `all_gather_into_tensor`
         # - coalesced `reduce_scatter_tensor`
         op0 = op_list[0].op
+        if any(op.op is not op0 for op in op_list):
+            raise RuntimeError(
+                "Coalescing manager requires all collectives to be the same type, "
+                f"but got mixed types: {set(op.op.__name__ for op in op_list)}"  # noqa: C401
+            )
+
         if op0 is all_reduce:
             tensors = [op.tensor for op in op_list]
             all_reduce_opts = AllreduceCoalescedOptions()
@@ -2763,7 +2840,9 @@ def _coalescing_manager(
                 outputs.append(not_none(op.dst_tensor))
             all_gather_opts = AllgatherOptions()
             all_gather_opts.asyncOp = async_ops
-            work = group.allgather_into_tensor_coalesced(outputs, inputs)
+            work = group.allgather_into_tensor_coalesced(
+                outputs, inputs, all_gather_opts
+            )
         elif op0 is reduce_scatter_tensor:
             inputs = []
             outputs = []
@@ -2923,6 +3002,15 @@ def peer_kwarg(op: P2POp) -> dict[str, int]:
         return reqs
 
 
+def _is_fp8(tensor):
+    return tensor.dtype in (
+        torch.float8_e4m3fn,
+        torch.float8_e4m3fnuz,
+        torch.float8_e5m2,
+        torch.float8_e5m2fnuz,
+    )
+
+
 @_exception_logger
 def broadcast(
     tensor: torch.Tensor,
@@ -2952,6 +3040,18 @@ def broadcast(
         None, if not async_op or if not part of the group
 
     """
+    relevant_args = (tensor,)
+    if has_torch_function(relevant_args):
+        return handle_torch_function(
+            broadcast,
+            relevant_args,
+            tensor,
+            src=src,
+            group=group,
+            async_op=async_op,
+            group_src=group_src,
+        )
+
     group = _group_or_default_group(group)
     group_src = _canonicalize_group_rank(group, src, group_src, return_global=False)
     _check_single_tensor(tensor, "tensor")
@@ -2963,8 +3063,14 @@ def broadcast(
     opts.rootRank = group_src
     opts.rootTensor = 0
     opts.asyncOp = async_op
+    sm90_or_more = not (
+        tensor.is_cuda and torch.cuda.get_device_capability(tensor.device)[0] >= 9
+    )
     if tensor.is_complex():
         tensor = torch.view_as_real(tensor)
+    elif _is_fp8(tensor) and not sm90_or_more:
+        # FP8 is supported by NCCL on sm90+, use workaround for older GPUs
+        tensor = tensor.view(torch.uint8)
     work = group.broadcast([tensor], opts)
     if async_op:
         return work
@@ -3119,6 +3225,17 @@ def all_reduce_coalesced(tensors, op=ReduceOp.SUM, group=None, async_op: bool =
     """
     if isinstance(tensors, torch.Tensor):
         tensors = [tensors]
+    relevant_args = tuple(tensors) if isinstance(tensors, (list, tuple)) else (tensors,)
+    if has_torch_function(relevant_args):
+        return handle_torch_function(
+            all_reduce_coalesced,
+            relevant_args,
+            tensors,
+            op=op,
+            group=group,
+            async_op=async_op,
+        )
+
     _check_tensor_list(tensors, "tensor")
     _ensure_all_tensors_same_dtype(tensors)
     if _rank_not_in_group(group):
@@ -3177,6 +3294,19 @@ def reduce(
         None, if not async_op or if not part of the group
 
     """
+    relevant_args = (tensor,)
+    if has_torch_function(relevant_args):
+        return handle_torch_function(
+            reduce,
+            relevant_args,
+            tensor,
+            dst=dst,
+            op=op,
+            group=group,
+            async_op=async_op,
+            group_dst=group_dst,
+        )
+
     group = _group_or_default_group(group)
     group_dst = _canonicalize_group_rank(group, dst, group_dst, return_global=False)
     _check_single_tensor(tensor, "tensor")
@@ -4277,6 +4407,21 @@ def all_gather_coalesced(
     performance improvements but users of this function should take extra care
     to ensure that each node passes in tensors whose shapes match across nodes.
     """
+    relevant_args = (
+        tuple(input_tensor_list)
+        if isinstance(input_tensor_list, (list, tuple))
+        else (input_tensor_list,)
+    )
+    if has_torch_function(relevant_args):
+        return handle_torch_function(
+            all_gather_coalesced,
+            relevant_args,
+            output_tensor_lists,
+            input_tensor_list,
+            group=group,
+            async_op=async_op,
+        )
+
     # We only check basic compatibility with C++ params here, C++ code will
     # do shape and type checking.
     if _rank_not_in_group(group):
@@ -4375,6 +4520,19 @@ def gather(
         None                                                                   # Rank 1
 
     """
+    relevant_args = (tensor,)
+    if has_torch_function(relevant_args):
+        return handle_torch_function(
+            gather,
+            relevant_args,
+            tensor,
+            gather_list=gather_list,
+            dst=dst,
+            group=group,
+            async_op=async_op,
+            group_dst=group_dst,
+        )
+
     _check_single_tensor(tensor, "tensor")
 
     # Parameter ``gather_list`` may be left unspecified on non-dst ranks.
@@ -4465,6 +4623,19 @@ def scatter(
         tensor([5., 5.], device='cuda:1') # Rank 1
 
     """
+    relevant_args = (tensor,)
+    if has_torch_function(relevant_args):
+        return handle_torch_function(
+            scatter,
+            relevant_args,
+            tensor,
+            scatter_list=scatter_list,
+            src=src,
+            group=group,
+            async_op=async_op,
+            group_src=group_src,
+        )
+
     _check_single_tensor(tensor, "tensor")
     # Parameter ``scatter_list`` may be left unspecified on non-src ranks.
     if scatter_list:
@@ -4536,6 +4707,18 @@ def reduce_scatter(
         None, if not async_op or if not part of the group.
 
     """
+    relevant_args = (output,)
+    if has_torch_function(relevant_args):
+        return handle_torch_function(
+            reduce_scatter,
+            relevant_args,
+            output,
+            input_list,
+            op=op,
+            group=group,
+            async_op=async_op,
+        )
+
     _check_single_tensor(output, "output")
     _check_tensor_list(input_list, "input_list")
     _ensure_all_tensors_same_dtype(output, input_list)
@@ -4934,6 +5117,21 @@ def all_to_all(
         [tensor([4+4j]), tensor([8+8j]), tensor([12+12j]), tensor([16+16j])]        # Rank 3
 
     """
+    relevant_args = (
+        tuple(input_tensor_list)
+        if isinstance(input_tensor_list, (list, tuple))
+        else (input_tensor_list,)
+    )
+    if has_torch_function(relevant_args):
+        return handle_torch_function(
+            all_to_all,
+            relevant_args,
+            output_tensor_list,
+            input_tensor_list,
+            group=group,
+            async_op=async_op,
+        )
+
     if _rank_not_in_group(group):
         _warn_not_in_group("all_to_all")
         return
@@ -4968,6 +5166,7 @@ def barrier(
     group: ProcessGroup | None = GroupMember.WORLD,
     async_op: bool = False,
     device_ids=None,
+    timeout: timedelta | None = None,
 ):
     """
     Synchronize all processes.
@@ -4980,6 +5179,8 @@ def barrier(
             the default process group will be used.
         async_op (bool, optional): Whether this op should be an async op
         device_ids ([int], optional): List of device/GPU ids. Only one id is expected.
+        timeout (datetime.timedelta, optional): Timeout for barrier.
+            If ``None``, the default process group timeout will be used.
 
     Returns:
         Async work handle, if async_op is set to True.
@@ -5000,25 +5201,24 @@ def barrier(
 
     opts = BarrierOptions()
     opts.asyncOp = async_op
+    if timeout is not None:
+        opts.timeout = timeout
     # Detect the accelerator on the machine. If no accelerator is available, it
     # returns CPU.
     device = torch._C._get_accelerator()
     if isinstance(device_ids, list):
         opts.device_ids = device_ids
         # use only the first device id
-        # pyrefly: ignore [read-only]
         opts.device = torch.device(device.type, device_ids[0])
     elif getattr(group, "bound_device_id", None) is not None:
         # Use device id from `init_process_group(device_id=...)`
         opts.device = group.bound_device_id  # type: ignore[assignment]
     elif device.type == "cpu" or _get_object_coll_device(group) == "cpu":
-        # pyrefly: ignore [read-only]
         opts.device = torch.device("cpu")
     else:
         # Use the current device set by the user. If user did not set any, this
         # may use default device 0, causing issues like hang or all processes
         # creating context on device 0.
-        # pyrefly: ignore [read-only]
         opts.device = device
         if group.rank() == 0:
             warnings.warn(  # warn only once
@@ -5263,10 +5463,18 @@ def split_group(
         )
 
     parent_group_rank = parent_global_to_group_ranks[global_rank]
-    parent_backend = parent_pg._get_backend(torch.device("cuda"))
+
+    if torch.accelerator.is_available():
+        parent_backend = parent_pg._get_backend(
+            torch.accelerator.current_accelerator()  # pyrefly: ignore[bad-argument-type]
+        )
+    else:
+        raise RuntimeError(
+            "No backend for the parent process group or its backend does not support splitting"
+        )
 
     # if the parent backend does not support splitting, raise error
-    # currently this API only support NCCL backend
+    # currently this API only support NCCL and XCCL backend
     if (
         not parent_backend or not parent_backend.supports_splitting
     ) and not _use_torchcomms_enabled():
@@ -5332,7 +5540,16 @@ def split_group(
 
     global_ranks_in_my_group = [parent_group_to_global_ranks[rank] for rank in my_group]
     split_pg.bound_device_id = device_id  # type: ignore[union-attr]
-    split_backend_class = split_pg._get_backend(torch.device("cuda"))
+
+    if torch.accelerator.is_available():
+        split_backend_class = split_pg._get_backend(
+            torch.accelerator.current_accelerator()  # pyrefly: ignore[bad-argument-type]
+        )
+    else:
+        raise RuntimeError(
+            "No backend for the parent process group or its backend does not support splitting"
+        )
+
     if not _use_torchcomms_enabled():
         split_backend_class._set_sequence_number_for_group()
     if split_pg.group_name != group_name:
@@ -5370,6 +5587,7 @@ def new_group(
     use_local_synchronization: bool = False,
     group_desc=None,
     device_id: torch.device | None = None,
+    sort_ranks: bool = True,
 ):
     """
     Create a new distributed group.
@@ -5424,6 +5642,11 @@ def new_group(
         device_id (torch.device, optional): a single, specific device
             to "bind" this process to,  The `new_group` call will try to initialize
             a communication backend immediately for the device if this field is given.
+        sort_ranks (bool, optional): If ``True`` (the default), sort the
+            ``ranks`` list before creating the group. If ``False``, preserve
+            the caller-provided rank order so that position in the ``ranks``
+            list determines group rank.  All processes must pass the identical
+            ``ranks`` list.  Default is ``True``.
 
     Returns:
         A handle of distributed group that can be given to collective calls or
@@ -5448,6 +5671,7 @@ def new_group(
         use_local_synchronization=use_local_synchronization,
         group_desc=group_desc,
         device_id=device_id,
+        sort_ranks=sort_ranks,
     )
 
 
@@ -5460,6 +5684,7 @@ def _new_group_with_tag(
     use_local_synchronization=False,
     group_desc=None,
     device_id: torch.device | None = None,
+    sort_ranks: bool = True,
 ):
     """
     Variant of ``new_group`` that exposes tag creation.
@@ -5504,7 +5729,12 @@ def _new_group_with_tag(
 
     # checks the input ranks
     if ranks is not None:
-        ranks = sorted(ranks)
+        if sort_ranks:
+            ranks = sorted(ranks)
+        if len(set(ranks)) != len(ranks):
+            raise ValueError(
+                f"ranks list must not contain duplicate entries, got {ranks}"
+            )
         group_world_size = len(ranks)
         if group_world_size > global_world_size:
             raise ValueError(
@@ -5516,8 +5746,8 @@ def _new_group_with_tag(
         for rank in ranks:
             if rank < 0 or rank >= global_world_size:
                 raise ValueError(
-                    "The new group's rank should be within "
-                    "the world_size set by init_process_group"
+                    f"Rank {rank} is out of range. Valid ranks are 0 to {global_world_size - 1} "
+                    f"(world_size={global_world_size})"
                 )
         if global_rank in ranks:
             group_rank = ranks.index(global_rank)
@@ -6352,3 +6582,28 @@ def _update_process_group_global_state(
         # Standard process group tag
         _world.tags_to_pg.setdefault(pg_tag, []).append(pg)
         _world.pg_to_tag[pg] = pg_tag
+
+
+@contextlib.contextmanager
+def record_comm(name: str):
+    """Context manager to set a custom profiling name for communication collectives.
+
+    When active, all c10d collectives issued within this context will use ``name``
+    as their profiling title in the Work base class, overriding the default
+    backend-specific name (e.g. ``nccl:all_reduce``). This works across all
+    backends without per-backend or per-collective changes.
+
+    Args:
+        name (str): The profiling name to associate with collectives.
+
+    Example::
+        >>> # xdoctest: +SKIP("undefined vars")
+        >>> with dist.record_comm("FSDP::all_gather (layer1)"):
+        ...     dist.all_gather_into_tensor(output, input, group=pg)
+    """
+    prev = torch._C._distributed_c10d._get_comm_profiling_name()
+    torch._C._distributed_c10d._set_comm_profiling_name(name)
+    try:
+        yield
+    finally:
+        torch._C._distributed_c10d._set_comm_profiling_name(prev)
diff --git a/torch/distributed/elastic/agent/server/api.py b/torch/distributed/elastic/agent/server/api.py
index 2575aa137a581..fe3dda0841a1d 100644
--- a/torch/distributed/elastic/agent/server/api.py
+++ b/torch/distributed/elastic/agent/server/api.py
@@ -105,8 +105,10 @@ class WorkerSpec:
     virtual_local_rank: bool = False
 
     def __post_init__(self):
-        assert self.local_world_size > 0
-        assert self.monitor_interval > 0
+        if self.local_world_size <= 0:
+            raise AssertionError
+        if self.monitor_interval <= 0:
+            raise AssertionError
 
         if self.fn:
             warnings.warn(
@@ -116,7 +118,8 @@ def __post_init__(self):
                 category=DeprecationWarning,
             )
             self.entrypoint = self.fn
-        assert self.entrypoint
+        if not self.entrypoint:
+            raise AssertionError
 
     def get_entrypoint_name(self):
         """Get the entry point name.
@@ -127,7 +130,8 @@ def get_entrypoint_name(self):
         if isinstance(self.entrypoint, str):
             return os.path.basename(self.entrypoint)
         else:
-            assert self.entrypoint is not None
+            if self.entrypoint is None:
+                raise AssertionError
             return self.entrypoint.__qualname__
 
 
@@ -455,12 +459,19 @@ class SimpleElasticAgent(ElasticAgent):
     such as one particular type of worker role.
     """
 
-    def __init__(self, spec: WorkerSpec, exit_barrier_timeout: float = 300):
+    def __init__(
+        self,
+        spec: WorkerSpec,
+        exit_barrier_timeout: float = 300,
+        shutdown_timeout: int = 30,
+    ):
         self._worker_group = WorkerGroup(spec)
         self._remaining_restarts = self._worker_group.spec.max_restarts
         self._store = None
         self._exit_barrier_timeout = exit_barrier_timeout
+        self._shutdown_timeout = shutdown_timeout
         self._total_execution_time = 0
+        self._in_exit_barrier: bool = False
 
     def get_worker_group(self, role: str = DEFAULT_ROLE) -> WorkerGroup:
         return self._worker_group
@@ -493,11 +504,14 @@ def _monitor_workers(self, worker_group: WorkerGroup) -> RunResult:
         raise NotImplementedError
 
     @abc.abstractmethod
-    def _shutdown(self, death_sig: signal.Signals = signal.SIGTERM) -> None:
+    def _shutdown(
+        self, death_sig: signal.Signals = signal.SIGTERM, timeout: int = 30
+    ) -> None:
         """Clean up any resources that were allocated during the agent's work.
 
         Args:
             death_sig: Signal to send to the child process, SIGTERM is default
+            timeout: Time to wait for graceful shutdown before sending SIGKILL
         """
         raise NotImplementedError
 
@@ -737,12 +751,12 @@ def run(self, role: str = DEFAULT_ROLE) -> RunResult:
             logger.info("Rendezvous gracefully exited: %s", e)  # noqa: G200
         except SignalException as e:
             logger.warning("Received %s death signal, shutting down workers", e.sigval)
-            self._shutdown(e.sigval)
+            self._shutdown(e.sigval, timeout=self._shutdown_timeout)
             shutdown_called = True
             raise
         finally:
             if not shutdown_called:
-                self._shutdown()
+                self._shutdown(timeout=self._shutdown_timeout)
             # record the execution time in case there were any exceptions during run.
             self._total_execution_time = int(time.monotonic() - start_time)
 
@@ -904,7 +918,8 @@ def _invoke_run(self, role: str = DEFAULT_ROLE) -> RunResult:
         rdzv_handler = spec.rdzv_handler
 
         while True:
-            assert self._worker_group.state != WorkerState.INIT
+            if self._worker_group.state == WorkerState.INIT:
+                raise AssertionError
             time.sleep(monitor_interval)
             run_result = self._monitor_workers(self._worker_group)
             state = run_result.state
@@ -974,6 +989,7 @@ def _exit_barrier(self):
             self._exit_barrier_timeout,
         )
         start = time.time()
+        self._in_exit_barrier = True
         try:
             store_util.barrier(
                 store=self._store,
@@ -993,3 +1009,5 @@ def _exit_barrier(self):
                 "Error waiting on exit barrier. Elapsed: %s seconds",
                 time.time() - start,
             )
+        finally:
+            self._in_exit_barrier = False
diff --git a/torch/distributed/elastic/agent/server/local_elastic_agent.py b/torch/distributed/elastic/agent/server/local_elastic_agent.py
index 50808edc0bd0a..fda632bb3aaba 100644
--- a/torch/distributed/elastic/agent/server/local_elastic_agent.py
+++ b/torch/distributed/elastic/agent/server/local_elastic_agent.py
@@ -7,6 +7,7 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+from __future__ import annotations
 
 import json
 import os
@@ -19,6 +20,7 @@
 from typing import Any, TYPE_CHECKING
 
 import torch.distributed.elastic.timer as timer
+from torch._utils_internal import justknobs_check
 from torch.distributed.elastic import events
 from torch.distributed.elastic.agent.server.api import (
     RunResult,
@@ -42,6 +44,8 @@
 
 
 if TYPE_CHECKING:
+    from collections.abc import Callable
+
     from torch.distributed.elastic.events.api import EventMetadataValue
 
 logger = get_logger(__name__)
@@ -58,6 +62,29 @@
 TORCHELASTIC_TIMER_FILE = "TORCHELASTIC_TIMER_FILE"
 
 
+class _AliveCallbackProxy:
+    """Mutable callback wrapper for the health check server.
+
+    The C++ pybind ``HealthCheckThriftServer`` binds its ``alive_callback``
+    at construction time and cannot update it afterward.  This proxy is
+    created *before* the health check server so it can be passed as the
+    callback.  Initially it returns ``time.time()`` (signalling "alive").
+    After the agent is constructed, :meth:`set_delegate` wires it to
+    ``agent._get_alive_time`` for real liveness tracking.
+    """
+
+    def __init__(self) -> None:
+        self._delegate: Callable[[], int] | None = None
+
+    def __call__(self) -> int:
+        if self._delegate is not None:
+            return self._delegate()
+        return int(time.time())
+
+    def set_delegate(self, delegate: Callable[[], int]) -> None:
+        self._delegate = delegate
+
+
 class LocalElasticAgent(SimpleElasticAgent):
     """An implementation of :py:class:`torchelastic.agent.server.ElasticAgent` that handles host-local workers.
 
@@ -154,15 +181,17 @@ def __init__(
         start_method="spawn",
         exit_barrier_timeout: float = 300,
         log_line_prefix_template: str | None = None,
+        shutdown_timeout: int = 30,
+        health_check_server: HealthCheckServer | None = None,
     ):
-        super().__init__(spec, exit_barrier_timeout)
+        super().__init__(spec, exit_barrier_timeout, shutdown_timeout)
         self._start_method = start_method
         self._pcontext: PContext | None = None
         self._rdzv_handler = spec.rdzv_handler
         self._log_line_prefix_template = log_line_prefix_template
         self._worker_watchdog: timer.FileTimerServer | None = None
         self._logs_specs = logs_specs
-        self._health_check_server: HealthCheckServer | None = None
+        self._health_check_server = health_check_server
 
     def _setup_local_watchdog(self, envs: dict[int, dict[str, str]]) -> None:
         enable_watchdog_env_name = TORCHELASTIC_ENABLE_FILE_TIMER
@@ -205,6 +234,27 @@ def _setup_local_watchdog(self, envs: dict[int, dict[str, str]]) -> None:
     def _get_current_time_secs() -> int:
         return int(time.time())
 
+    def _get_alive_time(self) -> int:
+        """Return the last progress time from the watchdog, or the current time.
+
+        This callback is passed to the health check server at startup and
+        is called on every TW health check poll. During initialization
+        (before rendezvous and worker launch), the watchdog does not exist
+        yet, so we return the current time to signal the agent is alive.
+        Once workers are running and the watchdog is active, we delegate
+        to the watchdog's ``get_last_progress_time`` for real liveness
+        tracking.
+
+        During the exit barrier wait, workers have finished and the watchdog
+        progress time is stale. We return the current time to prevent TW
+        from killing the task while agents coordinate shutdown.
+        """
+        if self._in_exit_barrier:
+            return int(time.time())
+        if self._worker_watchdog is not None:
+            return self._worker_watchdog.get_last_progress_time()
+        return int(time.time())
+
     def _setup_healthcheck(self) -> None:
         healthcheck_port_env_name = TORCHELASTIC_HEALTH_CHECK_PORT
         healthcheck_port = os.getenv(healthcheck_port_env_name)
@@ -214,13 +264,28 @@ def _setup_healthcheck(self) -> None:
                 healthcheck_port_env_name,
                 healthcheck_port,
             )
-            if self._worker_watchdog is None:
-                logger.info(
-                    "FileTimerServer doesn't exist, using current time as dummy callback"
-                )
-                alive_callback = LocalElasticAgent._get_current_time_secs
+
+            if justknobs_check(
+                "ai_infra/pytorch_distributed:torchelastic_enable_healthcheck_before_rendezvous",
+                default=False,
+            ):
+                # New behavior: idempotent guard + dynamic callback that
+                # returns current time before watchdog exists and delegates
+                # to watchdog once workers are running.
+                if self._health_check_server is not None:
+                    return
+                alive_callback = self._get_alive_time
             else:
-                alive_callback = self._worker_watchdog.get_last_progress_time
+                # Original behavior: pick callback based on watchdog state
+                # at call time (only called from _start_workers where
+                # watchdog is already set up).
+                if self._worker_watchdog is None:
+                    logger.info(
+                        "FileTimerServer doesn't exist, using current time as dummy callback"
+                    )
+                    alive_callback = LocalElasticAgent._get_current_time_secs
+                else:
+                    alive_callback = self._worker_watchdog.get_last_progress_time
 
             try:
                 healthcheck_port_as_int = int(healthcheck_port)
@@ -292,7 +357,8 @@ def _stop_workers(self, worker_group: WorkerGroup) -> None:
     def _start_workers(self, worker_group: WorkerGroup) -> dict[int, Any]:
         spec = worker_group.spec
         store = worker_group.store
-        assert store is not None
+        if store is None:
+            raise AssertionError
         restart_count = spec.max_restarts - self._remaining_restarts
 
         use_agent_store: bool = spec.rdzv_handler.use_agent_store
@@ -348,8 +414,10 @@ def _start_workers(self, worker_group: WorkerGroup) -> dict[int, Any]:
         self._setup_local_watchdog(envs=envs)
         self._setup_healthcheck()
 
-        assert spec.entrypoint is not None
-        assert self._logs_specs is not None
+        if spec.entrypoint is None:
+            raise AssertionError
+        if self._logs_specs is None:
+            raise AssertionError
         self._pcontext = start_processes(
             name=spec.role,
             entrypoint=spec.entrypoint,
@@ -405,7 +473,9 @@ def _set_local_rank_env(
         if "CUDA_VISIBLE_DEVICES" in os.environ:
             worker_env["CUDA_VISIBLE_DEVICES"] = os.environ["CUDA_VISIBLE_DEVICES"]
 
-    def _shutdown(self, death_sig: signal.Signals = signal.SIGTERM) -> None:
+    def _shutdown(
+        self, death_sig: signal.Signals = signal.SIGTERM, timeout: int = 30
+    ) -> None:
         if self._worker_watchdog is not None:
             self._worker_watchdog.stop()
             self._worker_watchdog = None
@@ -413,7 +483,7 @@ def _shutdown(self, death_sig: signal.Signals = signal.SIGTERM) -> None:
             self._health_check_server.stop()
             self._health_check_server = None
         if self._pcontext:
-            self._pcontext.close(death_sig)
+            self._pcontext.close(death_sig, timeout)
 
     # pyre-fixme[56]: Pyre was not able to infer the type of the decorator
     #  `torch.distributed.elastic.metrics.prof`.
@@ -421,7 +491,8 @@ def _shutdown(self, death_sig: signal.Signals = signal.SIGTERM) -> None:
     def _monitor_workers(self, worker_group: WorkerGroup) -> RunResult:
         role = worker_group.spec.role
         worker_pids = {w.id for w in worker_group.workers}
-        assert self._pcontext is not None
+        if self._pcontext is None:
+            raise AssertionError
         pc_pids = set(self._pcontext.pids().values())
         if worker_pids != pc_pids:
             logger.error(
diff --git a/torch/distributed/elastic/events/api.py b/torch/distributed/elastic/events/api.py
index 31afe29ff5f59..9638f98b402c4 100644
--- a/torch/distributed/elastic/events/api.py
+++ b/torch/distributed/elastic/events/api.py
@@ -15,7 +15,7 @@
 
 __all__ = ["EventSource", "Event", "NodeState", "RdzvEvent"]
 
-EventMetadataValue = Union[str, int, float, bool, None]
+EventMetadataValue = str | int | float | bool | None
 
 
 class EventSource(str, Enum):
diff --git a/torch/distributed/elastic/multiprocessing/api.py b/torch/distributed/elastic/multiprocessing/api.py
index 44a671c795a2c..1ce0d78ed25d7 100644
--- a/torch/distributed/elastic/multiprocessing/api.py
+++ b/torch/distributed/elastic/multiprocessing/api.py
@@ -779,7 +779,8 @@ def _is_done(self) -> bool:
         return len(self._return_values) == self.nprocs
 
     def _poll(self) -> RunProcsResult | None:
-        assert self._pc is not None  # assertion for mypy type checker
+        if self._pc is None:
+            raise AssertionError  # assertion for mypy type checker
 
         try:
             # torch.mp.ProcessContext Throws an Exception if some/all of
@@ -858,7 +859,8 @@ def _poll(self) -> RunProcsResult | None:
             )
 
     def pids(self) -> dict[int, int]:
-        assert self._pc is not None  # assertion for mypy type checking
+        if self._pc is None:
+            raise AssertionError  # assertion for mypy type checking
         return dict(enumerate(self._pc.pids()))
 
     def _close(self, death_sig: signal.Signals, timeout: int = 30) -> None:
diff --git a/torch/distributed/elastic/multiprocessing/errors/__init__.py b/torch/distributed/elastic/multiprocessing/errors/__init__.py
index f61c99dc5c777..25a60fce6b9c7 100644
--- a/torch/distributed/elastic/multiprocessing/errors/__init__.py
+++ b/torch/distributed/elastic/multiprocessing/errors/__init__.py
@@ -185,7 +185,7 @@ def timestamp_isoformat(self):
   time      : ${time}
   host      : ${hostname}
   rank      : ${rank} (local_rank: ${local_rank})
-  exitcode  : ${exitcode} (pid: ${pid})
+  exitcode  : ${exitcode} (pid: ${pid}) ${signal_name}
   error_file: ${error_file}
   traceback : ${message}"""
 
@@ -240,9 +240,9 @@ class ChildFailedError(Exception):
     def __init__(self, name: str, failures: dict[GlobalRank, ProcessFailure]):
         self.name = name
         self.failures = failures
-        assert (
-            self.failures
-        )  # does not make sense to create a ChildFaileError with no failures
+        # does not make sense to create a ChildFaileError with no failures
+        if not self.failures:
+            raise AssertionError
         super().__init__(self.format_msg())
 
     def get_first_failure(self) -> tuple[GlobalRank, ProcessFailure]:
@@ -294,6 +294,9 @@ def _format_failure(
                 .replace("\n", "\n  ")  # to properly indent the traceback
             )
 
+        signal_name = failure.signal_name()
+        signal_name_str = f" ({signal_name})" if signal_name != _NOT_AVAILABLE else ""
+
         fmt = Template(_FAILURE_FORMAT_TEMPLATE).substitute(
             idx=idx,
             time=failure.timestamp_isoformat(),
@@ -302,6 +305,7 @@ def _format_failure(
             local_rank=failure.local_rank,
             exitcode=failure.exitcode,
             pid=failure.pid,
+            signal_name=signal_name_str,
             error_file=failure.error_file,
             message=msg,
         )
@@ -356,7 +360,8 @@ def main():
     def wrap(f: Callable[_P, _R]) -> Callable[_P, _R | None]:
         @wraps(f)
         def wrapper(*args: _P.args, **kwargs: _P.kwargs):
-            assert error_handler is not None  # assertion for mypy type checker
+            if error_handler is None:
+                raise AssertionError  # assertion for mypy type checker
             error_handler.initialize()
             try:
                 return f(*args, **kwargs)
diff --git a/torch/distributed/elastic/multiprocessing/redirects.py b/torch/distributed/elastic/multiprocessing/redirects.py
index 057013fbb9e5b..cecdec12ffee2 100644
--- a/torch/distributed/elastic/multiprocessing/redirects.py
+++ b/torch/distributed/elastic/multiprocessing/redirects.py
@@ -24,12 +24,28 @@
 logger = logging.getLogger(__name__)
 
 
+_WIN32_STD_HANDLE = {
+    "stdout": -11,  # STD_OUTPUT_HANDLE
+    "stderr": -12,  # STD_ERROR_HANDLE
+}
+
+
 def get_libc():
-    if IS_WINDOWS or IS_MACOS:
-        logger.warning(
-            "NOTE: Redirects are currently not supported in Windows or MacOs."
-        )
+    if IS_MACOS:
+        logger.warning("NOTE: Redirects are currently not supported in MacOs.")
         return None
+    elif IS_WINDOWS:
+        for lib_name in ("ucrtbase", "msvcrt", "msvcr110", "msvcr100"):
+            try:
+                lib = ctypes.CDLL(lib_name)
+                logger.debug("Loaded Windows C runtime: %s", lib_name)
+                return lib
+            except OSError:
+                continue
+        raise RuntimeError(
+            "Could not load a C runtime DLL on Windows (tried: ucrtbase, msvcrt, "
+            "msvcr110, msvcr100). Redirects cannot function without a CRT."
+        )
     else:
         return ctypes.CDLL("libc.so.6")
 
@@ -38,6 +54,24 @@ def get_libc():
 
 
 def _c_std(stream: str):
+    if IS_WINDOWS:
+        stream_index = 2 if stream == "stderr" else 1
+        try:
+            iob_func = libc.__acrt_iob_func
+            iob_func.restype = ctypes.POINTER(ctypes.c_void_p)
+            iob_func.argtypes = [ctypes.c_uint]
+            return iob_func(stream_index)
+        except AttributeError:
+            pass
+        try:
+            legacy_index = 2 if stream == "stderr" else 1
+            iob = (ctypes.POINTER(ctypes.c_void_p) * 3).in_dll(libc, "_iob")
+            return iob[legacy_index]
+        except (AttributeError, OSError) as err:
+            raise RuntimeError(
+                f"Could not resolve C-runtime FILE* for '{stream}'. "
+                "Neither __acrt_iob_func nor _iob are available in the loaded CRT."
+            ) from err
     return ctypes.c_void_p.in_dll(libc, stream)
 
 
@@ -48,56 +82,147 @@ def _python_std(stream: str):
 _VALID_STD = {"stdout", "stderr"}
 
 
-@contextmanager
-def redirect(std: str, to_file: str):
-    """
-    Redirect ``std`` (one of ``"stdout"`` or ``"stderr"``) to a file in the path specified by ``to_file``.
-
-    This method redirects the underlying std file descriptor (not just python's ``sys.stdout|stderr``).
-    See usage for details.
-
-    Directory of ``dst_filename`` is assumed to exist and the destination file
-    is overwritten if it already exists.
-
-    .. note:: Due to buffering cross source writes are not guaranteed to
-              appear in wall-clock order. For instance in the example below
-              it is possible for the C-outputs to appear before the python
-              outputs in the log file.
-
-    Usage:
-
-    ::
-
-     # syntactic-sugar for redirect("stdout", "tmp/stdout.log")
-     with redirect_stdout("/tmp/stdout.log"):
-        print("python stdouts are redirected")
-        libc = ctypes.CDLL("libc.so.6")
-        libc.printf(b"c stdouts are also redirected"
-        os.system("echo system stdouts are also redirected")
-
-     print("stdout restored")
-
-    """
-    if std not in _VALID_STD:
-        raise ValueError(
-            f"unknown standard stream <{std}>, must be one of {_VALID_STD}"
-        )
-
-    c_std = _c_std(std)
-    python_std = _python_std(std)
-    std_fd = python_std.fileno()
-
-    def _redirect(dst):
-        libc.fflush(c_std)
-        python_std.flush()
-        os.dup2(dst.fileno(), std_fd)
-
-    with os.fdopen(os.dup(std_fd)) as orig_std, open(to_file, mode="w+b") as dst:
-        _redirect(dst)
-        try:
-            yield
-        finally:
-            _redirect(orig_std)
+if IS_WINDOWS:  # libc is None on macOS; all of the below is Windows-only
+    import io as _io
+    import msvcrt as _msvcrt
+
+    _kernel32 = ctypes.WinDLL("kernel32", use_last_error=True)  # type: ignore[attr-defined]
+
+    _crt_dup = libc._dup
+    _crt_dup2 = libc._dup2
+    _crt_dup.restype = ctypes.c_int
+    _crt_dup.argtypes = [ctypes.c_int]
+    _crt_dup2.restype = ctypes.c_int
+    _crt_dup2.argtypes = [ctypes.c_int, ctypes.c_int]
+
+    @contextmanager
+    def redirect(std: str, to_file: str):
+        """
+        Redirect ``std`` (one of ``"stdout"`` or ``"stderr"``) to a file at ``to_file``.
+
+        On Windows this performs a four-layer redirect:
+
+        1. ``sys.stdout``/``sys.stderr`` -- rewired to a new TextIOWrapper so
+           Python's ``print()`` writes to the destination file.
+        2. CRT fd (``_dup2``) -- captures C ``printf`` and UCRT ``FILE*`` writers.
+        3. Win32 ``SetStdHandle`` -- captures native code using ``WriteFile``/
+           ``WriteConsole`` directly, including HIP/ROCm.
+        4. ``fflush`` before each switch -- prevents lost output from CRT buffering.
+
+        .. note:: If ROCm/HIP caches the Win32 HANDLE before this redirect runs
+                  (e.g. at ``import torch`` time), set up the redirect *before*
+                  importing torch/ROCm to capture all output.
+
+        Directory of ``to_file`` is assumed to exist. The destination file is
+        overwritten if it already exists.
+        """
+        if std not in _VALID_STD:
+            raise ValueError(
+                f"unknown standard stream <{std}>, must be one of {_VALID_STD}"
+            )
+
+        std_fd = 1 if std == "stdout" else 2
+        win32_handle_id = _WIN32_STD_HANDLE[std]
+        orig_sys_std = getattr(sys, std)
+        orig_fd_dup = _crt_dup(std_fd)
+        if orig_fd_dup == -1:
+            raise OSError(f"CRT _dup failed for {std} (fd={std_fd})")
+        orig_win32_handle = _kernel32.GetStdHandle(win32_handle_id)
+
+        with open(to_file, mode="w+b") as dst:
+            dst_fd = dst.fileno()
+
+            try:
+                libc.fflush(_c_std(std))
+            except Exception:
+                pass
+            try:
+                orig_sys_std.flush()
+            except Exception:
+                pass
+
+            _kernel32.SetStdHandle(
+                win32_handle_id,
+                _msvcrt.get_osfhandle(dst_fd),  # pyrefly: ignore [missing-attribute]
+            )
+
+            if _crt_dup2(dst_fd, std_fd) == -1:
+                raise OSError(f"CRT _dup2 failed redirecting {std}")
+
+            new_sys_std = _io.TextIOWrapper(
+                open(dst_fd, mode="wb", closefd=False),  # noqa: SIM115
+                encoding=orig_sys_std.encoding or "utf-8",
+                errors="replace",
+                line_buffering=True,
+            )
+            setattr(sys, std, new_sys_std)
+
+            try:
+                yield
+            finally:
+                try:
+                    new_sys_std.flush()
+                except Exception:
+                    pass
+                try:
+                    libc.fflush(_c_std(std))
+                except Exception:
+                    pass
+
+                setattr(sys, std, orig_sys_std)
+                _crt_dup2(orig_fd_dup, std_fd)
+                os.close(orig_fd_dup)
+                _kernel32.SetStdHandle(win32_handle_id, orig_win32_handle)
+
+else:
+
+    @contextmanager
+    def redirect(std: str, to_file: str):
+        """
+        Redirect ``std`` (one of ``"stdout"`` or ``"stderr"``) to a file in the path specified by ``to_file``.
+
+        This method redirects the underlying std file descriptor (not just python's ``sys.stdout|stderr``).
+        See usage for details.
+
+        Directory of ``dst_filename`` is assumed to exist and the destination file
+        is overwritten if it already exists.
+
+        .. note:: Due to buffering cross source writes are not guaranteed to
+                  appear in wall-clock order. For instance in the example below
+                  it is possible for the C-outputs to appear before the python
+                  outputs in the log file.
+
+        Usage::
+
+            # syntactic-sugar for redirect("stdout", "tmp/stdout.log")
+            with redirect_stdout("/tmp/stdout.log"):
+                print("python stdouts are redirected")
+                libc = ctypes.CDLL("libc.so.6")
+                libc.printf(b"c stdouts are also redirected")
+                os.system("echo system stdouts are also redirected")
+
+            print("stdout restored")
+        """
+        if std not in _VALID_STD:
+            raise ValueError(
+                f"unknown standard stream <{std}>, must be one of {_VALID_STD}"
+            )
+
+        c_std = _c_std(std)
+        python_std = _python_std(std)
+        std_fd = python_std.fileno()
+
+        def _redirect(dst):
+            libc.fflush(c_std)
+            python_std.flush()
+            os.dup2(dst.fileno(), std_fd)
+
+        with os.fdopen(os.dup(std_fd)) as orig_std, open(to_file, mode="w+b") as dst:
+            _redirect(dst)
+            try:
+                yield
+            finally:
+                _redirect(orig_std)
 
 
 redirect_stdout = partial(redirect, "stdout")
diff --git a/torch/distributed/elastic/multiprocessing/tail_log.py b/torch/distributed/elastic/multiprocessing/tail_log.py
index 77d410cce55c0..d9d2edb81d688 100644
--- a/torch/distributed/elastic/multiprocessing/tail_log.py
+++ b/torch/distributed/elastic/multiprocessing/tail_log.py
@@ -105,7 +105,6 @@ def __init__(
         n = len(log_files)
         self._threadpool = None
         if n > 0:
-            # pyrefly: ignore [bad-assignment]
             self._threadpool = ThreadPoolExecutor(
                 max_workers=n,
                 thread_name_prefix=f"{self.__class__.__qualname__}_{name}",
diff --git a/torch/distributed/elastic/rendezvous/dynamic_rendezvous.py b/torch/distributed/elastic/rendezvous/dynamic_rendezvous.py
index 84adeea955731..312cd191cbc2b 100644
--- a/torch/distributed/elastic/rendezvous/dynamic_rendezvous.py
+++ b/torch/distributed/elastic/rendezvous/dynamic_rendezvous.py
@@ -1219,9 +1219,11 @@ def next_rendezvous(self) -> RendezvousInfo:
                 server_port=server_port,  # For non-0 rank, this is a no-op
             )
 
-        assert self._bootstrap_store_info is not None
+        if self._bootstrap_store_info is None:
+            raise AssertionError
         if rank == 0:
-            assert self._shared_tcp_store_server is not None
+            if self._shared_tcp_store_server is None:
+                raise AssertionError
 
         return RendezvousInfo(
             store,
diff --git a/torch/distributed/elastic/rendezvous/etcd_rendezvous.py b/torch/distributed/elastic/rendezvous/etcd_rendezvous.py
index 93a7073bed87a..96d82e03f9cf0 100644
--- a/torch/distributed/elastic/rendezvous/etcd_rendezvous.py
+++ b/torch/distributed/elastic/rendezvous/etcd_rendezvous.py
@@ -412,9 +412,8 @@ def join_phase(self, expected_version):
         active_version = self.wait_for_peers(expected_version)
         state = json.loads(active_version.value)
 
-        assert state["version"] == expected_version, (
-            "Logic error: failed to observe version mismatch"
-        )
+        if state["version"] != expected_version:
+            raise AssertionError("Logic error: failed to observe version mismatch")
 
         return self.confirm_phase(expected_version, this_rank)
 
@@ -532,9 +531,10 @@ def join_rendezvous(self, expected_version):
                     "Rendezvous version changed. Must try join the new one."
                 )
 
-            assert len(state["participants"]) < self._num_max_workers, (
-                "Logic error: joinable rendezvous should always have space left"
-            )
+            if len(state["participants"]) >= self._num_max_workers:
+                raise AssertionError(
+                    "Logic error: joinable rendezvous should always have space left"
+                )
 
             this_rank = len(state["participants"])
             state["participants"].append(this_rank)
@@ -954,7 +954,8 @@ def load_extra_data(self, rdzv_version, key, timeout=None):
 
             # Find the extra_data node, if it exists
             extra_data = [n for n in root.children if n.key == node]
-            assert len(extra_data) <= 1
+            if len(extra_data) > 1:
+                raise AssertionError
 
             # Node for extra_data exists, check the desired key inside it.
             if len(extra_data) == 1:
diff --git a/torch/distributed/elastic/rendezvous/etcd_server.py b/torch/distributed/elastic/rendezvous/etcd_server.py
index 347e7339d9a46..c3a6b2803d85d 100644
--- a/torch/distributed/elastic/rendezvous/etcd_server.py
+++ b/torch/distributed/elastic/rendezvous/etcd_server.py
@@ -177,7 +177,7 @@ def start(
                 curr_retries += 1
                 stop_etcd(self._etcd_proc)
                 logger.warning(  # noqa: G200
-                    "Failed to start etcd server, got error: %s, retrying", str(e)
+                    "Failed to start etcd server, got error: %s, retrying", e
                 )
                 if curr_retries >= num_retries:
                     shutil.rmtree(self._base_data_dir, ignore_errors=True)
diff --git a/torch/distributed/elastic/rendezvous/etcd_store.py b/torch/distributed/elastic/rendezvous/etcd_store.py
index faaf77587bc9d..09e588ff5b709 100644
--- a/torch/distributed/elastic/rendezvous/etcd_store.py
+++ b/torch/distributed/elastic/rendezvous/etcd_store.py
@@ -120,6 +120,7 @@ def add(self, key, num: int) -> int:
             except etcd.EtcdCompareFailed:
                 cas_delay()
 
+    # pyrefly: ignore [bad-override]
     def wait(self, keys, override_timeout: datetime.timedelta | None = None):
         """
         Wait until all of the keys are published, or until timeout.
diff --git a/torch/distributed/elastic/utils/data/elastic_distributed_sampler.py b/torch/distributed/elastic/utils/data/elastic_distributed_sampler.py
index 6dd024b32ecaf..4889ce0b07bf3 100644
--- a/torch/distributed/elastic/utils/data/elastic_distributed_sampler.py
+++ b/torch/distributed/elastic/utils/data/elastic_distributed_sampler.py
@@ -81,11 +81,13 @@ def __iter__(self) -> Iterator[T]:
 
         # add extra samples to make it evenly divisible
         indices += indices[: (self.total_size - len(indices))]
-        assert len(indices) == self.total_size
+        if len(indices) != self.total_size:
+            raise AssertionError
 
         # subsample
         indices = indices[self.rank : self.total_size : self.num_replicas]
-        assert len(indices) == self.num_samples
+        if len(indices) != self.num_samples:
+            raise AssertionError
 
         return iter(indices)
 
diff --git a/torch/distributed/elastic/utils/logging.py b/torch/distributed/elastic/utils/logging.py
index aadf37eb16b80..dbc0979d4f547 100644
--- a/torch/distributed/elastic/utils/logging.py
+++ b/torch/distributed/elastic/utils/logging.py
@@ -46,7 +46,8 @@ def _derive_module_name(depth: int = 1) -> str | None:
     """
     try:
         stack = inspect.stack()
-        assert depth < len(stack)
+        if depth >= len(stack):
+            raise AssertionError
         # FrameInfo is just a named tuple: (frame, filename, lineno, function, code_context, index)
         frame_info = stack[depth]
 
diff --git a/torch/distributed/elastic/utils/store.py b/torch/distributed/elastic/utils/store.py
index 598899e936aa0..00f37f0b12c2b 100644
--- a/torch/distributed/elastic/utils/store.py
+++ b/torch/distributed/elastic/utils/store.py
@@ -190,7 +190,8 @@ def barrier(
     """
 
     if rank is None:
-        assert rank_tracing_decoder is None, "Tracing requires rank information"
+        if rank_tracing_decoder is not None:
+            raise AssertionError("Tracing requires rank information")
 
     with store_timeout(store, barrier_timeout):
         last_member_key = _barrier_nonblocking(
diff --git a/torch/distributed/flight_recorder/components/builder.py b/torch/distributed/flight_recorder/components/builder.py
index 56736450e3f2a..663363e718257 100644
--- a/torch/distributed/flight_recorder/components/builder.py
+++ b/torch/distributed/flight_recorder/components/builder.py
@@ -126,12 +126,15 @@ def build_groups_memberships(
                 _memberships[pg_guid] = set(ranks)
             else:
                 # validation across ranks
-                assert _groups[pg_guid].desc == desc, (
-                    f"mismatch in desc {_groups[pg_guid].desc} vs {desc} for group {pg_guid}"
-                )
-                assert _memberships[pg_guid] == set(ranks), (
-                    f"mismatch in membership for group {pg_guid} {_memberships[pg_guid]} vs {set(ranks)}"
-                )
+                if _groups[pg_guid].desc != desc:
+                    raise AssertionError(
+                        f"mismatch in desc {_groups[pg_guid].desc} vs {desc} for group {pg_guid}"
+                    )
+                if _memberships[pg_guid] != set(ranks):
+                    raise AssertionError(
+                        f"mismatch in membership for group {pg_guid}"
+                        f" {_memberships[pg_guid]} vs {set(ranks)}"
+                    )
     return groups, _groups, memberships, _memberships, _pg_guids
 
 
@@ -248,14 +251,16 @@ def build_collectives(
                     op = Op(entry, _memberships, pg_name)
                     peer = None
                     if op.type == "send":
-                        assert op._src_g == curr, (
-                            f"Send src error: {curr} expected but {op._src_g} is set"
-                        )
+                        if op._src_g != curr:
+                            raise AssertionError(
+                                f"Send src error: {curr} expected but {op._src_g} is set"
+                            )
                         peer = op._dst_g
                     elif op.type == "recv":
-                        assert op._dst_g == curr, (
-                            f"Recv dst error: {curr} expected but {op._dst_g} is set"
-                        )
+                        if op._dst_g != curr:
+                            raise AssertionError(
+                                f"Recv dst error: {curr} expected but {op._dst_g} is set"
+                            )
                         peer = op._src_g
                     if peer and peer not in done_ranks:
                         candidate_ranks.add(peer)
@@ -306,7 +311,12 @@ def build_collectives(
                 # This extra cleanup is needed because we need to pop all collectives within a coalesced collective.
                 for i, k in idx_map.items():
                     for _ in range(1, num_coalesced_entries):
-                        all_entries[i].pop(k)
+                        try:
+                            all_entries[i].pop(k)
+                        except IndexError:
+                            # In the case of a missing rank symptom that a rank didn't schedule the coalesced collective,
+                            # we should not fail the analysis script here.
+                            pass
         else:
             # Iterate through all the ranks and check if there is a mismatch for the current entry.
             check_current_entry_match(
diff --git a/torch/distributed/flight_recorder/components/config_manager.py b/torch/distributed/flight_recorder/components/config_manager.py
index d1b1296658821..60de805f20808 100644
--- a/torch/distributed/flight_recorder/components/config_manager.py
+++ b/torch/distributed/flight_recorder/components/config_manager.py
@@ -94,15 +94,17 @@ def parse_args(self: "JobConfig", args: Sequence[str] | None) -> argparse.Namesp
         # pyrefly: ignore [missing-attribute]
         if args.selected_ranks is not None:
             # pyrefly: ignore [missing-attribute]
-            assert args.just_print_entries, (
-                "Not support selecting ranks without printing entries"
-            )
+            if not args.just_print_entries:
+                raise AssertionError(
+                    "Not support selecting ranks without printing entries"
+                )
         # pyrefly: ignore [missing-attribute]
         if args.pg_filters is not None:
             # pyrefly: ignore [missing-attribute]
-            assert args.just_print_entries, (
-                "Not support selecting pg filters without printing entries"
-            )
+            if not args.just_print_entries:
+                raise AssertionError(
+                    "Not support selecting pg filters without printing entries"
+                )
         # pyrefly: ignore [missing-attribute]
         if args.verbose:
             logger.set_log_level(logging.DEBUG)
diff --git a/torch/distributed/flight_recorder/components/loader.py b/torch/distributed/flight_recorder/components/loader.py
index ce361b103fe04..3e09353eea273 100644
--- a/torch/distributed/flight_recorder/components/loader.py
+++ b/torch/distributed/flight_recorder/components/loader.py
@@ -79,7 +79,8 @@ def read_dir(args: argparse.Namespace) -> tuple[dict[str, dict[str, Any]], str]:
     t0 = time.time()
     version = ""
     filecount = 0
-    assert os.path.isdir(args.trace_dir), f"folder {args.trace_dir} does not exist"
+    if not os.path.isdir(args.trace_dir):
+        raise AssertionError(f"folder {args.trace_dir} does not exist")
     for root, _, files in os.walk(args.trace_dir):
         if prefix is None:
             prefix = _determine_prefix(files)
@@ -91,8 +92,9 @@ def read_dir(args: argparse.Namespace) -> tuple[dict[str, dict[str, Any]], str]:
             if not version:
                 version = str(details[f]["version"])
     tb = time.time()
-    assert len(details) > 0, (
-        f"no files loaded from {args.trace_dir} with prefix {prefix}"
-    )
+    if len(details) <= 0:
+        raise AssertionError(
+            f"no files loaded from {args.trace_dir} with prefix {prefix}"
+        )
     logger.debug("loaded %s files in %ss", filecount, tb - t0)
     return details, version
diff --git a/torch/distributed/flight_recorder/components/types.py b/torch/distributed/flight_recorder/components/types.py
index e87fd5f4569d4..46ebda50dd506 100644
--- a/torch/distributed/flight_recorder/components/types.py
+++ b/torch/distributed/flight_recorder/components/types.py
@@ -209,16 +209,25 @@ class Database(NamedTuple):
     "reduce",
     "_reduce_oop",
     "all_gather",
+    "all_gather_single",
+    "all_gather_v",
     "all_reduce",
     "_all_gather_base",
     "all_gather_into_tensor_coalesced",
     "reduce_scatter",
+    "reduce_scatter_single",
+    "reduce_scatter_v",
     "reduce_scatter_tensor_coalesced",
     "_reduce_scatter_base",
     "gather",
     "scatter",
     "all_to_all",
+    "all_to_all_single",
+    "all_to_all_v_single",
     "all_reduce_barrier",
+    "barrier",
+    "split",
+    "new_window",
     "allreduce_coalesced",
     "ALLGATHER_coalesced",
     "REDUCE_SCATTER_coalesced",
@@ -283,7 +292,7 @@ def log(
         logger.info("input sizes: %s", self.input_sizes)
         logger.info("output sizes: %s", self.output_sizes)
         logger.info("world size: %d", len(self.expected_ranks))
-        logger.info("expected ranks: %s", str(self.expected_ranks))
+        logger.info("expected ranks: %s", self.expected_ranks)
         logger.info("collective state: %s", self.collective_state)
         if errors:
             self.errors = errors
@@ -320,8 +329,10 @@ def to_collective(
                 missing_ranks=getattr(self, "missing_ranks", None),
             )
         else:
-            assert idx_map is not None, "idx_map is None"
-            assert all_entries is not None, "all_entries is None"
+            if idx_map is None:
+                raise AssertionError("idx_map is None")
+            if all_entries is None:
+                raise AssertionError("all_entries is None")
             mismatch_collectives = {}
             for rank, error in errors:
                 idx = idx_map[rank]
@@ -407,9 +418,10 @@ def __init__(
     ):
         self.profiling_name = event["profiling_name"]
         comm_lib_backend, name = self.profiling_name.split(":")
-        assert comm_lib_backend in ["nccl", "xccl"], (
-            f"name formatting error? {comm_lib_backend} != 'nccl' or 'xccl'"
-        )
+        if comm_lib_backend not in ["nccl", "ncclx", "gloo", "xccl"]:
+            raise AssertionError(
+                f"name formatting error? {comm_lib_backend} not in supported backends"
+            )
         parts = name.split(" ")
         type = parts[0]
         meta = parts[1] if len(parts) == 2 else None
@@ -417,16 +429,17 @@ def __init__(
         # Store the hashed pg_name for accessing memberships, and original pg info for display
         self.pg_name = pg_name  # This is the hashed version used for memberships lookup
         self.original_pg_name, self.pg_desc = event["process_group"]
-        assert type in COLLECTIVES | P2P | {"coalesced"}, (
-            f"{type} is not a supported operation"
-        )
+        if type not in COLLECTIVES | P2P | {"coalesced"}:
+            raise AssertionError(f"{type} is not a supported operation")
         self.type = type
         if type == "send":
-            assert isinstance(meta, str)
+            if not isinstance(meta, str):
+                raise AssertionError
             s, d = meta.split("->")
             self._src, self._dst = int(s), int(d)
         elif type == "recv":
-            assert isinstance(meta, str)
+            if not isinstance(meta, str):
+                raise AssertionError
             d, s = meta.split("<-")
             self._dst, self._src = int(d), int(s)
         else:
@@ -454,12 +467,14 @@ def _init_global_src_dst(self, pg_ranks: set[Any]) -> None:
 
     @property
     def src(self) -> int:
-        assert self.type in P2P, "can't get src of non-p2p op"
+        if self.type not in P2P:
+            raise AssertionError("can't get src of non-p2p op")
         return self._src
 
     @property
     def dst(self) -> int:
-        assert self.type in P2P, "can't get dst of non-p2p op"
+        if self.type not in P2P:
+            raise AssertionError("can't get dst of non-p2p op")
         return self._dst
 
     def __repr__(self) -> str:
diff --git a/torch/distributed/flight_recorder/components/utils.py b/torch/distributed/flight_recorder/components/utils.py
index 8fcd14667c202..7700ac03155e9 100644
--- a/torch/distributed/flight_recorder/components/utils.py
+++ b/torch/distributed/flight_recorder/components/utils.py
@@ -60,7 +60,6 @@ def format_frame(frame: dict[str, str]) -> str:
 def format_frames(frames: list[dict[str, str]]) -> str:
     formatted_frames = []
     for frame in frames:
-        # pyrefly: ignore [bad-argument-type]
         formatted_frames.append(format_frame(frame))
     return "\n".join(formatted_frames)
 
@@ -163,7 +162,8 @@ def visualize_ops(
         if not op_list:
             # print("TODO- not sure if its valid for only some ranks in a PG to participate in a coalesced op?")
             return False
-        assert op_list[-1].type == "coalesced"
+        if op_list[-1].type != "coalesced":
+            raise AssertionError
         op_list.pop(-1)
 
     while all_ops:
@@ -602,7 +602,8 @@ def find_coalesced_group(
             break
 
     if len(found) > 1:
-        assert found[-1][1]["profiling_name"] == "nccl:coalesced"
+        if found[-1][1]["profiling_name"] != "nccl:coalesced":
+            raise AssertionError
         return found
     return []
 
@@ -701,20 +702,23 @@ def check_no_missing_dump_files(
     all_ranks = {int(membership.global_rank) for membership in memberships}
     dumps_ranks = {int(key) for key in entries}
     missing = all_ranks - dumps_ranks
-    assert len(missing) == 0, f"Missing dump files from ranks {missing}"
+    if len(missing) != 0:
+        raise AssertionError(f"Missing dump files from ranks {missing}")
 
 
 def check_version(version_by_ranks: dict[str, str], version: str) -> None:
     for rank, v in version_by_ranks.items():
-        assert v == version, (
-            f"Rank {rank} has different version {v} from the given version {version}"
-        )
+        if v != version:
+            raise AssertionError(
+                f"Rank {rank} has different version {v} from the given version {version}"
+            )
 
 
 def get_version_detail(version: str) -> tuple[int, int]:
     # pyrefly: ignore [bad-assignment]
     version = version.split(".")
-    assert len(version) == 2, f"Invalid version {version}"
+    if len(version) != 2:
+        raise AssertionError(f"Invalid version {version}")
     major, minor = map(int, version)
     return major, minor
 
diff --git a/torch/distributed/flight_recorder/fr_trace.py b/torch/distributed/flight_recorder/fr_trace.py
index ab338d1503ae0..29bc47cf13b4d 100644
--- a/torch/distributed/flight_recorder/fr_trace.py
+++ b/torch/distributed/flight_recorder/fr_trace.py
@@ -45,13 +45,15 @@ def main(args: Sequence[str] | None = None) -> None:
     # pyrefly: ignore [bad-assignment]
     args = config.parse_args(args)
     # pyrefly: ignore [missing-attribute]
-    assert args.trace_dir, "Trace directory trace_dir is required"
+    if not args.trace_dir:
+        raise AssertionError("Trace directory trace_dir is required")
     # pyrefly: ignore [bad-argument-type]
     details, version = read_dir(args)
     # pyrefly: ignore [missing-attribute]
     if args.transform_ft:
         # pyrefly: ignore [missing-attribute]
-        assert args.group_world_size, "World size is required for transform_ft"
+        if not args.group_world_size:
+            raise AssertionError("World size is required for transform_ft")
         # pyrefly: ignore [bad-argument-type]
         details = transform_ft(details, args.group_world_size)
     # pyrefly: ignore [bad-argument-type]
diff --git a/torch/distributed/fsdp/__init__.py b/torch/distributed/fsdp/__init__.py
index 1e4219250c39d..7c0d55e37d688 100644
--- a/torch/distributed/fsdp/__init__.py
+++ b/torch/distributed/fsdp/__init__.py
@@ -1,6 +1,7 @@
 from ._flat_param import FlatParameter as FlatParameter
 from ._fully_shard import (
     CPUOffloadPolicy,
+    DataParallelMeshDims,
     FSDPModule,
     fully_shard,
     MixedPrecisionPolicy,
@@ -49,6 +50,7 @@
     "StateDictType",
     # FSDP2
     "CPUOffloadPolicy",
+    "DataParallelMeshDims",
     "FSDPModule",
     "fully_shard",
     "MixedPrecisionPolicy",
@@ -60,6 +62,7 @@
 
 # Set namespace for exposed private names
 CPUOffloadPolicy.__module__ = "torch.distributed.fsdp"
+DataParallelMeshDims.__module__ = "torch.distributed.fsdp"
 FSDPModule.__module__ = "torch.distributed.fsdp"
 fully_shard.__module__ = "torch.distributed.fsdp"
 MixedPrecisionPolicy.__module__ = "torch.distributed.fsdp"
diff --git a/torch/distributed/fsdp/_common_utils.py b/torch/distributed/fsdp/_common_utils.py
index fba5ef316711f..ad6d543fc97f5 100644
--- a/torch/distributed/fsdp/_common_utils.py
+++ b/torch/distributed/fsdp/_common_utils.py
@@ -3,11 +3,12 @@
 This file includes private common utilities for FSDP.
 """
 
+import dataclasses
 import logging
 import traceback
 import warnings
 import weakref
-from collections.abc import Callable, Generator, Iterable
+from collections.abc import Callable, Generator, Iterable, Iterator
 from enum import auto, Enum
 from functools import partial
 from itertools import chain
@@ -40,6 +41,165 @@
 
     from ._flat_param import FlatParamHandle
 
+
+_MAX_TRAVERSE_DEPTH = 128
+
+
+def _is_namedtuple(obj: Any) -> bool:
+    # Mirrors torch.nn.parallel.scatter_gather._is_namedtuple
+    fields = getattr(type(obj), "_fields", None)
+    return (
+        isinstance(obj, tuple)
+        and hasattr(obj, "_asdict")
+        and isinstance(fields, tuple)
+        and all(isinstance(f, str) for f in fields)
+    )
+
+
+def collect_grad_tensors(output: Any) -> tuple[torch.Tensor, ...]:
+    """
+    Recursively collect tensors that require gradients from a nested structure.
+
+    Traverses dict, list, tuple, NamedTuple, and dataclass containers.
+    Sets and other iterables are *not* traversed (consistent with
+    ``tree_flatten``).  Uses the same traversal order as
+    :func:`replace_grad_tensors`.
+    """
+    tensors_list: list[torch.Tensor] = []
+    _collect_grad_tensors(output, tensors_list)
+    return tuple(tensors_list)
+
+
+def _collect_grad_tensors(
+    output: Any, out: list[torch.Tensor], _depth: int = 0
+) -> None:
+    """Collect grad-requiring tensors in the same order as _replace_grad_tensors."""
+    if _depth >= _MAX_TRAVERSE_DEPTH:
+        raise RuntimeError(
+            f"collect_grad_tensors exceeded max depth ({_MAX_TRAVERSE_DEPTH}), "
+            "likely due to a circular reference in the output structure"
+        )
+    # Branch order must mirror _replace_grad_tensors exactly.
+    # Only dict, list, tuple, NamedTuple, and dataclass are traversed;
+    # set and other iterables are intentionally skipped (matching tree_flatten).
+    if torch.is_tensor(output) and output.requires_grad:
+        out.append(output)
+    elif _is_namedtuple(output):
+        # NamedTuple before dataclass to match _replace_grad_tensors ordering.
+        for item in output:
+            _collect_grad_tensors(item, out, _depth + 1)
+    elif dataclasses.is_dataclass(output) and not isinstance(output, type):
+        for field in dataclasses.fields(output):
+            _collect_grad_tensors(getattr(output, field.name), out, _depth + 1)
+    elif isinstance(output, dict):
+        for v in output.values():
+            _collect_grad_tensors(v, out, _depth + 1)
+    elif isinstance(output, (list, tuple)):
+        for item in output:
+            _collect_grad_tensors(item, out, _depth + 1)
+
+
+def replace_grad_tensors(output: Any, tensor_iter: Iterator[torch.Tensor]) -> Any:
+    """
+    Replace grad-requiring tensors in a nested structure using replacements
+    from tensor_iter.
+
+    Tensors are consumed from tensor_iter in the same traversal order as
+    :func:`collect_grad_tensors`. Traverses dict, list, tuple, NamedTuple,
+    and dataclass containers; sets and other iterables are *not* traversed
+    (consistent with ``tree_flatten``).
+
+    Note: dataclass reconstruction uses ``dataclasses.replace()``, which calls
+    ``__init__``. Dataclasses with custom ``__init__`` validation,
+    ``__post_init__`` side effects, or non-standard dict subclass constructors
+    may not be compatible. In practice, FSDP module outputs are expected to be
+    shallowly nested, so recursion depth is not a concern.
+    """
+    result = _replace_grad_tensors(output, tensor_iter)
+    sentinel = object()
+    leftover = next(tensor_iter, sentinel)
+    if leftover is not sentinel:
+        # Count remaining without holding references to all of them
+        n = 1 + sum(1 for _ in tensor_iter)
+        raise RuntimeError(
+            f"{n} replacement tensors were not consumed while processing "
+            f"{type(output).__qualname__}"
+        )
+    return result
+
+
+def _replace_grad_tensors(
+    output: Any, tensor_iter: Iterator[torch.Tensor], _depth: int = 0
+) -> Any:
+    # Branch order must mirror _collect_grad_tensors exactly.
+    if _depth >= _MAX_TRAVERSE_DEPTH:
+        raise RuntimeError(
+            f"replace_grad_tensors exceeded max depth ({_MAX_TRAVERSE_DEPTH}), "
+            "likely due to a circular reference in the output structure"
+        )
+    if torch.is_tensor(output) and output.requires_grad:
+        return next(tensor_iter)
+    elif _is_namedtuple(output):
+        # NamedTuple before dataclass: a NamedTuple that is also a dataclass
+        # should be reconstructed via positional args, not dataclasses.replace.
+        new_items = []
+        any_changed = False
+        for item in output:
+            new_item = _replace_grad_tensors(item, tensor_iter, _depth + 1)
+            new_items.append(new_item)
+            if new_item is not item:
+                any_changed = True
+        if any_changed:
+            return type(output)(*new_items)
+        return output
+    elif dataclasses.is_dataclass(output) and not isinstance(output, type):
+        changes = {}
+        for field in dataclasses.fields(output):
+            old_val = getattr(output, field.name)
+            new_val = _replace_grad_tensors(old_val, tensor_iter, _depth + 1)
+            if new_val is not old_val:
+                changes[field.name] = new_val
+        if changes:
+            try:
+                return dataclasses.replace(output, **changes)
+            except TypeError as e:
+                raise TypeError(
+                    f"Failed to reconstruct dataclass {type(output).__qualname__} "
+                    f"via dataclasses.replace(). Dataclasses used as FSDP module "
+                    f"inputs/outputs must support dataclasses.replace(): {e}"
+                ) from None
+        return output
+    elif isinstance(output, dict):
+        new_dict = {}
+        any_changed = False
+        for k, v in output.items():
+            new_v = _replace_grad_tensors(v, tensor_iter, _depth + 1)
+            new_dict[k] = new_v
+            if new_v is not v:
+                any_changed = True
+        if any_changed:
+            return new_dict if type(output) is dict else type(output)(new_dict)
+        return output
+    elif isinstance(output, (list, tuple)):
+        new_items = []
+        any_changed = False
+        for item in output:
+            new_item = _replace_grad_tensors(item, tensor_iter, _depth + 1)
+            new_items.append(new_item)
+            if new_item is not item:
+                any_changed = True
+        if any_changed:
+            typ = type(output)
+            try:
+                return typ(new_items)
+            except TypeError:
+                # Fall back to base type for subclasses with custom __init__
+                return list(new_items) if isinstance(output, list) else tuple(new_items)
+        return output
+    else:
+        return output
+
+
 FSDP_WRAPPED_MODULE = "_fsdp_wrapped_module"
 FSDP_PREFIX = FSDP_WRAPPED_MODULE + "."
 FSDP_FLATTENED = "_fsdp_flattened"
@@ -65,7 +225,6 @@ def __init__(self, device: torch.device, backend: Any = None):
         if backend is None:
             try:
                 self.__backend = getattr(torch, device.type)
-                # pyrefly: ignore [read-only]
                 self.__device = device
             except AttributeError as exc:
                 raise AttributeError(
diff --git a/torch/distributed/fsdp/_flat_param.py b/torch/distributed/fsdp/_flat_param.py
index c56414268e774..190e2a0a88a34 100644
--- a/torch/distributed/fsdp/_flat_param.py
+++ b/torch/distributed/fsdp/_flat_param.py
@@ -565,7 +565,6 @@ def __init__(
         # Only align addresses for `use_orig_params=True` (for now)
         align_addresses = use_orig_params
         self._init_get_unflat_views_fn(align_addresses)
-        # pyrefly: ignore [read-only]
         self.device = device
         self._device_handle = _FSDPDeviceHandle.from_device(self.device)
         self.process_group = process_group
diff --git a/torch/distributed/fsdp/_fully_shard/__init__.py b/torch/distributed/fsdp/_fully_shard/__init__.py
index d4d0b341a3f82..f6aab00fc32b7 100644
--- a/torch/distributed/fsdp/_fully_shard/__init__.py
+++ b/torch/distributed/fsdp/_fully_shard/__init__.py
@@ -1,4 +1,9 @@
-from ._fsdp_api import CPUOffloadPolicy, MixedPrecisionPolicy, OffloadPolicy
+from ._fsdp_api import (
+    CPUOffloadPolicy,
+    DataParallelMeshDims,
+    MixedPrecisionPolicy,
+    OffloadPolicy,
+)
 from ._fully_shard import (
     FSDPModule,
     fully_shard,
@@ -10,6 +15,7 @@
 
 __all__ = [
     "CPUOffloadPolicy",
+    "DataParallelMeshDims",
     "FSDPModule",
     "fully_shard",
     "MixedPrecisionPolicy",
diff --git a/torch/distributed/fsdp/_fully_shard/_fsdp_api.py b/torch/distributed/fsdp/_fully_shard/_fsdp_api.py
index 1d60f7be8c60d..3f03d5707a9ba 100644
--- a/torch/distributed/fsdp/_fully_shard/_fsdp_api.py
+++ b/torch/distributed/fsdp/_fully_shard/_fsdp_api.py
@@ -2,13 +2,12 @@
 from abc import ABC, abstractmethod
 from collections.abc import Sequence
 from dataclasses import dataclass
-from typing import Union
 
 import torch
 import torch.distributed as dist
 
 
-_ReduceOp = Union[dist.ReduceOp, dist.ReduceOp.RedOpType]
+_ReduceOp = dist.ReduceOp | dist.ReduceOp.RedOpType
 
 
 @dataclass(frozen=True)
@@ -127,6 +126,49 @@ def __call__(
     ) -> dist.Work | None: ...
 
 
+@dataclass
+class DataParallelMeshDims:
+    """
+    Specifies which dimensions of a full SPMD :class:`DeviceMesh` correspond to
+    data parallelism when using :func:`fully_shard` whose parameters are already
+    DTensors on that mesh.
+
+    Attributes:
+        shard (Optional[Union[str, tuple[str, ...]]]): Mesh dimension name(s)
+            that FSDP shards parameters on. If a tuple of names, those dims
+            are flattened into a single shard dimension. At least one of
+            ``shard`` and ``replicate`` must be set.
+        replicate (Optional[Union[str, tuple[str, ...]]]): Mesh dimension
+            name(s) for HSDP or DDP replication. If a tuple of names, those
+            dims are flattened into a single replicate dimension.
+    """
+
+    shard: str | tuple[str, ...] | None = None
+    replicate: str | tuple[str, ...] | None = None
+
+    def __post_init__(self):
+        if self.shard is None and self.replicate is None:
+            raise ValueError(
+                "At least one of shard or replicate must be set in DataParallelMeshDims"
+            )
+
+    @property
+    def shard_names(self) -> tuple[str, ...]:
+        if self.shard is None:
+            return ()
+        if isinstance(self.shard, str):
+            return (self.shard,)
+        return tuple(self.shard)
+
+    @property
+    def replicate_names(self) -> tuple[str, ...]:
+        if self.replicate is None:
+            return ()
+        if isinstance(self.replicate, str):
+            return (self.replicate,)
+        return tuple(self.replicate)
+
+
 @dataclass
 class OffloadPolicy:
     """
diff --git a/torch/distributed/fsdp/_fully_shard/_fsdp_collectives.py b/torch/distributed/fsdp/_fully_shard/_fsdp_collectives.py
index 45b31c3799961..72fe919d84800 100644
--- a/torch/distributed/fsdp/_fully_shard/_fsdp_collectives.py
+++ b/torch/distributed/fsdp/_fully_shard/_fsdp_collectives.py
@@ -1,10 +1,11 @@
 import math
 from collections.abc import Callable, Sequence
 from itertools import chain
-from typing import Any, cast, NamedTuple
+from typing import Any, cast, Literal, NamedTuple
 
 import torch
 import torch.distributed as dist
+import torch.distributed._symmetric_memory as symm_mem
 from torch.distributed.device_mesh import _get_device_handle
 from torch.distributed.distributed_c10d import ReduceOp
 from torch.distributed.fsdp._fully_shard._fsdp_api import AllGather, ReduceScatter
@@ -15,7 +16,6 @@
     _get_dim0_padded_size,
     _raise_assert_with_print,
     _to_dtype_if_needed,
-    compiled_autograd_enabled,
 )
 from ._fsdp_param import FSDPParam, ShardedState
 
@@ -78,6 +78,36 @@ def allocate(
         return torch.empty(*size, dtype=dtype, device=device)
 
 
+class SymmMemAllocMixin:
+    def __init__(
+        self,
+        group: dist.ProcessGroup,
+        backend: Literal["NCCL"] = "NCCL",
+        *args: Any,
+        **kwargs: Any,
+    ):
+        self._group = group
+        symm_mem.set_backend(backend)
+        # Force initialization of communicator; otherwise, the rendezvous may
+        # see empty communicator.
+        # TODO: Remove this, maybe by warning user to perform eager dist init.
+        # For now, it is okay since it isjust a one-time cost at init.
+        dist.barrier(group=group)
+
+    def allocate(
+        self,
+        size: Sequence[int | torch.SymInt],
+        *,
+        dtype: torch.dtype,
+        device: torch.device,
+    ) -> torch.Tensor:
+        # Leverage MemPool to reuse the symmetric buffer, avoiding allocation
+        # and rendezvous overhead
+        mempool = symm_mem.get_mem_pool(device)
+        with torch.cuda.use_mem_pool(mempool):
+            return torch.empty(size, dtype=dtype, device=device)
+
+
 class DefaultAllGather(DefaultAllocMixin, AllGather):
     def __call__(
         self,
@@ -113,6 +143,35 @@ def __call__(
         )
 
 
+class SymmMemAllGather(SymmMemAllocMixin, AllGather):
+    def __init__(
+        self,
+        group: dist.ProcessGroup,
+        backend: Literal["NCCL"] = "NCCL",
+    ) -> None:
+        super().__init__(group, backend)
+
+    def __call__(
+        self,
+        output_tensor: torch.Tensor,
+        input_tensor: torch.Tensor,
+        group: dist.ProcessGroup,
+        async_op: bool = False,
+    ) -> dist.Work | None:
+        # We are doing inplace all-gather, so we need to rendezvous the output tensor only
+        symm_mem.rendezvous(output_tensor, group=group.group_name)
+        # Calling regular all-gather would already cause libraries like NCCL to
+        # use its optimized all-gather implementation for symmetric memory:
+        # - Copy Engine All-Gather (when zero-CTA policy is enabled)
+        # - Symmetric Kernel All-Gather (when zero-CTA policy is not enabled)
+        return dist.all_gather_into_tensor(
+            output_tensor,
+            input_tensor,
+            group=group,
+            async_op=async_op,
+        )
+
+
 class DefaultReduceScatter(DefaultAllocMixin, ReduceScatter):
     def __call__(
         self,
@@ -152,6 +211,35 @@ def __call__(
         )
 
 
+class SymmMemReduceScatter(SymmMemAllocMixin, ReduceScatter):
+    def __init__(
+        self,
+        group: dist.ProcessGroup,
+        backend: Literal["NCCL"] = "NCCL",
+    ) -> None:
+        super().__init__(group, backend)
+
+    def __call__(
+        self,
+        output_tensor: torch.Tensor,
+        input_tensor: torch.Tensor,
+        group: dist.ProcessGroup,
+        op: _ReduceOp,
+        async_op: bool = False,
+    ) -> dist.Work | None:
+        symm_mem.rendezvous(input_tensor, group=group.group_name)
+        symm_mem.rendezvous(output_tensor, group=group.group_name)
+        # Calling regular reduce-scatter would already cause libraries like NCCL to
+        # use its optimized reduce-scatter implementation for symmetric memory
+        return dist.reduce_scatter_tensor(
+            output=output_tensor,
+            input=input_tensor,
+            group=group,
+            op=op,
+            async_op=async_op,
+        )
+
+
 @torch.library.impl(lib, "all_gather_copy_in", "Meta")
 def all_gather_copy_in_meta(
     all_gather_inputs: list[torch.Tensor],
@@ -294,9 +382,6 @@ def foreach_all_gather(
 def _get_param_all_gather_inputs(
     fsdp_params: list[FSDPParam],
 ) -> list[list[torch.Tensor]]:
-    if compiled_autograd_enabled():
-        return [fsdp_param.all_gather_inputs for fsdp_param in fsdp_params]
-
     # Intentionally try to run a fast-path that bypasses abstractions for the
     # common FSDP case of bf16/fp32 mixed precision in order to use foreach
     # copy for lower CPU overhead and more efficient copying in eager
@@ -371,16 +456,13 @@ def foreach_all_gather_copy_out(
     ):
         # NOTE: Under compile, make sure we always recreate all_gather_outputs
         # per AllGather. See [Note: Invariants for torch.compile Traceable FSDP2].
-        force_recreate = compiled_autograd_enabled()
         fsdp_param.init_all_gather_outputs(
             all_gather_input_numels,
             all_gather_input_dtypes,
             world_size,
             device,
-            force_recreate=force_recreate,
         )
-        if not force_recreate:
-            fsdp_param.alloc_all_gather_outputs()
+        fsdp_param.alloc_all_gather_outputs()
         param_all_gather_outputs = fsdp_param.all_gather_outputs
         if fsdp_param.fsdp_placement.dim != 0:
             # Copy to a temporary and then chunk-cat into the final all-gather
@@ -398,13 +480,7 @@ def foreach_all_gather_copy_out(
         out = [t.view(world_size, -1) for t in split_with_sizes_out]
 
     # only avoid VC bump if we are not in inference mode
-    if torch._dynamo.is_compiling():
-        # For torch.compile, we turn off inference_mode for fake tensor
-        # propagation, and therefore graph break on is_inference. For `compile`,
-        # we don't care about VCs, so just skip the optimization.
-        non_inference_outs = []
-    else:
-        non_inference_outs = [o for o in out if not o.is_inference()]
+    non_inference_outs = [o for o in out if not o.is_inference()]
 
     if len(non_inference_outs) > 0:
         with torch.autograd._unsafe_preserve_version_counter(tuple(non_inference_outs)):
@@ -464,6 +540,7 @@ def foreach_reduce(
 ) -> tuple[
     torch.Tensor,
     torch.Event,
+    torch.Stream,
     torch.Event,
     torch.Tensor | None,
     torch.Event | None,
@@ -566,6 +643,7 @@ def foreach_reduce(
                 return (
                     reduce_scatter_input,
                     reduce_scatter_event,
+                    post_reduce_stream,
                     post_reduce_stream.record_event(),
                     all_reduce_input,
                     all_reduce_event,
@@ -584,6 +662,14 @@ def foreach_reduce(
                     group=all_reduce_group,
                     op=all_reduce_op,
                 )
+                # Keep refs to the reduce-dtype AR buffer + completion
+                # event so FSDPParamGroup._all_reduce_state can hold them
+                # across layers. This keeps the buffer off the caching
+                # allocator's free list; otherwise the next layer's
+                # reduce-scatter can reuse the same physical block while
+                # this layer's AR is still in flight, causing cross-layer
+                # gradient aliasing under slow AR. See PR #140044,
+                # regression test PR #180900.
                 all_reduce_input = reduce_output
                 all_reduce_event = all_reduce_stream.record_event()
     # -- END: ops in reduce_scatter stream
@@ -641,12 +727,11 @@ def foreach_reduce(
                     new_sharded_grad
                 )
                 fsdp_param.sharded_param.grad = new_sharded_dtensor_grad
-            if not compiled_autograd_enabled():
-                for hook in (
-                    getattr(fsdp_param.sharded_param, "_post_accumulate_grad_hooks", {})
-                    or {}
-                ).values():
-                    hook(fsdp_param.sharded_param)
+            for hook in (
+                getattr(fsdp_param.sharded_param, "_post_accumulate_grad_hooks", {})
+                or {}
+            ).values():
+                hook(fsdp_param.sharded_param)
             padded_sharded_numel = padded_unsharded_size.numel() // world_size
             flat_grad_offset += padded_sharded_numel
         post_reduce_event = post_reduce_stream.record_event()
@@ -657,6 +742,7 @@ def foreach_reduce(
     return (
         reduce_scatter_input,
         reduce_scatter_event,
+        post_reduce_stream,
         post_reduce_event,
         all_reduce_input,
         all_reduce_event,
diff --git a/torch/distributed/fsdp/_fully_shard/_fsdp_common.py b/torch/distributed/fsdp/_fully_shard/_fsdp_common.py
index 1a645ca7b7452..46b90e032df92 100644
--- a/torch/distributed/fsdp/_fully_shard/_fsdp_common.py
+++ b/torch/distributed/fsdp/_fully_shard/_fsdp_common.py
@@ -1,7 +1,8 @@
 # mypy: allow-untyped-defs
+import functools
 import math
 import traceback
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from enum import auto, Enum
 from typing import Any
 
@@ -9,31 +10,22 @@
 import torch.distributed as dist
 import torch.nn as nn
 from torch.distributed._composable.contract import _get_registry
-from torch.distributed.tensor import DeviceMesh, DTensor
+from torch.distributed.tensor import DeviceMesh, DTensor, Shard
 from torch.distributed.tensor._dtensor_spec import DTensorSpec
 
+from ._fsdp_api import DataParallelMeshDims
 
-_compiled_autograd_enabled: bool = False
 
+def _dynamo_disable(func):
+    """Disable dynamo tracing for FSDP hooks."""
 
-def detect_compiled_autograd():
-    if torch.compiler.is_compiling():
-        raise AssertionError(
-            "`detect_compiled_autograd()` is designed to be called in eager mode"
-        )
-    global _compiled_autograd_enabled
-    import torch._dynamo.compiled_autograd as ca
-
-    _compiled_autograd_enabled = (
-        ca.compiled_autograd_enabled
-        or ca.compiled_autograd_enabled_force_eager
-        or ca.in_compiled_autograd_region
-    )
+    @functools.wraps(func)
+    def wrapper(*args, **kwargs):
+        return torch._dynamo.disable(
+            func, recursive=True, reason="skipping FSDP hooks"
+        )(*args, **kwargs)
 
-
-def compiled_autograd_enabled():
-    global _compiled_autograd_enabled
-    return _compiled_autograd_enabled
+    return wrapper
 
 
 @dataclass
@@ -41,12 +33,19 @@ class DataParallelMeshInfo:
     mesh: DeviceMesh
     shard_mesh_dim: int | None = None
     replicate_mesh_dim: int | None = None
+    dp_mesh_dims: DataParallelMeshDims | None = None
+    # The full SPMD mesh (excluding PP dims) that params are distributed on.
+    # Must include all non-PP SPMD dims (e.g. DP + TP); passing a submesh
+    # that omits dims like TP will lead to incorrect behavior.
+    spmd_mesh: DeviceMesh | None = field(default=None, repr=False)
+    is_spmd_mesh: bool = field(default=False, init=False, repr=False)
 
     def __post_init__(self):
         if self.shard_mesh_dim is None and self.replicate_mesh_dim is None:
             raise AssertionError(
                 "At least one of shard_mesh_dim and replicate_mesh_dim must not be None"
             )
+        self.is_spmd_mesh = self.dp_mesh_dims is not None
 
 
 @dataclass
@@ -58,6 +57,10 @@ def __post_init__(self):
         self.shard_mesh_size: int = self.mesh.size(self.shard_mesh_dim)
         self.shard_process_group = self.mesh.get_group(self.shard_mesh_dim)
         self.shard_mesh_rank: int = self.shard_process_group.rank()
+        # Reduce-scatter shares the shard PG (one NCCL communicator) by default;
+        # FSDPModule.set_separate_reduce_scatter_group assigns a dedicated PG
+        # here so reduce-scatter can overlap all-gather in backward.
+        self.reduce_scatter_process_group: dist.ProcessGroup | None = None
 
 
 @dataclass
@@ -138,26 +141,16 @@ def _from_local_no_grad(
     This method is similar to ``DTensor.from_local()`` except that in eager mode
     it avoids some CPU overhead by avoiding default args and not being differentiable.
     """
-
-    if not compiled_autograd_enabled():
-        # pyrefly: ignore [bad-argument-type]
-        return DTensor(
-            # Use the local tensor directly instead of constructing a new tensor
-            # variable, e.g. with `view_as()`, since this is not differentiable
-            # pyrefly: ignore [bad-argument-count]
-            local_tensor,
-            sharding_spec,
-            # pyrefly: ignore [unexpected-keyword]
-            requires_grad=local_tensor.requires_grad,
-        )
-    else:
-        return DTensor.from_local(
-            local_tensor,
-            sharding_spec.mesh,
-            sharding_spec.placements,
-            shape=sharding_spec.shape,
-            stride=sharding_spec.stride,
-        )
+    # pyrefly: ignore [bad-argument-type]
+    return DTensor(
+        # Use the local tensor directly instead of constructing a new tensor
+        # variable, e.g. with `view_as()`, since this is not differentiable
+        # pyrefly: ignore [bad-argument-count]
+        local_tensor,
+        sharding_spec,
+        # pyrefly: ignore [unexpected-keyword]
+        requires_grad=local_tensor.requires_grad,
+    )
 
 
 def _to_dtype_if_needed(
@@ -180,3 +173,39 @@ def _cast_fp_tensor(dtype: torch.dtype, x: torch.Tensor) -> torch.Tensor:
 
 def is_bw() -> bool:
     return torch._C._current_graph_task_id() != -1
+
+
+@dataclass
+class ShardPlacementResult:
+    placement: Shard | None
+    mesh_info: FSDPMeshInfo
+
+
+ShardPlacementFnResult = Shard | ShardPlacementResult | None
+
+
+def resolve_shard_placement(
+    result: ShardPlacementFnResult,
+    default_mesh_info: FSDPMeshInfo,
+) -> ShardPlacementResult:
+    """Resolve the shard_placement_fn result to a ShardPlacementResult.
+
+    Handles different input types and applies defaults:
+    - None: Use default sharding (Shard(0)) on default mesh
+    - Shard: Use specified shard dimension on default mesh
+    - ShardPlacementResult: Use as-is
+
+    Args:
+        result: The return value from shard_placement_fn, or None if no fn provided.
+        default_mesh_info: The default FSDPMeshInfo to use if not specified.
+
+    Returns:
+        A ShardPlacementResult with placement and mesh_info.
+    """
+    if result is None:
+        return ShardPlacementResult(placement=None, mesh_info=default_mesh_info)
+    if isinstance(result, Shard):
+        return ShardPlacementResult(placement=result, mesh_info=default_mesh_info)
+    if isinstance(result, ShardPlacementResult):
+        return result
+    raise ValueError(f"Invalid shard_placement_fn result: {result}")
diff --git a/torch/distributed/fsdp/_fully_shard/_fsdp_init.py b/torch/distributed/fsdp/_fully_shard/_fsdp_init.py
index e0fe0455868c6..42f5293705cda 100644
--- a/torch/distributed/fsdp/_fully_shard/_fsdp_init.py
+++ b/torch/distributed/fsdp/_fully_shard/_fsdp_init.py
@@ -13,6 +13,7 @@
 from ._fsdp_common import (
     _is_composable_with_fsdp,
     DataParallelMeshInfo,
+    DDPMeshInfo,
     FSDPMeshInfo,
     HSDPMeshInfo,
 )
@@ -23,7 +24,8 @@
     from collections.abc import Callable
     from typing import Any
 
-    from ._fsdp_api import MixedPrecisionPolicy, OffloadPolicy
+    from ._fsdp_api import DataParallelMeshDims, MixedPrecisionPolicy, OffloadPolicy
+    from ._fsdp_common import ShardPlacementFnResult
     from ._fsdp_state import FSDPState
 
 
@@ -36,19 +38,44 @@ def _validate_module(module: nn.Module, func_name: str) -> None:
 
     Raises ValueError if the module is a container that doesn't implement forward.
     """
-    if isinstance(module, (nn.ModuleList, nn.ModuleDict)):
+    if (
+        isinstance(module, (nn.ModuleList, nn.ModuleDict))
+        and module.__class__.forward is nn.Module.forward
+    ):
         raise ValueError(
             f"{func_name} does not support containers that do not implement forward: {module}"
         )
 
 
-def _validate_mesh(mesh: "DeviceMesh") -> None:
+def _validate_mesh(
+    mesh: "DeviceMesh",
+    dp_mesh_dims: "DataParallelMeshDims | None" = None,
+) -> None:
     """
     Validate that the mesh can be used with fully_shard.
 
-    Raises ValueError if the mesh is not 1D or 2D.
-    Raises AssertionError if the mesh is 2D but mesh_dim_names is not specified.
+    When ``dp_mesh_dims`` is provided, validates that the named dims
+    exist in the mesh and at least one of shard/replicate is set.
+    Otherwise raises ValueError if the mesh is not 1D or 2D.
     """
+    if dp_mesh_dims is not None:
+        if dp_mesh_dims.shard is None and dp_mesh_dims.replicate is None:
+            raise ValueError(
+                "At least one of shard or replicate must be set in dp_mesh_dims"
+            )
+        if mesh.mesh_dim_names is None:
+            raise ValueError(
+                "mesh must have mesh_dim_names when dp_mesh_dims is provided"
+            )
+        names_to_check: list[str] = list(dp_mesh_dims.shard_names)
+        names_to_check.extend(dp_mesh_dims.replicate_names)
+        for name in names_to_check:
+            if name not in mesh.mesh_dim_names:
+                raise ValueError(
+                    f"Mesh dim name '{name}' not found in mesh.mesh_dim_names "
+                    f"{mesh.mesh_dim_names}"
+                )
+        return
     if mesh.ndim not in (1, 2):
         raise ValueError(f"fully_shard expects a 1D or 2D DeviceMesh but got {mesh}")
     if mesh.ndim == 2 and mesh.mesh_dim_names is None:
@@ -57,18 +84,71 @@ def _validate_mesh(mesh: "DeviceMesh") -> None:
         )
 
 
-def _get_mesh_info(mesh: "DeviceMesh") -> "FSDPMeshInfo":
+def _get_mesh_info(
+    mesh: "DeviceMesh",
+    dp_mesh_dims: "DataParallelMeshDims | None" = None,
+) -> "DataParallelMeshInfo":
     """
     Get the appropriate mesh info for the given mesh.
 
+    When ``dp_mesh_dims`` is provided, extracts the DP submesh from the
+    full SPMD mesh and returns FSDPMeshInfo, HSDPMeshInfo, or DDPMeshInfo
+    with ``dp_mesh_dims`` set and ``is_spmd_mesh`` as True.
+
     Returns FSDPMeshInfo for 1D mesh, HSDPMeshInfo for 2D mesh.
     """
+    if dp_mesh_dims is not None:
+        return _get_mesh_info_from_named_dims(mesh, dp_mesh_dims)
     if mesh.ndim == 1:
         return FSDPMeshInfo(mesh, shard_mesh_dim=0)
     else:
         return HSDPMeshInfo(mesh, shard_mesh_dim=1, replicate_mesh_dim=0)
 
 
+def _get_mesh_info_from_named_dims(
+    mesh: "DeviceMesh",
+    dp_mesh_dims: "DataParallelMeshDims",
+) -> "DataParallelMeshInfo":
+    shard_names = dp_mesh_dims.shard_names
+    replicate_names = dp_mesh_dims.replicate_names
+
+    def _get_submesh(names: tuple[str, ...]) -> "DeviceMesh":
+        if len(names) == 1:
+            return mesh[names[0]]
+        # Flatten multi-dim submesh into a single dim so FSDP's internal
+        # logic (which expects one shard and/or one replicate dim) works
+        # unchanged. This creates a new 1D DeviceMesh and ProcessGroup.
+        return mesh[names]._flatten("_".join(names))
+
+    if len(shard_names) == 0:  # DDP
+        dp_mesh = _get_submesh(replicate_names)
+        return DDPMeshInfo(
+            dp_mesh,
+            replicate_mesh_dim=0,
+            dp_mesh_dims=dp_mesh_dims,
+            spmd_mesh=mesh,
+        )
+    if len(replicate_names) == 0:  # FSDP
+        dp_mesh = _get_submesh(shard_names)
+        return FSDPMeshInfo(
+            dp_mesh,
+            shard_mesh_dim=0,
+            dp_mesh_dims=dp_mesh_dims,
+            spmd_mesh=mesh,
+        )
+    # HSDP
+    shard_mesh = _get_submesh(shard_names)
+    replicate_mesh = _get_submesh(replicate_names)
+    dp_mesh = DeviceMesh._concatenate([replicate_mesh, shard_mesh])
+    return HSDPMeshInfo(
+        dp_mesh,
+        shard_mesh_dim=1,
+        replicate_mesh_dim=0,
+        dp_mesh_dims=dp_mesh_dims,
+        spmd_mesh=mesh,
+    )
+
+
 def _get_post_forward_mesh_info(
     reshard_after_forward: bool | int, mesh_info: FSDPMeshInfo
 ) -> FSDPMeshInfo | None:
@@ -357,28 +437,99 @@ def _init_param_group(
     mesh_info: DataParallelMeshInfo,
     post_forward_mesh_info: FSDPMeshInfo | None,
     device: torch.device,
-    shard_placement_fn: "Callable[[nn.Parameter], Any] | None",
+    shard_placement_fn: "Callable[[nn.Parameter], ShardPlacementFnResult] | None",
     mp_policy: "MixedPrecisionPolicy",
     offload_policy: "OffloadPolicy",
+    reshard_after_forward: bool | int = True,
 ) -> None:
     """
-    Initialize the FSDP param group for the given state if there are params.
+    Initialize FSDP param groups for the given state.
 
-    This is shared between fully_shard and replicate.
+    Params are grouped by their process group (derived from ``mesh_info`` via
+    ``shard_placement_fn``). Each group becomes a separate ``FSDPParamGroup``.
+    When ``shard_placement_fn`` is ``None`` or returns the same mesh for all
+    params, this creates a single group.
     """
     # Import here to avoid circular imports
+    from ._fsdp_common import FSDPMeshInfo, resolve_shard_placement
     from ._fsdp_param_group import FSDPParamGroup
 
-    if params:
-        state._fsdp_param_group = FSDPParamGroup(
-            params,
-            modules,
+    if not params:
+        return
+
+    if shard_placement_fn is None:
+        # No shard_placement_fn means all params use the same mesh_info,
+        # so no grouping is needed. This also handles DDPMeshInfo from
+        # replicate_with_fsdp, which doesn't have shard_process_group.
+        state._fsdp_param_groups.append(
+            FSDPParamGroup(
+                params,
+                modules,
+                mesh_info,
+                post_forward_mesh_info,
+                device,
+                shard_placement_fn,
+                mp_policy,
+                offload_policy,
+            )
+        )
+        return
+
+    # Group params by their process group to support per-param mesh,
+    # e.g., expert params using ep_mesh vs regular params using dp_mesh.
+    # For HSDP, also key by replicate_process_group to avoid grouping
+    # FSDPMeshInfo params with HSDPMeshInfo params that share the same
+    # shard_process_group but require different gradient reduction behavior.
+    if not isinstance(mesh_info, FSDPMeshInfo):
+        raise ValueError(
+            "Per-param mesh via shard_placement_fn is not supported with "
+            f"{type(mesh_info).__name__}; it requires FSDPMeshInfo (or subclass)"
+        )
+    pg_to_group: dict[
+        tuple[dist.ProcessGroup, dist.ProcessGroup | None],
+        tuple[FSDPMeshInfo, list[nn.Parameter]],
+    ] = {}
+    for param in params:
+        param_mesh_info = resolve_shard_placement(
+            shard_placement_fn(param),
             mesh_info,
-            post_forward_mesh_info,
-            device,
-            shard_placement_fn,
-            mp_policy,
-            offload_policy,
+        ).mesh_info
+        shard_pg = param_mesh_info.shard_process_group
+        replicate_pg: dist.ProcessGroup | None = None
+        if isinstance(param_mesh_info, HSDPMeshInfo):
+            replicate_pg = param_mesh_info.replicate_process_group
+        key = (shard_pg, replicate_pg)
+        if key not in pg_to_group:
+            pg_to_group[key] = (param_mesh_info, [param])
+        else:
+            existing_mesh_info = pg_to_group[key][0]
+            if existing_mesh_info is not param_mesh_info:
+                raise ValueError(
+                    f"Params sharing the same process group must use the same "
+                    f"FSDPMeshInfo object, but got different objects: "
+                    f"{existing_mesh_info} vs {param_mesh_info}"
+                )
+            pg_to_group[key][1].append(param)
+
+    # Create a FSDPParamGroup per process group
+    for group_mesh_info, group_params in pg_to_group.values():
+        if group_mesh_info is not mesh_info:
+            group_post_forward = _get_post_forward_mesh_info(
+                reshard_after_forward, group_mesh_info
+            )
+        else:
+            group_post_forward = post_forward_mesh_info
+        state._fsdp_param_groups.append(
+            FSDPParamGroup(
+                group_params,
+                modules,
+                group_mesh_info,
+                group_post_forward,
+                device,
+                shard_placement_fn,
+                mp_policy,
+                offload_policy,
+            )
         )
 
 
diff --git a/torch/distributed/fsdp/_fully_shard/_fsdp_param.py b/torch/distributed/fsdp/_fully_shard/_fsdp_param.py
index b8b6162e87fa4..f09e7b8dcd00a 100644
--- a/torch/distributed/fsdp/_fully_shard/_fsdp_param.py
+++ b/torch/distributed/fsdp/_fully_shard/_fsdp_param.py
@@ -23,13 +23,24 @@
     _get_dim_chunked_size,
     _raise_assert_with_print,
     _to_dtype_if_needed,
-    compiled_autograd_enabled,
     DataParallelMeshInfo,
     FSDPMeshInfo,
     HSDPMeshInfo,
+    resolve_shard_placement,
+    ShardPlacementFnResult,
 )
 
 
+_orig_param_uid_counter = itertools.count()
+
+
+def _get_orig_param_uid(param: nn.Parameter) -> int:
+    if not hasattr(param, "_fsdp_orig_uid"):
+        uid = next(_orig_param_uid_counter)
+        param._fsdp_orig_uid = uid  # pyrefly: ignore[missing-attribute]
+    return param._fsdp_orig_uid  # pyrefly: ignore[missing-attribute]
+
+
 """
 [Note: FSDP tensors]
 FSDP considers the following tensors:
@@ -82,54 +93,6 @@ def copy_(tensor, data):
     tensor.copy_(data)
 
 
-"""
-[Note: Avoiding functionalization for fsdp.copy_ and inductor.resize_storage_bytes_]
-
-Currently we don't functionalize `fsdp.copy_` op or `inductor.resize_storage_bytes_` op
-(i.e. they show up as a mutation op in the middle of the AOT joint graph).
-
-Reason:
-Traceable FSDP2 compiled autograd BWD graph have the following traits:
-(1) Two inputs of the graph were aliased to each other (one from hook closed-over tensors, one from FWD saved tensors).
-(2) One of them is mutated (copy_ and resize_ to handle the all-gathered param).
-(3) They are both subclasses.
-The combination of these traits is not supported by AOTAutograd (it's difficult to reason about subclass aliasing).
-So this doesn't work at all for Traceable FSDP2.
-
-The compromise we use is to avoid functionalization for the FSDP2 copy_ and resize_ ops.
-This avoids the problem above, because from AOTAutograd point-of-view there are no mutations
-that functionalization needs to handle. (Although we need to be careful not to DCE those mutable ops.)
-
-We can avoid this functionalization because:
-(1) The nn.Parameter is never used before its .copy_() is called in eager code (i.e. no alias of it is created),
-so it's safe to call .copy_() in the middle of the graph to update its content and start using the nn.Parameter downstream.
-(2) We always re-allocate the buffer for nn.Parameter to store the AllGather output and to be used in downstream user ops.
-So calling resize-to-0 in the middle of the graph to free nn.Parameter memory after use should always be okay
-(since we always allocate anew next time we need it, we strictly don't need to keep the old tensor storage around anymore).
-
-Q: Wouldn't the extra resize_ and copy_ ops hurt both memory usage and performance?
-A: Yes it would. As an optimization, we have an Inductor post-grad FX pass to remove those resize_ and copy_ ops
-for unsharded params that have this pattern: resize_(full) -> copy_ -> resize_(0).
-
-TODO:
-Now that we are maintaining the invariant of "no aliased + mutated graph inputs" in both the forward and backward,
-it is now more feasible to functionalize all of the mutable FSDP ops. Some of the pros and cons are:
-
-Cons (of functionalizing those ops):
-(1) By not functionalizing them as we are today, we are making it more likely that they will run at the "correct" time
-in the generated code. If we start to functionalize them, we will need to make sure that Inductor reinplaces them
-in a way where it properly moves the mutations back to exactly where they should have run, or we risk suffering worse
-peak memory than eager. (We probably already need to do something similar in Inductor's reinplacing for copy_:
-https://github.com/pytorch/pytorch/issues/135305#issuecomment-2334888089)
-
-Pros (of functionalizing):
-(1) Better safety, we don't need to worry about the graph passes in inductor/partitioning handling input mutations
-mid-graph quite as much (to be fair we've already done some amount of auditing, but we might have to do some more).
-(2) Better perf: each mutation midway through the graph prevents Inductor from pattern matching across it.
-But maybe there are few enough mutations induced by FSDP for this to matter.
-"""
-
-
 @torch.library.impl(lib, "copy_", "Functionalize")
 def copy__functionalize(tensor, data):
     torch._sync(tensor)
@@ -213,12 +176,14 @@ class FSDPParam:
     _unsharded_param: nn.Parameter  # ND
     unsharded_accumulated_grad: torch.Tensor | None  # ND
     _sharding_spec: DTensorSpec
-    # DTensor attributes (only defined for DTensor `param`):
-    _tp_spec: DTensorSpec
+    _unsharded_dtensor_spec: (
+        DTensorSpec | None
+    )  # set for DTensor params (SPMD or TP/EP)
     all_gather_outputs: list[torch.Tensor]  # 1D
     # All-gather extension attributes
     _extensions_data: ExtensionsData
     _unsharded_inner_tensors: list[torch.Tensor]
+    _orig_param_uid: int
 
     def __init__(
         self,
@@ -227,14 +192,12 @@ def __init__(
         mesh_info: DataParallelMeshInfo,
         post_forward_mesh_info: FSDPMeshInfo | None,
         device: torch.device,
-        shard_placement_fn: Callable[[nn.Parameter], Shard | None] | None,
+        shard_placement_fn: Callable[[nn.Parameter], ShardPlacementFnResult] | None,
         mp_policy: MixedPrecisionPolicy,
         offload_policy: OffloadPolicy,
     ):
         self._module_info: ParamModuleInfo = module_info
-        self.mesh_info = mesh_info
         self.post_forward_mesh_info = post_forward_mesh_info
-        # pyrefly: ignore [read-only]
         self.device = device
         self.mp_policy = mp_policy
         self.offload_to_cpu: bool = isinstance(offload_policy, CPUOffloadPolicy)
@@ -242,7 +205,7 @@ def __init__(
             self.offload_to_cpu and cast(CPUOffloadPolicy, offload_policy).pin_memory
         )
         self.grad_offload_event: torch.Event | None = None
-        self._init_sharded_param(param, device, shard_placement_fn)
+        self._init_sharded_param(param, device, shard_placement_fn, mesh_info)
         if self.post_forward_mesh_info:
             self._init_sharded_post_forward_param_metadata(param)
         self._init_extensions()
@@ -262,8 +225,20 @@ def _init_sharded_param(
         self,
         param: nn.Parameter,
         device: torch.device,
-        shard_placement_fn: Callable | None,
+        shard_placement_fn: Callable[[nn.Parameter], ShardPlacementFnResult] | None,
+        mesh_info: DataParallelMeshInfo,
     ):
+        if callable(shard_placement_fn):
+            shard_result = resolve_shard_placement(
+                shard_placement_fn(param),
+                cast(FSDPMeshInfo, mesh_info),
+            )
+            self.mesh_info = shard_result.mesh_info
+            fsdp_placement = shard_result.placement
+        else:
+            self.mesh_info = mesh_info  # pyrefly: ignore[bad-assignment]
+            fsdp_placement = None
+        self._shard_mesh = self._init_shard_mesh()
         if param.device != device and param.device.type != "meta":
             raise AssertionError(
                 f"Expects the parameter to already be moved to device {device} but got {param.device}"
@@ -272,7 +247,6 @@ def _init_sharded_param(
             raise NotImplementedError(
                 f"FSDP does not support non-contiguous parameters yet: {param.shape=} {param.stride()=}"
             )
-        fsdp_placement = shard_placement_fn(param) if shard_placement_fn else None
         if fsdp_placement is None:
             fsdp_placement = Shard(0)
         elif fsdp_placement.dim < 0:
@@ -288,69 +262,8 @@ def _init_sharded_param(
         # TODO: Simplify the following sharded parameter padding logic after
         # https://github.com/pytorch/pytorch/issues/113045
         self.is_dtensor = isinstance(param, DTensor)
-        if self.is_dtensor:
-            self._tp_spec = cast(DTensor, param)._spec
-            dp_mesh, tp_mesh = (self.mesh_info.mesh, self._tp_spec.mesh)
-            if dp_mesh is None or tp_mesh is None:
-                raise AssertionError(
-                    "FSDP requires the DP and model parallel TP/EP mesh to be not None but got: \n"
-                    f"DP's mesh: {dp_mesh}\nTP/EP's mesh: {tp_mesh}"
-                )
-            self._spmd_mesh = DeviceMesh._concatenate([dp_mesh, tp_mesh])
-            if len(self._tp_spec.placements) > 2:
-                raise NotImplementedError(
-                    f"FSDP only supports 1D TP/EP or 2D EP+TP, not {self._tp_spec.placements}"
-                )
-            split_factor = self._tp_spec.num_shards_map[shard_dim]
-            if not (2 <= self._spmd_mesh.ndim <= 4):
-                raise AssertionError(
-                    "_spmd_mesh.ndim can only be 2 (FSDP+TP/EP), 3 (FSDP+EP+TP, HSDP+TP/EP), "
-                    f"or 4 (HSDP+EP+TP) but got {self._spmd_mesh.ndim}."
-                )
-            self._spmd_placements: tuple[Placement, ...]
-            if isinstance(self.mesh_info, FSDPMeshInfo):  # FSDP or HSDP
-                dp_shard_tp_placement = (
-                    (
-                        _StridedShard(shard_dim, split_factor=split_factor)
-                        if split_factor > 1
-                        else fsdp_placement
-                    ),
-                    *self._tp_spec.placements,
-                )
-            else:  # DDP
-                dp_shard_tp_placement = (
-                    (Replicate()),
-                    *self._tp_spec.placements,
-                )
-            if isinstance(self.mesh_info, HSDPMeshInfo):  # HSDP
-                if self.mesh_info.replicate_mesh_dim != 0:
-                    raise AssertionError(
-                        f"Expected replicate_mesh_dim to be 0, got {self.mesh_info.replicate_mesh_dim}"
-                    )
-                self._spmd_placements = (Replicate(),) + dp_shard_tp_placement
-            else:  # FSDP or DDP
-                self._spmd_placements = dp_shard_tp_placement
-
-            self._sharding_spec = DTensorSpec(
-                self._spmd_mesh,
-                self._spmd_placements,
-                tensor_meta=self._tp_spec.tensor_meta,
-            )
-            param_data = cast(DTensor, param)._local_tensor
-        else:
-            self._spmd_mesh = self.mesh_info.mesh
-            if isinstance(self.mesh_info, HSDPMeshInfo):  # HSDP
-                self._spmd_placements = (Replicate(), fsdp_placement)
-            elif isinstance(self.mesh_info, FSDPMeshInfo):  # FSDP
-                self._spmd_placements = (fsdp_placement,)
-            elif isinstance(self.mesh_info, DDPMeshInfo):  # DDP
-                self._spmd_placements = (Replicate(),)
-            self._sharding_spec = DTensorSpec(
-                self._spmd_mesh,
-                self._spmd_placements,
-                tensor_meta=TensorMeta(param.size(), param.stride(), param.dtype),
-            )
-            param_data = param
+        self._orig_param_uid = _get_orig_param_uid(param)
+        param_data = self._init_sharding_spec(param, fsdp_placement, shard_dim)
         if not param_data.is_contiguous():
             raise AssertionError(
                 f"Expected contiguous tensor, got {param_data.shape=} {param_data.stride()=}"
@@ -405,13 +318,189 @@ def _init_sharded_param(
             raise AssertionError(
                 f"Expected contiguous tensor with {self.fsdp_placement=}"
             )
-        self.sharded_param = nn.Parameter(self.to_sharded_dtensor(sharded_param))
-        self.sharded_param.requires_grad_(param.requires_grad)
+        self.sharded_param = nn.Parameter(
+            self.to_sharded_dtensor(sharded_param),
+            requires_grad=param.requires_grad,
+        )
         # Let `param_data` be freed normally when its ref count reaches 0 when
         # the `fully_shard` call returns to allow provided parameters to alias
         self._setattr_on_modules(self.sharded_param)
         self.sharded_state = ShardedState.SHARDED
 
+    def _init_sharding_spec(
+        self,
+        param: nn.Parameter,
+        fsdp_placement: Shard,
+        shard_dim: int,
+    ) -> torch.Tensor:
+        """
+        Build ``_sharding_spec``, ``_spmd_mesh``, and ``_spmd_placements`` and
+        return the local tensor data to be sharded.
+        """
+        self._unsharded_dtensor_spec = None
+        if self.mesh_info.is_spmd_mesh and not self.is_dtensor:
+            raise ValueError(
+                "When dp_mesh_dims is provided, all parameters must be "
+                "DTensors on the full SPMD mesh (e.g. via distribute_module). "
+                f"Got plain tensor for parameter '{self._module_info.param_name}'."
+            )
+        if self.is_dtensor and self.mesh_info.is_spmd_mesh:
+            return self._init_sharding_spec_spmd(param, fsdp_placement, shard_dim)
+        if self.is_dtensor:
+            return self._init_sharding_spec_tp(param, fsdp_placement, shard_dim)
+        return self._init_sharding_spec_plain(param, fsdp_placement)
+
+    def _init_sharding_spec_spmd(
+        self,
+        param: nn.Parameter,
+        fsdp_placement: Shard,
+        shard_dim: int,
+    ) -> torch.Tensor:
+        """SPMD path: param is a DTensor on the full SPMD mesh."""
+        self._unsharded_dtensor_spec = cast(DTensor, param)._spec
+        spmd_mesh = self._unsharded_dtensor_spec.mesh
+        dp_dim_names = self.mesh_info.dp_mesh_dims
+        if dp_dim_names is None:
+            raise AssertionError("dp_dim_names must not be None for SPMD mesh")
+        if spmd_mesh.mesh_dim_names is None:
+            raise AssertionError("spmd_mesh.mesh_dim_names must not be None")
+        if (
+            self.mesh_info.spmd_mesh is not None
+            and spmd_mesh is not self.mesh_info.spmd_mesh
+        ):
+            raise ValueError(
+                "Expected param's DTensor mesh to be the same mesh passed "
+                "to fully_shard, but got different mesh objects"
+            )
+
+        dp_shard_indices = [
+            spmd_mesh.mesh_dim_names.index(n) for n in dp_dim_names.shard_names
+        ]
+
+        orig_placements = self._unsharded_dtensor_spec.placements
+        for idx in dp_shard_indices:
+            if not isinstance(orig_placements[idx], Replicate):
+                raise ValueError(
+                    f"Expected Replicate() on DP shard dim "
+                    f"'{spmd_mesh.mesh_dim_names[idx]}' (index {idx}) "
+                    f"but got {orig_placements[idx]}"
+                )
+        dp_replicate_indices = []
+        for rep_name in dp_dim_names.replicate_names:
+            rep_idx = spmd_mesh.mesh_dim_names.index(rep_name)
+            dp_replicate_indices.append(rep_idx)
+            if not isinstance(orig_placements[rep_idx], Replicate):
+                raise ValueError(
+                    f"Expected Replicate() on DP replicate dim "
+                    f"'{spmd_mesh.mesh_dim_names[rep_idx]}' (index {rep_idx}) "
+                    f"but got {orig_placements[rep_idx]}"
+                )
+
+        # Cache DP dim indices so _get_grad_inner_tensor can skip
+        # redistribution on DP dims and let FSDP's reduce-scatter handle them.
+        self._dp_dim_indices: frozenset[int] = frozenset(
+            dp_shard_indices + dp_replicate_indices
+        )
+
+        new_placements = list(orig_placements)
+        for dp_idx in dp_shard_indices:
+            # split_factor = number of non-DP shards on shard_dim from
+            # mesh dims with higher index (the "right-side" dims that
+            # _StridedShard needs to interleave with)
+            sf = 1
+            for j in range(dp_idx + 1, spmd_mesh.ndim):
+                p = orig_placements[j]
+                if isinstance(p, (Shard, _StridedShard)) and p.dim == shard_dim:
+                    sf *= spmd_mesh.size(j)
+            new_placements[dp_idx] = (
+                _StridedShard(shard_dim, split_factor=sf) if sf > 1 else fsdp_placement
+            )
+
+        self._spmd_mesh = spmd_mesh
+        self._spmd_placements: tuple[Placement, ...] = tuple(new_placements)
+        self._sharding_spec = DTensorSpec(
+            self._spmd_mesh,
+            self._spmd_placements,
+            tensor_meta=self._unsharded_dtensor_spec.tensor_meta,
+        )
+        return cast(DTensor, param)._local_tensor
+
+    def _init_sharding_spec_tp(
+        self,
+        param: nn.Parameter,
+        fsdp_placement: Shard,
+        shard_dim: int,
+    ) -> torch.Tensor:
+        """TP/EP path: param is a DTensor, DP mesh is separate from TP mesh."""
+        self._unsharded_dtensor_spec = cast(DTensor, param)._spec
+        dp_mesh, tp_mesh = (self.mesh_info.mesh, self._unsharded_dtensor_spec.mesh)
+        if dp_mesh is None or tp_mesh is None:
+            raise AssertionError(
+                "FSDP requires the DP and model parallel TP/EP mesh to be not None but got: \n"
+                f"DP's mesh: {dp_mesh}\nTP/EP's mesh: {tp_mesh}"
+            )
+        self._spmd_mesh = DeviceMesh._concatenate([dp_mesh, tp_mesh])
+        if len(self._unsharded_dtensor_spec.placements) > 2:
+            raise NotImplementedError(
+                f"FSDP only supports 1D TP/EP or 2D EP+TP, not {self._unsharded_dtensor_spec.placements}"
+            )
+        split_factor = self._unsharded_dtensor_spec.num_shards_map[shard_dim]
+        if not (2 <= self._spmd_mesh.ndim <= 4):
+            raise AssertionError(
+                "_spmd_mesh.ndim can only be 2 (FSDP+TP/EP), 3 (FSDP+EP+TP, HSDP+TP/EP), "
+                f"or 4 (HSDP+EP+TP) but got {self._spmd_mesh.ndim}."
+            )
+        if isinstance(self.mesh_info, FSDPMeshInfo):
+            dp_shard_tp_placement = (
+                (
+                    _StridedShard(shard_dim, split_factor=split_factor)
+                    if split_factor > 1
+                    else fsdp_placement
+                ),
+                *self._unsharded_dtensor_spec.placements,
+            )
+        else:  # DDP
+            dp_shard_tp_placement = (
+                Replicate(),
+                *self._unsharded_dtensor_spec.placements,
+            )
+        self._spmd_placements: tuple[Placement, ...]
+        if isinstance(self.mesh_info, HSDPMeshInfo):
+            if self.mesh_info.replicate_mesh_dim != 0:
+                raise AssertionError(
+                    f"Expected replicate_mesh_dim to be 0, got {self.mesh_info.replicate_mesh_dim}"
+                )
+            self._spmd_placements = (Replicate(),) + dp_shard_tp_placement
+        else:
+            self._spmd_placements = dp_shard_tp_placement
+
+        self._sharding_spec = DTensorSpec(
+            self._spmd_mesh,
+            self._spmd_placements,
+            tensor_meta=self._unsharded_dtensor_spec.tensor_meta,
+        )
+        return cast(DTensor, param)._local_tensor
+
+    def _init_sharding_spec_plain(
+        self,
+        param: nn.Parameter,
+        fsdp_placement: Shard,
+    ) -> torch.Tensor:
+        """Plain tensor path: param is not a DTensor."""
+        self._spmd_mesh = self.mesh_info.mesh
+        if isinstance(self.mesh_info, HSDPMeshInfo):
+            self._spmd_placements = (Replicate(), fsdp_placement)
+        elif isinstance(self.mesh_info, FSDPMeshInfo):
+            self._spmd_placements = (fsdp_placement,)
+        elif isinstance(self.mesh_info, DDPMeshInfo):
+            self._spmd_placements = (Replicate(),)
+        self._sharding_spec = DTensorSpec(
+            self._spmd_mesh,
+            self._spmd_placements,
+            tensor_meta=TensorMeta(param.size(), param.stride(), param.dtype),
+        )
+        return param
+
     def _init_sharded_post_forward_param_metadata(self, param: torch.Tensor) -> None:
         mesh_info = self.post_forward_mesh_info
         if mesh_info is None:
@@ -443,8 +532,10 @@ def init_dtype_attrs(self, mp_policy: MixedPrecisionPolicy):
         # then we do not need extra casting
         if reduce_dtype == param_dtype:
             reduce_dtype = None
-        # Clamp `param_dtype` to `None` if no casting is required
-        if param_dtype == self.orig_dtype:
+        # Clamp `param_dtype` to `None` if no casting is required or if the
+        # parameter is non-floating-point (mixed precision is only meaningful
+        # for floating-point parameters)
+        if param_dtype == self.orig_dtype or not self.orig_dtype.is_floating_point:
             param_dtype = None
         self.param_dtype = param_dtype
         self.reduce_dtype = reduce_dtype
@@ -479,24 +570,7 @@ def init_all_gather_outputs(
         ]
 
     def init_unsharded_param(self):
-        """
-        [Note: Invariants for torch.compile Traceable FSDP2]
-        1. Under compile, we always re-populate the content of `self._unsharded_param`
-           per AllGather using the slow path.
-        2. Under compile, we always recreate `self.all_gather_outputs` per AllGather.
-           This is to ensure the buffer creation is internal to the graph and
-           avoid `self.all_gather_outputs` being captured as a graph input.
-        3. Under compile, at the end of `free_unsharded_param()`, we always clean up
-           `self.all_gather_outputs` and `self._unsharded_inner_tensors`,
-           to avoid them being captured as graph output.
-
-        With these invariants, only these tensors will be inputs to the graph:
-        - Sharded parameters
-        - Placeholders for the `self._unsharded_param` nn.Parameter
-        """
-        if not compiled_autograd_enabled() and hasattr(
-            self, "_unsharded_param"
-        ):  # after the 1st all-gather
+        if hasattr(self, "_unsharded_param"):  # after the 1st all-gather
             inner_tensor = self._sharded_local_tensor
             if not hasattr(inner_tensor, "fsdp_post_all_gather"):
                 return  # already initialized
@@ -512,9 +586,7 @@ def init_unsharded_param(self):
             self._extensions_data.clear()
             return
         inner_tensor = self._sharded_local_tensor
-        if not compiled_autograd_enabled() and hasattr(
-            inner_tensor, "fsdp_post_all_gather"
-        ):
+        if hasattr(inner_tensor, "fsdp_post_all_gather"):
             all_gather_outputs = self._unflatten_all_gather_outputs()
             (
                 unsharded_tensor,
@@ -539,27 +611,13 @@ def init_unsharded_param(self):
             self._contiguous_orig_stride,
             storage_offset=0,
         )
-        if self.is_dtensor:
-            unsharded_param = _from_local_no_grad(unsharded_param, self._tp_spec)
-        if hasattr(self, "_unsharded_param"):
-            if not compiled_autograd_enabled():
-                raise AssertionError("Expected compiled_autograd to be enabled")
-            with (
-                torch.no_grad(),
-                torch.autograd._unsafe_preserve_version_counter(self._unsharded_param),
-            ):
-                # NOTE: Under compile, if an unsharded param goes through
-                # resize_(full) -> copy_ -> resize_(0) pattern, we will remove those
-                # resize_ and copy_ ops in a compiler graph pass
-                # `remove_fsdp2_unsharded_param_graph_input_usage` to recover performance.
-                self._unsharded_param.untyped_storage().resize_(
-                    self._unsharded_param.numel() * self._unsharded_param.itemsize
-                )
-                torch.ops.fsdp.copy_(self._unsharded_param, unsharded_param)
-        else:
-            self._unsharded_param = nn.Parameter(
-                unsharded_param, requires_grad=self.sharded_param.requires_grad
+        if self._unsharded_dtensor_spec is not None:
+            unsharded_param = _from_local_no_grad(
+                unsharded_param, self._unsharded_dtensor_spec
             )
+        self._unsharded_param = nn.Parameter(
+            unsharded_param, requires_grad=self.sharded_param.requires_grad
+        )
 
     def _unflatten_all_gather_outputs(self) -> tuple[torch.Tensor, ...]:
         return tuple(
@@ -577,7 +635,7 @@ def to_sharded(self) -> None:
     def to_sharded_post_forward(self) -> None:
         if self.is_dtensor:
             raise NotImplementedError(
-                "Resharding to smaller mesh with TP is not supported yet"
+                "Resharding to smaller mesh is not supported for DTensor parameters yet"
             )
         self._assert_in_states(ShardedState.UNSHARDED)
         if self.post_forward_mesh_info is None:
@@ -593,7 +651,6 @@ def to_sharded_post_forward(self) -> None:
                 f"world size ({shard_world_size})"
             )
         shard_rank = self.post_forward_mesh_info.shard_mesh_rank
-        # pyrefly: ignore [unbound-name]
         sharded_numel = numel // shard_world_size
         self._sharded_post_forward_param_data = (
             self.all_gather_outputs[0].narrow(
@@ -607,7 +664,8 @@ def to_sharded_post_forward(self) -> None:
             storage_offset=0,
         )
         self._sharded_post_forward_param = nn.Parameter(
-            self.to_sharded_post_forward_dtensor(sharded_post_forward_tensor)
+            self.to_sharded_post_forward_dtensor(sharded_post_forward_tensor),
+            requires_grad=self.sharded_param.requires_grad,
         )
         self._setattr_on_modules(self._sharded_post_forward_param)
         self.free_unsharded_param()
@@ -693,34 +751,16 @@ def alloc_all_gather_outputs(self) -> None:
             alloc_storage(tensor)
 
     def free_unsharded_param(self) -> None:
-        if compiled_autograd_enabled():
-            """
-            Assumptions under compile:
-            - `self._unsharded_param` is NOT an alias of `self.all_gather_outputs`.
-            Instead, we resize `self._unsharded_param` storage size to full and then
-            explicitly *copy* the data from `self.all_gather_outputs` to `self._unsharded_param`
-            in `init_unsharded_param()`. (For full-graph FSDP2 case, we will then remove
-            the resize_ and copy_ ops in a compiler graph pass to recover performance.)
-            - `self.all_gather_outputs` and `self._unsharded_inner_tensors` are NOT
-            graph inputs. They are created within the graph and is guaranteed to be freed
-            by the end of the graph. They don't leak outside of the graph.
-            """
-            self._unsharded_param.untyped_storage().resize_(0)
-            self.all_gather_outputs = []
-            self._unsharded_inner_tensors = []
-        else:
-            for tensor in itertools.chain(
-                self.all_gather_outputs, self._unsharded_inner_tensors
-            ):
-                free_storage(tensor)
+        for tensor in itertools.chain(
+            self.all_gather_outputs, self._unsharded_inner_tensors
+        ):
+            free_storage(tensor)
 
     @property
     def all_gather_inputs(self) -> list[torch.Tensor]:  # 1D
         self._assert_in_states(ShardedState.SHARDED, ShardedState.SHARDED_POST_FORWARD)
         if self.sharded_state == ShardedState.SHARDED:
-            if not compiled_autograd_enabled() and hasattr(
-                self._sharded_local_tensor, "fsdp_pre_all_gather"
-            ):
+            if hasattr(self._sharded_local_tensor, "fsdp_pre_all_gather"):
                 sharded_local_tensor = self._sharded_local_tensor
                 if self.offload_to_cpu:
                     sharded_local_tensor = sharded_local_tensor.to(
@@ -786,9 +826,7 @@ def all_gather_inputs(self) -> list[torch.Tensor]:  # 1D
                 )
             return [_to_dtype_if_needed(sharded_param_data, self.param_dtype)]
         elif self.sharded_state == ShardedState.SHARDED_POST_FORWARD:
-            if not compiled_autograd_enabled() and hasattr(
-                self._sharded_local_tensor, "fsdp_pre_all_gather"
-            ):
+            if hasattr(self._sharded_local_tensor, "fsdp_pre_all_gather"):
                 raise NotImplementedError
             all_gather_input = _to_dtype_if_needed(
                 cast(torch.Tensor, self._sharded_post_forward_param_data),
@@ -821,13 +859,32 @@ def _get_grad_inner_tensor(self, grad: torch.Tensor) -> torch.Tensor:
                 grad = grad.wait()
             if not isinstance(grad, DTensor):
                 raise AssertionError(f"Expected DTensor, got {type(grad)}")
-            placements = self._tp_spec.placements
-            if placements != grad.placements:
-                if len(self._tp_spec.placements) != len(grad.placements):
-                    raise AssertionError(
-                        f"Expected same placement length: {self._tp_spec=} {grad.placements=}"
-                    )
-                grad = grad.redistribute(placements=placements)
+            if self._unsharded_dtensor_spec is None:
+                raise AssertionError(
+                    "Expected _unsharded_dtensor_spec for DTensor param"
+                )
+            placements = self._unsharded_dtensor_spec.placements
+            if self.mesh_info.is_spmd_mesh:
+                # Only redistribute non-DP dims; keep Partial on DP dims
+                # so FSDP's reduce-scatter handles them directly, avoiding
+                # a redundant all-reduce on the DP dimensions.
+                target_placements = tuple(
+                    grad.placements[i] if i in self._dp_dim_indices else placements[i]
+                    for i in range(len(placements))
+                )
+                if target_placements != grad.placements:
+                    if len(placements) != len(grad.placements):
+                        raise AssertionError(
+                            f"Expected same placement length: {placements=} {grad.placements=}"
+                        )
+                    grad = grad.redistribute(placements=target_placements)
+            else:
+                if placements != grad.placements:
+                    if len(placements) != len(grad.placements):
+                        raise AssertionError(
+                            f"Expected same placement length: {placements=} {grad.placements=}"
+                        )
+                    grad = grad.redistribute(placements=placements)
             grad = grad._local_tensor
         return grad
 
@@ -835,8 +892,7 @@ def _get_grad_inner_tensor(self, grad: torch.Tensor) -> torch.Tensor:
     def _sharded_local_tensor(self) -> torch.Tensor:
         return cast(DTensor, self.sharded_param)._local_tensor
 
-    @property
-    def shard_mesh(self):
+    def _init_shard_mesh(self) -> DeviceMesh:
         mesh = self.mesh_info.mesh
         if mesh.ndim == 1:
             return mesh
@@ -844,6 +900,10 @@ def shard_mesh(self):
             raise AssertionError("Expected mesh_dim_names to not be None")
         return mesh[mesh.mesh_dim_names[-1]]
 
+    @property
+    def shard_mesh(self):
+        return self._shard_mesh
+
     @property
     def shard_mesh_from_root(self):
         return self.shard_mesh
diff --git a/torch/distributed/fsdp/_fully_shard/_fsdp_param_group.py b/torch/distributed/fsdp/_fully_shard/_fsdp_param_group.py
index 14046f707424a..be453e5cffb52 100644
--- a/torch/distributed/fsdp/_fully_shard/_fsdp_param_group.py
+++ b/torch/distributed/fsdp/_fully_shard/_fsdp_param_group.py
@@ -1,17 +1,20 @@
 # mypy: allow-untyped-defs
+from __future__ import annotations
+
 import contextlib
 import logging
-from collections.abc import Callable
-from typing import Any, cast, NamedTuple
+from typing import Any, cast, Literal, NamedTuple, TYPE_CHECKING
 
 import torch
 import torch.distributed as dist
 import torch.nn as nn
 from torch.distributed.device_mesh import _get_device_handle
-from torch.distributed.fsdp._common_utils import _named_parameters_with_duplicates
-from torch.distributed.tensor import Shard
+from torch.distributed.fsdp._common_utils import (
+    _named_parameters_with_duplicates,
+    collect_grad_tensors,
+    replace_grad_tensors,
+)
 from torch.profiler import record_function
-from torch.utils._pytree import tree_flatten, tree_unflatten
 from torch.utils.hooks import RemovableHandle
 
 from ._fsdp_api import CPUOffloadPolicy, MixedPrecisionPolicy, OffloadPolicy
@@ -26,19 +29,26 @@
     ProcessGroupAllocAllGather,
     ProcessGroupAllocReduceScatter,
     ReduceScatter,
+    SymmMemAllGather,
+    SymmMemReduceScatter,
 )
 from ._fsdp_common import (
-    compiled_autograd_enabled,
+    _dynamo_disable,
     DataParallelMeshInfo,
     DDPMeshInfo,
     FSDPMeshInfo,
     HSDPMeshInfo,
     is_bw,
+    ShardPlacementFnResult,
     TrainingState,
 )
 from ._fsdp_param import alloc_storage, FSDPParam, ParamModuleInfo, ShardedState
 
 
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+
 logger = logging.getLogger("torch.distributed.fsdp.fully_shard")
 
 _ModuleToHandleDict = dict[nn.Module, RemovableHandle]  # for state dict
@@ -77,6 +87,12 @@ def lazy_init(self, device: torch.device):
         # Reduce-scatter stream gives separate execution "thread" for post-
         # backward logic like pre/post-gradient division and reduce-scatter
         self.reduce_scatter_stream = self.device_handle.Stream(priority=high_priority)
+        # The most recent post-reduce event across all param groups, recorded
+        # on reduce_scatter_stream (FSDP) or all_reduce_stream (HSDP). Since
+        # later events subsume earlier ones on the same stream, a single wait
+        # on the last event recorded on each stream used in the post backward
+        # hook covers all groups.
+        self._last_post_reduce_events: dict[torch.Stream, torch.Event] = dict()
         # Run the HSDP all-reduces concurrently with all-gather/reduce-scatter
         # since collectives use different network resources and can overlap
         # in the typical intra-node sharding / inter-node replication case
@@ -85,7 +101,12 @@ def lazy_init(self, device: torch.device):
         # tensors produced in one stream and used in another and accompanying
         # CUDA events for synchronization
         self.all_gather_state: AllGatherState | None = None
-        self.reduce_scatter_state: ReduceScatterState | None = None
+        self.reduce_scatter_states: list[ReduceScatterState] = []
+        # Effective cap on retained reduce_scatter_states, resolved from the
+        # per-group reduce_scatter_max_input_buffers in
+        # FSDPState._init_shared_state. It lives here, not per group, because
+        # reduce_scatter_states is a single shared list governed by one cap.
+        self.reduce_scatter_max_input_buffers: int = 1
         # Post-forward order for explicit backward prefetching
         self.post_forward_order: list[FSDPParamGroup] = []  # will cause ref cycles
 
@@ -131,7 +152,7 @@ def __init__(
         mesh_info: DataParallelMeshInfo,
         post_forward_mesh_info: FSDPMeshInfo | None,
         device: torch.device,
-        shard_placement_fn: Callable[[nn.Parameter], Shard | None] | None,
+        shard_placement_fn: Callable[[nn.Parameter], ShardPlacementFnResult] | None,
         mp_policy: MixedPrecisionPolicy,
         offload_policy: OffloadPolicy,
     ):
@@ -153,7 +174,6 @@ def __init__(
         ]
         self.mesh_info = mesh_info
         self.post_forward_mesh_info = post_forward_mesh_info
-        # pyrefly: ignore [read-only]
         self.device = device
         self.device_handle = _get_device_handle(device.type)
         self.mp_policy = mp_policy
@@ -180,6 +200,8 @@ def __init__(
 
         # - Communication and communication/computation overlap
         self.comm_ctx = FSDPCommContext()
+        self._param_group_index: int = 0
+        self._num_param_groups: int = 1
         # Group's indices in the shared post-forward order
         self._post_forward_indices: list[int] = []
         # Whether to reduce gradients at all (whether for FSDP or HSDP)
@@ -191,6 +213,12 @@ def __init__(
         # Whether to reshard parameters after backward (only useful for
         # gradient accumulation)
         self.reshard_after_backward: bool = True
+        # Per-group input for the reduce-scatter copy-in (chunk_cat) buffer
+        # cap-K, set via FSDPModule.set_reduce_scatter_max_input_buffers (see
+        # its docstring). FSDP keeps 1 by default. These per-group values are
+        # resolved into the single shared
+        # comm_ctx.reduce_scatter_max_input_buffers that post_backward reads.
+        self.reduce_scatter_max_input_buffers: int = 1
         # Optional custom factor for the gradient reduction op (e.g. to divide
         # by a factor other than the world size)
         self.gradient_divide_factor: float | None = None
@@ -263,11 +291,34 @@ def lazy_init(self):
             self._reset_sharded_params = True
         self._validate_no_meta_params()
         self._validate_cpu_offload_params()
+        self._validate_reduce_scatter_max_input_buffers()
         # Initialize mixed precision attributes lazily in case the user changes
         # the parameter dtypes after construction time but before forward
         self._init_mp_dtypes()
         self._register_state_dict_hooks()
 
+    def set_symm_mem(self, backend: Literal["NCCL"] = "NCCL") -> None:
+        if not isinstance(self._all_gather_comm, (DefaultAllGather | SymmMemAllGather)):
+            raise AssertionError(
+                "cannot call set_symm_mem() "
+                f"when all gather comm is custom: {self._all_gather_comm.__class__.__name__}"
+            )
+        self._all_gather_comm = SymmMemAllGather(
+            self._all_gather_process_group, backend
+        )
+        if not isinstance(
+            self._reduce_scatter_comm, (DefaultReduceScatter | SymmMemReduceScatter)
+        ):
+            raise AssertionError(
+                "cannot call set_symm_mem() "
+                f"when reduce scatter comm is custom: {self._reduce_scatter_comm.__class__.__name__}"
+            )
+        if self.force_sum_reduction_for_comms:
+            # As of NCCL 2.29.3, NCCL symmetric reduce-scatter only supports SUM reduction
+            self._reduce_scatter_comm = SymmMemReduceScatter(
+                self._reduce_scatter_process_group, backend
+            )
+
     def set_allocate_memory_from_process_group(self, enable: bool) -> None:
         """
         Whether to (try to) use the ProcessGroup's allocate_tensor method for
@@ -442,8 +493,7 @@ def reshard(self):
     def pre_forward(
         self, module: nn.Module, args: tuple[Any, ...], kwargs: dict[str, Any]
     ) -> tuple[tuple[Any, ...], dict[str, Any]]:
-        if not compiled_autograd_enabled():
-            logger.debug("%s", self._with_fqn("FSDP::pre_forward"))
+        logger.debug("%s", self._with_fqn("FSDP::pre_forward"))
         with record_function(self._with_fqn("FSDP::pre_forward")):
             self._training_state = TrainingState.FORWARD
             self.unshard(self.unshard_async_op)
@@ -452,16 +502,11 @@ def pre_forward(
             return args, kwargs
 
     def post_forward(self, module: nn.Module, input: Any, output: Any):
-        if not compiled_autograd_enabled():
-            logger.debug("%s", self._with_fqn("FSDP::post_forward"))
+        logger.debug("%s", self._with_fqn("FSDP::post_forward"))
         with record_function(self._with_fqn("FSDP::post_forward")):
-            if not compiled_autograd_enabled():
-                # for AC(fully_shard(model)), AC runs fsdp's _pre_forward
-                # it shouldn't change post_forward_order
-                if not is_bw():
-                    self.reshard()
-                    self._record_post_forward()
-            else:
+            # for AC(fully_shard(model)), AC runs fsdp's _pre_forward
+            # it shouldn't change post_forward_order
+            if not is_bw():
                 self.reshard()
                 self._record_post_forward()
             self._training_state = TrainingState.IDLE
@@ -474,30 +519,23 @@ def _record_post_forward(self) -> None:
         self.comm_ctx.post_forward_order.append(self)
         self._post_forward_indices.append(post_forward_index)
 
+    @_dynamo_disable
     def pre_backward(self, default_prefetch: bool, *unused: Any):
-        if (
-            compiled_autograd_enabled()
-            and self._training_state == TrainingState.PRE_BACKWARD
-        ):
-            # Traceable FSDP2 cannot trigger the param group's `post_backward` immediately after param usage;
-            # instead it relies on this to trigger the previously unexecuted `post_backward`.
-            self.post_backward()
         if self._training_state == TrainingState.PRE_BACKWARD:
             return
-        if not compiled_autograd_enabled():
-            logger.debug("%s", self._with_fqn("FSDP::pre_backward"))
+        logger.debug("%s", self._with_fqn("FSDP::pre_backward"))
         with record_function(self._with_fqn("FSDP::pre_backward")):
             self._training_state = TrainingState.PRE_BACKWARD
             self.unshard(self.unshard_async_op)  # no-op if prefetched
             self.wait_for_unshard()
-            if default_prefetch and not compiled_autograd_enabled():
+            if default_prefetch:
                 self._backward_prefetch()
 
+    @_dynamo_disable
     def post_backward(self, *unused: Any):
         # This method should be idempotent and safe to call even when this
         # FSDP parameter group was not used in backward (should be a no-op)
-        if not compiled_autograd_enabled():
-            logger.debug("%s", self._with_fqn("FSDP::post_backward"))
+        logger.debug("%s", self._with_fqn("FSDP::post_backward"))
         self._training_state = TrainingState.POST_BACKWARD
         with record_function(self._with_fqn("FSDP::post_backward_accumulate")):
             for fsdp_param in self.fsdp_params:
@@ -528,17 +566,29 @@ def post_backward(self, *unused: Any):
                     fsdp_param.unsharded_param.grad = None
             if self.reshard_after_backward:
                 self.reshard()
+        # Recycle prior modules' reduce-scatter input buffers, keeping at most
+        # `max_input_buffers` in flight: reclaim the oldest (wait on its
+        # reduce-scatter, then drop the keepalive ref that was deferring the
+        # allocator's reuse of its memory) until fewer than that remain, before
+        # this module's reduce appends one more. See
+        # set_reduce_scatter_max_input_buffers for the memory/overlap tradeoff.
+        # (Assumes backward fires groups N-1 first; if not, overlap degrades but
+        # correctness is preserved.)
+        max_input_buffers = self.comm_ctx.reduce_scatter_max_input_buffers
+        states = self.comm_ctx.reduce_scatter_states
+        if (
+            self._param_group_index == self._num_param_groups - 1
+            and len(states) >= max_input_buffers
+        ):
+            with record_function(f"FSDP::post_backward_rs_wait ({self._module_fqn})"):
+                while len(states) >= max_input_buffers:
+                    oldest = states.pop(0)
+                    if oldest.event is not None:
+                        self.device_handle.current_stream().wait_event(oldest.event)
+                    del oldest
         if len(fsdp_params_with_grad) == 0:
             return
         with record_function(self._with_fqn("FSDP::post_backward_reduce")):
-            if (
-                self.comm_ctx.reduce_scatter_state is not None
-                and self.comm_ctx.reduce_scatter_state.event is not None
-            ):
-                self.device_handle.current_stream().wait_event(
-                    self.comm_ctx.reduce_scatter_state.event
-                )
-            self.comm_ctx.reduce_scatter_state = None
             all_reduce_pg = (
                 self._all_reduce_process_group
                 if isinstance(self.mesh_info, DDPMeshInfo)
@@ -560,6 +610,7 @@ def post_backward(self, *unused: Any):
             (
                 reduce_scatter_input,
                 reduce_scatter_event,
+                post_reduce_stream,
                 self._post_reduce_event,
                 all_reduce_input,
                 all_reduce_event,
@@ -590,8 +641,11 @@ def post_backward(self, *unused: Any):
                 self._all_reduce_hook,
                 self.force_sum_reduction_for_comms,
             )
-            self.comm_ctx.reduce_scatter_state = ReduceScatterState(
-                reduce_scatter_input, reduce_scatter_event
+            self.comm_ctx._last_post_reduce_events[post_reduce_stream] = (
+                self._post_reduce_event
+            )
+            self.comm_ctx.reduce_scatter_states.append(
+                ReduceScatterState(reduce_scatter_input, reduce_scatter_event)
             )
             if all_reduce_input is not None:
                 if self.device.type != "cpu":
@@ -604,7 +658,11 @@ def post_backward(self, *unused: Any):
                 )
 
     def finalize_backward(self):
-        self._wait_for_post_backward()
+        for event in self.comm_ctx._last_post_reduce_events.values():
+            self.device_handle.current_stream().wait_event(event)
+        self.comm_ctx._last_post_reduce_events = dict()
+        self._post_reduce_event = None
+        self._all_reduce_state = None
         for fsdp_param in self.fsdp_params:
             if fsdp_param.grad_offload_event is not None:
                 fsdp_param.grad_offload_event.synchronize()
@@ -637,18 +695,40 @@ def _backward_prefetch(self) -> None:
                 # Can be cleared if running multiple `backward`s
                 return
             curr_index = self._post_forward_indices.pop()
-            if (target_index := curr_index - 1) < 0:
-                return
-            # Prefetch naively using the reverse post-forward order, which may
-            # have mistargeted prefetches if not all modules used in forward
-            # are used in this backward
-            # pyrefly: ignore [unbound-name]
-            target_fsdp_param_group = self.comm_ctx.post_forward_order[target_index]
-            self._prefetch_unshard(target_fsdp_param_group, "backward")
+            if self._num_param_groups > 1:
+                # Backward fires groups in reverse forward order:
+                # N-1, N-2, ..., 1, 0.  Index 1 is always the
+                # penultimate group regardless of N.  Prefetching here
+                # lets the next module's AG overlap with group 0's RS
+                # without holding unsharded params too long (as would
+                # happen if we prefetched from N-1).
+                if self._param_group_index != 1:
+                    return
+                # E.g. fully_shard(block, shard_placement_fn=...) creates two
+                # param groups per block (dense + moe), giving
+                # post_forward_order = [block0, block0.moe, block1, block1.moe].
+                # block1.moe walks back past block1 to prefetch block0.moe then block0.
+                curr_modules = self.modules
+                target_modules: tuple[nn.Module, ...] | None = None
+                for step in range(1, curr_index + 1):
+                    target = self.comm_ctx.post_forward_order[curr_index - step]
+                    if target.modules is curr_modules:
+                        continue
+                    if target_modules is None:
+                        target_modules = target.modules
+                    elif target.modules is not target_modules:
+                        break
+                    # Prefetch all groups of the target module in
+                    # reverse forward order (highest index first),
+                    # matching the explicit path in _pre_backward.
+                    self._prefetch_unshard(target, "backward")
+            elif curr_index > 0:
+                target = self.comm_ctx.post_forward_order[curr_index - 1]
+                self._prefetch_unshard(target, "backward")
 
     @staticmethod
     def _prefetch_unshard(
-        target_fsdp_param_group: "FSDPParamGroup", pass_type: str
+        target_fsdp_param_group: FSDPParamGroup, pass_type: str
     ) -> None:
         if pass_type == "backward":
             training_state = TrainingState.PRE_BACKWARD
@@ -708,31 +788,20 @@ def use_training_state(self, training_state: TrainingState):
     def _register_post_backward_hook(
         self, args: tuple[Any, ...], kwargs: dict[str, Any]
     ) -> tuple[tuple[Any, ...], dict[str, Any]]:
-        # Traceable FSDP2 relies on `root_post_backward_callback` to call each
-        # `FSDPParamGroup.post_backward`
-        if (not torch._dynamo.config.skip_fsdp_hooks) or compiled_autograd_enabled():
-            return args, kwargs
         if not torch.is_grad_enabled():
             return args, kwargs
-        args_list, args_spec = tree_flatten(args)
-        kwargs_list, kwargs_spec = tree_flatten(kwargs)
-        args_kwargs_list = list(args_list) + list(kwargs_list)
-        inp_tensor_indices: list[int] = []
-        inp_tensors: list[torch.Tensor] = []
-        for i, obj in enumerate(args_kwargs_list):
-            if torch.is_tensor(obj) and obj.requires_grad:
-                inp_tensor_indices.append(i)
-                inp_tensors.append(obj)
-        if len(inp_tensors) == 0:
-            return args, kwargs  # no tensors that require gradients
-        inp_tensors = RegisterPostBackwardFunction.apply(self, *inp_tensors)
-        for inp_tensor_idx, inp_tensor in zip(inp_tensor_indices, inp_tensors):
-            args_kwargs_list[inp_tensor_idx] = inp_tensor
-        args_list = args_kwargs_list[: len(args_list)]
-        kwargs_list = args_kwargs_list[len(args_list) :]
-        args = tree_unflatten(args_list, args_spec)
-        kwargs = tree_unflatten(kwargs_list, kwargs_spec)
-        return args, kwargs
+
+        # Collect all tensors that require gradients (including from dataclasses)
+        inp_tensors = collect_grad_tensors((args, kwargs))
+        if not inp_tensors:
+            return args, kwargs
+
+        # Apply RegisterPostBackwardFunction to all tensors at once
+        out_tensors = RegisterPostBackwardFunction.apply(self, *inp_tensors)
+
+        # Replace tensors in the structure (iterator order matches collect order)
+        new_args, new_kwargs = replace_grad_tensors((args, kwargs), iter(out_tensors))
+        return new_args, new_kwargs
 
     def _register_state_dict_hooks(self) -> None:
         num_pre_save_hooks = len(self._module_to_pre_save_state_dict_hook_handle)
@@ -793,7 +862,46 @@ def _reduce_scatter_process_group(self) -> dist.ProcessGroup:
             raise AssertionError(
                 f"Expected mesh_info to be FSDPMeshInfo, got {type(self.mesh_info)}"
             )
-        return self.mesh_info.shard_process_group
+        # Falls back to the shard PG (shared with all-gather) unless
+        # set_separate_reduce_scatter_group opted in to a dedicated PG.
+        return (
+            self.mesh_info.reduce_scatter_process_group
+            or self.mesh_info.shard_process_group
+        )
+
+    def _set_separate_reduce_scatter_group(
+        self, enable: bool, new_groups: dict[tuple[int, ...], dist.ProcessGroup]
+    ) -> None:
+        # Give reduce-scatter its own process group (separate from the shared
+        # all-gather/shard PG) so the two can overlap in backward, or reset to
+        # the shared group. ``new_groups`` caches groups by shard ranks within
+        # one set_separate_reduce_scatter_group call so meshes sharding over the
+        # same ranks share one communicator (typically one total) rather than one
+        # per mesh. HSDP ranks may create only their local shard subgroup, so
+        # use local synchronization to avoid waiting for nonmember ranks in
+        # new_group's post-init barrier.
+        mesh_info = self.mesh_info
+        if not isinstance(mesh_info, FSDPMeshInfo):
+            raise AssertionError(
+                f"Expected mesh_info to be FSDPMeshInfo, got {type(mesh_info)}"
+            )
+        if not enable:
+            mesh_info.reduce_scatter_process_group = None
+            return
+        ranks = tuple(dist.get_process_group_ranks(mesh_info.shard_process_group))
+        if ranks not in new_groups:
+            # Reuse the mesh's existing group if already enabled (idempotent),
+            # else create one for this rank set.
+            existing = mesh_info.reduce_scatter_process_group
+            if existing is not None:
+                new_groups[ranks] = existing
+            else:
+                new_groups[ranks] = dist.new_group(
+                    list(ranks),
+                    use_local_synchronization=True,
+                    group_desc="fsdp_reduce_scatter",
+                )
+        mesh_info.reduce_scatter_process_group = new_groups[ranks]
 
     @property
     def _all_reduce_process_group(self) -> dist.ProcessGroup:
@@ -805,7 +913,9 @@ def _all_reduce_process_group(self) -> dist.ProcessGroup:
 
     def _with_fqn(self, label: str) -> str:
         if self._module_fqn:
-            return f"{label} ({self._module_fqn})"
+            label = f"{label} ({self._module_fqn})"
+        if self._num_param_groups > 1 and isinstance(self.mesh_info, FSDPMeshInfo):
+            label = f"{label} [pg={self.mesh_info.shard_mesh_size}]"
         return label
 
     def __repr__(self):
@@ -841,6 +951,27 @@ def _validate_cpu_offload_params(self):
                 f"{[(fsdp_param._param_fqn, fsdp_param.sharded_param.device) for fsdp_param in fsdp_params_not_on_cpu]}\n"
             )
 
+    def _validate_reduce_scatter_max_input_buffers(self):
+        # Reject max_input_buffers > 1 with the symmetric-memory reduce-scatter
+        # comm: retaining >1 input buffers relies on the allocator handing out a
+        # fresh buffer while prior ones are ref-held, but the symmetric-memory
+        # pool reuses a single buffer, so retaining K would force K symmetric
+        # segments and can exhaust symmetric memory. Read the per-group cap (set
+        # in __init__) so this is safe to call before the comm context is lazily
+        # initialized (e.g. unshard before the first forward).
+        if self.reduce_scatter_max_input_buffers > 1 and isinstance(
+            self._reduce_scatter_comm, SymmMemReduceScatter
+        ):
+            raise ValueError(
+                "set_reduce_scatter_max_input_buffers(>1) is not supported with "
+                "the symmetric-memory reduce-scatter comm, but got "
+                f"{self.reduce_scatter_max_input_buffers} for module "
+                f"'{self._module_fqn}' using "
+                f"{self._reduce_scatter_comm.__class__.__name__}. Keep "
+                "max_input_buffers=1 with symmetric memory, or use the default "
+                "reduce-scatter comm."
+            )
+
 
 def _get_param_module_infos(
     params: list[nn.Parameter], modules: tuple[nn.Module, ...]
@@ -874,30 +1005,14 @@ def _get_param_module_infos(
 
 
 class RegisterPostBackwardFunction(torch.autograd.Function):
-    @staticmethod
-    def _assert_not_tracing_fsdp():
-        if compiled_autograd_enabled():
-            # TODO: Find a way to print the offending FSDP2 module.
-            msg = """\
-When Traceable FSDP2 is enabled, we should not be calling into `RegisterPostBackwardFunction`.
-Instead, we rely on the param group's next `pre_backward` hook to trigger its previously unexecuted
-`post_backward`, and we rely on FSDPState's `root_post_backward_callback` to trigger the resharding
-of any leftover unsharded param groups.
-If you are here, it means the forward part of this FSDP2 instance is not compiled, and you must also
-compile the forward part if you want to use Traceable FSDP2."""
-            torch._dynamo.comptime.comptime.print(msg)
-            raise RuntimeError(msg)
-
     @staticmethod
     # pyrefly: ignore [bad-override]
     def forward(ctx, param_group: FSDPParamGroup, *inputs: torch.Tensor):
         # All tensors in `inputs` should require gradient
-        RegisterPostBackwardFunction._assert_not_tracing_fsdp()
         ctx.param_group = param_group
         return inputs
 
     @staticmethod
     def backward(ctx, *grads: torch.Tensor):
-        RegisterPostBackwardFunction._assert_not_tracing_fsdp()
         ctx.param_group.post_backward()
         return (None,) + grads
diff --git a/torch/distributed/fsdp/_fully_shard/_fsdp_state.py b/torch/distributed/fsdp/_fully_shard/_fsdp_state.py
index 1db8ab7381edd..634a5ed191a68 100644
--- a/torch/distributed/fsdp/_fully_shard/_fsdp_state.py
+++ b/torch/distributed/fsdp/_fully_shard/_fsdp_state.py
@@ -16,16 +16,11 @@
     _State,
 )
 from torch.distributed.device_mesh import _get_device_handle
+from torch.distributed.fsdp._common_utils import collect_grad_tensors
 from torch.distributed.utils import _apply_to_tensors, _to_kwargs
-from torch.utils._pytree import tree_flatten
 
 from ._fsdp_api import MixedPrecisionPolicy
-from ._fsdp_common import (
-    _cast_fp_tensor,
-    compiled_autograd_enabled,
-    detect_compiled_autograd,
-    TrainingState,
-)
+from ._fsdp_common import _cast_fp_tensor, _dynamo_disable, TrainingState
 from ._fsdp_param_group import FSDPCommContext, FSDPParamGroup
 
 
@@ -57,28 +52,15 @@ def __init__(self) -> None:
         self.post_optim_event: torch.Event | None = None
 
 
-def disable_if_config_true(func):
-    @functools.wraps(func)
-    def fsdp_hook_wrapper(*args, **kwargs):
-        if torch._dynamo.config.skip_fsdp_hooks:
-            return torch._dynamo.disable(
-                func,
-                recursive=True,
-                reason="skipping FSDP hooks since torch._dynamo.config.skip_fsdp_hooks is set",
-            )(*args, **kwargs)
-        else:
-            return func(*args, **kwargs)
-
-    return fsdp_hook_wrapper
-
-
 class FSDPState(_State):
     # Name used in error messages; subclasses can override
     _state_name: str = "FSDP"
 
     def __init__(self) -> None:
         super().__init__()
-        self._fsdp_param_group: FSDPParamGroup | None = None
+        # Support multiple param groups for per-param mesh support.
+        # Each group has params with the same mesh_info.
+        self._fsdp_param_groups: list[FSDPParamGroup] = []
         self._is_root: bool | None = None  # root set during lazy init
         self._state_ctx = FSDPStateContext()
         self._comm_ctx = FSDPCommContext()
@@ -94,6 +76,27 @@ def _get_state_for_module(self, module: nn.Module) -> "FSDPState | None":
         """Get the state for a module. Subclasses can override to use different state getters."""
         return _get_module_fsdp_state(module)
 
+    @property
+    def _fsdp_param_group(self) -> FSDPParamGroup | None:
+        """
+        Returns the param group for backward compatibility.
+        This property is only valid when there is at most one param group.
+        For per-param mesh support with multiple param groups, use
+        ``_fsdp_param_groups`` instead.
+        """
+        if len(self._fsdp_param_groups) > 1:
+            group_fqns = [g._module_fqn for g in self._fsdp_param_groups]
+            raise AssertionError(
+                f"Expected at most 1 param group for backward compatibility, "
+                f"but got {len(self._fsdp_param_groups)} (fqns: {group_fqns}). "
+                f"Use `_fsdp_param_groups` (plural) to access all param groups "
+                f"when using per-param mesh via shard_placement_fn returning "
+                f"ShardPlacementResult."
+            )
+        if self._fsdp_param_groups:
+            return self._fsdp_param_groups[0]
+        return None
+
     # Define a separate init since `__init__` is called in the contract
     def init(
         self,
@@ -105,7 +108,6 @@ def init(
         for module in modules:
             _insert_module_state(module, self)
         self._modules = modules
-        # pyrefly: ignore [read-only]
         self._device = device
         self._device_handle = _get_device_handle(device.type)
         self._mp_policy = mp_policy
@@ -133,8 +135,7 @@ def _root_pre_forward(
         self._lazy_init()
         if self._state_ctx.iter_forward_root is not None:
             return args, kwargs
-        if not compiled_autograd_enabled():
-            logger.debug("FSDP::root_pre_forward")
+        logger.debug("FSDP::root_pre_forward")
         self._state_ctx.iter_forward_root = self
         with torch.profiler.record_function("FSDP::root_pre_forward"):
             # Wait for optimizer before implicitly prefetched all-gathers
@@ -174,7 +175,6 @@ def _lazy_init(self) -> None:
             raise RuntimeError(
                 f"{self._state_name} requires a single root module but got {self._modules}"
             )
-        detect_compiled_autograd()
         root_module = self._modules[0]
         visited_states: set[FSDPState] = set()
         for module_name, module in root_module.named_modules():
@@ -188,27 +188,62 @@ def _lazy_init(self) -> None:
                         "the root module first"
                     )
                 state._is_root = False
-            self._state_ctx.all_states.append(state)
+            # A single state can map to multiple modules (e.g.
+            # fully_shard([mod_a, mod_b, mod_c])), so dedup here.
+            if state not in visited_states:
+                self._state_ctx.all_states.append(state)
             visited_states.add(state)
-        if self._fsdp_param_group and self._auto_reshard_after_forward:
-            # For the root, do not reshard after forward since for training,
-            # the parameters would be freed and all-gathered immediately
-            self._fsdp_param_group.post_forward_mesh_info = None
+        # For the root, do not reshard after forward since for training,
+        # the parameters would be freed and all-gathered immediately
+        if self._auto_reshard_after_forward:
+            for fsdp_param_group in self._fsdp_param_groups:
+                fsdp_param_group.post_forward_mesh_info = None
         self._init_fqns()
         self._init_shared_state()
+        self._validate_no_duplicate_params()
         # Run parameter group lazy inits after initializing FQNs for improved
         # error messages
         for state in self._state_ctx.all_states:
-            if state._fsdp_param_group:
-                state._fsdp_param_group.lazy_init()
+            for fsdp_param_group in state._fsdp_param_groups:
+                fsdp_param_group.lazy_init()
+
+    def _validate_no_duplicate_params(self) -> None:
+        seen: set[int] = set()
+        for state in self._state_ctx.all_states:
+            for fsdp_param_group in state._fsdp_param_groups:
+                for fsdp_param in fsdp_param_group.fsdp_params:
+                    if fsdp_param._orig_param_uid in seen:
+                        raise ValueError(
+                            f"Parameter '{fsdp_param._param_fqn}' is shared with a "
+                            f"parameter already managed by another FSDP group. "
+                            f"For shared/tied parameters, use "
+                            f"fully_shard([module_a, module_b]) to place them in "
+                            f"the same FSDP group."
+                        )
+                    seen.add(fsdp_param._orig_param_uid)
 
     def _init_shared_state(self) -> None:
         self._comm_ctx.lazy_init(self._device)
         for state in self._state_ctx.all_states:
             state._state_ctx = self._state_ctx
             state._comm_ctx = self._comm_ctx
-            if fsdp_param_group := state._fsdp_param_group:
+            num_groups = len(state._fsdp_param_groups)
+            for i, fsdp_param_group in enumerate(state._fsdp_param_groups):
                 fsdp_param_group.comm_ctx = self._comm_ctx
+                fsdp_param_group._param_group_index = i
+                fsdp_param_group._num_param_groups = num_groups
+        # The reduce-scatter input-buffer cap governs the single shared
+        # reduce_scatter_states list, so resolve the per-group caps into one
+        # effective value (the most aggressive); a lower-cap group must not
+        # silently undo a higher-cap group's retention.
+        self._comm_ctx.reduce_scatter_max_input_buffers = max(
+            (
+                fsdp_param_group.reduce_scatter_max_input_buffers
+                for state in self._state_ctx.all_states
+                for fsdp_param_group in state._fsdp_param_groups
+            ),
+            default=1,
+        )
 
     def _init_fqns(self) -> None:
         """Sets module and parameter FQN attributes for debugging."""
@@ -216,30 +251,35 @@ def _init_fqns(self) -> None:
             raise AssertionError("Expected _is_root to be True")
         root_module = self._modules[0]
         param_to_fsdp_param: dict[nn.Parameter, FSDPParam] = {}
-        module_to_fsdp_param_group: dict[nn.Module, FSDPParamGroup] = {}
+        # Build a mapping from module to all its FSDPParamGroups (not just one)
+        module_to_fsdp_param_groups: dict[nn.Module, list[FSDPParamGroup]] = {}
         for state in self._state_ctx.all_states:
-            if fsdp_param_group := state._fsdp_param_group:
+            for fsdp_param_group in state._fsdp_param_groups:
                 for fsdp_param in fsdp_param_group.fsdp_params:
                     param_to_fsdp_param[fsdp_param.sharded_param] = fsdp_param
                 for module in fsdp_param_group.modules:
-                    module_to_fsdp_param_group[module] = fsdp_param_group
+                    if module not in module_to_fsdp_param_groups:
+                        module_to_fsdp_param_groups[module] = []
+                    module_to_fsdp_param_groups[module].append(fsdp_param_group)
         for param_name, param in root_module.named_parameters():
             if param in param_to_fsdp_param:
                 param_to_fsdp_param[param]._param_fqn = param_name
         for module_name, module in root_module.named_modules():
-            if module in module_to_fsdp_param_group:
-                module_fqn = module_to_fsdp_param_group[module]._module_fqn
-                if module_fqn is None:
-                    module_to_fsdp_param_group[module]._module_fqn = module_name
-                else:
-                    if not isinstance(module_fqn, str):
-                        raise AssertionError(
-                            f"Expected module_fqn to be str, got {type(module_fqn)}: {module_fqn}"
-                        )
-                    module_fqn += f", {module_name}"
-                    module_to_fsdp_param_group[module]._module_fqn = module_fqn
-
-    @disable_if_config_true
+            if module in module_to_fsdp_param_groups:
+                # Set FQN for all param groups associated with this module
+                for fsdp_param_group in module_to_fsdp_param_groups[module]:
+                    module_fqn = fsdp_param_group._module_fqn
+                    if module_fqn is None:
+                        fsdp_param_group._module_fqn = module_name
+                    else:
+                        if not isinstance(module_fqn, str):
+                            raise AssertionError(
+                                f"Expected module_fqn to be str, got {type(module_fqn)}: {module_fqn}"
+                            )
+                        module_fqn += f", {module_name}"
+                        fsdp_param_group._module_fqn = module_fqn
+
+    @_dynamo_disable
     def _pre_forward(
         self, module: nn.Module, args: tuple[Any, ...], kwargs: dict[str, Any]
     ) -> tuple[tuple[Any, ...], dict[str, Any]]:
@@ -249,9 +289,10 @@ def _pre_forward(
             # With nested FSDP and multiple forward passes before backward,
             # the params might have been resharded by a previous post_backward.
             # We need to ensure params are unsharded for AC recomputation.
-            if self._fsdp_param_group and not self._fsdp_param_group.is_unsharded:
-                self._fsdp_param_group.unshard()
-                self._fsdp_param_group.wait_for_unshard()
+            for fsdp_param_group in self._fsdp_param_groups:
+                if not fsdp_param_group.is_unsharded:
+                    fsdp_param_group.unshard()
+                    fsdp_param_group.wait_for_unshard()
             return args, kwargs
         self._training_state = TrainingState.FORWARD
         args, kwargs = self._root_pre_forward(module, args, kwargs)
@@ -264,21 +305,23 @@ def _pre_forward(
                     _apply_to_tensors(cast_fn, args),
                     _apply_to_tensors(cast_fn, kwargs),
                 )
-        if self._fsdp_param_group:
-            args, kwargs = self._fsdp_param_group.pre_forward(module, args, kwargs)
+        for fsdp_param_group in self._fsdp_param_groups:
+            args, kwargs = fsdp_param_group.pre_forward(module, args, kwargs)
         for fsdp_state in self._states_to_forward_prefetch:
-            if (target_param_group := fsdp_state._fsdp_param_group) is not None:
+            # Forward order (not reversed) to match forward execution order;
+            # contrast with reversed() in _pre_backward for backward order.
+            for target_param_group in fsdp_state._fsdp_param_groups:
                 FSDPParamGroup._prefetch_unshard(target_param_group, "forward")
         return args, kwargs
 
-    @disable_if_config_true
+    @_dynamo_disable
     def _post_forward(self, module: nn.Module, input: Any, output: Any) -> Any:
         # When composing with module-hook-based activation checkpointing, the
         # post-backward hook is responsible for the reshard
         if self._training_state == TrainingState.PRE_BACKWARD:
             return output
-        if self._fsdp_param_group:
-            output = self._fsdp_param_group.post_forward(module, input, output)
+        for fsdp_param_group in self._fsdp_param_groups:
+            output = fsdp_param_group.post_forward(module, input, output)
         output = self._register_pre_backward_hook(output)
         self._training_state = TrainingState.IDLE
         if self._state_ctx.iter_forward_root is self:
@@ -299,42 +342,49 @@ def _post_forward(self, module: nn.Module, input: Any, output: Any) -> Any:
                 )
         return output
 
+    @_dynamo_disable
     def _pre_backward(self, grad: torch.Tensor) -> torch.Tensor:
         self._training_state = TrainingState.PRE_BACKWARD
         self._register_root_post_backward_final_callback()
-        if self._fsdp_param_group:
-            default_prefetch = len(self._states_to_backward_prefetch) == 0
-            self._fsdp_param_group.pre_backward(default_prefetch)
+        default_prefetch = len(self._states_to_backward_prefetch) == 0
+        for fsdp_param_group in self._fsdp_param_groups:
+            fsdp_param_group.pre_backward(default_prefetch)
         for fsdp_state in self._states_to_backward_prefetch:
-            if (target_param_group := fsdp_state._fsdp_param_group) is not None:
+            # Reverse so higher-indexed groups are prefetched first,
+            # matching backward execution order (reverse of forward).
+            for target_param_group in reversed(fsdp_state._fsdp_param_groups):
                 FSDPParamGroup._prefetch_unshard(target_param_group, "backward")
         return grad
 
+    @_dynamo_disable
     def _root_post_backward_final_callback(self) -> None:
-        if not compiled_autograd_enabled():
-            logger.debug("FSDP::root_post_backward")
+        logger.debug("FSDP::root_post_backward")
         with torch.profiler.record_function("FSDP::root_post_backward_callback"):
             for state in self._state_ctx.all_states:
-                fsdp_param_group = state._fsdp_param_group
-                if (
-                    fsdp_param_group
-                    and fsdp_param_group._training_state != TrainingState.POST_BACKWARD
-                ):
-                    # Run post-backward in case forward inputs did not require
-                    # gradient so the autograd backward did not run
-                    fsdp_param_group.post_backward()
-                state._training_state = TrainingState.IDLE
-                if fsdp_param_group:
+                # Reverse so that the last param group (which gates the
+                # reduce-scatter wait/clear) fires first, matching the
+                # autograd backward order and preserving RS overlap for
+                # per-param-mesh modules whose inputs lack gradients.
+                for fsdp_param_group in reversed(state._fsdp_param_groups):
+                    if fsdp_param_group._training_state != TrainingState.POST_BACKWARD:
+                        # Run post-backward in case forward inputs did not require
+                        # gradient so the autograd backward did not run
+                        fsdp_param_group.post_backward()
                     fsdp_param_group._training_state = TrainingState.IDLE
+                state._training_state = TrainingState.IDLE
                 if self._state_ctx.is_last_backward:
                     state._finalize_backward()
             if self._state_ctx.is_last_backward:
                 self._comm_ctx.post_forward_order.clear()
-                if self._comm_ctx.reduce_scatter_state is not None:
-                    self._device_handle.current_stream().wait_event(
-                        self._comm_ctx.reduce_scatter_state.event
-                    )
-                    self._comm_ctx.reduce_scatter_state = None
+                # Wait on and release any retained reduce-scatter input buffers:
+                # the last module's (which no later module's rs_wait clears) and,
+                # when set_reduce_scatter_max_input_buffers retains more than
+                # one in flight, the rest. The compute stream (which reuses the
+                # memory) is ordered past each reduce-scatter first.
+                for rs_state in self._comm_ctx.reduce_scatter_states:
+                    if rs_state.event is not None:
+                        self._device_handle.current_stream().wait_event(rs_state.event)
+                self._comm_ctx.reduce_scatter_states.clear()
             self._state_ctx.post_backward_final_callback_queued = False
 
     def _finalize_backward(self) -> None:
@@ -350,16 +400,17 @@ def _finalize_backward(self) -> None:
             warning_once(logger, msg, stacklevel=2)
             # Clear since we want the next forward to run
             self._modules_to_run_forward.clear()
-        if self._fsdp_param_group:
-            self._fsdp_param_group.finalize_backward()
+        for fsdp_param_group in self._fsdp_param_groups:
+            fsdp_param_group.finalize_backward()
 
     def _register_pre_backward_hook(self, output: Any) -> Any:
         if not torch.is_grad_enabled():
             return output
-        flat_outputs, _ = tree_flatten(output)
-        for t in flat_outputs:
-            if torch.is_tensor(t) and t.requires_grad:
-                t.register_hook(self._pre_backward)
+        # output is the forward return value — pass directly without wrapping
+        # (unlike _register_post_backward_hook which wraps (args, kwargs))
+        tensors = collect_grad_tensors(output)
+        for t in tensors:
+            t.register_hook(self._pre_backward)
         return output
 
     def _register_root_post_backward_final_callback(self):
@@ -370,6 +421,37 @@ def _register_root_post_backward_final_callback(self):
             self._root_post_backward_final_callback
         )
 
+    def _reset_iter_state(self) -> None:
+        # Iteration-wide recovery after a mid-forward or mid-backward
+        # exception. Waits on in-flight collectives, reshards every param
+        # group, and clears per-iteration trackers so the next forward can
+        # start from a clean state. Any in-flight gradients (reduce-scatter
+        # results, HSDP partial reduce outputs, grad-accum state) are
+        # discarded: the failed iteration is treated as lost.
+        if self._is_root is False:
+            raise RuntimeError(
+                "reset_iter_state must be called on the root FSDP module"
+            )
+        current_stream = self._device_handle.current_stream()
+        if ag_state := self._comm_ctx.all_gather_state:
+            if ag_state.event is not None:
+                current_stream.wait_event(ag_state.event)
+            self._comm_ctx.all_gather_state = None
+        for rs_state in self._comm_ctx.reduce_scatter_states:
+            if rs_state.event is not None:
+                current_stream.wait_event(rs_state.event)
+        self._comm_ctx.reduce_scatter_states.clear()
+        for event in self._comm_ctx._last_post_reduce_events.values():
+            current_stream.wait_event(event)
+        self._comm_ctx._last_post_reduce_events.clear()
+        self._comm_ctx.post_forward_order.clear()
+        for state in self._state_ctx.all_states:
+            state._modules_to_run_forward.clear()
+            state._training_state = TrainingState.IDLE
+            for fsdp_param_group in state._fsdp_param_groups:
+                fsdp_param_group._reset_iter_state()
+        self._state_ctx.iter_forward_root = None
+        self._state_ctx.post_backward_final_callback_queued = False
 
 def _get_module_fsdp_state(module: nn.Module) -> FSDPState | None:
     state = _get_module_state(module)
@@ -391,14 +473,14 @@ def _register_group_forward_hooks(
     """
     modules_set = set(modules)
 
-    @disable_if_config_true
+    @_dynamo_disable
     @functools.wraps(pre_hook)
     def wrapped_pre_hook(*args: Any, **kwargs: Any):
         if len(modules_to_run) == 0:  # first to run
             modules_to_run.update(modules_set)
             return pre_hook(*args, **kwargs)
 
-    @disable_if_config_true
+    @_dynamo_disable
     def get_wrapped_post_hook(module: nn.Module):
         @functools.wraps(post_hook)
         def wrapped_post_hook(*args: Any, **kwargs: Any):
diff --git a/torch/distributed/fsdp/_fully_shard/_fully_shard.py b/torch/distributed/fsdp/_fully_shard/_fully_shard.py
index 64fdcf879ed4b..472e1fc326ee6 100644
--- a/torch/distributed/fsdp/_fully_shard/_fully_shard.py
+++ b/torch/distributed/fsdp/_fully_shard/_fully_shard.py
@@ -5,15 +5,21 @@
 
 import functools
 from contextlib import contextmanager
-from typing import Any, cast, NoReturn, overload, TYPE_CHECKING
+from typing import Any, cast, Literal, NoReturn, overload, TYPE_CHECKING
 from typing_extensions import deprecated
 
 import torch
 import torch.nn as nn
 from torch.distributed._composable import contract
 
-from ._fsdp_api import AllGather, MixedPrecisionPolicy, OffloadPolicy, ReduceScatter
-from ._fsdp_common import FSDPMeshInfo
+from ._fsdp_api import (
+    AllGather,
+    DataParallelMeshDims,
+    MixedPrecisionPolicy,
+    OffloadPolicy,
+    ReduceScatter,
+)
+from ._fsdp_common import FSDPMeshInfo, ShardPlacementFnResult
 from ._fsdp_init import (
     _apply_to_module,
     _get_device_from_mesh,
@@ -31,7 +37,7 @@
 if TYPE_CHECKING:
     from collections.abc import Callable, Iterable, Iterator
 
-    from torch.distributed.tensor import DeviceMesh, Shard
+    from torch.distributed.tensor import DeviceMesh
 
     from ._fsdp_param_group import FSDPParamGroup
 
@@ -60,10 +66,11 @@ def fully_shard(
     *,
     mesh: DeviceMesh | None = ...,
     reshard_after_forward: bool | int | None = ...,
-    shard_placement_fn: Callable[[nn.Parameter], Shard | None] | None = ...,
+    shard_placement_fn: Callable[[nn.Parameter], ShardPlacementFnResult] | None = ...,
     mp_policy: MixedPrecisionPolicy = ...,
     offload_policy: OffloadPolicy = ...,
     ignored_params: set[nn.Parameter] | None = ...,
+    dp_mesh_dims: DataParallelMeshDims | None = ...,
 ) -> FSDPModule: ...
 
 
@@ -74,10 +81,11 @@ def fully_shard(
     *,
     mesh: DeviceMesh | None = ...,
     reshard_after_forward: bool | int | None = ...,
-    shard_placement_fn: Callable[[nn.Parameter], Shard | None] | None = ...,
+    shard_placement_fn: Callable[[nn.Parameter], ShardPlacementFnResult] | None = ...,
     mp_policy: MixedPrecisionPolicy = ...,
     offload_policy: OffloadPolicy = ...,
     ignored_params: set[nn.Parameter] | None = ...,
+    dp_mesh_dims: DataParallelMeshDims | None = ...,
 ) -> list[FSDPModule]: ...
 
 
@@ -92,10 +100,11 @@ def fully_shard(
     *,
     mesh: DeviceMesh | None = None,
     reshard_after_forward: bool | int | None = None,
-    shard_placement_fn: Callable[[nn.Parameter], Shard | None] | None = None,
+    shard_placement_fn: Callable[[nn.Parameter], ShardPlacementFnResult] | None = None,
     mp_policy: MixedPrecisionPolicy = MixedPrecisionPolicy(),
     offload_policy: OffloadPolicy = OffloadPolicy(),
     ignored_params: set[nn.Parameter] | None = None,
+    dp_mesh_dims: DataParallelMeshDims | None = None,
 ):
     """
     Apply fully sharded data parallelism (FSDP) to ``module``, where FSDP
@@ -168,11 +177,20 @@ def fully_shard(
               between forward and backward, the registered parameters must be
               the sharded parameters. For ``False`` or an ``int``, this can be
               done by manually resharding via :meth:`reshard`.
-        shard_placement_fn (Optional[Callable[[nn.Parameter], Optional[Shard]]]):
-            This callable can be used to override the sharding placement for a
-            parameter to shard a parameter on a dimension other than dim-0. If
-            this callable returns a :class:`Shard` placement (not ``None``),
-            then FSDP will shard according to that placement (e.g. ``Shard(1)``).
+        shard_placement_fn (Optional[Callable[[nn.Parameter], Optional[Shard | ShardPlacementResult]]]):
+            This callable can be used to override the sharding placement and/or
+            mesh for a parameter. It can return:
+
+            - ``None``: Use default sharding (Shard(0)) on the mesh passed to
+              ``fully_shard``.
+            - :class:`Shard`: Shard the parameter on the specified dimension
+              using the mesh passed to ``fully_shard``.
+            - :class:`ShardPlacementResult`: Specify both the shard placement
+              and a custom :class:`FSDPMeshInfo`. This allows different
+              parameters to be sharded across different process groups, enabling
+              use cases like Mixture of Experts where expert params use a
+              different mesh than regular params.
+
             If sharding on a nonzero dim, we currently require even sharding,
             i.e. the tensor dim size on that dim must be divisible by the FSDP
             shard mesh size.
@@ -185,6 +203,12 @@ def fully_shard(
         ignored_params: Optional(Set[nn.Parameter]): The set of parameters to be
             ignored by FSDP. They will not be sharded, nor moved to the device
             during init, nor have their gradients reduced in backward.
+        dp_mesh_dims (Optional[DataParallelMeshDims]): When provided,
+            ``mesh`` is treated as the full SPMD mesh, and parameters should be
+            DTensors on this mesh with ``Replicate()`` on all DP dimensions.
+            The ``shard`` field names which dim(s) FSDP shards on (multiple
+            dims are flattened). The ``replicate`` field names the HSDP
+            replication dim(s) (multiple dims are flattened).
 
     Returns:
         FSDPModule: The module with FSDP applied (in-place).
@@ -192,21 +216,35 @@ def fully_shard(
     torch._C._log_api_usage_once("torch.distributed.fsdp.fully_shard")
     _validate_module(module, "fully_shard")
     mesh = mesh or _init_default_mesh()
-    _validate_mesh(mesh)
-    mesh_info = _get_mesh_info(mesh)
+    _validate_mesh(mesh, dp_mesh_dims)
+    mesh_info = _get_mesh_info(mesh, dp_mesh_dims)
     device = _get_device_from_mesh(mesh)
     auto_reshard_after_forward = reshard_after_forward is None
     # If the user does not provide ``reshard_after_forward``, we set it to True.
     # During lazy_init, we identify which module is the root and override its value to False
-    post_forward_mesh_info = _get_post_forward_mesh_info(
-        reshard_after_forward if not auto_reshard_after_forward else True,  # type: ignore[arg-type]
-        mesh_info,
-    )
+    if isinstance(mesh_info, FSDPMeshInfo):
+        if (
+            mesh_info.is_spmd_mesh
+            and not isinstance(reshard_after_forward, bool)
+            and isinstance(reshard_after_forward, int)
+        ):
+            raise NotImplementedError(
+                "reshard_after_forward as int is not yet supported with "
+                "SPMD mesh (dp_mesh_dims)"
+            )
+        post_forward_mesh_info = _get_post_forward_mesh_info(
+            reshard_after_forward if not auto_reshard_after_forward else True,  # type: ignore[arg-type]
+            mesh_info,
+        )
+    else:
+        # DDPMeshInfo: no sharding, so no post-forward resharding needed
+        post_forward_mesh_info = None
     arg_module, modules, managed_modules, params, buffers = _get_modules_and_states(
         module, device, ignored_params
     )
     state = fully_shard.state(modules[0])  # type: ignore[attr-defined]
     state.init(modules, device, mp_policy, auto_reshard_after_forward)
+
     _init_param_group(
         state,
         params,
@@ -217,6 +255,9 @@ def fully_shard(
         shard_placement_fn,
         mp_policy,
         offload_policy,
+        reshard_after_forward=reshard_after_forward
+        if not auto_reshard_after_forward
+        else True,
     )
 
     # For Dynamo
@@ -274,7 +315,7 @@ def reshard(self) -> None:
         module. This method is *not* recursive.
         """
         state = self._get_fsdp_state()
-        if fsdp_param_group := state._fsdp_param_group:
+        for fsdp_param_group in state._fsdp_param_groups:
             fsdp_param_group.reshard()
 
     def unshard(self, async_op: bool = False) -> UnshardHandle | None:
@@ -296,11 +337,12 @@ def unshard(self, async_op: bool = False) -> UnshardHandle | None:
             before pre-forward.
         """
         state = self._get_fsdp_state()
-        fsdp_param_group = state._fsdp_param_group
-        if fsdp_param_group is not None:
+        for fsdp_param_group in state._fsdp_param_groups:
             fsdp_param_group.lazy_init()
             fsdp_param_group.unshard(async_op=async_op)
-        handle = _UnshardHandleImpl(fsdp_param_group)
+        handle = _UnshardHandleImpl(
+            list(state._fsdp_param_groups) if state._fsdp_param_groups else None
+        )
         if async_op:
             return handle
         handle.wait()
@@ -336,7 +378,7 @@ def set_requires_gradient_sync(
         for module in modules:
             if isinstance(module, FSDPModule):
                 state = module._get_fsdp_state()
-                if fsdp_param_group := state._fsdp_param_group:
+                for fsdp_param_group in state._fsdp_param_groups:
                     fsdp_param_group.reduce_grads = requires_gradient_sync
                     fsdp_param_group.all_reduce_grads = requires_gradient_sync
 
@@ -353,7 +395,7 @@ def set_requires_all_reduce(
         for module in modules:
             if isinstance(module, FSDPModule):
                 state = module._get_fsdp_state()
-                if fsdp_param_group := state._fsdp_param_group:
+                for fsdp_param_group in state._fsdp_param_groups:
                     fsdp_param_group.all_reduce_grads = requires_all_reduce
 
     def set_reshard_after_forward(
@@ -383,8 +425,9 @@ def set_reshard_after_forward(
             if isinstance(module, FSDPModule):
                 state = module._get_fsdp_state()
                 state._auto_reshard_after_forward = False
-                if fsdp_param_group := state._fsdp_param_group:
-                    assert isinstance(fsdp_param_group.mesh_info, FSDPMeshInfo)
+                for fsdp_param_group in state._fsdp_param_groups:
+                    if not isinstance(fsdp_param_group.mesh_info, FSDPMeshInfo):
+                        raise AssertionError
                     fsdp_param_group.post_forward_mesh_info = (
                         _get_post_forward_mesh_info(
                             reshard_after_forward,
@@ -412,7 +455,7 @@ def set_reshard_after_backward(
         for module in modules:
             if isinstance(module, FSDPModule):
                 state = module._get_fsdp_state()
-                if fsdp_param_group := state._fsdp_param_group:
+                for fsdp_param_group in state._fsdp_param_groups:
                     fsdp_param_group.reshard_after_backward = reshard_after_backward
 
     def set_modules_to_forward_prefetch(self, modules: list[FSDPModule]) -> None:
@@ -465,7 +508,13 @@ def set_custom_all_gather(self, comm: AllGather) -> None:
             comm (AllGather): Custom all-gather communication.
         """
         state = self._get_fsdp_state()
-        if (fsdp_param_group := state._fsdp_param_group) is not None:
+        if len(state._fsdp_param_groups) > 1:
+            raise ValueError(
+                "set_custom_all_gather is not supported with multiple param "
+                "groups (from per-param mesh via shard_placement_fn). "
+                "The custom comm would be ambiguous across groups with different meshes."
+            )
+        for fsdp_param_group in state._fsdp_param_groups:
             fsdp_param_group._all_gather_comm = comm
 
     def set_custom_reduce_scatter(self, comm: ReduceScatter) -> None:
@@ -478,7 +527,13 @@ def set_custom_reduce_scatter(self, comm: ReduceScatter) -> None:
             comm (ReduceScatter): Custom reduce_scatter communication.
         """
         state = self._get_fsdp_state()
-        if (fsdp_param_group := state._fsdp_param_group) is not None:
+        if len(state._fsdp_param_groups) > 1:
+            raise ValueError(
+                "set_custom_reduce_scatter is not supported with multiple param "
+                "groups (from per-param mesh via shard_placement_fn). "
+                "The custom comm would be ambiguous across groups with different meshes."
+            )
+        for fsdp_param_group in state._fsdp_param_groups:
             fsdp_param_group._reduce_scatter_comm = comm
 
     def set_all_reduce_hook(
@@ -499,7 +554,13 @@ def set_all_reduce_hook(
                 all-reduce stream used by the native HSDP all-reduce.
         """
         state = self._get_fsdp_state()
-        if (fsdp_param_group := state._fsdp_param_group) is not None:
+        if len(state._fsdp_param_groups) > 1:
+            raise ValueError(
+                "set_all_reduce_hook is not supported with multiple param "
+                "groups (from per-param mesh via shard_placement_fn). "
+                "The hook would be ambiguous across groups with different meshes."
+            )
+        for fsdp_param_group in state._fsdp_param_groups:
             fsdp_param_group._all_reduce_hook = hook
             if stream is not None:
                 if fsdp_param_group._is_hsdp:
@@ -540,7 +601,7 @@ def set_gradient_divide_factor(self, factor: float) -> None:
             factor (float): Custom divide factor.
         """
         state = self._get_fsdp_state()
-        if (fsdp_param_group := state._fsdp_param_group) is not None:
+        for fsdp_param_group in state._fsdp_param_groups:
             fsdp_param_group.gradient_divide_factor = factor
 
     def set_force_sum_reduction_for_comms(self, enable: bool) -> None:
@@ -561,9 +622,125 @@ def set_force_sum_reduction_for_comms(self, enable: bool) -> None:
             enable (bool): Whether to only ever use ReduceOp.SUM for comms.
         """
         state = self._get_fsdp_state()
-        if (fsdp_param_group := state._fsdp_param_group) is not None:
+        for fsdp_param_group in state._fsdp_param_groups:
             fsdp_param_group.force_sum_reduction_for_comms = enable
 
+    def set_reduce_scatter_unused_params(
+        self, reduce_scatter_unused_params: bool, *, recurse: bool = True
+    ) -> None:
+        """
+        Sets whether to include zero gradients for parameters that did not
+        receive a gradient in backward. This is needed when different ranks
+        use different parameters due to conditional control flow (e.g.
+        multi-modal models, mixture of experts), causing mismatched
+        reduce-scatter collectives. Similar to DDP's
+        ``find_unused_parameters``.
+
+        Args:
+            reduce_scatter_unused_params (bool): Whether to include zero
+                gradients for unused parameters in gradient reduction.
+            recurse (bool): Whether to set for all FSDP submodules or just
+                the passed-in module.
+        """
+        self_module = cast(nn.Module, self)
+        modules = list(self_module.modules()) if recurse else [self_module]
+        for module in modules:
+            if isinstance(module, FSDPModule):
+                state = module._get_fsdp_state()
+                for fsdp_param_group in state._fsdp_param_groups:
+                    fsdp_param_group.reduce_scatter_unused_params = (
+                        reduce_scatter_unused_params
+                    )
+
+    def set_reduce_scatter_max_input_buffers(
+        self, max_input_buffers: int, *, recurse: bool = True
+    ) -> None:
+        """
+        Sets how many gradient reduce-scatter input buffers may be in flight at
+        once -- the copy-in (``chunk_cat``) buffer cap-K (experimental).
+
+        FSDP keeps **1** such buffer in flight by default, so the compute stream
+        must wait on the previous reduce-scatter before the next copy-in can
+        reuse that buffer. When the reduce-scatter is exposed (communication
+        slower than the backward compute meant to hide it), that recycle wait
+        stalls the compute stream. Raising the cap lets the next copy-in write a
+        **fresh** buffer instead of waiting -- removing the stall -- at the cost
+        of extra peak memory for the retained buffers. The copy-in stays on the
+        compute stream; there is no extra stream and no ``record_stream``. This
+        helps only when the reduce-scatter is exposed.
+
+        Args:
+            max_input_buffers (int): Max reduce-scatter input buffers retained in
+                flight (the memory<->overlap dial); must be ``>= 1``. ``1`` is
+                FSDP's default behavior (single buffer; the exposed-RS stall). A
+                small value (e.g. ``2``) bounds peak memory and adds no stall as
+                long as it is ``>=`` the reduce-scatter pipeline depth (otherwise
+                an exposed reduce-scatter trades back a tail stall); a larger
+                value retains more buffers for deeper overlap at higher peak
+                memory.
+            recurse (bool): Whether to set for all FSDP submodules or just the
+                passed-in module.
+        """
+        # bool is an int subclass; reject it so True does not silently mean 1.
+        if isinstance(max_input_buffers, bool) or not isinstance(
+            max_input_buffers, int
+        ):
+            raise TypeError(
+                "max_input_buffers must be an int, got "
+                f"{type(max_input_buffers).__name__}"
+            )
+        if max_input_buffers < 1:
+            raise ValueError(
+                f"max_input_buffers must be a positive int, got {max_input_buffers}"
+            )
+        self_module = cast(nn.Module, self)
+        modules = list(self_module.modules()) if recurse else [self_module]
+        for module in modules:
+            if isinstance(module, FSDPModule):
+                state = module._get_fsdp_state()
+                for fsdp_param_group in state._fsdp_param_groups:
+                    fsdp_param_group.reduce_scatter_max_input_buffers = (
+                        max_input_buffers
+                    )
+
+    def set_separate_reduce_scatter_group(
+        self, enable: bool = True, *, recurse: bool = True
+    ) -> None:
+        """
+        Enables (or disables) running gradient reduce-scatter on its own process
+        group so it can overlap with all-gather in the backward pass
+        (experimental).
+
+        By default FSDP runs all-gather and reduce-scatter on separate CUDA
+        streams but through the **same** process group -- one NCCL communicator,
+        which processes one collective at a time and so serializes them on the
+        wire. When enabled, FSDP creates a dedicated process group over the shard
+        ranks (``dist.new_group(..., use_local_synchronization=True)``) -- one
+        per distinct set of shard ranks, typically a single communicator -- so
+        the two collectives can progress concurrently when the network can
+        sustain it. This is collective for each shard rank set: like other FSDP
+        comm setup, call it consistently across ranks using this FSDP mesh.
+
+        Args:
+            enable (bool): ``True`` (default) gives reduce-scatter its own
+                process group; ``False`` resets it to the shared shard/all-gather
+                group.
+            recurse (bool): Whether to set for all FSDP submodules or just the
+                passed-in module.
+        """
+        self_module = cast(nn.Module, self)
+        modules = list(self_module.modules()) if recurse else [self_module]
+        # Cache created groups by shard ranks so meshes that shard over the same
+        # ranks share one communicator (typically one total), not one per mesh.
+        new_groups: dict = {}
+        for module in modules:
+            if isinstance(module, FSDPModule):
+                state = module._get_fsdp_state()
+                for fsdp_param_group in state._fsdp_param_groups:
+                    fsdp_param_group._set_separate_reduce_scatter_group(
+                        enable, new_groups
+                    )
+
     def set_unshard_in_backward(self, unshard_in_backward: bool) -> None:
         """
         Sets whether the FSDP module's parameters need to be unsharded in
@@ -572,7 +749,7 @@ def set_unshard_in_backward(self, unshard_in_backward: bool) -> None:
         backward computation (e.g. embedding).
         """
         state = self._get_fsdp_state()
-        if (fsdp_param_group := state._fsdp_param_group) is not None:
+        for fsdp_param_group in state._fsdp_param_groups:
             fsdp_param_group.unshard_in_backward = unshard_in_backward
 
     def set_allocate_memory_from_process_group_for_comm(self, enable: bool) -> None:
@@ -593,9 +770,43 @@ def set_allocate_memory_from_process_group_for_comm(self, enable: bool) -> None:
             enable (bool): Whether to turn on ProcessGroup allocation.
         """
         state = self._get_fsdp_state()
-        if (fsdp_param_group := state._fsdp_param_group) is not None:
+        for fsdp_param_group in state._fsdp_param_groups:
             fsdp_param_group.set_allocate_memory_from_process_group(enable)
 
+    def set_symm_mem_for_comm(self, backend: Literal["NCCL"] = "NCCL") -> None:
+        """
+        Sets the symmetric memory (``symm_mem``) backend for allocating the
+        staging buffers used in all-gather collectives. This allows NCCL to use
+        optimized all-gather implementations via symmetric memory. Such
+        optimization may depend on the topology of the system.  For single node,
+        Copy Engine All-Gather may be used. For multi-node, Symmetric Kernel
+        All-Gather may be used.
+
+        To enable Copy Engine All-Gather, you need to set the NCCL process group
+        with the zero-CTA policy.
+        ```python
+        opts = dist.ProcessGroupNCCL.Options()
+        opts.config.cta_policy = dist.ProcessGroupNCCL.NCCL_CTA_POLICY_ZERO
+        dist.init_process_group(backend="nccl", pg_options=opts, device_id=device)
+        ```
+        Alternatively, you can set the environment variable `NCCL_CTA_POLICY` to 2.
+        ```bash
+        export NCCL_CTA_POLICY=2
+        ```
+        For more details, see [Copy Engine
+        Collectives](https://docs.pytorch.org/docs/2.11/symmetric_memory.html#copy-engine-collectives).
+
+        This cannot be used together with :meth:`set_custom_all_gather` or
+        :meth:`set_custom_reduce_scatter`.
+
+        Args:
+            backend (str): The symmetric memory backend to use. Defaults to
+                ``"NCCL"``. Currently, only ``"NCCL"`` is supported.
+        """
+        state = self._get_fsdp_state()
+        for fsdp_param_group in state._fsdp_param_groups:
+            fsdp_param_group.set_symm_mem(backend)
+
     def _set_unshard_async_op(self, async_op: bool):
         """
         Sets whether to use ``async_op=True`` or ``False`` for the pre-forward
@@ -612,7 +823,7 @@ def _set_unshard_async_op(self, async_op: bool):
         for module in self_module.modules():
             if isinstance(module, FSDPModule):
                 state = module._get_fsdp_state()
-                if fsdp_param_group := state._fsdp_param_group:
+                for fsdp_param_group in state._fsdp_param_groups:
                     fsdp_param_group.unshard_async_op = async_op
 
     def _get_fsdp_state(self) -> FSDPState:
@@ -625,13 +836,14 @@ def _apply(self, *args: Any, **kwargs: Any) -> Any:
         self.reshard()
         ret = super()._apply(*args, **kwargs)  # type: ignore[misc]
         state = self._get_fsdp_state()
-        if not (fsdp_param_group := state._fsdp_param_group):
+        if not state._fsdp_param_groups:
             return ret
         # TODO: Remove this padding logic once DTensor pads the local tensor:
         # https://github.com/pytorch/pytorch/issues/113045
         with torch.no_grad():
-            for fsdp_param in fsdp_param_group.fsdp_params:
-                fsdp_param.reset_sharded_param()
+            for fsdp_param_group in state._fsdp_param_groups:
+                for fsdp_param in fsdp_param_group.fsdp_params:
+                    fsdp_param.reset_sharded_param()
         return ret
 
 
@@ -649,14 +861,15 @@ def wait(self) -> None:
 
 
 class _UnshardHandleImpl(UnshardHandle):
-    def __init__(self, fsdp_param_group: FSDPParamGroup | None):
-        self._fsdp_param_group = fsdp_param_group
+    def __init__(self, fsdp_param_groups: list[FSDPParamGroup] | None):
+        self._fsdp_param_groups = fsdp_param_groups
 
     def wait(self):
-        if self._fsdp_param_group is not None:
-            self._fsdp_param_group.wait_for_unshard()
+        if self._fsdp_param_groups is not None:
+            for fsdp_param_group in self._fsdp_param_groups:
+                fsdp_param_group.wait_for_unshard()
             # Avoid keeping a reference
-            self._fsdp_param_group = None
+            self._fsdp_param_groups = None
 
 
 def register_fsdp_forward_method(module: nn.Module, method_name: str) -> None:
@@ -721,7 +934,7 @@ def share_comm_ctx(modules: list[FSDPModule]) -> None:
     comm_ctx = fsdp_states[0]._comm_ctx
     for fsdp_state in fsdp_states[1:]:
         fsdp_state._comm_ctx = comm_ctx
-        if fsdp_param_group := fsdp_state._fsdp_param_group:
+        for fsdp_param_group in fsdp_state._fsdp_param_groups:
             fsdp_param_group.comm_ctx = comm_ctx
 
 
diff --git a/torch/distributed/fsdp/_init_utils.py b/torch/distributed/fsdp/_init_utils.py
index 23e353931dfeb..349ccfd767207 100644
--- a/torch/distributed/fsdp/_init_utils.py
+++ b/torch/distributed/fsdp/_init_utils.py
@@ -4,7 +4,7 @@
 import os
 import warnings
 from collections.abc import Callable, Generator, Iterable, Iterator
-from typing import Any, no_type_check, Optional, TYPE_CHECKING
+from typing import Any, no_type_check, TYPE_CHECKING
 
 import torch
 import torch.distributed as dist
@@ -12,6 +12,7 @@
 import torch.distributed.fsdp._traversal_utils as traversal_utils
 import torch.distributed.fsdp.fully_sharded_data_parallel as fsdp_file
 import torch.nn as nn
+from torch._opaque_base import OpaqueBase
 from torch.distributed.algorithms._comm_hooks import default_hooks
 from torch.distributed.device_mesh import DeviceMesh
 from torch.distributed.distributed_c10d import _get_default_group
@@ -61,7 +62,7 @@
 # Specification of process groups for hybrid sharding strategies.
 HybridShardProcessGroupType = tuple[dist.ProcessGroup, dist.ProcessGroup]
 # Overall specification of process group.
-ProcessGroupType = Optional[dist.ProcessGroup | HybridShardProcessGroupType]
+ProcessGroupType = dist.ProcessGroup | HybridShardProcessGroupType | None
 
 
 # TODO (awgu): Refactor this later
@@ -616,13 +617,26 @@ def _init_param_handle_from_module(
     managed_params = list(_get_orig_params(fully_sharded_module, state._ignored_params))
     _verify_managed_params(fully_sharded_module, managed_params)
     if sync_module_states:
-        _sync_module_params_and_buffers(
-            fully_sharded_module, managed_params, state.process_group
-        )
         if state.sharding_strategy in HYBRID_SHARDING_STRATEGIES:
+            # Broadcast inter-node first, then intra-node. The inter-node
+            # broadcast propagates rank 0's states to each node's local
+            # rank 0, so the subsequent intra-node broadcast has the
+            # correct source values on every node. Reversing this order
+            # causes local rank 0 on non-source nodes to broadcast
+            # uninitialized states (e.g. from meta-device materialization).
             _sync_module_params_and_buffers(
                 fully_sharded_module, managed_params, state._inter_node_pg
             )
+            # _sync_module_params_and_buffers marks each buffer with
+            # FSDP_SYNCED=True to avoid redundant syncs in nested
+            # wrapping. Reset the flag here so the intra-node broadcast
+            # below also includes buffers.
+            for buffer in fully_sharded_module.buffers():
+                if hasattr(buffer, FSDP_SYNCED):
+                    setattr(buffer, FSDP_SYNCED, False)
+        _sync_module_params_and_buffers(
+            fully_sharded_module, managed_params, state.process_group
+        )
     _init_param_handle_from_params(state, managed_params, fully_sharded_module)
     return state
 
@@ -1114,8 +1128,16 @@ def _sync_module_params_and_buffers(
                 # NOTE: Here we assume no nested subclasses, at most one level of subclass
                 # in both model's buffers and params
                 attrs, _ = detached_buffer.__tensor_flatten__()  # type: ignore[attr-defined]
-                inner_buffers = [getattr(detached_buffer, attr) for attr in attrs]
-                module_states.extend(inner_buffers)
+                for attr in attrs:
+                    match getattr(detached_buffer, attr):
+                        case torch.Tensor() as v:
+                            module_states.append(v)
+                        case OpaqueBase():
+                            pass
+                        case unexpected:
+                            raise AssertionError(
+                                f"expected Tensor or OpaqueBase, got {type(unexpected)}"
+                            )
             else:
                 module_states.append(detached_buffer)
 
@@ -1123,8 +1145,16 @@ def _sync_module_params_and_buffers(
         detached_param = param.detach()
         if is_traceable_wrapper_subclass(detached_param):
             attrs, _ = detached_param.__tensor_flatten__()  # type: ignore[attr-defined]
-            inner_params = [getattr(detached_param, attr) for attr in attrs]
-            module_states.extend(inner_params)
+            for attr in attrs:
+                match getattr(detached_param, attr):
+                    case torch.Tensor() as v:
+                        module_states.append(v)
+                    case OpaqueBase():
+                        pass
+                    case unexpected:
+                        raise AssertionError(
+                            f"expected Tensor or OpaqueBase, got {type(unexpected)}"
+                        )
         else:
             module_states.append(detached_param)
 
diff --git a/torch/distributed/fsdp/_optim_utils.py b/torch/distributed/fsdp/_optim_utils.py
index d2df57ee3dec6..b3e26d63ee276 100644
--- a/torch/distributed/fsdp/_optim_utils.py
+++ b/torch/distributed/fsdp/_optim_utils.py
@@ -988,7 +988,6 @@ def _get_param_id_to_param_from_optim_input(
     if optim_input is None:
         return dict(enumerate(model.parameters()))
     try:
-        # pyrefly: ignore [no-matching-overload]
         # pyrefly: ignore [redundant-cast]
         params = cast(list[nn.Parameter], list(optim_input))
     except TypeError as e:
@@ -2105,9 +2104,9 @@ def _set_optim_use_dtensor(
         state_dict_type = state_dict_settings.state_dict_type
         if state_dict_type == StateDictType.LOCAL_STATE_DICT:
             raise RuntimeError(
-                "Found state_dict_type LOCAL_STATE_DICT.",
-                "DeviceMesh is not compatible with LOCAL_STATE_DICT.",
-                "Please set state_dict_type to SHARDED_STATE_DICT to get DTensor state_dict.",
+                "Found state_dict_type LOCAL_STATE_DICT. "
+                "DeviceMesh is not compatible with LOCAL_STATE_DICT. "
+                "Please set state_dict_type to SHARDED_STATE_DICT to get DTensor state_dict."
             )
         else:
             state_dict_settings.optim_state_dict_config._use_dtensor = True
diff --git a/torch/distributed/fsdp/_runtime_utils.py b/torch/distributed/fsdp/_runtime_utils.py
index 42a07ebc4e653..31c23b06705ac 100644
--- a/torch/distributed/fsdp/_runtime_utils.py
+++ b/torch/distributed/fsdp/_runtime_utils.py
@@ -600,7 +600,7 @@ def _root_pre_forward(
             args_tuple, kwargs_tuple = _to_kwargs(
                 args, kwargs, state.compute_device, False
             )
-        args = args_tuple[0] if args_tuple else tuple()
+        args = args_tuple[0] if args_tuple else ()
         kwargs = kwargs_tuple[0] if kwargs_tuple else {}
 
         return _root_cast_forward_input(state, module, args, kwargs)
diff --git a/torch/distributed/fsdp/_state_dict_utils.py b/torch/distributed/fsdp/_state_dict_utils.py
index ec648ced837e1..9d9d5fec572f5 100644
--- a/torch/distributed/fsdp/_state_dict_utils.py
+++ b/torch/distributed/fsdp/_state_dict_utils.py
@@ -801,9 +801,9 @@ def _set_use_dtensor(fsdp_state: _FSDPState) -> None:
         state_dict_type = fsdp_state._state_dict_type
         if state_dict_type == StateDictType.LOCAL_STATE_DICT:
             raise RuntimeError(
-                "Found state_dict_type LOCAL_STATE_DICT",
-                "DeviceMesh is not compatible with LOCAL_STATE_DICT.",
-                "Please set state_dict_type to SHARDED_STATE_DICT to get DTensor state_dict.",
+                "Found state_dict_type LOCAL_STATE_DICT. "
+                "DeviceMesh is not compatible with LOCAL_STATE_DICT. "
+                "Please set state_dict_type to SHARDED_STATE_DICT to get DTensor state_dict."
             )
         else:
             fsdp_state._state_dict_config._use_dtensor = True
diff --git a/torch/distributed/launcher/api.py b/torch/distributed/launcher/api.py
index 2adf5549fecf1..5e8140477b8a5 100644
--- a/torch/distributed/launcher/api.py
+++ b/torch/distributed/launcher/api.py
@@ -15,10 +15,17 @@
 
 import torch
 import torch.distributed.elastic.rendezvous.registry as rdzv_registry
-from torch._utils_internal import get_default_numa_options
+from torch._utils_internal import get_default_numa_options, justknobs_check
 from torch.distributed.elastic import events, metrics
 from torch.distributed.elastic.agent.server.api import WorkerSpec
-from torch.distributed.elastic.agent.server.local_elastic_agent import LocalElasticAgent
+from torch.distributed.elastic.agent.server.health_check_server import (
+    create_healthcheck_server,
+)
+from torch.distributed.elastic.agent.server.local_elastic_agent import (
+    _AliveCallbackProxy,
+    LocalElasticAgent,
+    TORCHELASTIC_HEALTH_CHECK_PORT,
+)
 from torch.distributed.elastic.multiprocessing import (
     DefaultLogsSpecs,
     LogsSpecs,
@@ -79,6 +86,9 @@ class LaunchConfig:
                            When enabled, LOCAL_RANK is set to 0 for all workers and
                            CUDA_VISIBLE_DEVICES is adjusted so each worker accesses its
                            assigned GPU at device index 0.
+        shutdown_timeout: Time in seconds to wait for graceful shutdown of workers before
+                        sending SIGKILL. Can also be set via TORCH_ELASTIC_SHUTDOWN_TIMEOUT
+                        environment variable. Defaults to 30 seconds.
 
 
     .. note::
@@ -109,6 +119,7 @@ class LaunchConfig:
     duplicate_stdout_filters: list[str] | None = None
     duplicate_stderr_filters: list[str] | None = None
     virtual_local_rank: bool = False
+    shutdown_timeout: int | None = None
 
     def __post_init__(self):
         default_timeout = 900
@@ -130,6 +141,16 @@ def __post_init__(self):
             self.numa_options = get_default_numa_options()
             logger.info("Using default numa options = %r", self.numa_options)
 
+        # Set shutdown_timeout from environment variable if not explicitly set
+        if self.shutdown_timeout is None:
+            self.shutdown_timeout = int(
+                os.environ.get("TORCH_ELASTIC_SHUTDOWN_TIMEOUT", "30")
+            )
+        elif self.shutdown_timeout < 0:
+            raise ValueError(
+                f"shutdown_timeout must be non-negative, got {self.shutdown_timeout}"
+            )
+
 
 class elastic_launch:
     """
@@ -276,6 +297,34 @@ def launch_agent(
     # Set the signals to handle in the environment variable
     os.environ["TORCHELASTIC_SIGNALS_TO_HANDLE"] = config.signals_to_handle
 
+    # Start health check server before rendezvous so TW sees a healthy
+    # thrift port during the potentially long MAST rendezvous store barrier
+    # (10-22+ min for large jobs).  The _AliveCallbackProxy returns
+    # time.time() until wired to the agent after construction.
+    health_check_server = None
+    alive_callback_proxy = None
+    healthcheck_port = os.getenv(TORCHELASTIC_HEALTH_CHECK_PORT)
+    if healthcheck_port is not None and justknobs_check(
+        "ai_infra/pytorch_distributed:torchelastic_enable_healthcheck_before_rendezvous",
+        default=False,
+    ):
+        try:
+            alive_callback_proxy = _AliveCallbackProxy()
+            health_check_server = create_healthcheck_server(
+                alive_callback=alive_callback_proxy,
+                port=int(healthcheck_port),
+                timeout=60,
+            )
+            health_check_server.start()
+            logger.info(
+                "Started early health check server on port %s before rendezvous",
+                healthcheck_port,
+            )
+        except Exception:
+            logger.warning("Failed to start early health check server", exc_info=True)
+            health_check_server = None
+            alive_callback_proxy = None
+
     spec = WorkerSpec(
         role=config.role,
         local_world_size=config.nproc_per_node,
@@ -299,8 +348,13 @@ def launch_agent(
         logs_specs=config.logs_specs,  # type: ignore[arg-type]
         start_method=config.start_method,
         log_line_prefix_template=config.log_line_prefix_template,
+        shutdown_timeout=config.shutdown_timeout,  # type: ignore[arg-type]
+        health_check_server=health_check_server,
     )
 
+    if alive_callback_proxy is not None:
+        alive_callback_proxy.set_delegate(agent._get_alive_time)
+
     shutdown_rdzv = True
     try:
         metrics.initialize_metrics(metrics.MetricsConfig(config.metrics_cfg))
diff --git a/torch/distributed/nn/api/remote_module.py b/torch/distributed/nn/api/remote_module.py
index 68cc9eb23aaf8..15dfbb8d3e03a 100644
--- a/torch/distributed/nn/api/remote_module.py
+++ b/torch/distributed/nn/api/remote_module.py
@@ -5,7 +5,7 @@
 import sys
 import types
 from collections.abc import Callable, Iterator, Mapping
-from typing import Any, TypeVar, Union
+from typing import Any, TypeVar
 from typing_extensions import Self
 
 import torch
@@ -21,7 +21,7 @@
 
 __all__ = ["RemoteModule"]
 
-_grad_t = Union[tuple[Tensor, ...], Tensor]
+_grad_t = tuple[Tensor, ...] | Tensor
 # See https://mypy.readthedocs.io/en/latest/generics.html#generic-methods-and-generic-self for the use
 # of `T` to annotate `self`. Many methods of `Module` return `self` and we want those return values to be
 # the type of the subclass, not the looser type of `Module`.
@@ -457,7 +457,8 @@ def extra_repr(self) -> str:  # type: ignore[return]
     def _prepare_init(self, remote_device_str: str) -> bool:
         """Prepare the initialization and returns whether to enable automatically moving CPU tensors to CUDA devices."""
         # Sanity check.
-        assert rpc._is_current_rpc_agent_set(), "RemoteModule only works in RPC."
+        if not rpc._is_current_rpc_agent_set():
+            raise AssertionError("RemoteModule only works in RPC.")
 
         remote_device = _remote_device(remote_device_str)
         self.on = (
@@ -579,31 +580,23 @@ def init_from_module_rref(
 
         remote_module = object.__new__(RemoteModule)
 
-        # pyrefly: ignore [missing-attribute]
         enable_moving_cpu_tensors_to_cuda = remote_module._prepare_init(remote_device)
 
         if _module_interface_cls is not None:
             # Users reply on this field to know if this generated RemoteModule is TorchScript-able.
-            # pyrefly: ignore [missing-attribute]
             remote_module.is_scriptable = True
 
-            # pyrefly: ignore [missing-attribute]
             remote_module._init_template(
                 _module_interface_cls, enable_moving_cpu_tensors_to_cuda
             )
         else:
-            # pyrefly: ignore [missing-attribute]
             remote_module.is_scriptable = False
-            # pyrefly: ignore [missing-attribute]
             remote_module.generated_methods = (
                 _NON_SCRIPTABLE_REMOTE_MODULE_MODULE._generated_methods
             )
-        # pyrefly: ignore [missing-attribute]
         remote_module.module_rref = module_rref
 
-        # pyrefly: ignore [missing-attribute]
         remote_module._install_generated_methods()
-        # pyrefly: ignore [missing-attribute]
         remote_module._check_attribute_picklability()
 
         return remote_module
@@ -706,11 +699,9 @@ def _remote_module_receiver(
     m.__dict__.update(serialized_remote_module._asdict())
 
     # Unpickling the attribute `module_rref` must invoke RRef's `_deserialize()` method.
-    # pyrefly: ignore [missing-attribute]
     m.module_rref = rpc.PyRRef._deserialize(m.module_rref)
 
     # Install generated methods when unpickled.
-    # pyrefly: ignore [missing-attribute]
     for method in m.generated_methods:
         method_name = method.__name__
         method = torch.jit.export(method)
diff --git a/torch/distributed/nn/functional.py b/torch/distributed/nn/functional.py
index 0a23c5310000f..c8a83fd999185 100644
--- a/torch/distributed/nn/functional.py
+++ b/torch/distributed/nn/functional.py
@@ -1,4 +1,6 @@
 # mypy: allow-untyped-defs
+import warnings
+
 import torch
 import torch.distributed as dist
 from torch.autograd import Function
@@ -9,6 +11,24 @@
 from torch.distributed import group, ReduceOp
 
 
+def _not_supported_under_compile(name, *, suggestion=None):
+    msg = (
+        f"torch.distributed.nn.functional.{name} is not supported under torch.compile."
+    )
+    if suggestion:
+        msg += f" Use {suggestion} instead."
+    raise RuntimeError(msg)
+
+
+def _deprecated(name, suggestion):
+    warnings.warn(
+        f"torch.distributed.nn.functional.{name} is deprecated, "
+        f"use {suggestion} instead.",
+        category=FutureWarning,
+        stacklevel=3,
+    )
+
+
 def broadcast(tensor, src, group=group.WORLD):
     """
     Broadcasts the tensor to the whole group.
@@ -26,6 +46,12 @@ def broadcast(tensor, src, group=group.WORLD):
         Tensor: Received tensor from the broadcast op.
 
     """
+    if torch.compiler.is_compiling():
+        _not_supported_under_compile(
+            "broadcast",
+            suggestion="torch.distributed._functional_collectives.broadcast",
+        )
+    _deprecated("broadcast", "torch.distributed._functional_collectives.broadcast")
     return _Broadcast.apply(src, group, tensor)
 
 
@@ -41,6 +67,8 @@ def gather(tensor, dst=0, group=group.WORLD):
     Returns:
         tuple[Tensor]: List of appropriately-sized tensors with the gathered data.
     """
+    if torch.compiler.is_compiling():
+        _not_supported_under_compile("gather")
     return _Gather.apply(dst, group, tensor)
 
 
@@ -61,6 +89,8 @@ def scatter(tensors, src=0, group=group.WORLD):
         Tensor: Output tensor from the scatter operation.
 
     """
+    if torch.compiler.is_compiling():
+        _not_supported_under_compile("scatter")
     return _Scatter.apply(src, group, *tensors)
 
 
@@ -82,6 +112,8 @@ def reduce(tensor, dst, op=ReduceOp.SUM, group=group.WORLD):
         Tensor: Output of the collective.
 
     """
+    if torch.compiler.is_compiling():
+        _not_supported_under_compile("reduce")
     return _Reduce.apply(dst, op, group, tensor)
 
 
@@ -101,6 +133,15 @@ def reduce_scatter(output, input_list, op=ReduceOp.SUM, group=group.WORLD):
         Tensor: Output of the collective.
 
     """
+    if torch.compiler.is_compiling():
+        _not_supported_under_compile(
+            "reduce_scatter",
+            suggestion="torch.distributed._functional_collectives.reduce_scatter_tensor",
+        )
+    _deprecated(
+        "reduce_scatter",
+        "torch.distributed._functional_collectives.reduce_scatter_tensor",
+    )
     return _Reduce_Scatter.apply(op, group, output, *input_list)
 
 
@@ -116,6 +157,14 @@ def all_gather(tensor, group=group.WORLD):
         tuple([Tensor]): Output of the collective.
 
     """
+    if torch.compiler.is_compiling():
+        _not_supported_under_compile(
+            "all_gather",
+            suggestion="torch.distributed._functional_collectives.all_gather_tensor",
+        )
+    _deprecated(
+        "all_gather", "torch.distributed._functional_collectives.all_gather_tensor"
+    )
     return _AllGather.apply(group, tensor)
 
 
@@ -152,6 +201,8 @@ def _all_gather_base(output_tensor, input_tensor, group=group.WORLD):
         is correctly sized.
 
     """
+    if torch.compiler.is_compiling():
+        _not_supported_under_compile("_all_gather_base")
     return _AllGatherBase.apply(output_tensor, input_tensor, group)
 
 
@@ -168,6 +219,8 @@ def all_to_all(output_tensor_list, input_tensor_list, group=group.WORLD):
         tuple([Tensor]): Output of the collective.
 
     """
+    if torch.compiler.is_compiling():
+        _not_supported_under_compile("all_to_all")
     return _AlltoAll.apply(group, output_tensor_list, *input_tensor_list)
 
 
@@ -197,6 +250,15 @@ def all_to_all_single(
         Tensor: Output of the collective.
 
     """
+    if torch.compiler.is_compiling():
+        _not_supported_under_compile(
+            "all_to_all_single",
+            suggestion="torch.distributed._functional_collectives.all_to_all_single",
+        )
+    _deprecated(
+        "all_to_all_single",
+        "torch.distributed._functional_collectives.all_to_all_single",
+    )
     return _AlltoAllSingle.apply(
         group, output, output_split_sizes, input_split_sizes, input
     )
@@ -220,6 +282,12 @@ def all_reduce(tensor, op=ReduceOp.SUM, group=group.WORLD):
         Tensor: Output of the collective
 
     """
+    if torch.compiler.is_compiling():
+        _not_supported_under_compile(
+            "all_reduce",
+            suggestion="torch.distributed._functional_collectives.all_reduce",
+        )
+    _deprecated("all_reduce", "torch.distributed._functional_collectives.all_reduce")
     return _AllReduce.apply(op, group, tensor)
 
 
@@ -277,7 +345,8 @@ class _Scatter(Function):
     def forward(ctx, src, group, *tensors):
         ctx.src = src
         ctx.group = group
-        assert all(t.size() == tensors[0].size() for t in tensors)
+        if not all(t.size() == tensors[0].size() for t in tensors):
+            raise AssertionError
         output = torch.zeros_like(tensors[0])
         if dist.get_rank(group=group) == src:
             dist.scatter(output, list(tensors), src, group=group)
diff --git a/torch/distributed/nn/jit/instantiator.py b/torch/distributed/nn/jit/instantiator.py
index 7bc7f359fe834..3fe5526ba415e 100644
--- a/torch/distributed/nn/jit/instantiator.py
+++ b/torch/distributed/nn/jit/instantiator.py
@@ -14,20 +14,23 @@
 
 
 def get_arg_return_types_from_interface(module_interface):
-    assert getattr(module_interface, "__torch_script_interface__", False), (
-        "Expect a TorchScript class interface decorated by @torch.jit.interface."
-    )
+    if not getattr(module_interface, "__torch_script_interface__", False):
+        raise AssertionError(
+            "Expect a TorchScript class interface decorated by @torch.jit.interface."
+        )
     qualified_name = torch._jit_internal._qualified_name(module_interface)
     cu = torch.jit._state._python_cu
     module_interface_c = cu.get_interface(qualified_name)
-    assert "forward" in module_interface_c.getMethodNames(), (
-        f"Expect forward in interface methods, while it has {module_interface_c.getMethodNames()}"
-    )
+    if "forward" not in module_interface_c.getMethodNames():
+        raise AssertionError(
+            f"Expect forward in interface methods, while it has {module_interface_c.getMethodNames()}"
+        )
     method_schema = module_interface_c.getMethod("forward")
 
     arg_str_list = []
     arg_type_str_list = []
-    assert method_schema is not None
+    if method_schema is None:
+        raise AssertionError
     for argument in method_schema.arguments:
         arg_str_list.append(argument.name)
 
@@ -44,7 +47,8 @@ def get_arg_return_types_from_interface(module_interface):
     arg_type_str_list = arg_type_str_list[1:]  # Remove "self".
     arg_types_str = ", ".join(arg_type_str_list)
 
-    assert len(method_schema.returns) == 1
+    if len(method_schema.returns) != 1:
+        raise AssertionError
     argument = method_schema.returns[0]
     return_type_str = str(argument.type)
 
@@ -102,7 +106,8 @@ def _do_instantiate_remote_module_template(
     spec = importlib.util.spec_from_loader(
         generated_module_name, loader, origin="torch-git"
     )
-    assert spec is not None
+    if spec is None:
+        raise AssertionError
     module = importlib.util.module_from_spec(spec)
     sys.modules[generated_module_name] = module
     loader.exec_module(module)
diff --git a/torch/distributed/optim/named_optimizer.py b/torch/distributed/optim/named_optimizer.py
index a8432e198a083..2298f73505b48 100644
--- a/torch/distributed/optim/named_optimizer.py
+++ b/torch/distributed/optim/named_optimizer.py
@@ -106,8 +106,10 @@ def __init__(
     def _param_groups_check(self) -> None:
         if self.param_groups is not None:
             for param_group in self.param_groups:
-                assert isinstance(param_group, dict), "param group must be a dict"
-                assert "params" in param_group, "param group must contain key params"
+                if not isinstance(param_group, dict):
+                    raise AssertionError("param group must be a dict")
+                if "params" not in param_group:
+                    raise AssertionError("param group must contain key params")
                 params = param_group["params"]
                 if isinstance(params, torch.Tensor):
                     params = [params]
@@ -218,7 +220,8 @@ def load_state_dict(self, state_dict: dict[str, Any]) -> None:
 
                 src_state_val = state[param_key][state_key]
                 if isinstance(state_val, ShardedTensor):
-                    assert isinstance(src_state_val, ShardedTensor)
+                    if not isinstance(src_state_val, ShardedTensor):
+                        raise AssertionError
                     num_shards = len(state_val.local_shards())
                     num_new_shards = len(src_state_val.local_shards())
                     if num_shards != num_new_shards:
@@ -230,7 +233,8 @@ def load_state_dict(self, state_dict: dict[str, Any]) -> None:
                     ):
                         shard.tensor.detach().copy_(src_shard.tensor)
                 elif isinstance(state_val, torch.Tensor):
-                    assert isinstance(src_state_val, torch.Tensor)
+                    if not isinstance(src_state_val, torch.Tensor):
+                        raise AssertionError
                     state_val.detach().copy_(src_state_val)
                 else:
                     new_state[idx][state_key] = deepcopy(src_state_val)
@@ -275,7 +279,8 @@ def add_param_group(self, param_group: Mapping[str, Any]) -> None:
 
         Warning: This API is still in development and subject to change.
         """
-        assert isinstance(param_group, dict), "param group must be a dict"
+        if not isinstance(param_group, dict):
+            raise AssertionError("param group must be a dict")
 
         params = param_group["params"]
         if isinstance(params, torch.Tensor):
diff --git a/torch/distributed/optim/zero_redundancy_optimizer.py b/torch/distributed/optim/zero_redundancy_optimizer.py
index 6394b28b74872..02b16543892ac 100644
--- a/torch/distributed/optim/zero_redundancy_optimizer.py
+++ b/torch/distributed/optim/zero_redundancy_optimizer.py
@@ -114,7 +114,6 @@ def _broadcast_object(
         )
         # pyrefly: ignore [bad-argument-type]
         dist.broadcast(data_recv_tensor, src=src_rank, group=group, async_op=False)
-        # pyrefly: ignore [bad-argument-type]
         buffer = io.BytesIO(data_recv_tensor.cpu().numpy())
         obj = torch.load(buffer, map_location=device, weights_only=False)
     return obj
@@ -122,10 +121,11 @@ def _broadcast_object(
 
 class _ZeROJoinHook(JoinHook):
     def __init__(self, zero):
-        assert isinstance(zero, ZeroRedundancyOptimizer), (
-            "ZeRO join hook requires passing in a ZeroRedundancyOptimizer "
-            "instance as the state"
-        )
+        if not isinstance(zero, ZeroRedundancyOptimizer):
+            raise AssertionError(
+                "ZeRO join hook requires passing in a "
+                "ZeroRedundancyOptimizer instance as the state"
+            )
         self.zero = zero
         super().__init__()
 
@@ -172,7 +172,6 @@ def __init__(
         if len(self.parameters) == 0:
             raise ValueError("Empty bucket assignment")
         # DDP guarantees all parameters in the bucket have the same device
-        # pyrefly: ignore [read-only]
         self.device: torch.device = self.parameters[0].device
         self.tensor: torch.Tensor | None = None
 
@@ -270,9 +269,10 @@ def wait_for_broadcasts(self) -> None:
         meaning ``self.broadcast_handles`` is filled. This clears ``self.broadcast_handles``
         in preparation for the next iteration.
         """
-        assert len(self.broadcast_handles) == self.num_bucket_assignments, (
-            f"Missing at least one broadcast handle on rank {dist.get_rank()}"
-        )
+        if len(self.broadcast_handles) != self.num_bucket_assignments:
+            raise AssertionError(
+                f"Missing at least one broadcast handle on rank {dist.get_rank()}"
+            )
         _ = [x.wait() for x in self.broadcast_handles]
         self.broadcast_handles.clear()
 
@@ -703,10 +703,11 @@ def _partition_parameters(
             return self._partition_parameters_cache
 
         # Partition the parameters according to `params_per_rank`
-        assert len(self._partition_parameters_cache) == 0, (
-            "Specifying `params_per_rank` should only be done when the "
-            "parameters have not been partitioned yet"
-        )
+        if len(self._partition_parameters_cache) != 0:
+            raise AssertionError(
+                "Specifying `params_per_rank` should only be done when the "
+                "parameters have not been partitioned yet"
+            )
         if len(self.param_groups) != 1:
             raise RuntimeError(
                 "Specifying `params_per_rank` only supports a single parameter group"
@@ -767,11 +768,12 @@ def _broadcast_params_from_rank(self, rank: int):
             A :class:`list` of async work handles for the ``broadcast()`` s
             performed to synchronize the parameters.
         """
-        assert not self._overlap_with_ddp, (
-            "`_broadcast_params_from_rank()` should not be used if "
-            "`overlap_with_ddp=True`; instead, the broadcasting should "
-            "happen in the DDP communication hook"
-        )
+        if self._overlap_with_ddp:
+            raise AssertionError(
+                "`_broadcast_params_from_rank()` should not be used if "
+                "`overlap_with_ddp=True`; instead, the broadcasting should "
+                "happen in the DDP communication hook"
+            )
         handles = []
         if self.parameters_as_bucket_view:
             for dev_i_buckets in self._buckets:
@@ -849,10 +851,11 @@ def _device_to_params_per_rank(
             ...
         ...
         """
-        assert self.parameters_as_bucket_view, (
-            "`_device_to_params_per_rank` should only be used if "
-            "`parameters_as_bucket_view=True`"
-        )
+        if not self.parameters_as_bucket_view:
+            raise AssertionError(
+                "`_device_to_params_per_rank` should only be used if "
+                "`parameters_as_bucket_view=True`"
+            )
         if len(self._device_to_params_per_rank_cache) == 0:
             for rank, param_groups in enumerate(self._partition_parameters()):
                 for param_group in param_groups:
@@ -890,7 +893,8 @@ def _get_min_index(
             if value < min_value:
                 min_value = value
                 min_index = i
-        assert min_index >= 0, "All indices are disallowed"
+        if min_index < 0:
+            raise AssertionError("All indices are disallowed")
         return min_index
 
     def _assign_bucket_subset_to_rank(
@@ -942,21 +946,24 @@ def _bucket_assignments_per_rank(self) -> list[dict[int, _DDPBucketAssignment]]:
         mapping bucket indices to :class:`_DDPBucketAssignment` s for each
         rank.
         """
-        assert self._overlap_with_ddp, (
-            "`_bucket_assignments_per_rank` only be used if `overlap_with_ddp=True`"
-        )
+        if not self._overlap_with_ddp:
+            raise AssertionError(
+                "`_bucket_assignments_per_rank` only be used if `overlap_with_ddp=True`"
+            )
         if len(self._bucket_assignments_per_rank_cache) > 0:
             return self._bucket_assignments_per_rank_cache
 
         overlap_info = self._overlap_info
-        assert overlap_info.status == _OverlapStatus.INITIALIZED
+        if overlap_info.status != _OverlapStatus.INITIALIZED:
+            raise AssertionError
 
         self._bucket_assignments_per_rank_cache = [{} for _ in range(self.world_size)]
         params_per_bucket = overlap_info.params_per_bucket
 
         if overlap_info.shard_buckets:
             # Define the assignment threshold to approximate uniformity
-            assert overlap_info.total_size is not None, "`total_size` was not computed"
+            if overlap_info.total_size is None:
+                raise AssertionError("`total_size` was not computed")
             threshold = overlap_info.total_size / self.world_size  # type: ignore[operator]
             size_per_rank = [0 for _ in range(self.world_size)]
 
@@ -966,7 +973,8 @@ def _bucket_assignments_per_rank(self) -> list[dict[int, _DDPBucketAssignment]]:
         if not overlap_info.shard_buckets:
             # Assign each DDP bucket entirely to a single rank
             for bucket_index, bucket_params in enumerate(params_per_bucket):
-                assert len(bucket_params) > 0, "Empty bucket"
+                if len(bucket_params) <= 0:
+                    raise AssertionError("Empty bucket")
                 assigned_rank = self._get_assigned_rank(bucket_index)
                 self._assign_bucket_subset_to_rank(
                     bucket_index,
@@ -987,7 +995,8 @@ def _bucket_assignments_per_rank(self) -> list[dict[int, _DDPBucketAssignment]]:
                 enumerate(params_per_bucket), key=lambda x: sum(p.numel() for p in x[1])
             )
             for bucket_index, bucket_params in params_per_bucket_enum:
-                assert len(bucket_params) > 0, "Empty bucket"
+                if len(bucket_params) <= 0:
+                    raise AssertionError("Empty bucket")
                 bucket_offset = 0
                 assignment_size = 0
                 for param_index, param in enumerate(bucket_params):
@@ -1094,13 +1103,15 @@ def _local_step(
                 else self.optim.step(closure=closure, **kwargs)
             )
         else:
-            assert self._overlap_with_ddp, (
-                "Specifying `gradients` should not "
-                "be used when `overlap_with_ddp=False`"
-            )
-            assert closure is None, (
-                "`closure` is not supported when using a local functional optimizer"
-            )
+            if not self._overlap_with_ddp:
+                raise AssertionError(
+                    "Specifying `gradients` should not "
+                    "be used when `overlap_with_ddp=False`"
+                )
+            if closure is not None:
+                raise AssertionError(
+                    "`closure` is not supported when using a local functional optimizer"
+                )
             loss = self.optim.step(gradients=gradients)
 
         # Sync any updated attributes in the local optimizer to the exposed
@@ -1244,9 +1255,10 @@ def state_dict(self) -> dict[str, Any]:
         for rank, local_state_dict in enumerate(self._all_state_dicts):
             local_param_groups = local_state_dict["param_groups"]
             global_param_groups = self._partition_parameters()[rank]
-            assert len(local_param_groups) == len(global_param_groups), (
-                "Mismatch between number of local and global parameter groups"
-            )
+            if len(local_param_groups) != len(global_param_groups):
+                raise AssertionError(
+                    "Mismatch between number of local and global parameter groups"
+                )
 
             for local_param_group, global_param_group in zip(
                 local_param_groups, global_param_groups
@@ -1256,9 +1268,11 @@ def state_dict(self) -> dict[str, Any]:
                 local_param_indices = local_param_group["params"]
                 global_params = global_param_group["params"]
 
-                assert len(local_param_indices) == len(global_params), (
-                    "Mismatch between number of local and global parameters in parameter group"
-                )
+                if len(local_param_indices) != len(global_params):
+                    raise AssertionError(
+                        "Mismatch between number of local and global "
+                        "parameters in parameter group"
+                    )
                 for local_param_index, global_param in zip(
                     local_param_indices, global_params
                 ):
@@ -1291,9 +1305,10 @@ def _sync_param_groups(
             dst_param_groups (list[dict]): parameter groups giving the
                 attribute settings to set.
         """
-        assert len(src_param_groups) == len(dst_param_groups), (
-            "Mismatch between number of source and destination parameter groups"
-        )
+        if len(src_param_groups) != len(dst_param_groups):
+            raise AssertionError(
+                "Mismatch between number of source and destination parameter groups"
+            )
         for src_param_group, dst_param_group in zip(src_param_groups, dst_param_groups):
             # Sync all attributes except the parameters
             for attr in filter(lambda x: x != "params", src_param_group.keys()):
@@ -1380,14 +1395,16 @@ def _build_ddp_param_buckets(self) -> None:
                 bucket_size = 0
                 dtype = None
                 for param in params:
-                    assert _is_trainable(param), (
-                        "Model parameter "
-                        "corresponding to a gradient in a DDP bucket should "
-                        "require a gradient"
-                    )
+                    if not _is_trainable(param):
+                        raise AssertionError(
+                            "Model parameter "
+                            "corresponding to a gradient in a DDP bucket should "
+                            "require a gradient"
+                        )
                     bucket_size += param.numel()
                     dtype = param.dtype  # assumes all same dtype
-                assert bucket_size > 0, "Empty bucket"
+                if bucket_size <= 0:
+                    raise AssertionError("Empty bucket")
 
                 # Construct the bucket tensor (assuming all dense and same dtype)
                 tensor = torch.empty(
@@ -1502,19 +1519,19 @@ def _init_local_optimizer(self) -> None:
 
         The local optimizer is saved in ``self.optim``.
         """
-        assert self._optim_constructor is not None, (
-            "The local optimizer class has not been set"
-        )
+        if self._optim_constructor is None:
+            raise AssertionError("The local optimizer class has not been set")
 
         param_groups = self._partition_parameters()[self.rank]
         # `overlap_with_ddp=True` requires a local functional optimizer
         if self._overlap_with_ddp:
             # Functional optimizers only support a single parameter group and
             # require passing in the parameters as a list
-            assert len(param_groups) == 1, (
-                "Initializing the local "
-                "functional optimizer with more than one parameter group"
-            )
+            if len(param_groups) != 1:
+                raise AssertionError(
+                    "Initializing the local functional optimizer "
+                    "with more than one parameter group"
+                )
             params = param_groups[0]["params"]
             # Try to pass `_allow_empty_param_list=True` to avoid erroring
             if (
@@ -1564,20 +1581,22 @@ def _init_local_optimizer(self) -> None:
         # optimizer; remove this if/when the functional optimizers support
         # multiple parameter groups
         if self._overlap_with_ddp and not hasattr(self.optim, "param_groups"):
-            assert hasattr(self.optim, "param_group"), (
-                "The functional optimizer should set at least one of the "
-                "attributes `param_group` or `param_groups`"
-            )
+            if not hasattr(self.optim, "param_group"):
+                raise AssertionError(
+                    "The functional optimizer should set at least one of "
+                    "the attributes `param_group` or `param_groups`"
+                )
             self.optim.param_groups = [self.optim.param_group]  # type: ignore[attr-defined]
 
         self._sync_param_groups(self.optim.param_groups, self.param_groups)
 
     def _init_zero_for_overlap(self) -> None:
         r"""Perform a delayed initialization of the local optimizer and the supporting data structures."""
-        assert self._overlap_with_ddp, (
-            "`_init_zero_for_overlap()` should only be called when "
-            "`overlap_with_ddp=True`"
-        )
+        if not self._overlap_with_ddp:
+            raise AssertionError(
+                "`_init_zero_for_overlap()` should only be called when "
+                "`overlap_with_ddp=True`"
+            )
         self._overlap_info.status = _OverlapStatus.INITIALIZED
         self._clear_cache()
         self._partition_parameters(self._overlap_info.params_per_rank)
@@ -1592,11 +1611,12 @@ def _get_assigned_rank(self, bucket_index: int) -> int:
             bucket_index (int): index of the :class:`DistributedDataParallel`
                 bucket for which to get the assigned rank.
         """
-        assert not self._overlap_info.shard_buckets, (
-            "The bucket assignment requires global bucket information and "
-            "will be computed later; there should be no need to use this "
-            "method"
-        )
+        if self._overlap_info.shard_buckets:
+            raise AssertionError(
+                "The bucket assignment requires global bucket information "
+                "and will be computed later; there should be no need to "
+                "use this method"
+            )
         return bucket_index % self.world_size
 
     def _check_overlap_initialized(self):
diff --git a/torch/distributed/pipelining/_IR.py b/torch/distributed/pipelining/_IR.py
index 7d043f03d30ba..bf8c4613b0849 100644
--- a/torch/distributed/pipelining/_IR.py
+++ b/torch/distributed/pipelining/_IR.py
@@ -8,7 +8,7 @@
 from enum import Enum
 from inspect import Parameter, Signature, signature
 from types import MethodType
-from typing import Any, Union
+from typing import Any
 
 import torch
 import torch.fx as fx
@@ -364,7 +364,7 @@ class MultiUseParameterConfig(Enum):
     REPLICATE = 2
 
 
-MultiUseParamSpec = Union[MultiUseParameterConfig, dict[str, MultiUseParameterConfig]]
+MultiUseParamSpec = MultiUseParameterConfig | dict[str, MultiUseParameterConfig]
 
 
 class DetachExecutor(fx.Interpreter):
diff --git a/torch/distributed/pipelining/_backward.py b/torch/distributed/pipelining/_backward.py
index 1743e47c7ff9f..7e8f418a02d4a 100644
--- a/torch/distributed/pipelining/_backward.py
+++ b/torch/distributed/pipelining/_backward.py
@@ -2,7 +2,7 @@
 # Copyright (c) Meta Platforms, Inc. and affiliates
 import collections
 import logging
-from collections.abc import Iterator
+from collections.abc import Iterator, Sequence
 from typing import Any
 
 import torch
@@ -140,6 +140,40 @@ def get_param_groups(
     return unique_param_groups
 
 
+def _autograd_grad_for_inputs(
+    outputs: Sequence[torch.Tensor],
+    inputs: Sequence[torch.Tensor],
+    grad_outputs: Sequence[torch.Tensor | None] | None = None,
+    retain_graph: bool = False,
+    allow_unused: bool = False,
+) -> tuple[torch.Tensor | None, ...]:
+    """Compute input gradients, returning ``None`` for non-grad inputs."""
+    # Some inputs may not be used or may not require gradients, so we filter them out
+    # before calling autograd.grad and place None for those positions in the result.
+    grad_indices: list[int] = []
+    inputs_requiring_grad: list[torch.Tensor] = []
+    for i, inp in enumerate(inputs):
+        if isinstance(inp, torch.Tensor) and inp.requires_grad:
+            grad_indices.append(i)
+            inputs_requiring_grad.append(inp)
+
+    if not inputs_requiring_grad:
+        return tuple(None for _ in inputs)
+
+    grads = torch.autograd.grad(
+        outputs=outputs,
+        inputs=inputs_requiring_grad,
+        grad_outputs=grad_outputs,
+        retain_graph=retain_graph,
+        allow_unused=allow_unused,
+    )
+
+    result: list[torch.Tensor | None] = [None] * len(inputs)
+    for idx, g in zip(grad_indices, grads, strict=True):
+        result[idx] = g
+    return tuple(result)
+
+
 def stage_backward_input(
     stage_outputs_or_loss: list[torch.Tensor],
     output_grads: list[torch.Tensor] | None,
@@ -196,20 +230,20 @@ def hook(grad_inputs):
             torch.ones_like(stage_output) for stage_output in stage_outputs_or_loss
         ]
 
-    # Some inputs may not be used or may not require gradients, so we filter them out
-    input_values = [inp for inp in input_values if inp.requires_grad]
-    dinputs = torch.autograd.grad(
+    dinputs = _autograd_grad_for_inputs(
         stage_outputs_or_loss,
-        inputs=input_values,
-        grad_outputs=output_grads,
+        input_values,
+        output_grads,
         retain_graph=True,
     )
-    # Update the gradients for inputs
+
+    # Accumulate into .grad
     for inp, dinput in zip(input_values, dinputs):
-        if inp.grad is None:
-            inp.grad = dinput
-        else:
-            inp.grad += dinput
+        if isinstance(inp, torch.Tensor) and dinput is not None:
+            if inp.grad is None:
+                inp.grad = dinput
+            else:
+                inp.grad += dinput
 
     # stage_outputs_or_loss are not used in backwards after this point, so we can safely remove it from the autograd graph
     # this allows autograd to clear up the graph dedicated for this tensor and free up significant memory
@@ -231,7 +265,6 @@ def stage_backward_weight(
     weight_grads: list[torch.Tensor | None] = []
     for index, weight in enumerate(weights):
         grad_acc = _get_grad_fn_or_grad_acc(weight)
-        # pyrefly: ignore [unsupported-operation]
         grad_acc_to_weight[grad_acc] = weight, index
         weight_grads.append(weight.grad)
 
@@ -242,12 +275,11 @@ def stage_backward_weight(
         for grads_tuple, intermediate in zip(
             param_group["grads"], param_group["intermediates"]
         ):
-            non_none_grads = [g for g in grads_tuple if g is not None]
-            if non_none_grads:
-                summed_grad = sum(non_none_grads)
-                valid_edges.append(GradientEdge(intermediate, 0))
-                # pyrefly: ignore [bad-argument-type]
-                valid_grad_outputs.append(summed_grad)
+            for i, grad in enumerate(grads_tuple):
+                if grad is not None:
+                    valid_edges.append(GradientEdge(intermediate, i))
+                    # pyrefly: ignore [bad-argument-type]
+                    valid_grad_outputs.append(grad)
 
         # Break a reference cycle caused inside stage_backward_input->get_hook->hook
         # The summarized cycle is:
diff --git a/torch/distributed/pipelining/_schedule_visualizer.py b/torch/distributed/pipelining/_schedule_visualizer.py
index c810f95078429..5502e8503d936 100644
--- a/torch/distributed/pipelining/_schedule_visualizer.py
+++ b/torch/distributed/pipelining/_schedule_visualizer.py
@@ -87,7 +87,7 @@ def get_schedule_ops(
         raise ValueError(f"Invalid schedule: {schedule_class}")
 
     # Instantiate the schedule class
-    # pyrefly: ignore [bad-instantiation, bad-argument-type]
+    # pyrefly: ignore [bad-argument-type]
     schedule_instance = schedule_class(stages, num_microbatches)
     if schedule_instance.pipeline_order is None:
         raise AssertionError("Expected pipeline_order to not be None")
diff --git a/torch/distributed/pipelining/_utils.py b/torch/distributed/pipelining/_utils.py
index 79b74be406814..3da67258ecbd9 100644
--- a/torch/distributed/pipelining/_utils.py
+++ b/torch/distributed/pipelining/_utils.py
@@ -1,91 +1,506 @@
 # mypy: allow-untyped-defs
 # Copyright (c) Meta Platforms, Inc. and affiliates
 
+from __future__ import annotations
+
 import logging
-from dataclasses import dataclass
+import warnings
+from dataclasses import dataclass, field
+from enum import Enum
+from typing import cast, Literal, overload, Protocol, TYPE_CHECKING, TypeAlias
 
 import torch
 from torch import fx
+from torch.distributed._mesh_layout import _MeshLayout
+from torch.distributed.tensor import DTensor
+from torch.utils._pytree import tree_flatten, tree_unflatten
+
+
+if TYPE_CHECKING:
+    from torch.distributed.device_mesh import DeviceMesh
+    from torch.distributed.tensor.placement_types import Placement
 
 
 logger = logging.getLogger(__name__)
 
 
-def flatten_args_detach(args):
+class GetMeshCallback(Protocol):
+    """Callback to create/retrieve a DeviceMesh from its cache key components."""
+
+    def __call__(
+        self,
+        mesh_dim_names: tuple[str, ...],
+        mesh_layout: _MeshLayout | None,
+    ) -> DeviceMesh: ...
+
+
+# Key for mesh cache: (mesh_dim_names, mesh_layout)
+# mesh_layout is the _MeshLayout object containing shape and stride (not actual ranks).
+# This uniquely identifies a mesh within the same "universe" where all stages share
+# the same rank tensor.
+MeshCacheKey: TypeAlias = tuple[tuple[str, ...], _MeshLayout | None]
+
+
+class PipeliningMetadataError(RuntimeError):
+    """Raised on metadata mismatches during pipeline communication."""
+
+
+@dataclass(frozen=True, slots=True)
+class _TensorMeta:
+    """Tensor metadata for recv buffer allocation and validation.
+
+    For plain tensors, these are the tensor's actual attributes.
+    For DTensors, these are LOCAL shard attributes; global attributes
+    are stored in :class:`_DTensorMeta`.
     """
-    Flatten the args into a list form and detach the tensors from computational graph.
+
+    shape: torch.Size
+    stride: tuple[int, ...]
+    dtype: torch.dtype
+    requires_grad: bool
+
+    @staticmethod
+    def from_tensor(tensor: torch.Tensor) -> _TensorMeta:
+        """Create metadata from a plain tensor.
+
+        Args:
+            tensor: A plain ``torch.Tensor`` (not DTensor).
+
+        Returns:
+            Metadata capturing shape, stride, dtype, and requires_grad.
+
+        Raises:
+            TypeError: If ``tensor`` is a DTensor.
+        """
+        if isinstance(tensor, DTensor):
+            raise PipeliningMetadataError(
+                "Expected plain tensor, got DTensor. Use _DTensorMeta.from_dtensor instead."
+            )
+        return _TensorMeta(
+            shape=tensor.shape,
+            stride=tensor.stride(),
+            dtype=tensor.dtype,
+            requires_grad=tensor.requires_grad,
+        )
+
+    def to_tensor(self, device: torch.device | str) -> torch.Tensor:
+        """Reconstruct a tensor on ``device`` from this metadata.
+
+        Args:
+            device: Target device for the tensor.
+
+        Returns:
+            An empty strided tensor on ``device``.
+        """
+        t = _make_tensor_from_meta(self, device)
+        t.requires_grad_(self.requires_grad)
+        return t
+
+    def get_diff(self, other: _TensorMeta) -> list[str]:
+        """Return field-by-field differences with ``other``.
+
+        Args:
+            other: Metadata to compare against.
+
+        Returns:
+            List of human-readable difference strings (empty if equal).
+        """
+        if self == other:
+            return []
+
+        diffs = []
+        if self.shape != other.shape:
+            diffs.append(f"shape mismatch: {self.shape} vs {other.shape}")
+        if self.stride != other.stride:
+            diffs.append(f"stride mismatch: {self.stride} vs {other.stride}")
+        if self.dtype != other.dtype:
+            diffs.append(f"dtype mismatch: {self.dtype} vs {other.dtype}")
+        # requires_grad is intentionally excluded: it is a runtime concern
+        # determined by has_backward and grad context, not a metadata invariant.
+        return diffs
+
+
+@dataclass(frozen=True, slots=True)
+class _DTensorMeta(_TensorMeta):
+    """DTensor metadata extending :class:`_TensorMeta` with distribution info.
+
+    Inherited fields (shape, stride, etc.) are LOCAL shard attributes.
+    Additional fields capture global shape and placement information
+    needed to reconstruct a :class:`DTensor` via ``DTensor.from_local()``.
+
+    The :class:`DeviceMesh` is **not** stored (not serializable for P2P);
+    it is looked up from :class:`_MeshCache` using
+    ``(mesh_dim_names, mesh_layout)`` as the key.
     """
-    flat_detached_args = []
 
-    def extract_tensor_args(a):
-        nonlocal flat_detached_args
-        if isinstance(a, torch.Tensor):
-            val = a.detach().requires_grad_(a.requires_grad)
-            flat_detached_args.append(val)
-            return val
+    # Global DTensor properties (for reconstruction)
+    global_shape: torch.Size = field(default_factory=lambda: torch.Size([]))
+    global_stride: tuple[int, ...] = field(default=())
+
+    # DTensor distribution properties
+    placements: tuple[Placement, ...] = field(
+        default=()
+    )  # e.g., (Shard(0), Replicate())
+
+    # Mesh identification - used to look up the correct DeviceMesh from cache
+    mesh_dim_names: tuple[str, ...] = field(default=())  # e.g., ("tp",) or ("dp", "tp")
+    mesh_layout: _MeshLayout | None = field(
+        default=None
+    )  # _MeshLayout with shape/stride - uniquely identifies mesh within the same universe
+
+    @staticmethod
+    def from_dtensor(dtensor: DTensor) -> _DTensorMeta:
+        """Create metadata from a DTensor.
+
+        Args:
+            dtensor: The DTensor to extract metadata from.
+
+        Returns:
+            Metadata capturing both local and global attributes.
+        """
+        device_mesh = dtensor.device_mesh
+
+        return _DTensorMeta(
+            # Local tensor attributes (for recv buffer allocation)
+            shape=dtensor._local_tensor.shape,
+            stride=dtensor._local_tensor.stride(),
+            dtype=dtensor.dtype,
+            requires_grad=dtensor.requires_grad,
+            # Global DTensor attributes (for reconstruction)
+            global_shape=dtensor.shape,
+            global_stride=dtensor.stride(),
+            # Distribution info
+            placements=dtensor._spec.placements,
+            mesh_dim_names=(
+                tuple(device_mesh.mesh_dim_names) if device_mesh.mesh_dim_names else ()
+            ),
+            mesh_layout=device_mesh._layout,
+        )
+
+    @property
+    def mesh_cache_key(self) -> MeshCacheKey:
+        """Cache key ``(mesh_dim_names, mesh_layout)`` for mesh lookup."""
+        return (self.mesh_dim_names, self.mesh_layout)
+
+    def to_dtensor(self, device: torch.device | str, mesh: DeviceMesh) -> DTensor:
+        """Reconstruct a DTensor on ``device`` with placements.
+
+        Args:
+            device: Target device for the local tensor.
+            mesh: The ``DeviceMesh`` to attach.
+
+        Returns:
+            A DTensor on ``device``.
+        """
+        local_tensor = _make_tensor_from_meta(self, device)
+        # Set requires_grad after from_local() so that the from_local
+        # operation itself is not recorded in the autograd graph.
+        return cast(
+            DTensor,
+            DTensor.from_local(
+                local_tensor,
+                device_mesh=mesh,
+                placements=self.placements,
+                shape=self.global_shape,
+                stride=self.global_stride,
+                run_check=False,
+            ).requires_grad_(self.requires_grad),
+        )
+
+    def get_diff(self, other: _TensorMeta) -> list[str]:
+        """Return field-by-field differences, including DTensor-specific fields.
+
+        Args:
+            other: Metadata to compare against.
+
+        Returns:
+            List of human-readable difference strings (empty if equal).
+        """
+        if self == other:
+            return []
+
+        # Get base class differences (compares local shape/stride/dtype/requires_grad)
+        # NOTE: Use explicit class call instead of super() because
+        # @dataclass(slots=True) on both parent and child can break super().
+        diffs = _TensorMeta.get_diff(self, other)
+
+        # Add DTensor-specific comparisons if other is also _DTensorMeta
+        if isinstance(other, _DTensorMeta):
+            if self.global_shape != other.global_shape:
+                diffs.append(
+                    f"global_shape mismatch: {self.global_shape} vs {other.global_shape}"
+                )
+            if self.global_stride != other.global_stride:
+                diffs.append(
+                    f"global_stride mismatch: {self.global_stride} vs {other.global_stride}"
+                )
+            if self.placements != other.placements:
+                diffs.append(
+                    f"placements mismatch: {self.placements} vs {other.placements}"
+                )
+            if self.mesh_dim_names != other.mesh_dim_names:
+                diffs.append(
+                    f"mesh_dim_names mismatch: {self.mesh_dim_names} vs {other.mesh_dim_names}"
+                )
+            if self.mesh_layout != other.mesh_layout:
+                diffs.append(
+                    f"mesh_layout mismatch: {self.mesh_layout} vs {other.mesh_layout}"
+                )
         else:
-            flat_detached_args.append(a)
-            return a
+            diffs.append("type: _DTensorMeta vs _TensorMeta")
 
-    new_args = fx.node.map_aggregate(
-        args,
-        extract_tensor_args,
-    )
+        return diffs
 
-    return new_args, flat_detached_args
 
+# Type alias for union of tensor metadata types
+TensorMeta: TypeAlias = _TensorMeta | _DTensorMeta
+
+
+# Not frozen: fields are populated incrementally during forward and
+# backward metadata inference or from user provided static metadata
+@dataclass(slots=True)
+class _StageMeta:
+    """Consolidated tensor metadata for a pipeline stage's forward and backward passes."""
+
+    inputs: tuple[TensorMeta, ...] | None = None
+    outputs: tuple[TensorMeta, ...] | None = None
+    input_grads: tuple[TensorMeta | None, ...] | None = None
+    output_grads: tuple[TensorMeta | None, ...] | None = None
+
+    def has_any(self) -> bool:
+        """Check if any metadata field is populated."""
+        return any(
+            v is not None
+            for v in [self.inputs, self.outputs, self.input_grads, self.output_grads]
+        )
 
-def flatten_args(args):
+    def has_dtensors(self) -> bool:
+        """Check if any input/output metadata is DTensor type."""
+        for metas in [self.inputs, self.outputs]:
+            if metas and any(isinstance(m, _DTensorMeta) for m in metas if m):
+                return True
+        return False
+
+    def is_complete_for_forward(self) -> bool:
+        """Check if forward metadata is fully populated."""
+        return self.inputs is not None and self.outputs is not None
+
+
+@dataclass(frozen=True, slots=True)
+class _StageForwardMeta:
+    """Forward metadata transmitted from stage *i* to stage *i+1* during inference."""
+
+    forward_metas: tuple[TensorMeta, ...]  # Stage i's outputs → Stage i+1's inputs
+
+
+@dataclass(frozen=True, slots=True)
+class _StageBackwardMeta:
+    """Backward metadata transmitted from stage *i* to stage *i-1* during inference.
+
+    Gradient placements may differ from forward activations
+    (e.g., ``Replicate`` → ``Partial``).
     """
-    Flatten the args into a list form.
+
+    backward_metas: tuple[
+        TensorMeta | None, ...
+    ]  # Stage i's input_grads → Stage i-1's output_grads
+
+
+def _make_tensor_from_meta(
+    meta: _TensorMeta,
+    device: torch.device | str,
+) -> torch.Tensor:
+    """Create a tensor from metadata.
+
+    Args:
+        meta: Metadata with shape, stride, and dtype.
+        device: Target device for the tensor.
+
+    Returns:
+        Empty tensor preserving the exact memory layout.
     """
-    flat_args = []
+    return torch.empty_strided(
+        size=meta.shape,
+        stride=meta.stride,
+        dtype=meta.dtype,
+        device=device,
+    )
+
 
-    def extract_tensor_args(a):
-        nonlocal flat_args
-        flat_args.append(a)
-        return a
+def _derive_grad_metas(
+    tensor_metas: tuple[TensorMeta, ...],
+) -> tuple[_TensorMeta | None, ...]:
+    """Derive gradient metadata from tensor metadata.
 
-    fx.node.map_aggregate(
-        args,
-        extract_tensor_args,
+    Returns metadata with the same shape/stride/dtype but ``requires_grad=False``.
+    Entries where the source has ``requires_grad=False`` become ``None``.
+    """
+    return tuple(
+        _TensorMeta(shape=m.shape, stride=m.stride, dtype=m.dtype, requires_grad=False)
+        if m.requires_grad
+        else None
+        for m in tensor_metas
     )
 
-    return flat_args
 
+class _MeshCache:
+    """Cache for :class:`DeviceMesh` objects keyed by ``(mesh_dim_names, mesh_layout)``.
 
-class PipeliningShapeError(RuntimeError):
-    """Shape mismatch between configured and runtime values."""
+    Assumes all pipeline stages share the same rank tensor (true for
+    TorchTitan-style frameworks where meshes derive from a common world).
+    """
 
+    def __init__(self, get_mesh_cb: GetMeshCallback | None = None) -> None:
+        self._cache: dict[MeshCacheKey, DeviceMesh] = {}
+        self._get_mesh_cb = get_mesh_cb
 
-def validate_tensor_metadata(desc, expected, given):
-    if not expected.shape == given.shape:
-        raise PipeliningShapeError(
-            f"{desc} has a shape mismatch: expected {expected.shape} actual {given.shape}"
-        )
-    if not expected.dtype == given.dtype:
-        raise PipeliningShapeError(
-            f"{desc} has a dtype mismatch: expected {expected.dtype} actual {given.dtype}"
-        )
-    if not expected.stride() == given.stride():
-        raise PipeliningShapeError(
-            f"{desc} has a stride mismatch: expected {expected.stride()} actual {given.stride()}"
-        )
+    def get_mesh(self, key: MeshCacheKey) -> DeviceMesh:
+        """Return a cached mesh, or create one via the callback.
 
+        Args:
+            key: Cache key ``(mesh_dim_names, mesh_layout)``.
 
-def validate_tensors_metadata(
-    desc,
-    expected_tensors: list[torch.Tensor] | tuple[torch.Tensor, ...],
-    actual_tensors: list[torch.Tensor] | tuple[torch.Tensor, ...],
-):
-    if len(expected_tensors) != len(actual_tensors):
-        raise PipeliningShapeError(
-            f"{desc}: Number of values ({len(actual_tensors)}) does not match expected number ({len(expected_tensors)})"
-        )
-    for i in range(len(expected_tensors)):
-        validate_tensor_metadata(
-            f"{desc}: value {i}", expected_tensors[i], actual_tensors[i]
-        )
+        Returns:
+            The ``DeviceMesh``.
+
+        Raises:
+            PipeliningMetadataError: If not cached and no callback provided.
+        """
+        if key in self._cache:
+            return self._cache[key]
+
+        mesh_dim_names, mesh_layout = key
+
+        if self._get_mesh_cb is None:
+            raise PipeliningMetadataError(
+                f"Mesh not found in cache for mesh_dim_names={mesh_dim_names}, "
+                f"mesh_layout={mesh_layout}, and no get_mesh callback provided. "
+                f"Provide a get_mesh callback or use DTensors in static mode."
+            )
+
+        mesh = self._get_mesh_cb(mesh_dim_names, mesh_layout)
+        if mesh is None:
+            raise PipeliningMetadataError(
+                f"Mesh lookup failed: callback returned None for "
+                f"mesh_dim_names={mesh_dim_names}, mesh_layout={mesh_layout}. "
+                f"Ensure all stages use meshes from the same universe."
+            )
+        self._cache[key] = mesh
+        return mesh
+
+    def put(self, key: MeshCacheKey, mesh: DeviceMesh) -> None:
+        """Add a mesh to the cache."""
+        self._cache[key] = mesh
+
+    def update_from_tensors(self, tensors: tuple[torch.Tensor | None, ...]) -> None:
+        """Extract and cache meshes from any :class:`DTensor` instances in *tensors*."""
+        for tensor in tensors:
+            if isinstance(tensor, DTensor):
+                mesh = tensor.device_mesh
+                dim_names = tuple(mesh.mesh_dim_names) if mesh.mesh_dim_names else ()
+                mesh_layout = mesh._layout
+                key = (dim_names, mesh_layout)
+                if key not in self._cache:
+                    self._cache[key] = mesh
+
+    def __contains__(self, key: MeshCacheKey) -> bool:
+        return key in self._cache
+
+    def __len__(self) -> int:
+        return len(self._cache)
+
+
+# ============================================================================
+# Inference mode enum
+# ============================================================================
+
+
+class InferenceMode(Enum):
+    """Pipeline-level metadata inference mode, determined collectively across all PP ranks.
+
+    The mode is set by the schedule (not individual stages) because
+    ``has_backward`` is only known at schedule creation time and all
+    stages must agree to avoid P2P hangs.
+
+    .. attribute:: STATIC
+
+        All stages have sufficient metadata; runtime inference is skipped.
+
+    .. attribute:: DYNAMIC
+
+        At least one stage requires runtime metadata inference.
+    """
+
+    STATIC = "static"
+    DYNAMIC = "dynamic"
+
+    @classmethod
+    def needs_dynamic(cls, meta: _StageMeta, stage_has_backward: bool) -> bool:
+        """Determine whether dynamic metadata inference is needed for a stage.
+
+        Args:
+            meta: Stage metadata from user-provided args.
+            stage_has_backward: Whether a backward pass will be performed.
+
+        Returns:
+            ``True`` if dynamic inference is needed.
+        """
+        # Case 1: Forward metadata incomplete → needs DYNAMIC
+        if not meta.is_complete_for_forward():
+            return True
+
+        # Case 2: No DTensors → STATIC is fine (bwd metadata derivable from fwd metadata)
+        if not meta.has_dtensors():
+            return False
+
+        # Case 3: No backward needed → STATIC is fine (don't need grad metadata)
+        if not stage_has_backward:
+            return False
+
+        # Case 4: DTensors with backward but missing ANY grad metadata → needs DYNAMIC
+        # Both input_grads AND output_grads are required for static mode with DTensors
+        if meta.input_grads is None or meta.output_grads is None:
+            return True
+
+        # Case 5: DTensors with complete grads → STATIC is fine
+        return False
+
+
+# ============================================================================
+# Utility functions
+# ============================================================================
+
+
+def flatten_args(args, *, detach: bool = False):
+    """Flatten ``args`` into a list, optionally detaching tensors.
+
+    Args:
+        args: Nested arguments to flatten.
+        detach: If ``True``, detach tensors while preserving ``requires_grad``.
+
+    Returns:
+        ``(new_args, flat_detached_args)`` when ``detach=True``;
+        ``flat_args`` list otherwise.
+    """
+    flat_args, treespec = tree_flatten(args)
+
+    if detach:
+        flat_detached = [
+            a.detach().requires_grad_(a.requires_grad)
+            if isinstance(a, torch.Tensor)
+            else a
+            for a in flat_args
+        ]
+        new_args = tree_unflatten(flat_detached, treespec)
+        return new_args, flat_detached
+
+    return flat_args
+
+
+# Backward compatibility alias
+def flatten_args_detach(args):
+    """Flatten and detach. Deprecated: use ``flatten_args(args, detach=True)``."""
+    return flatten_args(args, detach=True)
 
 
 def generate_stage_to_rank_mapping(
@@ -148,7 +563,7 @@ def generate_rank_to_stage_mapping(
     return rank_to_stages
 
 
-@dataclass
+@dataclass(slots=True)
 class PipeInfo:
     """
     Captures information for a pipeline (`Pipe` object).
@@ -157,3 +572,367 @@ class PipeInfo:
     graph: fx.Graph
     num_stages: int
     has_loss_and_backward: bool
+
+
+# ============================================================================
+# Metadata extraction helpers
+# ============================================================================
+
+
+def extract_tensor_meta(tensor: torch.Tensor) -> TensorMeta:
+    """Extract metadata from a tensor.
+
+    Handles both plain Tensor and DTensor correctly: DTensors are
+    dispatched to ``_DTensorMeta.from_dtensor`` which captures local
+    shard attributes plus global shape/placement info, while plain
+    tensors use ``_TensorMeta.from_tensor``.
+
+    Args:
+        tensor: A plain tensor or DTensor.
+
+    Returns:
+        ``_TensorMeta`` for plain tensors, ``_DTensorMeta`` for DTensors.
+    """
+    if isinstance(tensor, DTensor):
+        return _DTensorMeta.from_dtensor(tensor)
+    else:
+        return _TensorMeta.from_tensor(tensor)
+
+
+@overload
+def extract_tensor_metas(
+    tensors: tuple[torch.Tensor, ...] | None,
+    *,
+    allow_none: Literal[False] = ...,
+) -> tuple[TensorMeta, ...] | None: ...
+
+
+@overload
+def extract_tensor_metas(
+    tensors: tuple[torch.Tensor | None, ...] | None,
+    *,
+    allow_none: Literal[True],
+) -> tuple[TensorMeta | None, ...] | None: ...
+
+
+def extract_tensor_metas(
+    tensors: tuple[torch.Tensor | None, ...] | tuple[torch.Tensor, ...] | None,
+    *,
+    allow_none: bool = False,
+) -> tuple[TensorMeta | None, ...] | None:
+    """Extract metadata from a tuple of tensors.
+
+    Args:
+        tensors: Tuple of tensors (may include ``None`` when ``allow_none=True``).
+        allow_none: If ``True``, preserve ``None`` elements (for gradients).
+
+    Returns:
+        Tuple of ``TensorMeta``, or ``None`` if ``tensors`` is ``None``.
+
+    Raises:
+        PipeliningMetadataError: If ``None`` found and ``allow_none=False``.
+    """
+    if tensors is None:
+        return None
+
+    metas_with_none: list[TensorMeta | None] = []
+    has_none = False
+    for t in tensors:
+        if isinstance(t, torch.Tensor):
+            metas_with_none.append(extract_tensor_meta(t))
+        else:
+            has_none = True
+            metas_with_none.append(None)
+    if not allow_none and has_none:
+        raise PipeliningMetadataError(
+            "None values are not allowed in tensor metadata tuples. "
+            "Use allow_none=True for optional values."
+        )
+    return tuple(metas_with_none)
+
+
+def to_local_if_dtensor(tensor: torch.Tensor, detach: bool = False) -> torch.Tensor:
+    """Convert a DTensor to its local shard, or return a plain tensor as-is.
+
+    When ``detach=True``, the tensor is detached before conversion —
+    this applies to both DTensors and plain tensors.
+
+    Args:
+        tensor: A tensor that may be a DTensor.
+        detach: If ``True``, detach before ``to_local()`` to avoid
+            redistribution during backward.
+
+    Returns:
+        The local tensor component.
+    """
+    maybe_detached_tensor = tensor.detach() if detach else tensor
+    if isinstance(maybe_detached_tensor, DTensor):
+        return maybe_detached_tensor.to_local()
+    return maybe_detached_tensor
+
+
+@overload
+def validate_and_normalize_to_tuple(
+    args: torch.Tensor | tuple[torch.Tensor, ...] | list[torch.Tensor] | None,
+    allow_none: Literal[False] = ...,
+) -> tuple[torch.Tensor, ...] | None: ...
+
+
+@overload
+def validate_and_normalize_to_tuple(
+    args: torch.Tensor
+    | tuple[torch.Tensor | None, ...]
+    | list[torch.Tensor | None]
+    | None,
+    allow_none: Literal[True] = ...,
+) -> tuple[torch.Tensor | None, ...] | None: ...
+
+
+def validate_and_normalize_to_tuple(
+    args: torch.Tensor
+    | tuple[torch.Tensor, ...]
+    | tuple[torch.Tensor | None, ...]
+    | list[torch.Tensor]
+    | list[torch.Tensor | None]
+    | None,
+    allow_none: bool = False,
+) -> tuple[torch.Tensor | None, ...] | tuple[torch.Tensor, ...] | None:
+    """Normalize ``args`` to a tuple and validate that all elements are tensors.
+
+    Args:
+        args: A single tensor, tuple/list of tensors, or ``None``.
+        allow_none: If ``True``, permit ``None`` elements (for gradients).
+
+    Returns:
+        Tuple of tensors, or ``None`` if ``args`` is ``None``.
+
+    Raises:
+        PipeliningMetadataError: On non-tensor values
+            (or ``None`` when ``allow_none=False``).
+    """
+    if args is None:
+        return None
+    elif isinstance(args, torch.Tensor):
+        return (args,)
+    elif isinstance(args, (tuple, list)):
+        for i, arg in enumerate(args):
+            if arg is None:
+                if not allow_none:
+                    raise PipeliningMetadataError(
+                        f"Stage arg[{i}] is None. "
+                        f"Stage args must be tensors. Use kwargs for optional values."
+                    )
+                continue
+            if not isinstance(arg, torch.Tensor):
+                raise PipeliningMetadataError(
+                    f"Stage arg[{i}] has type {type(arg).__name__}. "
+                    f"All stage args must be tensors. Use kwargs for non-tensor inputs."
+                )
+        # Normalize list to tuple
+        return tuple(args) if isinstance(args, list) else args
+    else:
+        raise PipeliningMetadataError(
+            f"Stage args must be a tensor, tuple, or list of tensors, got {type(args).__name__}."
+        )
+
+
+# ============================================================================
+# Validation functions
+# ============================================================================
+
+
+def validate_metadata(
+    desc: str,
+    expected: TensorMeta,
+    actual: torch.Tensor | TensorMeta,
+    *,
+    raise_on_mismatch: bool = False,
+    warn_on_mismatch: bool = False,
+) -> list[str]:
+    """
+    Compare expected metadata against actual tensor or metadata.
+
+    This is the unified validation/comparison function that uses get_diff() from
+    metadata classes. Works with both plain tensors and DTensors.
+
+    For plain tensors: compares shape/stride/dtype/requires_grad.
+    For DTensors: compares all properties including global shape and placements.
+
+    Args:
+        desc: Description for error/warning messages.
+        expected: Expected tensor metadata (_TensorMeta or _DTensorMeta).
+        actual: Actual tensor or metadata to compare against.
+        raise_on_mismatch: If True, raise PipeliningMetadataError on mismatch.
+        warn_on_mismatch: If True, issue a warning on mismatch.
+
+    Returns:
+        List of differences (empty if metadata matches).
+
+    Raises:
+        PipeliningMetadataError: If raise_on_mismatch=True and differences exist.
+    """
+    # Extract metadata if actual is a tensor
+    if isinstance(actual, torch.Tensor):
+        actual_meta = extract_tensor_meta(actual)
+    else:
+        actual_meta = actual
+
+    # Type check: ensure both are same type for meaningful comparison
+    if type(expected) is not type(actual_meta):
+        type_diff = [
+            f"type: expected {type(expected).__name__}, got {type(actual_meta).__name__}"
+        ]
+        if raise_on_mismatch:
+            raise PipeliningMetadataError(f"{desc}: {type_diff[0]}")
+        if warn_on_mismatch:
+            warnings.warn(
+                f"{desc}: Metadata type mismatch. {type_diff[0]}. "
+                f"Using dynamically inferred metadata instead.",
+                UserWarning,
+                stacklevel=2,
+            )
+        return type_diff
+
+    # Use get_diff() from the metadata class
+    diffs = expected.get_diff(actual_meta)
+
+    if diffs:
+        if raise_on_mismatch:
+            raise PipeliningMetadataError(f"{desc}: {'; '.join(diffs)}")
+        if warn_on_mismatch:
+            warnings.warn(
+                f"{desc}: Metadata mismatch. {'; '.join(diffs)}. "
+                f"Using dynamically inferred metadata instead.",
+                UserWarning,
+                stacklevel=2,
+            )
+
+    return diffs
+
+
+def validate_tensors_metadata(
+    desc: str,
+    expected: tuple[TensorMeta | None, ...],
+    actual: tuple[torch.Tensor | TensorMeta | None, ...],
+    *,
+    raise_on_mismatch: bool = True,
+    warn_on_mismatch: bool = False,
+) -> list[str]:
+    """Validate metadata for a tuple of tensors element-wise.
+
+    Args:
+        desc: Description prefix for error/warning messages.
+        expected: Tuple of expected metadata (may include ``None`` for grads).
+        actual: Tuple of actual tensors or metadata to compare against.
+        raise_on_mismatch: If ``True``, raise on the first mismatch.
+        warn_on_mismatch: If ``True``, issue warnings for mismatches.
+
+    Returns:
+        Aggregated list of difference strings.
+
+    Raises:
+        PipeliningMetadataError: If lengths differ or on mismatch.
+    """
+    if len(expected) != len(actual):
+        msg = f"{desc}: expected {len(expected)} tensors, got {len(actual)}"
+        if raise_on_mismatch:
+            raise PipeliningMetadataError(msg)
+        if warn_on_mismatch:
+            warnings.warn(msg, UserWarning, stacklevel=2)
+        return [msg]
+
+    all_diffs: list[str] = []
+    for i, (exp, act) in enumerate(zip(expected, actual, strict=True)):
+        if exp is None and act is None:
+            continue
+        if exp is None or act is None:
+            msg = (
+                f"{desc}[{i}]: expected {'None' if exp is None else 'metadata'}, "
+                f"got {'None' if act is None else 'metadata'}"
+            )
+            if raise_on_mismatch:
+                raise PipeliningMetadataError(msg)
+            if warn_on_mismatch:
+                warnings.warn(msg, UserWarning, stacklevel=2)
+            all_diffs.append(msg)
+            continue
+        diffs = validate_metadata(
+            f"{desc}[{i}]",
+            exp,
+            act,
+            raise_on_mismatch=raise_on_mismatch,
+            warn_on_mismatch=warn_on_mismatch,
+        )
+        all_diffs.extend(diffs)
+    return all_diffs
+
+
+def validate_static_arg_grad_correspondence(
+    stage_index: int,
+    args: tuple[torch.Tensor, ...],
+    grads: tuple[torch.Tensor | None, ...],
+    is_input: bool,
+) -> None:
+    """
+    Validate the args↔grads contract for static mode.
+
+    Enforces four rules for each (arg, grad) pair:
+      1. len(args) must equal len(grads).
+      2. If arg.requires_grad is False, grad must be None.
+      3. If arg.requires_grad is True and grad is None, emit a warning
+         (this is legal at pipeline boundaries but may indicate a bug).
+      4. If arg is a DTensor with requires_grad=True and grad is not None,
+         grad must also be a DTensor.
+
+    Args:
+        stage_index: The stage index for error messages.
+        args: Tuple of forward tensors.
+        grads: Tuple of gradient tensors (can include None).
+        is_input: True for input_args/input_grads, False for output_args/output_grads.
+
+    Raises:
+        PipeliningMetadataError: If any hard rule (1, 2, or 4) is violated.
+    """
+    kind = "input" if is_input else "output"
+    args_name = f"{kind}_args"
+    grads_name = f"{kind}_grads"
+
+    # Rule 1: lengths must match
+    if len(args) != len(grads):
+        raise PipeliningMetadataError(
+            f"Stage {stage_index}: {grads_name} length ({len(grads)}) does not match "
+            f"{args_name} length ({len(args)}). Each forward tensor must have a "
+            f"corresponding gradient entry (use None for tensors that don't require grad)."
+        )
+
+    for i, (arg, grad) in enumerate(zip(args, grads, strict=True)):
+        # Rule 2: no grad for a non-differentiable arg
+        if not arg.requires_grad and grad is not None:
+            raise PipeliningMetadataError(
+                f"Stage {stage_index}: {args_name}[{i}] has requires_grad=False, "
+                f"but {grads_name}[{i}] is not None ({type(grad).__name__}). "
+                f"Non-differentiable tensors must have None as their gradient entry."
+            )
+
+        # Rule 3: missing grad for a differentiable arg (warn, don't raise)
+        if arg.requires_grad and grad is None:
+            warnings.warn(
+                f"Stage {stage_index}: {args_name}[{i}] has requires_grad=True, "
+                f"but {grads_name}[{i}] is None. This is legal at pipeline boundaries "
+                f"but may indicate a missing gradient.",
+                UserWarning,
+                stacklevel=2,
+            )
+
+        # Rule 4: DTensor arg must have DTensor grad
+        if (
+            isinstance(arg, DTensor)
+            and arg.requires_grad
+            and grad is not None
+            and not isinstance(grad, DTensor)
+        ):
+            raise PipeliningMetadataError(
+                f"Stage {stage_index}: {args_name}[{i}] is a DTensor with requires_grad=True, "
+                f"but {grads_name}[{i}] is {type(grad).__name__}, expected DTensor or None. "
+                f"DTensor gradients may have different placements than forward tensors."
+            )
diff --git a/torch/distributed/pipelining/microbatch.py b/torch/distributed/pipelining/microbatch.py
index c1c353b191bb5..5b28c1942a59b 100644
--- a/torch/distributed/pipelining/microbatch.py
+++ b/torch/distributed/pipelining/microbatch.py
@@ -6,6 +6,8 @@
 from typing import Any
 
 import torch
+from torch.distributed.tensor import DTensor
+from torch.distributed.tensor.experimental import local_map
 from torch.fx.node import map_aggregate
 from torch.nn.attention.flex_attention import BlockMask
 from torch.utils._pytree import tree_flatten, tree_map, tree_unflatten
@@ -202,26 +204,60 @@ def _split_tensor(
         raise AssertionError(
             f"Tensor size {tensor.size(spec.split_dim)} is smaller than num_chunks"
         )
-    chunk_tensors = torch.tensor_split(tensor, num_chunks, spec.split_dim)
 
-    if not _debug_mask_minibatches:
-        return chunk_tensors
-
-    expanded_chunks = []
-    split_dim_idx = 0
-    for chunk_tensor in chunk_tensors:
-        new_val = torch.zeros_like(tensor)
-        upper_idx = split_dim_idx + chunk_tensor.size(spec.split_dim)
-
-        slice_indices = [slice(None, None, None)] * new_val.ndim
-        slice_indices[spec.split_dim] = slice(split_dim_idx, upper_idx)
-        new_val[slice_indices] = chunk_tensor
+    _is_dtensor = isinstance(tensor, DTensor)
+
+    if _is_dtensor:
+        # Use local_map to split locally and preserve placements.
+        # Going through DTensor dispatch would convert Shard(split_dim) to
+        # Replicate() via an implicit all-gather, which is both wasteful and
+        # semantically wrong for PP microbatch splitting.
+        placements = tensor.placements
+        split_fn = local_map(
+            lambda t: torch.tensor_split(t, num_chunks, spec.split_dim),
+            out_placements=(placements,) * num_chunks,
+            in_placements=(placements,),
+        )
+        chunk_tensors: Sequence[torch.Tensor] = split_fn(tensor)  # type: ignore[assignment]
+    else:
+        chunk_tensors = torch.tensor_split(tensor, num_chunks, spec.split_dim)
 
-        expanded_chunks.append(new_val)
+    # tensor_split on a leaf tensor produces non-leaf views that won't
+    # accumulate .grad during torch.autograd.backward().  Call retain_grad()
+    # on those views so that stage_backward() can read .grad from them.
+    if tensor.requires_grad and tensor.is_leaf:
+        for chunk in chunk_tensors:
+            chunk.retain_grad()
 
-        split_dim_idx += chunk_tensor.size(spec.split_dim)
+    if not _debug_mask_minibatches:
+        return chunk_tensors
 
-    return expanded_chunks
+    def _expand_chunks(
+        orig: torch.Tensor, *chunks: torch.Tensor
+    ) -> tuple[torch.Tensor, ...]:
+        expanded = []
+        idx = 0
+        for chunk in chunks:
+            new_val = torch.zeros_like(orig)
+            upper = idx + chunk.size(spec.split_dim)
+            slices: list[slice] = [slice(None)] * new_val.ndim
+            slices[spec.split_dim] = slice(idx, upper)
+            new_val[slices] = chunk
+            expanded.append(new_val)
+            idx += chunk.size(spec.split_dim)
+        return tuple(expanded)
+
+    if _is_dtensor:
+        placements = tensor.placements
+        n = len(chunk_tensors)
+        expand_fn = local_map(
+            _expand_chunks,
+            out_placements=(placements,) * n,
+            in_placements=(placements,) + (placements,) * n,
+        )
+        return list(expand_fn(tensor, *chunk_tensors))  # type: ignore[arg-type]
+    else:
+        return list(_expand_chunks(tensor, *chunk_tensors))
 
 
 def _shard_dict_of_args(
@@ -538,7 +574,31 @@ def merge_chunks(
             else:
                 values_to_cat = partial_values
 
-            args_flattened.append(torch.cat(values_to_cat, dim=arg.split_dim))
+            # Validate DTensor consistency: either all values are DTensors
+            # or none are. A mix indicates a bug in the pipeline stage.
+            dtensor_flags = [isinstance(v, DTensor) for v in values_to_cat]
+            if any(dtensor_flags):
+                if not all(dtensor_flags):
+                    raise AssertionError(
+                        "merge_chunks: expected all values to be DTensors or "
+                        "none to be DTensors, got a mix"
+                    )
+                # All DTensors must have matching placements.
+                placements = values_to_cat[0].placements
+                for i, v in enumerate(values_to_cat[1:], 1):
+                    if v.placements != placements:
+                        raise AssertionError(
+                            f"merge_chunks: placement mismatch at chunk {i}: "
+                            f"expected {placements}, got {v.placements}"
+                        )
+                cat_fn = local_map(
+                    lambda *chunks: torch.cat(chunks, dim=arg.split_dim),
+                    out_placements=(placements,),
+                    in_placements=tuple(placements for _ in range(len(values_to_cat))),
+                )
+                args_flattened.append(cat_fn(*values_to_cat))
+            else:
+                args_flattened.append(torch.cat(values_to_cat, dim=arg.split_dim))
         elif isinstance(arg, _CustomReducer):
             reduced_val = arg.init_value
 
diff --git a/torch/distributed/pipelining/schedules.py b/torch/distributed/pipelining/schedules.py
index 0367f599053e3..97ca83b019cc9 100644
--- a/torch/distributed/pipelining/schedules.py
+++ b/torch/distributed/pipelining/schedules.py
@@ -21,9 +21,13 @@
 from torch.nn.modules.loss import _Loss
 from torch.profiler import record_function
 
-from ._utils import generate_rank_to_stage_mapping, generate_stage_to_rank_mapping
+from ._utils import (
+    generate_rank_to_stage_mapping,
+    generate_stage_to_rank_mapping,
+    InferenceMode,
+)
 from .microbatch import merge_chunks, split_args_kwargs_into_chunks, TensorChunkSpec
-from .stage import _PipelineStageBase
+from .stage import _PipelineStageBase, PipelineStage
 
 
 __all__ = [
@@ -320,6 +324,135 @@ def _update_losses(self, stages, losses):
 
         self._internal_losses.clear()
 
+    def _warmup_p2p(
+        self,
+        stages: list[_PipelineStageBase],
+        has_backward: bool,
+        p2p_done: bool,
+    ) -> None:
+        """Run the P2P warm-up protocol for the given stages.
+
+        For ``PipelineStage`` instances this executes the forward/backward vote
+        protocol (which warms up 2-rank sub-communicators) and sets each
+        stage's ``_inference_mode``.  For other stage types it falls back to
+        the legacy ``_get_init_p2p_neighbors_ops`` + ``_batch_p2p`` path.
+
+        Args:
+            stages: The pipeline stages owned by this rank.
+            has_backward: Whether the schedule includes a backward pass.
+            p2p_done: ``True`` if P2P neighbours have already been initialised
+                (avoids redundant init on eval↔train mode switches).
+        """
+        if all(isinstance(stage, PipelineStage) for stage in stages):
+            acc: torch.Tensor | None = None
+            for stage in cast(list[PipelineStage], stages):
+                acc = stage._warmup_forward_vote(has_backward, received_acc=acc)
+            result: torch.Tensor | None = acc
+            determined_mode: InferenceMode | None = None
+            for stage in reversed(cast(list[PipelineStage], stages)):
+                result = stage._warmup_backward_result(received_result=result)
+                if result is None:
+                    raise RuntimeError("P2P warm-up voting failed")
+                determined_mode = (
+                    InferenceMode.STATIC
+                    if result.item() == 1
+                    else InferenceMode.DYNAMIC
+                )
+                stage._inference_mode = determined_mode
+            logger.debug(
+                "Rank determined inference_mode=%s for %d stage(s)",
+                determined_mode.value if determined_mode else "None",
+                len(stages),
+            )
+        elif not p2p_done:
+            all_ops: list[dist.P2POp] = []
+            for stage in stages:
+                all_ops.extend(stage._get_init_p2p_neighbors_ops())
+            _wait_batch_p2p(_batch_p2p(all_ops))
+
+        # TODO: STATIC mode group communicator warm-up gap
+        # The vote protocol above warms up 2-rank sub-communicators
+        # (used by `_batch_p2p` homogeneous fast-path).  In DYNAMIC mode,
+        # `_send_meta`/`_recv_meta` (called during `_prepare_forward_infra` →
+        # `_forward_metadata_inference`) also warm up the *group* communicator
+        # (used by `_batch_p2p` mixed-op path).  In STATIC mode, metadata
+        # inference is skipped, so the group communicator is NOT warmed up —
+        # it will be lazily created on the first mixed `_batch_p2p` call
+        # (e.g., 1F1B steady-state with both sends and recvs).
+        # Fix: run `_get_init_p2p_neighbors_ops` + `_batch_p2p` after the
+        # vote, gated by `not p2p_done`.
+
+    def _initialize_pp_stages(
+        self,
+        stages: list[_PipelineStageBase],
+        args: tuple[Any, ...] | Any,
+        kwargs: dict[str, Any] | None,
+        target: Any,
+        fwd_initialized: bool,
+        bwd_initialized: bool,
+    ) -> tuple[bool, bool]:
+        """Common stage initialization shared by Single and Multi schedules.
+
+        Handles mode-change detection (eval↔train), P2P warm-up, RNG forking,
+        forward / backward metadata inference, and FSDP cleanup.
+
+        Returns the updated ``(fwd_initialized, bwd_initialized)`` flags.
+        """
+        # Detect eval↔train mode switch: if has_backward changed since last
+        # init, re-initialize both fwd (recv buffers need different
+        # requires_grad) and bwd.  p2p_done avoids redundant P2P warm-up.
+        p2p_done = fwd_initialized
+        if fwd_initialized and (self._has_backward != bwd_initialized):
+            fwd_initialized = False
+            bwd_initialized = False
+
+        needs_fwd = not fwd_initialized
+        needs_bwd = self._has_backward and not bwd_initialized
+
+        if not needs_fwd and not needs_bwd:
+            return fwd_initialized, bwd_initialized
+
+        if needs_fwd:
+            self._warmup_p2p(stages, self._has_backward, p2p_done)
+
+        # Fork RNG so metadata inference doesn't perturb training RNG.
+        devices = list(
+            {
+                torch.device(stage.device)
+                for stage in stages
+                if torch.device(stage.device).type != "cpu"
+            }
+        )
+        with torch.random.fork_rng(devices=devices):
+            if needs_fwd:
+                next_stage_args: Any = None
+                for stage in stages:
+                    stage_args = args if stage.is_first else next_stage_args
+                    next_stage_args = stage._prepare_forward_infra(
+                        self._n_microbatches,
+                        stage_args,
+                        kwargs,
+                        has_backward=self._has_backward,
+                    )
+                fwd_initialized = True
+
+            if needs_bwd:
+                prev_stage_grad_meta: Any = None
+                for stage in reversed(stages):
+                    prev_stage_grad_meta = stage._prepare_backward_infra(
+                        self._n_microbatches,
+                        loss_fn=self._loss_fn,
+                        target=target,
+                        received_grad_meta=prev_stage_grad_meta,
+                    )
+                bwd_initialized = True
+
+        for stage in stages:
+            if isinstance(stage, PipelineStage):
+                stage._post_metadata_inference_cleanup()
+
+        return fwd_initialized, bwd_initialized
+
     @abstractmethod
     def _step_microbatches(
         self,
@@ -459,12 +592,29 @@ def _merge_outputs(self, output_chunks: list[Any]) -> Any:
 
 def _batch_p2p(p2p_ops: list[dist.P2POp], desc: str | None = None) -> list[dist.Work]:
     """
-    Simple wrapper over batch_isend_irecv from torch.distributed, which just adds a descriptive logger on top.
+    Wrapper over batch_isend_irecv that avoids coalescing for homogeneous
+    batches (all-send or all-recv).  Coalescing serializes ops on a single
+    CUDA stream, which causes head-of-line blocking when independent P2P ops
+    could otherwise overlap.  Mixed batches still use batch_isend_irecv for
+    deadlock avoidance.
     """
     if len(p2p_ops) == 0:
         return []
     desc_str = f"{desc}, " if desc else ""
     logger.debug("batch_p2p %s%s", desc_str, p2p_ops)
+
+    op_types = {p.op for p in p2p_ops}
+    if op_types == {dist.isend}:
+        return [
+            p.op(p.tensor, group=p.group, tag=p.tag, group_dst=p.group_peer)
+            for p in p2p_ops
+        ]
+    if op_types == {dist.irecv}:
+        return [
+            p.op(p.tensor, group=p.group, tag=p.tag, group_src=p.group_peer)
+            for p in p2p_ops
+        ]
+
     return dist.batch_isend_irecv(p2p_ops)
 
 
@@ -543,21 +693,18 @@ def __init__(
             self._get_pipeline_order()
         )
 
-    def _initialize_stage(self, args, kwargs):
-        if not self._stage_forward_initialized:
-            # Prepare the communication needed for the pipeline schedule execution
-            # This is needed because during execution we always perform a series of batch P2P ops
-            # The first call of the batched P2P needs to involve the global group
-            all_ops: list[dist.P2POp] = []
-            all_ops.extend(self._stage._get_init_p2p_neighbors_ops())
-            _wait_batch_p2p(_batch_p2p(all_ops))
-
-            self._stage._prepare_forward_infra(self._n_microbatches, args, kwargs)
-            self._stage_forward_initialized = True
-
-        if self._has_backward and not self._stage_backward_initialized:
-            self._stage._prepare_backward_infra(self._n_microbatches)
-            self._stage_backward_initialized = True
+    def _initialize_stage(self, args, kwargs, target=None):
+        (
+            self._stage_forward_initialized,
+            self._stage_backward_initialized,
+        ) = self._initialize_pp_stages(
+            [self._stage],
+            args,
+            kwargs,
+            target,
+            self._stage_forward_initialized,
+            self._stage_backward_initialized,
+        )
 
     def step(
         self,
@@ -653,7 +800,8 @@ def _step_microbatches(
             )
 
         arg_mbs, kwarg_mbs = self._check_inputs(arg_mbs, kwarg_mbs, target_mbs, losses)
-        self._initialize_stage(arg_mbs[0], kwarg_mbs[0])
+        maybe_first_target = target_mbs[0] if target_mbs is not None else None
+        self._initialize_stage(arg_mbs[0], kwarg_mbs[0], maybe_first_target)
 
         # Delay send waits
         fwd_sends_to_wait: list[list[dist.Work]] = []
@@ -704,7 +852,8 @@ def _step_microbatches(
             return_outputs: whether to return the outputs from the last stage.
         """
         arg_mbs, kwarg_mbs = self._check_inputs(arg_mbs, kwarg_mbs, target_mbs, losses)
-        self._initialize_stage(arg_mbs[0], kwarg_mbs[0])
+        maybe_first_target = target_mbs[0] if target_mbs is not None else None
+        self._initialize_stage(arg_mbs[0], kwarg_mbs[0], maybe_first_target)
 
         # Delay send waits
         fwd_sends_to_wait: list[list[dist.Work]] = []
@@ -848,7 +997,8 @@ def _step_microbatches(
             return_outputs: whether to return the outputs from the last stage.
         """
         arg_mbs, kwarg_mbs = self._check_inputs(arg_mbs, kwarg_mbs, target_mbs, losses)
-        self._initialize_stage(arg_mbs[0], kwarg_mbs[0])
+        maybe_first_target = target_mbs[0] if target_mbs is not None else None
+        self._initialize_stage(arg_mbs[0], kwarg_mbs[0], maybe_first_target)
 
         # Last stage has 1 warmup, second-to-last 2 warmups, ...
         # first stage `num_stages` warmups
@@ -1515,34 +1665,18 @@ def __init__(
                 "Simply stop passing it, and everything should still work fine."
             )
 
-    def _initialize_stages(self, args: tuple[Any, ...], kwargs):
-        if not self._stages_forward_initialized:
-            # Prepare the communication needed for the pipeline schedule execution
-            # This is needed because during execution we always perform a series of batch P2P ops
-            # The first call of the batched P2P needs to involve the global group
-            all_ops: list[dist.P2POp] = []
-            for stage in self._stages:
-                all_ops.extend(stage._get_init_p2p_neighbors_ops())
-            _wait_batch_p2p(_batch_p2p(all_ops))
-
-            # may be 'none' value (if this stage sends its output shapes to the next stage via P2P)
-            # or real value (if this stage and next stage are on the same device)
-            next_stage_args: tuple[Any, ...] = tuple()
-            for stage in self._stages:
-                if stage.is_first:
-                    next_stage_args = stage._prepare_forward_infra(
-                        self._n_microbatches, args, kwargs
-                    )
-                else:
-                    next_stage_args = stage._prepare_forward_infra(
-                        self._n_microbatches, next_stage_args, kwargs
-                    )
-            self._stages_forward_initialized = True
-
-        if self._has_backward and not self._stages_backward_initialized:
-            for stage in self._stages:
-                stage._prepare_backward_infra(self._n_microbatches)
-            self._stages_backward_initialized = True
+    def _initialize_stages(self, args: tuple[Any, ...], kwargs, target=None):
+        (
+            self._stages_forward_initialized,
+            self._stages_backward_initialized,
+        ) = self._initialize_pp_stages(
+            self._stages,
+            args,
+            kwargs,
+            target,
+            self._stages_forward_initialized,
+            self._stages_backward_initialized,
+        )
 
     def _validate_and_set_stage_mapping(
         self, actions: dict[int, list[_Action | None]]
@@ -1657,8 +1791,8 @@ def _step_microbatches(
         not support models with skip connections.
         """
         arg_mbs, kwarg_mbs = self._check_inputs(arg_mbs, kwarg_mbs, target_mbs, losses)
-
-        self._initialize_stages(arg_mbs[0], kwarg_mbs[0])
+        maybe_first_target = target_mbs[0] if target_mbs is not None else None
+        self._initialize_stages(arg_mbs[0], kwarg_mbs[0], maybe_first_target)
 
         # Based on the plan in Step 1 created in __init__:
         # 2. Perform communication based on the pipeline_order
@@ -1822,7 +1956,7 @@ def _step_microbatches(
 at time_step %s when running action %s",
                     self.rank,
                     self.__class__.__name__,
-                    str(e),
+                    e,
                     time_step,
                     action,
                 )
@@ -2043,7 +2177,8 @@ def _step_microbatches(
         not support models with skip connections.
         """
         arg_mbs, kwarg_mbs = self._check_inputs(arg_mbs, kwarg_mbs, target_mbs, losses)
-        self._initialize_stages(arg_mbs[0], kwarg_mbs[0])
+        maybe_first_target = target_mbs[0] if target_mbs is not None else None
+        self._initialize_stages(arg_mbs[0], kwarg_mbs[0], maybe_first_target)
 
         # Based on the plan in Step 1 created in __init__:
         # 2. Perform communication based on the pipeline_order
diff --git a/torch/distributed/pipelining/stage.py b/torch/distributed/pipelining/stage.py
index 7862a5e64a6ce..0fbe780714869 100644
--- a/torch/distributed/pipelining/stage.py
+++ b/torch/distributed/pipelining/stage.py
@@ -2,9 +2,10 @@
 # Copyright (c) Meta Platforms, Inc. and affiliates
 import logging
 import operator
+import warnings
 from abc import ABC, abstractmethod
 from collections.abc import Callable
-from typing import Any, cast, Union
+from typing import Any, cast
 
 import torch
 import torch.distributed as dist
@@ -13,13 +14,39 @@
 from torch._subclasses.fake_tensor import FakeTensor
 from torch.distributed._composable.replicate_with_fsdp import replicate, ReplicateModule
 from torch.distributed.fsdp import FSDPModule, fully_shard
+from torch.distributed.pipelining._utils import (
+    _derive_grad_metas,
+    _DTensorMeta,
+    _make_tensor_from_meta,
+    _MeshCache,
+    _StageBackwardMeta,
+    _StageForwardMeta,
+    _StageMeta,
+    _TensorMeta,
+    extract_tensor_meta,
+    extract_tensor_metas,
+    flatten_args,
+    GetMeshCallback,
+    InferenceMode,
+    PipeInfo,
+    PipeliningMetadataError,
+    TensorMeta,
+    to_local_if_dtensor,
+    validate_and_normalize_to_tuple,
+    validate_static_arg_grad_correspondence,
+    validate_tensors_metadata,
+)
+from torch.distributed.tensor import DTensor
 from torch.fx.node import Argument, map_aggregate
 from torch.nn.parallel import DistributedDataParallel
-from torch.utils._pytree import tree_map_only
 
-from ._backward import stage_backward, stage_backward_input, stage_backward_weight
+from ._backward import (
+    _autograd_grad_for_inputs,
+    stage_backward,
+    stage_backward_input,
+    stage_backward_weight,
+)
 from ._debug import map_debug_info
-from ._utils import flatten_args, PipeInfo, validate_tensors_metadata
 
 
 __all__ = [
@@ -63,61 +90,47 @@ def _normalize_model_output_as_tuple(output: Any) -> tuple[Any]:
     return output_tuple
 
 
-class _RootArgPlaceholder:
-    """
-    Placeholder for model-level inputs.
-    """
-
-    def __init__(self, tensor):
-        self.meta = tensor.to("meta")
-
-
 class _RecvInfo:
-    """
-    Represents a stage input.
+    """Input tensor descriptor for a pipeline stage.
+
+    Handles both received activations from a previous stage
+    (``is_root_arg=False``) and root-level model inputs provided
+    by the user (``is_root_arg=True``).
     """
 
     def __init__(
         self,
         input_name: str,
-        source: int,
-        buffer: torch.Tensor,
+        source: int | None,
+        buffer: torch.Tensor | None,
+        tensor_meta: TensorMeta | None,
+        *,
+        is_root_arg: bool = False,
     ):
         # Name of this input
         self.input_name = input_name
-        # Stage index of the source of this input
+        # Stage index of the source of this input (None for root args)
         self.source = source
-        # Buffer to receive the input into.
+        # Buffer to receive the input into (None for root args)
         self.buffer = buffer
+        # Tensor metadata for validation and DTensor reconstruction
+        self.tensor_meta = tensor_meta
+        # Whether this is a root-level model input (no recv needed)
+        self.is_root_arg = is_root_arg
 
     def __repr__(self):
-        return f"_RecvInfo(input={self.input_name}, source={self.source}, shape={self.buffer.size()})"
-
-
-# An input can be either a received activation or a model input
-InputInfo = Union[_RecvInfo, _RootArgPlaceholder]
-
-
-def _make_tensor_from_meta(
-    example: torch.Tensor | FakeTensor,
-    device: torch.device,
-) -> torch.Tensor:
-    """
-    Create a real tensor from a tensor.
-    """
-    return torch.empty(
-        example.size(),
-        dtype=example.dtype,
-        layout=example.layout,
-        device=device,
-    )
+        if self.is_root_arg:
+            return f"_RecvInfo(input={self.input_name}, root_arg=True)"
+        meta_type = type(self.tensor_meta).__name__ if self.tensor_meta else "None"
+        buffer_shape = self.buffer.size() if self.buffer is not None else "None"
+        return f"_RecvInfo(input={self.input_name}, source={self.source}, shape={buffer_shape}, meta={meta_type})"
 
 
 class _PipelineStageBase(ABC):
-    """
-    Base class for pipeline stages.
-    Defines or implements common methods used by the `_PipelineStage` used by
-    the tracing frontend and `PipelineStage` used by manual frontend.
+    """Base class for pipeline stages.
+
+    Defines common methods used by ``_PipelineStage`` (tracing frontend)
+    and ``PipelineStage`` (manual frontend).
     """
 
     def __init__(
@@ -131,20 +144,16 @@ def __init__(
     ):
         """
         Args:
-            submodule (torch.nn.Module): The module to be executed in this stage.
-            stage_index (int): The index of this stage.
-            num_stages (int): The total number of stages in this pipeline.
-            device (torch.device): The device to run this stage on.
-            group (Optional[dist.ProcessGroup]): The process group to use for communication.
-                If `None`, the default process group will be used.
-                Default: `None`.
-            dw_builder (Optional[Callable[[], Callable[..., None]]): If provided, dw_builder is a builder function
-                that will build a new dw_runner function that will run parts of module backward that were intentionally
-                skipped during the module's actual backward pass. The builder must be invoked by stage after stage runs
-                model backwards, and stage should save the latest dw_runner to run during weight pas (W).
-                If not provided, a dw_runner will be generated automatically by traversing the autograd graph.
-                When used with schedules that only have F and B steps, the fresh dw_runner function will be called as
-                part of I (input backwards). When used with F,I,W schedules, the dw_runner function implements 'W'.
+            submodule: The module to be executed in this stage.
+            stage_index: The index of this stage.
+            num_stages: The total number of stages in this pipeline.
+            device: The device to run this stage on.
+            group: Process group for communication. Defaults to the
+                default process group if ``None``.
+            dw_builder: Builder function that produces a ``dw_runner``
+                for deferred weight updates in F/I/W zero-bubble
+                schedules. If ``None``, a runner is generated
+                automatically via autograd graph traversal.
         """
         super().__init__()
         if stage_index >= num_stages:
@@ -155,7 +164,6 @@ def __init__(
         self.submod = submodule
         self.stage_index = stage_index
         self.num_stages = num_stages
-        # pyrefly: ignore [read-only]
         self.device = device
         self.group = group
 
@@ -176,7 +184,6 @@ def __init__(
             )
 
         # Run time states
-        self._outputs_meta: tuple[torch.Tensor, ...] | None = None
         # map microbatch ID to list of forward tensor args
         self.fwd_cache: dict[int, tuple[Any, list[torch.Tensor]]] = {}
         # map microbatch ID to list of backward grad tensor args
@@ -191,7 +198,7 @@ def __init__(
         self.log_prefix = f"[Stage {self.stage_index}]"
 
         # Forward infra
-        self.args_recv_info: dict[int, tuple[InputInfo, ...]] = {}
+        self.args_recv_info: dict[int, tuple[_RecvInfo, ...]] = {}
         self.act_send_info: dict[int, list] = {}
 
         # Backward infra will created lazily
@@ -204,6 +211,17 @@ def __init__(
             i: i % self.group_size for i in range(self.num_stages)
         }
 
+        # DTensor support: mesh cache for looking up DeviceMesh by (dim_names, layout)
+        self._mesh_cache = _MeshCache()
+
+        # Per-chunk runtime validation is expensive; only enable under
+        # TORCH_DISTRIBUTED_DEBUG=DETAIL for debugging shape/dtype mismatches.
+        self._runtime_validate = dist.get_debug_level() == dist.DebugLevel.DETAIL
+
+        # DTensor support: consolidated stage metadata container
+        # Contains inputs, outputs, input_grads, output_grads metadata
+        self._stage_meta = _StageMeta()
+
     @property
     def has_backward(self) -> bool:
         """
@@ -229,6 +247,21 @@ def is_last(self):
         """
         return self.stage_index == self.num_stages - 1
 
+    def _validate_stage_tensors(
+        self,
+        desc: str,
+        expected: tuple[TensorMeta | None, ...] | None,
+        actual: tuple[torch.Tensor | None, ...],
+    ) -> None:
+        """Validate actual tensors against expected metadata.
+
+        Raises:
+            PipeliningMetadataError: If metadata is missing or mismatched.
+        """
+        if expected is None:
+            raise PipeliningMetadataError(f"{desc}: no metadata available")
+        validate_tensors_metadata(desc, expected, actual)
+
     def _check_chunk_id(self, chunk_id: int):
         if self.chunks is None:
             raise RuntimeError(
@@ -239,27 +272,6 @@ def _check_chunk_id(self, chunk_id: int):
                 f"Chunk id {chunk_id} is out of range [0, {self.chunks})"
             )
 
-    def _configure_outputs_meta(self, outputs_meta: tuple[torch.Tensor, ...]):
-        """
-        Track the output shapes/dtype of this stage since they determine the send operation(s) which must match
-        recv operations of the next stage.  The next stage _will_ be freezing its recv buffers based on its initial
-        configuration, so it's important to also freeze/validate the output side to avoid any send/recv mismatches
-        which could show up as hangs, silent corruption, or other errors.
-        """
-        if self._outputs_meta is not None:
-            raise AssertionError(
-                "Attempting to reconfigure output_meta, which is not supported"
-            )
-        self._outputs_meta = tuple(outputs_meta)  # type: ignore[assignment]
-
-    def get_outputs_meta(self) -> tuple[torch.Tensor, ...]:
-        """Get the output metadata (meta tensors) reprensenting the outputs of this stage"""
-        if self._outputs_meta is None:
-            raise AssertionError(
-                "Attempted to get_outputs_meta() without configuring output meta"
-            )
-        return self._outputs_meta
-
     def _create_grad_send_info(
         self,
         args_recv_info: tuple,
@@ -273,12 +285,13 @@ def map_recv_to_send(a):
             # Note: we send gradients back to previous stage as long as in
             # forward it is a received input, regardless of whether it requires
             # grad. It is up to the previous stage to discard this gradient.
-            if isinstance(a, _RecvInfo):
-                grad_send_info.append(a.source)
-                return a.source
-            else:
+            if a.is_root_arg:
+                # Root args don't have a source stage to send gradients to
                 grad_send_info.append(None)
                 return None
+            else:
+                grad_send_info.append(a.source)
+                return a.source
 
         map_aggregate(args_recv_info, map_recv_to_send)
 
@@ -289,17 +302,30 @@ def map_recv_to_send(a):
     def _prepare_forward_infra(
         self,
         num_microbatches: int,
-        args: tuple[Any, ...],
+        args: tuple[Any, ...] | _StageForwardMeta | None,
         kwargs: dict[str, Any] | None = None,
-    ) -> tuple[Any, ...]:
+        has_backward: bool = False,
+    ) -> _StageForwardMeta | None:
+        raise NotImplementedError
+
+    @abstractmethod
+    def _prepare_backward_infra(
+        self,
+        num_microbatches: int,
+        loss_fn: Callable[..., torch.Tensor] | None = None,
+        target: torch.Tensor | None = None,
+        received_grad_meta: _StageBackwardMeta | None = None,
+    ) -> _StageBackwardMeta | None:
         raise NotImplementedError
 
-    def _prepare_backward_infra(self, num_microbatches: int):
+    def _setup_backward_recv_info(self, num_microbatches: int):
         # TODO: this is needed for backward_maybe_with_nosync
         self.chunks = num_microbatches
 
+        # IMPORTANT: _create_grad_recv_info reads self._stage_meta.output_grads
+        # to attach DTensor metadata to _RecvInfo objects. The clear below MUST
+        # happen after all _create_grad_recv_info calls complete.
         for mb_index in range(num_microbatches):
-            # `grad_recv_info` is a mirror of `act_send_info`
             self.grad_recv_info[mb_index] = self._create_grad_recv_info(
                 self.act_send_info
             )
@@ -311,9 +337,17 @@ def _create_grad_recv_info(
     ) -> tuple[_RecvInfo, ...]:
         raise NotImplementedError
 
+    def _resolve_peer_global_rank(self, stage_idx: int) -> int:
+        """Map a pipeline stage index to the corresponding global rank for P2P communication."""
+        peer_rank = self.stage_index_to_group_rank[stage_idx]
+        return dist.get_global_rank(
+            self.group or dist.distributed_c10d._get_default_group(),
+            peer_rank,
+        )
+
     def _get_recv_ops(
         self,
-        recv_infos: tuple[InputInfo, ...],
+        recv_infos: tuple[_RecvInfo, ...],
     ) -> list[dist.P2POp]:
         """
         Helper function shared by `get_fwd_recv_ops` and `get_bwd_recv_ops`.
@@ -321,15 +355,16 @@ def _get_recv_ops(
         """
         ops: list[dist.P2POp] = []
         for info in recv_infos:
-            if not isinstance(info, _RecvInfo):
+            if info.is_root_arg:
+                # Root args don't need recv operations
                 continue
-
-            peer_rank = self.stage_index_to_group_rank[info.source]
-            peer_global_rank = (
-                peer_rank
-                if self.group is None
-                else dist.get_global_rank(self.group, peer_rank)
-            )
+            # Skip entries with None buffer (None gradients)
+            if info.buffer is None:
+                assert info.tensor_meta is None  # noqa: S101
+                continue
+            # At this point, source and buffer are guaranteed non-None
+            assert info.source is not None  # noqa: S101
+            peer_global_rank = self._resolve_peer_global_rank(info.source)
             ops.append(
                 dist.P2POp(dist.irecv, info.buffer, peer_global_rank, self.group)
             )
@@ -338,46 +373,42 @@ def _get_recv_ops(
 
     """[Note: V-schedule special case]
 
-    V-Schedules have a special case where 2 stages with adjacent stage_id are on the same rank.
+    V-Schedules have a special case where 2 stages with adjacent stage_id
+    are on the same rank.
 
-    ex: 2 ranks, 4 stages forms a simple V:
-    rank0:  stage 0                   stage 3
-    rank1:          stage 1  stage 2
+    Example: 2 ranks, 4 stages forms a simple V::
 
-    stage 0,1 and 2,3 communicate activations using send/recv as usual, but stage 1,2 do not need to
-    use communication ops.  Instead, they should pass tensor data directly via function call.
+        rank0:  stage 0                   stage 3
+        rank1:          stage 1  stage 2
 
-    set_local_fwd_input and (get_local_bwd_output + set_local_bwd_input) facilitate this optimization, and
-    should be called at the appropriate time during the pipeline schedule (after forward or backward execution).
+    Stages 0/1 and 2/3 communicate via send/recv, but stages 1/2 pass
+    tensors directly via function call, avoiding communication ops.
     """
 
     def set_local_fwd_input(self, prev_stage_outputs: Any, mb_index: int) -> None:
+        """Pass outputs from a same-rank stage as forward inputs (V-schedule).
+
+        Detaches tensors and sets ``requires_grad`` so they serve as autograd
+        leaves. Handles DTensor activations transparently.
         """
-        Moves 'prev_stage_outputs' from another stage on the same rank into place as inputs for this stage. Avoids
-        copying tensor data or using send/recv op.  Detaches original tensor and sets requires_grad so the
-        tensor can serve as a leaf for autograd and gradients can be collected from it during backward.
-        """
-        recv_infos: tuple[InputInfo, ...] = self.args_recv_info[mb_index]
+        recv_infos: tuple[_RecvInfo, ...] = self.args_recv_info[mb_index]
 
         # See [Note: pipeline model output type]
         prev_stage_outputs = _normalize_model_output_as_tuple(prev_stage_outputs)
 
-        for info, tensor in zip(recv_infos, prev_stage_outputs):
+        for info, tensor in zip(recv_infos, prev_stage_outputs, strict=True):
             if not isinstance(tensor, torch.Tensor):
                 raise AssertionError(
                     f"expected tensor values as outputs from prev stage, got {type(tensor)}"
                 )
-            if not isinstance(info, _RecvInfo):
+            if info.is_root_arg:
                 raise AssertionError(
-                    "set_local_Fwd_input should only be called on non-first stage, which should always have RecvInfo"
+                    "set_local_fwd_input should only be called on non-first stage, which should always have non-root RecvInfo"
                 )
 
-            # We don't need to do a data copy here, since we can directly pass the activation tensor reference from
-            # one stage to the next.  However, we do need to mark the activation as a leaf tensor since it will serve
-            # as the input tensor for a fresh autograd graph, not part of the previous stage's autograd graph.
-            # TODO: confirm, do we use this activation as the root of the backward call for the previous stage? does
-            # detach have any affect on that?
-            info.buffer = tensor.detach().requires_grad_(True)
+            # Pass the activation tensor directly (same rank for local execution).
+            # Detach to create a new autograd leaf for the fresh autograd graph.
+            info.buffer = to_local_if_dtensor(tensor, detach=True)
 
     def get_local_bwd_output(self, mb_index):
         """
@@ -399,6 +430,7 @@ def set_local_bwd_input(
         """
         Moves 'grad input' tensors from the next stage to 'grad_output' on this stage, avoiding a copy or send/recv.
         Does not detach or set '_requires_grad'.
+        Handles DTensor gradients for V-schedule local passing.
         """
         if not isinstance(next_stage_bwd_outputs, tuple):
             raise AssertionError(f"Expected tuple, got {type(next_stage_bwd_outputs)}")
@@ -410,21 +442,27 @@ def set_local_bwd_input(
         if self.is_last:
             raise AssertionError("can't set bwd input if this stage is last")
         recv_infos = self.grad_recv_info[mb_index]
-        for info, tensor in zip(recv_infos, next_stage_bwd_outputs):
+        for info, tensor in zip(recv_infos, next_stage_bwd_outputs, strict=True):
+            if tensor is None:
+                continue
             if not isinstance(tensor, torch.Tensor):
                 raise AssertionError(
                     f"expected tensor values as outputs from prev stage, got {type(tensor)}"
                 )
-            if not isinstance(info, _RecvInfo):
-                raise AssertionError(f"Expected a recv info, got {type(info)}")
-            info.buffer = tensor
+            if info.is_root_arg:
+                raise AssertionError(
+                    "set_local_bwd_input should only be called with non-root RecvInfo"
+                )
+
+            # Extract local tensor for the buffer (handles DTensor or plain tensor)
+            info.buffer = to_local_if_dtensor(tensor)
 
     def get_fwd_recv_ops(self, fwd_chunk_id: int) -> list[dist.P2POp]:
         """
         Returns a list of ops that are needed to receive the input arguments
         for this stage.
         """
-        recv_infos: tuple[InputInfo, ...] = self.args_recv_info[fwd_chunk_id]
+        recv_infos: tuple[_RecvInfo, ...] = self.args_recv_info[fwd_chunk_id]
 
         return self._get_recv_ops(recv_infos)
 
@@ -442,6 +480,7 @@ def get_bwd_recv_ops(self, bwd_chunk_id: int) -> list[dist.P2POp]:
     def get_fwd_send_ops(self, fwd_chunk_id: int) -> list[dist.P2POp]:
         """
         Get the activation send ops for current stage's forward.
+        Handles DTensor outputs by extracting local tensors.
         """
         output_tuple, _ = self.fwd_cache[fwd_chunk_id]
 
@@ -452,25 +491,25 @@ def get_fwd_send_ops(self, fwd_chunk_id: int) -> list[dist.P2POp]:
             for dst in dst_stages:
                 if dst is None:
                     continue
+                # Extract local tensor if DTensor
+                send_tensor = to_local_if_dtensor(out, detach=True)
                 logger.debug(
                     "%s Sending tensor to Stage %s: %s",
                     self.log_prefix,
                     dst,
-                    out.size(),
+                    send_tensor.size(),
                 )
-                peer_rank = self.stage_index_to_group_rank[dst]
-                peer_global_rank = (
-                    peer_rank
-                    if self.group is None
-                    else dist.get_global_rank(self.group, peer_rank)
+                peer_global_rank = self._resolve_peer_global_rank(dst)
+                ops.append(
+                    dist.P2POp(dist.isend, send_tensor, peer_global_rank, self.group)
                 )
-                ops.append(dist.P2POp(dist.isend, out, peer_global_rank, self.group))
 
         return ops
 
     def get_bwd_send_ops(self, bwd_chunk_id: int) -> list[dist.P2POp]:
         """
         Get the gradient send ops for current stage's backward.
+        Handles DTensor gradients by extracting local tensors.
         """
         if not self.has_backward or self.is_first:
             return []
@@ -486,24 +525,24 @@ def get_bwd_send_ops(self, bwd_chunk_id: int) -> list[dist.P2POp]:
 
         ops: list[dist.P2POp] = []
         grads_input = self.bwd_cache.pop(bwd_chunk_id)
-        for grad, grad_recv_stage in zip(grads_input, self.grad_send_info):
+
+        for grad, grad_recv_stage in zip(grads_input, self.grad_send_info, strict=True):
             if isinstance(grad, torch.Tensor) and grad_recv_stage is not None:
+                # Extract local tensor if DTensor
+                send_tensor = to_local_if_dtensor(grad)
                 logger.debug(
                     "%s Sending gradient to Stage %s: %s",
                     self.log_prefix,
                     grad_recv_stage,
-                    grad.size(),
+                    send_tensor.size(),
                 )
-                peer_rank = self.stage_index_to_group_rank[grad_recv_stage]
-                peer_global_rank = (
-                    peer_rank
-                    if self.group is None
-                    else dist.get_global_rank(self.group, peer_rank)
+                peer_global_rank = self._resolve_peer_global_rank(grad_recv_stage)
+                ops.append(
+                    dist.P2POp(dist.isend, send_tensor, peer_global_rank, self.group)
                 )
-                ops.append(dist.P2POp(dist.isend, grad, peer_global_rank, self.group))
             else:
                 if grad is not None or grad_recv_stage is not None:
-                    raise RuntimeError(
+                    raise PipeliningMetadataError(
                         f"[{self.stage_index}] for chunk {bwd_chunk_id} has gradients {grad} "
                         f"and is expecting to send gradients to stage {grad_recv_stage}"
                     )
@@ -524,34 +563,90 @@ def clear_runtime_states(self) -> None:
         # don't want such accumulation.
         for recv_tuple in self.args_recv_info.values():  # iterate over all chunks
             for a in recv_tuple:  # iterate over all input args
-                if isinstance(a, _RecvInfo):
+                if not a.is_root_arg and a.buffer is not None:
                     # Set to None is the newer and recommended way to clear grads, compared to `zero_()`.
                     # See https://github.com/pytorch/pytorch/pull/92731
                     a.buffer.grad = None
 
     def _map_tensor_from_recv_info(
         self,
-        recv_infos: tuple[InputInfo, ...],
+        recv_infos: tuple[_RecvInfo, ...],
     ):
         """
         Map tensors from recv infos to a list.
         """
 
         def get_recv_tensor(info):
-            if isinstance(info, _RecvInfo):
-                return info.buffer
-            else:
-                raise AssertionError(f"Expected _RecvInfo but got {type(info)}")
+            if info.is_root_arg:
+                raise PipeliningMetadataError("Cannot get recv tensor from root arg")
+            return info.buffer
 
         return map_aggregate(cast(Argument, recv_infos), get_recv_tensor)
 
-    def _retrieve_recv_activations(self, fwd_chunk_id: int):
+    def _retrieve_recv_activations(
+        self,
+        fwd_chunk_id: int,
+    ):
         """
         Retrieve the activations received for the current stage during forward.
+        Reconstructs DTensors if the inputs were DTensors.
+        Also validates DTensor metadata against expected values.
         """
         recv_infos = self.args_recv_info[fwd_chunk_id]
-        activations = self._map_tensor_from_recv_info(recv_infos)
-        return activations
+
+        activations = []
+        for i, info in enumerate(recv_infos):
+            if not info.is_root_arg:
+                # Non-root args have valid buffer and tensor_meta
+                if info.buffer is None or info.tensor_meta is None:
+                    raise PipeliningMetadataError(
+                        f"Non-root arg '{info.input_name}' has None buffer or tensor_meta"
+                    )
+                # Effective requires_grad: metadata captures what the model
+                # produced, but the runtime context (has_backward, grad mode)
+                # determines whether we actually need gradients.
+                effective_requires_grad = (
+                    info.tensor_meta.requires_grad
+                    and self.has_backward
+                    and torch.is_grad_enabled()
+                )
+                if isinstance(info.tensor_meta, _DTensorMeta):
+                    # Buffer must not require grad so from_local stays out
+                    # of the autograd graph (no grad_placements needed).
+                    if info.buffer.requires_grad:
+                        raise PipeliningMetadataError(
+                            f"Stage {self.stage_index}: recv buffer "
+                            f"'{info.input_name}' unexpectedly requires grad "
+                            f"before DTensor reconstruction"
+                        )
+                    mesh = self._mesh_cache.get_mesh(info.tensor_meta.mesh_cache_key)
+                    activation = DTensor.from_local(
+                        info.buffer,
+                        device_mesh=mesh,
+                        placements=info.tensor_meta.placements,
+                        shape=info.tensor_meta.global_shape,
+                        stride=info.tensor_meta.global_stride,
+                        run_check=False,
+                    ).requires_grad_(effective_requires_grad)
+                else:
+                    activation = info.buffer.requires_grad_(effective_requires_grad)
+                # Activation must be a leaf so backward terminates here.
+                if effective_requires_grad and not activation.is_leaf:
+                    warnings.warn(
+                        f"Stage {self.stage_index}: activation "
+                        f"'{info.input_name}' is not a leaf "
+                        f"(grad_fn={activation.grad_fn}); using "
+                        f"retain_grad() as fallback",
+                        stacklevel=2,
+                    )
+                    activation.retain_grad()
+                activations.append(activation)
+            else:
+                raise PipeliningMetadataError(
+                    f"_retrieve_recv_activations expected non-root _RecvInfo but got root arg at index {i}"
+                )
+
+        return tuple(activations)
 
     def _retrieve_recv_grads(
         self,
@@ -559,10 +654,50 @@ def _retrieve_recv_grads(
     ):
         """
         Retrieve the gradients received for the current stage during backward.
+
+        Handles None gradients gracefully (for inputs that don't require grad).
         """
         recv_infos = self.grad_recv_info[bwd_chunk_id]
-        grads = self._map_tensor_from_recv_info(recv_infos)
-        return grads
+
+        grads: list[torch.Tensor | None] = []
+        for i, info in enumerate(recv_infos):
+            if not isinstance(info, _RecvInfo):
+                raise PipeliningMetadataError(
+                    f"Expected _RecvInfo but got {type(info)}"
+                )
+            if not info.is_root_arg:
+                # Gradients can be None for non-differentiable outputs
+                if info.buffer is None:
+                    if info.tensor_meta is not None:
+                        raise PipeliningMetadataError(
+                            f"Grad recv '{info.input_name}': buffer is None but tensor_meta is not None"
+                        )
+                    grads.append(None)
+                    continue
+                if info.tensor_meta is None:
+                    raise PipeliningMetadataError(
+                        f"Grad recv '{info.input_name}': buffer is not None but tensor_meta is None"
+                    )
+                if isinstance(info.tensor_meta, _DTensorMeta):
+                    # Reconstruct DTensor gradient from local tensor + metadata
+                    mesh = self._mesh_cache.get_mesh(info.tensor_meta.mesh_cache_key)
+                    grad = DTensor.from_local(
+                        info.buffer,
+                        device_mesh=mesh,
+                        placements=info.tensor_meta.placements,
+                        shape=info.tensor_meta.global_shape,
+                        stride=info.tensor_meta.global_stride,
+                        run_check=False,
+                    )
+                else:
+                    grad = info.buffer
+                grads.append(grad)
+            else:
+                raise PipeliningMetadataError(
+                    f"grad_recv_info should not contain root args, but found one at index {i}"
+                )
+
+        return tuple(grads)
 
     def forward_maybe_with_nosync(self, *args, **kwargs):
         # If submod is wrapped with DDP, we use the `no_sync` context manager to
@@ -692,7 +827,12 @@ def forward_one_chunk(
 
         composite_kwargs = kwargs or {}
 
-        self._validate_fwd_input(args, kwargs)
+        if self._runtime_validate:
+            self._validate_stage_tensors(
+                f"Stage {self.stage_index} forward inputs",
+                self._stage_meta.inputs,
+                composite_args,
+            )
 
         # Compute forward
         try:
@@ -728,7 +868,14 @@ def forward_one_chunk(
             fwd_chunk_id,
             map_debug_info(output),
         )
-        self._validate_fwd_outputs(output_tuple)
+        # Validate outputs before P2P send; skipped for last stage (outputs
+        # go to loss/user, not via send/recv).
+        if self._runtime_validate and not self.is_last:
+            self._validate_stage_tensors(
+                f"Stage {self.stage_index} forward outputs",
+                self._stage_meta.outputs,
+                output_tuple,
+            )
 
         # We return the original user-provided output, not normalized to tuple.
         # See [Note: pipeline model output type]
@@ -777,6 +924,13 @@ def backward_one_chunk(
         else:
             # Otherwise, receive gradients from next stage
             grads_output = self._retrieve_recv_grads(bwd_chunk_id)
+            if self._runtime_validate:
+                # Validate backward input (output gradients) for DTensor metadata
+                self._validate_stage_tensors(
+                    f"Stage {self.stage_index} backward input (output_grads)",
+                    self._stage_meta.output_grads,
+                    grads_output,
+                )
             # If an input to the pipeline requires gradient,
             # `torch.autograd.backward` will accumulate the gradient into the
             # `.grad` field of such input
@@ -829,8 +983,18 @@ def backward_one_chunk(
                 )
                 # Save a placeholder for the dw_runner
                 self.dw_runner[bwd_chunk_id] = lambda: None
-
-        self.bwd_cache[bwd_chunk_id] = grads_input
+        # Note: grads_input may contain gradients for both args and kwargs (from fwd_cache),
+        # Kwargs are local to each stage and don't need gradient transmission.
+        # Validate backward output (input gradients) for DTensor metadata
+        assert self._stage_meta.inputs is not None  # noqa: S101
+        num_fwd_args = len(self._stage_meta.inputs)
+        if self._runtime_validate and not self.is_first:
+            self._validate_stage_tensors(
+                f"Stage {self.stage_index} backward output (input_grads)",
+                self._stage_meta.input_grads,
+                grads_input[:num_fwd_args],
+            )
+        self.bwd_cache[bwd_chunk_id] = grads_input[:num_fwd_args]
 
         if self.is_last and not self.is_first:
             # Autograd dependencies:
@@ -888,48 +1052,6 @@ def backward_weight_one_chunk(self, bwd_chunk_id: int, last_backward=False):
                     "full", bwd_kwargs, last_backward=last_backward
                 )
 
-    def _validate_fwd_input(self, args, kwargs):
-        """Raises a RuntimeError if shapes of input args/kwargs do not match the shapes configured for this stage."""
-
-        if self.is_first:
-            # TODO why is there a separate recv_info for each pipeline chunk?
-            # kwen2501: to avoid passing a `fwd_chunk_id` to this function, we
-            # check all chunks against args_recv_info[0]
-            expected_args = self.args_recv_info[0]
-        else:
-            # We don't check inputs for non-0 stages assuming they don't accept
-            # user inputs in canonical pipeline scenarios
-            return
-
-        if len(kwargs):
-            # TODO- need a mapping of kwarg to position in self.args_recv_info
-            # Without it, we are not 100% sure how to match the args and
-            # expected_args.
-            return
-
-        # TODO- need a mapping of kwarg to position in self.args_recv_info
-        # maybe it's impossible to tell whether the len mismatches because
-        # (a) the user passed an extra arg or missed an arg
-        # (b) the user did not pass a kwarg, which has a default value baked into expected_args
-        expected_tensors_meta = [
-            e.meta if isinstance(e, _RootArgPlaceholder) else e.buffer
-            for e in expected_args
-        ]
-        validate_tensors_metadata(
-            f"Stage {self.stage_index} forward inputs", expected_tensors_meta, args
-        )
-
-    def _validate_fwd_outputs(self, outputs: tuple[torch.Tensor, ...]):
-        """Raises a RuntimeError if this stage produces an output of unexpected shape/dtype.
-        Most likely, this could be cause either by incorrect user specification of output shapes, or because
-        shape inference was done on the original model but then at runtime the model is wrapped with something like
-        mixed precision which changes output dtype.
-        """
-        expected_tensors_meta = self.get_outputs_meta()
-        validate_tensors_metadata(
-            f"Stage {self.stage_index} forward outputs", expected_tensors_meta, outputs
-        )
-
     def _get_init_p2p_neighbors_ops(self) -> list[dist.P2POp]:
         """
         Get the operations to initialize the p2p communicators between previous and next stages.
@@ -1007,8 +1129,8 @@ def perform_reduce_grad(self, grad_scale_factor: int):
                 distributed_state = fully_shard.state(fsdp_module)  # type: ignore[attr-defined]
 
             for state in distributed_state._state_ctx.all_states:
-                if state._fsdp_param_group:
-                    state._fsdp_param_group.post_backward()
+                for fsdp_param_group in state._fsdp_param_groups:
+                    fsdp_param_group.post_backward()
 
             # it would be much better if pipelining backward invoked .backward so autograd hooks
             # worked and modules like DDP/FSDP behaved as expected.  Working around this for the time being,
@@ -1055,7 +1177,7 @@ def __init__(
             node for node in pipe_info.graph.nodes if node.op == "call_module"
         ]
         if len(submod_nodes) != self.num_stages:
-            raise AssertionError(
+            raise PipeliningMetadataError(
                 f"Number of submodules in pipe graph {len(submod_nodes)} does not match number of stages {self.num_stages}"
             )
 
@@ -1093,21 +1215,85 @@ def _move_submod_to_device(self):
     def _prepare_forward_infra(
         self,
         num_microbatches: int,
-        args: tuple[Any, ...],
+        args: tuple[Any, ...] | _StageForwardMeta | None,
         kwargs: dict[str, Any] | None = None,
-    ) -> tuple[Any, ...]:
+        has_backward: bool = False,
+    ) -> _StageForwardMeta | None:
         """
-        Create send/recv infrastructures for activations (during forward)
+        Prepare forward infrastructure for traced pipeline.
+
+        Metadata is created directly from graph placeholders with correct
+        ``requires_grad`` — received activations get ``requires_grad=True``
+        when ``has_backward`` is set, fixing the fact that ``torch.export``
+        traces under ``no_grad()``.
+
+        ``_stage_meta.inputs`` is derived from recv infos and aligned with
+        ``forward_one_chunk``'s ``composite_args``: positional root inputs
+        on the first stage, received activations only on subsequent stages.
         """
-        # TODO(whc)
-        # this method should be deleted once lazy buffer allocation is implemented
-        # for now, it ignores args/kwargs because it should not need to do shape inference
+        # Step 1: Create recv info for each microbatch.
+        # _create_act_recv_info is self-contained: it creates _TensorMeta
+        # directly from graph placeholder values with correct requires_grad.
         for chunk in range(num_microbatches):
             self.args_recv_info[chunk] = self._create_act_recv_info()
 
-        # Send info during forward for each activation
+        # Step 2: Derive _stage_meta.inputs from recv infos.
+        # forward_one_chunk builds composite_args as:
+        #   - First stage:     args (positional root inputs, excludes kwargs)
+        #   - Non-first stages: received activations only (no root kwargs)
+        # _stage_meta.inputs must match composite_args for validation.
+        recv_infos = self.args_recv_info[0]
+        if self.is_first:
+            # All placeholders are root args.  Only the first len(args)
+            # correspond to positional inputs (composite_args); the rest
+            # are kwargs passed separately via composite_kwargs.
+            # First stage always receives real tensor args, never _StageForwardMeta.
+            if not isinstance(args, tuple):
+                raise AssertionError("First stage requires real tensor args")
+            n_positional = len(args)
+            self._stage_meta.inputs = tuple(
+                info.tensor_meta  # type: ignore[misc]
+                for info in recv_infos[:n_positional]
+            )
+        else:
+            self._stage_meta.inputs = tuple(
+                info.tensor_meta  # type: ignore[misc]
+                for info in recv_infos
+                if not info.is_root_arg
+            )
+
+        # Step 3: Create send info and output metadata.
         self.act_send_info = self._create_act_send_info()
-        return tuple()
+
+        return None
+
+    def _prepare_backward_infra(
+        self,
+        num_microbatches: int,
+        loss_fn: Callable[..., torch.Tensor] | None = None,
+        target: torch.Tensor | None = None,
+        received_grad_meta: _StageBackwardMeta | None = None,
+    ) -> _StageBackwardMeta | None:
+        """
+        Prepare backward infrastructure for traced pipeline.
+        Derives input_grads metadata from inputs (plain tensors only).
+
+        Note: DTensors are NOT supported in the traced frontend.
+        """
+        # Derive input_grads from inputs (for plain tensors, grad shape == input shape)
+        if self._stage_meta.inputs is None:
+            raise PipeliningMetadataError(
+                f"Stage {self.stage_index}: inputs metadata required for backward inference."
+            )
+
+        self._stage_meta.input_grads = _derive_grad_metas(self._stage_meta.inputs)
+
+        # Setup backward recv info (calls _create_grad_recv_info which sets output_grads).
+        # Note: grad_send_info is created lazily in get_bwd_send_ops() since
+        # it mirrors args_recv_info (already populated during forward).
+        self._setup_backward_recv_info(num_microbatches)
+
+        return None
 
     def get_stage_index_of_submod(
         self,
@@ -1117,7 +1303,7 @@ def get_stage_index_of_submod(
         Given a submodule name, return the stage index of the submodule.
         """
         if submod_name not in self.submod_to_stage_index:
-            raise AssertionError(f"Stage id of {submod_name} not found")
+            raise PipeliningMetadataError(f"Stage id of {submod_name} not found")
 
         return self.submod_to_stage_index[submod_name]
 
@@ -1126,38 +1312,65 @@ def _create_act_recv_info(
     ):
         """
         Create a tuple of `_RecvInfo` for inputs to the stage.
+
+        Self-contained: creates ``_TensorMeta`` directly from graph
+        placeholder values with correct ``requires_grad``.
+        ``torch.export`` traces under ``no_grad()`` so traced metadata
+        always has ``requires_grad=False``; for received activations we
+        set ``requires_grad=True`` when ``has_backward`` is set.
+
+        Note: DTensors are NOT supported in the traced frontend.
         """
 
         def create_recv_tensor(placeholder, arg_node):
-            """
-            Create a receive buffer for a placeholder.
-            """
             example_value = placeholder.meta["val"]
+
+            # Reject DTensors in traced frontend
+            if isinstance(example_value, DTensor):
+                raise PipeliningMetadataError(
+                    f"{self.log_prefix} DTensor detected in traced pipeline input "
+                    f"'{placeholder.name}'. DTensor metadata propagation is NOT "
+                    f"supported for the traced frontend (_PipelineStage). "
+                    f"Use the manual PipelineStage frontend for full DTensor support."
+                )
+
             if arg_node.op == "placeholder":
-                # This is a root level placeholder, thus an input argument to the entire model.
-                # We are likely at stage 0, hence no need to create a receive buffer.
-                return _RootArgPlaceholder(example_value)
+                # Root-level placeholder: an input argument to the entire
+                # model.  Keep original metadata from the trace.
+                return _RecvInfo(
+                    input_name=f"root_input_{placeholder.name}",
+                    source=None,
+                    buffer=None,
+                    tensor_meta=_TensorMeta.from_tensor(example_value),
+                    is_root_arg=True,
+                )
 
-            # Figure out the source stage of this input
+            # Received activation from a previous stage.
             while arg_node.target is operator.getitem:
-                # If the input is a getitem, we need to go deeper
                 arg_node = arg_node.args[0]
 
             if arg_node.op != "call_module":
-                raise AssertionError(f"Expecting call_module, got {arg_node.op}")
+                raise PipeliningMetadataError(
+                    f"Expecting call_module, got {arg_node.op}"
+                )
             src_stage = self.get_stage_index_of_submod(arg_node.name)
 
-            # Create a receive buffer for this placeholder
+            # Create metadata directly with correct requires_grad.
+            tensor_meta = _TensorMeta(
+                shape=example_value.shape,
+                stride=example_value.stride(),
+                dtype=example_value.dtype,
+                requires_grad=self.has_backward,
+            )
+
             logger.debug(
                 "%s Creating recv buffer for input '%s' : %s, %s",
                 self.log_prefix,
                 placeholder.name,
-                example_value.shape,
-                example_value.dtype,
+                tensor_meta.shape,
+                tensor_meta.dtype,
             )
-            buffer = _make_tensor_from_meta(example_value, self.device)
-            # In case there is backward pass, set requires_grad for receive buffers
-            # before first forward
+            buffer = _make_tensor_from_meta(tensor_meta, self.device)
             if self.has_backward:
                 buffer.requires_grad_(True)
 
@@ -1165,10 +1378,10 @@ def create_recv_tensor(placeholder, arg_node):
                 arg_node.name,
                 src_stage,
                 buffer,
+                tensor_meta,
             )
 
-        args_recv_info: list[InputInfo] = []
-        # Filter out placeholder nodes from `self.submod` (a GraphModule)
+        args_recv_info: list[_RecvInfo] = []
         placeholders = filter(  # type: ignore[var-annotated]
             lambda node: node.op == "placeholder",  # type: ignore[arg-type]
             self.submod.graph.nodes,  # type: ignore[arg-type,union-attr]
@@ -1176,15 +1389,12 @@ def create_recv_tensor(placeholder, arg_node):
         # `placeholders` are nodes internal to submod.
         # `self.node.args` are dependency nodes in the outer graph.
         # The two are 1:1.
-        for placeholder, arg_node in zip(placeholders, self.node.args):
-            # Create a receive buffer for this placeholder
-            recv_info = create_recv_tensor(placeholder, arg_node)
-            args_recv_info.append(recv_info)
+        for placeholder, arg_node in zip(placeholders, self.node.args, strict=True):
+            args_recv_info.append(create_recv_tensor(placeholder, arg_node))
 
         logger.debug(
             "%s Activation recv / args info: %s", self.log_prefix, args_recv_info
         )
-        # `args` is a Tuple, hence we will return a Tuple[InputInfo]
         return tuple(args_recv_info)
 
     def find_dst_rank(
@@ -1208,14 +1418,12 @@ def find_dst_rank(
 
     def _create_act_send_info(self):
         """
-        Create a dict of send info for activations.
-        The dict is of the form:
-        {
-            output_index: [dst_rank_0, dst_rank_1, ...],
-            ...
-        }
-        where the list of `dst_rank`s covers the case where an output value may
-        be consumed by multiple stages.
+        Create a dict of send info for activations and output metadata.
+
+        Output metadata is created directly with correct ``requires_grad``
+        (``torch.export`` traces under ``no_grad()``, so traced values
+        always have ``requires_grad=False``; at runtime, stage outputs
+        carry ``requires_grad=True`` when training).
         """
         # Output index: List of receiver ranks
         act_send_info: dict[int, list] = {}
@@ -1223,16 +1431,13 @@ def _create_act_send_info(self):
 
         for user in self.node.users:
             if user.target is operator.getitem:
-                # Recursively find the real destination
                 gi_dsts = act_send_info.setdefault(out_idx, [])
                 for gi_user in user.users:
                     dst_rank = self.find_dst_rank(gi_user)
                     if dst_rank is not None:
                         gi_dsts.append(dst_rank)
-                # Next `getitem` will point to the next output index
                 out_idx += 1
             else:
-                # In case of single output value, `out_idx` will not increase
                 dsts = act_send_info.setdefault(out_idx, [])
                 dst_rank = self.find_dst_rank(user)
                 if dst_rank is not None:
@@ -1242,7 +1447,25 @@ def _create_act_send_info(self):
         output_vals: tuple[torch.Tensor] = tuple(
             v.meta["val"] for v in flatten_args(output_node.args)
         )
-        self._configure_outputs_meta(output_vals)
+        # Reject DTensors and create output metadata directly with
+        # correct requires_grad.
+        output_metas: list[_TensorMeta] = []
+        for i, val in enumerate(output_vals):
+            if isinstance(val, DTensor):
+                raise PipeliningMetadataError(
+                    f"{self.log_prefix} DTensor detected in traced pipeline output index {i}. "
+                    f"DTensor metadata propagation is NOT supported for the traced frontend "
+                    f"(_PipelineStage). Use the manual PipelineStage frontend for full DTensor support."
+                )
+            output_metas.append(
+                _TensorMeta(
+                    shape=val.shape,
+                    stride=val.stride(),
+                    dtype=val.dtype,
+                    requires_grad=self.has_backward,
+                )
+            )
+        self._stage_meta.outputs = tuple(output_metas)
 
         logger.debug("%s Send info: %s", self.log_prefix, act_send_info)
         return act_send_info
@@ -1250,50 +1473,79 @@ def _create_act_send_info(self):
     def _get_output_node(self):
         output_nodes = [node for node in self.submod.graph.nodes if node.op == "output"]  # type: ignore[union-attr]
         if len(output_nodes) != 1:
-            raise AssertionError(f"Expected 1 output node, got {len(output_nodes)}")
+            raise PipeliningMetadataError(
+                f"Expected 1 output node, got {len(output_nodes)}"
+            )
         output_node = output_nodes[0]
         return output_node
 
-    def _create_grad_recv_info(
-        self,
-        act_send_info: dict,
-    ) -> tuple[_RecvInfo, ...]:
+    def _create_grad_recv_info(self, act_send_info: dict) -> tuple[_RecvInfo, ...]:
         """
         Create a tuple of `_RecvInfo` for gradients.
+        Reuses output metadata from _stage_meta.outputs (populated by _create_act_send_info).
         """
-        # Dict[output_index, _RecvInfo]
-        grad_recv_info: dict[int, _RecvInfo] = {}
-        output_node = self._get_output_node()
+        if self._stage_meta.outputs is None:
+            raise PipeliningMetadataError(
+                f"Stage {self.stage_index}: outputs metadata required for grad recv info. "
+                f"Ensure _create_act_send_info is called first."
+            )
 
-        # The output node may take multiple args, meaning the submod having multiple output values.
-        output_vals = flatten_args(output_node.args)
+        outputs_meta = self._stage_meta.outputs
+        output_grads_metas: list[TensorMeta | None] = []
+        grad_recv_infos: list[_RecvInfo] = []
 
-        for out_idx, dst_list in act_send_info.items():
-            if not dst_list:
-                # No actual receiver for activation so no grad coming back
-                continue
+        for out_idx, out_meta in enumerate(outputs_meta):
+            dst_list = act_send_info.get(out_idx, [])
 
-            output = output_vals[out_idx]
-            example_value = output.meta["val"]
-            logger.debug(
-                f"{self.log_prefix} Creating grad recv buffer for output {output.name} "  # noqa: G004
-                f": {example_value.shape}, {example_value.dtype}"
-            )
+            # Determine the source stage for gradients
+            grad_src = dst_list[0] if dst_list else self.stage_index + 1
 
-            # TODO: otherwise needs grad accumulation
-            if len(dst_list) != 1:
-                raise AssertionError("Backward of skip connections not supported yet")
-            grad_src = dst_list[0]
-            grad_recv_info[out_idx] = _RecvInfo(
-                f"{grad_src}",  # noqa: G004
-                grad_src,
-                _make_tensor_from_meta(example_value, self.device),
-            )
+            # Check if this output needs gradients
+            if not dst_list or not out_meta.requires_grad:
+                output_grads_metas.append(None)
+                grad_recv_infos.append(
+                    _RecvInfo(
+                        input_name=f"recv_grad_for_{self.stage_index}_none_{out_idx}",
+                        source=grad_src,
+                        buffer=None,
+                        tensor_meta=None,
+                    )
+                )
+            else:
+                # Derive grad metadata from output metadata (same shape, requires_grad=False)
+                grad_meta = _TensorMeta(
+                    shape=out_meta.shape,
+                    stride=out_meta.stride,
+                    dtype=out_meta.dtype,
+                    requires_grad=False,
+                )
+                output_grads_metas.append(grad_meta)
 
-        # Convert to tuple for convenience in get_ops and retrieve tensor
-        grad_recv_info_tuple = tuple(grad_recv_info.values())
-        logger.debug("%s Grad recv info: %s", self.log_prefix, grad_recv_info_tuple)
-        return grad_recv_info_tuple
+                if len(dst_list) != 1:
+                    raise PipeliningMetadataError(
+                        "Backward of skip connections not supported yet"
+                    )
+
+                logger.debug(
+                    "%s Creating grad recv buffer for output %s : %s, %s",
+                    self.log_prefix,
+                    out_idx,
+                    grad_meta.shape,
+                    grad_meta.dtype,
+                )
+
+                grad_recv_infos.append(
+                    _RecvInfo(
+                        input_name=f"recv_grad_for_{self.stage_index}_from_{grad_src}",
+                        source=grad_src,
+                        buffer=_make_tensor_from_meta(grad_meta, self.device),
+                        tensor_meta=grad_meta,
+                    )
+                )
+
+        self._stage_meta.output_grads = tuple(output_grads_metas)
+        logger.debug("%s Grad recv info: %s", self.log_prefix, grad_recv_infos)
+        return tuple(grad_recv_infos)
 
 
 # A helper function to create a pipeline stage based on traced pipeline information
@@ -1328,26 +1580,33 @@ def build_stage(
 
 
 class PipelineStage(_PipelineStageBase):
-    """
-    A class representing a pipeline stage in a pipeline parallelism setup.
+    """A pipeline stage for pipeline parallelism with sequential model partitioning.
+
+    Supports both **static** and **dynamic** metadata inference:
 
-    PipelineStage assumes sequential partitioning of the model, i.e. the model is split into chunks where outputs from
-    one chunk feed into inputs of the next chunk, with no skip connections.
+    Static mode:
+        All of ``input_args``, ``output_args`` (and ``input_grads``/``output_grads``
+        when DTensors are present) are provided at construction time.
 
-    PipelineStage performs runtime shape/dtype inference automatically by propagating the outputs from stage0 to
-    stage1 and so forth, in linear order.  To bypass shape inference, pass the `input_args` and `output_args` to each
-    PipelineStage instance.
+    Dynamic mode:
+        Metadata is inferred from the first microbatch at runtime; any
+        statically provided args are used for validation only.
 
     Args:
-        submodule (nn.Module): The PyTorch module wrapped by this stage.
-        stage_index (int): The ID of this stage.
-        num_stages (int): The total number of stages.
-        device (torch.device): The device where this stage is located.
-        input_args (Union[torch.Tensor, Tuple[torch.tensor]], optional): The input arguments for the submodule.
-        output_args (Union[torch.Tensor, Tuple[torch.tensor]], optional): The output arguments for the submodule.
-        group (dist.ProcessGroup, optional): The process group for distributed training. If None, default group.
-        dw_builder (Optional[Callable[[], Callable[..., None]]): If provided, dw_builder will build a new dw_runner function
-            that will the W action (input weights) for F, I, W (Fwd, Input, Weight) zero bubble schedules.
+        submodule: The ``nn.Module`` wrapped by this stage.
+        stage_index: Zero-based stage ID.
+        num_stages: Total number of stages in the pipeline.
+        device: Device this stage runs on.
+        input_args: Example input tensors (single tensor or tuple). Optional.
+        output_args: Example output tensors. Optional.
+        output_grads: Example output gradients (received from next stage). Optional.
+        input_grads: Example input gradients (sent to previous stage). Optional.
+        group: Process group for P2P communication. Defaults to the
+            world process group.
+        dw_builder: Builder for deferred weight-update runners used by
+            zero-bubble (F/I/W) schedules.
+        get_mesh: `GetMeshCallback` used during
+            dynamic DTensor inference. Ignored in fully static DTensor mode.
     """
 
     def __init__(
@@ -1358,247 +1617,633 @@ def __init__(
         device: torch.device,
         input_args: torch.Tensor | tuple[torch.Tensor, ...] | None = None,
         output_args: torch.Tensor | tuple[torch.Tensor, ...] | None = None,
+        output_grads: torch.Tensor | tuple[torch.Tensor | None, ...] | None = None,
+        input_grads: torch.Tensor | tuple[torch.Tensor | None, ...] | None = None,
         group: dist.ProcessGroup | None = None,
         dw_builder: Callable[[], Callable[..., None]] | None = None,
+        get_mesh: GetMeshCallback | None = None,
     ):
         super().__init__(submodule, stage_index, num_stages, device, group, dw_builder)
-        self.inputs: list[torch.Tensor] | None = None
-        self.inputs_meta: tuple[torch.Tensor, ...] | None = None
-        # Note: inputs and submod should ideally be on meta device. We decided not to assert this (yet) because it
-        # might be breaking for existing users.
-        if input_args is None:
-            if output_args is not None:
-                raise AssertionError(
-                    "If specifying output_args, input_args must also be specified. "
-                    "Otherwise, shape inference will be performed at runtime"
-                )
-        else:
-            self.inputs_meta = (
-                (input_args,) if isinstance(input_args, torch.Tensor) else input_args
-            )
-            if output_args is None:
-                logger.warning(
-                    "Deprecation warning: passing input_args and performing init-time shape inference is deprecated. "
-                    "PipelineStage now supports runtime shape inference using the real inputs provided to schedule step(). "
-                    "Either delete `input_args` arg to `PipelineStage` to opt-into runtime shape inference, "
-                    "or additionally pass `output_args` to `PipelineStage` to fully override shape inference. "
+
+        self._mesh_cache = _MeshCache(get_mesh_cb=get_mesh)
+        self._inference_mode: InferenceMode | None = None
+        self._fwd_outputs_for_bwd_meta: tuple[torch.Tensor, ...] | None = None
+        self._fwd_inputs_for_bwd_meta: tuple[torch.Tensor, ...] | None = None
+        self._fwd_kwargs_tensors_for_bwd_meta: tuple[torch.Tensor, ...] | None = None
+
+        # Validate and normalize args to tuples
+        inputs = validate_and_normalize_to_tuple(input_args)
+        outputs = validate_and_normalize_to_tuple(output_args)
+        in_grads = validate_and_normalize_to_tuple(input_grads, allow_none=True)
+        out_grads = validate_and_normalize_to_tuple(output_grads, allow_none=True)
+
+        self._user_meta = _StageMeta(
+            inputs=extract_tensor_metas(inputs),
+            outputs=extract_tensor_metas(outputs),
+            input_grads=extract_tensor_metas(in_grads, allow_none=True),
+            output_grads=extract_tensor_metas(out_grads, allow_none=True),
+        )
+
+        # Cache meshes from user-provided DTensors
+        for args in (inputs, outputs, in_grads, out_grads):
+            if args is not None:
+                self._mesh_cache.update_from_tensors(args)
+
+        # Validate DTensor↔grad correspondence independently for inputs and outputs
+        if self._user_meta.has_dtensors():
+            if inputs and in_grads:
+                validate_static_arg_grad_correspondence(
+                    self.stage_index, inputs, in_grads, is_input=True
                 )
-                try:
-                    with torch.no_grad():
-                        output_args = submodule(*self.inputs_meta)
-                    output_args = tree_map_only(
-                        torch.Tensor, lambda x: x.to("meta"), output_args
-                    )
-                except Exception as e:
-                    raise RuntimeError(
-                        "Failed to perform pipeline shape inference- are your inputs on the same device as your module?"
-                    ) from e
-            if output_args is None:
-                raise AssertionError(
-                    "If passing input_args, also pass output_args to override shape inference"
+            if outputs and out_grads:
+                validate_static_arg_grad_correspondence(
+                    self.stage_index, outputs, out_grads, is_input=False
                 )
-            self._configure_outputs_meta(
-                (output_args,) if isinstance(output_args, torch.Tensor) else output_args
+
+    def _recv_meta(self, src_stage: int) -> Any:
+        """Receive metadata object from a stage on a different rank via P2P."""
+        objects: list[Any] = [None]
+        dist.recv_object_list(
+            objects,
+            src=self._resolve_peer_global_rank(src_stage),
+            group=self.group,
+            device=self.device,
+            use_batch=True,
+        )
+        if len(objects) != 1:
+            raise PipeliningMetadataError(
+                f"Expected exactly one object to be received but got: {len(objects)}"
             )
+        return objects[0]
+
+    def _send_meta(self, meta: Any, dst_stage: int) -> None:
+        """Send metadata object to a stage on a different rank via P2P."""
+        dist.send_object_list(
+            [meta],
+            dst=self._resolve_peer_global_rank(dst_stage),
+            group=self.group,
+            device=self.device,
+            use_batch=True,
+        )
+
+    def _is_same_rank(self, other_stage: int) -> bool:
+        """Check if another stage is on the same rank as this stage."""
+        return self.stage_index_to_group_rank[other_stage] == self.group_rank
 
-        # these are the buffers used in backwards send/recv, they are allocated later
-        self.outputs_grad: list[torch.Tensor] = []
+    def _warmup_forward_vote(
+        self, has_backward: bool, received_acc: torch.Tensor | None = None
+    ) -> torch.Tensor:
+        """Forward phase of the warm-up vote protocol (stage 0 → N−1).
+
+        Each stage computes a vote (1 = STATIC, 0 = DYNAMIC) based on
+        ``InferenceMode.needs_dynamic``, multiplies it with the accumulated
+        product from the previous stage, and forwards the result to the next
+        stage.  The final product at stage N−1 is 1 iff *every* stage voted
+        STATIC.
+
+        Args:
+            has_backward: Whether the schedule includes a backward pass.
+            received_acc: Accumulated product tensor from the previous
+                same-rank stage (V-schedule), or ``None`` for the first
+                stage / cross-rank.
+
+        Returns:
+            The accumulated product tensor after this stage's vote.
+        """
+        my_vote = 0 if InferenceMode.needs_dynamic(self._user_meta, has_backward) else 1
+
+        my_vote_t = torch.tensor([my_vote], dtype=torch.int32, device=self.device)
+
+        if self.is_first:
+            acc = my_vote_t
+        elif self._is_same_rank(self.stage_index - 1):
+            assert received_acc is not None  # noqa: S101
+            acc = received_acc * my_vote_t
+        else:
+            peer_global = self._resolve_peer_global_rank(self.stage_index - 1)
+            acc = torch.zeros(1, dtype=torch.int32, device=self.device)
+            dist.recv(acc, src=peer_global, group=self.group)
+            acc = acc * my_vote_t
 
-        dbg_str = (
-            f"Finished pipeline stage init, {self.stage_index=}, {self.is_first=}, "  # noqa: G004
-            f"{self.is_last=}, {self.num_stages=}, "
+        if not self.is_last and not self._is_same_rank(self.stage_index + 1):
+            peer_global = self._resolve_peer_global_rank(self.stage_index + 1)
+            dist.send(acc, dst=peer_global, group=self.group)
+
+        return acc
+
+    def _warmup_backward_result(
+        self, received_result: torch.Tensor | None = None
+    ) -> torch.Tensor:
+        """Backward phase of the warm-up vote protocol (stage N−1 → 0).
+
+        Propagates the final accumulated product (computed in the forward
+        phase) back through the pipeline so every stage learns the global
+        inference mode.
+
+        Args:
+            received_result: Result tensor from the next same-rank stage
+                (V-schedule), or ``None`` for the last stage / cross-rank.
+
+        Returns:
+            The global vote result tensor for this stage.
+        """
+        if self.is_last or self._is_same_rank(self.stage_index + 1):
+            assert received_result is not None  # noqa: S101
+            result = received_result
+        else:
+            peer_global = self._resolve_peer_global_rank(self.stage_index + 1)
+            result = torch.zeros(1, dtype=torch.int32, device=self.device)
+            dist.recv(result, src=peer_global, group=self.group)
+
+        if not self.is_first and not self._is_same_rank(self.stage_index - 1):
+            peer_global = self._resolve_peer_global_rank(self.stage_index - 1)
+            dist.send(result, dst=peer_global, group=self.group)
+
+        return result
+
+    def _compute_outputs(
+        self,
+        *args: torch.Tensor,
+        module: torch.nn.Module,
+        **kwargs: Any,
+    ) -> torch.Tensor | tuple[torch.Tensor, ...] | list[torch.Tensor] | None:
+        """Compute outputs of the submodule."""
+        return module(*args, **kwargs)
+
+    def _compute_input_grads(
+        self,
+        outputs: list[torch.Tensor],
+        all_fwd_inputs: list[torch.Tensor],
+        grad_outputs: list[torch.Tensor | None] | None = None,
+    ) -> tuple[torch.Tensor | None, ...]:
+        """Compute input gradients via :func:`_autograd_grad_for_inputs`."""
+        return _autograd_grad_for_inputs(
+            outputs,
+            all_fwd_inputs,
+            grad_outputs,
         )
-        if self.inputs_meta is not None:
-            dbg_str += (
-                f"inputs: {[inp.shape for inp in self.inputs_meta]}, "
-                f"output: {[output.shape for output in self.get_outputs_meta()]}"
-            )
+
+    def _to_tensor(self, arg: torch.Tensor | TensorMeta) -> torch.Tensor:
+        """Convert a tensor or metadata to a real tensor on ``self.device``.
+
+        Real tensors are detached and re-set requires_grad to create a fresh
+        autograd leaf, isolating metadata inference from the user's graph.
+        TensorMeta is materialized as an empty tensor (or DTensor via mesh cache).
+        """
+        if isinstance(arg, torch.Tensor):
+            return arg.detach().requires_grad_(arg.requires_grad)
+        elif isinstance(arg, TensorMeta):
+            if isinstance(arg, _DTensorMeta):
+                mesh = self._mesh_cache.get_mesh(arg.mesh_cache_key)
+                return arg.to_dtensor(self.device, mesh)
+            else:
+                return arg.to_tensor(self.device)
         else:
-            dbg_str += " running shape-inference at runtime"
+            raise PipeliningMetadataError(
+                f"Unsupported type {type(arg)} for _to_tensor: {arg}"
+            )
 
-        logger.debug(dbg_str)
+    def _ones_from_metadata(self, meta: TensorMeta) -> torch.Tensor:
+        """Create a ones tensor from metadata for backward inference grad_outputs."""
+        local_ones = torch.ones(
+            meta.shape,
+            dtype=meta.dtype,
+            device=self.device,
+        )
+        if isinstance(meta, _DTensorMeta):
+            mesh = self._mesh_cache.get_mesh(meta.mesh_cache_key)
+            return DTensor.from_local(
+                local_ones,
+                device_mesh=mesh,
+                placements=meta.placements,
+                shape=meta.global_shape,
+                stride=meta.global_stride,
+                run_check=False,
+            )
+        return local_ones
 
-    def _shape_inference(
+    def _forward_metadata_inference(
         self,
-        args: tuple[Any, ...],
+        args: tuple[torch.Tensor, ...] | _StageForwardMeta | None,
         kwargs: dict[str, Any] | None = None,
-    ):
-        if kwargs is None:
-            kwargs = {}
-        if args is None:
-            raise AssertionError("Args may be an empty tuple but not None")
-
-        # We skip recv communication if we're the first stage, but also if the previous stage is on the same rank
-        # and can pass its output shapes in as args instead of using send/recv.
-        if (
-            self.is_first
-            # if not first stage, then check if prev stage is on the same rank
-            or self.stage_index_to_group_rank[self.stage_index - 1] == self.group_rank
-        ):
-            logger.debug(
-                "Shape inference: stage %s skipping recv, because shape info passed in via `args`",
-                self.stage_index,
-            )
-            args = tree_map_only(torch.Tensor, lambda x: x.to("meta"), args)
+        has_backward: bool = False,
+    ) -> _StageForwardMeta | None:
+        """Run forward metadata inference (Stage 0 → N).
+
+        Args:
+            args: Real tensors (first stage), ``_StageForwardMeta``
+                (same-rank), or ``None`` (cross-rank P2P).
+            kwargs: Keyword arguments forwarded to the submodule.
+            has_backward: Whether backward inference follows.
+
+        Returns:
+            ``_StageForwardMeta`` for the next stage, or ``None`` if sent via P2P.
+        """
+        kwargs = kwargs or {}
+
+        # === RECEIVE: Get input metadata and create meta tensors ===
+        if self.is_first:
+            # First stage: extract metadata from real tensors
+            if args is None or isinstance(args, _StageForwardMeta):
+                raise PipeliningMetadataError(
+                    f"Stage {self.stage_index}: First stage requires real tensors, "
+                    f"got {type(args).__name__}."
+                )
+            tensor_args = validate_and_normalize_to_tuple(args)
+            assert tensor_args is not None  # noqa: S101
+            self._stage_meta.inputs = extract_tensor_metas(tensor_args)
+            inference_args = tuple(self._to_tensor(a) for a in tensor_args)
+        elif self._is_same_rank(self.stage_index - 1):
+            # Same-rank: _StageForwardMeta passed via argument
+            if not isinstance(args, _StageForwardMeta):
+                raise PipeliningMetadataError(
+                    f"Stage {self.stage_index}: Expected _StageForwardMeta from same-rank "
+                    f"previous stage, got {type(args).__name__}."
+                )
+            self._stage_meta.inputs = args.forward_metas
+            inference_args = tuple(self._to_tensor(m) for m in args.forward_metas)
         else:
-            if len(args) != 0:
-                raise AssertionError(
-                    "Can't supply input args for shape inference on non-first stage"
+            # Cross-rank: receive _StageForwardMeta via P2P
+            recv_meta = self._recv_meta(self.stage_index - 1)
+            if not isinstance(recv_meta, _StageForwardMeta):
+                raise PipeliningMetadataError(
+                    f"Stage {self.stage_index}: Expected _StageForwardMeta from P2P, "
+                    f"got {type(recv_meta).__name__}."
                 )
-            objects = [None]
-            logger.debug(
-                "Shape inference: stage %s receiving from stage %s",
-                self.stage_index,
-                self.stage_index - 1,
-            )
-            dist.recv_object_list(
-                objects,
-                src=dist.get_global_rank(
-                    self.group or dist.distributed_c10d._get_default_group(),
-                    self.stage_index_to_group_rank[self.stage_index - 1],
-                ),
-                group=self.group,
-                device=self.device,
-                use_batch=True,
+            self._stage_meta.inputs = recv_meta.forward_metas
+            inference_args = tuple(self._to_tensor(m) for m in recv_meta.forward_metas)
+
+        inference_kwargs = {
+            k: self._to_tensor(v) if isinstance(v, torch.Tensor) else v
+            for k, v in kwargs.items()
+        }
+
+        # Isolate metadata inference from user's grad context.
+        # has_backward → enable_grad() so backward tracing sees grad_fn;
+        # no backward → no_grad() for cross-rank consistency.
+        ctx = torch.enable_grad() if has_backward else torch.no_grad()
+        with ctx:
+            outputs = self._compute_outputs(
+                *inference_args, module=self.submod, **inference_kwargs
             )
-            recv_args = objects[0]
-            if not isinstance(recv_args, tuple):
-                raise AssertionError(f"Expected tuple, got {type(recv_args)}")
-            args = recv_args
-
-        # cache input shapes for use during recv buffer allocation
-        self.inputs_meta = args
-        args = tree_map_only(
-            torch.Tensor, lambda x: torch.zeros_like(x, device=self.device), args
-        )
 
-        # set attributes needed for forward
-        with torch.no_grad():
-            outputs = self.submod(*args, **kwargs)
+        # Normalize outputs to tuple
+        outputs = validate_and_normalize_to_tuple(outputs)
 
-        # if single tensor, convert so it is always a list
-        if isinstance(outputs, torch.Tensor):
-            outputs = [outputs]
+        self._stage_meta.outputs = extract_tensor_metas(outputs)
 
-        # communicate meta outputs not real outputs for two reasons
-        # 1 - its faster (esp. since obj coll pickles tensor data!)
-        # 2 - avoid activating a cuda context for the src rank when unpickling on the recv end!
-        outputs_meta = tuple(
-            tree_map_only(torch.Tensor, lambda x: x.to("meta"), outputs)
-        )
-        logger.debug(
-            "Shape inference: stage %s inputs %s, outputs %s",
-            self.stage_index,
-            self.inputs_meta,
-            outputs_meta,
+        # Store for backward metadata inference (always, even during eval)
+        fwd_kwargs_tensors = tuple(
+            v for v in flatten_args(inference_kwargs) if isinstance(v, torch.Tensor)
         )
-        self._configure_outputs_meta(outputs_meta)
-
-        # Passing outputs to the next stage:
-        # two cases-
-        # 1. Usually: use send/recv communication to pass the output
-        # 2. Special case: for V-schedules, 2 'adjacent' stages (e.g. stage 3, 4 in an 8-stage 4-rank V)
-        #    pass their shape info via return value and function args rather than send/recv.
-        if (
-            self.is_last
-            # if not last stage, then check if next stage is on the same rank
-            or self.stage_index_to_group_rank[self.stage_index + 1] == self.group_rank
-        ):
-            # Case (2) above: pass shape info via return value and caller passes it as args to next stage's
-            # _shape_inference call
-            logger.debug(
-                "Shape inference: stage %s skipping send to next stage",
-                self.stage_index,
+        self._fwd_outputs_for_bwd_meta = outputs
+        self._fwd_inputs_for_bwd_meta = inference_args
+        self._fwd_kwargs_tensors_for_bwd_meta = fwd_kwargs_tensors
+
+        # === SEND: Pass output metadata to next stage ===
+        if self._stage_meta.outputs is None:
+            raise PipeliningMetadataError(
+                f"Stage {self.stage_index}: output metadata is required for forward inference."
             )
+        fwd_meta = _StageForwardMeta(forward_metas=self._stage_meta.outputs)
 
+        if self.is_last or self._is_same_rank(self.stage_index + 1):
+            # Same-rank or last: return for caller to pass
+            return fwd_meta
         else:
-            # Case (1): send shapes via send operation, and ensure not to return it to the caller
-            logger.debug(
-                "Shape inference: stage %s sending to stage %s",
-                self.stage_index,
-                self.stage_index + 1,
+            # Cross-rank: send via P2P
+            self._send_meta(fwd_meta, self.stage_index + 1)
+            return None
+
+    def _backward_metadata_inference(
+        self,
+        loss_fn: Callable[..., torch.Tensor] | None = None,
+        target: torch.Tensor | None = None,
+        received_grad_meta: _StageBackwardMeta | None = None,
+    ) -> _StageBackwardMeta | None:
+        """Run backward metadata inference (Stage N → 0).
+
+        Args:
+            loss_fn: Loss function (required for the last stage).
+            target: Target tensor (required for the last stage).
+            received_grad_meta: Grad metadata from next same-rank stage
+                (V-schedule only).
+
+        Returns:
+            ``_StageBackwardMeta`` for the previous stage, or ``None`` if sent via P2P.
+        """
+        fwd_outputs = self._fwd_outputs_for_bwd_meta
+        fwd_inputs = self._fwd_inputs_for_bwd_meta
+        if fwd_outputs is None or fwd_inputs is None:
+            raise PipeliningMetadataError(
+                "Backward metadata inference requires forward metadata inference to run first"
             )
-            dist.send_object_list(
-                [outputs_meta],
-                dst=dist.get_global_rank(
-                    self.group or dist.distributed_c10d._get_default_group(),
-                    self.stage_index_to_group_rank[self.stage_index + 1],
-                ),
-                group=self.group,
-                device=self.device,
-                use_batch=True,
+        kwargs_tensors = self._fwd_kwargs_tensors_for_bwd_meta or ()
+        all_fwd_inputs = list(fwd_inputs) + list(kwargs_tensors)
+        # Clear temporary storage early — local refs are sufficient from here
+        self._fwd_outputs_for_bwd_meta = None
+        self._fwd_inputs_for_bwd_meta = None
+        self._fwd_kwargs_tensors_for_bwd_meta = None
+        # === RECEIVE: Get output grad metadata (except last stage) ===
+        if self.is_last:
+            if loss_fn is None or target is None:
+                raise PipeliningMetadataError(
+                    f"Stage {self.stage_index}: loss_fn and target required for last stage"
+                )
+            inference_target = self._to_tensor(target)
+            loss = loss_fn(
+                fwd_outputs[0] if len(fwd_outputs) == 1 else fwd_outputs,
+                inference_target,
+            )
+            self._stage_meta.output_grads = None
+            all_input_grads = self._compute_input_grads(
+                [loss],
+                all_fwd_inputs,
             )
-            outputs_meta = tuple()
+        else:
+            # Non-last stage: receive grad metadata from next stage
+            if self._is_same_rank(self.stage_index + 1):
+                # Same-rank: _StageBackwardMeta passed via argument
+                if not isinstance(received_grad_meta, _StageBackwardMeta):
+                    raise PipeliningMetadataError(
+                        f"Stage {self.stage_index}: Expected _StageBackwardMeta from same-rank "
+                        f"next stage, got {type(received_grad_meta).__name__}."
+                    )
+                self._stage_meta.output_grads = received_grad_meta.backward_metas
+            else:
+                # Cross-rank: receive _StageBackwardMeta via P2P
+                recv_meta = self._recv_meta(self.stage_index + 1)
+                if not isinstance(recv_meta, _StageBackwardMeta):
+                    raise PipeliningMetadataError(
+                        f"Stage {self.stage_index}: Expected _StageBackwardMeta from P2P, "
+                        f"got {type(recv_meta).__name__}."
+                    )
+                self._stage_meta.output_grads = recv_meta.backward_metas
+
+            # === COMPUTE: Build grad_outputs and compute input grads ===
+            # Extract output tensors and corresponding grad_outputs from metadata
+            # Must iterate together to maintain alignment
+            if self._stage_meta.output_grads is None:
+                raise PipeliningMetadataError(
+                    f"Stage {self.stage_index}: output_grads metadata is required for backward inference."
+                )
+            stage_output_grad_metas = self._stage_meta.output_grads
+
+            filtered_fwd_outputs: list[torch.Tensor] = []
+            filtered_output_grads: list[torch.Tensor | None] = []
+
+            for idx, (fwd_out, grad_meta) in enumerate(
+                zip(fwd_outputs, stage_output_grad_metas, strict=True)
+            ):
+                # Match _backward.py behavior: skip if output doesn't require grad AND has no grad_fn
+                if not fwd_out.requires_grad:
+                    if grad_meta is not None:
+                        raise PipeliningMetadataError(
+                            f"Stage {self.stage_index}: output {idx} requires_grad=False, "
+                            f"but output_grads metadata is provided: {grad_meta}."
+                        )
+                    continue
+                filtered_fwd_outputs.append(fwd_out)
+                # For outputs that require grad, include them even if grad_meta is None
+                # (runtime passes None grad_outputs to autograd.backward in this case)
+                filtered_output_grads.append(
+                    self._ones_from_metadata(grad_meta) if grad_meta else None
+                )
 
-        return outputs_meta
+            if filtered_fwd_outputs:
+                all_input_grads = self._compute_input_grads(
+                    filtered_fwd_outputs, all_fwd_inputs, filtered_output_grads
+                )
+                # Free intermediate references early
+                filtered_fwd_outputs.clear()
+                filtered_output_grads.clear()
+                all_fwd_inputs.clear()
+                # Only positional input grads flow to previous stage
+            else:
+                all_input_grads = tuple(None for _ in range(len(all_fwd_inputs)))
+
+        input_grads = all_input_grads[: len(fwd_inputs)]
+        self._stage_meta.input_grads = tuple(
+            extract_tensor_meta(g) if isinstance(g, torch.Tensor) else None
+            for g in input_grads
+        )
+
+        # === SEND: Pass input grad metadata to previous stage ===
+        bwd_meta = _StageBackwardMeta(backward_metas=self._stage_meta.input_grads)
+
+        if self.is_first or self._is_same_rank(self.stage_index - 1):
+            # First rank or Same-rank: return for caller to pass
+            return bwd_meta
+        else:
+            # Cross-rank: send via P2P
+            self._send_meta(bwd_meta, self.stage_index - 1)
+            return None
+
+    def _post_metadata_inference_cleanup(self) -> None:
+        """Clean up FSDP side effects (unsharded params, stale grads, stored
+        tensors) after metadata inference with real tensors.
+        """
+        # Clear stored inference tensors (frees autograd graph + activations)
+        self._fwd_outputs_for_bwd_meta = None
+        self._fwd_inputs_for_bwd_meta = None
+        self._fwd_kwargs_tensors_for_bwd_meta = None
+
+        # Metadata inference runs real fwd/bwd, which unshards FSDP params and
+        # accumulates grads.  Reshard to free memory and clear stale grads.
+        for module in self.submod.modules():
+            if isinstance(module, FSDPModule):
+                module.reshard()
+                for param in module.parameters():
+                    param.grad = None
+
+    def _prepare_backward_infra(
+        self,
+        num_microbatches: int,
+        loss_fn: Callable[..., torch.Tensor] | None = None,
+        target: torch.Tensor | None = None,
+        received_grad_meta: "_StageBackwardMeta | None" = None,
+    ) -> "_StageBackwardMeta | None":
+        """Run backward metadata inference and prepare backward infrastructure.
+
+        Returns:
+            ``_StageBackwardMeta`` for the previous same-rank stage, or ``None``.
+        """
+        grad_meta_result: _StageBackwardMeta | None = None
+        if self._inference_mode == InferenceMode.DYNAMIC:
+            # DYNAMIC mode: run backward metadata inference
+            # received_grad_meta is used for same-rank V-schedule stages
+            grad_meta_result = self._backward_metadata_inference(
+                loss_fn=loss_fn,
+                target=target,
+                received_grad_meta=received_grad_meta,
+            )
+            # Validate dynamically inferred metadata against user-provided metadata
+            self._validate_inferred_metadata()
+        else:
+            # STATIC mode: metadata comes from user inputs, no validation needed
+            self._stage_meta.input_grads = self._user_meta.input_grads
+            self._stage_meta.output_grads = self._user_meta.output_grads
+            # For STATIC mode with plain tensors, if output_grads is not set but
+            # we have outputs, derive output_grads from outputs.
+            # (gradient shape == output shape, but requires_grad=False for gradients)
+            if self._stage_meta.output_grads is None:
+                if self._stage_meta.outputs is None:
+                    raise PipeliningMetadataError(
+                        f"Stage {self.stage_index}: output metadata is required for backward inference."
+                    )
+                self._stage_meta.output_grads = _derive_grad_metas(
+                    self._stage_meta.outputs
+                )
+            # Similarly, derive input_grads from inputs if not provided
+            if self._stage_meta.input_grads is None:
+                if self._stage_meta.inputs is None:
+                    raise PipeliningMetadataError(
+                        f"Stage {self.stage_index}: input metadata is required for backward inference."
+                    )
+                self._stage_meta.input_grads = _derive_grad_metas(
+                    self._stage_meta.inputs
+                )
+
+        # Note: grad_send_info is created lazily in get_bwd_send_ops() since
+        # it mirrors args_recv_info (already populated during forward).
+        self._setup_backward_recv_info(num_microbatches)
+        return grad_meta_result
+
+    def _validate_inferred_metadata(self) -> None:
+        """Validate dynamically inferred metadata against user-provided metadata."""
+        pairs = [
+            (self._user_meta.inputs, self._stage_meta.inputs, "input"),
+            (self._user_meta.outputs, self._stage_meta.outputs, "output"),
+            (self._user_meta.input_grads, self._stage_meta.input_grads, "input_grad"),
+            (
+                self._user_meta.output_grads,
+                self._stage_meta.output_grads,
+                "output_grad",
+            ),
+        ]
+        for user_val, stage_val, label in pairs:
+            if user_val and stage_val:
+                validate_tensors_metadata(
+                    f"Stage {self.stage_index} {label}",
+                    user_val,
+                    stage_val,
+                    warn_on_mismatch=True,
+                )
 
     def _prepare_forward_infra(
         self,
         num_microbatches: int,
-        args: tuple[Any, ...],
+        args: tuple[Any, ...] | _StageForwardMeta | None,
         kwargs: dict[str, Any] | None = None,
-    ) -> tuple[Any, ...]:
-        # TODO move self.device to an argument from step API (from its input tensors)?
-        if num_microbatches is None:
-            raise AssertionError("TODO fix num_microbatches")
-
-        outputs: tuple[Any, ...] = tuple()
-        if self.inputs_meta is None:
-            outputs = self._shape_inference(args, kwargs)
-
-        if self.inputs_meta is None:
-            raise AssertionError("Expected inputs_meta to be set after shape inference")
-        # Receive info during forward
-        # TODO: create args_recv_info lazily? (same needed for PipelineStage)
+        has_backward: bool = False,
+    ) -> _StageForwardMeta | None:
+        """Prepare the stage infrastructure for forward pass.
+
+        Returns:
+            ``_StageForwardMeta`` for next stage (same-rank), or ``None`` if sent via P2P.
+        """
+        if self._inference_mode is None:
+            raise PipeliningMetadataError(
+                f"Stage {self.stage_index}: inference mode not set. "
+                f"Run warmup vote protocol first."
+            )
+
+        fwd_meta_output: _StageForwardMeta | None = None
+
+        if self._inference_mode == InferenceMode.DYNAMIC:
+            # DYNAMIC mode: run forward metadata inference
+            # args may be _StageForwardMeta for same-rank V-schedule stages
+            fwd_meta_output = self._forward_metadata_inference(
+                args, kwargs, has_backward
+            )
+            # Validate dynamically inferred metadata against user-provided metadata
+            self._validate_inferred_metadata()
+        # STATIC mode: metadata comes from user inputs, no validation needed
+        else:
+            self._stage_meta.inputs = self._user_meta.inputs
+            self._stage_meta.outputs = self._user_meta.outputs
+
+        # Setup recv and send info
+        self._setup_forward_recv_info(num_microbatches, has_backward)
+        self._setup_forward_send_info()
+
+        return fwd_meta_output
+
+    def _setup_forward_recv_info(
+        self, num_microbatches: int, has_backward: bool
+    ) -> None:
+        """Setup receive info for forward pass."""
+        if self._stage_meta.inputs is None:
+            raise PipeliningMetadataError(
+                f"Stage {self.stage_index}: inputs metadata required for recv info."
+            )
         for chunk_id in range(num_microbatches):
-            if not self.is_first:
-                # We assume that we always receive from stage - 1
-                recv_infos = tuple(
+            if self.is_first:
+                # First stage: all inputs are root arguments (no recv needed)
+                self.args_recv_info[chunk_id] = tuple(
                     _RecvInfo(
-                        f"recv_for_{self.stage_index}_from_{self.stage_index - 1}",
-                        self.stage_index - 1,
-                        _make_tensor_from_meta(inp, self.device),
+                        input_name=f"root_input_{idx}",
+                        source=None,
+                        buffer=None,
+                        tensor_meta=meta,
+                        is_root_arg=True,
                     )
-                    for inp in self.inputs_meta
+                    for idx, meta in enumerate(self._stage_meta.inputs)
                 )
-                # In case there is backward pass, set requires_grad for receive buffers
-                if self.has_backward:
-                    for r in recv_infos:
-                        r.buffer.requires_grad_(True)
-
-                self.args_recv_info[chunk_id] = recv_infos
             else:
+                # Non-first stages: receive from previous stage
                 self.args_recv_info[chunk_id] = tuple(
-                    _RootArgPlaceholder(i) for i in self.inputs_meta
+                    _RecvInfo(
+                        input_name=f"recv_for_{self.stage_index}_from_{self.stage_index - 1}",
+                        source=self.stage_index - 1,
+                        buffer=_make_tensor_from_meta(meta, self.device),
+                        tensor_meta=meta,
+                    )
+                    for meta in self._stage_meta.inputs
                 )
 
-        # Send info during forward for each activation
-        # only need the rank that is being sent to
+    def _setup_forward_send_info(self) -> None:
+        """Setup send info for forward pass."""
         self.act_send_info: dict[int, list] = {}
-
-        for idx in range(len(self.get_outputs_meta())):
-            # We assume we always send to stage + 1
-            if not self.is_last:
-                self.act_send_info[idx] = [self.stage_index + 1]
-            else:
-                self.act_send_info[idx] = []
-
-        return outputs
+        if self._stage_meta.outputs is None:
+            raise PipeliningMetadataError(
+                f"Stage {self.stage_index}: outputs metadata required for recv info."
+            )
+        for idx in range(len(self._stage_meta.outputs)):
+            self.act_send_info[idx] = [self.stage_index + 1] if not self.is_last else []
 
     def _create_grad_recv_info(
         self,
         act_send_info: dict,
     ) -> tuple[_RecvInfo, ...]:
-        grad_recv_info: tuple[_RecvInfo, ...] = ()
+        grad_recv_infos: list[_RecvInfo] = []
         if not self.is_last:
+            # Ensure output_grads metadata is available
+            if self._stage_meta.output_grads is None:
+                raise PipeliningMetadataError(
+                    f"Stage {self.stage_index}: output_grads metadata is required for "
+                    f"creating grad recv info. Ensure backward metadata is populated."
+                )
+
             # Receiving gradients from multiple sources is not supported
             # hence we only take the first destination
-            grad_recv_info = tuple(
-                _RecvInfo(
-                    f"recv_grad_for_{self.stage_index}_from_{dst_list[0]}",
-                    dst_list[0],
-                    _make_tensor_from_meta(self.get_outputs_meta()[idx], self.device),
+            # Use a helper function to safely extract the metadata
+            output_grads = self._stage_meta.output_grads
+            for idx, dst_list in act_send_info.items():
+                if dst_list is None:
+                    raise PipeliningMetadataError(
+                        f"Stage {self.stage_index}: output {idx} is not sent to any stage."
+                    )
+                src = dst_list[0]
+                grad_meta = output_grads[idx]
+                grad_recv_infos.append(
+                    _RecvInfo(
+                        input_name=f"recv_grad_for_{self.stage_index}_from_{src}",
+                        source=src,
+                        buffer=_make_tensor_from_meta(grad_meta, self.device)
+                        if grad_meta
+                        else None,
+                        tensor_meta=grad_meta,
+                    )
                 )
-                for idx, dst_list in act_send_info.items()
-            )
-        return grad_recv_info
+        return tuple(grad_recv_infos)
diff --git a/torch/distributed/remote_device.py b/torch/distributed/remote_device.py
index 3ad0076f5e890..c03727a33584a 100644
--- a/torch/distributed/remote_device.py
+++ b/torch/distributed/remote_device.py
@@ -35,14 +35,12 @@ def __init__(self, remote_device: str | torch.device):
         elif isinstance(remote_device, str):
             fields = remote_device.split("/")
             if len(fields) == 2:
-                # pyrefly: ignore [bad-assignment]
                 self._worker_name, self._device = fields
             elif len(fields) == 1:
                 # Check if this is a valid device.
                 if _remote_device._is_valid_local_device(fields[0]):
                     self._device = fields[0]
                 else:
-                    # pyrefly: ignore [bad-assignment]
                     self._worker_name = fields[0]
                     self._device = "cpu"
             else:
@@ -64,7 +62,6 @@ def __init__(self, remote_device: str | torch.device):
                 # rank:<rank>/device format, extract rank
                 if fields[0] == "rank" and fields[1].isdigit():
                     self._rank = int(fields[1])  # type: ignore[assignment]
-                    # pyrefly: ignore [bad-assignment]
                     self._worker_name = None
                 else:
                     raise ValueError(PARSE_ERROR)
diff --git a/torch/distributed/rpc/__init__.py b/torch/distributed/rpc/__init__.py
index adf901d6b6e3e..d9cd701e41232 100644
--- a/torch/distributed/rpc/__init__.py
+++ b/torch/distributed/rpc/__init__.py
@@ -87,6 +87,7 @@ def is_available() -> bool:
     rendezvous_iterator: Generator[tuple[Store, int, int], None, None]
 
     __all__ += ["init_rpc", "BackendType", "TensorPipeRpcBackendOptions"]
+    # pyrefly: ignore [unresolvable-dunder-all]
     __all__ = __all__ + api.__all__ + backend_registry.__all__  # noqa: PLE0605
 
     def init_rpc(
diff --git a/torch/distributed/rpc/api.py b/torch/distributed/rpc/api.py
index 845ce0b7faf6c..74b3aba569537 100644
--- a/torch/distributed/rpc/api.py
+++ b/torch/distributed/rpc/api.py
@@ -136,13 +136,13 @@ def _gather_to_leader(sequence_id, worker_name, obj, worker_names=None):
     with _all_gather_dict_lock:
         if not worker_names:
             worker_names = _ALL_WORKER_NAMES
-            assert worker_name in worker_names, (
-                f"{worker_name} is not expected by leader."
-            )
+            if worker_name not in worker_names:
+                raise AssertionError(f"{worker_name} is not expected by leader.")
         states = _all_gather_sequence_id_to_states[sequence_id]
-        assert worker_name not in states.gathered_objects, (
-            f"{worker_name} reported intent sequence id {sequence_id} twice. "
-        )
+        if worker_name in states.gathered_objects:
+            raise AssertionError(
+                f"{worker_name} reported intent sequence id {sequence_id} twice. "
+            )
         states.gathered_objects[worker_name] = obj
         if worker_names == set(states.gathered_objects.keys()):
             states.proceed_signal.set()
@@ -152,9 +152,10 @@ def _broadcast_to_followers(sequence_id, objects_map):
     with _all_gather_dict_lock:
         states = _all_gather_sequence_id_to_states[sequence_id]
 
-    assert not states.proceed_signal.is_set(), (
-        f"Termination signal sequence id {sequence_id} got set twice."
-    )
+    if states.proceed_signal.is_set():
+        raise AssertionError(
+            f"Termination signal sequence id {sequence_id} got set twice."
+        )
     states.gathered_objects = objects_map
     states.proceed_signal.set()
 
@@ -201,9 +202,10 @@ def _all_gather(obj, worker_names=None, timeout: float = UNSET_RPC_TIMEOUT):
     function blocks until all workers have received the gathered results.
     """
     if not worker_names:
-        assert _ALL_WORKER_NAMES is not None, (
-            "`_ALL_WORKER_NAMES` is not initialized for `def _all_gather`."
-        )
+        if _ALL_WORKER_NAMES is None:
+            raise AssertionError(
+                "`_ALL_WORKER_NAMES` is not initialized for `def _all_gather`."
+            )
         worker_names = _ALL_WORKER_NAMES
     leader_name = min(worker_names)
 
@@ -530,7 +532,8 @@ def method(self, *args, **kwargs):
         owner, returns a reference to the local value.
     """
     docstring = getattr(method, "__doc__", None)
-    assert docstring is not None, "RRef user-facing methods should all have docstrings."
+    if docstring is None:
+        raise AssertionError("RRef user-facing methods should all have docstrings.")
 
     # Do surgery on pybind11 generated docstrings.
     docstring = docstring.replace(
@@ -688,8 +691,10 @@ def remote(to, func, args=None, kwargs=None, timeout=UNSET_RPC_TIMEOUT):
             )
         # attach profiling information
         if should_profile:
-            assert torch.autograd._profiler_enabled()
-            assert rf is not None
+            if not torch.autograd._profiler_enabled():
+                raise AssertionError
+            if rf is None:
+                raise AssertionError
             fut = rf._call_end_callbacks_on_future(rref._get_future())
             rref._set_profiling_future(fut)
 
@@ -744,8 +749,10 @@ def _invoke_rpc(
                 dst_worker_info, pickled_python_udf, tensors, rpc_timeout, is_async_exec
             )
         if should_profile:
-            assert torch.autograd._profiler_enabled()
-            assert rf is not None
+            if not torch.autograd._profiler_enabled():
+                raise AssertionError
+            if rf is None:
+                raise AssertionError
             # Schedule profiling callbacks to run when the future completes.
             # This returns a future that is completed when the original future
             # completes and the profiling callbacks have been completed as well,
diff --git a/torch/distributed/rpc/backend_registry.py b/torch/distributed/rpc/backend_registry.py
index f1c0a2d9bd51d..9e7185e6ae436 100644
--- a/torch/distributed/rpc/backend_registry.py
+++ b/torch/distributed/rpc/backend_registry.py
@@ -122,7 +122,8 @@ def _init_process_group(store, rank, world_size):
     # default group to be initialized.
     group = dist.ProcessGroupGloo(store, rank, world_size, process_group_timeout)
 
-    assert group is not None, "Failed to initialize default ProcessGroup."
+    if group is None:
+        raise AssertionError("Failed to initialize default ProcessGroup.")
 
     if (rank != -1) and (rank != group.rank()):
         raise RuntimeError(f"rank argument {rank} doesn't match pg rank {group.rank()}")
diff --git a/torch/distributed/rpc/internal.py b/torch/distributed/rpc/internal.py
index faef8afddfc2c..0cde066b2466b 100644
--- a/torch/distributed/rpc/internal.py
+++ b/torch/distributed/rpc/internal.py
@@ -274,7 +274,8 @@ def _start_record_function(exec_type, func_name, current_worker_name, dest_worke
     Returns:
         An instance of `torch.autograd._RecordFunction`.
     """
-    assert torch.autograd._profiler_enabled(), "Autograd profiler should be enabled."
+    if not torch.autograd._profiler_enabled():
+        raise AssertionError("Autograd profiler should be enabled.")
     profile_key = f"rpc_{exec_type.value}#{str(func_name)}({current_worker_name} -> {dest_worker_name})"
     rf = torch.autograd._RecordFunction()  # type: ignore[attr-defined]
     torch.autograd._run_before_callbacks(rf, profile_key)  # type: ignore[attr-defined]
diff --git a/torch/distributed/rpc/options.py b/torch/distributed/rpc/options.py
index c58a2bf923910..99471974e3e5f 100644
--- a/torch/distributed/rpc/options.py
+++ b/torch/distributed/rpc/options.py
@@ -1,12 +1,10 @@
 # mypy: allow-untyped-defs
-from typing import Union
-
 import torch
 
 from . import _is_tensorpipe_available, constants as rpc_contants
 
 
-DeviceType = Union[int, str, torch.device]
+DeviceType = int | str | torch.device
 
 __all__ = ["TensorPipeRpcBackendOptions"]
 
diff --git a/torch/distributed/rpc/server_process_global_profiler.py b/torch/distributed/rpc/server_process_global_profiler.py
index 29a916772d330..33ebd56ca95bb 100644
--- a/torch/distributed/rpc/server_process_global_profiler.py
+++ b/torch/distributed/rpc/server_process_global_profiler.py
@@ -176,13 +176,11 @@ def __exit__(self, exc_type, exc_val, exc_tb):
         flattened_function_events = list(
             itertools.chain.from_iterable(process_global_function_events)
         )
-        # pyrefly: ignore [bad-assignment]
         self.function_events = torch.autograd.profiler_util.EventList(
             flattened_function_events,
             use_device="cuda" if self.use_cuda else None,
             profile_memory=self.profile_memory,
         )
-        # pyrefly: ignore [missing-attribute]
         self.function_events._build_tree()
 
         self.process_global_function_events = process_global_function_events
diff --git a/torch/distributed/run.py b/torch/distributed/run.py
index 87eb663ccf615..87d00472edf4e 100644
--- a/torch/distributed/run.py
+++ b/torch/distributed/run.py
@@ -638,11 +638,11 @@ def comma_separated_list(value):
     parser.add_argument(
         "--master-port",
         "--master_port",
-        default=29500,
+        default=None,
         type=int,
         action=env,
         help="Port on the master node (rank 0) to be used for communication during distributed "
-        "training. It is only used for static rendezvous.",
+        "training. It is only used for static rendezvous. Defaults to 29500.",
     )
     parser.add_argument(
         "--local-addr",
@@ -685,6 +685,17 @@ def comma_separated_list(value):
         "Common additional signals: SIGUSR1,SIGUSR2 (used in SLURM environments).",
     )
 
+    parser.add_argument(
+        "--shutdown-timeout",
+        "--shutdown_timeout",
+        action=env,
+        type=int,
+        default=None,
+        help="Time in seconds to wait for graceful shutdown of worker processes before "
+        "sending SIGKILL. If not specified, uses TORCH_ELASTIC_SHUTDOWN_TIMEOUT environment "
+        "variable or defaults to 30 seconds.",
+    )
+
     parser.add_argument(
         "--virtual-local-rank",
         "--virtual_local_rank",
@@ -736,7 +747,7 @@ def determine_local_world_size(nproc_per_node: str):
         return int(nproc_per_node)
     except ValueError as e:
         if nproc_per_node == "cpu":
-            num_proc = os.cpu_count()
+            num_proc = torch._utils.cpu_count()
             device_type = "cpu"
         elif nproc_per_node == "gpu":
             if not torch.cuda.is_available():
@@ -758,7 +769,7 @@ def determine_local_world_size(nproc_per_node: str):
                 num_proc = torch.accelerator.device_count()
                 device_type = torch.accelerator.current_accelerator().type  # type: ignore[union-attr]
             else:
-                num_proc = os.cpu_count()
+                num_proc = torch._utils.cpu_count()
                 device_type = "cpu"
         else:
             raise ValueError(
@@ -817,7 +828,7 @@ def _get_logs_specs_class(logs_specs_name: str | None) -> type[LogsSpecs]:
             )
 
         logger.info(
-            "Using logs_spec '%s' mapped to %s", logs_specs_name, str(logs_specs_cls)
+            "Using logs_spec '%s' mapped to %s", logs_specs_name, logs_specs_cls
         )
     else:
         logs_specs_cls = DefaultLogsSpecs
@@ -915,6 +926,7 @@ def config_from_args(args) -> tuple[LaunchConfig, Callable | str, list[str]]:
         duplicate_stdout_filters=args.duplicate_stdout_filters,
         duplicate_stderr_filters=args.duplicate_stderr_filters,
         virtual_local_rank=args.virtual_local_rank,
+        shutdown_timeout=args.shutdown_timeout,
     )
 
     with_python = not args.no_python
@@ -976,6 +988,20 @@ def run(args):
             args.rdzv_endpoint,
             args.rdzv_id,
         )
+    elif (
+        args.rdzv_backend == "static"
+        and not args.rdzv_endpoint
+        and args.master_port is None
+    ):
+        _, max_nodes = parse_min_max_nnodes(args.nnodes)
+        if max_nodes == 1:
+            args.rdzv_backend = "c10d"
+            args.rdzv_endpoint = "localhost:0"
+            args.rdzv_id = str(uuid.uuid4())
+
+    # master_port is only used for the static rendezvous backend, not c10d
+    if args.master_port is None:
+        args.master_port = 29500
 
     config, cmd, cmd_args = config_from_args(args)
     elastic_launch(
diff --git a/torch/distributed/tensor/__init__.py b/torch/distributed/tensor/__init__.py
index 067d4c0917e9d..8795890f577fb 100644
--- a/torch/distributed/tensor/__init__.py
+++ b/torch/distributed/tensor/__init__.py
@@ -15,6 +15,7 @@
     zeros,
 )
 from torch.distributed.tensor.placement_types import (
+    _StridedShard,
     Partial,
     Placement,
     Replicate,
@@ -63,6 +64,7 @@
         Partial,
         Replicate,
         Shard,
+        _StridedShard,
     ]
 )
 
@@ -86,3 +88,9 @@
 rand.__module__ = "torch.distributed.tensor"
 randn.__module__ = "torch.distributed.tensor"
 zeros.__module__ = "torch.distributed.tensor"
+
+# Register DTensor dispatch for higher order operators
+from torch._higher_order_ops.print import _register_dtensor_impl
+
+
+_register_dtensor_impl()
diff --git a/torch/distributed/tensor/_api.py b/torch/distributed/tensor/_api.py
index c24efa3c47d59..00b984422926d 100644
--- a/torch/distributed/tensor/_api.py
+++ b/torch/distributed/tensor/_api.py
@@ -2,6 +2,7 @@
 # mypy: allow-untyped-defs
 # Copyright (c) Meta Platforms, Inc. and affiliates
 import copy
+import hashlib
 import inspect
 import warnings
 from collections.abc import Callable, Sequence
@@ -13,7 +14,11 @@
 import torch.distributed.tensor._random as random
 import torch.nn as nn
 from torch._export.wrappers import mark_subclass_constructor_exportable_experimental
-from torch.distributed.device_mesh import _mesh_resources, DeviceMesh
+from torch.distributed.device_mesh import (
+    _mesh_resources,
+    _register_distributed_opaque_types,
+    DeviceMesh,
+)
 from torch.distributed.tensor._collective_utils import check_tensor_meta, mesh_broadcast
 from torch.distributed.tensor._dtensor_spec import DTensorSpec, TensorMeta
 from torch.distributed.tensor._redistribute import (
@@ -98,6 +103,7 @@ def forward(  # type: ignore[override]
     ):
         ctx.dtensor_spec = input._spec
         ctx.grad_placements = grad_placements
+        ctx.set_materialize_grads(False)
         local_tensor = input._local_tensor
 
         # We need to return a fresh Tensor object there as autograd metadata
@@ -106,7 +112,10 @@ def forward(  # type: ignore[override]
         return local_tensor.view_as(local_tensor)
 
     @staticmethod
-    def backward(ctx, grad_output: torch.Tensor):  # type: ignore[override]
+    def backward(ctx, grad_output: torch.Tensor | None):  # type: ignore[override]
+        if grad_output is None:
+            return None, None
+
         dtensor_spec = ctx.dtensor_spec
         mesh = dtensor_spec.mesh
         grad_placements = ctx.grad_placements
@@ -141,13 +150,12 @@ def backward(ctx, grad_output: torch.Tensor):  # type: ignore[override]
                 ),
             )
         return (
-            # pyrefly: ignore [bad-argument-type]
-            DTensor(
-                # pyrefly: ignore [bad-argument-count]
+            DTensor.from_local(
                 grad_output,
-                grad_spec,
-                # pyrefly: ignore [unexpected-keyword]
-                requires_grad=grad_output.requires_grad,
+                grad_spec.device_mesh,
+                grad_spec.placements,
+                shape=grad_spec.shape,
+                stride=grad_spec.stride,
             ),
             None,
         )
@@ -163,9 +171,12 @@ def forward(  # type: ignore[override]
         run_check: bool,
         shape: torch.Size | None = None,
         stride: tuple[int, ...] | None = None,
+        grad_placements: tuple[Placement, ...] | None = None,
     ) -> "DTensor":
         ctx.forward_input_placements = placements
         ctx.forward_input_device_mesh = device_mesh
+        ctx.grad_placements = grad_placements
+        ctx.set_materialize_grads(False)
 
         if shape and stride:
             tensor_shape, tensor_stride = shape, stride
@@ -225,9 +236,13 @@ def forward(  # type: ignore[override]
         return dist_tensor
 
     @staticmethod
-    def backward(ctx, grad_output: "DTensor"):  # type: ignore[override]
+    def backward(ctx, grad_output: "DTensor | None"):  # type: ignore[override]
         forward_input_placements = ctx.forward_input_placements
         forward_input_device_mesh = ctx.forward_input_device_mesh
+        grad_placements = ctx.grad_placements
+
+        if grad_output is None:
+            return None, None, None, None, None, None, None
 
         # reshard to the placement when creating DistributedTensor
         # so that the gradient layout matches, and we could return
@@ -240,9 +255,15 @@ def backward(ctx, grad_output: "DTensor"):  # type: ignore[override]
         # was: Partial(fwd) - Partial(grad_output) - Partial(grad_input)
         # now: Partial(fwd) - Partial(grad_output) - Replicate(grad_input)
         # See DTensor.from_local docstring for gradient placement guarantees
-        normalized_placements: tuple[Placement, ...] = _normalize_placements_for_grad(
-            forward_input_placements
-        )
+
+        # user should provide grad_placements as there's no guarantee on input gradient placement
+        # if grad_placements is None, we provide default placement
+        if grad_placements is not None:
+            normalized_placements = grad_placements
+        else:
+            normalized_placements = _normalize_placements_for_grad(
+                forward_input_placements
+            )
 
         if grad_output.placements != normalized_placements:
             current_spec: DTensorSpec = grad_output._spec
@@ -259,11 +280,11 @@ def backward(ctx, grad_output: "DTensor"):  # type: ignore[override]
             )
             # TODO: return the redistributed local tensor directly without
             # differentiable backward. see if this make sense for all cases.
-            return output, None, None, None, None, None
+            return output, None, None, None, None, None, None
 
         # TODO: backward is also differentiable now, add a test
         # to test higher level gradients.
-        return grad_output.to_local(), None, None, None, None, None
+        return grad_output.to_local(), None, None, None, None, None, None
 
 
 class DTensor(torch.Tensor):
@@ -339,24 +360,32 @@ def __tensor_flatten__(self):
         protocol to inform how to flatten a DTensor to local tensor
         for PT2 tracing
         """
-        return ["_local_tensor"], (self._spec, self.requires_grad)
+        return ["_local_tensor", "device_mesh"], (
+            self._spec.placements,
+            self._spec.tensor_meta,
+            self._spec.shard_order,
+            self.requires_grad,
+        )
 
     @staticmethod
     def __tensor_unflatten__(inner_tensors, flatten_spec, outer_size, outer_stride):
-        assert flatten_spec is not None, (
-            "Expecting spec to be not None from `__tensor_flatten__` return value!"
-        )
+        if flatten_spec is None:
+            raise AssertionError(
+                "Expecting spec to be not None from `__tensor_flatten__` return value!"
+            )
         local_tensor = inner_tensors["_local_tensor"]
-        spec, requires_grad = flatten_spec
+        mesh = inner_tensors["device_mesh"]
+        placements, old_tensor_meta, shard_order, requires_grad = flatten_spec
         unflatten_tensor_meta = TensorMeta(
             shape=outer_size,
             stride=outer_stride,
-            dtype=spec.tensor_meta.dtype,
+            dtype=old_tensor_meta.dtype,
         )
         unflatten_spec = DTensorSpec(
-            spec.mesh,
-            spec.placements,
+            mesh,
+            placements,
             tensor_meta=unflatten_tensor_meta,
+            shard_order=shard_order,
         )
         # pyrefly: ignore [bad-argument-type]
         return DTensor(
@@ -367,6 +396,15 @@ def __tensor_unflatten__(inner_tensors, flatten_spec, outer_size, outer_stride):
             requires_grad=requires_grad,
         )
 
+    def _stable_hash_for_caching(self) -> str:
+        """
+        Return a stable hash for AOT autograd caching.
+        [See note: Tensor subclass stable hashing for AOT autograd cache]
+        """
+        # Combine spec's stable hash with requires_grad
+        cache_data = self._spec._stable_hash() + str(self.requires_grad)
+        return hashlib.blake2b(cache_data.encode(), digest_size=16).hexdigest()
+
     def __coerce_tangent_metadata__(self):
         if not any(isinstance(p, Partial) for p in self.placements):
             return self
@@ -379,22 +417,34 @@ def __coerce_same_metadata_as_tangent__(self, flatten_spec, expected_type=None):
         if expected_type is not None:
             return None
 
-        (spec, _) = flatten_spec  # Result of tensor_flatten()
+        (placements, _, _, _) = flatten_spec
         return self.redistribute(
             device_mesh=self.device_mesh,
-            placements=spec.placements,
+            placements=placements,
         )
 
     @classmethod
     def __torch_dispatch__(cls, func, types, args=(), kwargs=None):  # type: ignore[override]
-        # We just need to have an implementation here; the __torch_dispatch__ machinery
-        # calls into a specific C++ fast path that doesn't call here.
-        # See #167051 for details
-        # python_arg_parser.cpp: dispatch_on_subclass()
-        # -> python_variable.cpp: dispatchDTensorOp()
-        raise NotImplementedError(
-            "DTensor.__torch_dispatch__ should not actually get called"
-        )
+        # Base DTensor is normally dispatched via a C++ fast path (see #167051)
+        # and never reaches here. This implementation exists so that DTensor
+        # subclasses can delegate back via super().__torch_dispatch__().
+        # It unwraps subclass instances to base DTensor and re-calls the op,
+        # which re-enters dispatch and hits the C++ fast path.
+        def unwrap(t):
+            if isinstance(t, DTensor) and type(t) is not DTensor:
+                # pyrefly: ignore [bad-argument-type]
+                return DTensor(
+                    # pyrefly: ignore [bad-argument-count]
+                    t._local_tensor,
+                    t._spec,
+                    # pyrefly: ignore [unexpected-keyword]
+                    requires_grad=t.requires_grad,
+                )
+            return t
+
+        args = torch.utils._pytree.tree_map(unwrap, args)
+        kwargs = torch.utils._pytree.tree_map(unwrap, kwargs or {})
+        return func(*args, **kwargs)
 
     @staticmethod
     def from_local(
@@ -405,6 +455,7 @@ def from_local(
         run_check: bool = False,
         shape: torch.Size | None = None,
         stride: tuple[int, ...] | None = None,
+        grad_placements: Sequence[Placement] | None = None,
     ) -> "DTensor":
         """
         Create a :class:`DTensor` from a local torch.Tensor on each rank
@@ -433,6 +484,10 @@ def from_local(
             stride (tuple, optional): A List of int which specifies the stride of DTensor.
                 If not provided, ``stride`` will be computed assuming the given distributed
                 tensor is evenly sharded across ranks. default: None
+            grad_placements (List[:class:`Placement`], optional): specifies the expected
+                input gradient placements. The input gradient (a plain tensor) will be
+                redistributed to this placement before exiting DTensor. If not
+                specified, follows the default placement guarantees below. default: None
 
         Returns:
             A :class:`DTensor` object
@@ -515,6 +570,7 @@ def from_local(
             run_check,
             shape,
             stride,
+            tuple(grad_placements) if grad_placements is not None else None,
         )
 
     def to_local(
@@ -604,8 +660,9 @@ def redistribute(
         Returns:
             A :class:`DTensor` object
 
-        .. note:: ``redistribute`` is differentiable, which means user do not need to worry about
-            the backward formula of the redistribute operation.
+        .. note:: ``redistribute`` is twice-differentiable, which means user do not need to worry about
+            the backward formula of the redistribute operation, or its compatibility with autograd for
+            second-order gradients. Higher-order differentiation has not been tested (but may work).
 
         .. note:: ``redistribute`` currently only supports redistributing DTensor on the same DeviceMesh,
             Please file an issue if you need to redistribute DTensor to different DeviceMesh.
@@ -749,19 +806,6 @@ def __get_tensor_shard__(self, index):
         else:
             raise RuntimeError("Unsupported tensor type!")
 
-    @classmethod
-    def __metadata_guard__(
-        cls, orig: tuple[DTensorSpec, bool], other: tuple[DTensorSpec, bool]
-    ) -> bool:
-        # TODO - delete this - This is now unused after the PR -
-        # https://github.com/pytorch/pytorch/pull/165824
-        orig_spec, orig_requires_grad = orig
-        other_spec, other_requires_grad = other
-        return (
-            orig_spec._check_equals(other_spec, skip_shapes=True)
-            and orig_requires_grad == other_requires_grad
-        )
-
 
 def distribute_tensor(
     tensor: torch.Tensor,
@@ -914,7 +958,8 @@ def distribute_tensor(
             )
     placements = tuple(placements)
 
-    assert local_tensor is not None, "distributing a tensor should not be None"
+    if local_tensor is None:
+        raise AssertionError("distributing a tensor should not be None")
     # detach the local tensor passed to DTensor since after the construction
     # of DTensor, autograd would work on top of DTensor instead of local tensor
     spec = DTensorSpec(
@@ -1147,11 +1192,11 @@ def _dtensor_init_helper(  # type: ignore[no-untyped-def]
     placements = placements or tuple(Replicate() for _ in range(device_mesh.ndim))
 
     # check device_mesh against placements
-    assert device_mesh.ndim == len(placements), (
-        "mesh dimension does not match the length of placements"
-    )
+    if device_mesh.ndim != len(placements):
+        raise AssertionError("mesh dimension does not match the length of placements")
 
-    assert kwargs["layout"] == torch.strided, "layout value not supported!"
+    if kwargs["layout"] != torch.strided:
+        raise AssertionError("layout value not supported!")
     torch_stride = torch._prims_common.make_contiguous_strides_for(size)
 
     # get local tensor shape
@@ -1173,7 +1218,8 @@ def _dtensor_init_helper(  # type: ignore[no-untyped-def]
         if random.is_rng_supported_mesh(device_mesh) and not random._rng_tracker:
             random._rng_tracker = random.OffsetBasedRNGTracker(device_mesh)
 
-        assert random._rng_tracker is not None
+        if random._rng_tracker is None:
+            raise AssertionError
         with random._rng_tracker._distribute_region(spec):
             local_tensor = init_op(local_shape, **kwargs)
     else:
@@ -1459,3 +1505,12 @@ def zeros(  # type: ignore[no-untyped-def]
         device_mesh=device_mesh,
         placements=placements,
     )
+
+
+# Module-level alias for DTensor.from_local to support FX code generation.
+# FX uses __name__ instead of __qualname__ when generating code, which loses
+# the class context for static methods. This alias allows generated code like
+# `torch.distributed.tensor._api.from_local(...)` to work correctly.
+from_local = DTensor.from_local
+
+_register_distributed_opaque_types()
diff --git a/torch/distributed/tensor/_collective_utils.py b/torch/distributed/tensor/_collective_utils.py
index 007e951fcb874..08f8a8eaba3c8 100644
--- a/torch/distributed/tensor/_collective_utils.py
+++ b/torch/distributed/tensor/_collective_utils.py
@@ -10,37 +10,73 @@
 import torch.distributed.tensor._dtensor_spec as dtensor_spec
 from torch._C._distributed_c10d import _resolve_process_group
 from torch._logging import warning_once
+from torch.distributed._functional_collectives import _are_we_tracing
 from torch.distributed._local_tensor import (
     local_tensor_mode,
     maybe_run_for_local_tensor,
 )
 from torch.distributed.device_mesh import _mesh_resources, DeviceMesh
 from torch.distributed.distributed_c10d import (
-    _get_group_size_by_name,
     broadcast,
     get_group_rank,
     get_rank,
+    GroupName,
     ProcessGroup,
     scatter,
     Work,
 )
+from torch.fx.experimental.symbolic_shapes import guard_or_false
+from torch.types import IntLikeType
 
 
 logger = logging.getLogger(__name__)
 
+# Opaque types must be registered before defining schemas that reference them,
+# so the schema parser recognizes the type names and uses PyObjectType (which
+# wraps Python objects as ConcretePyObjectHolder) instead of AnyType (which
+# calls toTypeInferredIValue and fails for Python-only opaque types).
+from torch.distributed.device_mesh import _register_distributed_opaque_types
+
+
+_register_distributed_opaque_types()
+
+_dtensor_lib = torch.library.Library("_dtensor", "FRAGMENT")
+_dtensor_lib.define(
+    "mesh_get_process_group("
+    "torch.distributed.device_mesh.DeviceMesh mesh, int dim"
+    ") -> torch.distributed.distributed_c10d.ProcessGroup"
+)
+
+
+@torch.library.impl("_dtensor::mesh_get_process_group", "CompositeExplicitAutograd")
+def _mesh_get_process_group_impl(mesh, dim):
+    return mesh.get_group(dim)
+
+
+@torch.library.register_fake("_dtensor::mesh_get_process_group")
+def _mesh_get_process_group_fake(mesh, dim):
+    from torch._library.fake_class_registry import maybe_unwrap_fake_script_object
+
+    real_mesh = maybe_unwrap_fake_script_object(mesh)
+    return real_mesh.get_group(dim)
+
 
 @torch.library.register_fake("_dtensor::shard_dim_alltoall")
-def _shard_dim_alltoall_meta(input, gather_dim, shard_dim, group_name):
-    group_size = _get_group_size_by_name(group_name)
+def _shard_dim_alltoall_meta(
+    input, gather_dim, shard_dim, group_name: GroupName | ProcessGroup
+):
+    if isinstance(group_name, str):
+        # pyrefly: ignore[bad-argument-type]  # pyrefly bug
+        group_name = _resolve_process_group(group_name)
+    group_size = group_name.size()
     stacked_list = [torch.empty_like(input) for _ in range(group_size)]
-    group = _resolve_process_group(group_name)
-    group_rank = get_group_rank(group, get_rank())
+    group_rank = get_group_rank(group_name, get_rank())
 
-    return (
-        torch.cat(stacked_list, dim=gather_dim)
-        .chunk(group_size, dim=shard_dim)[group_rank]
-        .contiguous()
-    )
+    cat_tensor = torch.cat(stacked_list, dim=gather_dim)
+    # pyrefly: ignore [unsupported-operation]
+    chunk_size = cat_tensor.size(shard_dim) // group_size
+    chunk = torch.narrow(cat_tensor, shard_dim, group_rank * chunk_size, chunk_size)
+    return chunk.contiguous()
 
 
 def shard_dim_alltoall(input, gather_dim, shard_dim, mesh, mesh_dim):
@@ -54,15 +90,17 @@ def shard_dim_alltoall(input, gather_dim, shard_dim, mesh, mesh_dim):
         if isinstance(out, funcol.AsyncCollectiveTensor):
             # stick to the same behavior for the alltoall case, remove this once we enable alltoall async
             out = out.wait()
-        out = torch.chunk(out, mesh.size(mesh_dim), dim=shard_dim)[
+        from torch.distributed.tensor.placement_types import Shard
+
+        out = Shard._custom_chunk(out, mesh.size(mesh_dim), dim=shard_dim)[
             mesh.get_local_rank(mesh_dim)
         ]
         return out.contiguous()
 
-    group_name = funcol._resolve_group_name((mesh, mesh_dim))
+    group = funcol._resolve_group((mesh, mesh_dim))
     # TODO: enable async op for shard_dim_alltoall
     return torch.ops._dtensor.shard_dim_alltoall(
-        input, gather_dim, shard_dim, group_name
+        input, gather_dim, shard_dim, funcol._group_or_group_name(group)
     )
 
 
@@ -106,7 +144,8 @@ def mesh_scatter(
     if output.is_meta:
         return None
     dim_group = mesh.get_group(mesh_dim)
-    assert isinstance(dim_group, ProcessGroup)
+    if not isinstance(dim_group, ProcessGroup):
+        raise AssertionError
 
     if group_src == get_rank(dim_group):
         fut = scatter(
@@ -166,23 +205,35 @@ def mesh_broadcast(
     if tensor.is_meta:
         return None
     dim_group = mesh.get_group(mesh_dim)
-    assert isinstance(dim_group, ProcessGroup)
+    if not isinstance(dim_group, ProcessGroup):
+        raise AssertionError
 
     return broadcast(tensor, group=dim_group, async_op=async_op, group_src=group_src)
 
 
 @maybe_run_for_local_tensor
-def pad_tensor(tensor: torch.Tensor, pad_dim: int, pad_size: int) -> torch.Tensor:
-    if pad_size == 0:
+def pad_tensor(
+    tensor: torch.Tensor, pad_dim: int, pad_size: IntLikeType
+) -> torch.Tensor:
+    # During tracing, always emit the pad op even when pad_size=0 so all
+    # ranks produce identical FX graph structure (SPMD).
+    # guard_or_false returns False for symbolic sizes, so the pad is always
+    # emitted during tracing. In eager with concrete pad_size=0, it returns
+    # True and we skip the no-op pad.
+    if guard_or_false(pad_size == 0) and not _are_we_tracing():
         return tensor
     pad = [0, 0] * (tensor.ndim - pad_dim)
-    pad[-1] = pad_size
+    pad[-1] = pad_size  # pyrefly: ignore[unsupported-operation]
     return torch.nn.functional.pad(tensor, pad)
 
 
 @maybe_run_for_local_tensor
-def unpad_tensor(tensor: torch.Tensor, pad_dim: int, pad_size: int) -> torch.Tensor:
-    if pad_size == 0:
+def unpad_tensor(
+    tensor: torch.Tensor, pad_dim: int, pad_size: IntLikeType
+) -> torch.Tensor:
+    # During tracing, always emit the narrow op even when pad_size=0 so all
+    # ranks produce identical FX graph structure (SPMD).
+    if guard_or_false(pad_size == 0) and not _are_we_tracing():
         return tensor
     return tensor.narrow(
         pad_dim,
@@ -228,7 +279,8 @@ def check_tensor_meta(
 
 
 def spec_to_bytes(spec: "dtensor_spec.DTensorSpec") -> int:
-    assert spec.tensor_meta is not None, "spec should have tensor meta defined!"
+    if spec.tensor_meta is None:
+        raise AssertionError("spec should have tensor meta defined!")
     return spec.tensor_meta.dtype.itemsize * math.prod(spec.shape)
 
 
@@ -340,6 +392,8 @@ def _compute_placement_transition_cost(
 
     num_devices_on_mesh_dim = mesh_topo.mesh_dim_devices[mesh_dim]
 
+    # NOTE: is_shard() does not match _StridedShard; see _is_shard_like().
+    # Safe today: redistribute_cost bails with inf when shard_order is None.
     if current_placement.is_shard() and target_placement.is_replicate():
         # allgather gives larger comm bytes
         comm_bytes_gb *= num_devices_on_mesh_dim
@@ -414,7 +468,8 @@ def one_step_redistribute_cost(
     if mesh_dim == -1:
         return 0.0
 
-    assert current_placement is not None and target_placement is not None
+    if current_placement is None or target_placement is None:
+        raise AssertionError
 
     mesh_topo = MeshTopoInfo.build_from_mesh(current_spec.mesh)
     comm_bytes_gb = (
@@ -496,9 +551,8 @@ def redistribute_cost(
     else:
         transform_infos = _gen_transform_infos(current_spec, target_spec)
     for transform_info in transform_infos:
-        assert current_spec.tensor_meta is not None, (
-            "spec should have tensor meta defined!"
-        )
+        if current_spec.tensor_meta is None:
+            raise AssertionError("spec should have tensor meta defined!")
         current = transform_info.src_dst_placements[0]
         target = transform_info.src_dst_placements[1]
         mesh_dim = transform_info.mesh_dim
diff --git a/torch/distributed/tensor/_decompositions.py b/torch/distributed/tensor/_decompositions.py
index 14f413b8d856c..fdf298fb9acc5 100644
--- a/torch/distributed/tensor/_decompositions.py
+++ b/torch/distributed/tensor/_decompositions.py
@@ -7,18 +7,17 @@
 determine how placements propagate through the original op.
 """
 
+from __future__ import annotations
+
 import itertools
-from typing import Any
+from typing import Any, TYPE_CHECKING
 
 import torch
 from torch._decomp import decomposition_table
-from torch._ops import OpOverload
 from torch.distributed._functional_collectives import _are_we_tracing
 from torch.distributed.device_mesh import DeviceMesh
 from torch.distributed.tensor._dtensor_spec import DTensorSpec
 from torch.distributed.tensor._op_schema import OpSchema, OpStrategy, RuntimeSchemaInfo
-from torch.distributed.tensor._ops.utils import expand_to_full_mesh_op_strategy
-from torch.distributed.tensor._sharding_prop import ShardingPropagator
 from torch.distributed.tensor._utils import try_find_mesh_from_args
 from torch.distributed.tensor.placement_types import (
     _StridedShard,
@@ -63,6 +62,11 @@ def _infer_schema_info_from_op(op: OpOverload) -> RuntimeSchemaInfo:
 from torch.utils._pytree import tree_any, tree_flatten, tree_map, tree_map_only
 
 
+if TYPE_CHECKING:
+    from torch._ops import OpOverload
+    from torch.distributed.tensor._sharding_prop import ShardingPropagator
+
+
 def _extract_input_specs(op_schema: OpSchema) -> tuple[DTensorSpec | object, ...]:
     return op_schema.args_schema + tuple(op_schema.kwargs_schema.values())
 
@@ -109,8 +113,26 @@ def __torch_dispatch__(self, func, types, args=(), kwargs=None):
         else:
             output_sharding = self.sharding_prop.propagate_op_sharding(op_schema)
 
-        if output_sharding.needs_redistribute:  # pyrefly: ignore [missing-attribute]
-            raise RuntimeError(f"Decomposition requires redistribution for {func}")
+        if (
+            output_sharding.needs_redistribute  # pyrefly: ignore [missing-attribute]
+            and (
+                redistribute_schema
+                := output_sharding.redistribute_schema  # pyrefly: ignore [missing-attribute]
+            )
+            is not None
+        ):
+            # a pure .needs_redistribute check is too broad; we want to ban redistribution,
+            # but this flag is set for view ops that convert global shape -> local shape args.
+            # During decomposition tracing on meta tensors at global shape, the shape adjustment
+            # is irrelevant — only reject true redistribution.
+            for orig, desired in zip(
+                op_schema.args_spec,
+                redistribute_schema.args_spec,  # pyrefly: ignore [missing-attribute]
+            ):
+                if orig.placements != desired.placements:
+                    raise RuntimeError(
+                        f"Decomposition requires redistribution for {func}"
+                    )
 
         out = func(*args, **kwargs)
         # pyrefly: ignore [missing-attribute]
@@ -136,24 +158,39 @@ class DecompShardingStrategy:
     single-dimension strategies are then expanded to the full mesh.
     """
 
+    def __init__(self, sharding_prop: ShardingPropagator):
+        self.sharding_prop = sharding_prop
+        # Cache fake meshes per device type to avoid repeated allocation.
+        # A fake size-1 mesh ensures identical strategy computation across all ranks
+        # during decomposition tracing, avoiding potential SPMD divergence.
+        # False negatives are avoided (all sizes % 1 == 0), while false positives
+        # are caught on expansion to the real, multi-dim device mesh.
+        self._fake_meshes: dict[str, DeviceMesh] = {}
+
+    def _get_fake_mesh(self, device_type: str) -> DeviceMesh:
+        fake_mesh = self._fake_meshes.get(device_type)
+        if fake_mesh is None:
+            fake_mesh = DeviceMesh(device_type, [0], _init_backend=False, _rank=0)
+            self._fake_meshes[device_type] = fake_mesh
+        return fake_mesh
+
     @staticmethod
     def has_decomp(op: OpOverload) -> bool:
         # Check if op has a decomposition (explicit or CIA)
         return op in decomposition_table or op._can_decompose()
 
-    @staticmethod
-    def ensure_schema_info(op: OpOverload, sharding_prop: ShardingPropagator) -> None:
+    def ensure_schema_info(self, op: OpOverload) -> None:
         """
         Register schema_info for decomposition op on first invocation.
         Needed for correct shard prop cache key.
         """
-        if op not in sharding_prop.op_to_schema_info:
+        if op not in self.sharding_prop.op_to_schema_info:
             schema_info = _infer_schema_info_from_op(op)
-            sharding_prop.op_to_schema_info[op] = schema_info
+            self.sharding_prop.op_to_schema_info[op] = schema_info
 
-    @staticmethod
     def propagate_strategy(
-        op_schema: OpSchema, sharding_prop: ShardingPropagator
+        self,
+        op_schema: OpSchema,
     ) -> OpStrategy | None:
         if not tree_any(
             lambda x: isinstance(x, DTensorSpec),
@@ -161,29 +198,21 @@ def propagate_strategy(
         ):
             return None
 
-        candidate_placements = DecompShardingStrategy._get_candidate_placements(
-            op_schema
-        )
+        candidate_placements = self._get_candidate_placements(op_schema)
         mesh = try_find_mesh_from_args(
             op_schema.op,
             op_schema.args_schema + tuple(op_schema.kwargs_schema.values()),
         )
 
-        # Create a fake size-1 mesh where all ranks pretend to be rank 0.
-        # This ensures identical strategy computation across all ranks during
-        # decomposition tracing, avoiding potential SPMD divergence.
-        #
-        # Using a fake mesh could potentially cause false negatives/positives
-        # (in terms of valid sharding strategies). The size-1 mesh should theoretically avoid
-        # all false negatives (e.g. no unevenness problems; all sizes % 1 == 0), while false
-        # positives are meant to be caught on expandsion to to the real, multi-dim device mesh.
-        fake_mesh = DeviceMesh(mesh.device_type, [0], _init_backend=False, _rank=0)
+        fake_mesh = self._get_fake_mesh(mesh.device_type)
         single_dim_strategies = []
         output_placements: list[Placement | tuple[Placement, ...]] = []
         for input_placements in candidate_placements:
             try:
-                output = DecompShardingStrategy._propagate_through_decomp(
-                    op_schema, input_placements, fake_mesh, sharding_prop
+                output = self._propagate_through_decomp(
+                    op_schema,
+                    input_placements,
+                    fake_mesh,
                 )
             except NotImplementedError:
                 return None
@@ -205,17 +234,21 @@ def propagate_strategy(
             )
 
         n_outputs = len(output_placements)
-        strategy_schema = sharding_prop._wrap_with_op_strategy(op_schema)
+        strategy_schema = self.sharding_prop._wrap_with_op_strategy(op_schema)
+        # Import here to avoid circular import at module load time
+        from torch.distributed.tensor._ops.utils import (  # noqa: F811
+            expand_to_full_mesh_op_strategy,
+        )
+
         return expand_to_full_mesh_op_strategy(
             mesh, strategy_schema, single_dim_strategies, input_index=n_outputs
         )
 
-    @staticmethod
     def _propagate_through_decomp(
+        self,
         op_schema: OpSchema,
         placement: tuple[Placement | None],
         mesh: DeviceMesh,
-        sharding_prop: ShardingPropagator,
     ) -> Placement | tuple[Placement, ...]:
         op = op_schema.op
         if op in decomposition_table:
@@ -246,7 +279,7 @@ def to_meta(x):
             args_meta = tree_map(to_meta, op_schema.args_schema)
             kwargs_meta = tree_map(to_meta, op_schema.kwargs_schema)
 
-            with PlacementTrackingMode(sharding_prop, mesh):
+            with PlacementTrackingMode(self.sharding_prop, mesh):
                 output = decomp_fn(*args_meta, **kwargs_meta)
 
         def get_placement(t):
@@ -293,5 +326,5 @@ def _get_candidate_placements(
                         options |= {Shard(i) for i in range(spec.ndim)}
                 candidates.append(list(options))
 
-        # pyrefly: ignore [no-matching-overload]
+        # pyrefly: ignore [bad-argument-type, no-matching-overload]
         return list(itertools.product(*candidates))
diff --git a/torch/distributed/tensor/_dispatch.py b/torch/distributed/tensor/_dispatch.py
index ffc77ac92e755..d346c8e17e117 100644
--- a/torch/distributed/tensor/_dispatch.py
+++ b/torch/distributed/tensor/_dispatch.py
@@ -11,13 +11,11 @@
 import torch.distributed.tensor._random as random
 from torch._library.utils import fill_defaults
 from torch._logging import LazyString
+from torch._prims.rng_prims import run_dtensor_rng_op
 from torch.distributed._functional_collectives import _are_we_tracing
 from torch.distributed.device_mesh import DeviceMesh
 from torch.distributed.tensor._dtensor_spec import DTensorSpec, TensorMeta
-from torch.distributed.tensor._nonlinear_redux import (
-    argminmax_handler,
-    minmax_dim_handler,
-)
+from torch.distributed.tensor._nonlinear_redux import argminmax_handler
 from torch.distributed.tensor._op_schema import (
     OpInfo,
     OpSchema,
@@ -49,6 +47,18 @@
 aten = torch.ops.aten
 logger = logging.getLogger(__name__)
 
+# The C++ DTensor dispatch fast path caches whether debug logging is
+# enabled.  Wrap setLevel so the cached flag is reset automatically.
+_orig_setLevel = logger.setLevel
+
+
+def _setLevel_and_reinit(level: int) -> None:
+    _orig_setLevel(level)
+    torch._C._reinit_DTensor_dispatch_logger()
+
+
+logger.setLevel = _setLevel_and_reinit  # type: ignore[method-assign]
+
 
 def as_strided_handler(
     op_call: torch._ops.OpOverload,
@@ -56,7 +66,8 @@ def as_strided_handler(
     kwargs: dict[str, object],
 ):
     args, kwargs = fill_defaults(op_call._schema, args, kwargs)
-    assert not kwargs
+    if kwargs:
+        raise AssertionError
     tensor, size, stride, storage_offset = args
     if (
         tensor.size() == tuple(size)
@@ -77,6 +88,15 @@ def is_same_size_handler(
     return lhs.shape == rhs.shape
 
 
+def is_pinned_handler(
+    op_call: torch._ops.OpOverload,
+    args: tuple[object, ...],
+    kwargs: dict[str, object],
+) -> bool:
+    tensor = cast(dtensor.DTensor, args[0])
+    return tensor._local_tensor.is_pinned()
+
+
 def found_inf_reduce_handler(
     op_call: torch._ops.OpOverload,
     args: tuple[object, ...],
@@ -153,16 +173,20 @@ def __init__(self) -> None:
             aten.bernoulli.default,
             aten.bernoulli_.float,
         }
+        self._squeeze_inplace_ops = {
+            aten.squeeze_.dim,
+            aten.squeeze_.default,
+            aten.squeeze_.dims,
+        }
         self._custom_op_handlers = {
             aten.is_same_size.default: is_same_size_handler,
+            aten.is_pinned.default: is_pinned_handler,
             aten.convolution.default: convolution_handler,
             aten.convolution_backward.default: convolution_backward_handler,
             aten._amp_foreach_non_finite_check_and_unscale_.default: found_inf_reduce_handler,
             aten.as_strided.default: as_strided_handler,
             aten.argmin.default: argminmax_handler,
             aten.argmax.default: argminmax_handler,
-            aten.max.dim: minmax_dim_handler,
-            aten.min.dim: minmax_dim_handler,
         }
 
     # ********************************************************************************************
@@ -215,10 +239,11 @@ def _propagate_op_sharding_dispatch_slow_path(
         # NOTE: schema should always be populated when calling this function,
         # as it's only called from C++ after unwrap_to_op_info (create_schema=True).
         # See dispatchDTensorOp in python_variable.cpp line 1453-1460.
-        assert op_info.schema is not None, (
-            "op_info.schema should not be None in sharding propagation. "
-            "This function should only be called after unwrap_to_op_info."
-        )
+        if op_info.schema is None:
+            raise AssertionError(
+                "op_info.schema should not be None in sharding propagation. "
+                "This function should only be called after unwrap_to_op_info."
+            )
         try:
             # We have basically inlined propagate() here, but WITHOUT the
             # output_sharding assignment
@@ -243,7 +268,8 @@ def _propagate_op_sharding_dispatch_slow_path(
                 # When running under inference mode, CompositeImplicitAutograd ops show up in __torch_dispatch__,
                 # so we manually decompose them, here
                 out = op_call.decompose(*args, **kwargs)
-                assert out is not NotImplemented
+                if out is NotImplemented:
+                    raise AssertionError from None
                 return out
             else:
                 raise
@@ -259,8 +285,10 @@ def _dispatch_get_local_results_slow_path(
         op_info: OpInfo,
     ) -> object:
         output_sharding = op_info.output_sharding
-        assert output_sharding is not None, "output sharding should not be None"
-        assert op_info is not None, "op_info should never be None"
+        if output_sharding is None:
+            raise AssertionError("output sharding should not be None")
+        if op_info is None:
+            raise AssertionError("op_info should never be None")
 
         # Record output placements for debugging
         debug_mode = get_active_debug_mode()
@@ -275,7 +303,8 @@ def _dispatch_get_local_results_slow_path(
             if output_sharding.needs_redistribute:
                 # If sharding propagation decision needs redistribute, perform redistribute
                 # on args first, which could potentially modify args (i.e. allgather certain arg)
-                assert output_sharding.redistribute_schema is not None
+                if output_sharding.redistribute_schema is None:
+                    raise AssertionError
                 self.redistribute_local_args(
                     op_info,
                     output_sharding.redistribute_schema,
@@ -296,9 +325,20 @@ def _dispatch_get_local_results_slow_path(
             local_tensor_args = cast(tuple[object, ...], local_tensor_args)
             if op_call in self._random_ops:
                 if not random._rng_tracker and is_rng_supported_mesh(mesh):
-                    # Default to `OffsetBasedRNGTracker` if the parallelism API
-                    # did not already construct one
-                    random._rng_tracker = random.OffsetBasedRNGTracker(mesh)
+                    # Default to `OffsetBasedRNGTracker` if the parallelism API did not already construct one
+                    # Skip RNG state sync during tracing to avoid lazily initializing real RNG state under fake mode.
+                    run_state_sync = not _are_we_tracing()
+                    if not run_state_sync:
+                        logger.info(
+                            "DTensor RNG tracker is being lazily initialized during tracing. "
+                            "RNG states may not be synchronized across ranks, which can lead "
+                            "to silent incorrectness. Please call `torch.manual_seed()` with "
+                            "the same seed on all ranks before compiling DTensor random ops.",
+                            stacklevel=2,
+                        )
+                    random._rng_tracker = random.OffsetBasedRNGTracker(
+                        mesh, run_state_sync
+                    )
 
                 first_arg, first_local_arg = (
                     cast(dtensor.DTensor, args[0]),
@@ -309,24 +349,63 @@ def _dispatch_get_local_results_slow_path(
                 # so the op_call does not directly use it (we want op_call to fall back to the 'default' which is
                 # our RNG manager)
                 maybe_user_generator = op_info.local_kwargs.pop("generator", None)
-                assert maybe_user_generator is None or isinstance(
-                    maybe_user_generator, torch.Generator
-                )
-                # maybe_user_generator = None
-                rng_context = (
-                    random._rng_tracker._distribute_region(
-                        first_arg._spec, generator=maybe_user_generator
-                    )
-                    if random._rng_tracker and not first_local_arg.is_meta
-                    else contextlib.nullcontext()
-                )
-                # For DTensor random operator, run it within a RNGTracker context to
-                # ensure the random number generator is properly distributed.
-                with rng_context:
+                if not (
+                    maybe_user_generator is None
+                    or isinstance(maybe_user_generator, torch.Generator)
+                ):
+                    raise AssertionError
+
+                if (
+                    random._rng_tracker
+                    and not first_local_arg.is_meta
+                    and random._rng_tracker.distribute_region_enabled
+                ):
+                    if (
+                        maybe_user_generator is not None
+                        or first_local_arg.device.type != "cuda"
+                        or (
+                            not _are_we_tracing()
+                            and type(first_local_arg) is not torch.Tensor
+                        )
+                    ):
+                        with random._rng_tracker._distribute_region(
+                            first_arg._spec, generator=maybe_user_generator
+                        ):
+                            local_results = op_call(
+                                *local_tensor_args, **op_info.local_kwargs
+                            )
+                    else:
+                        # CUDA device without user generator, use HOP for traceability
+                        if not isinstance(
+                            random._rng_tracker, random.OffsetBasedRNGTracker
+                        ):
+                            raise AssertionError
+                        start_offset_incr, end_offset_incr = (
+                            random._rng_tracker._compute_rng_offsets(first_arg._spec)
+                        )
+                        local_results = run_dtensor_rng_op(
+                            start_offset_incr,
+                            end_offset_incr,
+                            op_call,
+                            *local_tensor_args,
+                            **op_info.local_kwargs,
+                        )
+                else:
+                    # No rng_tracker, meta tensor, or distribute_region disabled
                     local_results = op_call(*local_tensor_args, **op_info.local_kwargs)
             else:
                 # normal case, run local sharded op computation
-                local_results = op_call(*local_tensor_args, **op_info.local_kwargs)
+                if (
+                    output_sharding.needs_redistribute
+                    and output_sharding.redistribute_schema is not None
+                    and output_sharding.redistribute_schema.op != op_call
+                ):
+                    # Op was rewritten (e.g., squeeze.default → squeeze.dims)
+                    local_results = output_sharding.redistribute_schema.op(
+                        *local_tensor_args, **op_info.local_kwargs
+                    )
+                else:
+                    local_results = op_call(*local_tensor_args, **op_info.local_kwargs)
 
         else:
             # For a non-participating device (happens on rank that does not belong to
@@ -364,7 +443,8 @@ def default_tensor(spec: DTensorSpec) -> torch.Tensor:
                     local_results = [
                         default_tensor(s) if s is not None else None for s in spec
                     ]
-                    assert isinstance(local_results, list)
+                    if not isinstance(local_results, list):
+                        raise AssertionError
                     if None in local_results:
                         ret_type = str(ret_list[0].type)
                         raise NotImplementedError(
@@ -398,7 +478,8 @@ def _dispatch_fast_path_python_tail(
                 # The output of the equal op is a bool, by converting it into a
                 # a single value tensor, we can use all-reduce with min reduce op
                 # to simulate logical and.
-                assert local_results is None or isinstance(local_results, bool)
+                if not (local_results is None or isinstance(local_results, bool)):
+                    raise AssertionError
                 r = torch.tensor(
                     int(local_results) if local_results is not None else 1,
                     device=compute_mesh.device_type,
@@ -410,14 +491,14 @@ def _dispatch_fast_path_python_tail(
             # inplace op should return self instead of re-wrapping
             if output_sharding.output_spec is not None:
                 output_spec = output_sharding.output_spec
-                assert isinstance(output_spec, DTensorSpec)
-                assert isinstance(args[0], dtensor.DTensor)
-
-                # NOTE: aten.squeeze_.dim is an inplace op but it also may change
-                # the inplace argument's tensor meta. Here we choose to special case
-                # this op because as far as I know this is the only inplace op that
-                # has such as behavior. We can extend this special case if necessary.
-                if op_call == aten.squeeze_.dim:
+                if not isinstance(output_spec, DTensorSpec):
+                    raise AssertionError
+                if not isinstance(args[0], dtensor.DTensor):
+                    raise AssertionError
+
+                # NOTE: squeeze_ inplace ops may change the tensor's metadata
+                # (shape/strides). We special-case them to update the spec.
+                if op_call in self._squeeze_inplace_ops:
                     # update the spec to handle tensor meta changes
                     args[0]._spec = output_spec
                     # use return_and_correct_aliasing to match the outer and the inner
@@ -456,10 +537,12 @@ def _dispatch_fast_path_python_tail(
                     out_dts.append(out_dt)
                     spec_idx += 1
 
-            assert len(out_dts) >= 1, "out variant should have at least one out arg"
+            if len(out_dts) < 1:
+                raise AssertionError("out variant should have at least one out arg")
             return tuple(out_dts) if len(out_dts) > 1 else out_dts[0]
         else:
-            assert op_call == aten.equal.default, op_call
+            if op_call != aten.equal.default:
+                raise AssertionError(op_call)
             ret = self.wrap(local_results, output_sharding.output_spec)  # type: ignore[possibly-undefined]
             if participating and op_call._schema._is_view_op():
                 return return_and_correct_aliasing(op_call, args, kwargs, ret)
@@ -523,6 +606,13 @@ def redistribute_local_args(
                 else:
                     new_local_args.append(arg_spec)
 
+        # Append extra non-tensor args from rewritten schema (e.g., dims tuple).
+        if use_val_from_redistribute_schema:
+            for i in range(
+                len(op_info.flat_args_schema), len(flatten_args_schema_to_reshard)
+            ):
+                new_local_args.append(flatten_args_schema_to_reshard[i])
+
         op_info.local_args = tuple(new_local_args)
 
     def unwrap_to_op_info(
@@ -618,9 +708,10 @@ def _contains_tensor(arg: object) -> bool:
                 kwargs_schema[k] = v
                 local_kwargs[k] = v
 
-        assert compute_mesh is not None, (
-            f"found no DeviceMesh from dtensor args for {op_call}!"
-        )
+        if compute_mesh is None:
+            raise AssertionError(
+                f"found no DeviceMesh from dtensor args for {op_call}!"
+            )
         op_info = OpInfo(
             compute_mesh,
             OpSchema(
@@ -647,19 +738,22 @@ def _contains_tensor(arg: object) -> bool:
     def wrap(res: object, spec: OutputSpecType) -> object:
         if isinstance(res, torch.Tensor):
             if spec is not None:
-                assert isinstance(spec, DTensorSpec), (
-                    f"output spec does not match with output! Expected DTensorSpec, got {spec}."
-                )
+                if not isinstance(spec, DTensorSpec):
+                    raise AssertionError(
+                        f"output spec does not match with output! Expected DTensorSpec, got {spec}."
+                    )
                 # pyrefly: ignore [bad-argument-type, bad-argument-count, unexpected-keyword]
                 return dtensor.DTensor(res, spec, requires_grad=res.requires_grad)
             else:
                 # if output does not have a DTensorSpec due to specific ops, it must be a scalar tensor
-                assert res.ndim == 0, "output tensor should be scalar!"
+                if res.ndim != 0:
+                    raise AssertionError("output tensor should be scalar!")
                 return res
         elif isinstance(res, (list, tuple)):
-            assert spec is not None and isinstance(spec, (list, tuple)), (
-                f"output spec does not match with output! Expected list/tuple, got {spec}."
-            )
+            if not (spec is not None and isinstance(spec, (list, tuple))):
+                raise AssertionError(
+                    f"output spec does not match with output! Expected list/tuple, got {spec}."
+                )
             res_list = []
             for e, s in zip(res, spec):
                 # pyrefly: ignore [bad-argument-type]
diff --git a/torch/distributed/tensor/_dtensor_spec.py b/torch/distributed/tensor/_dtensor_spec.py
index 83e8c7fb132fb..d2456be1c7862 100644
--- a/torch/distributed/tensor/_dtensor_spec.py
+++ b/torch/distributed/tensor/_dtensor_spec.py
@@ -1,3 +1,4 @@
+import hashlib
 import itertools
 import math
 from collections import defaultdict
@@ -7,6 +8,7 @@
 import torch
 from torch.distributed.device_mesh import DeviceMesh
 from torch.distributed.tensor.placement_types import (
+    _is_shard_like,
     _MaskPartial,
     _StridedShard,
     Partial,
@@ -18,6 +20,13 @@
 from torch.utils._dtype_abbrs import dtype_abbrs
 
 
+# Defined here (not in placement_types.py) because decoding split_factor into
+# a shard order is a DTensorSpec concern — placement_types doesn't know about
+# shard orders.
+class _StridedShardNotDecodableError(ValueError):
+    """Raised when _StridedShard split_factor cannot be decoded into a shard order."""
+
+
 class ShardOrderEntry(NamedTuple):
     """
     Represents how a single tensor dimension is sharded across mesh dimensions.
@@ -129,10 +138,9 @@ def _normalize_placements_into_shard_order(
                 placements, mesh
             )
             if shard_order is None:
-                raise ValueError(
-                    "use_strided_shard_as_shard_order is True, but placements: "
-                    f"{placements} is unable to be interpreted into a corresponding "
-                    "shard_order"
+                raise _StridedShardNotDecodableError(
+                    f"_StridedShard placements {placements} cannot be decoded "
+                    "into a corresponding shard_order"
                 )
             normalized_placements = tuple(
                 [
@@ -162,12 +170,13 @@ def compute_default_shard_order(
         tensor_dim_to_mesh_dims: defaultdict[int, list[int]] = defaultdict(list)
         mesh_ndim = len(placements)
         for mesh_dim in range(mesh_ndim):
-            if isinstance(placements[mesh_dim], Shard | _StridedShard):
+            if _is_shard_like(placements[mesh_dim]):
                 placement = placements[mesh_dim]
                 shard_dim = placement.dim  # pyrefly: ignore [missing-attribute]
-                assert shard_dim >= 0, (
-                    f"Shard dim {shard_dim} in placements {placements} must be normalized"
-                )
+                if shard_dim < 0:
+                    raise AssertionError(
+                        f"Shard dim {shard_dim} in placements {placements} must be normalized"
+                    )
                 tensor_dim_to_mesh_dims[shard_dim].append(mesh_dim)
 
         # Convert dict into ShardOrderEntry tuples
@@ -312,9 +321,7 @@ def _maybe_convert_StridedShard_to_shard_order(
         """
         if not any(isinstance(p, _StridedShard) for p in placements):
             return DTensorSpec.compute_default_shard_order(placements)
-        max_tensor_dim = (
-            max([i.dim for i in placements if isinstance(i, Shard | _StridedShard)]) + 1
-        )
+        max_tensor_dim = max([i.dim for i in placements if _is_shard_like(i)]) + 1
         shard_order = []
 
         tensor_dim_to_mesh_dims_order: list[list[int]] = [
@@ -322,8 +329,7 @@ def _maybe_convert_StridedShard_to_shard_order(
         ]
         for mesh_dim in reversed(range(len(placements))):
             cur_placement = placements[mesh_dim]
-            # _StridedShard may not be a subclass of Shard in the future, so write in this way:
-            if isinstance(cur_placement, Shard | _StridedShard):
+            if _is_shard_like(cur_placement):
                 tensor_dim = cur_placement.dim
                 mesh_dims_order = tensor_dim_to_mesh_dims_order[tensor_dim]
                 cur_sf = 1
@@ -361,28 +367,34 @@ def _verify_shard_order(self, shard_order: ShardOrder) -> None:
         """Verify that the shard_order is valid and matches the placements."""
         total_shard = 0
         if any(isinstance(p, _StridedShard) for p in self.placements):
+            # _StridedShard shard_order validation not yet supported;
+            # the Shard-only checks below (line 390, 394) would fail.
             return
         prev_tensor_dim = -1
         for entry in shard_order:
             tensor_dim = entry.tensor_dim
             mesh_dims = entry.mesh_dims
-            assert len(mesh_dims) > 0, f"shard_order {shard_order} has empty mesh dim"
-            assert tensor_dim >= 0, (
-                f"shard_order {shard_order} has invalid tensor dim {tensor_dim}"
-            )
-            assert tensor_dim > prev_tensor_dim, (
-                "tensor dim should be sorted in shard_order"
-            )
+            if len(mesh_dims) <= 0:
+                raise AssertionError(f"shard_order {shard_order} has empty mesh dim")
+            if tensor_dim < 0:
+                raise AssertionError(
+                    f"shard_order {shard_order} has invalid tensor dim {tensor_dim}"
+                )
+            if tensor_dim <= prev_tensor_dim:
+                raise AssertionError("tensor dim should be sorted in shard_order")
             prev_tensor_dim = tensor_dim
             total_shard += len(mesh_dims)
             for mesh_dim in mesh_dims:
-                assert 0 <= mesh_dim < len(self.placements), (
-                    f"shard_order {shard_order} has invalid mesh dim {mesh_dims}"
-                )
-                assert self.placements[mesh_dim] == Shard(tensor_dim), (
-                    f"placement[{mesh_dim}] doesn't have a matching shard in shard_order"
-                )
-        assert total_shard == sum(1 for p in self.placements if isinstance(p, Shard))
+                if not (0 <= mesh_dim < len(self.placements)):
+                    raise AssertionError(
+                        f"shard_order {shard_order} has invalid mesh dim {mesh_dims}"
+                    )
+                if self.placements[mesh_dim] != Shard(tensor_dim):
+                    raise AssertionError(
+                        f"placement[{mesh_dim}] doesn't have a matching shard in shard_order"
+                    )
+        if total_shard != sum(1 for p in self.placements if isinstance(p, Shard)):
+            raise AssertionError
 
     def __setattr__(self, attr: str, value: Any) -> None:
         if attr == "shard_order" and value is not None:
@@ -407,7 +419,21 @@ def __setattr__(self, attr: str, value: Any) -> None:
             # TODO: the TensorMetadata arises from
             # test/distributed/tensor/experimental/test_tp_transform.py::TensorParallelTest::test_tp_transform_e2e
             # but I actually can't reproduce it, maybe it is also a bug!
-            assert isinstance(value, TensorMeta | TensorMetadata), value
+            if not isinstance(value, TensorMeta | TensorMetadata):
+                raise AssertionError(repr(value))
+
+    def _hash_key(self) -> tuple[Any, ...]:
+        """Return the tuple used for hashing. Used by both __hash__ and _stable_hash."""
+        if self.tensor_meta is not None:
+            return (
+                self.mesh,
+                self.placements,
+                self.shard_order,
+                self.tensor_meta.shape,
+                self.tensor_meta.stride,
+                self.tensor_meta.dtype,
+            )
+        return (self.mesh, self.placements, self.shard_order)
 
     def _hash_impl(self) -> int:
         # hashing and equality check for DTensorSpec are used to cache the sharding
@@ -415,18 +441,7 @@ def _hash_impl(self) -> int:
         # dtype and stride.
         # Caveat: we need to keep this in mind and sync hash and eq if we add more
         # fields to them.
-        if self.tensor_meta is not None:
-            return hash(
-                (
-                    self.mesh,
-                    self.placements,
-                    self.shard_order,
-                    self.tensor_meta.shape,
-                    self.tensor_meta.stride,
-                    self.tensor_meta.dtype,
-                )
-            )
-        return hash((self.mesh, self.placements, self.shard_order))
+        return hash(self._hash_key())
 
     def __hash__(self) -> int:
         # We lazily cache the spec to avoid recomputing the hash upon each
@@ -437,6 +452,17 @@ def __hash__(self) -> int:
             self._hash = self._hash_impl()
         return self._hash
 
+    def _stable_hash(self) -> str:
+        """
+        Return a stable hash for AOT autograd caching.
+        [See note: Tensor subclass stable hashing for AOT autograd cache]
+        """
+        # Get hash key, but replace mesh with its stable hash
+        key = self._hash_key()
+        # First element is mesh, replace with its stable hash
+        stable_key = (self.mesh._stable_hash(),) + key[1:]
+        return hashlib.blake2b(repr(stable_key).encode(), digest_size=16).hexdigest()
+
     def _check_equals(self, other: object, skip_shapes: bool = False) -> bool:
         if not (
             isinstance(other, DTensorSpec)
@@ -538,14 +564,15 @@ def format_shard_order_str(
         # native dtensor-style sharding representation: map from mesh
         # dim to tensor dim
         for mesh_dim, placement in enumerate(placements):
-            if isinstance(placement, (Shard, _StridedShard)):
+            if _is_shard_like(placement):
                 if shard_order is not None:
                     for entry in shard_order:
                         tensor_dim = entry.tensor_dim
                         mesh_dims = entry.mesh_dims
 
                         if placement.dim == tensor_dim:
-                            assert mesh_dim in mesh_dims
+                            if mesh_dim not in mesh_dims:
+                                raise AssertionError
                             if len(mesh_dims) > 1:
                                 out_str += f"{placement}[{mesh_dims.index(mesh_dim)}]"
                             else:
@@ -581,7 +608,7 @@ def ndim(self) -> int:
     def num_shards(self) -> int:
         num_shards = 1
         for i, placement in enumerate(self.placements):
-            if placement.is_shard():
+            if _is_shard_like(placement):
                 num_shards *= self.mesh.size(i)
         return num_shards
 
@@ -617,8 +644,8 @@ def dim_map(self) -> list[int]:
         # and int >=0 represent shard on that device mesh dim
         r = [-1] * self.ndim
         for i, placement in enumerate(self.placements):
-            if placement.is_shard():
-                shard_dim = cast(Shard, placement).dim
+            if _is_shard_like(placement):
+                shard_dim = placement.dim
                 if r[shard_dim] > -1:
                     raise ValueError(
                         f"Tensor dim {shard_dim} is already sharded on mesh dim {r[shard_dim]},"
@@ -645,9 +672,8 @@ def num_shards_map(self) -> list[int]:
         """
         r = [1] * self.ndim
         for i, placement in enumerate(self.placements):
-            if placement.is_shard():
-                shard_dim = cast(Shard, placement).dim
-                r[shard_dim] *= self.mesh.size(i)
+            if _is_shard_like(placement):
+                r[placement.dim] *= self.mesh.size(i)
 
         return r
 
@@ -696,7 +722,7 @@ def from_dim_map(
         for i, m in enumerate(dim_map):
             if m >= 0:
                 placement = placements[m]
-                if placement.is_shard():
+                if placement.is_shard():  # dim_map only produces Shard placements
                     placement = cast(Shard, placement)
                     raise RuntimeError(
                         f"DeviceMesh dimension can't be mapped to two dimension of the same tensor: {i} and {placement.dim}"
@@ -717,9 +743,9 @@ def is_replicated(self) -> bool:
 
     def is_sharded(self) -> bool:
         """
-        return True if the current DTensorSpec uses Shard() placement on any mesh dims (devices)
+        return True if the current DTensorSpec uses Shard() or _StridedShard() placement on any mesh dims (devices)
         """
-        return any(placement.is_shard() for placement in self.placements)
+        return any(_is_shard_like(placement) for placement in self.placements)
 
     def shallow_copy_with_tensor_meta(
         self, tensor_meta: TensorMeta | None
@@ -727,9 +753,11 @@ def shallow_copy_with_tensor_meta(
         """
         Shallow copy the DTensorSpec with a new tensor_meta.
         """
-        assert tensor_meta is not None, "shallow copy with no tensor_meta!"
+        if tensor_meta is None:
+            raise AssertionError("shallow copy with no tensor_meta!")
         return DTensorSpec(
             self.mesh,
             self.placements,
             tensor_meta=tensor_meta,
+            use_strided_shard_as_shard_order=self.use_strided_shard_as_shard_order,
         )
diff --git a/torch/distributed/tensor/_nonlinear_redux.py b/torch/distributed/tensor/_nonlinear_redux.py
index 720264e5ceb11..a94d6f91a8ebe 100644
--- a/torch/distributed/tensor/_nonlinear_redux.py
+++ b/torch/distributed/tensor/_nonlinear_redux.py
@@ -1,6 +1,6 @@
 import operator
 from functools import reduce
-from typing import cast, Optional
+from typing import cast
 
 import torch
 import torch.distributed._functional_collectives as funcol
@@ -32,7 +32,8 @@ def _get_output_sharding(
     op_info = dtensor.DTensor._op_dispatcher.unwrap_to_op_info(op_call, args, kwargs)
     dtensor.DTensor._op_dispatcher.sharding_propagator.propagate(op_info)
     output_sharding = op_info.output_sharding
-    assert output_sharding is not None, "output sharding should not be None"
+    if output_sharding is None:
+        raise AssertionError("output sharding should not be None")
     return output_sharding
 
 
@@ -45,7 +46,7 @@ def _prep_arguments(
     torch.Size,
     "torch.distributed.device_mesh.DeviceMesh",
     tuple[Placement, ...],
-    Optional[int],
+    int | None,
     bool,
 ]:
     """
@@ -60,7 +61,7 @@ def _prep_arguments(
         keepdim: Whether to keep the reduced dimension
     """
     input_dtensor = cast(dtensor.DTensor, args[0])
-    dim: Optional[int] = None
+    dim: int | None = None
     keepdim: bool = False
 
     if not isinstance(input_dtensor, dtensor.DTensor):
@@ -93,7 +94,7 @@ def _prep_arguments(
 
 
 def _get_expected_shape(
-    local_tensor: torch.Tensor, dim: Optional[int], keepdim: bool
+    local_tensor: torch.Tensor, dim: int | None, keepdim: bool
 ) -> torch.Size:
     """Compute the expected output shape after reduction."""
     input_shape = list(local_tensor.shape)
@@ -117,7 +118,7 @@ def _collect_shard_mesh_dims(
     op_call_repr: str,
     local_tensor: torch.Tensor,
     placements: tuple[Placement, ...],
-    dim: Optional[int],
+    dim: int | None,
 ) -> list[int]:
     """Collect mesh dimensions that are sharded along the reduction dimension."""
     shard_mesh_dims: list[int] = []
@@ -135,7 +136,7 @@ def _convert_to_global_idxs(
     global_shape: torch.Size,
     device_mesh: "torch.distributed.device_mesh.DeviceMesh",
     placements: tuple[Placement, ...],
-    dim: Optional[int],
+    dim: int | None,
 ) -> tuple[int, torch.Tensor]:
     """Convert local indices to global indices."""
     local_shape, global_offset = compute_local_shape_and_global_offset(
@@ -143,15 +144,17 @@ def _convert_to_global_idxs(
     )
 
     if dim is None:
-        local_coord = torch.unravel_index(local_idx, local_shape)
-        global_coord = torch.stack(local_coord)
+        # Convert flat local index → flat global index using arithmetic ops
+        # instead of torch.unravel_index, which doesn't support SymInt shapes.
+        gathered_idxs = torch.zeros_like(local_idx)
+        remaining = local_idx
+        for i in range(len(local_shape)):
+            local_stride = reduce(operator.mul, local_shape[i + 1 :], 1)
+            global_stride = reduce(operator.mul, global_shape[i + 1 :], 1)
+            coord = remaining // local_stride
+            remaining = remaining % local_stride
+            gathered_idxs = gathered_idxs + (coord + global_offset[i]) * global_stride
         gather_dim = 0
-        for i, offset in enumerate(global_offset):
-            global_coord[i] += offset
-        # compute with proper striding
-        gathered_idxs = torch.tensor(0, device=local_idx.device, dtype=torch.long)
-        for i, coord in enumerate(global_coord):
-            gathered_idxs += coord * reduce(operator.mul, global_shape[i + 1 :], 1)
     else:
         gather_dim = dim
         gathered_idxs = local_idx + global_offset[dim]
@@ -218,8 +221,10 @@ def argminmax_handler(
 
     # Compute local reduction
     if dim is None:
-        local_idx = op_call(local_tensor)
-        local_redux = local_tensor.flatten()[local_idx]
+        val_op = _ARGMINMAX_REDUCTION_OPS[op_call]
+        # unsqueeze scalars to 1-d so they can be allgathered
+        local_redux = val_op(local_tensor).unsqueeze(0)
+        local_idx = op_call(local_tensor).unsqueeze(0)
     else:
         val_op = _ARGMINMAX_REDUCTION_OPS[op_call]
         local_redux, local_idx = val_op(local_tensor, dim=dim, keepdim=True)
@@ -235,8 +240,10 @@ def argminmax_handler(
     gathered_redux, gather_idxs = _gather_tensors(
         gather_dim, gathered_idxs, local_redux, device_mesh, shard_mesh_dims
     )
-    # op_call here is argmin/argmax which returns indices only
-    rank_winner = op_call(gathered_redux, dim, True)
+    # Select the rank with the best value; use dim=0 when dim was None since
+    # the scalars were unsqueezed to 1-d for gathering
+    select_dim = 0 if dim is None else dim
+    rank_winner = op_call(gathered_redux, select_dim, True)
     final_idx = torch.gather(gather_idxs, dim=gather_dim, index=rank_winner)
 
     return dtensor.DTensor._op_dispatcher.wrap(
@@ -265,7 +272,8 @@ def minmax_dim_handler(
     )
 
     # Compute local reduction - min/max with dim always requires dim
-    assert dim is not None
+    if dim is None:
+        raise AssertionError
     local_redux, local_idx = op_call(local_tensor, dim=dim, keepdim=True)
 
     if not shard_mesh_dims:
diff --git a/torch/distributed/tensor/_op_schema.py b/torch/distributed/tensor/_op_schema.py
index 043535438e1d0..d413c67169171 100644
--- a/torch/distributed/tensor/_op_schema.py
+++ b/torch/distributed/tensor/_op_schema.py
@@ -71,7 +71,8 @@ def _rebuild_tensor_from_dtensor_meta(arg) -> object:
     """
     This is used to propagate tensor metadata, must be under fake mode
     """
-    assert arg.tensor_meta is not None, "DTensorSpec does not contain tensor_meta."
+    if arg.tensor_meta is None:
+        raise AssertionError("DTensorSpec does not contain tensor_meta.")
     return torch.empty_strided(
         arg.tensor_meta.shape,
         arg.tensor_meta.stride,
@@ -160,13 +161,15 @@ def mesh(self):
             return self.output_specs.mesh
         elif isinstance(self.output_specs, tuple):
             out_spec = self.output_specs[0]
-            assert isinstance(out_spec, DTensorSpec)
+            if not isinstance(out_spec, DTensorSpec):
+                raise AssertionError
             return out_spec.mesh
         elif self.output_specs is None:
             # For no-output ops, get mesh from input_specs
-            assert self.input_specs is not None and len(self.input_specs) > 0, (
-                "Cannot determine mesh: output_specs is None and input_specs is empty"
-            )
+            if self.input_specs is None or len(self.input_specs) <= 0:
+                raise AssertionError(
+                    "Cannot determine mesh: output_specs is None and input_specs is empty"
+                )
             return self.input_specs[0].mesh
         else:
             raise ValueError(
@@ -174,11 +177,13 @@ def mesh(self):
             )
 
     def input_spec(self, index: int = 0) -> DTensorSpec:
-        assert self.input_specs is not None, "input_specs of OpSpec is None!"
-        assert len(self.input_specs) > index, (
-            f"Invalid index {index} for input_specs of length "
-            f"{len(self.input_specs)}: {self.input_specs}"
-        )
+        if self.input_specs is None:
+            raise AssertionError("input_specs of OpSpec is None!")
+        if len(self.input_specs) <= index:
+            raise AssertionError(
+                f"Invalid index {index} for input_specs of length "
+                f"{len(self.input_specs)}: {self.input_specs}"
+            )
         return self.input_specs[index]
 
     def __str__(self) -> str:
@@ -257,7 +262,8 @@ def shape(self):
     @property
     def tensor_meta(self) -> TensorMeta:
         # TODO upstream this assert to DTensorSpec itself and fill any missing TensorMetas
-        assert self.strategies[0].output_spec.tensor_meta is not None
+        if self.strategies[0].output_spec.tensor_meta is None:
+            raise AssertionError
         return self.strategies[0].output_spec.tensor_meta
 
     def __hash__(self) -> int:
@@ -303,7 +309,8 @@ def childs(self) -> Sequence[StrategyType]:  # codespell:ignore childs
 
     def child_mesh(self, index: int) -> DeviceMesh:
         op_strategy = self.children[index]
-        assert isinstance(op_strategy, OpStrategy)
+        if not isinstance(op_strategy, OpStrategy):
+            raise AssertionError
         return op_strategy.mesh
 
     def __str__(self) -> str:
@@ -429,6 +436,9 @@ def convert_to_meta(item):
                 return item.tensor_meta
             elif isinstance(item, TupleStrategy):
                 return tuple(convert_to_meta(child) for child in item.children)
+            elif isinstance(item, (list, tuple)):
+                converted = [convert_to_meta(child) for child in item]
+                return type(item)(converted)
             else:
                 return item
 
@@ -443,6 +453,9 @@ def convert_to_meta(item):
                 return item.tensor_meta
             elif isinstance(item, TupleStrategy):
                 return tuple(convert_to_meta(child) for child in item.children)
+            elif isinstance(item, (list, tuple)):
+                converted = [convert_to_meta(child) for child in item]
+                return type(item)(converted)
             else:
                 return item
 
@@ -467,12 +480,14 @@ def __str__(self) -> str:
                 args_schema.append(str(arg))
                 device_mesh = arg.mesh
             elif isinstance(arg, OpStrategy):
-                assert len(arg.strategies) == 1
+                if len(arg.strategies) != 1:
+                    raise AssertionError
                 args_schema.append(_pretty_print_spec(arg.strategies[0].output_specs))
                 device_mesh = arg.mesh
             elif isinstance(arg, TupleStrategy):
                 first_op_strategy = arg.children[0]
-                assert isinstance(first_op_strategy, OpStrategy)
+                if not isinstance(first_op_strategy, OpStrategy):
+                    raise AssertionError
                 device_mesh = first_op_strategy.mesh
                 args_schema.append(str(arg))
             else:
@@ -529,18 +544,23 @@ def get_mesh_from_args(self, validate: bool = True) -> DeviceMesh:
             - for foreach like ops we need to check "zipped" inputs are on the same mesh
               for each index.
         """
-        first_arg = self.args_schema[0]
-        if isinstance(first_arg, (DTensorSpec, OpStrategy)):
-            mesh = first_arg.mesh
-        elif isinstance(first_arg, (list, tuple, TupleStrategy)):
-            first_elem = (
-                first_arg.children[0]
-                if isinstance(first_arg, TupleStrategy)
-                else first_arg[0]
-            )
-            assert isinstance(first_elem, (DTensorSpec, OpStrategy))
-            mesh = first_elem.mesh
-        else:
+        mesh = None
+        # Scan all args to find the first DTensorSpec/OpStrategy (not just the first arg)
+        for arg in self.args_schema:
+            if isinstance(arg, (DTensorSpec, OpStrategy)):
+                mesh = arg.mesh
+                break
+            elif isinstance(arg, (list, tuple, TupleStrategy)):
+                # Scan all elements in the list/tuple, not just the first one,
+                # to handle cases like List[Optional[Tensor]] where first elem may be None
+                elems = arg.children if isinstance(arg, TupleStrategy) else arg
+                for elem in elems:
+                    if isinstance(elem, (DTensorSpec, OpStrategy)):
+                        mesh = elem.mesh
+                        break
+                if mesh is not None:
+                    break
+        if mesh is None:
             raise ValueError(f"Cannot find device mesh from args for op : {self.op}.")
 
         if validate:
diff --git a/torch/distributed/tensor/_ops/_common_rules.py b/torch/distributed/tensor/_ops/_common_rules.py
index 77a6505f30937..5207f85947369 100644
--- a/torch/distributed/tensor/_ops/_common_rules.py
+++ b/torch/distributed/tensor/_ops/_common_rules.py
@@ -120,7 +120,8 @@ def merge_sharding(dim: str, a: int, b: int) -> int:
                 dim_to_sharding[dim] = merge_sharding(
                     dim, dim_to_sharding[dim], mesh_dim
                 )
-                assert dim_to_size[dim] == input_spec.shape[idx]
+                if dim_to_size[dim] != input_spec.shape[idx]:
+                    raise AssertionError
 
             # after merging sharding, we check if there're multiple
             # sharding on the same mesh dim.
@@ -165,7 +166,8 @@ def merge_sharding(dim: str, a: int, b: int) -> int:
                         d in input_dim
                         and input_spec.dim_map[input_dim.index(d)] == mesh_dim
                     ):
-                        assert input_spec.tensor_meta is not None
+                        if input_spec.tensor_meta is None:
+                            raise AssertionError
                         global_shape = input_spec.tensor_meta.shape
                         local_shape, _ = compute_local_shape_and_global_offset(
                             global_shape,
@@ -210,7 +212,8 @@ def merge_sharding(dim: str, a: int, b: int) -> int:
     # XXX: since we still need to have intermediate shape calculation, we need
     # to pass in the shape here. We should remove this once sharding decomp works
     # for ops like addmm
-    assert input_specs[0].tensor_meta is not None
+    if input_specs[0].tensor_meta is None:
+        raise AssertionError
     tensor_meta = TensorMeta(
         torch.Size(output_shape),
         input_specs[0].tensor_meta.stride,
diff --git a/torch/distributed/tensor/_ops/_conv_ops.py b/torch/distributed/tensor/_ops/_conv_ops.py
index df9b81ac5df6e..7e47bd552313d 100644
--- a/torch/distributed/tensor/_ops/_conv_ops.py
+++ b/torch/distributed/tensor/_ops/_conv_ops.py
@@ -1,10 +1,21 @@
 # Copyright (c) Meta Platforms, Inc. and affiliates
 # implement matrix related ops for distributed tensor
 
+from typing import Any
+
 import torch
 from torch.distributed.tensor._dtensor_spec import DTensorSpec, TensorMeta
-from torch.distributed.tensor._op_schema import OpSchema, OutputSharding
+from torch.distributed.tensor._op_schema import (
+    OpSchema,
+    OutputSharding,
+    RuntimeSchemaInfo,
+)
+from torch.distributed.tensor._ops.single_dim_strategy import (
+    _ShardingPlaceholder,
+    register_single_dim_strategy,
+)
 from torch.distributed.tensor._ops.utils import register_prop_rule
+from torch.distributed.tensor.placement_types import Partial, Placement, Replicate
 
 
 aten = torch.ops.aten
@@ -24,18 +35,26 @@ def convolution_rules(op_schema: OpSchema) -> OutputSharding:
         _groups,
     ) = op_schema.args_schema
 
-    assert isinstance(input_spec, DTensorSpec)
-    assert isinstance(weight_spec, DTensorSpec)
+    if not isinstance(input_spec, DTensorSpec):
+        raise AssertionError
+    if not isinstance(weight_spec, DTensorSpec):
+        raise AssertionError
     # bias_spec can be None (optional parameter in aten.convolution schema)
     if bias_spec is not None:
-        assert isinstance(bias_spec, DTensorSpec)
-    assert input_spec.tensor_meta is not None
-    assert weight_spec.tensor_meta is not None
+        if not isinstance(bias_spec, DTensorSpec):
+            raise AssertionError
+    if input_spec.tensor_meta is None:
+        raise AssertionError
+    if weight_spec.tensor_meta is None:
+        raise AssertionError
     in_shape = input_spec.tensor_meta.shape
     weight_shape = weight_spec.tensor_meta.shape
-    assert isinstance(stride, list), f"stride must be list, got {type(stride)}"
-    assert isinstance(padding, list), f"padding must be list, got {type(padding)}"
-    assert isinstance(dilation, list), f"dilation must be list, got {type(dilation)}"
+    if not isinstance(stride, list):
+        raise AssertionError(f"stride must be list, got {type(stride)}")
+    if not isinstance(padding, list):
+        raise AssertionError(f"padding must be list, got {type(padding)}")
+    if not isinstance(dilation, list):
+        raise AssertionError(f"dilation must be list, got {type(dilation)}")
     # weight_shape might not be torch.Size in all cases (e.g., SymIntArrayRef during tracing)
     # so we don't assert its type, just use it
     out_conv_shape = [
@@ -82,13 +101,18 @@ def convolution_backward_rules(op_schema: OpSchema) -> OutputSharding:
         _output_mask,
     ) = op_schema.args_schema
 
-    assert isinstance(grad_output_spec, DTensorSpec)
-    assert isinstance(input_spec, DTensorSpec)
-    assert isinstance(weight_spec, DTensorSpec)
+    if not isinstance(grad_output_spec, DTensorSpec):
+        raise AssertionError
+    if not isinstance(input_spec, DTensorSpec):
+        raise AssertionError
+    if not isinstance(weight_spec, DTensorSpec):
+        raise AssertionError
     # bias_shape_opt can be None (optional parameter in aten.convolution_backward schema)
     if bias_shape_opt is not None:
-        assert isinstance(bias_shape_opt, list)
-    assert input_spec.tensor_meta is not None
+        if not isinstance(bias_shape_opt, list):
+            raise AssertionError
+    if input_spec.tensor_meta is None:
+        raise AssertionError
     weight_tensor_meta = weight_spec.tensor_meta
 
     # Only create bias_tensor_meta if bias_shape_opt is not None
@@ -125,3 +149,54 @@ def convolution_backward_rules(op_schema: OpSchema) -> OutputSharding:
     # for a certain output Tensor. This also applies to the conv handler
     # in torch/distributed/tensor/_tp_conv.py
     return OutputSharding([grad_input_spec, grad_weight_spec, grad_bias_spec])
+
+
+# Single-dim strategies for autoparallel optimizer support.
+# These coexist with the prop_rules above — strategies take precedence
+# in the propagation path, while the prop_rules + custom handlers in
+# _tp_conv.py continue to handle runtime dispatch.
+
+
+@register_single_dim_strategy(
+    [aten.convolution.default],
+    schema_info=RuntimeSchemaInfo(2),
+)
+def convolution_single_dim_strategy(
+    op: torch._ops.OpOverload,
+    args_schema: tuple[Any, ...],
+    kwargs_schema: dict[str, Any],
+) -> list[list[Placement | _ShardingPlaceholder]]:
+    bias_meta = args_schema[2]
+    # [output, input, weight, (bias)]
+    rule: list[Placement | _ShardingPlaceholder] = [
+        _ShardingPlaceholder(0),  # output
+        _ShardingPlaceholder(0),  # input
+        Replicate(),  # weight
+    ]
+    if bias_meta is not None:
+        rule.append(Replicate())  # bias
+    return [rule]
+
+
+@register_single_dim_strategy(
+    [aten.convolution_backward.default],
+    schema_info=RuntimeSchemaInfo(3),
+)
+def convolution_backward_single_dim_strategy(
+    op: torch._ops.OpOverload,
+    args_schema: tuple[Any, ...],
+    kwargs_schema: dict[str, Any],
+) -> list[list[Placement | _ShardingPlaceholder | None]]:
+    bias_sizes = args_schema[3]
+    has_bias = bias_sizes is not None
+    # outputs: [grad_input, grad_weight, grad_bias]
+    # inputs: [grad_output, input, weight]
+    rule: list[Placement | _ShardingPlaceholder | None] = [
+        _ShardingPlaceholder(0),  # grad_input
+        Partial("sum"),  # grad_weight
+        Partial("sum") if has_bias else None,  # grad_bias
+        _ShardingPlaceholder(0),  # grad_output
+        _ShardingPlaceholder(0),  # input
+        Replicate(),  # weight
+    ]
+    return [rule]
diff --git a/torch/distributed/tensor/_ops/_einsum_strategy.py b/torch/distributed/tensor/_ops/_einsum_strategy.py
index 4a1c6033063a1..fc60980f4949d 100644
--- a/torch/distributed/tensor/_ops/_einsum_strategy.py
+++ b/torch/distributed/tensor/_ops/_einsum_strategy.py
@@ -30,8 +30,10 @@ def parse_equation(cls, equation: str) -> tuple[list[str], str]:
 
         # NOTE: only support at most two inputs, and single output
         # extend to support more inputs if needed in future
-        assert len(input_dims) <= 2, "Only support at most two inputs"
-        assert len(output_dims) == 1, "Only support single output"
+        if len(input_dims) > 2:
+            raise AssertionError("Only support at most two inputs")
+        if len(output_dims) != 1:
+            raise AssertionError("Only support single output")
         output_dim = output_dims[0]
         return input_dims, output_dim
 
@@ -63,9 +65,10 @@ def parse_dims(cls, input_dims: list[str], output_dim: str) -> "EinsumDims":
                 if is_batch_dim:
                     batch_dims.append(dim_char)
                 else:
-                    assert len(input_dims) == 2, (
-                        "free dimension only supported for two inputs!"
-                    )
+                    if len(input_dims) != 2:
+                        raise AssertionError(
+                            "free dimension only supported for two inputs!"
+                        )
                     lhs, rhs = input_dims
                     if dim_char in lhs:
                         lhs_out_only_dims.append(dim_char)
diff --git a/torch/distributed/tensor/_ops/_mask_buffer.py b/torch/distributed/tensor/_ops/_mask_buffer.py
index 26b0a713db42c..a28074184b205 100644
--- a/torch/distributed/tensor/_ops/_mask_buffer.py
+++ b/torch/distributed/tensor/_ops/_mask_buffer.py
@@ -15,7 +15,8 @@ def materialize_mask(self, mask):
         if self.refcount == 0:
             self.data = mask
         else:
-            assert self.data is not None
+            if self.data is None:
+                raise AssertionError
             if not torch.equal(self.data, mask):
                 raise RuntimeError(
                     "MaskBuffer has been materialized with conflicting data"
diff --git a/torch/distributed/tensor/_ops/_math_ops.py b/torch/distributed/tensor/_ops/_math_ops.py
index 1c437aea9dfb0..7f47c830a4cba 100644
--- a/torch/distributed/tensor/_ops/_math_ops.py
+++ b/torch/distributed/tensor/_ops/_math_ops.py
@@ -4,11 +4,11 @@
 from collections.abc import Sequence
 from dataclasses import dataclass
 from enum import Enum
-from typing import cast, Union
+from typing import Any, cast
 
 import torch
 from torch.distributed.device_mesh import DeviceMesh
-from torch.distributed.tensor._dtensor_spec import DTensorSpec
+from torch.distributed.tensor._dtensor_spec import DTensorSpec, TensorMeta
 from torch.distributed.tensor._op_schema import (
     OpSchema,
     OpSpec,
@@ -17,6 +17,10 @@
     RuntimeSchemaInfo,
     TupleStrategy,
 )
+from torch.distributed.tensor._ops.single_dim_strategy import (
+    _ShardingPlaceholder,
+    register_single_dim_strategy,
+)
 from torch.distributed.tensor._ops.utils import (
     as_list,
     expand_to_full_mesh_op_strategy,
@@ -29,6 +33,7 @@
 )
 from torch.distributed.tensor._utils import normalize_to_torch_size
 from torch.distributed.tensor.placement_types import (
+    _is_shard_like,
     _StridedShard,
     Partial,
     Placement,
@@ -38,6 +43,7 @@
 
 
 aten = torch.ops.aten
+prims = torch.ops.prims
 
 
 class Reduction(Enum):
@@ -51,7 +57,7 @@ class NormReduction:
     norm_type: int | float
 
 
-ReductionOpType = Union[NormReduction, str]
+ReductionOpType = NormReduction | str
 
 
 @dataclass(frozen=True)
@@ -154,7 +160,7 @@ def _replicate_dims_start_at(
 ) -> tuple[Placement, ...]:
     new_placements: list[Placement] = []
     for p in placements:
-        if p.is_partial() or (isinstance(p, Shard) and p.dim >= start_dim):
+        if p.is_partial() or (_is_shard_like(p) and p.dim >= start_dim):
             new_placements.append(Replicate())  # make it replicate
         else:
             new_placements.append(p)  # keep the placement
@@ -162,12 +168,16 @@ def _replicate_dims_start_at(
 
 
 # return new_placements which align with placements but skip the skipped_dim
+# Precondition: no shard-like placement on skipped_dim (callers must
+# replicate it first via replicate_reduction_dims).
 def _skip_dim(
     placements: tuple[Placement, ...], skipped_dim: int
 ) -> tuple[Placement, ...]:
     new_placements: list[Placement] = []
     for p in placements:
-        if isinstance(p, Shard) and p.dim >= skipped_dim:
+        if isinstance(p, _StridedShard) and p.dim >= skipped_dim:
+            new_placements.append(_StridedShard(p.dim - 1, split_factor=p.split_factor))
+        elif isinstance(p, Shard) and p.dim >= skipped_dim:
             new_placements.append(Shard(p.dim - 1))
         else:
             new_placements.append(p)
@@ -183,7 +193,7 @@ def replicate_reduction_dims(
     for p in placements:
         if p.is_partial():
             new_placements.append(Replicate())
-        elif isinstance(p, Shard) and p.dim in reduction_dims:
+        elif _is_shard_like(p) and p.dim in reduction_dims:
             new_placements.append(Replicate())
         else:
             new_placements.append(p)
@@ -205,7 +215,7 @@ def map_placements_after_reduction(
         if isinstance(placement, (Replicate, Partial)):
             new_placements.append(placement)
         else:
-            if not isinstance(placement, Shard | _StridedShard):
+            if not _is_shard_like(placement):
                 raise AssertionError(
                     f"Expected Shard/_StridedShard, got {type(placement)}"
                 )
@@ -216,14 +226,14 @@ def map_placements_after_reduction(
                 # (i.e. for the case where keepdims=True), we generate partial
                 new_placements.append(get_placement_from_reduction_op(reduction_op))
             else:
-                if isinstance(placement, Shard):
-                    new_placements.append(Shard(new_shard_dim))
-                else:
+                if isinstance(placement, _StridedShard):
                     new_placements.append(
                         _StridedShard(
                             new_shard_dim, split_factor=placement.split_factor
                         )
                     )
+                elif isinstance(placement, Shard):
+                    new_placements.append(Shard(new_shard_dim))
     return tuple(new_placements)
 
 
@@ -308,28 +318,34 @@ def common_reduction_strategy(
     aten.all.dim: "product",
     aten.sum.default: "sum",
     aten.sum.dim_IntList: "sum",
+    prims.sum.default: "sum",
     aten.any.default: "sum",
     aten.any.dim: "sum",
+    aten.any.dims: "sum",
     aten.any.out: "sum",
     # These are only valid when there is no padding
     aten.prod.default: "product",
     aten.prod.dim_int: "product",
     aten.prod.int_out: "product",
+    prims.prod.default: "product",
     # avg is only linear when there is no padding
     aten.mean.default: "avg",
     aten.mean.dim: "avg",
     aten.mean.out: "avg",
     aten.max.default: "max",
-    aten.max.dim: "max",
     aten.max.out: "max",
     aten.min.default: "min",
-    aten.min.dim: "min",
     aten.min.out: "min",
     aten.amax.default: "max",
     aten.amax.out: "max",
     aten.amin.default: "min",
     aten.amin.out: "min",
-    # argmax and argmin is using custom hanndler leveraging linear reduction of max and min
+    aten.nansum.default: "sum",
+}
+
+# argmax/argmin return indices which cannot be combined with P(max/min).
+# They need special handling that forces redistribution on reduction dims.
+ARGMAX_ARGMIN_OPS = {
     aten.argmax.default: "max",
     aten.argmin.default: "min",
 }
@@ -361,8 +377,79 @@ def linear_reduction_strategy(op_schema: OpSchema) -> OpStrategy:
     )
 
 
-@register_op_strategy(aten.cumsum.default, schema_info=RuntimeSchemaInfo(1))
-def cumsum_strategy(op_schema: OpSchema) -> OpStrategy:
+# max.dim/min.dim return (values, indices). Indices are local to each shard
+# and cannot be combined across ranks, so we force Replicate on reduction dims
+# (same approach as argmax/argmin).
+@register_single_dim_strategy(
+    [aten.max.dim, aten.min.dim], schema_info=RuntimeSchemaInfo(1)
+)
+def max_min_dim_single_dim_strategy(
+    op: torch._ops.OpOverload,
+    args_schema: tuple[Any, ...],
+    kwargs_schema: dict[str, Any],
+) -> list[list[Placement | _ShardingPlaceholder]]:
+    input_meta = args_schema[0]
+    if not isinstance(input_meta, TensorMeta):
+        raise AssertionError(f"Expected TensorMeta, got {type(input_meta)}")
+
+    ndim = len(input_meta.shape)
+    dim = normalize_dim(cast(int, args_schema[1]), ndim)
+    keep_dim = len(args_schema) > 2 and bool(args_schema[2])
+
+    strategies: list[list[Placement | _ShardingPlaceholder]] = []
+    for d in range(ndim):
+        if d == dim:
+            continue
+        out_d = d if keep_dim or d < dim else d - 1
+        # [values, indices, input]: shard on non-reduction dim
+        strategies.append(
+            [
+                _ShardingPlaceholder(out_d),
+                _ShardingPlaceholder(out_d),
+                _ShardingPlaceholder(d),
+            ]
+        )
+    return strategies
+
+
+@register_op_strategy(list(ARGMAX_ARGMIN_OPS.keys()), schema_info=RuntimeSchemaInfo(1))
+def argmax_argmin_strategy(op_schema: OpSchema) -> OpStrategy:
+    """
+    Strategy for argmax/argmin. These return indices, not values, so they cannot
+    use P(max/min) output placements. The indices are local to each shard and
+    cannot be meaningfully combined across ranks with a max/min reduction.
+    Force redistribution on reduction dimensions by using reduction_linear=False.
+    """
+    args_schema = op_schema.args_schema
+    input_strategy = args_schema[0]
+    if not isinstance(input_strategy, OpStrategy):
+        raise AssertionError(f"Expected OpStrategy, got {type(input_strategy)}")
+
+    dims = None
+    if len(op_schema.args_schema) > 1:
+        dims = _infer_reduction_dims(args_schema[1], input_strategy.ndim)
+
+    reduce_dims = list(range(input_strategy.ndim)) if dims is None else dims
+    keep_dim = len(op_schema.args_schema) > 2 and bool(op_schema.args_schema[2])
+    reduction_op = ARGMAX_ARGMIN_OPS[op_schema.op]
+    return common_reduction_strategy(
+        input_strategy,
+        reduce_dims,
+        keep_dim=keep_dim,
+        reduction_linear=False,  # Force redistribution - indices can't use P(max/min)
+        # reduction_op is effectively unused here: reduction_linear=False
+        # forces all reduction-dim Shard placements to Replicate before
+        # map_placements_after_reduction, so no Shard-on-reduction-dim
+        # remains to convert to Partial. Passed for consistency.
+        reduction_op=reduction_op,
+    )
+
+
+@register_op_strategy(
+    [aten.cumsum.default, aten.cumprod.default, aten.logcumsumexp.default],
+    schema_info=RuntimeSchemaInfo(1),
+)
+def scan_strategy(op_schema: OpSchema) -> OpStrategy:
     args_schema = op_schema.args_schema
     input_strategy = args_schema[0]
     if not isinstance(input_strategy, OpStrategy):
@@ -370,18 +457,105 @@ def cumsum_strategy(op_schema: OpSchema) -> OpStrategy:
     dim = args_schema[1]
     if not isinstance(dim, int):
         raise AssertionError(f"Expected int, got {type(dim)}")
-
     return common_reduction_strategy(
         input_strategy, [dim], keep_dim=True, reduction_linear=False
     )
 
 
+@register_op_strategy(
+    [aten.median.default, aten.nanmedian.default],
+    schema_info=RuntimeSchemaInfo(1),
+)
+def global_median_strategy(op_schema: OpSchema) -> OpStrategy:
+    input_strategy = cast(OpStrategy, op_schema.args_schema[0])
+    reduce_dims = list(range(input_strategy.ndim))
+    return common_reduction_strategy(
+        input_strategy, reduce_dims, reduction_linear=False
+    )
+
+
+@register_single_dim_strategy(
+    [aten.median.dim, aten.nanmedian.dim, aten.mode.default],
+    schema_info=RuntimeSchemaInfo(1),
+)
+def dim_reduction_with_indices_strategy(
+    op: torch._ops.OpOverload,
+    args_schema: tuple[Any, ...],
+    kwargs_schema: dict[str, Any],
+) -> list[list[Placement | _ShardingPlaceholder]]:
+    input_meta = args_schema[0]
+    if not isinstance(input_meta, TensorMeta):
+        raise AssertionError(f"Expected TensorMeta, got {type(input_meta)}")
+
+    ndim = len(input_meta.shape)
+    dim = normalize_dim(cast(int, args_schema[1]) if len(args_schema) > 1 else -1, ndim)
+    keep_dim = len(args_schema) > 2 and bool(args_schema[2])
+
+    strategies: list[list[Placement | _ShardingPlaceholder]] = []
+    for d in range(ndim):
+        if d == dim:
+            continue
+        out_d = d if keep_dim or d < dim else d - 1
+        strategies.append(
+            [
+                _ShardingPlaceholder(out_d),
+                _ShardingPlaceholder(out_d),
+                _ShardingPlaceholder(d),
+            ]
+        )
+    return strategies
+
+
+@register_single_dim_strategy(
+    [aten.kthvalue.default],
+    schema_info=RuntimeSchemaInfo(2),
+)
+def kthvalue_strategy(
+    op: torch._ops.OpOverload,
+    args_schema: tuple[Any, ...],
+    kwargs_schema: dict[str, Any],
+) -> list[list[Placement | _ShardingPlaceholder]]:
+    input_meta = args_schema[0]
+    if not isinstance(input_meta, TensorMeta):
+        raise AssertionError(f"Expected TensorMeta, got {type(input_meta)}")
+
+    ndim = len(input_meta.shape)
+    dim = normalize_dim(cast(int, args_schema[2]) if len(args_schema) > 2 else -1, ndim)
+    keep_dim = len(args_schema) > 3 and bool(args_schema[3])
+
+    strategies: list[list[Placement | _ShardingPlaceholder]] = []
+    for d in range(ndim):
+        if d == dim:
+            continue
+        out_d = d if keep_dim or d < dim else d - 1
+        strategies.append(
+            [
+                _ShardingPlaceholder(out_d),
+                _ShardingPlaceholder(out_d),
+                _ShardingPlaceholder(d),
+            ]
+        )
+    return strategies
+
+
+@register_op_strategy(
+    [aten.cummax.default, aten.cummin.default],
+    schema_info=RuntimeSchemaInfo(1),
+)
+def cummax_cummin_strategy(op_schema: OpSchema) -> OpStrategy:
+    dim = cast(int, op_schema.args_schema[1])
+    return sort_strategy(op_schema, dim)
+
+
 @register_op_strategy(
     [
         aten.std.correction,
         aten.std.correction_out,
         aten.var.correction,
         aten.var.correction_out,
+        aten.var_mean.correction,
+        aten.var_mean.correction_out,
+        prims.var.default,
     ],
     schema_info=RuntimeSchemaInfo(1, ["keepdim"]),
 )
@@ -414,12 +588,14 @@ def _get_norm_reduction_op(norm_type: int | float | str) -> ReductionOpType:
     elif norm_type in (float("-inf"), "-inf"):
         return "min"
     else:
-        assert isinstance(norm_type, (int, float))
+        if not isinstance(norm_type, (int, float)):
+            raise AssertionError
         return NormReduction(norm_type)
 
 
 @register_op_strategy(
-    [aten.linalg_vector_norm.default], schema_info=RuntimeSchemaInfo(1)
+    [aten.linalg_vector_norm.default, aten.norm.Scalar],
+    schema_info=RuntimeSchemaInfo(1),
 )
 def vector_norm_strategy(op_schema: OpSchema) -> OpStrategy:
     args_schema = op_schema.args_schema
@@ -564,12 +740,6 @@ def foreach_max_strategy(op_schema: OpSchema) -> TupleStrategy:
         aten.tril.default,
         aten.triu.default,
         aten._linalg_eigh.default,
-        aten.upsample_bicubic2d.default,
-        aten.upsample_bilinear2d.default,
-        aten.upsample_linear1d.default,
-        aten.upsample_nearest2d.default,
-        aten.upsample_trilinear3d.default,
-        # TODO: support the full F.interpolate set of options.
     ],
     schema_info=RuntimeSchemaInfo(1),
 )
@@ -604,6 +774,76 @@ def linalg_replicate_strategy(op_schema: OpSchema) -> OpStrategy:
     return OpStrategy(output_strategies)
 
 
+# Maps each pooling op to its spatial rank (number of spatial dimensions).
+# Batched inputs have layout (N, C, *spatial) with ndim = spatial_rank + 2;
+# unbatched inputs drop the batch dim giving ndim = spatial_rank + 1.
+POOL_SPATIAL_RANK: dict[torch._ops.OpOverload, int] = {
+    aten.avg_pool1d.default: 1,
+    aten.avg_pool2d.default: 2,
+    aten.avg_pool3d.default: 3,
+    aten.adaptive_avg_pool1d.default: 1,
+    aten._adaptive_avg_pool2d.default: 2,
+    aten._adaptive_avg_pool3d.default: 3,
+    aten.adaptive_max_pool1d.default: 1,
+    aten.adaptive_max_pool2d.default: 2,
+    aten.adaptive_max_pool3d.default: 3,
+    aten.fractional_max_pool2d.default: 2,
+    aten.fractional_max_pool3d.default: 3,
+    aten.max_pool1d_with_indices.default: 1,
+    aten.max_pool2d_with_indices.default: 2,
+    aten.max_pool3d_with_indices.default: 3,
+}
+
+AVG_POOL_OPS = [
+    aten.avg_pool1d.default,
+    aten.avg_pool2d.default,
+    aten.avg_pool3d.default,
+    aten.adaptive_avg_pool1d.default,
+    aten._adaptive_avg_pool2d.default,
+    aten._adaptive_avg_pool3d.default,
+]
+
+MAX_POOL_OPS = [
+    aten.adaptive_max_pool1d.default,
+    aten.adaptive_max_pool2d.default,
+    aten.adaptive_max_pool3d.default,
+    aten.fractional_max_pool2d.default,
+    aten.fractional_max_pool3d.default,
+    aten.max_pool1d_with_indices.default,
+    aten.max_pool2d_with_indices.default,
+    aten.max_pool3d_with_indices.default,
+]
+
+
+@register_op_strategy(
+    AVG_POOL_OPS + MAX_POOL_OPS,
+    schema_info=RuntimeSchemaInfo(1),
+)
+def pooling_strategy(op_schema: OpSchema) -> OpStrategy:
+    input_strategy = cast(OpStrategy, op_schema.args_schema[0])
+    mesh = input_strategy.mesh
+    num_outputs = 2 if op_schema.op in MAX_POOL_OPS else 1
+    num_inputs = len(op_schema.args_strategy) + len(op_schema.kwargs_strategy)
+    n = num_outputs + num_inputs
+    single_mesh_dim_strategies: list[PlacementList] = [
+        [Replicate()] * n,
+        [Shard(0)] * n,
+    ]
+    # avg_pool is linear: Partial(sum) and Partial(avg) pass through unchanged.
+    if op_schema.op in AVG_POOL_OPS:
+        single_mesh_dim_strategies.append([Partial("sum")] * n)
+        single_mesh_dim_strategies.append([Partial("avg")] * n)
+    # S(1) is safe when dim 1 is the channel dim (pooling never touches it).
+    # Batched inputs have layout (N, C, *spatial) with ndim = spatial_rank + 2.
+    spatial_rank = POOL_SPATIAL_RANK[op_schema.op]
+    is_batched = input_strategy.ndim >= spatial_rank + 2
+    if is_batched:
+        single_mesh_dim_strategies.append([Shard(1)] * n)
+    return expand_to_full_mesh_op_strategy(
+        mesh, op_schema, single_mesh_dim_strategies, input_index=num_outputs
+    )
+
+
 @register_op_strategy(
     [aten._log_softmax.default, aten._softmax.default, aten._safe_softmax.default],
     schema_info=RuntimeSchemaInfo(1),
@@ -919,13 +1159,17 @@ def nll_loss_backward_strategy(op_schema: OpSchema) -> OpStrategy:
                 generate_redistribute_costs(weight_strategy, weight_expected_spec)
             )
 
-        # total_weight should always be replicated
+        # total_weight is only used by the backward kernel for reduction='mean'.
+        # For reduction='sum' or 'none', it is unused, so no redistribution needed.
         total_weight_src_spec = total_weight_strategy.strategies[idx].output_spec
-        total_weight_expected_spec = DTensorSpec(
-            mesh=mesh,
-            placements=_replicate_dims_start_at(total_weight_src_spec.placements),
-            tensor_meta=total_weight_src_spec.tensor_meta,
-        )
+        if reduction == Reduction.MEAN.value:
+            total_weight_expected_spec = DTensorSpec(
+                mesh=mesh,
+                placements=_replicate_dims_start_at(total_weight_src_spec.placements),
+                tensor_meta=total_weight_src_spec.tensor_meta,
+            )
+        else:
+            total_weight_expected_spec = total_weight_src_spec
         op_args_target_specs.append(total_weight_expected_spec)
         redistribute_costs.append(
             generate_redistribute_costs(
@@ -945,358 +1189,145 @@ def nll_loss_backward_strategy(op_schema: OpSchema) -> OpStrategy:
     return grad_in_strategy
 
 
-def _common_norm_forward_strategy(
-    op_schema: OpSchema,
-    rms_norm: bool = False,
-) -> OpStrategy:
-    """Common forward strategy logic for layer_norm and rms_norm."""
-    mesh = op_schema.get_mesh_from_args()
-
-    if not rms_norm:
-        # layer_norm args: input, normalized_shape, weight, bias, eps
-        # for None weight and bias, their corresponding objects will
-        # be None as well. layer_norm_strategy returns one OpStrategy
-        # for the triple return values (out, mean, rstd).
-        if not len(op_schema.args_schema) == 5:
-            raise AssertionError(f"Expected 5 args, got {len(op_schema.args_schema)}")
-        (
-            input_strategy,
-            normalized_shape,
-            weight_strategy,
-            bias_strategy,
-            _,
-        ) = op_schema.args_schema
-    else:
-        # rms_norm args: input, normalized_shape, weight, eps
-        if not len(op_schema.args_schema) == 4:
-            raise AssertionError(f"Expected 4 args, got {len(op_schema.args_schema)}")
-        (
-            input_strategy,
-            normalized_shape,
-            weight_strategy,
-            _,
-        ) = op_schema.args_schema
-        bias_strategy = None
-
-    # the current norm implementation requires that all
-    # input DTensor's sharding must be in form of OpStrategy
-    if not isinstance(input_strategy, OpStrategy):
-        raise AssertionError(f"Expected OpStrategy, got {type(input_strategy)}")
-    if not isinstance(normalized_shape, (int, Sequence, torch.Size)):
-        raise AssertionError(
-            f"Expected int, Sequence, or torch.Size, got {type(normalized_shape)}"
-        )
-    normalized_size = normalize_to_torch_size(normalized_shape)
-
-    input_ndim = input_strategy.ndim
-    axis = input_ndim - len(normalized_size)
-
-    # we use OpStrategy because the output values (out, mean, rstd)
-    # should have the same placements
-    output_strategy = OpStrategy([])
-    for idx, input_placement_strategy in enumerate(input_strategy.strategies):
-        op_args_target_specs = []
-        redistribute_costs = []
-        input_src_spec = input_placement_strategy.output_spec
-
-        # for the input tensor, we replicate it on the inner dims if necessary
-        # TODO: we can avoid forcing the redistribution once we figure out
-        # how to decompose layer norm
-        input_target_spec = DTensorSpec(
-            mesh=mesh,
-            placements=_replicate_dims_start_at(input_src_spec.placements, axis),
-            tensor_meta=input_src_spec.tensor_meta,
-        )
-        op_args_target_specs.append(input_target_spec)
-        redistribute_costs.append(
-            generate_redistribute_costs(input_strategy, input_target_spec)
-        )
-
-        if weight_strategy is not None:
-            if not isinstance(weight_strategy, OpStrategy):
-                raise AssertionError(
-                    f"Expected OpStrategy, got {type(weight_strategy)}"
-                )
-            weight_src_spec = weight_strategy.strategies[idx].output_spec
-
-            # for the weight tensor, we replicate it on all dims if necessary
-            # TODO: we can avoid forcing the redistribution once we figure out
-            # how to decompose layer norm
-            weight_target_spec = DTensorSpec(
-                mesh=mesh,
-                placements=_replicate_dims_start_at(weight_src_spec.placements),
-                tensor_meta=weight_src_spec.tensor_meta,
-            )
-            op_args_target_specs.append(weight_target_spec)
-            redistribute_costs.append(
-                generate_redistribute_costs(weight_strategy, weight_target_spec)
-            )
-
-        if bias_strategy is not None:
-            if not isinstance(bias_strategy, OpStrategy):
-                raise AssertionError(f"Expected OpStrategy, got {type(bias_strategy)}")
-            bias_src_spec = bias_strategy.strategies[idx].output_spec
-
-            # for the bias tensor, we replicate it on all dims if necessary
-            # TODO: we can avoid forcing the redistribution once we figure out
-            # how to decompose layer norm
-            bias_target_spec = DTensorSpec(
-                mesh=mesh,
-                placements=_replicate_dims_start_at(bias_src_spec.placements),
-                tensor_meta=bias_src_spec.tensor_meta,
-            )
-            op_args_target_specs.append(bias_target_spec)
-            redistribute_costs.append(
-                generate_redistribute_costs(bias_strategy, bias_target_spec)
-            )
-
-        # the output spec is the same as input spec
-        output_target_spec = input_target_spec
-        output_strategy.strategies.append(
-            OpSpec(
-                output_specs=output_target_spec,
-                input_specs=op_args_target_specs,
-                redistribute_cost=redistribute_costs,
-            )
-        )
-
-    return output_strategy
-
-
-@register_op_strategy(
+@register_single_dim_strategy(
     [aten.native_layer_norm.default],
     schema_info=RuntimeSchemaInfo(1),
 )
-def layer_norm_strategy(op_schema: OpSchema) -> OpStrategy:
-    return _common_norm_forward_strategy(op_schema)
+def layer_norm_single_dim_strategy(
+    op: torch._ops.OpOverload,
+    args_schema: tuple[Any, ...],
+    kwargs_schema: dict[str, Any],
+) -> list[list[Placement | _ShardingPlaceholder]]:
+    input_meta = args_schema[0]
+    normalized_shape = args_schema[1]
+    weight_meta = args_schema[2]
+    bias_meta = args_schema[3]
+
+    axis = len(input_meta.shape) - len(normalize_to_torch_size(normalized_shape))
+
+    strategies: list[list[Placement | _ShardingPlaceholder]] = []
+    for dim in range(axis):
+        # [out, mean, rstd, input, weight?, bias?]
+        rule: list[Placement | _ShardingPlaceholder] = [
+            _ShardingPlaceholder(dim),  # out
+            _ShardingPlaceholder(dim),  # mean
+            _ShardingPlaceholder(dim),  # rstd
+            _ShardingPlaceholder(dim),  # input
+        ]
+        if weight_meta is not None:
+            rule.append(Replicate())
+        if bias_meta is not None:
+            rule.append(Replicate())
+        strategies.append(rule)
+    return strategies
 
 
-@register_op_strategy(
+@register_single_dim_strategy(
     [aten._fused_rms_norm.default],
     schema_info=RuntimeSchemaInfo(1),
 )
-def fused_rms_norm_strategy(op_schema: OpSchema) -> OpStrategy:
-    return _common_norm_forward_strategy(op_schema, rms_norm=True)
-
-
-def _common_norm_backward_strategy(
-    op_schema: OpSchema,
-    rms_norm: bool = False,
-) -> OpStrategy:
-    """Common backward strategy logic for layer_norm and rms_norm."""
-    # backward op does not need to validate the mesh since forward op has already done it
-    mesh = op_schema.get_mesh_from_args(validate=False)
-
-    if not rms_norm:
-        # layer_norm args: grad_out, input, normalized_shape, mean, rstd,
-        # weight, bias, output_mask. For None weight and bias, their
-        # corresponding objects will be None as well.
-        if not len(op_schema.args_schema) == 8:
-            raise AssertionError(f"Expected 8 args, got {len(op_schema.args_schema)}")
-        (
-            grad_out_strategy,
-            input_strategy,
-            normalized_shape,
-            mean_strategy,
-            rstd_strategy,
-            weight_strategy,
-            bias_strategy,
-            output_mask,
-        ) = op_schema.args_schema
-    else:
-        # rms_norm args: grad_out, input, normalized_shape, rstd,
-        if not len(op_schema.args_schema) == 6:
-            raise AssertionError(f"Expected 6 args, got {len(op_schema.args_schema)}")
-        (
-            grad_out_strategy,
-            input_strategy,
-            normalized_shape,
-            rstd_strategy,
-            weight_strategy,
-            output_mask,
-        ) = op_schema.args_schema
-        mean_strategy = None
-        bias_strategy = None
-
-    if not isinstance(grad_out_strategy, OpStrategy):
-        raise AssertionError(f"Expected OpStrategy, got {type(grad_out_strategy)}")
-    if not isinstance(input_strategy, OpStrategy):
-        raise AssertionError(f"Expected OpStrategy, got {type(input_strategy)}")
-    if not isinstance(rstd_strategy, OpStrategy):
-        raise AssertionError(f"Expected OpStrategy, got {type(rstd_strategy)}")
-    if mean_strategy is not None:
-        if not isinstance(mean_strategy, OpStrategy):
-            raise AssertionError(f"Expected OpStrategy, got {type(mean_strategy)}")
-
-    if not isinstance(normalized_shape, (int, Sequence, torch.Size)):
-        raise AssertionError(
-            f"Expected int, Sequence, or torch.Size, got {type(normalized_shape)}"
-        )
-    normalized_size = normalize_to_torch_size(normalized_shape)
-    input_ndim = input_strategy.ndim
-    axis = input_ndim - len(normalized_size)
-    outer_dims = list(range(axis))
-
-    if not rms_norm:
-        if not (isinstance(output_mask, list) and len(output_mask) == 3):
-            raise AssertionError(
-                f"Expected output_mask to be list of length 3, got {type(output_mask)} "
-                f"of length {len(output_mask) if isinstance(output_mask, list) else 'N/A'}"
-            )
-    else:
-        if not (isinstance(output_mask, list) and len(output_mask) == 2):
-            raise AssertionError(
-                f"Expected output_mask to be list of length 2, got {type(output_mask)} "
-                f"of length {len(output_mask) if isinstance(output_mask, list) else 'N/A'}"
-            )
-
-    # output tuple: (d_input, d_weight[, d_bias])
-    out_tuple_strategy = OpStrategy([])
-    for idx, input_placement_strategy in enumerate(input_strategy.strategies):
-        # args for OpSpec
-        output_specs_list: list[DTensorSpec | None] = []
-        input_specs_list: list[DTensorSpec] = []
-        redistribute_costs = []
-
-        input_src_spec = input_placement_strategy.output_spec
-        # arg: grad_out
-        # TODO: change the strategy to the following rule.
-        # d_input is basically a product of element-wise mul of
-        # grad_out, rstd, and normalized input, among which rstd
-        # and normalized input (x_hat) should have the same sharding
-        # placements, and grad_out's sharding is determined by the
-        # pointwise result of x_hat and weight/bias.
-        # TODO: now grad_out spec follows input spec. we may need
-        # to change it to apply a pointwise rule over grad_out,
-        # input, and weight.
-        grad_out_target_spec = DTensorSpec(
-            mesh=mesh,
-            placements=_replicate_dims_start_at(input_src_spec.placements, axis),
-            tensor_meta=input_src_spec.tensor_meta,
-        )
-        input_specs_list.append(grad_out_target_spec)
-        redistribute_costs.append(
-            generate_redistribute_costs(grad_out_strategy, grad_out_target_spec)
-        )
-        output_specs_list.append(grad_out_target_spec if output_mask[0] else None)
-
-        # arg: input
-        input_target_spec = DTensorSpec(
-            mesh=mesh,
-            placements=_replicate_dims_start_at(input_src_spec.placements, axis),
-            tensor_meta=input_src_spec.tensor_meta,
-        )
-        input_specs_list.append(input_target_spec)
-        redistribute_costs.append(
-            generate_redistribute_costs(input_strategy, input_target_spec)
-        )
-
-        # arg: mean
-        if not rms_norm:
-            if mean_strategy is None:
-                raise AssertionError("Expected mean_strategy to not be None")
-            mean_src_spec = mean_strategy.strategies[idx].output_spec
-            input_specs_list.append(mean_src_spec)
-            redistribute_costs.append([0.0 for _ in mean_strategy.strategies])
-
-        # arg: rstd
-        rstd_src_spec = rstd_strategy.strategies[idx].output_spec
-        input_specs_list.append(rstd_src_spec)
-        redistribute_costs.append([0.0 for _ in rstd_strategy.strategies])
-
-        def _add_target_input_spec(strategy) -> DTensorSpec:
-            # shared logic for setting the weight and bias target input specs
-            if not isinstance(strategy, OpStrategy):
-                raise AssertionError(f"Expected OpStrategy, got {type(strategy)}")
-            src_spec = strategy.strategies[idx].output_spec
-            # no need to redistribute since they should be replicated in forward pass
-            input_specs_list.append(src_spec)
-            redistribute_costs.append([0.0 for _ in strategy.strategies])
-            return src_spec
-
-        # arg: weight
-        # d_weight = sum(grad_out * (input - mean) / rstd, outer_dim, keepdim=False)
-        # For RMS norm, mean is 0, so it's just: sum(grad_out * input / rstd, outer_dim, keepdim=False)
-        if weight_strategy is not None:
-            weight_src_spec = _add_target_input_spec(weight_strategy)
-            # TODO: now d_weight spec follows input spec w/ a reduction.
-            # we may need to change to a pointwise rule over grad_out and
-            # input, then apply a reduction.
-            inp_placements = _replicate_dims_start_at(input_src_spec.placements, axis)
-            reduce_dims_map = _infer_reduce_dims_map(
-                outer_dims, input_src_spec.ndim, False
-            )
-            out_placements = map_placements_after_reduction(
-                inp_placements, outer_dims, reduce_dims_map, "sum"
-            )
-            weight_out_spec = DTensorSpec(
-                mesh=mesh,
-                placements=out_placements,
-                tensor_meta=weight_src_spec.tensor_meta,
-            )
-            output_specs_list.append(weight_out_spec if output_mask[1] else None)
-        else:
-            if not rms_norm:
-                error_msg = "output_mask[1] should not be `True` while weight argument is `None` in native_layer_norm_backward."
-            else:
-                error_msg = "output_mask[1] should not be `True` while weight argument is `None` in _fused_rms_norm_backward."
-            if output_mask[1] is not False:
-                raise AssertionError(error_msg)
-            output_specs_list.append(None)
-
-        # arg: bias
-        # d_bias = sum(grad_out, outer_dim, keepdim=False)
-        if not rms_norm:
-            if bias_strategy is not None:
-                bias_src_spec = _add_target_input_spec(bias_strategy)
-                # d_bias spec follows a reduction over grad_out
-                inp_placements = _replicate_dims_start_at(
-                    grad_out_target_spec.placements, axis
-                )
-                reduce_dims_map = _infer_reduce_dims_map(
-                    outer_dims, grad_out_target_spec.ndim, False
-                )
-                out_placements = map_placements_after_reduction(
-                    inp_placements, outer_dims, reduce_dims_map, "sum"
-                )
-                bias_out_spec = DTensorSpec(
-                    mesh=mesh,
-                    placements=out_placements,
-                    tensor_meta=bias_src_spec.tensor_meta,
-                )
-                output_specs_list.append(bias_out_spec if output_mask[2] else None)
-            else:
-                if output_mask[2] is not False:
-                    raise AssertionError(
-                        "output_mask[2] should not be `True` while bias argument is `None` in native_layer_norm_backward."
-                    )
-                output_specs_list.append(None)
-
-        out_tuple_strategy.strategies.append(
-            OpSpec(
-                output_specs=tuple(output_specs_list),
-                input_specs=input_specs_list,
-                redistribute_cost=redistribute_costs,
-            )
-        )
-
-    return out_tuple_strategy
+def rms_norm_single_dim_strategy(
+    op: torch._ops.OpOverload,
+    args_schema: tuple[Any, ...],
+    kwargs_schema: dict[str, Any],
+) -> list[list[Placement | _ShardingPlaceholder]]:
+    input_meta = args_schema[0]
+    normalized_shape = args_schema[1]
+    weight_meta = args_schema[2]
+
+    axis = len(input_meta.shape) - len(normalize_to_torch_size(normalized_shape))
+
+    strategies: list[list[Placement | _ShardingPlaceholder]] = []
+    for dim in range(axis):
+        # [out, rrms, input, weight?]
+        rule: list[Placement | _ShardingPlaceholder] = [
+            _ShardingPlaceholder(dim),  # out
+            _ShardingPlaceholder(dim),  # rrms
+            _ShardingPlaceholder(dim),  # input
+        ]
+        if weight_meta is not None:
+            rule.append(Replicate())
+        strategies.append(rule)
+    return strategies
 
 
-@register_op_strategy(
+@register_single_dim_strategy(
     [aten.native_layer_norm_backward.default],
     schema_info=RuntimeSchemaInfo(2),
 )
-def layer_norm_bwd_strategy(op_schema: OpSchema) -> OpStrategy:
-    return _common_norm_backward_strategy(op_schema)
+def layer_norm_bwd_single_dim_strategy(
+    op: torch._ops.OpOverload,
+    args_schema: tuple[Any, ...],
+    kwargs_schema: dict[str, Any],
+) -> list[list[Placement | _ShardingPlaceholder | None]]:
+    input_meta = args_schema[1]
+    normalized_shape = args_schema[2]
+    # mean = args_schema[3], rstd = args_schema[4]
+    weight_meta = args_schema[5]
+    bias_meta = args_schema[6]
+
+    axis = len(input_meta.shape) - len(normalize_to_torch_size(normalized_shape))
+
+    strategies: list[list[Placement | _ShardingPlaceholder | None]] = []
+    for dim in range(axis):
+        # outputs: [d_input, d_weight, d_bias] — always 3 per schema
+        # d_weight/d_bias use None when weight/bias are None
+        rule: list[Placement | _ShardingPlaceholder | None] = [
+            _ShardingPlaceholder(dim),  # d_input
+            Partial("sum") if weight_meta is not None else None,  # d_weight
+            Partial("sum") if bias_meta is not None else None,  # d_bias
+        ]
+        # inputs: [grad_out, input, mean, rstd, weight?, bias?]
+        rule.extend(
+            [
+                _ShardingPlaceholder(dim),  # grad_out
+                _ShardingPlaceholder(dim),  # input
+                _ShardingPlaceholder(dim),  # mean
+                _ShardingPlaceholder(dim),  # rstd
+            ]
+        )
+        if weight_meta is not None:
+            rule.append(Replicate())
+        if bias_meta is not None:
+            rule.append(Replicate())
+        strategies.append(rule)
 
+    return strategies
 
-@register_op_strategy(
+
+@register_single_dim_strategy(
     [aten._fused_rms_norm_backward.default],
     schema_info=RuntimeSchemaInfo(2),
 )
-def fused_rms_norm_bwd_strategy(op_schema: OpSchema) -> OpStrategy:
-    return _common_norm_backward_strategy(op_schema, rms_norm=True)
+def rms_norm_bwd_single_dim_strategy(
+    op: torch._ops.OpOverload,
+    args_schema: tuple[Any, ...],
+    kwargs_schema: dict[str, Any],
+) -> list[list[Placement | _ShardingPlaceholder | None]]:
+    input_meta = args_schema[1]
+    normalized_shape = args_schema[2]
+    # rstd = args_schema[3]
+    weight_meta = args_schema[4]
+
+    axis = len(input_meta.shape) - len(normalize_to_torch_size(normalized_shape))
+
+    strategies: list[list[Placement | _ShardingPlaceholder | None]] = []
+    for dim in range(axis):
+        # outputs: [d_input, d_weight] — always 2 per schema
+        # d_weight uses None when weight is None
+        # inputs: [grad_out, input, rstd, weight?]
+        rule: list[Placement | _ShardingPlaceholder | None] = [
+            _ShardingPlaceholder(dim),  # d_input
+            Partial("sum") if weight_meta is not None else None,  # d_weight
+            _ShardingPlaceholder(dim),  # grad_out
+            _ShardingPlaceholder(dim),  # input
+            _ShardingPlaceholder(dim),  # rstd
+        ]
+        if weight_meta is not None:
+            rule.append(Replicate())
+        strategies.append(rule)
+
+    return strategies
 
 
 def sort_strategy(op_schema: OpSchema, sort_dim: int) -> OpStrategy:
@@ -1420,3 +1451,299 @@ def logsumexp_strategy(op_schema: OpSchema) -> OpStrategy:
         keep_dim=keep_dim,
         reduction_linear=False,
     )
+
+
+_LINALG_NUM_PLACEMENTS = {
+    # 1 in 1 out
+    aten.cholesky.default: 2,
+    aten.cholesky_inverse.default: 2,
+    aten.linalg_matrix_exp.default: 2,
+    # 2 in 1 out
+    aten.cholesky_solve.default: 3,
+    aten.linalg_householder_product.default: 3,
+    aten.linalg_solve_triangular.default: 3,
+    # 3 in 1 out
+    aten.linalg_ldl_solve.default: 4,
+    aten.linalg_lu_solve.default: 4,
+    aten.ormqr.default: 4,
+    # 1 in 2 out
+    aten.geqrf.default: 3,
+    aten.linalg_cholesky_ex.default: 3,
+    aten.linalg_eig.default: 3,
+    aten.linalg_inv_ex.default: 3,
+    # 2 in 2 out
+    aten.triangular_solve.default: 4,
+    # 1 in 3 out
+    aten._linalg_det.default: 4,
+    aten.linalg_ldl_factor_ex.default: 4,
+    aten.linalg_lu.default: 4,
+    aten.linalg_lu_factor_ex.default: 4,
+    # 2 in 3 out
+    aten.lu_unpack.default: 5,
+    # 1 in 4 out
+    aten._linalg_slogdet.default: 5,
+    # 2 in 4 out
+    aten._linalg_solve_ex.default: 6,
+    # 1 in
+    aten._linalg_check_errors.default: 1,
+}
+
+
+def _linalg_batch_dim_strategies(
+    ndim: int, n_placements: int
+) -> list[list[Placement | _ShardingPlaceholder]]:
+    """Build single-dim strategies for linalg ops that operate on the last 1-2 dims.
+
+    Returns sharding on each batch dim (all dims except the last 2), with all
+    outputs and inputs sharded on the same dim.
+    """
+    strategies: list[list[Placement | _ShardingPlaceholder]] = []
+    for dim in range(ndim - 2):
+        strategies.append([_ShardingPlaceholder(dim)] * n_placements)
+    return strategies
+
+
+def _get_ndim(tensor_meta: Any) -> int:
+    if not isinstance(tensor_meta, TensorMeta):
+        raise AssertionError(f"Expected TensorMeta, got {type(tensor_meta)}")
+    return len(tensor_meta.shape)
+
+
+@register_single_dim_strategy(
+    [
+        aten.cholesky.default,
+        aten.cholesky_inverse.default,
+        aten.linalg_matrix_exp.default,
+        aten.cholesky_solve.default,
+        aten.linalg_householder_product.default,
+        aten.linalg_solve_triangular.default,
+        aten.linalg_ldl_solve.default,
+        aten.linalg_lu_solve.default,
+        aten.ormqr.default,
+        aten.geqrf.default,
+        aten.linalg_cholesky_ex.default,
+        aten.linalg_eig.default,
+        aten.linalg_inv_ex.default,
+        aten.triangular_solve.default,
+        aten._linalg_det.default,
+        aten.linalg_ldl_factor_ex.default,
+        aten.linalg_lu.default,
+        aten.linalg_lu_factor_ex.default,
+        aten.lu_unpack.default,
+        aten._linalg_slogdet.default,
+        aten._linalg_solve_ex.default,
+        aten._linalg_check_errors.default,
+    ],
+    schema_info=RuntimeSchemaInfo(1),
+)
+def linalg_batch_dim_strategy(
+    op: torch._ops.OpOverload,
+    args_schema: tuple[Any, ...],
+    kwargs_schema: dict[str, Any],
+) -> list[list[Placement | _ShardingPlaceholder]]:
+    ndim = _get_ndim(args_schema[0])
+    if op not in _LINALG_NUM_PLACEMENTS:
+        raise AssertionError(f"Expected op in _LINALG_NUM_PLACEMENTS, got {op}")
+
+    n_placements = _LINALG_NUM_PLACEMENTS[op]
+    strategies = _linalg_batch_dim_strategies(ndim, n_placements=n_placements)
+
+    if op == aten.linalg_solve_triangular.default:
+        # solve_triangular(A, B) -> result: linear in B
+        strategies.append([Partial(), Replicate(), Partial()])
+        strategies.append([Partial("avg"), Replicate(), Partial("avg")])
+        # A replicated, B sharded on batch dims (B may have more batch dims than A)
+        ndim_b = _get_ndim(args_schema[1])
+        for dim in range(ndim_b - 2):
+            strategies.append(
+                [_ShardingPlaceholder(dim), Replicate(), _ShardingPlaceholder(dim)]
+            )
+    elif op == aten.cholesky_solve.default:
+        # cholesky_solve(B, A) -> result  (B is arg0)
+        strategies.append([Partial(), Partial(), Replicate()])
+    elif op == aten.linalg_lu_solve.default:
+        # linalg_lu_solve(LU, pivots, B) -> result
+        strategies.append([Partial(), Replicate(), Replicate(), Partial()])
+    elif op == aten.linalg_ldl_solve.default:
+        # linalg_ldl_solve(LD, pivots, B) -> result
+        strategies.append([Partial(), Replicate(), Replicate(), Partial()])
+    elif op == aten.ormqr.default:
+        # ormqr(a, tau, C) -> result  (linear in C)
+        strategies.append([Partial(), Replicate(), Replicate(), Partial()])
+    elif op == aten._linalg_solve_ex.default:
+        # _linalg_solve_ex(A, B) -> (result, LU, pivots, info)
+        strategies.append(
+            [Partial(), Replicate(), Replicate(), Replicate(), Replicate(), Partial()]
+        )
+
+    return strategies
+
+
+# linalg_pinv has optional tensor kwargs atol, rtol (scalar tensors when present).
+# Schema: (Tensor self, *, Tensor? atol=None, Tensor? rtol=None, bool hermitian) -> Tensor
+# When atol/rtol are None, num_inputs=1; when present, they add to num_inputs.
+@register_single_dim_strategy(
+    [aten.linalg_pinv.atol_rtol_tensor],
+    schema_info=RuntimeSchemaInfo(1),
+)
+def linalg_pinv_strategy(
+    op: torch._ops.OpOverload,
+    args_schema: tuple[Any, ...],
+    kwargs_schema: dict[str, Any],
+) -> list[list[Placement | _ShardingPlaceholder]]:
+    ndim = _get_ndim(args_schema[0])
+    # Count optional tensor kwargs that are actually present
+    extra_tensors = sum(
+        isinstance(kwargs_schema.get(k), TensorMeta) for k in ("atol", "rtol")
+    )
+    strategies: list[list[Placement | _ShardingPlaceholder]] = []
+    for dim in range(ndim - 2):
+        s: list[Placement | _ShardingPlaceholder] = [
+            _ShardingPlaceholder(dim),
+            _ShardingPlaceholder(dim),
+        ]
+        # atol, rtol are scalar tensors — always Replicate
+        s.extend([Replicate()] * extra_tensors)
+        strategies.append(s)
+    return strategies
+
+
+# linalg_cross is pointwise on every dim except the cross-product dim (which
+# must be size 3).  Shard on any other dim.
+@register_single_dim_strategy(
+    [aten.linalg_cross.default],
+    schema_info=RuntimeSchemaInfo(1),
+)
+def linalg_cross_strategy(
+    op: torch._ops.OpOverload,
+    args_schema: tuple[Any, ...],
+    kwargs_schema: dict[str, Any],
+) -> list[list[Placement | _ShardingPlaceholder]]:
+    ndim = _get_ndim(args_schema[0])
+    cross_dim = kwargs_schema.get("dim", -1) % ndim
+    strategies: list[list[Placement | _ShardingPlaceholder]] = []
+    for dim in range(ndim):
+        if dim == cross_dim:
+            continue
+        strategies.append([_ShardingPlaceholder(dim)] * 3)
+    return strategies
+
+
+# ---------------------------------------------------------------------------
+# Interpolation / upsample / pooling ops
+#
+# These ops operate on spatial dims and are safely shardable on batch (dim 0)
+# and channel (dim 1). grid_sampler is batch-only because the grid tensor has
+# no channel dimension.
+# ---------------------------------------------------------------------------
+
+
+@register_single_dim_strategy(
+    [
+        aten.upsample_nearest1d.default,
+        aten.upsample_nearest2d.default,
+        aten.upsample_nearest3d.default,
+        aten._upsample_nearest_exact1d.default,
+        aten._upsample_nearest_exact2d.default,
+        aten._upsample_nearest_exact3d.default,
+        aten._upsample_bilinear2d_aa.default,
+        aten.upsample_bicubic2d.default,
+        aten.upsample_bilinear2d.default,
+        aten.upsample_linear1d.default,
+        aten.upsample_trilinear3d.default,
+    ],
+    schema_info=RuntimeSchemaInfo(1),
+)
+def interp_upsample_1out_1in_strategy(
+    op: torch._ops.OpOverload,
+    args_schema: tuple[Any, ...],
+    kwargs_schema: dict[str, Any],
+) -> list[list[Placement | _ShardingPlaceholder]]:
+    # 1 output + 1 input = 2 placements; shard on batch (0) and channel (1)
+    # Upsample is a linear transformation so Partial(sum/avg) is valid.
+    return [
+        [_ShardingPlaceholder(0)] * 2,
+        [_ShardingPlaceholder(1)] * 2,
+        [Partial("sum"), Partial("sum")],
+        [Partial("avg"), Partial("avg")],
+    ]
+
+
+@register_single_dim_strategy(
+    [
+        aten.max_unpool2d.default,
+        aten.max_unpool3d.default,
+        aten._adaptive_avg_pool2d_backward.default,
+    ],
+    schema_info=RuntimeSchemaInfo(1),
+)
+def interp_pool_1out_2in_strategy(
+    op: torch._ops.OpOverload,
+    args_schema: tuple[Any, ...],
+    kwargs_schema: dict[str, Any],
+) -> list[list[Placement | _ShardingPlaceholder]]:
+    # 1 output + 2 inputs = 3 placements; shard on batch (0) and channel (1)
+    return [
+        [_ShardingPlaceholder(0)] * 3,
+        [_ShardingPlaceholder(1)] * 3,
+    ]
+
+
+@register_single_dim_strategy(
+    [aten.max_pool2d_with_indices_backward.default],
+    schema_info=RuntimeSchemaInfo(1),
+)
+def pool_backward_strategy(
+    op: torch._ops.OpOverload,
+    args_schema: tuple[Any, ...],
+    kwargs_schema: dict[str, Any],
+) -> list[list[Placement | _ShardingPlaceholder]]:
+    # max_pool2d_with_indices_backward(grad_output, self, ..., indices) -> grad_input
+    # 1 output + 3 tensor inputs = 4 placements
+    # Order: [output, grad_output, self, indices]
+    input_meta = cast(TensorMeta, args_schema[0])
+    strategies: list[list[Placement | _ShardingPlaceholder]] = [
+        [_ShardingPlaceholder(0)] * 4,
+    ]
+    if len(input_meta.shape) >= 4:  # batched: (N, C, H, W)
+        strategies.append([_ShardingPlaceholder(1)] * 4)
+    # The backward is linear in grad_output, so P(sum/avg) pass through.
+    # indices must be replicated (integer positions, not reducible).
+    # self is only used for shape, so replicate it too.
+    r = Replicate()
+    for reduce_op in ("sum", "avg"):
+        p = Partial(reduce_op)
+        strategies.append([p, p, r, r])
+    return strategies
+
+
+@register_single_dim_strategy(
+    [aten.grid_sampler_2d.default, aten.grid_sampler_3d.default],
+    schema_info=RuntimeSchemaInfo(1),
+)
+def grid_sampler_strategy(
+    op: torch._ops.OpOverload,
+    args_schema: tuple[Any, ...],
+    kwargs_schema: dict[str, Any],
+) -> list[list[Placement | _ShardingPlaceholder]]:
+    # grid_sampler_{2,3}d(input[N,C,...], grid[N,...,{2,3}]) -> output[N,C,...]
+    # grid has no channel dim, so only batch sharding applies to both inputs.
+    # Linear in input: P(sum/avg) on input with replicated grid is valid.
+    return [
+        [_ShardingPlaceholder(0)] * 3,
+        [Partial("sum"), Partial("sum"), Replicate()],
+        [Partial("avg"), Partial("avg"), Replicate()],
+    ]
+
+
+@register_single_dim_strategy(
+    [aten.grid_sampler_2d_backward.default, aten.grid_sampler_3d_backward.default],
+    schema_info=RuntimeSchemaInfo(1),
+)
+def grid_sampler_backward_strategy(
+    op: torch._ops.OpOverload,
+    args_schema: tuple[Any, ...],
+    kwargs_schema: dict[str, Any],
+) -> list[list[Placement | _ShardingPlaceholder]]:
+    # grid_sampler_{2,3}d_backward: 2 outputs (grad_input, grad_grid) + 3 inputs = 5 placements, batch-only
+    return [[_ShardingPlaceholder(0)] * 5]
diff --git a/torch/distributed/tensor/_ops/_matrix_ops.py b/torch/distributed/tensor/_ops/_matrix_ops.py
index 4f40629939926..b1bda0d15fdc2 100644
--- a/torch/distributed/tensor/_ops/_matrix_ops.py
+++ b/torch/distributed/tensor/_ops/_matrix_ops.py
@@ -36,11 +36,13 @@
     compute_local_stride,
 )
 from torch.distributed.tensor.placement_types import (
+    _StridedShard,
     Partial,
     Placement,
     Replicate,
     Shard,
 )
+from torch.fx.experimental.symbolic_shapes import guard_or_false
 
 
 aten = torch.ops.aten
@@ -61,10 +63,16 @@ def transpose_strategy(op_schema: OpSchema) -> OpStrategy:
         if ndim <= 1:
             output_placements = list(input_spec.placements)
         else:
-            output_placements = [
-                Shard(1 - p.dim) if isinstance(p, Shard) else p
-                for p in input_spec.placements
-            ]
+            output_placements: list[Placement] = []
+            for p in input_spec.placements:
+                if isinstance(p, _StridedShard):
+                    output_placements.append(
+                        _StridedShard(1 - p.dim, split_factor=p.split_factor)
+                    )
+                elif isinstance(p, Shard):
+                    output_placements.append(Shard(1 - p.dim))
+                else:
+                    output_placements.append(p)
         transpose_strategy = OpSpec(
             output_specs=DTensorSpec(
                 mesh=input_strategy.mesh,
@@ -176,83 +184,42 @@ def _addmm_like_strategy(
     return mm_strategy
 
 
-def _scaled_mm_like_strategy(
-    mm_equation: str, mesh: DeviceMesh, op_schema: OpSchema
-) -> OpStrategy:
-    (
-        self_strategy,
-        mat2_strategy,
-        scale_self_strategy,
-        scale_mat2_strategy,
-        bias_strategy,
-        scale_result_strategy,
-        *_,
-    ) = op_schema.args_schema
-    if not isinstance(self_strategy, OpStrategy):
-        raise AssertionError(f"Expected OpStrategy, got {type(self_strategy)}")
-    if not isinstance(mat2_strategy, OpStrategy):
-        raise AssertionError(f"Expected OpStrategy, got {type(mat2_strategy)}")
-    if not isinstance(scale_self_strategy, OpStrategy):
-        raise AssertionError(f"Expected OpStrategy, got {type(scale_self_strategy)}")
-    if not isinstance(scale_mat2_strategy, OpStrategy):
-        raise AssertionError(f"Expected OpStrategy, got {type(scale_mat2_strategy)}")
-    # TODO: add support for these later
-    if bias_strategy is not None:
-        raise AssertionError("_scaled_mm on DTensors doesn't support bias")
-    if scale_result_strategy is not None:
-        raise AssertionError("_scaled_mm on DTensors doesn't support scale_result")
-    # generate all possible strategies for mm
-    mm_strategy = gen_einsum_strategies(mm_equation, mesh)
-    # filter out invalid strategies and associate costs
-    strategies = mm_strategy.strategies
-    filtered_strategies = []
-    for strtg in strategies:
-        if strtg.input_specs is None:
-            raise AssertionError(
-                f"Expected input_specs to be not None, got {strtg.input_specs}"
-            )
-        self_spec = strtg.input_specs[0]
-        mat2_spec = strtg.input_specs[1]
-        # propagate the operands' specs to their scales, except for tensor-wise
-        # scaling which can have any numbers of dims (legacy...), hence sharding
-        # dims won't map. for tensor-wise, anyways, we can only do replication.
-        scale_self_spec = (
-            DTensorSpec(self_spec.mesh, (Replicate(),))
-            if prod(scale_self_strategy.shape) == 1
-            else self_spec
-        )
-        scale_mat2_spec = (
-            DTensorSpec(mat2_spec.mesh, (Replicate(),))
-            if prod(scale_mat2_strategy.shape) == 1
-            else mat2_spec
-        )
-        strtg.input_specs = list(strtg.input_specs) + [scale_self_spec, scale_mat2_spec]
-        if (
-            is_tensor_shardable(
-                self_strategy.shape, self_spec, allow_unbacked_sharding=True
-            )
-            and is_tensor_shardable(
-                mat2_strategy.shape, mat2_spec, allow_unbacked_sharding=True
-            )
-            and is_tensor_shardable(
-                scale_self_strategy.shape, scale_self_spec, allow_unbacked_sharding=True
-            )
-            and is_tensor_shardable(
-                scale_mat2_strategy.shape, scale_mat2_spec, allow_unbacked_sharding=True
-            )
-        ):
-            redistribute_cost = [
-                generate_redistribute_costs(self_strategy, self_spec),
-                generate_redistribute_costs(mat2_strategy, mat2_spec),
-                generate_redistribute_costs(scale_self_strategy, scale_self_spec),
-                generate_redistribute_costs(scale_mat2_strategy, scale_mat2_spec),
-            ]
-            strtg.redistribute_cost = redistribute_cost
-            filtered_strategies.append(strtg)
+def _scaled_mm_scale_placement(
+    data_placement: Placement | _ShardingPlaceholder,
+    scale_shape: torch.Size,
+    contracting_dim: int,
+) -> Placement | _ShardingPlaceholder | None:
+    """
+    Derive scale placement from data operand placement for _scaled_mm.
 
-    mm_strategy.strategies = filtered_strategies
+    Handles three cases:
 
-    return mm_strategy
+    1. Tensor-wise scale (single element): always Replicate.
+    2. 2D (or higher) scale, e.g. row-wise [M,1]: copy data placement directly.
+    3. 1D blockwise scale, e.g. MX format [M*K/block_size]: map
+       non-contracting shard to Shard(0)/_ShardingPlaceholder(0), and reject
+       contracting-dim shards (returns None).
+    """
+    if prod(scale_shape) == 1:
+        return Replicate()
+
+    if len(scale_shape) != 1:
+        return data_placement
+
+    # 1D blockwise scale: Shard(>=1) is invalid on a 1D tensor, so we need
+    # to map the data operand's placement to a valid 1D placement.
+    if isinstance(data_placement, _ShardingPlaceholder):
+        if data_placement.dim == contracting_dim:
+            return None
+        return _ShardingPlaceholder(0)
+    # NOTE: isinstance(_, Shard) does not match _StridedShard; see _is_shard_like().
+    elif isinstance(data_placement, Shard):
+        if data_placement.dim == contracting_dim:
+            return None
+        return Shard(0)
+    elif isinstance(data_placement, (Replicate, Partial)):
+        return Replicate()
+    return data_placement
 
 
 @register_op_strategy(aten.dot.default)
@@ -296,7 +263,11 @@ def gen_single_dim_einsum_strategies(
        independently, one input can remain Partial while others are Replicate,
        producing a Partial output.
 
-    4. Bias input (optional): If bias_shape is provided, a bias placement
+    4. Batch-dimension linearity (all-Partial): When all dims are batch dims
+       (no contracting or free dims), the operation is element-wise and linear
+       in all inputs simultaneously, so all inputs can be Partial.
+
+    5. Bias input (optional): If bias_shape is provided, a bias placement
        is inserted after the output placement. The bias placement is derived from
        the output placement, accounting for broadcast semantics (based on ndim
        difference between output and bias). This is used for addmm-like ops
@@ -415,6 +386,20 @@ def _maybe_add_bias(
             _maybe_add_bias([output_placement, Replicate(), Partial(reduce_op)])
         )
 
+    # Batch-dimension linearity: when the einsum has no contracting dims and
+    # no free dims (all dims are batch dims), the operation is element-wise
+    # and linear in all inputs simultaneously. Add all-Partial strategies.
+    if (
+        not edims.contracting_dims
+        and not edims.lhs_out_only_dims
+        and not edims.rhs_out_only_dims
+    ):
+        for reduce_op in Partial.LINEAR_REDUCE_OPS:
+            linearity_placements: list[Placement | _ShardingPlaceholder] = [
+                Partial(reduce_op)
+            ] + [Partial(reduce_op) for _ in input_dims]
+            strategies_over_one_mesh_dim.append(_maybe_add_bias(linearity_placements))
+
     return strategies_over_one_mesh_dim
 
 
@@ -436,7 +421,8 @@ def addmm_single_dim_strategy(
     op: OpOverload, args_schema: ArgsType, kwargs_schema: KwargsType
 ) -> list[list[Placement | _ShardingPlaceholder]]:
     bias_meta = args_schema[0]
-    assert isinstance(bias_meta, TensorMeta)
+    if not isinstance(bias_meta, TensorMeta):
+        raise AssertionError
     return gen_single_dim_einsum_strategies("mk,kn->mn", bias_shape=bias_meta.shape)
 
 
@@ -464,14 +450,41 @@ def baddbmm_single_dim_strategy(
     op: OpOverload, args_schema: ArgsType, kwargs_schema: KwargsType
 ) -> list[list[Placement | _ShardingPlaceholder]]:
     bias_meta = args_schema[0]
-    assert isinstance(bias_meta, TensorMeta)
+    if not isinstance(bias_meta, TensorMeta):
+        raise AssertionError
     return gen_single_dim_einsum_strategies("bmk,bkn->bmn", bias_shape=bias_meta.shape)
 
 
-@register_op_strategy(aten._scaled_mm.default)
-def scaled_mm_strategy(op_schema: OpSchema) -> OpStrategy:
-    mesh = op_schema.get_mesh_from_args()
-    return _scaled_mm_like_strategy("mk,kn->mn", mesh, op_schema)
+@register_single_dim_strategy(aten._scaled_mm.default, allow_unbacked_sharding=True)
+def scaled_mm_single_dim_strategy(
+    op: OpOverload, args_schema: ArgsType, kwargs_schema: KwargsType
+) -> list[list[Placement | _ShardingPlaceholder]]:
+    scale_self_meta = args_schema[2]
+    scale_mat2_meta = args_schema[3]
+    if not isinstance(scale_self_meta, TensorMeta):
+        raise AssertionError
+    if not isinstance(scale_mat2_meta, TensorMeta):
+        raise AssertionError
+    if args_schema[4] is not None:
+        raise AssertionError("_scaled_mm on DTensors doesn't support bias")
+    if args_schema[5] is not None:
+        raise AssertionError("_scaled_mm on DTensors doesn't support scale_result")
+
+    # "mk,kn->mn": self_contracting_dim=1, mat2_contracting_dim=0
+    base_strategies = gen_single_dim_einsum_strategies("mk,kn->mn")
+    result = []
+    for strat in base_strategies:
+        # strat is [output, self, mat2]; derive scale placements
+        scale_self_p = _scaled_mm_scale_placement(
+            strat[1], scale_self_meta.shape, contracting_dim=1
+        )
+        scale_mat2_p = _scaled_mm_scale_placement(
+            strat[2], scale_mat2_meta.shape, contracting_dim=0
+        )
+        if scale_self_p is None or scale_mat2_p is None:
+            continue
+        result.append(strat + [scale_self_p, scale_mat2_p])
+    return result
 
 
 def _scaled_dot_product_flash_attention_base_strategies(
@@ -655,23 +668,52 @@ def scaled_dot_product_flash_attention_backward_strategy(
     )
 
 
-@register_op_strategy(aten.constant_pad_nd.default)
-def constant_pad_nd_strategy(op_schema: OpSchema) -> OpStrategy:
-    mesh = op_schema.get_mesh_from_args(validate=False)
+@register_single_dim_strategy(
+    aten.constant_pad_nd.default, schema_info=RuntimeSchemaInfo(1)
+)
+def constant_pad_nd_single_dim_strategy(
+    op: OpOverload, args_schema: ArgsType, kwargs_schema: KwargsType
+) -> list[list[Placement | _ShardingPlaceholder]]:
+    # Allow sharding on non-padded dimensions; ban sharding on dims
+    # that have non-zero padding (where the pad value must be inserted).
+    input_meta = args_schema[0]
+    if not isinstance(input_meta, TensorMeta):
+        raise AssertionError(f"Expected TensorMeta, got {type(input_meta)}")
+    ndim = len(input_meta.shape)
+    pad = args_schema[1]
+    if not isinstance(pad, (list, tuple)):
+        raise AssertionError(f"Expected list or tuple, got {type(pad)}")
+
+    # pad is [dim_{n-1}_left, dim_{n-1}_right, dim_{n-2}_left, ...] from
+    # the last dim backwards. Determine which dims have non-zero padding.
+    padded_dims = set()
+    for i in range(len(pad) // 2):
+        if not (
+            guard_or_false(pad[i * 2] == 0) and guard_or_false(pad[i * 2 + 1] == 0)
+        ):
+            padded_dims.add(ndim - 1 - i)
+
+    # Shard on any non-padded dim: output and input share the same placement.
+    # All-Replicate is added automatically by the framework.
+    strategies: list[list[Placement | _ShardingPlaceholder]] = []
+    for dim in range(ndim):
+        if dim not in padded_dims:
+            strategies.append([_ShardingPlaceholder(dim), _ShardingPlaceholder(dim)])
+
+    # Partial rules: at padded positions every rank writes the same constant v,
+    # so reduce(v, v, ..., v) = v for avg/max/min (idempotent). P(sum) only
+    # works when v=0 since sum(v, ..., v) = N*v != v otherwise.
+    # When all pad amounts are zero the op is a no-op, so all reduce ops hold.
+    value = args_schema[2] if len(args_schema) > 2 else 0
+    no_padding = all(guard_or_false(pad[i] == 0) for i in range(len(pad)))
+    if no_padding or guard_or_false(value == 0):
+        reduce_ops = ("sum", "avg", "max", "min")
+    else:
+        reduce_ops = ("avg", "max", "min")
+    for reduce_op in reduce_ops:
+        strategies.append([Partial(reduce_op), Partial(reduce_op)])
 
-    # TODO(d4l3k); implement a more correct strategy for constant_pad_nd
-    return OpStrategy(
-        [
-            OpSpec(
-                output_specs=DTensorSpec(mesh, (Replicate(),)),
-                input_specs=(
-                    DTensorSpec(mesh, (Replicate(),)),
-                    DTensorSpec(mesh, (Replicate(),)),
-                ),
-                redistribute_cost=[[1]],
-            )
-        ]
-    )
+    return strategies
 
 
 def _scaled_dot_product_efficient_attention_base_strategies(
@@ -1243,10 +1285,10 @@ def local_meta(spec: OpSpec, placements: tuple[Placement, ...]) -> TensorMeta:
                     f"Expected TensorMeta, got {type(spec.output_specs.tensor_meta)}"
                 )
             meta: TensorMeta = spec.output_specs.tensor_meta
-            local_stride = compute_local_stride(meta.stride, mesh, placements)
             local_shape, _ = compute_local_shape_and_global_offset(
                 meta.shape, mesh, placements, skip_offset=True
             )
+            local_stride = compute_local_stride(meta.stride, local_shape)
             return TensorMeta(torch.Size(local_shape), local_stride, meta.dtype)
 
         # pyrefly: ignore [missing-attribute]
diff --git a/torch/distributed/tensor/_ops/_pointwise_ops.py b/torch/distributed/tensor/_ops/_pointwise_ops.py
index 6bed06a617402..dfc253eff8e86 100644
--- a/torch/distributed/tensor/_ops/_pointwise_ops.py
+++ b/torch/distributed/tensor/_ops/_pointwise_ops.py
@@ -1,845 +1,517 @@
 # Copyright (c) Meta Platforms, Inc. and affiliates
-import functools
-from collections.abc import Callable, Sequence
-from typing import cast, Optional
+from collections.abc import Callable
 
 import torch
 from torch._ops import OpOverload
-from torch.distributed.tensor._dtensor_spec import DTensorSpec, TensorMeta
-from torch.distributed.tensor._op_schema import (
-    ArgsType,
-    KwargsType,
-    OpSchema,
-    OpSpec,
-    OpStrategy,
-    RuntimeSchemaInfo,
-    StrategyType,
-    TupleStrategy,
+from torch.distributed.tensor._dtensor_spec import TensorMeta
+from torch.distributed.tensor._op_schema import ArgsType, KwargsType, RuntimeSchemaInfo
+from torch.distributed.tensor._ops.single_dim_strategy import (
+    _ShardingPlaceholder,
+    register_single_dim_strategy,
 )
-from torch.distributed.tensor._ops._math_ops import _NormPartial
-from torch.distributed.tensor._ops.single_dim_strategy import _ShardingPlaceholder
-from torch.distributed.tensor._ops.utils import (
-    generate_redistribute_costs,
-    infer_broadcast_dims_map,
-    map_placements_after_broadcast,
-    normalize_dim,
-    register_op_strategy,
-)
-from torch.distributed.tensor.placement_types import (
-    _StridedShard,
-    Partial,
-    Placement,
-    Replicate,
-    Shard,
-)
-from torch.types import _Number
-from torch.utils._typing_utils import not_none
+from torch.distributed.tensor._ops.utils import infer_broadcast_dims_map
+from torch.distributed.tensor.placement_types import Partial, Placement, Replicate
 
 
 aten = torch.ops.aten
-# leave the remaining pointwise_ops list here for convenience,
-# Below ops are some pointwise ops that are yet to be supported,
-# they might not be a complete list.
-# pointwise_ops = [
-#     "fake_quantize_per_channel_affine",
-#     "fake_quantize_per_tensor_affine",
-#     "floor_divide",  # floor_divide is deprecated
-#     "frexp",  # multiple output pointwise op, need to add support
-#     "gradient",  #  need investigation on this op
-#     "imag",  # complex data type only
-#     "quantized_batch_norm",
-#     "quantized_max_pool1d",
-#     "quantized_max_pool2d",
-#     "real",  # complex data type only
-# ]
+prims = torch.ops.prims
 
+# Linear pointwise ops, split by linearity type.
+unary_linear_ops = [aten.to.dtype]
 
-pointwise_ops = [
-    # please keep the entries below alphabetically sorted
-    aten.__ilshift__.Scalar,
-    aten.__ilshift__.Tensor,
-    aten.__irshift__.Scalar,
-    aten.__irshift__.Tensor,
-    aten.__lshift__.Scalar,
-    aten.__lshift__.Tensor,
-    aten.__rshift__.Scalar,
-    aten.__rshift__.Tensor,
-    aten._conj.default,
-    aten.abs.default,
-    aten.abs.out,
-    aten.abs_.default,
-    aten.acos.default,
-    aten.acos.out,
-    aten.acos_.default,
-    aten.acosh.default,
-    aten.acosh.out,
-    aten.acosh_.default,
-    aten.add.Scalar,
+
+def _common_pointwise_single_dim_strategy(
+    partial_extra_rules: list[list[Placement | _ShardingPlaceholder]] | None = None,
+) -> Callable[
+    [OpOverload, ArgsType, KwargsType], list[list[Placement | _ShardingPlaceholder]]
+]:
+    """Factory for single-dim strategies that add partial placement rules.
+
+    Returns strategies shaped [output, *args] only.  Tensor kwarg placements
+    (e.g. ``out``, ``lr``) are appended by the wrapper in
+    ``_register_single_dim_pointwise``.
+    """
+
+    def strategy(
+        op: OpOverload,
+        args_schema: ArgsType,
+        kwargs_schema: KwargsType,
+    ) -> list[list[Placement | _ShardingPlaceholder]]:
+        tensor_arg_metas: list[TensorMeta] = [
+            arg for arg in args_schema if isinstance(arg, TensorMeta)
+        ]
+        common_shape = torch.broadcast_shapes(
+            *[arg.shape for arg in args_schema if isinstance(arg, TensorMeta)]
+        )
+        # For multi-output ops (e.g. frexp), all outputs share the same
+        # pointwise sharding, so replicate the output placement.
+        num_outputs = sum(1 for r in op._schema.returns if "Tensor" in str(r.type))
+        placements: list[list[Placement | _ShardingPlaceholder]] = []
+        for i in range(len(common_shape)):
+            shard_placements: list[Placement | _ShardingPlaceholder] = [
+                _ShardingPlaceholder(i)
+            ] * num_outputs
+            for arg in tensor_arg_metas:
+                common_dim_to_arg_dim = infer_broadcast_dims_map(
+                    common_shape, arg.shape
+                )
+                # If the output shard dim maps to an input dim, shard that
+                # input dim; otherwise it was broadcast, so replicate.
+                if common_dim_to_arg_dim[i] >= 0:
+                    shard_placements.append(
+                        _ShardingPlaceholder(common_dim_to_arg_dim[i])
+                    )
+                else:
+                    shard_placements.append(Replicate())
+            placements.append(shard_placements)
+        if partial_extra_rules:
+            n_tensors = len(tensor_arg_metas)
+            expected_len = num_outputs + n_tensors
+            for rule in partial_extra_rules:
+                # Filter rather than assert: some ops (e.g. mul.Tensor) mix
+                # unary rules (len 2, for scalar promotion) and binary rules
+                # (len 3, for tensor-tensor), so mismatched lengths are expected.
+                # see _MUL_RULES to see how _UNARY_LINEAR_RULES handles the
+                # scalar promotion case
+                if len(rule) == expected_len:
+                    placements.append(rule)
+        return placements
+
+    return strategy
+
+
+def _is_list_op(op: OpOverload) -> bool:
+    """Returns True if op is a foreach, amp_foreach, or fused op."""
+    name = op.name()
+    return name.startswith(("aten::_foreach_", "aten::_amp_foreach_", "aten::_fused_"))
+
+
+# The state_steps arg of fused adam / adamw is a Replicate scalar tensor, which will be put on
+# the compute_mesh of an op across all parameter groups, even when not all parameter groups
+# are on the same device mesh. This idx will help avoid hitting exceptions or unnecessary
+# redistribute during sharding propagation.
+_FUSED_OP_SCALAR_IDX = 5
+
+# Ops registered with extra Partial rules; populated by _register_single_dim_pointwise
+# when partial_extra_rules is not None, to avoid double-registration from tag discovery.
+_specially_registered_ops: set[OpOverload] = set()
+
+
+def _register_single_dim_pointwise(
+    op: OpOverload,
+    partial_extra_rules: list[list[Placement]] | None = None,
+    static_argnum: int = 0,
+) -> None:
+    if partial_extra_rules is not None:
+        _specially_registered_ops.add(op)
+    inner_fn = _common_pointwise_single_dim_strategy(
+        partial_extra_rules=partial_extra_rules  # pyrefly: ignore[bad-argument-type]
+    )
+
+    # Wrap to append tensor kwarg placements in schema declaration order.
+    # out = output placement (s[0]); everything else (e.g. lr) = Replicate.
+    # TODO: move kwargs handling upstream if this works
+    def strategy_fn(
+        op: OpOverload,
+        args: ArgsType,
+        kwargs: KwargsType,
+        _fn: Callable = inner_fn,
+    ) -> list[list[Placement | _ShardingPlaceholder]]:
+        strategies = _fn(op, args, kwargs)
+        kw_names = [k for k, v in kwargs.items() if isinstance(v, TensorMeta)]
+        if not kw_names:
+            return strategies
+        return [
+            s + [s[0] if name == "out" else Replicate() for name in kw_names]
+            for s in strategies
+        ]
+
+    if _is_list_op(op):
+        schema_info = RuntimeSchemaInfo(needs_pytree=True)
+    else:
+        schema_info = RuntimeSchemaInfo(static_argnum, static_kwargkey=["out"])
+    # Fused ops (e.g. _fused_adam_) have state_steps on a potentially different
+    # mesh; see the note in expand_to_full_mesh_op_strategy for details.
+    different_mesh_args: list[int] | None = None
+    if op.name().startswith("aten::_fused_"):
+        different_mesh_args = [_FUSED_OP_SCALAR_IDX]
+    register_single_dim_strategy(
+        op,
+        schema_info=schema_info,
+        allow_uneven_sharding=True,
+        allow_unbacked_sharding=True,
+        different_mesh_args=different_mesh_args,
+    )(strategy_fn)
+
+
+_UNARY_LINEAR_RULES: list[list[Placement]] = [
+    [Partial("sum"), Partial("sum")],
+    [Partial("avg"), Partial("avg")],
+]
+
+binary_additive_ops = [
+    aten.add.Tensor,
+    aten.add_.Tensor,
     aten.add.out,
-    aten.add_.Scalar,
-    aten.addcdiv.default,
-    aten.addcdiv.out,
-    aten.addcdiv_.default,
-    aten.addcmul.default,
-    aten.addcmul.out,
-    aten.addcmul_.default,
-    aten.angle.default,
-    aten.angle.out,
-    aten.asin.default,
-    aten.asin.out,
-    aten.asin_.default,
+    aten.sub.Tensor,
+    aten.sub_.Tensor,
+    aten.sub.out,
+    # foreach variants
+    aten._foreach_add.List,
+    aten._foreach_add_.List,
+    aten._foreach_sub.List,
+    aten._foreach_sub_.List,
+]
+
+_BINARY_ADDITIVE_RULES: list[list[Placement]] = [
+    [Partial("sum"), Partial("sum"), Partial("sum")],
+    [Partial("avg"), Partial("avg"), Partial("avg")],
+    # P(x), R -> P(x): adding/subtracting a replicated value preserves partial types
+    # avg, max, min. sum would result in R being added n times, n = num_ranks
+    # (the replicated value is constant across ranks, so reduce order is unaffected)
+    [Partial("avg"), Partial("avg"), Replicate()],
+    [Partial("max"), Partial("max"), Replicate()],
+    [Partial("min"), Partial("min"), Replicate()],
+    # R, P(avg) -> P(avg): avg is linear so this holds for any alpha
+    # (R, P(max/min) excluded: negative alpha would flip the ordering)
+    [Partial("avg"), Replicate(), Partial("avg")],
+]
+
+for op in binary_additive_ops:
+    _register_single_dim_pointwise(op, _BINARY_ADDITIVE_RULES)
+
+# mul: partials propagate through either arg. div: only through numerator.
+binary_mul_ops = [
+    aten.mul.Tensor,
+    aten.mul_.Tensor,
+    aten.mul.out,
+    # foreach variants
+    aten._foreach_mul.List,
+    aten._foreach_mul_.List,
+    aten._foreach_mul.Tensor,
+    aten._foreach_mul_.Tensor,
+]
+binary_div_ops = [
+    aten.div.Tensor,
+    aten.div_.Tensor,
+    aten.div.out,
+    # foreach variants
+    aten._foreach_div.List,
+    aten._foreach_div_.List,
+    aten._foreach_div.Tensor,
+    aten._foreach_div_.Tensor,
+]
+
+# _UNARY_LINEAR_RULES handles the scalar promotion case: Python's __mul__/__truediv__
+# promote scalars to 0-dim tensors, so aten.mul.Scalar dispatches as aten.mul.Tensor
+# with n_tensors=1, matching the length-2 unary rules.
+_MUL_RULES: list[list[Placement]] = [
+    [Partial("sum"), Partial("sum"), Replicate()],
+    [Partial("avg"), Partial("avg"), Replicate()],
+    [Partial("sum"), Replicate(), Partial("sum")],
+    [Partial("avg"), Replicate(), Partial("avg")],
+]
+
+_DIV_RULES: list[list[Placement]] = [
+    [Partial("sum"), Partial("sum"), Replicate()],
+    [Partial("avg"), Partial("avg"), Replicate()],
+]
+
+for op in binary_mul_ops:
+    _register_single_dim_pointwise(op, _UNARY_LINEAR_RULES + _MUL_RULES)
+
+for op in binary_div_ops:
+    _register_single_dim_pointwise(op, _UNARY_LINEAR_RULES + _DIV_RULES)
+
+scalar_linear_ops = [
+    aten.div.Scalar,
+    aten.div_.Scalar,
+    aten.mul.Scalar,
+    aten.mul_.Scalar,
+    # foreach variants
+    aten._foreach_div.Scalar,
+    aten._foreach_div_.Scalar,
+    aten._foreach_mul.Scalar,
+    aten._foreach_mul_.Scalar,
+    aten._foreach_div.ScalarList,
+    aten._foreach_div_.ScalarList,
+    aten._foreach_mul.ScalarList,
+    aten._foreach_mul_.ScalarList,
+]
+
+for op in scalar_linear_ops:
+    _register_single_dim_pointwise(op, _UNARY_LINEAR_RULES, static_argnum=1)
+
+
+# Non-decreasing unary ops: f(max(a,b)) = max(f(a),f(b)).
+# Only ops that are non-decreasing on their ENTIRE domain belong here.
+# Ops with restricted domains (e.g. log on (0,∞), asin on [-1,1]) do NOT qualify
+# because P(max) offsets can push inputs outside the valid domain.
+non_decreasing_unary_ops = [
     aten.asinh.default,
-    aten.asinh.out,
     aten.asinh_.default,
+    aten.asinh.out,
     aten.atan.default,
-    aten.atan.out,
-    aten.atan2.default,
-    aten.atan2.out,
-    aten.atan2_.default,
     aten.atan_.default,
-    aten.atanh.default,
-    aten.atanh.out,
-    aten.atanh_.default,
-    aten.bitwise_and.Scalar,
-    aten.bitwise_and.Scalar_Tensor,
-    aten.bitwise_and.Scalar_out,
-    aten.bitwise_and.Tensor,
-    aten.bitwise_and.Tensor_out,
-    aten.bitwise_and_.Scalar,
-    aten.bitwise_and_.Tensor,
-    aten.bitwise_left_shift.Scalar_Tensor,
-    aten.bitwise_left_shift.Tensor,
-    aten.bitwise_left_shift.Tensor_Scalar,
-    aten.bitwise_left_shift.Tensor_Scalar_out,
-    aten.bitwise_left_shift.Tensor_out,
-    aten.bitwise_left_shift_.Tensor,
-    aten.bitwise_left_shift_.Tensor_Scalar,
-    aten.bitwise_not.default,
-    aten.bitwise_not.out,
-    aten.bitwise_not_.default,
-    aten.bitwise_or.Scalar,
-    aten.bitwise_or.Scalar_Tensor,
-    aten.bitwise_or.Scalar_out,
-    aten.bitwise_or.Tensor,
-    aten.bitwise_or.Tensor_out,
-    aten.bitwise_or_.Scalar,
-    aten.bitwise_or_.Tensor,
-    aten.bitwise_right_shift.Scalar_Tensor,
-    aten.bitwise_right_shift.Tensor,
-    aten.bitwise_right_shift.Tensor_Scalar,
-    aten.bitwise_right_shift.Tensor_Scalar_out,
-    aten.bitwise_right_shift.Tensor_out,
-    aten.bitwise_right_shift_.Tensor,
-    aten.bitwise_right_shift_.Tensor_Scalar,
-    aten.bitwise_xor.Scalar,
-    aten.bitwise_xor.Scalar_Tensor,
-    aten.bitwise_xor.Scalar_out,
-    aten.bitwise_xor.Tensor,
-    aten.bitwise_xor.Tensor_out,
-    aten.bitwise_xor_.Scalar,
-    aten.bitwise_xor_.Tensor,
+    aten.atan.out,
     aten.ceil.default,
-    aten.ceil.out,
     aten.ceil_.default,
-    aten.clamp.default,
-    aten.clamp.Tensor,
-    aten.clamp.out,
-    aten.clamp_.default,
-    aten.clamp_.Tensor,
-    aten.clamp_min.default,
-    aten.clamp_min.Tensor,
-    aten.clamp_max.default,
-    aten.clamp_max.Tensor,
-    aten.clip.default,
-    aten.clip.out,
-    aten.clip_.default,
-    aten.conj_physical.default,
-    aten.conj_physical.out,
-    aten.conj_physical_.default,
-    aten.copysign.Scalar,
-    aten.copysign.Scalar_out,
-    aten.copysign.Tensor,
-    aten.copysign.out,
-    aten.copysign_.Scalar,
-    aten.copysign_.Tensor,
-    aten.cos.default,
-    aten.cos.out,
-    aten.cos_.default,
-    aten.cosh.default,
-    aten.cosh.out,
-    aten.cosh_.default,
+    aten.ceil.out,
     aten.deg2rad.default,
-    aten.deg2rad.out,
     aten.deg2rad_.default,
-    aten.digamma.default,
-    aten.digamma.out,
-    aten.digamma_.default,
-    aten.div.Tensor_mode,
-    aten.div.out,
-    aten.div.out_mode,
-    aten.div_.Tensor_mode,
-    aten.eq.Tensor,
-    aten.eq.Tensor_out,
-    aten.eq.Scalar,
-    aten.eq.Scalar_out,
+    aten.deg2rad.out,
     aten.erf.default,
-    aten.erf.out,
     aten.erf_.default,
-    aten.erfc.default,
-    aten.erfc.out,
-    aten.erfc_.default,
-    aten.erfinv.default,
-    aten.erfinv.out,
-    aten.erfinv_.default,
+    aten.erf.out,
     aten.exp.default,
+    aten.exp_.default,
     aten.exp.out,
     aten.exp2.default,
-    aten.exp2.out,
     aten.exp2_.default,
-    aten.exp_.default,
+    aten.exp2.out,
     aten.expm1.default,
-    aten.expm1.out,
     aten.expm1_.default,
-    aten.float_power.Scalar,
-    aten.float_power.Scalar_out,
-    aten.float_power.Tensor_Scalar,
-    aten.float_power.Tensor_Scalar_out,
-    aten.float_power.Tensor_Tensor,
-    aten.float_power.Tensor_Tensor_out,
-    aten.float_power_.Scalar,
-    aten.float_power_.Tensor,
+    aten.expm1.out,
     aten.floor.default,
-    aten.floor.out,
     aten.floor_.default,
-    aten.fmax.default,
-    aten.fmax.out,
-    aten.fmin.default,
-    aten.fmin.out,
-    aten.fmod.Scalar,
-    aten.fmod.Scalar_out,
-    aten.fmod.Tensor,
-    aten.fmod.Tensor_out,
-    aten.fmod_.Scalar,
-    aten.fmod_.Tensor,
-    aten.frac.default,
-    aten.frac.out,
-    aten.frac_.default,
-    aten.ge.Scalar,
-    aten.ge.Tensor,
-    aten.gelu.default,
-    aten.gt.Tensor,
-    aten.gt.Tensor_out,
-    aten.gt.Scalar,
-    aten.gt.Scalar_out,
-    aten.gt.Scalar,
-    aten.gt.Tensor,
-    aten.heaviside.default,
-    aten.heaviside.out,
-    aten.hypot.default,
-    aten.hypot.out,
-    aten.hypot_.default,
-    aten.i0.default,
-    aten.i0.out,
-    aten.i0_.default,
-    aten.igamma.default,
-    aten.igamma.out,
-    aten.igamma_.default,
-    aten.igammac.default,
-    aten.igammac.out,
-    aten.igammac_.default,
-    aten.isinf.default,
-    aten.isnan.default,
-    aten.isneginf.default,
-    aten.isneginf.out,
-    aten.isposinf.default,
-    aten.isposinf.out,
-    aten.ldexp.Tensor,
-    aten.ldexp.out,
-    aten.ldexp_.default,
-    aten.lt.Tensor,
-    aten.lt.Tensor_out,
-    aten.lt.Scalar,
-    aten.lt.Scalar_out,
-    aten.le.Scalar,
-    aten.le.Tensor,
-    aten.lerp.Scalar,
-    aten.lerp.Scalar_out,
-    aten.lerp.Tensor,
-    aten.lerp.Tensor_out,
-    aten.lerp_.Scalar,
-    aten.lerp_.Tensor,
-    aten.lgamma.default,
-    aten.lgamma.out,
-    aten.lgamma_.default,
-    aten.log.default,
-    aten.log.out,
-    aten.log10.default,
-    aten.log10.out,
-    aten.log10_.default,
-    aten.log1p.default,
-    aten.log1p.out,
-    aten.log1p_.default,
-    aten.log2.default,
-    aten.log2.out,
-    aten.log2_.default,
-    aten.log_.default,
-    aten.logaddexp.default,
-    aten.logaddexp.out,
-    aten.logaddexp2.default,
-    aten.logaddexp2.out,
-    aten.logical_and.default,
-    aten.logical_and.out,
-    aten.logical_and_.default,
-    aten.logical_not.default,
-    aten.logical_not.out,
-    aten.logical_not_.default,
-    aten.logical_or.default,
-    aten.logical_or.out,
-    aten.logical_or_.default,
-    aten.logical_xor.default,
-    aten.logical_xor.out,
-    aten.logical_xor_.default,
-    aten.logit.default,
-    aten.logit.out,
-    aten.logit_.default,
-    aten.masked_fill.Scalar,
-    aten.masked_fill_.Scalar,
-    aten.mul.out,
-    aten.mvlgamma.default,
-    aten.mvlgamma.out,
-    aten.mvlgamma_.default,
-    aten.native_dropout_backward.default,
-    aten.native_dropout_backward.out,
-    aten.nan_to_num.default,
-    aten.nan_to_num.out,
-    aten.nan_to_num_.default,
-    aten.ne.Scalar,
-    aten.neg.out,
-    aten.nextafter.default,
-    aten.nextafter.out,
-    aten.nextafter_.default,
-    aten.polygamma.default,
-    aten.polygamma.out,
-    aten.polygamma_.default,
-    aten.positive.default,
-    aten.pow.Scalar,
-    aten.pow.Scalar_out,
-    aten.pow.Tensor_Scalar,
-    aten.pow.Tensor_Scalar_out,
-    aten.pow.Tensor_Tensor,
-    aten.pow.Tensor_Tensor_out,
-    aten.pow_.Scalar,
-    aten.pow_.Tensor,
-    aten.reciprocal.default,
-    aten.reciprocal.out,
-    aten.reciprocal_.default,
+    aten.floor.out,
     aten.rad2deg.default,
-    aten.rad2deg.out,
     aten.rad2deg_.default,
+    aten.rad2deg.out,
     aten.relu.default,
     aten.relu_.default,
-    aten.remainder.Scalar,
-    aten.remainder.Scalar_Tensor,
-    aten.remainder.Scalar_out,
-    aten.remainder.Tensor,
-    aten.remainder.Tensor_out,
-    aten.remainder_.Scalar,
-    aten.remainder_.Tensor,
     aten.round.decimals,
-    aten.round.decimals_out,
     aten.round.default,
-    aten.round.out,
     aten.round_.decimals,
     aten.round_.default,
-    aten.rsqrt.default,
-    aten.rsqrt.out,
-    aten.rsqrt_.default,
-    aten.rsub.Scalar,
+    aten.round.decimals_out,
+    aten.round.out,
     aten.sgn.default,
-    aten.sgn.out,
     aten.sgn_.default,
+    aten.sgn.out,
     aten.sigmoid.default,
-    aten.sigmoid.out,
     aten.sigmoid_.default,
+    aten.sigmoid.out,
     aten.sign.default,
-    aten.sign.out,
     aten.sign_.default,
-    aten.signbit.default,
-    aten.signbit.out,
-    aten.silu.default,
-    aten.silu.out,
-    aten.sin.default,
-    aten.sin.out,
-    aten.sin_.default,
-    aten.sinc.default,
-    aten.sinc.out,
-    aten.sinc_.default,
+    aten.sign.out,
     aten.sinh.default,
-    aten.sinh.out,
     aten.sinh_.default,
-    aten.sqrt.default,
-    aten.sqrt.out,
-    aten.sqrt_.default,
-    aten.square.default,
-    aten.square.out,
-    aten.square_.default,
-    aten.sub.Scalar,
-    aten.sub.out,
-    aten.sub_.Scalar,
-    aten.tan.default,
-    aten.tan.out,
-    aten.tan_.default,
+    aten.sinh.out,
     aten.tanh.default,
-    aten.tanh.out,
     aten.tanh_.default,
-    aten.true_divide.Tensor,
+    aten.tanh.out,
     aten.trunc.default,
-    aten.trunc.out,
     aten.trunc_.default,
-    aten.where.self,
-    aten.where.self_out,
-    aten.xlogy.OutScalar_Self,
-    aten.xlogy.OutScalar_Other,
-    aten.xlogy.OutTensor,
-    aten.xlogy.Scalar_Other,
-    aten.xlogy.Scalar_Self,
-    aten.xlogy.Tensor,
-    aten.xlogy_.Scalar_Other,
-    aten.xlogy_.Tensor,
-    # backward point-wise ops
-    # please keep the entries below alphabetically sorted
-    aten.gelu_backward.default,
-    aten.sigmoid_backward.default,
-    aten.silu_backward.default,
-    aten.tanh_backward.default,
-    aten.threshold_backward.default,
+    aten.trunc.out,
+    # nan_to_num is non-decreasing on its entire domain (including nan/inf):
+    # it maps -inf→min, nan→0, inf→max, and is identity elsewhere.
+    aten.nan_to_num.default,
+    aten.nan_to_num_.default,
+    aten.nan_to_num.out,
+    # hardshrink: x if |x|>lambd else 0. Non-decreasing on entire domain.
+    aten.hardshrink.default,
+    # I1(x) is monotonically non-decreasing for all real x.
+    aten.special_modified_bessel_i1.default,
+    # threshold(x, t, v): x if x > t else v. Non-decreasing for v <= t (the
+    # common case, including the default v=0, t=0).
+    aten.threshold.default,
+    # foreach variants
+    aten._foreach_exp.default,
+    aten._foreach_exp_.default,
+    aten._foreach_clamp_max_.Scalar,
+    aten._foreach_clamp_min_.Scalar,
 ]
 
-# the linear pointwise ops map, key is op, value is the type of linearity
-linear_pointwise_ops = {
-    aten.to.dtype: 0,
-    aten.add.Tensor: 1,
-    aten.add_.Tensor: 1,
-    aten.sub.Tensor: 1,
-    aten.sub_.Tensor: 1,
-    aten.div.Scalar: 0,
-    aten.div_.Scalar: 0,
-    aten.div.Tensor: 2,
-    aten.div_.Tensor: 2,
-    aten.mul.Scalar: 0,
-    aten.mul_.Scalar: 0,
-    aten.mul.Tensor: 2,
-    aten.mul_.Tensor: 2,
-    # neg is linear: -(A1 + A2) = -A1 + -A2
-    aten.neg.default: 0,
-    aten.neg_.default: 0,
-}
-
-# Ops that preserve specific Partial types through the operation.
-# For example, torch.maximum preserves Partial("max") because
-# max(max(a), max(b)) == max(a, b).
-partial_preserving_ops: dict[torch._ops.OpOverload, str] = {
-    aten.maximum.default: "max",
-    aten.maximum.out: "max",
-    aten.minimum.default: "min",
-    aten.minimum.out: "min",
-}
-
-
-def pointwise_strategy(
-    op_schema: OpSchema,
-    linearity: int = -1,
-    preserve_partial: str | None = None,
-) -> OpStrategy:
-    followed_strategy_index = -1
-    max_shards = -1
-    max_ndim = -1
-
-    if op_schema.is_inplace_op():
-        # inplace op should follow the first arg strategy
-        followed_strategy = op_schema.args_schema[0]
-        followed_strategy_index = 0
-    elif op_schema.is_out_variant_op():
-        # out variant op should follow the out kwarg strategy
-        followed_strategy = op_schema.kwargs_schema["out"]
-        # out variant is technically a kwarg for the strategy to follow so it does not
-        # have an "index", we set it to a reasonably large number just to indicate it's
-        # not a valid index
-        followed_strategy_index = 100
-    else:
-        # normal pointwise op, we choose to follow the arg with
-        # the max shards in case operands needs reshard
-        # in case of multiple operands with max shard, we take
-        # the one with the max number of dimensions
-        for idx, arg_strategy in enumerate(op_schema.args_schema):
-            if not isinstance(arg_strategy, OpStrategy):
-                continue
-
-            arg_max_shards = arg_strategy.max_num_shards()
-            arg_max_ndim = arg_strategy.ndim
-            if (arg_max_shards > max_shards) or (
-                arg_max_shards == max_shards and arg_max_ndim > max_ndim
-            ):
-                followed_strategy_index = idx
-                max_shards = arg_max_shards
-                max_ndim = arg_max_ndim
-
-        followed_strategy = op_schema.args_schema[followed_strategy_index]
-
-    assert isinstance(followed_strategy, OpStrategy), (
-        f"no strategy to follow for {op_schema}!"
-    )
-    return common_pointwise_strategy(
-        op_schema.op,
-        op_schema.args_schema,
-        followed_strategy,
-        followed_strategy_index,
-        linearity,
-        preserve_partial=preserve_partial,
-    )
+_NON_DECREASING_RULES: list[list[Placement]] = [
+    [Partial("max"), Partial("max")],
+    [Partial("min"), Partial("min")],
+]
 
+for op in non_decreasing_unary_ops:
+    _register_single_dim_pointwise(op, _NON_DECREASING_RULES)
 
-def linear_pointwise_strategy(op_schema: OpSchema) -> StrategyType:
-    """
-    Linear pointwise operators can propagate pending reductions.
-    For example, c = add(a, b); if a is pending sum, then c will be
-    pending sum as well without any communication overhead.
-
-    Note that:
-    1. Only unary and binary operations are supported, out variant
-      ops are not supported.
-    2. There're multiple types of linearity, refer to the doc of
-      common_pointwise_strategy for more details.
-    """
-    linearity_type = linear_pointwise_ops.get(op_schema.op, -1)
-    return pointwise_strategy(op_schema, linearity=linearity_type)
+# Non-increasing unary ops: f(max(a,b)) = min(f(a),f(b)).
+# Note: acos excluded due to domain constraints [-1,1] causing validation failures
+non_increasing_unary_ops: list[OpOverload] = [
+    aten.erfc.default,
+    aten.erfc_.default,
+    aten.erfc.out,
+    aten.special_erfcx.default,
+    aten.special_erfcx.out,
+]
 
+_NON_INCREASING_RULES: list[list[Placement]] = [
+    [Partial("min"), Partial("max")],
+    [Partial("max"), Partial("min")],
+]
 
-def partial_preserving_pointwise_strategy(op_schema: OpSchema) -> StrategyType:
-    """
-    Strategy for pointwise ops that preserve specific Partial types.
+for op in non_increasing_unary_ops:
+    _register_single_dim_pointwise(op, _NON_INCREASING_RULES)
 
-    For example, torch.maximum preserves Partial("max") placements because
-    max(max(a), max(b)) == max(a, b). Similarly, torch.minimum preserves
-    Partial("min") placements.
-    """
-    preserve_partial = partial_preserving_ops.get(op_schema.op)
-    return pointwise_strategy(op_schema, preserve_partial=preserve_partial)
+# Bessel K functions are strictly decreasing for x > 0 but undefined at x <= 0.
+# Only P(min)->P(max) is safe: P(min) offsets add positive values to the
+# non-holding rank, keeping all inputs positive. P(max) offsets subtract,
+# which can push inputs to x <= 0 producing NaN.
+_POSITIVE_DOMAIN_NON_INCREASING_RULES: list[list[Placement]] = [
+    [Partial("max"), Partial("min")],
+]
 
+for op in [
+    aten.special_modified_bessel_k0.default,
+    aten.special_modified_bessel_k1.default,
+    aten.special_scaled_modified_bessel_k0.default,
+    aten.special_scaled_modified_bessel_k1.default,
+]:
+    _register_single_dim_pointwise(op, _POSITIVE_DOMAIN_NON_INCREASING_RULES)
 
-def single_mesh_dim_pointwise_strategy(
-    op: OpOverload,
-    args_schema: ArgsType,
-    kwargs_schema: KwargsType,
-    linearity: int = -1,
-) -> list[list[Placement | _ShardingPlaceholder]]:
-    return single_mesh_dim_common_pointwise_strategy(args_schema, linearity)
+# neg is linear: -(A1 + A2) = -A1 + -A2
+neg_ops = [
+    aten.neg.default,
+    aten.neg_.default,
+    aten.neg.out,
+    # foreach variants
+    aten._foreach_neg.default,
+    aten._foreach_neg_.default,
+]
 
+_NEG_RULES: list[list[Placement]] = _UNARY_LINEAR_RULES + _NON_INCREASING_RULES
 
-def single_mesh_dim_linear_pointwise_strategy(
-    linearity: int = -1,
-) -> Callable[
-    [OpOverload, ArgsType, KwargsType], list[list[Placement | _ShardingPlaceholder]]
-]:
-    return functools.partial(single_mesh_dim_pointwise_strategy, linearity=linearity)
-
-
-def single_mesh_dim_common_pointwise_strategy(
-    args_schema: ArgsType,
-    linearity: int = -1,
-    scalar_tensor_idx: Optional[int] = None,
-) -> list[list[Placement | _ShardingPlaceholder]]:
-    # TODO rename
-    tensor_arg_strategies: list[TensorMeta] = [
-        arg for arg in args_schema if isinstance(arg, TensorMeta)
-    ]
-    common_shape = torch.broadcast_shapes(
-        *[arg.shape for arg in args_schema if isinstance(arg, TensorMeta)]
-    )
-    placements_list: list[list[Placement | _ShardingPlaceholder]] = []
-    for i in range(len(common_shape)):
-        # Shard output dim i, and then shard the corresponding arguments if they have a corresponding (non broadcast) dim
-        shard_placements: list[Placement | _ShardingPlaceholder] = [
-            _ShardingPlaceholder(i)
-        ]
-        for arg in tensor_arg_strategies:
-            common_dim_to_arg_dim = infer_broadcast_dims_map(common_shape, arg.shape)
-            if common_dim_to_arg_dim[i] >= 0:
-                shard_placements.append(_ShardingPlaceholder(common_dim_to_arg_dim[i]))
-            else:
-                shard_placements.append(Replicate())
-
-        placements_list.append(shard_placements)
-
-    if linearity == 0:
-        # unary op (e.g. to_copy), and also binary ops like mul.scalar
-        # input, output can be partial
-        assert len(tensor_arg_strategies) == 1, (
-            "expected single tensor input for linearity==0 op"
-        )
-        placements_list.append([Partial("sum"), Partial("sum")])
-        # TODO: do i need to check scalar_tensor_index and assign a replicate to that one, or do i omit a placement for it
-        # TODO: can mul.scalar work with avg or only sum? i think only sum works. common_pointwise_strategy seems
-        # to support both.
-        # TODO: also, i'll be replacing 'Partial(sum)' here with some kind of 'PartialPlaceholder', not yet designed
-        placements_list.append([Partial("avg"), Partial("avg")])
-
-    elif linearity == 1:
-        # binary add ops
-        # (A1 + B1) + (A2 + B2) == (A1 + A2) + (B1 + B2)
-        assert len(tensor_arg_strategies) == 2, (
-            "expected two tensor inputs for linearity==1 op"
-        )
-        placements_list.append([Partial("sum"), Partial("sum"), Partial("sum")])
-    elif linearity == 2:
-        # binary mul ops (2 tensor inputs)
-        # (A * B1) + (A * B2) == A * (B1 + B2)
-        assert len(tensor_arg_strategies) == 2, (
-            "expected two tensor inputs for linearity==2 op"
-        )
-        placements_list.append([Partial("sum"), Partial("sum"), Replicate()])
-        placements_list.append([Partial("sum"), Replicate(), Partial("sum")])
+for op in neg_ops:
+    _register_single_dim_pointwise(op, _NEG_RULES)
 
-    # TODO: handle scalar_tensor_idx
-    return placements_list
+# xlog1py(x, y) = x * log1p(y). Linear in x with y replicated:
+# (a+b)*log1p(y) = a*log1p(y) + b*log1p(y).
+_XLOG1PY_RULES: list[list[Placement]] = [
+    [Partial("sum"), Partial("sum"), Replicate()],
+    [Partial("avg"), Partial("avg"), Replicate()],
+]
 
+for op in [aten.special_xlog1py.default, aten.special_xlog1py.other_scalar]:
+    _register_single_dim_pointwise(op, _XLOG1PY_RULES)
 
-def copy_strategy(op_schema: OpSchema) -> StrategyType:
-    """
-    Strategy for copy_ that preserves any Partial placement.
 
-    copy_ simply copies data and should preserve whatever Partial placement
-    the destination has, regardless of the reduce_op type (sum, avg, max, min, etc.).
-    """
-    return pointwise_strategy(op_schema, preserve_partial="all")
-
-
-def common_pointwise_strategy(
-    op,
-    args_schema: Sequence[object],
-    followed_strategy: OpStrategy,
-    followed_strategy_index: int,
-    linearity: int = -1,
-    scalar_tensor_idx: int | None = None,
-    preserve_partial: str | None = None,
-) -> OpStrategy:
-    """
-    Common strategy for pointwise operations.
-
-    Args:
-        args_schema: Input arguments schema
-        followed_strategy: Strategy to follow for output placement
-        followed_strategy_index: Index of the strategy being followed
-        linearity: depending on the operator, we support different types of linearity
-            -1: the operation does not support linearity
-            0: the unary operation that supports linearity, output propagates partial.
-            1: the binary operation supports add linearity, where it requires every operand
-                to be partial, output propagates partial.
-            2: the binary operation supports multiplicative linearity, where it requires
-                the primary operand to be partial, and the other operands to be replicate,
-                output propagates partial.
-        scalar_tensor_idx: Index of the Replicate scalar tensor for which we allow the mesh
-            to be different from the mesh of followed_strategy
-        preserve_partial: If set, Partial placements with this reduce_op will be preserved
-            through the operation (e.g., "max" for torch.maximum, "min" for torch.minimum).
-    """
-    # handle broadcasting
-    common_shape = torch.broadcast_shapes(
-        *[arg.shape for arg in args_schema if isinstance(arg, OpStrategy)]
-    )
-    pointwise_strategy = OpStrategy([])
-
-    for op_spec in followed_strategy.strategies:
-        spec_to_follow = op_spec.output_spec
-
-        out_placements: list[Placement] = []
-        for placement in spec_to_follow.placements:
-            if isinstance(placement, Shard | _StridedShard):
-                shard_dim = normalize_dim(placement.dim, len(spec_to_follow.shape))
-                common_ndim = len(common_shape)
-                new_shard_dim = common_ndim - len(spec_to_follow.shape) + shard_dim
-                if isinstance(placement, _StridedShard):
-                    out_placements.append(
-                        _StridedShard(
-                            new_shard_dim, split_factor=placement.split_factor
-                        )
-                    )
-                else:
-                    out_placements.append(Shard(new_shard_dim))
-            elif isinstance(placement, Partial):
-                is_scalar_arg = any(isinstance(arg, _Number) for arg in args_schema)
-                propagate_partial = False
-
-                # ordering matters here since NormPartial is a subclass of Partial
-                if isinstance(placement, _NormPartial):
-                    # explanation for args_schema[1] >= 0 can be found in summary
-                    # https://github.com/pytorch/pytorch/pull/170035
-                    propagate_partial = (
-                        op in norm_partial_avoidable_redistribute_ops
-                        and args_schema[1] >= 0  # pyre-ignore[unsupported-operation]
-                    )
+# All-partial-preserving unary ops: P(x)->P(x) for all x.
+# TODO: positive should be removed once CIA (Copy Is All) optimizes it away.
+all_partial_preserving_unary_ops = [
+    aten.to.dtype,
+    aten.positive.default,
+]
 
-                elif isinstance(placement, Partial):
-                    propagate_partial = not (
-                        op in p_sum_scalar_redistribute_ops and is_scalar_arg
-                    )
+_ALL_PARTIAL_PRESERVING_RULES: list[list[Placement]] = [
+    [Partial(r), Partial(r)] for r in ("sum", "avg", "max", "min")
+]
 
-                # Check if this partial type should be preserved
-                # preserve_partial="all" preserves any Partial type (used for copy_)
-                if preserve_partial == "all":
-                    out_placements.append(placement)
-                elif preserve_partial is not None and placement.is_partial(
-                    preserve_partial
-                ):
-                    out_placements.append(placement)
-                # note that only partial-sum and partial-avg are supported for linearity
-                elif (
-                    linearity >= 0
-                    and (placement.is_partial("sum") or placement.is_partial("avg"))
-                    and propagate_partial
-                ):
-                    # propagate the partial placement
-                    out_placements.append(placement)
-                else:
-                    # clear the partial placement if op does not support linearity
-                    # by default we just replicate the partial, need to see if this
-                    # is optimal for all cases
-                    out_placements.append(Replicate())
-            else:
-                out_placements.append(placement)
-
-        input_specs: list[DTensorSpec] = []
-        redistribute_costs: list[list[float]] = []
-        for input_idx, input_arg in enumerate(args_schema):
-            if isinstance(input_arg, OpStrategy):
-                input_arg_spec = input_arg.strategies[0].output_spec
-
-                # sanity check that all args that follow the same strategy
-                # are on the same DeviceMesh
-                if input_arg.mesh != followed_strategy.mesh:
-                    # For the scalar tensor arg in fused ops, do not follow followed_strategy;
-                    # instead, let the input mesh and the Replicate placements propagate through.
-                    if input_idx == scalar_tensor_idx:
-                        assert all(p == Replicate() for p in input_arg_spec.placements)
-                        input_arg_target_spec = DTensorSpec(
-                            mesh=input_arg.mesh,
-                            placements=input_arg_spec.placements,
-                            tensor_meta=input_arg_spec.tensor_meta,
-                        )
-                        input_specs.append(input_arg_target_spec)
-                        redistribute_costs.append(
-                            generate_redistribute_costs(
-                                input_arg, input_arg_target_spec
-                            )
-                        )
-                        continue
-                    else:
-                        raise ValueError(
-                            f"Could not run pointwise computation across different mesh: "
-                            f"Found {input_arg.mesh} and {followed_strategy.mesh}!"
-                        )
-
-                # every arg follow the out_placements, but need to handle broadcasting
-                input_arg_dims_map = infer_broadcast_dims_map(
-                    common_shape, input_arg_spec.shape
-                )
+for op in all_partial_preserving_unary_ops:
+    _register_single_dim_pointwise(op, _ALL_PARTIAL_PRESERVING_RULES)
 
-                # Determine if this input should convert Partial to Replicate based on linearity
-                should_convert_partial = (
-                    linearity == 2
-                    and input_idx
-                    != followed_strategy_index  # Don't convert the "followed" strategy
-                )
+all_partial_preserving_binary_ops = [
+    aten.copy_.default,
+    prims.copy_to.default,
+]
 
-                # For preserve_partial ops, check if non-followed input has incompatible
-                # Partial type. If so, it must be redistributed to Replicate first.
-                if (
-                    preserve_partial is not None
-                    and input_idx != followed_strategy_index
-                ):
-                    for out_p, in_p in zip(out_placements, input_arg_spec.placements):
-                        if (
-                            isinstance(out_p, Partial)
-                            and isinstance(in_p, Partial)
-                            and out_p != in_p
-                        ):
-                            should_convert_partial = True
-                            break
-
-                input_target_placements = map_placements_after_broadcast(
-                    tuple(out_placements),
-                    common_shape,
-                    input_arg_dims_map,
-                    partial_to_replicate=should_convert_partial,
-                )
+_ALL_PARTIAL_BINARY_PRESERVING_RULES: list[list[Placement]] = [
+    [Partial(r), Partial(r), Partial(r)] for r in ("sum", "avg", "max", "min")
+]
 
-                input_arg_target_spec = DTensorSpec(
-                    mesh=followed_strategy.mesh,
-                    placements=input_target_placements,
-                    tensor_meta=input_arg_spec.tensor_meta,
-                )
-                input_specs.append(input_arg_target_spec)
-                redistribute_costs.append(
-                    generate_redistribute_costs(input_arg, input_arg_target_spec)
-                )
+for op in all_partial_preserving_binary_ops:
+    _register_single_dim_pointwise(op, _ALL_PARTIAL_BINARY_PRESERVING_RULES)
 
-        pointwise_strategy.strategies.append(
-            OpSpec(
-                output_specs=DTensorSpec(
-                    mesh=followed_strategy.mesh,
-                    placements=tuple(out_placements),
-                ),
-                input_specs=input_specs,
-                redistribute_cost=redistribute_costs,
-            )
-        )
-    return pointwise_strategy
+# Monotonic increasing in both args but don't preserve any specific partial type.
+monotonic_binary_ops = [
+    aten.logaddexp.default,
+    aten.logaddexp.out,
+    aten.logaddexp2.default,
+    aten.logaddexp2.out,
+]
 
+_MONOTONE_BINARY_BASE_RULES: list[list[Placement]] = [
+    [Partial("max"), Partial("max"), Replicate()],
+    [Partial("max"), Replicate(), Partial("max")],
+    [Partial("min"), Partial("min"), Replicate()],
+    [Partial("min"), Replicate(), Partial("min")],
+]
 
-p_sum_scalar_redistribute_ops = {
-    aten.add.Tensor,
-    aten.add_.Tensor,
-    aten.sub.Tensor,
-    aten.sub_.Tensor,
-}
+for op in monotonic_binary_ops:
+    _register_single_dim_pointwise(op, _MONOTONE_BINARY_BASE_RULES)
 
-norm_partial_avoidable_redistribute_ops = {
-    aten.div.Scalar,
-    aten.div_.Scalar,
-    aten.mul.Scalar,
-    aten.mul_.Scalar,
-}
+# Binary ops monotonically increasing in both arguments.
+# max-preserving: P(max)+P(max)->P(max) because max(max(a),max(b)) = max(a,b)
+monotonic_max_preserving_binary_ops = [
+    aten.clamp_min.Tensor,
+    aten.fmax.default,
+    aten.fmax.out,
+    aten.maximum.default,
+    aten.maximum.out,
+    prims.fmax.default,
+    # foreach variants
+    aten._foreach_maximum_.List,
+]
 
-for op in linear_pointwise_ops:
-    if op in norm_partial_avoidable_redistribute_ops:
-        register_op_strategy(
-            op, schema_info=RuntimeSchemaInfo(1, static_kwargkey=["out"])
-        )(linear_pointwise_strategy)
-    else:
-        register_op_strategy(
-            op, schema_info=RuntimeSchemaInfo(static_kwargkey=["out"])
-        )(linear_pointwise_strategy)
+_MONOTONE_MAX_PRESERVING_BINARY_BASE_RULES: list[list[Placement]] = [
+    *_MONOTONE_BINARY_BASE_RULES,
+    [Partial("max"), Partial("max"), Partial("max")],
+]
 
-for op in partial_preserving_ops:
-    register_op_strategy(op, schema_info=RuntimeSchemaInfo(static_kwargkey=["out"]))(
-        partial_preserving_pointwise_strategy
-    )
+for op in monotonic_max_preserving_binary_ops:
+    _register_single_dim_pointwise(op, _MONOTONE_MAX_PRESERVING_BINARY_BASE_RULES)
+
+# min-preserving: P(min)+P(min)->P(min) because min(min(a),min(b)) = min(a,b)
+monotonic_min_preserving_binary_ops = [
+    aten.clamp_max.Tensor,
+    aten.fmin.default,
+    aten.fmin.out,
+    aten.minimum.default,
+    aten.minimum.out,
+    prims.fmin.default,
+]
 
-# Register copy_ with its custom strategy that preserves all Partial types
-register_op_strategy(
-    aten.copy_.default, schema_info=RuntimeSchemaInfo(static_kwargkey=["out"])
-)(copy_strategy)
+_MONOTONE_MIN_PRESERVING_BINARY_BASE_RULES: list[list[Placement]] = [
+    *_MONOTONE_BINARY_BASE_RULES,
+    [Partial("min"), Partial("min"), Partial("min")],
+]
+
+for op in monotonic_min_preserving_binary_ops:
+    _register_single_dim_pointwise(op, _MONOTONE_MIN_PRESERVING_BINARY_BASE_RULES)
 
-for op in pointwise_ops:
-    register_op_strategy(op, schema_info=RuntimeSchemaInfo(static_kwargkey=["out"]))(
-        pointwise_strategy
-    )
-    # register_single_dim_strategy(
-    #     op, schema_info=RuntimeSchemaInfo(static_kwargkey=["out"])
-    # )(single_mesh_dim_pointwise_strategy)
 
-# TODO: add all for_each ops
-for_each_ops = [
+# Ops that are pointwise for DTensor purposes but lack torch.Tag.pointwise.
+# TODO(pianpwk): add torch.Tag.pointwise to these ops in native_functions.yaml
+# so this list can be removed.
+_extra_pointwise_ops: list[OpOverload] = [
+    aten.__irshift__.Scalar,
+    aten.__irshift__.Tensor,
+    aten._conj.default,
+    aten.abs_.default,
+    aten.copysign_.Scalar,
+    aten.copysign_.Tensor,
+    aten.exponential_.default,
+    aten.float_power.Scalar,
+    aten.float_power.Scalar_out,
+    aten.float_power.Tensor_Scalar,
+    aten.float_power.Tensor_Scalar_out,
+    aten.float_power.Tensor_Tensor,
+    aten.float_power.Tensor_Tensor_out,
+    aten.masked_fill_.Scalar,
+    aten.native_dropout_backward.out,
+    aten.polygamma_.default,
+    aten.rrelu_with_noise.default,
+    aten.where.self_out,
+    aten.xlogy_.Scalar_Other,
+    prims.bessel_i0e.default,
+    prims.bessel_i1.default,
+    prims.bessel_i1e.default,
+    prims.bessel_j0.default,
+    prims.bessel_j1.default,
+    prims.div.default,
+    prims.erfcx.default,
+    prims.frexp.default,
+    prims.gcd.default,
+    prims.ndtri.default,
+    prims.ne.default,
+    prims.spherical_bessel_j0.default,
+    prims.zeta.default,
+    # foreach variants
     aten._foreach_abs.default,
     aten._foreach_abs_.default,
     aten._foreach_addcdiv_.Scalar,
@@ -849,144 +521,27 @@ def common_pointwise_strategy(
     aten._foreach_addcmul_.Scalar,
     aten._foreach_addcmul_.ScalarList,
     aten._foreach_addcmul_.Tensor,
-    aten._foreach_clamp_max_.Scalar,
-    aten._foreach_clamp_min_.Scalar,
-    aten._foreach_div_.List,
-    aten._foreach_div_.Scalar,
-    aten._foreach_div_.ScalarList,
-    aten._foreach_div_.Tensor,
-    aten._foreach_div.List,
-    aten._foreach_div.Scalar,
-    aten._foreach_div.ScalarList,
-    aten._foreach_div.Tensor,
     aten._foreach_lerp_.Scalar,
-    aten._foreach_maximum_.List,
-    aten._foreach_mul.Scalar,
-    aten._foreach_mul.ScalarList,
-    aten._foreach_mul.Tensor,
-    aten._foreach_mul.List,
-    aten._foreach_mul_.Scalar,
-    aten._foreach_mul_.ScalarList,
-    aten._foreach_mul_.Tensor,
-    aten._foreach_mul_.List,
     aten._foreach_pow.List,
     aten._foreach_pow.ScalarList,
-    aten._foreach_neg.default,
-    aten._foreach_neg_.default,
     aten._foreach_reciprocal_.default,
     aten._foreach_sub.Scalar,
     aten._foreach_sub_.Scalar,
-    aten._foreach_sub.List,
-    aten._foreach_sub_.List,
     aten._foreach_sub.ScalarList,
     aten._foreach_sub_.ScalarList,
     aten._foreach_sqrt.default,
     aten._foreach_sqrt_.default,
     aten._foreach_zero_.default,
-    aten._foreach_exp.default,
-    aten._foreach_exp_.default,
     aten._foreach_cos.default,
     aten._foreach_cos_.default,
     aten._foreach_log.default,
     aten._foreach_log_.default,
     aten._amp_foreach_non_finite_check_and_unscale_.default,
-]
-
-for_each_linearity_ops = [
+    # foreach linearity variants
     aten._foreach_add.Scalar,
     aten._foreach_add_.Scalar,
     aten._foreach_add_.ScalarList,
-    aten._foreach_add.List,
-    aten._foreach_add_.List,
-]
-
-
-def list_pointwise_strategy(
-    op_schema: OpSchema, linearity: bool = False
-) -> StrategyType:
-    """
-    Apply the pointwise strategy to the zipped arguments. For example, if we
-    run a foreach add of two lists l1 and l2, then we apply the pointwise
-    strategy on each pair (l1[i], l2[i]). If the first argument is a list but
-    the second (or later) one is a tensor, then we broadcast the tensor by
-    replicating it into a list with the length of the first argument.
-
-    Args:
-        mesh (DeviceMesh): device mesh for pointwise ops
-        op_schema (OpSchema): schema of the operator to generate strategy for
-        linearity (bool): specify whether op(a) + op(b) = op(a + b)
-
-    Returns:
-        OpStrategy: generated strategy
-    """
-
-    def args_tuple_strategies(
-        args_schema: tuple[object, ...],
-    ) -> list[TupleStrategy | None]:
-        first_arg = args_schema[0]
-        assert isinstance(first_arg, TupleStrategy)
-        strategy_len = len(first_arg.children)
-        tuple_strategies: list[TupleStrategy | None] = []
-        for arg_idx, arg in enumerate(args_schema):
-            if isinstance(arg, TupleStrategy):
-                # every tuple strategy should have the same length
-                assert len(arg.children) == strategy_len
-                tuple_strategies.append(arg)
-            elif isinstance(arg, OpStrategy):
-                if arg_idx > 0:  # implicitly broadcast
-                    tuple_strategies.append(
-                        TupleStrategy([arg for _ in range(strategy_len)])
-                    )
-                else:
-                    raise RuntimeError(
-                        f"list op only supports tuple strategy! {op_schema}"
-                    )
-            else:
-                # insert None as placeholder so that the idx of arg is kept
-                tuple_strategies.append(None)
-        return tuple_strategies
-
-    args_strategies = args_tuple_strategies(op_schema.args_schema)
-    follow_strategy: TupleStrategy = not_none(args_strategies[0])
-    list_strategy: list[OpStrategy] = []
-
-    for child_idx, child_strtgy in enumerate(follow_strategy.children):
-        assert isinstance(child_strtgy, OpStrategy)
-        args_schema: list[OpStrategy | None] = [
-            cast(OpStrategy, arg_strategy.children[child_idx]) if arg_strategy else None
-            for arg_strategy in args_strategies
-        ]
-        pointwise_strategy: OpStrategy = common_pointwise_strategy(
-            op_schema.op,
-            args_schema,
-            child_strtgy,
-            linearity,
-            scalar_tensor_idx=(
-                _FUSED_OP_SCALAR_IDX if op_schema.op in fused_ops else None
-            ),
-        )
-        list_strategy.append(pointwise_strategy)
-    return TupleStrategy(list_strategy)
-
-
-def list_linear_pointwise_strategy(op_schema: OpSchema) -> StrategyType:
-    """
-    for each list op stratgy that supports linearity
-    """
-    return list_pointwise_strategy(op_schema, linearity=True)
-
-
-for op in for_each_ops:
-    register_op_strategy(op, schema_info=RuntimeSchemaInfo(needs_pytree=True))(
-        list_pointwise_strategy
-    )
-
-for op in for_each_linearity_ops:
-    register_op_strategy(op, schema_info=RuntimeSchemaInfo(needs_pytree=True))(
-        list_linear_pointwise_strategy
-    )
-
-fused_ops = [
+    # fused optimizer ops
     aten._fused_adam_.default,
     aten._fused_adam.default,
     aten._fused_adam.tensor_lr,
@@ -998,13 +553,40 @@ def list_linear_pointwise_strategy(op_schema: OpSchema) -> StrategyType:
 ]
 
 
-# The state_steps arg of fused adam / adamw is a Replicate scalar tensor, which will be put on
-# the compute_mesh of an op across all parameter groups, even when not all parameter groups
-# are on the same device mesh. This idx will help avoid hitting exceptions or unnecessary
-# redistribute during sharding propagation.
-_FUSED_OP_SCALAR_IDX = 5
+def _get_pointwise_ops_from_tag() -> list[OpOverload]:
+    """
+    Auto-discover pointwise ops via torch.Tag.pointwise, from ops.aten, ops.prims.
+    """
+    ops = []
+    for ns in [torch.ops.aten, torch.ops.prims]:
+        for attr_name in dir(ns):
+            attr = getattr(ns, attr_name)
+            if isinstance(attr, torch._ops.OpOverloadPacket):
+                for overload_name in attr.overloads():
+                    op = getattr(attr, overload_name)
+                    if torch.Tag.pointwise in op.tags:
+                        ops.append(op)
+    return ops
 
-for op in fused_ops:
-    register_op_strategy(op, schema_info=RuntimeSchemaInfo(needs_pytree=True))(
-        list_pointwise_strategy
-    )
+
+pointwise_ops = [
+    op
+    for op in _get_pointwise_ops_from_tag() + _extra_pointwise_ops
+    if op not in _specially_registered_ops
+]
+
+
+for op in pointwise_ops:
+    _register_single_dim_pointwise(op)
+
+
+def register_inductor_prims() -> None:
+    """Register DTensor sharding strategies for inductor prims ops.
+
+    Called lazily because inductor prims are created via make_prim() in
+    torch._inductor.inductor_prims, which is imported after this module.
+    """
+    # TODO: handle other inductor prims ops that may need DTensor sharding
+    # strategies (e.g. mul_rn, div_rn). Those are more complicated and not
+    # necessarily pointwise.
+    _register_single_dim_pointwise(prims.fma.default)
diff --git a/torch/distributed/tensor/_ops/_random_ops.py b/torch/distributed/tensor/_ops/_random_ops.py
index 9db9b85e58d2d..76dd731636d90 100644
--- a/torch/distributed/tensor/_ops/_random_ops.py
+++ b/torch/distributed/tensor/_ops/_random_ops.py
@@ -23,7 +23,8 @@
 )
 def random_op_strategy(op_schema: OpSchema) -> StrategyType:
     self_strategy = op_schema.args_schema[0]
-    assert isinstance(self_strategy, OpStrategy)
+    if not isinstance(self_strategy, OpStrategy):
+        raise AssertionError
 
     random_strategy = OpStrategy([])
     for arg_strategy in self_strategy.strategies:
diff --git a/torch/distributed/tensor/_ops/_tensor_ops.py b/torch/distributed/tensor/_ops/_tensor_ops.py
index ad2b2704a0c7f..134dea27537e4 100644
--- a/torch/distributed/tensor/_ops/_tensor_ops.py
+++ b/torch/distributed/tensor/_ops/_tensor_ops.py
@@ -1,6 +1,6 @@
 # mypy: allow-untyped-defs
 # Copyright (c) Meta Platforms, Inc. and affiliates
-from collections.abc import Sequence, Sized
+from collections.abc import Callable, Sequence, Sized
 from typing import cast
 
 import torch
@@ -13,28 +13,28 @@
     OpSchema,
     OpSpec,
     OpStrategy,
-    OutputSharding,
     PlacementList,
     RuntimeSchemaInfo,
     StrategyType,
     TensorMeta,
     TupleStrategy,
 )
-from torch.distributed.tensor._ops._common_rules import pointwise_rule
-from torch.distributed.tensor._ops.single_dim_strategy import _ShardingPlaceholder
+from torch.distributed.tensor._ops.single_dim_strategy import (
+    _ShardingPlaceholder,
+    register_single_dim_strategy,
+)
 from torch.distributed.tensor._ops.utils import (
     expand_to_full_mesh_op_strategy,
     generate_redistribute_costs,
     is_tensor_dim_sharded,
-    is_tensor_evenly_shardable,
     is_tensor_partial,
     normalize_dim,
     register_op_strategy,
-    register_prop_rule,
     shift_shard_dims_after_insert,
     shift_shard_dims_after_remove,
 )
 from torch.distributed.tensor.placement_types import (
+    _is_shard_like,
     _MaskPartial,
     Partial,
     Placement,
@@ -45,6 +45,7 @@
 
 
 aten = torch.ops.aten
+prims = torch.ops.prims
 
 
 def propagate_single_input_strategy(op_schema: OpSchema) -> StrategyType:
@@ -94,13 +95,46 @@ def propagate_single_input_strategy(op_schema: OpSchema) -> StrategyType:
         aten.fill_.Scalar,
         aten.view.dtype,
         aten.zero_.default,
+        prims.view_of.default,
     ]
 )(propagate_single_input_strategy)
 
 
-register_op_strategy(
-    aten._to_copy.default, schema_info=RuntimeSchemaInfo(static_kwargkey=["dtype"])
-)(propagate_single_input_strategy)
+def _partial_needs_reduce_for_dtype_cast(
+    reduce_op: str,
+    src_dtype: torch.dtype,
+    target_dtype: torch.dtype | None,
+) -> bool:
+    """Return True when reduce_op does not commute with the dtype cast."""
+    if target_dtype is None or src_dtype == target_dtype:
+        return False
+    if target_dtype == torch.bool:
+        return True
+    if reduce_op in ("max", "min"):
+        return False
+    return src_dtype.is_floating_point and not target_dtype.is_floating_point
+
+
+@register_single_dim_strategy(
+    aten._to_copy.default,
+    schema_info=RuntimeSchemaInfo(static_kwargkey=["dtype"]),
+    allow_unbacked_sharding=True,
+    allow_uneven_sharding=True,
+)
+def _to_copy_single_dim_strategy(
+    op: OpOverload, args_schema: ArgsType, kwargs_schema: KwargsType
+) -> list[list[Placement | _ShardingPlaceholder]]:
+    input_meta = cast(TensorMeta, args_schema[0])
+    src_dtype = input_meta.dtype
+    target_dtype = cast(torch.dtype | None, kwargs_schema.get("dtype", None))
+
+    strategies: list[list[Placement | _ShardingPlaceholder]] = []
+    for dim in range(len(input_meta.shape)):
+        strategies.append([_ShardingPlaceholder(dim), _ShardingPlaceholder(dim)])
+    for reduce_op in Partial.ALL_REDUCE_OPS:
+        if not _partial_needs_reduce_for_dtype_cast(reduce_op, src_dtype, target_dtype):
+            strategies.append([Partial(reduce_op), Partial(reduce_op)])
+    return strategies
 
 
 @register_op_strategy(
@@ -249,14 +283,6 @@ def new_factory_strategy(op_schema: OpSchema) -> StrategyType:
         )
 
         if tuple(input_shape) == tuple(output_shape) and input_spec.is_sharded():
-            # NOTE: for new_empty_strided, currently the non-replicate sharding
-            #       is supported only when the shape is evenly shardable
-            if (
-                op_schema.op == aten.new_empty_strided.default
-                and not is_tensor_evenly_shardable(input_shape, input_spec)
-            ):
-                continue
-
             new_factory_strategy.strategies.append(
                 OpSpec(
                     output_specs=input_spec,
@@ -269,49 +295,34 @@ def new_factory_strategy(op_schema: OpSchema) -> StrategyType:
     return new_factory_strategy
 
 
-@register_op_strategy(aten.bucketize.Tensor)
-def gen_bucketize_strategy(op_schema: OpSchema) -> StrategyType:
-    """
-    Propagate input sharding to output, but expect replicated for boundaries input.
-    For Partial inputs, convert to Replicate since bucketize returns indices
-    which cannot be meaningfully combined with sum/avg reductions.
+@register_single_dim_strategy(aten.bucketize.Tensor)
+def bucketize_single_dim_strategy(
+    op: OpOverload, args_schema: ArgsType, kwargs_schema: KwargsType
+) -> list[list[Placement | _ShardingPlaceholder]]:
+    """Bucketize returns indices into a sorted boundary tensor.
+
+    Three families of strategies:
+    1. Shard the input (and output) on any dim, keep boundaries replicated.
+    2. Shard boundaries on dim 0, replicate input, output is Partial("sum").
+       Each rank counts how many of its local boundary values each input
+       element exceeds; summing across ranks gives the correct global index.
+    3. Partial("max") or Partial("min") input with replicated boundaries.
+       Bucketize is monotonically non-decreasing in its input, so reducing
+       local bucket indices with max (or min) across ranks gives the same
+       result as bucketizing the reduced input values.
     """
-    mesh = op_schema.get_mesh_from_args()
-    input_strategy, boundaries_strategy = op_schema.args_schema
-    bucketize_strategy = OpStrategy([])
-    if not isinstance(input_strategy, OpStrategy):
-        raise AssertionError(f"Expected OpStrategy, got {type(input_strategy)}")
-    if not isinstance(boundaries_strategy, OpStrategy):
-        raise AssertionError(f"Expected OpStrategy, got {type(boundaries_strategy)}")
-    for arg_strategy in input_strategy.strategies:
-        # Convert Partial placements to Replicate - bucketize returns indices
-        # which cannot be combined with sum/avg reductions
-        placements = tuple(
-            Replicate() if isinstance(p, Partial) else p
-            for p in arg_strategy.output_spec.placements
-        )
-        arg_spec = DTensorSpec(
-            mesh,
-            placements,
-            arg_strategy.output_spec.tensor_meta,
-        )
-        replica_spec = DTensorSpec(
-            mesh,
-            tuple([Replicate()] * mesh.ndim),
-            boundaries_strategy.strategies[0].output_spec.tensor_meta,
-        )
-        bucketize_strategy.strategies.append(
-            OpSpec(
-                output_specs=arg_spec,
-                input_specs=(arg_spec, replica_spec),
-                redistribute_cost=[
-                    generate_redistribute_costs(input_strategy, arg_spec),
-                    generate_redistribute_costs(boundaries_strategy, replica_spec),
-                ],
-            )
+    input_meta, _boundaries_meta = args_schema
+    if not isinstance(input_meta, TensorMeta):
+        raise AssertionError(f"Expected TensorMeta, got {type(input_meta)}")
+    strategies: list[list[Placement | _ShardingPlaceholder]] = []
+    for dim in range(len(input_meta.shape)):
+        strategies.append(
+            [_ShardingPlaceholder(dim), _ShardingPlaceholder(dim), Replicate()]
         )
-
-    return bucketize_strategy
+    strategies.append([Partial("sum"), Replicate(), _ShardingPlaceholder(0)])
+    for reduce_op in ("max", "min"):
+        strategies.append([Partial(reduce_op), Partial(reduce_op), Replicate()])
+    return strategies
 
 
 @register_op_strategy(aten.select.int, schema_info=RuntimeSchemaInfo(1))
@@ -488,7 +499,7 @@ def slice_backward_rules(op_schema: OpSchema) -> OpStrategy:
         new_placements: list[Placement] = []
         for placement in output_spec.placements:
             # Redistribute to replicate only if the dim is sharded and matches the slice dim
-            if isinstance(placement, Shard) and placement.dim == dim:
+            if _is_shard_like(placement) and placement.dim == dim:
                 new_placements.append(Replicate())
             else:
                 new_placements.append(placement)
@@ -506,7 +517,7 @@ def unshard_tensor_dim(
 ) -> tuple[Placement, ...]:
     """Disallow the given tensor dimension to be sharded."""
     return tuple(
-        p if (not isinstance(p, Shard) or p.dim != dim) else Replicate()
+        p if (not _is_shard_like(p) or p.dim != dim) else Replicate()
         for p in placements
     )
 
@@ -515,10 +526,8 @@ def replicate_tensor_dim(
     placements: Sequence[Placement], dim: int
 ) -> tuple[Placement, ...]:
     """Force the given tensor dimension to be replicated."""
-    # Not using p.is_shard() to avoid mypy complain about Placement not having
-    # attribute dim.
     return tuple(
-        Replicate() if p.is_partial() or isinstance(p, Shard) and p.dim == dim else p
+        Replicate() if p.is_partial() or (_is_shard_like(p) and p.dim == dim) else p
         for p in placements
     )
 
@@ -589,6 +598,70 @@ def gen_slice_scatter_strategy(op_schema: OpSchema) -> StrategyType:
     return slice_scatter_strategy
 
 
+@register_single_dim_strategy(
+    [aten.select_scatter.default],
+    schema_info=RuntimeSchemaInfo(1),
+)
+def select_scatter_single_dim_strategy(
+    op: OpOverload,
+    args_schema: ArgsType,
+    kwargs_schema: KwargsType,
+) -> list[list[Placement | _ShardingPlaceholder]]:
+    input_meta = args_schema[0]
+    if not isinstance(input_meta, TensorMeta):
+        raise AssertionError(f"Expected TensorMeta, got {type(input_meta)}")
+    ndim = len(input_meta.shape)
+    dim = normalize_dim(cast(int, args_schema[2]), ndim)
+    # [output, self, src] — src has the select dim removed
+    strategies: list[list[Placement | _ShardingPlaceholder]] = []
+    for d in range(ndim):
+        if d == dim:
+            continue
+        strategies.append(
+            [
+                _ShardingPlaceholder(d),
+                _ShardingPlaceholder(d),
+                _ShardingPlaceholder(d if d < dim else d - 1),
+            ]
+        )
+    return strategies
+
+
+@register_single_dim_strategy(
+    [aten.diagonal_scatter.default],
+    schema_info=RuntimeSchemaInfo(1),
+)
+def diagonal_scatter_single_dim_strategy(
+    op: OpOverload,
+    args_schema: ArgsType,
+    kwargs_schema: KwargsType,
+) -> list[list[Placement | _ShardingPlaceholder]]:
+    input_meta = args_schema[0]
+    if not isinstance(input_meta, TensorMeta):
+        raise AssertionError(f"Expected TensorMeta, got {type(input_meta)}")
+    ndim = len(input_meta.shape)
+    # schema: (self, src, offset=0, dim1=0, dim2=1)
+    dim1 = cast(int, args_schema[3]) if len(args_schema) > 3 else 0
+    dim2 = cast(int, args_schema[4]) if len(args_schema) > 4 else 1
+    dim1 = normalize_dim(dim1, ndim)
+    dim2 = normalize_dim(dim2, ndim)
+    min_d, max_d = min(dim1, dim2), max(dim1, dim2)
+    # [output, self, src] — src has dim1/dim2 removed and diagonal appended
+    strategies: list[list[Placement | _ShardingPlaceholder]] = []
+    for d in range(ndim):
+        if d in (dim1, dim2):
+            continue
+        removed = (1 if d > min_d else 0) + (1 if d > max_d else 0)
+        strategies.append(
+            [
+                _ShardingPlaceholder(d),
+                _ShardingPlaceholder(d),
+                _ShardingPlaceholder(d - removed),
+            ]
+        )
+    return strategies
+
+
 @register_op_strategy(aten._local_scalar_dense.default)
 def replica_only_strategy(op_schema: OpSchema) -> StrategyType:
     """Only allow replication on the input/output."""
@@ -730,11 +803,13 @@ def merge_placement(
         # check each placement for the current arg placement
         # to see if we want to merge/adjust the placement to follow
         # the priority: Partial -> Shard -> Replicate
+        # _StridedShard.__eq__ compares both dim and split_factor,
+        # so two _StridedShard with different split_factor won't match here.
         if cur_placement == new_placement:
             return cur_placement
 
         if cur_placement.is_partial():
-            if new_placement.is_shard():
+            if _is_shard_like(new_placement):
                 # follow new placement
                 return new_placement
             elif new_placement.is_partial():
@@ -743,8 +818,8 @@ def merge_placement(
             else:
                 # follow partial
                 return cur_placement
-        elif cur_placement.is_shard():
-            if new_placement.is_shard():
+        elif _is_shard_like(cur_placement):
+            if _is_shard_like(new_placement):
                 # cur/new placement are different sharding (i.e. different shard dim)
                 # currently fallback to replicate all args
                 return Replicate()
@@ -793,13 +868,15 @@ def stack_strategy(op_schema: OpSchema) -> StrategyType:
         raise AssertionError(f"Expected TupleStrategy, got {input_tuple_strategy}")
     input_strategies: list[OpStrategy] = []
     for child in input_tuple_strategy.children:
-        assert isinstance(child, OpStrategy), f"Expected OpStrategy, got {child}"
+        if not isinstance(child, OpStrategy):
+            raise AssertionError(f"Expected OpStrategy, got {child}")
         input_strategies.append(child)
     first_input_strategy = input_strategies[0]
     common_input_ndim = first_input_strategy.ndim
     dim = cast(int, args_schema[1]) if len(args_schema) > 1 else 0
-    # normalize the dim to be within the common input ndim
-    dim = normalize_dim(dim, common_input_ndim)
+    # normalize the dim to be within the output ndim (input ndim + 1),
+    # since stack inserts a new dimension
+    dim = normalize_dim(dim, common_input_ndim + 1)
 
     mesh = first_input_strategy.mesh
 
@@ -841,19 +918,23 @@ def cat_single_dim_strategy(
 ) -> list[list[Placement | _ShardingPlaceholder]]:
     input_list = args_schema[0]
     # unfortunate naming, but yes it's a TensorList input, and we represent it as a tuple of TensorMeta
-    assert isinstance(input_list, (tuple, list)), type(input_list)
-    assert all(isinstance(tm, TensorMeta) for tm in input_list)
+    if not isinstance(input_list, (tuple, list)):
+        raise AssertionError(type(input_list))
+    if not all(isinstance(tm, TensorMeta) for tm in input_list):
+        raise AssertionError
 
     if isinstance(input_list, list):
         input_list = tuple(input_list)
 
     num_inputs = len(input_list)
     ndim_set = {len(meta.shape) for meta in input_list}
-    assert len(ndim_set) in (1, 2), (
-        "Expected all cat inputs to be the same ndim, except empty tensors"
-    )
+    if len(ndim_set) not in (1, 2):
+        raise AssertionError(
+            "Expected all cat inputs to be the same ndim, except empty tensors"
+        )
     if len(ndim_set) == 2:
-        assert 0 in ndim_set
+        if 0 not in ndim_set:
+            raise AssertionError
     common_ndim = max(ndim_set)
     cat_dim = cast(int, args_schema[1]) if len(args_schema) > 1 else 0
     cat_dim = normalize_dim(cat_dim, common_ndim)
@@ -942,280 +1023,322 @@ def cat_strategy(op_schema: OpSchema) -> StrategyType:
     return op_strategy
 
 
-@register_prop_rule(aten.index_select.default, schema_info=RuntimeSchemaInfo(1))
-def prop_index_select(op_schema: OpSchema) -> OutputSharding:
-    values_spec, dim, indices_spec = op_schema.args_schema
-
-    if not isinstance(values_spec, DTensorSpec):
-        raise AssertionError(f"Expected DTensorSpec, got {type(values_spec)}")
+@register_single_dim_strategy(
+    aten.index_select.default, schema_info=RuntimeSchemaInfo(1)
+)
+def index_select_single_dim_strategy(
+    op: OpOverload, args_schema: ArgsType, kwargs_schema: KwargsType
+) -> list[list[Placement | _ShardingPlaceholder]]:
+    values_meta, dim, index_meta = args_schema
+    if not isinstance(values_meta, TensorMeta):
+        raise AssertionError(f"Expected TensorMeta, got {type(values_meta)}")
     if not isinstance(dim, int):
         raise AssertionError(f"Expected int, got {type(dim)}")
-    if not isinstance(indices_spec, DTensorSpec):
-        raise AssertionError(f"Expected DTensorSpec, got {type(indices_spec)}")
+    dim = normalize_dim(dim, len(values_meta.shape))
 
-    all_indices_spec: list[DTensorSpec | None] = [
-        indices_spec if dim == i else None for i in range(values_spec.ndim)
-    ]
+    strategies: list[list[Placement | _ShardingPlaceholder]] = []
 
-    result = prop_index(
-        OpSchema(
-            op=op_schema.op,
-            args_schema=(values_spec, all_indices_spec),
-            kwargs_schema=op_schema.kwargs_schema,
+    # Shard values on any non-indexed dim (output has same ndim)
+    for d in range(len(values_meta.shape)):
+        if d == dim:
+            continue
+        strategies.append(
+            [_ShardingPlaceholder(d), _ShardingPlaceholder(d), Replicate()]
         )
+
+    # Shard index → output sharded on the indexed dim
+    strategies.append([_ShardingPlaceholder(dim), Replicate(), _ShardingPlaceholder(0)])
+
+    # Partial passthrough from values
+    for reduce_op in Partial.ALL_REDUCE_OPS:
+        strategies.append([Partial(reduce_op), Partial(reduce_op), Replicate()])
+
+    return strategies
+
+
+@register_single_dim_strategy(
+    aten.index.Tensor, schema_info=RuntimeSchemaInfo(needs_pytree=True)
+)
+def index_single_dim_strategy(
+    op: OpOverload, args_schema: ArgsType, kwargs_schema: KwargsType
+) -> list[list[Placement | _ShardingPlaceholder]]:
+    values_meta, multi_indices_meta = args_schema
+    if not isinstance(values_meta, TensorMeta):
+        raise AssertionError(f"Expected TensorMeta, got {type(values_meta)}")
+    if not isinstance(multi_indices_meta, (list, tuple)):
+        raise AssertionError(f"Expected list or tuple, got {type(multi_indices_meta)}")
+
+    indexed_dims = [i for i, idx in enumerate(multi_indices_meta) if idx is not None]
+    non_indexed_dims = [
+        i for i in range(len(values_meta.shape)) if i not in set(indexed_dims)
+    ]
+
+    index_metas = [idx for idx in multi_indices_meta if idx is not None]
+    if not all(isinstance(m, TensorMeta) for m in index_metas):
+        raise AssertionError("Expected all index metas to be TensorMeta")
+    broadcast_ndim = max(len(m.shape) for m in index_metas)
+    num_indices = len(indexed_dims)
+
+    # Determine where index output dims are inserted in the result
+    all_consecutive = all(
+        indexed_dims[i + 1] - indexed_dims[i] == 1 for i in range(len(indexed_dims) - 1)
     )
-    if result.redistribute_schema:
-        schema_suggestion = result.redistribute_schema
-        result.redistribute_schema = OpSchema(
-            op=op_schema.op,
-            args_schema=(
-                schema_suggestion.args_schema[0],
-                dim,
-                schema_suggestion.args_schema[1][dim],  # type: ignore[index]
-            ),
-            kwargs_schema=op_schema.kwargs_schema,
-        )
-    return result
+    insert_dim = indexed_dims[0] if all_consecutive else 0
+
+    def values_dim_to_output_dim(d: int) -> int:
+        if d < insert_dim:
+            return d
+        return d + broadcast_ndim - sum(1 for idx_dim in indexed_dims if d > idx_dim)
+
+    strategies: list[list[Placement | _ShardingPlaceholder]] = []
+
+    # Shard values on a non-indexed dim, all indices replicated
+    for d in non_indexed_dims:
+        out_dim = values_dim_to_output_dim(d)
+        rule: list[Placement | _ShardingPlaceholder] = [_ShardingPlaceholder(out_dim)]
+        rule.append(_ShardingPlaceholder(d))
+        rule.extend([Replicate()] * num_indices)
+        strategies.append(rule)
+
+    # Shard indices on the same broadcast dim.  Each index tensor may
+    # have a different ndim, so we map broadcast dim → tensor dim via
+    # left-padding.  Tensors with size 1 on that dim are replicated
+    # (broadcast semantics).
+    for bd in range(broadcast_ndim):
+        per_tensor: list[tuple[int, int]] = []  # (tensor_dim, size)
+        for m in index_metas:
+            offset = broadcast_ndim - len(m.shape)
+            if bd < offset:
+                per_tensor.append((-1, 1))  # implicit broadcast
+            else:
+                td = bd - offset
+                per_tensor.append((td, m.shape[td]))
+        if all(s == 1 for _, s in per_tensor):
+            continue  # all broadcast-only, skip
+        out_dim = bd + insert_dim
+        rule: list[Placement | _ShardingPlaceholder] = [_ShardingPlaceholder(out_dim)]
+        rule.append(Replicate())
+        for td, s in per_tensor:
+            if s > 1:
+                rule.append(_ShardingPlaceholder(td))
+            else:
+                rule.append(Replicate())
+        strategies.append(rule)
+
+    # Partial passthrough from values
+    for reduce_op in Partial.LINEAR_REDUCE_OPS:
+        rule: list[Placement | _ShardingPlaceholder] = [
+            Partial(reduce_op),
+            Partial(reduce_op),
+        ]
+        rule.extend([Replicate()] * num_indices)
+        strategies.append(rule)
 
+    return strategies
 
-@register_op_strategy(
-    [
-        aten.index_put.default,
-        aten._index_put_impl_.default,
-    ],
+
+@register_single_dim_strategy(
+    [aten.index_put.default, aten.index_put_.default, aten._index_put_impl_.default],
     schema_info=RuntimeSchemaInfo(needs_pytree=True),
 )
-def prop_index_put(op_schema: OpSchema) -> StrategyType:
-    # We have 3 DTensor spec from argument `in`, `indices` and `values`
-    # accordingly.
-    in_spec, indices_spec, values_spec, *_ = op_schema.args_schema
-    if not isinstance(in_spec, OpStrategy):
-        raise AssertionError(f"Expected OpStrategy, got {type(in_spec)}")
-    # `indices`` is a tuple of scalar LongTensor, so we use TupleStrategy.
-    if not isinstance(indices_spec, TupleStrategy):
-        raise AssertionError(f"Expected TupleStrategy, got {type(indices_spec)}")
-    if not isinstance(values_spec, OpStrategy):
-        raise AssertionError(f"Expected OpStrategy, got {type(values_spec)}")
-    mesh = values_spec.mesh
-    op_strategy = OpStrategy([])
-    # 1. `indices` should all be replicated first.
-    indices_redistribute_costs = []
-    new_indices_spec: list[DTensorSpec | None] = []
-    for indices_spec_child in indices_spec.children:
-        if not isinstance(indices_spec_child, OpStrategy):
-            raise AssertionError(f"Expected OpStrategy, got {type(indices_spec_child)}")
-
-        replicated_spec = DTensorSpec(
-            mesh=mesh,
-            placements=tuple([Replicate()] * mesh.ndim),
-            tensor_meta=indices_spec_child.strategies[0].output_spec.tensor_meta,
-        )
-        new_indices_spec.append(replicated_spec)
-        child_costs = generate_redistribute_costs(indices_spec_child, replicated_spec)
-        indices_redistribute_costs.append(child_costs)
-
-    # 2. For placement rule of `values` and `in`, assume `values` shape =
-    # [a,b,c,d,e,f], `in` shape = [d,e,f]. Then `values`'s a,b,c (selected dim)
-    # must be replicated and d,e,f (nonselected dim) in both `values` and `in`
-    # should follow the same sharding (replicate or shard, but not partial).
-    size_offset = (
-        in_spec.strategies[0].output_spec.ndim
-        - values_spec.strategies[0].output_spec.ndim
-    )
-    # We can either let `values` follow `in`'s placements or reverse.
-    for exemplar_spec in [in_spec, values_spec]:
-        # use exemplar_spec as the target spec
-        for strategy in exemplar_spec.strategies:
-            in_spec_new_placements: list[Placement] = []
-            values_spec_new_placements: list[Placement] = []
-            placements = strategy.output_spec.placements
-            for placement in placements:
-                if placement.is_shard():
-                    if not isinstance(placement, Shard):
-                        raise AssertionError(f"Expected Shard, got {type(placement)}")
-                    if exemplar_spec is in_spec:
-                        # let `values_spce` follow `in_spec`
-                        if placement.dim < size_offset:
-                            # sharded on selected dim, need to change to replicate
-                            in_spec_new_placements.append(Replicate())
-                            values_spec_new_placements.append(Replicate())
-                        else:
-                            in_spec_new_placements.append(placement)
-                            values_spec_new_placements.append(
-                                Shard(placement.dim - size_offset)
-                            )
-                    else:
-                        # let `in_spec` follow `values_spec`
-                        in_spec_new_placements.append(
-                            Shard(placement.dim + size_offset)
-                        )
-                        values_spec_new_placements.append(placement)
-                else:
-                    in_spec_new_placements.append(Replicate())
-                    values_spec_new_placements.append(Replicate())
-            new_in_spec = DTensorSpec(
-                mesh=mesh,
-                placements=tuple(in_spec_new_placements),
-                tensor_meta=in_spec.strategies[0].output_spec.tensor_meta,
-            )
-            new_values_spec = DTensorSpec(
-                mesh=mesh,
-                placements=tuple(values_spec_new_placements),
-                tensor_meta=values_spec.strategies[0].output_spec.tensor_meta,
-            )
-            output_spec = DTensorSpec(
-                mesh=mesh,
-                placements=tuple(in_spec_new_placements),
-                tensor_meta=in_spec.strategies[0].output_spec.tensor_meta,
-            )
-            cost_in_spec = generate_redistribute_costs(in_spec, new_in_spec)
-            cost_values_spec = generate_redistribute_costs(values_spec, new_values_spec)
-            op_strategy.strategies.append(
-                OpSpec(
-                    # pyrefly: ignore [bad-argument-type]
-                    input_specs=(
-                        new_in_spec,
-                        *new_indices_spec,  # type: ignore[arg-type]
-                        new_values_spec,
-                    ),
-                    output_specs=output_spec,
-                    redistribute_cost=[
-                        cost_in_spec,
-                        *indices_redistribute_costs,
-                        cost_values_spec,
-                    ],
-                )
-            )
-    return op_strategy
+def index_put_single_dim_strategy(
+    op: OpOverload, args: ArgsType, kwargs: KwargsType
+) -> list[list[Placement | _ShardingPlaceholder]]:
+    """Single-dim sharding strategy for index_put(self, indices, values).
 
+    Strategy format: [output, input, *indices, value]
 
-@register_prop_rule(aten.index.Tensor, schema_info=RuntimeSchemaInfo(needs_pytree=True))
-def prop_index(op_schema: OpSchema) -> OutputSharding:
-    """
-    Expect replicated on the first input; _mostly_ pointwise on the second input.
+    How index_put works:
+
+      indices is a tuple of index tensors and Nones:
+      - an index tensor at entry i means self is indexed on dim i.
+      - a None at entry i means all elements along dim i are selected (like :).
+      - any trailing dims (if self.ndim > len(indices)) are also not indexed
+        (i.e. implicit trailing Nones).
+
+      All non-None index tensors are broadcast together to produce a
+      broadcasted indexing shape. Each position in this broadcasted shape
+      serves as an indexing coordinate into self. Each coordinate selects a
+      tensor element, or a slice (if non-indexed dims exist).
+
+      values is a tensor broadcastable to the indexing output shape.
+      When indexed dims are consecutive starting at dim k, this shape is
+      (*self[:k], *broadcast_shape, *self[k+n_indexed:]). When indexed
+      dims are non-consecutive, it is (*broadcast_shape, *non_indexed_dims).
+
+    Sharding rules (possibly conservative and incomplete):
+      - Index tensors: always Replicate (every rank needs all coordinates).
+      - Self cannot be sharded on indexed dims (local position != global position).
+      - Self and values CAN be sharded on non-indexed dims.
+        The exception is broadcasted value dimensions (size 1) - we require Replicate, but can shard self.
+      - Additionally, we allow the full Partial rule on non-indexing tensors.
 
-    TODO: exception: when the dtype of second input is "bool", then a torch.nonzero needs to be triggered first.
     """
-    # Current sharding constraints:
-    # For values:
-    #   1. We currently require that the dimension of values_spec be replicated or partial
-    #      if they are being indexed on.
-    #   2. Other dimensions of values_spec can remain sharded if they are so.
-    # For indices:
-    #   Indices can be either sharded or replicated. All index tensors need to be sharded
-    #   in a compatible way, following the pointwise rule (including resolving Partial
-    #   into either sharded or replicated)
-
-    values_spec, multi_indices_spec = op_schema.args_schema
-    if not isinstance(values_spec, DTensorSpec):
-        raise AssertionError(f"Expected DTensorSpec, got {type(values_spec)}")
-    if not isinstance(multi_indices_spec, list):
-        raise AssertionError(f"Expected list, got {type(multi_indices_spec)}")
-    multi_indices_spec = cast(list[DTensorSpec | None], multi_indices_spec)
-    valid_indices_spec: list[tuple[int, DTensorSpec]] = [
-        (i, a) for i, a in enumerate(multi_indices_spec) if a is not None
-    ]
+    self_meta = cast(TensorMeta, args[0])
+    indices_meta = cast(tuple[TensorMeta | None, ...], args[1])
+    values_meta = cast(TensorMeta, args[2])
+
+    # Determine indexed vs non-indexed dims of self.
+    indexed_dims = {i for i, idx in enumerate(indices_meta) if idx is not None}
+    non_indexed_dims = [d for d in range(len(self_meta.shape)) if d not in indexed_dims]
+    n_indexed = len(indexed_dims)
+    values_ndim = len(values_meta.shape)
+
+    # Explicitly compute the broadcast shape of the index tensors.
+    index_shapes = [idx.shape for idx in indices_meta if idx is not None]
+    broadcast_ndim = len(torch.broadcast_shapes(*index_shapes)) if index_shapes else 0
+
+    # Strategy format: [output, input, *indices, value]
+    # The infra flattens the indices list and drops None entries, so only
+    # non-None index tensors get a placement slot (all Replicate).
+    #
+    # Values dim mapping depends on whether indexed dims are contiguous:
+    #   Contiguous (e.g., (None, idx0, idx1)): broadcast replaces indexed block in-place.
+    #     values shape = (*non_indexed_before, *broadcast_shape, *non_indexed_after)
+    #   Non-contiguous (e.g., (idx0, None, idx1)): broadcast goes to front.
+    #     values shape = (*broadcast_shape, *non_indexed_dim_sizes)
+    indexed_dims_sorted = sorted(indexed_dims)
+    contiguous_indexed = len(indexed_dims_sorted) <= 1 or (
+        indexed_dims_sorted[-1] - indexed_dims_sorted[0] + 1 == len(indexed_dims_sorted)
+    )
 
-    # 1. All indices have to be sharded equally. Moreover, indices can be broadcast.
-    #    Here, we piggyback on the pointwise sharding rule for indices.
-    indices_out = pointwise_rule(
-        OpSchema(
-            op=op_schema.op,
-            args_schema=tuple(v[1] for v in valid_indices_spec),
-            kwargs_schema={},
+    strategies: list[list[Placement | _ShardingPlaceholder]] = []
+    for i, self_dim in enumerate(non_indexed_dims):
+        if contiguous_indexed and indexed_dims_sorted:
+            # Broadcast replaces the indexed block in-place.
+            first_indexed = indexed_dims_sorted[0]
+            if self_dim < first_indexed:
+                values_dim = self_dim
+            else:
+                values_dim = self_dim - n_indexed + broadcast_ndim
+        else:
+            # Broadcast goes to front (non-contiguous or no indexed dims).
+            values_dim = broadcast_ndim + i
+
+        # values_dim is the position in the result tensor, but values may
+        # have fewer dims (right-aligned broadcasting). Convert to the
+        # actual values tensor dimension.
+        result_ndim = broadcast_ndim + len(non_indexed_dims)
+        values_tensor_dim = values_dim - (result_ndim - values_ndim)
+
+        if values_tensor_dim < 0:
+            values_placement: Placement | _ShardingPlaceholder = Replicate()
+        elif values_meta.shape[values_tensor_dim] == 1:
+            values_placement = Replicate()
+        else:
+            values_placement = _ShardingPlaceholder(values_tensor_dim)
+
+        strategies.append(
+            [
+                _ShardingPlaceholder(self_dim),
+                _ShardingPlaceholder(self_dim),
+                *([Replicate()] * n_indexed),
+                values_placement,
+            ]
         )
+
+    # full-partial rule on non-indexing tensors
+    strategies.append(
+        [
+            Partial(),
+            Partial(),
+            *([Replicate()] * n_indexed),
+            Partial(),
+        ]
     )
-    need_reshard_on_indices = indices_out.output_spec is None
+    return strategies
 
-    if not need_reshard_on_indices:
-        # this means that our inputs are already sharded properly and we will use that as our indices_spec
-        if not isinstance(indices_out.output_spec, DTensorSpec):
-            raise AssertionError(
-                f"Expected DTensorSpec, got {type(indices_out.output_spec)}"
-            )
-        indices_spec: DTensorSpec = indices_out.output_spec
-    else:
-        if indices_out.redistribute_schema is None:
-            raise AssertionError("redistribute_schema should not be None")
-        valid_indices_suggestion = indices_out.redistribute_schema
-        for i, v in enumerate(valid_indices_suggestion.args_spec):
-            multi_indices_spec[valid_indices_spec[i][0]] = v
-        # we'll need to call pointwise_rule again to see what's our ideal indices_spec and then
-        # use that to compute our ideal values_spec
-        indices_output_spec = pointwise_rule(valid_indices_suggestion).output_spec
-        if not isinstance(indices_output_spec, DTensorSpec):
-            raise AssertionError(
-                f"Expected DTensorSpec, got {type(indices_output_spec)}"
-            )
-        indices_spec = indices_output_spec
 
-    lookup_dims = {v[0] for v in valid_indices_spec}
+def _index_dim_strategy(
+    args_schema: ArgsType,
+    shard_row: Callable[[int], list[Placement | _ShardingPlaceholder]],
+    partial_rules: list[list[Placement | _ShardingPlaceholder]] | None = None,
+) -> list[list[Placement | _ShardingPlaceholder]]:
+    """Common strategy for index ops that shard on all dims except the indexed dim.
 
-    need_reshard_on_values = tuple(
-        (isinstance(vp, Shard) and (vp.dim in lookup_dims or isinstance(ip, Shard)))
-        for vp, ip in zip(values_spec.placements, indices_spec.placements)
+    Args:
+        shard_row: given a dim d, returns the strategy row for sharding on that dim.
+        partial_rules: additional Partial passthrough strategies.
+    """
+    self_meta = cast(TensorMeta, args_schema[0])
+    ndim = len(self_meta.shape)
+    dim = normalize_dim(cast(int, args_schema[1]), ndim)
+    strategies: list[list[Placement | _ShardingPlaceholder]] = []
+    for d in range(ndim):
+        if d != dim:
+            strategies.append(shard_row(d))
+    if partial_rules:
+        strategies.extend(partial_rules)
+    return strategies
+
+
+@register_single_dim_strategy(
+    [aten.index_fill.int_Scalar, aten.index_fill_.int_Scalar],
+    schema_info=RuntimeSchemaInfo(1),
+)
+def index_fill_scalar_single_dim_strategy(
+    op: OpOverload, args_schema: ArgsType, kwargs_schema: KwargsType
+) -> list[list[Placement | _ShardingPlaceholder]]:
+    # index_fill(self, dim, index, value) — fills self[..., index, ...] with scalar value.
+    # Partial rules: each rank fills with the same scalar v, then reduces.
+    # Only idempotent reduces work: avg(v,v,...,v)=v, max(v,v,...,v)=v, min(v,v,...,v)=v.
+    # sum and product fail: sum(v,v,...,v)=nv, product(v,v,...,v)=v^n.
+    return _index_dim_strategy(
+        args_schema,
+        lambda d: [
+            _ShardingPlaceholder(d),  # result
+            _ShardingPlaceholder(d),  # self
+            Replicate(),  # value (scalar, same on all ranks)
+        ],
+        [[Partial(op), Partial(op), Replicate()] for op in ("avg", "max", "min")],
     )
 
-    if not need_reshard_on_indices and not any(need_reshard_on_values):
-        value_placements = values_spec.placements
 
-        all_dims_consecutive = all(
-            b[0] - a[0] == 1
-            for b, a in zip(valid_indices_spec[1:], valid_indices_spec[:-1])
-        )
-        if all_dims_consecutive:
-            # if all index vectors are consecutives, insert at the dimension of the first index
-            insert_dim: int = valid_indices_spec[0][0]
-        else:
-            # else, insert on the first dimension
-            insert_dim = 0
-
-        def place(vp: Placement, ip: Placement) -> Placement:
-            if isinstance(vp, Shard):
-                return Shard(
-                    vp.dim
-                    if vp.dim < insert_dim
-                    # accounts for the offset in output dimensions
-                    else vp.dim
-                    + indices_spec.ndim
-                    - sum(1 if vp.dim > v[0] else 0 for v in valid_indices_spec)
-                )
-            if isinstance(ip, Shard):
-                return Shard(ip.dim + insert_dim)
-            # Partial or Replicated
-            return vp
-
-        value_placements = tuple(
-            place(vp, ip)
-            for vp, ip in zip(values_spec.placements, indices_spec.placements)
-        )
-        result = OutputSharding(
-            output_spec=DTensorSpec(
-                mesh=values_spec.mesh,
-                placements=value_placements,
-            )
-        )
-        return result
-    else:
-        result = OutputSharding(
-            output_spec=None,
-            redistribute_schema=OpSchema(
-                op=op_schema.op,
-                args_schema=(
-                    DTensorSpec(
-                        mesh=values_spec.mesh,
-                        placements=tuple(
-                            Replicate() if need_reshard_on_values[i] else v
-                            for i, v in enumerate(values_spec.placements)
-                        ),
-                        tensor_meta=values_spec.tensor_meta,
-                    ),
-                    multi_indices_spec,
-                ),
-                kwargs_schema=op_schema.kwargs_schema,
-            ),
-        )
-        return result
+@register_single_dim_strategy(
+    [aten.index_fill.int_Tensor, aten.index_fill_.int_Tensor],
+    schema_info=RuntimeSchemaInfo(1),
+)
+def index_fill_tensor_single_dim_strategy(
+    op: OpOverload, args_schema: ArgsType, kwargs_schema: KwargsType
+) -> list[list[Placement | _ShardingPlaceholder]]:
+    # index_fill(self, dim, index, value) — fills self[..., index, ...] with 0-d tensor value.
+    # Partial rules: each rank fills with its partial value v_i, then reduces.
+    # All reduce ops work because reduce(v_0, ..., v_{n-1}) = V (the global value)
+    # regardless of op, since fill is a pure replacement (no mixing with self).
+    return _index_dim_strategy(
+        args_schema,
+        lambda d: [
+            _ShardingPlaceholder(d),  # result
+            _ShardingPlaceholder(d),  # self
+            Replicate(),  # index
+            Replicate(),  # value
+        ],
+        [
+            [Partial(op), Partial(op), Replicate(), Partial(op)]
+            for op in Partial.ALL_REDUCE_OPS
+        ],
+    )
+
+
+@register_single_dim_strategy(
+    [aten.index_reduce.default, aten.index_reduce_.default],
+    schema_info=RuntimeSchemaInfo(1),
+)
+def index_reduce_single_dim_strategy(
+    op: OpOverload, args_schema: ArgsType, kwargs_schema: KwargsType
+) -> list[list[Placement | _ShardingPlaceholder]]:
+    # index_reduce(self, dim, index, source, reduce) — reduces source into self at index positions.
+    # No partial rules: reduce ops are "mean"/"amax"/"amin"/"prod", which don't match
+    # any Partial reduce op names ("avg"/"max"/"min"/"product"/"sum").
+    return _index_dim_strategy(
+        args_schema,
+        lambda d: [
+            _ShardingPlaceholder(d),  # result
+            _ShardingPlaceholder(d),  # self
+            Replicate(),  # index
+            _ShardingPlaceholder(d),  # source
+        ],
+    )
 
 
 @register_op_strategy(
@@ -1330,9 +1453,8 @@ def eye_out_strategy(op_schema: OpSchema) -> OpStrategy:
     # eye.m_out has signature: eye(int n, int m, *, Tensor(a!) out) -> Tensor(a!)
     # The out kwarg is a DTensor that determines the sharding
     out_spec = op_schema.kwargs_schema["out"]
-    assert isinstance(out_spec, OpStrategy), (
-        f"Expected OpStrategy for out, got {type(out_spec)}"
-    )
+    if not isinstance(out_spec, OpStrategy):
+        raise AssertionError(f"Expected OpStrategy for out, got {type(out_spec)}")
 
     return OpStrategy(
         [
@@ -1344,3 +1466,66 @@ def eye_out_strategy(op_schema: OpSchema) -> OpStrategy:
             for strategy in out_spec.strategies
         ]
     )
+
+
+def _pass_through_partials(
+    num_inputs: int = 1,
+) -> list[list[Placement | _ShardingPlaceholder]]:
+    """Pass-through strategies for all supported reduce ops."""
+    return [[Partial(op)] * (1 + num_inputs) for op in ("sum", "avg", "max", "min")]
+
+
+def _shard_inactive_dims(
+    ndim: int, active_dims: set[int], num_inputs: int = 1
+) -> list[list[Placement | _ShardingPlaceholder]]:
+    """Single-dim strategies: shard on dims the op doesn't touch."""
+    strategies: list[list[Placement | _ShardingPlaceholder]] = []
+    for d in range(ndim):
+        if d not in active_dims:
+            strategies.append([_ShardingPlaceholder(d)] * (1 + num_inputs))
+    return strategies
+
+
+@register_single_dim_strategy(aten.roll.default, schema_info=RuntimeSchemaInfo(1))
+def roll_single_dim_strategy(
+    op: OpOverload, args_schema: ArgsType, kwargs_schema: KwargsType
+) -> list[list[Placement | _ShardingPlaceholder]]:
+    input_meta = args_schema[0]
+    if not isinstance(input_meta, TensorMeta):
+        raise AssertionError(f"Expected TensorMeta, got {type(input_meta)}")
+    ndim = len(input_meta.shape)
+    raw_dims = cast(list[int], args_schema[2]) if len(args_schema) > 2 else []
+    # When dims is empty, roll flattens the tensor — all dims are active
+    if not raw_dims:
+        raw_dims = list(range(ndim))
+    active_dims = {normalize_dim(d, ndim) for d in raw_dims}
+    return _shard_inactive_dims(ndim, active_dims) + _pass_through_partials()
+
+
+@register_single_dim_strategy(aten.flip.default, schema_info=RuntimeSchemaInfo(1))
+def flip_single_dim_strategy(
+    op: OpOverload, args_schema: ArgsType, kwargs_schema: KwargsType
+) -> list[list[Placement | _ShardingPlaceholder]]:
+    input_meta = args_schema[0]
+    if not isinstance(input_meta, TensorMeta):
+        raise AssertionError(f"Expected TensorMeta, got {type(input_meta)}")
+    ndim = len(input_meta.shape)
+    raw_dims = cast(list[int], args_schema[1])
+    active_dims = {normalize_dim(d, ndim) for d in raw_dims}
+    return _shard_inactive_dims(ndim, active_dims) + _pass_through_partials()
+
+
+@register_single_dim_strategy(
+    [aten._fft_c2c.default, aten._fft_r2c.default, aten._fft_c2r.default],
+    schema_info=RuntimeSchemaInfo(1),
+)
+def fft_single_dim_strategy(
+    op: OpOverload, args_schema: ArgsType, kwargs_schema: KwargsType
+) -> list[list[Placement | _ShardingPlaceholder]]:
+    input_meta = args_schema[0]
+    if not isinstance(input_meta, TensorMeta):
+        raise AssertionError(f"Expected TensorMeta, got {type(input_meta)}")
+    ndim = len(input_meta.shape)
+    raw_dims = cast(list[int], args_schema[1])
+    active_dims = {normalize_dim(d, ndim) for d in raw_dims}
+    return _shard_inactive_dims(ndim, active_dims) + _pass_through_partials()
diff --git a/torch/distributed/tensor/_ops/_view_ops.py b/torch/distributed/tensor/_ops/_view_ops.py
index 6ca4c9c9925dd..29d2c3d9e35d8 100644
--- a/torch/distributed/tensor/_ops/_view_ops.py
+++ b/torch/distributed/tensor/_ops/_view_ops.py
@@ -1,13 +1,14 @@
 # mypy: allow-untyped-defs
 # Copyright (c) Meta Platforms, Inc. and affiliates
+import math
 from collections.abc import Callable, Iterable, Sequence
 from dataclasses import dataclass
-from typing import cast
+from typing import cast, NamedTuple
 
 import torch
 from torch import Tensor
 from torch._prims_common import DimsType
-from torch.distributed.tensor._dtensor_spec import DTensorSpec
+from torch.distributed.tensor._dtensor_spec import DTensorSpec, TensorMeta
 from torch.distributed.tensor._op_schema import (
     OpSchema,
     OpSpec,
@@ -15,6 +16,10 @@
     RuntimeSchemaInfo,
     StrategyType,
 )
+from torch.distributed.tensor._ops.single_dim_strategy import (
+    _ShardingPlaceholder,
+    register_single_dim_strategy,
+)
 from torch.distributed.tensor._ops.utils import (
     generate_redistribute_costs,
     normalize_dim,
@@ -24,6 +29,7 @@
 )
 from torch.distributed.tensor.placement_types import (
     _StridedShard,
+    Partial,
     Placement,
     Replicate,
     Shard,
@@ -35,6 +41,13 @@
 Shape = tuple[int, ...]
 
 
+class ClaimedDim(NamedTuple):
+    """An (input_dim, output_dim) pair claimed by a mesh dim's _StridedShard rewrite."""
+
+    input_dim: int
+    output_dim: int
+
+
 @dataclass
 class DimSpec:
     """Specifies how an output dimension maps to an input dimension."""
@@ -52,12 +65,27 @@ class Singleton(DimSpec):
     """Output dimension is a singleton."""
 
 
-@dataclass
+@dataclass(eq=False)
 class InputDim(DimSpec):
     """Output dimension maps directly to an input dimension."""
 
     input_dim: int
 
+    def __eq__(self, other: object) -> bool:
+        """Raises TypeError for non-DimSpec comparisons to catch accidental
+        ``shard.dim == input_dim`` bugs where ``.input_dim`` was intended."""
+        if isinstance(other, InputDim):
+            return self.input_dim == other.input_dim
+        if not isinstance(other, DimSpec):
+            raise TypeError(
+                f"Cannot compare InputDim with {type(other).__name__}. "
+                f"Did you mean to use .input_dim?"
+            )
+        return NotImplemented
+
+    def __hash__(self) -> int:
+        return hash((InputDim, self.input_dim))
+
 
 @dataclass
 class Broadcast(DimSpec):
@@ -82,7 +110,9 @@ class NewDim(DimSpec):
 
     @classmethod
     def new(cls, size: int) -> DimSpec:
-        return Singleton() if size == 1 else NewDim(size)
+        from torch.fx.experimental.symbolic_shapes import guard_or_false
+
+        return Singleton() if guard_or_false(size == 1) else NewDim(size)
 
 
 @dataclass
@@ -94,7 +124,9 @@ class Repeat(DimSpec):
 
     @classmethod
     def new(cls, dim: DimSpec, times: int) -> DimSpec:
-        if times == 1:
+        from torch.fx.experimental.symbolic_shapes import guard_or_false
+
+        if guard_or_false(times == 1):
             return dim
         elif isinstance(dim, Singleton):
             # repeating a singleton is the same as broadcasting it
@@ -141,6 +173,8 @@ class Split(DimSpec):
 
     @classmethod
     def new(cls, dim: DimSpec, group_shape: tuple[int, ...], idx: int) -> DimSpec:
+        from torch.fx.experimental.symbolic_shapes import guard_or_false, guard_or_true
+
         if not len(group_shape) > 0:
             raise AssertionError(
                 f"Expected group_shape length > 0, got {len(group_shape)}"
@@ -150,13 +184,15 @@ def new(cls, dim: DimSpec, group_shape: tuple[int, ...], idx: int) -> DimSpec:
             if not idx == 0:
                 raise AssertionError(f"Expected idx == 0, got {idx}")
             return dim
-        elif group_shape[idx] == 1:
+        elif guard_or_false(group_shape[idx] == 1):
             return Singleton()
         else:
             # remove singletons from group
             # group_mapping = [(new_index, (shape, old_index)) ...]
             group_mapping = list(
-                enumerate((s, i) for i, s in enumerate(group_shape) if s != 1)
+                enumerate(
+                    (s, i) for i, s in enumerate(group_shape) if guard_or_true(s != 1)
+                )
             )
             new_group_shape = tuple(m[1][0] for m in group_mapping)
             new_idx = next(filter(lambda x: x[1][1] == idx, group_mapping))[0]
@@ -185,6 +221,8 @@ def dim_atleast_3d(ndim: int) -> DimMap:
 
 def expand(input_shape: Shape, shape: Shape) -> DimMap:
     """Implement broadcast on multiple dimensions."""
+    from torch.fx.experimental.symbolic_shapes import guard_or_false
+
     if not len(shape) >= len(input_shape):
         raise AssertionError(
             f"Expected len(shape) >= len(input_shape), got {len(shape)} < {len(input_shape)}"
@@ -203,21 +241,29 @@ def expand(input_shape: Shape, shape: Shape) -> DimMap:
             if not isinstance(p, InputDim):
                 raise AssertionError(f"DimSpec not supported in expand: {p}")
             actual_s = input_shape[p.input_dim]
-            if not (actual_s == 1 or desired_s == -1 or desired_s == actual_s):
+            if not (
+                guard_or_false(actual_s == 1)
+                or guard_or_false(desired_s == -1)
+                or guard_or_false(desired_s == actual_s)
+            ):
                 raise AssertionError(
                     f"Expected actual_s == 1 or desired_s == -1 or "
                     f"desired_s == actual_s, got actual_s={actual_s}, desired_s={desired_s}"
                 )
         mapping.append(
             p
-            if desired_s in (1, -1) or desired_s == actual_s
+            if (
+                guard_or_false(desired_s == 1)
+                or guard_or_false(desired_s == -1)
+                or guard_or_false(desired_s == actual_s)
+            )
             else Broadcast.new(p, desired_s)
         )
     return tuple(mapping)
 
 
 def normalize_sizes(sizes: Shape | tuple[Shape]) -> Shape:
-    if isinstance(sizes[0], int):
+    if isinstance(sizes[0], (int, torch.SymInt)):
         return cast(Shape, sizes)
     elif len(sizes) == 1:
         return sizes[0]
@@ -297,20 +343,24 @@ def infer_size(total_size: int, sizes: Shape) -> Shape:
 
     Infer the size of this dimension given the total_size.
     """
-    infers = [i for i, s in enumerate(sizes) if s == -1]
+    from torch.fx.experimental.symbolic_shapes import guard_or_false
+
+    infers = [i for i, s in enumerate(sizes) if guard_or_false(s == -1)]
     size = prod(sizes)
     if not len(infers) <= 1:
         raise AssertionError("can only infer one size")
     if infers:
         size = -size
         missing_size = total_size // size
-        if not total_size % size == 0:
-            raise AssertionError(
-                f"size inferred for -1 is not integral {sizes} should have {total_size} elements."
-            )
-        return tuple(s if s != -1 else missing_size for s in sizes)
-    if not size == total_size:
-        raise AssertionError(f"sizes do not match {total_size} vs {size}")
+        torch._check(
+            total_size % size == 0,
+            lambda: f"size inferred for -1 is not integral {sizes} should have {total_size} elements.",
+        )
+        return tuple(s if not guard_or_false(s == -1) else missing_size for s in sizes)
+    torch._check(
+        size == total_size,
+        lambda: f"sizes do not match {total_size} vs {size}",
+    )
     return sizes
 
 
@@ -343,18 +393,22 @@ def view_groups(from_size: Shape, to_size: Shape) -> DimMap:
     - in the above, input is flattened into a single dimension and then split
       into two separate dimensions with different sizes from the input.
     """
+    from torch.fx.experimental.symbolic_shapes import guard_or_false, guard_or_true
+
     from_nelem = prod(from_size)
     to_size = infer_size(from_nelem, normalize_sizes(to_size))
 
-    if not from_nelem == prod(to_size):
-        raise AssertionError("Total view shape does not add up")
+    torch._check(
+        from_nelem == prod(to_size),
+        lambda: "Total view shape does not add up",
+    )
 
     from_idx = 0
     to_idx = 0
     from_len = len(from_size)
     to_len = len(to_size)
 
-    result_pp = []
+    result_pp: list[DimSpec] = []
 
     while from_idx < from_len or to_idx < to_len:
         from_group_dim, to_group_shape = [], []
@@ -374,18 +428,20 @@ def view_groups(from_size: Shape, to_size: Shape) -> DimMap:
             to_idx += 1
 
         # if any of the groups is singleton, great, we need to backtrack though
-        if f == 1 and t != 1:
+        if guard_or_false(f == 1) and guard_or_true(t != 1):
             # produces ([1], [])
             to_idx -= 1
             to_group_shape = []
-        elif f != 1 and t == 1:
+        elif guard_or_true(f != 1) and guard_or_false(t == 1):
             # produces ([], [1])
             from_idx -= 1
             from_group_dim = []
         else:
             # produces ([1], [1]),  ([2], [2]), ([2,3], [6])
-            while f != t:
-                if f < t:
+            while guard_or_true(f != t):
+                if (
+                    t % f == 0 or t > f
+                ):  # for easier symbolic comparisons, e.g. u0*u1 > u0
                     nf = from_size[from_idx]
                     from_group_dim.append(from_idx)
                     from_idx += 1
@@ -398,7 +454,11 @@ def view_groups(from_size: Shape, to_size: Shape) -> DimMap:
 
         if len(to_group_shape) > 0:
             flattened = Flatten.new(
-                tuple(InputDim(fi) for fi in from_group_dim if from_size[fi] >= 1)
+                tuple(
+                    InputDim(fi)
+                    for fi in from_group_dim
+                    if guard_or_true(from_size[fi] >= 1)
+                )
             )
             result_pp += [
                 Split.new(flattened, tuple(to_group_shape), i)
@@ -428,15 +488,22 @@ def dim_transpose(ndim: int, dim1: int, dim2: int) -> DimMap:
     return tuple(dimmap)
 
 
-def dim_squeeze(shape: Shape, dim: int | None = None) -> DimMap:
-    # FIXME: this is wrong when dim=None and one of the dimensions
-    # equals size of the mesh. For example squeeze(DTensor(tensor(4), Shard[0])) could
-    # end up as squeeze(tensor(1)) if we have 4 devices; this would lead to
-    # removal of a dimension that is not actually a singleton.
+def dim_squeeze(shape: Shape, dim: DimsType | None = None) -> DimMap:
+    # Operates on local shape; sharding_prop rewrites squeeze ops to squeeze.dims
+    # with only globally-singleton dims before this is called.
+    from torch.fx.experimental.symbolic_shapes import guard_or_true
+
+    ndim = len(shape)
+    if dim is None:
+        target_dims = set(range(ndim))
+    elif isinstance(dim, int):
+        target_dims = {normalize_dim(dim, ndim)}
+    else:
+        target_dims = set(normalize_dims(dim, ndim))
     return tuple(
         InputDim(i)
         for i, s in enumerate(shape)
-        if s > 1 or (dim is not None and i != normalize_dim(dim, len(shape)))
+        if guard_or_true(s > 1) or i not in target_dims
     )
 
 
@@ -514,183 +581,667 @@ def propagate_shape_and_sharding(
 
     Sharding propagation follows mapped dimensions:
     - An output dimension that maps directly to an input dimension is sharded equally
-    - An output dimension that is a flattened set of input dimensions can only be
-      sharded if only the leftmost flattened dimension is sharded.
+    - An output dimension that is a flattened set of input dimensions can be sharded:
+      the first sharded dim stays as Shard, non-first sharded dims become _StridedShard
     - An output dimension that is a split of the input dimension can only be sharded
       if the leftmost split size is divisible by the mesh dimension
     """
-    if not len(input_src_placements) == len(mesh_sizes):
-        raise AssertionError(f"{input_src_placements} != {mesh_sizes}")
-    # for each input dim, for each mesh dim, provides a list of possible shardable dimensions
-    mesh_ndim = len(mesh_sizes)
-    shardable_dims: dict[int, list[bool]] = {}
-
-    # in case an input dimension disappears (e.g. collapsing, reduction)
-    # we cannot shard in that dimension (we need a replication fall-back rule)
-    seen_input_dims: set[int] = set()
+    propagator = _ViewShardingPropagator(
+        input_src_placements, global_input_shape, rule, mesh_sizes, strict_view
+    )
+    input_tgt_placements, input_to_output_tensor_dims = propagator.analyze()
+    output_placements = propagator.rewrite_output_placements(
+        input_tgt_placements, input_to_output_tensor_dims
+    )
+    return input_tgt_placements, output_placements
 
-    def collect_used_inputs(cmd: DimSpec) -> None:
-        if isinstance(cmd, InputDim):
-            seen_input_dims.add(cmd.input_dim)
-        for inp in cmd.inputs():
-            collect_used_inputs(inp)
 
-    for cmd in rule:
-        collect_used_inputs(cmd)
-    for dim in range(len(global_input_shape)):
-        shardable_dims[dim] = [dim in seen_input_dims] * mesh_ndim
+class _ViewShardingPropagator:
+    """Two-phase sharding propagator for view ops.
+
+    Phase 1 — ``analyze()``:
+      Walks the DimMap rule and returns:
+      - ``input_tgt_placements``: input placements with unshardable dims
+        demoted to Replicate.
+      - ``input_to_output_tensor_dims``: maps each input tensor dim to its
+        output dim(s).  Cardinality encodes the op type: 1→1 for InputDim,
+        N→1 for Flatten, 1→N for Split/unflatten.
+
+    Phase 2 — ``rewrite_output_placements()``:
+      Consumes both Phase 1 outputs.  Iterates mesh dims 0..n-1, maintaining:
+      - ``strided_shard_claimed_dims``: (input_dim, output_dim) pairs already assigned
+        to a mesh dim by _StridedShard rewriting.
+      - ``local_tensor_shapes``: global shape progressively divided by each
+        mesh dim's shard size.
+      For each surviving Shard/_StridedShard, looks up the output dim(s) and
+      produces the final output placement.
+    """
 
-    def maybe_get_shard_mesh_dim_and_placement(
-        input_dim: InputDim,
-    ) -> tuple[int | None, Shard | _StridedShard | None]:
-        # if input_dim is sharded, return the mesh_dim and shard placement
-        for i, placement in enumerate(input_src_placements):
+    def __init__(
+        self,
+        input_src_placements: Sequence[Placement],
+        global_input_shape: Shape,
+        rule: DimMap,
+        mesh_sizes: Shape,
+        strict_view: bool,
+    ) -> None:
+        self.input_src_placements = input_src_placements
+        self.global_input_shape = global_input_shape
+        self.rule = rule
+        self.mesh_sizes = mesh_sizes
+        self.strict_view = strict_view
+        self.mesh_ndim = len(mesh_sizes)
+
+        # shard_allowed[input_dim][mesh_dim]: whether input_dim can stay
+        # sharded on mesh_dim.  Populated by _analyze_dim and its helpers.
+        self.shard_allowed: dict[int, list[bool]] = {}
+        # Mesh dims whose _StridedShard has already been matched to an output dim.
+        # Populated by _analyze_split.
+        self.matched_strided_mesh_dims: set[int] = set()
+
+    # ------------------------------------------------------------------
+    # Public API: analyze → rewrite_output_placements
+    # ------------------------------------------------------------------
+
+    def analyze(
+        self,
+    ) -> tuple[Sequence[Placement], dict[int, list[int]]]:
+        """Phase 1: walk the DimMap rule, return (input_tgt_placements, input_to_output_tensor_dims)."""
+        input_dims_in_rule = self._input_dims_in_rule(self.rule)
+
+        # Default: shardable if the dim appears in the rule. Refined by _analyze_*.
+        for dim in range(len(self.global_input_shape)):
+            self.shard_allowed[dim] = [dim in input_dims_in_rule] * self.mesh_ndim
+
+        # Walk the rule to refine shard_allowed and build input_to_output_tensor_dims.
+        #
+        # Flatten example: view([2, 3, 4], [6, 4])
+        #   rule = (Flatten(InputDim(0), InputDim(1)), InputDim(2))
+        #   output_dim=0 (Flatten): hits the isinstance(cmd, Flatten) branch.
+        #     Maps input dims 0 and 1 to output dim 0.  Result: {0: [0], 1: [0]}
+        #   output_dim=1 (InputDim(2)): hits the len(in_dims) > 0 branch.
+        #     Maps input dim 2 to output dim 1.  Result: {0: [0], 1: [0], 2: [1]}
+        #
+        # Split example: view([6], [2, 3])
+        #   rule = (Split(InputDim(0), (2,3), 0), Split(InputDim(0), (2,3), 1))
+        #   output_dim=0 (split_id=0): hits the len(in_dims) > 0 branch.
+        #     Maps input dim 0 to output dim 0.  Result: {0: [0]}
+        #   output_dim=1 (split_id=1): hits the isinstance(cmd, Split) branch
+        #     because _analyze_split returns [] for split_id>0.  Chases root
+        #     InputDim(0) and appends output dim 1.  Result: {0: [0, 1]}
+        input_to_output_tensor_dims: dict[int, list[int]] = {}
+        for output_dim, cmd in enumerate(self.rule):
+            in_dims = self._analyze_dim(cmd)
+            if isinstance(cmd, Flatten):
+                for in_dim in in_dims:
+                    if in_dim.input_dim in input_to_output_tensor_dims:
+                        raise AssertionError(
+                            f"Input dim {in_dim.input_dim} already mapped to output dims "
+                            f"{input_to_output_tensor_dims[in_dim.input_dim]}"
+                        )
+                    input_to_output_tensor_dims[in_dim.input_dim] = [output_dim]
+            elif len(in_dims) > 0:
+                # InputDim (identity) or Split(split_id=0).
+                in_dim = in_dims[0]
+                if in_dim.input_dim not in input_to_output_tensor_dims:
+                    input_to_output_tensor_dims[in_dim.input_dim] = [output_dim]
+                else:
+                    input_to_output_tensor_dims[in_dim.input_dim].append(output_dim)
+            elif isinstance(cmd, Split):
+                # Split(split_id>0): _analyze_split returned [], so chase the
+                # root input dim and append this output dim to its existing entry.
+                #
+                # Flatten+Split example: view([2, 3], [3, 2])
+                #   rule = (Split(Flatten(InputDim(0), InputDim(1)), (3,2), 0),
+                #           Split(Flatten(InputDim(0), InputDim(1)), (3,2), 1))
+                #   output_dim=0 (split_id=0): same as Split example above.
+                #     Result: {0: [0]}
+                #   output_dim=1 (split_id=1): same as Split example, but
+                #     the chase unwraps the inner Flatten to find InputDim(0).
+                #     Result: {0: [0, 1]}
+                root_spec = cmd.input_dim
+                while isinstance(root_spec, (Flatten, Split)):
+                    if isinstance(root_spec, Flatten):
+                        # _analyze_flatten always returns input_dims[0] as
+                        # the first element (either as the only shardable dim
+                        # in non-strict mode, or as the fallback when nothing
+                        # is sharded), so split_id=0 uses it as the key in
+                        # input_to_output_tensor_dims. Use [0] here to match.
+                        root_spec = root_spec.input_dims[0]
+                    else:
+                        root_spec = root_spec.input_dim
+                root = root_spec if isinstance(root_spec, InputDim) else None
+                if root is not None and root.input_dim in input_to_output_tensor_dims:
+                    input_to_output_tensor_dims[root.input_dim].append(output_dim)
+
+        input_tgt_placements: list[Placement] = []
+        for mesh_dim, p in enumerate(self.input_src_placements):
             if (
-                isinstance(placement, Shard | _StridedShard)
-                and placement.dim == input_dim.input_dim
+                isinstance(p, Shard | _StridedShard)
+                and not self.shard_allowed[p.dim][mesh_dim]
             ):
-                return i, placement
-        return None, None
-
-    # NOTE: This function has three responsibilities:
-    # 1. determine "theoretically" if an output dimension can be sharded, i.e. fill the shardable_dims map
-    # 2. determine "theoretically" the corresponding input dimension to shard on, via return value
-    # 3. throw an error when strict_view is enabled and we cannot shard an output dimension
-    # 1 and 2 doesn't require the info of whether current input is sharded.
-    # 3 requires that info, to decide whether we can error out. Maybe we can refactor
-    # to make this function purely "theoretical".
-    def get_in_dim_to_shard(cmd: DimSpec) -> InputDim | None:
-        if isinstance(cmd, InputDim):
-            return cmd
-        elif isinstance(cmd, Flatten):
-            for i, dim in enumerate(cmd.input_dims):
-                # so far all Flatten is always composed of InputDims; revisit this if needed
-                if not isinstance(dim, InputDim):
-                    raise AssertionError(f"Expected InputDim, got {type(dim)}")
-                can_shard_dim = True
-                shard_mesh_dim, shard_placement = (
-                    maybe_get_shard_mesh_dim_and_placement(dim)
+                input_tgt_placements.append(Replicate())
+            else:
+                input_tgt_placements.append(p)
+        return input_tgt_placements, input_to_output_tensor_dims
+
+    def rewrite_output_placements(
+        self,
+        input_tgt_placements: Sequence[Placement],
+        input_to_output_tensor_dims: dict[int, list[int]],
+    ) -> list[Placement]:
+        """Phase 2: consume analyze() outputs, return final output placements."""
+        # (input_dim, output_dim) pairs claimed by earlier mesh dims
+        # (via _rewrite_strided_shard), to avoid double-assignment.
+        strided_shard_claimed_dims: set[ClaimedDim] = set()
+        # Starts as global_input_shape; each mesh dim divides its sharded dim.
+        local_tensor_shapes: list[int] = list(self.global_input_shape)
+
+        output_placements: list[Placement] = []
+        # Process mesh dims in order; _rewrite_*_shard relies on this for
+        # truncating division safety in local_tensor_shapes.
+        for mesh_dim, p in enumerate(input_tgt_placements):
+            if isinstance(p, Shard):
+                placement, local_tensor_shapes = self._rewrite_plain_shard(
+                    p,
+                    mesh_dim,
+                    input_tgt_placements,
+                    strided_shard_claimed_dims,
+                    local_tensor_shapes,
+                    input_to_output_tensor_dims,
                 )
-                input_sharded = shard_mesh_dim is not None
-                if i > 0:
-                    can_shard_dim = False
-                    if strict_view and input_sharded:
-                        raise RuntimeError(
-                            f"Attempted to flatten multiple dimensions, with dimension {dim.input_dim} being sharded. ",
-                            "It cannot be performed without redistribution, which is disallowed by the current operator.",
-                        )
-                elif input_sharded:
-                    if not (shard_placement is not None and shard_mesh_dim is not None):
-                        raise AssertionError(
-                            "Expected shard_placement and shard_mesh_dim to be not None"
-                        )
-                    tensor_dim_size = global_input_shape[shard_placement.dim]
-                    mesh_dim_size = mesh_sizes[shard_mesh_dim]
-                    if tensor_dim_size % mesh_dim_size != 0:
-                        can_shard_dim = False
-                        if strict_view:
-                            raise RuntimeError(
-                                f"Attempted to flatten unevenly sharded dimension {i}, "
-                                "which would require resharding the input. "
-                                "Please explicitly redistribute the tensor instead."
-                            )
-                shardable_dims[dim.input_dim] = [can_shard_dim] * mesh_ndim
+                output_placements.append(placement)
+            elif isinstance(p, _StridedShard):
+                placement, local_tensor_shapes = self._rewrite_strided_shard(
+                    p,
+                    mesh_dim,
+                    input_tgt_placements,
+                    strided_shard_claimed_dims,
+                    local_tensor_shapes,
+                    input_to_output_tensor_dims,
+                )
+                output_placements.append(placement)
+            else:
+                output_placements.append(p)
+        return output_placements
+
+    # ------------------------------------------------------------------
+    # Analysis phase helpers
+    # ------------------------------------------------------------------
+
+    @staticmethod
+    def _input_dims_in_rule(rule: DimMap) -> set[int]:
+        """Walk the DimMap rule tree and return all input dim indices that appear in it."""
+        seen: set[int] = set()
+
+        def _walk(cmd: DimSpec) -> None:
+            if isinstance(cmd, InputDim):
+                seen.add(cmd.input_dim)
+            for inp in cmd.inputs():
+                _walk(inp)
+
+        for cmd in rule:
+            _walk(cmd)
+        return seen
+
+    def _find_plain_shard(
+        self, input_dim: InputDim
+    ) -> tuple[int | None, Shard | _StridedShard | None]:
+        """Find the mesh dim with a plain Shard on ``input_dim``.
 
-            if not isinstance(cmd.input_dims[0], InputDim):
-                raise AssertionError(
-                    f"Expected InputDim, got {type(cmd.input_dims[0])}"
+        Only matches Shard, not _StridedShard.  Used by both _analyze_flatten
+        and _analyze_split.  _find_shard_for_split is the counterpart that
+        also matches _StridedShard with split_factor validation.
+        """
+        for mesh_dim, placement in enumerate(self.input_src_placements):
+            if isinstance(placement, Shard) and placement.dim == input_dim.input_dim:
+                return mesh_dim, placement
+        return None, None
+
+    def _find_shard_for_split(
+        self,
+        current_dim: int,
+        cmd: Split,
+        placements: Sequence[Placement],
+    ) -> tuple[int | None, Shard | _StridedShard | None]:
+        """Find the mesh dim and placement for an input dim in Split ops.
+
+        Matches both Shard and _StridedShard:
+        - Shard: plain unflatten, e.g. [6] Shard(0) → [2, 3].
+        - _StridedShard: unflatten after a prior flatten that produced
+          _StridedShard, e.g. [2,3,4] Shard(1) → flatten → [6,4]
+          _StridedShard(0,sf=2) → unflatten → [2,3,4].  Validates that
+          the split_factor matches the expected value for this split_id.
+        """
+        for mesh_dim, placement in enumerate(placements):
+            if not isinstance(placement, Shard | _StridedShard):
+                continue
+            if placement.dim != current_dim:
+                continue
+            if mesh_dim in self.matched_strided_mesh_dims:
+                continue
+
+            if isinstance(placement, _StridedShard):
+                expected_sf = self._expected_split_factor(
+                    cmd, current_dim, mesh_dim, placements
                 )
-            return cmd.input_dims[0]
-        elif isinstance(cmd, Split):
-            in_dim = get_in_dim_to_shard(cmd.input_dim)
-            out_size = cmd.group_shape[cmd.split_id]
-            if cmd.split_id == 0 and in_dim is not None:
-                # we need to check that the input dimension is divisible
-                # by the size of the submesh we're sharding it on
-                # NOTE: it would be possible to shard the same input dimension
-                # on more than one mesh dimension. In that case, the dimension
-                # needs to be divisible by the product of mesh sizes.
-                # In order to keep the problem more tractable, we will not consider
-                # double resharding as a suggestion (e.g. [Shard(0), Shard(0) ])
-                # but we will allow it if that's the input and it's compatible
-
-                # 1. is this dimension shardable on each individual mesh dim?
-                shardable_dims[in_dim.input_dim] = [
-                    out_size % mesh_dim_size == 0 for mesh_dim_size in mesh_sizes
-                ]
-
-                shard_mesh_dim, _ = maybe_get_shard_mesh_dim_and_placement(in_dim)
-                if strict_view and shard_mesh_dim is not None:
-                    if not shardable_dims[in_dim.input_dim][shard_mesh_dim]:
-                        raise RuntimeError(
-                            f"Attempted to split the sharded dimension {in_dim.input_dim} into multiple subdimensions. ",
-                            "It cannot be performed without redistribution, which is disallowed by the current operator.",
-                        )
+                if expected_sf == placement.split_factor:
+                    return mesh_dim, placement
+            else:
+                return mesh_dim, placement
+        return None, None
 
-                # 2. here we special case things like [Shard(0), Shard(0)]
-                submesh_size = 1
-                for size, shard in zip(mesh_sizes, input_src_placements):
-                    if isinstance(shard, Shard | _StridedShard) and shard.dim == in_dim:
-                        submesh_size *= size
-                if not out_size % submesh_size == 0:
-                    raise AssertionError(
-                        f"Resulting dimension size {out_size} is not divisible by its mesh dimension {submesh_size}."
+    def _analyze_flatten(self, cmd: Flatten) -> list[InputDim]:
+        """Fill self.shard_allowed for Flatten; return sharded input dims."""
+        from torch.fx.experimental.symbolic_shapes import guard_or_true
+
+        sharded_dims: list[InputDim] = []
+        num_input_dims = len(cmd.input_dims)
+        for i, dim in enumerate(cmd.input_dims):
+            if not isinstance(dim, InputDim):
+                raise AssertionError(f"Expected InputDim, got {type(dim)}")
+            shard_mesh_dim, shard_placement = self._find_plain_shard(dim)
+            if shard_mesh_dim is None or shard_placement is None:
+                continue  # default from analyze() already covers this
+            tensor_dim_size = self.global_input_shape[shard_placement.dim]
+            mesh_dim_size = self.mesh_sizes[shard_mesh_dim]
+            can_shard_dim = True
+            if self.strict_view:
+                is_last_input_dim = i == num_input_dims - 1
+                if not is_last_input_dim and guard_or_true(
+                    tensor_dim_size % mesh_dim_size != 0
+                ):
+                    raise RuntimeError(
+                        f"Cannot flatten unevenly sharded tensor: "
+                        f"dimension {dim.input_dim} (size {tensor_dim_size}) "
+                        f"is not evenly divisible by mesh dimension "
+                        f"{shard_mesh_dim} (size {mesh_dim_size}). "
+                        f"Please redistribute the tensor before this operation."
+                    )
+                sharded_dims.append(dim)
+            else:
+                # TODO: non-strict (reshape) should allow can_shard_dim = True
+                # for non-first flatten dims, since strict_view already does.
+                # Currently forces redistribution because the rewrite phase
+                # wasn't originally implemented for this case.
+                if i == 0:
+                    sharded_dims.append(dim)
+                    if guard_or_true(tensor_dim_size % mesh_dim_size != 0):
+                        can_shard_dim = False
+                else:
+                    can_shard_dim = False
+            self.shard_allowed[dim.input_dim] = [can_shard_dim] * self.mesh_ndim
+
+        if len(sharded_dims) > 0:
+            return sharded_dims
+        # No sharded dims: e.g. Flatten([InputDim(0), InputDim(1)]) where
+        # neither dim is sharded.  Return the first input dim so that
+        # input_to_output_tensor_dims is populated for identity rewrites.
+        if not isinstance(cmd.input_dims[0], InputDim):
+            raise AssertionError(f"Expected InputDim, got {type(cmd.input_dims[0])}")
+        return [cmd.input_dims[0]]
+
+    def _analyze_split(self, cmd: Split) -> list[InputDim]:
+        """Fill self.shard_allowed for Split; return shardable input dims."""
+        from torch.fx.experimental.symbolic_shapes import guard_or_false, guard_or_true
+
+        in_dims = self._analyze_dim(cmd.input_dim)
+        if len(in_dims) == 0:
+            return []
+        in_dim = in_dims[0]
+        out_size = cmd.group_shape[cmd.split_id]
+        shard_mesh_dim, input_src_placement = self._find_shard_for_split(
+            in_dim.input_dim, cmd, self.input_src_placements
+        )
+        # split_id == 0 sets the base shard_allowed for this input dim.
+        # Later split_ids (processed in subsequent rule iterations) refine
+        # individual mesh_dim entries via the _StridedShard branch below.
+        if cmd.split_id == 0:
+            self.shard_allowed[in_dim.input_dim] = [
+                guard_or_false(out_size % mesh_dim_size == 0)
+                for mesh_dim_size in self.mesh_sizes
+            ]
+            plain_mesh_dim, _ = self._find_plain_shard(in_dim)
+            # Non-strict silently redistributes via shard_allowed=False above;
+            # strict raises so the user knows to redistribute before view().
+            if self.strict_view and plain_mesh_dim is not None:
+                if not self.shard_allowed[in_dim.input_dim][plain_mesh_dim]:
+                    raise RuntimeError(
+                        f"Cannot unflatten unevenly sharded tensor: "
+                        f"output dimension {cmd.split_id} (size {out_size}) "
+                        f"is not evenly divisible by mesh dimension "
+                        f"{plain_mesh_dim} (size {self.mesh_sizes[plain_mesh_dim]}). "
+                        f"Please redistribute the tensor before this operation."
                     )
+        if shard_mesh_dim is not None and isinstance(
+            input_src_placement, _StridedShard
+        ):
+            # The last split dim doesn't require even divisibility because
+            # its local size is inferred: local_last = local_flat / product
+            # of earlier dims, and DTensor handles uneven local sizes.
+            # Non-last dims must be evenly divisible because they appear as
+            # fixed sizes in the local reshape — uneven division would make
+            # the stride pattern inconsistent across devices.
+            # E.g. [12] → [3, 4], _StridedShard targeting dim 1 (last),
+            # mesh=3: 4%3≠0, but local shapes [3,2],[3,1],[3,1] are valid.
+            is_last_split_dim = cmd.split_id == len(cmd.group_shape) - 1
+            if (
+                self.strict_view
+                and not is_last_split_dim
+                and guard_or_true(out_size % self.mesh_sizes[shard_mesh_dim] != 0)
+            ):
+                raise RuntimeError(
+                    f"Cannot unflatten unevenly sharded tensor: "
+                    f"output dimension {cmd.split_id} (size {out_size}) "
+                    f"is not evenly divisible by mesh dimension {shard_mesh_dim} "
+                    f"(size {self.mesh_sizes[shard_mesh_dim]}). "
+                    f"Please redistribute the tensor before this operation."
+                )
+            # Prevents _find_shard_for_split from matching this mesh dim
+            # again for a later split_id of the same Split group.
+            self.matched_strided_mesh_dims.add(shard_mesh_dim)
+            if in_dim.input_dim in self.shard_allowed:
+                self.shard_allowed[in_dim.input_dim][shard_mesh_dim] = (
+                    guard_or_false(out_size % self.mesh_sizes[shard_mesh_dim] == 0)
+                    or is_last_split_dim
+                )
+        # Only split_id==0 returns the input dim for input_to_output_tensor_dims.
+        # Later split_ids refine shard_allowed above but return [] — their
+        # output dims are linked via the root-input-dim chase in analyze().
+        return [in_dim] if cmd.split_id == 0 else []
 
-            # we will only shard our first component of the split
-            return in_dim if cmd.split_id == 0 else None
+    def _analyze_dim(self, cmd: DimSpec) -> list[InputDim]:
+        """Dispatch one DimSpec: update self.shard_allowed, return input dim(s) to shard on."""
+        if isinstance(cmd, InputDim):
+            return [cmd]
+        elif isinstance(cmd, Flatten):
+            return self._analyze_flatten(cmd)
+        elif isinstance(cmd, Split):
+            return self._analyze_split(cmd)
         elif isinstance(cmd, Repeat):
-            in_dim = get_in_dim_to_shard(cmd.input_dim)
-            if in_dim is not None:
-                shardable_dims[in_dim.input_dim] = [False] * mesh_ndim
-            return None
+            in_dims = self._analyze_dim(cmd.input_dim)
+            for d in in_dims:
+                self.shard_allowed[d.input_dim] = [False] * self.mesh_ndim
+            return []
         else:
-            return None
-
-    # for each output dim, find the corresponding input dim in terms of sharding prop
-    shard_dim_map = {}
-    for dim, cmd in enumerate(rule):
-        in_dim = get_in_dim_to_shard(cmd)
-        if in_dim is not None:
-            shard_dim_map[in_dim.input_dim] = dim
-
-    input_tgt_placements = [
-        (
-            Replicate()
-            if isinstance(p, Shard | _StridedShard)
-            and not shardable_dims[p.dim][mesh_dim]
-            else p
+            return []
+
+    # ------------------------------------------------------------------
+    # Rewrite phase helpers
+    # ------------------------------------------------------------------
+
+    @staticmethod
+    def _is_last_shard_in_flatten_range(
+        mesh_dim: int,
+        placements: Sequence[Placement],
+        flatten_start: int,
+        flatten_end: int,
+    ) -> bool:
+        """Check if no later mesh dim shards a dim within the flatten range at or above this one.
+
+        Uneven sharding on dim d breaks stride computation for all earlier dims
+        that flatten together with d. Only dims within [flatten_start, flatten_end)
+        matter; shards on dims outside the flatten range are independent.
+
+        Requires: placements[mesh_dim] must be Shard or _StridedShard.
+        """
+        p = placements[mesh_dim]
+        if not isinstance(p, (Shard, _StridedShard)):
+            raise AssertionError(
+                f"Expected Shard or _StridedShard at mesh_dim {mesh_dim}, got {type(p)}"
+            )
+        tensor_dim = p.dim
+        return not any(
+            isinstance(other_p, (Shard, _StridedShard))
+            and flatten_start <= other_p.dim < flatten_end
+            and other_p.dim >= tensor_dim
+            for other_p in placements[mesh_dim + 1 :]
         )
-        for mesh_dim, p in enumerate(input_src_placements)
-    ]
 
-    def _rewrite_shard_dim(p: Shard | _StridedShard):
+    def _expected_split_factor(
+        self,
+        cmd: Split,
+        sharded_dim: int,
+        mesh_dim: int,
+        placements: Sequence[Placement],
+    ) -> int | None:
+        """Compute the residual split factor for ``cmd`` after earlier mesh dims.
+
+        Starts from ``math.prod(cmd.group_shape[:cmd.split_id])`` and divides
+        out each earlier mesh dim that shards the same input dim.  Returns
+        ``None`` if any earlier mesh size doesn't divide evenly.
+        """
+        sf = math.prod(cmd.group_shape[: cmd.split_id])
+        for m in range(mesh_dim):
+            other_p = placements[m]
+            if (
+                isinstance(other_p, (_StridedShard, Shard))
+                and other_p.dim == sharded_dim
+            ):
+                if sf % self.mesh_sizes[m] != 0:
+                    return None
+                sf //= self.mesh_sizes[m]
+        return sf
+
+    def _find_keep_ss_dim(
+        self,
+        tgt_shard_dims: list[int],
+        p: _StridedShard,
+        mesh_dim: int,
+    ) -> int | None:
+        """Find an output dim where SS stays as SS.
+
+        Returns the first output dim whose Split can accommodate the combined
+        sharding (mesh_size * split_factor), or ``None`` if no dim fits.
         """
-        Rewrite the shard dim to the corresponding tensor dim in output.
-        For ``_StridedShard``, we can safely keep the placement type and
-        ``split_factor`` unchanged and only rewrite the ``dim`` because:
-        1. ``_StridedShard`` has no impact on sharding (i.e. how
-            tensor is partitioned) compared to ``Shard``. It only changes
-            how shards permute across the devices.
-        2. ``view()`` op on DTensor strictly forbids shard redistribution
-            which means if ``view()`` may cause shard permutation across
-            devices, it should be rejected. This is enforced in today's
-            sharding prop for ``view()``.
-        3. Since DTensor ``view()`` won't introduce any redistribution,
-            it's certain that ``placements`` won't change except the
-            inner ``dim`` attribute of ``Shard`` or ``_StridedShard``.
+        total_shard = self.mesh_sizes[mesh_dim] * p.split_factor
+        if self.global_input_shape[p.dim] % total_shard != 0:
+            return None
+        shard_size = self.global_input_shape[p.dim] // total_shard
+        for candidate_dim in tgt_shard_dims:
+            cmd = self.rule[candidate_dim]
+            if isinstance(cmd, Split):
+                inner_size = math.prod(cmd.group_shape[cmd.split_id + 1 :])
+                # When a Split wraps a Flatten, the per-shard chunk covers
+                # the sharded dim plus trailing dims flattened together.
+                trailing_size = 1
+                if isinstance(cmd.input_dim, Flatten):
+                    found = False
+                    for flat_dim in cmd.input_dim.input_dims:
+                        if not isinstance(flat_dim, InputDim):
+                            raise AssertionError(
+                                f"Expected InputDim, got {type(flat_dim)}"
+                            )
+                        if flat_dim.input_dim == p.dim:
+                            found = True
+                        elif found:
+                            trailing_size *= self.global_input_shape[flat_dim.input_dim]
+                flattened_shard_size = shard_size * trailing_size
+                if (
+                    flattened_shard_size >= inner_size
+                    and flattened_shard_size % inner_size == 0
+                ):
+                    return candidate_dim
+        return None
+
+    def _rewrite_plain_shard(
+        self,
+        p: Shard,
+        mesh_dim: int,
+        placements: Sequence[Placement],
+        strided_shard_claimed_dims: set[ClaimedDim],
+        local_tensor_shapes: list[int],
+        input_to_output_tensor_dims: dict[int, list[int]],
+    ) -> tuple[Placement, list[int]]:
+        """Given a plain Shard(dim=X) input placement on a specific mesh dim,
+        determine what output placement it maps to after the view op.
+
+        For identity and unflatten, produces Shard on the output dim.  For
+        flatten, Shard on the first flattened dim stays Shard, while Shard on
+        a non-first dim produces _StridedShard (consumed later by
+        _rewrite_strided_shard).
+
+        Returns the output placement and a new local_tensor_shapes with this
+        mesh dim's division applied.
         """
-        if isinstance(p, _StridedShard):
-            return _StridedShard(shard_dim_map[p.dim], split_factor=p.split_factor)
+        # Output dims that input dim p.dim maps to, filtering out any
+        # already claimed by _StridedShard rewriting on earlier mesh dims.
+        tgt_shard_dims = [
+            d
+            for d in input_to_output_tensor_dims[p.dim]
+            if ClaimedDim(p.dim, d) not in strided_shard_claimed_dims
+        ]
+        if len(tgt_shard_dims) == 0:
+            raise AssertionError(
+                f"No output dim available for Shard(dim={p.dim}) on mesh dim "
+                f"{mesh_dim}. All output dims already claimed by earlier mesh dims."
+            )
+        if len(tgt_shard_dims) == 1:
+            tgt_shard_dim = tgt_shard_dims[0]
         else:
-            return Shard(shard_dim_map[p.dim])
-
-    output_placements = [
-        _rewrite_shard_dim(p) if isinstance(p, Shard | _StridedShard) else p
-        for p in input_tgt_placements
-    ]
-
-    return input_tgt_placements, output_placements
+            # Unflatten: one input dim maps to multiple output dims
+            # (e.g. (24,) → (2, 3, 4) gives 3 splits). Plain Shard
+            # always targets the split_id=0 output dim.
+            tgt_shard_dim = next(
+                (
+                    d
+                    for d in tgt_shard_dims
+                    if isinstance(self.rule[d], Split)
+                    and cast(Split, self.rule[d]).split_id == 0
+                ),
+                None,
+            )
+            if tgt_shard_dim is None:
+                raise AssertionError(
+                    f"No Split(split_id=0) found among unclaimed output dims "
+                    f"{tgt_shard_dims} for Shard(dim={p.dim}) on mesh dim {mesh_dim}."
+                )
+        cmd = self.rule[tgt_shard_dim]
+        if isinstance(cmd, Split) and isinstance(cmd.input_dim, Flatten):
+            first_dim = cmd.input_dim.input_dims[0]
+            if isinstance(first_dim, InputDim) and p.dim != first_dim.input_dim:
+                raise RuntimeError(
+                    f"Shard(dim={p.dim}) through Split(Flatten(...), {cmd.group_shape}) "
+                    f"is not supported yet for non-first flatten dims."
+                )
+        if isinstance(cmd, (Split, InputDim)):
+            # Split/InputDim: 1:1 dim mapping, sharding transfers directly.
+            # Flatten needs stride computation below (multiple dims merge).
+            new_shapes = list(local_tensor_shapes)
+            new_shapes[p.dim] //= self.mesh_sizes[mesh_dim]
+            return Shard(tgt_shard_dim), new_shapes
+        if not isinstance(cmd, Flatten):
+            raise AssertionError(f"Expected Flatten, got {type(cmd)}")
+        first_dim = cmd.input_dims[0]
+        last_dim = cmd.input_dims[-1]
+        if not isinstance(first_dim, InputDim):
+            raise AssertionError(f"Expected InputDim, got {type(first_dim)}")
+        if not isinstance(last_dim, InputDim):
+            raise AssertionError(f"Expected InputDim, got {type(last_dim)}")
+        input_start_idx = first_dim.input_dim
+        if p.dim == input_start_idx:
+            output_placement: Placement = Shard(tgt_shard_dim)
+        else:
+            split_factor = math.prod(local_tensor_shapes[input_start_idx : p.dim])
+            output_placement = _StridedShard(tgt_shard_dim, split_factor=split_factor)
+        # Uneven sharding on a non-last flatten dim breaks _StridedShard:
+        # split_factor (number of groups) must be the same on all devices,
+        # but uneven division of a non-last dim makes group count vary.
+        # E.g. [3,4]→[12] Shard(0) mesh=2: device 0 has 2 groups of 4,
+        # device 1 has 1 group of 4 — no consistent split_factor.
+        # The last dim is exempt: only group *size* varies, not count.
+        flatten_end = last_dim.input_dim + 1
+        if local_tensor_shapes[p.dim] % self.mesh_sizes[
+            mesh_dim
+        ] != 0 and not self._is_last_shard_in_flatten_range(
+            mesh_dim, placements, input_start_idx, flatten_end
+        ):
+            raise RuntimeError(
+                f"Cannot shard unevenly distributed tensor: "
+                f"dimension {p.dim} (size {local_tensor_shapes[p.dim]}) "
+                f"is not evenly divisible by mesh dimension "
+                f"{mesh_dim} (size {self.mesh_sizes[mesh_dim]}). "
+                f"Please redistribute the tensor before this operation."
+            )
+        new_shapes = list(local_tensor_shapes)
+        new_shapes[p.dim] //= self.mesh_sizes[mesh_dim]
+        return output_placement, new_shapes
+
+    def _rewrite_strided_shard(
+        self,
+        p: _StridedShard,
+        mesh_dim: int,
+        placements: Sequence[Placement],
+        strided_shard_claimed_dims: set[ClaimedDim],
+        local_tensor_shapes: list[int],
+        input_to_output_tensor_dims: dict[int, list[int]],
+    ) -> tuple[Placement, list[int]]:
+        """Rewrite _StridedShard placement to target the correct output dim.
+
+        _StridedShard inputs arise from a prior flatten on a non-first dim
+        (produced by _rewrite_plain_shard above).  The interesting case is
+        unflatten (Split rule): the split_factor may resolve to contiguous
+        sharding (producing Shard) or stay as _StridedShard.  For
+        identity/flatten rules, falls through to the fallback and keeps the
+        placement as-is.
+
+        Returns the output placement and a new local_tensor_shapes with this
+        mesh dim's division applied.
+        """
+        tgt_shard_dims = [
+            d
+            for d in input_to_output_tensor_dims[p.dim]
+            if ClaimedDim(p.dim, d) not in strided_shard_claimed_dims
+        ]
+        # Phase 1: resolve SS → Shard.  If an output dim's Split has a
+        # group_shape prefix matching the split_factor, the strided pattern
+        # is fully captured by the Split, so SS simplifies to Shard.
+        # E.g. unflatten (6, 4) → (2, 3, 4) with SS(0, sf=2) on mesh (3):
+        # sf=2 means 2 groups of contiguous data in dim 0.  Split into
+        # (2, 3, 4) gives group_shape=(2, 3); prod(group_shape[:1])=2==sf,
+        # so the strided pattern lands exactly on output dim 1 → Shard(1).
+        for candidate_dim in tgt_shard_dims:
+            cmd = self.rule[candidate_dim]
+            if isinstance(cmd, Split):
+                expected_sf = self._expected_split_factor(
+                    cmd, p.dim, mesh_dim, placements
+                )
+                if expected_sf != p.split_factor:
+                    continue
+                strided_shard_claimed_dims.add(ClaimedDim(p.dim, candidate_dim))
+                new_shapes = list(local_tensor_shapes)
+                new_shapes[p.dim] //= self.mesh_sizes[mesh_dim]
+                return Shard(candidate_dim), new_shapes
+
+        # Phase 2: keep SS as SS.  Phase 1 is tried first because we prefer
+        # resolving to the simpler Shard when possible.
+        tgt_shard_dim = self._find_keep_ss_dim(tgt_shard_dims, p, mesh_dim)
+
+        if tgt_shard_dim is None:
+            if self.strict_view and any(
+                isinstance(self.rule[d], Split) for d in tgt_shard_dims
+            ):
+                raise RuntimeError(
+                    f"Cannot unflatten tensor with _StridedShard placement: "
+                    f"split_factor={p.split_factor} does not match any output "
+                    f"dimension. This typically means the _StridedShard placement "
+                    f"was constructed with a split_factor that is incompatible "
+                    f"with the unflatten shape. Please redistribute the tensor "
+                    f"before this operation."
+                )
+            if len(tgt_shard_dims) == 0:
+                raise AssertionError(
+                    f"No unclaimed output dims for _StridedShard(dim={p.dim}) "
+                    f"on mesh dim {mesh_dim}."
+                )
+            # Fallback for identity/flatten: tgt_shard_dims has exactly one
+            # element, so [0] is correct.  For Split rules this is unreachable
+            # in practice — the analysis phase rejects mismatched split_factors
+            # via shard_allowed, forcing redistribution before we get here.
+            tgt_shard_dim = tgt_shard_dims[0]
+        new_shapes = list(local_tensor_shapes)
+        new_shapes[p.dim] //= self.mesh_sizes[mesh_dim]
+        return _StridedShard(tgt_shard_dim, split_factor=p.split_factor), new_shapes
 
 
 def register_op_strategy_map(
@@ -740,12 +1291,17 @@ def reshape_strategy(op_schema: OpSchema) -> StrategyType:
                 placements=tuple(input_tgt_placements),
                 mesh=mesh,
                 tensor_meta=input_src_spec.tensor_meta,
+                use_strided_shard_as_shard_order=False,
             )
             redistribute_costs: list[list[float]] = [
                 generate_redistribute_costs(input_strategy, input_tgt_spec)
             ]
 
-            output_spec = DTensorSpec(mesh=mesh, placements=tuple(output_placements))
+            output_spec = DTensorSpec(
+                mesh=mesh,
+                placements=tuple(output_placements),
+                use_strided_shard_as_shard_order=False,
+            )
             output_strategy.strategies.append(
                 OpSpec(
                     output_specs=output_spec,
@@ -758,18 +1314,30 @@ def reshape_strategy(op_schema: OpSchema) -> StrategyType:
 
 
 register_op_strategy_map(aten.squeeze.default, torch.squeeze)
+register_op_strategy_map(aten.squeeze_.default, torch.squeeze)
 register_op_strategy_map(
     aten.squeeze_.dim, torch.squeeze, schema_info=RuntimeSchemaInfo(1)
 )
 register_op_strategy_map(
     aten.squeeze.dim, torch.squeeze, schema_info=RuntimeSchemaInfo(1)
 )
+register_op_strategy_map(
+    aten.squeeze.dims, torch.squeeze, schema_info=RuntimeSchemaInfo(1)
+)
+register_op_strategy_map(
+    aten.squeeze_.dims, torch.squeeze, schema_info=RuntimeSchemaInfo(1)
+)
 register_op_strategy_map(
     aten.view.default,
     Tensor.view,
     schema_info=RuntimeSchemaInfo(1),
     strict_view=True,
 )
+register_op_strategy_map(
+    aten.view_copy.default,
+    Tensor.view,
+    schema_info=RuntimeSchemaInfo(1),
+)
 register_op_strategy_map(
     aten.reshape.default, torch.reshape, schema_info=RuntimeSchemaInfo(1)
 )
@@ -785,6 +1353,9 @@ def reshape_strategy(op_schema: OpSchema) -> StrategyType:
 register_op_strategy_map(
     aten.expand.default, Tensor.expand, schema_info=RuntimeSchemaInfo(1)
 )
+register_op_strategy_map(
+    aten.expand_copy.default, Tensor.expand, schema_info=RuntimeSchemaInfo(1)
+)
 register_op_strategy_map(
     aten.permute.default, torch.permute, schema_info=RuntimeSchemaInfo(1)
 )
@@ -794,5 +1365,23 @@ def reshape_strategy(op_schema: OpSchema) -> StrategyType:
 register_op_strategy_map(
     aten.transpose.int, torch.transpose, schema_info=RuntimeSchemaInfo(1)
 )
-register_op_strategy_map(aten.view_as_complex.default, torch.view_as_complex)
+
+
+@register_single_dim_strategy(aten.view_as_complex.default)
+def view_as_complex_single_dim_strategy(op, args_schema, kwargs_schema):
+    # view_as_complex: float [..., 2] -> complex [...]
+    # Dims 0..ndim-2 map 1:1; last dim (real/imag pair) is consumed.
+    # P(max)/P(min) invalid: complex numbers have no total ordering.
+    input_meta = args_schema[0]
+    if not isinstance(input_meta, TensorMeta):
+        raise AssertionError(f"Expected TensorMeta, got {type(input_meta)}")
+    ndim = len(input_meta.shape)
+    strategies: list[list[Placement | _ShardingPlaceholder]] = []
+    for d in range(ndim - 1):
+        strategies.append([_ShardingPlaceholder(d), _ShardingPlaceholder(d)])
+    strategies.append([Partial("sum"), Partial("sum")])
+    strategies.append([Partial("avg"), Partial("avg")])
+    return strategies
+
+
 register_op_strategy_map(aten.view_as_real.default, torch.view_as_real)
diff --git a/torch/distributed/tensor/_ops/single_dim_strategy.py b/torch/distributed/tensor/_ops/single_dim_strategy.py
index fab302b953911..722cb8c24a8a2 100644
--- a/torch/distributed/tensor/_ops/single_dim_strategy.py
+++ b/torch/distributed/tensor/_ops/single_dim_strategy.py
@@ -1,19 +1,28 @@
 #  Copyright (c) Meta Platforms, Inc. and affiliates
 import functools
+import heapq
 import logging
+import math
+from collections import defaultdict
 from collections.abc import Callable, Sequence
 from dataclasses import dataclass, field
-from typing import Any, cast, Optional, TypeAlias, TypeVar, Union
+from itertools import count
+from typing import Any, cast, TypeAlias, TypeVar
+from typing_extensions import TypeIs
 
 import torch
 from torch._ops import OpOverload
+from torch.distributed.tensor._collective_utils import (
+    _compute_placement_transition_cost,
+    MeshTopoInfo,
+)
 from torch.distributed.tensor._dtensor_spec import DTensorSpec, TensorMeta
 from torch.distributed.tensor._op_schema import (
     ArgsType,
     KwargsType,
     OpSchema,
+    OpSpec,
     OpStrategy,
-    PlacementList,
     RuntimeSchemaInfo,
     StrategyType,
     TupleStrategy,
@@ -31,6 +40,10 @@
 logger = logging.getLogger(__name__)
 
 
+def _is_sharding(p: Placement) -> TypeIs[Shard | _StridedShard]:
+    return isinstance(p, (Shard, _StridedShard))
+
+
 class _ShardingPlaceholder:
     """
     A placeholder for a sharding placement that has a specified tensor dim, but the other
@@ -61,6 +74,11 @@ def __repr__(self) -> str:
 class _SingleDimStrategyInfo:
     func: _SingleDimStrategyFunc
     allow_unbacked_sharding: bool | None = field(default=None)
+    allow_uneven_sharding: bool = field(default=False)
+    # Positions (in args_schema) of args that may live on a different mesh
+    # than the op's compute mesh.  These args must be Replicate.
+    # See Note [Multi-mesh args] in expand_to_full_mesh_op_strategy.
+    different_mesh_args: list[int] | None = field(default=None)
 
     # Delegate to func so this can be used interchangeably with a raw
     # _SingleDimStrategyFunc (e.g. in tests that call strategy functions directly).
@@ -70,28 +88,39 @@ def __call__(self, *args, **kwargs):
 
 def _insert_single_dim_replication_strategy(
     single_dim_strategies_with_placeholders: list[
-        list[Placement | _ShardingPlaceholder]
+        list[Placement | _ShardingPlaceholder | None]
     ],
     num_outputs: int,
     num_input_tensors: int,
-) -> list[list[Placement | _ShardingPlaceholder]]:
+    output_tensor_meta: TensorMeta | Sequence[TensorMeta | None] | None = None,
+) -> list[list[Placement | _ShardingPlaceholder | None]]:
     """
     Inserts the [Replicate(), Replicate(), ...] strategy after asserting that such strategy does not yet exist.
+    For ops with masked-off outputs (e.g. backward ops with output_mask), output positions
+    where output_tensor_meta is None are set to None in the all-Replicate rule.
     """
     for strategy in single_dim_strategies_with_placeholders:
-        assert not all(isinstance(p, Replicate) for p in strategy)
-    single_dim_strategies_with_placeholders.insert(
-        0, [Replicate()] * (num_outputs + num_input_tensors)
-    )
+        if all(isinstance(p, Replicate) or p is None for p in strategy):
+            return single_dim_strategies_with_placeholders
+    total_len = num_outputs + num_input_tensors
+    replicate_rule: list[Placement | _ShardingPlaceholder | None] = [
+        Replicate()
+    ] * total_len
+    # Set None for masked-off output positions based on output_tensor_meta
+    if isinstance(output_tensor_meta, Sequence):
+        for i, meta in enumerate(output_tensor_meta):
+            if meta is None and i < num_outputs:
+                replicate_rule[i] = None
+    single_dim_strategies_with_placeholders.insert(0, replicate_rule)
     return single_dim_strategies_with_placeholders
 
 
 def _fill_single_dim_strategy_placeholders(
     unique_input_placements: set[Placement],
     single_dim_strategies_with_placeholders: list[
-        list[Placement | _ShardingPlaceholder]
+        list[Placement | _ShardingPlaceholder | None]
     ],
-) -> list[list[Placement]]:
+) -> list[list[Placement | None]]:
     """
     Replace any _ShardingPlaceholder with the specific Sharding types used by the inputs in op_schema.
     Supports implicit replication.
@@ -121,23 +150,29 @@ def _fill_single_dim_strategy_placeholders(
 
     # if any of the placements is a placeholder, we need to expand the strategy
     # to all possible combinations of placements
-    expanded_strategies_over_one_mesh_dim: list[list[Placement]] = []
+    expanded_strategies_over_one_mesh_dim: list[list[Placement | None]] = []
     for s in single_dim_strategies_with_placeholders:
         if any(isinstance(p, _ShardingPlaceholder) for p in s):
             for shard_builder in shard_builders.values():
-                expanded_strategy: list[Placement] = []
+                expanded_strategy: list[Placement | None] = []
                 for maybe_placeholder in s:
                     if isinstance(maybe_placeholder, _ShardingPlaceholder):
                         # we combine the tensor dim to shard from the placeholder
                         # with other metadata (e.g. split_factor) from the sharding class
                         expanded_strategy.append(shard_builder(maybe_placeholder.dim))
+                    elif maybe_placeholder is None:
+                        expanded_strategy.append(None)
                     else:
-                        assert isinstance(maybe_placeholder, Placement)
+                        if not isinstance(maybe_placeholder, Placement):
+                            raise AssertionError
                         expanded_strategy.append(maybe_placeholder)
                 expanded_strategies_over_one_mesh_dim.append(expanded_strategy)
         else:
-            assert all(isinstance(p, Placement) for p in s)
-            expanded_strategies_over_one_mesh_dim.append(cast(list[Placement], (s)))
+            if not all(isinstance(p, Placement) or p is None for p in s):
+                raise AssertionError
+            expanded_strategies_over_one_mesh_dim.append(
+                cast(list[Placement | None], (s))
+            )
 
     return expanded_strategies_over_one_mesh_dim
 
@@ -149,11 +184,15 @@ def _update_placements(obj: Any):
         if isinstance(obj, DTensorSpec):
             unique_placements.update(obj.placements)
         elif isinstance(obj, OpStrategy):
-            assert len(obj.strategies) == 1
+            if len(obj.strategies) != 1:
+                raise AssertionError
             unique_placements.update(obj.strategies[0].output_spec.placements)
         elif isinstance(obj, TupleStrategy):
             for child in obj.children:
                 _update_placements(child)
+        elif isinstance(obj, (list, tuple)):
+            for child in obj:
+                _update_placements(child)
 
     for obj in op_schema.args_schema:
         _update_placements(obj)
@@ -166,20 +205,300 @@ def _update_placements(obj: Any):
 
 def _get_num_tensor_inputs(op_schema: OpSchema) -> int:
     num_inputs = 0
-    for obj in op_schema.args_schema:
+
+    def _count(obj: Any) -> int:
         if isinstance(obj, OpStrategy):
-            num_inputs += 1
+            return 1
         elif isinstance(obj, TupleStrategy):
-            num_inputs += len(obj.children)
+            return sum(1 for child in obj.children if child is not None)
+        elif isinstance(obj, (list, tuple)):
+            return sum(_count(child) for child in obj)
+        return 0
+
+    for obj in op_schema.args_schema:
+        num_inputs += _count(obj)
     # Also count tensor kwargs (e.g., "out" for out-variant ops)
     for obj in op_schema.kwargs_schema.values():
-        if isinstance(obj, OpStrategy):
-            num_inputs += 1
-        elif isinstance(obj, TupleStrategy):
-            num_inputs += len(obj.children)
+        num_inputs += _count(obj)
     return num_inputs
 
 
+def _build_output_specs(
+    mesh: DeviceMesh,
+    per_mesh_dim_placements: list[tuple[Placement | None, ...]],
+    num_outputs: int,
+    output_metas: tuple[TensorMeta | None, ...],
+) -> DTensorSpec | tuple[DTensorSpec | None, ...]:
+    """Build output spec(s) by transposing per-mesh-dim placements to per-output.
+
+    per_mesh_dim_placements is indexed [mesh_dim][output_idx]. output_metas must
+    have exactly num_outputs elements. Outputs where output_metas[i] is None
+    (masked-off outputs) produce None specs.
+    """
+    if num_outputs <= 0:
+        raise AssertionError(f"Expected num_outputs > 0, got {num_outputs}")
+    if len(output_metas) != num_outputs:
+        raise AssertionError(
+            f"Expected {num_outputs} output_metas, got {len(output_metas)}"
+        )
+
+    def _spec_for_output(out_idx: int) -> DTensorSpec | None:
+        if output_metas[out_idx] is None:
+            return None
+        placements = tuple(
+            cast(Placement, out[out_idx]) for out in per_mesh_dim_placements
+        )
+        return DTensorSpec(mesh, placements, tensor_meta=output_metas[out_idx])
+
+    if num_outputs > 1:
+        return tuple(_spec_for_output(i) for i in range(num_outputs))
+    else:
+        spec = _spec_for_output(0)
+        if spec is None:
+            raise AssertionError("Single-output op cannot have None output meta")
+        return spec
+
+
+class _PreparedSingleDimStrategy:
+    """A single-dim strategy materialized for a specific op.
+
+    Expands a strategy function's placeholder-based rules into concrete
+    placement rules by filling in the actual shard/partial placements from
+    the op_schema. The result is a lookup table (strategy_lookup) that maps
+    input placements to output placements for one mesh dimension.
+
+    Provides try_propagate() for matching a multi-dim placement tuple against
+    the per-dim rules, and exposes allowed_sharding_per_input /
+    allowed_partial_per_input for graph search neighbor generation.
+    """
+
+    strategy_lookup: dict[tuple[Placement | None, ...], tuple[Placement | None, ...]]
+    expanded_strategies: list[list[Placement | None]]
+    num_outputs: int
+    num_inputs: int
+    output_metas: tuple[TensorMeta | None, ...]
+    allowed_sharding_per_input: dict[int, set[Shard | _StridedShard]]
+    allowed_partial_per_input: dict[int, set[Placement]]
+    allow_unbacked_sharding: bool | None
+
+    # many, but not all ops are able to support unevenly sharded tensors
+    # there are existing BC expectations even if we wanted to ban for
+    # simplicity, see why justification for why pointwise_ops always work
+    # with uneven sharding at
+    # https://github.com/pytorch/pytorch/pull/174874#issuecomment-3995152777
+    allow_uneven_sharding: bool
+
+    def __init__(
+        self,
+        strategy_fn: _SingleDimStrategyInfo
+        | Callable[
+            [OpOverload, ArgsType, KwargsType],
+            list[list[Placement | _ShardingPlaceholder]],
+        ],
+        op_schema: OpSchema,
+        output_tensor_meta: TensorMeta | Sequence[TensorMeta | None] | None,
+        num_inputs: int | None = None,
+    ) -> None:
+        # Note: circular import
+        from torch.distributed.tensor.placement_types import Partial
+
+        if isinstance(strategy_fn, _SingleDimStrategyInfo):
+            self.allow_unbacked_sharding = strategy_fn.allow_unbacked_sharding
+            self.allow_uneven_sharding = strategy_fn.allow_uneven_sharding
+            different_mesh_args = strategy_fn.different_mesh_args
+            func = strategy_fn.func
+        else:
+            self.allow_unbacked_sharding = None
+            self.allow_uneven_sharding = False
+            different_mesh_args = None
+            func = strategy_fn
+
+        # Determine element_mesh from the first OpStrategy arg.  For foreach
+        # per-element schemas the element's inputs may live on a smaller
+        # sub-mesh than the global compute_mesh.
+        self.element_mesh: DeviceMesh | None = None
+        for arg in op_schema.args_schema:
+            if isinstance(arg, OpStrategy):
+                self.element_mesh = arg.strategies[0].output_spec.mesh
+                break
+
+        # Validate that all inputs are on the same mesh (except
+        # different_mesh_args which are explicitly allowed to differ).
+        if self.element_mesh is not None:
+            allowed = set(different_mesh_args or [])
+            for i, arg in enumerate(op_schema.args_schema):
+                if isinstance(arg, OpStrategy) and i not in allowed:
+                    arg_mesh = arg.strategies[0].output_spec.mesh
+                    if arg_mesh != self.element_mesh:
+                        raise ValueError(
+                            f"Cannot run {op_schema.op} on inputs with different "
+                            f"meshes: got {self.element_mesh} and {arg_mesh}"
+                        )
+
+        # Remap different_mesh_args from args_schema positions to
+        # OpStrategy-only positions.  Non-OpStrategy args (e.g. empty lists)
+        # are filtered out by expand_to_full_mesh_op_strategy, shifting later
+        # indices.
+        self.remapped_different_mesh_args: list[int] | None = None
+        if different_mesh_args is not None:
+            schema_to_strategy: dict[int, int] = {}
+            strategy_pos = 0
+            for schema_pos, arg in enumerate(op_schema.args_schema):
+                if isinstance(arg, OpStrategy):
+                    schema_to_strategy[schema_pos] = strategy_pos
+                    strategy_pos += 1
+            self.remapped_different_mesh_args = [
+                schema_to_strategy[i]
+                for i in different_mesh_args
+                if i in schema_to_strategy
+            ]
+
+        if num_inputs is None:
+            num_inputs = _get_num_tensor_inputs(op_schema)
+        self.num_inputs = num_inputs
+
+        # Strategy functions may return None in output positions for masked-off
+        # outputs (e.g. backward ops with output_mask). Widen the type here.
+        strategies_with_placeholders = cast(
+            list[list[Placement | _ShardingPlaceholder | None]],
+            func(op_schema.op, op_schema.args_meta, op_schema.kwargs_meta),
+        )
+
+        # Validate strategy length against the op schema. The schema is the
+        # ground truth for num_outputs; combined with num_inputs (which counts
+        # all tensor args + kwargs), it gives the expected strategy length.
+        # A mismatch means the strategy is missing kwargs placements or has
+        # extra entries.
+        if len(strategies_with_placeholders) > 0:
+            schema_num_outputs = sum(
+                1 for r in op_schema.op._schema.returns if "Tensor" in str(r.type)
+            )
+            expected_len = schema_num_outputs + num_inputs
+            actual_len = len(strategies_with_placeholders[0])
+            if actual_len != expected_len:
+                raise AssertionError(
+                    f"Strategy length {actual_len} != expected {expected_len} "
+                    f"(schema_outputs={schema_num_outputs} + inputs={num_inputs}) "
+                    f"for {op_schema.op}. Strategies must include placements "
+                    f"for all outputs, args, and tensor kwargs."
+                )
+
+        # Compute num_outputs from strategy structure or output_tensor_meta
+        if len(strategies_with_placeholders) > 0:
+            num_outputs = len(strategies_with_placeholders[0]) - num_inputs
+        elif output_tensor_meta is None:
+            num_outputs = 0
+        elif isinstance(output_tensor_meta, TensorMeta):
+            num_outputs = 1
+        else:
+            num_outputs = len(output_tensor_meta)
+        self.num_outputs = num_outputs
+
+        strategies_with_placeholders = _insert_single_dim_replication_strategy(
+            strategies_with_placeholders,
+            num_outputs,
+            num_inputs,
+            output_tensor_meta,
+        )
+
+        unique_input_placements = _get_unique_placements(op_schema)
+        self.expanded_strategies = _fill_single_dim_strategy_placeholders(
+            unique_input_placements, strategies_with_placeholders
+        )
+
+        # Build strategy lookup: map input placements -> output placements
+        self.strategy_lookup = {}
+        for strategy in self.expanded_strategies:
+            input_key = tuple(strategy[num_outputs:])
+            if input_key not in self.strategy_lookup:
+                self.strategy_lookup[input_key] = tuple(strategy[:num_outputs])
+
+        # Precompute allowed placements per input from the expanded rules
+        self.allowed_sharding_per_input: dict[int, set[Shard | _StridedShard]] = (
+            defaultdict(set)
+        )
+        self.allowed_partial_per_input: dict[int, set[Placement]] = defaultdict(set)
+        for strategy in self.expanded_strategies:
+            for input_idx in range(num_inputs):
+                p = strategy[num_outputs + input_idx]
+                if p is None:
+                    continue
+                if _is_sharding(p):
+                    self.allowed_sharding_per_input[input_idx].add(p)
+                elif isinstance(p, Partial):
+                    self.allowed_partial_per_input[input_idx].add(p)
+
+        # Resolve output tensor_meta per output index
+        if output_tensor_meta is None:
+            self.output_metas = (None,) * max(num_outputs, 0)
+        elif isinstance(output_tensor_meta, TensorMeta):
+            self.output_metas = (output_tensor_meta,)
+        else:
+            self.output_metas = tuple(output_tensor_meta)
+
+    def try_propagate(
+        self,
+        mesh: DeviceMesh,
+        input_placements: tuple[tuple[Placement, ...], ...],
+        input_specs: list[DTensorSpec],
+    ) -> OpStrategy | None:
+        """Try to match input placements against single-dim strategy rules on every mesh dim.
+
+        Checks whether the given input placements independently match a rule in
+        strategy_lookup on each mesh dimension, and that all inputs are shardable
+        with those placements. If so, returns an OpStrategy with the matched output
+        placements and zero redistribute costs.
+        """
+        from torch.distributed.tensor._ops.utils import is_tensor_shardable
+
+        selected_output_placements: list[tuple[Placement | None, ...]] = []
+        for mesh_dim in range(mesh.ndim):
+            input_placements_for_dim = tuple(
+                placements[mesh_dim] for placements in input_placements
+            )
+            output_for_dim = self.strategy_lookup.get(input_placements_for_dim)
+            if output_for_dim is not None:
+                selected_output_placements.append(output_for_dim)
+            else:
+                return None
+
+        arg_specs = [
+            DTensorSpec(mesh, placements, tensor_meta=input_spec.tensor_meta)
+            for placements, input_spec in zip(input_placements, input_specs)
+        ]
+        if not all(
+            is_tensor_shardable(
+                spec.tensor_meta.shape,
+                spec,
+                allow_unbacked_sharding=self.allow_unbacked_sharding,
+            )
+            or (self.allow_uneven_sharding and input_spec.placements == spec.placements)
+            for spec, input_spec in zip(arg_specs, input_specs)
+            if spec.tensor_meta is not None
+        ):
+            return None
+
+        output_spec = (
+            _build_output_specs(
+                mesh,
+                selected_output_placements,
+                self.num_outputs,
+                self.output_metas,
+            )
+            if self.num_outputs > 0
+            else None
+        )
+        return OpStrategy(
+            [
+                OpSpec(
+                    output_specs=output_spec,
+                    input_specs=arg_specs,
+                    redistribute_cost=[[0.0] for _ in input_specs],
+                )
+            ]
+        )
+
+
 def _expand_single_dim_strategy_to_mesh(
     mesh: DeviceMesh,
     op_schema: OpSchema,
@@ -209,43 +528,8 @@ def _create_expanded_strategy_impl(
         def expanded_strategy(
             op: OpOverload, args_schema: ArgsType, kwargs_schema: KwargsType
         ) -> StrategyType:
-            # Note: op_schema vs [args_schema, kwargs_schema]
-            # -----------------------------------------------
-            # Inside `expanded_strategy function we purposefully have access to 2 similar structures.
-            # 1) (op, args_schema, kwargs_schema): This is all the single_dim_strategy is allowed to see.
-            # importantly, it does not contain information about input placements or meshes - just TensorMeta.
-            # 2) op_schema - captured from the parent scope, this contains the input placement and mesh info, needed
-            # to actually perform expansion.
-            unique_input_placements = _get_unique_placements(op_schema)
-            num_inputs = _get_num_tensor_inputs(op_schema)
-
-            # Compute num_outputs from output_tensor_meta
-            if output_tensor_meta is None:
-                num_outputs = 0
-            elif isinstance(output_tensor_meta, TensorMeta):
-                num_outputs = 1
-            else:
-                num_outputs = len(output_tensor_meta)
-
-            # Note: Trees vs Flat Lists
-            # -------------------------
-            # op_schema.args_schema may contain a TupleStrategy with child strategies for List[Tensor] inputs.
-            # args_schema has corresponding TupleStrategy, but with TensorSpec in place of child strategies.
-            # CURRENTLY: single_dim_strategy will return a flat list of Placements for each strategy, where any
-            # input tuple strategies have been inlined.  I'm not sure if we want to keep doing this, or preserve a pytree
-            # structure here.  I'm following the convention in the current DTensor sharding strategies for now.
-            # Inside expanded_strategy, we need to carefully align the OpStrategies / Specs from op_schema which are _not_
-            # flattened, with the flat Placement list returned from single_dim strategy.
-            strategies_over_one_mesh_dim = strategy_info.func(
-                op, args_schema, kwargs_schema
-            )
-            strategies_over_one_mesh_dim = _insert_single_dim_replication_strategy(
-                strategies_over_one_mesh_dim, num_outputs, num_inputs
-            )
-            expanded_strategies_over_one_mesh_dim = (
-                _fill_single_dim_strategy_placeholders(
-                    unique_input_placements, strategies_over_one_mesh_dim
-                )
+            prepared_strategy = _PreparedSingleDimStrategy(
+                strategy_info, op_schema, output_tensor_meta
             )
 
             # Detect inplace ops by checking if the base op name ends with '_'
@@ -253,14 +537,18 @@ def expanded_strategy(
             base_name = op_name.split("::")[1].split(".")[0]
             is_inplace = base_name.endswith("_")
 
+            element_mesh = prepared_strategy.element_mesh or mesh
+
             return expand_to_full_mesh_op_strategy(
-                mesh,
+                element_mesh,
                 op_schema,
-                cast(list[PlacementList], expanded_strategies_over_one_mesh_dim),
+                prepared_strategy.expanded_strategies,
                 output_tensor_meta=output_tensor_meta,
                 inplace_op=is_inplace,
-                input_index=num_outputs,
-                allow_unbacked_sharding=strategy_info.allow_unbacked_sharding,
+                input_index=prepared_strategy.num_outputs,
+                allow_unbacked_sharding=prepared_strategy.allow_unbacked_sharding,
+                allow_uneven_sharding=prepared_strategy.allow_uneven_sharding,
+                different_mesh_args=prepared_strategy.remapped_different_mesh_args,
             )
 
         return expanded_strategy
@@ -282,12 +570,14 @@ def _create_expanded_strategy(
             # Unhashable types (SymInts), skip caching
             return _create_expanded_strategy_impl(op_schema, output_tensor_meta)
 
-    def _translate_foreach_op_schema(
-        op_schema: OpSchema, output_tensor_meta: Sequence[TensorMeta], index: int
-    ) -> tuple[OpSchema, TensorMeta]:
-        """Translate foreach op to per-element version of schema."""
+    def _translate_list_op_schema(
+        op_schema: OpSchema,
+        output_tensor_meta: Sequence[TensorMeta] | None,
+        index: int,
+    ) -> tuple[OpSchema, TensorMeta | None]:
+        """Translate foreach/fused op to per-element version of schema."""
         op_parts = str(op_schema.op).split(".")
-        base_op_name = op_parts[-2].replace("_foreach_", "")
+        op_name = op_parts[-2]
         foreach_variant = op_parts[-1]
 
         # select per-element inputs, outputs
@@ -297,7 +587,30 @@ def _translate_foreach_op_schema(
             (op_schema.args_schema, op_schema.kwargs_schema),
             is_leaf=lambda x: isinstance(x, TupleStrategy),
         )
-        target_output_meta = output_tensor_meta[index]
+        # For inplace ops, output_tensor_meta is None
+        target_output_meta = (
+            output_tensor_meta[index] if output_tensor_meta is not None else None
+        )
+
+        # Strip the prefix to get the base op name and find the per-element op.
+        # Fused ops (e.g. _fused_adam) have no per-element ATen equivalent,
+        # so we keep the original op unchanged.
+        if op_name.startswith("_foreach_"):
+            base_op_name = op_name.replace("_foreach_", "", 1)
+        elif op_name.startswith("_amp_foreach_"):
+            base_op_name = op_name.replace("_amp_foreach_", "", 1)
+        else:
+            # Fused ops or unknown: keep original op, no translation
+            target_op = op_schema.op
+            op_schema = OpSchema(
+                target_op,  # type: ignore[arg-type]
+                args_schema=tuple(target_args),
+                kwargs_schema=op_schema.kwargs_schema,
+            )
+            return op_schema, target_output_meta
+
+        # Strip trailing underscore for inplace ops
+        base_op_name = base_op_name.removesuffix("_")
 
         # figure out target op variant
         variant_map = {
@@ -346,7 +659,7 @@ def expanded_foreach_strategy(
 
         child_strategies: list[StrategyType] = []
         for tensorlist_i in range(tensorlist_len):
-            per_index_schema, per_index_output_meta = _translate_foreach_op_schema(
+            per_index_schema, per_index_output_meta = _translate_list_op_schema(
                 op_schema,
                 output_tensor_meta,  # type: ignore[arg-type]
                 tensorlist_i,
@@ -363,16 +676,27 @@ def expanded_foreach_strategy(
         return TupleStrategy(children=child_strategies)
 
     # TODO maybe this could be helped by adding a new 'tag' to the OpOverload?
-    if op_schema.op.name().startswith("aten::_foreach_"):
+    # Only use the foreach path if the op has TupleStrategy inputs (i.e., actual
+    # list-of-tensor args). The name prefix alone is insufficient because ops like
+    # _fused_rms_norm share the "_fused_" prefix but are not foreach/fused-optimizer ops.
+    op_name = op_schema.op.name()
+    has_tuple_strategy = any(
+        isinstance(arg, TupleStrategy) for arg in op_schema.args_schema
+    )
+    if has_tuple_strategy and op_name.startswith(
+        ("aten::_foreach_", "aten::_amp_foreach_", "aten::_fused_")
+    ):
         return expanded_foreach_strategy
 
     return _create_expanded_strategy(op_schema, output_tensor_meta)
 
 
 def register_single_dim_strategy(
-    op: Union[torch._ops.OpOverload, list[torch._ops.OpOverload]],
-    schema_info: Optional[RuntimeSchemaInfo] = None,
+    op: torch._ops.OpOverload | list[torch._ops.OpOverload],
+    schema_info: RuntimeSchemaInfo | None = None,
     allow_unbacked_sharding: bool | None = None,
+    allow_uneven_sharding: bool = False,
+    different_mesh_args: list[int] | None = None,
 ) -> Callable[[_SingleDimStrategyFunc], _SingleDimStrategyFunc]:
     """
     Registers a single_dim_strategy function for the given op.
@@ -420,8 +744,341 @@ def wrapper(impl):
         info = _SingleDimStrategyInfo(
             func=impl,
             allow_unbacked_sharding=allow_unbacked_sharding,
+            allow_uneven_sharding=allow_uneven_sharding,
+            different_mesh_args=different_mesh_args,
         )
         registration_wrapper(info)
         return impl
 
     return wrapper
+
+
+@dataclass(order=True)
+class _PQEntry:
+    """Priority queue entry for the Dijkstra search in _dijkstra_expand_single_dim_strategy_to_mesh.
+
+    Ordered by (cost, counter) for heap comparison. The counter breaks ties
+    in FIFO order so that entries with equal cost are explored in insertion
+    order rather than by arbitrary tuple comparison on placements.
+    """
+
+    cost: float
+    counter: int
+    # Per-input placement tuples representing the current search state.
+    placements: tuple[tuple[Placement, ...], ...] = field(compare=False)
+    # History of (input_idx, mesh_dim, old_placement, new_placement) transitions
+    # from the initial state to this state, used for debugging.
+    transitions: list[tuple[int, int, Placement, Placement]] = field(compare=False)
+    # Accumulated redistribute cost per input (sum of incremental step costs).
+    per_input_costs: tuple[float, ...] = field(compare=False)
+    # Current communication bytes (in GB) per input, updated as placements change.
+    per_input_comm_bytes_gb: tuple[float, ...] = field(compare=False)
+
+
+def _get_neighbor_placements(
+    allowed_sharding: set[Shard | _StridedShard],
+    allowed_partial: set[Placement],
+    current: Placement,
+    input_placements: tuple[Placement, ...],
+    mesh_dim: int,
+) -> list[Placement]:
+    """Return valid one-shot placement transitions for one input on one mesh dim.
+
+    DTensor placements are applied left-to-right, so a tensor dim sharded on
+    multiple mesh dims has a specific nesting order. A one-shot collective on
+    mesh_dim M can only produce the correct data layout if no mesh dim to the
+    RIGHT of M already shards the same tensor dim. For example, going from
+    (R, S(0)) to (S(0), S(0)) via local chunk on mesh dim 0 produces a
+    strided-shard layout, not the correct left-to-right (S(0), S(0)) layout.
+
+    Transition rules:
+    - Replicate -> Shard(d): free local chunk, valid if d not sharded to the right
+    - Replicate -> Partial: local view, always valid
+    - Shard(d) -> Replicate: allgather, valid if d not sharded to the right
+    - Shard(d1) -> Shard(d2): all-to-all, valid if neither d1 nor d2 sharded right
+    - Partial -> Replicate: allreduce, always valid
+    - Partial -> Shard(d): reduce-scatter, valid if d not sharded to the right
+    """
+    # Note: circular import
+    from torch.distributed.tensor.placement_types import Partial
+
+    # Tensor dims sharded by mesh dims to the right of this one.
+    right_shard_dims: set[int] = set()
+    for i in range(mesh_dim + 1, len(input_placements)):
+        p = input_placements[i]
+        if _is_sharding(p):
+            right_shard_dims.add(p.dim)
+
+    neighbors: list[Placement] = []
+
+    if isinstance(current, Replicate):
+        neighbors.extend(s for s in allowed_sharding if s.dim not in right_shard_dims)
+        neighbors.extend(allowed_partial)
+
+    elif _is_sharding(current):
+        cur_dim_ok = current.dim not in right_shard_dims
+        if cur_dim_ok:
+            neighbors.append(Replicate())
+        for s in allowed_sharding:
+            if s != current and cur_dim_ok and s.dim not in right_shard_dims:
+                neighbors.append(s)
+
+    elif isinstance(current, Partial):
+        neighbors.append(Replicate())
+        neighbors.extend(s for s in allowed_sharding if s.dim not in right_shard_dims)
+
+    return neighbors
+
+
+def _dijkstra_expand_single_dim_strategy_to_mesh(
+    mesh: DeviceMesh,
+    op_schema: OpSchema,
+    single_dim_strategy: _SingleDimStrategyInfo
+    | Callable[
+        [OpOverload, ArgsType, KwargsType], list[list[Placement | _ShardingPlaceholder]]
+    ],
+    output_tensor_meta: TensorMeta | Sequence[TensorMeta | None] | None = None,
+    _collect_all_matches: set[tuple[tuple[Placement, ...], ...]] | None = None,
+) -> OpStrategy | None:
+    """
+    Find the lowest cost sharding for the given op_schema.
+
+    Uses a Dijkstra-like priority-queue search over input placement states. Each
+    state is a tuple of per-input placement tuples, and neighbors are generated
+    by changing one placement on one mesh dim for one input. The search
+    terminates when a state matches a single-dim strategy on every mesh dim.
+
+    This avoids the O(S^N) exhaustive expansion of _expand_single_dim_strategy_to_mesh
+    (S = single-dim strategies, N = mesh dims).  Benchmarks with mm on fake
+    process groups show:
+
+        1D(4):     S^N=8,   avg 0.2ms
+        2D(2,2):   S^N=64,  avg 2.3ms
+        3D(2,2,2): S^N=512, avg 41ms, worst 392ms
+
+    The step count is small (avg 0.6-2.0 pops) but per-step cost is dominated
+    by cost computation.  Each transition computes an incremental cost via
+    _compute_placement_transition_cost for the single changed placement, matching
+    the per-step costs used by graph-based transform info planning.
+
+    Returns None if any input has _StridedShard placement, signaling the caller
+    to fall back to full expansion.
+
+    Args:
+        _collect_all_matches: Testing-only. When non-None, exhaustively explores the
+            full transition graph, adding every shardable match to the set. Still
+            returns the optimal (first) match.
+    """
+    # Extract input DTensorSpecs from OpStrategy-wrapped args.
+    # Fall back for TupleStrategy (e.g. index tensors in index_put) since the PQ
+    # search doesn't model variable-length tuple inputs.
+    input_specs: list[DTensorSpec] = []
+    for arg in op_schema.args_schema:
+        if isinstance(arg, OpStrategy):
+            if len(arg.strategies) != 1:
+                raise AssertionError
+            input_specs.append(arg.strategies[0].output_spec)
+        elif isinstance(arg, TupleStrategy):
+            return None
+
+    # Fall back if any kwargs are tensor inputs — the PQ search only tracks
+    # positional tensor args and would miss redistribute costs for kwargs.
+    for kwarg in op_schema.kwargs_schema.values():
+        if isinstance(kwarg, (OpStrategy, TupleStrategy)):
+            return None
+
+    if len(input_specs) == 0:
+        raise AssertionError("broken input")
+    num_inputs = len(input_specs)
+
+    # Fall back to full expansion if any input has _StridedShard or symbolic shapes
+    # (symbolic shapes produce SymFloat costs that can't be compared in the PQ)
+    for spec in input_specs:
+        if any(isinstance(p, _StridedShard) for p in spec.placements):
+            return None
+        if spec.tensor_meta is not None and any(
+            isinstance(s, torch.SymInt) for s in spec.tensor_meta.shape
+        ):
+            return None
+
+    prepared_strategy = _PreparedSingleDimStrategy(
+        single_dim_strategy, op_schema, output_tensor_meta, num_inputs=num_inputs
+    )
+
+    initial_placements = tuple(spec.placements for spec in input_specs)
+    first_result: OpStrategy | None = None
+
+    # Fast path: if initial placements already match a strategy, skip search
+    fast_result = prepared_strategy.try_propagate(mesh, initial_placements, input_specs)
+    if fast_result is not None:
+        fast_result._pq_transitions = []  # type: ignore[attr-defined]
+        if _collect_all_matches is not None:
+            _collect_all_matches.add(initial_placements)
+            first_result = fast_result
+        else:
+            return fast_result
+
+    # Pre-compute mesh topology and per-input comm bytes for cost computation.
+    # comm_bytes_gb reflects the local shard size given current placements;
+    # it's tracked per PQ entry and updated as placements change.
+    mesh_topo = MeshTopoInfo.build_from_mesh(mesh)
+    initial_comm_bytes_gb: list[float] = []
+    for spec in input_specs:
+        if spec.tensor_meta is None:
+            raise AssertionError
+        total_bytes = spec.tensor_meta.dtype.itemsize * math.prod(
+            spec.tensor_meta.shape
+        )
+        # TODO: is_shard() misses _StridedShard, use spec.num_shards instead.
+        # Not fixing yet: the overestimate biases Dijkstra toward redistributing
+        # away from _StridedShard, which is the safer default until _StridedShard
+        # is fully validated.
+        num_shards = 1
+        for i, p in enumerate(spec.placements):
+            if p.is_shard():
+                num_shards *= mesh_topo.mesh_dim_devices[i]
+        initial_comm_bytes_gb.append(total_bytes / num_shards / (1024**3))
+
+    pq: list[_PQEntry] = []
+    visited: set[tuple[tuple[Placement, ...], ...]] = set()
+    next_counter = count()
+
+    initial_per_input_costs = (0.0,) * num_inputs
+    initial_per_input_comm_bytes = tuple(initial_comm_bytes_gb)
+    heapq.heappush(
+        pq,
+        _PQEntry(
+            0.0,
+            next(next_counter),
+            initial_placements,
+            [],
+            initial_per_input_costs,
+            initial_per_input_comm_bytes,
+        ),
+    )
+
+    def _push_neighbor(
+        input_idx: int,
+        mesh_dim: int,
+        new_placement: Placement,
+        source: _PQEntry,
+    ) -> None:
+        new_input_placements = [list(ps) for ps in source.placements]
+        old_placement = new_input_placements[input_idx][mesh_dim]
+        new_input_placements[input_idx][mesh_dim] = new_placement
+        candidate_placements = tuple(tuple(ps) for ps in new_input_placements)
+        if candidate_placements in visited:
+            return
+        # Check that the NET transition (original -> proposed) is feasible.
+        # Individual hops may each be valid (e.g. S->R then R->P) while the
+        # net redistribution (S->P) is unsupported by the runtime planner.
+        original_p = initial_placements[input_idx][mesh_dim]
+        net_cost, _ = _compute_placement_transition_cost(
+            original_p,
+            new_placement,
+            mesh_topo,
+            mesh_dim,
+            initial_comm_bytes_gb[input_idx],
+        )
+        if net_cost == float("inf"):
+            return
+        step_cost, new_comm_bytes = _compute_placement_transition_cost(
+            old_placement,
+            new_placement,
+            mesh_topo,
+            mesh_dim,
+            source.per_input_comm_bytes_gb[input_idx],
+        )
+        if step_cost == float("inf"):
+            return
+        changed_cost = source.per_input_costs[input_idx] + step_cost
+        new_per_input_costs = (
+            source.per_input_costs[:input_idx]
+            + (changed_cost,)
+            + source.per_input_costs[input_idx + 1 :]
+        )
+        new_per_input_comm_bytes = (
+            source.per_input_comm_bytes_gb[:input_idx]
+            + (new_comm_bytes,)
+            + source.per_input_comm_bytes_gb[input_idx + 1 :]
+        )
+        new_cost = sum(new_per_input_costs)
+        new_transitions = source.transitions + [
+            (input_idx, mesh_dim, old_placement, new_placement)
+        ]
+        heapq.heappush(
+            pq,
+            _PQEntry(
+                new_cost,
+                next(next_counter),
+                candidate_placements,
+                new_transitions,
+                new_per_input_costs,
+                new_per_input_comm_bytes,
+            ),
+        )
+
+    while pq:
+        candidate = heapq.heappop(pq)
+
+        if candidate.placements in visited:
+            continue
+        visited.add(candidate.placements)
+
+        match_result = prepared_strategy.try_propagate(
+            mesh, candidate.placements, input_specs
+        )
+        if match_result is not None:
+            # Use pre-computed per-input costs from the PQ search instead of
+            # recomputing via generate_redistribute_costs -> _gen_transform_infos.
+            match_spec = match_result.strategies[0]
+            if match_spec.input_specs is None:
+                raise AssertionError
+            op_spec = OpSpec(
+                output_specs=match_spec.output_specs,
+                input_specs=list(match_spec.input_specs),
+                redistribute_cost=[[cost] for cost in candidate.per_input_costs],
+            )
+
+            exhaustive = len(prepared_strategy.expanded_strategies) ** mesh.ndim
+            logger.debug(
+                "returning cost=%f %s, visited=%d, exhaustive=%d, transitions=%s",
+                candidate.cost,
+                op_spec,
+                len(visited),
+                exhaustive,
+                candidate.transitions,
+            )
+            result = OpStrategy([op_spec])
+            result._pq_transitions = candidate.transitions  # type: ignore[attr-defined]
+            if _collect_all_matches is not None:
+                _collect_all_matches.add(candidate.placements)
+                if first_result is None:
+                    first_result = result
+            else:
+                return result
+
+        # Generate neighbor states
+        for mesh_dim in range(mesh.ndim):
+            for input_idx in range(len(candidate.placements)):
+                current_p = candidate.placements[input_idx][mesh_dim]
+                for neighbor_p in _get_neighbor_placements(
+                    prepared_strategy.allowed_sharding_per_input[input_idx],
+                    prepared_strategy.allowed_partial_per_input[input_idx],
+                    current_p,
+                    candidate.placements[input_idx],
+                    mesh_dim,
+                ):
+                    _push_neighbor(input_idx, mesh_dim, neighbor_p, candidate)
+
+    if _collect_all_matches is not None and first_result is not None:
+        return first_result
+
+    logger.warning(
+        "Dijkstra search exhausted without finding a valid strategy for "
+        "%s on %s (explored %d combinations); falling back to full expansion",
+        op_schema,
+        mesh,
+        len(visited),
+    )
+    return None
diff --git a/torch/distributed/tensor/_ops/strategy_validation.py b/torch/distributed/tensor/_ops/strategy_validation.py
index e5f04d611031c..3d94de2d1455d 100644
--- a/torch/distributed/tensor/_ops/strategy_validation.py
+++ b/torch/distributed/tensor/_ops/strategy_validation.py
@@ -11,65 +11,70 @@
 
 Run as a module to compare DTensor rules against ground truth:
     python -m torch.distributed.tensor._ops.strategy_validation --op add
+    python -m torch.distributed.tensor._ops.strategy_validation --op relu,mul,add
+    python -m torch.distributed.tensor._ops.strategy_validation --op "nn.functional.*"
     python -m torch.distributed.tensor._ops.strategy_validation --all-registered
     python -m torch.distributed.tensor._ops.strategy_validation --op div --incorrect-only
+    python -m torch.distributed.tensor._ops.strategy_validation --op add --show-repro
 """
 
+import argparse
+import fnmatch
+import itertools
 import re
+import sys
+import time
+from collections import defaultdict
+from collections.abc import Callable
 from dataclasses import dataclass, field
 from typing import Any
 
 import torch
-from torch.distributed.tensor import Replicate
+import torch.distributed as dist
+import torch.testing._internal.common_methods_invocations as common_ops
+from torch._ops import OpOverload
+from torch.distributed._local_tensor import LocalTensor, LocalTensorMode
+from torch.distributed.device_mesh import DeviceMesh, init_device_mesh
+from torch.distributed.tensor import distribute_tensor, DTensor, Replicate
+from torch.distributed.tensor._decompositions import DecompShardingStrategy
+from torch.distributed.tensor._dtensor_spec import TensorMeta
+from torch.distributed.tensor._op_schema import (
+    DTensorSpec,
+    OpSchema,
+    OpSpec,
+    OpStrategy,
+)
+from torch.distributed.tensor._ops.single_dim_strategy import _ShardingPlaceholder
 from torch.distributed.tensor.placement_types import Partial, Placement, Shard
+from torch.testing._internal.common_methods_invocations import op_db, SampleInput
+from torch.testing._internal.opinfo import core as opinfo_core
 from torch.utils import _pytree as pytree
 
 
-# A combo key is (input_placement_strs, output_placement_str)
-ComboKey = tuple[tuple[str, ...], str]
+# A combo key is (input_placement_strs, output_placement_strs)
+# For single-output ops: (("S(0)",), ("P(min)",))
+# For multi-output ops:  (("S(0)",), ("P(min)", "P(min)"))
+ComboKey = tuple[tuple[str, ...], tuple[str, ...]]
 
 # Partial reduce ops to enumerate
 PARTIAL_REDUCE_OPS = ["sum", "avg", "min", "max"]
 
-# Ops to skip in validation because ground truth comparison is not meaningful.
-# These produce non-deterministic or uninitialized outputs.
-# TODO: can we get this list directly from opinfo tags, or something?
-SKIP_OPS = frozenset(
-    [
-        "bernoulli",  # Random sampling
-        "empty_like",  # Uninitialized memory
-        "new_empty",  # Uninitialized memory
-        "new_empty_strided",  # Uninitialized memory
-        "normal",  # Random sampling
-        "rand_like",  # Random sampling
-        "randint_like",  # Random sampling
-        "randn_like",  # Random sampling
-        "uniform",  # Random sampling
-    ]
-)
-
-
-@dataclass
-class PlacementCombination:
-    """Represents a combination of input and output placements."""
-
-    input_placements: tuple[Placement, ...]  # One placement per input tensor
-    output_placement: Placement  # Placement for the output tensor
-
-    def __hash__(self):
-        return hash(
-            (tuple(str(p) for p in self.input_placements), str(self.output_placement))
-        )
+SKIP_OPS: dict[str, str] = {
+    "bernoulli": "non-deterministic (random sampling)",
+    "empty_like": "uninitialized memory",
+    "exponential": "non-deterministic (random sampling)",
+    "new_empty": "uninitialized memory",
+    "new_empty_strided": "uninitialized memory",
+    "nn.functional.dropout": "non-deterministic (random masking)",
+    "normal": "non-deterministic (random sampling)",
+    "rand_like": "non-deterministic (random sampling)",
+    "randint_like": "non-deterministic (random sampling)",
+    "randn_like": "non-deterministic (random sampling)",
+    "uniform": "non-deterministic (random sampling)",
+}
 
-    def __eq__(self, other):
-        if not isinstance(other, PlacementCombination):
-            return NotImplemented
-        return tuple(str(p) for p in self.input_placements) == tuple(
-            str(p) for p in other.input_placements
-        ) and str(self.output_placement) == str(other.output_placement)
 
-    def __str__(self):
-        return f"inputs={placement_tuple_to_str(self.input_placements)}, output={placement_tuple_to_str((self.output_placement,))}"
+PlacementCombination = tuple[tuple[Placement, ...], tuple[Placement, ...]]
 
 
 @dataclass
@@ -77,15 +82,16 @@ class Discrepancy:
     """Represents a discrepancy between ground truth and DTensor's rules."""
 
     input_placements: tuple[str, ...]
-    output_placement: str
+    output_placements: tuple[str, ...]
     sample_idx: int
     input_shapes: tuple[tuple[int, ...], ...]
     discrepancy_type: str  # "false_positive" or "false_negative"
     error_msg: str = ""
     scalar_args: tuple[Any, ...] = ()
     scalar_kwargs: dict[str, Any] = field(default_factory=dict)
-    aten_op: Any = None
+    aten_op: OpOverload | None = None
     variant: str = ""
+    sample: SampleInput | None = None
 
 
 @dataclass
@@ -100,21 +106,31 @@ class ComparisonStats:
     false_negatives: list[Discrepancy] = field(
         default_factory=list
     )  # Ground truth valid, DTensor has no rule
+    total_samples: int = 0
+    total_combinations: int = 0
+    skip_reasons: dict[str, int] = field(default_factory=dict)
+    no_dtensor_support: bool = False
+    # Per aten op variant breakdown (e.g. "aten.min.dim" -> 5)
+    true_positives_by_op: dict[str, int] = field(default_factory=dict)
 
 
-def placement_tuple_to_str(placements: tuple[Placement, ...]) -> str:
-    """Convert a tuple of placements to a readable string."""
-    parts: list[str] = []
-    for p in placements:
-        if isinstance(p, Shard):
-            parts.append(f"S({p.dim})")
-        elif isinstance(p, Replicate):
-            parts.append("R")
-        elif isinstance(p, Partial):
-            parts.append(f"P({p.reduce_op})")
-        else:
-            parts.append(str(p))
-    return "(" + ", ".join(parts) + ")"
+@dataclass
+class _FalsePositiveMitigations:
+    """Bundle of sample variants used to detect false positive validations.
+
+    Contains negated and non-rounded variants of a sample. These are used to
+    re-test combinations that pass the primary validation, catching cases where
+    sign patterns or rounding modes mask real differences.
+    """
+
+    negated_sample: SampleInput | None = None
+    negated_tensors: list[tuple[str, torch.Tensor]] | None = None
+    negated_ground_truth: torch.Tensor | list[torch.Tensor] | None = None
+    non_rounded_sample: SampleInput | None = None
+    non_rounded_ground_truth: torch.Tensor | list[torch.Tensor] | None = None
+    non_rounded_negated_sample: SampleInput | None = None
+    non_rounded_negated_tensors: list[tuple[str, torch.Tensor]] | None = None
+    non_rounded_negated_ground_truth: torch.Tensor | list[torch.Tensor] | None = None
 
 
 def parse_placement(s: str) -> Placement | None:
@@ -143,6 +159,7 @@ def is_fully_replicated(placements: tuple[Placement, ...]) -> bool:
 
 def is_trivial_shard(p: Placement, tensor_shape: tuple[int, ...]) -> bool:
     """Check if placement is a Shard on a size-1 dimension."""
+    # NOTE: isinstance(_, Shard) does not match _StridedShard; see _is_shard_like().
     return (
         isinstance(p, Shard) and p.dim < len(tensor_shape) and tensor_shape[p.dim] == 1
     )
@@ -180,7 +197,7 @@ def normalize_placement_str(p_str: str, shape: tuple[int, ...]) -> str:
 def normalize_combo_key(
     combo_key: ComboKey,
     input_shapes: tuple[tuple[int, ...], ...],
-    output_shape: tuple[int, ...],
+    output_shapes: tuple[tuple[int, ...], ...],
 ) -> ComboKey:
     """
     Normalize a combo_key by converting trivial shards to Replicate.
@@ -190,25 +207,26 @@ def normalize_combo_key(
     - S(0), R -> R on input [1,4] becomes R, R -> R
 
     Args:
-        combo_key: (input_placement_strs, output_placement_str) tuple
+        combo_key: (input_placement_strs, output_placement_strs) tuple
         input_shapes: Shapes of input tensors
-        output_shape: Shape of output tensor
+        output_shapes: Shapes of output tensors
 
     Returns:
         Normalized combo_key with trivial shards converted to Replicate
     """
-    input_placement_strs, output_placement_str = combo_key
+    input_placement_strs, output_placement_strs = combo_key
 
-    # Normalize input placements
     normalized_inputs = tuple(
         normalize_placement_str(p_str, shape)
         for p_str, shape in zip(input_placement_strs, input_shapes)
     )
 
-    # Normalize output placement
-    normalized_output = normalize_placement_str(output_placement_str, output_shape)
+    normalized_outputs = tuple(
+        normalize_placement_str(p_str, shape)
+        for p_str, shape in zip(output_placement_strs, output_shapes)
+    )
 
-    return (normalized_inputs, normalized_output)
+    return (normalized_inputs, normalized_outputs)
 
 
 def get_1d_input_placements_for_tensor(
@@ -244,7 +262,9 @@ def get_1d_output_placements_for_tensor(t: torch.Tensor) -> list[Placement]:
     return placements
 
 
-def extract_tensors_from_sample(sample_input: Any) -> list[tuple[str, torch.Tensor]]:
+def extract_tensors_from_sample(
+    sample_input: SampleInput,
+) -> list[tuple[str, torch.Tensor]]:
     """
     Extract all tensor arguments from a SampleInput.
     Returns a list of (name, tensor) pairs.
@@ -267,3 +287,2018 @@ def _collect(x):
     pytree.tree_map(_collect, sample_input.kwargs)
 
     return tensors
+
+
+def _checkerboard_mask(
+    tensor: torch.Tensor, tensor_idx: int = 0, mask_shift: int = 0
+) -> torch.Tensor:
+    """Checkerboard mask that alternates in every dimension.
+
+    Unlike flat-index % 2, which can be uniform along even-stride dimensions
+    (causing all elements in a reduction group to get the same offset), the
+    checkerboard uses sum-of-coordinates mod 2 so adjacent elements differ
+    along every axis.
+
+    Returns a flat bool tensor of shape (numel,).
+    """
+    if tensor.ndim == 0:
+        return torch.tensor([(tensor_idx + mask_shift) % 2 == 0], device=tensor.device)
+    coords = [torch.arange(s, device=tensor.device) for s in tensor.shape]
+    grids = torch.meshgrid(*coords, indexing="ij")
+    coord_sum = grids[0].clone()
+    for g in grids[1:]:
+        coord_sum += g
+    return ((coord_sum + tensor_idx + mask_shift) % 2 == 0).flatten()
+
+
+def _create_partial_input(
+    tensor: torch.Tensor,
+    placement: Partial,
+    world_size: int,
+    tensor_idx: int = 0,
+    mask_shift: int = 0,
+) -> LocalTensor:
+    """
+    Create a LocalTensor with values that reduce to the original tensor.
+
+    Uses asymmetric splits to avoid coincidental matches when combining
+    different Partial types.
+
+    For each placement combination, it creates local tensors that would
+    reduce to the original (e.g., for P(sum), splits values across ranks so
+    they sum back), runs the op on those local tensors, wraps the output as
+    a DTensor, redistributes to Replicate, and compares against ground
+    truth.
+
+    The main challenge is avoiding false positives where a rule appears
+    valid on a specific input but is actually incorrect. Several techniques
+    are used:
+
+    Asymmetric splits for P(sum)/P(avg): instead of splitting evenly
+    (tensor/2 per rank), uses a 60/40 ratio (varied by tensor index) so
+    that ops which are not truly linear don't accidentally produce
+    matching outputs.
+
+    Sign-varying offsets for P(sum)/P(avg): adds an offset that
+    alternates sign across elements, so local tensors have mixed positive
+    and negative values. Without this, proportional splits preserve the
+    sign pattern of the original tensor, causing non-linear ops like abs
+    to falsely validate P(sum)->P(sum).
+
+    Distinct magnitudes for P(min) vs P(max): P(min) offsets non-holding
+    ranks by +(range*2+1) while P(max) offsets by -(range*2+1), where
+    range is the tensor's value range. Using adaptive offsets that exceed
+    the value range ensures that index-returning ops (argmin/argmax)
+    produce different results on different ranks, correctly rejecting
+    P(min)/P(max) inputs for those ops. Using different signs for min vs
+    max prevents accidental cancellation when both appear in the same
+    combination.
+
+    Alternating rank ownership for P(min)/P(max): a multi-dimensional
+    checkerboard mask (sum of coordinates mod 2) controls which rank holds
+    the true value vs the offset value. Unlike a flat-index mask which can
+    have uniform parity along an even-stride dimension, the checkerboard
+    guarantees alternation along EVERY dimension. The mask_shift parameter
+    allows re-validation with the complementary mask to catch ops where
+    the result coincidentally matches.
+
+    """
+    reduce_op = placement.reduce_op
+    local_tensors: dict[int, torch.Tensor] = {}
+
+    if reduce_op in ("sum", "avg"):
+        base_ratio = 0.6 + 0.1 * (tensor_idx % 3)
+
+        # See docstring above: "Sign-varying offsets"
+        flat = tensor.flatten()
+        offset_mag = flat.abs() + 1.0
+        signs = torch.ones_like(flat)
+        # Use checkerboard mask so offset sign alternates in every dimension,
+        # not just along flat index (which can be uniform along even-stride dims).
+        signs[_checkerboard_mask(tensor, tensor_idx, mask_shift)] = -1.0
+        offset = (offset_mag * signs).reshape(tensor.shape)
+
+        scale = world_size if reduce_op == "avg" else 1
+        for r in range(world_size):
+            if r == 0:
+                local_tensors[r] = tensor.clone() * base_ratio * scale + offset
+            else:
+                local_tensors[r] = tensor.clone() * (
+                    (1 - base_ratio) / (world_size - 1)
+                ) * scale - offset / (world_size - 1)
+
+    elif reduce_op == "min":
+        # See docstring above: "Distinct Magnitudes" and "Alternating Rank Ownership"
+        flat = tensor.flatten()
+        # Offset must exceed the tensor's value range so that the mask pattern
+        # determines argmin/argmax, not the original values.
+        value_range = (flat.max() - flat.min()).item()
+        min_offset = value_range * 2 + 1
+        mask = _checkerboard_mask(tensor, tensor_idx, mask_shift)
+        for r in range(world_size):
+            if r == 0:
+                r_offset = torch.where(
+                    mask, torch.zeros_like(flat), torch.full_like(flat, min_offset)
+                )
+            else:
+                r_offset = torch.where(
+                    mask, torch.full_like(flat, min_offset), torch.zeros_like(flat)
+                )
+            local_tensors[r] = (flat + r_offset).reshape(tensor.shape)
+
+    elif reduce_op == "max":
+        # See docstring above: "Distinct Magnitudes" and "Alternating Rank Ownership"
+        flat = tensor.flatten()
+        value_range = (flat.max() - flat.min()).item()
+        max_offset = -(value_range * 2 + 1)
+        mask = _checkerboard_mask(tensor, tensor_idx, mask_shift)
+        for r in range(world_size):
+            if r == 0:
+                r_offset = torch.where(
+                    mask, torch.zeros_like(flat), torch.full_like(flat, max_offset)
+                )
+            else:
+                r_offset = torch.where(
+                    mask, torch.full_like(flat, max_offset), torch.zeros_like(flat)
+                )
+            local_tensors[r] = (flat + r_offset).reshape(tensor.shape)
+
+    else:
+        for r in range(world_size):
+            local_tensors[r] = tensor.clone()
+
+    # pyrefly: ignore [bad-argument-type, bad-argument-count]
+    return LocalTensor(local_tensors)
+
+
+def _shard_tensors(
+    tensors: list[tuple[str, torch.Tensor]],
+    input_placements: tuple[Placement, ...],
+    world_size: int,
+    mesh: DeviceMesh,
+    mask_shift: int = 0,
+) -> list[LocalTensor | torch.Tensor]:
+    """Create sharded LocalTensors from tensors according to placements."""
+    local_tensors: list[LocalTensor | torch.Tensor] = []
+    for tensor_idx, ((name, tensor), placement) in enumerate(
+        zip(tensors, input_placements)
+    ):
+        if isinstance(placement, Partial):
+            local_tensor = _create_partial_input(
+                tensor, placement, world_size, tensor_idx, mask_shift
+            )
+        elif isinstance(placement, Replicate):
+            _tmp = {r: tensor.clone() for r in range(world_size)}
+            # pyrefly: ignore [bad-argument-type, bad-argument-count]
+            local_tensor = LocalTensor(_tmp)
+        elif isinstance(placement, Shard):
+            shard_dim = placement.dim
+            chunks = tensor.tensor_split(world_size, dim=shard_dim)
+            _tmp = {
+                r: chunks[r].clone(memory_format=torch.contiguous_format)
+                for r in range(world_size)
+            }
+            # pyrefly: ignore [bad-argument-type, bad-argument-count]
+            local_tensor = LocalTensor(_tmp)
+        else:
+            dt = distribute_tensor(tensor.clone(), mesh, (placement,))
+            local_tensor = dt.to_local()
+        local_tensors.append(local_tensor)
+    return local_tensors
+
+
+def _compare_outputs(
+    local_output: Any,
+    ground_truth: torch.Tensor | list[torch.Tensor],
+    output_placements: tuple[Placement, ...],
+    mesh: DeviceMesh,
+    world_size: int,
+) -> tuple[bool, str]:
+    """Compare op output (wrapped as DTensor) against ground truth."""
+    if isinstance(local_output, (list, tuple)):
+        local_outputs = list(local_output)
+    else:
+        local_outputs = [local_output]
+
+    if isinstance(ground_truth, list):
+        ground_truths = ground_truth
+    else:
+        ground_truths = [ground_truth]
+
+    if len(local_outputs) != len(ground_truths):
+        return (
+            False,
+            f"Output count mismatch: got {len(local_outputs)}, "
+            f"expected {len(ground_truths)}",
+        )
+
+    if len(local_outputs) != len(output_placements):
+        return (
+            False,
+            f"Output count mismatch with placements: "
+            f"got {len(local_outputs)}, expected {len(output_placements)}",
+        )
+
+    for i, (local_out, gt, out_plc) in enumerate(
+        zip(local_outputs, ground_truths, output_placements)
+    ):
+        if not isinstance(local_out, torch.Tensor):
+            return False, f"Local output[{i}] is not a tensor: {type(local_out)}"
+
+        if not isinstance(local_out, LocalTensor):
+            return False, f"LocalTensor inputs produced non-LocalTensor output[{i}]"
+
+        output_dt = DTensor.from_local(
+            local_out,
+            mesh,
+            (out_plc,),
+            shape=gt.shape,
+            stride=gt.stride(),
+        )
+
+        if isinstance(out_plc, Replicate):
+            local_values = [local_out._local_tensors[r] for r in range(world_size)]
+            all_same = all(
+                torch.allclose(local_values[0], lv, atol=1e-5, rtol=1e-5)
+                for lv in local_values[1:]
+            )
+            if not all_same:
+                return (
+                    False,
+                    f"Replicate output[{i}] but local values differ across ranks",
+                )
+
+        full_output = output_dt.redistribute(mesh, (Replicate(),)).to_local()
+
+        if isinstance(full_output, LocalTensor):
+            full_output = full_output._local_tensors[0]
+
+        if gt.shape != full_output.shape:
+            return (
+                False,
+                f"Shape mismatch[{i}]: expected {gt.shape}, got {full_output.shape}",
+            )
+
+        if not torch.allclose(gt, full_output, atol=1e-5, rtol=1e-5, equal_nan=True):
+            max_diff = (gt - full_output).abs().max().item()
+            return False, f"Value mismatch[{i}]: max_diff={max_diff:.6f}"
+
+    return True, ""
+
+
+def validate_combination(
+    op: Callable[..., Any],
+    sample_input: SampleInput,
+    tensors: list[tuple[str, torch.Tensor]],
+    combination: PlacementCombination,
+    ground_truth: torch.Tensor | list[torch.Tensor],
+    world_size: int = 2,
+    mesh: DeviceMesh | None = None,
+    mask_shift: int = 0,
+) -> tuple[bool | None, str]:
+    """
+    Validate a single placement combination against ground truth.
+
+    Returns (True, "") if valid, (False, error_msg) if invalid, or
+    (None, reason) if the combination cannot be tested (e.g. uneven shards).
+
+    The validation logic:
+    1. Shard inputs according to input placements to get local tensors
+    2. Run the raw op on local tensors (bypassing DTensor dispatch)
+    3. Wrap the local output in a DTensor with the claimed output placement
+    4. Redistribute to Replicate and compare with ground truth
+
+    Args:
+        op: The operator function
+        sample_input: The SampleInput with original arguments
+        tensors: List of (name, tensor) pairs extracted from sample
+        combination: The placement combination to validate
+        ground_truth: Expected output tensor(s). For multi-output ops, a list
+            of tensors where each element is validated independently against
+            the same output placement.
+        world_size: Number of simulated ranks
+        mesh: Optional pre-created device mesh (for performance)
+
+    Returns:
+        (is_valid, error_message)
+    """
+    try:
+        if mesh is None:
+            device = tensors[0][1].device.type if tensors else "cpu"
+            mesh = init_device_mesh(device, (world_size,))
+
+        # Uneven shards produce SymInt in LocalTensor's wrapper shape,
+        # which breaks C++ overload resolution before __torch_dispatch__
+        # can intercept. Return None to signal "untestable".
+        for (name, tensor), placement in zip(tensors, combination[0]):
+            if isinstance(placement, Shard):
+                if tensor.size(placement.dim) % world_size != 0:
+                    return None, "uneven shard"
+
+        local_tensors = _shard_tensors(
+            tensors, combination[0], world_size, mesh, mask_shift
+        )
+
+        local_idx = 0
+
+        def _replace_with_local(a):
+            nonlocal local_idx
+            if isinstance(a, torch.Tensor):
+                local = local_tensors[local_idx]
+                local_idx += 1
+                return local
+            return a
+
+        if isinstance(sample_input.input, torch.Tensor):
+            local_input = _replace_with_local(sample_input.input)
+        else:
+            local_input = pytree.tree_map(_replace_with_local, sample_input.input)
+
+        local_args = pytree.tree_map(_replace_with_local, sample_input.args)
+        local_kwargs = pytree.tree_map(_replace_with_local, sample_input.kwargs)
+
+        local_output = op(local_input, *local_args, **local_kwargs)
+
+        return _compare_outputs(
+            local_output, ground_truth, combination[1], mesh, world_size
+        )
+
+    except Exception as e:
+        # TODO: This is too broad. Consider: (1) explicit checks for shard dim
+        # validity and shape compatibility before calling tensor_split/from_local,
+        # (2) scoped try/except around op() and redistribute() that raise specific
+        # exceptions (e.g., UnsupportedRedistribute, OpError), and (3) only
+        # catching those here, letting real bugs propagate.
+        return False, f"Exception: {type(e).__name__}: {e}"
+
+
+def extract_tensors_from_args(
+    args: tuple[Any, ...],
+    kwargs: dict[str, Any],
+) -> list[tuple[str, torch.Tensor]]:
+    """Extract tensor arguments from captured aten args/kwargs.
+
+    Unlike extract_tensors_from_sample which walks SampleInput pytrees,
+    this walks the flat aten-level args and kwargs directly.
+    """
+    tensors: list[tuple[str, torch.Tensor]] = []
+    idx = 0
+
+    def _collect(x: Any) -> Any:
+        nonlocal idx
+        if isinstance(x, torch.Tensor):
+            tensors.append((f"tensor_{idx}", x))
+            idx += 1
+        return x
+
+    pytree.tree_map(_collect, args)
+    pytree.tree_map(_collect, kwargs)
+    return tensors
+
+
+def validate_aten_combination(
+    aten_op: OpOverload,
+    captured_args: tuple[Any, ...],
+    captured_kwargs: dict[str, Any],
+    ground_truth: torch.Tensor | list[torch.Tensor],
+    combination: PlacementCombination,
+    world_size: int,
+    mesh: DeviceMesh,
+    mask_shift: int = 0,
+) -> tuple[bool | None, str]:
+    """Validate a placement combination using aten-level captured args.
+
+    Works directly with aten op args/kwargs instead of SampleInput pytrees.
+    Replaces tensors in the flat args/kwargs with sharded LocalTensors,
+    calls the aten op, and compares output.
+
+    Returns (True, ""), (False, error_msg), or (None, reason) if untestable.
+    """
+    try:
+        tensors = extract_tensors_from_args(captured_args, captured_kwargs)
+        if not tensors:
+            return False, "No tensor args in captured aten call"
+
+        for (name, tensor), placement in zip(tensors, combination[0]):
+            if isinstance(placement, Shard):
+                if tensor.size(placement.dim) % world_size != 0:
+                    return None, "uneven shard"
+
+        local_tensors = _shard_tensors(
+            tensors, combination[0], world_size, mesh, mask_shift
+        )
+
+        local_idx = 0
+
+        def _replace_with_local(a: Any) -> Any:
+            nonlocal local_idx
+            if isinstance(a, torch.Tensor):
+                local = local_tensors[local_idx]
+                local_idx += 1
+                return local
+            return a
+
+        local_args = pytree.tree_map(_replace_with_local, captured_args)
+        local_kwargs = pytree.tree_map(_replace_with_local, captured_kwargs)
+
+        local_output = aten_op(*local_args, **local_kwargs)
+
+        return _compare_outputs(
+            local_output, ground_truth, combination[1], mesh, world_size
+        )
+
+    except Exception as e:
+        return False, f"Exception: {type(e).__name__}: {e}"
+
+
+def has_pmin_pmax(
+    input_placements: tuple[Placement, ...],
+    output_placements: tuple[Placement, ...],
+) -> bool:
+    """Check if any placement is Partial(min) or Partial(max)."""
+    for p in (*input_placements, *output_placements):
+        if isinstance(p, Partial) and p.reduce_op in ("min", "max"):
+            return True
+    return False
+
+
+def has_any_partial(
+    input_placements: tuple[Placement, ...],
+    output_placements: tuple[Placement, ...],
+) -> bool:
+    """Check if any placement is Partial (any reduce op)."""
+    for p in (*input_placements, *output_placements):
+        if isinstance(p, Partial):
+            return True
+    return False
+
+
+def negate_all_tensors(
+    tensors: list[tuple[str, torch.Tensor]],
+) -> list[tuple[str, torch.Tensor]]:
+    """Return a new list with all tensors negated."""
+    return [(name, -t) for name, t in tensors]
+
+
+def create_fully_negated_sample(sample: SampleInput) -> SampleInput:
+    """Create a sample with ALL tensors negated (for P(min)/P(max) sign testing)."""
+
+    def negate_tensor(x):
+        if isinstance(x, torch.Tensor):
+            return -x
+        return x
+
+    new_input = pytree.tree_map(negate_tensor, sample.input)
+    new_args = pytree.tree_map(negate_tensor, sample.args)
+    new_kwargs = pytree.tree_map(negate_tensor, sample.kwargs)
+
+    return SampleInput(new_input, args=new_args, kwargs=new_kwargs)
+
+
+def _run_op_on_sample(op: Callable[..., Any], sample: SampleInput) -> Any:
+    """Run an operator on a SampleInput, handling both tensor and tuple inputs."""
+    if isinstance(sample.input, torch.Tensor):
+        return op(sample.input, *sample.args, **sample.kwargs)
+    return op(*sample.input, *sample.args, **sample.kwargs)
+
+
+def _extract_rules_from_op_strategy(
+    op_strategy: Any,
+    input_shapes: tuple[tuple[int, ...], ...],
+    output_shapes: tuple[tuple[int, ...], ...],
+) -> set[ComboKey]:
+    """Extract normalized sharding rules from an OpStrategy.
+
+    Called during rule comparison to collect DTensor's claimed-valid placement
+    combinations. These are compared against ground truth (brute-force
+    validation) to find false positives (DTensor claims valid but wrong) and
+    false negatives (valid but DTensor has no rule).
+    """
+    rules: set[ComboKey] = set()
+    if not isinstance(op_strategy, OpStrategy):
+        return rules
+    for spec in op_strategy.strategies:
+        if spec.input_specs is None:
+            continue
+        if isinstance(spec.output_specs, tuple):
+            output_plcs: list[Placement] = []
+            has_none = False
+            for out_spec in spec.output_specs:
+                if out_spec is None:
+                    # None means the output placement is undefined for this
+                    # strategy (e.g. indices under P(max) reduction). Skip it.
+                    has_none = True
+                    break
+                output_plcs.append(out_spec.placements[0])
+            if has_none:
+                continue
+        else:
+            # Single DTensorSpec — the propagator duplicates it for all
+            # outputs of multi-output ops, so we do the same here.
+            output_plcs = [spec.output_spec.placements[0]] * len(output_shapes)
+        input_plcs = tuple(s.placements[0] for s in spec.input_specs)
+        rule_key: ComboKey = (
+            tuple(str(p) for p in input_plcs),
+            tuple(str(p) for p in output_plcs),
+        )
+        normalized_rule = normalize_combo_key(rule_key, input_shapes, output_shapes)
+        if not is_fully_replicated(
+            tuple(parse_placement(p) or Replicate() for p in normalized_rule[0])
+        ):
+            rules.add(normalized_rule)
+    return rules
+
+
+class _CaptureAtenOp(torch.utils._python_dispatch.TorchDispatchMode):
+    """Dispatch mode that captures aten ops called, their args, and return values."""
+
+    def __init__(self, target_op_name: str = ""):
+        self.target_op_name = target_op_name.lower()
+        self.all_ops: list[tuple[OpOverload, tuple[Any, ...], dict[str, Any], Any]] = []
+        self.best_match: OpOverload | None = None
+        self.best_match_args: tuple[Any, ...] | None = None
+        self.best_match_kwargs: dict[str, Any] | None = None
+        self.best_match_result: Any = None
+
+    def __torch_dispatch__(self, func, types, args=(), kwargs=None):
+        if kwargs is None:
+            kwargs = {}
+        result = func(*args, **kwargs)
+        if func.namespace == "aten":
+            self.all_ops.append((func, args, kwargs, result))
+            op_name = func.name().split("::")[1].split(".")[0].lower()
+            if self.target_op_name and self.target_op_name in op_name:
+                if self.best_match is None:
+                    self.best_match = func
+                    self.best_match_args = args
+                    self.best_match_kwargs = kwargs
+                    self.best_match_result = result
+        return result
+
+
+def get_aten_op_for_sample(
+    op: Callable[..., Any], sample: SampleInput, op_name: str = ""
+) -> _CaptureAtenOp:
+    """
+    Capture aten ops dispatched for a given sample.
+
+    Returns the _CaptureAtenOp object containing all captured ops with their
+    args, kwargs, and return values. Use best_match for the primary op or
+    all_ops for exhaustive iteration.
+    """
+    with _CaptureAtenOp(op_name) as capture:
+        try:
+            if isinstance(sample.input, torch.Tensor):
+                op(sample.input, *sample.args, **sample.kwargs)
+            else:
+                op(*sample.input, *sample.args, **sample.kwargs)
+        except Exception:
+            pass
+
+    # Populate best_match from first op if target match wasn't found
+    if capture.best_match is None and capture.all_ops:
+        first_op, first_args, first_kwargs, first_result = capture.all_ops[0]
+        capture.best_match = first_op
+        capture.best_match_args = first_args
+        capture.best_match_kwargs = first_kwargs
+        capture.best_match_result = first_result
+
+    return capture
+
+
+def query_single_dim_strategy(
+    op_overload: OpOverload,
+    captured_args: tuple[Any, ...],
+    captured_kwargs: dict[str, Any],
+) -> list[list[Placement]] | None:
+    """
+    Query DTensor's single-dim strategy for given input tensors.
+    Returns list of [output_placement, *input_placements] rules.
+    """
+    propagator = DTensor._op_dispatcher.sharding_propagator
+
+    if op_overload not in propagator.op_single_dim_strategy_funcs:
+        return None
+
+    strategy_func = propagator.op_single_dim_strategy_funcs[op_overload]
+
+    args_meta = tuple(
+        TensorMeta(shape=a.shape, stride=a.stride(), dtype=a.dtype)
+        if isinstance(a, torch.Tensor)
+        else a
+        for a in captured_args
+    )
+    kwargs_meta = {
+        k: TensorMeta(shape=v.shape, stride=v.stride(), dtype=v.dtype)
+        if isinstance(v, torch.Tensor)
+        else v
+        for k, v in captured_kwargs.items()
+    }
+
+    try:
+        result = strategy_func(op_overload, args_meta, kwargs_meta)
+
+        expanded_result: list[list[Placement]] = []
+        for combo in result:
+            expanded_combo: list[Placement] = []
+            for p in combo:
+                if isinstance(p, _ShardingPlaceholder):
+                    expanded_combo.append(Shard(p.dim))
+                else:
+                    expanded_combo.append(p)
+            expanded_result.append(expanded_combo)
+
+        return expanded_result
+    except Exception:
+        return None
+
+
+def get_opinfo_by_name(name: str) -> list[opinfo_core.OpInfo]:
+    """Find OpInfo entries by exact operator name."""
+    matches = [op for op in op_db if op.name == name]
+    if matches:
+        return matches
+
+    # Suggest alternatives
+    candidates = _find_opinfo_candidates(name)
+    if candidates:
+        suggestions = ", ".join(f'"{c}"' for c in candidates)
+        raise ValueError(f'No OpInfo found for "{name}", did you mean: {suggestions}?')
+    raise ValueError(
+        f'No OpInfo found for "{name}". OpInfo is required as it provides '
+        f"sample inputs for the operator."
+    )
+
+
+def _find_opinfo_candidates(name: str) -> list[str]:
+    """Find OpInfo names that plausibly match a short/incorrect name."""
+    candidates: list[str] = []
+    seen: set[str] = set()
+    # Match on aten_name (e.g., "relu" -> OpInfo with aten_name="relu")
+    for op in op_db:
+        if op.aten_name == name and op.name not in seen:
+            candidates.append(op.name)
+            seen.add(op.name)
+    # Suffix match: "relu" matches "nn.functional.relu"
+    suffix = "." + name
+    for op in op_db:
+        if op.name.endswith(suffix) and op.name not in seen:
+            candidates.append(op.name)
+            seen.add(op.name)
+    return candidates
+
+
+def resolve_op_names(patterns: list[str]) -> list[str]:
+    """Resolve user-provided op patterns to exact OpInfo names.
+
+    Supports exact names, comma separation, and glob patterns (e.g.,
+    "nn.functional.*"). Short names like "relu" are resolved unambiguously
+    or an error is raised with suggestions.
+    """
+    all_opinfo_names = sorted({op.name for op in op_db})
+    resolved: list[str] = []
+    seen: set[str] = set()
+
+    for pattern in patterns:
+        # Glob pattern
+        if "*" in pattern or "?" in pattern:
+            matches = fnmatch.filter(all_opinfo_names, pattern)
+            if not matches:
+                raise ValueError(f'No OpInfo names match pattern "{pattern}".')
+            for m in matches:
+                if m not in seen:
+                    resolved.append(m)
+                    seen.add(m)
+            continue
+
+        # Exact match
+        if pattern in {op.name for op in op_db}:
+            if pattern not in seen:
+                resolved.append(pattern)
+                seen.add(pattern)
+            continue
+
+        # Try to resolve shorthand
+        candidates = _find_opinfo_candidates(pattern)
+        if len(candidates) == 1:
+            name = candidates[0]
+            if name not in seen:
+                resolved.append(name)
+                seen.add(name)
+        elif len(candidates) > 1:
+            suggestions = ", ".join(f'"{c}"' for c in candidates)
+            raise ValueError(
+                f'"{pattern}" is ambiguous, matching: {suggestions}. '
+                f"Use the fully qualified name."
+            )
+        else:
+            raise ValueError(f'No OpInfo found for "{pattern}".')
+
+    return resolved
+
+
+def _is_tensor_output(result: Any) -> bool:
+    """Check if a result is a tensor or list/tuple of tensors."""
+    if isinstance(result, torch.Tensor):
+        return True
+    if isinstance(result, (list, tuple)):
+        has_tensor = any(isinstance(t, torch.Tensor) for t in result)
+        all_tensor = all(isinstance(t, torch.Tensor) for t in result)
+        if has_tensor and not all_tensor:
+            raise NotImplementedError(
+                f"Mixed tensor/non-tensor tuple outputs are not supported by the "
+                f"validator. Got types: {[type(t).__name__ for t in result]}"
+            )
+        return all_tensor
+    return False
+
+
+def _to_ground_truth(result: Any) -> torch.Tensor | list[torch.Tensor]:
+    """Convert an op result to the ground truth format (tensor or list of tensors)."""
+    if isinstance(result, torch.Tensor):
+        return result
+    return list(result)
+
+
+def _prepare_false_positive_mitigations(
+    op: Callable[..., Any],
+    sample: SampleInput,
+    tensors: list[tuple[str, torch.Tensor]],
+) -> _FalsePositiveMitigations:
+    """Create negated and non-rounded sample variants for false positive detection."""
+    m = _FalsePositiveMitigations()
+
+    try:
+        m.negated_sample = create_fully_negated_sample(sample)
+        m.negated_tensors = negate_all_tensors(tensors)
+        result = _run_op_on_sample(op, m.negated_sample)
+        if _is_tensor_output(result):
+            m.negated_ground_truth = _to_ground_truth(result)
+        else:
+            m.negated_sample = None
+    except Exception:
+        m.negated_sample = None
+        m.negated_tensors = None
+
+    if "rounding_mode" not in sample.kwargs:
+        return m
+
+    try:
+        non_rounded_kwargs = {
+            k: v for k, v in sample.kwargs.items() if k != "rounding_mode"
+        }
+        m.non_rounded_sample = SampleInput(
+            sample.input, args=sample.args, kwargs=non_rounded_kwargs
+        )
+        result = _run_op_on_sample(op, m.non_rounded_sample)
+        if not _is_tensor_output(result):
+            m.non_rounded_sample = None
+        else:
+            m.non_rounded_ground_truth = _to_ground_truth(result)
+            m.non_rounded_negated_sample = create_fully_negated_sample(
+                m.non_rounded_sample
+            )
+            m.non_rounded_negated_tensors = negate_all_tensors(tensors)
+            nr_neg_result = _run_op_on_sample(op, m.non_rounded_negated_sample)
+            if _is_tensor_output(nr_neg_result):
+                m.non_rounded_negated_ground_truth = _to_ground_truth(nr_neg_result)
+            else:
+                m.non_rounded_negated_sample = None
+    except Exception:
+        m.non_rounded_sample = None
+        m.non_rounded_ground_truth = None
+        m.non_rounded_negated_sample = None
+
+    return m
+
+
+def _query_dtensor_rules(
+    aten_op: OpOverload | None,
+    tensors: list[tuple[str, torch.Tensor]],
+    captured_args: tuple[Any, ...],
+    captured_kwargs: dict[str, Any],
+    input_shapes: tuple[tuple[int, ...], ...],
+    output_shapes: tuple[tuple[int, ...], ...],
+    world_size: int,
+    verbose: bool,
+) -> set[ComboKey]:
+    """Query DTensor's claimed sharding rules via single-dim, op_strategy, or decomp paths.
+
+    TODO: This reimplements strategy resolution logic from ShardingPropagator.
+    Refactor ShardingPropagator to expose a public API for querying sharding
+    rules given an op and tensor metadata, so this function can be replaced
+    with a single call.
+    """
+    if not aten_op:
+        return set()
+
+    num_tensors = len(tensors)
+    non_tensor_kwargs = {
+        k: v for k, v in captured_kwargs.items() if not isinstance(v, torch.Tensor)
+    }
+    n_outputs = len(output_shapes)
+    propagator = DTensor._op_dispatcher.sharding_propagator
+    rules: set[ComboKey] = set()
+
+    if aten_op in propagator.op_single_dim_strategy_funcs:
+        strategy_result = query_single_dim_strategy(
+            aten_op, captured_args, captured_kwargs
+        )
+        if strategy_result:
+            for combo in strategy_result:
+                if len(combo) >= n_outputs + num_tensors:
+                    output_plcs = combo[:n_outputs]
+                    input_plcs = tuple(combo[n_outputs : n_outputs + num_tensors])
+                    rule_key: ComboKey = (
+                        tuple(str(p) for p in input_plcs),
+                        tuple(str(p) for p in output_plcs),
+                    )
+                    normalized_rule = normalize_combo_key(
+                        rule_key, input_shapes, output_shapes
+                    )
+                    if not is_fully_replicated(
+                        tuple(
+                            parse_placement(p) or Replicate()
+                            for p in normalized_rule[0]
+                        )
+                    ):
+                        rules.add(normalized_rule)
+
+    elif aten_op in propagator.op_strategy_funcs:
+        try:
+            mesh = init_device_mesh("cpu", (world_size,))
+            # Build OpStrategy objects for each tensor, keyed by identity
+            tensor_to_strategy: dict[int, OpStrategy] = {}
+            for _, t in tensors:
+                input_placements = get_1d_input_placements_for_tensor(
+                    t, include_partial=True
+                )
+                specs = []
+                for p in input_placements:
+                    spec = DTensorSpec(
+                        mesh=mesh,
+                        placements=(p,),
+                        tensor_meta=TensorMeta(
+                            shape=t.shape, stride=t.stride(), dtype=t.dtype
+                        ),
+                    )
+                    specs.append(OpSpec(output_specs=spec, input_specs=tuple()))
+                tensor_to_strategy[id(t)] = OpStrategy(specs)
+            # Interleave strategies and non-tensor args at original positions
+            args_schema = [
+                tensor_to_strategy[id(a)] if isinstance(a, torch.Tensor) else a
+                for a in captured_args
+            ]
+            op_schema = OpSchema(aten_op, tuple(args_schema), non_tensor_kwargs)
+            strategy_func = propagator.op_strategy_funcs[aten_op]
+            output_strategy = strategy_func(op_schema)
+            rules |= _extract_rules_from_op_strategy(
+                output_strategy, input_shapes, output_shapes
+            )
+        except Exception as e:
+            if verbose:
+                print(f"        Error querying op_strategy: {e}")
+
+    else:
+        # Decomp-based strategy: only discovers rules reachable from a single
+        # seed (Shard(0) on the first input). Rules requiring other input
+        # placements (e.g., Shard(1), Partial, or sharding on non-first inputs)
+        # will not be found, so this under-reports DTensor's capabilities.
+        if DecompShardingStrategy.has_decomp(aten_op):
+            try:
+                mesh = init_device_mesh("cpu", (world_size,))
+                # Interleave DTensorSpec and non-tensor args at original positions
+                tensor_idx = 0
+                args_schema: list[Any] = []
+                for a in captured_args:
+                    if isinstance(a, torch.Tensor):
+                        # First tensor gets Shard(0) to seed candidate
+                        # placement generation in _get_candidate_placements
+                        plc = Shard(0) if tensor_idx == 0 else Replicate()
+                        spec = DTensorSpec(
+                            mesh=mesh,
+                            placements=(plc,),
+                            tensor_meta=TensorMeta(
+                                shape=a.shape, stride=a.stride(), dtype=a.dtype
+                            ),
+                        )
+                        args_schema.append(spec)
+                        tensor_idx += 1
+                    else:
+                        args_schema.append(a)
+                op_schema = OpSchema(aten_op, tuple(args_schema), non_tensor_kwargs)
+                propagator.decomp_strategy.ensure_schema_info(aten_op)
+                output_strategy = propagator.decomp_strategy.propagate_strategy(
+                    op_schema,
+                )
+                if output_strategy is not None:
+                    rules |= _extract_rules_from_op_strategy(
+                        output_strategy, input_shapes, output_shapes
+                    )
+            except Exception as e:
+                if verbose:
+                    print(f"        Error querying decomp strategy: {e}")
+
+    return rules
+
+
+def _validate_with_mitigations(
+    op: Callable[..., Any],
+    sample: SampleInput,
+    tensors: list[tuple[str, torch.Tensor]],
+    input_placements: tuple[Placement, ...],
+    output_placements: tuple[Placement, ...],
+    ground_truth: torch.Tensor | list[torch.Tensor],
+    world_size: int,
+    mesh: DeviceMesh,
+    mitigations: _FalsePositiveMitigations,
+) -> bool | None:
+    """Validate a combination, including false positive mitigation re-checks.
+
+    Returns True (valid), False (invalid), or None (untestable).
+    """
+    combo: PlacementCombination = (input_placements, output_placements)
+    is_valid, _ = validate_combination(
+        op, sample, tensors, combo, ground_truth, world_size, mesh
+    )
+    if is_valid is None:
+        return None
+
+    # Flipped-mask mitigation: the checkerboard mask that controls offset
+    # signs (for P(sum)/P(avg)) or rank ownership (for P(min)/P(max)) is
+    # deterministic per tensor_idx. Re-validate with the complementary mask
+    # to catch index-returning ops (argmin/argmax) where the result
+    # coincidentally matches because the dominant value happens to land on
+    # a position where both mask orientations preserve argmin/argmax.
+    if is_valid and has_any_partial(input_placements, output_placements):
+        is_valid, _ = validate_combination(
+            op,
+            sample,
+            tensors,
+            combo,
+            ground_truth,
+            world_size,
+            mesh,
+            mask_shift=1,
+        )
+
+    if (
+        is_valid
+        and mitigations.negated_sample
+        and has_pmin_pmax(input_placements, output_placements)
+    ):
+        if mitigations.negated_tensors is None:
+            raise AssertionError("negated_tensors is None")
+        if mitigations.negated_ground_truth is None:
+            raise AssertionError("negated_ground_truth is None")
+        is_valid, _ = validate_combination(
+            op,
+            mitigations.negated_sample,
+            mitigations.negated_tensors,
+            combo,
+            mitigations.negated_ground_truth,
+            world_size,
+            mesh,
+        )
+
+    if (
+        is_valid
+        and mitigations.non_rounded_sample
+        and has_any_partial(input_placements, output_placements)
+    ):
+        if mitigations.non_rounded_ground_truth is None:
+            raise AssertionError("non_rounded_ground_truth is None")
+        is_valid, _ = validate_combination(
+            op,
+            mitigations.non_rounded_sample,
+            tensors,
+            combo,
+            mitigations.non_rounded_ground_truth,
+            world_size,
+            mesh,
+        )
+
+    if (
+        is_valid
+        and mitigations.non_rounded_negated_sample
+        and has_pmin_pmax(input_placements, output_placements)
+    ):
+        if mitigations.non_rounded_negated_tensors is None:
+            raise AssertionError("non_rounded_negated_tensors is None")
+        if mitigations.non_rounded_negated_ground_truth is None:
+            raise AssertionError("non_rounded_negated_ground_truth is None")
+        is_valid, _ = validate_combination(
+            op,
+            mitigations.non_rounded_negated_sample,
+            mitigations.non_rounded_negated_tensors,
+            combo,
+            mitigations.non_rounded_negated_ground_truth,
+            world_size,
+            mesh,
+        )
+
+    return is_valid
+
+
+@dataclass
+class _AtenFalsePositiveMitigations:
+    """Bundle of negated variants for aten-level false positive detection.
+
+    Unlike _FalsePositiveMitigations, this works with captured aten args/kwargs
+    directly, without SampleInput. The rounding_mode mitigation is skipped
+    since it's an OpInfo-level concept that doesn't appear in captured aten kwargs.
+    """
+
+    negated_args: tuple[Any, ...] | None = None
+    negated_kwargs: dict[str, Any] | None = None
+    negated_ground_truth: torch.Tensor | list[torch.Tensor] | None = None
+
+
+def _negate_tensors_in_tree(tree: Any) -> Any:
+    """Negate all tensors in a pytree structure."""
+
+    def _negate(x: Any) -> Any:
+        if isinstance(x, torch.Tensor):
+            return -x
+        return x
+
+    return pytree.tree_map(_negate, tree)
+
+
+def _prepare_aten_mitigations(
+    aten_op: OpOverload,
+    captured_args: tuple[Any, ...],
+    captured_kwargs: dict[str, Any],
+) -> _AtenFalsePositiveMitigations:
+    """Create negated variants for aten-level false positive detection."""
+    m = _AtenFalsePositiveMitigations()
+    try:
+        m.negated_args = _negate_tensors_in_tree(captured_args)
+        m.negated_kwargs = _negate_tensors_in_tree(captured_kwargs)
+        result = aten_op(*m.negated_args, **m.negated_kwargs)
+        if _is_tensor_output(result):
+            m.negated_ground_truth = _to_ground_truth(result)
+        else:
+            m.negated_args = None
+            m.negated_kwargs = None
+    except Exception:
+        m.negated_args = None
+        m.negated_kwargs = None
+    return m
+
+
+def _validate_aten_with_mitigations(
+    aten_op: OpOverload,
+    captured_args: tuple[Any, ...],
+    captured_kwargs: dict[str, Any],
+    input_placements: tuple[Placement, ...],
+    output_placements: tuple[Placement, ...],
+    ground_truth: torch.Tensor | list[torch.Tensor],
+    world_size: int,
+    mesh: DeviceMesh,
+    mitigations: _AtenFalsePositiveMitigations,
+) -> bool | None:
+    """Validate an aten-level combination with false positive mitigations.
+
+    Returns True (valid), False (invalid), or None (untestable).
+    """
+    combo: PlacementCombination = (input_placements, output_placements)
+    is_valid, _ = validate_aten_combination(
+        aten_op,
+        captured_args,
+        captured_kwargs,
+        ground_truth,
+        combo,
+        world_size,
+        mesh,
+    )
+    if is_valid is None:
+        return None
+
+    if is_valid and has_any_partial(input_placements, output_placements):
+        is_valid, _ = validate_aten_combination(
+            aten_op,
+            captured_args,
+            captured_kwargs,
+            ground_truth,
+            combo,
+            world_size,
+            mesh,
+            mask_shift=1,
+        )
+
+    if (
+        is_valid
+        and mitigations.negated_args is not None
+        and has_pmin_pmax(input_placements, output_placements)
+    ):
+        if mitigations.negated_kwargs is None:
+            raise AssertionError("negated_kwargs must not be None")
+        if mitigations.negated_ground_truth is None:
+            raise AssertionError("negated_ground_truth must not be None")
+        is_valid, _ = validate_aten_combination(
+            aten_op,
+            mitigations.negated_args,
+            mitigations.negated_kwargs,
+            mitigations.negated_ground_truth,
+            combo,
+            world_size,
+            mesh,
+        )
+
+    return is_valid
+
+
+def _assert_keys_normalized(
+    keys: set[ComboKey],
+    input_shapes: tuple[tuple[int, ...], ...],
+    output_shapes: tuple[tuple[int, ...], ...],
+) -> None:
+    """Assert all combo keys have trivial shards already normalized to Replicate."""
+    for key in keys:
+        if key != normalize_combo_key(key, input_shapes, output_shapes):
+            raise AssertionError(
+                f"Key {key} contains un-normalized trivial shards; "
+                f"call normalize_combo_key before _compare_rules"
+            )
+
+
+def _compare_rules(
+    ground_truth_valid: set[ComboKey],
+    dtensor_rules: set[ComboKey],
+    input_shapes: tuple[tuple[int, ...], ...],
+    output_shapes: tuple[tuple[int, ...], ...],
+    sample_idx: int,
+    scalar_args: tuple[Any, ...],
+    scalar_kwargs: dict[str, Any],
+    aten_op: OpOverload | None,
+    variant: str,
+    stats: ComparisonStats,
+    sample: SampleInput | None = None,
+    untestable: set[ComboKey] | None = None,
+) -> None:
+    """Compare ground truth valid rules against DTensor claimed rules, updating stats."""
+    if not dtensor_rules:
+        return
+    if untestable is None:
+        untestable = set()
+
+    _assert_keys_normalized(ground_truth_valid, input_shapes, output_shapes)
+    _assert_keys_normalized(dtensor_rules, input_shapes, output_shapes)
+
+    op_str = str(aten_op)
+    for combo_key in ground_truth_valid:
+        if combo_key in dtensor_rules:
+            stats.true_positives += 1
+            stats.true_positives_by_op[op_str] = (
+                stats.true_positives_by_op.get(op_str, 0) + 1
+            )
+        elif combo_key not in untestable:
+            stats.false_negatives.append(
+                Discrepancy(
+                    input_placements=combo_key[0],
+                    output_placements=combo_key[1],
+                    sample_idx=sample_idx,
+                    input_shapes=input_shapes,
+                    discrepancy_type="false_negative",
+                    scalar_args=scalar_args,
+                    scalar_kwargs=scalar_kwargs,
+                    aten_op=aten_op,
+                    variant=variant,
+                    sample=sample,
+                )
+            )
+
+    for combo_key in dtensor_rules:
+        if combo_key not in ground_truth_valid and combo_key not in untestable:
+            stats.false_positives.append(
+                Discrepancy(
+                    input_placements=combo_key[0],
+                    output_placements=combo_key[1],
+                    sample_idx=sample_idx,
+                    input_shapes=input_shapes,
+                    discrepancy_type="false_positive",
+                    scalar_args=scalar_args,
+                    scalar_kwargs=scalar_kwargs,
+                    aten_op=aten_op,
+                    variant=variant,
+                    sample=sample,
+                )
+            )
+
+
+def _format_op_name(opinfo_name: str) -> str:
+    """Format an OpInfo name for display, adding aten. prefix for simple names."""
+    if "." not in opinfo_name:
+        return f"aten.{opinfo_name}"
+    return opinfo_name
+
+
+def _format_sample_repro(sample: SampleInput, aten_op: OpOverload | None = None) -> str:
+    """Format a SampleInput's values for repro output, using schema arg names."""
+    # Get argument names from the aten op schema if available.
+    # Schema args are ordered: the first is the input, then positional args,
+    # then keyword-only args (which appear in sample.kwargs).
+    arg_names: list[str] = []
+    if aten_op is not None:
+        try:
+            arg_names = [a.name for a in aten_op._schema.arguments]
+        except Exception:
+            pass
+
+    parts = []
+    # input is the first positional arg
+    name = arg_names[0] if arg_names else "input"
+    parts.append(f"{name}={sample.input!r}")
+    # remaining positional args
+    for i, arg in enumerate(sample.args):
+        name = arg_names[1 + i] if 1 + i < len(arg_names) else f"args[{i}]"
+        parts.append(f"{name}={arg!r}")
+    # kwargs — use their actual key names (already named)
+    for k, v in sample.kwargs.items():
+        parts.append(f"{k}={v!r}")
+    return ", ".join(parts)
+
+
+def _print_discrepancy_section(
+    title: str, discrepancies: list[Discrepancy], show_repro: int = 0
+) -> None:
+    """Print grouped discrepancies for a section (incorrect or missing)."""
+    if not discrepancies:
+        return
+    print(f"\n{title}")
+    by_op: dict[str, dict[ComboKey, list[Discrepancy]]] = defaultdict(
+        lambda: defaultdict(list)
+    )
+    for d in discrepancies:
+        op_str = str(d.aten_op)
+        key = (d.input_placements, d.output_placements)
+        by_op[op_str][key].append(d)
+
+    for op_str in sorted(by_op.keys()):
+        print(f"\n  [{op_str}]")
+        for (inp, out), discs in sorted(by_op[op_str].items(), key=str):
+            inp_str = ", ".join(inp)
+            out_str = out[0] if len(out) == 1 else "(" + ", ".join(out) + ")"
+            print(f"    {inp_str} -> {out_str}")
+            if show_repro:
+                limit = len(discs) if show_repro < 0 else show_repro
+                for d in discs[:limit]:
+                    if d.sample is not None:
+                        print(
+                            f"      Repro: {_format_sample_repro(d.sample, d.aten_op)}"
+                        )
+                if len(discs) > limit:
+                    print(f"      ... and {len(discs) - limit} more")
+
+
+def _print_comparison_summary(
+    stats: ComparisonStats,
+    show_repro: int = 0,
+) -> None:
+    """Print discrepancy details for an operator."""
+    # Per aten op variant breakdown
+    fp_by_op: dict[str, set[ComboKey]] = defaultdict(set)
+    for d in stats.false_positives:
+        op_str = str(d.aten_op)
+        fp_by_op[op_str].add((d.input_placements, d.output_placements))
+    fn_by_op: dict[str, set[ComboKey]] = defaultdict(set)
+    for d in stats.false_negatives:
+        op_str = str(d.aten_op)
+        fn_by_op[op_str].add((d.input_placements, d.output_placements))
+
+    all_ops = sorted(set(stats.true_positives_by_op) | set(fp_by_op) | set(fn_by_op))
+    if len(all_ops) > 1:
+        for op_str in all_ops:
+            tp = stats.true_positives_by_op.get(op_str, 0)
+            fp = len(fp_by_op.get(op_str, set()))
+            fn = len(fn_by_op.get(op_str, set()))
+            print(f"  {op_str}: {tp} correct, {fp} incorrect, {fn} missing")
+
+    _print_discrepancy_section(
+        "Incorrect (has rule but ground truth invalid)",
+        stats.false_positives,
+        show_repro,
+    )
+    _print_discrepancy_section(
+        "Possibly missing (valid in ground truth but no DTensor rule)",
+        stats.false_negatives,
+        show_repro,
+    )
+
+
+def _has_dtensor_support(aten_op: OpOverload) -> bool:
+    """Check if an aten op has any DTensor sharding strategy registered."""
+    propagator = DTensor._op_dispatcher.sharding_propagator
+    if aten_op in propagator.op_single_dim_strategy_funcs:
+        return True
+    if aten_op in propagator.op_strategy_funcs:
+        return True
+    return DecompShardingStrategy.has_decomp(aten_op)
+
+
+def _discover_aten_op(
+    opinfos: list[opinfo_core.OpInfo],
+    device: str,
+    dtype: torch.dtype,
+) -> OpOverload | None:
+    """Discover the aten op dispatched by the first valid sample."""
+    for opinfo in opinfos:
+        try:
+            samples = list(opinfo.sample_inputs(device, dtype))
+        except Exception:
+            continue
+        for sample in samples:
+            tensors = extract_tensors_from_sample(sample)
+            if not tensors or any(0 in t.shape for _, t in tensors):
+                continue
+            capture = get_aten_op_for_sample(opinfo.op, sample, opinfo.name)
+            aten_op = capture.best_match
+            if aten_op is not None:
+                return aten_op
+    return None
+
+
+def _check_ground_truth(
+    result: Any,
+) -> torch.Tensor | list[torch.Tensor] | None:
+    """Validate an op result is suitable as ground truth.
+
+    Returns the ground truth tensor(s) or None if the result should be skipped.
+    """
+    if isinstance(result, (list, tuple)):
+        if not all(isinstance(t, torch.Tensor) for t in result):
+            return None
+        gt = list(result)
+    elif isinstance(result, torch.Tensor):
+        gt = result
+    else:
+        return None
+
+    first_gt = gt[0] if isinstance(gt, list) else gt
+    if first_gt.numel() == 0:
+        return None
+    if (first_gt == 0).all():
+        return None
+    if first_gt.isnan().all():
+        return None
+    return gt
+
+
+def _validate_aten_op_for_sample(
+    aten_op: OpOverload,
+    captured_args: tuple[Any, ...],
+    captured_kwargs: dict[str, Any],
+    ground_truth: torch.Tensor | list[torch.Tensor],
+    world_size: int,
+    incorrect_only: bool,
+    verbose: bool,
+    sample_idx: int,
+    variant: str,
+    stats: ComparisonStats,
+    sample: SampleInput | None = None,
+) -> tuple[int, int]:
+    """Validate a single aten op with captured args against ground truth.
+
+    Shared logic used by both default (1:1) and allow_composite modes in
+    compare_operator. Returns (samples_counted, combinations_counted).
+    """
+    tensors = extract_tensors_from_args(captured_args, captured_kwargs)
+    if not tensors:
+        return 0, 0
+    if any(0 in t.shape for _, t in tensors):
+        return 0, 0
+
+    input_shapes = tuple(t.shape for _, t in tensors)
+    gt_list = ground_truth if isinstance(ground_truth, list) else [ground_truth]
+    output_shapes = tuple(tuple(gt.shape) for gt in gt_list)
+    n_outputs = len(gt_list)
+    first_gt = gt_list[0]
+
+    scalar_args = tuple(a for a in captured_args if not isinstance(a, torch.Tensor))
+    scalar_kwargs = {
+        k: v for k, v in captured_kwargs.items() if not isinstance(v, torch.Tensor)
+    }
+
+    mitigations = _prepare_aten_mitigations(aten_op, captured_args, captured_kwargs)
+
+    input_placement_options = [
+        get_1d_input_placements_for_tensor(t, include_partial=True) for _, t in tensors
+    ]
+    output_placement_options = get_1d_output_placements_for_tensor(first_gt)
+
+    dtensor_rules = _query_dtensor_rules(
+        aten_op,
+        tensors,
+        captured_args,
+        captured_kwargs,
+        input_shapes,
+        output_shapes,
+        world_size,
+        verbose,
+    )
+
+    ground_truth_valid: set[ComboKey] = set()
+    total_combinations = 0
+
+    tensor_device = tensors[0][1].device.type if tensors else "cpu"
+    with LocalTensorMode(frozenset(range(world_size))):
+        mesh = init_device_mesh(tensor_device, (world_size,))
+
+        if incorrect_only:
+            combinations_to_test = []
+            for combo_key in dtensor_rules:
+                input_plc_strs, output_plc_strs = combo_key
+                input_plcs_list: list[Placement] = []
+                all_valid = True
+                for s in input_plc_strs:
+                    p = parse_placement(s)
+                    if p is None:
+                        all_valid = False
+                        break
+                    input_plcs_list.append(p)
+                output_plcs_list: list[Placement] = []
+                for s in output_plc_strs:
+                    p = parse_placement(s)
+                    if p is None:
+                        all_valid = False
+                        break
+                    output_plcs_list.append(p)
+                if not all_valid:
+                    continue
+                combinations_to_test.append(
+                    (
+                        tuple(input_plcs_list),
+                        tuple(output_plcs_list),
+                        combo_key,
+                    )
+                )
+        else:
+            combinations_to_test = []
+            for input_placements in itertools.product(*input_placement_options):
+                if is_fully_replicated(input_placements):
+                    continue
+                for output_placement in output_placement_options:
+                    output_placements = tuple(
+                        output_placement for _ in range(n_outputs)
+                    )
+                    combo_key = (
+                        tuple(str(p) for p in input_placements),
+                        tuple(str(p) for p in output_placements),
+                    )
+                    combinations_to_test.append(
+                        (input_placements, output_placements, combo_key)
+                    )
+
+        untestable: set[ComboKey] = set()
+
+        for (
+            input_placements,
+            output_placements,
+            combo_key,
+        ) in combinations_to_test:
+            total_combinations += 1
+            is_valid = _validate_aten_with_mitigations(
+                aten_op,
+                captured_args,
+                captured_kwargs,
+                input_placements,
+                output_placements,
+                ground_truth,
+                world_size,
+                mesh,
+                mitigations,
+            )
+
+            if is_valid is None:
+                normalized_key = normalize_combo_key(
+                    combo_key, input_shapes, output_shapes
+                )
+                untestable.add(normalized_key)
+            elif is_valid:
+                normalized_key = normalize_combo_key(
+                    combo_key, input_shapes, output_shapes
+                )
+                if not is_fully_replicated(
+                    tuple(parse_placement(p) or Replicate() for p in normalized_key[0])
+                ):
+                    ground_truth_valid.add(normalized_key)
+
+    _compare_rules(
+        ground_truth_valid,
+        dtensor_rules,
+        input_shapes,
+        output_shapes,
+        sample_idx,
+        scalar_args,
+        scalar_kwargs,
+        aten_op,
+        variant,
+        stats,
+        sample,
+        untestable,
+    )
+
+    if verbose:
+        print(f"      Sample {sample_idx} [{aten_op}]: shapes={input_shapes}")
+        print(f"        Ground truth valid: {len(ground_truth_valid)}")
+        print(f"        DTensor rules: {len(dtensor_rules)}")
+
+    return 1, total_combinations
+
+
+def compare_operator(
+    op_name: str,
+    device: str = "cpu",
+    dtype: torch.dtype = torch.float32,
+    world_size: int = 2,
+    max_samples: int | None = None,
+    verbose: bool = False,
+    incorrect_only: bool = False,
+    allow_composite: bool = False,
+) -> ComparisonStats:
+    """
+    Compare DTensor's sharding rules against ground truth for an operator.
+
+    Args:
+        op_name: Name of the operator to test
+        device: Device to run on
+        dtype: Data type for tensors
+        world_size: Simulated world size
+        max_samples: Maximum number of samples to test per OpInfo
+        verbose: Print detailed output
+        incorrect_only: If True, only test DTensor's claimed rules for correctness.
+            Skips search for missing rules (much faster).
+        allow_composite: If True, validate each supported aten op individually for
+            samples that decompose into multiple aten calls. Default (False)
+            skips samples where the OpInfo doesn't map 1:1 to a single aten op.
+    """
+    if op_name in SKIP_OPS:
+        return ComparisonStats()
+
+    opinfos = get_opinfo_by_name(op_name)
+
+    stats = ComparisonStats()
+
+    if not allow_composite:
+        aten_op = _discover_aten_op(opinfos, device, dtype)
+        if aten_op is None or not _has_dtensor_support(aten_op):
+            if verbose:
+                print(f"  ATEN_OP_MAP: {op_name} -> {aten_op} [no_support]")
+            stats.no_dtensor_support = True
+            return stats
+        if verbose:
+            print(f"  ATEN_OP_MAP: {op_name} -> {aten_op} [supported]")
+
+    total_samples = 0
+    total_combinations = 0
+    skip_reasons: dict[str, int] = defaultdict(int)
+
+    for opinfo in opinfos:
+        variant = opinfo.variant_test_name
+        if variant and verbose:
+            print(f"\n  OpInfo variant: {variant}")
+
+        op = opinfo.op
+
+        try:
+            samples = list(opinfo.sample_inputs(device, dtype))
+        except Exception as e:
+            if verbose:
+                print(f"    Error generating samples: {e}")
+            continue
+
+        if max_samples:
+            samples = samples[:max_samples]
+
+        for sample_idx, sample in enumerate(samples):
+            # Check that SampleInput has tensor inputs and no zero-sized tensors
+            sample_tensors = extract_tensors_from_sample(sample)
+            if len(sample_tensors) == 0:
+                skip_reasons["no tensor inputs"] += 1
+                continue
+            if any(0 in t.shape for _, t in sample_tensors):
+                skip_reasons["zero-sized tensor"] += 1
+                continue
+
+            # Capture all aten ops dispatched for this sample
+            capture = get_aten_op_for_sample(op, sample, opinfo.name)
+            if capture.best_match is None:
+                skip_reasons["no aten op captured"] += 1
+                continue
+
+            # Count supported aten ops in the capture
+            supported_ops = [
+                (func, args, kwargs, result)
+                for func, args, kwargs, result in capture.all_ops
+                if _has_dtensor_support(func)
+            ]
+            num_supported = len(supported_ops)
+
+            if allow_composite:
+                # Validate each supported aten op individually
+                if num_supported == 0:
+                    skip_reasons["no supported aten ops"] += 1
+                    continue
+
+                for func, args, kwargs, result in supported_ops:
+                    gt = _check_ground_truth(result)
+                    if gt is None:
+                        skip_reasons["non-tensor/degenerate aten output"] += 1
+                        continue
+                    n_samples, n_combos = _validate_aten_op_for_sample(
+                        func,
+                        args,
+                        kwargs,
+                        gt,
+                        world_size,
+                        incorrect_only,
+                        verbose,
+                        sample_idx,
+                        variant,
+                        stats,
+                        sample,
+                    )
+                    total_samples += n_samples
+                    total_combinations += n_combos
+            else:
+                # Default: only validate samples with a single supported aten op
+                if num_supported != 1:
+                    skip_reasons["non-1:1 aten mapping"] += 1
+                    continue
+
+                func, args, kwargs, result = supported_ops[0]
+                gt = _check_ground_truth(result)
+                if gt is None:
+                    skip_reasons["non-tensor/degenerate aten output"] += 1
+                    continue
+
+                n_samples, n_combos = _validate_aten_op_for_sample(
+                    func,
+                    args,
+                    kwargs,
+                    gt,
+                    world_size,
+                    incorrect_only,
+                    verbose,
+                    sample_idx,
+                    variant,
+                    stats,
+                    sample,
+                )
+                total_samples += n_samples
+                total_combinations += n_combos
+
+    stats.total_samples = total_samples
+    stats.total_combinations = total_combinations
+    stats.skip_reasons = dict(skip_reasons)
+
+    # In allow_composite mode, check DTensor support after processing
+    if allow_composite and total_samples == 0 and not skip_reasons:
+        stats.no_dtensor_support = True
+
+    return stats
+
+
+def get_registered_op_names() -> list[str]:
+    """Get all op names that have DTensor sharding rules and also have OpInfo.
+
+    Returns OpInfo names (which may differ from aten base names, e.g.,
+    "nn.functional.relu" instead of "relu").
+    """
+
+    propagator = DTensor._op_dispatcher.sharding_propagator
+
+    # Get all registered aten ops
+    all_registered = set(propagator.op_single_dim_strategy_funcs.keys()) | set(
+        propagator.op_strategy_funcs.keys()
+    )
+
+    # Extract base names (aten.mul.Tensor -> mul)
+    base_names = set()
+    for op in all_registered:
+        parts = str(op).split(".")
+        if len(parts) >= 2:
+            base_names.add(parts[1])
+
+    # Build mappings from OpInfo: both by name and by aten_name
+    opinfo_by_name = {}
+    opinfo_by_aten_name: dict[str, list[str]] = {}
+    for op in op_db:
+        opinfo_by_name[op.name] = True
+        opinfo_by_aten_name.setdefault(op.aten_name, []).append(op.name)
+
+    result = set()
+    for base_name in base_names:
+        if base_name in opinfo_by_name:
+            # Direct match (e.g., "add" -> OpInfo named "add")
+            result.add(base_name)
+        elif base_name in opinfo_by_aten_name:
+            # Match via aten_name (e.g., "relu" -> OpInfo "nn.functional.relu"
+            # which has aten_name="relu")
+            result.update(opinfo_by_aten_name[base_name])
+
+    return sorted(result)
+
+
+def report_registrations(verbose: bool = False) -> None:
+    """Report the number (and optionally full list) of ops by registration method.
+
+    Directly registered categories (mutually exclusive):
+      - rule: ops registered via register_prop_rule
+      - op_strategy: ops registered via register_op_strategy
+      - single_dim_strategy: ops registered via register_single_dim_strategy
+
+    Ops not in any of the above may still be supported at runtime via
+    DecompShardingStrategy, which traces through the op's decomposition and
+    propagates placements through the decomposed sub-ops.  Whether this
+    actually works depends on every sub-op having a registered strategy.
+    We report the decomposition_table entries as a separate (untested) count.
+    """
+    from torch._decomp import decomposition_table
+
+    propagator = DTensor._op_dispatcher.sharding_propagator
+
+    rule_ops = sorted(propagator.op_to_rules.keys(), key=str)
+    strategy_ops = sorted(propagator.op_strategy_funcs.keys(), key=str)
+    single_dim_ops = sorted(propagator.op_single_dim_strategy_funcs.keys(), key=str)
+
+    directly_registered = (
+        set(propagator.op_to_rules.keys())
+        | set(propagator.op_strategy_funcs.keys())
+        | set(propagator.op_single_dim_strategy_funcs.keys())
+    )
+
+    # Ops from the explicit decomposition table that aren't directly registered.
+    # These *may* work via DecompShardingStrategy if all their sub-ops are
+    # supported, but we can't verify that without tracing each one.
+    decomp_only_ops = sorted(
+        (op for op in decomposition_table if op not in directly_registered),
+        key=str,
+    )
+
+    print("=" * 70)
+    print("DTensor operator registration report")
+    print("=" * 70)
+
+    print("\nDirectly registered:")
+    print(f"  rule (register_prop_rule):            {len(rule_ops):>4}")
+    print(f"  op_strategy (register_op_strategy):   {len(strategy_ops):>4}")
+    print(f"  single_dim_strategy:                  {len(single_dim_ops):>4}")
+    print(f"  total:                                {len(directly_registered):>4}")
+
+    print(f"\nDecomposition table (not directly registered): {len(decomp_only_ops)}")
+    print(
+        "  These ops have entries in torch._decomp.decomposition_table but no\n"
+        "  direct DTensor strategy. They may work at runtime via\n"
+        "  DecompShardingStrategy if all decomposed sub-ops are supported.\n"
+        "  Additional ops beyond this count may also be reachable via CIA\n"
+        "  (CompositeImplicitAutograd) decompositions."
+    )
+
+    if verbose:
+
+        def _print_ops(label: str, ops: list) -> None:
+            print(f"\n{label} ({len(ops)}):")
+            for op in ops:
+                print(f"  {op}")
+
+        _print_ops("rule", rule_ops)
+        _print_ops("op_strategy", strategy_ops)
+        _print_ops("single_dim_strategy", single_dim_ops)
+        _print_ops("decomp table (not directly registered)", decomp_only_ops)
+
+
+if __name__ == "__main__":
+    # Override common size variables to ensure even sharding across world_size=2.
+    # These are process-global mutations, but this is a CLI entry point so the
+    # process exits after running.
+    opinfo_core.L = 24  # pyrefly: ignore[bad-assignment]
+    opinfo_core.M = 12  # pyrefly: ignore[bad-assignment]
+    opinfo_core.S = 4  # pyrefly: ignore[bad-assignment]
+    opinfo_core.XS = 2  # pyrefly: ignore[bad-assignment]
+
+    common_ops.L = 24  # pyrefly: ignore[bad-assignment]
+    common_ops.M = 12  # pyrefly: ignore[bad-assignment]
+    common_ops.S = 4  # pyrefly: ignore[bad-assignment]
+    common_ops.XS = 2  # pyrefly: ignore[bad-assignment]
+
+    parser = argparse.ArgumentParser(
+        description="Compare DTensor rules against ground truth"
+    )
+    parser.add_argument(
+        "--op",
+        default=None,
+        help="Operator name(s) to compare (comma-separated, supports glob "
+        'patterns, e.g., "relu,add" or "nn.functional.*")',
+    )
+    parser.add_argument(
+        "--all-registered",
+        action="store_true",
+        help="Test all ops with DTensor sharding rules registered",
+    )
+    parser.add_argument(
+        "--incorrect-only",
+        action="store_true",
+        help="Only test DTensor's claimed rules (faster, skips missing detection)",
+    )
+    parser.add_argument(
+        "--allow-composite",
+        action="store_true",
+        help="Validate each supported aten op individually for decomposed ops "
+        "(default skips non-1:1 aten mappings)",
+    )
+    parser.add_argument("--device", default="cuda", help="Device to use")
+    parser.add_argument("--dtype", default="float32", help="Dtype to use")
+    parser.add_argument(
+        "--world-size", type=int, default=2, help="Simulated world size"
+    )
+    parser.add_argument(
+        "--max-samples", type=int, default=None, help="Max samples to test"
+    )
+    parser.add_argument(
+        "--show-repro",
+        nargs="?",
+        const=1,
+        type=int,
+        default=0,
+        metavar="N",
+        help="Show N sample repros per rule (default 1 if flag given, -1 for all)",
+    )
+    parser.add_argument(
+        "--report",
+        action="store_true",
+        help="Show registration statistics (op counts per registration method)",
+    )
+    parser.add_argument(
+        "--report-full",
+        action="store_true",
+        help="Show full op lists with registration statistics (implies --report)",
+    )
+    args = parser.parse_args()
+
+    dtype_map = {
+        "float32": torch.float32,
+        "float64": torch.float64,
+        "float16": torch.float16,
+        "bfloat16": torch.bfloat16,
+    }
+    dtype = dtype_map.get(args.dtype, torch.float32)
+
+    dist.init_process_group("fake", rank=0, world_size=args.world_size)
+
+    if args.report or args.report_full:
+        report_registrations(verbose=args.report_full)
+        if not args.all_registered and not args.op:
+            dist.destroy_process_group()
+            sys.exit(0)
+        print()
+
+    try:
+        if args.all_registered:
+            op_names = get_registered_op_names()
+        elif args.op:
+            patterns = [p.strip() for p in args.op.split(",")]
+            op_names = resolve_op_names(patterns)
+        else:
+            op_names = ["add"]
+
+        # Preamble
+        display_names = [_format_op_name(n) for n in op_names]
+        print(f"Testing ops: {', '.join(display_names)}")
+        if args.allow_composite:
+            print(
+                "Mode: allow-composite (validates each aten op in decomposed samples)"
+            )
+        if args.incorrect_only:
+            print("Mode: incorrect-only (fast)")
+        print(f"Device: {args.device}, Dtype: {dtype}, World size: {args.world_size}")
+
+        op_results: list[tuple[str, ComparisonStats, float]] = []
+
+        for i, op_name in enumerate(op_names):
+            display = _format_op_name(op_name)
+
+            if op_name in SKIP_OPS:
+                print(
+                    f"\n[{i + 1}/{len(op_names)}] {display}"
+                    f" — skipped: {SKIP_OPS[op_name]}"
+                )
+                continue
+
+            try:
+                op_start = time.time()
+                stats = compare_operator(
+                    op_name,
+                    args.device,
+                    dtype,
+                    args.world_size,
+                    args.max_samples,
+                    verbose=True,
+                    incorrect_only=args.incorrect_only,
+                    allow_composite=args.allow_composite,
+                )
+                elapsed = time.time() - op_start
+
+                if stats.no_dtensor_support:
+                    print(
+                        f"\n[{i + 1}/{len(op_names)}] {display}"
+                        f" — skipped: no DTensor support"
+                    )
+                    continue
+
+                if stats.total_samples == 0 and stats.skip_reasons:
+                    reasons = ", ".join(
+                        f"{r} ({n})" for r, n in stats.skip_reasons.items()
+                    )
+                    print(f"\n[{i + 1}/{len(op_names)}] {display} — skipped: {reasons}")
+                    continue
+
+                op_results.append((op_name, stats, elapsed))
+
+                skipped = sum(stats.skip_reasons.values())
+                skipped_str = f" ({skipped} skipped)" if skipped else ""
+                print(
+                    f"\n[{i + 1}/{len(op_names)}] {display}"
+                    f" — Samples: {stats.total_samples}{skipped_str},"
+                    f" Combinations: {stats.total_combinations}"
+                )
+                print("-" * 70)
+                _print_comparison_summary(stats, show_repro=args.show_repro)
+            except Exception as e:
+                print(f"\n[{i + 1}/{len(op_names)}] {display} — Error: {e}")
+
+        if len(op_results) >= 1:
+            # Summary table
+            print("\n" + "=" * 70)
+            print("Summary")
+            print("=" * 70)
+
+            formatted_results = [
+                (_format_op_name(name), stats, elapsed)
+                for name, stats, elapsed in op_results
+            ]
+
+            col_widths = {
+                "op": max(len(name) for name, _, _ in formatted_results),
+                "tp": 7,
+                "fp": 9,
+                "fn": 7,
+                "time": 6,
+            }
+            col_widths["op"] = max(col_widths["op"], 2)  # min width
+
+            header = (
+                f"{'Op':<{col_widths['op']}}  "
+                f"{'Correct':>{col_widths['tp']}}  "
+                f"{'Incorrect':>{col_widths['fp']}}  "
+                f"{'Missing':>{col_widths['fn']}}  "
+                f"{'Time':>{col_widths['time']}}"
+            )
+            print(header)
+            print("-" * len(header))
+
+            total_tp = 0
+            total_fp = 0
+            total_fn = 0
+            total_time = 0.0
+
+            for name, stats, elapsed in formatted_results:
+                fp_rules = {
+                    (d.input_placements, d.output_placements)
+                    for d in stats.false_positives
+                }
+                fn_rules = {
+                    (d.input_placements, d.output_placements)
+                    for d in stats.false_negatives
+                }
+                n_fp = len(fp_rules)
+                n_fn = len(fn_rules)
+                total_tp += stats.true_positives
+                total_fp += n_fp
+                total_fn += n_fn
+                total_time += elapsed
+
+                fp_str = str(n_fp) if n_fp else "0"
+                fn_str = str(n_fn) if n_fn else "0"
+
+                print(
+                    f"{name:<{col_widths['op']}}  "
+                    f"{stats.true_positives:>{col_widths['tp']}}  "
+                    f"{fp_str:>{col_widths['fp']}}  "
+                    f"{fn_str:>{col_widths['fn']}}  "
+                    f"{elapsed:>{col_widths['time']}.1f}s"
+                )
+
+            print("-" * len(header))
+            fp_str = str(total_fp) if total_fp else "0"
+            fn_str = str(total_fn) if total_fn else "0"
+            print(
+                f"{'Total':<{col_widths['op']}}  "
+                f"{total_tp:>{col_widths['tp']}}  "
+                f"{fp_str:>{col_widths['fp']}}  "
+                f"{fn_str:>{col_widths['fn']}}  "
+                f"{total_time:>{col_widths['time']}.1f}s"
+            )
+    finally:
+        dist.destroy_process_group()
diff --git a/torch/distributed/tensor/_ops/utils.py b/torch/distributed/tensor/_ops/utils.py
index 2fc0c410e1d44..1381a9abba6f8 100644
--- a/torch/distributed/tensor/_ops/utils.py
+++ b/torch/distributed/tensor/_ops/utils.py
@@ -4,7 +4,7 @@
 import itertools
 import operator
 from collections.abc import Callable, Iterable, Sequence
-from typing import cast, Optional, TypeAlias, TypeVar, Union
+from typing import TypeAlias, TypeVar
 
 import torch
 from torch._prims_common import DimsSequenceType, DimsType
@@ -22,6 +22,7 @@
 )
 from torch.distributed.tensor.device_mesh import DeviceMesh
 from torch.distributed.tensor.placement_types import (
+    _is_shard_like,
     _StridedShard,
     Partial,
     Placement,
@@ -32,9 +33,9 @@
 
 def _get_registration_wrapper(
     registration_fn,
-    op: Union[torch._ops.OpOverload, list[torch._ops.OpOverload]],
-    schema_info: Optional[RuntimeSchemaInfo],
-    arg_names_that_require_specializing_cache_strategy: Optional[list[str]],
+    op: torch._ops.OpOverload | list[torch._ops.OpOverload],
+    schema_info: RuntimeSchemaInfo | None,
+    arg_names_that_require_specializing_cache_strategy: list[str] | None,
 ):
     def wrapper(impl):
         overloads = op if isinstance(op, list) else [op]
@@ -194,7 +195,8 @@ def is_tensor_shardable(
     """
     from torch.fx.experimental.symbolic_shapes import guard_or_false, guard_or_true
 
-    assert allow_unbacked_sharding in [None, True, False]
+    if allow_unbacked_sharding not in [None, True, False]:
+        raise AssertionError
     guard_fn = {
         None: bool,
         True: guard_or_false,
@@ -204,7 +206,7 @@ def is_tensor_shardable(
     # number of shards in each tensor dimension
     num_shards = [1] * len(shape)
     for i, placement in enumerate(spec.placements):
-        if isinstance(placement, Shard | _StridedShard):
+        if _is_shard_like(placement):
             shard_dim = placement.dim
             if shard_dim >= len(shape):
                 return False
@@ -228,7 +230,7 @@ def is_tensor_evenly_shardable(shape: Sequence[int], spec: DTensorSpec) -> bool:
     # number of shards in each tensor dimension
     num_shards = [1] * len(shape)
     for i, placement in enumerate(spec.placements):
-        if isinstance(placement, Shard | _StridedShard):
+        if _is_shard_like(placement):
             shard_dim = placement.dim
             if shard_dim >= len(shape):
                 return False
@@ -253,17 +255,24 @@ def is_tensor_evenly_shardable_on_dim(
 
     num_shards = 1
     for i, placement in enumerate(spec.placements):
-        if placement.is_shard():
-            shard_dim = cast(Shard, placement).dim
-            if shard_dim == dim:
-                num_shards *= spec.mesh.size(i)
+        if _is_shard_like(placement) and placement.dim == dim:
+            num_shards *= spec.mesh.size(i)
+            if isinstance(placement, _StridedShard):
+                # _StridedShard._split_tensor first chunks into split_factor
+                # groups, then into num_shards within each group, so the dim
+                # must be divisible by the product of both.  This is stricter
+                # than the final num_shards check and implies it.  Note:
+                # num_shards already includes spec.mesh.size(i) from this
+                # iteration, so the check covers the full shard count.
+                if shape[dim] % (placement.split_factor * num_shards) != 0:
+                    return False
 
     return shape[dim] % num_shards == 0
 
 
 def is_tensor_dim_sharded(spec: DTensorSpec, dim: int) -> bool:
     """Return True if tensor dim is sharded."""
-    return any(p.is_shard(dim) for p in spec.placements)
+    return any(_is_shard_like(p) and p.dim == dim for p in spec.placements)
 
 
 def is_tensor_partial(spec: DTensorSpec) -> bool:
@@ -279,11 +288,13 @@ def infer_broadcast_dims_map(
     # e.g. if common_shape = [1, 2, 3, 4] and input_shape = [2, 3, 4],
     # broadcast_dims_map will be [-1, 0, 1, 2]
     # meaning that dim 0 in the output has no mapping to the input, and dim 1 in the output maps to dim 0 in the input
+    from torch.fx.experimental.symbolic_shapes import guard_or_false
+
     common_ndim = len(common_shape)
     input_ndim = len(input_shape)
     broadcast_dims_map = [-1] * common_ndim
     for idx in range(-1, -1 - input_ndim, -1):
-        if input_shape[idx] == common_shape[idx]:
+        if guard_or_false(input_shape[idx] == common_shape[idx]):
             broadcast_dims_map[common_ndim + idx] = input_ndim + idx
     return broadcast_dims_map
 
@@ -306,7 +317,8 @@ def map_placements_after_broadcast(
         elif isinstance(placement, Replicate):
             new_placements.append(placement)
         else:
-            assert isinstance(placement, Shard | _StridedShard)
+            if not _is_shard_like(placement):
+                raise AssertionError
             shard_dim = normalize_dim(placement.dim, len(shape))
             new_shard_dim = broadcast_dims_map[shard_dim]
             if new_shard_dim != -1:
@@ -358,10 +370,12 @@ def expand_to_full_mesh_op_strategy(
     input_index: int = 1,
     inplace_op: bool = False,
     allow_unbacked_sharding: bool | None = None,
+    allow_uneven_sharding: bool = False,
     is_valid_strategy_cb: Callable[
         [list[DTensorSpec], DTensorSpec | tuple[DTensorSpec | None, ...]], bool
     ]
     | None = None,
+    different_mesh_args: list[int] | None = None,
 ) -> OpStrategy:
     """
     Convenience function to allow writing a sharding strategy considering only a single mesh dimension,
@@ -398,6 +412,18 @@ def expand_to_full_mesh_op_strategy(
     args_strategy = op_schema.args_strategy
     kwargs_strategy = op_schema.kwargs_strategy
     input_args_strategy = args_strategy + kwargs_strategy
+
+    # Propagate use_strided_shard_as_shard_order from inputs so that
+    # strategy specs with _StridedShard get the correct flag (and thus
+    # correct shard_order) at construction time, avoiding shard_order
+    # mismatches in redistribute_cost computation.
+    _input_use_strided: bool | None = None
+    for input_strat in input_args_strategy:
+        input_spec = input_strat.strategies[0].output_spec
+        if any(isinstance(p, _StridedShard) for p in input_spec.placements):
+            _input_use_strided = input_spec.use_strided_shard_as_shard_order
+            break
+
     all_strategies = []
     # Track input placements if we skip strategies due to inplace placement mismatch
     blocking_inplace_input_placements: tuple[Placement, ...] | None = None
@@ -436,10 +462,45 @@ def expand_to_full_mesh_op_strategy(
                         input_strategy_counter += 1
 
                 # pyrefly: ignore [bad-argument-type]
-                spec_list.append(DTensorSpec(mesh, specs, tensor_meta=tensor_meta))
+                use_strided = (
+                    _input_use_strided
+                    if _input_use_strided is not None
+                    and any(isinstance(p, _StridedShard) for p in specs)
+                    else None
+                )
+                spec_list.append(
+                    DTensorSpec(
+                        mesh,
+                        specs,
+                        tensor_meta=tensor_meta,
+                        use_strided_shard_as_shard_order=use_strided,
+                    )
+                )
             else:
                 spec_list.append(None)
 
+        # Skip strategy combinations that would create mixed partial types
+        # (except sum+avg which commute with each other).
+        # We check (type, reduce_op) pairs rather than just reduce_op because
+        # Partial subclasses like _MaskPartial have different reduction semantics
+        # even when they share the same reduce_op string.
+        has_mixed_partial = False
+        for spec in spec_list:
+            if spec is not None:
+                partial_kinds = {
+                    (type(p), p.reduce_op)
+                    for p in spec.placements
+                    if isinstance(p, Partial)
+                }
+                if len(partial_kinds) > 1:
+                    reduce_ops = {ro for _, ro in partial_kinds}
+                    types = {t for t, _ in partial_kinds}
+                    if not (len(types) == 1 and reduce_ops == {"sum", "avg"}):
+                        has_mixed_partial = True
+                        break
+        if has_mixed_partial:
+            continue
+
         input_specs: list[DTensorSpec] = [
             s for s in spec_list[input_index:] if isinstance(s, DTensorSpec)
         ]
@@ -449,11 +510,62 @@ def expand_to_full_mesh_op_strategy(
                 f"input_specs({len(input_specs)}) != strategies({len(input_args_strategy)}: "
                 f"{len(args_strategy)} args + {len(kwargs_strategy)} kwargs)"
             )
+
+        # Note [Multi-mesh args]
+        #
+        # Some ops accept args whose DTensor lives on a different DeviceMesh
+        # than the op's primary compute mesh.  We call these "multi-mesh
+        # args".  They arise in fused optimizer ops (e.g. _fused_adam_)
+        # where *state_steps* is a per-rank scalar counter allocated on a
+        # smaller sub-mesh (e.g. 1-D DP) while params and grads live on a
+        # larger mesh (e.g. 2-D DP × TP).
+        #
+        # Why must these args be Replicate?
+        #   Sharding implies a specific partitioning of a tensor's data
+        #   across the ranks of a mesh.  If a tensor doesn't even *exist*
+        #   on the compute mesh, there is no meaningful way to interpret a
+        #   Shard placement for it.  Replicate, on the other hand, is
+        #   mesh-agnostic: every rank already holds the full data, so the
+        #   op can simply read the value regardless of which mesh owns it.
+        #
+        # What we do here:
+        #   We preserve the original mesh and Replicate placement for these
+        #   args so the propagator does not try to redistribute them onto
+        #   the compute mesh (which would fail or produce wrong results).
+        #
+        # This is distinct from the *element_mesh* handling in
+        # single_dim_strategy.py, which deals with foreach ops where
+        # different *elements* in a tensor list may live on different
+        # sub-meshes (e.g. param group A on 2-D mesh, param group B on
+        # 1-D mesh).
+        # TODO: refactor fused_ops handling so that there are no longer
+        # args on different meshes
+        if different_mesh_args is not None:
+            for idx in different_mesh_args:
+                if idx < len(input_args_strategy):
+                    cross_mesh_input = input_args_strategy[idx]
+                    original_spec = cross_mesh_input.strategies[0].output_spec
+                    if original_spec.mesh != mesh:
+                        if not all(p == Replicate() for p in original_spec.placements):
+                            raise RuntimeError(
+                                f"Cross-mesh input at index {idx} must be Replicate, "
+                                f"but got {original_spec.placements}"
+                            )
+                        input_specs[idx] = DTensorSpec(
+                            mesh=original_spec.mesh,
+                            placements=original_spec.placements,
+                            tensor_meta=original_spec.tensor_meta,
+                        )
         self_spec = input_args_strategy[0].strategies[0].output_spec
 
-        if inplace_op and self_spec.placements != input_specs[0].placements:
-            # if it's inplace op, we would only allow the OpSpec to be added when the
-            # input_spec matches the first argument's runtime sharding, otherwise we skip
+        redistribute_input = self_spec.placements != input_specs[0].placements
+        mismatching_input_output = (
+            spec_list[0] is not None and spec_list[0].placements != self_spec.placements
+        )
+        if inplace_op and (redistribute_input or mismatching_input_output):
+            # For inplace ops, both the proposed input[0] and the output must
+            # match self's runtime placement: input[0] because self can't be
+            # redistributed, output because the result IS self.
             if blocking_inplace_input_placements is None:
                 blocking_inplace_input_placements = self_spec.placements
             continue
@@ -487,6 +599,10 @@ def expand_to_full_mesh_op_strategy(
             is_tensor_shardable(
                 inp.shape, s, allow_unbacked_sharding=allow_unbacked_sharding
             )
+            or (
+                allow_uneven_sharding
+                and inp.strategies[0].output_spec.placements == s.placements
+            )
             for inp, s in zip(input_args_strategy, input_specs)
         ):
             continue
@@ -528,7 +644,11 @@ def shift_shard_dims_after_insert(
 ) -> Sequence[Placement]:
     normalized_placements: list[Placement] = []
     for placement in placements:
-        if isinstance(placement, Shard) and placement.dim >= insert_dim:
+        if isinstance(placement, _StridedShard) and placement.dim >= insert_dim:
+            normalized_placements.append(
+                _StridedShard(placement.dim + 1, split_factor=placement.split_factor)
+            )
+        elif isinstance(placement, Shard) and placement.dim >= insert_dim:
             normalized_placements.append(Shard(placement.dim + 1))
         else:
             normalized_placements.append(placement)
@@ -540,7 +660,11 @@ def shift_shard_dims_after_remove(
 ) -> Sequence[Placement]:
     normalized_placements: list[Placement] = []
     for placement in placements:
-        if isinstance(placement, Shard) and placement.dim > remove_dim:
+        if isinstance(placement, _StridedShard) and placement.dim > remove_dim:
+            normalized_placements.append(
+                _StridedShard(placement.dim - 1, split_factor=placement.split_factor)
+            )
+        elif isinstance(placement, Shard) and placement.dim > remove_dim:
             normalized_placements.append(Shard(placement.dim - 1))
         else:
             normalized_placements.append(placement)
diff --git a/torch/distributed/tensor/_random.py b/torch/distributed/tensor/_random.py
index 423b00b0a1a8c..ebbdd7a8d97ad 100644
--- a/torch/distributed/tensor/_random.py
+++ b/torch/distributed/tensor/_random.py
@@ -11,6 +11,7 @@
 from torch.distributed.device_mesh import _get_device_handle, DeviceMesh
 from torch.distributed.tensor._dtensor_spec import DTensorSpec
 from torch.distributed.tensor.placement_types import _StridedShard, Shard
+from torch.types import IntLikeType
 
 
 logger = getLogger(__name__)
@@ -130,7 +131,8 @@ def offset(self) -> torch.Tensor:
 
     @offset.setter
     def offset(self, offset: torch.Tensor) -> None:
-        assert offset.numel() == 1
+        if offset.numel() != 1:
+            raise AssertionError
         self._state[8:] = offset.view(torch.uint8)
 
     @property
@@ -139,7 +141,8 @@ def seed(self) -> torch.Tensor:
 
     @seed.setter
     def seed(self, seed: torch.Tensor) -> None:
-        assert seed.numel() == 1
+        if seed.numel() != 1:
+            raise AssertionError
         self._state[:8] = seed.view(torch.uint8)
 
 
@@ -153,7 +156,6 @@ class _RNGStateTracker:
     """
 
     def __init__(self, device: torch.device):
-        # pyrefly: ignore [read-only]
         self._device = device
         self._device_handle = _get_device_handle(self._device.type)
         if not (self._device_handle and self._device_handle.is_available()):
@@ -195,7 +197,8 @@ def __init__(
         run_state_sync: bool = True,
     ):
         super().__init__(_resolve_device(device_mesh=device_mesh))
-        assert self._device_handle is not None
+        if self._device_handle is None:
+            raise AssertionError
         # DTensor RNG tracker so far only supports CUDA/CUDA-like devices
         if self._device.type == "cpu":
             raise RuntimeError(
@@ -203,8 +206,8 @@ def __init__(
                 f"CUDA/CUDA-like/XPU device. Got {self._device.type} instead."
             )
 
-        rng_state = self._get_device_state()
         if run_state_sync:
+            rng_state = self._get_device_state()
             # synchronize RNG state using rank 0's current one
             torch.distributed.broadcast(rng_state, 0)
             my_rng_state = self._get_device_state()
@@ -266,7 +269,8 @@ def _distribute_region(
             with torch.random.fork_rng(
                 devices=[self._device], device_type=self._device.type
             ):
-                assert self._device_handle is not None
+                if self._device_handle is None:
+                    raise AssertionError
                 self._device_handle.set_rng_state(state.state)
                 try:
                     yield  # execute the region code
@@ -334,34 +338,8 @@ def _set_pre_op_offset(self, state: _PhiloxState, spec: DTensorSpec) -> None:
             The last value to calculate before obtaining the starting offset is the shard linear index.
             The starting offset for each rank will be its shard_linear_index * local_tensor_numel.
         """
-        mesh = spec.mesh
-        mesh_coordinate = mesh.get_coordinate()
-        assert mesh_coordinate is not None
-
-        # Compute shard index and total number of shards on each tensor dim
-        shard_idx_by_dim, total_num_shards_by_dim = _calc_shard_info(
-            mesh_coordinate, spec
-        )
-
-        # compute shard linear index
-        shard_linear_idx = self._calc_shard_linear_idx(
-            shard_idx_by_dim, total_num_shards_by_dim
-        )
-
-        # compute starting offset using the first shard's size
-        local_size_on_rank_0 = _calc_first_shard_size(spec)
-
-        from torch.distributed.tensor._ops.utils import prod
-
-        local_size = prod(local_size_on_rank_0)
-
-        # get current RNG offset
-        current_offset = state.offset
-
-        # pytorch: offset must be multiple of 4
-        # source: aten/src/ATen/cuda/CUDAGeneratorImpl.cpp
-        offset_incr = (shard_linear_idx * local_size + 3) // 4 * 4
-        state.offset = current_offset + offset_incr
+        start_offset_incr, _ = self._compute_rng_offsets(spec)
+        state.offset = state.offset + start_offset_incr
 
     def _set_post_op_offset(
         self, state: _PhiloxState, spec: DTensorSpec, old_offset: torch.Tensor
@@ -379,19 +357,41 @@ def _set_post_op_offset(
         Returns:
             None
         """
-        dtensor_shape = spec.shape
+        _, end_offset_incr = self._compute_rng_offsets(spec)
+        state.offset = old_offset + end_offset_incr
+
+    def _compute_rng_offsets(self, spec: DTensorSpec) -> tuple[int, int]:
+        """Compute the RNG offset increments for a distributed random op.
+
+        These values are derived from mesh topology, placements, and tensor shape,
+        and are static for a given compiled graph. They can be burned into the graph
+        as integer constants rather than keeping the DTensorSpec around at runtime.
 
+        Returns:
+            (start_offset_incr, end_offset_incr) — both aligned to multiples of 4.
+        """
         from torch.distributed.tensor._ops.utils import prod
 
-        numel = prod(dtensor_shape)
+        mesh = spec.mesh
+        mesh_coordinate = [mesh._sym_get_coordinate(i) for i in range(mesh.ndim)]
+
+        shard_idx_by_dim, total_num_shards_by_dim = _calc_shard_info(
+            mesh_coordinate, spec
+        )
+        shard_linear_idx = self._calc_shard_linear_idx(
+            shard_idx_by_dim, total_num_shards_by_dim
+        )
+        local_size = prod(_calc_first_shard_size(spec))
         # pytorch: offset must be multiple of 4
         # source: aten/src/ATen/cuda/CUDAGeneratorImpl.cpp
-        numel = (numel + 3) // 4 * 4
-        state.offset = old_offset + numel
+        start_offset_incr = (shard_linear_idx * local_size + 3) // 4 * 4
+        end_offset_incr = (prod(spec.shape) + 3) // 4 * 4
+
+        return start_offset_incr, end_offset_incr
 
     def _calc_shard_linear_idx(
-        self, shard_coord: list[int], shard_size: list[int]
-    ) -> int:
+        self, shard_coord: Sequence[IntLikeType], shard_size: Sequence[IntLikeType]
+    ) -> IntLikeType:
         return _calc_shard_linear_idx(shard_coord, shard_size)
 
 
@@ -410,8 +410,8 @@ def _calc_first_shard_size(spec: DTensorSpec) -> list[int]:
 
 
 def _calc_shard_info(
-    mesh_coordinate: Sequence[int], spec: DTensorSpec
-) -> tuple[list[int], list[int]]:
+    mesh_coordinate: Sequence[IntLikeType], spec: DTensorSpec
+) -> tuple[list[IntLikeType], list[IntLikeType]]:
     mesh = spec.mesh
     # note: dim_map does not allow double sharding which is the FSDP(fully_shard)+TP
     # case. Replace the custom logic with dim_map once we support it.
@@ -423,20 +423,22 @@ def _calc_shard_info(
                 dim_map[shard_dim] = [i]
             else:
                 mesh_dim_list = dim_map[shard_dim]
-                assert isinstance(mesh_dim_list, list)
+                if not isinstance(mesh_dim_list, list):
+                    raise AssertionError
                 mesh_dim_list.append(i)
 
     # Compute shard coordinate:
     # The coordinate on each tensor dim is a tuple (idx, range)
     # If a DTensor is partitioned on its dim i into n shards, and the current rank
     # holds the j-th, then its shard coordinate will be (idx=j, range=n) on dim i
-    assert mesh_coordinate is not None
     mesh_size = mesh.shape
     shard_idx_by_dim = []
-    total_num_shards_by_dim = []  # total number of shards on each tensor dim
+    total_num_shards_by_dim: list[
+        IntLikeType
+    ] = []  # total number of shards on each tensor dim
     for mesh_dim in dim_map:
-        shard_idx = 0
-        total_num_shards = 1
+        shard_idx: IntLikeType = 0
+        total_num_shards: IntLikeType = 1
         # the tensor dim is sharded on more than 1 mesh dim
         if isinstance(mesh_dim, list):
             rank_coord = [mesh_coordinate[d] for d in mesh_dim]
@@ -451,10 +453,12 @@ def _calc_shard_info(
     return shard_idx_by_dim, total_num_shards_by_dim
 
 
-def _calc_shard_linear_idx(shard_coord: list[int], shard_size: list[int]) -> int:
+def _calc_shard_linear_idx(
+    shard_coord: Sequence[IntLikeType], shard_size: Sequence[IntLikeType]
+) -> IntLikeType:
     # compute shard linear index
-    shard_linear_idx = 0
-    shard_coord_stride = 1
+    shard_linear_idx: IntLikeType = 0
+    shard_coord_stride: IntLikeType = 1
     for idx, size in zip(reversed(shard_coord), reversed(shard_size)):
         shard_linear_idx += idx * shard_coord_stride
         shard_coord_stride *= size
@@ -465,7 +469,8 @@ def _calc_shard_linear_idx(shard_coord: list[int], shard_size: list[int]) -> int
 def _resolve_device(device_mesh: DeviceMesh) -> torch.device:
     device_type = device_mesh.device_type
     device_handle = _get_device_handle(device_type)
-    assert device_handle is not None
+    if device_handle is None:
+        raise AssertionError
     device_idx = device_mesh.get_rank() % device_handle.device_count()
 
     @maybe_run_for_local_tensor
diff --git a/torch/distributed/tensor/_redistribute.py b/torch/distributed/tensor/_redistribute.py
index 9db77fe87afa3..3608e3edcbc2c 100644
--- a/torch/distributed/tensor/_redistribute.py
+++ b/torch/distributed/tensor/_redistribute.py
@@ -17,6 +17,7 @@
 from torch.distributed._functional_collectives import _are_we_tracing
 from torch.distributed.tensor._collective_utils import one_step_redistribute_cost
 from torch.distributed.tensor._dtensor_spec import (
+    _StridedShardNotDecodableError,
     DTensorSpec,
     ShardOrder,
     ShardOrderEntry,
@@ -25,12 +26,14 @@
 from torch.distributed.tensor._utils import assert_no_mixed_partial_types
 from torch.distributed.tensor.device_mesh import DeviceMesh
 from torch.distributed.tensor.placement_types import (
+    _is_shard_like,
     _StridedShard,
     Partial,
     Placement,
     Replicate,
     Shard,
 )
+from torch.types import IntLikeType
 from torch.utils._debug_mode import get_active_debug_mode
 
 
@@ -144,13 +147,15 @@ class _TransformInfo:
     mesh_dim: int
     src_dst_placements: tuple[Placement, Placement]
     # logical_shape on this mesh dimension
-    logical_shape: list[int]
+    logical_shape: Sequence[IntLikeType]
 
     def __post_init__(self):
-        assert self.mesh_dim >= 0
-        assert self.src_dst_placements[0] != self.src_dst_placements[1], (
-            "TransformInfo should only be created if it is an op with some effect, not a no-op"
-        )
+        if self.mesh_dim < 0:
+            raise AssertionError
+        if self.src_dst_placements[0] == self.src_dst_placements[1]:
+            raise AssertionError(
+                "TransformInfo should only be created if it is an op with some effect, not a no-op"
+            )
 
     def _comm_type_key(self) -> str | None:
         """
@@ -162,11 +167,11 @@ def _comm_type_key(self) -> str | None:
         src, dst = self.src_dst_placements
         if src.is_partial() and dst.is_replicate():
             return "all_reduce"
-        elif src.is_partial() and dst.is_shard():
+        elif src.is_partial() and _is_shard_like(dst):
             return "reduce_scatter"
-        elif src.is_shard() and dst.is_replicate():
+        elif _is_shard_like(src) and dst.is_replicate():
             return "all_gather"
-        elif src.is_shard() and dst.is_shard():
+        elif _is_shard_like(src) and _is_shard_like(dst):
             return "all_to_all"
         else:
             # Local ops (Replicate->Shard, Replicate->Partial, noop, etc.)
@@ -194,9 +199,10 @@ class _FlattenedTransformInfo(_TransformInfo):
     def __post_init__(self) -> None:
         _TransformInfo.__post_init__(self)
         if self.avg_scale is not None:
-            assert self.avg_scale > 1, (
-                f"avg_scale must be > 1 if set, got {self.avg_scale}"
-            )
+            if not self.avg_scale > 1:
+                raise AssertionError(
+                    f"avg_scale must be > 1 if set, got {self.avg_scale}"
+                )
 
 
 def _update_shard_order_and_placements(
@@ -248,18 +254,14 @@ def _update_shard_order_and_placements(
         current_placements[mesh_dim] = dst_placement
 
 
-def _get_flattened_mesh_by_layout(
+def _get_flattened_mesh_by_layout_impl(
     mesh: DeviceMesh, mesh_dims: tuple[int, ...]
 ) -> DeviceMesh | None:
     """
     Query for an explicitly created flattened mesh using layout comparison.
 
-    Args:
-        mesh: The DeviceMesh to query
-        mesh_dims: Tuple of mesh dimension indices to look for
-
-    Returns:
-        The flattened DeviceMesh if it was explicitly created, None otherwise.
+    Searches root_mesh._flatten_mapping for a mesh whose layout matches
+    the expected flattened layout for the given dims. Pure Python layout math.
     """
     root_mesh = mesh._get_root_mesh()
     mesh_dim_names = mesh.mesh_dim_names
@@ -285,6 +287,29 @@ def _get_flattened_mesh_by_layout(
     return None
 
 
+def _get_flattened_mesh_by_layout(
+    mesh: DeviceMesh, mesh_dims: tuple[int, ...]
+) -> DeviceMesh | None:
+    """
+    Query for an explicitly created flattened mesh using layout comparison.
+
+    When tracing with compile_on_one_rank, delegates to a custom op so the
+    flattened mesh appears as a call_function node derived from mesh (a graph
+    input) rather than as a get_attr constant holding an unpicklable
+    ProcessGroup.
+    """
+    if _are_we_tracing() and torch.distributed.config.compile_on_one_rank:
+        # Pre-check: the custom op can't return None (torch.library doesn't
+        # support Optional opaque return types), so guard here first.
+        if _get_flattened_mesh_by_layout_impl(mesh, mesh_dims) is None:
+            return None
+        from torch.distributed._ops import device_mesh as _  # noqa: F401
+
+        return torch.ops.device_mesh._get_flattened_submesh(mesh, list(mesh_dims))
+
+    return _get_flattened_mesh_by_layout_impl(mesh, mesh_dims)
+
+
 # Track (mesh_hash, mesh_dims, reason) we've already warned about to avoid repeated warnings
 _warned_flatten_issues: set[tuple[int, tuple[int, ...], str]] = set()
 
@@ -443,10 +468,13 @@ def try_create_flattened(
         # (e.g., can't merge Partial->Shard(0) with Partial->Shard(1))
         first_placements = infos[0].src_dst_placements
         comm_type = infos[0]._comm_type_key()
-        assert all(
+        if not all(
             are_placements_mergeable(info.src_dst_placements, first_placements)
             for info in infos
-        )
+        ):
+            raise AssertionError(
+                "All transforms must have mergeable src_dst_placements"
+            )
         mesh_dims = tuple(info.mesh_dim for info in infos)
         sorted_mesh_dims = tuple(sorted(mesh_dims))
 
@@ -736,7 +764,10 @@ def stringify_transform_infos(
         Returns:
             A string showing the sequence of DistState transitions, separated by '->'.
         """
-        assert len(src_placement) == mesh.ndim
+        if len(src_placement) != mesh.ndim:
+            raise AssertionError(
+                f"dimensions mismatch {len(src_placement)} vs {mesh.ndim}"
+            )
         if use_strided_shard_as_shard_order:
             src_placement, src_shard_order = (
                 DTensorSpec._normalize_placements_into_shard_order(
@@ -793,8 +824,10 @@ def __init__(
             dtensor_meta: TensorMeta of the DTensor to redistribute
         """
         self.device_mesh = device_mesh
-        assert device_mesh._is_current_rank_part_of_mesh()
-        assert dtensor_meta is not None
+        if not device_mesh._is_current_rank_part_of_mesh():
+            raise AssertionError
+        if dtensor_meta is None:
+            raise AssertionError
         self.dtensor_meta = dtensor_meta
         self.tensor_dimension = len(dtensor_meta.shape)
         self.strided_shard_placements_in_target: set[_StridedShard] = set()
@@ -1165,12 +1198,13 @@ def get_logical_shape(
         src_state: "DTensorRedistributePlanner.DistState",
         mesh_dim: int,
         full_tensor_shape: tuple[int, ...],
-    ) -> list[int]:
-        new_logical_shape = list(full_tensor_shape)
+    ) -> list[IntLikeType]:
+        new_logical_shape: list[IntLikeType] = list(full_tensor_shape)
         for entry in src_state.tensor_dim_to_mesh_dim:
             tensor_dim = entry.tensor_dim
             mesh_dims = entry.mesh_dims
-            assert len(mesh_dims) > 0
+            if len(mesh_dims) <= 0:
+                raise AssertionError
             for mdim in mesh_dims:
                 if mdim == mesh_dim:
                     continue
@@ -1339,6 +1373,11 @@ def generate_greedy_transform_infos(
                 target = target_placements[mesh_dim]
                 # If target is not Shard, we can directly redistribute since we
                 # are traversing from inner to outer placements here
+                # TODO: extend nested sharding detection to _StridedShard
+                # (isinstance check and is_shard() below miss it).
+                # Safe today: strategies convert _StridedShard to Replicate
+                # on ALL mesh dims for a given reduction dim, so misaligned
+                # nested _StridedShard targets can't arise.
                 if isinstance(target, Shard):
                     # If target is Shard, check for nested sharding on the
                     # tensor dim BEFORE the current mesh_dim
@@ -1418,15 +1457,29 @@ def _gen_transform_infos_non_cached(
         use_graph_based_transform = _FORCE_MIN_COST_REDISTRIBUTION_PLAN
     elif use_graph_based_transform is None:
         use_graph_based_transform = False
-    assert src_spec.tensor_meta is not None
+    if src_spec.tensor_meta is None:
+        raise AssertionError
     drp = get_redistribute_planner(
         device_mesh,
         src_spec.tensor_meta,
     )
     if use_graph_based_transform:
-        transform_infos = drp.generate_graph_based_transform_infos(
-            src_spec, dst_spec, src_spec.shape
-        )
+        # TODO(zpcore): Temporary workaround for the case where _StridedShard
+        # cannot be decoded into shard order. This happens when
+        # use_strided_shard_as_shard_order defaults to True (e.g. in
+        # Redistribute.forward where the target DTensorSpec is constructed from
+        # raw placements without the flag), but the split_factor doesn't
+        # correspond to any valid product of mesh dimension sizes (e.g. sf=2
+        # on a 1D mesh). A proper fix is to either pass
+        # use_strided_shard_as_shard_order through the Redistribute API, or
+        # migrate to explicit shard_order so _StridedShard is no longer
+        # overloaded for two purposes.
+        try:
+            transform_infos = drp.generate_graph_based_transform_infos(
+                src_spec, dst_spec, src_spec.shape
+            )
+        except _StridedShardNotDecodableError:
+            transform_infos = drp.generate_greedy_transform_infos(src_spec, dst_spec)
     else:
         transform_infos = drp.generate_greedy_transform_infos(src_spec, dst_spec)
     return transform_infos
@@ -1528,6 +1581,15 @@ def redistribute_local_tensor(
                 mesh_to_use = device_mesh
             i = transform_info.mesh_dim
             current, target = transform_info.src_dst_placements
+
+            # _StridedShard methods use device_mesh directly, not mesh_to_use.
+            # This is safe because _StridedShard.is_shard() returns False, so
+            # _comm_type_key() returns None and flattening is never attempted.
+            if isinstance(current, _StridedShard) or isinstance(target, _StridedShard):
+                assert mesh_to_use is device_mesh, (  # noqa: S101
+                    "_StridedShard redistribute assumes no flattened transforms"
+                )
+
             num_chunks = mesh_to_use.size(mesh_dim=i)
 
             if current == target:
@@ -1598,8 +1660,15 @@ def redistribute_local_tensor(
                             target_placement.dim,
                         )
                 elif isinstance(current, _StridedShard):
-                    raise NotImplementedError(
-                        "Redistribute from _StridedShard to Shard is not implemented yet"
+                    # _StridedShard -> Shard: go via Replicate as intermediate
+                    replicated = current._to_replicate_tensor(
+                        local_tensor, device_mesh, i, transform_info.logical_shape
+                    )
+                    new_local_tensor = target_placement._replicate_to_shard(
+                        replicated,
+                        mesh_to_use,
+                        i,
+                        mesh_to_use._sym_get_coordinate(i),
                     )
                 else:
                     raise ValueError(
@@ -1611,7 +1680,7 @@ def redistribute_local_tensor(
                     new_local_tensor = partial_spec._partition_value(
                         local_tensor, mesh_to_use, i
                     )
-                elif current.is_shard() or isinstance(current, _StridedShard):
+                elif _is_shard_like(current):
                     raise RuntimeError(
                         f"redistribute from {current} to {target} not supported yet"
                     )
@@ -1625,8 +1694,13 @@ def redistribute_local_tensor(
             elif isinstance(target, _StridedShard):
                 # Case 4: target is _StridedShard
                 if current.is_partial():
-                    raise NotImplementedError(
-                        "Redistribute from Partial to _StridedShard is not implemented yet"
+                    # Partial -> _StridedShard: reduce to Replicate, then strided shard
+                    partial_spec = cast(Partial, current)
+                    replicated = partial_spec._reduce_value(
+                        local_tensor, mesh_to_use, i
+                    )
+                    new_local_tensor = target._replicate_to_strided_shard(
+                        replicated, device_mesh, i, device_mesh._sym_get_coordinate(i)
                     )
                 elif current.is_replicate():
                     # split the tensor and return the corresponding local strided shard
@@ -1634,9 +1708,13 @@ def redistribute_local_tensor(
                         local_tensor, device_mesh, i, device_mesh._sym_get_coordinate(i)
                     )
                 elif current.is_shard():
-                    # Shard -> _StridedShard on potentially different dimensions
-                    raise NotImplementedError(
-                        "Redistribute from Shard to _StridedShard is not implemented yet"
+                    # Shard -> _StridedShard: all-gather to Replicate, then strided shard
+                    current_placement = cast(Shard, current)
+                    replicated = current_placement._to_replicate_tensor(
+                        local_tensor, mesh_to_use, i, transform_info.logical_shape
+                    )
+                    new_local_tensor = target._replicate_to_strided_shard(
+                        replicated, device_mesh, i, device_mesh._sym_get_coordinate(i)
                     )
                 elif isinstance(current, _StridedShard):
                     # _StridedShard -> _StridedShard: go through Replicate
@@ -1660,6 +1738,99 @@ def redistribute_local_tensor(
     return new_local_tensor
 
 
+def _redistribute_backward(
+    grad_output: "dtensor.DTensor",
+    previous_spec: DTensorSpec,
+    original_dtype: torch.dtype | None = None,
+    backward_dtype: torch.dtype | None = None,
+    async_op: bool = False,
+):
+    """
+    Common function for redistributing a distributed tensor during backward
+    and twice-backward backpropagation steps.
+
+    Args:
+        grad_output: The output gradient tensor.
+        previous_spec: DTensorSpec prior to redistribution.
+        original_dtype: Original output tensor dtype from forward pass (for type checking)
+        backward_dtype: Desired data type for backwards output.
+        async_op: whether to perform the DTensor redistribute operation
+                asynchronously or not. Default: False
+
+    Returns:
+        A :class:`torch.Tensor` object.
+        A :class:`DTensorSpec` object.
+    """
+    if backward_dtype is not None and backward_dtype != grad_output._local_tensor.dtype:
+        local_tensor = grad_output._local_tensor.to(dtype=backward_dtype)
+        current_spec = DTensorSpec(
+            mesh=grad_output._spec.device_mesh,
+            placements=grad_output._spec.placements,
+            tensor_meta=TensorMeta(
+                shape=grad_output.shape,
+                stride=grad_output.stride(),
+                # pyrefly: ignore [bad-argument-type]
+                dtype=backward_dtype,
+            ),
+            use_strided_shard_as_shard_order=grad_output._spec.use_strided_shard_as_shard_order,
+        )
+        previous_spec = DTensorSpec(
+            mesh=previous_spec.device_mesh,
+            placements=previous_spec.placements,
+            tensor_meta=current_spec.tensor_meta,
+            use_strided_shard_as_shard_order=previous_spec.use_strided_shard_as_shard_order,
+        )
+    else:
+        local_tensor = grad_output._local_tensor
+        current_spec = grad_output._spec
+    # skip the replicate to partial transformation when we are in backward pass
+    # In this case we keep the grad as replicate, this is because we don't
+    # want to convert the replicated gradients back to partial, although
+    # that's logically conform with the same layout, converting the gradients
+    # back to partial is actually useless as you would have to do reduce later
+    # which would be more expensive than keeping it replicate!
+
+    # for backward shard -> partial, we just do shard -> replicate
+    # for backward replicate -> partial, we skip the transformation
+    # NOTE: _is_shard_like covers _StridedShard defensively; currently
+    # unreachable because Partial -> _StridedShard is not implemented.
+    normalized_placements: list[Placement] = []
+    for current, target in zip(current_spec.placements, previous_spec.placements):
+        if (_is_shard_like(current) or current.is_replicate()) and target.is_partial():
+            normalized_placements.append(Replicate())
+        else:
+            normalized_placements.append(target)
+
+    previous_spec = DTensorSpec(
+        previous_spec.device_mesh,
+        placements=tuple(normalized_placements),
+        tensor_meta=previous_spec.tensor_meta,
+        use_strided_shard_as_shard_order=previous_spec.use_strided_shard_as_shard_order,
+    )
+
+    output = redistribute_local_tensor(
+        local_tensor,
+        current_spec,
+        previous_spec,
+        async_op=async_op,
+    )
+
+    if output.dtype != original_dtype:
+        output = output.to(original_dtype)
+
+    spec = DTensorSpec(
+        previous_spec.device_mesh,
+        tuple(normalized_placements),
+        tensor_meta=TensorMeta(
+            shape=grad_output.shape,
+            stride=grad_output.stride(),
+            dtype=output.dtype,
+        ),
+        use_strided_shard_as_shard_order=previous_spec.use_strided_shard_as_shard_order,
+    )
+    return output, spec
+
+
 class Redistribute(torch.autograd.Function):
     @staticmethod
     def forward(  # type: ignore[override]
@@ -1686,6 +1857,7 @@ def forward(  # type: ignore[override]
                     stride=input.stride(),
                     dtype=forward_dtype,
                 ),
+                use_strided_shard_as_shard_order=input._spec.use_strided_shard_as_shard_order,
             )
         else:
             local_tensor = input._local_tensor
@@ -1722,72 +1894,60 @@ def forward(  # type: ignore[override]
     @staticmethod
     def backward(ctx, grad_output: "dtensor.DTensor"):  # type: ignore[override]
         previous_spec = ctx.current_spec
-        async_op = ctx.async_op
-        backward_dtype = ctx.backward_dtype or ctx.original_dtype
+        output_dtensor = NestedRedistribute.apply(
+            grad_output,
+            previous_spec,
+            ctx.async_op,
+            ctx.backward_dtype,
+            ctx.original_dtype,
+        )
+        return (
+            output_dtensor,
+            None,
+            None,
+            None,
+            None,
+            None,
+        )
 
-        if backward_dtype != grad_output._local_tensor.dtype:
-            local_tensor = grad_output._local_tensor.to(dtype=backward_dtype)
-            current_spec = DTensorSpec(
-                mesh=grad_output._spec.device_mesh,
-                placements=grad_output._spec.placements,
-                tensor_meta=TensorMeta(
-                    shape=grad_output.shape,
-                    stride=grad_output.stride(),
-                    dtype=backward_dtype,
-                ),
-            )
-            previous_spec = DTensorSpec(
-                mesh=previous_spec.device_mesh,
-                placements=previous_spec.placements,
-                tensor_meta=current_spec.tensor_meta,
-            )
-        else:
-            local_tensor = grad_output._local_tensor
-            current_spec = grad_output._spec
-        # skip the replicate to partial transformation when we are in backward pass
-        # In this case we keep the grad as replicate, this is because we don't
-        # want to convert the replicated gradients back to partial, although
-        # that's logically conform with the same layout, converting the gradients
-        # back to partial is actually useless as you would have to do reduce later
-        # which would be more expensive than keeping it replicate!
-
-        # for backward shard -> partial, we just do shard -> replicate
-        # for backward replicate -> partial, we skip the transformation
-        normalized_placements: list[Placement] = []
-        for current, target in zip(current_spec.placements, previous_spec.placements):
-            if (current.is_shard() or current.is_replicate()) and target.is_partial():
-                normalized_placements.append(Replicate())
-            else:
-                normalized_placements.append(target)
 
-        previous_spec = DTensorSpec(
-            previous_spec.device_mesh,
-            placements=tuple(normalized_placements),
-            tensor_meta=previous_spec.tensor_meta,
-        )
+class NestedRedistribute(torch.autograd.Function):
+    """
+    This class is used to make the redistribution of a DTensor twice-differentiable.
+    This is called during the `Redistribute.forward`.
+    Therefore, `NestedRedistribute.forward` is called during the first backward pass,
+    and `NestedRedistribute.backward` is called during the second backward pass.
 
-        output = redistribute_local_tensor(
-            local_tensor,
-            current_spec,
+    Note: `NestedRedistribute.backward` is not differentiable, and therefore triple
+    backward is not yet supported.
+    """
+
+    @staticmethod
+    def forward(  # type: ignore[override]
+        # pyre-fixme[2]: Parameter must be annotated.
+        ctx,
+        grad_output: "dtensor.DTensor",
+        previous_spec: DTensorSpec,
+        async_op: bool = False,
+        forward_dtype: torch.dtype | None = None,
+        backward_dtype: torch.dtype | None = None,
+    ):
+        ctx.async_op = async_op
+        ctx.original_dtype = grad_output._local_tensor.dtype
+        ctx.backward_dtype = backward_dtype or ctx.original_dtype
+
+        output, spec = _redistribute_backward(
+            grad_output,
             previous_spec,
-            async_op=async_op,
-            is_explicit=True,
+            ctx.backward_dtype,
+            backward_dtype,
+            async_op,
         )
 
-        if output.dtype != ctx.original_dtype:
-            output = output.to(ctx.original_dtype)
+        ctx.current_spec = spec
 
-        spec = DTensorSpec(
-            previous_spec.device_mesh,
-            tuple(normalized_placements),
-            tensor_meta=TensorMeta(
-                shape=grad_output.shape,
-                stride=grad_output.stride(),
-                dtype=output.dtype,
-            ),
-        )
         # pyrefly: ignore [bad-argument-type]
-        output_dtensor = dtensor.DTensor(
+        return dtensor.DTensor(
             # pyrefly: ignore [bad-argument-count]
             output,
             spec,
@@ -1795,6 +1955,20 @@ def backward(ctx, grad_output: "dtensor.DTensor"):  # type: ignore[override]
             requires_grad=grad_output.requires_grad,
         )
 
+    @staticmethod
+    def backward(ctx, grad2_output: "dtensor.DTensor"):  # type: ignore[override]
+        previous_spec = ctx.current_spec
+        async_op = ctx.async_op
+        backward_dtype = ctx.backward_dtype or ctx.original_dtype
+
+        output_dtensor = NestedRedistribute.apply(
+            grad2_output,
+            previous_spec,
+            async_op,
+            backward_dtype,
+            ctx.original_dtype,
+        )
+
         return (
             output_dtensor,
             None,
diff --git a/torch/distributed/tensor/_sharding_prop.py b/torch/distributed/tensor/_sharding_prop.py
index 195404d4d6d72..463a0d4362fc0 100644
--- a/torch/distributed/tensor/_sharding_prop.py
+++ b/torch/distributed/tensor/_sharding_prop.py
@@ -2,16 +2,19 @@
 import logging
 import threading
 from collections.abc import Callable, Sequence
+from contextlib import nullcontext
 from functools import lru_cache
 from itertools import chain
-from typing import cast, Optional
+from typing import cast
 
 import torch
 from torch._guards import detect_fake_mode
+from torch._logging import LazyString
 from torch._ops import OpOverload
 from torch._subclasses import FakeTensorMode
 from torch.distributed._functional_collectives import _are_we_tracing
 from torch.distributed.device_mesh import DeviceMesh
+from torch.distributed.tensor._decompositions import DecompShardingStrategy
 from torch.distributed.tensor._dtensor_spec import DTensorSpec, TensorMeta
 from torch.distributed.tensor._op_schema import (
     OpInfo,
@@ -34,6 +37,7 @@
     try_find_mesh_from_args,
 )
 from torch.distributed.tensor.placement_types import _StridedShard, Shard
+from torch.utils._pytree import tree_map
 
 
 aten = torch.ops.aten
@@ -41,6 +45,56 @@
 log = logging.getLogger(__name__)
 
 
+def _propagate_use_strided_shard_flag(
+    op_strategy: OpStrategy,
+    op_schema: OpSchema,
+) -> None:
+    """Propagate use_strided_shard_as_shard_order from input specs to output specs.
+
+    When inputs carry _StridedShard with an explicit flag, all output (and input)
+    DTensorSpecs in the strategy that also contain _StridedShard must agree.
+    Strategy functions may forget to propagate the flag; this function fixes
+    them up centrally after the strategy is produced.
+    """
+    _use_strided: bool | None = None
+    for spec in op_schema.args_spec:
+        if any(isinstance(p, _StridedShard) for p in spec.placements):
+            val = spec.use_strided_shard_as_shard_order
+            if _use_strided is not None and _use_strided != val:
+                raise ValueError(
+                    "Conflicting use_strided_shard_as_shard_order across "
+                    f"input specs: got both {_use_strided} and {val}"
+                )
+            _use_strided = val
+
+    if _use_strided is None:
+        return
+
+    def _fixup(spec: DTensorSpec) -> None:
+        if not any(isinstance(p, _StridedShard) for p in spec.placements):
+            return
+        if spec.use_strided_shard_as_shard_order == _use_strided:
+            return
+        spec.use_strided_shard_as_shard_order = _use_strided
+        if _use_strided:
+            spec.shard_order = None  # pyrefly: ignore[bad-assignment]
+        else:
+            spec.shard_order = DTensorSpec.compute_default_shard_order(spec.placements)
+
+    for op_spec in op_strategy.strategies:
+        out = op_spec.output_specs
+        if out is not None:
+            if isinstance(out, DTensorSpec):
+                _fixup(out)
+            else:
+                for s in out:
+                    if s is not None:
+                        _fixup(s)
+        if op_spec.input_specs is not None:
+            for s in op_spec.input_specs:
+                _fixup(s)
+
+
 def _length(obj) -> int:
     if obj is None:
         return 0
@@ -49,14 +103,15 @@ def _length(obj) -> int:
     return len(obj)
 
 
-def _get_expected_num_tensor_outputs(op: OpOverload) -> int:
+def _get_expected_num_tensor_outputs(op: OpOverload) -> int | None:
     """
     Get the expected number of tensor outputs for an operator based on its schema.
 
     Returns:
         The number of tensor outputs expected. Returns 0 for ops that don't return tensors
         (e.g., _linalg_check_errors). Returns 1 for single tensor return, and >1 for
-        tuple returns where each element is a tensor.
+        tuple returns where each element is a tensor. Returns None for List[Tensor]
+        returns where the length is unknown at schema time.
     """
     return_types = op._schema.returns
     if len(return_types) == 0:
@@ -67,8 +122,8 @@ def _get_expected_num_tensor_outputs(op: OpOverload) -> int:
         # Could be single tensor or tuple of tensors
         return len(return_types)
     elif isinstance(first_return.type, torch.ListType):
-        # List[Tensor] - we don't know the length at schema time, treat as 1
-        return 1
+        # List[Tensor] - we don't know the length at schema time
+        return None
     else:
         # Not a tensor return type
         return 0
@@ -97,12 +152,23 @@ def _validate_tensor_meta_count(
     else:
         actual_outputs = len(tensor_meta)
 
-    assert actual_outputs == expected_outputs, (
-        f"Tensor meta count mismatch for {op_schema.op}: "
-        f"expected {expected_outputs} tensor output(s) based on op schema, "
-        f"but _propagate_tensor_meta returned {actual_outputs}. "
-        f"This usually indicates a bug in fake tensor propagation for this op."
-    )
+    if expected_outputs is None:
+        # List[Tensor] return type: length unknown at schema time, but
+        # tensor_meta must be a list of TensorMeta.
+        if not isinstance(tensor_meta, list):
+            raise AssertionError(
+                f"Tensor meta for {op_schema.op} should be a list[TensorMeta] "
+                f"(op returns List[Tensor]), but got {type(tensor_meta).__name__}"
+            )
+        return
+
+    if actual_outputs != expected_outputs:
+        raise AssertionError(
+            f"Tensor meta count mismatch for {op_schema.op}: "
+            f"expected {expected_outputs} tensor output(s) based on op schema, "
+            f"but _propagate_tensor_meta returned {actual_outputs}. "
+            f"This usually indicates a bug in fake tensor propagation for this op."
+        )
 
 
 class LocalLRUCache(threading.local):
@@ -135,6 +201,99 @@ def cache_clear(self):
         return self.cache.cache_clear()
 
 
+def _format_unbacked_hinting_log(
+    op_schema: OpSchema,
+    strategies: list[OpSpec],
+    strategy_index: int,
+    replacements: dict,
+) -> str:
+    """Format log message for unbacked hinting strategy selection (only called if debug logging enabled)."""
+    args_spec = tuple(str(spec) for spec in op_schema.args_schema)
+    strat = strategies[strategy_index]
+    if strat.input_specs is None:
+        placements_in = None
+    else:
+        placements_in = tuple(
+            spec.format_shard_order_str(spec.placements, spec.shard_order)
+            for spec in strat.input_specs
+        )
+    placements_out = tree_map(
+        lambda spec: spec.format_shard_order_str(spec.placements, spec.shard_order),
+        strat.output_specs,
+        is_leaf=lambda x: isinstance(x, DTensorSpec),
+    )
+    return (
+        f"Selected strategy {placements_in} -> {placements_out} "
+        f"for {op_schema.op} with input {args_spec}, using unbacked hints: {replacements}"
+    )
+
+
+def _select_min_redistribute_cost(
+    costs: list[torch.types.FloatLikeType],
+    strategies: list[OpSpec],
+    op_schema: OpSchema | None = None,
+) -> int:
+    """
+    Given a list of costs and corresponding op strategies, selects the minimum cost strategy, returning the index.
+    If unbacked symbols are involved, replaces them with known upper-bound values, falling back to hardcoded values.
+    """
+    from torch.fx.experimental.symbolic_shapes import (
+        free_unbacked_symbols,
+        is_concrete_float,
+    )
+    from torch.utils._sympy.interp import sympy_interp
+    from torch.utils._sympy.numbers import int_oo
+    from torch.utils._sympy.reference import PythonReferenceAnalysis
+
+    int_fallback = 8192
+    free_unbacked = list(set(chain(*[free_unbacked_symbols(cost) for cost in costs])))
+
+    # Easy path: no unbacked shapes involved, choose min cost strategy.
+    # Doing the hard path for backed could also make sense?
+    if all(is_concrete_float(c) for c in costs) or not free_unbacked:
+        return costs.index(min(costs))
+
+    # Figure out heuristic hints for unbacked shapes.
+    # If available, use shape upper bound. If not, fallback to some integer (inductor size-hinting style).
+    shape_env = next(iter(x for x in costs if not is_concrete_float(x))).node.shape_env  # type: ignore[arg-type]
+    replacements = {}
+    for sym in free_unbacked:
+        # TODO(laithsakka): unify with optimization_hint API
+        if (hint := shape_env.var_to_hint_override.get(sym)) is not None:
+            replacements[sym] = hint
+        elif (upper := shape_env.bound_sympy(sym).upper) is not int_oo:
+            replacements[sym] = upper
+        else:
+            replacements[sym] = int_fallback
+
+    # Use replacements for redistribute cost hints
+    proxy_costs = [
+        float(cost)
+        if is_concrete_float(cost)
+        else sympy_interp(
+            PythonReferenceAnalysis,
+            replacements,
+            cost.node.expr.xreplace(replacements),  # type: ignore[arg-type]
+        )
+        for cost in costs
+    ]
+    min_cost = min(proxy_costs)
+    strategy_index = proxy_costs.index(min_cost)
+
+    if op_schema:
+        log.debug(
+            "%s",
+            LazyString(
+                _format_unbacked_hinting_log,
+                op_schema,
+                strategies,
+                strategy_index,
+                replacements,
+            ),
+        )
+    return strategy_index
+
+
 def _select_min_cost_strategy(
     strategy: OpStrategy, op_schema: OpSchema | None = None
 ) -> OpSpec:
@@ -149,9 +308,8 @@ def _select_min_cost_strategy(
     negative_cost_index: int = -1
     zero_cost_index: int = -1
     for strategy_idx, op_spec in enumerate(strategy.strategies):
-        assert op_spec.redistribute_cost is not None, (
-            "must set redistribute cost each OpSpec!"
-        )
+        if op_spec.redistribute_cost is None:
+            raise AssertionError("must set redistribute cost each OpSpec!")
         redistribute_cost = sum(chain.from_iterable(op_spec.redistribute_cost))
         op_spec_costs.append(redistribute_cost)
 
@@ -196,13 +354,20 @@ def _select_min_cost_strategy(
         selected_strategy_index = zero_cost_index
     else:
         # default to choosing minimal redistribute cost
-        min_cost = min(op_spec_costs)
-        selected_strategy_index = op_spec_costs.index(min_cost)
+        selected_strategy_index = _select_min_redistribute_cost(
+            op_spec_costs, strategy.strategies, op_schema
+        )
 
     return strategy.strategies[selected_strategy_index]
 
 
 class ShardingPropagator:
+    # Lock to protect FakeTensorMode context during tensor meta propagation.
+    # By default this is a no-op (nullcontext) for performance. Multi-threaded
+    # tests should set this to threading.Lock() to prevent race conditions
+    # when multiple threads enter different FakeTensorMode contexts.
+    _fake_mode_lock = nullcontext()
+
     def __init__(self) -> None:
         self.op_to_rules: dict[OpOverload, Callable[[OpSchema], OutputSharding]] = {}
         self.op_strategy_funcs: dict[
@@ -222,6 +387,7 @@ def __init__(self) -> None:
         self.propagate_op_sharding = LocalLRUCache(
             self.propagate_op_sharding_non_cached
         )
+        self.decomp_strategy = DecompShardingStrategy(self)
         # op map to save indices of shape (and stride) args which may need to be
         # modified in sharding prop
         self.op_to_shape_and_stride_idx: dict[OpOverload, int | tuple[int, int]] = {
@@ -233,12 +399,23 @@ def __init__(self) -> None:
             aten.new_empty_strided.default: (1, 2),
             # view ops
             aten.expand.default: 1,
+            aten.expand_copy.default: 1,
             aten.reshape.default: 1,
             aten.view.default: 1,
+            aten.view_copy.default: 1,
             aten._unsafe_view.default: 1,
             aten.select_backward.default: 1,
             aten.slice_backward.default: 1,
         }
+        # squeeze ops that need dim arg rewritten to only globally-singleton dims
+        self.squeeze_op_to_dims_variant: dict[OpOverload, OpOverload] = {
+            aten.squeeze.default: aten.squeeze.dims,
+            aten.squeeze.dim: aten.squeeze.dims,
+            aten.squeeze.dims: aten.squeeze.dims,
+            aten.squeeze_.default: aten.squeeze_.dims,
+            aten.squeeze_.dim: aten.squeeze_.dims,
+            aten.squeeze_.dims: aten.squeeze_.dims,
+        }
 
     def register_sharding_prop_rule(
         self,
@@ -257,7 +434,7 @@ def register_single_dim_op_strategy(
         self,
         op_overload: OpOverload,
         strategy_info: _SingleDimStrategyInfo,
-        schema_info: Optional[RuntimeSchemaInfo] = None,
+        schema_info: RuntimeSchemaInfo | None = None,
     ):
         """
         Register a strategy over a single mesh-dim, relying on infra to automatically expand to the full mesh.
@@ -334,11 +511,15 @@ def _propagate_tensor_meta_non_cached(
 
         # NOTE: We must call the tracing in fake tensor mode so that it avoids
         # materializing memory.
-        fake_mode = detect_fake_mode() or FakeTensorMode()
-        with fake_mode:
-            fake_args = op_schema.gen_fake_args()
-            fake_kwargs = op_schema.gen_fake_kwargs()
-            fake_out = op_schema.op(*fake_args, **fake_kwargs)
+        # NOTE: Use _fake_mode_lock to serialize access when running in
+        # multi-threaded tests (lock must be set to threading.Lock()).
+        # This is a nullcontext by default.
+        with ShardingPropagator._fake_mode_lock:
+            fake_mode = detect_fake_mode() or FakeTensorMode()
+            with fake_mode:
+                fake_args = op_schema.gen_fake_args()
+                fake_kwargs = op_schema.gen_fake_kwargs()
+                fake_out = op_schema.op(*fake_args, **fake_kwargs)
 
         if isinstance(fake_out, torch.Tensor):
             return TensorMeta(
@@ -430,20 +611,11 @@ def _create_output_spec_with_new_tensor_meta(
                 if isinstance(spec, DTensorSpec):
                     output_tensor_meta_i = output_tensor_meta[i]
                     if not isinstance(output_tensor_meta_i, TensorMeta):
-                        # NOTE: aten.convolution_backward.default is an exception and it
-                        # needs extra handling because any Tensor in the output tuple
-                        # can be `None` depending on the output_mask parameter. This can
-                        # occur during double backpropagation or when certain gradients
-                        # are not needed (e.g., grad_input when input has requires_grad=False,
-                        # grad_weight/grad_bias when weight/bias have requires_grad=False,
-                        # or grad_bias when bias is None). We explicitly allow the
-                        # corresponding TensorMeta to be `None`.
-                        if (
-                            op == aten.convolution_backward.default
-                            and i in (0, 1, 2)
-                            and output_tensor_meta_i is None
-                        ):
-                            assert isinstance(output_specs, list)
+                        # Some ops (e.g. convolution_backward, native_layer_norm_backward,
+                        # _fused_rms_norm_backward) have an output_mask parameter that
+                        # controls which outputs are computed. When output_mask[i] is
+                        # False, the output at position i is None and has no TensorMeta.
+                        if output_tensor_meta_i is None:
                             new_specs.append(None)
                             continue
                         else:
@@ -460,7 +632,8 @@ def _create_output_spec_with_new_tensor_meta(
 
             return tuple(new_specs)
         else:
-            assert output_specs is None
+            if output_specs is not None:
+                raise AssertionError
             return output_specs
 
     def _wrap_with_op_strategy(self, op_schema: OpSchema) -> OpSchema:
@@ -473,17 +646,23 @@ def _wrap_with_op_strategy(self, op_schema: OpSchema) -> OpSchema:
         def spec_to_strategy(spec: object) -> object:
             if isinstance(spec, DTensorSpec):
                 return OpStrategy([OpSpec(spec)])
-            elif (
-                isinstance(spec, (list, tuple))
-                and len(spec) > 0
-                and isinstance(spec[0], DTensorSpec)
-            ):
-                # tensor list create tuple strategy
-                tuple_strategy = [spec_to_strategy(s) for s in spec]
-                tuple_strategy = cast(Sequence[StrategyType], tuple_strategy)
-                return TupleStrategy(
-                    tuple(tuple_strategy) if isinstance(spec, tuple) else tuple_strategy
-                )
+            elif isinstance(spec, (list, tuple)) and len(spec) > 0:
+                if all(isinstance(s, DTensorSpec) for s in spec):
+                    # tensor list create tuple strategy
+                    tuple_strategy = [spec_to_strategy(s) for s in spec]
+                    tuple_strategy = cast(Sequence[StrategyType], tuple_strategy)
+                    return TupleStrategy(
+                        tuple(tuple_strategy)
+                        if isinstance(spec, tuple)
+                        else tuple_strategy
+                    )
+                elif any(isinstance(s, DTensorSpec) for s in spec):
+                    # mixed list (e.g. [DTensorSpec, None, DTensorSpec]) for
+                    # ops like aten.index.Tensor; keep as list so pytree
+                    # flattening can extract OpStrategy items
+                    return [spec_to_strategy(s) for s in spec]
+                else:
+                    return spec
             else:
                 return spec
 
@@ -507,10 +686,11 @@ def propagate(self, op_info: OpInfo) -> None:
 
         # NOTE: schema should always be populated when calling this function,
         # as it's only called after unwrap_to_op_info (create_schema=True).
-        assert op_info.schema is not None, (
-            "op_info.schema should not be None in propagate. "
-            "This function should only be called after unwrap_to_op_info."
-        )
+        if op_info.schema is None:
+            raise AssertionError(
+                "op_info.schema should not be None in propagate. "
+                "This function should only be called after unwrap_to_op_info."
+            )
 
         # We cannot use an lru cache if we know that inputs will have dynamic shapes,
         # because SymInts are not hashable.
@@ -562,7 +742,8 @@ def propagate_op_sharding_non_cached(self, op_schema: OpSchema) -> OutputShardin
 
             if single_dim_strategy_info is not None:
                 mesh = try_find_mesh_from_args(op_schema.op, op_schema.args_schema)
-                assert isinstance(mesh, DeviceMesh), "Expected to find a valid mesh"
+                if not isinstance(mesh, DeviceMesh):
+                    raise AssertionError("Expected to find a valid mesh")
                 # expand to generate the full set of strategy combinations, each one
                 # with a redistribute cost, and then find the min strategy over those costs.
                 _expanded_strategy_fn = _expand_single_dim_strategy_to_mesh(
@@ -572,33 +753,27 @@ def propagate_op_sharding_non_cached(self, op_schema: OpSchema) -> OutputShardin
                     op_schema.op, strategy_schema.args_meta, strategy_schema.kwargs_meta
                 )
             else:
-                assert op_strategy_func is not None
+                if op_strategy_func is None:
+                    raise AssertionError
                 op_strategy = op_strategy_func(strategy_schema)
 
         else:
             # try operator decomposition path
-            from torch.distributed.tensor._decompositions import DecompShardingStrategy
 
-            # If the op has a CIA decomposition, we prioritize it over the decomposition flow,
-            # allowing decomposed ops to individually enter DTensor dispatch.
-            # TODO(pianpwk): maybe switch this back; the decomp flow could incur fewer comms.
             op_strategy = None
-            has_cia = torch._C._dispatch_has_kernel_for_dispatch_key(
-                op_schema.op.name(),
-                torch._C.DispatchKey.CompositeImplicitAutograd,
-            )
-            if not has_cia and DecompShardingStrategy.has_decomp(op_schema.op):
+            if DecompShardingStrategy.has_decomp(op_schema.op):
                 # Ensure schema_info is registered for proper cache key computation
-                DecompShardingStrategy.ensure_schema_info(op_schema.op, self)
+                self.decomp_strategy.ensure_schema_info(op_schema.op)
                 try:
-                    op_strategy = DecompShardingStrategy.propagate_strategy(
-                        op_schema, self
+                    op_strategy = self.decomp_strategy.propagate_strategy(
+                        op_schema,
                     )
                 except Exception as e:
                     decomp_exception = e
 
         if op_strategy is not None:
             if isinstance(op_strategy, OpStrategy):
+                _propagate_use_strided_shard_flag(op_strategy, op_schema)
                 # single Op strategy
                 output_strategy = _select_min_cost_strategy(op_strategy, op_schema)
 
@@ -612,7 +787,8 @@ def propagate_op_sharding_non_cached(self, op_schema: OpSchema) -> OutputShardin
                 # is a DTensorSpec, we use output_specs as the spec for each DTensor
                 # input arg.
                 if output_strategy.input_specs is None:
-                    assert isinstance(output_strategy.output_specs, DTensorSpec)
+                    if not isinstance(output_strategy.output_specs, DTensorSpec):
+                        raise AssertionError
 
                 for idx, input_spec in enumerate(op_schema.args_spec):
                     desired_spec = (
@@ -638,7 +814,8 @@ def propagate_op_sharding_non_cached(self, op_schema: OpSchema) -> OutputShardin
                 # shape and stride args need to be modified for
                 # view ops and new factory ops, potentially
                 if op_schema.op in self.op_to_shape_and_stride_idx:
-                    assert isinstance(output_strategy.output_spec, DTensorSpec)
+                    if not isinstance(output_strategy.output_spec, DTensorSpec):
+                        raise AssertionError
                     # It happens when the output has the same shape as the input
                     # and the input placements are not all Replicate().
                     if any(
@@ -646,13 +823,23 @@ def propagate_op_sharding_non_cached(self, op_schema: OpSchema) -> OutputShardin
                         for p in output_strategy.output_spec.placements
                     ):
                         schema = suggestion_schema or op_schema
-                        assert isinstance(out_tensor_meta, TensorMeta)
+                        if not isinstance(out_tensor_meta, TensorMeta):
+                            raise AssertionError
                         suggestion_schema = self._adjust_shape_and_stride_args(
                             out_tensor_meta, schema, output_strategy.output_spec
                         )
                         needs_redistribute = True
                         use_val_from_redistribute_schema = True
 
+                # rewrite squeeze to use only globally-singleton dims
+                if op_schema.op in self.squeeze_op_to_dims_variant:
+                    schema = suggestion_schema or op_schema
+                    adjusted = self._adjust_squeeze_to_global_singletons(schema)
+                    if adjusted is not None:
+                        suggestion_schema = adjusted
+                        needs_redistribute = True
+                        use_val_from_redistribute_schema = True
+
                 # construct output spec for the op
                 if op_schema.return_type_tuple_tensor_like():
                     # for ops that return multiple tensors and the output_specs is not
@@ -667,6 +854,7 @@ def propagate_op_sharding_non_cached(self, op_schema: OpSchema) -> OutputShardin
                                 mesh=output_specs.mesh,
                                 placements=output_specs.placements,
                                 tensor_meta=output_specs.tensor_meta,
+                                use_strided_shard_as_shard_order=output_specs.use_strided_shard_as_shard_order,
                             )
                             for _ in range(len(op_schema.op._schema.returns))
                         )
@@ -690,10 +878,13 @@ def propagate_op_sharding_non_cached(self, op_schema: OpSchema) -> OutputShardin
                 selected_strategies: list[OpSpec] = []
                 out_spec_list: list[DTensorSpec] = []
                 for strategy in op_strategy.children:
-                    assert isinstance(strategy, OpStrategy)
+                    if not isinstance(strategy, OpStrategy):
+                        raise AssertionError
+                    _propagate_use_strided_shard_flag(strategy, op_schema)
                     selected_strategy = _select_min_cost_strategy(strategy)
                     selected_strategies.append(selected_strategy)
-                    out_spec_list.append(selected_strategy.output_spec)
+                    if selected_strategy.output_specs is not None:
+                        out_spec_list.append(selected_strategy.output_spec)
 
                 needs_redistribute = False
                 suggestion_args: list[object] = []
@@ -827,14 +1018,59 @@ def _adjust_shape_and_stride_args(
         expected_input_schema = list(schema.args_schema)
         # adjust shape to be the same as that of the _local_tensor
         # of the DTensor input arg at index 0, which is inferred
-        expected_input_schema[shape_idx], _ = compute_local_shape_and_global_offset(
+        local_shape, _ = compute_local_shape_and_global_offset(
             out_tensor_meta.shape, spec.mesh, spec.placements, skip_offset=True
         )
+        expected_input_schema[shape_idx] = local_shape
 
         # adjust the stride arg for aten.new_empty_strided.default
         if stride_idx:
             expected_input_schema[stride_idx] = compute_local_stride(
-                out_tensor_meta.stride, spec.mesh, spec.placements
+                out_tensor_meta.stride, local_shape
             )
 
         return OpSchema(schema.op, tuple(expected_input_schema), schema.kwargs_schema)
+
+    def _adjust_squeeze_to_global_singletons(self, schema: OpSchema) -> OpSchema | None:
+        """
+        Rewrite squeeze ops to squeeze.dims with only globally-singleton dims.
+        Fixes bug where sharded dims with local size 1 get incorrectly squeezed.
+        Returns None if no rewrite is needed (already squeeze.dims with correct args).
+        """
+        from torch.fx.experimental.symbolic_shapes import guard_or_false
+
+        input_spec = cast(DTensorSpec, schema.args_schema[0])
+        tensor_meta = input_spec.tensor_meta
+        if tensor_meta is None:
+            raise RuntimeError("squeeze requires tensor metadata")
+        global_shape = tensor_meta.shape
+        ndim = len(global_shape)
+
+        def normalize(d: int) -> int:
+            return d if d >= 0 else d + ndim
+
+        def is_singleton(d: int) -> bool:
+            nd = normalize(d)
+            return 0 <= nd < ndim and guard_or_false(global_shape[nd] == 1)
+
+        # guard_or_false: conservatively keep dims when size is symbolic/unknown
+        if schema.op in (aten.squeeze.default, aten.squeeze_.default):
+            target_dims = tuple(
+                i for i, s in enumerate(global_shape) if guard_or_false(s == 1)
+            )
+        elif schema.op in (aten.squeeze.dim, aten.squeeze_.dim):
+            dim = normalize(schema.args_schema[1])  # type: ignore[arg-type]
+            target_dims = (dim,) if is_singleton(dim) else ()
+        else:
+            dims = cast(Sequence[int], schema.args_schema[1])
+            target_dims = tuple(  # type: ignore[union-attr]
+                normalize(d) for d in dims if is_singleton(d)
+            )
+
+        dims_variant = self.squeeze_op_to_dims_variant[schema.op]
+        # Skip rewrite if already targeting the right op with the same dims
+        if schema.op == dims_variant and len(schema.args_schema) > 1:
+            existing_dims = schema.args_schema[1]
+            if existing_dims == target_dims:
+                return None
+        return OpSchema(dims_variant, (input_spec, target_dims), {})
diff --git a/torch/distributed/tensor/_shards_wrapper.py b/torch/distributed/tensor/_shards_wrapper.py
index 1673dd7e34b99..2d71250dab72d 100644
--- a/torch/distributed/tensor/_shards_wrapper.py
+++ b/torch/distributed/tensor/_shards_wrapper.py
@@ -39,9 +39,10 @@ class LocalShardsWrapper(torch.Tensor):
     def __new__(
         cls, local_shards: list[torch.Tensor], local_offsets: list[tuple[int, ...]]
     ) -> "LocalShardsWrapper":
-        assert all(
+        if not all(
             tensor.device == local_shards[0].device for tensor in local_shards[1:]
-        )
+        ):
+            raise AssertionError
 
         # if empty shard, we create a empty tensor
         if len(local_shards) == 0:
@@ -146,10 +147,11 @@ def handle_view(args, kwargs) -> "LocalShardsWrapper":
         res_shards_list = []
         if len(args[0].local_shards()) > 1:
             if args[0].local_shards()[0].ndim == 2:
-                assert (
+                if not (
                     args[0].storage_metadata().size[0] == view_shape[0]
                     and args[0].storage_metadata().size[1] == view_shape[1]
-                )
+                ):
+                    raise AssertionError
                 # This accounts for a DTensor quirk, when multiple shards are present on a rank, DTensor on
                 # init calls view_as() on the global tensor shape
                 # will fail because the view shape is not applicable to individual shards.
@@ -158,7 +160,8 @@ def handle_view(args, kwargs) -> "LocalShardsWrapper":
                     for shard in args[0].local_shards()
                 ]
             elif args[0].local_shards()[0].ndim == 1:
-                assert args[0].storage_metadata().size[0] == view_shape[0]
+                if args[0].storage_metadata().size[0] != view_shape[0]:
+                    raise AssertionError
                 # This case is for optimizer sharding as regardless of sharding type, optimizer state is row wise sharded
                 res_shards_list = [
                     aten.view.default(shard, shard.shape, **kwargs)
diff --git a/torch/distributed/tensor/_tp_conv.py b/torch/distributed/tensor/_tp_conv.py
index 275cb07934b50..353bd2a0287a3 100644
--- a/torch/distributed/tensor/_tp_conv.py
+++ b/torch/distributed/tensor/_tp_conv.py
@@ -112,8 +112,10 @@ def tp_convolution(
     local_tensor_kwargs: dict[str, object],
     dim_map: list[int],
 ) -> object:
-    assert op_call == aten.convolution.default
-    assert len(local_tensor_args) == 9
+    if op_call != aten.convolution.default:
+        raise AssertionError
+    if len(local_tensor_args) != 9:
+        raise AssertionError
 
     rank = dist.get_rank()
     size = dist.get_world_size()
@@ -121,18 +123,23 @@ def tp_convolution(
     weight = cast(torch.Tensor, local_tensor_args[1])
     stride, padding, dilation = local_tensor_args[3:6]
 
-    assert _is_supported(in_tensor.shape, weight.shape, stride, padding, dilation)
-    assert isinstance(padding, list)
+    if not isinstance(padding, list):
+        raise AssertionError
 
     if not _requires_data_exchange(padding, dim_map):
         local_results = op_call(*local_tensor_args, **local_tensor_kwargs)
         return local_results
     else:
+        if not _is_supported(in_tensor.shape, weight.shape, stride, padding, dilation):
+            raise AssertionError(
+                "tp_convolution data exchange requires supported stride/padding/dilation"
+            )
         # step 0 compute the overlap pixels of the input tensor
         d = weight.shape[-1] - 1
         d1 = d // 2
         d2 = d - d1
-        assert d1 + d2 == d
+        if d1 + d2 != d:
+            raise AssertionError
         right = (rank + 1) % size
         left = (rank - 1 + size) % size
 
@@ -166,8 +173,10 @@ def tp_convolution_backward(
     local_tensor_kwargs: dict[str, object],
     dim_map: list[int],
 ) -> object:
-    assert op_call == aten.convolution_backward.default
-    assert len(local_tensor_args) == 11
+    if op_call != aten.convolution_backward.default:
+        raise AssertionError
+    if len(local_tensor_args) != 11:
+        raise AssertionError
 
     rank = dist.get_rank()
     size = dist.get_world_size()
@@ -176,18 +185,23 @@ def tp_convolution_backward(
     weight = cast(torch.Tensor, local_tensor_args[2])
     stride, padding, dilation = local_tensor_args[4:7]
 
-    assert _is_supported(in_tensor.shape, weight.shape, stride, padding, dilation)
-    assert isinstance(padding, list)
+    if not isinstance(padding, list):
+        raise AssertionError
 
     if not _requires_data_exchange(padding, dim_map):
         local_results = op_call(*local_tensor_args, **local_tensor_kwargs)
         return local_results
     else:
+        if not _is_supported(in_tensor.shape, weight.shape, stride, padding, dilation):
+            raise AssertionError(
+                "tp_convolution_backward data exchange requires supported stride/padding/dilation"
+            )
         # step 0 compute the overlap pixels of the input tensor
         d = weight.shape[3] - 1
         d1 = d // 2
         d2 = d - d1
-        assert d1 + d2 == d
+        if d1 + d2 != d:
+            raise AssertionError
         right = (rank + 1) % size
         left = (rank - 1 + size) % size
 
@@ -243,9 +257,11 @@ def convolution_handler(
     # sharding propagation
     dtensor.DTensor._op_dispatcher.sharding_propagator.propagate(op_info)
     output_sharding = op_info.output_sharding
-    assert output_sharding is not None, "output sharding should not be None"
+    if output_sharding is None:
+        raise AssertionError("output sharding should not be None")
     output_spec = output_sharding.output_spec
-    assert isinstance(output_spec, dtensor.DTensorSpec)
+    if not isinstance(output_spec, dtensor.DTensorSpec):
+        raise AssertionError
 
     # local propagation
     local_results = tp_convolution(
@@ -266,7 +282,10 @@ def convolution_backward_handler(
     # Redistribute grad_output tensor to the same placement as input tensor
     # pyrefly: ignore [bad-assignment]
     args = list(args)
-    assert isinstance(args[0], dtensor.DTensor) and isinstance(args[1], dtensor.DTensor)
+    if not (
+        isinstance(args[0], dtensor.DTensor) and isinstance(args[1], dtensor.DTensor)
+    ):
+        raise AssertionError
     # pyrefly: ignore [unsupported-operation]
     args[0] = args[0].redistribute(args[1].device_mesh, args[1].placements)
     args = tuple(args)
@@ -277,8 +296,10 @@ def convolution_backward_handler(
     # sharding propagation
     dtensor.DTensor._op_dispatcher.sharding_propagator.propagate(op_info)
     output_sharding = op_info.output_sharding
-    assert output_sharding is not None, "output sharding should not be None"
-    assert isinstance(op_info.flat_args_schema[0], dtensor.DTensorSpec)
+    if output_sharding is None:
+        raise AssertionError("output sharding should not be None")
+    if not isinstance(op_info.flat_args_schema[0], dtensor.DTensorSpec):
+        raise AssertionError
 
     # local propagation
     local_results = tp_convolution_backward(
diff --git a/torch/distributed/tensor/_utils.py b/torch/distributed/tensor/_utils.py
index fc53384e83cc6..89aede42e53f5 100644
--- a/torch/distributed/tensor/_utils.py
+++ b/torch/distributed/tensor/_utils.py
@@ -1,7 +1,7 @@
 import logging
 import threading
 from collections.abc import Callable, Sequence
-from typing import Any, cast
+from typing import Any
 
 import torch
 import torch.distributed._functional_collectives as funcol
@@ -15,6 +15,7 @@
 from torch.distributed.tensor._dtensor_spec import DTensorSpec
 from torch.distributed.tensor._op_schema import OpSchema
 from torch.distributed.tensor.placement_types import (
+    _is_shard_like,
     _StridedShard,
     Partial,
     Placement,
@@ -169,10 +170,12 @@ def _get_shard_size_and_offsets(
     if shard_size == 0:
         return shard_size, torch.arange(zero_global_offset, zero_global_offset + 1)
     if isinstance(placement, Shard) and not isinstance(placement, _StridedShard):
-        assert isinstance(shard_offsets, int)
+        if not isinstance(shard_offsets, int):
+            raise AssertionError
         index = torch.arange(shard_offsets, shard_offsets + shard_size)
     else:
-        assert isinstance(shard_offsets, list)
+        if not isinstance(shard_offsets, list):
+            raise AssertionError
         index = torch.tensor(shard_offsets)
     if previous_offsets is None:
         return shard_size, index
@@ -231,7 +234,8 @@ def _compute_local_shape_and_global_offset(
         def coordinate_lookup(dim: int) -> RankType:
             return _coord[dim]
     else:
-        assert my_coordinate is not None
+        if my_coordinate is None:
+            raise AssertionError
         coordinate_lookup = my_coordinate
 
     local_shape = list(global_shape)
@@ -253,9 +257,10 @@ def coordinate_lookup(dim: int) -> RankType:
             continue
         shard_dim = placement.dim
         zero_global_offset = global_shape[shard_dim]
-        assert shard_dim < len(local_shape), (
-            f"Sharding dim {shard_dim} greater than tensor ndim {len(local_shape)}"
-        )
+        if shard_dim >= len(local_shape):
+            raise AssertionError(
+                f"Sharding dim {shard_dim} greater than tensor ndim {len(local_shape)}"
+            )
         previous_offsets = shard_dim_to_global_offsets.get(shard_dim)
         shard_size, shard_offsets = _get_shard_size_and_offsets(
             local_shape[shard_dim],
@@ -310,23 +315,24 @@ def compute_local_tensor_info(
 
     for idx, placement in enumerate(placements):
         mesh_dim_size = mesh.size(idx)
-        if placement.is_shard():
-            shard_placement = cast(Shard, placement)
-            if shard_placement.dim < 0:
+        if _is_shard_like(placement):
+            if placement.dim < 0:
                 raise AssertionError(
                     "Shard placements should have negative dims normalized in "
-                    f"the user-facing APIs: {shard_placement}"
+                    f"the user-facing APIs: {placement}"
+                )
+            shard_dim = placement.dim
+            if shard_dim >= len(local_shape):
+                raise AssertionError(
+                    f"Sharding dim {shard_dim} greater than tensor ndim {len(local_shape)} "
+                    f"for placement number {idx}."
                 )
-            shard_dim = shard_placement.dim
-            assert shard_dim < len(local_shape), (
-                f"Sharding dim {shard_dim} greater than tensor ndim {len(local_shape)} "
-                f"for placement number {idx}."
-            )
 
             global_dim_size = local_shape[shard_dim]
-            assert global_dim_size % mesh_dim_size == 0, (
-                f"Global dim {global_dim_size} not divisible by mesh size {mesh_dim_size}"
-            )
+            if global_dim_size % mesh_dim_size != 0:
+                raise AssertionError(
+                    f"Global dim {global_dim_size} not divisible by mesh size {mesh_dim_size}"
+                )
             local_shape[shard_dim] = global_dim_size // mesh_dim_size
 
             # shrink strides that were scaled up globally
@@ -379,6 +385,7 @@ def compute_global_tensor_shape(
 
     if isinstance(placements[0], Replicate):
         return shape
+    # NOTE: isinstance(_, Shard) does not match _StridedShard; see _is_shard_like().
     elif isinstance(placements[0], Shard):
 
         @maybe_run_for_local_tensor
@@ -440,25 +447,26 @@ def try_find_mesh_from_args(
 
 
 def compute_local_stride(
-    global_stride: ShapeType, mesh: DeviceMesh, placements: Sequence[Placement]
+    global_stride: ShapeType, local_shape: ShapeType
 ) -> tuple[int, ...]:
     """
-    Compute the stride of a local tensor shard, given the global stride of the DTensor.
-    NOTE: Currently this function is assuming the DTensor is evenly shardable.
+    Compute the stride of a local tensor shard, given the global stride and local shape.
+
+    Derives strides by preserving the memory layout (dimension ordering) implied
+    by the global strides, then computing contiguous strides for the local shape
+    in that order.  Assumes the global tensor is non-overlapping and dense.
     """
-    stride_divisors = [1] * len(global_stride)
-    for mesh_idx, p in enumerate(placements):
-        if p.is_shard():
-            i = cast(Shard, p).dim
-            # tensor dimension i is sharded on mesh dimension mesh_idx,
-            # so we need to divide all the strides larger than stride[i]
-            # (by the submesh size)
-            for j in range(len(global_stride)):
-                if global_stride[j] > global_stride[i]:
-                    stride_divisors[j] *= mesh.size(mesh_idx)
-    return tuple(
-        global_stride[i] // stride_divisors[i] for i in range(len(global_stride))
-    )
+    ndim = len(global_stride)
+    # Sort dims by global stride descending to recover memory layout order.
+    # Stable sort preserves original dim order for ties, which only occur
+    # on size-1 dims where the stride value is semantically irrelevant.
+    perm = sorted(range(ndim), key=lambda d: global_stride[d], reverse=True)
+    local_strides = [0] * ndim
+    s = 1
+    for d in reversed(perm):
+        local_strides[d] = s
+        s *= local_shape[d]
+    return tuple(local_strides)
 
 
 def normalize_to_torch_size(size) -> torch.Size:  # type: ignore[no-untyped-def]
diff --git a/torch/distributed/tensor/debug/__init__.py b/torch/distributed/tensor/debug/__init__.py
index a367ce0370698..df99f7d750034 100644
--- a/torch/distributed/tensor/debug/__init__.py
+++ b/torch/distributed/tensor/debug/__init__.py
@@ -26,7 +26,6 @@ def _get_fast_path_sharding_prop_cache_stats():
     Get a tuple (hits, misses) for the fast path sharding propagation cache, used for debugging
     only.
     """
-    # pyrefly: ignore [missing-attribute]
     return torch._C._get_DTensor_sharding_propagator_cache_stats()
 
 
@@ -45,7 +44,6 @@ def _clear_fast_path_sharding_prop_cache():
     """
     Clears the cache for the fast path sharding propagation cache, used for debugging purpose only.
     """
-    # pyrefly: ignore [missing-attribute]
     torch._C._clear_DTensor_sharding_propagator_cache()
 
 
@@ -58,6 +56,15 @@ def _clear_sharding_prop_cache():
     _clear_fast_path_sharding_prop_cache()
 
 
+def _reinit_dispatch_logger():
+    """
+    Resets the cached DTensor dispatch logger state so that the next DTensor
+    dispatch re-checks whether debug logging is enabled. Call this after
+    changing the log level on the ``torch.distributed.tensor._dispatch`` logger.
+    """
+    torch._C._reinit_DTensor_dispatch_logger()
+
+
 # Set namespace for exposed private names
 CommDebugMode.__module__ = "torch.distributed.tensor.debug"
 visualize_sharding.__module__ = "torch.distributed.tensor.debug"
diff --git a/torch/distributed/tensor/debug/_comm_mode.py b/torch/distributed/tensor/debug/_comm_mode.py
index 66ec0bfff5f9c..d20b97867bac6 100644
--- a/torch/distributed/tensor/debug/_comm_mode.py
+++ b/torch/distributed/tensor/debug/_comm_mode.py
@@ -241,6 +241,7 @@ class CommDebugMode(TorchDispatchMode):
 
     def __init__(self):
         super().__init__()
+        self.supports_higher_order_operators = True
         self.comm_counts: dict[Any, int] = defaultdict(int)
         self.comm_module_counts = {}
         self.comm_module_operation_counts = {}
@@ -644,6 +645,12 @@ def __torch_dispatch__(self, func, types, args=(), kwargs=None):
         # Returning NotImplemented here gives us a chance to let DTensor
         # run and desugar into comms ops, before CommDebugMode sees them.
 
+        # Higher-order operators (e.g. run_dtensor_rng_op) don't have
+        # _overloadpacket and aren't collectives — just redispatch.
+        if isinstance(func, torch._ops.HigherOrderOperator):
+            kwargs = kwargs if kwargs else {}
+            return func(*args, **kwargs)
+
         # sets up operation-level collective count
         if self.advanced_module_tracker.name not in self.comm_module_operation_counts:
             # dictionary should hold module input and output shape, operations list and collective counter
diff --git a/torch/distributed/tensor/debug/_op_coverage.py b/torch/distributed/tensor/debug/_op_coverage.py
index 7315d64d697a8..c273bf8acd6ef 100644
--- a/torch/distributed/tensor/debug/_op_coverage.py
+++ b/torch/distributed/tensor/debug/_op_coverage.py
@@ -42,7 +42,8 @@ def get_inductor_decomp_graphs(model: nn.Module, args, kwargs):
     output.backward()
 
     # one fwd, one bwd graph
-    assert len(graphs) == 2
+    if len(graphs) != 2:
+        raise AssertionError
     return graphs
 
 
@@ -90,7 +91,6 @@ def print_op_coverage_summary(model: nn.Module, args, kwargs, *, output_csv=Fals
     op_infos.sort(key=itemgetter(count_idx), reverse=True)
 
     headers = ["Operator", "Schema", "Total Count", "Supported"]
-    # pyrefly: ignore [bad-argument-type]
     print(tabulate(op_infos, headers=headers))
 
     if output_csv:
@@ -102,5 +102,4 @@ def print_op_coverage_summary(model: nn.Module, args, kwargs, *, output_csv=Fals
             csv_writer.writerow(headers)
             # Write each table row to the CSV file
             for row in op_infos:
-                # pyrefly: ignore [bad-argument-type]
                 csv_writer.writerow(row)
diff --git a/torch/distributed/tensor/examples/comm_mode_features_example.py b/torch/distributed/tensor/examples/comm_mode_features_example.py
index f3bdf4cc1bd46..21480707c4290 100644
--- a/torch/distributed/tensor/examples/comm_mode_features_example.py
+++ b/torch/distributed/tensor/examples/comm_mode_features_example.py
@@ -692,7 +692,8 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
                         )
                     else:
                         x = block(x)
-                    assert x is not None
+                    if x is None:
+                        raise AssertionError
                     x = torch.nn.functional.relu(x)
                 return x
 
@@ -736,7 +737,8 @@ def run_example(world_size: int, rank: int, example_name: str) -> None:
     # this script is launched via torchrun which automatically manages ProcessGroup
     rank = int(os.environ["RANK"])
     world_size = int(os.environ["WORLD_SIZE"])
-    assert world_size == 4  # our example uses 4 worker ranks
+    if world_size != 4:
+        raise AssertionError  # our example uses 4 worker ranks
 
     parser = argparse.ArgumentParser(
         description="comm_mode_feature examples",
diff --git a/torch/distributed/tensor/examples/flex_attention_cp.py b/torch/distributed/tensor/examples/flex_attention_cp.py
index 8b309a6d2646e..4402ea6944042 100644
--- a/torch/distributed/tensor/examples/flex_attention_cp.py
+++ b/torch/distributed/tensor/examples/flex_attention_cp.py
@@ -39,7 +39,8 @@ def create_block_mask_cached(
 def flex_attn_example(world_size: int, rank: int) -> None:
     device_type = get_device_type()
     device_handle = getattr(torch, device_type, None)
-    assert device_handle is not None, f"Unsupported device type: {device_type}"
+    if device_handle is None:
+        raise AssertionError(f"Unsupported device type: {device_type}")
     num_devices_per_host = device_handle.device_count()
     device_handle.set_device(rank % num_devices_per_host)
     torch._dynamo.config.cache_size_limit = 1000
@@ -96,7 +97,8 @@ def causal_mask(b: int, h: int, q_idx: int, kv_idx: int) -> bool:
 
     q, k, v = qkv
     out = compiled_flex_attention(q, k, v, score_mod=None, block_mask=block_mask)
-    assert isinstance(out, torch.Tensor)
+    if not isinstance(out, torch.Tensor):
+        raise AssertionError
     expect_out = F.scaled_dot_product_attention(q, k, v, is_causal=True)
     torch.testing.assert_close(out, expect_out, atol=1e-1, rtol=1e-2)
 
@@ -137,7 +139,8 @@ def rewrite_mask_mod_for_cp(
         score_mod=None,
         block_mask=cp_block_mask,
     )
-    assert isinstance(cp_out, torch.Tensor)
+    if not isinstance(cp_out, torch.Tensor):
+        raise AssertionError
 
     # wrap the local output into a DTensor
     cp_out_dist = DTensor.from_local(cp_out, device_mesh, [Shard(seq_dim)])
@@ -171,7 +174,8 @@ def rewrite_mask_mod_for_cp(
     cp_out.backward(grad_out_dist.to_local())
 
     for cp_flex_grad_dist, expect_grad in zip([t.grad for t in qkv_dist], grad2):
-        assert isinstance(cp_flex_grad_dist, DTensor)
+        if not isinstance(cp_flex_grad_dist, DTensor):
+            raise AssertionError
         torch.testing.assert_close(
             cp_flex_grad_dist.full_tensor(), expect_grad, atol=1e-1, rtol=1e-2
         )
diff --git a/torch/distributed/tensor/examples/torchrec_sharding_example.py b/torch/distributed/tensor/examples/torchrec_sharding_example.py
index 492329f3d9007..0e6cd3698c23b 100644
--- a/torch/distributed/tensor/examples/torchrec_sharding_example.py
+++ b/torch/distributed/tensor/examples/torchrec_sharding_example.py
@@ -51,9 +51,12 @@ class LocalShardsWrapper(torch.Tensor):
     def __new__(
         cls, local_shards: list[torch.Tensor], offsets: list[torch.Size]
     ) -> "LocalShardsWrapper":
-        assert len(local_shards) > 0
-        assert len(local_shards) == len(offsets)
-        assert local_shards[0].ndim == 2
+        if len(local_shards) <= 0:
+            raise AssertionError
+        if len(local_shards) != len(offsets):
+            raise AssertionError
+        if local_shards[0].ndim != 2:
+            raise AssertionError
         # we calculate the total tensor size by "concat" on second tensor dimension
         cat_tensor_shape = list(local_shards[0].shape)
         if len(local_shards) > 1:  # column-wise sharding
@@ -91,10 +94,10 @@ def __torch_dispatch__(cls, func, types, args=(), kwargs=None):  # type: ignore[
         if func in supported_ops:
             res_shards_list = [
                 func(shard, *args[1:], **kwargs)
-                # pyrefly: ignore [bad-index, index-error]
+                # pyrefly: ignore [bad-index]
                 for shard in args[0].shards
             ]
-            # pyrefly: ignore [bad-index, index-error]
+            # pyrefly: ignore [bad-index]
             return LocalShardsWrapper(res_shards_list, args[0].shard_offsets)
         else:
             raise NotImplementedError(
@@ -182,11 +185,15 @@ def run_torchrec_row_wise_even_sharding_example(rank, world_size):
 
     # transform DTensor into LocalShardsWrapper
     dtensor_local_shards = dtensor.to_local()
-    assert isinstance(dtensor_local_shards, LocalShardsWrapper)
+    if not isinstance(dtensor_local_shards, LocalShardsWrapper):
+        raise AssertionError
     shard_tensor = dtensor_local_shards.shards[0]
-    assert torch.equal(shard_tensor, local_tensor)
-    assert dtensor_local_shards.shard_sizes[0] == local_shard_shape  # unwrap shape
-    assert dtensor_local_shards.shard_offsets[0] == local_shard_offset  # unwrap offset
+    if not torch.equal(shard_tensor, local_tensor):
+        raise AssertionError
+    if dtensor_local_shards.shard_sizes[0] != local_shard_shape:  # unwrap shape
+        raise AssertionError
+    if dtensor_local_shards.shard_offsets[0] != local_shard_offset:  # unwrap offset
+        raise AssertionError
 
 
 def run_torchrec_row_wise_uneven_sharding_example(rank, world_size):
@@ -250,18 +257,24 @@ def run_torchrec_row_wise_uneven_sharding_example(rank, world_size):
     # because it relies on offset computation which assumes even sharding.
     visualize_sharding(dtensor, header="Row-wise uneven sharding example in DTensor")
     # check the dtensor has the correct shape and stride on all ranks
-    assert dtensor.shape == emb_table_shape
-    assert dtensor.stride() == (embedding_dim, 1)
+    if dtensor.shape != emb_table_shape:
+        raise AssertionError
+    if dtensor.stride() != (embedding_dim, 1):
+        raise AssertionError
 
     ###########################################################################
     # example 2: transform DTensor into local_shards
     # note: DTensor.to_local() always returns a LocalShardsWrapper
     dtensor_local_shards = dtensor.to_local()
-    assert isinstance(dtensor_local_shards, LocalShardsWrapper)
+    if not isinstance(dtensor_local_shards, LocalShardsWrapper):
+        raise AssertionError
     shard_tensor = dtensor_local_shards.shards[0]
-    assert torch.equal(shard_tensor, local_tensor)
-    assert dtensor_local_shards.shard_sizes[0] == local_shard_shape  # unwrap shape
-    assert dtensor_local_shards.shard_offsets[0] == local_shard_offset  # unwrap offset
+    if not torch.equal(shard_tensor, local_tensor):
+        raise AssertionError
+    if dtensor_local_shards.shard_sizes[0] != local_shard_shape:  # unwrap shape
+        raise AssertionError
+    if dtensor_local_shards.shard_offsets[0] != local_shard_offset:  # unwrap offset
+        raise AssertionError
 
 
 def run_torchrec_table_wise_sharding_example(rank, world_size):
@@ -348,8 +361,10 @@ def run_torchrec_table_wise_sharding_example(rank, world_size):
             header=f"Table-wise sharding example in DTensor for Table {table_id}",
         )
         # check the dtensor has the correct shape and stride on all ranks
-        assert dtensor.shape == emb_table_shape
-        assert dtensor.stride() == (embedding_dim, 1)
+        if dtensor.shape != emb_table_shape:
+            raise AssertionError
+        if dtensor.stride() != (embedding_dim, 1):
+            raise AssertionError
 
     ###########################################################################
     # example 2: transform DTensor into torch.Tensor
@@ -358,17 +373,20 @@ def run_torchrec_table_wise_sharding_example(rank, world_size):
         # no matter what was passed to DTensor._local_tensor.
         dtensor_local_shards = table_to_dtensor[table_id].to_local()
         if rank == table_id:
-            assert isinstance(dtensor_local_shards, LocalShardsWrapper)
+            if not isinstance(dtensor_local_shards, LocalShardsWrapper):
+                raise AssertionError
             shard_tensor = dtensor_local_shards.shards[0]
-            assert torch.equal(shard_tensor, local_tensor)  # unwrap tensor
-            assert (
-                dtensor_local_shards.shard_sizes[0] == emb_table_shape
-            )  # unwrap shape
-            assert dtensor_local_shards.shard_offsets[0] == torch.Size(
+            if not torch.equal(shard_tensor, local_tensor):  # unwrap tensor
+                raise AssertionError
+            if dtensor_local_shards.shard_sizes[0] != emb_table_shape:  # unwrap shape
+                raise AssertionError
+            if dtensor_local_shards.shard_offsets[0] != torch.Size(
                 (0, 0)
-            )  # unwrap offset
+            ):  # unwrap offset
+                raise AssertionError
         else:
-            assert dtensor_local_shards.numel() == 0
+            if dtensor_local_shards.numel() != 0:
+                raise AssertionError
 
 
 def run_example(rank, world_size, example_name):
@@ -396,7 +414,8 @@ def run_example(rank, world_size, example_name):
     # this script is launched via torchrun which automatically manages ProcessGroup
     rank = int(os.environ["RANK"])
     world_size = int(os.environ["WORLD_SIZE"])
-    assert world_size == 4  # our example uses 4 worker ranks
+    if world_size != 4:  # our example uses 4 worker ranks
+        raise AssertionError
     # parse the arguments
     parser = argparse.ArgumentParser(
         description="torchrec sharding examples",
diff --git a/torch/distributed/tensor/examples/visualize_sharding_example.py b/torch/distributed/tensor/examples/visualize_sharding_example.py
index 7c0ab3adfffae..e879b2bd3e457 100644
--- a/torch/distributed/tensor/examples/visualize_sharding_example.py
+++ b/torch/distributed/tensor/examples/visualize_sharding_example.py
@@ -14,7 +14,8 @@
 import torch.distributed.tensor.debug
 
 
-assert int(os.getenv("WORLD_SIZE", "1")) >= 4, "We need at least 4 devices"
+if int(os.getenv("WORLD_SIZE", "1")) < 4:
+    raise AssertionError("We need at least 4 devices")
 rank = int(os.environ["RANK"])
 
 
diff --git a/torch/distributed/tensor/experimental/_context_parallel/_attention.py b/torch/distributed/tensor/experimental/_context_parallel/_attention.py
index 88ed34d274418..5edf94eaa43f4 100644
--- a/torch/distributed/tensor/experimental/_context_parallel/_attention.py
+++ b/torch/distributed/tensor/experimental/_context_parallel/_attention.py
@@ -127,7 +127,8 @@ def _partial_update(
     The result is a tensor that is the same size as ``original``.
     """
     chunks = list(original.chunk(n_chunks, dim=dim))
-    assert chunks[idx].shape == new.shape, (original.shape, new.shape, idx)
+    if chunks[idx].shape != new.shape:
+        raise AssertionError((original.shape, new.shape, idx))
     if add:
         chunks[idx] += new
     else:
@@ -156,15 +157,18 @@ def _merge_one(
         if len(block_lse.shape) < len(block_out.shape):
             block_lse = block_lse.unsqueeze(dim=-1)
             self._should_lse_squeeze = True
-        assert len(block_lse.shape) == len(block_out.shape)
+        if len(block_lse.shape) != len(block_out.shape):
+            raise AssertionError
 
         if self._lse is None:
             self._lse = block_lse
             self._out = block_out
         else:
             ROUND_ROBIN_CYCLE = 2
-            assert self._lse is not None
-            assert self._out is not None
+            if self._lse is None:
+                raise AssertionError
+            if self._out is None:
+                raise AssertionError
             lse = (
                 self._lse.chunk(ROUND_ROBIN_CYCLE, dim=self._seq_dim)[1]
                 if partial
@@ -213,8 +217,10 @@ def step(self, out: torch.Tensor, lse: torch.Tensor, partial: bool) -> None:
         self._merge_one(out, lse, partial)
 
     def results(self) -> tuple[torch.Tensor, torch.Tensor]:
-        assert self._out is not None
-        assert self._lse is not None
+        if self._out is None:
+            raise AssertionError
+        if self._lse is None:
+            raise AssertionError
         out = self._out.to(self._out_dtype)
         if self._should_lse_squeeze:
             lse = self._lse.squeeze(-1).to(self._lse_dtype)
@@ -259,7 +265,8 @@ def exchange_buffers(self, curr_buffer: torch.Tensor) -> None:
         self._buffer = ft_c.permute_tensor(curr_buffer, dsts, self._pg)
 
     def next_buffer(self) -> torch.Tensor:
-        assert self._buffer is not None
+        if self._buffer is None:
+            raise AssertionError
         return _maybe_wait(self._buffer)
 
 
@@ -287,7 +294,8 @@ def next_buffer(self) -> torch.Tensor:
         rank = dist.get_rank(self._pg)
         idx = rank - self._idx
 
-        assert self._aggregated_buffer is not None
+        if self._aggregated_buffer is None:
+            raise AssertionError
         self._aggregated_buffer = _maybe_wait(self._aggregated_buffer)
         return self._aggregated_buffer.chunk(dist.get_world_size(self._pg))[idx]
 
@@ -397,9 +405,8 @@ def _templated_ring_attention(
     if not is_causal and _cp_options.enable_load_balance:
         raise RuntimeError("Load balancing requires `is_causal=True`.")
 
-    assert isinstance(group, dist.ProcessGroup), (
-        "process group must be single dimension"
-    )
+    if not isinstance(group, dist.ProcessGroup):
+        raise AssertionError("process group must be single dimension")
     rank = dist.get_rank(group)
     size = dist.get_world_size(group)
 
@@ -631,8 +638,10 @@ def _templated_ring_attention_backward(
                 add=True,
             )
 
-    assert grad_key_ is not None
-    assert grad_value_ is not None
+    if grad_key_ is None:
+        raise AssertionError
+    if grad_value_ is None:
+        raise AssertionError
     grad_query = grad_query.to(query.dtype)
     next_grad_kv = dkv_rotater.next_buffer().to(key.dtype)
     grad_key = next_grad_kv[: grad_key.numel()].reshape(grad_key.shape)
@@ -900,8 +909,10 @@ def _sdpa_handler(
     # propagate.
     DTensor._op_dispatcher.sharding_propagator.propagate(op_info)
     output_sharding = op_info.output_sharding
-    assert output_sharding is not None, "output sharding should not be None"
-    assert not output_sharding.needs_redistribute, "inputs need to be redistributed"
+    if output_sharding is None:
+        raise AssertionError("output sharding should not be None")
+    if output_sharding.needs_redistribute:
+        raise AssertionError("inputs need to be redistributed")
 
     call_maps: dict[torch._ops.OpOverload, Callable] = {
         aten._scaled_dot_product_flash_attention.default: _scaled_dot_product_ring_flash_attention,
@@ -1080,10 +1091,11 @@ def _context_parallel_buffers(
     # generate the index tensor for rearranging the buffer if a load-balance
     # is available
     load_balance_indices = load_balancer._generate_indices() if load_balancer else None
-    assert load_balance_indices is None or load_balance_indices.ndim == 2, (
-        "load balance index expects shape (1, seq_len) or (B, seq_len) "
-        f"but got {load_balance_indices.shape}."
-    )
+    if not (load_balance_indices is None or load_balance_indices.ndim == 2):
+        raise AssertionError(
+            "load balance index expects shape (1, seq_len) or (B, seq_len) "
+            f"but got {load_balance_indices.shape}."
+        )
 
     new_buffers = []
     sharded_buffer: torch.Tensor | BlockMask
@@ -1229,10 +1241,11 @@ def _rewrite_mask_mod(
         local_q_size: int,
         qkv_rearrange_indices: torch.Tensor | None = None,
     ) -> _mask_mod_signature:
-        assert qkv_rearrange_indices is None or qkv_rearrange_indices.ndim == 2, (
-            "load balance index expects shape (1, seq_len) or (B, seq_len) "
-            f"but got {qkv_rearrange_indices.shape}."
-        )
+        if not (qkv_rearrange_indices is None or qkv_rearrange_indices.ndim == 2):
+            raise AssertionError(
+                "load balance index expects shape (1, seq_len) or (B, seq_len) "
+                f"but got {qkv_rearrange_indices.shape}."
+            )
 
         def qkv_idx_restore(
             b: torch.Tensor, idx_post_rearrange: torch.Tensor
@@ -1343,9 +1356,12 @@ def flex_input_fn(
                 args_list.append(kwargs.pop(name, None))
 
         query, key, value = args_list[: len(expected_arg_names)]
-        assert isinstance(query, torch.Tensor)
-        assert isinstance(key, torch.Tensor)
-        assert isinstance(value, torch.Tensor)
+        if not isinstance(query, torch.Tensor):
+            raise AssertionError
+        if not isinstance(key, torch.Tensor):
+            raise AssertionError
+        if not isinstance(value, torch.Tensor):
+            raise AssertionError
 
         key = key.contiguous()
         value = value.contiguous()
@@ -1374,7 +1390,8 @@ def sdpa_input_fn(
         for arg in itertools.chain(args, kwargs.values()):
             if isinstance(arg, torch.Tensor):
                 if isinstance(arg, DTensor):
-                    assert arg._spec.placements == placement
+                    if arg._spec.placements != placement:
+                        raise AssertionError
                 else:
                     arg = DTensor.from_local(arg, mesh, placement, run_check=False)
 
@@ -1463,11 +1480,11 @@ def _context_parallel_shard(
         device = flat_buffers[0].kv_num_blocks.device
     for buffer in flat_buffers:
         if isinstance(buffer, torch.Tensor):
-            assert device == buffer.device, "All buffers must be on the same device"
+            if device != buffer.device:
+                raise AssertionError("All buffers must be on the same device")
         else:
-            assert device == buffer.kv_num_blocks.device, (
-                "All buffers must be on the same device"
-            )
+            if device != buffer.kv_num_blocks.device:
+                raise AssertionError("All buffers must be on the same device")
 
     flat_sharded_buffers = _context_parallel_buffers(
         mesh, flat_buffers, flat_seq_dims, load_balancer
@@ -1568,7 +1585,8 @@ def context_parallel(
         load_balancer,
     )
     for buffer, shard in zip(buffers, shards):
-        assert isinstance(shard, torch.Tensor), "ContextParallel only supports Tensor"
+        if not isinstance(shard, torch.Tensor):
+            raise AssertionError("ContextParallel only supports Tensor")
         shard = shard.clone()
         buffer.resize_(shard.shape)
         buffer.copy_(shard)
@@ -1630,10 +1648,11 @@ def context_parallel_unshard(
         load_balancer._generate_indices(restore=True) if load_balancer else None
     )
 
-    assert restore_indices is None or restore_indices.ndim == 2, (
-        "load balance restore index expects shape (1, seq_len) or (B, seq_len) "
-        f"but got {restore_indices.shape}."
-    )
+    if not (restore_indices is None or restore_indices.ndim == 2):
+        raise AssertionError(
+            "load balance restore index expects shape (1, seq_len) or (B, seq_len) "
+            f"but got {restore_indices.shape}."
+        )
     unsharded_buffers = []
     for b, dim in zip(buffers, seq_dims):
         b = b.contiguous()
diff --git a/torch/distributed/tensor/experimental/_context_parallel/_load_balancer.py b/torch/distributed/tensor/experimental/_context_parallel/_load_balancer.py
index a945b840d38b3..f5d982d5a6561 100644
--- a/torch/distributed/tensor/experimental/_context_parallel/_load_balancer.py
+++ b/torch/distributed/tensor/experimental/_context_parallel/_load_balancer.py
@@ -145,29 +145,19 @@ def _generate_indices(self, restore: bool = False) -> Tensor:
         """
         seq_length = self.seq_length
         world_size = self.world_size
-        assert seq_length % (world_size * 2) == 0
+        if seq_length % (world_size * 2) != 0:
+            raise AssertionError
         chunk_size = seq_length // (world_size * 2)
-        all_indices = []
 
-        for rank in range(world_size):
-            # Generate indices for first chunk of the cp rank
-            first_chunk_start = rank * chunk_size
-            first_chunk_indices = list(
-                range(first_chunk_start, first_chunk_start + chunk_size)
-            )
-
-            # Second chunk: positions from the complementary chunk
-            second_chunk_idx = world_size * 2 - rank - 1
-            second_chunk_start = second_chunk_idx * chunk_size
-            second_chunk_indices = list(
-                range(second_chunk_start, second_chunk_start + chunk_size)
-            )
-            # combine the indices for this rank
-            all_indices.extend(first_chunk_indices + second_chunk_indices)
+        # Split sequence into 2*world_size chunks, then pair chunk r with
+        # chunk (2*world_size - 1 - r) for each rank.
+        indices = torch.arange(seq_length, dtype=torch.int, device=self.device)
+        chunks = indices.view(world_size * 2, chunk_size)
+        head_idx = torch.arange(world_size, device=self.device)
+        tail_idx = 2 * world_size - 1 - head_idx
+        paired = torch.stack([chunks[head_idx], chunks[tail_idx]], dim=1)
+        all_indices_tensor = paired.reshape(-1)
 
-        all_indices_tensor = torch.tensor(
-            all_indices, dtype=torch.int, device=self.device
-        )
         if restore:
             all_indices_tensor = torch.argsort(all_indices_tensor)
 
@@ -261,9 +251,10 @@ def _generate_indices(self, restore: bool = False) -> Tensor:
     def _generate_indices_for_batch(self, seq_length_per_doc, restore) -> Tensor:  # type: ignore[no-untyped-def]
         world_size = self.world_size
         device = self.device
-        assert all(
+        if not all(
             seq_length % (2 * world_size) == 0 for seq_length in seq_length_per_doc
-        )
+        ):
+            raise AssertionError
         chunk_length_per_doc = [
             seq_length // (2 * world_size) for seq_length in seq_length_per_doc
         ]
@@ -344,7 +335,8 @@ def ptrr_scheduling(process_time: Tensor, group_size: int) -> Tensor:
                 [0, 5, 6, 10]       # values = [9, 15, 8, 15], sum = 47
             ]
         """
-        assert process_time.ndim == 1
+        if process_time.ndim != 1:
+            raise AssertionError
 
         num_tasks = process_time.size(0)
 
@@ -460,9 +452,8 @@ def _generate_indices(self, restore: bool = False) -> Tensor:
 
         # NOTE: only support the case where the qkv block size are equal
         q_blk_size, kv_blk_size = block_mask.BLOCK_SIZE
-        assert q_blk_size == kv_blk_size, (
-            "for now only support q_blk_size == kv_blk_size"
-        )
+        if q_blk_size != kv_blk_size:
+            raise AssertionError("for now only support q_blk_size == kv_blk_size")
 
         indices = torch.arange(
             q_blk_size * ptrr_indices.size(1), device=ptrr_indices.device
diff --git a/torch/distributed/tensor/experimental/_func_map.py b/torch/distributed/tensor/experimental/_func_map.py
index ebc30bda2c3aa..0c4bd4c48b537 100644
--- a/torch/distributed/tensor/experimental/_func_map.py
+++ b/torch/distributed/tensor/experimental/_func_map.py
@@ -2,7 +2,6 @@
 # Copyright (c) Meta Platforms, Inc. and affiliates
 import functools
 from collections.abc import Callable, Sequence
-from typing import Optional, Union
 
 import torch
 from torch.distributed._functional_collectives import AsyncCollectiveTensor
@@ -18,9 +17,9 @@
 
 __all__ = ["local_map"]
 
-PlacementType = Optional[Sequence[Placement]]
-InputPlacements = Optional[tuple[PlacementType, ...]]
-OutputPlacements = Union[PlacementType, tuple[PlacementType, ...]]
+PlacementType = Sequence[Placement] | None
+InputPlacements = tuple[PlacementType, ...] | None
+OutputPlacements = PlacementType | tuple[PlacementType, ...]
 
 
 def local_map(
@@ -171,10 +170,11 @@ def _local_map_wrapped(
     # process input args
     flat_args, args_spec = pytree.tree_flatten(args)
     if in_placements is not None:
-        assert len(in_placements) == len(flat_args), (
-            f"in_placements length {len(in_placements)} does not match the number "
-            f"of input args {len(flat_args)}!"
-        )
+        if len(in_placements) != len(flat_args):
+            raise AssertionError(
+                f"in_placements length {len(in_placements)} does not match the number "
+                f"of input args {len(flat_args)}!"
+            )
 
     # we assume every DTensor object is placed on the same device mesh
     flat_local_args = []
@@ -192,9 +192,10 @@ def _local_map_wrapped(
 
             if in_placements is not None:
                 spec = in_placements[idx]
-                assert spec is not None, (
-                    f"DTensor input {arg} expects placements but received {spec}!"
-                )
+                if spec is None:
+                    raise AssertionError(
+                        f"DTensor input {arg} expects placements but received {spec}!"
+                    )
 
                 if not isinstance(spec, tuple):
                     spec = tuple(spec)
@@ -214,9 +215,10 @@ def _local_map_wrapped(
 
             if in_grad_placements is not None:
                 spec = in_grad_placements[idx]
-                assert spec is not None, (
-                    f"DTensor input {arg} expects in grad placements but received {spec}!"
-                )
+                if spec is None:
+                    raise AssertionError(
+                        f"DTensor input {arg} expects in grad placements but received {spec}!"
+                    )
                 if not isinstance(spec, tuple):
                     spec = tuple(spec)
                 local_arg = arg.to_local(grad_placements=spec)
@@ -231,10 +233,11 @@ def _local_map_wrapped(
             # Non-Tensor input must have None in `in_placements`
             if in_placements is not None and not isinstance(arg, torch.Tensor):
                 spec = in_placements[idx]
-                assert spec is None, (
-                    f"Non-Tensor input {arg} expects None placements "
-                    f"but received {spec}!"
-                )
+                if spec is not None:
+                    raise AssertionError(
+                        f"Non-Tensor input {arg} expects None placements "
+                        f"but received {spec}!"
+                    )
 
             flat_local_args.append(arg)
 
@@ -251,25 +254,28 @@ def _local_map_wrapped(
         out_placements_tuple = (
             out_placements if isinstance(out_placements, tuple) else (out_placements,)
         )
-        assert len(flat_out) == len(out_placements_tuple), (
-            "local_map requires one PlacementType be provided for each output value,"
-            f" received {len(out_placements_tuple)} out_placements but"
-            f" {len(flat_out)} is expected!"
-        )
+        if len(flat_out) != len(out_placements_tuple):
+            raise AssertionError(
+                "local_map requires one PlacementType be provided for each output value,"
+                f" received {len(out_placements_tuple)} out_placements but"
+                f" {len(flat_out)} is expected!"
+            )
         for out, spec in zip(flat_out, out_placements_tuple):
             if isinstance(out, torch.Tensor):
-                assert not isinstance(out, DTensor), (
-                    f"torch.Tensor output expected but received {type(out)}: {out}"
-                )
+                if isinstance(out, DTensor):
+                    raise AssertionError(
+                        f"torch.Tensor output expected but received {type(out)}: {out}"
+                    )
 
                 flat_dist_out.append(
                     # pyrefly: ignore [bad-argument-type]
                     DTensor.from_local(out, device_mesh, spec, run_check=False)
                 )
             else:
-                assert spec is None, (
-                    f"Non-tensor output {out} expects None placements but received {spec}!"
-                )
+                if spec is not None:
+                    raise AssertionError(
+                        f"Non-tensor output {out} expects None placements but received {spec}!"
+                    )
 
                 flat_dist_out.append(out)
 
diff --git a/torch/distributed/tensor/experimental/_tp_transform.py b/torch/distributed/tensor/experimental/_tp_transform.py
index 7b945e8507b24..adde891c4b68c 100644
--- a/torch/distributed/tensor/experimental/_tp_transform.py
+++ b/torch/distributed/tensor/experimental/_tp_transform.py
@@ -60,7 +60,8 @@ def tensor_parallel_transformation(
             exported_program.graph_signature,
             parallel_strategies,
         )(gm)
-        assert res is not None
+        if res is None:
+            raise AssertionError
         gm = res.graph_module
 
     return exported_program._update(gm, sig, state_dict=state_dict)
@@ -116,7 +117,8 @@ def _generate_parameter_and_buffer_placements(
     for linear_fqn, parallel_style in parallel_strategies.items():
         weight_fqn = f"{linear_fqn}.weight"
         bias_fqn = f"{linear_fqn}.bias"
-        assert weight_fqn in params_and_buffers
+        if weight_fqn not in params_and_buffers:
+            raise AssertionError
         parameter_placements[weight_fqn] = (
             Shard(0) if parallel_style == ColwiseParallel else Shard(1)
         )
@@ -205,9 +207,11 @@ def _mark_sharding(
         elif node.op == "call_function":
             if node.target is operator.getitem:
                 input_nodes = node.all_input_nodes
-                assert len(input_nodes) == 1, (
-                    f"non-compute op only support one input now, found node: {node} with length of inputs: {len(node.args)}"
-                )
+                if len(input_nodes) != 1:
+                    raise AssertionError(
+                        f"non-compute op only support one input now, found node: {node} "
+                        f"with length of inputs: {len(node.args)}"
+                    )
                 arg_strategy = placement_strategies[input_nodes[0]]
                 placement_strategies[node] = _create_placement_strategy(
                     node,
@@ -264,8 +268,10 @@ def _get_output_spec_from_output_sharding(
         return output_sharding.output_spec
     else:
         # For ops that return multiple outputs, the outputs should have the same output spec
-        assert isinstance(output_sharding.output_spec, Sequence)
-        assert output_sharding.output_spec[0] is not None
+        if not isinstance(output_sharding.output_spec, Sequence):
+            raise AssertionError
+        if output_sharding.output_spec[0] is None:
+            raise AssertionError
         output_sharding.output_spec[0].tensor_meta = None
         return output_sharding.output_spec[0]
 
@@ -295,16 +301,19 @@ def _populate_tensor_meta(node: Node, output_spec: OutputSpecType) -> None:
     Util function to populate tensor meta of output_spec based on node metadata.
     """
     if isinstance(node.meta["val"], Sequence):
-        assert isinstance(output_spec, Sequence)
+        if not isinstance(output_spec, Sequence):
+            raise AssertionError
         for spec, fake_tensor in zip(output_spec, node.meta["val"]):
-            assert spec is not None
+            if spec is None:
+                raise AssertionError
             spec.tensor_meta = TensorMeta(
                 shape=fake_tensor.shape,
                 stride=fake_tensor.stride(),
                 dtype=fake_tensor.dtype,
             )
     else:
-        assert isinstance(output_spec, DTensorSpec)
+        if not isinstance(output_spec, DTensorSpec):
+            raise AssertionError
         output_spec.tensor_meta = TensorMeta(
             shape=node.meta["val"].shape,
             stride=node.meta["val"].stride(),
@@ -421,11 +430,13 @@ def _partition_val(val: Any, spec: DTensorSpec) -> Any:
             return local_shard
 
         for idx, placement in enumerate(spec.placements):
+            # NOTE: is_shard() does not match _StridedShard; see _is_shard_like().
             if placement.is_shard():
                 placement = cast(Shard, placement)
                 num_chunks = spec.mesh.size(mesh_dim=idx)
                 my_coord = spec.mesh.get_coordinate()
-                assert my_coord is not None, "current rank not in mesh!"
+                if my_coord is None:
+                    raise AssertionError("current rank not in mesh!")
                 my_coord_on_mesh_dim = my_coord[idx]
                 local_shard = placement._select_split_tensor(
                     local_shard,
@@ -507,7 +518,8 @@ def _get_input_node_specs(
     for input_arg in node.all_input_nodes:
         if input_arg in placement_strategies:
             output_spec = placement_strategies[input_arg].output_specs
-            assert isinstance(output_spec, DTensorSpec)
+            if not isinstance(output_spec, DTensorSpec):
+                raise AssertionError
             input_specs_list.append(output_spec)
         else:
             raise ValueError(f"{input_arg} does not have output_spec populated.")
@@ -547,7 +559,8 @@ def _shard_state_dict(
             fqn = graph_signature.inputs_to_buffers[node.name]
         else:
             continue
-        assert fqn in state_dict, f"{fqn} not found in state dict: {state_dict.keys()}"
+        if fqn not in state_dict:
+            raise AssertionError(f"{fqn} not found in state dict: {state_dict.keys()}")
 
         original_param = state_dict[fqn]
         dtensor_param = distribute_tensor(
diff --git a/torch/distributed/tensor/parallel/ddp.py b/torch/distributed/tensor/parallel/ddp.py
index 19c1d3ca5477e..a1331f8bd8c87 100644
--- a/torch/distributed/tensor/parallel/ddp.py
+++ b/torch/distributed/tensor/parallel/ddp.py
@@ -29,7 +29,8 @@ def _update_module_param(param_list: list[tuple[nn.Module, str, nn.Parameter]]):
     """
     for item in param_list:
         parent_module, module_path, t = item
-        assert hasattr(parent_module, module_path)
+        if not hasattr(parent_module, module_path):
+            raise AssertionError
         delattr(parent_module, module_path)
         setattr(parent_module, module_path, t)
 
diff --git a/torch/distributed/tensor/parallel/fsdp.py b/torch/distributed/tensor/parallel/fsdp.py
index 9e68ed6b1dba5..5f6facf74542a 100644
--- a/torch/distributed/tensor/parallel/fsdp.py
+++ b/torch/distributed/tensor/parallel/fsdp.py
@@ -30,12 +30,14 @@
 
 def _get_box(tensor: DTensor) -> tuple[torch.Size, torch.Size]:
     device_mesh = tensor.device_mesh
-    assert device_mesh.ndim == 1, "Only 1D DeviceMeshes currently handled"
+    if device_mesh.ndim != 1:
+        raise AssertionError("Only 1D DeviceMeshes currently handled")
 
     placement = tensor.placements[0]
     offsets = [0] * len(tensor.size())
     num_chunks = device_mesh.size(mesh_dim=0)
 
+    # NOTE: is_shard() does not match _StridedShard; see _is_shard_like().
     if tensor.placements[0].is_shard():
         shard_dim = cast(DShard, placement).dim
         chunk_size = tensor.size(shard_dim) // num_chunks
@@ -52,13 +54,15 @@ def _get_box_for(tensor: DTensor, idx: int) -> tuple[torch.Size, torch.Size]:
 def _get_local_box(tensor: DTensor) -> tuple[torch.Size, torch.Size]:
     device_mesh = tensor.device_mesh
     coord = device_mesh.get_coordinate()
-    assert coord is not None
+    if coord is None:
+        raise AssertionError
     return _get_box_for(tensor, coord[0])
 
 
 def _create_shard_md_from_dt(dt: DTensor, current_rank: int) -> ShardMetadata:
     mesh = dt.device_mesh
-    assert mesh.ndim == 1, "Only 1D DeviceMeshes currently handled"
+    if mesh.ndim != 1:
+        raise AssertionError("Only 1D DeviceMeshes currently handled")
 
     offsets, sizes = _get_local_box(dt)
     return ShardMetadata(
@@ -78,6 +82,7 @@ def _create_sharded_tensor_md_from_dt(
     my_rank = dist.get_rank(dt_pg)
     scapegoat_rank = 0 if my_rank > 0 else 1
 
+    # NOTE: is_shard() does not match _StridedShard; see _is_shard_like().
     if dt.placements[0].is_shard():
         shard_count = dt_pg.size()
     else:
@@ -109,7 +114,8 @@ def _create_sharded_tensor_md_from_dt(
 
 def _get_dt_pg(dt: DTensor) -> c10d.ProcessGroup:
     mesh = dt.device_mesh
-    assert mesh.ndim == 1, "Only 1D DeviceMeshes currently handled"
+    if mesh.ndim != 1:
+        raise AssertionError("Only 1D DeviceMeshes currently handled")
     return mesh.get_group()
 
 
@@ -152,7 +158,8 @@ def _chunk_tensor(
     pg: dist.ProcessGroup,
 ) -> torch.Tensor:
     if type(tensor) is ShardedTensor:
-        assert len(tensor.local_shards()) == 1
+        if len(tensor.local_shards()) != 1:
+            raise AssertionError
 
         inner_param = tensor.local_tensor()
         inner_st = _create_chunk_sharded_tensor(
@@ -179,7 +186,8 @@ def _chunk_tensor(
         return st_outer
     elif type(tensor) is DTensor:
         device_mesh = tensor.device_mesh
-        assert device_mesh.ndim == 1, "Only 1D DeviceMeshes currently handled"
+        if device_mesh.ndim != 1:
+            raise AssertionError("Only 1D DeviceMeshes currently handled")
 
         inner_param = tensor._local_tensor
 
@@ -300,7 +308,8 @@ def _all_gather_dtensor(
     parent_mesh: DeviceMesh | None,
 ) -> torch.Tensor:
     """All gather a DTensor in its FSDP dimension and return the local tensor."""
-    assert parent_mesh == tensor.device_mesh
+    if parent_mesh != tensor.device_mesh:
+        raise AssertionError
 
     placements = list(copy.deepcopy(tensor.placements))
     # FSDP + TP: [Shard(0), tp_placement] -> [Replicate(), tp_placement]
diff --git a/torch/distributed/tensor/parallel/input_reshard.py b/torch/distributed/tensor/parallel/input_reshard.py
index 81e25621e040a..5e60cfcc7fb8f 100644
--- a/torch/distributed/tensor/parallel/input_reshard.py
+++ b/torch/distributed/tensor/parallel/input_reshard.py
@@ -4,6 +4,7 @@
 
 import torch
 from torch.distributed.tensor import DeviceMesh, DTensor, Replicate, Shard
+from torch.distributed.tensor.placement_types import _is_shard_like
 
 
 __all__ = [
@@ -87,7 +88,7 @@ def _unpack_hook_tp(mesh: DeviceMesh, input_reshard_dim: int, x: Any) -> torch.T
     if (
         isinstance(x, DTensor)
         and len(x._spec.placements) == 1
-        and x._spec.placements[0].is_shard()
+        and _is_shard_like(x._spec.placements[0])
     ):
         return x.redistribute(device_mesh=mesh, placements=[Replicate()])
     elif (
diff --git a/torch/distributed/tensor/parallel/loss.py b/torch/distributed/tensor/parallel/loss.py
index d0b4ba3ef6b47..b5c8a4980ea33 100644
--- a/torch/distributed/tensor/parallel/loss.py
+++ b/torch/distributed/tensor/parallel/loss.py
@@ -112,9 +112,10 @@ def _propagate_tensor_meta(
     kwargs: dict[str, object],
 ) -> TensorMeta:
     op_info = DTensor._op_dispatcher.unwrap_to_op_info(op_call, args, kwargs)
-    assert op_info.schema is not None, (
-        "op_info.schema should not be None after unwrap_to_op_info"
-    )
+    if op_info.schema is None:
+        raise AssertionError(
+            "op_info.schema should not be None after unwrap_to_op_info"
+        )
     tensor_meta = DTensor._op_dispatcher.sharding_propagator._propagate_tensor_meta(
         op_info.schema
     )
@@ -131,7 +132,8 @@ def _propagate_tensor_meta(
 # with all_reduce manually inserted to perform distributed computation.
 def _log_softmax(x, dim, half_to_float, mesh, mesh_dim):
     if half_to_float:
-        assert x.dtype == torch.half
+        if x.dtype != torch.half:
+            raise AssertionError
     computation_dtype, result_dtype = utils.elementwise_dtypes(
         x, type_promotion_kind=utils.ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT
     )
@@ -232,7 +234,8 @@ def _weight_view(weight: Tensor) -> Tensor:
 
     if weight is not None:
         w = _weight_view(weight)
-        assert local_weight is not None
+        if local_weight is None:
+            raise AssertionError
         local_w = _weight_view(local_weight)
         x = x * local_w
     safe_target = torch.where(target != ignore_index, target, 0)
@@ -308,7 +311,8 @@ def _nll_loss_forward_handler(
             Shard(0) if i == mesh_dim else Replicate() for i in range(spec.mesh.ndim)
         ]
         local_weight = weight.redistribute(spec.mesh, sharded_placements)._local_tensor
-        assert local_weight.shape[0] == x._local_tensor.shape[channel_dim]
+        if local_weight.shape[0] != x._local_tensor.shape[channel_dim]:
+            raise AssertionError
 
     if reduction == Reduction.NONE.value:
         output_placements = target_placements
@@ -383,7 +387,8 @@ def _nll_loss_and_log_softmax_backward(
     safe_target = safe_target.squeeze(channel_dim).flatten()
     masked_safe_target = partial_placement._partition_value(safe_target, mesh, mesh_dim)
     # only update grad_input to -1 if not masked
-    assert partial_placement.mask_buffer.data is not None
+    if partial_placement.mask_buffer.data is None:
+        raise AssertionError
     grad_update = partial_placement.mask_buffer.data.to(grad_input.dtype) - 1.0
     arange_1d = torch.arange(
         masked_safe_target.shape[0], device=masked_safe_target.device
diff --git a/torch/distributed/tensor/parallel/style.py b/torch/distributed/tensor/parallel/style.py
index c4a035f792bc1..22169c2a4e428 100644
--- a/torch/distributed/tensor/parallel/style.py
+++ b/torch/distributed/tensor/parallel/style.py
@@ -105,7 +105,10 @@ def _prepare_input_fn(
         input_tensor = inputs[0]
         if not isinstance(input_tensor, DTensor):
             input_tensor = DTensor.from_local(
-                input_tensor, device_mesh, input_layouts, run_check=False
+                input_tensor,
+                device_mesh,
+                input_layouts,
+                run_check=False,
             )
 
         # transform the input layouts to the desired layouts of ColwiseParallel
@@ -162,7 +165,9 @@ def _apply(self, module: nn.Module, device_mesh: DeviceMesh) -> nn.Module:
             device_mesh,
             partition_fn,
             partial(
-                self._prepare_input_fn, self.input_layouts, self.desired_input_layouts
+                self._prepare_input_fn,
+                self.input_layouts,
+                self.desired_input_layouts,
             ),
             partial(
                 self._prepare_output_fn, self.output_layouts, self.use_local_output
@@ -230,7 +235,10 @@ def _prepare_input_fn(
         input_tensor = inputs[0]
         if not isinstance(input_tensor, DTensor):
             input_tensor = DTensor.from_local(
-                input_tensor, device_mesh, input_layouts, run_check=False
+                input_tensor,
+                device_mesh,
+                input_layouts,
+                run_check=False,
             )
 
         if input_layouts != desired_input_layouts:
@@ -310,7 +318,9 @@ def _apply(self, module: nn.Module, device_mesh: DeviceMesh) -> nn.Module:
             device_mesh,
             partition_fn,
             partial(
-                self._prepare_input_fn, self.input_layouts, self.desired_input_layouts
+                self._prepare_input_fn,
+                self.input_layouts,
+                self.desired_input_layouts,
             ),
             partial(
                 self._prepare_output_fn, self.output_layouts, self.use_local_output
@@ -494,21 +504,20 @@ def __init__(
         )
         self.use_local_output = use_local_output
         if self.input_layouts is not None:
-            assert self.desired_input_layouts is not None, (
-                "desired module inputs should not be None!"
-            )
-            assert len(self.input_layouts) == len(self.desired_input_layouts), (
-                "input_layouts and desired_input_layouts should have same length!"
-            )
+            if self.desired_input_layouts is None:
+                raise AssertionError("desired module inputs should not be None!")
+            if len(self.input_layouts) != len(self.desired_input_layouts):
+                raise AssertionError(
+                    "input_layouts and desired_input_layouts should have same length!"
+                )
         self.with_kwargs = input_kwarg_layouts is not None
         self.input_kwarg_layouts = input_kwarg_layouts or {}
         self.desired_input_kwarg_layouts = desired_input_kwarg_layouts or {}
         if self.with_kwargs:
-            assert len(self.input_kwarg_layouts) == len(
-                self.desired_input_kwarg_layouts
-            ), (
-                "input_kwarg_layouts and desired_input_kwarg_layouts should have same length!"
-            )
+            if len(self.input_kwarg_layouts) != len(self.desired_input_kwarg_layouts):
+                raise AssertionError(
+                    "input_kwarg_layouts and desired_input_kwarg_layouts should have same length!"
+                )
 
     def _prepare_input_arg(
         self,
@@ -523,9 +532,8 @@ def _prepare_input_arg(
                 # assert inp.placements[0] == input_layout
                 dt_inp = input
             else:
-                assert isinstance(input, torch.Tensor), (
-                    "expecting input to be a torch.Tensor!"
-                )
+                if not isinstance(input, torch.Tensor):
+                    raise AssertionError("expecting input to be a torch.Tensor!")
                 dt_inp = DTensor.from_local(
                     input, mesh, (input_layout,), run_check=False
                 )
@@ -546,9 +554,8 @@ def _prepare_input_fn(self, inputs, device_mesh):
         if len(inputs) != len(self.input_layouts):
             raise ValueError("module inputs and input_layouts should have same length!")
 
-        assert self.desired_input_layouts is not None, (
-            "desired module inputs should not be None!"
-        )
+        if self.desired_input_layouts is None:
+            raise AssertionError("desired module inputs should not be None!")
 
         for inp, input_layout, desired_layout in zip(
             inputs, self.input_layouts, self.desired_input_layouts
@@ -653,9 +660,10 @@ def __init__(
             else desired_output_layouts
         )
         self.use_local_output = use_local_output
-        assert len(self.output_layouts) == len(self.desired_output_layouts), (
-            "output_layouts and desired_output_layouts should have same length!"
-        )
+        if len(self.output_layouts) != len(self.desired_output_layouts):
+            raise AssertionError(
+                "output_layouts and desired_output_layouts should have same length!"
+            )
 
     def _prepare_out_fn(self, outputs, device_mesh):
         prepared_outputs = []
diff --git a/torch/distributed/tensor/placement_types.py b/torch/distributed/tensor/placement_types.py
index 975e49eec5a70..67ab7b3f12754 100644
--- a/torch/distributed/tensor/placement_types.py
+++ b/torch/distributed/tensor/placement_types.py
@@ -1,14 +1,14 @@
 # mypy: allow-untyped-defs
 # Copyright (c) Meta Platforms, Inc. and affiliates
 
-import typing
+import functools
+from collections.abc import Sequence
 from dataclasses import dataclass, field
-from typing import cast, TypeVar
+from typing import cast, TypeGuard, TypeVar
 
 import torch
 import torch._C
 import torch.distributed._functional_collectives as funcol
-from torch import sym_min
 from torch._C._distributed import Placement
 from torch.distributed import RankType
 from torch.distributed._local_tensor import maybe_run_for_local_tensor
@@ -22,6 +22,7 @@
     unpad_tensor,
 )
 from torch.distributed.tensor._ops._mask_buffer import MaskBuffer
+from torch.types import IntLikeType
 
 
 __all__ = ["Placement", "Shard", "Replicate", "Partial"]
@@ -33,6 +34,18 @@
 Placement.__module__ = "torch.distributed.tensor.placement_types"
 
 
+def _raise_error(method_name, _):
+    raise RuntimeError(
+        f"Placement method '{method_name}' should not be called as "
+        "it should be overridden by the subclass"
+    )
+
+
+Placement.__eq__ = functools.partial(_raise_error, method_name="__eq__")
+Placement.__hash__ = functools.partial(_raise_error, method_name="__hash__")
+Placement.__fx_repr__ = functools.partial(_raise_error, method_name="__fx_repr__")
+
+
 class Shard(torch._C._distributed.Shard):
     """
     The ``Shard(dim)`` placement describes the DTensor sharding on tensor dimension
@@ -49,6 +62,10 @@ class Shard(torch._C._distributed.Shard):
 
     .. warning:: sharding on a tensor dimension where the tensor dimension size is not
         evenly divisible on a DeviceMesh dimension is currently experimental and subject to change.
+
+    .. note:: When checking whether a placement is shard-like, use
+        :func:`_is_shard_like` instead of ``isinstance(p, Shard)`` to also
+        match :class:`_StridedShard`.
     """
 
     def _split_tensor(
@@ -80,15 +97,13 @@ def _split_tensor_helper(
         contiguous: bool,
         dim: int,
     ) -> tuple[list[torch.Tensor], list[int]]:
-        assert dim <= tensor.ndim, (
-            f"Sharding dim {dim} greater than tensor ndim {tensor.ndim}"
-        )
+        if dim > tensor.ndim:
+            raise AssertionError(
+                f"Sharding dim {dim} greater than tensor ndim {tensor.ndim}"
+            )
 
         # chunk tensor over dimension `dim` into n slices
-        tensor_list = list(torch.chunk(tensor, num_chunks, dim=dim))
-        tensor_list = fill_empty_tensor_to_shards(
-            tensor_list, dim, num_chunks - len(tensor_list)
-        )
+        tensor_list = Shard._custom_chunk(tensor, num_chunks, dim=dim)
 
         # compute the chunk size inline with ``torch.chunk`` to calculate padding
         full_chunk_size = (tensor.size(dim) + num_chunks - 1) // num_chunks
@@ -145,9 +160,10 @@ def _select_split_tensor(
         # For the SymInt implementation just compute the value for the tensor we
         # want rather than computing all of them.
 
-        assert self.dim <= tensor.ndim, (
-            f"Sharding dim {self.dim} greater than tensor ndim {tensor.ndim}"
-        )
+        if self.dim > tensor.ndim:
+            raise AssertionError(
+                f"Sharding dim {self.dim} greater than tensor ndim {tensor.ndim}"
+            )
 
         # chunk tensor over dimension `dim` into n slices
         dim_size = tensor.size(self.dim)
@@ -164,10 +180,44 @@ def _select_split_tensor(
             result = result.contiguous()
         return result
 
+    @staticmethod
+    def _custom_chunk(
+        tensor: torch.Tensor, num_chunks: int, dim: int
+    ) -> list[torch.Tensor]:
+        """
+        Returns list of tensor chunks along dim.
+        Uses torch.chunk in eager mode, but torch.narrow under tracing to be unbacked-symint safe.
+        Also handles uneven/zero-sharding cases.
+        """
+        from torch.distributed._functional_collectives import _are_we_tracing
+        from torch.fx.experimental.symbolic_shapes import has_free_unbacked_symbols
+
+        if tensor.dim() <= 0:
+            raise AssertionError(f"Expected tensor.dim() > 0, got {tensor.dim()}")
+        if num_chunks <= 0:
+            raise AssertionError(f"Expected num_chunks > 0, got {num_chunks}")
+
+        # TODO(pianpwk): remove the unbacked symbols check and fix AsyncTP pattern matching
+        # for test_micro_pipeline_tp.py.
+        if not _are_we_tracing() or not has_free_unbacked_symbols(tensor):
+            tensor_list = list(torch.chunk(tensor, num_chunks, dim=dim))
+            return fill_empty_tensor_to_shards(
+                tensor_list, dim, num_chunks - len(tensor_list)
+            )
+        else:
+            dim_size = tensor.size(dim)
+            split_size = (dim_size + num_chunks - 1) // num_chunks
+            chunks = []
+            for i in range(num_chunks):
+                start = torch.sym_min(split_size * i, dim_size)
+                end = torch.sym_min(split_size * (i + 1), dim_size)
+                chunks.append(tensor.narrow(dim, start, end - start))
+            return chunks
+
     @staticmethod
     @maybe_run_for_local_tensor
     def local_shard_size_and_offset(
-        curr_local_size: int,
+        curr_local_size: IntLikeType,
         num_chunks: int,
         rank: _RankTypeT,
     ) -> tuple[_RankTypeT, _RankTypeT]:
@@ -182,8 +232,11 @@ def local_shard_size_and_offset(
         Returns (new local shard size, offset)
 
         """
+        from torch.fx.experimental.symbolic_shapes import guard_or_false
+
         # Compute the chunk size inline with ``torch.chunk``
-        if curr_local_size % num_chunks == 0:
+        # Assume uneven sharding as general case for unbacked sizes.
+        if guard_or_false(curr_local_size % num_chunks == 0):
             full_chunk_size = curr_local_size // num_chunks
             # pyrefly: ignore[bad-assignment] # pyrefly bug?
             shard_starting_idx: _RankTypeT = full_chunk_size * rank
@@ -194,15 +247,11 @@ def local_shard_size_and_offset(
         # pyrefly: ignore[bad-assignment] # pyrefly bug?
         shard_starting_idx: _RankTypeT = full_chunk_size * rank
 
-        if curr_local_size < shard_starting_idx:
-            # pyrefly: ignore[bad-return]
-            return 0, typing.cast(_RankTypeT, curr_local_size)
-        else:
-            local_shard_size = (
-                sym_min(curr_local_size, shard_starting_idx + full_chunk_size)
-                - shard_starting_idx
-            )
-            return local_shard_size, shard_starting_idx
+        shard_end_idx = torch.sym_min(
+            curr_local_size, shard_starting_idx + full_chunk_size
+        )
+        local_shard_size = torch.sym_max(0, shard_end_idx - shard_starting_idx)
+        return local_shard_size, torch.sym_min(curr_local_size, shard_starting_idx)
 
     def _local_shard_size_and_offset(
         self,
@@ -210,7 +259,7 @@ def _local_shard_size_and_offset(
         num_chunks: int,
         rank: RankType,
     ) -> tuple[int, RankType]:
-        # pyrefly: ignore[bad-argument-type]  # pyrefly bug
+        # pyrefly: ignore [bad-argument-type, bad-return]
         return Shard.local_shard_size_and_offset(curr_local_size, num_chunks, rank)
 
     @staticmethod
@@ -218,8 +267,11 @@ def _local_shard_size_and_offset(
     def _maybe_unpad_tensor_with_sizes(
         dim, local_tensor, pad_sizes, mesh_dim_local_rank, make_contiguous
     ) -> torch.Tensor:
+        from torch.fx.experimental.symbolic_shapes import guard_or_true
+
         # Only unpad if the local_tensor was padded on the dimension.
-        if pad_sizes[mesh_dim_local_rank] > 0:
+        # Assume padding (uneven sharding) as general case for unbacked sizes.
+        if guard_or_true(pad_sizes[mesh_dim_local_rank] > 0):
             local_tensor = unpad_tensor(
                 local_tensor, dim, pad_sizes[mesh_dim_local_rank]
             )
@@ -273,7 +325,8 @@ def _shard_tensor(
         first = next(it)
         # Tensors in the scatter list are expected to have the same shape because
         # split is requested with padding.
-        assert all(first.shape == v.shape for v in it)
+        if not all(first.shape == v.shape for v in it):
+            raise AssertionError
 
         output = torch.empty_like(first)
 
@@ -308,13 +361,18 @@ def _reduce_shard_tensor(
         """
         reduce and scatter a tensor on a mesh dimension
         """
-        if not mesh._is_current_rank_part_of_mesh():
+        from torch.fx.experimental.symbolic_shapes import guard_or_true
+
+        my_coordinate = mesh.get_coordinate()
+        num_chunks = mesh.size(mesh_dim=mesh_dim)
+
+        if my_coordinate is None:
             # if rank is not part of mesh, we simply return local_tensor,
             # which should be an empty tensor
             return tensor
 
-        num_chunks = mesh.size(mesh_dim=mesh_dim)
-        is_padded = tensor.size(self.dim) % num_chunks != 0
+        # Assume padding (uneven sharding) as general case for unbacked sizes.
+        is_padded = guard_or_true(tensor.size(self.dim) % num_chunks != 0)
         pad_sizes = None
         if is_padded:
             scattered_list, pad_sizes = self._split_tensor(
@@ -329,7 +387,8 @@ def _reduce_shard_tensor(
         )
 
         if is_padded:
-            assert pad_sizes is not None
+            if pad_sizes is None:
+                raise AssertionError
             output = Shard._maybe_unpad_tensor_with_sizes(
                 self.dim, output, pad_sizes, mesh._sym_get_coordinate(mesh_dim), False
             )
@@ -339,10 +398,13 @@ def _reduce_shard_tensor(
     def _maybe_pad_tensor(
         self,
         local_tensor: torch.Tensor,
-        logical_dim_size: int,
+        logical_dim_size: IntLikeType,
         num_chunks: int,
     ) -> torch.Tensor:
-        is_padded = logical_dim_size % num_chunks != 0
+        from torch.fx.experimental.symbolic_shapes import guard_or_true
+
+        # Assume padding (uneven sharding) as general case for unbacked sizes.
+        is_padded = guard_or_true(logical_dim_size % num_chunks != 0)
 
         if is_padded:
             full_chunk_size = (logical_dim_size + num_chunks - 1) // num_chunks
@@ -358,16 +420,35 @@ def _maybe_pad_tensor(
     def _maybe_unpad_tensor(
         self,
         local_tensor: torch.Tensor,
-        logical_dim_size: int,
+        logical_dim_size: IntLikeType,
         num_chunks: int,
     ) -> torch.Tensor:
-        is_padded = logical_dim_size % num_chunks != 0
+        from torch.fx.experimental.symbolic_shapes import guard_or_false, guard_or_true
+
+        # Assume padding (uneven sharding) as general case for unbacked sizes.
+        is_padded = guard_or_true(logical_dim_size % num_chunks != 0)
 
         if is_padded:
             full_chunk_size = (logical_dim_size + num_chunks - 1) // num_chunks
             unpad_size = full_chunk_size * num_chunks - logical_dim_size  # type: ignore[possibly-undefined]
             local_tensor = unpad_tensor(local_tensor, self.dim, unpad_size)
 
+        # Bind derived symbolic sizes (e.g. 2*(s//2)) back to the original
+        # symbol - needed for correct shape propagation and dynamo generation
+        if local_tensor.size(self.dim) is not logical_dim_size:
+            orig_size = local_tensor.size(self.dim)
+            torch._check(orig_size >= logical_dim_size)
+            local_tensor = local_tensor.narrow(self.dim, 0, logical_dim_size)
+
+            # Safety check: the narrow should never change the concrete size.
+            # Use guard_or_false so we don't trigger data-dependent guards
+            # on unbacked symints.
+            if guard_or_false(local_tensor.size(self.dim) != orig_size):
+                raise RuntimeError(
+                    f"narrow unexpectedly changed concrete size on dim {self.dim}: "
+                    f"{orig_size} -> {local_tensor.size(self.dim)}"
+                )
+
         return local_tensor
 
     def _to_replicate_tensor(
@@ -375,7 +456,7 @@ def _to_replicate_tensor(
         local_tensor: torch.Tensor,
         mesh: DeviceMesh,
         mesh_dim: int,
-        current_logical_shape: list[int],
+        current_logical_shape: Sequence[IntLikeType],
     ) -> torch.Tensor:
         """
         This function all_gather all shards and return a tensor that
@@ -403,7 +484,7 @@ def _replicate_to_shard(
         local_tensor: torch.Tensor,
         mesh: DeviceMesh,
         mesh_dim: int,
-        shard_index: int,
+        shard_index: IntLikeType,
     ) -> torch.Tensor:
         """
         transform from replicated tensor to a sharded tensor on
@@ -430,15 +511,17 @@ def _get_shard_pad_size(
 
     @staticmethod
     def _compute_padding_info(
-        current_logical_shape: list[int],
+        current_logical_shape: Sequence[IntLikeType],
         num_chunks: int,
         old_shard_dim: int,
         new_shard_dim: int,
-    ) -> tuple[bool, int, int, bool, int, int]:
+    ) -> tuple[bool, IntLikeType, int, bool, IntLikeType, int]:
+        from torch.fx.experimental.symbolic_shapes import guard_or_true
+
         results = []
         for shard_dim in [old_shard_dim, new_shard_dim]:
             dim_logical_size = current_logical_shape[shard_dim]
-            dim_padding = dim_logical_size % num_chunks != 0
+            dim_padding = guard_or_true(dim_logical_size % num_chunks != 0)
             dim_full_chunk_size = (dim_logical_size + num_chunks - 1) // num_chunks
             results.append((dim_padding, dim_logical_size, dim_full_chunk_size))
 
@@ -447,7 +530,7 @@ def _compute_padding_info(
     @staticmethod
     @maybe_run_for_local_tensor
     def _pad_for_new_shard_dim(
-        current_logical_shape: list[int],
+        current_logical_shape: Sequence[IntLikeType],
         local_tensor: torch.Tensor,
         num_chunks: int,
         old_shard_dim: int,
@@ -482,7 +565,7 @@ def _pad_for_new_shard_dim(
     @staticmethod
     @maybe_run_for_local_tensor
     def _unpad_for_new_shard_dim(
-        current_logical_shape: list[int],
+        current_logical_shape: Sequence[IntLikeType],
         local_tensor: torch.Tensor,
         num_chunks: int,
         old_shard_dim: int,
@@ -521,7 +604,7 @@ def _to_new_shard_dim(
         local_tensor: torch.Tensor,
         mesh: DeviceMesh,
         mesh_dim: int,
-        current_logical_shape: list[int],
+        current_logical_shape: Sequence[IntLikeType],
         new_shard_dim: int,
     ) -> torch.Tensor:
         """
@@ -564,6 +647,13 @@ def __repr__(self) -> str:
         """
         return f"Shard(dim={self.dim})"
 
+    def __fx_repr__(self):
+        """
+        Returns FX-evaluable repr and required globals for Shard placement.
+        Needed for passing this type as an opaque object input to a custom op.
+        """
+        return f"torch.distributed.tensor.placement_types.Shard(dim={self.dim})", {}
+
     def __str__(self) -> str:
         """human readable representation of the Shard placement"""
         return f"S({self.dim})"
@@ -650,6 +740,16 @@ def __str__(self) -> str:
     def _select_shard(shards: list[torch.Tensor], shard_index) -> torch.Tensor:
         return shards[shard_index].clone()
 
+    def __fx_repr__(self):
+        """
+        Returns FX-evaluable repr and required globals for Shard placement.
+        Needed for passing this type as an opaque object input to a custom op.
+        """
+        return (
+            f"torch.distributed.tensor.placement_types._StridedShard(dim={self.dim}, sf={self.split_factor})",  # noqa: B950
+            {},
+        )
+
     @classmethod
     def _make_shard_tensor(
         cls,
@@ -709,7 +809,8 @@ def _shard_tensor(
         first = next(it)
         # Tensors in the scatter list are expected to have the same shape because
         # split is requested with padding.
-        assert all(first.shape == v.shape for v in it)
+        if not all(first.shape == v.shape for v in it):
+            raise AssertionError
 
         output = torch.empty_like(first)
 
@@ -730,9 +831,10 @@ def _split_tensor(
         with_padding: bool = True,
         contiguous: bool = True,
     ) -> tuple[list[torch.Tensor], list[int]]:
-        assert self.dim <= tensor.ndim, (
-            f"Sharding dim {self.dim} greater than tensor ndim {tensor.ndim}"
-        )
+        if self.dim > tensor.ndim:
+            raise AssertionError(
+                f"Sharding dim {self.dim} greater than tensor ndim {tensor.ndim}"
+            )
 
         # Essentially _StridedShard express the right-to-left sharding in the
         # reversed order. Here we perform first_split as the virtual "right" sharding,
@@ -777,7 +879,7 @@ def _select_split_tensor(
         self,
         tensor: torch.Tensor,
         num_chunks: int,
-        index: int,
+        index: IntLikeType,
         *,
         with_padding: bool = True,
         contiguous: bool = True,
@@ -811,7 +913,7 @@ def _to_replicate_tensor(
         local_tensor: torch.Tensor,
         mesh: DeviceMesh,
         mesh_dim: int,
-        current_logical_shape: list[int],
+        current_logical_shape: Sequence[IntLikeType],
     ) -> torch.Tensor:
         """
         Replay the replicate-to-shard process to understand how to stitch shards back.
@@ -906,6 +1008,7 @@ def _to_replicate_tensor(
         # indices_tensor is 1D torch.arange(logical_dim_size) unsqueezed
         # so that we can reuse self._split_tensor which splits on self.dim
         shape = [1] * self.dim + [logical_dim_size]
+        # pyrefly: ignore [no-matching-overload]
         indices_tensor = torch.arange(
             logical_dim_size, device=local_tensor.device
         ).view(shape)
@@ -970,7 +1073,7 @@ def _replicate_to_strided_shard(
         local_tensor: torch.Tensor,
         mesh: DeviceMesh,
         mesh_dim: int,
-        shard_index: int,
+        shard_index: IntLikeType,
     ) -> torch.Tensor:
         """
         Transform from replicated tensor to a strided-sharded tensor on the current rank.
@@ -1017,7 +1120,7 @@ def _local_shard_size_and_offset(
     @maybe_run_for_local_tensor
     def local_shard_size_and_offset(
         self,
-        curr_local_size: int,
+        curr_local_size: IntLikeType,
         num_chunks: int,
         rank: RankType,
         return_first_offset: bool = True,
@@ -1047,6 +1150,7 @@ def local_shard_size_and_offset(
         # indices_tensor is 1D torch.arange(logical_dim_size) unsqueezed
         # so that we can reuse self._split_tensor which splits on self.dim
         shape = [1] * self.dim + [curr_local_size]
+        # pyrefly: ignore [no-matching-overload]
         indices_tensor = torch.arange(
             curr_local_size,
         ).view(shape)
@@ -1074,6 +1178,17 @@ def local_shard_size_and_offset(
         return local_shard_size, offsets
 
 
+def _is_shard_like(p: "Placement") -> TypeGuard[Shard | _StridedShard]:
+    """Check if a placement is Shard or _StridedShard.
+
+    Use this instead of ``isinstance(p, Shard)`` to avoid silently missing
+    ``_StridedShard``.  When ``_StridedShard`` is unified with ``Shard``
+    (see TODO on the class), this helper can be collapsed to a single
+    ``isinstance`` check.
+    """
+    return isinstance(p, Shard | _StridedShard)
+
+
 class Replicate(torch._C._distributed.Replicate):
     """
     The ``Replicate()`` placement describes the DTensor replicating on a corresponding
@@ -1092,6 +1207,13 @@ def __repr__(self) -> str:
         """
         return "Replicate()"
 
+    def __fx_repr__(self):
+        """
+        Returns FX-evaluable repr and required globals for Replicate placement.
+        Needed for passing this type as an opaque object input to a custom op.
+        """
+        return "torch.distributed.tensor.placement_types.Replicate()", {}
+
     def __str__(self) -> str:
         """
         human readable representation of the Replicate placement
@@ -1136,6 +1258,7 @@ class Partial(torch._C._distributed.Partial):
     # reduce_ops that distribute over addition, enabling per-input linearity
     # for bilinear ops like mm: reduce_op(A_i @ B) = reduce_op(A_i) @ B
     LINEAR_REDUCE_OPS: tuple[str, ...] = ("sum", "avg")
+    ALL_REDUCE_OPS: tuple[str, ...] = ("sum", "avg", "min", "max", "product")
 
     """
     The ``Partial(reduce_op)`` placement describes the DTensor that is pending
@@ -1239,6 +1362,16 @@ def __repr__(self) -> str:
         """
         return f"Partial({self.reduce_op})"
 
+    def __fx_repr__(self):
+        """
+        Returns FX-evaluable repr and required globals for Partial placement.
+        Needed for passing this type as an input to a custom op.
+        """
+        return (
+            f"torch.distributed.tensor.placement_types.Partial({self.reduce_op!r})",
+            {},
+        )
+
     def __str__(self) -> str:
         """
         human readable representation of the Partial placement
@@ -1286,7 +1419,9 @@ def __init__(
     @staticmethod
     @maybe_run_for_local_tensor
     def _mask_tensor(
-        tensor: torch.Tensor, local_offset_on_dim: int, local_shard_size: int
+        tensor: torch.Tensor,
+        local_offset_on_dim: IntLikeType,
+        local_shard_size: IntLikeType,
     ) -> tuple[torch.Tensor, torch.Tensor]:
         # Build the input mask and save it for the current partial placement
         # this is so that the output of embedding op can reuse the same partial
@@ -1295,20 +1430,22 @@ def _mask_tensor(
             tensor >= local_offset_on_dim + local_shard_size
         )
         # mask the input tensor
+        # pyrefly: ignore [unsupported-operation]
         masked_tensor = tensor.clone() - local_offset_on_dim
         masked_tensor[mask] = 0
+        # pyrefly: ignore [bad-return]
         return mask, masked_tensor
 
     def _partition_value(
         self, tensor: torch.Tensor, mesh: DeviceMesh, mesh_dim: int
     ) -> torch.Tensor:
-        assert mesh._is_current_rank_part_of_mesh(), "rank is not part of mesh"
+        if not mesh._is_current_rank_part_of_mesh():
+            raise AssertionError("rank is not part of mesh")
         # override parent logic to perform partial mask for embedding
         num_chunks = mesh.size(mesh_dim)
         # get local shard size and offset on the embedding_dim
-        assert self.offset_shape is not None, (
-            "offset_shape needs to be set for _MaskPartial"
-        )
+        if self.offset_shape is None:
+            raise AssertionError("offset_shape needs to be set for _MaskPartial")
         local_shard_size, local_offset_on_dim = Shard.local_shard_size_and_offset(
             self.offset_shape[self.offset_dim],
             num_chunks,
@@ -1325,7 +1462,8 @@ def _reduce_value(
         self, tensor: torch.Tensor, mesh: DeviceMesh, mesh_dim: int
     ) -> torch.Tensor:
         # by the time we need reduction, we should have already saved the mask
-        assert self.mask_buffer.data is not None
+        if self.mask_buffer.data is None:
+            raise AssertionError
 
         # apply the mask to the tensor that pending reduction
         self.mask_buffer.apply_mask(tensor)
@@ -1346,7 +1484,8 @@ def _reduce_shard_value(
         shard_spec: Placement,
     ) -> torch.Tensor:
         # by the time we need reduction, we should have already saved the mask
-        assert self.mask_buffer.data is not None
+        if self.mask_buffer.data is None:
+            raise AssertionError
 
         # apply the mask to the tensor that pending reduction
         self.mask_buffer.apply_mask(tensor)
@@ -1390,3 +1529,47 @@ def __str__(self) -> str:
         human readable representation of the _MaskPartial placement
         """
         return f"MaskP({self.reduce_op}, {self.offset_shape}, {self.offset_dim})"
+
+    def __fx_repr__(self):
+        """
+        Returns FX-evaluable repr and required globals for Partial placement.
+        Needed for passing this type as an input to a custom op.
+        """
+        return (
+            f"torch.distributed.tensor.placement_types.MaskPartial(reduce_op={self.reduce_op}, offset_shape={self.offset_shape}, offset_dim={self.offset_dim})",  # noqa: B950
+            {},
+        )
+
+
+def _register_placements_as_opaque():
+    from torch._library.opaque_object import MemberType, register_opaque_type
+
+    allowed_members = {
+        "is_shard": MemberType.USE_REAL,
+        "is_partial": MemberType.USE_REAL,
+        "is_replicate": MemberType.USE_REAL,
+        "__eq__": MemberType.USE_REAL,
+    }
+    register_opaque_type(Placement, typ="value", members=allowed_members)
+    register_opaque_type(
+        Shard, typ="value", members=allowed_members | {"dim": MemberType.USE_REAL}
+    )
+    register_opaque_type(Replicate, typ="value", members=allowed_members)
+    register_opaque_type(
+        Partial,
+        typ="value",
+        members=allowed_members | {"reduce_op": MemberType.USE_REAL},
+    )
+    register_opaque_type(
+        _StridedShard,
+        typ="value",
+        members=allowed_members | {"dim": MemberType.USE_REAL},
+    )
+    register_opaque_type(
+        _MaskPartial,
+        typ="value",
+        members=allowed_members | {"reduce_op": MemberType.USE_REAL},
+    )
+
+
+_register_placements_as_opaque()
diff --git a/torch/distributed/utils.py b/torch/distributed/utils.py
index 28a18428c673d..d5651926b102a 100644
--- a/torch/distributed/utils.py
+++ b/torch/distributed/utils.py
@@ -134,16 +134,16 @@ def to_map(obj):
         from torch.nn.parallel.scatter_gather import _is_namedtuple
 
         if _is_namedtuple(obj):
-            # pyrefly: ignore [no-matching-overload]
+            # pyrefly: ignore [bad-argument-type, no-matching-overload]
             return [type(obj)(*args) for args in zip(*map(to_map, obj))]
         if isinstance(obj, tuple) and len(obj) > 0:
-            # pyrefly: ignore [no-matching-overload]
+            # pyrefly: ignore [bad-argument-type, no-matching-overload]
             return list(zip(*map(to_map, obj)))
         if isinstance(obj, list) and len(obj) > 0:
-            # pyrefly: ignore [no-matching-overload]
+            # pyrefly: ignore [bad-argument-type, no-matching-overload]
             return [list(i) for i in zip(*map(to_map, obj))]
         if isinstance(obj, dict) and len(obj) > 0:
-            # pyrefly: ignore [no-matching-overload]
+            # pyrefly: ignore [bad-argument-type, no-matching-overload]
             return [type(obj)(i) for i in zip(*map(to_map, obj.items()))]
         return [obj]
 
diff --git a/torch/distributions/generalized_pareto.py b/torch/distributions/generalized_pareto.py
index e04f38b87c707..6b5dabf542dd2 100644
--- a/torch/distributions/generalized_pareto.py
+++ b/torch/distributions/generalized_pareto.py
@@ -131,7 +131,6 @@ def variance(self):
         concentration = self.concentration
         valid = concentration < 0.5
         safe_conc = torch.where(valid, concentration, 0.25)
-        # pyrefly: ignore [unsupported-operation]
         result = self.scale**2 / ((1 - safe_conc) ** 2 * (1 - 2 * safe_conc))
         return torch.where(valid, result, nan)
 
diff --git a/torch/distributions/kl.py b/torch/distributions/kl.py
index 85932828d21af..43991b0ac265a 100644
--- a/torch/distributions/kl.py
+++ b/torch/distributions/kl.py
@@ -970,4 +970,5 @@ def _add_kl_info():
         )
     kl_info = "\n\t".join(rows)
     if kl_divergence.__doc__:
+        # pyrefly: ignore [missing-attribute]
         kl_divergence.__doc__ += kl_info
diff --git a/torch/distributions/normal.py b/torch/distributions/normal.py
index 4607d8d1cc4ec..6c50cdfc29f73 100644
--- a/torch/distributions/normal.py
+++ b/torch/distributions/normal.py
@@ -88,7 +88,6 @@ def log_prob(self, value):
         if self._validate_args:
             self._validate_sample(value)
         # compute the variance
-        # pyrefly: ignore [unsupported-operation]
         var = self.scale**2
         log_scale = (
             math.log(self.scale)
diff --git a/torch/distributions/transforms.py b/torch/distributions/transforms.py
index 96939173ca23a..0ba84fe5e8bd8 100644
--- a/torch/distributions/transforms.py
+++ b/torch/distributions/transforms.py
@@ -890,7 +890,6 @@ def _call(self, x):
         # apply stick-breaking on the squared values
         # Note that y = sign(r) * sqrt(z * z1m_cumprod)
         #             = (sign(r) * sqrt(z)) * sqrt(z1m_cumprod) = r * sqrt(z1m_cumprod)
-        # pyrefly: ignore [unsupported-operation]
         z = r**2
         z1m_cumprod_sqrt = (1 - z).sqrt().cumprod(-1)
         # Diagonal elements must be 1.
diff --git a/torch/distributions/von_mises.py b/torch/distributions/von_mises.py
index 72642c1926198..552b5c5f6662a 100644
--- a/torch/distributions/von_mises.py
+++ b/torch/distributions/von_mises.py
@@ -163,10 +163,8 @@ def _concentration(self) -> Tensor:
     @lazy_property
     def _proposal_r(self) -> Tensor:
         kappa = self._concentration
-        # pyrefly: ignore [unsupported-operation]
         tau = 1 + (1 + 4 * kappa**2).sqrt()
         rho = (tau - (2 * tau).sqrt()) / (2 * kappa)
-        # pyrefly: ignore [unsupported-operation]
         _proposal_r = (1 + rho**2) / (2 * rho)
         # second order Taylor expansion around 0 for small kappa
         _proposal_r_taylor = 1 / kappa + kappa
diff --git a/torch/export/_leakage_detection_utils.py b/torch/export/_leakage_detection_utils.py
index fe211e1dc079c..722a756431ea1 100644
--- a/torch/export/_leakage_detection_utils.py
+++ b/torch/export/_leakage_detection_utils.py
@@ -2,8 +2,9 @@
 import types
 import typing
 import weakref
+from typing_extensions import TypeIs
 
-import torch
+from torch.fx.experimental.symbolic_shapes import TrackedFake
 
 
 """
@@ -37,8 +38,8 @@ def _is_globals_or_locals(obj: typing.Any) -> bool:
     return obj is globals() or obj is locals()
 
 
-def _is_tracked_fake(obj: typing.Any) -> bool:
-    return isinstance(obj, torch.fx.experimental.symbolic_shapes.TrackedFake)
+def _is_tracked_fake(obj: typing.Any) -> TypeIs[TrackedFake]:
+    return isinstance(obj, TrackedFake)
 
 
 def _is_gm_meta_like_dict(d: dict, o: typing.Any) -> bool:
diff --git a/torch/export/_trace.py b/torch/export/_trace.py
index 5c827b5a90a18..b4753adda5b82 100644
--- a/torch/export/_trace.py
+++ b/torch/export/_trace.py
@@ -71,7 +71,8 @@
     aot_export_joint_with_descriptors,
 )
 from torch._guards import detect_fake_mode, tracing, TracingContext
-from torch._library.fake_class_registry import FakeScriptObject
+from torch._library.fake_class_registry import FakeScriptObject, maybe_to_fake_obj
+from torch._library.opaque_object import is_opaque_type
 from torch._logging import dtrace_structured
 from torch._subclasses.fake_tensor import FakeTensorMode
 from torch._utils_internal import compile_time_strobelight_meta, log_export_usage
@@ -1629,9 +1630,12 @@ def _strict_export(
                     raise AssertionError(
                         "Cannot find dynamo_fake_mode. This could be due to the exported graph module have no placeholders."
                     )
-                node.meta["val"] = dynamo_fake_mode.from_tensor(
-                    attr, static_shapes=True
-                )
+                if is_opaque_type(type(attr)):
+                    node.meta["val"] = maybe_to_fake_obj(dynamo_fake_mode, attr)
+                else:
+                    node.meta["val"] = dynamo_fake_mode.from_tensor(
+                        attr, static_shapes=True
+                    )
 
     # Fix the graph output signature to be tuple if scalar
     wrap_tuple = False
diff --git a/torch/export/experimental/_utils.py b/torch/export/experimental/_utils.py
index 9684e3530f89f..45c26f7de517d 100644
--- a/torch/export/experimental/_utils.py
+++ b/torch/export/experimental/_utils.py
@@ -196,7 +196,7 @@ def _get_make_file(package_name: str, model_names: list[str], device_type: str)
             "cmake_minimum_required(VERSION 3.10)",
             "project(TestProject)",
             "",
-            "set(CMAKE_CXX_STANDARD 17)",
+            "set(CMAKE_CXX_STANDARD 20)",
             "",
         ]
     )
diff --git a/torch/export/graph_signature.py b/torch/export/graph_signature.py
index 678cd99e40747..e8d4bdbf9634d 100644
--- a/torch/export/graph_signature.py
+++ b/torch/export/graph_signature.py
@@ -2,7 +2,7 @@
 import dataclasses
 from collections.abc import Collection, Mapping
 from enum import auto, Enum
-from typing import TYPE_CHECKING, Union
+from typing import TYPE_CHECKING
 
 from torch._library.fake_class_registry import FakeScriptObject
 from torch._library.opaque_object import get_opaque_type_name, is_opaque_type
@@ -67,15 +67,15 @@ class ConstantArgument:
     value: int | float | bool | str | None
 
 
-ArgumentSpec = Union[
-    TensorArgument,
-    SymIntArgument,
-    SymFloatArgument,
-    SymBoolArgument,
-    ConstantArgument,
-    CustomObjArgument,
-    TokenArgument,
-]
+ArgumentSpec = (
+    TensorArgument
+    | SymIntArgument
+    | SymFloatArgument
+    | SymBoolArgument
+    | ConstantArgument
+    | CustomObjArgument
+    | TokenArgument
+)
 
 
 class InputKind(Enum):
diff --git a/torch/export/pt2_archive/_package.py b/torch/export/pt2_archive/_package.py
index 8d7cb8972cec6..2de31f40b20cb 100644
--- a/torch/export/pt2_archive/_package.py
+++ b/torch/export/pt2_archive/_package.py
@@ -7,6 +7,7 @@
 import zipfile
 from dataclasses import dataclass
 from typing import Any, IO, TYPE_CHECKING, TypeAlias
+from typing_extensions import TypeIs
 
 import torch
 import torch.utils._pytree as pytree
@@ -331,7 +332,7 @@ def _package_aoti_files(
             logger.debug(weights_config)
 
 
-def _is_fake_tensor(t: torch.Tensor) -> bool:
+def _is_fake_tensor(t: torch.Tensor) -> TypeIs[FakeTensor]:
     return isinstance(t, FakeTensor)
 
 
@@ -602,6 +603,8 @@ def _package_exported_programs(
             ep,
             opset_version,
             pickle_protocol,
+            serialize_state_dict=False,
+            serialize_constants=False,
         )
 
         archive_writer.write_bytes(
diff --git a/torch/export/unflatten.py b/torch/export/unflatten.py
index c76c52c37e152..d67c29ed46db4 100644
--- a/torch/export/unflatten.py
+++ b/torch/export/unflatten.py
@@ -1650,7 +1650,10 @@ def _reorder_submodules(
         fqn = prefix + name
         _reorder_submodules(child, fqn_order, prefix=fqn.split("@")[0] + ".")
         delattr(parent, name)
-        children.append((fqn_order[fqn], name, child))
+        base_fqn = fqn.split("@")[0]
+        children.append(
+            (fqn_order.get(fqn, fqn_order.get(base_fqn, len(fqn_order))), name, child)
+        )
     children.sort(key=operator.itemgetter(0))
     for _, name, child in children:
         parent.register_module(name, child)
diff --git a/torch/func/_random.py b/torch/func/_random.py
new file mode 100644
index 0000000000000..a0b6fe6b5ab1f
--- /dev/null
+++ b/torch/func/_random.py
@@ -0,0 +1,263 @@
+"""Stateless PRNG APIs.
+
+These are experimental and subject to change without notice.
+Access via ``torch.func._random``.
+"""
+
+from collections.abc import Sequence
+
+import torch
+
+
+def key(
+    seed: int, impl: str = "philox4x32-10", device: torch.device | None = None
+) -> torch.Tensor:
+    r"""Create a PRNG key from a seed.
+
+    A key is a tensor that encodes the state needed to deterministically
+    produce random values. Keys are consumed by generation functions to produce
+    reproducible random tensors without any global state. The internal
+    representation of the key depends on the chosen PRNG algorithm.
+
+    Args:
+        seed (int): The seed value for the PRNG.
+        impl (str): PRNG algorithm. Currently only ``"philox4x32-10"`` is
+            supported.
+        device (:class:`torch.device`, optional): The desired device for the
+            returned key. Default: ``cpu``.
+
+    Returns:
+        A tensor representing the PRNG key.
+
+    .. note::
+
+        For the ``"philox4x32-10"`` algorithm, the key is a uint64 tensor of
+        shape ``(2,)`` encoding a ``(seed, offset)`` pair. The offset determines
+        the starting position in the Philox output stream.
+
+    Example::
+
+        >>> key = torch.func._random.key(42, device="cuda")  # doctest: +SKIP
+    """
+    if impl != "philox4x32-10":
+        raise NotImplementedError(f"key() does not support PRNG impl '{impl}'")
+
+    # (seed, offset)
+    return torch.tensor([seed, 0], dtype=torch.uint64, device=device)
+
+
+def split(key: torch.Tensor, num: int = 2) -> torch.Tensor:
+    r"""Split a PRNG key into ``num`` new independent keys.
+
+    Each returned key produces a different, deterministic random sequence.
+    This is the primary mechanism for deriving multiple independent keys from
+    a single parent key without mutating any state.
+
+    Supports batched keys: if ``key`` has shape ``(*batch, K)``, each key in the
+    batch is split independently and the result has shape ``(num, *batch, K)``.
+
+    Args:
+        key (Tensor): A PRNG key returned by :func:`key`, :func:`split`, or
+            :func:`fold_in`.
+        num (int): Number of keys to produce. Default: ``2``.
+
+    Returns:
+        A tensor of shape ``(num, *key.shape)`` containing the derived keys.
+
+    Example::
+
+        >>> key = torch.func._random.key(42, device="cuda")  # doctest: +SKIP
+        >>> k1, k2 = torch.func._random.split(key)  # doctest: +SKIP
+    """
+    return torch.ops.aten._philox_key_split(key, num)
+
+
+def fold_in(key: torch.Tensor, data: int) -> torch.Tensor:
+    r"""Deterministically derive a new key by folding in an integer.
+
+    Equivalent to ``split(key, data + 1)[data]``, but more efficient when
+    only a single derived key is needed. Useful for associating a key with
+    a loop iteration, layer index, or other integer identifier.
+
+    Supports batched keys: if ``key`` has shape ``(*batch, K)``, each key in
+    the batch is folded independently.
+
+    Args:
+        key (Tensor): A PRNG key returned by :func:`key`, :func:`split`, or
+            :func:`fold_in`.
+        data (int): An integer to fold into the key, interpreted as uint64.
+
+    Returns:
+        A new key tensor with the same shape as ``key``.
+
+    Example::
+
+        >>> key = torch.func._random.key(42, device="cuda")  # doctest: +SKIP
+        >>> k0 = torch.func._random.fold_in(key, 0)  # doctest: +SKIP
+        >>> k1 = torch.func._random.fold_in(key, 1)  # doctest: +SKIP
+        >>> # Equivalent to split:
+        >>> keys = torch.func._random.split(key, 2)  # doctest: +SKIP
+        >>> assert torch.equal(k0, keys[0])  # doctest: +SKIP
+        >>> assert torch.equal(k1, keys[1])  # doctest: +SKIP
+    """
+    return torch.ops.aten._philox_key_fold_in(key, data)
+
+
+def normal_(
+    key: torch.Tensor,
+    result: torch.Tensor,
+    *,
+    mean: float = 0.0,
+    std: float = 1.0,
+) -> torch.Tensor:
+    r"""Fill ``result`` in-place with normal random values from a PRNG key.
+
+    The values are drawn from a normal distribution with the specified ``mean``
+    and ``std``. The output is fully determined by the key, so calling with the
+    same key always produces the same result.
+
+    Supports batched keys: if ``key`` has shape ``(*batch, K)``, the leading
+    dimensions of ``result`` must be broadcastable with ``*batch`` and each key
+    independently generates its slice of the output.
+
+    Args:
+        key (Tensor): A PRNG key returned by :func:`key`, :func:`split`, or
+            :func:`fold_in`.
+        result (Tensor): The output tensor to fill in-place.
+        mean (float): Mean of the normal distribution. Default: ``0.0``.
+        std (float): Standard deviation of the normal distribution. Default: ``1.0``.
+
+    Returns:
+        ``result``, filled with normal random values.
+
+    Example::
+
+        >>> key = torch.func._random.key(42, device="cuda")  # doctest: +SKIP
+        >>> result = torch.empty(1000, device="cuda")  # doctest: +SKIP
+        >>> torch.func._random.normal_(key, result)  # doctest: +SKIP
+    """
+    return torch.ops.aten._philox_normal_(result, key, mean, std)
+
+
+def normal(
+    key: torch.Tensor,
+    *shape: tuple[int, ...],
+    mean: float = 0.0,
+    std: float = 1.0,
+    dtype: torch.dtype | None = None,
+) -> torch.Tensor:
+    r"""Generate normally distributed random values from a PRNG key.
+
+    Produces a tensor of the given shape filled with values drawn from a normal
+    distribution with the specified ``mean`` and ``std``. The output is fully
+    determined by the key, so calling with the same key always returns the same
+    result. The output is placed on the same device as ``key``.
+
+    Supports batched keys: if ``key`` has shape ``(*batch, K)``, the leading
+    dimensions of ``shape`` must be broadcastable with ``*batch`` and each key
+    independently generates its slice of the output.
+
+    Args:
+        key (Tensor): A PRNG key returned by :func:`key`, :func:`split`, or
+            :func:`fold_in`.
+        *shape (int): The desired output shape.
+        mean (float): Mean of the normal distribution. Default: ``0.0``.
+        std (float): Standard deviation of the normal distribution. Default: ``1.0``.
+        dtype (:class:`torch.dtype`, optional): The desired dtype. Default: ``torch.float32``.
+
+    Returns:
+        A tensor of the given shape filled with normal random values.
+
+    Example::
+
+        >>> key = torch.func._random.key(42, device="cuda")  # doctest: +SKIP
+        >>> torch.func._random.normal(key, (1000,))  # doctest: +SKIP
+    """
+    if len(shape) == 1 and isinstance(shape[0], Sequence):
+        # pyrefly: ignore [bad-argument-type]
+        shape = tuple(shape[0])
+    if dtype is None:
+        dtype = torch.float32
+    # pyrefly: ignore [no-matching-overload]
+    result = torch.empty(shape, dtype=dtype, device=key.device)
+    return normal_(key, result, mean=mean, std=std)
+
+
+def uniform_(
+    key: torch.Tensor,
+    result: torch.Tensor,
+    *,
+    low: float = 0.0,
+    high: float = 1.0,
+) -> torch.Tensor:
+    r"""Fill ``result`` in-place with uniform random values from a PRNG key.
+
+    The values are drawn uniformly from the interval ``[low, high)``. The output
+    is fully determined by the key, so calling with the same key always produces
+    the same result.
+
+    Supports batched keys: if ``key`` has shape ``(*batch, K)``, the leading
+    dimensions of ``result`` must be broadcastable with ``*batch`` and each key
+    independently generates its slice of the output.
+
+    Args:
+        key (Tensor): A PRNG key returned by :func:`key`, :func:`split`, or
+            :func:`fold_in`.
+        result (Tensor): The output tensor to fill in-place.
+        low (float): Lower bound (inclusive) of the uniform distribution. Default: ``0.0``.
+        high (float): Upper bound (exclusive) of the uniform distribution. Default: ``1.0``.
+
+    Returns:
+        ``result``, filled with uniform random values.
+
+    Example::
+
+        >>> key = torch.func._random.key(42, device="cuda")  # doctest: +SKIP
+        >>> result = torch.empty(1000, device="cuda")  # doctest: +SKIP
+        >>> torch.func._random.uniform_(key, result)  # doctest: +SKIP
+    """
+    return torch.ops.aten._philox_uniform_(result, key, low, high)
+
+
+def uniform(
+    key: torch.Tensor,
+    *shape: tuple[int, ...],
+    low: float = 0.0,
+    high: float = 1.0,
+    dtype: torch.dtype | None = None,
+) -> torch.Tensor:
+    r"""Generate uniformly distributed random values from a PRNG key.
+
+    Produces a tensor of the given shape filled with values drawn uniformly
+    from the interval ``[low, high)``. The output is fully determined by the
+    key, so calling with the same key always returns the same result. The output
+    is placed on the same device as ``key``.
+
+    Supports batched keys: if ``key`` has shape ``(*batch, K)``, the leading
+    dimensions of ``shape`` must be broadcastable with ``*batch`` and each key
+    independently generates its slice of the output.
+
+    Args:
+        key (Tensor): A PRNG key returned by :func:`key`, :func:`split`, or
+            :func:`fold_in`.
+        *shape (int): The desired output shape.
+        low (float): Lower bound (inclusive) of the uniform distribution. Default: ``0.0``.
+        high (float): Upper bound (exclusive) of the uniform distribution. Default: ``1.0``.
+        dtype (:class:`torch.dtype`, optional): The desired dtype. Default: ``torch.float32``.
+
+    Returns:
+        A tensor of the given shape filled with uniform random values.
+
+    Example::
+
+        >>> key = torch.func._random.key(42, device="cuda")  # doctest: +SKIP
+        >>> torch.func._random.uniform(key, (1000,))  # doctest: +SKIP
+    """
+    if len(shape) == 1 and isinstance(shape[0], Sequence):
+        # pyrefly: ignore [bad-argument-type]
+        shape = tuple(shape[0])
+    if dtype is None:
+        dtype = torch.float32
+    # pyrefly: ignore [no-matching-overload]
+    result = torch.empty(shape, dtype=dtype, device=key.device)
+    return uniform_(key, result, low=low, high=high)
diff --git a/torch/futures/__init__.py b/torch/futures/__init__.py
index 220a53c4bcbea..6a36da2d5d6c0 100644
--- a/torch/futures/__init__.py
+++ b/torch/futures/__init__.py
@@ -2,7 +2,7 @@
 # pylint: disable=useless-parent-delegation
 from __future__ import annotations
 
-from typing import cast, Generic, Optional, TYPE_CHECKING, TypeVar, Union
+from typing import cast, Generic, TYPE_CHECKING, TypeVar
 
 import torch
 
@@ -27,9 +27,7 @@ class Future(torch._C.Future, Generic[T]):
     .. warning:: GPU support is a beta feature, subject to changes.
     """
 
-    def __init__(
-        self, *, devices: Optional[list[Union[int, str, torch.device]]] = None
-    ):
+    def __init__(self, *, devices: list[int | str | torch.device] | None = None):
         r"""
         Create an empty unset ``Future``. If the future is intended to hold
         values containing CUDA tensors, (a superset of) their CUDA devices must
diff --git a/torch/fx/_compatibility.py b/torch/fx/_compatibility.py
index c07dd1b51bc05..265369cb84c21 100644
--- a/torch/fx/_compatibility.py
+++ b/torch/fx/_compatibility.py
@@ -16,6 +16,7 @@ def compatibility(is_backward_compatible: bool) -> Callable[[_T], _T]:
         def mark_back_compat(fn: _T) -> _T:
             docstring = textwrap.dedent(getattr(fn, "__doc__", None) or "")
             docstring += """
+
 .. note::
     Backwards-compatibility for this API is guaranteed.
 """
@@ -30,6 +31,7 @@ def mark_back_compat(fn: _T) -> _T:
         def mark_not_back_compat(fn: _T) -> _T:
             docstring = textwrap.dedent(getattr(fn, "__doc__", None) or "")
             docstring += """
+
 .. warning::
     This API is experimental and is *NOT* backward-compatible.
 """
diff --git a/torch/fx/_graph_pickler.py b/torch/fx/_graph_pickler.py
index 645635cd43e75..3eb8e74c4f51b 100644
--- a/torch/fx/_graph_pickler.py
+++ b/torch/fx/_graph_pickler.py
@@ -1,11 +1,13 @@
+import contextlib
 import dataclasses
 import importlib
 import io
 import itertools
 import pickle
+import weakref
 from abc import abstractmethod
-from collections.abc import Callable
-from typing import Any, NewType, Optional, TypeVar, Union
+from collections.abc import Callable, Generator
+from typing import Any, NewType, TypeVar
 from typing_extensions import override, Self
 
 from torch.utils._import_utils import import_dill
@@ -19,6 +21,7 @@
 import torch.utils._pytree as pytree
 from torch._guards import TracingContext
 from torch._inductor.standalone_compile import AOTCompiledArtifact
+from torch._library.fake_class_registry import FakeScriptObject
 from torch._subclasses.fake_tensor import FakeTensor, FakeTensorMode, Tensor
 from torch._subclasses.meta_utils import (
     MetaConverter,
@@ -59,10 +62,41 @@ def _node_metadata_key_filter_safe(key: str) -> bool:
 class Options:
     # A filter for which ops will cause the pickler to raise a
     # BypassFxGraphCache exception. If None then all ops are allowed.
-    ops_filter: Optional[Callable[[str], bool]] = _ops_filter_safe
-    node_metadata_key_filter: Optional[Callable[[str], bool]] = (
+    ops_filter: Callable[[str], bool] | None = _ops_filter_safe
+    node_metadata_key_filter: Callable[[str], bool] | None = (
         _node_metadata_key_filter_safe
     )
+    # If True, raw torch.fx.Node objects encountered during pickling will be
+    # silently replaced with None instead of raising an AssertionError.
+    ignore_raw_node: bool = False
+
+
+def _unpickle_as_none() -> None:
+    return None
+
+
+def _unpickle_as_weakref(referent: object) -> weakref.ref[object]:
+    return weakref.ref(referent)
+
+
+def _unpickle_as_dead_weakref() -> Callable[[], None]:
+    return lambda: None
+
+
+@contextlib.contextmanager
+def patch_pytree_map_over_slice() -> Generator[None]:
+    if slice in pytree.SUPPORTED_NODES:
+        yield
+        return
+
+    pytree._private_register_pytree_node(
+        slice, lambda x: ([x.start, x.stop, x.step], None), lambda x, c: slice(*x)
+    )
+
+    try:
+        yield
+    finally:
+        pytree._deregister_pytree_node(slice)
 
 
 # pyrefly: ignore [invalid-inheritance]
@@ -72,7 +106,7 @@ class GraphPickler(pickle.Pickler):
     GraphModule.
     """
 
-    def __init__(self, file: io.BytesIO, options: Optional[Options] = None) -> None:
+    def __init__(self, file: io.BytesIO, options: Options | None = None) -> None:
         if dill is not None:
             super().__init__(file, byref=True)
         else:
@@ -123,9 +157,29 @@ def reducer_override(
             return _SymNodePickleData.reduce_helper(self, obj)
         elif isinstance(obj, torch._guards.TracingContext):
             return _TracingContextPickleData.reduce_helper(self, obj)
+        elif isinstance(obj, FakeScriptObject):
+            from torch._library.opaque_object import is_opaque_value_type
+
+            real_obj = object.__getattribute__(obj, "real_obj")
+            if real_obj is not None and is_opaque_value_type(type(real_obj)):
+                # Use default pickling; value-type opaques are picklable.
+                return NotImplemented
+            # Reference-type FakeScriptObjects can't be default-pickled.
+            return (_unpickle_as_none, ())
+        elif isinstance(obj, weakref.ref):
+            # Serialize weakrefs properly: if the referent is alive,
+            # serialize it and reconstruct the weakref on unpickle.
+            # If the referent is dead, unpickle as a dead-weakref-like callable.
+            referent = obj()
+            if referent is not None:
+                return (_unpickle_as_weakref, (referent,))
+            else:
+                return (_unpickle_as_dead_weakref, ())
         else:
             # We should never get a raw Node!
             if isinstance(obj, torch.fx.Node):
+                if self.options.ignore_raw_node:
+                    return (_unpickle_as_none, ())
                 raise AssertionError("Unexpected raw Node during pickling")
             if reduce := _TorchNumpyPickleData.reduce_helper(self, obj):
                 return reduce
@@ -136,18 +190,18 @@ def reducer_override(
 
     @override
     # pyrefly: ignore [bad-override]
-    def persistent_id(self, obj: object) -> Optional[str]:
+    def persistent_id(self, obj: object) -> str | None:
         if obj is self._unpickle_state:
             return "unpickle_state"
         else:
             return None
 
     @classmethod
-    def dumps(cls, obj: object, options: Optional[Options] = None) -> bytes:
+    def dumps(cls, obj: object, options: Options | None = None) -> bytes:
         """
         Pickle an object.
         """
-        with io.BytesIO() as stream:
+        with patch_pytree_map_over_slice(), io.BytesIO() as stream:
             pickler = cls(stream, options)
             pickler.dump(obj)
             return stream.getvalue()
@@ -157,10 +211,13 @@ def loads(data: bytes, fake_mode: FakeTensorMode) -> object:
         """
         Unpickle an object.
         """
-        state = _UnpickleState(fake_mode)
-        with io.BytesIO(data) as stream:
-            unpickler = _GraphUnpickler(stream, state)
-            return unpickler.load()
+        from torch._dynamo.utils import dynamo_timed
+
+        with patch_pytree_map_over_slice(), dynamo_timed("GraphPickler.loads"):
+            state = _UnpickleState(fake_mode)
+            with io.BytesIO(data) as stream:
+                unpickler = _GraphUnpickler(stream, state)
+                return unpickler.load()
 
     @classmethod
     def debug_dumps(
@@ -171,7 +228,7 @@ def debug_dumps(
         max_depth: int = 80,
         max_iter_items: int = 50,
         verbose: bool = True,
-    ) -> Optional[str]:
+    ) -> str | None:
         """
         Find the first leaf that GraphPickler.dumps cannot serialize and return its path.
 
@@ -201,14 +258,14 @@ def log(msg: str) -> None:
             if verbose:
                 print(msg)
 
-        def fail_exc(o: Any) -> Optional[BaseException]:
+        def fail_exc(o: Any) -> BaseException | None:
             try:
                 cls.dumps(o, options)
                 return None
             except Exception as e:
                 return e
 
-        def walk(o: Any, path: str, depth: int) -> Optional[str]:
+        def walk(o: Any, path: str, depth: int) -> str | None:
             if depth > max_depth:
                 log(f"{'  ' * depth}Depth limit at {path} ({type(o)})")
                 return path + " (depth_limit)"
@@ -497,7 +554,7 @@ def unpickle(self, unpickle_state: _UnpickleState) -> FakeTensor:
             metadata = dataclasses.replace(metadata, base=new_base)
 
         def with_fake(
-            make_meta_t: Callable[[], torch.Tensor], device: Union[torch.device, str]
+            make_meta_t: Callable[[], torch.Tensor], device: torch.device | str
         ) -> FakeTensor:
             with no_dispatch():
                 return FakeTensor(
@@ -520,11 +577,12 @@ class _TorchNumpyPickleData:
     @classmethod
     def reduce_helper(
         cls, pickler: GraphPickler, obj: object
-    ) -> Optional[
+    ) -> (
         tuple[
             Callable[[Self, _UnpickleState], object], tuple[Self, _UnpickleStateToken]
         ]
-    ]:
+        | None
+    ):
         if data := cls.from_object(obj):
             return (cls.unpickle, (data, pickler._unpickle_state))
         else:
@@ -539,7 +597,7 @@ def unpickle(self, unpickle_state: _UnpickleState) -> Callable[..., object]:
         return torch._dynamo.variables.misc.get_np_to_tnp_map()[np]
 
     @classmethod
-    def from_object(cls, tnp: object) -> Optional[Self]:
+    def from_object(cls, tnp: object) -> Self | None:
         if not callable(tnp):
             return None
 
@@ -685,9 +743,7 @@ def pickle(cls, op: object, options: Options) -> "_OpPickleData":
     @staticmethod
     def _pickle_op(
         name: str,
-        datacls: Union[
-            type["_OpOverloadPickleData"], type["_OpOverloadPacketPickleData"]
-        ],
+        datacls: type["_OpOverloadPickleData"] | type["_OpOverloadPacketPickleData"],
         options: Options,
     ) -> "_OpPickleData":
         if (ops_filter := options.ops_filter) and not ops_filter(name):
diff --git a/torch/fx/_lazy_graph_module.py b/torch/fx/_lazy_graph_module.py
index 7808dec3cee13..20a7837a66618 100644
--- a/torch/fx/_lazy_graph_module.py
+++ b/torch/fx/_lazy_graph_module.py
@@ -1,5 +1,6 @@
-# mypy: allow-untyped-defs
+from collections.abc import Iterator
 from contextlib import contextmanager
+from typing import Any, TYPE_CHECKING
 
 from torch.fx.graph_module import (
     _format_import_block,
@@ -12,13 +13,17 @@
 from ._compatibility import compatibility
 
 
+if TYPE_CHECKING:
+    from torch.fx.graph import PythonCode
+
+
 _use_lazy_graph_module_flag = False
 _force_skip_lazy_graph_module_flag = False
 
 
 @compatibility(is_backward_compatible=False)
 @contextmanager
-def _force_skip_lazy_graph_module():
+def _force_skip_lazy_graph_module() -> Iterator[None]:
     """
     Skip using lazy graph module disregarding the setting of _use_lazy_graph_module.
     Use to skip _LazyGraphModule when testing inductor torchscript related backend.
@@ -37,7 +42,7 @@ def _force_skip_lazy_graph_module():
 
 @compatibility(is_backward_compatible=False)
 @contextmanager
-def _use_lazy_graph_module(should_use: bool):
+def _use_lazy_graph_module(should_use: bool) -> Iterator[None]:
     try:
         global _use_lazy_graph_module_flag
         prior = _use_lazy_graph_module_flag
@@ -50,11 +55,13 @@ def _use_lazy_graph_module(should_use: bool):
 
 
 @compatibility(is_backward_compatible=False)
-def _get_graph_module_cls():
+def _get_graph_module_cls() -> type[GraphModule]:
     return _LazyGraphModule if _use_lazy_graph_module_flag else GraphModule
 
 
-def _make_graph_module(*args, graph_module_cls=None, **kwargs):
+def _make_graph_module(
+    *args: Any, graph_module_cls: type[GraphModule] | None = None, **kwargs: Any
+) -> GraphModule:
     if graph_module_cls is None:
         graph_module_cls = _get_graph_module_cls()
 
@@ -88,14 +95,14 @@ class _LazyGraphModule(GraphModule):
     """
 
     @classmethod
-    def from_graphmodule(cls, gm: GraphModule):
+    def from_graphmodule(cls, gm: GraphModule) -> GraphModule:
         if isinstance(gm, _LazyGraphModule):
             return gm
         else:
             return _LazyGraphModule(gm, gm.graph)
 
     @staticmethod
-    def force_recompile(gm):
+    def force_recompile(gm: GraphModule) -> None:
         """
         Sometimes we need force a recompile as a workaround
         - we want to do the real recompilation before symbolic_trace to avoid error:
@@ -104,15 +111,15 @@ def force_recompile(gm):
         if isinstance(gm, _LazyGraphModule):
             gm.real_recompile()
 
-    def real_recompile(self):
+    def real_recompile(self) -> None:
         if self._needs_recompile():
             self._real_recompile()
 
     @classmethod
-    def _needs_recompile(cls):
+    def _needs_recompile(cls) -> bool:
         return cls.forward is cls._lazy_forward
 
-    def _lazy_forward(self, *args, **kwargs):
+    def _lazy_forward(self, *args: Any, **kwargs: Any) -> Any:
         # Call self.real_recompile() rather than self._real_recompile() here.
         # The _lazy_forward method may be saved and call repeatedly.
         # Calling self.real_recompile can make sure we skip recompilation if
@@ -128,7 +135,9 @@ def _lazy_forward(self, *args, **kwargs):
 
     forward = _lazy_forward
 
-    def __reduce_package__(self, exporter: PackageExporter):
+    def __reduce_package__(
+        self, exporter: PackageExporter
+    ) -> tuple[Any, tuple[Any, str]]:
         """
         Follow GraphModule.__reduce__ but call 'self._real_recompile' rather
         than 'self.recompile' since for a _LazyGraphModule, self.recompile just
@@ -148,7 +157,7 @@ def __reduce_package__(self, exporter: PackageExporter):
             (dict_without_graph, generated_module_name),
         )
 
-    def __reduce__(self):
+    def __reduce__(self) -> tuple[Any, tuple[Any, str]]:
         """
         Follow GraphModule.__reduce__ but call 'self._real_recompile' rather
         than 'self.recompile' since for a _LazyGraphModule, self.recompile just
@@ -160,11 +169,11 @@ def __reduce__(self):
         del dict_without_graph["_graph"]
         return (reduce_graph_module, (dict_without_graph, import_block))
 
-    def _real_recompile(self):
+    def _real_recompile(self) -> "PythonCode":
         return super().recompile()
 
     @classmethod
-    def recompile(cls):
+    def recompile(cls) -> None:  # pyrefly: ignore[bad-override]
         cls.forward = cls._lazy_forward
 
     @property
diff --git a/torch/fx/_pytree.py b/torch/fx/_pytree.py
index bfb62f871eb6f..4d923d528e710 100644
--- a/torch/fx/_pytree.py
+++ b/torch/fx/_pytree.py
@@ -1,17 +1,21 @@
 from collections import namedtuple
 from collections.abc import Callable
-from typing import Any, Optional, TypeVar
+from typing import Any, TypeVar
 from typing_extensions import NamedTuple
 
 import torch.return_types
 from torch.utils._pytree import PyTree, tree_flatten, TreeSpec
 
 
-FlattenFuncSpec = Callable[[PyTree, TreeSpec], list]
-FlattenFuncExactMatchSpec = Callable[[PyTree, TreeSpec], bool]
+FlattenFnSpec = Callable[[PyTree, TreeSpec], list[Any]]
+FlattenFnExactMatchSpec = Callable[[PyTree, TreeSpec], bool]
 
-SUPPORTED_NODES: dict[type[Any], FlattenFuncSpec] = {}
-SUPPORTED_NODES_EXACT_MATCH: dict[type[Any], Optional[FlattenFuncExactMatchSpec]] = {}
+# Keep deprecated alias for backward compatibility
+FlattenFuncSpec = FlattenFnSpec  # deprecated
+FlattenFuncExactMatchSpec = FlattenFnExactMatchSpec  # deprecated
+
+SUPPORTED_NODES: dict[type[Any], FlattenFnSpec] = {}
+SUPPORTED_NODES_EXACT_MATCH: dict[type[Any], FlattenFnExactMatchSpec | None] = {}
 
 _T = TypeVar("_T")
 _K = TypeVar("_K")
@@ -20,8 +24,8 @@
 
 def register_pytree_flatten_spec(
     cls: type[Any],
-    flatten_fn_spec: FlattenFuncSpec,
-    flatten_fn_exact_match_spec: Optional[FlattenFuncExactMatchSpec] = None,
+    flatten_fn_spec: FlattenFnSpec,
+    flatten_fn_exact_match_spec: FlattenFnExactMatchSpec | None = None,
 ) -> None:
     SUPPORTED_NODES[cls] = flatten_fn_spec
     SUPPORTED_NODES_EXACT_MATCH[cls] = flatten_fn_exact_match_spec
@@ -48,7 +52,7 @@ def tree_flatten_spec(
     if spec.type in SUPPORTED_NODES:
         flatten_fn_spec = SUPPORTED_NODES[spec.type]
         child_pytrees = flatten_fn_spec(pytree, spec)
-        result = []
+        result: list[Any] = []
         for child, child_spec in zip(child_pytrees, spec.children()):
             flat = tree_flatten_spec(child, child_spec)
             result += flat
diff --git a/torch/fx/_symbolic_trace.py b/torch/fx/_symbolic_trace.py
index 5fbac2b4086a9..3a5541811cb6e 100644
--- a/torch/fx/_symbolic_trace.py
+++ b/torch/fx/_symbolic_trace.py
@@ -1,4 +1,3 @@
-# mypy: allow-untyped-defs
 import builtins
 import collections
 import contextlib
@@ -9,10 +8,10 @@
 import math
 import os
 import warnings
-from collections.abc import Callable
+from collections.abc import Callable, Iterable, Iterator
 from itertools import chain
-from types import CodeType, FunctionType, ModuleType
-from typing import Any, get_args, NamedTuple, Optional, TypeAlias, Union
+from types import CodeType, FunctionType, ModuleType, TracebackType
+from typing import Any, get_args, NamedTuple, overload, ParamSpec, TypeAlias, TypeVar
 
 import torch
 import torch.utils._pytree as pytree
@@ -32,24 +31,28 @@
 
 HAS_VARSTUFF = inspect.CO_VARARGS | inspect.CO_VARKEYWORDS
 
+_F = TypeVar("_F", bound=Callable[..., Any])
+_P = ParamSpec("_P")
+_T = TypeVar("_T")
+
 # These need to run in global scope to handle nested calls correctly
-_orig_module_call: Callable = torch.nn.Module.__call__
-_orig_module_getattr: Callable = torch.nn.Module.__getattr__
+_orig_module_call: Callable[..., Any] = torch.nn.Module.__call__
+_orig_module_getattr: Callable[..., Any] = torch.nn.Module.__getattr__
 
 _proxyable_classes: dict[type, None] = {}
 
 _is_fx_tracing_flag = False
 
-_ConstantAttributeType: TypeAlias = Union[
-    torch.Tensor, torch.ScriptObject, FakeScriptObject, pytree.TreeSpec
-]
+_ConstantAttributeType: TypeAlias = (
+    torch.Tensor | torch.ScriptObject | FakeScriptObject | pytree.TreeSpec
+)
 
 _constant_attribute_types = get_args(_ConstantAttributeType)
 
 
 # We only want to print this once to avoid flooding logs
 @functools.lru_cache
-def is_fx_tracing_warning():
+def is_fx_tracing_warning() -> None:
     log.warning(
         "is_fx_tracing will return true for both fx.symbolic_trace and "
         "torch.export. Please use "
@@ -58,12 +61,12 @@ def is_fx_tracing_warning():
     )
 
 
-def is_fx_tracing():
+def is_fx_tracing() -> bool:
     is_fx_tracing_warning()
     return _is_fx_tracing_flag
 
 
-def is_fx_symbolic_tracing():
+def is_fx_symbolic_tracing() -> bool:
     return _is_fx_tracing_flag and not torch.compiler.is_compiling()
 
 
@@ -116,20 +119,22 @@ def forward(self, x : __main___TensorPair, y : torch.Tensor):
     tracing.
     """
 
-    def __init__(cls, name, bases, attrs):
+    def __init__(
+        cls, name: str, bases: tuple[type, ...], attrs: dict[str, Any]
+    ) -> None:
         _proxyable_classes.setdefault(cls)
         super().__init__(name, bases, attrs)
 
-    def __call__(cls, *args, **kwargs):
+    def __call__(cls, *args: Any, **kwargs: Any) -> Any:
         instance = cls.__new__(cls)  # type: ignore[call-overload]
 
         if not is_fx_tracing():
             cls.__init__(instance, *args, **kwargs)  # type: ignore[misc]
             return instance
 
-        found_proxies = []
+        found_proxies: list[Proxy] = []
 
-        def check_proxy(a):
+        def check_proxy(a: object) -> None:
             if isinstance(a, Proxy):
                 found_proxies.append(a)
 
@@ -147,7 +152,7 @@ def check_proxy(a):
 def _patch_function(fn: FunctionType, nargs: int) -> FunctionType:
     co = fn.__code__
     co_flags = co.co_flags & ~HAS_VARSTUFF
-    co_args: tuple
+    co_args: tuple[Any, ...]
     if hasattr(co, "co_qualname"):
         # Python-3.11+ code signature
         co_args = (
@@ -223,7 +228,7 @@ class PHBase:
     Object representing an input placeholder to `concrete_args`
     """
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         return "PH"
 
 
@@ -236,14 +241,14 @@ class PHWithMeta(PHBase):
     Object representing an input placeholder to `concrete_args`
     """
 
-    def __init__(self, ph_key: Optional[str] = None):
+    def __init__(self, ph_key: str | None = None) -> None:
         super().__init__()
 
         # Provide a hey for user to identify placeholder node during analysis
         self.ph_key = ph_key
 
 
-def _transfer_attrs(fr, to):
+def _transfer_attrs(fr: object, to: object) -> None:
     for attr_name in dir(fr):
         attr_val = getattr(fr, attr_name)
         if (
@@ -279,7 +284,7 @@ class Tracer(TracerBase):
     def __init__(
         self,
         autowrap_modules: tuple[ModuleType] = (math,),
-        autowrap_functions: tuple[Callable, ...] = (),
+        autowrap_functions: tuple[Callable[..., Any], ...] = (),
         param_shapes_constant: bool = False,
     ) -> None:
         # This method's signature is overridden by the first line of this class'
@@ -326,7 +331,7 @@ def __init__(
         self._autowrap_search: list[ModuleType] = list(autowrap_modules)
         self.param_shapes_constant = param_shapes_constant
 
-        self.submodule_paths: Optional[dict[torch.nn.Module, str]] = None
+        self.submodule_paths: dict[torch.nn.Module, str] | None = None
         self.root_module_name: str = ""
         # Maps the containing module's name to the operator name
         self.scope = Scope("", None)
@@ -425,7 +430,7 @@ def create_arg(self, a: Any) -> "Argument":
         if isinstance(a, _constant_attribute_types) or (
             is_opaque_reference_type(type(a))
         ):
-            qualname: Optional[str] = self.tensor_attrs.get(a)
+            qualname: str | None = self.tensor_attrs.get(a)
 
             # Tensor was not found in the Module hierarchy, stow it away in a
             # special attribute and set the qualname to refer to that
@@ -581,7 +586,9 @@ def call_module(
         return ret_val
 
     @compatibility(is_backward_compatible=False)
-    def getattr(self, attr: str, attr_val: Any, parameter_proxy_cache: dict[str, Any]):
+    def getattr(
+        self, attr: str, attr_val: Any, parameter_proxy_cache: dict[str, Proxy]
+    ) -> Any:
         """
         Method that specifies the behavior of this ``Tracer`` when we call getattr
         on a call to an ``nn.Module`` instance.
@@ -605,8 +612,10 @@ def getattr(self, attr: str, attr_val: Any, parameter_proxy_cache: dict[str, Any
         """
 
         def maybe_get_proxy_for_attr(
-            attr_val, collection_to_search, parameter_proxy_cache
-        ):
+            attr_val: Any,
+            collection_to_search: Iterable[tuple[str, Any]],
+            parameter_proxy_cache: dict[str, Proxy],
+        ) -> Proxy | None:
             for n, p in collection_to_search:
                 if attr_val is p:
                     if n not in parameter_proxy_cache:
@@ -646,7 +655,12 @@ def maybe_get_proxy_for_attr(
 
     # This method will be refactored
     @compatibility(is_backward_compatible=False)
-    def create_args_for_root(self, root_fn, is_module, concrete_args=None):
+    def create_args_for_root(
+        self,
+        root_fn: Callable[..., Any],
+        is_module: bool,
+        concrete_args: dict[str, Any] | tuple[Any, ...] | None = None,
+    ) -> tuple[Any, list[Any]]:
         """
         Create ``placeholder`` nodes corresponding to the signature of the ``root``
         Module. This method introspects root's signature and emits those
@@ -708,7 +722,7 @@ def create_args_for_root(self, root_fn, is_module, concrete_args=None):
                 )
             concrete_args = dict(zip(arg_names, concrete_args))
 
-        def proxy_placeholder(name):
+        def proxy_placeholder(name: str) -> Any:
             return self._proxy_placeholder(name, concrete_args, sig, fn_for_analysis)
 
         args.extend(proxy_placeholder(names) for names in arg_names)
@@ -730,7 +744,9 @@ def proxy_placeholder(name):
                 _PyTreeInfo(orig_args[:total_args], in_spec, None)
             )
 
-            def flatten_fn(*args):
+            # TODO: annotate return type. inspect.get_annotations(flatten_fn)
+            # leaks the return annotation into the generated forward() code.
+            def flatten_fn(*args: Any):  # pyrefly: ignore[unannotated-parameter]
                 tree_args = pytree.tree_unflatten(list(args), in_spec)
                 tree_out = root_fn(*tree_args)
                 out_args, out_spec = pytree.tree_flatten(tree_out)
@@ -750,8 +766,8 @@ def flatten_fn(*args):
     @compatibility(is_backward_compatible=True)
     def trace(
         self,
-        root: Union[torch.nn.Module, Callable[..., Any]],
-        concrete_args: Optional[dict[str, Any]] = None,
+        root: torch.nn.Module | Callable[..., Any],
+        concrete_args: dict[str, Any] | None = None,
     ) -> Graph:
         """
         Trace ``root`` and return the corresponding FX ``Graph`` representation. ``root``
@@ -787,7 +803,9 @@ def trace(
                 # without this.
                 from torch.fx._lazy_graph_module import _LazyGraphModule
 
-                _LazyGraphModule.force_recompile(root)
+                _LazyGraphModule.force_recompile(
+                    root  # pyrefly: ignore[bad-argument-type]
+                )
 
                 self.root = root
 
@@ -804,7 +822,7 @@ def trace(
                 self.root = torch.nn.Module()
                 fn = root
 
-            tracer_cls: Optional[type[Tracer]] = getattr(self, "__class__", None)
+            tracer_cls: type[Tracer] | None = getattr(self, "__class__", None)
             self.graph = Graph(tracer_cls=tracer_cls)
             if hasattr(fn, "__code__"):
                 code = fn.__code__
@@ -823,7 +841,9 @@ def trace(
                 str,
             ] = {}
 
-            def collect_tensor_attrs(m: torch.nn.Module, prefix_atoms: list[str]):
+            def collect_tensor_attrs(
+                m: torch.nn.Module, prefix_atoms: list[str]
+            ) -> None:
                 for k, v in m.__dict__.items():
                     if isinstance(v, _constant_attribute_types):
                         self.tensor_attrs[v] = ".".join(prefix_atoms + [k])
@@ -847,13 +867,15 @@ def collect_tensor_attrs(m: torch.nn.Module, prefix_atoms: list[str]):
             # Method dispatch on parameters is not recorded unless it's directly used.
             # Thus, we need to insert a proxy when __getattr__ requests a parameter.
             @functools.wraps(_orig_module_getattr)
-            def module_getattr_wrapper(mod, attr):
+            def module_getattr_wrapper(mod: torch.nn.Module, attr: str) -> Any:
                 attr_val = _orig_module_getattr(mod, attr)
                 return self.getattr(attr, attr_val, parameter_proxy_cache)
 
             @functools.wraps(_orig_module_call)
-            def module_call_wrapper(mod, *args, **kwargs):
-                def forward(*args, **kwargs):
+            def module_call_wrapper(
+                mod: torch.nn.Module, *args: Any, **kwargs: Any
+            ) -> Any:
+                def forward(*args: Any, **kwargs: Any) -> Any:
                     return _orig_module_call(mod, *args, **kwargs)
 
                 _autowrap_check(
@@ -907,7 +929,7 @@ def forward(*args, **kwargs):
             _is_fx_tracing_flag = old_is_fx_tracing_flag
         return self.graph
 
-    def __deepcopy__(self, memo):
+    def __deepcopy__(self, memo: dict[int, Any]) -> "Tracer":
         # _autowrap_search contains modules, which cannot be deepcopied.
         new_tracer = Tracer.__new__(Tracer)
 
@@ -921,11 +943,17 @@ def __deepcopy__(self, memo):
 
         return new_tracer
 
-    def _proxy_placeholder(self, name, concrete_args, sig, fn_for_analysis):
+    def _proxy_placeholder(
+        self,
+        name: str,
+        concrete_args: dict[str, Any] | None,
+        sig: inspect.Signature,
+        fn_for_analysis: Callable[..., Any],
+    ) -> Any:
         if concrete_args is not None and name in concrete_args:
             cnt = 0
 
-            def replace_ph(x):
+            def replace_ph(x: object) -> object:
                 nonlocal cnt
                 cnt += 1
                 param = sig.parameters[name]
@@ -993,7 +1021,7 @@ def replace_ph(x):
 # the purposes of the wrap() API.
 # We key by the globals dict id and function name to ensure we're wrapping a given
 # function only once.
-_wrapped_fns_to_patch: dict[tuple[int, str], dict] = {}
+_wrapped_fns_to_patch: dict[tuple[int, str], dict[str, Any]] = {}
 
 # List of methods on classes to wrap (class type, function name)
 # this currently only works for Tensor.* methods that aren't traced properly
@@ -1008,25 +1036,26 @@ def replace_ph(x):
     _wrapped_methods_to_patch.append((torch.Tensor, "__getitem__"))
 
 
-def _find_proxy(*objects_to_search):
+def _find_proxy(*objects_to_search: object) -> Proxy | None:
     """
     Recursively search a data structure for a Proxy() and return it,
     return None if not found.
     """
     proxy = None
 
-    def find_proxy(x):
+    def find_proxy(x: object) -> None:
         nonlocal proxy
         if isinstance(x, Proxy):
             proxy = x
 
+    # pyrefly: ignore[bad-specialization]
     map_aggregate(objects_to_search, find_proxy)
     return proxy
 
 
-def _create_wrapped_func(orig_fn):
+def _create_wrapped_func(orig_fn: Callable[_P, _T]) -> Callable[_P, _T]:
     @functools.wraps(orig_fn)
-    def wrapped(*args, **kwargs):
+    def wrapped(*args: _P.args, **kwargs: _P.kwargs) -> Any:
         """
         Given an closed-over ``orig_function`` to invoke, search the args and kwargs for
         a Proxy object. If there is one, emit a ``call_function`` node to preserve the
@@ -1045,11 +1074,11 @@ def wrapped(*args, **kwargs):
     return wrapped
 
 
-def _create_wrapped_method(cls, name):
+def _create_wrapped_method(cls: type, name: str) -> Callable[..., Any]:
     orig_fn = getattr(cls, name)
 
     @functools.wraps(orig_fn)
-    def wrapped(*args, **kwargs):
+    def wrapped(*args: Any, **kwargs: Any) -> Any:
         """
         Search the args and kwargs for a Proxy object. If there is one,
         emit a ``call_method`` node to preserve the call to this method
@@ -1065,39 +1094,39 @@ def wrapped(*args, **kwargs):
 
 
 class _PatchedFn(NamedTuple):
-    frame_dict: Any
+    frame_dict: Any  # dict[str, Any] for SetItem/Del, type for SetAttr
     fn_name: str
-    orig_fn: Any
-    new_fn: Any
+    orig_fn: Callable[..., Any] | None
+    new_fn: Callable[..., Any]
 
-    def revert(self):
+    def revert(self) -> None:
         raise NotImplementedError
 
-    def patch(self):
+    def patch(self) -> None:
         raise NotImplementedError
 
 
 class _PatchedFnSetItem(_PatchedFn):
-    def revert(self):
+    def revert(self) -> None:
         self.frame_dict[self.fn_name] = self.orig_fn
 
-    def patch(self):
+    def patch(self) -> None:
         self.frame_dict[self.fn_name] = self.new_fn
 
 
 class _PatchedFnDel(_PatchedFn):
-    def revert(self):
+    def revert(self) -> None:
         del self.frame_dict[self.fn_name]
 
-    def patch(self):
+    def patch(self) -> None:
         self.frame_dict[self.fn_name] = self.new_fn
 
 
 class _PatchedFnSetAttr(_PatchedFn):
-    def revert(self):
+    def revert(self) -> None:
         setattr(self.frame_dict, self.fn_name, self.orig_fn)
 
-    def patch(self):
+    def patch(self) -> None:
         setattr(self.frame_dict, self.fn_name, self.new_fn)
 
 
@@ -1111,9 +1140,9 @@ def patch(
         self,
         frame_dict: dict[str, Any],
         name: str,
-        new_fn: Callable,
+        new_fn: Callable[..., Any],
         deduplicate: bool = True,
-    ):
+    ) -> None:
         """
         Replace frame_dict[name] with new_fn until we exit the context manager.
         """
@@ -1130,8 +1159,8 @@ def patch(
             self.patches_made[-1].patch()
 
     def patch_method(
-        self, cls: type, name: str, new_fn: Callable, deduplicate: bool = True
-    ):
+        self, cls: type, name: str, new_fn: Callable[..., Any], deduplicate: bool = True
+    ) -> None:
         """
         Replace object_or_dict.name with new_fn until we exit the context manager.
         """
@@ -1142,7 +1171,7 @@ def patch_method(
         self.patches_made.append(_PatchedFnSetAttr(cls, name, orig_fn, new_fn))
         self.patches_made[-1].patch()
 
-    def visit_once(self, thing: Any):
+    def visit_once(self, thing: object) -> bool:
         """Return True on the first call to with thing, otherwise false"""
         idx = id(thing)
         if idx in self.visited:
@@ -1150,7 +1179,7 @@ def visit_once(self, thing: Any):
         self.visited.add(idx)
         return True
 
-    def revert_all_patches(self):
+    def revert_all_patches(self) -> list[_PatchedFn]:
         """
         Remove all the stored patcheds. It doesn't modify patches_made.
         """
@@ -1158,7 +1187,7 @@ def revert_all_patches(self):
             patch.revert()
         return self.patches_made
 
-    def reapply_all_patches(self):
+    def reapply_all_patches(self) -> list[_PatchedFn]:
         """
         Patch all the stored patcheds. It doesn't modify patches_made.
         """
@@ -1166,10 +1195,15 @@ def reapply_all_patches(self):
             patch.patch()
         return self.patches_made
 
-    def __enter__(self):
+    def __enter__(self) -> "_Patcher":
         return self
 
-    def __exit__(self, exc_type, exc_val, exc_tb):
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: TracebackType | None,
+    ) -> None:
         """
         Undo all the changes made via self.patch() and self.patch_method()
         """
@@ -1179,11 +1213,11 @@ def __exit__(self, exc_type, exc_val, exc_tb):
         self.visited.clear()
 
 
-CURRENT_PATCHER: Optional[_Patcher] = None
+CURRENT_PATCHER: _Patcher | None = None
 
 
 @contextlib.contextmanager
-def _new_patcher():
+def _new_patcher() -> Iterator[_Patcher]:
     global CURRENT_PATCHER
     prior_patcher = CURRENT_PATCHER
     try:
@@ -1198,7 +1232,7 @@ def _new_patcher():
 
 
 @contextlib.contextmanager
-def _maybe_revert_all_patches():
+def _maybe_revert_all_patches() -> Iterator[None]:
     current_patcher = CURRENT_PATCHER
     patches_made = None
     patches_removed = None
@@ -1215,7 +1249,7 @@ def _maybe_revert_all_patches():
             )
 
 
-def _patch_wrapped_functions(patcher: _Patcher):
+def _patch_wrapped_functions(patcher: _Patcher) -> None:
     """
     Go through ``_wrapped_fn_patch_table`` and, for each frame object, wrap
     the listed global functions in the `_create_wrapped_func` wrapper.
@@ -1233,7 +1267,7 @@ def _patch_wrapped_functions(patcher: _Patcher):
 
 def _autowrap_check(
     patcher: _Patcher, frame_dict: dict[str, Any], function_ids: set[int]
-):
+) -> None:
     """
     Some methods, like `math.sqrt` are common enough we want to automatically wrap them as we see them.
     This method searches a scope for them and patches them if found.
@@ -1248,8 +1282,12 @@ def _autowrap_check(
                 patcher.patch(frame_dict, name, _create_wrapped_func(value))
 
 
+@overload
+def wrap(fn_or_name: _F) -> _F: ...
+@overload
+def wrap(fn_or_name: str) -> str: ...
 @compatibility(is_backward_compatible=True)
-def wrap(fn_or_name: Union[str, Callable]):
+def wrap(fn_or_name: str | Callable[..., Any]) -> str | Callable[..., Any]:
     """
     This function can be called at module-level scope to register fn_or_name as a "leaf function".
     A "leaf function" will be preserved as a CallFunction node in the FX trace instead of being
@@ -1319,8 +1357,8 @@ def my_custom_function(x, y):
 
 @compatibility(is_backward_compatible=True)
 def symbolic_trace(
-    root: Union[torch.nn.Module, Callable[..., Any]],
-    concrete_args: Optional[dict[str, Any]] = None,
+    root: torch.nn.Module | Callable[..., Any],
+    concrete_args: dict[str, Any] | None = None,
 ) -> GraphModule:
     """
     Symbolic tracing API
@@ -1382,6 +1420,6 @@ def f(x):
 
 
 @wrap
-def _assert_is_none(value, msg):
+def _assert_is_none(value: object, msg: str) -> None:
     if value is not None:
         raise AssertionError(msg)
diff --git a/torch/fx/_utils.py b/torch/fx/_utils.py
index 25f1c51171734..c58c36502b68d 100644
--- a/torch/fx/_utils.py
+++ b/torch/fx/_utils.py
@@ -1,17 +1,18 @@
-# mypy: allow-untyped-defs
 import sys
-from typing import Optional
+from typing import Any
 
 import torch
 from torch._logging import LazyString
 
 
-def lazy_format_graph_code(name, gm, maybe_id=None, **kwargs):
+def lazy_format_graph_code(
+    name: str, gm: torch.fx.GraphModule, maybe_id: int | None = None, **kwargs: Any
+) -> LazyString:
     """
     Returns a LazyString that formats the graph code.
     """
 
-    def format_name():
+    def format_name() -> str:
         if maybe_id is not None:
             return f"{name} {maybe_id}"
         else:
@@ -36,14 +37,14 @@ def format_name():
     )
 
 
-def _format_graph_code(name, filename, graph_str):
+def _format_graph_code(name: str, filename: str, graph_str: str) -> str:
     """
     Returns a string that formats the graph code.
     """
     return f"TRACED GRAPH\n {name} {filename} {graph_str}\n"
 
 
-def first_call_function_nn_module_stack(graph: torch.fx.Graph) -> Optional[dict]:
+def first_call_function_nn_module_stack(graph: torch.fx.Graph) -> dict[str, Any] | None:
     """
     Returns the nn_module_stack of the first call_function node.
     """
@@ -53,14 +54,15 @@ def first_call_function_nn_module_stack(graph: torch.fx.Graph) -> Optional[dict]
     return None
 
 
-def get_node_context(node, num_nodes=2) -> str:
+def get_node_context(node: torch.fx.Node, num_nodes: int = 2) -> str:
     """
     Returns a string of the last num_nodes nodes in the graph.
     """
     node_contexts = []
     cur = node
     for _ in range(num_nodes):
-        node_contexts.append(cur.format_node())
+        # cast to str to handle None return value
+        node_contexts.append(str(cur.format_node()))
         if cur.op == "root":
             break
         cur = cur.prev
diff --git a/torch/fx/annotate.py b/torch/fx/annotate.py
index b3c5056066251..fd90605d8e3b5 100644
--- a/torch/fx/annotate.py
+++ b/torch/fx/annotate.py
@@ -1,11 +1,15 @@
-# mypy: allow-untyped-defs
+from typing import Any
+
 from torch.fx.proxy import Proxy
 
 from ._compatibility import compatibility
 
 
+__all__ = ["annotate"]
+
+
 @compatibility(is_backward_compatible=False)
-def annotate(val, type):
+def annotate(val: Any, type: type) -> Any:
     """
     Annotates a Proxy object with a given type.
 
diff --git a/torch/fx/config.py b/torch/fx/config.py
index db06176c43e13..9ebdb86499174 100644
--- a/torch/fx/config.py
+++ b/torch/fx/config.py
@@ -4,3 +4,8 @@
 
 # If True this also shows the node names in each pass, for small models this is great but larger models it's quite noisy
 verbose_progress = False
+
+# When True, skip collecting stack traces during tracing. This avoids the cost
+# of CapturedTraceback.extract() and symbolization for every FX node, but means
+# node.meta["stack_trace"] will be unset, degrading error messages and debugging info.
+do_not_emit_stack_traces = False
diff --git a/torch/fx/experimental/_config.py b/torch/fx/experimental/_config.py
index 1f2a534794638..2bea364341ba7 100644
--- a/torch/fx/experimental/_config.py
+++ b/torch/fx/experimental/_config.py
@@ -1,6 +1,6 @@
+import enum
 import os
 import sys
-from typing import Optional
 
 from torch.utils._config_module import Config, install_config_module
 
@@ -83,7 +83,7 @@
 # issued (as we test if we've hit the limit on-the-fly, whereas we may
 # do further simplifications at final guard issuance time that make guards
 # irrelevant.)
-symbol_guard_limit_before_specialize: Optional[int] = None
+symbol_guard_limit_before_specialize: int | None = None
 
 # This flag changes whether we should use the same symbolic variable to represent input sizes that are the same.
 use_duck_shape = True
@@ -115,5 +115,20 @@
 # code.
 soft_pending_unbacked_not_found_error = False
 
+# When True, aggressively return fallback values in guard_or opting into
+# guard-free semantics. This optimizes tracing time when symbolic reasoning
+# is expensive. Since guard_or_X already have a general path to take, we
+# can skip expensive static evaluation and just return the fallback value directly.
+# This is usually safe because the fallback represents a valid code path that
+# could be taken anyway.
+# See AggressiveGuardFreeMode below for valid values.
+aggressive_guard_free_semantics = 0
+
 
 install_config_module(sys.modules[__name__])
+
+
+class AggressiveGuardFreeMode(enum.IntEnum):
+    DISABLED = 0
+    VALUE_RANGE_ANALYSIS = 1  # use bound_sympy before returning fallback
+    SKIP_RANGE_ANALYSIS = 2  # skip range analysis entirely, just return fallback_value
diff --git a/torch/fx/experimental/_dynamism.py b/torch/fx/experimental/_dynamism.py
index f6f30779ecc28..186e8de3898c0 100644
--- a/torch/fx/experimental/_dynamism.py
+++ b/torch/fx/experimental/_dynamism.py
@@ -1,13 +1,13 @@
 import re
 from collections.abc import Callable
-from typing import Any, Union
+from typing import Any
 
 import torch
 from torch.utils._pytree import tree_flatten_with_path, tree_map
 
 
 KeyPath = tuple[Any, ...]
-NonTensorShapeFn = Callable[[Union[int, float]], tuple[Any, ...]]
+NonTensorShapeFn = Callable[[int | float], tuple[Any, ...]]
 
 __all__ = [
     "normalize_source_name",
diff --git a/torch/fx/experimental/_size_hinting.py b/torch/fx/experimental/_size_hinting.py
new file mode 100644
index 0000000000000..578054fb139cf
--- /dev/null
+++ b/torch/fx/experimental/_size_hinting.py
@@ -0,0 +1,444 @@
+"""
+Size hinting utilities for symbolic shape expressions.
+
+This module contains the core logic for resolving symbolic expressions to
+concrete integer hints. Two strategies are provided:
+
+- _guarding_hint_or_throw_base: strict, only uses backed symbol hints, throws on
+  unbacked symbols. Use for correctness-critical guarding decisions.
+- _optimization_hint_base: permissive, uses heuristics and fallbacks for unbacked
+  symbols. Use for performance optimization decisions.
+"""
+
+from __future__ import annotations
+
+import logging
+import sys
+from typing import Any, TYPE_CHECKING
+
+import sympy
+
+from torch.utils._sympy.numbers import int_oo
+
+
+log = logging.getLogger(__name__)
+
+# Maximum number of free symbols in an expression before we skip
+# sympy.factor() in optimization_hint process for unbacked.
+# Factoring polynomials with many variables is expensive.
+SYMPY_FACTOR_MAX_FREE_SYMBOLS = 50
+
+if TYPE_CHECKING:
+    from torch.fx.experimental.symbolic_shapes import ShapeEnv
+
+
+def _sympy_subs(expr: sympy.Basic, replacements: dict[sympy.Expr, Any]) -> sympy.Basic:
+    """
+    When the passed replacement symbol v is a string, it is converted to a symbol with name v that
+    have the same replaced expression integer and nonnegative properties.
+    """
+
+    def to_symbol(replaced: sympy.Expr, replacement: sympy.Expr | str) -> sympy.Symbol:
+        if not isinstance(replaced, sympy.Expr):
+            raise AssertionError(
+                f"Expected sympy.Expr key, got {type(replaced)}: {replaced}"
+            )
+        if isinstance(replacement, str):
+            return sympy.Symbol(
+                replacement,
+                integer=replaced.is_integer,  # type: ignore[attr-defined]
+                nonnegative=replaced.is_nonnegative,  # type: ignore[attr-defined]
+            )
+        else:
+            return replacement
+
+    # xreplace is faster than subs, but is way more picky
+    return sympy.sympify(expr).xreplace(
+        {k: to_symbol(k, v) for k, v in replacements.items()}
+    )
+
+
+def _maybe_realize_expr(
+    expr: sympy.Basic, nan_fallback: int | None
+) -> int | bool | None:
+    """
+    Handle special sympy values in hinting APIs.
+
+    Returns:
+        - True/False for sympy.true/sympy.false (preserves bool type)
+        - Raises ValueError for complex numbers
+        - sys.maxsize for positive infinity
+        - -sys.maxsize for negative infinity
+        - fallback for NaN
+        - None if no special handling needed
+    """
+    if expr is sympy.true:
+        return True
+    if expr is sympy.false:
+        return False
+
+    try:
+        return int(expr)
+    except (TypeError, ValueError):
+        pass
+
+    if isinstance(expr, sympy.Expr):
+        if expr.has(sympy.I):
+            raise ValueError(
+                f"_maybe_realize_expr received a complex expression: {expr}. "
+                "Tensor dimensions cannot be complex numbers."
+            )
+        if expr in (int_oo, sympy.oo):
+            return sys.maxsize
+        if expr in (-int_oo, -sympy.oo):
+            return -sys.maxsize
+        if nan_fallback is not None and (expr is sympy.nan or expr.has(sympy.nan)):
+            return nan_fallback
+
+    return None
+
+
+def _guarding_hint_or_throw_base(
+    shape_env: ShapeEnv,
+    expr: sympy.Expr | sympy.Basic | int | bool,
+    precomputed_replacements: dict[sympy.Expr, sympy.Symbol],
+) -> int | bool:
+    """
+    Return a concrete integer hint for an expression that is safe to use for guarding.
+
+    This function evaluates the expression using only backed-symbols hints. Unlike
+    _optimization_hint_base(), this function does NOT use heuristics or fallback values
+    for unbacked symbols.
+
+    Use this when you need a hint value that will be used for a guarding decision.
+
+    Args:
+        shape_env: The ShapeEnv instance.
+        expr: A sympy expression or integer to evaluate.
+        precomputed_replacements: Precomputed replacements for PRECOMPUTED_SIZE symbols.
+
+    Returns:
+        The concrete integer value of the expression based on backed symbol hints.
+
+    Raises:
+        GuardOnDataDependentSymNode: If the expression contains unbacked symbols
+        (data-dependent values) that cannot be resolved to concrete values.
+
+    See Also:
+        _optimization_hint_base: For cases where fallback/heuristic values are acceptable
+            for unbacked symbols.
+    """
+    from torch.fx.experimental.symbolic_shapes import (
+        has_free_unbacked_symbols,
+        symbol_is_type,
+        SymT,
+    )
+
+    # sympy.expand() doesn't work with boolean expressions like Or/And
+    if isinstance(expr, sympy.Expr):
+        expr = sympy.expand(expr).xreplace(shape_env.replacements)
+    else:
+        expr = sympy.sympify(expr).xreplace(shape_env.replacements)
+
+    if isinstance(expr, sympy.Expr):
+        expr = expr.expand(identity=True)
+
+    result = _maybe_realize_expr(expr, None)
+    if result is not None:
+        return result
+
+    if not isinstance(expr, sympy.Basic):
+        raise RuntimeError("isinstance(expr, sympy.Basic)", expr, type(expr))
+
+    if any(symbol_is_type(s, SymT.PRECOMPUTED_SIZE) for s in expr.free_symbols):  # type: ignore[attr-defined]
+        expr = _sympy_subs(expr, precomputed_replacements)
+
+    # TODO do we need sympy_subs, or just xreplace
+    expr = _sympy_subs(expr, shape_env.backed_var_to_val)
+    if isinstance(expr, sympy.Expr):
+        expr = expr.expand(identity=True)
+
+    if has_free_unbacked_symbols(expr):
+        # Note: we could do better here and call
+        # _maybe_evaluate_static(orig_expr, compute_hint=True)
+        # but is it worth the overhead? probably not.
+        raise shape_env._make_data_dependent_error(expr, expr)
+
+    result = _maybe_realize_expr(expr, None)
+    if result is None:
+        raise RuntimeError("unexpected None!", expr)
+    return result
+
+
+def _get_unbacked_replacements(shape_env: ShapeEnv) -> dict[sympy.Expr, sympy.Expr]:
+    """Builds a mapping from unbacked expressions to canonical equivalents
+    using a union-find algorithm over deferred runtime asserts.
+    Used by optimization_hint to resolve unbacked symbols to consistent values."""
+    from collections import defaultdict
+
+    from torch.fx.experimental.symbolic_shapes import has_free_unbacked_symbols
+    from torch.utils._ordered_set import OrderedSet
+
+    if shape_env._unbacked_replacements is not None:
+        return shape_env._unbacked_replacements
+
+    class CanonicalExprFinder:
+        """
+        A disjoint-set/union-find data structure that can return the
+        "canonical" expression for a group of equivalent expressions.
+        - The canonical expression must come from the input eq_graph.
+        - The heuristics used to choose a leader determines which
+        expression becomes the canonical expression.
+        """
+
+        def __init__(self, eq_graph: dict[sympy.Expr, OrderedSet[sympy.Expr]]):
+            self.eq_graph = eq_graph
+            self.expressions = list(eq_graph.keys())
+            self.reverse_expressions = {
+                expr: i for i, expr in enumerate(self.expressions)
+            }
+            self.leader = list(range(len(self.expressions)))
+            self.size = [1] * len(self.expressions)
+            self._build_canonical_expr_mapping()
+
+        def _build_canonical_expr_mapping(self):
+            for expr, edges in self.eq_graph.items():
+                for adj in edges:
+                    self.union_expr(expr, adj)
+
+        def union_expr(self, a: sympy.Expr, b: sympy.Expr):
+            return self.union(self.reverse_expressions[a], self.reverse_expressions[b])
+
+        def union(self, a: int, b: int):
+            rootA = self.find(a)
+            rootB = self.find(b)
+            if rootA == rootB:
+                return False
+            leader, other = self.choose_leader(rootA, rootB)
+            self.leader[other] = leader
+            self.size[leader] += self.size[other]
+            return True
+
+        def find_expr(self, expr: sympy.Expr):
+            parent = self.find(self.reverse_expressions[expr])
+            return self.expressions[parent]
+
+        def find(self, x: int):
+            if self.leader[x] != x:
+                self.leader[x] = self.find(self.leader[x])
+            return self.leader[x]
+
+        def choose_leader(self, a: int, b: int):
+            """
+            The leader will become the canonical expression.
+            Returns a (leader, follower) tuple.
+
+            Heuristics:
+            1. Backed expression or constants preferred over unbacked expr
+            2. Simpler sub-expr when one contains the other
+            3. Higher frequency across equalities from deferred runtime assertions
+            4. Size of the set
+            5. Fallback to sympy.Basic.compare
+            """
+
+            def _choose(x: int, y: int) -> bool:
+                lhs, rhs = self.expressions[x], self.expressions[y]
+
+                any_unbacked_lhs = has_free_unbacked_symbols(lhs)
+                any_unbacked_rhs = has_free_unbacked_symbols(rhs)
+                if any_unbacked_lhs != any_unbacked_rhs:
+                    return bool(any_unbacked_rhs)
+
+                if lhs.has(rhs):
+                    return False
+                elif rhs.has(lhs):
+                    return True
+
+                degrees_lhs = len(self.eq_graph[lhs])
+                degrees_rhs = len(self.eq_graph[rhs])
+                if degrees_lhs != degrees_rhs:
+                    return degrees_lhs > degrees_rhs
+
+                if self.size[x] != self.size[y]:
+                    return self.size[x] > self.size[y]
+
+                return lhs.compare(rhs) == -1
+
+            if _choose(a, b):
+                return a, b
+            return b, a
+
+    # Build an undirected graph using ShapeEnv's deferred runtime assertions.
+    shape_env._equality_graph = defaultdict(OrderedSet)
+    for assertions in shape_env.deferred_runtime_asserts.values():
+        for assertion in assertions:
+            if not isinstance(assertion.expr, sympy.Equality):
+                continue
+            lhs = sympy.sympify(assertion.expr.lhs)
+            rhs = sympy.sympify(assertion.expr.rhs)
+            shape_env._equality_graph[lhs].add(rhs)
+            shape_env._equality_graph[rhs].add(lhs)
+
+    uf = CanonicalExprFinder(shape_env._equality_graph)
+
+    shape_env._unbacked_replacements = {}
+    for expr in shape_env._equality_graph:
+        canonical_expr = uf.find_expr(expr)
+        if expr != canonical_expr:
+            shape_env._unbacked_replacements[expr] = canonical_expr
+
+    return shape_env._unbacked_replacements
+
+
+def _sub_unbacked_exprs(shape_env: ShapeEnv, expr: sympy.Expr) -> sympy.Expr:
+    """Substitute unbacked expressions with canonical equivalents.
+    Used by optimization_hint to maximize consistency when hinting unbacked symbols."""
+    replacements = _get_unbacked_replacements(shape_env)
+
+    # consider making this threshold configurable
+    sub_cnt_limit = 30
+    sub_cnt = 0
+    while sub_cnt < sub_cnt_limit:
+        new_expr = expr.subs(replacements)
+        if new_expr == expr:
+            break
+        if len(new_expr.free_symbols) <= SYMPY_FACTOR_MAX_FREE_SYMBOLS:
+            expr = sympy.factor(new_expr)
+        else:
+            expr = new_expr
+        sub_cnt += 1
+    else:
+        log.warning("Substitution limit (%d) reached w/ %s", sub_cnt_limit, expr)
+
+    expr = _sympy_subs(expr, shape_env.backed_var_to_val)
+    expr = _sympy_subs(expr, shape_env.var_to_hint_override)
+    return expr
+
+
+def _optimization_hint_base(
+    shape_env: ShapeEnv,
+    expr: sympy.Expr | int,
+    precomputed_replacements: dict[sympy.Expr, sympy.Symbol],
+    fallback: int | None = None,
+) -> int:
+    """
+    Return a concrete integer hint for an expression using heuristics.
+
+    This function should be used for non-guarding based optimizations.
+    It will hint unbacked symbols using user provided optimization hints.
+    If not provided, fallback will be used along with some heuristics
+    that try to maximize consistency with the shape environment.
+
+    Args:
+        shape_env: The ShapeEnv instance.
+        expr: A sympy expression or integer to evaluate.
+        precomputed_replacements: Precomputed replacements for PRECOMPUTED_SIZE symbols.
+        fallback: Fallback value for unbacked symbols. If None, reads from config.
+
+    Returns:
+        A concrete integer hint for the expression.
+    """
+    from torch.fx.experimental.symbolic_shapes import (
+        has_free_unbacked_symbols,
+        symbol_is_type,
+        SymT,
+    )
+
+    # Read config at call time to respect runtime patches (e.g., in tests)
+    if fallback is None:
+        from torch._inductor.config import unbacked_symint_fallback
+
+        fallback = unbacked_symint_fallback
+
+    # to have expanded (Identity free) expr stored in original
+    if isinstance(expr, sympy.Expr):
+        expr = expr.expand(identity=True)
+
+    original = expr
+    # sympy.expand() doesn't work with boolean expressions like Or/And
+    if isinstance(expr, sympy.Expr):
+        expr = expr.xreplace(shape_env.replacements)
+    else:
+        expr = sympy.sympify(expr).xreplace(shape_env.replacements)
+
+    result = _maybe_realize_expr(expr, fallback)
+    if result is not None:
+        return result
+
+    if isinstance(expr, sympy.Expr):
+        expr = expr.expand(identity=True)
+
+    # Replace backed symbols with their hints, leaving unbacked symbols alone.
+    result = _maybe_realize_expr(expr, None)
+    if result is not None:
+        return result
+
+    if not isinstance(expr, sympy.Expr):
+        raise RuntimeError("isinstance(expr, sympy.Expr)", expr)
+
+    if any(symbol_is_type(s, SymT.PRECOMPUTED_SIZE) for s in expr.free_symbols):  # type: ignore[attr-defined]
+        expr = _sympy_subs(expr, precomputed_replacements)
+
+    expr = _sympy_subs(expr, shape_env.backed_var_to_val)
+    if isinstance(expr, sympy.Expr):
+        expr = expr.expand(identity=True)
+
+    result = _maybe_realize_expr(expr, fallback)
+    if result is not None:
+        return result
+
+    expr = _sympy_subs(expr, shape_env.var_to_hint_override)
+
+    result = _maybe_realize_expr(expr, fallback)
+    if result is not None:
+        return result
+
+    # If unbacked symbols remain, try to substitute them using heuristics
+    # that maximize consistency with the shape environment.
+    if has_free_unbacked_symbols(expr):
+        # Make sure to substitute with the factored version
+        # e.g. 10*(s0 + u0) instead of 10*s0 + 10*u0
+        if (
+            isinstance(original, sympy.Expr)
+            and len(original.free_symbols) <= SYMPY_FACTOR_MAX_FREE_SYMBOLS
+        ):
+            original = sympy.factor(original)
+        expr = _sub_unbacked_exprs(shape_env, original)
+
+    # For multiple expressions that depend on an unbacked symint,
+    # we want to compute them consistently for a size hint we have chosen.
+    # So, recursively compute expressions via size hints of contained symbols.
+    # For example: u1 * u2 - 10 ==> fallback * fallback - 10
+
+    if not isinstance(expr, sympy.Expr):
+        raise RuntimeError(f"Expected sympy Expr, got {type(expr)}: {expr}")
+    free_symbols = expr.free_symbols
+
+    # Constrain fallback per-symbol based on var_to_range bounds
+    size_dict = {}
+    for s in free_symbols:
+        sym_fallback = fallback
+        vr = shape_env.var_to_range.get(s, None)
+        if vr is not None:
+            if isinstance(vr.lower, (int, sympy.Integer)):
+                sym_fallback = max(sym_fallback, int(vr.lower))
+            if isinstance(vr.upper, (int, sympy.Integer)):
+                sym_fallback = min(sym_fallback, int(vr.upper))
+        size_dict[s] = sym_fallback
+
+    try:
+        final_result = expr.subs(size_dict)
+    except ZeroDivisionError:
+        # Expressions like ModularIndexing(x, u1, 4) crash during subs()
+        # when u1 is substituted with 0, because sympy eagerly evaluates
+        # (x // 0) % 4.  This can happen when an unbacked symbol with
+        # var_to_range lower=0 is used as a divisor (e.g. from
+        # _dynamic_reshape_indexer) and the fallback also maps to 0.
+        # Return fallback in that case.
+        return fallback if fallback is not None else 0
+
+    final_result = _maybe_realize_expr(final_result, fallback)
+    if final_result is None:
+        raise RuntimeError(f"Failed to realize expression to int: {expr}")
+
+    return final_result
diff --git a/torch/fx/experimental/const_fold.py b/torch/fx/experimental/const_fold.py
index caf29c34cd3e6..f44293e2242a7 100644
--- a/torch/fx/experimental/const_fold.py
+++ b/torch/fx/experimental/const_fold.py
@@ -1,7 +1,6 @@
 # mypy: allow-untyped-defs
 import re
 from collections.abc import Callable
-from typing import Optional, Union
 
 import torch.fx
 from torch.fx.node import map_arg
@@ -29,8 +28,8 @@ def __init__(
         self,
         root: torch.nn.Module,
         graph: torch.fx.Graph,
-        const_subgraph: Optional[torch.fx.Graph] = None,
-        fx_const_folded_attrs_name: Optional[str] = None,
+        const_subgraph: torch.fx.Graph | None = None,
+        fx_const_folded_attrs_name: str | None = None,
         device_for_folded_attrs: str = "cuda",
     ):
         super().__init__(root, graph)
@@ -193,8 +192,8 @@ def get_unique_attr_name_in_module(mod_traced: torch.fx.GraphModule, name: str)
 
 
 def split_const_subgraphs(
-    module: Union[torch.nn.Module, torch.fx.GraphModule],
-    skip_folding_node_fn: Optional[Callable[[torch.fx.Node], bool]] = None,
+    module: torch.nn.Module | torch.fx.GraphModule,
+    skip_folding_node_fn: Callable[[torch.fx.Node], bool] | None = None,
     device_for_folded_attrs: str = "cpu",
 ) -> FoldedGraphModule:
     """
diff --git a/torch/fx/experimental/merge_matmul.py b/torch/fx/experimental/merge_matmul.py
index c6a51918f930f..3c200a2a206f9 100644
--- a/torch/fx/experimental/merge_matmul.py
+++ b/torch/fx/experimental/merge_matmul.py
@@ -92,12 +92,15 @@ def merge_matmul(in_mod: torch.nn.Module):
     """
     A graph transformation that merges matrix multiplication operations that share the same right-hand
     side operand into one large matrix multiplication.
-               ____      _________        _________
-      ----    |    |    |         |     M|  A * C  |
-    M| A  |  T| B  | * K|    C    | =    |---------|
-      ---- ,  |    |    |         |     T|  B * C  |
-       K       ----      ---------        ---------
-                K            R                R
+
+    ::
+
+                   ____      _________        _________
+          ----    |    |    |         |     M|  A * C  |
+        M| A  |  T| B  | * K|    C    | =    |---------|
+          ---- ,  |    |    |         |     T|  B * C  |
+           K       ----      ---------        ---------
+                    K            R                R
     """
     gm = symbolic_trace(in_mod)
 
diff --git a/torch/fx/experimental/meta_tracer.py b/torch/fx/experimental/meta_tracer.py
index 7e9140e2f0fe4..472ed0860b3de 100644
--- a/torch/fx/experimental/meta_tracer.py
+++ b/torch/fx/experimental/meta_tracer.py
@@ -3,7 +3,7 @@
 import functools
 import warnings
 from collections.abc import Callable
-from typing import Any, Optional, Union
+from typing import Any
 
 import torch
 import torch.fx
@@ -317,9 +317,9 @@ def trace(self, root, meta_args: dict[str, torch.Tensor], concrete_args=None):
 
 
 def symbolic_trace(
-    root: Union[torch.nn.Module, Callable[..., Any]],
-    meta_args: Optional[dict[str, torch.Tensor]] = None,
-    concrete_args: Optional[dict[str, Any]] = None,
+    root: torch.nn.Module | Callable[..., Any],
+    meta_args: dict[str, torch.Tensor] | None = None,
+    concrete_args: dict[str, Any] | None = None,
 ) -> torch.fx.GraphModule:
     tracer = MetaTracer()
     graph = tracer.trace(root, meta_args, concrete_args)  # type: ignore[arg-type]
diff --git a/torch/fx/experimental/migrate_gradual_types/constraint.py b/torch/fx/experimental/migrate_gradual_types/constraint.py
index ce724d39f228f..2a296203c52c0 100644
--- a/torch/fx/experimental/migrate_gradual_types/constraint.py
+++ b/torch/fx/experimental/migrate_gradual_types/constraint.py
@@ -1,4 +1,40 @@
-# mypy: allow-untyped-defs
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, TypeAlias
+
+
+if TYPE_CHECKING:
+    from collections.abc import Sequence
+
+__all__ = [
+    "ApplyBroadcasting",
+    "BinConstraintD",
+    "BinConstraintT",
+    "BinaryConstraint",
+    "BVar",
+    "CalcConv",
+    "CalcMaxPool",
+    "CalcProduct",
+    "CanReshape",
+    "Conj",
+    "Constraint",
+    "DGreatestUpperBound",
+    "Disj",
+    "DVar",
+    "F",
+    "GetItem",
+    "GetItemTensor",
+    "IndexSelect",
+    "Prod",
+    "T",
+    "TGreatestUpperBound",
+    "Transpose",
+    "TVar",
+    "is_algebraic_expression",
+    "is_bool_expr",
+    "is_dim",
+]
+
 from torch.fx.experimental.migrate_gradual_types.operation import (
     op_add,
     op_div,
@@ -10,7 +46,7 @@
     op_neq,
     op_sub,
 )
-from torch.fx.tensor_type import Dyn, TensorType
+from torch.fx.tensor_type import _DynType, Dyn, TensorType
 
 
 class Constraint:
@@ -18,55 +54,53 @@ class Constraint:
 
 
 class Conj(Constraint):
-    def __init__(self, conjuncts):
+    def __init__(self, conjuncts: Sequence[Constraint]) -> None:
         """
         :param conjuncts: Conjunction of constraints
         """
-        self.conjucts = conjuncts
+        self.conjucts = list(conjuncts)
 
-    def __eq__(self, other):
+    def __eq__(self, other: object) -> bool:
         if isinstance(other, Conj):
-            return self.conjucts == other.conjucts and self.conjucts == other.conjucts
+            return self.conjucts == other.conjucts
         else:
             return False
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         return f"And({self.conjucts})"
 
 
 class Disj(Constraint):
-    def __init__(self, disjuncts):
+    def __init__(self, disjuncts: Sequence[Constraint]) -> None:
         """
         :param disjuncts: Disjunction of constraints
         """
-        self.disjuncts = disjuncts
+        self.disjuncts = list(disjuncts)
 
-    def __eq__(self, other):
+    def __eq__(self, other: object) -> bool:
         if isinstance(other, Disj):
-            return (
-                self.disjuncts == other.disjuncts and self.disjuncts == other.disjuncts
-            )
+            return self.disjuncts == other.disjuncts
         else:
             return False
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         return f"Or({self.disjuncts})"
 
 
 class Prod(Constraint):
-    def __init__(self, products):
+    def __init__(self, products: Sequence[DVar | int | _DynType]) -> None:
         """
         :param products: lists of dimensions to multiply
         """
-        self.products = products
+        self.products = list(products)
 
-    def __eq__(self, other):
+    def __eq__(self, other: object) -> bool:
         if isinstance(other, Prod):
-            return self.products == other.products and self.products == other.products
+            return self.products == other.products
         else:
             return False
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         return f"Product({self.products})"
 
 
@@ -78,10 +112,10 @@ class T(Constraint):
     def __init__(self) -> None:
         pass
 
-    def __eq__(self, other):
+    def __eq__(self, other: object) -> bool:
         return isinstance(other, T)
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         return "True"
 
 
@@ -93,10 +127,10 @@ class F(Constraint):
     def __init__(self) -> None:
         pass
 
-    def __eq__(self, other):
+    def __eq__(self, other: object) -> bool:
         return isinstance(other, F)
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         return "False"
 
 
@@ -105,7 +139,7 @@ class BinaryConstraint(Constraint):
     Represents all binary operations
     """
 
-    def __init__(self, lhs, rhs, op):
+    def __init__(self, lhs: _Operand, rhs: _Operand, op: str | None) -> None:
         """
         :param lhs: lhs of the constraint
         :param rhs: rhs of the constraint
@@ -115,7 +149,7 @@ def __init__(self, lhs, rhs, op):
         self.rhs = rhs
         self.op = op
 
-    def __eq__(self, other):
+    def __eq__(self, other: object) -> bool:
         if isinstance(other, BinaryConstraint):
             return (
                 self.lhs == other.lhs and self.rhs == other.rhs and self.op == other.op
@@ -123,7 +157,7 @@ def __eq__(self, other):
         else:
             return False
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         return f"({self.lhs} {self.op} {self.rhs})"
 
 
@@ -132,7 +166,7 @@ class BinConstraintT(BinaryConstraint):
     Binary constraints about tensors
     """
 
-    def __init__(self, lhs, rhs, op):
+    def __init__(self, lhs: _Operand, rhs: _Operand, op: str | None) -> None:
         if not (
             (isinstance(lhs, (TVar, TensorType, int)) or lhs == Dyn)
             and (isinstance(rhs, (TVar, TensorType, int)) or rhs == Dyn)
@@ -146,7 +180,7 @@ class BinConstraintD(BinaryConstraint):
     Binary constraints about dimensions
     """
 
-    def __init__(self, lhs, rhs, op):
+    def __init__(self, lhs: _Operand, rhs: _Operand, op: str | None) -> None:
         if not (is_algebraic_expression(lhs) or is_dim(lhs) or is_bool_expr(lhs)):
             raise AssertionError(f"Invalid lhs type: {type(lhs)}")
         if not (is_algebraic_expression(rhs) or is_dim(rhs) or is_bool_expr(rhs)):
@@ -160,7 +194,7 @@ class TGreatestUpperBound(Constraint):
     Greatest Upper bound for tensors with dynamic type
     """
 
-    def __init__(self, res, rhs1, rhs2):
+    def __init__(self, res: TVar, rhs1: TVar, rhs2: TVar) -> None:
         """
         :param res: tensor variable that stores the result of the output
         :param rhs1: tensor or tensor variable
@@ -170,10 +204,10 @@ def __init__(self, res, rhs1, rhs2):
         self.rhs1 = rhs1
         self.rhs2 = rhs2
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         return f"{self.res} = {self.rhs1}\u2294*{self.rhs2}"
 
-    def __eq__(self, other):
+    def __eq__(self, other: object) -> bool:
         if isinstance(other, TGreatestUpperBound):
             return (
                 self.res == other.res
@@ -189,7 +223,12 @@ class DGreatestUpperBound(Constraint):
     Greatest Upper bound for dimensions
     """
 
-    def __init__(self, res, rhs1, rhs2):
+    def __init__(
+        self,
+        res: DVar | int | _DynType,
+        rhs1: DVar | int | _DynType,
+        rhs2: DVar | int | _DynType,
+    ) -> None:
         """
         :param res: Dimension variable to store the result
         :param rhs1: dimension variable 1
@@ -206,10 +245,10 @@ def __init__(self, res, rhs1, rhs2):
         self.rhs1 = rhs1
         self.rhs2 = rhs2
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         return f"{self.res} = {self.rhs1}\u2294{self.rhs2}"
 
-    def __eq__(self, other):
+    def __eq__(self, other: object) -> bool:
         if isinstance(other, DGreatestUpperBound):
             return (
                 self.res == other.res
@@ -225,7 +264,7 @@ class CanReshape(Constraint):
     can_reshape constraint
     """
 
-    def __init__(self, src, target):
+    def __init__(self, src: TVar, target: TensorType) -> None:
         """
         :param src: tensor variable
         :param target: tensor
@@ -233,10 +272,10 @@ def __init__(self, src, target):
         self.src = src
         self.target = target
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         return f"can-reshape({self.src}, {self.target})"
 
-    def __eq__(self, other):
+    def __eq__(self, other: object) -> bool:
         if isinstance(other, CanReshape):
             return self.src == other.src and self.target == other.target
         else:
@@ -244,7 +283,14 @@ def __eq__(self, other):
 
 
 class IndexSelect(Constraint):
-    def __init__(self, tensor_size, input_var, dim_replace, index, output):
+    def __init__(
+        self,
+        tensor_size: int,
+        input_var: TVar,
+        dim_replace: DVar | _DynType,
+        index: int,
+        output: TVar,
+    ) -> None:
         """
         Args:
             input_var: input to index_select
@@ -268,7 +314,7 @@ def __init__(self, tensor_size, input_var, dim_replace, index, output):
         self.index = index
         self.output = output
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         return (
             f" {self.output} = "
             f"IndexSelect({self.input_var}, "
@@ -277,7 +323,7 @@ def __repr__(self):
             f"{self.index})"
         )
 
-    def __eq__(self, other):
+    def __eq__(self, other: object) -> bool:
         if isinstance(other, IndexSelect):
             return (
                 self.tensor_size == other.tensor_size
@@ -291,7 +337,9 @@ def __eq__(self, other):
 
 
 class Transpose(Constraint):
-    def __init__(self, tensor_size, input_var, index1, index2, output):
+    def __init__(
+        self, tensor_size: int, input_var: TVar, index1: int, index2: int, output: TVar
+    ) -> None:
         """
         Args:
             tensor_size: current tensor size
@@ -315,7 +363,7 @@ def __init__(self, tensor_size, input_var, index1, index2, output):
         self.index2 = index2
         self.output = output
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         return (
             f" {self.output} = "
             f"Transpose({self.input_var}, "
@@ -324,7 +372,7 @@ def __repr__(self):
             f"{self.index2})"
         )
 
-    def __eq__(self, other):
+    def __eq__(self, other: object) -> bool:
         if isinstance(other, Transpose):
             return (
                 self.tensor_size == other.tensor_size
@@ -338,7 +386,9 @@ def __eq__(self, other):
 
 
 class GetItem(Constraint):
-    def __init__(self, tensor_size, index, res, input_var):
+    def __init__(
+        self, tensor_size: int, index: int, res: DVar, input_var: TVar
+    ) -> None:
         """
         Constraint for getting item given a tensor size
         :param tensor_size: actual number
@@ -354,10 +404,10 @@ def __init__(self, tensor_size, index, res, input_var):
         self.index = index
         self.input_var = input_var
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         return f" {self.res} = GetItem({self.input_var}, tensor_size: {self.tensor_size}, {self.index})"
 
-    def __eq__(self, other):
+    def __eq__(self, other: object) -> bool:
         if isinstance(other, GetItem):
             return (
                 self.res == other.res
@@ -370,7 +420,13 @@ def __eq__(self, other):
 
 
 class GetItemTensor(Constraint):
-    def __init__(self, tensor_size, index_tuple, res, input_var):
+    def __init__(
+        self,
+        tensor_size: int,
+        index_tuple: tuple[None | slice, ...],
+        res: TVar,
+        input_var: TVar,
+    ) -> None:
         """
         Constraint for getting item given a tensor size
         However, when the argument is a tuple, we will
@@ -388,10 +444,10 @@ def __init__(self, tensor_size, index_tuple, res, input_var):
         self.index_tuple = index_tuple
         self.input_var = input_var
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         return f" {self.res} = GetItemT({self.input_var}, tensor_size: {self.tensor_size}, {self.index_tuple})"
 
-    def __eq__(self, other):
+    def __eq__(self, other: object) -> bool:
         if isinstance(other, GetItemTensor):
             return (
                 self.res == other.res
@@ -406,15 +462,15 @@ def __eq__(self, other):
 class CalcConv(Constraint):
     def __init__(
         self,
-        conv_result,
-        input_var,
-        c_out,
-        kernel,
-        padding,
-        stride,
-        dilation,
-        matching_constraint_vars,
-    ):
+        conv_result: TVar,
+        input_var: TVar,
+        c_out: int,
+        kernel: int | tuple[int, int],
+        padding: int | tuple[int, int],
+        stride: int | tuple[int, int],
+        dilation: int | tuple[int, int],
+        matching_constraint_vars: list[DVar],
+    ) -> None:
         """
         :param conv_result: the convolution result
         :param input_var: input to convolution
@@ -430,7 +486,7 @@ def __init__(
         self.dilation = dilation
         self.matching_constraint = matching_constraint_vars
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         return (
             f"{self.conv_result} ="
             f" calc-conv({self.input_var},"
@@ -439,7 +495,7 @@ def __repr__(self):
             f" {self.dilation})"
         )
 
-    def __eq__(self, other):
+    def __eq__(self, other: object) -> bool:
         if isinstance(other, CalcConv):
             return (
                 self.conv_result == other.conv_result
@@ -458,14 +514,14 @@ def __eq__(self, other):
 class CalcMaxPool(Constraint):
     def __init__(
         self,
-        maxpool_result,
-        input_var,
-        kernel,
-        padding,
-        stride,
-        dilation,
-        matching_constraint_vars,
-    ):
+        maxpool_result: TVar,
+        input_var: TVar,
+        kernel: int | tuple[int, int],
+        padding: int | tuple[int, int],
+        stride: int | tuple[int, int],
+        dilation: int | tuple[int, int],
+        matching_constraint_vars: list[DVar],
+    ) -> None:
         """
         :param maxpool_result: the result of maxpool
         :param input_var: input to convolution
@@ -479,7 +535,7 @@ def __init__(
         self.dilation = dilation
         self.matching_constraint = matching_constraint_vars
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         return (
             f"{self.maxpool_result} ="
             f" calc-maxpool({self.input_var},"
@@ -488,7 +544,7 @@ def __repr__(self):
             f" {self.dilation})"
         )
 
-    def __eq__(self, other):
+    def __eq__(self, other: object) -> bool:
         if isinstance(other, CalcMaxPool):
             return (
                 self.maxpool_result == other.maxpool_result
@@ -504,7 +560,7 @@ def __eq__(self, other):
 
 
 class ApplyBroadcasting(Constraint):
-    def __init__(self, res1, res2, input1, input2):
+    def __init__(self, res1: TVar, res2: TVar, input1: TVar, input2: TVar) -> None:
         """
         :param res1: resulting tensor 1
         :param res2: resulting tensor 2
@@ -516,7 +572,7 @@ def __init__(self, res1, res2, input1, input2):
         self.input1 = input1
         self.input2 = input2
 
-    def __eq__(self, other):
+    def __eq__(self, other: object) -> bool:
         if isinstance(other, ApplyBroadcasting):
             return (
                 self.res1 == other.res1
@@ -527,7 +583,7 @@ def __eq__(self, other):
         else:
             return False
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         return (
             f"{self.res1}, {self.res2} ="
             f" apply-broadcasting({self.input1},"
@@ -540,7 +596,9 @@ class CalcProduct(Constraint):
     Given correct dimensions, calculate the product for flatten accounting for Dyn
     """
 
-    def __init__(self, start, end, flattened, dims_to_flatten):
+    def __init__(
+        self, start: int, end: int, flattened: TVar, dims_to_flatten: list[DVar]
+    ) -> None:
         """
         :param start: start index
         :param end: end index
@@ -561,7 +619,7 @@ def __init__(self, start, end, flattened, dims_to_flatten):
         self.dims_to_flatten = dims_to_flatten
         self.flattened = flattened
 
-    def __eq__(self, other):
+    def __eq__(self, other: object) -> bool:
         if isinstance(other, CalcProduct):
             return (
                 self.start == other.start
@@ -573,7 +631,7 @@ def __eq__(self, other):
         else:
             return False
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         return f"{self.flattened} = CalcProduct({self.start}, {self.end}, {self.dims_to_flatten})"
 
 
@@ -582,16 +640,16 @@ class TVar:
     Tensor variable with no tensor constructor
     """
 
-    def __init__(self, tvar):
+    def __init__(self, tvar: int) -> None:
         """
         :param tvar: tensor variable
         """
         self.tvar = tvar
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         return f"TV({self.tvar})"
 
-    def __eq__(self, other):
+    def __eq__(self, other: object) -> bool:
         if isinstance(other, TVar):
             return self.tvar == other.tvar
         else:
@@ -603,16 +661,16 @@ class DVar:
     Dimension variable
     """
 
-    def __init__(self, c):
+    def __init__(self, c: int) -> None:
         """
         :param c: character or number
         """
         self.c = c
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         return f"DV({self.c})"
 
-    def __eq__(self, other):
+    def __eq__(self, other: object) -> bool:
         if isinstance(other, DVar):
             return self.c == other.c
         else:
@@ -624,35 +682,52 @@ class BVar:
     Boolean variable
     """
 
-    def __init__(self, c):
+    def __init__(self, c: int) -> None:
         """
         :param c: character or number
         """
         self.c = c
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         return f"BV({self.c})"
 
-    def __eq__(self, other):
+    def __eq__(self, other: object) -> bool:
         if isinstance(other, BVar):
             return self.c == other.c
         else:
             return False
 
 
-def is_algebraic_expression(constraint):
+_Operand: TypeAlias = (
+    TVar
+    | TensorType
+    | DVar
+    | int
+    | float
+    | bool
+    | _DynType
+    | BinConstraintD
+    | Prod
+    | BVar
+    | Conj
+    | Disj
+    | None
+)
+
+
+def is_algebraic_expression(constraint: object) -> bool:
     if isinstance(constraint, BinConstraintD):
         return constraint.op in [op_add, op_sub, op_div, op_mul, op_mod]
     else:
         return isinstance(constraint, Prod)
 
 
-def is_bool_expr(constraint):
+def is_bool_expr(constraint: object) -> bool:
     if isinstance(constraint, BinConstraintD):
         return constraint.op in [op_gt, op_lt, op_neq, op_eq]
     else:
         return isinstance(constraint, (BVar, Conj, Disj))
 
 
-def is_dim(d):
+def is_dim(d: object) -> bool:
     return isinstance(d, (DVar, int)) or d == Dyn
diff --git a/torch/fx/experimental/migrate_gradual_types/constraint_generator.py b/torch/fx/experimental/migrate_gradual_types/constraint_generator.py
index 9f5a7e8064a2b..2bc3e128ea83c 100644
--- a/torch/fx/experimental/migrate_gradual_types/constraint_generator.py
+++ b/torch/fx/experimental/migrate_gradual_types/constraint_generator.py
@@ -1,8 +1,7 @@
-# mypy: allow-untyped-defs
 import operator
 import warnings
-from collections.abc import Callable, Iterable
-from typing import TypeVar
+from collections.abc import Callable, Iterable, Sequence
+from typing import TypeAlias, TypeVar
 from typing_extensions import ParamSpec
 
 import torch
@@ -11,11 +10,13 @@
     ApplyBroadcasting,
     BinConstraintD,
     BinConstraintT,
+    BVar,
     CalcConv,
     CalcMaxPool,
     CalcProduct,
     CanReshape,
     Conj,
+    Constraint,
     DGreatestUpperBound,
     Disj,
     DVar,
@@ -49,6 +50,7 @@
     gen_tensor_dims,
     gen_tvar,
 )
+from torch.fx.graph import Graph
 from torch.fx.node import Node, Target
 from torch.fx.tensor_type import Dyn, TensorType
 from torch.nn.modules.batchnorm import BatchNorm2d
@@ -58,7 +60,9 @@
 _T = TypeVar("_T")
 _P = ParamSpec("_P")
 
-_INFERENCE_RULES: dict[Target, Callable] = {}
+_SymbolDict: TypeAlias = dict[Node, TVar | DVar | BVar]
+
+_INFERENCE_RULES: dict[Target, Callable[..., tuple[list[Constraint], int]]] = {}
 
 MAX_TENSOR_RANK = 4
 
@@ -117,13 +121,15 @@ def register_inference_rule(
     def register(fn: Callable[_P, _T]) -> Callable[_P, _T]:
         if call_target in _INFERENCE_RULES:
             raise RuntimeError(f"Inference rule already registered for {call_target}!")
-        _INFERENCE_RULES[call_target] = fn
+        _INFERENCE_RULES[call_target] = fn  # pyrefly: ignore[unsupported-operation]
         return fn
 
     return register
 
 
-def generate_flatten_constraints(start_dim, end_dim, input, flattened, n, counter):
+def generate_flatten_constraints(
+    start_dim: int, end_dim: int, input: TVar, flattened: TVar, n: int, counter: int
+) -> tuple[Conj, int]:
     d, counter = gen_tensor_dims(n, counter)
     c1 = BinConstraintT(input, TensorType(d), op_eq)
     start_dim = n if start_dim == -1 else abs(start_dim)
@@ -134,7 +140,9 @@ def generate_flatten_constraints(start_dim, end_dim, input, flattened, n, counte
 
 
 @register_inference_rule(getattr)
-def get_attr_inference_rule(n: Node, symbols, constraints, counter):
+def get_attr_inference_rule(
+    n: Node, symbols: _SymbolDict, constraints: list[Constraint], counter: int
+) -> tuple[list[Constraint], int]:
     """
     If the attribute is "device" then the tensor shape is preserved
     """
@@ -155,7 +163,9 @@ def get_attr_inference_rule(n: Node, symbols, constraints, counter):
 
 
 @register_inference_rule(torch.bmm)
-def bmm_inference_rule(n: Node, symbols, constraints, counter):
+def bmm_inference_rule(
+    n: Node, symbols: _SymbolDict, constraints: list[Constraint], counter: int
+) -> tuple[list[Constraint], int]:
     """
     Constraints that match the input to a size 3 tensor
     and switch the dimensions according to the rules
@@ -227,7 +237,9 @@ def bmm_inference_rule(n: Node, symbols, constraints, counter):
 
 
 @register_inference_rule("index_select")
-def index_select_inference_rule(n: Node, symbols, constraints, counter):
+def index_select_inference_rule(
+    n: Node, symbols: _SymbolDict, constraints: list[Constraint], counter: int
+) -> tuple[list[Constraint], int]:
     """
     We constrain the second argument to a vector or Dyn.
     The output replaces the input with the shape of the vector
@@ -256,7 +268,13 @@ def index_select_inference_rule(n: Node, symbols, constraints, counter):
             Disj(
                 [
                     IndexSelect(
-                        i + 1, symbols[n.args[0]], dims[0], n.args[1], index_select
+                        i + 1,
+                        symbols[  # pyrefly: ignore[bad-argument-type, bad-index]
+                            n.args[0]
+                        ],
+                        dims[0],
+                        n.args[1],
+                        index_select,
                     )
                     for i in range(MAX_TENSOR_RANK)
                 ]
@@ -268,7 +286,15 @@ def index_select_inference_rule(n: Node, symbols, constraints, counter):
             is_dyn,
             Disj(
                 [
-                    IndexSelect(i + 1, symbols[n.args[0]], Dyn, n.args[1], index_select)
+                    IndexSelect(
+                        i + 1,
+                        symbols[  # pyrefly: ignore[bad-argument-type, bad-index]
+                            n.args[0]
+                        ],
+                        Dyn,
+                        n.args[1],
+                        index_select,
+                    )
                     for i in range(MAX_TENSOR_RANK)
                 ]
             ),
@@ -279,7 +305,9 @@ def index_select_inference_rule(n: Node, symbols, constraints, counter):
 
 
 @register_inference_rule("expand")
-def expand_inference_rule(n: Node, symbols, constraints, counter):
+def expand_inference_rule(
+    n: Node, symbols: _SymbolDict, constraints: list[Constraint], counter: int
+) -> tuple[list[Constraint], int]:
     """
     We generate the exact constraints as we do for tensor additions but we constraint
     the rank of this expression to be equal to len(n.args[1:]) so that only
@@ -308,13 +336,22 @@ def expand_inference_rule(n: Node, symbols, constraints, counter):
     e2_constraint = BinConstraintT(
         e2,
         TensorType(
-            [arg if isinstance(arg, int) else symbols[arg] for arg in n.args[1:]]
+            [
+                arg
+                if isinstance(arg, int)
+                else symbols[arg]  # pyrefly: ignore[bad-index]
+                for arg in n.args[1:]
+            ]
         ),
         op_eq,
     )
 
     constraints, counter = gen_broadcasting_constraints(
-        e1, e2, symbols, counter, expand
+        e1,  # pyrefly: ignore[bad-argument-type]
+        e2,
+        symbols,
+        counter,
+        expand,
     )
 
     # constraint the output size
@@ -341,7 +378,9 @@ def expand_inference_rule(n: Node, symbols, constraints, counter):
 @register_inference_rule("contiguous")
 @register_inference_rule(torch.ones)
 @register_inference_rule(torch.zeros)
-def equality_inference_rule(n: Node, symbols, constraints, counter):
+def equality_inference_rule(
+    n: Node, symbols: _SymbolDict, constraints: list[Constraint], counter: int
+) -> tuple[list[Constraint], int]:
     """
     We generate the constraint: input = output
     """
@@ -356,23 +395,27 @@ def equality_inference_rule(n: Node, symbols, constraints, counter):
         # then we have dimension variables
         else:
             for arg in n.args:
-                if not isinstance(symbols[arg], DVar):
-                    raise AssertionError(f"Expected DVar, got {type(symbols[arg])}")
-        my_size = [symbols[arg] for arg in n.args]
+                if not isinstance(symbols[arg], DVar):  # pyrefly: ignore[bad-index]
+                    raise AssertionError(
+                        f"Expected DVar, got {type(symbols[arg])}"  # pyrefly: ignore[bad-index]
+                    )
+        my_size = [symbols[arg] for arg in n.args]  # pyrefly: ignore[bad-index]
         return [BinConstraintT(output, TensorType(my_size), op_eq)], counter
 
     elif isinstance(n.args[0], tuple):
         # then the tuple is the size
         if len(n.args[0]) > 4:
             raise AssertionError(f"Expected len <= 4, got {len(n.args[0])}")
-        my_size = [symbols[arg] for arg in n.args[0]]
+        my_size = [symbols[arg] for arg in n.args[0]]  # pyrefly: ignore[bad-index]
         return [BinConstraintT(output, TensorType(my_size), op_eq)], counter
     else:
         raise NotImplementedError("Method not yet implemented")
 
 
 @register_inference_rule("transpose")
-def transpose_inference_rule(n: Node, symbols, constraints, counter):
+def transpose_inference_rule(
+    n: Node, symbols: _SymbolDict, constraints: list[Constraint], counter: int
+) -> tuple[list[Constraint], int]:
     """
     Can be considered as a sequence of two index selects, so we generate constraints accordingly
     """
@@ -407,7 +450,9 @@ def transpose_inference_rule(n: Node, symbols, constraints, counter):
 
 
 @register_inference_rule("type_as")
-def type_inference_rule(n: Node, symbols, constraints, counter):
+def type_inference_rule(
+    n: Node, symbols: _SymbolDict, constraints: list[Constraint], counter: int
+) -> tuple[list[Constraint], int]:
     """
     We generate the constraint: input = output
     """
@@ -434,7 +479,9 @@ def type_inference_rule(n: Node, symbols, constraints, counter):
 
 
 @register_inference_rule("masked_fill_")
-def masked_fill_inference_rule(n: Node, symbols, constraints, counter):
+def masked_fill_inference_rule(
+    n: Node, symbols: _SymbolDict, constraints: list[Constraint], counter: int
+) -> tuple[list[Constraint], int]:
     """
     Similar to addition. For now we implement the constraints when
     the argument is a boolean tensor. There is also a case for when
@@ -463,11 +510,13 @@ def masked_fill_inference_rule(n: Node, symbols, constraints, counter):
 
 
 @register_inference_rule(torch.nn.functional.embedding)
-def embedding_inference_rule_functional(n: Node, symbols, constraints, counter):
+def embedding_inference_rule_functional(
+    n: Node, symbols: _SymbolDict, constraints: list[Constraint], counter: int
+) -> tuple[list[Constraint], int]:
     if not isinstance(n.args[0], Node):
         raise AssertionError(f"Expected Node, got {type(n.args[0])}")
 
-    embedding_dim_weights = symbols[n.args[1]]
+    embedding_dim_weights = symbols[n.args[1]]  # pyrefly: ignore[bad-index]
 
     # will treat this as a static shape. So we will not use matching.
     weight_dims, counter = gen_tensor_dims(2, counter)
@@ -480,7 +529,13 @@ def embedding_inference_rule_functional(n: Node, symbols, constraints, counter):
 
 
 @register_inference_rule(torch.nn.modules.sparse.Embedding)
-def embedding_inference_rule(n: Node, module_instance, symbols, constraints, counter):
+def embedding_inference_rule(
+    n: Node,
+    module_instance: torch.nn.Embedding,
+    symbols: _SymbolDict,
+    constraints: list[Constraint],
+    counter: int,
+) -> tuple[list[Constraint], int]:
     """
     The output shape differs from the input shape in the last dimension
     """
@@ -489,10 +544,12 @@ def embedding_inference_rule(n: Node, module_instance, symbols, constraints, cou
     return gen_embedding_rules(n, symbols, module_instance.embedding_dim, counter)
 
 
-def gen_embedding_rules(n: Node, symbols, embedding_dim, counter):
+def gen_embedding_rules(
+    n: Node, symbols: _SymbolDict, embedding_dim: int | DVar, counter: int
+) -> tuple[list[Constraint], int]:
     embedding_output, counter = gen_tvar(counter)
     symbols[n] = embedding_output
-    embedding_input = symbols[n.args[0]]
+    embedding_input = symbols[n.args[0]]  # pyrefly: ignore[bad-index]
 
     input_dyn = BinConstraintT(embedding_input, Dyn, op_eq)
     output_dyn = BinConstraintT(embedding_output, Dyn, op_eq)
@@ -520,18 +577,22 @@ def gen_embedding_rules(n: Node, symbols, embedding_dim, counter):
 
 
 @register_inference_rule(torch.tensor)
-def tensor_inference_rule(n: Node, symbols, constraints, counter):
+def tensor_inference_rule(
+    n: Node, symbols: _SymbolDict, constraints: list[Constraint], counter: int
+) -> tuple[list[Constraint], int]:
     """
     If the tensor is a scalar, we will skip it since we
     do not support scalars yet. We will add support in the future
     if it's needed. For our examples so far, scalars are not needed.
     """
-    return [], counter
+    return [], counter  # pyrefly: ignore[implicit-any]
 
 
 @register_inference_rule("reshape")
 @register_inference_rule("view")
-def view_inference_rule(n: Node, symbols, constraints, counter):
+def view_inference_rule(
+    n: Node, symbols: _SymbolDict, constraints: list[Constraint], counter: int
+) -> tuple[list[Constraint], int]:
     """
     Similar to reshape but with an extra condition on the strides
     """
@@ -564,7 +625,7 @@ def view_inference_rule(n: Node, symbols, constraints, counter):
     t2_type = TensorType(t2_type)  # type: ignore[assignment]
 
     c1 = BinConstraintT(my_view, t2_type, op_eq)
-    c2 = CanReshape(src_var, t2_type)
+    c2 = CanReshape(src_var, t2_type)  # pyrefly: ignore[bad-argument-type]
 
     # TODO: add the extra check mentioned here:
     # https://pytorch.org/docs/stable/generated/torch.Tensor.view.html#torch.Tensor.view
@@ -573,7 +634,9 @@ def view_inference_rule(n: Node, symbols, constraints, counter):
 
 
 @register_inference_rule("size")
-def size_inference_rule(n: Node, symbols, constraints, counter):
+def size_inference_rule(
+    n: Node, symbols: _SymbolDict, constraints: list[Constraint], counter: int
+) -> tuple[list[Constraint], int]:
     """
     The constraint is just lhs = rhs.
     Ex: size = input_ids.size()
@@ -583,7 +646,7 @@ def size_inference_rule(n: Node, symbols, constraints, counter):
         # generate the new variable
         size, counter = gen_tvar(counter)
         symbols[n] = size
-        input = symbols[n.args[0]]
+        input = symbols[n.args[0]]  # pyrefly: ignore[bad-index]
         c = BinConstraintT(input, size, op_eq)
         return [c], counter
 
@@ -593,9 +656,14 @@ def size_inference_rule(n: Node, symbols, constraints, counter):
             # generate the new variable
             size_index, counter = gen_dvar(counter)
             symbols[n] = size_index
-            input = symbols[n.args[0]]
+            input = symbols[n.args[0]]  # pyrefly: ignore[bad-index]
             c2 = [
-                GetItem(i + 1, n.args[1], size_index, input)
+                GetItem(
+                    i + 1,
+                    n.args[1],
+                    size_index,
+                    input,  # pyrefly: ignore[bad-argument-type]
+                )
                 for i in range(MAX_TENSOR_RANK)
             ]
             c3 = BinConstraintD(0, size_index, op_leq)
@@ -613,7 +681,7 @@ def size_inference_rule(n: Node, symbols, constraints, counter):
         raise NotImplementedError
 
 
-def range_check(i, n):
+def range_check(i: int, n: int) -> T | F:
     """
     Checks if an index i is within range of a size n list
     Args:
@@ -629,7 +697,9 @@ def range_check(i, n):
 
 
 @register_inference_rule(torch.cumsum)
-def cumsum_inference_rule(n: Node, symbols, constraints, counter):
+def cumsum_inference_rule(
+    n: Node, symbols: _SymbolDict, constraints: list[Constraint], counter: int
+) -> tuple[list[Constraint], int]:
     """
     Input and output shapes should be equal
     We should verify that the index is valid
@@ -668,14 +738,18 @@ def cumsum_inference_rule(n: Node, symbols, constraints, counter):
 
 
 @register_inference_rule(_assert_is_none)
-def assert_inference_rule(n: Node, symbols, constraints, counter):
+def assert_inference_rule(
+    n: Node, symbols: _SymbolDict, constraints: list[Constraint], counter: int
+) -> tuple[list[Constraint], int]:
     if len(n.users) != 0:
         raise AssertionError(f"Expected no users, got {len(n.users)}")
-    return [], counter
+    return [], counter  # pyrefly: ignore[implicit-any]
 
 
 @register_inference_rule(operator.getitem)
-def getitem_inference_rule(n: Node, symbols, constraints, counter):
+def getitem_inference_rule(
+    n: Node, symbols: _SymbolDict, constraints: list[Constraint], counter: int
+) -> tuple[list[Constraint], int]:
     if not isinstance(n.args[0], Node):
         raise AssertionError(f"Expected Node, got {type(n.args[0])}")
 
@@ -732,7 +806,7 @@ def getitem_inference_rule(n: Node, symbols, constraints, counter):
             ]
         else:
             # TODO: we should figure out why there is a key-error here.
-            return [], counter
+            return [], counter  # pyrefly: ignore[implicit-any]
 
         return [Disj([c1, *c2])], counter
 
@@ -741,7 +815,9 @@ def getitem_inference_rule(n: Node, symbols, constraints, counter):
 
 
 @register_inference_rule(operator.gt)
-def gt_inference_rule(n: Node, symbols, constraints, counter):
+def gt_inference_rule(
+    n: Node, symbols: _SymbolDict, constraints: list[Constraint], counter: int
+) -> tuple[list[Constraint], int]:
     if not isinstance(n.args[0], (Node, int)):
         raise AssertionError(f"Expected Node or int, got {type(n.args[0])}")
     if not isinstance(n.args[1], (Node, int)):
@@ -804,7 +880,9 @@ def gt_inference_rule(n: Node, symbols, constraints, counter):
 
 
 @register_inference_rule(operator.eq)
-def eq_inference_rule(n: Node, symbols, constraints, counter):
+def eq_inference_rule(
+    n: Node, symbols: _SymbolDict, constraints: list[Constraint], counter: int
+) -> tuple[list[Constraint], int]:
     if not isinstance(n.args[0], (Node, int)):
         raise AssertionError(f"Expected Node or int, got {type(n.args[0])}")
     if not isinstance(n.args[1], (Node, int)):
@@ -845,7 +923,9 @@ def eq_inference_rule(n: Node, symbols, constraints, counter):
 
 
 @register_inference_rule(operator.ne)
-def neq_inference_rule(n: Node, symbols, constraints, counter):
+def neq_inference_rule(
+    n: Node, symbols: _SymbolDict, constraints: list[Constraint], counter: int
+) -> tuple[list[Constraint], int]:
     """
     Translates to inconsistent in gradual types.
     To prove inequality, we should prove that
@@ -973,7 +1053,9 @@ def neq_inference_rule(n: Node, symbols, constraints, counter):
 
 
 @register_inference_rule(operator.lt)
-def lt_inference_rule(n: Node, symbols, constraints, counter):
+def lt_inference_rule(
+    n: Node, symbols: _SymbolDict, constraints: list[Constraint], counter: int
+) -> tuple[list[Constraint], int]:
     if not isinstance(n.args[0], (Node, int)):
         raise AssertionError(f"Expected Node or int, got {type(n.args[0])}")
     if not isinstance(n.args[1], (Node, int)):
@@ -1018,7 +1100,9 @@ def lt_inference_rule(n: Node, symbols, constraints, counter):
 
 
 @register_inference_rule(torch.full)
-def full_inference_rule(n: Node, symbols, constraints, counter):
+def full_inference_rule(
+    n: Node, symbols: _SymbolDict, constraints: list[Constraint], counter: int
+) -> tuple[list[Constraint], int]:
     full, counter = gen_tvar(counter)
     symbols[n] = full
     res = []
@@ -1026,7 +1110,9 @@ def full_inference_rule(n: Node, symbols, constraints, counter):
     if not isinstance(n.args[0], Iterable):
         raise AssertionError(f"Expected Iterable, got {type(n.args[0])}")
     for arg in n.args[0]:
-        dim = arg if isinstance(arg, int) else symbols[arg]
+        dim = (
+            arg if isinstance(arg, int) else symbols[arg]  # pyrefly: ignore[bad-index]
+        )
         res.append(dim)
     c = BinConstraintT(full, TensorType(list(res)), op_eq)  # type: ignore[arg-type]
     return [c], counter
@@ -1034,12 +1120,14 @@ def full_inference_rule(n: Node, symbols, constraints, counter):
 
 # TODO normalize index
 @register_inference_rule(torch.arange)
-def arange_inference_rule(n: Node, symbols, constraints, counter):
+def arange_inference_rule(
+    n: Node, symbols: _SymbolDict, constraints: list[Constraint], counter: int
+) -> tuple[list[Constraint], int]:
     start = 0
     step = 1
 
     if len(n.args) == 1:
-        end = symbols[n.args[0]]
+        end = symbols[n.args[0]]  # pyrefly: ignore[bad-index]
     else:
         raise NotImplementedError("Not yet implemented")
 
@@ -1078,7 +1166,9 @@ def arange_inference_rule(n: Node, symbols, constraints, counter):
     ], counter
 
 
-def gen_broadcasting_constraints(e1, e2, symbols, counter, output_var):
+def gen_broadcasting_constraints(
+    e1: TVar, e2: TVar, symbols: _SymbolDict, counter: int, output_var: TVar
+) -> tuple[list[Constraint], int]:
     # additional vars that don't correspond to expressions
     e11, counter = gen_tvar(counter)
     e22, counter = gen_tvar(counter)
@@ -1095,7 +1185,9 @@ def gen_broadcasting_constraints(e1, e2, symbols, counter, output_var):
 @register_inference_rule("ne")
 @register_inference_rule(torch.add)
 @register_inference_rule(operator.add)
-def broadcasting_inference_rule(n: Node, symbols, constraints, counter):
+def broadcasting_inference_rule(
+    n: Node, symbols: _SymbolDict, constraints: list[Constraint], counter: int
+) -> tuple[list[Constraint], int]:  # pyrefly: ignore[bad-return]
     op_code = None
     if n.target is operator.add or n.target is torch.add:
         op_code = op_add
@@ -1111,7 +1203,13 @@ def broadcasting_inference_rule(n: Node, symbols, constraints, counter):
             e1 = symbols[n.args[0]]
             e2 = symbols[n.args[1]]
 
-            return gen_broadcasting_constraints(e1, e2, symbols, counter, my_output)
+            return gen_broadcasting_constraints(
+                e1,  # pyrefly: ignore[bad-argument-type]
+                e2,  # pyrefly: ignore[bad-argument-type]
+                symbols,
+                counter,
+                my_output,
+            )
         else:
             raise NotImplementedError("Method not yet implemented")
 
@@ -1168,7 +1266,9 @@ def broadcasting_inference_rule(n: Node, symbols, constraints, counter):
 
 
 @register_inference_rule(torch.flatten)
-def flatten_inference_rule(n: Node, symbols, constraints, counter):
+def flatten_inference_rule(
+    n: Node, symbols: _SymbolDict, constraints: list[Constraint], counter: int
+) -> tuple[list[Constraint], int]:
     if not isinstance(n.args[0], Node):
         raise AssertionError(f"Expected Node, got {type(n.args[0])}")
 
@@ -1199,7 +1299,12 @@ def flatten_inference_rule(n: Node, symbols, constraints, counter):
     const = []
     for i in range(1, MAX_TENSOR_RANK + 1):
         c, counter = generate_flatten_constraints(
-            start_dim, end_dim, input, flattened, i, counter
+            start_dim,
+            end_dim,
+            input,  # pyrefly: ignore[bad-argument-type]
+            flattened,
+            i,
+            counter,
         )
         const.append(c)
 
@@ -1207,17 +1312,30 @@ def flatten_inference_rule(n: Node, symbols, constraints, counter):
 
 
 @register_inference_rule(torch.nn.functional.layer_norm)
-def layer_norm_functional(n: Node, symbols, constraints, counter):
+def layer_norm_functional(
+    n: Node, symbols: _SymbolDict, constraints: list[Constraint], counter: int
+) -> tuple[list[Constraint], int]:
     """
     We generate the constraint: input = output
     """
     if not isinstance(n.args[0], Node):
         raise AssertionError(f"Expected Node, got {type(n.args[0])}")
-    return gen_layer_norm_constraints(n, n.args[1], symbols, counter)
+    return gen_layer_norm_constraints(
+        n,
+        n.args[1],  # pyrefly: ignore[bad-argument-type]
+        symbols,
+        counter,
+    )
 
 
 @register_inference_rule(torch.nn.LayerNorm)
-def layer_norm_inference_rule(n: Node, module_instance, symbols, constraints, counter):
+def layer_norm_inference_rule(
+    n: Node,
+    module_instance: torch.nn.LayerNorm,
+    symbols: _SymbolDict,
+    constraints: list[Constraint],
+    counter: int,
+) -> tuple[list[Constraint], int]:
     """
     Input and output shapes should be equal.
     Input should be consistent with the normalized_shape
@@ -1229,10 +1347,12 @@ def layer_norm_inference_rule(n: Node, module_instance, symbols, constraints, co
     )
 
 
-def gen_layer_norm_constraints(n: Node, normalized_shape, symbols, counter):
+def gen_layer_norm_constraints(
+    n: Node, normalized_shape: Sequence[int], symbols: _SymbolDict, counter: int
+) -> tuple[list[Constraint], int]:
     output, counter = gen_tvar(counter)
     symbols[n] = output
-    input = symbols[n.args[0]]
+    input = symbols[n.args[0]]  # pyrefly: ignore[bad-index]
 
     input_dyn = BinConstraintT(input, Dyn, op_eq)
     output_dyn = BinConstraintT(output, Dyn, op_eq)
@@ -1258,7 +1378,13 @@ def gen_layer_norm_constraints(n: Node, normalized_shape, symbols, counter):
 
 @register_inference_rule(torch.nn.Dropout)
 @register_inference_rule(torch.nn.ReLU)
-def relu_inference_rule(n: Node, module_instance, symbols, constraints, counter):
+def relu_inference_rule(
+    n: Node,
+    module_instance: torch.nn.Module,
+    symbols: _SymbolDict,
+    constraints: list[Constraint],
+    counter: int,
+) -> tuple[list[Constraint], int]:
     """
     Input and output shapes should be equal.
     """
@@ -1273,7 +1399,13 @@ def relu_inference_rule(n: Node, module_instance, symbols, constraints, counter)
 
 
 @register_inference_rule(torch.nn.Linear)
-def linear_inference_rule(n: Node, module_instance, symbols, constraints, counter):
+def linear_inference_rule(
+    n: Node,
+    module_instance: torch.nn.Linear,
+    symbols: _SymbolDict,
+    constraints: list[Constraint],
+    counter: int,
+) -> tuple[list[Constraint], int]:
     """
     Input and output sizes should be the same except for the last dimension
     If the input is Dyn, then so should the output
@@ -1286,7 +1418,9 @@ def linear_inference_rule(n: Node, module_instance, symbols, constraints, counte
 
 
 @register_inference_rule("dim")
-def torch_dim_inference_rule(n: Node, symbols, constraints, counter):
+def torch_dim_inference_rule(
+    n: Node, symbols: _SymbolDict, constraints: list[Constraint], counter: int
+) -> tuple[list[Constraint], int]:
     if not isinstance(n.args[0], Node):
         raise AssertionError(f"Expected Node, got {type(n.args[0])}")
     my_dim, counter = gen_dvar(counter)
@@ -1313,12 +1447,16 @@ def torch_dim_inference_rule(n: Node, symbols, constraints, counter):
 
 
 @register_inference_rule(torch._C._nn.linear)
-def torch_linear_inference_rule(n: Node, symbols, constraints, counter):
+def torch_linear_inference_rule(
+    n: Node, symbols: _SymbolDict, constraints: list[Constraint], counter: int
+) -> tuple[list[Constraint], int]:
     if not isinstance(n.args[0], Node):
         raise AssertionError(f"Expected Node, got {type(n.args[0])}")
     weight_dims, counter = gen_tensor_dims(2, counter)
     equality_constraint = BinConstraintT(
-        symbols[n.args[1]], TensorType(weight_dims), op_eq
+        symbols[n.args[1]],  # pyrefly: ignore[bad-index]
+        TensorType(weight_dims),
+        op_eq,
     )
     constraints, counter = linear_constraints(
         n, weight_dims[1], weight_dims[0], symbols, counter
@@ -1326,10 +1464,16 @@ def torch_linear_inference_rule(n: Node, symbols, constraints, counter):
     return [equality_constraint] + constraints, counter
 
 
-def linear_constraints(n: Node, in_features, out_features, symbols, counter):
+def linear_constraints(
+    n: Node,
+    in_features: int | DVar,
+    out_features: int | DVar,
+    symbols: _SymbolDict,
+    counter: int,
+) -> tuple[list[Constraint], int]:
     linear_output, counter = gen_tvar(counter)
     symbols[n] = linear_output
-    linear_input = symbols[n.args[0]]
+    linear_input = symbols[n.args[0]]  # pyrefly: ignore[bad-index]
 
     input_dyn = BinConstraintT(linear_input, Dyn, op_eq)
     output_dyn = BinConstraintT(linear_output, Dyn, op_eq)
@@ -1357,7 +1501,9 @@ def linear_constraints(n: Node, in_features, out_features, symbols, counter):
     return [Disj([c1, Disj(c2)])], counter
 
 
-def add_layer_norm_constraints(input_dim, normalized_dim):
+def add_layer_norm_constraints(
+    input_dim: list[DVar], normalized_dim: list[int]
+) -> list[Constraint]:
     """
     The constraints say that the type has te form: [*, 1024, 1024]
      while the normalized_dim have the form [1024, 1024]
@@ -1372,16 +1518,21 @@ def add_layer_norm_constraints(input_dim, normalized_dim):
         return [F()]
 
     else:
-        constraints = []
+        constraints: list[Constraint] = []
         for i, n in zip(reversed(input_dim), reversed(normalized_dim)):
             constraints.append(BinConstraintD(i, n, op_consistency))
         return constraints
 
 
-def add_linear_constraints(dims1, dims2, in_features, out_features):
+def add_linear_constraints(
+    dims1: list[DVar],
+    dims2: list[DVar],
+    in_features: int | DVar,
+    out_features: int | DVar,
+) -> list[Constraint]:
     if len(dims1) != len(dims2):
         raise AssertionError(f"Expected same length, got {len(dims1)} vs {len(dims2)}")
-    constraints = []
+    constraints: list[Constraint] = []
     for i in range(len(dims1)):
         if i == len(dims1) - 1:
             constraints.append(BinConstraintD(dims1[i], in_features, op_consistency))
@@ -1393,7 +1544,9 @@ def add_linear_constraints(dims1, dims2, in_features, out_features):
 
 
 @register_inference_rule(torch.reshape)
-def reshape_inference_rule(n: Node, symbols, constraints, counter):
+def reshape_inference_rule(
+    n: Node, symbols: _SymbolDict, constraints: list[Constraint], counter: int
+) -> tuple[list[Constraint], int]:
     if not isinstance(n.args[0], Node):
         raise AssertionError(f"Expected Node, got {type(n.args[0])}")
 
@@ -1405,13 +1558,19 @@ def reshape_inference_rule(n: Node, symbols, constraints, counter):
     t2 = n.args[1]
     t2_type = TensorType([Dyn if elem == -1 else elem for elem in t2])  # type: ignore[union-attr]
     c1 = BinConstraintT(my_reshape, t2_type, op_eq)  # type: ignore[union-attr]
-    c2 = CanReshape(src_var, t2_type)
+    c2 = CanReshape(src_var, t2_type)  # pyrefly: ignore[bad-argument-type]
 
     return [c1, c2], counter
 
 
 @register_inference_rule(BatchNorm2d)
-def batchnorm_inference_rule(n: Node, module_instance, symbols, constraints, counter):
+def batchnorm_inference_rule(
+    n: Node,
+    module_instance: BatchNorm2d,
+    symbols: _SymbolDict,
+    constraints: list[Constraint],
+    counter: int,
+) -> tuple[list[Constraint], int]:
     if not isinstance(n.args[0], Node):
         raise AssertionError(f"Expected Node, got {type(n.args[0])}")
 
@@ -1434,7 +1593,13 @@ def batchnorm_inference_rule(n: Node, module_instance, symbols, constraints, cou
 
 
 @register_inference_rule(torch.nn.AdaptiveAvgPool2d)
-def adaptive_inference_rule(n: Node, module_instance, symbols, constraints, counter):
+def adaptive_inference_rule(
+    n: Node,
+    module_instance: torch.nn.AdaptiveAvgPool2d,
+    symbols: _SymbolDict,
+    constraints: list[Constraint],
+    counter: int,
+) -> tuple[list[Constraint], int]:
     if not isinstance(n.args[0], Node):
         raise AssertionError(f"Expected Node, got {type(n.args[0])}")
 
@@ -1453,7 +1618,16 @@ def adaptive_inference_rule(n: Node, module_instance, symbols, constraints, coun
     c2 = BinConstraintT(
         avg_pool,
         TensorType(
-            [d1, d2, module_instance.output_size[0], module_instance.output_size[1]]
+            [
+                d1,
+                d2,
+                module_instance.output_size[  # pyrefly: ignore[bad-index, unsupported-operation]
+                    0
+                ],
+                module_instance.output_size[  # pyrefly: ignore[bad-index, unsupported-operation]
+                    1
+                ],
+            ]
         ),
         op_eq,
     )
@@ -1462,7 +1636,13 @@ def adaptive_inference_rule(n: Node, module_instance, symbols, constraints, coun
 
 
 @register_inference_rule(Conv2d)
-def conv2d_inference_rule(n: Node, module_instance, symbols, constraints, counter):
+def conv2d_inference_rule(
+    n: Node,
+    module_instance: Conv2d,
+    symbols: _SymbolDict,
+    constraints: list[Constraint],
+    counter: int,
+) -> tuple[list[Constraint], int]:
     if not isinstance(n.args[0], Node):
         raise AssertionError(f"Expected Node, got {type(n.args[0])}")
 
@@ -1481,12 +1661,12 @@ def conv2d_inference_rule(n: Node, module_instance, symbols, constraints, counte
 
     c3 = CalcConv(
         my_conv,
-        input_var,
+        input_var,  # pyrefly: ignore[bad-argument-type]
         module_instance.out_channels,
-        module_instance.kernel_size,
-        module_instance.padding,
-        module_instance.stride,
-        module_instance.dilation,
+        module_instance.kernel_size,  # pyrefly: ignore[bad-argument-type]
+        module_instance.padding,  # pyrefly: ignore[bad-argument-type]
+        module_instance.stride,  # pyrefly: ignore[bad-argument-type]
+        module_instance.dilation,  # pyrefly: ignore[bad-argument-type]
         [d1, d2, d3, d4],
     )
 
@@ -1496,7 +1676,13 @@ def conv2d_inference_rule(n: Node, module_instance, symbols, constraints, counte
 
 
 @register_inference_rule(torch.nn.MaxPool2d)
-def maxpool_inference_rule(n: Node, module_instance, symbols, constraints, counter):
+def maxpool_inference_rule(
+    n: Node,
+    module_instance: torch.nn.MaxPool2d,
+    symbols: _SymbolDict,
+    constraints: list[Constraint],
+    counter: int,
+) -> tuple[list[Constraint], int]:
     if not isinstance(n.args[0], Node):
         raise AssertionError(f"Expected Node, got {type(n.args[0])}")
     maxpool, counter = gen_tvar(counter)
@@ -1510,7 +1696,7 @@ def maxpool_inference_rule(n: Node, module_instance, symbols, constraints, count
 
     c2 = CalcMaxPool(
         maxpool,
-        input_var,
+        input_var,  # pyrefly: ignore[bad-argument-type]
         module_instance.kernel_size,
         module_instance.padding,
         module_instance.stride,
@@ -1524,21 +1710,21 @@ def maxpool_inference_rule(n: Node, module_instance, symbols, constraints, count
 
 
 class ConstraintGenerator:
-    def __init__(self, traced, graph=None):
+    def __init__(self, traced: torch.nn.Module, graph: Graph | None = None) -> None:
         self.traced = traced  # traced or tracer.root
         self.traced_params = dict(self.traced.named_parameters())
         self.constraints = []
         self.symbol_dict = {}
         self.graph = traced.graph if hasattr(traced, "graph") else graph
 
-    def generate_constraints(self, counter=0):
+    def generate_constraints(self, counter: int = 0) -> tuple[Conj, int]:
         """
         Iterate through every node and generate constraints
         Effect: self.constraints will be populated with the final constraints
         """
         graph = self.graph
 
-        all_constraints = []
+        all_constraints: list[Constraint] = []
 
         # pyrefly: ignore [missing-attribute]
         for n in graph.nodes:
@@ -1547,7 +1733,9 @@ def generate_constraints(self, counter=0):
 
         return Conj(all_constraints), counter
 
-    def generate_constraints_node(self, n: Node, counter):
+    def generate_constraints_node(
+        self, n: Node, counter: int
+    ) -> tuple[list[Constraint], int]:
         """
         Generate constraints the given node:
         Currently supported operations:
@@ -1586,7 +1774,9 @@ def generate_constraints_node(self, n: Node, counter):
                 )
 
         elif n.op == "call_module":
-            module_instance = self.traced.get_submodule(n.target)
+            module_instance = self.traced.get_submodule(
+                n.target  # pyrefly: ignore[bad-argument-type]
+            )
             if type(module_instance) in _INFERENCE_RULES:
                 return _INFERENCE_RULES[type(module_instance)](
                     n, module_instance, self.symbol_dict, self.constraints, counter
@@ -1607,7 +1797,9 @@ def generate_constraints_node(self, n: Node, counter):
                 )
 
         elif n.op == "get_attr":
-            t = self.traced_params.get(n.target, None)
+            t = self.traced_params.get(  # pyrefly: ignore[no-matching-overload]
+                n.target, None
+            )
 
             if isinstance(t, torch.Tensor):
                 if len(t.shape) > 0:
@@ -1618,12 +1810,12 @@ def generate_constraints_node(self, n: Node, counter):
                     return [BinConstraintT(output, attr_type, op_eq)], counter
                 else:
                     # scalar?
-                    return [], counter
+                    return [], counter  # pyrefly: ignore[implicit-any]
             else:
-                return [], counter
+                return [], counter  # pyrefly: ignore[implicit-any]
 
         elif n.op == "output":
-            return [], counter
+            return [], counter  # pyrefly: ignore[implicit-any]
 
         else:
             raise NotImplementedError(f"Method {n.op} not yet implemented")
diff --git a/torch/fx/experimental/migrate_gradual_types/constraint_transformation.py b/torch/fx/experimental/migrate_gradual_types/constraint_transformation.py
index 9ec1006e7fedf..20060b786c0d4 100644
--- a/torch/fx/experimental/migrate_gradual_types/constraint_transformation.py
+++ b/torch/fx/experimental/migrate_gradual_types/constraint_transformation.py
@@ -1,7 +1,6 @@
-# mypy: ignore-errors
 import copy
 import itertools
-from collections.abc import Callable
+from collections.abc import Callable, Sequence
 
 from torch.fx.experimental.migrate_gradual_types.constraint import (
     ApplyBroadcasting,
@@ -47,14 +46,54 @@
     gen_nat_constraints,
     gen_tensor_dims,
 )
-from torch.fx.tensor_type import Dyn, TensorType
-
-
-_TRANSFORMATION_RULES: dict[Constraint, Callable] = {}
-
-
-def register_transformation_rule(call_target):
-    def register(fn):
+from torch.fx.tensor_type import _DynType, Dyn, TensorType
+
+
+__all__ = [
+    "apply_padding",
+    "broadcast_dim",
+    "calc_last_two_dims",
+    "create_equality_constraints_for_broadcasting",
+    "gen_all_reshape_possibilities",
+    "gen_broadcasting_constraints",
+    "gen_consistency_constraints",
+    "gen_greatest_upper_bound",
+    "gen_lists_of_dims",
+    "generate_all_broadcasting_possibilities_no_padding",
+    "generate_all_int_dyn_dim_possibilities",
+    "generate_binconstraint_d",
+    "generate_binconstraint_t",
+    "generate_broadcasting",
+    "generate_calc_conv",
+    "generate_calc_maxpool",
+    "generate_calc_product",
+    "generate_conj",
+    "generate_d_gub",
+    "generate_disj",
+    "generate_gub",
+    "generate_reshape",
+    "is_dim_div_by_target",
+    "is_target_div_by_dim",
+    "no_broadcast_dim_with_index",
+    "register_transformation_rule",
+    "transform_constraint",
+    "transform_get_item",
+    "transform_get_item_tensor",
+    "transform_index_select",
+    "transform_transpose",
+    "valid_index",
+    "valid_index_tensor",
+]
+
+
+_TransformFn = Callable[[Constraint, int], tuple[Constraint, int]]
+_TRANSFORMATION_RULES: dict[type, _TransformFn] = {}
+
+
+def register_transformation_rule(
+    call_target: type[Constraint],
+) -> Callable[[_TransformFn], _TransformFn]:
+    def register(fn: _TransformFn) -> _TransformFn:
         if call_target in _TRANSFORMATION_RULES:
             raise RuntimeError(
                 f"Transformation rule already registered for {call_target}!"
@@ -65,7 +104,7 @@ def register(fn):
     return register
 
 
-def valid_index(index, dims):
+def valid_index(index: int, dims: list[DVar]) -> Constraint:
     """
     Given a list of dimensions, checks if an index is valid in the list
     """
@@ -77,10 +116,12 @@ def valid_index(index, dims):
 
 
 @register_transformation_rule(Transpose)
-def transform_transpose(constraint, counter):
+def transform_transpose(constraint: Constraint, counter: int) -> tuple[Constraint, int]:
     """
     Similar to a sequence of two index-selects
     """
+    if not isinstance(constraint, Transpose):
+        raise TypeError(type(constraint))
     dims, counter = gen_tensor_dims(constraint.tensor_size, counter)
     is_valid_index1 = valid_index(constraint.index1, dims)
     is_valid_index2 = valid_index(constraint.index2, dims)
@@ -104,21 +145,25 @@ def transform_transpose(constraint, counter):
 
 
 @register_transformation_rule(IndexSelect)
-def transform_index_select(constraint, counter):
+def transform_index_select(
+    constraint: Constraint, counter: int
+) -> tuple[Constraint, int]:
     """
     The constraints consider the given tensor size, checks if the index is valid
     and if so, generates a constraint for replacing the input dimension
     with the required dimension
     """
+    if not isinstance(constraint, IndexSelect):
+        raise TypeError(type(constraint))
     dims, counter = gen_tensor_dims(constraint.tensor_size, counter)
     is_valid_index = valid_index(constraint.index, dims)
     nat_constraints = gen_nat_constraints(dims)
 
     # if the index is valid then replace the input dimension with the new dimension
     # otherwise the dimension will not be replaced and the clause will contain False
+    new_dims = copy.deepcopy(dims)
     if is_valid_index == T():
-        new_dims = copy.deepcopy(dims)
-        new_dims[constraint.index] = constraint.dim_replace
+        new_dims[constraint.index] = constraint.dim_replace  # type: ignore[unsupported-operation]
 
     transformed_constraint = Conj(
         [
@@ -134,7 +179,7 @@ def transform_index_select(constraint, counter):
 
 
 @register_transformation_rule(GetItem)
-def transform_get_item(constraint, counter):
+def transform_get_item(constraint: Constraint, counter: int) -> tuple[Constraint, int]:
     """
     generate an equality of the form:
     t = [a1, ..., an]
@@ -149,6 +194,8 @@ def transform_get_item(constraint, counter):
     Returns: simplified constraints for GetItem
 
     """
+    if not isinstance(constraint, GetItem):
+        raise TypeError(type(constraint))
     dims, counter = gen_tensor_dims(constraint.tensor_size, counter)
     nat_constraints = gen_nat_constraints(dims)
 
@@ -170,7 +217,7 @@ def transform_get_item(constraint, counter):
     return Conj(all_constraints), counter
 
 
-def valid_index_tensor(index, dims):
+def valid_index_tensor(index: tuple[None | slice, ...], dims: list[DVar]) -> Constraint:
     """
     if the slice instances exceed the length of the dimensions
     then this is a type error so we return False
@@ -186,7 +233,9 @@ def valid_index_tensor(index, dims):
 
 
 @register_transformation_rule(GetItemTensor)
-def transform_get_item_tensor(constraint, counter):
+def transform_get_item_tensor(
+    constraint: Constraint, counter: int
+) -> tuple[Constraint, int]:
     """
     When the index is a tuple, then the output will be a tensor
     TODO: we have to check if this is the case for all HF models
@@ -200,6 +249,8 @@ def transform_get_item_tensor(constraint, counter):
 
      slice with default arguments does not change the rank
     """
+    if not isinstance(constraint, GetItemTensor):
+        raise TypeError(type(constraint))
     if not isinstance(constraint.index_tuple, tuple):
         raise AssertionError(
             f"Expected tuple for index_tuple, got {type(constraint.index_tuple)}"
@@ -212,7 +263,8 @@ def transform_get_item_tensor(constraint, counter):
     # generate a place-holder list of the right rank
     # where "slice" does not contribute to the rank and "None" does
     none_c = constraint.index_tuple.count(None)
-    resulting_tensor_dims = (none_c + len(dims)) * [None]
+    # list invariance: [None] * n types as list[None], but elements are reassigned to int/DVar
+    resulting_tensor_dims: list[int | DVar | None] = [None] * (none_c + len(dims))  # type: ignore[assignment]
 
     dim_index = 0
     for i in range(len(constraint.index_tuple)):
@@ -251,10 +303,14 @@ def transform_get_item_tensor(constraint, counter):
 
 
 @register_transformation_rule(BinConstraintT)
-def generate_binconstraint_t(constraint, counter):
+def generate_binconstraint_t(
+    constraint: Constraint, counter: int
+) -> tuple[Constraint, int]:
     """
     Transform binary constraints for tensors
     """
+    if not isinstance(constraint, BinConstraintT):
+        raise TypeError(type(constraint))
 
     # precision constraints
     if constraint.op == op_precision:
@@ -280,6 +336,8 @@ def generate_binconstraint_t(constraint, counter):
                     + [BinConstraintD(1, new_dim, op_leq) for new_dim in new_dims]
                 )
                 return Conj(new_dim_constraints), counter
+        else:
+            return constraint, counter
 
     # matching
     elif constraint.op == op_matching:
@@ -342,15 +400,21 @@ def generate_binconstraint_t(constraint, counter):
 
 
 @register_transformation_rule(BinConstraintD)
-def generate_binconstraint_d(constraint, counter):
+def generate_binconstraint_d(
+    constraint: Constraint, counter: int
+) -> tuple[Constraint, int]:
     """
     Transform binary constraints for dimensions
     """
+    if not isinstance(constraint, BinConstraintD):
+        raise TypeError(type(constraint))
     if constraint.op == op_precision:
         if isinstance(constraint.lhs, int):
             return BinConstraintD(constraint.lhs, constraint.rhs, op_eq), counter
         elif constraint.lhs == Dyn:
             return T(), counter
+        else:
+            return constraint, counter
 
     elif constraint.op == op_consistency:
         return (
@@ -369,10 +433,12 @@ def generate_binconstraint_d(constraint, counter):
 
 
 @register_transformation_rule(Conj)
-def generate_conj(constraint, counter):
+def generate_conj(constraint: Constraint, counter: int) -> tuple[Constraint, int]:
     """
     Transform conjunctions
     """
+    if not isinstance(constraint, Conj):
+        raise TypeError(type(constraint))
     new = []
     for c in constraint.conjucts:
         new_c, counter = transform_constraint(c, counter)
@@ -381,10 +447,12 @@ def generate_conj(constraint, counter):
 
 
 @register_transformation_rule(Disj)
-def generate_disj(constraint, counter):
+def generate_disj(constraint: Constraint, counter: int) -> tuple[Constraint, int]:
     """
     Transform disjunctions
     """
+    if not isinstance(constraint, Disj):
+        raise TypeError(type(constraint))
     new = []
     for c in constraint.disjuncts:
         new_c, counter = transform_constraint(c, counter)
@@ -393,11 +461,13 @@ def generate_disj(constraint, counter):
 
 
 @register_transformation_rule(TGreatestUpperBound)
-def generate_gub(constraint, counter):
+def generate_gub(constraint: Constraint, counter: int) -> tuple[Constraint, int]:
     """
     Transform greatest upper bound for tensors. Results in equality and Greatest Upper Bound
     on dimensions
     """
+    if not isinstance(constraint, TGreatestUpperBound):
+        raise TypeError(type(constraint))
     c1 = Conj(
         [
             Disj(
@@ -416,10 +486,12 @@ def generate_gub(constraint, counter):
 
 
 @register_transformation_rule(DGreatestUpperBound)
-def generate_d_gub(constraint, counter):
+def generate_d_gub(constraint: Constraint, counter: int) -> tuple[Constraint, int]:
     """
     Transform greatest upper bound for dimensions into equality constraints
     """
+    if not isinstance(constraint, DGreatestUpperBound):
+        raise TypeError(type(constraint))
     c1 = Conj(
         [
             BinConstraintD(constraint.rhs1, Dyn, op_eq),
@@ -442,7 +514,9 @@ def generate_d_gub(constraint, counter):
 
 
 @register_transformation_rule(CalcConv)
-def generate_calc_conv(constraint, counter):
+def generate_calc_conv(constraint: Constraint, counter: int) -> tuple[Constraint, int]:
+    if not isinstance(constraint, CalcConv):
+        raise TypeError(type(constraint))
     d, counter = gen_tensor_dims(4, counter)
     conv_result = TensorType([d[0], d[1], d[2], d[3]])
 
@@ -475,10 +549,14 @@ def generate_calc_conv(constraint, counter):
 
 
 @register_transformation_rule(CalcMaxPool)
-def generate_calc_maxpool(constraint, counter):
+def generate_calc_maxpool(
+    constraint: Constraint, counter: int
+) -> tuple[Constraint, int]:
     """
     Transform maxpool constraints
     """
+    if not isinstance(constraint, CalcMaxPool):
+        raise TypeError(type(constraint))
     d, counter = gen_tensor_dims(4, counter)
     maxpool_result = TensorType([d[0], d[1], d[2], d[3]])
 
@@ -503,10 +581,14 @@ def generate_calc_maxpool(constraint, counter):
 
 
 @register_transformation_rule(CalcProduct)
-def generate_calc_product(constraint, counter):
+def generate_calc_product(
+    constraint: Constraint, counter: int
+) -> tuple[Constraint, int]:
     """
     Transform flatten constraints
     """
+    if not isinstance(constraint, CalcProduct):
+        raise TypeError(type(constraint))
     start = constraint.start
     end = constraint.end
     dims = constraint.dims_to_flatten
@@ -524,7 +606,7 @@ def generate_calc_product(constraint, counter):
 
     all_possibilities = generate_all_int_dyn_dim_possibilities(mid)
 
-    all_constraints = []
+    all_constraints: list[Constraint] = []
 
     for p in all_possibilities:
         p = list(p)
@@ -575,10 +657,12 @@ def generate_calc_product(constraint, counter):
 
 
 @register_transformation_rule(CanReshape)
-def generate_reshape(constraint, counter):
+def generate_reshape(constraint: Constraint, counter: int) -> tuple[Constraint, int]:
     """
     Transform reshape constraints
     """
+    if not isinstance(constraint, CanReshape):
+        raise TypeError(type(constraint))
     d, counter = gen_tensor_dims(4, counter)
 
     d1 = d[0]
@@ -710,10 +794,14 @@ def generate_reshape(constraint, counter):
 
 
 @register_transformation_rule(ApplyBroadcasting)
-def generate_broadcasting(constraint, counter):
+def generate_broadcasting(
+    constraint: Constraint, counter: int
+) -> tuple[Constraint, int]:
     """
     Transform broadcasting constraints
     """
+    if not isinstance(constraint, ApplyBroadcasting):
+        raise TypeError(type(constraint))
     e11, e12 = constraint.res1, constraint.res2
     e1, e2 = constraint.input1, constraint.input2
 
@@ -784,7 +872,9 @@ def generate_broadcasting(constraint, counter):
     )
 
 
-def transform_constraint(constraint: Constraint, counter: int):
+def transform_constraint(
+    constraint: Constraint, counter: int
+) -> tuple[Constraint, int]:
     """
     Transforms a constraint into a simpler constraint.
     Ex: precision and consistency are transformed to equality
@@ -802,7 +892,9 @@ def transform_constraint(constraint: Constraint, counter: int):
         return constraint, counter
 
 
-def calc_last_two_dims(constraint, d: list[DVar]):
+def calc_last_two_dims(
+    constraint: CalcConv | CalcMaxPool, d: list[DVar]
+) -> tuple[Constraint, Constraint]:
     """
     Generates constraints for the last two dimensions of a convolution or a maxpool output
     Args:
@@ -874,7 +966,9 @@ def calc_last_two_dims(constraint, d: list[DVar]):
     return c4, c5
 
 
-def generate_all_int_dyn_dim_possibilities(my_list: list[DVar]):
+def generate_all_int_dyn_dim_possibilities(
+    my_list: list[DVar],
+) -> list[tuple[BinConstraintD, ...]]:
     """
     Generate all possibilities of being equal or not equal to dyn for my_list
     Args:
@@ -896,7 +990,9 @@ def generate_all_int_dyn_dim_possibilities(my_list: list[DVar]):
     return all_possibilities
 
 
-def is_target_div_by_dim(target: list[int], dim: list[DVar]):
+def is_target_div_by_dim(
+    target: Sequence[DVar | int | _DynType], dim: DVar | Prod
+) -> BinConstraintD:
     """
     Generate constraints to check if the target dimensions are divisible by the input dimensions
     Args:
@@ -909,7 +1005,9 @@ def is_target_div_by_dim(target: list[int], dim: list[DVar]):
     return BinConstraintD(BinConstraintD(Prod(target), dim, op_mod), 0, op_eq)
 
 
-def is_dim_div_by_target(target: list[int], dim: list[DVar]):
+def is_dim_div_by_target(
+    target: Sequence[DVar | int | _DynType], dim: DVar | Prod
+) -> BinConstraintD:
     """
     Generate constraints to check if the input dimensions is divisible by the target dimensions
     Args:
@@ -922,7 +1020,9 @@ def is_dim_div_by_target(target: list[int], dim: list[DVar]):
     return BinConstraintD(BinConstraintD(dim, Prod(target), op_mod), 0, op_eq)
 
 
-def gen_all_reshape_possibilities(list_of_dims, target):
+def gen_all_reshape_possibilities(
+    list_of_dims: list[DVar], target: Sequence[DVar | int | _DynType]
+) -> Constraint:
     """
     Consider all possibilities what the input dimensions could be (number or dynamic)
     Then generate the appropriate constraints using multiplication or mod depending on the possibility
@@ -942,7 +1042,7 @@ def gen_all_reshape_possibilities(list_of_dims, target):
     all_constraints = []
 
     for p in all_possibilities:
-        to_multiply = []
+        to_multiply: list[DVar] = []
 
         p = list(p)
 
@@ -950,7 +1050,7 @@ def gen_all_reshape_possibilities(list_of_dims, target):
             if not isinstance(constraint, BinConstraintD):
                 raise AssertionError(f"Expected BinConstraintD, got {type(constraint)}")
             if constraint.op == op_neq:
-                to_multiply.append(constraint.lhs)
+                to_multiply.append(constraint.lhs)  # type: ignore[arg-type]
 
         if not to_multiply:
             all_constraints.append(Conj(p))
@@ -967,7 +1067,14 @@ def gen_all_reshape_possibilities(list_of_dims, target):
     return Disj(all_constraints)
 
 
-def broadcast_dim(tensor_input1, tensor_input2, res1, res2, index, padding=False):
+def broadcast_dim(
+    tensor_input1: Sequence[DVar | None],
+    tensor_input2: Sequence[DVar],
+    res1: Sequence[DVar],
+    res2: Sequence[DVar],
+    index: int,
+    padding: bool = False,
+) -> Constraint:
     """
     Apply broadcasting to the 'index' dimension of tensor_input1.
     Args:
@@ -989,7 +1096,7 @@ def broadcast_dim(tensor_input1, tensor_input2, res1, res2, index, padding=False
         # then the inputs are the same length so they all have dimensions at "index"
         return Conj(
             [
-                BinConstraintD(tensor_input1[index], 1, op_eq),
+                BinConstraintD(tensor_input1[index], 1, op_eq),  # type: ignore[arg-type]
                 BinConstraintD(res1[index], res2[index], op_eq),
                 BinConstraintD(res2[index], tensor_input2[index], op_eq),
             ]
@@ -1014,7 +1121,7 @@ def apply_padding(
     d11: list[DVar],
     d12: list[DVar],
     counter: int,
-):
+) -> tuple[Constraint, int]:
     """
     We are considering the possibility where one input has less dimensions than
     another input, so we apply padding to the broadcasted results
@@ -1080,7 +1187,7 @@ def apply_padding(
 
 def no_broadcast_dim_with_index(
     d1: list[DVar], d2: list[DVar], d3: list[DVar], d4: list[DVar], i: int
-):
+) -> Constraint:
     """
     Args:
         d1: input 1
@@ -1115,7 +1222,9 @@ def no_broadcast_dim_with_index(
     )
 
 
-def gen_lists_of_dims(num_tensors: int, dim_size: int, counter: int):
+def gen_lists_of_dims(
+    num_tensors: int, dim_size: int, counter: int
+) -> tuple[list[list[DVar]], int]:
     """
     Generate lists of DVar to represent tensor dimensions
     Args:
@@ -1144,7 +1253,7 @@ def create_equality_constraints_for_broadcasting(
     d2: list[DVar],
     d11: list[DVar],
     d12: list[DVar],
-):
+) -> list[BinConstraintT]:
     """
     Create equality constraints for when no broadcasting occurs
     Args:
@@ -1168,7 +1277,9 @@ def create_equality_constraints_for_broadcasting(
     return [e1_tensor, e11_tensor, e2_tensor, e12_tensor]
 
 
-def gen_consistency_constraints(constraint: Constraint, counter: int):
+def gen_consistency_constraints(
+    constraint: BinConstraintT, counter: int
+) -> tuple[list[Constraint], int]:
     """
     Args:
         constraint: Consistency constraint on tensors
@@ -1178,7 +1289,7 @@ def gen_consistency_constraints(constraint: Constraint, counter: int):
 
     """
 
-    all_constraints = []
+    all_constraints: list[Constraint] = []
 
     for i in range(1, MAX_TENSOR_RANK + 1):
         new_dims_rhs_1, counter = gen_tensor_dims(i, counter)
@@ -1203,7 +1314,9 @@ def gen_consistency_constraints(constraint: Constraint, counter: int):
     return all_constraints, counter
 
 
-def gen_greatest_upper_bound(constraint: TGreatestUpperBound, counter: int):
+def gen_greatest_upper_bound(
+    constraint: TGreatestUpperBound, counter: int
+) -> tuple[list[Constraint], int]:
     """
     Args:
         constraint: Greatest upper bound on tensors
@@ -1213,10 +1326,10 @@ def gen_greatest_upper_bound(constraint: TGreatestUpperBound, counter: int):
 
     """
 
-    all_constraints = []
+    all_constraints: list[Constraint] = []
 
     for i in range(1, MAX_TENSOR_RANK + 1):
-        c = []
+        c: list[Constraint] = []
         dims1, counter = gen_tensor_dims(i, counter)
         c1tensor = TensorType(dims1)
 
@@ -1249,7 +1362,7 @@ def gen_greatest_upper_bound(constraint: TGreatestUpperBound, counter: int):
 
 def generate_all_broadcasting_possibilities_no_padding(
     d1: list[DVar], d2: list[DVar], d11: list[DVar], d12: list[DVar]
-):
+) -> Constraint:
     """
     Generate broadcasting constraints assuming no padding. Broadcasting can happen at any dimension.
     We look at all combinations for all dimensions in d1 and d2
@@ -1279,7 +1392,7 @@ def generate_all_broadcasting_possibilities_no_padding(
 
 def gen_broadcasting_constraints(
     e1: TVar, e2: TVar, e11: TVar, e12: TVar, i: int, counter: int
-):
+) -> tuple[Constraint, Constraint, Constraint, list[BinConstraintD], int]:
     """
     Simulates broadcasting on e1 and e2 and returns the results
     respectively in e11 and e12. Because of gradual types,
diff --git a/torch/fx/experimental/migrate_gradual_types/transform_to_z3.py b/torch/fx/experimental/migrate_gradual_types/transform_to_z3.py
index 6feccb3483fd7..73414d628bb77 100644
--- a/torch/fx/experimental/migrate_gradual_types/transform_to_z3.py
+++ b/torch/fx/experimental/migrate_gradual_types/transform_to_z3.py
@@ -1,9 +1,29 @@
-# mypy: allow-untyped-defs
+from typing import Any, TypeAlias
+
+import torch
+
+
+__all__ = [
+    "evaluate_conditional_with_constraints",
+    "iterate_till_fixed_point",
+    "transform_algebraic_expression",
+    "transform_all_constraints",
+    "transform_all_constraints_trace_time",
+    "transform_dimension",
+    "transform_to_z3",
+    "transform_var",
+]
+
+
+# z3 is an optional dependency with no type stubs, so we use aliases for its types.
+_Z3Expr: TypeAlias = Any
+_Z3Result: TypeAlias = Any
 from torch.fx.experimental.migrate_gradual_types.constraint import (
     BinConstraintD,
     BinConstraintT,
     BVar,
     Conj,
+    Constraint,
     Disj,
     DVar,
     F,
@@ -32,7 +52,9 @@
     op_neq,
     op_sub,
 )
-from torch.fx.tensor_type import Dyn, TensorType
+from torch.fx.graph import Graph
+from torch.fx.node import Node
+from torch.fx.tensor_type import _DynType, Dyn, TensorType
 
 
 try:
@@ -46,7 +68,9 @@
 
     HAS_Z3 = True
 
-    def transform_to_z3(constraint, counter, dimension_dict):
+    def transform_to_z3(
+        constraint: Constraint, counter: int, dimension_dict: dict[int, int]
+    ) -> tuple[_Z3Expr, int]:
         if isinstance(constraint, Conj):
             conjuncts = []
             for c in constraint.conjucts:
@@ -69,8 +93,16 @@ def transform_to_z3(constraint, counter, dimension_dict):
 
         elif isinstance(constraint, BinConstraintT):
             if constraint.op == op_eq:
-                lhs, counter = transform_var(constraint.lhs, counter, dimension_dict)
-                rhs, counter = transform_var(constraint.rhs, counter, dimension_dict)
+                lhs, counter = transform_var(
+                    constraint.lhs,  # pyrefly: ignore[bad-argument-type]
+                    counter,
+                    dimension_dict,
+                )
+                rhs, counter = transform_var(
+                    constraint.rhs,  # pyrefly: ignore[bad-argument-type]
+                    counter,
+                    dimension_dict,
+                )
                 return (lhs == rhs), counter
 
             else:
@@ -80,7 +112,9 @@ def transform_to_z3(constraint, counter, dimension_dict):
             if constraint.op == op_eq:
                 if isinstance(constraint.lhs, BVar) and is_bool_expr(constraint.rhs):
                     transformed_rhs, counter = transform_to_z3(
-                        constraint.rhs, counter, dimension_dict
+                        constraint.rhs,  # pyrefly: ignore[bad-argument-type]
+                        counter,
+                        dimension_dict,
                     )
                     transformed_lhs = z3.Bool(constraint.lhs.c)
                     return transformed_lhs == transformed_rhs, counter
@@ -88,10 +122,14 @@ def transform_to_z3(constraint, counter, dimension_dict):
                 elif is_dim(constraint.lhs) and is_dim(constraint.rhs):
                     # with dimension transformations we consider the encoding
                     lhs, counter = transform_dimension(
-                        constraint.lhs, counter, dimension_dict
+                        constraint.lhs,  # pyrefly: ignore[bad-argument-type]
+                        counter,
+                        dimension_dict,
                     )
                     rhs, counter = transform_dimension(
-                        constraint.rhs, counter, dimension_dict
+                        constraint.rhs,  # pyrefly: ignore[bad-argument-type]
+                        counter,
+                        dimension_dict,
                     )
                     return lhs == rhs, counter
 
@@ -99,10 +137,14 @@ def transform_to_z3(constraint, counter, dimension_dict):
                     # then we have an algebraic expression which means that we disregard the
                     # first element of the encoding
                     lhs, counter = transform_algebraic_expression(
-                        constraint.lhs, counter, dimension_dict
+                        constraint.lhs,  # pyrefly: ignore[bad-argument-type]
+                        counter,
+                        dimension_dict,
                     )
                     rhs, counter = transform_algebraic_expression(
-                        constraint.rhs, counter, dimension_dict
+                        constraint.rhs,  # pyrefly: ignore[bad-argument-type]
+                        counter,
+                        dimension_dict,
                     )
                     return lhs == rhs, counter
 
@@ -113,15 +155,19 @@ def transform_to_z3(constraint, counter, dimension_dict):
                 if not is_dim(constraint.rhs):
                     raise AssertionError("Expected rhs to be a dimension")
                 lhs, counter = transform_dimension(
-                    constraint.lhs, counter, dimension_dict
+                    constraint.lhs,  # pyrefly: ignore[bad-argument-type]
+                    counter,
+                    dimension_dict,
                 )
                 rhs, counter = transform_dimension(
-                    constraint.rhs, counter, dimension_dict
+                    constraint.rhs,  # pyrefly: ignore[bad-argument-type]
+                    counter,
+                    dimension_dict,
                 )
                 if constraint.rhs == Dyn or constraint.lhs == Dyn:
                     if constraint.rhs == Dyn:
                         return lhs.arg(0) == 1, counter
-                    elif constraint.lhs == Dyn:
+                    else:
                         return rhs.arg(0) == 1, counter
 
                 # if one of the instances is a number
@@ -137,7 +183,7 @@ def transform_to_z3(constraint, counter, dimension_dict):
                             counter,
                         )
 
-                    elif isinstance(constraint.rhs, int):
+                    else:
                         return (
                             z3.Or(
                                 [
@@ -173,10 +219,14 @@ def transform_to_z3(constraint, counter, dimension_dict):
                 if not (is_dim(constraint.lhs) and is_dim(constraint.rhs)):
                     raise AssertionError("Expected both lhs and rhs to be dimensions")
                 lhs, counter = transform_algebraic_expression(
-                    constraint.lhs, counter, dimension_dict
+                    constraint.lhs,  # pyrefly: ignore[bad-argument-type]
+                    counter,
+                    dimension_dict,
                 )
                 rhs, counter = transform_algebraic_expression(
-                    constraint.rhs, counter, dimension_dict
+                    constraint.rhs,  # pyrefly: ignore[bad-argument-type]
+                    counter,
+                    dimension_dict,
                 )
                 return lhs <= rhs, counter
 
@@ -184,10 +234,14 @@ def transform_to_z3(constraint, counter, dimension_dict):
                 if not (is_dim(constraint.lhs) and is_dim(constraint.rhs)):
                     raise AssertionError("Expected both lhs and rhs to be dimensions")
                 lhs, counter = transform_algebraic_expression(
-                    constraint.lhs, counter, dimension_dict
+                    constraint.lhs,  # pyrefly: ignore[bad-argument-type]
+                    counter,
+                    dimension_dict,
                 )
                 rhs, counter = transform_algebraic_expression(
-                    constraint.rhs, counter, dimension_dict
+                    constraint.rhs,  # pyrefly: ignore[bad-argument-type]
+                    counter,
+                    dimension_dict,
                 )
                 return lhs > rhs, counter
 
@@ -195,10 +249,14 @@ def transform_to_z3(constraint, counter, dimension_dict):
                 if not (is_dim(constraint.lhs) and is_dim(constraint.rhs)):
                     raise AssertionError("Expected both lhs and rhs to be dimensions")
                 lhs, counter = transform_algebraic_expression(
-                    constraint.lhs, counter, dimension_dict
+                    constraint.lhs,  # pyrefly: ignore[bad-argument-type]
+                    counter,
+                    dimension_dict,
                 )
                 rhs, counter = transform_algebraic_expression(
-                    constraint.rhs, counter, dimension_dict
+                    constraint.rhs,  # pyrefly: ignore[bad-argument-type]
+                    counter,
+                    dimension_dict,
                 )
                 return lhs < rhs, counter
 
@@ -208,7 +266,11 @@ def transform_to_z3(constraint, counter, dimension_dict):
         else:
             raise NotImplementedError("Operation not yet implemented")
 
-    def transform_var(tensor, counter, dimension_dict):
+    def transform_var(
+        tensor: TVar | TensorType | _DynType,
+        counter: int,
+        dimension_dict: dict[int, int],
+    ) -> tuple[_Z3Expr, int]:
         """
         Transforms tensor variables to a format understood by z3
         Args:
@@ -217,7 +279,7 @@ def transform_var(tensor, counter, dimension_dict):
 
         """
         if isinstance(tensor, TensorType):
-            res = []
+            res: list[_Z3Expr] = []
             for t in tensor.__args__:
                 transformed, counter = transform_dimension(t, counter, dimension_dict)
                 res.append(transformed)
@@ -232,6 +294,10 @@ def transform_var(tensor, counter, dimension_dict):
                 return tensor_type.tensor3(res[0], res[1], res[2]), counter
             elif len(tensor.__args__) == 4:
                 return tensor_type.tensor4(res[0], res[1], res[2], res[3]), counter
+            else:
+                raise AssertionError(
+                    f"Unexpected tensor args length: {len(tensor.__args__)}"
+                )
 
         elif tensor == Dyn:
             return z3_dyn, counter
@@ -239,7 +305,12 @@ def transform_var(tensor, counter, dimension_dict):
         elif isinstance(tensor, TVar):
             return z3.Const(tensor.tvar, tensor_type), counter
 
-    def transform_dimension(dimension, counter, dimension_dict):
+        else:
+            raise NotImplementedError(f"Unsupported tensor type: {type(tensor)}")
+
+    def transform_dimension(
+        dimension: DVar | int | _DynType, counter: int, dimension_dict: dict[int, int]
+    ) -> tuple[_Z3Expr, int]:
         """
         Takes a dimension variable or a number and transforms it to a tuple
         according to our scheme
@@ -266,7 +337,14 @@ def transform_dimension(dimension, counter, dimension_dict):
                 dimension_dict[dimension.c] = counter
                 return D(z3.Int(counter), z3.Int(dimension.c)), counter
 
-    def transform_algebraic_expression(expr, counter, dimension_dict):
+        else:
+            raise NotImplementedError(f"Unsupported dimension type: {type(dimension)}")
+
+    def transform_algebraic_expression(
+        expr: DVar | int | _DynType | Prod | BinConstraintD,
+        counter: int,
+        dimension_dict: dict[int, int],
+    ) -> tuple[_Z3Expr, int]:
         """
         Transforms an algebraic expression to z3 format
         Args:
@@ -280,7 +358,11 @@ def transform_algebraic_expression(expr, counter, dimension_dict):
             raise AssertionError("Expected algebraic expression or dimension")
 
         if is_dim(expr):
-            transformed, counter = transform_dimension(expr, counter, dimension_dict)
+            transformed, counter = transform_dimension(
+                expr,  # pyrefly: ignore[bad-argument-type]
+                counter,
+                dimension_dict,
+            )
             return transformed.arg(1), counter
 
         elif isinstance(expr, Prod):
@@ -294,13 +376,17 @@ def transform_algebraic_expression(expr, counter, dimension_dict):
 
         elif is_algebraic_expression(expr):
             lhs, counter = transform_algebraic_expression(
-                expr.lhs, counter, dimension_dict
+                expr.lhs,  # pyrefly: ignore[missing-attribute]
+                counter,
+                dimension_dict,
             )
             rhs, counter = transform_algebraic_expression(
-                expr.rhs, counter, dimension_dict
+                expr.rhs,  # pyrefly: ignore[missing-attribute]
+                counter,
+                dimension_dict,
             )
 
-            if expr.op == op_sub:
+            if expr.op == op_sub:  # pyrefly: ignore[missing-attribute]
                 c = lhs - rhs
 
             elif expr.op == op_add:
@@ -323,31 +409,25 @@ def transform_algebraic_expression(expr, counter, dimension_dict):
         else:
             raise RuntimeError
 
-    def transform_all_constraints(traced, counter=0):
+    def transform_all_constraints(traced: torch.nn.Module, counter: int = 0) -> _Z3Expr:
         """
         Given a trace, generates constraints and transforms them to z3 format
 
         """
-        dimension_dict = {}  # type: ignore[var-annotated]
+        dimension_dict: dict[int, int] = {}
 
         generator = ConstraintGenerator(traced)
         new_constraints, counter = generator.generate_constraints(counter)
 
-        # print(new_constraints.conjucts[0])
-        # print(*new_constraints.conjucts, sep='\n')
-
-        # transform precision, matching, consistency till obtaining a fixed point
         new_constraints, counter = iterate_till_fixed_point(new_constraints, counter)
-        # print(new_constraints)
-        # print(new_constraints.conjucts)
-        # new_constraints.conjucts = new_constraints.conjucts[:-1]
-        # print(*new_constraints.conjucts, sep='\n')
 
         transformed, counter = transform_to_z3(new_constraints, counter, dimension_dict)
         # print(transformed)
         return transformed
 
-    def iterate_till_fixed_point(constraints, counter):
+    def iterate_till_fixed_point(
+        constraints: Constraint, counter: int
+    ) -> tuple[Constraint, int]:
         """
         Transform constraints till reaching a fixed point
         """
@@ -357,7 +437,9 @@ def iterate_till_fixed_point(constraints, counter):
             constraints, counter = transform_constraint(constraints, counter)
         return constraints, counter
 
-    def transform_all_constraints_trace_time(tracer_root, graph, node, counter=0):
+    def transform_all_constraints_trace_time(
+        tracer_root: torch.nn.Module, graph: Graph, node: Node, counter: int = 0
+    ) -> tuple[_Z3Expr, _Z3Expr]:
         """
         Takes a node and a graph and generates two sets of constraints.
         One set constraints the node's constraints and another set
@@ -373,7 +455,7 @@ def transform_all_constraints_trace_time(tracer_root, graph, node, counter=0):
         its negation.
 
         """
-        dimension_dict = {}  # type: ignore[var-annotated]
+        dimension_dict: dict[int, int] = {}
 
         generator = ConstraintGenerator(tracer_root, graph)
         new_constraints, counter = generator.generate_constraints(counter)
@@ -394,10 +476,14 @@ def transform_all_constraints_trace_time(tracer_root, graph, node, counter=0):
         # we make sure the constraint is of the form:
         # c = b where b is a boolean expression
         # and we consider b (constraint.rhs) for transformation
+        if not isinstance(condition_constraint, BinConstraintD):
+            raise TypeError(type(condition_constraint))
         if not isinstance(condition_constraint.lhs, BVar):
             raise AssertionError(f"Expected BVar, got {type(condition_constraint.lhs)}")
         if not is_bool_expr(condition_constraint.rhs):
             raise AssertionError("Expected bool expression for rhs")
+        if not isinstance(condition_constraint.rhs, Constraint):
+            raise TypeError(type(condition_constraint.rhs))
         condition_constraint_rhs = condition_constraint.rhs
 
         # transform the condition constraint
@@ -420,8 +506,12 @@ def transform_all_constraints_trace_time(tracer_root, graph, node, counter=0):
         )
 
     def evaluate_conditional_with_constraints(
-        tracer_root, graph, node, counter=0, user_constraints=None
-    ):
+        tracer_root: torch.nn.Module,
+        graph: Graph,
+        node: Node,
+        counter: int = 0,
+        user_constraints: _Z3Expr | None = None,
+    ) -> tuple[_Z3Result, _Z3Result]:
         """
         Given an IR and a node representing a conditional, evaluate the conditional
         and its negation
diff --git a/torch/fx/experimental/normalize.py b/torch/fx/experimental/normalize.py
index 31634206c348e..e490919b3eca1 100644
--- a/torch/fx/experimental/normalize.py
+++ b/torch/fx/experimental/normalize.py
@@ -1,7 +1,7 @@
 # mypy: allow-untyped-defs
 import operator
 from collections.abc import Callable
-from typing import Any, Optional
+from typing import Any
 
 import torch
 import torch.fx
@@ -70,8 +70,8 @@ def call_function(
         target: Target,
         args: tuple[Argument, ...],
         kwargs: dict[str, Any],
-        arg_types: Optional[tuple[Any, ...]] = None,
-        kwarg_types: Optional[dict[str, Any]] = None,
+        arg_types: tuple[Any, ...] | None = None,
+        kwarg_types: dict[str, Any] | None = None,
     ):
         if not callable(target):
             raise AssertionError(f"Expected callable target, got {type(target)}")
diff --git a/torch/fx/experimental/optimization.py b/torch/fx/experimental/optimization.py
index de4d5f1ceaf1a..105528da26639 100644
--- a/torch/fx/experimental/optimization.py
+++ b/torch/fx/experimental/optimization.py
@@ -6,7 +6,7 @@
 from collections import defaultdict
 from collections.abc import Iterable
 from enum import Enum
-from typing import Any, cast, Optional
+from typing import Any, cast
 
 import torch
 import torch.fx as fx
@@ -291,7 +291,7 @@ def use_mkl_length(graph: MklSubgraph) -> bool:
 
 class UnionFind:
     def __init__(self, n):
-        self.parent: list[Optional[int]] = [None] * n
+        self.parent: list[int | None] = [None] * n
         self.size: list[int] = [0] * n
 
     def make_set(self, v: int):
@@ -319,7 +319,7 @@ def join(self, a: int, b: int):
 
 def optimize_for_inference(
     model: torch.nn.Module,
-    pass_config: Optional[dict[str, Any]] = None,
+    pass_config: dict[str, Any] | None = None,
     tracer: type[fx.Tracer] = fx.Tracer,
 ) -> torch.nn.Module:
     """
diff --git a/torch/fx/experimental/proxy_tensor.py b/torch/fx/experimental/proxy_tensor.py
index 4a04cafc8e13d..28787fd3057ba 100644
--- a/torch/fx/experimental/proxy_tensor.py
+++ b/torch/fx/experimental/proxy_tensor.py
@@ -1,4 +1,3 @@
-# mypy: allow-untyped-decorators
 # Copyright (c) Facebook, Inc. and its affiliates.
 # All rights reserved.
 #
@@ -7,6 +6,7 @@
 
 from __future__ import annotations
 
+import contextvars
 import functools
 import inspect
 import logging
@@ -14,7 +14,6 @@
 import threading
 import typing
 import typing_extensions
-import weakref
 from collections import defaultdict, OrderedDict
 from collections.abc import Callable, Generator, Mapping, Sequence
 from contextlib import _GeneratorContextManager, contextmanager, ExitStack, nullcontext
@@ -22,7 +21,7 @@
 from typing import (
     Any,
     Concatenate,
-    Optional,
+    Literal,
     overload,
     Protocol,
     TYPE_CHECKING,
@@ -41,9 +40,16 @@
 from torch import SymBool, SymInt, Tensor
 from torch._dispatch.python import enable_python_dispatcher
 from torch._library.fake_class_registry import FakeScriptObject
-from torch._library.opaque_object import is_opaque_value, OpaqueType
+from torch._library.opaque_object import (
+    get_reconstruct_fn,
+    is_opaque_reference_type,
+    is_opaque_value,
+    is_opaque_value_type,
+    OpaqueBase,
+    should_hoist,
+)
 from torch._logging import trace_structured
-from torch._ops import HigherOrderOperator
+from torch._ops import HigherOrderOperator, OpOverload
 from torch._subclasses.fake_impls import fast_detach
 from torch._subclasses.fake_tensor import (
     FakeTensor,
@@ -86,8 +92,6 @@
     import sympy
 
     from torch._higher_order_ops.utils import FunctionalizeCtxWrapper
-    from torch._ops import OpOverload
-    from torch.fx._symbolic_trace import PHBase
     from torch.types import BoolLikeType, FloatLikeType, IntLikeType
 
 __all__ = [
@@ -105,9 +109,10 @@
 ]
 
 _ProxyTracer = Union["PythonKeyTracer", "_GraphAppendingTracerEx"]
+_TracingMode: TypeAlias = Literal["real", "fake", "symbolic"]
 
 _AnyScriptObject = (torch.ScriptObject, FakeScriptObject)
-_AnyScriptObjectType = Union[torch.ScriptObject, FakeScriptObject]
+_AnyScriptObjectType = torch.ScriptObject | FakeScriptObject
 
 aten = torch.ops.aten
 prim = torch.ops.prim
@@ -115,7 +120,9 @@
 log = logging.getLogger(__name__)
 not_implemented_log = torch._logging.getArtifactLogger(__name__, "not_implemented")
 
-CURRENT_DECOMPOSITION_TABLE: Mapping[OpOverload, Callable] = {}
+CURRENT_DECOMPOSITION_TABLE: contextvars.ContextVar[
+    Mapping[OpOverload, Callable[..., Any]]
+] = contextvars.ContextVar("CURRENT_DECOMPOSITION_TABLE")
 
 CONSTANT_NUMEL_LIMIT = 1
 
@@ -125,7 +132,6 @@
 R = TypeVar("R")
 _Ts = TypeVarTuple("_Ts")
 
-null_ctx_type = type(nullcontext)
 # We currently convert all SymInt to proxies before we use them.
 # This could plausibly be handled at the Dynamo level.
 pytree.register_pytree_node(
@@ -156,15 +162,14 @@ def fake_signature(fn: Callable[_P, R], nargs: int) -> Callable[_P, R]:
 
 @contextmanager
 def decompose(
-    decomposition_table: Optional[Mapping[OpOverload, Callable]],
-) -> Generator[Mapping[OpOverload, Callable], None, None]:
-    global CURRENT_DECOMPOSITION_TABLE
-    old_decomposition_table = CURRENT_DECOMPOSITION_TABLE
-    CURRENT_DECOMPOSITION_TABLE = decomposition_table or {}
+    decomposition_table: Mapping[OpOverload, Callable[..., Any]] | None,
+) -> Generator[Mapping[OpOverload, Callable[..., Any]], None, None]:
+    table = decomposition_table or {}
+    token = CURRENT_DECOMPOSITION_TABLE.set(table)
     try:
-        yield CURRENT_DECOMPOSITION_TABLE
+        yield table
     finally:
-        CURRENT_DECOMPOSITION_TABLE = old_decomposition_table
+        CURRENT_DECOMPOSITION_TABLE.reset(token)
 
 
 # ensure we cannot collide with other properties
@@ -258,7 +263,7 @@ def set_proxy_slot(obj: Tensor, tracer: _ProxyTracer, proxy: _ProxyTensor) -> No
 
 @overload
 def set_proxy_slot(
-    obj: _AnyScriptObjectType | OpaqueType, tracer: _ProxyTracer, proxy: Proxy
+    obj: _AnyScriptObjectType | OpaqueBase, tracer: _ProxyTracer, proxy: Proxy
 ) -> None: ...
 
 
@@ -269,7 +274,7 @@ def set_proxy_slot(
 
 
 def set_proxy_slot(
-    obj: Union[PySymType, _AnyScriptObjectType, Tensor, OpaqueType],
+    obj: PySymType | _AnyScriptObjectType | Tensor | OpaqueBase,
     tracer: _ProxyTracer,
     proxy: object,
 ) -> None:
@@ -286,10 +291,39 @@ def set_proxy_slot(
         ):
             tracer.tensor_tracker[obj] = proxy
     elif isinstance(obj, (_AnyScriptObject)) or is_opaque_value(obj):
-        # We DO want to clobber proxies, with a similar rationale as for tensors.
         if not isinstance(proxy, Proxy):
             raise AssertionError(f"Expected Proxy, got {type(proxy)}")
-        tracer.script_object_tracker[obj] = proxy
+        # ScriptObject (actual C++ torchbind) uses _WeakHashRef-keyed tracker
+        # because the same C++ IValue can produce different Python wrappers.
+        # FakeScriptObject/OpaqueBase uses WeakIdRef-keyed tracker because
+        # value-equal objects (e.g. primal vs tangent) must be tracked separately.
+        if isinstance(obj, torch.ScriptObject):
+            tracer.script_object_tracker[obj] = proxy
+        else:
+            # NB: Never clobber a pre-existing proxy for the same
+            # underlying real object.  Multiple FakeScriptObject wrappers
+            # can share the same real_obj (e.g. primal vs tangent
+            # placeholders during joint graph tracing).  We always keep the
+            # first proxy registered, with the same rationale as the
+            # symnode_tracker first-one-wins policy below: primals are
+            # registered first, so this avoids spurious tangent dependencies
+            # in forward outputs (which would break the partitioner).
+            real_obj = None
+            if isinstance(obj, FakeScriptObject):
+                try:
+                    real_obj = object.__getattribute__(obj, "real_obj")
+                except AttributeError:
+                    pass
+
+            if real_obj is not None:
+                existing = tracer._opaque_real_obj_proxy.get(id(real_obj))
+                if existing is not None:
+                    tracer.opaque_tracker[obj] = existing
+                else:
+                    tracer.opaque_tracker[obj] = proxy
+                    tracer._opaque_real_obj_proxy[id(real_obj)] = proxy
+            else:
+                tracer.opaque_tracker[obj] = proxy
     else:
         # NB: Never clobber pre-existing proxy.  Although the proxies
         # are in principle equivalent, when we do graph partitioning
@@ -331,7 +365,7 @@ def has_proxy_slot(obj: Tensor, tracer: _ProxyTracer) -> bool:
 
 
 _PySymProxyType: TypeAlias = Thunk[Proxy]
-_OpaqueObjectProxyType: TypeAlias = Thunk[Proxy]
+_OpaqueObjectProxyType: TypeAlias = Proxy
 
 
 @overload
@@ -346,7 +380,7 @@ def get_proxy_slot(
     obj: Tensor,
     tracer: _ProxyTracer,
     default: U,
-) -> Union[_ProxyTensor, U]: ...
+) -> _ProxyTensor | U: ...
 
 
 @overload
@@ -355,7 +389,7 @@ def get_proxy_slot(
     tracer: _ProxyTracer,
     default: U,
     transform: Callable[[_ProxyTensor], R],
-) -> Union[R, U]: ...
+) -> R | U: ...
 
 
 @overload
@@ -370,7 +404,7 @@ def get_proxy_slot(
     obj: _AnyScriptObjectType,
     tracer: _ProxyTracer,
     default: U,
-) -> Union[Proxy, U]: ...
+) -> Proxy | U: ...
 
 
 @overload
@@ -379,7 +413,7 @@ def get_proxy_slot(
     tracer: _ProxyTracer,
     default: U,
     transform: Callable[[Proxy], R],
-) -> Union[R, U]: ...
+) -> R | U: ...
 
 
 @overload
@@ -394,15 +428,15 @@ def get_proxy_slot(
     obj: PySymType,
     tracer: _ProxyTracer,
     default: T,
-) -> Union[T, _PySymProxyType]: ...
+) -> T | _PySymProxyType: ...
 
 
 @overload
 def get_proxy_slot(
-    obj: OpaqueType,
+    obj: OpaqueBase,
     tracer: _ProxyTracer,
     default: T,
-) -> Union[T, _OpaqueObjectProxyType]: ...
+) -> T | _OpaqueObjectProxyType: ...
 
 
 @overload
@@ -411,23 +445,26 @@ def get_proxy_slot(
     tracer: _ProxyTracer,
     default: U,
     transform: Callable[[_PySymProxyType], R],
-) -> Union[R, U]: ...
+) -> R | U: ...
 
 
 # the default argument is what to return if the slot is not set.
 # the transform argument is handy if you need to extract a subfield from
 # the successfully looked up result (but NOT the default.)
 def get_proxy_slot(
-    obj: Union[Tensor, _AnyScriptObjectType, PySymType, OpaqueType],
+    obj: Tensor | _AnyScriptObjectType | PySymType | OpaqueBase,
     tracer: _ProxyTracer,
     default: object = no_default,
-    transform: Callable = lambda x: x,
+    transform: Callable[..., Any] = lambda x: x,
 ) -> object:
     tracker: Any
     if isinstance(obj, Tensor):
         tracker = tracer.tensor_tracker
     elif isinstance(obj, _AnyScriptObject) or is_opaque_value(obj):
-        tracker = tracer.script_object_tracker
+        if isinstance(obj, torch.ScriptObject):
+            tracker = tracer.script_object_tracker
+        else:
+            tracker = tracer.opaque_tracker
     else:
         if not isinstance(obj, py_sym_types):
             raise AssertionError(f"Expected py_sym_types, got {type(obj)}")
@@ -447,9 +484,16 @@ def get_proxy_slot(
             else:
                 # Attempt to build it from first principles.
                 _build_proxy_for_sym_expr(tracer, obj.node.expr, obj)
-                # pyrefly: ignore [no-matching-overload]
+                # pyrefly: ignore [bad-argument-type, no-matching-overload]
                 value = tracker.get(obj)
 
+    if value is None and isinstance(obj, FakeScriptObject):
+        # A new FakeScriptObject wrapping the same real_obj may have been
+        # created (e.g. output flattening in unwrap_tensor_subclasses calls
+        # maybe_to_fake_obj which always mints a fresh wrapper).  Fall back
+        # to the real-object dedup map that set_proxy_slot maintains.
+        value = tracer._opaque_real_obj_proxy.get(id(obj.real_obj))
+
     if value is None:
         # We don't know this value - return the default.
         if isinstance(default, _NoDefault):
@@ -488,9 +532,11 @@ def _get_proxies(t: torch.Tensor) -> list[Proxy]:
 @functools.cache
 def _sympy_handlers() -> dict[type[sympy.Expr], Callable[..., Any]]:
     """
-    Returns a dict converting sympy functions to python operators
-    (i.e. `sympy.Mul` -> `operator.mul`)
+    Returns a dict mapping sympy types to Python callables
+    (e.g. ``sympy.Mul`` -> ``operator.mul``, ``sympy.Add`` -> ``torch.sym_sum``).
     """
+    import sympy
+
     import torch.utils._sympy.interp
 
     handlers = {}
@@ -498,6 +544,11 @@ def _sympy_handlers() -> dict[type[sympy.Expr], Callable[..., Any]]:
         op = getattr(operator, v, None)
         if op is not None:
             handlers[k] = op
+
+    # sympy.Add is n-ary (e.g. Add(a, b, c)) but operator.add is binary.
+    # torch.sym_sum handles n-ary integer addition and accepts both
+    # sym_sum([a, b, c]) and sym_sum(a, b, c).
+    handlers[sympy.Add] = torch.sym_sum
     return handlers
 
 
@@ -592,21 +643,19 @@ def _build_proxy_for_sym_expr(
         if (arg_value := _build_proxy_for_sym_expr(tracer, arg)) is None:
             return None
         args.append(arg_value)
-    args = tuple(args)
 
     func: OpOverload | None = _sympy_handlers().get(expr.func)  # type: ignore[assignment]
     if not func:
-        # Handler not found
         return None
 
     if out is None:
         out = func(*args)
     else:
-        _sym_register(tracer, func, args, out)
+        _sym_register(tracer, func, tuple(args), out)
     return out
 
 
-def snapshot_fake(val: Tensor, include_real: bool = False) -> Optional[Tensor]:
+def snapshot_fake(val: Tensor, include_real: bool = False) -> Tensor | None:
     # val.detach() will also eventually call fast_detach(),
     # but this saves us a full trip into __torch_dispatch__
     # (snapshot_fake is called a lot)
@@ -619,7 +668,7 @@ def snapshot_fake(val: Tensor, include_real: bool = False) -> Optional[Tensor]:
 _ExtractValType: TypeAlias = (
     None
     | PySymType
-    | OpaqueType
+    | OpaqueBase
     | _AnyScriptObjectType
     | BackwardState
     | list["_ExtractValType"]
@@ -637,7 +686,7 @@ def extract_val(val: _ExtractValType, include_real: bool = False) -> _ExtractVal
         return snapshot_fake(val, include_real=include_real)
     elif isinstance(val, py_sym_types):
         return val
-    elif isinstance(val, _AnyScriptObject):
+    elif isinstance(val, (_AnyScriptObject, OpaqueBase)):
         return val
     elif isinstance(val, BackwardState):
         return val
@@ -763,7 +812,7 @@ def thunkify(
 
 
 def track_tensor(
-    tensor: Tensor, proxy: Proxy, *, constant: Optional[Tensor], tracer: _ProxyTracer
+    tensor: Tensor, proxy: Proxy, *, constant: Tensor | None, tracer: _ProxyTracer
 ) -> None:
     def try_set_proxy_slot(
         outer_s: IntLikeType,
@@ -836,10 +885,10 @@ def try_set_proxy_slot(
     set_proxy_slot(tensor, tracer, _ProxyTensor(proxy, constant))
 
 
-_NestedProxys = Union[
+_NestedProxys = Union[  # noqa: UP007
     Proxy, Sequence["_NestedProxys"], Mapping[object, "_NestedProxys"]
 ]
-_NestedTensors = Union[
+_NestedTensors = Union[  # noqa: UP007
     Tensor, Sequence["_NestedTensors"], Mapping[object, "_NestedTensors"]
 ]
 
@@ -848,7 +897,7 @@ def track_tensor_tree(
     inner_res: T,
     proxy_res: _NestedProxys,
     *,
-    constant: Optional[_NestedTensors],
+    constant: _NestedTensors | None,
     tracer: _ProxyTracer,
 ) -> T:
     # NB: We call set_unbacked_bindings only on the *topmost* call to
@@ -867,7 +916,7 @@ def track_tensor_tree(
     _set_unbacked_bindings(inner_res, proxy_res)
 
     def wrap_with_proxy(
-        e: object, proxy: _NestedProxys, constant: Optional[_NestedTensors]
+        e: object, proxy: _NestedProxys, constant: _NestedTensors | None
     ) -> None:
         if isinstance(e, Tensor):
             if not isinstance(proxy, Proxy):
@@ -885,6 +934,15 @@ def wrap_with_proxy(
         elif isinstance(e, _AnyScriptObject) or is_opaque_value(e):
             if not isinstance(proxy, Proxy):
                 raise AssertionError(f"Expected Proxy, got {type(proxy)}")
+            # Non-hoisted opaque value types should be baked as constants
+            # in the graph, not tracked as proxy references. This matches
+            # dynamo's behavior where non-hoisted values are not graph inputs.
+            if (
+                is_opaque_value_type(type(e))  # pyrefly: ignore[bad-argument-type]
+                and not should_hoist(type(e))
+            ):
+                set_meta(proxy, e)
+                return
             set_proxy_slot(e, tracer, proxy)
             set_meta(proxy, e)
         elif isinstance(e, (tuple, list)):
@@ -893,8 +951,8 @@ def wrap_with_proxy(
                 set_meta(proxy, e)
 
             def get_constant(
-                c: Optional[_NestedTensors], idx: int
-            ) -> Optional[_NestedTensors]:
+                c: _NestedTensors | None, idx: int
+            ) -> _NestedTensors | None:
                 if c is None:
                     return None
                 else:
@@ -942,13 +1000,13 @@ def get_constant(
 @dataclass
 class _ProxyTensor:
     proxy: Proxy
-    constant: Optional[Tensor]
+    constant: Tensor | None
 
 
 def fetch_sym_proxy(
     tracer: _ProxyTracer,
-) -> Callable[[PySymType], Union[bool, int, float, Proxy]]:
-    def inner(e: PySymType) -> Union[int, bool, float, Proxy]:
+) -> Callable[[PySymType], bool | int | float | Proxy]:
+    def inner(e: PySymType) -> int | bool | float | Proxy:
         n = e.node
         if n.constant is not None:
             return n.constant
@@ -968,32 +1026,30 @@ def inner(e: PySymType) -> Union[int, bool, float, Proxy]:
 
 
 @overload
-def fetch_object_proxy(
-    tracer: _ProxyTracer, t: Tensor
-) -> Union[_ProxyTensor, Tensor]: ...
+def fetch_object_proxy(tracer: _ProxyTracer, t: Tensor) -> _ProxyTensor | Tensor: ...
 
 
 @overload
 def fetch_object_proxy(
     tracer: _ProxyTracer, t: _AnyScriptObjectType
-) -> Union[Proxy, _AnyScriptObjectType]: ...
+) -> Proxy | _AnyScriptObjectType: ...
 
 
 @overload
 def fetch_object_proxy(
     tracer: _ProxyTracer, t: PySymType
-) -> Union[_PySymProxyType, PySymType]: ...
+) -> _PySymProxyType | PySymType: ...
 
 
 @overload
 def fetch_object_proxy(
-    tracer: _ProxyTracer, t: OpaqueType
-) -> Union[_OpaqueObjectProxyType, PySymType]: ...
+    tracer: _ProxyTracer, t: OpaqueBase
+) -> _OpaqueObjectProxyType | PySymType: ...
 
 
 def fetch_object_proxy(
     tracer: _ProxyTracer,
-    t: Tensor | _AnyScriptObjectType | PySymType | OpaqueType,
+    t: Tensor | _AnyScriptObjectType | PySymType | OpaqueBase,
 ) -> object:
     return get_proxy_slot(t, tracer, t)
 
@@ -1034,7 +1090,7 @@ def _maybe_record_pointwise_barrier(
 
 
 def _fetch_proxies_and_all_constant_flag(
-    flat_args_kwargs: Union[list[object], tuple[object, ...]], tracer: _ProxyTracer
+    flat_args_kwargs: list[object] | tuple[object, ...], tracer: _ProxyTracer
 ) -> tuple[list[object], tuple[object, ...], bool]:
     """
     Given flat arguments, fetch the proxies and whether they are all constants.
@@ -1273,7 +1329,13 @@ def tensor_numel_in_limit(t: Tensor) -> bool:
     else:
         constant = None
 
-    track_tensor_tree(out, proxy_out, constant=constant, tracer=tracer)
+    track_tensor_tree(
+        out,
+        proxy_out,
+        # pyrefly: ignore[bad-argument-type]
+        constant=constant,
+        tracer=tracer,
+    )
     _maybe_record_pointwise_barrier(func, proxy_mode)
     return out
 
@@ -1296,7 +1358,7 @@ def __contains__(self, key: PySymType) -> bool:
         return key.node in self.sym_node_dict
 
     def get(
-        self, key: PySymType, default: Optional[_PySymProxyType] = None
+        self, key: PySymType, default: _PySymProxyType | None = None
     ) -> _PySymProxyType:
         # dict.get()'s annotation doesn't accept `None` when the value type
         # isn't Optional.
@@ -1315,29 +1377,42 @@ class _SympyExprTrackerValue:
     value: PySymType
 
 
+def _init_proxy_trackers(tracer: PythonKeyTracer | _GraphAppendingTracerEx) -> None:
+    """Initialize the tracker dictionaries shared by PythonKeyTracer and _GraphAppendingTracerEx."""
+    tracer.tensor_tracker = WeakTensorKeyDictionary()
+    tracer.symnode_tracker = _SymNodeDict()
+    tracer.script_object_tracker = WeakIdKeyDictionary(dict=None, ref_type=_WeakHashRef)
+    tracer.opaque_tracker = WeakIdKeyDictionary()
+    tracer._opaque_real_obj_proxy = {}
+    tracer.sympy_expr_tracker = {}
+    # Stores the torch function that was called during tracing
+    tracer.torch_fn_metadata = None
+    # Stores the counts for every torch function called. This is to help
+    # distinguish between different calls to the same torch function.
+    tracer.torch_fn_counts = {}
+    tracer.enable_thunkify = False
+
+
 class PythonKeyTracer(Tracer):
-    script_object_tracker: MutableMapping[_AnyScriptObjectType | OpaqueType, Proxy]
+    # ScriptObject uses _WeakHashRef because the same C++ IValue can produce
+    # different Python wrapper objects, so Python id() won't match.
+    script_object_tracker: MutableMapping[torch.ScriptObject, Proxy]
+    # FakeScriptObject/OpaqueBase uses WeakIdRef because distinct objects that
+    # are value-equal (e.g. primal vs tangent opaques) must be tracked separately.
+    opaque_tracker: MutableMapping[FakeScriptObject | OpaqueBase, Proxy]
+    # Maps id(real_obj) -> proxy for opaque FSOs, so that multiple FSO wrappers
+    # of the same real object (e.g. primal vs tangent) resolve to one proxy.
+    _opaque_real_obj_proxy: dict[int, Proxy]
     symnode_tracker: _SymNodeDict
     sympy_expr_tracker: dict[sympy.Symbol, _SympyExprTrackerValue]
     tensor_tracker: MutableMapping[Tensor, _ProxyTensor]
+    torch_fn_metadata: OpOverload | None
     torch_fn_counts: dict[OpOverload, int]
     enable_thunkify: bool = False
 
     def __init__(self) -> None:
         super().__init__(autowrap_modules=())  # type: ignore[arg-type]
-        self.tensor_tracker = WeakTensorKeyDictionary()
-        self.symnode_tracker = _SymNodeDict()
-        self.script_object_tracker = WeakIdKeyDictionary(
-            dict=None, ref_type=_WeakHashRef
-        )
-        self.sympy_expr_tracker = {}
-
-        # Stores the torch function that was called during tracing
-        self.torch_fn_metadata = None
-        # Stores the counts for every torch function called. This is to help
-        # distinguish between different calls to the same torch function.
-        self.torch_fn_counts = {}
-        self.enable_thunkify = False
+        _init_proxy_trackers(self)
 
     # In general, we don't want to make modules leaves. In principle, users of
     # this tracer might want to override this in order to turn a couple specific
@@ -1371,18 +1446,83 @@ def create_arg(self, a: object) -> fx.node.Node:
             if a.node.constant is None:
                 raise AssertionError("a.node.constant should not be None")
             return a.node.constant
+
+        # Try reconstructing untracked opaque reference types from existing
+        # graph inputs (e.g. derive a DeviceMesh submesh from its root mesh).
+        if isinstance(a, (FakeScriptObject, OpaqueBase)):
+            node = self._try_reconstruct_opaque(a)
+            if node is not None:
+                return node
+
         return super().create_arg(a)  # type: ignore[return-value]
 
+    def _try_reconstruct_opaque(
+        self, a: FakeScriptObject | OpaqueBase
+    ) -> fx.node.Node | None:
+        """Try to reconstruct an opaque object from existing graph inputs.
+
+        When make_fx encounters an untracked opaque reference type (e.g. a
+        DeviceMesh submesh captured by a backward closure), this method checks
+        if the type has a registered reconstruct_fn that can derive the object
+        from inputs already in the graph.  Returns an FX Node on success,
+        None on failure (falls back to get_attr constant).
+        """
+        real_obj: OpaqueBase = a.real_obj if isinstance(a, FakeScriptObject) else a
+
+        if not is_opaque_reference_type(type(real_obj)):
+            return None
+
+        reconstruct_fn = get_reconstruct_fn(type(real_obj))
+        if reconstruct_fn is None:
+            return None
+
+        def get_tracked_proxy(obj: OpaqueBase) -> Proxy | None:
+            proxy = self._opaque_real_obj_proxy.get(id(obj))
+            if proxy is not None:
+                return proxy
+            proxy = get_proxy_slot(obj, self, None)
+            if proxy is not None:
+                return proxy
+            # The object may be identity-different but equal to a tracked
+            # FSO's real_obj. This happens because maybe_to_fake_obj creates
+            # a new DeviceMesh wrapper, so the FSO's real_obj is a distinct
+            # object from the one held by submesh._root_mesh. Equality
+            # comparison is needed, not identity.
+            for tracked_obj, p in self.opaque_tracker.items():
+                if not isinstance(tracked_obj, FakeScriptObject):
+                    continue
+                if tracked_obj.real_obj == obj:
+                    return p
+            return None
+
+        result = reconstruct_fn(real_obj, get_tracked_proxy, self)
+        if result is None:
+            return None
+
+        # The reconstruct_fn dispatches through a custom op with a Proxy
+        # argument, which goes through Proxy.__torch_function__.  That path
+        # creates a graph node but does NOT populate meta["val"] (unlike the
+        # ProxyTorchDispatchMode.__torch_dispatch__ path which calls
+        # track_tensor_tree).  Set it here so downstream consumers (e.g. the
+        # min-cut partitioner) can classify the node correctly.
+        if "val" not in result.node.meta:
+            set_meta(result, a)
+
+        # Also register for the *input* object so dedup works for re-encounters.
+        set_proxy_slot(a, self, result)
+        if id(real_obj) not in self._opaque_real_obj_proxy:
+            self._opaque_real_obj_proxy[id(real_obj)] = result
+
+        return result.node
+
     @overload
-    def unwrap_proxy(self, e: Tensor) -> Union[Proxy, Tensor]: ...
+    def unwrap_proxy(self, e: Tensor) -> Proxy | Tensor: ...
 
     @overload
-    def unwrap_proxy(self, e: PySymType) -> Union[Proxy, PySymType]: ...
+    def unwrap_proxy(self, e: PySymType) -> Proxy | PySymType: ...
 
     @overload
-    def unwrap_proxy(
-        self, e: _AnyScriptObjectType
-    ) -> Union[Proxy, _AnyScriptObjectType]: ...
+    def unwrap_proxy(self, e: _AnyScriptObjectType) -> Proxy | _AnyScriptObjectType: ...
 
     def unwrap_proxy(self, e: T) -> object:
         if isinstance(e, Tensor):
@@ -1400,8 +1540,8 @@ def create_node(
         target: Target,
         args: tuple[Argument, ...],
         kwargs: dict[str, Argument],
-        name: Optional[str] = None,
-        type_expr: Optional[Any] = None,
+        name: str | None = None,
+        type_expr: Any | None = None,
     ) -> torch.fx.Node:
         node = super().create_node(kind, target, args, kwargs, name, type_expr)  # type: ignore[arg-type]
 
@@ -1416,7 +1556,7 @@ def create_node(
                 with disable_proxy_modes_tracing():
                     node.meta["val"] = extract_val(attr)
 
-        def map_fn(v: Any) -> Optional[_ExtractValType]:
+        def map_fn(v: Any) -> _ExtractValType | None:
             if not isinstance(v, torch.fx.Node) or "val" not in v.meta:
                 return None
             val = v.meta["val"]
@@ -1440,7 +1580,7 @@ def map_fn(v: Any) -> Optional[_ExtractValType]:
 
 def _should_save_eager_input_vals(
     target: Any,
-    args_kwargs: Optional[tuple[tuple[Argument, ...], dict[str, Argument]]] = None,
+    args_kwargs: tuple[tuple[Argument, ...], dict[str, Argument]] | None = None,
 ) -> bool:
     from torch._higher_order_ops.invoke_subgraph import InvokeSubgraphHOP
 
@@ -1487,9 +1627,9 @@ def _should_save_eager_input_vals(
 
 def _make_temp_remove_mode_context_manager(
     mode_ty: type[TorchFunctionMode],
-) -> Callable[[], _GeneratorContextManager[Optional[TorchFunctionMode]]]:
+) -> Callable[[], _GeneratorContextManager[TorchFunctionMode | None]]:
     @contextmanager
-    def context_manager_fn() -> Generator[Optional[TorchFunctionMode], None, None]:
+    def context_manager_fn() -> Generator[TorchFunctionMode | None, None, None]:
         from torch.overrides import _len_torch_function_stack, _pop_mode, _push_mode
 
         temp_elements = []
@@ -1526,9 +1666,9 @@ def context_manager_fn() -> Generator[Optional[TorchFunctionMode], None, None]:
 
 @torch._disable_dynamo
 def dispatch_trace(
-    root: Union[Module, Callable],
+    root: Module | Callable[..., Any],
     tracer: Tracer,
-    concrete_args: Optional[tuple[Any, ...]] = None,
+    concrete_args: tuple[object, ...] | None = None,
 ) -> GraphModule:
     graph = tracer.trace(root, concrete_args)  # type: ignore[arg-type]
 
@@ -1597,7 +1737,7 @@ def wrapped(*proxies: _P.args, **_unused: _P.kwargs) -> R:
                 for t, p in zip(tensors, proxies)  # type: ignore[arg-type]
             ]
 
-        def get_tensor_proxy_slot(t: Tensor) -> Union[Tensor, Proxy]:
+        def get_tensor_proxy_slot(t: Tensor) -> Tensor | Proxy:
             return get_proxy_slot(t, tracer, t, lambda x: x.proxy)  # type: ignore[attr-defined]
 
         out = f(*tensors)  # type:ignore[call-arg]
@@ -1605,6 +1745,9 @@ def get_tensor_proxy_slot(t: Tensor) -> Union[Tensor, Proxy]:
         out = pytree.tree_map_only(
             _AnyScriptObject, lambda t: get_proxy_slot(t, tracer, t, lambda x: x), out
         )
+        out = pytree.tree_map_only(
+            OpaqueBase, lambda t: get_proxy_slot(t, tracer, t, lambda x: x), out
+        )
 
         def get_sym_proxy_slot(t: PySymType) -> Proxy:
             return get_proxy_slot(t, tracer).force()
@@ -1616,21 +1759,22 @@ def get_sym_proxy_slot(t: PySymType) -> Proxy:
 
 
 # TODO: Make downstream users of this work with OperatorBase
-ORIGINAL_ATEN: Optional[object] = None
+ORIGINAL_ATEN: contextvars.ContextVar[object | None] = contextvars.ContextVar(
+    "ORIGINAL_ATEN", default=None
+)
 
 
 @contextmanager
 def set_original_aten_op(
     func: OpOverload | torch._ops.HigherOrderOperator,
 ) -> Generator[None, None, None]:
-    global ORIGINAL_ATEN
-    if ORIGINAL_ATEN is None and fx_traceback.has_preserved_node_meta():
-        ORIGINAL_ATEN = func
+    if ORIGINAL_ATEN.get() is None and fx_traceback.has_preserved_node_meta():
+        token = ORIGINAL_ATEN.set(func)
         fx_traceback.current_meta["original_aten"] = func
         try:
             yield
         finally:
-            ORIGINAL_ATEN = None
+            ORIGINAL_ATEN.reset(token)
             fx_traceback.current_meta["original_aten"] = None
     else:
         yield
@@ -1645,7 +1789,7 @@ def __torch_function__(
         func: OpOverload,
         types: tuple[torch._C._TensorMeta, ...],
         args: tuple[object, ...] = (),
-        kwargs: Optional[dict[str, object]] = None,
+        kwargs: dict[str, object] | None = None,
     ) -> object:
         kwargs = kwargs or {}
         # pyrefly: ignore [bad-assignment]
@@ -1672,10 +1816,10 @@ def __init__(self, tracer: _ProxyTracer) -> None:
 
     def __torch_function__(
         self,
-        func: Union[OpOverload, Callable],
+        func: OpOverload | Callable[..., Any],
         types: tuple[torch._C._TensorMeta, ...],
         args: tuple[object, ...] = (),
-        kwargs: Optional[dict[str, object]] = None,
+        kwargs: dict[str, object] | None = None,
     ) -> object:
         kwargs = kwargs or {}
         if func in _side_effectful_need_to_be_preserved_pre_dispatch:
@@ -1739,7 +1883,7 @@ def enable_tracing(self) -> bool:
     def __init__(
         self,
         tracer: _ProxyTracer,
-        tracing_mode: str,
+        tracing_mode: _TracingMode,
         pre_dispatch: bool = False,
         _allow_fake_constant: bool = False,
         _error_on_data_dependent_ops: bool = True,
@@ -1757,7 +1901,7 @@ def __init__(
         # Every time we enter a mode, we maintain a stack telling us what the previous
         # ProxyTorchDispatchMode state was (if there was any).
         # This lets us properly reset the state on exit.
-        self.enter_stack: list[Optional[ProxyTorchDispatchMode]] = []
+        self.enter_stack: list[ProxyTorchDispatchMode | None] = []
         self.decomp_layers: int = 0
         # See invoke_subgraph
         self._invoke_subgraph_names: set[str] = set()
@@ -1774,7 +1918,7 @@ def __torch_dispatch__(
         func: OpOverload,
         types: tuple[torch._C._TensorMeta, ...],
         args: tuple[object, ...] = (),
-        kwargs: Optional[dict[str, object]] = None,
+        kwargs: dict[str, object] | None = None,
     ) -> object:
         with set_original_aten_op(func):
             kwargs = kwargs or {}
@@ -1792,10 +1936,10 @@ def __enter__(self) -> Self:
 
     def __exit__(
         self,
-        exc_type: Optional[type[BaseException]],
-        exc_value: Optional[BaseException],
-        traceback: Optional[types.TracebackType],
-    ) -> Optional[bool]:
+        exc_type: type[BaseException] | None,
+        exc_value: BaseException | None,
+        traceback: types.TracebackType | None,
+    ) -> bool | None:
         b = super().__exit__(exc_type, exc_value, traceback)
 
         # Re-enable the previous proxy mode, if there was one.
@@ -1883,27 +2027,21 @@ def _compute_proxy(
 
 
 class _GraphAppendingTracerEx(fx.proxy.GraphAppendingTracer):
-    script_object_tracker: MutableMapping[_AnyScriptObjectType, Proxy]
-    symnode_tracker: MutableMapping[PySymType, _PySymProxyType]
+    script_object_tracker: MutableMapping[torch.ScriptObject, Proxy]
+    opaque_tracker: MutableMapping[FakeScriptObject | OpaqueBase, Proxy]
+    # Maps id(real_obj) -> proxy for opaque FSOs, so that multiple FSO wrappers
+    # of the same real object (e.g. primal vs tangent) resolve to one proxy.
+    _opaque_real_obj_proxy: dict[int, Proxy]
+    symnode_tracker: _SymNodeDict
     tensor_tracker: MutableMapping[Tensor, _ProxyTensor]
     sympy_expr_tracker: dict[sympy.Symbol, _SympyExprTrackerValue]
-    torch_fn_metadata: Optional[OpOverload]
+    torch_fn_metadata: OpOverload | None
     torch_fn_counts: dict[OpOverload, int]
     enable_thunkify: bool = False
 
     def __init__(self, graph: fx.graph.Graph) -> None:
         super().__init__(graph)
-        self.symnode_tracker = weakref.WeakKeyDictionary()
-        self.tensor_tracker = WeakTensorKeyDictionary()
-        self.sympy_expr_tracker = {}
-        self.script_object_tracker = WeakIdKeyDictionary(
-            dict=None, ref_type=_WeakHashRef
-        )
-        # Stores the torch function that was called during tracing
-        self.torch_fn_metadata = None
-        # Stores the counts for every torch function called. This is to help
-        # distinguish between different calls to the same torch function.
-        self.torch_fn_counts = {}
+        _init_proxy_trackers(self)
 
 
 # TODO: I'm not sure what the point of this class is; you can just
@@ -1913,7 +2051,7 @@ def __init__(
         self,
         module: fx.GraphModule,
         new_graph: fx.Graph,
-        decomposition_table: Optional[Mapping[OpOverload, Callable]] = None,
+        decomposition_table: Mapping[OpOverload, Callable[..., Any]] | None = None,
         **kwargs: object,
     ) -> None:
         super().__init__(module, **kwargs)  # type: ignore[arg-type]
@@ -1962,7 +2100,7 @@ def output(
         def get_proxy_node(x: _ProxyTensor) -> fx.node.Node:
             return x.proxy.node
 
-        def unwrap(e: Tensor) -> Union[Tensor, fx.Node]:
+        def unwrap(e: Tensor) -> Tensor | fx.Node:
             return get_proxy_slot(e, self.tracer, e, get_proxy_node)
 
         self.new_graph.output(pytree.tree_map(unwrap, out))
@@ -1980,7 +2118,7 @@ def __init__(
         self,
         module: fx.GraphModule,
         should_decompose: Callable[[fx.Node], bool],
-        decomposition_table: Mapping[OpOverload, Callable],
+        decomposition_table: Mapping[OpOverload, Callable[..., Any]] | None,
         **kwargs: object,
     ) -> None:
         """
@@ -1995,7 +2133,7 @@ def __init__(
     def recursive_wrap(
         gm: fx.GraphModule,
         should_decompose: Callable[[fx.Node], bool],
-        decomposition_table: Mapping[OpOverload, Callable],
+        decomposition_table: Mapping[OpOverload, Callable[..., Any]] | None,
         **kwargs: object,
     ) -> _SelectiveDecomposeInterpreter:
         """
@@ -2026,7 +2164,7 @@ def recursive_wrap(
             gm, should_decompose, decomposition_table, **kwargs
         )
 
-    def run_node(self, n):
+    def run_node(self, n: fx.Node) -> Any:
         if self.should_decompose(n):
             with decompose(self.decomposition_table):
                 result = super().run_node(n)
@@ -2037,9 +2175,9 @@ def run_node(self, n):
 
 def selective_decompose(
     joint_gm: fx.GraphModule,
-    *args,
-    decomposition,
-    should_decompose,
+    *args: object,
+    decomposition: Mapping[OpOverload, Callable[..., Any]] | None,
+    should_decompose: Callable[..., bool],
     trace_joint_graph: bool,
 ) -> fx.GraphModule:
     """Retrace a joint graph module and selectively apply decomposition."""
@@ -2048,13 +2186,13 @@ def selective_decompose(
         # the arg name, primals and tangents, are important.
         # make_fx keeps the name in the traced graph and partitioner later relies
         # on the name to partition joint graph correctly.
-        def wrap_fn(primals: list[Any], tangents: list[Any]):
+        def wrap_fn(primals: list[Any], tangents: list[Any]) -> Any:
             return _SelectiveDecomposeInterpreter.recursive_wrap(
                 joint_gm, should_decompose, decomposition
             ).run(*args)
     else:
 
-        def wrap_fn(*args):
+        def wrap_fn(*args: Any) -> Any:
             return _SelectiveDecomposeInterpreter.recursive_wrap(
                 joint_gm, should_decompose, decomposition
             ).run(*args)
@@ -2117,7 +2255,7 @@ class _ModuleStackTracer(PythonKeyTracer):
 
     def __init__(self, scope_root: GraphModule) -> None:
         super().__init__()
-        self.record_stack_traces = True
+        self.record_stack_traces = not fx.config.do_not_emit_stack_traces
         self._record_forward_stack_traces_only = True
         self.scope_root = scope_root
         self.enable_attr_proxy = False
@@ -2148,7 +2286,7 @@ def __init__(self, scope_root: GraphModule) -> None:
         tracer = self
 
         class AttrProxy(_AttrProxy):
-            def __init__(self, base: Union[Module, _AttrProxy], path: str) -> None:
+            def __init__(self, base: Module | _AttrProxy, path: str) -> None:
                 if isinstance(base, _AttrProxy):
                     base = base.get_base()  # type: ignore[attr-defined]
 
@@ -2187,7 +2325,7 @@ def __getattr__(self, name: str) -> AttrProxy:
             def get_base(self) -> Module:
                 return tracer.proxy_modules[self]
 
-            def __getitem__(self, idx: Union[int, slice]) -> AttrProxy:
+            def __getitem__(self, idx: int | slice) -> AttrProxy:
                 if isinstance(idx, slice):
                     if isinstance(self, torch.nn.Sequential):
                         # Copied from nn/modules/container.py
@@ -2259,7 +2397,7 @@ def getattr(
         return self.attr_proxy_map[attr_val]
 
     def trace(  # type: ignore[override]
-        self, root: Union[Module, Callable], concrete_args: Optional[dict[str, object]]
+        self, root: Module | Callable[..., Any], concrete_args: dict[str, object] | None
     ) -> fx.Graph:
         res = super().trace(root, concrete_args)
 
@@ -2330,7 +2468,7 @@ def _delete_proxy_attr(obj: Module, target: str) -> bool:
     def call_module(
         self,
         m: Module,
-        forward: Callable,
+        forward: Callable[..., Any],
         args: tuple[object, ...],
         kwargs: dict[str, object],
     ) -> None:
@@ -2357,7 +2495,7 @@ def call_module(
                 "This might be because the module was not properly registered "
                 "as a submodule, which is not good practice. We will trace "
                 "through the module without recording stack information.",
-                str(m),
+                m,
             )
             return forward(*args, **kwargs)
 
@@ -2389,12 +2527,14 @@ def create_node(self, *args: object, **kwargs: object) -> fx.node.Node:
         # torch_fn
         if (
             node.op == "call_function"
-            and self.torch_fn_metadata is not None
+            and (torch_fn := self.torch_fn_metadata) is not None
             and "torch_fn" not in node.meta
         ):
             node.meta["torch_fn"] = (
-                f"{self.torch_fn_metadata.__name__}_{self.torch_fn_counts[self.torch_fn_metadata]}",
-                f"{self.torch_fn_metadata.__class__.__name__}.{self.torch_fn_metadata.__name__}",
+                # pyrefly: ignore[missing-attribute,bad-index]
+                f"{torch_fn.__name__}_{self.torch_fn_counts[torch_fn]}",
+                # pyrefly: ignore[missing-attribute]
+                f"{torch_fn.__class__.__name__}.{torch_fn.__name__}",
             )
 
         return node
@@ -2403,27 +2543,27 @@ def create_node(self, *args: object, **kwargs: object) -> fx.node.Node:
 class _MakefxTracer:
     def __init__(
         self,
-        decomposition_table: Optional[Mapping[OpOverload, Callable]],
-        tracing_mode: str,
+        decomposition_table: Mapping[OpOverload, Callable[..., Any]] | None,
+        tracing_mode: _TracingMode,
         _allow_non_fake_inputs: bool,
         pre_dispatch: bool,
         record_module_stack: bool,
         _allow_fake_constant: bool,
         _error_on_data_dependent_ops: bool,
         record_stack_traces: bool = False,
-        parent_tracer: Optional[_MakefxTracer] = None,
+        parent_tracer: _MakefxTracer | None = None,
         proxy_module_inputs: bool = False,
         _disable_torch_fn_metadata_mode: bool = False,
     ) -> None:
         # Configurations that are used to initialize the context managers and their states.
         # Should not modify them during tracing.
-        self.decomposition_table: dict[OpOverload, Callable] = dict(
+        self.decomposition_table: dict[OpOverload, Callable[..., Any]] = dict(
             decomposition_table or {}
         )
         self.decomposition_table.setdefault(
             torch.ops.aten.sym_numel.default, torch._decomp.decompositions.sym_numel
         )
-        self.tracing_mode: str = tracing_mode
+        self.tracing_mode: _TracingMode = tracing_mode
         self._allow_non_fake_inputs: bool = _allow_non_fake_inputs
         self.pre_dispatch: bool = pre_dispatch
         self.record_module_stack: bool = record_module_stack
@@ -2434,18 +2574,18 @@ def __init__(
         # and configurations. After tracing, their states should be cleaned except for shape_env.
         # Remember to specify how to initialize it from user inputs and from parent tracer whenever
         # adding new modes in _MakefxTracer.
-        self.fake_tensor_mode: Optional[FakeTensorMode] = None
-        self.proxy_mode: Union[nullcontext, ProxyTorchDispatchMode] = nullcontext()
-        self.proxy_function_mode: Union[nullcontext, PreDispatchTorchFunctionMode] = (
+        self.fake_tensor_mode: FakeTensorMode | None = None
+        self.proxy_mode: nullcontext[None] | ProxyTorchDispatchMode = nullcontext()
+        self.proxy_function_mode: nullcontext[None] | PreDispatchTorchFunctionMode = (
             nullcontext()
         )
-        self.fx_tracer: Optional[PythonKeyTracer] = None
-        self.python_dispatcher_mode: Union[nullcontext, Any] = nullcontext()
-        self.torch_fn_metadata_mode: Union[nullcontext, TorchFunctionMetadataMode] = (
+        self.fx_tracer: PythonKeyTracer | None = None
+        self.python_dispatcher_mode: nullcontext[None] | Any = nullcontext()
+        self.torch_fn_metadata_mode: nullcontext[None] | TorchFunctionMetadataMode = (
             nullcontext()
         )
         self.record_stack_traces = record_stack_traces
-        self.parent_tracer: Optional[_MakefxTracer] = parent_tracer
+        self.parent_tracer: _MakefxTracer | None = parent_tracer
         self.proxy_module_inputs = proxy_module_inputs
         self._disable_torch_fn_metadata_mode = _disable_torch_fn_metadata_mode
 
@@ -2461,12 +2601,12 @@ def _checkpoint_modes(self) -> list[Any]:
 
     def _restore_modes(
         self,
-        prev_fake_tensor_mode: Optional[FakeTensorMode],
-        prev_proxy_mode: Union[nullcontext, ProxyTorchDispatchMode],
-        prev_proxy_function_mode: Union[nullcontext, PreDispatchTorchFunctionMode],
-        prev_fx_tracer: Optional[PythonKeyTracer],
-        prev_python_dispatcher_mode: Union[nullcontext, Any],
-        prev_torch_fn_metadata_mode: Union[nullcontext, TorchFunctionMetadataMode],
+        prev_fake_tensor_mode: FakeTensorMode | None,
+        prev_proxy_mode: nullcontext[None] | ProxyTorchDispatchMode,
+        prev_proxy_function_mode: nullcontext[None] | PreDispatchTorchFunctionMode,
+        prev_fx_tracer: PythonKeyTracer | None,
+        prev_python_dispatcher_mode: nullcontext[None] | Any,
+        prev_torch_fn_metadata_mode: nullcontext[None] | TorchFunctionMetadataMode,
     ) -> None:
         self.fake_tensor_mode = prev_fake_tensor_mode
         self.proxy_mode = prev_proxy_mode
@@ -2477,7 +2617,7 @@ def _restore_modes(
 
     @contextmanager
     def _init_modes_from_inputs(
-        self, f: Callable, args: tuple[object, ...]
+        self, f: Callable[..., Any], args: tuple[object, ...]
     ) -> Generator[None, None, None]:
         prev_modes = self._checkpoint_modes()
         try:
@@ -2491,8 +2631,10 @@ def _init_modes_from_inputs(
                 self.fx_tracer = _ModuleStackTracer(scope_root)
             else:
                 self.fx_tracer = PythonKeyTracer()
-                self.fx_tracer.record_stack_traces = self.record_stack_traces
-                if self.record_stack_traces:
+                self.fx_tracer.record_stack_traces = (
+                    self.record_stack_traces and not fx.config.do_not_emit_stack_traces
+                )
+                if self.fx_tracer.record_stack_traces:
                     self.fx_tracer._record_forward_stack_traces_only = True
 
             if self.tracing_mode == "fake":
@@ -2590,7 +2732,57 @@ def _create_sub_fx_tracer(parent_tracer: _ProxyTracer) -> PythonKeyTracer:
         finally:
             self._restore_modes(*prev_modes)
 
-    def _trace_inner(self, f: Callable, *args: object) -> GraphModule:
+    def _convert_args_to_fake(self, args: T) -> T:
+        if self.tracing_mode == "real":
+            return args
+
+        arg_count = 0
+
+        def inner_wrap_fake(x: object) -> object:
+            nonlocal arg_count
+            # TODO: it would be nice to line these up with the names
+            # FX will choose for the placeholders, but we don't
+            # actually know what the names will be at this point yet
+            # NB: the Source here is actually meaningless
+            from torch._dynamo.source import ConstantSource
+
+            if self.fake_tensor_mode is None:
+                raise AssertionError("fake_tensor_mode should not be None")
+            source = ConstantSource(f"input{arg_count}")
+            if isinstance(x, Tensor):
+                arg_count += 1
+                return self.fake_tensor_mode.from_tensor(x, source=source)
+            # NB: don't match on bools
+            elif type(x) is int and self.tracing_mode == "symbolic":
+                if self.fake_tensor_mode.shape_env is None:
+                    raise AssertionError(
+                        "shape_env should be set if tracing with 'symbolic'"
+                    )
+                return self.fake_tensor_mode.shape_env.create_symintnode(
+                    self.fake_tensor_mode.shape_env.create_symbol(
+                        x, source, positive=None
+                    ),
+                    hint=x,
+                    source=source,
+                )
+            elif isinstance(x, torch.ScriptObject) or is_opaque_value(x):
+                if is_opaque_value_type(
+                    type(x)  # pyrefly: ignore[bad-argument-type]
+                ):
+                    return x
+                return torch._library.fake_class_registry.maybe_to_fake_obj(
+                    self.fake_tensor_mode, x
+                )
+
+            if isinstance(x, FakeScriptObject):
+                raise AssertionError(
+                    f"ScriptObject {x} has been fakified. Cannot wrap_fake it again."
+                )
+            return x
+
+        return pytree.tree_map(inner_wrap_fake, args)
+
+    def _trace_inner(self, f: Callable[..., Any], *args: object) -> GraphModule:
         # TODO: We need to explicitly import torch._dynamo before calling dispatch_trace,
         # because dispatch_trace will introduce the lazy import of torch._dynamo,
         # and some contexts set before calling dispatch_trace will cause problems with the import of torch._dynamo,
@@ -2600,66 +2792,16 @@ def _trace_inner(self, f: Callable, *args: object) -> GraphModule:
 
         phs = pytree.tree_map(lambda _: torch.fx._symbolic_trace.PH, args)
 
-        def _wrap_fake(args: T) -> T:
-            arg_count = 0
-
-            def inner_wrap_fake(x: object) -> object:
-                nonlocal arg_count
-                # TODO: it would be nice to line these up with the names
-                # FX will choose for the placeholders, but we don't
-                # actually know what the names will be at this point yet
-                # NB: the Source here is actually meaningless
-                from torch._dynamo.source import ConstantSource
-
-                if self.fake_tensor_mode is None:
-                    raise AssertionError("fake_tensor_mode should not be None")
-                source = ConstantSource(f"input{arg_count}")
-                if isinstance(x, Tensor):
-                    arg_count += 1
-                    return self.fake_tensor_mode.from_tensor(x, source=source)
-                # NB: don't match on bools
-                elif type(x) is int and self.tracing_mode == "symbolic":
-                    if self.fake_tensor_mode.shape_env is None:
-                        raise AssertionError(
-                            "shape_env should be set if tracing with 'symbolic'"
-                        )
-                    return self.fake_tensor_mode.shape_env.create_symintnode(
-                        self.fake_tensor_mode.shape_env.create_symbol(
-                            x, source, positive=None
-                        ),
-                        hint=x,
-                        source=source,
-                    )
-                elif isinstance(x, torch.ScriptObject) or is_opaque_value(x):
-                    return torch._library.fake_class_registry.maybe_to_fake_obj(
-                        self.fake_tensor_mode, x
-                    )
+        args = self._convert_args_to_fake(args)
 
-                if isinstance(x, FakeScriptObject):
-                    raise AssertionError(
-                        f"ScriptObject {x} has been fakified. Cannot wrap_fake it again."
-                    )
-                return x
-
-            wrap_fn_map = {
-                "real": lambda x: x,
-                "fake": inner_wrap_fake,
-                "symbolic": inner_wrap_fake,
-            }
-            return pytree.tree_map(wrap_fn_map[self.tracing_mode], args)
-
-        def _wrap_func(f: Callable[_P, R], phs: Sequence[PHBase]) -> Callable[_P, R]:
-            if (
-                not hasattr(inspect.unwrap(f), "__code__")
-                or inspect.unwrap(f).__code__.co_flags & inspect.CO_VARARGS
-            ):
-                # FX doesn't support varargs, so we gotta fake up a wrapper
-                # TODO: Would be nice to fix this at the source...
-                return fake_signature(f, len(phs))
-            return f
-
-        args = _wrap_fake(args)
-        func = _wrap_func(f, phs)
+        # FX doesn't support varargs, so we gotta fake up a wrapper
+        # TODO: Would be nice to fix this at the source...
+        func: Callable[..., Any] = f
+        if (
+            not hasattr(inspect.unwrap(f), "__code__")
+            or inspect.unwrap(f).__code__.co_flags & inspect.CO_VARARGS
+        ):
+            func = fake_signature(f, len(phs))
         # We disable the autocast cache as the autocast cache causes type conversions on parameters to
         # check a cache, which introduces untracked tensors into the graph
         #
@@ -2720,18 +2862,21 @@ def _wrap_func(f: Callable[_P, R], phs: Sequence[PHBase]) -> Callable[_P, R]:
             t.shape_env = self.fake_tensor_mode.shape_env  # type: ignore[assignment]
         return t
 
-    def trace(self, f: Callable, *args: object) -> fx.GraphModule:
+    def trace(self, f: Callable[..., Any], *args: object) -> fx.GraphModule:
         with self._init_modes_from_inputs(f, args):
             return self._trace_inner(f, *args)
 
     def is_hop_subgraph_tracer(self) -> bool:
         return self.parent_tracer is not None
 
-    def trace_subgraph(self, f: Callable, *args: object) -> GraphModule:
-        # Create a new tracer based on parent's config
-        sub_tracer = _MakefxTracer(
-            self.decomposition_table,
-            "real",
+    def _make_sub_tracer(
+        self,
+        decomp_table: Mapping[OpOverload, Callable[..., Any]] | None = None,
+        tracing_mode: _TracingMode = "real",
+    ) -> _MakefxTracer:
+        return _MakefxTracer(
+            decomp_table if decomp_table is not None else self.decomposition_table,
+            tracing_mode,
             self._allow_non_fake_inputs,
             self.pre_dispatch,
             self.record_module_stack,
@@ -2739,30 +2884,26 @@ def trace_subgraph(self, f: Callable, *args: object) -> GraphModule:
             self._error_on_data_dependent_ops,
             parent_tracer=self,
         )
+
+    def trace_subgraph(self, f: Callable[..., Any], *args: object) -> GraphModule:
+        sub_tracer = self._make_sub_tracer()
         with sub_tracer._init_modes_from_parent(self):
             return sub_tracer._trace_inner(f, *args)
 
     def trace_subgraph_custom_decomp(
-        self, f: Callable, decomp_table: Mapping[OpOverload, Callable], *args
+        self,
+        f: Callable[..., Any],
+        decomp_table: Mapping[OpOverload, Callable[..., Any]],
+        *args: object,
     ) -> GraphModule:
         if not isinstance(decomp_table, Mapping):
             raise AssertionError(f"Expected Mapping, got {type(decomp_table)}")
-        # Create a new tracer based on parent's config, but use a different decomposition table
-        sub_tracer = _MakefxTracer(
-            decomp_table,
-            "real",
-            self._allow_non_fake_inputs,
-            self.pre_dispatch,
-            self.record_module_stack,
-            self._allow_fake_constant,
-            self._error_on_data_dependent_ops,
-            parent_tracer=self,
-        )
+        sub_tracer = self._make_sub_tracer(decomp_table=decomp_table)
         with sub_tracer._init_modes_from_parent(self):
             return sub_tracer._trace_inner(f, *args)
 
 
-_CURRENT_MAKE_FX_TRACER: Optional[_MakefxTracer] = None
+_CURRENT_MAKE_FX_TRACER: _MakefxTracer | None = None
 
 
 @contextmanager
@@ -2777,9 +2918,9 @@ def _set_make_fx_tracer(tracer: _MakefxTracer) -> Generator[None, None, None]:
 
 
 def make_fx(
-    f: Callable,
-    decomposition_table: Optional[Mapping[OpOverload, Callable]] = None,
-    tracing_mode: str = "real",
+    f: Callable[..., Any],
+    decomposition_table: Mapping[OpOverload, Callable[..., Any]] | None = None,
+    tracing_mode: _TracingMode = "real",
     _allow_non_fake_inputs: bool = False,
     *,
     pre_dispatch: bool = False,
@@ -2832,11 +2973,11 @@ def get_torch_dispatch_modes() -> list[TorchDispatchMode]:
 
 # TODO: this is a legacy name, there is only ever one proxy mode as it's an
 # infra mode
-def get_innermost_proxy_mode() -> Optional[ProxyTorchDispatchMode]:
+def get_innermost_proxy_mode() -> ProxyTorchDispatchMode | None:
     return get_proxy_mode()
 
 
-def get_proxy_mode() -> Optional[ProxyTorchDispatchMode]:
+def get_proxy_mode() -> ProxyTorchDispatchMode | None:
     """
     Current the currently active proxy tracing mode, or None if
     we are not currently tracing.  This includes pre-dispatch proxy
@@ -2885,7 +3026,8 @@ def maybe_handle_decomp(
 ) -> object:
     from torch._inductor.compiler_bisector import CompilerBisector
 
-    if op in CURRENT_DECOMPOSITION_TABLE:
+    decomp_table = CURRENT_DECOMPOSITION_TABLE.get({})
+    if op in decomp_table:
         if CompilerBisector.disable_subsystem(
             "aot_eager_decomp_partition", "decomposition", lambda: repr(op)
         ):
@@ -2893,7 +3035,7 @@ def maybe_handle_decomp(
 
         with proxy_mode:
             proxy_mode.decomp_layers += 1
-            out = CURRENT_DECOMPOSITION_TABLE[op](*args, **kwargs)
+            out = decomp_table[op](*args, **kwargs)
             proxy_mode.decomp_layers -= 1
             return out
 
@@ -2901,11 +3043,11 @@ def maybe_handle_decomp(
 
 
 def get_isolated_graphmodule(
-    func: Callable,
+    func: Callable[..., Any],
     args: tuple[object, ...],
     kwargs: dict[str, object],
-    tracing_mode: str = "real",
-    decomposition_table: Optional[Mapping[OpOverload, Callable]] = None,
+    tracing_mode: _TracingMode = "real",
+    decomposition_table: Mapping[OpOverload, Callable[..., Any]] | None = None,
 ) -> GraphModule:
     """A helper function used to get the GraphModule for the given func.
 
diff --git a/torch/fx/experimental/recording.py b/torch/fx/experimental/recording.py
index f4bec90fd9bbe..0d401df3d18ad 100644
--- a/torch/fx/experimental/recording.py
+++ b/torch/fx/experimental/recording.py
@@ -1,16 +1,26 @@
-# mypy: allow-untyped-defs
+from __future__ import annotations
+
 import functools
 import inspect
 import itertools
 import logging
-from collections.abc import Callable
 from dataclasses import dataclass
-from typing import Any, Optional, Union
+from typing import Any, ParamSpec, TYPE_CHECKING, TypeVar
+
+
+_P = ParamSpec("_P")
+_R = TypeVar("_R")
 
 import torch
 import torch.utils._pytree as pytree
 
 
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+    from torch.fx.experimental.symbolic_shapes import ShapeEnv, TrackedFake
+
+
 log = logging.getLogger(__name__)
 trace_shape_events_log = torch._logging.getArtifactLogger(
     __name__, "trace_shape_events"
@@ -81,21 +91,21 @@
 @dataclass
 class ShapeEnvEvent:
     # ShapeEnv method.
-    f: Callable
+    f: Callable[..., Any]
 
     # Arguments and keyword arguments called with.
-    args: Optional[list[Any]] = None
-    kwargs: Optional[dict[str, Any]] = None
+    args: list[object] | None = None
+    kwargs: dict[str, Any] | None = None
 
     # List of tracked_fakes at the time the method was called.
-    tracked_fakes: Optional[list[Any]] = None
+    tracked_fakes: list[TrackedFake] | None = None
 
     # Name of the captured event.
     # Used for special handling of particular methods.
-    name: Optional[str] = None
+    name: str | None = None
 
     # Replay itself, but using shape_env as self.
-    def run(self, shape_env=None) -> Any:
+    def run(self, shape_env: ShapeEnv | None = None) -> Any:
         from torch.fx.experimental.symbolic_shapes import (
             is_symbolic,
             ShapeEnv,
@@ -147,7 +157,7 @@ def maybe_convert_node(x: Any) -> Any:
             return name_to_node[x.name]
 
         # Replaces the value of an specific argument by the result of fn.
-        def replacearg(index: int, key: str, fn: Callable):
+        def replacearg(index: int, key: str, fn: Callable[..., Any]) -> None:
             if index < len(args):
                 args[index] = fn(args[index])
             if key in kwargs:
@@ -196,10 +206,12 @@ def is_defer_runtime_assert(self) -> bool:
 #   2. SymInt, SymFloat, or SymBool arguments
 # If we find more than one object of any of the above types, we
 # also check that the ShapeEnv instance is the same for all of them.
-def _extract_shape_env_and_assert_equal(args, kwargs):
+def _extract_shape_env_and_assert_equal(
+    args: tuple[object, ...] | list[object], kwargs: dict[str, object]
+) -> ShapeEnv | None:
     from torch.fx.experimental.symbolic_shapes import is_symbolic, ShapeEnv, SymTypes
 
-    def assert_equal(old: Optional[ShapeEnv], new: ShapeEnv) -> ShapeEnv:
+    def assert_equal(old: ShapeEnv | None, new: ShapeEnv) -> ShapeEnv:
         if old is not None:
             if old is not new:
                 raise AssertionError("call with different ShapeEnv")
@@ -240,9 +252,9 @@ def assert_equal(old: Optional[ShapeEnv], new: ShapeEnv) -> ShapeEnv:
 #   - ShapeEnv.evaluate_expr
 #   - ShapeEnv.guard_or_defer_runtime_assert
 def record_shapeenv_event(
-    *, save_tracked_fakes: bool = False, name: Optional[str] = None
-) -> Callable:
-    def decorator(fn: Callable) -> Callable:
+    *, save_tracked_fakes: bool = False, name: str | None = None
+) -> Callable[[Callable[_P, _R]], Callable[_P, _R]]:
+    def decorator(fn: Callable[_P, _R]) -> Callable[_P, _R]:
         if not callable(fn):
             raise AssertionError(f"Expected callable, got {type(fn)}")
         args = inspect.getfullargspec(fn).args
@@ -256,7 +268,7 @@ def decorator(fn: Callable) -> Callable:
             name = fn.__name__
 
         @functools.wraps(fn)
-        def wrapper(*args, **kwargs):
+        def wrapper(*args: _P.args, **kwargs: _P.kwargs) -> _R:
             from torch.fx.experimental.symbolic_shapes import ShapeEnv
 
             if not isinstance(args[0], ShapeEnv):
@@ -269,7 +281,7 @@ def wrapper(*args, **kwargs):
             )
             NEST += 1
 
-            def retlog(r):
+            def retlog(r: _R) -> _R:
                 trace_shape_events_log.debug("%s-> %s", " " * (NEST - 1), r)
                 return r
 
@@ -346,7 +358,7 @@ def retlog(r):
 # It assumes the first event is the constructor call.
 #
 # fn: transforms an old FX node into one corresponding to the newly created ShapeEnv.
-def replay_shape_env_events(events):
+def replay_shape_env_events(events: list[ShapeEnvEvent]) -> ShapeEnv:
     from torch.fx.experimental.symbolic_shapes import ShapeEnv
 
     constructor_event = events[0]
@@ -376,25 +388,25 @@ def replay_shape_env_events(events):
 # ShapeEnv.produce_guards.
 @dataclass
 class FakeTensorMeta:
-    tensor_size: tuple[Union[int, torch.SymInt], ...]
-    tensor_stride: tuple[Union[int, torch.SymInt], ...]
-    tensor_storage_offset: Union[int, torch.SymInt]
+    tensor_size: tuple[int | torch.SymInt, ...]
+    tensor_stride: tuple[int | torch.SymInt, ...]
+    tensor_storage_offset: int | torch.SymInt
     is_nested: bool
 
-    def size(self) -> tuple[Union[int, torch.SymInt], ...]:
+    def size(self) -> tuple[int | torch.SymInt, ...]:
         return self.tensor_size
 
-    def stride(self) -> tuple[Union[int, torch.SymInt], ...]:
+    def stride(self) -> tuple[int | torch.SymInt, ...]:
         return self.tensor_stride
 
-    def storage_offset(self) -> Union[int, torch.SymInt]:
+    def storage_offset(self) -> int | torch.SymInt:
         return self.tensor_storage_offset
 
     def dim(self) -> int:
         return len(self.tensor_size)
 
     @staticmethod
-    def from_fake(fake) -> "FakeTensorMeta":
+    def from_fake(fake: torch.Tensor) -> FakeTensorMeta:
         return FakeTensorMeta(
             fake.size(), fake.stride(), fake.storage_offset(), fake.is_nested
         )
@@ -446,7 +458,12 @@ def from_fake(fake) -> "FakeTensorMeta":
 
 # Checks whether the state of two ShapeEnv are equal w.r.t. the guards
 # returned by ShapeEnv.produce_guards.
-def shape_env_check_state_equal(env1, env2, non_state_variable_names, map_value):
+def shape_env_check_state_equal(
+    env1: ShapeEnv,
+    env2: ShapeEnv,
+    non_state_variable_names: tuple[str, ...],
+    map_value: Callable[[str, object], object],
+) -> None:
     # Collect and remove variables that don't necessarily represent the state
     # of a ShapeEnv. Note: we copy the dictionary so that we don't modify the
     # instance itself.
@@ -476,7 +493,7 @@ def value_to_str(value: Any) -> str:
     # Here, we allow the value of each field to be mapped, so that we appropriately
     # compare the two values.
     def compare_vars(
-        map_value: Callable[[str, Any], Any],
+        map_value: Callable[[str, object], object],
     ) -> list[tuple[str, str, str]]:
         env1_set, env2_set = set(env1_vars), set(env2_vars)
 
diff --git a/torch/fx/experimental/rewriter.py b/torch/fx/experimental/rewriter.py
index 327b6e9ff0c03..73e9abaa8e30f 100644
--- a/torch/fx/experimental/rewriter.py
+++ b/torch/fx/experimental/rewriter.py
@@ -7,7 +7,7 @@
 import textwrap
 from collections.abc import Callable
 from types import FunctionType
-from typing import Any, cast, Optional, Union
+from typing import Any, cast
 
 import torch
 from torch._sources import normalize_source_lines
@@ -115,13 +115,13 @@ def visit_AnnAssign(self, node):
 class RewritingTracer(Tracer):
     def trace(
         self,
-        root: Union[torch.nn.Module, Callable],
-        concrete_args: Optional[dict[str, Any]] = None,
+        root: torch.nn.Module | Callable,
+        concrete_args: dict[str, Any] | None = None,
     ) -> Graph:
         return super().trace(_rewrite(root), concrete_args)
 
 
-def _rewrite(fn: Union[torch.nn.Module, Callable]) -> Union[torch.nn.Module, Callable]:
+def _rewrite(fn: torch.nn.Module | Callable) -> torch.nn.Module | Callable:
     if isinstance(fn, torch.nn.Module):
         # Rewrite this module's `forward` as well as the `forward`s of
         # all of this module's recursive descendents. Return the new,
diff --git a/torch/fx/experimental/schema_type_annotation.py b/torch/fx/experimental/schema_type_annotation.py
index b37317ac462e3..9377f0bafdfcd 100644
--- a/torch/fx/experimental/schema_type_annotation.py
+++ b/torch/fx/experimental/schema_type_annotation.py
@@ -1,15 +1,21 @@
 # mypy: allow-untyped-defs
+from __future__ import annotations
+
 import inspect
-from typing import Any, Optional
+from typing import Any, TYPE_CHECKING
 
 import torch
 import torch.fx
 from torch._jit_internal import boolean_dispatched
 from torch.fx import Transformer
-from torch.fx.node import Argument, Target
 from torch.fx.operator_schemas import _torchscript_type_to_python_type
 
 
+if TYPE_CHECKING:
+    from torch.fx.graph_module import GraphModule
+    from torch.fx.node import Argument, Target
+
+
 class AnnotateTypesWithSchema(Transformer):
     """
     Use Python function signatures to annotate types for `Nodes` within an FX graph.
@@ -31,7 +37,7 @@ class AnnotateTypesWithSchema(Transformer):
 
     def __init__(
         self,
-        module: torch.nn.Module,
+        module: GraphModule,
         annotate_functionals: bool = True,
         annotate_modules: bool = True,
         annotate_get_attrs: bool = True,
@@ -121,7 +127,7 @@ def get_attr(
 
         return attr_proxy
 
-    def _extract_python_return_type(self, target: Target) -> Optional[Any]:
+    def _extract_python_return_type(self, target: Target) -> Any | None:
         """
         Given a Python call target, try to extract the Python return annotation
         if it is available, otherwise return None
diff --git a/torch/fx/experimental/shape_inference/infer_symbol_values.py b/torch/fx/experimental/shape_inference/infer_symbol_values.py
index 1ee6d2a939ae2..646da82d12678 100644
--- a/torch/fx/experimental/shape_inference/infer_symbol_values.py
+++ b/torch/fx/experimental/shape_inference/infer_symbol_values.py
@@ -1,6 +1,6 @@
 import re
 from collections import defaultdict
-from typing import Any, Union
+from typing import Any
 
 import numpy as np
 import sympy as sp
@@ -14,10 +14,10 @@
 
 
 def infer_symbol_values(
-    symints: list[Union[torch.SymInt, int]],
-    init_symints: list[Union[torch.SymInt, int]],
+    symints: list[torch.SymInt | int],
+    init_symints: list[torch.SymInt | int],
     symbol_idx_dict: dict[str, int],
-    padding_constraints: defaultdict[torch.SymInt, list[Union[sp.Expr, int]]],
+    padding_constraints: defaultdict[torch.SymInt, list[sp.Expr | int]],
     constraint: str,
 ) -> None:
     if constraint.find("non-singleton") != -1:
@@ -82,9 +82,9 @@ def infer_symbol_values(
 
 
 def calculate_value(
-    left_expression: Union[str, Any, None],
-    right_expression: Union[str, Any, None],
-    symints: list[Union[torch.SymInt, int]],
+    left_expression: str | Any | None,
+    right_expression: str | Any | None,
+    symints: list[torch.SymInt | int],
     symbol_idx_dict: dict[str, int],
 ) -> None:
     var, val = solve_equation(left_expression, right_expression)
@@ -94,8 +94,8 @@ def calculate_value(
 
 
 def solve_equation(
-    left_expression: Union[str, Any, None],
-    right_expression: Union[str, Any, None],
+    left_expression: str | Any | None,
+    right_expression: str | Any | None,
 ) -> tuple[str, int]:
     expression = f"{left_expression} - {right_expression}"
     var = re.findall(s_pattern, expression)[0]
@@ -117,9 +117,9 @@ def solve_equation(
 
 
 def update_equation(
-    symints: list[Union[torch.SymInt, int]],
-    init_symints: list[Union[torch.SymInt, int]],
-    padding_constraints: defaultdict[torch.SymInt, list[Union[sp.Expr, int]]],
+    symints: list[torch.SymInt | int],
+    init_symints: list[torch.SymInt | int],
+    padding_constraints: defaultdict[torch.SymInt, list[sp.Expr | int]],
     init_eq: sp.Expr,
     new_mod_num: int,
     var: torch.SymInt,
diff --git a/torch/fx/experimental/sym_node.py b/torch/fx/experimental/sym_node.py
index aafef3d2333d8..a175c92df0526 100644
--- a/torch/fx/experimental/sym_node.py
+++ b/torch/fx/experimental/sym_node.py
@@ -1,5 +1,3 @@
-# mypy: allow-untyped-defs
-
 from __future__ import annotations
 
 
@@ -23,7 +21,7 @@
 import operator
 import sys
 from functools import lru_cache, update_wrapper
-from typing import Optional, TYPE_CHECKING, Union
+from typing import Any, overload, TYPE_CHECKING
 
 import torch
 import torch._logging.structured as structured
@@ -43,6 +41,11 @@
 
 
 if TYPE_CHECKING:
+    from collections.abc import Callable
+    from typing import Self
+
+    import sympy
+
     from torch.fx.experimental.symbolic_shapes import ShapeEnv
 
 log = logging.getLogger(__name__)
@@ -62,7 +65,15 @@
 from torch.types import py_sym_types as SymTypes
 
 
-def _to_symtype(t):
+@overload
+def _to_symtype(t: type[bool]) -> type[SymBool]: ...
+@overload
+def _to_symtype(t: type[int]) -> type[SymInt]: ...
+@overload
+def _to_symtype(t: type[float]) -> type[SymFloat]: ...
+@overload
+def _to_symtype(t: type) -> type: ...
+def _to_symtype(t: type) -> type:
     if t is bool:
         return SymBool
     if t is int:
@@ -92,18 +103,20 @@ class SymNode:
 
     def __init__(
         self,
-        expr,
-        shape_env,
-        pytype,
-        hint: Optional[Union[int, float, bool]],
-        constant=None,
-        fx_node=None,
-        optimized_summation=False,
-    ):
+        expr: object,
+        shape_env: ShapeEnv | None,
+        pytype: type,
+        hint: HintType | object,
+        constant: int | float | bool | None = None,
+        fx_node: object = None,
+        optimized_summation: bool = False,
+    ) -> None:
         self._expr = expr
         self.shape_env = shape_env
         self.pytype = pytype
         self._optimized_summation = optimized_summation
+        self._expr_ver = -1
+        self._expr_cache = None
 
         # What's the difference between hint and constant?
         #
@@ -132,7 +145,7 @@ def __init__(
         # potential refinements to unbacked symints this got harder to keep
         # in sync, so we've deleted it for now.)
 
-        def compute_hint():
+        def compute_hint() -> HintType | SymInt | SymFloat | SymBool:
             from torch.fx.experimental.symbolic_shapes import has_free_unbacked_symbols
 
             # This occasionally gets exercised by, e.g.,
@@ -142,6 +155,8 @@ def compute_hint():
             # expensive.
             if has_free_unbacked_symbols(self.expr):
                 return None
+            if self.shape_env is None:
+                raise RuntimeError("shape_env is required to compute hint")
             hint = self.shape_env._maybe_evaluate_static(self.expr, compute_hint=True)
             if hint is not None:
                 hint = self.pytype(hint) if not isinstance(hint, SymTypes) else hint
@@ -165,7 +180,7 @@ def compute_hint():
         else:
             hint = compute_hint()
         self._hint = hint
-        self.constant: Optional[Union[int, float, bool]] = constant
+        self.constant: int | float | bool | None = constant
 
         # Record the FX node of the current node if we are doing translation
         # validation. They will be used for building the input assertions for
@@ -195,50 +210,39 @@ def _value_hash(self) -> int:
         return hash((self._expr, self.pytype, self._hint, self.constant, self.fx_node))
 
     @property
-    def expr(self):
-        return self.shape_env.replace(self._expr)
+    def expr(self) -> sympy.Basic:
+        if (
+            isinstance(self._expr, int)
+            or self._expr.is_number  # pyrefly: ignore[missing-attribute]
+        ):
+            return self._expr
+        if self.shape_env is None:
+            raise AssertionError("shape_env is required to access expr")
+        ver = self.shape_env._replacements_version_counter
+        if ver == 0:
+            return self._expr
+        if self._expr_cache is not None and ver == self._expr_ver:
+            return self._expr_cache
+        result = self.shape_env.replace(self._expr)
+        self._expr_ver = ver
+        self._expr_cache = result
+        return result
 
     @property
-    def hint(self):
+    def hint(self) -> HintType | SymInt | SymFloat | SymBool:
         return self._hint
 
-    def has_hint(self):
+    def has_hint(self) -> bool:
         return self._hint is not None
 
-    def require_hint(self, fallback=None):
-        from torch.fx.experimental.symbolic_shapes import free_unbacked_symbols
-
-        if self._hint is None:
-            if fallback is not None:
-                # Say we have some expr like 2*u0 + s0
-                # The hint will be None, since the expr contains at least 1 unbacked.
-                # We will:
-                # - replace every backed free symbol with its corresponding hint
-                # - replace every unbacked free symbol with the fallback
-                # - regenerate the expression with those symbol replacements
-                # Note: this is not really complete either, since right now
-                # this logic does not take into account any value ranges
-                # for the unbacked symints, we may need to beef it up at some point.
-                unbacked_symbols = free_unbacked_symbols(self.expr)
-                replacements = {
-                    s: fallback
-                    if s in unbacked_symbols
-                    else self.shape_env.backed_var_to_val[s]
-                    for s in self.expr.free_symbols
-                }
-                return int(self.expr.xreplace(replacements))
-            # NB: we expect this to raise
-            return self.shape_env.size_hint(self.expr)
-        return self._hint
-
-    def maybe_as_int(self):
+    def maybe_as_int(self) -> int | None:
         if self.expr.is_number:
             return int(self.expr)
         else:
             return None
 
     # NB: This does conversions, not sure if this is good or not
-    def maybe_as_float(self):
+    def maybe_as_float(self) -> float | None:
         import sympy
 
         if isinstance(self.expr, sympy.Float):
@@ -246,7 +250,7 @@ def maybe_as_float(self):
         else:
             return None
 
-    def maybe_as_bool(self):
+    def maybe_as_bool(self) -> bool | None:
         import sympy
 
         if self.expr is sympy.true:
@@ -256,16 +260,16 @@ def maybe_as_bool(self):
         else:
             return None
 
-    def is_int(self):
+    def is_int(self) -> bool:
         return self.pytype is int
 
-    def is_float(self):
+    def is_float(self) -> bool:
         return self.pytype is float
 
-    def is_bool(self):
+    def is_bool(self) -> bool:
         return self.pytype is bool
 
-    def is_nested_int(self):
+    def is_nested_int(self) -> bool:
         # Unbacked SymInts cannot be nested int today
         return (
             self._hint is not None
@@ -273,7 +277,7 @@ def is_nested_int(self):
             and self._hint.node.is_nested_int()
         )
 
-    def wrap_int(self, num):
+    def wrap_int(self, num: int) -> SymNode:
         if type(num) is not int:
             raise AssertionError(f"Expected int, got {type(num)}")
         import sympy
@@ -282,7 +286,7 @@ def wrap_int(self, num):
             sympy.Integer(num), self.shape_env, int, num, constant=num, fx_node=num
         )
 
-    def wrap_float(self, num):
+    def wrap_float(self, num: float) -> SymNode:
         if type(num) is not float:
             raise AssertionError(f"Expected float, got {type(num)}")
         import sympy
@@ -291,7 +295,7 @@ def wrap_float(self, num):
             sympy.Float(num), self.shape_env, float, num, constant=num, fx_node=num
         )
 
-    def wrap_bool(self, num):
+    def wrap_bool(self, num: bool) -> SymNode:
         if type(num) is not bool:
             raise AssertionError(f"Expected bool, got {type(num)}")
         import sympy
@@ -305,16 +309,16 @@ def wrap_bool(self, num):
             fx_node=num,
         )
 
-    def clone(self):
+    def clone(self) -> SymNode:
         return self
 
-    def str(self):
+    def str(self) -> builtins.str:
         return f"{self.expr}"
 
-    def __str__(self):
+    def __str__(self) -> builtins.str:
         return self.str()
 
-    def __repr__(self):
+    def __repr__(self) -> builtins.str:
         rep = [
             f"SymNode({self._expr}, shape_env={self.shape_env}, pytype={self.pytype}",
         ]
@@ -338,70 +342,70 @@ def abs(self) -> SymNode:
     def pos(self) -> SymNode:
         return self._pos()  # type: ignore[attr-defined]
 
-    def round(self, ndigits=None) -> SymNode:
+    def round(self, ndigits: int | None = None) -> SymNode:
         return self._round(ndigits)  # type: ignore[attr-defined]
 
     def trunc(self) -> SymNode:
         return self._trunc()  # type: ignore[attr-defined]
 
-    def add(self, other) -> SymNode:
+    def add(self, other: SymNode) -> SymNode:
         return self._add(other)  # type: ignore[attr-defined]
 
-    def sub(self, other) -> SymNode:
+    def sub(self, other: SymNode) -> SymNode:
         return self._sub(other)  # type: ignore[attr-defined]
 
-    def mul(self, other) -> SymNode:
+    def mul(self, other: SymNode) -> SymNode:
         return self._mul(other)  # type: ignore[attr-defined]
 
-    def mod(self, other) -> SymNode:
+    def mod(self, other: SymNode) -> SymNode:
         return self._mod(other)  # type: ignore[attr-defined]
 
-    def float_pow(self, other) -> SymNode:
+    def float_pow(self, other: SymNode) -> SymNode:
         return self._float_pow(other)  # type: ignore[attr-defined]
 
-    def pow_by_natural(self, other) -> SymNode:
+    def pow_by_natural(self, other: SymNode) -> SymNode:
         return self._pow_by_natural(other)  # type: ignore[attr-defined]
 
-    def and_(self, other) -> SymNode:
+    def and_(self, other: SymNode) -> SymNode:
         return self._and_(other)  # type: ignore[attr-defined]
 
-    def or_(self, other) -> SymNode:
+    def or_(self, other: SymNode) -> SymNode:
         return self._or_(other)  # type: ignore[attr-defined]
 
-    def float_truediv(self, other) -> SymNode:
+    def float_truediv(self, other: SymNode) -> SymNode:
         return self._float_truediv(other)  # type: ignore[attr-defined]
 
-    def int_truediv(self, other) -> SymNode:
+    def int_truediv(self, other: SymNode) -> SymNode:
         return self._int_truediv(other)  # type: ignore[attr-defined]
 
-    def int_floordiv(self, other) -> SymNode:
+    def int_floordiv(self, other: SymNode) -> SymNode:
         return self._int_floordiv(other)  # type: ignore[attr-defined]
 
-    def lshift(self, other) -> SymNode:
+    def lshift(self, other: SymNode) -> SymNode:
         return self._lshift(other)  # type: ignore[attr-defined]
 
-    def rshift(self, other) -> SymNode:
+    def rshift(self, other: SymNode) -> SymNode:
         return self._rshift(other)  # type: ignore[attr-defined]
 
     def sym_not(self) -> SymNode:  # noqa: F811
         return self._sym_not()  # type: ignore[attr-defined]
 
-    def eq(self, other) -> SymNode:
+    def eq(self, other: SymNode) -> SymNode:
         return self._eq(other)  # type: ignore[attr-defined]
 
-    def ne(self, other) -> SymNode:
+    def ne(self, other: SymNode) -> SymNode:
         return self._ne(other)  # type: ignore[attr-defined]
 
-    def gt(self, other) -> SymNode:
+    def gt(self, other: SymNode) -> SymNode:
         return self._gt(other)  # type: ignore[attr-defined]
 
-    def lt(self, other) -> SymNode:
+    def lt(self, other: SymNode) -> SymNode:
         return self._lt(other)  # type: ignore[attr-defined]
 
-    def le(self, other) -> SymNode:
+    def le(self, other: SymNode) -> SymNode:
         return self._le(other)  # type: ignore[attr-defined]
 
-    def ge(self, other) -> SymNode:
+    def ge(self, other: SymNode) -> SymNode:
         return self._ge(other)  # type: ignore[attr-defined]
 
     def floor(self) -> SymNode:
@@ -422,74 +426,86 @@ def ceil(self) -> SymNode:
     def neg(self) -> SymNode:
         return self._neg()  # type: ignore[attr-defined]
 
-    def sym_min(self, other) -> SymNode:  # noqa: F811
+    def sym_min(self, other: SymNode) -> SymNode:  # noqa: F811
         return self._sym_min(other)  # type: ignore[attr-defined]
 
-    def sym_max(self, other) -> SymNode:  # noqa: F811
+    def sym_max(self, other: SymNode) -> SymNode:  # noqa: F811
         return self._sym_max(other)  # type: ignore[attr-defined]
 
-    def sym_ite(self, then_val, else_val) -> SymNode:
+    def sym_ite(self, then_val: SymNode, else_val: SymNode) -> SymNode:
         return self._sym_ite(then_val, else_val)  # type: ignore[attr-defined]
 
-    def is_contiguous(self, sizes, strides) -> SymNode:
+    def is_contiguous(self, sizes: list[SymNode], strides: list[SymNode]) -> SymNode:
         return self._is_contiguous(sizes, strides)  # type: ignore[attr-defined]
 
-    def is_channels_last_contiguous_2d(self, sizes, strides) -> SymNode:
+    def is_channels_last_contiguous_2d(
+        self, sizes: list[SymNode], strides: list[SymNode]
+    ) -> SymNode:
         return self._is_channels_last_contiguous_2d(sizes, strides)  # type: ignore[attr-defined]
 
-    def is_channels_last_contiguous_3d(self, sizes, strides) -> SymNode:
+    def is_channels_last_contiguous_3d(
+        self, sizes: list[SymNode], strides: list[SymNode]
+    ) -> SymNode:
         return self._is_channels_last_contiguous_3d(sizes, strides)  # type: ignore[attr-defined]
 
-    def is_channels_last_strides_2d(self, sizes, strides) -> SymNode:
+    def is_channels_last_strides_2d(
+        self, sizes: list[SymNode], strides: list[SymNode]
+    ) -> SymNode:
         return self._is_channels_last_strides_2d(sizes, strides)  # type: ignore[attr-defined]
 
-    def is_channels_last_strides_3d(self, sizes, strides) -> SymNode:
+    def is_channels_last_strides_3d(
+        self, sizes: list[SymNode], strides: list[SymNode]
+    ) -> SymNode:
         return self._is_channels_last_strides_3d(sizes, strides)  # type: ignore[attr-defined]
 
-    def is_non_overlapping_and_dense_indicator(self, sizes, strides) -> SymNode:
+    def is_non_overlapping_and_dense_indicator(
+        self, sizes: list[SymNode], strides: list[SymNode]
+    ) -> SymNode:
         return self._is_non_overlapping_and_dense_indicator(sizes, strides)  # type: ignore[attr-defined]
 
     # Make C++ happy
-    def sym_or(self, other):
+    def sym_or(self, other: SymNode) -> SymNode:
         return self.or_(other)
 
-    def sym_and(self, other):
+    def sym_and(self, other: SymNode) -> SymNode:
         return self.and_(other)
 
     # Integer bitwise ops
-    def bitwise_and(self, other):
+    def bitwise_and(self, other: SymNode) -> SymNode:
         return self._bitwise_and(other)  # type: ignore[attr-defined]
 
-    def bitwise_or(self, other):
+    def bitwise_or(self, other: SymNode) -> SymNode:
         return self._bitwise_or(other)  # type: ignore[attr-defined]
 
-    def bitwise_xor(self, other):
+    def bitwise_xor(self, other: SymNode) -> SymNode:
         return self._bitwise_xor(other)  # type: ignore[attr-defined]
 
     # There is no int_truediv available from C++
-    def truediv(self, other):
+    def truediv(self, other: SymNode) -> SymNode:
         return self.float_truediv(other)
 
-    def floordiv(self, other) -> SymNode:
+    def floordiv(self, other: SymNode) -> SymNode:
         return self.int_floordiv(other)
 
     # We didn't bind integer pow in C++
-    def pow(self, other):
+    def pow(self, other: SymNode) -> SymNode:
         return self.float_pow(other)
 
-    def is_non_overlapping_and_dense(self, sizes, strides):
+    def is_non_overlapping_and_dense(
+        self, sizes: list[SymNode], strides: list[SymNode]
+    ) -> SymNode:
         return self.is_non_overlapping_and_dense_indicator(sizes, strides).eq(
             to_node(self, 1)
         )  # type: ignore[attr-defined]
 
-    def int_(self):
+    def int_(self) -> int:
         return self.guard_int("", 0)  # NB: uses Python backtrace
 
     # This one is currently done by hand, but if we add other variadic
     # functions consider factoring it out to be metaprogrammed too.  Note that
     # some load bearing logic is directly in torch.sym_sum
 
-    def sym_sum(self, args) -> SymNode:
+    def sym_sum(self, args: list[SymNode]) -> SymNode:
         import sympy
 
         # Inner impl
@@ -511,14 +527,16 @@ def sym_sum(self, args) -> SymNode:
         out = sympy.Add(*exprs)
 
         size_hints = []
-        out_hint = None
+        out_hint: object = _NO_HINT
         for a in args:
             if a.hint is None:
                 break
             size_hints.append(a.hint)
         else:
-            out_hint = sum(size_hints)
+            out_hint = sum(size_hints)  # pyrefly: ignore[no-matching-overload]
 
+        if self.shape_env is None:
+            raise RuntimeError("shape_env is required for sym_sum")
         fx_node, _ = self.shape_env._create_fx_call_function(
             torch.sym_sum, (tuple(a.fx_node for a in args),)
         )
@@ -526,11 +544,13 @@ def sym_sum(self, args) -> SymNode:
         # NB: Only for integers!
         return SymNode(out, self.shape_env, int, out_hint, fx_node=fx_node)
 
-    def evaluate(self, size_oblivious=False):
+    def evaluate(self, size_oblivious: bool = False) -> bool | int | float:
+        if self.shape_env is None:
+            raise RuntimeError("shape_env is required to evaluate")
         return self.shape_env.evaluate_sym_node(self, size_oblivious)
 
     # You can manually trigger a guard with this function
-    def guard_int(self, file, line):
+    def guard_int(self, file: builtins.str, line: int) -> int:
         # TODO: use the file/line for some useful diagnostic on why a
         # guard occurred
         r = self.evaluate()
@@ -540,7 +560,7 @@ def guard_int(self, file, line):
             log.warning("Failed to convert to int: %s", r)
             raise
 
-    def guard_float(self, file, line):
+    def guard_float(self, file: builtins.str, line: int) -> float:
         # TODO: use the file/line for some useful diagnostic on why a
         # guard occurred
         r = self.evaluate()
@@ -550,7 +570,7 @@ def guard_float(self, file, line):
             log.warning("Failed to convert to float: %s", r)
             raise
 
-    def guard_bool(self, file, line):
+    def guard_bool(self, file: builtins.str, line: int) -> bool:
         # TODO: use the file/line for some useful diagnostic on why a
         # guard occurred
         r = self.evaluate()
@@ -560,9 +580,11 @@ def guard_bool(self, file, line):
             log.warning("Failed to convert to bool: %s", r)
             raise
 
-    def expect_true(self, file, line):
+    def expect_true(self, file: builtins.str, line: int) -> bool:
         from torch.fx.experimental.symbolic_shapes import free_unbacked_symbols
 
+        if self.shape_env is None:
+            raise RuntimeError("shape_env is required for expect_true")
         if (
             self.has_hint()
             and not free_unbacked_symbols(self.expr)
@@ -578,14 +600,14 @@ def expect_true(self, file, line):
             self.expr, f"{file}:{line}", fx_node=self.fx_node
         )
 
-    def statically_known_true(self, file, line):
+    def statically_known_true(self, file: builtins.str, line: int) -> bool:
         from torch.fx.experimental.symbolic_shapes import statically_known_true
 
         if not self.is_bool():
             raise AssertionError("Expected bool type")
         return statically_known_true(SymBool(self))
 
-    def guard_size_oblivious(self, file, line):
+    def guard_size_oblivious(self, file: builtins.str, line: int) -> bool:
         """
         Like guard_bool, but if we encounter unbacked symbols, if those symbols
         are size-like, we will treat them as >= 2 for the purposes of the analysis.
@@ -605,35 +627,35 @@ def guard_size_oblivious(self, file, line):
             log.warning("Failed to convert to bool: %s", r)
             raise
 
-    def guard_or_false(self, file, line):
+    def guard_or_false(self, file: builtins.str, line: int) -> bool:
         from torch.fx.experimental.symbolic_shapes import guard_or_false
 
         if not self.is_bool():
             raise AssertionError("Expected bool type")
         return guard_or_false(SymBool(self))
 
-    def guard_or_true(self, file, line):
+    def guard_or_true(self, file: builtins.str, line: int) -> bool:
         from torch.fx.experimental.symbolic_shapes import guard_or_true
 
         if not self.is_bool():
             raise AssertionError("Expected bool type")
         return guard_or_true(SymBool(self))
 
-    def bool_(self):
+    def bool_(self) -> bool:
         return self.guard_bool("", 0)
 
-    def is_symbolic(self):
+    def is_symbolic(self) -> bool:
         return True
 
-    def nested_int(self):
+    def nested_int(self) -> None:
         return None
 
-    def is_constant(self):
+    def is_constant(self) -> bool:
         return False
 
 
 class _DynamicScalar:
-    def __new__(cls, *args):
+    def __new__(cls, *args: object) -> Self:
         if cls is _DynamicScalar:
             raise TypeError("_DynamicScalar is an abstract base class, use DynamicInt.")
         return super().__new__(cls, *args)
@@ -651,21 +673,44 @@ class DynamicInt(_DynamicScalar, int):
         fn(x)  # compiles x as a dynamic integer input; returns f(4)
     """
 
-    def __new__(cls, val):
+    def __new__(cls, val: int) -> Self:
         if not isinstance(val, int):
             raise AssertionError(f"Expected int, got {type(val)}")
         obj = super().__new__(cls, int(val))
         return obj
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         return f"DynamicInt({self.real})"
 
-    def __floordiv__(self, other):  # // was casting to int without these overrides?
+    def __floordiv__(
+        self, other: int
+    ) -> DynamicInt:  # // was casting to int without these overrides?
         return DynamicInt(self.real // other)
 
-    def __rfloordiv__(self, other):
+    def __rfloordiv__(self, other: int) -> DynamicInt:
         return DynamicInt(other // self.real)
 
+    def __pow__(self, other, modulo=None):
+        if modulo is not None:
+            result = pow(self.real, other, modulo)
+        else:
+            result = self.real**other
+        # Only create DynamicInt if result is int, otherwise return plain value
+        # (e.g., negative exponent produces float)
+        if isinstance(result, int):
+            return DynamicInt(result)
+        return result
+
+    def __rpow__(self, other, modulo=None):
+        if modulo is not None:
+            result = pow(other, self.real, modulo)
+        else:
+            result = other**self.real
+        # Only create DynamicInt if result is int, otherwise return plain value
+        if isinstance(result, int):
+            return DynamicInt(result)
+        return result
+
 
 # TODO: this probably needs the sizes-strides eval functions
 METHOD_TO_OPERATOR = {
@@ -720,8 +765,8 @@ def __rfloordiv__(self, other):
 
 
 # Adding math ops: sqrt, cos, sin, ...
-def _get_sym_node_fn(name):
-    def fn(self):
+def _get_sym_node_fn(name: str) -> Callable[[SymNode], SymNode]:
+    def fn(self: SymNode) -> SymNode:
         return getattr(self, f"_sym_{name}")()
 
     return fn
@@ -746,6 +791,7 @@ def fn(self):
     setattr(SymNode, sym_name, _get_sym_node_fn(name))
     METHOD_TO_OPERATOR[sym_name] = getattr(torch, priv_sym_name)
     unary_magic_methods.add(sym_name)
+    # pyrefly: ignore [unresolvable-dunder-all]
     __all__.append(sym_name)
 
 
@@ -799,25 +845,25 @@ def fn(self):
 # Methods that have a `__foo__` as well as `__rfoo__`
 
 
-def _sympy_float_truediv(a, b):
+def _sympy_float_truediv(a: sympy.Basic, b: sympy.Basic) -> sympy.Basic:
     from torch.utils._sympy.functions import FloatTrueDiv
 
     return FloatTrueDiv(a, b)
 
 
-def _sympy_int_truediv(a, b):
+def _sympy_int_truediv(a: sympy.Basic, b: sympy.Basic) -> sympy.Basic:
     from torch.utils._sympy.functions import IntTrueDiv
 
     return IntTrueDiv(a, b)
 
 
-def _sympy_floordiv(a, b):
+def _sympy_floordiv(a: sympy.Basic, b: sympy.Basic) -> sympy.Basic:
     from torch.utils._sympy.functions import FloorDiv
 
     return FloorDiv(a, b)
 
 
-def _sympy_mod(a, b):
+def _sympy_mod(a: sympy.Basic, b: sympy.Basic) -> sympy.Basic:
     from torch.utils._sympy.functions import Mod, PythonMod
 
     if a.is_nonnegative and b.is_nonnegative:
@@ -826,43 +872,45 @@ def _sympy_mod(a, b):
         return PythonMod(a, b)
 
 
-def _sympy_pow_by_natural(a, b):
+def _sympy_pow_by_natural(a: sympy.Basic, b: sympy.Basic) -> sympy.Basic:
     from torch.utils._sympy.functions import PowByNatural
 
     return PowByNatural(a, b)
 
 
-def _sympy_float_pow(a, b):
+def _sympy_float_pow(a: sympy.Basic, b: sympy.Basic) -> sympy.Basic:
     from torch.utils._sympy.functions import FloatPow
 
     return FloatPow(a, b)
 
 
-def _sympy_and(a, b):
+def _sympy_and(a: sympy.Basic, b: sympy.Basic) -> sympy.Basic:
     import sympy
 
     return sympy.And(a, b)
 
 
-def _sympy_or(a, b):
+def _sympy_or(a: sympy.Basic, b: sympy.Basic) -> sympy.Basic:
     import sympy
 
     return sympy.Or(a, b)
 
 
-def _sympy_lshift(a, b):
+def _sympy_lshift(a: sympy.Basic, b: sympy.Basic) -> sympy.Basic:
     from torch.utils._sympy.functions import LShift
 
     return LShift(a, b)
 
 
-def _sympy_rshift(a, b):
+def _sympy_rshift(a: sympy.Basic, b: sympy.Basic) -> sympy.Basic:
     from torch.utils._sympy.functions import RShift
 
     return RShift(a, b)
 
 
-def _binary_search_insert_arg(ordered_args, new_arg):
+def _binary_search_insert_arg(
+    ordered_args: list[sympy.Basic], new_arg: sympy.Basic
+) -> list[sympy.Basic] | None:
     """
     If new_arg is found in ordered_args None is returned, else the new
     ordered_args with new_arg inserted
@@ -897,8 +945,11 @@ def _binary_search_insert_arg(ordered_args, new_arg):
 
 
 def _optimized_add(
-    lhs, rhs, lhs_is_optimized_summation=False, rhs_is_optimized_summation=False
-):
+    lhs: sympy.Basic,
+    rhs: sympy.Basic,
+    lhs_is_optimized_summation: bool = False,
+    rhs_is_optimized_summation: bool = False,
+) -> tuple[bool, sympy.Basic]:
     """
     Custom optimization for Add used to optimize incremental binary summations of certain properties. The idea
     is when we know the expression is a summation of unique symbols all we need to know is the correct order of symbols,
@@ -912,7 +963,7 @@ def _optimized_add(
     import sympy
     from sympy.core.basic import _args_sortkey as sortkey
 
-    def make_optimized(ordered_args):
+    def make_optimized(ordered_args: list[sympy.Basic]) -> tuple[bool, sympy.Basic]:
         if ordered_args is None:
             raise AssertionError("ordered_args is None")
         # Use _from_args directly to bypass _exec_constructor_postprocessors
@@ -964,19 +1015,19 @@ def make_optimized(ordered_args):
     return (_is_symbols_binary_summation(result), result)
 
 
-def _bitwise_and(a, b):
+def _bitwise_and(a: sympy.Basic, b: sympy.Basic) -> sympy.Basic:
     from torch.utils._sympy.functions import BitwiseFn_bitwise_and
 
     return BitwiseFn_bitwise_and(a, b)
 
 
-def _bitwise_or(a, b):
+def _bitwise_or(a: sympy.Basic, b: sympy.Basic) -> sympy.Basic:
     from torch.utils._sympy.functions import BitwiseFn_bitwise_or
 
     return BitwiseFn_bitwise_or(a, b)
 
 
-def _bitwise_xor(a, b):
+def _bitwise_xor(a: sympy.Basic, b: sympy.Basic) -> sympy.Basic:
     from torch.utils._sympy.functions import BitwiseFn_bitwise_xor
 
     return BitwiseFn_bitwise_xor(a, b)
@@ -1002,7 +1053,7 @@ def _bitwise_xor(a, b):
 }
 
 
-def _floor_ceil_helper(a, fn):
+def _floor_ceil_helper(a: sympy.Basic, fn: Callable[..., sympy.Basic]) -> sympy.Basic:
     import sympy
 
     if isinstance(a, sympy.Mul):
@@ -1020,7 +1071,7 @@ def _floor_ceil_helper(a, fn):
     return fn(a)
 
 
-def _sympy_floor(a):
+def _sympy_floor(a: sympy.Basic) -> sympy.Basic:
     from torch.utils._sympy.functions import FloorToInt
 
     return FloorToInt(a)
@@ -1028,67 +1079,67 @@ def _sympy_floor(a):
 
 # NB: this is Python trunc semantics which returns an int.  Do NOT use this to
 # represent torch.trunc (which is float to float)
-def _sympy_trunc(a):
+def _sympy_trunc(a: sympy.Basic) -> sympy.Basic:
     from torch.utils._sympy.functions import TruncToInt
 
     return TruncToInt(a)
 
 
-def _sympy_ceil(a):
+def _sympy_ceil(a: sympy.Basic) -> sympy.Basic:
     from torch.utils._sympy.functions import CeilToInt
 
     return CeilToInt(a)
 
 
-def _sympy_eq(a, b):
+def _sympy_eq(a: sympy.Basic, b: sympy.Basic) -> sympy.Basic:
     import sympy
 
     return sympy.Eq(a, b)
 
 
-def _sympy_ne(a, b):
+def _sympy_ne(a: sympy.Basic, b: sympy.Basic) -> sympy.Basic:
     import sympy
 
     return sympy.Ne(a, b)
 
 
-def _sympy_gt(a, b):
+def _sympy_gt(a: sympy.Basic, b: sympy.Basic) -> sympy.Basic:
     import sympy
 
     return sympy.Gt(a, b)
 
 
-def _sympy_lt(a, b):
+def _sympy_lt(a: sympy.Basic, b: sympy.Basic) -> sympy.Basic:
     import sympy
 
     return sympy.Lt(a, b)
 
 
-def _sympy_le(a, b):
+def _sympy_le(a: sympy.Basic, b: sympy.Basic) -> sympy.Basic:
     import sympy
 
     return sympy.Le(a, b)
 
 
-def _sympy_ge(a, b):
+def _sympy_ge(a: sympy.Basic, b: sympy.Basic) -> sympy.Basic:
     import sympy
 
     return sympy.Ge(a, b)
 
 
-def _sympy_min(a, b):
+def _sympy_min(a: sympy.Basic, b: sympy.Basic) -> sympy.Basic:
     from torch.utils._sympy.functions import Min
 
     return Min(a, b)
 
 
-def _sympy_max(a, b):
+def _sympy_max(a: sympy.Basic, b: sympy.Basic) -> sympy.Basic:
     from torch.utils._sympy.functions import Max
 
     return Max(a, b)
 
 
-def _sympy_ite(a, t, f):
+def _sympy_ite(a: sympy.Basic, t: sympy.Basic, f: sympy.Basic) -> sympy.Basic:
     import sympy
 
     return sympy.Piecewise((t, a), (f, True))
@@ -1097,8 +1148,8 @@ def _sympy_ite(a, t, f):
 current_module = sys.modules[__name__]
 
 
-def _get_sym_math_fn(name):
-    def fn(a):
+def _get_sym_math_fn(name: str) -> Callable[[sympy.Basic], sympy.Basic]:
+    def fn(a: sympy.Basic) -> sympy.Basic:
         import torch.utils._sympy.functions
 
         return getattr(torch.utils._sympy.functions, f"OpaqueUnaryFn_{name}")(a)
@@ -1115,13 +1166,15 @@ def fn(a):
 del fn, name, priv_sympy_name  # type: ignore[possibly-undefined]
 
 
-def _sympy_abs(a):
+def _sympy_abs(a: sympy.Basic) -> sympy.Basic:
     import sympy
 
     return sympy.Abs(a)
 
 
-def _sympy_round(number, ndigits=None):
+def _sympy_round(
+    number: sympy.Basic, ndigits: sympy.Basic | None = None
+) -> sympy.Basic:
     from torch.utils._sympy.functions import RoundDecimal, RoundToInt
 
     if ndigits is None:
@@ -1130,7 +1183,7 @@ def _sympy_round(number, ndigits=None):
         return RoundDecimal(number, ndigits)
 
 
-def _sympy_sym_float(a):
+def _sympy_sym_float(a: sympy.Basic) -> sympy.Basic:
     from torch.utils._sympy.functions import ToFloat
 
     # NB: Cannot use a * 1.0 here, because 0 * 1.0 is 0 which incorrectly
@@ -1138,7 +1191,7 @@ def _sympy_sym_float(a):
     return ToFloat(a)
 
 
-def _sympy_is_integer(a):
+def _sympy_is_integer(a: sympy.Basic) -> sympy.Basic:
     import sympy
 
     from torch.utils._sympy.functions import ToFloat
@@ -1177,12 +1230,16 @@ def _sympy_is_integer(a):
 del name, sym_name, math_op_names, current_module  # type: ignore[possibly-undefined]
 
 
-def sympy_is_contiguous(sizes, strides):
+def sympy_is_contiguous(
+    sizes: list[sympy.Basic], strides: list[sympy.Basic]
+) -> sympy.Basic:
     dim = len(sizes)
     return sympy_is_contiguous_generic(sizes, strides, list(range(dim - 1, -1, -1)))
 
 
-def sympy_is_contiguous_generic(sizes, strides, dim_order):
+def sympy_is_contiguous_generic(
+    sizes: list[sympy.Basic], strides: list[sympy.Basic], dim_order: list[int]
+) -> sympy.Basic:
     import sympy
 
     dim = len(sizes)
@@ -1206,15 +1263,21 @@ def sympy_is_contiguous_generic(sizes, strides, dim_order):
 # happens you will need to refactor this
 
 
-def sympy_is_channels_last_contiguous_2d(sizes, strides):
+def sympy_is_channels_last_contiguous_2d(
+    sizes: list[sympy.Basic], strides: list[sympy.Basic]
+) -> sympy.Basic:
     return sympy_is_contiguous_generic(sizes, strides, [1, 3, 2, 0])
 
 
-def sympy_is_channels_last_contiguous_3d(sizes, strides):
+def sympy_is_channels_last_contiguous_3d(
+    sizes: list[sympy.Basic], strides: list[sympy.Basic]
+) -> sympy.Basic:
     return sympy_is_contiguous_generic(sizes, strides, [1, 4, 3, 2, 0])
 
 
-def sympy_is_channels_last_strides_generic(sizes, strides, dim_order):
+def sympy_is_channels_last_strides_generic(
+    sizes: list[sympy.Basic], strides: list[sympy.Basic], dim_order: list[int]
+) -> sympy.Basic:
     import sympy
 
     from torch.utils._sympy.functions import Max
@@ -1254,15 +1317,21 @@ def sympy_is_channels_last_strides_generic(sizes, strides, dim_order):
     return r
 
 
-def sympy_is_channels_last_strides_2d(sizes, strides):
+def sympy_is_channels_last_strides_2d(
+    sizes: list[sympy.Basic], strides: list[sympy.Basic]
+) -> sympy.Basic:
     return sympy_is_channels_last_strides_generic(sizes, strides, [1, 3, 2, 0])
 
 
-def sympy_is_channels_last_strides_3d(sizes, strides):
+def sympy_is_channels_last_strides_3d(
+    sizes: list[sympy.Basic], strides: list[sympy.Basic]
+) -> sympy.Basic:
     return sympy_is_channels_last_strides_generic(sizes, strides, [1, 4, 3, 2, 0])
 
 
-def _sympy_is_non_overlapping_and_dense_indicator(sizes, strides):
+def _sympy_is_non_overlapping_and_dense_indicator(
+    sizes: list[sympy.Basic], strides: list[sympy.Basic]
+) -> sympy.Basic:
     from torch.utils._sympy.functions import IsNonOverlappingAndDenseIndicator
 
     return IsNonOverlappingAndDenseIndicator(*sizes, *strides)
@@ -1280,7 +1349,7 @@ def _sympy_is_non_overlapping_and_dense_indicator(sizes, strides):
 }
 
 
-def to_node(self, num):
+def to_node(self: SymNode, num: object) -> SymNode:
     if isinstance(num, SymTypes):
         return num.node
     elif type(num) is bool:
@@ -1295,7 +1364,7 @@ def to_node(self, num):
         return NotImplemented
 
 
-def wrap_node(x):
+def wrap_node(x: SymNode) -> SymInt | SymFloat | SymBool | int | float | bool:
     # TODO: let C++ also take advantage of this
     if isinstance(x, SymNode) and x.constant is not None:
         return x.constant
@@ -1309,11 +1378,11 @@ def wrap_node(x):
         raise AssertionError(f"unrecognized return type {x}")
 
 
-def method_to_operator(method):
+def method_to_operator(method: str) -> Callable[..., object]:
     return METHOD_TO_OPERATOR[method]
 
 
-def _make_node_magic(method, func):
+def _make_node_magic(method: str, func: Callable[..., sympy.Basic]) -> None:
     func = lru_cache(256)(func)
 
     if method in magic_methods_on_operator_with_trailing_underscore:
@@ -1338,9 +1407,9 @@ def uninteresting_files() -> set[str]:
             | {"<string>"}
         )
 
-    def capture_provenance(fn):
+    def capture_provenance(fn: Callable[..., SymNode]) -> Callable[..., SymNode]:
         @functools.wraps(fn)
-        def wrapper(self, other=None):
+        def wrapper(self: SymNode, other: SymNode | None = None) -> SymNode:
             if other is None:
                 result = fn(self)
             else:
@@ -1351,7 +1420,7 @@ def wrapper(self, other=None):
                 else:
                     arguments = [self]
 
-                def get_id(sym_node) -> Optional[int]:
+                def get_id(sym_node: SymNode) -> int | None:
                     # We don't want to return an ID if the input is a constant
                     import sympy
 
@@ -1385,7 +1454,7 @@ def get_id(sym_node) -> Optional[int]:
         return wrapper
 
     @capture_provenance
-    def binary_magic_impl(self, other):
+    def binary_magic_impl(self: SymNode, other: SymNode) -> SymNode:
         from torch.fx.experimental.proxy_tensor import (
             get_proxy_mode,
             handle_sym_dispatch,
@@ -1411,6 +1480,8 @@ def binary_magic_impl(self, other):
                 # Special handling for mod that requires access to the value
                 # ranges
                 shape_env = self.shape_env
+                if shape_env is None:
+                    raise AssertionError("shape_env is required for mod")
                 if (
                     self.expr.is_nonnegative
                     or shape_env.bound_sympy(self.expr).lower >= 0
@@ -1494,6 +1565,8 @@ def binary_magic_impl(self, other):
 
         # Create a FX node that corresponds to the operation being applied to
         # this node.
+        if self.shape_env is None:
+            raise RuntimeError("shape_env is required for binary op")
         fx_node, _ = self.shape_env._create_fx_call_function(
             op, (self.fx_node, other.fx_node)
         )
@@ -1509,7 +1582,7 @@ def binary_magic_impl(self, other):
         return result
 
     @capture_provenance
-    def unary_magic_impl(self):
+    def unary_magic_impl(self: SymNode) -> SymNode:
         from torch.fx.experimental.proxy_tensor import (
             get_proxy_mode,
             handle_sym_dispatch,
@@ -1520,6 +1593,8 @@ def unary_magic_impl(self):
             return to_node(self, handle_sym_dispatch(op, (wrap_node(self),), {}))
         # TODO: consider constant prop here
         expr = self.expr
+        if self.shape_env is None:
+            raise RuntimeError("shape_env is required for unary op")
         if method == "floor" or method == "ceiling":
             expr = self.shape_env._simplify_floor_div(expr)
 
@@ -1549,13 +1624,20 @@ def unary_magic_impl(self):
         setattr(SymNode, f"_{method_attr}", unary_magic_impl)
     elif method == "sym_ite":
 
-        def sym_ite_impl(pred_node, then_node, else_node):
+        def sym_ite_impl(
+            pred_node: SymNode, then_node: SymNode, else_node: SymNode
+        ) -> SymNode:
             from torch.fx.experimental.proxy_tensor import (
                 get_proxy_mode,
                 handle_sym_dispatch,
             )
 
-            out_hint = then_node.hint if pred_node.hint else else_node.hint
+            if pred_node.hint is None:
+                out_hint = None
+            elif pred_node.hint:
+                out_hint = then_node.hint
+            else:
+                out_hint = else_node.hint
             if get_proxy_mode():
                 return to_node(
                     pred_node,
@@ -1582,6 +1664,8 @@ def sym_ite_impl(pred_node, then_node, else_node):
                 )
                 raise
 
+            if pred_node.shape_env is None:
+                raise RuntimeError("shape_env is required for sym_ite")
             fx_node, _ = pred_node.shape_env._create_fx_call_function(
                 sym_ite, (pred_node.fx_node, then_node.fx_node, else_node.fx_node)
             )
@@ -1592,7 +1676,7 @@ def sym_ite_impl(pred_node, then_node, else_node):
         setattr(SymNode, f"_{method_attr}", sym_ite_impl)
     elif method == "round":
 
-        def round_impl(self, ndigits=None):
+        def round_impl(self: SymNode, ndigits: int | None = None) -> SymNode:
             from torch.fx.experimental.proxy_tensor import (
                 get_proxy_mode,
                 handle_sym_dispatch,
@@ -1618,7 +1702,9 @@ def round_impl(self, ndigits=None):
 
             out_hint = None
             if self.hint is not None:
-                out_hint = op(self.hint, ndigits)
+                out_hint = op(  # pyrefly: ignore[no-matching-overload]
+                    self.hint, ndigits
+                )
 
             # Internally, None is used as sentinel to indicate that a something is not a node on an FX graph. At the
             # same time, there is no way to wrap a plain None into an FX node. Thus, there is no way to pass None here
@@ -1630,6 +1716,8 @@ def round_impl(self, ndigits=None):
             args = [self.fx_node]
             if ndigits is not None:
                 args.append(ndigits)
+            if self.shape_env is None:
+                raise RuntimeError("shape_env is required for round")
             fx_node, _ = self.shape_env._create_fx_call_function(op, tuple(args))
             return SymNode(out, self.shape_env, pytype, out_hint, fx_node=fx_node)
 
@@ -1638,10 +1726,12 @@ def round_impl(self, ndigits=None):
         setattr(SymNode, f"_{method_attr}", binary_magic_impl)
 
 
-def _make_node_sizes_strides(method, func):
+def _make_node_sizes_strides(method: str, func: Callable[..., sympy.Basic]) -> None:
     # NB: don't LRU cache, lots of arguments
 
-    def sizes_strides_impl(self, sizes, strides):
+    def sizes_strides_impl(
+        self: SymNode, sizes: list[SymNode], strides: list[SymNode]
+    ) -> SymNode:
         from torch.fx.experimental.proxy_tensor import (
             get_proxy_mode,
             handle_sym_dispatch,
@@ -1694,7 +1784,9 @@ def sizes_strides_impl(self, sizes, strides):
     # TODO: This is technically hotpath, but in the ideal end state
     # guards on this will resolve at a higher level so you never
     # spend time in this code
-    def sizes_strides_user(sizes, strides):
+    def sizes_strides_user(
+        sizes: list[object], strides: list[object]
+    ) -> SymInt | SymFloat | SymBool | int | float | bool:
         import sympy
 
         from torch.fx.experimental.symbolic_shapes import (
@@ -1710,7 +1802,10 @@ def sizes_strides_user(sizes, strides):
                     )
                 )
         if method == "is_non_overlapping_and_dense_indicator":
-            return eval_is_non_overlapping_and_dense(sizes, strides)
+            return eval_is_non_overlapping_and_dense(
+                sizes,  # pyrefly: ignore[bad-argument-type]
+                strides,  # pyrefly: ignore[bad-argument-type]
+            )
         else:
             # TODO: this is an awful implementation
             return bool(
@@ -1732,7 +1827,7 @@ def sizes_strides_user(sizes, strides):
     _make_node_sizes_strides(method, func)
 
 
-def _make_user_magic(method, user_type):
+def _make_user_magic(method: str, user_type: type) -> None:
     # User magic takes care of wrapping the other operand into a node,
     # so that our internal logic can assume everything is nodes
     if method in magic_methods_on_operator_with_trailing_underscore:
@@ -1740,7 +1835,9 @@ def _make_user_magic(method, user_type):
     else:
         method_attr = method
 
-    def get_constant(x: Union[SymInt, int, SymFloat, float, SymBool, bool]):
+    def get_constant(
+        x: SymInt | int | SymFloat | float | SymBool | bool,
+    ) -> int | float | bool:
         if isinstance(x, (int, float, bool)):
             return x
         if isinstance(x, SymInt):
@@ -1749,7 +1846,7 @@ def get_constant(x: Union[SymInt, int, SymFloat, float, SymBool, bool]):
             return x.node.guard_bool("", 0)
         raise AssertionError("expect to be called with constant SymBools")
 
-    def is_constant(x):
+    def is_constant(x: SymInt | int | SymFloat | float | SymBool | bool) -> bool:
         if isinstance(x, (int, float, bool)):
             return True
         if isinstance(x, (SymInt, SymFloat, SymBool)):
@@ -1784,7 +1881,7 @@ def is_constant(x):
 
     if method in bool_becomes_int_magic_methods:
 
-        def promote(x):
+        def promote(x: object) -> Any:
             """Implements True+True=2, which works in python but not sympy"""
             if isinstance(x, SymBool):
                 return SymInt(x.node.wrap_int(int(x)))
@@ -1792,10 +1889,10 @@ def promote(x):
 
     else:
 
-        def promote(x):
+        def promote(x: object) -> Any:
             return x
 
-    def promote2(self, other):
+    def promote2(self: object, other: object) -> tuple[Any, Any]:
         # TODO: Remove eq and other relations from this list.
         # CPython has fancy implementations for these to get as much precision
         # as possible instead of just promoting to float64 and praying, so we
@@ -1837,13 +1934,13 @@ def promote2(self, other):
     # Alternatively, we could also rewrap into constant Symbool (i.e. by
     # implementing wrap_bool in ConstantSymNodeImpl), but we're not doing that
     # today for no particular reason.
-    def unary_magic_impl(self):
+    def unary_magic_impl(self: object) -> Any:
         self = promote(self)
         if is_constant(self):
             return (method_to_operator(method))(get_constant(self))
         return wrap_node(getattr(self.node, method_attr)())
 
-    def binary_magic_impl(self, other):
+    def binary_magic_impl(self: object, other: object) -> Any:
         if not isinstance(other, (int, float, bool, SymInt, SymFloat, SymBool)):
             return NotImplemented
         sym_node_log.debug("MAGIC %s %s %s", method, self, other)
@@ -1860,7 +1957,7 @@ def binary_magic_impl(self, other):
         ret = wrap_node(getattr(self.node, method_attr)(other_node))
         return get_constant(ret) if is_constant(ret) else ret
 
-    def rbinary_magic_impl(self, other):
+    def rbinary_magic_impl(self: object, other: object) -> Any:
         if not isinstance(other, (int, float, bool, SymInt, SymFloat, SymBool)):
             return NotImplemented
         self = promote(self)
@@ -1876,7 +1973,7 @@ def rbinary_magic_impl(self, other):
         ret = wrap_node(getattr(other_node, method_attr)(self.node))
         return get_constant(ret) if is_constant(ret) else ret
 
-    def setattrs(user_type, attr, symnode_impl):
+    def setattrs(user_type: type, attr: str, symnode_impl: object) -> None:
         """
         Registers the SymNode magic method on SymInt/Float/Bool,
         and optionally registers a corresponding wrapped method on DynamicInt.
@@ -1886,8 +1983,10 @@ def setattrs(user_type, attr, symnode_impl):
         setattr(user_type, attr, symnode_impl)
 
         # DynamicInt impl
-        def dynamic_int_impl(*args):
-            args = [x.real if isinstance(x, DynamicInt) else x for x in args]
+        def dynamic_int_impl(*args: object) -> Any:
+            args = [  # pyrefly: ignore[bad-assignment]
+                x.real if isinstance(x, DynamicInt) else x for x in args
+            ]
             out = getattr(int, attr)(*args)
             if isinstance(out, int) and not isinstance(out, bool):
                 return DynamicInt(out)
@@ -1903,7 +2002,9 @@ def dynamic_int_impl(*args):
         setattrs(user_type, method, update_wrapper(unary_magic_impl, orig))
     elif method == "sym_ite":
 
-        def sym_ite_magic_impl(pred, then_val, else_val):
+        def sym_ite_magic_impl(
+            pred: SymBool, then_val: object, else_val: object
+        ) -> Any:
             pred_node = pred.node
             then_node = to_node(pred_node, then_val)
             else_node = to_node(pred_node, else_val)
@@ -1918,12 +2019,17 @@ def sym_ite_magic_impl(pred, then_val, else_val):
                     "then_node and else_node must be SymNodes with same pytype"
                 )
             ret = wrap_node(getattr(pred.node, method_attr)(then_node, else_node))
-            return get_constant(ret) if ret.node.is_constant() else ret
+            return (
+                get_constant(ret)
+                # pyrefly: ignore[missing-attribute]
+                if ret.node.is_constant()
+                else ret
+            )
 
         setattrs(user_type, f"__{method}__", sym_ite_magic_impl)
     elif method == "round":
 
-        def round_magic_impl(self, ndigits=None):
+        def round_magic_impl(self: SymFloat, ndigits: int | None = None) -> Any:
             if is_constant(self):
                 return builtins.round(get_constant(self), ndigits)
 
diff --git a/torch/fx/experimental/symbolic_shapes.py b/torch/fx/experimental/symbolic_shapes.py
index 74e330a7e5b9c..d239bae2a5628 100644
--- a/torch/fx/experimental/symbolic_shapes.py
+++ b/torch/fx/experimental/symbolic_shapes.py
@@ -42,12 +42,10 @@
     Generic,
     NamedTuple,
     NoReturn,
-    Optional,
     TYPE_CHECKING,
     TypeAlias,
     TypeGuard,
     TypeVar,
-    Union,
 )
 from typing_extensions import deprecated, ParamSpec
 
@@ -63,9 +61,11 @@
 from torch._library.fake_class_registry import FakeScriptObject
 from torch._library.opaque_object import is_opaque_value
 from torch._logging import dtrace_structured, LazyString, structured, trace_structured
+from torch._opaque_base import OpaqueBase
 from torch._subclasses.meta_utils import is_sparse_any
 from torch._utils_internal import signpost_event
 from torch.fx.experimental import _config as config
+from torch.fx.experimental._config import AggressiveGuardFreeMode
 from torch.fx.experimental.recording import (
     FakeTensorMeta,
     record_shapeenv_event,
@@ -119,6 +119,54 @@
 log = logging.getLogger(__name__)
 
 
+from torch.fx.experimental._size_hinting import (
+    _guarding_hint_or_throw_base,
+    _optimization_hint_base,
+)
+
+
+def guarding_hint_or_throw(
+    a: torch.SymInt | torch.SymBool | int | bool | SymNode,
+) -> int | bool:
+    """
+    Return a concrete hint for a symbolic value, for use in guarding decisions.
+
+    Returns Python bool (True/False) for boolean inputs (SymBool, bool),
+    and Python int for integer inputs (SymInt, int).
+    """
+    if isinstance(a, SymNode):
+        if a._hint is not None:
+            return a._hint  # pyrefly: ignore[bad-return]
+        if a.shape_env is None:
+            raise AssertionError("shape_env is required for guarding_hint_or_throw")
+        hint = a.shape_env.guarding_hint_or_throw(a.expr)
+        a._hint = hint
+        return hint
+    if isinstance(a, (torch.SymInt, torch.SymBool)):
+        return guarding_hint_or_throw(a.node)
+    if isinstance(a, bool):
+        return a
+    if type(a) is not int:
+        raise AssertionError(f"Expected int, got {type(a)}")
+    return a
+
+
+def optimization_hint(a: torch.SymInt | int, fallback: int | None = None) -> int:
+    """
+    Return a concrete hint for a symbolic integer, for use in optimization decisions.
+
+    Unlike guarding_hint_or_throw, this function does not add guards and is intended
+    for optimization purposes only (e.g., memory estimation).
+    """
+    if isinstance(a, torch.SymInt):
+        if a.node._hint is not None:
+            return a.node._hint
+        return a.node.shape_env.optimization_hint(a.node.expr, fallback=fallback)
+    if type(a) is not int:
+        raise AssertionError(f"Expected int, got {type(a)}")
+    return a
+
+
 class GuardOnDataDependentSymNode(RuntimeError):
     cond: sympy.Basic
 
@@ -131,10 +179,15 @@ class PendingUnbackedSymbolNotFound(RuntimeError):
     pass
 
 
+class _ShapeEnvGuardError(RuntimeError):
+    """Raised when a guard is attempted while the ShapeEnv is in error-on-guard mode."""
+
+
 aten = torch._ops.ops.aten  # type: ignore[has-type]
 
 __all__ = [
-    "size_hint",
+    "optimization_hint",
+    "guarding_hint_or_throw",
     "guard_or_false",
     "guard_or_true",
     "has_symbolic_sizes_strides",
@@ -148,7 +201,6 @@ class PendingUnbackedSymbolNotFound(RuntimeError):
     "guard_float",
     "guard_scalar",
     "canonicalize_bool_expr",
-    "hint_int",
     "SYMPY_INTERP",
     "free_symbols",
     "is_symbol_binding_fx_node",
@@ -227,13 +279,13 @@ class SymIntEqByExpr:
     """
 
     @staticmethod
-    def _extract(val: Union[torch.SymInt, int]) -> sympy.Expr:
+    def _extract(val: torch.SymInt | int) -> sympy.Expr:
         if isinstance(val, torch.SymInt):
             return val.node.expr
         else:
             return sympy.Integer(val)
 
-    def __init__(self, val: Union[torch.SymInt, int]) -> None:
+    def __init__(self, val: torch.SymInt | int) -> None:
         self.val: sympy.Expr = SymIntEqByExpr._extract(val)
 
     def __repr__(self) -> str:
@@ -262,7 +314,7 @@ def _nested_int_aware_sort(
 
 # Wrapper on lru_cache that reports statistics at process end
 def lru_cache(
-    maxsize: Optional[int],
+    maxsize: int | None,
 ) -> Callable[[Callable[..., _T]], functools._lru_cache_wrapper[_T]]:
     def inner(f: Callable[..., _T]) -> functools._lru_cache_wrapper[_T]:
         wrapped_f = functools.lru_cache(maxsize)(f)
@@ -356,7 +408,7 @@ def has_symbolic_sizes_strides(elem: torch.Tensor) -> bool:
     return elem._has_symbolic_sizes_strides
 
 
-Int: TypeAlias = Union[torch.SymInt, int]
+Int: TypeAlias = torch.SymInt | int
 
 
 def create_contiguous(shape: Sequence[Int]) -> list[Int]:
@@ -366,29 +418,16 @@ def create_contiguous(shape: Sequence[Int]) -> list[Int]:
     return list(reversed(strides))
 
 
-@deprecated("used size_hint instead of hint_int", category=FutureWarning)
-def hint_int(a: Union[torch.SymInt, int], fallback: Optional[int] = None) -> int:
-    return size_hint(a, fallback)
+Scalar: TypeAlias = torch.SymInt | torch.SymFloat | torch.SymBool | int | float | bool
 
 
-def size_hint(a: Union[torch.SymInt, int], fallback: Optional[int] = None) -> int:
-    """
-    Retrieve the hint for an int (based on the underlying real values as observed
-    at runtime).  If no hint is available (e.g., because data dependent shapes),
-    if fallback is not None, use that instead to hint each unbacked symbol individually
-    (otherwise raise an error).
+def has_guarding_hint(a: Scalar) -> bool:
     """
-    if isinstance(a, torch.SymInt):
-        return a.node.require_hint(fallback)
-    if type(a) is not int:
-        raise AssertionError(f"Expected int, got {type(a)}")
-    return a
-
-
-Scalar: TypeAlias = Union[torch.SymInt, torch.SymFloat, torch.SymBool, int, float, bool]
-
+    Check if a symbolic value has a hint available for guarding.
 
-def has_hint(a: Scalar) -> bool:
+    Returns True if the value is concrete or if the symbolic node has a hint,
+    False otherwise.
+    """
     if isinstance(a, SymTypes):
         return a.node.has_hint()
     return True
@@ -458,7 +497,7 @@ def is_concrete_bool(a: BoolLikeType) -> bool:
     return False
 
 
-def has_static_value(a: Union[SymBool, SymFloat, SymInt, bool, float, int]) -> bool:
+def has_static_value(a: SymBool | SymFloat | SymInt | bool | float | int) -> bool:
     """
     User-code friendly utility to check if a value is static or dynamic.
     Returns true if given a constant, or a symbolic expression with a fixed value.
@@ -488,7 +527,7 @@ def has_static_value(a: Union[SymBool, SymFloat, SymInt, bool, float, int]) -> b
     potentially utilizing guard_or_false, guard_or_true, or statically_known_true",
     category=FutureWarning,
 )
-def guard_size_oblivious(expr: Union[torch.SymBool, bool]) -> bool:
+def guard_size_oblivious(expr: torch.SymBool | bool) -> bool:
     """
     Perform a guard on a symbolic boolean expression in a size oblivious way.
     This is typically used when a non-oblivious test would result in a guard
@@ -535,9 +574,9 @@ def check_consistent(new: _T, old: _T) -> None:
 
 
 def resolve_unbacked_bindings(
-    shape_env: Optional[ShapeEnv],
-    bindings: Optional[dict[sympy.Symbol, pytree.KeyPath]],
-) -> Optional[dict[sympy.Symbol, pytree.KeyPath]]:
+    shape_env: ShapeEnv | None,
+    bindings: dict[sympy.Symbol, pytree.KeyPath] | None,
+) -> dict[sympy.Symbol, pytree.KeyPath] | None:
     """
     When we do fake tensor prop, we oftentimes will allocate new unbacked symints.
     We then run proxy tensor mode, which populates node.meta["unbacked_bindings"]
@@ -553,11 +592,11 @@ def resolve_unbacked_bindings(
     return {shape_env.unbacked_renamings.get(k, k): v for k, v in bindings.items()}
 
 
-Result: TypeAlias = Union[torch.Tensor, tuple[torch.Tensor, ...]]
+Result: TypeAlias = torch.Tensor | tuple[torch.Tensor, ...]
 
 
 def rebind_unbacked(
-    shape_env: Optional[ShapeEnv], n: torch.fx.Node, result: Result
+    shape_env: ShapeEnv | None, n: torch.fx.Node, result: Result
 ) -> None:
     """
     Suppose we are retracing a pre-existing FX graph that previously had
@@ -757,10 +796,10 @@ def canonicalize_bool_expr(expr: _T) -> _T:
 
 
 def _sympy_from_args(
-    cls: type[Union[sympy.Add, sympy.Mul]],
+    cls: type[sympy.Add | sympy.Mul],
     args: list[sympy.Expr],
     sort: bool = True,
-    is_commutative: Optional[bool] = None,
+    is_commutative: bool | None = None,
 ) -> sympy.Expr:
     """
     Create a sympy expression from a list of arguments, optimizing for performance.
@@ -820,7 +859,7 @@ def _canonicalize_bool_expr_impl(expr: SympyBoolean) -> SympyBoolean:
         return type(expr)(*map(canonicalize_bool_expr, expr.args))
 
     opposite = {sympy.Gt: sympy.Lt, sympy.Ge: sympy.Le}
-    t: Union[type[Any]]
+    t: type[Any]
     if isinstance(expr, tuple(opposite.keys())):
         rhs = expr.lhs - expr.rhs  # type: ignore[attr-defined]
         t = opposite[type(expr)]  # type: ignore[index]
@@ -906,14 +945,14 @@ def div_by_factor(x: sympy.Expr, factor: int) -> sympy.Expr:
     return expr
 
 
-def is_nested_int(s: IntLikeType) -> TypeGuard[SymInt]:
+def is_nested_int(s: IntLikeType | FloatLikeType) -> TypeGuard[SymInt]:
     return isinstance(s, torch.SymInt) and s.node.is_nested_int()
 
 
-IterateExprsAtom: TypeAlias = Union[
-    SymInt, SymFloat, SymBool, int, float, bool, sympy.Basic, torch.Tensor
-]
-IterateExprs: TypeAlias = Union[IterateExprsAtom, Sequence[IterateExprsAtom]]
+IterateExprsAtom: TypeAlias = (
+    SymInt | SymFloat | SymBool | int | float | bool | sympy.Basic | torch.Tensor
+)
+IterateExprs: TypeAlias = IterateExprsAtom | Sequence[IterateExprsAtom]
 
 
 def _iterate_exprs(val: IterateExprs) -> Iterator[sympy.Basic]:
@@ -1063,7 +1102,7 @@ def _free_non_source_unbacked_symbols(
 
 # WARNING: Don't use this on Dynamo produced graphs, they don't have meta
 # setup!
-def is_symbol_binding_fx_node(node: torch.fx.Node) -> Optional[sympy.Symbol]:
+def is_symbol_binding_fx_node(node: torch.fx.Node) -> sympy.Symbol | None:
     """
     Check if a given FX node is a symbol binding node.
 
@@ -1105,7 +1144,7 @@ def find_symbol_binding_fx_nodes(
     Returns:
         A dictionary mapping from sympy Symbols to their binding FX nodes
     """
-    r = {}
+    r: dict[sympy.Symbol, torch.fx.Node] = {}
     # NB: Prefer first occurrence of symbol
     for node in graph.nodes:
         if (s := is_symbol_binding_fx_node(node)) is not None and s not in r:
@@ -1124,7 +1163,7 @@ class Specialization:
     """
 
     source: TensorPropertySource
-    check_fn: Callable
+    check_fn: Callable[[int], bool]
 
 
 # Analogous to ConvertIntSource
@@ -1177,9 +1216,9 @@ def get(self, o: int) -> int:
 def _free_unbacked_symbols_with_path(
     a: object,
     path: pytree.KeyPath,
-    real: Optional[object] = None,
-    shape_env: Optional[ShapeEnv] = None,
-    pending: Optional[set[sympy.Symbol]] = None,
+    real: object | None = None,
+    shape_env: ShapeEnv | None = None,
+    pending: set[sympy.Symbol] | None = None,
     simplify: bool = False,
 ) -> dict[sympy.Symbol, pytree.KeyPath]:
     """
@@ -1207,7 +1246,7 @@ def _free_unbacked_symbols_with_path(
         simplify=simplify,
     )
 
-    def expr(s: Union[SymInt, SymFloat, SymBool]) -> sympy.Expr:
+    def expr(s: SymInt | SymFloat | SymBool) -> sympy.Expr:
         if simplify:
             return s.node.expr
         # (When called from compute_unbacked_bindings)
@@ -1219,7 +1258,7 @@ def expr(s: Union[SymInt, SymFloat, SymBool]) -> sympy.Expr:
         pending = set()
     r = {}
 
-    def match_tensor(a: torch.Tensor, real_tensor: Optional[torch.Tensor] = None):
+    def match_tensor(a: torch.Tensor, real_tensor: torch.Tensor | None = None) -> None:
         r.update(
             go(
                 a.size(),
@@ -1369,11 +1408,11 @@ def _symint_wrap(s: sympy.Symbol) -> SymInt:
 
 
 def compute_unbacked_bindings(
-    shape_env: Optional[ShapeEnv],
+    shape_env: ShapeEnv | None,
     example_value: object,
-    old_example_value: Optional[object] = None,
+    old_example_value: object | None = None,
     peek: bool = False,
-) -> Optional[dict[sympy.Symbol, pytree.KeyPath]]:
+) -> dict[sympy.Symbol, pytree.KeyPath] | None:
     """
     After having run fake tensor propagation and producing example_value
     result, traverse example_value looking for freshly bound unbacked
@@ -1530,7 +1569,7 @@ def guard_or_true(a: BoolLikeType) -> bool:
     return _guard_or(a, True)
 
 
-def _static_eval_sym_bool(x: SymBool) -> Optional[bool]:
+def _static_eval_sym_bool(x: SymBool) -> bool | None:
     if not isinstance(x, SymBool):
         raise AssertionError(f"Expected SymBool, got {type(x)}")
     expr = x.node.expr
@@ -1633,8 +1672,8 @@ def sym_or(x: BoolLikeType, *others: BoolLikeType) -> BoolLikeType:
 
 
 def guard_scalar(
-    a: Union[SymBool, SymInt, SymFloat, int, bool, float],
-) -> Union[bool, int, float]:
+    a: SymBool | SymInt | SymFloat | int | bool | float,
+) -> bool | int | float:
     """
     Guard a scalar value, which can be a symbolic or concrete boolean, integer, or float.
 
@@ -1696,6 +1735,7 @@ def _advise_is_size(a: SymInt) -> None:
         isinstance(a, SymInt)
         and isinstance(a.node, SymNode)
         and isinstance(a.node.expr, sympy.Symbol)
+        and a.node.shape_env is not None
         and a.node.shape_env.is_unbacked_symint(a.node.expr)
     ):
         _constrain_range_for_size(a)
@@ -1706,6 +1746,7 @@ def _advise_is_bounded(a: SymInt, upper_bound: IntLikeType) -> None:
         isinstance(a, SymInt)
         and isinstance(a.node, SymNode)
         and isinstance(a.node.expr, sympy.Symbol)
+        and a.node.shape_env is not None
         and a.node.shape_env.is_unbacked_symint(a.node.expr)
         and isinstance(upper_bound, int)  # TODO: relax
     ):
@@ -1713,7 +1754,7 @@ def _advise_is_bounded(a: SymInt, upper_bound: IntLikeType) -> None:
 
 
 def _constrain_range_for_size(
-    a: SymInt, min: Optional[int] = None, max: Optional[int] = None
+    a: SymInt, min: int | None = None, max: int | None = None
 ) -> None:
     """
     This function is NOT INTENDED to be used by itself.
@@ -1731,9 +1772,7 @@ def _constrain_range_for_size(
 
 
 # inclusive both ways
-def constrain_range(
-    a: SymInt, *, min: Optional[int], max: Optional[int] = None
-) -> None:
+def constrain_range(a: SymInt, *, min: int | None, max: int | None = None) -> None:
     """
     Applies a constraint that the passed in SymInt must lie between min-max
     inclusive-inclusive, WITHOUT introducing a guard on the SymInt (meaning
@@ -1973,7 +2012,7 @@ class StrictMinMaxConstraint(Constraint):
     for N=0/1 too.
     """
 
-    vr: ValueRanges
+    vr: ValueRanges[sympy.Expr]
 
     def render(self, source: Source) -> str:
         """Format the constrain equation"""
@@ -2009,7 +2048,7 @@ def render(self, source: Source) -> str:
 # NB: None here indicates the client constraint is whatever is implicitly
 # inferred by guards from tracing, and that a backend can add whatever guards
 # it wants (including fully specializing the value).
-DimConstraint = Union[StrictMinMaxConstraint, RelaxedUnspecConstraint, None]
+DimConstraint = StrictMinMaxConstraint | RelaxedUnspecConstraint | None
 
 
 @dataclass(frozen=True, slots=True)
@@ -2034,7 +2073,7 @@ class EqualityConstraint(Constraint):
 
     source_pairs: list[tuple[Source, Source]]
     derived_equalities: list[
-        tuple[Source, Union[Source, sympy.Symbol], Callable[[sympy.Expr], sympy.Expr]]
+        tuple[Source, Source | sympy.Symbol, Callable[[sympy.Expr], sympy.Expr]]
     ]
     phantom_symbols: list[sympy.Symbol]
     relaxed_sources: set[Source]
@@ -2192,13 +2231,15 @@ class StatelessSymbolicContext(SymbolicContext, Generic[_P1, _T1]):
     dynamic_strides: DimList[DimDynamic] = None  # type: ignore[assignment]
     constraint_sizes: DimList[DimConstraint] = None  # type: ignore[assignment]
     constraint_strides: DimList[DimConstraint] = None  # type: ignore[assignment]
-    specialize_on: Optional[list[list[Callable[_P1, _T1]]]] = None
+    specialize_on: list[list[Callable[_P1, _T1]]] | None = None
     # If the tensor is a view, this should be populated for the base. It contains
     # information on how to allocate symbols when recursively fakeifying the base
     # during view fake-ification.
-    view_base_context: Optional[SymbolicContext] = None
+    view_base_context: SymbolicContext | None = None
     # Maps dimension index to shape_id.
-    shape_ids: Optional[dict[int, Optional[str]]] = None
+    shape_ids: dict[int, str | None] | None = None
+    # Maps dimension index to (min, max) bounds for unbacked dimensions.
+    unbacked_bounds: dict[int, tuple[int | None, int | None]] | None = None
     # TODO: add storage offset and stride symbolic_context
 
     def __post_init__(self) -> None:
@@ -2254,7 +2295,7 @@ def __post_init__(self) -> None:
 # is used.
 # TODO(voz): Shape env validation
 @dataclass(frozen=True, slots=True, kw_only=True)
-class StatefulSymbolicContext(StatelessSymbolicContext):
+class StatefulSymbolicContext(StatelessSymbolicContext[..., Any]):
     """
     Create symbols in ``create_symbolic_sizes_strides_storage_offset`` via
     a symbolic_context determination as given by a cache of Source:Symbol. A cache hit
@@ -2281,6 +2322,7 @@ class StatefulSymbolicContext(StatelessSymbolicContext):
     shape_env_to_source_to_symbol_cache: dict[int, dict[str, sympy.Expr]] = field(
         default_factory=dict
     )
+    excluded_sizes: tuple[int | None, ...] | None = None
 
 
 @dataclass(frozen=True, slots=True)
@@ -2315,8 +2357,8 @@ def __eq__(self, other: object) -> bool:
 
 
 def is_symbolic(
-    val: Union[int, SymInt, float, SymFloat, bool, SymBool],
-) -> TypeGuard[Union[SymInt, SymFloat, SymBool]]:
+    val: int | SymInt | float | SymFloat | bool | SymBool,
+) -> TypeGuard[SymInt | SymFloat | SymBool]:
     if isinstance(val, (int, float, bool)):
         return False
     return val.node.is_symbolic()
@@ -2342,7 +2384,8 @@ def _expandsums(args: list[sympy.Expr]) -> tuple[sympy.Expr, bool]:
         - A boolean indicating whether expansion occurred (True if multiple additive
           expressions were present or if there was at least one additive and one other expression)
     """
-    adds, other = [], []
+    adds: list[sympy.Expr] = []
+    other: list[sympy.Expr] = []
     for arg in args:
         if arg.is_Add:
             adds.append(arg)
@@ -2436,8 +2479,8 @@ def safe_expand(r: _SympyT) -> _SympyT:
 
 class _SymbolInfo(NamedTuple):
     k: sympy.Symbol
-    vr: Optional[ValueRanges]
-    val: Optional[sympy.Integer]
+    vr: ValueRanges[sympy.Expr] | None
+    val: sympy.Integer | None
     is_size_like: bool
 
 
@@ -2448,7 +2491,7 @@ def _maybe_evaluate_static_worker(
     symbol_info: tuple[_SymbolInfo, ...],
     unbacked_only: bool,
     size_oblivious: bool,
-) -> Optional[_SympyT]:
+) -> _SympyT | None:
     """
     This variant of ShapeEnv._maybe_evaluate_static has no dependence on
     ShapeEnv and thus can be cached indefinitely.  It does the "heavy" lifting
@@ -2601,8 +2644,8 @@ def _sympy_cast_symbool_to_symint_guardless(x: SympyBoolean) -> sympy.Expr:
 
 
 def cast_symbool_to_symint_guardless(
-    symbool: Union[bool, torch.SymBool],
-) -> Union[int, torch.SymInt]:
+    symbool: bool | torch.SymBool,
+) -> int | torch.SymInt:
     """
     Converts a SymBool or bool to a SymInt or int without introducing guards.
 
@@ -2621,12 +2664,21 @@ def cast_symbool_to_symint_guardless(
         return 1 if symbool else 0
     int_sym = _sympy_cast_symbool_to_symint_guardless(symbool.node.expr)
     return symbool.node.shape_env.create_symintnode(
-        int_sym, hint=int(symbool.node.require_hint()) if has_hint(symbool) else None
+        int_sym,
+        hint=guarding_hint_or_throw(symbool) if has_guarding_hint(symbool) else None,
     )
 
 
+def _eval_is_non_overlapping_and_dense_flat(*args: int) -> int:
+    # Guard code strings print IsNonOverlappingAndDenseIndicator with flat args
+    # (s0, s1, ..., stride0, stride1, ...) but eval_is_non_overlapping_and_dense
+    # expects two sequences (sizes, strides). This wrapper bridges the gap.
+    dim = len(args) // 2
+    return eval_is_non_overlapping_and_dense(list(args[:dim]), list(args[dim:]))
+
+
 SYMPY_INTERP = {
-    "IsNonOverlappingAndDenseIndicator": eval_is_non_overlapping_and_dense,
+    "IsNonOverlappingAndDenseIndicator": _eval_is_non_overlapping_and_dense_flat,
     "cast_symbool_to_symint_guardless": cast_symbool_to_symint_guardless,
     "math": math,
     "torch": torch,
@@ -2634,7 +2686,7 @@ def cast_symbool_to_symint_guardless(
 
 
 def _lru_cache(
-    fn: Callable[..., _T], maxsize: Optional[int] = None
+    fn: Callable[..., _T], maxsize: int | None = None
 ) -> functools._lru_cache_wrapper[_T]:
     """
     Wrapper around lru_cache that clears when new info about shapes has been
@@ -2905,7 +2957,7 @@ class _CppShapeGuardsHelper(_ShapeGuardsHelper):
 
 
 class LoggingShapeGuardPrinter(ShapeGuardPythonPrinter):
-    def __init__(self, var_to_sources: Mapping[sympy.Symbol, list[Source]]):
+    def __init__(self, var_to_sources: Mapping[sympy.Symbol, list[Source]]) -> None:
         super().__init__(var_to_sources, lambda n: n.name, var_to_sources)
 
 
@@ -2922,7 +2974,7 @@ def __init__(
         self,
         symbol_to_source: dict[sympy.Symbol, list[Source]],
         source_name_to_debug_name: Mapping[str, str],
-    ):
+    ) -> None:
         super().__init__()
         self.symbol_to_source = symbol_to_source
         self.source_name_to_debug_name = source_name_to_debug_name
@@ -3163,7 +3215,7 @@ def add_equality(self, source: Source, expr: sympy.Expr) -> None:
     def _reduce_congruences(self) -> dict[sympy.Symbol, set[sympy.Expr]]:
         reduced_congruences: dict[sympy.Symbol, set[sympy.Expr]] = {}
         for s, congruences in self._congruences.items():
-            remainder_modulus_pairs = []
+            remainder_modulus_pairs: list[tuple[sympy.Integer, sympy.Integer]] = []
             congruences_to_check = set()
             for congruence in congruences:
                 base, divisor = congruence.args
@@ -3542,7 +3594,7 @@ def _check_same_range(c: Mapping[str, int], dim: object) -> bool:
     def prettify_results(
         self,
         original_signature: inspect.Signature,
-        dynamic_shapes: Union[dict[str, Any], tuple[Any], list[Any]],
+        dynamic_shapes: dict[str, Any] | tuple[Any] | list[Any],
         constraint_violation_error: object,
         forced_specializations: dict[str, str],
     ) -> str:
@@ -3640,8 +3692,8 @@ def relation_with_digit(expr: str, op: str, digit: int) -> None:
 
         self._process_derived_dim_roots(results, name_to_dim)
 
-        dims = []
-        others = []
+        dims: list[str] = []
+        others: list[str] = []
 
         # order results by source name
         results2 = {
@@ -3712,7 +3764,7 @@ class ValueRangesSLoc:
 
 
 @contextmanager
-def _suppress_guards(shape_env: ShapeEnv) -> Iterator[None]:
+def _suppress_guards(shape_env: ShapeEnv) -> Generator[None, None, None]:
     shape_env._suppress_guards_enter()
     try:
         yield
@@ -3722,7 +3774,7 @@ def _suppress_guards(shape_env: ShapeEnv) -> Iterator[None]:
 
 @dataclass(slots=True)
 class _FrameLocalResult:
-    loc: Optional[str] = None
+    loc: str | None = None
     locals: dict[str, Any] = field(default_factory=dict)
     symbols: dict[str, str] = field(default_factory=dict)
 
@@ -3741,8 +3793,8 @@ class ShapeEnv:
     def __init__(
         self,
         *,
-        should_record_events: Optional[bool] = None,
-        tracked_fakes: Optional[list[Any]] = None,
+        should_record_events: bool | None = None,
+        tracked_fakes: list[Any] | None = None,
         **kwargs: Any,
     ) -> None:
         self._init(**kwargs)
@@ -3824,9 +3876,9 @@ def _init(
         specialize_zero_one: bool = True,
         # When True, assume input sizes which have the same size are
         # symbolically equal.
-        duck_shape: Optional[bool] = None,
+        duck_shape: bool | None = None,
         # For debugging
-        co_fields: Optional[dict[str, str]] = None,
+        co_fields: dict[str, str] | None = None,
         # When True, whenever safe, we will generate a deferred runtime assert
         # instead of a guard whenever we know that an expression must be True,
         # otherwise it would be an error, even for backed SymInts (where we
@@ -3875,7 +3927,7 @@ def _init(
         # are conservative: the int MUST fall in the range, but the
         # range may contain ints which may not actually appear in
         # practice
-        self.var_to_range: dict[sympy.Symbol, ValueRanges] = {}
+        self.var_to_range: dict[sympy.Symbol, ValueRanges[sympy.Expr]] = {}
         self.var_to_range_sloc: dict[sympy.Symbol, ValueRangesSLoc] = {}
         self.source_name_to_debug_name: dict[str, str] = {}
         self.var_to_sources: dict[sympy.Symbol, list[Source]] = {}
@@ -3898,12 +3950,17 @@ def _init(
         self.unbacked_renamings: dict[sympy.Symbol, sympy.Symbol] = {}
         # Set holds a % b expressions that evaluate to 0.
         self.divisible: set[sympy.Expr] = set()
+        # Exclusion constraints from automatic_dynamic transitions.
+        # Each (symbol, excluded_value) pair represents one dim/scalar that
+        # transitioned static → dynamic. All pairs are combined into a single
+        # Or(Ne(...), ...) guard in produce_guards_verbose.
+        self.exclusion_constraints: list[tuple[sympy.Symbol, int]] = []
         # Set that holds "size-like" symbols.  When we perform
         # "size-oblivious" tests, these can be assumed to be >= 2.
         self.size_like: set[sympy.Symbol] = set()
         # Duck-shaping says that if two input tensors have the same size,
         # they get assigned the same symbolic variable
-        self.val_to_var: dict[int, sympy.Symbol] = {}
+        self.val_to_var: dict[IntLikeType | FloatLikeType, sympy.Symbol] = {}
         self.unbacked_symfloat_counter = 0
         self.unbacked_symint_counter = 0
         # Similar to guards, but these MUST evaluate to true and can
@@ -3935,7 +3992,7 @@ def _init(
         #     latest key, an assert will only show up at the moment when
         #     we can actually codegen it.
         self.deferred_runtime_asserts: dict[
-            Optional[sympy.Symbol], list[RuntimeAssert]
+            sympy.Symbol | None, list[RuntimeAssert]
         ] = {}
         # This exists so we can efficiently invalidate the cache (it's used as
         # part of the cache key); otherwise we'd have to iterate through
@@ -3944,8 +4001,9 @@ def _init(
         self.log = log
         self.log.info("create_env")
         self.frozen = False
+        self._error_on_new_guards = False
         self.runtime_asserts_frozen = False
-        self.dim_constraints: Optional[DimConstraints] = None
+        self.dim_constraints: DimConstraints | None = None
         self.counter: Counter[str] = collections.Counter()
         # Mapping from sympy.Symbol to the number of guards which mention this
         # symbol
@@ -3980,6 +4038,9 @@ def _init(
         # Version counter used to invalidate cached values
         self._prev_cache_key = self._get_key()
         self._version_counter = 0
+        # Separate counter tracking only replacement changes, used by
+        # SymNode.expr
+        self._replacements_version_counter = 0
 
         # Each time divisible is changed this should be set to True, this is set in _update_version_counter.
         self._resimplify_floor_div_axioms = True
@@ -3990,7 +4051,9 @@ def _init(
         #   2. list of arguments
         # This drastically reduces the size of the FX graph, avoiding
         # duplicated nodes.
-        self.fx_node_cache: dict[tuple[Callable, tuple[Any, ...]], torch.fx.Node] = {}
+        self.fx_node_cache: dict[
+            tuple[Callable[..., Any], tuple[Any, ...]], torch.fx.Node
+        ] = {}
         self.source_to_symbol: dict[str, sympy.Symbol] = {}
 
         # Suppose you want to replace an unbacked symbol with another
@@ -4022,6 +4085,11 @@ def _init(
 
         self.specialization_stacks: dict[Source, traceback.StackSummary] = {}
 
+        # Used by _get_unbacked_replacements / _sub_unbacked_exprs for
+        # optimization_hint canonicalization of unbacked expressions.
+        self._equality_graph: dict[sympy.Expr, OrderedSet[sympy.Expr]] | None = None
+        self._unbacked_replacements: dict[sympy.Expr, sympy.Expr] | None = None
+
         self.trace_asserts = trace_asserts
 
         self.specializations: OrderedSet[Specialization] = OrderedSet()
@@ -4082,7 +4150,7 @@ def prefer_deferred_runtime_asserts_over_guards(self) -> bool:
     @contextmanager
     def patch_source_specialization(
         self, source: Source, check_fn: Callable[[sympy.Symbol], sympy.Expr]
-    ) -> Iterator[None]:
+    ) -> Generator[None, None, None]:
         """
         Temporarily add symbol-level axioms to the ShapeEnv. This is useful when you want to "fork"
         and have parallel universes of ShapeEnvs. For example, we use this when doing multi-graph
@@ -4151,9 +4219,13 @@ def check_equal(self, other: ShapeEnv) -> None:
             # source locations are OK to diverge
             "var_to_range_sloc",
             "replacements_slocs",
+            "_replacements_version_counter",
             "_resimplify_floor_div_axioms",
             "_expr_sym_node_id",
             "specialization_stacks",
+            # Cached state for optimization_hint unbacked canonicalization
+            "_equality_graph",
+            "_unbacked_replacements",
         )
 
         # Mapping of the value of each to-be-compared field into the values that
@@ -4184,7 +4256,7 @@ def map_value(key: str, value: Any) -> Any:
 
         shape_env_check_state_equal(self, other, non_state_variable_names, map_value)
 
-    def _snapshot_tracked_fakes(self) -> Optional[list[Any]]:
+    def _snapshot_tracked_fakes(self) -> list[Any] | None:
         if self.tracked_fakes is None:
             return None
 
@@ -4208,7 +4280,7 @@ def _last_event_index(self) -> int:
         return len(self.events) - 1
 
     @contextmanager
-    def _recording(self) -> Iterator[None]:
+    def _recording(self) -> Generator[None, None, None]:
         self.is_recording = True
         try:
             yield
@@ -4220,7 +4292,9 @@ def _eliminate_unbacked(self, orig_s: sympy.Symbol, new_s: sympy.Expr) -> None:
         self._set_replacement(orig_s, new_s, "eliminate_unbacked")
 
     @record_shapeenv_event()
-    def set_real_tensor_prop_unbacked_vals(self, k: sympy.Symbol, v: int) -> None:
+    def set_real_tensor_prop_unbacked_vals(
+        self, k: sympy.Symbol, v: int | float | torch.Tensor
+    ) -> None:
         """Used only when propagate_real_tensors; registers a value for an
         unbacked symbol, which can be used last resort to resolve hints."""
         log.info("set_real_tensor_prop_unbacked_vals %s = %s", k, v)
@@ -4257,7 +4331,7 @@ def _constrain_is_bounded(self, a: sympy.Symbol, upper_bound: int) -> None:
 
     @record_shapeenv_event()
     def _constrain_range_for_size(
-        self, a: sympy.Symbol, min: Optional[int] = None, max: Optional[int] = None
+        self, a: sympy.Symbol, min: int | None = None, max: int | None = None
     ) -> None:
         if min is None:
             min = 0
@@ -4346,7 +4420,7 @@ def _ignore_fresh_unbacked_symbols_set(self, b: bool) -> bool:
         return prev
 
     @contextmanager
-    def ignore_fresh_unbacked_symbols(self) -> Iterator[None]:
+    def ignore_fresh_unbacked_symbols(self) -> Generator[None, None, None]:
         """
         Indicates that the newly allocated unbacked SymInts are being
         discarded
@@ -4378,7 +4452,7 @@ def freeze_runtime_asserts(self) -> None:
         # self.prefer_deferred_runtime_asserts_over_guards = False
         self.runtime_asserts_frozen = True
 
-    def _create_symbol_for_source(self, source: Source) -> Optional[sympy.Symbol]:
+    def _create_symbol_for_source(self, source: Source) -> sympy.Symbol | None:
         if not self._translation_validation_enabled:
             return None
         srcname = source.name
@@ -4405,9 +4479,9 @@ def _check_translation_validate(self) -> None:
     @record_shapeenv_event()
     def _create_fx_call_function(
         self,
-        op: Callable,
-        args: tuple,
-    ) -> tuple[Optional[torch.fx.Node], bool]:
+        op: Callable[..., object],
+        args: tuple[Any, ...],
+    ) -> tuple[torch.fx.Node | None, bool]:
         # Cache this tuple in order to avoid duplicated nodes.
         node_key = (op, args)
         # Flags whether the returned node was cached or not.
@@ -4439,7 +4513,7 @@ def _create_fx_placeholder_and_z3var(
         self,
         symbol: sympy.Symbol,
         type: type,
-    ) -> Optional[torch.fx.Node]:
+    ) -> torch.fx.Node | None:
         if not self._translation_validation_enabled:
             return None
 
@@ -4463,7 +4537,7 @@ def _create_fx_placeholder_and_z3var(
 
         return self.fx_node_cache[node_key]
 
-    def _remove_fx_node(self, node: Optional[torch.fx.Node]) -> None:
+    def _remove_fx_node(self, node: torch.fx.Node | None) -> None:
         if self._translation_validation_enabled and node is not None:
             self.name_to_node.pop(node.name)
             self.graph.erase_node(node)
@@ -4497,9 +4571,27 @@ def _suppress_guards_exit(self) -> None:
         TLS.suppress_guards = old
 
     def suppress_guards(self) -> _GeneratorContextManager[None]:
-        """Context manager to ignore all guards generated inside"""
+        """Context manager to ignore all guards generated inside."""
         return _suppress_guards(self)
 
+    @contextmanager
+    def error_on_new_guards(self) -> Generator[None, None, None]:
+        """Context manager that raises _ShapeEnvGuardError if a guard is attempted.
+
+        Temporarily freezes the ShapeEnv and makes _check_frozen raise
+        instead of warn, so that guard-installing code paths produce an
+        exception that is not cached by the _inner_evaluate_expr LRU cache.
+        """
+        old_frozen = self.frozen
+        old_error = self._error_on_new_guards
+        self.frozen = True
+        self._error_on_new_guards = True
+        try:
+            yield
+        finally:
+            self.frozen = old_frozen
+            self._error_on_new_guards = old_error
+
     def _get_key(self) -> tuple[int, int, int, int]:
         """
         Defines the current "state" of the guards we've accumulated in this ShapeEnv.
@@ -4546,7 +4638,7 @@ def _produce_dyn_sizes_from_int_tuple(
         tensor_size: Sequence[IntLikeType],
         source: Source,
         symbolic_context: SymbolicContext,
-        hint_overrides: Optional[dict[int, int]] = None,
+        hint_overrides: dict[int, int] | None = None,
     ) -> list[sympy.Expr]:
         if not all(not is_symbolic(val) for val in tensor_size):
             raise AssertionError(
@@ -4595,7 +4687,7 @@ def create_symbolic_sizes_strides_storage_offset(
         ex: torch.Tensor,
         source: Source,
         *,
-        symbolic_context: Optional[SymbolicContext] = None,
+        symbolic_context: SymbolicContext | None = None,
     ) -> tuple[
         tuple[IntLikeType, ...],
         tuple[IntLikeType, ...],
@@ -4632,7 +4724,7 @@ def create_symbolic_sizes_strides_storage_offset(
     # Case 1: when SymInt is backed, dynamo can proceed with FakeTensors that have concrete shape.
     # produce_guards will trigger specializations on the outer stuff
 
-    # Case 2: when the SymInt is unbacked, we will throw an data dependent error in require_hint().
+    # Case 2: when the SymInt is unbacked, we will throw a data dependent error in guarding_hint_or_throw().
     #
     # It's probably good for now but it's important to note that this approach has implications for
     # the original shape_env when checking guards in different order.
@@ -4670,7 +4762,7 @@ def _maybe_specialize_sym_int_with_hint(
                 raise AssertionError(
                     "expect the symbol is created from an shape env other than current one."
                 )
-            return maybe_sym.node.require_hint()
+            return guarding_hint_or_throw(maybe_sym.node)
         return maybe_sym
 
     @record_shapeenv_event()
@@ -4684,8 +4776,8 @@ def _create_symbolic_sizes_strides_storage_offset(
         is_dim_dynamic: Sequence[bool],
         source: Source,
         *,
-        symbolic_context: Optional[SymbolicContext] = None,
-        hint_overrides: Optional[dict[int, int]] = None,
+        symbolic_context: SymbolicContext | None = None,
+        hint_overrides: dict[int, int] | None = None,
     ) -> tuple[
         tuple[IntLikeType, ...],
         tuple[IntLikeType, ...],
@@ -4756,6 +4848,27 @@ def _create_symbolic_sizes_strides_storage_offset(
         size: list[sympy.Expr] = self._produce_dyn_sizes_from_int_tuple(
             ex_size, source, symbolic_context, hint_overrides=hint_overrides
         )
+        # Record tensor exclusion constraints for stable graph selection.
+        # The ndim check guards against stale excluded_sizes from graph
+        # breaks where the resumed tensor may have different dimensionality.
+        # Skip dims with hint overrides: the overridden hint in
+        # backed_var_to_val would mismatch the excluded value, causing the
+        # not-all check in produce_guards_verbose to emit a guard that
+        # immediately fails.
+        excluded_sizes = getattr(symbolic_context, "excluded_sizes", None)
+        if (
+            excluded_sizes
+            and len(excluded_sizes) == dim
+            and any(v is not None for v in excluded_sizes)
+        ):
+            for i in range(dim):
+                ev = excluded_sizes[i]
+                if (
+                    ev is not None
+                    and isinstance(size[i], sympy.Symbol)
+                    and i not in (hint_overrides or {})
+                ):
+                    self._record_exclusion_constraint(size[i], ev)
         stride = self._compute_symbolic_stride(
             source,
             size,
@@ -4821,14 +4934,14 @@ def _compute_symbolic_stride(
         ex_stride: Sequence[IntLikeType],
         dynamic_strides: Sequence[DimDynamic],
         constraint_strides: Sequence[
-            Optional[Union[StrictMinMaxConstraint, RelaxedUnspecConstraint]]
+            StrictMinMaxConstraint | RelaxedUnspecConstraint | None
         ],
         are_sizes_static: bool,
         symbolic_context: SymbolicContext,
     ) -> list[sympy.Expr]:
         from torch._dynamo.source import TensorProperty, TensorPropertySource
 
-        stride: list[Optional[sympy.Expr]] = [None] * len(size)
+        stride: list[sympy.Expr | None] = [None] * len(size)
         candidates: dict[IntLikeType, sympy.Expr] = {}
 
         # iterate over unbound strides in val ascending order with
@@ -4877,8 +4990,8 @@ def create_symintnode(
         self,
         sym: sympy.Expr,
         *,
-        hint: Optional[int],
-        source: Optional[Source] = None,
+        hint: int | float | bool | torch.SymInt | None,
+        source: Source | None = None,
     ) -> IntLikeType:
         """Create a SymInt value from a symbolic expression
 
@@ -4923,8 +5036,8 @@ def create_symfloatnode(
         self,
         sym: sympy.Expr,
         *,
-        hint: Optional[int | float | bool],
-        source: Optional[Source] = None,
+        hint: int | float | bool | torch.SymInt | None,
+        source: Source | None = None,
     ) -> FloatLikeType:
         """Create a SymFloat value from a symbolic expression"""
         if self._translation_validation_enabled and source is not None:
@@ -4960,17 +5073,28 @@ def create_symfloatnode(
             out = SymFloat(SymNode(sym, self, float, hint, fx_node=fx_node))
         return out
 
+    @record_shapeenv_event()
+    def _record_exclusion_constraint(self, sym: sympy.Symbol, val: int) -> None:
+        self.exclusion_constraints.append((sym, val))
+
     @record_shapeenv_event()
     def create_unspecified_symint_and_symbol(
-        self, value: int, source: Source, dynamic_dim: DimDynamic
+        self,
+        value: int,
+        source: Source,
+        dynamic_dim: DimDynamic,
+        excluded_value: int | None = None,
     ) -> IntLikeType:
         """Create a SymInt wrapping a new unspecified symbol"""
+        sym = self.create_unspecified_symbol(
+            value,
+            source=source,
+            dynamic_dim=dynamic_dim,
+        )
+        if excluded_value is not None:
+            self._record_exclusion_constraint(sym, excluded_value)
         return self.create_symintnode(
-            self.create_unspecified_symbol(
-                value,
-                source=source,
-                dynamic_dim=dynamic_dim,
-            ),
+            sym,
             hint=value,
             source=source,
         )
@@ -4985,14 +5109,14 @@ def _log_create_unbacked_symbol(
         self,
         prefix: str,
         symbol: sympy.Symbol,
-        vr: ValueRanges,
-        source: Optional[Source] = None,
-        sym_node: Optional[SymNode] = None,
+        vr: ValueRanges[sympy.Expr],
+        source: Source | None = None,
+        sym_node: SymNode | None = None,
     ) -> None:
         is_debug = config.extended_debug_create_symbol is not None and str(
             symbol
         ) in config.extended_debug_create_symbol.split(",")
-        sloc: Union[str, SLoc]
+        sloc: str | SLoc
         if source is None:
             sloc, maybe_extra_debug = self._get_stack_summary(is_debug)
         else:
@@ -5046,7 +5170,7 @@ def create_unbacked_symfloat(self) -> SymFloat:
         return SymFloat(sym_node)
 
     @record_shapeenv_event()
-    def create_unbacked_symint(self, source: Optional[Source] = None) -> SymInt:
+    def create_unbacked_symint(self, source: Source | None = None) -> SymInt:
         """Create a symbolic integer without a hint value"""
         symbol: sympy.Symbol = make_symbol(
             SymT.UNBACKED_INT, self.unbacked_symint_counter, integer=True
@@ -5105,11 +5229,11 @@ def create_unbacked_symbool(self) -> SymBool:
     @record_shapeenv_event()
     def create_unspecified_symbol(
         self,
-        val: Union[int, SymInt, float, SymFloat],
+        val: int | SymInt | float | SymFloat,
         source: Source,
         dynamic_dim: DimDynamic = DimDynamic.DUCK,
         constraint_dim: DimConstraint = None,  # NB: includes None
-        symbolic_context: Optional[StatelessSymbolicContext] = None,
+        symbolic_context: SymbolicContext | None = None,
     ) -> sympy.Expr:
         """
         Create a symbol with an unspecified value
@@ -5135,13 +5259,13 @@ def create_unspecified_symbol(
     @record_shapeenv_event()
     def create_symbol(
         self,
-        val: int,
+        val: IntLikeType | FloatLikeType,
         source: Source,
         dynamic_dim: DimDynamic = DimDynamic.DUCK,
         constraint_dim: DimConstraint = None,  # NB: includes None
-        positive: Optional[bool] = True,
+        positive: bool | None = True,
         do_not_specialize_zero_one: bool = False,
-        symbolic_context: Optional[StatelessSymbolicContext] = None,
+        symbolic_context: SymbolicContext | None = None,
     ) -> sympy.Expr:
         """Create a new symbol which is tracked by this ShapeEnv"""
         # check if constraint_dim is actually static integer
@@ -5155,7 +5279,7 @@ def create_symbol(
                     f"Static shape constraint of {constraint_dim.vr.lower} does not match input size of {val}, "
                     f"for {source.name}"
                 )
-            if symbolic_context:
+            if isinstance(symbolic_context, StatelessSymbolicContext):
                 from torch._dynamo.source import TensorPropertySource
 
                 if not isinstance(source, TensorPropertySource):
@@ -5194,6 +5318,8 @@ def create_symbol(
             # If so, we allocate a fresh symbol but add a runtime equality check
             # via torch._check against the existing symbols with the same shape_id.
             shape_id = None
+            unbacked_min = None
+            unbacked_max = None
             if (
                 isinstance(symbolic_context, StatelessSymbolicContext)
                 and symbolic_context.shape_ids is not None
@@ -5203,10 +5329,30 @@ def create_symbol(
                 if isinstance(source, TensorPropertySource) and source.idx is not None:
                     shape_id = symbolic_context.shape_ids.get(source.idx)
 
+            # Check for unbacked bounds
+            if (
+                isinstance(symbolic_context, StatelessSymbolicContext)
+                and symbolic_context.unbacked_bounds is not None
+            ):
+                from torch._dynamo.source import TensorPropertySource
+
+                if isinstance(source, TensorPropertySource) and source.idx is not None:
+                    bounds = symbolic_context.unbacked_bounds.get(source.idx)
+                    if bounds is not None:
+                        unbacked_min, unbacked_max = bounds
+
             # Always allocate a fresh unbacked symbol
             out = self.create_unbacked_symint(source).node.expr
             self._constrain_range_for_size(out)
 
+            # Apply min/max bounds via torch._check if specified
+            if unbacked_min is not None or unbacked_max is not None:
+                out_symint = self.create_symintnode(out, hint=None)
+                if unbacked_min is not None:
+                    torch._check(out_symint >= unbacked_min)
+                if unbacked_max is not None:
+                    torch._check(out_symint <= unbacked_max)
+
             # Add runtime equality check for shape_id if applicable
             if shape_id is not None:
                 if shape_id in self._shape_id_to_unbacked_symbol:
@@ -5456,7 +5602,7 @@ def _debug_name(self, source: Source) -> str:
         return self.source_name_to_debug_name.get(src_name, src_name)
 
     def _render_range_for_constraint_violation(
-        self, source: Source, c: Union[StrictMinMaxConstraint, RelaxedUnspecConstraint]
+        self, source: Source, c: StrictMinMaxConstraint | RelaxedUnspecConstraint
     ) -> str:
         if isinstance(c, StrictMinMaxConstraint):
             lower, upper = c.vr.lower, c.vr.upper
@@ -5490,11 +5636,11 @@ def produce_guards_verbose(
         sources: Sequence[Source],
         source_ref: Callable[[Source], str] = lambda n: n.name,
         *,
-        guards: Optional[list[ShapeGuard]] = None,
-        input_contexts: Optional[DimList[SymbolicContext]] = None,
+        guards: list[ShapeGuard] | None = None,
+        input_contexts: DimList[SymbolicContext] | None = None,
         # Encodes user-specified input shape equations of the form s = s' and s = fn(s').
         # (See docs on EqualityConstraint for details of the encoding.)
-        equalities_inputs: Optional[EqualityConstraint] = None,
+        equalities_inputs: EqualityConstraint | None = None,
         _simplified: bool = False,
         # Indicates if we should produce guards for known static values.
         ignore_static: bool = True,
@@ -5540,7 +5686,9 @@ def produce_guards_verbose(
             raise AssertionError(f"len({placeholders}) != len({sources})")
         Tensorlike = (torch.Tensor, FakeTensorMeta)
 
-        def _create_no_constraints_context(t: Tensor) -> StatelessSymbolicContext:
+        def _create_no_constraints_context(
+            t: Tensor,
+        ) -> StatelessSymbolicContext[..., Any]:
             return StatelessSymbolicContext(
                 # Ignored; only the constraints part is relevant below.
                 dynamic_sizes=[DimDynamic.DYNAMIC] * t.dim(),
@@ -5666,7 +5814,7 @@ def record_constraint_violation(
             warn_only: bool,
             debug_name: str,
             msg: str,
-            hint: Optional[Callable[[], str]] = None,
+            hint: Callable[[], str] | None = None,
         ) -> None:
             constraint_violations.append(
                 (warn_only, debug_name, lambda: f"{msg}{hint()}" if hint else msg)
@@ -5864,7 +6012,7 @@ def track_symfloat(source: Source, val: FloatLikeType) -> None:
                 s = sympy.Float(val)
                 input_guards.append((source, s))
 
-        # pyrefly: ignore [no-matching-overload]
+        # pyrefly: ignore [bad-argument-type, no-matching-overload]
         for t, source, context in zip(placeholders, sources, input_contexts):
             if isinstance(source, str):
                 from torch._dynamo.source import LocalSource
@@ -5901,16 +6049,23 @@ def track_symfloat(source: Source, val: FloatLikeType) -> None:
                 ]
                 attrs, _ = t.__tensor_flatten__()
                 for attr in attrs:
-                    inner_t = getattr(t, attr)
-                    inner_context = context.inner_contexts[attr]
-                    sources_tensors_constraints.append(
-                        (
-                            AttrSource(source, attr),
-                            inner_t,
-                            inner_context.constraint_sizes,  # type: ignore[attr-defined]
-                            inner_context.constraint_strides,  # type: ignore[attr-defined]
-                        )
-                    )
+                    match getattr(t, attr):
+                        case torch.Tensor() as inner_t:
+                            inner_context = context.inner_contexts[attr]
+                            sources_tensors_constraints.append(
+                                (
+                                    AttrSource(source, attr),
+                                    inner_t,
+                                    inner_context.constraint_sizes,  # type: ignore[attr-defined]
+                                    inner_context.constraint_strides,  # type: ignore[attr-defined]
+                                )
+                            )
+                        case OpaqueBase():
+                            pass
+                        case unexpected:
+                            raise AssertionError(
+                                f"expected Tensor or OpaqueBase, got {type(unexpected)}"
+                            )
             else:
                 sources_tensors_constraints = [
                     (source, t, context.constraint_sizes, context.constraint_strides)  # type: ignore[attr-defined]
@@ -6153,7 +6308,7 @@ def issue_guard(guard: ShapeGuard) -> None:
 
             if not sources:
                 raise AssertionError(f"sources must not be empty for symbol {symbol}")
-            bounds = []
+            bounds: list[sympy.Basic] = []
             rf = source_ref(sources[0])
             verbose_expr = ""
             if r.lower not in (-sympy.oo, -int_oo):
@@ -6223,6 +6378,93 @@ def issue_guard(guard: ShapeGuard) -> None:
                     else:
                         raise NotImplementedError(f"Unimplemented for lang: {lang}")
 
+        # Exclusion guard for stable graph selection with automatic dynamic.
+        #
+        # When automatic_dynamic promotes a static dim to dynamic, the new
+        # (more general) graph is inserted *before* the old (specialized) graph
+        # in the guard cache.  Without an exclusion guard, inputs that exactly
+        # match the old graph's static sizes would be captured by the new
+        # dynamic graph instead, violating the invariant "once an input is
+        # served by graph X it is always served by graph X". This condition
+        # is true iff there is no branching on dynamic shapes.
+        #
+        # Soundness argument (cache-flip / LIFO order):
+        #   Graph_new sits before Graph_old in the cache.  Graph_old accepts
+        #   only inputs whose sizes match its static constraints exactly.
+        #   Graph_new must therefore reject exactly that set of inputs so they
+        #   fall through to Graph_old.  The excluded values are the static
+        #   sizes from Graph_old, so the guard
+        #       Or(Ne(s0, v0), Ne(s1, v1), ...)
+        #   passes iff at least one dim differs from the old sizes — i.e. the
+        #   input does NOT fully match Graph_old.  Conversely, when every dim
+        #   matches the old sizes the guard fails and the input falls through
+        #   to Graph_old, which is guaranteed to accept it.
+        #
+        # Theorem: For graphs G0, ..., Gn compiled via progressive dynamism
+        # (one dim per step), each input is accepted by at most one graph.
+        #
+        #   Setup: G0 is all-static with shape S. Gk is created by making
+        #   dim d_k dynamic, with exclusion guard d_k != S[d_k].
+        #
+        #   Proof by induction on n:
+        #
+        #   Base case (n=0): Only G0, all-static. Trivially unique.
+        #
+        #   Inductive step: Assume the property holds for G0, ..., G_{n-1}.
+        #   We add Gn with newly-dynamic dim d_n and exclusion d_n != S[d_n].
+        #
+        #   For any input X that passes Gn's shape guards, exactly one of:
+        #
+        #   Case A — exclusion passes (X[d_n] != S[d_n]):
+        #     Dim d_n is static in all G0, ..., G_{n-1} with value S[d_n],
+        #     so X fails all prior graphs on that dim. Only Gn accepts X.
+        #
+        #   Case B — exclusion rejects (X[d_n] == S[d_n]):
+        #     X matches Gn's shape guards on all other dims, and matches
+        #     the static value for d_n. So X satisfies G_{n-1}'s shape
+        #     guards. By the inductive hypothesis, exactly one of
+        #     G0, ..., G_{n-1} accepts X. Gn rejects X.
+        #
+        #   Corollary: Evaluation order does not affect correctness.
+        #
+        # All exclusion pairs across all tensors and scalars are flattened
+        # into a single list — each pair is just (symbol, excluded_int),
+        # and the multi-tensor case is the same logic as multi-dim within
+        # one tensor.  The combined Or rejects only when ALL pairs match
+        # simultaneously, which is the exact condition for Graph_old to
+        # accept.  If the current concrete values already match every
+        # excluded value the guard is skipped (it would fail on creation).
+        import torch._dynamo.config as dynamo_config
+
+        if (
+            dynamo_config.automatic_dynamic_exclusion_guard
+            and not dynamo_config.enable_compiler_collectives
+            and self.exclusion_constraints
+        ):
+            all_pairs = [
+                (sym, val)
+                for sym, val in self.exclusion_constraints
+                if symbol_to_source.get(sym)
+            ]
+            if all_pairs and not all(
+                self.backed_var_to_val.get(sym) == val for sym, val in all_pairs
+            ):
+                if len(all_pairs) == 1:
+                    excl_expr = sympy.Ne(
+                        all_pairs[0][0], all_pairs[0][1], evaluate=False
+                    )
+                else:
+                    excl_expr = sympy.Or(
+                        *[sympy.Ne(sym, val, evaluate=False) for sym, val in all_pairs]
+                    )
+                for exprs, printer, lang in zip(all_exprs, printers, langs):
+                    guard_expr = printer.doprint(excl_expr)
+                    if lang == "verbose_python":
+                        guard_expr = (
+                            f"{guard_expr}  # exclusion guard for automatic dynamic"
+                        )
+                    exprs.append(guard_expr)
+
         if constraint_violations:
             warn_msgs: list[str] = []
             error_msgs: list[str] = []
@@ -6305,11 +6547,11 @@ def issue_guard(guard: ShapeGuard) -> None:
 
     def produce_guards_expression(
         self,
-        placeholders: Sequence[Union[SymInt, FakeTensor]],
+        placeholders: Sequence[SymInt | FakeTensor],
         *,
-        guards: Optional[list[ShapeGuard]] = None,
+        guards: list[ShapeGuard] | None = None,
         ignore_static: bool = True,
-    ) -> Optional[str]:
+    ) -> str | None:
         """
         Expected to be used with evaluate_guards_expression(). Produces the guards
         for the given placeholders and returns a string expression to be evaluated
@@ -6328,14 +6570,14 @@ def produce_guards_expression(
             return " and ".join(produced_guards)
         return None
 
-    def evaluate_symexpr(self, code: str) -> Union[int, float, bool]:
+    def evaluate_symexpr(self, code: str) -> int | float | bool:
         """
         To be used by compile_fx to evaluate symexprs
         """
         args = {str(e): val for e, val in self.backed_var_to_val.items()}
         return eval(code, SYMPY_INTERP, args)
 
-    def deserialize_symexpr(self, code: str) -> Union[SymInt, SymFloat, SymBool]:
+    def deserialize_symexpr(self, code: str) -> SymInt | SymFloat | SymBool:
         """
         To be used by compile_fx to deserialize symexprs
         """
@@ -6457,7 +6699,7 @@ def format_guards(self, verbose: bool = False) -> str:
 
     def bound_sympy(
         self, expr: sympy.Expr, size_oblivious: bool = False
-    ) -> ValueRanges:
+    ) -> ValueRanges[sympy.Expr]:
         """Given a sympy expression, computes a ValueRanges bound for what values it can be"""
         # TODO: maybe it's guaranteed x in is var_to_range?
         var_to_range = {x: self.var_to_range.get(x, None) for x in expr.free_symbols}
@@ -6476,7 +6718,7 @@ def bound_sympy(
     @_lru_cache
     def get_axioms(
         self,
-        symbols: Optional[tuple[sympy.Symbol]] = None,
+        symbols: tuple[sympy.Symbol] | None = None,
         compute_hint: bool = False,
     ) -> tuple[SympyBoolean, ...]:
         """
@@ -6569,7 +6811,7 @@ def _is_nonneg_sum(self, expr: sympy.Expr) -> bool:
 
         return True
 
-    def _maybe_fast_eval_comparison(self, expr: sympy.Basic) -> Optional[sympy.Basic]:
+    def _maybe_fast_eval_comparison(self, expr: sympy.Basic) -> sympy.Basic | None:
         """
         Fast path for trivial comparisons: sum of non-negative terms >= 0.
         Returns sympy.true if pattern matches, None otherwise.
@@ -6592,6 +6834,25 @@ def _maybe_fast_eval_comparison(self, expr: sympy.Basic) -> Optional[sympy.Basic
 
         return None
 
+    def _maybe_evaluate_range_only(
+        self,
+        expr: sympy.Basic,
+        fallback: sympy.Basic | None = None,
+    ) -> sympy.Basic | None:
+        """
+        Lightweight range-based evaluation using only bound_sympy (value range
+        analysis), without expensive simplification, axiom matching, or symbol
+        reallocation.
+
+        Returns the resolved value if range analysis determines it, otherwise
+        returns fallback.
+        """
+        var_ranges = {x: self.var_to_range.get(x) for x in expr.free_symbols}
+        out = bound_sympy(expr, var_ranges)  # type: ignore[arg-type]
+        if out.is_singleton():
+            return out.lower
+        return fallback
+
     @_lru_cache
     def _maybe_evaluate_static(
         self,
@@ -6600,9 +6861,9 @@ def _maybe_evaluate_static(
         unbacked_only: bool = False,
         compute_hint: bool = False,
         size_oblivious: bool = False,
-        axioms: Optional[tuple[SympyBoolean]] = None,
-        var_to_range: Optional[tuple[tuple[sympy.Symbol, ValueRanges]]] = None,
-    ) -> Optional[sympy.Basic]:
+        axioms: tuple[SympyBoolean] | None = None,
+        var_to_range: tuple[tuple[sympy.Symbol, ValueRanges[sympy.Expr]]] | None = None,
+    ) -> sympy.Basic | None:
         """
         Tries to evaluate expr without introducing guards
 
@@ -6687,8 +6948,14 @@ def resimplify_floor_div(axioms: dict[sympy.Expr, sympy.Expr]) -> None:
     def replace(self, expr: _SympyT) -> _SympyT:
         """
         Apply symbol replacements to any symbols in the given expression.
+
+        IMPORTANT: The output of this method MUST depend only on
+        self.replacements and the input expr. Do not add dependencies on other
+        mutable state. SymNode.expr uses _replacements_version_counter (which
+        tracks only replacement changes) to cache calls to this method, so
+        depending on other state would cause stale cache results.
         """
-        replacements = {}
+        replacements: dict[sympy.Basic, sympy.Basic] = {}
         # pyrefly: ignore [missing-attribute]
         for s in expr.free_symbols:
             r = self._find(s)
@@ -6724,7 +6991,7 @@ def simplify(self, expr: _SympyT, size_oblivious: bool = False) -> _SympyT:
         # Simplify max(0/1, x) to x when x >= 0/1. max(1, x) is a commonly introduced
         # expression when creating contiguous strides.
         if not size_oblivious:
-            min_max_replacements = {}
+            min_max_replacements: dict[sympy.Basic, sympy.Basic] = {}
             for atom in expr.atoms(Max):  # type: ignore[has-type]
                 if len(atom.args) > 2:
                     continue
@@ -6740,11 +7007,11 @@ def simplify(self, expr: _SympyT, size_oblivious: bool = False) -> _SympyT:
                 expr = expr.xreplace(min_max_replacements)
 
         if expr.has(TruncToInt):
-            trunc_replacements = {}
+            trunc_replacements: dict[sympy.Basic, sympy.Basic] = {}
             for atom in expr.atoms(TruncToInt):
                 if isinstance(atom.args[0], IntTrueDiv):
                     base, divisor = atom.args[0].args
-                    if base % divisor == 0:
+                    if Mod(base, divisor) == 0:
                         trunc_replacements[atom] = CleanDiv(base, divisor)
                     else:
                         # TruncToInt(IntTrueDiv(a,b)) == FloorDiv(a, b)
@@ -6760,7 +7027,7 @@ def simplify(self, expr: _SympyT, size_oblivious: bool = False) -> _SympyT:
         # for now just do a separate pass to catch common nested case
         if expr.has(FloorDiv):
             self._update_divisible()
-            div_replacements = {}
+            div_replacements: dict[sympy.Basic, sympy.Basic] = {}
             for atom in expr.atoms(FloorDiv):
                 base, divisor = atom.args
                 if isinstance(divisor, FloorDiv):
@@ -6775,7 +7042,7 @@ def simplify(self, expr: _SympyT, size_oblivious: bool = False) -> _SympyT:
                 expr = expr.xreplace(div_replacements)
                 expr = safe_expand(expr)
         if expr.has(FloorDiv):
-            div_replacements = {}
+            div_replacements: dict[sympy.Basic, sympy.Basic] = {}
             pows = expr.atoms(sympy.Pow)
             rationals = expr.atoms(sympy.Rational).difference(expr.atoms(sympy.Integer))
             for fd in expr.atoms(FloorDiv):
@@ -6795,10 +7062,14 @@ def simplify(self, expr: _SympyT, size_oblivious: bool = False) -> _SympyT:
         return expr
 
     # TODO: overload for allow_none literal
+    @deprecated(
+        "use guarding_hint_or_throw or optimization_hint instead",
+        category=FutureWarning,
+    )
     @lru_cache(256)
     def size_hint(
         self, expr: sympy.Basic, *, allow_none: bool = False
-    ) -> Optional[sympy.Basic]:
+    ) -> sympy.Basic | None:
         """
         Gets a size hint for a given expression from the underlying shapes we had.
         Does not introduce a guard, so only use this when you can guarantee that
@@ -6841,13 +7112,46 @@ def size_hint(
             raise self._make_data_dependent_error(result_expr, expr)
         return result_expr
 
-    # NB: keep in sync with size_hint
     @lru_cache(256)
-    def has_hint(self, expr: sympy.Expr) -> bool:
-        result_expr = safe_expand(expr).xreplace(self.backed_var_to_val)
-        return (
-            result_expr.is_number
-            or self._maybe_evaluate_static(result_expr) is not None
+    def guarding_hint_or_throw(self, expr: sympy.Expr | int) -> int | bool:
+        """
+        Return a concrete hint for an expression.
+
+        Returns Python bool (True/False) for boolean expressions (e.g. Eq, Ne),
+        and Python int for integer expressions.
+        """
+        return _guarding_hint_or_throw_base(self, expr, {})
+
+    @lru_cache(256)
+    def has_guarding_hint(self, expr: sympy.Expr) -> bool:
+        try:
+            self.guarding_hint_or_throw(expr)
+        except GuardOnDataDependentSymNode:
+            return False
+        return True
+
+    def optimization_hint(
+        self, expr: sympy.Expr | int, fallback: int | None = None
+    ) -> int:
+        """
+        Return a concrete integer hint for an expression.
+
+        This function should be used for non-guarding based optimizations. If you
+        want a hint that you can guard on, use the guarding_hint API instead.
+
+        This function will hint unbacked symbols using user provided optimization
+        hints. If not provided, fallback will be used along with some heuristics
+        that try to maximize consistency with the shape environment.
+
+        Special cases:
+
+        - Complex numbers (containing sympy.I): raises an error since tensor
+          dimensions cannot be complex.
+        - Infinity (int_oo, sympy.oo): returns sys.maxsize.
+        - NaN (sympy.nan): returns the fallback value.
+        """
+        return _optimization_hint_base(
+            self, expr, precomputed_replacements={}, fallback=fallback
         )
 
     def _make_data_dependent_error(
@@ -6855,7 +7159,7 @@ def _make_data_dependent_error(
         expr: sympy.Basic,
         unhinted_expr: sympy.Basic,
         *,
-        expr_sym_node_id: Optional[int] = None,
+        expr_sym_node_id: int | None = None,
     ) -> GuardOnDataDependentSymNode:
         # TODO: in a Dynamo context, having user code, and having the
         # name of the local, will be much better
@@ -6912,8 +7216,8 @@ def _make_data_dependent_error(
     def _update_var_to_range(
         self,
         symbol: sympy.Symbol,
-        vr: ValueRanges,
-        vr_sloc: Optional[ValueRangesSLoc] = None,
+        vr: ValueRanges[sympy.Expr],
+        vr_sloc: ValueRangesSLoc | None = None,
         *,
         is_constraint: bool = False,
     ) -> None:
@@ -7127,6 +7431,7 @@ def _set_replacement(self, a: sympy.Symbol, tgt: sympy.Expr, msg: str) -> None:
         # them)
         if a not in self.replacements_slocs:
             self.replacements_slocs[a] = self._get_sloc()
+        self._replacements_version_counter += 1
         self._update_version_counter()
 
         # When specializing 'a == tgt', the equality should be also conveyed to
@@ -7146,6 +7451,15 @@ def _find(self, a: sympy.Symbol) -> sympy.Expr:
 
         a: b + c
         c: d
+
+        IMPORTANT: The output of this method MUST depend only on
+        self.replacements and the input symbol. Do not add dependencies on other
+        mutable state. SymNode.expr uses _replacements_version_counter (which
+        tracks only replacement changes) to cache calls to replace() (and
+        transitively this method), so depending on other state would cause
+        stale cache results. (Note: _set_replacement,  may read other fields
+        like var_to_range, but those are side effects that do not affect the
+        returned value.)
         """
         if a not in self.replacements:
             return a
@@ -7290,11 +7604,11 @@ def trivial_solve(lhs: sympy.Expr, rhs: sympy.Expr) -> bool:
     # See: Note - On 0/1 specialization
     def _default_value_range(
         self, do_not_specialize_zero_one: bool = False
-    ) -> ValueRanges:
+    ) -> ValueRanges[sympy.Expr]:
         lower = 0 if (do_not_specialize_zero_one or not self.specialize_zero_one) else 2
         return ValueRanges(lower, int_oo)
 
-    def _default_unspecified_value_range(self) -> ValueRanges:
+    def _default_unspecified_value_range(self) -> ValueRanges[sympy.Expr]:
         return ValueRanges.unknown_int()
 
     @_lru_cache
@@ -7312,8 +7626,12 @@ def _simplify_floor_div(self, expr: sympy.Expr) -> sympy.Expr:
         return self.simplify(expr)
 
     # We're about to add a guard/runtime assert, check if the ShapeEnv is frozen
-    # and if so issue a warning
+    # and if so issue a warning (or raise if error_on_new_guards is set)
     def _check_frozen(self, expr: sympy.Basic, concrete_val: sympy.Basic) -> None:
+        if self._error_on_new_guards:
+            raise _ShapeEnvGuardError(
+                f"Guard attempted while ShapeEnv guards are frozen: {expr} == {concrete_val}"
+            )
         if self.frozen:
             self.counter["ignored_backward_guard"] += 1
             signpost_event(
@@ -7335,7 +7653,7 @@ def _check_frozen(self, expr: sympy.Basic, concrete_val: sympy.Basic) -> None:
                 stack_info=log.getEffectiveLevel() < logging.WARNING,
             )
 
-    def _get_user_frame(self) -> Optional[types.FrameType]:
+    def _get_user_frame(self) -> types.FrameType | None:
         frame = inspect.currentframe()
         while frame is not None:
             if frame.f_code.co_filename not in uninteresting_files():
@@ -7344,9 +7662,9 @@ def _get_user_frame(self) -> Optional[types.FrameType]:
         return frame
 
     def _get_stack_summary(
-        self, is_debug: bool = False, framework_loc: Optional[str] = None
+        self, is_debug: bool = False, framework_loc: str | None = None
     ) -> tuple[SLoc, str]:
-        floc: Optional[Union[str, traceback.FrameSummary]] = framework_loc
+        floc: str | traceback.FrameSummary | None = framework_loc
         if floc is None:
             frame = self._get_user_frame()
             try:
@@ -7387,7 +7705,7 @@ def _get_stack_summary(
         return SLoc(floc, maybe_user_loc), maybe_extra_debug
 
     # Pass in framework_loc to override the framework location info
-    def _get_sloc(self, framework_loc: Optional[str] = None) -> SLoc:
+    def _get_sloc(self, framework_loc: str | None = None) -> SLoc:
         sloc, _ = self._get_stack_summary(framework_loc=framework_loc)
         return sloc
 
@@ -7430,7 +7748,7 @@ def _find_frame_locals(self) -> _FrameLocalResult:
             return _FrameLocalResult()
 
         # track involved locals and free symbols
-        def go(x: Any) -> Optional[str]:
+        def go(x: Any) -> str | None:
             if isinstance(x, torch.Tensor):
                 for y in x.size():
                     go(y)
@@ -7526,13 +7844,13 @@ def _log_guard(self, prefix: str, g: SympyBoolean, forcing_spec: bool) -> None:
     # A local variable to evaluate_expr stored in the class to avoid
     # using it for the lru_cache that is on top of it since it does
     # not effect the results. When needed its read directly.
-    _expr_sym_node_id: Optional[int] = None
+    _expr_sym_node_id: int | None = None
 
     def evaluate_sym_node(
         self,
         sym_node: SymNode,
         size_oblivious: bool = False,
-        fallback_value: Optional[bool] = None,
+        fallback_value: bool | None = None,
     ) -> sympy.Basic:
         """
         Given a a SymNode, evaluates sym_node.expr, adding guards if necessary.
@@ -7541,8 +7859,8 @@ def evaluate_sym_node(
         self._expr_sym_node_id = id(sym_node)
         return self.evaluate_expr(
             sym_node.expr,
-            sym_node.hint,
-            sym_node.fx_node,
+            sym_node.hint,  # pyrefly: ignore[bad-argument-type]
+            sym_node.fx_node,  # pyrefly: ignore[bad-argument-type]
             size_oblivious,
             fallback_value=fallback_value,
         )
@@ -7621,10 +7939,10 @@ def _log_real_tensor_propagation(
     def evaluate_expr(
         self,
         orig_expr: sympy.Basic,
-        hint: Optional[Union[int, bool, float]] = None,
-        fx_node: Optional[torch.fx.Node] = None,
+        hint: int | bool | float | None = None,
+        fx_node: torch.fx.Node | None = None,
         size_oblivious: bool = False,
-        fallback_value: Optional[bool] = None,
+        fallback_value: bool | None = None,
         *,
         forcing_spec: bool = False,
     ) -> sympy.Basic:
@@ -7650,12 +7968,12 @@ def evaluate_expr(
     def _inner_evaluate_expr(
         self,
         orig_expr: sympy.Basic,
-        hint: Optional[Union[int, bool, float]],
-        fx_node: Optional[torch.fx.Node],
+        hint: int | bool | float | None,
+        fx_node: torch.fx.Node | None,
         size_oblivious: bool,
         forcing_spec: bool,
         _suppress_guards_tls: bool,
-        fallback_value: Optional[bool] = None,
+        fallback_value: bool | None = None,
     ) -> sympy.Basic:
         try:
             return self._evaluate_expr(
@@ -7692,10 +8010,10 @@ def _log_suppressed_dde(self, a: SymBool, assumed_value: bool) -> None:
     def _evaluate_expr(
         self,
         orig_expr: sympy.Basic,
-        hint: Optional[Union[bool, int, float]] = None,
-        fx_node: Optional[torch.fx.Node] = None,
+        hint: bool | int | float | None = None,
+        fx_node: torch.fx.Node | None = None,
         size_oblivious: bool = False,
-        fallback_value: Optional[bool] = None,
+        fallback_value: bool | None = None,
         *,
         forcing_spec: bool = False,
     ) -> sympy.Basic:
@@ -7712,15 +8030,25 @@ def _evaluate_expr(
         def compute_concrete_val() -> sympy.Basic:
             if hint is None:
                 # This is only ever called for expressions WITHOUT unbacked
-                # symbols
-                r = self.size_hint(orig_expr)
-                if r is None:
-                    raise AssertionError("r must not be None")
-                return r
+                # symbols.  guarding_hint_or_throw returns Python bool for
+                # boolean expressions and int for integer expressions;
+                # sympify converts them to the proper sympy types
+                # (True -> sympy.true, 5 -> Integer(5)).
+                try:
+                    return sympy.sympify(self.guarding_hint_or_throw(orig_expr))
+                except GuardOnDataDependentSymNode:
+                    # guarding_hint_or_throw only does backed-symbol replacement.
+                    # For expressions with unbacked symbols resolvable via axioms
+                    # (e.g. Eq(x, 0) when torch._check(Ne(x, 0)) was previously
+                    # asserted), fall back to static evaluation with compute_hint.
+                    r = self._maybe_evaluate_static(orig_expr, compute_hint=True)
+                    if r is not None:
+                        return r
+                    raise
             else:
                 return sympy.sympify(hint)
 
-        concrete_val: Optional[sympy.Basic]
+        concrete_val: sympy.Basic | None
 
         # Check if:
         #   1. 'translation_validation' is set
@@ -7788,9 +8116,31 @@ def compute_concrete_val() -> sympy.Basic:
 
             # Try to quickly evaluate trivially true/false comparisons
             # using var_to_range, before calling expensive _maybe_evaluate_static.
-            fast_result = self._maybe_fast_eval_comparison(expr)
-            if fast_result is not None:
-                return fast_result
+            if (
+                torch.fx.experimental._config.aggressive_guard_free_semantics
+                < AggressiveGuardFreeMode.SKIP_RANGE_ANALYSIS
+            ):
+                fast_result = self._maybe_fast_eval_comparison(expr)
+                if fast_result is not None:
+                    return fast_result
+
+            # Aggressive guard-free semantics:
+            # VALUE_RANGE_ANALYSIS: use value range analysis (bound_sympy) before returning fallback
+            # SKIP_RANGE_ANALYSIS: skip range analysis entirely, just return fallback_value
+            aggressive_level = (
+                torch.fx.experimental._config.aggressive_guard_free_semantics
+            )
+            if hint is None and aggressive_level > 0 and fallback_value is not None:
+                if aggressive_level >= AggressiveGuardFreeMode.SKIP_RANGE_ANALYSIS:
+                    # Skip range analysis entirely
+                    self._log_suppressed_dde(orig_expr, fallback_value)
+                    return fallback_value
+                else:
+                    # Level 1: try range analysis first
+                    range_result = self._maybe_evaluate_range_only(expr, fallback_value)
+                    if range_result is fallback_value:
+                        self._log_suppressed_dde(orig_expr, fallback_value)
+                    return range_result
 
             static_expr = self._maybe_evaluate_static(
                 expr, size_oblivious=size_oblivious
@@ -7990,7 +8340,7 @@ def _should_skip_static_eval(self, expr: SympyBoolean) -> bool:
     @lru_cache(256)
     @record_shapeenv_event(save_tracked_fakes=True)
     def guard_or_defer_runtime_assert(
-        self, orig_expr: SympyBoolean, msg: str, fx_node: Optional[torch.fx.Node] = None
+        self, orig_expr: SympyBoolean, msg: str, fx_node: torch.fx.Node | None = None
     ) -> bool:
         """
         Adds a guard that orig_expr is True if we can or fall back to adding an assert
@@ -8192,7 +8542,7 @@ def run_node(self, n: torch.fx.Node) -> Result:
         return result
 
 
-def _find_user_code_frame() -> Optional[types.FrameType]:
+def _find_user_code_frame() -> types.FrameType | None:
     frame = inspect.currentframe()
     while frame is not None:
         if not frame.f_code.co_filename.startswith(
@@ -8357,6 +8707,8 @@ def _remove_effect_token_unbacked_bindings(
 # that we apply is unbacked renaming.
 def _get_placeholder_expr(sym_node: SymNode) -> sympy.Expr:
     shape_env = sym_node.shape_env
+    if shape_env is None:
+        raise AssertionError("shape_env is required for _get_placeholder_expr")
     result = sym_node._expr
     if result in shape_env.unbacked_renamings:
         return shape_env.unbacked_renamings[result]
diff --git a/torch/fx/experimental/unification/multipledispatch/dispatcher.py b/torch/fx/experimental/unification/multipledispatch/dispatcher.py
index 55ab9ea12e297..6ecb4502b978e 100644
--- a/torch/fx/experimental/unification/multipledispatch/dispatcher.py
+++ b/torch/fx/experimental/unification/multipledispatch/dispatcher.py
@@ -29,16 +29,19 @@ class MDNotImplementedError(NotImplementedError):
 
 
 def ambiguity_warn(dispatcher, ambiguities):
-    """Raise warning when ambiguity is detected
+    """Raise warning when ambiguity is detected.
+
     Parameters
     ----------
     dispatcher : Dispatcher
         The dispatcher on which the ambiguity was detected
     ambiguities : set
         Set of type signature pairs that are ambiguous within this dispatcher
-    See Also:
-        Dispatcher.add
-        warning_text
+
+    See Also
+    --------
+    Dispatcher.add
+    warning_text
     """
     warn(warning_text(dispatcher.name, ambiguities), AmbiguityWarning)
 
@@ -62,21 +65,23 @@ def restart_ordering(on_ambiguity=ambiguity_warn):
 
 def variadic_signature_matches_iter(types, full_signature):
     """Check if a set of input types matches a variadic signature.
+
     Notes
     -----
     The algorithm is as follows:
-    Initialize the current signature to the first in the sequence
-    For each type in `types`:
-        If the current signature is variadic
-            If the type matches the signature
-                yield True
-            Else
-                Try to get the next signature
-                If no signatures are left we can't possibly have a match
-                    so yield False
-        Else
-            yield True if the type matches the current signature
-            Get the next signature
+
+    Initialize the current signature to the first in the sequence.
+    For each type in ``types``:
+
+    - If the current signature is variadic
+
+      - If the type matches the signature, yield True
+      - Else, try to get the next signature.
+        If no signatures are left we can't possibly have a match,
+        so yield False.
+
+    - Else, yield True if the type matches the current signature.
+      Get the next signature.
     """
     sigiter = iter(full_signature)
     sig = next(sigiter)
diff --git a/torch/fx/experimental/unification/multipledispatch/utils.py b/torch/fx/experimental/unification/multipledispatch/utils.py
index 4f005773514db..f89d31aaef25a 100644
--- a/torch/fx/experimental/unification/multipledispatch/utils.py
+++ b/torch/fx/experimental/unification/multipledispatch/utils.py
@@ -65,11 +65,14 @@ def _toposort(edges):
 
 
 def reverse_dict(d):
-    """Reverses direction of dependence dict
+    """Reverses direction of dependence dict.
+
     >>> d = {"a": (1, 2), "b": (2, 3), "c": ()}
     >>> reverse_dict(d)  # doctest: +SKIP
     {1: ('a',), 2: ('a', 'b'), 3: ('b',)}
-    :note: dict order are not deterministic. As we iterate on the
+
+    .. note::
+        dict order are not deterministic. As we iterate on the
         input dict, it make the output of this function depend on the
         dict order. So this function output order should be considered
         as undeterministic.
diff --git a/torch/fx/experimental/unification/utils.py b/torch/fx/experimental/unification/utils.py
index f9fa8778b31f9..8f38ab7f23d06 100644
--- a/torch/fx/experimental/unification/utils.py
+++ b/torch/fx/experimental/unification/utils.py
@@ -67,11 +67,14 @@ def _toposort(edges):
 
 
 def reverse_dict(d):
-    """Reverses direction of dependence dict
+    """Reverses direction of dependence dict.
+
     >>> d = {"a": (1, 2), "b": (2, 3), "c": ()}
     >>> reverse_dict(d)  # doctest: +SKIP
     {1: ('a',), 2: ('a', 'b'), 3: ('b',)}
-    :note: dict order are not deterministic. As we iterate on the
+
+    .. note::
+        dict order are not deterministic. As we iterate on the
         input dict, it make the output of this function depend on the
         dict order. So this function output order should be considered
         as undeterministic.
diff --git a/torch/fx/experimental/validator.py b/torch/fx/experimental/validator.py
index d0db1fd4ec00b..34046d1958580 100644
--- a/torch/fx/experimental/validator.py
+++ b/torch/fx/experimental/validator.py
@@ -6,7 +6,7 @@
 import operator
 from collections.abc import Callable
 from dataclasses import dataclass
-from typing import Any, Optional, Union
+from typing import Any
 
 import sympy
 
@@ -327,7 +327,7 @@ def wrapper(*args):
             torch.sym_max: lift(ops.max),
             torch.sym_min: lift(ops.min),
             torch.sym_sum: lift(ops.sym_sum),
-            torch.sym_ite: lift(lambda b, t, f: t if b else f),
+            torch.sym_ite: lift(lambda b, t, f: z3.If(b, t, f)),
             torch._sym_sqrt: lift(ops.sqrt),  # type: ignore[attr-defined]
             # Not lifted because we only use this function as a
             # marker for adding the expression as validator input.
@@ -562,7 +562,7 @@ def add_target_expr(self, e: "sympy.logic.boolalg.Boolean") -> None:
                 log.debug("add target guard: %s", z3str(z3expr))
             self._target_exprs.add(z3expr)
 
-        def add_assertion(self, e: Union[z3.BoolRef, sympy.Basic]) -> None:
+        def add_assertion(self, e: z3.BoolRef | sympy.Basic) -> None:
             if isinstance(e, sympy.Basic):
                 self._check_freesymbols(e)
                 ref = self.to_z3_boolean_expr(e)
@@ -784,8 +784,8 @@ def new_with_shape_env(shape_env: ShapeEnv, fake) -> Any:
 
     # Checks whether the given shape_env fails when produce_guards is called.
     def check_shapeenv_fails(
-        shape_env: ShapeEnv, tracked_fakes: Optional[list[Any]]
-    ) -> Optional[ValidationException]:
+        shape_env: ShapeEnv, tracked_fakes: list[Any] | None
+    ) -> ValidationException | None:
         if tracked_fakes is None:
             raise AssertionError("tracked_fakes is None")
         try:
@@ -803,7 +803,7 @@ def check_shapeenv_fails(
 
     # Checks whether the ShapeEnv reconstructed by replaying the events until
     # node is created fails when produce_guards is called.
-    def check_node_fails(node: torch.fx.Node) -> Optional[ValidationException]:
+    def check_node_fails(node: torch.fx.Node) -> ValidationException | None:
         number = node.meta[SHAPEENV_EVENT_KEY]
         # Reconstruct shape_env until the event at event_number.
         shape_env = replay_shape_env_events(events[: number + 1])
diff --git a/torch/fx/graph.py b/torch/fx/graph.py
index d9da8ab585f6c..665a6ddd74094 100644
--- a/torch/fx/graph.py
+++ b/torch/fx/graph.py
@@ -160,7 +160,7 @@ def __init__(self):
         self._used_names: set[str] = set()
         self._base_count: dict[str, int] = {}
 
-    def create_name(self, candidate: str, obj: Optional[Any]) -> str:
+    def create_name(self, candidate: str, obj: Any | None) -> str:
         """Create a unique name.
 
         Arguments:
@@ -190,7 +190,10 @@ def create_name(self, candidate: str, obj: Optional[Any]) -> str:
 
         base, num = match.group(1, 2)
         if num is None or candidate in self._used_names:
-            num = self._base_count.get(candidate, 0)
+            # Look up `base` to match the key used in the store on line below;
+            # using `candidate` misses when it has a numeric suffix, making
+            # the while-loop quadratic.
+            num = self._base_count.get(base, 0)
             if _illegal_names.get(candidate, obj) is not obj:
                 num += 1
                 candidate = f"{base}_{num}"
@@ -237,7 +240,7 @@ class PythonCode:
     globals: dict[str, Any]
     # Optional mapping from the forward function's line number to
     # node index. Line number starts at the prologue (i.e. forward()).
-    _lineno_map: Optional[dict[int, Optional[int]]]
+    _lineno_map: dict[int, int | None] | None
     # The line number of prologue in fn_code
     _prologue_start: int = 0
 
@@ -291,7 +294,7 @@ class _PyTreeInfo(NamedTuple):
 
     orig_args: list[str]
     in_spec: pytree.TreeSpec
-    out_spec: Optional[pytree.TreeSpec]
+    out_spec: pytree.TreeSpec | None
 
 
 @dataclass(frozen=True)
@@ -311,7 +314,7 @@ def get_summary_str(self):
 
 # get File:lineno code from stack_trace
 def _parse_stack_trace(
-    stack_trace: str, filter_fn: Optional[Callable[[str, str, str], bool]] = None
+    stack_trace: str, filter_fn: Callable[[str, str, str], bool] | None = None
 ):
     if stack_trace is None:
         return None
@@ -341,7 +344,7 @@ class CodeGen:
     _sym_repr: Callable[["torch.types.PySymType"], str] = lambda x: repr(x)
 
     def __init__(self):
-        self._body_transformer: Optional[TransformCodeFunc] = None
+        self._body_transformer: TransformCodeFunc | None = None
         self._func_name: str = "forward"
 
     def _format_multiline_args(self, args: list[str]) -> str:
@@ -360,15 +363,20 @@ def _get_delimiters(self, container) -> tuple[str, str]:
         """Helper to get opening and closing delimiters for containers."""
         return ("(", ")") if isinstance(container, tuple) else ("[", "]")
 
-    def _format_multiline_container(self, items, descs=None, prefix="") -> str:
+    def _format_multiline_container(
+        self, items, descs=None, prefix="", repr_fn=None
+    ) -> str:
         """Helper to format containers (lists/tuples) in multiline format."""
         ldelim, rdelim = self._get_delimiters(items)
         desc_trailers = self._get_desc_trailers(items, descs)
+        if repr_fn is None:
+            repr_fn = repr
 
         return (
             f"{prefix}{ldelim}\n"
             + "".join(
-                f"    {item},{trailer}\n" for item, trailer in zip(items, desc_trailers)
+                f"    {repr_fn(item)},{trailer}\n"
+                for item, trailer in zip(items, desc_trailers)
             )
             + f"{rdelim}"
         )
@@ -411,16 +419,24 @@ def gen_fn_def(
             return f"def {self._func_name}({', '.join(free_vars)}){maybe_return_annotation}:"
 
     def generate_output(
-        self, output_args: Argument, *, descs: Optional[Any] = None
+        self,
+        output_args: Argument,
+        *,
+        descs: Any | None = None,
+        repr_fn: Any | None = None,
     ) -> str:
         """
         Given the output arguments, generates the return statement of the FX function.
         Note: The returned statement should not be indented.
         """
+        if repr_fn is None:
+            repr_fn = repr
         if descs is not None and isinstance(output_args, (list, tuple)):
-            return self._format_multiline_container(output_args, descs, "return ")
+            return self._format_multiline_container(
+                output_args, descs, "return ", repr_fn=repr_fn
+            )
         else:
-            return f"return {repr(output_args)}"
+            return f"return {repr_fn(output_args)}"
 
     def process_inputs(self, *args: Any) -> Any:
         """
@@ -461,7 +477,7 @@ def _gen_python_code(
         # Render each argument on its own line
         expanded_def: bool = False,
         record_func: bool = False,
-        additional_meta: Optional[list[str]] = None,
+        additional_meta: list[str] | None = None,
     ) -> PythonCode:
         free_vars: list[str] = []
         body: list[str] = []
@@ -897,6 +913,7 @@ def _tensor_annotation(t: torch.Tensor) -> str:
                         self.generate_output,
                         node.args[0],
                         descs=desc if expanded_def else None,
+                        repr_fn=_get_repr,
                     )
                 )
                 return
@@ -958,7 +975,7 @@ def _tensor_annotation(t: torch.Tensor) -> str:
         )
 
         # remove counter and generate lineno to node index mapping
-        lineno_map: dict[int, Optional[int]] = {}
+        lineno_map: dict[int, int | None] = {}
         prologue_len = prologue.count("\n") + 1
         new_lines: list[str] = []
         cur_idx = None
@@ -1136,21 +1153,26 @@ def gen_fn_def(
             fn_definition += self.gen_var_bindings(fn_args, free_vars, expanded_def)
         return fn_definition
 
-    def generate_output(self, output_args, *, descs: Optional[Any] = None):
+    def generate_output(
+        self, output_args, *, descs: Any | None = None, repr_fn: Any | None = None
+    ):
+        if repr_fn is None:
+            repr_fn = repr
         if self.pytree_info and self.pytree_info.out_spec:
             if descs is not None and isinstance(output_args, (list, tuple)):
                 return (
                     self._format_multiline_container(
-                        output_args, descs, "return pytree.tree_unflatten("
+                        output_args,
+                        descs,
+                        "return pytree.tree_unflatten(",
+                        repr_fn=repr_fn,
                     )
                     + ", self._out_spec)"
                 )
             else:
-                return (
-                    f"return pytree.tree_unflatten({repr(output_args)}, self._out_spec)"
-                )
+                return f"return pytree.tree_unflatten({repr_fn(output_args)}, self._out_spec)"
         else:
-            return super().generate_output(output_args, descs=descs)
+            return super().generate_output(output_args, descs=descs, repr_fn=repr_fn)
 
 
 class _ExportCodeGen(_PyTreeCodeGen):
@@ -1160,7 +1182,7 @@ def __init__(
         in_shuffle_graph: "GraphModule",
         out_shuffle_graph: "GraphModule",
         tree_leaf_names: list[str],
-        root: Optional[torch.nn.Module],
+        root: torch.nn.Module | None,
     ):
         super().__init__(pytree_info)
         self.in_shuffle_graph = in_shuffle_graph
@@ -1205,11 +1227,11 @@ class _FindNodesLookupTable:
     """
 
     def __init__(self):
-        self.table: dict[tuple[str, Optional[Target]], dict[Node, None]] = defaultdict(
+        self.table: dict[tuple[str, Target | None], dict[Node, None]] = defaultdict(
             dict
         )
 
-    def _key(self, node) -> tuple[str, Optional[Target]]:
+    def _key(self, node) -> tuple[str, Target | None]:
         return (node.op, node.target if node.op == "call_function" else None)
 
     def __contains__(self, node) -> bool:
@@ -1287,8 +1309,8 @@ def forward(self, x):
     def __init__(
         self,
         owning_module: Optional["GraphModule"] = None,
-        tracer_cls: Optional[type["Tracer"]] = None,
-        tracer_extras: Optional[dict[str, Any]] = None,
+        tracer_cls: type["Tracer"] | None = None,
+        tracer_extras: dict[str, Any] | None = None,
     ):
         """
         Construct an empty Graph.
@@ -1364,7 +1386,7 @@ def find_nodes(
     @compatibility(is_backward_compatible=True)
     def graph_copy(
         self, g: "Graph", val_map: dict[Node, Node], return_output_node=False
-    ) -> "Optional[Argument]":
+    ) -> "Argument | None":
         """
         Copy all nodes from a given graph into ``self``.
 
@@ -1422,10 +1444,10 @@ def create_node(
         self,
         op: str,
         target: "Target",
-        args: Optional[tuple["Argument", ...]] = None,
-        kwargs: Optional[dict[str, "Argument"]] = None,
-        name: Optional[str] = None,
-        type_expr: Optional[Any] = None,
+        args: tuple["Argument", ...] | None = None,
+        kwargs: dict[str, "Argument"] | None = None,
+        name: str | None = None,
+        type_expr: Any | None = None,
     ) -> Node:
         """
         Create a ``Node`` and add it to the ``Graph`` at the current insert-point.
@@ -1535,7 +1557,7 @@ def erase_node(self, to_erase: Node) -> None:
         )
 
     @compatibility(is_backward_compatible=True)
-    def inserting_before(self, n: Optional[Node] = None):
+    def inserting_before(self, n: Node | None = None):
         """Set the point at which create_node and companion methods will insert into the graph.
         When used within a 'with' statement, this will temporary set the insert point and
         then restore it when the with statement exits::
@@ -1560,7 +1582,7 @@ def inserting_before(self, n: Optional[Node] = None):
         return _InsertPoint(self, n.prepend)
 
     @compatibility(is_backward_compatible=True)
-    def inserting_after(self, n: Optional[Node] = None):
+    def inserting_after(self, n: Node | None = None):
         """Set the point at which create_node and companion methods will insert into the graph.
         When used within a 'with' statement, this will temporary set the insert point and
         then restore it when the with statement exits::
@@ -1588,7 +1610,7 @@ def inserting_after(self, n: Optional[Node] = None):
     def placeholder(
         self,
         name: str,
-        type_expr: Optional[Any] = None,
+        type_expr: Any | None = None,
         default_value: Any = inspect.Signature.empty,
     ) -> Node:
         """
@@ -1618,7 +1640,7 @@ def placeholder(
         return self.create_node("placeholder", name, args=args, type_expr=type_expr)
 
     @compatibility(is_backward_compatible=True)
-    def get_attr(self, qualified_name: str, type_expr: Optional[Any] = None) -> Node:
+    def get_attr(self, qualified_name: str, type_expr: Any | None = None) -> Node:
         """
         Insert a ``get_attr`` node into the Graph. A ``get_attr`` ``Node`` represents the
         fetch of an attribute from the ``Module`` hierarchy.
@@ -1689,9 +1711,9 @@ def _get_attr_reference_exists(
     def call_module(
         self,
         module_name: str,
-        args: Optional[tuple["Argument", ...]] = None,
-        kwargs: Optional[dict[str, "Argument"]] = None,
-        type_expr: Optional[Any] = None,
+        args: tuple["Argument", ...] | None = None,
+        kwargs: dict[str, "Argument"] | None = None,
+        type_expr: Any | None = None,
     ) -> Node:
         """
         Insert a ``call_module`` ``Node`` into the ``Graph``. A ``call_module`` node
@@ -1739,9 +1761,9 @@ def call_module(
     def call_method(
         self,
         method_name: str,
-        args: Optional[tuple["Argument", ...]] = None,
-        kwargs: Optional[dict[str, "Argument"]] = None,
-        type_expr: Optional[Any] = None,
+        args: tuple["Argument", ...] | None = None,
+        kwargs: dict[str, "Argument"] | None = None,
+        type_expr: Any | None = None,
     ) -> Node:
         """
         Insert a ``call_method`` ``Node`` into the ``Graph``. A ``call_method`` node
@@ -1778,10 +1800,10 @@ def call_method(
     def call_function(
         self,
         the_function: Callable[..., Any],
-        args: Optional[tuple["Argument", ...]] = None,
-        kwargs: Optional[dict[str, "Argument"]] = None,
-        type_expr: Optional[Any] = None,
-        name: Optional[str] = None,
+        args: tuple["Argument", ...] | None = None,
+        kwargs: dict[str, "Argument"] | None = None,
+        type_expr: Any | None = None,
+        name: str | None = None,
     ) -> Node:
         """
         Insert a ``call_function`` ``Node`` into the ``Graph``. A ``call_function`` node
@@ -1854,7 +1876,7 @@ def node_copy(
         return result_node
 
     @compatibility(is_backward_compatible=True)
-    def output(self, result: "Argument", type_expr: Optional[Any] = None):
+    def output(self, result: "Argument", type_expr: Any | None = None):
         """
         Insert an ``output`` ``Node`` into the ``Graph``. An ``output`` node represents
         a ``return`` statement in Python code. ``result`` is the value that should
@@ -1876,7 +1898,7 @@ def output(self, result: "Argument", type_expr: Optional[Any] = None):
             op="output", target="output", args=(result,), type_expr=type_expr
         )
 
-    def _target_to_str(self, target: Optional[Target]) -> str:
+    def _target_to_str(self, target: Target | None) -> str:
         if callable(target):
             op = target.__name__
         else:
@@ -1899,7 +1921,7 @@ def python_code(
         colored: bool = False,
         expanded_def: bool = False,
         record_func: bool = False,
-        additional_meta: Optional[list[str]] = None,
+        additional_meta: list[str] | None = None,
     ) -> PythonCode:
         """
         Turn this ``Graph`` into valid Python code.
@@ -1982,7 +2004,7 @@ def _python_code(
         colored: bool = False,
         expanded_def: bool = False,
         record_func: bool = False,
-        additional_meta: Optional[list[str]] = None,
+        additional_meta: list[str] | None = None,
     ) -> PythonCode:
         return self._codegen._gen_python_code(
             self.nodes,
@@ -2049,7 +2071,7 @@ def lint(self):
         """
 
         # Check topo order
-        def check_arg(arg: Node, n: Optional[Node] = None) -> None:
+        def check_arg(arg: Node, n: Node | None = None) -> None:
             context_str = f" of Node '{n}' " if n else " "
             if arg.graph is not self:
                 raise RuntimeError(
@@ -2119,7 +2141,7 @@ def check_arg(arg: Node, n: Optional[Node] = None) -> None:
 
     @compatibility(is_backward_compatible=True)
     def eliminate_dead_code(
-        self, is_impure_node: Optional[Callable[[Node], bool]] = None
+        self, is_impure_node: Callable[[Node], bool] | None = None
     ) -> bool:
         """
         Remove all dead code from the graph, based on each node's number of
@@ -2213,7 +2235,7 @@ def set_codegen(self, codegen: CodeGen):
     @compatibility(is_backward_compatible=False)
     def on_generate_code(
         self,
-        make_transformer: Callable[[Optional[TransformCodeFunc]], TransformCodeFunc],
+        make_transformer: Callable[[TransformCodeFunc | None], TransformCodeFunc],
     ):
         """Register a transformer function when python code is generated
 
diff --git a/torch/fx/graph_module.py b/torch/fx/graph_module.py
index f08bf19b24b43..9ca2984fb9b85 100644
--- a/torch/fx/graph_module.py
+++ b/torch/fx/graph_module.py
@@ -11,7 +11,7 @@
 import warnings
 from collections.abc import Callable
 from pathlib import Path
-from typing import Any, Optional, Union
+from typing import Any
 
 import torch
 import torch.nn as nn
@@ -87,7 +87,7 @@ def cache(self, src: str, globals: dict[str, Any], co_fields=None):
 
     # Part of the loader protocol (PEP 302)
     # linecache will use this method when trying to find source code
-    def get_source(self, module_name) -> Optional[str]:
+    def get_source(self, module_name) -> str | None:
         if module_name in self.eval_cache:
             return self.eval_cache[module_name]
         return None
@@ -498,7 +498,7 @@ class GraphModuleImpl(cls):  # type: ignore[misc, valid-type]
     @compatibility(is_backward_compatible=True)
     def __init__(
         self,
-        root: Union[torch.nn.Module, dict[str, Any]],
+        root: torch.nn.Module | dict[str, Any],
         graph: Graph,
         class_name: str = "GraphModule",
     ):
@@ -634,7 +634,7 @@ def graph(self, g: Graph) -> None:
         self.recompile()
 
     @compatibility(is_backward_compatible=False)
-    def to_folder(self, folder: Union[str, os.PathLike], module_name: str = "FxModule"):
+    def to_folder(self, folder: str | os.PathLike, module_name: str = "FxModule"):
         """Dumps out module to ``folder`` with ``module_name`` so that it can be
         imported with ``from <folder> import <module_name>``
 
@@ -660,7 +660,7 @@ def __init__(self):
         super().__init__()
 """
 
-        def _gen_model_repr(module_name: str, module: torch.nn.Module) -> Optional[str]:
+        def _gen_model_repr(module_name: str, module: torch.nn.Module) -> str | None:
             safe_reprs = [
                 nn.Linear,
                 nn.Conv1d,
@@ -1061,7 +1061,7 @@ def print_readable(
         # but may result in less-readable output.
         fast_sympy_print: bool = False,
         expanded_def: bool = False,
-        additional_meta: Optional[list[str]] = None,
+        additional_meta: list[str] | None = None,
     ):
         """
         Return the Python code generated for current GraphModule and its children GraphModules.
diff --git a/torch/fx/interpreter.py b/torch/fx/interpreter.py
index 5d301f116716b..7cd4ba1776567 100644
--- a/torch/fx/interpreter.py
+++ b/torch/fx/interpreter.py
@@ -1,8 +1,8 @@
-# mypy: allow-untyped-defs
 import inspect
 import logging
+from collections.abc import Iterator
 from contextlib import contextmanager
-from typing import Any, Optional, TYPE_CHECKING, Union
+from typing import Any
 
 import torch
 import torch.fx.traceback as fx_traceback
@@ -19,15 +19,12 @@
 from .proxy import Proxy
 
 
-if TYPE_CHECKING:
-    from collections.abc import Iterator
-
 log = logging.getLogger(__name__)
 
 __all__ = ["Interpreter", "Transformer"]
 
 
-def _format_fx_node(n):
+def _format_fx_node(n: Node) -> str:
     """
     Format a torch.fx.Node into a human-readable string for debug logging.
 
@@ -117,8 +114,8 @@ def __init__(
         self,
         module: torch.nn.Module,
         garbage_collect_values: bool = True,
-        graph: Optional[Graph] = None,
-    ):
+        graph: Graph | None = None,
+    ) -> None:
         self.module = module
         self.submodules = dict(self.module.named_modules())
         if graph is not None:
@@ -138,7 +135,7 @@ def __init__(
             node_to_last_use: dict[Node, Node] = {}
             self.user_to_last_uses: dict[Node, list[Node]] = {}
 
-            def register_last_uses(n: Node, user: Node):
+            def register_last_uses(n: Node, user: Node) -> None:
                 if n not in node_to_last_use:
                     node_to_last_use[n] = user
                     self.user_to_last_uses.setdefault(user, []).append(n)
@@ -150,8 +147,8 @@ def register_last_uses(n: Node, user: Node):
     @compatibility(is_backward_compatible=True)
     def run(
         self,
-        *args,
-        initial_env: Optional[dict[Node, Any]] = None,
+        *args: Any,
+        initial_env: dict[Node, Any] | None = None,
         enable_io_processing: bool = True,
     ) -> Any:
         """
@@ -240,7 +237,7 @@ def run(
                 )
 
     @compatibility(is_backward_compatible=True)
-    def boxed_run(self, args_list):
+    def boxed_run(self, args_list: list[Any]) -> Any:
         """
         Run `module` via interpretation and return the result.  This uses the "boxed"
         calling convention, where you pass a list of arguments, which will be cleared
@@ -267,7 +264,7 @@ def boxed_run(self, args_list):
         return self.run(initial_env=env)
 
     @contextmanager
-    def _set_current_node(self, node):
+    def _set_current_node(self, node: Node) -> Iterator[None]:
         with fx_traceback.set_current_meta(
             node, f"Interpreter_{self.__class__.__name__}"
         ):
@@ -452,7 +449,7 @@ def output(
 
     # Helper methods
     @compatibility(is_backward_compatible=True)
-    def fetch_attr(self, target: str):
+    def fetch_attr(self, target: str) -> Any:
         """
         Fetch an attribute from the ``Module`` hierarchy of ``self.module``.
 
@@ -473,7 +470,9 @@ def fetch_attr(self, target: str):
         return attr_itr
 
     @compatibility(is_backward_compatible=True)
-    def fetch_args_kwargs_from_env(self, n: Node) -> tuple[tuple, dict]:
+    def fetch_args_kwargs_from_env(
+        self, n: Node
+    ) -> tuple[tuple[Any, ...], dict[str, Any]]:
         """
         Fetch the concrete values of ``args`` and ``kwargs`` of node ``n``
         from the current execution environment.
@@ -568,18 +567,18 @@ def fn(x):
     """
 
     @compatibility(is_backward_compatible=True)
-    def __init__(self, module):
+    def __init__(self, module: GraphModule) -> None:
         super().__init__(module)
         self.new_graph = Graph()
         self.new_graph.set_codegen(module.graph._codegen)
 
         class TransformerTracer(Tracer):
-            def __init__(self, graph: Graph):
+            def __init__(self, graph: Graph) -> None:
                 super().__init__()
                 self.graph = graph
                 self.tensor_attrs: dict[torch.Tensor, str] = {}  # type: ignore[assignment]
 
-            def is_leaf_module(self, _, __) -> bool:
+            def is_leaf_module(self, _: torch.nn.Module, __: str) -> bool:
                 return True
 
         self.tracer = TransformerTracer(self.new_graph)
@@ -655,7 +654,7 @@ def transform(self) -> GraphModule:
             result = super().run(enable_io_processing=False)
         if result is not None:
 
-            def strip_proxy(a: Union[Argument, Proxy]) -> Any:
+            def strip_proxy(a: Argument | Proxy) -> Any:
                 return a.node if isinstance(a, Proxy) else a
 
             new_output_node = self.new_graph.output(map_aggregate(result, strip_proxy))
diff --git a/torch/fx/node.py b/torch/fx/node.py
index 683eaa6d5c030..ecb9fbd5ab27b 100644
--- a/torch/fx/node.py
+++ b/torch/fx/node.py
@@ -29,7 +29,7 @@
 
 log = logging.getLogger(__name__)
 
-BaseArgumentTypes = Union[
+BaseArgumentTypes = Union[  # noqa: UP007
     str,
     int,
     float,
@@ -47,9 +47,9 @@
 ]
 base_types = typing.get_args(BaseArgumentTypes)
 
-Target: TypeAlias = Union[Callable[..., Any], str]
+Target: TypeAlias = Callable[..., Any] | str
 
-Argument = Optional[
+Argument = Optional[  # noqa: UP007, UP045
     Union[
         tuple["Argument", ...],
         Sequence["Argument"],
@@ -185,6 +185,7 @@ def _get_qualified_name(func: Callable[..., Any]) -> str:
     ):
         return f"torch.Tensor.{func.__name__}"
     name = func.__name__
+
     if name == "<lambda>":
         # For lambdas, try to get their defining name in the module
         try:
@@ -291,10 +292,10 @@ class Node(_NodeBase):
     # generated function return type. (Note this is a special case. ``return``
     # does not produce a value, it's more of a notation. Thus, this value
     # describes the type of args[0] in the ``return`` node.
-    type: Optional[Any]
+    type: Any | None
     _sort_key: Any
     # If set, use this fn to print this node
-    _repr_fn: Optional[Callable[["Node"], str]]
+    _repr_fn: Callable[["Node"], str] | None
     # Dictionary to store metadata passes need to do their
     # transformations. This metadata is preserved across node copies
     meta: dict[str, Any]
@@ -308,7 +309,7 @@ def __init__(
         target: "Target",
         args: tuple["Argument", ...],
         kwargs: dict[str, "Argument"],
-        return_type: Optional[Any] = None,
+        return_type: Any | None = None,
     ) -> None:
         """
         Instantiate an instance of ``Node``. Note: most often, you want to use the
@@ -543,7 +544,7 @@ def update_kwarg(self, key: str, arg: Argument) -> None:
         self.kwargs = {**self.kwargs, key: arg}
 
     @property
-    def stack_trace(self) -> Optional[str]:
+    def stack_trace(self) -> str | None:
         """
         Return the Python stack trace that was recorded during tracing, if any.
         When traced with fx.Tracer, this property is usually populated by
@@ -557,7 +558,7 @@ def stack_trace(self) -> Optional[str]:
         return self.meta.get("stack_trace", None)
 
     @stack_trace.setter
-    def stack_trace(self, trace: Optional[str]) -> None:
+    def stack_trace(self, trace: str | None) -> None:
         self.meta["stack_trace"] = trace
 
     def __repr__(self) -> str:
@@ -593,11 +594,11 @@ def _pretty_print_target(target: object) -> str:
     @compatibility(is_backward_compatible=True)
     def format_node(
         self,
-        placeholder_names: Optional[list[str]] = None,
-        maybe_return_typename: Optional[list[str]] = None,
+        placeholder_names: list[str] | None = None,
+        maybe_return_typename: list[str] | None = None,
         *,
         include_tensor_metadata: bool = False,
-    ) -> Optional[str]:
+    ) -> str | None:
         """
         Return a descriptive string representation of ``self``.
 
@@ -654,7 +655,7 @@ def format_node(
             return f"return {self.args[0]}"
         else:
 
-            def stringify_shape(shape: Iterable) -> str:
+            def stringify_shape(shape: Iterable[Any]) -> str:
                 return f"[{', '.join([str(x) for x in shape])}]"
 
             meta_val = self.meta.get(
@@ -691,7 +692,7 @@ def stringify_shape(shape: Iterable) -> str:
     def replace_all_uses_with(
         self,
         replace_with: "Node",
-        delete_user_cb: Optional[Callable[["Node"], bool]] = None,
+        delete_user_cb: Callable[["Node"], bool] | None = None,
         *,
         propagate_meta: bool = False,
     ) -> list["Node"]:
@@ -785,10 +786,10 @@ def is_impure(self, impure_random: bool = True) -> bool:
     def normalized_arguments(
         self,
         root: torch.nn.Module,
-        arg_types: Optional[tuple[Any]] = None,
-        kwarg_types: Optional[dict[str, Any]] = None,
+        arg_types: tuple[Any] | None = None,
+        kwarg_types: dict[str, Any] | None = None,
         normalize_to_only_use_kwargs: bool = False,
-    ) -> Optional[ArgsKwargsPair]:
+    ) -> ArgsKwargsPair | None:
         """
         Returns normalized arguments to Python targets. This means that
         `args/kwargs` will be matched up to the module/functional's
diff --git a/torch/fx/operator_schemas.py b/torch/fx/operator_schemas.py
index 9073b3f26f056..532c8543ba888 100644
--- a/torch/fx/operator_schemas.py
+++ b/torch/fx/operator_schemas.py
@@ -1,4 +1,3 @@
-# mypy: allow-untyped-defs
 import enum
 import inspect
 import numbers
@@ -6,11 +5,12 @@
 import typing
 import warnings
 from collections.abc import Callable
-from typing import Any, cast, NamedTuple, Optional, TYPE_CHECKING
+from typing import Any, cast, Literal, NamedTuple, overload, TYPE_CHECKING
 
 import torch
 from torch._jit_internal import boolean_dispatched
 from torch._ops import OpOverload, OpOverloadPacket
+from torch.utils._inspect import _fast_bind
 
 from ._compatibility import compatibility
 
@@ -39,18 +39,18 @@ class ArgsKwargsPair(NamedTuple):
     kwargs: dict[str, Any]
 
 
-_manual_overrides: dict[Callable, list[inspect.Signature]] = {}
+_manual_overrides: dict[Callable[..., Any], list[inspect.Signature]] = {}
 
 
-def _nonzero_schemas():
+def _nonzero_schemas() -> list[inspect.Signature]:
     signatures = []
 
-    def nonzero(self):
+    def nonzero(self: torch.Tensor) -> None:
         pass
 
     signatures.append(inspect.signature(nonzero))
 
-    def nonzero(self, *, as_tuple: bool):  # type: ignore[no-redef]
+    def nonzero(self: torch.Tensor, *, as_tuple: bool) -> None:  # type: ignore[no-redef]
         pass
 
     signatures.append(inspect.signature(nonzero))
@@ -62,7 +62,7 @@ def nonzero(self, *, as_tuple: bool):  # type: ignore[no-redef]
 
 
 class _FakeGlobalNamespace:
-    def __getattr__(self, name):
+    def __getattr__(self, name: str) -> types.ModuleType:
         if name == "torch":
             return torch
         raise RuntimeError("Expected a torch namespace lookup")
@@ -171,24 +171,26 @@ def _torchscript_schema_to_signature(
 
 @compatibility(is_backward_compatible=False)
 def check_for_mutable_operation(
-    target: Callable, args: tuple["Argument", ...], kwargs: dict[str, "Argument"]
-):
+    target: Callable[..., Any],
+    args: tuple["Argument", ...],
+    kwargs: dict[str, "Argument"],
+) -> None:
     signatures, schemas = get_signature_for_torch_op(target, return_schemas=True)
 
     if signatures and schemas:
-        matched_schemas = []
+        matched_schemas: list[tuple[inspect.Signature, torch._C.FunctionSchema]] = []
 
         # Iterate through all of the schema until we find one that matches
         # If one matches, populate `new_args_and_kwargs` with the new args/kwargs
         # values. If none matches, `new_args_and_kwargs` will be None
         for candidate_signature, schema in zip(signatures, schemas):
             try:
-                candidate_signature.bind(*args, **kwargs)
+                _fast_bind(candidate_signature, *args, **kwargs)
                 matched_schemas.append((candidate_signature, schema))
             except TypeError:
                 continue
 
-        def throw_if_mutable(schema):
+        def throw_if_mutable(schema: torch._C.FunctionSchema) -> None:
             if schema.is_mutable:
                 raise RuntimeError(
                     f"Tried to trace mutable operation {schema}. FX only supports functional "
@@ -209,8 +211,26 @@ def throw_if_mutable(schema):
             pass
 
 
+@overload
+def get_signature_for_torch_op(
+    op: Callable[..., Any], return_schemas: Literal[True]
+) -> tuple[list[inspect.Signature] | None, list[torch._C.FunctionSchema] | None]: ...
+
+
+@overload
+def get_signature_for_torch_op(
+    op: Callable[..., Any], return_schemas: Literal[False] = ...
+) -> list[inspect.Signature] | None: ...
+
+
 @compatibility(is_backward_compatible=False)
-def get_signature_for_torch_op(op: Callable, return_schemas: bool = False):
+def get_signature_for_torch_op(
+    op: Callable[..., Any], return_schemas: bool = False
+) -> (
+    list[inspect.Signature]
+    | tuple[list[inspect.Signature] | None, list[torch._C.FunctionSchema] | None]
+    | None
+):
     """
     Given an operator on the `torch` namespace, return a list of `inspect.Signature`
     objects corresponding to the overloads of that op.. May return `None` if a signature
@@ -245,7 +265,7 @@ def get_signature_for_torch_op(op: Callable, return_schemas: bool = False):
 
 
 @compatibility(is_backward_compatible=False)
-def create_type_hint(x):
+def create_type_hint(x: object) -> object:
     """
     Produces a type hint for the given argument.
 
@@ -262,12 +282,12 @@ def create_type_hint(x):
             # todo(chilli): Figure out the right way for mypy to handle this
             if isinstance(x, list):
 
-                def ret_type(x):
+                def ret_type(x: Any) -> Any:
                     return list[x]  # type: ignore[valid-type]
 
             else:
 
-                def ret_type(x):
+                def ret_type(x: Any) -> Any:
                     return tuple[x, ...]  # type: ignore[valid-type]
 
             if len(x) == 0:
@@ -290,7 +310,7 @@ def ret_type(x):
 
 
 @compatibility(is_backward_compatible=False)
-def type_matches(signature_type: Any, argument_type: Any):
+def type_matches(signature_type: Any, argument_type: Any) -> bool:
     sig_origin_type = getattr(signature_type, "__origin__", signature_type)
 
     if signature_type is argument_type:
@@ -317,11 +337,11 @@ def type_matches(signature_type: Any, argument_type: Any):
         if getattr(argument_type, "__origin__", None) is list:
             return issubclass(argument_type.__args__[0], sig_el_type)
 
-        def is_homogeneous_tuple(t):
-            if getattr(t, "__origin__", None) is not tuple:
+        def is_homogeneous_tuple(t: object) -> bool:
+            if typing.get_origin(t) is not tuple:
                 return False
-            contained = t.__args__
-            if t.__args__ == ((),):  # Tuple[()].__args__ == ((),) for some reason
+            contained = typing.get_args(t)
+            if contained == ((),):  # Tuple[()].__args__ == ((),) for some reason
                 return True
             return all((c is Ellipsis) or issubclass(c, sig_el_type) for c in contained)
 
@@ -342,11 +362,11 @@ def is_homogeneous_tuple(t):
 
 @compatibility(is_backward_compatible=False)
 def _normalize_function_or_error(
-    target: Callable,
+    target: Callable[..., Any],
     args: tuple[Any, ...],
-    kwargs: Optional[dict[str, Any]] = None,
-    arg_types: Optional[tuple[Any]] = None,
-    kwarg_types: Optional[dict[str, Any]] = None,
+    kwargs: dict[str, Any] | None = None,
+    arg_types: tuple[Any] | None = None,
+    kwarg_types: dict[str, Any] | None = None,
     normalize_to_only_use_kwargs: bool = False,
 ) -> ArgsKwargsPair:
     """
@@ -366,13 +386,13 @@ def _normalize_function_or_error(
 
 @compatibility(is_backward_compatible=False)
 def normalize_function(
-    target: Callable,
+    target: Callable[..., Any],
     args: tuple[Any, ...],
-    kwargs: Optional[dict[str, Any]] = None,
-    arg_types: Optional[tuple[Any]] = None,
-    kwarg_types: Optional[dict[str, Any]] = None,
+    kwargs: dict[str, Any] | None = None,
+    arg_types: tuple[Any] | None = None,
+    kwarg_types: dict[str, Any] | None = None,
     normalize_to_only_use_kwargs: bool = False,
-) -> Optional[ArgsKwargsPair]:
+) -> ArgsKwargsPair | None:
     """
     Returns normalized arguments to PyTorch functions. This means that
     `args/kwargs` will be matched up to the functional's
@@ -442,14 +462,14 @@ def normalize_function(
         if not callable(target):
             raise AssertionError(f"target must be callable, got {type(target)}")
         torch_op_schemas = get_signature_for_torch_op(target)
-        matched_schemas = []
+        matched_schemas: list[inspect.Signature] = []
         if torch_op_schemas:
             # Iterate through all of the schema until we find one that matches
             # If one matches, populate `new_args_and_kwargs` with the new args/kwargs
             # values. If none matches, `new_args_and_kwargs` will be None
             for candidate_signature in torch_op_schemas:
                 try:
-                    candidate_signature.bind(*args, **kwargs)
+                    _fast_bind(candidate_signature, *args, **kwargs)
                     matched_schemas.append(candidate_signature)
                 except TypeError:
                     continue
@@ -469,8 +489,8 @@ def normalize_function(
                     for candidate_signature in torch_op_schemas:
                         sig_matches = True
                         try:
-                            bound_types = candidate_signature.bind(
-                                *arg_types, **kwarg_types
+                            bound_types = _fast_bind(
+                                candidate_signature, *arg_types, **kwarg_types
                             )
                             for arg_name, arg_type in bound_types.arguments.items():
                                 param = candidate_signature.parameters[arg_name]
@@ -509,9 +529,9 @@ def normalize_module(
     root: torch.nn.Module,
     target: str,
     args: tuple[Any],
-    kwargs: Optional[dict[str, Any]] = None,
+    kwargs: dict[str, Any] | None = None,
     normalize_to_only_use_kwargs: bool = False,
-) -> Optional[ArgsKwargsPair]:
+) -> ArgsKwargsPair | None:
     """
     Returns normalized arguments to PyTorch modules. This means that
     `args/kwargs` will be matched up to the functional's
@@ -556,7 +576,7 @@ def _args_kwargs_to_normalized_args_kwargs(
     args: tuple[Any, ...],
     kwargs: dict[str, Any],
     normalize_to_only_use_kwargs: bool,
-) -> Optional[ArgsKwargsPair]:
+) -> ArgsKwargsPair | None:
     """
     Given a call target, args, and kwargs, return the arguments normalized into
     an ArgsKwargsPair, or None if the type signature is not supported by
@@ -589,7 +609,7 @@ def _args_kwargs_to_normalized_args_kwargs(
         if list(sig.parameters.keys()) != ["input", "from", "to", "generator"]:
             return None
 
-    bound_args = sig.bind(*args, **kwargs)
+    bound_args = _fast_bind(sig, *args, **kwargs)
     bound_args.apply_defaults()
 
     new_kwargs: dict[str, Any] = {}
diff --git a/torch/fx/passes/_tensorify_python_scalars.py b/torch/fx/passes/_tensorify_python_scalars.py
index 235274b82f415..3ccbac7bc9de3 100644
--- a/torch/fx/passes/_tensorify_python_scalars.py
+++ b/torch/fx/passes/_tensorify_python_scalars.py
@@ -2,7 +2,7 @@
 
 import logging
 import os
-from typing import Any, TYPE_CHECKING, Union
+from typing import Any, TYPE_CHECKING
 
 from sympy import Integer, Number, Symbol
 from sympy.logic.boolalg import BooleanAtom
@@ -111,7 +111,6 @@ def tensorify_python_scalars(
     Returns:
         None
     """
-    import sympy
 
     knob = True
     if (env := os.getenv("TENSORIFY_PYTHON_SCALARS")) is not None:
@@ -122,6 +121,23 @@ def tensorify_python_scalars(
     if not knob:
         return None
 
+    # This pass uses MetaProxy which relies on __torch_function__.
+    # DisableTorchFunctionSubclass may be active here (see #177088),
+    # so re-enable dispatch for MetaProxy ops.
+    with torch._C._EnableTorchFunction():
+        return _tensorify_impl(gm, shape_env, fake_mode)
+
+
+def _tensorify_impl(
+    gm: GraphModule,
+    shape_env: ShapeEnv,
+    fake_mode: fake_tensor.FakeTensorMode,
+) -> None:
+    """Helper fn in tensorify_python_scalars so the caller can wrap
+    with _EnableTorchFunction (#180906).
+    """
+    import sympy
+
     graph = gm.graph
     tracer = fx.proxy.GraphAppendingTracer(graph)
     expr_to_sym_proxy: dict[sympy.Expr, MetaProxy] = {}
@@ -155,7 +171,7 @@ def _sympy_interp(expr: sympy.Expr) -> MetaProxy:
         # cache constants, why not
         if isinstance(expr, (Integer, Number, BooleanAtom)):
             dtype = None
-            c: Union[bool, int, float]
+            c: bool | int | float
             if isinstance(expr, BooleanAtom):
                 dtype = torch.bool
                 c = bool(expr)
@@ -209,7 +225,8 @@ def _sympy_interp(expr: sympy.Expr) -> MetaProxy:
                 and node.op == "call_function"
                 and node.target is torch.ops.aten._local_scalar_dense.default
             ):
-                dtype = node.args[0].meta["val"].dtype
+                source_tensor = node.args[0].meta["val"]
+                dtype = source_tensor.dtype
 
                 if not isinstance(node.args[0], fx.Node):
                     raise AssertionError(f"Expected fx.Node, got {node.args[0]}")
@@ -227,6 +244,15 @@ def _sympy_interp(expr: sympy.Expr) -> MetaProxy:
                 expr_to_tensor_proxy[s] = MetaProxy(
                     node.args[0], tracer=tracer, fake_mode=fake_mode
                 )
+                if len(source_tensor.shape) != 0:
+                    # .item() always produces a scalar value, even when it is
+                    # called on a size-1 tensor with rank > 0. Preserve that 0-d
+                    # semantics before tensorifying the scalar expression so
+                    # later tensor math and autograd tangents do not keep an
+                    # accidental length-1 dimension.
+                    expr_to_tensor_proxy[s] = torch.ops.aten.reshape.default(
+                        expr_to_tensor_proxy[s], []
+                    )
                 # Upcast the float tensor to torch.float64 to avoid precision problem
                 expr_to_tensor_proxy[s] = torch.ops.prims.convert_element_type.default(
                     expr_to_tensor_proxy[s], torch.float64
@@ -337,7 +363,7 @@ def _sympy_interp(expr: sympy.Expr) -> MetaProxy:
                     ):
                         failed_tensorify_ops.update(str(node.target))
 
-                        log.info("Failed to tensorify %s", str(node.target))
+                        log.info("Failed to tensorify %s", node.target)
 
     # Now do one more pass that specializes all symfloats we didn't manage
     # to tensorify away.
diff --git a/torch/fx/passes/fake_tensor_prop.py b/torch/fx/passes/fake_tensor_prop.py
index a4df194946edc..45507632228ef 100644
--- a/torch/fx/passes/fake_tensor_prop.py
+++ b/torch/fx/passes/fake_tensor_prop.py
@@ -1,5 +1,4 @@
 # mypy: allow-untyped-defs
-from typing import Optional
 
 import torch.fx
 from torch._subclasses.fake_tensor import FakeTensor, FakeTensorMode
@@ -29,7 +28,7 @@ class FakeTensorProp(torch.fx.Interpreter):
     """
 
     def __init__(
-        self, module: torch.fx.GraphModule, mode: Optional[FakeTensorMode] = None
+        self, module: torch.fx.GraphModule, mode: FakeTensorMode | None = None
     ):
         super().__init__(module)
         if mode is None:
diff --git a/torch/fx/passes/graph_drawer.py b/torch/fx/passes/graph_drawer.py
index 0f0e2e59c50af..f94d4cd174f53 100644
--- a/torch/fx/passes/graph_drawer.py
+++ b/torch/fx/passes/graph_drawer.py
@@ -3,7 +3,7 @@
 import hashlib
 from itertools import chain
 from types import ModuleType
-from typing import Any, Optional, TYPE_CHECKING
+from typing import Any, TYPE_CHECKING
 
 import torch
 import torch.fx
@@ -19,7 +19,7 @@
 
     HAS_PYDOT = True
 else:
-    pydot: Optional[ModuleType]
+    pydot: ModuleType | None
     try:
         import pydot
 
@@ -83,7 +83,7 @@ def __init__(
             ignore_parameters_and_buffers: bool = False,
             skip_node_names_in_args: bool = True,
             parse_stack_trace: bool = False,
-            dot_graph_shape: Optional[str] = None,
+            dot_graph_shape: str | None = None,
             normalize_args: bool = False,
         ):
             self._name = name
@@ -498,7 +498,7 @@ def __init__(
                 ignore_parameters_and_buffers: bool = False,
                 skip_node_names_in_args: bool = True,
                 parse_stack_trace: bool = False,
-                dot_graph_shape: Optional[str] = None,
+                dot_graph_shape: str | None = None,
                 normalize_args: bool = False,
             ):
                 raise RuntimeError(
diff --git a/torch/fx/passes/graph_manipulation.py b/torch/fx/passes/graph_manipulation.py
index a558f962b8b5b..3c617634101d1 100644
--- a/torch/fx/passes/graph_manipulation.py
+++ b/torch/fx/passes/graph_manipulation.py
@@ -1,5 +1,5 @@
 # mypy: allow-untyped-defs
-from typing import Any, NamedTuple, Optional
+from typing import Any, NamedTuple
 
 import torch
 from torch.fx._compatibility import compatibility
@@ -26,8 +26,10 @@ def replace_target_nodes_with(
     new_op: str,
     new_target: Target,
 ):
-    """Modifies all nodes in fx_module.graph.nodes which match the specified op code and target,
-    and updates them to match the new op code and target"""
+    """
+    Modifies all nodes in fx_module.graph.nodes which match the specified op code
+    and target, and updates them to match the new op code and target.
+    """
     new_graph = Graph()
     val_map: dict[Node, Node] = {}
     for node in fx_module.graph.nodes:
@@ -54,7 +56,7 @@ class size_bytes(NamedTuple):
 
 @compatibility(is_backward_compatible=False)
 def get_size_of_all_nodes(
-    fx_module: GraphModule, args: Optional[list[torch.Tensor]] = None
+    fx_module: GraphModule, args: list[torch.Tensor] | None = None
 ) -> None:
     """Given a fx graph module, update each node with its total size (weights + bias + output)
     and its output_size(output). For a non-module node, the total size is the output size.
diff --git a/torch/fx/passes/graph_transform_observer.py b/torch/fx/passes/graph_transform_observer.py
index 6c42a614e1b48..f91877b00a057 100644
--- a/torch/fx/passes/graph_transform_observer.py
+++ b/torch/fx/passes/graph_transform_observer.py
@@ -1,7 +1,7 @@
 # mypy: allow-untyped-defs
 import os
 from collections.abc import Callable
-from typing import Optional, TypeVar
+from typing import TypeVar
 
 from torch.fx import Graph, Node
 from torch.fx._compatibility import compatibility
@@ -26,8 +26,8 @@ def __init__(
         self,
         gm: GraphModule,
         passname: str,
-        subsystem: Optional[str] = None,
-        log_url: Optional[str] = None,
+        subsystem: str | None = None,
+        log_url: str | None = None,
     ):
         """
         log_url is inferred to be torch._inductor.config.trace.log_url_for_graph_xform unless otherwise specified
@@ -76,7 +76,7 @@ def __init__(
     def get_current_pass_count(cls):
         return cls.__pass_count
 
-    def apply_gm_pass(self, pass_fn: Callable[[GraphModule], T]) -> Optional[T]:
+    def apply_gm_pass(self, pass_fn: Callable[[GraphModule], T]) -> T | None:
         from torch._dynamo.utils import dynamo_timed
 
         with self:
@@ -89,7 +89,7 @@ def apply_gm_pass(self, pass_fn: Callable[[GraphModule], T]) -> Optional[T]:
             ):
                 return pass_fn(self.gm)
 
-    def apply_graph_pass(self, pass_fn: Callable[[Graph], T]) -> Optional[T]:
+    def apply_graph_pass(self, pass_fn: Callable[[Graph], T]) -> T | None:
         from torch._dynamo.utils import dynamo_timed
 
         with self:
diff --git a/torch/fx/passes/infra/partitioner.py b/torch/fx/passes/infra/partitioner.py
index d612854dc7bc8..550c6d57be865 100644
--- a/torch/fx/passes/infra/partitioner.py
+++ b/torch/fx/passes/infra/partitioner.py
@@ -4,7 +4,6 @@
 import logging
 import operator
 from collections.abc import Iterable, Sequence
-from typing import Optional
 
 from torch.fx.graph_module import GraphModule
 from torch.fx.node import _get_qualified_name, Node
@@ -19,12 +18,12 @@
 class Partition:
     def __init__(
         self,
-        id: Optional[int] = None,
-        nodes: Optional[Iterable[Node]] = None,
-        node_orders: Optional[Iterable[int]] = None,
+        id: int | None = None,
+        nodes: Iterable[Node] | None = None,
+        node_orders: Iterable[int] | None = None,
     ):
         self.id = id
-        self.nodes: dict[Node, Optional[int]] = {}
+        self.nodes: dict[Node, int | None] = {}
         if nodes is not None:
             if node_orders is None:
                 self.nodes = dict.fromkeys(nodes, None)
@@ -40,7 +39,7 @@ def __init__(
     def __repr__(self) -> str:
         return str(self.nodes)
 
-    def add_node(self, node: Node, node_order: Optional[int] = None):
+    def add_node(self, node: Node, node_order: int | None = None):
         self.nodes.update({node: node_order})
 
     def remove_node(self, node: Node):
@@ -70,8 +69,8 @@ def __init__(
         graph_module: GraphModule,
         operator_support: OperatorSupportBase,
         allows_single_node_partition: bool = False,
-        non_compute_ops: Optional[Sequence[str]] = None,
-        allowed_single_node_partition_ops: Optional[Sequence[str]] = None,
+        non_compute_ops: Sequence[str] | None = None,
+        allowed_single_node_partition_ops: Sequence[str] | None = None,
     ) -> None:
         self.graph_module = graph_module
         self.operator_support = operator_support
@@ -186,7 +185,7 @@ def dfs_iter_find_cycle(all_user_nodes: set[Node]):
 
             return merge_id, True
 
-        def merge_single_node(node: Node, node_order: Optional[int], id: Optional[int]):
+        def merge_single_node(node: Node, node_order: int | None, id: int | None):
             def _update_partition_map(node: Node, id: int):
                 # Iterate through all the users of this node and update the partition map to indicate
                 # that there is a path from the partition id of this node to the target partition id.
diff --git a/torch/fx/passes/infra/pass_base.py b/torch/fx/passes/infra/pass_base.py
index 32c641031b31f..109f62f1488bb 100644
--- a/torch/fx/passes/infra/pass_base.py
+++ b/torch/fx/passes/infra/pass_base.py
@@ -1,7 +1,6 @@
 # mypy: allow-untyped-defs
 import abc
 from collections import namedtuple
-from typing import Optional
 
 from torch.fx._compatibility import compatibility
 from torch.fx.graph_module import GraphModule
@@ -38,7 +37,7 @@ class PassBase(abc.ABC):
     the PassManager's `passes` attribute.
     """
 
-    def __call__(self, graph_module: GraphModule) -> Optional[PassResult]:
+    def __call__(self, graph_module: GraphModule) -> PassResult | None:
         """
         Runs the precondition check, the pass itself, and the postcondition check.
         """
@@ -49,7 +48,7 @@ def __call__(self, graph_module: GraphModule) -> Optional[PassResult]:
         return res
 
     @abc.abstractmethod
-    def call(self, graph_module: GraphModule) -> Optional[PassResult]:
+    def call(self, graph_module: GraphModule) -> PassResult | None:
         """
         The pass that is run through the given graph module. To implement a
         pass, it is required to implement this function.
diff --git a/torch/fx/passes/infra/pass_manager.py b/torch/fx/passes/infra/pass_manager.py
index 87fb6e70037f9..484ce248dfd2d 100644
--- a/torch/fx/passes/infra/pass_manager.py
+++ b/torch/fx/passes/infra/pass_manager.py
@@ -124,22 +124,21 @@ def _topological_sort_passes(
 @compatibility(is_backward_compatible=False)
 def this_before_that_pass_constraint(this: Callable, that: Callable) -> Callable:
     """
-    Defines a partial order ('depends on' function) where `this` must occur
-    before `that`.
+    Defines a partial order ('depends on' function) where ``this`` must occur
+    before ``that``.
 
-    For example, the following pass list and constraint list would be invalid.
-    ```
-    passes = [pass_b, pass_a]
+    For example, the following pass list and constraint list would be invalid::
 
-    constraints = [this_before_that_pass_constraint(pass_a, pass_b)]
-    ```
+        passes = [pass_b, pass_a]
+
+        constraints = [this_before_that_pass_constraint(pass_a, pass_b)]
 
     Args:
         this (Callable): pass which should occur first
         that (Callable): pass which should occur later
 
     Returns:
-        depends_on (Callable[[Object, Object], bool]
+        depends_on (Callable[[Object, Object], bool])
     """
 
     def depends_on(a: Callable, b: Callable):
diff --git a/torch/fx/passes/net_min_base.py b/torch/fx/passes/net_min_base.py
index 1b05cb7496a43..34636c6f3664a 100644
--- a/torch/fx/passes/net_min_base.py
+++ b/torch/fx/passes/net_min_base.py
@@ -2,7 +2,7 @@
 import logging
 from collections.abc import Callable
 from dataclasses import dataclass
-from typing import Any, cast, Optional
+from typing import Any, cast
 
 import torch
 import torch.fx
@@ -110,10 +110,9 @@ def __init__(
             [TensorOrTensors, TensorOrTensors, Names], tuple[float, bool]
         ],
         settings: _MinimizerSettingBase,
-        module_exporter: Optional[
-            Callable[[Tensors, torch.fx.GraphModule, str], None]
-        ] = None,
-        exclusion_fn: Optional[Callable[[NodeList, int, int], None]] = None,
+        module_exporter: Callable[[Tensors, torch.fx.GraphModule, str], None]
+        | None = None,
+        exclusion_fn: Callable[[NodeList, int, int], None] | None = None,
     ):
         if not isinstance(module, torch.fx.GraphModule):
             raise AssertionError(f"Expected GraphModule, got {type(module)}")
@@ -549,7 +548,7 @@ def _sequential_traverse(self, nodes: NodeList) -> NodeSet:
 
     def _block_traverse_impl(
         self, nodes: NodeList, start_idx: int, end_idx: int, find_last_node: bool
-    ) -> Optional[int]:
+    ) -> int | None:
         """
         Recursive block search implementation.
         find_last_node: If True, search for the last node which result in numerics difference
@@ -653,9 +652,7 @@ def _block_traverse_impl(
             else:
                 return self._block_traverse_impl(nodes, start_idx, mid, find_last_node)
 
-    def _block_traverse(
-        self, nodes: NodeList, find_last_node: Optional[bool]
-    ) -> NodeSet:
+    def _block_traverse(self, nodes: NodeList, find_last_node: bool | None) -> NodeSet:
         """
         Traverse topologically sorted node list
         Find minimum block (start_idx, end_idx) which contains the culprit
@@ -675,8 +672,8 @@ def _block_traverse(
         start_idx = 0
         end_idx = len(nodes) - 1
 
-        final_start_idx: Optional[int] = start_idx
-        final_end_idx: Optional[int] = end_idx
+        final_start_idx: int | None = start_idx
+        final_end_idx: int | None = end_idx
 
         run_both = find_last_node is None
 
@@ -860,7 +857,7 @@ def _skip_traverse(self, all_nodes: NodeList, skip_nodes: list) -> NodeSet:
 
         return culprits
 
-    def _collect_nodes(self, start: Optional[str], end: Optional[str]) -> NodeList:
+    def _collect_nodes(self, start: str | None, end: str | None) -> NodeList:
         """
         Collect nodes in the model that between nodes with name of `start` and `end`.
         These two nodes are also included.
@@ -883,7 +880,7 @@ def _collect_nodes(self, start: Optional[str], end: Optional[str]) -> NodeList:
 
         return nodes
 
-    def run_nodes(self, start: Optional[str] = None, end: Optional[str] = None):
+    def run_nodes(self, start: str | None = None, end: str | None = None):
         """
         Run part of the model from `start` node to `end` node. If `start` is None
         then we start from the beginning of the model. If `end` is None then we
@@ -930,10 +927,10 @@ def print_reports(self):
 
     def minimize(
         self,
-        start: Optional[str] = None,
-        end: Optional[str] = None,
-        skip_nodes: Optional[list] = None,
-        find_last_node: Optional[bool] = None,
+        start: str | None = None,
+        end: str | None = None,
+        skip_nodes: list | None = None,
+        find_last_node: bool | None = None,
     ) -> NodeSet:
         """
         Minimizing the model from node with name `start` to node with name `end` base
diff --git a/torch/fx/passes/operator_support.py b/torch/fx/passes/operator_support.py
index 2af56c6dc0d71..813e22090ba7a 100644
--- a/torch/fx/passes/operator_support.py
+++ b/torch/fx/passes/operator_support.py
@@ -23,12 +23,13 @@
 TargetTypeName = str
 
 # Arguments' dtypes for a given node, see `OperatorSupport`
-SupportedArgumentDTypes = t.Optional[
+SupportedArgumentDTypes = (
     tuple[
         t.Sequence[t.Sequence[torch.dtype]],
         dict[str, t.Sequence[torch.dtype]],
     ]
-]
+    | None
+)
 
 SupportDict = t.Mapping[TargetTypeName, SupportedArgumentDTypes]
 
@@ -66,7 +67,7 @@ class OperatorSupport(OperatorSupportBase):
 
     _support_dict: SupportDict
 
-    def __init__(self, support_dict: t.Optional[SupportDict] = None):
+    def __init__(self, support_dict: SupportDict | None = None):
         self._support_dict = support_dict or {}
 
     def is_node_supported(
diff --git a/torch/fx/passes/pass_manager.py b/torch/fx/passes/pass_manager.py
index 9e1501cfcd5fc..8c18ccdea9df7 100644
--- a/torch/fx/passes/pass_manager.py
+++ b/torch/fx/passes/pass_manager.py
@@ -3,7 +3,6 @@
 from collections.abc import Callable
 from functools import wraps
 from inspect import unwrap
-from typing import Optional
 
 
 logger = logging.getLogger(__name__)
@@ -44,22 +43,21 @@ def log_hook(fn: Callable, level=logging.INFO) -> Callable:
     """
     Logs callable output.
 
-    This is useful for logging output of passes. Note inplace_wrapper replaces
+    This is useful for logging output of passes. Note ``inplace_wrapper`` replaces
     the pass output with the modified object. If we want to log the original
-    output, apply this wrapper before inplace_wrapper.
+    output, apply this wrapper before ``inplace_wrapper``.
 
+    Example::
 
-    ```
-    def my_pass(d: Dict) -> bool:
-        changed = False
-        if "foo" in d:
-            d["foo"] = "bar"
-            changed = True
-        return changed
+        def my_pass(d: Dict) -> bool:
+            changed = False
+            if "foo" in d:
+                d["foo"] = "bar"
+                changed = True
+            return changed
 
 
-    pm = PassManager(passes=[inplace_wrapper(log_hook(my_pass))])
-    ```
+        pm = PassManager(passes=[inplace_wrapper(log_hook(my_pass))])
 
     Args:
         fn (Callable[Type1, Type2])
@@ -80,8 +78,8 @@ def wrapped_fn(gm):
 
 def loop_pass(
     base_pass: Callable,
-    n_iter: Optional[int] = None,
-    predicate: Optional[Callable] = None,
+    n_iter: int | None = None,
+    predicate: Callable | None = None,
 ):
     """
     Convenience wrapper for passes which need to be applied multiple times.
@@ -148,25 +146,24 @@ def depends_on(a: Callable, b: Callable):
 
 def these_before_those_pass_constraint(these: Callable, those: Callable):
     """
-    Defines a partial order ('depends on' function) where `these` must occur
-    before `those`. Where the inputs are 'unwrapped' before comparison.
+    Defines a partial order ('depends on' function) where ``these`` must occur
+    before ``those``. Where the inputs are 'unwrapped' before comparison.
 
-    For example, the following pass list and constraint list would be invalid.
-    ```
-    passes = [
-        loop_pass(pass_b, 3),
-        loop_pass(pass_a, 5),
-    ]
+    For example, the following pass list and constraint list would be invalid::
 
-    constraints = [these_before_those_pass_constraint(pass_a, pass_b)]
-    ```
+        passes = [
+            loop_pass(pass_b, 3),
+            loop_pass(pass_a, 5),
+        ]
+
+        constraints = [these_before_those_pass_constraint(pass_a, pass_b)]
 
     Args:
         these (Callable): pass which should occur first
         those (Callable): pass which should occur later
 
     Returns:
-        depends_on (Callable[[Object, Object], bool]
+        depends_on (Callable[[Object, Object], bool])
     """
 
     def depends_on(a: Callable, b: Callable):
diff --git a/torch/fx/passes/regional_inductor.py b/torch/fx/passes/regional_inductor.py
index aa830c705ef72..f980ab4f8dcdb 100644
--- a/torch/fx/passes/regional_inductor.py
+++ b/torch/fx/passes/regional_inductor.py
@@ -1,7 +1,9 @@
 # mypy: allow-untyped-defs
 
+import contextlib
 import functools
 import logging
+from collections.abc import Iterator
 
 import torch
 from torch.fx._compatibility import compatibility
@@ -23,23 +25,15 @@ def inner(*args, **kwargs):
     return inner
 
 
-def _partition_by_supported_nodes(gm, supported_ops, prefix):
-    from torch.fx.passes.infra.partitioner import CapabilityBasedPartitioner
-    from torch.fx.passes.utils.fuser_utils import fuse_by_partitions
-
-    partitioner = CapabilityBasedPartitioner(
-        gm, supported_ops, allows_single_node_partition=True
-    )
-
-    candidate_partitions = partitioner.propose_partitions()
-    partitioned_gm = fuse_by_partitions(
-        partitioner.graph_module,
-        [partition.nodes for partition in candidate_partitions],
-        prefix=prefix,
-        always_return_tuple=True,
-    )
-
-    return partitioned_gm
+@contextlib.contextmanager
+def _disable_remat_for_regional_subcompile() -> Iterator[None]:
+    # In torch.compile, regional_inductor subcompiles run after the enclosing
+    # non-strict full graph has already been partitioned, so any graph-SAC
+    # remat pass has already run before we reach this nested compile.
+    # Rerunning remat here can see stage-2-reordered backward nodes that
+    # violate remat's contiguous-backward-region assumption.
+    with torch._functorch.config.patch(remat_using_tags_for_fwd_loss_bwd_graph=False):
+        yield
 
 
 def _compile_submod(gm, prefix):
@@ -92,7 +86,10 @@ def _compile_submod(gm, prefix):
                         f"Available config keys can be found in torch._inductor.config"
                     )
 
-            with inductor_config.patch(inductor_options):
+            with (
+                inductor_config.patch(inductor_options),
+                _disable_remat_for_regional_subcompile(),
+            ):
                 compiled_fn = torch._inductor.standalone_compile(
                     submod,
                     fake_inputs,
@@ -134,29 +131,59 @@ class _RegionScooper:
 
     @staticmethod
     def scoop_regions(gm):
-        from torch.fx.passes.operator_support import OperatorSupport
-
-        found_marked_node = False
+        from torch.fx.passes.infra.partitioner import CapabilityBasedPartitioner
+        from torch.fx.passes.operator_support import create_op_support
+        from torch.fx.passes.utils.fuser_utils import fuse_by_partitions
+
+        # Group tagged nodes by region ID.  The region ID comes from the
+        # optional "inductor_region" key inside the compile_with_inductor
+        # annotation. When absent, all tagged nodes share a single default region
+        _DEFAULT_REGION = object()
+        regions: dict[object, set[torch.fx.Node]] = {}
         for node in gm.graph.nodes:
             if _needs_inductor_compile(node):
-                found_marked_node = True
-                break
+                compile_value = node.meta["custom"]["compile_with_inductor"]
+                if (
+                    isinstance(compile_value, dict)
+                    and "inductor_region" in compile_value
+                ):
+                    rid = compile_value["inductor_region"]
+                else:
+                    rid = _DEFAULT_REGION
+                regions.setdefault(rid, set()).add(node)
 
-        if not found_marked_node:
+        if not regions:
             logger.info("No inductor marked nodes found")
             return gm
 
-        class InductorMarkedNodes(OperatorSupport):
-            def is_node_supported(self, submodules, node: torch.fx.Node) -> bool:
-                return _needs_inductor_compile(node)
+        # Run CapabilityBasedPartitioner per region to get cycle-safe partitions
+        # without merging across region boundaries.
+        def _is_in_region(region_nodes):
+            def is_node_supported(_submodules, node):
+                return node in region_nodes
+
+            return is_node_supported
 
-        marked_nodes = InductorMarkedNodes()
-        return _partition_by_supported_nodes(
-            gm, marked_nodes, "__marked_inductor_submod"
+        all_partitions: list[dict[torch.fx.Node, int | None]] = []
+        for region_nodes in regions.values():
+            support = create_op_support(_is_in_region(region_nodes))
+            partitioner = CapabilityBasedPartitioner(
+                gm, support, allows_single_node_partition=True
+            )
+            for partition in partitioner.propose_partitions():
+                all_partitions.append(partition.nodes)
+
+        return fuse_by_partitions(
+            gm,
+            all_partitions,
+            prefix="__marked_inductor_submod",
+            always_return_tuple=True,
         )
 
     @staticmethod
-    def recursively_scoop_regions(gm):
+    def recursively_scoop_regions(gm, _processed=None):
+        if _processed is None:
+            _processed = set()
         for node in gm.graph.find_nodes(op="get_attr"):
             if _needs_inductor_compile(node):
                 # If the get_attr itself is marked for compile, the outer graph will
@@ -164,8 +191,13 @@ def recursively_scoop_regions(gm):
                 # regional inductor compiles that do not work well.
                 continue
             submod = getattr(gm, node.target)
-            if isinstance(submod, torch.fx.GraphModule):
-                _RegionScooper.recursively_scoop_regions(submod)
+            # Track by id: multiple get_attr nodes may reference the same GraphModule
+            if (
+                isinstance(submod, torch.fx.GraphModule)
+                and id(submod) not in _processed
+            ):
+                _processed.add(id(submod))
+                _RegionScooper.recursively_scoop_regions(submod, _processed)
 
         return _RegionScooper.scoop_regions(gm)
 
@@ -228,15 +260,19 @@ def regional_inductor(gm, *example_args):
     """
     Scoops out inductor marked regions and compiles them with inductor.
 
-    Inductor options should be provided via the annotation API:
-    with fx_traceback.annotate({
-        "compile_with_inductor": {
-            "inductor_configs": {
-                "max_autotune": True,
-                "triton.cudagraphs": False
+    Inductor options should be provided via the annotation API::
+
+        with fx_traceback.annotate(
+            {
+                "compile_with_inductor": {
+                    "inductor_configs": {
+                        "max_autotune": True,
+                        "triton.cudagraphs": False,
+                    }
+                }
             }
-        }
-    }):
+        ):
+            ...
     """
 
     # fuser utils create new nodes using create_proxy which retains the seq_nr
diff --git a/torch/fx/passes/regional_inductor_invoke_subgraph.py b/torch/fx/passes/regional_inductor_invoke_subgraph.py
index b2a6e5c3a9b28..8a106ace3f077 100644
--- a/torch/fx/passes/regional_inductor_invoke_subgraph.py
+++ b/torch/fx/passes/regional_inductor_invoke_subgraph.py
@@ -8,7 +8,10 @@
 from torch._inductor.standalone_compile import AOTCompiledArtifact
 from torch.compiler._cache import CacheArtifactManager
 from torch.fx._compatibility import compatibility
-from torch.fx.passes.regional_inductor import _dummy_wrapper
+from torch.fx.passes.regional_inductor import (
+    _disable_remat_for_regional_subcompile,
+    _dummy_wrapper,
+)
 
 
 logger = logging.getLogger(__name__)
@@ -67,6 +70,7 @@ def get_compiled_fn():
             torch._guards.tracing(context),
             CacheArtifactManager.with_fresh_cache(),
             torch._functorch.config.patch("bundled_autograd_cache", True),
+            _disable_remat_for_regional_subcompile(),
         ):
             # compile_fx can mutate gm
             gm = copy.deepcopy(submod)
diff --git a/torch/fx/passes/reinplace.py b/torch/fx/passes/reinplace.py
index 17dde9256f5fa..07623a9362a69 100644
--- a/torch/fx/passes/reinplace.py
+++ b/torch/fx/passes/reinplace.py
@@ -309,104 +309,108 @@ def matching_view_metadata(a, b):
 
 @compatibility(is_backward_compatible=True)
 def reinplace(gm, *sample_args):
-    """
+    r"""
     Given an fx.GraphModule, modifies it to perform "reinplacing",
     mutating the nodes of the graph.
-    We look for out-of-place op call sites like `b = a.add(...)`,
-    and convert them to be inplace (`b = a.add_(...)`),
+    We look for out-of-place op call sites like ``b = a.add(...)``,
+    and convert them to be inplace (``b = a.add_(...)``),
     as long as the input to the current operator ("a") isn't reused
     anywhere later in the graph.
 
     This pass currently expects to operate on a **functional, ATen** graph.
-    This can be obtained by running `make_fx(functionalize(f))`.
+    This can be obtained by running ``make_fx(functionalize(f))``.
 
     Sample inputs are needed to determine aliasing relationships of the inputs.
-    In general, we can't reinplace node `b = a.add(...)` if "a" aliases any of the
+    In general, we can't reinplace node ``b = a.add(...)`` if "a" aliases any of the
     inputs to the program.
 
-    Given a node "b = foo(a, args...) the algorithm for re-inplacing is as follows:
+    Given a node ``b = foo(a, args...)`` the algorithm for re-inplacing is as follows:
+
+    **(1)** Perform some initial checks on the metadata of "a" and "args..."
+    that can disqualify them from being reinplaced.
+
+    - **(1a)** Check that the self argument we're attempting to reinplace
+      has acceptable dtype/size metadata to reinplace with.
+
+      For example, if we have::
+
+        a = torch.ones(1)
+        b = torch.ones(10)
+        out = torch.add(a, b)
+
+      We can't turn that into ``a.add_(b)`` because that would require resizing "a".
 
-    (1) Perform some initial checks on the metadata of "a" and "args..."
-        that can disqualify them from being reinplaced.
+      Similarly, we can't convert ``torch.ge(a, b)`` into ``a.ge_(b)``,
+      because that would require changing a's dtype (from e.g. float32 to bool).
+      Note that in this specific example, we could technically do better..
 
-      (1a) Check that the self argument we're attempting to reinplace
-           has acceptable dtype/size metadata to reinplace with.
+      If we see the pattern::
 
-           For example, if we have:
-             a = torch.ones(1)
-             b = torch.ones(10)
-             out = torch.add(a, b)
-           We can't turn that into
-             a.add_(b)
-           Because that would require resizing "a".
+        a_1 = a.ge(b)
+        a_2 = aten._to_copy(a_1, a.dtype)
 
-           Similarly, we can't convert torch.ge(a, b) into a.ge_(b),
-           because that would require changing a's dtype (from e.g. float32 to bool).
-           Note that in this specific example, we could technically do better..
+      Then this should be valid to completely re-inplace
+      (this is exactly what functionalization will emit when it sees ``a.ge_(b)``).
 
-           If we see the pattern:
-             a_1 = a.ge(b)
-             a_2 = aten._to_copy(a_1, a.dtype)
-           Then we this should be valid to completely re-inplace
-           (this is exactly what functionalization will emit when it sees a.ge_(b)).
+      This optimization is only really important for user programs
+      that directly use inplace comparison ops though.
 
-           This optimization is only really important for user programs
-           that directly use inplace comparison ops though.
+      We also cannot re-inplace on tensors that have overlapping memory,
+      e.g. ``torch.ones(1).expand(4, 4).add_(1)``.
 
-           We also cannot re-inplace on tensors that have overlapping memory,
-           e.g. torch.ones(1).expand(4, 4).add_(1)
+    - **(1b)** Check if "a" is an alias of any of the program inputs.
 
-      (1b) Check if "a" is an alias of any of the program inputs.
+      If it is, skip and move to the next node.
+      Inplace'ing an op that would cause it to mutate a program is not sound,
+      because that would be a side effect visible to the user.
 
-          If it is, skip and move to the next node.
-          Inplace'ing an op that would cause it to mutate a program is not sound,
-          because that would be a side effect visible to the user.
+      NOTE: there's a future optimization that we should make:
+      if "a" is a (alias of a)  program input, but later in the program
+      there is a node that looks like ``a.copy_(...)``,
+      then re-inplacing is ok to do - we are temporarily reusing a's buffer,
+      which will later be overwritten by the ``copy_()`` call.
 
-          NOTE: there's a future optimization that we should make:
-          if "a" is a (alias of a)  program input, but later in the program
-          there is a node that looks like "a.copy_(...)",
-          Then re-inplacing is ok to do - we are temporarily reusing a's buffer,
-          which will later be overwritten by the copy_() call.
+      This will be an important optimization to have for programs that mutate
+      their inputs. It currently isn't implemented though.
 
-          This will be an important optimization to have for programs that mutate
-          their inputs. It currently isn't implemented though.
+    - **(1c)** Check if "a" and "args..." alias.
 
-      (1c) Check if "a" and "args..." alias
+      For example, re-inplacing to create code like the below
+      isn't guaranteed to be sound::
 
-          For example, re-inplacing to create code like the below
-          isn't guaranteed to be sound:
+        aten.mul_(a, a)
 
-            aten.mul_(a, a)
+    **(2)** Check that "a" and all of its outstanding aliases are not used anywhere
+    later in the graph. If this is the case, then it's safe to re-inplace
+    to ``b = foo_(a)``.
 
-    (2) Check that "a" and all of its outstanding aliases are not used anywhere
-        later in the graph. If this is the case, then it's safe to re-inplace
-        to "b = foo_(a)".
+    There are a few caveats to this, explained in more detail below:
 
-        There are a few caveats to this, explained in more detail below:
-        (a) If "a" is used later as an argument to a view op, that is okay.
-            It's only a problem if "a" (or that view) is later passed
-            into a normal operator, or if it is returned as the program output.
-        (b) If "a" is a repeat argument in `foo()`, then don't reinplace.
-            Most ATen kernels don't make any guarantees that this is sound,
-            e.g. if you do aten.mul_(a, a).
-            So we'll just ban re-inplacing in this case.
-            It's only a problem if "a" (or that view) is later passed
-        (c) If "a" is used as an input into a view "inverse" / "scatter"
-            operator, it is potentially fine to re-inplace
-            (and remove that scatter operator from the graph).
-            See below for a more detailed example.
+    - (a) If "a" is used later as an argument to a view op, that is okay.
+      It's only a problem if "a" (or that view) is later passed
+      into a normal operator, or if it is returned as the program output.
+    - (b) If "a" is a repeat argument in ``foo()``, then don't reinplace.
+      Most ATen kernels don't make any guarantees that this is sound,
+      e.g. if you do ``aten.mul_(a, a)``.
+      So we'll just ban re-inplacing in this case.
+    - (c) If "a" is used as an input into a view "inverse" / "scatter"
+      operator, it is potentially fine to re-inplace
+      (and remove that scatter operator from the graph).
+      See below for a more detailed example.
 
-        NOTE: there is an optimization in this step that is crucial
-        to fully recovering performance from functionalization.
+    NOTE: there is an optimization in this step that is crucial
+    to fully recovering performance from functionalization.
+
+    Given this program::
 
-        Given this program:
         def f(x):
             a = torch.ops.aten.add(x, x)
             b = torch.ops.aten.diagonal(a)
             torch.ops.aten.fill_(b, 0)
             return d
 
-        Functionalization will emit the following:
+    Functionalization will emit the following::
+
         def f(x):
             a = torch.ops.aten.add(x, x)
             b = torch.ops.aten.diagonal(a, 0, 1)
@@ -414,99 +418,115 @@ def f(x):
             a_updated = torch.ops.aten.diagonal_scatter(a, b_updated, 0, 1)
             return a_updated
 
-        Ordinarily, we would not be able to reinplace the fill,
-        because "b" aliases with "a" which is used by the diagonal_scatter call.
-
-        "re-inplacing" is on the hook for figuring out that it is ok to
-        completely, the expensive diagonal_scatter call, if we re-inplace the add().
-
-        So, for every `alias in alias_set(a)`, instead of checking
-        that "alias" is not used anywhere later in the graph,
-        we check that
-            EITHER:
-          (a) alias is not used anywhere later in the graph
-            OR:
-          (b) alias is used exactly once later on in the graph,
-              in the following op:
-
-                out = foo_scatter(alias, x, args...)
-
-              where the following must hold:
-                (i) "foo_scatter" is the "inverse" operator for foo.
-                    This only applies to "foo" ops that are view operators,
-                    which view into a subset of the original tensor's memory.
-                    In practice, there are ~4 operators where this applies:
-                      diagonal -> diagonal_scatter
-                      slice -> slice_scatter
-                      select -> select_scatter
-                      as_strided -> as_strided_scatter
-                (ii) "args..." are the same between the foo() and foo_scatter() calls.
-
-    (3) Perform the actual re-inplacing on foo!
-
-      (3b) is the common case, but special care is needed for {view}_scatter (3a)
-
-      (3a) {view}_scatter ops.
-
-        Consider this program:
-          a = torch.zeros(2, 2)
-          b = torch.ones(2)
-          a[0] = b
-
-        Post functionalization, that will look like:
-          a = torch.zeros(2)
-          b = torch.ones(1)
-          a_updated = torch.select_scatter(a, b, 0, 0)
-
-        In this case though, there is no "functional" op to re-inplace!
-        Instead, we'd like to directly remove toe select_scatter call.
-        We already know from (3) that this is valid,
-        because "a" has no later usages in the graph.
-
-        We perform the re-inplacing on the {view}_scatter op like so
-        Before:
-          a_updated = torch.select_scatter(a, b, args...)
-        After:
-          a_slice = a.select(a, args...)
-          a_slice.copy_(b)
-
-      (3b) Otherwise, replace the functional op with its inplace variant.
-        Before:
-          b = foo(a, args...)
-        After:
-          a.foo_(args...)
-
-    (4) Finally, after converting either:
-          Before:
-            b = foo(a)
-          After:
-            foo_(a)
-        or
-          Before:
-            b = {slice}_scatter(a, mutated_slice, args...)
-          After:
-            slice = {slice}(a, args...)
-            slice.copy_(mutated_slice)
-
-        We now need to find all later nodes that use "b" as an argument
-        and update them to take in "a" instead.
-
-        Note that for the majority of inplace ops, this isn't actually necessary
-        (because most inplace ops return "self" as their output).
-        This isn't generally true for all mutable ops though, which is why
-        we need to actually replace all of the arguments.
-
-        We also need to update our metadata of Dict[StorageWeakRef, Set[Node]],
-        That maps a given tensor storage to the set of all nodes that take in that storage
-        as an input.
-        Specifically, re-inplacing `b = foo(a)` causes "a" and "b"'s sets to get fused
-        together.
-
-    (5) Any "view_inverse/scatter" nodes that were identified as "it's ok to ignore them"
-        during step (3) get manually deleted from the graph.
-        Their outputs are no longer used, so technically standard DCE would be able
-        to do this, but we can no longer run FX's DCE pass now that we have mutable
-        ops in the graph.
+    Ordinarily, we would not be able to reinplace the fill,
+    because "b" aliases with "a" which is used by the diagonal_scatter call.
+
+    "re-inplacing" is on the hook for figuring out that it is ok to
+    completely remove the expensive diagonal_scatter call, if we re-inplace
+    the add().
+
+    So, for every ``alias in alias_set(a)``, instead of checking
+    that "alias" is not used anywhere later in the graph,
+    we check that EITHER:
+
+    - (a) alias is not used anywhere later in the graph, OR
+    - (b) alias is used exactly once later on in the graph,
+      in the following op::
+
+        out = foo_scatter(alias, x, args...)
+
+      where the following must hold:
+
+      - (i) ``foo_scatter`` is the "inverse" operator for foo.
+        This only applies to "foo" ops that are view operators,
+        which view into a subset of the original tensor's memory.
+        In practice, there are ~4 operators where this applies::
+
+          diagonal -> diagonal_scatter
+          slice -> slice_scatter
+          select -> select_scatter
+          as_strided -> as_strided_scatter
+
+      - (ii) "args..." are the same between the ``foo()`` and
+        ``foo_scatter()`` calls.
+
+    **(3)** Perform the actual re-inplacing on foo!
+
+    (3b) is the common case, but special care is needed for
+    ``{view}_scatter`` (3a).
+
+    - **(3a)** ``{view}_scatter`` ops.
+
+      Consider this program::
+
+        a = torch.zeros(2, 2)
+        b = torch.ones(2)
+        a[0] = b
+
+      Post functionalization, that will look like::
+
+        a = torch.zeros(2)
+        b = torch.ones(1)
+        a_updated = torch.select_scatter(a, b, 0, 0)
+
+      In this case though, there is no "functional" op to re-inplace!
+      Instead, we'd like to directly remove the select_scatter call.
+      We already know from (3) that this is valid,
+      because "a" has no later usages in the graph.
+
+      We perform the re-inplacing on the ``{view}_scatter`` op like so.
+
+      Before::
+
+        a_updated = torch.select_scatter(a, b, args...)
+
+      After::
+
+        a_slice = a.select(a, args...)
+        a_slice.copy_(b)
+
+    - **(3b)** Otherwise, replace the functional op with its inplace variant.
+
+      Before::
+
+        b = foo(a, args...)
+
+      After::
+
+        a.foo_(args...)
+
+    **(4)** Finally, after converting either::
+
+        # Before:                              # After:
+        b = foo(a)                             foo_(a)
+
+    or::
+
+        # Before:
+        b = {slice}_scatter(a, mutated_slice, args...)
+        # After:
+        slice = {slice}(a, args...)
+        slice.copy_(mutated_slice)
+
+    We now need to find all later nodes that use "b" as an argument
+    and update them to take in "a" instead.
+
+    Note that for the majority of inplace ops, this isn't actually necessary
+    (because most inplace ops return "self" as their output).
+    This isn't generally true for all mutable ops though, which is why
+    we need to actually replace all of the arguments.
+
+    We also need to update our metadata of ``Dict[StorageWeakRef, Set[Node]]``,
+    that maps a given tensor storage to the set of all nodes that take in that
+    storage as an input.
+    Specifically, re-inplacing ``b = foo(a)`` causes "a" and "b"'s sets to get
+    fused together.
+
+    **(5)** Any ``view_inverse/scatter`` nodes that were identified as
+    "it's ok to ignore them" during step (3) get manually deleted from the graph.
+    Their outputs are no longer used, so technically standard DCE would be able
+    to do this, but we can no longer run FX's DCE pass now that we have mutable
+    ops in the graph.
     """
     _FunctionalizationMetadataProp(gm).propagate(*sample_args)
 
@@ -751,6 +771,7 @@ def replace_arg(a):
                         and len(node_res_storage) == 1
                         and old_res_storage == node_res_storage
                     ):
+                        # pyrefly: ignore [missing-attribute]
                         new_flattened_res = pytree.tree_leaves(new.meta["fake_result"])
                         new_res_storage = {
                             StorageWeakRef(x._typed_storage())
diff --git a/torch/fx/passes/runtime_assert.py b/torch/fx/passes/runtime_assert.py
index cc89b31366c86..5a4ff5c5a2ea9 100644
--- a/torch/fx/passes/runtime_assert.py
+++ b/torch/fx/passes/runtime_assert.py
@@ -31,7 +31,7 @@
 graph_code_log = torch._logging.getArtifactLogger(__name__, "graph_code_verbose")
 
 
-def _get_example_value(node: fx.Node) -> Optional[str]:
+def _get_example_value(node: fx.Node) -> str | None:
     """
     Get the example value key for a node, since dynamo uses "example_value"
     while non-strict export uses "val.
@@ -98,7 +98,6 @@ def insert_deferred_runtime_asserts(
         _get_placeholder_expr,
         _has_uninterpretable_sympy_function,
         CallMethodKey,
-        cast_symbool_to_symint_guardless,
         ConvertIntKey,
         DivideByKey,
         free_symbols,
@@ -161,38 +160,39 @@ def _is_intermediate_tensor_sym_call(node: fx.Node) -> bool:
         elif "val" in node.meta:
             break
 
+    # Note: DO NOT register one _set_node_metadata_hook(_node_metadata_hook)
+    # for each nodes in the graph.
+    # _set_node_metadata_hook is expensive and this can cause compile
+    # time to regress significantly.
     def _node_metadata_hook(
         node: torch.fx.Node,
-        stack_trace: Optional[str] = None,
-        nn_module_stack: Optional[dict[str, Any]] = None,
-        custom: Optional[dict[str, Any]] = None,
+        stack_trace: str | None = None,
+        nn_module_stack: dict[str, Any] | None = None,
+        custom: dict[str, Any] | None = None,
+        skip_val: bool = False,
     ) -> None:
-        fake_args = pytree.tree_map(
-            lambda arg: (
-                _get_example_value(arg) if isinstance(arg, torch.fx.Node) else arg
-            ),
-            node.args,
-        )
-        try:
-            target = node.target
-            if node.op == "call_method":
-                if not isinstance(node.target, str):
-                    raise AssertionError(
-                        f"Expected str target, got {type(node.target)}"
-                    )
-                target = getattr(fake_args[0], node.target)
-                fake_args = fake_args[1:]
-            node.meta[val_key] = target(*fake_args)  # type: ignore[operator]
-        except NotImplementedError:
-            # This can happen when attempting to reify a symbol with an unsupported call_function node,
-            # e.g. with NestedTensors + sym_size.int via match_symbol().
-            # This seems to be fine, as the node gets CSE'd and deleted later in favor of a SymInt graph input.
-            pass
-        except torch.fx.experimental.symbolic_shapes.GuardOnDataDependentSymNode:
-            # This can happen when node args are symints
-            # e.g. test/dynamo/test_export.py -k test_export_preserve_constraints_as_metadata_tensor
-            # aten.sym_constrain_range_for_size(u0)
-            pass
+        if not skip_val:
+            fake_args = pytree.tree_map(
+                lambda arg: (
+                    _get_example_value(arg) if isinstance(arg, torch.fx.Node) else arg
+                ),
+                node.args,
+            )
+            try:
+                target = node.target
+                if node.op == "call_method":
+                    if not isinstance(node.target, str):
+                        raise AssertionError(
+                            f"Expected str target, got {type(node.target)}"
+                        )
+                    target = getattr(fake_args[0], node.target)
+                    fake_args = fake_args[1:]
+                node.meta[val_key] = target(*fake_args)  # type: ignore[operator]
+            except NotImplementedError:
+                # This can happen when attempting to reify a symbol with an unsupported call_function node,
+                # e.g. with NestedTensors + sym_size.int via match_symbol().
+                # This seems to be fine, as the node gets CSE'd and deleted later in favor of a SymInt graph input.
+                pass
         if stack_trace is not None:
             node.meta["stack_trace"] = stack_trace
         if nn_module_stack is not None:
@@ -268,17 +268,28 @@ def add_runtime_asserts(ras):
             else:
                 # Convert the sympy expression into a sequence of FX
                 # nodes
-                res = _sympy_interp(expr_to_proxy, ra.expr).node
-
-                graph.call_function(
-                    torch.ops.aten._assert_scalar.default,
-                    # TODO: use ra.msg here, but it's pretty
-                    # useless right now
-                    (
-                        res,
-                        f"Runtime assertion failed for expression {ra.expr} on node '{res}'",
+                with _set_node_metadata_hook(
+                    gm,
+                    functools.partial(
+                        _node_metadata_hook,
+                        stack_trace=node.meta.get("stack_trace"),
+                        nn_module_stack=node.meta.get("nn_module_stack"),
+                        # nodes added in `apply_runtime_assertion_pass` will have the same annotation
+                        # as the input node to the assertion
+                        custom=node.meta.get("custom"),
                     ),
-                )
+                ):
+                    res = _sympy_interp(expr_to_proxy, ra.expr).node
+
+                    graph.call_function(
+                        torch.ops.aten._assert_scalar.default,
+                        # TODO: use ra.msg here, but it's pretty
+                        # useless right now
+                        (
+                            res,
+                            f"Runtime assertion failed for expression {ra.expr} on node '{res}'",
+                        ),
+                    )
                 added_asserts.add(ra.expr)
 
     nodes = list(graph.nodes)
@@ -286,19 +297,8 @@ def add_runtime_asserts(ras):
         # Placeholders can match symbols, but when we destructure them
         # with size we have to make sure we insert the nodes after all
         # the placeholders
-        with (
-            graph.inserting_before(
-                nodes[i + 1] if node not in placeholders else first_non_placeholder
-            ),
-            _set_node_metadata_hook(
-                gm,
-                functools.partial(
-                    _node_metadata_hook,
-                    stack_trace=node.meta.get("stack_trace"),
-                    nn_module_stack=node.meta.get("nn_module_stack"),
-                    custom=node.meta.get("custom"),
-                ),
-            ),
+        with graph.inserting_before(
+            nodes[i + 1] if node not in placeholders else first_non_placeholder
         ):
             # Unfortunately, this logic still must remain because manual
             # make_fx calls may not explicitly bind all symbolic ints as
@@ -318,7 +318,18 @@ def match_symbol(symint, cb):
                         )
                         and s not in expr_to_proxy
                     ):
-                        expr_to_proxy[s] = fx.Proxy(cb(), tracer=tracer)
+                        with _set_node_metadata_hook(
+                            gm,
+                            functools.partial(
+                                _node_metadata_hook,
+                                stack_trace=node.meta.get("stack_trace"),
+                                nn_module_stack=node.meta.get("nn_module_stack"),
+                                # nodes added in `apply_runtime_assertion_pass` will have the same annotation
+                                # as the input node to the assertion
+                                custom=node.meta.get("custom"),
+                            ),
+                        ):
+                            expr_to_proxy[s] = fx.Proxy(cb(), tracer=tracer)
 
                         log.debug("expr_to_proxy[%s] = %s", s, expr_to_proxy[s])
 
@@ -413,10 +424,21 @@ def has_new_unbacked_bindings():
                     if _is_intermediate_tensor_sym_call(
                         node
                     ):  # reify from input shapes
-                        expr_to_proxy[sym_expr] = _sympy_interp(
-                            expr_to_proxy,
-                            sym_expr,
-                        )  # type: ignore[arg-type]
+                        with _set_node_metadata_hook(
+                            gm,
+                            functools.partial(
+                                _node_metadata_hook,
+                                stack_trace=node.meta.get("stack_trace"),
+                                nn_module_stack=node.meta.get("nn_module_stack"),
+                                # nodes added in `apply_runtime_assertion_pass` will have the same annotation
+                                # as the input node to the assertion
+                                custom=node.meta.get("custom"),
+                            ),
+                        ):
+                            expr_to_proxy[sym_expr] = _sympy_interp(
+                                expr_to_proxy,
+                                sym_expr,
+                            )  # type: ignore[arg-type]
                         # won't try DCE-ing tensor compute here
                     hash_node = expr_to_proxy[sym_expr].node  # type: ignore[arg-type]
                     node.replace_all_uses_with(hash_node)
@@ -512,9 +534,7 @@ def go(node, keypath):
                             )
                         elif isinstance(keypath[0], ConvertIntKey):
                             return go(
-                                graph.call_function(
-                                    cast_symbool_to_symint_guardless, (node,)
-                                ),
+                                graph.call_function(torch.sym_ite, (node, 1, 0)),
                                 keypath[1:],
                             )
                         elif isinstance(keypath[0], DivideByKey):
@@ -536,7 +556,20 @@ def go(node, keypath):
                             raise AssertionError(f"unrecognized keypath {keypath}")
 
                     if s not in expr_to_proxy:
-                        expr_to_proxy[s] = fx.Proxy(go(node, keypath), tracer=tracer)
+                        with _set_node_metadata_hook(
+                            gm,
+                            functools.partial(
+                                _node_metadata_hook,
+                                stack_trace=node.meta.get("stack_trace"),
+                                nn_module_stack=node.meta.get("nn_module_stack"),
+                                # nodes added in `apply_runtime_assertion_pass` will have the same annotation
+                                # as the input node to the assertion
+                                custom=node.meta.get("custom"),
+                            ),
+                        ):
+                            expr_to_proxy[s] = fx.Proxy(
+                                go(node, keypath), tracer=tracer
+                            )
                         log.debug("expr_to_proxy[%s] = %s", s, expr_to_proxy[s])
 
             for i0 in defs:
@@ -601,13 +634,17 @@ def convert(s):
                         except TypeError:
                             return None
 
-                    if (
-                        expr_to_proxy[i0].node.target
-                        is not cast_symbool_to_symint_guardless
+                    with _set_node_metadata_hook(
+                        gm,
+                        functools.partial(
+                            _node_metadata_hook,
+                            stack_trace=node.meta.get("stack_trace"),
+                            nn_module_stack=node.meta.get("nn_module_stack"),
+                            # nodes added in `apply_runtime_assertion_pass` will have the same annotation
+                            # as the input node to the assertion
+                            custom=node.meta.get("custom"),
+                        ),
                     ):
-                        # TODO(pianpwk): calling sym_constrain_range_for_size or adding bound asserts
-                        # raises AOTAutograd errors on cast_symbool_to_symint_guardless
-
                         if (min_val := convert(vr.lower)) is not None:
                             ge = _sympy_interp(expr_to_proxy, i0 >= min_val).node
                             graph.call_function(
diff --git a/torch/fx/passes/shape_prop.py b/torch/fx/passes/shape_prop.py
index 62ea218356138..63e35cae175cd 100644
--- a/torch/fx/passes/shape_prop.py
+++ b/torch/fx/passes/shape_prop.py
@@ -1,7 +1,7 @@
 # mypy: ignore-errors
 
 import traceback
-from typing import Any, NamedTuple, Optional
+from typing import Any, NamedTuple
 
 import torch
 import torch.fx
@@ -26,7 +26,7 @@ class TensorMetadata(NamedTuple):
     dtype: torch.dtype
     requires_grad: bool
     stride: tuple[int, ...]
-    memory_format: Optional[torch.memory_format]
+    memory_format: torch.memory_format | None
 
     # Quantization metadata
     is_quantized: bool
diff --git a/torch/fx/passes/split_module.py b/torch/fx/passes/split_module.py
index e1b0716add840..115c6ad93c643 100644
--- a/torch/fx/passes/split_module.py
+++ b/torch/fx/passes/split_module.py
@@ -3,10 +3,11 @@
 import logging
 from collections import OrderedDict
 from collections.abc import Callable
-from typing import Any, Optional
+from typing import Any
 
 import torch
 from torch.fx._compatibility import compatibility
+from torch.fx._lazy_graph_module import _make_graph_module
 from torch.fx._utils import lazy_format_graph_code
 from torch.fx.graph_module import GraphModule
 from torch.fx.node import Node
@@ -56,12 +57,13 @@ def split_module(
     m: GraphModule,
     root_m: torch.nn.Module,
     split_callback: Callable[[Node], int],
-    qualname_map: Optional[dict[str, str]] = None,
-    keep_original_order: Optional[bool] = False,
-    keep_original_node_name: Optional[bool] = False,
+    qualname_map: dict[str, str] | None = None,
+    keep_original_order: bool | None = False,
+    keep_original_node_name: bool | None = False,
     keep_original_input_name: bool = True,
     *,
-    partition_affix: Optional[str] = None,
+    partition_affix: str | None = None,
+    tuple_return: bool = False,
 ):
     """
     Creates subgraphs out of main graph
@@ -86,6 +88,9 @@ def split_module(
             have the same input names as the original graph.
         partition_affix: Optional[str]: If specified, the submodules' names will contain
             the affix, e.g. "submod_<affix>_<idx>".
+        tuple_return: bool: If True, submodule outputs are always wrapped in a tuple,
+            even when there is only a single output value.  This makes all subgraphs
+            conform to the convention expected by ``torch._inductor.compile_fx``.
 
     Returns:
         GraphModule: the module after split.
@@ -209,7 +214,7 @@ def construct_graph(
     orig_nodes: dict[str, Node] = {}
     symbol_to_node: dict[sympy.Symbol, Node] = {}
 
-    def record_cross_partition_use(def_node: Node, use_node: Optional[Node]):
+    def record_cross_partition_use(def_node: Node, use_node: Node | None):
         from torch.fx.experimental.symbolic_shapes import free_symbols
 
         defined = getattr(def_node, "_fx_partition", None)
@@ -300,7 +305,7 @@ def instantiate_node_partition_mapping(node):
     # 3. last region: we will only insert _enter at the beginning
     # We will do so in the order in which the autocasts were instantiated.
     autocast_regions: OrderedDict[Node, set[int]] = OrderedDict()
-    autocast_exits: dict[Node, Optional[Node]] = {}
+    autocast_exits: dict[Node, Node | None] = {}
 
     active_grad = None
     active_autocasts = set()
@@ -603,15 +608,10 @@ def add_placeholder():
             partition.environment[orig_nodes[name]] for name in partition.outputs
         )
 
-        # skip output node generation if there are no output values
-        num_output_vals = len(output_vals)
-        if num_output_vals == 1:
+        if len(output_vals) == 1 and not tuple_return:
             partition.graph.output(output_vals[0])
-        elif num_output_vals > 1:
-            partition.graph.output(output_vals)
         else:
-            # Invariant - Graph should always have an output node.
-            partition.graph.output(())
+            partition.graph.output(output_vals)
 
         if keep_original_order:
             # first get the attr nodes required by this partition
@@ -638,7 +638,7 @@ def add_placeholder():
                 )
                 already_constructed_attr_nodes.add(node)
 
-        base_mod_attrs[partition.submod_name] = torch.fx.graph_module.GraphModule(
+        base_mod_attrs[partition.submod_name] = _make_graph_module(
             partition.targets, partition.graph
         )  # noqa: B950
 
@@ -649,8 +649,8 @@ def add_placeholder():
         )
 
         num_outputs = len(partition.outputs)
-        if num_outputs > 1:
-            # Unpack multiple return values from submodule
+        if num_outputs > 1 or (num_outputs == 1 and tuple_return):
+            # Unpack return values from submodule
             output_val_proxy = torch.fx.proxy.Proxy(output_val)
             for i, output_name in enumerate(partition.outputs):
                 base_mod_env[output_name] = output_val_proxy[i].node  # type: ignore[index]
@@ -674,7 +674,7 @@ def add_placeholder():
                 torch.fx.graph.map_arg(node.args[0], lambda n: base_mod_env[n.name])
             )  # noqa: B950
 
-    ret = torch.fx.graph_module.GraphModule(base_mod_attrs, base_mod_graph)
+    ret = _make_graph_module(base_mod_attrs, base_mod_graph)
     log.debug(
         "%s",
         lazy_format_graph_code("post split_module", ret, colored=True),
diff --git a/torch/fx/passes/split_utils.py b/torch/fx/passes/split_utils.py
index c828d90b41ec7..08a4c5a685ec7 100644
--- a/torch/fx/passes/split_utils.py
+++ b/torch/fx/passes/split_utils.py
@@ -1,7 +1,6 @@
 # mypy: allow-untyped-defs
 import copy
 from dataclasses import dataclass, field
-from typing import Optional, Union
 
 import torch.fx
 from torch.fx._compatibility import compatibility
@@ -67,7 +66,7 @@ class Component:
     # Mapping from get_attr node in original graph to get_attr node in `graph`.
     getattr_maps: dict[torch.fx.Node, torch.fx.Node] = field(default_factory=dict)
     constructor_args: list[str] = field(default_factory=list)
-    gm: Optional[torch.fx.GraphModule] = None
+    gm: torch.fx.GraphModule | None = None
 
 
 @compatibility(is_backward_compatible=False)
@@ -77,7 +76,7 @@ def split_by_tags(
     return_fqn_mapping: bool = False,
     return_tuple: bool = False,
     GraphModuleCls: type[torch.fx.GraphModule] = torch.fx.GraphModule,
-) -> Union[torch.fx.GraphModule, tuple[torch.fx.GraphModule, dict[str, str]]]:
+) -> torch.fx.GraphModule | tuple[torch.fx.GraphModule, dict[str, str]]:
     """
     Splits a GraphModule using tags on its graph nodes. We honor the order of
     tags. For example, we have tags = ["a", "b", "c"], the function will create
@@ -166,7 +165,7 @@ def flatten(x: torch.fx.node.Argument) -> NodeList:
     main_remapping: dict[torch.fx.Node, torch.fx.Node] = {}
 
     # Output node of original module.
-    output_node: Optional[torch.fx.Node] = None
+    output_node: torch.fx.Node | None = None
 
     # Create a component for each tag, we don't expect to create other components afterwards.
     for tag in tags:
diff --git a/torch/fx/passes/splitter_base.py b/torch/fx/passes/splitter_base.py
index 9347cfad4a897..3de6753228a82 100644
--- a/torch/fx/passes/splitter_base.py
+++ b/torch/fx/passes/splitter_base.py
@@ -6,7 +6,7 @@
 from collections import defaultdict
 from collections.abc import Iterable, Sequence
 from dataclasses import dataclass
-from typing import Any, Literal, NamedTuple, Optional
+from typing import Any, Literal, NamedTuple
 
 import torch
 from torch._logging import trace_structured
@@ -158,7 +158,7 @@ class NodeEvent:
     """
 
     def __init__(
-        self, source: torch.fx.Node, desc: str, dep: Optional[torch.fx.Node] = None
+        self, source: torch.fx.Node, desc: str, dep: torch.fx.Node | None = None
     ):
         self.source = source
         self.desc = desc
@@ -187,7 +187,7 @@ def __init__(self, tracker_mode, dump_prefix):
         self.node_events = {}
         self.writer = print
 
-    def add(self, node: torch.fx.Node, desc: str, dep: Optional[torch.fx.Node] = None):
+    def add(self, node: torch.fx.Node, desc: str, dep: torch.fx.Node | None = None):
         """
         Add a new event to the tracker.
         """
@@ -411,7 +411,7 @@ class FxNetSplitterInternalError(Exception):
 class Subgraph:
     is_acc: bool
     nodes: NodeList
-    device_ordinal: Optional[int] = None
+    device_ordinal: int | None = None
 
 
 @compatibility(is_backward_compatible=False)
@@ -534,7 +534,7 @@ def __init__(
         settings: _SplitterSettingBase,
         non_acc_submodule_name: str = "_run_on_cpu_",
         return_tuple: bool = False,
-        nodes_finder: Optional[FxNetAccNodesFinder] = None,
+        nodes_finder: FxNetAccNodesFinder | None = None,
     ):
         """
         Preprocesses graph before splitting:
@@ -563,6 +563,7 @@ def __init__(
             self.fusions = {}
         else:
             self.fusions = FxNetAccFusionsFinder(module, self.acc_nodes)()
+            self._merge_overlapping_fusions()
 
         # Modify deps to add more deps for fused nodes
         self.deps = self.find_deps()
@@ -620,6 +621,71 @@ def update_deps_for_fusions(self):
                     if user not in fusion:
                         self.deps[user].add(node)
 
+    def _merge_overlapping_fusions(self):
+        """
+        Merge fusion groups that share nodes.
+
+        FxNetAccFusionsFinder can produce overlapping fusion groups when a
+        node is absorbed into multiple groups during expansion. When
+        update_deps_for_fusions() later propagates deps, shared nodes act
+        as bridges between groups, creating bidirectional dependency cycles
+        that crash the splitter with "Subgraph can't be empty".
+
+        This method detects overlapping groups via union-find and merges
+        them into single groups before dep propagation.
+        """
+        if os.environ.get("_SPLITTER_MERGE_OVERLAPPING_FUSIONS", "0") != "1":
+            return
+
+        if not self.fusions:
+            return
+
+        # Collect unique groups by identity.
+        unique_groups: dict[int, NodeSet] = {}
+        for group in self.fusions.values():
+            unique_groups[id(group)] = group
+
+        if len(unique_groups) <= 1:
+            return
+
+        # Map each node to all group IDs it belongs to.
+        node_to_gids: dict[torch.fx.Node, list[int]] = defaultdict(list)
+        for gid, group in unique_groups.items():
+            for node in group:
+                node_to_gids[node].append(gid)
+
+        # Union-find: merge groups that share nodes.
+        parent: dict[int, int] = {gid: gid for gid in unique_groups}
+
+        def find(x: int) -> int:
+            while parent[x] != x:
+                parent[x] = parent[parent[x]]
+                x = parent[x]
+            return x
+
+        needs_merge = False
+        for gids in node_to_gids.values():
+            if len(gids) > 1:
+                root = find(gids[0])
+                for i in range(1, len(gids)):
+                    other = find(gids[i])
+                    if other != root:
+                        parent[other] = root
+                        needs_merge = True
+
+        if not needs_merge:
+            return
+
+        # Build merged groups.
+        merged: dict[int, NodeSet] = defaultdict(set)
+        for gid, group in unique_groups.items():
+            merged[find(gid)].update(group)
+
+        # Rebuild self.fusions so every node points to its merged group.
+        for merged_group in merged.values():
+            for node in merged_group:
+                self.fusions[node] = merged_group
+
     # ===============================================================
     # Helpers for preview
     # ===============================================================
@@ -850,7 +916,7 @@ def get_bytes(node: torch.fx.Node):
     # ===============================================================
 
     def find_reverse_deps(
-        self, tag_id: Optional[int] = None
+        self, tag_id: int | None = None
     ) -> dict[torch.fx.Node, NodeSet]:
         """
         Builds reversed topological node dependencies, if tag_id is specified,
diff --git a/torch/fx/passes/tools_common.py b/torch/fx/passes/tools_common.py
index 17d17b1996cf1..75f76a6d0da18 100644
--- a/torch/fx/passes/tools_common.py
+++ b/torch/fx/passes/tools_common.py
@@ -4,7 +4,7 @@
 import operator
 from collections.abc import Mapping
 from dataclasses import dataclass
-from typing import Any, Optional, Union
+from typing import Any
 
 import torch
 import torch.fx
@@ -21,8 +21,8 @@
     "stable_topological_sort",
 ]
 
-Tensors = Union[tuple[torch.Tensor], list[torch.Tensor]]
-TensorOrTensors = Union[torch.Tensor, Tensors]
+Tensors = tuple[torch.Tensor] | list[torch.Tensor]
+TensorOrTensors = torch.Tensor | Tensors
 NodeList = list[torch.fx.Node]
 NodeSet = set[torch.fx.Node]
 Names = list[str]
@@ -146,8 +146,8 @@ def add_node(self, node):
     def recursive_add_node(
         self,
         fusion_group: "FxNetAccFusionsFinder.FusionGroup",
-        inputs: Union[NodeSet, NodeList],
-        visited: Optional[NodeSet] = None,
+        inputs: NodeSet | NodeList,
+        visited: NodeSet | None = None,
     ):
         """
         Start from inputs and going reverse topological order. If any upstream node
diff --git a/torch/fx/passes/utils/fuser_utils.py b/torch/fx/passes/utils/fuser_utils.py
index 6d80dd6357d1e..4930278afffbd 100644
--- a/torch/fx/passes/utils/fuser_utils.py
+++ b/torch/fx/passes/utils/fuser_utils.py
@@ -1,6 +1,5 @@
 import copy
-from queue import SimpleQueue
-from typing import Optional as _Optional
+import heapq
 
 import torch.fx
 from torch.fx._compatibility import compatibility
@@ -13,27 +12,30 @@
 
 @compatibility(is_backward_compatible=False)
 def topo_sort(nodes: NodeList) -> NodeList:
-    # sort nodes according to the topological order
+    # Stable topological sort: among nodes with no dependency between them,
+    # preserve their relative order in the input list. This uses a min-heap
+    # keyed by original position instead of a FIFO queue.
     indegree_map = dict.fromkeys(nodes, 0)
-    candidates: SimpleQueue[Node] = SimpleQueue()
+    position = {node: i for i, node in enumerate(nodes)}
+    candidates: list[tuple[int, Node]] = []
 
     for node in nodes:
         for n in node.all_input_nodes:
             if n in indegree_map:
                 indegree_map[node] += 1
         if indegree_map[node] == 0:
-            candidates.put(node)
+            heapq.heappush(candidates, (position[node], node))
 
     sorted_nodes: NodeList = []
-    while not candidates.empty():
-        node = candidates.get()
+    while candidates:
+        _, node = heapq.heappop(candidates)
         sorted_nodes.append(node)
 
         for n in node.users:
             if n in indegree_map:
                 indegree_map[n] -= 1
                 if indegree_map[n] == 0:
-                    candidates.put(n)
+                    heapq.heappush(candidates, (position[n], n))
 
     if len(nodes) != len(sorted_nodes):
         raise AssertionError(
@@ -97,7 +99,7 @@ def fuse_as_graphmodule(
     gm: GraphModule,
     nodes: NodeList,
     module_name: str,
-    partition_lookup_table: _Optional[dict[Node, _Optional[int]]] = None,
+    partition_lookup_table: dict[Node, int | None] | None = None,
     *,
     always_return_tuple: bool = False,
 ) -> tuple[GraphModule, tuple[Node, ...], tuple[Node, ...]]:
@@ -220,42 +222,48 @@ def insert_subgm(
     sub_gm: GraphModule,
     orig_inputs: tuple[Node, ...],
     orig_outputs: tuple[Node, ...],
+    insertion_point: Node | None = None,
 ) -> GraphModule:
     # add sub_gm into gm
     submodule_name = sub_gm.__class__.__name__
     gm.add_submodule(submodule_name, sub_gm)
 
-    def last_node(target_nodes: tuple[Node, ...]) -> Node | None:
+    # Use provided insertion point, or fall back to last output node for backwards compat
+    if insertion_point is None:
         for node in reversed(gm.graph.nodes):
-            if node in target_nodes:
-                return node
-        return None
-
-    last_output_node: Node | None = last_node(orig_outputs)
-    if last_output_node is None:
-        raise AssertionError("last_output_node is None")
+            if node in orig_outputs:
+                insertion_point = node
+                break
+        if insertion_point is None:
+            raise AssertionError(
+                "Cannot determine insertion point: no insertion_point provided and "
+                "orig_outputs is empty. Pass the last partition node as insertion_point."
+            )
 
     # Create a call_module node in main graph.
-    with gm.graph.inserting_after(last_output_node):
+    with gm.graph.inserting_after(insertion_point):
         module_node = gm.graph.call_module(
             submodule_name, args=orig_inputs, kwargs=None
         )
         output_node = sub_gm.graph.output_node()
 
-    next_node = module_node.next
-    with gm.graph.inserting_before(next_node):
-        if len(orig_outputs) == 1 and not isinstance(output_node.args[0], tuple):
-            # main_remapping[comp.orig_outputs[0]] = module_node
-            orig_outputs[0].replace_all_uses_with(module_node, propagate_meta=True)
-        else:
-            for i, orig_output in enumerate(orig_outputs):
-                # Use Proxy to record getitem access.
-                proxy_out = torch.fx.Proxy(module_node)[i].node  # type: ignore[index]
-                orig_output.replace_all_uses_with(proxy_out, propagate_meta=True)
-
-            module_node.meta["val"] = tuple(
-                orig_output.meta.get("val", None) for orig_output in orig_outputs
-            )
+    # Replace uses of original outputs with the fused module outputs.
+    # If there are no external outputs, skip replacement (nothing to replace).
+    if orig_outputs:
+        next_node = module_node.next
+        with gm.graph.inserting_before(next_node):
+            if len(orig_outputs) == 1 and not isinstance(output_node.args[0], tuple):
+                # main_remapping[comp.orig_outputs[0]] = module_node
+                orig_outputs[0].replace_all_uses_with(module_node, propagate_meta=True)
+            else:
+                for i, orig_output in enumerate(orig_outputs):
+                    # Use Proxy to record getitem access.
+                    proxy_out = torch.fx.Proxy(module_node)[i].node  # type: ignore[index]
+                    orig_output.replace_all_uses_with(proxy_out, propagate_meta=True)
+
+                module_node.meta["val"] = tuple(
+                    orig_output.meta.get("val", None) for orig_output in orig_outputs
+                )
     return gm
 
 
@@ -269,7 +277,7 @@ def erase_nodes(gm: GraphModule, nodes: NodeList) -> None:
 @compatibility(is_backward_compatible=False)
 def fuse_by_partitions(
     gm: GraphModule,
-    partitions: list[dict[Node, _Optional[int]]],
+    partitions: list[dict[Node, int | None]],
     prefix: str = "fused_",
     always_return_tuple: bool = False,
 ) -> GraphModule:
@@ -285,7 +293,7 @@ def fuse_by_partitions(
             always_return_tuple=always_return_tuple,
         )
 
-        insert_subgm(gm, sub_gm, orig_inputs, orig_outputs)
+        insert_subgm(gm, sub_gm, orig_inputs, orig_outputs, sorted_nodes[-1])
 
         erase_nodes(gm, sorted_nodes)
 
diff --git a/torch/fx/passes/utils/matcher_utils.py b/torch/fx/passes/utils/matcher_utils.py
index 961d5c7a43bb4..e0ebe687e95f5 100644
--- a/torch/fx/passes/utils/matcher_utils.py
+++ b/torch/fx/passes/utils/matcher_utils.py
@@ -4,7 +4,7 @@
 import os
 from collections import defaultdict
 from dataclasses import dataclass, field
-from typing import Any, Union
+from typing import Any
 
 import torch
 from torch.fx import Graph, Node
@@ -251,7 +251,7 @@ def _match_nodes(
         # match for `gn`
         match_found = True
 
-        def _match_args(args1: Union[list, tuple], args2: Union[list, tuple]) -> bool:
+        def _match_args(args1: list | tuple, args2: list | tuple) -> bool:
             if len(args1) != len(args2):
                 return False
 
diff --git a/torch/fx/passes/utils/source_matcher_utils.py b/torch/fx/passes/utils/source_matcher_utils.py
index 82259b8a36ab7..ab0b372ee27a1 100644
--- a/torch/fx/passes/utils/source_matcher_utils.py
+++ b/torch/fx/passes/utils/source_matcher_utils.py
@@ -2,7 +2,7 @@
 import os
 from collections.abc import Callable
 from dataclasses import dataclass, field
-from typing import Any, Optional
+from typing import Any
 
 from torch.fx._compatibility import compatibility
 from torch.fx.graph import Graph
@@ -56,7 +56,7 @@ class SourcePartition:
 def get_source_partitions(
     graph: Graph,
     wanted_sources: list[Any],
-    filter_fn: Optional[Callable[[Node], bool]] = None,
+    filter_fn: Callable[[Node], bool] | None = None,
 ) -> dict[Any, list[SourcePartition]]:
     """
     Args:
@@ -72,6 +72,11 @@ def get_source_partitions(
     """
     modules: dict[type, dict[str, list[Node]]] = {}
 
+    def add_to_partition(src: Any, fqn: str, node: Node) -> None:
+        diff_modules = modules.setdefault(src, {})
+        partition = diff_modules.setdefault(fqn, [])
+        partition.append(node)
+
     for node in graph.nodes:
         # The metadata source_fn should contain a tuple of a unique name for the
         # source, and the source function if the node is decomposed from a
@@ -82,22 +87,45 @@ def get_source_partitions(
         # be different from "source_fn_stack", for example for the add_ node
         # decomposed from batch norm. We should remove the check on "source_fn_stack"
         # after we fix "torch_fn". T199561090
-        if (source_fn_st := node.meta.get("source_fn_stack", None)) is None and (
-            torch_fn := node.meta.get("torch_fn", None)
-        ) is not None:
-            node_fqn, source_fn = torch_fn
-            source_fn_name = source_fn.split(".")[1]
-            if source_fn_name in wanted_sources:
-                diff_modules = modules.setdefault(source_fn_name, {})
-                partition = diff_modules.setdefault(node_fqn, [])
-                partition.append(node)
-
-        if (source_fn_st := node.meta.get("source_fn_stack", None)) is not None:
+        source_fn_st = node.meta.get("source_fn_stack", None)
+        if source_fn_st is None:
+            matched = False
+            torch_fn = node.meta.get("torch_fn", None)
+            if torch_fn is not None:
+                node_fqn, source_fn = torch_fn
+                source_fn_name = source_fn.split(".")[1]
+                if source_fn_name in wanted_sources:
+                    add_to_partition(source_fn_name, node_fqn, node)
+                    matched = True
+            # Fallback: when source_fn_stack is not populated (e.g. strict=False export),
+            # use nn_module_stack to resolve the originating module type.
+            # Only apply to call_function nodes to avoid incorrectly including
+            # placeholder, get_attr, or output nodes in partitions.
+            if not matched and node.op == "call_function":
+                nn_module_stack = node.meta.get("nn_module_stack", None)
+                if nn_module_stack:
+                    # Get the innermost module (last entry in the ordered dict)
+                    innermost_fqn, innermost_cls = list(nn_module_stack.values())[-1]
+                    for src in wanted_sources:
+                        if isinstance(src, type):
+                            if isinstance(innermost_cls, type) and issubclass(
+                                innermost_cls, src
+                            ):
+                                add_to_partition(src, innermost_fqn, node)
+                                break
+                            elif isinstance(innermost_cls, str):
+                                src_str = src.__module__ + "." + src.__qualname__
+                                if innermost_cls == src_str:
+                                    add_to_partition(src, innermost_fqn, node)
+                                    break
+                        elif innermost_cls == src:
+                            add_to_partition(src, innermost_fqn, node)
+                            break
+
+        if source_fn_st is not None:
             source_fn = source_fn_st[-1]
             if source_fn[1] in wanted_sources:
-                diff_modules = modules.setdefault(source_fn[1], {})
-                partition = diff_modules.setdefault(source_fn[0], [])
-                partition.append(node)
+                add_to_partition(source_fn[1], source_fn[0], node)
 
     def make_partition(nodes: list[Node], module_type: type) -> SourcePartition:
         input_nodes = set()
diff --git a/torch/fx/proxy.py b/torch/fx/proxy.py
index 9fb76f5dc2503..811f58e5317d3 100644
--- a/torch/fx/proxy.py
+++ b/torch/fx/proxy.py
@@ -1,5 +1,3 @@
-# mypy: ignore-errors
-
 import collections
 import copy
 import dis
@@ -9,10 +7,12 @@
 import operator
 import sys
 import traceback
+import types
 from collections import OrderedDict
 from collections.abc import Callable, Iterator
 from dataclasses import fields, is_dataclass
-from typing import Any, Optional, TypeVar
+from typing import Any, cast, TypeVar
+from typing_extensions import Never
 
 import torch
 import torch.fx.traceback as fx_traceback
@@ -117,7 +117,7 @@ def forward(self, x):
 
     """
 
-    def __init__(self, module_path: str, module_type: Any):
+    def __init__(self, module_path: str, module_type: Any) -> None:
         super().__init__()
         self.module_path = module_path
         self.module_type = module_type
@@ -134,7 +134,7 @@ def __init__(
         self,
         scope: Scope,
         current_scope: Scope,
-    ):
+    ) -> None:
         super().__init__()
         # Keep a copy of prev scope to restore on exit
         self._prev_scope = copy.copy(scope)
@@ -144,10 +144,10 @@ def __init__(
         # Save a reference so we can restore it
         self._scope = scope
 
-    def __enter__(self):
+    def __enter__(self) -> Scope:
         return self._scope
 
-    def __exit__(self, *args):
+    def __exit__(self, *args: object) -> None:
         self._scope.module_path = self._prev_scope.module_path
         self._scope.module_type = self._prev_scope.module_type
         return
@@ -205,8 +205,8 @@ def create_node(
         target: Target,
         args: tuple[Argument, ...],
         kwargs: dict[str, Argument],
-        name: Optional[str] = None,
-        type_expr: Optional[Any] = None,
+        name: str | None = None,
+        type_expr: Any | None = None,
     ) -> Node:
         """
         Inserts a graph node given target, args, kwargs, and name.
@@ -217,7 +217,8 @@ def create_node(
         """
 
         if kind == "call_function" and self.check_mutable_operations:
-            check_for_mutable_operation(target, args, kwargs)
+            target_fn = cast(Callable[..., Any], target)
+            check_for_mutable_operation(target_fn, args, kwargs)
 
         node = self.graph.create_node(kind, target, args, kwargs, name, type_expr)
         # TODO node_name_to_scope will be depreciated in favor of
@@ -237,6 +238,7 @@ def create_node(
 
                 if fx_traceback.GRADIENT_ACC_SPECIAL_STACK in stack_trace:
                     node.meta["is_gradient_acc"] = True
+                    node.meta["autograd_backward"] = True
 
             # Explicitly set the stack_trace, nn_module_stack and source_fn on the node.meta
             # If other meta fields are needed, they can be added here
@@ -254,7 +256,7 @@ def create_node(
             # See Note [Functionalization View Replay Annotation]
             # Overriding some node meta with the original node meta of the
             # regenerated node.
-            replay_node: Node = fx_traceback.get_current_replay_node()
+            replay_node: Node | None = fx_traceback.get_current_replay_node()
             if replay_node is not None:
                 node.meta["is_functional_regenerated"] = True
                 if "custom" in replay_node.meta:
@@ -262,6 +264,9 @@ def create_node(
                 if "stack_trace" in replay_node.meta:
                     node.stack_trace = replay_node.meta.get("stack_trace")
 
+            if current_meta.get("autograd_backward", False):
+                node.meta["autograd_backward"] = True
+
         elif self.module_stack:
             node.meta["nn_module_stack"] = copy.copy(self.module_stack)
 
@@ -280,7 +285,7 @@ def _filter_traceback_frames(
     ) -> traceback.StackSummary:
         # This method can be overridden to customize the frame filtering logic
         # for the recorded stack trace
-        user_frames = []
+        user_frames: list[traceback.FrameSummary] = []
         if self._record_forward_stack_traces_only:
             user_frames = [
                 frame
@@ -301,7 +306,7 @@ def _filter_traceback_frames(
             # Not having a "forward" call in the stacktrace implies the
             # stacktrace will probably be irrelevant
             if first_forward == -1:
-                user_frames = []
+                user_frames: list[traceback.FrameSummary] = []
 
         from torch.fx.experimental.symbolic_shapes import uninteresting_files
 
@@ -324,11 +329,10 @@ def create_proxy(
         target: Target,
         args: tuple[Any, ...],
         kwargs: dict[str, Any],
-        name: Optional[str] = None,
-        type_expr: Optional[Any] = None,
-        # fix noqa when updating bc tests
-        proxy_factory_fn: Callable[[Node], "Proxy"] = None,  # noqa: RUF013
-    ):
+        name: str | None = None,
+        type_expr: Any | None = None,
+        proxy_factory_fn: Callable[[Node], "Proxy"] | None = None,
+    ) -> "Proxy":
         """
         Create a Node from the given arguments, then return the Node
         wrapped in a Proxy object.
@@ -355,7 +359,7 @@ def create_proxy(
 
         return proxy
 
-    def _find_user_frame(self):
+    def _find_user_frame(self) -> types.FrameType | None:
         """
         Find the Python stack frame executing the user code during
         symbolic tracing.
@@ -440,14 +444,14 @@ def create_arg(self, a: Any) -> Argument:
             )
 
         elif isinstance(a, range):
-            return range(
-                self.create_arg(a.start),
-                self.create_arg(a.stop),
-                self.create_arg(a.step),
+            return range(  # pyrefly: ignore[no-matching-overload]
+                self.create_arg(a.start),  # pyrefly: ignore[bad-argument-type]
+                self.create_arg(a.stop),  # pyrefly: ignore[bad-argument-type]
+                self.create_arg(a.step),  # pyrefly: ignore[bad-argument-type]
             )
 
         elif isinstance(a, (torch._ops.OpOverload, torch._ops.HigherOrderOperator)):
-            return a
+            return a  # pyrefly: ignore[bad-return]
 
         elif is_opaque_value_type(type(a)):
             return a
@@ -460,7 +464,7 @@ def create_arg(self, a: Any) -> Argument:
             return self.create_node("call_function", a.__class__, (), kwargs)
 
         elif isinstance(a, (*base_types, enum.Enum)) or a is None or a is ...:
-            return a
+            return a  # pyrefly: ignore[bad-return]
 
         raise NotImplementedError(f"argument of type: {type(a)}")
 
@@ -494,7 +498,7 @@ def iter(self, obj: "Proxy") -> Iterator:
         )
 
     @compatibility(is_backward_compatible=True)
-    def keys(self, obj: "Proxy") -> Any:
+    def keys(self, obj: "Proxy") -> "Proxy":
         """Called when a proxy object is has the keys() method called.
         This is what happens when ** is called on a proxy. This should return an
         iterator it ** is suppose to work in your custom tracer.
@@ -502,7 +506,7 @@ def keys(self, obj: "Proxy") -> Any:
         return Attribute(obj, "keys")()
 
 
-def _get_seq_nr(node_name: str = ""):
+def _get_seq_nr(node_name: str = "") -> int | None:
     """
     Returns the seq_nr node meta for the current proxy node that we're creating.
     The seq_nr number in node meta is related to but not the same as the "sequence number"
@@ -547,7 +551,7 @@ def _get_seq_nr(node_name: str = ""):
         # See Note [Functionalization View Replay Annotation]
         # Overriding some node meta with the original node meta of the
         # regenerated node.
-        replay_node: Node = fx_traceback.get_current_replay_node()
+        replay_node: Node | None = fx_traceback.get_current_replay_node()
         if replay_node is not None:
             if "seq_nr" in replay_node.meta:
                 annotation_log.debug("%s: seq_nr from replay_node", node_name)
@@ -559,16 +563,16 @@ def _get_seq_nr(node_name: str = ""):
 # used in Proxy object when just appending to the graph while not tracing.
 @compatibility(is_backward_compatible=True)
 class GraphAppendingTracer(TracerBase):
-    def __init__(self, graph: Graph):
+    def __init__(self, graph: Graph) -> None:
         super().__init__()
         self.graph = graph
         self.scope = Scope("", None)
-        self.module_stack = collections.OrderedDict()
-        self.node_name_to_scope = {}
+        self.module_stack: OrderedDict[str, tuple[str, Any]] = collections.OrderedDict()
+        self.node_name_to_scope: dict[str, tuple[str, type]] = {}
 
 
 @compatibility(is_backward_compatible=False)
-def assert_fn(x):
+def assert_fn(x: object) -> None:
     if not x:
         raise AssertionError("Assertion failed")
 
@@ -609,7 +613,7 @@ class Proxy:
     """
 
     @compatibility(is_backward_compatible=True)
-    def __init__(self, node: Node, tracer: "Optional[TracerBase]" = None):
+    def __init__(self, node: Node, tracer: "TracerBase | None" = None) -> None:
         if tracer is None:
             # This allows you to create a Proxy object around a raw Node
             tracer = GraphAppendingTracer(node.graph)
@@ -619,21 +623,21 @@ def __init__(self, node: Node, tracer: "Optional[TracerBase]" = None):
     def __repr__(self) -> str:
         return f"Proxy({self.node.name})"
 
-    def __getattr__(self, k) -> "Attribute":
+    def __getattr__(self, k: str) -> "Attribute":
         # note: not added to the graph yet, if this is a method call
         # we peephole optimize to the method invocation
         return Attribute(self, k)
 
-    def __getstate__(self) -> dict:
+    def __getstate__(self) -> dict[str, Any]:
         return self.__dict__
 
-    def __deepcopy__(self, memo) -> dict:
+    def __deepcopy__(self, memo: dict[int, Any]) -> "Proxy":
         # We have to explicitly override this method, because otherwise deepcopy
         # will go to __getattr__(self, "__deepcopy__") and return a
         # Attribute(__deepcopy__), and may go into an infinite loop in some cases.
         import copy
 
-        new_dict = {}
+        new_dict: dict[str, Any] = {}
         for k, v in self.__dict__.items():
             try:
                 new_obj = copy.deepcopy(v, memo)
@@ -655,11 +659,11 @@ def __deepcopy__(self, memo) -> dict:
             new_proxy.__dict__[k] = v
         return new_proxy
 
-    def __setstate__(self, d):
+    def __setstate__(self, d: dict[str, Any]) -> None:
         # This is called when being unpickled/loaded.
         self.__dict__ = d
 
-    def __call__(self, *args, **kwargs) -> "Proxy":
+    def __call__(self, *args: Any, **kwargs: Any) -> "Proxy":
         return self.tracer.create_proxy(
             "call_method", "__call__", (self,) + args, kwargs
         )
@@ -686,7 +690,7 @@ def __iter__(self) -> Iterator["Proxy"]:
 
         return self.tracer.iter(self)
 
-    def __abs__(self):
+    def __abs__(self) -> "Proxy":
         return self.tracer.create_proxy("call_function", operator.abs, (self,), {})
 
     def __bool__(self) -> bool:
@@ -725,10 +729,10 @@ def __bool__(self) -> bool:
         return self.tracer.to_bool(self)
 
     @compatibility(is_backward_compatible=True)
-    def keys(self):
+    def keys(self) -> "Proxy":
         return self.tracer.keys(self)
 
-    def __len__(self):
+    def __len__(self) -> int:
         raise RuntimeError(
             "'len' is not supported in symbolic tracing by default. If you want "
             "this call to be recorded, please call torch.fx.wrap('len') at "
@@ -736,13 +740,19 @@ def __len__(self):
         )
 
     @classmethod
-    def __torch_function__(cls, orig_method, types, args=None, kwargs=None):
+    def __torch_function__(
+        cls,
+        orig_method: Callable[..., Any],
+        types: tuple[type, ...],
+        args: tuple[Any, ...] | None = None,
+        kwargs: dict[str, Any] | None = None,
+    ) -> "Proxy":
         args = args if args else ()
         kwargs = kwargs if kwargs else {}
 
-        tracers: dict[Any, None] = {}
+        tracers: dict[TracerBase, None] = {}
 
-        def find_tracer(a):
+        def find_tracer(a: Any) -> None:
             if isinstance(a, cls):
                 tracers[a.tracer] = None
 
@@ -789,8 +799,8 @@ class MetaProxy(Proxy):
     """
 
     def __init__(
-        self, node: Node, tracer: "Optional[TracerBase]" = None, fake_mode=None
-    ):
+        self, node: Node, tracer: "TracerBase | None" = None, fake_mode: Any = None
+    ) -> None:
         super().__init__(node, tracer)
         self.fake_mode = fake_mode
 
@@ -798,7 +808,13 @@ def __repr__(self) -> str:
         return f"MetaProxy({self.node.name})"
 
     @classmethod
-    def __torch_function__(cls, orig_method, types, args=None, kwargs=None):
+    def __torch_function__(
+        cls,
+        orig_method: Callable[..., Any],
+        types: tuple[type, ...],
+        args: tuple[Any, ...] | None = None,
+        kwargs: dict[str, Any] | None = None,
+    ) -> "MetaProxy":
         args = args if args else ()
         kwargs = kwargs if kwargs else {}
 
@@ -825,14 +841,14 @@ def __torch_function__(cls, orig_method, types, args=None, kwargs=None):
 @compatibility(is_backward_compatible=True)
 class Attribute(Proxy):
     @compatibility(is_backward_compatible=True)
-    def __init__(self, root: Proxy, attr: str):
+    def __init__(self, root: Proxy, attr: str) -> None:
         self.root = root
         self.attr = attr
         self.tracer = root.tracer
-        self._node: Optional[Node] = None
+        self._node: Node | None = None
 
     @property
-    def node(self):
+    def node(self) -> Node:  # pyrefly: ignore[bad-override]
         # the node for attributes is added lazily, since most will just be method calls
         # which do not rely on the getitem call
         if self._node is None:
@@ -841,7 +857,7 @@ def node(self):
             ).node
         return self._node
 
-    def __call__(self, *args, **kwargs):
+    def __call__(self, *args: Any, **kwargs: Any) -> "Proxy":
         return self.tracer.create_proxy(
             "call_method", self.attr, (self.root,) + args, kwargs
         )
@@ -855,7 +871,9 @@ class ParameterProxy(Proxy):
     so that conditional tests on these attributes will not throw exception during tracing
     """
 
-    def __init__(self, tracer: TracerBase, node: Node, name, param):
+    def __init__(
+        self, tracer: TracerBase, node: Node, name: str, param: torch.nn.Parameter
+    ) -> None:
         super().__init__(node, tracer)
         if not isinstance(param, torch.nn.Parameter):
             raise AssertionError(f"Expected Parameter, got {type(param)}")
@@ -866,30 +884,30 @@ def __repr__(self) -> str:
         return f"ParameterProxy({self.name})"
 
     @property
-    def shape(self):
+    def shape(self) -> torch.Size:
         return self.param.shape
 
-    def size(self):
+    def size(self) -> torch.Size:
         return self.param.size()
 
-    def dim(self):
+    def dim(self) -> int:
         return self.param.dim()
 
     @property
-    def ndim(self):
+    def ndim(self) -> int:
         return self.param.ndim
 
-    def numel(self):
+    def numel(self) -> int:
         return self.param.numel()
 
-    def nelement(self):
+    def nelement(self) -> int:
         return self.param.nelement()
 
 
 for method in magic_methods:
 
-    def _scope(method):
-        def impl(*args, **kwargs):
+    def _scope(method: str) -> None:
+        def impl(*args: Any, **kwargs: Any) -> "Proxy":
             tracer = args[0].tracer
             target = getattr(operator, method)
             return tracer.create_proxy("call_function", target, args, kwargs)
@@ -901,10 +919,10 @@ def impl(*args, **kwargs):
     _scope(method)
 
 
-def _define_reflectable(orig_method_name):
+def _define_reflectable(orig_method_name: str) -> None:
     method_name = f"__r{orig_method_name.strip('_')}__"
 
-    def impl(self, rhs):
+    def impl(self: "Proxy", rhs: Any) -> "Proxy":
         target = getattr(operator, orig_method_name)
         return self.tracer.create_proxy("call_function", target, (rhs, self), {})
 
@@ -917,15 +935,15 @@ def impl(self, rhs):
     _define_reflectable(orig_method_name)
 
 
-def _no_nodes_error(arg):
+def _no_nodes_error(arg: Argument) -> Never:
     raise RuntimeError(
         "Keys for dictionaries used as an argument cannot contain a "
         f"Node. Got key: {arg}"
     )
 
 
-def _create_arg_dict(self, a):
-    r = {}
+def _create_arg_dict(self: TracerBase, a: dict[Any, Any]) -> dict[Any, Argument]:
+    r: dict[Any, Argument] = {}
     for k, v in a.items():
         if not isinstance(k, str):
             # Check for invalid dict keys. We do not want a Proxy to appear
diff --git a/torch/fx/subgraph_rewriter.py b/torch/fx/subgraph_rewriter.py
index bf69fb6e12288..71d884cc86e70 100644
--- a/torch/fx/subgraph_rewriter.py
+++ b/torch/fx/subgraph_rewriter.py
@@ -1,7 +1,7 @@
 import copy
 from collections.abc import Callable
 from dataclasses import dataclass
-from typing import Any, NamedTuple, Optional, TYPE_CHECKING, Union
+from typing import Any, NamedTuple, TYPE_CHECKING
 
 import torch
 
@@ -48,7 +48,7 @@ def _replace_attributes(gm: GraphModule, replacement: torch.nn.Module) -> None:
     if isinstance(replacement, GraphModule):
         replacement.graph.lint()
 
-    def try_get_attr(gm: torch.nn.Module, target: str) -> Optional[Any]:
+    def try_get_attr(gm: torch.nn.Module, target: str) -> Any | None:
         module_path, _, attr_name = target.rpartition(".")
         try:
             mod: torch.nn.Module = gm.get_submodule(module_path)
@@ -95,8 +95,8 @@ def try_get_attr(gm: torch.nn.Module, target: str) -> Optional[Any]:
 @compatibility(is_backward_compatible=True)
 def replace_pattern(
     gm: GraphModule,
-    pattern: Union[Callable, GraphModule],
-    replacement: Union[Callable, GraphModule],
+    pattern: Callable[..., Any] | GraphModule,
+    replacement: Callable[..., Any] | GraphModule,
 ) -> list[Match]:
     """
     Matches all possible non-overlapping sets of operators and their
@@ -225,16 +225,13 @@ def forward(self, x, w1, w2):
 @compatibility(is_backward_compatible=False)
 def replace_pattern_with_filters(
     gm: GraphModule,
-    pattern: Union[Callable, Graph, GraphModule],
-    replacement: Union[Callable, Graph, GraphModule, None] = None,
-    match_filters: Optional[
-        list[Callable[["InternalMatch", Graph, Graph], bool]]
-    ] = None,
+    pattern: Callable[..., Any] | Graph | GraphModule,
+    replacement: Callable[..., Any] | Graph | GraphModule | None = None,
+    match_filters: list[Callable[["InternalMatch", Graph, Graph], bool]] | None = None,
     ignore_literals: bool = False,
     # Placed at the end to avoid breaking backward compatibility
-    replacement_callback: Optional[
-        Callable[["InternalMatch", Graph, Graph], Graph]
-    ] = None,
+    replacement_callback: Callable[["InternalMatch", Graph, Graph], Graph]
+    | None = None,
     node_name_match: str = "",
 ) -> list[ReplacedPatterns]:
     """
@@ -264,16 +261,13 @@ def replace_pattern_with_filters(
 
 def _replace_pattern(
     gm: GraphModule,
-    pattern: Union[Callable, Graph, GraphModule],
-    replacement: Union[Callable, Graph, GraphModule, None] = None,
-    match_filters: Optional[
-        list[Callable[["InternalMatch", Graph, Graph], bool]]
-    ] = None,
+    pattern: Callable[..., Any] | Graph | GraphModule,
+    replacement: Callable[..., Any] | Graph | GraphModule | None = None,
+    match_filters: list[Callable[["InternalMatch", Graph, Graph], bool]] | None = None,
     ignore_literals: bool = False,
     # Placed at the end to avoid breaking backward compatibility
-    replacement_callback: Optional[
-        Callable[["InternalMatch", Graph, Graph], Graph]
-    ] = None,
+    replacement_callback: Callable[["InternalMatch", Graph, Graph], Graph]
+    | None = None,
     node_name_match: str = "",
 ) -> list[ReplacedPatterns]:
     from torch.fx.passes.utils.matcher_utils import InternalMatch, SubgraphMatcher
@@ -423,7 +417,9 @@ def _replace_pattern(
                 f"{len(copied_returning_nodes)}"  # pyrefly: ignore [bad-argument-type]
             )
         for gn, copied_node in zip(match.returning_nodes, copied_returning_nodes):  # type: ignore[arg-type]
+            # pyrefly: ignore [bad-argument-type]
             gn.replace_all_uses_with(copied_node)
+            # pyrefly: ignore [unsupported-operation]
             match_changed_node[gn] = copied_node
         # Remove the original nodes
         for node in reversed(pattern_graph.nodes):
diff --git a/torch/fx/tensor_type.py b/torch/fx/tensor_type.py
index 4f375e461ef28..6e04c7b1cc858 100644
--- a/torch/fx/tensor_type.py
+++ b/torch/fx/tensor_type.py
@@ -1,9 +1,21 @@
-# mypy: allow-untyped-defs
+from __future__ import annotations
+
+from typing import Any, TYPE_CHECKING
+
 from torch.fx.experimental.unification import Var  # type: ignore[attr-defined]
 
 from ._compatibility import compatibility
 
 
+if TYPE_CHECKING:
+    from collections.abc import Sequence
+
+    from torch.fx.experimental.migrate_gradual_types.constraint import DVar
+
+
+__all__ = ["Dyn", "TensorType", "is_consistent", "is_more_precise"]
+
+
 @compatibility(is_backward_compatible=False)
 class TensorType:
     """
@@ -14,21 +26,23 @@ def forward(self, x:TensorType((1,2,3, Dyn)), y:TensorType((1,2,3, Dyn))):
                 return torch.add(x, y)
     """
 
-    def __init__(self, dim):
+    __args__: Sequence[DVar | int | _DynType]
+
+    def __init__(self, dim: Sequence[Any]) -> None:
         self.__origin__ = TensorType
         self.__args__ = dim
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         return f"TensorType[{self.__args__}]"
 
-    def __eq__(self, other):
+    def __eq__(self, other: object) -> bool:
         if isinstance(other, self.__class__):
             return list(self.__args__) == list(other.__args__)
         else:
             return False
 
     @staticmethod
-    def __class_getitem__(*args):
+    def __class_getitem__(*args: object) -> TensorType:
         if len(args) == 1 and isinstance(args[0], tuple):
             args = args[0]
         return TensorType(tuple(args))
@@ -42,13 +56,13 @@ class _DynType:
     def __init__(self) -> None:
         self.__name__ = "_DynType"
 
-    def __eq__(self, other):
+    def __eq__(self, other: object) -> bool:
         return isinstance(other, self.__class__)
 
-    def __str__(self):
+    def __str__(self) -> str:
         return "Dyn"
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         return "Dyn"
 
 
@@ -56,7 +70,7 @@ def __repr__(self):
 
 
 @compatibility(is_backward_compatible=False)
-def is_consistent(t1, t2):
+def is_consistent(t1: object, t2: object) -> bool:
     """
     A binary relation denoted by ~ that determines if t1 is consistent with t2.
     The relation is reflexive, symmetric but not transitive.
@@ -84,7 +98,7 @@ def is_consistent(t1, t2):
 
 
 @compatibility(is_backward_compatible=False)
-def is_more_precise(t1, t2):
+def is_more_precise(t1: object, t2: object) -> bool:
     """
     A binary relation denoted by <= that determines if t1 is more precise than t2.
     The relation is reflexive and transitive.
diff --git a/torch/fx/traceback.py b/torch/fx/traceback.py
index 0889c586089b4..485d2fbfa6a00 100644
--- a/torch/fx/traceback.py
+++ b/torch/fx/traceback.py
@@ -1,11 +1,11 @@
-# mypy: allow-untyped-defs
 import copy
 import logging
 import traceback
 from collections import defaultdict
+from collections.abc import Callable, Iterator
 from contextlib import contextmanager
 from enum import Enum
-from typing import Any, Optional, Union
+from typing import Any, Optional, ParamSpec, TypeVar, Union
 
 from torch._utils_internal import signpost_event
 
@@ -15,6 +15,9 @@
 from .node import Node
 
 
+_P = ParamSpec("_P")
+_R = TypeVar("_R")
+
 log = logging.getLogger(__name__)
 
 __all__ = [
@@ -36,7 +39,7 @@
 ]
 
 current_meta: dict[str, Any] = {}
-current_replay_node: Optional[Node] = None
+current_replay_node: Node | None = None
 # Preserve the node meta fields in torch.fx.proxy._COPY_META_FIELDS
 should_preserve_node_meta = False
 # Preserve the "seq_nr" node meta field
@@ -82,7 +85,7 @@ class NodeSource:
     """
 
     class NodeInfo:
-        def __init__(self, name: str, target: str, graph_id: int):
+        def __init__(self, name: str, target: str, graph_id: int) -> None:
             self.name = name
             self.target = target
             self.graph_id = graph_id
@@ -91,15 +94,15 @@ def __init__(self, name: str, target: str, graph_id: int):
     action: list["NodeSourceAction"]
     from_node: list["NodeSource"]
     node_info: Optional["NodeInfo"]
-    _dict: Optional[dict[str, Any]]
-    _action_string: Optional[str]
+    _dict: dict[str, Any] | None
+    _action_string: str | None
 
     def __init__(
         self,
-        node: Optional[Node],
+        node: Node | None,
         pass_name: str = "",
-        action: Optional[Union["NodeSourceAction", list["NodeSourceAction"]]] = None,
-    ):
+        action: Union["NodeSourceAction", list["NodeSourceAction"]] | None = None,
+    ) -> None:
         self.pass_name = pass_name
 
         if action is None:
@@ -124,8 +127,8 @@ def __init__(
             self.from_node = []
 
         # cache the action string and dict representation for performance.
-        self._action_string: Optional[str] = None
-        self._dict: Optional[dict[str, Any]] = None
+        self._action_string: str | None = None
+        self._dict: dict[str, Any] | None = None
 
     @property
     def name(self) -> str:
@@ -139,15 +142,15 @@ def target(self) -> str:
     def graph_id(self) -> int:
         return self.node_info.graph_id if self.node_info else -1
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         return self.print_readable()
 
-    def _get_action_string(self):
+    def _get_action_string(self) -> str:
         if self._action_string is None:
             self._action_string = "+".join([a.name.lower() for a in self.action])
         return self._action_string
 
-    def print_readable(self, indent=0):
+    def print_readable(self, indent: int = 0) -> str:
         if indent > 9:
             return ""
         result = ""
@@ -160,7 +163,7 @@ def print_readable(self, indent=0):
             result += item.print_readable(indent + 1)
         return result
 
-    def to_dict(self) -> dict:
+    def to_dict(self) -> dict[str, Any]:
         if self._dict is None:
             # Convert the object to a dictionary
             action_string = self._get_action_string()
@@ -177,15 +180,15 @@ def to_dict(self) -> dict:
             raise AssertionError("_dict is None after initialization")
         return self._dict
 
-    def __eq__(self, other: object):
+    def __eq__(self, other: object) -> bool:
         if not isinstance(other, NodeSource):
             return False
         return self.to_dict() == other.to_dict()
 
-    def __hash__(self):
+    def __hash__(self) -> int:
         # Create a hash based on the dictionary representation
         # We need to convert the dict to a hashable form
-        def _make_hashable(obj):
+        def _make_hashable(obj: Any) -> Any:
             if isinstance(obj, dict):
                 return tuple(sorted((k, _make_hashable(v)) for k, v in obj.items()))
             elif isinstance(obj, list):
@@ -196,7 +199,7 @@ def _make_hashable(obj):
         return hash(_make_hashable(self.to_dict()))
 
     @classmethod
-    def _from_dict(cls, d: Optional[dict]) -> Optional["NodeSource"]:
+    def _from_dict(cls, d: dict[str, Any] | None) -> Optional["NodeSource"]:
         """
         Recursively deserialize from_node metadata from dictionary data.
         It is used to deserialize the from_node field from serialized metadata.
@@ -253,7 +256,7 @@ def _from_dict(cls, d: Optional[dict]) -> Optional["NodeSource"]:
 
 @compatibility(is_backward_compatible=False)
 @contextmanager
-def preserve_node_meta(enable=True):
+def preserve_node_meta(enable: bool = True) -> Iterator[None]:
     global should_preserve_node_meta
     global current_meta
     saved_should_preserve_node_meta = should_preserve_node_meta
@@ -268,7 +271,7 @@ def preserve_node_meta(enable=True):
 
 
 @contextmanager
-def _preserve_node_seq_nr(preserve_seq_nr=True):
+def _preserve_node_seq_nr(preserve_seq_nr: bool = True) -> Iterator[None]:
     """
     Temporarily enables or disables the preservation of node.meta["seq_nr"] in the
     tracing context.
@@ -284,16 +287,21 @@ def _preserve_node_seq_nr(preserve_seq_nr=True):
 
 
 @compatibility(is_backward_compatible=False)
-def set_stack_trace(stack: list[str]):
+def set_stack_trace(stack: list[str]) -> None:
     global current_meta
 
-    if should_preserve_node_meta and stack:
-        current_meta["stack_trace"] = "".join(stack)
+    if should_preserve_node_meta:
+        if stack:
+            current_meta["stack_trace"] = "".join(stack)
+        else:
+            # when the stack is empty, we explicitly clear the stack_trace to avoid
+            # propagating it to future node.˙
+            current_meta.pop("stack_trace", None)
 
 
 @compatibility(is_backward_compatible=False)
 @contextmanager
-def annotate(annotation_dict: dict):
+def annotate(annotation_dict: dict[str, Any]) -> Iterator[None]:
     """
     Temporarily adds custom annotations to the current tracing context.
     The fx_node produced from this tracing context will have the
@@ -333,7 +341,7 @@ def annotate(annotation_dict: dict):
 
     try:
         if not has_custom:
-            current_meta["custom"] = {}
+            current_meta["custom"] = dict[str, Any]()
 
         # Update with all key-value pairs from the input dict
         current_meta["custom"].update(annotation_dict)
@@ -347,7 +355,9 @@ def annotate(annotation_dict: dict):
 
 
 @compatibility(is_backward_compatible=False)
-def annotate_fn(annotation_dict: dict):
+def annotate_fn(
+    annotation_dict: dict[str, Any],
+) -> Callable[[Callable[_P, _R]], Callable[_P, _R]]:
     """
     A decorator that wraps a function with the annotate context manager.
     Use this when you want to annotate an entire function instead of a specific code block.
@@ -372,19 +382,53 @@ def annotate_fn(annotation_dict: dict):
     """
     from functools import wraps
 
-    def decorator(func):
+    def decorator(func: Callable[_P, _R]) -> Callable[_P, _R]:
         @wraps(func)
-        def wrapper(*args, **kwargs):
+        # NB: Do not annotate with _P.args/_P.kwargs here. Dynamo guards on
+        # the identity of ParamSpec annotation objects, causing guard failures.
+        def wrapper(*args: Any, **kwargs: Any) -> Any:
             with annotate(annotation_dict):
                 return func(*args, **kwargs)
 
-        return wrapper
+        return wrapper  # type: ignore[return-value]
 
     return decorator
 
 
+@contextmanager
+def _set_autograd_backward(enable: bool = True) -> Iterator[None]:
+    global current_meta
+
+    had_autograd_backward = "autograd_backward" in current_meta
+    old_autograd_backward = current_meta.get("autograd_backward", False)
+
+    if enable:
+        _mark_autograd_backward()
+    try:
+        yield
+    finally:
+        if had_autograd_backward:
+            current_meta["autograd_backward"] = old_autograd_backward
+        else:
+            _reset_autograd_backward()
+
+
+@compatibility(is_backward_compatible=False)
+def _mark_autograd_backward() -> None:
+    global current_meta
+
+    current_meta["autograd_backward"] = True
+
+
 @compatibility(is_backward_compatible=False)
-def set_grad_fn_seq_nr(seq_nr):
+def _reset_autograd_backward() -> None:
+    global current_meta
+
+    current_meta.pop("autograd_backward", None)
+
+
+@compatibility(is_backward_compatible=False)
+def set_grad_fn_seq_nr(seq_nr: int) -> None:
     global current_meta
 
     if should_preserve_node_meta:
@@ -396,7 +440,7 @@ def set_grad_fn_seq_nr(seq_nr):
 
 
 @compatibility(is_backward_compatible=False)
-def reset_grad_fn_seq_nr():
+def reset_grad_fn_seq_nr() -> None:
     # NB: reset state properly, this would be helpful towards supporting
     #     reentrant autograd if we actually wanted to do that.
     global current_meta
@@ -432,7 +476,7 @@ def _is_preserving_node_seq_nr() -> bool:
 
 @compatibility(is_backward_compatible=False)
 @contextmanager
-def set_current_meta(node, pass_name=""):
+def set_current_meta(node: Node, pass_name: str = "") -> Iterator[None]:
     global current_meta
     if should_preserve_node_meta and node.meta:
         saved_meta = current_meta
@@ -460,7 +504,7 @@ def get_current_meta() -> dict[str, Any]:
 
 @compatibility(is_backward_compatible=False)
 @contextmanager
-def set_current_replay_node(node):
+def set_current_replay_node(node: Node | None) -> Iterator[None]:
     """
     Set the currently replay node. If `current_replay_node` is not None,
     then we're re-generating the `current_replay_node` in FunctionalTensorMode.
@@ -476,7 +520,7 @@ def set_current_replay_node(node):
 
 
 @compatibility(is_backward_compatible=False)
-def get_current_replay_node():
+def get_current_replay_node() -> Node | None:
     """
     Get the currently replay node
     """
@@ -517,7 +561,7 @@ def _get_custom_metadata(gm: GraphModule) -> str:
     if not isinstance(gm, GraphModule):
         raise AssertionError(f"Expected GraphModule, got {type(gm)}")
 
-    def helper(gm: GraphModule):
+    def helper(gm: GraphModule) -> list[Any]:
         custom_metadata = []
         for node in gm.graph.nodes:
             if hasattr(node, "meta") and node.meta.get("custom", None):
@@ -525,14 +569,17 @@ def helper(gm: GraphModule):
             if node.op == "get_attr" and isinstance(
                 getattr(gm, node.target), GraphModule
             ):
-                custom_metadata.append(helper(getattr(gm, node.target)))
+                custom_metadata.append(
+                    # pyrefly: ignore[bad-argument-type]
+                    helper(getattr(gm, node.target))
+                )
         return custom_metadata
 
     return "\n".join(str(x) for x in helper(gm))
 
 
 def _get_ordered_seq_nr_groups(
-    gm: Union[GraphModule, list[GraphModule]],
+    gm: GraphModule | list[GraphModule],
 ) -> list[list[str]]:
     """
     Group call_function nodes by seq_nr, order by seq_nr value,
diff --git a/torch/headeronly/CMakeLists.txt b/torch/headeronly/CMakeLists.txt
index 93d2d7802b528..1f880929341bd 100644
--- a/torch/headeronly/CMakeLists.txt
+++ b/torch/headeronly/CMakeLists.txt
@@ -2,7 +2,7 @@ cmake_minimum_required(VERSION 3.27 FATAL_ERROR)
 
 project(headeronly CXX)
 
-set(CMAKE_CXX_STANDARD 17 CACHE STRING "The C++ standard whose features are requested to build this target.")
+set(CMAKE_CXX_STANDARD 20 CACHE STRING "The C++ standard whose features are requested to build this target.")
 set(CMAKE_EXPORT_COMPILE_COMMANDS ON)
 
 # Main build file for torch/headeronly, except there's no build cuz this lib is header-only!
@@ -30,6 +30,8 @@ add_library(headeronly INTERFACE ${HEADERONLY_HEADERS})
 
 install(FILES ${CMAKE_BINARY_DIR}/torch/headeronly/macros/cmake_macros.h
         DESTINATION include/torch/headeronly/macros)
+install(FILES ${CMAKE_BINARY_DIR}/torch/headeronly/core/enum_tag.h
+        DESTINATION include/torch/headeronly/core)
 
 if(NOT BUILD_LIBTORCHLESS)
   # ---[ Installation copied from c10/CMakeLists.txt
diff --git a/torch/headeronly/build.bzl b/torch/headeronly/build.bzl
index c44a8dc563d04..2efabb1d0d172 100644
--- a/torch/headeronly/build.bzl
+++ b/torch/headeronly/build.bzl
@@ -25,13 +25,71 @@ def define_targets(rules):
 
     rules.genrule(**genrule_args)
 
-    rules.cc_library(
-        name = "torch_headeronly",
-        hdrs = rules.glob([
-            "**/*.h"
-        ]) + ["version.h.in"],
-        visibility = ["//visibility:public"],
-        deps = [
-            "//torch/headeronly/macros",
-        ],
-    )
+    # Generate core/enum_tag.h from tags.yaml using the real torchgen codegen.
+    # The main generate_aten rule (in //:BUILD.bazel) also produces this file
+    # via headeronly_fm, but Bazel can't declare genrule outputs across package
+    # boundaries.
+    if is_buck:
+        # In Buck, aten_srcs_for_headeronly doesn't exist.  Use the xplat
+        # aten_src_path target which is available from any cell.
+        genrule_args = {
+            "name": "enum_tag_h",
+            "srcs": ["fbsource//xplat/caffe2:aten_src_path"],
+            "outs": ["core/enum_tag.h"],
+            "cmd": "mkdir -p $TMP/__aten_unused $TMP/__hdr && " +
+                   "$(execpath //torchgen:gen) " +
+                   "--source-path $(location fbsource//xplat/caffe2:aten_src_path)/aten/src/ATen " +
+                   "--install_dir $TMP/__aten_unused " +
+                   "--headeronly-install-dir $TMP/__hdr " +
+                   "--generate headers && " +
+                   "cp $TMP/__hdr/enum_tag.h $OUT",
+            "tools": ["//torchgen:gen"],
+        }
+    else:
+        genrule_args = {
+            "name": "enum_tag_h",
+            "srcs": ["//:aten_srcs_for_headeronly"],
+            "outs": ["core/enum_tag.h"],
+            "cmd": "$(execpath //torchgen:gen) " +
+                   "--source-path aten/src/ATen " +
+                   "--install_dir $(RULEDIR)/__aten_unused " +
+                   "--headeronly-install-dir $(@D) " +
+                   "--generate headers",
+            "tools": ["//torchgen:gen"],
+        }
+
+    if not is_buck:
+        genrule_args["visibility"] = ["//visibility:public"]
+
+    rules.genrule(**genrule_args)
+
+    if is_buck:
+        # Use Buck-native cxx_library so we can set header_namespace.
+        # The Bazel cc_library maps genrule outputs by target name rather
+        # than output path, breaking the torch/headeronly/core/enum_tag.h
+        # include path.
+        native.cxx_library(
+            name = "torch_headeronly",
+            header_namespace = "torch/headeronly",
+            exported_headers = {
+                h: h for h in native.glob(["**/*.h"], exclude = ["version.h"])
+            } | {
+                "version.h": ":version_h",
+                "core/enum_tag.h": ":enum_tag_h",
+            },
+            visibility = ["PUBLIC"],
+            exported_deps = [
+                "//caffe2/torch/headeronly/macros:macros",
+            ],
+        )
+    else:
+        rules.cc_library(
+            name = "torch_headeronly",
+            hdrs = rules.glob([
+                "**/*.h"
+            ]) + ["version.h.in", ":enum_tag_h"],
+            visibility = ["//visibility:public"],
+            deps = [
+                "//torch/headeronly/macros",
+            ],
+        )
diff --git a/torch/headeronly/core/Dispatch_v2.h b/torch/headeronly/core/Dispatch_v2.h
index 47ca18ce79c54..aa1bc1d9e9f72 100644
--- a/torch/headeronly/core/Dispatch_v2.h
+++ b/torch/headeronly/core/Dispatch_v2.h
@@ -54,6 +54,11 @@
   torch::headeronly::ScalarType::UInt16,     \
       torch::headeronly::ScalarType::UInt32, \
       torch::headeronly::ScalarType::UInt64
+#define AT_OPAQUE_TYPES                                                       \
+  torch::headeronly::ScalarType::Byte, torch::headeronly::ScalarType::UInt16, \
+      torch::headeronly::ScalarType::UInt32,                                  \
+      torch::headeronly::ScalarType::UInt64,                                  \
+      torch::headeronly::ScalarType::ComplexDouble
 #define AT_INTEGRAL_TYPES_V2 \
   AT_EXPAND(AT_INTEGRAL_TYPES), AT_EXPAND(AT_BAREBONES_UNSIGNED_TYPES)
 #define AT_COMPLEX_TYPES                        \
diff --git a/torch/headeronly/macros/Macros.h b/torch/headeronly/macros/Macros.h
index 63aa0d20d8e54..cef99df3f566f 100644
--- a/torch/headeronly/macros/Macros.h
+++ b/torch/headeronly/macros/Macros.h
@@ -325,41 +325,88 @@ constexpr uint32_t CUDA_THREADS_PER_BLOCK_FALLBACK = 256;
 #define C10_HIP_HOST_DEVICE
 #endif
 
-#if defined(USE_ROCM)
 // C10_WARP_SIZE is only allowed for device code.
-// Host code _must_ use at::cuda::warp_size()
+// Host code dynamically-sized launch configs _must_ use at::cuda::warp_size().
+// Host or device statically-sized arrays _must_ use either
+// C10_WARP_SIZE_UPPER_BOUND or C10_WARP_SIZE_LOWER_BOUND, as needed.
+//
 // HIP header used to define warpSize as a constexpr that was either 32 or 64
 // depending on the target device, and then always set it to 64 for host code.
-// Host pass of HIP compiler needs C10_WARP_SIZE defined to _something_ so we
-// set it to something unreasonable to trigger obvious host code errors.
-
+// For a time, that allowed C10_WARP_SIZE to be defined like so:
+//
+// #ifdef USE_ROCM
+// #define C10_WARP_SIZE warpSize
+// #else
+// #define C10_WARP_SIZE 32
+// #endif
+//
+// In ROCm 7, warpSize is no longer constexpr, matching CUDA behavior.
+// We can now only use warpSize for C10_WARP_SIZE in device code and this is
+// enforced by using __device__ in its definition.  In host code where
+// C10_WARP_SIZE was previously used as a compile-time constant, this will now
+// cause a compile-time error.
+//
+// If an array was previously expected to be sized at compile-time using
+// C10_WARP_SIZE, users must now use either C10_WARP_SIZE_UPPER_BOUND or
+// C10_WARP_SIZE_LOWER_BOUND depending on the situation.
+//
+// If C10_WARP_SIZE was previously used to determine kernel launch sizes, users
+// must now use at::cuda::warp_size() for the dynamic runtime query.
+//
+// Unfortunately, C10_WARP_SIZE has been public and available for both host and
+// device since approximately 2019, so forcing it to be device-only would break
+// existing code in the wild.
+#if defined(USE_ROCM)
 namespace at::cuda {
 TORCH_CUDA_CPP_API int warp_size();
 }
-#ifdef __HIPCC__
-static inline int __host__ C10_WARP_SIZE_INTERNAL() {
+#if defined(__HIPCC__)
+static __host__ inline int C10_WARP_SIZE_INTERNAL() {
   return at::cuda::warp_size();
 }
-
-static inline constexpr int __device__ C10_WARP_SIZE_INTERNAL() {
+// NOTE: __device__ C10_WARP_SIZE_INTERNAL
+// For __SPIRV__, we must use dynamic warpSize. When not targeting __SPIRV__,
+// we can use constexpr. This matches prior behavior. We preserve this for
+// backward compatibility instead of forcing old code to use dynamic warpSize
+// and losing constexpr. However, compiling for --offload-arch=amdgcnspirv
+// could expose where C10_WARP_SIZE was used incorrectly where the dynamic
+// warpSize is not allowed.
+#if defined(__SPIRV__)
+static __device__ inline int C10_WARP_SIZE_INTERNAL() {
+  return warpSize;
+}
+#else // __SPIRV__
+static __device__ inline constexpr int C10_WARP_SIZE_INTERNAL() {
 #if defined(__GFX9__)
   return 64;
 #else // __GFX9__
   return 32;
 #endif // __GFX9__
 }
-#else // __HIPCC__
+#endif // __SPIRV__
+#if defined(__SPIRV__)
+#define C10_WARP_SIZE_LOWER_BOUND 32
+#define C10_WARP_SIZE_UPPER_BOUND 64
+#elif defined(__GFX9__)
+#define C10_WARP_SIZE_LOWER_BOUND 64
+#define C10_WARP_SIZE_UPPER_BOUND 64
+#else
+#define C10_WARP_SIZE_LOWER_BOUND 32
+#define C10_WARP_SIZE_UPPER_BOUND 32
+#endif
+#else // !__HIPCC__
 static inline int C10_WARP_SIZE_INTERNAL() {
   return at::cuda::warp_size();
 }
+#define C10_WARP_SIZE_LOWER_BOUND 32
+#define C10_WARP_SIZE_UPPER_BOUND 64
 #endif // __HIPCC__
-
 #define C10_WARP_SIZE (C10_WARP_SIZE_INTERNAL())
-#define C10_WARP_SIZE_STATIC 64
-
-#else // defined(USE_ROCM)
+#else // !USE_ROCM
 #define C10_WARP_SIZE 32
-#endif
+#define C10_WARP_SIZE_LOWER_BOUND 32
+#define C10_WARP_SIZE_UPPER_BOUND 32
+#endif // USE_ROCM
 
 #if defined(_MSC_VER) && _MSC_VER <= 1900
 #define __func__ __FUNCTION__
@@ -629,7 +676,7 @@ __host__ __device__
 // This macro is used to find older C++ compilers
 // that don't support move optimization for return values.
 
-#if (defined(__GNUC__) && __GNUC__ < 13) || \
+#if (defined(__GNUC__) && __GNUC__ < 13 && __cplusplus < 202002L) || \
     (defined(__clang_major__) && __clang_major__ < 13)
 #define C10_RETURN_MOVE_IF_OLD_COMPILER 1
 #else
diff --git a/torch/headeronly/ovrsource_defs.bzl b/torch/headeronly/ovrsource_defs.bzl
index 56fa9596585d7..aca50b6ed1a93 100644
--- a/torch/headeronly/ovrsource_defs.bzl
+++ b/torch/headeronly/ovrsource_defs.bzl
@@ -97,6 +97,10 @@ def define_ovrsource_targets():
     oxx_static_library(
         name = "torch_headeronly_ovrsource",
         compatible_with = cpu_supported_platforms,
+        header_namespace = "",
+        public_generated_headers = {
+            "torch/headeronly/core/enum_tag.h": ":generate-enum-tag-header[enum_tag.h]",
+        },
         exported_deps = select({
             "DEFAULT": [":torch_headeronly_full_ovrsource"],
             "ovr_config//os:android": [":torch_headeronly_mobile_ovrsource"],
diff --git a/torch/headeronly/util/BFloat16.h b/torch/headeronly/util/BFloat16.h
index 64479ba36f125..9aa08c265bd2c 100644
--- a/torch/headeronly/util/BFloat16.h
+++ b/torch/headeronly/util/BFloat16.h
@@ -12,7 +12,7 @@
 #include <iosfwd>
 #include <ostream>
 
-#if defined(__CUDACC__) && !defined(USE_ROCM)
+#if defined(__CUDACC__) && (!defined(USE_ROCM) || (TORCH_HIP_VERSION >= 702))
 #include <cuda_bf16.h>
 #endif
 
@@ -46,7 +46,7 @@ struct alignas(2) BFloat16 {
   /* implicit */ inline C10_HOST_DEVICE BFloat16(float value);
   inline C10_HOST_DEVICE operator float() const;
 
-#if defined(__CUDACC__) && !defined(USE_ROCM)
+#if defined(__CUDACC__) && (!defined(USE_ROCM) || (TORCH_HIP_VERSION >= 702))
   inline C10_HOST_DEVICE BFloat16(const __nv_bfloat16& value);
   explicit inline C10_HOST_DEVICE operator __nv_bfloat16() const;
 #endif
@@ -124,8 +124,9 @@ C10_CLANG_DIAGNOSTIC_IGNORE("-Wimplicit-int-float-conversion")
 /// Constructors
 inline C10_HOST_DEVICE BFloat16::BFloat16(float value)
     :
-#if defined(__CUDACC__) && !defined(USE_ROCM) && defined(__CUDA_ARCH__) && \
-    __CUDA_ARCH__ >= 800
+#if defined(__CUDACC__) &&                                                   \
+    (!defined(USE_ROCM) && defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800 || \
+     defined(USE_ROCM) && (TORCH_HIP_VERSION >= 702))
       x(__bfloat16_as_ushort(__float2bfloat16(value)))
 #elif defined(__SYCL_DEVICE_ONLY__) && \
     defined(SYCL_EXT_ONEAPI_BFLOAT16_MATH_FUNCTIONS)
@@ -139,7 +140,7 @@ inline C10_HOST_DEVICE BFloat16::BFloat16(float value)
 
 /// Implicit conversions
 inline C10_HOST_DEVICE BFloat16::operator float() const {
-#if defined(__CUDACC__) && !defined(USE_ROCM)
+#if defined(__CUDACC__) && (!defined(USE_ROCM) || (TORCH_HIP_VERSION >= 702))
   return __bfloat162float(*reinterpret_cast<const __nv_bfloat16*>(&x));
 #elif defined(__SYCL_DEVICE_ONLY__) && \
     defined(SYCL_EXT_ONEAPI_BFLOAT16_MATH_FUNCTIONS)
@@ -149,7 +150,7 @@ inline C10_HOST_DEVICE BFloat16::operator float() const {
 #endif
 }
 
-#if defined(__CUDACC__) && !defined(USE_ROCM)
+#if defined(__CUDACC__) && (!defined(USE_ROCM) || (TORCH_HIP_VERSION >= 702))
 inline C10_HOST_DEVICE BFloat16::BFloat16(const __nv_bfloat16& value) {
   x = *reinterpret_cast<const unsigned short*>(&value);
 }
diff --git a/torch/headeronly/util/Float8_e4m3fn.h b/torch/headeronly/util/Float8_e4m3fn.h
index b35fb1a7aa85d..06d70e6b4982c 100644
--- a/torch/headeronly/util/Float8_e4m3fn.h
+++ b/torch/headeronly/util/Float8_e4m3fn.h
@@ -49,6 +49,7 @@ struct alignas(1) Float8_e4m3fn {
   inline C10_HOST_DEVICE Float8_e4m3fn(float value);
   inline C10_HOST_DEVICE operator float() const;
   inline C10_HOST_DEVICE bool isnan() const;
+  inline C10_HOST_DEVICE bool isinf() const;
 };
 
 inline std::ostream& operator<<(std::ostream& out, const Float8_e4m3fn& value) {
@@ -256,6 +257,11 @@ inline C10_HOST_DEVICE bool Float8_e4m3fn::isnan() const {
   return (x & 0b01111111) == 0b01111111;
 }
 
+inline C10_HOST_DEVICE bool Float8_e4m3fn::isinf() const {
+  // Note: fp8e4m3fn does not have infinity, so this always returns false.
+  return false;
+}
+
 /// Arithmetic
 
 inline C10_HOST_DEVICE Float8_e4m3fn
diff --git a/torch/headeronly/util/Float8_e4m3fnuz.h b/torch/headeronly/util/Float8_e4m3fnuz.h
index e361a2f92a2a5..4c11b3be05593 100644
--- a/torch/headeronly/util/Float8_e4m3fnuz.h
+++ b/torch/headeronly/util/Float8_e4m3fnuz.h
@@ -52,6 +52,7 @@ struct alignas(1) Float8_e4m3fnuz {
   inline C10_HOST_DEVICE Float8_e4m3fnuz(float value);
   inline C10_HOST_DEVICE operator float() const;
   inline C10_HOST_DEVICE bool isnan() const;
+  inline C10_HOST_DEVICE bool isinf() const;
 };
 
 inline std::ostream& operator<<(
@@ -160,6 +161,11 @@ inline C10_HOST_DEVICE bool Float8_e4m3fnuz::isnan() const {
   return x == 0b10000000;
 }
 
+inline C10_HOST_DEVICE bool Float8_e4m3fnuz::isinf() const {
+  // Note: fp8e4m3fnuz does not have infinity, so this always returns false.
+  return false;
+}
+
 /// Arithmetic
 
 inline C10_HOST_DEVICE Float8_e4m3fnuz
diff --git a/torch/jit/_builtins.py b/torch/jit/_builtins.py
index f61fadf375bf1..34caf60599a94 100644
--- a/torch/jit/_builtins.py
+++ b/torch/jit/_builtins.py
@@ -3,7 +3,6 @@
 import math
 import warnings
 from collections import OrderedDict
-from typing import Optional
 
 import torch
 import torch.backends.cudnn as cudnn
@@ -16,7 +15,7 @@
 )
 
 
-_builtin_table: Optional[dict[int, str]] = None
+_builtin_table: dict[int, str] | None = None
 
 _modules_containing_builtins = (
     torch,
diff --git a/torch/jit/_decompositions.py b/torch/jit/_decompositions.py
index c4bbffbcb9ad6..0e3d5a7f9824c 100644
--- a/torch/jit/_decompositions.py
+++ b/torch/jit/_decompositions.py
@@ -7,7 +7,7 @@
 import inspect
 import warnings
 from collections.abc import Callable
-from typing import Optional, TypeVar
+from typing import TypeVar
 from typing_extensions import ParamSpec
 
 from torch.types import Number
@@ -68,7 +68,7 @@ def signatures_match(decomposition_sig, torch_op_sig):
 
 def register_decomposition(
     aten_op: torch._ops.OpOverload,
-    registry: Optional[dict[str, torch.jit.ScriptFunction]] = None,
+    registry: dict[str, torch.jit.ScriptFunction] | None = None,
 ) -> Callable[[Callable[_P, _T]], Callable[_P, _T]]:
     def decomposition_decorator(f: Callable[_P, _T]) -> Callable[_P, _T]:
         nonlocal registry
@@ -104,8 +104,8 @@ def decomposition_decorator(f: Callable[_P, _T]) -> Callable[_P, _T]:
 @register_decomposition(aten.var.correction)
 def var_decomposition(
     input: Tensor,
-    dim: Optional[list[int]] = None,
-    correction: Optional[Number] = None,
+    dim: list[int] | None = None,
+    correction: Number | None = None,
     keepdim: bool = False,
 ) -> Tensor:
     if dim is None:
diff --git a/torch/jit/_freeze.py b/torch/jit/_freeze.py
index 6dac4d2de21b0..e2344dc8a7bec 100644
--- a/torch/jit/_freeze.py
+++ b/torch/jit/_freeze.py
@@ -6,14 +6,13 @@
 """
 
 import warnings
-from typing import Optional
 
 import torch
 from torch.jit._script import RecursiveScriptModule, ScriptModule
 
 
 def freeze(
-    mod, preserved_attrs: Optional[list[str]] = None, optimize_numerics: bool = True
+    mod, preserved_attrs: list[str] | None = None, optimize_numerics: bool = True
 ):
     r"""Freeze ScriptModule, inline submodules, and attributes as constants.
 
@@ -132,7 +131,7 @@ def forward(self, input):
 
 
 def run_frozen_optimizations(
-    mod, optimize_numerics: bool = True, preserved_methods: Optional[list[str]] = None
+    mod, optimize_numerics: bool = True, preserved_methods: list[str] | None = None
 ) -> None:
     r"""
     Run a series of optimizations looking for patterns that occur in frozen graphs.
@@ -191,7 +190,7 @@ def run_frozen_optimizations(
 
 
 def optimize_for_inference(
-    mod: ScriptModule, other_methods: Optional[list[str]] = None
+    mod: ScriptModule, other_methods: list[str] | None = None
 ) -> ScriptModule:
     """
     Perform a set of optimization passes to optimize a model for the purposes of inference.
diff --git a/torch/jit/_ir_utils.py b/torch/jit/_ir_utils.py
index 7a775717de07f..3b2c55288ebde 100644
--- a/torch/jit/_ir_utils.py
+++ b/torch/jit/_ir_utils.py
@@ -1,5 +1,4 @@
 from types import TracebackType
-from typing import Optional, Union
 
 import torch
 
@@ -8,7 +7,7 @@ class _InsertPoint:
     def __init__(
         self,
         insert_point_graph: torch._C.Graph,
-        insert_point: Union[torch._C.Node, torch._C.Block],
+        insert_point: torch._C.Node | torch._C.Block,
     ) -> None:
         self.insert_point = insert_point
         self.g = insert_point_graph
@@ -20,14 +19,14 @@ def __enter__(self) -> None:
 
     def __exit__(
         self,
-        exc_type: Optional[type[BaseException]],
-        exc_val: Optional[BaseException],
-        exc_tb: Optional[TracebackType],
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: TracebackType | None,
     ) -> None:
         self.g.setInsertPoint(self.prev_insert_point)
 
 
 def insert_point_guard(
-    self: torch._C.Graph, insert_point: Union[torch._C.Node, torch._C.Block]
+    self: torch._C.Graph, insert_point: torch._C.Node | torch._C.Block
 ) -> _InsertPoint:
     return _InsertPoint(self, insert_point)
diff --git a/torch/jit/_monkeytype_config.py b/torch/jit/_monkeytype_config.py
index 196c21ac4525d..c4cf1089494ce 100644
--- a/torch/jit/_monkeytype_config.py
+++ b/torch/jit/_monkeytype_config.py
@@ -6,7 +6,6 @@
 from collections.abc import Iterable
 from pathlib import Path
 from types import CodeType
-from typing import Optional
 
 import torch
 
@@ -105,7 +104,7 @@ def add(self, traces: Iterable[CallTrace]) -> None:
         def filter(
             self,
             qualified_name: str,
-            qualname_prefix: Optional[str] = None,
+            qualname_prefix: str | None = None,
             limit: int = 2000,
         ) -> list[CallTraceThunk]:
             return self.trace_records[qualified_name]
@@ -154,7 +153,7 @@ def trace_logger(self) -> JitTypeTraceStoreLogger:
         def trace_store(self) -> CallTraceStore:
             return self.s
 
-        def code_filter(self) -> Optional[CodeFilter]:
+        def code_filter(self) -> CodeFilter | None:
             return jit_code_filter
 
 else:
diff --git a/torch/jit/_pickle.py b/torch/jit/_pickle.py
index 9d6b64a52a983..0d73254f74994 100644
--- a/torch/jit/_pickle.py
+++ b/torch/jit/_pickle.py
@@ -7,8 +7,6 @@
 # a type attached and restored via `restore_type_tag` below. The legacy
 # functions should stick around for backwards-compatibility.
 
-from typing import Union
-
 
 def build_intlist(data: list[int]) -> list[int]:
     return data
@@ -26,7 +24,7 @@ def build_boollist(data: list[bool]) -> list[bool]:
     return data
 
 
-def build_tensor_from_id(data: Union[int, object]) -> Union[int, None]:
+def build_tensor_from_id(data: int | object) -> int | None:
     if isinstance(data, int):
         # just the id, can't really do anything
         return data
diff --git a/torch/jit/_script.py b/torch/jit/_script.py
index d6c404756f82e..fb723c62ff658 100644
--- a/torch/jit/_script.py
+++ b/torch/jit/_script.py
@@ -16,7 +16,7 @@
 import sys
 import warnings
 from collections.abc import Callable, Iterator, Mapping, Sequence
-from typing import Any, TypeVar, Union
+from typing import Any, TypeVar
 from typing_extensions import deprecated, Self
 
 import torch
@@ -1136,7 +1136,7 @@ def _script_impl(
     optimize=None,
     _frames_up=0,
     _rcb=None,
-    example_inputs: Union[list[tuple], dict[Callable, list[tuple]], None] = None,
+    example_inputs: list[tuple] | dict[Callable, list[tuple]] | None = None,
 ):
     global type_trace_db
 
@@ -1275,7 +1275,7 @@ def script(
     optimize: None = None,
     _frames_up: int = 0,
     _rcb: Callable[[str], Any] | None = None,
-    example_inputs: Union[list[tuple], dict[Callable, list[tuple]], None] = None,
+    example_inputs: list[tuple] | dict[Callable, list[tuple]] | None = None,
 ) -> Any:
     r"""Script the function.
 
diff --git a/torch/jit/_trace.py b/torch/jit/_trace.py
index d195e24a43211..ca2d85a76bc3a 100644
--- a/torch/jit/_trace.py
+++ b/torch/jit/_trace.py
@@ -19,7 +19,7 @@
 import warnings
 from collections.abc import Callable
 from enum import Enum
-from typing import Any, Optional, TypeVar
+from typing import Any, TypeVar
 from typing_extensions import ParamSpec
 
 import torch
@@ -1037,7 +1037,7 @@ def forward(self, x):
     return traced_func
 
 
-_trace_module_map: Optional[dict[Any, Any]] = None
+_trace_module_map: dict[Any, Any] | None = None
 
 
 def trace_module(
diff --git a/torch/lib/libshm/CMakeLists.txt b/torch/lib/libshm/CMakeLists.txt
index c3cd26fea7bf3..688e7e7929147 100644
--- a/torch/lib/libshm/CMakeLists.txt
+++ b/torch/lib/libshm/CMakeLists.txt
@@ -21,7 +21,7 @@ target_include_directories(shm PUBLIC
 set_target_properties(shm PROPERTIES
   PREFIX "lib"
   IMPORT_PREFIX "lib"
-  CXX_STANDARD 17)
+  CXX_STANDARD 20)
 target_link_libraries(shm PRIVATE ${TORCH_CPU_LIB})
 
 if(UNIX AND NOT APPLE)
diff --git a/torch/lib/libshm/manager.cpp b/torch/lib/libshm/manager.cpp
index 5d2c318c25142..3084f794154a9 100644
--- a/torch/lib/libshm/manager.cpp
+++ b/torch/lib/libshm/manager.cpp
@@ -46,12 +46,7 @@ static void register_fd(int fd) {
 }
 
 static void unregister_fd(int fd) {
-  pollfds.erase(
-      std::remove_if(
-          pollfds.begin(),
-          pollfds.end(),
-          [fd](const struct pollfd& pfd) { return pfd.fd == fd; }),
-      pollfds.end());
+  std::erase_if(pollfds, [fd](const struct pollfd& pfd) { return pfd.fd == fd; });
   client_sessions.erase(fd);
 }
 
diff --git a/torch/library.py b/torch/library.py
index c13ae5bf8b009..377c464e76a16 100644
--- a/torch/library.py
+++ b/torch/library.py
@@ -4,7 +4,6 @@
 import inspect
 import re
 import sys
-import traceback
 import weakref
 from collections.abc import Callable, Sequence
 from typing import Any, overload, TYPE_CHECKING, TypeVar, Union
@@ -88,16 +87,15 @@ def __init__(self, ns, kind, dispatch_key=""):
         from torch.fx.operator_schemas import _SCHEMA_TO_SIGNATURE_CACHE
 
         if kind not in ("IMPL", "DEF", "FRAGMENT"):
-            raise ValueError("Unsupported kind: ", kind)
+            raise ValueError(f"Unsupported kind: {kind}")
 
         if ns in _reserved_namespaces and (kind == "DEF" or kind == "FRAGMENT"):
             raise ValueError(
-                ns,
-                " is a reserved namespace. Please try creating a library with another name.",
+                f"{ns} is a reserved namespace. Please try creating a library with another name."
             )
 
-        frame = traceback.extract_stack(limit=2)[0]
-        filename, lineno = frame.filename, frame.lineno
+        f = sys._getframe(1)
+        filename, lineno = f.f_code.co_filename, f.f_lineno
         self.m: Any | None = torch._C._dispatch_library(
             kind, ns, dispatch_key, filename, lineno
         )
diff --git a/torch/masked/_ops.py b/torch/masked/_ops.py
index 99a9685964666..509cf6a056e23 100644
--- a/torch/masked/_ops.py
+++ b/torch/masked/_ops.py
@@ -1,7 +1,7 @@
 # mypy: allow-untyped-defs
 import warnings
 from collections.abc import Callable
-from typing import Any, Optional, TYPE_CHECKING, TypeAlias, TypeVar
+from typing import Any, TYPE_CHECKING, TypeAlias, TypeVar
 from typing_extensions import ParamSpec
 
 import torch
@@ -20,7 +20,7 @@
 else:
     # The JIT doesn't understand Union, nor torch.dtype here
     DType = int
-    DimOrDims = Optional[tuple[int, ...]]
+    DimOrDims = tuple[int, ...] | None
 
 
 __all__: list[str] = []
diff --git a/torch/masked/maskedtensor/_ops_refs.py b/torch/masked/maskedtensor/_ops_refs.py
index d9add0a1dfbae..fc6ee102a66b5 100644
--- a/torch/masked/maskedtensor/_ops_refs.py
+++ b/torch/masked/maskedtensor/_ops_refs.py
@@ -93,7 +93,7 @@ def backward(ctx, grad_output):
             return grad_output.to_sparse_csr()
         elif layout == torch.strided:
             return grad_output.to_dense()
-        raise ValueError("to_dense: Unsupported input layout: ", layout)
+        raise ValueError(f"to_dense: Unsupported input layout: {layout}")
 
 
 class _MaskedToSparse(torch.autograd.Function):
diff --git a/torch/mtia/__init__.py b/torch/mtia/__init__.py
index 10769f87718be..23fb8e4f79d7f 100644
--- a/torch/mtia/__init__.py
+++ b/torch/mtia/__init__.py
@@ -168,13 +168,13 @@ def synchronize(device: Device = None) -> None:
 
 def device_count() -> int:
     r"""Return the number of MTIA devices available."""
-    # TODO: Update _accelerator_hooks_device_count to abstract a MTIA device count API
     return torch._C._mtia_getDeviceCount()
 
 
 def current_device() -> int:
     r"""Return the index of a currently selected device."""
-    return torch._C._accelerator_hooks_get_current_device()
+    # pyrefly: ignore [missing-attribute]
+    return torch._C._mtia_getDevice()
 
 
 def current_stream(device: Device = None) -> Stream:
@@ -186,7 +186,8 @@ def current_stream(device: Device = None) -> Stream:
             by :func:`~torch.mtia.current_device`, if :attr:`device` is ``None``
             (default).
     """
-    return torch._C._mtia_getCurrentStream(_get_device_index(device, optional=True))
+    device_idx = _get_device_index(device, optional=True)
+    return torch._C._mtia_getCurrentStream(device_idx)
 
 
 def default_stream(device: Device = None) -> Stream:
@@ -277,7 +278,8 @@ def set_device(device: Device) -> None:
     """
     device = _get_device_index(device)
     if device >= 0:
-        torch._C._accelerator_hooks_set_current_device(device)
+        # pyrefly: ignore [missing-attribute]
+        torch._C._mtia_setDevice(device)
 
 
 def get_device_properties(device: Device = None) -> dict[str, Any]:
@@ -304,10 +306,12 @@ def __init__(self, device: Any):
         self.prev_idx = -1
 
     def __enter__(self):
-        self.prev_idx = torch._C._accelerator_hooks_maybe_exchange_device(self.idx)
+        # pyrefly: ignore [missing-attribute]
+        self.prev_idx = torch._C._mtia_maybeExchangeDevice(self.idx)
 
     def __exit__(self, type: Any, value: Any, traceback: Any):
-        self.idx = torch._C._accelerator_hooks_maybe_exchange_device(self.prev_idx)
+        # pyrefly: ignore [missing-attribute]
+        self.idx = torch._C._mtia_maybeExchangeDevice(self.prev_idx)
         return False
 
 
diff --git a/torch/mtia/_utils.py b/torch/mtia/_utils.py
index 090e26f321232..16710c244c61c 100644
--- a/torch/mtia/_utils.py
+++ b/torch/mtia/_utils.py
@@ -2,9 +2,6 @@
 
 import torch
 
-# The _get_device_index has been moved to torch.utils._get_device_index
-from torch._utils import _get_device_index as _torch_get_device_index
-
 
 def _get_device_index(
     device: Any, optional: bool = False, allow_cpu: bool = False
@@ -22,17 +19,29 @@ def _get_device_index(
     If :attr:`device` is ``None``, this will return the current default MTIA
     device if :attr:`optional` is ``True``.
     """
+
+    if device is None and optional:
+        # If device is None (frequent), then we can can short-circuit the logic
+        return torch._C._mtia_getDevice()
     if isinstance(device, int):
         return device
-    if isinstance(device, str):
-        device = torch.device(device)
-    if isinstance(device, torch.device):
-        if allow_cpu:
-            if device.type not in ["mtia", "cpu"]:
-                raise ValueError(f"Expected a mtia or cpu device, but got: {device}")
-        elif device.type != "mtia":
-            raise ValueError(f"Expected a mtia device, but got: {device}")
     if not torch.jit.is_scripting():
         if isinstance(device, torch.mtia.device):
             return device.idx
-    return _torch_get_device_index(device, optional, allow_cpu)
+    if isinstance(device, str):
+        device = torch.device(device)
+    device_idx: int | None = None
+    if isinstance(device, torch.device):
+        if not allow_cpu and device.type == "cpu":
+            raise ValueError(f"Expected a non cpu device, but got: {device}")
+        if device.type not in ["mtia", "cpu"]:
+            raise ValueError(f"Expected a mtia or cpu device, but got: {device}")
+        device_idx = -1 if device.type == "cpu" else device.index
+    if device_idx is None:
+        if optional:
+            device_idx = torch._C._mtia_getDevice()
+        else:
+            raise ValueError(
+                f"Expected a torch.device with a specified index or an integer, but got: {device}"
+            )
+    return device_idx
diff --git a/torch/nativert/executor/ConstantFolder.cpp b/torch/nativert/executor/ConstantFolder.cpp
index 34946aa09161a..0cad2841e2391 100644
--- a/torch/nativert/executor/ConstantFolder.cpp
+++ b/torch/nativert/executor/ConstantFolder.cpp
@@ -102,7 +102,7 @@ void ConstantFolder::unlinkConstants(
         if (user == output) {
           continue;
         }
-        if (foldable.find(user) == foldable.end()) {
+        if (!foldable.contains(user)) {
           constFoldableCandidates.push(user);
         }
       }
@@ -115,7 +115,7 @@ void ConstantFolder::unlinkConstants(
         // we only store folded values if there is a non-foldable user
         if (const auto& users = value->users();
             std::any_of(users.begin(), users.end(), [&](const auto* u) {
-              return foldable.find(u) == foldable.end();
+              return !foldable.contains(u);
             })) {
           foldedOutputValueIds_.insert(value->id());
         }
@@ -123,19 +123,16 @@ void ConstantFolder::unlinkConstants(
     }
   }
 
-  for (const auto& f : foldables_) {
-    VLOG(1) << "Const-folded node: " << *f.node;
+  if (VLOG_IS_ON(1)) {
+    for (const auto& f : foldables_) {
+      VLOG(1) << "Const-folded node: " << *f.node;
+    }
+    VLOG(1) << "Const-folded " << foldables_.size() << " nodes";
   }
-  LOG(INFO) << "Const-folded " << foldables_.size() << " nodes";
 
   // remove moved (i.e., associated w/ const-folded nodes) kernels
   // from the input kernel vector
-  kernels.erase(
-      std::remove_if(
-          kernels.begin(),
-          kernels.end(),
-          [](const auto& k) { return k == nullptr; }),
-      kernels.end());
+  std::erase(kernels, nullptr);
 
   graph_.renumberValues();
   graph_.finalize();
diff --git a/torch/nativert/executor/ExecutionFrame.h b/torch/nativert/executor/ExecutionFrame.h
index 1f256bb6b534b..e5d46449553dc 100644
--- a/torch/nativert/executor/ExecutionFrame.h
+++ b/torch/nativert/executor/ExecutionFrame.h
@@ -111,6 +111,18 @@ class ExecutionFrame {
     }
   }
 
+  // Clears all non-persistent, non-managed intermediate values.
+  // This is used when returning frames to the pool to prevent stale
+  // tensor data from persisting across frame reuses (e.g. during
+  // rebatching with varying batch sizes).
+  void clearNonPersistentValues() {
+    for (size_t i = 0; i < allValues_.size(); ++i) {
+      if (!persistent_[i] && !isManagedValue(static_cast<ValueId>(i))) {
+        allValues_[i] = c10::IValue();
+      }
+    }
+  }
+
   void destroyBorrowedIValues() {
     for (const auto& id : borrowedValueIds_) {
       c10::MaybeOwnedTraits<c10::IValue>::destroyBorrow(getIValue(id));
diff --git a/torch/nativert/executor/Executor.cpp b/torch/nativert/executor/Executor.cpp
index 50ad86daf32ac..c1c0c77bb6e6f 100644
--- a/torch/nativert/executor/Executor.cpp
+++ b/torch/nativert/executor/Executor.cpp
@@ -283,6 +283,13 @@ void Executor::returnExecutorFrameToPool(
 
   try {
     frame->destroyBorrowedIValues();
+    // When tryFreeUnmanagedValuesAfterUse is false, intermediate values are
+    // not freed during execution. Clear them here to prevent stale tensor
+    // data from persisting across frame reuses (e.g. during rebatching with
+    // varying batch sizes, stale tensors can cause out-of-bounds accesses).
+    if (!executorConfig_.tryFreeUnmanagedValuesAfterUse) {
+      frame->clearNonPersistentValues();
+    }
     // Always return to active execution frame pool, indicating that frame was
     // used in the previous time interval
     TORCH_CHECK(
diff --git a/torch/nativert/executor/GraphExecutorBase.cpp b/torch/nativert/executor/GraphExecutorBase.cpp
index a623d5873ea56..5cec883376e1f 100644
--- a/torch/nativert/executor/GraphExecutorBase.cpp
+++ b/torch/nativert/executor/GraphExecutorBase.cpp
@@ -79,6 +79,21 @@ ProfileMetrics GraphExecutorBase::benchmarkIndividualNodes(
     }
   }
 
+  // Capture input element counts per node (frame is populated after warmup).
+  results.inputElementsPerNode.resize(numNodes, 0);
+  for (const auto i : c10::irange(numNodes)) {
+    int64_t total_elements = 0;
+    for (const auto& input : nodeKernels_[i]->node()->inputs()) {
+      if (input.value && input.value->type().kind() == Type::Kind::Tensor) {
+        const auto& iv = executionFrame.getIValue(input.value->id());
+        if (iv.isTensor()) {
+          total_elements += iv.toTensor().numel();
+        }
+      }
+    }
+    results.inputElementsPerNode[i] = total_elements;
+  }
+
   // Execute kernels
   caffe2::Timer timer;
   executionFrame.withManagedMemory([&](auto) {
diff --git a/torch/nativert/executor/GraphExecutorBase.h b/torch/nativert/executor/GraphExecutorBase.h
index dfe020ebae29e..d78d28253e201 100644
--- a/torch/nativert/executor/GraphExecutorBase.h
+++ b/torch/nativert/executor/GraphExecutorBase.h
@@ -21,6 +21,7 @@ struct ProfileMetrics {
   std::unordered_map<std::string, int> instancesPerNodeType;
   std::unordered_set<std::string> staticDispatchNodes;
   std::unordered_set<std::string> primNodes;
+  std::vector<int> inputElementsPerNode;
   float totalTime{0};
   std::string name;
 };
diff --git a/torch/nativert/executor/Placement.cpp b/torch/nativert/executor/Placement.cpp
index 0432ecdc2a7c3..698d012029910 100644
--- a/torch/nativert/executor/Placement.cpp
+++ b/torch/nativert/executor/Placement.cpp
@@ -1,6 +1,7 @@
 #include <torch/nativert/executor/Placement.h>
 
 #include <fmt/ostream.h>
+#include <algorithm>
 #include <ostream>
 
 namespace torch::nativert {
diff --git a/torch/nativert/executor/PlacementUtils.cpp b/torch/nativert/executor/PlacementUtils.cpp
index 7dc6e761a0c94..97a2642323363 100644
--- a/torch/nativert/executor/PlacementUtils.cpp
+++ b/torch/nativert/executor/PlacementUtils.cpp
@@ -22,6 +22,5 @@ bool isSameDevice(const c10::Device& a, const c10::Device& b) {
     return b.is_mtia();
   }
   TORCH_CHECK(false, "isSameDevice: Unsupported device type ", a, " and ", b);
-  return false;
 }
 } // namespace torch::nativert
diff --git a/torch/nativert/executor/triton/CudaTritonKernelManager.cpp b/torch/nativert/executor/triton/CudaTritonKernelManager.cpp
index 83c70ffc810ad..3a1b55e033507 100644
--- a/torch/nativert/executor/triton/CudaTritonKernelManager.cpp
+++ b/torch/nativert/executor/triton/CudaTritonKernelManager.cpp
@@ -14,12 +14,12 @@ const at::cuda::NVRTC& get_nvrtc() {
 }
 } // namespace
 
-#define CU_LOG_ERROR(fn, result, ...)                   \
-  {                                                     \
-    LOG(ERROR) << #fn << " returned error: " << result; \
-    const char* errMsg = nullptr;                       \
-    get_nvrtc().cuGetErrorString(result, &errMsg);      \
-    LOG(ERROR) << "cuGetErrorString: " << errMsg;       \
+#define CU_LOG_ERROR(fn, result, ...)                    \
+  {                                                      \
+    LOG(ERROR) << #fn << " returned error: " << result;  \
+    const char* errMsg = nullptr;                        \
+    (void)get_nvrtc().cuGetErrorString(result, &errMsg); \
+    LOG(ERROR) << "cuGetErrorString: " << errMsg;        \
   }
 
 namespace torch::nativert {
diff --git a/torch/nativert/graph/Graph.cpp b/torch/nativert/graph/Graph.cpp
index b2be28da1f4d7..3531f12c6a685 100644
--- a/torch/nativert/graph/Graph.cpp
+++ b/torch/nativert/graph/Graph.cpp
@@ -194,6 +194,9 @@ std::ostream& operator<<(std::ostream& out, const Type& ty) {
             case Type::Kind::TensorList:
               out << "TensorList";
               break;
+            case Type::Kind::NestedTensorList:
+              out << "NestedTensorList";
+              break;
             case Type::Kind::OptionalTensorList:
               out << "OptionalTensorList";
               break;
@@ -417,6 +420,10 @@ Node* Graph::createListPack(std::vector<Value*> inputs, const Type& inputType) {
     node->addOutput(name, Type::Kind::TensorList);
   } else if (inputType == Type::Kind::SymInt) {
     node->addOutput(name, Type::Kind::SymIntList);
+  } else if (inputType == Type::Kind::TensorList) {
+    // For nested tensor lists (List[List[Tensor]]), the inner lists are
+    // TensorList type. We output a NestedTensorList type.
+    node->addOutput(name, Type::Kind::NestedTensorList);
   }
 
   return node;
@@ -546,6 +553,24 @@ bool Graph::cleanupDeadNodes() {
 
   const bool mutated = !toRemove.empty();
 
+  if (mutated && VLOG_IS_ON(1)) {
+    c10::FastMap<std::string_view, int> removedByTarget;
+    for (const auto* n : toRemove) {
+      removedByTarget[n->target()]++;
+    }
+    VLOG(1) << "cleanupDeadNodes: removing " << toRemove.size()
+            << " dead nodes. Breakdown by op:";
+    // Sort by count descending for readability
+    std::vector<std::pair<std::string_view, int>> sorted(
+        removedByTarget.begin(), removedByTarget.end());
+    std::sort(sorted.begin(), sorted.end(), [](const auto& a, const auto& b) {
+      return a.second > b.second;
+    });
+    for (const auto& [target, count] : sorted) {
+      VLOG(1) << "  " << target << ": " << count;
+    }
+  }
+
   // Remove nodes in reverse order to handle input/output dependencies
   for (auto it = toRemove.rbegin(); it != toRemove.rend(); ++it) {
     removeNode(*it);
@@ -605,7 +630,23 @@ void Graph::finalize() {
       userOutputs_.emplace_back(getValue(*outputName));
     } else {
       if (constantIndex < constantOutputs_.size()) {
-        userOutputs_.emplace_back(std::move(constantOutputs_[constantIndex]));
+        // Copy the constant rather than moving it, because finalize() may be
+        // called multiple times (e.g. after constant folding). Moving would
+        // leave constantOutputs_ entries in a moved-from state, causing
+        // subsequent calls to produce empty strings/vectors.
+        // Constant is non-copyable due to unique_ptr<Graph>, so we use
+        // std::visit to copy each alternative individually.
+        userOutputs_.emplace_back(std::visit(
+            [](const auto& val) -> Constant {
+              using T = std::decay_t<decltype(val)>;
+              if constexpr (is_same_v<T, std::unique_ptr<Graph>>) {
+                TORCH_CHECK(false, "Graph constant outputs cannot be copied");
+                return Constant(None{});
+              } else {
+                return Constant(val);
+              }
+            },
+            constantOutputs_[constantIndex]));
         constantIndex++;
       } else {
         TORCH_CHECK(false, "No more constant outputs available");
@@ -923,10 +964,7 @@ void Value::addUser(Node* node) {
 }
 
 void Value::eraseUser(Node* node) {
-  users_.erase(
-      std::remove_if(
-          users_.begin(), users_.end(), [&](Node* el) { return el == node; }),
-      users_.end());
+  std::erase(users_, node);
 }
 
 std::vector<const Value*> Value::getListElements() const {
@@ -1042,6 +1080,15 @@ std::ostream& operator<<(std::ostream& out, const Constant& constant) {
             out << fmt::format("{}", fmt::streamed(inner_list));
           }
           out << ']';
+        } else if constexpr (is_same_v<T, vector<vector<double>>>) {
+          out << '[';
+          for (const auto& [idx, inner_list] : c10::enumerate(arg)) {
+            if (idx > 0) {
+              out << ", ";
+            }
+            out << fmt::format("{}", fmt::streamed(inner_list));
+          }
+          out << ']';
         } else if constexpr (is_same_v<T, unique_ptr<Graph>>) {
           out << fmt::format("<subgraph>");
           VLOG(0) << "Subgraph pretty print is not implemented";
diff --git a/torch/nativert/graph/Graph.h b/torch/nativert/graph/Graph.h
index 7ac14d3890350..6d8c449d15c14 100644
--- a/torch/nativert/graph/Graph.h
+++ b/torch/nativert/graph/Graph.h
@@ -28,6 +28,7 @@ class Type {
     None,
     Tensor,
     TensorList,
+    NestedTensorList,
     OptionalTensorList,
     SymInt,
     SymIntList,
@@ -99,6 +100,7 @@ using Constant = std::variant<
     std::vector<bool>,
     std::vector<std::string>,
     std::vector<std::vector<int64_t>>,
+    std::vector<std::vector<double>>,
     std::unique_ptr<Graph>>;
 
 c10::IValue constantToIValue(const Constant& constant);
diff --git a/torch/nativert/graph/GraphSignature.cpp b/torch/nativert/graph/GraphSignature.cpp
index 855dfff07fd77..84b4c10192564 100644
--- a/torch/nativert/graph/GraphSignature.cpp
+++ b/torch/nativert/graph/GraphSignature.cpp
@@ -17,6 +17,8 @@ bool isSymbolicOutput(torch::_export::Argument::Tag t) {
   switch (t) {
     case torch::_export::Argument::Tag::AS_TENSOR:
     case torch::_export::Argument::Tag::AS_TENSORS:
+    case torch::_export::Argument::Tag::AS_NESTED_TENSORS:
+    case torch::_export::Argument::Tag::AS_OPTIONAL_TENSOR:
     case torch::_export::Argument::Tag::AS_OPTIONAL_TENSORS:
     case torch::_export::Argument::Tag::AS_SYM_BOOL:
     case torch::_export::Argument::Tag::AS_SYM_BOOLS:
@@ -196,6 +198,77 @@ GraphSignature::GraphSignature(const torch::_export::GraphSignature& storage) {
             userInputArg.tag() ==
             torch::_export::Argument::Tag::AS_CUSTOM_OBJ) {
           userInputs_.emplace_back(userInputArg.get_as_custom_obj().get_name());
+        } else if (
+            userInputArg.tag() == torch::_export::Argument::Tag::AS_TENSORS) {
+          // Handle list of tensors
+          for (const auto& tensor : userInputArg.get_as_tensors()) {
+            userInputs_.emplace_back(tensor.get_name());
+          }
+        } else if (
+            userInputArg.tag() ==
+            torch::_export::Argument::Tag::AS_OPTIONAL_TENSORS) {
+          // Handle list of optional tensors
+          for (const auto& optTensor : userInputArg.get_as_optional_tensors()) {
+            if (optTensor.tag() ==
+                torch::_export::OptionalTensorArgument::Tag::AS_TENSOR) {
+              userInputs_.emplace_back(optTensor.get_as_tensor().get_name());
+            }
+            // Skip None tensors
+          }
+        } else if (
+            userInputArg.tag() ==
+            torch::_export::Argument::Tag::AS_OPTIONAL_TENSOR) {
+          // Handle single optional tensor
+          const auto& optTensor = userInputArg.get_as_optional_tensor();
+          if (optTensor.tag() ==
+              torch::_export::OptionalTensorArgument::Tag::AS_TENSOR) {
+            userInputs_.emplace_back(optTensor.get_as_tensor().get_name());
+          }
+          // Skip if None
+        } else if (
+            userInputArg.tag() == torch::_export::Argument::Tag::AS_SYM_INT) {
+          // Handle symbolic int input
+          const auto& symInt = userInputArg.get_as_sym_int();
+          if (symInt.tag() == torch::_export::SymIntArgument::Tag::AS_NAME) {
+            userInputs_.emplace_back(symInt.get_as_name());
+          }
+          // Skip AS_INT (constant) symints
+        } else if (
+            userInputArg.tag() == torch::_export::Argument::Tag::AS_SYM_INTS) {
+          // Handle list of symbolic ints
+          for (const auto& symInt : userInputArg.get_as_sym_ints()) {
+            if (symInt.tag() == torch::_export::SymIntArgument::Tag::AS_NAME) {
+              userInputs_.emplace_back(symInt.get_as_name());
+            }
+            // Skip AS_INT (constant) symints
+          }
+        } else if (
+            userInputArg.tag() == torch::_export::Argument::Tag::AS_NONE ||
+            userInputArg.tag() == torch::_export::Argument::Tag::AS_INT ||
+            userInputArg.tag() == torch::_export::Argument::Tag::AS_INTS ||
+            userInputArg.tag() == torch::_export::Argument::Tag::AS_FLOAT ||
+            userInputArg.tag() == torch::_export::Argument::Tag::AS_FLOATS ||
+            userInputArg.tag() == torch::_export::Argument::Tag::AS_BOOL ||
+            userInputArg.tag() == torch::_export::Argument::Tag::AS_BOOLS ||
+            userInputArg.tag() == torch::_export::Argument::Tag::AS_STRING ||
+            userInputArg.tag() == torch::_export::Argument::Tag::AS_STRINGS ||
+            userInputArg.tag() == torch::_export::Argument::Tag::AS_SYM_BOOL ||
+            userInputArg.tag() == torch::_export::Argument::Tag::AS_SYM_BOOLS ||
+            userInputArg.tag() == torch::_export::Argument::Tag::AS_SYM_FLOAT ||
+            userInputArg.tag() ==
+                torch::_export::Argument::Tag::AS_SYM_FLOATS ||
+            userInputArg.tag() == torch::_export::Argument::Tag::AS_INT_LISTS ||
+            userInputArg.tag() ==
+                torch::_export::Argument::Tag::AS_FLOAT_LISTS ||
+            userInputArg.tag() ==
+                torch::_export::Argument::Tag::AS_SCALAR_TYPE ||
+            userInputArg.tag() ==
+                torch::_export::Argument::Tag::AS_MEMORY_FORMAT ||
+            userInputArg.tag() == torch::_export::Argument::Tag::AS_LAYOUT ||
+            userInputArg.tag() == torch::_export::Argument::Tag::AS_DEVICE ||
+            userInputArg.tag() == torch::_export::Argument::Tag::AS_COMPLEX) {
+          // Non-tensor inputs are constant values, skip them for now
+          // These don't map to named graph inputs in the same way tensors do
         } else {
           // TODO: handle other types
           TORCH_CHECK(false, "Non tensor inputs not implemented yet.");
@@ -256,21 +329,96 @@ GraphSignature::GraphSignature(const torch::_export::GraphSignature& storage) {
       case torch::_export::OutputSpec::Tag::LOSS_OUTPUT:
         lossOutput_ = outputSpec.get_loss_output().get_arg().get_name();
         break;
-      case torch::_export::OutputSpec::Tag::USER_OUTPUT:
-        if (isSymbolicOutput(outputSpec.get_user_output().get_arg().tag())) {
-          switch (outputSpec.get_user_output().get_arg().tag()) {
+      case torch::_export::OutputSpec::Tag::USER_OUTPUT: {
+        const auto& userOutputArg = outputSpec.get_user_output().get_arg();
+        if (isSymbolicOutput(userOutputArg.tag())) {
+          switch (userOutputArg.tag()) {
             case torch::_export::Argument::Tag::AS_TENSOR: {
-              userOutputs_.emplace_back(outputSpec.get_user_output()
-                                            .get_arg()
-                                            .get_as_tensor()
-                                            .get_name());
+              userOutputs_.emplace_back(
+                  userOutputArg.get_as_tensor().get_name());
+              break;
+            }
+            case torch::_export::Argument::Tag::AS_TENSORS: {
+              // Handle list of tensors - each tensor is a separate output
+              for (const auto& tensor : userOutputArg.get_as_tensors()) {
+                userOutputs_.emplace_back(tensor.get_name());
+              }
+              break;
+            }
+            case torch::_export::Argument::Tag::AS_OPTIONAL_TENSORS: {
+              // Handle list of optional tensors
+              for (const auto& optTensor :
+                   userOutputArg.get_as_optional_tensors()) {
+                if (optTensor.tag() ==
+                    torch::_export::OptionalTensorArgument::Tag::AS_TENSOR) {
+                  userOutputs_.emplace_back(
+                      optTensor.get_as_tensor().get_name());
+                } else {
+                  // None tensor - no name
+                  userOutputs_.emplace_back(std::nullopt);
+                }
+              }
+              break;
+            }
+            case torch::_export::Argument::Tag::AS_OPTIONAL_TENSOR: {
+              // Handle single optional tensor
+              const auto& optTensor = userOutputArg.get_as_optional_tensor();
+              if (optTensor.tag() ==
+                  torch::_export::OptionalTensorArgument::Tag::AS_TENSOR) {
+                userOutputs_.emplace_back(optTensor.get_as_tensor().get_name());
+              } else {
+                // None tensor - no name
+                userOutputs_.emplace_back(std::nullopt);
+              }
+              break;
+            }
+            case torch::_export::Argument::Tag::AS_CUSTOM_OBJ: {
+              userOutputs_.emplace_back(
+                  userOutputArg.get_as_custom_obj().get_name());
               break;
             }
             case torch::_export::Argument::Tag::AS_SYM_INT: {
-              userOutputs_.emplace_back(outputSpec.get_user_output()
-                                            .get_arg()
-                                            .get_as_sym_int()
-                                            .get_as_name());
+              userOutputs_.emplace_back(
+                  userOutputArg.get_as_sym_int().get_as_name());
+              break;
+            }
+            case torch::_export::Argument::Tag::AS_SYM_INTS: {
+              for (const auto& symInt : userOutputArg.get_as_sym_ints()) {
+                if (symInt.tag() ==
+                    torch::_export::SymIntArgument::Tag::AS_NAME) {
+                  userOutputs_.emplace_back(symInt.get_as_name());
+                }
+                // Skip AS_INT (constant) symints
+              }
+              break;
+            }
+            case torch::_export::Argument::Tag::AS_SYM_BOOL: {
+              userOutputs_.emplace_back(
+                  userOutputArg.get_as_sym_bool().get_as_name());
+              break;
+            }
+            case torch::_export::Argument::Tag::AS_SYM_BOOLS: {
+              for (const auto& symBool : userOutputArg.get_as_sym_bools()) {
+                if (symBool.tag() ==
+                    torch::_export::SymBoolArgument::Tag::AS_NAME) {
+                  userOutputs_.emplace_back(symBool.get_as_name());
+                }
+                // Skip AS_BOOL (constant) symbools
+              }
+              break;
+            }
+            case torch::_export::Argument::Tag::AS_SYM_FLOAT: {
+              // SymFloat doesn't have get_as_name in all versions
+              // For now, treat as unnamed symbolic output
+              userOutputs_.emplace_back(std::nullopt);
+              break;
+            }
+            case torch::_export::Argument::Tag::AS_SYM_FLOATS: {
+              // SymFloats - treat each as unnamed for now
+              for (size_t i = 0; i < userOutputArg.get_as_sym_floats().size();
+                   ++i) {
+                userOutputs_.emplace_back(std::nullopt);
+              }
               break;
             }
             default: {
@@ -283,6 +431,7 @@ GraphSignature::GraphSignature(const torch::_export::GraphSignature& storage) {
           userOutputs_.emplace_back(std::nullopt);
         }
         break;
+      }
       case torch::_export::OutputSpec::Tag::BUFFER_MUTATION:
         buffersToMutate_.emplace(
             outputSpec.get_buffer_mutation().get_arg().get_name(),
diff --git a/torch/nativert/graph/Serialization.cpp b/torch/nativert/graph/Serialization.cpp
index 8d846de6601af..ada34f15026dd 100644
--- a/torch/nativert/graph/Serialization.cpp
+++ b/torch/nativert/graph/Serialization.cpp
@@ -28,6 +28,25 @@ Value* symbolicToValue(
           graph.createListPack(std::move(listValue), Type::Kind::Tensor);
       return graph.insertBefore(listPack, insertBefore)->outputs()[0];
     }
+    case torch::_export::Argument::Tag::AS_NESTED_TENSORS: {
+      // Handle nested tensor lists (List[List[Tensor]])
+      // Create inner ListPack nodes for each inner list, then pack them
+      // together
+      std::vector<Value*> outerListValues;
+      for (const auto& innerList : arg.get_as_nested_tensors()) {
+        std::vector<Value*> innerListValues;
+        for (const auto& tensorArg : innerList) {
+          innerListValues.push_back(graph.getValue(tensorArg.get_name()));
+        }
+        auto innerListPack = graph.createListPack(
+            std::move(innerListValues), Type::Kind::Tensor);
+        outerListValues.push_back(
+            graph.insertBefore(innerListPack, insertBefore)->outputs()[0]);
+      }
+      auto outerListPack = graph.createListPack(
+          std::move(outerListValues), Type::Kind::TensorList);
+      return graph.insertBefore(outerListPack, insertBefore)->outputs()[0];
+    }
     case torch::_export::Argument::Tag::AS_OPTIONAL_TENSORS: {
       // Need to insert a list pack node
       std::vector<Value*> listValue;
@@ -200,6 +219,86 @@ std::unique_ptr<Graph> jsonToSubgraph(
           graph->addInput(name, Type::Kind::Tensor);
           break;
         }
+        case torch::_export::Argument::Tag::AS_TENSORS: {
+          // Handle list of tensors - each tensor becomes a separate input
+          for (const auto& tensor : input.get_as_tensors()) {
+            graph->addInput(tensor.get_name(), Type::Kind::Tensor);
+          }
+          break;
+        }
+        case torch::_export::Argument::Tag::AS_OPTIONAL_TENSOR: {
+          // Handle single optional tensor
+          const auto& optTensor = input.get_as_optional_tensor();
+          if (optTensor.tag() ==
+              torch::_export::OptionalTensorArgument::Tag::AS_TENSOR) {
+            graph->addInput(
+                optTensor.get_as_tensor().get_name(), Type::Kind::Tensor);
+          }
+          // Skip if None
+          break;
+        }
+        case torch::_export::Argument::Tag::AS_OPTIONAL_TENSORS: {
+          // Handle list of optional tensors
+          for (const auto& optTensor : input.get_as_optional_tensors()) {
+            if (optTensor.tag() ==
+                torch::_export::OptionalTensorArgument::Tag::AS_TENSOR) {
+              graph->addInput(
+                  optTensor.get_as_tensor().get_name(), Type::Kind::Tensor);
+            }
+            // Skip None tensors
+          }
+          break;
+        }
+        case torch::_export::Argument::Tag::AS_SYM_INT: {
+          const auto& symInt = input.get_as_sym_int();
+          if (symInt.tag() == torch::_export::SymIntArgument::Tag::AS_NAME) {
+            graph->addInput(symInt.get_as_name(), Type::Kind::SymInt);
+          }
+          // Skip constant symints
+          break;
+        }
+        case torch::_export::Argument::Tag::AS_SYM_INTS: {
+          for (const auto& symInt : input.get_as_sym_ints()) {
+            if (symInt.tag() == torch::_export::SymIntArgument::Tag::AS_NAME) {
+              graph->addInput(symInt.get_as_name(), Type::Kind::SymInt);
+            }
+            // Skip constant symints
+          }
+          break;
+        }
+        case torch::_export::Argument::Tag::AS_SYM_BOOL: {
+          const auto& symBool = input.get_as_sym_bool();
+          if (symBool.tag() == torch::_export::SymBoolArgument::Tag::AS_NAME) {
+            graph->addInput(symBool.get_as_name(), Type::Kind::SymBool);
+          }
+          // Skip constant symbools
+          break;
+        }
+        case torch::_export::Argument::Tag::AS_SYM_BOOLS: {
+          for (const auto& symBool : input.get_as_sym_bools()) {
+            if (symBool.tag() ==
+                torch::_export::SymBoolArgument::Tag::AS_NAME) {
+              graph->addInput(symBool.get_as_name(), Type::Kind::SymBool);
+            }
+            // Skip constant symbools
+          }
+          break;
+        }
+        case torch::_export::Argument::Tag::AS_SYM_FLOAT: {
+          // SymFloat inputs - add as SymFloat type
+          graph->addInput(
+              fmt::format("sym_float_{}", graph->numValues()),
+              Type::Kind::SymFloat);
+          break;
+        }
+        case torch::_export::Argument::Tag::AS_SYM_FLOATS: {
+          for (size_t i = 0; i < input.get_as_sym_floats().size(); ++i) {
+            graph->addInput(
+                fmt::format("sym_float_{}_{}", graph->numValues(), i),
+                Type::Kind::SymFloat);
+          }
+          break;
+        }
         case torch::_export::Argument::Tag::AS_CUSTOM_OBJ: {
           const auto& asCustomObj = input.get_as_custom_obj();
           const std::string& name = asCustomObj.get_name();
@@ -339,6 +438,39 @@ std::unique_ptr<Graph> jsonToSubgraph(
           graph->addOutput(outputValue);
           break;
         }
+        case torch::_export::Argument::Tag::AS_TENSORS: {
+          // Handle list of tensors - each tensor becomes a separate output
+          for (const auto& tensor : output.get_as_tensors()) {
+            Value* outputValue = graph->getValue(tensor.get_name());
+            graph->addOutput(outputValue);
+          }
+          break;
+        }
+        case torch::_export::Argument::Tag::AS_OPTIONAL_TENSOR: {
+          // Handle single optional tensor
+          const auto& optTensor = output.get_as_optional_tensor();
+          if (optTensor.tag() ==
+              torch::_export::OptionalTensorArgument::Tag::AS_TENSOR) {
+            Value* outputValue =
+                graph->getValue(optTensor.get_as_tensor().get_name());
+            graph->addOutput(outputValue);
+          }
+          // Skip None tensors
+          break;
+        }
+        case torch::_export::Argument::Tag::AS_OPTIONAL_TENSORS: {
+          // Handle list of optional tensors
+          for (const auto& optTensor : output.get_as_optional_tensors()) {
+            if (optTensor.tag() ==
+                torch::_export::OptionalTensorArgument::Tag::AS_TENSOR) {
+              Value* outputValue =
+                  graph->getValue(optTensor.get_as_tensor().get_name());
+              graph->addOutput(outputValue);
+            }
+            // Skip None tensors
+          }
+          break;
+        }
         case torch::_export::Argument::Tag::AS_SYM_INT: {
           const auto& asSymInt = output.get_as_sym_int();
           TORCH_CHECK(
@@ -348,6 +480,52 @@ std::unique_ptr<Graph> jsonToSubgraph(
           graph->addOutput(outputValue);
           break;
         }
+        case torch::_export::Argument::Tag::AS_SYM_INTS: {
+          for (const auto& symInt : output.get_as_sym_ints()) {
+            if (symInt.tag() == torch::_export::SymIntArgument::Tag::AS_NAME) {
+              Value* outputValue = graph->getValue(symInt.get_as_name());
+              graph->addOutput(outputValue);
+            }
+          }
+          break;
+        }
+        case torch::_export::Argument::Tag::AS_SYM_BOOL: {
+          const auto& symBool = output.get_as_sym_bool();
+          if (symBool.tag() == torch::_export::SymBoolArgument::Tag::AS_NAME) {
+            Value* outputValue = graph->getValue(symBool.get_as_name());
+            graph->addOutput(outputValue);
+          }
+          break;
+        }
+        case torch::_export::Argument::Tag::AS_SYM_BOOLS: {
+          for (const auto& symBool : output.get_as_sym_bools()) {
+            if (symBool.tag() ==
+                torch::_export::SymBoolArgument::Tag::AS_NAME) {
+              Value* outputValue = graph->getValue(symBool.get_as_name());
+              graph->addOutput(outputValue);
+            }
+          }
+          break;
+        }
+        case torch::_export::Argument::Tag::AS_SYM_FLOAT: {
+          const auto& symFloat = output.get_as_sym_float();
+          Value* outputValue = graph->getValue(symFloat.get_as_name());
+          graph->addOutput(outputValue);
+          break;
+        }
+        case torch::_export::Argument::Tag::AS_SYM_FLOATS: {
+          for (const auto& symFloat : output.get_as_sym_floats()) {
+            Value* outputValue = graph->getValue(symFloat.get_as_name());
+            graph->addOutput(outputValue);
+          }
+          break;
+        }
+        case torch::_export::Argument::Tag::AS_CUSTOM_OBJ: {
+          const auto& asCustomObj = output.get_as_custom_obj();
+          Value* outputValue = graph->getValue(asCustomObj.get_name());
+          graph->addOutput(outputValue);
+          break;
+        }
         default:
           TORCH_CHECK(
               false,
@@ -449,6 +627,7 @@ bool isSymbolic(const torch::_export::Argument& arg) {
   switch (arg.tag()) {
     case torch::_export::Argument::Tag::AS_TENSOR:
     case torch::_export::Argument::Tag::AS_TENSORS:
+    case torch::_export::Argument::Tag::AS_NESTED_TENSORS:
     case torch::_export::Argument::Tag::AS_OPTIONAL_TENSORS:
     case torch::_export::Argument::Tag::AS_SYM_INT:
     case torch::_export::Argument::Tag::AS_SYM_INTS:
@@ -552,6 +731,17 @@ Constant constantToValue(
       }
       return ret;
     }
+    case torch::_export::Argument::Tag::AS_FLOAT_LISTS: {
+      std::vector<std::vector<double>> ret;
+      for (const auto& inner_list : jsonArg.get_as_float_lists()) {
+        std::vector<double> inner_ret;
+        for (const auto& val : inner_list) {
+          inner_ret.push_back(val.get());
+        }
+        ret.push_back(inner_ret);
+      }
+      return ret;
+    }
     case torch::_export::Argument::Tag::AS_STRING_TO_ARGUMENT:
       return None();
     default:
diff --git a/torch/nativert/graph/passes/pass_manager/PassManager.cpp b/torch/nativert/graph/passes/pass_manager/PassManager.cpp
index 4dbb0012877d8..bed89e57b41c3 100644
--- a/torch/nativert/graph/passes/pass_manager/PassManager.cpp
+++ b/torch/nativert/graph/passes/pass_manager/PassManager.cpp
@@ -26,10 +26,28 @@ bool GraphPassManager::run(Graph* graph) {
 bool GraphPassManager::run_pass(Graph* graph, const GraphPassIdentifier& name) {
   const auto& pass = GraphPassRegistry::get().get_pass(name);
 
+  size_t nodesBefore = 0;
+  if (VLOG_IS_ON(1)) {
+    nodesBefore = graph->nodes().size();
+  }
+
   bool changed = pass_pre_run_hook(graph, pass);
   changed |= (pass.get())(graph);
   changed |= pass_post_run_hook(graph, pass);
 
+  if (VLOG_IS_ON(1)) {
+    size_t nodesAfter = graph->nodes().size();
+    if (changed) {
+      VLOG(1) << "Pass " << name << ": " << nodesBefore << " -> " << nodesAfter
+              << " nodes (delta: "
+              << static_cast<int64_t>(nodesAfter) -
+              static_cast<int64_t>(nodesBefore)
+              << ")";
+    } else {
+      VLOG(1) << "Pass " << name << ": no change (" << nodesAfter << " nodes)";
+    }
+  }
+
   return changed;
 }
 
diff --git a/torch/nativert/kernels/GeneratedStaticDispatchKernels.cpp b/torch/nativert/kernels/GeneratedStaticDispatchKernels.cpp
index f919639f48def..c5b9b3580f2c6 100644
--- a/torch/nativert/kernels/GeneratedStaticDispatchKernels.cpp
+++ b/torch/nativert/kernels/GeneratedStaticDispatchKernels.cpp
@@ -1574,6 +1574,19 @@ REGISTER_CPU_KERNEL(
       at::cpu::bitwise_and_out(out, self, other);
     })
 
+// __and__ is the Python dunder alias for bitwise_and
+REGISTER_CPU_KERNEL("torch.ops.aten.__and__.Tensor", aten___and___Tensor, {
+  const auto& self = KernelInput(0).toTensor();
+  const auto& other = KernelInput(1).toTensor();
+  if (KernelOutput(0).isNone()) {
+    KernelOutput(0) = at::cpu::bitwise_and(self, other);
+    return;
+  }
+  auto& out = KernelOutput(0).toTensor();
+  fastResizeToZero(out);
+  at::cpu::bitwise_and_out(out, self, other);
+})
+
 REGISTER_CPU_KERNEL(
     "torch.ops.aten.bitwise_or.Tensor",
     aten_bitwise_or_Tensor,
@@ -1589,6 +1602,32 @@ REGISTER_CPU_KERNEL(
       at::cpu::bitwise_or_out(out, self, other);
     })
 
+// __or__ is the Python dunder alias for bitwise_or
+REGISTER_CPU_KERNEL("torch.ops.aten.__or__.Tensor", aten___or___Tensor, {
+  const auto& self = KernelInput(0).toTensor();
+  const auto& other = KernelInput(1).toTensor();
+  if (KernelOutput(0).isNone()) {
+    KernelOutput(0) = at::cpu::bitwise_or(self, other);
+    return;
+  }
+  auto& out = KernelOutput(0).toTensor();
+  fastResizeToZero(out);
+  at::cpu::bitwise_or_out(out, self, other);
+})
+
+REGISTER_NATIVE_CPU_KERNEL(
+    "torch.ops.aten.scalar_tensor.default",
+    aten_scalar_tensor_default,
+    {
+      const auto s = KernelInput(0).toScalar();
+      const auto dtype = KernelInput(1).toOptional<at::ScalarType>();
+      const auto layout = KernelInput(2).toOptional<at::Layout>();
+      const auto device = KernelInput(3).toOptional<at::Device>();
+      const auto pin_memory = KernelInput(4).toOptional<bool>();
+      KernelOutput(0) =
+          at::native::scalar_tensor(s, dtype, layout, device, pin_memory);
+    })
+
 REGISTER_CPU_KERNEL(
     "torch.ops.aten.bitwise_xor.Tensor",
     aten_bitwise_xor_Tensor,
diff --git a/torch/nativert/kernels/PrimKernelRegistry.cpp b/torch/nativert/kernels/PrimKernelRegistry.cpp
index 21d61203548ca..030b97affc11f 100644
--- a/torch/nativert/kernels/PrimKernelRegistry.cpp
+++ b/torch/nativert/kernels/PrimKernelRegistry.cpp
@@ -21,6 +21,10 @@ class OpKernel_prim_listpack : public OpKernel {
       case Type::Kind::TensorList:
         type_ = c10::TensorType::get();
         break;
+      case Type::Kind::NestedTensorList:
+        // For List[List[Tensor]], each element is a List[Tensor]
+        type_ = c10::ListType::create(c10::TensorType::get());
+        break;
       case Type::Kind::SymIntList:
         type_ = c10::IntType::get();
         break;
diff --git a/torch/nativert/python/Bindings.cpp b/torch/nativert/python/Bindings.cpp
index 0cb861cd3567d..de86306adea0a 100644
--- a/torch/nativert/python/Bindings.cpp
+++ b/torch/nativert/python/Bindings.cpp
@@ -3,6 +3,7 @@
 #include <torch/csrc/jit/python/pybind_utils.h>
 #include <torch/csrc/utils/pybind.h>
 
+#ifdef FBCODE_CAFFE2
 #include <torch/nativert/ModelRunner.h>
 
 namespace py = pybind11;
@@ -27,12 +28,14 @@ void initModelRunnerPybind(py::module& m) {
              py::args pyargs,
              const py::kwargs& pykwargs) {
             std::vector<c10::IValue> args;
+            args.reserve(pyargs.size());
             for (const auto i : c10::irange(pyargs.size())) {
               auto ivalue =
                   torch::jit::toIValue(pyargs[i], c10::AnyType::get());
               args.push_back(std::move(ivalue));
             }
             std::unordered_map<std::string, c10::IValue> kwargs;
+            kwargs.reserve(pykwargs.size());
             for (const auto& [key, pyarg] : pykwargs) {
               auto ivalue = torch::jit::toIValue(pyarg, c10::AnyType::get());
               kwargs[py::str(key)] = std::move(ivalue);
@@ -46,12 +49,14 @@ void initModelRunnerPybind(py::module& m) {
              py::args pyargs,
              const py::kwargs& pykwargs) {
             std::vector<c10::IValue> args;
+            args.reserve(pyargs.size());
             for (const auto i : c10::irange(pyargs.size())) {
               auto ivalue =
                   torch::jit::toIValue(pyargs[i], c10::AnyType::get());
               args.push_back(std::move(ivalue));
             }
             std::unordered_map<std::string, c10::IValue> kwargs;
+            kwargs.reserve(pykwargs.size());
             for (const auto& [key, pyarg] : pykwargs) {
               auto ivalue = torch::jit::toIValue(pyarg, c10::AnyType::get());
               kwargs[py::str(key)] = std::move(ivalue);
@@ -63,6 +68,7 @@ void initModelRunnerPybind(py::module& m) {
           "run_with_flat_inputs_and_outputs",
           [](torch::nativert::ModelRunner& self, py::args pyargs) {
             std::vector<c10::IValue> args;
+            args.reserve(pyargs.size());
             for (const auto i : c10::irange(pyargs.size())) {
               auto ivalue =
                   torch::jit::toIValue(pyargs[i], c10::AnyType::get());
@@ -76,3 +82,22 @@ void initModelRunnerPybind(py::module& m) {
 }
 
 } // namespace torch::nativert
+
+#else // !FBCODE_CAFFE2
+
+namespace py = pybind11;
+
+namespace torch::nativert {
+
+class StubModelRunner {};
+
+// PyModelRunner is referenced from
+// https://github.com/pytorch/benchmark/blob/b8d35ba51a3149b7212888b4010ddee97f19947f/userbenchmark/dynamo/dynamobench/common.py#L45
+void initModelRunnerPybind(py::module& m) {
+  py::class_<StubModelRunner, std::shared_ptr<StubModelRunner>>(
+      m, "PyModelRunner");
+}
+
+} // namespace torch::nativert
+
+#endif // FBCODE_CAFFE2
diff --git a/torch/nn/attention/_fa3.py b/torch/nn/attention/_fa3.py
index 0fa1ed9dc51e2..3ea065134c434 100644
--- a/torch/nn/attention/_fa3.py
+++ b/torch/nn/attention/_fa3.py
@@ -98,6 +98,11 @@ def _fa3_register_kernels() -> Library:
     lib.impl(
         "_flash_attention_forward", _fa3_flash_attention_forward_impl_default, "CUDA"
     )
+    lib.impl(
+        "_flash_attention_forward_no_dropout_inplace",
+        _fa3_flash_attention_forward_no_dropout_inplace_impl,
+        "CUDA",
+    )
     lib.impl(
         "_scaled_dot_product_flash_attention",
         _fa3_scaled_dot_product_flash_attention_forward_impl_default,
@@ -260,6 +265,8 @@ def _fa3_run_forward(
     q_descale: torch.Tensor | None = None,
     k_descale: torch.Tensor | None = None,
     v_descale: torch.Tensor | None = None,
+    block_table: torch.Tensor | None = None,
+    num_splits: int | None = None,
 ) -> tuple[torch.Tensor, torch.Tensor]:
     """
     Run the FA3 forward pass by calling the C++ kernel directly.
@@ -280,6 +287,7 @@ def _fa3_run_forward(
     cu_seqlens_q = _maybe_contiguous(cu_seq_q)
     cu_seqlens_k = _maybe_contiguous(cu_seq_k)
     seqused_k = _maybe_contiguous(seqused_k)
+    block_table = _maybe_contiguous(block_table)
 
     out, softmax_lse, out_accum, softmax_lse_accum = _FA3_CUDA_FWD(
         q,
@@ -296,7 +304,7 @@ def _fa3_run_forward(
         seqused_k,  # seqused_k
         max_q,  # max_seqlen_q
         max_k,  # max_seqlen_k
-        None,  # page_table,
+        block_table,  # block_table,
         None,  # kv_batch_idx,
         None,  # leftpad_k,
         None,  # rotary_cos,
@@ -313,7 +321,8 @@ def _fa3_run_forward(
         0.0,  # softcap,
         True,  # rotary_interleaved,
         None,  # scheduler_metadata,
-        1 if torch.are_deterministic_algorithms_enabled() else 0,  # num_splits,
+        num_splits
+        or (1 if torch.are_deterministic_algorithms_enabled() else 0),  # num_splits,
         None,  # pack_gqa,
         torch._C._get_sm_carveout_experimental() or 0,  # sm_margin,
     )
@@ -400,6 +409,9 @@ def _fa3_flash_attention_forward_impl(
     seqused_k: torch.Tensor | None = None,
     alibi_slopes: torch.Tensor | None = None,
     out: torch.Tensor | None = None,
+    block_table: torch.Tensor | None = None,
+    compute_auxiliary: bool = True,
+    num_splits: int | None = None,
 ):
     error = _fa3_forward_support_error(
         query,
@@ -433,13 +445,68 @@ def _fa3_flash_attention_forward_impl(
         q_descale,
         k_descale,
         v_descale,
+        block_table,
+        num_splits,
     )
-    rng_state = torch.zeros((2,), dtype=torch.uint64, device=query.device)
-    philox_offset = torch.zeros((), dtype=torch.uint64, device=query.device)
-    debug_mask = torch.empty(0, dtype=query.dtype, device=query.device)
+    if compute_auxiliary:
+        rng_state = torch.zeros((2,), dtype=torch.uint64, device=query.device)
+        philox_offset = torch.zeros((), dtype=torch.uint64, device=query.device)
+        debug_mask = torch.empty(0, dtype=query.dtype, device=query.device)
+    else:
+        rng_state = None
+        philox_offset = None
+        debug_mask = None
     return out, lse, rng_state, philox_offset, debug_mask
 
 
+def _fa3_flash_attention_forward_no_dropout_inplace_impl(
+    out: torch.Tensor,
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    cum_seq_q: torch.Tensor | None,
+    cum_seq_k: torch.Tensor | None,
+    max_q: int,
+    max_k: int,
+    dropout_p: float,
+    is_causal: bool,
+    return_debug_mask: bool,
+    *,
+    scale: float | None = None,
+    window_size_left: int = -1,
+    window_size_right: int = -1,
+    seqused_k: torch.Tensor | None = None,
+    alibi_slopes: torch.Tensor | None = None,
+    block_table: torch.Tensor | None = None,
+    num_splits: int | None = None,
+):
+    _, lse, _, _, _ = _fa3_flash_attention_forward_impl(
+        query,
+        key,
+        value,
+        cum_seq_q,
+        cum_seq_k,
+        max_q,
+        max_k,
+        dropout_p,
+        is_causal,
+        return_debug_mask,
+        None,
+        None,
+        None,
+        scale=scale,
+        window_size_left=window_size_left,
+        window_size_right=window_size_right,
+        seqused_k=seqused_k,
+        alibi_slopes=alibi_slopes,
+        out=out,
+        block_table=block_table,
+        compute_auxiliary=False,
+        num_splits=num_splits,
+    )
+    return lse
+
+
 def _fa3_flash_attention_forward_impl_default(
     query: torch.Tensor,
     key: torch.Tensor,
@@ -457,7 +524,9 @@ def _fa3_flash_attention_forward_impl_default(
     window_size_right: int = -1,
     seqused_k: torch.Tensor | None = None,
     alibi_slopes: torch.Tensor | None = None,
+    block_table: torch.Tensor | None = None,
     out: torch.Tensor | None = None,
+    num_splits: int | None = None,
 ):
     return _fa3_flash_attention_forward_impl(
         query,
@@ -479,6 +548,8 @@ def _fa3_flash_attention_forward_impl_default(
         seqused_k=seqused_k,
         alibi_slopes=alibi_slopes,
         out=out,
+        block_table=block_table,
+        num_splits=num_splits,
     )
 
 
diff --git a/torch/nn/attention/_fa4.py b/torch/nn/attention/_fa4.py
index f0ea99a463532..c2786d9dceefe 100644
--- a/torch/nn/attention/_fa4.py
+++ b/torch/nn/attention/_fa4.py
@@ -67,6 +67,11 @@ def _fa4_import_module(module_path: str) -> ModuleType:
 def _fa4_register_kernels() -> Library:
     lib = Library("aten", "IMPL", "CUDA")  # noqa: TOR901
     lib.impl("_flash_attention_forward", _fa4_flash_attention_forward_impl, "CUDA")
+    lib.impl(
+        "_flash_attention_forward_no_dropout_inplace",
+        _fa4_flash_attention_forward_no_dropout_inplace_impl,
+        "CUDA",
+    )
     lib.impl("_flash_attention_backward", _fa4_flash_attention_backward_impl, "CUDA")
     lib.impl(
         "_scaled_dot_product_flash_attention",
@@ -116,6 +121,8 @@ def _fa4_forward_support_error(
     alibi_slopes: torch.Tensor | None,
     seqused_k: torch.Tensor | None,
     cum_seq_q: torch.Tensor | None,
+    block_table: torch.Tensor | None = None,
+    num_splits: int | None = None,
 ) -> str | None:
     if dropout_p != 0.0:
         return "dropout_p must be 0"
@@ -128,6 +135,11 @@ def _fa4_forward_support_error(
             return "seqused_k must be int32"
         if not seqused_k.is_cuda:
             return "seqused_k must be CUDA"
+    major = _get_device_major(query.device)
+    if block_table is not None and major != 10:
+        return f"paged KV (block_table) not supported on SM {major}0"
+    if num_splits is not None and num_splits > 1 and major != 10:
+        return f"SplitKV (num_splits > 1) not supported on SM {major}0"
     error = _fa4_common_support_error(
         query,
         (query, key, value),
@@ -149,13 +161,9 @@ def _fa4_backward_support_error(
     logsumexp: torch.Tensor,
     dropout_p: float,
     cum_seq_q: torch.Tensor | None,
-    window_size_left: int | None,
-    window_size_right: int | None,
 ) -> str | None:
     if dropout_p != 0.0:
         return "dropout_p must be 0"
-    if window_size_left is not None or window_size_right is not None:
-        return "windowed attention not supported"
     error = _fa4_common_support_error(
         query,
         (grad_out, query, key, value, out, logsumexp),
@@ -167,6 +175,11 @@ def _fa4_backward_support_error(
     return None
 
 
+def _aten_to_fa4_window_size(val: int | None) -> int | None:
+    """need to convert -1 to None for FA4"""
+    return None if val == -1 else val
+
+
 Ts = TypeVarTuple("Ts")
 
 
@@ -180,12 +193,16 @@ def _fa4_run_forward(
     value: torch.Tensor,
     cu_seq_q: torch.Tensor | None,
     cu_seq_k: torch.Tensor | None,
+    max_q: int | None,
+    max_k: int | None,
     scale: float | None,
     is_causal: bool,
     window_size_left: int | None,
     window_size_right: int | None,
     seqused_k: torch.Tensor | None,
     out: torch.Tensor | None = None,
+    block_table: torch.Tensor | None = None,
+    num_splits: int | None = None,
 ) -> tuple[torch.Tensor, torch.Tensor]:
     if _FA4_MODULE_PATH is None:
         raise RuntimeError("FA4 not registered")
@@ -194,15 +211,18 @@ def _fa4_run_forward(
     kwargs: dict[str, Any] = {
         "softmax_scale": scale,
         "causal": is_causal,
-        "window_size_left": window_size_left,
-        "window_size_right": window_size_right,
+        "window_size_left": _aten_to_fa4_window_size(window_size_left),
+        "window_size_right": _aten_to_fa4_window_size(window_size_right),
         "return_lse": True,
         "cu_seqlens_q": cu_seq_q,
         "cu_seqlens_k": cu_seq_k,
+        "max_seqlen_q": max_q,
+        "max_seqlen_k": max_k,
         "seqused_k": seqused_k.contiguous() if seqused_k is not None else None,
+        "page_table": block_table,
+        "num_splits": num_splits or 1,
+        "out": out,
     }
-    if out is not None:
-        kwargs["out"] = out
     out, lse = module._flash_attn_fwd(query, key, value, **kwargs)
     return out, lse.contiguous()
 
@@ -218,6 +238,8 @@ def _fa4_run_backward(
     cu_seq_k: torch.Tensor | None,
     scale: float | None,
     is_causal: bool,
+    window_size_left: int | None,
+    window_size_right: int | None,
     deterministic: bool = False,
 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
     if _FA4_MODULE_PATH is None:
@@ -232,6 +254,8 @@ def _fa4_run_backward(
         logsumexp.contiguous(),
         softmax_scale=scale,
         causal=is_causal,
+        window_size_left=_aten_to_fa4_window_size(window_size_left),
+        window_size_right=_aten_to_fa4_window_size(window_size_right),
         cu_seqlens_q=cu_seq_q,
         cu_seqlens_k=cu_seq_k,
         deterministic=deterministic,
@@ -257,6 +281,9 @@ def _fa4_flash_attention_forward_impl(
     seqused_k: torch.Tensor | None = None,
     alibi_slopes: torch.Tensor | None = None,
     out: torch.Tensor | None = None,
+    block_table: torch.Tensor | None = None,
+    compute_auxiliary: bool = True,
+    num_splits: int | None = None,
 ):
     error = _fa4_forward_support_error(
         query,
@@ -267,6 +294,8 @@ def _fa4_flash_attention_forward_impl(
         alibi_slopes,
         seqused_k,
         cum_seq_q,
+        block_table,
+        num_splits,
     )
     if error is not None:
         raise RuntimeError(f"FA4 flash_attention forward unsupported: {error}")
@@ -276,19 +305,73 @@ def _fa4_flash_attention_forward_impl(
         value,
         cum_seq_q,
         cum_seq_k,
+        max_q,
+        max_k,
         scale,
         is_causal,
         window_size_left,
         window_size_right,
         seqused_k,
         out,
+        block_table,
+        num_splits,
     )
-    rng_state = torch.zeros((2,), dtype=torch.uint64, device=query.device)
-    philox_offset = torch.zeros((), dtype=torch.uint64, device=query.device)
-    debug_mask = torch.empty(0, dtype=query.dtype, device=query.device)
+    if compute_auxiliary:
+        rng_state = torch.zeros((2,), dtype=torch.uint64, device=query.device)
+        philox_offset = torch.zeros((), dtype=torch.uint64, device=query.device)
+        debug_mask = torch.empty(0, dtype=query.dtype, device=query.device)
+    else:
+        rng_state = None
+        philox_offset = None
+        debug_mask = None
     return out, lse, rng_state, philox_offset, debug_mask
 
 
+def _fa4_flash_attention_forward_no_dropout_inplace_impl(
+    out: torch.Tensor,
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    cum_seq_q: torch.Tensor | None,
+    cum_seq_k: torch.Tensor | None,
+    max_q: int,
+    max_k: int,
+    dropout_p: float,
+    is_causal: bool,
+    return_debug_mask: bool,
+    *,
+    scale: float | None = None,
+    window_size_left: int | None = None,
+    window_size_right: int | None = None,
+    seqused_k: torch.Tensor | None = None,
+    alibi_slopes: torch.Tensor | None = None,
+    block_table: torch.Tensor | None = None,
+    num_splits: int | None = None,
+):
+    _, lse, _, _, _ = _fa4_flash_attention_forward_impl(
+        query,
+        key,
+        value,
+        cum_seq_q,
+        cum_seq_k,
+        max_q,
+        max_k,
+        dropout_p,
+        is_causal,
+        return_debug_mask,
+        scale=scale,
+        window_size_left=window_size_left,
+        window_size_right=window_size_right,
+        seqused_k=seqused_k,
+        alibi_slopes=alibi_slopes,
+        out=out,
+        block_table=block_table,
+        compute_auxiliary=False,
+        num_splits=num_splits,
+    )
+    return lse
+
+
 def _fa4_flash_attention_backward_impl(
     grad_out: torch.Tensor,
     query: torch.Tensor,
@@ -318,8 +401,6 @@ def _fa4_flash_attention_backward_impl(
         logsumexp,
         dropout_p,
         cum_seq_q,
-        window_size_left,
-        window_size_right,
     )
     if error is not None:
         raise RuntimeError(f"FA4 flash_attention backward unsupported: {error}")
@@ -335,6 +416,8 @@ def _fa4_flash_attention_backward_impl(
         cum_seq_k,
         scale,
         is_causal,
+        window_size_left,
+        window_size_right,
         deterministic,
     )
     return dq, dk, dv
@@ -428,8 +511,6 @@ def _fa4_scaled_dot_product_flash_attention_backward_impl(
         logsumexp,
         dropout_p,
         None,
-        None,
-        None,
     )
     if error is not None:
         raise RuntimeError(f"FA4 SDPA backward unsupported: {error}")
diff --git a/torch/nn/attention/experimental/_scaled_dot_product_attention_quantized.py b/torch/nn/attention/experimental/_scaled_dot_product_attention_quantized.py
index 93145d576c165..47f9b92670a6d 100644
--- a/torch/nn/attention/experimental/_scaled_dot_product_attention_quantized.py
+++ b/torch/nn/attention/experimental/_scaled_dot_product_attention_quantized.py
@@ -6,7 +6,6 @@
 
 import warnings
 from enum import IntEnum
-from typing import Optional
 
 import torch
 from torch import Tensor
@@ -27,7 +26,7 @@ class DescaleType(IntEnum):
 
 
 def _validate_descale(
-    descale: Optional[Tensor],
+    descale: Tensor | None,
     name: str,
     query: Tensor,
     key: Tensor,
@@ -92,10 +91,10 @@ def _scaled_dot_product_attention_quantized(
     key: Tensor,
     value: Tensor,
     is_causal: bool = False,
-    scale: Optional[float] = None,
-    q_descale: Optional[Tensor] = None,
-    k_descale: Optional[Tensor] = None,
-    v_descale: Optional[Tensor] = None,
+    scale: float | None = None,
+    q_descale: Tensor | None = None,
+    k_descale: Tensor | None = None,
+    v_descale: Tensor | None = None,
     q_descale_type: DescaleType = DescaleType.PER_HEAD,
     k_descale_type: DescaleType = DescaleType.PER_HEAD,
     v_descale_type: DescaleType = DescaleType.PER_HEAD,
diff --git a/torch/nn/attention/flex_attention.py b/torch/nn/attention/flex_attention.py
index 5c41d30a4cd7b..5cf2c42e2b93c 100644
--- a/torch/nn/attention/flex_attention.py
+++ b/torch/nn/attention/flex_attention.py
@@ -2,25 +2,38 @@
 # flake8: noqa: B950
 """This module implements the user facing API for flex_attention in PyTorch."""
 
+from __future__ import annotations
+
 import functools
 import inspect
 import itertools
 import math
 import operator
+import types
 import typing
 import warnings
 from collections.abc import Callable
 from enum import Enum
-from typing import Any, Literal, NamedTuple, TypeAlias
-from typing_extensions import NotRequired, TypedDict
+from typing import Any, cast, Literal, NamedTuple, overload, TypeAlias, TypeVar
+from typing_extensions import deprecated, Never, NotRequired, Self, TypedDict
 
 import torch
 from torch import Tensor
 from torch._higher_order_ops.flex_attention import flex_attention as flex_attention_hop
 from torch._higher_order_ops.utils import setup_compilation_env
-from torch._prims_common import DeviceLikeType
 from torch.nn.attention._utils import _validate_sdpa_input
-from torch.utils._pytree import GetAttrKey, tree_map_only
+from torch.utils._pytree import (
+    GetAttrKey,
+    tree_flatten,
+    tree_map_only,
+    tree_unflatten,
+    TreeSpec,
+)
+
+
+if typing.TYPE_CHECKING:
+    from torch._prims_common import DeviceLikeType
+    from torch.fx.node import BaseArgumentTypes
 
 
 # Private debug flag to disable internal compilation wrapping for debugging purposes.
@@ -70,6 +83,7 @@ def _warn_once(
 _score_mod_signature = Callable[[Tensor, Tensor, Tensor, Tensor, Tensor], Tensor]
 _mask_mod_signature = Callable[[Tensor, Tensor, Tensor, Tensor], Tensor]
 _Backend: TypeAlias = Literal["AUTO", "TRITON", "FLASH", "TRITON_DECODE"]
+_R = TypeVar("_R")
 
 
 class FlexKernelOptions(TypedDict, total=False):
@@ -206,6 +220,11 @@ class FlexKernelOptions(TypedDict, total=False):
     """
 
 
+class _KernelOptionsWithInternals(FlexKernelOptions, total=False):
+    OUTPUT_LOGSUMEXP: bool
+    OUTPUT_MAX: bool
+
+
 class AuxRequest(NamedTuple):
     """Request which auxiliary outputs to compute from flex_attention.
 
@@ -238,7 +257,9 @@ class _ModificationType(Enum):
     UNKNOWN = 3
 
 
-def _get_mod_type(fn: Callable) -> _ModificationType:
+def _get_mod_type(
+    fn: _score_mod_signature | _mask_mod_signature | Callable[..., Any],
+) -> _ModificationType:
     """Get the type of modification function.
     This function inspects the number of positional arguments of the function to determine
     the type of modification function. If the function has 5 positional arguments, it is
@@ -273,12 +294,12 @@ def _get_mod_type(fn: Callable) -> _ModificationType:
 
 # Need to define it here so that Dynamo doesn't skip it
 def _vmap_for_bhqkv(
-    fn: Callable,
+    fn: Callable[..., _R],
     prefix: tuple[int | None, ...],
     suffix: tuple[int | None, ...] = (),
     out_dims: int | list[int | None] = 0,
     group_dim: bool = False,
-):
+) -> Callable[..., _R]:
     """Used to vmap both score_mods and mask_mods over 4-dimensional/5-dimension inputs.
     Mapping over the [b, hq, q_idx, kv_idx] or [b, hkv, g, q_idx, kv_idx] dimensions.
 
@@ -343,7 +364,7 @@ def _sliced_mask_mod_error(
     head: Tensor,
     token_q: Tensor,
     token_kv: Tensor,
-) -> Tensor:
+) -> Never:
     """
     Raises helpful error when using mask_mod from a sliced BlockMask.
 
@@ -369,7 +390,7 @@ def _sliced_mask_mod_error(
 _LARGE_SPARSE_BLOCK_SIZE = 1 << 30
 
 
-def _ordered_to_dense(num_blocks_in_row: Tensor, col_indices: Tensor):
+def _ordered_to_dense(num_blocks_in_row: Tensor, col_indices: Tensor) -> Tensor:
     num_rows = col_indices.shape[-2]
     num_cols = col_indices.shape[-1]
     batch_dims = num_blocks_in_row.shape[:-1]
@@ -399,7 +420,7 @@ def create_dense_one(kv_num_blocks, kv_indices):
     return out
 
 
-def _dense_to_ordered(dense_mask) -> tuple[Tensor, Tensor]:
+def _dense_to_ordered(dense_mask: Tensor) -> tuple[Tensor, Tensor]:
     dense_mask = dense_mask.to(dtype=torch.int32)
     num_blocks_in_row = dense_mask.sum(dim=-1)
     col_indices = torch.argsort(dense_mask, dim=-1, descending=True, stable=True)
@@ -409,7 +430,9 @@ def _dense_to_ordered(dense_mask) -> tuple[Tensor, Tensor]:
     )
 
 
-def _transpose_ordered(num_blocks_in_row: Tensor, col_indices: Tensor):
+def _transpose_ordered(
+    num_blocks_in_row: Tensor, col_indices: Tensor
+) -> tuple[Tensor, Tensor]:
     dense = _ordered_to_dense(num_blocks_in_row, col_indices)
     return _dense_to_ordered(dense.transpose(-2, -1))
 
@@ -419,7 +442,7 @@ def _adjust_num_blocks_and_indices(
     indices: Tensor,
     new_num_rows: int,
     new_num_cols: int,
-):
+) -> tuple[Tensor, Tensor]:
     indices = indices[:, :, :new_num_rows, :new_num_cols]
     num_blocks = num_blocks[:, :, :new_num_rows]
     num_blocks = torch.where(num_blocks < new_num_cols, num_blocks, new_num_cols)
@@ -427,46 +450,228 @@ def _adjust_num_blocks_and_indices(
     return num_blocks, indices
 
 
-def _closure_contents(fn: object) -> tuple[object, ...]:
-    """Extract closure cell contents for comparison."""
-    closure = getattr(fn, "__closure__", None)
-    if closure is None:
-        return ()
-    return tuple(cell.cell_contents for cell in closure)
+# TreeSpec for an empty tuple — used as the sentinel when there are no closure leaves.
+_EMPTY_CLOSURE_SPEC = tree_flatten(())[1]
+
+
+class _ExtractedLeaf:
+    """Sentinel in _StrippedClosure.leaf_entries marking a position that is
+    filled from the extracted pytree leaves list during reconstruction."""
+
+    __slots__ = ()
+
+    def __repr__(self) -> str:
+        return "_EXTRACTED_LEAF"
+
+
+_EXTRACTED_LEAF = _ExtractedLeaf()
+
+
+class _FunctionLeaf(typing.NamedTuple):
+    """Entry in _StrippedClosure.leaf_entries for a recursively processed
+    function.  Stores enough information to reconstruct the function from
+    the extracted leaves during unflattening."""
+
+    stripped: _StrippedClosure | Callable[..., Any]
+    closure_spec: TreeSpec
+    n_extracted: int  # number of extracted leaves this function contributes
+
+
+class _StrippedClosure(typing.NamedTuple):
+    """Data container holding the parts of a function needed for reconstruction.
+
+    Created by _extract_closure_pytree when closure tensors are lifted into
+    pytree leaves.  Unlike a FunctionType with None-filled cells, this is not
+    callable — it is pure data stored in the pytree context.
+    """
+
+    code: types.CodeType
+    globals_dict: dict[str, Any]
+    name: str
+    qualname: str
+    defaults: tuple[Any, ...] | None
+    kwdefaults: dict[str, Any] | None
+    extra_dict: dict[str, Any]
+    # Per-position info for the closure's flattened leaves.
+    # _EXTRACTED_LEAF → position filled from the extracted leaves list.
+    # _FunctionLeaf   → recursively processed function (not a valid pytree leaf).
+    leaf_entries: tuple[_ExtractedLeaf | _FunctionLeaf, ...]
+
+
+def _extract_closure_pytree(
+    fn, _seen: set[int] | None = None
+) -> tuple[
+    tuple[BaseArgumentTypes, ...], TreeSpec, _StrippedClosure | Callable[..., Any]
+]:
+    """Extract closure contents as a flattened sub-pytree.
+
+    Returns (extracted_leaves, closure_spec, fn_or_stripped) where:
+    - extracted_leaves: flattened non-function contents from the closure,
+      plus any tensors/scalars recursively extracted from nested function
+      closures
+    - closure_spec: TreeSpec describing how to reconstruct the closure contents
+    - fn_or_stripped: either the original fn (no extraction) or a
+      _StrippedClosure carrying the function parts needed for reconstruction
+
+    Functions found among the closure leaves are recursively processed: their
+    own closure tensors are extracted into the leaves list, and their skeleton
+    is stored in _StrippedClosure.leaf_entries as a _FunctionLeaf.  All other
+    values (tensors, scalars, None, etc.) remain as extracted leaves.
+
+    If fn is not a plain function, has no closure, or has empty cells, returns
+    the original function unchanged with no closure leaves.
+
+    Skipped under Dynamo tracing (torch.compiler.is_compiling) because Dynamo
+    can't trace through closure cell introspection and handles freevars via its
+    own lifting mechanism.
+    """
+    if not inspect.isfunction(fn) or torch.compiler.is_compiling():
+        return (), _EMPTY_CLOSURE_SPEC, fn
+
+    # Cycle detection for self-referencing closures.
+    if _seen is None:
+        _seen = set()
+    if id(fn) in _seen:
+        return (), _EMPTY_CLOSURE_SPEC, fn
+    _seen.add(id(fn))
+
+    closure = fn.__closure__
+    if not closure:
+        return (), _EMPTY_CLOSURE_SPEC, fn
+
+    try:
+        contents = tuple(cell.cell_contents for cell in closure)
+    except ValueError:
+        # Empty cell (created but not yet assigned) — can't extract
+        return (), _EMPTY_CLOSURE_SPEC, fn
+
+    closure_leaves, closure_spec = tree_flatten(contents)
+
+    extracted: list[BaseArgumentTypes] = []
+    leaf_entries: list[_ExtractedLeaf | _FunctionLeaf] = []
+    for leaf in closure_leaves:
+        if inspect.isfunction(leaf):
+            child_extracted, child_spec, child_stripped = _extract_closure_pytree(
+                leaf, _seen
+            )
+            extracted.extend(child_extracted)
+            leaf_entries.append(
+                _FunctionLeaf(child_stripped, child_spec, len(child_extracted))
+            )
+        else:
+            extracted.append(leaf)
+            leaf_entries.append(_EXTRACTED_LEAF)
+
+    stripped = _StrippedClosure(
+        code=fn.__code__,
+        globals_dict=fn.__globals__,
+        name=fn.__name__,
+        qualname=fn.__qualname__,
+        defaults=fn.__defaults__,
+        kwdefaults=fn.__kwdefaults__,
+        extra_dict=dict(fn.__dict__) if fn.__dict__ else {},
+        leaf_entries=tuple(leaf_entries),
+    )
+
+    return tuple(extracted), closure_spec, stripped
+
+
+def _reconstruct_closure_fn(stripped, extracted_leaves, closure_spec):
+    """Rebuild a function from a _StrippedClosure and flattened extracted leaves."""
+    if not isinstance(stripped, _StrippedClosure):
+        return stripped
+
+    all_leaves: list[BaseArgumentTypes | Callable[..., Any]] = []
+    idx = 0
+    for entry in stripped.leaf_entries:
+        if isinstance(entry, _FunctionLeaf):
+            child_fn = _reconstruct_closure_fn(
+                entry.stripped,
+                extracted_leaves[idx : idx + entry.n_extracted],
+                entry.closure_spec,
+            )
+            all_leaves.append(child_fn)
+            idx += entry.n_extracted
+        else:
+            # _EXTRACTED_LEAF — take from extracted leaves
+            all_leaves.append(extracted_leaves[idx])
+            idx += 1
+
+    contents = tree_unflatten(all_leaves, closure_spec)
+    new_cells = tuple(types.CellType(v) for v in contents)
+
+    restored = types.FunctionType(
+        stripped.code,
+        stripped.globals_dict,
+        stripped.name,
+        stripped.defaults,
+        new_cells,
+    )
+    restored.__qualname__ = stripped.qualname
+    if stripped.kwdefaults:
+        restored.__kwdefaults__ = stripped.kwdefaults
+    if stripped.extra_dict:
+        restored.__dict__.update(stripped.extra_dict)
+
+    return restored
 
 
 class _MaskModWrapper:
-    """Wraps a mask_mod function with value-based equality.
+    """Wraps a mask_mod or _StrippedClosure with value-based equality.
 
     BlockMask stores an arbitrary callable (mask_mod) in its pytree context.
     The default __eq__ for functions uses identity comparison, which is too
     strict when the same closure is recreated (e.g., defined inside forward()).
-    This wrapper compares functions by their code object and closure contents.
+
+    When closure tensors have been extracted (by _extract_closure_pytree), fn
+    is a _StrippedClosure (pure data, not callable).  Equality compares the
+    code objects + closure_spec — no tensor dispatch is triggered.
+
+    When extraction is skipped (e.g., under Dynamo), fn is the original
+    callable and equality compares code objects + closure contents (for plain
+    functions) or delegates to __eq__ (for callable objects).
     """
 
-    __slots__ = ("fn",)
+    __slots__ = ("fn", "closure_spec")
 
-    def __init__(self, fn: _mask_mod_signature) -> None:
+    def __init__(self, fn, closure_spec=None) -> None:
         self.fn = fn
+        self.closure_spec = closure_spec
 
-    def __call__(self, *args, **kwargs):
-        return self.fn(*args, **kwargs)
+    def __call__(self, b: Tensor, h: Tensor, q_idx: Tensor, kv_idx: Tensor) -> Tensor:
+        if isinstance(self.fn, _StrippedClosure):
+            raise RuntimeError(
+                "_MaskModWrapper with _StrippedClosure is not callable — "
+                "use _reconstruct_closure_fn to rebuild the function first"
+            )
+        return self.fn(b, h, q_idx, kv_idx)
 
     def __eq__(self, other: object) -> bool:
         if not isinstance(other, _MaskModWrapper):
             return False
-        if self.fn is other.fn:
+        if self.fn is other.fn and self.closure_spec is other.closure_spec:
             return True
-        if (
-            inspect.isfunction(self.fn)
-            and inspect.isfunction(other.fn)
-            and self.fn.__code__ == other.fn.__code__
-            and _closure_contents(self.fn) == _closure_contents(other.fn)
+        # Extracted case: _StrippedClosure — compare code + closure_spec
+        if isinstance(self.fn, _StrippedClosure) and isinstance(
+            other.fn, _StrippedClosure
         ):
-            return True
+            return (
+                self.fn.code == other.fn.code
+                and self.closure_spec == other.closure_spec
+            )
+        # Non-extracted plain functions: compare code + closure contents
+        if inspect.isfunction(self.fn) and inspect.isfunction(other.fn):
+            return self.fn.__code__ == other.fn.__code__
+        # Callable objects: delegate to their __eq__
+        if not isinstance(self.fn, _StrippedClosure) and not isinstance(
+            other.fn, _StrippedClosure
+        ):
+            return self.fn == other.fn
         return False
 
     def __hash__(self) -> int:
+        if isinstance(self.fn, _StrippedClosure):
+            return hash(self.fn.code)
         if inspect.isfunction(self.fn):
             return hash(self.fn.__code__)
         return hash(self.fn)
@@ -618,7 +823,7 @@ def from_kv_blocks(
         mask_mod: _mask_mod_signature | None = None,
         seq_lengths: tuple[int, int] | None = None,
         compute_q_blocks: bool = True,
-    ):
+    ) -> Self:
         """
         Creates a BlockMask instance from key-value block information.
 
@@ -683,7 +888,43 @@ def from_kv_blocks(
             mask_mod=mask_mod,
         )
 
-    def as_tuple(self, flatten: bool = True):
+    @overload
+    def as_tuple(
+        self, flatten: Literal[True] = ...
+    ) -> tuple[
+        int,
+        int,
+        Tensor,
+        Tensor,
+        Tensor | None,
+        Tensor | None,
+        Tensor | None,
+        Tensor | None,
+        Tensor | None,
+        Tensor | None,
+        int,
+        int,
+        _mask_mod_signature,
+    ]: ...
+
+    @overload
+    def as_tuple(
+        self, flatten: Literal[False]
+    ) -> tuple[
+        tuple[int, int],
+        Tensor,
+        Tensor,
+        Tensor | None,
+        Tensor | None,
+        Tensor | None,
+        Tensor | None,
+        Tensor | None,
+        Tensor | None,
+        int | tuple[int, int],
+        _mask_mod_signature,
+    ]: ...
+
+    def as_tuple(self, flatten: bool = True) -> tuple[Any, ...]:
         """
         Returns a tuple of the attributes of the BlockMask.
 
@@ -713,7 +954,7 @@ def as_tuple(self, flatten: bool = True):
         )
 
     @property
-    def shape(self):
+    def shape(self) -> tuple[int, ...]:
         *batch_dims, _, _ = self.kv_indices.shape
         return tuple(batch_dims) + self.seq_lengths
 
@@ -724,7 +965,9 @@ def __str__(self) -> str:
         s += "\n)"
         return s
 
-    def __getitem__(self, index) -> "BlockMask":
+    def __getitem__(
+        self, index: int | slice | Tensor | tuple[int | slice | Tensor, ...]
+    ) -> Self:
         """
         Returns a new BlockMask instance by getting the mask for the given index position.
 
@@ -817,7 +1060,7 @@ def shape_or_none(x: torch.Tensor | None):
             f")"
         )
 
-    def _adjust(self, new_q_len: int, new_kv_len: int):
+    def _adjust(self, new_q_len: int, new_kv_len: int) -> Self:
         new_num_rows = (new_q_len + self.BLOCK_SIZE[0] - 1) // self.BLOCK_SIZE[0]
         new_num_cols = (new_kv_len + self.BLOCK_SIZE[1] - 1) // self.BLOCK_SIZE[1]
         new_kv_num_blocks, new_kv_indices = _adjust_num_blocks_and_indices(
@@ -847,7 +1090,7 @@ def _adjust(self, new_q_len: int, new_kv_len: int):
             self.mask_mod,
         )
 
-    def numel(self):
+    def numel(self) -> int:
         """Returns the number of elements (not accounting for sparsity) in the mask."""
         shape = self.shape
 
@@ -879,7 +1122,9 @@ def to_dense(self) -> Tensor:
             )
         return partial_dense
 
-    def to_string(self, grid_size=(20, 20), limit=4):
+    def to_string(
+        self, grid_size: int | tuple[int, int] = (20, 20), limit: int = 4
+    ) -> str:
         """Returns a string representation of the block mask. Quite nifty.
 
         If grid_size is -1, prints out an uncompressed version. Warning, it can be quite big!
@@ -945,7 +1190,7 @@ def cdiv(a, b):
 
         return "\n".join(total_vis)
 
-    def to(self, device: torch.device | str) -> "BlockMask":
+    def to(self, device: torch.device | str) -> BlockMask:
         """Moves the BlockMask to the specified device.
 
         Args:
@@ -983,58 +1228,87 @@ def _unwrap_context_value(attr: str, value: Any) -> Any:
             return value.fn
         return value
 
-    def _flatten(self):
+    def _flatten(
+        self,
+    ) -> tuple[tuple[BaseArgumentTypes | None, ...], tuple[Any, ...]]:
         """Flatten BlockMask into a list of tensors and context.
 
-        Wraps mask_mod in _MaskModWrapper for value-based comparison in TreeSpec.
+        Closure tensors from mask_mod are extracted into the leaves via
+        _extract_closure_pytree so they are visible to the tracing
+        infrastructure (instead of being hidden in the pytree context).
         """
         tensors = tuple(getattr(self, attr) for attr in self._TENSOR_ATTRS)
+        closure_leaves, closure_spec, stripped = _extract_closure_pytree(self.mask_mod)
+        all_leaves = tensors + closure_leaves
         context = tuple(
             self._wrap_context_value(attr, getattr(self, attr))
+            if attr != "mask_mod"
+            else _MaskModWrapper(stripped, closure_spec)
             for attr in self._CONTEXT_ATTRS
         )
-        return tensors, context
+        return all_leaves, context
 
     @classmethod
-    def _unflatten(cls, tensors, context):
-        """Unflatten tensors and context back into a BlockMask."""
-        kwargs = {
-            attr: cls._unwrap_context_value(attr, val)
-            for attr, val in zip(cls._CONTEXT_ATTRS, context)
-        }
-        kwargs.update(zip(cls._TENSOR_ATTRS, tensors))
-        # pyrefly: ignore [bad-argument-type]
+    def _unflatten(
+        cls,
+        leaves: tuple[Any, ...],
+        context: tuple[Any, ...],
+    ) -> Self:
+        """Unflatten leaves and context back into a BlockMask."""
+        n_regular = len(cls._TENSOR_ATTRS)
+        regular_leaves = leaves[:n_regular]
+        closure_leaves = leaves[n_regular:]
+        kwargs = {}
+        for attr, val in zip(cls._CONTEXT_ATTRS, context):
+            if attr == "mask_mod" and isinstance(val, _MaskModWrapper):
+                kwargs[attr] = _reconstruct_closure_fn(
+                    val.fn, closure_leaves, val.closure_spec
+                )
+            else:
+                kwargs[attr] = cls._unwrap_context_value(attr, val)
+        kwargs.update(zip(cls._TENSOR_ATTRS, regular_leaves))
         return cls(**kwargs)
 
-    def _flatten_with_keys(self):
+    def _flatten_with_keys(
+        self,
+    ) -> tuple[tuple[tuple[GetAttrKey, Any], ...], tuple[tuple[GetAttrKey, Any], ...]]:
         """Flatten BlockMask with keys for better tracing.
 
-        Wraps mask_mod in _MaskModWrapper for value-based comparison in TreeSpec.
+        Closure tensors from mask_mod are extracted into the leaves via
+        _extract_closure_pytree so they are visible to the tracing
+        infrastructure (instead of being hidden in the pytree context).
         """
         tensors = tuple(
             (GetAttrKey(attr), getattr(self, attr)) for attr in self._TENSOR_ATTRS
         )
+        closure_leaves, closure_spec, stripped = _extract_closure_pytree(self.mask_mod)
+        closure_with_keys = tuple(
+            (GetAttrKey(f"_closure_{i}"), leaf) for i, leaf in enumerate(closure_leaves)
+        )
+        all_leaves = tensors + closure_with_keys
         context = tuple(
             (GetAttrKey(attr), self._wrap_context_value(attr, getattr(self, attr)))
+            if attr != "mask_mod"
+            else (GetAttrKey(attr), _MaskModWrapper(stripped, closure_spec))
             for attr in self._CONTEXT_ATTRS
         )
-        return tensors, context
+        return all_leaves, context
 
 
-def _broadcast_to_dim(x, dim):
+def _broadcast_to_dim(x: Tensor, dim: int) -> Tensor:
     while x.dim() < dim:
         x = x.unsqueeze(0)
     return x
 
 
-def _round_up_to_multiple(x, multiple):
+def _round_up_to_multiple(x: int, multiple: int) -> int:
     return (x + multiple - 1) // multiple * multiple
 
 
 def _convert_mask_to_block_mask(
     mask: Tensor,
-    Q_BLOCK_SIZE=_DEFAULT_SPARSE_BLOCK_SIZE,
-    KV_BLOCK_SIZE=_DEFAULT_SPARSE_BLOCK_SIZE,
+    Q_BLOCK_SIZE: int = _DEFAULT_SPARSE_BLOCK_SIZE,
+    KV_BLOCK_SIZE: int = _DEFAULT_SPARSE_BLOCK_SIZE,
     separate_full_blocks: bool = False,
 ) -> tuple[Tensor, Tensor | None]:
     if mask.dtype != torch.bool:
@@ -1089,7 +1363,7 @@ def or_masks(*mask_mods: _mask_mod_signature) -> _mask_mod_signature:
     if not all(callable(arg) for arg in mask_mods):
         raise RuntimeError(f"All inputs should be callable mask_mods: {mask_mods}")
 
-    def or_mask(b, h, q_idx, kv_idx):
+    def or_mask(b: Tensor, h: Tensor, q_idx: Tensor, kv_idx: Tensor) -> Tensor:
         result = b.new_zeros((), dtype=torch.bool)
         for mask in mask_mods:
             result = result | mask(b, h, q_idx, kv_idx)
@@ -1103,7 +1377,7 @@ def and_masks(*mask_mods: _mask_mod_signature) -> _mask_mod_signature:
     if not all(callable(arg) for arg in mask_mods):
         raise RuntimeError(f"All inputs should be callable mask_mods: {mask_mods}")
 
-    def and_mask(b, h, q_idx, kv_idx):
+    def and_mask(b: Tensor, h: Tensor, q_idx: Tensor, kv_idx: Tensor) -> Tensor:
         result = b.new_ones((), dtype=torch.bool)
         for mask in mask_mods:
             result = result & mask(b, h, q_idx, kv_idx)
@@ -1113,9 +1387,9 @@ def and_mask(b, h, q_idx, kv_idx):
 
 
 def _convert_block_mask_to_mask(
-    block_mask,
-    KV_BLOCK_SIZE=_DEFAULT_SPARSE_BLOCK_SIZE,
-    Q_BLOCK_SIZE=_DEFAULT_SPARSE_BLOCK_SIZE,
+    block_mask: Tensor,
+    KV_BLOCK_SIZE: int = _DEFAULT_SPARSE_BLOCK_SIZE,
+    Q_BLOCK_SIZE: int = _DEFAULT_SPARSE_BLOCK_SIZE,
 ) -> Tensor:
     if block_mask.dim() != 4:
         raise AssertionError(f"block_mask.dim() must be 4, got {block_mask.dim()}")
@@ -1129,7 +1403,7 @@ def _convert_block_mask_to_mask(
 
 def _create_sparse_block_from_block_mask(
     block_mask: tuple[Tensor, Tensor | None],
-    mask_mod: Callable | None,
+    mask_mod: _mask_mod_signature | None,
     seq_lengths: tuple[int, int],
     Q_BLOCK_SIZE: int = _DEFAULT_SPARSE_BLOCK_SIZE,
     KV_BLOCK_SIZE: int = _DEFAULT_SPARSE_BLOCK_SIZE,
@@ -1309,10 +1583,13 @@ def _apply_kernel_options(
     key: Tensor,
     value: Tensor,
     return_lse: bool,
-    kernel_options,
+    kernel_options: FlexKernelOptions | None,
     return_aux: AuxRequest | None = None,
-):
-    kernel_options = {} if kernel_options is None else dict(kernel_options)
+) -> _KernelOptionsWithInternals:
+    kernel_options = cast(
+        _KernelOptionsWithInternals,
+        {} if kernel_options is None else dict(kernel_options),
+    )
 
     if "BACKEND" in kernel_options and kernel_options.get(
         "FORCE_USE_FLEX_ATTENTION", False
@@ -1370,6 +1647,11 @@ def _apply_kernel_options(
     # If forward kernel needs to return max is decided by this rule internally.
     if "OUTPUT_MAX" in kernel_options:
         raise AssertionError("OUTPUT_MAX must not be in kernel_options")
+    if kernel_options["BACKEND"] == "FLASH" and output_max:
+        raise NotImplementedError(
+            "Returning max scores is not supported with BACKEND='FLASH'. "
+            "Use return_aux=AuxRequest(lse=True) or omit max_scores."
+        )
     kernel_options["OUTPUT_MAX"] = output_max
     if any_inputs_on_cpu_device and output_max:
         # CPU doesn't support returning max yet
@@ -1465,6 +1747,75 @@ def is_col_major(tensor: Tensor) -> bool:
     return query, key, value
 
 
+@overload
+def flex_attention(
+    query: Tensor,
+    key: Tensor,
+    value: Tensor,
+    score_mod: _score_mod_signature | None = ...,
+    block_mask: BlockMask | None = ...,
+    scale: float | None = ...,
+    enable_gqa: bool = ...,
+    return_lse: Literal[False] = ...,
+    kernel_options: FlexKernelOptions | None = ...,
+    *,
+    return_aux: None = ...,
+) -> Tensor: ...
+
+
+@overload
+@deprecated(
+    "return_lse is deprecated and will be removed in v2.10. "
+    "Use return_aux=AuxRequest(lse=True) instead.",
+    category=FutureWarning,
+)
+def flex_attention(
+    query: Tensor,
+    key: Tensor,
+    value: Tensor,
+    score_mod: _score_mod_signature | None = ...,
+    block_mask: BlockMask | None = ...,
+    scale: float | None = ...,
+    enable_gqa: bool = ...,
+    return_lse: Literal[True] = ...,
+    kernel_options: FlexKernelOptions | None = ...,
+    *,
+    return_aux: None = ...,
+) -> tuple[Tensor, Tensor]: ...
+
+
+@overload
+def flex_attention(
+    query: Tensor,
+    key: Tensor,
+    value: Tensor,
+    score_mod: _score_mod_signature | None = ...,
+    block_mask: BlockMask | None = ...,
+    scale: float | None = ...,
+    enable_gqa: bool = ...,
+    return_lse: bool = ...,
+    kernel_options: FlexKernelOptions | None = ...,
+    *,
+    return_aux: AuxRequest,
+) -> tuple[Tensor, AuxOutput]: ...
+
+
+@overload
+def flex_attention(
+    query: Tensor,
+    key: Tensor,
+    value: Tensor,
+    score_mod: _score_mod_signature | None = ...,
+    block_mask: BlockMask | None = ...,
+    scale: float | None = ...,
+    enable_gqa: bool = ...,
+    return_lse: Literal[True] = ...,
+    kernel_options: FlexKernelOptions | None = ...,
+    *,
+    return_aux: AuxRequest,
+) -> Never: ...
+
+
 def flex_attention(
     query: Tensor,
     key: Tensor,
@@ -1733,7 +2084,7 @@ def _flex_attention_hop_wrapper(*args, **kwargs):
             key,
             value,
             score_mod,
-            block_mask.as_tuple(),  # type: ignore[union-attr]
+            block_mask.as_tuple(),
             scale,
             kernel_options,
         )
diff --git a/torch/nn/attention/varlen.py b/torch/nn/attention/varlen.py
index 5be971e1bd96c..454bb82c45e2a 100644
--- a/torch/nn/attention/varlen.py
+++ b/torch/nn/attention/varlen.py
@@ -14,7 +14,7 @@
 
 log = logging.getLogger(__name__)
 
-__all__ = ["varlen_attn", "AuxRequest"]
+__all__ = ["varlen_attn", "varlen_attn_out", "AuxRequest"]
 
 
 def _normalize_window_size(window_size: list[int] | None) -> list[int]:
@@ -48,12 +48,16 @@ def _varlen_attn(
     key: torch.Tensor,
     value: torch.Tensor,
     cu_seq_q: torch.Tensor,
-    cu_seq_k: torch.Tensor,
+    cu_seq_k: torch.Tensor | None,
     max_q: int,
     max_k: int,
     is_causal: bool = False,
     scale: float | None = None,
     window_size: list[int] | None = None,
+    enable_gqa: bool = False,
+    seqused_k: torch.Tensor | None = None,
+    block_table: torch.Tensor | None = None,
+    num_splits: int | None = None,
 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
     """
     Private custom op for variable-length attention.
@@ -67,10 +71,23 @@ def _varlen_attn(
     if use_cudnn:
         log.info("Using cuDNN backend for varlen_attn")
 
+        if enable_gqa:
+            # TODO: check this
+            raise RuntimeError("GQA is not supported with the cuDNN backend.")
+        if num_splits is not None:
+            # TODO: check this
+            raise RuntimeError("num_splits is not supported with the cuDNN backend.")
         if window_size[0] != -1 or window_size[1] != -1:
             raise RuntimeError(
                 "cuDNN backend does not support window attention. Please use Flash Attention backend."
             )
+        if seqused_k is not None or block_table is not None:
+            # TODO: cuDNN supports per-sequence KV lengths via SEQ_LEN_KV + padding_mask,
+            # but _cudnn_attention_forward doesn't expose it yet.
+            raise RuntimeError(
+                "seqused_k/block_table is not yet supported with the cuDNN backend."
+            )
+
         result = torch.ops.aten._cudnn_attention_forward(
             query,
             key,
@@ -104,6 +121,9 @@ def _varlen_attn(
             scale=scale,
             window_size_left=window_size[0],
             window_size_right=window_size[1],
+            seqused_k=seqused_k,
+            block_table=block_table,
+            num_splits=num_splits,
         )
 
     rng_state_ = torch.zeros(
@@ -118,12 +138,16 @@ def _varlen_attn_fake(
     key: torch.Tensor,
     value: torch.Tensor,
     cu_seq_q: torch.Tensor,
-    cu_seq_k: torch.Tensor,
+    cu_seq_k: torch.Tensor | None,
     max_q: int,
     max_k: int,
     is_causal: bool = False,
     scale: float | None = None,
     window_size: list[int] | None = None,
+    enable_gqa: bool = False,
+    seqused_k: torch.Tensor | None = None,
+    block_table: torch.Tensor | None = None,
+    num_splits: int | None = None,
 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
     """
     Fake implementation for meta tensor computation and tracing.
@@ -140,16 +164,9 @@ def _varlen_attn_fake(
     # For varlen path: logsumexp shape is (num_heads, total_q)
     total_q = query.size(0)
     num_heads = query.size(1)
-    if torch.version.hip:
-        # ROCm uses batched format: [batch_size, num_heads, max_q]
-        batch_size = cu_seq_q.size(0) - 1
-        logsumexp = torch.empty(
-            (batch_size, num_heads, max_q), dtype=torch.float, device=query.device
-        )
-    else:
-        logsumexp = torch.empty(
-            (num_heads, total_q), dtype=torch.float, device=query.device
-        )
+    logsumexp = torch.empty(
+        (num_heads, total_q), dtype=torch.float, device=query.device
+    )
 
     rng_state = torch.empty((2,), dtype=torch.uint64, device=query.device)
 
@@ -161,23 +178,29 @@ def varlen_attn(
     key: torch.Tensor,
     value: torch.Tensor,
     cu_seq_q: torch.Tensor,
-    cu_seq_k: torch.Tensor,
+    cu_seq_k: torch.Tensor | None,
     max_q: int,
     max_k: int,
     *,
     return_aux: AuxRequest | None = None,
     scale: float | None = None,
     window_size: tuple[int, int] = (-1, -1),
+    enable_gqa: bool = False,
+    seqused_k: torch.Tensor | None = None,
+    block_table: torch.Tensor | None = None,
+    num_splits: int | None = None,
 ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
-    """
-    Compute variable-length attention using Flash Attention.
+    r"""Compute variable-length attention using Flash Attention.
+
     This function is similar to scaled_dot_product_attention but optimized for
     variable-length sequences using cumulative sequence position tensors.
 
     Args:
-        query (Tensor): Query tensor; shape :math:`(T_q, H, D)`
-        key (Tensor): Key tensor; shape :math:`(T_k, H, D)`
-        value (Tensor): Value tensor; shape :math:`(T_k, H, D)`
+        query (Tensor): Query tensor; shape :math:`(T_q, H_q, D)`
+        key (Tensor): Key tensor; shape :math:`(T_k, H_{kv}, D)`, or
+            :math:`(\text{total\_pages}, \text{page\_size}, H_{kv}, D)` when ``block_table`` is provided.
+        value (Tensor): Value tensor; shape :math:`(T_k, H_{kv}, D)`, or
+            :math:`(\text{total\_pages}, \text{page\_size}, H_{kv}, D)` when ``block_table`` is provided.
         cu_seq_q (Tensor): Cumulative sequence positions for queries; shape :math:`(N+1,)`
         cu_seq_k (Tensor): Cumulative sequence positions for keys/values; shape :math:`(N+1,)`
         max_q (int): Maximum query sequence length in the batch.
@@ -187,18 +210,50 @@ def varlen_attn(
         window_size (tuple[int, int], optional): Window size for sliding window attention as (left, right).
             Use (-1, -1) for full attention (default), (-1, 0) for causal attention,
             or (W, 0) for causal attention with sliding window of size W.
+        enable_gqa (bool): If set to True, enables Grouped Query Attention (GQA)
+            and allows key/value to have fewer heads than query.
+            Each KV head is shared by a group of :math:`H_q / H_{kv}` query heads,
+            so :math:`H_q` must be divisible by :math:`H_{kv}`.
+            Default is False.
+        seqused_k (Tensor, optional): Number of valid KV tokens per batch element; shape :math:`(N,)`.
+            When set, only the first ``seqused_k[i]`` tokens in the key/value sequence for batch
+            element *i* participate in attention. Useful for KV-cache decoding where the cache slot
+            is larger than the actual sequence. Inference-only (not supported in backward).
+        block_table (Tensor, optional): Block table for paged KV cache; shape
+            :math:`(N, \text{max\_pages\_per\_seq})`, dtype ``int32``.
+            Requires ``seqused_k``. Inference-only (not supported in backward).
+
+            When ``block_table`` is provided, ``key`` and ``value`` are a "pool" of
+            pages of tokens of KV data and the pages belong to any sequence/order.
+            The ``block_table`` is what maps each sequence's logical chunks
+            back to physical pages in this pool.
+
+            ``seqused_k[i]`` tells the kernel how many tokens in sequence *i* are
+            actually valid, since the last page is typically only partially filled.
+        num_splits (int, optional): Number of splits for split-KV. Set to ``1``
+            to disable split-KV which enables batch invariance. Split-KV
+            parallelizes the key/value sequence dimension across multiple thread
+            blocks and combines partial results. The split decision depends
+            on ``max_k`` (the longest sequence in the batch), so different batch
+            compositions can change the reduction order and produce different
+            floating-point results for the same sequence. When this is disabled,
+            bitwise identical outputs are guaranteed for a given sequence
+            regardless of what other sequences are in the batch, at the
+            cost of lower GPU utilization when there are few queries. When
+            ``None`` (default), the kernel chooses automatically.
 
     Returns:
-        output (Tensor): Output tensor from attention computation; shape :math:`(T_q, H, D)`.
+        output (Tensor): Output tensor from attention computation; shape :math:`(T_q, H_q, D)`.
 
         If ``return_aux`` is not None and ``return_aux.lse`` is True:
-            lse (Tensor): Log-sum-exp of attention scores; shape :math:`(T_q, H)`.
+            lse (Tensor): Log-sum-exp of attention scores; shape :math:`(T_q, H_q)`.
 
     Shape legend:
         - :math:`N`: Batch size
         - :math:`T_q`: Total number of query tokens in the batch (sum of all query sequence lengths)
         - :math:`T_k`: Total number of key/value tokens in the batch (sum of all key/value sequence lengths)
-        - :math:`H`: Number of attention heads
+        - :math:`H_q`: Number of query attention heads
+        - :math:`H_{kv}`: Number of key/value attention heads (equal to :math:`H_q` unless GQA is enabled)
         - :math:`D`: Head dimension
 
     Example::
@@ -234,6 +289,21 @@ def varlen_attn(
         ...     query, key, value, cu_seq, cu_seq, max_len, max_len
         ... )
     """
+
+    num_heads_q = query.size(1)
+    num_heads_k = key.size(2) if block_table is not None else key.size(1)
+    if not enable_gqa and num_heads_q != num_heads_k:
+        raise ValueError(
+            f"Expect query and key/value to have the same number of heads "
+            f"but got Hq={num_heads_q} and Hkv={num_heads_k}. "
+            f"Try setting enable_gqa=True for GQA."
+        )
+    if enable_gqa and num_heads_q % num_heads_k != 0:
+        raise ValueError(
+            f"Expect number of query heads to be a multiple of kv heads for GQA "
+            f"but got Hq={num_heads_q} and Hkv={num_heads_k}."
+        )
+
     is_causal = window_size == (-1, 0)
     out, lse, _ = torch.ops.torch_attn._varlen_attn(
         query,
@@ -246,6 +316,155 @@ def varlen_attn(
         is_causal,
         scale,
         list(window_size),
+        enable_gqa,
+        seqused_k,
+        block_table,
+        num_splits,
+    )
+    if return_aux is not None and return_aux.lse:
+        return out, lse
+    return out
+
+
+@torch.library.custom_op("torch_attn::_varlen_attn_out", mutates_args={"out"})
+def _varlen_attn_out(
+    out: torch.Tensor,
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    cu_seq_q: torch.Tensor,
+    cu_seq_k: torch.Tensor | None,
+    max_q: int,
+    max_k: int,
+    is_causal: bool = False,
+    scale: float | None = None,
+    window_size: list[int] | None = None,
+    enable_gqa: bool = False,
+    seqused_k: torch.Tensor | None = None,
+    block_table: torch.Tensor | None = None,
+    num_splits: int | None = None,
+) -> torch.Tensor:
+    """
+    Private custom op for variable-length attention with pre-allocated output.
+    Same as _varlen_attn but writes the attention output into the provided out tensor.
+    """
+    window_size = _normalize_window_size(window_size)
+
+    use_cudnn = query.is_cuda and _should_use_cudnn(query.device.index)
+
+    if use_cudnn:
+        # TODO: look into this
+        raise RuntimeError("cuDNN backend does not support out variant.")
+
+    log.info("Using Flash Attention backend for varlen_attn_out")
+    softmax_lse = torch.ops.aten._flash_attention_forward_no_dropout_inplace(
+        out,
+        query,
+        key,
+        value,
+        cu_seq_q,
+        cu_seq_k,
+        max_q,
+        max_k,
+        0.0,  # dropout_p hardcoded to 0.0
+        is_causal,
+        False,  # return_debug_mask
+        scale=scale,
+        window_size_left=window_size[0],
+        window_size_right=window_size[1],
+        seqused_k=seqused_k,
+        block_table=block_table,
+        num_splits=num_splits,
+    )
+
+    return softmax_lse
+
+
+@_varlen_attn_out.register_fake
+def _varlen_attn_out_fake(
+    out: torch.Tensor,
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    cu_seq_q: torch.Tensor,
+    cu_seq_k: torch.Tensor | None,
+    max_q: int,
+    max_k: int,
+    is_causal: bool = False,
+    scale: float | None = None,
+    window_size: list[int] | None = None,
+    enable_gqa: bool = False,
+    seqused_k: torch.Tensor | None = None,
+    block_table: torch.Tensor | None = None,
+    num_splits: int | None = None,
+) -> torch.Tensor:
+    """
+    Fake implementation for meta tensor computation and tracing.
+    """
+    total_q = query.size(0)
+    num_heads = query.size(1)
+    logsumexp = torch.empty(
+        (num_heads, total_q), dtype=torch.float, device=query.device
+    )
+
+    return logsumexp
+
+
+def varlen_attn_out(
+    out: torch.Tensor,
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    cu_seq_q: torch.Tensor,
+    cu_seq_k: torch.Tensor | None,
+    max_q: int,
+    max_k: int,
+    *,
+    return_aux: AuxRequest | None = None,
+    scale: float | None = None,
+    window_size: tuple[int, int] = (-1, -1),
+    enable_gqa: bool = False,
+    seqused_k: torch.Tensor | None = None,
+    block_table: torch.Tensor | None = None,
+    num_splits: int | None = None,
+) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
+    r"""Compute variable-length attention using Flash Attention with a pre-allocated output tensor.
+
+    Same as :func:`varlen_attn` but writes the attention output into the provided ``out`` tensor
+    instead of allocating a new one.
+
+    """
+    num_heads_q = query.size(1)
+    num_heads_k = key.size(2) if block_table is not None else key.size(1)
+    if not enable_gqa and num_heads_q != num_heads_k:
+        raise ValueError(
+            f"Expect query and key/value to have the same number of heads "
+            f"but got Hq={num_heads_q} and Hkv={num_heads_k}. "
+            f"Try setting enable_gqa=True for GQA."
+        )
+    if enable_gqa and num_heads_q % num_heads_k != 0:
+        raise ValueError(
+            f"Expect number of query heads to be a multiple of kv heads for GQA "
+            f"but got Hq={num_heads_q} and Hkv={num_heads_k}."
+        )
+
+    is_causal = window_size == (-1, 0)
+    lse = torch.ops.torch_attn._varlen_attn_out(
+        out,
+        query,
+        key,
+        value,
+        cu_seq_q,
+        cu_seq_k,
+        max_q,
+        max_k,
+        is_causal,
+        scale,
+        list(window_size),
+        enable_gqa,
+        seqused_k,
+        block_table,
+        num_splits,
     )
     if return_aux is not None and return_aux.lse:
         return out, lse
@@ -264,9 +483,18 @@ def _setup_context(ctx: Any, inputs: tuple[Any, ...], output: Any) -> None:
         is_causal,
         scale,
         window_size,
+        enable_gqa,
+        seqused_k,
+        block_table,
+        num_splits,
     ) = inputs
     out, lse, rng_state = output
 
+    if seqused_k is not None:
+        raise RuntimeError("seqused_k is an inference-only parameter.")
+    if block_table is not None:
+        raise RuntimeError("block_table is an inference-only parameter.")
+
     ctx.save_for_backward(query, key, value, cu_seq_q, cu_seq_k, out, lse, rng_state)
 
     ctx.max_q = max_q
@@ -401,7 +629,26 @@ def _backward(
         scale,
         window_size,
     )
-    return dq, dk, dv, None, None, None, None, None, None, None
+    # cu_seq_q, cu_seq_k, max_q, max_k, is_causal, scale, window_size, \
+    # enable_gqa, seqused_k, block_table, num_splits
+    num_params = 11
+    return (dq, dk, dv, *((None,) * num_params))
 
 
 _varlen_attn.register_autograd(_backward, setup_context=_setup_context)
+
+torch._dynamo.disallow_in_graph(
+    torch.ops.aten._flash_attention_forward_no_dropout_inplace
+)
+
+from torch.utils.flop_counter import (
+    _varlen_attn_backward_flop,
+    _varlen_attn_forward_flop,
+    _varlen_attn_out_flop,
+    flop_registry,
+)
+
+
+flop_registry[torch.ops.torch_attn._varlen_attn] = _varlen_attn_forward_flop
+flop_registry[torch.ops.torch_attn._varlen_attn_out] = _varlen_attn_out_flop
+flop_registry[torch.ops.torch_attn._varlen_attn_backward] = _varlen_attn_backward_flop
diff --git a/torch/nn/functional.py b/torch/nn/functional.py
index 04467c65ef013..0906e86ee623c 100644
--- a/torch/nn/functional.py
+++ b/torch/nn/functional.py
@@ -439,10 +439,10 @@
 def fractional_max_pool2d_with_indices(
     input: Tensor,
     kernel_size: BroadcastingList2[int],
-    output_size: Optional[BroadcastingList2[int]] = None,
-    output_ratio: Optional[BroadcastingList2[float]] = None,
+    output_size: Optional[BroadcastingList2[int]] = None,  # noqa: UP045
+    output_ratio: Optional[BroadcastingList2[float]] = None,  # noqa: UP045
     return_indices: bool = False,
-    _random_samples: Optional[Tensor] = None,
+    _random_samples: Tensor | None = None,
 ) -> tuple[Tensor, Tensor]:  # noqa: D400
     r"""
     fractional_max_pool2d(input, kernel_size, output_size=None, output_ratio=None, return_indices=False, _random_samples=None)
@@ -517,10 +517,10 @@ def fractional_max_pool2d_with_indices(
 def _fractional_max_pool2d(
     input: Tensor,
     kernel_size: BroadcastingList2[int],
-    output_size: Optional[BroadcastingList2[int]] = None,
-    output_ratio: Optional[BroadcastingList2[float]] = None,
+    output_size: Optional[BroadcastingList2[int]] = None,  # noqa: UP045
+    output_ratio: Optional[BroadcastingList2[float]] = None,  # noqa: UP045
     return_indices: bool = False,
-    _random_samples: Optional[Tensor] = None,
+    _random_samples: Tensor | None = None,
 ) -> Tensor:
     if has_torch_function_variadic(input, _random_samples):
         return handle_torch_function(
@@ -552,10 +552,10 @@ def _fractional_max_pool2d(
 def fractional_max_pool3d_with_indices(
     input: Tensor,
     kernel_size: BroadcastingList3[int],
-    output_size: Optional[BroadcastingList3[int]] = None,
-    output_ratio: Optional[BroadcastingList3[float]] = None,
+    output_size: Optional[BroadcastingList3[int]] = None,  # noqa: UP045
+    output_ratio: Optional[BroadcastingList3[float]] = None,  # noqa: UP045
     return_indices: bool = False,
-    _random_samples: Optional[Tensor] = None,
+    _random_samples: Tensor | None = None,
 ) -> tuple[Tensor, Tensor]:  # noqa: D400
     r"""
     fractional_max_pool3d(input, kernel_size, output_size=None, output_ratio=None, return_indices=False, _random_samples=None)
@@ -634,10 +634,10 @@ def fractional_max_pool3d_with_indices(
 def _fractional_max_pool3d(
     input: Tensor,
     kernel_size: BroadcastingList3[int],
-    output_size: Optional[BroadcastingList3[int]] = None,
-    output_ratio: Optional[BroadcastingList3[float]] = None,
+    output_size: Optional[BroadcastingList3[int]] = None,  # noqa: UP045
+    output_ratio: Optional[BroadcastingList3[float]] = None,  # noqa: UP045
     return_indices: bool = False,
-    _random_samples: Optional[Tensor] = None,
+    _random_samples: Tensor | None = None,
 ) -> Tensor:
     if has_torch_function_variadic(input, _random_samples):
         return handle_torch_function(
@@ -669,7 +669,7 @@ def _fractional_max_pool3d(
 def max_pool1d_with_indices(
     input: Tensor,
     kernel_size: BroadcastingList1[int],
-    stride: Optional[BroadcastingList1[int]] = None,
+    stride: Optional[BroadcastingList1[int]] = None,  # noqa: UP045
     padding: BroadcastingList1[int] = 0,
     dilation: BroadcastingList1[int] = 1,
     ceil_mode: bool = False,
@@ -722,7 +722,7 @@ def max_pool1d_with_indices(
 def _max_pool1d(
     input: Tensor,
     kernel_size: BroadcastingList1[int],
-    stride: Optional[BroadcastingList1[int]] = None,
+    stride: Optional[BroadcastingList1[int]] = None,  # noqa: UP045
     padding: BroadcastingList1[int] = 0,
     dilation: BroadcastingList1[int] = 1,
     ceil_mode: bool = False,
@@ -759,7 +759,7 @@ def _max_pool1d(
 def max_pool2d_with_indices(
     input: Tensor,
     kernel_size: BroadcastingList2[int],
-    stride: Optional[BroadcastingList2[int]] = None,
+    stride: Optional[BroadcastingList2[int]] = None,  # noqa: UP045
     padding: BroadcastingList2[int] = 0,
     dilation: BroadcastingList2[int] = 1,
     ceil_mode: bool = False,
@@ -812,7 +812,7 @@ def max_pool2d_with_indices(
 def _max_pool2d(
     input: Tensor,
     kernel_size: BroadcastingList2[int],
-    stride: Optional[BroadcastingList2[int]] = None,
+    stride: Optional[BroadcastingList2[int]] = None,  # noqa: UP045
     padding: BroadcastingList2[int] = 0,
     dilation: BroadcastingList2[int] = 1,
     ceil_mode: bool = False,
@@ -849,7 +849,7 @@ def _max_pool2d(
 def max_pool3d_with_indices(
     input: Tensor,
     kernel_size: BroadcastingList3[int],
-    stride: Optional[BroadcastingList3[int]] = None,
+    stride: Optional[BroadcastingList3[int]] = None,  # noqa: UP045
     padding: BroadcastingList3[int] = 0,
     dilation: BroadcastingList3[int] = 1,
     ceil_mode: bool = False,
@@ -902,7 +902,7 @@ def max_pool3d_with_indices(
 def _max_pool3d(
     input: Tensor,
     kernel_size: BroadcastingList3[int],
-    stride: Optional[BroadcastingList3[int]] = None,
+    stride: Optional[BroadcastingList3[int]] = None,  # noqa: UP045
     padding: BroadcastingList3[int] = 0,
     dilation: BroadcastingList3[int] = 1,
     ceil_mode: bool = False,
@@ -941,7 +941,7 @@ def _unpool_output_size(
     kernel_size: list[int],
     stride: list[int],
     padding: list[int],
-    output_size: Optional[list[int]],
+    output_size: list[int] | None,
 ) -> list[int]:
     input_size = input.size()
     default_size = torch.jit.annotate(list[int], [])
@@ -977,9 +977,9 @@ def max_unpool1d(
     input: Tensor,
     indices: Tensor,
     kernel_size: BroadcastingList1[int],
-    stride: Optional[BroadcastingList1[int]] = None,
+    stride: Optional[BroadcastingList1[int]] = None,  # noqa: UP045
     padding: BroadcastingList1[int] = 0,
-    output_size: Optional[BroadcastingList1[int]] = None,
+    output_size: Optional[BroadcastingList1[int]] = None,  # noqa: UP045
 ) -> Tensor:
     r"""Compute a partial inverse of :class:`MaxPool1d`.
 
@@ -1016,9 +1016,9 @@ def max_unpool2d(
     input: Tensor,
     indices: Tensor,
     kernel_size: BroadcastingList2[int],
-    stride: Optional[BroadcastingList2[int]] = None,
+    stride: Optional[BroadcastingList2[int]] = None,  # noqa: UP045
     padding: BroadcastingList2[int] = 0,
-    output_size: Optional[BroadcastingList2[int]] = None,
+    output_size: Optional[BroadcastingList2[int]] = None,  # noqa: UP045
 ) -> Tensor:
     r"""Compute a partial inverse of :class:`MaxPool2d`.
 
@@ -1049,9 +1049,9 @@ def max_unpool3d(
     input: Tensor,
     indices: Tensor,
     kernel_size: BroadcastingList3[int],
-    stride: Optional[BroadcastingList3[int]] = None,
+    stride: Optional[BroadcastingList3[int]] = None,  # noqa: UP045
     padding: BroadcastingList3[int] = 0,
-    output_size: Optional[BroadcastingList3[int]] = None,
+    output_size: Optional[BroadcastingList3[int]] = None,  # noqa: UP045
 ) -> Tensor:
     r"""Compute a partial inverse of :class:`MaxPool3d`.
 
@@ -1082,7 +1082,7 @@ def lp_pool3d(
     input: Tensor,
     norm_type: int | float,
     kernel_size: BroadcastingList3[int],
-    stride: Optional[BroadcastingList3[int]] = None,
+    stride: Optional[BroadcastingList3[int]] = None,  # noqa: UP045
     ceil_mode: bool = False,
 ) -> Tensor:
     r"""
@@ -1123,7 +1123,7 @@ def lp_pool2d(
     input: Tensor,
     norm_type: int | float,
     kernel_size: BroadcastingList2[int],
-    stride: Optional[BroadcastingList2[int]] = None,
+    stride: Optional[BroadcastingList2[int]] = None,  # noqa: UP045
     ceil_mode: bool = False,
 ) -> Tensor:
     r"""
@@ -1162,7 +1162,7 @@ def lp_pool1d(
     input: Tensor,
     norm_type: int | float,
     kernel_size: int,
-    stride: Optional[BroadcastingList1[int]] = None,
+    stride: Optional[BroadcastingList1[int]] = None,  # noqa: UP045
     ceil_mode: bool = False,
 ) -> Tensor:
     r"""Apply a 1D power-average pooling over an input signal composed of several input planes.
@@ -2088,9 +2088,9 @@ def _get_softmax_dim(name: str, ndim: int, stacklevel: int) -> int:
 
 def softmin(
     input: Tensor,
-    dim: Optional[int] = None,
+    dim: int | None = None,
     _stacklevel: int = 3,
-    dtype: Optional[DType] = None,
+    dtype: DType | None = None,
 ) -> Tensor:
     r"""Apply a softmin function.
 
@@ -2121,9 +2121,9 @@ def softmin(
 
 def softmax(
     input: Tensor,
-    dim: Optional[int] = None,
+    dim: int | None = None,
     _stacklevel: int = 3,
-    dtype: Optional[DType] = None,
+    dtype: DType | None = None,
 ) -> Tensor:
     r"""Apply a softmax function.
 
@@ -2238,9 +2238,9 @@ def gumbel_softmax(
 
 def log_softmax(
     input: Tensor,
-    dim: Optional[int] = None,
+    dim: int | None = None,
     _stacklevel: int = 3,
-    dtype: Optional[DType] = None,
+    dtype: DType | None = None,
 ) -> Tensor:
     r"""Apply a softmax followed by a logarithm.
 
@@ -2455,8 +2455,8 @@ def _no_grad_embedding_renorm_(
 def embedding(
     input: Tensor,
     weight: Tensor,
-    padding_idx: Optional[int] = None,
-    max_norm: Optional[float] = None,
+    padding_idx: int | None = None,
+    max_norm: float | None = None,
     norm_type: float = 2.0,
     scale_grad_by_freq: bool = False,
     sparse: bool = False,
@@ -2570,15 +2570,15 @@ def embedding(
 def embedding_bag(
     input: Tensor,
     weight: Tensor,
-    offsets: Optional[Tensor] = None,
-    max_norm: Optional[float] = None,
+    offsets: Tensor | None = None,
+    max_norm: float | None = None,
     norm_type: float = 2,
     scale_grad_by_freq: bool = False,
     mode: str = "mean",
     sparse: bool = False,
-    per_sample_weights: Optional[Tensor] = None,
+    per_sample_weights: Tensor | None = None,
     include_last_offset: bool = False,
-    padding_idx: Optional[int] = None,
+    padding_idx: int | None = None,
 ) -> Tensor:
     r"""Compute sums, means or maxes of `bags` of embeddings.
 
@@ -2810,10 +2810,10 @@ def _verify_batch_size(size: list[int]) -> None:
 
 def batch_norm(
     input: Tensor,
-    running_mean: Optional[Tensor],
-    running_var: Optional[Tensor],
-    weight: Optional[Tensor] = None,
-    bias: Optional[Tensor] = None,
+    running_mean: Tensor | None,
+    running_var: Tensor | None,
+    weight: Tensor | None = None,
+    bias: Tensor | None = None,
     training: bool = False,
     momentum: float = 0.1,
     eps: float = 1e-5,
@@ -2840,8 +2840,12 @@ def batch_norm(
         # pyrefly: ignore [bad-argument-type]
         _verify_batch_size(input.size())
 
-    if eps <= 0.0:
-        raise ValueError(f"batch_norm eps must be positive, but got {eps}")
+    if training and eps <= 0.0:
+        raise ValueError(
+            f"batch_norm eps must be positive during training, but got {eps}"
+        )
+    elif eps < 0.0:
+        raise ValueError(f"batch_norm eps must be non-negative, but got {eps}")
 
     return torch.batch_norm(
         input,
@@ -2869,10 +2873,10 @@ def _verify_spatial_size(size: list[int]) -> None:
 
 def instance_norm(
     input: Tensor,
-    running_mean: Optional[Tensor] = None,
-    running_var: Optional[Tensor] = None,
-    weight: Optional[Tensor] = None,
-    bias: Optional[Tensor] = None,
+    running_mean: Tensor | None = None,
+    running_var: Tensor | None = None,
+    weight: Tensor | None = None,
+    bias: Tensor | None = None,
     use_input_stats: bool = True,
     momentum: float = 0.1,
     eps: float = 1e-5,
@@ -2914,8 +2918,8 @@ def instance_norm(
 def layer_norm(
     input: Tensor,
     normalized_shape: list[int],
-    weight: Optional[Tensor] = None,
-    bias: Optional[Tensor] = None,
+    weight: Tensor | None = None,
+    bias: Tensor | None = None,
     eps: float = 1e-5,
 ) -> Tensor:
     r"""Apply Layer Normalization for last certain number of dimensions.
@@ -2940,8 +2944,8 @@ def layer_norm(
 def rms_norm(
     input: Tensor,
     normalized_shape: list[int],
-    weight: Optional[Tensor] = None,
-    eps: Optional[float] = None,
+    weight: Tensor | None = None,
+    eps: float | None = None,
 ) -> Tensor:
     r"""Apply Root Mean Square Layer Normalization.
 
@@ -2957,8 +2961,8 @@ def rms_norm(
 def group_norm(
     input: Tensor,
     num_groups: int,
-    weight: Optional[Tensor] = None,
-    bias: Optional[Tensor] = None,
+    weight: Tensor | None = None,
+    bias: Tensor | None = None,
     eps: float = 1e-5,
 ) -> Tensor:
     r"""Apply Group Normalization for last certain number of dimensions.
@@ -3122,10 +3126,10 @@ def ctc_loss(
 def nll_loss(
     input: Tensor,
     target: Tensor,
-    weight: Optional[Tensor] = None,
-    size_average: Optional[bool] = None,
+    weight: Tensor | None = None,
+    size_average: bool | None = None,
     ignore_index: int = -100,
-    reduce: Optional[bool] = None,
+    reduce: bool | None = None,
     reduction: str = "mean",
 ) -> Tensor:
     r"""Compute the negative log likelihood loss.
@@ -3191,9 +3195,9 @@ def poisson_nll_loss(
     target: Tensor,
     log_input: bool = True,
     full: bool = False,
-    size_average: Optional[bool] = None,
+    size_average: bool | None = None,
     eps: float = 1e-8,
-    reduce: Optional[bool] = None,
+    reduce: bool | None = None,
     reduction: str = "mean",
 ) -> Tensor:
     r"""Compute the Poisson negative log likelihood loss.
@@ -3328,7 +3332,6 @@ def gaussian_nll_loss(
         var.clamp_(min=eps)
 
     # Calculate the loss
-    # pyrefly: ignore [unsupported-operation]
     loss = 0.5 * (torch.log(var) + (input - target) ** 2 / var)
     if full:
         loss += 0.5 * math.log(2 * math.pi)
@@ -3344,8 +3347,8 @@ def gaussian_nll_loss(
 def kl_div(
     input: Tensor,
     target: Tensor,
-    size_average: Optional[bool] = None,
-    reduce: Optional[bool] = None,
+    size_average: bool | None = None,
+    reduce: bool | None = None,
     reduction: str = "mean",
     log_target: bool = False,
 ) -> Tensor:
@@ -3421,10 +3424,10 @@ def kl_div(
 def cross_entropy(
     input: Tensor,
     target: Tensor,
-    weight: Optional[Tensor] = None,
-    size_average: Optional[bool] = None,
+    weight: Tensor | None = None,
+    size_average: bool | None = None,
     ignore_index: int = -100,
-    reduce: Optional[bool] = None,
+    reduce: bool | None = None,
     reduction: str = "mean",
     label_smoothing: float = 0.0,
 ) -> Tensor:
@@ -3515,9 +3518,9 @@ def cross_entropy(
 def binary_cross_entropy(
     input: Tensor,
     target: Tensor,
-    weight: Optional[Tensor] = None,
-    size_average: Optional[bool] = None,
-    reduce: Optional[bool] = None,
+    weight: Tensor | None = None,
+    size_average: bool | None = None,
+    reduce: bool | None = None,
     reduction: str = "mean",
 ) -> Tensor:
     r"""Compute Binary Cross Entropy between the target and input probabilities.
@@ -3577,11 +3580,11 @@ def binary_cross_entropy(
 def binary_cross_entropy_with_logits(
     input: Tensor,
     target: Tensor,
-    weight: Optional[Tensor] = None,
-    size_average: Optional[bool] = None,
-    reduce: Optional[bool] = None,
+    weight: Tensor | None = None,
+    size_average: bool | None = None,
+    reduce: bool | None = None,
     reduction: str = "mean",
-    pos_weight: Optional[Tensor] = None,
+    pos_weight: Tensor | None = None,
 ) -> Tensor:
     r"""Compute Binary Cross Entropy between target and input logits.
 
@@ -3647,8 +3650,8 @@ def binary_cross_entropy_with_logits(
 def smooth_l1_loss(
     input: Tensor,
     target: Tensor,
-    size_average: Optional[bool] = None,
-    reduce: Optional[bool] = None,
+    size_average: bool | None = None,
+    reduce: bool | None = None,
     reduction: str = "mean",
     beta: float = 1.0,
 ) -> Tensor:
@@ -3720,7 +3723,7 @@ def huber_loss(
     target: Tensor,
     reduction: str = "mean",
     delta: float = 1.0,
-    weight: Optional[Tensor] = None,
+    weight: Tensor | None = None,
 ) -> Tensor:
     r"""Compute the Huber loss, with optional weighting.
 
@@ -3806,10 +3809,10 @@ def huber_loss(
 def l1_loss(
     input: Tensor,
     target: Tensor,
-    size_average: Optional[bool] = None,
-    reduce: Optional[bool] = None,
+    size_average: bool | None = None,
+    reduce: bool | None = None,
     reduction: str = "mean",
-    weight: Optional[Tensor] = None,
+    weight: Tensor | None = None,
 ) -> Tensor:  # noqa: D400,D402
     r"""Compute the L1 loss, with optional weighting.
 
@@ -3882,10 +3885,10 @@ def l1_loss(
 def mse_loss(
     input: Tensor,
     target: Tensor,
-    size_average: Optional[bool] = None,
-    reduce: Optional[bool] = None,
+    size_average: bool | None = None,
+    reduce: bool | None = None,
     reduction: str = "mean",
-    weight: Optional[Tensor] = None,
+    weight: Tensor | None = None,
 ) -> Tensor:
     r"""Compute the element-wise mean squared error, with optional weighting.
 
@@ -3962,8 +3965,8 @@ def margin_ranking_loss(
     input2: Tensor,
     target: Tensor,
     margin: float = 0,
-    size_average: Optional[bool] = None,
-    reduce: Optional[bool] = None,
+    size_average: bool | None = None,
+    reduce: bool | None = None,
     reduction: str = "mean",
 ) -> Tensor:  # noqa: D400,D402
     r"""Compute the margin ranking loss.
@@ -4012,8 +4015,8 @@ def hinge_embedding_loss(
     input: Tensor,
     target: Tensor,
     margin: float = 1.0,
-    size_average: Optional[bool] = None,
-    reduce: Optional[bool] = None,
+    size_average: bool | None = None,
+    reduce: bool | None = None,
     reduction: str = "mean",
 ) -> Tensor:  # noqa: D400,D402
     r"""Compute the hinge embedding loss.
@@ -4055,8 +4058,8 @@ def hinge_embedding_loss(
 def multilabel_margin_loss(
     input: Tensor,
     target: Tensor,
-    size_average: Optional[bool] = None,
-    reduce: Optional[bool] = None,
+    size_average: bool | None = None,
+    reduce: bool | None = None,
     reduction: str = "mean",
 ) -> Tensor:  # noqa: D400,D402
     r"""Compute the multilabel margin loss.
@@ -4097,8 +4100,8 @@ def multilabel_margin_loss(
 def soft_margin_loss(
     input: Tensor,
     target: Tensor,
-    size_average: Optional[bool] = None,
-    reduce: Optional[bool] = None,
+    size_average: bool | None = None,
+    reduce: bool | None = None,
     reduction: str = "mean",
 ) -> Tensor:  # noqa: D400,D402
     r"""Compute the soft margin loss.
@@ -4139,9 +4142,9 @@ def soft_margin_loss(
 def multilabel_soft_margin_loss(
     input: Tensor,
     target: Tensor,
-    weight: Optional[Tensor] = None,
-    size_average: Optional[bool] = None,
-    reduce: Optional[bool] = None,
+    weight: Tensor | None = None,
+    size_average: bool | None = None,
+    reduce: bool | None = None,
     reduction: str = "mean",
 ) -> Tensor:  # noqa: D400,D402
     r"""Compute the multilabel soft margin loss.
@@ -4201,8 +4204,8 @@ def cosine_embedding_loss(
     input2: Tensor,
     target: Tensor,
     margin: float = 0,
-    size_average: Optional[bool] = None,
-    reduce: Optional[bool] = None,
+    size_average: bool | None = None,
+    reduce: bool | None = None,
     reduction: str = "mean",
 ) -> Tensor:  # noqa: D400,D402
     r"""Compute the cosine embedding loss.
@@ -4248,9 +4251,9 @@ def multi_margin_loss(
     target: Tensor,
     p: int = 1,
     margin: float = 1.0,
-    weight: Optional[Tensor] = None,
-    size_average: Optional[bool] = None,
-    reduce: Optional[bool] = None,
+    weight: Tensor | None = None,
+    size_average: bool | None = None,
+    reduce: bool | None = None,
     reduction: str = "mean",
 ) -> Tensor:  # noqa: D400,D402
     r"""Compute the multi margin loss, with optional weighting.
@@ -4445,10 +4448,10 @@ def multi_margin_loss(
 @_overload
 def upsample(  # noqa: F811
     input: Tensor,
-    size: Optional[int] = None,
-    scale_factor: Optional[float] = None,
+    size: int | None = None,
+    scale_factor: float | None = None,
     mode: str = "nearest",
-    align_corners: Optional[bool] = None,
+    align_corners: bool | None = None,
     # pyrefly: ignore [bad-return]
 ) -> Tensor:  # noqa: B950
     pass
@@ -4457,10 +4460,10 @@ def upsample(  # noqa: F811
 @_overload
 def upsample(  # noqa: F811
     input: Tensor,
-    size: Optional[list[int]] = None,
-    scale_factor: Optional[float] = None,
+    size: list[int] | None = None,
+    scale_factor: float | None = None,
     mode: str = "nearest",
-    align_corners: Optional[bool] = None,
+    align_corners: bool | None = None,
     # pyrefly: ignore [bad-return]
 ) -> Tensor:  # noqa: B950
     pass
@@ -4494,7 +4497,7 @@ def upsample(  # noqa: F811
     `mini-batch x channels x [optional depth] x [optional height] x width`.
 
     The modes available for upsampling are: `nearest`, `linear` (3D-only),
-    `bilinear`, `bicubic` (4D-only), `trilinear` (5D-only)
+    `bilinear`, `bicubic` (4D-only), `trilinear` (5D-only), `lanczos` (4D-only)
 
     Args:
         input (Tensor): the input tensor
@@ -4503,7 +4506,7 @@ def upsample(  # noqa: F811
         scale_factor (float or Tuple[float]): multiplier for spatial size. Has to match input size if it is a tuple.
         mode (str): algorithm used for upsampling:
             ``'nearest'`` | ``'linear'`` | ``'bilinear'`` | ``'bicubic'`` |
-            ``'trilinear'``. Default: ``'nearest'``
+            ``'trilinear'`` | ``'lanczos'``. Default: ``'nearest'``
         align_corners (bool, optional): Geometrically, we consider the pixels of the
             input and output as squares rather than points.
             If set to ``True``, the input and output tensors are aligned by the
@@ -4516,7 +4519,7 @@ def upsample(  # noqa: F811
             Default: ``False``
 
     .. note::
-        With ``mode='bicubic'``, it's possible to cause overshoot, in other words it can produce
+        With ``mode='bicubic'`` or ``mode='lanczos'``, it's possible to cause overshoot, in other words it can produce
         negative values or values greater than 255 for images.
         Explicitly call ``result.clamp(min=0, max=255)`` if you want to reduce the overshoot
         when displaying the image.
@@ -4558,11 +4561,11 @@ def _is_integer(x) -> bool:
 @_overload
 def interpolate(  # noqa: F811
     input: Tensor,
-    size: Optional[int] = None,
-    scale_factor: Optional[list[float]] = None,
+    size: int | None = None,
+    scale_factor: list[float] | None = None,
     mode: str = "nearest",
-    align_corners: Optional[bool] = None,
-    recompute_scale_factor: Optional[bool] = None,
+    align_corners: bool | None = None,
+    recompute_scale_factor: bool | None = None,
     antialias: bool = False,
     # pyrefly: ignore [bad-return]
 ) -> Tensor:  # noqa: B950
@@ -4572,11 +4575,11 @@ def interpolate(  # noqa: F811
 @_overload
 def interpolate(  # noqa: F811
     input: Tensor,
-    size: Optional[list[int]] = None,
-    scale_factor: Optional[list[float]] = None,
+    size: list[int] | None = None,
+    scale_factor: list[float] | None = None,
     mode: str = "nearest",
-    align_corners: Optional[bool] = None,
-    recompute_scale_factor: Optional[bool] = None,
+    align_corners: bool | None = None,
+    recompute_scale_factor: bool | None = None,
     antialias: bool = False,
     # pyrefly: ignore [bad-return]
 ) -> Tensor:  # noqa: B950
@@ -4586,11 +4589,11 @@ def interpolate(  # noqa: F811
 @_overload
 def interpolate(  # noqa: F811
     input: Tensor,
-    size: Optional[int] = None,
-    scale_factor: Optional[float] = None,
+    size: int | None = None,
+    scale_factor: float | None = None,
     mode: str = "nearest",
-    align_corners: Optional[bool] = None,
-    recompute_scale_factor: Optional[bool] = None,
+    align_corners: bool | None = None,
+    recompute_scale_factor: bool | None = None,
     antialias: bool = False,
     # pyrefly: ignore [bad-return]
 ) -> Tensor:  # noqa: B950
@@ -4600,11 +4603,11 @@ def interpolate(  # noqa: F811
 @_overload
 def interpolate(  # noqa: F811
     input: Tensor,
-    size: Optional[list[int]] = None,
-    scale_factor: Optional[float] = None,
+    size: list[int] | None = None,
+    scale_factor: float | None = None,
     mode: str = "nearest",
-    align_corners: Optional[bool] = None,
-    recompute_scale_factor: Optional[bool] = None,
+    align_corners: bool | None = None,
+    recompute_scale_factor: bool | None = None,
     antialias: bool = False,
     # pyrefly: ignore [bad-return]
 ) -> Tensor:
@@ -4613,11 +4616,11 @@ def interpolate(  # noqa: F811
 
 def interpolate(  # noqa: F811
     input: Tensor,
-    size: Optional[int] = None,
-    scale_factor: Optional[list[float]] = None,
+    size: int | None = None,
+    scale_factor: list[float] | None = None,
     mode: str = "nearest",
-    align_corners: Optional[bool] = None,
-    recompute_scale_factor: Optional[bool] = None,
+    align_corners: bool | None = None,
+    recompute_scale_factor: bool | None = None,
     antialias: bool = False,
 ) -> Tensor:  # noqa: B950
     r"""Down/up samples the input.
@@ -4634,7 +4637,7 @@ def interpolate(  # noqa: F811
     `mini-batch x channels x [optional depth] x [optional height] x width`.
 
     The modes available for resizing are: `nearest`, `linear` (3D-only),
-    `bilinear`, `bicubic` (4D-only), `trilinear` (5D-only), `area`, `nearest-exact`
+    `bilinear`, `bicubic` (4D-only), `trilinear` (5D-only), `lanczos` (4D-only, CPU only), `area`, `nearest-exact`
 
     Args:
         input (Tensor): the input tensor
@@ -4644,7 +4647,7 @@ def interpolate(  # noqa: F811
             its length has to match the number of spatial dimensions; `input.dim() - 2`.
         mode (str): algorithm used for upsampling:
             ``'nearest'`` | ``'linear'`` | ``'bilinear'`` | ``'bicubic'`` |
-            ``'trilinear'`` | ``'area'`` | ``'nearest-exact'``. Default: ``'nearest'``
+            ``'trilinear'`` | ``'lanczos'`` | ``'area'`` | ``'nearest-exact'``. Default: ``'nearest'``
         align_corners (bool, optional): Geometrically, we consider the pixels of the
             input and output as squares rather than points.
             If set to ``True``, the input and output tensors are aligned by the
@@ -4666,14 +4669,19 @@ def interpolate(  # noqa: F811
             be used directly for interpolation. Default: ``None``.
         antialias (bool, optional): flag to apply anti-aliasing. Default: ``False``. Using anti-alias
             option together with ``align_corners=False``, interpolation result would match Pillow
-            result for downsampling operation. Supported modes: ``'bilinear'``, ``'bicubic'``.
+            result for downsampling operation. Supported modes: ``'bilinear'``, ``'bicubic'``, ``'lanczos'``.
 
     .. note::
-        With ``mode='bicubic'``, it's possible to cause overshoot. For some dtypes, it can produce
+        With ``mode='bicubic'`` or ``mode='lanczos'``, it's possible to cause overshoot. For some dtypes, it can produce
         negative values or values greater than 255 for images. Explicitly call ``result.clamp(min=0,max=255)``
         if you want to reduce the overshoot when displaying the image.
         For ``uint8`` inputs, it already performs saturating cast operation. So, no manual `clamp` operation is needed.
 
+    .. note::
+        Mode ``mode='lanczos'`` uses a Lanczos-3 windowed sinc filter (6 taps) and requires
+        ``antialias=True``. It only supports 4-D input (i.e. 2D spatial) and CPU. With ``antialias=True``
+        and ``align_corners=False``, the result matches PIL's ``Image.LANCZOS`` resampling filter.
+
     .. note::
         Mode ``mode='nearest-exact'`` matches Scikit-Image and PIL nearest neighbours interpolation
         algorithms and fixes known issues with ``mode='nearest'``. This mode is introduced to keep
@@ -4804,9 +4812,11 @@ def interpolate(  # noqa: F811
             ]
         scale_factors = None
 
-    if antialias and not (mode in ("bilinear", "bicubic") and input.ndim == 4):
+    if antialias and not (
+        mode in ("bilinear", "bicubic", "lanczos") and input.ndim == 4
+    ):
         raise ValueError(
-            "Anti-alias option is restricted to bilinear and bicubic modes and requires a 4-D tensor as input"
+            "Anti-alias option is restricted to bilinear, bicubic, and lanczos modes and requires a 4-D tensor as input"
         )
 
     if input.dim() == 3 and mode == "nearest":
@@ -4920,6 +4930,21 @@ def interpolate(  # noqa: F811
             scale_factors,
         )
 
+    if input.dim() == 4 and mode == "lanczos":
+        if align_corners is None:
+            raise AssertionError("align_corners is unexpectedly None")
+        if align_corners:
+            raise ValueError("Lanczos mode does not support align_corners=True")
+        if not antialias:
+            raise ValueError("Lanczos mode requires antialias=True")
+        return torch._C._nn._upsample_lanczos2d_aa(
+            input,
+            # pyrefly: ignore [bad-argument-type]
+            output_size,
+            align_corners,
+            scale_factors,
+        )
+
     if input.dim() == 3 and mode == "bilinear":
         raise NotImplementedError("Got 3D input, but bilinear mode needs 4D input")
     if input.dim() == 3 and mode == "trilinear":
@@ -4935,7 +4960,7 @@ def interpolate(  # noqa: F811
 
     raise NotImplementedError(
         "Input Error: Only 3D, 4D and 5D input Tensors supported"
-        f" (got {input.dim()}D) for the modes: nearest | linear | bilinear | bicubic | trilinear | area | nearest-exact"
+        f" (got {input.dim()}D) for the modes: nearest | linear | bilinear | bicubic | trilinear | lanczos | area | nearest-exact"
         f" (got {mode})"
     )
 
@@ -4947,8 +4972,8 @@ def interpolate(  # noqa: F811
 @_overload
 def upsample_nearest(  # noqa: F811
     input: Tensor,
-    size: Optional[int] = None,
-    scale_factor: Optional[float] = None,
+    size: int | None = None,
+    scale_factor: float | None = None,
     # pyrefly: ignore [bad-return]
 ) -> Tensor:
     pass
@@ -4957,8 +4982,8 @@ def upsample_nearest(  # noqa: F811
 @_overload
 def upsample_nearest(  # noqa: F811
     input: Tensor,
-    size: Optional[list[int]] = None,
-    scale_factor: Optional[float] = None,
+    size: list[int] | None = None,
+    scale_factor: float | None = None,
     # pyrefly: ignore [bad-return]
 ) -> Tensor:
     pass
@@ -4999,8 +5024,8 @@ def upsample_nearest(input, size=None, scale_factor=None):  # noqa: F811
 @_overload
 def upsample_bilinear(  # noqa: F811
     input: Tensor,
-    size: Optional[int] = None,
-    scale_factor: Optional[float] = None,
+    size: int | None = None,
+    scale_factor: float | None = None,
     # pyrefly: ignore [bad-return]
 ) -> Tensor:
     pass
@@ -5009,8 +5034,8 @@ def upsample_bilinear(  # noqa: F811
 @_overload
 def upsample_bilinear(  # noqa: F811
     input: Tensor,
-    size: Optional[list[int]] = None,
-    scale_factor: Optional[float] = None,
+    size: list[int] | None = None,
+    scale_factor: float | None = None,
     # pyrefly: ignore [bad-return]
 ) -> Tensor:
     pass
@@ -5019,8 +5044,8 @@ def upsample_bilinear(  # noqa: F811
 @_overload
 def upsample_bilinear(  # noqa: F811
     input: Tensor,
-    size: Optional[int] = None,
-    scale_factor: Optional[list[float]] = None,
+    size: int | None = None,
+    scale_factor: list[float] | None = None,
     # pyrefly: ignore [bad-return]
 ) -> Tensor:
     pass
@@ -5029,8 +5054,8 @@ def upsample_bilinear(  # noqa: F811
 @_overload
 def upsample_bilinear(  # noqa: F811
     input: Tensor,
-    size: Optional[list[int]] = None,
-    scale_factor: Optional[list[float]] = None,
+    size: list[int] | None = None,
+    scale_factor: list[float] | None = None,
     # pyrefly: ignore [bad-return]
 ) -> Tensor:
     pass
@@ -5087,7 +5112,7 @@ def grid_sample(
     grid: Tensor,
     mode: str = "bilinear",
     padding_mode: str = "zeros",
-    align_corners: Optional[bool] = None,
+    align_corners: bool | None = None,
 ) -> Tensor:
     r"""Compute grid sample.
 
@@ -5245,7 +5270,7 @@ def grid_sample(
 def affine_grid(
     theta: Tensor,
     size: list[int],
-    align_corners: Optional[bool] = None,
+    align_corners: bool | None = None,
 ) -> Tensor:
     r"""Generate 2D or 3D flow field (sampling grid), given a batch of affine matrices :attr:`theta`.
 
@@ -5350,7 +5375,7 @@ def pad(
     input: Tensor,
     pad: list[int],
     mode: str = "constant",
-    value: Optional[float] = None,
+    value: float | None = None,
 ) -> Tensor:
     r"""
     pad(input, pad, mode="constant", value=None) -> Tensor
@@ -5422,6 +5447,13 @@ def pad(
         ):
             if mode == "replicate":
                 # Use slow decomp whose backward will be in terms of index_put.
+                if torch.compiler.is_compiling():
+                    # nonstrict_trace makes Dynamo skip the function body
+                    # (which contains Dynamo-untraceable code) while
+                    # AOTAutograd still traces into it for the backward.
+                    return torch._dynamo.decorators.nonstrict_trace(
+                        torch._decomp.decompositions._replication_pad
+                    )(input, pad)
                 # importlib is required because the import cannot be top level
                 # (cycle) and cannot be nested (TS doesn't support)
                 return importlib.import_module(
@@ -5563,8 +5595,8 @@ def triplet_margin_loss(
     p: float = 2,
     eps: float = 1e-6,
     swap: bool = False,
-    size_average: Optional[bool] = None,
-    reduce: Optional[bool] = None,
+    size_average: bool | None = None,
+    reduce: bool | None = None,
     reduction: str = "mean",
 ) -> Tensor:
     r"""Compute the triplet loss between given input tensors and a margin greater than 0.
@@ -5602,7 +5634,7 @@ def triplet_margin_with_distance_loss(
     positive: Tensor,
     negative: Tensor,
     *,
-    distance_function: Optional[Callable[[Tensor, Tensor], Tensor]] = None,
+    distance_function: Callable[[Tensor, Tensor], Tensor] | None = None,
     margin: float = 1.0,
     swap: bool = False,
     reduction: str = "mean",
@@ -5679,7 +5711,7 @@ def normalize(
     p: float = 2.0,
     dim: int = 1,
     eps: float = 1e-12,
-    out: Optional[Tensor] = None,
+    out: Tensor | None = None,
 ) -> Tensor:
     r"""Perform :math:`L_p` normalization of inputs over specified dimension.
 
@@ -5800,7 +5832,7 @@ def _in_projection_packed(
     k: Tensor,
     v: Tensor,
     w: Tensor,
-    b: Optional[Tensor] = None,
+    b: Tensor | None = None,
 ) -> list[Tensor]:
     r"""Perform the in-projection step of the attention operation, using packed weights.
 
@@ -5880,9 +5912,9 @@ def _in_projection(
     w_q: Tensor,
     w_k: Tensor,
     w_v: Tensor,
-    b_q: Optional[Tensor] = None,
-    b_k: Optional[Tensor] = None,
-    b_v: Optional[Tensor] = None,
+    b_q: Tensor | None = None,
+    b_k: Tensor | None = None,
+    b_v: Tensor | None = None,
 ) -> tuple[Tensor, Tensor, Tensor]:
     r"""Perform the in-projection step of the attention operation.
 
@@ -5964,7 +5996,7 @@ def scaled_dot_product_attention(query, key, value, attn_mask=None, dropout_p=0.
             attn_bias = torch.zeros(L, S, dtype=query.dtype, device=query.device)
             if is_causal:
                 assert attn_mask is None
-                temp_mask = torch.ones(L, S, dtype=torch.bool).tril(diagonal=0)
+                temp_mask = torch.ones(L, S, dtype=torch.bool, device=query.device).tril(diagonal=0)
                 attn_bias.masked_fill_(temp_mask.logical_not(), float("-inf"))
 
             if attn_mask is not None:
@@ -6068,7 +6100,7 @@ def forward(self, ...):
         key (Tensor): Key tensor; shape :math:`(N, ..., H, S, E)`.
         value (Tensor): Value tensor; shape :math:`(N, ..., H, S, Ev)`.
         attn_mask (optional Tensor): Attention mask; shape must be broadcastable to the shape of attention weights,
-            which is :math:`(N,..., L, S)`. Two types of masks are supported.
+            which is :math:`(N,..., Hq, L, S)`. Two types of masks are supported.
             A boolean mask where a value of True indicates that the element *should* take part in attention.
             A float mask of the same type as query, key, value that is added to the attention score.
         dropout_p (float): Dropout probability; if greater than 0.0, dropout is applied
@@ -6124,8 +6156,8 @@ def _mha_shape_check(
     query: Tensor,
     key: Tensor,
     value: Tensor,
-    key_padding_mask: Optional[Tensor],
-    attn_mask: Optional[Tensor],
+    key_padding_mask: Tensor | None,
+    attn_mask: Tensor | None,
     num_heads: int,
 ):
     # Verifies the expected shape for `query, `key`, `value`, `key_padding_mask` and `attn_mask`
@@ -6190,13 +6222,13 @@ def _mha_shape_check(
 
 
 def _canonical_mask(
-    mask: Optional[Tensor],
+    mask: Tensor | None,
     mask_name: str,
-    other_type: Optional[DType],
+    other_type: DType | None,
     other_name: str,
     target_type: DType,
     check_other: bool = True,
-) -> Optional[Tensor]:
+) -> Tensor | None:
     if mask is not None:
         _mask_dtype = mask.dtype
         _mask_is_float = torch.is_floating_point(mask)
@@ -6218,7 +6250,7 @@ def _canonical_mask(
     return mask
 
 
-def _none_or_dtype(input: Optional[Tensor]) -> Optional[DType]:
+def _none_or_dtype(input: Tensor | None) -> DType | None:
     if input is None:
         return None
     elif isinstance(input, torch.Tensor):
@@ -6247,27 +6279,27 @@ def multi_head_attention_forward(
     value: Tensor,
     embed_dim_to_check: int,
     num_heads: int,
-    in_proj_weight: Optional[Tensor],
-    in_proj_bias: Optional[Tensor],
-    bias_k: Optional[Tensor],
-    bias_v: Optional[Tensor],
+    in_proj_weight: Tensor | None,
+    in_proj_bias: Tensor | None,
+    bias_k: Tensor | None,
+    bias_v: Tensor | None,
     add_zero_attn: bool,
     dropout_p: float,
     out_proj_weight: Tensor,
-    out_proj_bias: Optional[Tensor],
+    out_proj_bias: Tensor | None,
     training: bool = True,
-    key_padding_mask: Optional[Tensor] = None,
+    key_padding_mask: Tensor | None = None,
     need_weights: bool = True,
-    attn_mask: Optional[Tensor] = None,
+    attn_mask: Tensor | None = None,
     use_separate_proj_weight: bool = False,
-    q_proj_weight: Optional[Tensor] = None,
-    k_proj_weight: Optional[Tensor] = None,
-    v_proj_weight: Optional[Tensor] = None,
-    static_k: Optional[Tensor] = None,
-    static_v: Optional[Tensor] = None,
+    q_proj_weight: Tensor | None = None,
+    k_proj_weight: Tensor | None = None,
+    v_proj_weight: Tensor | None = None,
+    static_k: Tensor | None = None,
+    static_v: Tensor | None = None,
     average_attn_weights: bool = True,
     is_causal: bool = False,
-) -> tuple[Tensor, Optional[Tensor]]:
+) -> tuple[Tensor, Tensor | None]:
     r"""Forward method for MultiHeadAttention.
 
     See :class:`torch.nn.MultiheadAttention` for details.
@@ -6550,10 +6582,10 @@ def multi_head_attention_forward(
     #
     # reshape q, k, v for multihead attention and make them batch first
     #
-    # pyrefly: ignore [no-matching-overload]
+    # pyrefly: ignore [bad-argument-type, no-matching-overload]
     q = q.view(tgt_len, bsz * num_heads, head_dim).transpose(0, 1)
     if static_k is None:
-        # pyrefly: ignore [no-matching-overload]
+        # pyrefly: ignore [bad-argument-type, no-matching-overload]
         k = k.view(k.shape[0], bsz * num_heads, head_dim).transpose(0, 1)
     else:
         # TODO finish disentangling control flow so we don't do in-projections when statics are passed
@@ -6567,7 +6599,7 @@ def multi_head_attention_forward(
             )
         k = static_k
     if static_v is None:
-        # pyrefly: ignore [no-matching-overload]
+        # pyrefly: ignore [bad-argument-type, no-matching-overload]
         v = v.view(v.shape[0], bsz * num_heads, head_dim).transpose(0, 1)
     else:
         # TODO finish disentangling control flow so we don't do in-projections when statics are passed
@@ -6640,14 +6672,17 @@ def multi_head_attention_forward(
             )
         else:
             attn_output_weights = torch.bmm(q_scaled, k.transpose(-2, -1))
+        if not torch.jit.is_scripting():
+            del q_scaled, k
         attn_output_weights = softmax(attn_output_weights, dim=-1)
         if dropout_p > 0.0:
             attn_output_weights = dropout(attn_output_weights, p=dropout_p)
 
         attn_output = torch.bmm(attn_output_weights, v)
+        if not torch.jit.is_scripting():
+            del v
 
         attn_output = (
-            # pyrefly: ignore [no-matching-overload]
             attn_output.transpose(0, 1).contiguous().view(tgt_len * bsz, embed_dim)
         )
         attn_output = linear(attn_output, out_proj_weight, out_proj_bias)
@@ -6673,17 +6708,23 @@ def multi_head_attention_forward(
             else:
                 attn_mask = attn_mask.view(bsz, num_heads, -1, src_len)
 
+        # pyrefly: ignore [bad-argument-type]
         q = q.view(bsz, num_heads, tgt_len, head_dim)
-        # pyrefly: ignore [no-matching-overload]
+        # pyrefly: ignore [bad-argument-type]
         k = k.view(bsz, num_heads, src_len, head_dim)
-        # pyrefly: ignore [no-matching-overload]
+        # pyrefly: ignore [bad-argument-type]
         v = v.view(bsz, num_heads, src_len, head_dim)
 
         attn_output = scaled_dot_product_attention(
             q, k, v, attn_mask, dropout_p, is_causal
         )
+        # Free q, k, v and their backing projection storage before the
+        # .contiguous() call below allocates.  In self-attention the three
+        # tensors are views of a single packed projection, so releasing all
+        # references here lets the allocator reclaim that memory immediately.
+        if not torch.jit.is_scripting():
+            del q, k, v
         attn_output = (
-            # pyrefly: ignore [no-matching-overload]
             attn_output.permute(2, 0, 1, 3).contiguous().view(bsz * tgt_len, embed_dim)
         )
 
@@ -6699,9 +6740,9 @@ def grouped_mm(
     mat_a: Tensor,
     mat_b: Tensor,
     *,
-    offs: Optional[Tensor] = None,
-    bias: Optional[Tensor] = None,
-    out_dtype: Optional[torch.dtype] = None,
+    offs: Tensor | None = None,
+    bias: Tensor | None = None,
+    out_dtype: torch.dtype | None = None,
 ) -> Tensor:
     r"""
     grouped_mm(mat_a, mat_b, *, offs=None, bias=None, out_dtype=None)
@@ -6750,8 +6791,8 @@ def scaled_mm(
     scale_recipe_b: ScalingType | list[ScalingType],
     swizzle_a: SwizzleType | list[SwizzleType] | None = None,
     swizzle_b: SwizzleType | list[SwizzleType] | None = None,
-    bias: Optional[Tensor] = None,
-    output_dtype: Optional[torch.dtype] = torch.bfloat16,
+    bias: Tensor | None = None,
+    output_dtype: torch.dtype | None = torch.bfloat16,
     contraction_dim: list[int] | tuple[int, ...] = (),
     use_fast_accum: bool = False,
 ) -> Tensor:
@@ -6834,9 +6875,9 @@ def scaled_grouped_mm(
     scale_recipe_b: ScalingType | list[ScalingType],
     swizzle_a: SwizzleType | list[SwizzleType] | None = None,
     swizzle_b: SwizzleType | list[SwizzleType] | None = None,
-    bias: Optional[Tensor] = None,
-    offs: Optional[Tensor] = None,
-    output_dtype: Optional[torch.dtype] = torch.bfloat16,
+    bias: Tensor | None = None,
+    offs: Tensor | None = None,
+    output_dtype: torch.dtype | None = torch.bfloat16,
     contraction_dim: list[int] | tuple[int, ...] = (),
     use_fast_accum: bool = False,
 ) -> Tensor:
diff --git a/torch/nn/init.py b/torch/nn/init.py
index 900b2d34bc08f..982845a1cfec0 100644
--- a/torch/nn/init.py
+++ b/torch/nn/init.py
@@ -91,9 +91,11 @@ def _no_grad_trunc_normal_(
     b: float,
     generator: torch.Generator | None = None,
 ) -> Tensor:
-    # Method based on https://people.sc.fsu.edu/~jburkardt/presentations/truncated_normal.pdf
+    # Meta tensors have no storage, so sampling is a no-op.
+    if tensor.is_meta:
+        return tensor
+
     def norm_cdf(x: float) -> float:
-        # Computes standard normal cumulative distribution function
         return (1.0 + math.erf(x / math.sqrt(2.0))) / 2.0
 
     if (mean < a - 2 * std) or (mean > b + 2 * std):
@@ -104,26 +106,57 @@ def norm_cdf(x: float) -> float:
         )
 
     with torch.no_grad():
-        # Values are generated by using a truncated uniform distribution and
-        # then using the inverse CDF for the normal distribution.
-        # Get upper and lower cdf values
-        l = norm_cdf((a - mean) / std)
-        u = norm_cdf((b - mean) / std)
-
-        # Uniformly fill tensor with values from [l, u], then translate to
-        # [2l-1, 2u-1].
-        tensor.uniform_(2 * l - 1, 2 * u - 1, generator=generator)
-
-        # Use inverse cdf transform for normal distribution to get truncated
-        # standard normal
-        tensor.erfinv_()
-
-        # Transform to proper mean, std
-        tensor.mul_(std * math.sqrt(2.0))
-        tensor.add_(mean)
-
-        # Clamp to ensure it's in the proper range
-        tensor.clamp_(min=a, max=b)
+        p = norm_cdf((b - mean) / std) - norm_cdf((a - mean) / std)
+
+        if p > 0.3:
+            # Cast bounds to tensor dtype so the rejection mask is consistent
+            # with the tensor's representable values.
+            lo = tensor.new_tensor(a, device="cpu").item()
+            hi = tensor.new_tensor(b, device="cpu").item()
+            result = tensor.normal_(mean, std, generator=generator)
+            while True:
+                mask = (result < lo) | (result > hi)
+                if not mask.any():
+                    break
+                result = torch.where(
+                    mask,
+                    torch.empty_like(result).normal_(mean, std, generator=generator),
+                    result,
+                )
+            if tensor is not result:
+                tensor.copy_(result)
+        else:
+            mode = max(a, min(mean, b))
+            log_peak = -0.5 * ((mode - mean) / std) ** 2
+
+            candidates = torch.empty_like(tensor)
+            accept_buf = torch.empty_like(tensor)
+
+            # First iteration: sample directly into tensor to avoid
+            # a where() + copy_() if all samples are accepted.
+            tensor.uniform_(a, b, generator=generator)
+            candidates.copy_(tensor)
+            # log_pdf = -0.5 * ((candidates - mean) / std) ** 2
+            candidates.sub_(mean).div_(std).pow_(2).mul_(-0.5).sub_(log_peak)
+            pending = accept_buf.uniform_(generator=generator).log_().gt(candidates)
+            if not pending.any():
+                pass
+            else:
+                result = tensor
+                while True:
+                    candidates.uniform_(a, b, generator=generator)
+                    result = torch.where(pending, candidates, result)
+                    # log_pdf = -0.5 * ((candidates - mean) / std) ** 2
+                    candidates.sub_(mean).div_(std).pow_(2).mul_(-0.5).sub_(log_peak)
+                    pending = torch.where(
+                        pending,
+                        accept_buf.uniform_(generator=generator).log_().gt(candidates),
+                        pending,
+                    )
+                    if not pending.any():
+                        break
+                tensor.copy_(result)
+
         return tensor
 
 
@@ -281,6 +314,11 @@ def trunc_normal_(
     the bounds. The method used for generating the random values works
     best when :math:`a \leq \text{mean} \leq b`.
 
+    For reduced-precision types (``torch.float16`` and ``torch.bfloat16``),
+    sampling quality depends on the underlying ``normal_()`` and ``uniform_()``
+    implementations which operate at higher internal precision to avoid
+    quantization artifacts.
+
     Args:
         tensor: an n-dimensional `torch.Tensor`
         mean: the mean of the normal distribution
@@ -386,6 +424,9 @@ def dirac_(tensor: Tensor, groups: int = 1) -> Tensor:
     if sizes[0] % groups != 0:
         raise ValueError("dim 0 must be divisible by groups")
 
+    if tensor.is_meta:
+        return tensor
+
     out_chans_per_grp = sizes[0] // groups
     min_dim = min(out_chans_per_grp, sizes[1])
 
@@ -653,7 +694,7 @@ def orthogonal_(
     if tensor.ndimension() < 2:
         raise ValueError("Only tensors with 2 or more dimensions are supported")
 
-    if tensor.numel() == 0:
+    if tensor.numel() == 0 or tensor.is_meta:
         # no-op
         return tensor
     rows = tensor.size(0)
@@ -705,6 +746,9 @@ def sparse_(
     if tensor.ndimension() != 2:
         raise ValueError("Only tensors with 2 dimensions are supported")
 
+    if tensor.is_meta:
+        return tensor
+
     rows, cols = tensor.shape
     num_zeros = math.ceil(sparsity * rows)
 
diff --git a/torch/nn/modules/_functions.py b/torch/nn/modules/_functions.py
index 408e6ef42f128..59352264b3216 100644
--- a/torch/nn/modules/_functions.py
+++ b/torch/nn/modules/_functions.py
@@ -124,6 +124,7 @@ def forward(
             return torch.empty_like(input)
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def backward(self, grad_output):
         if not (
             grad_output.is_contiguous(memory_format=torch.channels_last)
diff --git a/torch/nn/modules/activation.py b/torch/nn/modules/activation.py
index 658dd0a4d4fa3..739b2a95fc3c1 100644
--- a/torch/nn/modules/activation.py
+++ b/torch/nn/modules/activation.py
@@ -1394,6 +1394,7 @@ def forward(
         elif torch.is_autocast_enabled():
             why_not_fast_path = "autocast is enabled"
 
+        fast_path_blocked_by_tracing = False
         if not why_not_fast_path:
             tensor_args = (
                 query,
@@ -1408,8 +1409,6 @@ def forward(
             # generator expressions.
             if torch.overrides.has_torch_function(tensor_args):
                 why_not_fast_path = "some Tensor argument has_torch_function"
-            elif _is_make_fx_tracing():
-                why_not_fast_path = "we are running make_fx tracing"
             elif not all(_check_arg_device(x) for x in tensor_args):
                 why_not_fast_path = (
                     "some Tensor argument's device is neither one of "
@@ -1422,6 +1421,9 @@ def forward(
                     "grad is enabled and at least one of query or the "
                     "input/output projection weights or biases requires_grad"
                 )
+            elif _is_make_fx_tracing():
+                why_not_fast_path = "we are running make_fx tracing"
+                fast_path_blocked_by_tracing = True
             if not why_not_fast_path:
                 merged_mask, mask_type = self.merge_masks(
                     attn_mask, key_padding_mask, query
@@ -1511,7 +1513,11 @@ def forward(
                 is_causal=is_causal,
             )
         if self.batch_first and is_batched:
-            return attn_output.transpose(1, 0), attn_output_weights
+            attn_output = attn_output.transpose(1, 0)
+            if fast_path_blocked_by_tracing:
+                # Keep the traced slowpath layout aligned with eager fastpath.
+                attn_output = attn_output.contiguous()
+            return attn_output, attn_output_weights
         else:
             return attn_output, attn_output_weights
 
diff --git a/torch/nn/modules/batchnorm.py b/torch/nn/modules/batchnorm.py
index 91e8de79855a3..e56a747e14190 100644
--- a/torch/nn/modules/batchnorm.py
+++ b/torch/nn/modules/batchnorm.py
@@ -44,6 +44,8 @@ def __init__(
         track_running_stats: bool = True,
         device=None,
         dtype=None,
+        *,
+        bias: bool = True,
     ) -> None:
         factory_kwargs = {"device": device, "dtype": dtype}
         super().__init__()
@@ -54,7 +56,10 @@ def __init__(
         self.track_running_stats = track_running_stats
         if self.affine:
             self.weight = Parameter(torch.empty(num_features, **factory_kwargs))
-            self.bias = Parameter(torch.empty(num_features, **factory_kwargs))
+            if bias:
+                self.bias = Parameter(torch.empty(num_features, **factory_kwargs))
+            else:
+                self.register_parameter("bias", None)
         else:
             self.register_parameter("weight", None)
             self.register_parameter("bias", None)
@@ -95,7 +100,8 @@ def reset_parameters(self) -> None:
         self.reset_running_stats()
         if self.affine:
             init.ones_(self.weight)
-            init.zeros_(self.bias)
+            if self.bias is not None:
+                init.zeros_(self.bias)
 
     def _check_input_dim(self, input):
         raise NotImplementedError
@@ -103,7 +109,9 @@ def _check_input_dim(self, input):
     def extra_repr(self):
         return (
             "{num_features}, eps={eps}, momentum={momentum}, affine={affine}, "
-            "track_running_stats={track_running_stats}".format(**self.__dict__)
+            "bias={use_bias}, track_running_stats={track_running_stats}".format(
+                **self.__dict__, use_bias=self.bias is not None
+            )
         )
 
     def _load_from_state_dict(
@@ -151,10 +159,18 @@ def __init__(
         track_running_stats: bool = True,
         device=None,
         dtype=None,
+        *,
+        bias: bool = True,
     ) -> None:
         factory_kwargs = {"device": device, "dtype": dtype}
         super().__init__(
-            num_features, eps, momentum, affine, track_running_stats, **factory_kwargs
+            num_features,
+            eps,
+            momentum,
+            affine,
+            track_running_stats,
+            **factory_kwargs,
+            bias=bias,
         )
 
     def forward(self, input: Tensor) -> Tensor:
@@ -220,11 +236,13 @@ def __init__(
         track_running_stats=True,
         device=None,
         dtype=None,
+        *,
+        bias=True,
     ) -> None:
         factory_kwargs = {"device": device, "dtype": dtype}
         # pyrefly: ignore [bad-argument-type]
         super().__init__(
-            # affine and track_running_stats are hardcoded to False to
+            # affine, bias and track_running_stats are hardcoded to False to
             # avoid creating tensors that will soon be overwritten.
             0,
             eps,
@@ -232,14 +250,16 @@ def __init__(
             False,
             False,
             **factory_kwargs,
+            bias=False,
         )
         self.affine = affine
         self.track_running_stats = track_running_stats
         if self.affine:
             # pyrefly: ignore [unexpected-keyword]
             self.weight = UninitializedParameter(**factory_kwargs)
-            # pyrefly: ignore [unexpected-keyword]
-            self.bias = UninitializedParameter(**factory_kwargs)
+            if bias:
+                # pyrefly: ignore  # bad-argument-type
+                self.bias = UninitializedParameter(**factory_kwargs)
         if self.track_running_stats:
             # pyrefly: ignore [unexpected-keyword]
             self.running_mean = UninitializedBuffer(**factory_kwargs)
@@ -266,10 +286,13 @@ def initialize_parameters(self, input) -> None:  # type: ignore[override]
                     raise AssertionError(
                         "self.weight must be an UninitializedParameter"
                     )
-                if not isinstance(self.bias, UninitializedParameter):
-                    raise AssertionError("self.bias must be an UninitializedParameter")
                 self.weight.materialize((self.num_features,))
-                self.bias.materialize((self.num_features,))
+                if self.bias is not None:
+                    if not isinstance(self.bias, UninitializedParameter):
+                        raise AssertionError(
+                            "self.bias must be an UninitializedParameter"
+                        )
+                    self.bias.materialize((self.num_features,))
             if self.track_running_stats:
                 self.running_mean.materialize(  # type:ignore[union-attr]
                     (self.num_features,)
@@ -335,6 +358,8 @@ class BatchNorm1d(_BatchNorm):
             buffers :attr:`running_mean` and :attr:`running_var` as ``None``.
             When these buffers are ``None``, this module always uses batch statistics.
             in both training and eval modes. Default: ``True``
+        bias: If set to ``False``, the layer will not learn an additive bias (only relevant if
+            :attr:`affine` is ``True``). Default: ``True``
 
     Shape:
         - Input: :math:`(N, C)` or :math:`(N, C, L)`, where :math:`N` is the batch size,
@@ -381,6 +406,8 @@ class LazyBatchNorm1d(_LazyNormBase, _BatchNorm):
             buffers :attr:`running_mean` and :attr:`running_var` as ``None``.
             When these buffers are ``None``, this module always uses batch statistics.
             in both training and eval modes. Default: ``True``
+        bias: If set to ``False``, the layer will not learn an additive bias (only relevant if
+            :attr:`affine` is ``True``). Default: ``True``
     """
 
     cls_to_become = BatchNorm1d  # type: ignore[assignment]
@@ -447,6 +474,8 @@ class BatchNorm2d(_BatchNorm):
             buffers :attr:`running_mean` and :attr:`running_var` as ``None``.
             When these buffers are ``None``, this module always uses batch statistics.
             in both training and eval modes. Default: ``True``
+        bias: If set to ``False``, the layer will not learn an additive bias (only relevant if
+            :attr:`affine` is ``True``). Default: ``True``
 
     Shape:
         - Input: :math:`(N, C, H, W)`
@@ -492,6 +521,8 @@ class LazyBatchNorm2d(_LazyNormBase, _BatchNorm):
             buffers :attr:`running_mean` and :attr:`running_var` as ``None``.
             When these buffers are ``None``, this module always uses batch statistics.
             in both training and eval modes. Default: ``True``
+        bias: If set to ``False``, the layer will not learn an additive bias (only relevant if
+            :attr:`affine` is ``True``). Default: ``True``
     """
 
     cls_to_become = BatchNorm2d  # type: ignore[assignment]
@@ -558,6 +589,8 @@ class BatchNorm3d(_BatchNorm):
             buffers :attr:`running_mean` and :attr:`running_var` as ``None``.
             When these buffers are ``None``, this module always uses batch statistics.
             in both training and eval modes. Default: ``True``
+        bias: If set to ``False``, the layer will not learn an additive bias (only relevant if
+            :attr:`affine` is ``True``). Default: ``True``
 
     Shape:
         - Input: :math:`(N, C, D, H, W)`
@@ -603,6 +636,8 @@ class LazyBatchNorm3d(_LazyNormBase, _BatchNorm):
             buffers :attr:`running_mean` and :attr:`running_var` as ``None``.
             When these buffers are ``None``, this module always uses batch statistics.
             in both training and eval modes. Default: ``True``
+        bias: If set to ``False``, the layer will not learn an additive bias (only relevant if
+            :attr:`affine` is ``True``). Default: ``True``
     """
 
     cls_to_become = BatchNorm3d  # type: ignore[assignment]
@@ -677,6 +712,8 @@ class SyncBatchNorm(_BatchNorm):
         process_group: synchronization of stats happen within each process group
             individually. Default behavior is synchronization across the whole
             world
+        bias: If set to ``False``, the layer will not learn an additive bias (only relevant if
+            :attr:`affine` is ``True``). Default: ``True``
 
     Shape:
         - Input: :math:`(N, C, +)`
@@ -725,10 +762,18 @@ def __init__(
         process_group: Any | None = None,
         device=None,
         dtype=None,
+        *,
+        bias: bool = True,
     ) -> None:
         factory_kwargs = {"device": device, "dtype": dtype}
         super().__init__(
-            num_features, eps, momentum, affine, track_running_stats, **factory_kwargs
+            num_features,
+            eps,
+            momentum,
+            affine,
+            track_running_stats,
+            **factory_kwargs,
+            bias=bias,
         )
         self.process_group = process_group
 
@@ -886,6 +931,7 @@ def convert_sync_batchnorm(cls, module, process_group=None):
                 module.affine,
                 module.track_running_stats,
                 process_group,
+                bias=module.bias is not None,
             )
             if module.affine:
                 with torch.no_grad():
diff --git a/torch/nn/modules/container.py b/torch/nn/modules/container.py
index f93ed643e73b8..a1a1d7cc990d1 100644
--- a/torch/nn/modules/container.py
+++ b/torch/nn/modules/container.py
@@ -778,11 +778,9 @@ def extra_repr(self) -> str:
                     size_str,
                     device_str,
                 )
-                # pyrefly: ignore [bad-argument-type]
                 child_lines.append("  (" + str(k) + "): " + parastr)
             else:
                 child_lines.append(
-                    # pyrefly: ignore [bad-argument-type]
                     "  (" + str(k) + "): Object of type: " + type(p).__name__
                 )
 
@@ -1015,11 +1013,9 @@ def extra_repr(self) -> str:
                     size_str,
                     device_str,
                 )
-                # pyrefly: ignore [bad-argument-type]
                 child_lines.append("  (" + str(k) + "): " + parastr)
             else:
                 child_lines.append(
-                    # pyrefly: ignore [bad-argument-type]
                     "  (" + str(k) + "): Object of type: " + type(p).__name__
                 )
         tmpstr = "\n".join(child_lines)
diff --git a/torch/nn/modules/conv.py b/torch/nn/modules/conv.py
index f7a656c98b280..bcc39e3e24f05 100644
--- a/torch/nn/modules/conv.py
+++ b/torch/nn/modules/conv.py
@@ -1,6 +1,6 @@
 # mypy: allow-untyped-defs
 import math
-from typing import Literal, Optional
+from typing import Literal
 from typing_extensions import deprecated
 
 import torch
@@ -64,7 +64,7 @@ class _ConvNd(Module):
         "out_channels",
         "kernel_size",
     ]
-    __annotations__ = {"bias": Optional[torch.Tensor]}
+    __annotations__ = {"bias": torch.Tensor | None}
 
     def _conv_forward(  # type: ignore[empty-body]
         self, input: Tensor, weight: Tensor, bias: Tensor | None
@@ -180,14 +180,24 @@ def __init__(
         self.reset_parameters()
 
     def reset_parameters(self) -> None:
-        # Setting a=sqrt(5) in kaiming_uniform is the same as initializing with
-        # uniform(-1/sqrt(k), 1/sqrt(k)), where k = weight.size(1) * prod(*kernel_size)
-        # For more details see: https://github.com/pytorch/pytorch/issues/15314#issuecomment-477448573
-        init.kaiming_uniform_(self.weight, a=math.sqrt(5))
-        if self.bias is not None:
-            fan_in, _ = init._calculate_fan_in_and_fan_out(self.weight)
-            if fan_in != 0:
-                bound = 1 / math.sqrt(fan_in)
+        # Use torch.no_grad() to allow in-place modification of weights/bias
+        with torch.no_grad():
+            # Check if the weight tensor is in channels_last_3d format to ensure numerical consistency
+            if not self.weight.is_contiguous():
+                # Use an empty contiguous buffer to optimize memory and ensure consistent initialization
+                temp_weight = torch.empty_like(
+                    self.weight, memory_format=torch.contiguous_format
+                )
+                init.kaiming_uniform_(temp_weight, a=math.sqrt(5))
+                # Copy the initialized values back to the original tensor
+                self.weight.copy_(temp_weight)
+            else:
+                # Standard initialization for default contiguous memory format
+                init.kaiming_uniform_(self.weight, a=math.sqrt(5))
+            if self.bias is not None:
+                # Bias initialization remains independent of weight memory format
+                fan_in, _ = init._calculate_fan_in_and_fan_out(self.weight)
+                bound = 1 / math.sqrt(fan_in) if fan_in > 0 else 0
                 init.uniform_(self.bias, -bound, bound)
 
     def extra_repr(self):
@@ -390,9 +400,11 @@ class Conv2d(_ConvNd):
 
 
     where :math:`\star` is the valid 2D `cross-correlation`_ operator,
-    :math:`N` is a batch size, :math:`C` denotes a number of channels,
-    :math:`H` is a height of input planes in pixels, and :math:`W` is
-    width in pixels.
+    :math:`N` is a batch size, :math:`C_{\text{in}}` and :math:`C_{\text{out}}` correspond to
+    :attr:`in_channels` and :attr:`out_channels` respectively,
+    :math:`H` and :math:`W` are the input height and width in pixels.
+    See the Shape section below for how :math:`H_{\text{out}}` and :math:`W_{\text{out}}`
+    are derived from :attr:`kernel_size`, :attr:`stride`, :attr:`padding`, and :attr:`dilation`.
     """
         + r"""
 
@@ -1537,11 +1549,11 @@ def __init__(
             padding_mode,
             **factory_kwargs,
         )
-        # pyrefly: ignore [bad-argument-type, bad-override, unexpected-keyword]
+        # pyrefly: ignore [bad-override, unexpected-keyword]
         self.weight = UninitializedParameter(**factory_kwargs)
         self.out_channels = out_channels
         if bias:
-            # pyrefly: ignore [bad-argument-type, bad-override, unexpected-keyword]
+            # pyrefly: ignore [bad-override, unexpected-keyword]
             self.bias = UninitializedParameter(**factory_kwargs)
 
     def _get_num_spatial_dims(self) -> int:
@@ -1609,11 +1621,11 @@ def __init__(
             padding_mode,
             **factory_kwargs,
         )
-        # pyrefly: ignore [bad-argument-type, bad-override, unexpected-keyword]
+        # pyrefly: ignore [bad-override, unexpected-keyword]
         self.weight = UninitializedParameter(**factory_kwargs)
         self.out_channels = out_channels
         if bias:
-            # pyrefly: ignore [bad-argument-type, bad-override, unexpected-keyword]
+            # pyrefly: ignore [bad-override, unexpected-keyword]
             self.bias = UninitializedParameter(**factory_kwargs)
 
     def _get_num_spatial_dims(self) -> int:
@@ -1682,11 +1694,11 @@ def __init__(
             padding_mode,
             **factory_kwargs,
         )
-        # pyrefly: ignore [bad-argument-type, bad-override, unexpected-keyword]
+        # pyrefly: ignore [bad-override, unexpected-keyword]
         self.weight = UninitializedParameter(**factory_kwargs)
         self.out_channels = out_channels
         if bias:
-            # pyrefly: ignore [bad-argument-type, bad-override, unexpected-keyword]
+            # pyrefly: ignore [bad-override, unexpected-keyword]
             self.bias = UninitializedParameter(**factory_kwargs)
 
     def _get_num_spatial_dims(self) -> int:
@@ -1754,11 +1766,11 @@ def __init__(
             padding_mode,
             **factory_kwargs,
         )
-        # pyrefly: ignore [bad-argument-type, bad-override, unexpected-keyword]
+        # pyrefly: ignore [bad-override, unexpected-keyword]
         self.weight = UninitializedParameter(**factory_kwargs)
         self.out_channels = out_channels
         if bias:
-            # pyrefly: ignore [bad-argument-type, bad-override, unexpected-keyword]
+            # pyrefly: ignore [bad-override, unexpected-keyword]
             self.bias = UninitializedParameter(**factory_kwargs)
 
     def _get_num_spatial_dims(self) -> int:
@@ -1826,11 +1838,11 @@ def __init__(
             padding_mode,
             **factory_kwargs,
         )
-        # pyrefly: ignore [bad-argument-type, bad-override, unexpected-keyword]
+        # pyrefly: ignore [bad-override, unexpected-keyword]
         self.weight = UninitializedParameter(**factory_kwargs)
         self.out_channels = out_channels
         if bias:
-            # pyrefly: ignore [bad-argument-type, bad-override, unexpected-keyword]
+            # pyrefly: ignore [bad-override, unexpected-keyword]
             self.bias = UninitializedParameter(**factory_kwargs)
 
     def _get_num_spatial_dims(self) -> int:
@@ -1898,11 +1910,11 @@ def __init__(
             padding_mode,
             **factory_kwargs,
         )
-        # pyrefly: ignore [bad-argument-type, bad-override, unexpected-keyword]
+        # pyrefly: ignore [bad-override, unexpected-keyword]
         self.weight = UninitializedParameter(**factory_kwargs)
         self.out_channels = out_channels
         if bias:
-            # pyrefly: ignore [bad-argument-type, bad-override, unexpected-keyword]
+            # pyrefly: ignore [bad-override, unexpected-keyword]
             self.bias = UninitializedParameter(**factory_kwargs)
 
     def _get_num_spatial_dims(self) -> int:
diff --git a/torch/nn/modules/instancenorm.py b/torch/nn/modules/instancenorm.py
index 058ffb3ed9aa9..76a343afbd0dd 100644
--- a/torch/nn/modules/instancenorm.py
+++ b/torch/nn/modules/instancenorm.py
@@ -28,10 +28,18 @@ def __init__(
         track_running_stats: bool = False,
         device=None,
         dtype=None,
+        *,
+        bias: bool = True,  # for backward compatibility
     ) -> None:
         factory_kwargs = {"device": device, "dtype": dtype}
         super().__init__(
-            num_features, eps, momentum, affine, track_running_stats, **factory_kwargs
+            num_features,
+            eps,
+            momentum,
+            affine,
+            track_running_stats,
+            **factory_kwargs,
+            bias=bias,
         )
 
     def _check_input_dim(self, input):
@@ -174,11 +182,13 @@ class InstanceNorm1d(_InstanceNorm):
         momentum: the value used for the running_mean and running_var computation. Default: 0.1
         affine: a boolean value that when set to ``True``, this module has
             learnable affine parameters, initialized the same way as done for batch normalization.
-            Default: ``False``.
+            Default: ``False``
         track_running_stats: a boolean value that when set to ``True``, this
             module tracks the running mean and variance, and when set to ``False``,
             this module does not track such statistics and always uses batch
             statistics in both training and eval modes. Default: ``False``
+        bias: If set to ``False``, the layer will not learn an additive bias (only relevant if
+            :attr:`affine` is ``True``). Default: ``True``
 
     Shape:
         - Input: :math:`(N, C, L)` or :math:`(C, L)`
@@ -218,11 +228,13 @@ class LazyInstanceNorm1d(_LazyNormBase, _InstanceNorm):
         momentum: the value used for the running_mean and running_var computation. Default: 0.1
         affine: a boolean value that when set to ``True``, this module has
             learnable affine parameters, initialized the same way as done for batch normalization.
-            Default: ``False``.
+            Default: ``False``
         track_running_stats: a boolean value that when set to ``True``, this
             module tracks the running mean and variance, and when set to ``False``,
             this module does not track such statistics and always uses batch
             statistics in both training and eval modes. Default: ``False``
+        bias: If set to ``False``, the layer will not learn an additive bias (only relevant if
+            :attr:`affine` is ``True``). Default: ``True``
 
     Shape:
         - Input: :math:`(N, C, L)` or :math:`(C, L)`
@@ -290,11 +302,13 @@ class InstanceNorm2d(_InstanceNorm):
         momentum: the value used for the running_mean and running_var computation. Default: 0.1
         affine: a boolean value that when set to ``True``, this module has
             learnable affine parameters, initialized the same way as done for batch normalization.
-            Default: ``False``.
+            Default: ``False``
         track_running_stats: a boolean value that when set to ``True``, this
             module tracks the running mean and variance, and when set to ``False``,
             this module does not track such statistics and always uses batch
             statistics in both training and eval modes. Default: ``False``
+        bias: If set to ``False``, the layer will not learn an additive bias (only relevant if
+            :attr:`affine` is ``True``). Default: ``True``
 
     Shape:
         - Input: :math:`(N, C, H, W)` or :math:`(C, H, W)`
@@ -335,11 +349,13 @@ class LazyInstanceNorm2d(_LazyNormBase, _InstanceNorm):
         momentum: the value used for the running_mean and running_var computation. Default: 0.1
         affine: a boolean value that when set to ``True``, this module has
             learnable affine parameters, initialized the same way as done for batch normalization.
-            Default: ``False``.
+            Default: ``False``
         track_running_stats: a boolean value that when set to ``True``, this
             module tracks the running mean and variance, and when set to ``False``,
             this module does not track such statistics and always uses batch
             statistics in both training and eval modes. Default: ``False``
+        bias: If set to ``False``, the layer will not learn an additive bias (only relevant if
+            :attr:`affine` is ``True``). Default: ``True``
 
     Shape:
         - Input: :math:`(N, C, H, W)` or :math:`(C, H, W)`
@@ -406,11 +422,13 @@ class InstanceNorm3d(_InstanceNorm):
         momentum: the value used for the running_mean and running_var computation. Default: 0.1
         affine: a boolean value that when set to ``True``, this module has
             learnable affine parameters, initialized the same way as done for batch normalization.
-            Default: ``False``.
+            Default: ``False``
         track_running_stats: a boolean value that when set to ``True``, this
             module tracks the running mean and variance, and when set to ``False``,
             this module does not track such statistics and always uses batch
             statistics in both training and eval modes. Default: ``False``
+        bias: If set to ``False``, the layer will not learn an additive bias (only relevant if
+            :attr:`affine` is ``True``). Default: ``True``
 
     Shape:
         - Input: :math:`(N, C, D, H, W)` or :math:`(C, D, H, W)`
@@ -451,11 +469,13 @@ class LazyInstanceNorm3d(_LazyNormBase, _InstanceNorm):
         momentum: the value used for the running_mean and running_var computation. Default: 0.1
         affine: a boolean value that when set to ``True``, this module has
             learnable affine parameters, initialized the same way as done for batch normalization.
-            Default: ``False``.
+            Default: ``False``
         track_running_stats: a boolean value that when set to ``True``, this
             module tracks the running mean and variance, and when set to ``False``,
             this module does not track such statistics and always uses batch
             statistics in both training and eval modes. Default: ``False``
+        bias: If set to ``False``, the layer will not learn an additive bias (only relevant if
+            :attr:`affine` is ``True``). Default: ``True``
 
     Shape:
         - Input: :math:`(N, C, D, H, W)` or :math:`(C, D, H, W)`
diff --git a/torch/nn/modules/linear.py b/torch/nn/modules/linear.py
index 650e7e7c80ff9..30d693b3278b2 100644
--- a/torch/nn/modules/linear.py
+++ b/torch/nn/modules/linear.py
@@ -298,11 +298,11 @@ def __init__(
         # that will soon be overwritten.
         # pyrefly: ignore [bad-argument-type]
         super().__init__(0, 0, False)
-        # pyrefly: ignore [bad-argument-type, unexpected-keyword]
+        # pyrefly: ignore [unexpected-keyword]
         self.weight = UninitializedParameter(**factory_kwargs)
         self.out_features = out_features
         if bias:
-            # pyrefly: ignore [bad-argument-type, unexpected-keyword]
+            # pyrefly: ignore [unexpected-keyword]
             self.bias = UninitializedParameter(**factory_kwargs)
 
     def reset_parameters(self) -> None:
diff --git a/torch/nn/modules/loss.py b/torch/nn/modules/loss.py
index f1fe46c773fd8..9af79927c1b31 100644
--- a/torch/nn/modules/loss.py
+++ b/torch/nn/modules/loss.py
@@ -1919,7 +1919,6 @@ def __init__(
                 f"TripletMarginWithDistanceLoss: expected margin to be greater than 0, got {margin} instead"
             )
         self.distance_function: Callable[[Tensor, Tensor], Tensor] | None = (
-            # pyrefly: ignore [bad-assignment]
             distance_function if distance_function is not None else PairwiseDistance()
         )
         self.margin = margin
diff --git a/torch/nn/modules/module.py b/torch/nn/modules/module.py
index 744c1b367a4bd..e004ef8d14f13 100644
--- a/torch/nn/modules/module.py
+++ b/torch/nn/modules/module.py
@@ -30,7 +30,7 @@
     "Module",
 ]
 
-_grad_t = Union[tuple[Tensor, ...], Tensor]
+_grad_t = tuple[Tensor, ...] | Tensor
 # See https://mypy.readthedocs.io/en/latest/generics.html#generic-methods-and-generic-self for the use
 # of `T` to annotate `self`. Many methods of `Module` return `self` and we want those return values to be
 # the type of the subclass, not the looser type of `Module`.
@@ -44,7 +44,6 @@ class _IncompatibleKeys(
     __slots__ = ()
 
     def __repr__(self) -> str:
-        # pyrefly: ignore [missing-attribute]
         if not self.missing_keys and not self.unexpected_keys:
             return "<All keys matched successfully>"
         return super().__repr__()
@@ -73,7 +72,7 @@ def _addindent(s_, numSpaces):
 
 
 class _WrappedHook:
-    def __init__(self, hook: Callable, module: Optional["Module"] = None) -> None:
+    def __init__(self, hook: Callable, module: Optional["Module"] = None) -> None:  # noqa: UP045
         self.hook: Callable = hook
         functools.update_wrapper(self, hook)
 
@@ -94,7 +93,7 @@ def __call__(self, *args: Any, **kwargs: Any) -> Any:
     def __getstate__(self) -> dict:
         result = {"hook": self.hook, "with_module": self.with_module}
         if self.with_module:
-            # pyrefly: ignore [bad-typed-dict-key, unsupported-operation]
+            # pyrefly: ignore [bad-typed-dict-key]
             result["module"] = self.module()
 
         return result
@@ -476,7 +475,7 @@ def forward(self, x):
     _load_state_dict_pre_hooks: dict[int, Callable]
     _state_dict_pre_hooks: dict[int, Callable]
     _load_state_dict_post_hooks: dict[int, Callable]
-    _modules: dict[str, Optional["Module"]]
+    _modules: dict[str, Optional["Module"]]  # noqa: UP045
     call_super_init: bool = False
     _compiled_call_impl: Callable | None = None
 
@@ -640,7 +639,7 @@ def register_parameter(self, name: str, param: Parameter | None) -> None:
                     param = output
             self._parameters[name] = param
 
-    def add_module(self, name: str, module: Optional["Module"]) -> None:
+    def add_module(self, name: str, module: Optional["Module"]) -> None:  # noqa: UP045
         r"""Add a child module to the current module.
 
         The module can be accessed as an attribute using the given name.
@@ -668,7 +667,7 @@ def add_module(self, name: str, module: Optional["Module"]) -> None:
                 module = output
         self._modules[name] = module
 
-    def register_module(self, name: str, module: Optional["Module"]) -> None:
+    def register_module(self, name: str, module: Optional["Module"]) -> None:  # noqa: UP045
         r"""Alias for :func:`add_module`."""
         self.add_module(name, module)
 
@@ -818,7 +817,7 @@ def set_submodule(
                 raise AttributeError("`" + atoms[-1] + "` is not an nn.Module")
         setattr(parent, atoms[-1], module)
 
-    def get_parameter(self, target: str) -> "Parameter":
+    def get_parameter(self, target: str) -> Parameter:
         """Return the parameter given by ``target`` if it exists, otherwise throw an error.
 
         See the docstring for ``get_submodule`` for a more detailed
@@ -854,7 +853,7 @@ def get_parameter(self, target: str) -> "Parameter":
 
         return param
 
-    def get_buffer(self, target: str) -> "Tensor":
+    def get_buffer(self, target: str) -> Tensor:
         """Return the buffer given by ``target`` if it exists, otherwise throw an error.
 
         See the docstring for ``get_submodule`` for a more detailed
@@ -2836,7 +2835,7 @@ def modules(self, remove_duplicate: bool = True) -> Iterator["Module"]:
 
     def named_modules(
         self,
-        memo: set["Module"] | None = None,
+        memo: set["Module"] | None = None,  # noqa: UP007
         prefix: str = "",
         remove_duplicate: bool = True,
     ):
diff --git a/torch/nn/modules/normalization.py b/torch/nn/modules/normalization.py
index d278aabcfbb44..c4df21528cae2 100644
--- a/torch/nn/modules/normalization.py
+++ b/torch/nn/modules/normalization.py
@@ -1,6 +1,5 @@
 # mypy: allow-untyped-defs
 import numbers
-from typing import Union
 
 import torch
 from torch import Size, Tensor
@@ -100,7 +99,7 @@ def extra_repr(self) -> str:
         return "{size}, alpha={alpha}, beta={beta}, k={k}".format(**self.__dict__)
 
 
-_shape_t = Union[int, list[int], Size]
+_shape_t = int | list[int] | Size
 
 
 class LayerNorm(Module):
@@ -143,9 +142,9 @@ class LayerNorm(Module):
         eps: a value added to the denominator for numerical stability. Default: 1e-5
         elementwise_affine: a boolean value that when set to ``True``, this module
             has learnable per-element affine parameters initialized to ones (for weights)
-            and zeros (for biases). Default: ``True``.
+            and zeros (for biases). Default: ``True``
         bias: If set to ``False``, the layer will not learn an additive bias (only relevant if
-            :attr:`elementwise_affine` is ``True``). Default: ``True``.
+            :attr:`elementwise_affine` is ``True``). Default: ``True``
 
     Attributes:
         weight: the learnable weights of the module of shape
@@ -232,8 +231,8 @@ def forward(self, input: Tensor) -> Tensor:
 
     def extra_repr(self) -> str:
         return (
-            "{normalized_shape}, eps={eps}, "
-            "elementwise_affine={elementwise_affine}".format(**self.__dict__)
+            "{normalized_shape}, eps={eps}, elementwise_affine={elementwise_affine}, "
+            "bias={use_bias}".format(**self.__dict__, use_bias=self.bias is not None)
         )
 
 
@@ -264,7 +263,9 @@ class GroupNorm(Module):
         eps: a value added to the denominator for numerical stability. Default: 1e-5
         affine: a boolean value that when set to ``True``, this module
             has learnable per-channel affine parameters initialized to ones (for weights)
-            and zeros (for biases). Default: ``True``.
+            and zeros (for biases). Default: ``True``
+        bias: If set to ``False``, the layer will not learn an additive bias (only relevant if
+            :attr:`affine` is ``True``). Default: ``True``
 
     Shape:
         - Input: :math:`(N, C, *)` where :math:`C=\text{num\_channels}`
@@ -297,6 +298,8 @@ def __init__(
         affine: bool = True,
         device=None,
         dtype=None,
+        *,
+        bias: bool = True,
     ) -> None:
         factory_kwargs = {"device": device, "dtype": dtype}
         super().__init__()
@@ -311,7 +314,10 @@ def __init__(
         self.affine = affine
         if self.affine:
             self.weight = Parameter(torch.empty(num_channels, **factory_kwargs))
-            self.bias = Parameter(torch.empty(num_channels, **factory_kwargs))
+            if bias:
+                self.bias = Parameter(torch.empty(num_channels, **factory_kwargs))
+            else:
+                self.register_parameter("bias", None)
         else:
             self.register_parameter("weight", None)
             self.register_parameter("bias", None)
@@ -321,14 +327,16 @@ def __init__(
     def reset_parameters(self) -> None:
         if self.affine:
             init.ones_(self.weight)
-            init.zeros_(self.bias)
+            if self.bias is not None:
+                init.zeros_(self.bias)
 
     def forward(self, input: Tensor) -> Tensor:
         return F.group_norm(input, self.num_groups, self.weight, self.bias, self.eps)
 
     def extra_repr(self) -> str:
-        return "{num_groups}, {num_channels}, eps={eps}, affine={affine}".format(
-            **self.__dict__
+        return (
+            "{num_groups}, {num_channels}, eps={eps}, affine={affine}, "
+            "bias={use_bias}".format(**self.__dict__, use_bias=self.bias is not None)
         )
 
 
@@ -357,10 +365,10 @@ class RMSNorm(Module):
 
             If a single integer is used, it is treated as a singleton list, and this module will
             normalize over the last dimension which is expected to be of that specific size.
-        eps: a value added to the denominator for numerical stability. If not specified,
-            uses the machine epsilon of the computation (opmath) type: fp16/bf16 and
-            fp32 inputs use ``torch.finfo(torch.float32).eps``, while fp64 inputs use
-            ``torch.finfo(torch.float64).eps``.
+        eps (float, optional): a value added to the denominator for numerical stability.
+            If not specified, uses the machine epsilon of the computation (opmath) type:
+            fp16/bf16 and fp32 inputs use ``torch.finfo(torch.float32).eps``, while fp64
+            inputs use ``torch.finfo(torch.float64).eps``. Default: ``None``
         elementwise_affine: a boolean value that when set to ``True``, this module
             has learnable per-element affine parameters initialized to ones (for weights). Default: ``True``.
 
diff --git a/torch/nn/modules/rnn.py b/torch/nn/modules/rnn.py
index 76bfa4ea96475..6f6e3b6154725 100644
--- a/torch/nn/modules/rnn.py
+++ b/torch/nn/modules/rnn.py
@@ -655,12 +655,10 @@ def __init__(self, *args, **kwargs):
 
     @overload
     @torch._jit_internal._overload_method  # noqa: F811
-    # pyrefly: ignore [bad-override]
     def forward(
         self,
         input: Tensor,
         hx: Tensor | None = None,
-        # pyrefly: ignore [bad-return]
     ) -> tuple[Tensor, Tensor]:
         pass
 
@@ -670,7 +668,6 @@ def forward(
         self,
         input: PackedSequence,
         hx: Tensor | None = None,
-        # pyrefly: ignore [bad-return]
     ) -> tuple[PackedSequence, Tensor]:
         pass
 
@@ -802,7 +799,6 @@ def forward(self, input, hx=None):  # noqa: F811
         if isinstance(orig_input, PackedSequence):
             output_packed = PackedSequence(
                 output,
-                # pyrefly: ignore [bad-argument-type]
                 batch_sizes,
                 sorted_indices,
                 unsorted_indices,
@@ -1063,12 +1059,10 @@ def permute_hidden(  # type: ignore[override]
     # Same as above, see torch/nn/modules/module.py::_forward_unimplemented
     @overload  # type: ignore[override]
     @torch._jit_internal._overload_method  # noqa: F811
-    # pyrefly: ignore [bad-override]
     def forward(
         self,
         input: Tensor,
         hx: tuple[Tensor, Tensor] | None = None,
-        # pyrefly: ignore [bad-return]
     ) -> tuple[Tensor, tuple[Tensor, Tensor]]:  # noqa: F811
         pass
 
@@ -1079,7 +1073,6 @@ def forward(
         self,
         input: PackedSequence,
         hx: tuple[Tensor, Tensor] | None = None,
-        # pyrefly: ignore [bad-return]
     ) -> tuple[PackedSequence, tuple[Tensor, Tensor]]:  # noqa: F811
         pass
 
@@ -1196,7 +1189,6 @@ def forward(self, input, hx=None):  # noqa: F811
         if isinstance(orig_input, PackedSequence):
             output_packed = PackedSequence(
                 output,
-                # pyrefly: ignore [bad-argument-type]
                 batch_sizes,
                 sorted_indices,
                 unsorted_indices,
@@ -1365,12 +1357,10 @@ def __init__(self, *args, **kwargs):
 
     @overload  # type: ignore[override]
     @torch._jit_internal._overload_method  # noqa: F811
-    # pyrefly: ignore [bad-override]
     def forward(
         self,
         input: Tensor,
         hx: Tensor | None = None,
-        # pyrefly: ignore [bad-return]
     ) -> tuple[Tensor, Tensor]:  # noqa: F811
         pass
 
@@ -1380,7 +1370,6 @@ def forward(
         self,
         input: PackedSequence,
         hx: Tensor | None = None,
-        # pyrefly: ignore [bad-return]
     ) -> tuple[PackedSequence, Tensor]:  # noqa: F811
         pass
 
@@ -1477,7 +1466,6 @@ def forward(self, input, hx=None):  # noqa: F811
         if isinstance(orig_input, PackedSequence):
             output_packed = PackedSequence(
                 output,
-                # pyrefly: ignore [bad-argument-type]
                 batch_sizes,
                 sorted_indices,
                 unsorted_indices,
diff --git a/torch/nn/modules/transformer.py b/torch/nn/modules/transformer.py
index 6841e85ed6d2e..7f47013c14585 100644
--- a/torch/nn/modules/transformer.py
+++ b/torch/nn/modules/transformer.py
@@ -138,7 +138,6 @@ def __init__(
                 d_model,
                 eps=layer_norm_eps,
                 bias=bias,
-                # pyrefly: ignore [bad-argument-type]
                 **factory_kwargs,
             )
             self.encoder = TransformerEncoder(
@@ -164,7 +163,6 @@ def __init__(
                 d_model,
                 eps=layer_norm_eps,
                 bias=bias,
-                # pyrefly: ignore [bad-argument-type]
                 **factory_kwargs,
             )
             self.decoder = TransformerDecoder(
@@ -770,9 +768,7 @@ def __init__(
         self.linear2 = Linear(dim_feedforward, d_model, bias=bias, **factory_kwargs)
 
         self.norm_first = norm_first
-        # pyrefly: ignore [bad-argument-type]
         self.norm1 = LayerNorm(d_model, eps=layer_norm_eps, bias=bias, **factory_kwargs)
-        # pyrefly: ignore [bad-argument-type]
         self.norm2 = LayerNorm(d_model, eps=layer_norm_eps, bias=bias, **factory_kwargs)
         self.dropout1 = Dropout(dropout)
         self.dropout2 = Dropout(dropout)
@@ -1065,11 +1061,8 @@ def __init__(
         self.linear2 = Linear(dim_feedforward, d_model, bias=bias, **factory_kwargs)
 
         self.norm_first = norm_first
-        # pyrefly: ignore [bad-argument-type]
         self.norm1 = LayerNorm(d_model, eps=layer_norm_eps, bias=bias, **factory_kwargs)
-        # pyrefly: ignore [bad-argument-type]
         self.norm2 = LayerNorm(d_model, eps=layer_norm_eps, bias=bias, **factory_kwargs)
-        # pyrefly: ignore [bad-argument-type]
         self.norm3 = LayerNorm(d_model, eps=layer_norm_eps, bias=bias, **factory_kwargs)
         self.dropout1 = Dropout(dropout)
         self.dropout2 = Dropout(dropout)
diff --git a/torch/nn/parallel/data_parallel.py b/torch/nn/parallel/data_parallel.py
index 8eede2b979133..16537624663b1 100644
--- a/torch/nn/parallel/data_parallel.py
+++ b/torch/nn/parallel/data_parallel.py
@@ -161,7 +161,6 @@ def __init__(
         self.module = module
         self.device_ids = [_get_device_index(x, True) for x in device_ids]
         self.output_device = _get_device_index(output_device, True)
-        # pyrefly: ignore [read-only]
         self.src_device_obj = torch.device(device_type, self.device_ids[0])
 
         if device_type == "cuda":
@@ -260,7 +259,7 @@ def data_parallel(
 
     device_ids = [_get_device_index(x, True) for x in device_ids]
     output_device = _get_device_index(output_device, True)
-    # pyrefly: ignore [no-matching-overload]
+    # pyrefly: ignore [bad-argument-type, no-matching-overload]
     src_device_obj = torch.device(device_type, device_ids[0])
 
     # pyrefly: ignore [bad-argument-type]
diff --git a/torch/nn/parallel/distributed.py b/torch/nn/parallel/distributed.py
index 8199208957046..50f230b995fba 100644
--- a/torch/nn/parallel/distributed.py
+++ b/torch/nn/parallel/distributed.py
@@ -26,6 +26,142 @@
 
 
 RPC_AVAILABLE = False
+
+# Default bucket size in MiB for gradient reduction
+_DEFAULT_BUCKET_CAP_MB = 25
+# Conversion factor from MiB to bytes
+_MB_TO_BYTES = 1024 * 1024
+
+
+@dataclass(frozen=True)
+class _BucketCapacityConfig:
+    """Configuration for DDP gradient reduction bucket capacities.
+
+    This immutable dataclass encapsulates bucket size settings for
+    DistributedDataParallel. Use the `create()` factory method to construct.
+
+    The bucket capacity determines how parameters are grouped for AllReduce:
+    - Smaller buckets: more frequent, smaller AllReduce operations
+    - Larger buckets: less frequent, larger AllReduce operations
+
+    Attributes:
+        bucket_bytes_cap: The primary bucket size limit in bytes. This is passed
+            to the C++ Reducer as the default bucket size.
+
+        per_bucket_bytes_caps: Per-bucket size limits in bytes when user specifies
+            bucket_cap_mb_list. Empty tuple means use uniform bucket_bytes_cap.
+
+        first_bucket_bytes_cap: Size limit for the first bucket in bytes. Can be
+            smaller than bucket_bytes_cap to reduce latency (first bucket contains
+            parameters whose gradients complete last, so a smaller first bucket
+            allows AllReduce to start sooner).
+    """
+
+    bucket_bytes_cap: int
+    per_bucket_bytes_caps: tuple[int, ...]
+    first_bucket_bytes_cap: int
+
+    @classmethod
+    def create(
+        cls,
+        bucket_cap_mb: int | None,
+        bucket_cap_mb_list: list[int] | None,
+        use_python_reducer: bool,
+    ) -> "_BucketCapacityConfig":
+        """Factory method to create a BucketCapacityConfig from user inputs.
+
+        Args:
+            bucket_cap_mb: Single bucket size limit in MiB. If None, uses default
+                of 25 MiB. Ignored if bucket_cap_mb_list is provided.
+            bucket_cap_mb_list: Per-bucket size limits in MiB. When provided, each
+                bucket gets its own size limit.
+            use_python_reducer: Whether Python reducer is being used (for validation).
+
+        Returns:
+            A new BucketCapacityConfig instance.
+
+        Raises:
+            AssertionError: If bucket_cap_mb_list is provided with Python reducer.
+        """
+        is_using_default = bucket_cap_mb is None
+        effective_bucket_cap_mb = (
+            bucket_cap_mb if bucket_cap_mb is not None else _DEFAULT_BUCKET_CAP_MB
+        )
+
+        # Process per-bucket size list if provided
+        per_bucket_bytes_caps: tuple[int, ...] = ()
+        if bucket_cap_mb_list:
+            if use_python_reducer:
+                raise AssertionError(
+                    "when using bucket_cap_mb_list, python reducer is not supported"
+                )
+            per_bucket_bytes_caps = tuple(
+                int(cap_mb * _MB_TO_BYTES) for cap_mb in bucket_cap_mb_list
+            )
+            effective_bucket_cap_mb = max(bucket_cap_mb_list)
+            is_using_default = False
+
+        bucket_bytes_cap = int(effective_bucket_cap_mb * _MB_TO_BYTES)
+
+        # First bucket optimization: use smaller size when using defaults
+        # to reduce latency for early-computed gradients
+        first_bucket_bytes_cap = (
+            dist._DEFAULT_FIRST_BUCKET_BYTES if is_using_default else bucket_bytes_cap
+        )
+
+        return cls(
+            bucket_bytes_cap=bucket_bytes_cap,
+            per_bucket_bytes_caps=per_bucket_bytes_caps,
+            first_bucket_bytes_cap=first_bucket_bytes_cap,
+        )
+
+    @property
+    def bucket_cap_mb(self) -> int:
+        """Return bucket size in MiB (for backward compatibility)."""
+        return self.bucket_bytes_cap // _MB_TO_BYTES
+
+    @property
+    def has_custom_per_bucket_caps(self) -> bool:
+        """Return True if user specified per-bucket size limits."""
+        return len(self.per_bucket_bytes_caps) > 0
+
+    def compute_bucket_size_limits(
+        self,
+        static_graph: bool,
+        find_unused_parameters: bool,
+    ) -> tuple[list[int], list[int]]:
+        """Compute bucket size limits for initial bucketing and rebuilding.
+
+        Args:
+            static_graph: Whether the computation graph is static.
+            find_unused_parameters: Whether to find unused parameters.
+
+        Returns:
+            A tuple of (bucket_size_limits, bucket_size_limits_for_rebuilding):
+            - bucket_size_limits: Used for initial bucket assignment
+            - bucket_size_limits_for_rebuilding: Passed to C++ Reducer (empty list
+              means use C++ Reducer's default logic)
+        """
+        # Case 1: User provided explicit per-bucket size limits
+        if self.has_custom_per_bucket_caps:
+            limits = list(self.per_bucket_bytes_caps)
+            return (limits, limits)
+
+        # Case 2: Compute default bucket size limits
+        # When static_graph or not finding unused params, disable initial bucketing
+        # by using maxsize (put all params in one bucket initially)
+        if static_graph or not find_unused_parameters:
+            bucket_size_limits = [sys.maxsize]
+        elif self.first_bucket_bytes_cap < self.bucket_bytes_cap:
+            # Use smaller first bucket for latency optimization
+            bucket_size_limits = [self.first_bucket_bytes_cap, self.bucket_bytes_cap]
+        else:
+            bucket_size_limits = [self.bucket_bytes_cap]
+
+        # Empty list lets C++ Reducer use its default logic
+        return (bucket_size_limits, [])
+
+
 if dist.is_available():
     from torch.distributed.distributed_c10d import (
         _get_default_group,
@@ -518,7 +654,7 @@ class DistributedDataParallel(Module, Joinable):
         If you plan on using this module with a ``nccl`` backend or a ``gloo``
         backend (that uses Infiniband), together with a DataLoader that uses
         multiple workers, please change the multiprocessing start method to
-        ``forkserver`` (Python 3 only) or ``spawn``. Unfortunately
+        ``forkserver`` or ``spawn``. Unfortunately
         Gloo (that uses Infiniband) and NCCL2 are not fork safe, and you will
         likely experience deadlocks if you don't change this setting.
 
@@ -637,6 +773,17 @@ class DistributedDataParallel(Module, Joinable):
                     This requires that unused parameters remain the same across all ranks throughout
                     the entire training process. If this condition is not met, it may cause
                     desynchronization and result in training hang.
+        batched_grad_copy (bool): When set to ``True``, individual per-parameter
+                    gradient-to-bucket copy and division operations are deferred
+                    and flushed as a single ``_foreach_copy_`` plus one flat
+                    ``div_`` when a bucket becomes ready. This reduces per-parameter
+                    kernel launches down to 2 kernels per bucket, which can improve
+                    throughput for models with many small parameters. The
+                    optimization is most effective with
+                    ``optimizer.zero_grad(set_to_none=True)`` (the default), where
+                    ``gradient_as_bucket_view`` alone cannot avoid copies because
+                    the bucket view alias is destroyed every iteration.
+                    (default: ``False``)
 
 
     Attributes:
@@ -671,7 +818,8 @@ def __init__(
         mixed_precision: _MixedPrecision | None = None,
         device_mesh=None,
         skip_all_reduce_unused_params=False,
-        bucket_cap_mb_list: Optional[list[int]] = None,
+        bucket_cap_mb_list: list[int] | None = None,
+        batched_grad_copy=False,
     ):
         super().__init__()
         Joinable.__init__(self)
@@ -772,24 +920,23 @@ def __init__(
             or self.is_multi_device_module
         ):
             if device_ids or output_device:
+                devices = {p.device for p in self._module_parameters}
                 self._log_and_throw(
                     ValueError,
                     "DistributedDataParallel device_ids and output_device arguments "
                     "only work with single-device/multiple-device GPU modules or CPU modules, "
                     f"but got device_ids {device_ids}, output_device {output_device}, "
-                    f"and module parameters { ({p.device for p in self._module_parameters}) }.",
+                    f"and module parameters {devices}.",
                 )
 
             self.device_ids = None
             self.output_device = None
         else:
-            # pyrefly: ignore [bad-assignment]
             self.device_ids = [_get_device_index(x, True) for x in device_ids]
 
             if output_device is None:
                 output_device = device_ids[0]
 
-            # pyrefly: ignore [bad-assignment]
             self.output_device = _get_device_index(output_device, True)
 
         self.static_graph = False
@@ -801,6 +948,7 @@ def __init__(
         self.require_backward_grad_sync = True
         self.require_forward_param_sync = True
         self.gradient_as_bucket_view = gradient_as_bucket_view
+        self.batched_grad_copy = batched_grad_copy
         self.mixed_precision = mixed_precision
         if self.mixed_precision is not None:
             logger.warning("Received mixed precision config %s", self.mixed_precision)
@@ -827,30 +975,21 @@ def __init__(
         # used for intra-node param sync and inter-node sync as well
         self.broadcast_bucket_size = 250 * 1024 * 1024
 
-        # reduction bucket size
-        if bucket_cap_mb is None:
-            # default case (bucket cap is 25 MiB)
-            bucket_cap_mb = 25
-            self.bucket_bytes_cap_default = True
-        else:
-            self.bucket_bytes_cap_default = False
-
-        self.bucket_bytes_cap_list = []
-        if bucket_cap_mb_list:
-            if self._use_python_reducer:
-                raise AssertionError(
-                    "when using bucket_cap_mb_list, python reducer is not supported"
-                )
-            self.bucket_bytes_cap_list = [
-                int(bucket_cap_mb * 1024 * 1024) for bucket_cap_mb in bucket_cap_mb_list
-            ]
-            # This is not supposed to be used later when custom bucket_cap_mb_list is passed,
-            # just a safe placeholder for backward compatibility
-            self.bucket_bytes_cap_default = False
-            bucket_cap_mb = max(bucket_cap_mb_list)
-        if not bucket_cap_mb:
-            raise AssertionError("bucket_cap_mb should be set by now")
-        self.bucket_bytes_cap = int(bucket_cap_mb * 1024 * 1024)
+        # Initialize bucket capacity configuration using the factory method
+        self._bucket_config = _BucketCapacityConfig.create(
+            bucket_cap_mb=bucket_cap_mb,
+            bucket_cap_mb_list=bucket_cap_mb_list,
+            use_python_reducer=self._use_python_reducer,
+        )
+        # Expose config values as instance attributes for backward compatibility
+        # TODO: Remove in the future
+        self.bucket_bytes_cap = self._bucket_config.bucket_bytes_cap
+        self.bucket_bytes_cap_list = list(self._bucket_config.per_bucket_bytes_caps)
+        self.bucket_bytes_cap_default = (
+            self._bucket_config.first_bucket_bytes_cap
+            < self._bucket_config.bucket_bytes_cap
+        )
+        bucket_cap_mb = self._bucket_config.bucket_cap_mb
 
         # Whether to perform input tensor CPU to GPU copies on a side-stream
         self.use_side_stream_for_tensor_copies = (
@@ -1215,25 +1354,16 @@ def _ddp_init_helper(
         # After the first iteration, it's OK to rebuild buckets,
         # because "bucket rebuild" bucketizes parameters based on its real execution order in backward graph.
 
-        # Can remove this branching once #73732 is landed.
-        if self.bucket_bytes_cap_list:
-            bucket_size_limits = self.bucket_bytes_cap_list
-            # When bucket_cap_mb_list is provided, use it for rebuilding buckets
-            bucket_size_limits_for_rebuilding = self.bucket_bytes_cap_list
-        else:
-            if static_graph is True or self.find_unused_parameters is False:
-                bucket_size_limits = [sys.maxsize]
-            else:
-                if self.bucket_bytes_cap_default:
-                    bucket_size_limits = [
-                        dist._DEFAULT_FIRST_BUCKET_BYTES,
-                        self.bucket_bytes_cap,
-                    ]
-                else:
-                    bucket_size_limits = [self.bucket_bytes_cap]
-            # When bucket_cap_mb_list is not provided, pass empty list
-            # to let C++ Reducer use the original logic (first_bucket_bytes_cap_ and bucket_bytes_cap_)
-            bucket_size_limits_for_rebuilding = []
+        # Compute bucket size limits for initial bucketing and rebuilding
+        # See BucketCapacityConfig.compute_bucket_size_limits for detailed explanation
+        (
+            bucket_size_limits,
+            bucket_size_limits_for_rebuilding,
+        ) = self._bucket_config.compute_bucket_size_limits(
+            static_graph=static_graph,
+            find_unused_parameters=self.find_unused_parameters,
+        )
+
         (
             bucket_indices,
             per_bucket_size_limits,
@@ -1267,16 +1397,13 @@ def _ddp_init_helper(
             self.find_unused_parameters,
             self.gradient_as_bucket_view,
             param_to_name_mapping,
-            # User can set dist._DEFAULT_FIRST_BUCKET_BYTES to tune DDP first
-            # bucket.
-            (
-                dist._DEFAULT_FIRST_BUCKET_BYTES
-                if self.bucket_bytes_cap_default
-                else self.bucket_bytes_cap
-            ),
+            # First bucket can be smaller for reduced latency. See
+            # BucketCapacityConfig.first_bucket_bytes_cap for details.
+            self._bucket_config.first_bucket_bytes_cap,
             self.skip_all_reduce_unused_params,
             self._use_python_reducer,
             bucket_size_limits_for_rebuilding,
+            self.batched_grad_copy,
         )
 
         self.logger = dist.Logger(self.reducer)
@@ -1510,11 +1637,17 @@ def _get_active_ddp_module(cls):
     @contextmanager
     @torch._disable_dynamo(recursive=False)
     def _inside_ddp_forward(self):
+        # Save and restore the previous _active_ddp_module to handle nested
+        # DDP correctly (e.g., TorchRec wraps embeddings in an inner DDP inside
+        # an outer DDP).  Without this, the inner DDP's exit would clear the
+        # flag to None, causing DDPOptimizer to miss compiled regions that run
+        # after the inner forward.
+        old = DistributedDataParallel._active_ddp_module
         DistributedDataParallel._active_ddp_module = self
         try:
             yield
         finally:
-            DistributedDataParallel._active_ddp_module = None
+            DistributedDataParallel._active_ddp_module = old
 
     def _run_ddp_forward(self, *inputs, **kwargs):
         if self._use_python_reducer:
@@ -2228,7 +2361,10 @@ def _sync_buffers(self):
             # Update self.modules_buffers in case any buffers were
             # reassigned.
             self._assign_modules_buffers()
-            self._sync_module_buffers(authoritative_rank)
+            with torch.autograd._unsafe_preserve_version_counter(
+                tuple(self.modules_buffers)
+            ):
+                self._sync_module_buffers(authoritative_rank)
 
     def _sync_module_buffers(self, authoritative_rank):
         if not hasattr(self, "buffer_hook"):
diff --git a/torch/nn/parallel/scatter_gather.py b/torch/nn/parallel/scatter_gather.py
index 8954302f2a44e..bcd975a4fe2f7 100644
--- a/torch/nn/parallel/scatter_gather.py
+++ b/torch/nn/parallel/scatter_gather.py
@@ -58,19 +58,19 @@ def scatter_map(obj):
         if _is_namedtuple(obj):
             return [
                 type(obj)(*args)
-                # pyrefly: ignore [no-matching-overload]
+                # pyrefly: ignore [bad-argument-type, no-matching-overload]
                 for args in zip(*map(scatter_map, obj), strict=False)
             ]
         if isinstance(obj, tuple) and len(obj) > 0:
-            # pyrefly: ignore [no-matching-overload]
+            # pyrefly: ignore [bad-argument-type, no-matching-overload]
             return list(zip(*map(scatter_map, obj), strict=False))
         if isinstance(obj, list) and len(obj) > 0:
-            # pyrefly: ignore [no-matching-overload]
+            # pyrefly: ignore [bad-argument-type, no-matching-overload]
             return [list(i) for i in zip(*map(scatter_map, obj), strict=False)]
         if isinstance(obj, dict) and len(obj) > 0:
             return [
                 type(obj)(i)
-                # pyrefly: ignore [no-matching-overload]
+                # pyrefly: ignore [bad-argument-type, no-matching-overload]
                 for i in zip(*map(scatter_map, obj.items()), strict=False)
             ]
         return [obj for _ in target_gpus]
@@ -136,9 +136,9 @@ def gather_map(outputs):
             # pyrefly: ignore [not-callable]
             return type(out)((k, gather_map([d[k] for d in outputs])) for k in out)
         if _is_namedtuple(out):
-            # pyrefly: ignore [no-matching-overload]
+            # pyrefly: ignore [bad-argument-type]
             return type(out)._make(map(gather_map, zip(*outputs, strict=True)))
-        # pyrefly: ignore [no-matching-overload]
+        # pyrefly: ignore [bad-argument-type]
         return type(out)(map(gather_map, zip(*outputs, strict=True)))
 
     # Recursive function calls like this create reference cycles.
diff --git a/torch/nn/parameter.py b/torch/nn/parameter.py
index a832c5323cc3c..f5af0581d61fe 100644
--- a/torch/nn/parameter.py
+++ b/torch/nn/parameter.py
@@ -101,6 +101,7 @@ def __reduce_ex__(self, proto):
             (self.data, self.requires_grad, hooks, state),
         )
 
+    # pyrefly: ignore [bad-override]
     __torch_function__ = _disabled_torch_function_impl
 
 
@@ -144,7 +145,7 @@ def materialize(self, shape, device=None, dtype=None) -> None:
         if dtype is None:
             dtype = self.data.dtype
         self.data = torch.empty(shape, device=device, dtype=dtype)
-        # pyrefly: ignore [bad-override, missing-attribute]
+        # pyrefly: ignore [missing-attribute]
         self.__class__ = self.cls_to_become
 
     @property
@@ -273,6 +274,7 @@ def __new__(cls, data=None, *, persistent=True):
         t._is_buffer = True
         return t
 
+    # pyrefly: ignore [bad-override]
     __torch_function__ = _disabled_torch_function_impl
 
 
diff --git a/torch/nn/utils/_expanded_weights/expanded_weights_impl.py b/torch/nn/utils/_expanded_weights/expanded_weights_impl.py
index 8c3a1a5c9d7aa..f3a5999f392c3 100644
--- a/torch/nn/utils/_expanded_weights/expanded_weights_impl.py
+++ b/torch/nn/utils/_expanded_weights/expanded_weights_impl.py
@@ -131,7 +131,7 @@ def __torch_function__(cls, func, _, args=(), kwargs=None):
             # in aten, choosing the input or data variants is done by parsing logic. This mimics some of that
             decomp_opts = expanded_weights_rnn_decomps[func]
             use_input_variant = isinstance(
-                # pyrefly: ignore [bad-index, index-error]
+                # pyrefly: ignore [bad-index]
                 args[2],
                 list,
             )  # data variant uses a list here
diff --git a/torch/nn/utils/clip_grad.py b/torch/nn/utils/clip_grad.py
index 21667f20d2074..c7b2d6fcddfca 100644
--- a/torch/nn/utils/clip_grad.py
+++ b/torch/nn/utils/clip_grad.py
@@ -38,7 +38,6 @@ def _no_grad(func: Callable[_P, _R]) -> Callable[_P, _R]:
 
     def _no_grad_wrapper(*args, **kwargs):
         with torch.no_grad():
-            # pyrefly: ignore [invalid-param-spec]
             return func(*args, **kwargs)
 
     functools.update_wrapper(_no_grad_wrapper, func)
diff --git a/torch/nn/utils/parametrize.py b/torch/nn/utils/parametrize.py
index de25bba41eda1..81a3ad796abeb 100644
--- a/torch/nn/utils/parametrize.py
+++ b/torch/nn/utils/parametrize.py
@@ -9,6 +9,8 @@
 import torch
 from torch import Tensor
 from torch.__future__ import get_swap_module_params_on_conversion
+from torch._library.opaque_object import is_opaque_reference_type
+from torch._opaque_base import OpaqueBase
 from torch.nn.modules.container import Module, ModuleDict, ModuleList
 from torch.nn.parameter import Parameter
 from torch.utils._python_dispatch import is_traceable_wrapper_subclass
@@ -204,20 +206,29 @@ def __init__(
             _register_parameter_or_buffer(self, "original", original)
         else:
             for i, originali in enumerate(new):
-                if not isinstance(originali, Tensor):
-                    raise ValueError(
-                        "'right_inverse' must return a Tensor or a Sequence of tensors "
-                        "(list, tuple...). "
-                        f"Got element {i} of the sequence with type {type(originali).__name__}."
-                    )
-
-                # If the original tensor was a Parameter that required grad, we expect the user to
-                # add the new parameters to the optimizer after registering the parametrization
-                # (this is documented)
-                if isinstance(original, Parameter):
-                    originali = Parameter(originali, original.requires_grad)
-                originali.requires_grad_(original.requires_grad)
-                _register_parameter_or_buffer(self, f"original{i}", originali)
+                match originali:
+                    case OpaqueBase():
+                        if not is_opaque_reference_type(type(originali)):
+                            raise ValueError(
+                                f"'right_inverse' must return a Tensor or a reference-type "
+                                f"opaque. Got element {i} of the sequence with type "
+                                f"{type(originali).__name__}."
+                            )
+                        setattr(self, f"original{i}", originali)
+                    case Tensor():
+                        # If the original tensor was a Parameter that required grad, we expect the user to
+                        # add the new parameters to the optimizer after registering the parametrization
+                        # (this is documented)
+                        if isinstance(original, Parameter):
+                            originali = Parameter(originali, original.requires_grad)
+                        originali.requires_grad_(original.requires_grad)
+                        _register_parameter_or_buffer(self, f"original{i}", originali)
+                    case _:
+                        raise ValueError(
+                            "'right_inverse' must return a Tensor or a Sequence of tensors "
+                            "(list, tuple...). "
+                            f"Got element {i} of the sequence with type {type(originali).__name__}."
+                        )
 
         if not self.unsafe:
             # Consistency checks:
@@ -292,17 +303,25 @@ def right_inverse(self, value: Tensor) -> None:
                     )
                 for i, tensor in enumerate(value):
                     original_i = getattr(self, f"original{i}")
-                    if not isinstance(tensor, Tensor):
-                        raise ValueError(
-                            f"`right_inverse` must return a sequence of tensors. "
-                            f"Got element {i} of type {type(tensor).__name__}"
-                        )
-                    if original_i.dtype != tensor.dtype:
-                        raise ValueError(
-                            f"Tensor {i} returned by `right_inverse` has dtype {tensor.dtype} "
-                            f"while `original{i}` has dtype {original_i.dtype}"
-                        )
-                    _maybe_set(original_i, tensor)
+                    match tensor:
+                        case OpaqueBase():
+                            if is_opaque_reference_type(type(tensor)):
+                                setattr(self, f"original{i}", tensor)
+                                continue
+                            # Fall-through
+                        case Tensor():
+                            if original_i.dtype != tensor.dtype:
+                                raise ValueError(
+                                    f"Tensor {i} returned by `right_inverse` has dtype {tensor.dtype} "
+                                    f"while `original{i}` has dtype {original_i.dtype}"
+                                )
+                            _maybe_set(original_i, tensor)
+                            continue
+                    raise ValueError(
+                        f"'right_inverse' must return a sequence of tensors "
+                        f"or reference-type opaques. Got element {i} of type "
+                        f"{type(tensor).__name__}."
+                    )
 
     def forward(self) -> Tensor:
         if torch.jit.is_scripting():
diff --git a/torch/nn/utils/spectral_norm.py b/torch/nn/utils/spectral_norm.py
index b785e080b79e6..79fea5a50b793 100644
--- a/torch/nn/utils/spectral_norm.py
+++ b/torch/nn/utils/spectral_norm.py
@@ -185,8 +185,6 @@ def apply(
         return fn
 
 
-# This is a top level class because Py2 pickle doesn't like inner class nor an
-# instancemethod.
 class SpectralNormLoadStateDictPreHook:
     # See docstring of SpectralNorm._version on the changes to spectral_norm.
     def __init__(self, fn) -> None:
@@ -244,8 +242,6 @@ def __call__(
                 state_dict[weight_key + "_v"] = v
 
 
-# This is a top level class because Py2 pickle doesn't like inner class nor an
-# instancemethod.
 class SpectralNormStateDictHook:
     # See docstring of SpectralNorm._version on the changes to spectral_norm.
     def __init__(self, fn) -> None:
diff --git a/torch/onnx/_internal/exporter/_building.py b/torch/onnx/_internal/exporter/_building.py
index 557ab0d97e559..52e0e215fea2c 100644
--- a/torch/onnx/_internal/exporter/_building.py
+++ b/torch/onnx/_internal/exporter/_building.py
@@ -14,7 +14,7 @@
 import inspect
 import logging
 from collections.abc import Iterable, Mapping, Sequence
-from typing import Any, TYPE_CHECKING, Union
+from typing import Any, TYPE_CHECKING
 
 from onnxscript import evaluator
 
@@ -29,13 +29,18 @@
 
 logger = logging.getLogger(__name__)
 
-ValidAttributeType = Union[
-    ir.TensorProtocol, int, float, bool, str, Sequence[int], Sequence[float], None
-]
+ValidAttributeType = (
+    ir.TensorProtocol
+    | int
+    | float
+    | bool
+    | str
+    | Sequence[int]
+    | Sequence[float]
+    | None
+)
 
-AllowedArgType = Union[
-    ir.Value, Sequence[Union[ir.Value, ValidAttributeType]], ValidAttributeType
-]
+AllowedArgType = ir.Value | Sequence[ir.Value | ValidAttributeType] | ValidAttributeType
 
 
 # Logic for adapting inputs from general Python or PyTorch inputs to ONNX ir.Value
@@ -271,7 +276,6 @@ def _get_or_create_constant(
     # float representation of complex numbers
     if isinstance(arg, complex):
         # Convert the complex number to a float
-        # pyrefly: ignore [bad-assignment]
         arg = (arg.real, arg.imag)
 
     if isinstance(arg, list):
@@ -440,7 +444,6 @@ def _process_python_sequences(
             # when the expected input type is INT64
             # We assume this only happens for 0D cases
             if all(isinstance(val, ir.Value) for val in arg):
-                # pyrefly: ignore
                 expanded_args = [_reshape_to_1d_tensor(opset, val) for val in arg]
                 named_inputs[name] = opset.Concat(*expanded_args, axis=0)
                 continue
diff --git a/torch/onnx/_internal/exporter/_core.py b/torch/onnx/_internal/exporter/_core.py
index f26b1779e7ed6..3b1e1dfc21620 100644
--- a/torch/onnx/_internal/exporter/_core.py
+++ b/torch/onnx/_internal/exporter/_core.py
@@ -133,10 +133,8 @@ def numpy(self) -> npt.NDArray:
         # view the tensor as that dtype so that it is convertible to NumPy,
         # and then view it back to the proper dtype (using ml_dtypes obtained by
         # calling dtype.numpy()).
-        # pyrefly: ignore [missing-attribute]
         if self.dtype == ir.DataType.BFLOAT16:
             return (
-                # pyrefly: ignore [missing-attribute]
                 self.raw.view(torch.uint16).numpy(force=True).view(self.dtype.numpy())
             )
         if self.dtype in {
@@ -145,11 +143,9 @@ def numpy(self) -> npt.NDArray:
             ir.DataType.FLOAT8E5M2,
             ir.DataType.FLOAT8E5M2FNUZ,
         }:
-            # pyrefly: ignore [missing-attribute]
             return self.raw.view(torch.uint8).numpy(force=True).view(self.dtype.numpy())
         if self.dtype == ir.DataType.FLOAT4E2M1:
             return _type_casting.unpack_float4x2_as_uint8(self.raw).view(
-                # pyrefly: ignore [missing-attribute]
                 self.dtype.numpy()
             )
 
@@ -171,7 +167,6 @@ def _get_cbytes(self):
 
         if isinstance(tensor, torch._subclasses.fake_tensor.FakeTensor):
             raise TypeError(
-                # pyrefly: ignore [missing-attribute]
                 f"Cannot take content out from the FakeTensor ('{self.name}'). Please replace the tensor "
                 "with a tensor backed by real data using ONNXProgram.apply_weights() "
                 "or save the model without initializers by setting include_initializers=False."
diff --git a/torch/onnx/_internal/exporter/_input_observer.py b/torch/onnx/_internal/exporter/_input_observer.py
index 2a4392919e120..820db724f288a 100644
--- a/torch/onnx/_internal/exporter/_input_observer.py
+++ b/torch/onnx/_internal/exporter/_input_observer.py
@@ -27,11 +27,11 @@ def _flatten_unflatten_for_dynamic_shapes(
 
     Args:
         obj: Object from a custom class.
-        change_function: Function to modify the tensor in the structure itself,
-            like replace them by a shape.
+        change_function: If not None, this function is called to modify the tensors
+            in the structure itself, like replace them by a shape.
 
     Returns:
-        the serialized object
+        The flattened object.
     """
     if isinstance(obj, torch.Tensor):
         return change_function(obj) if change_function else obj
@@ -77,10 +77,10 @@ def _infer_dynamic_dimensions(
 
     Args:
         shape_list:
-            list of shapes, they must all have the same length
+            List of shapes, they must all have the same length.
         set_batch_dimension:
-            forces the first dimension to be treated as dynamic,
-            even if all shapes have the same value for that dimension
+            Forces the first dimension to be treated as dynamic,
+            even if all shapes have the same value for that dimension.
 
     Returns:
         list of dynamic dimensions
@@ -88,7 +88,7 @@ def _infer_dynamic_dimensions(
     unique_ranks = {len(shape) for shape in shape_list}
     torch._check(
         len(unique_ranks) == 1,
-        lambda: "all shapes in shape_list must have the same rank",
+        lambda: f"All shapes in shape_list must have the same rank but {shape_list=}.",
     )
     rank = unique_ranks.pop()
     dynamic = []
@@ -102,6 +102,7 @@ def _infer_dynamic_dimensions(
 class InputCandidate:
     """Retains one set of inputs given to the forward method or any
     other method the class :class:`InputObserver` is stealing from.
+    Any class is allowed as long as it can be flattened.
 
     Args:
         args: Positional arguments.
@@ -110,6 +111,9 @@ class InputCandidate:
             may be modified inplace, the original value must be retained.
         cst_kwargs: Any optional arguments constant over multiple calls.
             int, float, str, bool values must be stored here.
+
+    The constructor flattens the received arguments.
+    Any necessary flattening function should have been registered first.
     """
 
     def __init__(
@@ -283,18 +287,36 @@ class InputObserverInfo:
             to be the same in the ordered dictionaries `add_inputs` receive.
         default_values: Default values defined by the signature of the function,
             any value equal to that is ignored to simplify the export.
+        value_if_missing: If an argument is missing,
+            a default value will be taken in this dictionary,
+            this is used when after the prefill step, an argument
+            disappears (such as `pixel_values`) and another one
+            is added (such as `past_key_values`).
+            The values are only to infer dynamic shapes and arguments,
+            not to run the model.
+        args_name_and_position: Name of parameter `*args`
+            and its position if it exists.
+        kwargs_name: Name of the variable keyword parameter `**kwargs` if it exists.
+
+    This is used by class :class:`InputObserver`.
     """
 
     def __init__(
         self,
         signature_names: list[str],
         default_values: dict[str, int | bool | str | float],
+        value_if_missing: dict[str | int, Any],
+        args_name_and_position: tuple[str, int] | None,
+        kwargs_name: str | None,
     ):
         self.default_values = default_values
+        self.value_if_missing = value_if_missing
         self.inputs: list[InputCandidate] = []
         self.outputs_specs: list[torch.utils._pytree.PyTreeSpec] = []
         self.flat_outputs: list[list[torch.Tensor | None]] = []
         self.latencies: list[float] = []
+        self.args_name_and_position = args_name_and_position
+        self.kwargs_name = kwargs_name
         self.signature_names = signature_names
         self._best_candidate: InputCandidate | None = None
         self._captured_inputs: dict[int | str, int] | None = None
@@ -316,6 +338,7 @@ def add_inputs(self, args: tuple[Any, ...], kwargs: dict[str, Any]):
             if k in self.signature_names
             and isinstance(v, (int, float, bool, str))
             and v != self.default_values.get(k, None)
+            and self.default_values.get(k, None) is not None
         }
         kwargs = {
             k: v
@@ -323,6 +346,50 @@ def add_inputs(self, args: tuple[Any, ...], kwargs: dict[str, Any]):
             if v is not None and not isinstance(v, (int, float, bool, str))
         }
 
+        # adds value_if_missing attributes
+        for k, v in self.value_if_missing.items():
+            if isinstance(k, str):
+                if k not in kwargs:
+                    # Validate that `value_if_missing` keys are compatible
+                    # with the observed signature.
+                    # If the function does not accept **kwargs,
+                    # all value_if_missing keys must be
+                    # present in the observed signature names.
+                    if k not in self.signature_names and not self.kwargs_name:
+                        raise ValueError(
+                            f"Unexpected keyword argument {k!r} "
+                            f"provided as a value_if_missing input "
+                            "for a function that does not accept it. "
+                            f"All value_if_missing keys must "
+                            f"be in the observed signature: {tuple(self.signature_names)}."
+                        )
+                    kwargs[k] = v
+            elif isinstance(k, int):
+                if k >= len(self.signature_names):
+                    raise ValueError(
+                        f"Unexpected keyword argument {k=} "
+                        f"provided as a value_if_missing input "
+                        "for a function that does not accept it. "
+                        f"All value_if_missing indices must "
+                        f"be in the observed signature: {tuple(self.signature_names)}."
+                    )
+                if k >= len(args):
+                    raise NotImplementedError(
+                        f"Unexpected keyword argument {k=} "
+                        f"provided as a value_if_missing input "
+                        "for a function that does not accept it. "
+                        f"All value_if_missing indices must "
+                        f"be in the observed signature: {tuple(self.signature_names)}, "
+                        f"only {len(args)} were given."
+                    )
+                list_args = list(args)
+                list_args[k] = v
+                args = tuple(list_args)
+            else:
+                raise TypeError(
+                    f"Unexpected type {type(k)} for a missing value. The key is {k!r}."
+                )
+
         # kwargs may come in a different order each time.
         # dictionaries are ordered and torch.export.export expects
         # dynamic shapes and kwargs to follow the same order.
@@ -458,27 +525,61 @@ def _set_batch_dimension_for_flat_index(index) -> bool:
         flat_dynamic_shapes = [dict.fromkeys(dims, cst) for dims in dynamic_shapes]
         if return_flat:
             return tuple(flat_dynamic_shapes)
+
+        # Let's regroup.
         if len(flat_dynamic_shapes) == len(self._best_candidate.args) + len(
             self._best_candidate.kwargs
         ):
             # It means forward method is called with tensors only.
-            if not self._best_candidate.kwargs and not self._best_candidate.cst_kwargs:
+            if (
+                not self._best_candidate.kwargs
+                and not self._best_candidate.cst_kwargs
+                and not self.args_name_and_position
+            ):
                 # only positional arguments
                 return tuple(flat_dynamic_shapes)
             if not self._best_candidate.args:
                 # only named arguments
                 ds = dict(zip(list(self._best_candidate.kwargs), flat_dynamic_shapes))
-                return {**ds, **dict.fromkeys(self._best_candidate.cst_kwargs, None)}
+                return self._post_process_for_kwargs(
+                    {**ds, **dict.fromkeys(self._best_candidate.cst_kwargs, None)}
+                )
+            if not self.args_name_and_position:
+                # positional arguments needs to be moved to the named arguments
+                n_args = len(self._best_candidate.args)
+                pos_names = self.signature_names[:n_args]
+                return self._post_process_for_kwargs(
+                    {
+                        **dict(zip(pos_names, flat_dynamic_shapes[:n_args])),
+                        **dict(
+                            zip(
+                                list(self._best_candidate.kwargs),
+                                flat_dynamic_shapes[n_args:],
+                            )
+                        ),
+                        **dict.fromkeys(self._best_candidate.cst_kwargs, None),
+                    }
+                )
             # positional arguments needs to be moved to the named arguments
-            n_args = len(self._best_candidate.args)
+            n_args = min(len(self._best_candidate.args), self.args_name_and_position[1])
+            i_kwargs = max(
+                len(self._best_candidate.args), self.args_name_and_position[1]
+            )
+            var_pos = self.args_name_and_position[0]
             pos_names = self.signature_names[:n_args]
-            return {
-                **dict(zip(pos_names, flat_dynamic_shapes[:n_args])),
-                **dict(
-                    zip(list(self._best_candidate.kwargs), flat_dynamic_shapes[n_args:])
-                ),
-                **dict.fromkeys(self._best_candidate.cst_kwargs, None),
-            }
+            return self._post_process_for_kwargs(
+                {
+                    **dict(zip(pos_names, flat_dynamic_shapes[:n_args])),
+                    var_pos: tuple(flat_dynamic_shapes[n_args:i_kwargs]),
+                    **dict(
+                        zip(
+                            list(self._best_candidate.kwargs),
+                            flat_dynamic_shapes[i_kwargs:],
+                        )
+                    ),
+                    **dict.fromkeys(self._best_candidate.cst_kwargs, None),
+                }
+            )
 
         # nested types, here comes the fun part because the shapes cannot be unflattened,
         # custom classes must appear in their flattened shape.
@@ -518,20 +619,62 @@ def change_function(t):
                 **ds_kwargs,
                 **dict.fromkeys(self._best_candidate.cst_kwargs, None),
             }
-        if not ds_kwargs:
+        if not ds_kwargs and not self.args_name_and_position:
             return tuple(ds_args)
         if not ds_args:
-            return ds_kwargs
-        pos_names = self.signature_names[: len(ds_args)]
-        return {**dict(zip(pos_names, ds_args)), **ds_kwargs}
+            return self._post_process_for_kwargs(ds_kwargs)
+
+        if not self.args_name_and_position:
+            pos_names = self.signature_names[: len(ds_args)]
+            return self._post_process_for_kwargs(
+                {**dict(zip(pos_names, ds_args)), **ds_kwargs}
+            )
+
+        n_args = min(len(ds_args), self.args_name_and_position[1])
+        pos_names = self.signature_names[:n_args]
+        return self._post_process_for_kwargs(
+            {
+                **dict(zip(pos_names, ds_args[:n_args])),
+                self.args_name_and_position[0]: tuple(ds_args[n_args:]),
+                **ds_kwargs,
+            }
+        )
 
     def infer_arguments(
         self,
         index_or_candidate: InputCandidate | int | None = None,
         /,
         flat: bool = False,
-    ) -> list[torch.Tensor | None] | tuple[torch.Tensor, ...] | dict[str, torch.Tensor]:
-        """Infers arguments based on the collected tensors."""
+        as_args_kwargs: bool = False,
+    ) -> (
+        list[torch.Tensor | None]
+        | tuple[torch.Tensor, ...]
+        | dict[str, torch.Tensor]
+        | tuple[list[torch.Tensor] | tuple[torch.Tensor, ...], dict[str, torch.Tensor]]
+    ):
+        """Infers arguments based on the collected tensors.
+
+        Args:
+            index_or_candidate: If missing, the method selects one set of inputs
+                among the available ones, usually the set of inputs containing
+                with the highest number of tensors.
+                It then replaces None values and missing tensors with empty tensors.
+                If not missing, it can be an integer to fetch one of the stored set
+                or some inputs.
+            flat: If True, it returns a flattened list of tensors,
+                if False, it returns a tuple or a dictionary preserving
+                the nested structures. The flat version is used internally.
+                It produces a single list of tensors easier to process or modify
+                rather than a nested structure holding the same tensors.
+                The original structure can be restored with
+                ``torch.utils._pytree.tree_unflatten(flat_list, self.aligned_spec)``.
+                This mechanism is used to replace None values by empty tensors.
+            as_args_kwargs: If True, the method always returns `(args, kwargs)`,
+                otherwise, it returns either a tuple (only args) or a dictionary
+                (only kwargs) or raises an exception if it cannot do so.
+        Returns:
+            Inferred arguments, every optional tensor is replaced by an empty tensor.
+        """
         # This is already checked by _build_inputs_completed_with_none_values
         # but this is not always well captured by tools checking types.
         self.align_inputs_none_values()
@@ -540,9 +683,9 @@ def infer_arguments(
         if index_or_candidate is None:
             for cand in self.inputs:
                 args, kwargs = cand.args, cand.kwargs
-                if len(args) == len(self._best_candidate.args) and len(kwargs) == len(
-                    self._best_candidate.kwargs
-                ):
+                if len(args) == len(self._best_candidate.args or ()) and len(
+                    kwargs
+                ) == len(self._best_candidate.kwargs or {}):
                     candidate = cand
                     break
         elif isinstance(index_or_candidate, int):
@@ -592,7 +735,6 @@ def infer_arguments(
                     aligned_flat_list[index] = tensor
                     continue
                 if not shape:
-                    # pyrefly: ignore[unsupported-operation]
                     aligned_flat_list[index] = torch.zeros(
                         tensor.shape, dtype=tensor.dtype, device=tensor.device
                     )
@@ -607,7 +749,6 @@ def infer_arguments(
                 )
                 new_shape = list(tensor.shape)
                 new_shape[dim] = 0
-                # pyrefly: ignore[unsupported-operation]
                 aligned_flat_list[index] = torch.empty(
                     tuple(new_shape), dtype=tensor.dtype, device=tensor.device
                 )
@@ -619,16 +760,52 @@ def infer_arguments(
             candidate.aligned_spec,
         )
         if self._best_candidate.cst_kwargs:
-            # pyrefly: ignore[invalid-argument]
             kwargs = {**kwargs, **self._best_candidate.cst_kwargs}
 
-        if not kwargs:
-            return args
-        if not args:
+        if not as_args_kwargs:
+            if not kwargs:
+                return args
+            if not args:
+                return kwargs
+
+            # We need to move args to kwargs
+            if self.args_name_and_position:
+                raise RuntimeError(
+                    "Cannot return arguments "
+                    "as a single tuple or a single dictionary "
+                    "because of '*args' in the function signature. "
+                    "You need to set `as_args_kwargs=True`."
+                )
+            n_args = len(args)
+            pos_names = self.signature_names[:n_args]
+            return {**dict(zip(pos_names, args[:n_args])), **kwargs}
+
+        # Generic case.
+        return tuple(args), kwargs
+
+    def _post_process_for_kwargs(self, kwargs: dict[str, Any]) -> dict[str, Any]:
+        """:func:`torch.export.export` requires dynamic shapes and keyword arguments
+        that are not part of the explicit function signature but are collected via
+        ``**<kwargs_name>`` to be wrapped under the corresponding parameter name
+        (``self.kwargs_name``) as ``{<kwargs_name>: {'param': shape or tensor}}``.
+        This function ensures this wrapping is performed when ``self.kwargs_name`` is set.
+        """
+        if not self.kwargs_name:
+            # Nothing to do here.
             return kwargs
-        # We need to move args to kwargs
-        pos_names = self.signature_names[: len(args)]
-        return {**dict(zip(pos_names, args)), **kwargs}
+        to_be_moved = {k for k in kwargs if k not in self.signature_names}
+        if not to_be_moved:
+            return kwargs
+        keywords = {k: v for k, v in kwargs.items() if k in to_be_moved}
+        new_kwargs = {k: v for k, v in kwargs.items() if k not in to_be_moved}
+        if self.kwargs_name in new_kwargs:
+            raise ValueError(
+                f"Keyword argument name collision: received a keyword argument "
+                f"'{self.kwargs_name}' which conflicts with the **{self.kwargs_name} "
+                "parameter used to collect extra keyword arguments. "
+                "Passing a keyword argument with this name is not supported."
+            )
+        return {**new_kwargs, self.kwargs_name: keywords}
 
 
 class InputObserver:
@@ -636,6 +813,15 @@ class InputObserver:
     This information is used to infer dynamic shapes and
     export arguments.
 
+    Args:
+        value_if_missing: If an argument is missing,
+            a default value will be taken in this dictionary,
+            this is used when after the prefill step, an argument
+            disappears (such as `pixel_values`) and another one
+            is added (such as `past_key_values`).
+            The values are only to infer dynamic shapes and arguments,
+            not to run the model.
+
     Examples
     --------
     >>> input_observer = InputObserver()
@@ -660,11 +846,59 @@ class InputObserver:
     >>>     dynamic_shapes.input_observer.infer_dynamic_shapes(),
     >>> )
 
+    The last example considers an LLM taking images and text as inputs.
+    The first call to the forward method which we try to export has `pixel_values`
+    but no `past_key_values`. The next calls do not have `pixel_values` but
+    `past_key_values`. The observer understands `pixel_values` and `past_key_values`
+    are needed but they may not be both specified at the same time.
+    Since `pixel_values` only appears in the first call, the observer cannot
+    tell how to infer an empty tensor for this argument. That's what the argument
+    `value_if_missing` is for. The following example is more than a dummy example
+    but shows how to use it with ``transformers``.
+
+    .. code-block:: python
+
+        from transformers import pipeline
+
+        model_id = "tiny-random/gemma-3"
+        pipe = pipeline(
+            "image-text-to-text",
+            model=model_id,
+            device="cpu",
+            trust_remote_code=True,
+            max_new_tokens=3,
+            dtype=torch.float16,
+        )
+        messages = [
+            {
+                "role": "system",
+                "content": [{"type": "text", "text": "You are a helpful assistant."}],
+            },
+            {
+                "role": "user",
+                "content": [
+                    {
+                        "type": "image",
+                        "url": "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/p-blog/candy.JPG",
+                    },
+                    {"type": "text", "text": "What animal is on the candy?"},
+                ],
+            },
+        ]
+        observer = InputObserver(
+            value_if_missing=dict(
+                pixel_values=torch.empty((0, 3, 896, 896), dtype=torch.float16)
+            )
+        )
+        with observer(pipe.model):
+            pipe(text=messages, max_new_tokens=4)
+
     .. versionadded:: 2.11.0
     """
 
-    def __init__(self):
+    def __init__(self, value_if_missing: dict[str | int, Any] | None = None):
         self.info: InputObserverInfo | None = None
+        self.value_if_missing = value_if_missing or {}
 
     def _replaced_method(
         self,
@@ -716,6 +950,16 @@ def __call__(
         captured_method = getattr(model, method_name)
         sig = inspect.signature(captured_method)
         if self.info is None:
+            kwargs_names = [
+                p
+                for p in sig.parameters
+                if sig.parameters[p].kind == inspect.Parameter.VAR_KEYWORD
+            ]
+            args_names = [
+                (p, i)
+                for (i, p) in enumerate(sig.parameters)
+                if sig.parameters[p].kind == inspect.Parameter.VAR_POSITIONAL
+            ]
             self.info = InputObserverInfo(
                 signature_names=list(sig.parameters),
                 default_values={
@@ -724,6 +968,9 @@ def __call__(
                     if p.default != inspect.Parameter.empty
                     and isinstance(p.default, (int, bool, str, float))
                 },
+                value_if_missing=self.value_if_missing,
+                args_name_and_position=args_names[0] if args_names else None,
+                kwargs_name=kwargs_names[0] if kwargs_names else None,
             )
         n_already_stored = len(self.info)
         lambda_method = lambda *args, _cm=captured_method, _snc=(  # noqa: E731
@@ -777,7 +1024,13 @@ def infer_arguments(
         self,
         index_or_args_or_kwargs: tuple[Any] | dict[str, Any] | int | None = None,
         flat: bool = False,
-    ) -> list[torch.Tensor | None] | tuple[torch.Tensor, ...] | dict[str, torch.Tensor]:
+        as_args_kwargs: bool = False,
+    ) -> (
+        list[torch.Tensor | None]
+        | tuple[torch.Tensor, ...]
+        | dict[str, torch.Tensor]
+        | tuple[list[torch.Tensor] | tuple[torch.Tensor, ...], dict[str, torch.Tensor]]
+    ):
         """Infers arguments based on the collected tensors.
 
         Args:
@@ -789,8 +1042,15 @@ def infer_arguments(
                 or some inputs.
             flat: If True, it returns a flattened list of tensors,
                 if False, it returns a tuple or a dictionary preserving
-                the nested structures.
-
+                the nested structures. The flat version is used internally.
+                It produces a single list of tensors easier to process or modify
+                rather than a nested structure holding the same tensors.
+                The original structure can be restored with
+                ``torch.utils._pytree.tree_unflatten(flat_list, self.aligned_spec)``.
+                This mechanism is used to replace None values by empty tensors.
+            as_args_kwargs: If True, the method always returns `(args, kwargs)`,
+                otherwise, it returns either a tuple (only args) or a dictionary
+                (only kwargs) or raises an exception if it cannot do so.
         Returns:
             Inferred arguments, every optional tensor is replaced by an empty tensor.
         """
@@ -832,7 +1092,11 @@ def infer_arguments(
                 self.info._captured_inputs,
                 self.info.signature_names,
             )
-        return self.info.infer_arguments(index_or_candidate, flat=flat)
+        return self.info.infer_arguments(
+            index_or_candidate,
+            flat=flat,
+            as_args_kwargs=as_args_kwargs,
+        )
 
     def check_discrepancies(
         self,
@@ -843,25 +1107,26 @@ def check_discrepancies(
         initializer: Callable[
             [str | bytes], ort.InferenceSession
         ] = _onnx_program._ort_session_initializer,
-    ) -> list[dict[str, str | int | float]]:
+        skip_none: bool = True,
+    ) -> list[dict[str, str | int | float | bool]]:
         """Computes the discrepancies between the saved inputs and outputs
         with the saved onnx model.
 
         Args:
-            onnx_program:
-                Exported Model to verify.
-            atol:
-                Absolute tolerance, recommended values, 1e-4 for float, 1e-2 for float16.
-            rtol:
-                Relative tolerance.
-            progress_bar:
-                Shows a progress bar (requires `tqdm`).
-            initializer: The function to initialize the ONNX Runtime inference
+            onnx_program: Exported Model to verify.
+            atol: Absolute tolerance, recommended values, 1e-4 for float, 1e-2 for float16.
+            rtol: Relative tolerance.
+            progress_bar: Shows a progress bar (requires `tqdm`).
+            initializer: The function called to initialize the ONNX Runtime inference
                 session with the specified model. By default, it uses the
                 `_ort_session_initializer` function.
+            skip_none: Does not check discrepancies when an output is None.
 
         Returns:
             A list of dictionaries, ready to be consumed by a dataframe.
+
+        The function catches exceptions, it shows the error in the returned
+        summary.
         """
         # For big models, we should consider taking a filename to avoid the users
         # creating the model proto twice.
@@ -895,7 +1160,6 @@ def check_discrepancies(
 
             begin = time.perf_counter()
             try:
-                # pyrefly: ignore[bad-unpacking]
                 ort_outputs = onnx_program(**feeds)
                 error = None
             except Exception as e:
@@ -904,7 +1168,9 @@ def check_discrepancies(
 
             duration = time.perf_counter() - begin
             if error:
-                diff: dict[str, Any] = dict(error=error, SUCCESS=False)
+                diff: dict[str, str | int | float | bool] = dict(
+                    error=error, SUCCESS=False
+                )
             elif ort_outputs is None or len(outputs) != len(ort_outputs):
                 diff = dict(SUCCESS=False, error="not the same number of outputs")
             else:
@@ -912,10 +1178,9 @@ def check_discrepancies(
                 err_abs = 0.0
                 err_rel = 0.0
                 error = ""
-                # pyrefly: ignore[no-matching-overload]
                 for torch_tensor, ort_tensor in zip(outputs, ort_outputs):
                     if torch_tensor is None or ort_tensor is None:
-                        if type(torch_tensor) is not type(ort_tensor):
+                        if type(torch_tensor) is not type(ort_tensor) and not skip_none:
                             success = False
                             error = "missing output"
                             break
diff --git a/torch/onnx/_internal/exporter/_ir_passes.py b/torch/onnx/_internal/exporter/_ir_passes.py
index 8e5f15e456910..ee8fca22437e0 100644
--- a/torch/onnx/_internal/exporter/_ir_passes.py
+++ b/torch/onnx/_internal/exporter/_ir_passes.py
@@ -116,7 +116,6 @@ def rename_axis(
                 continue
             dim_name = dim.value
             if dim_name in sorted_rename_mapping:
-                # pyrefly: ignore [bad-index]
                 new_shape.append(sorted_rename_mapping[dim_name])
                 changed = True
             elif dim_name is not None:
diff --git a/torch/onnx/_internal/exporter/_isolated.py b/torch/onnx/_internal/exporter/_isolated.py
index b33ca002a577e..543a1e20d027f 100644
--- a/torch/onnx/_internal/exporter/_isolated.py
+++ b/torch/onnx/_internal/exporter/_isolated.py
@@ -5,7 +5,7 @@
 import multiprocessing
 import os
 import warnings
-from typing import Any, TYPE_CHECKING, TypeVar, Union
+from typing import Any, TYPE_CHECKING, TypeVar
 from typing_extensions import ParamSpec, TypeVarTuple, Unpack
 
 
@@ -22,7 +22,7 @@
 
 def _call_function_and_return_exception(
     func: Callable[[Unpack[_Ts]], _R], args: tuple[Unpack[_Ts]], kwargs: dict[str, Any]
-) -> Union[_R, Exception]:
+) -> _R | Exception:
     """Call function and return a exception if there is one."""
 
     try:
diff --git a/torch/onnx/_internal/exporter/_onnx_program.py b/torch/onnx/_internal/exporter/_onnx_program.py
index 2ffa88bac6d6c..5d7484d85bed5 100644
--- a/torch/onnx/_internal/exporter/_onnx_program.py
+++ b/torch/onnx/_internal/exporter/_onnx_program.py
@@ -157,7 +157,7 @@ def _to_ort_value(input: torch.Tensor | int | float | str | bool) -> ort.OrtValu
             int: np.int64,
             float: np.float32,
         }
-        # pyrefly: ignore [no-matching-overload]
+        # pyrefly: ignore [bad-argument-type, no-matching-overload]
         dtype = dtype_mapping.get(type(input))
         return ort.OrtValue.ortvalue_from_numpy(np.array(input, dtype=dtype))
 
@@ -254,7 +254,6 @@ def __call__(self, *args, **kwargs) -> Sequence[torch.Tensor]:
         run_options = ort.RunOptions()
         run_options.log_severity_level = 3  # 3: Error
         logger.debug("Running the inference session with %s arguments.", len(ort_input))
-        # pyrefly: ignore [missing-attribute]
         outputs = self._inference_session.run_with_ort_values(
             None, ort_input, run_options=run_options
         )
diff --git a/torch/onnx/_internal/exporter/_registration.py b/torch/onnx/_internal/exporter/_registration.py
index c754704e05185..d22186944e1fe 100644
--- a/torch/onnx/_internal/exporter/_registration.py
+++ b/torch/onnx/_internal/exporter/_registration.py
@@ -19,7 +19,7 @@
 import operator
 import types
 from collections.abc import Callable
-from typing import Literal, TypeAlias, Union
+from typing import Literal, TypeAlias
 
 import torch
 import torch._ops
@@ -28,7 +28,7 @@
 from torch.onnx._internal.exporter._torchlib import _torchlib_registry
 
 
-TorchOp: TypeAlias = Union[torch._ops.OpOverload, types.BuiltinFunctionType, Callable]
+TorchOp: TypeAlias = torch._ops.OpOverload | types.BuiltinFunctionType | Callable
 
 logger = logging.getLogger(__name__)
 
@@ -64,11 +64,8 @@ def __post_init__(self) -> None:
                 if isinstance(self.onnx_function, onnxscript.OnnxFunction):
                     signature = _schemas.op_signature_from_function(
                         self.onnx_function,
-                        # pyrefly: ignore [missing-attribute]
                         self.onnx_function.function_ir.domain,
-                        # pyrefly: ignore [missing-attribute]
                         self.onnx_function.name,
-                        # pyrefly: ignore [missing-attribute]
                         since_version=self.onnx_function.opset.version,
                     )
                 else:
@@ -296,12 +293,10 @@ def _cleanup_registry_based_on_opset_version(self) -> None:
                 max_opset = max(d.opset_introduced for d in decomps)
 
                 # Keep all decompositions with the maximum opset_introduced
-                # pyrefly: ignore [unsupported-operation]
                 cleaned_functions[target_or_name] = [
                     d for d in decomps if d.opset_introduced == max_opset
                 ]
 
-        # pyrefly: ignore [bad-assignment]
         self.functions = cleaned_functions
 
     def __repr__(self) -> str:
diff --git a/torch/onnx/_internal/exporter/_schemas.py b/torch/onnx/_internal/exporter/_schemas.py
index dbc26d50bd006..9b702baeb6351 100644
--- a/torch/onnx/_internal/exporter/_schemas.py
+++ b/torch/onnx/_internal/exporter/_schemas.py
@@ -302,7 +302,7 @@ def op_signature_from_function(
             if (
                 return_param_name := _get_type_constraint_name(return_type_i)
             ) in type_constraints:
-                # pyrefly: ignore [bad-index, index-error]
+                # pyrefly: ignore [bad-index]
                 type_constraint = type_constraints[return_param_name]
             else:
                 return_param_name = f"TReturn{i}"
@@ -315,7 +315,6 @@ def op_signature_from_function(
                 type_constraints[return_param_name] = type_constraint
             outputs.append(
                 ir.schemas.Parameter(
-                    # pyrefly: ignore [bad-argument-type]
                     name=return_param_name,
                     type_constraint=type_constraint,
                     required=True,
diff --git a/torch/onnx/_internal/exporter/_tensors.py b/torch/onnx/_internal/exporter/_tensors.py
index ef7ea0ba64284..981e2fef97de0 100644
--- a/torch/onnx/_internal/exporter/_tensors.py
+++ b/torch/onnx/_internal/exporter/_tensors.py
@@ -35,16 +35,13 @@ def __init__(
 
     @property
     def rank(self) -> int | None:
-        # pyrefly: ignore [missing-attribute]
         if self.shape is None:
             return None
-        # pyrefly: ignore [bad-argument-type]
         return len(self.shape)
 
     # TODO: Implement indexing
 
     def __mod__(self, other):
-        # pyrefly: ignore [missing-attribute]
         if self.dtype in {
             ir.DataType.FLOAT,
             ir.DataType.DOUBLE,
diff --git a/torch/onnx/_internal/exporter/_torchlib/_tensor_typing.py b/torch/onnx/_internal/exporter/_torchlib/_tensor_typing.py
index c0faf24f6f269..3212cf32468e1 100644
--- a/torch/onnx/_internal/exporter/_torchlib/_tensor_typing.py
+++ b/torch/onnx/_internal/exporter/_torchlib/_tensor_typing.py
@@ -24,7 +24,7 @@
 # NOTE: We do not care about unsigned types beyond UINT8 because PyTorch does not us them.
 # More detail can be found: https://pytorch.org/docs/stable/tensors.html
 
-TensorType = Union[
+TensorType = Union[  # noqa: UP007
     BFLOAT16,
     BOOL,
     COMPLEX64,
@@ -38,9 +38,9 @@
     INT64,
     UINT8,
 ]
-_FloatType = Union[FLOAT16, FLOAT, DOUBLE, BFLOAT16]
-IntType = Union[INT8, INT16, INT32, INT64]
-RealType = Union[
+_FloatType = Union[FLOAT16, FLOAT, DOUBLE, BFLOAT16]  # noqa: UP007
+IntType = Union[INT8, INT16, INT32, INT64]  # noqa: UP007
+RealType = Union[  # noqa: UP007
     BFLOAT16,
     FLOAT16,
     FLOAT,
@@ -55,19 +55,21 @@
 # Duplicate TTensor for inputs/outputs that accept the same set of types as TTensor
 # but do not constrain the type to be the same as the other inputs/outputs
 TTensor2 = TypeVar("TTensor2", bound=TensorType)
-TTensorOrString = TypeVar("TTensorOrString", bound=Union[TensorType, STRING])
+TTensorOrString = TypeVar("TTensorOrString", bound=Union[TensorType, STRING])  # noqa: UP007
 TFloat = TypeVar("TFloat", bound=_FloatType)
 TFloatOrUInt8 = TypeVar(
-    "TFloatOrUInt8", bound=Union[FLOAT, FLOAT16, DOUBLE, INT8, UINT8]
+    "TFloatOrUInt8",
+    bound=Union[FLOAT, FLOAT16, DOUBLE, INT8, UINT8],  # noqa: UP007
 )
 TInt = TypeVar("TInt", bound=IntType)
 TReal = TypeVar("TReal", bound=RealType)
 TRealUnlessInt16OrInt8 = TypeVar(
     "TRealUnlessInt16OrInt8",
-    bound=Union[FLOAT16, FLOAT, DOUBLE, BFLOAT16, INT32, INT64],
+    bound=Union[FLOAT16, FLOAT, DOUBLE, BFLOAT16, INT32, INT64],  # noqa: UP007
 )
 TRealUnlessFloat16OrInt8 = TypeVar(
-    "TRealUnlessFloat16OrInt8", bound=Union[DOUBLE, FLOAT, INT16, INT32, INT64]
+    "TRealUnlessFloat16OrInt8",
+    bound=Union[DOUBLE, FLOAT, INT16, INT32, INT64],  # noqa: UP007
 )
-TRealOrUInt8 = TypeVar("TRealOrUInt8", bound=Union[RealType, UINT8])
-TFloatHighPrecision = TypeVar("TFloatHighPrecision", bound=Union[FLOAT, DOUBLE])
+TRealOrUInt8 = TypeVar("TRealOrUInt8", bound=Union[RealType, UINT8])  # noqa: UP007
+TFloatHighPrecision = TypeVar("TFloatHighPrecision", bound=Union[FLOAT, DOUBLE])  # noqa: UP007
diff --git a/torch/onnx/_internal/exporter/_torchlib/ops/nn.py b/torch/onnx/_internal/exporter/_torchlib/ops/nn.py
index 1c0d1ca227bf3..14c3ad6d2a3ab 100644
--- a/torch/onnx/_internal/exporter/_torchlib/ops/nn.py
+++ b/torch/onnx/_internal/exporter/_torchlib/ops/nn.py
@@ -6,7 +6,7 @@
 
 from __future__ import annotations
 
-from typing import Optional, Sequence, TYPE_CHECKING  # noqa: UP035
+from typing import Sequence, TYPE_CHECKING  # noqa: UP035
 
 from onnxscript.onnx_opset import (  # type: ignore[attr-defined]
     opset20 as op20,
@@ -39,8 +39,8 @@ def aten_gelu_opset20(
 def aten_group_norm(
     input: TFloat,
     num_groups: int,
-    weight: Optional[TFloat] = None,
-    bias: Optional[TFloat] = None,
+    weight: TFloat | None = None,
+    bias: TFloat | None = None,
     eps: float = 1e-05,
     cudnn_enabled: bool = True,
 ) -> TFloat:
@@ -60,8 +60,8 @@ def aten_group_norm(
 def aten_rms_norm(
     input: TFloat,
     normalized_shape: Sequence[int],
-    weight: Optional[TFloat] = None,
-    eps: Optional[float] = None,
+    weight: TFloat | None = None,
+    eps: float | None = None,
 ) -> TFloat:
     """rms_norm(Tensor input, SymInt[] normalized_shape, Tensor? weight=None, float? eps=None) -> Tensor"""
 
@@ -92,10 +92,10 @@ def aten_scaled_dot_product_attention_23(
     query: TFloat,
     key: TFloat,
     value: TFloat,
-    attn_mask: Optional[TFloat] = None,
+    attn_mask: TFloat | None = None,
     dropout_p: float = 0.0,
     is_causal: bool = False,
-    scale: Optional[float] = None,
+    scale: float | None = None,
     enable_gqa: bool = False,
 ) -> TFloat:
     """scaled_dot_product_attention(Tensor query, Tensor key, Tensor value, Tensor? attn_mask=None, float dropout_p=0.0, bool is_causal=False, *, float? scale=None, bool enable_gqa=False) -> Tensor
diff --git a/torch/onnx/_internal/fx/passes/type_promotion.py b/torch/onnx/_internal/fx/passes/type_promotion.py
index 71d299a29f4fa..1f500e170cf74 100644
--- a/torch/onnx/_internal/fx/passes/type_promotion.py
+++ b/torch/onnx/_internal/fx/passes/type_promotion.py
@@ -1540,7 +1540,6 @@ def _maybe_promote_arg(
         elif isinstance(fx_arg, (tuple, list)):
             logger.info("Argument %s is a tuple/list. Promoting each element.", fx_arg)
             return type(fx_arg)(
-                # pyrefly: ignore [bad-argument-type]
                 self._maybe_promote_arg(node, fx_arg_elem, dtype)
                 for fx_arg_elem in fx_arg
             )
diff --git a/torch/onnx/_internal/fx/type_utils.py b/torch/onnx/_internal/fx/type_utils.py
index bd32799bc4e2b..ca434fa213f79 100644
--- a/torch/onnx/_internal/fx/type_utils.py
+++ b/torch/onnx/_internal/fx/type_utils.py
@@ -4,11 +4,14 @@
 from __future__ import annotations
 
 from typing import Any
+from typing_extensions import TypeIs
 
 import torch
 
 
-def is_torch_symbolic_type(value: Any) -> bool:
+def is_torch_symbolic_type(
+    value: Any,
+) -> TypeIs[torch.SymBool | torch.SymInt | torch.SymFloat]:
     return isinstance(value, (torch.SymBool, torch.SymInt, torch.SymFloat))
 
 
diff --git a/torch/onnx/_internal/torchscript_exporter/_experimental.py b/torch/onnx/_internal/torchscript_exporter/_experimental.py
index 0fac4450a71c8..6f1fee027331a 100644
--- a/torch/onnx/_internal/torchscript_exporter/_experimental.py
+++ b/torch/onnx/_internal/torchscript_exporter/_experimental.py
@@ -2,7 +2,6 @@
 
 import dataclasses
 from collections.abc import Mapping, Sequence
-from typing import Optional, Union
 
 import torch
 import torch._C._onnx as _C_onnx
@@ -17,12 +16,12 @@ class ExportOptions:
     export_params: bool = True
     verbose: bool = False
     training: _C_onnx.TrainingMode = _C_onnx.TrainingMode.EVAL
-    input_names: Optional[Sequence[str]] = None
-    output_names: Optional[Sequence[str]] = None
+    input_names: Sequence[str] | None = None
+    output_names: Sequence[str] | None = None
     operator_export_type: _C_onnx.OperatorExportTypes = _C_onnx.OperatorExportTypes.ONNX
-    opset_version: Optional[int] = None
+    opset_version: int | None = None
     do_constant_folding: bool = True
-    dynamic_axes: Optional[Mapping[str, Union[Mapping[int, str], Sequence[int]]]] = None
-    keep_initializers_as_inputs: Optional[bool] = None
-    custom_opsets: Optional[Mapping[str, int]] = None
-    export_modules_as_functions: Union[bool, set[type[torch.nn.Module]]] = False
+    dynamic_axes: Mapping[str, Mapping[int, str] | Sequence[int]] | None = None
+    keep_initializers_as_inputs: bool | None = None
+    custom_opsets: Mapping[str, int] | None = None
+    export_modules_as_functions: bool | set[type[torch.nn.Module]] = False
diff --git a/torch/onnx/_internal/torchscript_exporter/_type_utils.py b/torch/onnx/_internal/torchscript_exporter/_type_utils.py
index 9a253bbc3934a..c0ffa56a38e3b 100644
--- a/torch/onnx/_internal/torchscript_exporter/_type_utils.py
+++ b/torch/onnx/_internal/torchscript_exporter/_type_utils.py
@@ -153,7 +153,7 @@ def from_dtype(cls, dtype: torch.dtype | None) -> JitScalarType:
         """
         if dtype not in _DTYPE_TO_SCALAR_TYPE:
             raise errors.OnnxExporterError(f"Unknown dtype: {dtype}")
-        # pyrefly: ignore [bad-index, index-error]
+        # pyrefly: ignore [bad-index]
         return _DTYPE_TO_SCALAR_TYPE[dtype]
 
     @classmethod
@@ -173,6 +173,7 @@ def from_onnx_type(
         """
         if onnx_type not in _ONNX_TO_SCALAR_TYPE:
             raise errors.OnnxExporterError(f"Unknown onnx_type: {onnx_type}")
+        # pyrefly: ignore [redundant-cast]
         return _ONNX_TO_SCALAR_TYPE[typing.cast(_C_onnx.TensorProtoDataType, onnx_type)]
 
     @classmethod
diff --git a/torch/onnx/_internal/torchscript_exporter/registration.py b/torch/onnx/_internal/torchscript_exporter/registration.py
index e35903e6823c9..3ac1fdb048ea1 100644
--- a/torch/onnx/_internal/torchscript_exporter/registration.py
+++ b/torch/onnx/_internal/torchscript_exporter/registration.py
@@ -3,7 +3,7 @@
 
 import warnings
 from collections.abc import Callable, Collection, Sequence
-from typing import Generic, Optional, TypeVar, Union
+from typing import Generic, TypeVar
 from typing_extensions import ParamSpec
 
 from torch.onnx import _constants, errors
@@ -14,7 +14,7 @@
 
 def _dispatch_opset_version(
     target: OpsetVersion, registered_opsets: Collection[OpsetVersion]
-) -> Optional[OpsetVersion]:
+) -> OpsetVersion | None:
     """Finds the registered opset given a target opset version and the available opsets.
 
     Args:
@@ -96,7 +96,7 @@ def overridden(self, key: _K) -> bool:
     def __getitem__(self, key: _K) -> _V:
         return self._merged[key]
 
-    def get(self, key: _K, default: Optional[_V] = None):
+    def get(self, key: _K, default: _V | None = None):
         return self._merged.get(key, default)
 
     def __contains__(self, key: object) -> bool:
@@ -143,7 +143,7 @@ def __getitem__(self, key: OpsetVersion) -> Callable:
 
     # TODO(justinchuby): Add @functools.lru_cache(maxsize=None) if lookup time becomes
     # a problem.
-    def get(self, opset: OpsetVersion) -> Optional[Callable]:
+    def get(self, opset: OpsetVersion) -> Callable | None:
         """Find the most recent version of the function."""
         version = _dispatch_opset_version(opset, self._functions)
         if version is None:
@@ -245,7 +245,7 @@ def unregister(self, name: str, opset: OpsetVersion) -> None:
             return
         self._registry[name].remove_custom(opset)
 
-    def get_function_group(self, name: str) -> Optional[_SymbolicFunctionGroup]:
+    def get_function_group(self, name: str) -> _SymbolicFunctionGroup | None:
         """Returns the function group for the given name."""
         return self._registry.get(name)
 
@@ -263,8 +263,8 @@ def all_functions(self) -> set[str]:
 
 def onnx_symbolic(
     name: str,
-    opset: Union[OpsetVersion, Sequence[OpsetVersion]],
-    decorate: Optional[Sequence[Callable]] = None,
+    opset: OpsetVersion | Sequence[OpsetVersion],
+    decorate: Sequence[Callable] | None = None,
     custom: bool = False,
 ) -> Callable:
     """Registers a symbolic function.
@@ -314,8 +314,8 @@ def wrapper(func: Callable[_P, _R]) -> Callable[_P, _R]:
 
 def custom_onnx_symbolic(
     name: str,
-    opset: Union[OpsetVersion, Sequence[OpsetVersion]],
-    decorate: Optional[Sequence[Callable]] = None,
+    opset: OpsetVersion | Sequence[OpsetVersion],
+    decorate: Sequence[Callable] | None = None,
 ) -> Callable:
     """Registers a custom symbolic function.
 
diff --git a/torch/onnx/_internal/torchscript_exporter/symbolic_helper.py b/torch/onnx/_internal/torchscript_exporter/symbolic_helper.py
index d52e2a6ee9249..0e00a7a5cd263 100644
--- a/torch/onnx/_internal/torchscript_exporter/symbolic_helper.py
+++ b/torch/onnx/_internal/torchscript_exporter/symbolic_helper.py
@@ -110,7 +110,7 @@
     NoReturn,
     TypeVar as _TypeVar,
 )
-from typing_extensions import ParamSpec as _ParamSpec
+from typing_extensions import ParamSpec as _ParamSpec, TypeIs as _TypeIs
 
 import torch
 import torch._C._onnx as _C_onnx
@@ -561,7 +561,7 @@ def _is_none(x: Any) -> bool:
     return x is None or (x.node().mustBeNone() if isinstance(x, _C.Value) else False)
 
 
-def _is_value(x: Any) -> bool:
+def _is_value(x: Any) -> _TypeIs[_C.Value]:
     return isinstance(x, _C.Value)
 
 
diff --git a/torch/onnx/_internal/torchscript_exporter/symbolic_opset14.py b/torch/onnx/_internal/torchscript_exporter/symbolic_opset14.py
index 2fcbdad947082..7beb3de13f272 100644
--- a/torch/onnx/_internal/torchscript_exporter/symbolic_opset14.py
+++ b/torch/onnx/_internal/torchscript_exporter/symbolic_opset14.py
@@ -167,7 +167,7 @@ def scaled_dot_product_attention(
     # NOTE: onnx-script has different logic here, because the attribute perms in
     # transpose needs list of ints
     key_shape_builtin = symbolic_helper._get_tensor_rank(key)
-    # pyrefly: ignore [no-matching-overload]
+    # pyrefly: ignore [bad-argument-type, no-matching-overload]
     key_transposed_axes = list(range(key_shape_builtin))
     key_transposed_axes[-1], key_transposed_axes[-2] = (
         key_transposed_axes[-2],
diff --git a/torch/onnx/_internal/torchscript_exporter/symbolic_opset17.py b/torch/onnx/_internal/torchscript_exporter/symbolic_opset17.py
index 3ad38e0540e8c..bb17d8f72942b 100644
--- a/torch/onnx/_internal/torchscript_exporter/symbolic_opset17.py
+++ b/torch/onnx/_internal/torchscript_exporter/symbolic_opset17.py
@@ -19,7 +19,6 @@
 
 import functools
 from collections.abc import Sequence
-from typing import Optional
 
 import torch
 from torch import _C
@@ -108,13 +107,13 @@ def stft(
     g: jit_utils.GraphContext,
     input: _C.Value,
     n_fft: int,
-    hop_length: Optional[int] = None,
-    win_length: Optional[int] = None,
-    window: Optional[_C.Value] = None,
+    hop_length: int | None = None,
+    win_length: int | None = None,
+    window: _C.Value | None = None,
     normalized: bool = False,
-    onesided: Optional[bool] = True,
-    return_complex: Optional[bool] = False,
-    align_to_window: Optional[bool] = None,
+    onesided: bool | None = True,
+    return_complex: bool | None = False,
+    align_to_window: bool | None = None,
 ) -> _C.Value:
     """Associates `torch.stft` with the `STFT` ONNX operator.
     Note that torch.stft calls _VF.stft, without centering or padding options.
diff --git a/torch/onnx/_internal/torchscript_exporter/symbolic_opset18.py b/torch/onnx/_internal/torchscript_exporter/symbolic_opset18.py
index 0951ea4304f0e..1d622bdd141e3 100644
--- a/torch/onnx/_internal/torchscript_exporter/symbolic_opset18.py
+++ b/torch/onnx/_internal/torchscript_exporter/symbolic_opset18.py
@@ -21,7 +21,6 @@
 
 import functools
 from collections.abc import Sequence
-from typing import Optional
 
 import torch
 from torch import _C
@@ -266,7 +265,7 @@ def linalg_vector_norm(
     g: jit_utils.GraphContext,
     self: torch._C.Value,
     ord: float,
-    dim: Optional[Sequence[int]],
+    dim: Sequence[int] | None,
     keepdim: bool,
     dtype: torch._C.Value,
 ):
diff --git a/torch/onnx/_internal/torchscript_exporter/symbolic_opset9.py b/torch/onnx/_internal/torchscript_exporter/symbolic_opset9.py
index 38979c324729e..d781c0dfbdf63 100644
--- a/torch/onnx/_internal/torchscript_exporter/symbolic_opset9.py
+++ b/torch/onnx/_internal/torchscript_exporter/symbolic_opset9.py
@@ -2403,7 +2403,6 @@ def _convolution_mode(
         "group_i": groups,
     }
 
-    # pyrefly: ignore [bad-argument-type]
     n = g.op("Conv", *args, **kwargs)
 
     if (
diff --git a/torch/onnx/_internal/torchscript_exporter/verification.py b/torch/onnx/_internal/torchscript_exporter/verification.py
index 16be695de7fef..57f764e106ba3 100644
--- a/torch/onnx/_internal/torchscript_exporter/verification.py
+++ b/torch/onnx/_internal/torchscript_exporter/verification.py
@@ -19,7 +19,7 @@
 import tempfile
 import warnings
 from collections.abc import Mapping, Sequence
-from typing import Any, Union
+from typing import Any
 
 import numpy as np
 import numpy.typing as npt
@@ -34,11 +34,11 @@
 
 _ORT_PROVIDERS = ("CPUExecutionProvider",)
 
-_NumericType = Union[Number, torch.Tensor, np.ndarray]
-_ModelType = Union[torch.nn.Module, torch.jit.ScriptModule]
-_InputArgsType = Union[torch.Tensor, tuple[Any, ...]]
+_NumericType = Number | torch.Tensor | np.ndarray
+_ModelType = torch.nn.Module | torch.jit.ScriptModule
+_InputArgsType = torch.Tensor | tuple[Any, ...]
 _InputKwargsType = Mapping[str, Any]
-_OutputsType = Union[Sequence[_NumericType], Sequence]
+_OutputsType = Sequence[_NumericType] | Sequence
 
 
 class OnnxBackend(enum.Enum):
diff --git a/torch/onnx/ops/_impl.py b/torch/onnx/ops/_impl.py
index fb2c509eb436c..be50d953b619b 100644
--- a/torch/onnx/ops/_impl.py
+++ b/torch/onnx/ops/_impl.py
@@ -8,7 +8,7 @@
 # flake8: noqa: B950
 import math
 from collections.abc import Callable
-from typing import Optional, TypeVar
+from typing import TypeVar
 from typing_extensions import ParamSpec
 
 import torch
@@ -54,7 +54,7 @@ def _rotary_embedding_23_fake_impl(
     x: torch.Tensor,
     cos_cache: torch.Tensor,
     sin_cache: torch.Tensor,
-    position_ids: Optional[torch.Tensor] = None,
+    position_ids: torch.Tensor | None = None,
     *,
     interleaved: bool = False,
     num_heads: int = 0,
@@ -69,7 +69,7 @@ def rotary_embedding_23(
     x: torch.Tensor,
     cos_cache: torch.Tensor,
     sin_cache: torch.Tensor,
-    position_ids: Optional[torch.Tensor] = None,
+    position_ids: torch.Tensor | None = None,
     *,
     interleaved: bool = False,
     num_heads: int = 0,
@@ -200,7 +200,7 @@ def rotary_embedding_23(
     return torch.permute(output, (0, 2, 1, 3))
 
 
-def _get_scale_factor(scale: Optional[float], head_size: int) -> float:
+def _get_scale_factor(scale: float | None, head_size: int) -> float:
     """Get the scale factor for attention computation."""
     return scale if scale is not None else (1.0 / math.sqrt(head_size))
 
@@ -223,7 +223,7 @@ def _get_qk_output_for_aten_spda(
     K: torch.Tensor,
     current_q_num_heads: int,
     current_kv_num_heads: int,
-    scale: Optional[float],
+    scale: float | None,
     qk_matmul_output_mode: int,
 ) -> torch.Tensor:
     """Get QK output tensor based on the specified mode."""
@@ -251,7 +251,7 @@ def _compute_qk_output_for_mode_0(
     K: torch.Tensor,
     current_q_num_heads: int,
     current_kv_num_heads: int,
-    scale: Optional[float],
+    scale: float | None,
 ) -> torch.Tensor:
     """Helper function to compute QK output for qk_matmul_output_mode == 0."""
     # Handle GQA manually for QK output
@@ -272,17 +272,17 @@ def _attention_23_fake_impl(
     Q: torch.Tensor,
     K: torch.Tensor,
     V: torch.Tensor,
-    attn_mask: Optional[torch.Tensor] = None,
-    past_key: Optional[torch.Tensor] = None,
-    past_value: Optional[torch.Tensor] = None,
+    attn_mask: torch.Tensor | None = None,
+    past_key: torch.Tensor | None = None,
+    past_value: torch.Tensor | None = None,
     *,
     is_causal: bool = False,
     kv_num_heads: int = 0,
     q_num_heads: int = 0,
     qk_matmul_output_mode: int = 0,
-    scale: Optional[float] = None,
+    scale: float | None = None,
     softcap: float = 0.0,
-    softmax_precision: Optional[int] = None,
+    softmax_precision: int | None = None,
 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
     """Fake implementation for Attention-23 for torch.compile purposes."""
     batch_size = Q.shape[0]
@@ -357,17 +357,17 @@ def attention_23(
     Q: torch.Tensor,
     K: torch.Tensor,
     V: torch.Tensor,
-    attn_mask: Optional[torch.Tensor] = None,
-    past_key: Optional[torch.Tensor] = None,
-    past_value: Optional[torch.Tensor] = None,
+    attn_mask: torch.Tensor | None = None,
+    past_key: torch.Tensor | None = None,
+    past_value: torch.Tensor | None = None,
     *,
     is_causal: bool = False,
     kv_num_heads: int = 0,
     q_num_heads: int = 0,
     qk_matmul_output_mode: int = 0,
-    scale: Optional[float] = None,
+    scale: float | None = None,
     softcap: float = 0.0,
-    softmax_precision: Optional[int] = None,
+    softmax_precision: int | None = None,
 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
     """Attention-23 https://onnx.ai/onnx/operators/onnx__Attention.html#attention-23"""
 
diff --git a/torch/onnx/ops/_symbolic_impl.py b/torch/onnx/ops/_symbolic_impl.py
index 669d6ef57c6a4..85963dd85b1d2 100644
--- a/torch/onnx/ops/_symbolic_impl.py
+++ b/torch/onnx/ops/_symbolic_impl.py
@@ -14,7 +14,6 @@
 # flake8: noqa: B950
 import dataclasses
 from collections.abc import Sequence
-from typing import Optional, Union
 
 import torch
 from torch.onnx.ops import _dtype_mappings
@@ -58,16 +57,14 @@ def from_dict(
         cls,
         attrs: dict[
             str,
-            Union[
-                int,
-                float,
-                str,
-                bool,
-                Sequence[int],
-                Sequence[float],
-                Sequence[str],
-                Sequence[bool],
-            ],
+            int
+            | float
+            | str
+            | bool
+            | Sequence[int]
+            | Sequence[float]
+            | Sequence[str]
+            | Sequence[bool],
         ],
     ) -> "EncodedAttrs":
         encoded = cls(
@@ -131,26 +128,12 @@ def to_dict(
         self,
     ) -> dict[
         str,
-        Union[
-            int,
-            float,
-            str,
-            list[int],
-            list[float],
-            list[str],
-        ],
+        int | float | str | list[int] | list[float] | list[str],
     ]:
         """Convert the encoded attributes back to a dictionary for creating an ONNX node."""
         attrs: dict[
             str,
-            Union[
-                int,
-                float,
-                str,
-                list[int],
-                list[float],
-                list[str],
-            ],
+            int | float | str | list[int] | list[float] | list[str],
         ] = {}
         for i, key in enumerate(self.attr_keys):
             attr_type = self.attr_types[i]
@@ -183,11 +166,11 @@ def to_dict(
     ),
 )
 def _symbolic(
-    inputs: Sequence[Optional[torch.Tensor]],
+    inputs: Sequence[torch.Tensor | None],
     op_type: str,
     onnx_dtype: int,
     *,
-    shape: Sequence[Union[int, torch.SymInt]],
+    shape: Sequence[int | torch.SymInt],
     attr_keys: Sequence[str],
     attr_types: Sequence[str],
     attr_pos: Sequence[tuple[int, int]],
@@ -197,7 +180,7 @@ def _symbolic(
     metadata_props_keys: Sequence[str] = (),
     metadata_props_values: Sequence[str] = (),
     domain: str = "",
-    version: Optional[int] = None,
+    version: int | None = None,
 ) -> torch.Tensor:
     torch._check(
         onnx_dtype in _dtype_mappings.ONNX_DTYPE_TO_TORCH_DTYPE,
@@ -214,7 +197,7 @@ def _(
     op_type: str,
     onnx_dtype: int,
     *,
-    shape: Sequence[Union[int, torch.SymInt]],
+    shape: Sequence[int | torch.SymInt],
     attr_keys: Sequence[str],
     attr_types: Sequence[str],
     attr_pos: Sequence[tuple[int, int]],
@@ -224,7 +207,7 @@ def _(
     metadata_props_keys: Sequence[str] = (),
     metadata_props_values: Sequence[str] = (),
     domain: str = "",
-    version: Optional[int] = None,
+    version: int | None = None,
 ) -> torch.Tensor:
     torch._check(
         onnx_dtype in _dtype_mappings.ONNX_DTYPE_TO_TORCH_DTYPE,
@@ -249,11 +232,11 @@ def _(
     ),
 )
 def _symbolic_multi_out(
-    inputs: Sequence[Optional[torch.Tensor]],
+    inputs: Sequence[torch.Tensor | None],
     op_type: str,
     onnx_dtypes: Sequence[int],
     *,
-    shapes: Sequence[Sequence[Union[int, torch.SymInt]]],
+    shapes: Sequence[Sequence[int | torch.SymInt]],
     attr_keys: Sequence[str],
     attr_types: Sequence[str],
     attr_pos: Sequence[tuple[int, int]],
@@ -263,7 +246,7 @@ def _symbolic_multi_out(
     metadata_props_keys: Sequence[str] = (),
     metadata_props_values: Sequence[str] = (),
     domain: str = "",
-    version: Optional[int] = None,
+    version: int | None = None,
 ) -> list[torch.Tensor]:
     outputs = []
     torch._check(
@@ -289,7 +272,7 @@ def _(
     op_type: str,
     onnx_dtypes: Sequence[int],
     *,
-    shapes: Sequence[Sequence[Union[int, torch.SymInt]]],
+    shapes: Sequence[Sequence[int | torch.SymInt]],
     attr_keys: Sequence[str],
     attr_types: Sequence[str],
     attr_pos: Sequence[tuple[int, int]],
@@ -299,7 +282,7 @@ def _(
     metadata_props_keys: Sequence[str] = (),
     metadata_props_values: Sequence[str] = (),
     domain: str = "",
-    version: Optional[int] = None,
+    version: int | None = None,
 ) -> list[torch.Tensor]:
     outputs = []
     torch._check(
diff --git a/torch/optim/_multi_tensor/__init__.py b/torch/optim/_multi_tensor/__init__.py
index face68d0bc504..b6818e5a50f3b 100644
--- a/torch/optim/_multi_tensor/__init__.py
+++ b/torch/optim/_multi_tensor/__init__.py
@@ -13,7 +13,6 @@
 
 def partialclass(cls, *args, **kwargs):  # noqa: D103
     class NewCls(cls):
-        # pyrefly: ignore [not-iterable]
         __init__ = partialmethod(cls.__init__, *args, **kwargs)
 
     return NewCls
diff --git a/torch/optim/_muon.py b/torch/optim/_muon.py
index e441c8b911b2f..2e45e07c4a596 100644
--- a/torch/optim/_muon.py
+++ b/torch/optim/_muon.py
@@ -273,6 +273,29 @@ def step(self, closure=None):
         adjust_lr_fn (str, optional): function to adjust learning rate. One of "original" and "match_rms_adamw".
             If not specified, we will default to use "original". (default: None)
 
+    Example:
+        >>> # xdoctest: +SKIP
+        >>> # Muon only supports 2D params; use a standard optimizer
+        >>> # such as AdamW for biases, embeddings, and other non-2D
+        >>> # parameters.
+        >>> muon_params = [
+        ...     p for p in model.parameters() if p.ndim == 2
+        ... ]
+        >>> other_params = [
+        ...     p for p in model.parameters() if p.ndim != 2
+        ... ]
+        >>> optim_muon = torch.optim.Muon(
+        ...     muon_params, lr=0.02, momentum=0.95
+        ... )
+        >>> optim_adamw = torch.optim.AdamW(
+        ...     other_params, lr=3e-4, weight_decay=0.01
+        ... )
+        >>> optim_muon.zero_grad()
+        >>> optim_adamw.zero_grad()
+        >>> loss_fn(model(input), target).backward()
+        >>> optim_muon.step()
+        >>> optim_adamw.step()
+
     .. _Muon\: An optimizer for hidden layers in neural networks:
         https://kellerjordan.github.io/posts/muon/
     .. _Muon is Scalable for LLM Training:
diff --git a/torch/optim/adagrad.py b/torch/optim/adagrad.py
index a6a57fb61b8ba..d18e32fe17a98 100644
--- a/torch/optim/adagrad.py
+++ b/torch/optim/adagrad.py
@@ -16,6 +16,7 @@
     _to_scalar,
     _use_grad_for_differentiable,
     _view_as_real,
+    DeviceDict,
     Optimizer,
     ParamsT,
 )
@@ -73,6 +74,7 @@ def __init__(
             if foreach:
                 raise RuntimeError("`fused` and `foreach` cannot be `True` together.")
             self._need_device_dtype_check_for_fused = True
+            self._step_supports_amp_scaling = True
 
         for group in self.param_groups:
             for p in group["params"]:
@@ -106,6 +108,20 @@ def __setstate__(self, state):
             group.setdefault("differentiable", False)
             fused = group.setdefault("fused", None)
 
+            for p in group["params"]:
+                p_state = self.state.get(p, [])
+                if len(p_state) != 0 and not torch.is_tensor(p_state["step"]):
+                    step_val = float(p_state["step"])
+                    p_state["step"] = (
+                        torch.tensor(
+                            step_val,
+                            dtype=_get_scalar_dtype(is_fused=fused),
+                            device=p.device,
+                        )
+                        if group["fused"]
+                        else torch.tensor(step_val, dtype=_get_scalar_dtype())
+                    )
+
         state_values = list(self.state.values())
         step_is_tensor = (len(state_values) != 0) and torch.is_tensor(
             state_values[0]["step"]
@@ -132,13 +148,38 @@ def _init_group(self, group, params_with_grad, grads, state_sums, state_steps):
                     "_need_device_dtype_check_for_fused",
                     True,
                 ):
-                    _device_dtype_check_for_fused(p, cuda_unsupported=True)
+                    _device_dtype_check_for_fused(p)
                     self._need_device_dtype_check_for_fused = False
                 has_sparse_grad |= p.grad.is_sparse
                 has_complex |= torch.is_complex(p)
                 params_with_grad.append(p)
                 grads.append(p.grad)
                 state = self.state[p]
+                if len(state) == 0:
+                    if group["fused"]:
+                        _device_dtype_check_for_fused(p)
+
+                    state["step"] = (
+                        torch.zeros(
+                            (),
+                            dtype=_get_scalar_dtype(is_fused=group["fused"]),
+                            device=p.device,
+                        )
+                        if group["fused"]
+                        else torch.tensor(0.0, dtype=_get_scalar_dtype())
+                    )
+
+                    initial_accumulator_value = self.defaults[
+                        "initial_accumulator_value"
+                    ]
+                    init_value = (
+                        complex(initial_accumulator_value, initial_accumulator_value)
+                        if torch.is_complex(p)
+                        else initial_accumulator_value
+                    )
+                    state["sum"] = torch.full_like(
+                        p, init_value, memory_format=torch.preserve_format
+                    )
                 state_sums.append(state["sum"])
                 state_steps.append(state["step"])
 
@@ -230,7 +271,7 @@ def step(self, closure=None):
         {_foreach_doc}
         {_maximize_doc}
         {_differentiable_doc}
-        fused (bool, optional): whether the fused implementation (CPU only) is used.
+        fused (bool, optional): whether the fused implementation (CPU and CUDA only) is used.
             Currently, `torch.float64`, `torch.float32`, `torch.float16`, and `torch.bfloat16`
             are supported. (default: None). Please note that the fused implementations does not
             support sparse or complex gradients.
@@ -507,7 +548,7 @@ def _fused_adagrad(
     grad_scale: Tensor | None,
     found_inf: Tensor | None,
     *,
-    lr: float,
+    lr: float | Tensor,
     weight_decay: float,
     lr_decay: float,
     eps: float,
@@ -526,12 +567,15 @@ def _fused_adagrad(
             "adagrad with fused=True does not support differentiable=True"
         )
 
-    lr = _to_scalar(lr)
-
-    grad_scale_dict = (
-        {grad_scale.device: grad_scale} if grad_scale is not None else None
+    grad_scale_dict: DeviceDict = (
+        {grad_scale.device: grad_scale} if grad_scale is not None else {}
+    )
+    found_inf_dict: DeviceDict = (
+        {found_inf.device: found_inf} if found_inf is not None else {}
+    )
+    lr_dict: DeviceDict | None = (
+        {lr.device: lr} if isinstance(lr, Tensor) and str(lr.device) != "cpu" else None
     )
-    found_inf_dict = {found_inf.device: found_inf} if found_inf is not None else None
 
     grouped_tensors = Optimizer._group_tensors_by_device_and_dtype(
         [params, grads, state_sums, state_steps]  # type: ignore[list-item]
@@ -551,14 +595,17 @@ def _fused_adagrad(
         device_state_steps = cast(list[Tensor], device_state_steps_)
 
         device_grad_scale, device_found_inf = None, None
-        if grad_scale is not None and grad_scale_dict is not None:
-            if device not in grad_scale_dict:
-                grad_scale_dict[device] = grad_scale.to(device, non_blocking=True)  # type: ignore[index]
-            device_grad_scale = grad_scale_dict[device]  # type: ignore[index]
-        if found_inf is not None and found_inf_dict is not None:
-            if found_inf not in found_inf_dict:
-                found_inf_dict[device] = found_inf.to(device, non_blocking=True)  # type: ignore[index]
-            device_found_inf = found_inf_dict[device]  # type: ignore[index]
+        if grad_scale is not None:
+            device_grad_scale = grad_scale_dict.setdefault(
+                device, grad_scale.to(device, non_blocking=True)
+            )
+        if found_inf is not None:
+            device_found_inf = found_inf_dict.setdefault(
+                device, found_inf.to(device, non_blocking=True)
+            )
+        if lr_dict is not None and device not in lr_dict:
+            lr_dict[device] = lr.to(device=device, non_blocking=True)  # type: ignore[union-attr]
+            lr = lr_dict[device]
         torch._foreach_add_(device_state_steps, 1)
         torch._fused_adagrad_(
             device_params,
diff --git a/torch/optim/adam.py b/torch/optim/adam.py
index 3351e095ee84c..ca3750484c464 100644
--- a/torch/optim/adam.py
+++ b/torch/optim/adam.py
@@ -316,7 +316,7 @@ def step(self, closure=None):
         lr (float, Tensor, optional): learning rate (default: 1e-3). A tensor LR
             is not yet supported for all our implementations. Please use a float
             LR if you are not also specifying fused=True or capturable=True.
-        betas (tuple[Union[float, Tensor], Union[float, Tensor]], optional):
+        betas (tuple[float | Tensor, float | Tensor], optional):
             coefficients used for computing running averages of gradient and
             its square. If a tensor is provided, must be 1-element. (default: (0.9, 0.999))
         eps (float, optional): term added to the denominator to improve
@@ -423,7 +423,6 @@ def _single_tensor_adam(
                     if weight_decay.requires_grad:
                         grad = grad.addcmul_(param.clone(), weight_decay)
                     else:
-                        # pyrefly: ignore [bad-argument-type]
                         grad = grad.add(param, alpha=weight_decay)
                 else:
                     grad = grad.add(param, alpha=weight_decay)
diff --git a/torch/optim/adamw.py b/torch/optim/adamw.py
index aa3b922cf90b4..fbb3e45327469 100644
--- a/torch/optim/adamw.py
+++ b/torch/optim/adamw.py
@@ -101,7 +101,7 @@ def __setstate__(self, state):
         lr (float, Tensor, optional): learning rate (default: 1e-3). A tensor LR
             is not yet supported for all our implementations. Please use a float
             LR if you are not also specifying fused=True or capturable=True.
-        betas (tuple[Union[float, Tensor], Union[float, Tensor]], optional):
+        betas (tuple[float | Tensor, float | Tensor], optional):
             coefficients used for computing running averages of gradient and
             its square. If a tensor is provided, must be 1-element. (default: (0.9, 0.999))
         eps (float, optional): term added to the denominator to improve
diff --git a/torch/optim/asgd.py b/torch/optim/asgd.py
index 21e119b2a13fc..d36e6e0ec8c70 100644
--- a/torch/optim/asgd.py
+++ b/torch/optim/asgd.py
@@ -264,7 +264,6 @@ def _single_tensor_asgd(
             ax.copy_(param)
 
         if capturable:
-            # pyrefly: ignore [unsupported-operation]
             eta.copy_(lr / ((1 + lambd * lr * step_t) ** alpha))
             mu.copy_(1 / torch.maximum(step_t - t0, torch.ones_like(step_t)))
         else:
diff --git a/torch/optim/lbfgs.py b/torch/optim/lbfgs.py
index ed4cf1a8b2e88..af6b6e8b02bb0 100644
--- a/torch/optim/lbfgs.py
+++ b/torch/optim/lbfgs.py
@@ -112,13 +112,11 @@ def _strong_wolfe(
 
         # compute new trial value
         t = _cubic_interpolate(
-            # pyrefly: ignore [index-error]
             # pyrefly: ignore [unbound-name]
             bracket[0],
             # pyrefly: ignore [unbound-name]
             bracket_f[0],
             bracket_gtd[0],  # type: ignore[possibly-undefined]
-            # pyrefly: ignore [index-error]
             # pyrefly: ignore [unbound-name]
             bracket[1],
             # pyrefly: ignore [unbound-name]
@@ -164,7 +162,6 @@ def _strong_wolfe(
         # pyrefly: ignore [unbound-name]
         if f_new > (f + c1 * t * gtd) or f_new >= bracket_f[low_pos]:
             # Armijo condition not satisfied or not lower than lowest point
-            # pyrefly: ignore [unsupported-operation]
             # pyrefly: ignore [unbound-name]
             bracket[high_pos] = t
             # pyrefly: ignore [unbound-name]
@@ -178,11 +175,9 @@ def _strong_wolfe(
             if abs(gtd_new) <= -c2 * gtd:
                 # Wolfe conditions satisfied
                 done = True
-            # pyrefly: ignore [index-error]
             # pyrefly: ignore [unbound-name]
             elif gtd_new * (bracket[high_pos] - bracket[low_pos]) >= 0:
                 # old high becomes new low
-                # pyrefly: ignore [unsupported-operation]
                 # pyrefly: ignore [unbound-name]
                 bracket[high_pos] = bracket[low_pos]
                 # pyrefly: ignore [unbound-name]
@@ -192,7 +187,6 @@ def _strong_wolfe(
                 bracket_gtd[high_pos] = bracket_gtd[low_pos]
 
             # new point becomes new low
-            # pyrefly: ignore [unsupported-operation]
             # pyrefly: ignore [unbound-name]
             bracket[low_pos] = t
             # pyrefly: ignore [unbound-name]
@@ -281,7 +275,6 @@ def __init__(
 
     def _numel(self):
         if self._numel_cache is None:
-            # pyrefly: ignore [bad-assignment]
             self._numel_cache = sum(
                 2 * p.numel() if torch.is_complex(p) else p.numel()
                 for p in self._params
diff --git a/torch/optim/lr_scheduler.py b/torch/optim/lr_scheduler.py
index 208a182bb1770..b2049c73fc472 100644
--- a/torch/optim/lr_scheduler.py
+++ b/torch/optim/lr_scheduler.py
@@ -413,7 +413,6 @@ def state_dict(self) -> dict[str, Any]:
 
         for idx, fn in enumerate(self.lr_lambdas):
             if not isinstance(fn, types.FunctionType):
-                # pyrefly: ignore [unsupported-operation]
                 state_dict["lr_lambdas"][idx] = fn.__dict__.copy()
 
         return state_dict
@@ -533,7 +532,6 @@ def state_dict(self) -> dict[str, Any]:
 
         for idx, fn in enumerate(self.lr_lambdas):
             if not isinstance(fn, types.FunctionType):
-                # pyrefly: ignore [unsupported-operation]
                 state_dict["lr_lambdas"][idx] = fn.__dict__.copy()
 
         return state_dict
@@ -1668,7 +1666,6 @@ def __init__(
             self.default_min_lr = None
             self.min_lrs = list(min_lr)
         else:
-            # pyrefly: ignore [bad-assignment]
             self.default_min_lr = min_lr
             self.min_lrs = [min_lr] * len(optimizer.param_groups)
 
@@ -1728,7 +1725,6 @@ def _reduce_lr(self, epoch) -> None:
                     "of the `optimizer` param groups."
                 )
             else:
-                # pyrefly: ignore [bad-assignment]
                 self.min_lrs = [self.default_min_lr] * len(self.optimizer.param_groups)
 
         for i, param_group in enumerate(self.optimizer.param_groups):
diff --git a/torch/optim/optimizer.py b/torch/optim/optimizer.py
index 6c4f56f1fde90..cf9a3c53bf8e8 100644
--- a/torch/optim/optimizer.py
+++ b/torch/optim/optimizer.py
@@ -8,7 +8,7 @@
 from copy import deepcopy
 from itertools import chain
 from typing import Any, cast, overload, TypeAlias, TypeVar
-from typing_extensions import ParamSpec, Self
+from typing_extensions import deprecated, ParamSpec, Self
 
 import torch
 import torch.utils.hooks as hooks
@@ -60,7 +60,6 @@ def _use_grad_for_differentiable(func: Callable[_P, _T]) -> Callable[_P, _T]:
     def _use_grad(*args: _P.args, **kwargs: _P.kwargs) -> _T:
         import torch._dynamo
 
-        # pyrefly: ignore [unsupported-operation]
         self = cast(Optimizer, args[0])  # assume first positional arg is `self`
         prev_grad = torch.is_grad_enabled()
         try:
@@ -134,13 +133,11 @@ def maybe_fallback(*args: _P.args, **kwargs: _P.kwargs):
             if torch.compiler.is_compiling() and (
                 not kwargs.get("capturable", False)
                 and has_state_steps
-                # pyrefly: ignore [unsupported-operation]
                 and (arg := args[state_steps_ind])
                 and isinstance(arg, Sequence)
                 and arg[0].is_cuda
                 or (
                     "state_steps" in kwargs
-                    # pyrefly: ignore [unsupported-operation]
                     and (kwarg := kwargs["state_steps"])
                     and isinstance(kwarg, Sequence)
                     and kwarg[0].is_cuda
@@ -495,6 +492,11 @@ def _accelerator_graph_capture_health_check(self) -> None:
                 )
                 self._warned_capturable_if_run_uncaptured = True
 
+    # Backward compatibility alias for internal callers still using the old name.
+    _cuda_graph_capture_health_check = deprecated(
+        "Use _accelerator_graph_capture_health_check instead",
+    )(_accelerator_graph_capture_health_check)
+
     def _optimizer_step_code(self) -> None:
         """Entry point for `torch.profile.profiler`.
 
diff --git a/torch/optim/radam.py b/torch/optim/radam.py
index 3560b3979df77..d414b9c7edbee 100644
--- a/torch/optim/radam.py
+++ b/torch/optim/radam.py
@@ -323,7 +323,6 @@ def _single_tensor_radam(
         rho_t = rho_inf - 2 * step * (beta2**step) / bias_correction2
 
         def _compute_rect():
-            # pyrefly: ignore [unsupported-operation]
             return (
                 (rho_t - 4)
                 * (rho_t - 2)
@@ -338,7 +337,6 @@ def _compute_adaptive_lr():
             else:
                 exp_avg_sq_sqrt = exp_avg_sq_sqrt.add_(eps)
 
-            # pyrefly: ignore [unsupported-operation]
             return (bias_correction2**0.5) / exp_avg_sq_sqrt
 
         # Compute the variance rectification term and update parameters accordingly
diff --git a/torch/optim/sgd.py b/torch/optim/sgd.py
index 8044d853f0b4e..7c8783fd23182 100644
--- a/torch/optim/sgd.py
+++ b/torch/optim/sgd.py
@@ -350,7 +350,6 @@ def _single_tensor_sgd(
                     # usually this is the differentiable path, which is why the param.clone() is needed
                     grad = grad.addcmul_(param.clone(), weight_decay)
                 else:
-                    # pyrefly: ignore [bad-argument-type]
                     grad = grad.add(param, alpha=weight_decay)
             else:
                 grad = grad.add(param, alpha=weight_decay)
diff --git a/torch/optim/swa_utils.py b/torch/optim/swa_utils.py
index d9cdd39963a98..ec60e81aff9b0 100644
--- a/torch/optim/swa_utils.py
+++ b/torch/optim/swa_utils.py
@@ -1,12 +1,13 @@
 # mypy: allow-untyped-defs
 r"""Implementation for Stochastic Weight Averaging implementation."""
 
+from __future__ import annotations
+
 import itertools
 import math
 import warnings
-from collections.abc import Callable, Iterable
 from copy import deepcopy
-from typing import Any, cast, Literal, Union
+from typing import Any, cast, Literal, TYPE_CHECKING
 from typing_extensions import override
 
 import torch
@@ -15,7 +16,11 @@
 from torch.optim.lr_scheduler import _format_param, LRScheduler
 from torch.utils._foreach_utils import _get_foreach_kernels_supported_devices
 
-from .optimizer import Optimizer
+
+if TYPE_CHECKING:
+    from collections.abc import Callable, Iterable
+
+    from .optimizer import Optimizer
 
 
 __all__ = [
@@ -31,7 +36,7 @@
 from torch.utils._foreach_utils import _group_tensors_by_device_and_dtype
 
 
-PARAM_LIST = Union[tuple[Tensor, ...], list[Tensor]]
+PARAM_LIST = tuple[Tensor, ...] | list[Tensor]
 
 
 def get_ema_multi_avg_fn(decay=0.999):
diff --git a/torch/overrides.py b/torch/overrides.py
index 3470f95e8230e..67920aef7d853 100644
--- a/torch/overrides.py
+++ b/torch/overrides.py
@@ -30,7 +30,7 @@
 import warnings
 from collections.abc import Callable, Iterable
 from functools import wraps
-from typing import Any, TypeVar
+from typing import Any, cast, TypeVar
 from typing_extensions import ParamSpec
 
 import torch
@@ -132,6 +132,7 @@ def get_ignored_functions() -> set[Callable]:
         torch.manual_seed,
         torch.initial_seed,
         torch.seed,
+        torch.thread_safe_generator,
         torch.save,
         torch.load,
         torch.set_printoptions,
@@ -315,9 +316,6 @@ def get_ignored_functions() -> set[Callable]:
         torch.unify_type_list,
         torch.is_warn_always_enabled,
         torch.set_warn_always,
-        torch.vitals_enabled,
-        torch.set_vital,
-        torch.read_vitals,
         torch.vmap,
         torch.cond,
         torch.frombuffer,
@@ -384,6 +382,8 @@ def get_ignored_functions() -> set[Callable]:
         Tensor._addmm_activation,
         Tensor.to_padded_tensor,
         Tensor._use_count,
+        Tensor._philox_normal_,
+        Tensor._philox_uniform_,
     }
 
     if sys.version_info >= (3, 14):
@@ -1567,10 +1567,41 @@ def get_testing_overrides() -> dict[Callable, Callable]:
                 ret2[func] = v
 
     ret.update(ret2)
+
+    # Distributed functions are added after the auto-generation loop above
+    # to avoid generating spurious Tensor method entries (e.g., dist.reduce
+    # would otherwise generate __reduce__ on Tensor).
+    if torch.distributed.is_available():
+        import torch.distributed as dist
+
+        ret.update(
+            {
+                dist.broadcast: lambda tensor, src=None, group=None, async_op=False, group_src=None: -1,
+                dist.all_reduce: lambda tensor, op=None, group=None, async_op=False: -1,
+                dist.reduce: lambda tensor, dst=None, op=None, group=None, async_op=False, group_dst=None: -1,
+                dist.all_reduce_coalesced: lambda tensors, op=None, group=None, async_op=False: -1,
+                dist.all_gather: lambda tensor_list, tensor, group=None, async_op=False: -1,
+                dist.all_gather_into_tensor: lambda output_tensor, input_tensor, group=None, async_op=False: -1,
+                dist.all_gather_coalesced: lambda output_tensor_lists, input_tensor_list, group=None, async_op=False: -1,
+                dist.gather: lambda tensor, gather_list=None, dst=None, group=None, async_op=False, group_dst=None: -1,
+                dist.scatter: lambda tensor, scatter_list=None, src=None, group=None, async_op=False, group_src=None: -1,
+                dist.reduce_scatter: lambda output, input_list, op=None, group=None, async_op=False: -1,
+                dist.reduce_scatter_tensor: lambda output, input, op=None, group=None, async_op=False: -1,
+                dist.all_to_all_single: lambda output, input, output_split_sizes=None, input_split_sizes=None, group=None, async_op=False: -1,  # noqa: B950
+                dist.all_to_all: lambda output_tensor_list, input_tensor_list, group=None, async_op=False: -1,
+                dist.isend: lambda tensor, dst=None, group=None, tag=0, group_dst=None: -1,
+                dist.irecv: lambda tensor, src=None, group=None, tag=0, group_src=None: -1,
+                dist.send: lambda tensor, dst=None, group=None, tag=0, group_dst=None: -1,
+                dist.recv: lambda tensor, src=None, group=None, tag=0, group_src=None: -1,
+            }
+        )  # fmt: skip
+
     return ret
 
 
-def wrap_torch_function(dispatcher: Callable):
+def wrap_torch_function(
+    dispatcher: Callable[_P, Iterable[Any]],
+) -> Callable[[Callable[_P, _R]], Callable[_P, _R]]:
     """Wraps a given function with ``__torch_function__`` -related functionality.
 
     Parameters
@@ -1594,16 +1625,18 @@ def wrap_torch_function(dispatcher: Callable):
     ...     return a + 0
     """
 
-    def inner(func):
+    def inner(func: Callable[_P, _R]) -> Callable[_P, _R]:
         @functools.wraps(func)
-        def wrapped(*args, **kwargs):
+        def wrapped(*args: _P.args, **kwargs: _P.kwargs) -> _R:
             relevant_args = dispatcher(*args, **kwargs)
             if has_torch_function(relevant_args):
-                return handle_torch_function(wrapped, relevant_args, *args, **kwargs)
+                return handle_torch_function(
+                    cast(Callable[_P, _R], wrapped), relevant_args, *args, **kwargs
+                )
 
             return func(*args, **kwargs)
 
-        return wrapped
+        return cast(Callable[_P, _R], wrapped)
 
     return inner
 
@@ -1688,11 +1721,11 @@ def _get_overloaded_args(
 
 
 def handle_torch_function(
-    public_api: Callable,
+    public_api: Callable[_P, _R],
     relevant_args: Iterable[Any],
-    *args,
-    **kwargs,
-) -> Any:
+    *args: _P.args,
+    **kwargs: _P.kwargs,
+) -> _R:
     """Implement a function with checks for ``__torch_function__`` overrides.
 
     See torch::autograd::handle_torch_function for the equivalent of this
diff --git a/torch/package/file_structure_representation.py b/torch/package/file_structure_representation.py
index 2dae130ed6007..b926645e8e84b 100644
--- a/torch/package/file_structure_representation.py
+++ b/torch/package/file_structure_representation.py
@@ -120,10 +120,10 @@ def _create_directory_from_file_list(
 
         file_list (List[str]): List of files to add to the top-level directory.
 
-        include (Union[List[str], str]): An optional pattern that limits what is included from the file_list to
+        include (list[str] | str): An optional pattern that limits what is included from the file_list to
             files whose name matches the pattern.
 
-        exclude (Union[List[str], str]): An optional pattern that excludes files whose name match the pattern.
+        exclude (list[str] | str): An optional pattern that excludes files whose name match the pattern.
 
     Returns:
             :class:`Directory`: a :class:`Directory` file structure representation created from a list of files.
diff --git a/torch/package/glob_group.py b/torch/package/glob_group.py
index 986938cd256ec..3f20d7cbc6542 100644
--- a/torch/package/glob_group.py
+++ b/torch/package/glob_group.py
@@ -1,10 +1,11 @@
 # mypy: allow-untyped-defs
+from __future__ import annotations
+
 import re
 from collections.abc import Iterable
-from typing import Union
 
 
-GlobPattern = Union[str, Iterable[str]]
+GlobPattern = str | Iterable[str]
 
 
 class GlobGroup:
@@ -27,10 +28,10 @@ class GlobGroup:
     none of the ``exclude`` patterns.
 
     Args:
-        include (Union[str, Iterable[str]]): A string or list of strings,
+        include (str | Iterable[str]): A string or list of strings,
             each representing a pattern to be matched against. A candidate
             will match if it matches *any* include pattern
-        exclude (Union[str, Iterable[str]]): A string or list of strings,
+        exclude (str | Iterable[str]): A string or list of strings,
             each representing a pattern to be matched against. A candidate
             will be excluded from matching if it matches *any* exclude pattern.
         separator (str): A string that delimits segments in candidates and
diff --git a/torch/package/package_exporter.py b/torch/package/package_exporter.py
index a18a0bed5881f..fffbaa6648434 100644
--- a/torch/package/package_exporter.py
+++ b/torch/package/package_exporter.py
@@ -822,10 +822,10 @@ def intern(
         included in the package and have its dependencies processed recursively.
 
         Args:
-            include (Union[List[str], str]): A string e.g. "my_package.my_subpackage", or list of strings
+            include (list[str] | str): A string e.g. "my_package.my_subpackage", or list of strings
                 for the names of the modules to be externed. This can also be a glob-style pattern, as described in :meth:`mock`.
 
-            exclude (Union[List[str], str]): An optional pattern that excludes some patterns that match the include string.
+            exclude (list[str] | str): An optional pattern that excludes some patterns that match the include string.
 
             allow_empty (bool): An optional flag that specifies whether the intern modules specified by this call
                 to the ``intern`` method must be matched to some module during packaging. If an ``intern`` module glob
@@ -851,7 +851,7 @@ def mock(
         Use this function to mock this functionality out without having to modify the original code.
 
         Args:
-            include (Union[List[str], str]): A string e.g. ``"my_package.my_subpackage"``, or list of strings
+            include (list[str] | str): A string e.g. ``"my_package.my_subpackage"``, or list of strings
                 for the names of the modules to be mocked out. Strings can also be a glob-style pattern
                 string that may match multiple modules. Any required dependencies that match this pattern
                 string will be mocked out automatically.
@@ -863,7 +863,7 @@ def mock(
                     ``'torch.*'`` -- matches ``'torch.nn'`` or ``'torch.functional'``, but not
                     ``'torch.nn.functional'``
 
-            exclude (Union[List[str], str]): An optional pattern that excludes some patterns that match the include string.
+            exclude (list[str] | str): An optional pattern that excludes some patterns that match the include string.
                 e.g. ``include='torch.**', exclude='torch.foo'`` will mock all torch packages except ``'torch.foo'``,
                 Default: is ``[]``.
 
@@ -891,11 +891,11 @@ def extern(
         Code for extern modules must also exist in the process loading the package.
 
         Args:
-            include (Union[List[str], str]): A string e.g. ``"my_package.my_subpackage"``, or list of strings
+            include (list[str] | str): A string e.g. ``"my_package.my_subpackage"``, or list of strings
                 for the names of the modules to be externed. This can also be a glob-style pattern, as
                 described in :meth:`mock`.
 
-            exclude (Union[List[str], str]): An optional pattern that excludes some patterns that match the
+            exclude (list[str] | str): An optional pattern that excludes some patterns that match the
                 include string.
 
             allow_empty (bool): An optional flag that specifies whether the extern modules specified by this call
@@ -914,10 +914,10 @@ def deny(self, include: "GlobPattern", *, exclude: "GlobPattern" = ()):
         If a dependency on any matching packages is found, a :class:`PackagingError` is raised.
 
         Args:
-            include (Union[List[str], str]): A string e.g. ``"my_package.my_subpackage"``, or list of strings
+            include (list[str] | str): A string e.g. ``"my_package.my_subpackage"``, or list of strings
                 for the names of the modules to be externed. This can also be a glob-style pattern, as described in :meth:`mock`.
 
-            exclude (Union[List[str], str]): An optional pattern that excludes some patterns that match the include string.
+            exclude (list[str] | str): An optional pattern that excludes some patterns that match the include string.
         """
         self.patterns[GlobGroup(include, exclude=exclude)] = _PatternInfo(
             _ModuleProviderAction.DENY, allow_empty=True
diff --git a/torch/package/package_importer.py b/torch/package/package_importer.py
index c87fc7a060b63..a42af10794fa5 100644
--- a/torch/package/package_importer.py
+++ b/torch/package/package_importer.py
@@ -154,6 +154,7 @@ def __init__(
         # used for torch.serialization._load
         self.Unpickler = lambda *args, **kwargs: PackageUnpickler(self, *args, **kwargs)
 
+    # pyrefly: ignore [bad-override]
     def import_module(self, name: str, package=None):
         """Load a module from the package if it hasn't already been loaded, and then return
         the module. Modules are loaded locally
@@ -328,11 +329,11 @@ def file_structure(
         """Returns a file structure representation of package's zipfile.
 
         Args:
-            include (Union[List[str], str]): An optional string e.g. ``"my_package.my_subpackage"``, or optional list of strings
+            include (list[str] | str): An optional string e.g. ``"my_package.my_subpackage"``, or optional list of strings
                 for the names of the files to be included in the zipfile representation. This can also be
                 a glob-style pattern, as described in :meth:`PackageExporter.mock`
 
-            exclude (Union[List[str], str]): An optional pattern that excludes files whose name match the pattern.
+            exclude (list[str] | str): An optional pattern that excludes files whose name match the pattern.
 
         Returns:
             :class:`Directory`
@@ -348,7 +349,7 @@ def python_version(self):
         file later on.
 
         Returns:
-            :class:`Optional[str]` a python version e.g. 3.8.9 or None if no version was stored with this package
+            :class:`str | None` a python version e.g. 3.8.9 or None if no version was stored with this package
         """
         python_version_path = ".data/python_version"
         return (
diff --git a/torch/profiler/_memory_profiler.py b/torch/profiler/_memory_profiler.py
index 3caa490d174fe..39e9452360657 100644
--- a/torch/profiler/_memory_profiler.py
+++ b/torch/profiler/_memory_profiler.py
@@ -1,15 +1,14 @@
 # mypy: allow-untyped-defs
+from __future__ import annotations
+
 import collections
 import dataclasses
 import enum
 import itertools as it
 import logging
-from collections.abc import Iterator
-from typing import Any, cast, Literal, Optional
+from typing import Any, cast, Literal, TYPE_CHECKING
 
 import torch
-from torch._C import FunctionSchema
-from torch._C._autograd import _ProfilerResult
 from torch._C._profiler import (
     _EventType,
     _ExtraFields_Allocation,
@@ -22,6 +21,13 @@
 from torch.profiler import _utils
 
 
+if TYPE_CHECKING:
+    from collections.abc import Iterator
+
+    from torch._C import FunctionSchema
+    from torch._C._autograd import _ProfilerResult
+
+
 KeyAndID = tuple["Key", int]
 TensorAndID = tuple["TensorKey", int]
 
@@ -105,7 +111,7 @@ class TensorKey(Key):
     def __repr__(self) -> str:
         return f"id={self.id}: {repr(self.storage):<24} ({self.device})"
 
-    def __lt__(self, other: "TensorKey") -> bool:
+    def __lt__(self, other: TensorKey) -> bool:
         return self._as_sortable < other._as_sortable
 
     @staticmethod
@@ -114,7 +120,7 @@ def _make(
         storage_ptr: int | None,
         allocation_id: int | None,
         device: torch.device,
-    ) -> Optional["TensorKey"]:
+    ) -> TensorKey | None:
         if (
             tensor_id is not None
             and storage_ptr is not None
@@ -124,11 +130,11 @@ def _make(
         return None
 
     @classmethod
-    def from_allocation(cls, alloc: _ExtraFields_Allocation) -> Optional["TensorKey"]:
+    def from_allocation(cls, alloc: _ExtraFields_Allocation) -> TensorKey | None:
         return cls._make(alloc.id, alloc.ptr, alloc.allocation_id, alloc.device)
 
     @classmethod
-    def from_tensor(cls, t: _TensorMetadata | None) -> Optional["TensorKey"]:
+    def from_tensor(cls, t: _TensorMetadata | None) -> TensorKey | None:
         if t is not None:
             return cls._make(t.id, t.storage_data_ptr, t.allocation_id, t.device)
         return None
@@ -420,7 +426,7 @@ def is_deletion(self) -> bool:
 
 
 class DataFlowNode:
-    def __init__(self, event: _ProfilerEvent, graph: "DataFlowGraph") -> None:
+    def __init__(self, event: _ProfilerEvent, graph: DataFlowGraph) -> None:
         self._event = event
         self._graph = graph
         self._edges: dict[TensorKey, DataFlowEdge] = self._determine_edges()
diff --git a/torch/profiler/profiler.py b/torch/profiler/profiler.py
index 8404083545366..2dc5a63cb33f4 100644
--- a/torch/profiler/profiler.py
+++ b/torch/profiler/profiler.py
@@ -1,14 +1,15 @@
 # mypy: allow-untyped-defs
+from __future__ import annotations
+
 import gzip
 import json
 import os
 import shutil
 import tempfile
 from abc import ABC, abstractmethod
-from collections.abc import Callable, Iterable
 from enum import Enum
 from functools import partial
-from typing import Any, Optional
+from typing import Any, TYPE_CHECKING
 from typing_extensions import deprecated, Self
 from warnings import warn
 
@@ -28,6 +29,10 @@
 from torch.profiler._memory_profiler import MemoryProfile, MemoryProfileTimeline
 
 
+if TYPE_CHECKING:
+    from collections.abc import Callable, Iterable
+
+
 __all__ = [
     "supported_activities",
     "ProfilerAction",
@@ -100,6 +105,33 @@ def cleanup(self):
         pass
 
 
+def _parse_activities(
+    activities: Iterable[ProfilerActivity | dict[ProfilerActivity, list[str]]],
+) -> tuple[set[ProfilerActivity], dict[ProfilerActivity, set[str]]]:
+    """Parse a mixed activities list into a set of activities and a filter dict.
+
+    Each item is either a bare ``ProfilerActivity`` (collect all defaults) or a
+    ``dict[ProfilerActivity, list[str]]`` (collect only the named subset).
+    An empty list value (e.g. ``{CUDA: []}``) means collect nothing for that group.
+    """
+    parsed_activities: set[ProfilerActivity] = set()
+    activity_filters: dict[ProfilerActivity, set[str]] = {}
+    for item in activities:
+        if isinstance(item, ProfilerActivity):
+            if item in parsed_activities:
+                raise ValueError(f"Activity {item} specified more than once")
+            parsed_activities.add(item)
+        elif isinstance(item, dict):
+            for key, val in item.items():
+                if key in parsed_activities:
+                    raise ValueError(f"Activity {key} specified more than once")
+                parsed_activities.add(key)
+                activity_filters[key] = set(val)
+        else:
+            raise TypeError(f"Expected ProfilerActivity or dict, got {type(item)}")
+    return parsed_activities, activity_filters
+
+
 class _KinetoProfile:
     """Low-level profiler wrap the autograd profile
 
@@ -109,6 +141,14 @@ class _KinetoProfile:
             ``torch.profiler.ProfilerActivity.XPU``.
             Default value: ProfilerActivity.CPU and (when available) ProfilerActivity.CUDA
             or (when available) ProfilerActivity.XPU.
+
+            Each item can be a ``ProfilerActivity`` enum (collects all default
+            activity types for that group) or a ``dict`` mapping a ``ProfilerActivity``
+            to a list of individual activity type names to collect, e.g.
+            ``{ProfilerActivity.CUDA: ["GPU_MEMCPY", "CUDA_RUNTIME"]}``.
+            An empty list (e.g. ``{ProfilerActivity.CUDA: []}``) means collect
+            nothing for that group.
+            The same activity group must not appear more than once.
         record_shapes (bool): save information about operator's input shapes.
         profile_memory (bool): track tensor memory allocation/deallocation (see ``export_memory_timeline``
             for more details).
@@ -147,7 +187,8 @@ class _KinetoProfile:
     def __init__(
         self,
         *,
-        activities: Iterable[ProfilerActivity] | None = None,
+        activities: Iterable[ProfilerActivity | dict[ProfilerActivity, list[str]]]
+        | None = None,
         record_shapes: bool = False,
         profile_memory: bool = False,
         with_stack: bool = False,
@@ -159,7 +200,11 @@ def __init__(
         custom_trace_id_callback: Callable[[], str] | None = None,
         post_processing_timeout_s: float | None = None,
     ) -> None:
-        self.activities = set(activities) if activities else supported_activities()
+        if activities is not None:
+            self.activities, self.activity_filters = _parse_activities(activities)
+        else:
+            self.activities = supported_activities()
+            self.activity_filters: dict[ProfilerActivity, set[str]] = {}
         self.record_shapes = record_shapes
         self.with_flops = with_flops
         self.profile_memory = profile_memory
@@ -219,6 +264,9 @@ def prepare_trace(self) -> None:
                 acc_events=self.acc_events,
                 custom_trace_id_callback=self.custom_trace_id_callback,
                 post_processing_timeout_s=self.post_processing_timeout_s,
+                activity_filters=self.activity_filters
+                if self.activity_filters
+                else None,
             )
         if (self.profiler is not None) and (not self.acc_events):
             _warn_once(
@@ -607,6 +655,14 @@ class profile(_KinetoProfile):
             ``torch.profiler.ProfilerActivity.XPU``.
             Default value: ProfilerActivity.CPU and (when available) ProfilerActivity.CUDA
             or (when available) ProfilerActivity.XPU.
+
+            Each item can be a ``ProfilerActivity`` enum (collects all default
+            activity types for that group) or a ``dict`` mapping a ``ProfilerActivity``
+            to a list of individual activity type names to collect, e.g.
+            ``{ProfilerActivity.CUDA: ["GPU_MEMCPY", "CUDA_RUNTIME"]}``.
+            An empty list (e.g. ``{ProfilerActivity.CUDA: []}``) means collect
+            nothing for that group.
+            The same activity group must not appear more than once.
         schedule (Callable): callable that takes step (int) as a single parameter and returns
             ``ProfilerAction`` value that specifies the profiler action to perform at each step.
         on_trace_ready (Callable): callable that is called at each step when ``schedule``
@@ -735,7 +791,8 @@ def trace_handler(prof):
     def __init__(
         self,
         *,
-        activities: Iterable[ProfilerActivity] | None = None,
+        activities: Iterable[ProfilerActivity | dict[ProfilerActivity, list[str]]]
+        | None = None,
         schedule: Callable[[int], ProfilerAction] | None = None,
         on_trace_ready: Callable[..., Any] | None = None,
         record_shapes: bool = False,
@@ -751,7 +808,16 @@ def __init__(
         custom_trace_id_callback: Callable[[], str] | None = None,
         post_processing_timeout_s: float | None = None,
     ) -> None:
-        activities_set = set(activities) if activities else supported_activities()
+        # Extract activities for the use_cuda deprecation check.
+        if activities is not None:
+            activities_set: set[ProfilerActivity] = set()
+            for item in activities:
+                if isinstance(item, ProfilerActivity):
+                    activities_set.add(item)
+                elif isinstance(item, dict):
+                    activities_set.update(item.keys())
+        else:
+            activities_set = supported_activities()
         if use_cuda is not None:
             warn(
                 "`use_cuda` is deprecated, use `activities` argument instead",
@@ -774,8 +840,7 @@ def __init__(
             with_modules=with_modules,
             experimental_config=experimental_config,
             execution_trace_observer=execution_trace_observer
-            if execution_trace_observer
-            else ExecutionTraceObserver.build_execution_trace_obs_from_env(),
+            or ExecutionTraceObserver.build_execution_trace_obs_from_env(),
             acc_events=acc_events,
             custom_trace_id_callback=custom_trace_id_callback,
             post_processing_timeout_s=post_processing_timeout_s,
@@ -966,7 +1031,7 @@ def __del__(self) -> None:
         self.unregister_callback()
 
     @staticmethod
-    def build_execution_trace_obs_from_env() -> Optional["ExecutionTraceObserver"]:
+    def build_execution_trace_obs_from_env() -> ExecutionTraceObserver | None:
         """
         Returns an ExecutionTraceObserver instance if the environment variable
         ENABLE_PYTORCH_EXECUTION_TRACE is set to 1, otherwise returns None.
diff --git a/torch/random.py b/torch/random.py
index 67b649802c659..dd137d3a2d4fb 100644
--- a/torch/random.py
+++ b/torch/random.py
@@ -2,8 +2,25 @@
 import contextlib
 import warnings
 from collections.abc import Generator
+from typing import TYPE_CHECKING
 
 import torch
+
+
+__all__ = [
+    "set_rng_state",
+    "get_rng_state",
+    "manual_seed",
+    "seed",
+    "initial_seed",
+    "fork_rng",
+    "thread_safe_generator",
+]
+
+
+if TYPE_CHECKING:
+    from torch.utils.data._utils.worker import WorkerInfo
+
 from torch._C import default_generator
 
 
@@ -215,3 +232,39 @@ def fork_rng(
         torch.set_rng_state(cpu_rng_state)
         for device, device_rng_state in zip(devices, device_rng_states):
             device_mod.set_rng_state(device_rng_state, device)
+
+
+def thread_safe_generator() -> torch.Generator | None:
+    """Returns a thread-safe random number generator for use in DataLoader workers.
+    This function provides a convenient way for transforms and user code to use
+    thread-safe random number generation without manually checking worker context.
+    When called in a DataLoader thread worker, returns the worker's thread-local
+    :class:`torch.Generator`. When called in the main process or process workers,
+    returns ``None`` (which causes PyTorch functions to use the default global RNG).
+    Returns:
+        Optional[torch.Generator]: Thread-local generator in thread workers, None otherwise.
+    Example::
+        >>> from torch.random import thread_safe_generator
+        >>> generator = thread_safe_generator()
+        >>> torch.randint(0, 10, (5,), generator=generator)
+    Example with transforms::
+        >>> from torch.random import thread_safe_generator
+        >>> class MyRandomTransform:
+        ...     def __call__(self, img):
+        ...         generator = thread_safe_generator()
+        ...         offset = torch.randint(0, 10, (2,), generator=generator)
+        ...         return img[..., offset[0]:, offset[1]:]
+    """
+    # Lazy import to avoid circular dependency during torch module initialization
+    # torch.__init__ loads torch.random early, but torch.utils.data triggers
+    # torch.distributed which needs torch to be fully initialized
+    from torch.utils.data import get_worker_info
+
+    worker_info: WorkerInfo | None = get_worker_info()
+    if (
+        worker_info is not None
+        and worker_info.worker_method == "thread"
+        and worker_info.rng is not None
+    ):
+        return worker_info.rng.torch_generator
+    return None
diff --git a/torch/return_types.py b/torch/return_types.py
index d456742be4b88..4c74fe01537e7 100644
--- a/torch/return_types.py
+++ b/torch/return_types.py
@@ -39,6 +39,7 @@ def structseq_unflatten(values, context):
     globals()[name] = _attr
 
     if not name.startswith("_"):
+        # pyrefly: ignore [unresolvable-dunder-all]
         __all__.append(name)
         all_return_types.append(_attr)
 
diff --git a/torch/serialization.py b/torch/serialization.py
index f6ebe0fce7e4d..32314ae0723e3 100644
--- a/torch/serialization.py
+++ b/torch/serialization.py
@@ -1384,7 +1384,7 @@ def load(
             second step is a no-op if the final location is CPU. When the ``mmap`` flag is set, instead of copying the
             tensor storages from disk to CPU memory in the first step, ``f`` is mapped, which means tensor storages
             will be lazily loaded when their data is accessed.
-        pickle_load_args: (Python 3 only) optional keyword arguments passed over to
+        pickle_load_args: optional keyword arguments passed over to
             :func:`pickle_module.load` and :func:`pickle_module.Unpickler`,
             only works if :attr:`weights_only=False`, e.g., :attr:`errors=...`.
 
diff --git a/torch/sparse/__init__.py b/torch/sparse/__init__.py
index 92ed11591c420..78846265e9669 100644
--- a/torch/sparse/__init__.py
+++ b/torch/sparse/__init__.py
@@ -1,7 +1,7 @@
 # mypy: allow-untyped-defs
 # The Tensor classes are added to this module by python_tensor.cpp
 # A workaround to support both TorchScript and MyPy:
-from typing import Any, Optional, TYPE_CHECKING, Union
+from typing import Any, TYPE_CHECKING
 
 import torch
 from torch import Tensor
@@ -19,11 +19,11 @@
 if TYPE_CHECKING:
     from torch.types import _dtype as DType
 
-    DimOrDims = Optional[int | tuple[int, ...] | list[int]]
+    DimOrDims = int | tuple[int, ...] | list[int] | None
 else:
     # The JIT doesn't understand Union, nor torch.dtype here
     DType = int
-    DimOrDims = Optional[tuple[int]]
+    DimOrDims = tuple[int] | None
 
 
 __all__ = [
@@ -642,20 +642,27 @@ def convert_to_strided_representation(args):
                     if obj.layout is torch.sparse_coo:
                         # pyrefly: ignore [no-matching-overload]
                         d.update(
-                            indices=obj._indices(), is_coalesced=obj.is_coalesced()
+                            # pyrefly: ignore [bad-argument-type]
+                            indices=obj._indices(),
+                            # pyrefly: ignore [bad-argument-type]
+                            is_coalesced=obj.is_coalesced(),
                         )
                         values = obj._values()
                     elif obj.layout in {torch.sparse_csr, torch.sparse_bsr}:
                         # pyrefly: ignore [no-matching-overload]
                         d.update(
+                            # pyrefly: ignore [bad-argument-type]
                             compressed_indices=obj.crow_indices(),
+                            # pyrefly: ignore [bad-argument-type]
                             plain_indices=obj.col_indices(),
                         )
                         values = obj.values()
                     else:
                         # pyrefly: ignore [no-matching-overload]
                         d.update(
+                            # pyrefly: ignore [bad-argument-type]
                             compressed_indices=obj.ccol_indices(),
+                            # pyrefly: ignore [bad-argument-type]
                             plain_indices=obj.row_indices(),
                         )
                         values = obj.values()
diff --git a/torch/sparse/_semi_structured_ops.py b/torch/sparse/_semi_structured_ops.py
index 18e4c6abfc638..0254a33a4bc62 100644
--- a/torch/sparse/_semi_structured_ops.py
+++ b/torch/sparse/_semi_structured_ops.py
@@ -15,6 +15,9 @@
     "semi_sparse_addmm",
     "semi_sparse_linear",
     "semi_sparse_scaled_mm",
+    "semi_sparse_clone",
+    "semi_sparse_to",
+    "semi_sparse_to_copy",
 ]
 
 
@@ -136,20 +139,14 @@ def semi_sparse_mm(func, types, args=(), kwargs=None) -> torch.Tensor:
             "`SparseSemiStructuredTensor` matmul: Broadcasting is not implemented"
         )
     if isinstance(A, torch.sparse.SparseSemiStructuredTensor):
-        row, col = B.shape
-        B_padded = A._pad_dense_input(B)
-        res = A._mm(B_padded)
-        return res[:, :col]
+        return A._mm(B)
     else:
         B_t = B.t()
         if not isinstance(B_t, torch.sparse.SparseSemiStructuredTensor):
             raise AssertionError(
                 f"expected SparseSemiStructuredTensor, got {type(B_t).__name__}"
             )
-        row, col = A.shape
-        A_padded = B._pad_dense_input(A)
-        res = B_t._mm(A_padded.t()).t()
-        return res[:row, :]
+        return B_t._mm(A, should_transpose_dense=True).t()
 
 
 def semi_sparse_addmm(func, types, args=(), kwargs=None) -> torch.Tensor:
@@ -174,9 +171,7 @@ def semi_sparse_addmm(func, types, args=(), kwargs=None) -> torch.Tensor:
             f"expected SparseSemiStructuredTensor, got {type(B_t).__name__}"
         )
     row, _col = A.shape
-    A_padded = B_t._pad_dense_input(A)
-    result = B_t._mm(A_padded.t(), bias=bias).t()
-    return result[:row, :]
+    return B_t._mm(A, bias=bias, should_transpose_dense=True).t()
 
 
 def semi_sparse_linear(func, types, args=(), kwargs=None) -> torch.Tensor:
@@ -187,7 +182,6 @@ def semi_sparse_linear(func, types, args=(), kwargs=None) -> torch.Tensor:
 
     shape = A.shape
     A_2d = A.view(-1, shape[-1])
-
     if bias is None:
         res = A_2d @ B.t()
     else:
@@ -196,7 +190,6 @@ def semi_sparse_linear(func, types, args=(), kwargs=None) -> torch.Tensor:
             types=None,
             args=[bias, A_2d, B.t()],
         )
-
     return res.view(*shape[:-1], -1)
 
 
@@ -234,3 +227,74 @@ def semi_sparse_scaled_mm(func, types, args=(), kwargs=None) -> torch.Tensor:
         out_dtype=out_dtype,
     )
     return sparse_result
+
+
+def semi_sparse_clone(func, types, args=(), kwargs=None) -> torch.Tensor:
+    if len(args) != 1:
+        raise AssertionError(f"expected 1 arg, got {len(args)}")
+
+    self = args[0]
+    if not isinstance(self, torch.sparse.SparseSemiStructuredTensor):
+        raise AssertionError(
+            f"expected SparseSemiStructuredTensor, got {type(self).__name__}"
+        )
+
+    # pyrefly: ignore [no-matching-overload]
+    return self.__class__(
+        shape=self.shape,
+        packed=None if self.packed is None else self.packed.clone(),
+        meta=None if self.meta is None else self.meta.clone(),
+        packed_t=None if self.packed_t is None else self.packed_t.clone(),
+        meta_t=None if self.meta_t is None else self.meta_t.clone(),
+        compressed_swizzled_bitmask=(
+            None
+            if self.compressed_swizzled_bitmask is None
+            else self.compressed_swizzled_bitmask.clone()
+        ),
+        fuse_transpose_cusparselt=self.fuse_transpose_cusparselt,
+        alg_id_cusparselt=self.alg_id_cusparselt,
+        requires_grad=self.requires_grad,
+    )
+
+
+def semi_sparse_to_copy(func, types, args, kwargs=None) -> torch.Tensor:
+    self = args[0]
+    kwargs = kwargs or {}
+
+    device = kwargs.get("device", None)
+
+    if device is not None and torch.device(device).type == "cpu":
+        dense = self.to_dense()
+        return func(dense, **kwargs)
+
+    raise NotImplementedError(
+        f"`_to_copy()` with kwargs={kwargs} is not implemented "
+        "for SparseSemiStructuredTensor. Only converting to CPU is supported currently."
+    )
+
+
+def semi_sparse_to(func, types, args, kwargs=None) -> torch.Tensor:
+    self = args[0]
+    remaining_args = args[1:]
+    kwargs = kwargs or {}
+
+    # Determine the target device from args/kwargs
+    device = None
+    if remaining_args:
+        first_arg = remaining_args[0]
+        if isinstance(first_arg, (torch.device, str)):
+            try:
+                device = torch.device(first_arg)
+            except RuntimeError:
+                pass
+    if "device" in kwargs:
+        device = torch.device(kwargs["device"])
+
+    if device is not None and device.type == "cpu":
+        dense = self.to_dense()
+        return func(dense, *remaining_args, **kwargs)
+
+    raise NotImplementedError(
+        f"`to()` with args={remaining_args}, kwargs={kwargs} is not implemented "
+        "for SparseSemiStructuredTensor. Only `to('cpu')` is supported currently."
+    )
diff --git a/torch/sparse/_triton_ops.py b/torch/sparse/_triton_ops.py
index c50a633f123bd..106306c6e4cde 100644
--- a/torch/sparse/_triton_ops.py
+++ b/torch/sparse/_triton_ops.py
@@ -1339,28 +1339,17 @@ def kernel(grid, *sliced_tensors):
         # pyrefly: ignore [unsupported-operation]
         _bsr_strided_addmm_kernel[grid](
             *ptr_stride_extractor(*sliced_tensors),
-            # pyrefly: ignore [bad-argument-count]
             beta,
             alpha,
-            # pyrefly: ignore [bad-keyword-argument, bad-argument-type]
             beta_is_one=beta == 1,
-            # pyrefly: ignore [bad-keyword-argument, bad-argument-type]
             beta_is_nonzero=beta != 0,
-            # pyrefly: ignore [bad-keyword-argument, bad-argument-type]
             alpha_is_one=alpha == 1,
-            # pyrefly: ignore [bad-keyword-argument, bad-argument-type]
             left_alpha_is_one=left_alpha_is_one,
-            # pyrefly: ignore [bad-keyword-argument, bad-argument-type]
             right_alpha_is_one=right_alpha_is_one,
-            # pyrefly: ignore [bad-keyword-argument, bad-argument-type]
             BLOCKSIZE_ROW=BM,
-            # pyrefly: ignore [bad-keyword-argument, bad-argument-type]
             BLOCKSIZE_INNER=BK,
-            # pyrefly: ignore [bad-keyword-argument]
             BLOCKSIZE_COL=BN,
-            # pyrefly: ignore [bad-keyword-argument, bad-argument-type]
             allow_tf32=dot_out_dtype == tl.float32,
-            # pyrefly: ignore [bad-keyword-argument, bad-argument-type]
             acc_dtype=dot_out_dtype,
             **meta,
         )
@@ -1477,7 +1466,7 @@ def _sampled_addmm_kernel(
 
                 mat1_block = tl.load(
                     mat1_block_ptrs + mat1_col_block_stride * k_offsets[None, :],
-                    # pyrefly: ignore [bad-index, index-error]
+                    # pyrefly: ignore [bad-index]
                     mask=mask_k[None, :],
                     other=0.0,
                 )
@@ -1486,7 +1475,7 @@ def _sampled_addmm_kernel(
                     mat2_block_ptrs
                     + mat2_tiled_col_stride * col_block
                     + mat2_row_block_stride * k_offsets[:, None],
-                    # pyrefly: ignore [bad-index, index-error]
+                    # pyrefly: ignore [bad-index]
                     mask=mask_k[:, None],
                     other=0.0,
                 )
@@ -1681,17 +1670,12 @@ def kernel(grid, *sliced_tensors):
                 beta,
                 is_beta_zero,
                 *blocksize,
-                # pyrefly: ignore [bad-argument-count]
                 k,
                 tile_k,
                 *ptr_stride_extractor(*sliced_tensors),
-                # pyrefly: ignore [bad-keyword-argument, bad-argument-type]
                 acc_dtype=acc_dtype,
-                # pyrefly: ignore [bad-keyword-argument, bad-argument-type]
                 allow_tf32=allow_tf32,
-                # pyrefly: ignore [unexpected-keyword]
                 num_stages=1,
-                # pyrefly: ignore [unexpected-keyword]
                 num_warps=4,
             )
 
@@ -1976,7 +1960,6 @@ def bsr_softmax(input, max_row_nnz=None):
         def kernel(grid, *sliced_tensors):
             _bsr_softmax_kernel[grid](
                 *ptr_stride_extractor(*sliced_tensors),
-                # pyrefly: ignore [bad-argument-count]
                 row_block,
                 col_block,
                 max_row_nnz,
@@ -2151,11 +2134,8 @@ def grid(META):
         if "allow_tf32" not in meta:
             meta.update(allow_tf32=dot_out_dtype == tl.float32)
         _scatter_mm2_kernel[grid](
-            # pyrefly: ignore [bad-argument-type]
             M,
-            # pyrefly: ignore [bad-argument-type]
             K,
-            # pyrefly: ignore [bad-argument-type]
             N,
             blocks,
             blocks.stride(0),
@@ -2174,9 +2154,7 @@ def grid(META):
             pq_indices,
             pq_indices.stride(0),
             pq_indices.stride(1),
-            # pyrefly: ignore [bad-argument-type]
             dot_out_dtype=dot_out_dtype,
-            # pyrefly: ignore [bad-argument-type]
             **meta,
         )
 
@@ -2373,7 +2351,6 @@ def grid(META):
         _scatter_mm6_kernel[grid](
             B,
             Ms,
-            # pyrefly: ignore [bad-argument-type]
             Ks,
             N,
             blocks,
@@ -2392,7 +2369,6 @@ def grid(META):
             r_offsets,
             p_offsets,
             q_offsets,
-            # pyrefly: ignore [bad-argument-type]
             dot_out_dtype=dot_out_dtype,
             **meta,
         )
diff --git a/torch/sparse/_triton_ops_meta.py b/torch/sparse/_triton_ops_meta.py
index 9a9999460cab3..64f9ca15ca8da 100644
--- a/torch/sparse/_triton_ops_meta.py
+++ b/torch/sparse/_triton_ops_meta.py
@@ -241,10 +241,8 @@ def sort_key(key):
     part2 = current_content[end_data_index:]
     data_part = []
     for op_key in sorted(_operation_device_version_data, key=sort_key):
-        # pyrefly: ignore [bad-argument-type]
         data_part.append("    " + repr(op_key).replace("'", '"') + ": {")
         op_data = _operation_device_version_data[op_key]
-        # pyrefly: ignore [bad-argument-type]
         data_part.extend(f"        {key}: {op_data[key]}," for key in sorted(op_data))
         data_part.append("    },")
     new_content = part1 + "\n".join(data_part) + "\n" + part2
diff --git a/torch/sparse/semi_structured.py b/torch/sparse/semi_structured.py
index 0ddfd3fe09298..4d99dac9eadef 100644
--- a/torch/sparse/semi_structured.py
+++ b/torch/sparse/semi_structured.py
@@ -12,12 +12,15 @@
 from torch.sparse._semi_structured_ops import (
     fallback_dispatcher,
     semi_sparse_addmm,
+    semi_sparse_clone,
     semi_sparse_detach,
     semi_sparse_indices,
     semi_sparse_linear,
     semi_sparse_mm,
     semi_sparse_scaled_mm,
     semi_sparse_t,
+    semi_sparse_to,
+    semi_sparse_to_copy,
     semi_sparse_values,
     semi_sparse_view,
 )
@@ -230,8 +233,10 @@ def _load_dispatch_table(cls, custom_dispatch_table=None) -> None:
                 torch.ops.aten.matmul: semi_sparse_mm,
                 torch.ops.aten.addmm: semi_sparse_addmm,
                 torch.ops.aten.linear: semi_sparse_linear,
-                torch.ops.aten._to_copy: fallback_dispatcher,
+                torch.ops.aten._to_copy: semi_sparse_to_copy,
                 torch.ops.aten._scaled_mm: semi_sparse_scaled_mm,
+                torch.ops.aten.clone: semi_sparse_clone,
+                torch.ops.aten.to: semi_sparse_to,
             }
             if custom_dispatch_table is not None:
                 cls.SPARSE_DISPATCH.update(custom_dispatch_table)
@@ -279,35 +284,16 @@ def _validate_device_dim_dtype_shape(cls, original_tensor: torch.Tensor) -> None
                 f"Both dimensions must be larger or equal than and a multiple of ({min_rows}, {min_cols})"
             )
 
-    @classmethod
-    def _pad_dense_input(cls, dense_input: torch.Tensor) -> torch.Tensor:
-        """
-        Calculates padding for dense tensor and pads tensor if necessary.
-        If padding is not required, this function returns the original tensor.
-        """
-        # only 2d matmul
-        if dense_input.dim() != 2:
-            raise AssertionError(f"dense_input must be 2D, got {dense_input.dim()}D")
-
-        # check shape
-        m, n = dense_input.shape
-        min_rows = cls._DTYPE_SHAPE_CONSTRAINTS[dense_input.dtype].dense_min_rows
-        min_cols = cls._DTYPE_SHAPE_CONSTRAINTS[dense_input.dtype].dense_min_cols
-
-        # calculate padding
-        to_pad_m = -m % min_rows if m < min_rows or m % min_rows else 0
-        to_pad_n = -n % min_cols if n < min_cols or n % min_rows else 0
-        if to_pad_m or to_pad_n:
-            return torch.nn.functional.pad(dense_input, (0, to_pad_n, 0, to_pad_m))
-        else:
-            return dense_input
-
     def to_dense(self):  # type:ignore[override]
         col = self.shape[-1]
         return torch.mm(self, torch.eye(col, dtype=self.dtype, device=self.device))
 
     @classmethod
-    def from_dense(cls, original_tensor: torch.Tensor) -> "SparseSemiStructuredTensor":
+    def from_dense(
+        cls,
+        original_tensor: torch.Tensor,
+        alg_id: int = _DEFAULT_ALG_ID,
+    ) -> "SparseSemiStructuredTensor":
         raise NotImplementedError
 
     def _mm(
@@ -323,6 +309,7 @@ def _mm(
 def to_sparse_semi_structured(
     original_tensor: torch.Tensor,
     transposed: bool = False,
+    alg_id: int = SparseSemiStructuredTensor._DEFAULT_ALG_ID,
 ) -> SparseSemiStructuredTensor:
     """
     This function converts a dense tensor into a sparse semi-structured tensor.
@@ -336,6 +323,8 @@ def to_sparse_semi_structured(
     Args:
         original_tensor (Tensor): the dense tensor to convert
         transposed (bool, optional): deprecated arg to be removed in another release. Do not use.
+        alg_id (int, optional): the algorithm id to use for cuSPARSELt matmul. Defaults to 0.
+            Can be obtained via ``torch._cslt_sparse_mm_search``.
     Returns:
         SparseSemiStructuredTensor: A sparse semi-structured tensor created from the given original_tensor
     Raises:
@@ -385,7 +374,7 @@ def to_sparse_semi_structured(
         else torch.sparse.SparseSemiStructuredTensorCUSPARSELT
     )
 
-    return SPARSE_SUBCLASS.from_dense(original_tensor)
+    return SPARSE_SUBCLASS.from_dense(original_tensor, alg_id=alg_id)
 
 
 class SparseSemiStructuredTensorCUTLASS(SparseSemiStructuredTensor):
@@ -410,7 +399,9 @@ class SparseSemiStructuredTensorCUTLASS(SparseSemiStructuredTensor):
 
     @classmethod
     def from_dense(
-        cls, original_tensor: torch.Tensor
+        cls,
+        original_tensor: torch.Tensor,
+        alg_id: int = SparseSemiStructuredTensor._DEFAULT_ALG_ID,
     ) -> "SparseSemiStructuredTensorCUTLASS":
         cls._validate_device_dim_dtype_shape(original_tensor)
         (
@@ -517,7 +508,12 @@ def prune_dense_static_sort(
         )
 
     def _mm(
-        self, B: torch.Tensor, *, bias: torch.Tensor | None = None, **kwargs
+        self,
+        B: torch.Tensor,
+        *,
+        bias: torch.Tensor | None = None,
+        should_transpose_dense: bool = False,
+        **kwargs,
     ) -> torch.Tensor:
         if isinstance(B, SparseSemiStructuredTensor):
             raise ValueError(
@@ -533,13 +529,18 @@ def _mm(
                 f"`{cls_name}` matmul: operation is not supported"
             )
         else:
-            if bias is None:
-                res = torch._sparse_semi_structured_mm(self.packed, self.meta, B)
-            else:
-                res = torch._sparse_semi_structured_addmm(
-                    bias, self.packed, self.meta, B
-                )
-            return res[: self.shape[0]]
+            _ensure_cutlass_mm_registered()
+            constraints = self._DTYPE_SHAPE_CONSTRAINTS[B.dtype]
+            return torch.ops.semi_structured.cutlass_mm(
+                B,
+                self.packed,
+                self.meta,
+                bias,
+                self.shape[0],
+                constraints.dense_min_rows,
+                constraints.dense_min_cols,
+                should_transpose_dense,
+            )
 
 
 class SparseSemiStructuredTensorCUSPARSELT(SparseSemiStructuredTensor):
@@ -564,7 +565,9 @@ class SparseSemiStructuredTensorCUSPARSELT(SparseSemiStructuredTensor):
 
     @classmethod
     def from_dense(
-        cls, original_tensor: torch.Tensor
+        cls,
+        original_tensor: torch.Tensor,
+        alg_id: int = SparseSemiStructuredTensor._DEFAULT_ALG_ID,
     ) -> "SparseSemiStructuredTensorCUSPARSELT":
         cls._validate_device_dim_dtype_shape(original_tensor)
         # pyrefly: ignore [no-matching-overload]
@@ -576,7 +579,7 @@ def from_dense(
             meta_t=None,
             compressed_swizzled_bitmask=None,
             fuse_transpose_cusparselt=SparseSemiStructuredTensor._FUSE_TRANSPOSE,
-            alg_id_cusparselt=SparseSemiStructuredTensor._DEFAULT_ALG_ID,
+            alg_id_cusparselt=alg_id,
             requires_grad=original_tensor.requires_grad,
         )
 
@@ -646,7 +649,12 @@ def prune_dense_static_sort(
         )
 
     def _mm(
-        self, B: torch.Tensor, *, bias: torch.Tensor | None = None, **kwargs
+        self,
+        B: torch.Tensor,
+        *,
+        bias: torch.Tensor | None = None,
+        should_transpose_dense: bool = False,
+        **kwargs,
     ) -> torch.Tensor:
         if isinstance(B, SparseSemiStructuredTensor):
             raise ValueError(
@@ -680,11 +688,153 @@ def _mm(
                 f"`{self.__class__.__name__}` matmul: operation is not supported"
             )
         else:
-            res = torch._cslt_sparse_mm(
-                self.packed,
+            _ensure_cusparselt_mm_registered()
+            constraints = self._DTYPE_SHAPE_CONSTRAINTS[B.dtype]
+            return torch.ops.semi_structured.cusparselt_mm(
                 B,
-                bias=bias,
-                transpose_result=self.fuse_transpose_cusparselt,
-                alg_id=self.alg_id_cusparselt,
+                self.packed,
+                bias,
+                self.shape[0],
+                constraints.dense_min_rows,
+                constraints.dense_min_cols,
+                self.fuse_transpose_cusparselt,
+                self.alg_id_cusparselt,
+                should_transpose_dense,
             )
-            return res.t() if self.fuse_transpose_cusparselt else res
+
+
+_cutlass_mm_registered = False
+
+
+def _ensure_cutlass_mm_registered():
+    """Lazily register the cutlass_mm custom op.
+
+    Registration is deferred to avoid importing torch.library at module load
+    time, since torch.sparse is imported early during ``import torch``.
+    """
+    global _cutlass_mm_registered
+    if _cutlass_mm_registered:
+        return
+    _cutlass_mm_registered = True
+
+    from torch.library import custom_op
+
+    @custom_op("semi_structured::cutlass_mm", mutates_args=())
+    def cutlass_mm(
+        dense: torch.Tensor,
+        packed: torch.Tensor,
+        meta: torch.Tensor,
+        bias: torch.Tensor | None,
+        out_features: int,
+        min_rows: int,
+        min_cols: int,
+        should_transpose_dense: bool,
+    ) -> torch.Tensor:
+        m, n = dense.shape
+        to_pad_m = (-m) % min_rows
+        to_pad_n = (-n) % min_cols
+        need_pad = to_pad_m != 0 or to_pad_n != 0
+        dense_padded = dense
+        if need_pad:
+            dense_padded = torch.nn.functional.pad(dense, (0, to_pad_n, 0, to_pad_m))
+        mm_input = dense_padded.t() if should_transpose_dense else dense_padded
+        if bias is None:
+            res = torch._sparse_semi_structured_mm(packed, meta, mm_input)
+        else:
+            res = torch._sparse_semi_structured_addmm(bias, packed, meta, mm_input)
+        if need_pad:
+            out_cols = m if should_transpose_dense else n
+            return (
+                res[:out_features]
+                .narrow(1, 0, out_cols)
+                .clone(memory_format=torch.contiguous_format)
+            )
+        return res.contiguous()
+
+    @cutlass_mm.register_fake
+    def _cutlass_mm_fake(
+        dense: torch.Tensor,
+        packed: torch.Tensor,
+        meta: torch.Tensor,
+        bias: torch.Tensor | None,
+        out_features: int,
+        min_rows: int,
+        min_cols: int,
+        transpose_dense: bool,
+    ) -> torch.Tensor:
+        out_cols = dense.shape[0] if transpose_dense else dense.shape[1]
+        return torch.empty(
+            out_features,
+            out_cols,
+            dtype=dense.dtype,
+            device=dense.device,
+        )
+
+
+_cusparselt_mm_registered = False
+
+
+def _ensure_cusparselt_mm_registered():
+    """Lazily register the cusparselt_mm custom op."""
+    global _cusparselt_mm_registered
+    if _cusparselt_mm_registered:
+        return
+    _cusparselt_mm_registered = True
+
+    from torch.library import custom_op
+
+    @custom_op("semi_structured::cusparselt_mm", mutates_args=())
+    def cusparselt_mm(
+        dense: torch.Tensor,
+        packed: torch.Tensor,
+        bias: torch.Tensor | None,
+        out_features: int,
+        min_rows: int,
+        min_cols: int,
+        fuse_transpose: bool,
+        alg_id: int,
+        should_transpose_dense: bool = False,
+    ) -> torch.Tensor:
+        m, n = dense.shape
+        to_pad_m = (-m) % min_rows
+        to_pad_n = (-n) % min_cols
+        need_pad = to_pad_m != 0 or to_pad_n != 0
+        dense_padded = dense
+        if need_pad:
+            dense_padded = torch.nn.functional.pad(dense, (0, to_pad_n, 0, to_pad_m))
+        mm_input = dense_padded.t() if should_transpose_dense else dense_padded
+        res = torch._cslt_sparse_mm(
+            packed,
+            mm_input,
+            bias=bias,
+            transpose_result=fuse_transpose,
+            alg_id=alg_id,
+        )
+        if fuse_transpose:
+            res = res.t()
+        if need_pad:
+            out_cols = m if should_transpose_dense else n
+            return res.narrow(1, 0, out_cols).clone(
+                memory_format=torch.contiguous_format
+            )
+        return res.contiguous()
+
+    @cusparselt_mm.register_fake
+    def _cusparselt_mm_fake(
+        dense: torch.Tensor,
+        packed: torch.Tensor,
+        bias: torch.Tensor | None,
+        out_features: int,
+        min_rows: int,
+        min_cols: int,
+        fuse_transpose: bool,
+        alg_id: int,
+        should_transpose_dense: bool,
+    ) -> torch.Tensor:
+        out_cols = dense.shape[0] if should_transpose_dense else dense.shape[1]
+        return torch.empty(
+            out_features,
+            out_cols,
+            dtype=dense.dtype,
+            device=dense.device,
+        )
diff --git a/torch/testing/_comparison.py b/torch/testing/_comparison.py
index 0f3f7f3e02aa7..ef60dd19f7c3e 100644
--- a/torch/testing/_comparison.py
+++ b/torch/testing/_comparison.py
@@ -4,7 +4,7 @@
 import collections.abc
 import contextlib
 from collections.abc import Callable, Collection, Sequence
-from typing import Any, NoReturn, Optional, Union
+from typing import Any, NoReturn
 from typing_extensions import deprecated
 
 import torch
@@ -18,6 +18,37 @@
     HAS_NUMPY = False
     np = None  # type: ignore[assignment]
 
+_HAS_DTENSOR = torch.distributed.is_available()
+
+
+def _unwrap_dtensor_for_comparison(actual, expected):
+    """Handle DTensor inputs for assertEqual/assert_close."""
+    if not _HAS_DTENSOR:
+        return actual, expected
+    from torch.distributed.tensor import DTensor
+
+    actual_dt = isinstance(actual, DTensor)
+    expected_dt = isinstance(expected, DTensor)
+    if actual_dt and expected_dt:
+        if actual.placements != expected.placements:
+            raise AssertionError(
+                f"DTensor placements do not match: "
+                f"{actual.placements} != {expected.placements}"
+            )
+        if actual.device_mesh != expected.device_mesh:
+            raise AssertionError(
+                f"DTensor device meshes do not match: "
+                f"{actual.device_mesh} != {expected.device_mesh}"
+            )
+        return actual.to_local(), expected.to_local()
+    elif actual_dt != expected_dt:
+        raise TypeError(
+            "Comparing a DTensor to a non-DTensor is ambiguous. "
+            "Call .full_tensor() to compare the full logical tensor "
+            "or .to_local() to compare the local shard."
+        )
+    return actual, expected
+
 
 class ErrorMeta(Exception):
     """Internal testing exception that makes that carries error metadata."""
@@ -35,9 +66,7 @@ def __init__(
         self.msg = msg
         self.id = id
 
-    def to_error(
-        self, msg: Optional[Union[str, Callable[[str], str]]] = None
-    ) -> Exception:
+    def to_error(self, msg: str | Callable[[str], str] | None = None) -> Exception:
         if not isinstance(msg, str):
             generated_msg = self.msg
             if self.id:
@@ -71,8 +100,8 @@ def to_error(
 
 
 def default_tolerances(
-    *inputs: Union[torch.Tensor, torch.dtype],
-    dtype_precisions: Optional[dict[torch.dtype, tuple[float, float]]] = None,
+    *inputs: torch.Tensor | torch.dtype,
+    dtype_precisions: dict[torch.dtype, tuple[float, float]] | None = None,
 ) -> tuple[float, float]:
     """Returns the default absolute and relative testing tolerances for a set of inputs based on the dtype.
 
@@ -99,9 +128,9 @@ def default_tolerances(
 
 
 def get_tolerances(
-    *inputs: Union[torch.Tensor, torch.dtype],
-    rtol: Optional[float],
-    atol: Optional[float],
+    *inputs: torch.Tensor | torch.dtype,
+    rtol: float | None,
+    atol: float | None,
     id: tuple[Any, ...] = (),
 ) -> tuple[float, float]:
     """Gets absolute and relative to be used for numeric comparisons.
@@ -133,9 +162,9 @@ def get_tolerances(
 def _make_bitwise_mismatch_msg(
     *,
     default_identifier: str,
-    identifier: Optional[Union[str, Callable[[str], str]]] = None,
-    extra: Optional[str] = None,
-    first_mismatch_idx: Optional[tuple[int, ...]] = None,
+    identifier: str | Callable[[str], str] | None = None,
+    extra: str | None = None,
+    first_mismatch_idx: tuple[int, ...] | None = None,
 ):
     """Makes a mismatch error message for bitwise values.
 
@@ -164,13 +193,13 @@ def _make_bitwise_mismatch_msg(
 def _make_mismatch_msg(
     *,
     default_identifier: str,
-    identifier: Optional[Union[str, Callable[[str], str]]] = None,
-    extra: Optional[str] = None,
+    identifier: str | Callable[[str], str] | None = None,
+    extra: str | None = None,
     abs_diff: float,
-    abs_diff_idx: Optional[Union[int, tuple[int, ...]]] = None,
+    abs_diff_idx: int | tuple[int, ...] | None = None,
     atol: float,
     rel_diff: float,
-    rel_diff_idx: Optional[Union[int, tuple[int, ...]]] = None,
+    rel_diff_idx: int | tuple[int, ...] | None = None,
     rtol: float,
 ) -> str:
     """Makes a mismatch error message for numeric values.
@@ -196,7 +225,7 @@ def make_diff_msg(
         *,
         type: str,
         diff: float,
-        idx: Optional[Union[int, tuple[int, ...]]],
+        idx: int | tuple[int, ...] | None,
         tol: float,
     ) -> str:
         if idx is None:
@@ -224,12 +253,12 @@ def make_diff_msg(
 
 
 def make_scalar_mismatch_msg(
-    actual: Union[bool, int, float, complex],
-    expected: Union[bool, int, float, complex],
+    actual: bool | int | float | complex,
+    expected: bool | int | float | complex,
     *,
     rtol: float,
     atol: float,
-    identifier: Optional[Union[str, Callable[[str], str]]] = None,
+    identifier: str | Callable[[str], str] | None = None,
 ) -> str:
     """Makes a mismatch error message for scalars.
 
@@ -263,7 +292,7 @@ def make_tensor_mismatch_msg(
     *,
     rtol: float,
     atol: float,
-    identifier: Optional[Union[str, Callable[[str], str]]] = None,
+    identifier: str | Callable[[str], str] | None = None,
 ):
     """Makes a mismatch error message for tensors.
 
@@ -374,7 +403,7 @@ def _inputs_not_supported() -> NoReturn:
         raise UnsupportedInputs
 
     @staticmethod
-    def _check_inputs_isinstance(*inputs: Any, cls: Union[type, tuple[type, ...]]):
+    def _check_inputs_isinstance(*inputs: Any, cls: type | tuple[type, ...]):
         """Checks if all inputs are instances of a given class and raise :class:`UnsupportedInputs` otherwise."""
         if not all(isinstance(input, cls) for input in inputs):
             Pair._inputs_not_supported()
@@ -395,7 +424,7 @@ def _fail(
     def compare(self) -> None:
         """Compares the inputs and raises an :class`ErrorMeta` in case they mismatch."""
 
-    def extra_repr(self) -> Sequence[Union[str, tuple[str, Any]]]:
+    def extra_repr(self) -> Sequence[str | tuple[str, Any]]:
         """Returns extra information that will be included in the representation.
 
         Should be overwritten by all subclasses that use additional options. The representation of the object will only
@@ -561,8 +590,8 @@ def __init__(
         expected: Any,
         *,
         id: tuple[Any, ...] = (),
-        rtol: Optional[float] = None,
-        atol: Optional[float] = None,
+        rtol: float | None = None,
+        atol: float | None = None,
         equal_nan: bool = False,
         check_dtype: bool = False,
         **other_parameters: Any,
@@ -589,7 +618,7 @@ def _supported_types(self) -> tuple[type, ...]:
 
     def _process_inputs(
         self, actual: Any, expected: Any, *, id: tuple[Any, ...]
-    ) -> tuple[Union[int, float, complex], Union[int, float, complex]]:
+    ) -> tuple[int | float | complex, int | float | complex]:
         self._check_inputs_isinstance(actual, expected, cls=self._supported_types)
         actual, expected = (
             self._to_number(number_like, id=id) for number_like in (actual, expected)
@@ -598,7 +627,7 @@ def _process_inputs(
 
     def _to_number(
         self, number_like: Any, *, id: tuple[Any, ...]
-    ) -> Union[int, float, complex]:
+    ) -> int | float | complex:
         # pyrefly: ignore [missing-attribute]
         if HAS_NUMPY and isinstance(number_like, np.number):
             return number_like.item()
@@ -673,8 +702,8 @@ def __init__(
         *,
         id: tuple[Any, ...] = (),
         allow_subclasses: bool = True,
-        rtol: Optional[float] = None,
-        atol: Optional[float] = None,
+        rtol: float | None = None,
+        atol: float | None = None,
         equal_nan: bool = False,
         check_device: bool = True,
         check_dtype: bool = True,
@@ -880,7 +909,7 @@ def bitwise_comp(
                 rtol: float,
                 atol: float,
                 equal_nan: bool,
-                identifier: Optional[Union[str, Callable[[str], str]]] = None,
+                identifier: str | Callable[[str], str] | None = None,
             ) -> None:
                 if rtol != 0.0 or atol != 0.0:
                     raise ErrorMeta(
@@ -1063,7 +1092,7 @@ def _compare_regular_values_equal(
         expected: torch.Tensor,
         *,
         equal_nan: bool = False,
-        identifier: Optional[Union[str, Callable[[str], str]]] = None,
+        identifier: str | Callable[[str], str] | None = None,
     ) -> None:
         """Checks if the values of two tensors are equal."""
         self._compare_regular_values_close(
@@ -1078,7 +1107,7 @@ def _compare_regular_values_close(
         rtol: float,
         atol: float,
         equal_nan: bool,
-        identifier: Optional[Union[str, Callable[[str], str]]] = None,
+        identifier: str | Callable[[str], str] | None = None,
     ) -> None:
         """Checks if the values of two tensors are close up to a desired tolerance."""
         matches = torch.isclose(
@@ -1330,14 +1359,14 @@ def assert_close(
     expected: Any,
     *,
     allow_subclasses: bool = True,
-    rtol: Optional[float] = None,
-    atol: Optional[float] = None,
+    rtol: float | None = None,
+    atol: float | None = None,
     equal_nan: bool = False,
     check_device: bool = True,
     check_dtype: bool = True,
     check_layout: bool = True,
     check_stride: bool = False,
-    msg: Optional[Union[str, Callable[[str], str]]] = None,
+    msg: str | Callable[[str], str] | None = None,
 ):
     r"""Asserts that ``actual`` and ``expected`` are close.
 
@@ -1575,6 +1604,8 @@ def assert_close(
     # Hide this function from `pytest`'s traceback
     __tracebackhide__ = True
 
+    actual, expected = _unwrap_dtensor_for_comparison(actual, expected)
+
     error_metas = not_close_error_metas(
         actual,
         expected,
@@ -1609,8 +1640,8 @@ def assert_close(
 def assert_allclose(
     actual: Any,
     expected: Any,
-    rtol: Optional[float] = None,
-    atol: Optional[float] = None,
+    rtol: float | None = None,
+    atol: float | None = None,
     equal_nan: bool = True,
     msg: str = "",
 ) -> None:
diff --git a/torch/testing/_creation.py b/torch/testing/_creation.py
index 23d80d6ceae4f..6b212d4a84b1d 100644
--- a/torch/testing/_creation.py
+++ b/torch/testing/_creation.py
@@ -6,7 +6,7 @@
 import functools
 import math
 import warnings
-from typing import cast, Optional, Union
+from typing import cast
 
 import torch
 
@@ -43,15 +43,15 @@ def _uniform_random_(t: torch.Tensor, low: float, high: float) -> torch.Tensor:
 
 
 def make_tensor(
-    *shape: Union[int, torch.Size, list[int], tuple[int, ...]],
+    *shape: int | torch.Size | list[int] | tuple[int, ...],
     dtype: torch.dtype,
-    device: Union[str, torch.device],
-    low: Optional[float] = None,
-    high: Optional[float] = None,
+    device: str | torch.device,
+    low: float | None = None,
+    high: float | None = None,
     requires_grad: bool = False,
     noncontiguous: bool = False,
     exclude_zero: bool = False,
-    memory_format: Optional[torch.memory_format] = None,
+    memory_format: torch.memory_format | None = None,
 ) -> torch.Tensor:
     r"""Creates a tensor with the given :attr:`shape`, :attr:`device`, and :attr:`dtype`, and filled with
     values uniformly drawn from ``[low, high)``.
@@ -125,8 +125,8 @@ def make_tensor(
     """
 
     def modify_low_high(
-        low: Optional[float],
-        high: Optional[float],
+        low: float | None,
+        high: float | None,
         *,
         lowest_inclusive: float,
         highest_exclusive: float,
diff --git a/torch/testing/_internal/autocast_test_lists.py b/torch/testing/_internal/autocast_test_lists.py
index b3616fede6ce6..2f0f02e507dd7 100644
--- a/torch/testing/_internal/autocast_test_lists.py
+++ b/torch/testing/_internal/autocast_test_lists.py
@@ -146,6 +146,7 @@ def __init__(self, dev):
             ("softmax", pointwise0_fp16 + (0,)),
             ("log_softmax", pointwise0_fp16 + (0,)),
             ("layer_norm", pointwise0_fp16 + ((pointwise0_fp16[0].numel(),),)),
+            ("rms_norm", pointwise0_fp16 + ((pointwise0_fp16[0].numel(),),)),
             ("group_norm", mat0_fp16 + (1,)),
             ("norm", pointwise0_fp16),
             ("norm", pointwise0_fp16, {"dim": 0}),
diff --git a/torch/testing/_internal/common_cuda.py b/torch/testing/_internal/common_cuda.py
index 3f539d586e8bc..2724b2b316eea 100644
--- a/torch/testing/_internal/common_cuda.py
+++ b/torch/testing/_internal/common_cuda.py
@@ -5,7 +5,7 @@
 import functools
 import torch
 import torch.cuda
-from torch.testing._internal.common_utils import LazyVal, TEST_NUMBA, TEST_WITH_ROCM, TEST_CUDA, IS_WINDOWS, IS_MACOS
+from torch.testing._internal.common_utils import LazyVal, TEST_NUMBA, TEST_WITH_ROCM, TEST_CUDA, IS_WINDOWS, IS_MACOS, TEST_XPU
 import inspect
 import contextlib
 import os
@@ -36,13 +36,23 @@
 SM100OrLater = LazyVal(lambda: torch.cuda.is_available() and torch.cuda.get_device_capability() >= (10, 0))
 SM120OrLater = LazyVal(lambda: torch.cuda.is_available() and torch.cuda.get_device_capability() >= (12, 0))
 
-IS_THOR = LazyVal(lambda: torch.cuda.is_available() and
+IS_THOR = LazyVal(lambda: torch.cuda.is_available() and torch.version.cuda is not None and
                   ((torch.cuda.get_device_capability() == (11, 0) and int(torch.version.cuda[:2]) >= 13) or
                    (torch.cuda.get_device_capability() == (10, 1) and int(torch.version.cuda[:2]) < 13)))
 IS_JETSON = LazyVal(lambda: torch.cuda.is_available() and (torch.cuda.get_device_capability() in [(7, 2), (8, 7)] or IS_THOR))
 IS_SM89 = LazyVal(lambda: torch.cuda.is_available() and torch.cuda.get_device_capability() == (8, 9))
 IS_SM90 = LazyVal(lambda: torch.cuda.is_available() and torch.cuda.get_device_capability() == (9, 0))
 IS_SM100 = LazyVal(lambda: torch.cuda.is_available() and torch.cuda.get_device_capability() == (10, 0))
+IS_SM12X = LazyVal(lambda: torch.cuda.is_available() and torch.cuda.get_device_capability()[0] == 12)
+
+@contextlib.contextmanager
+def blas_library_context(backend):
+    prev_backend = torch.backends.cuda.preferred_blas_library()
+    torch.backends.cuda.preferred_blas_library(backend)
+    try:
+        yield
+    finally:
+        torch.backends.cuda.preferred_blas_library(prev_backend)
 
 def evaluate_gfx_arch_within(arch_list):
     if not torch.cuda.is_available():
@@ -53,30 +63,57 @@ def evaluate_gfx_arch_within(arch_list):
     # Hence the matching should be done reversely
     return any(arch in effective_arch for arch in arch_list)
 
+# Per-generation gfx targets, ordered oldest -> newest. Each "OrLater" helper
+# below unions its own generation with every newer one, so the predicates are
+# nested by construction: CDNA5OrLater => CDNA3OrLater => CDNA2OrLater.
+_CDNA2_ARCHS = ["gfx90a"]
+_CDNA3_ARCHS = ["gfx942", "gfx950"]
+# GFX1250 (CDNA 5)
+_CDNA5_ARCHS = ["gfx1250"]
+
+def CDNA5OrLater():
+    return evaluate_gfx_arch_within(_CDNA5_ARCHS)
+
 def CDNA3OrLater():
-    return evaluate_gfx_arch_within(["gfx942", "gfx950"])
+    return evaluate_gfx_arch_within(_CDNA3_ARCHS + _CDNA5_ARCHS)
 
 def CDNA2OrLater():
-    return evaluate_gfx_arch_within(["gfx90a", "gfx942"])
+    return evaluate_gfx_arch_within(_CDNA2_ARCHS + _CDNA3_ARCHS + _CDNA5_ARCHS)
+
+# Archs that take the opportunistic_fastAtomicAdd path (packed 2x16 atomics + DPP
+# lane coalescing) in ScatterGatherKernel.cu. Keep in sync with that kernel's arch
+# gate; this is intentionally not CDNA3OrLater (gfx1250 uses plain fastAtomicAdd).
+def gfx_arch_supports_opportunistic_fastatomics():
+    return evaluate_gfx_arch_within(["gfx942", "gfx950"])
 
 def evaluate_platform_supports_flash_attention():
     if TEST_WITH_ROCM:
-        arch_list = ["gfx90a", "gfx942", "gfx1100", "gfx1201", "gfx950"]
+        arch_list = ["gfx90a", "gfx942", "gfx1100", "gfx1201", "gfx950", "gfx1250"]
         if os.environ.get("TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL", "0") != "0":
             arch_list += ["gfx1101", "gfx1102", "gfx1150", "gfx1151", "gfx1200"]
         return evaluate_gfx_arch_within(arch_list)
     if TEST_CUDA:
         return not IS_WINDOWS and SM80OrLater
+    if TEST_XPU:
+        return True
     return False
 
+def evaluate_platform_supports_ck_sdpa():
+    if TEST_WITH_ROCM:
+        return torch.backends.cuda.is_ck_sdpa_available()
+    else:
+        return False
+
 def evaluate_platform_supports_efficient_attention():
     if TEST_WITH_ROCM:
-        arch_list = ["gfx90a", "gfx942", "gfx1100", "gfx1201", "gfx950"]
+        arch_list = ["gfx90a", "gfx942", "gfx1100", "gfx1201", "gfx950", "gfx1250"]
         if os.environ.get("TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL", "0") != "0":
             arch_list += ["gfx1101", "gfx1102", "gfx1150", "gfx1151", "gfx1200"]
         return evaluate_gfx_arch_within(arch_list)
     if TEST_CUDA:
         return True
+    if TEST_XPU:
+        return True
     return False
 
 def evaluate_platform_supports_cudnn_attention():
@@ -102,12 +139,16 @@ def evaluate_platform_supports_green_context():
 
 PLATFORM_SUPPORTS_FUSED_SDPA: bool = TEST_CUDA and not TEST_WITH_ROCM
 
+PLATFORM_SUPPORTS_CK_SDPA: bool = LazyVal(lambda: evaluate_platform_supports_ck_sdpa())
+
 
 def evaluate_platform_supports_bf16():
     if torch.version.cuda:
         return SM80OrLater
     elif torch.version.hip:
         return True
+    elif TEST_XPU:
+        return True
     return False
 
 
@@ -131,6 +172,18 @@ def evaluate_platform_supports_half_atomics():
 
 PLATFORM_SUPPORTS_GREEN_CONTEXT: bool = LazyVal(lambda: evaluate_platform_supports_green_context())
 
+def evaluate_platform_supports_workqueue_config():
+    if IS_WINDOWS:
+        return False
+    if not _get_torch_cuda_version() >= (13, 1):
+        return False
+    driver_version = torch.utils.collect_env.get_nvidia_driver_version(torch.utils.collect_env.run)
+    if driver_version is None:
+        return False
+    return int(driver_version.split('.')[0]) >= 590
+
+PLATFORM_SUPPORTS_WORKQUEUE_CONFIG: bool = LazyVal(lambda: evaluate_platform_supports_workqueue_config())
+
 def evaluate_platform_supports_fp8():
     if torch.cuda.is_available():
         if torch.version.hip:
@@ -139,18 +192,26 @@ def evaluate_platform_supports_fp8():
                 archs.extend(['gfx120'])
             if ROCM_VERSION >= (6, 5):
                 archs.append('gfx95')
+            if ROCM_VERSION >= (7, 14):
+                archs.append('gfx1250')
             for arch in archs:
                 if arch in torch.cuda.get_device_properties(0).gcnArchName:
                     return True
+            return False
         else:
             return SM90OrLater or torch.cuda.get_device_capability() == (8, 9)
-    return False
+    if torch.xpu.is_available():
+        return True
+    # As CPU supports FP8 and is always available, return True.
+    return True
 
 def evaluate_platform_supports_fp8_grouped_gemm():
     if torch.cuda.is_available():
         if torch.version.hip:
             if "USE_MSLK" not in torch.__config__.show():
                 return False
+            # gfx1250 omitted: MSLK only builds gfx942/gfx950 kernels (see the arch
+            # filter in aten/src/ATen/CMakeLists.txt). Add gfx1250 here once MSLK does.
             archs = ['gfx942', 'gfx950']
             for arch in archs:
                 if arch in torch.cuda.get_device_properties(0).gcnArchName:
@@ -163,7 +224,8 @@ def evaluate_platform_supports_mx_gemm():
     if torch.cuda.is_available():
         if torch.version.hip:
             if ROCM_VERSION >= (7, 0):
-                return 'gfx950' in torch.cuda.get_device_properties(0).gcnArchName
+                gcn_name = torch.cuda.get_device_properties(0).gcnArchName
+                return 'gfx950' in gcn_name or ('gfx1250' in gcn_name and ROCM_VERSION >= (7, 14))
         else:
             return SM100OrLater
     return False
@@ -174,8 +236,21 @@ def evaluate_platform_supports_mxfp8_grouped_gemm():
         return built_with_mslk and IS_SM100
     return False
 
+def evaluate_platform_supports_fp8_sparse():
+    if torch.cuda.is_available():
+        if torch.version.hip:
+            return 'gfx950' in torch.cuda.get_device_properties(0).gcnArchName
+        else:
+            return (
+                (SM90OrLater or torch.cuda.get_device_capability() == (8, 9))
+                and torch.backends.cusparselt.is_available()
+                and torch.backends.cusparselt.version() >= 602
+            )
+    return False
+
 PLATFORM_SUPPORTS_MX_GEMM: bool = LazyVal(lambda: evaluate_platform_supports_mx_gemm())
 PLATFORM_SUPPORTS_FP8: bool = LazyVal(lambda: evaluate_platform_supports_fp8())
+PLATFORM_SUPPORTS_FP8_SPARSE: bool = LazyVal(lambda: evaluate_platform_supports_fp8_sparse())
 PLATFORM_SUPPORTS_FP8_GROUPED_GEMM: bool = LazyVal(lambda: evaluate_platform_supports_fp8_grouped_gemm())
 PLATFORM_SUPPORTS_MXFP8_GROUPED_GEMM: bool = LazyVal(lambda: evaluate_platform_supports_mxfp8_grouped_gemm())
 
@@ -183,7 +258,7 @@ def evaluate_platform_supports_mxfp8_grouped_gemm():
     try:
         import numba.cuda
         TEST_NUMBA_CUDA = numba.cuda.is_available()
-    except (ImportError, RuntimeError):
+    except (ImportError, RuntimeError, OSError):
         TEST_NUMBA_CUDA = False
         TEST_NUMBA = False
 else:
@@ -425,6 +500,9 @@ def xfailIfSM100OrLater(func):
 def xfailIfSM120OrLater(func):
     return func if not SM120OrLater else unittest.expectedFailure(func)
 
+def xfailIfSM12X(func):
+    return func if not IS_SM12X else unittest.expectedFailure(func)
+
 def xfailIfDistributedNotSupported(func):
     return func if not (IS_MACOS or IS_JETSON) else unittest.expectedFailure(func)
 
diff --git a/torch/testing/_internal/common_device_type.py b/torch/testing/_internal/common_device_type.py
index 3fe2620cc8323..8fc9f2cd9682f 100644
--- a/torch/testing/_internal/common_device_type.py
+++ b/torch/testing/_internal/common_device_type.py
@@ -12,11 +12,12 @@
 from collections.abc import Callable, Iterable, Sequence
 from enum import Enum
 from functools import partial, wraps
-from typing import Any, ClassVar, Optional, TypeVar, Union
+from typing import Any, ClassVar, TypeVar
 from typing_extensions import ParamSpec
 
 import torch
 from torch._inductor.utils import GPU_TYPES
+from torch._utils import _is_privateuse1_backend_available
 from torch.testing._internal.common_cuda import (
     _get_torch_cuda_version,
     _get_torch_hipblaslt_version,
@@ -33,7 +34,6 @@
     get_tracked_input,
     IS_FBCODE,
     IS_MACOS,
-    is_privateuse1_backend_available,
     IS_REMOTE_GPU,
     IS_S390X,
     IS_SANDCASTLE,
@@ -318,6 +318,15 @@ def _update_param_kwargs(param_kwargs, name, value):
 class DeviceTypeTestBase(TestCase):
     device_type: str = "generic_device_type"
 
+    # When True, @onlyOn-based decorators (@onlyCUDA, @onlyMPS, etc.) will not
+    # skip tests for this device type. This is a pragmatic short-term solution to
+    # allow PrivateUse1 backends to run tests that are currently gated behind
+    # device-specific decorators. It is intended to be used together with the
+    # skip mechanism (see https://github.com/pytorch/pytorch/issues/177253).
+    # In the longer term, we are incrementally migrating accelerator tests to be
+    # device-generic and removing @onlyCUDA on tests that should be device-generic.
+    bypass_device_restrictions: bool = False
+
     # Flag to disable test suite early due to unrecoverable error such as CUDA error.
     _stop_test_suite = False
 
@@ -670,6 +679,7 @@ class PrivateUse1TestBase(DeviceTypeTestBase):
     primary_device: ClassVar[str]
     device_mod = None
     device_type = "privateuse1"
+    bypass_device_restrictions = False
 
     @classmethod
     def get_primary_device(cls):
@@ -723,7 +733,7 @@ def get_device_type_test_bases():
         if torch.cuda.is_available():
             test_bases.append(CUDATestBase)
 
-        if is_privateuse1_backend_available():
+        if _is_privateuse1_backend_available():
             test_bases.append(PrivateUse1TestBase)
         # Disable MPS testing in generic device testing temporarily while we're
         # ramping up support.
@@ -750,7 +760,7 @@ def filter_desired_device_types(device_type_test_bases, except_for=None, only_fo
     # This handles the case where PrivateUse1TestBase.device_type has been
     # changed from "privateuse1" to the actual backend name (e.g., "openreg")
     # by setUpClass being called during previous instantiate_device_type_tests calls
-    if is_privateuse1_backend_available():
+    if _is_privateuse1_backend_available():
         privateuse1_backend_name = torch._C._get_privateuse1_backend_name()
 
         def func_replace(x: str) -> str:
@@ -1064,8 +1074,8 @@ def __init__(
         self,
         op_list,
         *,
-        dtypes: Union[OpDTypes, Sequence[torch.dtype]] = OpDTypes.supported,
-        allowed_dtypes: Optional[Sequence[torch.dtype]] = None,
+        dtypes: OpDTypes | Sequence[torch.dtype] = OpDTypes.supported,
+        allowed_dtypes: Sequence[torch.dtype] | None = None,
         skip_if_dynamo=True,
     ):
         self.op_list = list(op_list)
@@ -1087,7 +1097,7 @@ def _parametrize_test(self, test, generic_cls, device_cls):
         op = check_exhausted_iterator = object()
         for op in self.op_list:
             # Determine the set of dtypes to use.
-            dtypes: Union[set[torch.dtype], set[None]]
+            dtypes: set[torch.dtype] | set[None]
             if isinstance(self.opinfo_dtypes, Sequence):
                 dtypes = set(self.opinfo_dtypes)
             elif self.opinfo_dtypes == OpDTypes.unsupported_backward:
@@ -1104,7 +1114,7 @@ def _parametrize_test(self, test, generic_cls, device_cls):
                 dtypes = set(op.supported_dtypes(device_cls.device_type))
             elif self.opinfo_dtypes == OpDTypes.any_one:
                 # Tries to pick a dtype that supports both forward or backward
-                supported = op.supported_dtypes(device_cls.device_type)
+                supported = set(op.supported_dtypes(device_cls.device_type))
                 supported_backward = op.supported_backward_dtypes(
                     device_cls.device_type
                 )
@@ -1437,13 +1447,15 @@ def efail_fn(slf, *args, **kwargs):
 
 
 class onlyOn:
-    def __init__(self, device_type: Union[str, list]):
+    def __init__(self, device_type: str | list):
         self.device_type = device_type
 
     def __call__(self, fn):
         @wraps(fn)
         def only_fn(slf, *args, **kwargs):
             if slf.device_type not in self.device_type:
+                if getattr(slf, "bypass_device_restrictions", False):
+                    return fn(slf, *args, **kwargs)
                 reason = f"Only runs on {self.device_type}"
                 if IS_SANDCASTLE or IS_FBCODE:
                     print(
@@ -1828,6 +1840,15 @@ def has_hipsolver():
     return rocm_version >= (5, 3)
 
 
+# Skips a test on CUDA if cuSOLVER is not available,
+# and on ROCm if MAGMA is not available.
+def skipCUDAIfNoCusolverROCMIfNoMagma(fn):
+    if TEST_WITH_ROCM:
+        return skipCUDAIfNoMagma(fn)
+    else:
+        return skipCUDAIfNoCusolver(fn)
+
+
 # Skips a test on CUDA/ROCM if cuSOLVER/hipSOLVER is not available
 def skipCUDAIfNoCusolver(fn):
     return skipCUDAIf(
@@ -1932,7 +1953,7 @@ def skipCUDAIfNotMiopenSuggestNHWC(fn):
 
 
 # Skips a test for specified CUDA versions, given in the form of a list of [major, minor]s.
-def skipCUDAVersionIn(versions: Optional[list[tuple[int, int]]] = None):
+def skipCUDAVersionIn(versions: list[tuple[int, int]] | None = None):
     def dec_fn(fn):
         @wraps(fn)
         def wrap_fn(self, *args, **kwargs):
@@ -1950,7 +1971,7 @@ def wrap_fn(self, *args, **kwargs):
 
 
 # Skips a test for CUDA versions less than specified, given in the form of [major, minor].
-def skipCUDAIfVersionLessThan(versions: Optional[tuple[int, int]] = None):
+def skipCUDAIfVersionLessThan(versions: tuple[int, int] | None = None):
     def dec_fn(fn):
         @wraps(fn)
         def wrap_fn(self, *args, **kwargs):
diff --git a/torch/testing/_internal/common_distributed.py b/torch/testing/_internal/common_distributed.py
index 1705cd1398e56..ff0b5f4ad89f5 100644
--- a/torch/testing/_internal/common_distributed.py
+++ b/torch/testing/_internal/common_distributed.py
@@ -24,7 +24,7 @@
 from enum import Enum
 from functools import partial, reduce, wraps
 from io import StringIO
-from typing import Any, NamedTuple, Optional, Union
+from typing import Any, NamedTuple
 from unittest.mock import patch
 
 import torch
@@ -341,7 +341,7 @@ def with_nccl_blocking_wait(func):
     def wrapper(*args, **kwargs):
         # Save and unset TORCH_NCCL_ASYNC_ERROR_HANDLING
         try:
-            cached_nccl_async_error_handling: Union[str, None] = os.environ[
+            cached_nccl_async_error_handling: str | None = os.environ[
                 "TORCH_NCCL_ASYNC_ERROR_HANDLING"
             ]
             del os.environ["TORCH_NCCL_ASYNC_ERROR_HANDLING"]
@@ -351,7 +351,7 @@ def wrapper(*args, **kwargs):
 
         # Save val of TORCH_NCCL_BLOCKING_WAIT and set it.
         try:
-            cached_nccl_blocking_wait: Union[str, None] = os.environ[
+            cached_nccl_blocking_wait: str | None = os.environ[
                 "TORCH_NCCL_BLOCKING_WAIT"
             ]
         except KeyError:
@@ -609,7 +609,10 @@ def create_tcp_store(
     TIMEOUT_DEFAULT = 500
 else:
     TIMEOUT_DEFAULT = int(os.getenv("DISTRIBUTED_TESTS_DEFAULT_TIMEOUT", "300"))
-TIMEOUT_OVERRIDE = {"test_ddp_uneven_inputs": 400}
+TIMEOUT_OVERRIDE = {
+    "test_ddp_uneven_inputs": 400,
+    "test_DistributedDataParallel": 500,
+}
 
 
 # https://github.com/pytorch/pytorch/issues/75665
@@ -712,10 +715,10 @@ def init_multigpu_helper(world_size: int, backend: str):
     return rank_to_GPU
 
 
-tmp_dir: Optional[tempfile.TemporaryDirectory] = None
+tmp_dir: tempfile.TemporaryDirectory | None = None
 
 
-def initialize_temp_directories(init_method: Optional[str] = None) -> None:
+def initialize_temp_directories(init_method: str | None = None) -> None:
     global tmp_dir
     tmp_dir = tempfile.TemporaryDirectory()
     os.environ["TEMP_DIR"] = tmp_dir.name
@@ -737,6 +740,32 @@ def cleanup_temp_dir() -> None:
         tmp_dir.cleanup()
 
 
+def retrieve_result_from_completion_queue(
+    process: torch.multiprocessing.Process,
+    completion_queue: torch.multiprocessing.Queue,
+    timeout: int | None = None,
+) -> Any:
+    """Get result from the completion_queue associated with process.
+
+    When the process finished without putting a result or the timeout expired an exception instance will be returned"""
+    queue_timeout = 120 if timeout is None else max(10, min(120, timeout // 4))
+    start_time = time.time()
+    # Periodically check the process for liveness
+    while True:
+        try:
+            return completion_queue.get(timeout=queue_timeout)
+        except queue.Empty:
+            # If the process is no longer alive we cannot get a result from the queue unless it is there right now.
+            # This can happen if the timeout occurred just before the process put its result and terminated.
+            # So do a last check for emptiness before considering it as a failure.
+            if not process.is_alive() and completion_queue.empty():
+                return RuntimeError(f"Exited with {process.exitcode}")
+        if timeout is not None:
+            elapsed = time.time() - start_time
+            if elapsed > timeout:
+                return RuntimeError(f"Process timed out out after {elapsed}s")
+
+
 # Most tests operate with this worldsize
 if TEST_WITH_ROCM:
     DEFAULT_WORLD_SIZE = min(4, max(2, torch.cuda.device_count()))
@@ -945,7 +974,7 @@ def run_test(self, test_name: str, parent_pipe) -> None:
                 "Process %s skipping test %s for following reason: %s",
                 self.rank,
                 test_name,
-                str(se),
+                se,
             )
             sys.exit(TEST_SKIPS["generic"].exit_code)
         except Exception:
@@ -1496,7 +1525,7 @@ def _check_return_codes(cls, failed_ranks, timeout, fn):
                     "Thread %s skipping test %s for following reason: %s",
                     rank,
                     fn,
-                    str(exc),
+                    exc,
                 )
                 if skip_code < 0:
                     skip_code = TEST_SKIPS["generic"].exit_code
@@ -1695,7 +1724,7 @@ class MultiProcContinuousTest(TestCase):
     # rank of the current process
     rank: int = -2  # unset state
     # Rendezvous file
-    rdvz_file: Optional[str] = None
+    rdvz_file: str | None = None
     # timeout configured per class
     timeout: timedelta = timedelta(seconds=120)
     # Poison pill for rest of tests if one of them fails
@@ -1704,7 +1733,7 @@ class MultiProcContinuousTest(TestCase):
     _processes_spawned: bool = False
 
     @classmethod
-    def backend_str(cls) -> Optional[str]:
+    def backend_str(cls) -> str | None:
         """
         ProcessGroup backend str.
         To be customized by sub test classes, e.g. "nccl".
@@ -1919,14 +1948,24 @@ def _ensure_processes_spawned(cls):
         This supports instantiate_device_type_tests which calls setUpClass during
         class creation (before any tests run), when spawning would be premature.
         """
-        if cls._processes_spawned:
+        # Check the class's own dict: a subclass must not inherit the flag from
+        # a base test class that already ran and tore down its workers, else it
+        # would dispatch tests to the base class's dead worker processes.
+        if cls.__dict__.get("_processes_spawned", False):
             return
 
-        # Handle both method and string attribute for device_type
+        # Handle method, property, and string attribute for device_type
         # (instantiate_device_type_tests sets device_type as a string attribute,
         # making this compatible as a drop-in replacement for MultiProcessTestCase)
-        device_type_attr = cls.device_type
-        if callable(device_type_attr):
+        device_type_attr = cls.__dict__.get("device_type", cls.device_type)
+        if isinstance(device_type_attr, classmethod):
+            device_type = device_type_attr.__func__(cls)
+        elif isinstance(device_type_attr, property):
+            # Note: fget expects an instance but we pass cls since no instance
+            # exists yet. This works because DTensorTestMixin.device_type only
+            # accesses class-level attributes (world_size, module constants).
+            device_type = device_type_attr.fget(cls)
+        elif callable(device_type_attr):
             device_type = device_type_attr()
         else:
             device_type = device_type_attr
@@ -1960,8 +1999,10 @@ def tearDownClass(cls):
         Class-scope test fixture. Run once for entire test class, after all tests finish.
         Tear down the process group if spawned.
         """
-        # If processes were never spawned (e.g., all tests were skipped), nothing to tear down
-        if not cls._processes_spawned:
+        # If processes were never spawned (e.g., all tests were skipped), nothing to tear down.
+        # Check the class's own dict to avoid tearing down an already-finished
+        # base class's workers (see _ensure_processes_spawned).
+        if not cls.__dict__.get("_processes_spawned", False):
             super().tearDownClass()
             return
 
@@ -2009,21 +2050,29 @@ def _worker_run_main_wait(self, fn):
         def wrapper(self):
             if self.rank == self.MAIN_PROCESS_RANK:
                 logger.debug(f"Waiting for workers to finish {self.id()}")  # noqa: G004
-                # Wait for the workers to finish the test
-                for i, completion_queue in enumerate(self.completion_queues):
-                    rv = completion_queue.get()
+                # Drain all completion queues before raising any exception,
+                # so stale results don't desync subsequent tests.
+                deferred_exception = None
+                for i, (p, completion_queue) in enumerate(
+                    zip(self.processes, self.completion_queues)
+                ):
+                    rv = retrieve_result_from_completion_queue(
+                        p, completion_queue, timeout=get_timeout(self.id())
+                    )
+                    if deferred_exception is not None:
+                        # Already captured an exception; just drain
+                        continue
                     if isinstance(rv, unittest.SkipTest):
-                        raise rv
+                        deferred_exception = rv
+                        continue
                     if isinstance(rv, BaseException):
-                        # Hit an exception, re-raise it in the main process.
                         logger.warning(
                             f"Detected failure from Rank {i} in: {self.id()}, "  # noqa: G004
                             f"skipping rest of tests in Test class: {self.__class__.__name__}"  # noqa: G004
                         )
-                        # Poison rest of tests (because ProcessGroup may be not
-                        # reusable now)
                         self.__class__.poison_pill = True
-                        raise rv
+                        deferred_exception = rv
+                        continue
 
                     # Success
                     if rv != self.id():
@@ -2033,6 +2082,9 @@ def wrapper(self):
                     logger.debug(
                         f"Main proc detected rank {i} finished {self.id()}"  # noqa: G004
                     )
+
+                if deferred_exception is not None:
+                    raise deferred_exception
             else:
                 # Worker just runs the test
                 fn()
diff --git a/torch/testing/_internal/common_dtype.py b/torch/testing/_internal/common_dtype.py
index 20c8d88b46ccf..6e62009b6b601 100644
--- a/torch/testing/_internal/common_dtype.py
+++ b/torch/testing/_internal/common_dtype.py
@@ -222,6 +222,20 @@ def get_all_qint_dtypes() -> list[torch.dtype]:
     return [torch.qint8, torch.quint8, torch.qint32, torch.quint4x2, torch.quint2x4]
 
 
+def highest_precision_float(device):
+    if torch.device(device).type == "mps":
+        return torch.float32
+    else:
+        return torch.float64
+
+
+def highest_precision_complex(device):
+    if torch.device(device).type == "mps":
+        return torch.complex64
+    else:
+        return torch.complex128
+
+
 float_to_corresponding_complex_type_map = {
     torch.float16: torch.complex32,
     torch.float32: torch.complex64,
diff --git a/torch/testing/_internal/common_fsdp.py b/torch/testing/_internal/common_fsdp.py
index 0cc39e4e3f5dc..cebeb1a4c0197 100644
--- a/torch/testing/_internal/common_fsdp.py
+++ b/torch/testing/_internal/common_fsdp.py
@@ -14,7 +14,7 @@
 from copy import deepcopy
 from enum import auto, Enum
 from functools import wraps
-from typing import Any, cast, no_type_check, Optional, Union
+from typing import Any, cast, no_type_check
 from unittest import mock
 
 import torch
@@ -337,10 +337,10 @@ def init(
         group: dist.ProcessGroup,
         fsdp_init_mode: FSDPInitMode,
         device_init_mode: DEVICEInitMode,
-        fsdp_kwargs: Optional[dict[str, Any]] = None,
+        fsdp_kwargs: dict[str, Any] | None = None,
         deterministic: bool = False,
         add_bn: bool = True,
-    ) -> Union[nn.Module, FSDP]:
+    ) -> nn.Module | FSDP:
         """
         Initializes a :class:`TransformerWithSharedParams` instance.
 
@@ -466,7 +466,7 @@ def init(
         group: dist.ProcessGroup,
         fsdp_init_mode: FSDPInitMode,
         device_init_mode: DEVICEInitMode,
-        fsdp_kwargs: Optional[dict[str, Any]] = None,
+        fsdp_kwargs: dict[str, Any] | None = None,
         deterministic: bool = False,
     ) -> nn.Module:
         """
@@ -514,7 +514,7 @@ def init(
         group: dist.ProcessGroup,
         fsdp_init_mode: FSDPInitMode,
         device_init_mode: DEVICEInitMode,
-        fsdp_kwargs: Optional[dict[str, Any]] = None,
+        fsdp_kwargs: dict[str, Any] | None = None,
         deterministic: bool = False,
     ):
         """
@@ -596,7 +596,7 @@ def init(
         group: dist.ProcessGroup,
         fsdp_init_mode: FSDPInitMode,
         device_init_mode: DEVICEInitMode,
-        fsdp_kwargs: Optional[dict[str, Any]] = None,
+        fsdp_kwargs: dict[str, Any] | None = None,
         deterministic: bool = False,
     ):
         """
@@ -721,7 +721,7 @@ def init(  # type: ignore[override]
         group: dist.ProcessGroup,
         fsdp_init_mode: FSDPInitMode,
         device_init_mode: DEVICEInitMode = DEVICEInitMode.DEVICE_AFTER,
-        fsdp_kwargs: Optional[dict[str, Any]] = None,
+        fsdp_kwargs: dict[str, Any] | None = None,
         deterministic: bool = False,
         delay_after_loss_ms: int = 0,
         delay_before_reduction_ms: int = 0,
@@ -841,7 +841,7 @@ def init(
         group: dist.ProcessGroup,
         fsdp_init_mode: FSDPInitMode,
         device_init_mode: DEVICEInitMode,
-        fsdp_kwargs: Optional[dict[str, Any]] = None,
+        fsdp_kwargs: dict[str, Any] | None = None,
         deterministic: bool = False,
         delay_before_free_ms: int = 0,
     ):
@@ -892,7 +892,7 @@ class MLP(nn.Module):
     def __init__(
         self,
         dim: int,
-        device: Optional[torch.device] = None,
+        device: torch.device | None = None,
         *,
         bias: bool = True,
         with_buffer: bool = False,
@@ -980,7 +980,7 @@ def __init__(self, dim: int, use_second_linear: bool = True):
 
     def forward(
         self, x: torch.Tensor
-    ) -> Union[tuple[torch.Tensor, torch.Tensor], torch.Tensor]:
+    ) -> tuple[torch.Tensor, torch.Tensor] | torch.Tensor:
         if self.use_second_linear:
             return self.relu(self.lin1(x)), self.relu(self.lin2(x))
         return self.relu(self.lin1(x))
@@ -1271,12 +1271,12 @@ def _train_for_several_steps(
         num_steps: int,
         autocast: bool,
         lr: float = 0.01,
-        fsdp_cpu_offload: Optional[CPUOffload] = None,
+        fsdp_cpu_offload: CPUOffload | None = None,
         save_model: bool = False,
-        mixed_precision: Optional[MixedPrecision] = None,
+        mixed_precision: MixedPrecision | None = None,
         enable_sharded_grad_scaler: bool = False,
         use_pure_fp16: bool = False,
-        sharded_grad_scaler_kwargs: Optional[dict[str, Any]] = None,
+        sharded_grad_scaler_kwargs: dict[str, Any] | None = None,
     ):
         cpu_offload_params = fsdp_cpu_offload and fsdp_cpu_offload.offload_params
 
@@ -1361,19 +1361,19 @@ def _test_fsdp_parity(
         model_class: type[FSDPTestModel],
         fsdp_init_mode: FSDPInitMode,
         device_init_mode: DEVICEInitMode,
-        ref_init_fn: Optional[Callable] = None,
+        ref_init_fn: Callable | None = None,
         num_iters: int = 2,
         save_model: bool = True,
         cpu_offload: CPUOffload = CPUOffload(),
-        backward_prefetch: Optional[BackwardPrefetch] = None,
-        sharding_strategy: Optional[ShardingStrategy] = None,
-        mixed_precision: Optional[MixedPrecision] = None,
+        backward_prefetch: BackwardPrefetch | None = None,
+        sharding_strategy: ShardingStrategy | None = None,
+        mixed_precision: MixedPrecision | None = None,
         forward_prefetch: bool = False,
         use_orig_params: bool = False,
         enable_sharded_grad_scaler: bool = False,
         use_pure_fp16: bool = False,
-        init_kwargs: Optional[dict[str, Any]] = None,
-        sharded_grad_scaler_kwargs: Optional[dict[str, Any]] = None,
+        init_kwargs: dict[str, Any] | None = None,
+        sharded_grad_scaler_kwargs: dict[str, Any] | None = None,
         **fsdp_kwargs,
     ):
         """
@@ -1664,7 +1664,7 @@ def process_group(self):
         return self.__class__.pg
 
 
-def compiled_fsdp_test(compile_compute_on_module: Optional[type] = None):
+def compiled_fsdp_test(compile_compute_on_module: type | None = None):
     def fully_shard_with_compiled_compute(*args, **kwargs):
         torch.distributed.fsdp.fully_shard(*args, **kwargs)  # type: ignore[operator]
         if compile_compute_on_module is None or isinstance(
@@ -1687,14 +1687,12 @@ def wrapper(*args, **kwargs):
                     )
                     continue
                 # barrier to ensure thread reading the same value
-                original_skip_fsdp_hooks = torch._dynamo.config.skip_fsdp_hooks
                 original_compile_threads = torch._inductor.config.compile_threads
                 torch.distributed.barrier()
 
                 if mode == FullyShardMode.EAGER:
                     fully_shard_patch = original_fully_shard
                 elif mode == FullyShardMode.COMPILED_COMPUTE:
-                    torch._dynamo.config.skip_fsdp_hooks = True
                     torch._inductor.config.compile_threads = 1
                     fully_shard_patch = fully_shard_with_compiled_compute  # type: ignore[assignment]
                 else:
@@ -1709,7 +1707,6 @@ def wrapper(*args, **kwargs):
                 # other threads use patched func before this thread restores
                 torch.distributed.barrier()
                 func.__globals__[original_fully_shard.__name__] = original_fully_shard
-                torch._dynamo.config.skip_fsdp_hooks = original_skip_fsdp_hooks
                 torch._inductor.config.compile_threads = original_compile_threads
 
         return wrapper
diff --git a/torch/testing/_internal/common_jit.py b/torch/testing/_internal/common_jit.py
index ac6e851d7e28b..5381752d64142 100644
--- a/torch/testing/_internal/common_jit.py
+++ b/torch/testing/_internal/common_jit.py
@@ -16,7 +16,6 @@
 
 # Standard library
 from itertools import chain
-from typing import Union
 from torch._C import TensorType
 
 import io
@@ -62,7 +61,7 @@ def clone_tensor(t, preserve_requires_grad):
         return t.detach().clone().requires_grad_(require_grad)
 
     def clone_inputs(preserve_requires_grad: bool):
-        inputs: list[Union[torch.Tensor, list[torch.Tensor]]] = []
+        inputs: list[torch.Tensor | list[torch.Tensor]] = []
 
         for arg in args:
             if isinstance(arg, torch.Tensor):
@@ -284,7 +283,7 @@ def assertAutodiffNode(self, graph, should_autodiff_node, nonfusible_nodes, fusi
             self.assertEqual(should_autodiff_node,
                              found_all_nonfusible_nodes and found_all_fusible_nodes, err_msg)
 
-    def checkShapeAnalysis(self, out_sizes: Union[list[int], list[list[int]]],
+    def checkShapeAnalysis(self, out_sizes: list[int] | list[list[int]],
                            traced_graph, assert_propagation, constant_prop=True):
         # repropagte input shapes provided by tracing,
         prev_symbolic_shapes_test_enabled = torch._C._jit_symbolic_shapes_test_mode_enabled()
diff --git a/torch/testing/_internal/common_methods_invocations.py b/torch/testing/_internal/common_methods_invocations.py
index 7a9e9df8a519d..9b1665c67d0b3 100644
--- a/torch/testing/_internal/common_methods_invocations.py
+++ b/torch/testing/_internal/common_methods_invocations.py
@@ -16,13 +16,15 @@
 import numpy.typing as npt
 from torch import inf, nan
 
-from typing import Any, Union
+from typing import Any
 from collections.abc import Sequence
 from torch.testing import make_tensor
 from torch.testing._internal.common_dtype import (
     _dispatch_dtypes, floating_types, floating_types_and, complex_types, floating_and_complex_types,
     floating_and_complex_types_and, all_types_and_complex_and, all_types_and, all_types_and_complex, integral_types_and,
     empty_types, complex_types_and, integral_types, custom_types, all_types_complex_float8_and, float8_types,
+    highest_precision_complex,
+    highest_precision_float,
 )
 from torch.testing._internal.common_device_type import (
     onlyCPU, onlyCUDA, onlyNativeDeviceTypes, disablecuDNN, skipCUDAIfNoMagma, skipCUDAIfNoMagmaAndNoCusolver,
@@ -39,7 +41,9 @@
 )
 from torch.testing._internal.common_utils import (
     make_fullrank_matrices_with_distinct_singular_values,
-    TEST_WITH_ROCM, IS_FBCODE, IS_WINDOWS, IS_MACOS, IS_S390X, TEST_SCIPY,
+    IS_ARM64,
+    IS_CPU_EXT_SVE_SUPPORTED,
+    TEST_WITH_ROCM, IS_FBCODE, IS_WINDOWS, IS_MACOS, TEST_SCIPY,
     torch_to_numpy_dtype_dict, numpy_to_torch_dtype, TEST_WITH_ASAN,
     GRADCHECK_NONDET_TOL, slowTest, TEST_WITH_SLOW,
     TEST_WITH_TORCHINDUCTOR,
@@ -1555,7 +1559,8 @@ def error_inputs_multi_margin_loss(op, device, **kwargs):
                      error_regex=r'Expected non-empty vector or matrix with optional 0-dim batch size, but got: \[0\]')
     # invalid target
     yield ErrorInput(SampleInput(make_input(5, 4), args=(make_input(5, 4),), kwargs={}),
-                     error_type=RuntimeError, error_regex=r'inconsistent target size, expected 5 but got \[5, 4\]')
+                     error_type=RuntimeError,
+                     error_regex=r'target tensor should be 1-D with size equal to.*Expected target size \[5\].*but got \[5, 4\]')
     # invalid target dtype
     yield ErrorInput(SampleInput(make_input(5, 4), args=(make_input(5,),), kwargs={}),
                      error_type=RuntimeError, error_regex='expected scalar type Long but found Float')
@@ -1612,7 +1617,7 @@ def sample_inputs_like_fns(self, device, dtype, requires_grad, **kwargs):
         ((S,), {'dtype': dtype, 'device': device}),
         # Hard-code some dtypes/devices. We want to test cases where the
         # (dtype, device) is different from the input's (dtype, device)
-        ((S,), {'dtype': torch.double if device != 'mps:0' else torch.float}),
+        ((S,), {'dtype': highest_precision_float(device)}),
         ((S,), {'device': 'cpu'}),
         ((S,), {'dtype': torch.double, 'device': 'cpu'}),
     ]
@@ -1806,7 +1811,7 @@ def sample_inputs_new_fns(self, device, dtype, requires_grad, *, is_strided=Fals
         ((S,), (2, 3), (7, 8), {'dtype': dtype, 'device': device}),
         # Hard-code some dtypes/devices. We want to test cases where the
         # (dtype, device) is different from the input's (dtype, device)
-        ((S,), (10,), (S,), {'dtype': torch.double if device != 'mps:0' else torch.float}),
+        ((S,), (10,), (S,), {'dtype': highest_precision_float(device)}),
         ((S,), (1, 1, 12), (S, L, M), {'device': 'cpu'}),
         ((S,), (2, 2, 2), (L, M, S), {'dtype': torch.double, 'device': 'cpu'}),
     ]
@@ -1929,7 +1934,7 @@ def sample_inputs_full_like(self, device, dtype, requires_grad, **kwargs):
     def get_val(dtype):
         return make_tensor([], dtype=dtype, device="cpu").item()
 
-    double_dtype = torch.double if device != "mps:0" else torch.float
+    double_dtype = highest_precision_float(device)
     inputs = [
         ((), get_val(dtype), {}),
         ((S, S), get_val(dtype), {}),
@@ -2043,10 +2048,11 @@ def sample_inputs_logcumsumexp(self, device, dtype, requires_grad, **kwargs):
             yield SampleInput(t, dim)
 
 def sample_inputs_trace(self, device, dtype, requires_grad, **kwargs):
-    yield SampleInput(
-        make_tensor((S, S), dtype=dtype, device=device,
-                    low=None, high=None,
-                    requires_grad=requires_grad))
+    make_arg = partial(make_tensor, device=device, dtype=dtype,
+                       requires_grad=requires_grad, low=None, high=None)
+    # Square, tall (rows > cols), wide (rows < cols), single row/col (#171704)
+    for shape in ((S, S), (S + 2, S), (S, S + 2), (1, S), (S, 1)):
+        yield SampleInput(make_arg(shape))
 
 
 def error_inputs_trace(op, device):
@@ -2527,7 +2533,7 @@ def reference_inputs_cat(op, device, dtype, requires_grad, **kwargs):
 
     # Noncontiguous type promoting tensors
     a = make_arg((3, 4, 2))
-    b = make_arg((3, 2, 2), noncontiguous=True, dtype=torch.double)
+    b = make_arg((3, 2, 2), noncontiguous=True, dtype=highest_precision_float(device))
     c = make_arg((3, 3, 2), dtype=torch.float16).permute(1, 0, 2)
 
     yield SampleInput((a, b, c), kwargs={'dim': 1})
@@ -2686,7 +2692,7 @@ def error_inputs_gather(op_info, device, **kwargs):
     # Creates new src & idx since SampleInputs can't share tensors
     src = torch.tensor(((1, 2), (3, 4)), device=device, dtype=torch.float32)
     idx = torch.tensor(((0, 0), (1, 0)), device=device, dtype=torch.long)
-    out = torch.empty((2, 2), device=device, dtype=torch.float64)
+    out = torch.empty((2, 2), device=device, dtype=torch.float16 if torch.device(device).type == 'mps' else torch.float64)
     yield ErrorInput(SampleInput(src, args=(1, idx), kwargs={'out': out}),
                      error_regex="Expected out tensor to have dtype")
 
@@ -2749,7 +2755,7 @@ def error_inputs_scatter_and_scatter_add(op_info, device, **kwargs):
     # Error when self.dtype != src.dtype (and src is not a scalar)
     src = make_tensor((2, 5), device=device, dtype=torch.float32)
     idx = torch.tensor(((0, 1), (1, 2)), device=device, dtype=torch.long)
-    dst = torch.zeros((3, 5), device=device, dtype=torch.double)
+    dst = torch.zeros((3, 5), device=device, dtype=torch.float16 if torch.device(device).type == 'mps' else torch.double)
     yield ErrorInput(SampleInput(dst, args=(0, idx, src)),
                      error_regex="Expected self.dtype to be equal to src.dtype")
 
@@ -2825,7 +2831,8 @@ def error_inputs_t(op_info, device, **kwargs):
 
 
 def error_inputs_multinomial(op_info, device, **kwargs):
-    x = torch.empty(1, 2, 3, dtype=torch.double, device=device)
+    dtype = highest_precision_float(device)
+    x = torch.empty(1, 2, 3, dtype=dtype, device=device)
     yield ErrorInput(SampleInput(x, args=(2,)),
                      error_regex="prob_dist must be 1 or 2 dim")
 
@@ -2833,24 +2840,24 @@ def error_inputs_multinomial(op_info, device, **kwargs):
     yield ErrorInput(SampleInput(x, args=(2,)),
                      error_regex="multinomial only supports floating-point dtypes for input")
 
-    x = torch.empty(1, 2, dtype=torch.double, device=device)
-    y = torch.empty(1, 2, dtype=torch.double, device=device)
+    x = torch.empty(1, 2, dtype=dtype, device=device)
+    y = torch.empty(1, 2, dtype=dtype, device=device)
     yield ErrorInput(SampleInput(x, args=(2,), kwargs=dict(out=y)),
                      error_regex="multinomial expects Long tensor out")
 
-    x = torch.empty(2, dtype=torch.double, device=device)
+    x = torch.empty(2, dtype=dtype, device=device)
     yield ErrorInput(SampleInput(x, args=(0,)),
                      error_regex="cannot sample n_sample <= 0 samples")
 
-    x = torch.empty(2, dtype=torch.double, device=device)
+    x = torch.empty(2, dtype=dtype, device=device)
     yield ErrorInput(SampleInput(x, args=(-1,)),
                      error_regex="cannot sample n_sample <= 0 samples")
 
-    x = torch.empty(2, dtype=torch.double, device=device)
+    x = torch.empty(2, dtype=dtype, device=device)
     yield ErrorInput(SampleInput(x, args=(3, False,)),
                      error_regex="cannot sample n_sample > prob_dist")
 
-    x = torch.empty(16777217, dtype=torch.double, device=device)
+    x = torch.empty(16777217, dtype=dtype, device=device)
     yield ErrorInput(SampleInput(x, args=(3,)),
                      error_regex="number of categories cannot exceed")
 
@@ -4052,8 +4059,8 @@ def sample_inputs_conv1d(op_info, device, dtype, requires_grad, **kwargs):
 
 
 def error_inputs_conv1d(opinfo, device, **kwargs):
-    dtype = torch.float64 if device != 'mps:0' else torch.float32
-    cdtype = torch.complex128 if device != 'mps:0' else torch.complex64
+    dtype = highest_precision_float(device)
+    cdtype = highest_precision_complex(device)
     make_arg = partial(make_tensor, device=device, dtype=dtype)
     make_int_arg = partial(make_tensor, device=device, dtype=torch.int64)
     make_complex_arg = partial(make_tensor, device=device, dtype=cdtype)
@@ -4114,8 +4121,8 @@ def error_inputs_conv1d(opinfo, device, **kwargs):
 
 
 def error_inputs_conv2d(opinfo, device, **kwargs):
-    dtype = torch.float64 if device != 'mps:0' else torch.float32
-    cdtype = torch.complex128 if device != 'mps:0' else torch.complex64
+    dtype = highest_precision_float(device)
+    cdtype = highest_precision_complex(device)
     make_arg = partial(make_tensor, device=device, dtype=dtype)
     make_int_arg = partial(make_tensor, device=device, dtype=torch.int64)
     make_complex_arg = partial(make_tensor, device=device, dtype=cdtype)
@@ -4253,8 +4260,8 @@ def sample_inputs_conv3d(opinfo, device, dtype, requires_grad, **kwargs):
 
 
 def error_inputs_conv3d(opinfo, device, **kwargs):
-    dtype = torch.float64 if device != 'mps:0' else torch.float32
-    cdtype = torch.complex128 if device != 'mps:0' else torch.complex64
+    dtype = highest_precision_float(device)
+    cdtype = highest_precision_complex(device)
     make_arg = partial(make_tensor, device=device, dtype=dtype)
     make_int_arg = partial(make_tensor, device=device, dtype=torch.int64)
     make_complex_arg = partial(make_tensor, device=device, dtype=cdtype)
@@ -5136,7 +5143,7 @@ def sample_inputs_avgpool1d(op_info, device, dtype, requires_grad, **kwargs):
     make_arg = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad)
 
     # Order: input_shape, kernel_size, kwargs
-    cases: list[tuple[tuple[int, ...], Union[int, tuple[int, ...]], dict]] = [
+    cases: list[tuple[tuple[int, ...], int | tuple[int, ...], dict]] = [
         ((2, 3, 9), (3,), {}),
         ((1, 3, 9), 3, dict(stride=1, padding=1, ceil_mode=True, count_include_pad=False)),
         ((1, 3, 9), (6,), dict(stride=(3,), padding=(2,), ceil_mode=True, count_include_pad=True)),
@@ -5155,7 +5162,7 @@ def sample_inputs_avgpool3d(op_info, device, dtype, requires_grad, **kwargs):
     make_arg = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad)
 
     # Order: input_shape, kernel_size, stride, padding, ceil_mode, count_include_pad, divisor_override
-    cases: list[tuple[tuple[int, ...], Union[int, tuple[int, ...]], dict]] = [
+    cases: list[tuple[tuple[int, ...], int | tuple[int, ...], dict]] = [
         ((2, 3, 3, 4, 4), (2, 2, 2), {}),
         ((1, 2, 4, 4, 4), 2, dict(stride=1, padding=1, ceil_mode=True,
                                   count_include_pad=False, divisor_override=2)),
@@ -6927,14 +6934,12 @@ def sample_inputs_meshgrid(op_info: OpInfo, device: torch.device, dtype: torch.d
                            *, variant: str, **kwargs) -> list[SampleInput]:
     if variant == 'variadic':
         def make_inputs(
-                tensors: list[torch.Tensor]) -> tuple[Union[torch.Tensor,
-                                                            list[torch.Tensor]],
+                tensors: list[torch.Tensor]) -> tuple[torch.Tensor | list[torch.Tensor],
                                                       tuple[torch.Tensor, ...]]:
             return tensors
     elif variant == 'list':
         def make_inputs(
-                tensors: list[torch.Tensor]) -> tuple[Union[torch.Tensor,
-                                                            list[torch.Tensor]],
+                tensors: list[torch.Tensor]) -> tuple[torch.Tensor | list[torch.Tensor],
                                                       tuple[torch.Tensor, ...]]:
             return [tensors]
     else:
@@ -6996,8 +7001,6 @@ def skips_mvlgamma(skip_redundant=False):
         DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs',
                      'test_reference_numerics_small',
                      dtypes=(torch.int8,)),
-        # NotImplementedError: The operator 'aten::mvlgamma.out' is not currently implemented for the MPS device
-        DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_variant_consistency_eager', device_type='mps'),
     )
     if skip_redundant:
         # Redundant tests
@@ -8916,6 +8919,32 @@ def make_mat(size: tuple[int], scale: float, fp8_dtype: torch.dtype) -> torch.Te
     scale_tensor2 = make_scale(scale2, torch.float8_e4m3fn)
     samples.append(SampleInput(mat1, mat2, scale_tensor1, scale_tensor2))
 
+    # Case 4: MXFP4 (float4_e2m1fn_x2) with E8M0 blockwise scaling
+    # Regression test: E8M0 blockwise scale size validation must account for
+    # packed FP4 format where self.size(1) = K/2.
+    # Only supported on MI350 (gfx950).
+    if device == 'cuda' and torch.version.hip:
+        if 'gfx950' in torch.cuda.get_device_properties(0).gcnArchName:
+            mxfp4_M, mxfp4_K, mxfp4_N = 256, 256, 256
+            block_size_k = 32
+            block_size_mn = 128
+            num_k_blocks = math.ceil(mxfp4_K / block_size_k)
+            padded_num_k_blocks = math.ceil(num_k_blocks / 4) * 4
+            scale_a_size = block_size_mn * math.ceil(mxfp4_M / block_size_mn) * padded_num_k_blocks
+            scale_b_size = block_size_mn * math.ceil(mxfp4_N / block_size_mn) * padded_num_k_blocks
+            mat1 = _bfloat16_to_float4_e2m1fn_x2(
+                torch.randn(mxfp4_M, mxfp4_K, device=device, dtype=torch.bfloat16)
+            )
+            mat2 = _bfloat16_to_float4_e2m1fn_x2(
+                torch.randn(mxfp4_N, mxfp4_K, device=device, dtype=torch.bfloat16)
+            ).t()
+            scale_tensor1 = torch.ones(scale_a_size, device=device, dtype=torch.float8_e8m0fnu)
+            scale_tensor2 = torch.ones(scale_b_size, device=device, dtype=torch.float8_e8m0fnu)
+            samples.append(SampleInput(
+                mat1, mat2, scale_tensor1, scale_tensor2,
+                out_dtype=torch.bfloat16,
+            ))
+
     yield from samples
 
 def sample_inputs_scaled_mm_v2(op_info, device, dtype, requires_grad, **kwargs):
@@ -9084,10 +9113,10 @@ def sample_inputs_scaled_mm_v2(op_info, device, dtype, requires_grad, **kwargs):
                     mat2_fp4,
                     [scale1, global_scale1],
                     [ScalingType.BlockWise1x16, ScalingType.TensorWise],
-                    [SwizzleType.SWIZZLE_32_4_4, ],
+                    [SwizzleType.SWIZZLE_32_4_4, SwizzleType.NO_SWIZZLE],
                     [scale2, global_scale2],
                     [ScalingType.BlockWise1x16, ScalingType.TensorWise],
-                    [SwizzleType.SWIZZLE_32_4_4, ],
+                    [SwizzleType.SWIZZLE_32_4_4, SwizzleType.NO_SWIZZLE],
                     None,  # bias
                     torch.bfloat16,  # out_dtype
                 )
@@ -9384,8 +9413,9 @@ def sample_inputs_l1_loss(op_info, device, dtype, requires_grad, **kwargs):
     # test COMPLEX_TO_FLOAT promotion
     if dtype.is_complex:
         make = partial(make_tensor, (), device=device, requires_grad=requires_grad)
-        yield SampleInput(make(dtype=dtype), args=(make(dtype=torch.double),))
-        yield SampleInput(make(dtype=torch.double), args=(make(dtype=dtype),))
+        other_dtype = highest_precision_float(device)
+        yield SampleInput(make(dtype=dtype), args=(make(dtype=other_dtype),))
+        yield SampleInput(make(dtype=other_dtype), args=(make(dtype=dtype),))
 
 def error_inputs_l1_loss(op_info, device, **kwargs):
     make = partial(make_tensor, device=device, dtype=torch.float32)
@@ -9749,6 +9779,8 @@ def _should_disable_fastpath(self, opinfo, rightmost_arg, rightmost_arg_type, dt
         if self.arity == 1:
             if "foreach_abs" in opinfo.name and dtype in complex_types():
                 return True
+            if "foreach_clone" in opinfo.name:
+                return False
             # unary
             if opinfo.ref in (torch.abs, torch.neg):
                 return False
@@ -9927,7 +9959,7 @@ def sample_zero_size_tensor_inputs(self, opinfo, device, dtype, requires_grad, *
         for ord in (0, 1, 2, -1, -2, float('inf'), float('-inf')):
             input = sample_inputs_foreach(None, device, dtype, NUM_SIZE0_TENSORS, zero_size=True, **_foreach_inputs_kwargs)
             disable_fastpath = True
-            if ord in (1, 2, float('inf')) and dtype in floating_types_and(torch.half, torch.bfloat16):
+            if ord in (0, 1, 2, float('inf')) and dtype in floating_types_and(torch.half, torch.bfloat16):
                 disable_fastpath = False
             yield ForeachSampleInput(input, ord=ord, disable_fastpath=disable_fastpath)
 
@@ -9942,7 +9974,7 @@ def __call__(self, opinfo, device, dtype, requires_grad, **kwargs):
         for num_tensors, ord, out_dtype, intersperse_empty_tensors in product(
             num_input_tensors,
             (0, 1, 2, -1, -2, float('inf'), float('-inf')),
-            (None,) + (torch.complex128,) if dtype in complex_types() else (torch.float64,),
+            (None,) + (highest_precision_complex(device),) if dtype in complex_types() else (highest_precision_float(device),),
             (True, False),
         ):
             # inf norm and negative norms on empty tensors is not supported by our reference func vector norm:
@@ -9953,7 +9985,7 @@ def __call__(self, opinfo, device, dtype, requires_grad, **kwargs):
             _foreach_inputs_kwargs["intersperse_empty_tensors"] = intersperse_empty_tensors
             input = sample_inputs_foreach(None, device, dtype, num_tensors, zero_size=False, **_foreach_inputs_kwargs)
             disable_fastpath = True
-            if ord in (1, 2, float('inf')) and dtype in floating_types_and(torch.half, torch.bfloat16):
+            if ord in (0, 1, 2, float('inf')) and dtype in floating_types_and(torch.half, torch.bfloat16):
                 disable_fastpath = False
             yield ForeachSampleInput(input, ord=ord, disable_fastpath=disable_fastpath, dtype=out_dtype)
 
@@ -9970,7 +10002,7 @@ def __call__(self, opinfo, device, dtype, requires_grad, **kwargs):
             for input in nan_inputs:
                 x = torch.tensor(input, device=device)
                 disable_fastpath = True
-                if ord in (1, 2, float('inf')) and dtype in floating_types_and(torch.half, torch.bfloat16):
+                if ord in (0, 1, 2, float('inf')) and dtype in floating_types_and(torch.half, torch.bfloat16):
                     disable_fastpath = False
                 yield ForeachSampleInput([x], ord=ord, disable_fastpath=disable_fastpath)
 
@@ -11216,6 +11248,13 @@ def __call__(self, opinfo, device, dtype, requires_grad, **kwargs):
             ),
         ),
     ),
+    ForeachFuncInfo(
+        "clone",
+        sample_inputs_func=foreach_inputs_sample_func(1, False, False),
+        supports_forward_ad=True,
+        supports_autograd=True,
+        supports_inplace_autograd=True,
+    ),
 ]
 
 foreach_binary_op_db: list[OpInfo] = [
@@ -12057,6 +12096,20 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
     yield SampleInput(make_tensor((S,), dtype=dtype, device=device, requires_grad=requires_grad))
     yield SampleInput(make_tensor((), dtype=dtype, device=device, requires_grad=requires_grad))
 
+def sample_inputs_abs(op_info, device, dtype, requires_grad, op_kwargs=None, **kwargs):
+    yield from sample_inputs_elementwise_unary(op_info, device, dtype, requires_grad, op_kwargs=None, **kwargs)
+    if dtype == torch.cfloat:
+        yield SampleInput(torch.tensor(
+            [
+                1e-30 + 1e-30j,
+                1e30 + 1e30j,
+                1e-30 + 1e30j,
+                1e30 + 1e-30j,
+            ],
+            device=device,
+            dtype=dtype,
+            requires_grad=requires_grad,
+        ))
 
 # Operator database (sorted alphabetically)
 op_db: list[OpInfo] = [
@@ -12066,6 +12119,7 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
                    dtypes=all_types_and_complex_and(torch.half, torch.bfloat16, torch.chalf),
                    dtypesIfCUDA=all_types_and_complex_and(torch.bool, torch.half, torch.bfloat16, torch.chalf),
                    dtypesIfHpu=custom_types(torch.float32, torch.bfloat16),
+                   sample_inputs_func=sample_inputs_abs,
                    skips=(
                        DecorateInfo(unittest.skip("In-place abs not supported for complex tensors"), 'TestBwdGradients',
                                     'test_inplace_grad', dtypes=(torch.cdouble,)),
@@ -12091,6 +12145,28 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
                                     dtypes=(torch.cdouble, torch.cfloat, torch.chalf)),
                        DecorateInfo(unittest.expectedFailure, 'TestMeta', 'test_dispatch_symbolic_meta_inplace_all_strides',
                                     dtypes=(torch.cdouble, torch.cfloat, torch.chalf)),
+                       # ValueError: Expected 2D tensor but got tensor with dimension: 1.
+                       DecorateInfo(
+                           unittest.expectedFailure, 'TestSparseCSR', 'test_sparse_csr_unary_inplace',
+                           device_type='cpu', dtypes=(torch.complex64,)
+                       ),
+                       DecorateInfo(
+                           unittest.expectedFailure, 'TestSparseCSR', 'test_sparse_csr_unary_out',
+                           device_type='cpu', dtypes=(torch.complex64,)
+                       ),
+                       DecorateInfo(
+                           unittest.expectedFailure, 'TestSparseCSR', 'test_sparse_csr_unary_inplace',
+                           device_type='cuda', dtypes=(torch.complex64,)
+                       ),
+                       DecorateInfo(
+                           unittest.expectedFailure, 'TestSparseCSR', 'test_sparse_csr_unary_out',
+                           device_type='cuda', dtypes=(torch.complex64,)
+                       ),
+                       # AssertionError: Tensor-likes are not close!
+                       DecorateInfo(
+                           unittest.expectedFailure, 'TestOutputConsistencyFullGraph', 'test_complex_output_match_opinfo_',
+                           device_type='cpu', dtypes=(torch.complex64,)
+                       ),
                    ),
                    supports_fwgrad_bwgrad=True,
                    assert_autodiffed=True,
@@ -12712,8 +12788,6 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
                    'TestSchemaCheckModeOpInfo',
                    'test_schema_correctness',
                    dtypes=(torch.complex64, torch.complex128)),
-               # NotImplementedError: "dot" not implemented for 'Bool'
-               DecorateInfo(unittest.expectedFailure, 'TestConsistency', device_type='mps', dtypes=(torch.bool,)),
            )),
     OpInfo('vdot',
            dtypes=all_types_and_complex_and(torch.float16, torch.bfloat16),
@@ -12731,11 +12805,6 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
                    'TestSchemaCheckModeOpInfo',
                    'test_schema_correctness',
                    dtypes=(torch.complex64, torch.complex128)),
-               # UnboundLocalError: cannot access local variable 'cpu_out' where it is not associated with a value
-               DecorateInfo(
-                   unittest.expectedFailure, 'TestConsistency', 'test_output_match',
-                   device_type='mps', dtypes=(torch.bool,),
-               ),
            )),
     OpInfo('bmm',
            dtypes=all_types_and_complex_and(torch.float16, torch.bfloat16),
@@ -12987,6 +13056,9 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
                        DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_large',
                                     device_type='cuda', dtypes=[torch.cfloat],
                                     active_if=IS_WINDOWS),
+                       DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_extremal',
+                                    device_type='xpu', dtypes=[torch.cfloat, torch.cdouble],
+                                    active_if=IS_WINDOWS),
                        DecorateInfo(unittest.skip("Skipped! sparse backward not supported"),
                                     'TestSparseUnaryUfuncs', 'test_sparse_fn_grad'),
                    )),
@@ -13144,6 +13216,7 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
            assert_autodiffed=False,
            sample_inputs_func=sample_inputs_cdist,
            skips=(
+               # NotImplementedError: The operator 'aten::_cdist_backward' is not currently implemented for the MPS device
                DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_noncontiguous_samples', device_type='mps'),
                DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_variant_consistency_eager', device_type='mps'),
                DecorateInfo(toleranceOverride(
@@ -13165,8 +13238,6 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
                                     'TestNNCOpInfo',
                                     'test_nnc_correctness',
                                     dtypes=tuple(t for t in integral_types() if t != torch.uint8)),
-                       # NotImplementedError: "ceil_vml_cpu" not implemented for 'Bool'
-                       DecorateInfo(unittest.expectedFailure, 'TestConsistency', device_type='mps', dtypes=(torch.bool,)),
                    ),
                    supports_sparse=True,
                    supports_sparse_csr=True,
@@ -13244,16 +13315,14 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
     OpInfo('cholesky_solve',
            op=torch.cholesky_solve,
            dtypes=floating_and_complex_types(),
+           dtypesIfMPS=custom_types(torch.float32),
+           backward_dtypesIfMPS=custom_types(torch.float32),
            sample_inputs_func=sample_inputs_cholesky_solve,
            check_batched_gradgrad=False,
            supports_forward_ad=True,
            supports_fwgrad_bwgrad=True,
            gradcheck_wrapper=lambda *args, **kwargs: gradcheck_wrapper_triangular_input(*args, idx=1, **kwargs),
-           decorators=[skipCUDAIfNoMagma, skipCPUIfNoLapack],
-           skips=(
-               # Error: The operator 'aten::_cholesky_solve_helper' is not currently implemented for the MPS device
-               DecorateInfo(unittest.expectedFailure, 'TestCommon', device_type='mps'),
-           )),
+           decorators=[skipCUDAIfNoMagma, skipCPUIfNoLapack]),
     OpInfo('chunk',
            dtypes=all_types_and_complex_and(torch.bool, torch.bfloat16, torch.float16, torch.chalf),
            dtypesIfHpu=custom_types(torch.float32, torch.bfloat16, torch.int32),
@@ -13491,6 +13560,9 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
                        DecorateInfo(unittest.expectedFailure, 'TestUnaryUfuncs', 'test_reference_numerics_large',
                                     device_type='cuda',
                                     dtypes=(torch.chalf,), active_if=IS_WINDOWS),
+                       DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_large',
+                                    device_type='xpu',
+                                    dtypes=(torch.chalf,), active_if=IS_WINDOWS),
                    )),
     UnaryUfuncInfo('cosh',
                    ref=np_unary_ufunc_integer_promotion_wrapper(np.cosh),
@@ -13523,6 +13595,9 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
                        DecorateInfo(unittest.expectedFailure, 'TestUnaryUfuncs', 'test_reference_numerics_large',
                                     device_type='cuda',
                                     dtypes=(torch.chalf,), active_if=IS_WINDOWS),
+                       DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_large',
+                                    device_type='xpu',
+                                    dtypes=(torch.chalf,), active_if=IS_WINDOWS),
                    )),
     OpInfo('cov',
            dtypes=all_types_and_complex_and(torch.half, torch.bfloat16),
@@ -13564,23 +13639,7 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
            supports_fwgrad_bwgrad=True,
            supports_out=True,
            supports_forward_ad=True,
-           skips=(
-               # The following dtypes did not work in forward but are listed by the OpInfo: {torch.complex64}
-               DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_dtypes', device_type='mps'),
-               # RuntimeError: Failed to create function state object for: cross_float2
-               DecorateInfo(
-                   unittest.expectedFailure, 'TestCommon', 'test_noncontiguous_samples',
-                   device_type='mps', dtypes=(torch.complex64,)
-               ),
-               DecorateInfo(
-                   unittest.expectedFailure, 'TestCommon', 'test_out_requires_grad_error',
-                   device_type='mps', dtypes=(torch.complex64,)
-               ),
-               DecorateInfo(
-                   unittest.expectedFailure, 'TestCommon', 'test_variant_consistency_eager',
-                   device_type='mps', dtypes=(torch.complex64,)
-               ),
-           ),),
+           ),
     OpInfo('cumsum',
            dtypes=all_types_and_complex_and(torch.half, torch.bfloat16),
            dtypesIfHpu=custom_types(torch.float32, torch.bfloat16, torch.int32),
@@ -13589,21 +13648,8 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
            skips=(
                # cumsum does not handle correctly out= dtypes
                DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out'),
-               # The following dtypes did not work in forward but are listed by the OpInfo: {torch.complex64}
+               # The following dtypes did not work in forward but are listed by the OpInfo: {torch.bool}
                DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_dtypes', device_type='mps'),
-               # RuntimeError: cumulative ops are not yet supported for complex
-               DecorateInfo(
-                   unittest.expectedFailure, 'TestCommon', 'test_noncontiguous_samples',
-                   device_type='mps', dtypes=(torch.complex64,)
-               ),
-               DecorateInfo(
-                   unittest.expectedFailure, 'TestCommon', 'test_out_requires_grad_error',
-                   device_type='mps', dtypes=(torch.complex64,)
-               ),
-               DecorateInfo(
-                   unittest.expectedFailure, 'TestCommon', 'test_variant_consistency_eager',
-                   device_type='mps', dtypes=(torch.complex64,)
-               ),
            ),
            sample_inputs_func=sample_inputs_cumulative_ops),
     OpInfo('cumprod',
@@ -13614,21 +13660,8 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
            skips=(
                # cumprod does not handle correctly out= dtypes
                DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out'),
-               # The following dtypes did not work in forward but are listed by the OpInfo: {torch.complex64}
+               # The following dtypes did not work in forward but are listed by the OpInfo: {torch.bool}
                DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_dtypes', device_type='mps'),
-               # RuntimeError: cumulative ops are not yet supported for complex
-               DecorateInfo(
-                   unittest.expectedFailure, 'TestCommon', 'test_noncontiguous_samples',
-                   device_type='mps', dtypes=(torch.complex64,)
-               ),
-               DecorateInfo(
-                   unittest.expectedFailure, 'TestCommon', 'test_out_requires_grad_error',
-                   device_type='mps', dtypes=(torch.complex64,)
-               ),
-               DecorateInfo(
-                   unittest.expectedFailure, 'TestCommon', 'test_variant_consistency_eager',
-                   device_type='mps', dtypes=(torch.complex64,)
-               ),
            ),
            # gradgradcheck fails in fast_mode=True: #56275
            sample_inputs_func=sample_inputs_cumprod,
@@ -13780,6 +13813,8 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
                        # Reference: https://github.com/pytorch/pytorch/issues/48010
                        DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_extremal',
                                     device_type='cpu', dtypes=[torch.cfloat, torch.cdouble], active_if=IS_WINDOWS),
+                       DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_large',
+                                    device_type='xpu', dtypes=(torch.chalf,), active_if=IS_WINDOWS),
                    ),
                    assert_autodiffed=True,
                    supports_forward_ad=True,
@@ -13930,8 +13965,6 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
                                      'test_fn_grad',
                                      dtypes=(torch.float64,),
                                      device_type='cpu'),
-                        # NotImplementedError: "fmod_cpu" not implemented for 'Bool'
-                        DecorateInfo(unittest.expectedFailure, 'TestConsistency', device_type='mps', dtypes=(torch.bool,)),
                     )),
     BinaryUfuncInfo('remainder',
                     ref=np.remainder,
@@ -14046,15 +14079,6 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
                DecorateInfo(unittest.expectedFailure, 'TestBwdGradients', 'test_fn_grad'),
                # Pre-existing condition (calls .item); needs to be fixed
                DecorateInfo(unittest.expectedFailure, 'TestCompositeCompliance', 'test_backward'),
-               # RuntimeError: Failed to create function state object for: unfold_backward_float2
-               DecorateInfo(
-                   unittest.expectedFailure, 'TestCommon', 'test_noncontiguous_samples',
-                   device_type='mps', dtypes=(torch.complex64,)
-               ),
-               DecorateInfo(
-                   unittest.expectedFailure, 'TestCommon', 'test_variant_consistency_eager',
-                   device_type='mps', dtypes=(torch.complex64,)
-               ),
            )),
     UnaryUfuncInfo('floor',
                    ref=np.floor,
@@ -14068,8 +14092,6 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
                                     'TestNNCOpInfo',
                                     'test_nnc_correctness',
                                     dtypes=tuple(t for t in integral_types() if t != torch.uint8)),
-                       # NotImplementedError: "floor_vml_cpu" not implemented for 'Bool'
-                       DecorateInfo(unittest.expectedFailure, 'TestConsistency', device_type='mps', dtypes=(torch.bool,)),
                    ),
                    supports_sparse=True,
                    supports_sparse_csr=True,
@@ -14234,8 +14256,6 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
                                      dtypes=(torch.float16,)),
                         DecorateInfo(toleranceOverride({torch.float16: tol(atol=1e-3, rtol=5e-3)}),
                                      'TestBinaryUfuncs', 'test_reference_numerics'),
-                        # NotImplementedError: "div_floor_cpu" not implemented for 'Bool'
-                        DecorateInfo(unittest.expectedFailure, 'TestConsistency', device_type='mps', dtypes=(torch.bool,)),
                     )),
     UnaryUfuncInfo('frexp',
                    op=torch.frexp,
@@ -14720,8 +14740,6 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
                DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out_warning', device_type='mps'),
                # Exception: linalg.solve.triangular(); Only float is supported!
                DecorateInfo(unittest.expectedFailure, 'TestCommon', device_type='mps', dtypes=(torch.complex64,)),
-               # AssertionError: Tensor-likes are not close!
-               DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_noncontiguous_samples', device_type='mps'),
                DecorateInfo(unittest.skip("Tests different backward paths"),
                             "TestCommon", "test_floating_inputs_are_differentiable"),),
            decorators=[skipCPUIfNoLapack, skipCUDAIfNoMagmaAndNoCusolver]),
@@ -14750,15 +14768,9 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
                             dtypes=[torch.bool], active_if=TEST_WITH_ROCM),
                # Exception: RuntimeError not raised : inplace variant either incorrectly allowed resizing
                DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_variant_consistency_eager', device_type='mps'),
-               # Exception: Tensor-likes are not equal!
-               DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_noncontiguous_samples', device_type='mps'),
                # Exception: "masked_scatter_ only supports boolean masks" does not match
                # "masked_scatter: expected BoolTensor or ByteTensor for mask"
                DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_errors', device_type='mps'),
-               # The following dtypes did not work in backward but are listed by
-               # the OpInfo: {torch.float32, torch.float16, torch.complex64,
-               # torch.bfloat16}.
-               DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_dtypes', device_type='mps'),
            )),
     OpInfo('masked_select',
            dtypes=all_types_and_complex_and(torch.bool, torch.half, torch.bfloat16),
@@ -15400,18 +15412,15 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
            dtypes=floating_and_complex_types_and(torch.half, torch.bfloat16),
            sample_inputs_func=sample_inputs_normalize,
            supports_forward_ad=True,
-           skips=(
-               # Exception: norm ops are not supported for complex yet
-               DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_dtypes', device_type='mps'),
-               DecorateInfo(unittest.expectedFailure, 'TestCommon', device_type='mps', dtypes=(torch.complex64,)),
-           ),
            supports_fwgrad_bwgrad=True),
     OpInfo('aminmax',
            ref=lambda x, dim=None, keepdim=False: (np.amin(x, axis=dim, keepdims=keepdim), np.amax(x, axis=dim, keepdims=keepdim)),
            dtypes=all_types_and(torch.bool, torch.float16, torch.bfloat16),
            dtypesIfHpu=custom_types(torch.float32, torch.bfloat16, torch.int32, torch.int8),
            decorators=(onlyNativeDeviceTypes,),
-           supports_autograd=False,
+           supports_autograd=True,
+           supports_forward_ad=True,
+           supports_fwgrad_bwgrad=True,
            sample_inputs_func=sample_inputs_aminmax,
            skips=(
                # Exception: MPS supports tensors with dimensions <= 16, but got 65.
@@ -15544,11 +15553,6 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
                DecorateInfo(unittest.skip("Unsupported on MPS for now"), 'TestCommon', 'test_numpy_ref_mps'),
                DecorateInfo(toleranceOverride({torch.float32: tol(atol=2e-03, rtol=5e-03)}),
                             "TestDecomp", "test_comprehensive", device_type="cpu"),
-               # AssertionError: Tensor-likes are not close!
-               DecorateInfo(
-                   unittest.expectedFailure, 'TestCommon', 'test_noncontiguous_samples',
-                   device_type='mps', dtypes=(torch.float32,)
-               ),
                # See https://github.com/pytorch/pytorch/issues/173525
                DecorateInfo(toleranceOverride({torch.float32: tol(atol=2e-4, rtol=2e-5)}),
                             'TestConsistency', 'test_output_grad_match', device_type='mps'),
@@ -16052,6 +16056,10 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
                DecorateInfo(
                    toleranceOverride({torch.chalf: tol(atol=9e-2, rtol=9e-2), }),
                    'TestCommon', 'test_complex_half_reference_testing'),
+               DecorateInfo(
+                   toleranceOverride({torch.float32: tol(atol=5e-5, rtol=5e-6)}),
+                   'TestOperators', 'test_vjpvmap', device_type='cuda'
+               ),
                DecorateInfo(
                    toleranceOverride({torch.half: tol(atol=9e-3, rtol=2e-1), }),
                    'TestInductorOpInfo', 'test_comprehensive', device_type='cpu')],
@@ -16148,6 +16156,10 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
                    toleranceOverride({torch.float16: tol(atol=5e-3, rtol=1e-3)}),
                    'TestInductorOpInfo', 'test_comprehensive',
                ),
+               DecorateInfo(
+                   toleranceOverride({torch.float32: tol(atol=1e-5, rtol=1e-5)}),
+                   'TestCommon', 'test_noncontiguous_samples', device_type='mps',
+               ),
            ),
            skips=(
                # RuntimeError: !lhs.isAliasOf(rhs)INTERNAL ASSERT FAILED at
@@ -16162,10 +16174,9 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
                             'test_nnc_correctness', dtypes=(torch.complex64, torch.complex128)),
                # RuntimeError: Convolution is supported only for Floating types
                DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_dtypes', device_type='mps'),
-               # AssertionError: Tensor-likes are not close!
                DecorateInfo(
-                   unittest.expectedFailure, 'TestCommon', 'test_noncontiguous_samples',
-                   device_type='mps', dtypes=(torch.float32, torch.int64)
+                   unittest.expectedFailure, 'TestCommon',
+                   device_type='mps', dtypes=(torch.int64,)
                ),
            ),
            supports_expanded_weight=True,
@@ -16205,6 +16216,18 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
                    toleranceOverride({torch.float32: tol(atol=5e-5, rtol=5e-6)}),
                    'TestOperators', 'test_vjpvmap',
                ),
+               DecorateInfo(
+                   toleranceOverride({torch.float32: tol(atol=5e-5, rtol=5e-6)}),
+                   'TestOperators', 'test_jvpvjp', device_type="cuda"
+               ),
+               DecorateInfo(
+                   toleranceOverride({torch.float32: tol(atol=5e-5, rtol=5e-6)}),
+                   'TestOperators', 'test_vjp', device_type="cuda"
+               ),
+               DecorateInfo(
+                   toleranceOverride({torch.float32: tol(atol=5e-5, rtol=5e-6)}),
+                   'TestCompositeCompliance', 'test_backward', device_type="cuda"
+               ),
                DecorateInfo(
                    toleranceOverride({torch.float16: tol(atol=5e-3, rtol=1e-3)}),
                    'TestInductorOpInfo', 'test_comprehensive',
@@ -16235,7 +16258,7 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
            aliases=('group_norm',),
            ref=reference_group_norm,
            dtypes=floating_types_and(torch.float16, torch.bfloat16),
-           dtypesIfMPS=floating_types_and(torch.float16, torch.bfloat16, torch.int32, torch.int16),
+           dtypesIfMPS=floating_types_and(torch.float16, torch.bfloat16, torch.int32, torch.int16, torch.int8, torch.uint8),
            supports_out=False,
            supports_forward_ad=True,
            supports_fwgrad_bwgrad=True,
@@ -16250,6 +16273,9 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
                    "test_comprehensive",
                    device_type="cpu"
                ),
+               # MPS supports int8/uint8 but CPU does not, so consistency test cannot run
+               DecorateInfo(unittest.expectedFailure, 'TestConsistency', 'test_output_match',
+                            device_type='mps', dtypes=(torch.int8, torch.uint8)),
            ],
            sample_inputs_func=sample_inputs_group_norm,
            reference_inputs_func=reference_inputs_group_norm,
@@ -16285,11 +16311,6 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
                    'TestCommon', 'test_numpy_refs'
                ),
                DecorateInfo(unittest.skip("Bug in MPS backend!"), 'TestCommon', 'test_numpy_ref_mps'),
-               # AssertionError: Tensor-likes are not close!
-               DecorateInfo(
-                   unittest.expectedFailure, 'TestCommon', 'test_noncontiguous_samples',
-                   device_type='mps', dtypes=(torch.float32,)
-               ),
                # See https://github.com/pytorch/pytorch/issues/173525
                DecorateInfo(toleranceOverride({torch.float32: tol(atol=2e-4, rtol=2e-5)}),
                             'TestConsistency', 'test_output_grad_match', device_type='mps'),
@@ -16609,7 +16630,6 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
            supports_out=False,
            skips=(
                DecorateInfo(unittest.expectedFailure, 'TestJit', 'test_variant_consistency_jit'),
-               DecorateInfo(unittest.expectedFailure, 'TestDTensorOps', 'test_dtensor_op_db'),
                DecorateInfo(unittest.expectedFailure, 'TestInductorOpInfo', 'test_comprehensive'),
                DecorateInfo(unittest.expectedFailure, 'TestMathBits', 'test_neg_view'),
                # Error: The operator 'aten::_upsample_bilinear2d_aa_backward.grad_input'
@@ -17036,9 +17056,8 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
            supports_expanded_weight=True,
            skips=(
                # RuntimeError: MPS device does not support linear for non-float inputs
-               # RuntimeError: mps linear does not support complex types
                DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_dtypes', device_type='mps'),
-               DecorateInfo(unittest.expectedFailure, 'TestCommon', device_type='mps', dtypes=(torch.complex64, torch.int64)),
+               DecorateInfo(unittest.expectedFailure, 'TestCommon', device_type='mps', dtypes=(torch.int64,)),
            ),
            decorators=(
                # Strides are not the same!
@@ -17595,6 +17614,9 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
                          'TestUnaryUfuncs', "test_reference_numerics_extremal",
                          dtypes=(torch.complex64, torch.complex128), device_type='cpu',
                          active_if=(IS_MACOS or IS_WINDOWS)),
+            DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_extremal',
+                         device_type='xpu', dtypes=[torch.cfloat, torch.cdouble],
+                         active_if=IS_WINDOWS),
         ),
         # tan(j * pi/2 * odd_number) is nan which also make tanhshrink nan.
         reference_numerics_filter=NumericsFilter(
@@ -17630,11 +17652,6 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
         supports_out=False,
         supports_forward_ad=True,
         supports_fwgrad_bwgrad=True,
-        skips=(
-            # RuntimeError: norm ops are not supported for complex yet
-            DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_dtypes', device_type='mps'),
-            DecorateInfo(unittest.expectedFailure, 'TestCommon', device_type='mps', dtypes=(torch.complex64,)),
-        ),
     ),
     OpInfo(
         "nn.functional.triplet_margin_with_distance_loss",
@@ -17657,9 +17674,6 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
                 "TestNormalizeOperators",
                 "test_normalize_operator_exhaustive",
             ),
-            # RuntimeError: norm ops are not supported for complex yet
-            DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_dtypes', device_type='mps'),
-            DecorateInfo(unittest.expectedFailure, 'TestCommon', device_type='mps', dtypes=(torch.complex64,)),
         ),
     ),
     BinaryUfuncInfo('nextafter',
@@ -17987,30 +18001,15 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
            sample_inputs_func=sample_inputs_mode,),
     make_mvlgamma_opinfo(variant_test_name='mvlgamma_p_1',
                          domain=(1, None),
-                         skips=skips_mvlgamma() + (
-                             # NotImplementedError: The operator 'aten::mvlgamma.out' is not currently
-                             # implemented for the MPS device
-                             DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out', device_type='mps'),
-                             DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out_warning', device_type='mps'),
-                         ),
+                         skips=skips_mvlgamma(),
                          sample_kwargs=lambda device, dtype, input: ({'p': 1}, {'d': 1})),
     make_mvlgamma_opinfo(variant_test_name='mvlgamma_p_3',
                          domain=(2, None),
-                         skips=skips_mvlgamma() + (
-                             # NotImplementedError: The operator 'aten::mvlgamma.out' is not currently
-                             # implemented for the MPS device
-                             DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out', device_type='mps'),
-                             DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out_warning', device_type='mps'),
-                         ),
+                         skips=skips_mvlgamma(),
                          sample_kwargs=lambda device, dtype, input: ({'p': 3}, {'d': 3})),
     make_mvlgamma_opinfo(variant_test_name='mvlgamma_p_5',
                          domain=(3, None),
-                         skips=skips_mvlgamma() + (
-                             # NotImplementedError: The operator 'aten::mvlgamma.out' is not currently
-                             # implemented for the MPS device
-                             DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out', device_type='mps'),
-                             DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out_warning', device_type='mps'),
-                         ),
+                         skips=skips_mvlgamma(),
                          sample_kwargs=lambda device, dtype, input: ({'p': 5}, {'d': 5})),
     BinaryUfuncInfo('ne',
                     ref=np.not_equal,
@@ -18038,9 +18037,8 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
     OpInfo('narrow_copy',
            dtypes=all_types_and_complex_and(torch.bool, torch.bfloat16, torch.float16, torch.chalf),
            supports_out=True,
-           supports_forward_ad=False,
-           supports_fwgrad_bwgrad=False,
-           supports_autograd=False,
+           supports_forward_ad=True,
+           supports_fwgrad_bwgrad=True,
            # https://github.com/pytorch/pytorch/issues/86931
            sample_inputs_func=partial(sample_inputs_narrow_narrow_copy, is_narrow=False),
            reference_inputs_func=partial(reference_inputs_narrow_narrow_copy, is_narrow=False),
@@ -18098,18 +18096,6 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
            # Could not allocate memory to change Tensor SizesAndStrides!
            check_batched_forward_grad=False,
            supports_fwgrad_bwgrad=True,
-           skips=(
-               # norm ops are not supported for complex yet
-               DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_dtypes', device_type='mps'),
-               DecorateInfo(
-                   unittest.expectedFailure, 'TestCommon', 'test_noncontiguous_samples',
-                   device_type='mps', dtypes=(torch.complex64,)
-               ),
-               DecorateInfo(
-                   unittest.expectedFailure, 'TestCommon', 'test_variant_consistency_eager',
-                   device_type='mps', dtypes=(torch.complex64,)
-               ),
-           ),
            sample_inputs_func=sample_inputs_dist),
     OpInfo('outer',
            op=torch.outer,
@@ -18270,17 +18256,14 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
     OpInfo('qr',
            op=torch.qr,
            dtypes=floating_and_complex_types(),
+           dtypesIfMPS=floating_types(),
            sample_inputs_func=sample_inputs_linalg_qr_geqrf,
            supports_forward_ad=True,
            supports_fwgrad_bwgrad=True,
            # In-place ops
            check_batched_gradgrad=False,
            decorators=[skipCUDAIfNoCusolver, skipCPUIfNoLapack],
-           skips=(
-               # NotImplementedError: The operator 'aten::linalg_qr.out' is not
-               # currently implemented for the MPS device
-               DecorateInfo(unittest.expectedFailure, 'TestCommon', device_type='mps'),
-           )),
+           ),
     UnaryUfuncInfo('rad2deg',
                    ref=np.degrees,
                    decorators=(precisionOverride({torch.bfloat16: 7e-1,
@@ -18433,6 +18416,10 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
                                     dtypes=(torch.cfloat, torch.cdouble,), device_type='cpu', active_if=IS_WINDOWS),
                        DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_large',
                                     dtypes=(torch.cfloat, torch.cdouble,), device_type='cpu', active_if=IS_WINDOWS),
+                       DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_extremal',
+                                    device_type='xpu', dtypes=(torch.cfloat, torch.cdouble), active_if=IS_WINDOWS),
+                       DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_large',
+                                    device_type='xpu', dtypes=(torch.chalf,), active_if=IS_WINDOWS),
                        DecorateInfo(unittest.skip("Skipped! sparse backward not supported"),
                                     'TestSparseUnaryUfuncs', 'test_sparse_fn_grad'),
                        DecorateInfo(toleranceOverride({torch.float16: tol(atol=1e-3, rtol=2e-3)}),
@@ -18465,6 +18452,11 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
                        DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_extremal',
                                     device_type='cpu', dtypes=[torch.cfloat, torch.cdouble],
                                     active_if=(IS_MACOS or IS_WINDOWS)),
+                       DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_extremal',
+                                    device_type='xpu', dtypes=[torch.cfloat, torch.cdouble],
+                                    active_if=IS_WINDOWS),
+                       DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_large',
+                                    device_type='xpu', dtypes=(torch.chalf,), active_if=IS_WINDOWS),
                        DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_large',
                                     device_type='cpu', dtypes=[torch.cfloat, torch.cdouble],
                                     active_if=(IS_MACOS or IS_WINDOWS)),
@@ -18714,8 +18706,6 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
                         # TODO: FIXME tolerance is too high
                         DecorateInfo(unittest.skip('Skipped!'), 'TestFwdGradients'),
                         DecorateInfo(unittest.skip('Skipped!'), 'TestBwdGradients'),
-                        # Exception: "pow" not implemented for 'Bool'
-                        DecorateInfo(unittest.expectedFailure, 'TestConsistency', device_type='mps', dtypes=(torch.bool,)),
                     ),
                     assert_autodiffed=True,
                     autodiff_nonfusible_nodes=['aten::pow'],),
@@ -18805,6 +18795,11 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
                        DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_large',
                                     device_type='cpu', dtypes=[torch.cfloat, torch.cdouble],
                                     active_if=(IS_MACOS or IS_WINDOWS)),
+                       DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_extremal',
+                                    device_type='xpu', dtypes=[torch.cfloat, torch.cdouble],
+                                    active_if=IS_WINDOWS),
+                       DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_large',
+                                    device_type='xpu', dtypes=(torch.chalf,), active_if=IS_WINDOWS),
                        DecorateInfo(unittest.skip("Skipped! sparse backward not supported"),
                                     'TestSparseUnaryUfuncs', 'test_sparse_fn_grad'),
                        # FIXME:
@@ -18854,6 +18849,11 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
                        DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_large',
                                     device_type='cpu', dtypes=[torch.cfloat, torch.cdouble],
                                     active_if=(IS_MACOS or IS_WINDOWS)),
+                       DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_extremal',
+                                    device_type='xpu', dtypes=[torch.cfloat, torch.cdouble],
+                                    active_if=IS_WINDOWS),
+                       DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_large',
+                                    device_type='xpu', dtypes=(torch.chalf,), active_if=IS_WINDOWS),
                        DecorateInfo(unittest.skip("Skipped! sparse backward not supported"),
                                     'TestSparseUnaryUfuncs', 'test_sparse_fn_grad'),
                        DecorateInfo(toleranceOverride({torch.complex64: tol(atol=3e-5, rtol=7e-6)}),
@@ -18938,10 +18938,14 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
                # RuntimeError: linalg.solve.triangular(); Only float is supported!
                DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_dtypes', device_type='mps'),
                DecorateInfo(unittest.expectedFailure, 'TestCommon', device_type='mps', dtypes=(torch.complex64,)),
-               # AssertionError: Tensor-likes are not close!
+               # see https://github.com/pytorch/pytorch/issues/177251
                DecorateInfo(
-                   unittest.expectedFailure, 'TestCommon', 'test_noncontiguous_samples',
-                   device_type='mps', dtypes=(torch.float32,)
+                   unittest.expectedFailure,
+                   'TestOperators',
+                   'test_jvp',
+                   device_type='cpu',
+                   dtypes=[torch.float32],
+                   active_if=IS_ARM64 and IS_CPU_EXT_SVE_SUPPORTED,
                ),
            )),
     UnaryUfuncInfo('trunc',
@@ -19097,6 +19101,8 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
                        # tensor(inf+nanj, device='cuda:0')
                        DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_extremal',
                                     device_type='cuda', dtypes=[torch.cfloat, torch.cdouble]),
+                       DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_extremal',
+                                    device_type='xpu', dtypes=[torch.cfloat, torch.cdouble]),
                        DecorateInfo(unittest.expectedFailure, 'TestMeta', 'test_meta_inplace',
                                     dtypes=[torch.bool]),
                        DecorateInfo(unittest.expectedFailure, 'TestMeta', 'test_dispatch_meta_inplace',
@@ -19227,22 +19233,6 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
                    'TestSchemaCheckModeOpInfo',
                    'test_schema_correctness',
                    dtypes=(torch.complex64, torch.complex128)),
-               DecorateInfo(unittest.skip("Skipped!"), 'TestCommon', 'test_out',
-                            device_type='mps', dtypes=[torch.float32]),
-               DecorateInfo(unittest.skip("Skipped!"), 'TestCommon', 'test_variant_consistency_eager',
-                            device_type='mps', dtypes=[torch.float32]),
-               DecorateInfo(unittest.skip("Skipped!"), 'TestJit', 'test_variant_consistency_jit',
-                            device_type='mps', dtypes=[torch.float32]),
-               # The operator 'aten::take' is not currently implemented for the MPS device
-               DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out_warning', device_type='mps'),
-               # RuntimeError: svd_backward: The singular vectors in the complex
-               # case are specified up to multiplication by e^{i phi}. The
-               # specified loss function depends on this phase term, making it
-               # ill-defined.
-               DecorateInfo(
-                   unittest.expectedFailure, 'TestCommon', 'test_noncontiguous_samples',
-                   device_type='mps', dtypes=(torch.complex64,)
-               ),
            )),
     OpInfo('svd_lowrank',
            op=lambda *args, **kwargs: wrapper_set_seed(
@@ -19250,6 +19240,7 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
                *args, **kwargs
            ),
            dtypes=floating_and_complex_types(),
+           dtypesIfMPS=floating_types(),
            # Runs very slowly on slow gradcheck - alternatively reduce input sizes
            gradcheck_fast_mode=True,
            supports_out=False,
@@ -19281,12 +19272,6 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
                DecorateInfo(unittest.expectedFailure, 'TestSchemaCheckModeOpInfo', 'test_schema_correctness',
                             dtypes=(torch.complex64, torch.complex128)),
                DecorateInfo(slowTest, 'TestCompositeCompliance', 'test_forward_ad'),
-               # NotImplementedError: The operator 'aten::linalg_qr.out' is not currently implemented for the MPS device
-               DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_dtypes', device_type='mps'),
-               DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out', device_type='mps'),
-               DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out_warning', device_type='mps'),
-               DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_variant_consistency_eager', device_type='mps'),
-               DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_noncontiguous_samples', device_type='mps')
            )),
     OpInfo('pca_lowrank',
            op=lambda *args, **kwargs: wrapper_set_seed(
@@ -19294,6 +19279,7 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
                *args, **kwargs
            ),
            dtypes=floating_and_complex_types(),
+           dtypesIfMPS=floating_types(),
            # Runs very slowly on slow gradcheck - alternatively reduce input sizes
            gradcheck_fast_mode=True,
            supports_out=False,
@@ -19319,6 +19305,9 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
                        DecorateInfo(
                            toleranceOverride({torch.float32: tol(atol=3e-5, rtol=1e-3)}),
                            'TestInductorOpInfo', 'test_comprehensive', device_type='cuda'),
+                       DecorateInfo(
+                           toleranceOverride({torch.float32: tol(atol=5e-5, rtol=3e-4)}),
+                           'TestConsistency', 'test_output_grad_match', device_type='mps'),
                        ],
            skips=(
                # test does not work with passing lambda for op
@@ -19328,9 +19317,6 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
                DecorateInfo(unittest.expectedFailure, 'TestSchemaCheckModeOpInfo', 'test_schema_correctness',
                             dtypes=(torch.complex64, torch.complex128)),
                DecorateInfo(unittest.skip('output is non-deterministic'), 'TestCommon', 'test_compare_cpu'),
-               # NotImplementedError: The operator 'aten::linalg_qr.out' is not
-               # currently implemented for the MPS device
-               DecorateInfo(unittest.expectedFailure, 'TestCommon', device_type='mps'),
            )),
     BinaryUfuncInfo('polar',
                     dtypes=floating_types(),
@@ -19371,15 +19357,6 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
                    sample_inputs_func=sample_inputs_polygamma,
                    skips=(
                        DecorateInfo(unittest.expectedFailure, "TestNormalizeOperators", "test_normalize_operator_exhaustive"),
-                       # AssertionError: Tensor-likes are not close!
-                       DecorateInfo(
-                           unittest.expectedFailure, 'TestCommon', 'test_noncontiguous_samples',
-                           device_type='mps', dtypes=(torch.float32, torch.int64)
-                       ),
-                       DecorateInfo(
-                           unittest.expectedFailure, 'TestCommon', 'test_out',
-                           device_type='mps', dtypes=(torch.float32,)
-                       ),
                    ),
                    sample_kwargs=lambda device, dtype, input: ({'n': 0}, {'n': 0}),
                    # polygamma functions have multiple singularities at x having non-positive integer value
@@ -19586,12 +19563,6 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
            supports_forward_ad=True,
            supports_fwgrad_bwgrad=True,
            error_inputs_func=error_inputs_gather,
-           skips=(
-               # RuntimeError: gather(): Yet not supported for complex
-               DecorateInfo(unittest.expectedFailure, 'TestCommon', device_type='mps', dtypes=(torch.complex64,)),
-               # TypeError: Cannot convert a MPS Tensor to float64 dtype as the MPS framework doesn't support float64
-               DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_errors', device_type='mps'),
-           ),
            ),
     OpInfo('index_fill',
            dtypes=all_types_and_complex_and(torch.bool, torch.float16, torch.bfloat16, torch.complex32),
@@ -19606,10 +19577,6 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
                DecorateInfo(unittest.expectedFailure, 'TestFakeTensor', 'test_fake_crossref_backward_no_amp'),
                # RuntimeError: Mismatch on aten._unique.default: Shapes torch.Size([2]) and torch.Size([1]) are not equal!
                DecorateInfo(unittest.expectedFailure, 'TestFakeTensor', 'test_fake_crossref_backward_amp'),
-               # RuntimeError: index_fill_(): Complex types are yet not supported
-               # NotImplementedError: "_local_scalar_dense_mps" not implemented for 'ComplexHalf'
-               DecorateInfo(unittest.expectedFailure, 'TestCommon', device_type='mps', dtypes=(torch.complex64, torch.complex32)),
-               DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_dtypes', device_type='mps'),
            ),
            sample_inputs_func=sample_inputs_index,
            reference_inputs_func=partial(sample_inputs_index, reference=True)),
@@ -19622,18 +19589,6 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
            check_batched_forward_grad=False,
            sample_inputs_func=sample_inputs_index,
            reference_inputs_func=partial(sample_inputs_index, reference=True),
-           skips=(
-               # Exception: index_fill_(): Complex types are yet not supported
-               DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_dtypes', device_type='mps'),
-               DecorateInfo(
-                   unittest.expectedFailure, 'TestCommon', 'test_variant_consistency_eager',
-                   device_type='mps', dtypes=(torch.complex64,)
-               ),
-               DecorateInfo(
-                   unittest.expectedFailure, 'TestCommon', 'test_noncontiguous_samples',
-                   device_type='mps', dtypes=(torch.complex64,)
-               ),
-           ),
            gradcheck_nondet_tol=GRADCHECK_NONDET_TOL),
     OpInfo('index_select',
            dtypes=all_types_and_complex_and(torch.bool, torch.float16, torch.bfloat16, torch.chalf),
@@ -19675,11 +19630,13 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
     *(OpInfo('index_reduce',
              variant_test_name=reduction_type,
              dtypes=all_types_and(torch.float16, torch.bfloat16),
+             dtypesIfMPS=custom_types(
+                 torch.float32, torch.bfloat16, torch.float16, torch.int32,
+                 torch.int16, torch.int8, torch.uint8,
+             ),
              skips=(
                  DecorateInfo(toleranceOverride({torch.float16: tol(atol=2e-3, rtol=3e-3)}),
                               'TestInductorOpInfo', 'test_comprehensive'),
-                 # Error: The operator 'aten::index_reduce.out' is not currently implemented for the MPS device
-                 DecorateInfo(unittest.expectedFailure, 'TestCommon', device_type='mps'),
              ),
              supports_out=True,
              sample_inputs_func=sample_inputs_index_reduce,
@@ -19810,11 +19767,6 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
                # Compiler issue on ROCm. Regression started in ROCm 6.4.
                DecorateInfo(unittest.skip('Skipped!'), 'TestCommon', 'test_non_standard_bool_values',
                             dtypes=[torch.bool], active_if=TEST_WITH_ROCM),
-               # RuntimeError: scatter(): Yet not supported for complex
-               DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_dtypes', device_type='mps'),
-               DecorateInfo(unittest.expectedFailure, 'TestCommon', device_type='mps', dtypes=(torch.complex64,)),
-               # TypeError: Cannot convert a MPS Tensor to float64 dtype as the MPS framework doesn't support float64
-               DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_errors', device_type='mps'),
            )),
     UnaryUfuncInfo(
         'bfloat16',
@@ -20360,8 +20312,6 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
                DecorateInfo(unittest.expectedFailure, 'TestMeta', 'test_dispatch_meta_outplace'),
                DecorateInfo(unittest.expectedFailure, 'TestMeta', 'test_dispatch_symbolic_meta_outplace'),
                DecorateInfo(unittest.expectedFailure, 'TestMeta', 'test_dispatch_symbolic_meta_outplace_all_strides'),
-               # MPS: The following dtypes worked in forward but are not listed by the OpInfo: {torch.complex128}.
-               DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_dtypes', device_type='mps'),
            )),
     OpInfo('empty',
            dtypes=all_types_and_complex_and(torch.bool, torch.half, torch.bfloat16, torch.chalf),
@@ -20402,8 +20352,6 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
                DecorateInfo(unittest.skip("Expected: empty is not comparable"),
                             'TestCommon', 'test_complex_half_reference_testing'),
                DecorateInfo(unittest.skip('output is non-deterministic'), 'TestCommon', 'test_compare_cpu'),
-               # MPS: The following dtypes worked in forward but are not listed by the OpInfo: {torch.complex128}.
-               DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_dtypes', device_type='mps'),
            )),
     OpInfo('eye',
            dtypes=all_types_complex_float8_and(torch.bool, torch.half, torch.bfloat16),
@@ -20473,8 +20421,6 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
                DecorateInfo(unittest.skip("Expected: empty_permuted is not comparable"),
                             'TestCommon', 'test_complex_half_reference_testing'),
                DecorateInfo(unittest.skip('output is non-deterministic'), 'TestCommon', 'test_compare_cpu'),
-               # MPS: The following dtypes worked in forward but are not listed by the OpInfo: {torch.complex128}.
-               DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_dtypes', device_type='mps'),
            )),
     OpInfo('scalar_tensor',
            dtypes=all_types_and_complex_and(torch.bool, torch.half, torch.bfloat16, torch.chalf),
@@ -20493,8 +20439,6 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
                DecorateInfo(unittest.skip("Skipped!"), 'TestMathBits', 'test_conj_view'),
                DecorateInfo(unittest.skip("Skipped!"), 'TestMathBits', 'test_neg_conj_view'),
                DecorateInfo(unittest.skip("Skipped!"), 'TestMathBits', 'test_neg_view'),
-               # MPS: The following dtypes worked in forward but are not listed by the OpInfo: {torch.complex128}.
-               DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_dtypes', device_type='mps'),
            )),
     OpInfo('new_full',
            op=lambda x, *args, **kwargs: x.new_full(*args, **kwargs),
@@ -20516,8 +20460,6 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
            error_inputs_func=error_inputs_multinomial,
            skips=(
                DecorateInfo(unittest.expectedFailure, "TestNormalizeOperators", "test_normalize_operator_exhaustive"),
-               # TypeError: Cannot convert a MPS Tensor to float64 dtype as the MPS framework doesn't support float64
-               DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_errors', device_type='mps'),
                # Strides are not the same!
                # This may not be reproducible in CI
                DecorateInfo(unittest.skip("Skipped!"), 'TestCommon', 'test_out'),
@@ -20624,7 +20566,7 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
                # NotImplementedError: "bernoulli_tensor_cpu_p_" not implemented for *
                DecorateInfo(
                    unittest.expectedFailure, 'TestConsistency', device_type='mps',
-                   dtypes=(torch.bool, torch.uint8, torch.int8, torch.int16, torch.int32, torch.int64)
+                   dtypes=(torch.uint8, torch.int8, torch.int16, torch.int32, torch.int64)
                ),
                DecorateInfo(unittest.skip('output is non-deterministic'), 'TestCommon', 'test_compare_cpu'))),
     OpInfo('scatter_add',
@@ -20638,11 +20580,6 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
                # Compiler issue on ROCm. Regression started in ROCm 6.4.
                DecorateInfo(unittest.skip('Skipped!'), 'TestCommon', 'test_non_standard_bool_values',
                             dtypes=[torch.bool], active_if=TEST_WITH_ROCM),
-               # RuntimeError: scatter(): Yet not supported for complex
-               DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_dtypes', device_type='mps'),
-               DecorateInfo(unittest.expectedFailure, 'TestCommon', device_type='mps', dtypes=(torch.complex64,)),
-               # TypeError: Cannot convert a MPS Tensor to float64 dtype as the MPS framework doesn't support float64
-               DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_errors', device_type='mps'),
            )),
     OpInfo('stack',
            dtypes=all_types_and_complex_and(torch.complex32, torch.bool, torch.float16, torch.bfloat16),
@@ -20692,11 +20629,6 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
                DecorateInfo(unittest.expectedFailure, 'TestJit', 'test_variant_consistency_jit'),
                # Not Implemented on XLA.
                DecorateInfo(unittest.skip("Skipped!"), 'TestOpInfo', device_type='xla'),
-               # RuntimeError: histogramdd(): weight should be contiguous on MPS
-               DecorateInfo(
-                   unittest.expectedFailure, 'TestCommon', 'test_noncontiguous_samples',
-                   device_type='mps', dtypes=(torch.float32,)
-               ),
            )),
     OpInfo('histogramdd',
            dtypes=floating_types(),
@@ -20710,11 +20642,6 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
                # JIT tests don't work with Tensor keyword arguments
                # https://github.com/pytorch/pytorch/issues/58507
                DecorateInfo(unittest.expectedFailure, 'TestJit', 'test_variant_consistency_jit'),
-               # RuntimeError: histogramdd(): weight should be contiguous on MPS
-               DecorateInfo(
-                   unittest.expectedFailure, 'TestCommon', 'test_noncontiguous_samples',
-                   device_type='mps', dtypes=(torch.float32,)
-               ),
            )),
     OpInfo('histc',
            dtypes=floating_types_and(torch.bfloat16, torch.float16),
@@ -20727,8 +20654,6 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
                # "AssertionError: RuntimeError not raised : Expected RuntimeError when doing an unsafe cast
                # from a result of dtype torch.float32 into an out= with dtype torch.long"
                DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out', device_type='cuda'),
-               # The following dtypes did not work in forward but are listed by the OpInfo: {torch.bfloat16, torch.float16}.
-               DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_dtypes', device_type='mps'),
            )),
     OpInfo('bincount',
            dtypes=integral_types_and(),
@@ -20786,8 +20711,6 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
            check_batched_forward_grad=False,
            assert_autodiffed=True,
            skips=(
-               # https://github.com/pytorch/pytorch/issues/89353
-               DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_numpy_ref_mps'),
                # RuntimeError: Arguments for call not valid.
                #               Expected a value of type 'List[Tensor]' for argument
                #               'tensors' but instead found type 'Tensor (inferred)'.
@@ -20797,8 +20720,6 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
                # see https://github.com/pytorch/pytorch/issues/99806
                # RuntimeError: The size of tensor a (25) must match the size of tensor b (0) at non-singleton dimension 0.
                DecorateInfo(unittest.expectedFailure, 'TestBwdGradients', 'test_fn_gradgrad'),
-               # TypeError: Cannot convert a MPS Tensor to float64 dtype as the MPS framework doesn't support float64
-               DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_numpy_ref', device_type='mps', dtypes=(torch.int64,)),
            )),
     OpInfo('unbind',
            dtypes=all_types_and_complex_and(torch.complex32, torch.bool, torch.float16, torch.bfloat16),
@@ -20857,8 +20778,6 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
                DecorateInfo(unittest.expectedFailure, "TestNormalizeOperators", "test_normalize_operator_exhaustive"),
                # Skip operator schema test because this is a functional and not an operator
                DecorateInfo(unittest.expectedFailure, 'TestOperatorSignatures', 'test_get_torch_func_signature_exhaustive'),
-               # RuntimeError: Failed to create function state object for: unfold_backward_float2
-               DecorateInfo(unittest.expectedFailure, 'TestCommon', device_type='mps', dtypes=(torch.complex64,)),
            ),
            sample_inputs_func=sample_inputs_unfold),
     OpInfo('unfold_copy',
@@ -20872,17 +20791,6 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
            check_batched_gradgrad=False,
            # See https://github.com/pytorch/pytorch/issues/66357
            check_batched_forward_grad=False,
-           skips=(
-               # RuntimeError: Failed to create function state object for: unfold_backward_float2
-               DecorateInfo(
-                   unittest.expectedFailure, 'TestCommon', 'test_variant_consistency_eager',
-                   device_type='mps', dtypes=(torch.complex64,)
-               ),
-               DecorateInfo(
-                   unittest.expectedFailure, 'TestCommon', 'test_noncontiguous_samples',
-                   device_type='mps', dtypes=(torch.complex64,)
-               ),
-           ),
            sample_inputs_func=sample_inputs_unfold),
     OpInfo('msort',
            dtypes=all_types_and(torch.bool, torch.float16, torch.bfloat16),
@@ -20917,7 +20825,7 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
                             device_type='cpu', dtypes=(torch.float16,)),
                DecorateInfo(toleranceOverride({torch.float16: tol(atol=3e-4, rtol=3e-6)}),
                             "TestConsistency", "test_output_match", device_type="mps"),
-               # RuntimeError: norm ops are not supported for complex yet
+               # RuntimeError: Failed to create function state object for: renorm_float2
                DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_dtypes', device_type='mps'),
                DecorateInfo(unittest.expectedFailure, 'TestCommon', device_type='mps', dtypes=(torch.complex64,)),
            )),
@@ -20933,9 +20841,6 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
                   sample_inputs_func=sample_repeat_tile,
                   skips=(
                       DecorateInfo(unittest.expectedFailure, "TestNormalizeOperators", "test_normalize_operator_exhaustive"),
-                      # Exception: repeat(): Not supported for complex yet!
-                      DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_dtypes', device_type='mps'),
-                      DecorateInfo(unittest.expectedFailure, 'TestCommon', device_type='mps', dtypes=(torch.complex64,)),
                   )),
     OpInfo('squeeze',
            ref=_squeeze_ref,
@@ -21049,9 +20954,6 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
            decorators=(
                # RuntimeError: view size is not compatible with input tensor's size and stride
                DecorateInfo(unittest.expectedFailure, "TestMeta", "test_dispatch_symbolic_meta_outplace_all_strides"),
-               # MPS: gather(): Yet not supported for complex
-               DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_dtypes', device_type='mps'),
-               DecorateInfo(unittest.expectedFailure, 'TestCommon', device_type='mps', dtypes=(torch.complex64,)),
            )),
     ShapeFuncInfo('tile',
                   ref=np.tile,
@@ -21061,11 +20963,6 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
                   supports_out=False,
                   supports_forward_ad=True,
                   supports_fwgrad_bwgrad=True,
-                  skips=(
-                      # RuntimeError: repeat(): Not supported for complex yet!
-                      DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_dtypes', device_type='mps'),
-                      DecorateInfo(unittest.expectedFailure, 'TestCommon', device_type='mps', dtypes=(torch.complex64,)),
-                  ),
                   sample_inputs_func=sample_repeat_tile),
     OpInfo('trapz',  # TODO: in the future, 'trapz' should be made a proper alias of 'trapezoid'
            dtypes=all_types_and_complex_and(torch.float16, torch.bfloat16),
@@ -21107,16 +21004,6 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
                    toleranceOverride({torch.float16: tol(atol=4e-3, rtol=4e-3)}),
                    'TestInductorOpInfo', 'test_comprehensive',
                ),
-               # RuntimeError: cumulative ops are not yet supported for complex
-               DecorateInfo(
-                   unittest.expectedFailure, 'TestCommon', 'test_variant_consistency_eager',
-                   device_type='mps', dtypes=(torch.complex64,)
-               ),
-               DecorateInfo(
-                   unittest.expectedFailure, 'TestCommon', 'test_noncontiguous_samples',
-                   device_type='mps', dtypes=(torch.complex64,)
-               ),
-               DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_dtypes', device_type='mps'),
            ),
            sample_inputs_func=sample_cumulative_trapezoid,),
     OpInfo('unsqueeze',
@@ -21190,17 +21077,13 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
            reference_inputs_func=reference_inputs_logsumexp),
     OpInfo('trace',
            dtypes=all_types_and_complex(),
+           dtypesIfMPS=all_types_and_complex_and(torch.bool, torch.half, torch.bfloat16),
            dtypesIfCUDA=all_types_and_complex_and(torch.chalf, torch.bool, torch.half, torch.bfloat16),
            error_inputs_func=error_inputs_trace,
            supports_inplace_autograd=False,
            supports_out=False,
            supports_forward_ad=True,
            supports_fwgrad_bwgrad=True,
-           skips=(
-               # RuntimeError: index_fill_(): Complex types are yet not supported
-               DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_dtypes', device_type='mps'),
-               DecorateInfo(unittest.expectedFailure, 'TestCommon', device_type='mps', dtypes=(torch.complex64,)),
-           ),
            sample_inputs_func=sample_inputs_trace),
     OpInfo('transpose',
            ref=_numpy_ref_transpose,
@@ -21421,8 +21304,6 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
                    "test_comprehensive",
                    device_type="cuda"
                ),
-               # RuntimeError: Failed to create function state object for: logcumsumexp_outer_float2
-               DecorateInfo(unittest.expectedFailure, 'TestCommon', device_type='mps', dtypes=(torch.complex64,)),
                # AssertionError: UserWarning not triggered : Resized a non-empty tensor but did not warn about it.
                DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out_warning', device_type='mps'),
            ),
@@ -21462,8 +21343,6 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
                    supports_fwgrad_bwgrad=True,
                    skips=(
                        # AssertionError: Tensor-likes are not close!
-                       DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out', device_type='mps', dtypes=(torch.float32,)),
-                       DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_noncontiguous_samples', device_type='mps'),
                        DecorateInfo(toleranceOverride({torch.float32: tol(atol=1e-5, rtol=3.1e-6)}),
                                     'TestConsistency', 'test_output_grad_match', device_type='mps'),
                    ),
@@ -21556,9 +21435,6 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
                 "test_variant_consistency_jit",
                 dtypes=(torch.float32,),
             ),
-            # Cannot convert a MPS Tensor to float64 dtype as the MPS framework doesn't support float64
-            DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_dtypes', device_type='mps'),
-            DecorateInfo(unittest.expectedFailure, 'TestCommon', device_type='mps', dtypes=(torch.complex64,)),
         ),
     ),
     UnaryUfuncInfo('lgamma',
@@ -21579,9 +21455,6 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
                                     dtypes=[torch.float32, torch.float64], active_if=IS_WINDOWS),
                        DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_large',
                                     dtypes=[torch.float32, torch.float64], active_if=IS_WINDOWS),
-                       # AssertionError: Tensor-likes are not close!
-                       DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out', device_type='mps', dtypes=(torch.float32,)),
-                       DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_noncontiguous_samples', device_type='mps'),
                    ),
                    # lgamma have multiple singularities at x <= 0
                    reference_numerics_filter=NumericsFilter(condition=lambda x: x < 0.1, safe_val=1)),
@@ -21688,7 +21561,6 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
            sample_inputs_func=sample_inputs_nonzero_static,
            supports_out=False,
            supports_autograd=False,
-           decorators=[onlyCPU],
            skips=(
                DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out'),
                DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out_warning'),
@@ -21879,15 +21751,6 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
         supports_forward_ad=True,
         supports_fwgrad_bwgrad=True,
         skips=(
-            # AssertionError: RuntimeError not raised : Expected RuntimeError when doing an unsafe cast
-            # from a result of dtype torch.float32 into an out= with dtype torch.long
-            DecorateInfo(
-                unittest.expectedFailure, 'TestCommon', 'test_out',
-                device_type='mps', dtypes=(torch.float32,)
-            ),
-            # Error: norm ops are not supported for complex yet
-            DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_dtypes', device_type='mps'),
-            DecorateInfo(unittest.expectedFailure, 'TestCommon', device_type='mps', dtypes=(torch.cfloat, torch.chalf)),
             # Dispatches in Python to vector_norm. Not sure how to make this test happy
             # Happens to pass on complex64. Also a mystery
             DecorateInfo(unittest.expectedFailure, 'TestJit', 'test_variant_consistency_jit',
@@ -21933,12 +21796,6 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
                    'test_output_match',
 
                ),
-               # RuntimeError: norm ops are not supported for complex yet
-               DecorateInfo(unittest.expectedFailure, 'TestCommon', device_type='mps', dtypes=(torch.complex64,)),
-               DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_dtypes', device_type='mps'),
-               # AssertionError: RuntimeError not raised : Expected RuntimeError when doing an
-               # unsafe cast from a result of dtype torch.float32 into an out= with dtype torch.long
-               DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out', device_type='mps', dtypes=(torch.float32,)),
                # Issue with conj and torch dispatch, see https://github.com/pytorch/pytorch/issues/82479
                DecorateInfo(
                    unittest.skip("Skipped!"),
@@ -21961,18 +21818,10 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
         # fast gradcheck produces NaNs
         gradcheck_fast_mode=False,
         skips=(
-            # RuntimeError: norm ops are not supported for complex yet
-            DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_dtypes', device_type='mps'),
-            # AssertionError: RuntimeError not raised : Expected RuntimeError
-            # when doing an unsafe cast from a result of dtype torch.float32
-            # into an out= with dtype torch.long
-            DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out', device_type='mps', dtypes=(torch.float32,)),
             DecorateInfo(
                 toleranceOverride({torch.float16: tol(atol=2e-3, rtol=1e-3)}),
                 'TestInductorOpInfo', 'test_comprehensive', device_type='cuda',
             ),
-            # Error: norm ops are not supported for complex yet
-            DecorateInfo(unittest.expectedFailure, 'TestCommon', device_type='mps', dtypes=(torch.cfloat, torch.chalf)),
             # Dispatches in Python to vector_norm. Not sure how to make this test happy
             # Happens to pass on complex64. Also a mystery
             DecorateInfo(unittest.expectedFailure, 'TestJit', 'test_variant_consistency_jit',
@@ -22309,7 +22158,7 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
         supports_out=False,
         sample_inputs_func=sample_inputs_grid_sample,
         reference_inputs_func=reference_inputs_grid_sample,
-        supports_gradgrad=False,
+        supports_gradgrad=True,
         skips=(
             # Exception: The operator 'aten::grid_sampler_2d_backward' is not currently implemented for the MPS device
             DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_dtypes', device_type='mps'),
@@ -22323,7 +22172,7 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
         dtypes=floating_types_and(torch.float16, torch.bfloat16),
         supports_out=False,
         sample_inputs_func=sample_inputs_grid_sampler_2d,
-        supports_gradgrad=False,
+        supports_gradgrad=True,
         gradcheck_nondet_tol=1e-15,
         skips=(
             DecorateInfo(slowTest, 'TestDecomp', 'test_comprehensive', dtypes=(torch.float32, torch.float64),
@@ -22340,7 +22189,7 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
         dtypes=floating_types_and(torch.float16, torch.bfloat16),
         supports_out=False,
         sample_inputs_func=sample_inputs_grid_sampler_3d,
-        supports_gradgrad=False,
+        supports_gradgrad=True,
         gradcheck_nondet_tol=1e-15,
         skips=(
             # NOTE: Only run on MPS
@@ -22351,7 +22200,6 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
             # Error: The operator 'aten::grid_sampler_3d_backward' is not currently implemented for the MPS device.
             DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_dtypes', device_type='mps'),
             DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_variant_consistency_eager', device_type='mps'),
-            # AssertionError: Tensor-likes are not close!
             DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_noncontiguous_samples', device_type='mps'),
             DecorateInfo(toleranceOverride({torch.float32: tol(atol=1e-4, rtol=1e-4),
                                             torch.float16: tol(atol=1e-4, rtol=1e-4),
@@ -22534,6 +22382,10 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
                          dtypes=[torch.float16]),
             DecorateInfo(unittest.skip("Skipped!"), 'TestReductions', 'test_ref_extremal_values',
                          device_type='cuda', dtypes=[torch.complex64]),
+            # Skipped on XPU because complex mean with extremal values (Inf/NaN) exhibits backend-dependent
+            # IEEE-754 behavior that differs from the CPU reference.
+            DecorateInfo(unittest.skip("Skipped!"), 'TestReductions', 'test_ref_extremal_values',
+                         device_type='xpu', dtypes=[torch.complex64]),
             # Driver issue of XPU, see https://github.com/intel/torch-xpu-ops/issues/2295
             DecorateInfo(
                 unittest.skip('Skipped!'),
@@ -22759,8 +22611,6 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
                 unittest.skip('Skipped!'), 'TestReductions', 'test_ref_small_input',
                 device_type='xpu',
                 dtypes=[torch.complex128, torch.int8, torch.int16, torch.int32, torch.int64]),
-            # RuntimeError: cumulative ops are not yet supported for complex
-            DecorateInfo(unittest.expectedFailure, 'TestCommon', device_type='mps', dtypes=(torch.complex64,)),
         ),
     ),
     ReductionOpInfo(
@@ -23070,11 +22920,6 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
                 "test_variant_consistency_jit",
                 dtypes=(torch.float32, torch.complex64),
             ),
-            # Cannot convert a MPS Tensor to float64 dtype as the MPS framework doesn't support float64
-            # RuntimeError: norm ops are not supported for complex yet
-            DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_dtypes', device_type='mps'),
-            DecorateInfo(unittest.expectedFailure, 'TestCommon', device_type='mps', dtypes=(torch.complex64,)),
-            DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_numpy_ref', device_type='mps', dtypes=(torch.int64,)),
         ),
     ),
     OpInfo(
@@ -23285,15 +23130,6 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
             DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs',
                          'test_reference_numerics_small',
                          dtypes=[torch.int8], active_if=TEST_WITH_ASAN),
-            # TypeError: Trying to convert ComplexDouble to the MPS backend but it does not have support for that dtype.
-            DecorateInfo(
-                unittest.expectedFailure, 'TestCommon', 'test_python_ref',
-                device_type='mps', dtypes=(torch.complex32,)
-            ),
-            DecorateInfo(
-                unittest.expectedFailure, 'TestCommon', 'test_python_ref_torch_fallback',
-                device_type='mps', dtypes=(torch.complex32,)
-            ),
         ),
     ),
     ElementwiseUnaryPythonRefInfo(
@@ -23408,16 +23244,17 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
         "_refs.lerp",
         torch_opinfo_name="lerp",
         skips=(
-            # TypeError: Cannot convert a MPS Tensor to float64 dtype
+            # Exception: Dtypes torch.float32 and * are not equal!
             DecorateInfo(
                 unittest.expectedFailure, 'TestCommon', 'test_python_ref', device_type='mps',
-                dtypes=(torch.bfloat16, torch.bool, torch.float16, torch.int16, torch.int32, torch.int64, torch.int8, torch.uint8)
+                dtypes=(torch.bool, torch.int16, torch.int32, torch.int64, torch.int8, torch.uint8)
             ),
-            DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_python_ref_meta', device_type='mps', dtypes=(torch.bool,)),
             DecorateInfo(
                 unittest.expectedFailure, 'TestCommon', 'test_python_ref_torch_fallback', device_type='mps',
-                dtypes=(torch.bfloat16, torch.bool, torch.float16, torch.int16, torch.int32, torch.int64, torch.int8, torch.uint8)
+                dtypes=(torch.bool, torch.int16, torch.int32, torch.int64, torch.int8, torch.uint8)
             ),
+            # RuntimeError: Failed to create function state object for: abs_dense_bool_bool
+            DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_python_ref_meta', device_type='mps', dtypes=(torch.bool,)),
         ),
     ),
     PythonRefInfo(
@@ -23649,15 +23486,6 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
             DecorateInfo(unittest.expectedFailure, 'TestMathBits', 'test_neg_view'),
             DecorateInfo(unittest.expectedFailure, 'TestMathBits', 'test_conj_view'),
             DecorateInfo(unittest.expectedFailure, 'TestMathBits', 'test_neg_conj_view'),
-            # TypeError: Cannot convert a MPS Tensor to float64 dtype as the MPS framework doesn't support float64
-            DecorateInfo(
-                unittest.expectedFailure, 'TestCommon', 'test_python_ref', device_type='mps',
-                dtypes=(torch.bfloat16, torch.float16,)
-            ),
-            DecorateInfo(
-                unittest.expectedFailure, 'TestCommon', 'test_python_ref_torch_fallback', device_type='mps',
-                dtypes=(torch.bfloat16, torch.float16,)
-            ),
             # RuntimeError: value cannot be converted to type uint8_t without overflow
             DecorateInfo(
                 unittest.expectedFailure, 'TestCommon', 'test_python_ref', device_type='mps',
@@ -23741,12 +23569,13 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
             DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_python_ref_executor',
                          dtypes=(torch.uint8, torch.int8, torch.int16, torch.int32, torch.int64),
                          device_type="cuda"),
-            # Cannot convert a MPS Tensor to float64 dtype as the MPS framework doesn't support float64
-            DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out', device_type='mps'),
-            DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out_warning', device_type='mps'),
+            # RuntimeError: no _refs support for aten.copy.default
             DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_python_ref', device_type='mps'),
-            DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_python_ref_meta', device_type='mps'),
-            DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_python_ref_torch_fallback', device_type='mps'),
+            # AssertionError: Tensor-likes are not equal!
+            DecorateInfo(
+                unittest.expectedFailure, 'TestCommon', 'test_python_ref_torch_fallback', device_type='mps',
+                dtypes=(torch.bool, torch.int16, torch.int32, torch.int64, torch.int8, torch.uint8)
+            ),
         ),
     ),
     PythonRefInfo(
@@ -23832,9 +23661,6 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
             DecorateInfo(unittest.expectedFailure,
                          'TestCommon',
                          'test_python_ref'),
-            # MPS: gather(): Yet not supported for complex
-            DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_dtypes', device_type='mps'),
-            DecorateInfo(unittest.expectedFailure, 'TestCommon', device_type='mps', dtypes=(torch.complex64,)),
         ),
     ),
     PythonRefInfo(
@@ -23939,6 +23765,10 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
                          'test_reference_numerics_large',
                          device_type='cuda', dtypes=[torch.cfloat, torch.cdouble],
                          active_if=IS_WINDOWS),
+            DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs',
+                         'test_reference_numerics_extremal',
+                         device_type='xpu', dtypes=[torch.cfloat, torch.cdouble],
+                         active_if=IS_WINDOWS),
         ),
     ),
     ElementwiseUnaryPythonRefInfo(
@@ -23963,6 +23793,10 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
                          'test_reference_numerics_large',
                          device_type='cuda', dtypes=[torch.cfloat],
                          active_if=IS_WINDOWS),
+            DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs',
+                         'test_reference_numerics_extremal',
+                         device_type='xpu', dtypes=[torch.cfloat, torch.cdouble],
+                         active_if=IS_WINDOWS),
         ),
     ),
     ElementwiseUnaryPythonRefInfo(
@@ -24015,6 +23849,10 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
                          'test_reference_numerics_large',
                          device_type='cuda',
                          dtypes=(torch.chalf,), active_if=IS_WINDOWS),
+            DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs',
+                         'test_reference_numerics_large',
+                         device_type='xpu',
+                         dtypes=(torch.chalf,), active_if=IS_WINDOWS),
         ),
     ),
     ElementwiseUnaryPythonRefInfo(
@@ -24049,27 +23887,15 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
                          'test_reference_numerics_large',
                          device_type='cuda',
                          dtypes=(torch.chalf,), active_if=IS_WINDOWS),
+            DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs',
+                         'test_reference_numerics_large',
+                         device_type='xpu',
+                         dtypes=(torch.chalf,), active_if=IS_WINDOWS),
         ),
     ),
     ElementwiseUnaryPythonRefInfo(
         "_refs.digamma",
         torch_opinfo_name="digamma",
-        skips=(
-            # TypeError: Cannot convert a MPS Tensor to float64 dtype as the MPS framework doesn't support float64
-            # AssertionError: Tensor-likes are not close!
-            DecorateInfo(
-                unittest.expectedFailure, 'TestCommon', 'test_python_ref', device_type='mps', dtypes=(
-                    torch.uint8, torch.int8, torch.int64, torch.int32,
-                    torch.int16, torch.float16, torch.bool, torch.bfloat16,
-                )
-            ),
-            DecorateInfo(
-                unittest.expectedFailure, 'TestCommon', 'test_python_ref_torch_fallback', device_type='mps', dtypes=(
-                    torch.uint8, torch.int8, torch.int64, torch.int32,
-                    torch.int16, torch.float16, torch.bool, torch.bfloat16,
-                )
-            ),
-        ),
     ),
     ElementwiseUnaryPythonRefInfo(
         "_refs.erf",
@@ -24109,6 +23935,10 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
             DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs',
                          'test_reference_numerics_extremal',
                          device_type='cpu', dtypes=[torch.cfloat, torch.cdouble], active_if=IS_WINDOWS),
+            DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs',
+                         'test_reference_numerics_large',
+                         device_type='xpu', dtypes=[torch.chalf],
+                         active_if=IS_WINDOWS),
         ),
     ),
     ElementwiseUnaryPythonRefInfo(
@@ -24126,6 +23956,10 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
             DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs',
                          'test_reference_numerics_extremal',
                          device_type='cpu', dtypes=[torch.cfloat, torch.cdouble], active_if=IS_WINDOWS),
+            DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs',
+                         'test_reference_numerics_large',
+                         device_type='xpu',
+                         dtypes=(torch.chalf,), active_if=IS_WINDOWS),
         ),
     ),
     ElementwiseUnaryPythonRefInfo(
@@ -24219,20 +24053,6 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
             DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs',
                          'test_reference_numerics_large',
                          dtypes=[torch.float32, torch.float64], active_if=IS_WINDOWS),
-            # TypeError: Cannot convert a MPS Tensor to float64 dtype as the MPS framework doesn't support float64
-            # AssertionError: Tensor-likes are not close!
-            DecorateInfo(
-                unittest.expectedFailure, 'TestCommon', 'test_python_ref', device_type='mps', dtypes=(
-                    torch.uint8, torch.int8, torch.int64, torch.int32,
-                    torch.int16, torch.float16, torch.bool, torch.bfloat16,
-                )
-            ),
-            DecorateInfo(
-                unittest.expectedFailure, 'TestCommon', 'test_python_ref_torch_fallback', device_type='mps', dtypes=(
-                    torch.uint8, torch.int8, torch.int64, torch.int32,
-                    torch.int16, torch.float16, torch.bool, torch.bfloat16,
-                )
-            ),
         ),
     ),
     ElementwiseUnaryPythonRefInfo(
@@ -24251,33 +24071,13 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
         "_refs.special.multigammaln",
         torch_opinfo_name="mvlgamma",
         torch_opinfo_variant_name="mvlgamma_p_3",
-        skips=skips_mvlgamma() + (
-            # TypeError: Cannot convert a MPS Tensor to float64 dtype as the MPS framework doesn't support float64
-            DecorateInfo(
-                unittest.expectedFailure, 'TestCommon', 'test_python_ref',
-                device_type='mps', dtypes=(torch.bfloat16, torch.float16)
-            ),
-            DecorateInfo(
-                unittest.expectedFailure, 'TestCommon', 'test_python_ref_torch_fallback',
-                device_type='mps', dtypes=(torch.bfloat16, torch.float16)
-            ),
-        ),
+        skips=skips_mvlgamma(),
     ),
     ElementwiseUnaryPythonRefInfo(
         "_refs.special.multigammaln",
         torch_opinfo_name="mvlgamma",
         torch_opinfo_variant_name="mvlgamma_p_5",
-        skips=skips_mvlgamma() + (
-            # TypeError: Cannot convert a MPS Tensor to float64 dtype as the MPS framework doesn't support float64
-            DecorateInfo(
-                unittest.expectedFailure, 'TestCommon', 'test_python_ref',
-                device_type='mps', dtypes=(torch.bfloat16, torch.float16)
-            ),
-            DecorateInfo(
-                unittest.expectedFailure, 'TestCommon', 'test_python_ref_torch_fallback',
-                device_type='mps', dtypes=(torch.bfloat16, torch.float16)
-            ),
-        ),
+        skips=skips_mvlgamma(),
     ),
     ElementwiseUnaryPythonRefInfo(
         "_refs.log",
@@ -24326,14 +24126,15 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
         torch_opinfo_name="log_softmax",
         torch_opinfo_variant_name="with_dtype",
         skips=(
-            # TypeError: Cannot convert a MPS Tensor to float64 dtype as the MPS framework doesn't support float64
+            # AssertionError: Tensor-likes are not close!
+            # RuntimeError: softmax only supported for floating types
             DecorateInfo(
                 unittest.expectedFailure, 'TestCommon', 'test_python_ref', device_type='mps',
-                dtypes=(torch.float32, torch.float16, torch.complex64, torch.complex32)
+                dtypes=(torch.float32, torch.complex64, torch.complex32)
             ),
             DecorateInfo(
                 unittest.expectedFailure, 'TestCommon', 'test_python_ref_torch_fallback', device_type='mps',
-                dtypes=(torch.float32, torch.float16, torch.complex64, torch.complex32)
+                dtypes=(torch.float32, torch.complex64, torch.complex32)
             ),
         ),
     ),
@@ -24468,6 +24269,14 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
                          'test_reference_numerics_large',
                          dtypes=(torch.cfloat, torch.cdouble,), device_type='cpu',
                          active_if=IS_WINDOWS),
+            DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs',
+                         'test_reference_numerics_extremal',
+                         device_type='xpu', dtypes=[torch.cfloat, torch.cdouble],
+                         active_if=IS_WINDOWS),
+            DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs',
+                         'test_reference_numerics_large',
+                         device_type='xpu',
+                         dtypes=(torch.chalf,), active_if=IS_WINDOWS),
         ),
     ),
     ElementwiseUnaryPythonRefInfo(
@@ -24480,15 +24289,6 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
             DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs',
                          'test_reference_numerics_small',
                          dtypes=[torch.cfloat]),
-            # TypeError: Trying to convert ComplexDouble to the MPS backend but it does not have support for that dtype.
-            DecorateInfo(
-                unittest.expectedFailure, 'TestCommon', 'test_python_ref',
-                device_type='mps', dtypes=(torch.complex64,)
-            ),
-            DecorateInfo(
-                unittest.expectedFailure, 'TestCommon', 'test_python_ref_torch_fallback',
-                device_type='mps', dtypes=(torch.complex64,)
-            ),
         ),
     ),
     ElementwiseUnaryPythonRefInfo(
@@ -24511,6 +24311,14 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
             DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs',
                          'test_reference_numerics_large',
                          device_type='cpu', dtypes=[torch.int8]),
+            DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs',
+                         'test_reference_numerics_extremal',
+                         device_type='xpu', dtypes=[torch.cfloat, torch.cdouble],
+                         active_if=IS_WINDOWS),
+            DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs',
+                         'test_reference_numerics_large',
+                         device_type='xpu',
+                         dtypes=(torch.chalf,), active_if=IS_WINDOWS),
         ),
     ),
     PythonRefInfo(
@@ -24518,15 +24326,15 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
         torch_opinfo_name="softmax",
         torch_opinfo_variant_name="with_dtype",
         skips=(
-            # TypeError: Cannot convert a MPS Tensor to float64 dtype as the MPS framework doesn't support float64
-            # TypeError: Trying to convert ComplexDouble to the MPS backend but it does not have support for that dtype.
+            # AssertionError: Tensor-likes are not close!
+            # RuntimeError: softmax only supported for floating types
             DecorateInfo(
                 unittest.expectedFailure, 'TestCommon', 'test_python_ref', device_type='mps',
-                dtypes=(torch.float32, torch.float16, torch.complex64),
+                dtypes=(torch.float32, torch.complex64),
             ),
             DecorateInfo(
                 unittest.expectedFailure, 'TestCommon', 'test_python_ref_torch_fallback', device_type='mps',
-                dtypes=(torch.float32, torch.float16, torch.complex64),
+                dtypes=(torch.float32, torch.complex64),
             ),
         ),
     ),
@@ -24584,6 +24392,14 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
                          'test_reference_numerics_large',
                          device_type='cpu', dtypes=[torch.cfloat, torch.cdouble],
                          active_if=(IS_MACOS or IS_WINDOWS)),
+            DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs',
+                         'test_reference_numerics_extremal',
+                         device_type='xpu', dtypes=[torch.cfloat, torch.cdouble],
+                         active_if=IS_WINDOWS),
+            DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs',
+                         'test_reference_numerics_large',
+                         device_type='xpu',
+                         dtypes=(torch.chalf,), active_if=IS_WINDOWS),
         )
     ),
     ElementwiseUnaryPythonRefInfo(
@@ -24603,6 +24419,14 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
                          'test_reference_numerics_large',
                          device_type='cpu', dtypes=[torch.cfloat, torch.cdouble],
                          active_if=(IS_MACOS or IS_WINDOWS)),
+            DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs',
+                         'test_reference_numerics_extremal',
+                         device_type='xpu', dtypes=[torch.cfloat, torch.cdouble],
+                         active_if=IS_WINDOWS),
+            DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs',
+                         'test_reference_numerics_large',
+                         device_type='xpu',
+                         dtypes=(torch.chalf,), active_if=IS_WINDOWS),
         ),
     ),
     ElementwiseUnaryPythonRefInfo(
@@ -24617,10 +24441,11 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
         torch_opinfo_variant_name="with_dtype",
         supports_out=False,
         skips=(
-            # TypeError: Cannot convert a MPS Tensor to float64 dtype as the MPS framework doesn't support float64
+            # AssertionError: Tensor-likes are not close!
+            # NotImplementedError: log_softmax for complex is not supported for MPS
             DecorateInfo(
                 unittest.expectedFailure, 'TestCommon', 'test_python_ref', device_type='mps',
-                dtypes=(torch.float32, torch.float16, torch.complex64, torch.complex32)
+                dtypes=(torch.float32, torch.complex64, torch.complex32)
             ),
         ),
     ),
@@ -24630,11 +24455,11 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
         torch_opinfo_variant_name="with_dtype",
         supports_out=False,
         skips=(
-            # TypeError: Cannot convert a MPS Tensor to float64 dtype as the MPS framework doesn't support float64
-            # TypeError: Trying to convert ComplexDouble to the MPS backend but it does not have support for that dtype.
+            # Exception: softmax only supported for floating types
+            # AssertionError: Tensor-likes are not close!
             DecorateInfo(
                 unittest.expectedFailure, 'TestCommon', 'test_python_ref', device_type='mps',
-                dtypes=(torch.float32, torch.float16, torch.complex64),
+                dtypes=(torch.float32, torch.complex64),
             ),
         ),
     ),
@@ -24644,17 +24469,6 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
     ElementwiseUnaryPythonRefInfo(
         "_refs.special.logit",
         torch_opinfo_name="logit",
-        skips=(
-            # TypeError: Cannot convert a MPS Tensor to float64 dtype as the MPS framework doesn't support float64
-            DecorateInfo(
-                unittest.expectedFailure, 'TestCommon', 'test_python_ref',
-                device_type='mps', dtypes=(torch.complex64, torch.float16)
-            ),
-            DecorateInfo(
-                unittest.expectedFailure, 'TestCommon', 'test_python_ref_torch_fallback',
-                device_type='mps', dtypes=(torch.complex64, torch.float16)
-            ),
-        ),
     ),
     #
     # Elementwise Unary nn.functional OpInfos
@@ -24793,12 +24607,6 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
         "_refs.nn.functional.pairwise_distance",
         torch_opinfo_name="nn.functional.pairwise_distance",
         supports_out=True,
-        skips=(
-            # Cannot convert a MPS Tensor to float64 dtype as the MPS framework doesn't support float64
-            # RuntimeError: norm ops are not supported for complex yet
-            DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_dtypes', device_type='mps'),
-            DecorateInfo(unittest.expectedFailure, 'TestCommon', device_type='mps', dtypes=(torch.complex64,)),
-        ),
     ),
     PythonRefInfo(
         "_refs.nn.functional.pdist",
@@ -24831,10 +24639,11 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
         torch_opinfo_variant_name="with_dtype",
         supports_out=False,
         skips=(
-            # TypeError: Cannot convert a MPS Tensor to float64 dtype as the MPS framework doesn't support float64
+            # NotImplementedError: log_softmax for complex is not supported for MPS
+            # AssertionError: Tensor-likes are not close!
             DecorateInfo(
                 unittest.expectedFailure, 'TestCommon', 'test_python_ref', device_type='mps',
-                dtypes=(torch.float32, torch.float16, torch.complex64, torch.complex32)
+                dtypes=(torch.float32, torch.complex64, torch.complex32)
             ),
         ),
     ),
@@ -24849,17 +24658,6 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
     PythonRefInfo(
         "_refs.nn.functional.poisson_nll_loss",
         torch_opinfo_name="nn.functional.poisson_nll_loss",
-        skips=(
-            # TypeError: Cannot convert a MPS Tensor to float64 dtype as the MPS framework doesn't support float64
-            DecorateInfo(
-                unittest.expectedFailure, 'TestCommon', 'test_python_ref',
-                device_type='mps', dtypes=(torch.float16, torch.bfloat16)
-            ),
-            DecorateInfo(
-                unittest.expectedFailure, 'TestCommon', 'test_python_ref_torch_fallback',
-                device_type='mps', dtypes=(torch.float16, torch.bfloat16)
-            ),
-        ),
     ),
     ElementwiseUnaryPythonRefInfo(
         "_refs.nn.functional.prelu",
@@ -24934,11 +24732,11 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
         torch_opinfo_variant_name="with_dtype",
         supports_out=False,
         skips=(
-            # TypeError: Cannot convert a MPS Tensor to float64 dtype as the MPS framework doesn't support float64
-            # TypeError: Trying to convert ComplexDouble to the MPS backend but it does not have support for that dtype.
+            # RuntimeError: softmax only supported for floating types
+            # AssertionError: Tensor-likes are not close!
             DecorateInfo(
                 unittest.expectedFailure, 'TestCommon', 'test_python_ref', device_type='mps',
-                dtypes=(torch.float32, torch.float16, torch.complex64),
+                dtypes=(torch.float32, torch.complex64),
             ),
         ),
     ),
@@ -24948,11 +24746,11 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
         torch_opinfo_variant_name="with_dtype",
         supports_out=False,
         skips=(
-            # TypeError: Trying to convert ComplexDouble to the MPS backend but it does not have support for that dtype.
-            # TypeError: Cannot convert a MPS Tensor to float64 dtype as the MPS framework doesn't support float64
+            # RuntimeError: softmax only supported for floating types
+            # AssertionError: Tensor-likes are not close!
             DecorateInfo(
                 unittest.expectedFailure, 'TestCommon', 'test_python_ref', device_type='mps',
-                dtypes=(torch.complex64, torch.float16, torch.float32),
+                dtypes=(torch.complex64, torch.float32),
             ),
         ),
     ),
@@ -24980,11 +24778,6 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
     PythonRefInfo(
         "_refs.nn.functional.l1_loss",
         torch_opinfo_name="nn.functional.l1_loss",
-        skips=(
-            # Cannot convert a MPS Tensor to float64 dtype as the MPS framework doesn't support float64
-            DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_dtypes', device_type='mps'),
-            DecorateInfo(unittest.expectedFailure, 'TestCommon', device_type='mps', dtypes=(torch.complex64,)),
-        ),
     ),
     PythonRefInfo(
         "_refs.nn.functional.margin_ranking_loss",
@@ -25001,13 +24794,6 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
     PythonRefInfo(
         "_refs.nn.functional.hinge_embedding_loss",
         torch_opinfo_name="nn.functional.hinge_embedding_loss",
-        skips=(
-            # TypeError: Cannot convert a MPS Tensor to float64 dtype as the MPS framework doesn't support float64
-            DecorateInfo(
-                unittest.expectedFailure, 'TestCommon', 'test_python_ref', device_type='mps',
-                dtypes=(torch.bfloat16, torch.float16,),
-            ),
-        ),
     ),
     PythonRefInfo(
         "_refs.nn.functional.nll_loss",
@@ -25088,15 +24874,10 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
                          dtypes=(torch.complex64, torch.complex128),
                          device_type='cpu',
                          active_if=(IS_MACOS or IS_WINDOWS)),
-            # TypeError: Cannot convert a MPS Tensor to float64 dtype as the MPS framework doesn't support float64
-            DecorateInfo(
-                unittest.expectedFailure, 'TestCommon', 'test_python_ref',
-                device_type='mps', dtypes=(torch.float16, torch.bfloat16)
-            ),
-            DecorateInfo(
-                unittest.expectedFailure, 'TestCommon', 'test_python_ref_torch_fallback',
-                device_type='mps', dtypes=(torch.float16, torch.bfloat16)
-            ),
+            DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs',
+                         'test_reference_numerics_extremal',
+                         device_type='xpu', dtypes=[torch.cfloat, torch.cdouble],
+                         active_if=IS_WINDOWS),
         ),
     ),
     ElementwiseUnaryPythonRefInfo(
@@ -25137,15 +24918,6 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
                          'TestBinaryUfuncs',
                          'test_reference_numerics_extremal_values',
                          dtypes=(torch.complex64, torch.complex128)),
-            # TypeError: Cannot convert a MPS Tensor to float64 dtype as the MPS framework doesn't support float64
-            DecorateInfo(
-                unittest.expectedFailure, 'TestCommon', 'test_python_ref', device_type='mps',
-                dtypes=(torch.bfloat16,)
-            ),
-            DecorateInfo(
-                unittest.expectedFailure, 'TestCommon', 'test_python_ref_torch_fallback', device_type='mps',
-                dtypes=(torch.float16,)
-            ),
         ),
     ),
     ElementwiseBinaryPythonRefInfo(
@@ -25273,15 +25045,6 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
         decorators=(
             # See https://github.com/pytorch/pytorch/issues/111126
             DecorateInfo(unittest.expectedFailure, 'TestBinaryUfuncs', 'test_type_promotion'),
-            # TypeError: Cannot convert a MPS Tensor to float64 dtype as the MPS framework doesn't support float64
-            DecorateInfo(
-                unittest.expectedFailure, 'TestCommon', 'test_python_ref',
-                device_type='mps', dtypes=(torch.bfloat16, torch.float16)
-            ),
-            DecorateInfo(
-                unittest.expectedFailure, 'TestCommon', 'test_python_ref_torch_fallback',
-                device_type='mps', dtypes=(torch.bfloat16, torch.float16)
-            ),
         ),
     ),
     ElementwiseBinaryPythonRefInfo(
@@ -25294,16 +25057,6 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
         decorators=(
             # See https://github.com/pytorch/pytorch/issues/111126
             DecorateInfo(unittest.expectedFailure, 'TestBinaryUfuncs', 'test_type_promotion'),
-            # Reference result was farther (nan) from the precise computation than the
-            # torch result was (inf)!
-            DecorateInfo(
-                unittest.expectedFailure,
-                "TestCommon",
-                "test_python_ref",
-                dtypes=(torch.bfloat16,),
-                device_type="cpu",
-                active_if=not IS_S390X,
-            ),
         ),
     ),
     ElementwiseBinaryPythonRefInfo(
@@ -25497,17 +25250,6 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
     ElementwiseUnaryPythonRefInfo(
         "_refs.logical_not",
         torch_opinfo_name="logical_not",
-        skips=(
-            # RuntimeError: Undefined type ComplexDouble
-            DecorateInfo(
-                unittest.expectedFailure, 'TestCommon', 'test_python_ref_torch_fallback',
-                device_type='mps', dtypes=(torch.complex64,)
-            ),
-            DecorateInfo(
-                unittest.expectedFailure, 'TestCommon', 'test_python_ref',
-                device_type='mps', dtypes=(torch.complex64,)
-            ),
-        ),
     ),
     ElementwiseBinaryPythonRefInfo(
         "_refs.logical_or",
@@ -25702,16 +25444,6 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
                          'TestBinaryUfuncs',
                          'test_reference_numerics_small_values',
                          dtypes=(torch.uint8,)),
-            # TypeError: Cannot convert a MPS Tensor to float64 dtype as the MPS framework doesn't support float64
-            # NotImplementedError: "_local_scalar_dense_mps" not implemented for 'ComplexHalf'
-            DecorateInfo(
-                unittest.expectedFailure, 'TestCommon', 'test_python_ref_torch_fallback',
-                device_type='mps', dtypes=(torch.float16,)
-            ),
-            DecorateInfo(
-                unittest.expectedFailure, 'TestCommon', 'test_python_ref',
-                device_type='mps', dtypes=(torch.float16, torch.bfloat16)
-            ),
         ),
     ),
     ElementwiseBinaryPythonRefInfo(
@@ -25747,17 +25479,6 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
     PythonRefInfo(
         "_refs.addcdiv",
         torch_opinfo_name="addcdiv",
-        skips=(
-            # TypeError: Cannot convert a MPS Tensor to float64 dtype as the MPS framework doesn't support float64
-            DecorateInfo(
-                unittest.expectedFailure, 'TestCommon', 'test_python_ref', device_type='mps',
-                dtypes=(torch.bfloat16, torch.float16,)
-            ),
-            DecorateInfo(
-                unittest.expectedFailure, 'TestCommon', 'test_python_ref_torch_fallback', device_type='mps',
-                dtypes=(torch.bfloat16, torch.float16,)
-            ),
-        ),
     ),
     PythonRefInfo(
         "_refs.addcmul",
@@ -25771,18 +25492,17 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
                          dtypes=(torch.float16,), device_type="cpu"),
             DecorateInfo(unittest.skip("Skipped!"), 'TestCommon', 'test_python_ref_torch_fallback',
                          dtypes=(torch.float16,), device_type="cpu"),
-            # TypeError: Cannot convert a MPS Tensor to float64 dtype as the MPS framework doesn't support float64
             # AssertionError: Tensor-likes are not close!
             DecorateInfo(
                 unittest.expectedFailure, 'TestCommon', 'test_python_ref', device_type='mps', dtypes=(
                     torch.uint8, torch.int8, torch.int64, torch.int32,
-                    torch.int16, torch.float16, torch.bfloat16,
+                    torch.int16,
                 )
             ),
             DecorateInfo(
                 unittest.expectedFailure, 'TestCommon', 'test_python_ref_torch_fallback', device_type='mps', dtypes=(
                     torch.uint8, torch.int8, torch.int64, torch.int32,
-                    torch.int16, torch.float16, torch.bfloat16,
+                    torch.int16,
                 )
             ),
         ),
@@ -25793,11 +25513,6 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
         skips=(
             # test error disabled since rhs non-tensor python scalar is supported
             DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_python_ref_errors'),
-            # TypeError: Cannot convert a MPS Tensor to float64 dtype as the MPS framework doesn't support float64
-            DecorateInfo(
-                unittest.expectedFailure, 'TestCommon', 'test_python_ref',
-                device_type='mps', dtypes=(torch.float32, torch.float16,)
-            ),
         ),
     ),
     ElementwiseBinaryPythonRefInfo(
@@ -25806,11 +25521,6 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
         skips=(
             # test error disabled since rhs non-tensor python scalar is supported
             DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_python_ref_errors'),
-            # TypeError: Cannot convert a MPS Tensor to float64 dtype as the MPS framework doesn't support float64
-            DecorateInfo(
-                unittest.expectedFailure, 'TestCommon', 'test_python_ref',
-                device_type='mps', dtypes=(torch.float32, torch.float16,)
-            ),
         ),
     ),
     PythonRefInfo(
@@ -25828,9 +25538,6 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
             # Greatest relative difference: 8.519846983548175e-06 at index (4,) (up to 1.3e-06 allowed)
             DecorateInfo(unittest.skip("Skipped!"), 'TestCommon', 'test_python_ref',
                          dtypes=(torch.uint8,), device_type="cpu"),
-            # RuntimeError: norm ops are not supported for complex yet
-            DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_dtypes', device_type='mps'),
-            DecorateInfo(unittest.expectedFailure, 'TestCommon', device_type='mps', dtypes=(torch.complex64,)),
         )
     ),
     ElementwiseBinaryPythonRefInfo(
@@ -26116,10 +25823,15 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
         skips=(
             # FIXME: AssertionError: RuntimeError not raised
             DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_python_ref_errors'),
-            # TypeError: Cannot convert a MPS Tensor to float64 dtype as the MPS framework doesn't support float64
-            DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_python_ref', device_type='mps'),
-            DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_python_ref_meta', device_type='mps'),
-            DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_python_ref_torch_fallback', device_type='mps'),
+            # RuntimeError: Failed to create function state object for: cat_int32_t_*
+            DecorateInfo(
+                unittest.expectedFailure, 'TestCommon', 'test_python_ref', device_type='mps',
+                dtypes=(torch.complex64, torch.complex32)
+            ),
+            DecorateInfo(
+                unittest.expectedFailure, 'TestCommon', 'test_python_ref_torch_fallback', device_type='mps',
+                dtypes=(torch.complex64, torch.complex32)
+            ),
         ),
     ),
     PythonRefInfo(
@@ -26250,11 +25962,12 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
         skips=(
             # RuntimeError: mean(): could not infer output dtype. Input dtype must be either a floating point or complex dtype
             DecorateInfo(
-                unittest.expectedFailure, 'TestCommon', 'test_python_ref', device_type='mps', dtypes=(torch.int32, torch.int16)
+                unittest.expectedFailure, 'TestCommon', 'test_python_ref', device_type='mps',
+                dtypes=(torch.int32, torch.int16, torch.int8, torch.uint8)
             ),
             DecorateInfo(
                 unittest.expectedFailure, 'TestCommon', 'test_python_ref_meta', device_type='mps',
-                dtypes=(torch.int32, torch.int16)
+                dtypes=(torch.int32, torch.int16, torch.int8, torch.uint8)
             ),
         )
     ),
@@ -26302,23 +26015,7 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
         "_refs.renorm",
         torch_opinfo_name="renorm",
         skips=(
-            # RuntimeError: norm ops are not supported for complex yet
-            DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_dtypes', device_type='mps'),
-            DecorateInfo(unittest.expectedFailure, 'TestCommon', device_type='mps', dtypes=(torch.complex64,)),
-            # Cannot convert a MPS Tensor to float64 dtype as the MPS framework doesn't support float64
-            DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_python_ref', device_type='mps', dtypes=(torch.float16,)),
-            DecorateInfo(
-                unittest.expectedFailure, 'TestCommon', 'test_python_ref_torch_fallback',
-                device_type='mps', dtypes=(torch.float16,)
-            ),
-        ),
-    ),
-    PythonRefInfo(
-        "_refs.repeat",
-        torch_opinfo_name="repeat",
-        validate_view_consistency=False,
-        skips=(
-            # Exception: repeat(): Not supported for complex yet!
+            # RuntimeError: Failed to create function state object for: renorm_float2
             DecorateInfo(
                 unittest.expectedFailure, 'TestCommon', 'test_python_ref',
                 device_type='mps', dtypes=(torch.complex64,)
@@ -26329,6 +26026,11 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
             ),
         ),
     ),
+    PythonRefInfo(
+        "_refs.repeat",
+        torch_opinfo_name="repeat",
+        validate_view_consistency=False,
+    ),
     PythonRefInfo(
         "_refs.reshape",
         torch_opinfo_name="reshape",
@@ -26680,15 +26382,6 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
         skips=(
             # doesn't test out behavior properly for this operator
             DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out'),
-            # RuntimeError: cumulative ops are not yet supported for complex
-            DecorateInfo(
-                unittest.expectedFailure, 'TestCommon', 'test_python_ref', device_type='mps',
-                dtypes=(torch.complex64,)
-            ),
-            DecorateInfo(
-                unittest.expectedFailure, 'TestCommon', 'test_python_ref_torch_fallback', device_type='mps',
-                dtypes=(torch.complex64,)
-            ),
         ),
     ),
     PythonRefInfo(
@@ -26698,15 +26391,6 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
         skips=(
             # doesn't test out behavior properly for this operator
             DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out'),
-            # RuntimeError: cumulative ops are not yet supported for complex
-            DecorateInfo(
-                unittest.expectedFailure, 'TestCommon', 'test_python_ref', device_type='mps',
-                dtypes=(torch.complex64,)
-            ),
-            DecorateInfo(
-                unittest.expectedFailure, 'TestCommon', 'test_python_ref_torch_fallback', device_type='mps',
-                dtypes=(torch.complex64,)
-            ),
         ),
     ),
     PythonRefInfo(
@@ -26751,6 +26435,8 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
             # FIXME: improve precision
             DecorateInfo(
                 unittest.skip("Skipped!"), 'TestReductions', 'test_ref_small_input'),
+            DecorateInfo(
+                unittest.skip("Skipped!"), 'TestReductions', 'test_ref_duplicate_values'),
             # torch._subclasses.fake_tensor.MetadataMismatchError: Dtypes torch.float32 and torch.complex64 are not equal!
             DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_python_ref', device_type='mps', dtypes=(torch.complex64,)),
             DecorateInfo(
@@ -26797,14 +26483,11 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
         decorators=(
             DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_python_ref',),
             # RuntimeError: MPS device does not support addr for non-float input
-            # TypeError: Cannot convert a MPS Tensor to float64 dtype as the MPS
-            # framework doesn't support float64.
             DecorateInfo(
                 unittest.expectedFailure, 'TestCommon', 'test_python_ref_torch_fallback', device_type='mps',
                 dtypes=(
                     torch.uint8, torch.int8, torch.int64, torch.int32,
-                    torch.int16, torch.float16, torch.complex64, torch.bool,
-                    torch.bfloat16
+                    torch.int16, torch.complex64, torch.bool,
                 )
             ),
         ),
@@ -26820,19 +26503,6 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
         # Uses vector_norm inside and vector_norm is affected by
         # https://github.com/pytorch/pytorch/issues/77216
         validate_view_consistency=False,
-        skips=(
-            # TypeError: Cannot convert a MPS Tensor to float64 dtype as the MPS framework doesn't support float64
-            DecorateInfo(
-                unittest.expectedFailure, 'TestCommon', 'test_python_ref',
-                device_type='mps', dtypes=(torch.float16,)
-            ),
-            # RuntimeError: norm ops are not supported for complex yet
-            DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_dtypes', device_type='mps'),
-            DecorateInfo(
-                unittest.expectedFailure, 'TestCommon', device_type='mps',
-                dtypes=(torch.complex32, torch.complex64)
-            ),
-        ),
     ),
     #
     # Tensor Creation Reference OpInfos
@@ -27091,11 +26761,6 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
         torch_opinfo_name="index_fill",
         # empty_strided
         skips=(
-            # RuntimeError: index_fill_(): Complex types are yet not supported
-            DecorateInfo(
-                unittest.expectedFailure, 'TestCommon', 'test_python_ref_torch_fallback',
-                device_type='mps', dtypes=(torch.complex64, torch.complex32)
-            ),
             # no _refs support for Tensor.__setitem__
             DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_python_ref'),)
     ),
@@ -27143,15 +26808,6 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
                 'test_python_ref_executor',
                 dtypes=(torch.complex64, torch.complex128),
             ),
-            # RuntimeError: Failed to create function state object for: unfold_backward_float2
-            DecorateInfo(
-                unittest.expectedFailure, 'TestCommon', 'test_python_ref_meta',
-                device_type='mps', dtypes=(torch.complex64,)
-            ),
-            DecorateInfo(
-                unittest.expectedFailure, 'TestCommon', 'test_python_ref_torch_fallback',
-                device_type='mps', dtypes=(torch.complex64,)
-            ),
         ],
     ),
     PythonRefInfo(
@@ -27169,7 +26825,6 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
 ops_and_refs = op_db + python_ref_db
 unary_ufuncs = [op for op in ops_and_refs if isinstance(op, UnaryUfuncInfo)]
 binary_ufuncs = [op for op in ops_and_refs if isinstance(op, BinaryUfuncInfo)]
-binary_ufuncs_and_refs = tuple(op for op in ops_and_refs if isinstance(op, BinaryUfuncInfo))
 spectral_funcs = [op for op in ops_and_refs if isinstance(op, SpectralFuncInfo)]
 sparse_unary_ufuncs = [op for op in op_db if isinstance(op, UnaryUfuncInfo) and op.supports_sparse]
 sparse_csr_unary_ufuncs = [op for op in op_db if isinstance(op, UnaryUfuncInfo) and op.supports_sparse_csr]
diff --git a/torch/testing/_internal/common_modules.py b/torch/testing/_internal/common_modules.py
index f77491705f979..3e5d4af0d2a58 100644
--- a/torch/testing/_internal/common_modules.py
+++ b/torch/testing/_internal/common_modules.py
@@ -25,7 +25,7 @@
     nllloss_reference, nlllossNd_reference, smoothl1loss_reference, softmarginloss_reference, get_reduction)
 from torch.testing._internal.common_utils import (
     freeze_rng_state, skipIfMPS, GRADCHECK_NONDET_TOL, TEST_WITH_ROCM, IS_WINDOWS,
-    skipIfTorchDynamo)
+    skipIfTorchDynamo, skipIfXpu)
 from types import ModuleType
 import operator
 
@@ -913,7 +913,10 @@ def module_inputs_torch_nn_BatchNorm1d(module_info, device, dtype, requires_grad
                     desc='3d_input_not_affine'),
         ModuleInput(constructor_input=FunctionInput(5, 1e-3, 0.3, False),
                     forward_input=FunctionInput(make_input((0, 5, 9))),
-                    desc='zero_batch')]
+                    desc='zero_batch'),
+        ModuleInput(constructor_input=FunctionInput(10, bias=False),
+                    forward_input=FunctionInput(make_input((4, 10))),
+                    desc='affine_not_bias'),]
 
 
 def module_inputs_torch_nn_BatchNorm2d(module_info, device, dtype, requires_grad, training, **kwargs):
@@ -936,7 +939,10 @@ def module_inputs_torch_nn_BatchNorm2d(module_info, device, dtype, requires_grad
                     desc='not_tracking_stats'),
         ModuleInput(constructor_input=FunctionInput(5, 1e-3, 0.3, False),
                     forward_input=FunctionInput(make_input((0, 5, 2, 2))),
-                    desc='zero_batch')]
+                    desc='zero_batch'),
+        ModuleInput(constructor_input=FunctionInput(3, bias=False),
+                    forward_input=FunctionInput(make_input((2, 3, 6, 6))),
+                    desc='affine_not_bias'),]
 
 
 def module_inputs_torch_nn_BatchNorm3d(module_info, device, dtype, requires_grad, training, **kwargs):
@@ -959,7 +965,10 @@ def module_inputs_torch_nn_BatchNorm3d(module_info, device, dtype, requires_grad
                     desc='not_tracking_stats'),
         ModuleInput(constructor_input=FunctionInput(5, 1e-3, 0.3, False),
                     forward_input=FunctionInput(make_input((0, 5, 2, 2, 2))),
-                    desc='zero_batch')]
+                    desc='zero_batch'),
+        ModuleInput(constructor_input=FunctionInput(3, bias=False),
+                    forward_input=FunctionInput(make_input((2, 3, 4, 4, 4))),
+                    desc='affine_not_bias'),]
 
 
 def module_error_inputs_torch_nn_BatchNorm1d_2d_3d(module_info, device, dtype, requires_grad, training, **kwargs):
@@ -1666,6 +1675,104 @@ def reference_fn(m, p, i, t, reduction=reduction, constructor_kwargs=constructor
     return module_inputs
 
 
+def module_error_inputs_torch_nn_CrossEntropyLoss(module_info, device, dtype, requires_grad, training, **kwargs):
+    make_input = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad)
+    make_weight = partial(make_tensor, device=device, dtype=dtype, requires_grad=False)
+
+    return [
+        # Non-floating-point target with same shape as input (soft-target path)
+        ErrorModuleInput(
+            ModuleInput(
+                constructor_input=FunctionInput(),
+                forward_input=FunctionInput(
+                    make_input((2, 3)),
+                    torch.zeros((2, 3), device=device, dtype=torch.long),
+                ),
+            ),
+            error_on=ModuleErrorEnum.FORWARD_ERROR,
+            error_type=RuntimeError,
+            error_regex="Expected floating point type for target with class probabilities",
+        ),
+        # Non-long target with different shape from input (class-index path)
+        ErrorModuleInput(
+            ModuleInput(
+                constructor_input=FunctionInput(),
+                forward_input=FunctionInput(
+                    make_input((2, 3)),
+                    make_input((2,)),
+                ),
+            ),
+            error_on=ModuleErrorEnum.FORWARD_ERROR,
+            error_type=RuntimeError,
+            error_regex="expected target dtype to be Long or Byte",
+        ),
+        # Batch size mismatch
+        ErrorModuleInput(
+            ModuleInput(
+                constructor_input=FunctionInput(),
+                forward_input=FunctionInput(
+                    make_input((2, 3)),
+                    torch.tensor([0, 1, 2], device=device, dtype=torch.long),
+                ),
+            ),
+            error_on=ModuleErrorEnum.FORWARD_ERROR,
+            error_type=ValueError,
+            error_regex="Expected input batch_size .* to match target batch_size",
+        ),
+        # Weight wrong shape
+        ErrorModuleInput(
+            ModuleInput(
+                constructor_input=FunctionInput(weight=make_weight((2,))),
+                forward_input=FunctionInput(
+                    make_input((2, 3)),
+                    torch.tensor([0, 1], device=device, dtype=torch.long),
+                ),
+            ),
+            error_on=ModuleErrorEnum.FORWARD_ERROR,
+            error_type=RuntimeError,
+            error_regex="weight tensor should be defined either for all .* classes or no classes",
+        ),
+        # Prob target with ignore_index
+        ErrorModuleInput(
+            ModuleInput(
+                constructor_input=FunctionInput(ignore_index=1),
+                forward_input=FunctionInput(
+                    make_input((2, 3)),
+                    make_input((2, 3)).softmax(dim=1).detach(),
+                ),
+            ),
+            error_on=ModuleErrorEnum.FORWARD_ERROR,
+            error_type=RuntimeError,
+            error_regex="ignore_index is not supported for floating point target",
+        ),
+        # label_smoothing out of range
+        ErrorModuleInput(
+            ModuleInput(
+                constructor_input=FunctionInput(label_smoothing=1.5),
+                forward_input=FunctionInput(
+                    make_input((2, 3)),
+                    torch.tensor([0, 1], device=device, dtype=torch.long),
+                ),
+            ),
+            error_on=ModuleErrorEnum.FORWARD_ERROR,
+            error_type=RuntimeError,
+            error_regex="label_smoothing must be between 0.0 and 1.0",
+        ),
+        # Target spatial size mismatch (3D input)
+        ErrorModuleInput(
+            ModuleInput(
+                constructor_input=FunctionInput(),
+                forward_input=FunctionInput(
+                    make_input((2, 3, 5)),
+                    torch.zeros((2, 4), device=device, dtype=torch.long),
+                ),
+            ),
+            error_on=ModuleErrorEnum.FORWARD_ERROR,
+            error_type=RuntimeError,
+            error_regex="Expected target size",
+        ),
+    ]
+
 
 def module_inputs_torch_nn_CTCLoss(module_info, device, dtype, requires_grad, training, **kwargs):
     make_input = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad)
@@ -1739,6 +1846,10 @@ def module_inputs_torch_nn_GroupNorm(module_info, device, dtype, requires_grad,
             constructor_input=FunctionInput(3, 6, 1e-3),
             forward_input=FunctionInput(make_input((4, 6, 5))),
             desc='1d_affine'),
+        ModuleInput(
+            constructor_input=FunctionInput(3, 6, 1e-3, bias=False),
+            forward_input=FunctionInput(make_input((4, 6, 5))),
+            desc='1d_affine_not_bias'),
         ModuleInput(
             constructor_input=FunctionInput(3, 12, 1e-3),
             forward_input=FunctionInput(make_input((4, 12))),
@@ -1759,6 +1870,10 @@ def module_inputs_torch_nn_GroupNorm(module_info, device, dtype, requires_grad,
             constructor_input=FunctionInput(3, 6, 1e-3),
             forward_input=FunctionInput(make_input((4, 6, 2, 3))),
             desc='2d_affine'),
+        ModuleInput(
+            constructor_input=FunctionInput(3, 9, 1e-3, bias=False),
+            forward_input=FunctionInput(make_input((4, 9, 2, 3))),
+            desc='2d_affine_not_bias'),
         ModuleInput(
             constructor_input=FunctionInput(3, 3, 1e-3, False),
             forward_input=FunctionInput(make_input((4, 3, 2, 3))),
@@ -1766,8 +1881,7 @@ def module_inputs_torch_nn_GroupNorm(module_info, device, dtype, requires_grad,
         ModuleInput(
             constructor_input=FunctionInput(1, 3, 1e-3, False),
             forward_input=FunctionInput(make_input((4, 3, 2, 3))),
-            desc='2d_no_affine_LN'),
-    ]
+            desc='2d_no_affine_LN'),]
 
 
 def module_error_inputs_torch_nn_GroupNorm(module_info, device, dtype, requires_grad, training, **kwargs):
@@ -1956,8 +2070,21 @@ def module_inputs_torch_nn_InstanceNormNd(module_info, device, dtype, requires_g
             ),
             forward_input=FunctionInput(make_input(input_no_batch_shape)),
             reference_fn=no_batch_dim_reference_fn,
-            desc='no_batch_dim')
-    ]
+            desc='no_batch_dim'),
+        ModuleInput(
+            constructor_input=(
+                FunctionInput(eps, momentum, affine=True) if lazy else
+                FunctionInput(num_features, eps, momentum, affine=True)
+            ),
+            forward_input=FunctionInput(make_input(input_batch_shape)),
+            desc='affine'),
+        ModuleInput(
+            constructor_input=(
+                FunctionInput(eps, momentum, affine=True, bias=False) if lazy else
+                FunctionInput(num_features, eps, momentum, affine=True, bias=False)
+            ),
+            forward_input=FunctionInput(make_input(input_batch_shape)),
+            desc='affine_not_bias'),]
 
 def module_inputs_torch_nn_LayerNorm(module_info, device, dtype, requires_grad, training, **kwargs):
     make_input = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad)
@@ -2158,6 +2285,47 @@ def module_inputs_torch_nn_MaxPool2d(module_info, device, dtype, requires_grad,
             desc='return_indices'),
     ]
 
+
+def module_error_inputs_torch_nn_MaxPool2d(module_info, device, dtype, requires_grad, training, **kwargs):
+    """
+    Error inputs for MaxPool2d that test error messages for invalid inputs.
+    """
+    make_input = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad)
+
+    return [
+        # Wrong input dimensions: 2D input instead of 3D/4D
+        ErrorModuleInput(
+            ModuleInput(
+                constructor_input=FunctionInput(2),
+                forward_input=FunctionInput(make_input((3, 4))),  # 2D input
+            ),
+            error_on=ModuleErrorEnum.FORWARD_ERROR,
+            error_type=RuntimeError,
+            error_regex=r"non-empty 3D or 4D \(batch mode\) tensor expected for input"
+        ),
+        # Wrong input dimensions: 5D input
+        ErrorModuleInput(
+            ModuleInput(
+                constructor_input=FunctionInput(2),
+                forward_input=FunctionInput(make_input((1, 2, 3, 4, 5))),  # 5D input
+            ),
+            error_on=ModuleErrorEnum.FORWARD_ERROR,
+            error_type=RuntimeError,
+            error_regex=r"non-empty 3D or 4D \(batch mode\) tensor expected for input"
+        ),
+        # Invalid padding: padding > kernel_size / 2
+        ErrorModuleInput(
+            ModuleInput(
+                constructor_input=FunctionInput(3, padding=5),  # kernel=3, pad=5 > 3/2
+                forward_input=FunctionInput(make_input((1, 1, 10, 10))),
+            ),
+            error_on=ModuleErrorEnum.FORWARD_ERROR,
+            error_type=RuntimeError,
+            error_regex=r"pad should be at most half of effective kernel size"
+        ),
+    ]
+
+
 def module_inputs_torch_nn_MaxPool3d(module_info, device, dtype, requires_grad, training, **kwargs):
     make_input = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad)
 
@@ -2398,6 +2566,31 @@ def reference_fn(m, p, i, t, constructor_kwargs=constructor_kwargs):
     return module_inputs
 
 
+def module_error_inputs_torch_nn_MultiMarginLoss(module_info, device, dtype, requires_grad, training, **kwargs):
+    """
+    Error inputs for MultiMarginLoss that test the improved error message
+    for inconsistent target size.
+    Regression test for issue #106251.
+    """
+    make_input = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad)
+
+    return [
+        # Test: target size doesn't match batch size (5 vs 3)
+        ErrorModuleInput(
+            ModuleInput(
+                constructor_input=FunctionInput(),
+                forward_input=FunctionInput(
+                    make_input((5, 10)),  # input: batch_size=5, num_classes=10
+                    torch.tensor([0, 1, 2], device=device, dtype=torch.long),  # target: wrong size (3)
+                ),
+            ),
+            error_on=ModuleErrorEnum.FORWARD_ERROR,
+            error_type=RuntimeError,
+            error_regex=r"target tensor should be 1-D with size equal to.*Expected target size \[5\].*but got \[3\]"
+        ),
+    ]
+
+
 def module_inputs_torch_nn_MultiLabelSoftMarginLoss(module_info, device, dtype, requires_grad, training, **kwargs):
     make_input = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad)
     make_target = partial(make_tensor, device=device, dtype=torch.long, requires_grad=False)
@@ -2692,6 +2885,57 @@ def module_inputs_torch_nn_Embedding(module_info, device, dtype, requires_grad,
     ]
 
 
+def module_error_inputs_torch_nn_Embedding(module_info, device, dtype, requires_grad, training, **kwargs):
+    """
+    Error inputs for Embedding that test error messages for invalid inputs.
+    """
+    samples = []
+
+    # Out of range indices: index exceeds num_embeddings
+    # Only test on CPU - CUDA triggers kernel assertion instead of Python exception
+    if torch.device(device).type == 'cpu':
+        samples.append(
+            ErrorModuleInput(
+                ModuleInput(
+                    constructor_input=FunctionInput(num_embeddings=10, embedding_dim=3),
+                    forward_input=FunctionInput(torch.tensor([0, 5, 15], device=device, dtype=torch.long)),
+                ),
+                error_on=ModuleErrorEnum.FORWARD_ERROR,
+                error_type=IndexError,
+                error_regex=r"index out of range in self"
+            )
+        )
+
+    # Float indices: wrong dtype for indices (works on all devices)
+    samples.append(
+        ErrorModuleInput(
+            ModuleInput(
+                constructor_input=FunctionInput(num_embeddings=10, embedding_dim=3),
+                forward_input=FunctionInput(torch.tensor([1.5, 2.5], device=device, dtype=torch.float32)),
+            ),
+            error_on=ModuleErrorEnum.FORWARD_ERROR,
+            error_type=RuntimeError,
+            error_regex=r"Expected tensor for argument.*indices.*to have.*scalar type.*Long.*Int"
+        )
+    )
+
+    # Negative num_embeddings (construction error, device-independent)
+    samples.append(
+        ErrorModuleInput(
+            ModuleInput(
+                constructor_input=FunctionInput(num_embeddings=-1, embedding_dim=3),
+                forward_input=FunctionInput(),
+            ),
+            error_on=ModuleErrorEnum.CONSTRUCTION_ERROR,
+            error_type=RuntimeError,
+            error_regex=r"Trying to create tensor with negative dimension"
+        )
+    )
+
+    return samples
+
+
+
 def module_inputs_torch_nn_MultiheadAttention(module_info, device, dtype, requires_grad, training, **kwargs):
     # Currently all samples below are for validating the no-batch-dim support.
     make_input = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad)
@@ -3733,6 +3977,8 @@ def module_error_inputs_torch_nn_Pad3d(module_info, device, dtype, requires_grad
                    # Not implemented for chalf on CPU
                    DecorateInfo(unittest.expectedFailure, 'TestModule', 'test_cpu_gpu_parity',
                                 dtypes=(torch.chalf,), device_type='cuda'),
+                   DecorateInfo(skipIfXpu, 'TestModule', 'test_cpu_gpu_parity',
+                                dtypes=(torch.chalf,), device_type='xpu'),
                ),
                decorators=(
                    DecorateInfo(precisionOverride({torch.float32: 1e-04}), 'TestModule', 'test_memory_format'),
@@ -3752,7 +3998,7 @@ def module_error_inputs_torch_nn_Pad3d(module_info, device, dtype, requires_grad
                    # Not implemented for chalf on CPU
                    DecorateInfo(unittest.expectedFailure, 'TestModule', 'test_cpu_gpu_parity',
                                 dtypes=(torch.chalf,), device_type='cuda'),
-                   DecorateInfo(unittest.expectedFailure, 'TestModule', 'test_cpu_gpu_parity',
+                   DecorateInfo(skipIfXpu, 'TestModule', 'test_cpu_gpu_parity',
                                 dtypes=(torch.chalf,), device_type='xpu'),
                ),
                decorators=(
@@ -3951,6 +4197,7 @@ def module_error_inputs_torch_nn_Pad3d(module_info, device, dtype, requires_grad
                ),
     ModuleInfo(torch.nn.MaxPool2d,
                module_inputs_func=module_inputs_torch_nn_MaxPool2d,
+               module_error_inputs_func=module_error_inputs_torch_nn_MaxPool2d,
                ),
     ModuleInfo(torch.nn.MaxPool3d,
                module_inputs_func=module_inputs_torch_nn_MaxPool3d,
@@ -3993,6 +4240,7 @@ def module_error_inputs_torch_nn_Pad3d(module_info, device, dtype, requires_grad
                ),
     ModuleInfo(torch.nn.MultiMarginLoss,
                module_inputs_func=module_inputs_torch_nn_MultiMarginLoss,
+               module_error_inputs_func=module_error_inputs_torch_nn_MultiMarginLoss,
                skips=(
                    # No channels_last support for loss functions.
                    DecorateInfo(unittest.skip("Skipped!"), 'TestModule', 'test_memory_format'),
@@ -4069,6 +4317,7 @@ def module_error_inputs_torch_nn_Pad3d(module_info, device, dtype, requires_grad
                ),
     ModuleInfo(torch.nn.CrossEntropyLoss,
                module_inputs_func=module_inputs_torch_nn_CrossEntropyLoss,
+               module_error_inputs_func=module_error_inputs_torch_nn_CrossEntropyLoss,
                dtypes=get_all_fp_dtypes(include_half=True, include_bfloat16=False),
                decorators=(
                    # No channels_last support for loss functions.
@@ -4235,6 +4484,7 @@ def module_error_inputs_torch_nn_Pad3d(module_info, device, dtype, requires_grad
                ),
     ModuleInfo(torch.nn.Embedding,
                module_inputs_func=module_inputs_torch_nn_Embedding,
+               module_error_inputs_func=module_error_inputs_torch_nn_Embedding,
                gradcheck_nondet_tol=GRADCHECK_NONDET_TOL,
                decorators=[
                    DecorateInfo(toleranceOverride({torch.float32: tol(atol=1e-4, rtol=1e-4)}),
diff --git a/torch/testing/_internal/common_mps.py b/torch/testing/_internal/common_mps.py
index 926bd819f1108..f8fc0f4437d8b 100644
--- a/torch/testing/_internal/common_mps.py
+++ b/torch/testing/_internal/common_mps.py
@@ -1,6 +1,5 @@
 import unittest
 from collections.abc import Sequence
-from typing import Optional
 
 import torch
 
@@ -13,7 +12,7 @@
     def mps_ops_modifier(
         ops: Sequence[OpInfo],
         device_type: str = "mps",
-        xfail_exclusion: Optional[list[str]] = None,
+        xfail_exclusion: list[str] | None = None,
         sparse: bool = False,
     ) -> Sequence[OpInfo]:
         if xfail_exclusion is None:
@@ -55,12 +54,17 @@ def mps_ops_modifier(
             "contiguous",
             "cos",
             "cosh",
+            "cross",
+            "cumsum",
+            "cumprod",
+            "cumulative_trapezoid",
             "diag",
             "diag_embed",
             "diagflat",
             "diagonal",
             "diagonal_copy",
             "diagonal_scatter",
+            "dist",
             "divno_rounding_mode",
             "dsplit",
             "empty",
@@ -72,6 +76,7 @@ def mps_ops_modifier(
             "expand",
             "expand_as",
             "expand_copy",
+            "gather",
             "flatten",
             "fill",
             "full",
@@ -81,23 +86,29 @@ def mps_ops_modifier(
             "imag",
             "index_add",
             "index_copy",
+            "index_fill",
             "index_select",
             "index_put",
             "isfinite",
             "isinf",
             "isreal",
+            "istft",
             "item",
             "kron",
+            "linalg.cross",
             "linalg.diagonal",
             "linalg.householder_product",
             "linalg.svd",
+            "linalg.vander",
             "linalg.vecdot",
+            "linalg.vector_norm",
             "log10",
             "log1p",
             "log2",
             "log",
             "logaddexp",
             "logaddexp2",
+            "logcumsumexp",
             "mH",
             "mT",
             "masked_fill",
@@ -119,12 +130,22 @@ def mps_ops_modifier(
             "nn.functional.conv_transpose2d",
             "nn.functional.conv_transpose3d",
             "nn.functional.feature_alpha_dropoutwithout_train",
+            "nn.functional.l1_loss",
+            "nn.functional.linear",
+            "nn.functional.normalize",
             "nn.functional.padcircular",
+            "nn.functional.pairwise_distance",
             "nn.functional.softminwith_dtype",
             "nn.functional.softsign",
             "nn.functional.tanhshrink",
+            "nn.functional.triplet_margin_loss",
+            "nn.functional.triplet_margin_with_distance_loss",
             "nn.functional.unfold",
             "nonzero",
+            "nonzero_static",
+            "norm",
+            "normfro",
+            "norminf",
             "ones",
             "ones_like",
             "outer",
@@ -134,6 +155,7 @@ def mps_ops_modifier(
             "randn",
             "ravel",
             "real",
+            "repeat",
             "repeat_interleave",
             "reshape_as",
             "reshape",
@@ -142,6 +164,8 @@ def mps_ops_modifier(
             "rsqrt",
             "rsub",
             "scalar_tensor",
+            "scatter",
+            "scatter_add",
             "select",
             "sgn",
             "sigmoid",
@@ -166,9 +190,11 @@ def mps_ops_modifier(
             "svd",
             "t",
             "t_copy",
+            "take_along_dim",
             "tanh",
             "tan",
             "tensor_split",
+            "tile",
             "transpose",
             "transpose_copy",
             "tril",
@@ -262,7 +288,10 @@ def mps_ops_modifier(
             "logical_xor",
             "logsumexp",
             "long",
+            "masked.cumsum",
+            "masked.cumprod",
             "masked.mean",
+            "masked.normalize",
             "masked.prod",
             "masked.std",
             "masked.sum",
@@ -306,23 +335,18 @@ def mps_ops_modifier(
         }
 
         # Those ops are not expected to work
-        UNIMPLEMENTED_XFAILLIST: dict[str, Optional[list]] = {
+        UNIMPLEMENTED_XFAILLIST: dict[str, list | None] = {
             # Failures due to lack of op implementation on MPS backend
             "logspace": None,
             "logspacetensor_overload": None,
             "linalg.eig": None,
             "linalg.eigvals": None,
             "put": None,
-            "cholesky_solve": None,
             "frexp": None,
             "geqrf": None,
             "nn.functional.grid_sample": None,  # Unsupported Border padding mode
             "hash_tensor": None,
             "heaviside": None,
-            "index_reduceprod": None,
-            "index_reducemean": None,
-            "index_reduceamax": None,
-            "index_reduceamin": None,
             # "kthvalue": None,
             "lcm": None,
             "linalg.cond": None,
@@ -336,7 +360,6 @@ def mps_ops_modifier(
             "linalg.matrix_norm": [torch.float32],
             "linalg.norm": [torch.float32],
             "linalg.normsubgradients_at_zero": [torch.float32],
-            "linalg.qr": None,
             "linalg.svdvals": None,
             "masked.median": None,
             "matrix_exp": None,
@@ -346,7 +369,6 @@ def mps_ops_modifier(
                 torch.int32,
                 torch.int64,
                 torch.uint8,
-                torch.bool,
             ],
             "median": [torch.bool],
             "mode": None,
@@ -377,7 +399,6 @@ def mps_ops_modifier(
                 torch.int16,
                 torch.int32,
                 torch.uint8,
-                torch.bool,
                 torch.int8,
             ],
             "nn.functional.batch_norm": [
@@ -526,8 +547,6 @@ def mps_ops_modifier(
                 torch.bool,
                 torch.int8,
             ],
-            "nn.functional.padreflect": [torch.bool],
-            "nn.functional.padreplicate": [torch.bool],
             "nn.functional.padreplicate_negative": [torch.bool],
             "nn.functional.pdist": None,
             "nn.functional.relu": [torch.bool],
@@ -536,7 +555,6 @@ def mps_ops_modifier(
                 torch.int16,
                 torch.int32,
                 torch.uint8,
-                torch.bool,
                 torch.int8,
             ],
             "nn.functional.softplus": [
@@ -547,12 +565,7 @@ def mps_ops_modifier(
                 torch.int16,
             ],
             "nn.functional.norm": None,
-            "nn.functional.threshold": [torch.bool],
             "ormqr": None,
-            "pca_lowrank": None,
-            "pow": [torch.bool],
-            "qr": None,
-            "remainder": [torch.bool],
             "rounddecimals_0": [
                 torch.uint8,
                 torch.int8,
@@ -588,7 +601,6 @@ def mps_ops_modifier(
             "symeig": None,
             "take": None,
             "to": None,
-            "trunc": [torch.bool],
             "var_meanunbiased": [
                 torch.uint8,
                 torch.int8,
@@ -621,8 +633,6 @@ def mps_ops_modifier(
             "float_power": None,
             "linalg.matrix_rankhermitian": None,
             "linalg.pinvhermitian": None,
-            "linalg.pinvsingular": None,  # Missing `aten::linalg_qr.out`.
-            "nonzero_static": None,
             # MPS: input sizes must be divisible by output sizes
             "nn.functional.adaptive_avg_pool1d": None,
             "nn.functional.adaptive_avg_pool2d": None,
@@ -638,7 +648,6 @@ def mps_ops_modifier(
                 torch.float16,
             ],
             # Unsupported dtypes
-            "histc": [torch.float16, torch.bfloat16],
             # GEMM on MPS is not supported for integral types
             "nn.functional.linear": [
                 torch.int16,
@@ -657,7 +666,7 @@ def mps_ops_modifier(
                 torch.int8,
             ],
         }
-        UNIMPLEMENTED_XFAILLIST_SPARSE: dict[str, Optional[list]] = {
+        UNIMPLEMENTED_XFAILLIST_SPARSE: dict[str, list | None] = {
             "logspace": None,
             "logspacetensor_overload": None,
             "linalg.eig": None,
@@ -675,7 +684,7 @@ def mps_ops_modifier(
         if sparse:
             UNIMPLEMENTED_XFAILLIST.update(UNIMPLEMENTED_XFAILLIST_SPARSE)
 
-        UNDEFINED_XFAILLIST: dict[str, Optional[list]] = {
+        UNDEFINED_XFAILLIST: dict[str, list | None] = {
             # Top 60 operators
             # topk fails with duplicate indices
             "topk": [
@@ -769,7 +778,7 @@ def mps_ops_modifier(
             ],
         }
 
-        ON_MPS_XFAILLIST: dict[str, Optional[list]] = {
+        ON_MPS_XFAILLIST: dict[str, list | None] = {
             # Failures due to lack of implementation of downstream functions on MPS backend
             # TODO: remove these once downstream function 'aten::_linalg_svd.U' have been implemented
             "linalg.matrix_rank": None,
@@ -921,16 +930,6 @@ def mps_ops_grad_modifier(ops: Sequence[OpInfo]) -> Sequence[OpInfo]:
             "grid_sampler_3d": None,
             "igamma": None,  # currently not supported for any device
             "igammac": None,  # currently not supported for any device
-            "linalg.solve": [torch.float16, torch.float32],  # missing `aten::lu_solve`.
-            "linalg.solve_ex": [
-                torch.float16,
-                torch.float32,
-            ],  # missing `aten::lu_solve`.
-            "linalg.tensorsolve": [
-                torch.float16,
-                torch.float32,
-            ],  # missing `aten::lu_solve`.
-            "aminmax": [torch.float32, torch.float16],
             "special.i1": [torch.float16],  # "i1_backward" not implemented for 'Half'
             "special.i1e": [torch.float16],  # "i1e_backward" not implemented for 'Half'
             # Correctness issues
@@ -951,8 +950,6 @@ def mps_ops_grad_modifier(ops: Sequence[OpInfo]) -> Sequence[OpInfo]:
             "nextafter": None,
             # derivative for aten::floor_divide is not implemented on CPU
             "floor_divide": [torch.float16, torch.float32],
-            # derivative for aten::narrow_copy is not implemented on CPU
-            "narrow_copy": [torch.float16, torch.float32],
             # derivative for aten::_histogramdd_from_bin_cts is not implemented on CPU
             "histogramdd": [torch.float16, torch.float32],
             # derivative for aten::histogram is not implemented
@@ -1035,11 +1032,6 @@ def mps_ops_error_inputs_modifier(ops: Sequence[OpInfo]) -> Sequence[OpInfo]:
             "clamp_max",
             "clamp_min",
             "masked_scatter",
-            # unsupported float64 dtype
-            "multinomial",
-            "gather",
-            "scatter",
-            "scatter_add",
             # MPS does not support tensor dimensions > 16
             "amax",
             "amin",
@@ -1060,7 +1052,7 @@ def addDecorator(op: OpInfo, d: DecorateInfo) -> None:
     def mps_ops_modifier(
         ops: Sequence[OpInfo],
         device_type: str = "mps",
-        xfail_exclusion: Optional[list[str]] = None,
+        xfail_exclusion: list[str] | None = None,
         sparse: bool = False,
     ) -> Sequence[OpInfo]:
         return ops
diff --git a/torch/testing/_internal/common_nn.py b/torch/testing/_internal/common_nn.py
index dd97912d2495a..cf04c33e4ef87 100644
--- a/torch/testing/_internal/common_nn.py
+++ b/torch/testing/_internal/common_nn.py
@@ -17,14 +17,14 @@
 from torch.nn import _reduction as _Reduction
 from torch.testing._internal import common_utils
 from torch.testing._internal.common_utils import TestCase, to_gpu, freeze_rng_state, is_iterable, \
-    gradcheck, gradgradcheck, set_default_dtype, skipIfTorchDynamo, TEST_WITH_ROCM
+    gradcheck, gradgradcheck, set_default_dtype, skipIfTorchDynamo, skipIfRocm, TEST_WITH_ROCM
 from torch.testing._internal.common_cuda import TEST_CUDA, SM90OrLater
 from torch.autograd.gradcheck import _get_numerical_jacobian, _iter_tensors
 from torch.autograd import Variable
 from torch.types import _TensorOrTensors
 import torch.backends.cudnn
 
-from typing import Union, Any
+from typing import Any
 from collections.abc import Callable
 from collections.abc import Sequence
 
@@ -508,7 +508,7 @@ def nllloss_no_reduce_test():
 
 def nllloss_no_reduce_ignore_index_test():
     t = Variable(torch.empty(15).uniform_().mul(10).floor().long())
-    kwargs: dict[str, Union[int, str]] = {'ignore_index': 2, 'reduction': 'none'}
+    kwargs: dict[str, int | str] = {'ignore_index': 2, 'reduction': 'none'}
     return dict(
         fullname='NLLLoss_no_reduce_ignore_index',
         constructor=wrap_functional(
@@ -611,7 +611,7 @@ def nllloss2d_no_reduce_test():
 
 def nllloss2d_no_reduce_ignore_index_test():
     t = Variable(torch.rand(2, 5, 5).mul(3).floor().long())
-    kwargs: dict[str, Union[int, str]] = {'ignore_index': 1, 'reduction': 'none'}
+    kwargs: dict[str, int | str] = {'ignore_index': 1, 'reduction': 'none'}
     return dict(
         fullname='NLLLoss2d_no_reduce_ignore_index',
         constructor=wrap_functional(
@@ -668,7 +668,7 @@ def nlllossNd_no_reduce_test():
 
 def nlllossNd_no_reduce_ignore_index_test():
     t = Variable(torch.rand(2, 5, 5, 2, 2).mul(3).floor().long())
-    kwargs: dict[str, Union[int, str]] = {'ignore_index': 1, 'reduction': 'none'}
+    kwargs: dict[str, int | str] = {'ignore_index': 1, 'reduction': 'none'}
     return dict(
         fullname='NLLLossNd_no_reduce_ignore_index',
         constructor=wrap_functional(
@@ -1727,7 +1727,9 @@ def get_new_module_tests():
             check_gradgrad=False,
             desc='discontiguous',
             default_dtype=torch.double,
-            decorator=skipIfTorchDynamo("https://github.com/pytorch/pytorch/issues/117971")
+            decorator=lambda fn: skipIfTorchDynamo("https://github.com/pytorch/pytorch/issues/117971")(
+                skipIfRocm(msg="https://github.com/pytorch/pytorch/issues/186910")(fn)
+            ),
         ),
         dict(
             module_name='EmbeddingBag',
@@ -3231,7 +3233,7 @@ def _zero_grad_parameters(self, module: nn.Module) -> None:
     @abstractmethod
     def _backward(self, module: nn.Module,
                   input: _TensorOrTensors, output: torch.Tensor,
-                  grad_output: Union[torch.Tensor, Sequence[torch.Tensor]],
+                  grad_output: torch.Tensor | Sequence[torch.Tensor],
                   create_graph: bool = False):
         raise NotImplementedError
 
diff --git a/torch/testing/_internal/common_ops_unbacked.py b/torch/testing/_internal/common_ops_unbacked.py
index ee290e6943d18..766126de04375 100644
--- a/torch/testing/_internal/common_ops_unbacked.py
+++ b/torch/testing/_internal/common_ops_unbacked.py
@@ -40,7 +40,6 @@ def skip(op_name, variant_name="", *, device_type=None, dtypes=None):
     xfail("cumulative_trapezoid"),
     xfail("diagonal_scatter"),
     xfail("diff"),
-    xfail("dist"),
     xfail("dsplit"),
     xfail("equal"),
     xfail("exponential"),
@@ -97,12 +96,9 @@ def skip(op_name, variant_name="", *, device_type=None, dtypes=None):
     xfail("linalg.solve"),
     xfail("linalg.solve_ex"),
     xfail("linalg.solve_triangular"),
-    xfail("linalg.svd"),
-    xfail("linalg.svdvals"),
     xfail("linalg.tensorinv"),
     xfail("linalg.tensorsolve"),
     xfail("linalg.vander"),
-    xfail("linalg.vector_norm"),
     xfail("log_normal"),
     xfail("logdet"),
     xfail("logsumexp"),
@@ -128,7 +124,6 @@ def skip(op_name, variant_name="", *, device_type=None, dtypes=None):
     xfail("masked.var"),
     xfail("max_pool2d_with_indices_backward"),
     xfail("multinomial"),
-    xfail("nanquantile"),
     xfail("nn.functional.adaptive_avg_pool1d"),
     xfail("nn.functional.adaptive_avg_pool2d"),
     xfail("nn.functional.adaptive_avg_pool3d"),
@@ -144,12 +139,6 @@ def skip(op_name, variant_name="", *, device_type=None, dtypes=None):
     xfail("nn.functional.binary_cross_entropy"),
     xfail("nn.functional.binary_cross_entropy_with_logits"),
     xfail("nn.functional.channel_shuffle"),
-    xfail("nn.functional.conv1d"),
-    xfail("nn.functional.conv2d"),
-    xfail("nn.functional.conv3d"),
-    xfail("nn.functional.conv_transpose1d"),
-    xfail("nn.functional.conv_transpose2d"),
-    xfail("nn.functional.conv_transpose3d"),
     xfail("nn.functional.cosine_similarity"),
     xfail("nn.functional.cross_entropy"),
     xfail("nn.functional.ctc_loss"),
@@ -163,7 +152,6 @@ def skip(op_name, variant_name="", *, device_type=None, dtypes=None):
     xfail("nn.functional.fractional_max_pool2d"),
     xfail("nn.functional.fractional_max_pool3d"),
     xfail("nn.functional.gaussian_nll_loss"),
-    xfail("nn.functional.glu"),
     xfail("nn.functional.grid_sample"),
     xfail("nn.functional.group_norm"),
     xfail("nn.functional.huber_loss"),
@@ -200,16 +188,12 @@ def skip(op_name, variant_name="", *, device_type=None, dtypes=None):
     xfail("nn.functional.smooth_l1_loss"),
     xfail("nn.functional.unfold"),
     xfail("nn.functional.upsample_bilinear"),
-    xfail("norm"),
-    xfail("norm", "fro"),
-    xfail("norm", "nuc"),
     xfail("normal"),
     xfail("normal", "in_place"),
     xfail("normal", "number_mean"),
     xfail("ormqr"),
     xfail("pca_lowrank"),
     xfail("pinverse"),
-    xfail("quantile"),
     xfail("qr"),
     xfail("rand_like"),
     xfail("randint_like"),
@@ -222,7 +206,6 @@ def skip(op_name, variant_name="", *, device_type=None, dtypes=None):
     xfail("sparse.mm", "reduce"),
     xfail("split"),
     xfail("stft"),
-    xfail("svd"),
     xfail("svd_lowrank"),
     xfail("sum_to_size"),
     xfail("take"),
diff --git a/torch/testing/_internal/common_optimizers.py b/torch/testing/_internal/common_optimizers.py
index fbf56b9010f27..fb67c4d845021 100644
--- a/torch/testing/_internal/common_optimizers.py
+++ b/torch/testing/_internal/common_optimizers.py
@@ -6,7 +6,7 @@
 import unittest
 from copy import deepcopy
 from enum import Enum
-from typing import Any, Union
+from typing import Any
 
 import torch
 from torch import Tensor
@@ -58,9 +58,7 @@ class OptimizerInput:
 
     def __init__(
         self,
-        params: Union[
-            list[Parameter], list[Tensor], dict[Any, Any], list[dict[str, Any]]
-        ],
+        params: list[Parameter] | list[Tensor] | dict[Any, Any] | list[dict[str, Any]],
         kwargs: dict[str, Any],
         desc: str = "",
     ):
@@ -508,6 +506,28 @@ def optim_error_inputs_func_adagrad(device, dtype):
                 error_regex="Invalid lr_decay value: -0.5",
             ),
         ]
+    if _get_device_type(device) == "cuda":
+        sample_tensor = torch.empty((), device=device, dtype=dtype)
+        error_inputs += [
+            ErrorOptimizerInput(
+                OptimizerInput(
+                    params=[sample_tensor],
+                    kwargs={"foreach": True, "fused": True},
+                    desc="`fused` and `foreach` cannot be `True` together",
+                ),
+                error_type=RuntimeError,
+                error_regex="`fused` and `foreach` cannot be `True` together",
+            ),
+            ErrorOptimizerInput(
+                OptimizerInput(
+                    params=[sample_tensor],
+                    kwargs={"fused": True, "differentiable": True},
+                    desc="`fused` does not support `differentiable`",
+                ),
+                error_type=RuntimeError,
+                error_regex="`fused` does not support `differentiable`",
+            ),
+        ]
     return error_inputs
 
 
@@ -1329,7 +1349,7 @@ def optim_error_inputs_func_sparseadam(device, dtype):
     return error_inputs
 
 
-def _get_device_type(device: Union[str, torch.device]) -> str:
+def _get_device_type(device: str | torch.device) -> str:
     # Returns the device type as a string, e.g., "cpu" or "cuda"
     if isinstance(device, torch.device):
         device = str(device.type)
@@ -1537,6 +1557,7 @@ def _get_optim_inputs_including_global_cliquey_kwargs(
                 "CompiledOptimizerParityTests",
                 "test_correctness",
                 device_type="xpu",
+                active_if=lambda kwargs: kwargs.get("use_closure", False),
             ),
             DecorateInfo(
                 skipIfTorchDynamo("See #133268 regarding dtype being None"),
@@ -1600,7 +1621,7 @@ def _get_optim_inputs_including_global_cliquey_kwargs(
             "maximize",
             "capturable",
         ),
-        supports_fused_on=("cpu",),
+        supports_fused_on=("cpu", "cuda"),
         supports_sparse=True,
         metadata_for_sparse=(
             {"lr": 0.1, "weight_decay": 0, "lr_decay": 0},
@@ -1625,6 +1646,16 @@ def _get_optim_inputs_including_global_cliquey_kwargs(
                 "TestOptimRenewed",
                 "test_fused_matches_forloop",
             ),
+            DecorateInfo(
+                toleranceOverride(
+                    {  # https://github.com/pytorch/pytorch/issues/116202
+                        torch.float32: tol(atol=5e-04, rtol=0.015),
+                    }
+                ),
+                "TestOptimRenewed",
+                "test_mixed_device_dtype",
+                active_if=TEST_WITH_TORCHDYNAMO,
+            ),
         ),
         skips=(
             DecorateInfo(
diff --git a/torch/testing/_internal/common_quantization.py b/torch/testing/_internal/common_quantization.py
index a3ea1d769d454..abd69aa09e28c 100644
--- a/torch/testing/_internal/common_quantization.py
+++ b/torch/testing/_internal/common_quantization.py
@@ -74,7 +74,7 @@
 import os
 
 import unittest
-from typing import Any, Optional, Union
+from typing import Any
 from collections.abc import Callable
 
 import numpy as np
@@ -1067,7 +1067,7 @@ def assert_types_for_matched_subgraph_pairs(
 
             def _get_underlying_op_type(
                 node: Node, gm: GraphModule
-            ) -> Union[Callable, str]:
+            ) -> Callable | str:
                 if node.op == "call_module":
                     mod = getattr(gm, node.target)
                     return type(mod)
@@ -2478,8 +2478,8 @@ def __init__(self) -> None:
     def forward(
         self,
         input: torch.Tensor,
-        offsets: Optional[torch.Tensor] = None,
-        per_sample_weights: Optional[torch.Tensor] = None,
+        offsets: torch.Tensor | None = None,
+        per_sample_weights: torch.Tensor | None = None,
     ):
         x = self.emb(input, offsets, per_sample_weights)
         x = self.quant(x)
diff --git a/torch/testing/_internal/common_utils.py b/torch/testing/_internal/common_utils.py
index c533ee687c5bd..193d625c5e9ce 100644
--- a/torch/testing/_internal/common_utils.py
+++ b/torch/testing/_internal/common_utils.py
@@ -49,9 +49,7 @@
 from statistics import mean
 from typing import (
     Any,
-    Optional,
     TypeVar,
-    Union,
 )
 from collections.abc import Callable
 from collections.abc import Iterable, Iterator
@@ -84,17 +82,19 @@
 )
 from torch.testing import make_tensor
 from torch.testing._comparison import (
+    _unwrap_dtensor_for_comparison,
     BooleanPair,
     NonePair,
+    not_close_error_metas,
     NumberPair,
     Pair,
     TensorLikePair,
 )
-from torch.testing._comparison import not_close_error_metas
 from torch.testing._internal.common_dtype import get_all_dtypes
 from torch.utils._import_utils import _check_module_exists
 import torch.utils._pytree as pytree
 from torch.utils import cpp_extension
+from torch._utils import _is_privateuse1_backend_available
 try:
     import pytest  # type: ignore[import-not-found]
     has_pytest = True
@@ -116,7 +116,7 @@ class ProfilingMode(Enum):
 
 # Set by parse_cmd_line_args() if called
 DISABLED_TESTS_FILE = ""
-GRAPH_EXECUTOR : Optional[ProfilingMode] = None
+GRAPH_EXECUTOR : ProfilingMode | None = None
 LOG_SUFFIX = ""
 PYTEST_SINGLE_TEST = ""
 REPEAT_COUNT = 0
@@ -304,7 +304,7 @@ def repro_env_var_prefix() -> str:
 )
 
 # possibly restrict OpInfo tests to a single sample input
-OPINFO_SAMPLE_INPUT_INDEX: Optional[int] = TestEnvironment.def_setting(
+OPINFO_SAMPLE_INPUT_INDEX: int | None = TestEnvironment.def_setting(
     "OPINFO_SAMPLE_INPUT_INDEX",
     env_var="PYTORCH_OPINFO_SAMPLE_INPUT_INDEX",
     default=None,
@@ -355,7 +355,7 @@ def wrapper(*args, **kwargs):
 
 # Tries to extract the current test function by crawling the stack.
 # If unsuccessful, return None.
-def extract_test_fn() -> Optional[Callable]:
+def extract_test_fn() -> Callable | None:
     try:
         stack = inspect.stack()
         for frame_info in stack:
@@ -382,7 +382,7 @@ class TrackedInput:
 
 # Attempt to pull out tracked input information from the test function.
 # A TrackedInputIter is used to insert this information.
-def get_tracked_input() -> Optional[TrackedInput]:
+def get_tracked_input() -> TrackedInput | None:
     test_fn = extract_test_fn()
     if test_fn is None:
         return None
@@ -1410,15 +1410,9 @@ def printErrors(self) -> None:
 IS_X86 = platform.machine() in ('x86_64', 'i386')
 IS_ARM64 = platform.machine() in ('arm64', 'aarch64')
 IS_S390X = platform.machine() == "s390x"
-
-def is_avx512_vnni_supported():
-    if sys.platform != 'linux':
-        return False
-    with open("/proc/cpuinfo", encoding="ascii") as f:
-        lines = f.read()
-    return "vnni" in lines
-
-IS_AVX512_VNNI_SUPPORTED = is_avx512_vnni_supported()
+IS_AVX512_VNNI_SUPPORTED = torch.cpu.get_capabilities().get("avx512_vnni", False)
+IS_CPU_EXT_SVE_SUPPORTED = torch.cpu.get_capabilities().get("sve", False)
+IS_CPU_CAPABILITY_SVE256 = torch._C._get_cpu_capability() == "SVE256"
 
 if IS_WINDOWS:
     @contextmanager
@@ -1460,12 +1454,6 @@ def TemporaryDirectoryName(suffix=None):
             yield d
 
 
-def is_privateuse1_backend_available():
-    privateuse1_backend_name = torch._C._get_privateuse1_backend_name()
-    privateuse1_backend_module = getattr(torch, privateuse1_backend_name, None)
-    return (is_available := getattr(privateuse1_backend_module, "is_available", None)) and is_available()
-
-
 def make_lazy_class(cls):
 
     def lazy_init(self, cb):
@@ -1519,7 +1507,7 @@ class LazyVal:
 TEST_ACCELERATOR = LazyVal(lambda: torch.accelerator.is_available())  # type: ignore[call-arg]
 TEST_MULTIACCELERATOR = LazyVal(lambda: torch.accelerator.device_count() > 1)  # type: ignore[call-arg]
 custom_device_mod = getattr(torch, torch._C._get_privateuse1_backend_name(), None)
-TEST_PRIVATEUSE1 = is_privateuse1_backend_available()
+TEST_PRIVATEUSE1 = _is_privateuse1_backend_available()
 TEST_PRIVATEUSE1_DEVICE_TYPE = torch._C._get_privateuse1_backend_name()
 TEST_NUMBA = _check_module_exists('numba')
 TEST_TRANSFORMERS = _check_module_exists('transformers')
@@ -1531,6 +1519,44 @@ class LazyVal:
 
 TEST_Z3 = _check_module_exists('z3')
 
+# DSL availability (lazy evaluation to avoid import overhead)
+class LazyDSLCheck:
+    """Lazy DSL availability checker to avoid import-time overhead"""
+    def __init__(self):
+        self._registry = None
+        self._import_attempted = False
+
+    def _get_registry(self):
+        if not self._import_attempted:
+            self._import_attempted = True
+            try:
+                from torch._native.dsl_registry import dsl_registry
+                self._registry = dsl_registry
+            except ImportError:
+                self._registry = None
+        return self._registry
+
+    def is_available(self, dsl_name: str) -> bool:
+        """Check if specific DSL is available"""
+        registry = self._get_registry()
+        return registry.is_dsl_available(dsl_name) if registry is not None else False
+
+    def list_available(self) -> list[str]:
+        """Get list of available DSLs"""
+        registry = self._get_registry()
+        return list(registry.list_available_dsls()) if registry is not None else []
+
+    def list_all(self) -> list[str]:
+        """Get list of all registered DSLs"""
+        registry = self._get_registry()
+        return list(registry.list_all_dsls()) if registry is not None else []
+
+_dsl_checker = LazyDSLCheck()
+
+# Lazy constants to avoid import-time overhead
+TEST_TRITON_DSL = LazyVal(lambda: _dsl_checker.is_available('triton'))
+TEST_CUTEDSL = LazyVal(lambda: _dsl_checker.is_available('cutedsl'))
+
 def split_if_not_empty(x: str):
     return x.split(",") if len(x) != 0 else []
 
@@ -1538,6 +1564,34 @@ def split_if_not_empty(x: str):
 
 skipIfNoDill = unittest.skipIf(not TEST_DILL, "no dill")
 
+# DSL skip decorators (following existing pattern)
+skipIfNoTritonDSL = unittest.skipIf(not TEST_TRITON_DSL, "Triton DSL not available")
+skipIfNoCuteDSL = unittest.skipIf(not TEST_CUTEDSL, "CuTeDSL not available")
+
+def skipIfDSLUnavailable(dsl_name: str, reason: str | None = None):
+    """Skip test if specific DSL is not available"""
+    available = _dsl_checker.is_available(dsl_name)
+    msg = reason or f"{dsl_name} DSL not available"
+    return unittest.skipIf(not available, msg)
+
+def skipUnlessDSLAvailable(dsl_name: str, reason: str | None = None):
+    """Skip test unless specific DSL is available"""
+    available = _dsl_checker.is_available(dsl_name)
+    msg = reason or f"{dsl_name} DSL required"
+    return unittest.skipUnless(available, msg)
+
+def get_available_dsls() -> list[str]:
+    """Get list of available DSL names for test parameterization"""
+    return _dsl_checker.list_available()
+
+def is_dsl_available(dsl_name: str) -> bool:
+    """Check if specific DSL is available for conditional testing"""
+    return _dsl_checker.is_available(dsl_name)
+
+def get_all_dsls() -> list[str]:
+    """Get all registered DSL names (available or not) for comprehensive testing"""
+    return _dsl_checker.list_all()
+
 
 NO_MULTIPROCESSING_SPAWN: bool = False
 TEST_WITH_ASAN: bool = TestEnvironment.def_flag(
@@ -1559,6 +1613,7 @@ def split_if_not_empty(x: str):
 TEST_WITH_ROCM: bool = TestEnvironment.def_flag(
     "TEST_WITH_ROCM",
     env_var="PYTORCH_TEST_WITH_ROCM",
+    implied_by_fn=lambda: torch.version.hip is not None,
 )
 TEST_WITH_MTIA: bool = TestEnvironment.def_flag(
     "TEST_WITH_MTIA",
@@ -1605,6 +1660,12 @@ def split_if_not_empty(x: str):
 )
 
 TEST_CUDA_CUDSS = TEST_CUDA and (torch.version.cuda and int(torch.version.cuda.split(".")[0]) >= 12)
+TEST_CUDA_GRAPH_CONDITIONAL_NODES = TEST_CUDA_GRAPH and (
+    torch.version.cuda and (
+        (int(torch.version.cuda.split(".")[0]) >= 12 and int(torch.version.cuda.split(".")[1]) >= 4) or
+        (int(torch.version.cuda.split(".")[0]) >= 13)
+    )
+)
 
 TEST_CUDA_PYTHON_BINDINGS = _check_module_exists("cuda.bindings") and (
     torch.version.cuda and int(torch.version.cuda.split(".")[0]) >= 12
@@ -1896,7 +1957,7 @@ def wrapper(*args, **kwargs):
 
 
 def make_dynamo_test(
-    fn: Optional[Callable[..., Any]] = None
+    fn: Callable[..., Any] | None = None
 ) -> Callable[..., Any]:
     """
     Decorator function to create a dynamo test case. A function annotate with
@@ -2004,7 +2065,7 @@ def has_corresponding_torch_dtype(np_dtype):
 
 def skipIfNNModuleInlined(
     msg="test doesn't currently work with nn module inlining",
-    condition=torch._dynamo.config.inline_inbuilt_nn_modules,
+    condition=True,
 ):
     def decorator(fn):
         if not isinstance(fn, type):
@@ -2110,12 +2171,31 @@ def wrapper(*args, **kwargs):
     return dec_fn
 
 def skipIfMPS(fn):
+    sig = inspect.signature(fn)
+    has_device_arg = "device" in sig.parameters
+
+    if not has_device_arg:
+        warnings.warn(
+            f"skipIfMPS applied to {fn.__qualname__} which has no 'device' parameter. "
+            "Consider using device-generic tests with instantiate_device_type_tests instead.",
+            stacklevel=2,
+        )
+
     @wraps(fn)
     def wrapper(*args, **kwargs):
-        if TEST_MPS:
+        if has_device_arg:
+            # For device-generic tests, only skip when actually running on MPS
+            slf = args[0] if args else None
+            if slf is not None:
+                device_type = getattr(slf, "device_type", None) or getattr(
+                    slf, "device", None
+                )
+                if isinstance(device_type, str) and device_type == "mps":
+                    raise unittest.SkipTest("test doesn't currently work with MPS")
+        elif TEST_MPS:
             raise unittest.SkipTest("test doesn't currently work with MPS")
-        else:
-            fn(*args, **kwargs)
+        return fn(*args, **kwargs)
+
     return wrapper
 
 
@@ -3084,8 +3164,8 @@ class UnittestPair(Pair):
 
     Define the :attr:`UnittestPair.CLS` in a subclass to indicate which class(es) of the inputs the pair should support.
     """
-    CLS: Union[type, tuple[type, ...]]
-    TYPE_NAME: Optional[str] = None
+    CLS: type | tuple[type, ...]
+    TYPE_NAME: str | None = None
 
     def __init__(self, actual, expected, **other_parameters):
         self._check_inputs_isinstance(actual, expected, cls=self.CLS)
@@ -4276,10 +4356,10 @@ def assertEqual(
             self,
             x,
             y,
-            msg: Optional[Union[str, Callable[[str], str]]] = None,
+            msg: str | Callable[[str], str] | None = None,
             *,
-            atol: Optional[float] = None,
-            rtol: Optional[float] = None,
+            atol: float | None = None,
+            rtol: float | None = None,
             equal_nan=True,
             exact_dtype=True,
             # TODO: default this to True
@@ -4317,6 +4397,8 @@ def to_list(input):
         if isinstance(y, torch.Tensor) and y.is_nested and y.layout == torch.strided:
             y = y.unbind()
 
+        x, y = _unwrap_dtensor_for_comparison(x, y)
+
         error_metas = not_close_error_metas(
             x,
             y,
@@ -4363,8 +4445,8 @@ def to_list(input):
                 (lambda generated_msg: f"{generated_msg}\n{msg}") if isinstance(msg, str) and self.longMessage else msg
             )
 
-    def assertNotEqual(self, x, y, msg: Optional[str] = None, *,                                       # type: ignore[override]
-                       atol: Optional[float] = None, rtol: Optional[float] = None, **kwargs) -> None:
+    def assertNotEqual(self, x, y, msg: str | None = None, *,                                       # type: ignore[override]
+                       atol: float | None = None, rtol: float | None = None, **kwargs) -> None:
         with self.assertRaises(AssertionError, msg=msg):
             self.assertEqual(x, y, msg, atol=atol, rtol=rtol, **kwargs)
 
@@ -4384,7 +4466,7 @@ def assertObjectIn(self, obj: Any, iterable: Iterable[Any]) -> None:
     # _ignore_not_implemented_error is True
     def assertRaises(self, expected_exception, *args, **kwargs):
         if self._ignore_not_implemented_error:
-            context: Optional[AssertRaisesContextIgnoreNotImplementedError] = \
+            context: AssertRaisesContextIgnoreNotImplementedError | None = \
                 AssertRaisesContextIgnoreNotImplementedError(expected_exception, self)  # type: ignore[call-arg]
             try:
                 return context.handle('assertRaises', args, kwargs)  # type: ignore[union-attr, arg-type]
@@ -4698,7 +4780,7 @@ def _attempt_load_from_subprocess(
         self,
         file: pathlib.Path,
         import_string: str,
-        expected_failure_message: Optional[str] = None
+        expected_failure_message: str | None = None
     ) -> None:
         """
         Attempts weights_only `torch.load` in a subprocess. This is used to test that
@@ -6044,3 +6126,17 @@ def wrapper(self, *args, **kwargs):
 
         return wrapper
     return decorator
+
+def get_gcc_major_version():
+    """
+    Return GCC major version as int, or None if GCC is not available.
+    """
+    try:
+        out = subprocess.check_output(
+            ["gcc", "-dumpfullversion", "-dumpversion"],
+            stderr=subprocess.STDOUT,
+            text=True,
+        ).strip()
+        return int(out.split(".")[0])
+    except Exception:
+        return None
diff --git a/torch/testing/_internal/common_xpu.py b/torch/testing/_internal/common_xpu.py
new file mode 100644
index 0000000000000..ee4f8c921eac5
--- /dev/null
+++ b/torch/testing/_internal/common_xpu.py
@@ -0,0 +1,88 @@
+import enum
+import functools
+
+import torch
+import torch.xpu
+from torch.testing._internal.common_utils import IS_WINDOWS, LazyVal, TEST_XPU
+
+
+XPU_ALREADY_INITIALIZED_ON_IMPORT = torch.xpu.is_initialized()
+
+
+class XPUCodename(enum.Enum):
+    PVC = "PVC"  # Intel® Data Center GPU Max Series
+    BMG = "BMG"  # Intel® Arc™ Pro Battlemage Graphics
+
+
+class XPUArch(enum.IntEnum):
+    Unknown = 0
+    Xe = 1  # Xe HPC
+    Xe2 = 2
+
+
+# device_id -> GPU codename
+# From https://github.com/intel/intel-graphics-compiler/blob/master/inc/common/igfxfmid.h
+_DEVICE_ID_TO_CODENAME = {
+    0x0BD0: XPUCodename.PVC,
+    0x0BD4: XPUCodename.PVC,
+    0x0BD5: XPUCodename.PVC,
+    0x0BD6: XPUCodename.PVC,
+    0x0BD7: XPUCodename.PVC,
+    0x0BD8: XPUCodename.PVC,
+    0x0BD9: XPUCodename.PVC,
+    0x0BDA: XPUCodename.PVC,
+    0x0BDB: XPUCodename.PVC,
+    0x0B69: XPUCodename.PVC,
+    0x0B6E: XPUCodename.PVC,
+    0xE202: XPUCodename.BMG,
+    0xE20B: XPUCodename.BMG,
+    0xE20C: XPUCodename.BMG,
+    0xE20D: XPUCodename.BMG,
+    0xE210: XPUCodename.BMG,
+    0xE212: XPUCodename.BMG,
+    0xE215: XPUCodename.BMG,
+    0xE216: XPUCodename.BMG,
+    0xE220: XPUCodename.BMG,
+    0xE221: XPUCodename.BMG,
+    0xE222: XPUCodename.BMG,
+    0xE223: XPUCodename.BMG,
+}
+
+# GPU codename -> architecture
+_CODENAME_TO_ARCH = {
+    XPUCodename.PVC: XPUArch.Xe,
+    XPUCodename.BMG: XPUArch.Xe2,
+}
+
+
+@functools.lru_cache(1)
+def get_xpu_codename() -> XPUCodename | None:
+    device_id = torch.xpu.get_device_capability()["device_id"]
+    return _DEVICE_ID_TO_CODENAME.get(device_id)
+
+
+@functools.lru_cache(1)
+def get_xpu_arch() -> XPUArch | None:
+    codename = get_xpu_codename()
+    return _CODENAME_TO_ARCH.get(codename, XPUArch.Unknown)
+
+
+Xe2_Or_Later = LazyVal(
+    lambda: torch.xpu.is_available() and get_xpu_arch() >= XPUArch.Xe2
+)
+
+
+def evaluate_platform_supports_flash_attention():
+    if TEST_XPU:
+        return not IS_WINDOWS and Xe2_Or_Later
+    return False
+
+
+PLATFORM_SUPPORTS_FLASH_ATTENTION_XPU: bool = LazyVal(
+    lambda: evaluate_platform_supports_flash_attention()
+)
+
+# Importing this module should NOT eagerly initialize XPU
+if not XPU_ALREADY_INITIALIZED_ON_IMPORT:
+    if torch.xpu.is_initialized():
+        raise AssertionError("XPU should not be initialized on import")
diff --git a/torch/testing/_internal/distributed/_tensor/common_dtensor.py b/torch/testing/_internal/distributed/_tensor/common_dtensor.py
index 29f9781ea2a2f..68c6baa863e98 100644
--- a/torch/testing/_internal/distributed/_tensor/common_dtensor.py
+++ b/torch/testing/_internal/distributed/_tensor/common_dtensor.py
@@ -7,16 +7,22 @@
 import functools
 import itertools
 import sys
+import threading
 import types
+import unittest
 from collections.abc import Callable, Iterator, Sequence
 from dataclasses import dataclass
 from functools import partial, wraps
-from typing import Any, cast, Optional, TypeVar, Union
+from typing import Any, cast, TypeVar
 
 import torch
 import torch.distributed as dist
 import torch.nn as nn
 import torch.nn.functional as F
+from torch.distributed._functional_collectives import (
+    all_gather_tensor_autograd,
+    reduce_scatter_tensor_autograd,
+)
 from torch.distributed._local_tensor import (
     LocalIntNode,
     LocalTensor,
@@ -26,9 +32,11 @@
 )
 from torch.distributed.tensor import (
     DeviceMesh,
+    distribute_module,
     distribute_tensor,
     DTensor,
     init_device_mesh,
+    Partial,
     Placement,
     Replicate,
     Shard,
@@ -38,6 +46,7 @@
 from torch.distributed.tensor.parallel import (
     ColwiseParallel,
     parallelize_module,
+    ParallelStyle,
     PrepareModuleInput,
     RowwiseParallel,
     SequenceParallel,
@@ -58,6 +67,7 @@
     TEST_WITH_ROCM,
     TEST_XPU,
 )
+from torch.testing._internal.distributed.fake_pg import FakeStore
 from torch.utils._pytree import tree_flatten, tree_unflatten, TreeSpec
 
 
@@ -137,6 +147,7 @@ class ModelArgs:
     use_attn_mask: bool = True
     weight_tying: bool = True
     checkpoint_activations: bool = False
+    num_experts: int = 0
 
 
 class Attention(nn.Module):
@@ -192,19 +203,247 @@ def forward(self, x):
         return self.resid_dropout(self.w2(self.gelu(self.w1(x))))
 
 
+class Experts(nn.Module):
+    def __init__(self, dim: int, hidden_dim: int, num_experts: int):
+        super().__init__()
+        self.num_experts = num_experts
+        self.w1 = nn.Parameter(torch.empty(num_experts, hidden_dim, dim))
+        self.w2 = nn.Parameter(torch.empty(num_experts, dim, hidden_dim))
+        self.reset_parameters()
+        self.gelu = nn.GELU()
+
+    def reset_parameters(self):
+        nn.init.normal_(self.w1, std=0.02)
+        nn.init.normal_(self.w2, std=0.02)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        # Weights are DTensors (sharded by EP/TP) but x is a plain tensor
+        # (dispatched by EP hooks), so extract local shards for bmm.
+        if isinstance(self.w1, DTensor):
+            w1, w2 = self.w1.to_local(), self.w2.to_local()
+        else:
+            w1, w2 = self.w1, self.w2
+        E = w1.shape[0]
+        x_exp = x.unsqueeze(0).expand(E, -1, -1)
+        h = self.gelu(torch.bmm(x_exp, w1.transpose(-2, -1)))
+        out = torch.bmm(h, w2.transpose(-2, -1))
+        return out.sum(dim=0)
+
+
+class ExpertLayer(nn.Module):
+    def __init__(self, num_experts: int, dim: int, hidden_dim: int):
+        super().__init__()
+        self.num_experts = num_experts
+        self.experts = Experts(dim=dim, hidden_dim=hidden_dim, num_experts=num_experts)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        bs, slen, dim = x.shape
+        x_flat = x.view(-1, dim)
+        expert_out = self.experts(x_flat)
+        expert_out = expert_out / self.num_experts
+        return expert_out.view(bs, slen, dim)
+
+
+class TensorParallelForExpert(ParallelStyle):
+    """TP for Experts: shard w1 colwise (Shard(1)), w2 rowwise (Shard(2)).
+
+    For seq parallel, set input_layouts=Shard(0) and output_layouts=Shard(0)
+    to all-gather tokens before computation and reduce-scatter after.
+    """
+
+    def __init__(self, *, input_layouts=None, output_layouts=None):
+        super().__init__()
+        self.input_layouts = input_layouts
+        self.output_layouts = output_layouts or Replicate()
+
+    def _partition_fn(self, name, mod, device_mesh):
+        for pn, p in mod.named_parameters(recurse=False):
+            if pn == "w1":
+                placement = [Shard(1)]
+            elif pn == "w2":
+                placement = [Shard(2)]
+            else:
+                continue
+            mod.register_parameter(
+                pn, nn.Parameter(distribute_tensor(p, device_mesh, placement))
+            )
+
+    def _input_fn(self, mod, inputs, device_mesh):
+        if self.input_layouts is not None:
+            x = inputs[0]
+            x = (
+                DTensor.from_local(x, device_mesh, [self.input_layouts])
+                .redistribute(device_mesh, [Replicate()])
+                .to_local()
+            )
+            return (x,)
+        return inputs
+
+    def _output_fn(self, mod, outputs, device_mesh):
+        return (
+            DTensor.from_local(outputs, device_mesh, [Partial()])
+            .redistribute(device_mesh, [self.output_layouts])
+            .to_local()
+        )
+
+    def _apply(self, module, device_mesh):
+        return distribute_module(
+            module,
+            device_mesh,
+            partition_fn=self._partition_fn,
+            input_fn=self._input_fn,
+            output_fn=self._output_fn,
+        )
+
+
+class ExpertParallel(ParallelStyle):
+    """Distributes experts across ranks with Shard(0) on the expert dimension.
+
+    Dispatch: all-gather tokens so every rank sees all tokens.
+    Combine: reduce-scatter (sum) expert outputs back to token owners.
+    """
+
+    def _partition_fn(self, name: str, mod: nn.Module, device_mesh: DeviceMesh) -> None:
+        for param_name, param in mod.named_parameters(recurse=False):
+            dist_param = nn.Parameter(distribute_tensor(param, device_mesh, [Shard(0)]))
+            mod.register_parameter(param_name, dist_param)
+
+    def _token_dispatch(
+        self, mod: nn.Module, inputs: tuple, device_mesh: DeviceMesh
+    ) -> tuple[torch.Tensor]:
+        (x,) = inputs
+        x_gathered = all_gather_tensor_autograd(
+            x,
+            gather_dim=0,
+            group=device_mesh.get_group(),
+        )
+        x_gathered = torch.ops._c10d_functional.wait_tensor(x_gathered)
+        return (x_gathered,)
+
+    def _token_combine(
+        self, mod: nn.Module, output: torch.Tensor, device_mesh: DeviceMesh
+    ) -> torch.Tensor:
+        result = reduce_scatter_tensor_autograd(
+            output,
+            "sum",
+            scatter_dim=0,
+            group=device_mesh.get_group(),
+        )
+        result = torch.ops._c10d_functional.wait_tensor(result)
+        return result
+
+    def _apply(self, module: nn.Module, device_mesh: DeviceMesh) -> nn.Module:
+        return distribute_module(
+            module,
+            device_mesh,
+            partition_fn=self._partition_fn,
+            input_fn=self._token_dispatch,
+            output_fn=self._token_combine,
+        )
+
+
+class ExpertParallelWithTP(ParallelStyle):
+    """Combined EP + TP for experts.
+
+    Applied to ExpertLayer. Distributes expert params on a 2D (ep, tp) mesh
+    with [Shard(0), Shard(1/2)]. Token dispatch/combine hooks are registered
+    on ExpertLayer (outer), TP reduction hook on Experts (inner), so forward
+    execution is: EP dispatch -> TP input -> forward -> TP reduce -> EP combine.
+    """
+
+    def __init__(
+        self,
+        ep_mesh: DeviceMesh,
+        tp_mesh: DeviceMesh,
+    ):
+        super().__init__()
+        self.ep_mesh = ep_mesh
+        self.tp_mesh = tp_mesh
+        self.ep_tp_mesh = DeviceMesh._concatenate([ep_mesh, tp_mesh])
+
+    def _apply(self, module: nn.Module, device_mesh: DeviceMesh) -> nn.Module:
+        experts = module.experts  # type: ignore[attr-defined]
+
+        # Partition expert weights on 2D (ep, tp) mesh
+        for pn, p in experts.named_parameters(recurse=False):
+            if pn == "w1":
+                placements = [Shard(0), Shard(1)]
+            elif pn == "w2":
+                placements = [Shard(0), Shard(2)]
+            else:
+                continue
+            experts.register_parameter(
+                pn, nn.Parameter(distribute_tensor(p, self.ep_tp_mesh, placements))
+            )
+
+        # EP dispatch/combine hooks on ExpertLayer (outer module)
+        ep_mesh = self.ep_mesh
+
+        def ep_dispatch(mod, inputs):
+            (x,) = inputs
+            x = all_gather_tensor_autograd(x, gather_dim=0, group=ep_mesh.get_group())
+            return (torch.ops._c10d_functional.wait_tensor(x),)
+
+        def ep_combine(mod, inputs, output):
+            out = reduce_scatter_tensor_autograd(
+                output, "sum", scatter_dim=0, group=ep_mesh.get_group()
+            )
+            return torch.ops._c10d_functional.wait_tensor(out)
+
+        module.register_forward_pre_hook(ep_dispatch)
+        module.register_forward_hook(ep_combine)
+
+        # TP reduction hook on Experts (inner module)
+        tp_mesh = self.tp_mesh
+
+        def tp_allreduce_input_grad(mod, inputs):
+            (x,) = inputs
+            return (
+                DTensor.from_local(x, tp_mesh, [Replicate()]).to_local(
+                    grad_placements=[Partial()]
+                ),
+            )
+
+        experts.register_forward_pre_hook(tp_allreduce_input_grad)
+
+        def tp_reduce(mod, inputs, output):
+            return (
+                DTensor.from_local(output, tp_mesh, [Partial()])
+                .redistribute(tp_mesh, [Replicate()])
+                .to_local()
+            )
+
+        experts.register_forward_hook(tp_reduce)
+        return module
+
+
 class TransformerBlock(nn.Module):
     def __init__(self, args: ModelArgs):
         super().__init__()
         self.attention_norm = nn.LayerNorm(args.dim)
         self.attention = Attention(args)
         self.ffn_norm = nn.LayerNorm(args.dim)
-        self.feed_forward = FeedForward(
-            args.dim, hidden_dim=4 * args.dim, dropout_p=args.dropout_p
-        )
+
+        self.has_experts = args.num_experts > 0
+        if self.has_experts:
+            self.expert_layer = ExpertLayer(
+                args.num_experts,
+                dim=args.dim,
+                hidden_dim=4 * args.dim,
+            )
+        else:
+            self.feed_forward = FeedForward(
+                args.dim,
+                hidden_dim=4 * args.dim,
+                dropout_p=args.dropout_p,
+            )
 
     def forward(self, x):
         h = x + self.attention(self.attention_norm(x))
-        out = h + self.feed_forward(self.ffn_norm(h))
+        if self.has_experts:
+            out = h + self.expert_layer(self.ffn_norm(h))
+        else:
+            out = h + self.feed_forward(self.ffn_norm(h))
         return out
 
 
@@ -254,97 +493,130 @@ def forward(self, tokens):
     @staticmethod
     def parallelize(
         module: "Transformer",
-        device_mesh: DeviceMesh,
+        tp_mesh: DeviceMesh | None,
         use_seq_parallel: bool,
         local_output_for_attn: bool = False,
+        ep_mesh: DeviceMesh | None = None,
     ) -> nn.Module:
         if not isinstance(module, Transformer):
             raise AssertionError(f"Requires Transformer but got {module}")
-        # Parallelize the root submodules.
-        if use_seq_parallel:
-            root_plan = {
-                "tok_embeddings": RowwiseParallel(
-                    input_layouts=Replicate(), output_layouts=Shard(1)
-                ),
-                "pos_embeddings": RowwiseParallel(
-                    input_layouts=Replicate(), output_layouts=Shard(0)
-                ),
-                "norm": SequenceParallel(),
-            }
-        else:
-            root_plan = {
-                "tok_embeddings": RowwiseParallel(
-                    input_layouts=Replicate(), output_layouts=Replicate()
-                ),
-                "pos_embeddings": RowwiseParallel(
-                    input_layouts=Replicate(), output_layouts=Replicate()
-                ),
-            }
+        if tp_mesh is None and ep_mesh is None:
+            raise ValueError("At least one of tp_mesh or ep_mesh must be provided")
 
-        module_tp = parallelize_module(module, device_mesh, root_plan)
-        # Parallelize the attention and feed forward submodules.
-        for layer in module_tp.layers:
-            layer_parallelize_plan = {}
+        # Parallelize the root submodules with TP.
+        if tp_mesh is not None:
             if use_seq_parallel:
-                layer_parallelize_plan["attention"] = PrepareModuleInput(
-                    input_layouts=Shard(1),
-                    desired_input_layouts=Replicate(),
+                root_plan = {
+                    "tok_embeddings": RowwiseParallel(
+                        input_layouts=Replicate(), output_layouts=Shard(1)
+                    ),
+                    "pos_embeddings": RowwiseParallel(
+                        input_layouts=Replicate(), output_layouts=Shard(0)
+                    ),
+                    "norm": SequenceParallel(),
+                }
+            else:
+                root_plan = {
+                    "tok_embeddings": RowwiseParallel(
+                        input_layouts=Replicate(), output_layouts=Replicate()
+                    ),
+                    "pos_embeddings": RowwiseParallel(
+                        input_layouts=Replicate(), output_layouts=Replicate()
+                    ),
+                }
+            parallelize_module(module, tp_mesh, root_plan)
+
+        # Parallelize the attention and feed forward submodules.
+        for layer in module.layers:
+            if tp_mesh is not None:
+                layer_parallelize_plan = {}
+                if use_seq_parallel:
+                    layer_parallelize_plan["attention"] = PrepareModuleInput(
+                        input_layouts=Shard(1),
+                        desired_input_layouts=Replicate(),
+                    )
+                    # shard the RMSNorms
+                    layer_parallelize_plan["attention_norm"] = SequenceParallel()
+                    layer_parallelize_plan["ffn_norm"] = SequenceParallel()
+                layer_parallelize_plan["attention.wq"] = ColwiseParallel(
+                    use_local_output=local_output_for_attn
+                )
+                layer_parallelize_plan["attention.wk"] = ColwiseParallel(
+                    use_local_output=local_output_for_attn
+                )
+                layer_parallelize_plan["attention.wv"] = ColwiseParallel(
+                    use_local_output=local_output_for_attn
+                )
+                layer_parallelize_plan["attention.wo"] = (
+                    RowwiseParallel(output_layouts=Shard(1))
+                    if use_seq_parallel
+                    else RowwiseParallel()
                 )
-                # shard the RMSNorms
-                layer_parallelize_plan["attention_norm"] = SequenceParallel()
-                layer_parallelize_plan["ffn_norm"] = SequenceParallel()
-            layer_parallelize_plan["attention.wq"] = ColwiseParallel(
-                use_local_output=local_output_for_attn
-            )
-            layer_parallelize_plan["attention.wk"] = ColwiseParallel(
-                use_local_output=local_output_for_attn
-            )
-            layer_parallelize_plan["attention.wv"] = ColwiseParallel(
-                use_local_output=local_output_for_attn
-            )
-            layer_parallelize_plan["attention.wo"] = (
-                RowwiseParallel(output_layouts=Shard(1))
-                if use_seq_parallel
-                else RowwiseParallel()
-            )
 
-            layer_parallelize_plan["feed_forward.w1"] = (
-                ColwiseParallel(input_layouts=Shard(1))
-                if use_seq_parallel
-                else ColwiseParallel()
-            )
-            layer_parallelize_plan["feed_forward.w2"] = (
-                RowwiseParallel(output_layouts=Shard(1))
-                if use_seq_parallel
-                else RowwiseParallel()
-            )
+                if not layer.has_experts:
+                    layer_parallelize_plan["feed_forward.w1"] = (
+                        ColwiseParallel(input_layouts=Shard(1))
+                        if use_seq_parallel
+                        else ColwiseParallel()
+                    )
+                    layer_parallelize_plan["feed_forward.w2"] = (
+                        RowwiseParallel(output_layouts=Shard(1))
+                        if use_seq_parallel
+                        else RowwiseParallel()
+                    )
+                elif ep_mesh is None:
+                    # No EP mesh provided, use TP for experts
+                    layer_parallelize_plan["expert_layer.experts"] = (
+                        TensorParallelForExpert(
+                            input_layouts=Shard(0),
+                            output_layouts=Shard(0),
+                        )
+                        if use_seq_parallel
+                        else TensorParallelForExpert()
+                    )
 
-            parallelize_module(layer, device_mesh, layer_parallelize_plan)
+                parallelize_module(layer, tp_mesh, layer_parallelize_plan)
 
-        # Parallelize the output submodule. If weight tying is enabled, we need to
-        # make sure output.weight is sharded consistently as tok_embeddings.weight,
-        # at the cost of the all_reduce operation using RowwiseParallel.
-        output_parallelize_plan = (
-            ColwiseParallel(
-                input_layouts=Shard(1),
-                output_layouts=Replicate(),
-            )
-            if use_seq_parallel
-            else ColwiseParallel(output_layouts=Replicate())
-        )
-        parallelize_module(module_tp.output, device_mesh, output_parallelize_plan)
+            # EP (+ optional TP) for experts
+            if ep_mesh is not None and layer.has_experts:
+                if tp_mesh is not None:
+                    parallelize_module(
+                        layer.expert_layer,
+                        ep_mesh,
+                        ExpertParallelWithTP(ep_mesh, tp_mesh),
+                    )
+                else:
+                    parallelize_module(
+                        layer.expert_layer.experts, ep_mesh, ExpertParallel()
+                    )
 
-        if local_output_for_attn:
-            for layer in module_tp.layers:
-                layer.attention.n_heads = (
-                    module_tp.model_args.n_heads // device_mesh.size()
+        if tp_mesh is not None:
+            # Parallelize the output submodule. If weight tying is enabled,
+            # we need to make sure output.weight is sharded consistently as
+            # tok_embeddings.weight, at the cost of the all_reduce operation
+            # using RowwiseParallel.
+            output_parallelize_plan = (
+                ColwiseParallel(
+                    input_layouts=Shard(1),
+                    output_layouts=Replicate(),
                 )
+                if use_seq_parallel
+                else ColwiseParallel(output_layouts=Replicate())
+            )
+            parallelize_module(module.output, tp_mesh, output_parallelize_plan)
 
-        # Manually set output.weight so that parameters and gradients are shared.
-        if module_tp.model_args.weight_tying:
-            module_tp.output.weight = module_tp.tok_embeddings.weight
+            if local_output_for_attn:
+                for layer in module.layers:
+                    layer.attention.n_heads = (
+                        module.model_args.n_heads // tp_mesh.size()
+                    )
 
-        return module_tp
+            # Manually set output.weight so that parameters and gradients
+            # are shared.
+            if module.model_args.weight_tying:
+                module.output.weight = module.tok_embeddings.weight
+
+        return module
 
 
 def skip_unless_torch_gpu(method: T) -> T:
@@ -360,18 +632,69 @@ def skip_unless_torch_gpu(method: T) -> T:
     return cast(T, skip_if_lt_x_gpu(NUM_DEVICES)(method))
 
 
-class DTensorContinuousTestBase(MultiProcContinuousTest):
-    @classmethod
-    def device_type(cls) -> str:
-        # if enough GPU/XPU/HPU we can use those devices, otherwise we fallback to CPU
+class DTensorTestMixin:
+    """Shared test helpers for DTensorTestBase and DTensorContinuousTestBase."""
+
+    @property
+    def is_local_tensor_enabled(self) -> bool:
+        return False
+
+    @property
+    def device_type(self) -> str:
         if (
             not (TEST_CUDA or TEST_XPU or TEST_HPU or TEST_PRIVATEUSE1)
-            or DEVICE_COUNT < cls.world_size
+            or DEVICE_COUNT < self.world_size
         ):
             return "cpu"
         else:
             return DEVICE_TYPE
 
+    def build_device_mesh(self) -> DeviceMesh:
+        return init_device_mesh(self.device_type, (self.world_size,))
+
+    def init_manual_seed_for_rank(self) -> None:
+        torch.manual_seed(self.rank)
+
+    def _test_op_on_dtensor(self, op_call, *args, **kwargs) -> None:
+        """
+        Checks ``op_call(dtensor).full_tensor() == op_call(dtensor.full_tensor())``.
+        Unlike _test_op where the DTensor sharding is generated by DTensorConverter,
+        this function takes in DTensor object directly as argument and test the equality
+        of calling op on full_tensor() and DTensor.
+        """
+        args_flattened, args_spec = tree_flatten(args)
+        full_tensor_args_flattened = tuple(
+            arg.full_tensor().detach().clone() if isinstance(arg, DTensor) else arg
+            for arg in args_flattened
+        )
+        full_tensor_args = tree_unflatten(full_tensor_args_flattened, args_spec)
+        full_tensor_kwargs = {
+            k: v.full_tensor() if isinstance(v, DTensor) else v
+            for k, v in kwargs.items()
+        }
+
+        out_flattened, _ = tree_flatten(
+            op_call(*full_tensor_args, **full_tensor_kwargs)
+        )
+        d_out_flattened, _ = tree_flatten(op_call(*args, **kwargs))
+        d_out_full_tensor_flattened = [dt.full_tensor() for dt in d_out_flattened]
+        self.assertEqual(out_flattened, d_out_full_tensor_flattened)
+
+    # pyre-ignore[2]:
+    def _test_op(self, mesh: DeviceMesh, op_call, *args, **kwargs) -> None:
+        out = op_call(*args, **kwargs)
+        dtc = DTensorConverter(mesh, args, kwargs)
+        for d_args, d_kwargs in dtc:
+            # pyre can't find assertTrue anymore?
+            self.assertEqual(dtc.successful(), True)
+            d_out = op_call(*d_args, **d_kwargs)
+            self.assertEqual(d_out.full_tensor(), out)
+
+    def run_subtests(self, *args, **kwargs):
+        return run_subtests(self, *args, **kwargs)
+
+
+class DTensorContinuousTestBase(DTensorTestMixin, MultiProcContinuousTest):
     @classmethod
     def backend_str(cls) -> str:
         backend = dist.get_default_backend_for_device(DEVICE_TYPE)
@@ -392,38 +715,86 @@ def _init_pg(cls, rank, world_size, rdvz_file):
         super()._init_pg(rank, world_size, rdvz_file)
 
 
-class DTensorTestBase(MultiProcessTestCase):
+class LocalDTensorContinuousTestBase(DTensorContinuousTestBase):
     @property
     def is_local_tensor_enabled(self) -> bool:
-        return False
+        return True
+
+    def _handle_test_skip(self, msg: str) -> None:
+        self.skipTest(msg)
+
+    def _get_local_tensor_mode(self):
+        return LocalTensorMode(frozenset(range(self.world_size)))
+
+    @classmethod
+    def _ensure_processes_spawned(cls):
+        if cls._processes_spawned:
+            return
+        if cls.world_size == -2:
+            cls.world_size = NUM_DEVICES
+        store = FakeStore()
+        dist.init_process_group(
+            backend="fake",
+            world_size=cls.world_size,
+            rank=0,
+            store=store,
+        )
+        cls.processes = []
+        cls.task_queues = []
+        cls.completion_queues = []
+        cls._processes_spawned = True
+
+    @classmethod
+    def tearDownClass(cls):
+        if cls._processes_spawned:
+            dist.destroy_process_group()
+            cls._processes_spawned = False
+        unittest.TestCase.tearDownClass()
+
+    def setUp(self):
+        unittest.TestCase.setUp(self)
+        self.__class__._ensure_processes_spawned()
+        torch.autograd._enable_record_function(False)
+
+    def tearDown(self):
+        from torch.distributed.tensor import _random as random
+
+        random._rng_tracker = None
+        unittest.TestCase.tearDown(self)
+        torch.autograd._enable_record_function(True)
+
+    def __init__(self, method_name="runTest", methodName="runTest"):
+        if methodName != "runTest":
+            method_name = methodName
+        unittest.TestCase.__init__(self, method_name)
 
     @property
-    def world_size(self) -> int:
-        return NUM_DEVICES
+    def rank(self):
+        return torch.SymInt(LocalIntNode({r: r for r in range(self.world_size)}))
 
+    @rank.setter
+    def rank(self, rank):
+        pass
+
+    def build_device_mesh(self) -> DeviceMesh:
+        with maybe_disable_local_tensor_mode():
+            return super().build_device_mesh()
+
+    def init_manual_seed_for_rank(self) -> None:
+        torch.manual_seed(0)
+
+
+class DTensorTestBase(DTensorTestMixin, MultiProcessTestCase):
     @property
-    def device_type(self) -> str:
-        # if enough GPU/XPU/HPU we can use those devices, otherwise we fallback to CPU
-        if (
-            not (TEST_CUDA or TEST_XPU or TEST_HPU or TEST_PRIVATEUSE1)
-            or DEVICE_COUNT < self.world_size
-        ):
-            return "cpu"
-        else:
-            return DEVICE_TYPE
+    def world_size(self) -> int:
+        return NUM_DEVICES
 
     @property
     def backend(self) -> str:
         backend = dist.get_default_backend_for_device(self.device_type)
         return backend
 
-    def init_manual_seed_for_rank(self) -> None:
-        torch.manual_seed(self.rank)
-
-    def build_device_mesh(self) -> DeviceMesh:
-        return init_device_mesh(self.device_type, (self.world_size,))
-
-    def init_pg(self, eager_init, backend: Optional[str] = None) -> None:
+    def init_pg(self, eager_init, backend: str | None = None) -> None:
         if backend is None:
             backend = self.backend
 
@@ -471,7 +842,7 @@ def init_pg(self, eager_init, backend: Optional[str] = None) -> None:
             device_id=device_id,
         )
 
-    def destroy_pg(self, device_id: Optional[int] = None) -> None:
+    def destroy_pg(self, device_id: int | None = None) -> None:
         # Wait for all ranks to reach here before starting shutdown.
         # FIXME dist.barrier deadlocks with multiple threads and NCCL: https://github.com/pytorch/pytorch/issues/95895
         # dist.all_reduce(torch.zeros((1,), device="cuda" if TEST_CUDA else "cpu"))
@@ -499,55 +870,16 @@ def setUp(self) -> None:
         super().setUp()
         self._spawn_processes()
 
-    def _test_op_on_dtensor(self, op_call, *args, **kwargs) -> None:
-        """
-        This function checks ``op_call(dtensor).full_tensor() == op_call(dtensor.full_tensor())``.
-        Unlike _test_op where the DTensor sharding is generated by DTensorConverter,
-        this function takes in DTensor object directly as argument and test the equality
-        of calling op on full_tensor() and DTensor.
-        """
-        # call full_tensor() on DTensor args/kwargs
-        args_flattened, args_spec = tree_flatten(args)
-        full_tensor_args_flattened = tuple(
-            arg.full_tensor().detach().clone() if isinstance(arg, DTensor) else arg
-            for arg in args_flattened
-        )
-        full_tensor_args = tree_unflatten(full_tensor_args_flattened, args_spec)
-        full_tensor_kwargs = {
-            k: v.full_tensor() if isinstance(v, DTensor) else v
-            for k, v in kwargs.items()
-        }
-
-        out_flattened, _ = tree_flatten(
-            op_call(*full_tensor_args, **full_tensor_kwargs)
-        )
-        d_out_flattened, _ = tree_flatten(op_call(*args, **kwargs))
-        d_out_full_tensor_flattened = [dt.full_tensor() for dt in d_out_flattened]
-        self.assertEqual(out_flattened, d_out_full_tensor_flattened)
-
-    # pyre-ignore[2]:
-    def _test_op(self, mesh: DeviceMesh, op_call, *args, **kwargs) -> None:
-        out = op_call(*args, **kwargs)
-        dtc = DTensorConverter(mesh, args, kwargs)
-        for d_args, d_kwargs in dtc:
-            # pyre can't find assertTrue anymore?
-            self.assertEqual(dtc.successful(), True)
-            d_out = op_call(*d_args, **d_kwargs)
-            self.assertEqual(d_out.full_tensor(), out)
-
-    def run_subtests(self, *args, **kwargs):
-        return run_subtests(self, *args, **kwargs)
-
 
 TestFunc = Callable[[...], object]
 
 
 # wrapper to initialize comms (processgroup)
 def with_comms(
-    eager_init: Union[TestFunc, bool] = False,
-    backend: Optional[str] = None,
+    eager_init: TestFunc | bool = False,
+    backend: str | None = None,
 ) -> TestFunc:
-    def decorator(func, eager_init: bool = False, backend: Optional[str] = None):
+    def decorator(func, eager_init: bool = False, backend: str | None = None):
         @wraps(func)  # pyre-ignore[6]
         def wrapper(
             self,
@@ -593,8 +925,21 @@ def build_device_mesh(self):
 
     def setUp(self) -> None:
         super().setUp()
+        # Enable thread-safe lock for ShardingPropagator since we run
+        # multi-threaded tests.
+        from torch.distributed.tensor._sharding_prop import ShardingPropagator
+
+        self._orig_fake_mode_lock = ShardingPropagator._fake_mode_lock
+        ShardingPropagator._fake_mode_lock = threading.Lock()
         self._spawn_threads()
 
+    def tearDown(self) -> None:
+        # Restore the original (no-op) lock
+        from torch.distributed.tensor._sharding_prop import ShardingPropagator
+
+        ShardingPropagator._fake_mode_lock = self._orig_fake_mode_lock
+        super().tearDown()
+
 
 # This is a class for converting args/kwargs of an op into distributed args/kwargs
 class DTensorConverter:
@@ -795,11 +1140,11 @@ def build_device_mesh(self) -> DeviceMesh:
         with maybe_disable_local_tensor_mode():
             return super().build_device_mesh()
 
-    def init_pg(self, eager_init, backend: Optional[str] = None) -> None:
+    def init_pg(self, eager_init, backend: str | None = None) -> None:
         dist.init_process_group("fake", rank=0, world_size=self.world_size)
         self._pg = dist.distributed_c10d._get_default_group()
 
-    def destroy_pg(self, device_id: Optional[int] = None) -> None:
+    def destroy_pg(self, device_id: int | None = None) -> None:
         dist.destroy_process_group(self._pg)
         self._pg = None
 
@@ -857,11 +1202,11 @@ def build_device_mesh(self) -> DeviceMesh:
         with maybe_disable_local_tensor_mode():
             return super().build_device_mesh()
 
-    def init_pg(self, eager_init, backend: Optional[str] = None) -> None:
+    def init_pg(self, eager_init, backend: str | None = None) -> None:
         dist.init_process_group("fake", rank=0, world_size=self.world_size)
         self._pg = dist.distributed_c10d._get_default_group()
 
-    def destroy_pg(self, device_id: Optional[int] = None) -> None:
+    def destroy_pg(self, device_id: int | None = None) -> None:
         dist.destroy_process_group(self._pg)
         self._pg = None
 
@@ -1104,8 +1449,61 @@ def _to_local_shard(a):
     # run and compare
     ref_output = op(*full_args, **full_kwargs)
     local_output = op(*local_args, **local_kwargs)
-    output_dt = DTensor.from_local(local_output, device_mesh, output_placements)
-    full_output = output_dt.redistribute(device_mesh, (Replicate(),)).to_local()
-    return ref_output.shape == full_output.shape and torch.allclose(
-        ref_output, full_output, atol=1e-5, rtol=1e-5
-    )
+
+    ref_tensors = [
+        t for t in pytree.tree_leaves(ref_output) if isinstance(t, torch.Tensor)
+    ]
+    local_tensors = [
+        t for t in pytree.tree_leaves(local_output) if isinstance(t, torch.Tensor)
+    ]
+
+    for ref, local, plc in zip(ref_tensors, local_tensors, output_placements):
+        dt = DTensor.from_local(local, device_mesh, (plc,))
+        full = dt.redistribute(device_mesh, (Replicate(),)).to_local()
+        if ref.shape != full.shape or not torch.allclose(
+            ref, full, atol=1e-5, rtol=1e-5, equal_nan=True
+        ):
+            return False
+    return True
+
+
+@contextlib.contextmanager
+def op_strategy_context(op_overload, strategy_func, schema_info=None):
+    """
+    Context manager for setting and clearing op strategies.
+    Args:
+        op_overload: The operator overload to set or clear the strategy for.
+        strategy_func: The strategy function to set for the operator overload.
+        schema_info: Optional schema information for the operator overload.
+    Yields:
+        None
+    """
+    from torch.distributed.tensor._ops.utils import register_op_strategy
+    from torch.distributed.tensor.debug import _clear_sharding_prop_cache
+
+    propagator = DTensor._op_dispatcher.sharding_propagator
+    _origin_op_strategy_funcs = None
+    _origin_op_strategy_schema = None
+    try:
+        # register the op strategy
+        if op_overload in propagator.op_strategy_funcs:
+            _origin_op_strategy_funcs = propagator.op_strategy_funcs[op_overload]
+            del propagator.op_strategy_funcs[op_overload]
+        if op_overload in propagator.op_to_schema_info:
+            _origin_op_strategy_schema = propagator.op_to_schema_info[op_overload]
+            del propagator.op_to_schema_info[op_overload]
+        register_op_strategy(op_overload, schema_info=schema_info)(strategy_func)
+        yield
+    finally:
+        # clear this op strategy cache
+        if _origin_op_strategy_funcs is None:
+            if op_overload in propagator.op_strategy_funcs:
+                del propagator.op_strategy_funcs[op_overload]
+        else:
+            propagator.op_strategy_funcs[op_overload] = _origin_op_strategy_funcs
+        if _origin_op_strategy_schema is None:
+            if op_overload in propagator.op_to_schema_info:
+                del propagator.op_to_schema_info[op_overload]
+        else:
+            propagator.op_to_schema_info[op_overload] = _origin_op_strategy_schema
+        _clear_sharding_prop_cache()
diff --git a/torch/testing/_internal/distributed/checkpoint_utils.py b/torch/testing/_internal/distributed/checkpoint_utils.py
index 29c39b261052c..91e79a5dbf1c2 100644
--- a/torch/testing/_internal/distributed/checkpoint_utils.py
+++ b/torch/testing/_internal/distributed/checkpoint_utils.py
@@ -9,7 +9,7 @@
 import tempfile
 from collections.abc import Callable
 from functools import wraps
-from typing import Any, cast, IO, Optional
+from typing import Any, cast, IO
 
 # introduced as collections.abc.Buffer in Python 3.12
 from typing_extensions import Buffer
@@ -64,7 +64,7 @@ def __init__(self, output: IO[bytes]) -> None:
             def writeable(self) -> bool:
                 return True
 
-            def write(self, b: Buffer) -> Optional[int]:
+            def write(self, b: Buffer) -> int | None:
                 # Don't mutate the input
                 chunk = bytearray(b)
                 Rot13Example._rot13bytes(chunk, len(chunk))
@@ -83,7 +83,7 @@ def __init__(self, input: IO[bytes]) -> None:
             def readable(self) -> bool:
                 return True
 
-            def readinto(self, b: Buffer) -> Optional[int]:
+            def readinto(self, b: Buffer) -> int | None:
                 if hasattr(self.input, "readinto"):
                     count = self.input.readinto(b)
                 else:
@@ -122,8 +122,8 @@ def get_test_extension_registry() -> ExtensionRegistry:
 
 
 def with_temp_dir(
-    func: Optional[Callable] = None,
-) -> Optional[Callable]:
+    func: Callable | None = None,
+) -> Callable | None:
     """
     Wrapper to initialize temp directory for distributed checkpoint.
     """
@@ -163,10 +163,10 @@ def wrapper(self, *args: tuple[object], **kwargs: dict[str, Any]) -> None:
 
 
 def with_checkpoint_logging(
-    func: Optional[Callable] = None,
+    func: Callable | None = None,
     logger_name: str = "torch.distributed.checkpoint",
     level: int = logging.INFO,
-) -> Optional[Callable]:
+) -> Callable | None:
     """
     Wrapper to configure checkpoint logging for distributed tests.
 
diff --git a/torch/testing/_internal/distributed/distributed_test.py b/torch/testing/_internal/distributed/distributed_test.py
index 757717253adad..3f96956170976 100644
--- a/torch/testing/_internal/distributed/distributed_test.py
+++ b/torch/testing/_internal/distributed/distributed_test.py
@@ -18,7 +18,7 @@
 from dataclasses import dataclass
 from datetime import timedelta
 from functools import reduce
-from typing import Any, NamedTuple, Union
+from typing import Any, NamedTuple
 
 import numpy as np
 
@@ -54,7 +54,13 @@
     _verify_param_shape_across_processes,
 )
 from torch.nn.parallel import DistributedDataParallel
-from torch.nn.parallel.distributed import _dump_DDP_relevant_env_vars, _MixedPrecision
+from torch.nn.parallel.distributed import (
+    _BucketCapacityConfig,
+    _DEFAULT_BUCKET_CAP_MB,
+    _dump_DDP_relevant_env_vars,
+    _MB_TO_BYTES,
+    _MixedPrecision,
+)
 from torch.profiler import ExecutionTraceObserver, ProfilerActivity
 from torch.testing._internal.common_distributed import (
     captured_output,
@@ -192,9 +198,6 @@ class TestNamedTupleInput_1(NamedTuple):
 BACKEND = os.environ["BACKEND"]
 INIT_METHOD = os.getenv("INIT_METHOD", "env://")
 
-DEFAULT_TIMEOUT = 300
-CUSTOMIZED_TIMEOUT = {"test_DistributedDataParallel": 500}
-
 
 def get_profiling_event(event_name, profiler, dedup_gpu_user_annotation=False):
     event_list = (
@@ -251,7 +254,7 @@ def get_profiler_nccl_meta(prof):
 class DDPUnevenTestInput(NamedTuple):
     name: str
     model: nn.Module
-    inp: Union[torch.tensor, tuple]
+    inp: torch.Tensor | tuple
     sync_interval: int
     throw_on_early_termination: bool = False
     hook: Callable = None
@@ -396,14 +399,6 @@ def forward(self, x):
             return F.relu(self.lin1(x))
 
 
-def get_timeout(test_id):
-    test_name = test_id.split(".")[-1]
-    if test_name in CUSTOMIZED_TIMEOUT:
-        return CUSTOMIZED_TIMEOUT[test_name]
-    else:
-        return DEFAULT_TIMEOUT
-
-
 default_pg_timeout = 60
 
 CUSTOM_PG_TIMEOUT = {
@@ -899,6 +894,24 @@ def test_barrier_timeout_full_group(self):
             if group_id is not None:
                 self._test_barrier_timeout(group_id, timeout)
 
+        @skip_but_pass_in_sandcastle_if(
+            BACKEND != "gloo", "Only gloo backend supports timeouts"
+        )
+        def test_barrier_timeout_arg(self):
+            """Test that the timeout argument to barrier() overrides PG default.
+
+            Create a PG with a large default timeout, then have only rank 0
+            call barrier with a tiny timeout. The barrier should time out using
+            the per-call timeout (1ms) rather than the PG default (300s).
+            """
+            pg = dist.new_group(timeout=timedelta(seconds=300))
+
+            if dist.get_rank() == 0:
+                with self.assertRaisesRegex(RuntimeError, "Timed out waiting 1ms"):
+                    dist.barrier(group=pg, timeout=timedelta(seconds=0.001))
+
+            dist.destroy_process_group(pg)
+
         @skip_but_pass_in_sandcastle_if(
             BACKEND not in DistTestCases.backend_feature["subgroup"],
             f"The {BACKEND} backend does not support creating subgroups on CUDA devices",
@@ -1012,7 +1025,7 @@ def test_new_subgroups_by_enumeration_input_rank_exceeds_world_size(self):
 
             with self.assertRaisesRegex(
                 ValueError,
-                "The new group's rank should be within the world_size set by init_process_group",
+                f"Rank {world_size} is out of range",
             ):
                 dist.new_subgroups_by_enumeration(
                     ranks_per_subgroup_list=[[0, 1], [world_size, 2]]
@@ -1028,7 +1041,7 @@ def test_new_subgroups_by_enumeration_negative_input_rank(self):
 
             with self.assertRaisesRegex(
                 ValueError,
-                "The new group's rank should be within the world_size set by init_process_group",
+                r"Rank -\d+ is out of range",
             ):
                 dist.new_subgroups_by_enumeration(
                     ranks_per_subgroup_list=[[-1, -2], [-3, -4]]
@@ -10491,4 +10504,333 @@ def test_skip_all_reduce_unused_parameters(self):
                 self.assertEqual(i, j)
 
 
+# =============================================================================
+# Unit tests for _BucketCapacityConfig dataclass
+# These tests verify pure Python logic without requiring multi-process setup.
+# =============================================================================
+class TestBucketCapacityConfig(unittest.TestCase):
+    """Unit tests for _BucketCapacityConfig dataclass."""
+
+    def test_create_with_default_bucket_cap(self):
+        """Test _BucketCapacityConfig.create() with default bucket_cap_mb (None)."""
+        config = _BucketCapacityConfig.create(
+            bucket_cap_mb=None,
+            bucket_cap_mb_list=None,
+            use_python_reducer=False,
+        )
+
+        expected_bytes = _DEFAULT_BUCKET_CAP_MB * _MB_TO_BYTES
+        self.assertEqual(config.bucket_bytes_cap, expected_bytes)
+        self.assertEqual(config.per_bucket_bytes_caps, ())
+        self.assertEqual(
+            config.first_bucket_bytes_cap, dist._DEFAULT_FIRST_BUCKET_BYTES
+        )
+        self.assertEqual(config.bucket_cap_mb, _DEFAULT_BUCKET_CAP_MB)
+
+    def test_create_with_custom_bucket_cap_mb(self):
+        """Test _BucketCapacityConfig.create() with custom bucket_cap_mb."""
+        custom_cap_mb = 100
+        config = _BucketCapacityConfig.create(
+            bucket_cap_mb=custom_cap_mb,
+            bucket_cap_mb_list=None,
+            use_python_reducer=False,
+        )
+
+        expected_bytes = custom_cap_mb * _MB_TO_BYTES
+        self.assertEqual(config.bucket_bytes_cap, expected_bytes)
+        self.assertEqual(config.per_bucket_bytes_caps, ())
+        self.assertEqual(config.first_bucket_bytes_cap, expected_bytes)
+        self.assertEqual(config.bucket_cap_mb, custom_cap_mb)
+
+    def test_create_with_bucket_cap_mb_list(self):
+        """Test _BucketCapacityConfig.create() with bucket_cap_mb_list."""
+        cap_list = [10, 25, 50]
+        config = _BucketCapacityConfig.create(
+            bucket_cap_mb=None,
+            bucket_cap_mb_list=cap_list,
+            use_python_reducer=False,
+        )
+
+        expected_list = tuple(cap * _MB_TO_BYTES for cap in cap_list)
+        self.assertEqual(config.per_bucket_bytes_caps, expected_list)
+        expected_max_bytes = max(cap_list) * _MB_TO_BYTES
+        self.assertEqual(config.bucket_bytes_cap, expected_max_bytes)
+        self.assertEqual(config.first_bucket_bytes_cap, expected_max_bytes)
+        self.assertTrue(config.has_custom_per_bucket_caps)
+
+    def test_create_with_bucket_cap_mb_list_overrides_bucket_cap_mb(self):
+        """Test that bucket_cap_mb_list takes precedence over bucket_cap_mb."""
+        cap_list = [10, 25, 50]
+        config = _BucketCapacityConfig.create(
+            bucket_cap_mb=100,
+            bucket_cap_mb_list=cap_list,
+            use_python_reducer=False,
+        )
+
+        expected_max_bytes = max(cap_list) * _MB_TO_BYTES
+        self.assertEqual(config.bucket_bytes_cap, expected_max_bytes)
+
+    def test_create_raises_for_python_reducer_with_list(self):
+        """Test that AssertionError is raised when using bucket_cap_mb_list with Python reducer."""
+        with self.assertRaises(AssertionError) as context:
+            _BucketCapacityConfig.create(
+                bucket_cap_mb=None,
+                bucket_cap_mb_list=[10, 25, 50],
+                use_python_reducer=True,
+            )
+
+        self.assertIn("python reducer is not supported", str(context.exception))
+
+    def test_has_custom_per_bucket_caps_property(self):
+        """Test has_custom_per_bucket_caps property."""
+        config_no_list = _BucketCapacityConfig.create(
+            bucket_cap_mb=25,
+            bucket_cap_mb_list=None,
+            use_python_reducer=False,
+        )
+        self.assertFalse(config_no_list.has_custom_per_bucket_caps)
+
+        config_with_list = _BucketCapacityConfig.create(
+            bucket_cap_mb=None,
+            bucket_cap_mb_list=[10, 25],
+            use_python_reducer=False,
+        )
+        self.assertTrue(config_with_list.has_custom_per_bucket_caps)
+
+    def test_bucket_cap_mb_property(self):
+        """Test bucket_cap_mb property correctly derives MiB from bytes."""
+        config = _BucketCapacityConfig.create(
+            bucket_cap_mb=100,
+            bucket_cap_mb_list=None,
+            use_python_reducer=False,
+        )
+        self.assertEqual(config.bucket_cap_mb, 100)
+
+        config_default = _BucketCapacityConfig.create(
+            bucket_cap_mb=None,
+            bucket_cap_mb_list=None,
+            use_python_reducer=False,
+        )
+        self.assertEqual(config_default.bucket_cap_mb, _DEFAULT_BUCKET_CAP_MB)
+
+    def test_immutability(self):
+        """Test that _BucketCapacityConfig is immutable (frozen dataclass)."""
+        config = _BucketCapacityConfig.create(
+            bucket_cap_mb=25,
+            bucket_cap_mb_list=None,
+            use_python_reducer=False,
+        )
+
+        with self.assertRaises(AttributeError):
+            config.bucket_bytes_cap = 100
+
+        with self.assertRaises(AttributeError):
+            config.per_bucket_bytes_caps = (100,)
+
+
+class TestBucketCapacityConfigComputeLimits(unittest.TestCase):
+    """Unit tests for _BucketCapacityConfig.compute_bucket_size_limits()."""
+
+    def test_compute_limits_with_custom_per_bucket_caps(self):
+        """Test compute_bucket_size_limits returns custom caps when provided."""
+        cap_list = [10, 25, 50]
+        config = _BucketCapacityConfig.create(
+            bucket_cap_mb=None,
+            bucket_cap_mb_list=cap_list,
+            use_python_reducer=False,
+        )
+
+        limits, rebuild_limits = config.compute_bucket_size_limits(
+            static_graph=False,
+            find_unused_parameters=True,
+        )
+
+        expected_limits = [cap * _MB_TO_BYTES for cap in cap_list]
+        self.assertEqual(limits, expected_limits)
+        self.assertEqual(rebuild_limits, expected_limits)
+
+    def test_compute_limits_static_graph_uses_maxsize(self):
+        """Test that static_graph=True uses sys.maxsize for initial bucketing."""
+        config = _BucketCapacityConfig.create(
+            bucket_cap_mb=25,
+            bucket_cap_mb_list=None,
+            use_python_reducer=False,
+        )
+
+        limits, rebuild_limits = config.compute_bucket_size_limits(
+            static_graph=True,
+            find_unused_parameters=False,
+        )
+
+        self.assertEqual(limits, [sys.maxsize])
+        self.assertEqual(rebuild_limits, [])
+
+    def test_compute_limits_find_unused_false_uses_maxsize(self):
+        """Test that find_unused_parameters=False uses sys.maxsize."""
+        config = _BucketCapacityConfig.create(
+            bucket_cap_mb=25,
+            bucket_cap_mb_list=None,
+            use_python_reducer=False,
+        )
+
+        limits, rebuild_limits = config.compute_bucket_size_limits(
+            static_graph=False,
+            find_unused_parameters=False,
+        )
+
+        self.assertEqual(limits, [sys.maxsize])
+        self.assertEqual(rebuild_limits, [])
+
+    def test_compute_limits_default_with_find_unused_true(self):
+        """Test default config with find_unused_parameters=True uses smaller first bucket."""
+        config = _BucketCapacityConfig.create(
+            bucket_cap_mb=None,
+            bucket_cap_mb_list=None,
+            use_python_reducer=False,
+        )
+
+        limits, rebuild_limits = config.compute_bucket_size_limits(
+            static_graph=False,
+            find_unused_parameters=True,
+        )
+
+        expected_first = dist._DEFAULT_FIRST_BUCKET_BYTES
+        expected_main = _DEFAULT_BUCKET_CAP_MB * _MB_TO_BYTES
+        self.assertEqual(limits, [expected_first, expected_main])
+        self.assertEqual(rebuild_limits, [])
+
+    def test_compute_limits_custom_cap_with_find_unused_true(self):
+        """Test custom bucket_cap_mb with find_unused_parameters=True uses uniform bucket."""
+        config = _BucketCapacityConfig.create(
+            bucket_cap_mb=100,
+            bucket_cap_mb_list=None,
+            use_python_reducer=False,
+        )
+
+        limits, rebuild_limits = config.compute_bucket_size_limits(
+            static_graph=False,
+            find_unused_parameters=True,
+        )
+
+        expected_bytes = 100 * _MB_TO_BYTES
+        self.assertEqual(limits, [expected_bytes])
+        self.assertEqual(rebuild_limits, [])
+
+    def test_compute_limits_empty_rebuild_for_non_list_config(self):
+        """Test that rebuild_limits is empty when not using bucket_cap_mb_list."""
+        config = _BucketCapacityConfig.create(
+            bucket_cap_mb=50,
+            bucket_cap_mb_list=None,
+            use_python_reducer=False,
+        )
+
+        _, rebuild_limits = config.compute_bucket_size_limits(
+            static_graph=False,
+            find_unused_parameters=True,
+        )
+
+        self.assertEqual(rebuild_limits, [])
+
+
+class TestBucketCapacityConfigEdgeCases(unittest.TestCase):
+    """Edge case tests for _BucketCapacityConfig."""
+
+    def test_single_element_bucket_cap_mb_list(self):
+        """Test bucket_cap_mb_list with a single element."""
+        config = _BucketCapacityConfig.create(
+            bucket_cap_mb=None,
+            bucket_cap_mb_list=[50],
+            use_python_reducer=False,
+        )
+
+        self.assertEqual(config.per_bucket_bytes_caps, (50 * _MB_TO_BYTES,))
+        self.assertEqual(config.bucket_bytes_cap, 50 * _MB_TO_BYTES)
+        self.assertTrue(config.has_custom_per_bucket_caps)
+
+    def test_large_bucket_cap_mb(self):
+        """Test with large bucket_cap_mb value."""
+        large_cap = 1000
+        config = _BucketCapacityConfig.create(
+            bucket_cap_mb=large_cap,
+            bucket_cap_mb_list=None,
+            use_python_reducer=False,
+        )
+
+        expected_bytes = large_cap * _MB_TO_BYTES
+        self.assertEqual(config.bucket_bytes_cap, expected_bytes)
+        self.assertEqual(config.first_bucket_bytes_cap, expected_bytes)
+
+    def test_small_bucket_cap_mb(self):
+        """Test with small bucket_cap_mb value."""
+        small_cap = 1
+        config = _BucketCapacityConfig.create(
+            bucket_cap_mb=small_cap,
+            bucket_cap_mb_list=None,
+            use_python_reducer=False,
+        )
+
+        expected_bytes = small_cap * _MB_TO_BYTES
+        self.assertEqual(config.bucket_bytes_cap, expected_bytes)
+
+    def test_empty_bucket_cap_mb_list(self):
+        """Test with empty bucket_cap_mb_list (should use default)."""
+        config = _BucketCapacityConfig.create(
+            bucket_cap_mb=None,
+            bucket_cap_mb_list=[],
+            use_python_reducer=False,
+        )
+
+        expected_bytes = _DEFAULT_BUCKET_CAP_MB * _MB_TO_BYTES
+        self.assertEqual(config.bucket_bytes_cap, expected_bytes)
+        self.assertEqual(config.per_bucket_bytes_caps, ())
+        self.assertFalse(config.has_custom_per_bucket_caps)
+
+    def test_bucket_cap_mb_list_with_duplicates(self):
+        """Test bucket_cap_mb_list with duplicate values."""
+        cap_list = [25, 25, 25]
+        config = _BucketCapacityConfig.create(
+            bucket_cap_mb=None,
+            bucket_cap_mb_list=cap_list,
+            use_python_reducer=False,
+        )
+
+        expected_list = tuple(25 * _MB_TO_BYTES for _ in range(3))
+        self.assertEqual(config.per_bucket_bytes_caps, expected_list)
+        self.assertEqual(config.bucket_bytes_cap, 25 * _MB_TO_BYTES)
+
+    def test_bucket_cap_mb_list_ascending_order(self):
+        """Test bucket_cap_mb_list in ascending order."""
+        cap_list = [5, 10, 15, 20]
+        config = _BucketCapacityConfig.create(
+            bucket_cap_mb=None,
+            bucket_cap_mb_list=cap_list,
+            use_python_reducer=False,
+        )
+
+        self.assertEqual(config.bucket_bytes_cap, 20 * _MB_TO_BYTES)
+
+    def test_bucket_cap_mb_list_descending_order(self):
+        """Test bucket_cap_mb_list in descending order."""
+        cap_list = [100, 50, 25, 10]
+        config = _BucketCapacityConfig.create(
+            bucket_cap_mb=None,
+            bucket_cap_mb_list=cap_list,
+            use_python_reducer=False,
+        )
+
+        self.assertEqual(config.bucket_bytes_cap, 100 * _MB_TO_BYTES)
+
+
+class TestBucketCapacityConfigConstants(unittest.TestCase):
+    """Tests for module-level constants."""
+
+    def test_default_bucket_cap_mb_value(self):
+        """Test that _DEFAULT_BUCKET_CAP_MB is 25."""
+        self.assertEqual(_DEFAULT_BUCKET_CAP_MB, 25)
+
+    def test_mb_to_bytes_value(self):
+        """Test that _MB_TO_BYTES is 1024 * 1024."""
+        self.assertEqual(_MB_TO_BYTES, 1024 * 1024)
+
+
 instantiate_parametrized_tests(DistributedTest._DistTestBase)
diff --git a/torch/testing/_internal/distributed/multi_threaded_pg.py b/torch/testing/_internal/distributed/multi_threaded_pg.py
index ffd989444d0f2..a60a1479a91a6 100644
--- a/torch/testing/_internal/distributed/multi_threaded_pg.py
+++ b/torch/testing/_internal/distributed/multi_threaded_pg.py
@@ -5,7 +5,6 @@
 import weakref
 from dataclasses import dataclass
 from functools import partial, reduce
-from typing import Optional, Union
 
 import torch
 import torch.distributed as dist
@@ -132,7 +131,7 @@ def work(self, data):
     def _size_cumsum(
         self,
         buf_size: int,
-        sizes: Union[torch.Tensor, list[int], None],
+        sizes: torch.Tensor | list[int] | None,
         world_size: int,
     ) -> torch.Tensor:
         if sizes is None or len(sizes) == 0:
@@ -403,8 +402,8 @@ def alltoall_base(
         self,
         output_buffer: torch.Tensor,
         input_buffer: torch.Tensor,
-        output_split_sizes: Optional[list[int]],
-        input_split_sizes: Optional[list[int]],
+        output_split_sizes: list[int] | None,
+        input_split_sizes: list[int] | None,
         opts=AllToAllOptions(),
     ) -> torch.Tensor:
         coll = ProcessLocalGroup._start_coll(AllToAllBase(), self)
@@ -552,14 +551,15 @@ def _create_threaded_pg(prefix_store, rank, world_size, timeout):
 @dataclass
 class WorldData:
     default_pg: dist.ProcessGroup
-    pg_map: dict[dist.ProcessGroup, tuple[str, Optional[Store]]]
+    pg_map: dict[dist.ProcessGroup, tuple[str, Store | None]]
     pg_names: dict[dist.ProcessGroup, str]
     pg_group_ranks: dict[dist.ProcessGroup, dict[int, int]]
     pg_backend_config: dict[dist.ProcessGroup, str]
     group_count: int
     tags_to_pg: dict[str, list[dist.ProcessGroup]]
     pg_to_tag: dict[dist.ProcessGroup, str]
-    pg_coalesce_state: dict[dist.ProcessGroup, list[Union[_CollOp, P2POp]]]
+    pg_coalesce_state: dict[dist.ProcessGroup, list[_CollOp | P2POp]]
+    comms: list
 
 
 class ThreadLocalWorld:
@@ -568,7 +568,7 @@ class ThreadLocalWorld:
     def _get_world(self) -> WorldData:
         if not hasattr(ThreadLocalWorld._world, "world"):
             ThreadLocalWorld._world.world = WorldData(
-                None, {}, {}, {}, {}, 0, {}, {}, {}
+                None, {}, {}, {}, {}, 0, {}, {}, {}, []
             )
         return ThreadLocalWorld._world.world
 
@@ -613,9 +613,13 @@ def pg_to_tag(self):
         return self._get_world().pg_to_tag
 
     @property
-    def pg_coalesce_state(self) -> dict[dist.ProcessGroup, list[Union[_CollOp, P2POp]]]:
+    def pg_coalesce_state(self) -> dict[dist.ProcessGroup, list[_CollOp | P2POp]]:
         return self._get_world().pg_coalesce_state
 
+    @property
+    def comms(self):
+        return self._get_world().comms
+
 
 _old_pg_world = None
 _ctx_manager = None
@@ -632,4 +636,9 @@ def _install_threaded_pg():
 
 
 def _uninstall_threaded_pg():
+    global _ctx_manager
     dist.distributed_c10d._world = _old_pg_world
+    # Restore autograd multithreading state that was disabled in _install_threaded_pg
+    if _ctx_manager is not None:
+        _ctx_manager.__exit__(None, None, None)
+        _ctx_manager = None
diff --git a/torch/testing/_internal/distributed/rpc/rpc_test.py b/torch/testing/_internal/distributed/rpc/rpc_test.py
index e81994ce760cb..e3b2452d5eea6 100644
--- a/torch/testing/_internal/distributed/rpc/rpc_test.py
+++ b/torch/testing/_internal/distributed/rpc/rpc_test.py
@@ -6015,11 +6015,12 @@ def test_owner_rref_forward_synchronization4(self):
 
     @staticmethod
     def _return_tensor_view(i):
-        x = torch.ones(1000, 200).cuda(0) * i
-        torch.cuda._sleep(10 * FIFTY_MIL_CYCLES)
-        # serialization of the return value will create a new tensor from the
-        # view, which is done outside of the user function.
-        return x.split(100)[0]
+        with torch.cuda.stream(torch.cuda.current_stream(0)):
+            x = torch.ones(1000, 200).cuda(0) * i
+            torch.cuda._sleep(10 * FIFTY_MIL_CYCLES)
+            # serialization of the return value will create a new tensor from the
+            # view, which is done outside of the user function.
+            return x.split(100)[0]
 
     @skip_if_lt_x_gpu(1)
     def test_tensor_view_as_return_value(self):
diff --git a/torch/testing/_internal/dynamo_test_failures.py b/torch/testing/_internal/dynamo_test_failures.py
index cdc69b7920cf0..f56a7b07f63a3 100644
--- a/torch/testing/_internal/dynamo_test_failures.py
+++ b/torch/testing/_internal/dynamo_test_failures.py
@@ -24,10 +24,9 @@
 import logging
 import os
 import sys
-from typing import Optional
 
 
-def find_test_dir() -> Optional[str]:
+def find_test_dir() -> str | None:
     # Find the path to the dynamo expected failure and skip files.
     from os.path import abspath, basename, dirname, exists, join, normpath
 
diff --git a/torch/testing/_internal/fake_config_module.py b/torch/testing/_internal/fake_config_module.py
index 1e93c41de72a7..0dc2a5d7df406 100644
--- a/torch/testing/_internal/fake_config_module.py
+++ b/torch/testing/_internal/fake_config_module.py
@@ -1,5 +1,4 @@
 import sys
-from typing import Optional
 
 from torch.utils._config_module import Config, install_config_module
 
@@ -12,8 +11,8 @@
 e_set = {1}
 e_tuple = (1,)
 e_dict = {1: 2}
-e_none: Optional[bool] = None
-e_optional: Optional[bool] = True
+e_none: bool | None = None
+e_optional: bool | None = True
 e_ignored = True
 _e_ignored = True
 magic_cache_config_ignored = True
@@ -32,6 +31,15 @@
 e_aliased_bool: bool = Config(
     alias="torch.testing._internal.fake_config_module2.e_aliasing_bool"
 )
+e_deprecated: bool = Config(
+    default=True, deprecated=True, deprecation_message="is no longer needed"
+)
+e_not_deprecated: bool = Config(default=False)
+e_deprecated_alias: bool = Config(
+    alias="torch.testing._internal.fake_config_module.e_not_deprecated",
+    deprecated=True,
+    deprecation_message="use something else instead",
+)
 
 
 class nested:
diff --git a/torch/testing/_internal/fake_config_module3.py b/torch/testing/_internal/fake_config_module3.py
index ff4118438e74c..41e2c3fddba2d 100644
--- a/torch/testing/_internal/fake_config_module3.py
+++ b/torch/testing/_internal/fake_config_module3.py
@@ -1,11 +1,11 @@
 import sys
-from typing import Callable, Optional  # noqa: UP035
+from typing import Callable  # noqa: UP035
 
 from torch.utils._config_module import install_config_module
 
 
 e_list = [1]
 e_set = {1}
-e_func: Optional[Callable] = None
+e_func: Callable | None = None
 
 install_config_module(sys.modules[__name__])
diff --git a/torch/testing/_internal/hop_db.py b/torch/testing/_internal/hop_db.py
index ae7c263a38d90..80a249b8d909f 100644
--- a/torch/testing/_internal/hop_db.py
+++ b/torch/testing/_internal/hop_db.py
@@ -5,8 +5,16 @@
 
 import torch
 from functorch.experimental.control_flow import map
-from torch.nn.attention.flex_attention import _create_empty_block_mask, flex_attention
+from torch._higher_order_ops.flex_attention import (
+    flex_attention as flex_attention_hop,
+)
+from torch.nn.attention.flex_attention import (
+    _create_empty_block_mask,
+    create_block_mask,
+    flex_attention,
+)
 from torch.testing import make_tensor
+from torch._higher_order_ops.inline_asm_elementwise import inline_asm_elementwise
 from torch.testing._internal.common_device_type import onlyCUDA
 from torch.testing._internal.common_dtype import all_types_and, custom_types
 from torch.testing._internal.opinfo.core import DecorateInfo, OpInfo, SampleInput
@@ -77,6 +85,7 @@ def f2(x, y0, y1):
     "autograd_function_apply",
     "run_and_save_rng_state",
     "run_with_rng_state",
+    "run_dtensor_rng_op",
     "graphsafe_run_with_rng_state",
     "out_dtype",
     "trace_wrapped",
@@ -194,6 +203,99 @@ def score_mod(score, b, h, m, n):
     yield SampleInput(q, k, v, score_mod, block_mask)
 
 
+def sample_inputs_flex_attention_backward(
+    opinfo, device, dtype, requires_grad, **kwargs
+):
+    make_arg = functools.partial(
+        make_tensor, device=device, dtype=dtype, requires_grad=False
+    )
+
+    def score_mod(score, b, h, m, n):
+        return score
+
+    def mask_mod(b, h, m, n):
+        return m >= n
+
+    q, k, v = (make_arg(2, 2, 128, 16, low=0.1, high=2) for _ in range(3))
+    block_mask = create_block_mask(mask_mod, B=2, H=2, Q_LEN=128, KV_LEN=128, device=device)
+    scale = 1.0 / q.size(-1) ** 0.5
+    out, logsumexp, _ = flex_attention_hop(
+        q, k, v, score_mod, block_mask.as_tuple(), scale, {},
+    )
+    yield SampleInput(
+        q,
+        args=(
+            k, v, out.detach(), logsumexp.detach(), torch.rand_like(out), None,
+            score_mod, None, block_mask.as_tuple(),
+            scale, {}, (), (),
+        ),
+    )
+
+
+def sample_inputs_flex_attention_backward_explicit_buffers(
+    opinfo, device, dtype, requires_grad, **kwargs
+):
+    make_arg = functools.partial(
+        make_tensor, device=device, dtype=dtype, requires_grad=False
+    )
+    mask_offset = torch.full((), 128, device=device, dtype=torch.int32)
+
+    def score_mod(score, b, h, m, n):
+        return score
+
+    def mask_mod(b, h, m, n):
+        return m + mask_offset >= n
+
+    q, k, v = (make_arg(2, 2, 128, 16, low=0.1, high=2) for _ in range(3))
+    block_mask = create_block_mask(mask_mod, B=2, H=2, Q_LEN=128, KV_LEN=128, device=device)
+    scale = 1.0 / q.size(-1) ** 0.5
+    out, logsumexp, _ = flex_attention_hop(
+        q, k, v, score_mod, block_mask.as_tuple(), scale, {},
+    )
+    yield SampleInput(
+        q,
+        args=(
+            k, v, out.detach(), logsumexp.detach(), torch.rand_like(out), None,
+            score_mod, None, block_mask.as_tuple(),
+            scale, {}, (), (),
+        ),
+    )
+
+
+def simple_flex_attention_backward(
+    query,
+    key,
+    value,
+    out,
+    logsumexp,
+    grad_out,
+    grad_logsumexp,
+    fw_graph,
+    joint_graph,
+    block_mask,
+    scale,
+    kernel_options,
+    score_mod_other_buffers,
+    mask_mod_other_buffers,
+):
+    return torch.ops.higher_order.flex_attention_backward(
+        query,
+        key,
+        value,
+        out,
+        logsumexp,
+        grad_out,
+        grad_logsumexp,
+        fw_graph,
+        joint_graph,
+        block_mask,
+        scale,
+        kernel_options,
+        score_mod_other_buffers,
+        mask_mod_other_buffers,
+    )
+
+
 def sample_inputs_while_loop(opinfo, device, dtype, requires_grad, **kwargs):
     make_arg = functools.partial(
         make_tensor, device=device, dtype=dtype, requires_grad=False
@@ -291,6 +393,27 @@ def fn(x):
     return invoke_quant_packed(fn, x)[0] * 2.0
 
 
+def sample_inputs_inline_asm(opinfo, device, dtype, requires_grad, **kwargs):
+    make_arg = functools.partial(
+        make_tensor, device=device, dtype=dtype, requires_grad=requires_grad
+    )
+    yield SampleInput(make_arg(2, 2, 2, low=0.1, high=2))
+
+
+def simple_inline_asm(x):
+    if torch.version.hip:
+        return inline_asm_elementwise(
+            x,
+            asm_str="v_mov_b32_e32 $0, $1",
+            constraints="=v, v",
+            dtype=torch.float32,
+        )
+
+    return inline_asm_elementwise(
+        x, asm_str="mov.f32 $0, $1;", constraints="=f,f", dtype=torch.float32
+    )
+
+
 hop_db = [
     OpInfo(
         name="scan",
@@ -475,14 +598,39 @@ def fn(x):
     OpInfo(
         name="flex_attention_backward",
         variant_test_name="simple",
-        op=flex_attention,
-        sample_inputs_func=sample_inputs_flex_attention,
+        op=simple_flex_attention_backward,
+        sample_inputs_func=sample_inputs_flex_attention_backward,
         dtypes=custom_types(torch.float16, torch.float32),
         supports_out=False,
         check_batched_grad=False,
         check_batched_gradgrad=False,
         check_batched_forward_grad=False,
         check_inplace_batched_forward_grad=False,
+        supports_autograd=False,
+        supports_gradgrad=False,
+        skips=(
+            DecorateInfo(unittest.expectedFailure, "TestHOP", "test_aot_export"),
+            DecorateInfo(
+                unittest.expectedFailure, "TestHOP", "test_pre_dispatch_export"
+            ),
+            DecorateInfo(unittest.expectedFailure, "TestHOP", "test_serialize_export"),
+            DecorateInfo(unittest.expectedFailure, "TestHOP", "test_retrace_export"),
+        ),
+        decorators=[onlyCUDA],
+    ),
+    OpInfo(
+        name="flex_attention_backward",
+        variant_test_name="explicit_buffers",
+        op=simple_flex_attention_backward,
+        sample_inputs_func=sample_inputs_flex_attention_backward_explicit_buffers,
+        dtypes=custom_types(torch.float16, torch.float32),
+        supports_out=False,
+        check_batched_grad=False,
+        check_batched_gradgrad=False,
+        check_batched_forward_grad=False,
+        check_inplace_batched_forward_grad=False,
+        supports_autograd=False,
+        supports_gradgrad=False,
         skips=(
             DecorateInfo(unittest.expectedFailure, "TestHOP", "test_aot_export"),
             DecorateInfo(
@@ -519,4 +667,18 @@ def fn(x):
             ),
         ],
     ),
+    OpInfo(
+        name="inline_asm_elementwise",
+        variant_test_name="simple",
+        op=simple_inline_asm,
+        sample_inputs_func=sample_inputs_inline_asm,
+        dtypes=custom_types(torch.float32),
+        supports_out=False,
+        check_batched_grad=False,
+        check_batched_gradgrad=False,
+        check_batched_forward_grad=False,
+        check_inplace_batched_forward_grad=False,
+        supports_autograd=False,
+        decorators=[onlyCUDA],
+    ),
 ]
diff --git a/torch/testing/_internal/inductor_utils.py b/torch/testing/_internal/inductor_utils.py
index 43269268660b4..dae4c952c6a81 100644
--- a/torch/testing/_internal/inductor_utils.py
+++ b/torch/testing/_internal/inductor_utils.py
@@ -43,7 +43,6 @@
 )
 from torch.testing._internal.common_utils import (
     IS_CI,
-    IS_FBCODE,
     IS_WINDOWS,
     LazyVal,
     TestCase,
@@ -57,7 +56,7 @@
 def test_cpu():
     try:
         CppCodeCache.load("")
-        return not IS_FBCODE
+        return True
     except (
         CalledProcessError,
         OSError,
@@ -182,14 +181,19 @@ def skip_windows_ci(name: str, file: str) -> None:
 requires_helion = functools.partial(unittest.skipIf, not HAS_HELION, "requires helion")
 
 
-def requires_cuda_with_enough_memory(min_mem_required):
+def requires_gpu_with_enough_memory(min_mem_required):
     def inner(fn):
+        total_memory = sys.maxsize
+        if torch.xpu.is_available():
+            total_memory = torch.xpu.get_device_properties().total_memory
+        elif torch.cuda.is_available():
+            total_memory = torch.cuda.get_device_properties().total_memory
         if (
-            not torch.cuda.is_available()
-            or torch.cuda.get_device_properties().total_memory < min_mem_required
+            not (torch.cuda.is_available() or torch.xpu.is_available())
+            or total_memory < min_mem_required
         ):
             return unittest.skip(
-                f"Only if the CUDA device has at least {min_mem_required / 1e9:.3f}GB memory to be safe"
+                f"Only if the GPU device has at least {min_mem_required / 1e9:.3f}GB memory to be safe"
             )(fn)
         else:
             return fn
diff --git a/torch/testing/_internal/jit_metaprogramming_utils.py b/torch/testing/_internal/jit_metaprogramming_utils.py
index bcb8d137c34fa..43894c733829c 100644
--- a/torch/testing/_internal/jit_metaprogramming_utils.py
+++ b/torch/testing/_internal/jit_metaprogramming_utils.py
@@ -13,7 +13,7 @@
 
 import collections
 from copy import deepcopy
-from typing import Any, Union
+from typing import Any
 import math  # noqa: F401
 
 # Testing utils
@@ -363,7 +363,7 @@ def get_constant(x):
 
 def get_script_args(args):
     formals: list[str] = []
-    tensors: list[Union[torch.Tensor, list[torch.Tensor]]] = []
+    tensors: list[torch.Tensor | list[torch.Tensor]] = []
     actuals: list[str] = []
     for arg in args:
         if isinstance(arg, torch.Tensor):
diff --git a/torch/testing/_internal/jit_utils.py b/torch/testing/_internal/jit_utils.py
index e75454a29b4b8..7381e1583dcf2 100644
--- a/torch/testing/_internal/jit_utils.py
+++ b/torch/testing/_internal/jit_utils.py
@@ -39,7 +39,7 @@
 import tempfile
 import textwrap
 from importlib.abc import Loader
-from typing import Any, Union
+from typing import Any
 
 RUN_CUDA = torch.cuda.is_available()
 RUN_CUDA_MULTI_GPU = RUN_CUDA and torch.cuda.device_count() > 1
@@ -800,7 +800,7 @@ def restore(self):
         torch._C._jit_set_te_must_use_llvm_cpu(self.old_te_must_use_llvm_cpu)
 
 def clone_inputs(args):
-    inputs: list[Union[torch.Tensor, list[torch.Tensor]]] = []
+    inputs: list[torch.Tensor | list[torch.Tensor]] = []
 
     for arg in args:
         if isinstance(arg, torch.Tensor):
diff --git a/torch/testing/_internal/logging_tensor.py b/torch/testing/_internal/logging_tensor.py
index e71f0f4685475..0edbe8ced6dd4 100644
--- a/torch/testing/_internal/logging_tensor.py
+++ b/torch/testing/_internal/logging_tensor.py
@@ -2,7 +2,6 @@
 
 import torch
 from torch.utils._pytree import tree_map
-from typing import Optional
 from collections.abc import Iterator
 import logging
 import contextlib
@@ -84,7 +83,7 @@ class LoggingTensorReentrant(LoggingTensor):
 class LoggingTensorHandler(logging.Handler):
     def __init__(
             self, log_list: list[str], use_shortid_for_all_tensors: bool,
-            with_type: bool, tracebacks_list: Optional[list]) -> None:
+            with_type: bool, tracebacks_list: list | None) -> None:
         logging.Handler.__init__(self)
         self.log_list = log_list
         self.use_shortid_for_all_tensors = use_shortid_for_all_tensors
diff --git a/torch/testing/_internal/opinfo/core.py b/torch/testing/_internal/opinfo/core.py
index 59639eabc5449..acd9cad916441 100644
--- a/torch/testing/_internal/opinfo/core.py
+++ b/torch/testing/_internal/opinfo/core.py
@@ -13,7 +13,7 @@
 from enum import Enum
 from functools import partial
 from itertools import product
-from typing import Any, Optional, TypeVar, Union
+from typing import Any, TypeVar
 
 import torch
 from torch.testing import make_tensor
@@ -652,7 +652,7 @@ class OpInfo:
 
     # An optional reference function that accepts ndarrays (AKA "NumPy arrays").
     # If given, the op will be compared with its reference on each of its sample inputs.
-    ref: Optional[Callable] = None
+    ref: Callable | None = None
 
     # the following metadata describes the operator, its variants, and its aliases, if any
 
@@ -837,11 +837,11 @@ def __setattr__(self, name: str, value: Any) -> None:
 
     # If `supports_cow_input_no_materialize_forward == True`, this list contains
     # the arg indices or kwarg names of inputs that are expected to materialize
-    allow_cow_input_materialize_forward: list[Union[int, str]] = None
+    allow_cow_input_materialize_forward: list[int | str] = None
 
     # If `supports_cow_input_no_materialize_backward == True`, this list contains
     # the arg indices or kwarg names of inputs that are expected to materialize
-    allow_cow_input_materialize_backward: list[Union[int, str]] = None
+    allow_cow_input_materialize_backward: list[int | str] = None
 
     # wrapper function for gradcheck
     gradcheck_wrapper: Callable = lambda op, *args, **kwargs: op(*args, **kwargs)
@@ -876,10 +876,10 @@ def __setattr__(self, name: str, value: Any) -> None:
     aten_name: str = None
 
     # if this is a composite implicit autograd op, the decomposed op
-    decomp_aten_name: Optional[str] = None
+    decomp_aten_name: str | None = None
 
     # name of the corresponding aten:: operator for backwards
-    aten_backward_name: Optional[str] = None
+    aten_backward_name: str | None = None
 
     # if a op's aten::node is expected to be symbolically autodiffed
     assert_autodiffed: bool = False
@@ -1819,11 +1819,11 @@ def __init__(
         name,
         *,
         # The identity value for the operator if it has one.
-        identity: Optional[Any] = None,
+        identity: Any | None = None,
         # The nan policy for the operator if it implements one.
         # - propagate: NaN values are propagated to the output
         # - omit: NaN values are discarded during the reduction
-        nan_policy: Optional[str] = None,
+        nan_policy: str | None = None,
         # Whether the operator supports reducing multiple dimensions.
         supports_multiple_dims: bool = True,
         # Whether the operator promotes integral to floating point dtypes.
@@ -1833,7 +1833,7 @@ def __init__(
         # If a specific dtype is given, then the operator always returns that
         # dtype irrespective of the input dtype. If None, the operator returns
         # the dtype according to the type promotion rules above.
-        result_dtype: Optional[torch.dtype] = None,
+        result_dtype: torch.dtype | None = None,
         # Casts complex results to real (e.g. linalg.norm or torch.var)
         complex_to_real: bool = False,
         # ReductionOpInfo tests generate their own input, dim and keepdim
diff --git a/torch/testing/_internal/opinfo/definitions/_masked.py b/torch/testing/_internal/opinfo/definitions/_masked.py
index a983073d78183..670b34c64ecdf 100644
--- a/torch/testing/_internal/opinfo/definitions/_masked.py
+++ b/torch/testing/_internal/opinfo/definitions/_masked.py
@@ -456,6 +456,13 @@ def sample_inputs_masked_normalize(op_info, device, dtype, requires_grad, **kwar
                 "test_reference_masked",
                 dtypes=(torch.bool, torch.int8, torch.int16, torch.int32),
             ),
+            # FIXME: improve precision
+            DecorateInfo(
+                unittest.skip("Skipped!"),
+                "TestReductions",
+                "test_reference_masked",
+                dtypes=(torch.float16,),
+            ),
             DecorateInfo(
                 unittest.expectedFailure,
                 "TestNormalizeOperators",
@@ -634,16 +641,10 @@ def sample_inputs_masked_normalize(op_info, device, dtype, requires_grad, **kwar
             DecorateInfo(
                 unittest.skip("Skipped!"), "TestJit", "test_variant_consistency_jit"
             ),
-            # Exception: cumulative ops are not yet supported for complex
+            # The following dtypes worked in forward but are not listed by the OpInfo: {torch.bool}.
             DecorateInfo(
                 unittest.expectedFailure, "TestCommon", "test_dtypes", device_type="mps"
             ),
-            DecorateInfo(
-                unittest.expectedFailure,
-                "TestCommon",
-                device_type="mps",
-                dtypes=(torch.complex64,),
-            ),
         ),
         # Can reuse the same inputs; dim is required in both
         sample_inputs_func=sample_inputs_masked_cumops,
@@ -681,16 +682,10 @@ def sample_inputs_masked_normalize(op_info, device, dtype, requires_grad, **kwar
                 "test_comprehensive",
                 device_type="cuda",
             ),
-            # Exception: cumulative ops are not yet supported for complex
+            # The following dtypes worked in forward but are not listed by the OpInfo: {torch.bool}.
             DecorateInfo(
                 unittest.expectedFailure, "TestCommon", "test_dtypes", device_type="mps"
             ),
-            DecorateInfo(
-                unittest.expectedFailure,
-                "TestCommon",
-                device_type="mps",
-                dtypes=(torch.complex64,),
-            ),
         ),
         # Can reuse the same inputs; dim is required in both
         sample_inputs_func=sample_inputs_masked_cumops,
@@ -1308,16 +1303,6 @@ def sample_inputs_masked_normalize(op_info, device, dtype, requires_grad, **kwar
             DecorateInfo(
                 unittest.expectedFailure, "TestJit", "test_variant_consistency_jit"
             ),
-            # Exception: norm ops are not supported for complex yet
-            DecorateInfo(
-                unittest.expectedFailure, "TestCommon", "test_dtypes", device_type="mps"
-            ),
-            DecorateInfo(
-                unittest.expectedFailure,
-                "TestCommon",
-                device_type="mps",
-                dtypes=(torch.complex64,),
-            ),
         ),
         gradcheck_wrapper=gradcheck_wrapper_masked_operation,
         # Runs very slowly on slow gradcheck - alternatively reduce input sizes
diff --git a/torch/testing/_internal/opinfo/definitions/linalg.py b/torch/testing/_internal/opinfo/definitions/linalg.py
index 70c650dc327d8..b96a66af8db95 100644
--- a/torch/testing/_internal/opinfo/definitions/linalg.py
+++ b/torch/testing/_internal/opinfo/definitions/linalg.py
@@ -29,9 +29,13 @@
     all_types_and_complex_and,
     floating_and_complex_types,
     floating_and_complex_types_and,
+    floating_types,
 )
 from torch.testing._internal.common_utils import (
     GRADCHECK_NONDET_TOL,
+    IS_ARM64,
+    IS_CPU_CAPABILITY_SVE256,
+    IS_LINUX,
     make_fullrank_matrices_with_distinct_singular_values,
     skipIfSlowGradcheckEnv,
     slowTest,
@@ -1169,16 +1173,6 @@ def make_input():
                 "TestCommon",
                 "test_numpy_ref_mps",
             ),
-            # RuntimeError: Failed to create function state object for: cross_float2
-            DecorateInfo(
-                unittest.expectedFailure, "TestCommon", "test_dtypes", device_type="mps"
-            ),
-            DecorateInfo(
-                unittest.expectedFailure,
-                "TestCommon",
-                device_type="mps",
-                dtypes=(torch.complex64,),
-            ),
         ),
     ),
     OpInfo(
@@ -1608,6 +1602,15 @@ def make_input():
             ),
             # Exception: The operator 'aten::linalg_lstsq.out' is not currently implemented for the MPS device
             DecorateInfo(unittest.expectedFailure, "TestCommon", device_type="mps"),
+            # see https://github.com/pytorch/pytorch/issues/177249
+            DecorateInfo(
+                unittest.expectedFailure,
+                "TestJit",
+                "test_variant_consistency_jit",
+                device_type="cpu",
+                dtypes=[torch.complex64],
+                active_if=IS_LINUX and IS_ARM64 and not IS_CPU_CAPABILITY_SVE256,
+            ),
         ),
     ),
     OpInfo(
@@ -1798,14 +1801,6 @@ def make_input():
                 "test_noncontiguous_samples",
                 device_type="mps",
             ),
-            # RuntimeError: norm ops are not supported for complex yet
-            DecorateInfo(
-                unittest.expectedFailure,
-                "TestCommon",
-                "test_out_requires_grad_error",
-                device_type="mps",
-                dtypes=(torch.complex64,),
-            ),
         ),
     ),
     OpInfo(
@@ -1870,14 +1865,6 @@ def make_input():
                 "test_noncontiguous_samples",
                 device_type="mps",
             ),
-            # RuntimeError: norm ops are not supported for complex yet
-            DecorateInfo(
-                unittest.expectedFailure,
-                "TestCommon",
-                "test_out_requires_grad_error",
-                device_type="mps",
-                dtypes=(torch.complex64,),
-            ),
         ),
     ),
     OpInfo(
@@ -1929,16 +1916,6 @@ def make_input():
                 "test_variant_consistency_eager",
                 device_type="mps",
             ),
-            # Exception: norm ops are not supported for complex yet
-            DecorateInfo(
-                unittest.expectedFailure, "TestCommon", "test_dtypes", device_type="mps"
-            ),
-            DecorateInfo(
-                unittest.expectedFailure,
-                "TestCommon",
-                device_type="mps",
-                dtypes=(torch.complex64,),
-            ),
         ),
     ),
     OpInfo(
@@ -1946,15 +1923,13 @@ def make_input():
         aten_name="linalg_qr",
         op=torch.linalg.qr,
         dtypes=floating_and_complex_types(),
+        dtypesIfMPS=floating_types(),
         supports_forward_ad=True,
         supports_fwgrad_bwgrad=True,
         # In-place ops
         check_batched_gradgrad=False,
         sample_inputs_func=sample_inputs_linalg_qr_geqrf,
         decorators=[skipCUDAIfNoCusolver, skipCPUIfNoLapack],
-        skips=(
-            DecorateInfo(unittest.expectedFailure, "TestCommon", device_type="mps"),
-        ),
     ),
     OpInfo(
         "linalg.slogdet",
@@ -1994,16 +1969,6 @@ def make_input():
                 "TestCommon",
                 "test_numpy_ref_mps",
             ),
-            # Exception: cumulative ops are not yet supported for complex
-            DecorateInfo(
-                unittest.expectedFailure, "TestCommon", "test_dtypes", device_type="mps"
-            ),
-            DecorateInfo(
-                unittest.expectedFailure,
-                "TestCommon",
-                device_type="mps",
-                dtypes=(torch.complex64,),
-            ),
         ),
     ),
     ReductionOpInfo(
@@ -2021,18 +1986,6 @@ def make_input():
         dtypes=floating_and_complex_types_and(torch.float16, torch.bfloat16),
         generate_args_kwargs=sample_kwargs_vector_norm,
         aten_name="linalg_vector_norm",
-        skips=(
-            # Exception: norm ops are not supported for complex yet
-            DecorateInfo(
-                unittest.expectedFailure, "TestCommon", "test_dtypes", device_type="mps"
-            ),
-            DecorateInfo(
-                unittest.expectedFailure,
-                "TestCommon",
-                device_type="mps",
-                dtypes=(torch.complex64,),
-            ),
-        ),
     ),
     OpInfo(
         "linalg.lu_factor",
@@ -2169,13 +2122,6 @@ def make_input():
             DecorateInfo(
                 unittest.expectedFailure, "TestCommon", "test_out", device_type="mps"
             ),
-            # Exception: Tensor-likes are not close!
-            DecorateInfo(
-                unittest.expectedFailure,
-                "TestCommon",
-                "test_noncontiguous_samples",
-                device_type="mps",
-            ),
             # RuntimeError: linalg.solve.triangular(); Only float is supported!
             DecorateInfo(
                 unittest.expectedFailure,
@@ -2294,13 +2240,6 @@ def make_input():
                 "test_noncontiguous_samples",
                 device_type="cpu",
             ),
-            # IndexError: index -1085850505 is out of bounds for dimension 0 with size 5
-            DecorateInfo(
-                unittest.expectedFailure,
-                "TestCommon",
-                "test_noncontiguous_samples",
-                device_type="mps",
-            ),
             # Exception: linalg.lu_factor(): MPS only supports floats.
             DecorateInfo(
                 unittest.expectedFailure, "TestCommon", "test_dtypes", device_type="mps"
@@ -2376,13 +2315,6 @@ def make_input():
                 device_type="mps",
                 dtypes=[torch.float32],
             ),
-            # IndexError: index -1085850505 is out of bounds for dimension 0 with size 5
-            DecorateInfo(
-                unittest.expectedFailure,
-                "TestCommon",
-                "test_noncontiguous_samples",
-                device_type="mps",
-            ),
             # Exception: linalg.lu_factor(): MPS only supports floats.
             DecorateInfo(
                 unittest.expectedFailure, "TestCommon", "test_dtypes", device_type="mps"
@@ -2405,12 +2337,6 @@ def make_input():
         skips=(
             skipCPUIfNoLapack,
             # AssertionError: Tensor-likes are not close!
-            DecorateInfo(
-                unittest.expectedFailure,
-                "TestCommon",
-                "test_noncontiguous_samples",
-                device_type="mps",
-            ),
             DecorateInfo(
                 unittest.expectedFailure, "TestCommon", "test_out", device_type="mps"
             ),
@@ -2554,6 +2480,7 @@ def make_input():
         # to avoid any rank changes caused by the perturbations in the gradcheck
         op=lambda a, b: torch.linalg.pinv(a @ b.mT),
         dtypes=floating_and_complex_types(),
+        dtypesIfMPS=floating_types(),
         supports_out=False,
         check_batched_grad=False,
         check_batched_gradgrad=False,
@@ -2583,12 +2510,6 @@ def make_input():
                 device_type="cuda",
                 dtypes=[torch.cdouble],
             ),
-            # NotImplementedError: The operator 'aten::linalg_qr.out' is not currently implemented for the MPS device
-            DecorateInfo(
-                unittest.expectedFailure,
-                "TestCommon",
-                device_type="mps",
-            ),
         ),
     ),
     OpInfo(
@@ -2665,6 +2586,15 @@ def make_input():
                 "test_variant_consistency_eager",
                 device_type="mps",
             ),
+            # see https://github.com/pytorch/pytorch/issues/177264
+            DecorateInfo(
+                unittest.expectedFailure,
+                "TestEagerFusionOpInfo",
+                "test_aot_autograd_symbolic_exhaustive",
+                device_type="cpu",
+                dtypes=[torch.float32],
+                active_if=IS_ARM64 and IS_LINUX,
+            ),
         ),
     ),
     OpInfo(
@@ -2684,27 +2614,6 @@ def make_input():
         sample_inputs_func=sample_inputs_svd,
         decorators=[skipCUDAIfNoMagmaAndNoCusolver, skipCPUIfNoLapack, with_tf32_off],
         skips=(
-            DecorateInfo(
-                unittest.skip("Skipped!"),
-                "TestCommon",
-                "test_out",
-                device_type="mps",
-                dtypes=[torch.float32],
-            ),
-            DecorateInfo(
-                unittest.skip("Skipped!"),
-                "TestCommon",
-                "test_variant_consistency_eager",
-                device_type="mps",
-                dtypes=[torch.float32],
-            ),
-            DecorateInfo(
-                unittest.skip("Skipped!"),
-                "TestJit",
-                "test_variant_consistency_jit",
-                device_type="mps",
-                dtypes=[torch.float32],
-            ),
             DecorateInfo(
                 unittest.skip("Skipped!"),
                 "TestFakeTensor",
@@ -2721,24 +2630,6 @@ def make_input():
                 dtypes=[torch.float32],
                 active_if=TEST_WITH_ROCM,
             ),
-            # MPS: AssertionError: The values for attribute 'shape' do not match: torch.Size([0, 0]) != torch.Size([0, 1]).
-            DecorateInfo(
-                unittest.expectedFailure,
-                "TestCommon",
-                "test_out_warning",
-                device_type="mps",
-            ),
-            # MPS: RuntimeError: svd_backward: The singular vectors in the
-            # complex case are specified up to multiplication by e^{i phi}. The
-            # specified loss function depends on this phase term, making it
-            # ill-defined.
-            DecorateInfo(
-                unittest.expectedFailure,
-                "TestCommon",
-                "test_noncontiguous_samples",
-                device_type="mps",
-                dtypes=(torch.complex64,),
-            ),
         ),
     ),
     OpInfo(
@@ -2831,27 +2722,8 @@ def make_input():
                 "test_output_match",
                 device_type="mps",
             ),
-            # Exception: index ????? is out of bounds for dimension 0 with size 6
-            DecorateInfo(
-                unittest.expectedFailure,
-                "TestCommon",
-                "test_variant_consistency_eager",
-                device_type="mps",
-            ),
-            # Exception: Tensor-likes are not close!
-            DecorateInfo(
-                unittest.expectedFailure,
-                "TestCommon",
-                "test_noncontiguous_samples",
-                device_type="mps",
-            ),
         ],
         skips=(
-            DecorateInfo(
-                unittest.skip("Unsupported on MPS for now"),
-                "TestCommon",
-                "test_numpy_ref_mps",
-            ),
             # Exception: linalg.lu_factor(): MPS only supports floats.
             DecorateInfo(
                 unittest.expectedFailure, "TestCommon", "test_dtypes", device_type="mps"
@@ -2880,20 +2752,6 @@ def make_input():
             DecorateInfo(
                 unittest.expectedFailure, "TestCommon", "test_python_ref_errors"
             ),
-            DecorateInfo(
-                unittest.expectedFailure,
-                "TestCommon",
-                "test_python_ref",
-                device_type="mps",
-                dtypes=(torch.bfloat16, torch.complex64, torch.float16),
-            ),
-            DecorateInfo(
-                unittest.expectedFailure,
-                "TestCommon",
-                "test_python_ref_torch_fallback",
-                device_type="mps",
-                dtypes=(torch.bfloat16, torch.complex64, torch.float16),
-            ),
         ),
     ),
     PythonRefInfo(
@@ -2906,61 +2764,12 @@ def make_input():
         "_refs.linalg.vecdot",
         torch_opinfo_name="linalg.vecdot",
         op_db=op_db,
-        skips=(
-            # TypeError: Cannot convert a MPS Tensor to float64 dtype as the MPS framework doesn't support float64
-            DecorateInfo(
-                unittest.expectedFailure,
-                "TestCommon",
-                "test_python_ref_torch_fallback",
-                device_type="mps",
-                dtypes=(torch.float16, torch.bfloat16),
-            ),
-            DecorateInfo(
-                unittest.expectedFailure,
-                "TestCommon",
-                "test_python_ref",
-                device_type="mps",
-                dtypes=(torch.float16, torch.bfloat16),
-            ),
-        ),
     ),
     ReductionPythonRefInfo(
         "_refs.linalg.vector_norm",
         torch_opinfo_name="linalg.vector_norm",
         supports_out=True,
         op_db=op_db,
-        skips=(
-            # TypeError: Cannot convert a MPS Tensor to float64 dtype as the MPS framework doesn't support float64
-            DecorateInfo(
-                unittest.expectedFailure,
-                "TestCommon",
-                "test_python_ref",
-                device_type="mps",
-                dtypes=(torch.float16,),
-            ),
-            DecorateInfo(
-                unittest.expectedFailure,
-                "TestCommon",
-                "test_python_ref_torch_fallback",
-                device_type="mps",
-                dtypes=(torch.float16,),
-            ),
-            # Exception: norm ops are not supported for complex yet
-            DecorateInfo(
-                unittest.expectedFailure,
-                "TestCommon",
-                "test_python_ref_torch_fallback",
-                device_type="mps",
-                dtypes=(torch.complex64,),
-            ),
-            DecorateInfo(
-                unittest.expectedFailure,
-                "TestCommon",
-                "test_python_ref",
-                device_type="mps",
-                dtypes=(torch.complex64,),
-            ),
-        ),
     ),
     PythonRefInfo(
         "_refs.linalg.matrix_norm",
@@ -2977,29 +2786,14 @@ def make_input():
                 "TestCommon",
                 "test_python_ref_torch_fallback",
                 device_type="mps",
-                dtypes=(torch.float32,),
-            ),
-            DecorateInfo(
-                unittest.expectedFailure,
-                "TestCommon",
-                "test_python_ref",
-                device_type="mps",
-                dtypes=(torch.float32,),
+                dtypes=(torch.float32, torch.complex64),
             ),
-            # Exception: norm ops are not supported for complex yet
             DecorateInfo(
                 unittest.expectedFailure,
                 "TestCommon",
                 "test_python_ref",
                 device_type="mps",
-                dtypes=(torch.complex64,),
-            ),
-            DecorateInfo(
-                unittest.expectedFailure,
-                "TestCommon",
-                "test_python_ref_torch_fallback",
-                device_type="mps",
-                dtypes=(torch.complex64,),
+                dtypes=(torch.float32, torch.complex64),
             ),
         ),
     ),
diff --git a/torch/testing/_internal/opinfo/definitions/nested.py b/torch/testing/_internal/opinfo/definitions/nested.py
index a709a8aff158e..bdb4d942aba06 100644
--- a/torch/testing/_internal/opinfo/definitions/nested.py
+++ b/torch/testing/_internal/opinfo/definitions/nested.py
@@ -4,7 +4,6 @@
 from copy import copy
 from dataclasses import dataclass
 from functools import partial
-from typing import Optional
 
 import torch
 from torch.fx.experimental.symbolic_shapes import is_nested_int
@@ -47,7 +46,7 @@ class ExtraOpData:
     # Returns: tuple of (single dim argname if available, dim list argname if available)
     # If the op doesn't support dim-related args at all OR this op only has overloads
     # with multiple dim args (e.g. transpose()), then this returns (None, None).
-    def get_dim_argnames(self) -> tuple[Optional[str], Optional[str]]:
+    def get_dim_argnames(self) -> tuple[str | None, str | None]:
         if self.dim_args is None:
             return (None, None)
 
diff --git a/torch/testing/_internal/opinfo/definitions/sparse.py b/torch/testing/_internal/opinfo/definitions/sparse.py
index 7f8781ce87b24..c04009f1f2ac9 100644
--- a/torch/testing/_internal/opinfo/definitions/sparse.py
+++ b/torch/testing/_internal/opinfo/definitions/sparse.py
@@ -4,6 +4,7 @@
 
 import torch
 from torch.testing import make_tensor  # noqa: F401
+from torch.testing._internal.common_dtype import highest_precision_float
 from torch.testing._internal.opinfo.core import (  # noqa: F401
     BinaryUfuncInfo,
     ErrorInput,
@@ -769,8 +770,9 @@ def _sample_inputs_sparse_like_fns(
             tensor, args=(), kwargs=dict(device=device, dtype=dtype, layout=layout)
         )
 
-        if dtype is not torch.float64:
-            yield SampleInput(tensor, args=(), kwargs=dict(dtype=torch.float64))
+        hpf = highest_precision_float(device)
+        if dtype is not hpf:
+            yield SampleInput(tensor, args=(), kwargs=dict(dtype=hpf))
 
         if torch.cuda.is_available():
             other_device = "cuda" if tensor.device.type == "cpu" else "cpu"
diff --git a/torch/testing/_internal/opinfo/definitions/special.py b/torch/testing/_internal/opinfo/definitions/special.py
index 8ca4383afa349..74b3b03ce138b 100644
--- a/torch/testing/_internal/opinfo/definitions/special.py
+++ b/torch/testing/_internal/opinfo/definitions/special.py
@@ -223,21 +223,6 @@ def sample_inputs_erfcx(op_info, device, dtype, requires_grad, **kwargs):
                 "TestNormalizeOperators",
                 "test_normalize_operator_exhaustive",
             ),
-            # AssertionError: Tensor-likes are not close!
-            DecorateInfo(
-                unittest.expectedFailure,
-                "TestCommon",
-                "test_noncontiguous_samples",
-                device_type="mps",
-                dtypes=(torch.float32, torch.int64),
-            ),
-            DecorateInfo(
-                unittest.expectedFailure,
-                "TestCommon",
-                "test_out",
-                device_type="mps",
-                dtypes=(torch.float32,),
-            ),
         ),
         sample_kwargs=lambda device, dtype, input: ({"n": 0}, {"n": 0}),
         # polygamma functions have multiple singularities at x having non-positive integer value
@@ -381,7 +366,7 @@ def sample_inputs_erfcx(op_info, device, dtype, requires_grad, **kwargs):
         ),
         dtypes=all_types_and(torch.bool, *_unsigned_int_types),
         dtypesIfMPS=all_types_and(torch.bool),
-        ref=lambda x: scipy.special.airy(x)[0] if TEST_SCIPY else None,
+        ref=(lambda x: scipy.special.airy(x)[0]) if TEST_SCIPY else None,
         skips=(
             DecorateInfo(
                 unittest.skip("Skipped!"),
@@ -862,7 +847,7 @@ def sample_inputs_erfcx(op_info, device, dtype, requires_grad, **kwargs):
         ),
         dtypes=all_types_and(torch.bool, *_unsigned_int_types),
         dtypesIfMPS=all_types_and(torch.bool, torch.float16, torch.bfloat16),
-        ref=lambda x: scipy.special.spherical_jn(0, x) if TEST_SCIPY else None,
+        ref=(lambda x: scipy.special.spherical_jn(0, x)) if TEST_SCIPY else None,
         supports_autograd=False,
         skips=(
             DecorateInfo(
@@ -893,23 +878,6 @@ def sample_inputs_erfcx(op_info, device, dtype, requires_grad, **kwargs):
                 },
             ),
         ),
-        skips=(
-            # TypeError: Cannot convert a MPS Tensor to float64 dtype as the MPS framework doesn't support float64
-            DecorateInfo(
-                unittest.expectedFailure,
-                "TestCommon",
-                "test_python_ref",
-                device_type="mps",
-                dtypes=(torch.float16, torch.bfloat16),
-            ),
-            DecorateInfo(
-                unittest.expectedFailure,
-                "TestCommon",
-                "test_python_ref_torch_fallback",
-                device_type="mps",
-                dtypes=(torch.float16, torch.bfloat16),
-            ),
-        ),
     ),
     ElementwiseUnaryPythonRefInfo(
         "_refs.special.bessel_j1",
@@ -923,23 +891,6 @@ def sample_inputs_erfcx(op_info, device, dtype, requires_grad, **kwargs):
                 },
             ),
         ),
-        skips=(
-            # TypeError: Cannot convert a MPS Tensor to float64 dtype as the MPS framework doesn't support float64
-            DecorateInfo(
-                unittest.expectedFailure,
-                "TestCommon",
-                "test_python_ref",
-                device_type="mps",
-                dtypes=(torch.float16, torch.bfloat16),
-            ),
-            DecorateInfo(
-                unittest.expectedFailure,
-                "TestCommon",
-                "test_python_ref_torch_fallback",
-                device_type="mps",
-                dtypes=(torch.float16, torch.bfloat16),
-            ),
-        ),
     ),
     ElementwiseUnaryPythonRefInfo(
         "_refs.special.entr",
@@ -1010,21 +961,6 @@ def sample_inputs_erfcx(op_info, device, dtype, requires_grad, **kwargs):
                 "test_reference_numerics_large",
                 dtypes=(torch.int8,),
             ),
-            # TypeError: Cannot convert a MPS Tensor to float64 dtype as the MPS framework doesn't support float64
-            DecorateInfo(
-                unittest.expectedFailure,
-                "TestCommon",
-                "test_python_ref",
-                device_type="mps",
-                dtypes=(torch.float16,),
-            ),
-            DecorateInfo(
-                unittest.expectedFailure,
-                "TestCommon",
-                "test_python_ref_torch_fallback",
-                device_type="mps",
-                dtypes=(torch.float16,),
-            ),
         ),
     ),
     ElementwiseUnaryPythonRefInfo(
@@ -1056,23 +992,6 @@ def sample_inputs_erfcx(op_info, device, dtype, requires_grad, **kwargs):
         "_refs.special.ndtr",
         torch_opinfo_name="special.ndtr",
         op_db=op_db,
-        skips=(
-            # TypeError: Cannot convert a MPS Tensor to float64 dtype as the MPS framework doesn't support float64
-            DecorateInfo(
-                unittest.expectedFailure,
-                "TestCommon",
-                "test_python_ref",
-                device_type="mps",
-                dtypes=(torch.float16, torch.bfloat16),
-            ),
-            DecorateInfo(
-                unittest.expectedFailure,
-                "TestCommon",
-                "test_python_ref_torch_fallback",
-                device_type="mps",
-                dtypes=(torch.float16, torch.bfloat16),
-            ),
-        ),
     ),
     ElementwiseUnaryPythonRefInfo(
         "_refs.special.ndtri",
diff --git a/torch/testing/_internal/optests/generate_tests.py b/torch/testing/_internal/optests/generate_tests.py
index 7e3a7b0408333..68ace6688bfeb 100644
--- a/torch/testing/_internal/optests/generate_tests.py
+++ b/torch/testing/_internal/optests/generate_tests.py
@@ -11,7 +11,7 @@
 import threading
 import unittest
 from collections.abc import Callable, Sequence
-from typing import Any, Optional, Union
+from typing import Any
 
 import torch
 import torch._dynamo
@@ -50,8 +50,8 @@ def safe_schema_check(
     kwargs: dict[str, Any],
     *,
     copy_inputs: bool = True,
-    rtol: Optional[float] = None,
-    atol: Optional[float] = None,
+    rtol: float | None = None,
+    atol: float | None = None,
 ) -> Any:
     if copy_inputs:
         args, kwargs = deepcopy_tensors((args, kwargs))
@@ -68,8 +68,8 @@ def safe_autograd_registration_check(
     kwargs: dict[str, Any],
     *,
     copy_inputs: bool = True,
-    rtol: Optional[float] = None,
-    atol: Optional[float] = None,
+    rtol: float | None = None,
+    atol: float | None = None,
 ) -> None:
     if pytree.tree_any_only(torch.Tensor, is_abstract, (args, kwargs)):
         return
@@ -89,8 +89,8 @@ def safe_fake_check(
     kwargs: dict[str, Any],
     *,
     copy_inputs: bool = True,
-    rtol: Optional[float] = None,
-    atol: Optional[float] = None,
+    rtol: float | None = None,
+    atol: float | None = None,
 ) -> None:
     if pytree.tree_any_only(torch.Tensor, is_abstract, (args, kwargs)):
         return None
@@ -106,8 +106,8 @@ def safe_aot_autograd_check(
     dynamic: bool,
     *,
     copy_inputs: bool = True,
-    rtol: Optional[float] = None,
-    atol: Optional[float] = None,
+    rtol: float | None = None,
+    atol: float | None = None,
 ) -> Any:
     # NB: copy_inputs does nothing for aot_autograd_check: it always needs to copy
     # inputs.
@@ -178,8 +178,8 @@ def deepcopy_tensors(inputs: Any) -> Any:
 def generate_opcheck_tests(
     testcase: Any,
     namespaces: list[str],
-    failures_dict_path: Optional[str] = None,
-    additional_decorators: Optional[dict[str, Callable]] = None,
+    failures_dict_path: str | None = None,
+    additional_decorators: dict[str, Callable] | None = None,
     test_utils: list[str] = DEFAULT_TEST_UTILS,
 ) -> None:
     """Given an existing TestCase, use the existing tests to generate
@@ -649,14 +649,14 @@ def should_print_better_repro() -> None:
 
 
 def opcheck(
-    op: Union[torch._ops.OpOverload, torch._ops.OpOverloadPacket, CustomOpDef],
+    op: torch._ops.OpOverload | torch._ops.OpOverloadPacket | CustomOpDef,
     args: tuple[Any, ...],
-    kwargs: Optional[dict[str, Any]] = None,
+    kwargs: dict[str, Any] | None = None,
     *,
-    test_utils: Union[str, Sequence[str]] = DEFAULT_TEST_UTILS,
+    test_utils: str | Sequence[str] = DEFAULT_TEST_UTILS,
     raise_exception: bool = True,
-    rtol: Optional[float] = None,
-    atol: Optional[float] = None,
+    rtol: float | None = None,
+    atol: float | None = None,
 ) -> dict[str, str]:
     """See torch.library.opcheck for docstring"""
 
@@ -817,7 +817,7 @@ def load(path, *, create_file=False) -> "FailuresDict":
                     )
         return FailuresDict(path, dct["data"])
 
-    def _save(self, to_str=False) -> Optional[str]:
+    def _save(self, to_str=False) -> str | None:
         to_dump = {
             "_description": DESCRIPTION,
             "data": self.data,
@@ -849,7 +849,7 @@ def set_status(
         test_name: str,
         status: str,
         *,
-        comment: Optional[str] = None,
+        comment: str | None = None,
     ):
         if qualname not in self.data:
             self.data[qualname] = {}
diff --git a/torch/testing/_internal/py312_intrinsics.py b/torch/testing/_internal/py312_intrinsics.py
new file mode 100644
index 0000000000000..ae6dc3ddd7244
--- /dev/null
+++ b/torch/testing/_internal/py312_intrinsics.py
@@ -0,0 +1,29 @@
+import torch
+
+
+# This is a very odd way to define a test case that uses CALL_INTRINSIC_1 4 and
+# CALL_INTRINSIC_2 7. The problem is those intrinsics only exists on Python
+# 3.12+ and the Python code to produce them is not backwards compatible with
+# previous versions.
+class Foo:
+    @classmethod
+    def _default_update(cls):
+        def f[T](a: "This is a new annotation"):  # noqa: F722
+            """This is a test"""
+
+        f.attr = "This is also a test"
+        f.__wrapped__ = "This is a bald faced lie"
+
+        def wrapper(b: "This is the prior annotation"):  # noqa: F722
+            pass
+
+        return wrapper, f
+
+    def test_default_update(self):
+        @torch.compile(backend="eager", fullgraph=True)
+        def fn(x):
+            wrapper, f = Foo._default_update()
+            return x + 1
+
+        x = torch.randn(2)
+        fn(x)
diff --git a/torch/testing/_internal/subclasses.py b/torch/testing/_internal/subclasses.py
index 93bcf98e15fc4..03326233a8fbb 100644
--- a/torch/testing/_internal/subclasses.py
+++ b/torch/testing/_internal/subclasses.py
@@ -1,5 +1,5 @@
 # mypy: ignore-errors
-from typing import Any, Optional
+from typing import Any
 
 import torch
 import torch.utils._pytree as pytree
@@ -71,7 +71,7 @@ def __torch_dispatch__(cls, func, types, args, kwargs):
             return return_and_correct_aliasing(func, args, kwargs, out)
 
     def __coerce_same_metadata_as_tangent__(
-        self, expected_metadata: Any, expected_type: Optional[type] = None
+        self, expected_metadata: Any, expected_type: type | None = None
     ):
         if expected_type is type(self.a):
             return self.a
diff --git a/torch/testing/_internal/torchbind_impls.py b/torch/testing/_internal/torchbind_impls.py
index e5162ba0d6cb6..02ec2623271c9 100644
--- a/torch/testing/_internal/torchbind_impls.py
+++ b/torch/testing/_internal/torchbind_impls.py
@@ -1,14 +1,13 @@
 # mypy: allow-untyped-defs
 import contextlib
 from pathlib import Path
-from typing import Optional
 
 import torch
 
 
 _TORCHBIND_IMPLS_INITIALIZED = False
 
-_TENSOR_QUEUE_GLOBAL_TEST: Optional[torch.ScriptObject] = None
+_TENSOR_QUEUE_GLOBAL_TEST: torch.ScriptObject | None = None
 
 
 def init_torchbind_implementations():
diff --git a/torch/testing/_internal/triton_utils.py b/torch/testing/_internal/triton_utils.py
index e08820dd5b825..9961330e57bb9 100644
--- a/torch/testing/_internal/triton_utils.py
+++ b/torch/testing/_internal/triton_utils.py
@@ -13,6 +13,9 @@
 requires_cuda_and_triton = unittest.skipUnless(
     HAS_CUDA_AND_TRITON, "requires cuda and triton"
 )
+requires_xpu_and_triton = unittest.skipUnless(
+    HAS_XPU_AND_TRITON, "requires xpu and triton"
+)
 requires_gpu_and_triton = unittest.skipUnless(
     HAS_XPU_AND_TRITON or HAS_CUDA_AND_TRITON, "requires gpu and triton"
 )
@@ -1055,6 +1058,26 @@ def add_kernel_with_boolean_param(
             output = x
         tl.store(out_ptr + offsets, output, mask=mask)
 
+    @triton.jit
+    def masked_add_kernel_with_bool_tensor(
+        in_ptr0,
+        in_ptr1,
+        mask_ptr,
+        out_ptr,
+        n_elements,
+        BLOCK_SIZE: "tl.constexpr",
+    ):
+        """Kernel that loads a bool tensor and uses it as a mask."""
+        pid = tl.program_id(axis=0)
+        block_start = pid * BLOCK_SIZE
+        offsets = block_start + tl.arange(0, BLOCK_SIZE)
+        valid = offsets < n_elements
+        x = tl.load(in_ptr0 + offsets, mask=valid)
+        y = tl.load(in_ptr1 + offsets, mask=valid)
+        keep = tl.load(mask_ptr + offsets, mask=valid, other=0) != 0
+        output = tl.where(keep, x + y, x)
+        tl.store(out_ptr + offsets, output, mask=valid)
+
     # support the old (experimental) and new (tensor_descriptor) APIs
     def create_tensor_descriptor_shim(
         tensor, block_sizes: list[int], new_api: bool = True
diff --git a/torch/testing/_utils.py b/torch/testing/_utils.py
index b86edfdd67f3c..02eeb1f46f5c8 100644
--- a/torch/testing/_utils.py
+++ b/torch/testing/_utils.py
@@ -33,8 +33,12 @@ def freeze_rng_state():
     # which we need to disable to get and set rng state
     with torch.utils._mode_utils.no_dispatch(), torch._C._DisableFuncTorch():
         rng_state = torch.get_rng_state()
-        if torch.cuda.is_available():
-            cuda_rng_state = torch.cuda.get_rng_state()
+        if torch.accelerator.is_available():
+            accelerator = torch.accelerator.current_accelerator(check_available=True)
+            if accelerator is not None:
+                accelerator_rng_state = torch.get_device_module(
+                    accelerator.type
+                ).get_rng_state()
     try:
         yield
     finally:
@@ -47,6 +51,12 @@ def freeze_rng_state():
         #
         # NB: Mode disable is to avoid running cross-ref tests on this seeding
         with torch.utils._mode_utils.no_dispatch(), torch._C._DisableFuncTorch():
-            if torch.cuda.is_available():
-                torch.cuda.set_rng_state(cuda_rng_state)  # type: ignore[possibly-undefined]
+            if torch.accelerator.is_available():
+                accelerator = torch.accelerator.current_accelerator(
+                    check_available=True
+                )
+                if accelerator is not None:
+                    torch.get_device_module(accelerator.type).set_rng_state(
+                        accelerator_rng_state  # type: ignore[possibly-undefined]
+                    )
             torch.set_rng_state(rng_state)
diff --git a/torch/types.py b/torch/types.py
index 29a6a65b5ebdc..d495a129f40c7 100644
--- a/torch/types.py
+++ b/torch/types.py
@@ -12,7 +12,7 @@
     str as _str,
 )
 from collections.abc import Sequence
-from typing import Any, IO, Optional, TYPE_CHECKING, TypeAlias, Union
+from typing import Any, IO, TYPE_CHECKING, TypeAlias, Union
 from typing_extensions import Self
 
 # `as` imports have better static analysis support than assignment `ExposedType: TypeAlias = HiddenType`
@@ -39,7 +39,7 @@
 # Convenience aliases for common composite types that we need
 # to talk about in PyTorch
 _TensorOrTensors: TypeAlias = Tensor | Sequence[Tensor]  # noqa: PYI047
-_TensorOrOptionalTensors: TypeAlias = Union[Tensor, Sequence[Optional[Tensor]]]  # noqa: PYI047
+_TensorOrOptionalTensors: TypeAlias = Tensor | Sequence[Tensor | None]  # noqa: PYI047
 _TensorOrTensorsOrGradEdge: TypeAlias = Union[  # noqa: PYI047
     Tensor,
     Sequence[Tensor],
diff --git a/torch/utils/_config_module.py b/torch/utils/_config_module.py
index c962f20b33ceb..99e7cf8f9cc77 100644
--- a/torch/utils/_config_module.py
+++ b/torch/utils/_config_module.py
@@ -9,11 +9,11 @@
 import tokenize
 import unittest
 from collections.abc import Callable
+from contextvars import ContextVar
 from dataclasses import dataclass
 from types import FunctionType, ModuleType
 from typing import Any, Generic, NoReturn, Optional, TYPE_CHECKING, TypeVar
 from typing_extensions import deprecated
-from unittest import mock
 
 from torch._utils_internal import justknobs_check
 
@@ -21,6 +21,10 @@
 # Types saved/loaded in configs
 CONFIG_TYPES = (int, float, bool, type(None), str, list, set, tuple, dict)
 
+# Immutable scalar types that don't need deepcopy when returned from configs.
+# Everything else is defensively copied to prevent accidental mutation.
+_IMMUTABLE_CONFIG_TYPES = (int, float, bool, type(None), str, tuple)
+
 
 # Duplicated, because mypy needs these types statically
 T = TypeVar("T", bound=int | float | bool | str | list | set | tuple | dict | None)
@@ -46,7 +50,7 @@ class _Config(Generic[T]):
             If multiple env variables are given, the precedence order is from
             left to right.
         user_override: If a user sets a value (i.e. foo.bar=True), that
-            has precedence over everything after this.
+            has precedence over everything after this.  User overrides are thread-local.
         env_name_default: If set, this environment variable will override everything
             after this.
             If multiple env variables are given, the precedence order is from
@@ -75,6 +79,9 @@ class _Config(Generic[T]):
     env_name_force: list[str] | None = None
     value_type: type | None = None
     alias: str | None = None
+    # Deprecation support
+    deprecated: bool = False
+    deprecation_message: str | None = None
 
     def __post_init__(self) -> None:
         self.env_name_default = _Config.string_or_list_of_string_to_list(
@@ -122,6 +129,9 @@ def Config(
         env_name_force: str | list[str] | None = None,
         value_type: type | None = None,
         alias: str | None = None,
+        # Deprecation support
+        deprecated: bool = False,
+        deprecation_message: str | None = None,
     ) -> T: ...
 
 else:
@@ -133,6 +143,9 @@ def Config(
         env_name_force: str | list[str] | None = None,
         value_type: type | None = None,
         alias: str | None = None,
+        # Deprecation support
+        deprecated: bool = False,
+        deprecation_message: str | None = None,
     ) -> _Config[T]:
         return _Config(
             default=default,
@@ -141,6 +154,9 @@ def Config(
             env_name_force=env_name_force,
             value_type=value_type,
             alias=alias,
+            # Deprecation support
+            deprecated=deprecated,
+            deprecation_message=deprecation_message,
         )
 
 
@@ -191,7 +207,7 @@ def visit(
             annotated_type = type_hints.get(key, None)
             if isinstance(value, CONFIG_TYPES):
                 config[name] = _ConfigEntry(
-                    _Config(default=value, value_type=annotated_type)
+                    _Config(default=value, value_type=annotated_type), name
                 )
                 if dest is module:
                     delattr(module, key)
@@ -199,7 +215,7 @@ def visit(
                 if annotated_type is not None and value.value_type is None:
                     value.value_type = annotated_type
 
-                config[name] = _ConfigEntry(value)
+                config[name] = _ConfigEntry(value, name)
 
                 if dest is module:
                     delattr(module, key)
@@ -279,7 +295,7 @@ class _ConfigEntry:
     value_type: type
     # The value specified by the user when they overrode the configuration
     # _UNSET_SENTINEL indicates the value is not set.
-    user_override: Any = _UNSET_SENTINEL
+    user_override: ContextVar[object]
     # The justknob to check for this config
     justknob: str | None = None
     # environment variables are read at install time
@@ -298,14 +314,24 @@ class _ConfigEntry:
     # upstream bug - python/cpython#126886
     hide: bool = False
     alias: str | None = None
+    # Deprecation support
+    deprecated: bool = False
+    deprecation_message: str | None = None
+    _deprecation_warned: bool = False
 
-    def __init__(self, config: _Config) -> None:
+    def __init__(self, config: _Config, name: str) -> None:
         self.default = config.default
         self.value_type = (
             config.value_type if config.value_type is not None else type(self.default)
         )
         self.justknob = config.justknob
         self.alias = config.alias
+        # Deprecation fields
+        self.deprecated = config.deprecated
+        self.deprecation_message = config.deprecation_message
+        self._deprecation_warned = False
+
+        self.user_override = ContextVar(name, default=_UNSET_SENTINEL)
         if config.env_name_default is not None:
             for val in config.env_name_default:
                 if (env_value := _read_env_variable(val)) is not None:
@@ -329,8 +355,8 @@ def __init__(self, config: _Config) -> None:
             if self.value_type not in (
                 bool,
                 str,
-                Optional[bool],
-                Optional[str],
+                Optional[bool],  # noqa: UP045
+                Optional[str],  # noqa: UP045
             ):
                 raise AssertionError(
                     f"envvar configs only support (optional) booleans or strings, {self.value_type} is neither"
@@ -354,17 +380,34 @@ def __init__(self) -> None:
             f"use {__name__}.install_config_module(sys.modules[__name__])"
         )
 
+    def _warn_if_deprecated(self, name: str, config: _ConfigEntry) -> None:
+        """Issue deprecation warning for config if not already warned."""
+        if config.deprecated and not config._deprecation_warned:
+            import warnings
+
+            msg = f"{self.__name__}.{name} is deprecated"
+            if config.deprecation_message:
+                msg += f" and {config.deprecation_message}"
+            msg += ". It will be removed in a future version of PyTorch."
+            warnings.warn(msg, FutureWarning, stacklevel=3)
+            config._deprecation_warned = True
+
     def __setattr__(self, name: str, value: object) -> None:
         if name in self._bypass_keys:
             super().__setattr__(name, value)
         elif name not in self._config:
             raise AttributeError(f"{self.__name__}.{name} does not exist")
-        elif self._config[name].alias is not None:
-            self._set_alias_val(self._config[name], value)
         else:
-            self._config[name].user_override = value
-            self._is_dirty = True
-            self._config[name].hide = False
+            # Issue deprecation warning on write (once per config)
+            config = self._config[name]
+            self._warn_if_deprecated(name, config)
+
+            if config.alias is not None:
+                self._set_alias_val(config, value)
+            else:
+                config.user_override.set(value)
+                self._is_dirty = True
+                config.hide = False
 
     def __getattr__(self, name: str) -> Any:
         try:
@@ -373,6 +416,9 @@ def __getattr__(self, name: str) -> Any:
             if config.hide:
                 raise AttributeError(f"{self.__name__}.{name} does not exist")
 
+            # Issue deprecation warning on read (once per config)
+            self._warn_if_deprecated(name, config)
+
             alias_val = self._get_alias_val(config)
             if alias_val is not _UNSET_SENTINEL:
                 return alias_val
@@ -380,8 +426,9 @@ def __getattr__(self, name: str) -> Any:
             if config.env_value_force is not _UNSET_SENTINEL:
                 return config.env_value_force
 
-            if config.user_override is not _UNSET_SENTINEL:
-                return config.user_override
+            user_override = config.user_override.get()
+            if user_override is not _UNSET_SENTINEL:
+                return user_override
 
             if config.env_value_default is not _UNSET_SENTINEL:
                 return config.env_value_default
@@ -390,12 +437,11 @@ def __getattr__(self, name: str) -> Any:
                 # JK only supports bools and ints
                 return justknobs_check(name=config.justknob, default=config.default)
 
-            # Note that reference types can still be modified, so we
-            # copy them to user_overrides in case the user overrides
-            # them
-            if isinstance(config.default, (list, set, dict)):
-                config.user_override = copy.deepcopy(config.default)
-                return config.user_override
+            # Reference types can still be modified, so copy them to
+            # user_overrides to prevent accidental mutation of defaults.
+            if not isinstance(config.default, _IMMUTABLE_CONFIG_TYPES):
+                config.user_override.set(copy.deepcopy(config.default))
+                return config.user_override.get()
             return config.default
 
         except KeyError as e:
@@ -406,7 +452,7 @@ def __delattr__(self, name: str) -> None:
         self._is_dirty = True
         # must support delete because unittest.mock.patch deletes
         # then recreate things
-        self._config[name].user_override = _UNSET_SENTINEL
+        self._config[name].user_override.set(_UNSET_SENTINEL)
         self._config[name].hide = True
 
     def _get_alias_module_and_name(
@@ -457,10 +503,10 @@ def _is_default(self, name: str) -> bool:
             or config_val.env_value_force == config_val.default
         )
 
-        unset = config_val.user_override is _UNSET_SENTINEL
+        unset = config_val.user_override.get() is _UNSET_SENTINEL
         # Handle reference types specially to avoid spammy warnings
-        if isinstance(config_val.default, (list, set, dict)):
-            unset = unset or config_val.user_override == config_val.default
+        if not isinstance(config_val.default, _IMMUTABLE_CONFIG_TYPES):
+            unset = unset or config_val.user_override.get() == config_val.default
         return unset and not_set_env_default and not_set_env_force
 
     def _get_dict(
@@ -486,7 +532,9 @@ def _get_dict(
                 it skips it.
         """
         config: dict[str, Any] = {}
-        for key in self._config:
+        for key, entry in self._config.items():
+            if entry.alias is not None:
+                continue
             if ignored_keys and key in ignored_keys:
                 continue
             if ignored_prefixes:
@@ -494,9 +542,23 @@ def _get_dict(
                     continue
             if skip_default and self._is_default(key):
                 continue
-            if self._config[key].alias is not None:
-                continue
-            config[key] = copy.deepcopy(getattr(self, key))
+
+            # Read value directly, bypassing __getattr__ overhead
+            # (deprecation warnings, alias resolution).
+            user_override = entry.user_override.get()
+            if entry.env_value_force is not _UNSET_SENTINEL:
+                val = entry.env_value_force
+            elif user_override is not _UNSET_SENTINEL:
+                val = user_override
+            elif entry.env_value_default is not _UNSET_SENTINEL:
+                val = entry.env_value_default
+            elif entry.justknob is not None:
+                val = justknobs_check(name=entry.justknob, default=entry.default)
+            else:
+                val = entry.default
+            if not isinstance(val, _IMMUTABLE_CONFIG_TYPES):
+                val = copy.deepcopy(val)
+            config[key] = val
 
         return config
 
@@ -519,7 +581,19 @@ def save_config_portable(
         if ignore_private_configs:
             prefixes.append("_")
         prefixes.extend(getattr(self, "_cache_config_ignore_prefix", []))
-        return self._get_dict(ignored_prefixes=prefixes)
+        config = self._get_dict(ignored_prefixes=prefixes)
+        for key in getattr(self, "_cache_config_factory_keys", []):
+            if key in config and config[key] is not None:
+                instance = config[key]()
+                if hasattr(instance, "uuid"):
+                    config[key] = instance.uuid()
+                else:
+                    raise RuntimeError(
+                        f"Config '{key}' is set to {config[key]} which does not "
+                        f"implement uuid(). Implement uuid() for cache key "
+                        f"participation."
+                    )
+        return config
 
     def codegen_config(self) -> str:
         """Convert config to Python statements that replicate current config.
@@ -642,7 +716,7 @@ def patch(
         **kwargs: dict[str, Any],
     ) -> "ContextDecorator":
         """
-        Decorator and/or context manager to make temporary changes to a config.
+        Decorator and/or context manager to make temporary changes to a config.  Note that patched settings are thread-local.
 
         As a decorator:
 
@@ -686,28 +760,40 @@ def foo(...):
                 )
         if not isinstance(changes, dict):
             raise AssertionError(f"expected `dict` got {type(changes)}")
-        prior: dict[str, Any] = {}
         config = self
 
         class ConfigPatch(ContextDecorator):
             def __init__(self) -> None:
                 self.changes = changes
+                self._prior: ContextVar[tuple[dict[str, Any], ...]] = ContextVar(
+                    f"{config.__name__}.ConfigPatch[{id(self)}]",
+                    default=(),
+                )
 
             def __enter__(self) -> None:
-                if prior:
-                    raise AssertionError(
-                        "prior should be empty when entering ConfigPatch"
-                    )
+                prior: dict[str, Any] = {}
                 for key in self.changes:
                     # KeyError on invalid entry
                     prior[key] = config.__getattr__(key)
-                for k, v in self.changes.items():
-                    config.__setattr__(k, v)
+                prior_stack = self._prior.get()
+                self._prior.set((*prior_stack, prior))
+                try:
+                    for k, v in self.changes.items():
+                        config.__setattr__(k, v)
+                except Exception:
+                    self._prior.set(prior_stack)
+                    raise
 
             def __exit__(self, exc_type, exc_val, exc_tb):  # type: ignore[no-untyped-def]
+                prior_stack = self._prior.get()
+                if not prior_stack:
+                    raise AssertionError(
+                        "prior should not be empty when exiting ConfigPatch"
+                    )
+                prior = prior_stack[-1]
+                self._prior.set(prior_stack[:-1])
                 for k, v in prior.items():
                     config.__setattr__(k, v)
-                prior.clear()
 
         return ConfigPatch()
 
@@ -732,13 +818,13 @@ def _make_closure_patcher(self, **changes: dict[str, Any]) -> Any:
         config = self._config
 
         def change() -> Callable[[], None]:
-            prior = {k: config[k].user_override for k in changes}
+            prior = {k: config[k].user_override.get() for k in changes}
             for k, v in changes.items():
-                self._config[k].user_override = v
+                self._config[k].user_override.set(v)
 
             def revert() -> None:
                 for k, v in prior.items():
-                    self._config[k].user_override = v
+                    self._config[k].user_override.set(v)
 
             return revert
 
@@ -807,15 +893,6 @@ def __delattr__(self, name: str) -> None:
         return self._config.__delattr__(self._prefix + name)
 
 
-def patch_object(obj: object, name: str, value: object) -> object:
-    """
-    Workaround `mock.patch.object` issue with ConfigModule
-    """
-    if isinstance(obj, ConfigModule):
-        return obj.patch(name, value)
-    return mock.patch.object(obj, name, value)
-
-
 def get_tristate_env(name: str, default: Any = None) -> bool | None:
     value = os.environ.get(name)
     if value == "1":
diff --git a/torch/utils/_cxx_pytree.py b/torch/utils/_cxx_pytree.py
index d378d02edbf74..d13cc7eb156f1 100644
--- a/torch/utils/_cxx_pytree.py
+++ b/torch/utils/_cxx_pytree.py
@@ -1,7 +1,9 @@
+# Owner(s): ["module: pytree"]
+
 """
 Contains utility functions for working with nested python data structures.
 
-A *pytree* is Python nested data structure. It is a tree in the sense that
+A *pytree* is a Python nested data structure. It is a tree in the sense that
 nodes are Python collections (e.g., list, tuple, dict) and the leaves are
 Python values. Furthermore, a pytree should not contain reference cycles.
 
@@ -22,13 +24,21 @@
 import torch.utils._pytree as python_pytree
 from torch.torch_version import TorchVersion as _TorchVersion
 from torch.utils._pytree import (
+    Context,
+    DumpableContext,
+    FlattenFn,
+    FlattenWithKeysFn,
+    FromDumpableContextFn,
     is_namedtuple,
     is_namedtuple_class,
     is_namedtuple_instance,
     is_structseq,
     is_structseq_class,
     is_structseq_instance,
-    KeyEntry,
+    KeyPath,
+    PyTree,
+    ToDumpableContextFn,
+    UnflattenFn,
 )
 
 
@@ -49,8 +59,8 @@
 __all__ = [
     "PyTree",
     "Context",
-    "FlattenFunc",
-    "UnflattenFunc",
+    "FlattenFn",
+    "UnflattenFn",
     "DumpableContext",
     "ToDumpableContextFn",
     "FromDumpableContextFn",
@@ -101,22 +111,16 @@
 U = TypeVar("U")
 R = TypeVar("R")
 
-
 TreeSpec: TypeAlias = PyTreeSpec
+OpTreeUnflattenFn = Callable[[Context, Iterable[Any]], PyTree]
 
-Context = Any
-PyTree = Any
-FlattenFunc = Callable[[PyTree], tuple[list[Any], Context]]
-UnflattenFunc = Callable[[Iterable[Any], Context], PyTree]
-OpTreeUnflattenFunc = Callable[[Context, Iterable[Any]], PyTree]
-DumpableContext = Any  # Any json dumpable text
-ToDumpableContextFn = Callable[[Context], DumpableContext]
-FromDumpableContextFn = Callable[[DumpableContext], Context]
-KeyPath = tuple[KeyEntry, ...]
-FlattenWithKeysFunc = Callable[[PyTree], tuple[list[tuple[KeyEntry, Any]], Any]]
+# Keep deprecated alias for backward compatibility
+FlattenFunc = FlattenFn  # deprecated
+UnflattenFunc = UnflattenFn  # deprecated
+FlattenWithKeysFunc = FlattenWithKeysFn  # deprecated
 
 
-def _reverse_args(func: UnflattenFunc) -> OpTreeUnflattenFunc:
+def _reverse_args(func: UnflattenFn) -> OpTreeUnflattenFn:
     @functools.wraps(func)
     def wrapped(*args: Any, **kwargs: Any) -> Any:
         return func(*reversed(args), **kwargs)
@@ -126,13 +130,13 @@ def wrapped(*args: Any, **kwargs: Any) -> Any:
 
 def register_pytree_node(
     cls: type[Any],
-    flatten_fn: FlattenFunc,
-    unflatten_fn: UnflattenFunc,
+    flatten_fn: FlattenFn,
+    unflatten_fn: UnflattenFn,
     *,
     serialized_type_name: str | None = None,
     to_dumpable_context: ToDumpableContextFn | None = None,
     from_dumpable_context: FromDumpableContextFn | None = None,
-    flatten_with_keys_fn: FlattenWithKeysFunc | None = None,
+    flatten_with_keys_fn: FlattenWithKeysFn | None = None,
 ) -> None:
     """Register a container-like type as pytree node.
 
@@ -158,7 +162,7 @@ def register_pytree_node(
     Example::
 
         >>> # xdoctest: +SKIP
-        >>> # Registry a Python type with lambda functions
+        >>> # Register a Python type with lambda functions
         >>> register_pytree_node(
         ...     set,
         ...     lambda s: (sorted(s), None, None),
@@ -194,8 +198,8 @@ def register_pytree_node(
 )
 def _register_pytree_node(
     cls: type[Any],
-    flatten_fn: FlattenFunc,
-    unflatten_fn: UnflattenFunc,
+    flatten_fn: FlattenFn,
+    unflatten_fn: UnflattenFn,
     *,
     serialized_type_name: str | None = None,
     to_dumpable_context: ToDumpableContextFn | None = None,
@@ -245,8 +249,8 @@ def _register_pytree_node(
 
 def _private_register_pytree_node(
     cls: type[Any],
-    flatten_fn: FlattenFunc,
-    unflatten_fn: UnflattenFunc,
+    flatten_fn: FlattenFn,
+    unflatten_fn: UnflattenFn,
     *,
     serialized_type_name: str | None = None,
     to_dumpable_context: ToDumpableContextFn | None = None,
@@ -1017,16 +1021,22 @@ def treespec_loads(serialized: str) -> TreeSpec:
     return treespec
 
 
-class _DummyLeaf:
+class _Asterisk(str):
+    __slots__ = ()
+
+    def __new__(cls) -> Self:
+        return super().__new__(cls, "*")
+
     def __repr__(self) -> str:
-        return "*"
+        return "*"  # no quotes
+
+
+_asterisk = _Asterisk()
+del _Asterisk
 
 
 def treespec_pprint(treespec: TreeSpec) -> str:
-    dummy_tree = tree_unflatten(
-        [_DummyLeaf() for _ in range(treespec.num_leaves)],
-        treespec,
-    )
+    dummy_tree = tree_unflatten([_asterisk] * treespec.num_leaves, treespec)
     return repr(dummy_tree)
 
 
diff --git a/torch/utils/_debug_mode/__init__.py b/torch/utils/_debug_mode/__init__.py
new file mode 100644
index 0000000000000..fd40af116d139
--- /dev/null
+++ b/torch/utils/_debug_mode/__init__.py
@@ -0,0 +1,30 @@
+# mypy: allow-untyped-defs
+"""
+DebugMode: a debugging TorchDispatchMode that intercepts and logs runtime calls.
+
+See torch.utils._debug_mode._mode for the full implementation and docstring.
+"""
+
+from torch.utils._debug_mode._calls import (
+    _AnnotateCall,
+    _DebugCall,
+    _get_call_name,
+    _OpCall,
+    _OutputPlacementCall,
+    _RedistributeCall,
+    _TritonKernelCall,
+)
+from torch.utils._debug_mode._mode import (
+    DebugInterpreter,
+    DebugMode,
+    get_active_debug_mode,
+)
+from torch.utils._debug_mode._utils import (
+    _stringify_shape,
+    hash_tensor_fn,
+    norm_hash_fn,
+    TensorIdTracker,
+)
+
+
+__all__ = ["DebugMode", "get_active_debug_mode"]
diff --git a/torch/utils/_debug_mode/_calls.py b/torch/utils/_debug_mode/_calls.py
new file mode 100644
index 0000000000000..8da15e03bbb46
--- /dev/null
+++ b/torch/utils/_debug_mode/_calls.py
@@ -0,0 +1,354 @@
+# mypy: allow-untyped-defs
+"""
+Call record classes for DebugMode: _DebugCall and its subclasses.
+"""
+
+from typing import Any, TYPE_CHECKING
+
+import torch
+
+# Import _utils module for mutable globals
+from torch.utils._debug_mode import _utils
+from torch.utils._debug_mode._utils import (
+    _arg_to_str,
+    _get_op_name,
+    _get_stack_trace,
+    _maybe_get_autograd_trace,
+    REDISTRIBUTE_FUNC,
+    TensorIdTracker,
+)
+from torch.utils._pytree import tree_all, tree_map
+
+
+if TYPE_CHECKING:
+    from torch._dynamo.device_interface import DeviceInterface
+
+
+class _DebugCall:
+    """Base class for tracking operator calls in DebugMode"""
+
+    def __init__(
+        self,
+        call_depth: int,
+        record: dict[str, Any] | None = None,
+        log: dict[str, Any] | None = None,
+        stack: bool = False,
+    ) -> None:
+        self.call_depth = call_depth
+        if stack:
+            self.stack_trace = _get_stack_trace()
+            self.fwd_stack_trace = _maybe_get_autograd_trace()
+
+        # results from dispatch hooks
+        self.record = record
+        self.log = log
+        self.output_str: str | None = None
+
+    def stringify_args(
+        self, attributes: list[str], tensor_memo: TensorIdTracker | None = None
+    ) -> None:
+        """
+        To reduce memory consumption, this method stringifies args/kwargs, stores the result, and deletes original args/kwargs.
+        """
+        raise NotImplementedError(
+            "Subclasses must implement stringify_args(), even if no-op"
+        )
+
+    def stringify_output(
+        self,
+        output: Any,
+        attributes: list[str],
+        tensor_memo: TensorIdTracker | None = None,
+    ) -> None:
+        """Store stringified version of call output in self.output_str"""
+        if tree_all(lambda x: x is None, output):
+            return
+        output_str = tree_map(lambda x: _arg_to_str(x, attributes, tensor_memo), output)
+        self.output_str = f"  ->  {str(output_str)}"
+
+    def render(self, attributes: list[str]) -> str:
+        raise NotImplementedError("Subclasses must implement string render()")
+
+    def __repr__(self) -> str:
+        return self.render([])
+
+
+class _OpCall(_DebugCall):
+    """Normal operator call"""
+
+    def __init__(
+        self,
+        op,
+        args: tuple,
+        kwargs: dict,
+        call_depth: int,
+        stack: bool = False,
+    ) -> None:
+        super().__init__(call_depth, stack=stack)
+        self.op = op
+        self.args = args
+        self.kwargs = kwargs
+
+        self.args_str: str | None = None
+        self.kwargs_str: str | None = None
+
+    def stringify_args(
+        self, attributes: list[str], tensor_memo: TensorIdTracker | None = None
+    ) -> None:
+        self.args_str = ", ".join(
+            _arg_to_str(arg, attributes, tensor_memo) for arg in self.args
+        )
+        if self.kwargs:
+            self.kwargs_str = ", " + ", ".join(
+                f"{k}={_arg_to_str(v, attributes, tensor_memo)}"
+                for k, v in self.kwargs.items()
+            )
+        else:
+            self.kwargs_str = ""
+        del self.args
+        del self.kwargs
+
+    def render(self, attributes: list[str]) -> str:
+        if self.args_str is not None:
+            args_str = self.args_str
+        else:
+            args_str = ", ".join(_arg_to_str(arg, attributes) for arg in self.args)
+
+        if self.kwargs_str is not None:
+            kwargs_str = self.kwargs_str
+        else:
+            if self.kwargs:
+                kwargs_str = ", " + ", ".join(
+                    f"{k}={_arg_to_str(v, attributes)}" for k, v in self.kwargs.items()
+                )
+            else:
+                kwargs_str = ""
+
+        if isinstance(self.op, torch._ops.OpOverload):
+            op_name = self.op.__qualname__
+        elif hasattr(self.op, "__module__") and hasattr(self.op, "__name__"):
+            op_name = f"{self.op.__module__}.{self.op.__name__}"
+        else:
+            op_name = str(self.op)
+
+        base_str = f"{op_name}({args_str}{kwargs_str})"
+
+        if self.output_str:
+            base_str += self.output_str
+        if self.log:
+            base_str += f"  # {self.log}"
+        return base_str
+
+    def __iter__(self):
+        # for BC; tuple(self) returns (op, args, kwargs, call_depth)
+        if self.args_str is not None:
+            yield from [self.op, self.args_str, self.kwargs_str, self.call_depth]
+        else:
+            yield from [self.op, self.args, self.kwargs, self.call_depth]
+
+
+class _RedistributeCall(_DebugCall):
+    def __init__(
+        self,
+        arg,
+        src_placement,
+        dst_placement,
+        transform_info_str,
+        call_depth,
+        stack=False,
+        is_explicit=False,
+    ) -> None:
+        super().__init__(call_depth, stack=stack)
+        self.arg = arg
+        self.src_placement = src_placement
+        self.dst_placement = dst_placement
+        self.transform_info_str = transform_info_str
+        self.is_explicit = is_explicit
+        self.is_outer_call = isinstance(arg, int)
+
+        self.arg_str: str | None = None
+
+    def stringify_args(
+        self, attributes: list[str], tensor_memo: TensorIdTracker | None = None
+    ) -> None:
+        self.arg_str = f"{_arg_to_str(self.arg, attributes, tensor_memo)}"
+        del self.arg
+
+    def render(self, attributes: list[str]) -> str:
+        if self.arg_str is not None:
+            arg_str = self.arg_str
+        else:
+            arg_str = f"{_arg_to_str(self.arg, attributes)}"
+
+        if self.transform_info_str is not None:  # prioritize over src/dst placements
+            placement_str = f"trace: {self.transform_info_str}"
+        else:
+            src_placement_str = _arg_to_str(self.src_placement, attributes)
+            dst_placement_str = _arg_to_str(self.dst_placement, attributes)
+            placement_str = f"{src_placement_str} -> {dst_placement_str}"
+
+        # DebugMode will add redistribute_input logs at 2 levels,
+        # once per redistribute decision, and once per redistributed input.
+        # We only annotate [implicit/explicit] logs on the former (outer-level call).
+        if self.is_outer_call:
+            annotation = " [implicit] "
+        elif self.is_explicit:
+            annotation = " [explicit] "
+        else:
+            annotation = ""
+
+        base_str = f"{REDISTRIBUTE_FUNC}{annotation}({arg_str}, {placement_str})"
+        if self.output_str:
+            base_str += self.output_str
+        return base_str
+
+    def __iter__(self):
+        # for BC; tuple(self) returns (op, placement info, kwargs, call_depth)
+        if self.arg_str is not None:
+            arg = self.arg_str
+        else:
+            arg = self.arg
+
+        yield REDISTRIBUTE_FUNC
+        if self.transform_info_str:
+            yield [arg, self.transform_info_str]
+        else:
+            yield [arg, self.src_placement, self.dst_placement]
+        yield {}
+        yield self.call_depth
+
+
+class _OutputPlacementCall(_DebugCall):
+    """Records output placement for a DTensor op."""
+
+    def __init__(self, placements_str: str, call_depth: int) -> None:
+        super().__init__(call_depth)
+        self.placements_str = placements_str
+
+    def stringify_args(
+        self, attributes: list[str], tensor_memo: TensorIdTracker | None = None
+    ) -> None:
+        pass  # Already stringified
+
+    def render(self, attributes: list[str]) -> str:
+        return f"-> output: {self.placements_str}"
+
+
+class _TritonKernelCall(_DebugCall):
+    """Triton kernel call from Inductor"""
+
+    def __init__(
+        self,
+        kernel_name: str,
+        kwargs: dict[str, Any],
+        call_depth: int,
+    ):
+        super().__init__(call_depth)
+        self.kernel_name = kernel_name
+        self.kwargs = kwargs
+        self.kwargs_str: str | None = None
+
+        self.pre_hashes: dict[str, Any] | None = None
+        self.post_hashes: dict[str, Any] | None = None
+
+    def stringify_args(
+        self, attributes: list[str], tensor_memo: TensorIdTracker | None = None
+    ) -> None:
+        # Optionally hash kernel inputs before launch
+        if hash_fn := _utils._TRITON_INPUT_HASH_FN:
+            self.pre_hashes = {
+                k: hash_fn(v)
+                for k, v in self.kwargs.items()
+                if isinstance(v, torch.Tensor)
+            }
+
+        if self.kwargs:
+            self.kwargs_str = ", ".join(
+                f"{k}={_arg_to_str(v, attributes, tensor_memo)}"
+                for k, v in self.kwargs.items()
+            )
+        else:
+            self.kwargs_str = ""
+
+    def render(self, attributes: list[str]) -> str:
+        base_str = f"[triton] {self.kernel_name}({self.kwargs_str})"
+        if self.pre_hashes:
+            pre_hashes_str = ", ".join(f"{k}: {v}" for k, v in self.pre_hashes.items())
+            pre_hashes_str = (
+                "\n  "
+                + "  " * self.call_depth
+                + f"# pre-kernel hashes: {{{pre_hashes_str}}}"
+            )
+        else:
+            pre_hashes_str = ""
+        if self.post_hashes:
+            post_hashes_str = ", ".join(
+                f"{k}: {v}" for k, v in self.post_hashes.items()
+            )
+            post_hashes_str = (
+                "\n  "
+                + "  " * self.call_depth
+                + f"# post-kernel hashes: {{{post_hashes_str}}}"
+            )
+        else:
+            post_hashes_str = ""
+        return f"{base_str}{pre_hashes_str}{post_hashes_str}\n"
+
+    def finalize(self, device_interface: "DeviceInterface"):
+        # synchronize -> hash/store kernel results
+        device_interface.synchronize(device_interface.current_device())
+        if _utils._RECORD_TRITON_OUTPUTS:
+            self.record = {
+                "output": {
+                    k: v.clone() if isinstance(v, torch.Tensor) else v
+                    for k, v in self.kwargs.items()
+                }
+            }
+        if hash_fn := _utils._TRITON_OUTPUT_HASH_FN:
+            self.post_hashes = {
+                k: hash_fn(v)
+                for k, v in self.kwargs.items()
+                if isinstance(v, torch.Tensor)
+            }
+
+        # don't store tensors
+        del self.kwargs
+
+    def __iter__(self):
+        yield from [self.kernel_name, (), self.kwargs_str, self.call_depth]
+
+
+class _AnnotateCall(_DebugCall):
+    """Custom annotation call"""
+
+    def __init__(
+        self, tag: Any, header: str, call_depth: int, stack: bool = False
+    ) -> None:
+        super().__init__(call_depth, stack=stack)
+        self.tag = tag
+        self.header = header
+
+    def render(self, attributes: list[str]) -> str:
+        return f"[{self.header}] {self.tag}"
+
+    def __iter__(self):
+        yield from [
+            f"[{self.header}] {self.tag}",
+            (),
+            {},
+            self.call_depth,
+        ]
+
+
+def _get_call_name(call: _DebugCall) -> str:
+    """String identifying _DebugCall (e.g. func, kernel, module name)"""
+    if isinstance(call, _OpCall):
+        return _get_op_name(call.op)
+    elif isinstance(call, _TritonKernelCall):
+        return call.kernel_name
+    elif isinstance(call, _AnnotateCall):
+        return f"[{call.header}] {call.tag}"
+    elif isinstance(call, _RedistributeCall):
+        return REDISTRIBUTE_FUNC
+    else:
+        return str(call)
diff --git a/torch/utils/_debug_mode.py b/torch/utils/_debug_mode/_mode.py
similarity index 65%
rename from torch/utils/_debug_mode.py
rename to torch/utils/_debug_mode/_mode.py
index 429617cb47a8f..0d73b0847b8e8 100644
--- a/torch/utils/_debug_mode.py
+++ b/torch/utils/_debug_mode/_mode.py
@@ -34,19 +34,34 @@
 
 import contextlib
 import functools
-import inspect
 import logging
-import os
-import traceback
-import weakref
 from collections.abc import Callable
 from typing import Any, TYPE_CHECKING
 
 import torch
 from torch._logging import warning_once
 from torch._subclasses.fake_tensor import FakeTensor, FakeTensorMode
-from torch.fx.graph import _parse_stack_trace
-from torch.utils._dtype_abbrs import dtype_abbrs
+
+# Import _utils module for mutable globals
+from torch.utils._debug_mode import _utils
+from torch.utils._debug_mode._calls import (
+    _AnnotateCall,
+    _get_call_name,
+    _OpCall,
+    _OutputPlacementCall,
+    _RedistributeCall,
+    _TritonKernelCall,
+)
+from torch.utils._debug_mode._utils import (
+    _compute_rel_diff,
+    _get_user_stack_trace,
+    _run_dispatch_hooks,
+    _run_dispatch_pre_log_hooks,
+    _stringify_dtensor_spec,
+    hash_tensor_fn,
+    norm_hash_fn,
+    TensorIdTracker,
+)
 from torch.utils._python_dispatch import (
     _get_current_dispatch_mode,
     _get_current_dispatch_mode_stack,
@@ -59,210 +74,18 @@
     tree_map_only,
     tree_map_with_path,
 )
-from torch.utils._traceback import CapturedTraceback
-from torch.utils.weak import WeakIdRef
 
 
 if TYPE_CHECKING:
-    from torch._dynamo.device_interface import DeviceInterface
     from torch.distributed._tools.mod_tracker import ModTracker
 
 
 log = logging.getLogger(__name__)
 
-__all__ = ["DebugMode", "get_active_debug_mode"]
-
-
-REDISTRIBUTE_FUNC = "redistribute_input"
-# registered dispatch call hooks
-_DISPATCH_RECORD_HOOKS: list[Callable] = []
-_DISPATCH_LOG_HOOKS: list[Callable] = []
-_DISPATCH_PRE_LOG_HOOKS: list[Callable] = []
-# Tracks if we're in inductor benchmarking, and temporarily disables logging
-# (for ignoring autotuning kernel launches which don't affect the user-facing result)
-_IN_INDUCTOR_BENCHMARK = False
-# For record_outputs, log_tensor_hashes hooks for triton kernels.
-# Stores kernel outputs in call.record["output"]
-_RECORD_TRITON_OUTPUTS = False
-# Annotates kernel output hashes, and stores them in call.post_hashes
-_TRITON_OUTPUT_HASH_FN = None
-# Annotates kernel input hashes, and stores them in call.pre_hashes
-_TRITON_INPUT_HASH_FN = None
+
 # Counter for active DebugMode instances (fast path for get_active_debug_mode)
 _ACTIVE_DEBUG_MODE_COUNT = 0
 
-
-def _stringify_shape(shape) -> str:
-    return f"[{', '.join([str(x) for x in shape])}]"
-
-
-def _stringify_device_mesh(mesh) -> str:
-    return f"DM({', '.join([str(s) for s in mesh.shape])})"
-
-
-def _stringify_placement(placement) -> str:
-    return f"[{', '.join([str(p) for p in placement])}]"
-
-
-def _stringify_attributes(tensor, attributes) -> str:
-    pairs = {}
-    for attr in attributes:
-        if hasattr(tensor, attr):
-            pairs[attr] = getattr(tensor, attr)
-    if len(pairs) == 0:
-        return ""
-    return f"{{{', '.join([f'{k}={v}' for k, v in pairs.items()])}}}"
-
-
-def _stringify_dtensor_spec(spec) -> str:
-    from torch.distributed.tensor._dtensor_spec import DTensorSpec
-
-    return DTensorSpec.format_shard_order_str(spec.placements, spec.shard_order)
-
-
-class TensorIdTracker:
-    def __init__(self) -> None:
-        self.tensor_memo: dict[WeakIdRef, int] = {}
-        self.next_tensor_id = 0
-
-    def _id(self, tensor) -> int:
-        with torch._C._DisablePythonDispatcher():
-            o = WeakIdRef(tensor)
-
-            def del_memo() -> None:
-                self.tensor_memo.pop(o, None)
-
-            weakref.finalize(tensor, del_memo)
-            if o not in self.tensor_memo:
-                self.tensor_memo[o] = self.next_tensor_id
-                self.next_tensor_id += 1
-            return self.tensor_memo[o]
-
-
-def _tensor_debug_string(tensor, attributes, tensor_memo=None) -> str:
-    """Convert tensor to debug string representation."""
-
-    if isinstance(tensor, torch.Tensor):
-        tensor_debug_str = f"{dtype_abbrs[tensor.dtype]}{_stringify_shape(tensor.shape)}{_stringify_attributes(tensor, attributes)}"
-        id_str = f"${tensor_memo._id(tensor)}" if tensor_memo is not None else ""
-        if isinstance(tensor, torch.distributed.tensor.DTensor):
-            # omitted device mesh
-            return f"dt{id_str}: {tensor_debug_str}| {_stringify_dtensor_spec(tensor._spec)}"
-        elif isinstance(tensor, FakeTensor):
-            return f"ft{id_str}: {tensor_debug_str}"
-        else:
-            return f"t{id_str}: {tensor_debug_str}"
-    else:
-        raise RuntimeError(f"Unsupported tensor type: {type(tensor)}")
-
-
-def _arg_to_str(arg, attributes, tensor_memo=None) -> str:
-    from torch.distributed.tensor._dtensor_spec import DTensorSpec
-
-    def to_str(x):
-        if isinstance(x, torch.Tensor):
-            return _tensor_debug_string(x, attributes, tensor_memo)
-        elif isinstance(x, DTensorSpec):
-            return _stringify_dtensor_spec(x)
-        return x
-
-    arg = tree_map(to_str, arg)
-    return str(arg)
-
-
-def norm_hash_fn(t: torch.Tensor, use_scalar: bool = False) -> torch.Tensor | float:
-    """
-    from Observer. Computes a hash for a tensor by converting it to float (if needed), making it contiguous,
-    replacing NaN/inf values with fixed numbers, and then computing the L1 norm in float64 or complex128.
-    This is used to generate a deterministic summary value for tensor comparison.
-    """
-    with torch._C._DisablePythonDispatcher():
-        if not (t.is_floating_point() or t.is_complex()):
-            t = t.float()
-        t = t.contiguous()
-
-        if t.is_complex():
-            t_float = t.to(dtype=torch.complex128)
-        else:
-            t_float = t.to(dtype=torch.float64)
-
-        out = t_float.norm(p=1)
-        if use_scalar:
-            return out.item()
-        return out
-
-
-def _compute_rel_diff(hash1, hash2):
-    # Relative difference: |hash1 - hash2| / max(|hash1|, |hash2|, eps)
-    numerator = abs(hash1 - hash2)
-    denominator = max(abs(hash1), abs(hash2), 1e-10)
-    return numerator / denominator
-
-
-def hash_tensor_fn(t: torch.Tensor, use_scalar: bool = False) -> torch.Tensor | int:
-    """
-    wrapper over torch.hash_tensor
-    """
-    if isinstance(t, torch.distributed.tensor.DTensor):
-        t = t.to_local()
-
-    if t.is_floating_point():
-        t_clean = t.to(dtype=torch.float64)
-    elif t.is_complex():
-        t_clean = t.to(dtype=torch.complex128).view(torch.float64)
-    else:
-        t_clean = t.to(dtype=torch.int64)
-
-    if t.numel() > 0:
-        out = torch.hash_tensor(t_clean)
-    else:
-        out = torch.zeros((), device=t_clean.device, dtype=torch.uint64)
-
-    if use_scalar:
-        return out.item()  # type: ignore[attribute]
-    return out
-
-
-def _get_stack_trace() -> str:
-    from torch.fx.experimental.symbolic_shapes import uninteresting_files
-
-    summary = CapturedTraceback.extract().summary()
-    summary = summary[:-4]  # filter out DebugMode frames
-    summary = [
-        frame for frame in summary if frame.filename not in uninteresting_files()
-    ]
-    summary = traceback.StackSummary.from_list(summary)
-    return "".join(summary.format())
-
-
-def _get_user_stack_trace(stack_trace_str: str) -> str | None:
-    # Extract user code stack trace, filtering out torch internals.
-    torch_dir = os.path.dirname(inspect.getfile(torch))
-    filter_fn = lambda file, name, code: not file.startswith(torch_dir + os.path.sep)  # noqa: E731
-    trace = _parse_stack_trace(stack_trace_str, filter_fn=filter_fn)
-    if trace:
-        return f"File: {trace.file}:{trace.lineno} in {trace.name}, code: {trace.code}"
-    return None
-
-
-def _maybe_get_autograd_trace() -> str | None:
-    if torch._C._current_autograd_node() is not None:
-        tb = torch._C._current_autograd_node().metadata.get("traceback_")  # type: ignore[attr-defined]
-        if tb:
-            return "".join(tb)
-    return None
-
-
-def _get_op_name(op) -> str:
-    if isinstance(op, torch._ops.OpOverload):
-        op_name = op.__qualname__
-    elif hasattr(op, "__module__") and hasattr(op, "__name__"):
-        op_name = f"{op.__module__}.{op.__name__}"
-    else:
-        op_name = str(op)
-    return op_name
-
-
 _annotate_decorated = False
 
 
@@ -290,378 +113,6 @@ def _annotate_lowering(tag: str) -> None:
         _annotate_decorated = True
 
 
-class _DebugCall:
-    """Base class for tracking operator calls in DebugMode"""
-
-    def __init__(
-        self,
-        call_depth: int,
-        record: dict[str, Any] | None = None,
-        log: dict[str, Any] | None = None,
-        stack: bool = False,
-    ) -> None:
-        self.call_depth = call_depth
-        if stack:
-            self.stack_trace = _get_stack_trace()
-            self.fwd_stack_trace = _maybe_get_autograd_trace()
-
-        # results from dispatch hooks
-        self.record = record
-        self.log = log
-        self.output_str: str | None = None
-
-    def stringify_args(
-        self, attributes: list[str], tensor_memo: TensorIdTracker | None = None
-    ) -> None:
-        """
-        To reduce memory consumption, this method stringifies args/kwargs, stores the result, and deletes original args/kwargs.
-        """
-        raise NotImplementedError(
-            "Subclasses must implement stringify_args(), even if no-op"
-        )
-
-    def stringify_output(
-        self,
-        output: Any,
-        attributes: list[str],
-        tensor_memo: TensorIdTracker | None = None,
-    ) -> None:
-        """Store stringified version of call output in self.output_str"""
-        if tree_all(lambda x: x is None, output):
-            return
-        output_str = tree_map(lambda x: _arg_to_str(x, attributes, tensor_memo), output)
-        self.output_str = f"  ->  {str(output_str)}"
-
-    def render(self, attributes: list[str]) -> str:
-        raise NotImplementedError("Subclasses must implement string render()")
-
-    def __repr__(self) -> str:
-        return self.render([])
-
-
-class _OpCall(_DebugCall):
-    """Normal operator call"""
-
-    def __init__(
-        self,
-        op,
-        args: tuple,
-        kwargs: dict,
-        call_depth: int,
-        stack: bool = False,
-    ) -> None:
-        super().__init__(call_depth, stack=stack)
-        self.op = op
-        self.args = args
-        self.kwargs = kwargs
-
-        self.args_str: str | None = None
-        self.kwargs_str: str | None = None
-
-    def stringify_args(
-        self, attributes: list[str], tensor_memo: TensorIdTracker | None = None
-    ) -> None:
-        self.args_str = ", ".join(
-            _arg_to_str(arg, attributes, tensor_memo) for arg in self.args
-        )
-        if self.kwargs:
-            self.kwargs_str = ", " + ", ".join(
-                f"{k}={_arg_to_str(v, attributes, tensor_memo)}"
-                for k, v in self.kwargs.items()
-            )
-        else:
-            self.kwargs_str = ""
-        del self.args
-        del self.kwargs
-
-    def render(self, attributes: list[str]) -> str:
-        if self.args_str is not None:
-            args_str = self.args_str
-        else:
-            args_str = ", ".join(_arg_to_str(arg, attributes) for arg in self.args)
-
-        if self.kwargs_str is not None:
-            kwargs_str = self.kwargs_str
-        else:
-            if self.kwargs:
-                kwargs_str = ", " + ", ".join(
-                    f"{k}={_arg_to_str(v, attributes)}" for k, v in self.kwargs.items()
-                )
-            else:
-                kwargs_str = ""
-
-        if isinstance(self.op, torch._ops.OpOverload):
-            op_name = self.op.__qualname__
-        elif hasattr(self.op, "__module__") and hasattr(self.op, "__name__"):
-            op_name = f"{self.op.__module__}.{self.op.__name__}"
-        else:
-            op_name = str(self.op)
-
-        base_str = f"{op_name}({args_str}{kwargs_str})"
-
-        if self.output_str:
-            base_str += self.output_str
-        if self.log:
-            base_str += f"  # {self.log}"
-        return base_str
-
-    def __iter__(self):
-        # for BC; tuple(self) returns (op, args, kwargs, call_depth)
-        if self.args_str is not None:
-            yield from [self.op, self.args_str, self.kwargs_str, self.call_depth]
-        else:
-            yield from [self.op, self.args, self.kwargs, self.call_depth]
-
-
-class _RedistributeCall(_DebugCall):
-    def __init__(
-        self,
-        arg,
-        src_placement,
-        dst_placement,
-        transform_info_str,
-        call_depth,
-        stack=False,
-        is_explicit=False,
-    ) -> None:
-        super().__init__(call_depth, stack=stack)
-        self.arg = arg
-        self.src_placement = src_placement
-        self.dst_placement = dst_placement
-        self.transform_info_str = transform_info_str
-        self.is_explicit = is_explicit
-        self.is_outer_call = isinstance(arg, int)
-
-        self.arg_str: str | None = None
-
-    def stringify_args(
-        self, attributes: list[str], tensor_memo: TensorIdTracker | None = None
-    ) -> None:
-        self.arg_str = f"{_arg_to_str(self.arg, attributes, tensor_memo)}"
-        del self.arg
-
-    def render(self, attributes: list[str]) -> str:
-        if self.arg_str is not None:
-            arg_str = self.arg_str
-        else:
-            arg_str = f"{_arg_to_str(self.arg, attributes)}"
-
-        if self.transform_info_str is not None:  # prioritize over src/dst placements
-            placement_str = f"trace: {self.transform_info_str}"
-        else:
-            src_placement_str = _arg_to_str(self.src_placement, attributes)
-            dst_placement_str = _arg_to_str(self.dst_placement, attributes)
-            placement_str = f"{src_placement_str} -> {dst_placement_str}"
-
-        # DebugMode will add redistribute_input logs at 2 levels,
-        # once per redistribute decision, and once per redistributed input.
-        # We only annotate [implicit/explicit] logs on the former (outer-level call).
-        if self.is_outer_call:
-            annotation = " [implicit] "
-        elif self.is_explicit:
-            annotation = " [explicit] "
-        else:
-            annotation = ""
-
-        base_str = f"{REDISTRIBUTE_FUNC}{annotation}({arg_str}, {placement_str})"
-        if self.output_str:
-            base_str += self.output_str
-        return base_str
-
-    def __iter__(self):
-        # for BC; tuple(self) returns (op, placement info, kwargs, call_depth)
-        if self.arg_str is not None:
-            arg = self.arg_str
-        else:
-            arg = self.arg
-
-        yield REDISTRIBUTE_FUNC
-        if self.transform_info_str:
-            yield [arg, self.transform_info_str]
-        else:
-            yield [arg, self.src_placement, self.dst_placement]
-        yield {}
-        yield self.call_depth
-
-
-class _OutputPlacementCall(_DebugCall):
-    """Records output placement for a DTensor op."""
-
-    def __init__(self, placements_str: str, call_depth: int) -> None:
-        super().__init__(call_depth)
-        self.placements_str = placements_str
-
-    def stringify_args(
-        self, attributes: list[str], tensor_memo: TensorIdTracker | None = None
-    ) -> None:
-        pass  # Already stringified
-
-    def render(self, attributes: list[str]) -> str:
-        return f"-> output: {self.placements_str}"
-
-
-class _TritonKernelCall(_DebugCall):
-    """Triton kernel call from Inductor"""
-
-    def __init__(
-        self,
-        kernel_name: str,
-        kwargs: dict[str, Any],
-        call_depth: int,
-    ):
-        super().__init__(call_depth)
-        self.kernel_name = kernel_name
-        self.kwargs = kwargs
-        self.kwargs_str: str | None = None
-
-        self.pre_hashes: dict[str, Any] | None = None
-        self.post_hashes: dict[str, Any] | None = None
-
-    def stringify_args(
-        self, attributes: list[str], tensor_memo: TensorIdTracker | None = None
-    ) -> None:
-        # Optionally hash kernel inputs before launch
-        global _TRITON_INPUT_HASH_FN
-        if hash_fn := _TRITON_INPUT_HASH_FN:
-            self.pre_hashes = {
-                k: hash_fn(v)
-                for k, v in self.kwargs.items()
-                if isinstance(v, torch.Tensor)
-            }
-
-        if self.kwargs:
-            self.kwargs_str = ", ".join(
-                f"{k}={_arg_to_str(v, attributes, tensor_memo)}"
-                for k, v in self.kwargs.items()
-            )
-        else:
-            self.kwargs_str = ""
-
-    def render(self, attributes: list[str]) -> str:
-        base_str = f"[triton] {self.kernel_name}({self.kwargs_str})"
-        if self.pre_hashes:
-            pre_hashes_str = ", ".join(f"{k}: {v}" for k, v in self.pre_hashes.items())
-            pre_hashes_str = (
-                "\n  "
-                + "  " * self.call_depth
-                + f"# pre-kernel hashes: {{{pre_hashes_str}}}"
-            )
-        else:
-            pre_hashes_str = ""
-        if self.post_hashes:
-            post_hashes_str = ", ".join(
-                f"{k}: {v}" for k, v in self.post_hashes.items()
-            )
-            post_hashes_str = (
-                "\n  "
-                + "  " * self.call_depth
-                + f"# post-kernel hashes: {{{post_hashes_str}}}"
-            )
-        else:
-            post_hashes_str = ""
-        return f"{base_str}{pre_hashes_str}{post_hashes_str}\n"
-
-    def finalize(self, device_interface: "DeviceInterface"):
-        # synchronize -> hash/store kernel results
-        global _RECORD_TRITON_OUTPUTS, _TRITON_OUTPUT_HASH_FN
-        device_interface.synchronize(device_interface.current_device())
-        if _RECORD_TRITON_OUTPUTS:
-            self.record = {
-                "output": {
-                    k: v.clone() if isinstance(v, torch.Tensor) else v
-                    for k, v in self.kwargs.items()
-                }
-            }
-        if hash_fn := _TRITON_OUTPUT_HASH_FN:
-            self.post_hashes = {
-                k: hash_fn(v)
-                for k, v in self.kwargs.items()
-                if isinstance(v, torch.Tensor)
-            }
-
-        # don't store tensors
-        del self.kwargs
-
-    def __iter__(self):
-        yield from [self.kernel_name, (), self.kwargs_str, self.call_depth]
-
-
-class _AnnotateCall(_DebugCall):
-    """Custom annotation call"""
-
-    def __init__(
-        self, tag: Any, header: str, call_depth: int, stack: bool = False
-    ) -> None:
-        super().__init__(call_depth, stack=stack)
-        self.tag = tag
-        self.header = header
-
-    def render(self, attributes: list[str]) -> str:
-        return f"[{self.header}] {self.tag}"
-
-    def __iter__(self):
-        yield from [
-            f"[{self.header}] {self.tag}",
-            (),
-            {},
-            self.call_depth,
-        ]
-
-
-def _run_hook(hook, *args):
-    out = hook(*args)
-    if out is not None and not isinstance(out, dict):
-        raise AssertionError(f"hook must return None or dict, got {type(out).__name__}")
-    return out
-
-
-def _run_dispatch_pre_log_hooks(call: _DebugCall, func, types, args, kwargs) -> None:
-    global _DISPATCH_PRE_LOG_HOOKS
-    if _DISPATCH_PRE_LOG_HOOKS:
-        for hook in _DISPATCH_PRE_LOG_HOOKS:
-            hook_out = _run_hook(hook, func, types, args, kwargs, call)
-            if hook_out is not None:
-                # Store pre-hook results in call.log
-                if call.log is None:
-                    call.log = {}
-                call.log.update(hook_out)
-
-
-def _run_dispatch_hooks(call: _DebugCall, func, types, args, kwargs, result) -> None:
-    global _DISPATCH_RECORD_HOOKS, _DISPATCH_LOG_HOOKS
-    if _DISPATCH_RECORD_HOOKS:
-        record = {}
-        for hook in _DISPATCH_RECORD_HOOKS:
-            hook_out = _run_hook(hook, func, types, args, kwargs, result)
-            if hook_out is not None:
-                record.update(hook_out)
-        if record:
-            call.record = record
-
-    if _DISPATCH_LOG_HOOKS:
-        # Preserve existing log from pre-hooks (e.g., input_hash)
-        if call.log is None:
-            call.log = {}
-        for hook in _DISPATCH_LOG_HOOKS:
-            hook_out = _run_hook(hook, func, types, args, kwargs, result)
-            if hook_out is not None:
-                call.log.update(hook_out)
-
-
-def _get_call_name(call: _DebugCall) -> str:
-    """String identifying _DebugCall (e.g. func, kernel, module name)"""
-    if isinstance(call, _OpCall):
-        return _get_op_name(call.op)
-    elif isinstance(call, _TritonKernelCall):
-        return call.kernel_name
-    elif isinstance(call, _AnnotateCall):
-        return f"[{call.header}] {call.tag}"
-    elif isinstance(call, _RedistributeCall):
-        return REDISTRIBUTE_FUNC
-    else:
-        return str(call)
-
-
 @torch.library.custom_op("debug_mode_ops::annotate", mutates_args=())
 def _annotate(tag: str) -> None:
     # This is special-cased in DebugMode.__torch_dispatch__
@@ -832,15 +283,14 @@ def _track_op_output(self, op_index, result) -> None:
         self._output_info[op_index] = result
 
     # Without this override, running torch.compile under DebugMode
-    # will force torch.compile to always use the “eager” backend
+    # will force torch.compile to always use the "eager" backend
     # With this, DebugMode will not take effect on torch.compile
     @classmethod
     def ignore_compile_internals(cls) -> bool:
         return True
 
     def _record_call(self, call) -> None:
-        global _IN_INDUCTOR_BENCHMARK
-        if _IN_INDUCTOR_BENCHMARK:
+        if _utils._IN_INDUCTOR_BENCHMARK:
             return
 
         if str(call).startswith("profiler::_record_function"):
@@ -932,6 +382,7 @@ def __torch_dispatch__(self, func, types, args=(), kwargs=None):
             self._handle_annotate(args[0])
             return
 
+        from torch.distributed._functional_collectives import AsyncCollectiveTensor
         from torch.distributed._local_tensor import LocalTensor
 
         # Record the operation with its call depth
@@ -966,6 +417,17 @@ def __torch_dispatch__(self, func, types, args=(), kwargs=None):
                     stack=self.record_stack_trace,
                 )
                 self._record_call(call)
+        elif AsyncCollectiveTensor in types:
+            # Record AsyncCollectiveTensor operations so debugging/tracing tools can see them
+            if self.record_realtensor:
+                call = _OpCall(
+                    func,
+                    args,
+                    kwargs,
+                    self.call_depth + 1,
+                    stack=self.record_stack_trace,
+                )
+                self._record_call(call)
         elif len(types) == 0:
             if self.record_realtensor:
                 call = _OpCall(
@@ -1220,23 +682,21 @@ def dispatch_hooks(
         pre_log_hook signature is (func, types, args, kwargs, call) and is executed before
         the operation. It allows capturing state before in-place mutations.
         """
-        global _DISPATCH_RECORD_HOOKS, _DISPATCH_LOG_HOOKS, _DISPATCH_PRE_LOG_HOOKS
-
         if record_hook:
-            _DISPATCH_RECORD_HOOKS.append(record_hook)
+            _utils._DISPATCH_RECORD_HOOKS.append(record_hook)
         if log_hook:
-            _DISPATCH_LOG_HOOKS.append(log_hook)
+            _utils._DISPATCH_LOG_HOOKS.append(log_hook)
         if pre_log_hook:
-            _DISPATCH_PRE_LOG_HOOKS.append(pre_log_hook)
+            _utils._DISPATCH_PRE_LOG_HOOKS.append(pre_log_hook)
         try:
             yield
         finally:
             if record_hook:
-                _DISPATCH_RECORD_HOOKS.pop()
+                _utils._DISPATCH_RECORD_HOOKS.pop()
             if log_hook:
-                _DISPATCH_LOG_HOOKS.pop()
+                _utils._DISPATCH_LOG_HOOKS.pop()
             if pre_log_hook:
-                _DISPATCH_PRE_LOG_HOOKS.pop()
+                _utils._DISPATCH_PRE_LOG_HOOKS.pop()
 
     @staticmethod
     @contextlib.contextmanager
@@ -1251,14 +711,13 @@ def dispatch_hook(func, types, args, kwargs, result):
             )
             return {"output": out}
 
-        global _RECORD_TRITON_OUTPUTS
         try:
-            _old_record_triton = _RECORD_TRITON_OUTPUTS
-            _RECORD_TRITON_OUTPUTS = True
+            _old_record_triton = _utils._RECORD_TRITON_OUTPUTS
+            _utils._RECORD_TRITON_OUTPUTS = True
             with DebugMode.dispatch_hooks(record_hook=dispatch_hook):
                 yield
         finally:
-            _RECORD_TRITON_OUTPUTS = _old_record_triton
+            _utils._RECORD_TRITON_OUTPUTS = _old_record_triton
 
     @staticmethod
     @contextlib.contextmanager
@@ -1332,13 +791,12 @@ def _dispatch_post_hook(func, types, args, kwargs, result):
                 return None
             return out
 
-        global _TRITON_INPUT_HASH_FN, _TRITON_OUTPUT_HASH_FN
         try:
             if hash_inputs:
-                _old_input_hfn = _TRITON_INPUT_HASH_FN
-                _TRITON_INPUT_HASH_FN = fn
-            _old_output_hfn = _TRITON_OUTPUT_HASH_FN
-            _TRITON_OUTPUT_HASH_FN = fn
+                _old_input_hfn = _utils._TRITON_INPUT_HASH_FN
+                _utils._TRITON_INPUT_HASH_FN = fn
+            _old_output_hfn = _utils._TRITON_OUTPUT_HASH_FN
+            _utils._TRITON_OUTPUT_HASH_FN = fn
             with DebugMode.dispatch_hooks(
                 log_hook=_dispatch_post_hook,
                 pre_log_hook=_dispatch_pre_log_hook if hash_inputs else None,
@@ -1346,8 +804,8 @@ def _dispatch_post_hook(func, types, args, kwargs, result):
                 yield
         finally:
             if hash_inputs:
-                _TRITON_INPUT_HASH_FN = _old_input_hfn  # type: ignore[assignment]
-            _TRITON_OUTPUT_HASH_FN = _old_output_hfn
+                _utils._TRITON_INPUT_HASH_FN = _old_input_hfn  # type: ignore[assignment]
+            _utils._TRITON_OUTPUT_HASH_FN = _old_output_hfn
 
     @staticmethod
     @contextlib.contextmanager
@@ -1356,12 +814,11 @@ def _benchmarking_inductor():
         Context manager for disabling logging during inductor benchmarking,
         so logs don't contain all kernels launched from autotuning.
         """
-        global _IN_INDUCTOR_BENCHMARK
         try:
-            _IN_INDUCTOR_BENCHMARK = True
+            _utils._IN_INDUCTOR_BENCHMARK = True
             yield
         finally:
-            _IN_INDUCTOR_BENCHMARK = False
+            _utils._IN_INDUCTOR_BENCHMARK = False
 
     @property
     def logs(self):
diff --git a/torch/utils/_debug_mode/_utils.py b/torch/utils/_debug_mode/_utils.py
new file mode 100644
index 0000000000000..0eea2b7cabcd1
--- /dev/null
+++ b/torch/utils/_debug_mode/_utils.py
@@ -0,0 +1,251 @@
+# mypy: allow-untyped-defs
+"""
+Utility functions for DebugMode: tensor formatting, hashing, stack traces, and hook runners.
+"""
+
+import inspect
+import os
+import traceback
+import weakref
+from collections.abc import Callable
+from typing import TYPE_CHECKING
+
+import torch
+from torch._subclasses.fake_tensor import FakeTensor
+from torch.fx.graph import _parse_stack_trace
+from torch.utils._dtype_abbrs import dtype_abbrs
+from torch.utils._pytree import tree_map
+from torch.utils._traceback import CapturedTraceback
+from torch.utils.weak import WeakIdRef
+
+
+if TYPE_CHECKING:
+    from torch.utils._debug_mode._calls import _DebugCall
+
+
+REDISTRIBUTE_FUNC = "redistribute_input"
+
+# Tracks if we're in inductor benchmarking, and temporarily disables logging
+# (for ignoring autotuning kernel launches which don't affect the user-facing result)
+_IN_INDUCTOR_BENCHMARK: bool = False
+# For record_outputs, log_tensor_hashes hooks for triton kernels.
+# Stores kernel outputs in call.record["output"]
+_RECORD_TRITON_OUTPUTS: bool = False
+# Annotates kernel output hashes, and stores them in call.post_hashes
+_TRITON_OUTPUT_HASH_FN: Callable | None = None
+# Annotates kernel input hashes, and stores them in call.pre_hashes
+_TRITON_INPUT_HASH_FN: Callable | None = None
+
+# registered dispatch call hooks
+_DISPATCH_RECORD_HOOKS: list[Callable] = []
+_DISPATCH_LOG_HOOKS: list[Callable] = []
+_DISPATCH_PRE_LOG_HOOKS: list[Callable] = []
+
+
+def _stringify_shape(shape) -> str:
+    return f"[{', '.join([str(x) for x in shape])}]"
+
+
+def _stringify_device_mesh(mesh) -> str:
+    return f"DM({', '.join([str(s) for s in mesh.shape])})"
+
+
+def _stringify_placement(placement) -> str:
+    return f"[{', '.join([str(p) for p in placement])}]"
+
+
+def _stringify_attributes(tensor, attributes) -> str:
+    pairs = {}
+    for attr in attributes:
+        if hasattr(tensor, attr):
+            pairs[attr] = getattr(tensor, attr)
+    if len(pairs) == 0:
+        return ""
+    return f"{{{', '.join([f'{k}={v}' for k, v in pairs.items()])}}}"
+
+
+def _stringify_dtensor_spec(spec) -> str:
+    from torch.distributed.tensor._dtensor_spec import DTensorSpec
+
+    return DTensorSpec.format_shard_order_str(spec.placements, spec.shard_order)
+
+
+class TensorIdTracker:
+    def __init__(self) -> None:
+        self.tensor_memo: dict[WeakIdRef, int] = {}
+        self.next_tensor_id = 0
+
+    def _id(self, tensor) -> int:
+        with torch._C._DisablePythonDispatcher():
+            o = WeakIdRef(tensor)
+
+            def del_memo() -> None:
+                self.tensor_memo.pop(o, None)
+
+            weakref.finalize(tensor, del_memo)
+            if o not in self.tensor_memo:
+                self.tensor_memo[o] = self.next_tensor_id
+                self.next_tensor_id += 1
+            return self.tensor_memo[o]
+
+
+def _tensor_debug_string(tensor, attributes, tensor_memo=None) -> str:
+    """Convert tensor to debug string representation."""
+
+    if isinstance(tensor, torch.Tensor):
+        tensor_debug_str = f"{dtype_abbrs[tensor.dtype]}{_stringify_shape(tensor.shape)}{_stringify_attributes(tensor, attributes)}"
+        id_str = f"${tensor_memo._id(tensor)}" if tensor_memo is not None else ""
+        if isinstance(tensor, torch.distributed.tensor.DTensor):
+            # omitted device mesh
+            return f"dt{id_str}: {tensor_debug_str}| {_stringify_dtensor_spec(tensor._spec)}"
+        elif isinstance(tensor, FakeTensor):
+            return f"ft{id_str}: {tensor_debug_str}"
+        else:
+            return f"t{id_str}: {tensor_debug_str}"
+    else:
+        raise RuntimeError(f"Unsupported tensor type: {type(tensor)}")
+
+
+def _arg_to_str(arg, attributes, tensor_memo=None) -> str:
+    from torch.distributed.tensor._dtensor_spec import DTensorSpec
+
+    def to_str(x):
+        if isinstance(x, torch.Tensor):
+            return _tensor_debug_string(x, attributes, tensor_memo)
+        elif isinstance(x, DTensorSpec):
+            return _stringify_dtensor_spec(x)
+        return x
+
+    arg = tree_map(to_str, arg)
+    return str(arg)
+
+
+def norm_hash_fn(t: torch.Tensor, use_scalar: bool = False) -> torch.Tensor | float:
+    """
+    from Observer. Computes a hash for a tensor by converting it to float (if needed), making it contiguous,
+    replacing NaN/inf values with fixed numbers, and then computing the L1 norm in float64 or complex128.
+    This is used to generate a deterministic summary value for tensor comparison.
+    """
+    with torch._C._DisablePythonDispatcher():
+        if not (t.is_floating_point() or t.is_complex()):
+            t = t.float()
+        t = t.contiguous()
+
+        if t.is_complex():
+            t_float = t.to(dtype=torch.complex128)
+        else:
+            t_float = t.to(dtype=torch.float64)
+
+        out = t_float.norm(p=1)
+        if use_scalar:
+            return out.item()
+        return out
+
+
+def _compute_rel_diff(hash1, hash2):
+    # Relative difference: |hash1 - hash2| / max(|hash1|, |hash2|, eps)
+    numerator = abs(hash1 - hash2)
+    denominator = max(abs(hash1), abs(hash2), 1e-10)
+    return numerator / denominator
+
+
+def hash_tensor_fn(t: torch.Tensor, use_scalar: bool = False) -> torch.Tensor | int:
+    """
+    wrapper over torch.hash_tensor
+    """
+    if isinstance(t, torch.distributed.tensor.DTensor):
+        t = t.to_local()
+
+    if t.is_floating_point():
+        t_clean = t.to(dtype=torch.float64)
+    elif t.is_complex():
+        t_clean = t.to(dtype=torch.complex128).view(torch.float64)
+    else:
+        t_clean = t.to(dtype=torch.int64)
+
+    if t.numel() > 0:
+        out = torch.hash_tensor(t_clean)
+    else:
+        out = torch.zeros((), device=t_clean.device, dtype=torch.uint64)
+
+    if use_scalar:
+        return out.item()  # type: ignore[attribute]
+    return out
+
+
+def _get_stack_trace() -> str:
+    from torch.fx.experimental.symbolic_shapes import uninteresting_files
+
+    summary = CapturedTraceback.extract().summary()
+    summary = summary[:-4]  # filter out DebugMode frames
+    summary = [
+        frame for frame in summary if frame.filename not in uninteresting_files()
+    ]
+    summary = traceback.StackSummary.from_list(summary)
+    return "".join(summary.format())
+
+
+def _get_user_stack_trace(stack_trace_str: str) -> str | None:
+    # Extract user code stack trace, filtering out torch internals.
+    torch_dir = os.path.dirname(inspect.getfile(torch))
+    filter_fn = lambda file, name, code: not file.startswith(torch_dir + os.path.sep)  # noqa: E731
+    trace = _parse_stack_trace(stack_trace_str, filter_fn=filter_fn)
+    if trace:
+        return f"File: {trace.file}:{trace.lineno} in {trace.name}, code: {trace.code}"
+    return None
+
+
+def _maybe_get_autograd_trace() -> str | None:
+    if torch._C._current_autograd_node() is not None:
+        tb = torch._C._current_autograd_node().metadata.get("traceback_")  # type: ignore[attr-defined]
+        if tb:
+            return "".join(tb)
+    return None
+
+
+def _get_op_name(op) -> str:
+    if isinstance(op, torch._ops.OpOverload):
+        op_name = op.__qualname__
+    elif hasattr(op, "__module__") and hasattr(op, "__name__"):
+        op_name = f"{op.__module__}.{op.__name__}"
+    else:
+        op_name = str(op)
+    return op_name
+
+
+def _run_hook(hook, *args):
+    out = hook(*args)
+    if out is not None and not isinstance(out, dict):
+        raise AssertionError(f"hook must return None or dict, got {type(out).__name__}")
+    return out
+
+
+def _run_dispatch_pre_log_hooks(call: "_DebugCall", func, types, args, kwargs) -> None:
+    if _DISPATCH_PRE_LOG_HOOKS:
+        for hook in _DISPATCH_PRE_LOG_HOOKS:
+            hook_out = _run_hook(hook, func, types, args, kwargs, call)
+            if hook_out is not None:
+                # Store pre-hook results in call.log
+                if call.log is None:
+                    call.log = {}
+                call.log.update(hook_out)
+
+
+def _run_dispatch_hooks(call: "_DebugCall", func, types, args, kwargs, result) -> None:
+    if _DISPATCH_RECORD_HOOKS:
+        record = {}
+        for hook in _DISPATCH_RECORD_HOOKS:
+            hook_out = _run_hook(hook, func, types, args, kwargs, result)
+            if hook_out is not None:
+                record.update(hook_out)
+        if record:
+            call.record = record
+
+    if _DISPATCH_LOG_HOOKS:
+        # Preserve existing log from pre-hooks (e.g., input_hash)
+        if call.log is None:
+            call.log = {}
+        for hook in _DISPATCH_LOG_HOOKS:
+            hook_out = _run_hook(hook, func, types, args, kwargs, result)
+            if hook_out is not None:
+                call.log.update(hook_out)
diff --git a/torch/utils/_device.py b/torch/utils/_device.py
index 427fc878dd7a5..8b93bbef8193a 100644
--- a/torch/utils/_device.py
+++ b/torch/utils/_device.py
@@ -1,6 +1,5 @@
 # mypy: allow-untyped-defs
 import functools
-from typing import Optional
 
 import torch
 from torch._C import _len_torch_function_stack
@@ -54,15 +53,22 @@ def _device_constructors():
         torch.tensor,
         torch.as_tensor,
         torch.scalar_tensor,
+        # *_like may contain device kwarg, but the user implicitly
+        # expects a specific device even when kwarg unused.
+        # torch.zeros_like,
+        # torch.randint_like,
+        # torch.randn_like,
+        # torch.ones_like,
+        # torch.full_like,
+        # torch.empty_like,
     }
 
 
 # NB: This is directly called from C++ in torch/csrc/Device.cpp
 class DeviceContext(TorchFunctionMode):
     def __init__(self, device) -> None:
-        # pyrefly: ignore [read-only]
         self.device = torch.device(device)
-        self.prev_mode: Optional[DeviceContext] = None
+        self.prev_mode: DeviceContext | None = None
 
     def __enter__(self):
         global CURRENT_DEVICE
diff --git a/torch/utils/_functools.py b/torch/utils/_functools.py
index dd3a460efb1cc..215aa5103e445 100644
--- a/torch/utils/_functools.py
+++ b/torch/utils/_functools.py
@@ -36,12 +36,10 @@ def wrap(self: _C, *args: _P.args, **kwargs: _P.kwargs) -> _T:
         if not (cache := getattr(self, cache_name, None)):
             cache = {}
             setattr(self, cache_name, cache)
-        # pyrefly: ignore [unbound-name]
         cached_value = cache.get(args, _cache_sentinel)
         if cached_value is not _cache_sentinel:
             return cached_value
         value = f(self, *args, **kwargs)
-        # pyrefly: ignore [unbound-name]
         cache[args] = value
         return value
 
diff --git a/torch/utils/_get_clean_triton.py b/torch/utils/_get_clean_triton.py
index fbbabc3f50e62..a3f0de4f64aef 100644
--- a/torch/utils/_get_clean_triton.py
+++ b/torch/utils/_get_clean_triton.py
@@ -156,8 +156,8 @@ def get_clean_triton(
     """Run experiments and output results to file
 
     Args:
-        input_path (Optional[Path]): Path to inductor generated output codede
-        output_path (Optional[Path]): Path to write out the new python file
+        input_path (Path): Path to inductor generated output codede
+        output_path (Path): Path to write out the new python file
         auto_generate_params (bool): Whether to automatically generate launch_params if missing
     """
     return process_file(str(input_path), str(output_path), auto_generate_params)
diff --git a/torch/utils/_inspect.py b/torch/utils/_inspect.py
new file mode 100644
index 0000000000000..67d29b9bc7069
--- /dev/null
+++ b/torch/utils/_inspect.py
@@ -0,0 +1,90 @@
+import inspect
+from typing import Any
+
+
+def _signature_metadata(
+    sig: inspect.Signature,
+) -> tuple[tuple[inspect.Parameter, ...], bool, int]:
+    """
+    Returns tuple(sig.parameters.values()), if any has VAR_POSITIONAL or VAR_KEYWORD, and the max_positional
+    """
+    params = tuple(sig.parameters.values())
+    has_var_args = False
+    max_positional = 0
+
+    for p in params:
+        kind = p.kind
+        if kind in (inspect.Parameter.VAR_POSITIONAL, inspect.Parameter.VAR_KEYWORD):
+            has_var_args = True
+        if kind in (
+            inspect.Parameter.POSITIONAL_ONLY,
+            inspect.Parameter.POSITIONAL_OR_KEYWORD,
+        ):
+            max_positional += 1
+
+    return params, has_var_args, max_positional
+
+
+def _fast_bind(
+    sig: inspect.Signature, *args: Any, **kwargs: Any
+) -> inspect.BoundArguments:
+    """
+    Fast path for inspect.Signature.bind() for signatures without
+    VAR_POSITIONAL or VAR_KEYWORD parameters. Falls back to sig.bind()
+    for signatures that contain *args or **kwargs.
+    """
+    params, has_var_args, max_positional = _signature_metadata(sig)
+
+    # fallback for complex signatures
+    if has_var_args:
+        return sig.bind(*args, **kwargs)
+
+    len_args = len(args)
+
+    if len_args > max_positional:
+        raise TypeError(
+            f"Too many positional arguments: expected max {max_positional}, got {len_args}"
+        )
+
+    arguments: dict[str, Any] = {}
+    arg_i = 0
+
+    for p in params:
+        name = p.name
+        kind = p.kind
+
+        if kind is inspect.Parameter.POSITIONAL_ONLY:
+            if name in kwargs:
+                raise TypeError(
+                    f"Got some positional-only arguments passed as keyword arguments: '{name}'"
+                )
+            if arg_i < len_args:
+                arguments[name] = args[arg_i]
+                arg_i += 1
+            elif p.default is inspect.Parameter.empty:
+                raise TypeError(f"Missing required argument '{name}'")
+
+        elif kind is inspect.Parameter.POSITIONAL_OR_KEYWORD:
+            if arg_i < len_args:
+                if name in kwargs:
+                    raise TypeError(f"Multiple values for argument '{name}'")
+                arguments[name] = args[arg_i]
+                arg_i += 1
+            elif name in kwargs:
+                arguments[name] = kwargs[name]
+            elif p.default is inspect.Parameter.empty:
+                raise TypeError(f"Missing required argument '{name}'")
+
+        elif kind is inspect.Parameter.KEYWORD_ONLY:
+            if name in kwargs:
+                arguments[name] = kwargs[name]
+            elif p.default is inspect.Parameter.empty:
+                raise TypeError(f"Missing required argument '{name}'")
+
+    # disallow extra keyword arguments not in the signature
+    # cause kwargs have been processed by sig.bind at the beginning
+    for name in kwargs:
+        if name not in sig.parameters:
+            raise TypeError(f"Got an unexpected keyword argument '{name}'")
+
+    return inspect.BoundArguments(sig, arguments)  # type: ignore[arg-type]
diff --git a/torch/utils/_ordered_set.py b/torch/utils/_ordered_set.py
index f00b4ac31ef74..3d218395308f3 100644
--- a/torch/utils/_ordered_set.py
+++ b/torch/utils/_ordered_set.py
@@ -77,7 +77,6 @@ def _wrap_iter_in_set(cls, other: Any) -> Any:
     def pop(self) -> T:
         if not self:
             raise KeyError("pop from an empty set")
-        # pyrefly: ignore [bad-return]
         return self._dict.popitem()[0]
 
     def copy(self) -> OrderedSet[T]:
diff --git a/torch/utils/_pallas.py b/torch/utils/_pallas.py
index a53a40c8331de..84e04d2835e7b 100644
--- a/torch/utils/_pallas.py
+++ b/torch/utils/_pallas.py
@@ -83,12 +83,14 @@ def has_jax_tpu_backend() -> bool:
 
 @functools.cache
 def has_torch_tpu() -> bool:
-    """Check if torch_tpu is available."""
+    """Check if torch_tpu is installed and available."""
     try:
-        import torch_tpu  # noqa: F401  # type: ignore[import-not-found]
+        import torch_tpu.api  # noqa: F401  # type: ignore[import]
 
+        # Verify hardware/runtime access
+        torch_tpu.api.tpu_device()
         return True
-    except ImportError:
+    except (ImportError, RuntimeError):
         return False
 
 
@@ -107,7 +109,7 @@ def has_cuda_pallas() -> bool:
 @functools.cache
 def has_tpu_pallas() -> bool:
     """Checks for a full Pallas-on-TPU environment."""
-    return has_pallas_package() and has_torch_tpu()
+    return has_pallas_package() and has_jax_tpu_backend() and has_torch_tpu()
 
 
 @functools.cache
diff --git a/torch/utils/_pytree.py b/torch/utils/_pytree.py
index 3918576dc4f11..de90d98f77951 100644
--- a/torch/utils/_pytree.py
+++ b/torch/utils/_pytree.py
@@ -1,7 +1,9 @@
+# Owner(s): ["module: pytree"]
+
 """
 Contains utility functions for working with nested python data structures.
 
-A *pytree* is Python nested data structure. It is a tree in the sense that
+A *pytree* is a Python nested data structure. It is a tree in the sense that
 nodes are Python collections (e.g., list, tuple, dict) and the leaves are
 Python values. Furthermore, a pytree should not contain reference cycles.
 
@@ -20,6 +22,7 @@
 import importlib
 import importlib.metadata
 import json
+import logging
 import sys
 import threading
 import types
@@ -39,13 +42,15 @@
     TYPE_CHECKING,
     TypeAlias,
     TypeVar,
-    Union,
 )
 from typing_extensions import deprecated, NamedTuple, Self, TypeIs
 
 from torch.torch_version import TorchVersion as _TorchVersion
 
 
+log = logging.getLogger(__name__)
+
+
 if TYPE_CHECKING:
     import torch.utils._cxx_pytree as cxx_pytree
 
@@ -53,8 +58,8 @@
 __all__ = [
     "PyTree",
     "Context",
-    "FlattenFunc",
-    "UnflattenFunc",
+    "FlattenFn",
+    "UnflattenFn",
     "DumpableContext",
     "ToDumpableContextFn",
     "FromDumpableContextFn",
@@ -126,15 +131,20 @@ def default(self, obj: object) -> str | dict[str, Any]:
 
 Context = Any
 PyTree = Any
-FlattenFunc = Callable[[PyTree], tuple[list[Any], Context]]
-UnflattenFunc = Callable[[Iterable[Any], Context], PyTree]
+FlattenFn = Callable[[PyTree], tuple[list[Any], Context]]
+UnflattenFn = Callable[[Iterable[Any], Context], PyTree]
 DumpableContext = Any  # Any json dumpable text
 ToDumpableContextFn = Callable[[Context], DumpableContext]
 FromDumpableContextFn = Callable[[DumpableContext], Context]
-ToStrFunc = Callable[["TreeSpec", list[str]], str]
-MaybeFromStrFunc = Callable[[str], tuple[Any, Context, str] | None]
+ToStrFunc = Callable[["TreeSpec", list[str]], str]  # deprecated
+MaybeFromStrFunc = Callable[[str], tuple[Any, Context, str] | None]  # deprecated
 KeyPath = tuple[KeyEntry, ...]
-FlattenWithKeysFunc = Callable[[PyTree], tuple[list[tuple[KeyEntry, Any]], Any]]
+FlattenWithKeysFn = Callable[[PyTree], tuple[list[tuple[KeyEntry, Any]], Any]]
+
+# Keep deprecated alias for backward compatibility
+FlattenFunc = FlattenFn  # deprecated
+UnflattenFunc = UnflattenFn  # deprecated
+FlattenWithKeysFunc = FlattenWithKeysFn  # deprecated
 
 
 # A NodeDef holds two callables:
@@ -148,9 +158,9 @@ def default(self, obj: object) -> str | dict[str, Any]:
 #   pytree and returns a list of (keypath, value) pairs and a context.
 class NodeDef(NamedTuple):
     type: type[Any]
-    flatten_fn: FlattenFunc
-    unflatten_fn: UnflattenFunc
-    flatten_with_keys_fn: FlattenWithKeysFunc | None
+    flatten_fn: FlattenFn
+    unflatten_fn: UnflattenFn
+    flatten_with_keys_fn: FlattenWithKeysFn | None
 
 
 _NODE_REGISTRY_LOCK = threading.RLock()
@@ -201,13 +211,13 @@ class _SerializeNodeDef(NamedTuple):
 
 def register_pytree_node(
     cls: type[Any],
-    flatten_fn: FlattenFunc,
-    unflatten_fn: UnflattenFunc,
+    flatten_fn: FlattenFn,
+    unflatten_fn: UnflattenFn,
     *,
     serialized_type_name: str | None = None,
     to_dumpable_context: ToDumpableContextFn | None = None,
     from_dumpable_context: FromDumpableContextFn | None = None,
-    flatten_with_keys_fn: FlattenWithKeysFunc | None = None,
+    flatten_with_keys_fn: FlattenWithKeysFn | None = None,
 ) -> None:
     """Register a container-like type as pytree node.
 
@@ -527,15 +537,15 @@ def _register_namedtuple(
 )
 def _register_pytree_node(
     cls: type[Any],
-    flatten_fn: FlattenFunc,
-    unflatten_fn: UnflattenFunc,
+    flatten_fn: FlattenFn,
+    unflatten_fn: UnflattenFn,
     to_str_fn: ToStrFunc | None = None,  # deprecated
     maybe_from_str_fn: MaybeFromStrFunc | None = None,  # deprecated
     *,
     serialized_type_name: str | None = None,
     to_dumpable_context: ToDumpableContextFn | None = None,
     from_dumpable_context: FromDumpableContextFn | None = None,
-    flatten_with_keys_fn: FlattenWithKeysFunc | None = None,
+    flatten_with_keys_fn: FlattenWithKeysFn | None = None,
 ) -> None:
     """Register a container-like type as pytree node for the Python pytree only.
 
@@ -596,13 +606,13 @@ def _deregister_pytree_node(
 
 def _private_register_pytree_node(
     cls: type[Any],
-    flatten_fn: FlattenFunc,
-    unflatten_fn: UnflattenFunc,
+    flatten_fn: FlattenFn,
+    unflatten_fn: UnflattenFn,
     *,
     serialized_type_name: str | None = None,
     to_dumpable_context: ToDumpableContextFn | None = None,
     from_dumpable_context: FromDumpableContextFn | None = None,
-    flatten_with_keys_fn: FlattenWithKeysFunc | None = None,
+    flatten_with_keys_fn: FlattenWithKeysFn | None = None,
 ) -> None:
     """This is an internal function that is used to register a pytree node type
     for the Python pytree only. End-users should use :func:`register_pytree_node`
@@ -610,11 +620,25 @@ def _private_register_pytree_node(
     """
     from torch._library.opaque_object import is_opaque_type
 
-    if is_opaque_type(cls):
-        raise ValueError(
-            f"{cls} cannot be registered as a pytree as it has been "
-            "registered as an opaque object. Opaque objects must be pytree leaves."
-        )
+    if isinstance(cls, type) and is_opaque_type(cls):
+        # TODO: remove this allowance once downstream callers stop calling
+        # register_constant on Enum subclasses. Enums are now natively
+        # supported as opaque value types and don't need pytree registration.
+        import enum
+
+        if issubclass(cls, enum.Enum):
+            log.warning(
+                "%s is an Enum subclass and is now natively supported by "
+                "torch.compile as an opaque value type. Calling "
+                "register_constant() on Enum subclasses is deprecated and "
+                "will be an error in a future release.",
+                cls,
+            )
+        else:
+            raise ValueError(
+                f"{cls} cannot be registered as a pytree as it has been "
+                "registered as an opaque object. Opaque objects must be pytree leaves."
+            )
 
     with _NODE_REGISTRY_LOCK:
         if cls in SUPPORTED_NODES:
@@ -1025,6 +1049,7 @@ def _deque_unflatten(values: Iterable[T], context: Context) -> deque[T]:
 STANDARD_DICT_TYPES: frozenset[type] = frozenset({dict, OrderedDict, defaultdict})
 # pyrefly: ignore [no-matching-overload]
 BUILTIN_TYPES: frozenset[type] = frozenset(
+    # pyrefly: ignore [bad-argument-type]
     {
         tuple,
         list,
@@ -1101,6 +1126,8 @@ def _is_leaf(tree: PyTree, is_leaf: Callable[[PyTree], bool] | None = None) -> b
 #   is_leaf(): whether the root Node is a leaf
 @dataclasses.dataclass(init=False, frozen=True, eq=True, repr=False, slots=True)
 class TreeSpec:
+    """Representing the structure of the pytree."""
+
     type: Any
     _context: Context
     _children: list[Self]
@@ -1184,15 +1211,20 @@ def children_specs(self) -> list[Self]:
         return self._children
 
     def is_leaf(self) -> bool:
+        """Test whether the treespec represents a leaf."""
         return self.num_nodes == 1 and self.num_leaves == 1
 
     def children(self) -> list[Self]:
+        """Get all the child treespecs."""
         return self._children.copy()
 
     def child(self, index: int) -> Self:
+        """Get the child treespec at the given index."""
         return self._children[index]
 
     def flatten_up_to(self, tree: PyTree) -> list[PyTree]:
+        """Flatten the subtrees in ``tree`` up to the structure of this treespec and return a list of subtrees."""
+
         def helper(treespec: TreeSpec, node: PyTree, subtrees: list[PyTree]) -> None:
             if treespec.is_leaf():
                 subtrees.append(node)
@@ -1276,6 +1308,7 @@ def helper(treespec: TreeSpec, node: PyTree, subtrees: list[PyTree]) -> None:
         return subtrees
 
     def unflatten(self, leaves: Iterable[Any]) -> PyTree:
+        """Reconstruct a pytree from the leaves."""
         if not isinstance(leaves, (list, tuple)):
             leaves = list(leaves)
         if len(leaves) != self.num_leaves:
@@ -1333,14 +1366,20 @@ def __hash__(self) -> int:
     "use `isinstance(treespec, TreeSpec) and treespec.is_leaf()` instead.",
     category=FutureWarning,
 )
-@dataclasses.dataclass(init=True, frozen=True, eq=False, repr=False, slots=True)
+@dataclasses.dataclass(init=False, frozen=True, eq=False, repr=False, slots=True)
 class LeafSpec(TreeSpec):
     type: Any = dataclasses.field(default=None, init=False)
     _context: Context = dataclasses.field(default=None, init=False)
     _children: list[Self] = dataclasses.field(default_factory=list, init=False)
 
-    def __post_init__(self) -> None:
-        # Override `__post_init__` for `num_leaves` derivation.
+    def __init__(self) -> None:
+        # On Python 3.10.0, the dataclass-generated __init__ for frozen+slots
+        # dataclasses doesn't populate init=False fields with simple defaults,
+        # causing copy.deepcopy to fail with AttributeError. Avoid the bug by
+        # providing a custom __init__ (init=False above) instead.
+        object.__setattr__(self, "type", None)
+        object.__setattr__(self, "_context", None)
+        object.__setattr__(self, "_children", [])
         object.__setattr__(self, "num_nodes", 1)
         object.__setattr__(self, "num_leaves", 1)
         object.__setattr__(self, "num_children", 0)
@@ -1385,7 +1424,7 @@ def treespec_dict(
 
 def _is_pytreespec_instance(
     obj: Any,
-) -> TypeIs[Union[TreeSpec, "cxx_pytree.PyTreeSpec"]]:
+) -> TypeIs["TreeSpec | cxx_pytree.PyTreeSpec"]:
     if isinstance(obj, TreeSpec):
         return True
     if "torch.utils._cxx_pytree" in sys.modules:
@@ -1407,7 +1446,7 @@ def _is_pytreespec_instance(
 
 
 def _ensure_python_treespec_instance(
-    treespec: Union[TreeSpec, "cxx_pytree.PyTreeSpec"],
+    treespec: "TreeSpec | cxx_pytree.PyTreeSpec",
 ) -> TreeSpec:
     if isinstance(treespec, TreeSpec):
         return treespec
@@ -2054,16 +2093,22 @@ def treespec_loads(serialized: str) -> TreeSpec:
     )
 
 
-class _DummyLeaf:
+class _Asterisk(str):
+    __slots__ = ()
+
+    def __new__(cls) -> Self:
+        return super().__new__(cls, "*")
+
     def __repr__(self) -> str:
-        return "*"
+        return "*"  # no quotes
+
+
+_asterisk = _Asterisk()
+del _Asterisk
 
 
 def treespec_pprint(treespec: TreeSpec) -> str:
-    dummy_tree = tree_unflatten(
-        [_DummyLeaf() for _ in range(treespec.num_leaves)],
-        treespec,
-    )
+    dummy_tree = tree_unflatten([_asterisk] * treespec.num_leaves, treespec)
     return repr(dummy_tree)
 
 
diff --git a/torch/utils/_runtime_estimation.py b/torch/utils/_runtime_estimation.py
index 9efe613cefc85..b2d56dcfc7f79 100644
--- a/torch/utils/_runtime_estimation.py
+++ b/torch/utils/_runtime_estimation.py
@@ -1,8 +1,7 @@
 import torch
 from torch._inductor.utils import get_device_tflops, get_gpu_dram_gbps
 from torch.fx.experimental.symbolic_shapes import (
-    has_hint,
-    size_hint,
+    optimization_hint,
     statically_known_true,
 )
 from torch.utils._ordered_set import OrderedSet
@@ -121,12 +120,9 @@ def get_num_bytes(t: torch.Tensor) -> int:
     """
     real_numel = 1
     for size, stride in zip(t.shape, t.stride()):
-        if not has_hint(size) or not has_hint(stride):
-            return 0
-
         # For dims with stride=0 (expanded/broadcast), only 1 element accessed
         if not statically_known_true(stride == 0):
-            real_numel *= size_hint(size)
+            real_numel *= optimization_hint(size, fallback=0)
 
     return real_numel * t.element_size()
 
diff --git a/torch/utils/_sympy/functions.py b/torch/utils/_sympy/functions.py
index ee04c2461f3c6..1f95d5337f6a1 100644
--- a/torch/utils/_sympy/functions.py
+++ b/torch/utils/_sympy/functions.py
@@ -200,12 +200,10 @@ class FloorDiv(sympy.Function):
 
     @property
     def base(self) -> sympy.Basic:
-        # pyrefly: ignore [missing-attribute]
         return self.args[0]
 
     @property
     def divisor(self) -> sympy.Basic:
-        # pyrefly: ignore [missing-attribute]
         return self.args[1]
 
     def _sympystr(self, printer: sympy.printing.StrPrinter) -> str:
@@ -305,7 +303,6 @@ def eval(cls, base: sympy.Integer, divisor: sympy.Integer) -> sympy.Basic | None
         return None
 
     def _eval_is_nonnegative(self) -> bool | None:
-        # pyrefly: ignore [missing-attribute]
         p, q = self.args[:2]
         if all([p.is_integer, q.is_integer, p.is_nonnegative, q.is_nonnegative]):
             return True
@@ -374,7 +371,6 @@ def eval(
         return None
 
     def _eval_is_nonnegative(self) -> bool | None:
-        # pyrefly: ignore [missing-attribute]
         p, q = self.args[:2]
         return fuzzy_eq(p.is_nonnegative, q.is_nonnegative)  # type: ignore[attr-defined]
 
@@ -470,11 +466,8 @@ def _eval_is_nonpositive(self) -> bool | None:
         return True if self.args[1].is_negative else None  # type: ignore[attr-defined]
 
     def _ccode(self, printer) -> str:
-        # pyrefly: ignore [missing-attribute]
         p = printer.parenthesize(self.args[0], PRECEDENCE["Atom"] - 0.5)
-        # pyrefly: ignore [missing-attribute]
         q = printer.parenthesize(self.args[1], PRECEDENCE["Atom"] - 0.5)
-        # pyrefly: ignore [missing-attribute]
         abs_q = str(q) if self.args[1].is_positive else f"abs({q})"
         return f"({p} % {q}) < 0 ? {p} % {q} + {abs_q} : {p} % {q}"
 
@@ -557,7 +550,6 @@ def eval(cls, number):
             return sympy.Integer(math.ceil(float(number)))
 
     def _ccode(self, printer) -> str:
-        # pyrefly: ignore [missing-attribute]
         number = printer.parenthesize(self.args[0], self.args[0].precedence - 0.5)
         return f"ceil({number})"
 
@@ -723,7 +715,7 @@ def _unique_symbols(
         Return seen_symbols if all atoms in all args are all unique symbols,
         else returns None. initial_set can be used to represent initial value for seen_symbols
         """
-        seen_symbols = set() if initial_set is None else initial_set
+        seen_symbols = set() if initial_set is None else initial_set.copy()
         for arg in args:
             for element in arg.atoms():
                 if not isinstance(element, sympy.core.symbol.Symbol):
@@ -828,7 +820,6 @@ def do(ai, a):
             if not cond:
                 return ai.func(*[do(i, a) for i in ai.args], evaluate=False)
             if isinstance(ai, cls):
-                # pyrefly: ignore [missing-attribute]
                 return ai.func(*[do(i, a) for i in ai.args if i != a], evaluate=False)
             return a
 
@@ -1006,7 +997,6 @@ def _eval_is_nonnegative(self):  # type:ignore[override]
         return fuzzy_or(a.is_nonnegative for a in self.args)  # type: ignore[attr-defined]
 
     def _eval_is_negative(self):  # type:ignore[override]
-        # pyrefly: ignore [missing-attribute]
         return fuzzy_and(a.is_negative for a in self.args)
 
 
@@ -1025,7 +1015,6 @@ def _eval_is_nonnegative(self):  # type:ignore[override]
         return fuzzy_and(a.is_nonnegative for a in self.args)  # type: ignore[attr-defined]
 
     def _eval_is_negative(self):  # type:ignore[override]
-        # pyrefly: ignore [missing-attribute]
         return fuzzy_or(a.is_negative for a in self.args)
 
 
@@ -1160,9 +1149,7 @@ def eval(cls, base, divisor):
             return sympy.Float(int(base) / int(divisor))
 
     def _ccode(self, printer) -> str:
-        # pyrefly: ignore [missing-attribute]
         base = printer.parenthesize(self.args[0], PRECEDENCE["Atom"] - 0.5)
-        # pyrefly: ignore [missing-attribute]
         divisor = printer.parenthesize(self.args[1], PRECEDENCE["Atom"] - 0.5)
         return f"((int){base}/(int){divisor})"
 
@@ -1329,16 +1316,13 @@ class Identity(sympy.Function):
     precedence = 10
 
     def __repr__(self) -> str:  # type: ignore[override]
-        # pyrefly: ignore [missing-attribute]
         return f"Identity({self.args[0]})"
 
     def _sympystr(self, printer) -> str:
         """Controls how sympy's StrPrinter prints this"""
-        # pyrefly: ignore [missing-attribute]
         return f"({printer.doprint(self.args[0])})"
 
     def _eval_is_real(self):
-        # pyrefly: ignore [missing-attribute]
         return self.args[0].is_real
 
     def _eval_is_integer(self):
@@ -1348,26 +1332,53 @@ def _eval_is_integer(self):
     def is_number(self):
         # Treat Identity as numeric only when the argument is comparable.
         # This avoids creating numeric non-comparable Identity(I) terms.
-        # pyrefly: ignore [missing-attribute]
         return bool(self.args[0].is_number and self.args[0].is_comparable)
 
     @property
     def is_comparable(self):
         # Delegate comparability to the wrapped argument.
-        # pyrefly: ignore [missing-attribute]
         return bool(self.args[0].is_comparable)
 
     def _eval_expand_identity(self, **hints):
         # Removes the identity op.
-        # pyrefly: ignore [missing-attribute]
         return self.args[0]
 
     def __int__(self) -> int:
-        # pyrefly: ignore [missing-attribute]
         return int(self.args[0])
 
+    def _identity_atom_compare(self, other, op):
+        """
+        Fast path for comparing wrapped numeric atomics against other numeric atomics.
+        Keep compound expressions on SymPy's default symbolic path.
+        """
+        arg = self.args[0]
+        if isinstance(other, int):
+            other = sympy.Integer(other)
+        if not isinstance(other, sympy.Expr):
+            return None
+        if not (arg.is_Atom and arg.is_number and arg.is_comparable):
+            return None
+        if not (other.is_Atom and other.is_number and other.is_comparable):
+            return None
+        return sympy.S.true if op(arg, other) else sympy.S.false
+
+    def __ge__(self, other):
+        out = self._identity_atom_compare(other, lambda a, b: a >= b)
+        return out if out is not None else super().__ge__(other)
+
+    def __gt__(self, other):
+        out = self._identity_atom_compare(other, lambda a, b: a > b)
+        return out if out is not None else super().__gt__(other)
+
+    def __le__(self, other):
+        out = self._identity_atom_compare(other, lambda a, b: a <= b)
+        return out if out is not None else super().__le__(other)
+
+    def __lt__(self, other):
+        out = self._identity_atom_compare(other, lambda a, b: a < b)
+        return out if out is not None else super().__lt__(other)
+
     def __float__(self) -> float:
-        # pyrefly: ignore [missing-attribute]
         return float(self.args[0])
 
 
diff --git a/torch/utils/_sympy/printers.py b/torch/utils/_sympy/printers.py
index 61617c8aececf..e9ca80812b04d 100644
--- a/torch/utils/_sympy/printers.py
+++ b/torch/utils/_sympy/printers.py
@@ -70,7 +70,6 @@ def _print_Float(self, expr: sympy.Expr) -> str:
     # NB: this pow by natural, you should never have used builtin sympy.pow
     # for FloatPow, and a symbolic exponent should be PowByNatural.  These
     # means exp is guaranteed to be integer.
-    # pyrefly: ignore [bad-override]
     def _print_Pow(self, expr: sympy.Expr) -> str:
         base, exp = expr.args
         if exp != int(exp):
@@ -98,6 +97,9 @@ def _print_NegativeInfinity(self, expr: sympy.Expr) -> str:
             f"_print_NegativeInfinity not implemented for {type(self)}"
         )
 
+    def _print_NaN(self, expr: sympy.Expr) -> str:
+        raise NotImplementedError(f"_print_NaN not implemented for {type(self)}")
+
     def _print_FloorDiv(self, expr: sympy.Expr) -> str:
         raise NotImplementedError(f"_print_FloorDiv not implemented for {type(self)}")
 
@@ -177,6 +179,9 @@ def _print_Infinity(self, expr: sympy.Expr) -> str:
     def _print_NegativeInfinity(self, expr: sympy.Expr) -> str:
         return "-math.inf"
 
+    def _print_NaN(self, expr: sympy.Expr) -> str:
+        return "math.nan"
+
     # WARNING: this is dangerous for Triton, which has C-style modulus
     def _print_PythonMod(self, expr: sympy.Expr) -> str:
         return self.stringify(expr.args, " % ", PRECEDENCE["Atom"] - 0.5)
@@ -192,8 +197,17 @@ def _print_IntTrueDiv(self, expr: sympy.Expr) -> str:
         return self.stringify(expr.args, " / ", PRECEDENCE["Atom"] - 0.5)
 
     def _helper_sqrt(self, expr: sympy.Expr) -> str:
-        # pyrefly: ignore [missing-attribute]
-        return f"math.sqrt({self._print(expr)})"
+        # NB: We use torch._sym_sqrt here instead of math.sqrt because the
+        # guard expression may be evaluated with SymInt/SymFloat inputs (e.g.
+        # during cache hit re-evaluation in evaluate_guards_expression).
+        # math.sqrt on a SymFloat triggers evaluate_expr which forces
+        # concretization/specialization of the symbol, creating spurious
+        # guards that didn't exist in the original program.
+        # torch._sym_sqrt properly propagates through the symbolic system
+        # without forcing specialization.
+        # See https://github.com/pytorch/pytorch/issues/152435
+        # pyrefly: ignore [missing-attribute]
+        return f"torch._sym_sqrt({self._print(expr)})"
 
     def _print_OpaqueUnaryFn_sqrt(self, expr: sympy.Expr) -> str:
         return self._helper_sqrt(expr.args[0])
@@ -373,6 +387,9 @@ def _print_Where(self, expr: sympy.Expr) -> str:
         )
         return f"{c} ? {p} : {q}"
 
+    def _print_Or(self, expr: sympy.Expr) -> str:
+        return self.stringify(expr.args, " || ", precedence(expr))
+
     def _print_Piecewise(self, expr: sympy.Expr) -> str:
         # Convert Piecewise(expr_cond_pairs) to nested ternary operators
         # Piecewise((e1, c1), (e2, c2), ..., (eN, cN))
@@ -644,3 +661,6 @@ def _print_Infinity(self, expr: sympy.Expr) -> str:
 
     def _print_NegativeInfinity(self, expr: sympy.Expr) -> str:
         return f"-{self._print_Infinity(expr)}"
+
+    def _print_NaN(self, expr: sympy.Expr) -> str:
+        return "std::numeric_limits<double>::quiet_NaN()"
diff --git a/torch/utils/_sympy/reference.py b/torch/utils/_sympy/reference.py
index 3118e89ebfcf0..3879cef698154 100644
--- a/torch/utils/_sympy/reference.py
+++ b/torch/utils/_sympy/reference.py
@@ -177,7 +177,7 @@ def sqrt(x):
 
     @staticmethod
     def pow(a, b):
-        # pyrefly: ignore [bad-argument-type]
+        # pyrefly: ignore [bad-argument-count, bad-argument-type]
         return _keep_float(FloatPow)(a, b)
 
     @staticmethod
@@ -340,6 +340,20 @@ def bitwise_or(a, b):
     def bitwise_xor(a, b):
         return a ^ b
 
+    @staticmethod
+    def expr_cond_pair(expr, cond):
+        return (expr, cond)
+
+    @staticmethod
+    def piecewise(*pairs):
+        # Build nested sym_ite from right to left.
+        # Piecewise((e1, c1), (e2, c2), ..., (en, True)) becomes
+        # sym_ite(c1, e1, sym_ite(c2, e2, ... en))
+        result = pairs[-1][0]
+        for expr, cond in reversed(pairs[:-1]):
+            result = torch.sym_ite(cond, expr, result)
+        return result
+
 
 # Like PythonReferenceAnalysis, but some export-unfriendly choices of
 # operators to make things faster
diff --git a/torch/utils/_triton.py b/torch/utils/_triton.py
index 98de7bbcc5868..075aa5a332234 100644
--- a/torch/utils/_triton.py
+++ b/torch/utils/_triton.py
@@ -1,5 +1,6 @@
 import functools
 import hashlib
+import os
 from typing import Any
 
 
@@ -193,6 +194,24 @@ def triton_backend() -> Any:
     return make_backend(target)
 
 
+def _extern_libs_key(backend: Any) -> str:
+    """Return a cache key fragment for extern libs (e.g. libdevice.10.bc).
+
+    These files affect codegen but are not covered by triton_key() (Python
+    sources only) or backend.hash() (ptxas version and arch only).
+    """
+    opts = backend.parse_options({})
+    extern_libs = getattr(opts, "extern_libs", None)
+    if not extern_libs:
+        return ""
+    parts = []
+    for name, path in sorted(extern_libs):
+        if os.path.isfile(path):
+            with open(path, "rb") as f:
+                parts.append(f"{name}-{hashlib.sha256(f.read()).hexdigest()}")
+    return "-".join(parts)
+
+
 @functools.cache
 def triton_hash_with_backend() -> str:
     from torch._inductor.runtime.triton_compat import triton_key
diff --git a/torch/utils/backend_registration.py b/torch/utils/backend_registration.py
index 0e1eb82f8869e..f775cc0430d3c 100644
--- a/torch/utils/backend_registration.py
+++ b/torch/utils/backend_registration.py
@@ -186,9 +186,10 @@ def _generate_module_methods_for_privateuse1_backend(custom_backend_name: str) -
         )
 
     def wrap_module_to(
+        # pyrefly: ignore [invalid-type-var]
         self: torch.nn.modules.module.T,
         device: int | torch.device | None = None,
-    ) -> torch.nn.modules.module.T:
+    ) -> torch.nn.modules.module.T:  # pyrefly: ignore [invalid-type-var]
         r"""Move all model parameters and buffers to the custom device.
 
         This also makes associated parameters and buffers different objects. So
diff --git a/torch/utils/benchmark/utils/compile.py b/torch/utils/benchmark/utils/compile.py
index b96ca3c86769b..dd15a582a2749 100644
--- a/torch/utils/benchmark/utils/compile.py
+++ b/torch/utils/benchmark/utils/compile.py
@@ -89,14 +89,11 @@ def benchmark_compile(
             try:
                 torch._dynamo.reset()
                 compile_counter_with_backend = CompileCounterWithBackend(backend)
-                # pyrefly: ignore [no-matching-overload]
                 opt_model = torch.compile(model, backend=compile_counter_with_backend, mode=mode)
 
                 # Compilation only happens after the first inference
-                # pyrefly: ignore [bad-argument-type]
                 compilation_time = bench_loop(opt_model, sample_input, 1, optimizer, loss_fn)
 
-                # pyrefly: ignore [bad-argument-type]
                 running_time = bench_loop(opt_model, sample_input, num_iters, optimizer, loss_fn)
 
                 if compile_counter_with_backend.frame_count == 0:
@@ -112,7 +109,6 @@ def benchmark_compile(
         else:
             opt_model = model
             compilation_time = None
-            # pyrefly: ignore [bad-argument-type]
             running_time = bench_loop(opt_model, sample_input, num_iters, optimizer, loss_fn)
 
         compilation_time = round(compilation_time, 2) if compilation_time else None
diff --git a/torch/utils/benchmark/utils/sparse_fuzzer.py b/torch/utils/benchmark/utils/sparse_fuzzer.py
index a2a573b9b44fd..cf38c11240b27 100644
--- a/torch/utils/benchmark/utils/sparse_fuzzer.py
+++ b/torch/utils/benchmark/utils/sparse_fuzzer.py
@@ -91,7 +91,6 @@ def sparse_tensor_constructor(size, dtype, sparse_dim, nnz, is_coalesced):
         return x
 
     def _make_tensor(self, params, state):
-        # pyrefly: ignore [missing-attribute]
         size, _, _ = self._get_size_and_steps(params)
         density = params['density']
         nnz = math.ceil(sum(size) * density)
@@ -101,10 +100,8 @@ def _make_tensor(self, params, state):
         is_coalesced = params['coalesced']
         sparse_dim = params['sparse_dim'] if self._sparse_dim else len(size)
         sparse_dim = min(sparse_dim, len(size))
-        # pyrefly: ignore [missing-attribute]
         tensor = self.sparse_tensor_constructor(size, self._dtype, sparse_dim, nnz, is_coalesced)
 
-        # pyrefly: ignore [missing-attribute]
         if self._cuda:
             tensor = tensor.cuda()
         sparse_dim = tensor.sparse_dim()
@@ -120,7 +117,6 @@ def _make_tensor(self, params, state):
             "sparse_dim": sparse_dim,
             "dense_dim": dense_dim,
             "is_hybrid": is_hybrid,
-            # pyrefly: ignore [missing-attribute]
             "dtype": str(self._dtype),
         }
         return tensor, properties
diff --git a/torch/utils/benchmark/utils/valgrind_wrapper/timer_interface.py b/torch/utils/benchmark/utils/valgrind_wrapper/timer_interface.py
index c3ee8894f3ebd..02738aa26aa10 100644
--- a/torch/utils/benchmark/utils/valgrind_wrapper/timer_interface.py
+++ b/torch/utils/benchmark/utils/valgrind_wrapper/timer_interface.py
@@ -1,4 +1,6 @@
 """Intermediate layer between `Timer` and `valgrind`."""
+from __future__ import annotations
+
 import collections
 import enum
 import dataclasses
@@ -12,15 +14,17 @@
 import textwrap
 from typing import (
     cast, Any, NamedTuple,
-    Union, TYPE_CHECKING)
-from collections.abc import Callable
-from collections.abc import Iterator
+    TYPE_CHECKING)
 
 import torch
 from torch.utils.benchmark.utils import common, cpp_jit
-from torch.utils.benchmark.utils._stubs import CallgrindModuleType
 import operator
 
+if TYPE_CHECKING:
+    from collections.abc import Callable, Iterator
+
+    from torch.utils.benchmark.utils._stubs import CallgrindModuleType
+
 
 __all__ = ["FunctionCount", "FunctionCounts", "CallgrindStats", "CopyIfCallgrind"]
 
@@ -63,10 +67,9 @@ def __iter__(self) -> Iterator[FunctionCount]:
     def __len__(self) -> int:
         return len(self._data)
 
-    def __getitem__(self, item: Any) -> Union[FunctionCount, "FunctionCounts"]:
+    def __getitem__(self, item: Any) -> FunctionCount | FunctionCounts:
         data: FunctionCount | tuple[FunctionCount, ...] = self._data[item]
         return (
-            # pyrefly: ignore [bad-return]
             FunctionCounts(cast(tuple[FunctionCount, ...], data), self.inclusive, truncate_rows=False)
             if isinstance(data, tuple) else data
         )
@@ -96,22 +99,22 @@ def __repr__(self) -> str:
 
     def __add__(
         self,
-        other: "FunctionCounts",
-    ) -> "FunctionCounts":
+        other: FunctionCounts,
+    ) -> FunctionCounts:
         return self._merge(other, lambda c: c)
 
     def __sub__(
         self,
-        other: "FunctionCounts",
-    ) -> "FunctionCounts":
+        other: FunctionCounts,
+    ) -> FunctionCounts:
         return self._merge(other, operator.neg)
 
-    def __mul__(self, other: int | float) -> "FunctionCounts":
+    def __mul__(self, other: int | float) -> FunctionCounts:
         return self._from_dict({
             fn: int(c * other) for c, fn in self._data
         }, self.inclusive)
 
-    def transform(self, map_fn: Callable[[str], str]) -> "FunctionCounts":
+    def transform(self, map_fn: Callable[[str], str]) -> FunctionCounts:
         """Apply `map_fn` to all of the function names.
 
         This can be used to regularize function names (e.g. stripping irrelevant
@@ -124,14 +127,14 @@ def transform(self, map_fn: Callable[[str], str]) -> "FunctionCounts":
 
         return self._from_dict(counts, self.inclusive)
 
-    def filter(self, filter_fn: Callable[[str], bool]) -> "FunctionCounts":
+    def filter(self, filter_fn: Callable[[str], bool]) -> FunctionCounts:
         """Keep only the elements where `filter_fn` applied to function name returns True."""
         return FunctionCounts(tuple(i for i in self if filter_fn(i.function)), self.inclusive)
 
     def sum(self) -> int:
         return sum(c for c, _ in self)
 
-    def denoise(self) -> "FunctionCounts":
+    def denoise(self) -> FunctionCounts:
         """Remove known noisy instructions.
 
         Several instructions in the CPython interpreter are rather noisy. These
@@ -143,9 +146,9 @@ def denoise(self) -> "FunctionCounts":
 
     def _merge(
         self,
-        second: "FunctionCounts",
+        second: FunctionCounts,
         merge_fn: Callable[[int], int]
-    ) -> "FunctionCounts":
+    ) -> FunctionCounts:
         if self.inclusive != second.inclusive:
             raise AssertionError("Cannot merge inclusive and exclusive counts.")
         counts: collections.defaultdict[str, int] = collections.defaultdict(int)
@@ -158,7 +161,7 @@ def _merge(
         return self._from_dict(counts, self.inclusive)
 
     @staticmethod
-    def _from_dict(counts: dict[str, int], inclusive: bool) -> "FunctionCounts":
+    def _from_dict(counts: dict[str, int], inclusive: bool) -> FunctionCounts:
         flat_counts = (FunctionCount(c, fn) for fn, c in counts.items() if c)
         return FunctionCounts(tuple(sorted(flat_counts, reverse=True)), inclusive)
 
@@ -223,7 +226,7 @@ def counts(self, *, denoise: bool = False) -> int:
     # FIXME: Once 3.7 is the minimum version, type annotate `other` per PEP 563
     def delta(
         self,
-        other: "CallgrindStats",
+        other: CallgrindStats,
         inclusive: bool = False,
     ) -> FunctionCounts:
         """Diff two sets of counts.
@@ -238,7 +241,7 @@ def delta(
         """
         return self.stats(inclusive=inclusive) - other.stats(inclusive=inclusive)
 
-    def as_standardized(self) -> "CallgrindStats":
+    def as_standardized(self) -> CallgrindStats:
         """Strip library names and some prefixes from function strings.
 
         When comparing two different sets of instruction counts, on stumbling
@@ -450,13 +453,11 @@ def construct(self) -> str:
         load_lines = []
         for name, wrapped_value in self._globals.items():
             if wrapped_value.setup is not None:
-                # pyrefly: ignore [bad-argument-type]
                 load_lines.append(textwrap.dedent(wrapped_value.setup))
 
             if wrapped_value.serialization == Serialization.PICKLE:
                 path = os.path.join(self._data_dir, f"{name}.pkl")
                 load_lines.append(
-                    # pyrefly: ignore [bad-argument-type]
                     f"with open({repr(path)}, 'rb') as f:\n    {name} = pickle.load(f)")
                 with open(path, "wb") as f:
                     pickle.dump(wrapped_value.value, f)
@@ -466,13 +467,11 @@ def construct(self) -> str:
                 # TODO: Figure out if we can use torch.serialization.add_safe_globals here
                 # Using weights_only=False after the change in
                 # https://dev-discuss.pytorch.org/t/bc-breaking-change-torch-load-is-being-flipped-to-use-weights-only-true-by-default-in-the-nightlies-after-137602/2573
-                # pyrefly: ignore [bad-argument-type]
                 load_lines.append(f"{name} = torch.load({repr(path)}, weights_only=False)")
                 torch.save(wrapped_value.value, path)
 
             elif wrapped_value.serialization == Serialization.TORCH_JIT:
                 path = os.path.join(self._data_dir, f"{name}.pt")
-                # pyrefly: ignore [bad-argument-type]
                 load_lines.append(f"{name} = torch.jit.load({repr(path)})")
                 with open(path, "wb") as f:
                     torch.jit.save(wrapped_value.value, f)  # type: ignore[no-untyped-call]
diff --git a/torch/utils/checkpoint.py b/torch/utils/checkpoint.py
index 8e99c8a45dafd..d5f2bb039262b 100644
--- a/torch/utils/checkpoint.py
+++ b/torch/utils/checkpoint.py
@@ -1,11 +1,15 @@
 # mypy: allow-untyped-defs
+from __future__ import annotations
+
 import contextlib
+import itertools
 import platform
 import uuid
 import warnings
 import weakref
 from collections import defaultdict
 from typing import *  # noqa: F403
+from typing_extensions import Self
 import enum
 from weakref import ReferenceType
 
@@ -14,6 +18,7 @@
 from torch.utils._pytree import tree_map
 from torch.testing._internal.logging_tensor import capture_logs, LoggingTensorMode
 from torch.utils._python_dispatch import TorchDispatchMode
+from torch._C._autograd import _make_saved_tensor, SavedTensor
 from typing import NoReturn
 
 __all__ = [
@@ -38,11 +43,11 @@
 
 _DEFAULT_DETERMINISM_MODE = "default"
 
-_checkpoint_debug_enabled: Optional[bool] = None
+_checkpoint_debug_enabled: bool | None = None
 
 
 @contextlib.contextmanager
-def set_checkpoint_debug_enabled(enabled: Optional[bool]):
+def set_checkpoint_debug_enabled(enabled: bool | None):
     """
     Context manager that sets whether checkpoint should print additional debug
     information when running. See the ``debug`` flag for
@@ -77,8 +82,7 @@ def detach_variable(inputs: Tuple[Any, ...]) -> Tuple[torch.Tensor, ...]:
         return tuple(out)
     else:
         raise RuntimeError(
-            "Only tuple of tensors is supported. Got Unsupported input type: ",
-            type(inputs).__name__,
+            f"Only tuple of tensors is supported. Got Unsupported input type: {type(inputs).__name__}"
         )
 
 
@@ -106,7 +110,7 @@ class DefaultDeviceType:
     to save and restore for recomputation.
     """
 
-    _default_device_type: Optional[str] = None
+    _default_device_type: str | None = None
 
     @staticmethod
     def set_device_type(device: str = "cuda") -> None:
@@ -135,10 +139,11 @@ def get_device_type() -> str:
 def _infer_device_type(*args):
     device_types = []
 
-    def add_device_types(arg) -> None:
+    def add_device_types(arg):
         nonlocal device_types
         if isinstance(arg, torch.Tensor) and arg.device.type != "cpu":
             device_types.append(arg.device.type)
+        return arg
     tree_map(add_device_types, args)
 
     device_types_set = set(device_types)
@@ -171,10 +176,11 @@ def get_device_states(*args) -> Tuple[List[int], List[torch.Tensor]]:
     # the conditionals short-circuit.
     fwd_device_ids = []
 
-    def add_device_ids(arg) -> None:
+    def add_device_ids(arg):
         nonlocal fwd_device_ids
         if isinstance(arg, torch.Tensor) and arg.device.type not in {"cpu", "meta"}:
             fwd_device_ids.append(arg.get_device())
+        return arg
     tree_map(add_device_ids, args)
 
     fwd_device_states = []
@@ -349,7 +355,7 @@ def noop_context_fn():
 def checkpoint(
     function,
     *args,
-    use_reentrant: Optional[bool] = None,
+    use_reentrant: bool | None = None,
     context_fn: Callable[[], Tuple[ContextManager, ContextManager]] = noop_context_fn,
     determinism_check: str = _DEFAULT_DETERMINISM_MODE,
     debug: bool = False,
@@ -748,7 +754,7 @@ def _internal_assert(cond) -> None:
 #    by holder=None. We skip over them. We still save x at (4) (since its holder
 #    is still alive.)
 
-_enable_checkpoint_early_stop: Optional[bool] = None
+_enable_checkpoint_early_stop: bool | None = None
 
 
 @contextlib.contextmanager
@@ -788,51 +794,13 @@ class _Handle:
 
 class _Holder:
     def __init__(self) -> None:
-        self.handles: Dict[int, Optional[_Handle]] = {}
-
-
-class _NoopSaveInputs(torch.autograd.Function):
-    @staticmethod
-    # pyrefly: ignore [bad-override]
-    def forward(*args):
-        return torch.empty((0,))
-
-    @staticmethod
-    def setup_context(ctx: Any, inputs: Tuple[Any, ...], output: Any) -> None:
-        # Only tensors can be saved with ctx.save_for_backward, everything else
-        # is captured by get_args, which is saved directly on ctx
-        tensor_indices, tensors = zip(
-            *[(i, o) for i, o in enumerate(inputs) if isinstance(o, torch.Tensor)], strict=False
-        )
-        idx2saved_idx = {b: a for a, b in enumerate(tensor_indices)}
-        # args but with tensors replaced with None as placeholders
-        args = [None if isinstance(o, torch.Tensor) else o for o in inputs]
-
-        def get_args(saved_tensors):
-            # restore the placeholders with the original tensors grabbed from
-            # ctx.saved_tensors (which may be saved on a parent checkpoint if
-            # this checkpoint is nested, and that would trigger a recursive
-            # unpack!)
-            ret = [
-                saved_tensors[idx2saved_idx[i]] if i in tensor_indices else o
-                for i, o in enumerate(args)
-            ]
-            # grab the tail since we also saved the dummy to avoid having to explicitly
-            # handle the case where there are no tensor inputs
-            return ret[1:]
-
-        ctx.get_args = get_args
-        ctx.save_for_backward(*tensors)
-
-    @staticmethod
-    def backward(ctx, *grad_outputs) -> NoReturn:
-        raise AssertionError("Did not expect to backward on this graph")
+        self.handles: dict[int, _Handle | None] = {}
 
 
 class _CheckpointFrame:
     def __init__(self, recompute_fn, early_stop, unpack_error_cb, metadata_fn) -> None:
         self.recompute_fn = recompute_fn
-        self.input_saver = None
+        self.saved_args: List[Any] = []
         self.weak_holders: List[ReferenceType] = []
         # We store this as a weakkeydictionary so that in the case of a partial
         # backward, the entries in the dict are cleared alongside the Holder
@@ -855,6 +823,19 @@ def __init__(self, recompute_fn, early_stop, unpack_error_cb, metadata_fn) -> No
         self.forward_completed = False
         self.ignore_saved_mismatch = False
 
+    def save_inputs(self, *args):
+        self.saved_args = [
+            _make_saved_tensor(arg, is_output=False)
+            if isinstance(arg, torch.Tensor) else arg
+            for arg in args
+        ]
+
+    def get_inputs(self):
+        return [
+            arg.unpack() if isinstance(arg, SavedTensor) else arg
+            for arg in self.saved_args
+        ]
+
     def check_recomputed_tensors_match(self, gid) -> None:
         if self.ignore_saved_mismatch:
             # TODO: we can probably make this check stricter by checking that
@@ -931,6 +912,18 @@ def check_recomputed_tensors_match(self, gid) -> None:
 `torch.utils.checkpoint.checkpoint(...)` or wrap the code block
 with `with torch.utils.checkpoint.set_checkpoint_debug_enabled(True):` to
 enable checkpoint‑debug mode globally.
+
+If this error occurs under torch.compile with automatic_dynamic_shapes enabled,
+it may be because the recomputation selected a different compiled graph than the
+forward pass (e.g., a dynamic graph instead of the original static graph).
+To fix this, either:
+  - Use torch._dynamo.mark_dynamic() to explicitly mark varying dimensions as
+    dynamic upfront, avoiding the static-to-dynamic transition.
+  - Call torch._C._dynamo.eval_frame._set_lru_cache(False) to disable LRU cache
+    reordering, which can change which graph is checked first between forward
+    and recompute.
+See https://github.com/pytorch/pytorch/issues/166926 for more details and
+workaround examples.
 """
 
 
@@ -1003,8 +996,7 @@ def _get_debug_context_and_cb() -> Tuple[Callable[[], Any], Callable[[Checkpoint
     # checkpointing mechanism. error_cb is invoked when an error is detected
     # during unpack.
 
-    # record_context_cpp is not support on non-linux non-x86_64 platforms
-    cpp_tb = platform.machine() == 'x86_64' and platform.system() == 'Linux'
+    cpp_tb = platform.machine() in ('x86_64', 'aarch64', 'arm64') and platform.system() == 'Linux'
 
     class CaptureLogs:
         def __init__(self) -> None:
@@ -1016,7 +1008,6 @@ def get_context_manager(self):
             def logging_mode():
                 with LoggingTensorMode(), \
                      capture_logs(True, python_tb=True, script_tb=True, cpp_tb=cpp_tb) as logs_and_tb:
-                    # pyrefly: ignore [bad-assignment]
                     self.logs, self.tbs = logs_and_tb
                     yield logs_and_tb
             return logging_mode()
@@ -1079,7 +1070,11 @@ class _StopRecomputationError(Exception):
 
 
 class _recomputation_hook(torch.autograd.graph.saved_tensors_hooks):
-    def __init__(self, target_frame_ref: ReferenceType, gid: Union["GraphExecGroup", int]) -> None:
+    def __init__(self, target_frame_ref: ReferenceType, gid: GraphExecGroup | int) -> None:
+        # Dynamo guards on WeakKeyDictionary internals are unstable here
+        # (dict length/keys change every call), causing recompilation storms.
+        # with `.compile()` so we disable
+        @torch._dynamo.disable
         def pack_hook(x):
             x = x.detach() if x.requires_grad else x
             target_frame = target_frame_ref()
@@ -1153,7 +1148,7 @@ def pack_hook(x):
 
         def unpack_hook(holder):
             # First check if we're inside a GraphExecGroup context
-            gid: Union[GraphExecGroup, None, int] = GraphExecGroup._get_current_group()
+            gid: GraphExecGroup | None | int = GraphExecGroup._get_current_group()
             if gid is None:
                 # Fallback to using the current graph task id
                 gid = torch._C._current_graph_task_id()
@@ -1162,8 +1157,7 @@ def unpack_hook(holder):
                     gid = int(uuid.uuid4())
 
             if not frame.is_recomputed[gid]:
-                ctx = frame.input_saver.grad_fn
-                args = ctx.get_args(ctx.saved_tensors)
+                args = frame.get_inputs()
 
                 try:
                     with _recomputation_hook(
@@ -1208,16 +1202,18 @@ def unpack_hook_with_error_cb(holder):
 
 
 def _is_compiling(func, args, kwargs):
-    # Check if we are under AOTAutograd tracing
-    # Checking that a functional mode is active should always do what we want
+    # Check if we are under AOTAutograd tracing or export tracing
+    # Checking that a proxy mode is active should always do what we want
+    if torch.compiler._is_non_strict_tracing():
+        return False
     return torch._C._get_dispatch_mode(torch._C._TorchDispatchModeKey.PROXY) is not None
 
 
 class _VersionWrapper:
     # Check that cached tensors are not mutated.
     def __init__(self, val) -> None:
-        self.val: Union[torch.Tensor, Any] = val
-        self.version: Optional[int] = val._version if isinstance(val, torch.Tensor) else None
+        self.val: torch.Tensor | Any = val
+        self.version: int | None = val._version if isinstance(val, torch.Tensor) else None
 
     def get_val(self, allow_cache_entry_mutation):
         if self.version is not None and not allow_cache_entry_mutation:
@@ -1270,8 +1266,9 @@ class SelectiveCheckpointContext:
         >>>     context_fn=context_fn,
         >>> )
     """
-    def __init__(self, *, is_recompute) -> None:
+    def __init__(self, *, is_recompute, op_output=None) -> None:
         self.is_recompute = is_recompute
+        self.op_output = op_output
 
 
 class CheckpointPolicy(enum.Enum):
@@ -1328,25 +1325,29 @@ def ignore_compile_internals(cls):
         return True
 
     # Used together with _CachedTorchDispatchMode to implement SAC.
-    def __init__(self, policy_fn, storage) -> None:
+    def __init__(self, policy_fn, storage, ac_graph_id=None) -> None:
         self.policy_fn = policy_fn
         self.storage = storage
+        self.ac_graph_id = ac_graph_id
 
     def __torch_dispatch__(self, func, types, args=(), kwargs=None):
-        if func in SAC_IGNORED_OPS:
-            return func(*args, **kwargs)
-
         kwargs = {} if kwargs is None else kwargs
-        policy = self.policy_fn(SelectiveCheckpointContext(is_recompute=False),
-                                func, *args, **kwargs)
-        if isinstance(policy, bool):
-            policy = _policy_from_bool(policy)
-
         is_compiling = _is_compiling(func, args, kwargs)
 
         if is_compiling:
-            # Overwrite each node's "recompute" tag to add in the user annotation.
-            fx_traceback.current_meta["recompute"] = policy
+            fx_traceback.current_meta["ac_graph_id"] = self.ac_graph_id
+            fx_traceback.current_meta["recompute"] = CheckpointPolicy.PREFER_RECOMPUTE
+
+        if func in SAC_IGNORED_OPS:
+            return func(*args, **kwargs)
+
+        proxy_mode = None
+        graph_len_before = 0
+        if is_compiling:
+            from torch.fx.experimental.proxy_tensor import get_proxy_mode
+            proxy_mode = get_proxy_mode()
+            if proxy_mode is not None:
+                graph_len_before = len(proxy_mode.tracer.graph.nodes)
 
         out = func(*args, **kwargs)
 
@@ -1358,6 +1359,18 @@ def __torch_dispatch__(self, func, types, args=(), kwargs=None):
         else:
             any_ret_has_alias_info = any(ret.alias_info is not None for ret in func._schema.returns)
 
+        policy = self.policy_fn(SelectiveCheckpointContext(is_recompute=False, op_output=out),
+                                func, *args, **kwargs)
+        if isinstance(policy, bool):
+            policy = _policy_from_bool(policy)
+
+        if is_compiling:
+            if proxy_mode is not None:
+                graph = proxy_mode.tracer.graph
+                num_new = len(graph.nodes) - graph_len_before
+                for node in itertools.islice(reversed(graph.nodes), num_new):
+                    node.meta["recompute"] = policy
+
         if policy in (CheckpointPolicy.MUST_SAVE, CheckpointPolicy.PREFER_SAVE) or is_compiling:
             self.storage[func].append(tree_map(lambda x: _VersionWrapper(_maybe_detach(x, any_ret_has_alias_info)), out))
         return out
@@ -1584,9 +1597,10 @@ def _checkpoint_without_reentrant_generator(
         ),
         contextlib.nullcontext(),
     )
+    error_on_nested_fx_trace = torch._dynamo.config.error_on_nested_fx_trace
+    is_non_strict_tracing = torch.compiler._is_non_strict_tracing()
 
-    def recompute_fn(*inputs) -> None:
-        kwargs, *args = inputs
+    def recompute_fn(*args) -> None:
         # This will be called later during recomputation. This wrapping enables
         # the necessary global state to be captured.
         rng_devices = []
@@ -1603,7 +1617,20 @@ def recompute_fn(*inputs) -> None:
             device_autocast_ctx = torch.amp.autocast(
                 device_type=device_type, **device_autocast_kwargs
             ) if torch.amp.is_autocast_available(device_type) else contextlib.nullcontext()
-            with device_autocast_ctx, torch.amp.autocast("cpu", **cpu_autocast_kwargs), recompute_context, device_ctx:  # type: ignore[attr-defined]
+            nested_fx_trace_ctx = (
+                torch._dynamo.config.patch(
+                    error_on_nested_fx_trace=error_on_nested_fx_trace
+                )
+                if is_non_strict_tracing
+                else contextlib.nullcontext()
+            )
+            with (
+                device_autocast_ctx,
+                torch.amp.autocast("cpu", **cpu_autocast_kwargs),
+                recompute_context,
+                device_ctx,
+                nested_fx_trace_ctx,
+            ):  # type: ignore[attr-defined]
                 fn(*args, **kwargs)
 
     new_frame = _CheckpointFrame(
@@ -1612,14 +1639,13 @@ def recompute_fn(*inputs) -> None:
         unpack_error_cb,
         metadata_fn
     )
-    dummy = torch.empty((0,), requires_grad=True)
-    new_frame.input_saver = _NoopSaveInputs.apply(dummy, kwargs, *args)
 
-    # When ambient grad_mode is False
-    if new_frame.input_saver.grad_fn is None:
+    if not torch.is_grad_enabled():
         yield
         return
 
+    new_frame.save_inputs(*args)
+
     with _checkpoint_hook(new_frame), forward_context:
         yield
     new_frame.forward_completed = True
@@ -1652,7 +1678,7 @@ class GraphExecGroup:
         is a no-op otherwise.
     """
 
-    def __enter__(self) -> "GraphExecGroup":
+    def __enter__(self) -> Self:
         if torch._C._get_graph_exec_group() is not None:
             raise RuntimeError(
                 "GraphExecGroup contexts cannot be nested. "
@@ -1665,7 +1691,7 @@ def __exit__(self, *args: object) -> None:
         torch._C._set_graph_exec_group(None)
 
     @classmethod
-    def _get_current_group(cls) -> Optional["GraphExecGroup"]:
+    def _get_current_group(cls) -> GraphExecGroup | None:
         # Private API to be used by utils like AC
         return torch._C._get_graph_exec_group()
 
diff --git a/torch/utils/collect_env.py b/torch/utils/collect_env.py
index c5932dd9369a6..b566cd94e5d03 100644
--- a/torch/utils/collect_env.py
+++ b/torch/utils/collect_env.py
@@ -319,7 +319,6 @@ def get_linux_pkg_version(run_lambda, pkg_name):
         },
     }
 
-    # pyrefly: ignore [redundant-cast]
     field_index: int = int(_cast(int, grep_version[pkg_mgr]["field_index"]))
     cmd: str = str(grep_version[pkg_mgr]["command"])
     cmd = cmd.format(pkg_name)
diff --git a/torch/utils/cpp_extension.py b/torch/utils/cpp_extension.py
index a63bff50d5ec3..3e9d5d3823cbe 100644
--- a/torch/utils/cpp_extension.py
+++ b/torch/utils/cpp_extension.py
@@ -71,8 +71,8 @@
     '12.5': ((6, 0, 0), (14, 0)),
     '12.6': ((6, 0, 0), (14, 0)),
     '12.7': ((6, 0, 0), (14, 0)),
-    '12.8': ((6, 0, 0), (14, 0)),
-    '12.9': ((6, 0, 0), (14, 0)),
+    '12.8': ((6, 0, 0), (15, 0)),
+    '12.9': ((6, 0, 0), (15, 0)),
     '13.0': ((6, 0, 0), (16, 0)),
 }
 
@@ -635,7 +635,7 @@ class BuildExtension(build_ext):
     A custom :mod:`setuptools` build extension .
 
     This :class:`setuptools.build_ext` subclass takes care of passing the
-    minimum required compiler flags (e.g. ``-std=c++17``) as well as mixed
+    minimum required compiler flags (e.g. ``-std=c++20``) as well as mixed
     C++/CUDA/SYCL compilation (and support for CUDA/SYCL files in general).
 
     When using :class:`BuildExtension`, it is allowed to supply a dictionary
@@ -767,7 +767,7 @@ def append_std17_if_no_std_present(cflags) -> None:
             # overriding the option if the user explicitly passed it.
             cpp_format_prefix = '/{}:' if self.compiler.compiler_type == 'msvc' else '-{}='
             cpp_flag_prefix = cpp_format_prefix.format('std')
-            cpp_flag = cpp_flag_prefix + 'c++17'
+            cpp_flag = cpp_flag_prefix + 'c++20'
             if not any(flag.startswith(cpp_flag_prefix) for flag in cflags):
                 cflags.append(cpp_flag)
 
@@ -947,6 +947,20 @@ def win_cuda_flags(cflags):
         def win_hip_flags(cflags):
             return (COMMON_HIPCC_FLAGS + COMMON_HIP_FLAGS + cflags + _get_rocm_arch_flags(cflags))
 
+        def win_filter_msvc_include_dirs(pp_opts) -> list[str]:
+            """Filter out MSVC include dirs from pp_opts for oneAPI 2025.3+."""
+            # oneAPI 2025.3+ changed include path ordering to match MSVC behavior.
+            # Filter out MSVC headers to avoid conflicting declarations with oneAPI's std headers.
+            icpx_version = int(_get_icpx_version())
+            if icpx_version >= 20250300:
+                vc_tools_dir = os.path.normcase(os.environ.get('VCToolsInstallDir', ''))
+                if vc_tools_dir:
+                    pp_opts = [
+                        path for path in pp_opts
+                        if vc_tools_dir not in os.path.normcase(path)
+                    ]
+            return pp_opts
+
         def win_wrap_single_compile(sources,
                                     output_dir=None,
                                     macros=None,
@@ -997,7 +1011,7 @@ def spawn(cmd):
                         if IS_HIP_EXTENSION:
                             cflags = win_hip_flags(cflags)
                         else:
-                            cflags = win_cuda_flags(cflags) + ['-std=c++17', '--use-local-env']
+                            cflags = win_cuda_flags(cflags) + ['-std=c++20', '--use-local-env']
                             for ignore_warning in MSVC_IGNORE_CUDAFE_WARNINGS:
                                 cflags = ['-Xcudafe', '--diag_suppress=' + ignore_warning] + cflags
                         for flag in COMMON_MSVC_FLAGS:
@@ -1084,7 +1098,7 @@ def win_wrap_ninja_compile(sources,
             cuda_post_cflags = None
             cuda_cflags = None
             if with_cuda:
-                cuda_cflags = ['-std=c++17']
+                cuda_cflags = ['-std=c++20']
                 for common_cflag in common_cflags:
                     cuda_cflags.append('-Xcompiler')
                     cuda_cflags.append(common_cflag)
@@ -1116,7 +1130,7 @@ def win_wrap_ninja_compile(sources,
             sycl_post_cflags = None
             sycl_dlink_post_cflags = None
             if with_sycl:
-                sycl_cflags = common_cflags + pp_opts + _COMMON_SYCL_FLAGS
+                sycl_cflags = common_cflags + win_filter_msvc_include_dirs(pp_opts) + _COMMON_SYCL_FLAGS
                 if isinstance(extra_postargs, dict):
                     sycl_post_cflags = extra_postargs['sycl']
                 else:
@@ -1946,7 +1960,7 @@ def build_precompile_header(pch_cmd) -> None:
     if not is_standalone:
         common_cflags += ['-DTORCH_API_INCLUDE_EXTENSION_H']
 
-    common_cflags += ['-std=c++17', '-fPIC']
+    common_cflags += ['-std=c++20', '-fPIC']
     common_cflags_str = listToString(common_cflags)
 
     pch_cmd = format_precompiler_header_cmd(compiler, head_file, head_file_pch, common_cflags_str, torch_include_dirs_str, extra_cflags_str, extra_include_paths_str)
@@ -2247,8 +2261,11 @@ def _jit_compile(name,
                         hipified_sources = set()
                         for source in sources:
                             s_abs = os.path.abspath(source)
-                            hipified_sources.add(hipify_result[s_abs].hipified_path if s_abs in hipify_result else s_abs)
-
+                            if s_abs in hipify_result and hipify_result[s_abs].hipified_path is not None:
+                                hipified_s_abs = hipify_result[s_abs].hipified_path
+                            else:
+                                hipified_s_abs = s_abs
+                            hipified_sources.add(hipified_s_abs)
                         sources = list(hipified_sources)
 
                     _write_ninja_file_and_build_library(
@@ -2636,19 +2653,24 @@ def _get_rocm_arch_flags(cflags: list[str] | None = None) -> list[str]:
     # Allow env var to override, just like during initial cmake build.
     _archs = os.environ.get('PYTORCH_ROCM_ARCH', None)
     if not _archs:
-        archFlags = torch._C._cuda_getArchFlags()
-        if archFlags:
-            archs = archFlags.split()
-        else:
-            archs = []
-            logger.warning(
-                "Failed to auto-detect ROCm architecture. Extensions will be compiled "
-                "without architecture-specific optimizations. Set PYTORCH_ROCM_ARCH "
-                "environment variable to specify target architectures "
-                "(e.g., export PYTORCH_ROCM_ARCH='gfx90a;gfx942')."
-            )
+        arch_set = set()
+        # the assumption is that the extension should run on any of the currently visible cards,
+        # which could be of different types - therefore all archs for visible cards should be included
+        for i in range(torch.cuda.device_count()):
+            device_properties = torch.cuda.get_device_properties(i)
+            if hasattr(device_properties, "gcnArchName"):
+                device_arch = (device_properties.gcnArchName).split(":", 1)[0]
+                arch_set.add(device_arch)
+
+        archs = ";".join(arch_set)
+
+        logger.warning(
+            "The environment variable `PYTORCH_ROCM_ARCH` is not set, all archs for visible cards are included for compilation (%s).\n"
+            "If this is not desired, please set the environment variable `PYTORCH_ROCM_ARCH` to specific architectures.", archs)
     else:
-        archs = _archs.replace(' ', ';').split(';')
+        archs = _archs.replace(' ', ';')
+
+    archs = archs.split(';')
     flags = [f'--offload-arch={arch}' for arch in archs]
     flags += [] if has_gpu_rdc_flag else ['-fno-gpu-rdc']
     return flags
@@ -2855,15 +2877,15 @@ def _write_ninja_file_to_build_library(path,
 
     if IS_WINDOWS:
         COMMON_HIP_FLAGS.extend(['-fms-runtime-lib=dll'])
-        cflags = common_cflags + ['/std:c++17'] + extra_cflags
+        cflags = common_cflags + ['/std:c++20'] + extra_cflags
         cflags += COMMON_MSVC_FLAGS + (COMMON_HIP_FLAGS if IS_HIP_EXTENSION else [])
         cflags = _nt_quote_args(cflags)
     else:
-        cflags = common_cflags + ['-fPIC', '-std=c++17'] + extra_cflags
+        cflags = common_cflags + ['-fPIC', '-std=c++20'] + extra_cflags
 
     if with_cuda and IS_HIP_EXTENSION:
         cuda_flags = ['-DWITH_HIP'] + common_cflags + extra_cflags + COMMON_HIP_FLAGS + COMMON_HIPCC_FLAGS
-        cuda_flags = cuda_flags + ['-std=c++17']
+        cuda_flags = cuda_flags + ['-std=c++20']
         cuda_flags += _get_rocm_arch_flags(cuda_flags)
         cuda_flags += extra_cuda_cflags
         if IS_WINDOWS:
@@ -2875,14 +2897,14 @@ def _write_ninja_file_to_build_library(path,
                 cuda_flags = ['-Xcompiler', flag] + cuda_flags
             for ignore_warning in MSVC_IGNORE_CUDAFE_WARNINGS:
                 cuda_flags = ['-Xcudafe', '--diag_suppress=' + ignore_warning] + cuda_flags
-            cuda_flags = cuda_flags + ['-std=c++17']
+            cuda_flags = cuda_flags + ['-std=c++20']
             cuda_flags = _nt_quote_args(cuda_flags)
             cuda_flags += _nt_quote_args(extra_cuda_cflags)
         else:
             cuda_flags += ['--compiler-options', "'-fPIC'"]
             cuda_flags += extra_cuda_cflags
             if not any(flag.startswith('-std=') for flag in cuda_flags):
-                cuda_flags.append('-std=c++17')
+                cuda_flags.append('-std=c++20')
             cc_env = os.getenv("CC")
             if cc_env is not None:
                 cuda_flags = ['-ccbin', cc_env] + cuda_flags
diff --git a/torch/utils/data/_utils/worker.py b/torch/utils/data/_utils/worker.py
index 611aee4766bf4..850365ee9b19b 100644
--- a/torch/utils/data/_utils/worker.py
+++ b/torch/utils/data/_utils/worker.py
@@ -5,11 +5,13 @@
 static methods.
 """
 
+from __future__ import annotations
+
 import os
 import queue
 import random
 from dataclasses import dataclass
-from typing import Optional, TYPE_CHECKING
+from typing import TYPE_CHECKING
 
 import torch
 from torch._utils import ExceptionWrapper
@@ -70,32 +72,28 @@ def is_alive(self) -> bool:
             return not self.manager_dead
 
 
-_worker_info: Optional["WorkerInfo"] = None
+_worker_info: WorkerInfo | None = None
 
 
+@dataclass(frozen=True, slots=True)
 class WorkerInfo:
+    """Information about the current DataLoader worker process or thread.
+
+    Attributes:
+        id: The current worker id (0 to num_workers - 1)
+        num_workers: Total number of workers
+        seed: Random seed set for this worker
+        dataset: Copy of the dataset object in this worker
+        rng: Optional RNG state container. Defaults to None.
+        worker_method: Optional worker method ("multiprocessing" or "thread"). Defaults to "multiprocessing".
+    """
+
     id: int
     num_workers: int
     seed: int
-    dataset: "Dataset"
-    __initialized = False
-
-    def __init__(self, **kwargs) -> None:
-        for k, v in kwargs.items():
-            setattr(self, k, v)
-        self.__keys = tuple(kwargs.keys())
-        self.__initialized = True
-
-    def __setattr__(self, key, val) -> None:
-        if self.__initialized:
-            raise RuntimeError(
-                f"Cannot assign attributes to {self.__class__.__name__} objects"
-            )
-        return super().__setattr__(key, val)
-
-    def __repr__(self) -> str:
-        items = [f"{k}={getattr(self, k)}" for k in self.__keys]
-        return f"{self.__class__.__name__}({', '.join(items)})"
+    dataset: Dataset
+    rng: _RNG | None = None
+    worker_method: str | None = "multiprocessing"
 
 
 def get_worker_info() -> WorkerInfo | None:
@@ -143,6 +141,24 @@ class _ResumeIteration:
     seed: int | None = None
 
 
+@dataclass(frozen=True, slots=True)
+class _RNG:
+    """Container for thread-local random number generator state.
+
+    Used by thread workers to maintain separate RNG state per worker thread
+    to avoid race conditions.
+
+    Attributes:
+        random_generator: Python random.Random generator for this thread
+        torch_generator: PyTorch Generator for this thread
+        numpy_generator: NumPy Generator for this thread (None if numpy not available)
+    """
+
+    random_generator: random.Random
+    torch_generator: torch.Generator
+    numpy_generator: object | None = None
+
+
 # The function `_generate_state` is adapted from `numpy.random.SeedSequence`
 # from https://github.com/numpy/numpy/blob/main/numpy/random/bit_generator.pyx
 # It's MIT licensed, here is the copyright:
diff --git a/torch/utils/data/dataloader.py b/torch/utils/data/dataloader.py
index 0b51561ca6e7f..4eaa755980787 100644
--- a/torch/utils/data/dataloader.py
+++ b/torch/utils/data/dataloader.py
@@ -423,8 +423,6 @@ def __init__(
 
         self.check_worker_number_rationality()
 
-        torch.set_vital("Dataloader", "enabled", "True")  # type: ignore[attr-defined]
-
     def _get_iterator(self) -> _BaseDataLoaderIter:
         if self.num_workers == 0:
             return _SingleProcessDataLoaderIter(self)
@@ -606,31 +604,10 @@ def _create_warning_msg(num_worker_suggest, num_worker_created, cpuset_checked):
             return
 
         # try to compute a suggested max number of worker based on system's resource
-        max_num_worker_suggest = None
-        cpuset_checked = False
-        if hasattr(os, "sched_getaffinity"):
-            try:
-                max_num_worker_suggest = len(os.sched_getaffinity(0))
-                cpuset_checked = True
-            except Exception:
-                pass
-        if max_num_worker_suggest is None:
-            # os.cpu_count() could return Optional[int]
-            # get cpu count first and check None in order to satisfy mypy check
-            cpu_count = os.cpu_count()
-            if cpu_count is not None:
-                max_num_worker_suggest = cpu_count
-
-        if max_num_worker_suggest is None:
-            warnings.warn(
-                _create_warning_msg(
-                    max_num_worker_suggest, self.num_workers, cpuset_checked
-                ),
-                stacklevel=2,
-            )
-            return
+        max_num_worker_suggest = torch._utils.cpu_count()
+        cpuset_checked = hasattr(os, "sched_getaffinity")
 
-        if self.num_workers > max_num_worker_suggest:
+        if max_num_worker_suggest is None or self.num_workers > max_num_worker_suggest:
             warnings.warn(
                 _create_warning_msg(
                     max_num_worker_suggest, self.num_workers, cpuset_checked
diff --git a/torch/utils/data/datapipes/_typing.py b/torch/utils/data/datapipes/_typing.py
index c72eab3fdb364..d44f2f6d67c83 100644
--- a/torch/utils/data/datapipes/_typing.py
+++ b/torch/utils/data/datapipes/_typing.py
@@ -284,7 +284,6 @@ def __new__(cls, name, bases, namespace, **kwargs):
         return super().__new__(cls, name, bases, namespace, **kwargs)  # type: ignore[call-overload]
 
         # TODO: the statements below are not reachable by design as there is a bug and typing is low priority for now.
-        # pyrefly: ignore [no-access]
         cls.__origin__ = None
         if "type" in namespace:
             return super().__new__(cls, name, bases, namespace, **kwargs)  # type: ignore[call-overload]
diff --git a/torch/utils/data/datapipes/dataframe/dataframes.py b/torch/utils/data/datapipes/dataframe/dataframes.py
index 5361c29b48224..d2608796dc75e 100644
--- a/torch/utils/data/datapipes/dataframe/dataframes.py
+++ b/torch/utils/data/datapipes/dataframe/dataframes.py
@@ -81,7 +81,6 @@ def __str__(self) -> str:
 
     def _ops_str(self):
         res = ""
-        # pyrefly: ignore [not-iterable]
         for op in self.ctx["operations"]:
             if len(res) > 0:
                 res += "\n"
@@ -91,7 +90,6 @@ def _ops_str(self):
     def __getstate__(self):
         # TODO(VitalyFedyunin): Currently can't pickle (why?)
         self.ctx["schema_df"] = None
-        # pyrefly: ignore [not-iterable]
         for var in self.ctx["variables"]:
             var.calculated_value = None
         state = {}
@@ -115,13 +113,11 @@ def __getitem__(self, key):
         return CaptureGetItem(self, key, ctx=self.ctx)
 
     def __setitem__(self, key, value) -> None:
-        # pyrefly: ignore [missing-attribute]
         self.ctx["operations"].append(CaptureSetItem(self, key, value, ctx=self.ctx))
 
     def __add__(self, add_val):
         res = CaptureAdd(self, add_val, ctx=self.ctx)
         var = CaptureVariable(res, ctx=self.ctx)
-        # pyrefly: ignore [missing-attribute]
         self.ctx["operations"].append(
             CaptureVariableAssign(variable=var, value=res, ctx=self.ctx)
         )
@@ -130,7 +126,6 @@ def __add__(self, add_val):
     def __sub__(self, add_val):
         res = CaptureSub(self, add_val, ctx=self.ctx)
         var = CaptureVariable(res, ctx=self.ctx)
-        # pyrefly: ignore [missing-attribute]
         self.ctx["operations"].append(
             CaptureVariableAssign(variable=var, value=res, ctx=self.ctx)
         )
@@ -140,19 +135,15 @@ def __mul__(self, add_val):
         res = CaptureMul(self, add_val, ctx=self.ctx)
         var = CaptureVariable(res, ctx=self.ctx)
         t = CaptureVariableAssign(variable=var, value=res, ctx=self.ctx)
-        # pyrefly: ignore [missing-attribute]
         self.ctx["operations"].append(t)
         return var
 
     def _is_context_empty(self):
-        # pyrefly: ignore [bad-argument-type]
         return len(self.ctx["operations"]) == 0 and len(self.ctx["variables"]) == 0
 
     def apply_ops_2(self, dataframe) -> None:
         # TODO(VitalyFedyunin): Make this calculation thread safe (as currently it updates pointer)
-        # pyrefly: ignore [unsupported-operation]
         self.ctx["variables"][0].calculated_value = dataframe
-        # pyrefly: ignore [not-iterable]
         for op in self.ctx["operations"]:
             op.execute()
 
@@ -185,7 +176,6 @@ def __call__(self, *args, **kwargs):
         res = CaptureCall(self, ctx=self.ctx, args=args, kwargs=kwargs)
         var = CaptureVariable(None, ctx=self.ctx)
         t = CaptureVariableAssign(ctx=self.ctx, variable=var, value=res)
-        # pyrefly: ignore [missing-attribute]
         self.ctx["operations"].append(t)
         return var
 
@@ -287,9 +277,7 @@ def execute(self):
 
     def apply_ops(self, dataframe):
         # TODO(VitalyFedyunin): Make this calculation thread safe (as currently it updates pointer)
-        # pyrefly: ignore [unsupported-operation]
         self.ctx["variables"][0].calculated_value = dataframe
-        # pyrefly: ignore [not-iterable]
         for op in self.ctx["operations"]:
             op.execute()
         return self.calculated_value
@@ -411,7 +399,6 @@ class CaptureDataFrame(CaptureInitial):
 
 class CaptureDataFrameWithDataPipeOps(CaptureDataFrame):
     def as_datapipe(self):
-        # pyrefly: ignore [unsupported-operation]
         return DataFrameTracedOps(self.ctx["variables"][0].source_datapipe, self)
 
     def raw_iterator(self):
diff --git a/torch/utils/data/datapipes/iter/grouping.py b/torch/utils/data/datapipes/iter/grouping.py
index b113f94bf5f1a..dd3a377523519 100644
--- a/torch/utils/data/datapipes/iter/grouping.py
+++ b/torch/utils/data/datapipes/iter/grouping.py
@@ -219,14 +219,12 @@ def __init__(
         if group_size is not None and buffer_size is not None:
             if not (0 < group_size <= buffer_size):
                 raise AssertionError("group_size must be > 0 and <= buffer_size")
-            # pyrefly: ignore [bad-assignment]
             self.guaranteed_group_size = group_size
         if guaranteed_group_size is not None:
             if group_size is None or not (0 < guaranteed_group_size <= group_size):
                 raise AssertionError(
                     "guaranteed_group_size must be > 0 and <= group_size and group_size must be set"
                 )
-            # pyrefly: ignore [bad-assignment]
             self.guaranteed_group_size = guaranteed_group_size
         self.drop_remaining = drop_remaining
         self.wrapper_class = DataChunk
diff --git a/torch/utils/data/graph.py b/torch/utils/data/graph.py
index 9820aac095ae3..3f8dbb17ceaea 100644
--- a/torch/utils/data/graph.py
+++ b/torch/utils/data/graph.py
@@ -11,6 +11,7 @@
 __all__ = ["traverse", "traverse_dps"]
 
 DataPipe = IterDataPipe | MapDataPipe
+# pyrefly: ignore [invalid-type-alias]
 DataPipeGraph = dict[int, tuple[DataPipe, "DataPipeGraph"]]
 
 
diff --git a/torch/utils/data/typing.ipynb b/torch/utils/data/typing.ipynb
index b25d82d421d9b..27ee65677fb35 100644
--- a/torch/utils/data/typing.ipynb
+++ b/torch/utils/data/typing.ipynb
@@ -16,7 +16,7 @@
    "outputs": [],
    "source": [
     "from torch.utils.data import IterDataPipe\n",
-    "from typing import Any, TypeVar, Union\n",
+    "from typing import Any, TypeVar\n",
     "from collections.abc import Iterator\n",
     "import sys\n",
     "\n",
@@ -739,11 +739,11 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 30,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
-    "class DP(IterDataPipe[Union[int, str]]):\n",
+    "class DP(IterDataPipe[int | str]):\n",
     "    def __init__(self, label) -> None:\n",
     "        if label == 'int':\n",
     "            self.reinforce_type(int)\n",
@@ -829,4 +829,4 @@
  },
  "nbformat": 4,
  "nbformat_minor": 4
-}
+}
\ No newline at end of file
diff --git a/torch/utils/file_baton.py b/torch/utils/file_baton.py
index 5b4f55d8c88dd..74de1495c8fc5 100644
--- a/torch/utils/file_baton.py
+++ b/torch/utils/file_baton.py
@@ -31,7 +31,6 @@ def try_acquire(self) -> bool | None:
             True if the file could be created, else False.
         """
         try:
-            # pyrefly: ignore [bad-assignment]
             self.fd = os.open(self.lock_file_path, os.O_CREAT | os.O_EXCL)
             return True
         except FileExistsError:
diff --git a/torch/utils/flop_counter.py b/torch/utils/flop_counter.py
index 918a22cbefa60..5f015a32f9c31 100644
--- a/torch/utils/flop_counter.py
+++ b/torch/utils/flop_counter.py
@@ -282,18 +282,27 @@ def sdpa_flop_count(query_shape, key_shape, value_shape):
     """
     Count flops for self-attention.
 
-    NB: We can assume that value_shape == key_shape
+    Supports GQA (grouped-query attention) where key/value have fewer heads
+    than the query. The kernel broadcasts KV heads to match query heads.
     """
-    b, h, s_q, d_q = query_shape
-    _b2, _h2, s_k, _d2 = key_shape
+    b, h_q, s_q, d_q = query_shape
+    _b2, h_kv, s_k, _d2 = key_shape
     _b3, _h3, _s3, d_v = value_shape
-    if not b == _b2 == _b3 or not h == _h2 == _h3 or not d_q == _d2 or not s_k == _s3 or not d_q == _d2:
-        raise AssertionError("sdpa_flop_count: query/key/value shapes are incompatible")
+    if not (b == _b2 == _b3 and h_kv == _h3 and d_q == _d2 and s_k == _s3):
+        raise AssertionError(
+            f"sdpa_flop_count: query/key/value shapes are incompatible: "
+            f"q={query_shape}, k={key_shape}, v={value_shape}"
+        )
+    if h_q < h_kv or h_q % h_kv != 0:
+        raise AssertionError(
+            f"sdpa_flop_count: query heads ({h_q}) must be a multiple of "
+            f"key/value heads ({h_kv})"
+        )
     total_flops = 0
-    # q: [b, h, s_q, d_q] @ k: [b, h, d_q, s_k] -> scores: [b, h, s_q, s_k]
-    total_flops += bmm_flop((b * h, s_q, d_q), (b * h, d_q, s_k))
-    # scores: [b, h, s_q, s_k] @ v: [b, h, s_k, d_v] -> out: [b, h, s_q, d_v]
-    total_flops += bmm_flop((b * h, s_q, s_k), (b * h, s_k, d_v))
+    # q: [b, h_q, s_q, d_q] @ k: [b, h_q, d_q, s_k] -> scores: [b, h_q, s_q, s_k]
+    total_flops += bmm_flop((b * h_q, s_q, d_q), (b * h_q, d_q, s_k))
+    # scores: [b, h_q, s_q, s_k] @ v: [b, h_q, s_k, d_v] -> out: [b, h_q, s_q, d_v]
+    total_flops += bmm_flop((b * h_q, s_q, s_k), (b * h_q, s_k, d_v))
     return total_flops
 
 
@@ -490,31 +499,39 @@ def _efficient_attention_forward_flop(
 
 
 def sdpa_backward_flop_count(grad_out_shape, query_shape, key_shape, value_shape):
-    total_flops = 0
-    b, h, s_q, d_q = query_shape
-    _b2, _h2, s_k, _d2 = key_shape
+    b, h_q, s_q, d_q = query_shape
+    _b2, h_kv, s_k, _d2 = key_shape
     _b3, _h3, _s3, d_v = value_shape
     _b4, _h4, _s4, _d4 = grad_out_shape
-    if not b == _b2 == _b3 == _b4 or not h == _h2 == _h3 == _h4 or not d_q == _d2:
-        raise AssertionError("sdpa_backward_flop_count: batch/heads/dimension mismatch among tensors")
-    if not d_v == _d4 or not s_k == _s3 or not s_q == _s4:
-        raise AssertionError("sdpa_backward_flop_count: grad_out/value/key/query shapes are incompatible")
+    if not (b == _b2 == _b3 == _b4 and h_kv == _h3 and h_q == _h4):
+        raise AssertionError(
+            "sdpa_backward_flop_count: batch/heads mismatch among tensors"
+        )
+    if h_q < h_kv or h_q % h_kv != 0:
+        raise AssertionError(
+            f"sdpa_backward_flop_count: query heads ({h_q}) must be a multiple of "
+            f"key/value heads ({h_kv})"
+        )
+    if not (d_q == _d2 and d_v == _d4 and s_k == _s3 and s_q == _s4):
+        raise AssertionError(
+            "sdpa_backward_flop_count: grad_out/value/key/query shapes are incompatible"
+        )
     total_flops = 0
     # Step 1: We recompute the scores matrix.
-    # q: [b, h, s_q, d_q] @ k: [b, h, d_q, s_k] -> scores: [b, h, s_q, s_k]
-    total_flops += bmm_flop((b * h, s_q, d_q), (b * h, d_q, s_k))
+    # q: [b, h_q, s_q, d_q] @ k: [b, h_q, d_q, s_k] -> scores: [b, h_q, s_q, s_k]
+    total_flops += bmm_flop((b * h_q, s_q, d_q), (b * h_q, d_q, s_k))
 
     # Step 2: We propagate the gradients through the score @ v operation.
-    # gradOut: [b, h, s_q, d_v] @ v: [b, h, d_v, s_k] -> gradScores: [b, h, s_q, s_k]
-    total_flops += bmm_flop((b * h, s_q, d_v), (b * h, d_v, s_k))
-    # scores: [b, h, s_k, s_q] @ gradOut: [b, h, s_q, d_v] -> gradV: [b, h, s_k, d_v]
-    total_flops += bmm_flop((b * h, s_k, s_q), (b * h, s_q, d_v))
+    # gradOut: [b, h_q, s_q, d_v] @ v: [b, h_q, d_v, s_k] -> gradScores: [b, h_q, s_q, s_k]
+    total_flops += bmm_flop((b * h_q, s_q, d_v), (b * h_q, d_v, s_k))
+    # scores: [b, h_q, s_k, s_q] @ gradOut: [b, h_q, s_q, d_v] -> gradV: [b, h_q, s_k, d_v]
+    total_flops += bmm_flop((b * h_q, s_k, s_q), (b * h_q, s_q, d_v))
 
     # Step 3: We propagate th gradients through the k @ v operation
-    # gradScores: [b, h, s_q, s_k] @ k: [b, h, s_k, d_q] -> gradQ: [b, h, s_q, d_q]
-    total_flops += bmm_flop((b * h, s_q, s_k), (b * h, s_k, d_q))
-    # q: [b, h, d_q, s_q] @ gradScores: [b, h, s_q, s_k] -> gradK: [b, h, d_q, s_k]
-    total_flops += bmm_flop((b * h, d_q, s_q), (b * h, s_q, s_k))
+    # gradScores: [b, h_q, s_q, s_k] @ k: [b, h_q, s_k, d_q] -> gradQ: [b, h_q, s_q, d_q]
+    total_flops += bmm_flop((b * h_q, s_q, s_k), (b * h_q, s_k, d_q))
+    # q: [b, h_q, d_q, s_q] @ gradScores: [b, h_q, s_q, s_k] -> gradK: [b, h_q, d_q, s_k]
+    total_flops += bmm_flop((b * h_q, d_q, s_q), (b * h_q, s_q, s_k))
     return total_flops
 
 
@@ -591,6 +608,85 @@ def _efficient_attention_backward_flop(
     )
 
 
+def _varlen_attn_forward_flop(
+    query,
+    key,
+    value,
+    cu_seq_q,
+    cu_seq_k,
+    max_q,
+    max_k,
+    *args,
+    out_val=None,
+    **kwargs,
+) -> int:
+    """Count flops for varlen_attn forward."""
+    sizes = _unpack_flash_attention_nested_shapes(
+        query=query,
+        key=key,
+        value=value,
+        cum_seq_q=cu_seq_q,
+        cum_seq_k=cu_seq_k if cu_seq_k is not None else cu_seq_q,
+        max_q=max_q,
+        max_k=max_k,
+    )
+    return sum(
+        sdpa_flop_count(query_shape, key_shape, value_shape)
+        for query_shape, key_shape, value_shape, _ in sizes
+    )
+
+
+def _varlen_attn_out_flop(
+    out,
+    query,
+    key,
+    value,
+    cu_seq_q,
+    cu_seq_k,
+    max_q,
+    max_k,
+    *args,
+    out_val=None,
+    **kwargs,
+) -> int:
+    """Count flops for varlen_attn_out forward."""
+    return _varlen_attn_forward_flop(
+        query, key, value, cu_seq_q, cu_seq_k, max_q, max_k,
+    )
+
+
+def _varlen_attn_backward_flop(
+    grad_out,
+    query,
+    key,
+    value,
+    out,
+    lse,
+    cu_seq_q,
+    cu_seq_k,
+    max_q,
+    max_k,
+    *args,
+    out_val=None,
+    **kwargs,
+) -> int:
+    """Count flops for varlen_attn backward."""
+    sizes = _unpack_flash_attention_nested_shapes(
+        query=query,
+        key=key,
+        value=value,
+        grad_out=grad_out,
+        cum_seq_q=cu_seq_q,
+        cum_seq_k=cu_seq_k,
+        max_q=max_q,
+        max_k=max_k,
+    )
+    return sum(
+        sdpa_backward_flop_count(grad_out_shape, query_shape, key_shape, value_shape)
+        for query_shape, key_shape, value_shape, grad_out_shape in sizes
+    )
+
+
 flop_registry = {
     aten.mm: mm_flop,
     aten.addmm: addmm_flop,
diff --git a/torch/utils/hipify/cuda_to_hip_mappings.py b/torch/utils/hipify/cuda_to_hip_mappings.py
index 3dff877946cea..dfef0e363db93 100644
--- a/torch/utils/hipify/cuda_to_hip_mappings.py
+++ b/torch/utils/hipify/cuda_to_hip_mappings.py
@@ -24,7 +24,6 @@
     ("std::frexp", ("::frexp")),
 ])
 
-# pyrefly: ignore  # no-matching-overload
 CUDA_TYPE_NAME_MAP = collections.OrderedDict([
     ("CUresult", "hipError_t"),
     ("cudaError_t", "hipError_t"),
@@ -350,7 +349,6 @@
     ("tensorpipe/tensorpipe_cuda.h", "tensorpipe/tensorpipe_hip.h"),
 ])
 
-# pyrefly: ignore  # no-matching-overload
 CUDA_IDENTIFIER_MAP = collections.OrderedDict([
     ("__CUDACC__", "__HIPCC__"),
     ("CUDA_ERROR_INVALID_CONTEXT", "hipErrorInvalidContext"),
@@ -768,6 +766,7 @@
     ("CU_MEM_HANDLE_TYPE_WIN32", "hipMemHandleTypeWin32"),
     ("CU_MEM_HANDLE_TYPE_WIN32_KMT", "hipMemHandleTypeWin32Kmt"),
     ("CU_MEM_LOCATION_TYPE_DEVICE", "hipMemLocationTypeDevice"),
+    ("CU_MEM_LOCATION_TYPE_HOST", "hipMemLocationTypeHost"),
     ("CU_MEM_LOCATION_TYPE_INVALID", "hipMemLocationTypeInvalid"),
     ("CU_MEM_OPERATION_TYPE_MAP", "hipMemOperationTypeMap"),
     ("CU_MEM_OPERATION_TYPE_UNMAP", "hipMemOperationTypeUnmap"),
@@ -1355,6 +1354,7 @@
     ("cudaEventCreateWithFlags", "hipEventCreateWithFlags"),
     ("cudaEventDestroy", "hipEventDestroy"),
     ("cudaEventRecord", "hipEventRecord"),
+    ("cudaEventRecordWithFlags", "hipEventRecordWithFlags"),
     ("cudaEventElapsedTime", "hipEventElapsedTime"),
     ("cudaEventSynchronize", "hipEventSynchronize"),
     ("cudaEventQuery", "hipEventQuery"),
@@ -1362,6 +1362,10 @@
     ("cudaEventBlockingSync", "hipEventBlockingSync"),
     ("cudaEventDisableTiming", "hipEventDisableTiming"),
     ("cudaEventInterprocess", "hipEventInterprocess"),
+    ("cudaEventRecordDefault", "hipEventRecordDefault"),
+    ("cudaEventRecordExternal", "hipEventRecordExternal"),
+    ("cudaEventWaitDefault", "hipEventWaitDefault"),
+    ("cudaEventWaitExternal", "hipEventWaitExternal"),
     ("cudaStreamCreate", "hipStreamCreate"),
     ("cudaStreamCreateWithFlags", "hipStreamCreateWithFlags"),
     ("cudaStreamCreateWithPriority", "hipStreamCreateWithPriority"),
@@ -3394,6 +3398,23 @@
     ("cudnnTensorDescriptor_t ", "miopenTensorDescriptor_t "),
     ("CUDNN_ENFORCE", "MIOPEN_ENFORCE"),
     ("CUDNN_CHECK", "MIOPEN_CHECK"),
+    # NVSHMEM → rocSHMEM mappings (only symbols used in hipified files:
+    # NVSHMEMSymmetricMemory.cpp and nvshmem_team_manager.hpp).
+    ("NVSHMEM_TEAM_INVALID", "rocshmem::ROCSHMEM_TEAM_INVALID"),
+    ("NVSHMEM_TEAM_WORLD", "rocshmem::ROCSHMEM_TEAM_WORLD"),
+    ("NVSHMEMX_INIT_WITH_UNIQUEID", "rocshmem::ROCSHMEM_INIT_WITH_UNIQUEID"),
+
+    ("nvshmem_malloc", "rocshmem::rocshmem_malloc"),
+    ("nvshmem_free", "rocshmem::rocshmem_free"),
+    ("nvshmem_ptr", "rocshmem::rocshmem_ptr"),
+    ("nvshmem_team_t", "rocshmem::rocshmem_team_t"),
+    ("nvshmem_team_split_strided", "rocshmem::rocshmem_team_split_strided"),
+
+    ("nvshmemx_uniqueid_t", "rocshmem::rocshmem_uniqueid_t"),
+    ("nvshmemx_get_uniqueid", "rocshmem::rocshmem_get_uniqueid"),
+    ("nvshmemx_init_attr", "rocshmem::rocshmem_init_attr"),
+    ("nvshmemx_init_attr_t", "rocshmem::rocshmem_init_attr_t"),
+    ("nvshmemx_set_attr_uniqueid_args", "rocshmem::rocshmem_set_attr_uniqueid_args"),
 ])
 
 C10_MAPPINGS = collections.OrderedDict([
@@ -3417,6 +3438,7 @@
     ("c10/cuda/CUDAEvent.h", "c10/hip/HIPEvent.h"),
     ("c10/cuda/impl/CUDAGuardImpl.h", "c10/hip/impl/HIPGuardImpl.h"),
     ("c10/cuda/impl/CUDATest.h", "c10/hip/impl/HIPTest.h"),
+    ("CUDATest.hpp", "HIPTest.hpp"),
     ("c10/cuda/impl/cuda_cmake_macros.h", "c10/hip/impl/hip_cmake_macros.h"),
     # TODO: Remove these. They were necessary for Meta-internal builds.
     ("c10::hip::c10_hip_check_implementation", "c10::cuda::c10_cuda_check_implementation"),
diff --git a/torch/utils/hipify/hipify_python.py b/torch/utils/hipify/hipify_python.py
index 557f8c128a905..3919d61bbaa0c 100755
--- a/torch/utils/hipify/hipify_python.py
+++ b/torch/utils/hipify/hipify_python.py
@@ -647,6 +647,8 @@ def is_pytorch_file(rel_filepath) -> bool:
         return True
     if rel_filepath.startswith("tools/autograd/templates/"):
         return True
+    if rel_filepath.startswith("test/cpp/c10d/"):
+        return True
     return False
 
 
diff --git a/torch/utils/hooks.py b/torch/utils/hooks.py
index 8e89d3ec9b3a0..1e3f6fb9ab09d 100644
--- a/torch/utils/hooks.py
+++ b/torch/utils/hooks.py
@@ -145,7 +145,6 @@ def hook(grad_input, _):
 
                 res = out
 
-            # pyrefly: ignore [bad-assignment]
             self.grad_outputs = None
 
             return self._unpack_none(self.input_tensors_index, res)
diff --git a/torch/utils/tensorboard/_onnx_graph.py b/torch/utils/tensorboard/_onnx_graph.py
index 2405e1320a19f..59d1d5b5f0a5c 100644
--- a/torch/utils/tensorboard/_onnx_graph.py
+++ b/torch/utils/tensorboard/_onnx_graph.py
@@ -23,7 +23,6 @@ def parse(graph):
         print(node.name)
         shapeproto = TensorShapeProto(
             dim=[
-                # pyrefly: ignore [missing-attribute]
                 TensorShapeProto.Dim(size=d.dim_value)
                 for d in node.type.tensor_type.shape.dim
             ]
diff --git a/torch/utils/tensorboard/_proto_graph.py b/torch/utils/tensorboard/_proto_graph.py
index 7815a2aa8d61d..8a9dc8536755c 100644
--- a/torch/utils/tensorboard/_proto_graph.py
+++ b/torch/utils/tensorboard/_proto_graph.py
@@ -6,7 +6,6 @@
 from tensorboard.compat.proto.tensor_shape_pb2 import TensorShapeProto
 
 
-# pyrefly: ignore [not-a-type]
 def attr_value_proto(dtype: object, shape: Sequence[int] | None, s: str | None) -> dict[str, AttrValue]:
     """Create a dict of objects matching a NodeDef's attr field.
 
@@ -19,18 +18,15 @@ def attr_value_proto(dtype: object, shape: Sequence[int] | None, s: str | None)
         attr["attr"] = AttrValue(s=s.encode(encoding="utf_8"))
     if shape is not None:
         shapeproto = tensor_shape_proto(shape)
-        # pyrefly: ignore [missing-attribute]
         attr["_output_shapes"] = AttrValue(list=AttrValue.ListValue(shape=[shapeproto]))
     return attr
 
 
-# pyrefly: ignore [not-a-type]
 def tensor_shape_proto(outputsize: Sequence[int]) -> TensorShapeProto:
     """Create an object matching a tensor_shape field.
 
     Follows https://github.com/tensorflow/tensorboard/blob/master/tensorboard/compat/proto/tensor_shape.proto .
     """
-    # pyrefly: ignore [missing-attribute]
     return TensorShapeProto(dim=[TensorShapeProto.Dim(size=d) for d in outputsize])
 
 
@@ -42,7 +38,6 @@ def node_proto(
     shape: tuple[int, ...] | None = None,
     outputsize: Sequence[int] | None = None,
     attributes: str = "",
-    # pyrefly: ignore [not-a-type]
 ) -> NodeDef:
     """Create an object matching a NodeDef.
 
diff --git a/torch/utils/tensorboard/summary.py b/torch/utils/tensorboard/summary.py
index 2a29e87c1d8b2..f5d86875387bf 100644
--- a/torch/utils/tensorboard/summary.py
+++ b/torch/utils/tensorboard/summary.py
@@ -260,7 +260,6 @@ def hparams(hparam_dict=None, metric_dict=None, hparam_domain_discrete=None):
             hps.append(
                 HParamInfo(
                     name=k,
-                    # pyrefly: ignore [missing-attribute]
                     type=DataType.Value("DATA_TYPE_FLOAT64"),
                     domain_discrete=domain_discrete,
                 )
@@ -283,7 +282,6 @@ def hparams(hparam_dict=None, metric_dict=None, hparam_domain_discrete=None):
             hps.append(
                 HParamInfo(
                     name=k,
-                    # pyrefly: ignore [missing-attribute]
                     type=DataType.Value("DATA_TYPE_STRING"),
                     domain_discrete=domain_discrete,
                 )
@@ -306,7 +304,6 @@ def hparams(hparam_dict=None, metric_dict=None, hparam_domain_discrete=None):
             hps.append(
                 HParamInfo(
                     name=k,
-                    # pyrefly: ignore [missing-attribute]
                     type=DataType.Value("DATA_TYPE_BOOL"),
                     domain_discrete=domain_discrete,
                 )
@@ -316,7 +313,6 @@ def hparams(hparam_dict=None, metric_dict=None, hparam_domain_discrete=None):
         if isinstance(v, torch.Tensor):
             v = make_np(v)[0]
             ssi.hparams[k].number_value = v
-            # pyrefly: ignore [missing-attribute]
             hps.append(HParamInfo(name=k, type=DataType.Value("DATA_TYPE_FLOAT64")))
             continue
         raise ValueError(
@@ -325,12 +321,10 @@ def hparams(hparam_dict=None, metric_dict=None, hparam_domain_discrete=None):
 
     content = HParamsPluginData(session_start_info=ssi, version=PLUGIN_DATA_VERSION)
     smd = SummaryMetadata(
-        # pyrefly: ignore [missing-attribute]
         plugin_data=SummaryMetadata.PluginData(
             plugin_name=PLUGIN_NAME, content=content.SerializeToString()
         )
     )
-    # pyrefly: ignore [missing-attribute]
     ssi = Summary(value=[Summary.Value(tag=SESSION_START_INFO_TAG, metadata=smd)])
 
     mts = [MetricInfo(name=MetricName(tag=k)) for k in metric_dict]
@@ -339,24 +333,19 @@ def hparams(hparam_dict=None, metric_dict=None, hparam_domain_discrete=None):
 
     content = HParamsPluginData(experiment=exp, version=PLUGIN_DATA_VERSION)
     smd = SummaryMetadata(
-        # pyrefly: ignore [missing-attribute]
         plugin_data=SummaryMetadata.PluginData(
             plugin_name=PLUGIN_NAME, content=content.SerializeToString()
         )
     )
-    # pyrefly: ignore [missing-attribute]
     exp = Summary(value=[Summary.Value(tag=EXPERIMENT_TAG, metadata=smd)])
 
-    # pyrefly: ignore [missing-attribute]
     sei = SessionEndInfo(status=Status.Value("STATUS_SUCCESS"))
     content = HParamsPluginData(session_end_info=sei, version=PLUGIN_DATA_VERSION)
     smd = SummaryMetadata(
-        # pyrefly: ignore [missing-attribute]
         plugin_data=SummaryMetadata.PluginData(
             plugin_name=PLUGIN_NAME, content=content.SerializeToString()
         )
     )
-    # pyrefly: ignore [missing-attribute]
     sei = Summary(value=[Summary.Value(tag=SESSION_END_INFO_TAG, metadata=smd)])
 
     return exp, ssi, sei
@@ -390,12 +379,10 @@ def scalar(name, tensor, collections=None, new_style=False, double_precision=Fal
         if double_precision:
             tensor_proto = TensorProto(double_val=[scalar], dtype="DT_DOUBLE")
 
-        # pyrefly: ignore [missing-attribute]
         plugin_data = SummaryMetadata.PluginData(plugin_name="scalars")
         smd = SummaryMetadata(plugin_data=plugin_data)
         return Summary(
             value=[
-                # pyrefly: ignore [missing-attribute]
                 Summary.Value(
                     tag=name,
                     tensor=tensor_proto,
@@ -404,7 +391,6 @@ def scalar(name, tensor, collections=None, new_style=False, double_precision=Fal
             ]
         )
     else:
-        # pyrefly: ignore [missing-attribute]
         return Summary(value=[Summary.Value(tag=name, simple_value=scalar)])
 
 
@@ -432,7 +418,6 @@ def tensor_proto(tag, tensor):
             **{
                 "dtype": dtype,
                 "tensor_shape": TensorShapeProto(
-                    # pyrefly: ignore [missing-attribute]
                     dim=[TensorShapeProto.Dim(size=x) for x in tensor.shape]
                 ),
                 field_name: conversion_fn(tensor),
@@ -441,10 +426,8 @@ def tensor_proto(tag, tensor):
     else:
         raise ValueError(f"{tag} has unsupported tensor dtype {tensor.dtype}")
 
-    # pyrefly: ignore [missing-attribute]
     plugin_data = SummaryMetadata.PluginData(plugin_name="tensor")
     smd = SummaryMetadata(plugin_data=plugin_data)
-    # pyrefly: ignore [missing-attribute]
     return Summary(value=[Summary.Value(tag=tag, metadata=smd, tensor=tensor_proto)])
 
 
@@ -478,7 +461,6 @@ def histogram_raw(name, min, max, num, sum, sum_squares, bucket_limits, bucket_c
         bucket_limit=bucket_limits,
         bucket=bucket_counts,
     )
-    # pyrefly: ignore [missing-attribute]
     return Summary(value=[Summary.Value(tag=name, histo=hist)])
 
 
@@ -501,7 +483,6 @@ def histogram(name, values, bins, max_bins=None):
     """
     values = make_np(values)
     hist = make_histogram(values.astype(float), bins, max_bins)
-    # pyrefly: ignore [missing-attribute]
     return Summary(value=[Summary.Value(tag=name, histo=hist)])
 
 
@@ -595,7 +576,6 @@ def image(tag, tensor, rescale=1, dataformats="NCHW"):
     tensor = tensor.astype(np.float32)
     tensor = (tensor * scale_factor).clip(0, 255).astype(np.uint8)
     image = make_image(tensor, rescale=rescale)
-    # pyrefly: ignore [missing-attribute]
     return Summary(value=[Summary.Value(tag=tag, image=image)])
 
 
@@ -613,7 +593,6 @@ def image_boxes(
         rois=tensor_boxes,
         labels=labels,
     )
-    # pyrefly: ignore [missing-attribute]
     return Summary(value=[Summary.Value(tag=tag, image=image)])
 
 
@@ -652,7 +631,6 @@ def make_image(tensor, rescale=1, rois=None, labels=None):
     image.save(output, format="PNG")
     image_string = output.getvalue()
     output.close()
-    # pyrefly: ignore [missing-attribute]
     return Summary.Image(
         height=height,
         width=width,
@@ -669,7 +647,6 @@ def video(tag, tensor, fps=4):
     tensor = tensor.astype(np.float32)
     tensor = (tensor * scale_factor).clip(0, 255).astype(np.uint8)
     video = make_video(tensor, fps)
-    # pyrefly: ignore [missing-attribute]
     return Summary(value=[Summary.Value(tag=tag, image=video)])
 
 
@@ -707,7 +684,6 @@ def make_video(tensor, fps):
         f.seek(0)
         tensor_string = f.read()
 
-        # pyrefly: ignore [missing-attribute]
         return Summary.Image(
             height=h, width=w, colorspace=c, encoded_image_string=tensor_string
         )
@@ -735,7 +711,6 @@ def audio(tag, tensor, sample_rate=44100):
         wave_write.writeframes(array.data)
     audio_string = fio.getvalue()
     fio.close()
-    # pyrefly: ignore [missing-attribute]
     audio = Summary.Audio(
         sample_rate=sample_rate,
         num_channels=1,
@@ -743,7 +718,6 @@ def audio(tag, tensor, sample_rate=44100):
         encoded_audio_string=audio_string,
         content_type="audio/wav",
     )
-    # pyrefly: ignore [missing-attribute]
     return Summary(value=[Summary.Value(tag=tag, audio=audio)])
 
 
@@ -758,7 +732,6 @@ def custom_scalars(layout):
                     raise AssertionError("len(tags) != 3")
                 mgcc = layout_pb2.MarginChartContent(
                     series=[
-                        # pyrefly: ignore [missing-attribute]
                         layout_pb2.MarginChartContent.Series(
                             value=tags[0], lower=tags[1], upper=tags[2]
                         )
@@ -772,7 +745,6 @@ def custom_scalars(layout):
         categories.append(layout_pb2.Category(title=k, chart=charts))
 
     layout = layout_pb2.Layout(category=categories)
-    # pyrefly: ignore [missing-attribute]
     plugin_data = SummaryMetadata.PluginData(plugin_name="custom_scalars")
     smd = SummaryMetadata(plugin_data=plugin_data)
     tensor = TensorProto(
@@ -782,14 +754,12 @@ def custom_scalars(layout):
     )
     return Summary(
         value=[
-            # pyrefly: ignore [missing-attribute]
             Summary.Value(tag="custom_scalars__config__", tensor=tensor, metadata=smd)
         ]
     )
 
 
 def text(tag, text):
-    # pyrefly: ignore [missing-attribute]
     plugin_data = SummaryMetadata.PluginData(
         plugin_name="text", content=TextPluginData(version=0).SerializeToString()
     )
@@ -797,11 +767,9 @@ def text(tag, text):
     tensor = TensorProto(
         dtype="DT_STRING",
         string_val=[text.encode(encoding="utf_8")],
-        # pyrefly: ignore [missing-attribute]
         tensor_shape=TensorShapeProto(dim=[TensorShapeProto.Dim(size=1)]),
     )
     return Summary(
-        # pyrefly: ignore [missing-attribute]
         value=[Summary.Value(tag=tag + "/text_summary", metadata=smd, tensor=tensor)]
     )
 
@@ -815,7 +783,6 @@ def pr_curve_raw(
     pr_curve_plugin_data = PrCurvePluginData(
         version=0, num_thresholds=num_thresholds
     ).SerializeToString()
-    # pyrefly: ignore [missing-attribute]
     plugin_data = SummaryMetadata.PluginData(
         plugin_name="pr_curves", content=pr_curve_plugin_data
     )
@@ -825,14 +792,11 @@ def pr_curve_raw(
         float_val=data.reshape(-1).tolist(),
         tensor_shape=TensorShapeProto(
             dim=[
-                # pyrefly: ignore [missing-attribute]
                 TensorShapeProto.Dim(size=data.shape[0]),
-                # pyrefly: ignore [missing-attribute]
                 TensorShapeProto.Dim(size=data.shape[1]),
             ]
         ),
     )
-    # pyrefly: ignore [missing-attribute]
     return Summary(value=[Summary.Value(tag=tag, metadata=smd, tensor=tensor)])
 
 
@@ -845,7 +809,6 @@ def pr_curve(tag, labels, predictions, num_thresholds=127, weights=None):
     pr_curve_plugin_data = PrCurvePluginData(
         version=0, num_thresholds=num_thresholds
     ).SerializeToString()
-    # pyrefly: ignore [missing-attribute]
     plugin_data = SummaryMetadata.PluginData(
         plugin_name="pr_curves", content=pr_curve_plugin_data
     )
@@ -855,14 +818,11 @@ def pr_curve(tag, labels, predictions, num_thresholds=127, weights=None):
         float_val=data.reshape(-1).tolist(),
         tensor_shape=TensorShapeProto(
             dim=[
-                # pyrefly: ignore [missing-attribute]
                 TensorShapeProto.Dim(size=data.shape[0]),
-                # pyrefly: ignore [missing-attribute]
                 TensorShapeProto.Dim(size=data.shape[1]),
             ]
         ),
     )
-    # pyrefly: ignore [missing-attribute]
     return Summary(value=[Summary.Value(tag=tag, metadata=smd, tensor=tensor)])
 
 
@@ -947,17 +907,13 @@ def _get_tensor_summary(
         float_val=tensor.reshape(-1).tolist(),
         tensor_shape=TensorShapeProto(
             dim=[
-                # pyrefly: ignore [missing-attribute]
                 TensorShapeProto.Dim(size=tensor.shape[0]),
-                # pyrefly: ignore [missing-attribute]
                 TensorShapeProto.Dim(size=tensor.shape[1]),
-                # pyrefly: ignore [missing-attribute]
                 TensorShapeProto.Dim(size=tensor.shape[2]),
             ]
         ),
     )
 
-    # pyrefly: ignore [missing-attribute]
     tensor_summary = Summary.Value(
         tag=metadata.get_instance_name(name, content_type),
         tensor=tensor,
@@ -1005,11 +961,8 @@ def mesh(
 
     summaries = []
     tensors = [
-        # pyrefly: ignore [missing-attribute]
         (vertices, MeshPluginData.VERTEX),
-        # pyrefly: ignore [missing-attribute]
         (faces, MeshPluginData.FACE),
-        # pyrefly: ignore [missing-attribute]
         (colors, MeshPluginData.COLOR),
     ]
     tensors = [tensor for tensor in tensors if tensor[0] is not None]
diff --git a/torch/utils/tensorboard/writer.py b/torch/utils/tensorboard/writer.py
index 008ae59e94e6a..2f1ccce77cf72 100644
--- a/torch/utils/tensorboard/writer.py
+++ b/torch/utils/tensorboard/writer.py
@@ -1,9 +1,10 @@
 # mypy: allow-untyped-defs
 """Provide an API for writing protocol buffers to event files to be consumed by TensorBoard for visualization."""
+from __future__ import annotations
 
 import os
 import time
-from typing import TYPE_CHECKING, Union
+from typing import TYPE_CHECKING
 
 import torch
 
@@ -166,7 +167,6 @@ def reopen(self) -> None:
         The events will go into a new events file.
         Does nothing if the EventFileWriter was not closed.
         """
-        # pyrefly: ignore [missing-attribute]
         self.event_writer.reopen()
 
 
@@ -255,9 +255,7 @@ def __init__(
         buckets = []
         neg_buckets = []
         while v < 1e20:
-            # pyrefly: ignore [bad-argument-type]
             buckets.append(v)
-            # pyrefly: ignore [bad-argument-type]
             neg_buckets.append(-v)
             v *= 1.1
         self.default_bins = neg_buckets[::-1] + [0] + buckets
@@ -281,7 +279,6 @@ def _get_file_writer(self):
                 self.file_writer.add_event(
                     Event(
                         step=most_recent_step,
-                        # pyrefly: ignore [missing-attribute]
                         session_log=SessionLog(status=SessionLog.START),
                     )
                 )
@@ -732,7 +729,7 @@ def add_image_with_boxes(
     def add_figure(
         self,
         tag: str,
-        figure: Union["Figure", list["Figure"]],
+        figure: Figure | list[Figure],
         global_step: int | None = None,
         close: bool = True,
         walltime: float | None = None,
diff --git a/torch/utils/viz/MemoryViz.js b/torch/utils/viz/MemoryViz.js
index 20b08b801eade..4518c7755724d 100644
--- a/torch/utils/viz/MemoryViz.js
+++ b/torch/utils/viz/MemoryViz.js
@@ -1,3 +1,61 @@
+/**
+ * ================================================================================
+ * MemoryViz.js - PyTorch Memory Visualization Tool
+ * ================================================================================
+ *
+ * OVERVIEW:
+ * ---------
+ * This file contains the core visualization logic for PyTorch's memory profiler.
+ * It renders memory allocation timelines, stack traces, and provides interactive
+ * exploration of memory snapshots captured during model execution.
+ *
+ * KEY FEATURES:
+ * - Multiple visualization tabs/views for different memory analysis perspectives
+ * - Interactive stack trace display (supports both click and hover modes)
+ * - Zoom and brush controls for navigating large memory timelines
+ * - Support for loading memory snapshot files (.pickle format)
+ *
+ * ================================================================================
+ * TESTING INSTRUCTIONS FOR ENGINEERS & AGENTS
+ * ================================================================================
+ *
+ * 1. LOCAL TESTING SETUP:
+ *    - Create a simple HTML file that references this JS file:
+ *
+ *      <!DOCTYPE html>
+ *      <html>
+ *      <head><title>MemoryViz Test</title></head>
+ *      <body>
+ *        <script type="module" src="MemoryViz.js"></script>
+ *      </body>
+ *      </html>
+ *
+ *    - Serve locally using: python3 -m http.server 8888
+ *    - Open http://localhost:8888 in your browser
+ *
+ * 2. WHAT TO TEST:
+ *    - Ensure ALL tabs/views render correctly and switch properly
+ *    - Verify BOTH interaction modes work:
+ *      * Click mode: stack traces appear on click
+ *      * Hover mode: stack traces appear on mouseover
+ *    - Test zoom and brush controls for timeline navigation
+ *    - Verify memory allocation blocks are rendered and interactive
+ *
+ * 3. TEST DATA REQUIREMENTS:
+ *    - DO NOT just test with small dummy .pickle files
+ *    - Use realistic, decent-sized .pickle files (10-100+ MB range)
+ *    - Large files stress-test rendering performance and memory handling
+ *    - Test with snapshots from real model training/inference runs
+ *
+ * 4. COMMON ISSUES TO WATCH FOR:
+ *    - Performance degradation with large snapshots
+ *    - Stack trace popups not appearing or positioning incorrectly
+ *    - Tab switching not updating the visualization properly
+ *    - Zoom/brush state not persisting across interactions
+ *
+ * ================================================================================
+ */
+
 'use strict';
 
 import * as d3 from "https://cdn.jsdelivr.net/npm/d3@7/+esm";
@@ -5,6 +63,9 @@ import {axisLeft} from "https://cdn.jsdelivr.net/npm/d3-axis@3/+esm";
 import {scaleLinear} from "https://cdn.jsdelivr.net/npm/d3-scale@4/+esm";
 import {zoom, zoomIdentity} from "https://cdn.jsdelivr.net/npm/d3-zoom@3/+esm";
 import {brushX} from "https://cdn.jsdelivr.net/npm/d3-brush@3/+esm";
+import {process_alloc_data, isPrivatePoolId, formatSize, formatAddr,
+        elideRepeats, frameFilter, format_user_metadata,
+        format_forward_frames, format_frames} from "./process_alloc_data.js";
 
 // Global configuration for trace interaction mode
 // 'hover' = show trace on hover (default)
@@ -48,8 +109,8 @@ function version_space() {
   };
 }
 
-function Segment(addr, size, stream, frames, version, user_metadata) {
-  return {addr, size, stream, version, frames, user_metadata};
+function Segment(addr, size, stream, frames, version, user_metadata, segment_pool_id) {
+  return {addr, size, stream, version, frames, user_metadata, segment_pool_id};
 }
 
 function Block(addr, size, requested_size, frames, free_requested, version, user_metadata) {
@@ -115,25 +176,6 @@ function EventSelector(outer, events, stack_info, memory_view) {
   return es;
 }
 
-function formatSize(num, showBytes = true) {
-  const orig = num;
-  // https://stackoverflow.com/questions/1094841/get-human-readable-version-of-file-size
-  const units = ['', 'Ki', 'Mi', 'Gi', 'Ti', 'Pi', 'Ei', 'Zi'];
-  for (const unit of units) {
-    if (Math.abs(num) < 1024.0) {
-      if (showBytes) {
-        return `${num.toFixed(1)}${unit}B (${orig} bytes)`;
-      }
-      return `${num.toFixed(1)}${unit}B`;
-    }
-    num /= 1024.0;
-  }
-  return `${num.toFixed(1)}YiB`;
-}
-function formatAddr(event) {
-  const prefix = event.action.startsWith('segment') ? 's\'' : 'b\'';
-  return `${prefix}${event.addr.toString(16)}_${event.version}`;
-}
 function formatEvent(event) {
   const stream =
     event.stream === null ? '' : `\n              (stream ${event.stream})`;
@@ -246,6 +288,7 @@ function MemoryView(outer, stack_info, snapshot, device) {
         seg.frames || [],
         seg.version,
         seg.user_metadata,
+        seg.segment_pool_id,
       ),
     );
     for (const b of seg.blocks) {
@@ -490,13 +533,17 @@ function MemoryView(outer, stack_info, snapshot, device) {
           const user_metadata_str = format_user_metadata(t.user_metadata);
           const frames_str = format_frames(t.frames);
           const forward_frames_str = format_forward_frames(t.forward_frames);
+          let pool_str = '';
+          if (isPrivatePoolId(t.segment_pool_id)) {
+            pool_str = `, pool_id (${t.segment_pool_id[0]}, ${t.segment_pool_id[1]})`;
+          }
           return (
             `s${t.addr.toString(16)}_${t.version}: segment ${formatSize(
               t.size,
             )} allocated, ` +
             `${formatSize(free)} free${internal} (stream ${
               t.stream
-            })\n` +
+            }${pool_str})\n` +
             (user_metadata_str ? user_metadata_str + '\n' : '') +
             frames_str +
             forward_frames_str
@@ -564,11 +611,15 @@ function MemoryView(outer, stack_info, snapshot, device) {
           const user_metadata_str = format_user_metadata(t.user_metadata);
           const frames_str = format_frames(t.frames);
           const forward_frames_str = format_forward_frames(t.forward_frames);
+          let pool_str = '';
+          if (isPrivatePoolId(t.segment?.segment_pool_id)) {
+            pool_str = `, pool_id (${t.segment.segment_pool_id[0]}, ${t.segment.segment_pool_id[1]})`;
+          }
           return (
             `b${t.addr.toString(16)}_${t.version} ` +
             `${formatSize(t.requested_size)} allocation${requested} (stream ${
               t.segment.stream
-            })\n` +
+            }${pool_str})\n` +
             (user_metadata_str ? user_metadata_str + '\n' : '') +
             frames_str +
             forward_frames_str
@@ -778,6 +829,7 @@ function annotate_snapshot(snapshot) {
         }
       }
       b.version = snapshot.block_version(b.addr, false);
+      b.segment_pool_id = seg.segment_pool_id;
       // Note [BigInt and Number Safe Arithmetic]
       // Device pointer addresses may be represented as either Number or BigInt.
       // Use explicit conversions to perform arithmetic safely and avoid mixing
@@ -794,364 +846,6 @@ function annotate_snapshot(snapshot) {
   }
 }
 
-function elideRepeats(frames) {
-  const result = [];
-  const length = frames.length;
-  for (let i = 0; i < length; ) {
-    let j = i + 1;
-    const f = frames[i];
-    while (j < length && f === frames[j]) {
-      j++;
-    }
-    switch (j - i) {
-      case 1:
-        result.push(f);
-        break;
-      case 2:
-        result.push(f, f);
-        break;
-      default:
-        result.push(f, `<repeats ${j - i - 1} times>`);
-        break;
-    }
-    i = j;
-  }
-  return result;
-}
-function frameFilter({name, filename}) {
-  const omitFunctions = [
-    'unwind::unwind',
-    'CapturedTraceback::gather',
-    'gather_with_cpp',
-    '_start',
-    '__libc_start_main',
-    'PyEval_',
-    'PyObject_',
-    'PyFunction_',
-  ];
-
-  const omitFilenames = [
-    'core/boxing',
-    '/Register',
-    '/Redispatch',
-    'pythonrun.c',
-    'Modules/main.c',
-    'Objects/call.c',
-    'Objects/methodobject.c',
-    'pycore_ceval.h',
-    'ceval.c',
-    'cpython/abstract.h',
-  ];
-
-  for (const of of omitFunctions) {
-    if (name.includes(of)) {
-      return false;
-    }
-  }
-
-  for (const of of omitFilenames) {
-    if (filename.includes(of)) {
-      return false;
-    }
-  }
-
-  return true;
-}
-
-function format_user_metadata(user_metadata) {
-  if (!user_metadata) {
-    return '';
-  }
-  // Handle string metadata
-  if (typeof user_metadata === 'string') {
-    return `User Metadata:\n  ${user_metadata}`;
-  }
-  // Handle object metadata
-  if (typeof user_metadata === 'object' && Object.keys(user_metadata).length === 0) {
-    return '';
-  }
-  const metadata_lines = Object.entries(user_metadata)
-    .map(([key, value]) => `  ${key}: ${value}`);
-  return 'User Metadata:\n' + metadata_lines.join('\n');
-}
-
-function format_forward_frames(forward_frames) {
-  if (!forward_frames || forward_frames.length === 0) {
-    return '';
-  }
-  // forward_frames is a list of strings (each string is a frame line from the forward pass)
-  // Each frame string already includes newlines, so we just join them directly
-  let frames_str = forward_frames.join('');
-  // Ensure we don't have a trailing newline that could cause display issues
-  frames_str = frames_str.trimEnd();
-  return `\n\n=== Forward Pass Stack Trace (where this tensor was created) ===\n${frames_str}`;
-}
-
-function format_frames(frames) {
-  if (frames.length === 0) {
-    return (
-      `This block has no frames. Potential causes:\n` +
-      `1) This block was allocated before _record_memory_history was enabled.\n` +
-      `2) The context or stacks passed to _record_memory_history does not include this block. Consider changing context to 'state', 'alloc', or 'all', or changing stacks to 'all'.\n` +
-      `3) This event occurred during backward, which has no python frames, and memory history did not include C++ frames. Use stacks='all' to record both C++ and python frames.`
-    );
-  }
-  const frame_strings = frames
-    .filter(frameFilter)
-    .map(f => {
-      let frame_str = `${f.filename}:${f.line}:${f.name}`;
-
-      // Add FX debug information if available
-      if (f.fx_node_op || f.fx_node_name || f.fx_node_target) {
-        const fx_parts = [];
-        if (f.fx_node_name) fx_parts.push(`node=${f.fx_node_name}`);
-        if (f.fx_node_op) fx_parts.push(`op=${f.fx_node_op}`);
-        if (f.fx_node_target) fx_parts.push(`target=${f.fx_node_target}`);
-        frame_str += `\n    >> FX: ${fx_parts.join(', ')}`;
-      }
-
-      if (f.fx_original_trace) {
-        frame_str += `\n    >> Original Model Code:`;
-        const original_lines = f.fx_original_trace.trim().split('\n');
-        // Show all lines of the original trace
-        for (const line of original_lines) {
-          frame_str += `\n       ${line}`;
-        }
-      }
-
-      return frame_str;
-    });
-  return elideRepeats(frame_strings).join('\n');
-}
-
-function process_alloc_data(snapshot, device, plot_segments, max_entries) {
-  const elements = [];
-  const initially_allocated = [];
-  const actions = [];
-  const addr_to_alloc = {};
-
-  const alloc = plot_segments ? 'segment_alloc' : 'alloc';
-  const [free, free_completed] = plot_segments
-    ? ['segment_free', 'segment_free']
-    : ['free', 'free_completed'];
-  for (const e of snapshot.device_traces[device]) {
-    switch (e.action) {
-      case alloc:
-        elements.push(e);
-        addr_to_alloc[e.addr] = elements.length - 1;
-        actions.push(elements.length - 1);
-        break;
-      case free:
-      case free_completed:
-        if (e.addr in addr_to_alloc) {
-          actions.push(addr_to_alloc[e.addr]);
-          delete addr_to_alloc[e.addr];
-        } else {
-          elements.push(e);
-          initially_allocated.push(elements.length - 1);
-          actions.push(elements.length - 1);
-        }
-        break;
-      default:
-        break;
-    }
-  }
-  for (const seg of snapshot.segments) {
-    if (seg.device !== device) {
-      continue;
-    }
-    if (plot_segments) {
-      if (!(seg.address in addr_to_alloc)) {
-        const element = {
-          action: 'alloc',
-          addr: seg.address,
-          size: seg.total_size,
-          frames: [],
-          stream: seg.stream,
-          version: seg.version,
-        };
-        elements.push(element);
-        initially_allocated.push(elements.length - 1);
-      }
-    } else {
-      for (const b of seg.blocks) {
-        if (b.state === 'active_allocated' && !(b.addr in addr_to_alloc)) {
-          const element = {
-            action: 'alloc',
-            addr: b.addr,
-            size: b.requested_size,
-            frames: b.frames,
-            stream: seg.stream,
-            version: b.version,
-          };
-          elements.push(element);
-          initially_allocated.push(elements.length - 1);
-        }
-      }
-    }
-  }
-  initially_allocated.reverse();
-  // if there are no actions, the graph will be blank,
-  // but if there are existing allocations we do not want to hide them
-  // by having just one allocate action it will show a flat graph with all segments
-  if (actions.length === 0 && initially_allocated.length > 0) {
-    actions.push(initially_allocated.pop());
-  }
-
-  const current = [];
-  const current_data = [];
-  const data = [];
-  let max_size = 0;
-
-  let total_mem = 0;
-  let total_summarized_mem = 0;
-  let timestep = 0;
-
-  const max_at_time = [];
-
-  const summarized_mem = {
-    elem: 'summarized',
-    timesteps: [],
-    offsets: [total_mem],
-    size: [],
-    color: 0,
-  };
-  const summarized_elems = {};
-
-  function advance(n) {
-    summarized_mem.timesteps.push(timestep);
-    summarized_mem.offsets.push(total_mem);
-    summarized_mem.size.push(total_summarized_mem);
-    timestep += n;
-    for (let i = 0; i < n; i++) {
-      max_at_time.push(total_mem + total_summarized_mem);
-    }
-  }
-
-  const sizes = elements
-    .map((x, i) => [x.size, i])
-    .sort(([x, _xi], [y, _yi]) => y - x);
-
-  const draw_elem = {};
-  for (const [_s, e] of sizes.slice(0, max_entries)) {
-    draw_elem[e] = true;
-  }
-
-  function add_allocation(elem) {
-    const element_obj = elements[elem];
-    const size = element_obj.size;
-    current.push(elem);
-    let color = elem;
-    if (snapshot.categories.length > 0) {
-      color = snapshot.categories.indexOf(element_obj.category || 'unknown');
-    }
-    const e = {
-      elem,
-      timesteps: [timestep],
-      offsets: [total_mem],
-      size,
-      color,
-    };
-    current_data.push(e);
-    data.push(e);
-    total_mem += size;
-    element_obj.max_allocated_mem = total_mem + total_summarized_mem;
-  }
-
-  for (const elem of initially_allocated) {
-    if (elem in draw_elem) {
-      add_allocation(elem);
-    } else {
-      total_summarized_mem += elements[elem].size;
-      summarized_elems[elem] = true;
-    }
-  }
-
-  for (const elem of actions) {
-    const size = elements[elem].size;
-    if (!(elem in draw_elem)) {
-      if (elem in summarized_elems) {
-        advance(1);
-        total_summarized_mem -= size;
-        summarized_elems[elem] = null;
-      } else {
-        total_summarized_mem += size;
-        summarized_elems[elem] = true;
-        advance(1);
-      }
-      continue;
-    }
-    const idx = current.findLastIndex(x => x === elem);
-    // first time we see an action we add it
-    // second time we remove it
-    if (idx === -1) {
-      add_allocation(elem);
-      advance(1);
-    } else {
-      advance(1);
-      const removed = current_data[idx];
-      removed.timesteps.push(timestep);
-      removed.offsets.push(removed.offsets.at(-1));
-      current.splice(idx, 1);
-      current_data.splice(idx, 1);
-
-      if (idx < current.length) {
-        for (let j = idx; j < current.length; j++) {
-          const e = current_data[j];
-          e.timesteps.push(timestep);
-          e.offsets.push(e.offsets.at(-1));
-          e.timesteps.push(timestep + 3);
-          e.offsets.push(e.offsets.at(-1) - size);
-        }
-        advance(3);
-      }
-      total_mem -= size;
-    }
-    max_size = Math.max(total_mem + total_summarized_mem, max_size);
-  }
-
-  for (const elem of current_data) {
-    elem.timesteps.push(timestep);
-    elem.offsets.push(elem.offsets.at(-1));
-  }
-  data.push(summarized_mem);
-
-  return {
-    max_size,
-    allocations_over_time: data,
-    max_at_time,
-    summarized_mem,
-    elements_length: elements.length,
-    context_for_id: id => {
-      const elem = elements[id];
-      let text = `Addr: ${formatAddr(elem)}`;
-      text = `${text}, Size: ${formatSize(elem.size)} allocation`;
-      text = `${text}, Total memory used after allocation: ${formatSize(
-        elem.max_allocated_mem,
-      )}`;
-      const context = elem?.compile_context ?? 'None';
-      text = `${text}, Compile context: ${context}`;
-      if (elem.stream !== null) {
-        text = `${text}, stream ${elem.stream}`;
-      }
-      if (elem.timestamp !== null) {
-        var d = new Date(elem.time_us / 1000);
-        text = `${text}, timestamp ${d}`;
-      }
-      if (!elem.action.includes('alloc')) {
-        text = `${text}\nalloc not recorded, stack trace for free:`;
-      }
-      const user_metadata_str = format_user_metadata(elem.user_metadata);
-      if (user_metadata_str) {
-        text = `${text}\n${user_metadata_str}`;
-      }
-      text = `${text}\n${format_frames(elem.frames)}`;
-      text = `${text}${format_forward_frames(elem.forward_frames)}`;
-      return text;
-    },
-  };
-}
-
 function MemoryPlot(
   svg,
   data,
@@ -1213,7 +907,8 @@ function MemoryPlot(
     .enter()
     .append('polygon')
     .attr('points', format_points)
-    .attr('fill', d => colors[d.color % colors.length]);
+    .attr('fill', d => colors[d.color % colors.length])
+    .attr('opacity', d => d.opacity ?? 1);
 
   const axis = plot_coordinate_space.append('g').call(yaxis);
 
@@ -1272,11 +967,34 @@ function MemoryPlot(
 function ContextViewer(text, data) {
   let current_selected = null;
 
+  function restore_search_highlight(d) {
+    if (!d) return;
+    const addr = d.attr('data-search-match') === 'true';
+    const frame = d.attr('data-frame-match') === 'true';
+    if (addr && frame) {
+      d.attr('stroke', '#ff00ff')
+        .attr('stroke-width', 3)
+        .attr('stroke-dasharray', '6,3')
+        .attr('vector-effect', 'non-scaling-stroke');
+    } else if (addr) {
+      d.attr('stroke', 'red')
+        .attr('stroke-width', 2)
+        .attr('stroke-dasharray', null)
+        .attr('vector-effect', 'non-scaling-stroke');
+    } else if (frame) {
+      d.attr('stroke', '#2196F3')
+        .attr('stroke-width', 2)
+        .attr('stroke-dasharray', null)
+        .attr('vector-effect', 'non-scaling-stroke');
+    }
+  }
+
   return {
     default_selected: null,
     set_selected: d => {
       if (current_selected !== null) {
-        current_selected.attr('stroke', null).attr('stroke-width', null);
+        current_selected.attr('stroke', null).attr('stroke-width', null).attr('stroke-dasharray', null);
+        restore_search_highlight(current_selected);
       }
       if (d === null) {
         text.text('');
@@ -1287,6 +1005,10 @@ function ContextViewer(text, data) {
             'Small tensors that were not plotted to cutdown on render time.\n' +
               'Use detail slider to see smaller allocations.',
           );
+        } else if (typeof dd.elem === 'string' && dd.elem.startsWith('pool:')) {
+          const pool_key = dd.elem.slice(5);
+          const capacity = Array.isArray(dd.size) ? dd.size.at(-1) : dd.size;
+          text.text(`Private Pool (${pool_key}): capacity ${formatSize(capacity)}`);
         } else {
           text.text(`${dd.elem} ${data.context_for_id(dd.elem)}`);
         }
@@ -1387,13 +1109,20 @@ function create_trace_view(
   device,
   plot_segments = false,
   max_entries = 15000,
+  include_private_inactive = false,
 ) {
   const left_pad = 70;
-  const data = process_alloc_data(snapshot, device, plot_segments, max_entries);
+  const data = process_alloc_data(snapshot, device, plot_segments, max_entries, include_private_inactive);
   dst.selectAll('svg').remove();
   dst.selectAll('div').remove();
 
   max_entries = Math.min(max_entries, data.elements_length);
+  if (include_private_inactive) {
+    dst.append('div')
+      .attr('style', 'padding: 4px 8px; background: #fff3cd; border: 1px solid #ffc107; font-size: 13px; margin-bottom: 4px;')
+      .text('Note: Private pool memory (the gray bar) is shown as allocated until the pool\'s segment is freed. '
+          + 'This view requires that MemPools are not deleted before torch.cuda.memory._snapshot() is called.');
+  }
   const d = dst.append('div');
   d.append('input')
     .attr('type', 'range')
@@ -1401,12 +1130,28 @@ function create_trace_view(
     .attr('max', data.elements_length)
     .attr('value', max_entries)
     .on('change', function () {
-      create_trace_view(dst, snapshot, device, plot_segments, this.value);
+      create_trace_view(dst, snapshot, device, plot_segments, this.value, include_private_inactive);
     });
   d.append('label').text(
     `Detail: ${max_entries} of ${data.elements_length} entries`,
   );
 
+  d.append('span').text('  |  ');
+  const search_input = d.append('input')
+    .attr('type', 'text')
+    .attr('placeholder', 'Search address (hex)...')
+    .attr('style', 'width: 180px; margin-left: 4px; font-family: monospace;');
+  const search_label = d.append('label')
+    .attr('style', 'margin-left: 4px;');
+
+  d.append('span').text('  |  ');
+  const frame_input = d.append('input')
+    .attr('type', 'text')
+    .attr('placeholder', 'Search stack frame...')
+    .attr('style', 'width: 200px; margin-left: 4px; font-family: monospace;');
+  const frame_label = d.append('label')
+    .attr('style', 'margin-left: 4px;');
+
   const grid_container = dst
     .append('div')
     .attr(
@@ -1443,6 +1188,61 @@ function create_trace_view(
     );
   const delegate = ContextViewer(context_div.append('pre').text('none'), data);
   plot.set_delegate(delegate);
+
+  function apply_search_highlights() {
+    const addr_query = search_input.node().value.toLowerCase().trim();
+    const frame_query = frame_input.node().value.toLowerCase().trim();
+    const polygons = plot_svg.selectAll('polygon');
+    let addr_matches = 0;
+    let frame_matches = 0;
+    polygons.each(function () {
+      const dd = d3.select(this).datum();
+      if (!dd || typeof dd.elem !== 'number') {
+        d3.select(this)
+          .attr('data-search-match', null)
+          .attr('data-frame-match', null);
+        return;
+      }
+      const ctx = data.context_for_id(dd.elem);
+      const ctx_lower = ctx.toLowerCase();
+      const addr_hit = addr_query && ctx_lower.includes(addr_query);
+      const frame_hit = frame_query && ctx_lower.includes(frame_query);
+      d3.select(this)
+        .attr('data-search-match', addr_hit ? 'true' : null)
+        .attr('data-frame-match', frame_hit ? 'true' : null);
+      if (addr_hit && frame_hit) {
+        d3.select(this)
+          .attr('stroke', '#ff00ff')
+          .attr('stroke-width', 3)
+          .attr('stroke-dasharray', '6,3')
+          .attr('vector-effect', 'non-scaling-stroke');
+      } else if (addr_hit) {
+        d3.select(this)
+          .attr('stroke', 'red')
+          .attr('stroke-width', 2)
+          .attr('stroke-dasharray', null)
+          .attr('vector-effect', 'non-scaling-stroke');
+      } else if (frame_hit) {
+        d3.select(this)
+          .attr('stroke', '#2196F3')
+          .attr('stroke-width', 2)
+          .attr('stroke-dasharray', null)
+          .attr('vector-effect', 'non-scaling-stroke');
+      } else {
+        d3.select(this)
+          .attr('stroke', null)
+          .attr('stroke-width', null)
+          .attr('stroke-dasharray', null);
+      }
+      if (addr_hit) addr_matches++;
+      if (frame_hit) frame_matches++;
+    });
+    search_label.text(addr_query ? `${addr_matches} match${addr_matches !== 1 ? 'es' : ''}` : '');
+    frame_label.text(frame_query ? `${frame_matches} match${frame_matches !== 1 ? 'es' : ''}` : '');
+  }
+
+  search_input.on('input', apply_search_highlights);
+  frame_input.on('input', apply_search_highlights);
 }
 
 function create_settings_view(dst, snapshot, device) {
@@ -1764,6 +1564,8 @@ function decode_base64(input) {
 
 const kinds = {
   'Active Memory Timeline': create_trace_view,
+  'Allocated Memory (incl. Private Pools)': (dst, snapshot, device) =>
+    create_trace_view(dst, snapshot, device, false, 15000, true),
   'Allocator State History': create_segment_view,
   'Active Cached Segment Timeline': (dst, snapshot, device) =>
     create_trace_view(dst, snapshot, device, true),
diff --git a/torch/utils/viz/_cycles.py b/torch/utils/viz/_cycles.py
index df4bf34db2114..2f68cf75dee4d 100644
--- a/torch/utils/viz/_cycles.py
+++ b/torch/utils/viz/_cycles.py
@@ -467,7 +467,6 @@ def to_html(nodes):
         if n.context is None:
             continue
         s = _listener_template.format(id=str(i + 1), stack=escape(f'{n.label}:\n{n.context}'))
-        # pyrefly: ignore [bad-argument-type]
         listeners.append(s)
     dot = to_dot(nodes)
     return _template.replace('$DOT', repr(dot)).replace('$LISTENERS', '\n'.join(listeners))
diff --git a/torch/utils/viz/process_alloc_data.js b/torch/utils/viz/process_alloc_data.js
new file mode 100644
index 0000000000000..6e5feda1046ea
--- /dev/null
+++ b/torch/utils/viz/process_alloc_data.js
@@ -0,0 +1,902 @@
+// Pure data-processing functions for PyTorch memory visualization.
+// Extracted from MemoryViz.js so they can be tested independently (no d3/DOM deps).
+//
+// This file is the single source of truth for these functions:
+//   - MemoryViz.js imports them via ESM: import {...} from "./process_alloc_data.js"
+//   - Node.js tests load this file by stripping the export line and eval-ing
+//
+// TRACE EVENT ACTIONS (from c10/core/CachingDeviceAllocator.h TraceEntry::Action):
+//
+//   "alloc"          - Sub-allocation returned to user from the caching allocator.
+//                      Recorded in alloc_found_block() (CUDACachingAllocator.cpp:1834).
+//
+//   "free_requested" - User code called free (tensor out of scope). The block may not
+//                      be immediately returned to the free pool if it's in use on
+//                      another stream via record_stream.
+//                      Recorded in free() (CUDACachingAllocator.cpp:2123).
+//
+//   "free_completed" - Block actually returned to the allocator's free pool. For simple
+//                      cases this fires immediately after free_requested. For cross-stream
+//                      blocks, deferred until CUDA events confirm all streams are done.
+//                      Recorded in free_block() (CUDACachingAllocator.cpp:3148).
+//
+//   "segment_alloc"  - New segment allocated from OS via cudaMalloc (or cuMemCreate for
+//                      expandable segments).
+//                      Recorded in alloc_from_expandable_segment() (CUDACachingAllocator.cpp:3548).
+//
+//   "segment_free"   - Segment returned to OS via cudaFree. Happens during empty_cache()
+//                      or defragmentation. Only for non-expandable segments.
+//                      Recorded in release_block() (CUDACachingAllocator.cpp:3686).
+//
+//   "segment_map"    - Physical pages mapped into an expandable segment via cuMemMap.
+//                      The segment grows. Only with expandable segments enabled.
+//                      Recorded in alloc_from_expandable_segment() (CUDACachingAllocator.cpp:3092).
+//
+//   "segment_unmap"  - Physical pages unmapped from an expandable segment via cuMemUnmap.
+//                      Virtual address range retained, physical memory returned to OS.
+//                      Only with expandable segments. Causes "pool_id unknown" for any
+//                      trace events whose addresses fall in the unmapped range, since
+//                      the segment no longer exists at snapshot time.
+//                      Recorded in unmap_block() (CUDACachingAllocator.cpp:3790).
+//
+//   "snapshot"       - A call to torch.cuda.memory._snapshot(). Timestamp marker to
+//                      correlate trace events with snapshot state. addr=0.
+//                      Recorded in snapshot() (CUDACachingAllocator.cpp:2689).
+//
+//   "oom"            - Allocator failed to satisfy an allocation after all retries.
+//                      addr=device_free (bytes free on GPU), size=requested allocation.
+//                      Recorded in malloc() (CUDACachingAllocator.cpp:1629).
+//
+// HOW SEGMENT EVENTS ARE USED IN VISUALIZATION:
+//
+//   The snapshot pickle contains two separate data sources:
+//     1. device_traces  - Ring buffer of TraceEntry actions (alloc, free, segment_map, etc.)
+//     2. segments       - Point-in-time dump of all segments/blocks at _snapshot() time
+//
+//   Block-level views ("Active Memory Timeline", "Allocated Memory (incl. Private Pools)"):
+//     - process_alloc_data matches "alloc" and "free_completed" from device_traces.
+//     - segment_alloc/segment_free/segment_map/segment_unmap are ignored (skipped in switch).
+//     - The segments snapshot is used only to resolve pool_id via find_pool_id().
+//
+//   Segment-level view ("Active Cached Segment Timeline"):
+//     - process_alloc_data is called with plot_segments=true.
+//     - Matches "segment_alloc" and "segment_free" instead of alloc/free.
+//     - segment_map/segment_unmap are NOT matched (they don't appear in the switch).
+//     - Segments from the snapshot that weren't seen in the trace are added as
+//       initially_allocated (Phase 2).
+//
+//   Allocator State History ("Allocator State History"):
+//     - EventSelector lists ALL trace events including segment_map/segment_unmap.
+//     - MemoryView renders the segment/block layout from the segments snapshot.
+//     - Clicking an event in the list redraws the layout at that point in time.
+//
+//   Ring buffer overflow:
+//     - All trace event types share the same ring buffer. When it overflows, older
+//       events are overwritten. The allocator_settings.trace_alloc_overflowed flag
+//       indicates this happened, and trace_alloc_max_entries gives the buffer size.
+//     - Segment snapshot data (segments array) is NOT affected by ring buffer overflow.
+//     - The segment snapshot is always complete regardless of overflow.
+
+/**
+ * Returns true if pool_id represents a private (user-created) memory pool,
+ * as opposed to the default pool [0, 0].
+ *
+ * @param {number[]|null} pool_id - Two-element array [owner_id, pool_id] from
+ *   the CUDA caching allocator. The default pool is [0, 0]; any other non-null
+ *   value is a private pool (e.g. FSDP's MemPool).
+ * @returns {boolean}
+ */
+function isPrivatePoolId(pool_id) {
+  return pool_id && !(pool_id[0] === 0 && pool_id[1] === 0);
+}
+
+/**
+ * Formats a byte count as a human-readable string (e.g. "1.5GiB (1610612736 bytes)").
+ *
+ * @param {number} num - Size in bytes.
+ * @param {boolean} [showBytes=true] - Whether to include the raw byte count in parentheses.
+ * @returns {string}
+ */
+function formatSize(num, showBytes = true) {
+  const orig = num;
+  // https://stackoverflow.com/questions/1094841/get-human-readable-version-of-file-size
+  const units = ['', 'Ki', 'Mi', 'Gi', 'Ti', 'Pi', 'Ei', 'Zi'];
+  for (const unit of units) {
+    if (Math.abs(num) < 1024.0) {
+      if (showBytes) {
+        return `${num.toFixed(1)}${unit}B (${orig} bytes)`;
+      }
+      return `${num.toFixed(1)}${unit}B`;
+    }
+    num /= 1024.0;
+  }
+  return `${num.toFixed(1)}YiB`;
+}
+
+/**
+ * Formats a trace event's address as a display string like "b'7f4c00000_3".
+ * Segment-level events get an "s'" prefix, block-level events get "b'".
+ *
+ * @param {{action: string, addr: number|BigInt, version: number}} event
+ * @returns {string}
+ */
+function formatAddr(event) {
+  const prefix = event.action.startsWith('segment') ? 's\'' : 'b\'';
+  return `${prefix}${event.addr.toString(16)}_${event.version}`;
+}
+
+/**
+ * Collapses consecutive duplicate strings in an array. If a string appears
+ * N > 2 times in a row, it's replaced with [str, "<repeats N-1 times>"].
+ * Used to compress repetitive stack frames in the display.
+ *
+ * @param {string[]} frames
+ * @returns {string[]}
+ */
+function elideRepeats(frames) {
+  const result = [];
+  const length = frames.length;
+  for (let i = 0; i < length; ) {
+    let j = i + 1;
+    const f = frames[i];
+    while (j < length && f === frames[j]) {
+      j++;
+    }
+    switch (j - i) {
+      case 1:
+        result.push(f);
+        break;
+      case 2:
+        result.push(f, f);
+        break;
+      default:
+        result.push(f, `<repeats ${j - i - 1} times>`);
+        break;
+    }
+    i = j;
+  }
+  return result;
+}
+
+/**
+ * Returns false for stack frames that are internal runtime noise
+ * (e.g. Python interpreter internals, C++ dispatch machinery).
+ * Used as a filter predicate on frame arrays.
+ *
+ * @param {{name: string, filename: string}} frame
+ * @returns {boolean}
+ */
+function frameFilter({name, filename}) {
+  const omitFunctions = [
+    'unwind::unwind',
+    'CapturedTraceback::gather',
+    'gather_with_cpp',
+    '_start',
+    '__libc_start_main',
+    'PyEval_',
+    'PyObject_',
+    'PyFunction_',
+  ];
+
+  const omitFilenames = [
+    'core/boxing',
+    '/Register',
+    '/Redispatch',
+    'pythonrun.c',
+    'Modules/main.c',
+    'Objects/call.c',
+    'Objects/methodobject.c',
+    'pycore_ceval.h',
+    'ceval.c',
+    'cpython/abstract.h',
+  ];
+
+  for (const of of omitFunctions) {
+    if (name.includes(of)) {
+      return false;
+    }
+  }
+
+  for (const of of omitFilenames) {
+    if (filename.includes(of)) {
+      return false;
+    }
+  }
+
+  return true;
+}
+
+/**
+ * Formats user-attached metadata (from torch.cuda.memory._record_memory_history)
+ * as a display string. Returns '' if no metadata is present.
+ *
+ * @param {string|Object|null|undefined} user_metadata
+ * @returns {string}
+ */
+function format_user_metadata(user_metadata) {
+  if (!user_metadata) {
+    return '';
+  }
+  if (typeof user_metadata === 'string') {
+    return `User Metadata:\n  ${user_metadata}`;
+  }
+  if (typeof user_metadata === 'object' && Object.keys(user_metadata).length === 0) {
+    return '';
+  }
+  const metadata_lines = Object.entries(user_metadata)
+    .map(([key, value]) => `  ${key}: ${value}`);
+  return 'User Metadata:\n' + metadata_lines.join('\n');
+}
+
+/**
+ * Formats the forward-pass stack trace (captured via torch.autograd) as a
+ * display string showing where a tensor was originally created.
+ *
+ * @param {string[]|null|undefined} forward_frames
+ * @returns {string}
+ */
+function format_forward_frames(forward_frames) {
+  if (!forward_frames || forward_frames.length === 0) {
+    return '';
+  }
+  let frames_str = forward_frames.join('');
+  frames_str = frames_str.trimEnd();
+  return `\n\n=== Forward Pass Stack Trace (where this tensor was created) ===\n${frames_str}`;
+}
+
+/**
+ * Formats an array of stack frames into a human-readable string.
+ * Filters out runtime noise via frameFilter, annotates FX graph debug info
+ * when available, and collapses consecutive duplicate frames.
+ *
+ * @param {{filename: string, line: number, name: string,
+ *          fx_node_op?: string, fx_node_name?: string,
+ *          fx_node_target?: string, fx_original_trace?: string}[]} frames
+ * @returns {string}
+ */
+function format_frames(frames) {
+  if (frames.length === 0) {
+    return (
+      `This block has no frames. Potential causes:\n` +
+      `1) This block was allocated before _record_memory_history was enabled.\n` +
+      `2) The context or stacks passed to _record_memory_history does not include this block. Consider changing context to 'state', 'alloc', or 'all', or changing stacks to 'all'.\n` +
+      `3) This event occurred during backward, which has no python frames, and memory history did not include C++ frames. Use stacks='all' to record both C++ and python frames.\n` +
+      `4) This block was reconstructed from the allocator's segment snapshot (not from a trace event). The snapshot records which blocks exist at the moment _snapshot() is called, but does not carry stack frames. This typically happens for blocks that were allocated before tracing started and never freed, or for inactive blocks in private memory pools.\n` +
+      `5) The original alloc event was evicted from the trace ring buffer (older entries are overwritten when the buffer is full). Increase the max_entries argument to _record_memory_history to retain more events.`
+    );
+  }
+  const frame_strings = frames
+    .filter(frameFilter)
+    .map(f => {
+      let frame_str = `${f.filename}:${f.line}:${f.name}`;
+
+      if (f.fx_node_op || f.fx_node_name || f.fx_node_target) {
+        const fx_parts = [];
+        if (f.fx_node_name) fx_parts.push(`node=${f.fx_node_name}`);
+        if (f.fx_node_op) fx_parts.push(`op=${f.fx_node_op}`);
+        if (f.fx_node_target) fx_parts.push(`target=${f.fx_node_target}`);
+        frame_str += `\n    >> FX: ${fx_parts.join(', ')}`;
+      }
+
+      if (f.fx_original_trace) {
+        frame_str += `\n    >> Original Model Code:`;
+        const original_lines = f.fx_original_trace.trim().split('\n');
+        for (const line of original_lines) {
+          frame_str += `\n       ${line}`;
+        }
+      }
+
+      return frame_str;
+    });
+  return elideRepeats(frame_strings).join('\n');
+}
+
+/**
+ * Transforms a memory snapshot into a stacked-area timeline suitable for
+ * rendering by MemoryPlot. This is the core data-processing function behind
+ * the "Active Memory Timeline" and "Allocated Memory (incl. Private Pools)"
+ * visualization tabs.
+ *
+ * HIGH-LEVEL ALGORITHM:
+ *
+ * 1. TRACE EVENT MATCHING: Scans device_traces to pair alloc events with their
+ *    corresponding free_completed events (by address). Events whose matching
+ *    alloc was lost (e.g. ring buffer wrap) become "initially_allocated" —
+ *    blocks assumed to exist at the start of the trace.
+ *
+ * 2. SEGMENT SNAPSHOT: Supplements trace data with the current segment state.
+ *    Blocks marked active_allocated (or inactive in private pools, when
+ *    include_private_inactive=true) that weren't seen in the trace are also
+ *    added as initially_allocated.
+ *
+ * 3. DETAIL LIMITING: Only the largest max_entries elements get individual
+ *    rectangles in the plot. Smaller elements are aggregated into a single
+ *    "summarized" band to keep rendering fast.
+ *
+ * 4. STACKED AREA CONSTRUCTION: Replays alloc/free events in order, building
+ *    a stacked-area dataset where each element has timesteps, y-offsets, and
+ *    a size. Elements are stacked bottom-to-top; frees remove from the stack
+ *    and shift elements above downward.
+ *
+ * 5. PRIVATE POOL ENVELOPES (include_private_inactive=true): Each private pool
+ *    (e.g. FSDP's MemPool) gets a single gray "envelope" rectangle whose
+ *    height is the pool's high-water mark. Active blocks within the pool are
+ *    rendered as colored stripes inside the envelope. The envelope only grows
+ *    (never shrinks), representing reserved capacity.
+ *
+ *    Initially-allocated private pool blocks are PRE-LOADED into pool state
+ *    so that when their free event appears in the trace, they are correctly
+ *    recognized as frees (not misinterpreted as new allocations).
+ *
+ * NOTE ON FREE EVENT MATCHING: The C++ allocator emits 'free_requested' and
+ * 'free_completed' for each deallocation. This function matches against 'free'
+ * (which no longer appears in modern traces — effectively dead code) and
+ * 'free_completed'. Only free_completed does the actual matching. Matching
+ * both 'free_requested' AND 'free_completed' would cause double-processing
+ * since they share the same address.
+ *
+ * @param {Object} snapshot - Memory snapshot from torch.cuda.memory._snapshot().
+ * @param {Object[]} snapshot.segments - Current allocator segment state.
+ * @param {Object[][]} snapshot.device_traces - Per-device arrays of trace events.
+ *   Each event has {action, addr, size, frames, stream, segment_pool_id?, ...}.
+ * @param {string[]} snapshot.categories - Category names for color-coding.
+ * @param {number} device - Device index into snapshot.device_traces.
+ * @param {boolean} plot_segments - If true, plot segment-level (cudaMalloc)
+ *   events instead of sub-allocation events.
+ * @param {number} max_entries - Maximum number of elements to render individually.
+ *   Elements beyond this limit are aggregated into the "summarized" band.
+ * @param {boolean} [include_private_inactive=false] - If true, include inactive
+ *   blocks from private pools and render pool envelopes. Used by the
+ *   "Allocated Memory (incl. Private Pools)" tab.
+ *
+ * @returns {{
+ *   max_size: number,
+ *   allocations_over_time: Object[],
+ *   max_at_time: number[],
+ *   summarized_mem: Object,
+ *   elements_length: number,
+ *   context_for_id: function(number): string
+ * }}
+ *   - max_size: peak total memory observed during the action replay (used for
+ *     y-axis scaling). Note: this is only updated inside the action loop, so
+ *     the initial state from initially_allocated may not be reflected here
+ *     (use max_at_time for the true peak).
+ *   - allocations_over_time: array of stacked-area data objects, each with
+ *     {elem, timesteps[], offsets[], size, color}.
+ *   - max_at_time: total memory at each timestep (for minimap rendering).
+ *   - summarized_mem: the aggregated band for small elements.
+ *   - elements_length: total number of unique allocation elements.
+ *   - context_for_id: function that returns a human-readable description
+ *     string for a given element index (address, size, stack trace, etc.).
+ */
+function process_alloc_data(snapshot, device, plot_segments, max_entries, include_private_inactive = false) {
+  const elements = [];
+  // Contains two types of blocks
+  // 1. free without alloc in trace
+  // 2. actively allocated in segments, but no matching alloc in trace
+  const initially_allocated = [];
+  const actions = [];
+  const addr_to_alloc = {};
+
+  const device_segments = snapshot.segments
+    .filter(s => s.device === device)
+    .sort((a, b) => {
+      if (a.address === b.address) return 0;
+      return a.address < b.address ? -1 : 1;
+    });
+
+  // Binary search to find which segment contains a given address.
+  function find_pool_id(addr) {
+    let left = 0;
+    let right = device_segments.length - 1;
+    while (left <= right) {
+      const mid = Math.floor((left + right) / 2);
+      const seg = device_segments[mid];
+      const seg_end = seg.address + (typeof seg.address === "bigint" ? BigInt(seg.total_size) : seg.total_size);
+      if (addr < seg.address) {
+        right = mid - 1;
+      } else if (addr >= seg_end) {
+        left = mid + 1;
+      } else {
+        return seg.segment_pool_id;
+      }
+    }
+    return null;
+  }
+
+  const alloc = plot_segments ? 'segment_alloc' : 'alloc';
+  const [free, free_completed] = plot_segments
+    ? ['segment_free', 'segment_free']
+    : ['free', 'free_completed'];
+  for (const e of snapshot.device_traces[device]) {
+    switch (e.action) {
+      case alloc:
+        elements.push(e);
+        addr_to_alloc[e.addr] = elements.length - 1;
+        actions.push(elements.length - 1);
+        break;
+      case free:
+      case free_completed:
+        if (e.addr in addr_to_alloc) {
+          // Matched: reuse the element from the alloc event
+          actions.push(addr_to_alloc[e.addr]);
+          delete addr_to_alloc[e.addr];
+        } else {
+          // Unmatched free: alloc happened before recording (or was evicted
+          // from the ring buffer). Create a new element from the free event;
+          // its stack trace will show the free site, not the alloc site.
+          elements.push(e);
+          initially_allocated.push(elements.length - 1);
+          actions.push(elements.length - 1);
+        }
+        break;
+      default:
+        break;
+    }
+  }
+
+  // --- Phase 2: Add elements from the snapshot ---
+  for (const seg of snapshot.segments) {
+    if (seg.device !== device) {
+      continue;
+    }
+    if (plot_segments) {
+      if (!(seg.address in addr_to_alloc)) {
+        const element = {
+          action: 'alloc',
+          addr: seg.address,
+          size: seg.total_size,
+          frames: [],
+          stream: seg.stream,
+          version: seg.version,
+        };
+        elements.push(element);
+        initially_allocated.push(elements.length - 1);
+      }
+    } else {
+      for (const b of seg.blocks) {
+        const addr = b.addr ?? b.address;
+        if (b.state === 'active_allocated' && !(addr in addr_to_alloc)) {
+          const element = {
+            action: 'alloc',
+            addr,
+            size: b.requested_size,
+            frames: b.frames,
+            stream: seg.stream,
+            version: b.version,
+            segment_pool_id: seg.segment_pool_id,
+            ghost: true,
+          };
+          elements.push(element);
+          initially_allocated.push(elements.length - 1);
+        }
+      }
+    }
+  }
+
+  // Resolve pool IDs for trace elements by looking up which segment they fall in
+  for (const elem of elements) {
+    if (!elem.segment_pool_id) {
+      elem.segment_pool_id = find_pool_id(elem.addr);
+    }
+  }
+
+  initially_allocated.reverse();
+  // If there are no trace actions but there are existing allocations,
+  // show a flat graph with the initial state
+  if (actions.length === 0 && initially_allocated.length > 0) {
+    actions.push(initially_allocated.pop());
+  }
+
+  // --- Phase 3: Build the stacked-area timeline ---
+  const current = [];      // stack of element indices (bottom to top)
+  const current_data = []; // parallel array of visualization data objects
+  const data = [];         // all data objects (including completed ones)
+  let max_size = 0;
+
+  let total_mem = 0;
+  let total_summarized_mem = 0;
+  let timestep = 0;
+
+  const max_at_time = [];
+
+  const summarized_mem = {
+    elem: 'summarized',
+    timesteps: [],
+    offsets: [total_mem],
+    size: [],
+    color: 0,
+  };
+  const summarized_elems = {};
+
+  // Record the current memory state and advance time by n steps
+  function advance(n) {
+    summarized_mem.timesteps.push(timestep);
+    summarized_mem.offsets.push(total_mem);
+    summarized_mem.size.push(total_summarized_mem);
+    timestep += n;
+    for (let i = 0; i < n; i++) {
+      max_at_time.push(total_mem + total_summarized_mem);
+    }
+  }
+
+  // Only render the largest max_entries elements individually;
+  // everything else goes into the summarized band
+  const sizes = elements
+    .map((x, i) => [x.size, i])
+    .sort(([x, _xi], [y, _yi]) => y - x);
+
+  const draw_elem = {};
+  for (const [_s, e] of sizes.slice(0, max_entries)) {
+    draw_elem[e] = true;
+  }
+
+  // Push an element onto the memory stack
+  function add_allocation(elem) {
+    const element_obj = elements[elem];
+    const size = element_obj.size;
+    current.push(elem);
+    let color = elem;
+    if (snapshot.categories.length > 0) {
+      color = snapshot.categories.indexOf(element_obj.category || 'unknown');
+    }
+    const e = {
+      elem,
+      timesteps: [timestep],
+      offsets: [total_mem],
+      size,
+      color,
+    };
+    if (element_obj.ghost) e.ghost = true;
+    current_data.push(e);
+    data.push(e);
+    total_mem += size;
+    element_obj.max_allocated_mem = total_mem + total_summarized_mem;
+  }
+
+  // --- Pool envelope tracking (only when include_private_inactive=true) ---
+  // Each private pool gets a gray envelope whose height = high-water mark.
+  // Active blocks are rendered as colored stripes within the envelope.
+  const pools = {};
+  const pool_active_elems = {};
+
+  function get_pool_key(elem_idx) {
+    const pid = elements[elem_idx].segment_pool_id;
+    if (!isPrivatePoolId(pid)) return null;
+    const stream = elements[elem_idx].stream;
+    return `${pid[0]},${pid[1]},s${stream}`;
+  }
+
+  function get_or_create_pool(pool_key) {
+    if (!(pool_key in pools)) {
+      pools[pool_key] = {
+        max: 0, active: 0, envelope_data: null,
+        block_stack: [],  // [{elem, size, inner_offset, stripe_data}]
+      };
+    }
+    return pools[pool_key];
+  }
+
+  function elem_color(elem_idx) {
+    if (snapshot.categories.length > 0) {
+      return snapshot.categories.indexOf(elements[elem_idx].category || 'unknown');
+    }
+    return elem_idx;
+  }
+
+  function shift_pool_stripes(pool, delta) {
+    for (const block of pool.block_stack) {
+      const s = block.stripe_data;
+      s.timesteps.push(timestep);
+      s.offsets.push(s.offsets.at(-1));
+      s.timesteps.push(timestep + 3);
+      s.offsets.push(s.offsets.at(-1) + delta);
+    }
+  }
+
+  // Animate shifting all elements above idx by delta (used when an element
+  // is inserted or removed from the middle of the stack)
+  function shift_elements_above(idx, delta) {
+    for (let j = idx; j < current.length; j++) {
+      const e = current_data[j];
+      e.timesteps.push(timestep);
+      e.offsets.push(e.offsets.at(-1));
+      e.timesteps.push(timestep + 3);
+      e.offsets.push(e.offsets.at(-1) + delta);
+      if (Array.isArray(e.size)) {
+        e.size.push(e.size.at(-1));
+        e.size.push(e.size.at(-1));
+      }
+      const pk = typeof current[j] === 'string' && current[j].startsWith('pool:')
+        ? current[j].slice(5) : null;
+      if (pk && pk in pools) {
+        shift_pool_stripes(pools[pk], delta);
+      }
+    }
+  }
+
+  // Grow a pool envelope to accommodate new_active bytes.
+  // The envelope only grows (never shrinks) — it represents reserved capacity.
+  function grow_pool_envelope(pool, pool_key, new_active) {
+    if (new_active <= pool.max) return;
+    const delta = new_active - pool.max;
+    pool.max = new_active;
+    const env = pool.envelope_data;
+    env.timesteps.push(timestep);
+    env.offsets.push(env.offsets.at(-1));
+    env.size.push(env.size.at(-1));
+    env.timesteps.push(timestep + 3);
+    env.offsets.push(env.offsets.at(-1));
+    env.size.push(pool.max);
+    const pidx = current.indexOf(`pool:${pool_key}`);
+    if (pidx >= 0) {
+      shift_elements_above(pidx + 1, delta);
+    }
+    total_mem += delta;
+    advance(3);
+  }
+
+  // --- Process initially_allocated elements ---
+  // Private pool blocks are pre-loaded into pool state at timestep 0 (no
+  // animation) so that their envelope starts at the correct initial size
+  // and free events are correctly recognized as frees.
+  for (const elem of initially_allocated) {
+    if (include_private_inactive && get_pool_key(elem)) {
+      const pk = get_pool_key(elem);
+      const size = elements[elem].size;
+      const pool = get_or_create_pool(pk);
+      pool_active_elems[elem] = pk;
+
+      if (pool.envelope_data === null) {
+        const env = {
+          elem: `pool:${pk}`,
+          timesteps: [0],
+          offsets: [total_mem],
+          size: [0],
+          color: 9,
+        };
+        pool.envelope_data = env;
+        current.push(`pool:${pk}`);
+        current_data.push(env);
+        data.push(env);
+      }
+
+      const inner_offset = pool.active;
+      pool.active += size;
+
+      // Grow envelope directly without animation — these blocks pre-exist
+      if (pool.active > pool.max) {
+        const delta = pool.active - pool.max;
+        pool.max = pool.active;
+        const env = pool.envelope_data;
+        env.size[env.size.length - 1] = pool.max;
+        total_mem += delta;
+        const pidx = current.indexOf(`pool:${pk}`);
+        if (pidx >= 0) {
+          for (let j = pidx + 1; j < current.length; j++) {
+            const e = current_data[j];
+            e.offsets[e.offsets.length - 1] += delta;
+          }
+        }
+      }
+
+      const stripe = {
+        elem,
+        timesteps: [0],
+        offsets: [pool.envelope_data.offsets.at(-1) + inner_offset],
+        size,
+        color: elem_color(elem),
+        opacity: 0.5,
+        ghost: elements[elem].ghost || false,
+      };
+      pool.block_stack.push({elem, size, inner_offset, stripe_data: stripe});
+      data.push(stripe);
+      continue;
+    }
+    if (elem in draw_elem) {
+      add_allocation(elem);
+    } else {
+      total_summarized_mem += elements[elem].size;
+      summarized_elems[elem] = true;
+    }
+  }
+
+  // Fix up pool stripe offsets — stripes are not in current_data so they
+  // don't get shifted when other pools grow during initially_allocated
+  // processing. Recompute from the envelope's final offset.
+  for (const pk in pools) {
+    const p = pools[pk];
+    if (!p.envelope_data) continue;
+    const env_offset = p.envelope_data.offsets.at(-1);
+    for (const block of p.block_stack) {
+      const s = block.stripe_data;
+      for (let i = 0; i < s.offsets.length; i++) {
+        s.offsets[i] = env_offset + block.inner_offset;
+      }
+    }
+  }
+
+  // --- Replay alloc/free actions to build the timeline ---
+  for (const elem of actions) {
+    const size = elements[elem].size;
+    const pool_key = include_private_inactive ? get_pool_key(elem) : null;
+
+    if (pool_key) {
+      // --- Private pool element ---
+      if (!(elem in pool_active_elems)) {
+        // Pool alloc: add to pool, grow envelope if needed
+        pool_active_elems[elem] = pool_key;
+        const pool = get_or_create_pool(pool_key);
+
+        if (pool.envelope_data === null) {
+          const env = {
+            elem: `pool:${pool_key}`,
+            timesteps: [timestep],
+            offsets: [total_mem],
+            size: [0],
+            color: 9,
+          };
+          pool.envelope_data = env;
+          current.push(`pool:${pool_key}`);
+          current_data.push(env);
+          data.push(env);
+        }
+
+        const inner_offset = pool.active;
+        pool.active += size;
+
+        if (pool.active > pool.max) {
+          grow_pool_envelope(pool, pool_key, pool.active);
+        }
+
+        const stripe = {
+          elem,
+          timesteps: [timestep],
+          offsets: [pool.envelope_data.offsets.at(-1) + inner_offset],
+          size,
+          color: elem_color(elem),
+          opacity: 0.5,
+        };
+        pool.block_stack.push({elem, size, inner_offset, stripe_data: stripe});
+        data.push(stripe);
+        advance(1);
+        elements[elem].max_allocated_mem = total_mem + total_summarized_mem;
+      } else {
+        // Pool free: end stripe, shift stripes above down within the pool.
+        // The envelope stays at its high-water mark (never shrinks).
+        const pool = pools[pool_key];
+        const block_idx = pool.block_stack.findIndex(b => b.elem === elem);
+        if (block_idx >= 0) {
+          advance(1);
+          const block = pool.block_stack[block_idx];
+          block.stripe_data.timesteps.push(timestep);
+          block.stripe_data.offsets.push(block.stripe_data.offsets.at(-1));
+
+          pool.block_stack.splice(block_idx, 1);
+          pool.active -= size;
+
+          if (block_idx < pool.block_stack.length) {
+            for (let j = block_idx; j < pool.block_stack.length; j++) {
+              const b = pool.block_stack[j];
+              b.inner_offset -= size;
+              const s = b.stripe_data;
+              s.timesteps.push(timestep);
+              s.offsets.push(s.offsets.at(-1));
+              s.timesteps.push(timestep + 3);
+              s.offsets.push(pool.envelope_data.offsets.at(-1) + b.inner_offset);
+            }
+            advance(3);
+          }
+        } else {
+          pool.active -= size;
+          advance(1);
+        }
+        delete pool_active_elems[elem];
+      }
+      max_size = Math.max(total_mem + total_summarized_mem, max_size);
+      continue;
+    }
+
+    // --- Non-pool element ---
+    if (!(elem in draw_elem)) {
+      // Too small to render individually — goes into the summarized band
+      if (elem in summarized_elems) {
+        advance(1);
+        total_summarized_mem -= size;
+        summarized_elems[elem] = null;
+      } else {
+        total_summarized_mem += size;
+        summarized_elems[elem] = true;
+        advance(1);
+      }
+      continue;
+    }
+    const idx = current.findLastIndex(x => x === elem);
+    if (idx === -1) {
+      // First appearance → alloc
+      add_allocation(elem);
+      advance(1);
+    } else {
+      // Second appearance → free: remove from stack, shift elements above down
+      advance(1);
+      const removed = current_data[idx];
+      removed.timesteps.push(timestep);
+      removed.offsets.push(removed.offsets.at(-1));
+      current.splice(idx, 1);
+      current_data.splice(idx, 1);
+
+      if (idx < current.length) {
+        shift_elements_above(idx, -size);
+        advance(3);
+      }
+      total_mem -= size;
+    }
+    max_size = Math.max(total_mem + total_summarized_mem, max_size);
+  }
+
+  // --- Finalize: close all still-active elements ---
+  for (const elem of current_data) {
+    elem.timesteps.push(timestep);
+    elem.offsets.push(elem.offsets.at(-1));
+    if (Array.isArray(elem.size)) {
+      elem.size.push(elem.size.at(-1));
+    }
+  }
+  for (const pk in pools) {
+    for (const block of pools[pk].block_stack) {
+      const s = block.stripe_data;
+      s.timesteps.push(timestep);
+      s.offsets.push(s.offsets.at(-1));
+    }
+  }
+  data.push(summarized_mem);
+
+  return {
+    max_size,
+    allocations_over_time: data,
+    max_at_time,
+    summarized_mem,
+    elements_length: elements.length,
+    context_for_id: id => {
+      const elem = elements[id];
+      let text = `Addr: ${formatAddr(elem)}`;
+      text = `${text}, Size: ${formatSize(elem.size)} allocation`;
+      text = `${text}, Total memory used after allocation: ${formatSize(
+        elem.max_allocated_mem,
+      )}`;
+      const context = elem?.compile_context ?? 'None';
+      text = `${text}, Compile context: ${context}`;
+      if (elem.stream !== null) {
+        text = `${text}, stream ${elem.stream}`;
+      }
+      if (elem.segment_pool_id) {
+        text = `${text}, pool_id (${elem.segment_pool_id[0]}, ${elem.segment_pool_id[1]})`;
+      } else {
+        text = `${text}, pool_id unknown`;
+      }
+      if (elem.timestamp !== null) {
+        var d = new Date(elem.time_us / 1000);
+        text = `${text}, timestamp ${d}`;
+      }
+      if (!elem.action.includes('alloc')) {
+        text = `${text}\nalloc not recorded, stack trace for free:`;
+      }
+      if (elem.ghost) {
+        text = `${text}\n[Ghost block] This block exists in the segment snapshot but has no alloc trace events. ` +
+          `It was allocated before _record_memory_history() was called, or its alloc event was evicted ` +
+          `from the trace ring buffer. The block is still active (not freed) at snapshot time.`;
+      }
+      const user_metadata_str = format_user_metadata(elem.user_metadata);
+      if (user_metadata_str) {
+        text = `${text}\n${user_metadata_str}`;
+      }
+      text = `${text}\n${format_frames(elem.frames)}`;
+      text = `${text}${format_forward_frames(elem.forward_frames)}`;
+      return text;
+    },
+  };
+}
+
+export { process_alloc_data, isPrivatePoolId, formatSize, formatAddr,
+         elideRepeats, frameFilter, format_user_metadata,
+         format_forward_frames, format_frames };
diff --git a/torch/xpu/__init__.py b/torch/xpu/__init__.py
index bc3cf07c88e91..0c77d45201c09 100644
--- a/torch/xpu/__init__.py
+++ b/torch/xpu/__init__.py
@@ -7,16 +7,22 @@
 :func:`is_available()` to determine if your system supports XPU.
 """
 
+from __future__ import annotations
+
 import threading
 import traceback
-from collections.abc import Callable
 from functools import lru_cache
-from typing import Any, NewType, Optional
+from typing import Any, NewType, TYPE_CHECKING
 
 import torch
 import torch._C
 from torch._utils import _dummy_type, _LazySeedTracker
-from torch.types import Device
+
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+    from torch.types import Device
 
 from ._utils import _get_device_index
 from .graphs import (
@@ -266,7 +272,7 @@ def get_device_capability(device: Device = None) -> dict[str, Any]:
 
 def get_device_properties(
     device: Device = None,
-) -> _XpuDeviceProperties:  # pyrefly: ignore  # not-a-type
+) -> _XpuDeviceProperties:
     r"""Get the properties of a device. Returns _XpuDeviceProperties containing the following device properties:
 
     - ``name`` (str): device name.
@@ -279,6 +285,8 @@ def get_device_properties(
     - ``gpu_eu_count`` (int): number of EUs (Execution Unit).
     - ``max_work_group_size``: (int): maximum number of work-items permitted in a work-group.
     - ``max_num_sub_groups`` (int): maximum number of sub-groups supported in a work-group.
+    - ``memory_clock_rate`` (int) maximum clock rate of device's global memory in MHz.
+    - ``memory_bus_width`` (int) maximum bus width between device and memory in bits.
     - ``sub_group_sizes``: (list[int]): a list of supported sub-group sizes.
     - ``local_mem_size`` (int): device local memory capacity that can be allocated per work-group in bytes.
     - ``has_fp16`` (bool): whether float16 dtype is supported.
@@ -353,9 +361,9 @@ class StreamContext:
     .. note:: Streams are per-device.
     """
 
-    cur_stream: Optional["torch.xpu.Stream"]
+    cur_stream: torch.xpu.Stream | None
 
-    def __init__(self, stream: Optional["torch.xpu.Stream"]) -> None:
+    def __init__(self, stream: torch.xpu.Stream | None) -> None:
         self.stream = stream
         self.idx = _get_device_index(None, True)
         if self.idx is None:
@@ -384,7 +392,7 @@ def __exit__(self, type: Any, value: Any, traceback: Any):
         torch.xpu.set_stream(self.src_prev_stream)
 
 
-def stream(stream: Optional["torch.xpu.Stream"]) -> StreamContext:
+def stream(stream: torch.xpu.Stream | None) -> StreamContext:
     r"""Wrap around the Context-manager StreamContext that selects a given stream.
 
     Arguments:
diff --git a/torch/xpu/graphs.py b/torch/xpu/graphs.py
index 415408f62407b..51780050f5937 100644
--- a/torch/xpu/graphs.py
+++ b/torch/xpu/graphs.py
@@ -2,7 +2,7 @@
 
 import typing
 from collections.abc import Callable
-from typing import Optional, overload, TYPE_CHECKING, TypeAlias, Union
+from typing import overload, TYPE_CHECKING, TypeAlias
 from typing_extensions import ParamSpec, Self, TypeVar
 
 import torch
@@ -34,7 +34,6 @@
         "_xpu_isCurrentStreamCapturing"
     )
 
-# pyrefly: ignore [missing-module-attribute]
 from torch._C import _xpu_graph_pool_handle, _xpu_isCurrentStreamCapturing, _XPUGraph
 
 
@@ -71,7 +70,7 @@ class XPUGraph(_XPUGraph):
     def __new__(cls, keep_graph: bool = False) -> Self:
         return super().__new__(cls, keep_graph)
 
-    def capture_begin(self, pool: Optional[_POOL_HANDLE] = None) -> None:
+    def capture_begin(self, pool: _POOL_HANDLE | None = None) -> None:
         r"""Begin capturing XPU work on the current xpu stream.
 
         Typically, you shouldn't call ``capture_begin`` yourself.
@@ -165,13 +164,13 @@ class graph:
 
     """  # noqa: B950
 
-    default_capture_stream: Optional[torch.xpu.Stream] = None
+    default_capture_stream: torch.xpu.Stream | None = None
 
     def __init__(
         self,
         xpu_graph: XPUGraph,
-        pool: Optional[_POOL_HANDLE] = None,
-        stream: Optional[torch.xpu.Stream] = None,
+        pool: _POOL_HANDLE | None = None,
+        stream: torch.xpu.Stream | None = None,
     ):
         # Lazy-init of default_capture_stream helps avoid circular-import errors.
         # Not thread safe, but graphs already have the general (explicitly documented)
@@ -179,9 +178,7 @@ def __init__(
         if self.__class__.default_capture_stream is None:
             self.__class__.default_capture_stream = torch.xpu.Stream()
 
-        self.pool: Union[tuple[()], tuple[_POOL_HANDLE]] = (
-            () if pool is None else (pool,)
-        )
+        self.pool: tuple[()] | tuple[_POOL_HANDLE] = () if pool is None else (pool,)
         self.capture_stream = (
             stream if stream is not None else self.__class__.default_capture_stream
         )
@@ -204,7 +201,7 @@ def __exit__(self, *args: object) -> None:
         self.stream_ctx.__exit__(*args)
 
 
-_ModuleOrCallable: TypeAlias = Union["torch.nn.Module", Callable[..., object]]
+_ModuleOrCallable: TypeAlias = torch.nn.Module | Callable[..., object]
 
 
 @overload
@@ -213,7 +210,7 @@ def make_graphed_callables(
     sample_args: tuple[Tensor, ...],
     num_warmup_iters: int = 3,
     allow_unused_input: bool = False,
-    pool: Optional[_POOL_HANDLE] = None,
+    pool: _POOL_HANDLE | None = None,
 ) -> _ModuleOrCallable: ...
 
 
@@ -223,17 +220,17 @@ def make_graphed_callables(
     sample_args: tuple[tuple[Tensor, ...], ...],
     num_warmup_iters: int = 3,
     allow_unused_input: bool = False,
-    pool: Optional[_POOL_HANDLE] = None,
+    pool: _POOL_HANDLE | None = None,
 ) -> tuple[_ModuleOrCallable, ...]: ...
 
 
 def make_graphed_callables(
-    callables: Union[_ModuleOrCallable, tuple[_ModuleOrCallable, ...]],
-    sample_args: Union[tuple[Tensor, ...], tuple[tuple[Tensor, ...], ...]],
+    callables: _ModuleOrCallable | tuple[_ModuleOrCallable, ...],
+    sample_args: tuple[Tensor, ...] | tuple[tuple[Tensor, ...], ...],
     num_warmup_iters: int = 3,
     allow_unused_input: bool = False,
-    pool: Optional[_POOL_HANDLE] = None,
-) -> Union[_ModuleOrCallable, tuple[_ModuleOrCallable, ...]]:
+    pool: _POOL_HANDLE | None = None,
+) -> _ModuleOrCallable | tuple[_ModuleOrCallable, ...]:
     r"""Accept callables (functions or :class:`nn.Module<torch.nn.Module>`\ s) and returns graphed versions.
 
     Each graphed callable's forward pass runs its source callable's
@@ -445,7 +442,7 @@ def make_graphed_autograd_function(
         output_unflatten_spec: torch.utils._pytree.TreeSpec,
         static_input_surface: tuple[Tensor, ...],
         static_outputs: tuple[Tensor, ...],
-        static_grad_outputs: tuple[Optional[Tensor], ...],
+        static_grad_outputs: tuple[Tensor | None, ...],
         static_grad_inputs: tuple[Tensor, ...],
     ) -> Callable[..., object]:
         class Graphed(torch.autograd.Function):
@@ -478,9 +475,7 @@ def backward(ctx: object, *grads: Tensor) -> tuple[Tensor, ...]:
                 if not isinstance(static_grad_inputs, tuple):
                     raise RuntimeError("static_grad_inputs must be a tuple")
                 return tuple(
-                    # pyrefly: ignore [bad-argument-type]
-                    b.detach() if b is not None else b
-                    for b in static_grad_inputs
+                    b.detach() if b is not None else b for b in static_grad_inputs
                 )
 
         def functionalized(*user_args: object) -> object:
diff --git a/torch/xpu/memory.py b/torch/xpu/memory.py
index 9e70381e17b4b..04ca0dd9fc397 100644
--- a/torch/xpu/memory.py
+++ b/torch/xpu/memory.py
@@ -251,7 +251,6 @@ def set_per_process_memory_fraction(fraction: float, device: Device = None) -> N
     device = _get_device_index(device, optional=True)
     if not isinstance(fraction, float):
         raise TypeError("Invalid type for fraction argument, must be `float`")
-    # pyrefly: ignore [missing-attribute]
     torch._C._xpu_setMemoryFraction(fraction, device)
 
 
@@ -271,7 +270,6 @@ def memory_snapshot(
     """
     if not is_initialized():
         return []
-    # pyrefly: ignore [missing-attribute]
     return torch._C._xpu_memorySnapshot(mempool_id)["segments"]
 
 
@@ -358,7 +356,6 @@ class TraceEntry(TypedDict):
     Returns:
         The Snapshot dictionary object
     """
-    # pyrefly: ignore [missing-attribute]
     s = torch._C._xpu_memorySnapshot(None)
     if augment_with_fx_traces:
         s = _augment_memory_snapshot_stack_traces(s)  # type: ignore[assignment, arg-type]
@@ -467,7 +464,6 @@ def _record_memory_history(
 
             Defaults to ``None`` (record all actions).
     """
-    # pyrefly: ignore [missing-attribute]
     torch._C._xpu_recordMemoryHistory(
         enabled,
         context,
@@ -481,7 +477,6 @@ def _record_memory_history(
 class _XPUAllocator:
     r"""Wrapper over internal XPU memory allocators."""
 
-    # pyrefly: ignore [missing-attribute]
     def __init__(self, allocator: torch._C._xpu_XPUAllocator):
         self._allocator = allocator
 
@@ -528,7 +523,6 @@ def __init__(self, path_to_lib_file: str, alloc_fn_name: str, free_fn_name: str)
                 "Failed to load allocator symbols from the shared library."
             )
 
-        # pyrefly: ignore [missing-attribute]
         self._allocator = torch._C._xpu_customAllocator(alloc_fn_addr, free_fn_addr)
 
 
@@ -541,7 +535,6 @@ def change_current_allocator(allocator: _XPUAllocator) -> None:
     Arguments:
         allocator (torch.xpu.memory._XPUAllocator): allocator to be set as the active one.
     """
-    # pyrefly: ignore [missing-attribute]
     torch._C._xpu_changeCurrentAllocator(allocator.allocator())
 
 
@@ -551,7 +544,6 @@ def _get_current_allocator() -> _XPUAllocator:
     Returns:
         _XPUAllocator: the allocator being currently used.
     """
-    # pyrefly: ignore [missing-attribute]
     return _XPUAllocator(torch._C._xpu_getAllocator())
 
 
diff --git a/torch/xpu/streams.py b/torch/xpu/streams.py
index 87c953ce7c323..3e5068ee3b44b 100644
--- a/torch/xpu/streams.py
+++ b/torch/xpu/streams.py
@@ -121,7 +121,6 @@ class Event(torch._C._XpuEventBase):
     def __new__(cls, enable_timing=False):
         return super().__new__(cls, enable_timing=enable_timing)
 
-    # pyrefly: ignore [bad-override]
     def record(self, stream: Stream | torch.Stream | None = None) -> None:
         r"""Record the event in a given stream.
 
@@ -131,7 +130,6 @@ def record(self, stream: Stream | torch.Stream | None = None) -> None:
         """
         if stream is None:
             stream = torch.xpu.current_stream()
-        # pyrefly: ignore [bad-argument-type]
         super().record(stream)
 
     def wait(self, stream: Stream | torch.Stream | None = None) -> None:
@@ -142,7 +140,6 @@ def wait(self, stream: Stream | torch.Stream | None = None) -> None:
         """
         if stream is None:
             stream = torch.xpu.current_stream()
-        # pyrefly: ignore [bad-argument-type]
         super().wait(stream)
 
     def query(self) -> bool:
diff --git a/torchgen/_autoheuristic/benchmark_runner.py b/torchgen/_autoheuristic/benchmark_runner.py
index 117058b3373f3..67e24aec82d36 100644
--- a/torchgen/_autoheuristic/benchmark_runner.py
+++ b/torchgen/_autoheuristic/benchmark_runner.py
@@ -55,12 +55,15 @@ def add_base_arguments(self) -> None:
     def run(self) -> None:
         torch.set_default_device("cuda")
         args = self.parser.parse_args()
+        # Set environment variables to control autoheuristic behavior
+        import os
+
         if args.use_heuristic:
-            torch._inductor.config.autoheuristic_use = self.name
-            torch._inductor.config.autoheuristic_collect = ""
+            os.environ["TORCHINDUCTOR_AUTOHEURISTIC_USE"] = self.name
+            os.environ["TORCHINDUCTOR_AUTOHEURISTIC_COLLECT"] = ""
         else:
-            torch._inductor.config.autoheuristic_use = ""
-            torch._inductor.config.autoheuristic_collect = self.name
+            os.environ["TORCHINDUCTOR_AUTOHEURISTIC_USE"] = ""
+            os.environ["TORCHINDUCTOR_AUTOHEURISTIC_COLLECT"] = self.name
         torch._inductor.config.autoheuristic_log_path = args.o
         if args.device is not None:
             torch.cuda.set_device(args.device)
diff --git a/torchgen/_autoheuristic/pad_mm/collect_known_mm_shapes.py b/torchgen/_autoheuristic/pad_mm/collect_known_mm_shapes.py
new file mode 100644
index 0000000000000..45a5854044b33
--- /dev/null
+++ b/torchgen/_autoheuristic/pad_mm/collect_known_mm_shapes.py
@@ -0,0 +1,235 @@
+import argparse
+import csv
+import sys
+from pathlib import Path
+
+
+# Add parent directory to path for imports
+sys.path.append(str(Path(__file__).absolute().parents[1]))
+sys.path.append(
+    str(
+        Path(__file__).absolute().parents[3]
+        / "benchmarks"
+        / "dynamo"
+        / "microbenchmarks"
+    )
+)
+
+from operator_inp_utils import (  # type: ignore[import-not-found]
+    deserialize_args,
+    OperatorInputsLoader,
+)
+
+import torch
+from torch._inductor.fx_passes.pad_mm import (
+    get_alignment_size_dtype,  # type: ignore[import-not-found]
+)
+from torch._subclasses.fake_tensor import FakeTensorMode
+
+
+def is_aligned(dim: int, align_size: int) -> bool:
+    """Check if dimension is aligned to the given alignment size."""
+    return dim % align_size == 0
+
+
+def extract_mm_shapes_from_loader(
+    loader: OperatorInputsLoader,
+) -> list[tuple[int, int, int, torch.dtype, torch.dtype]]:
+    """Extract matrix multiplication shapes from an OperatorInputsLoader using deserialize_args with FakeTensorMode."""
+    shapes = []
+
+    # Matrix multiplication operators to look for
+    mm_operators = ["aten.mm.default", "aten.addmm.default", "aten.bmm.default"]
+
+    # Use FakeTensorMode to avoid instantiating actual tensors
+    with FakeTensorMode():
+        for op_name in mm_operators:
+            if op_name not in loader.operator_db:
+                continue
+
+            # Count shapes extracted from this operator
+            shape_count = 0
+
+            # Access the raw string data directly from operator_db and reuse existing parsing
+            for input_str in loader.operator_db[op_name]:
+                try:
+                    # Use deserialize_args to parse inputs - will create fake tensors
+                    args, kwargs = deserialize_args(input_str)
+
+                    if op_name == "aten.mm.default":
+                        # mm(input, mat2) -> result
+                        if len(args) >= 2:
+                            a, b = args[0], args[1]
+                            if isinstance(a, torch.Tensor) and isinstance(
+                                b, torch.Tensor
+                            ):
+                                a_shape, a_dtype = tuple(a.shape), a.dtype
+                                b_shape, b_dtype = tuple(b.shape), b.dtype
+                                if len(a_shape) == 2 and len(b_shape) == 2:
+                                    m, k = a_shape
+                                    k2, n = b_shape
+                                    if k == k2:  # Valid matrix multiplication
+                                        shapes.append((m, k, n, a_dtype, b_dtype))
+                                        shape_count += 1
+
+                    elif op_name == "aten.addmm.default":
+                        # addmm(bias, input, mat2) -> result
+                        if len(args) >= 3:
+                            _, a, b = args[0], args[1], args[2]
+                            if isinstance(a, torch.Tensor) and isinstance(
+                                b, torch.Tensor
+                            ):
+                                a_shape, a_dtype = tuple(a.shape), a.dtype
+                                b_shape, b_dtype = tuple(b.shape), b.dtype
+                                if len(a_shape) == 2 and len(b_shape) == 2:
+                                    m, k = a_shape
+                                    k2, n = b_shape
+                                    if k == k2:  # Valid matrix multiplication
+                                        shapes.append((m, k, n, a_dtype, b_dtype))
+                                        shape_count += 1
+
+                    elif op_name == "aten.bmm.default":
+                        # bmm(input, mat2) -> result (batch matrix multiplication)
+                        if len(args) >= 2:
+                            a, b = args[0], args[1]
+                            if isinstance(a, torch.Tensor) and isinstance(
+                                b, torch.Tensor
+                            ):
+                                a_shape, a_dtype = tuple(a.shape), a.dtype
+                                b_shape, b_dtype = tuple(b.shape), b.dtype
+                                if len(a_shape) == 3 and len(b_shape) == 3:
+                                    batch1, m, k = a_shape
+                                    batch2, k2, n = b_shape
+                                    if (
+                                        batch1 == batch2 and k == k2
+                                    ):  # Valid batch matrix multiplication
+                                        shapes.append((m, k, n, a_dtype, b_dtype))
+                                        shape_count += 1
+
+                except Exception:
+                    # Skip invalid inputs
+                    continue
+
+            print(f"    Extracted {shape_count} shapes from {op_name}")
+
+    return shapes
+
+
+def filter_unaligned_shapes(
+    shapes: list[tuple[int, int, int, torch.dtype, torch.dtype]],
+) -> list[tuple[int, int, int, torch.dtype, torch.dtype]]:
+    """Filter shapes to keep only those that are not completely aligned (so padding is relevant)."""
+    filtered_shapes = []
+
+    for m, k, n, dtype1, dtype2 in shapes:
+        # Use the primary dtype for alignment calculation (assume both dtypes are similar for alignment purposes)
+        dtype = dtype1
+        try:
+            align_size = get_alignment_size_dtype(dtype)
+
+            # Only keep shapes where not all dimensions are aligned
+            if not all(is_aligned(dim, align_size) for dim in [m, k, n]):
+                filtered_shapes.append((m, k, n, dtype1, dtype2))
+
+        except Exception:
+            # If we can't get alignment size, skip this shape
+            continue
+
+    return filtered_shapes
+
+
+def collect_known_mm_shapes() -> list[tuple[int, int, int, torch.dtype, torch.dtype]]:
+    """
+    Collect known matrix multiplication shapes from HuggingFace, TIMM, and TorchBench datasets.
+
+    Returns:
+        List of tuples containing (m, k, n, dtype1, dtype2) for matrix multiplication shapes
+        that are not completely aligned (so padding is relevant).
+    """
+    all_shapes = []
+
+    loaders = []
+
+    # Try to load each dataset
+    try:
+        hf_loader = OperatorInputsLoader.get_huggingface_loader()
+        loaders.append(("HuggingFace", hf_loader))
+    except Exception as e:
+        print(f"Warning: Could not load HuggingFace dataset: {e}")
+
+    try:
+        timm_loader = OperatorInputsLoader.get_timm_loader()
+        loaders.append(("TIMM", timm_loader))
+    except Exception as e:
+        print(f"Warning: Could not load TIMM dataset: {e}")
+
+    try:
+        torchbench_loader = OperatorInputsLoader.get_torchbench_loader()
+        loaders.append(("TorchBench", torchbench_loader))
+    except Exception as e:
+        print(f"Warning: Could not load TorchBench dataset: {e}")
+
+    # Extract shapes from each loader
+    for dataset_name, loader in loaders:
+        print(f"Extracting shapes from {dataset_name}...")
+
+        shapes = extract_mm_shapes_from_loader(loader)
+        print(f"Found {len(shapes)} matrix multiplication shapes from {dataset_name}")
+        all_shapes.extend(shapes)
+
+    # Remove duplicates
+    unique_shapes = list(set(all_shapes))
+    print(f"Total unique shapes before filtering: {len(unique_shapes)}")
+
+    # Filter for unaligned shapes only
+    filtered_shapes = filter_unaligned_shapes(unique_shapes)
+    print(f"Shapes after filtering for unaligned: {len(filtered_shapes)}")
+
+    return filtered_shapes
+
+
+def main(output_file="mm_shapes.csv"):
+    shapes = collect_known_mm_shapes()
+
+    print(f"\nCollected {len(shapes)} real-world matrix multiplication shapes")
+
+    # Convert dtype objects to strings and filter for desired dtypes
+    dtype_map = {
+        torch.float16: "float16",
+        torch.bfloat16: "bfloat16",
+        torch.float32: "float32",
+    }
+
+    # Convert to desired format and filter dtypes
+    csv_rows = []
+    for m, k, n, dtype1, dtype2 in shapes:
+        # Use the first dtype and convert to string
+        if dtype1 in dtype_map:
+            dtype_str = dtype_map[dtype1]
+            csv_rows.append([m, k, n, dtype_str])
+
+    # Save to CSV file
+    with open(output_file, "w", newline="") as csvfile:
+        writer = csv.writer(csvfile)
+        # Write header
+        writer.writerow(["M", "K", "N", "dtype"])
+        # Write data rows
+        writer.writerows(csv_rows)
+
+    print(f"Saved matrix multiplication shapes to {output_file}")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(
+        description="Collect matrix multiplication shapes from real-world datasets and save to CSV"
+    )
+    parser.add_argument(
+        "--output",
+        "-o",
+        type=str,
+        default="mm_shapes.csv",
+        help="Output CSV filename (default: mm_shapes.csv)",
+    )
+
+    args = parser.parse_args()
+    main(args.output)
diff --git a/torchgen/_autoheuristic/pad_mm/evaluate_pad_mm_heuristics.py b/torchgen/_autoheuristic/pad_mm/evaluate_pad_mm_heuristics.py
new file mode 100644
index 0000000000000..42b37d74c3092
--- /dev/null
+++ b/torchgen/_autoheuristic/pad_mm/evaluate_pad_mm_heuristics.py
@@ -0,0 +1,426 @@
+#!/usr/bin/env python3
+
+import argparse
+import csv
+import functools
+
+import torch
+from torch._inductor.autoheuristic.autoheuristic_utils import AHContext, AHMetadata
+from torch._inductor.fx_passes.pad_mm import get_alignment_size_dtype
+from torch._inductor.runtime.benchmarking import benchmarker
+from torch._inductor.utils import get_gpu_shared_memory
+
+
+def fits_in_memory(dtype, m: int, k: int, n: int) -> bool:
+    threshold_memory = torch.cuda.get_device_properties(0).total_memory / 4
+    return dtype.itemsize * (m * k + k * n + m * n) < threshold_memory
+
+
+def set_precision(dtype, float32_precision: str = "highest") -> None:
+    precision = float32_precision if dtype == torch.float32 else "high"
+    torch.set_float32_matmul_precision(precision)
+
+
+def get_heuristic_decision(m: int, k: int, n: int, dtype: torch.dtype) -> str | None:
+    from torch._inductor.autoheuristic.autoheuristic import AutoHeuristic, LocalFeedback
+    from torch._inductor.fx_passes.pad_mm import (
+        get_alignment_size,
+        get_context,
+        get_padded_length,
+        pad_mm_operations,
+        pad_mm_precondition,
+    )
+
+    torch._inductor.config.autoheuristic_use.pad_mm = True
+
+    if not torch._inductor.config.run_autoheuristic("pad_mm"):
+        return None
+
+    a = torch.randn(m, k, dtype=dtype, device="cuda")
+    b = torch.randn(k, n, dtype=dtype, device="cuda")
+
+    m_padded_length = get_padded_length(m, get_alignment_size(a))
+    k_padded_length = get_padded_length(k, get_alignment_size(a))
+    n_padded_length = get_padded_length(n, get_alignment_size(b))
+
+    context = get_context(
+        a,
+        b,
+        mat1_pre_padded=False,
+        mat2_pre_padded=False,
+        m_padded_length=m_padded_length,
+        k_padded_length=k_padded_length,
+        n_padded_length=n_padded_length,
+    )
+
+    def dummy_feedback(choice: str) -> float:
+        return 1.0
+
+    def fallback() -> str:
+        return "no_decision"
+
+    autoheuristic = AutoHeuristic(
+        fallback=fallback,
+        choices=["orig", "pad"],
+        feedback=LocalFeedback(dummy_feedback),
+        context=context,
+        name="pad_mm",
+        augment_context=pad_mm_operations(),
+        precondition=pad_mm_precondition,
+    )
+
+    choice = autoheuristic.get_choice()
+    return choice
+
+
+def benchmark_both_choices(
+    m: int,
+    k: int,
+    n: int,
+    dtype: torch.dtype,
+    num_reps: int = 3,
+    float32_precision: str = "highest",
+) -> tuple[float, float]:
+    set_precision(dtype, float32_precision)
+    a = torch.randn(m, k, dtype=dtype, device="cuda")
+    b = torch.randn(k, n, dtype=dtype, device="cuda")
+
+    # Use existing benchmarking infrastructure with proper cache management
+    # benchmarker returns time in milliseconds, so convert to seconds for consistency
+    orig_time_ms = benchmarker.benchmark(
+        torch.mm, fn_args=(a, b), rep=num_reps, is_vetted_benchmarking=True
+    )
+    orig_time = orig_time_ms / 1000.0  # Convert ms to seconds
+
+    from torch._inductor.fx_passes.pad_mm import (
+        get_alignment_size,
+        get_padded_length,
+        pad_mm,
+    )
+
+    m_padded_length = get_padded_length(a.shape[0], get_alignment_size(a))
+    k_padded_length = get_padded_length(a.shape[1], get_alignment_size(a))
+    n_padded_length = get_padded_length(b.shape[1], get_alignment_size(b))
+
+    if m_padded_length == 0 and k_padded_length == 0 and n_padded_length == 0:
+        return orig_time, orig_time
+
+    pad_time_ms = benchmarker.benchmark(
+        pad_mm,
+        fn_args=(a, b, m_padded_length, k_padded_length, n_padded_length),
+        rep=num_reps,
+        is_vetted_benchmarking=True,
+    )
+    pad_time = pad_time_ms / 1000.0  # Convert ms to seconds
+
+    return orig_time, pad_time
+
+
+def load_shapes_from_csv(csv_file: str) -> list:
+    shapes = []
+    with open(csv_file) as f:
+        reader = csv.DictReader(f)
+        for row in reader:
+            m, k, n = int(row["M"]), int(row["K"]), int(row["N"])
+            dtype_str = row["dtype"]
+
+            if dtype_str == "float16":
+                dtype = torch.float16
+            elif dtype_str == "bfloat16":
+                dtype = torch.bfloat16
+            elif dtype_str == "float32":
+                dtype = torch.float32
+            else:
+                continue
+
+            shapes.append((m, k, n, dtype))
+
+    print(f"Loaded {len(shapes)} shapes from {csv_file}")
+    return shapes
+
+
+@functools.cache
+def get_shared_mem_size():
+    return get_gpu_shared_memory()
+
+
+def check_shape_passes_precondition(m: int, k: int, n: int, dtype: torch.dtype) -> bool:
+    """
+    Check if a shape passes the same precondition used by the actual pad_mm AutoHeuristics.
+
+    This uses the exact same pad_mm_precondition function that the AutoHeuristic system
+    uses, avoiding hardcoded magic numbers by delegating to the source of truth.
+    """
+    from torch._inductor.autoheuristic.autoheuristic_utils import pad_mm_precondition
+
+    shared_memory = get_shared_mem_size()
+    device_capa = torch.cuda.get_device_capability()
+
+    # Create the same metadata and context that AutoHeuristics uses
+    metadata = AHMetadata(
+        shared_memory=shared_memory,
+        device_capa=device_capa,
+        choices=["orig", "pad"],  # Required but not used for precondition check
+        name="pad_mm",  # Required but not used for precondition check
+    )
+
+    context = AHContext()
+    context.add_feature("m", m)
+    context.add_feature("k", k)
+    context.add_feature("n", n)
+
+    # Use the actual pad_mm_precondition function - no hardcoded values!
+    return pad_mm_precondition(metadata, context)
+
+
+def filter_shapes(shapes: list) -> list:
+    filtered = []
+    aligned_count = 0
+    precondition_failed_count = 0
+    memory_count = 0
+
+    for m, k, n, dtype in shapes:
+        # Check if already aligned
+        align_size = get_alignment_size_dtype(dtype)
+        is_aligned = all((dim % align_size == 0) for dim in [m, k, n])
+
+        if is_aligned:
+            aligned_count += 1
+            continue
+
+        # Check if passes the actual precondition used by pad_mm AutoHeuristics
+        if not check_shape_passes_precondition(m, k, n, dtype):
+            precondition_failed_count += 1
+            continue
+
+        # Check if fits in memory
+        if not fits_in_memory(dtype, m, k, n):
+            memory_count += 1
+            continue
+
+        # This shape is suitable for evaluation
+        filtered.append((m, k, n, dtype))
+
+    print("Filtering results:")
+    print(f"  Already aligned (skipped): {aligned_count}")
+    print(f"  Failed pad_mm_precondition (skipped): {precondition_failed_count}")
+    print(f"  Too large for memory (skipped): {memory_count}")
+    print(f"  Suitable for evaluation: {len(filtered)}")
+
+    return filtered
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Evaluate trained AutoHeuristics for pad_mm optimization"
+    )
+    parser.add_argument("csv_file", help="Path to CSV file with M,K,N,dtype columns")
+    parser.add_argument(
+        "--num-reps", type=int, default=3, help="Benchmark repetitions (default: 3)"
+    )
+    parser.add_argument(
+        "--device", type=int, default=None, help="CUDA device (default: current)"
+    )
+    parser.add_argument(
+        "--max-shapes",
+        type=int,
+        default=10000,
+        help="Max shapes to test (default: 10000)",
+    )
+    parser.add_argument(
+        "--float32_matmul_precision",
+        type=str,
+        choices=["high", "highest"],
+        default="highest",
+        help="Matmul precision for float32 (default: highest). Non-fp32 always uses 'high'.",
+    )
+
+    args = parser.parse_args()
+
+    torch.set_default_device("cuda")
+    if args.device is not None:
+        torch.cuda.set_device(args.device)
+
+    print(f"Using CUDA device: {torch.cuda.current_device()}")
+    print()
+
+    shapes = load_shapes_from_csv(args.csv_file)
+    if not shapes:
+        print("No shapes found!")
+        return
+
+    shapes = filter_shapes(shapes)
+    if not shapes:
+        print("No suitable shapes found!")
+        return
+
+    if len(shapes) > args.max_shapes:
+        shapes = shapes[: args.max_shapes]
+        print(f"Limited to first {args.max_shapes} shapes")
+
+    print(f"Evaluating {len(shapes)} shapes with {args.num_reps} reps each")
+    print()
+
+    total_decisions = 0
+    correct_decisions = 0
+    true_positives = 0  # Chose pad, should pad
+    true_negatives = 0  # Chose orig, should orig
+    false_positives = 0  # Chose pad, should orig
+    false_negatives = 0  # Chose orig, should pad
+    no_decision_shapes = 0
+
+    tp_speedups = []  # Speed-up percentages for true positives
+    fp_slowdowns = []  # Speed-down percentages for false positives
+
+    # Track non-confident decisions and confident decisions by dtype
+    no_decision_shape_list = []  # List of (M, K, N, dtype) where heuristic chose no_decision
+    confident_by_dtype = {}  # Count of confident decisions by dtype
+
+    for i, (m, k, n, dtype) in enumerate(shapes, 1):
+        print(f"Shape {i}/{len(shapes)}: M={m}, K={k}, N={n}, dtype={dtype}")
+
+        heuristic_choice = get_heuristic_decision(m, k, n, dtype)
+        print(f"  Heuristic: {heuristic_choice}")
+
+        orig_time, pad_time = benchmark_both_choices(
+            m, k, n, dtype, args.num_reps, args.float32_matmul_precision
+        )
+        ground_truth = "pad" if pad_time < orig_time else "orig"
+
+        print(f"  Times: orig={orig_time:.6f}s, pad={pad_time:.6f}s")
+        print(f"  Ground truth: {ground_truth}")
+
+        if heuristic_choice == "no_decision":
+            # Heuristic punted to benchmarking - this is correct behavior for small/uncertain shapes
+            no_decision_shapes += 1
+            no_decision_shape_list.append((m, k, n, dtype))
+            print("  Heuristic chose to benchmark (conservative)")
+        else:
+            # Heuristic made a confident decision - evaluate accuracy
+            total_decisions += 1
+            # Track confident decisions by dtype
+            dtype_str = str(dtype).replace("torch.", "")
+            confident_by_dtype[dtype_str] = confident_by_dtype.get(dtype_str, 0) + 1
+            if heuristic_choice == ground_truth:
+                correct_decisions += 1
+                print("  ✓ CORRECT")
+                if heuristic_choice == "pad":
+                    true_positives += 1  # Correctly chose pad
+                    # Calculate speed-up: (orig_time - pad_time) / orig_time * 100
+                    speedup = (orig_time - pad_time) / orig_time * 100
+                    tp_speedups.append(speedup)
+                    print(f"    Speed-up: {speedup:.1f}%")
+                else:
+                    true_negatives += 1  # Correctly chose orig
+            else:
+                print("  ✗ WRONG")
+                if heuristic_choice == "pad" and ground_truth == "orig":
+                    false_positives += 1
+                    # Calculate speed-down: (pad_time - orig_time) / orig_time * 100
+                    slowdown = (pad_time - orig_time) / orig_time * 100
+                    fp_slowdowns.append(slowdown)
+                    print(f"    Speed-down: {slowdown:.1f}%")
+                elif heuristic_choice == "orig" and ground_truth == "pad":
+                    false_negatives += 1
+
+        print(f"  Confidence Rate: {total_decisions}/{i}")
+        if total_decisions > 0:
+            accuracy = correct_decisions / total_decisions * 100
+            tp_rate = true_positives / total_decisions * 100
+            tn_rate = true_negatives / total_decisions * 100
+            fp_rate = false_positives / total_decisions * 100
+            fn_rate = false_negatives / total_decisions * 100
+
+            # Compute average speedup/slowdown
+            avg_tp_speedup = sum(tp_speedups) / len(tp_speedups) if tp_speedups else 0
+            avg_fp_slowdown = (
+                sum(fp_slowdowns) / len(fp_slowdowns) if fp_slowdowns else 0
+            )
+
+            print(
+                f"  Accuracy: {correct_decisions}/{total_decisions} ({accuracy:.1f}%) "
+                f"| TP: {tp_rate:.1f}% (avg speedup: {avg_tp_speedup:.1f}%) "
+                f"| TN: {tn_rate:.1f}% "
+                f"| FP: {fp_rate:.1f}% (avg slowdown: {avg_fp_slowdown:.1f}%)"
+                f"| FN: {fn_rate:.1f}%"
+            )
+
+        print()
+
+    print("=== FINAL RESULTS ===")
+    print(f"Confident decisions: {total_decisions}")
+    print(f"#Shapes without confident decisions: {no_decision_shapes}")
+
+    if total_decisions > 0:
+        accuracy = correct_decisions / total_decisions * 100
+        tp_rate = true_positives / total_decisions * 100
+        tn_rate = true_negatives / total_decisions * 100
+        fp_rate = false_positives / total_decisions * 100
+        fn_rate = false_negatives / total_decisions * 100
+
+        avg_tp_speedup = sum(tp_speedups) / len(tp_speedups) if tp_speedups else 0
+        avg_fp_slowdown = sum(fp_slowdowns) / len(fp_slowdowns) if fp_slowdowns else 0
+
+        print(
+            f"\nConfident decision accuracy: {accuracy:.1f}% ({correct_decisions}/{total_decisions})"
+        )
+
+        if tp_speedups:
+            print(
+                f"True Positives (chose pad, should pad): {tp_rate:.1f}% ({true_positives}) "
+                f"| Avg speed-up: {avg_tp_speedup:.1f}%"
+            )
+        else:
+            print(
+                f"True Positives (chose pad, should pad): {tp_rate:.1f}% ({true_positives})"
+            )
+
+        print(
+            f"True Negatives (chose orig, should orig): {tn_rate:.1f}% ({true_negatives})"
+        )
+
+        if fp_slowdowns:
+            print(
+                f"False Positives (chose pad, should orig): {fp_rate:.1f}% ({false_positives}) "
+                f"| Avg speed-down: {avg_fp_slowdown:.1f}%"
+            )
+        else:
+            print(
+                f"False Positives (chose pad, should orig): {fp_rate:.1f}% ({false_positives})"
+            )
+
+        print(
+            f"False Negatives (chose orig, should pad): {fn_rate:.1f}% ({false_negatives})"
+        )
+    else:
+        print("No confident decisions made!")
+
+    total_evaluated = total_decisions + no_decision_shapes
+    if total_evaluated > 0:
+        print(
+            f"\nConfidence rate: ({total_decisions}/{total_evaluated} made confident decisions)"
+        )
+
+    # Print shapes where AutoHeuristics did not make a confident decision
+    print(f"\n=== NON-CONFIDENT DECISIONS ({len(no_decision_shape_list)}) ===")
+    if no_decision_shape_list:
+        print("Shapes where AutoHeuristics chose 'no_decision' (non-confident):")
+        for m, k, n, dtype in no_decision_shape_list:
+            dtype_str = str(dtype).replace("torch.", "")
+            print(f"  M={m}, K={k}, N={n}, dtype={dtype_str}")
+    else:
+        print("All shapes had confident decisions!")
+
+    # Print confident decisions by dtype
+    print("\n=== CONFIDENT DECISIONS BY DTYPE ===")
+    if confident_by_dtype:
+        print("Number of confident decisions per dtype:")
+        for dtype_str, count in sorted(confident_by_dtype.items()):
+            print(f"  {dtype_str}: {count} confident decisions")
+        print(f"Total confident decisions: {sum(confident_by_dtype.values())}")
+    else:
+        print("No confident decisions made!")
+
+
+if __name__ == "__main__":
+    main()
diff --git a/torchgen/_autoheuristic/pad_mm/gen_data_pad_mm.py b/torchgen/_autoheuristic/pad_mm/gen_data_pad_mm.py
index b476bacfb67db..eba7dd19b53b2 100644
--- a/torchgen/_autoheuristic/pad_mm/gen_data_pad_mm.py
+++ b/torchgen/_autoheuristic/pad_mm/gen_data_pad_mm.py
@@ -1,5 +1,7 @@
+import csv
 import random
 import sys
+from collections.abc import Generator
 from pathlib import Path
 from typing import Any
 
@@ -10,9 +12,13 @@
 from benchmark_utils import (  # type: ignore[import-not-found]
     fits_in_memory,
     get_mm_tensors,
+    get_random_between_pow2,
     set_precision,
     transpose_tensors,
 )
+from collect_known_mm_shapes import (
+    collect_known_mm_shapes,  # type: ignore[import-not-found]
+)
 
 import torch
 from torch._inductor.fx_passes.pad_mm import (  # type: ignore[import-not-found]
@@ -29,10 +35,121 @@ class BenchmarkRunnerPadMM(BenchmarkRunner):  # type: ignore[misc, no-any-unimpo
     def __init__(self) -> None:
         super().__init__("pad_mm")
 
+        # Add CLI argument for additional shape CSV files
+        self.parser.add_argument(
+            "--additional-shape-csv",
+            nargs="*",
+            default=[],
+            help="List of CSV files containing additional matrix multiplication shapes (M,K,N,dtype format)",
+        )
+
+        # Initialize additional_shape_collections
+        self.additional_shape_collections: list[
+            list[tuple[int, int, int, torch.dtype, torch.dtype]]
+        ] = []
+
+        # Initialize the shape generator (will be set up after parsing args)
+        self.shape_generator = None
+
+    def load_shapes_from_csv(
+        self, csv_file: str
+    ) -> list[tuple[int, int, int, torch.dtype, torch.dtype]]:
+        """Load matrix multiplication shapes from a CSV file in M,K,N,dtype format."""
+        shapes = []
+        dtype_map = {
+            "float16": torch.float16,
+            "bfloat16": torch.bfloat16,
+            "float32": torch.float32,
+        }
+
+        try:
+            with open(csv_file) as f:
+                reader = csv.DictReader(f)
+                for row in reader:
+                    m = int(row["M"])
+                    k = int(row["K"])
+                    n = int(row["N"])
+                    dtype_str = row["dtype"]
+
+                    if dtype_str in dtype_map:
+                        dtype = dtype_map[dtype_str]
+                        # Store as (m, k, n, dtype1, dtype2) with same dtype for both
+                        shapes.append((m, k, n, dtype, dtype))
+                    else:
+                        print(
+                            f"Warning: Unknown dtype '{dtype_str}' in {csv_file}, skipping row"
+                        )
+
+            print(f"Loaded {len(shapes)} shapes from {csv_file}")
+        except Exception as e:
+            print(f"Error loading shapes from {csv_file}: {e}")
+
+        return shapes
+
+    def setup_shape_collections(self, csv_files: list[str]) -> None:
+        """Setup additional shape collections from CSV files and built-in collection."""
+        self.additional_shape_collections = []
+
+        # Load shapes from provided CSV files first
+        for csv_file in csv_files:
+            shapes = self.load_shapes_from_csv(csv_file)
+            if shapes:
+                self.additional_shape_collections.append(shapes)
+
+        self.additional_shape_collections.append(collect_known_mm_shapes())
+        self.shape_generator = self.generate_mm_shapes()
+
+    def generate_mm_shapes(self) -> Generator[tuple[int, int, int, Any], None, None]:
+        """Generator that yields (m, k, n, dtype) tuples for matrix multiplication.
+
+        First exhausts all shapes from additional_shape_collections, then generates random shapes.
+        Only yields unaligned shapes since external CSV shapes may not be pre-filtered.
+        """
+        # Phase 1: Use all shapes from additional shape collections
+        for collection in self.additional_shape_collections:
+            for m, k, n, dtype1, _ in collection:
+                # Filter for unaligned shapes only (external CSVs may not be pre-filtered)
+                align_size = get_alignment_size_dtype(dtype1)
+                if not all(self.is_aligned(dim, align_size) for dim in [m, k, n]):
+                    # Check if it fits in memory
+                    if fits_in_memory(dtype1, m, k, n):
+                        yield (m, k, n, dtype1)
+
+        # Phase 2: Generate infinite random shapes
+
+        while True:
+            # Generate random dtype
+            dtype_choices = [torch.float16, torch.bfloat16, torch.float32]
+            dtype = random.choices(dtype_choices)[0]
+
+            # Generate random shape for this dtype
+            uniform = random.choices([True, False])[0]
+            align_size = get_alignment_size_dtype(dtype)
+
+            # Keep trying until we get a valid unaligned shape that fits in memory
+            while True:
+                if uniform:
+                    m = random.randint(1, 65536)
+                    k = random.randint(1, 65536)
+                    n = random.randint(1, 65536)
+                else:
+                    m = self.get_random_dim()
+                    k = self.get_random_dim()
+                    n = self.get_random_dim()
+
+                # Skip if all dimensions are aligned (we need unaligned for padding to be relevant)
+                if all(self.is_aligned(dim, align_size) for dim in [m, k, n]):
+                    continue
+
+                # Check if it fits in memory
+                if fits_in_memory(dtype, m, k, n):
+                    yield (m, k, n, dtype)
+                    break
+
     def create_input(self) -> tuple[Any, ...]:
-        dtype = self.get_dtype()
+        # Get the next shape from the generator
+        m, k, n, dtype = next(self.shape_generator)
         set_precision(dtype)
-        m, k, n = self.get_m_k_n(dtype)
 
         (transpose_left, transpose_right) = transpose_tensors()
         prepadded_left = self.prepadded()
@@ -107,40 +224,45 @@ def get_random_dim(
             return 2 ** random.randint(min_power2, max_power2)  # type: ignore[no-any-return]
         else:
             # choose a random number between 2^i and 2^(i+1)
-            return self.get_random_between_pow2(min_power2, max_power2)  # type: ignore[no-any-return]
+            return get_random_between_pow2(min_power2, max_power2)  # type: ignore[no-any-return]
 
     def is_aligned(self, dim: int, align_size: int) -> bool:
         return dim % align_size == 0
 
-    def get_m_k_n(self, dtype: Any) -> tuple[int, int, int]:
-        uniform = random.choices([True, False])[0]
-        align_size = get_alignment_size_dtype(dtype)
+    def prepadded(self, p_prepadded: float = 0.2) -> bool:
+        # p_prepadded: probability that a tensor is "prepadded", i.e. pad_mm excludes time it takes to pad from benchmarking
+        return random.choices([True, False], [p_prepadded, 1 - p_prepadded])[0]
+
+    def run(self) -> None:
+        """Override run to setup shape collections before running."""
+        import time
 
-        # repeat until tensors fit in memory
-        while True:
-            if uniform:
-                m = random.randint(1, 65536)
-                k = random.randint(1, 65536)
-                n = random.randint(1, 65536)
-            else:
-                m = self.get_random_dim()
-                k = self.get_random_dim()
-                n = self.get_random_dim()
+        from tqdm import tqdm
 
-            if all(self.is_aligned(dim, align_size) for dim in [m, k, n]):
-                # skip if already aligned
-                continue
+        torch.set_default_device("cuda")
+        args = self.parser.parse_args()
 
-            if fits_in_memory(dtype, m, k, n):
-                return (m, k, n)
+        # Setup shape collections based on CLI arguments
+        self.setup_shape_collections(args.additional_shape_csv)
 
-    def prepadded(self, p_prepadded: float = 0.2) -> bool:
-        # p_prepadded: probability that a tensor is "prepadded", i.e. pad_mm excludes time it takes to pad from benchmarking
-        return random.choices([True, False], [p_prepadded, 1 - p_prepadded])[0]
+        # Set up torch configuration (copied from parent run method)
+
+        if args.use_heuristic:
+            torch._inductor.config.autoheuristic_use.pad_mm = True
+            torch._inductor.config.autoheuristic_collect.pad_mm = False
+        else:
+            torch._inductor.config.autoheuristic_use.pad_mm = False
+            torch._inductor.config.autoheuristic_collect.pad_mm = True
+        torch._inductor.config.autoheuristic_log_path = args.o
+        if args.device is not None:
+            torch.cuda.set_device(args.device)
+        random.seed(time.time())
 
-    def get_dtype(self) -> Any:
-        dtype_choices = [torch.float16, torch.bfloat16, torch.float32]
-        return random.choices(dtype_choices)[0]
+        # Run the main benchmarking loop
+        for _ in tqdm(range(args.num_samples)):
+            input = self.create_input()
+            for _ in range(args.num_reps):
+                self.run_benchmark(*input)
 
 
 if __name__ == "__main__":
diff --git a/torchgen/aoti/fallback_ops.py b/torchgen/aoti/fallback_ops.py
index 93a2b3362efe2..84d413d426d19 100644
--- a/torchgen/aoti/fallback_ops.py
+++ b/torchgen/aoti/fallback_ops.py
@@ -36,16 +36,18 @@
     "aten._fft_c2c.default": {},
     "aten._fft_r2c.default": {},
     "aten._flash_attention_backward.default": {},
-    "aten._flash_attention_forward.default": {},
+    "aten._flash_attention_forward.default": {"v2": ["block_table", "num_splits"]},
+    "aten._flash_attention_forward_no_dropout_inplace.default": {"v2": ["num_splits"]},
     "aten._flash_attention_forward.quantized": {},
     "aten._fused_moving_avg_obs_fq_helper_functional.default": {},
     "aten._fused_moving_avg_obs_fq_helper.default": {},
     "aten._fused_rms_norm.default": {},
+    "aten._grouped_mm.default": {},
     "aten._histogramdd_from_bin_cts.default": {},
     "aten._int_mm.out": {},
     "aten._pdist_backward.default": {},
     "aten._pdist_forward.default": {},
-    "aten._scaled_dot_product_attention_math_for_mps.default": {},
+    "aten._scaled_dot_product_attention_math_for_mps.default": {"v2": ["enable_gqa"]},
     "aten._scaled_dot_product_cudnn_attention_backward.default": {},
     "aten._scaled_dot_product_cudnn_attention.default": {},
     "aten._scaled_dot_product_efficient_attention_backward.default": {},
@@ -150,6 +152,12 @@
     "aten.randn.default": {},
     "aten.randn.generator": {},
     "aten.randperm.default": {},
+    "aten.rand_like.default": {},
+    "aten.rand_like.generator": {},
+    "aten.randint_like.default": {},
+    "aten.randint_like.low_dtype": {},
+    "aten.randn_like.default": {},
+    "aten.randn_like.generator": {},
     "aten.repeat_interleave.Tensor": {},
     "aten.replication_pad1d_backward.default": {},
     "aten.replication_pad2d_backward.default": {},
diff --git a/torchgen/api/python.py b/torchgen/api/python.py
index ca971e854b234..254d7c1ee9b43 100644
--- a/torchgen/api/python.py
+++ b/torchgen/api/python.py
@@ -288,7 +288,7 @@ def argument_str_pyi(
             name += "_"
 
         # pyi merges the _out and functional variants into the same signature, with an optional out arg
-        if name == "out" and type_str == "Tensor" and not deprecated:
+        if name == "out" and not deprecated:
             type_str = f"{type_str} | None".replace(" | None | None", " | None")
 
         # pyi deprecated signatures don't get defaults for their out arg
@@ -975,14 +975,17 @@ def argument_type_str_pyi(t: Type) -> str:
         if str(t.elem) == "int":
             ret = "_int | _size" if t.size is not None else "_size"
         elif t.is_tensor_like():
-            # TODO: this doesn't seem right...
-            # Tensor?[] currently translates to tuple[Tensor, ...] | list[Tensor] | None
-            # It should probably translate to   tuple[Tensor | None, ...] | list[Tensor | None]
-            add_optional = True
+            # Tensor?[] translates to tuple[Tensor | None, ...] | list[Tensor | None] | None
+            # Tensor[] translates to tuple[Tensor, ...] | list[Tensor]
+            if isinstance(t.elem, OptionalType):
+                add_optional = True
+                elem_str = "Tensor | None"
+            else:
+                elem_str = "Tensor"
             ret = (
-                "Tensor | tuple[Tensor, ...] | list[Tensor]"
+                f"Tensor | tuple[{elem_str}, ...] | list[{elem_str}]"
                 if t.size is not None
-                else "tuple[Tensor, ...] | list[Tensor]"
+                else f"tuple[{elem_str}, ...] | list[{elem_str}]"
             )
         elif str(t.elem) == "float":
             ret = "Sequence[_float]"
diff --git a/torchgen/api/types/types_base.py b/torchgen/api/types/types_base.py
index 08085fa0fa2bf..322ae1c39c1ed 100644
--- a/torchgen/api/types/types_base.py
+++ b/torchgen/api/types/types_base.py
@@ -18,7 +18,7 @@
 from abc import ABC, abstractmethod
 from dataclasses import dataclass
 from enum import auto, Enum
-from typing import TYPE_CHECKING, Union
+from typing import TYPE_CHECKING
 
 
 if TYPE_CHECKING:
@@ -35,7 +35,7 @@ class SpecialArgName(Enum):
     possibly_redundant_memory_format = auto()
 
 
-ArgName = Union[str, SpecialArgName]
+ArgName = str | SpecialArgName
 
 
 # This class shouldn't be created directly; instead, use/create one of the singletons below.
diff --git a/torchgen/gen.py b/torchgen/gen.py
index e8147a86d7ad5..10f727b9dfba0 100644
--- a/torchgen/gen.py
+++ b/torchgen/gen.py
@@ -2101,6 +2101,7 @@ def gen_headers(
     static_dispatch_idx: list[BackendIndex],
     selector: SelectiveBuilder,
     backend_indices: dict[DispatchKey, BackendIndex],
+    headeronly_fm: FileManager,
     core_fm: FileManager,
     cpu_fm: FileManager,
     device_fms: dict[str, FileManager],
@@ -2227,7 +2228,7 @@ def gen_aten_interned_strings() -> dict[str, str]:
     def gen_tags_enum() -> dict[str, str]:
         return {"enum_of_valid_tags": (",\n".join(sorted(valid_tags)))}
 
-    core_fm.write("enum_tag.h", gen_tags_enum)
+    headeronly_fm.write("enum_tag.h", gen_tags_enum)
 
 
 def gen_source_files(
@@ -2782,6 +2783,14 @@ def main() -> None:
         help="output directory for AOTInductor shim",
         default="torch/csrc/inductor/aoti_torch/generated",
     )
+    parser.add_argument(
+        "--headeronly-install-dir",
+        "--headeronly_install_dir",
+        help="output directory for header-only generated files (e.g. enum_tag.h). "
+        "Defaults to `<install-dir>/core` when --install-dir is set, otherwise "
+        "`build/torch/headeronly/core`.",
+        default=None,
+    )
     parser.add_argument(
         "--rocm",
         action="store_true",
@@ -2979,12 +2988,23 @@ def main() -> None:
     aoti_install_dir = f"{options.aoti_install_dir}"
     Path(aoti_install_dir).mkdir(parents=True, exist_ok=True)
 
+    if options.headeronly_install_dir is not None:
+        headeronly_install_dir = options.headeronly_install_dir
+    elif options.install_dir is not None:
+        headeronly_install_dir = f"{options.install_dir}/core"
+    else:
+        headeronly_install_dir = "build/torch/headeronly/core"
+    Path(headeronly_install_dir).mkdir(parents=True, exist_ok=True)
+
     core_fm = make_file_manager(options=options, install_dir=core_install_dir)
     cpu_fm = make_file_manager(options=options)
     cpu_vec_fm = make_file_manager(options=options)
     cuda_fm = make_file_manager(options=options)
     ops_fm = make_file_manager(options=options, install_dir=ops_install_dir)
     aoti_fm = make_file_manager(options=options, install_dir=aoti_install_dir)
+    headeronly_fm = make_file_manager(
+        options=options, install_dir=headeronly_install_dir
+    )
     device_fms = {"cuda": cuda_fm}
     if options.xpu:
         device_fms["xpu"] = make_file_manager(options=options)
@@ -3034,6 +3054,7 @@ def main() -> None:
             static_dispatch_idx=static_dispatch_idx,
             selector=selector,
             backend_indices=backend_indices,
+            headeronly_fm=headeronly_fm,
             core_fm=core_fm,
             cpu_fm=cpu_fm,
             device_fms=device_fms,
diff --git a/torchgen/gen_aoti_c_shim.py b/torchgen/gen_aoti_c_shim.py
index 3f626955108f6..dde6a6c8eda98 100644
--- a/torchgen/gen_aoti_c_shim.py
+++ b/torchgen/gen_aoti_c_shim.py
@@ -214,7 +214,16 @@ def gen_arguments(
     callsite_exprs: list[str] = []
     for arg in flat_arguments:
         if arg.name in skipped_args:
-            callsite_exprs.append("std::nullopt")
+            # Pass the arg's schema default when available (e.g. "false" for
+            # a bool arg with default=False), so non-optional args with defaults
+            # can be versioned too. Fall back to std::nullopt for optional args
+            # with no default (matches historical behavior).
+            if arg.default is not None:
+                from torchgen.api.cpp import default_expr
+
+                callsite_exprs.append(default_expr(arg.default, arg.type, symint=False))
+            else:
+                callsite_exprs.append("std::nullopt")
             continue
         new_types, names, _, new_callsite_exprs = convert_arg_type_and_name(
             arg.type, arg.name, arg.is_write
diff --git a/torchgen/gen_backend_stubs.py b/torchgen/gen_backend_stubs.py
index 6580613a60f4b..efe63a80249eb 100644
--- a/torchgen/gen_backend_stubs.py
+++ b/torchgen/gen_backend_stubs.py
@@ -549,9 +549,13 @@ def gen_dispatcher_registrations(
 def run(
     source_yaml: str, output_dir: str, dry_run: bool, impl_path: str | None = None
 ) -> None:
-    # Assumes that this file lives at PYTORCH_ROOT/torchgen/gen_backend_stubs.py
-    pytorch_root = Path(__file__).absolute().parent.parent
-    template_dir = os.path.join(pytorch_root, "aten/src/ATen/templates")
+    # Assumes that this file lives at torchgen/gen_backend_stubs.py
+    root = Path(__file__).absolute().parent.parent
+    common_dir = os.path.join(root, "aten/src")  # Assumes root is pytorch_root
+    if not os.path.exists(common_dir):  # This file is out-of-tree.
+        common_dir = os.path.join(root, "torchgen/packaged")
+
+    template_dir = os.path.join(common_dir, "ATen/templates")
 
     def make_file_manager(install_dir: str) -> FileManager:
         return FileManager(
@@ -560,10 +564,8 @@ def make_file_manager(install_dir: str) -> FileManager:
 
     fm = make_file_manager(output_dir)
 
-    native_yaml_path = os.path.join(
-        pytorch_root, "aten/src/ATen/native/native_functions.yaml"
-    )
-    tags_yaml_path = os.path.join(pytorch_root, "aten/src/ATen/native/tags.yaml")
+    native_yaml_path = os.path.join(common_dir, "ATen/native/native_functions.yaml")
+    tags_yaml_path = os.path.join(common_dir, "ATen/native/tags.yaml")
     parsed_yaml = parse_native_yaml(native_yaml_path, tags_yaml_path)
     native_functions, backend_indices = (
         parsed_yaml.native_functions,
diff --git a/torchgen/gen_functionalization_type.py b/torchgen/gen_functionalization_type.py
index ffd0fc3d45fad..251ba64248a3c 100644
--- a/torchgen/gen_functionalization_type.py
+++ b/torchgen/gen_functionalization_type.py
@@ -37,6 +37,7 @@
     NativeFunction,
     NativeFunctionsGroup,
     NativeFunctionsViewGroup,
+    OperatorName,
     Return,
     SchemaKind,
     SelfArgument,
@@ -71,9 +72,21 @@
         "resize_as_",
         # This function is used as for testing purposes only.
         "_fill_mem_eff_dropout_mask_",
+        # Inference-only op called behind a custom op graph break.
+        "_flash_attention_forward_no_dropout_inplace",
     ]
 )
 
+# Eager cumulative out variants compute in the out dtype when dtype is omitted.
+# Functionalization normally lowers mutable ops through their functional variants,
+# so these need to thread the out dtype explicitly to preserve eager semantics.
+CUMULATIVE_OUT_OPS_PRESERVING_OUT_DTYPE = {
+    OperatorName.parse("cumsum.out"),
+    OperatorName.parse("cumprod.out"),
+    OperatorName.parse("cumsum.dimname_out"),
+    OperatorName.parse("cumprod.dimname_out"),
+}
+
 # This file contains codegen that relates to the functionalization pass.
 # It includes:
 # - gen_functionalization_definition
@@ -605,6 +618,40 @@ def wrap_propagate_mutations_and_return(
     {returns_str}"""
 
 
+def maybe_replace_cumulative_out_dtype_exprs(
+    f: NativeFunction,
+    functional_sig: DispatcherSignature,
+    functional_exprs: list[str],
+) -> list[str]:
+    if (
+        f.func.kind() != SchemaKind.out
+        or f.func.name not in CUMULATIVE_OUT_OPS_PRESERVING_OUT_DTYPE
+    ):
+        return functional_exprs
+
+    if len(f.func.arguments.out) != 1:
+        raise AssertionError(
+            f"Expected a single out argument for cumulative out op: {f.func.name}"
+        )
+
+    dtype_arg_idx = next(
+        (i for i, arg in enumerate(functional_sig.arguments()) if arg.name == "dtype"),
+        None,
+    )
+    if dtype_arg_idx is None:
+        raise AssertionError(
+            f"Expected dtype argument for cumulative out op: {f.func.name}"
+        )
+
+    adjusted_exprs = functional_exprs.copy()
+    dtype_expr = adjusted_exprs[dtype_arg_idx]
+    adjusted_exprs[dtype_arg_idx] = (
+        f"{dtype_expr}.has_value() ? {dtype_expr} : "
+        f"std::optional<at::ScalarType>({f.func.arguments.out[0].name}_.scalar_type())"
+    )
+    return adjusted_exprs
+
+
 # Generates the Functionalization kernel for:
 # - mutation ops (inplace and out= ops)
 @with_native_function_and
@@ -676,6 +723,9 @@ def emit_inplace_functionalization_body(
         e.expr
         for e in translate(unwrapped_args_ctx, functional_sig.arguments(), method=False)
     ]
+    functional_exprs = maybe_replace_cumulative_out_dtype_exprs(
+        f, functional_sig, functional_exprs
+    )
 
     meta_conversion_str, meta_call_ctx = convert_to_meta_tensors(dispatcher_sig)
     # We don't want to run the inplace meta func for ops like .set_(), because:
diff --git a/torchgen/static_runtime/generator.py b/torchgen/static_runtime/generator.py
index e15b2514830ac..d8aba4d13bcde 100644
--- a/torchgen/static_runtime/generator.py
+++ b/torchgen/static_runtime/generator.py
@@ -270,7 +270,7 @@ def is_supported(g: NativeFunctionsGroup | NativeFunctionsViewGroup) -> bool:
         # the string, just test the dang thing directly
         if "at::Tensor" != cpp.returns_type(func.returns, symint=False).cpp_type():
             # Returns a non-Tensor value.
-            logger.info("NON-TENSOR RET TYPE: %s", str(func))
+            logger.info("NON-TENSOR RET TYPE: %s", func)
             return False
         return True
 
diff --git a/version.txt b/version.txt
index f925b7d0ce58a..d8b698973a491 100644
--- a/version.txt
+++ b/version.txt
@@ -1 +1 @@
-2.11.0a0
+2.12.0